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) 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->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->gIdx, maps->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->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->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->num_elements; 87 h_maps.num_face = maps->num_face; 88 h_maps.num_reduced = maps->num_reduced; 89 h_maps.deviceType = maps->deviceType; 90 h_maps.Nf = Nf; 91 h_maps.Nq = Nq; 92 h_maps.c_maps = (pointInterpolationP4est (*)[LANDAU_MAX_Q_FACE]) d_points->data(); 93 maps->vp1 = (void*)d_points; 94 h_maps.gIdx = (LandauIdx (*)[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ]) d_gidx->data(); 95 maps->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->data = d_maps_k->data(); 101 maps->vp3 = (void*)d_maps_k; 102 } 103 PetscFunctionReturn(0); 104 } 105 PetscErrorCode LandauKokkosDestroyMatMaps(P4estVertexMaps *maps) 106 { 107 PetscFunctionBegin; 108 Kokkos::View<pointInterpolationP4est*[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight> *a = (Kokkos::View<pointInterpolationP4est*[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight>*)maps->vp1; 109 Kokkos::View<LandauIdx*[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ], Kokkos::LayoutRight>*b = (Kokkos::View<LandauIdx*[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ], Kokkos::LayoutRight>*)maps->vp2; 110 Kokkos::View<P4estVertexMaps*> *c = (Kokkos::View<P4estVertexMaps*>*)maps->vp3; 111 delete a; delete b; delete c; 112 PetscFunctionReturn(0); 113 } 114 115 PetscErrorCode LandauKokkosStaticDataSet(DM plex, const PetscInt Nq, PetscReal nu_alpha[], PetscReal nu_beta[], PetscReal a_invMass[], PetscReal a_invJ[], PetscReal a_mass_w[], 116 PetscReal a_x[], PetscReal a_y[], PetscReal a_z[], PetscReal a_w[], LandauGeomData *SData_d) 117 { 118 PetscReal *BB,*DD; 119 PetscErrorCode ierr; 120 PetscTabulation *Tf; 121 LandauCtx *ctx; 122 PetscInt *Nbf,dim,Nf,Nb,nip,cStart,cEnd; 123 PetscDS prob; 124 125 PetscFunctionBegin; 126 ierr = DMGetApplicationContext(plex, &ctx);CHKERRQ(ierr); 127 if (!ctx) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "no context"); 128 ierr = DMGetDimension(plex, &dim);CHKERRQ(ierr); 129 ierr = DMPlexGetHeightStratum(plex,0,&cStart,&cEnd);CHKERRQ(ierr); 130 nip = (cEnd - cStart)*Nq; 131 ierr = DMGetDS(plex, &prob);CHKERRQ(ierr); 132 ierr = PetscDSGetNumFields(prob, &Nf);CHKERRQ(ierr); 133 ierr = PetscDSGetDimensions(prob, &Nbf);CHKERRQ(ierr); Nb = Nbf[0]; 134 if (Nq != Nb) SETERRQ2(PETSC_COMM_SELF, PETSC_ERR_PLIB, "Nq != Nb. %D %D",Nq,Nb); 135 if (LANDAU_DIM != dim) SETERRQ2(PETSC_COMM_WORLD, PETSC_ERR_PLIB, "dim %D != LANDAU_DIM %d",dim,LANDAU_DIM); 136 ierr = PetscDSGetTabulation(prob, &Tf);CHKERRQ(ierr); 137 BB = Tf[0]->T[0]; DD = Tf[0]->T[1]; 138 ierr = PetscKokkosInitializeCheck();CHKERRQ(ierr); 139 { 140 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_alpha (nu_alpha, Nf); 141 auto alpha = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("alpha", Nf); 142 SData_d->alpha = static_cast<void*>(alpha); 143 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_beta (nu_beta, Nf); 144 auto beta = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("beta", Nf); 145 SData_d->beta = static_cast<void*>(beta); 146 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_invMass (a_invMass,Nf); 147 auto invMass = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("invMass", Nf); 148 SData_d->invMass = static_cast<void*>(invMass); 149 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_BB (BB,Nq*Nb); 150 auto B = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("B", Nq*Nb); 151 SData_d->B = static_cast<void*>(B); 152 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_DD (DD,Nq*Nb*dim); 153 auto D = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("D", Nq*Nb*dim); 154 SData_d->D = static_cast<void*>(D); 155 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_mass_w (a_mass_w, nip); 156 auto mass_w = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("mass_w", nip); 157 SData_d->mass_w = static_cast<void*>(mass_w); 158 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_invJ (a_invJ, nip*dim*dim); 159 auto invJ = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("invJ", nip*dim*dim); 160 SData_d->invJ = static_cast<void*>(invJ); 161 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_x (a_x, nip); 162 auto x = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("x", nip); 163 SData_d->x = static_cast<void*>(x); 164 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_y (a_y, nip); 165 auto y = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("y", nip); 166 SData_d->y = static_cast<void*>(y); 167 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_w (a_w, nip); 168 auto w = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("w", nip); 169 SData_d->w = static_cast<void*>(w); 170 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_z (a_z , dim==3 ? nip : 0); 171 auto z = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("z", dim==3 ? nip : 0); 172 SData_d->z = static_cast<void*>(z); 173 174 Kokkos::deep_copy (*mass_w, h_mass_w); 175 Kokkos::deep_copy (*alpha, h_alpha); 176 Kokkos::deep_copy (*beta, h_beta); 177 Kokkos::deep_copy (*invMass, h_invMass); 178 Kokkos::deep_copy (*B, h_BB); 179 Kokkos::deep_copy (*D, h_DD); 180 Kokkos::deep_copy (*invJ, h_invJ); 181 Kokkos::deep_copy (*x, h_x); 182 Kokkos::deep_copy (*y, h_y); 183 Kokkos::deep_copy (*z, h_z); 184 Kokkos::deep_copy (*w, h_w); 185 186 auto Eq_m = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("Eq_m",Nf); 187 SData_d->Eq_m = static_cast<void*>(Eq_m); 188 auto IPf = new Kokkos::View<PetscScalar*, Kokkos::LayoutLeft> ("IPf",nip*Nf); // Nq==Nb 189 SData_d->IPf = static_cast<void*>(IPf); 190 } 191 PetscFunctionReturn(0); 192 } 193 194 PetscErrorCode LandauKokkosStaticDataClear(LandauGeomData *SData_d) 195 { 196 PetscFunctionBegin; 197 if (SData_d->alpha) { 198 auto alpha = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->alpha); 199 delete alpha; 200 auto beta = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->beta); 201 delete beta; 202 auto invMass = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->invMass); 203 delete invMass; 204 auto B = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->B); 205 delete B; 206 auto D = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->D); 207 delete D; 208 auto mass_w = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->mass_w); 209 delete mass_w; 210 auto invJ = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->invJ); 211 delete invJ; 212 auto x = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->x); 213 delete x; 214 auto y = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->y); 215 delete y; 216 auto z = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->z); 217 delete z; 218 auto w = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->w); 219 delete w; 220 auto Eq_m = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->Eq_m); 221 delete Eq_m; 222 if (SData_d->IPf) { 223 auto IPf = static_cast<Kokkos::View<PetscScalar*, Kokkos::LayoutLeft>*>(SData_d->IPf); 224 delete IPf; 225 } 226 } 227 PetscFunctionReturn(0); 228 } 229 230 PetscErrorCode LandauKokkosJacobian(DM plex, const PetscInt Nq, PetscReal a_Eq_m[], PetscScalar a_IPf[], const PetscInt N, const PetscScalar a_xarray[], LandauGeomData *SData_d, 231 const PetscInt num_sub_blocks, PetscReal shift, const PetscLogEvent events[], Mat JacP) 232 { 233 using scr_mem_t = Kokkos::DefaultExecutionSpace::scratch_memory_space; 234 using fieldMats_scr_t = Kokkos::View<PetscScalar**, Kokkos::LayoutRight, scr_mem_t>; 235 using idx_scr_t = Kokkos::View<PetscInt**, Kokkos::LayoutRight, scr_mem_t>; 236 using scale_scr_t = Kokkos::View<PetscReal**, Kokkos::LayoutRight, scr_mem_t>; 237 using g2_scr_t = Kokkos::View<PetscReal***, Kokkos::LayoutRight, scr_mem_t>; 238 using g3_scr_t = Kokkos::View<PetscReal****, Kokkos::LayoutRight, scr_mem_t>; 239 PetscErrorCode ierr; 240 PetscInt *Nbf,Nb,cStart,cEnd,Nf,dim,numCells,totDim,global_elem_mat_sz,nip,nfaces=0; 241 PetscDS prob; 242 LandauCtx *ctx; 243 PetscReal *d_Eq_m=NULL; 244 PetscScalar *d_IPf=NULL; 245 P4estVertexMaps *d_maps=NULL; 246 PetscSplitCSRDataStructure d_mat=NULL; 247 const int conc = Kokkos::DefaultExecutionSpace().concurrency(), openmp = !!(conc < 1000), team_size = (openmp==0) ? Nq : 1; 248 auto d_alpha_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->alpha); //static data 249 const PetscReal *d_alpha = d_alpha_k->data(); 250 auto d_beta_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->beta); 251 const PetscReal *d_beta = d_beta_k->data(); 252 auto d_invMass_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->invMass); 253 const PetscReal *d_invMass = d_invMass_k->data(); 254 auto d_B_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->B); 255 const PetscReal *d_BB = d_B_k->data(); 256 auto d_D_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->D); 257 const PetscReal *d_DD = d_D_k->data(); 258 auto d_mass_w_k = *static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->mass_w); 259 auto d_invJ_k = *static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->invJ); // use Kokkos vector in kernels 260 auto d_x_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->x); //static data 261 const PetscReal *d_x = d_x_k->data(); 262 auto d_y_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->y); //static data 263 const PetscReal *d_y = d_y_k->data(); 264 auto d_z_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->z); //static data 265 const PetscReal *d_z = d_z_k->data(); 266 auto d_w_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->w); //static data 267 const PetscReal *d_w = d_w_k->data(); 268 // dynamic data pointers 269 auto d_Eq_m_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->Eq_m); //static data 270 auto d_IPf_k = static_cast<Kokkos::View<PetscScalar*, Kokkos::LayoutLeft>*>(SData_d->IPf); //static data 271 272 PetscFunctionBegin; 273 ierr = PetscLogEventBegin(events[3],0,0,0,0);CHKERRQ(ierr); 274 ierr = DMGetApplicationContext(plex, &ctx);CHKERRQ(ierr); 275 if (!ctx) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "no context"); 276 ierr = DMGetDimension(plex, &dim);CHKERRQ(ierr); 277 ierr = DMPlexGetHeightStratum(plex,0,&cStart,&cEnd);CHKERRQ(ierr); 278 numCells = cEnd - cStart; 279 nip = numCells*Nq; 280 ierr = DMGetDS(plex, &prob);CHKERRQ(ierr); 281 ierr = PetscDSGetNumFields(prob, &Nf);CHKERRQ(ierr); 282 ierr = PetscDSGetDimensions(prob, &Nbf);CHKERRQ(ierr); Nb = Nbf[0]; 283 if (Nq != Nb) SETERRQ2(PETSC_COMM_SELF, PETSC_ERR_PLIB, "Nq != Nb. %D %D",Nq,Nb); 284 if (LANDAU_DIM != dim) SETERRQ2(PETSC_COMM_WORLD, PETSC_ERR_PLIB, "dim %D != LANDAU_DIM %d",dim,LANDAU_DIM); 285 ierr = PetscDSGetTotalDimension(prob, &totDim);CHKERRQ(ierr); 286 if (ctx->gpu_assembly) { 287 PetscContainer container; 288 ierr = PetscObjectQuery((PetscObject) JacP, "assembly_maps", (PetscObject *) &container);CHKERRQ(ierr); 289 if (container) { // not here first call 290 #if defined(PETSC_HAVE_KOKKOS_KERNELS) 291 P4estVertexMaps *h_maps=NULL; 292 ierr = PetscContainerGetPointer(container, (void **) &h_maps);CHKERRQ(ierr); 293 if (h_maps->data) { 294 d_maps = h_maps->data; 295 nfaces = h_maps->num_face; 296 } else { 297 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "GPU assembly but no metadata in container"); 298 } 299 // this does the setup the first time called 300 ierr = MatKokkosGetDeviceMatWrite(JacP,&d_mat);CHKERRQ(ierr); 301 global_elem_mat_sz = 0; 302 #else 303 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "GPU assembly w/o kokkos kernels -- should not be here"); 304 #endif 305 } else { // kernel output - first call assembled on device 306 global_elem_mat_sz = numCells; 307 } 308 } else { 309 global_elem_mat_sz = numCells; // no device assembly 310 } 311 ierr = PetscKokkosInitializeCheck();CHKERRQ(ierr); 312 ierr = PetscLogEventEnd(events[3],0,0,0,0);CHKERRQ(ierr); 313 314 Kokkos::View<PetscScalar**, Kokkos::LayoutRight> d_elem_mats("element matrices", global_elem_mat_sz, totDim*totDim); 315 if (a_IPf || a_xarray) { // Jacobian 316 const int num_face = nfaces; 317 const int scr_bytes = 2*(g2_scr_t::shmem_size(dim,Nf,Nq) + g3_scr_t::shmem_size(dim,dim,Nf,Nq)) + fieldMats_scr_t::shmem_size(Nb,Nb) + idx_scr_t::shmem_size(Nb,num_face) + scale_scr_t::shmem_size(Nb,num_face); 318 #if defined(LANDAU_LAYOUT_LEFT) 319 Kokkos::View<PetscReal***, Kokkos::LayoutLeft > d_fdf_k("df", dim+1, Nf, (a_IPf || a_xarray) ? nip : 0); 320 #else 321 Kokkos::View<PetscReal***, Kokkos::LayoutRight > d_fdf_k("df", dim+1, Nf, (a_IPf || a_xarray) ? nip : 0); 322 #endif 323 const Kokkos::View<PetscScalar*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >h_IPf_k (a_IPf, (a_IPf || a_xarray) ? nip*Nf : 0); 324 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_Eq_m_k (a_Eq_m, (a_IPf || a_xarray) ? Nf : 0); 325 Kokkos::deep_copy (*d_Eq_m_k, h_Eq_m_k); 326 d_Eq_m = d_Eq_m_k->data(); 327 // copy dynamic data to device 328 ierr = PetscLogEventBegin(events[1],0,0,0,0);CHKERRQ(ierr); 329 if (a_IPf) { 330 Kokkos::deep_copy (*d_IPf_k, h_IPf_k); 331 d_IPf = d_IPf_k->data(); 332 } else { 333 d_IPf = (PetscScalar*)a_xarray; 334 } 335 ierr = PetscLogEventEnd(events[1],0,0,0,0);CHKERRQ(ierr); 336 #define KOKKOS_SHARED_LEVEL 1 337 // get f and df 338 ierr = PetscLogEventBegin(events[8],0,0,0,0);CHKERRQ(ierr); 339 ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr); 340 Kokkos::parallel_for("f, df", Kokkos::TeamPolicy<>(numCells, team_size, /* Kokkos::AUTO */ 16), KOKKOS_LAMBDA (const team_member team) { 341 const PetscInt elem = team.league_rank(); 342 const PetscScalar *coef; 343 PetscScalar coef_buff[LANDAU_MAX_SPECIES*LANDAU_MAX_NQ]; 344 // un pack IPData 345 if (!d_maps) { 346 coef = &d_IPf[elem*Nb*Nf]; 347 } else { 348 coef = coef_buff; 349 for (int f = 0; f < Nf; ++f) { 350 LandauIdx *const Idxs = &d_maps->gIdx[elem][f][0]; 351 for (int b = 0; b < Nb; ++b) { 352 PetscInt idx = Idxs[b]; 353 if (idx >= 0) { 354 coef_buff[f*Nb+b] = d_IPf[idx]; 355 } else { 356 idx = -idx - 1; 357 coef_buff[f*Nb+b] = 0; 358 for (int q = 0; q < d_maps->num_face; q++) { 359 PetscInt id = d_maps->c_maps[idx][q].gid; 360 PetscScalar scale = d_maps->c_maps[idx][q].scale; 361 coef_buff[f*Nb+b] += scale*d_IPf[id]; 362 } 363 } 364 } 365 } 366 } 367 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nq), [=] (int myQi) { 368 const PetscInt ipidx = myQi + elem * Nq; 369 const PetscReal *const invJj = &d_invJ_k(ipidx*dim*dim); 370 const PetscReal *Bq = &d_BB[myQi*Nb], *Dq = &d_DD[myQi*Nb*dim]; 371 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,(int)Nf), [=] (int f) { 372 PetscInt b, e, d; 373 PetscReal refSpaceDer[LANDAU_DIM]; 374 d_fdf_k(0,f,ipidx) = 0.0; 375 for (d = 0; d < LANDAU_DIM; ++d) refSpaceDer[d] = 0.0; 376 for (b = 0; b < Nb; ++b) { 377 const PetscInt cidx = b; 378 d_fdf_k(0,f,ipidx) += Bq[cidx]*PetscRealPart(coef[f*Nb+cidx]); 379 for (d = 0; d < dim; ++d) refSpaceDer[d] += Dq[cidx*dim+d]*PetscRealPart(coef[f*Nb+cidx]); 380 } 381 for (d = 0; d < dim; ++d) { 382 for (e = 0, d_fdf_k(d+1,f,ipidx) = 0.0; e < dim; ++e) { 383 d_fdf_k(d+1,f,ipidx) += invJj[e*dim+d]*refSpaceDer[e]; 384 } 385 } 386 }); // Nf 387 }); // Nq 388 }); // elems 389 #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) 390 ierr = PetscLogGpuFlops(nip*(PetscLogDouble)(2*Nb*(1+dim)));CHKERRQ(ierr); 391 #else 392 ierr = PetscLogFlops(nip*(PetscLogDouble)(2*Nb*(1+dim)));CHKERRQ(ierr); 393 #endif 394 ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr); 395 ierr = PetscLogEventEnd(events[8],0,0,0,0);CHKERRQ(ierr); 396 // Jacobian 397 ierr = PetscLogEventBegin(events[4],0,0,0,0);CHKERRQ(ierr); 398 ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr); 399 #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) 400 ierr = PetscLogGpuFlops(nip*(PetscLogDouble)((nip*(11*Nf+ 4*dim*dim) + 6*Nf*dim*dim*dim + 10*Nf*dim*dim + 4*Nf*dim + Nb*Nf*Nb*Nq*dim*dim*5)));CHKERRQ(ierr); 401 if (ctx->deviceType == LANDAU_CPU) PetscInfo(plex, "Warning: Landau selected CPU but no support for Kokkos using CPU\n"); 402 #else 403 ierr = PetscLogFlops(nip*(PetscLogDouble)((nip*(11*Nf+ 4*dim*dim) + 6*Nf*dim*dim*dim + 10*Nf*dim*dim + 4*Nf*dim + Nb*Nf*Nb*Nq*dim*dim*5)));CHKERRQ(ierr); 404 #endif 405 ierr = PetscInfo5(plex, "Jacobian shared memory size: %d bytes in level %d conc=%D team size=%D #face=%D\n",scr_bytes,KOKKOS_SHARED_LEVEL,conc,team_size,nfaces);CHKERRQ(ierr); 406 Kokkos::parallel_for("Landau Jacobian", Kokkos::TeamPolicy<>(numCells, team_size, /*Kokkos::AUTO*/ 16).set_scratch_size(KOKKOS_SHARED_LEVEL, Kokkos::PerTeam(scr_bytes)), KOKKOS_LAMBDA (const team_member team) { 407 const PetscInt elem = team.league_rank(); 408 g2_scr_t g2(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,Nf,Nq); 409 g3_scr_t g3(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,dim,Nf,Nq); 410 g2_scr_t gg2(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,Nf,Nq); 411 g3_scr_t gg3(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,dim,Nf,Nq); 412 // get g2[] & g3[] 413 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nq), [=] (int myQi) { 414 using Kokkos::parallel_reduce; 415 const PetscInt jpidx = myQi + elem * Nq; 416 const PetscReal* const invJj = &d_invJ_k(jpidx*dim*dim); 417 const PetscReal vj[3] = {d_x[jpidx], d_y[jpidx], d_z ? d_z[jpidx] : 0}, wj = d_w[jpidx]; 418 landau_inner_red::TensorValueType gg_temp; // reduce on part of gg2 and g33 for IP jpidx 419 Kokkos::parallel_reduce(Kokkos::ThreadVectorRange (team, (int)nip), [=] (const int& ipidx, landau_inner_red::TensorValueType & ggg) { 420 const PetscReal wi = d_w[ipidx], x = d_x[ipidx], y = d_y[ipidx]; 421 PetscReal temp1[3] = {0, 0, 0}, temp2 = 0; 422 PetscInt fieldA,d2,d3; 423 #if LANDAU_DIM==2 424 PetscReal Ud[2][2], Uk[2][2]; 425 LandauTensor2D(vj, x, y, Ud, Uk, (ipidx==jpidx) ? 0. : 1.); 426 #else 427 PetscReal U[3][3], z = d_z[ipidx]; 428 LandauTensor3D(vj, x, y, z, U, (ipidx==jpidx) ? 0. : 1.); 429 #endif 430 for (fieldA = 0; fieldA < Nf; ++fieldA) { 431 temp1[0] += d_fdf_k(1,fieldA,ipidx)*d_beta[fieldA]*d_invMass[fieldA]; 432 temp1[1] += d_fdf_k(2,fieldA,ipidx)*d_beta[fieldA]*d_invMass[fieldA]; 433 #if LANDAU_DIM==3 434 temp1[2] += d_fdf_k(3,fieldA,ipidx)*d_beta[fieldA]*d_invMass[fieldA]; 435 #endif 436 temp2 += d_fdf_k(0,fieldA,ipidx)*d_beta[fieldA]; 437 } 438 temp1[0] *= wi; 439 temp1[1] *= wi; 440 #if LANDAU_DIM==3 441 temp1[2] *= wi; 442 #endif 443 temp2 *= wi; 444 #if LANDAU_DIM==2 445 for (d2 = 0; d2 < 2; d2++) { 446 for (d3 = 0; d3 < 2; ++d3) { 447 /* K = U * grad(f): g2=e: i,A */ 448 ggg.gg2[d2] += Uk[d2][d3]*temp1[d3]; 449 /* D = -U * (I \kron (fx)): g3=f: i,j,A */ 450 ggg.gg3[d2][d3] += Ud[d2][d3]*temp2; 451 } 452 } 453 #else 454 for (d2 = 0; d2 < 3; ++d2) { 455 for (d3 = 0; d3 < 3; ++d3) { 456 /* K = U * grad(f): g2 = e: i,A */ 457 ggg.gg2[d2] += U[d2][d3]*temp1[d3]; 458 /* D = -U * (I \kron (fx)): g3 = f: i,j,A */ 459 ggg.gg3[d2][d3] += U[d2][d3]*temp2; 460 } 461 } 462 #endif 463 }, Kokkos::Sum<landau_inner_red::TensorValueType>(gg_temp)); 464 // add alpha and put in gg2/3 465 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)Nf), [&] (const int& fieldA) { 466 PetscInt d2,d3; 467 for (d2 = 0; d2 < dim; d2++) { 468 gg2(d2,fieldA,myQi) = gg_temp.gg2[d2]*d_alpha[fieldA]; 469 for (d3 = 0; d3 < dim; d3++) { 470 gg3(d2,d3,fieldA,myQi) = -gg_temp.gg3[d2][d3]*d_alpha[fieldA]*d_invMass[fieldA]; 471 } 472 } 473 }); 474 /* add electric field term once per IP */ 475 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)Nf), [&] (const int& fieldA) { 476 gg2(dim-1,fieldA,myQi) += d_Eq_m[fieldA]; 477 }); 478 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)Nf), [=] (const int& fieldA) { 479 int d,d2,d3,dp; 480 /* Jacobian transform - g2, g3 - per thread (2D) */ 481 for (d = 0; d < dim; ++d) { 482 g2(d,fieldA,myQi) = 0; 483 for (d2 = 0; d2 < dim; ++d2) { 484 g2(d,fieldA,myQi) += invJj[d*dim+d2]*gg2(d2,fieldA,myQi); 485 g3(d,d2,fieldA,myQi) = 0; 486 for (d3 = 0; d3 < dim; ++d3) { 487 for (dp = 0; dp < dim; ++dp) { 488 g3(d,d2,fieldA,myQi) += invJj[d*dim + d3]*gg3(d3,dp,fieldA,myQi)*invJj[d2*dim + dp]; 489 } 490 } 491 g3(d,d2,fieldA,myQi) *= wj; 492 } 493 g2(d,fieldA,myQi) *= wj; 494 } 495 }); 496 }); // Nq 497 team.team_barrier(); 498 { /* assemble */ 499 fieldMats_scr_t s_fieldMats(team.team_scratch(KOKKOS_SHARED_LEVEL),Nq,Nq); // Only used for GPU assembly (ie, not first pass) 500 idx_scr_t s_idx(team.team_scratch(KOKKOS_SHARED_LEVEL),Nq,num_face); 501 scale_scr_t s_scale(team.team_scratch(KOKKOS_SHARED_LEVEL),Nq,num_face); 502 PetscInt fieldA; 503 for (fieldA = 0; fieldA < Nf; fieldA++) { 504 /* assemble */ 505 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nb), [=] (int f) { 506 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,Nb), [=] (int g) { 507 PetscScalar t = 0; 508 for (int qj = 0 ; qj < Nq ; qj++) { // look at others integration points 509 const PetscReal *BJq = &d_BB[qj*Nb], *DIq = &d_DD[qj*Nb*dim]; 510 for (int d = 0; d < dim; ++d) { 511 t += DIq[f*dim+d]*g2(d,fieldA,qj)*BJq[g]; 512 for (int d2 = 0; d2 < dim; ++d2) { 513 t += DIq[f*dim + d]*g3(d,d2,fieldA,qj)*DIq[g*dim + d2]; 514 } 515 } 516 } 517 if (global_elem_mat_sz) { 518 const PetscInt fOff = (fieldA*Nb + f)*totDim + fieldA*Nb + g; 519 d_elem_mats(elem,fOff) = t; 520 } else s_fieldMats(f,g) = t; 521 }); 522 }); 523 if (!global_elem_mat_sz) { 524 const LandauIdx *const Idxs = &d_maps->gIdx[elem][fieldA][0]; 525 team.team_barrier(); 526 Kokkos::parallel_for(Kokkos::TeamVectorRange(team,0,Nb), [=] (int f) { 527 PetscInt q,idx = Idxs[f]; 528 if (idx >= 0) { 529 s_idx(f,0) = idx; 530 s_scale(f,0) = 1.; 531 } else { 532 idx = -idx - 1; 533 for (q = 0; q < nfaces; q++) { 534 s_idx(f,q) = d_maps->c_maps[idx][q].gid; 535 s_scale(f,q) = d_maps->c_maps[idx][q].scale; 536 } 537 } 538 }); 539 team.team_barrier(); 540 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nb), [=] (int f) { 541 PetscInt nr,idx = Idxs[f]; 542 if (idx >= 0) { 543 nr = 1; 544 } else { 545 nr = nfaces; 546 } 547 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,Nb), [=] (int g) { 548 PetscScalar vals[LANDAU_MAX_Q_FACE*LANDAU_MAX_Q_FACE]; 549 PetscInt q,d,nc,idx = Idxs[g]; 550 if (idx >= 0) { 551 nc = 1; 552 } else { 553 nc = nfaces; 554 } 555 for (q = 0; q < nr; q++) { 556 for (d = 0; d < nc; d++) { 557 vals[q*nc + d] = s_scale(f,q)*s_scale(g,d)*s_fieldMats(f,g); 558 } 559 } 560 MatSetValuesDevice(d_mat,nr,&s_idx(f,0),nc,&s_idx(g,0),vals,ADD_VALUES); 561 }); 562 }); 563 } 564 } 565 } 566 }); 567 ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr); 568 ierr = PetscLogEventEnd(events[4],0,0,0,0);CHKERRQ(ierr); 569 } else { // mass 570 int scr_bytes = fieldMats_scr_t::shmem_size(Nq,Nq) + idx_scr_t::shmem_size(Nb,nfaces) + scale_scr_t::shmem_size(Nb,nfaces); 571 ierr = PetscLogEventBegin(events[4],0,0,0,0);CHKERRQ(ierr); 572 ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr); 573 #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) 574 ierr = PetscLogGpuFlops(nip*(PetscLogDouble)(Nb*Nf*Nb*Nq*4));CHKERRQ(ierr); 575 if (ctx->deviceType == LANDAU_CPU) PetscInfo(plex, "Warning: Landau selected CPU but no support for Kokkos using CPU\n"); 576 #else 577 ierr = PetscLogFlops(nip*(PetscLogDouble)(Nb*Nf*Nb*Nq*4));CHKERRQ(ierr); 578 #endif 579 ierr = PetscInfo5(plex, "Mass shared memory size: %d bytes in level %d conc=%D team size=%D #face=%D\n",scr_bytes,KOKKOS_SHARED_LEVEL,conc,team_size,nfaces);CHKERRQ(ierr); 580 Kokkos::parallel_for("Landau mass", Kokkos::TeamPolicy<>(numCells, team_size, /*Kokkos::AUTO*/ 16).set_scratch_size(KOKKOS_SHARED_LEVEL, Kokkos::PerTeam(scr_bytes)), KOKKOS_LAMBDA (const team_member team) { 581 const PetscInt elem = team.league_rank(); 582 fieldMats_scr_t s_fieldMats(team.team_scratch(KOKKOS_SHARED_LEVEL),Nb,Nb); 583 idx_scr_t s_idx(team.team_scratch(KOKKOS_SHARED_LEVEL),Nb,nfaces); 584 scale_scr_t s_scale(team.team_scratch(KOKKOS_SHARED_LEVEL),Nb,nfaces); 585 for (int fieldA = 0; fieldA < Nf; fieldA++) { 586 /* assemble */ 587 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nb), [=] (int f) { 588 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,Nb), [=] (int g) { 589 PetscScalar t = 0; 590 for (int qj = 0 ; qj < Nq ; qj++) { // look at others integration points 591 const PetscReal *BJq = &d_BB[qj*Nb]; 592 const PetscInt jpidx = qj + elem * Nq; 593 t += BJq[f] * d_mass_w_k(jpidx) * shift * BJq[g]; 594 } 595 if (global_elem_mat_sz) { 596 const PetscInt fOff = (fieldA*Nb + f)*totDim + fieldA*Nb + g; 597 d_elem_mats(elem,fOff) = t; 598 } else s_fieldMats(f,g) = t; 599 }); 600 }); 601 if (!global_elem_mat_sz) { 602 const LandauIdx *const Idxs = &d_maps->gIdx[elem][fieldA][0]; 603 team.team_barrier(); // needed ???? 604 Kokkos::parallel_for(Kokkos::TeamVectorRange(team,0,Nb), [=] (int f) { 605 PetscInt q,idx = Idxs[f]; 606 if (idx >= 0) { 607 s_idx(f,0) = idx; 608 s_scale(f,0) = 1.; 609 } else { 610 idx = -idx - 1; 611 for (q = 0; q < nfaces; q++) { 612 s_idx(f,q) = d_maps->c_maps[idx][q].gid; 613 s_scale(f,q) = d_maps->c_maps[idx][q].scale; 614 } 615 } 616 }); 617 team.team_barrier(); 618 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nb), [=] (int f) { 619 PetscInt nr,idx = Idxs[f]; 620 if (idx >= 0) { 621 nr = 1; 622 } else { 623 nr = nfaces; 624 } 625 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,Nb), [=] (int g) { 626 PetscScalar vals[LANDAU_MAX_Q_FACE*LANDAU_MAX_Q_FACE]; 627 PetscInt q,d,nc,idx = Idxs[g]; 628 if (idx >= 0) { 629 nc = 1; 630 } else { 631 nc = nfaces; 632 } 633 for (q = 0; q < nr; q++) { 634 for (d = 0; d < nc; d++) { 635 vals[q*nc + d] = s_scale(f,q)*s_scale(g,d)*s_fieldMats(f,g); 636 } 637 } 638 MatSetValuesDevice(d_mat,nr,&s_idx(f,0),nc,&s_idx(g,0),vals,ADD_VALUES); 639 }); 640 }); 641 } 642 } 643 }); 644 ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr); 645 ierr = PetscLogEventEnd(events[4],0,0,0,0);CHKERRQ(ierr); 646 } 647 Kokkos::fence(); 648 649 if (global_elem_mat_sz) { 650 PetscSection section, globalSection; 651 Kokkos::View<PetscScalar**, Kokkos::LayoutRight>::HostMirror h_elem_mats = Kokkos::create_mirror_view(d_elem_mats); 652 ierr = PetscLogEventBegin(events[5],0,0,0,0);CHKERRQ(ierr); 653 ierr = DMGetLocalSection(plex, §ion);CHKERRQ(ierr); 654 ierr = DMGetGlobalSection(plex, &globalSection);CHKERRQ(ierr); 655 Kokkos::deep_copy (h_elem_mats, d_elem_mats); 656 ierr = PetscLogEventEnd(events[5],0,0,0,0);CHKERRQ(ierr); 657 ierr = PetscLogEventBegin(events[6],0,0,0,0);CHKERRQ(ierr); 658 PetscInt ej; 659 for (ej = cStart ; ej < cEnd; ++ej) { 660 const PetscScalar *elMat = &h_elem_mats(ej-cStart,0); 661 ierr = DMPlexMatSetClosure(plex, section, globalSection, JacP, ej, elMat, ADD_VALUES);CHKERRQ(ierr); 662 if (ej==-1) { 663 int d,f; 664 PetscPrintf(PETSC_COMM_SELF,"Kokkos Element matrix %d/%d\n",1,(int)numCells); 665 for (d = 0; d < totDim; ++d){ 666 for (f = 0; f < totDim; ++f) PetscPrintf(PETSC_COMM_SELF," %12.5e", PetscRealPart(elMat[d*totDim + f])); 667 PetscPrintf(PETSC_COMM_SELF,"\n"); 668 } 669 exit(14); 670 } 671 } 672 ierr = PetscLogEventEnd(events[6],0,0,0,0);CHKERRQ(ierr); 673 // transition to use of maps for VecGetClosure 674 if (ctx->gpu_assembly) { 675 auto IPf = static_cast<Kokkos::View<PetscScalar*, Kokkos::LayoutLeft>*>(SData_d->IPf); 676 delete IPf; 677 SData_d->IPf = NULL; 678 if (!(a_IPf || a_xarray)) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "transition without Jacobian"); 679 } 680 } 681 PetscFunctionReturn(0); 682 } 683 } // extern "C" 684