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