Actual source code: landau.kokkos.cxx

  1: /*
  2:   Implements the Kokkos kernel
  3: */
  4: #include <petscconf.h>
  5: #include <petscvec_kokkos.hpp>
  6: #include <petsc/private/dmpleximpl.h>
  7: #include <petsc/private/deviceimpl.h>
  8: #include <petsclandau.h>
  9: #include <petscts.h>

 11: #include <Kokkos_Core.hpp>
 12: #include <cstdio>
 13: typedef Kokkos::TeamPolicy<>::member_type team_member;
 14: #include "../land_tensors.h"
 15: #include <petscaijdevice.h>

 17: namespace landau_inner_red
 18: { // namespace helps with name resolution in reduction identity
 19: template <class ScalarType>
 20: struct array_type {
 21:   ScalarType gg2[LANDAU_DIM];
 22:   ScalarType gg3[LANDAU_DIM][LANDAU_DIM];

 24:   KOKKOS_INLINE_FUNCTION // Default constructor - Initialize to 0's
 25:   array_type()
 26:   {
 27:     for (int j = 0; j < LANDAU_DIM; j++) {
 28:       gg2[j] = 0;
 29:       for (int k = 0; k < LANDAU_DIM; k++) gg3[j][k] = 0;
 30:     }
 31:   }
 32:   KOKKOS_INLINE_FUNCTION // Copy Constructor
 33:   array_type(const array_type &rhs)
 34:   {
 35:     for (int j = 0; j < LANDAU_DIM; j++) {
 36:       gg2[j] = rhs.gg2[j];
 37:       for (int k = 0; k < LANDAU_DIM; k++) gg3[j][k] = rhs.gg3[j][k];
 38:     }
 39:   }
 40:   KOKKOS_INLINE_FUNCTION // add operator
 41:     array_type &
 42:     operator+=(const array_type &src)
 43:   {
 44:     for (int j = 0; j < LANDAU_DIM; j++) {
 45:       gg2[j] += src.gg2[j];
 46:       for (int k = 0; k < LANDAU_DIM; k++) gg3[j][k] += src.gg3[j][k];
 47:     }
 48:     return *this;
 49:   }
 50:   KOKKOS_INLINE_FUNCTION // volatile add operator
 51:     void
 52:     operator+=(const volatile array_type &src) volatile
 53:   {
 54:     for (int j = 0; j < LANDAU_DIM; j++) {
 55:       gg2[j] += src.gg2[j];
 56:       for (int k = 0; k < LANDAU_DIM; k++) gg3[j][k] += src.gg3[j][k];
 57:     }
 58:   }
 59: };
 60: typedef array_type<PetscReal> TensorValueType; // used to simplify code below
 61: } // namespace landau_inner_red

 63: namespace Kokkos
 64: { //reduction identity must be defined in Kokkos namespace
 65: template <>
 66: struct reduction_identity<landau_inner_red::TensorValueType> {
 67:   KOKKOS_FORCEINLINE_FUNCTION static landau_inner_red::TensorValueType sum() { return landau_inner_red::TensorValueType(); }
 68: };
 69: } // namespace Kokkos

 71: extern "C" {
 72: PetscErrorCode LandauKokkosCreateMatMaps(P4estVertexMaps maps[], pointInterpolationP4est (*pointMaps)[LANDAU_MAX_Q_FACE], PetscInt Nf[], PetscInt Nq, PetscInt grid)
 73: {
 74:   P4estVertexMaps                                                                                                                                     h_maps; /* host container */
 75:   const Kokkos::View<pointInterpolationP4est *[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>>   h_points((pointInterpolationP4est *)pointMaps, maps[grid].num_reduced);
 76:   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);
 77:   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);
 78:   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);

 80:   Kokkos::deep_copy(*d_gidx, h_gidx);
 81:   Kokkos::deep_copy(*d_points, h_points);
 82:   h_maps.num_elements = maps[grid].num_elements;
 83:   h_maps.num_face     = maps[grid].num_face;
 84:   h_maps.num_reduced  = maps[grid].num_reduced;
 85:   h_maps.deviceType   = maps[grid].deviceType;
 86:   h_maps.numgrids     = maps[grid].numgrids;
 87:   h_maps.Nf           = Nf[grid];
 88:   h_maps.c_maps       = (pointInterpolationP4est(*)[LANDAU_MAX_Q_FACE])d_points->data();
 89:   maps[grid].vp1      = (void *)d_points;
 90:   h_maps.gIdx         = (LandauIdx(*)[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ])d_gidx->data();
 91:   maps[grid].vp2      = (void *)d_gidx;
 92:   {
 93:     Kokkos::View<P4estVertexMaps, Kokkos::HostSpace> h_maps_k(&h_maps);
 94:     Kokkos::View<P4estVertexMaps>                   *d_maps_k = new Kokkos::View<P4estVertexMaps>(Kokkos::create_mirror(Kokkos::DefaultExecutionSpace::memory_space(), h_maps_k));
 95:     Kokkos::deep_copy(*d_maps_k, h_maps_k);
 96:     maps[grid].d_self = d_maps_k->data();
 97:     maps[grid].vp3    = (void *)d_maps_k;
 98:   }
 99:   return 0;
100: }
101: PetscErrorCode LandauKokkosDestroyMatMaps(P4estVertexMaps maps[], PetscInt num_grids)
102: {
103:   for (PetscInt grid = 0; grid < num_grids; grid++) {
104:     auto a = static_cast<Kokkos::View<pointInterpolationP4est *[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight> *>(maps[grid].vp1);
105:     auto b = static_cast<Kokkos::View<LandauIdx *[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ], Kokkos::LayoutRight> *>(maps[grid].vp2);
106:     auto c = static_cast<Kokkos::View<P4estVertexMaps *> *>(maps[grid].vp3);
107:     delete a;
108:     delete b;
109:     delete c;
110:   }
111:   return 0;
112: }

114: 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)
115: {
116:   PetscReal       *BB, *DD;
117:   PetscTabulation *Tf;
118:   PetscInt         dim;
119:   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;
120:   PetscDS          prob;

122:   DMGetDimension(plex, &dim);
123:   DMGetDS(plex, &prob);
125:   PetscDSGetTabulation(prob, &Tf);
126:   BB           = Tf[0]->T[0];
127:   DD           = Tf[0]->T[1];
128:   ip_offset[0] = ipf_offset[0] = elem_offset[0] = 0;
129:   nip                                           = 0;
130:   IPf_sz                                        = 0;
131:   for (PetscInt grid = 0; grid < num_grids; grid++) {
132:     PetscInt nfloc        = a_species_offset[grid + 1] - a_species_offset[grid];
133:     elem_offset[grid + 1] = elem_offset[grid] + a_numCells[grid];
134:     nip += a_numCells[grid] * Nq;
135:     ip_offset[grid + 1] = nip;
136:     IPf_sz += Nq * nfloc * a_numCells[grid];
137:     ipf_offset[grid + 1] = IPf_sz;
138:   }
139:   Nftot = a_species_offset[num_grids];
140:   PetscKokkosInitializeCheck();
141:   {
142:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_alpha(a_nu_alpha, Nftot);
143:     auto                                                                                                            alpha = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("alpha", Nftot);
144:     SData_d->alpha                                                                                                        = static_cast<void *>(alpha);
145:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_beta(a_nu_beta, Nftot);
146:     auto                                                                                                            beta = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("beta", Nftot);
147:     SData_d->beta                                                                                                        = static_cast<void *>(beta);
148:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_invMass(a_invMass, Nftot);
149:     auto                                                                                                            invMass = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("invMass", Nftot);
150:     SData_d->invMass                                                                                                        = static_cast<void *>(invMass);
151:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_BB(BB, Nq * Nb);
152:     auto                                                                                                            B = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("B", Nq * Nb);
153:     SData_d->B                                                                                                        = static_cast<void *>(B);
154:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_DD(DD, Nq * Nb * dim);
155:     auto                                                                                                            D = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("D", Nq * Nb * dim);
156:     SData_d->D                                                                                                        = static_cast<void *>(D);
157:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_invJ(a_invJ, nip * dim * dim);
158:     auto                                                                                                            invJ = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("invJ", nip * dim * dim);
159:     SData_d->invJ                                                                                                        = static_cast<void *>(invJ);
160:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_x(a_x, nip);
161:     auto                                                                                                            x = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("x", nip);
162:     SData_d->x                                                                                                        = static_cast<void *>(x);
163:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_y(a_y, nip);
164:     auto                                                                                                            y = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("y", nip);
165:     SData_d->y                                                                                                        = static_cast<void *>(y);
166:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_w(a_w, nip);
167:     auto                                                                                                            w = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("w", nip);
168:     SData_d->w                                                                                                        = static_cast<void *>(w);

170:     Kokkos::deep_copy(*alpha, h_alpha);
171:     Kokkos::deep_copy(*beta, h_beta);
172:     Kokkos::deep_copy(*invMass, h_invMass);
173:     Kokkos::deep_copy(*B, h_BB);
174:     Kokkos::deep_copy(*D, h_DD);
175:     Kokkos::deep_copy(*invJ, h_invJ);
176:     Kokkos::deep_copy(*x, h_x);
177:     Kokkos::deep_copy(*y, h_y);
178:     Kokkos::deep_copy(*w, h_w);

180:     if (dim == 3) {
181:       const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_z(a_z, nip);
182:       auto                                                                                                            z = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("z", nip);
183:       SData_d->z                                                                                                        = static_cast<void *>(z);
184:       Kokkos::deep_copy(*z, h_z);
185:     } else SData_d->z = NULL;

187:     //
188:     const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_NCells(a_numCells, num_grids);
189:     auto                                                                                                           nc = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("NCells", num_grids);
190:     SData_d->NCells                                                                                                   = static_cast<void *>(nc);
191:     Kokkos::deep_copy(*nc, h_NCells);

193:     const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_species_offset(a_species_offset, num_grids + 1);
194:     auto                                                                                                           soff = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("species_offset", num_grids + 1);
195:     SData_d->species_offset                                                                                             = static_cast<void *>(soff);
196:     Kokkos::deep_copy(*soff, h_species_offset);

198:     const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_mat_offset(a_mat_offset, num_grids + 1);
199:     auto                                                                                                           moff = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("mat_offset", num_grids + 1);
200:     SData_d->mat_offset                                                                                                 = static_cast<void *>(moff);
201:     Kokkos::deep_copy(*moff, h_mat_offset);

203:     const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_ip_offset(ip_offset, num_grids + 1);
204:     auto                                                                                                           ipoff = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("ip_offset", num_grids + 1);
205:     SData_d->ip_offset                                                                                                   = static_cast<void *>(ipoff);
206:     Kokkos::deep_copy(*ipoff, h_ip_offset);

208:     const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_elem_offset(elem_offset, num_grids + 1);
209:     auto                                                                                                           eoff = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("elem_offset", num_grids + 1);
210:     SData_d->elem_offset                                                                                                = static_cast<void *>(eoff);
211:     Kokkos::deep_copy(*eoff, h_elem_offset);

213:     const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_ipf_offset(ipf_offset, num_grids + 1);
214:     auto                                                                                                           ipfoff = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("ipf_offset", num_grids + 1);
215:     SData_d->ipf_offset                                                                                                   = static_cast<void *>(ipfoff);
216:     Kokkos::deep_copy(*ipfoff, h_ipf_offset);
217: #if defined(LANDAU_LAYOUT_LEFT) // preallocate dynamic data, no copy
218:     auto ipfdf_data = new Kokkos::View<PetscReal ***, Kokkos::LayoutLeft>("fdf", batch_sz, dim + 1, IPf_sz);
219: #else
220:     auto ipfdf_data = new Kokkos::View<PetscReal ***, Kokkos::LayoutRight>("fdf", batch_sz, dim + 1, IPf_sz);
221: #endif
222:     SData_d->ipfdf_data = static_cast<void *>(ipfdf_data);
223:     auto Eq_m           = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("Eq_m", Nftot); // allocate but do not set, same for whole batch
224:     SData_d->Eq_m       = static_cast<void *>(Eq_m);
225:     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);
226:     auto                                                                                                            coo_elem_offsets = new Kokkos::View<LandauIdx *, Kokkos::LayoutLeft>("coo_elem_offsets", SData_d->coo_n_cellsTot + 1);
227:     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);
228:     auto                                                                                                            coo_elem_fullNb = new Kokkos::View<LandauIdx *, Kokkos::LayoutLeft>("coo_elem_offsets", SData_d->coo_n_cellsTot);
229:     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));
230:     auto coo_elem_point_offsets = new Kokkos::View<LandauIdx *, Kokkos::LayoutLeft>("coo_elem_point_offsets", SData_d->coo_n_cellsTot * (LANDAU_MAX_NQ + 1));
231:     // assign & copy
232:     Kokkos::deep_copy(*coo_elem_offsets, h_coo_elem_offsets);
233:     Kokkos::deep_copy(*coo_elem_fullNb, h_coo_elem_fullNb);
234:     Kokkos::deep_copy(*coo_elem_point_offsets, h_coo_elem_point_offsets);
235:     // need to free this now and use pointer space
236:     PetscFree3(SData_d->coo_elem_offsets, SData_d->coo_elem_fullNb, SData_d->coo_elem_point_offsets);
237:     SData_d->coo_elem_offsets       = static_cast<void *>(coo_elem_offsets);
238:     SData_d->coo_elem_fullNb        = static_cast<void *>(coo_elem_fullNb);
239:     SData_d->coo_elem_point_offsets = static_cast<void *>(coo_elem_point_offsets);
240:     auto coo_vals                   = new Kokkos::View<PetscScalar *, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace>("coo_vals", SData_d->coo_size);
241:     SData_d->coo_vals               = static_cast<void *>(coo_vals);
242:   }
243:   SData_d->maps = NULL; // not used
244:   return 0;
245: }

247: PetscErrorCode LandauKokkosStaticDataClear(LandauStaticData *SData_d)
248: {
249:   if (SData_d->alpha) {
250:     auto alpha = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->alpha);
251:     delete alpha;
252:     SData_d->alpha = NULL;
253:     auto beta      = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->beta);
254:     delete beta;
255:     auto invMass = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->invMass);
256:     delete invMass;
257:     auto B = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->B);
258:     delete B;
259:     auto D = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->D);
260:     delete D;
261:     auto invJ = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->invJ);
262:     delete invJ;
263:     auto x = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->x);
264:     delete x;
265:     auto y = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->y);
266:     delete y;
267:     if (SData_d->z) {
268:       auto z = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->z);
269:       delete z;
270:     }
271: #if defined(LANDAU_LAYOUT_LEFT) // preallocate dynamic data, no copy
272:     auto z = static_cast<Kokkos::View<PetscReal ***, Kokkos::LayoutLeft> *>(SData_d->ipfdf_data);
273: #else
274:     auto z          = static_cast<Kokkos::View<PetscReal ***, Kokkos::LayoutRight> *>(SData_d->ipfdf_data);
275: #endif
276:     delete z;
277:     auto w = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->w);
278:     delete w;
279:     auto Eq_m = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->Eq_m);
280:     delete Eq_m;
281:     // offset
282:     auto nc = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->NCells);
283:     delete nc;
284:     auto soff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->species_offset);
285:     delete soff;
286:     auto moff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->mat_offset);
287:     delete moff;
288:     auto ipoff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->ip_offset);
289:     delete ipoff;
290:     auto eoff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->elem_offset);
291:     delete eoff;
292:     auto ipfoff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->ipf_offset);
293:     delete ipfoff;
294:     auto coo_elem_offsets = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>((void *)SData_d->coo_elem_offsets);
295:     delete coo_elem_offsets;
296:     auto coo_elem_fullNb = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>((void *)SData_d->coo_elem_fullNb);
297:     delete coo_elem_fullNb;
298:     auto coo_elem_point_offsets = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>((void *)SData_d->coo_elem_point_offsets);
299:     delete coo_elem_point_offsets;
300:     SData_d->coo_elem_offsets       = NULL;
301:     SData_d->coo_elem_point_offsets = NULL;
302:     SData_d->coo_elem_fullNb        = NULL;
303:     auto coo_vals                   = static_cast<Kokkos::View<PetscScalar *, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace> *>((void *)SData_d->coo_vals);
304:     delete coo_vals;
305:   }
306:   return 0;
307: }

309: #define KOKKOS_SHARED_LEVEL 0 // 0 is shared, 1 is global

311: KOKKOS_INLINE_FUNCTION
312: 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)
313: {
314:   PetscInt               idx, q, nr, nc, d;
315:   const LandauIdx *const Idxs                         = &d_maps->gIdx[elem][fieldA][0];
316:   PetscScalar            row_scale[LANDAU_MAX_Q_FACE] = {0}, col_scale[LANDAU_MAX_Q_FACE] = {0};
317:   if (coo_vals) { // mirror (i,j) in CreateStaticGPUData
318:     const int fullNb = coo_elem_fullNb[glb_elem_idx], fullNb2 = fullNb * fullNb;
319:     idx = Idxs[f];
320:     if (idx >= 0) {
321:       nr           = 1;
322:       row_scale[0] = 1.;
323:     } else {
324:       idx = -idx - 1;
325:       for (q = 0, nr = 0; q < d_maps->num_face; q++, nr++) {
326:         if (d_maps->c_maps[idx][q].gid < 0) break;
327:         row_scale[q] = d_maps->c_maps[idx][q].scale;
328:       }
329:     }
330:     idx = Idxs[g];
331:     if (idx >= 0) {
332:       nc           = 1;
333:       col_scale[0] = 1.;
334:     } else {
335:       idx = -idx - 1;
336:       nc  = d_maps->num_face;
337:       for (q = 0, nc = 0; q < d_maps->num_face; q++, nc++) {
338:         if (d_maps->c_maps[idx][q].gid < 0) break;
339:         col_scale[q] = d_maps->c_maps[idx][q].scale;
340:       }
341:     }
342:     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];
343:     for (int q = 0, idx2 = idx0; q < nr; q++) {
344:       for (int d = 0; d < nc; d++, idx2++) coo_vals[idx2] = row_scale[q] * col_scale[d] * Aij;
345:     }
346:   } else {
347:     PetscScalar vals[LANDAU_MAX_Q_FACE * LANDAU_MAX_Q_FACE] = {0};
348:     PetscInt    rows[LANDAU_MAX_Q_FACE], cols[LANDAU_MAX_Q_FACE];
349:     idx = Idxs[f];
350:     if (idx >= 0) {
351:       nr           = 1;
352:       rows[0]      = idx;
353:       row_scale[0] = 1.;
354:     } else {
355:       idx = -idx - 1;
356:       for (q = 0, nr = 0; q < d_maps->num_face; q++, nr++) {
357:         if (d_maps->c_maps[idx][q].gid < 0) break;
358:         rows[q]      = d_maps->c_maps[idx][q].gid;
359:         row_scale[q] = d_maps->c_maps[idx][q].scale;
360:       }
361:     }
362:     idx = Idxs[g];
363:     if (idx >= 0) {
364:       nc           = 1;
365:       cols[0]      = idx;
366:       col_scale[0] = 1.;
367:     } else {
368:       idx = -idx - 1;
369:       for (q = 0, nc = 0; q < d_maps->num_face; q++, nc++) {
370:         if (d_maps->c_maps[idx][q].gid < 0) break;
371:         cols[q]      = d_maps->c_maps[idx][q].gid;
372:         col_scale[q] = d_maps->c_maps[idx][q].scale;
373:       }
374:     }

376:     for (q = 0; q < nr; q++) rows[q] = rows[q] + moffset;
377:     for (q = 0; q < nc; q++) cols[q] = cols[q] + moffset;
378:     for (q = 0; q < nr; q++) {
379:       for (d = 0; d < nc; d++) vals[q * nc + d] = row_scale[q] * col_scale[d] * Aij;
380:     }
381:     MatSetValuesDevice(d_mat, nr, rows, nc, cols, vals, ADD_VALUES);
382:   }
383:   return 0;
384: }

386: 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)
387: {
388:   using scr_mem_t                   = Kokkos::DefaultExecutionSpace::scratch_memory_space;
389:   using real2_scr_t                 = Kokkos::View<PetscScalar **, Kokkos::LayoutRight, scr_mem_t>;
390:   using g2_scr_t                    = Kokkos::View<PetscReal ***, Kokkos::LayoutRight, scr_mem_t>;
391:   using g3_scr_t                    = Kokkos::View<PetscReal ****, Kokkos::LayoutRight, scr_mem_t>;
392:   PetscInt                   Nb     = Nq, dim, num_cells_max, Nf_max, num_cells_batch;
393:   int                        nfaces = 0, vector_size = 512 / Nq;
394:   LandauCtx                 *ctx;
395:   PetscReal                 *d_Eq_m     = NULL;
396:   PetscScalar               *d_vertex_f = NULL;
397:   P4estVertexMaps           *maps[LANDAU_MAX_GRIDS]; // this gets captured
398:   PetscSplitCSRDataStructure d_mat;
399:   PetscContainer             container;
400:   const int                  conc = Kokkos::DefaultExecutionSpace().concurrency(), openmp = !!(conc < 1000), team_size = (openmp == 0) ? Nq : 1;
401:   const PetscInt             coo_sz_batch = SData_d->coo_size / batch_sz;                                                 // capture
402:   auto                       d_alpha_k    = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->alpha); //static data
403:   const PetscReal           *d_alpha      = d_alpha_k->data();
404:   const PetscInt             Nftot        = d_alpha_k->size(); // total number of species
405:   auto                       d_beta_k     = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->beta);
406:   const PetscReal           *d_beta       = d_beta_k->data();
407:   auto                       d_invMass_k  = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->invMass);
408:   const PetscReal           *d_invMass    = d_invMass_k->data();
409:   auto                       d_B_k        = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->B);
410:   const PetscReal           *d_BB         = d_B_k->data();
411:   auto                       d_D_k        = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->D);
412:   const PetscReal           *d_DD         = d_D_k->data();
413:   auto                       d_invJ_k     = *static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->invJ); // use Kokkos vector in kernels
414:   auto                       d_x_k        = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->x);     //static data
415:   const PetscReal           *d_x          = d_x_k->data();
416:   auto                       d_y_k        = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->y); //static data
417:   const PetscReal           *d_y          = d_y_k->data();
418:   auto                       d_z_k        = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->z); //static data
419:   const PetscReal           *d_z          = (LANDAU_DIM == 3) ? d_z_k->data() : NULL;
420:   auto                       d_w_k        = *static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->w); //static data
421:   const PetscReal           *d_w          = d_w_k.data();
422:   // grid offsets - single vertex grid data
423:   auto            d_numCells_k       = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->NCells);
424:   const PetscInt *d_numCells         = d_numCells_k->data();
425:   auto            d_species_offset_k = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->species_offset);
426:   const PetscInt *d_species_offset   = d_species_offset_k->data();
427:   auto            d_mat_offset_k     = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->mat_offset);
428:   const PetscInt *d_mat_offset       = d_mat_offset_k->data();
429:   auto            d_ip_offset_k      = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->ip_offset);
430:   const PetscInt *d_ip_offset        = d_ip_offset_k->data();
431:   auto            d_ipf_offset_k     = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->ipf_offset);
432:   const PetscInt *d_ipf_offset       = d_ipf_offset_k->data();
433:   auto            d_elem_offset_k    = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->elem_offset);
434:   const PetscInt *d_elem_offset      = d_elem_offset_k->data();
435: #if defined(LANDAU_LAYOUT_LEFT) // preallocate dynamic data, including batched vertices
436:   Kokkos::View<PetscReal ***, Kokkos::LayoutLeft> d_fdf_k = *static_cast<Kokkos::View<PetscReal ***, Kokkos::LayoutLeft> *>(SData_d->ipfdf_data);
437: #else
438:   Kokkos::View<PetscReal ***, Kokkos::LayoutRight> d_fdf_k = *static_cast<Kokkos::View<PetscReal ***, Kokkos::LayoutRight> *>(SData_d->ipfdf_data);
439: #endif
440:   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
441:   // COO
442:   auto       d_coo_elem_offsets_k                         = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>(SData_d->coo_elem_offsets);
443:   LandauIdx *d_coo_elem_offsets                           = (SData_d->coo_size == 0) ? NULL : d_coo_elem_offsets_k->data();
444:   auto       d_coo_elem_fullNb_k                          = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>(SData_d->coo_elem_fullNb);
445:   LandauIdx *d_coo_elem_fullNb                            = (SData_d->coo_size == 0) ? NULL : d_coo_elem_fullNb_k->data();
446:   auto       d_coo_elem_point_offsets_k                   = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>(SData_d->coo_elem_point_offsets);
447:   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();
448:   auto         d_coo_vals_k                               = static_cast<Kokkos::View<PetscScalar *, Kokkos::LayoutRight> *>(SData_d->coo_vals);
449:   PetscScalar *d_coo_vals                                 = (SData_d->coo_size == 0) ? NULL : d_coo_vals_k->data();

451:   while (vector_size & (vector_size - 1)) vector_size = vector_size & (vector_size - 1);
452:   if (vector_size > 16) vector_size = 16; // printf("DEBUG\n");
453:   PetscLogEventBegin(events[3], 0, 0, 0, 0);
454:   DMGetApplicationContext(plex[0], &ctx);
456:   DMGetDimension(plex[0], &dim);
458:   if (ctx->gpu_assembly) {
459:     PetscObjectQuery((PetscObject)JacP, "assembly_maps", (PetscObject *)&container);
460:     if (container) {
461:       P4estVertexMaps *h_maps = NULL;
462:       PetscContainerGetPointer(container, (void **)&h_maps);
463:       for (PetscInt grid = 0; grid < num_grids; grid++) {
464:         if (h_maps[grid].d_self) {
465:           maps[grid] = h_maps[grid].d_self;
466:           nfaces     = h_maps[grid].num_face; // nface=0 for CPU assembly
467:         } else {
468:           SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "GPU assembly but no metadata in container");
469:         }
470:       }
471:       if (!d_coo_vals) {
472:         // this does the setup the first time called
473:         MatKokkosGetDeviceMatWrite(JacP, &d_mat);
474:       } else {
475:         d_mat = NULL;
476:       }
477:     } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "GPU assembly but no metadata in container");
478:   } else {
479:     for (PetscInt grid = 0; grid < num_grids; grid++) maps[grid] = NULL;
480:     nfaces    = 0;
481:     d_mat     = NULL;
482:     container = NULL;
483:   }
484:   num_cells_batch = Nf_max = num_cells_max = 0;
485:   for (PetscInt grid = 0; grid < num_grids; grid++) {
486:     int Nfloc = a_species_offset[grid + 1] - a_species_offset[grid];
487:     if (Nfloc > Nf_max) Nf_max = Nfloc;
488:     if (a_numCells[grid] > num_cells_max) num_cells_max = a_numCells[grid];
489:     num_cells_batch += a_numCells[grid]; // we don't have a host element offset here (but in ctx)
490:   }
491:   const int elem_mat_num_cells_max_grid = container ? 0 : num_cells_max;
492: #ifdef LAND_SUPPORT_CPU_ASS
493:   const int                                           totDim_max = Nf_max * Nq, elem_mat_size_max = totDim_max * totDim_max;
494:   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)
495: #endif
496:   const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_Eq_m_k(a_Eq_m, Nftot);
497:   if (a_elem_closure || a_xarray) {
498:     Kokkos::deep_copy(*d_Eq_m_k, h_Eq_m_k);
499:     d_Eq_m = d_Eq_m_k->data();
500:   } else d_Eq_m = NULL;
501:   PetscKokkosInitializeCheck();
502:   PetscLogEventEnd(events[3], 0, 0, 0, 0);
503:   if (a_elem_closure || a_xarray) { // Jacobian, create f & df
504:     Kokkos::View<PetscScalar *, Kokkos::LayoutLeft> *d_vertex_f_k = NULL;
505:     PetscLogEventBegin(events[1], 0, 0, 0, 0);
506:     if (a_elem_closure) {
507:       PetscInt closure_sz = 0; // argh, don't have this on the host!!!
508:       for (PetscInt grid = 0; grid < num_grids; grid++) {
509:         PetscInt nfloc = a_species_offset[grid + 1] - a_species_offset[grid];
510:         closure_sz += Nq * nfloc * a_numCells[grid];
511:       }
512:       closure_sz *= batch_sz;
513:       d_vertex_f_k = new Kokkos::View<PetscScalar *, Kokkos::LayoutLeft>("closure", closure_sz);
514:       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
515:       Kokkos::deep_copy(*d_vertex_f_k, h_closure_k);
516:       d_vertex_f = d_vertex_f_k->data();
517:     } else {
518:       d_vertex_f = (PetscScalar *)a_xarray;
519:     }
520:     PetscLogEventEnd(events[1], 0, 0, 0, 0);
521:     PetscLogEventBegin(events[8], 0, 0, 0, 0);
522:     PetscLogGpuTimeBegin();

524:     const int scr_bytes_fdf = real2_scr_t::shmem_size(Nf_max, Nb);
525:     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);
526:     Kokkos::parallel_for(
527:       "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) {
528:         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];
529:         // find my grid
530:         PetscInt grid = 0;
531:         while (b_elem_idx >= d_elem_offset[grid + 1]) grid++;
532:         {
533:           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];
534:           const PetscInt moffset = LAND_MOFFSET(b_id, grid, batch_sz, num_grids, d_mat_offset);
535:           {
536:             real2_scr_t      s_coef_k(team.team_scratch(KOKKOS_SHARED_LEVEL), Nf_max, Nb);
537:             PetscScalar     *coef;
538:             const PetscReal *invJe = &d_invJ_k((d_ip_offset[grid] + loc_elem * Nq) * dim * dim);
539:             // un pack IPData
540:             if (!maps[grid]) {
541:               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
542:             } else {
543:               coef = s_coef_k.data();
544:               Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, (int)loc_Nf), [=](int f) {
545:                 //for (int f = 0; f < loc_Nf; ++f) {
546:                 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, 0, (int)Nb), [=](int b) {
547:                   //for (int b = 0; b < Nb; ++b) {
548:                   LandauIdx idx = maps[grid]->gIdx[loc_elem][f][b];
549:                   if (idx >= 0) {
550:                     coef[f * Nb + b] = d_vertex_f[idx + moffset]; // xarray data, not IP, need raw array to deal with CPU assemble case (not used)
551:                   } else {
552:                     idx              = -idx - 1;
553:                     coef[f * Nb + b] = 0;
554:                     for (int q = 0; q < maps[grid]->num_face; q++) {
555:                       PetscInt    id    = maps[grid]->c_maps[idx][q].gid;
556:                       PetscScalar scale = maps[grid]->c_maps[idx][q].scale;
557:                       if (id >= 0) coef[f * Nb + b] += scale * d_vertex_f[id + moffset];
558:                     }
559:                   }
560:                 });
561:               });
562:             }
563:             team.team_barrier();
564:             Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, Nq), [=](int myQi) {
565:               const PetscReal *const invJ = &invJe[myQi * dim * dim]; // b_elem_idx: batch element index
566:               const PetscReal       *Bq = &d_BB[myQi * Nb], *Dq = &d_DD[myQi * Nb * dim];
567:               Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, 0, (int)loc_Nf), [=](int f) {
568:                 PetscInt       b, e, d;
569:                 PetscReal      refSpaceDer[LANDAU_DIM];
570:                 const PetscInt idx    = d_ipf_offset[grid] + f * loc_nip + loc_elem * Nq + myQi;
571:                 d_fdf_k(b_id, 0, idx) = 0.0;
572:                 for (d = 0; d < LANDAU_DIM; ++d) refSpaceDer[d] = 0.0;
573:                 for (b = 0; b < Nb; ++b) {
574:                   d_fdf_k(b_id, 0, idx) += Bq[b] * PetscRealPart(coef[f * Nb + b]);
575:                   for (d = 0; d < dim; ++d) refSpaceDer[d] += Dq[b * dim + d] * PetscRealPart(coef[f * Nb + b]);
576:                 }
577:                 for (d = 0; d < dim; ++d) {
578:                   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];
579:                 }
580:               }); // f
581:             });   // myQi
582:           }
583:           team.team_barrier();
584:         } // 'grid' scope
585:       }); // global elems - fdf
586:     Kokkos::fence();
587:     PetscLogGpuTimeEnd();
588:     PetscLogEventEnd(events[8], 0, 0, 0, 0);
589:     // Jacobian
590:     int         maximum_shared_mem_size = 64000;
591:     PetscDevice device;

593:     PetscDeviceGetDefault_Internal(&device);
594:     PetscDeviceGetAttribute(device, PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK, &maximum_shared_mem_size);
595:     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));
596:     const int jac_shared_level = (jac_scr_bytes > maximum_shared_mem_size) ? 1 : KOKKOS_SHARED_LEVEL;
597:     auto      jac_lambda       = KOKKOS_LAMBDA(const team_member team)
598:     {
599:       const PetscInt b_Nelem = d_elem_offset[num_grids], b_elem_idx = team.league_rank() % b_Nelem, b_id = team.league_rank() / b_Nelem;
600:       // find my grid
601:       PetscInt grid = 0;
602:       while (b_elem_idx >= d_elem_offset[grid + 1]) grid++;
603:       {
604:         const PetscInt loc_Nf = d_species_offset[grid + 1] - d_species_offset[grid], loc_elem = b_elem_idx - d_elem_offset[grid];
605:         const PetscInt moffset = LAND_MOFFSET(b_id, grid, batch_sz, num_grids, d_mat_offset);
606:         const PetscInt f_off   = d_species_offset[grid];
607: #ifdef LAND_SUPPORT_CPU_ASS
608:         const PetscInt totDim = loc_Nf * Nq;
609: #endif
610:         g2_scr_t g2(team.team_scratch(jac_shared_level), dim, loc_Nf, Nq);
611:         g3_scr_t g3(team.team_scratch(jac_shared_level), dim, dim, loc_Nf, Nq);
612:         g2_scr_t gg2(team.team_scratch(jac_shared_level), dim, loc_Nf, Nq);
613:         g3_scr_t gg3(team.team_scratch(jac_shared_level), dim, dim, loc_Nf, Nq);
614:         // get g2[] & g3[]
615:         Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, Nq), [=](int myQi) {
616:           using Kokkos::parallel_reduce;
617:           const PetscInt                    jpidx_glb = d_ip_offset[grid] + loc_elem * Nq + myQi;
618:           const PetscReal                  *invJ      = &d_invJ_k(jpidx_glb * dim * dim);
619:           const PetscReal                   vj[3] = {d_x[jpidx_glb], d_y[jpidx_glb], d_z ? d_z[jpidx_glb] : 0}, wj = d_w[jpidx_glb];
620:           landau_inner_red::TensorValueType gg_temp; // reduce on part of gg2 and g33 for IP jpidx_g
621:           Kokkos::parallel_reduce(
622:             Kokkos::ThreadVectorRange(team, (int)d_ip_offset[num_grids]),
623:             [=](const int &ipidx, landau_inner_red::TensorValueType &ggg) {
624:               const PetscReal wi = d_w[ipidx], x = d_x[ipidx], y = d_y[ipidx];
625:               PetscReal       temp1[3] = {0, 0, 0}, temp2 = 0;
626:               PetscInt        fieldA, d2, d3, f_off_r, grid_r, ipidx_g, nip_loc_r, loc_Nf_r;
627: #if LANDAU_DIM == 2
628:               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.;
629:               LandauTensor2D(vj, x, y, Ud, Uk, mask);
630: #else
631:                 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.;
632:                 LandauTensor3D(vj, x, y, z, U, mask);
633: #endif
634:               grid_r = 0;
635:               while (ipidx >= d_ip_offset[grid_r + 1]) grid_r++; // yuck search for grid
636:               f_off_r   = d_species_offset[grid_r];
637:               ipidx_g   = ipidx - d_ip_offset[grid_r];
638:               nip_loc_r = d_numCells[grid_r] * Nq;
639:               loc_Nf_r  = d_species_offset[grid_r + 1] - d_species_offset[grid_r];
640:               for (fieldA = 0; fieldA < loc_Nf_r; ++fieldA) {
641:                 const PetscInt idx = d_ipf_offset[grid_r] + fieldA * nip_loc_r + ipidx_g;
642:                 temp1[0] += d_fdf_k(b_id, 1, idx) * d_beta[fieldA + f_off_r] * d_invMass[fieldA + f_off_r];
643:                 temp1[1] += d_fdf_k(b_id, 2, idx) * d_beta[fieldA + f_off_r] * d_invMass[fieldA + f_off_r];
644: #if LANDAU_DIM == 3
645:                 temp1[2] += d_fdf_k(b_id, 3, idx) * d_beta[fieldA + f_off_r] * d_invMass[fieldA + f_off_r];
646: #endif
647:                 temp2 += d_fdf_k(b_id, 0, idx) * d_beta[fieldA + f_off_r];
648:               }
649:               temp1[0] *= wi;
650:               temp1[1] *= wi;
651: #if LANDAU_DIM == 3
652:               temp1[2] *= wi;
653: #endif
654:               temp2 *= wi;
655: #if LANDAU_DIM == 2
656:               for (d2 = 0; d2 < 2; d2++) {
657:                 for (d3 = 0; d3 < 2; ++d3) {
658:                   /* K = U * grad(f): g2=e: i,A */
659:                   ggg.gg2[d2] += Uk[d2][d3] * temp1[d3];
660:                   /* D = -U * (I \kron (fx)): g3=f: i,j,A */
661:                   ggg.gg3[d2][d3] += Ud[d2][d3] * temp2;
662:                 }
663:               }
664: #else
665:                 for (d2 = 0; d2 < 3; ++d2) {
666:                   for (d3 = 0; d3 < 3; ++d3) {
667:                     /* K = U * grad(f): g2 = e: i,A */
668:                     ggg.gg2[d2] += U[d2][d3]*temp1[d3];
669:                     /* D = -U * (I \kron (fx)): g3 = f: i,j,A */
670:                     ggg.gg3[d2][d3] += U[d2][d3]*temp2;
671:                   }
672:                 }
673: #endif
674:             },
675:             Kokkos::Sum<landau_inner_red::TensorValueType>(gg_temp));
676:           // add alpha and put in gg2/3
677:           Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, (int)loc_Nf), [&](const int &fieldA) {
678:             PetscInt d2, d3;
679:             for (d2 = 0; d2 < dim; d2++) {
680:               gg2(d2, fieldA, myQi) = gg_temp.gg2[d2] * d_alpha[fieldA + f_off];
681:               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];
682:             }
683:           });
684:           /* add electric field term once per IP */
685:           Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, (int)loc_Nf), [&](const int &fieldA) { gg2(dim - 1, fieldA, myQi) += d_Eq_m[fieldA + f_off]; });
686:           Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, (int)loc_Nf), [=](const int &fieldA) {
687:             int d, d2, d3, dp;
688:             /* Jacobian transform - g2, g3 - per thread (2D) */
689:             for (d = 0; d < dim; ++d) {
690:               g2(d, fieldA, myQi) = 0;
691:               for (d2 = 0; d2 < dim; ++d2) {
692:                 g2(d, fieldA, myQi) += invJ[d * dim + d2] * gg2(d2, fieldA, myQi);
693:                 g3(d, d2, fieldA, myQi) = 0;
694:                 for (d3 = 0; d3 < dim; ++d3) {
695:                   for (dp = 0; dp < dim; ++dp) g3(d, d2, fieldA, myQi) += invJ[d * dim + d3] * gg3(d3, dp, fieldA, myQi) * invJ[d2 * dim + dp];
696:                 }
697:                 g3(d, d2, fieldA, myQi) *= wj;
698:               }
699:               g2(d, fieldA, myQi) *= wj;
700:             }
701:           });
702:         }); // Nq
703:         team.team_barrier();
704:         { /* assemble */
705:           for (PetscInt fieldA = 0; fieldA < loc_Nf; fieldA++) {
706:             /* assemble */
707:             Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, Nb), [=](int f) {
708:               Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, 0, Nb), [=](int g) {
709:                 PetscScalar t = 0;
710:                 for (int qj = 0; qj < Nq; qj++) { // look at others integration points
711:                   const PetscReal *BJq = &d_BB[qj * Nb], *DIq = &d_DD[qj * Nb * dim];
712:                   for (int d = 0; d < dim; ++d) {
713:                     t += DIq[f * dim + d] * g2(d, fieldA, qj) * BJq[g];
714:                     for (int d2 = 0; d2 < dim; ++d2) t += DIq[f * dim + d] * g3(d, d2, fieldA, qj) * DIq[g * dim + d2];
715:                   }
716:                 }
717:                 if (elem_mat_num_cells_max_grid) { // CPU assembly
718: #ifdef LAND_SUPPORT_CPU_ASS
719:                   const PetscInt fOff                     = (fieldA * Nb + f) * totDim + fieldA * Nb + g;
720:                   d_elem_mats(b_id, grid, loc_elem, fOff) = t;
721: #endif
722:                 } else {
723:                   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);
724:                 }
725:               });
726:             });
727:           }
728:         }
729:       } // scope with 'grid'
730:     };
731:     PetscLogEventBegin(events[4], 0, 0, 0, 0);
732:     PetscLogGpuTimeBegin();
733:     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);
734:     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>
735:     Kokkos::fence();
736:     PetscLogGpuTimeEnd();
737:     PetscLogEventEnd(events[4], 0, 0, 0, 0);
738:     if (d_vertex_f_k) delete d_vertex_f_k;
739:   } else { // mass
740:     PetscLogEventBegin(events[16], 0, 0, 0, 0);
741:     PetscLogGpuTimeBegin();
742:     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");
743:     Kokkos::parallel_for(
744:       "Mass", Kokkos::TeamPolicy<>(num_cells_batch * batch_sz, team_size, vector_size), KOKKOS_LAMBDA(const team_member team) { // Kokkos::LaunchBounds<512,4>
745:         const PetscInt b_Nelem = d_elem_offset[num_grids], b_elem_idx = team.league_rank() % b_Nelem, b_id = team.league_rank() / b_Nelem;
746:         // find my grid
747:         PetscInt grid = 0;
748:         while (b_elem_idx >= d_elem_offset[grid + 1]) grid++;
749:         {
750:           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;
751: #ifdef LAND_SUPPORT_CPU_ASS
752:           const PetscInt totDim = loc_Nf * Nq;
753: #endif
754:           const PetscInt moffset = LAND_MOFFSET(b_id, grid, batch_sz, num_grids, d_mat_offset);
755:           for (int fieldA = 0; fieldA < loc_Nf; fieldA++) {
756:             /* assemble */
757:             Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, Nb), [=](int f) {
758:               Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, 0, Nb), [=](int g) {
759:                 PetscScalar t = 0;
760:                 for (int qj = 0; qj < Nq; qj++) { // look at others integration points
761:                   const PetscReal *BJq       = &d_BB[qj * Nb];
762:                   const PetscInt   jpidx_glb = jpidx_0 + qj;
763:                   if (dim == 2) {
764:                     t += BJq[f] * d_w_k(jpidx_glb) * shift * BJq[g] * 2. * PETSC_PI;
765:                   } else {
766:                     t += BJq[f] * d_w_k(jpidx_glb) * shift * BJq[g];
767:                   }
768:                 }
769:                 if (elem_mat_num_cells_max_grid) {
770: #ifdef LAND_SUPPORT_CPU_ASS
771:                   const PetscInt fOff                     = (fieldA * Nb + f) * totDim + fieldA * Nb + g;
772:                   d_elem_mats(b_id, grid, loc_elem, fOff) = t;
773: #endif
774:                 } else {
775:                   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);
776:                 }
777:               });
778:             });
779:           } // field
780:         }   // grid
781:       });
782:     Kokkos::fence();
783:     PetscLogGpuTimeEnd();
784:     PetscLogEventEnd(events[16], 0, 0, 0, 0);
785:   }
786:   if (d_coo_vals) {
787:     MatSetValuesCOO(JacP, d_coo_vals, ADD_VALUES);
788: #ifndef LAND_SUPPORT_CPU_ASS
789:     Kokkos::fence(); // for timers
790: #endif
791:   } else if (elem_mat_num_cells_max_grid) { // CPU assembly
792: #ifdef LAND_SUPPORT_CPU_ASS
793:     Kokkos::View<PetscScalar ****, Kokkos::LayoutRight>::HostMirror h_elem_mats = Kokkos::create_mirror_view(d_elem_mats);
794:     Kokkos::deep_copy(h_elem_mats, d_elem_mats);
796:     for (PetscInt b_id = 0; b_id < batch_sz; b_id++) { // OpenMP (once)
797:       for (PetscInt grid = 0; grid < num_grids; grid++) {
798:         PetscSection       section, globalSection;
799:         PetscInt           cStart, cEnd;
800:         Mat                B       = subJ[LAND_PACK_IDX(b_id, grid)];
801:         PetscInt           moffset = LAND_MOFFSET(b_id, grid, batch_sz, num_grids, a_mat_offset), nloc, nzl, colbuf[1024], row;
802:         const PetscInt    *cols;
803:         const PetscScalar *vals;
804:         PetscLogEventBegin(events[5], 0, 0, 0, 0);
805:         DMPlexGetHeightStratum(plex[grid], 0, &cStart, &cEnd);
806:         DMGetLocalSection(plex[grid], &section);
807:         DMGetGlobalSection(plex[grid], &globalSection);
808:         PetscLogEventEnd(events[5], 0, 0, 0, 0);
809:         PetscLogEventBegin(events[6], 0, 0, 0, 0);
810:         for (PetscInt ej = cStart; ej < cEnd; ++ej) {
811:           const PetscScalar *elMat = &h_elem_mats(b_id, grid, ej - cStart, 0);
812:           DMPlexMatSetClosure(plex[grid], section, globalSection, B, ej, elMat, ADD_VALUES);
813:           if (grid == 0 && ej == -1) {
814:             const PetscInt loc_Nf = a_species_offset[grid + 1] - a_species_offset[grid], totDim = loc_Nf * Nq;
815:             int            d, f;
816:             PetscPrintf(PETSC_COMM_SELF, "Kokkos Element matrix %d/%d\n", 1, (int)a_numCells[grid]);
817:             for (d = 0; d < totDim; ++d) {
818:               for (f = 0; f < totDim; ++f) PetscPrintf(PETSC_COMM_SELF, " %12.5e", PetscRealPart(elMat[d * totDim + f]));
819:               PetscPrintf(PETSC_COMM_SELF, "\n");
820:             }
821:             exit(14);
822:           }
823:         }
824:         // move nest matrix to global JacP
825:         MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY);
826:         MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY);
827:         MatGetSize(B, &nloc, NULL);
828:         for (int i = 0; i < nloc; i++) {
829:           MatGetRow(B, i, &nzl, &cols, &vals);
831:           for (int j = 0; j < nzl; j++) colbuf[j] = cols[j] + moffset;
832:           row = i + moffset;
833:           MatSetValues(JacP, 1, &row, nzl, colbuf, vals, ADD_VALUES);
834:           MatRestoreRow(B, i, &nzl, &cols, &vals);
835:         }
836:         MatDestroy(&subJ[LAND_PACK_IDX(b_id, grid)]);
837:         PetscLogEventEnd(events[6], 0, 0, 0, 0);
838:       } // grids
839:     }
840: #else
841:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "CPU assembly not supported");
842: #endif
843:   }
844:   return 0;
845: }
846: } // extern "C"