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 // debug 105 // PetscErrorCode ierr; 106 // PetscInt ej,q,s; 107 // ierr = PetscPrintf(PETSC_COMM_SELF,"=== c_maps size --> %D. Kokkos::View<P4estVertexMaps*>=%p maps->data=%p P4estVertexMaps *maps=%p\n",maps->num_reduced, d_maps_k, maps->data, maps);CHKERRQ(ierr); 108 // for (ej = 0; ej < maps->num_reduced; ++ej) { 109 // ierr = PetscPrintf(PETSC_COMM_SELF,"\t\tconstrant %3D) ",ej); 110 // for (q = 0; q < maps->num_face; ++q) { 111 // ierr = PetscPrintf(PETSC_COMM_SELF,"\t %3D - %13.5e ; ", maps->data->c_maps[ej][q].gid, maps->data->c_maps[ej][q].scale);CHKERRQ(ierr); 112 // } 113 // ierr = PetscPrintf(PETSC_COMM_SELF,"\n"); 114 // } 115 // ierr = PetscPrintf(PETSC_COMM_SELF,"================= gIDx \n"); 116 // for (ej = 0; ej < maps->num_elements; ++ej) { 117 // for (q = 0; q < Nq; ++q) { 118 // for (s = 0; s < Nf; ++s) { 119 // ierr = PetscPrintf(PETSC_COMM_SELF," %3D ", maps->data->gIdx[ej][s][q]); 120 // } 121 // ierr = PetscPrintf(PETSC_COMM_SELF," ; "); 122 // } 123 // ierr = PetscPrintf(PETSC_COMM_SELF,"\n"); 124 // } 125 } 126 PetscFunctionReturn(0); 127 } 128 PetscErrorCode LandauKokkosDestroyMatMaps(P4estVertexMaps *maps) 129 { 130 Kokkos::View<pointInterpolationP4est*[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight> *a = (Kokkos::View<pointInterpolationP4est*[LANDAU_MAX_Q_FACE], Kokkos::LayoutRight>*)maps->vp1; 131 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; 132 Kokkos::View<P4estVertexMaps*> *c = (Kokkos::View<P4estVertexMaps*>*)maps->vp3; 133 delete a; delete b; delete c; 134 PetscFunctionReturn(0); 135 } 136 137 PetscErrorCode LandauKokkosJacobian(DM plex, const PetscInt Nq, PetscReal nu_alpha[], PetscReal nu_beta[], PetscReal invMass[], PetscReal Eq_m[], 138 const LandauIPData *const IPData, PetscReal a_invJ[], const PetscInt num_sub_blocks, PetscReal *mass_w, PetscReal shift, 139 const PetscLogEvent events[], Mat JacP) 140 { 141 PetscErrorCode ierr; 142 PetscInt *Nbf,Nb,cStart,cEnd,Nf,dim,numCells,totDim,ipdatasz,global_elem_mat_sz,NfJac,nip; 143 PetscTabulation *Tf; 144 PetscDS prob; 145 PetscSection section, globalSection; 146 PetscReal *BB,*DD; 147 LandauCtx *ctx; 148 P4estVertexMaps *d_maps=NULL; 149 PetscSplitCSRDataStructure *d_mat=NULL; 150 151 PetscFunctionBegin; 152 ierr = DMGetApplicationContext(plex, &ctx);CHKERRQ(ierr); 153 if (!ctx) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "no context"); 154 ierr = DMGetDimension(plex, &dim);CHKERRQ(ierr); 155 ierr = DMPlexGetHeightStratum(plex,0,&cStart,&cEnd);CHKERRQ(ierr); 156 numCells = cEnd - cStart; 157 ierr = DMGetDS(plex, &prob);CHKERRQ(ierr); 158 ierr = PetscDSGetNumFields(prob, &Nf);CHKERRQ(ierr); 159 ierr = PetscDSGetDimensions(prob, &Nbf);CHKERRQ(ierr); Nb = Nbf[0]; 160 if (Nq != Nb) SETERRQ2(PETSC_COMM_SELF, PETSC_ERR_PLIB, "Nq != Nb. %D %D",Nq,Nb); 161 if (LANDAU_DIM != dim) SETERRQ2(PETSC_COMM_WORLD, PETSC_ERR_PLIB, "dim %D != LANDAU_DIM %d",dim,LANDAU_DIM); 162 ierr = PetscDSGetTotalDimension(prob, &totDim);CHKERRQ(ierr); 163 ierr = PetscDSGetTabulation(prob, &Tf);CHKERRQ(ierr); 164 if (ctx->gpu_assembly) { 165 PetscContainer container; 166 ierr = PetscObjectQuery((PetscObject) JacP, "assembly_maps", (PetscObject *) &container);CHKERRQ(ierr); 167 if (container) { // not here first call 168 #if defined(PETSC_HAVE_KOKKOS_KERNELS) 169 P4estVertexMaps *h_maps=NULL; 170 ierr = PetscContainerGetPointer(container, (void **) &h_maps);CHKERRQ(ierr); 171 ierr = PetscInfo2(JacP, "Have container maps=%p maps->data=%p\n", h_maps, h_maps ? h_maps->data : NULL);CHKERRQ(ierr); 172 if (h_maps->data) { 173 d_maps = h_maps->data; 174 } else { 175 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "GPU assembly but no metadata in container"); 176 } 177 // this does the setup the first time called 178 ierr = MatKokkosGetDeviceMatWrite(JacP,&d_mat);CHKERRQ(ierr); 179 global_elem_mat_sz = 0; 180 #else 181 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "GPU assembly w/o kokkos kernels -- should not be here"); 182 #endif 183 } else { // kernel output - first call assembled on device 184 global_elem_mat_sz = numCells; 185 } 186 } else { 187 global_elem_mat_sz = numCells; // no device assembly 188 } 189 BB = Tf[0]->T[0]; DD = Tf[0]->T[1]; 190 ierr = DMGetLocalSection(plex, §ion);CHKERRQ(ierr); 191 ierr = DMGetGlobalSection(plex, &globalSection);CHKERRQ(ierr); 192 if (mass_w) { 193 ipdatasz = 0; NfJac = 0; nip = numCells*Nq; 194 } else { 195 ipdatasz = LandauGetIPDataSize(IPData); NfJac = Nf; nip = IPData->nip_; 196 } 197 ierr = PetscKokkosInitializeCheck();CHKERRQ(ierr); 198 { 199 using scr_mem_t = Kokkos::DefaultExecutionSpace::scratch_memory_space; 200 using g2_scr_t = Kokkos::View<PetscReal***, Kokkos::LayoutRight, scr_mem_t>; 201 using g3_scr_t = Kokkos::View<PetscReal****, Kokkos::LayoutRight, scr_mem_t>; 202 203 const int scr_bytes = 2*(g2_scr_t::shmem_size(dim,Nf,Nq) + g3_scr_t::shmem_size(dim,dim,Nf,Nq)); 204 int conc, team_size; 205 ierr = PetscLogEventBegin(events[3],0,0,0,0);CHKERRQ(ierr); 206 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_alpha (nu_alpha, NfJac); 207 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_alpha ("nu_alpha", NfJac); 208 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_beta (nu_beta, NfJac); 209 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_beta ("nu_beta", NfJac); 210 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_invMass (invMass,NfJac); 211 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_invMass ("invMass", NfJac); 212 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_Eq_m (Eq_m,NfJac); 213 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_Eq_m ("Eq_m", NfJac); 214 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_BB (BB,Nq*Nb); 215 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_BB ("BB", Nq*Nb); 216 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_DD (DD,Nq*Nb*dim); 217 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_DD ("DD", Nq*Nb*dim); 218 const Kokkos::View<LandauIPReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_ipdata_raw (IPData->w, ipdatasz); // IPData->w==0 for mass 219 Kokkos::View<LandauIPReal*, Kokkos::LayoutLeft> d_ipdata_raw ("ipdata", ipdatasz); 220 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_invJ (a_invJ, mass_w ? 0 : nip*dim*dim); 221 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_invJ ("invJ", mass_w ? 0 : nip*dim*dim); 222 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_mass_w (mass_w, mass_w ? nip : 0); 223 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_mass_w ("mass_w", mass_w ? nip : 0); 224 Kokkos::View<PetscScalar**, Kokkos::LayoutRight> d_elem_mats("element matrices", global_elem_mat_sz, totDim*totDim); 225 Kokkos::View<PetscScalar**, Kokkos::LayoutRight> d_f("element matrices", NfJac, nip); 226 Kokkos::View<PetscScalar***, Kokkos::LayoutRight> d_df("element matrices", dim, NfJac, nip); 227 Kokkos::deep_copy (d_mass_w, h_mass_w); 228 Kokkos::deep_copy (d_ipdata_raw, h_ipdata_raw); 229 Kokkos::deep_copy (d_alpha, h_alpha); 230 Kokkos::deep_copy (d_beta, h_beta); 231 Kokkos::deep_copy (d_invMass, h_invMass); 232 Kokkos::deep_copy (d_Eq_m, h_Eq_m); 233 Kokkos::deep_copy (d_BB, h_BB); 234 Kokkos::deep_copy (d_DD, h_DD); 235 Kokkos::deep_copy (d_invJ, h_invJ); 236 237 ierr = PetscLogEventEnd(events[3],0,0,0,0);CHKERRQ(ierr); 238 #define KOKKOS_SHARED_LEVEL 1 239 conc = Kokkos::DefaultExecutionSpace().concurrency(), team_size = conc > Nq ? Nq : 1; 240 // 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(),num_sub_blocks,team_size); 241 // get f and df 242 if (!mass_w) { 243 ierr = PetscLogEventBegin(events[8],0,0,0,0);CHKERRQ(ierr); 244 ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr); 245 Kokkos::parallel_for("df_d_elements", Kokkos::TeamPolicy<>(numCells, team_size, num_sub_blocks), KOKKOS_LAMBDA (const team_member team) { 246 const PetscInt myelem = team.league_rank(); 247 // un pack IPData 248 LandauIPReal *IPData_coefs = &d_ipdata_raw[nip*(dim+1)]; 249 LandauIPReal *coef = &IPData_coefs[myelem*Nb*Nf]; 250 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nq), [=] (int myQi) { 251 const PetscInt ipidx = myQi + myelem * Nq; 252 const PetscReal *const invJj = &d_invJ(ipidx*dim*dim); 253 const PetscReal *Bq = &d_BB[myQi*Nb], *Dq = &d_DD[myQi*Nb*dim]; 254 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,(int)Nf), [=] (int f) { 255 PetscInt b, e, d; 256 PetscScalar refSpaceDer[LANDAU_DIM]; 257 d_f(f,ipidx) = 0.0; 258 for (d = 0; d < LANDAU_DIM; ++d) refSpaceDer[d] = 0.0; 259 for (b = 0; b < Nb; ++b) { 260 const PetscInt cidx = b; 261 d_f(f,ipidx) += Bq[cidx]*coef[f*Nb+cidx]; 262 for (d = 0; d < dim; ++d) refSpaceDer[d] += Dq[cidx*dim+d]*coef[f*Nb+cidx]; 263 } 264 for (d = 0; d < dim; ++d) { 265 for (e = 0, d_df(d,f,ipidx) = 0.0; e < dim; ++e) { 266 d_df(d,f,ipidx) += invJj[e*dim+d]*refSpaceDer[e]; 267 } 268 } 269 }); // Nf 270 }); // Nq 271 }); // elems 272 Kokkos::fence(); 273 #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) 274 ierr = PetscLogGpuFlops(nip*(PetscLogDouble)(2*Nb*(1+dim)));CHKERRQ(ierr); 275 #else 276 ierr = PetscLogFlops(nip*(PetscLogDouble)(2*Nb*(1+dim)));CHKERRQ(ierr); 277 #endif 278 ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr); 279 ierr = PetscLogEventEnd(events[8],0,0,0,0);CHKERRQ(ierr); 280 } 281 ierr = PetscLogEventBegin(events[4],0,0,0,0);CHKERRQ(ierr); 282 ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr); 283 #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) 284 ierr = PetscLogGpuFlops(nip*(PetscLogDouble)(mass_w ? (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); 285 if (ctx->deviceType == LANDAU_CPU) PetscInfo(plex, "Warning: Landau selected CPU but no support for Kokkos using CPU\n"); 286 #else 287 ierr = PetscLogFlops(nip*(PetscLogDouble)(mass_w ? (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); 288 #endif 289 Kokkos::parallel_for("Landau_elements", Kokkos::TeamPolicy<>(numCells, team_size, num_sub_blocks).set_scratch_size(KOKKOS_SHARED_LEVEL, Kokkos::PerTeam(scr_bytes)), KOKKOS_LAMBDA (const team_member team) { 290 const PetscInt myelem = team.league_rank(); 291 g2_scr_t g2(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,Nf,Nq); // we don't use these for mass matrix 292 g3_scr_t g3(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,dim,Nf,Nq); 293 if (!d_mass_w.size()) { 294 g2_scr_t gg2(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,Nf,Nq); 295 g3_scr_t gg3(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,dim,Nf,Nq); 296 LandauIPData d_IPData; 297 // un pack IPData 298 d_IPData.w = &d_ipdata_raw[0]; 299 d_IPData.x = &d_ipdata_raw[1*nip]; 300 d_IPData.y = &d_ipdata_raw[2*nip]; 301 if (dim==2) d_IPData.z = NULL; 302 else d_IPData.z = &d_ipdata_raw[3*nip]; 303 // get g2[] & g3[] 304 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nq), [=] (int myQi) { 305 using Kokkos::parallel_reduce; 306 const PetscInt jpidx = myQi + myelem * Nq; 307 const PetscReal* const invJj = &d_invJ(jpidx*dim*dim); 308 const PetscReal vj[3] = {d_IPData.x[jpidx], d_IPData.y[jpidx], d_IPData.z ? d_IPData.z[jpidx] : 0}, wj = d_IPData.w[jpidx]; 309 landau_inner_red::TensorValueType gg_temp; // reduce on part of gg2 and g33 for IP jpidx 310 Kokkos::parallel_reduce(Kokkos::ThreadVectorRange (team, (int)nip), [=] (const int& ipidx, landau_inner_red::TensorValueType & ggg) { 311 const PetscReal wi = d_IPData.w[ipidx], x = d_IPData.x[ipidx], y = d_IPData.y[ipidx]; 312 PetscReal temp1[3] = {0, 0, 0}, temp2 = 0; 313 PetscInt fieldA,d2,d3; 314 #if LANDAU_DIM==2 315 PetscReal Ud[2][2], Uk[2][2]; 316 LandauTensor2D(vj, x, y, Ud, Uk, (ipidx==jpidx) ? 0. : 1.); 317 #else 318 PetscReal U[3][3], z = d_IPData.z[ipidx]; 319 LandauTensor3D(vj, x, y, z, U, (ipidx==jpidx) ? 0. : 1.); 320 #endif 321 for (fieldA = 0; fieldA < Nf; ++fieldA) { 322 temp1[0] += d_df(0,fieldA,ipidx)*d_beta[fieldA]*d_invMass[fieldA]; 323 temp1[1] += d_df(1,fieldA,ipidx)*d_beta[fieldA]*d_invMass[fieldA]; 324 #if LANDAU_DIM==3 325 temp1[2] += d_df(2,fieldA,ipidx)*d_beta[fieldA]*d_invMass[fieldA]; 326 #endif 327 temp2 += d_f(fieldA,ipidx)*d_beta[fieldA]; 328 } 329 temp1[0] *= wi; 330 temp1[1] *= wi; 331 #if LANDAU_DIM==3 332 temp1[2] *= wi; 333 #endif 334 temp2 *= wi; 335 #if LANDAU_DIM==2 336 for (d2 = 0; d2 < 2; d2++) { 337 for (d3 = 0; d3 < 2; ++d3) { 338 /* K = U * grad(f): g2=e: i,A */ 339 ggg.gg2[d2] += Uk[d2][d3]*temp1[d3]; 340 /* D = -U * (I \kron (fx)): g3=f: i,j,A */ 341 ggg.gg3[d2][d3] += Ud[d2][d3]*temp2; 342 } 343 } 344 #else 345 for (d2 = 0; d2 < 3; ++d2) { 346 for (d3 = 0; d3 < 3; ++d3) { 347 /* K = U * grad(f): g2 = e: i,A */ 348 ggg.gg2[d2] += U[d2][d3]*temp1[d3]; 349 /* D = -U * (I \kron (fx)): g3 = f: i,j,A */ 350 ggg.gg3[d2][d3] += U[d2][d3]*temp2; 351 } 352 } 353 #endif 354 }, Kokkos::Sum<landau_inner_red::TensorValueType>(gg_temp)); 355 // 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]); 356 // add alpha and put in gg2/3 357 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)Nf), [&] (const int& fieldA) { 358 PetscInt d2,d3; 359 for (d2 = 0; d2 < dim; d2++) { 360 gg2(d2,fieldA,myQi) = gg_temp.gg2[d2]*d_alpha[fieldA]; 361 //if (myelem==0 && fieldA==1) printf("\t\t:%d.%d) gg2[%d]=%e (+= %e)\n",myelem,myQi,d2,gg2(d2,fieldA,myQi),gg_temp.gg2[d2]*d_alpha[fieldA]); 362 //gg2[d2][myQi][fieldA] += gg_temp.gg2[d2]*d_alpha[fieldA]; 363 for (d3 = 0; d3 < dim; d3++) { 364 //gg3[d2][d3][myQi][fieldA] -= gg_temp.gg3[d2][d3]*d_alpha[fieldA]*s_invMass[fieldA]; 365 gg3(d2,d3,fieldA,myQi) = -gg_temp.gg3[d2][d3]*d_alpha[fieldA]*d_invMass[fieldA]; 366 //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)); 367 } 368 } 369 }); 370 /* add electric field term once per IP */ 371 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)Nf), [&] (const int& fieldA) { 372 //gg.gg2[fieldA][dim-1] += d_Eq_m[fieldA]; 373 gg2(dim-1,fieldA,myQi) += d_Eq_m[fieldA]; 374 }); 375 // Kokkos::single(Kokkos::PerThread(team), [&]() { 376 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)Nf), [=] (const int& fieldA) { 377 int d,d2,d3,dp; 378 //printf("%d %d %d gg2[][1]=%18.10e\n",myelem,myQi,fieldA,gg.gg2[fieldA][dim-1]); 379 /* Jacobian transform - g2, g3 - per thread (2D) */ 380 for (d = 0; d < dim; ++d) { 381 g2(d,fieldA,myQi) = 0; 382 for (d2 = 0; d2 < dim; ++d2) { 383 g2(d,fieldA,myQi) += invJj[d*dim+d2]*gg2(d2,fieldA,myQi); 384 //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]); 385 g3(d,d2,fieldA,myQi) = 0; 386 for (d3 = 0; d3 < dim; ++d3) { 387 for (dp = 0; dp < dim; ++dp) { 388 g3(d,d2,fieldA,myQi) += invJj[d*dim + d3]*gg3(d3,dp,fieldA,myQi)*invJj[d2*dim + dp]; 389 //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]); 390 } 391 } 392 g3(d,d2,fieldA,myQi) *= wj; 393 } 394 g2(d,fieldA,myQi) *= wj; 395 } 396 }); 397 }); // Nq 398 team.team_barrier(); 399 } // Jacobian 400 /* assemble */ 401 //Kokkos::single(Kokkos::PerTeam(team), [&]() { 402 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nb), [=] (int blk_i) { 403 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,(int)Nf), [=] (int fieldA) { 404 int blk_j,qj,d,d2; 405 const PetscInt i = fieldA*Nb + blk_i; /* Element matrix row */ 406 for (blk_j = 0; blk_j < Nb; ++blk_j) { 407 const PetscInt j = fieldA*Nb + blk_j; /* Element matrix column */ 408 const PetscInt fOff = i*totDim + j; 409 PetscScalar t = global_elem_mat_sz ? d_elem_mats(myelem,fOff) : 0; 410 for (qj = 0 ; qj < Nq ; qj++) { // look at others integration points 411 const PetscReal *BJq = &d_BB[qj*Nb], *DIq = &d_DD[qj*Nb*dim]; 412 if (!d_mass_w.size()) { 413 for (d = 0; d < dim; ++d) { 414 t += DIq[blk_i*dim+d]*g2(d,fieldA,qj)*BJq[blk_j]; 415 for (d2 = 0; d2 < dim; ++d2) { 416 t += DIq[blk_i*dim + d]*g3(d,d2,fieldA,qj)*DIq[blk_j*dim + d2]; 417 } 418 } 419 } else { 420 const PetscInt jpidx = qj + myelem * Nq; 421 t += BJq[blk_i] * d_mass_w(jpidx)*shift * BJq[blk_j]; 422 //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(jpidx),shift,BJq[blk_j]); 423 } 424 } 425 if (global_elem_mat_sz) d_elem_mats(myelem,fOff) = t; // can set this because local element matrix[fOff] 426 else { 427 PetscErrorCode ierr = 0; 428 PetscScalar vals[LANDAU_MAX_Q*LANDAU_MAX_Q],row_scale[LANDAU_MAX_Q],col_scale[LANDAU_MAX_Q]; 429 PetscInt q,idx,nr,nc,rows0[LANDAU_MAX_Q],cols0[LANDAU_MAX_Q],rows[LANDAU_MAX_Q],cols[LANDAU_MAX_Q]; 430 const LandauIdx *const Idxs = &d_maps->gIdx[myelem][fieldA][0]; 431 idx = Idxs[blk_i]; 432 if (idx >= 0) { 433 nr = 1; 434 rows0[0] = idx; 435 row_scale[0] = 1.; 436 } else { 437 idx = -idx - 1; 438 nr = d_maps->num_face; 439 for (q = 0; q < d_maps->num_face; q++) { 440 rows0[q] = d_maps->c_maps[idx][q].gid; 441 row_scale[q] = d_maps->c_maps[idx][q].scale; 442 } 443 } 444 idx = Idxs[blk_j]; 445 if (idx >= 0) { 446 nc = 1; 447 cols0[0] = idx; 448 col_scale[0] = 1.; 449 } else { 450 idx = -idx - 1; 451 nc = d_maps->num_face; 452 for (q = 0; q < d_maps->num_face; q++) { 453 cols0[q] = d_maps->c_maps[idx][q].gid; 454 col_scale[q] = d_maps->c_maps[idx][q].scale; 455 } 456 } 457 for (q = 0; q < nr; q++) rows[q] = rows0[q]; 458 for (q = 0; q < nc; q++) cols[q] = cols0[q]; 459 for (q = 0; q < nr; q++) { 460 for (d = 0; d < nc; d++) { 461 vals[q*nc + d] = row_scale[q]*col_scale[d]*t; 462 } 463 } 464 MatSetValuesDevice(d_mat,nr,rows,nc,cols,vals,ADD_VALUES,&ierr); 465 if (ierr) return; 466 } 467 } 468 }); 469 }); 470 }); 471 Kokkos::fence(); 472 ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr); 473 ierr = PetscLogEventEnd(events[4],0,0,0,0);CHKERRQ(ierr); 474 if (global_elem_mat_sz) { 475 Kokkos::View<PetscScalar**, Kokkos::LayoutRight>::HostMirror h_elem_mats = Kokkos::create_mirror_view(d_elem_mats); 476 ierr = PetscLogEventBegin(events[5],0,0,0,0);CHKERRQ(ierr); 477 Kokkos::deep_copy (h_elem_mats, d_elem_mats); 478 ierr = PetscLogEventEnd(events[5],0,0,0,0);CHKERRQ(ierr); 479 ierr = PetscLogEventBegin(events[6],0,0,0,0);CHKERRQ(ierr); 480 PetscInt ej; 481 for (ej = cStart ; ej < cEnd; ++ej) { 482 const PetscScalar *elMat = &h_elem_mats(ej-cStart,0); 483 ierr = DMPlexMatSetClosure(plex, section, globalSection, JacP, ej, elMat, ADD_VALUES);CHKERRQ(ierr); 484 if (ej==-1) { 485 int d,f; 486 PetscPrintf(PETSC_COMM_SELF,"Kokkos Element matrix %d/%d\n",1,(int)numCells); 487 for (d = 0; d < totDim; ++d){ 488 for (f = 0; f < totDim; ++f) PetscPrintf(PETSC_COMM_SELF," %12.5e", PetscRealPart(elMat[d*totDim + f])); 489 PetscPrintf(PETSC_COMM_SELF,"\n"); 490 } 491 } 492 } 493 ierr = PetscLogEventEnd(events[6],0,0,0,0);CHKERRQ(ierr); 494 } 495 } 496 PetscFunctionReturn(0); 497 } 498 } // extern "C" 499