1 /* 2 Implements the Kokkos kernel 3 */ 4 #include <petscconf.h> 5 #if defined(PETSC_HAVE_CUDA_CLANG) 6 #include <petsclandau.h> 7 #define LANDAU_NOT_IMPLEMENTED SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "Not supported with CLANG") 8 PetscErrorCode LandauKokkosJacobian(DM[], const PetscInt, const PetscInt, const PetscInt, const PetscInt[], PetscReal[], PetscScalar[], const PetscScalar[], const LandauStaticData *, const PetscReal, const PetscLogEvent[], const PetscInt[], const PetscInt[], Mat[], Mat) 9 { 10 LANDAU_NOT_IMPLEMENTED; 11 } 12 PetscErrorCode LandauKokkosCreateMatMaps(P4estVertexMaps *, pointInterpolationP4est (*)[LANDAU_MAX_Q_FACE], PetscInt[], PetscInt, PetscInt) 13 { 14 LANDAU_NOT_IMPLEMENTED; 15 } 16 PetscErrorCode LandauKokkosDestroyMatMaps(P4estVertexMaps *, PetscInt) 17 { 18 LANDAU_NOT_IMPLEMENTED; 19 } 20 PetscErrorCode LandauKokkosStaticDataSet(DM, const PetscInt, const PetscInt, const PetscInt, PetscInt[], PetscInt[], PetscInt[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], LandauStaticData *) 21 { 22 LANDAU_NOT_IMPLEMENTED; 23 } 24 PetscErrorCode LandauKokkosStaticDataClear(LandauStaticData *) 25 { 26 LANDAU_NOT_IMPLEMENTED; 27 } 28 #else 29 #include <petscvec_kokkos.hpp> 30 #include <petsclandau.h> 31 #include <petsc/private/dmpleximpl.h> /*I "petscdmplex.h" I*/ 32 #include <petsc/private/deviceimpl.h> 33 #include <petscts.h> 34 35 #include <Kokkos_Core.hpp> 36 #include <cstdio> 37 typedef Kokkos::TeamPolicy<>::member_type team_member; 38 #include "../land_tensors.h" 39 #include <petscaijdevice.h> 40 41 namespace landau_inner_red 42 { // namespace helps with name resolution in reduction identity 43 template <class ScalarType> 44 struct array_type { 45 ScalarType gg2[LANDAU_DIM]; 46 ScalarType gg3[LANDAU_DIM][LANDAU_DIM]; 47 48 KOKKOS_INLINE_FUNCTION // Default constructor - Initialize to 0's 49 array_type() 50 { 51 for (int j = 0; j < LANDAU_DIM; j++) { 52 gg2[j] = 0; 53 for (int k = 0; k < LANDAU_DIM; k++) gg3[j][k] = 0; 54 } 55 } 56 KOKKOS_INLINE_FUNCTION // Copy Constructor 57 array_type(const array_type &rhs) 58 { 59 for (int j = 0; j < LANDAU_DIM; j++) { 60 gg2[j] = rhs.gg2[j]; 61 for (int k = 0; k < LANDAU_DIM; k++) gg3[j][k] = rhs.gg3[j][k]; 62 } 63 } 64 KOKKOS_INLINE_FUNCTION // add operator 65 array_type & 66 operator+=(const array_type &src) 67 { 68 for (int j = 0; j < LANDAU_DIM; j++) { 69 gg2[j] += src.gg2[j]; 70 for (int k = 0; k < LANDAU_DIM; k++) gg3[j][k] += src.gg3[j][k]; 71 } 72 return *this; 73 } 74 KOKKOS_INLINE_FUNCTION // volatile add operator 75 void 76 operator+=(const volatile array_type &src) volatile 77 { 78 for (int j = 0; j < LANDAU_DIM; j++) { 79 gg2[j] += src.gg2[j]; 80 for (int k = 0; k < LANDAU_DIM; k++) gg3[j][k] += src.gg3[j][k]; 81 } 82 } 83 }; 84 typedef array_type<PetscReal> TensorValueType; // used to simplify code below 85 } // namespace landau_inner_red 86 87 namespace Kokkos 88 { //reduction identity must be defined in Kokkos namespace 89 template <> 90 struct reduction_identity<landau_inner_red::TensorValueType> { 91 KOKKOS_FORCEINLINE_FUNCTION static landau_inner_red::TensorValueType sum() { return landau_inner_red::TensorValueType(); } 92 }; 93 } // namespace Kokkos 94 95 extern "C" { 96 PetscErrorCode LandauKokkosCreateMatMaps(P4estVertexMaps maps[], pointInterpolationP4est (*pointMaps)[LANDAU_MAX_Q_FACE], PetscInt Nf[], PetscInt Nq, PetscInt grid) 97 { 98 P4estVertexMaps h_maps; /* host container */ 99 const Kokkos::View<pointInterpolationP4est *[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_points((pointInterpolationP4est *)pointMaps, maps[grid].num_reduced); 100 const Kokkos::View<LandauIdx *[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ], Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_gidx((LandauIdx *)maps[grid].gIdx, maps[grid].num_elements); 101 Kokkos::View<pointInterpolationP4est *[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight> *d_points = new Kokkos::View<pointInterpolationP4est *[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight>("points", maps[grid].num_reduced); 102 Kokkos::View<LandauIdx *[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ], Kokkos::LayoutRight> *d_gidx = new Kokkos::View<LandauIdx *[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ], Kokkos::LayoutRight>("gIdx", maps[grid].num_elements); 103 104 PetscFunctionBegin; 105 Kokkos::deep_copy(*d_gidx, h_gidx); 106 Kokkos::deep_copy(*d_points, h_points); 107 h_maps.num_elements = maps[grid].num_elements; 108 h_maps.num_face = maps[grid].num_face; 109 h_maps.num_reduced = maps[grid].num_reduced; 110 h_maps.deviceType = maps[grid].deviceType; 111 h_maps.numgrids = maps[grid].numgrids; 112 h_maps.Nf = Nf[grid]; 113 h_maps.c_maps = (pointInterpolationP4est(*)[LANDAU_MAX_Q_FACE])d_points->data(); 114 maps[grid].vp1 = (void *)d_points; 115 h_maps.gIdx = (LandauIdx(*)[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ])d_gidx->data(); 116 maps[grid].vp2 = (void *)d_gidx; 117 { 118 Kokkos::View<P4estVertexMaps, Kokkos::HostSpace> h_maps_k(&h_maps); 119 Kokkos::View<P4estVertexMaps> *d_maps_k = new Kokkos::View<P4estVertexMaps>(Kokkos::create_mirror(Kokkos::DefaultExecutionSpace::memory_space(), h_maps_k)); 120 Kokkos::deep_copy(*d_maps_k, h_maps_k); 121 maps[grid].d_self = d_maps_k->data(); 122 maps[grid].vp3 = (void *)d_maps_k; 123 } 124 PetscFunctionReturn(PETSC_SUCCESS); 125 } 126 PetscErrorCode LandauKokkosDestroyMatMaps(P4estVertexMaps maps[], PetscInt num_grids) 127 { 128 PetscFunctionBegin; 129 for (PetscInt grid = 0; grid < num_grids; grid++) { 130 auto a = static_cast<Kokkos::View<pointInterpolationP4est *[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight> *>(maps[grid].vp1); 131 auto b = static_cast<Kokkos::View<LandauIdx *[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ], Kokkos::LayoutRight> *>(maps[grid].vp2); 132 auto c = static_cast<Kokkos::View<P4estVertexMaps *> *>(maps[grid].vp3); 133 delete a; 134 delete b; 135 delete c; 136 } 137 PetscFunctionReturn(PETSC_SUCCESS); 138 } 139 140 PetscErrorCode LandauKokkosStaticDataSet(DM plex, const PetscInt Nq, const PetscInt batch_sz, const PetscInt num_grids, PetscInt a_numCells[], PetscInt a_species_offset[], PetscInt a_mat_offset[], PetscReal a_nu_alpha[], PetscReal a_nu_beta[], PetscReal a_invMass[], PetscReal a_invJ[], PetscReal a_x[], PetscReal a_y[], PetscReal a_z[], PetscReal a_w[], LandauStaticData *SData_d) 141 { 142 PetscReal *BB, *DD; 143 PetscTabulation *Tf; 144 PetscInt dim; 145 PetscInt Nb = Nq, ip_offset[LANDAU_MAX_GRIDS + 1], ipf_offset[LANDAU_MAX_GRIDS + 1], elem_offset[LANDAU_MAX_GRIDS + 1], nip, IPf_sz, Nftot; 146 PetscDS prob; 147 148 PetscFunctionBegin; 149 PetscCall(DMGetDimension(plex, &dim)); 150 PetscCall(DMGetDS(plex, &prob)); 151 PetscCheck(LANDAU_DIM == dim, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "dim %" PetscInt_FMT " != LANDAU_DIM %d", dim, LANDAU_DIM); 152 PetscCall(PetscDSGetTabulation(prob, &Tf)); 153 BB = Tf[0]->T[0]; 154 DD = Tf[0]->T[1]; 155 ip_offset[0] = ipf_offset[0] = elem_offset[0] = 0; 156 nip = 0; 157 IPf_sz = 0; 158 for (PetscInt grid = 0; grid < num_grids; grid++) { 159 PetscInt nfloc = a_species_offset[grid + 1] - a_species_offset[grid]; 160 elem_offset[grid + 1] = elem_offset[grid] + a_numCells[grid]; 161 nip += a_numCells[grid] * Nq; 162 ip_offset[grid + 1] = nip; 163 IPf_sz += Nq * nfloc * a_numCells[grid]; 164 ipf_offset[grid + 1] = IPf_sz; 165 } 166 Nftot = a_species_offset[num_grids]; 167 PetscCall(PetscKokkosInitializeCheck()); 168 { 169 const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_alpha(a_nu_alpha, Nftot); 170 auto alpha = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("alpha", Nftot); 171 SData_d->alpha = static_cast<void *>(alpha); 172 const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_beta(a_nu_beta, Nftot); 173 auto beta = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("beta", Nftot); 174 SData_d->beta = static_cast<void *>(beta); 175 const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_invMass(a_invMass, Nftot); 176 auto invMass = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("invMass", Nftot); 177 SData_d->invMass = static_cast<void *>(invMass); 178 const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_BB(BB, Nq * Nb); 179 auto B = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("B", Nq * Nb); 180 SData_d->B = static_cast<void *>(B); 181 const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_DD(DD, Nq * Nb * dim); 182 auto D = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("D", Nq * Nb * dim); 183 SData_d->D = static_cast<void *>(D); 184 const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_invJ(a_invJ, nip * dim * dim); 185 auto invJ = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("invJ", nip * dim * dim); 186 SData_d->invJ = static_cast<void *>(invJ); 187 const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_x(a_x, nip); 188 auto x = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("x", nip); 189 SData_d->x = static_cast<void *>(x); 190 const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_y(a_y, nip); 191 auto y = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("y", nip); 192 SData_d->y = static_cast<void *>(y); 193 const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_w(a_w, nip); 194 auto w = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("w", nip); 195 SData_d->w = static_cast<void *>(w); 196 197 Kokkos::deep_copy(*alpha, h_alpha); 198 Kokkos::deep_copy(*beta, h_beta); 199 Kokkos::deep_copy(*invMass, h_invMass); 200 Kokkos::deep_copy(*B, h_BB); 201 Kokkos::deep_copy(*D, h_DD); 202 Kokkos::deep_copy(*invJ, h_invJ); 203 Kokkos::deep_copy(*x, h_x); 204 Kokkos::deep_copy(*y, h_y); 205 Kokkos::deep_copy(*w, h_w); 206 207 if (dim == 3) { 208 const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_z(a_z, nip); 209 auto z = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("z", nip); 210 SData_d->z = static_cast<void *>(z); 211 Kokkos::deep_copy(*z, h_z); 212 } else SData_d->z = NULL; 213 214 // 215 const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_NCells(a_numCells, num_grids); 216 auto nc = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("NCells", num_grids); 217 SData_d->NCells = static_cast<void *>(nc); 218 Kokkos::deep_copy(*nc, h_NCells); 219 220 const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_species_offset(a_species_offset, num_grids + 1); 221 auto soff = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("species_offset", num_grids + 1); 222 SData_d->species_offset = static_cast<void *>(soff); 223 Kokkos::deep_copy(*soff, h_species_offset); 224 225 const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_mat_offset(a_mat_offset, num_grids + 1); 226 auto moff = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("mat_offset", num_grids + 1); 227 SData_d->mat_offset = static_cast<void *>(moff); 228 Kokkos::deep_copy(*moff, h_mat_offset); 229 230 const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_ip_offset(ip_offset, num_grids + 1); 231 auto ipoff = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("ip_offset", num_grids + 1); 232 SData_d->ip_offset = static_cast<void *>(ipoff); 233 Kokkos::deep_copy(*ipoff, h_ip_offset); 234 235 const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_elem_offset(elem_offset, num_grids + 1); 236 auto eoff = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("elem_offset", num_grids + 1); 237 SData_d->elem_offset = static_cast<void *>(eoff); 238 Kokkos::deep_copy(*eoff, h_elem_offset); 239 240 const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_ipf_offset(ipf_offset, num_grids + 1); 241 auto ipfoff = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("ipf_offset", num_grids + 1); 242 SData_d->ipf_offset = static_cast<void *>(ipfoff); 243 Kokkos::deep_copy(*ipfoff, h_ipf_offset); 244 #if defined(LANDAU_LAYOUT_LEFT) // preallocate dynamic data, no copy 245 auto ipfdf_data = new Kokkos::View<PetscReal ***, Kokkos::LayoutLeft>("fdf", batch_sz, dim + 1, IPf_sz); 246 #else 247 auto ipfdf_data = new Kokkos::View<PetscReal ***, Kokkos::LayoutRight>("fdf", batch_sz, dim + 1, IPf_sz); 248 #endif 249 SData_d->ipfdf_data = static_cast<void *>(ipfdf_data); 250 auto Eq_m = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("Eq_m", Nftot); // allocate but do not set, same for whole batch 251 SData_d->Eq_m = static_cast<void *>(Eq_m); 252 const Kokkos::View<LandauIdx *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_coo_elem_offsets((LandauIdx *)SData_d->coo_elem_offsets, SData_d->coo_n_cellsTot + 1); 253 auto coo_elem_offsets = new Kokkos::View<LandauIdx *, Kokkos::LayoutLeft>("coo_elem_offsets", SData_d->coo_n_cellsTot + 1); 254 const Kokkos::View<LandauIdx *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_coo_elem_fullNb((LandauIdx *)SData_d->coo_elem_fullNb, SData_d->coo_n_cellsTot); 255 auto coo_elem_fullNb = new Kokkos::View<LandauIdx *, Kokkos::LayoutLeft>("coo_elem_offsets", SData_d->coo_n_cellsTot); 256 const Kokkos::View<LandauIdx *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_coo_elem_point_offsets((LandauIdx *)SData_d->coo_elem_point_offsets, SData_d->coo_n_cellsTot * (LANDAU_MAX_NQ + 1)); 257 auto coo_elem_point_offsets = new Kokkos::View<LandauIdx *, Kokkos::LayoutLeft>("coo_elem_point_offsets", SData_d->coo_n_cellsTot * (LANDAU_MAX_NQ + 1)); 258 // assign & copy 259 Kokkos::deep_copy(*coo_elem_offsets, h_coo_elem_offsets); 260 Kokkos::deep_copy(*coo_elem_fullNb, h_coo_elem_fullNb); 261 Kokkos::deep_copy(*coo_elem_point_offsets, h_coo_elem_point_offsets); 262 // need to free this now and use pointer space 263 PetscCall(PetscFree3(SData_d->coo_elem_offsets, SData_d->coo_elem_fullNb, SData_d->coo_elem_point_offsets)); 264 SData_d->coo_elem_offsets = static_cast<void *>(coo_elem_offsets); 265 SData_d->coo_elem_fullNb = static_cast<void *>(coo_elem_fullNb); 266 SData_d->coo_elem_point_offsets = static_cast<void *>(coo_elem_point_offsets); 267 auto coo_vals = new Kokkos::View<PetscScalar *, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace>("coo_vals", SData_d->coo_size); 268 SData_d->coo_vals = static_cast<void *>(coo_vals); 269 } 270 SData_d->maps = NULL; // not used 271 PetscFunctionReturn(PETSC_SUCCESS); 272 } 273 274 PetscErrorCode LandauKokkosStaticDataClear(LandauStaticData *SData_d) 275 { 276 PetscFunctionBegin; 277 if (SData_d->alpha) { 278 auto alpha = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->alpha); 279 delete alpha; 280 SData_d->alpha = NULL; 281 auto beta = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->beta); 282 delete beta; 283 auto invMass = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->invMass); 284 delete invMass; 285 auto B = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->B); 286 delete B; 287 auto D = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->D); 288 delete D; 289 auto invJ = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->invJ); 290 delete invJ; 291 auto x = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->x); 292 delete x; 293 auto y = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->y); 294 delete y; 295 if (SData_d->z) { 296 auto z = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->z); 297 delete z; 298 } 299 #if defined(LANDAU_LAYOUT_LEFT) // preallocate dynamic data, no copy 300 auto z = static_cast<Kokkos::View<PetscReal ***, Kokkos::LayoutLeft> *>(SData_d->ipfdf_data); 301 #else 302 auto z = static_cast<Kokkos::View<PetscReal ***, Kokkos::LayoutRight> *>(SData_d->ipfdf_data); 303 #endif 304 delete z; 305 auto w = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->w); 306 delete w; 307 auto Eq_m = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->Eq_m); 308 delete Eq_m; 309 // offset 310 auto nc = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->NCells); 311 delete nc; 312 auto soff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->species_offset); 313 delete soff; 314 auto moff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->mat_offset); 315 delete moff; 316 auto ipoff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->ip_offset); 317 delete ipoff; 318 auto eoff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->elem_offset); 319 delete eoff; 320 auto ipfoff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->ipf_offset); 321 delete ipfoff; 322 auto coo_elem_offsets = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>((void *)SData_d->coo_elem_offsets); 323 delete coo_elem_offsets; 324 auto coo_elem_fullNb = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>((void *)SData_d->coo_elem_fullNb); 325 delete coo_elem_fullNb; 326 auto coo_elem_point_offsets = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>((void *)SData_d->coo_elem_point_offsets); 327 delete coo_elem_point_offsets; 328 SData_d->coo_elem_offsets = NULL; 329 SData_d->coo_elem_point_offsets = NULL; 330 SData_d->coo_elem_fullNb = NULL; 331 auto coo_vals = static_cast<Kokkos::View<PetscScalar *, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace> *>((void *)SData_d->coo_vals); 332 delete coo_vals; 333 } 334 PetscFunctionReturn(PETSC_SUCCESS); 335 } 336 337 #define KOKKOS_SHARED_LEVEL 0 // 0 is shared, 1 is global 338 339 KOKKOS_INLINE_FUNCTION 340 PetscErrorCode landau_mat_assemble(PetscSplitCSRDataStructure d_mat, PetscScalar *coo_vals, const PetscScalar Aij, const PetscInt f, const PetscInt g, const PetscInt Nb, PetscInt moffset, const PetscInt elem, const PetscInt fieldA, const P4estVertexMaps *d_maps, const LandauIdx coo_elem_offsets[], const LandauIdx coo_elem_fullNb[], const LandauIdx (*coo_elem_point_offsets)[LANDAU_MAX_NQ + 1], const PetscInt glb_elem_idx, const PetscInt bid_coo_sz_batch) 341 { 342 PetscInt idx, q, nr, nc, d; 343 const LandauIdx *const Idxs = &d_maps->gIdx[elem][fieldA][0]; 344 PetscScalar row_scale[LANDAU_MAX_Q_FACE] = {0}, col_scale[LANDAU_MAX_Q_FACE] = {0}; 345 if (coo_vals) { // mirror (i,j) in CreateStaticGPUData 346 const int fullNb = coo_elem_fullNb[glb_elem_idx], fullNb2 = fullNb * fullNb; 347 idx = Idxs[f]; 348 if (idx >= 0) { 349 nr = 1; 350 row_scale[0] = 1.; 351 } else { 352 idx = -idx - 1; 353 for (q = 0, nr = 0; q < d_maps->num_face; q++, nr++) { 354 if (d_maps->c_maps[idx][q].gid < 0) break; 355 row_scale[q] = d_maps->c_maps[idx][q].scale; 356 } 357 } 358 idx = Idxs[g]; 359 if (idx >= 0) { 360 nc = 1; 361 col_scale[0] = 1.; 362 } else { 363 idx = -idx - 1; 364 nc = d_maps->num_face; 365 for (q = 0, nc = 0; q < d_maps->num_face; q++, nc++) { 366 if (d_maps->c_maps[idx][q].gid < 0) break; 367 col_scale[q] = d_maps->c_maps[idx][q].scale; 368 } 369 } 370 const int idx0 = bid_coo_sz_batch + coo_elem_offsets[glb_elem_idx] + fieldA * fullNb2 + fullNb * coo_elem_point_offsets[glb_elem_idx][f] + nr * coo_elem_point_offsets[glb_elem_idx][g]; 371 for (int q = 0, idx2 = idx0; q < nr; q++) { 372 for (int d = 0; d < nc; d++, idx2++) coo_vals[idx2] = row_scale[q] * col_scale[d] * Aij; 373 } 374 } else { 375 PetscScalar vals[LANDAU_MAX_Q_FACE * LANDAU_MAX_Q_FACE] = {0}; 376 PetscInt rows[LANDAU_MAX_Q_FACE], cols[LANDAU_MAX_Q_FACE]; 377 idx = Idxs[f]; 378 if (idx >= 0) { 379 nr = 1; 380 rows[0] = idx; 381 row_scale[0] = 1.; 382 } else { 383 idx = -idx - 1; 384 for (q = 0, nr = 0; q < d_maps->num_face; q++, nr++) { 385 if (d_maps->c_maps[idx][q].gid < 0) break; 386 rows[q] = d_maps->c_maps[idx][q].gid; 387 row_scale[q] = d_maps->c_maps[idx][q].scale; 388 } 389 } 390 idx = Idxs[g]; 391 if (idx >= 0) { 392 nc = 1; 393 cols[0] = idx; 394 col_scale[0] = 1.; 395 } else { 396 idx = -idx - 1; 397 for (q = 0, nc = 0; q < d_maps->num_face; q++, nc++) { 398 if (d_maps->c_maps[idx][q].gid < 0) break; 399 cols[q] = d_maps->c_maps[idx][q].gid; 400 col_scale[q] = d_maps->c_maps[idx][q].scale; 401 } 402 } 403 404 for (q = 0; q < nr; q++) rows[q] = rows[q] + moffset; 405 for (q = 0; q < nc; q++) cols[q] = cols[q] + moffset; 406 for (q = 0; q < nr; q++) { 407 for (d = 0; d < nc; d++) vals[q * nc + d] = row_scale[q] * col_scale[d] * Aij; 408 } 409 static_cast<void>(MatSetValuesDevice(d_mat, nr, rows, nc, cols, vals, ADD_VALUES)); 410 } 411 return PETSC_SUCCESS; 412 } 413 414 PetscErrorCode LandauKokkosJacobian(DM plex[], const PetscInt Nq, const PetscInt batch_sz, const PetscInt num_grids, const PetscInt a_numCells[], PetscReal a_Eq_m[], PetscScalar a_elem_closure[], const PetscScalar a_xarray[], const LandauStaticData *SData_d, const PetscReal shift, const PetscLogEvent events[], const PetscInt a_mat_offset[], const PetscInt a_species_offset[], Mat subJ[], Mat JacP) 415 { 416 using scr_mem_t = Kokkos::DefaultExecutionSpace::scratch_memory_space; 417 using real2_scr_t = Kokkos::View<PetscScalar **, Kokkos::LayoutRight, scr_mem_t>; 418 using g2_scr_t = Kokkos::View<PetscReal ***, Kokkos::LayoutRight, scr_mem_t>; 419 using g3_scr_t = Kokkos::View<PetscReal ****, Kokkos::LayoutRight, scr_mem_t>; 420 PetscInt Nb = Nq, dim, num_cells_max, Nf_max, num_cells_batch; 421 int nfaces = 0, vector_size = 512 / Nq; 422 LandauCtx *ctx; 423 PetscReal *d_Eq_m = NULL; 424 PetscScalar *d_vertex_f = NULL; 425 P4estVertexMaps *maps[LANDAU_MAX_GRIDS]; // this gets captured 426 PetscSplitCSRDataStructure d_mat; 427 PetscContainer container; 428 const int conc = Kokkos::DefaultExecutionSpace().concurrency(), openmp = !!(conc < 1000), team_size = (openmp == 0) ? Nq : 1; 429 const PetscInt coo_sz_batch = SData_d->coo_size / batch_sz; // capture 430 auto d_alpha_k = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->alpha); //static data 431 const PetscReal *d_alpha = d_alpha_k->data(); 432 const PetscInt Nftot = d_alpha_k->size(); // total number of species 433 auto d_beta_k = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->beta); 434 const PetscReal *d_beta = d_beta_k->data(); 435 auto d_invMass_k = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->invMass); 436 const PetscReal *d_invMass = d_invMass_k->data(); 437 auto d_B_k = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->B); 438 const PetscReal *d_BB = d_B_k->data(); 439 auto d_D_k = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->D); 440 const PetscReal *d_DD = d_D_k->data(); 441 auto d_invJ_k = *static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->invJ); // use Kokkos vector in kernels 442 auto d_x_k = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->x); //static data 443 const PetscReal *d_x = d_x_k->data(); 444 auto d_y_k = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->y); //static data 445 const PetscReal *d_y = d_y_k->data(); 446 auto d_z_k = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->z); //static data 447 const PetscReal *d_z = (LANDAU_DIM == 3) ? d_z_k->data() : NULL; 448 auto d_w_k = *static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->w); //static data 449 const PetscReal *d_w = d_w_k.data(); 450 // grid offsets - single vertex grid data 451 auto d_numCells_k = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->NCells); 452 const PetscInt *d_numCells = d_numCells_k->data(); 453 auto d_species_offset_k = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->species_offset); 454 const PetscInt *d_species_offset = d_species_offset_k->data(); 455 auto d_mat_offset_k = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->mat_offset); 456 const PetscInt *d_mat_offset = d_mat_offset_k->data(); 457 auto d_ip_offset_k = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->ip_offset); 458 const PetscInt *d_ip_offset = d_ip_offset_k->data(); 459 auto d_ipf_offset_k = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->ipf_offset); 460 const PetscInt *d_ipf_offset = d_ipf_offset_k->data(); 461 auto d_elem_offset_k = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->elem_offset); 462 const PetscInt *d_elem_offset = d_elem_offset_k->data(); 463 #if defined(LANDAU_LAYOUT_LEFT) // preallocate dynamic data, including batched vertices 464 Kokkos::View<PetscReal ***, Kokkos::LayoutLeft> d_fdf_k = *static_cast<Kokkos::View<PetscReal ***, Kokkos::LayoutLeft> *>(SData_d->ipfdf_data); 465 #else 466 Kokkos::View<PetscReal ***, Kokkos::LayoutRight> d_fdf_k = *static_cast<Kokkos::View<PetscReal ***, Kokkos::LayoutRight> *>(SData_d->ipfdf_data); 467 #endif 468 auto d_Eq_m_k = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->Eq_m); // static storage, dynamci data - E(t), copy later, single vertex 469 // COO 470 auto d_coo_elem_offsets_k = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>(SData_d->coo_elem_offsets); 471 LandauIdx *d_coo_elem_offsets = (SData_d->coo_size == 0) ? NULL : d_coo_elem_offsets_k->data(); 472 auto d_coo_elem_fullNb_k = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>(SData_d->coo_elem_fullNb); 473 LandauIdx *d_coo_elem_fullNb = (SData_d->coo_size == 0) ? NULL : d_coo_elem_fullNb_k->data(); 474 auto d_coo_elem_point_offsets_k = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>(SData_d->coo_elem_point_offsets); 475 LandauIdx(*d_coo_elem_point_offsets)[LANDAU_MAX_NQ + 1] = (SData_d->coo_size == 0) ? NULL : (LandauIdx(*)[LANDAU_MAX_NQ + 1]) d_coo_elem_point_offsets_k->data(); 476 auto d_coo_vals_k = static_cast<Kokkos::View<PetscScalar *, Kokkos::LayoutRight> *>(SData_d->coo_vals); 477 PetscScalar *d_coo_vals = (SData_d->coo_size == 0) ? NULL : d_coo_vals_k->data(); 478 479 PetscFunctionBegin; 480 while (vector_size & (vector_size - 1)) vector_size = vector_size & (vector_size - 1); 481 if (vector_size > 16) vector_size = 16; // printf("DEBUG\n"); 482 PetscCall(PetscLogEventBegin(events[3], 0, 0, 0, 0)); 483 PetscCall(DMGetApplicationContext(plex[0], &ctx)); 484 PetscCheck(ctx, PETSC_COMM_SELF, PETSC_ERR_PLIB, "no context"); 485 PetscCall(DMGetDimension(plex[0], &dim)); 486 PetscCheck(LANDAU_DIM == dim, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "dim %" PetscInt_FMT " != LANDAU_DIM %d", dim, LANDAU_DIM); 487 if (ctx->gpu_assembly) { 488 PetscCall(PetscObjectQuery((PetscObject)JacP, "assembly_maps", (PetscObject *)&container)); 489 if (container) { 490 P4estVertexMaps *h_maps = NULL; 491 PetscCall(PetscContainerGetPointer(container, (void **)&h_maps)); 492 for (PetscInt grid = 0; grid < num_grids; grid++) { 493 if (h_maps[grid].d_self) { 494 maps[grid] = h_maps[grid].d_self; 495 nfaces = h_maps[grid].num_face; // nface=0 for CPU assembly 496 } else { 497 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "GPU assembly but no metadata in container"); 498 } 499 } 500 if (!d_coo_vals) { 501 // this does the setup the first time called 502 PetscCall(MatKokkosGetDeviceMatWrite(JacP, &d_mat)); 503 } else { 504 d_mat = NULL; 505 } 506 } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "GPU assembly but no metadata in container"); 507 } else { 508 for (PetscInt grid = 0; grid < num_grids; grid++) maps[grid] = NULL; 509 nfaces = 0; 510 d_mat = NULL; 511 container = NULL; 512 } 513 num_cells_batch = Nf_max = num_cells_max = 0; 514 for (PetscInt grid = 0; grid < num_grids; grid++) { 515 int Nfloc = a_species_offset[grid + 1] - a_species_offset[grid]; 516 if (Nfloc > Nf_max) Nf_max = Nfloc; 517 if (a_numCells[grid] > num_cells_max) num_cells_max = a_numCells[grid]; 518 num_cells_batch += a_numCells[grid]; // we don't have a host element offset here (but in ctx) 519 } 520 const int elem_mat_num_cells_max_grid = container ? 0 : num_cells_max; 521 #ifdef LAND_SUPPORT_CPU_ASS 522 const int totDim_max = Nf_max * Nq, elem_mat_size_max = totDim_max * totDim_max; 523 Kokkos::View<PetscScalar ****, Kokkos::LayoutRight> d_elem_mats("element matrices", batch_sz, num_grids, elem_mat_num_cells_max_grid, elem_mat_size_max); // not used (cpu assembly) 524 #endif 525 const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_Eq_m_k(a_Eq_m, Nftot); 526 if (a_elem_closure || a_xarray) { 527 Kokkos::deep_copy(*d_Eq_m_k, h_Eq_m_k); 528 d_Eq_m = d_Eq_m_k->data(); 529 } else d_Eq_m = NULL; 530 PetscCall(PetscKokkosInitializeCheck()); 531 PetscCall(PetscLogEventEnd(events[3], 0, 0, 0, 0)); 532 if (a_elem_closure || a_xarray) { // Jacobian, create f & df 533 Kokkos::View<PetscScalar *, Kokkos::LayoutLeft> *d_vertex_f_k = NULL; 534 PetscCall(PetscLogEventBegin(events[1], 0, 0, 0, 0)); 535 if (a_elem_closure) { 536 PetscInt closure_sz = 0; // argh, don't have this on the host!!! 537 for (PetscInt grid = 0; grid < num_grids; grid++) { 538 PetscInt nfloc = a_species_offset[grid + 1] - a_species_offset[grid]; 539 closure_sz += Nq * nfloc * a_numCells[grid]; 540 } 541 closure_sz *= batch_sz; 542 d_vertex_f_k = new Kokkos::View<PetscScalar *, Kokkos::LayoutLeft>("closure", closure_sz); 543 const Kokkos::View<PetscScalar *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_closure_k(a_elem_closure, closure_sz); // Vertex data for each element 544 Kokkos::deep_copy(*d_vertex_f_k, h_closure_k); 545 d_vertex_f = d_vertex_f_k->data(); 546 } else { 547 d_vertex_f = (PetscScalar *)a_xarray; 548 } 549 PetscCall(PetscLogEventEnd(events[1], 0, 0, 0, 0)); 550 PetscCall(PetscLogEventBegin(events[8], 0, 0, 0, 0)); 551 PetscCall(PetscLogGpuTimeBegin()); 552 553 const int scr_bytes_fdf = real2_scr_t::shmem_size(Nf_max, Nb); 554 PetscCall(PetscInfo(plex[0], "df & f shared memory size: %d bytes in level, %d num cells total=%d, team size=%d, vector size=%d, #face=%d, Nf_max=%d\n", scr_bytes_fdf, KOKKOS_SHARED_LEVEL, num_cells_batch * batch_sz, team_size, vector_size, nfaces, Nf_max)); 555 Kokkos::parallel_for( 556 "f, df", Kokkos::TeamPolicy<>(num_cells_batch * batch_sz, team_size, /* Kokkos::AUTO */ vector_size).set_scratch_size(KOKKOS_SHARED_LEVEL, Kokkos::PerTeam(scr_bytes_fdf)), KOKKOS_LAMBDA(const team_member team) { 557 const PetscInt b_Nelem = d_elem_offset[num_grids], b_elem_idx = team.league_rank() % b_Nelem, b_id = team.league_rank() / b_Nelem, IPf_sz_glb = d_ipf_offset[num_grids]; 558 // find my grid 559 PetscInt grid = 0; 560 while (b_elem_idx >= d_elem_offset[grid + 1]) grid++; 561 { 562 const PetscInt loc_nip = d_numCells[grid] * Nq, loc_Nf = d_species_offset[grid + 1] - d_species_offset[grid], loc_elem = b_elem_idx - d_elem_offset[grid]; 563 const PetscInt moffset = LAND_MOFFSET(b_id, grid, batch_sz, num_grids, d_mat_offset); 564 { 565 real2_scr_t s_coef_k(team.team_scratch(KOKKOS_SHARED_LEVEL), Nf_max, Nb); 566 PetscScalar *coef; 567 const PetscReal *invJe = &d_invJ_k((d_ip_offset[grid] + loc_elem * Nq) * dim * dim); 568 // un pack IPData 569 if (!maps[grid]) { 570 coef = &d_vertex_f[b_id * IPf_sz_glb + d_ipf_offset[grid] + loc_elem * Nb * loc_Nf]; // closure and IP indexing are the same 571 } else { 572 coef = s_coef_k.data(); 573 Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, (int)loc_Nf), [=](int f) { 574 //for (int f = 0; f < loc_Nf; ++f) { 575 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, 0, (int)Nb), [=](int b) { 576 //for (int b = 0; b < Nb; ++b) { 577 LandauIdx idx = maps[grid]->gIdx[loc_elem][f][b]; 578 if (idx >= 0) { 579 coef[f * Nb + b] = d_vertex_f[idx + moffset]; // xarray data, not IP, need raw array to deal with CPU assemble case (not used) 580 } else { 581 idx = -idx - 1; 582 coef[f * Nb + b] = 0; 583 for (int q = 0; q < maps[grid]->num_face; q++) { 584 PetscInt id = maps[grid]->c_maps[idx][q].gid; 585 PetscScalar scale = maps[grid]->c_maps[idx][q].scale; 586 if (id >= 0) coef[f * Nb + b] += scale * d_vertex_f[id + moffset]; 587 } 588 } 589 }); 590 }); 591 } 592 team.team_barrier(); 593 Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, Nq), [=](int myQi) { 594 const PetscReal *const invJ = &invJe[myQi * dim * dim]; // b_elem_idx: batch element index 595 const PetscReal *Bq = &d_BB[myQi * Nb], *Dq = &d_DD[myQi * Nb * dim]; 596 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, 0, (int)loc_Nf), [=](int f) { 597 PetscInt b, e, d; 598 PetscReal refSpaceDer[LANDAU_DIM]; 599 const PetscInt idx = d_ipf_offset[grid] + f * loc_nip + loc_elem * Nq + myQi; 600 d_fdf_k(b_id, 0, idx) = 0.0; 601 for (d = 0; d < LANDAU_DIM; ++d) refSpaceDer[d] = 0.0; 602 for (b = 0; b < Nb; ++b) { 603 d_fdf_k(b_id, 0, idx) += Bq[b] * PetscRealPart(coef[f * Nb + b]); 604 for (d = 0; d < dim; ++d) refSpaceDer[d] += Dq[b * dim + d] * PetscRealPart(coef[f * Nb + b]); 605 } 606 for (d = 0; d < dim; ++d) { 607 for (e = 0, d_fdf_k(b_id, d + 1, idx) = 0.0; e < dim; ++e) d_fdf_k(b_id, d + 1, idx) += invJ[e * dim + d] * refSpaceDer[e]; 608 } 609 }); // f 610 }); // myQi 611 } 612 team.team_barrier(); 613 } // 'grid' scope 614 }); // global elems - fdf 615 Kokkos::fence(); 616 PetscCall(PetscLogGpuTimeEnd()); 617 PetscCall(PetscLogEventEnd(events[8], 0, 0, 0, 0)); 618 // Jacobian 619 int maximum_shared_mem_size = 64000; 620 PetscDevice device; 621 622 PetscCall(PetscDeviceGetDefault_Internal(&device)); 623 PetscCall(PetscDeviceGetAttribute(device, PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK, &maximum_shared_mem_size)); 624 const int jac_scr_bytes = 2 * (g2_scr_t::shmem_size(dim, Nf_max, Nq) + g3_scr_t::shmem_size(dim, dim, Nf_max, Nq)); 625 const int jac_shared_level = (jac_scr_bytes > maximum_shared_mem_size) ? 1 : KOKKOS_SHARED_LEVEL; 626 auto jac_lambda = KOKKOS_LAMBDA(const team_member team) 627 { 628 const PetscInt b_Nelem = d_elem_offset[num_grids], b_elem_idx = team.league_rank() % b_Nelem, b_id = team.league_rank() / b_Nelem; 629 // find my grid 630 PetscInt grid = 0; 631 while (b_elem_idx >= d_elem_offset[grid + 1]) grid++; 632 { 633 const PetscInt loc_Nf = d_species_offset[grid + 1] - d_species_offset[grid], loc_elem = b_elem_idx - d_elem_offset[grid]; 634 const PetscInt moffset = LAND_MOFFSET(b_id, grid, batch_sz, num_grids, d_mat_offset); 635 const PetscInt f_off = d_species_offset[grid]; 636 #ifdef LAND_SUPPORT_CPU_ASS 637 const PetscInt totDim = loc_Nf * Nq; 638 #endif 639 g2_scr_t g2(team.team_scratch(jac_shared_level), dim, loc_Nf, Nq); 640 g3_scr_t g3(team.team_scratch(jac_shared_level), dim, dim, loc_Nf, Nq); 641 g2_scr_t gg2(team.team_scratch(jac_shared_level), dim, loc_Nf, Nq); 642 g3_scr_t gg3(team.team_scratch(jac_shared_level), dim, dim, loc_Nf, Nq); 643 // get g2[] & g3[] 644 Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, Nq), [=](int myQi) { 645 using Kokkos::parallel_reduce; 646 const PetscInt jpidx_glb = d_ip_offset[grid] + loc_elem * Nq + myQi; 647 const PetscReal *invJ = &d_invJ_k(jpidx_glb * dim * dim); 648 const PetscReal vj[3] = {d_x[jpidx_glb], d_y[jpidx_glb], d_z ? d_z[jpidx_glb] : 0}, wj = d_w[jpidx_glb]; 649 landau_inner_red::TensorValueType gg_temp; // reduce on part of gg2 and g33 for IP jpidx_g 650 Kokkos::parallel_reduce( 651 Kokkos::ThreadVectorRange(team, (int)d_ip_offset[num_grids]), 652 [=](const int &ipidx, landau_inner_red::TensorValueType &ggg) { 653 const PetscReal wi = d_w[ipidx], x = d_x[ipidx], y = d_y[ipidx]; 654 PetscReal temp1[3] = {0, 0, 0}, temp2 = 0; 655 PetscInt fieldA, d2, d3, f_off_r, grid_r, ipidx_g, nip_loc_r, loc_Nf_r; 656 #if LANDAU_DIM == 2 657 PetscReal Ud[2][2], Uk[2][2], mask = (PetscAbs(vj[0] - x) < 100 * PETSC_SQRT_MACHINE_EPSILON && PetscAbs(vj[1] - y) < 100 * PETSC_SQRT_MACHINE_EPSILON) ? 0. : 1.; 658 LandauTensor2D(vj, x, y, Ud, Uk, mask); 659 #else 660 PetscReal U[3][3], z = d_z[ipidx], mask = (PetscAbs(vj[0]-x) < 100*PETSC_SQRT_MACHINE_EPSILON && PetscAbs(vj[1]-y) < 100*PETSC_SQRT_MACHINE_EPSILON && PetscAbs(vj[2]-z) < 100*PETSC_SQRT_MACHINE_EPSILON) ? 0. : 1.; 661 LandauTensor3D(vj, x, y, z, U, mask); 662 #endif 663 grid_r = 0; 664 while (ipidx >= d_ip_offset[grid_r + 1]) grid_r++; // yuck search for grid 665 f_off_r = d_species_offset[grid_r]; 666 ipidx_g = ipidx - d_ip_offset[grid_r]; 667 nip_loc_r = d_numCells[grid_r] * Nq; 668 loc_Nf_r = d_species_offset[grid_r + 1] - d_species_offset[grid_r]; 669 for (fieldA = 0; fieldA < loc_Nf_r; ++fieldA) { 670 const PetscInt idx = d_ipf_offset[grid_r] + fieldA * nip_loc_r + ipidx_g; 671 temp1[0] += d_fdf_k(b_id, 1, idx) * d_beta[fieldA + f_off_r] * d_invMass[fieldA + f_off_r]; 672 temp1[1] += d_fdf_k(b_id, 2, idx) * d_beta[fieldA + f_off_r] * d_invMass[fieldA + f_off_r]; 673 #if LANDAU_DIM == 3 674 temp1[2] += d_fdf_k(b_id, 3, idx) * d_beta[fieldA + f_off_r] * d_invMass[fieldA + f_off_r]; 675 #endif 676 temp2 += d_fdf_k(b_id, 0, idx) * d_beta[fieldA + f_off_r]; 677 } 678 temp1[0] *= wi; 679 temp1[1] *= wi; 680 #if LANDAU_DIM == 3 681 temp1[2] *= wi; 682 #endif 683 temp2 *= wi; 684 #if LANDAU_DIM == 2 685 for (d2 = 0; d2 < 2; d2++) { 686 for (d3 = 0; d3 < 2; ++d3) { 687 /* K = U * grad(f): g2=e: i,A */ 688 ggg.gg2[d2] += Uk[d2][d3] * temp1[d3]; 689 /* D = -U * (I \kron (fx)): g3=f: i,j,A */ 690 ggg.gg3[d2][d3] += Ud[d2][d3] * temp2; 691 } 692 } 693 #else 694 for (d2 = 0; d2 < 3; ++d2) { 695 for (d3 = 0; d3 < 3; ++d3) { 696 /* K = U * grad(f): g2 = e: i,A */ 697 ggg.gg2[d2] += U[d2][d3]*temp1[d3]; 698 /* D = -U * (I \kron (fx)): g3 = f: i,j,A */ 699 ggg.gg3[d2][d3] += U[d2][d3]*temp2; 700 } 701 } 702 #endif 703 }, 704 Kokkos::Sum<landau_inner_red::TensorValueType>(gg_temp)); 705 // add alpha and put in gg2/3 706 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, (int)loc_Nf), [&](const int &fieldA) { 707 PetscInt d2, d3; 708 for (d2 = 0; d2 < dim; d2++) { 709 gg2(d2, fieldA, myQi) = gg_temp.gg2[d2] * d_alpha[fieldA + f_off]; 710 for (d3 = 0; d3 < dim; d3++) gg3(d2, d3, fieldA, myQi) = -gg_temp.gg3[d2][d3] * d_alpha[fieldA + f_off] * d_invMass[fieldA + f_off]; 711 } 712 }); 713 /* add electric field term once per IP */ 714 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, (int)loc_Nf), [&](const int &fieldA) { gg2(dim - 1, fieldA, myQi) += d_Eq_m[fieldA + f_off]; }); 715 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, (int)loc_Nf), [=](const int &fieldA) { 716 int d, d2, d3, dp; 717 /* Jacobian transform - g2, g3 - per thread (2D) */ 718 for (d = 0; d < dim; ++d) { 719 g2(d, fieldA, myQi) = 0; 720 for (d2 = 0; d2 < dim; ++d2) { 721 g2(d, fieldA, myQi) += invJ[d * dim + d2] * gg2(d2, fieldA, myQi); 722 g3(d, d2, fieldA, myQi) = 0; 723 for (d3 = 0; d3 < dim; ++d3) { 724 for (dp = 0; dp < dim; ++dp) g3(d, d2, fieldA, myQi) += invJ[d * dim + d3] * gg3(d3, dp, fieldA, myQi) * invJ[d2 * dim + dp]; 725 } 726 g3(d, d2, fieldA, myQi) *= wj; 727 } 728 g2(d, fieldA, myQi) *= wj; 729 } 730 }); 731 }); // Nq 732 team.team_barrier(); 733 { /* assemble */ 734 for (PetscInt fieldA = 0; fieldA < loc_Nf; fieldA++) { 735 /* assemble */ 736 Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, Nb), [=](int f) { 737 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, 0, Nb), [=](int g) { 738 PetscScalar t = 0; 739 for (int qj = 0; qj < Nq; qj++) { // look at others integration points 740 const PetscReal *BJq = &d_BB[qj * Nb], *DIq = &d_DD[qj * Nb * dim]; 741 for (int d = 0; d < dim; ++d) { 742 t += DIq[f * dim + d] * g2(d, fieldA, qj) * BJq[g]; 743 for (int d2 = 0; d2 < dim; ++d2) t += DIq[f * dim + d] * g3(d, d2, fieldA, qj) * DIq[g * dim + d2]; 744 } 745 } 746 if (elem_mat_num_cells_max_grid) { // CPU assembly 747 #ifdef LAND_SUPPORT_CPU_ASS 748 const PetscInt fOff = (fieldA * Nb + f) * totDim + fieldA * Nb + g; 749 d_elem_mats(b_id, grid, loc_elem, fOff) = t; 750 #endif 751 } else { 752 static_cast<void>(landau_mat_assemble(d_mat, d_coo_vals, t, f, g, Nb, moffset, loc_elem, fieldA, maps[grid], d_coo_elem_offsets, d_coo_elem_fullNb, d_coo_elem_point_offsets, b_elem_idx, b_id * coo_sz_batch)); 753 } 754 }); 755 }); 756 } 757 } 758 } // scope with 'grid' 759 }; 760 PetscCall(PetscLogEventBegin(events[4], 0, 0, 0, 0)); 761 PetscCall(PetscLogGpuTimeBegin()); 762 PetscCall(PetscInfo(plex[0], "Jacobian shared memory size: %d bytes in level %s (max shared=%d), num cells total=%d, team size=%d, vector size=%d, #face=%d, Nf_max=%d\n", jac_scr_bytes, jac_shared_level == 0 ? "local" : "global", maximum_shared_mem_size, num_cells_batch * batch_sz, team_size, vector_size, nfaces, Nf_max)); 763 Kokkos::parallel_for("Jacobian", Kokkos::TeamPolicy<>(num_cells_batch * batch_sz, team_size, vector_size).set_scratch_size(jac_shared_level, Kokkos::PerTeam(jac_scr_bytes)), jac_lambda); // Kokkos::LaunchBounds<512,2> 764 Kokkos::fence(); 765 PetscCall(PetscLogGpuTimeEnd()); 766 PetscCall(PetscLogEventEnd(events[4], 0, 0, 0, 0)); 767 if (d_vertex_f_k) delete d_vertex_f_k; 768 } else { // mass 769 PetscCall(PetscLogEventBegin(events[16], 0, 0, 0, 0)); 770 PetscCall(PetscLogGpuTimeBegin()); 771 PetscCall(PetscInfo(plex[0], "Mass team size=%d vector size=%d #face=%d Nb=%" PetscInt_FMT ", %s assembly\n", team_size, vector_size, nfaces, Nb, d_coo_vals ? "COO" : "CSR")); 772 Kokkos::parallel_for( 773 "Mass", Kokkos::TeamPolicy<>(num_cells_batch * batch_sz, team_size, vector_size), KOKKOS_LAMBDA(const team_member team) { // Kokkos::LaunchBounds<512,4> 774 const PetscInt b_Nelem = d_elem_offset[num_grids], b_elem_idx = team.league_rank() % b_Nelem, b_id = team.league_rank() / b_Nelem; 775 // find my grid 776 PetscInt grid = 0; 777 while (b_elem_idx >= d_elem_offset[grid + 1]) grid++; 778 { 779 const PetscInt loc_Nf = d_species_offset[grid + 1] - d_species_offset[grid], loc_elem = b_elem_idx - d_elem_offset[grid], jpidx_0 = d_ip_offset[grid] + loc_elem * Nq; 780 #ifdef LAND_SUPPORT_CPU_ASS 781 const PetscInt totDim = loc_Nf * Nq; 782 #endif 783 const PetscInt moffset = LAND_MOFFSET(b_id, grid, batch_sz, num_grids, d_mat_offset); 784 for (int fieldA = 0; fieldA < loc_Nf; fieldA++) { 785 /* assemble */ 786 Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, Nb), [=](int f) { 787 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, 0, Nb), [=](int g) { 788 PetscScalar t = 0; 789 for (int qj = 0; qj < Nq; qj++) { // look at others integration points 790 const PetscReal *BJq = &d_BB[qj * Nb]; 791 const PetscInt jpidx_glb = jpidx_0 + qj; 792 if (dim == 2) { 793 t += BJq[f] * d_w_k(jpidx_glb) * shift * BJq[g] * 2. * PETSC_PI; 794 } else { 795 t += BJq[f] * d_w_k(jpidx_glb) * shift * BJq[g]; 796 } 797 } 798 if (elem_mat_num_cells_max_grid) { 799 #ifdef LAND_SUPPORT_CPU_ASS 800 const PetscInt fOff = (fieldA * Nb + f) * totDim + fieldA * Nb + g; 801 d_elem_mats(b_id, grid, loc_elem, fOff) = t; 802 #endif 803 } else { 804 static_cast<void>(landau_mat_assemble(d_mat, d_coo_vals, t, f, g, Nb, moffset, loc_elem, fieldA, maps[grid], d_coo_elem_offsets, d_coo_elem_fullNb, d_coo_elem_point_offsets, b_elem_idx, b_id * coo_sz_batch)); 805 } 806 }); 807 }); 808 } // field 809 } // grid 810 }); 811 Kokkos::fence(); 812 PetscCall(PetscLogGpuTimeEnd()); 813 PetscCall(PetscLogEventEnd(events[16], 0, 0, 0, 0)); 814 } 815 if (d_coo_vals) { 816 PetscCall(MatSetValuesCOO(JacP, d_coo_vals, ADD_VALUES)); 817 #ifndef LAND_SUPPORT_CPU_ASS 818 Kokkos::fence(); // for timers 819 #endif 820 } else if (elem_mat_num_cells_max_grid) { // CPU assembly 821 #ifdef LAND_SUPPORT_CPU_ASS 822 Kokkos::View<PetscScalar ****, Kokkos::LayoutRight>::HostMirror h_elem_mats = Kokkos::create_mirror_view(d_elem_mats); 823 Kokkos::deep_copy(h_elem_mats, d_elem_mats); 824 PetscCheck(!container, PETSC_COMM_SELF, PETSC_ERR_PLIB, "?????"); 825 for (PetscInt b_id = 0; b_id < batch_sz; b_id++) { // OpenMP (once) 826 for (PetscInt grid = 0; grid < num_grids; grid++) { 827 PetscSection section, globalSection; 828 PetscInt cStart, cEnd; 829 Mat B = subJ[LAND_PACK_IDX(b_id, grid)]; 830 PetscInt moffset = LAND_MOFFSET(b_id, grid, batch_sz, num_grids, a_mat_offset), nloc, nzl, colbuf[1024], row; 831 const PetscInt *cols; 832 const PetscScalar *vals; 833 PetscCall(PetscLogEventBegin(events[5], 0, 0, 0, 0)); 834 PetscCall(DMPlexGetHeightStratum(plex[grid], 0, &cStart, &cEnd)); 835 PetscCall(DMGetLocalSection(plex[grid], §ion)); 836 PetscCall(DMGetGlobalSection(plex[grid], &globalSection)); 837 PetscCall(PetscLogEventEnd(events[5], 0, 0, 0, 0)); 838 PetscCall(PetscLogEventBegin(events[6], 0, 0, 0, 0)); 839 for (PetscInt ej = cStart; ej < cEnd; ++ej) { 840 const PetscScalar *elMat = &h_elem_mats(b_id, grid, ej - cStart, 0); 841 PetscCall(DMPlexMatSetClosure(plex[grid], section, globalSection, B, ej, elMat, ADD_VALUES)); 842 if (grid == 0 && ej == -1) { 843 const PetscInt loc_Nf = a_species_offset[grid + 1] - a_species_offset[grid], totDim = loc_Nf * Nq; 844 int d, f; 845 PetscPrintf(PETSC_COMM_SELF, "Kokkos Element matrix %d/%d\n", 1, (int)a_numCells[grid]); 846 for (d = 0; d < totDim; ++d) { 847 for (f = 0; f < totDim; ++f) PetscPrintf(PETSC_COMM_SELF, " %12.5e", PetscRealPart(elMat[d * totDim + f])); 848 PetscPrintf(PETSC_COMM_SELF, "\n"); 849 } 850 exit(14); 851 } 852 } 853 // move nest matrix to global JacP 854 PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY)); 855 PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY)); 856 PetscCall(MatGetSize(B, &nloc, NULL)); 857 for (int i = 0; i < nloc; i++) { 858 PetscCall(MatGetRow(B, i, &nzl, &cols, &vals)); 859 PetscCheck(nzl <= 1024, PetscObjectComm((PetscObject)B), PETSC_ERR_PLIB, "Row too big: %" PetscInt_FMT, nzl); 860 for (int j = 0; j < nzl; j++) colbuf[j] = cols[j] + moffset; 861 row = i + moffset; 862 PetscCall(MatSetValues(JacP, 1, &row, nzl, colbuf, vals, ADD_VALUES)); 863 PetscCall(MatRestoreRow(B, i, &nzl, &cols, &vals)); 864 } 865 PetscCall(MatDestroy(&subJ[LAND_PACK_IDX(b_id, grid)])); 866 PetscCall(PetscLogEventEnd(events[6], 0, 0, 0, 0)); 867 } // grids 868 } 869 #else 870 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "CPU assembly not supported"); 871 #endif 872 } 873 PetscFunctionReturn(PETSC_SUCCESS); 874 } 875 } // extern "C" 876 #endif 877