1 #if !defined(__CUSPARSEMATIMPL) 2 #define __CUSPARSEMATIMPL 3 4 #include <petscpkg_version.h> 5 #include <petsc/private/cudavecimpl.h> 6 7 #include <cusparse_v2.h> 8 9 #include <algorithm> 10 #include <vector> 11 12 #include <thrust/device_vector.h> 13 #include <thrust/device_ptr.h> 14 #include <thrust/device_malloc_allocator.h> 15 #include <thrust/transform.h> 16 #include <thrust/functional.h> 17 #include <thrust/sequence.h> 18 19 #if (CUSPARSE_VER_MAJOR > 10 || CUSPARSE_VER_MAJOR == 10 && CUSPARSE_VER_MINOR >= 2) /* According to cuda/10.1.168 on OLCF Summit */ 20 #define CHKERRCUSPARSE(stat) \ 21 do { \ 22 if (PetscUnlikely(stat)) { \ 23 const char *name = cusparseGetErrorName(stat); \ 24 const char *descr = cusparseGetErrorString(stat); \ 25 SETERRQ3(PETSC_COMM_SELF,PETSC_ERR_GPU,"cuSPARSE error %d (%s) : %s",(int)stat,name,descr); \ 26 } \ 27 } while (0) 28 #else 29 #define CHKERRCUSPARSE(stat) do {if (PetscUnlikely(stat)) SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_GPU,"cusparse error %d",(int)stat);} while (0) 30 #endif 31 32 #if defined(PETSC_USE_COMPLEX) 33 #if defined(PETSC_USE_REAL_SINGLE) 34 const cuComplex PETSC_CUSPARSE_ONE = {1.0f, 0.0f}; 35 const cuComplex PETSC_CUSPARSE_ZERO = {0.0f, 0.0f}; 36 #elif defined(PETSC_USE_REAL_DOUBLE) 37 const cuDoubleComplex PETSC_CUSPARSE_ONE = {1.0, 0.0}; 38 const cuDoubleComplex PETSC_CUSPARSE_ZERO = {0.0, 0.0}; 39 #endif 40 #else 41 const PetscScalar PETSC_CUSPARSE_ONE = 1.0; 42 const PetscScalar PETSC_CUSPARSE_ZERO = 0.0; 43 #endif 44 45 #if PETSC_PKG_CUDA_VERSION_GE(9,0,0) 46 #define cusparse_create_analysis_info cusparseCreateCsrsv2Info 47 #define cusparse_destroy_analysis_info cusparseDestroyCsrsv2Info 48 #if defined(PETSC_USE_COMPLEX) 49 #if defined(PETSC_USE_REAL_SINGLE) 50 #define cusparse_get_svbuffsize(a,b,c,d,e,f,g,h,i,j) cusparseCcsrsv2_bufferSize(a,b,c,d,e,(cuComplex*)(f),g,h,i,j) 51 #define cusparse_analysis(a,b,c,d,e,f,g,h,i,j,k) cusparseCcsrsv2_analysis(a,b,c,d,e,(const cuComplex*)(f),g,h,i,j,k) 52 #define cusparse_solve(a,b,c,d,e,f,g,h,i,j,k,l,m,n) cusparseCcsrsv2_solve(a,b,c,d,(const cuComplex*)(e),f,(const cuComplex*)(g),h,i,j,(const cuComplex*)(k),(cuComplex*)(l),m,n) 53 #elif defined(PETSC_USE_REAL_DOUBLE) 54 #define cusparse_get_svbuffsize(a,b,c,d,e,f,g,h,i,j) cusparseZcsrsv2_bufferSize(a,b,c,d,e,(cuDoubleComplex*)(f),g,h,i,j) 55 #define cusparse_analysis(a,b,c,d,e,f,g,h,i,j,k) cusparseZcsrsv2_analysis(a,b,c,d,e,(const cuDoubleComplex*)(f),g,h,i,j,k) 56 #define cusparse_solve(a,b,c,d,e,f,g,h,i,j,k,l,m,n) cusparseZcsrsv2_solve(a,b,c,d,(const cuDoubleComplex*)(e),f,(const cuDoubleComplex*)(g),h,i,j,(const cuDoubleComplex*)(k),(cuDoubleComplex*)(l),m,n) 57 #endif 58 #else /* not complex */ 59 #if defined(PETSC_USE_REAL_SINGLE) 60 #define cusparse_get_svbuffsize cusparseScsrsv2_bufferSize 61 #define cusparse_analysis cusparseScsrsv2_analysis 62 #define cusparse_solve cusparseScsrsv2_solve 63 #elif defined(PETSC_USE_REAL_DOUBLE) 64 #define cusparse_get_svbuffsize cusparseDcsrsv2_bufferSize 65 #define cusparse_analysis cusparseDcsrsv2_analysis 66 #define cusparse_solve cusparseDcsrsv2_solve 67 #endif 68 #endif 69 #else 70 #define cusparse_create_analysis_info cusparseCreateSolveAnalysisInfo 71 #define cusparse_destroy_analysis_info cusparseDestroySolveAnalysisInfo 72 #if defined(PETSC_USE_COMPLEX) 73 #if defined(PETSC_USE_REAL_SINGLE) 74 #define cusparse_solve(a,b,c,d,e,f,g,h,i,j,k) cusparseCcsrsv_solve((a),(b),(c),(cuComplex*)(d),(e),(cuComplex*)(f),(g),(h),(i),(cuComplex*)(j),(cuComplex*)(k)) 75 #define cusparse_analysis(a,b,c,d,e,f,g,h,i) cusparseCcsrsv_analysis((a),(b),(c),(d),(e),(cuComplex*)(f),(g),(h),(i)) 76 #elif defined(PETSC_USE_REAL_DOUBLE) 77 #define cusparse_solve(a,b,c,d,e,f,g,h,i,j,k) cusparseZcsrsv_solve((a),(b),(c),(cuDoubleComplex*)(d),(e),(cuDoubleComplex*)(f),(g),(h),(i),(cuDoubleComplex*)(j),(cuDoubleComplex*)(k)) 78 #define cusparse_analysis(a,b,c,d,e,f,g,h,i) cusparseZcsrsv_analysis((a),(b),(c),(d),(e),(cuDoubleComplex*)(f),(g),(h),(i)) 79 #endif 80 #else /* not complex */ 81 #if defined(PETSC_USE_REAL_SINGLE) 82 #define cusparse_solve cusparseScsrsv_solve 83 #define cusparse_analysis cusparseScsrsv_analysis 84 #elif defined(PETSC_USE_REAL_DOUBLE) 85 #define cusparse_solve cusparseDcsrsv_solve 86 #define cusparse_analysis cusparseDcsrsv_analysis 87 #endif 88 #endif 89 #endif 90 91 #if PETSC_PKG_CUDA_VERSION_GE(11,0,0) 92 #define cusparse_csr2csc cusparseCsr2cscEx2 93 #if defined(PETSC_USE_COMPLEX) 94 #if defined(PETSC_USE_REAL_SINGLE) 95 #define cusparse_scalartype CUDA_C_32F 96 #elif defined(PETSC_USE_REAL_DOUBLE) 97 #define cusparse_scalartype CUDA_C_64F 98 #endif 99 #else /* not complex */ 100 #if defined(PETSC_USE_REAL_SINGLE) 101 #define cusparse_scalartype CUDA_R_32F 102 #elif defined(PETSC_USE_REAL_DOUBLE) 103 #define cusparse_scalartype CUDA_R_64F 104 #endif 105 #endif 106 #else 107 #if defined(PETSC_USE_COMPLEX) 108 #if defined(PETSC_USE_REAL_SINGLE) 109 #define cusparse_csr_spmv(a,b,c,d,e,f,g,h,i,j,k,l,m) cusparseCcsrmv((a),(b),(c),(d),(e),(cuComplex*)(f),(g),(cuComplex*)(h),(i),(j),(cuComplex*)(k),(cuComplex*)(l),(cuComplex*)(m)) 110 #define cusparse_csr_spmm(a,b,c,d,e,f,g,h,i,j,k,l,m,n,o,p) cusparseCcsrmm((a),(b),(c),(d),(e),(f),(cuComplex*)(g),(h),(cuComplex*)(i),(j),(k),(cuComplex*)(l),(m),(cuComplex*)(n),(cuComplex*)(o),(p)) 111 #define cusparse_csr2csc(a,b,c,d,e,f,g,h,i,j,k,l) cusparseCcsr2csc((a),(b),(c),(d),(cuComplex*)(e),(f),(g),(cuComplex*)(h),(i),(j),(k),(l)) 112 #define cusparse_hyb_spmv(a,b,c,d,e,f,g,h) cusparseChybmv((a),(b),(cuComplex*)(c),(d),(e),(cuComplex*)(f),(cuComplex*)(g),(cuComplex*)(h)) 113 #define cusparse_csr2hyb(a,b,c,d,e,f,g,h,i,j) cusparseCcsr2hyb((a),(b),(c),(d),(cuComplex*)(e),(f),(g),(h),(i),(j)) 114 #define cusparse_hyb2csr(a,b,c,d,e,f) cusparseChyb2csr((a),(b),(c),(cuComplex*)(d),(e),(f)) 115 #elif defined(PETSC_USE_REAL_DOUBLE) 116 #define cusparse_csr_spmv(a,b,c,d,e,f,g,h,i,j,k,l,m) cusparseZcsrmv((a),(b),(c),(d),(e),(cuDoubleComplex*)(f),(g),(cuDoubleComplex*)(h),(i),(j),(cuDoubleComplex*)(k),(cuDoubleComplex*)(l),(cuDoubleComplex*)(m)) 117 #define cusparse_csr_spmm(a,b,c,d,e,f,g,h,i,j,k,l,m,n,o,p) cusparseZcsrmm((a),(b),(c),(d),(e),(f),(cuDoubleComplex*)(g),(h),(cuDoubleComplex*)(i),(j),(k),(cuDoubleComplex*)(l),(m),(cuDoubleComplex*)(n),(cuDoubleComplex*)(o),(p)) 118 #define cusparse_csr2csc(a,b,c,d,e,f,g,h,i,j,k,l) cusparseZcsr2csc((a),(b),(c),(d),(cuDoubleComplex*)(e),(f),(g),(cuDoubleComplex*)(h),(i),(j),(k),(l)) 119 #define cusparse_hyb_spmv(a,b,c,d,e,f,g,h) cusparseZhybmv((a),(b),(cuDoubleComplex*)(c),(d),(e),(cuDoubleComplex*)(f),(cuDoubleComplex*)(g),(cuDoubleComplex*)(h)) 120 #define cusparse_csr2hyb(a,b,c,d,e,f,g,h,i,j) cusparseZcsr2hyb((a),(b),(c),(d),(cuDoubleComplex*)(e),(f),(g),(h),(i),(j)) 121 #define cusparse_hyb2csr(a,b,c,d,e,f) cusparseZhyb2csr((a),(b),(c),(cuDoubleComplex*)(d),(e),(f)) 122 #endif 123 #else 124 #if defined(PETSC_USE_REAL_SINGLE) 125 #define cusparse_csr_spmv cusparseScsrmv 126 #define cusparse_csr_spmm cusparseScsrmm 127 #define cusparse_csr2csc cusparseScsr2csc 128 #define cusparse_hyb_spmv cusparseShybmv 129 #define cusparse_csr2hyb cusparseScsr2hyb 130 #define cusparse_hyb2csr cusparseShyb2csr 131 #elif defined(PETSC_USE_REAL_DOUBLE) 132 #define cusparse_csr_spmv cusparseDcsrmv 133 #define cusparse_csr_spmm cusparseDcsrmm 134 #define cusparse_csr2csc cusparseDcsr2csc 135 #define cusparse_hyb_spmv cusparseDhybmv 136 #define cusparse_csr2hyb cusparseDcsr2hyb 137 #define cusparse_hyb2csr cusparseDhyb2csr 138 #endif 139 #endif 140 #endif 141 142 #define THRUSTINTARRAY32 thrust::device_vector<int> 143 #define THRUSTINTARRAY thrust::device_vector<PetscInt> 144 #define THRUSTARRAY thrust::device_vector<PetscScalar> 145 146 /* A CSR matrix structure */ 147 struct CsrMatrix { 148 PetscInt num_rows; 149 PetscInt num_cols; 150 PetscInt num_entries; 151 THRUSTINTARRAY32 *row_offsets; 152 THRUSTINTARRAY32 *column_indices; 153 THRUSTARRAY *values; 154 }; 155 156 /* This is struct holding the relevant data needed to a MatSolve */ 157 struct Mat_SeqAIJCUSPARSETriFactorStruct { 158 /* Data needed for triangular solve */ 159 cusparseMatDescr_t descr; 160 cusparseOperation_t solveOp; 161 CsrMatrix *csrMat; 162 #if PETSC_PKG_CUDA_VERSION_GE(9,0,0) 163 csrsv2Info_t solveInfo; 164 #else 165 cusparseSolveAnalysisInfo_t solveInfo; 166 #endif 167 cusparseSolvePolicy_t solvePolicy; /* whether level information is generated and used */ 168 int solveBufferSize; 169 void *solveBuffer; 170 size_t csr2cscBufferSize; /* to transpose the triangular factor (only used for CUDA >= 11.0) */ 171 void *csr2cscBuffer; 172 }; 173 174 /* This is a larger struct holding all the triangular factors for a solve, transpose solve, and any indices used in a reordering */ 175 struct Mat_SeqAIJCUSPARSETriFactors { 176 Mat_SeqAIJCUSPARSETriFactorStruct *loTriFactorPtr; /* pointer for lower triangular (factored matrix) on GPU */ 177 Mat_SeqAIJCUSPARSETriFactorStruct *upTriFactorPtr; /* pointer for upper triangular (factored matrix) on GPU */ 178 Mat_SeqAIJCUSPARSETriFactorStruct *loTriFactorPtrTranspose; /* pointer for lower triangular (factored matrix) on GPU for the transpose (useful for BiCG) */ 179 Mat_SeqAIJCUSPARSETriFactorStruct *upTriFactorPtrTranspose; /* pointer for upper triangular (factored matrix) on GPU for the transpose (useful for BiCG)*/ 180 THRUSTINTARRAY *rpermIndices; /* indices used for any reordering */ 181 THRUSTINTARRAY *cpermIndices; /* indices used for any reordering */ 182 THRUSTARRAY *workVector; 183 cusparseHandle_t handle; /* a handle to the cusparse library */ 184 PetscInt nnz; /* number of nonzeros ... need this for accurate logging between ICC and ILU */ 185 }; 186 187 struct Mat_CusparseSpMV { 188 PetscBool initialized; /* Don't rely on spmvBuffer != NULL to test if the struct is initialized, */ 189 size_t spmvBufferSize; /* since I'm not sure if smvBuffer can be NULL even after cusparseSpMV_bufferSize() */ 190 void *spmvBuffer; 191 #if PETSC_PKG_CUDA_VERSION_GE(11,0,0) /* these are present from CUDA 10.1, but PETSc code makes use of them from CUDA 11 on */ 192 cusparseDnVecDescr_t vecXDescr,vecYDescr; /* descriptor for the dense vectors in y=op(A)x */ 193 #endif 194 }; 195 196 /* This is struct holding the relevant data needed to a MatMult */ 197 struct Mat_SeqAIJCUSPARSEMultStruct { 198 void *mat; /* opaque pointer to a matrix. This could be either a cusparseHybMat_t or a CsrMatrix */ 199 cusparseMatDescr_t descr; /* Data needed to describe the matrix for a multiply */ 200 THRUSTINTARRAY *cprowIndices; /* compressed row indices used in the parallel SpMV */ 201 PetscScalar *alpha_one; /* pointer to a device "scalar" storing the alpha parameter in the SpMV */ 202 PetscScalar *beta_zero; /* pointer to a device "scalar" storing the beta parameter in the SpMV as zero*/ 203 PetscScalar *beta_one; /* pointer to a device "scalar" storing the beta parameter in the SpMV as one */ 204 #if PETSC_PKG_CUDA_VERSION_GE(11,0,0) 205 cusparseSpMatDescr_t matDescr; /* descriptor for the matrix, used by SpMV and SpMM */ 206 Mat_CusparseSpMV cuSpMV[3]; /* different Mat_CusparseSpMV structs for non-transpose, transpose, conj-transpose */ 207 Mat_SeqAIJCUSPARSEMultStruct() : matDescr(NULL) { 208 for (int i=0; i<3; i++) cuSpMV[i].initialized = PETSC_FALSE; 209 } 210 #endif 211 }; 212 213 /* This is a larger struct holding all the matrices for a SpMV, and SpMV Tranpose */ 214 struct Mat_SeqAIJCUSPARSE { 215 Mat_SeqAIJCUSPARSEMultStruct *mat; /* pointer to the matrix on the GPU */ 216 Mat_SeqAIJCUSPARSEMultStruct *matTranspose; /* pointer to the matrix on the GPU (for the transpose ... useful for BiCG) */ 217 THRUSTARRAY *workVector; /* pointer to a workvector to which we can copy the relevant indices of a vector we want to multiply */ 218 THRUSTINTARRAY32 *rowoffsets_gpu; /* rowoffsets on GPU in non-compressed-row format. It is used to convert CSR to CSC */ 219 PetscInt nrows; /* number of rows of the matrix seen by GPU */ 220 MatCUSPARSEStorageFormat format; /* the storage format for the matrix on the device */ 221 cudaStream_t stream; /* a stream for the parallel SpMV ... this is not owned and should not be deleted */ 222 cusparseHandle_t handle; /* a handle to the cusparse library ... this may not be owned (if we're working in parallel i.e. multiGPUs) */ 223 PetscObjectState nonzerostate; /* track nonzero state to possibly recreate the GPU matrix */ 224 PetscBool transgen; /* whether or not to generate explicit transpose for MatMultTranspose operations */ 225 #if PETSC_PKG_CUDA_VERSION_GE(11,0,0) 226 size_t csr2cscBufferSize; /* stuff used to compute the matTranspose above */ 227 void *csr2cscBuffer; /* This is used as a C struct and is calloc'ed by PetscNewLog() */ 228 cusparseCsr2CscAlg_t csr2cscAlg; /* algorithms can be selected from command line options */ 229 cusparseSpMVAlg_t spmvAlg; 230 cusparseSpMMAlg_t spmmAlg; 231 #endif 232 PetscSplitCSRDataStructure *deviceMat; /* Matrix on device for, eg, assembly */ 233 THRUSTINTARRAY *cooPerm; 234 THRUSTINTARRAY *cooPerm_a; 235 THRUSTARRAY *cooPerm_v; 236 THRUSTARRAY *cooPerm_w; 237 }; 238 239 PETSC_INTERN PetscErrorCode MatCUSPARSECopyToGPU(Mat); 240 PETSC_INTERN PetscErrorCode MatCUSPARSESetStream(Mat, const cudaStream_t stream); 241 PETSC_INTERN PetscErrorCode MatCUSPARSESetHandle(Mat, const cusparseHandle_t handle); 242 PETSC_INTERN PetscErrorCode MatCUSPARSEClearHandle(Mat); 243 #endif 244