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