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