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