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