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, int Nf > 18 struct array_type { 19 ScalarType gg2[Nf][LANDAU_DIM]; 20 ScalarType gg3[Nf][LANDAU_DIM][LANDAU_DIM]; 21 22 KOKKOS_INLINE_FUNCTION // Default constructor - Initialize to 0's 23 array_type() { 24 for (int i = 0; i < Nf; i++){ 25 for (int j = 0; j < LANDAU_DIM; j++){ 26 gg2[i][j] = 0; 27 for (int k = 0; k < LANDAU_DIM; k++){ 28 gg3[i][j][k] = 0; 29 } 30 } 31 } 32 } 33 KOKKOS_INLINE_FUNCTION // Copy Constructor 34 array_type(const array_type & rhs) { 35 for (int i = 0; i < Nf; i++){ 36 for (int j = 0; j < LANDAU_DIM; j++){ 37 gg2[i][j] = rhs.gg2[i][j]; 38 for (int k = 0; k < LANDAU_DIM; k++){ 39 gg3[i][j][k] = rhs.gg3[i][j][k]; 40 } 41 } 42 } 43 } 44 KOKKOS_INLINE_FUNCTION // add operator 45 array_type& operator += (const array_type& src) { 46 for (int i = 0; i < Nf; i++){ 47 for (int j = 0; j < LANDAU_DIM; j++){ 48 gg2[i][j] += src.gg2[i][j]; 49 for (int k = 0; k < LANDAU_DIM; k++){ 50 gg3[i][j][k] += src.gg3[i][j][k]; 51 } 52 } 53 } 54 return *this; 55 } 56 KOKKOS_INLINE_FUNCTION // volatile add operator 57 void operator += (const volatile array_type& src) volatile { 58 for (int i = 0; i < Nf; i++){ 59 for (int j = 0; j < LANDAU_DIM; j++){ 60 gg2[i][j] += src.gg2[i][j]; 61 for (int k = 0; k < LANDAU_DIM; k++){ 62 gg3[i][j][k] += src.gg3[i][j][k]; 63 } 64 } 65 } 66 } 67 }; 68 typedef array_type<PetscReal,LANDAU_MAX_SPECIES> ValueType; // used to simplify code below 69 } 70 71 namespace Kokkos { //reduction identity must be defined in Kokkos namespace 72 template<> 73 struct reduction_identity< landau_inner_red::ValueType > { 74 KOKKOS_FORCEINLINE_FUNCTION static landau_inner_red::ValueType sum() { 75 return landau_inner_red::ValueType(); 76 } 77 }; 78 } 79 80 extern "C" { 81 PetscErrorCode LandauKokkosJacobian(DM plex, const PetscInt Nq, PetscReal nu_alpha[], PetscReal nu_beta[], 82 PetscReal invMass[], PetscReal Eq_m[], PetscReal * const IPDataGlobal, 83 PetscReal wiGlobal[], PetscReal invJ[], const PetscInt num_sub_blocks, const PetscLogEvent events[], PetscBool quarter3DDomain, 84 Mat JacP) 85 { 86 PetscErrorCode ierr; 87 PetscInt *Nbf,Nb,cStart,cEnd,Nf,dim,numCells,totDim,nip; 88 PetscTabulation *Tf; 89 PetscDS prob; 90 PetscSection section, globalSection; 91 PetscLogDouble flops; 92 PetscReal *BB,*DD; 93 LandauCtx *ctx; 94 95 PetscFunctionBegin; 96 ierr = DMGetApplicationContext(plex, &ctx);CHKERRQ(ierr); 97 if (!ctx) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "no context"); 98 ierr = DMGetDimension(plex, &dim);CHKERRQ(ierr); 99 if (dim!=LANDAU_DIM) SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "LANDAU_DIM != dim"); 100 ierr = DMPlexGetHeightStratum(plex,0,&cStart,&cEnd);CHKERRQ(ierr); 101 numCells = cEnd - cStart; 102 nip = numCells*Nq; 103 ierr = DMGetDS(plex, &prob);CHKERRQ(ierr); 104 ierr = PetscDSGetNumFields(prob, &Nf);CHKERRQ(ierr); 105 ierr = PetscDSGetDimensions(prob, &Nbf);CHKERRQ(ierr); Nb = Nbf[0]; 106 if (Nq != Nb) SETERRQ2(PETSC_COMM_SELF, PETSC_ERR_PLIB, "Nq != Nb. %D %D",Nq,Nb); 107 ierr = PetscDSGetTotalDimension(prob, &totDim);CHKERRQ(ierr); 108 ierr = PetscDSGetTabulation(prob, &Tf);CHKERRQ(ierr); 109 BB = Tf[0]->T[0]; DD = Tf[0]->T[1]; 110 ierr = DMGetLocalSection(plex, §ion);CHKERRQ(ierr); 111 ierr = DMGetGlobalSection(plex, &globalSection);CHKERRQ(ierr); 112 flops = (PetscLogDouble)numCells*Nq*(5*dim*dim*Nf*Nf + 165); 113 if (!Kokkos::is_initialized()){ 114 int argc = 1; 115 char string[1][32], *argv[1] = {string[0]}; 116 ierr = PetscStrcpy(string[0],"landau");CHKERRQ(ierr); 117 Kokkos::initialize(argc, argv); 118 } 119 #if defined(KOKKOS_ENABLE_CXX11_DISPATCH_LAMBDA) 120 #else 121 SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_PLIB, "no KOKKOS_ENABLE_CXX11_DISPATCH_LAMBDA"); 122 #endif 123 { 124 using scr_mem_t = Kokkos::DefaultExecutionSpace::scratch_memory_space; 125 using g2_scr_t = Kokkos::View<PetscReal***, Kokkos::LayoutRight, scr_mem_t>; 126 using g3_scr_t = Kokkos::View<PetscReal****, Kokkos::LayoutRight, scr_mem_t>; 127 const int scr_bytes = g2_scr_t::shmem_size(Nf,Nq,dim) + g3_scr_t::shmem_size(Nf,Nq,dim,dim); 128 ierr = PetscLogEventBegin(events[3],0,0,0,0);CHKERRQ(ierr); 129 Kokkos::View<PetscScalar**, Kokkos::LayoutRight> d_elem_mats("element matrices", numCells, totDim*totDim); 130 Kokkos::View<PetscScalar**, Kokkos::LayoutRight>::HostMirror h_elem_mats = Kokkos::create_mirror_view(d_elem_mats); 131 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_alpha (nu_alpha, Nf); 132 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_alpha ("nu_alpha", Nf); 133 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_beta (nu_beta, Nf); 134 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_beta ("nu_beta", Nf); 135 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_invMass (invMass,Nf); 136 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_invMass ("invMass", Nf); 137 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_Eq_m (Eq_m,Nf); 138 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_Eq_m ("Eq_m", Nf); 139 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_BB (BB,Nq*Nb); 140 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_BB ("BB", Nq*Nb); 141 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_DD (DD,Nq*Nb*dim); 142 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_DD ("DD", Nq*Nb*dim); 143 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_wiGlobal (wiGlobal,nip); 144 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_wiGlobal ("wiGlobal", nip); 145 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_ipdata (IPDataGlobal,nip*(dim + Nf*(dim+1))); 146 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_ipdata ("ipdata", nip*(dim + Nf*(dim+1))); 147 const Kokkos::View<PetscReal*, Kokkos::LayoutLeft, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged> > h_invJ (invJ,nip*dim*dim); 148 Kokkos::View<PetscReal*, Kokkos::LayoutLeft> d_invJ ("invJ", nip*dim*dim); 149 150 Kokkos::deep_copy (d_alpha, h_alpha); 151 Kokkos::deep_copy (d_beta, h_beta); 152 Kokkos::deep_copy (d_invMass, h_invMass); 153 Kokkos::deep_copy (d_Eq_m, h_Eq_m); 154 Kokkos::deep_copy (d_BB, h_BB); 155 Kokkos::deep_copy (d_DD, h_DD); 156 Kokkos::deep_copy (d_wiGlobal, h_wiGlobal); 157 Kokkos::deep_copy (d_ipdata, h_ipdata); 158 Kokkos::deep_copy (d_invJ, h_invJ); 159 ierr = PetscLogEventEnd(events[3],0,0,0,0);CHKERRQ(ierr); 160 ierr = PetscLogEventBegin(events[4],0,0,0,0);CHKERRQ(ierr); 161 #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_VIENNACL) 162 ierr = PetscLogGpuFlops(flops*nip);CHKERRQ(ierr); 163 if (ctx->deviceType == LANDAU_CPU) PetscInfo(plex, "Warning: Landau selected CPU but no support for Kokkos using GPU\n"); 164 #else 165 ierr = PetscLogFlops(flops*nip);CHKERRQ(ierr); 166 #endif 167 #define KOKKOS_SHARED_LEVEL 1 168 //PetscInfo2(plex, "shared memory size: %D kB in level %d\n",Nf*Nq*dim*(dim+1)*sizeof(PetscReal)/1024,KOKKOS_SHARED_LEVEL); 169 int conc = Kokkos::DefaultExecutionSpace().concurrency(), team_size = conc > Nq ? Nq : 1; 170 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) { 171 const PetscInt myelem = team.league_rank(); 172 g2_scr_t g2(team.team_scratch(KOKKOS_SHARED_LEVEL),Nf,Nq,dim); 173 g3_scr_t g3(team.team_scratch(KOKKOS_SHARED_LEVEL),Nf,Nq,dim,dim); 174 // get g2[] & g3[] 175 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nq), [=] (int myQi) { 176 using Kokkos::parallel_reduce; 177 const PetscInt jpidx = myQi + myelem * Nq; 178 const PetscReal * const invJj = &d_invJ(jpidx*dim*dim); 179 const PetscInt ipdata_sz = (dim + Nf*(1+dim)); 180 const LandauPointData * const fplpt_j = (LandauPointData*)(&d_ipdata(jpidx*ipdata_sz)); 181 const PetscReal * const vj = fplpt_j->crd, wj = d_wiGlobal[jpidx]; 182 // reduce on g22 and g33 for IP jpidx 183 landau_inner_red::ValueType gg; 184 Kokkos::parallel_reduce(Kokkos::ThreadVectorRange (team, nip), [=] (const int& ipidx, landau_inner_red::ValueType & ggg) { 185 const LandauPointData * const fplpt = (LandauPointData*)(&d_ipdata(ipidx*ipdata_sz)); 186 const LandauFDF * const fdf = &fplpt->fdf[0]; 187 const PetscReal wi = d_wiGlobal[ipidx]; 188 PetscInt fieldA,fieldB,d2,d3; 189 #if LANDAU_DIM==2 190 PetscReal Ud[2][2], Uk[2][2]; 191 LandauTensor2D(vj, fplpt->crd[0], fplpt->crd[1], Ud, Uk, (ipidx==jpidx) ? 0. : 1.); 192 for (fieldA = 0; fieldA < Nf; ++fieldA) { 193 for (fieldB = 0; fieldB < Nf; ++fieldB) { 194 for (d2 = 0; d2 < 2; ++d2) { 195 for (d3 = 0; d3 < 2; ++d3) { 196 /* K = U * grad(f): g2=e: i,A */ 197 ggg.gg2[fieldA][d2] += d_alpha[fieldA]*d_beta[fieldB] * d_invMass[fieldB] * Uk[d2][d3] * fdf[fieldB].df[d3] * wi; 198 /* D = -U * (I \kron (fx)): g3=f: i,j,A */ 199 ggg.gg3[fieldA][d2][d3] -= d_alpha[fieldA]*d_beta[fieldB] * d_invMass[fieldA] * Ud[d2][d3] * fdf[fieldB].f * wi; 200 } 201 } 202 } 203 } 204 #else 205 PetscReal U[3][3]; 206 LandauTensor3D(vj, fplpt->crd[0], fplpt->crd[1], fplpt->crd[2], U, (ipidx==jpidx) ? 0. : 1.); 207 for (fieldA = 0; fieldA < Nf; ++fieldA) { 208 for (fieldB = 0; fieldB < Nf; ++fieldB) { 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[fieldA][d2] += d_alpha[fieldA]*d_beta[fieldB] * d_invMass[fieldB] * U[d2][d3] * fplpt->fdf[fieldB].df[d3] * wi; 213 /* D = -U * (I \kron (fx)): g3 = f: i,j,A */ 214 ggg.gg3[fieldA][d2][d3] -= d_alpha[fieldA]*d_beta[fieldB] * d_invMass[fieldA] * U[d2][d3] * fplpt->fdf[fieldB].f * wi; 215 } 216 } 217 } 218 } 219 #endif 220 }, Kokkos::Sum<landau_inner_red::ValueType>(gg)); 221 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, Nf), [&] (const int& fieldA) { 222 gg.gg2[fieldA][dim-1] += d_Eq_m[fieldA]; 223 }); 224 //kkos::single(Kokkos::PerThread(team), [&]() { 225 Kokkos::parallel_for(Kokkos::ThreadVectorRange (team, Nf), [=] (const int& fieldA) { 226 int d,d2,d3,dp; 227 //printf("%d %d %d gg2[][1]=%18.10e\n",myelem,myQi,fieldA,gg.gg2[fieldA][dim-1]); 228 /* Jacobian transform - g2, g3 - per thread (2D) */ 229 for (d = 0; d < dim; ++d) { 230 g2(fieldA,myQi,d) = 0; 231 for (d2 = 0; d2 < dim; ++d2) { 232 g2(fieldA,myQi,d) += invJj[d*dim+d2]*gg.gg2[fieldA][d2]; 233 //printf("\t\t%d %d %d %d %d g2[%d][%d][%d]=%g\n",myelem,myQi,fieldA,d,d2,fieldA,myQi,d,g2(fieldA,myQi,d)); 234 g3(fieldA,myQi,d,d2) = 0; 235 for (d3 = 0; d3 < dim; ++d3) { 236 for (dp = 0; dp < dim; ++dp) { 237 g3(fieldA,myQi,d,d2) += invJj[d*dim + d3]*gg.gg3[fieldA][d3][dp]*invJj[d2*dim + dp]; 238 //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]); 239 } 240 } 241 g3(fieldA,myQi,d,d2) *= wj; 242 } 243 g2(fieldA,myQi,d) *= wj; 244 } 245 }); 246 }); 247 team.team_barrier(); 248 /* assemble - on the diagonal (I,I) */ 249 //Kokkos::single(Kokkos::PerTeam(team), [&]() { 250 { // int fieldA,blk_i; 251 Kokkos::parallel_for(Kokkos::TeamThreadRange(team,0,Nb), [=] (int blk_i) { 252 Kokkos::parallel_for(Kokkos::ThreadVectorRange(team,0,Nf), [=] (int fieldA) { 253 //for (fieldA = 0; fieldA < Nf; ++fieldA) { 254 //for (blk_i = 0; blk_i < Nb; ++blk_i) { 255 int blk_j,qj,d,d2; 256 const PetscInt i = fieldA*Nb + blk_i; /* Element matrix row */ 257 for (blk_j = 0; blk_j < Nb; ++blk_j) { 258 const PetscInt j = fieldA*Nb + blk_j; /* Element matrix column */ 259 const PetscInt fOff = i*totDim + j; 260 for (qj = 0 ; qj < Nq ; qj++) { // look at others integration points 261 const PetscReal *BJq = &d_BB[qj*Nb], *DIq = &d_DD[qj*Nb*dim]; 262 for (d = 0; d < dim; ++d) { 263 d_elem_mats(myelem,fOff) += DIq[blk_i*dim+d]*g2(fieldA,qj,d)*BJq[blk_j]; 264 //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]); 265 for (d2 = 0; d2 < dim; ++d2) { 266 d_elem_mats(myelem,fOff) += DIq[blk_i*dim + d]*g3(fieldA,qj,d,d2)*DIq[blk_j*dim + d2]; 267 } 268 } 269 } 270 } 271 }); 272 }); 273 } 274 }); 275 Kokkos::fence(); 276 ierr = PetscLogEventEnd(events[4],0,0,0,0);CHKERRQ(ierr); 277 278 ierr = PetscLogEventBegin(events[5],0,0,0,0);CHKERRQ(ierr); 279 Kokkos::deep_copy (h_elem_mats, d_elem_mats); 280 ierr = PetscLogEventEnd(events[5],0,0,0,0);CHKERRQ(ierr); 281 282 ierr = PetscLogEventBegin(events[6],0,0,0,0);CHKERRQ(ierr); 283 #if defined(PETSC_HAVE_OPENMP) 284 { 285 PetscContainer container = NULL; 286 ierr = PetscObjectQuery((PetscObject)JacP,"coloring",(PetscObject*)&container);CHKERRQ(ierr); 287 if (!container) { 288 ierr = PetscLogEventBegin(events[8],0,0,0,0);CHKERRQ(ierr); 289 ierr = LandauCreateColoring(JacP, plex, &container);CHKERRQ(ierr); 290 ierr = PetscLogEventEnd(events[8],0,0,0,0);CHKERRQ(ierr); 291 } 292 ierr = LandauAssembleOpenMP(cStart, cEnd, totDim, plex, section, globalSection, JacP, &h_elem_mats(0,0), container);CHKERRQ(ierr); 293 } 294 #else 295 { 296 PetscInt ej; 297 for (ej = cStart ; ej < cEnd; ++ej) { 298 const PetscScalar *elMat = &h_elem_mats(ej-cStart,0); 299 ierr = DMPlexMatSetClosure(plex, section, globalSection, JacP, ej, elMat, ADD_VALUES);CHKERRQ(ierr); 300 if (ej==-1) { 301 int d,f; 302 printf("Kokkos Element matrix %d/%d\n",1,numCells); 303 for (d = 0; d < totDim; ++d){ 304 for (f = 0; f < totDim; ++f) printf(" %17.9e", PetscRealPart(elMat[d*totDim + f])); 305 printf("\n"); 306 } 307 } 308 } 309 } 310 #endif 311 ierr = PetscLogEventEnd(events[6],0,0,0,0);CHKERRQ(ierr); 312 } 313 PetscFunctionReturn(0); 314 } 315 } // extern "C" 316