Actual source code: landau.kokkos.cxx

      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, 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)
     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, const PetscInt, PetscInt[], PetscInt[], PetscInt[], PetscReal[], 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>
     32: #include <petsc/private/deviceimpl.h>
     33: #include <petscts.h>

     35:   #include <Kokkos_Core.hpp>
     36:   #include <cstdio>
     37: typedef Kokkos::TeamPolicy<>::member_type team_member;
     38:   #include "../land_tensors.h"

     40: namespace landau_inner_red
     41: { // namespace helps with name resolution in reduction identity
     42: template <class ScalarType>
     43: struct array_type {
     44:   ScalarType gg2[LANDAU_DIM];
     45:   ScalarType gg3[LANDAU_DIM][LANDAU_DIM];

     47:   KOKKOS_INLINE_FUNCTION // Default constructor - Initialize to 0's
     48:   array_type()
     49:   {
     50:     for (int j = 0; j < LANDAU_DIM; j++) {
     51:       gg2[j] = 0;
     52:       for (int k = 0; k < LANDAU_DIM; k++) gg3[j][k] = 0;
     53:     }
     54:   }
     55:   KOKKOS_INLINE_FUNCTION // Copy Constructor
     56:   array_type(const array_type &rhs)
     57:   {
     58:     for (int j = 0; j < LANDAU_DIM; j++) {
     59:       gg2[j] = rhs.gg2[j];
     60:       for (int k = 0; k < LANDAU_DIM; k++) gg3[j][k] = rhs.gg3[j][k];
     61:     }
     62:   }
     63:   KOKKOS_INLINE_FUNCTION // add operator
     64:     array_type &
     65:     operator+=(const array_type &src)
     66:   {
     67:     for (int j = 0; j < LANDAU_DIM; j++) {
     68:       gg2[j] += src.gg2[j];
     69:       for (int k = 0; k < LANDAU_DIM; k++) gg3[j][k] += src.gg3[j][k];
     70:     }
     71:     return *this;
     72:   }
     73:   KOKKOS_INLINE_FUNCTION // volatile add operator
     74:     void
     75:     operator+=(const volatile array_type &src) volatile
     76:   {
     77:     for (int j = 0; j < LANDAU_DIM; j++) {
     78:       gg2[j] += src.gg2[j];
     79:       for (int k = 0; k < LANDAU_DIM; k++) gg3[j][k] += src.gg3[j][k];
     80:     }
     81:   }
     82: };
     83: typedef array_type<PetscReal> TensorValueType; // used to simplify code below
     84: } // namespace landau_inner_red

     86: namespace Kokkos
     87: { //reduction identity must be defined in Kokkos namespace
     88: template <>
     89: struct reduction_identity<landau_inner_red::TensorValueType> {
     90:   KOKKOS_FORCEINLINE_FUNCTION static landau_inner_red::TensorValueType sum() { return landau_inner_red::TensorValueType(); }
     91: };
     92: } // namespace Kokkos

     94: extern "C" {
     95: PetscErrorCode LandauKokkosCreateMatMaps(P4estVertexMaps maps[], pointInterpolationP4est (*pointMaps)[LANDAU_MAX_Q_FACE], PetscInt Nf[], PetscInt grid)
     96: {
     97:   P4estVertexMaps                                                                                                                                       h_maps; /* host container */
     98:   const Kokkos::View<pointInterpolationP4est *[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>>     h_points((pointInterpolationP4est *)pointMaps, maps[grid].num_reduced);
     99:   const Kokkos::View<LandauIdx *[LANDAU_MAX_SPECIES][LANDAU_MAX_NQND], Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_gidx((LandauIdx *)maps[grid].gIdx, maps[grid].num_elements);
    100:   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);
    101:   Kokkos::View<LandauIdx *[LANDAU_MAX_SPECIES][LANDAU_MAX_NQND], Kokkos::LayoutRight> *d_gidx   = new Kokkos::View<LandauIdx *[LANDAU_MAX_SPECIES][LANDAU_MAX_NQND], Kokkos::LayoutRight>("gIdx", maps[grid].num_elements);

    103:   PetscFunctionBegin;
    104:   Kokkos::deep_copy(*d_gidx, h_gidx);
    105:   Kokkos::deep_copy(*d_points, h_points);
    106:   h_maps.num_elements = maps[grid].num_elements;
    107:   h_maps.num_face     = maps[grid].num_face;
    108:   h_maps.num_reduced  = maps[grid].num_reduced;
    109:   h_maps.deviceType   = maps[grid].deviceType;
    110:   h_maps.numgrids     = maps[grid].numgrids;
    111:   h_maps.Nf           = Nf[grid];
    112:   h_maps.c_maps       = (pointInterpolationP4est(*)[LANDAU_MAX_Q_FACE])d_points->data();
    113:   maps[grid].vp1      = (void *)d_points;
    114:   h_maps.gIdx         = (LandauIdx(*)[LANDAU_MAX_SPECIES][LANDAU_MAX_NQND])d_gidx->data();
    115:   maps[grid].vp2      = (void *)d_gidx;
    116:   {
    117:     Kokkos::View<P4estVertexMaps, Kokkos::HostSpace> h_maps_k(&h_maps);
    118:     Kokkos::View<P4estVertexMaps>                   *d_maps_k = new Kokkos::View<P4estVertexMaps>(Kokkos::create_mirror(Kokkos::DefaultExecutionSpace::memory_space(), h_maps_k));
    119:     Kokkos::deep_copy(*d_maps_k, h_maps_k);
    120:     maps[grid].d_self = d_maps_k->data();
    121:     maps[grid].vp3    = (void *)d_maps_k;
    122:   }
    123:   PetscFunctionReturn(PETSC_SUCCESS);
    124: }
    125: PetscErrorCode LandauKokkosDestroyMatMaps(P4estVertexMaps maps[], PetscInt num_grids)
    126: {
    127:   PetscFunctionBegin;
    128:   for (PetscInt grid = 0; grid < num_grids; grid++) {
    129:     auto a = static_cast<Kokkos::View<pointInterpolationP4est *[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight> *>(maps[grid].vp1);
    130:     auto b = static_cast<Kokkos::View<LandauIdx *[LANDAU_MAX_SPECIES][LANDAU_MAX_NQND], Kokkos::LayoutRight> *>(maps[grid].vp2);
    131:     auto c = static_cast<Kokkos::View<P4estVertexMaps *> *>(maps[grid].vp3);
    132:     delete a;
    133:     delete b;
    134:     delete c;
    135:   }
    136:   PetscFunctionReturn(PETSC_SUCCESS);
    137: }

    139: PetscErrorCode LandauKokkosStaticDataSet(DM plex, const PetscInt Nq, const PetscInt Nb, 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)
    140: {
    141:   PetscReal       *BB, *DD;
    142:   PetscTabulation *Tf;
    143:   PetscInt         dim;
    144:   PetscInt         ip_offset[LANDAU_MAX_GRIDS + 1], ipf_offset[LANDAU_MAX_GRIDS + 1], elem_offset[LANDAU_MAX_GRIDS + 1], nip, IPf_sz, Nftot;
    145:   PetscDS          prob;

    147:   PetscFunctionBegin;
    148:   PetscCall(DMGetDimension(plex, &dim));
    149:   PetscCall(DMGetDS(plex, &prob));
    150:   PetscCheck(LANDAU_DIM == dim, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "dim %" PetscInt_FMT " != LANDAU_DIM %d", dim, LANDAU_DIM);
    151:   PetscCall(PetscDSGetTabulation(prob, &Tf));
    152:   BB           = Tf[0]->T[0];
    153:   DD           = Tf[0]->T[1];
    154:   ip_offset[0] = ipf_offset[0] = elem_offset[0] = 0;
    155:   nip                                           = 0;
    156:   IPf_sz                                        = 0;
    157:   for (PetscInt grid = 0; grid < num_grids; grid++) {
    158:     PetscInt nfloc        = a_species_offset[grid + 1] - a_species_offset[grid];
    159:     elem_offset[grid + 1] = elem_offset[grid] + a_numCells[grid];
    160:     nip += a_numCells[grid] * Nq;
    161:     ip_offset[grid + 1] = nip;
    162:     IPf_sz += Nq * nfloc * a_numCells[grid];
    163:     ipf_offset[grid + 1] = IPf_sz;
    164:   }
    165:   Nftot = a_species_offset[num_grids];
    166:   PetscCall(PetscKokkosInitializeCheck());
    167:   {
    168:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_alpha(a_nu_alpha, Nftot);
    169:     auto                                                                                                            alpha = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("alpha", Nftot);
    170:     SData_d->alpha                                                                                                        = static_cast<void *>(alpha);
    171:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_beta(a_nu_beta, Nftot);
    172:     auto                                                                                                            beta = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("beta", Nftot);
    173:     SData_d->beta                                                                                                        = static_cast<void *>(beta);
    174:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_invMass(a_invMass, Nftot);
    175:     auto                                                                                                            invMass = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("invMass", Nftot);
    176:     SData_d->invMass                                                                                                        = static_cast<void *>(invMass);
    177:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_lambdas(a_lambdas, LANDAU_MAX_GRIDS * LANDAU_MAX_GRIDS);
    178:     auto                                                                                                            lambdas = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("lambdas", LANDAU_MAX_GRIDS * LANDAU_MAX_GRIDS);
    179:     SData_d->lambdas                                                                                                        = static_cast<void *>(lambdas);
    180:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_BB(BB, Nq * Nb);
    181:     auto                                                                                                            B = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("B", Nq * Nb);
    182:     SData_d->B                                                                                                        = static_cast<void *>(B);
    183:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_DD(DD, Nq * Nb * dim);
    184:     auto                                                                                                            D = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("D", Nq * Nb * dim);
    185:     SData_d->D                                                                                                        = static_cast<void *>(D);
    186:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_invJ(a_invJ, nip * dim * dim);
    187:     auto                                                                                                            invJ = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("invJ", nip * dim * dim);
    188:     SData_d->invJ                                                                                                        = static_cast<void *>(invJ);
    189:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_x(a_x, nip);
    190:     auto                                                                                                            x = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("x", nip);
    191:     SData_d->x                                                                                                        = static_cast<void *>(x);
    192:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_y(a_y, nip);
    193:     auto                                                                                                            y = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("y", nip);
    194:     SData_d->y                                                                                                        = static_cast<void *>(y);
    195:     const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_w(a_w, nip);
    196:     auto                                                                                                            w = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("w", nip);
    197:     SData_d->w                                                                                                        = static_cast<void *>(w);

    199:     Kokkos::deep_copy(*alpha, h_alpha);
    200:     Kokkos::deep_copy(*beta, h_beta);
    201:     Kokkos::deep_copy(*invMass, h_invMass);
    202:     Kokkos::deep_copy(*lambdas, h_lambdas);
    203:     Kokkos::deep_copy(*B, h_BB);
    204:     Kokkos::deep_copy(*D, h_DD);
    205:     Kokkos::deep_copy(*invJ, h_invJ);
    206:     Kokkos::deep_copy(*x, h_x);
    207:     Kokkos::deep_copy(*y, h_y);
    208:     Kokkos::deep_copy(*w, h_w);

    210:     if (dim == 3) {
    211:       const Kokkos::View<PetscReal *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_z(a_z, nip);
    212:       auto                                                                                                            z = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("z", nip);
    213:       SData_d->z                                                                                                        = static_cast<void *>(z);
    214:       Kokkos::deep_copy(*z, h_z);
    215:     } else SData_d->z = NULL;

    217:     //
    218:     const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_NCells(a_numCells, num_grids);
    219:     auto                                                                                                           nc = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("NCells", num_grids);
    220:     SData_d->NCells                                                                                                   = static_cast<void *>(nc);
    221:     Kokkos::deep_copy(*nc, h_NCells);

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

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

    233:     const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_ip_offset(ip_offset, num_grids + 1);
    234:     auto                                                                                                           ipoff = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("ip_offset", num_grids + 1);
    235:     SData_d->ip_offset                                                                                                   = static_cast<void *>(ipoff);
    236:     Kokkos::deep_copy(*ipoff, h_ip_offset);

    238:     const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_elem_offset(elem_offset, num_grids + 1);
    239:     auto                                                                                                           eoff = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("elem_offset", num_grids + 1);
    240:     SData_d->elem_offset                                                                                                = static_cast<void *>(eoff);
    241:     Kokkos::deep_copy(*eoff, h_elem_offset);

    243:     const Kokkos::View<PetscInt *, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_ipf_offset(ipf_offset, num_grids + 1);
    244:     auto                                                                                                           ipfoff = new Kokkos::View<PetscInt *, Kokkos::LayoutLeft>("ipf_offset", num_grids + 1);
    245:     SData_d->ipf_offset                                                                                                   = static_cast<void *>(ipfoff);
    246:     Kokkos::deep_copy(*ipfoff, h_ipf_offset);
    247:   #if defined(LANDAU_LAYOUT_LEFT) // preallocate dynamic data, no copy
    248:     auto ipfdf_data = new Kokkos::View<PetscReal ***, Kokkos::LayoutLeft>("fdf", batch_sz, dim + 1, IPf_sz);
    249:   #else
    250:     auto ipfdf_data = new Kokkos::View<PetscReal ***, Kokkos::LayoutRight>("fdf", batch_sz, dim + 1, IPf_sz);
    251:   #endif
    252:     SData_d->ipfdf_data = static_cast<void *>(ipfdf_data);
    253:     auto Eq_m           = new Kokkos::View<PetscReal *, Kokkos::LayoutLeft>("Eq_m", Nftot); // allocate but do not set, same for whole batch
    254:     SData_d->Eq_m       = static_cast<void *>(Eq_m);
    255:     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);
    256:     auto                                                                                                            coo_elem_offsets = new Kokkos::View<LandauIdx *, Kokkos::LayoutLeft>("coo_elem_offsets", SData_d->coo_n_cellsTot + 1);
    257:     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);
    258:     auto                                                                                                            coo_elem_fullNb = new Kokkos::View<LandauIdx *, Kokkos::LayoutLeft>("coo_elem_offsets", SData_d->coo_n_cellsTot);
    259:     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_NQND + 1));
    260:     auto coo_elem_point_offsets = new Kokkos::View<LandauIdx *, Kokkos::LayoutLeft>("coo_elem_point_offsets", SData_d->coo_n_cellsTot * (LANDAU_MAX_NQND + 1));
    261:     // assign & copy
    262:     Kokkos::deep_copy(*coo_elem_offsets, h_coo_elem_offsets);
    263:     Kokkos::deep_copy(*coo_elem_fullNb, h_coo_elem_fullNb);
    264:     Kokkos::deep_copy(*coo_elem_point_offsets, h_coo_elem_point_offsets);
    265:     // need to free this now and use pointer space
    266:     PetscCall(PetscFree3(SData_d->coo_elem_offsets, SData_d->coo_elem_fullNb, SData_d->coo_elem_point_offsets));
    267:     SData_d->coo_elem_offsets       = static_cast<void *>(coo_elem_offsets);
    268:     SData_d->coo_elem_fullNb        = static_cast<void *>(coo_elem_fullNb);
    269:     SData_d->coo_elem_point_offsets = static_cast<void *>(coo_elem_point_offsets);
    270:     auto coo_vals                   = new Kokkos::View<PetscScalar *, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace>("coo_vals", SData_d->coo_size);
    271:     SData_d->coo_vals               = static_cast<void *>(coo_vals);
    272:   }
    273:   SData_d->maps = NULL; // not used
    274:   PetscFunctionReturn(PETSC_SUCCESS);
    275: }

    277: PetscErrorCode LandauKokkosStaticDataClear(LandauStaticData *SData_d)
    278: {
    279:   PetscFunctionBegin;
    280:   if (SData_d->alpha) {
    281:     auto alpha = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->alpha);
    282:     delete alpha;
    283:     SData_d->alpha = NULL;
    284:     auto beta      = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->beta);
    285:     delete beta;
    286:     auto invMass = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->invMass);
    287:     delete invMass;
    288:     auto lambdas = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->lambdas);
    289:     delete lambdas;
    290:     auto B = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->B);
    291:     delete B;
    292:     auto D = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->D);
    293:     delete D;
    294:     auto invJ = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->invJ);
    295:     delete invJ;
    296:     auto x = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->x);
    297:     delete x;
    298:     auto y = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->y);
    299:     delete y;
    300:     if (SData_d->z) {
    301:       auto z = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->z);
    302:       delete z;
    303:     }
    304:   #if defined(LANDAU_LAYOUT_LEFT) // preallocate dynamic data, no copy
    305:     auto z = static_cast<Kokkos::View<PetscReal ***, Kokkos::LayoutLeft> *>(SData_d->ipfdf_data);
    306:   #else
    307:     auto z = static_cast<Kokkos::View<PetscReal ***, Kokkos::LayoutRight> *>(SData_d->ipfdf_data);
    308:   #endif
    309:     delete z;
    310:     auto w = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->w);
    311:     delete w;
    312:     auto Eq_m = static_cast<Kokkos::View<PetscReal *, Kokkos::LayoutLeft> *>(SData_d->Eq_m);
    313:     delete Eq_m;
    314:     // offset
    315:     auto nc = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->NCells);
    316:     delete nc;
    317:     auto soff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->species_offset);
    318:     delete soff;
    319:     auto moff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->mat_offset);
    320:     delete moff;
    321:     auto ipoff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->ip_offset);
    322:     delete ipoff;
    323:     auto eoff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->elem_offset);
    324:     delete eoff;
    325:     auto ipfoff = static_cast<Kokkos::View<PetscInt *, Kokkos::LayoutLeft> *>(SData_d->ipf_offset);
    326:     delete ipfoff;
    327:     auto coo_elem_offsets = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>((void *)SData_d->coo_elem_offsets);
    328:     delete coo_elem_offsets;
    329:     auto coo_elem_fullNb = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>((void *)SData_d->coo_elem_fullNb);
    330:     delete coo_elem_fullNb;
    331:     auto coo_elem_point_offsets = static_cast<Kokkos::View<LandauIdx *, Kokkos::LayoutLeft> *>((void *)SData_d->coo_elem_point_offsets);
    332:     delete coo_elem_point_offsets;
    333:     SData_d->coo_elem_offsets       = NULL;
    334:     SData_d->coo_elem_point_offsets = NULL;
    335:     SData_d->coo_elem_fullNb        = NULL;
    336:     auto coo_vals                   = static_cast<Kokkos::View<PetscScalar *, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace> *>((void *)SData_d->coo_vals);
    337:     delete coo_vals;
    338:   }
    339:   PetscFunctionReturn(PETSC_SUCCESS);
    340: }

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

    344: KOKKOS_INLINE_FUNCTION
    345: PetscErrorCode landau_mat_assemble(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_NQND + 1], const PetscInt glb_elem_idx, const PetscInt bid_coo_sz_batch)
    346: {
    347:   PetscInt               idx, q, nr, nc;
    348:   const LandauIdx *const Idxs                         = &d_maps->gIdx[elem][fieldA][0];
    349:   PetscScalar            row_scale[LANDAU_MAX_Q_FACE] = {0}, col_scale[LANDAU_MAX_Q_FACE] = {0};
    350:   { // mirror (i,j) in CreateStaticGPUData
    351:     const int fullNb = coo_elem_fullNb[glb_elem_idx], fullNb2 = fullNb * fullNb;
    352:     idx = Idxs[f];
    353:     if (idx >= 0) {
    354:       nr           = 1;
    355:       row_scale[0] = 1.;
    356:     } else {
    357:       idx = -idx - 1;
    358:       for (q = 0, nr = 0; q < d_maps->num_face; q++, nr++) {
    359:         if (d_maps->c_maps[idx][q].gid < 0) break;
    360:         row_scale[q] = d_maps->c_maps[idx][q].scale;
    361:       }
    362:     }
    363:     idx = Idxs[g];
    364:     if (idx >= 0) {
    365:       nc           = 1;
    366:       col_scale[0] = 1.;
    367:     } else {
    368:       idx = -idx - 1;
    369:       nc  = d_maps->num_face;
    370:       for (q = 0, nc = 0; q < d_maps->num_face; q++, nc++) {
    371:         if (d_maps->c_maps[idx][q].gid < 0) break;
    372:         col_scale[q] = d_maps->c_maps[idx][q].scale;
    373:       }
    374:     }
    375:     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];
    376:     for (int q = 0, idx2 = idx0; q < nr; q++) {
    377:       for (int d = 0; d < nc; d++, idx2++) coo_vals[idx2] = row_scale[q] * col_scale[d] * Aij;
    378:     }
    379:   }
    380:   return PETSC_SUCCESS;
    381: }

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

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

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

    586:     PetscCall(PetscDeviceGetDefault_Internal(&device));
    587:     PetscCall(PetscDeviceGetAttribute(device, PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK, &maximum_shared_mem_size));
    588:     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));
    589:     const int jac_shared_level = (jac_scr_bytes > maximum_shared_mem_size) ? 1 : KOKKOS_SHARED_LEVEL;
    590:     // device function/lambda
    591:     auto jac_lambda = KOKKOS_LAMBDA(const team_member team)
    592:     {
    593:       const PetscInt b_Nelem = d_elem_offset[num_grids], b_elem_idx = team.league_rank() % b_Nelem, b_id = team.league_rank() / b_Nelem;
    594:       // find my grid
    595:       PetscInt grid = 0;
    596:       while (b_elem_idx >= d_elem_offset[grid + 1]) grid++;
    597:       {
    598:         const PetscInt loc_Nf = d_species_offset[grid + 1] - d_species_offset[grid], loc_elem = b_elem_idx - d_elem_offset[grid];
    599:         const PetscInt moffset = LAND_MOFFSET(b_id, grid, batch_sz, num_grids, d_mat_offset);
    600:         const PetscInt f_off   = d_species_offset[grid];
    601:   #ifdef LAND_SUPPORT_CPU_ASS
    602:         const PetscInt totDim = loc_Nf * Nb;
    603:   #endif
    604:         g2_scr_t g2(team.team_scratch(jac_shared_level), dim, loc_Nf, Nq);
    605:         g3_scr_t g3(team.team_scratch(jac_shared_level), dim, dim, loc_Nf, Nq);
    606:         g2_scr_t gg2(team.team_scratch(jac_shared_level), dim, loc_Nf, Nq);
    607:         g3_scr_t gg3(team.team_scratch(jac_shared_level), dim, dim, loc_Nf, Nq);
    608:         // get g2[] & g3[]
    609:         Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, Nq), [=](int myQi) {
    610:           using Kokkos::parallel_reduce;
    611:           const PetscInt                    jpidx_glb = d_ip_offset[grid] + loc_elem * Nq + myQi;
    612:           const PetscReal                  *invJ      = &d_invJ_k(jpidx_glb * dim * dim);
    613:           const PetscReal                   vj[3] = {d_x[jpidx_glb], d_y[jpidx_glb], d_z ? d_z[jpidx_glb] : 0}, wj = d_w[jpidx_glb];
    614:           landau_inner_red::TensorValueType gg_temp; // reduce on part of gg2 and g33 for IP jpidx_g
    615:           Kokkos::parallel_reduce(
    616:             Kokkos::ThreadVectorRange(team, (int)d_ip_offset[num_grids]),
    617:             [=](const int &ipidx, landau_inner_red::TensorValueType &ggg) {
    618:               const PetscReal wi = d_w[ipidx], x = d_x[ipidx], y = d_y[ipidx];
    619:               PetscReal       temp1[3] = {0, 0, 0}, temp2 = 0;
    620:               PetscInt        fieldB, d2, d3, f_off_r, grid_r, ipidx_g, nip_loc_r, loc_Nf_r;
    621:   #if LANDAU_DIM == 2
    622:               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.;
    623:               LandauTensor2D(vj, x, y, Ud, Uk, mask);
    624:   #else
    625:                 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.;
    626:                 LandauTensor3D(vj, x, y, z, U, mask);
    627:   #endif
    628:               grid_r = 0;
    629:               while (ipidx >= d_ip_offset[grid_r + 1]) grid_r++; // yuck search for grid
    630:               f_off_r   = d_species_offset[grid_r];
    631:               ipidx_g   = ipidx - d_ip_offset[grid_r];
    632:               nip_loc_r = d_numCells[grid_r] * Nq;
    633:               loc_Nf_r  = d_species_offset[grid_r + 1] - d_species_offset[grid_r];
    634:               for (fieldB = 0; fieldB < loc_Nf_r; ++fieldB) { // fieldB is \beta  d_lambdas[grid][grid_r]
    635:                 const PetscInt  idx = d_ipf_offset[grid_r] + fieldB * nip_loc_r + ipidx_g;
    636:                 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;
    637:                 temp1[0] += d_fdf_k(b_id, 1, idx) * ff2;
    638:                 temp1[1] += d_fdf_k(b_id, 2, idx) * ff2;
    639:   #if LANDAU_DIM == 3
    640:                 temp1[2] += d_fdf_k(b_id, 3, idx) * ff2;
    641:   #endif
    642:                 temp2 += d_fdf_k(b_id, 0, idx) * ff1;
    643:               }
    644:               temp1[0] *= wi;
    645:               temp1[1] *= wi;
    646:   #if LANDAU_DIM == 3
    647:               temp1[2] *= wi;
    648:   #endif
    649:               temp2 *= wi;
    650:   #if LANDAU_DIM == 2
    651:               for (d2 = 0; d2 < 2; d2++) {
    652:                 for (d3 = 0; d3 < 2; ++d3) {
    653:                   /* K = U * grad(f): g2=e: i,A */
    654:                   ggg.gg2[d2] += Uk[d2][d3] * temp1[d3];
    655:                   /* D = -U * (I \kron (fx)): g3=f: i,j,A */
    656:                   ggg.gg3[d2][d3] += Ud[d2][d3] * temp2;
    657:                 }
    658:               }
    659:   #else
    660:                 for (d2 = 0; d2 < 3; ++d2) {
    661:                   for (d3 = 0; d3 < 3; ++d3) {
    662:                     /* K = U * grad(f): g2 = e: i,A */
    663:                     ggg.gg2[d2] += U[d2][d3]*temp1[d3];
    664:                     /* D = -U * (I \kron (fx)): g3 = f: i,j,A */
    665:                     ggg.gg3[d2][d3] += U[d2][d3]*temp2;
    666:                   }
    667:                 }
    668:   #endif
    669:             },
    670:             Kokkos::Sum<landau_inner_red::TensorValueType>(gg_temp));
    671:           // add alpha and put in gg2/3
    672:           Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, (int)loc_Nf), [&](const int &fieldA) { // \alpha
    673:             PetscInt d2, d3;
    674:             for (d2 = 0; d2 < dim; d2++) {
    675:               gg2(d2, fieldA, myQi) = gg_temp.gg2[d2] * d_alpha[fieldA + f_off];
    676:               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];
    677:             }
    678:           });
    679:           /* add electric field term once per IP */
    680:           Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, (int)loc_Nf), [&](const int &fieldA) { gg2(dim - 1, fieldA, myQi) += d_Eq_m[fieldA + f_off]; });
    681:           Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, (int)loc_Nf), [=](const int &fieldA) {
    682:             int d, d2, d3, dp;
    683:             /* Jacobian transform - g2, g3 - per thread (2D) */
    684:             for (d = 0; d < dim; ++d) {
    685:               g2(d, fieldA, myQi) = 0;
    686:               for (d2 = 0; d2 < dim; ++d2) {
    687:                 g2(d, fieldA, myQi) += invJ[d * dim + d2] * gg2(d2, fieldA, myQi);
    688:                 g3(d, d2, fieldA, myQi) = 0;
    689:                 for (d3 = 0; d3 < dim; ++d3) {
    690:                   for (dp = 0; dp < dim; ++dp) g3(d, d2, fieldA, myQi) += invJ[d * dim + d3] * gg3(d3, dp, fieldA, myQi) * invJ[d2 * dim + dp];
    691:                 }
    692:                 g3(d, d2, fieldA, myQi) *= wj;
    693:               }
    694:               g2(d, fieldA, myQi) *= wj;
    695:             }
    696:           });
    697:         }); // Nq
    698:         team.team_barrier();
    699:         { /* assemble */
    700:           for (PetscInt fieldA = 0; fieldA < loc_Nf; fieldA++) {
    701:             /* assemble */
    702:             Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, Nb), [=](int f) {
    703:               Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, 0, Nb), [=](int g) {
    704:                 PetscScalar t = 0;
    705:                 for (int qj = 0; qj < Nq; qj++) { // look at others integration points
    706:                   const PetscReal *BJq = &d_BB[qj * Nb], *DIq = &d_DD[qj * Nb * dim];
    707:                   for (int d = 0; d < dim; ++d) {
    708:                     t += DIq[f * dim + d] * g2(d, fieldA, qj) * BJq[g];
    709:                     for (int d2 = 0; d2 < dim; ++d2) t += DIq[f * dim + d] * g3(d, d2, fieldA, qj) * DIq[g * dim + d2];
    710:                   }
    711:                 }
    712:                 if (elem_mat_num_cells_max_grid) { // CPU assembly
    713:   #ifdef LAND_SUPPORT_CPU_ASS
    714:                   const PetscInt fOff                     = (fieldA * Nb + f) * totDim + fieldA * Nb + g;
    715:                   d_elem_mats(b_id, grid, loc_elem, fOff) = t;
    716:   #endif
    717:                 } else {
    718:                   static_cast<void>(landau_mat_assemble(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));
    719:                 }
    720:               });
    721:             });
    722:           }
    723:         }
    724:       } // scope with 'grid'
    725:     };
    726:   #if defined(PETSC_HAVE_HIP)
    727:     const int lbound2 = 1;
    728:   #else
    729:     const int lbound2 = 2;
    730:   #endif
    731:     PetscCall(PetscLogEventBegin(events[4], 0, 0, 0, 0));
    732:     PetscCall(PetscLogGpuTimeBegin());
    733:     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));
    734:     Kokkos::parallel_for("Jacobian", Kokkos::TeamPolicy<Kokkos::LaunchBounds<256, lbound2>>(num_cells_batch * batch_sz, team_size, vector_size).set_scratch_size(jac_shared_level, Kokkos::PerTeam(jac_scr_bytes)), jac_lambda);
    735:     Kokkos::fence();
    736:     PetscCall(PetscLogGpuTimeEnd());
    737:     PetscCall(PetscLogEventEnd(events[4], 0, 0, 0, 0));
    738:     if (d_vertex_f_k) delete d_vertex_f_k;
    739:   } else { // mass
    740:     PetscCall(PetscLogEventBegin(events[16], 0, 0, 0, 0));
    741:     PetscCall(PetscLogGpuTimeBegin());
    742:     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"));
    743:     Kokkos::parallel_for(
    744:       "Mass", Kokkos::TeamPolicy<>(num_cells_batch * batch_sz, team_size, vector_size), KOKKOS_LAMBDA(const team_member team) {
    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 * Nb;
    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:                   static_cast<void>(landau_mat_assemble(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:     PetscCall(PetscLogGpuTimeEnd());
    784:     PetscCall(PetscLogEventEnd(events[16], 0, 0, 0, 0));
    785:   }
    786:   if (d_coo_vals) {
    787:     PetscCall(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);
    795:     PetscCheck(!container, PETSC_COMM_SELF, PETSC_ERR_PLIB, "?????");
    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:         PetscCall(PetscLogEventBegin(events[5], 0, 0, 0, 0));
    805:         PetscCall(DMPlexGetHeightStratum(plex[grid], 0, &cStart, &cEnd));
    806:         PetscCall(DMGetLocalSection(plex[grid], &section));
    807:         PetscCall(DMGetGlobalSection(plex[grid], &globalSection));
    808:         PetscCall(PetscLogEventEnd(events[5], 0, 0, 0, 0));
    809:         PetscCall(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:           PetscCall(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 * Nb;
    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:         PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
    826:         PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
    827:         PetscCall(MatGetSize(B, &nloc, NULL));
    828:         for (int i = 0; i < nloc; i++) {
    829:           PetscCall(MatGetRow(B, i, &nzl, &cols, &vals));
    830:           PetscCheck(nzl <= 1024, PetscObjectComm((PetscObject)B), PETSC_ERR_PLIB, "Row too big: %" PetscInt_FMT, nzl);
    831:           for (int j = 0; j < nzl; j++) colbuf[j] = cols[j] + moffset;
    832:           row = i + moffset;
    833:           PetscCall(MatSetValues(JacP, 1, &row, nzl, colbuf, vals, ADD_VALUES));
    834:           PetscCall(MatRestoreRow(B, i, &nzl, &cols, &vals));
    835:         }
    836:         PetscCall(MatDestroy(&subJ[LAND_PACK_IDX(b_id, grid)]));
    837:         PetscCall(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:   PetscFunctionReturn(PETSC_SUCCESS);
    845: }
    846: } // extern "C"
    847: #endif