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.Nq = Nq; 93 h_maps.c_maps = (pointInterpolationP4est (*)[LANDAU_MAX_Q_FACE]) d_points->data(); 94 maps[grid].vp1 = (void*)d_points; 95 h_maps.gIdx = (LandauIdx (*)[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ]) d_gidx->data(); 96 maps[grid].vp2 = (void*)d_gidx; 97 { 98 Kokkos::View<P4estVertexMaps, Kokkos::HostSpace> h_maps_k(&h_maps); 99 Kokkos::View<P4estVertexMaps> *d_maps_k = new Kokkos::View<P4estVertexMaps>(Kokkos::create_mirror(Kokkos::DefaultExecutionSpace::memory_space(),h_maps_k)); 100 Kokkos::deep_copy (*d_maps_k, h_maps_k); 101 maps[grid].d_self = d_maps_k->data(); 102 maps[grid].vp3 = (void*)d_maps_k; 103 } 104 PetscFunctionReturn(0); 105 } 106 PetscErrorCode LandauKokkosDestroyMatMaps(P4estVertexMaps maps[], PetscInt num_grids) 107 { 108 PetscFunctionBegin; 109 for (PetscInt grid=0;grid<num_grids;grid++) { 110 auto a = static_cast<Kokkos::View<pointInterpolationP4est*[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight>*>(maps[grid].vp1); 111 auto b = static_cast<Kokkos::View<LandauIdx*[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ], Kokkos::LayoutRight>*>(maps[grid].vp2); 112 auto c = static_cast<Kokkos::View<P4estVertexMaps*>*>(maps[grid].vp3); 113 delete a; delete b; delete c; 114 } 115 PetscFunctionReturn(0); 116 } 117 118 PetscErrorCode LandauKokkosStaticDataSet(DM plex, const PetscInt Nq, const PetscInt batch_sz, const PetscInt num_grids, 119 PetscInt a_numCells[], PetscInt a_species_offset[], PetscInt a_mat_offset[], 120 PetscReal a_nu_alpha[], PetscReal a_nu_beta[], PetscReal a_invMass[], PetscReal a_invJ[], 121 PetscReal a_x[], PetscReal a_y[], PetscReal a_z[], PetscReal a_w[], LandauStaticData *SData_d) 122 { 123 PetscReal *BB,*DD; 124 PetscErrorCode ierr; 125 PetscTabulation *Tf; 126 PetscInt dim; 127 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; 128 PetscDS prob; 129 130 PetscFunctionBegin; 131 ierr = DMGetDimension(plex, &dim);CHKERRQ(ierr); 132 ierr = DMGetDS(plex, &prob);CHKERRQ(ierr); 133 if (LANDAU_DIM != dim) SETERRQ2(PETSC_COMM_WORLD, PETSC_ERR_PLIB, "dim %D != LANDAU_DIM %d",dim,LANDAU_DIM); 134 ierr = PetscDSGetTabulation(prob, &Tf);CHKERRQ(ierr); 135 BB = Tf[0]->T[0]; DD = Tf[0]->T[1]; 136 ip_offset[0] = ipf_offset[0] = elem_offset[0] = 0; 137 nip = 0; 138 IPf_sz = 0; 139 for (PetscInt grid=0 ; grid<num_grids ; grid++) { 140 PetscInt nfloc = a_species_offset[grid+1] - a_species_offset[grid]; 141 elem_offset[grid+1] = elem_offset[grid] + a_numCells[grid]; 142 nip += a_numCells[grid]*Nq; 143 ip_offset[grid+1] = nip; 144 IPf_sz += Nq*nfloc*a_numCells[grid]; 145 ipf_offset[grid+1] = IPf_sz; 146 } 147 Nftot = a_species_offset[num_grids]; 148 ierr = PetscKokkosInitializeCheck();CHKERRQ(ierr); 149 { 150 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_alpha (a_nu_alpha, Nftot); 151 auto alpha = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("alpha", Nftot); 152 SData_d->alpha = static_cast<void*>(alpha); 153 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_beta (a_nu_beta, Nftot); 154 auto beta = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("beta", Nftot); 155 SData_d->beta = static_cast<void*>(beta); 156 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_invMass (a_invMass,Nftot); 157 auto invMass = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("invMass", Nftot); 158 SData_d->invMass = static_cast<void*>(invMass); 159 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_BB (BB,Nq*Nb); 160 auto B = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("B", Nq*Nb); 161 SData_d->B = static_cast<void*>(B); 162 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_DD (DD,Nq*Nb*dim); 163 auto D = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("D", Nq*Nb*dim); 164 SData_d->D = static_cast<void*>(D); 165 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_invJ (a_invJ, nip*dim*dim); 166 auto invJ = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("invJ", nip*dim*dim); 167 SData_d->invJ = static_cast<void*>(invJ); 168 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_x (a_x, nip); 169 auto x = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("x", nip); 170 SData_d->x = static_cast<void*>(x); 171 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_y (a_y, nip); 172 auto y = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("y", nip); 173 SData_d->y = static_cast<void*>(y); 174 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_w (a_w, nip); 175 auto w = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("w", nip); 176 SData_d->w = static_cast<void*>(w); 177 178 Kokkos::deep_copy (*alpha, h_alpha); 179 Kokkos::deep_copy (*beta, h_beta); 180 Kokkos::deep_copy (*invMass, h_invMass); 181 Kokkos::deep_copy (*B, h_BB); 182 Kokkos::deep_copy (*D, h_DD); 183 Kokkos::deep_copy (*invJ, h_invJ); 184 Kokkos::deep_copy (*x, h_x); 185 Kokkos::deep_copy (*y, h_y); 186 Kokkos::deep_copy (*w, h_w); 187 188 if (dim==3) { 189 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_z (a_z , dim==3 ? nip : 0); 190 auto z = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("z", nip); 191 SData_d->z = static_cast<void*>(z); 192 Kokkos::deep_copy (*z, h_z); 193 } else SData_d->z = NULL; 194 195 // 196 const Kokkos::View<PetscInt*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_NCells (a_numCells, num_grids); 197 auto nc = new Kokkos::View<PetscInt*, Kokkos::LayoutLeft> ("NCells",num_grids); 198 SData_d->NCells = static_cast<void*>(nc); 199 Kokkos::deep_copy (*nc, h_NCells); 200 201 const Kokkos::View<PetscInt*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_species_offset (a_species_offset, num_grids+1); 202 auto soff = new Kokkos::View<PetscInt*, Kokkos::LayoutLeft> ("species_offset",num_grids+1); 203 SData_d->species_offset = static_cast<void*>(soff); 204 Kokkos::deep_copy (*soff, h_species_offset); 205 206 const Kokkos::View<PetscInt*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_mat_offset (a_mat_offset, num_grids+1); 207 auto moff = new Kokkos::View<PetscInt*, Kokkos::LayoutLeft> ("mat_offset",num_grids+1); 208 SData_d->mat_offset = static_cast<void*>(moff); 209 Kokkos::deep_copy (*moff, h_mat_offset); 210 211 const Kokkos::View<PetscInt*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_ip_offset (ip_offset, num_grids+1); 212 auto ipoff = new Kokkos::View<PetscInt*, Kokkos::LayoutLeft> ("ip_offset",num_grids+1); 213 SData_d->ip_offset = static_cast<void*>(ipoff); 214 Kokkos::deep_copy (*ipoff, h_ip_offset); 215 216 const Kokkos::View<PetscInt*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_elem_offset (elem_offset, num_grids+1); 217 auto eoff = new Kokkos::View<PetscInt*, Kokkos::LayoutLeft> ("elem_offset",num_grids+1); 218 SData_d->elem_offset = static_cast<void*>(eoff); 219 Kokkos::deep_copy (*eoff, h_elem_offset); 220 221 const Kokkos::View<PetscInt*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_ipf_offset (ipf_offset, num_grids+1); 222 auto ipfoff = new Kokkos::View<PetscInt*, Kokkos::LayoutLeft> ("ipf_offset",num_grids+1); 223 SData_d->ipf_offset = static_cast<void*>(ipfoff); 224 Kokkos::deep_copy (*ipfoff, h_ipf_offset); 225 #if defined(LANDAU_LAYOUT_LEFT) // preallocate dynamic data, no copy 226 auto ipfdf_data = new Kokkos::View<PetscReal***, Kokkos::LayoutLeft > ("fdf", batch_sz, dim+1, IPf_sz); 227 #else 228 auto ipfdf_data = new Kokkos::View<PetscReal***, Kokkos::LayoutRight > ("fdf",batch_sz, dim+1, IPf_sz); 229 #endif 230 SData_d->ipfdf_data = static_cast<void*>(ipfdf_data); 231 auto Eq_m = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("Eq_m",Nftot); // allocate but do not set, same for whole batch 232 SData_d->Eq_m = static_cast<void*>(Eq_m); 233 } 234 SData_d->maps = NULL; // not used 235 PetscFunctionReturn(0); 236 } 237 238 PetscErrorCode LandauKokkosStaticDataClear(LandauStaticData *SData_d) 239 { 240 PetscFunctionBegin; 241 if (SData_d->alpha) { 242 auto alpha = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->alpha); 243 delete alpha; 244 SData_d->alpha = NULL; 245 auto beta = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->beta); 246 delete beta; 247 auto invMass = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->invMass); 248 delete invMass; 249 auto B = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->B); 250 delete B; 251 auto D = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->D); 252 delete D; 253 auto invJ = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->invJ); 254 delete invJ; 255 auto x = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->x); 256 delete x; 257 auto y = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->y); 258 delete y; 259 if (SData_d->z) { 260 auto z = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->z); 261 delete z; 262 } 263 #if defined(LANDAU_LAYOUT_LEFT) // preallocate dynamic data, no copy 264 auto z = static_cast<Kokkos::View<PetscReal***, Kokkos::LayoutLeft>*>(SData_d->ipfdf_data); 265 #else 266 auto z = static_cast<Kokkos::View<PetscReal***, Kokkos::LayoutRight>*>(SData_d->ipfdf_data); 267 #endif 268 delete z; 269 auto w = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->w); 270 delete w; 271 auto Eq_m = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->Eq_m); 272 delete Eq_m; 273 // offset 274 auto nc = static_cast<Kokkos::View<PetscInt*, Kokkos::LayoutLeft>*>(SData_d->NCells); 275 delete nc; 276 auto soff = static_cast<Kokkos::View<PetscInt*, Kokkos::LayoutLeft>*>(SData_d->species_offset); 277 delete soff; 278 auto moff = static_cast<Kokkos::View<PetscInt*, Kokkos::LayoutLeft>*>(SData_d->mat_offset); 279 delete moff; 280 auto ipoff = static_cast<Kokkos::View<PetscInt*, Kokkos::LayoutLeft>*>(SData_d->ip_offset); 281 delete ipoff; 282 auto eoff = static_cast<Kokkos::View<PetscInt*, Kokkos::LayoutLeft>*>(SData_d->elem_offset); 283 delete eoff; 284 auto ipfoff = static_cast<Kokkos::View<PetscInt*, Kokkos::LayoutLeft>*>(SData_d->ipf_offset); 285 delete ipfoff; 286 } 287 PetscFunctionReturn(0); 288 } 289 290 #define KOKKOS_SHARED_LEVEL 1 291 KOKKOS_INLINE_FUNCTION 292 PetscErrorCode landau_mat_assemble(PetscSplitCSRDataStructure d_mat, const team_member team, 293 Kokkos::View<PetscScalar**, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace::scratch_memory_space> s_fieldMats, 294 Kokkos::View<PetscInt**, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace::scratch_memory_space> s_idx, 295 Kokkos::View<PetscReal**, Kokkos::LayoutRight, Kokkos::DefaultExecutionSpace::scratch_memory_space> s_scale, 296 const PetscInt Nb, const PetscInt Nq, const PetscInt nfaces, const PetscInt moffset, const PetscInt elem, const PetscInt fieldA, const P4estVertexMaps *d_maps) 297 { 298 const LandauIdx *const Idxs = &d_maps->gIdx[elem][fieldA][0]; 299 team.team_barrier(); 300 Kokkos::parallel_for(Kokkos::TeamVectorRange(team,0,Nb), [=] (int f) { 301 PetscInt q, idx = Idxs[f]; 302 if (idx >= 0) { 303 s_idx(f,0) = idx + moffset; 304 s_scale(f,0) = 1.; 305 } else { 306 idx = -idx - 1; 307 for (q = 0; q < nfaces; q++) { 308 s_idx(f,q) = d_maps->c_maps[idx][q].gid + moffset; 309 s_scale(f,q) = d_maps->c_maps[idx][q].scale; 310 } 311 } 312 }); 313 team.team_barrier(); 314 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nb), [=] (int f) { 315 PetscInt nr,idx = Idxs[f]; 316 if (idx >= 0) { 317 nr = 1; 318 } else { 319 nr = nfaces; 320 } 321 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,Nb), [=] (int g) { 322 PetscScalar vals[LANDAU_MAX_Q_FACE*LANDAU_MAX_Q_FACE]; 323 PetscInt q,d,nc,idx = Idxs[g]; 324 if (idx >= 0) { 325 nc = 1; 326 } else { 327 nc = nfaces; 328 } 329 for (q = 0; q < nr; q++) { 330 for (d = 0; d < nc; d++) { 331 vals[q*nc + d] = s_scale(f,q)*s_scale(g,d)*s_fieldMats(f,g); 332 } 333 } 334 MatSetValuesDevice(d_mat,nr,&s_idx(f,0),nc,&s_idx(g,0),vals,ADD_VALUES); 335 }); 336 }); 337 return 0; 338 } 339 340 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[], 341 const PetscInt N, const PetscScalar a_xarray[], const LandauStaticData *SData_d, const PetscInt num_sub_blocks, const PetscReal shift, 342 const PetscLogEvent events[], const PetscInt a_mat_offset[], const PetscInt a_species_offset[], Mat subJ[], Mat JacP) 343 { 344 using scr_mem_t = Kokkos::DefaultExecutionSpace::scratch_memory_space; 345 using fieldMats_scr_t = Kokkos::View<PetscScalar**, Kokkos::LayoutRight, scr_mem_t>; 346 using idx_scr_t = Kokkos::View<PetscInt**, Kokkos::LayoutRight, scr_mem_t>; 347 using scale_scr_t = Kokkos::View<PetscReal**, Kokkos::LayoutRight, scr_mem_t>; 348 using g2_scr_t = Kokkos::View<PetscReal***, Kokkos::LayoutRight, scr_mem_t>; 349 using g3_scr_t = Kokkos::View<PetscReal****, Kokkos::LayoutRight, scr_mem_t>; 350 PetscErrorCode ierr; 351 PetscInt Nb=Nq,dim,num_cells_max,Nf_max,num_cells_batch; 352 int nfaces=0; 353 LandauCtx *ctx; 354 PetscReal *d_Eq_m=NULL; 355 PetscScalar *d_vertex_f=NULL; 356 P4estVertexMaps *maps[LANDAU_MAX_GRIDS]; // this gets captured 357 PetscSplitCSRDataStructure d_mat; 358 PetscContainer container; 359 const int conc = Kokkos::DefaultExecutionSpace().concurrency(), openmp = !!(conc < 1000), team_size = (openmp==0) ? Nq : 1; 360 auto d_alpha_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->alpha); //static data 361 const PetscReal *d_alpha = d_alpha_k->data(); 362 const PetscInt Nftot = d_alpha_k->size(); // total number of species 363 auto d_beta_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->beta); 364 const PetscReal *d_beta = d_beta_k->data(); 365 auto d_invMass_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->invMass); 366 const PetscReal *d_invMass = d_invMass_k->data(); 367 auto d_B_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->B); 368 const PetscReal *d_BB = d_B_k->data(); 369 auto d_D_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->D); 370 const PetscReal *d_DD = d_D_k->data(); 371 auto d_invJ_k = *static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->invJ); // use Kokkos vector in kernels 372 auto d_x_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->x); //static data 373 const PetscReal *d_x = d_x_k->data(); 374 auto d_y_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->y); //static data 375 const PetscReal *d_y = d_y_k->data(); 376 auto d_z_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->z); //static data 377 const PetscReal *d_z = (LANDAU_DIM==3) ? d_z_k->data() : NULL; 378 auto d_w_k = *static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->w); //static data 379 const PetscReal *d_w = d_w_k.data(); 380 // grid offsets - single vertex grid data 381 auto d_numCells_k = static_cast<Kokkos::View<PetscInt*, Kokkos::LayoutLeft>*>(SData_d->NCells); 382 const PetscInt *d_numCells = d_numCells_k->data(); 383 auto d_species_offset_k = static_cast<Kokkos::View<PetscInt*, Kokkos::LayoutLeft>*>(SData_d->species_offset); 384 const PetscInt *d_species_offset = d_species_offset_k->data(); 385 auto d_mat_offset_k = static_cast<Kokkos::View<PetscInt*, Kokkos::LayoutLeft>*>(SData_d->mat_offset); 386 const PetscInt *d_mat_offset = d_mat_offset_k->data(); 387 auto d_ip_offset_k = static_cast<Kokkos::View<PetscInt*, Kokkos::LayoutLeft>*>(SData_d->ip_offset); 388 const PetscInt *d_ip_offset = d_ip_offset_k->data(); 389 auto d_ipf_offset_k = static_cast<Kokkos::View<PetscInt*, Kokkos::LayoutLeft>*>(SData_d->ipf_offset); 390 const PetscInt *d_ipf_offset = d_ipf_offset_k->data(); 391 auto d_elem_offset_k = static_cast<Kokkos::View<PetscInt*, Kokkos::LayoutLeft>*>(SData_d->elem_offset); 392 const PetscInt *d_elem_offset = d_elem_offset_k->data(); 393 #if defined(LANDAU_LAYOUT_LEFT) // preallocate dynamic data, including batched vertices 394 Kokkos::View<PetscReal***, Kokkos::LayoutLeft > d_fdf_k = *static_cast<Kokkos::View<PetscReal***, Kokkos::LayoutLeft >*>(SData_d->ipfdf_data); 395 #else 396 Kokkos::View<PetscReal***, Kokkos::LayoutRight > d_fdf_k = *static_cast<Kokkos::View<PetscReal***, Kokkos::LayoutRight >*>(SData_d->ipfdf_data); 397 #endif 398 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 399 400 PetscFunctionBegin; 401 ierr = PetscLogEventBegin(events[3],0,0,0,0);CHKERRQ(ierr); 402 ierr = DMGetApplicationContext(plex[0], &ctx);CHKERRQ(ierr); 403 if (!ctx) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "no context"); 404 ierr = DMGetDimension(plex[0], &dim);CHKERRQ(ierr); 405 if (LANDAU_DIM != dim) SETERRQ2(PETSC_COMM_WORLD, PETSC_ERR_PLIB, "dim %D != LANDAU_DIM %d",dim,LANDAU_DIM); 406 if (ctx->gpu_assembly) { 407 ierr = PetscObjectQuery((PetscObject) JacP, "assembly_maps", (PetscObject *) &container);CHKERRQ(ierr); 408 if (container) { // not here first call 409 P4estVertexMaps *h_maps=NULL; 410 ierr = PetscContainerGetPointer(container, (void **) &h_maps);CHKERRQ(ierr); 411 for (PetscInt grid=0 ; grid<num_grids ; grid++) { 412 if (h_maps[grid].d_self) { 413 maps[grid] = h_maps[grid].d_self; 414 nfaces = h_maps[grid].num_face; // nface=0 for CPU assembly 415 } else { 416 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "GPU assembly but no metadata in container"); 417 } 418 } 419 // this does the setup the first time called 420 ierr = MatKokkosGetDeviceMatWrite(JacP,&d_mat);CHKERRQ(ierr); 421 } else { // kernel output - first call assembled on device 422 for (PetscInt grid=0 ; grid<num_grids ; grid++) maps[grid] = NULL; 423 nfaces = 0; 424 d_mat = NULL; 425 } 426 } else { 427 for (PetscInt grid=0 ; grid<num_grids ; grid++) maps[grid] = NULL; 428 nfaces = 0; 429 d_mat = NULL; 430 container = NULL; 431 } 432 num_cells_batch = Nf_max = num_cells_max = 0; 433 for (PetscInt grid=0 ; grid<num_grids ; grid++) { 434 int Nfloc = a_species_offset[grid+1] - a_species_offset[grid]; 435 if (Nfloc > Nf_max) Nf_max = Nfloc; 436 if (a_numCells[grid] > num_cells_max) num_cells_max = a_numCells[grid]; 437 num_cells_batch += a_numCells[grid]; // we don't have a host element offset here (but in ctx) 438 } 439 const PetscInt totDim_max = Nf_max*Nq, elem_mat_size_max = totDim_max*totDim_max; 440 const PetscInt elem_mat_num_cells_max_grid = container ? 0 : num_cells_max; 441 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 442 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_Eq_m_k (a_Eq_m, Nftot); 443 if (a_elem_closure || a_xarray) { 444 Kokkos::deep_copy (*d_Eq_m_k, h_Eq_m_k); 445 d_Eq_m = d_Eq_m_k->data(); 446 } else d_Eq_m = NULL; 447 ierr = PetscKokkosInitializeCheck();CHKERRQ(ierr); 448 ierr = PetscLogEventEnd(events[3],0,0,0,0);CHKERRQ(ierr); 449 if (a_elem_closure || a_xarray) { // Jacobian, create f & df 450 Kokkos::View<PetscScalar*, Kokkos::LayoutLeft> *d_vertex_f_k = NULL; 451 ierr = PetscLogEventBegin(events[1],0,0,0,0);CHKERRQ(ierr); 452 if (a_elem_closure) { 453 PetscInt closure_sz = 0; // argh, don't have this on the host!!! 454 for (PetscInt grid=0 ; grid<num_grids ; grid++) { 455 PetscInt nfloc = a_species_offset[grid+1] - a_species_offset[grid]; 456 closure_sz += Nq*nfloc*a_numCells[grid]; 457 } 458 closure_sz *= batch_sz; 459 d_vertex_f_k = new Kokkos::View<PetscScalar*, Kokkos::LayoutLeft> ("closure",closure_sz); 460 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 461 Kokkos::deep_copy (*d_vertex_f_k, h_closure_k); 462 d_vertex_f = d_vertex_f_k->data(); 463 } else { 464 d_vertex_f = (PetscScalar*)a_xarray; 465 } 466 ierr = PetscLogEventEnd(events[1],0,0,0,0);CHKERRQ(ierr); 467 ierr = PetscLogEventBegin(events[8],0,0,0,0);CHKERRQ(ierr); 468 ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr); 469 Kokkos::parallel_for("f, df", Kokkos::TeamPolicy<>(num_cells_batch*batch_sz, team_size, /* Kokkos::AUTO */ 16), KOKKOS_LAMBDA (const team_member team) { 470 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]; 471 // find my grid 472 PetscInt grid = 0; 473 while (b_elem_idx >= d_elem_offset[grid+1]) grid++; 474 { 475 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]; 476 const PetscInt moffset = LAND_MOFFSET(b_id,grid,batch_sz,num_grids,d_mat_offset); 477 const PetscScalar *coef; 478 PetscScalar coef_buff[LANDAU_MAX_SPECIES*LANDAU_MAX_NQ]; 479 const PetscReal *invJe = &d_invJ_k( (d_ip_offset[grid] + loc_elem*Nq)*dim*dim); 480 // un pack IPData 481 if (!maps[grid]) { 482 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 483 } else { 484 coef = coef_buff; 485 for (int f = 0; f < loc_Nf; ++f) { 486 LandauIdx *const Idxs = &maps[grid]->gIdx[loc_elem][f][0]; 487 for (int b = 0; b < Nb; ++b) { 488 PetscInt idx = Idxs[b]; 489 if (idx >= 0) { 490 coef_buff[f*Nb+b] = d_vertex_f[idx+moffset]; // xarray data, not IP 491 } else { 492 idx = -idx - 1; 493 coef_buff[f*Nb+b] = 0; 494 for (int q = 0; q < maps[grid]->num_face; q++) { 495 PetscInt id = maps[grid]->c_maps[idx][q].gid; 496 PetscScalar scale = maps[grid]->c_maps[idx][q].scale; 497 coef_buff[f*Nb+b] += scale*d_vertex_f[id+moffset]; 498 } 499 } 500 } 501 } 502 } 503 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nq), [=] (int myQi) { 504 const PetscReal *const invJ = &invJe[myQi*dim*dim]; // b_elem_idx: batch element index 505 const PetscReal *Bq = &d_BB[myQi*Nb], *Dq = &d_DD[myQi*Nb*dim]; 506 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,(int)loc_Nf), [=] (int f) { 507 PetscInt b, e, d; 508 PetscReal refSpaceDer[LANDAU_DIM]; 509 const PetscInt idx = d_ipf_offset[grid] + f*loc_nip + loc_elem*Nq + myQi; 510 d_fdf_k(b_id,0,idx) = 0.0; 511 for (d = 0; d < LANDAU_DIM; ++d) refSpaceDer[d] = 0.0; 512 for (b = 0; b < Nb; ++b) { 513 const PetscInt cidx = b; 514 d_fdf_k(b_id,0,idx) += Bq[cidx]*PetscRealPart(coef[f*Nb+cidx]); 515 for (d = 0; d < dim; ++d) refSpaceDer[d] += Dq[cidx*dim+d]*PetscRealPart(coef[f*Nb+cidx]); 516 } 517 //printf("\t%d.%d.%d.%d.%d) d_fdf_k(b_id,0,%d) = %g\n",b_id,grid,loc_elem,myQi,f,idx,d_fdf_k(b_id,0,idx)); 518 for (d = 0; d < dim; ++d) { 519 for (e = 0, d_fdf_k(b_id,d+1,idx) = 0.0; e < dim; ++e) { 520 d_fdf_k(b_id,d+1,idx) += invJ[e*dim+d]*refSpaceDer[e]; 521 } 522 //printf("\t\t\t%d.%d.%d.%d) d_fdf_k(b_id,%d,%d) = %g\n",b_id,grid,loc_elem,myQi,d+1,idx,d_fdf_k(b_id,d+1,idx)); 523 } 524 }); // f 525 }); // myQi 526 } 527 }); // global elems - fdf 528 // #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) 529 // ierr = PetscLogGpuFlops(nip_loc*(PetscLogDouble)(2*Nb*(1+dim)));CHKERRQ(ierr); 530 // #else 531 // ierr = PetscLogFlops(nip_loc*(PetscLogDouble)(2*Nb*(1+dim)));CHKERRQ(ierr); 532 // #endif 533 Kokkos::fence(); 534 ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr); // is this a fence? 535 ierr = PetscLogEventEnd(events[8],0,0,0,0);CHKERRQ(ierr); 536 if (d_vertex_f_k) delete d_vertex_f_k; 537 // Jacobian 538 ierr = PetscLogEventBegin(events[4],0,0,0,0);CHKERRQ(ierr); 539 ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr); 540 // #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) 541 // ierr = PetscLogGpuFlops(nip_loc*(PetscLogDouble)((nip_loc*(11*loc_Nf+ 4*dim*dim) + 6*loc_Nf*dim*dim*dim + 10*loc_Nf*dim*dim + 4*loc_Nf*dim + Nb*loc_Nf*Nb*Nq*dim*dim*5)));CHKERRQ(ierr); 542 // #else 543 // ierr = PetscLogFlops(nip_loc*(PetscLogDouble)((nip_loc*(11*loc_Nf+ 4*dim*dim) + 6*loc_Nf*dim*dim*dim + 10*loc_Nf*dim*dim + 4*loc_Nf*dim + Nb*loc_Nf*Nb*Nq*dim*dim*5)));CHKERRQ(ierr); 544 // #endif 545 const int scr_bytes = 2*(g2_scr_t::shmem_size(dim,Nf_max,Nq) + g3_scr_t::shmem_size(dim,dim,Nf_max,Nq))+fieldMats_scr_t::shmem_size(Nb,Nb)+idx_scr_t::shmem_size(Nb,nfaces)+scale_scr_t::shmem_size(Nb,nfaces); 546 ierr = PetscInfo6(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);CHKERRQ(ierr); 547 Kokkos::parallel_for("Jacobian", 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) { 548 const PetscInt b_Nelem = d_elem_offset[num_grids], b_elem_idx = team.league_rank()%b_Nelem, b_id = team.league_rank()/b_Nelem; 549 // find my grid 550 PetscInt grid = 0; 551 while (b_elem_idx >= d_elem_offset[grid+1]) grid++; 552 { 553 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; 554 const PetscInt moffset = LAND_MOFFSET(b_id,grid,batch_sz,num_grids,d_mat_offset); 555 const PetscInt f_off = d_species_offset[grid]; 556 g2_scr_t g2(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,loc_Nf,Nq); 557 g3_scr_t g3(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,dim,loc_Nf,Nq); 558 g2_scr_t gg2(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,loc_Nf,Nq); 559 g3_scr_t gg3(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,dim,loc_Nf,Nq); 560 // get g2[] & g3[] 561 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nq), [=] (int myQi) { 562 using Kokkos::parallel_reduce; 563 const PetscInt jpidx_glb = d_ip_offset[grid] + loc_elem * Nq + myQi; 564 const PetscReal *invJ = &d_invJ_k( jpidx_glb*dim*dim); 565 const PetscReal vj[3] = {d_x[jpidx_glb], d_y[jpidx_glb], d_z ? d_z[jpidx_glb] : 0}, wj = d_w[jpidx_glb]; 566 landau_inner_red::TensorValueType gg_temp; // reduce on part of gg2 and g33 for IP jpidx_g 567 Kokkos::parallel_reduce(Kokkos::ThreadVectorRange (team, (int)d_ip_offset[num_grids]), [=] (const int& ipidx, landau_inner_red::TensorValueType & ggg) { 568 const PetscReal wi = d_w[ipidx], x = d_x[ipidx], y = d_y[ipidx]; 569 PetscReal temp1[3] = {0, 0, 0}, temp2 = 0; 570 PetscInt fieldA,d2,d3,f_off_r,grid_r,ipidx_g,nip_loc_r,loc_Nf_r; 571 #if LANDAU_DIM==2 572 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.; 573 LandauTensor2D(vj, x, y, Ud, Uk, mask); 574 #else 575 PetscReal U[3][3], z = d_z[jpidx_glb], 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.; 576 LandauTensor3D(vj, x, y, z, U, mask); 577 #endif 578 grid_r = 0; 579 while (ipidx >= d_ip_offset[grid_r+1]) grid_r++; // yuck search for grid 580 f_off_r = d_species_offset[grid_r]; 581 ipidx_g = ipidx - d_ip_offset[grid_r]; 582 nip_loc_r = d_numCells[grid_r]*Nq; 583 loc_Nf_r = d_species_offset[grid_r+1] - d_species_offset[grid_r]; 584 for (fieldA = 0; fieldA < loc_Nf_r; ++fieldA) { 585 const PetscInt idx = d_ipf_offset[grid_r] + fieldA*nip_loc_r + ipidx_g; 586 temp1[0] += d_fdf_k(b_id,1,idx)*d_beta[fieldA+f_off_r]*d_invMass[fieldA+f_off_r]; 587 temp1[1] += d_fdf_k(b_id,2,idx)*d_beta[fieldA+f_off_r]*d_invMass[fieldA+f_off_r]; 588 #if LANDAU_DIM==3 589 temp1[2] += d_fdf_k(b_id,3,idx)*d_beta[fieldA+f_off_r]*d_invMass[fieldA+f_off_r]; 590 #endif 591 temp2 += d_fdf_k(b_id,0,idx)*d_beta[fieldA+f_off_r]; 592 } 593 temp1[0] *= wi; 594 temp1[1] *= wi; 595 #if LANDAU_DIM==3 596 temp1[2] *= wi; 597 #endif 598 temp2 *= wi; 599 #if LANDAU_DIM==2 600 for (d2 = 0; d2 < 2; d2++) { 601 for (d3 = 0; d3 < 2; ++d3) { 602 /* K = U * grad(f): g2=e: i,A */ 603 ggg.gg2[d2] += Uk[d2][d3]*temp1[d3]; 604 /* D = -U * (I \kron (fx)): g3=f: i,j,A */ 605 ggg.gg3[d2][d3] += Ud[d2][d3]*temp2; 606 } 607 } 608 #else 609 for (d2 = 0; d2 < 3; ++d2) { 610 for (d3 = 0; d3 < 3; ++d3) { 611 /* K = U * grad(f): g2 = e: i,A */ 612 ggg.gg2[d2] += U[d2][d3]*temp1[d3]; 613 /* D = -U * (I \kron (fx)): g3 = f: i,j,A */ 614 ggg.gg3[d2][d3] += U[d2][d3]*temp2; 615 } 616 } 617 #endif 618 }, Kokkos::Sum<landau_inner_red::TensorValueType>(gg_temp)); 619 // add alpha and put in gg2/3 620 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)loc_Nf), [&] (const int& fieldA) { 621 PetscInt d2,d3; 622 for (d2 = 0; d2 < dim; d2++) { 623 gg2(d2,fieldA,myQi) = gg_temp.gg2[d2]*d_alpha[fieldA+f_off]; 624 //printf("\t\t\t%d.%d.%d.%d) gg2(%d,%d)) = %g\n",b_id,grid,loc_elem,myQi,d2, fieldA, gg2(d2,fieldA,myQi)); 625 for (d3 = 0; d3 < dim; d3++) { 626 gg3(d2,d3,fieldA,myQi) = -gg_temp.gg3[d2][d3]*d_alpha[fieldA+f_off]*d_invMass[fieldA+f_off]; 627 } 628 } 629 }); 630 /* add electric field term once per IP */ 631 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)loc_Nf), [&] (const int& fieldA) { 632 gg2(dim-1,fieldA,myQi) += d_Eq_m[fieldA+f_off]; 633 }); 634 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)loc_Nf), [=] (const int& fieldA) { 635 int d,d2,d3,dp; 636 /* Jacobian transform - g2, g3 - per thread (2D) */ 637 for (d = 0; d < dim; ++d) { 638 g2(d,fieldA,myQi) = 0; 639 for (d2 = 0; d2 < dim; ++d2) { 640 g2(d,fieldA,myQi) += invJ[d*dim+d2]*gg2(d2,fieldA,myQi); 641 g3(d,d2,fieldA,myQi) = 0; 642 for (d3 = 0; d3 < dim; ++d3) { 643 for (dp = 0; dp < dim; ++dp) { 644 g3(d,d2,fieldA,myQi) += invJ[d*dim + d3]*gg3(d3,dp,fieldA,myQi)*invJ[d2*dim + dp]; 645 } 646 } 647 g3(d,d2,fieldA,myQi) *= wj; 648 } 649 g2(d,fieldA,myQi) *= wj; 650 } 651 }); 652 }); // Nq 653 team.team_barrier(); 654 { /* assemble */ 655 fieldMats_scr_t s_fieldMats(team.team_scratch(KOKKOS_SHARED_LEVEL),Nb,Nb); // Only used for GPU assembly (ie, not first pass) 656 idx_scr_t s_idx(team.team_scratch(KOKKOS_SHARED_LEVEL),Nb,nfaces); 657 scale_scr_t s_scale(team.team_scratch(KOKKOS_SHARED_LEVEL),Nb,nfaces); 658 for (PetscInt fieldA = 0; fieldA < loc_Nf; fieldA++) { 659 /* assemble */ 660 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nb), [=] (int f) { 661 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,Nb), [=] (int g) { 662 PetscScalar t = 0; 663 for (int qj = 0 ; qj < Nq ; qj++) { // look at others integration points 664 const PetscReal *BJq = &d_BB[qj*Nb], *DIq = &d_DD[qj*Nb*dim]; 665 for (int d = 0; d < dim; ++d) { 666 t += DIq[f*dim+d]*g2(d,fieldA,qj)*BJq[g]; 667 for (int d2 = 0; d2 < dim; ++d2) { 668 t += DIq[f*dim + d]*g3(d,d2,fieldA,qj)*DIq[g*dim + d2]; 669 //printf("\t\t%d.%d.%d) t = %g : %g %g %g\n",b_id,grid,loc_elem,t,DIq[f*dim + d],g3(d,d2,fieldA,qj),DIq[g*dim + d2]); 670 } 671 } 672 } 673 if (elem_mat_num_cells_max_grid) { // CPU assembly 674 const PetscInt fOff = (fieldA*Nb + f)*totDim + fieldA*Nb + g; 675 d_elem_mats(b_id,grid,loc_elem,fOff) = t; 676 } else s_fieldMats(f,g) = t; 677 }); 678 }); 679 if (!elem_mat_num_cells_max_grid) { // GPU assembly 680 landau_mat_assemble (d_mat, team, s_fieldMats, s_idx, s_scale, Nb, Nq, nfaces, moffset, loc_elem, fieldA, maps[grid]); 681 } 682 } 683 } 684 } 685 }); 686 ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr); 687 ierr = PetscLogEventEnd(events[4],0,0,0,0);CHKERRQ(ierr); 688 } else { // mass 689 ierr = PetscLogEventBegin(events[4],0,0,0,0);CHKERRQ(ierr); 690 ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr); 691 // #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) 692 // ierr = PetscLogGpuFlops(nip_loc*(PetscLogDouble)(Nb*loc_Nf*Nb*Nq*4));CHKERRQ(ierr); 693 // if (ctx->deviceType == LANDAU_CPU) PetscInfo(plex[grid], "Warning: Landau selected CPU but no support for Kokkos using CPU\n"); 694 // #else 695 // ierr = PetscLogFlops(nip_loc*(PetscLogDouble)(Nb*loc_Nf*Nb*Nq*4));CHKERRQ(ierr); 696 // #endif 697 int scr_bytes = fieldMats_scr_t::shmem_size(Nq,Nq) + idx_scr_t::shmem_size(Nb,nfaces) + scale_scr_t::shmem_size(Nb,nfaces); 698 ierr = PetscInfo6(plex[0], "Mass shared memory size: %d bytes in level %d conc=%D team size=%D #face=%D Nb=%D\n",scr_bytes,KOKKOS_SHARED_LEVEL,conc,team_size,nfaces,Nb);CHKERRQ(ierr); 699 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) { 700 const PetscInt b_Nelem = d_elem_offset[num_grids], b_elem_idx = team.league_rank()%b_Nelem, b_id = team.league_rank()/b_Nelem; 701 fieldMats_scr_t s_fieldMats(team.team_scratch(KOKKOS_SHARED_LEVEL),Nb,Nb); 702 idx_scr_t s_idx(team.team_scratch(KOKKOS_SHARED_LEVEL),Nb,nfaces); 703 scale_scr_t s_scale(team.team_scratch(KOKKOS_SHARED_LEVEL),Nb,nfaces); 704 // find my grid 705 PetscInt grid = 0; 706 while (b_elem_idx >= d_elem_offset[grid+1]) grid++; 707 { 708 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; 709 const PetscInt moffset = LAND_MOFFSET(b_id,grid,batch_sz,num_grids,d_mat_offset); 710 for (int fieldA = 0; fieldA < loc_Nf; fieldA++) { 711 /* assemble */ 712 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nb), [=] (int f) { 713 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,Nb), [=] (int g) { 714 PetscScalar t = 0; 715 for (int qj = 0 ; qj < Nq ; qj++) { // look at others integration points 716 const PetscReal *BJq = &d_BB[qj*Nb]; 717 const PetscInt jpidx_glb = jpidx_0 + qj; 718 if (dim==2) { 719 t += BJq[f] * d_w_k(jpidx_glb) * shift * BJq[g] * 2. * PETSC_PI; 720 } else { 721 t += BJq[f] * d_w_k(jpidx_glb) * shift * BJq[g]; 722 } 723 } 724 if (elem_mat_num_cells_max_grid) { 725 const PetscInt fOff = (fieldA*Nb + f)*totDim + fieldA*Nb + g; 726 d_elem_mats(b_id,grid,loc_elem,fOff) = t; 727 } else s_fieldMats(f,g) = t; 728 }); 729 }); 730 if (!elem_mat_num_cells_max_grid) { // device assembly 731 landau_mat_assemble (d_mat, team, s_fieldMats, s_idx, s_scale, Nb, Nq, nfaces, moffset, loc_elem, fieldA, maps[grid]); 732 } // else not using GPU assembly 733 } 734 } 735 }); 736 ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr); 737 ierr = PetscLogEventEnd(events[4],0,0,0,0);CHKERRQ(ierr); 738 } 739 Kokkos::fence(); 740 if (elem_mat_num_cells_max_grid) { // CPU assembly 741 Kokkos::View<PetscScalar****, Kokkos::LayoutRight>::HostMirror h_elem_mats = Kokkos::create_mirror_view(d_elem_mats); 742 Kokkos::deep_copy (h_elem_mats, d_elem_mats); 743 if (container) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "?????"); 744 for (PetscInt b_id = 0 ; b_id < batch_sz ; b_id++) { // OpenMP (once) 745 for (PetscInt grid=0 ; grid<num_grids ; grid++) { 746 PetscSection section, globalSection; 747 PetscInt cStart,cEnd; 748 Mat B = subJ[ LAND_PACK_IDX(b_id,grid) ]; 749 PetscInt moffset = LAND_MOFFSET(b_id,grid,batch_sz,num_grids,a_mat_offset), nloc, nzl, colbuf[1024], row; 750 const PetscInt *cols; 751 const PetscScalar *vals; 752 ierr = PetscLogEventBegin(events[5],0,0,0,0);CHKERRQ(ierr); 753 ierr = DMPlexGetHeightStratum(plex[grid],0,&cStart,&cEnd);CHKERRQ(ierr); 754 ierr = DMGetLocalSection(plex[grid], §ion);CHKERRQ(ierr); 755 ierr = DMGetGlobalSection(plex[grid], &globalSection);CHKERRQ(ierr); 756 ierr = PetscLogEventEnd(events[5],0,0,0,0);CHKERRQ(ierr); 757 ierr = PetscLogEventBegin(events[6],0,0,0,0);CHKERRQ(ierr); 758 for (PetscInt ej = cStart ; ej < cEnd; ++ej) { 759 const PetscScalar *elMat = &h_elem_mats(b_id,grid,ej-cStart,0); 760 ierr = DMPlexMatSetClosure(plex[grid], section, globalSection, B, ej, elMat, ADD_VALUES);CHKERRQ(ierr); 761 if (grid==0 && ej==-1) { 762 const PetscInt loc_Nf = a_species_offset[grid+1]-a_species_offset[grid], totDim = loc_Nf*Nq; 763 int d,f; 764 PetscPrintf(PETSC_COMM_SELF,"Kokkos Element matrix %d/%d\n",1,(int)a_numCells[grid]); 765 for (d = 0; d < totDim; ++d){ 766 for (f = 0; f < totDim; ++f) PetscPrintf(PETSC_COMM_SELF," %12.5e", PetscRealPart(elMat[d*totDim + f])); 767 PetscPrintf(PETSC_COMM_SELF,"\n"); 768 } 769 exit(14); 770 } 771 } 772 //printf("\t\t\t%d.%d: a_mat_offset = %d %d. moffset=%d B=%p\n",b_id,grid,a_mat_offset[0],a_mat_offset[1],moffset, B); 773 // move nest matrix to global JacP 774 ierr = MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 775 ierr = MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 776 ierr = MatGetSize(B, &nloc, NULL);CHKERRQ(ierr); 777 for (int i=0 ; i<nloc ; i++) { 778 ierr = MatGetRow(B,i,&nzl,&cols,&vals);CHKERRQ(ierr); 779 if (nzl>1024) SETERRQ1(PetscObjectComm((PetscObject) B), PETSC_ERR_PLIB, "Row too big: %D",nzl); 780 for (int j=0; j<nzl; j++) colbuf[j] = cols[j] + moffset; 781 row = i + moffset; 782 ierr = MatSetValues(JacP,1,&row,nzl,colbuf,vals,ADD_VALUES);CHKERRQ(ierr); 783 ierr = MatRestoreRow(B,i,&nzl,&cols,&vals);CHKERRQ(ierr); 784 } 785 ierr = MatDestroy(&subJ[ LAND_PACK_IDX(b_id,grid) ]);CHKERRQ(ierr); 786 ierr = PetscLogEventEnd(events[6],0,0,0,0);CHKERRQ(ierr); 787 } // grids 788 } 789 } 790 PetscFunctionReturn(0); 791 } 792 } // extern "C" 793