1: #include <petscdevice_cuda.h>
  2: #include <petsc/private/cupmblasinterface.hpp>
  3: #include <petsc/private/petsclegacycupmblas.h>
  4: #include <../src/ksp/pc/impls/pbjacobi/pbjacobi.h>

  6: #if PETSC_PKG_CUDA_VERSION_LT(11, 7, 0)
  7: __global__ static void MatMultBatched(PetscInt bs, PetscInt mbs, const PetscScalar *A, const PetscScalar *x, PetscScalar *y, PetscBool transpose)
  8: {
  9:   const PetscInt gridSize = gridDim.x * blockDim.x;
 10:   PetscInt       row      = blockIdx.x * blockDim.x + threadIdx.x;
 11:   const PetscInt bs2      = bs * bs;

 13:   /* One row per thread. The blocks are stored in column-major order */
 14:   for (; row < bs * mbs; row += gridSize) {
 15:     const PetscScalar *Ap, *xp;
 16:     PetscScalar       *yp;
 17:     PetscInt           i, j, k;

 19:     k  = row / bs;                               /* k-th block */
 20:     i  = row % bs;                               /* this thread deals with i-th row of the block */
 21:     Ap = &A[bs2 * k + i * (transpose ? bs : 1)]; /* Ap points to the first entry of i-th row */
 22:     xp = &x[bs * k];
 23:     yp = &y[bs * k];
 24:     /* multiply i-th row (column) with x */
 25:     yp[i] = 0.0;
 26:     for (j = 0; j < bs; j++) {
 27:       yp[i] += Ap[0] * xp[j];
 28:       Ap += (transpose ? 1 : bs); /* block is in column major order */
 29:     }
 30:   }
 31: }
 32: #endif

 34: static PetscErrorCode PCApplyOrTranspose_PBJacobi_CUDA(PC pc, Vec x, Vec y, cublasOperation_t op)
 35: {
 36:   const PetscScalar *xx;
 37:   PetscScalar       *yy;
 38:   cublasHandle_t     handle;
 39:   PC_PBJacobi       *jac = (PC_PBJacobi *)pc->data;
 40:   const PetscScalar *A   = (const PetscScalar *)jac->spptr;
 41:   const PetscInt     bs = jac->bs, mbs = jac->mbs;

 43:   PetscFunctionBegin;
 44:   PetscCall(VecCUDAGetArrayRead(x, &xx));
 45:   PetscCall(VecCUDAGetArrayWrite(y, &yy));
 46:   PetscCall(PetscCUBLASGetHandle(&handle));
 47:   {
 48:     const Petsc::device::cupm::impl::BlasInterface<Petsc::device::cupm::DeviceType::CUDA>::CUPMBlasPointerModeGuard guard(handle, CUBLAS_POINTER_MODE_HOST);

 50: #if PETSC_PKG_CUDA_VERSION_GE(11, 7, 0)
 51:     /* y = alpha op(A) x + beta y */
 52:     const PetscScalar alpha = 1.0, beta = 0.0;
 53:     PetscCallCUBLAS(cublasXgemvStridedBatched(handle, op, bs, bs, &alpha, A, bs, bs * bs, xx, 1, bs, &beta, yy, 1, bs, mbs));
 54: #else
 55:     PetscInt gridSize = PetscMin((bs * mbs + 255) / 256, 2147483647); /* <= 2^31-1 */
 56:     MatMultBatched<<<gridSize, 256>>>(bs, mbs, A, xx, yy, op == CUBLAS_OP_T ? PETSC_TRUE : PETSC_FALSE);
 57:     PetscCallCUDA(cudaGetLastError());
 58: #endif
 59:   }
 60:   PetscCall(VecCUDARestoreArrayRead(x, &xx));
 61:   PetscCall(VecCUDARestoreArrayWrite(y, &yy));
 62:   PetscCall(PetscLogGpuFlops(bs * bs * mbs * 2));
 63:   PetscFunctionReturn(PETSC_SUCCESS);
 64: }

 66: static PetscErrorCode PCApply_PBJacobi_CUDA(PC pc, Vec x, Vec y)
 67: {
 68:   PetscFunctionBegin;
 69:   PetscCall(PCApplyOrTranspose_PBJacobi_CUDA(pc, x, y, CUBLAS_OP_N)); // No transpose
 70:   PetscFunctionReturn(PETSC_SUCCESS);
 71: }

 73: static PetscErrorCode PCApplyTranspose_PBJacobi_CUDA(PC pc, Vec x, Vec y)
 74: {
 75:   PetscFunctionBegin;
 76:   PetscCall(PCApplyOrTranspose_PBJacobi_CUDA(pc, x, y, CUBLAS_OP_T)); // Transpose
 77:   PetscFunctionReturn(PETSC_SUCCESS);
 78: }

 80: static PetscErrorCode PCDestroy_PBJacobi_CUDA(PC pc)
 81: {
 82:   PC_PBJacobi *jac = (PC_PBJacobi *)pc->data;

 84:   PetscFunctionBegin;
 85:   PetscCallCUDA(cudaFree(jac->spptr));
 86:   PetscCall(PCDestroy_PBJacobi(pc));
 87:   PetscFunctionReturn(PETSC_SUCCESS);
 88: }

 90: PETSC_INTERN PetscErrorCode PCSetUp_PBJacobi_CUDA(PC pc, Mat diagVB)
 91: {
 92:   PC_PBJacobi *jac = (PC_PBJacobi *)pc->data;
 93:   size_t       size;

 95:   PetscFunctionBegin;
 96:   PetscCall(PCSetUp_PBJacobi_Host(pc, diagVB)); /* Compute the inverse on host now. Might worth doing it on device directly */
 97:   size = sizeof(PetscScalar) * jac->bs * jac->bs * jac->mbs;

 99:   /* PBJacobi_CUDA is simple so that we use jac->spptr as if it is diag_d */
100:   if (!jac->spptr) PetscCallCUDAVoid(cudaMalloc(&jac->spptr, size));
101:   PetscCallCUDAVoid(cudaMemcpy(jac->spptr, jac->diag, size, cudaMemcpyHostToDevice));
102:   PetscCall(PetscLogCpuToGpu(size));

104:   pc->ops->apply          = PCApply_PBJacobi_CUDA;
105:   pc->ops->applytranspose = PCApplyTranspose_PBJacobi_CUDA;
106:   pc->ops->destroy        = PCDestroy_PBJacobi_CUDA;
107:   PetscFunctionReturn(PETSC_SUCCESS);
108: }