1 #include <petscvec_kokkos.hpp> 2 #include <../src/vec/vec/impls/seq/kokkos/veckokkosimpl.hpp> 3 #include <petscdevice.h> 4 #include <../src/ksp/pc/impls/vpbjacobi/vpbjacobi.h> 5 6 /* A class that manages helper arrays assisting parallel PCApply() with Kokkos */ 7 struct PC_VPBJacobi_Kokkos { 8 /* Cache the old sizes to check if we need realloc */ 9 PetscInt n; /* number of rows of the local matrix */ 10 PetscInt nblocks; /* number of point blocks */ 11 PetscInt nsize; /* sum of sizes of the point blocks */ 12 13 /* Helper arrays that are pre-computed on host and then copied to device. 14 bs: [nblocks+1], "csr" version of bsizes[] 15 bs2: [nblocks+1], "csr" version of squares of bsizes[] 16 matIdx: [n], row i of the local matrix belongs to the matIdx_d[i] block 17 */ 18 PetscIntKokkosDualView bs_dual, bs2_dual, matIdx_dual; 19 PetscScalarKokkosDualView diag_dual; 20 21 PC_VPBJacobi_Kokkos(PetscInt n,PetscInt nblocks,PetscInt nsize,const PetscInt *bsizes,MatScalar *diag_ptr_h) 22 : n(n),nblocks(nblocks),nsize(nsize), 23 bs_dual("bs_dual",nblocks+1),bs2_dual("bs2_dual",nblocks+1),matIdx_dual("matIdx_dual",n) 24 { 25 PetscScalarKokkosViewHost diag_h(diag_ptr_h,nsize); 26 27 auto diag_d = Kokkos::create_mirror_view(DefaultMemorySpace(),diag_h); 28 diag_dual = PetscScalarKokkosDualView(diag_d,diag_h); 29 PetscCallVoid(UpdateOffsetsOnDevice(bsizes,diag_ptr_h)); 30 } 31 32 PetscErrorCode UpdateOffsetsOnDevice(const PetscInt *bsizes,MatScalar *diag_ptr_h) 33 { 34 PetscFunctionBegin; 35 PetscCheck(diag_dual.view_host().data() == diag_ptr_h, PETSC_COMM_SELF,PETSC_ERR_PLIB,"Host pointer has changed since last call"); 36 PetscCall(ComputeOffsetsOnHost(bsizes)); 37 38 PetscCallCXX(bs_dual.modify_host()); 39 PetscCallCXX(bs2_dual.modify_host()); 40 PetscCallCXX(matIdx_dual.modify_host()); 41 PetscCallCXX(diag_dual.modify_host()); 42 43 PetscCallCXX(bs_dual.sync_device()); 44 PetscCallCXX(bs2_dual.sync_device()); 45 PetscCallCXX(matIdx_dual.sync_device()); 46 PetscCallCXX(diag_dual.sync_device()); 47 PetscFunctionReturn(0); 48 } 49 50 private: 51 PetscErrorCode ComputeOffsetsOnHost(const PetscInt *bsizes) 52 { 53 PetscInt *bs_h = bs_dual.view_host().data(); 54 PetscInt *bs2_h = bs2_dual.view_host().data(); 55 PetscInt *matIdx_h = matIdx_dual.view_host().data(); 56 57 PetscFunctionBegin; 58 bs_h[0] = bs2_h[0] = 0; 59 for (PetscInt i=0; i<nblocks; i++) { 60 bs_h[i+1] = bs_h[i] + bsizes[i]; 61 bs2_h[i+1] = bs2_h[i] + bsizes[i]*bsizes[i]; 62 for (PetscInt j=0; j<bsizes[i]; j++) matIdx_h[bs_h[i]+j] = i; 63 } 64 PetscFunctionReturn(0); 65 } 66 }; 67 68 static PetscErrorCode PCApply_VPBJacobi_Kokkos(PC pc,Vec x,Vec y) 69 { 70 PC_VPBJacobi *jac = (PC_VPBJacobi*)pc->data; 71 PC_VPBJacobi_Kokkos *pckok = static_cast<PC_VPBJacobi_Kokkos*>(jac->spptr); 72 ConstPetscScalarKokkosView xv; 73 PetscScalarKokkosView yv; 74 PetscScalarKokkosView diag = pckok->diag_dual.view_device(); 75 PetscIntKokkosView bs = pckok->bs_dual.view_device(); 76 PetscIntKokkosView bs2 = pckok->bs2_dual.view_device(); 77 PetscIntKokkosView matIdx = pckok->matIdx_dual.view_device(); 78 79 PetscFunctionBegin; 80 VecErrorIfNotKokkos(x); 81 VecErrorIfNotKokkos(y); 82 PetscCall(VecGetKokkosView(x,&xv)); 83 PetscCall(VecGetKokkosViewWrite(y,&yv)); 84 PetscCallCXX(Kokkos::parallel_for("PCApply_VPBJacobi_Kokkos",pckok->n,KOKKOS_LAMBDA(PetscInt row) { 85 const PetscScalar *Ap,*xp; 86 PetscScalar *yp; 87 PetscInt i,j,k,m; 88 89 k = matIdx(row); /* k-th block/matrix */ 90 m = bs(k+1) - bs(k); /* block size of the k-th block */ 91 i = row - bs(k); /* i-th row of the block */ 92 Ap = &diag(bs2(k) + i); /* Ap points to the first entry of i-th row */ 93 xp = &xv(bs(k)); 94 yp = &yv(bs(k)); 95 96 yp[i] = 0.0; 97 for (j=0; j<m; j++) {yp[i] += Ap[0]*xp[j]; Ap += m;} 98 })); 99 PetscCall(VecRestoreKokkosView(x,&xv)); 100 PetscCall(VecRestoreKokkosViewWrite(y,&yv)); 101 PetscFunctionReturn(0); 102 } 103 104 static PetscErrorCode PCDestroy_VPBJacobi_Kokkos(PC pc) 105 { 106 PC_VPBJacobi *jac = (PC_VPBJacobi*)pc->data; 107 108 PetscFunctionBegin; 109 PetscCallCXX(delete static_cast<PC_VPBJacobi_Kokkos*>(jac->spptr)); 110 PetscCall(PCDestroy_VPBJacobi(pc)); 111 PetscFunctionReturn(0); 112 } 113 114 PETSC_INTERN PetscErrorCode PCSetUp_VPBJacobi_Kokkos(PC pc) 115 { 116 PC_VPBJacobi *jac = (PC_VPBJacobi*)pc->data; 117 PC_VPBJacobi_Kokkos *pckok = static_cast<PC_VPBJacobi_Kokkos*>(jac->spptr); 118 PetscInt i,n,nblocks,nsize = 0; 119 const PetscInt *bsizes; 120 121 PetscFunctionBegin; 122 PetscCall(PCSetUp_VPBJacobi_Host(pc)); /* Compute the inverse on host now. Might worth doing it on device directly */ 123 PetscCall(MatGetVariableBlockSizes(pc->pmat,&nblocks,&bsizes)); 124 for (i=0; i<nblocks; i++) nsize += bsizes[i]*bsizes[i]; 125 PetscCall(MatGetLocalSize(pc->pmat,&n,NULL)); 126 127 /* If one calls MatSetVariableBlockSizes() multiple times and sizes have been changed (is it allowed?), we delete the old and rebuild anyway */ 128 if (pckok && (pckok->n != n || pckok->nblocks != nblocks || pckok->nsize != nsize)) { 129 PetscCallCXX(delete pckok); 130 pckok = nullptr; 131 } 132 133 if (!pckok) { /* allocate the struct along with the helper arrays from the scatch */ 134 PetscCallCXX(jac->spptr = new PC_VPBJacobi_Kokkos(n,nblocks,nsize,bsizes,jac->diag)); 135 } else { /* update the value only */ 136 PetscCall(pckok->UpdateOffsetsOnDevice(bsizes,jac->diag)); 137 } 138 139 pc->ops->apply = PCApply_VPBJacobi_Kokkos; 140 pc->ops->destroy = PCDestroy_VPBJacobi_Kokkos; 141 PetscFunctionReturn(0); 142 } 143