1 /* 2 Implements the Kokkos kernel 3 */ 4 #define PETSC_SKIP_CXX_COMPLEX_FIX 5 #include <petscconf.h> 6 #include <petsc/private/dmpleximpl.h> /*I "petscdmplex.h" I*/ 7 #include <petsclandau.h> 8 #include <petscts.h> 9 #include <petscveckokkos.hpp> 10 11 #include <Kokkos_Core.hpp> 12 #include <cstdio> 13 typedef Kokkos::TeamPolicy<>::member_type team_member; 14 #define PETSC_DEVICE_FUNC_DECL KOKKOS_INLINE_FUNCTION 15 #include "../land_tensors.h" 16 #define atomicAdd(e, f) Kokkos::atomic_fetch_add(e, f) 17 #include <petscaijdevice.h> 18 #undef atomicAdd 19 20 namespace landau_inner_red { // namespace helps with name resolution in reduction identity 21 template< class ScalarType > 22 struct array_type { 23 ScalarType gg2[LANDAU_DIM]; 24 ScalarType gg3[LANDAU_DIM][LANDAU_DIM]; 25 26 KOKKOS_INLINE_FUNCTION // Default constructor - Initialize to 0's 27 array_type() { 28 for (int j = 0; j < LANDAU_DIM; j++){ 29 gg2[j] = 0; 30 for (int k = 0; k < LANDAU_DIM; k++){ 31 gg3[j][k] = 0; 32 } 33 } 34 } 35 KOKKOS_INLINE_FUNCTION // Copy Constructor 36 array_type(const array_type & rhs) { 37 for (int j = 0; j < LANDAU_DIM; j++){ 38 gg2[j] = rhs.gg2[j]; 39 for (int k = 0; k < LANDAU_DIM; k++){ 40 gg3[j][k] = rhs.gg3[j][k]; 41 } 42 } 43 } 44 KOKKOS_INLINE_FUNCTION // add operator 45 array_type& operator += (const array_type& src) { 46 for (int j = 0; j < LANDAU_DIM; j++){ 47 gg2[j] += src.gg2[j]; 48 for (int k = 0; k < LANDAU_DIM; k++){ 49 gg3[j][k] += src.gg3[j][k]; 50 } 51 } 52 return *this; 53 } 54 KOKKOS_INLINE_FUNCTION // volatile add operator 55 void operator += (const volatile array_type& src) volatile { 56 for (int j = 0; j < LANDAU_DIM; j++){ 57 gg2[j] += src.gg2[j]; 58 for (int k = 0; k < LANDAU_DIM; k++){ 59 gg3[j][k] += src.gg3[j][k]; 60 } 61 } 62 } 63 }; 64 typedef array_type<PetscReal> TensorValueType; // used to simplify code below 65 } 66 67 namespace Kokkos { //reduction identity must be defined in Kokkos namespace 68 template<> 69 struct reduction_identity< landau_inner_red::TensorValueType > { 70 KOKKOS_FORCEINLINE_FUNCTION static landau_inner_red::TensorValueType sum() { 71 return landau_inner_red::TensorValueType(); 72 } 73 }; 74 } 75 76 extern "C" { 77 PetscErrorCode LandauKokkosCreateMatMaps(P4estVertexMaps *maps, pointInterpolationP4est (*pointMaps)[LANDAU_MAX_Q_FACE], PetscInt Nf, PetscInt Nq) 78 { 79 P4estVertexMaps h_maps; /* host container */ 80 const Kokkos::View<pointInterpolationP4est*[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_points ((pointInterpolationP4est*)pointMaps, maps->num_reduced); 81 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); 82 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); 83 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); 84 85 PetscFunctionBegin; 86 Kokkos::deep_copy (*d_gidx, h_gidx); 87 Kokkos::deep_copy (*d_points, h_points); 88 h_maps.num_elements = maps->num_elements; 89 h_maps.num_face = maps->num_face; 90 h_maps.num_reduced = maps->num_reduced; 91 h_maps.deviceType = maps->deviceType; 92 h_maps.Nf = Nf; 93 h_maps.Nq = Nq; 94 h_maps.c_maps = (pointInterpolationP4est (*)[LANDAU_MAX_Q_FACE]) d_points->data(); 95 maps->vp1 = (void*)d_points; 96 h_maps.gIdx = (LandauIdx (*)[LANDAU_MAX_SPECIES][LANDAU_MAX_NQ]) d_gidx->data(); 97 maps->vp2 = (void*)d_gidx; 98 { 99 Kokkos::View<P4estVertexMaps, Kokkos::HostSpace> h_maps_k(&h_maps); 100 Kokkos::View<P4estVertexMaps> *d_maps_k = new Kokkos::View<P4estVertexMaps>(Kokkos::create_mirror(DeviceMemorySpace(),h_maps_k)); 101 Kokkos::deep_copy (*d_maps_k, h_maps_k); 102 maps->data = d_maps_k->data(); 103 maps->vp3 = (void*)d_maps_k; 104 } 105 PetscFunctionReturn(0); 106 } 107 PetscErrorCode LandauKokkosDestroyMatMaps(P4estVertexMaps *maps) 108 { 109 Kokkos::View<pointInterpolationP4est*[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight> *a = (Kokkos::View<pointInterpolationP4est*[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight>*)maps->vp1; 110 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; 111 Kokkos::View<P4estVertexMaps*> *c = (Kokkos::View<P4estVertexMaps*>*)maps->vp3; 112 delete a; delete b; delete c; 113 PetscFunctionReturn(0); 114 } 115 116 PetscErrorCode LandauKokkosStaticDataSet(DM plex, const PetscInt Nq, PetscReal nu_alpha[], PetscReal nu_beta[], PetscReal a_invMass[], PetscReal a_invJ[], PetscReal a_mass_w[], 117 PetscReal a_x[], PetscReal a_y[], PetscReal a_z[], PetscReal a_w[], LandauGeomData *SData_d) 118 { 119 PetscReal *BB,*DD; 120 PetscErrorCode ierr; 121 PetscTabulation *Tf; 122 LandauCtx *ctx; 123 PetscInt *Nbf,dim,Nf,Nb,nip,cStart,cEnd; 124 PetscDS prob; 125 126 PetscFunctionBegin; 127 ierr = DMGetApplicationContext(plex, &ctx);CHKERRQ(ierr); 128 if (!ctx) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "no context"); 129 ierr = DMGetDimension(plex, &dim);CHKERRQ(ierr); 130 ierr = DMPlexGetHeightStratum(plex,0,&cStart,&cEnd);CHKERRQ(ierr); 131 nip = (cEnd - cStart)*Nq; 132 ierr = DMGetDS(plex, &prob);CHKERRQ(ierr); 133 ierr = PetscDSGetNumFields(prob, &Nf);CHKERRQ(ierr); 134 ierr = PetscDSGetDimensions(prob, &Nbf);CHKERRQ(ierr); Nb = Nbf[0]; 135 if (Nq != Nb) SETERRQ2(PETSC_COMM_SELF, PETSC_ERR_PLIB, "Nq != Nb. %D %D",Nq,Nb); 136 if (LANDAU_DIM != dim) SETERRQ2(PETSC_COMM_WORLD, PETSC_ERR_PLIB, "dim %D != LANDAU_DIM %d",dim,LANDAU_DIM); 137 ierr = PetscDSGetTabulation(prob, &Tf);CHKERRQ(ierr); 138 BB = Tf[0]->T[0]; DD = Tf[0]->T[1]; 139 ierr = PetscKokkosInitializeCheck();CHKERRQ(ierr); 140 { 141 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_alpha (nu_alpha, Nf); 142 auto alpha = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("alpha", Nf); 143 SData_d->alpha = static_cast<void*>(alpha); 144 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_beta (nu_beta, Nf); 145 auto beta = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("beta", Nf); 146 SData_d->beta = static_cast<void*>(beta); 147 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_invMass (a_invMass,Nf); 148 auto invMass = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("invMass", Nf); 149 SData_d->invMass = static_cast<void*>(invMass); 150 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_BB (BB,Nq*Nb); 151 auto B = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("B", Nq*Nb); 152 SData_d->B = static_cast<void*>(B); 153 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_DD (DD,Nq*Nb*dim); 154 auto D = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("D", Nq*Nb*dim); 155 SData_d->D = static_cast<void*>(D); 156 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_mass_w (a_mass_w, nip); 157 auto mass_w = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("mass_w", nip); 158 SData_d->mass_w = static_cast<void*>(mass_w); 159 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_invJ (a_invJ, nip*dim*dim); 160 auto invJ = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("invJ", nip*dim*dim); 161 SData_d->invJ = static_cast<void*>(invJ); 162 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_x (a_x, nip); 163 auto x = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("x", nip); 164 SData_d->x = static_cast<void*>(x); 165 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_y (a_y, nip); 166 auto y = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("y", nip); 167 SData_d->y = static_cast<void*>(y); 168 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_w (a_w, nip); 169 auto w = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("w", nip); 170 SData_d->w = static_cast<void*>(w); 171 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_z (a_z , dim==3 ? nip : 0); 172 auto z = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("z", dim==3 ? nip : 0); 173 SData_d->z = static_cast<void*>(z); 174 175 Kokkos::deep_copy (*mass_w, h_mass_w); 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 (*z, h_z); 185 Kokkos::deep_copy (*w, h_w); 186 187 auto Eq_m = new Kokkos::View<PetscReal*, Kokkos::LayoutLeft> ("Eq_m",Nf); 188 SData_d->Eq_m = static_cast<void*>(Eq_m); 189 auto IPf = new Kokkos::View<PetscScalar*, Kokkos::LayoutLeft> ("IPf",nip*Nf); // Nq==Nb 190 SData_d->IPf = static_cast<void*>(IPf); 191 } 192 PetscFunctionReturn(0); 193 } 194 195 PetscErrorCode LandauKokkosStaticDataClear(LandauGeomData *SData_d) 196 { 197 PetscFunctionBegin; 198 if (SData_d->alpha) { 199 auto alpha = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->alpha); 200 delete alpha; 201 auto beta = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->beta); 202 delete beta; 203 auto invMass = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->invMass); 204 delete invMass; 205 auto B = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->B); 206 delete B; 207 auto D = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->D); 208 delete D; 209 auto mass_w = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->mass_w); 210 delete mass_w; 211 auto invJ = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->invJ); 212 delete invJ; 213 auto x = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->x); 214 delete x; 215 auto y = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->y); 216 delete y; 217 auto z = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->z); 218 delete z; 219 auto w = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->w); 220 delete w; 221 auto Eq_m = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->Eq_m); 222 delete Eq_m; 223 auto IPf = static_cast<Kokkos::View<PetscScalar*, Kokkos::LayoutLeft>*>(SData_d->IPf); 224 delete IPf; 225 } 226 PetscFunctionReturn(0); 227 } 228 229 PetscErrorCode LandauKokkosJacobian(DM plex, const PetscInt Nq, PetscReal a_Eq_m[], PetscScalar a_IPf[], LandauGeomData *SData_d, const PetscInt num_sub_blocks, PetscReal shift, 230 const PetscLogEvent events[], Mat JacP) 231 { 232 using scr_mem_t = Kokkos::DefaultExecutionSpace::scratch_memory_space; 233 using g2_scr_t = Kokkos::View<PetscReal***, Kokkos::LayoutRight, scr_mem_t>; 234 using g3_scr_t = Kokkos::View<PetscReal****, Kokkos::LayoutRight, scr_mem_t>; 235 PetscErrorCode ierr; 236 PetscInt *Nbf,Nb,cStart,cEnd,Nf,dim,numCells,totDim,global_elem_mat_sz,nip; 237 PetscDS prob; 238 LandauCtx *ctx; 239 PetscReal *d_Eq_m=NULL; 240 PetscScalar *d_IPf=NULL; 241 P4estVertexMaps *d_maps=NULL; 242 PetscSplitCSRDataStructure *d_mat=NULL; 243 const int conc = Kokkos::DefaultExecutionSpace().concurrency(), team_size = conc > Nq ? Nq : 1; 244 int scr_bytes,nnn = 256/Nq; 245 auto d_alpha_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->alpha); //static data 246 const PetscReal *d_alpha = d_alpha_k->data(); 247 auto d_beta_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->beta); 248 const PetscReal *d_beta = d_beta_k->data(); 249 auto d_invMass_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->invMass); 250 const PetscReal *d_invMass = d_invMass_k->data(); 251 auto d_B_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->B); 252 const PetscReal *d_BB = d_B_k->data(); 253 auto d_D_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->D); 254 const PetscReal *d_DD = d_D_k->data(); 255 auto d_mass_w_k = *static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->mass_w); 256 auto d_invJ_k = *static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->invJ); // use Kokkos vector in kernels 257 auto d_x_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->x); //static data 258 const PetscReal *d_x = d_x_k->data(); 259 auto d_y_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->y); //static data 260 const PetscReal *d_y = d_y_k->data(); 261 auto d_z_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->z); //static data 262 const PetscReal *d_z = d_z_k->data(); 263 auto d_w_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->w); //static data 264 const PetscReal *d_w = d_w_k->data(); 265 // dynamic data pointers 266 auto d_Eq_m_k = static_cast<Kokkos::View<PetscReal*, Kokkos::LayoutLeft>*>(SData_d->Eq_m); //static data 267 auto d_IPf_k = static_cast<Kokkos::View<PetscScalar*, Kokkos::LayoutLeft>*>(SData_d->IPf); //static data 268 269 PetscFunctionBegin; 270 while (nnn & (nnn - 1)) nnn = nnn & (nnn - 1); 271 if (nnn>16) nnn = 16; 272 ierr = PetscLogEventBegin(events[3],0,0,0,0);CHKERRQ(ierr); 273 ierr = DMGetApplicationContext(plex, &ctx);CHKERRQ(ierr); 274 if (!ctx) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "no context"); 275 ierr = DMGetDimension(plex, &dim);CHKERRQ(ierr); 276 ierr = DMPlexGetHeightStratum(plex,0,&cStart,&cEnd);CHKERRQ(ierr); 277 numCells = cEnd - cStart; 278 nip = numCells*Nq; 279 ierr = DMGetDS(plex, &prob);CHKERRQ(ierr); 280 ierr = PetscDSGetNumFields(prob, &Nf);CHKERRQ(ierr); 281 ierr = PetscDSGetDimensions(prob, &Nbf);CHKERRQ(ierr); Nb = Nbf[0]; 282 if (Nq != Nb) SETERRQ2(PETSC_COMM_SELF, PETSC_ERR_PLIB, "Nq != Nb. %D %D",Nq,Nb); 283 if (LANDAU_DIM != dim) SETERRQ2(PETSC_COMM_WORLD, PETSC_ERR_PLIB, "dim %D != LANDAU_DIM %d",dim,LANDAU_DIM); 284 scr_bytes = 2*(g2_scr_t::shmem_size(dim,Nf,Nq) + g3_scr_t::shmem_size(dim,dim,Nf,Nq)); 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 } else { 296 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "GPU assembly but no metadata in container"); 297 } 298 // this does the setup the first time called 299 ierr = MatKokkosGetDeviceMatWrite(JacP,&d_mat);CHKERRQ(ierr); 300 global_elem_mat_sz = 0; 301 #else 302 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "GPU assembly w/o kokkos kernels -- should not be here"); 303 #endif 304 } else { // kernel output - first call assembled on device 305 global_elem_mat_sz = numCells; 306 } 307 } else { 308 global_elem_mat_sz = numCells; // no device assembly 309 } 310 ierr = PetscKokkosInitializeCheck();CHKERRQ(ierr); 311 ierr = PetscLogEventEnd(events[3],0,0,0,0);CHKERRQ(ierr); 312 { 313 Kokkos::View<PetscReal**, Kokkos::LayoutRight> d_f_k ("f", Nf, a_IPf ? nip : 0); // allocated each run, computed and used on GPU 314 Kokkos::View<PetscReal***, Kokkos::LayoutRight> d_df_k("df", dim, Nf, a_IPf ? nip : 0); 315 const Kokkos::View<PetscScalar*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> >h_IPf_k (a_IPf, a_IPf ? nip*Nf : 0); 316 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_Eq_m_k (a_Eq_m, a_IPf ? Nf : 0); 317 Kokkos::View<PetscScalar**, Kokkos::LayoutRight> d_elem_mats("element matrices", global_elem_mat_sz, totDim*totDim); 318 // copy dynamic data to device 319 if (a_IPf) { 320 ierr = PetscLogEventBegin(events[1],0,0,0,0);CHKERRQ(ierr); 321 Kokkos::deep_copy (*d_IPf_k, h_IPf_k); 322 Kokkos::deep_copy (*d_Eq_m_k, h_Eq_m_k); 323 d_Eq_m = d_Eq_m_k->data(); 324 d_IPf = d_IPf_k->data(); 325 ierr = PetscLogEventEnd(events[1],0,0,0,0);CHKERRQ(ierr); 326 #define KOKKOS_SHARED_LEVEL 1 327 // PetscInfo5(plex, "shared memory size: %d bytes in level %d. mass_w.size=%d. #threads=%D team size=%D\n",scr_bytes,KOKKOS_SHARED_LEVEL,d_mass_w.size(),nnn,team_size); 328 // get f and df 329 ierr = PetscLogEventBegin(events[8],0,0,0,0);CHKERRQ(ierr); 330 ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr); 331 Kokkos::parallel_for("f, df", Kokkos::TeamPolicy<>(numCells, team_size, nnn>Nf ? Nf : nnn), KOKKOS_LAMBDA (const team_member team) { 332 const PetscInt myelem = team.league_rank(); 333 // un pack IPData 334 PetscScalar *coef = &d_IPf[myelem*Nb*Nf]; 335 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nq), [=] (int myQi) { 336 const PetscInt ipidx = myQi + myelem * Nq; 337 const PetscReal *const invJj = &d_invJ_k(ipidx*dim*dim); 338 const PetscReal *Bq = &d_BB[myQi*Nb], *Dq = &d_DD[myQi*Nb*dim]; 339 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,(int)Nf), [=] (int f) { 340 PetscInt b, e, d; 341 PetscReal refSpaceDer[LANDAU_DIM]; 342 d_f_k(f,ipidx) = 0.0; 343 for (d = 0; d < LANDAU_DIM; ++d) refSpaceDer[d] = 0.0; 344 for (b = 0; b < Nb; ++b) { 345 const PetscInt cidx = b; 346 d_f_k(f,ipidx) += Bq[cidx]*coef[f*Nb+cidx]; 347 for (d = 0; d < dim; ++d) refSpaceDer[d] += Dq[cidx*dim+d]*coef[f*Nb+cidx]; 348 } 349 for (d = 0; d < dim; ++d) { 350 for (e = 0, d_df_k(d,f,ipidx) = 0.0; e < dim; ++e) { 351 d_df_k(d,f,ipidx) += invJj[e*dim+d]*refSpaceDer[e]; 352 } 353 } 354 }); // Nf 355 }); // Nq 356 }); // elems 357 Kokkos::fence(); 358 #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) 359 ierr = PetscLogGpuFlops(nip*(PetscLogDouble)(2*Nb*(1+dim)));CHKERRQ(ierr); 360 #else 361 ierr = PetscLogFlops(nip*(PetscLogDouble)(2*Nb*(1+dim)));CHKERRQ(ierr); 362 #endif 363 ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr); 364 ierr = PetscLogEventEnd(events[8],0,0,0,0);CHKERRQ(ierr); 365 } 366 ierr = PetscLogEventBegin(events[4],0,0,0,0);CHKERRQ(ierr); 367 ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr); 368 #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) 369 ierr = PetscLogGpuFlops(nip*(PetscLogDouble)(!a_IPf ? (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) : Nb*Nf*Nb*Nq*4));CHKERRQ(ierr); 370 if (ctx->deviceType == LANDAU_CPU) PetscInfo(plex, "Warning: Landau selected CPU but no support for Kokkos using CPU\n"); 371 #else 372 ierr = PetscLogFlops(nip*(PetscLogDouble)(!a_IPf ? (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) : Nb*Nf*Nb*Nq*4));CHKERRQ(ierr); 373 #endif 374 Kokkos::parallel_for("Landau elements", Kokkos::TeamPolicy<>(numCells, team_size, nnn).set_scratch_size(KOKKOS_SHARED_LEVEL, Kokkos::PerTeam(scr_bytes)), KOKKOS_LAMBDA (const team_member team) { 375 const PetscInt myelem = team.league_rank(); 376 g2_scr_t g2(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,Nf,Nq); // we don't use these for mass matrix 377 g3_scr_t g3(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,dim,Nf,Nq); 378 if (d_IPf) { 379 g2_scr_t gg2(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,Nf,Nq); 380 g3_scr_t gg3(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,dim,Nf,Nq); 381 // get g2[] & g3[] 382 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nq), [=] (int myQi) { 383 using Kokkos::parallel_reduce; 384 const PetscInt jpidx = myQi + myelem * Nq; 385 const PetscReal* const invJj = &d_invJ_k(jpidx*dim*dim); 386 const PetscReal vj[3] = {d_x[jpidx], d_y[jpidx], d_z ? d_z[jpidx] : 0}, wj = d_w[jpidx]; 387 landau_inner_red::TensorValueType gg_temp; // reduce on part of gg2 and g33 for IP jpidx 388 Kokkos::parallel_reduce(Kokkos::ThreadVectorRange (team, (int)nip), [=] (const int& ipidx, landau_inner_red::TensorValueType & ggg) { 389 const PetscReal wi = d_w[ipidx], x = d_x[ipidx], y = d_y[ipidx]; 390 //printf("\t\t:%d.%d) ip %d, w=%e x=%e y=%e f[0]=%e\n",myelem,myQi,ipidx,wi,x,y,d_df_k(0,0,ipidx)); 391 PetscReal temp1[3] = {0, 0, 0}, temp2 = 0; 392 PetscInt fieldA,d2,d3; 393 #if LANDAU_DIM==2 394 PetscReal Ud[2][2], Uk[2][2]; 395 LandauTensor2D(vj, x, y, Ud, Uk, (ipidx==jpidx) ? 0. : 1.); 396 #else 397 PetscReal U[3][3], z = d_z[ipidx]; 398 LandauTensor3D(vj, x, y, z, U, (ipidx==jpidx) ? 0. : 1.); 399 #endif 400 for (fieldA = 0; fieldA < Nf; ++fieldA) { 401 temp1[0] += d_df_k(0,fieldA,ipidx)*d_beta[fieldA]*d_invMass[fieldA]; 402 temp1[1] += d_df_k(1,fieldA,ipidx)*d_beta[fieldA]*d_invMass[fieldA]; 403 #if LANDAU_DIM==3 404 temp1[2] += d_df_k(2,fieldA,ipidx)*d_beta[fieldA]*d_invMass[fieldA]; 405 #endif 406 temp2 += d_f_k(fieldA,ipidx)*d_beta[fieldA]; 407 } 408 temp1[0] *= wi; 409 temp1[1] *= wi; 410 #if LANDAU_DIM==3 411 temp1[2] *= wi; 412 #endif 413 temp2 *= wi; 414 #if LANDAU_DIM==2 415 for (d2 = 0; d2 < 2; d2++) { 416 for (d3 = 0; d3 < 2; ++d3) { 417 /* K = U * grad(f): g2=e: i,A */ 418 ggg.gg2[d2] += Uk[d2][d3]*temp1[d3]; 419 /* D = -U * (I \kron (fx)): g3=f: i,j,A */ 420 ggg.gg3[d2][d3] += Ud[d2][d3]*temp2; 421 } 422 } 423 #else 424 for (d2 = 0; d2 < 3; ++d2) { 425 for (d3 = 0; d3 < 3; ++d3) { 426 /* K = U * grad(f): g2 = e: i,A */ 427 ggg.gg2[d2] += U[d2][d3]*temp1[d3]; 428 /* D = -U * (I \kron (fx)): g3 = f: i,j,A */ 429 ggg.gg3[d2][d3] += U[d2][d3]*temp2; 430 } 431 } 432 #endif 433 }, Kokkos::Sum<landau_inner_red::TensorValueType>(gg_temp)); 434 // if (myelem==0) printf("\t:%d.%d) temp gg3=%e %e %e %e\n",myelem,myQi,gg_temp.gg3[0][0],gg_temp.gg3[1][0],gg_temp.gg3[0][1],gg_temp.gg3[1][1]); 435 // add alpha and put in gg2/3 436 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)Nf), [&] (const int& fieldA) { 437 PetscInt d2,d3; 438 for (d2 = 0; d2 < dim; d2++) { 439 gg2(d2,fieldA,myQi) = gg_temp.gg2[d2]*d_alpha[fieldA]; 440 //if (myelem==0 && fieldA==1) printf("\t\t:%d.%d) gg2[%d]=%e = %e %e\n",myelem,myQi,d2,gg2(d2,fieldA,myQi),gg_temp.gg2[d2],d_alpha[fieldA]); 441 //gg2[d2][myQi][fieldA] += gg_temp.gg2[d2]*d_alpha[fieldA]; 442 for (d3 = 0; d3 < dim; d3++) { 443 //gg3[d2][d3][myQi][fieldA] -= gg_temp.gg3[d2][d3]*d_alpha[fieldA]*s_invMass[fieldA]; 444 gg3(d2,d3,fieldA,myQi) = -gg_temp.gg3[d2][d3]*d_alpha[fieldA]*d_invMass[fieldA]; 445 //if (myelem==0 && fieldA==1) printf("\t\t\t:%d.%d) gg3[%d][%d]=%e\n",myelem,myQi,d2,d3,gg3(d2,d3,fieldA,myQi)); 446 } 447 } 448 }); 449 /* add electric field term once per IP */ 450 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)Nf), [&] (const int& fieldA) { 451 //gg.gg2[fieldA][dim-1] += d_Eq_m[fieldA]; 452 gg2(dim-1,fieldA,myQi) += d_Eq_m[fieldA]; 453 }); 454 // Kokkos::single(Kokkos::PerThread(team), [&]() { 455 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)Nf), [=] (const int& fieldA) { 456 int d,d2,d3,dp; 457 //printf("%d %d %d gg2[][1]=%18.10e\n",myelem,myQi,fieldA,gg.gg2[fieldA][dim-1]); 458 /* Jacobian transform - g2, g3 - per thread (2D) */ 459 for (d = 0; d < dim; ++d) { 460 g2(d,fieldA,myQi) = 0; 461 for (d2 = 0; d2 < dim; ++d2) { 462 g2(d,fieldA,myQi) += invJj[d*dim+d2]*gg2(d2,fieldA,myQi); 463 //if (myelem==0 && myQi==0) printf("\t:g2[%d][%d][%d]=%e. %e %e\n",(int)myQi,(int)fieldA,(int)d,g2(fieldA,myQi,d),invJj[d*dim+d2],gg.gg2[fieldA][d2]); 464 g3(d,d2,fieldA,myQi) = 0; 465 for (d3 = 0; d3 < dim; ++d3) { 466 for (dp = 0; dp < dim; ++dp) { 467 g3(d,d2,fieldA,myQi) += invJj[d*dim + d3]*gg3(d3,dp,fieldA,myQi)*invJj[d2*dim + dp]; 468 //printf("\t%d %d %d %d %d %d %d g3=%g wj=%g g3 = %g * %g * %g\n",myelem,myQi,fieldA,d,d2,d3,dp,g3(fieldA,myQi,d,d2),wj,invJj[d*dim + d3],gg.gg3[fieldA][d3][dp],invJj[d2*dim + dp]); 469 } 470 } 471 g3(d,d2,fieldA,myQi) *= wj; 472 } 473 g2(d,fieldA,myQi) *= wj; 474 } 475 }); 476 }); // Nq 477 team.team_barrier(); 478 } // Jacobian 479 /* assemble */ 480 //Kokkos::single(Kokkos::PerTeam(team), [&]() { 481 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nb), [=] (int blk_i) { 482 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,(int)Nf), [=] (int fieldA) { 483 int blk_j,qj,d,d2; 484 const PetscInt i = fieldA*Nb + blk_i; /* Element matrix row */ 485 for (blk_j = 0; blk_j < Nb; ++blk_j) { 486 const PetscInt j = fieldA*Nb + blk_j; /* Element matrix column */ 487 const PetscInt fOff = i*totDim + j; 488 PetscScalar t = global_elem_mat_sz ? d_elem_mats(myelem,fOff) : 0; 489 for (qj = 0 ; qj < Nq ; qj++) { // look at others integration points 490 const PetscReal *BJq = &d_BB[qj*Nb], *DIq = &d_DD[qj*Nb*dim]; 491 if (d_IPf) { 492 for (d = 0; d < dim; ++d) { 493 t += DIq[blk_i*dim+d]*g2(d,fieldA,qj)*BJq[blk_j]; 494 for (d2 = 0; d2 < dim; ++d2) { 495 t += DIq[blk_i*dim + d]*g3(d,d2,fieldA,qj)*DIq[blk_j*dim + d2]; 496 } 497 } 498 } else { 499 const PetscInt jpidx = qj + myelem * Nq; 500 t += BJq[blk_i] * d_mass_w_k(jpidx) * shift * BJq[blk_j]; 501 //printf("\tmat[%d %d %d %d %d]: B=%g w=%g shift=%g B=%g\n",myelem,fOff,fieldA,qj,d,BJq[blk_i],d_mass_w_k(jpidx),shift,BJq[blk_j]); 502 } 503 } 504 if (global_elem_mat_sz) d_elem_mats(myelem,fOff) = t; // can set this because local element matrix[fOff] 505 else { 506 PetscErrorCode ierr = 0; 507 PetscScalar vals[LANDAU_MAX_Q_FACE*LANDAU_MAX_Q_FACE],row_scale[LANDAU_MAX_Q_FACE],col_scale[LANDAU_MAX_Q_FACE]; 508 PetscInt q,idx,nr,nc,rows0[LANDAU_MAX_Q_FACE],cols0[LANDAU_MAX_Q_FACE],rows[LANDAU_MAX_Q_FACE],cols[LANDAU_MAX_Q_FACE]; 509 const LandauIdx *const Idxs = &d_maps->gIdx[myelem][fieldA][0]; 510 idx = Idxs[blk_i]; 511 if (idx >= 0) { 512 nr = 1; 513 rows0[0] = idx; 514 row_scale[0] = 1.; 515 } else { 516 idx = -idx - 1; 517 nr = d_maps->num_face; 518 for (q = 0; q < d_maps->num_face; q++) { 519 rows0[q] = d_maps->c_maps[idx][q].gid; 520 row_scale[q] = d_maps->c_maps[idx][q].scale; 521 } 522 } 523 idx = Idxs[blk_j]; 524 if (idx >= 0) { 525 nc = 1; 526 cols0[0] = idx; 527 col_scale[0] = 1.; 528 } else { 529 idx = -idx - 1; 530 nc = d_maps->num_face; 531 for (q = 0; q < d_maps->num_face; q++) { 532 cols0[q] = d_maps->c_maps[idx][q].gid; 533 col_scale[q] = d_maps->c_maps[idx][q].scale; 534 } 535 } 536 for (q = 0; q < nr; q++) rows[q] = rows0[q]; 537 for (q = 0; q < nc; q++) cols[q] = cols0[q]; 538 for (q = 0; q < nr; q++) { 539 for (d = 0; d < nc; d++) { 540 vals[q*nc + d] = row_scale[q]*col_scale[d]*t; 541 } 542 } 543 MatSetValuesDevice(d_mat,nr,rows,nc,cols,vals,ADD_VALUES,&ierr); 544 if (ierr) return; 545 } 546 } 547 }); 548 }); 549 }); 550 ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr); 551 Kokkos::fence(); 552 ierr = PetscLogEventEnd(events[4],0,0,0,0);CHKERRQ(ierr); 553 if (global_elem_mat_sz) { 554 PetscSection section, globalSection; 555 Kokkos::View<PetscScalar**, Kokkos::LayoutRight>::HostMirror h_elem_mats = Kokkos::create_mirror_view(d_elem_mats); 556 ierr = PetscLogEventBegin(events[5],0,0,0,0);CHKERRQ(ierr); 557 ierr = DMGetLocalSection(plex, §ion);CHKERRQ(ierr); 558 ierr = DMGetGlobalSection(plex, &globalSection);CHKERRQ(ierr); 559 Kokkos::deep_copy (h_elem_mats, d_elem_mats); 560 ierr = PetscLogEventEnd(events[5],0,0,0,0);CHKERRQ(ierr); 561 ierr = PetscLogEventBegin(events[6],0,0,0,0);CHKERRQ(ierr); 562 PetscInt ej; 563 for (ej = cStart ; ej < cEnd; ++ej) { 564 const PetscScalar *elMat = &h_elem_mats(ej-cStart,0); 565 ierr = DMPlexMatSetClosure(plex, section, globalSection, JacP, ej, elMat, ADD_VALUES);CHKERRQ(ierr); 566 if (ej==-1) { 567 int d,f; 568 PetscPrintf(PETSC_COMM_SELF,"Kokkos Element matrix %d/%d\n",1,(int)numCells); 569 for (d = 0; d < totDim; ++d){ 570 for (f = 0; f < totDim; ++f) PetscPrintf(PETSC_COMM_SELF," %12.5e", PetscRealPart(elMat[d*totDim + f])); 571 PetscPrintf(PETSC_COMM_SELF,"\n"); 572 } 573 exit(14); 574 } 575 } 576 ierr = PetscLogEventEnd(events[6],0,0,0,0);CHKERRQ(ierr); 577 } 578 } 579 PetscFunctionReturn(0); 580 } 581 } // extern "C" 582