1 #include <petscdevice.h> 2 #include <../src/ksp/pc/impls/vpbjacobi/vpbjacobi.h> 3 4 /* A class that manages helper arrays assisting parallel PCApply() with CUDA */ 5 struct PC_VPBJacobi_CUDA { 6 /* Cache the old sizes to check if we need realloc */ 7 PetscInt n; /* number of rows of the local matrix */ 8 PetscInt nblocks; /* number of point blocks */ 9 PetscInt nsize; /* sum of sizes of the point blocks */ 10 11 /* Helper arrays that are pre-computed on host and then copied to device. 12 bs: [nblocks+1], "csr" version of bsizes[], with bs[0] = 0, bs[nblocks] = n. 13 bs2: [nblocks+1], "csr" version of squares of bsizes[], with bs2[0] = 0, bs2[nblocks] = nsize. 14 matIdx: [n], row i of the local matrix belongs to the matIdx_d[i] block 15 */ 16 PetscInt *bs_h, *bs2_h, *matIdx_h; 17 PetscInt *bs_d, *bs2_d, *matIdx_d; 18 19 MatScalar *diag_d; /* [nsize], store inverse of the point blocks on device */ 20 21 PC_VPBJacobi_CUDA(PetscInt n, PetscInt nblocks, PetscInt nsize, const PetscInt *bsizes, MatScalar *diag_h) : n(n), nblocks(nblocks), nsize(nsize) { 22 /* malloc memory on host and device, and then update */ 23 PetscCallVoid(PetscMalloc3(nblocks + 1, &bs_h, nblocks + 1, &bs2_h, n, &matIdx_h)); 24 PetscCallCUDAVoid(cudaMalloc(&bs_d, sizeof(PetscInt) * (nblocks + 1))); 25 PetscCallCUDAVoid(cudaMalloc(&bs2_d, sizeof(PetscInt) * (nblocks + 1))); 26 PetscCallCUDAVoid(cudaMalloc(&matIdx_d, sizeof(PetscInt) * n)); 27 PetscCallCUDAVoid(cudaMalloc(&diag_d, sizeof(MatScalar) * nsize)); 28 PetscCallVoid(UpdateOffsetsOnDevice(bsizes, diag_h)); 29 } 30 31 PetscErrorCode UpdateOffsetsOnDevice(const PetscInt *bsizes, MatScalar *diag_h) { 32 PetscFunctionBegin; 33 PetscCall(ComputeOffsetsOnHost(bsizes)); 34 PetscCallCUDA(cudaMemcpy(bs_d, bs_h, sizeof(PetscInt) * (nblocks + 1), cudaMemcpyHostToDevice)); 35 PetscCallCUDA(cudaMemcpy(bs2_d, bs2_h, sizeof(PetscInt) * (nblocks + 1), cudaMemcpyHostToDevice)); 36 PetscCallCUDA(cudaMemcpy(matIdx_d, matIdx_h, sizeof(PetscInt) * n, cudaMemcpyHostToDevice)); 37 PetscCallCUDA(cudaMemcpy(diag_d, diag_h, sizeof(MatScalar) * nsize, cudaMemcpyHostToDevice)); 38 PetscCall(PetscLogCpuToGpu(sizeof(PetscInt) * (2 * nblocks + 2 + n) + sizeof(MatScalar) * nsize)); 39 PetscFunctionReturn(0); 40 } 41 42 ~PC_VPBJacobi_CUDA() { 43 PetscCallVoid(PetscFree3(bs_h, bs2_h, matIdx_h)); 44 PetscCallCUDAVoid(cudaFree(bs_d)); 45 PetscCallCUDAVoid(cudaFree(bs2_d)); 46 PetscCallCUDAVoid(cudaFree(matIdx_d)); 47 PetscCallCUDAVoid(cudaFree(diag_d)); 48 } 49 50 private: 51 PetscErrorCode ComputeOffsetsOnHost(const PetscInt *bsizes) { 52 PetscFunctionBegin; 53 bs_h[0] = bs2_h[0] = 0; 54 for (PetscInt i = 0; i < nblocks; i++) { 55 bs_h[i + 1] = bs_h[i] + bsizes[i]; 56 bs2_h[i + 1] = bs2_h[i] + bsizes[i] * bsizes[i]; 57 for (PetscInt j = 0; j < bsizes[i]; j++) matIdx_h[bs_h[i] + j] = i; 58 } 59 PetscFunctionReturn(0); 60 } 61 }; 62 63 /* Like cublasDgemvBatched() but with variable-sized blocks 64 65 Input Parameters: 66 + n - number of rows of the local matrix 67 . bs - [nblocks+1], prefix sum of bsizes[] 68 . bs2 - [nblocks+1], prefix sum of squares of bsizes[] 69 . matIdx - [n], store block/matrix index for each row 70 . A - blocks of the matrix back to back in column-major order 71 - x - the input vector 72 73 Output Parameter: 74 . y - the output vector 75 */ 76 __global__ static void MatMultBatched(PetscInt n, const PetscInt *bs, const PetscInt *bs2, const PetscInt *matIdx, const MatScalar *A, const PetscScalar *x, PetscScalar *y) { 77 const PetscInt gridSize = gridDim.x * blockDim.x; 78 PetscInt tid = blockIdx.x * blockDim.x + threadIdx.x; 79 PetscInt i, j, k, m; 80 81 /* One row per thread. The blocks/matrices are stored in column-major order */ 82 for (; tid < n; tid += gridSize) { 83 k = matIdx[tid]; /* k-th block */ 84 m = bs[k + 1] - bs[k]; /* block size of the k-th block */ 85 i = tid - bs[k]; /* i-th row of the block */ 86 A += bs2[k] + i; /* advance A to the first entry of i-th row */ 87 x += bs[k]; 88 y += bs[k]; 89 90 y[i] = 0.0; 91 for (j = 0; j < m; j++) { 92 y[i] += A[0] * x[j]; 93 A += m; 94 } 95 } 96 } 97 98 static PetscErrorCode PCApply_VPBJacobi_CUDA(PC pc, Vec x, Vec y) { 99 PC_VPBJacobi *jac = (PC_VPBJacobi *)pc->data; 100 PC_VPBJacobi_CUDA *pcuda = static_cast<PC_VPBJacobi_CUDA *>(jac->spptr); 101 const PetscScalar *xx; 102 PetscScalar *yy; 103 PetscInt n; 104 105 PetscFunctionBegin; 106 PetscCall(PetscLogGpuTimeBegin()); 107 if (PetscDefined(USE_DEBUG)) { 108 PetscBool isCuda; 109 PetscCall(PetscObjectTypeCompareAny((PetscObject)x, &isCuda, VECSEQCUDA, VECMPICUDA, "")); 110 if (isCuda) PetscCall(PetscObjectTypeCompareAny((PetscObject)y, &isCuda, VECSEQCUDA, VECMPICUDA, "")); 111 PetscCheck(isCuda, PETSC_COMM_SELF, PETSC_ERR_SUP, "PC: applying a CUDA pmat to non-cuda vectors"); 112 } 113 114 PetscCall(MatGetLocalSize(pc->pmat, &n, NULL)); 115 if (n) { 116 PetscInt gridSize = PetscMin((n + 255) / 256, 2147483647); /* <= 2^31-1 */ 117 PetscCall(VecCUDAGetArrayRead(x, &xx)); 118 PetscCall(VecCUDAGetArrayWrite(y, &yy)); 119 MatMultBatched<<<gridSize, 256>>>(n, pcuda->bs_d, pcuda->bs2_d, pcuda->matIdx_d, pcuda->diag_d, xx, yy); 120 PetscCallCUDA(cudaGetLastError()); 121 PetscCall(VecCUDARestoreArrayRead(x, &xx)); 122 PetscCall(VecCUDARestoreArrayWrite(y, &yy)); 123 } 124 PetscCall(PetscLogGpuFlops(pcuda->nsize * 2)); /* FMA on entries in all blocks */ 125 PetscCall(PetscLogGpuTimeEnd()); 126 PetscFunctionReturn(0); 127 } 128 129 static PetscErrorCode PCDestroy_VPBJacobi_CUDA(PC pc) { 130 PC_VPBJacobi *jac = (PC_VPBJacobi *)pc->data; 131 132 PetscFunctionBegin; 133 PetscCallCXX(delete static_cast<PC_VPBJacobi_CUDA *>(jac->spptr)); 134 PCDestroy_VPBJacobi(pc); 135 PetscFunctionReturn(0); 136 } 137 138 PETSC_INTERN PetscErrorCode PCSetUp_VPBJacobi_CUDA(PC pc) { 139 PC_VPBJacobi *jac = (PC_VPBJacobi *)pc->data; 140 PC_VPBJacobi_CUDA *pcuda = static_cast<PC_VPBJacobi_CUDA *>(jac->spptr); 141 PetscInt i, n, nblocks, nsize = 0; 142 const PetscInt *bsizes; 143 144 PetscFunctionBegin; 145 PetscCall(PCSetUp_VPBJacobi_Host(pc)); /* Compute the inverse on host now. Might worth doing it on device directly */ 146 PetscCall(MatGetVariableBlockSizes(pc->pmat, &nblocks, &bsizes)); 147 for (i = 0; i < nblocks; i++) nsize += bsizes[i] * bsizes[i]; 148 PetscCall(MatGetLocalSize(pc->pmat, &n, NULL)); 149 150 /* If one calls MatSetVariableBlockSizes() multiple times and sizes have been changed (is it allowed?), we delete the old and rebuild anyway */ 151 if (pcuda && (pcuda->n != n || pcuda->nblocks != nblocks || pcuda->nsize != nsize)) { 152 PetscCallCXX(delete pcuda); 153 pcuda = nullptr; 154 } 155 156 if (!pcuda) { /* allocate the struct along with the helper arrays from the scatch */ 157 PetscCallCXX(jac->spptr = new PC_VPBJacobi_CUDA(n, nblocks, nsize, bsizes, jac->diag)); 158 } else { /* update the value only */ 159 PetscCall(pcuda->UpdateOffsetsOnDevice(bsizes, jac->diag)); 160 } 161 162 pc->ops->apply = PCApply_VPBJacobi_CUDA; 163 pc->ops->destroy = PCDestroy_VPBJacobi_CUDA; 164 PetscFunctionReturn(0); 165 } 166