1 /* 2 Implements the Kokkos kernel 3 */ 4 5 #define PETSC_SKIP_CXX_COMPLEX_FIX 6 #include <petscconf.h> 7 #include <petsc/private/dmpleximpl.h> /*I "petscdmplex.h" I*/ 8 #include <petsclandau.h> 9 #include <petscts.h> 10 #include <Kokkos_Core.hpp> 11 #include <cstdio> 12 typedef Kokkos::TeamPolicy<>::member_type team_member; 13 #define PETSC_DEVICE_FUNC_DECL KOKKOS_INLINE_FUNCTION 14 #include "../land_tensors.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 for (int j = 0; j < LANDAU_DIM; j++){ 43 gg2[j] += src.gg2[j]; 44 for (int k = 0; k < LANDAU_DIM; k++){ 45 gg3[j][k] += src.gg3[j][k]; 46 } 47 } 48 return *this; 49 } 50 KOKKOS_INLINE_FUNCTION // volatile add operator 51 void operator += (const volatile array_type& src) volatile { 52 for (int j = 0; j < LANDAU_DIM; j++){ 53 gg2[j] += src.gg2[j]; 54 for (int k = 0; k < LANDAU_DIM; k++){ 55 gg3[j][k] += src.gg3[j][k]; 56 } 57 } 58 } 59 }; 60 typedef array_type<PetscReal> TensorValueType; // used to simplify code below 61 } 62 63 namespace Kokkos { //reduction identity must be defined in Kokkos namespace 64 template<> 65 struct reduction_identity< landau_inner_red::TensorValueType > { 66 KOKKOS_FORCEINLINE_FUNCTION static landau_inner_red::TensorValueType sum() { 67 return landau_inner_red::TensorValueType(); 68 } 69 }; 70 } 71 72 extern "C" { 73 PetscErrorCode LandauKokkosJacobian(DM plex, const PetscInt Nq, PetscReal nu_alpha[], PetscReal nu_beta[], PetscReal invMass[], PetscReal Eq_m[], 74 const LandauIPData *const IPData, PetscReal invJ[], const PetscInt num_sub_blocks, const PetscLogEvent events[], Mat JacP) 75 { 76 PetscErrorCode ierr; 77 PetscInt *Nbf,Nb,cStart,cEnd,Nf,dim,numCells,totDim,ipdatasz; 78 PetscTabulation *Tf; 79 PetscDS prob; 80 PetscSection section, globalSection; 81 PetscLogDouble flops; 82 PetscReal *BB,*DD; 83 LandauCtx *ctx; 84 85 PetscFunctionBegin; 86 ierr = DMGetApplicationContext(plex, &ctx);CHKERRQ(ierr); 87 if (!ctx) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "no context"); 88 ierr = DMGetDimension(plex, &dim);CHKERRQ(ierr); 89 ierr = DMPlexGetHeightStratum(plex,0,&cStart,&cEnd);CHKERRQ(ierr); 90 numCells = cEnd - cStart; 91 ierr = DMGetDS(plex, &prob);CHKERRQ(ierr); 92 ierr = PetscDSGetNumFields(prob, &Nf);CHKERRQ(ierr); 93 ierr = PetscDSGetDimensions(prob, &Nbf);CHKERRQ(ierr); Nb = Nbf[0]; 94 if (Nq != Nb) SETERRQ2(PETSC_COMM_SELF, PETSC_ERR_PLIB, "Nq != Nb. %D %D",Nq,Nb); 95 if (LANDAU_DIM != dim) SETERRQ2(PETSC_COMM_WORLD, PETSC_ERR_PLIB, "dim %D != LANDAU_DIM %d",dim,LANDAU_DIM); 96 ierr = PetscDSGetTotalDimension(prob, &totDim);CHKERRQ(ierr); 97 ierr = PetscDSGetTabulation(prob, &Tf);CHKERRQ(ierr); 98 BB = Tf[0]->T[0]; DD = Tf[0]->T[1]; 99 ierr = DMGetLocalSection(plex, §ion);CHKERRQ(ierr); 100 ierr = DMGetGlobalSection(plex, &globalSection);CHKERRQ(ierr); 101 flops = (PetscLogDouble)numCells*Nq*(5*dim*dim*Nf*Nf + 165); 102 ipdatasz = LandauGetIPDataSize(IPData); 103 ierr = PetscKokkosInitializeCheck();CHKERRQ(ierr); 104 { 105 using scr_mem_t = Kokkos::DefaultExecutionSpace::scratch_memory_space; 106 using g2_scr_t = Kokkos::View<PetscReal***, Kokkos::LayoutRight, scr_mem_t>; 107 using g3_scr_t = Kokkos::View<PetscReal****, Kokkos::LayoutRight, scr_mem_t>; 108 const int scr_bytes = 2*(g2_scr_t::shmem_size(dim,Nf,Nq) + g3_scr_t::shmem_size(dim,dim,Nf,Nq)); 109 ierr = PetscLogEventBegin(events[3],0,0,0,0);CHKERRQ(ierr); 110 Kokkos::View<PetscScalar**, Kokkos::LayoutRight> d_elem_mats("element matrices", numCells, totDim*totDim); 111 Kokkos::View<PetscScalar**, Kokkos::LayoutRight>::HostMirror h_elem_mats = Kokkos::create_mirror_view(d_elem_mats); 112 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_alpha (nu_alpha, Nf); 113 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_alpha ("nu_alpha", Nf); 114 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_beta (nu_beta, Nf); 115 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_beta ("nu_beta", Nf); 116 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_invMass (invMass,Nf); 117 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_invMass ("invMass", Nf); 118 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_Eq_m (Eq_m,Nf); 119 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_Eq_m ("Eq_m", Nf); 120 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_BB (BB,Nq*Nb); 121 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_BB ("BB", Nq*Nb); 122 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_DD (DD,Nq*Nb*dim); 123 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_DD ("DD", Nq*Nb*dim); 124 const Kokkos::View<LandauIPReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_ipdata_raw (IPData->w_data,ipdatasz); 125 Kokkos::View<LandauIPReal*, Kokkos::LayoutLeft> d_ipdata_raw ("ipdata", ipdatasz); 126 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_invJ (invJ,IPData->nip_*dim*dim); 127 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_invJ ("invJ", IPData->nip_*dim*dim); 128 129 Kokkos::deep_copy (d_ipdata_raw, h_ipdata_raw); 130 Kokkos::deep_copy (d_alpha, h_alpha); 131 Kokkos::deep_copy (d_beta, h_beta); 132 Kokkos::deep_copy (d_invMass, h_invMass); 133 Kokkos::deep_copy (d_Eq_m, h_Eq_m); 134 Kokkos::deep_copy (d_BB, h_BB); 135 Kokkos::deep_copy (d_DD, h_DD); 136 Kokkos::deep_copy (d_invJ, h_invJ); 137 138 ierr = PetscLogEventEnd(events[3],0,0,0,0);CHKERRQ(ierr); 139 ierr = PetscLogEventBegin(events[4],0,0,0,0);CHKERRQ(ierr); 140 #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_VIENNACL) 141 ierr = PetscLogGpuFlops(flops*IPData->nip_);CHKERRQ(ierr); 142 if (ctx->deviceType == LANDAU_CPU) PetscInfo(plex, "Warning: Landau selected CPU but no support for Kokkos using GPU\n"); 143 #else 144 ierr = PetscLogFlops(flops*IPData->nip_);CHKERRQ(ierr); 145 #endif 146 #define KOKKOS_SHARED_LEVEL 1 147 // PetscInfo2(plex, "shared memory size: %d bytes in level %d\n",scr_bytes,KOKKOS_SHARED_LEVEL); 148 int conc = Kokkos::DefaultExecutionSpace().concurrency(), team_size = conc > Nq ? Nq : 1; 149 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) { 150 const PetscInt myelem = team.league_rank(); 151 g2_scr_t g2(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,Nf,Nq); 152 g3_scr_t g3(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,dim,Nf,Nq); 153 g2_scr_t gg2(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,Nf,Nq); 154 g3_scr_t gg3(team.team_scratch(KOKKOS_SHARED_LEVEL),dim,dim,Nf,Nq); 155 LandauIPData d_IPData; 156 // pack IPData 157 d_IPData.w_data = &d_ipdata_raw[0]; 158 d_IPData.x = &d_ipdata_raw[1*IPData->nip_]; 159 d_IPData.y = &d_ipdata_raw[2*IPData->nip_]; 160 d_IPData.z = &d_ipdata_raw[3*IPData->nip_]; 161 d_IPData.f = &d_ipdata_raw[IPData->nip_*((dim+1) + 0)]; 162 d_IPData.dfx = &d_ipdata_raw[IPData->nip_*((dim+1) + 1*Nf)]; 163 d_IPData.dfy = &d_ipdata_raw[IPData->nip_*((dim+1) + 2*Nf)]; 164 if (dim==2) d_IPData.z = d_IPData.dfz = NULL; 165 else d_IPData.dfz = &d_ipdata_raw[IPData->nip_*((dim+1) + 3*Nf)]; 166 167 // get g2[] & g3[] 168 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nq), [=] (int myQi) { 169 using Kokkos::parallel_reduce; 170 const PetscInt jpidx = myQi + myelem * Nq; 171 const PetscReal* const invJj = &d_invJ(jpidx*dim*dim); 172 const PetscReal vj[3] = {d_IPData.x[jpidx], d_IPData.y[jpidx], d_IPData.z ? d_IPData.z[jpidx] : 0}, wj = d_IPData.w_data[jpidx]; 173 landau_inner_red::TensorValueType gg_temp; // reduce on part of gg2 and g33 for IP jpidx 174 Kokkos::parallel_reduce(Kokkos::ThreadVectorRange (team, (int)IPData->nip_), [=] (const int& ipidx, landau_inner_red::TensorValueType & ggg) { 175 const PetscReal wi = d_IPData.w_data[ipidx], x = d_IPData.x[ipidx], y = d_IPData.y[ipidx]; 176 PetscReal temp1[3] = {0, 0, 0}, temp2 = 0; 177 PetscInt fieldA,d2,d3; 178 #if LANDAU_DIM==2 179 PetscReal Ud[2][2], Uk[2][2]; 180 LandauTensor2D(vj, x, y, Ud, Uk, (ipidx==jpidx) ? 0. : 1.); 181 #else 182 PetscReal U[3][3], z = d_IPData.z[ipidx]; 183 LandauTensor3D(vj, x, y, z, U, (ipidx==jpidx) ? 0. : 1.); 184 #endif 185 for (fieldA = 0; fieldA < Nf; ++fieldA) { 186 temp1[0] += d_IPData.dfx[ipidx + fieldA*IPData->nip_]*d_beta[fieldA]*d_invMass[fieldA]; 187 temp1[1] += d_IPData.dfy[ipidx + fieldA*IPData->nip_]*d_beta[fieldA]*d_invMass[fieldA]; 188 #if LANDAU_DIM==3 189 temp1[2] += d_IPData.dfz[ipidx + fieldA*IPData->nip_]*d_beta[fieldA]*d_invMass[fieldA]; 190 #endif 191 temp2 += d_IPData.f[ipidx + fieldA*IPData->nip_]*d_beta[fieldA]; 192 } 193 temp1[0] *= wi; 194 temp1[1] *= wi; 195 #if LANDAU_DIM==3 196 temp1[2] *= wi; 197 #endif 198 temp2 *= wi; 199 #if LANDAU_DIM==2 200 for (d2 = 0; d2 < 2; d2++) { 201 for (d3 = 0; d3 < 2; ++d3) { 202 /* K = U * grad(f): g2=e: i,A */ 203 ggg.gg2[d2] += Uk[d2][d3]*temp1[d3]; 204 /* D = -U * (I \kron (fx)): g3=f: i,j,A */ 205 ggg.gg3[d2][d3] += Ud[d2][d3]*temp2; 206 } 207 } 208 #else 209 for (d2 = 0; d2 < 3; ++d2) { 210 for (d3 = 0; d3 < 3; ++d3) { 211 /* K = U * grad(f): g2 = e: i,A */ 212 ggg.gg2[d2] += U[d2][d3]*temp1[d3]; 213 /* D = -U * (I \kron (fx)): g3 = f: i,j,A */ 214 ggg.gg3[d2][d3] += U[d2][d3]*temp2; 215 } 216 } 217 #endif 218 }, Kokkos::Sum<landau_inner_red::TensorValueType>(gg_temp)); 219 //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]); 220 // add alpha and put in gg2/3 221 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)Nf), [&] (const int& fieldA) { 222 PetscInt d2,d3; 223 for (d2 = 0; d2 < dim; d2++) { 224 gg2(d2,fieldA,myQi) = gg_temp.gg2[d2]*d_alpha[fieldA]; 225 //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]); 226 //gg2[d2][myQi][fieldA] += gg_temp.gg2[d2]*d_alpha[fieldA]; 227 for (d3 = 0; d3 < dim; d3++) { 228 //gg3[d2][d3][myQi][fieldA] -= gg_temp.gg3[d2][d3]*d_alpha[fieldA]*s_invMass[fieldA]; 229 gg3(d2,d3,fieldA,myQi) = -gg_temp.gg3[d2][d3]*d_alpha[fieldA]*d_invMass[fieldA]; 230 //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)); 231 } 232 } 233 }); 234 235 /* add electric field term once per IP */ 236 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)Nf), [&] (const int& fieldA) { 237 //gg.gg2[fieldA][dim-1] += d_Eq_m[fieldA]; 238 gg2(dim-1,fieldA,myQi) += d_Eq_m[fieldA]; 239 }); 240 // Kokkos::single(Kokkos::PerThread(team), [&]() { 241 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, (int)Nf), [=] (const int& fieldA) { 242 int d,d2,d3,dp; 243 //printf("%d %d %d gg2[][1]=%18.10e\n",myelem,myQi,fieldA,gg.gg2[fieldA][dim-1]); 244 /* Jacobian transform - g2, g3 - per thread (2D) */ 245 for (d = 0; d < dim; ++d) { 246 g2(d,fieldA,myQi) = 0; 247 for (d2 = 0; d2 < dim; ++d2) { 248 g2(d,fieldA,myQi) += invJj[d*dim+d2]*gg2(d2,fieldA,myQi); 249 //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]); 250 g3(d,d2,fieldA,myQi) = 0; 251 for (d3 = 0; d3 < dim; ++d3) { 252 for (dp = 0; dp < dim; ++dp) { 253 g3(d,d2,fieldA,myQi) += invJj[d*dim + d3]*gg3(d3,dp,fieldA,myQi)*invJj[d2*dim + dp]; 254 //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]); 255 } 256 } 257 g3(d,d2,fieldA,myQi) *= wj; 258 } 259 g2(d,fieldA,myQi) *= wj; 260 } 261 }); 262 }); // Nq 263 team.team_barrier(); 264 /* assemble - on the diagonal (I,I) */ 265 //Kokkos::single(Kokkos::PerTeam(team), [&]() { 266 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nb), [=] (int blk_i) { 267 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,(int)Nf), [=] (int fieldA) { 268 //for (fieldA = 0; fieldA < Nf; ++fieldA) { 269 //for (blk_i = 0; blk_i < Nb; ++blk_i) { 270 int blk_j,qj,d,d2; 271 const PetscInt i = fieldA*Nb + blk_i; /* Element matrix row */ 272 for (blk_j = 0; blk_j < Nb; ++blk_j) { 273 const PetscInt j = fieldA*Nb + blk_j; /* Element matrix column */ 274 const PetscInt fOff = i*totDim + j; 275 for (qj = 0 ; qj < Nq ; qj++) { // look at others integration points 276 const PetscReal *BJq = &d_BB[qj*Nb], *DIq = &d_DD[qj*Nb*dim]; 277 for (d = 0; d < dim; ++d) { 278 d_elem_mats(myelem,fOff) += DIq[blk_i*dim+d]*g2(d,fieldA,qj)*BJq[blk_j]; 279 //printf("\tmat[%d %d %d %d %d]=%g D[%d]=%g g2[%d][%d][%d]=%g B=%g\n",myelem,fOff,fieldA,qj,d,d_elem_mats(myelem,fOff),blk_i*dim+d,DIq[blk_i*dim+d],fieldA,qj,d,g2(fieldA,qj,d),BJq[blk_j]); 280 for (d2 = 0; d2 < dim; ++d2) { 281 d_elem_mats(myelem,fOff) += DIq[blk_i*dim + d]*g3(d,d2,fieldA,qj)*DIq[blk_j*dim + d2]; 282 } 283 } 284 } 285 } 286 }); 287 }); 288 }); 289 Kokkos::fence(); 290 ierr = PetscLogEventEnd(events[4],0,0,0,0);CHKERRQ(ierr); 291 ierr = PetscLogEventBegin(events[5],0,0,0,0);CHKERRQ(ierr); 292 Kokkos::deep_copy (h_elem_mats, d_elem_mats); 293 ierr = PetscLogEventEnd(events[5],0,0,0,0);CHKERRQ(ierr); 294 ierr = PetscLogEventBegin(events[6],0,0,0,0);CHKERRQ(ierr); 295 #if defined(PETSC_HAVE_OPENMP) 296 { 297 PetscContainer container = NULL; 298 ierr = PetscObjectQuery((PetscObject)JacP,"coloring",(PetscObject*)&container);CHKERRQ(ierr); 299 if (!container) { 300 ierr = PetscLogEventBegin(events[8],0,0,0,0);CHKERRQ(ierr); 301 ierr = LandauCreateColoring(JacP, plex, &container);CHKERRQ(ierr); 302 ierr = PetscLogEventEnd(events[8],0,0,0,0);CHKERRQ(ierr); 303 } 304 ierr = LandauAssembleOpenMP(cStart, cEnd, totDim, plex, section, globalSection, JacP, &h_elem_mats(0,0), container);CHKERRQ(ierr); 305 } 306 #else 307 { 308 PetscInt ej; 309 for (ej = cStart ; ej < cEnd; ++ej) { 310 const PetscScalar *elMat = &h_elem_mats(ej-cStart,0); 311 ierr = DMPlexMatSetClosure(plex, section, globalSection, JacP, ej, elMat, ADD_VALUES);CHKERRQ(ierr); 312 if (ej==-1) { 313 int d,f; 314 PetscPrintf(PETSC_COMM_SELF,"Kokkos Element matrix %d/%d\n",1,(int)numCells); 315 for (d = 0; d < totDim; ++d){ 316 for (f = 0; f < totDim; ++f) PetscPrintf(PETSC_COMM_SELF," %12.5e", PetscRealPart(elMat[d*totDim + f])); 317 PetscPrintf(PETSC_COMM_SELF,"\n"); 318 } 319 } 320 } 321 } 322 #endif 323 ierr = PetscLogEventEnd(events[6],0,0,0,0);CHKERRQ(ierr); 324 } 325 PetscFunctionReturn(0); 326 } 327 } // extern "C" 328