Actual source code: aijcusparse.cu
1: /*
2: Defines the basic matrix operations for the AIJ (compressed row)
3: matrix storage format using the CUSPARSE library,
4: */
5: #define PETSC_SKIP_IMMINTRIN_H_CUDAWORKAROUND 1
7: #include <petscconf.h>
8: #include <../src/mat/impls/aij/seq/aij.h>
9: #include <../src/mat/impls/sbaij/seq/sbaij.h>
10: #include <../src/vec/vec/impls/dvecimpl.h>
11: #include <petsc/private/vecimpl.h>
12: #undef VecType
13: #include <../src/mat/impls/aij/seq/seqcusparse/cusparsematimpl.h>
14: #include <../src/mat/impls/aij/seq/cupm/aijcupm.hpp>
15: #include <thrust/adjacent_difference.h>
16: #if PETSC_CPP_VERSION >= 14
17: #define PETSC_HAVE_THRUST_ASYNC 1
18: // thrust::for_each(thrust::cuda::par.on()) requires C++14
19: #endif
20: #include <thrust/iterator/constant_iterator.h>
21: #include <thrust/remove.h>
22: #include <thrust/sort.h>
23: #include <thrust/tuple.h>
24: #include <thrust/unique.h>
25: #include <thrust/gather.h>
26: #include <thrust/binary_search.h> // for thrust::lower_bound
27: #include <thrust/version.h>
28: #if THRUST_VERSION >= 200800
29: #include <cuda/std/functional>
30: #endif
31: #if CCCL_VERSION >= 3004000
32: #include <cuda/iterator>
33: #endif
35: const char *const MatCUSPARSEStorageFormats[] = {"CSR", "ELL", "HYB", "MatCUSPARSEStorageFormat", "MAT_CUSPARSE_", 0};
36: /*
37: The following are copied from cusparse.h in CUDA-11.0. In MatCUSPARSESpMVAlgorithms[] etc, we copy them in
38: 0-based integer value order, since we want to use PetscOptionsEnum() to parse user command line options for them.
39: */
40: const char *const MatCUSPARSESpMVAlgorithms[] = {"MV_ALG_DEFAULT", "COOMV_ALG", "CSRMV_ALG1", "CSRMV_ALG2", "cusparseSpMVAlg_t", "CUSPARSE_", 0};
41: const char *const MatCUSPARSESpMMAlgorithms[] = {"ALG_DEFAULT", "COO_ALG1", "COO_ALG2", "COO_ALG3", "CSR_ALG1", "COO_ALG4", "CSR_ALG2", "cusparseSpMMAlg_t", "CUSPARSE_SPMM_", 0};
42: const char *const MatCUSPARSECsr2CscAlgorithms[] = {"INVALID" /*cusparse does not have enum 0! We created one*/, "ALG1", "ALG2", "cusparseCsr2CscAlg_t", "CUSPARSE_CSR2CSC_", 0};
44: static PetscErrorCode MatICCFactorSymbolic_SeqAIJCUSPARSE(Mat, Mat, IS, const MatFactorInfo *);
45: static PetscErrorCode MatCholeskyFactorSymbolic_SeqAIJCUSPARSE(Mat, Mat, IS, const MatFactorInfo *);
46: static PetscErrorCode MatCholeskyFactorNumeric_SeqAIJCUSPARSE(Mat, Mat, const MatFactorInfo *);
47: static PetscErrorCode MatILUFactorSymbolic_SeqAIJCUSPARSE(Mat, Mat, IS, IS, const MatFactorInfo *);
48: static PetscErrorCode MatSetFromOptions_SeqAIJCUSPARSE(Mat, PetscOptionItems PetscOptionsObject);
49: static PetscErrorCode MatAXPY_SeqAIJCUSPARSE(Mat, PetscScalar, Mat, MatStructure);
50: static PetscErrorCode MatScale_SeqAIJCUSPARSE(Mat, PetscScalar);
51: static PetscErrorCode MatDiagonalScale_SeqAIJCUSPARSE(Mat, Vec, Vec);
52: static PetscErrorCode MatMult_SeqAIJCUSPARSE(Mat, Vec, Vec);
53: static PetscErrorCode MatMultAdd_SeqAIJCUSPARSE(Mat, Vec, Vec, Vec);
54: static PetscErrorCode MatMultTranspose_SeqAIJCUSPARSE(Mat, Vec, Vec);
55: static PetscErrorCode MatMultTransposeAdd_SeqAIJCUSPARSE(Mat, Vec, Vec, Vec);
56: static PetscErrorCode MatMultHermitianTranspose_SeqAIJCUSPARSE(Mat, Vec, Vec);
57: static PetscErrorCode MatMultHermitianTransposeAdd_SeqAIJCUSPARSE(Mat, Vec, Vec, Vec);
58: static PetscErrorCode MatMultAddKernel_SeqAIJCUSPARSE(Mat, Vec, Vec, Vec, PetscBool, PetscBool);
60: static PetscErrorCode CsrMatrix_Destroy(CsrMatrix **);
61: static PetscErrorCode MatSeqAIJCUSPARSEMultStruct_Destroy(Mat_SeqAIJCUSPARSEMultStruct **, MatCUSPARSEStorageFormat);
62: static PetscErrorCode MatSeqAIJCUSPARSETriFactors_Destroy(Mat_SeqAIJCUSPARSETriFactors **);
63: static PetscErrorCode MatSeqAIJCUSPARSE_Destroy(Mat);
65: static PetscErrorCode MatSeqAIJCUSPARSECopyFromGPU(Mat);
66: static PetscErrorCode MatSeqAIJCUSPARSEInvalidateTranspose(Mat, PetscBool);
68: static PetscErrorCode MatSeqAIJCopySubArray_SeqAIJCUSPARSE(Mat, PetscInt, const PetscInt[], PetscScalar[]);
69: static PetscErrorCode MatSetPreallocationCOO_SeqAIJCUSPARSE(Mat, PetscCount, PetscInt[], PetscInt[]);
70: static PetscErrorCode MatSetValuesCOO_SeqAIJCUSPARSE(Mat, const PetscScalar[], InsertMode);
71: PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJCUSPARSE(Mat, MatType, MatReuse, Mat *);
73: // cusparseCreateCsr() separates types for row offsets and column indices in prototype, but requires them to have the same type at runtime!
74: const cusparseIndexType_t csrRowOffsetsType = PetscDefined(USE_64BIT_INDICES) ? CUSPARSE_INDEX_64I : CUSPARSE_INDEX_32I;
75: const cusparseIndexType_t csrColIndType = PetscDefined(USE_64BIT_INDICES) ? CUSPARSE_INDEX_64I : CUSPARSE_INDEX_32I;
77: using Csr2coo = Petsc::mat::aij::cupm::impl::Csr2coo;
78: using PetscIntToCInt = Petsc::mat::aij::cupm::impl::PetscIntToCInt;
80: PETSC_INTERN PetscErrorCode MatCUSPARSESetFormat_SeqAIJCUSPARSE(Mat A, MatCUSPARSEFormatOperation op, MatCUSPARSEStorageFormat format)
81: {
82: Mat_SeqAIJCUSPARSE *cusparsestruct = (Mat_SeqAIJCUSPARSE *)A->spptr;
84: PetscFunctionBegin;
85: switch (op) {
86: case MAT_CUSPARSE_MULT:
87: cusparsestruct->format = format;
88: break;
89: case MAT_CUSPARSE_ALL:
90: cusparsestruct->format = format;
91: break;
92: default:
93: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "unsupported operation %d for MatCUSPARSEFormatOperation. MAT_CUSPARSE_MULT and MAT_CUSPARSE_ALL are currently supported.", op);
94: }
95: PetscFunctionReturn(PETSC_SUCCESS);
96: }
98: /*@
99: MatCUSPARSESetFormat - Sets the storage format of `MATSEQCUSPARSE` matrices for a particular
100: operation. Only the `MatMult()` operation can use different GPU storage formats
102: Not Collective
104: Input Parameters:
105: + A - Matrix of type `MATSEQAIJCUSPARSE`
106: . op - `MatCUSPARSEFormatOperation`. `MATSEQAIJCUSPARSE` matrices support `MAT_CUSPARSE_MULT` and `MAT_CUSPARSE_ALL`.
107: `MATMPIAIJCUSPARSE` matrices support `MAT_CUSPARSE_MULT_DIAG`,`MAT_CUSPARSE_MULT_OFFDIAG`, and `MAT_CUSPARSE_ALL`.
108: - format - `MatCUSPARSEStorageFormat` (one of `MAT_CUSPARSE_CSR`, `MAT_CUSPARSE_ELL`, `MAT_CUSPARSE_HYB`.)
110: Level: intermediate
112: .seealso: [](ch_matrices), `Mat`, `MATSEQAIJCUSPARSE`, `MatCUSPARSEStorageFormat`, `MatCUSPARSEFormatOperation`
113: @*/
114: PetscErrorCode MatCUSPARSESetFormat(Mat A, MatCUSPARSEFormatOperation op, MatCUSPARSEStorageFormat format)
115: {
116: PetscFunctionBegin;
118: PetscTryMethod(A, "MatCUSPARSESetFormat_C", (Mat, MatCUSPARSEFormatOperation, MatCUSPARSEStorageFormat), (A, op, format));
119: PetscFunctionReturn(PETSC_SUCCESS);
120: }
122: PETSC_INTERN PetscErrorCode MatCUSPARSESetUseCPUSolve_SeqAIJCUSPARSE(Mat A, PetscBool use_cpu)
123: {
124: Mat_SeqAIJCUSPARSE *cusparsestruct = (Mat_SeqAIJCUSPARSE *)A->spptr;
126: PetscFunctionBegin;
127: cusparsestruct->use_cpu_solve = use_cpu;
128: PetscFunctionReturn(PETSC_SUCCESS);
129: }
131: /*@
132: MatCUSPARSESetUseCPUSolve - Sets to use CPU `MatSolve()`.
134: Input Parameters:
135: + A - Matrix of type `MATSEQAIJCUSPARSE`
136: - use_cpu - set flag for using the built-in CPU `MatSolve()`
138: Level: intermediate
140: Note:
141: The NVIDIA cuSPARSE LU solver currently computes the factors with the built-in CPU method
142: and moves the factors to the GPU for the solve. We have observed better performance keeping the data on the CPU and performing the solve there.
143: This method to specify if the solve is done on the CPU or GPU (GPU is the default).
145: .seealso: [](ch_matrices), `Mat`, `MatSolve()`, `MATSEQAIJCUSPARSE`, `MatCUSPARSEStorageFormat`, `MatCUSPARSEFormatOperation`
146: @*/
147: PetscErrorCode MatCUSPARSESetUseCPUSolve(Mat A, PetscBool use_cpu)
148: {
149: PetscFunctionBegin;
151: PetscTryMethod(A, "MatCUSPARSESetUseCPUSolve_C", (Mat, PetscBool), (A, use_cpu));
152: PetscFunctionReturn(PETSC_SUCCESS);
153: }
155: static PetscErrorCode MatSetOption_SeqAIJCUSPARSE(Mat A, MatOption op, PetscBool flg)
156: {
157: PetscFunctionBegin;
158: switch (op) {
159: case MAT_FORM_EXPLICIT_TRANSPOSE:
160: /* need to destroy the transpose matrix if present to prevent from logic errors if flg is set to true later */
161: if (A->form_explicit_transpose && !flg) PetscCall(MatSeqAIJCUSPARSEInvalidateTranspose(A, PETSC_TRUE));
162: A->form_explicit_transpose = flg;
163: break;
164: default:
165: PetscCall(MatSetOption_SeqAIJ(A, op, flg));
166: break;
167: }
168: PetscFunctionReturn(PETSC_SUCCESS);
169: }
171: static PetscErrorCode MatSetFromOptions_SeqAIJCUSPARSE(Mat A, PetscOptionItems PetscOptionsObject)
172: {
173: MatCUSPARSEStorageFormat format;
174: PetscBool flg;
175: Mat_SeqAIJCUSPARSE *cusparsestruct = (Mat_SeqAIJCUSPARSE *)A->spptr;
177: PetscFunctionBegin;
178: PetscOptionsHeadBegin(PetscOptionsObject, "SeqAIJCUSPARSE options");
179: if (A->factortype == MAT_FACTOR_NONE) {
180: PetscCall(PetscOptionsEnum("-mat_cusparse_mult_storage_format", "sets storage format of (seq)aijcusparse gpu matrices for SpMV", "MatCUSPARSESetFormat", MatCUSPARSEStorageFormats, (PetscEnum)cusparsestruct->format, (PetscEnum *)&format, &flg));
181: if (flg) PetscCall(MatCUSPARSESetFormat(A, MAT_CUSPARSE_MULT, format));
183: PetscCall(PetscOptionsEnum("-mat_cusparse_storage_format", "sets storage format of (seq)aijcusparse gpu matrices for SpMV and TriSolve", "MatCUSPARSESetFormat", MatCUSPARSEStorageFormats, (PetscEnum)cusparsestruct->format, (PetscEnum *)&format, &flg));
184: if (flg) PetscCall(MatCUSPARSESetFormat(A, MAT_CUSPARSE_ALL, format));
185: PetscCall(PetscOptionsBool("-mat_cusparse_use_cpu_solve", "Use CPU (I)LU solve", "MatCUSPARSESetUseCPUSolve", cusparsestruct->use_cpu_solve, &cusparsestruct->use_cpu_solve, &flg));
186: if (flg) PetscCall(MatCUSPARSESetUseCPUSolve(A, cusparsestruct->use_cpu_solve));
187: PetscCall(PetscOptionsEnum("-mat_cusparse_spmv_alg", "sets cuSPARSE algorithm used in sparse-mat dense-vector multiplication (SpMV)", "cusparseSpMVAlg_t", MatCUSPARSESpMVAlgorithms, (PetscEnum)cusparsestruct->spmvAlg, (PetscEnum *)&cusparsestruct->spmvAlg, &flg));
188: /* If user did use this option, check its consistency with cuSPARSE, since PetscOptionsEnum() sets enum values based on their position in MatCUSPARSESpMVAlgorithms[] */
189: PetscCheck(!flg || CUSPARSE_SPMV_CSR_ALG1 == 2, PETSC_COMM_SELF, PETSC_ERR_SUP, "cuSPARSE enum cusparseSpMVAlg_t has been changed but PETSc has not been updated accordingly");
190: PetscCall(PetscOptionsEnum("-mat_cusparse_spmm_alg", "sets cuSPARSE algorithm used in sparse-mat dense-mat multiplication (SpMM)", "cusparseSpMMAlg_t", MatCUSPARSESpMMAlgorithms, (PetscEnum)cusparsestruct->spmmAlg, (PetscEnum *)&cusparsestruct->spmmAlg, &flg));
191: PetscCheck(!flg || CUSPARSE_SPMM_CSR_ALG1 == 4, PETSC_COMM_SELF, PETSC_ERR_SUP, "cuSPARSE enum cusparseSpMMAlg_t has been changed but PETSc has not been updated accordingly");
192: PetscCall(
193: PetscOptionsEnum("-mat_cusparse_csr2csc_alg", "sets cuSPARSE algorithm used in converting CSR matrices to CSC matrices", "cusparseCsr2CscAlg_t", MatCUSPARSECsr2CscAlgorithms, (PetscEnum)cusparsestruct->csr2cscAlg, (PetscEnum *)&cusparsestruct->csr2cscAlg, &flg));
194: PetscCheck(!flg || CUSPARSE_CSR2CSC_ALG1 == 1, PETSC_COMM_SELF, PETSC_ERR_SUP, "cuSPARSE enum cusparseCsr2CscAlg_t has been changed but PETSc has not been updated accordingly");
195: }
196: PetscOptionsHeadEnd();
197: PetscFunctionReturn(PETSC_SUCCESS);
198: }
200: static PetscErrorCode MatSeqAIJCUSPARSEBuildFactoredMatrix_LU(Mat A)
201: {
202: Mat_SeqAIJ *a = static_cast<Mat_SeqAIJ *>(A->data);
203: PetscInt m = A->rmap->n;
204: Mat_SeqAIJCUSPARSETriFactors *fs = static_cast<Mat_SeqAIJCUSPARSETriFactors *>(A->spptr);
205: const PetscInt *Ai = a->i, *Aj = a->j, *adiag;
206: const MatScalar *Aa = a->a;
207: PetscInt *Mi, *Mj, Mnz;
208: PetscScalar *Ma;
210: PetscFunctionBegin;
211: PetscCall(MatGetDiagonalMarkers_SeqAIJ(A, &adiag, NULL));
212: if (A->offloadmask == PETSC_OFFLOAD_CPU) { // A's latest factors are on CPU
213: if (!fs->csrRowPtr) { // Is this the first time we are doing setup? Use csrRowPtr since it is not null even when m=0
214: // Re-arrange the (skewed) factored matrix and put the result into M, a regular csr matrix on host
215: Mnz = (Ai[m] - Ai[0]) + (adiag[0] - adiag[m]); // Lnz (without the unit diagonal) + Unz (with the non-unit diagonal)
216: PetscCall(PetscMalloc1(m + 1, &Mi));
217: PetscCall(PetscMalloc1(Mnz, &Mj)); // Mj is temp
218: PetscCall(PetscMalloc1(Mnz, &Ma));
219: Mi[0] = 0;
220: for (PetscInt i = 0; i < m; i++) {
221: PetscInt llen = Ai[i + 1] - Ai[i];
222: PetscInt ulen = adiag[i] - adiag[i + 1];
223: PetscCall(PetscArraycpy(Mj + Mi[i], Aj + Ai[i], llen)); // entries of L
224: Mj[Mi[i] + llen] = i; // diagonal entry
225: PetscCall(PetscArraycpy(Mj + Mi[i] + llen + 1, Aj + adiag[i + 1] + 1, ulen - 1)); // entries of U on the right of the diagonal
226: Mi[i + 1] = Mi[i] + llen + ulen;
227: }
228: // Copy M (L,U) from host to device
229: PetscCallCUDA(cudaMalloc(&fs->csrRowPtr, sizeof(*fs->csrRowPtr) * (m + 1)));
230: PetscCallCUDA(cudaMalloc(&fs->csrColIdx, sizeof(*fs->csrColIdx) * Mnz));
231: PetscCallCUDA(cudaMalloc(&fs->csrVal, sizeof(*fs->csrVal) * Mnz));
232: PetscCallCUDA(cudaMemcpy(fs->csrRowPtr, Mi, sizeof(*fs->csrRowPtr) * (m + 1), cudaMemcpyHostToDevice));
233: PetscCallCUDA(cudaMemcpy(fs->csrColIdx, Mj, sizeof(*fs->csrColIdx) * Mnz, cudaMemcpyHostToDevice));
235: // Create descriptors for L, U. See https://docs.nvidia.com/cuda/cusparse/index.html#cusparseDiagType_t
236: // cusparseDiagType_t: This type indicates if the matrix diagonal entries are unity. The diagonal elements are always
237: // assumed to be present, but if CUSPARSE_DIAG_TYPE_UNIT is passed to an API routine, then the routine assumes that
238: // all diagonal entries are unity and will not read or modify those entries. Note that in this case the routine
239: // assumes the diagonal entries are equal to one, regardless of what those entries are actually set to in memory.
240: cusparseFillMode_t fillMode = CUSPARSE_FILL_MODE_LOWER;
241: cusparseDiagType_t diagType = CUSPARSE_DIAG_TYPE_UNIT;
243: PetscCallCUSPARSE(cusparseCreateCsr(&fs->spMatDescr_L, m, m, Mnz, fs->csrRowPtr, fs->csrColIdx, fs->csrVal, csrRowOffsetsType, csrColIndType, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
244: PetscCallCUSPARSE(cusparseSpMatSetAttribute(fs->spMatDescr_L, CUSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
245: PetscCallCUSPARSE(cusparseSpMatSetAttribute(fs->spMatDescr_L, CUSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));
247: fillMode = CUSPARSE_FILL_MODE_UPPER;
248: diagType = CUSPARSE_DIAG_TYPE_NON_UNIT;
249: PetscCallCUSPARSE(cusparseCreateCsr(&fs->spMatDescr_U, m, m, Mnz, fs->csrRowPtr, fs->csrColIdx, fs->csrVal, csrRowOffsetsType, csrColIndType, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
250: PetscCallCUSPARSE(cusparseSpMatSetAttribute(fs->spMatDescr_U, CUSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
251: PetscCallCUSPARSE(cusparseSpMatSetAttribute(fs->spMatDescr_U, CUSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));
253: // Allocate work vectors in SpSv
254: PetscCallCUDA(cudaMalloc((void **)&fs->X, sizeof(*fs->X) * m));
255: PetscCallCUDA(cudaMalloc((void **)&fs->Y, sizeof(*fs->Y) * m));
257: PetscCallCUSPARSE(cusparseCreateDnVec(&fs->dnVecDescr_X, m, fs->X, cusparse_scalartype));
258: PetscCallCUSPARSE(cusparseCreateDnVec(&fs->dnVecDescr_Y, m, fs->Y, cusparse_scalartype));
260: // Query buffer sizes for SpSV and then allocate buffers, temporarily assuming opA = CUSPARSE_OPERATION_NON_TRANSPOSE
261: PetscCallCUSPARSE(cusparseSpSV_createDescr(&fs->spsvDescr_L));
262: PetscCallCUSPARSE(cusparseSpSV_bufferSize(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, &fs->spsvBufferSize_L));
263: PetscCallCUSPARSE(cusparseSpSV_createDescr(&fs->spsvDescr_U));
264: PetscCallCUSPARSE(cusparseSpSV_bufferSize(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_U, &fs->spsvBufferSize_U));
265: PetscCallCUDA(cudaMalloc((void **)&fs->spsvBuffer_U, fs->spsvBufferSize_U));
266: PetscCallCUDA(cudaMalloc((void **)&fs->spsvBuffer_L, fs->spsvBufferSize_L));
268: // Record for reuse
269: fs->csrRowPtr_h = Mi;
270: fs->csrVal_h = Ma;
271: PetscCall(PetscFree(Mj));
272: }
273: // Copy the value
274: Mi = fs->csrRowPtr_h;
275: Ma = fs->csrVal_h;
276: Mnz = Mi[m];
277: for (PetscInt i = 0; i < m; i++) {
278: PetscInt llen = Ai[i + 1] - Ai[i];
279: PetscInt ulen = adiag[i] - adiag[i + 1];
280: PetscCall(PetscArraycpy(Ma + Mi[i], Aa + Ai[i], llen)); // entries of L
281: Ma[Mi[i] + llen] = (MatScalar)1.0 / Aa[adiag[i]]; // recover the diagonal entry
282: PetscCall(PetscArraycpy(Ma + Mi[i] + llen + 1, Aa + adiag[i + 1] + 1, ulen - 1)); // entries of U on the right of the diagonal
283: }
284: PetscCallCUDA(cudaMemcpy(fs->csrVal, Ma, sizeof(*Ma) * Mnz, cudaMemcpyHostToDevice));
286: #if PETSC_PKG_CUDA_VERSION_GE(12, 1, 1)
287: if (fs->updatedSpSVAnalysis) { // have done cusparseSpSV_analysis before, and only matrix values changed?
288: // Otherwise cusparse would error out: "On entry to cusparseSpSV_updateMatrix() parameter number 3 (newValues) had an illegal value: NULL pointer"
289: if (fs->csrVal) PetscCallCUSPARSE(cusparseSpSV_updateMatrix(fs->handle, fs->spsvDescr_L, fs->csrVal, CUSPARSE_SPSV_UPDATE_GENERAL));
290: if (fs->csrVal) PetscCallCUSPARSE(cusparseSpSV_updateMatrix(fs->handle, fs->spsvDescr_U, fs->csrVal, CUSPARSE_SPSV_UPDATE_GENERAL));
291: } else
292: #endif
293: {
294: // Do cusparseSpSV_analysis(), which is numeric and requires valid and up-to-date matrix values
295: PetscCallCUSPARSE(cusparseSpSV_analysis(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, fs->spsvBuffer_L));
297: PetscCallCUSPARSE(cusparseSpSV_analysis(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_U, fs->spsvBuffer_U));
298: fs->updatedSpSVAnalysis = PETSC_TRUE;
299: fs->updatedTransposeSpSVAnalysis = PETSC_FALSE;
300: }
301: }
302: PetscFunctionReturn(PETSC_SUCCESS);
303: }
305: static PetscErrorCode MatSeqAIJCUSPARSEILUAnalysisAndCopyToGPU(Mat A)
306: {
307: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;
308: Mat_SeqAIJCUSPARSETriFactors *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors *)A->spptr;
309: IS isrow = a->row, isicol = a->icol;
310: PetscBool row_identity, col_identity;
311: PetscInt n = A->rmap->n;
313: PetscFunctionBegin;
314: PetscCheck(cusparseTriFactors, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing cusparseTriFactors");
315: PetscCall(MatSeqAIJCUSPARSEBuildFactoredMatrix_LU(A));
317: cusparseTriFactors->nnz = a->nz;
319: A->offloadmask = PETSC_OFFLOAD_BOTH; // factored matrix is sync'ed to GPU
320: /* lower triangular indices */
321: PetscCall(ISIdentity(isrow, &row_identity));
322: if (!row_identity && !cusparseTriFactors->rpermIndices) {
323: const PetscInt *r;
325: PetscCall(ISGetIndices(isrow, &r));
326: cusparseTriFactors->rpermIndices = new THRUSTINTARRAY(n);
327: cusparseTriFactors->rpermIndices->assign(r, r + n);
328: PetscCall(ISRestoreIndices(isrow, &r));
329: PetscCall(PetscLogCpuToGpu(n * sizeof(PetscInt)));
330: }
332: /* upper triangular indices */
333: PetscCall(ISIdentity(isicol, &col_identity));
334: if (!col_identity && !cusparseTriFactors->cpermIndices) {
335: const PetscInt *c;
337: PetscCall(ISGetIndices(isicol, &c));
338: cusparseTriFactors->cpermIndices = new THRUSTINTARRAY(n);
339: cusparseTriFactors->cpermIndices->assign(c, c + n);
340: PetscCall(ISRestoreIndices(isicol, &c));
341: PetscCall(PetscLogCpuToGpu(n * sizeof(PetscInt)));
342: }
343: PetscFunctionReturn(PETSC_SUCCESS);
344: }
346: static PetscErrorCode MatSeqAIJCUSPARSEBuildFactoredMatrix_Cholesky(Mat A)
347: {
348: Mat_SeqAIJ *a = static_cast<Mat_SeqAIJ *>(A->data);
349: PetscInt m = A->rmap->n;
350: Mat_SeqAIJCUSPARSETriFactors *fs = static_cast<Mat_SeqAIJCUSPARSETriFactors *>(A->spptr);
351: const PetscInt *Ai = a->i, *Aj = a->j, *adiag;
352: const MatScalar *Aa = a->a;
353: PetscInt *Mj, Mnz;
354: PetscScalar *Ma, *D;
356: PetscFunctionBegin;
357: PetscCall(MatGetDiagonalMarkers_SeqAIJ(A, &adiag, NULL));
358: if (A->offloadmask == PETSC_OFFLOAD_CPU) { // A's latest factors are on CPU
359: if (!fs->csrRowPtr) { // Is this the first time we are doing setup? Use csrRowPtr since it is not null even m=0
360: // Re-arrange the (skewed) factored matrix and put the result into M, a regular csr matrix on host.
361: // See comments at MatICCFactorSymbolic_SeqAIJ() on the layout of the factored matrix (U) on host.
362: Mnz = Ai[m]; // Unz (with the unit diagonal)
363: PetscCall(PetscMalloc1(Mnz, &Ma));
364: PetscCall(PetscMalloc1(Mnz, &Mj)); // Mj[] is temp
365: PetscCall(PetscMalloc1(m, &D)); // the diagonal
366: for (PetscInt i = 0; i < m; i++) {
367: PetscInt ulen = Ai[i + 1] - Ai[i];
368: Mj[Ai[i]] = i; // diagonal entry
369: PetscCall(PetscArraycpy(Mj + Ai[i] + 1, Aj + Ai[i], ulen - 1)); // entries of U on the right of the diagonal
370: }
371: // Copy M (U) from host to device
372: PetscCallCUDA(cudaMalloc(&fs->csrRowPtr, sizeof(*fs->csrRowPtr) * (m + 1)));
373: PetscCallCUDA(cudaMalloc(&fs->csrColIdx, sizeof(*fs->csrColIdx) * Mnz));
374: PetscCallCUDA(cudaMalloc(&fs->csrVal, sizeof(*fs->csrVal) * Mnz));
375: PetscCallCUDA(cudaMalloc(&fs->diag, sizeof(*fs->diag) * m));
376: PetscCallCUDA(cudaMemcpy(fs->csrRowPtr, Ai, sizeof(*Ai) * (m + 1), cudaMemcpyHostToDevice));
377: PetscCallCUDA(cudaMemcpy(fs->csrColIdx, Mj, sizeof(*Mj) * Mnz, cudaMemcpyHostToDevice));
379: // Create descriptors for L, U. See https://docs.nvidia.com/cuda/cusparse/index.html#cusparseDiagType_t
380: // cusparseDiagType_t: This type indicates if the matrix diagonal entries are unity. The diagonal elements are always
381: // assumed to be present, but if CUSPARSE_DIAG_TYPE_UNIT is passed to an API routine, then the routine assumes that
382: // all diagonal entries are unity and will not read or modify those entries. Note that in this case the routine
383: // assumes the diagonal entries are equal to one, regardless of what those entries are actually set to in memory.
384: cusparseFillMode_t fillMode = CUSPARSE_FILL_MODE_UPPER;
385: cusparseDiagType_t diagType = CUSPARSE_DIAG_TYPE_UNIT; // U is unit diagonal
387: PetscCallCUSPARSE(cusparseCreateCsr(&fs->spMatDescr_U, m, m, Mnz, fs->csrRowPtr, fs->csrColIdx, fs->csrVal, csrRowOffsetsType, csrColIndType, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
388: PetscCallCUSPARSE(cusparseSpMatSetAttribute(fs->spMatDescr_U, CUSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
389: PetscCallCUSPARSE(cusparseSpMatSetAttribute(fs->spMatDescr_U, CUSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));
391: // Allocate work vectors in SpSv
392: PetscCallCUDA(cudaMalloc((void **)&fs->X, sizeof(*fs->X) * m));
393: PetscCallCUDA(cudaMalloc((void **)&fs->Y, sizeof(*fs->Y) * m));
395: PetscCallCUSPARSE(cusparseCreateDnVec(&fs->dnVecDescr_X, m, fs->X, cusparse_scalartype));
396: PetscCallCUSPARSE(cusparseCreateDnVec(&fs->dnVecDescr_Y, m, fs->Y, cusparse_scalartype));
398: // Query buffer sizes for SpSV and then allocate buffers
399: PetscCallCUSPARSE(cusparseSpSV_createDescr(&fs->spsvDescr_U));
400: PetscCallCUSPARSE(cusparseSpSV_bufferSize(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_U, &fs->spsvBufferSize_U));
401: PetscCallCUDA(cudaMalloc((void **)&fs->spsvBuffer_U, fs->spsvBufferSize_U));
403: PetscCallCUSPARSE(cusparseSpSV_createDescr(&fs->spsvDescr_Ut)); // Ut solve uses the same matrix (spMatDescr_U), but different descr and buffer
404: PetscCallCUSPARSE(cusparseSpSV_bufferSize(fs->handle, CUSPARSE_OPERATION_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Ut, &fs->spsvBufferSize_Ut));
405: PetscCallCUDA(cudaMalloc((void **)&fs->spsvBuffer_Ut, fs->spsvBufferSize_Ut));
407: // Record for reuse
408: fs->csrVal_h = Ma;
409: fs->diag_h = D;
410: PetscCall(PetscFree(Mj));
411: }
412: // Copy the value
413: Ma = fs->csrVal_h;
414: D = fs->diag_h;
415: Mnz = Ai[m];
416: for (PetscInt i = 0; i < m; i++) {
417: D[i] = Aa[adiag[i]]; // actually Aa[adiag[i]] is the inverse of the diagonal
418: Ma[Ai[i]] = (MatScalar)1.0; // set the unit diagonal, which is cosmetic since cusparse does not really read it given CUSPARSE_DIAG_TYPE_UNIT
419: for (PetscInt k = 0; k < Ai[i + 1] - Ai[i] - 1; k++) Ma[Ai[i] + 1 + k] = -Aa[Ai[i] + k];
420: }
421: PetscCallCUDA(cudaMemcpy(fs->csrVal, Ma, sizeof(*Ma) * Mnz, cudaMemcpyHostToDevice));
422: PetscCallCUDA(cudaMemcpy(fs->diag, D, sizeof(*D) * m, cudaMemcpyHostToDevice));
424: #if PETSC_PKG_CUDA_VERSION_GE(12, 1, 1)
425: if (fs->updatedSpSVAnalysis) {
426: if (fs->csrVal) PetscCallCUSPARSE(cusparseSpSV_updateMatrix(fs->handle, fs->spsvDescr_U, fs->csrVal, CUSPARSE_SPSV_UPDATE_GENERAL));
427: if (fs->csrVal) PetscCallCUSPARSE(cusparseSpSV_updateMatrix(fs->handle, fs->spsvDescr_Ut, fs->csrVal, CUSPARSE_SPSV_UPDATE_GENERAL));
428: } else
429: #endif
430: {
431: // Do cusparseSpSV_analysis(), which is numeric and requires valid and up-to-date matrix values
432: PetscCallCUSPARSE(cusparseSpSV_analysis(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_U, fs->spsvBuffer_U));
433: PetscCallCUSPARSE(cusparseSpSV_analysis(fs->handle, CUSPARSE_OPERATION_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Ut, fs->spsvBuffer_Ut));
434: fs->updatedSpSVAnalysis = PETSC_TRUE;
435: }
436: }
437: PetscFunctionReturn(PETSC_SUCCESS);
438: }
440: // Solve Ut D U x = b
441: static PetscErrorCode MatSolve_SeqAIJCUSPARSE_Cholesky(Mat A, Vec b, Vec x)
442: {
443: Mat_SeqAIJCUSPARSETriFactors *fs = static_cast<Mat_SeqAIJCUSPARSETriFactors *>(A->spptr);
444: Mat_SeqAIJ *aij = static_cast<Mat_SeqAIJ *>(A->data);
445: const PetscScalar *barray;
446: PetscScalar *xarray;
447: thrust::device_ptr<const PetscScalar> bGPU;
448: thrust::device_ptr<PetscScalar> xGPU;
449: const cusparseSpSVAlg_t alg = CUSPARSE_SPSV_ALG_DEFAULT;
450: PetscInt m = A->rmap->n;
452: PetscFunctionBegin;
453: PetscCall(PetscLogGpuTimeBegin());
454: PetscCall(VecCUDAGetArrayWrite(x, &xarray));
455: PetscCall(VecCUDAGetArrayRead(b, &barray));
456: xGPU = thrust::device_pointer_cast(xarray);
457: bGPU = thrust::device_pointer_cast(barray);
459: // Reorder b with the row permutation if needed, and wrap the result in fs->X
460: if (fs->rpermIndices) {
461: PetscCallThrust(thrust::copy(thrust::cuda::par.on(PetscDefaultCudaStream), thrust::make_permutation_iterator(bGPU, fs->rpermIndices->begin()), thrust::make_permutation_iterator(bGPU, fs->rpermIndices->end()), thrust::device_pointer_cast(fs->X)));
462: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, fs->X));
463: } else {
464: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, (void *)barray));
465: }
467: // Solve Ut Y = X
468: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_Y, fs->Y));
469: PetscCallCUSPARSE(cusparseSpSV_solve(fs->handle, CUSPARSE_OPERATION_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, alg, fs->spsvDescr_Ut));
471: // Solve diag(D) Z = Y. Actually just do Y = Y*D since D is already inverted in MatCholeskyFactorNumeric_SeqAIJ().
472: // It is basically a vector element-wise multiplication, but cublas does not have it!
473: #if CCCL_VERSION >= 3001000
474: auto multiplies = cuda::std::multiplies<PetscScalar>();
475: #else
476: auto multiplies = thrust::multiplies<PetscScalar>();
477: #endif
478: PetscCallThrust(thrust::transform(thrust::cuda::par.on(PetscDefaultCudaStream), thrust::device_pointer_cast(fs->Y), thrust::device_pointer_cast(fs->Y + m), thrust::device_pointer_cast(fs->diag), thrust::device_pointer_cast(fs->Y), multiplies));
480: // Solve U X = Y
481: if (fs->cpermIndices) { // if need to permute, we need to use the intermediate buffer X
482: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, fs->X));
483: } else {
484: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, xarray));
485: }
486: PetscCallCUSPARSE(cusparseSpSV_solve(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_Y, fs->dnVecDescr_X, cusparse_scalartype, alg, fs->spsvDescr_U));
488: // Reorder X with the column permutation if needed, and put the result back to x
489: if (fs->cpermIndices) {
490: PetscCallThrust(thrust::copy(thrust::cuda::par.on(PetscDefaultCudaStream), thrust::make_permutation_iterator(thrust::device_pointer_cast(fs->X), fs->cpermIndices->begin()),
491: thrust::make_permutation_iterator(thrust::device_pointer_cast(fs->X + m), fs->cpermIndices->end()), xGPU));
492: }
494: PetscCall(VecCUDARestoreArrayRead(b, &barray));
495: PetscCall(VecCUDARestoreArrayWrite(x, &xarray));
496: PetscCall(PetscLogGpuTimeEnd());
497: PetscCall(PetscLogGpuFlops(4.0 * aij->nz - A->rmap->n));
498: PetscFunctionReturn(PETSC_SUCCESS);
499: }
501: static PetscErrorCode MatSeqAIJCUSPARSEICCAnalysisAndCopyToGPU(Mat A)
502: {
503: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;
504: Mat_SeqAIJCUSPARSETriFactors *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors *)A->spptr;
505: IS ip = a->row;
506: PetscBool perm_identity;
507: PetscInt n = A->rmap->n;
509: PetscFunctionBegin;
510: PetscCheck(cusparseTriFactors, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing cusparseTriFactors");
512: PetscCall(MatSeqAIJCUSPARSEBuildFactoredMatrix_Cholesky(A));
513: cusparseTriFactors->nnz = (a->nz - n) * 2 + n;
515: A->offloadmask = PETSC_OFFLOAD_BOTH;
517: /* lower triangular indices */
518: PetscCall(ISIdentity(ip, &perm_identity));
519: if (!perm_identity) {
520: IS iip;
521: const PetscInt *irip, *rip;
523: PetscCall(ISInvertPermutation(ip, PETSC_DECIDE, &iip));
524: PetscCall(ISGetIndices(iip, &irip));
525: PetscCall(ISGetIndices(ip, &rip));
526: cusparseTriFactors->rpermIndices = new THRUSTINTARRAY(n);
527: cusparseTriFactors->rpermIndices->assign(rip, rip + n);
528: cusparseTriFactors->cpermIndices = new THRUSTINTARRAY(n);
529: cusparseTriFactors->cpermIndices->assign(irip, irip + n);
530: PetscCall(ISRestoreIndices(iip, &irip));
531: PetscCall(ISDestroy(&iip));
532: PetscCall(ISRestoreIndices(ip, &rip));
533: PetscCall(PetscLogCpuToGpu(2. * n * sizeof(PetscInt)));
534: }
535: PetscFunctionReturn(PETSC_SUCCESS);
536: }
538: static PetscErrorCode MatCholeskyFactorNumeric_SeqAIJCUSPARSE(Mat B, Mat A, const MatFactorInfo *info)
539: {
540: PetscFunctionBegin;
541: PetscCall(MatSeqAIJCUSPARSECopyFromGPU(A));
542: PetscCall(MatCholeskyFactorNumeric_SeqAIJ(B, A, info));
543: B->offloadmask = PETSC_OFFLOAD_CPU;
544: B->ops->solve = MatSolve_SeqAIJCUSPARSE_Cholesky;
545: B->ops->solvetranspose = MatSolve_SeqAIJCUSPARSE_Cholesky; // since symmetric
546: B->ops->matsolve = NULL;
547: B->ops->matsolvetranspose = NULL;
548: /* get the triangular factors */
549: PetscCall(MatSeqAIJCUSPARSEICCAnalysisAndCopyToGPU(B));
550: PetscFunctionReturn(PETSC_SUCCESS);
551: }
553: static PetscErrorCode MatSeqAIJCUSPARSEFormExplicitTranspose(Mat A)
554: {
555: Mat_SeqAIJCUSPARSE *cusparsestruct = (Mat_SeqAIJCUSPARSE *)A->spptr;
556: Mat_SeqAIJCUSPARSEMultStruct *matstruct, *matstructT;
557: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;
558: cusparseIndexBase_t indexBase;
560: PetscFunctionBegin;
561: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
562: matstruct = (Mat_SeqAIJCUSPARSEMultStruct *)cusparsestruct->mat;
563: PetscCheck(matstruct, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing mat struct");
564: matstructT = (Mat_SeqAIJCUSPARSEMultStruct *)cusparsestruct->matTranspose;
565: PetscCheck(!A->transupdated || matstructT, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing matTranspose struct");
566: if (A->transupdated) PetscFunctionReturn(PETSC_SUCCESS);
567: PetscCall(PetscLogEventBegin(MAT_CUSPARSEGenerateTranspose, A, 0, 0, 0));
568: PetscCall(PetscLogGpuTimeBegin());
569: if (cusparsestruct->format != MAT_CUSPARSE_CSR) PetscCall(MatSeqAIJCUSPARSEInvalidateTranspose(A, PETSC_TRUE));
570: if (!cusparsestruct->matTranspose) { /* create cusparse matrix */
571: matstructT = new Mat_SeqAIJCUSPARSEMultStruct;
572: PetscCallCUSPARSE(cusparseCreateMatDescr(&matstructT->descr));
573: indexBase = cusparseGetMatIndexBase(matstruct->descr);
574: PetscCallCUSPARSE(cusparseSetMatIndexBase(matstructT->descr, indexBase));
575: PetscCallCUSPARSE(cusparseSetMatType(matstructT->descr, CUSPARSE_MATRIX_TYPE_GENERAL));
577: /* set alpha and beta */
578: PetscCallCUDA(cudaMalloc((void **)&matstructT->alpha_one, sizeof(PetscScalar)));
579: PetscCallCUDA(cudaMalloc((void **)&matstructT->beta_zero, sizeof(PetscScalar)));
580: PetscCallCUDA(cudaMalloc((void **)&matstructT->beta_one, sizeof(PetscScalar)));
581: PetscCallCUDA(cudaMemcpy(matstructT->alpha_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
582: PetscCallCUDA(cudaMemcpy(matstructT->beta_zero, &PETSC_CUSPARSE_ZERO, sizeof(PetscScalar), cudaMemcpyHostToDevice));
583: PetscCallCUDA(cudaMemcpy(matstructT->beta_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
585: if (cusparsestruct->format == MAT_CUSPARSE_CSR) {
586: CsrMatrix *matrixT = new CsrMatrix;
587: matstructT->mat = matrixT;
588: matrixT->num_rows = A->cmap->n;
589: matrixT->num_cols = A->rmap->n;
590: matrixT->num_entries = a->nz;
591: matrixT->row_offsets = new THRUSTINTARRAY(matrixT->num_rows + 1);
592: matrixT->column_indices = new THRUSTINTARRAY(a->nz);
593: matrixT->values = new THRUSTARRAY(a->nz);
595: if (!cusparsestruct->rowoffsets_gpu) cusparsestruct->rowoffsets_gpu = new THRUSTINTARRAY(A->rmap->n + 1);
596: cusparsestruct->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
597: PetscCallCUSPARSE(cusparseCreateCsr(&matstructT->matDescr, matrixT->num_rows, matrixT->num_cols, matrixT->num_entries, matrixT->row_offsets->data().get(), matrixT->column_indices->data().get(), matrixT->values->data().get(), csrRowOffsetsType, csrColIndType, indexBase, cusparse_scalartype));
598: } else if (cusparsestruct->format == MAT_CUSPARSE_ELL || cusparsestruct->format == MAT_CUSPARSE_HYB) {
599: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MAT_CUSPARSE_ELL and MAT_CUSPARSE_HYB are not supported since CUDA-11.0");
600: }
601: }
602: if (cusparsestruct->format == MAT_CUSPARSE_CSR) { /* transpose mat struct may be already present, update data */
603: CsrMatrix *matrix = (CsrMatrix *)matstruct->mat;
604: CsrMatrix *matrixT = (CsrMatrix *)matstructT->mat;
605: PetscCheck(matrix, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrix");
606: PetscCheck(matrix->row_offsets, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrix rows");
607: PetscCheck(matrix->column_indices, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrix cols");
608: PetscCheck(matrix->values, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrix values");
609: PetscCheck(matrixT, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrixT");
610: PetscCheck(matrixT->row_offsets, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrixT rows");
611: PetscCheck(matrixT->column_indices, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrixT cols");
612: PetscCheck(matrixT->values, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrixT values");
613: if (!cusparsestruct->rowoffsets_gpu) { /* this may be absent when we did not construct the transpose with csr2csc */
614: cusparsestruct->rowoffsets_gpu = new THRUSTINTARRAY(A->rmap->n + 1);
615: cusparsestruct->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
616: PetscCall(PetscLogCpuToGpu((A->rmap->n + 1) * sizeof(PetscInt)));
617: }
618: if (!cusparsestruct->csr2csc_i) { // not using cusparseCsr2cscEx2() because it requires 32-bit indices
619: THRUSTINTARRAY row_indices(matrix->num_entries);
621: // Transpose the matrix via COO, i.e., by putting the row indices in column_indices[] and the column indices in row_indices[]
622: cusparsestruct->csr2csc_i = new THRUSTINTARRAY(matrix->num_entries); // will store the matrix to matrixT permutation, i.e., entry matrixT[i] is matrix[csr2csc_i[i]]
623: PetscCallThrust(thrust::sequence(thrust::device, cusparsestruct->csr2csc_i->begin(), cusparsestruct->csr2csc_i->end()));
624: PetscCallThrust(thrust::for_each(thrust::device, thrust::counting_iterator<PetscInt>(0), thrust::counting_iterator<PetscInt>(A->rmap->n), Csr2coo(cusparsestruct->rowoffsets_gpu->data().get(), matrixT->column_indices->data().get())));
625: row_indices = *matrix->column_indices;
626: // Sort the COO by row then column, and get the permutation csr2csc_i[]
627: PetscCallThrust(thrust::sort_by_key(thrust::device, thrust::make_zip_iterator(thrust::make_tuple(row_indices.begin(), matrixT->column_indices->begin())), thrust::make_zip_iterator(thrust::make_tuple(row_indices.end(), matrixT->column_indices->end())),
628: cusparsestruct->csr2csc_i->begin()));
629: // Finalize matrixT's row_offsets by looking up row_indices[]
630: PetscCallThrust(thrust::lower_bound(thrust::device, row_indices.begin(), row_indices.end(), thrust::counting_iterator<PetscInt>(0), thrust::counting_iterator<PetscInt>(A->cmap->n + 1), matrixT->row_offsets->begin()));
631: }
632: PetscCallThrust(thrust::gather(thrust::device, cusparsestruct->csr2csc_i->begin(), cusparsestruct->csr2csc_i->end(), matrix->values->begin(), matrixT->values->begin()));
633: }
634: PetscCall(PetscLogGpuTimeEnd());
635: PetscCall(PetscLogEventEnd(MAT_CUSPARSEGenerateTranspose, A, 0, 0, 0));
636: /* the compressed row indices is not used for matTranspose */
637: matstructT->cprowIndices = NULL;
638: /* assign the pointer */
639: ((Mat_SeqAIJCUSPARSE *)A->spptr)->matTranspose = matstructT;
640: A->transupdated = PETSC_TRUE;
641: PetscFunctionReturn(PETSC_SUCCESS);
642: }
644: static PetscErrorCode MatSolve_SeqAIJCUSPARSE_LU(Mat A, Vec b, Vec x)
645: {
646: const PetscScalar *barray;
647: PetscScalar *xarray;
648: thrust::device_ptr<const PetscScalar> bGPU;
649: thrust::device_ptr<PetscScalar> xGPU;
650: Mat_SeqAIJCUSPARSETriFactors *fs = static_cast<Mat_SeqAIJCUSPARSETriFactors *>(A->spptr);
651: const Mat_SeqAIJ *aij = static_cast<Mat_SeqAIJ *>(A->data);
652: const cusparseOperation_t op = CUSPARSE_OPERATION_NON_TRANSPOSE;
653: const cusparseSpSVAlg_t alg = CUSPARSE_SPSV_ALG_DEFAULT;
654: PetscInt m = A->rmap->n;
656: PetscFunctionBegin;
657: PetscCall(PetscLogGpuTimeBegin());
658: PetscCall(VecCUDAGetArrayWrite(x, &xarray));
659: PetscCall(VecCUDAGetArrayRead(b, &barray));
660: xGPU = thrust::device_pointer_cast(xarray);
661: bGPU = thrust::device_pointer_cast(barray);
663: // Reorder b with the row permutation if needed, and wrap the result in fs->X
664: if (fs->rpermIndices) {
665: PetscCallThrust(thrust::copy(thrust::cuda::par.on(PetscDefaultCudaStream), thrust::make_permutation_iterator(bGPU, fs->rpermIndices->begin()), thrust::make_permutation_iterator(bGPU, fs->rpermIndices->end()), thrust::device_pointer_cast(fs->X)));
666: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, fs->X));
667: } else {
668: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, (void *)barray));
669: }
671: // Solve L Y = X
672: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_Y, fs->Y));
673: // Note that cusparseSpSV_solve() secretly uses the external buffer used in cusparseSpSV_analysis()!
674: PetscCallCUSPARSE(cusparseSpSV_solve(fs->handle, op, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, alg, fs->spsvDescr_L));
676: // Solve U X = Y
677: if (fs->cpermIndices) {
678: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, fs->X));
679: } else {
680: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, xarray));
681: }
682: PetscCallCUSPARSE(cusparseSpSV_solve(fs->handle, op, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_Y, fs->dnVecDescr_X, cusparse_scalartype, alg, fs->spsvDescr_U));
684: // Reorder X with the column permutation if needed, and put the result back to x
685: if (fs->cpermIndices) {
686: PetscCallThrust(thrust::copy(thrust::cuda::par.on(PetscDefaultCudaStream), thrust::make_permutation_iterator(thrust::device_pointer_cast(fs->X), fs->cpermIndices->begin()),
687: thrust::make_permutation_iterator(thrust::device_pointer_cast(fs->X + m), fs->cpermIndices->end()), xGPU));
688: }
689: PetscCall(VecCUDARestoreArrayRead(b, &barray));
690: PetscCall(VecCUDARestoreArrayWrite(x, &xarray));
691: PetscCall(PetscLogGpuTimeEnd());
692: PetscCall(PetscLogGpuFlops(2.0 * aij->nz - m));
693: PetscFunctionReturn(PETSC_SUCCESS);
694: }
696: static PetscErrorCode MatSolveTranspose_SeqAIJCUSPARSE_LU(Mat A, Vec b, Vec x)
697: {
698: Mat_SeqAIJCUSPARSETriFactors *fs = static_cast<Mat_SeqAIJCUSPARSETriFactors *>(A->spptr);
699: Mat_SeqAIJ *aij = static_cast<Mat_SeqAIJ *>(A->data);
700: const PetscScalar *barray;
701: PetscScalar *xarray;
702: thrust::device_ptr<const PetscScalar> bGPU;
703: thrust::device_ptr<PetscScalar> xGPU;
704: const cusparseOperation_t opA = CUSPARSE_OPERATION_TRANSPOSE;
705: const cusparseSpSVAlg_t alg = CUSPARSE_SPSV_ALG_DEFAULT;
706: PetscInt m = A->rmap->n;
708: PetscFunctionBegin;
709: PetscCall(PetscLogGpuTimeBegin());
710: if (!fs->createdTransposeSpSVDescr) { // Call MatSolveTranspose() for the first time
711: PetscCallCUSPARSE(cusparseSpSV_createDescr(&fs->spsvDescr_Lt));
712: PetscCallCUSPARSE(cusparseSpSV_bufferSize(fs->handle, opA, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, /* The matrix is still L. We only do transpose solve with it */
713: fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, alg, fs->spsvDescr_Lt, &fs->spsvBufferSize_Lt));
715: PetscCallCUSPARSE(cusparseSpSV_createDescr(&fs->spsvDescr_Ut));
716: PetscCallCUSPARSE(cusparseSpSV_bufferSize(fs->handle, opA, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, alg, fs->spsvDescr_Ut, &fs->spsvBufferSize_Ut));
717: PetscCallCUDA(cudaMalloc((void **)&fs->spsvBuffer_Lt, fs->spsvBufferSize_Lt));
718: PetscCallCUDA(cudaMalloc((void **)&fs->spsvBuffer_Ut, fs->spsvBufferSize_Ut));
719: fs->createdTransposeSpSVDescr = PETSC_TRUE;
720: }
722: if (!fs->updatedTransposeSpSVAnalysis) {
723: PetscCallCUSPARSE(cusparseSpSV_analysis(fs->handle, opA, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, alg, fs->spsvDescr_Lt, fs->spsvBuffer_Lt));
725: PetscCallCUSPARSE(cusparseSpSV_analysis(fs->handle, opA, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, alg, fs->spsvDescr_Ut, fs->spsvBuffer_Ut));
726: fs->updatedTransposeSpSVAnalysis = PETSC_TRUE;
727: }
729: PetscCall(VecCUDAGetArrayWrite(x, &xarray));
730: PetscCall(VecCUDAGetArrayRead(b, &barray));
731: xGPU = thrust::device_pointer_cast(xarray);
732: bGPU = thrust::device_pointer_cast(barray);
734: // Reorder b with the row permutation if needed, and wrap the result in fs->X
735: if (fs->rpermIndices) {
736: PetscCallThrust(thrust::copy(thrust::cuda::par.on(PetscDefaultCudaStream), thrust::make_permutation_iterator(bGPU, fs->rpermIndices->begin()), thrust::make_permutation_iterator(bGPU, fs->rpermIndices->end()), thrust::device_pointer_cast(fs->X)));
737: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, fs->X));
738: } else {
739: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, (void *)barray));
740: }
742: // Solve Ut Y = X
743: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_Y, fs->Y));
744: PetscCallCUSPARSE(cusparseSpSV_solve(fs->handle, opA, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, alg, fs->spsvDescr_Ut));
746: // Solve Lt X = Y
747: if (fs->cpermIndices) { // if need to permute, we need to use the intermediate buffer X
748: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, fs->X));
749: } else {
750: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, xarray));
751: }
752: PetscCallCUSPARSE(cusparseSpSV_solve(fs->handle, opA, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_Y, fs->dnVecDescr_X, cusparse_scalartype, alg, fs->spsvDescr_Lt));
754: // Reorder X with the column permutation if needed, and put the result back to x
755: if (fs->cpermIndices) {
756: PetscCallThrust(thrust::copy(thrust::cuda::par.on(PetscDefaultCudaStream), thrust::make_permutation_iterator(thrust::device_pointer_cast(fs->X), fs->cpermIndices->begin()),
757: thrust::make_permutation_iterator(thrust::device_pointer_cast(fs->X + m), fs->cpermIndices->end()), xGPU));
758: }
760: PetscCall(VecCUDARestoreArrayRead(b, &barray));
761: PetscCall(VecCUDARestoreArrayWrite(x, &xarray));
762: PetscCall(PetscLogGpuTimeEnd());
763: PetscCall(PetscLogGpuFlops(2.0 * aij->nz - A->rmap->n));
764: PetscFunctionReturn(PETSC_SUCCESS);
765: }
767: static PetscErrorCode MatILUFactorNumeric_SeqAIJCUSPARSE_ILU0(Mat fact, Mat A, const MatFactorInfo *)
768: {
769: Mat_SeqAIJCUSPARSETriFactors *fs = (Mat_SeqAIJCUSPARSETriFactors *)fact->spptr;
770: Mat_SeqAIJ *aij = (Mat_SeqAIJ *)fact->data;
771: Mat_SeqAIJCUSPARSE *Acusp = (Mat_SeqAIJCUSPARSE *)A->spptr;
772: CsrMatrix *Acsr;
773: PetscInt m, nz;
774: PetscBool flg;
776: PetscFunctionBegin;
777: if (PetscDefined(USE_DEBUG)) {
778: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJCUSPARSE, &flg));
779: PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Expected MATSEQAIJCUSPARSE, but input is %s", ((PetscObject)A)->type_name);
780: }
782: /* Copy A's value to fact */
783: m = fact->rmap->n;
784: nz = aij->nz;
785: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
786: Acsr = (CsrMatrix *)Acusp->mat->mat;
787: PetscCallCUDA(cudaMemcpyAsync(fs->csrVal, Acsr->values->data().get(), sizeof(PetscScalar) * nz, cudaMemcpyDeviceToDevice, PetscDefaultCudaStream));
789: PetscCall(PetscLogGpuTimeBegin());
790: /* Factorize fact inplace */
791: if (m)
792: PetscCallCUSPARSE(cusparseXcsrilu02(fs->handle, m, nz, /* cusparseXcsrilu02 errors out with empty matrices (m=0) */
793: fs->matDescr_M, fs->csrVal, fs->csrRowPtr32, fs->csrColIdx32, fs->ilu0Info_M, fs->policy_M, fs->factBuffer_M));
794: if (PetscDefined(USE_DEBUG)) {
795: int numerical_zero;
796: cusparseStatus_t status;
797: status = cusparseXcsrilu02_zeroPivot(fs->handle, fs->ilu0Info_M, &numerical_zero);
798: PetscAssert(CUSPARSE_STATUS_ZERO_PIVOT != status, PETSC_COMM_SELF, PETSC_ERR_USER_INPUT, "Numerical zero pivot detected in csrilu02: A(%d,%d) is zero", numerical_zero, numerical_zero);
799: }
801: #if PETSC_PKG_CUDA_VERSION_GE(12, 1, 1)
802: if (fs->updatedSpSVAnalysis) {
803: if (fs->csrVal) PetscCallCUSPARSE(cusparseSpSV_updateMatrix(fs->handle, fs->spsvDescr_L, fs->csrVal, CUSPARSE_SPSV_UPDATE_GENERAL));
804: if (fs->csrVal) PetscCallCUSPARSE(cusparseSpSV_updateMatrix(fs->handle, fs->spsvDescr_U, fs->csrVal, CUSPARSE_SPSV_UPDATE_GENERAL));
805: } else
806: #endif
807: {
808: /* cusparseSpSV_analysis() is numeric, i.e., it requires valid matrix values, therefore, we do it after cusparseXcsrilu02()
809: See discussion at https://github.com/NVIDIA/CUDALibrarySamples/issues/78
810: */
811: PetscCallCUSPARSE(cusparseSpSV_analysis(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, fs->spsvBuffer_L));
813: PetscCallCUSPARSE(cusparseSpSV_analysis(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_U, fs->spsvBuffer_U));
815: fs->updatedSpSVAnalysis = PETSC_TRUE;
816: /* L, U values have changed, reset the flag to indicate we need to redo cusparseSpSV_analysis() for transpose solve */
817: fs->updatedTransposeSpSVAnalysis = PETSC_FALSE;
818: }
820: fact->offloadmask = PETSC_OFFLOAD_GPU;
821: fact->ops->solve = MatSolve_SeqAIJCUSPARSE_LU; // spMatDescr_L/U uses 32-bit indices, but cusparseSpSV_solve() supports both 32 and 64. The info is encoded in cusparseSpMatDescr_t.
822: fact->ops->solvetranspose = MatSolveTranspose_SeqAIJCUSPARSE_LU;
823: fact->ops->matsolve = NULL;
824: fact->ops->matsolvetranspose = NULL;
825: PetscCall(PetscLogGpuTimeEnd());
826: PetscCall(PetscLogGpuFlops(fs->numericFactFlops));
827: PetscFunctionReturn(PETSC_SUCCESS);
828: }
830: static PetscErrorCode MatILUFactorSymbolic_SeqAIJCUSPARSE_ILU0(Mat fact, Mat A, IS, IS, const MatFactorInfo *info)
831: {
832: Mat_SeqAIJCUSPARSETriFactors *fs = (Mat_SeqAIJCUSPARSETriFactors *)fact->spptr;
833: Mat_SeqAIJ *aij = (Mat_SeqAIJ *)fact->data;
834: PetscInt m, nz;
836: PetscFunctionBegin;
837: if (PetscDefined(USE_DEBUG)) {
838: PetscBool flg, diagDense;
840: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJCUSPARSE, &flg));
841: PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Expected MATSEQAIJCUSPARSE, but input is %s", ((PetscObject)A)->type_name);
842: PetscCheck(A->rmap->n == A->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Must be square matrix, rows %" PetscInt_FMT " columns %" PetscInt_FMT, A->rmap->n, A->cmap->n);
843: PetscCall(MatGetDiagonalMarkers_SeqAIJ(A, NULL, &diagDense));
844: PetscCheck(diagDense, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix is missing a diagonal entry");
845: }
847: /* Free the old stale stuff */
848: PetscCall(MatSeqAIJCUSPARSETriFactors_Reset(&fs));
850: /* Copy over A's meta data to fact. Note that we also allocated fact's i,j,a on host,
851: but they will not be used. Allocate them just for easy debugging.
852: */
853: PetscCall(MatDuplicateNoCreate_SeqAIJ(fact, A, MAT_DO_NOT_COPY_VALUES, PETSC_TRUE /*malloc*/));
855: fact->offloadmask = PETSC_OFFLOAD_BOTH;
856: fact->factortype = MAT_FACTOR_ILU;
857: fact->info.factor_mallocs = 0;
858: fact->info.fill_ratio_given = info->fill;
859: fact->info.fill_ratio_needed = 1.0;
861: aij->row = NULL;
862: aij->col = NULL;
864: /* ====================================================================== */
865: /* Copy A's i, j to fact and also allocate the value array of fact. */
866: /* We'll do in-place factorization on fact */
867: /* ====================================================================== */
868: const PetscInt *Ai, *Aj;
870: m = fact->rmap->n;
871: nz = aij->nz;
873: PetscCallCUDA(cudaMalloc((void **)&fs->csrRowPtr32, sizeof(*fs->csrRowPtr32) * (m + 1)));
874: PetscCallCUDA(cudaMalloc((void **)&fs->csrColIdx32, sizeof(*fs->csrColIdx32) * nz));
875: PetscCallCUDA(cudaMalloc((void **)&fs->csrVal, sizeof(*fs->csrVal) * nz));
876: PetscCall(MatSeqAIJCUSPARSEGetIJ(A, PETSC_FALSE, &Ai, &Aj)); // Ai is uncompressed
878: PetscCheck(nz <= INT_MAX && m <= INT_MAX, PETSC_COMM_SELF, PETSC_ERR_SUP, "nnz %" PetscInt_FMT " and rows %" PetscInt_FMT " overflow C int", nz, m);
879: PetscCallThrust(thrust::transform(thrust::cuda::par.on(PetscDefaultCudaStream), Ai, Ai + m + 1, fs->csrRowPtr32, PetscIntToCInt()));
880: PetscCallThrust(thrust::transform(thrust::cuda::par.on(PetscDefaultCudaStream), Aj, Aj + nz, fs->csrColIdx32, PetscIntToCInt()));
882: /* ====================================================================== */
883: /* Create descriptors for M, L, U */
884: /* ====================================================================== */
885: cusparseFillMode_t fillMode;
886: cusparseDiagType_t diagType;
888: PetscCallCUSPARSE(cusparseCreateMatDescr(&fs->matDescr_M));
889: PetscCallCUSPARSE(cusparseSetMatIndexBase(fs->matDescr_M, CUSPARSE_INDEX_BASE_ZERO));
890: PetscCallCUSPARSE(cusparseSetMatType(fs->matDescr_M, CUSPARSE_MATRIX_TYPE_GENERAL));
892: /* https://docs.nvidia.com/cuda/cusparse/index.html#cusparseDiagType_t
893: cusparseDiagType_t: This type indicates if the matrix diagonal entries are unity. The diagonal elements are always
894: assumed to be present, but if CUSPARSE_DIAG_TYPE_UNIT is passed to an API routine, then the routine assumes that
895: all diagonal entries are unity and will not read or modify those entries. Note that in this case the routine
896: assumes the diagonal entries are equal to one, regardless of what those entries are actually set to in memory.
897: */
898: fillMode = CUSPARSE_FILL_MODE_LOWER;
899: diagType = CUSPARSE_DIAG_TYPE_UNIT;
900: PetscCallCUSPARSE(cusparseCreateCsr(&fs->spMatDescr_L, m, m, nz, fs->csrRowPtr32, fs->csrColIdx32, fs->csrVal, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
901: PetscCallCUSPARSE(cusparseSpMatSetAttribute(fs->spMatDescr_L, CUSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
902: PetscCallCUSPARSE(cusparseSpMatSetAttribute(fs->spMatDescr_L, CUSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));
904: fillMode = CUSPARSE_FILL_MODE_UPPER;
905: diagType = CUSPARSE_DIAG_TYPE_NON_UNIT;
906: PetscCallCUSPARSE(cusparseCreateCsr(&fs->spMatDescr_U, m, m, nz, fs->csrRowPtr32, fs->csrColIdx32, fs->csrVal, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
907: PetscCallCUSPARSE(cusparseSpMatSetAttribute(fs->spMatDescr_U, CUSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
908: PetscCallCUSPARSE(cusparseSpMatSetAttribute(fs->spMatDescr_U, CUSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));
910: /* ========================================================================= */
911: /* Query buffer sizes for csrilu0, SpSV and allocate buffers */
912: /* ========================================================================= */
913: PetscCallCUSPARSE(cusparseCreateCsrilu02Info(&fs->ilu0Info_M));
914: if (m)
915: PetscCallCUSPARSE(cusparseXcsrilu02_bufferSize(fs->handle, m, nz, /* cusparseXcsrilu02 errors out with empty matrices (m=0) */
916: fs->matDescr_M, fs->csrVal, fs->csrRowPtr32, fs->csrColIdx32, fs->ilu0Info_M, &fs->factBufferSize_M));
918: PetscCallCUDA(cudaMalloc((void **)&fs->X, sizeof(PetscScalar) * m));
919: PetscCallCUDA(cudaMalloc((void **)&fs->Y, sizeof(PetscScalar) * m));
921: PetscCallCUSPARSE(cusparseCreateDnVec(&fs->dnVecDescr_X, m, fs->X, cusparse_scalartype));
922: PetscCallCUSPARSE(cusparseCreateDnVec(&fs->dnVecDescr_Y, m, fs->Y, cusparse_scalartype));
924: PetscCallCUSPARSE(cusparseSpSV_createDescr(&fs->spsvDescr_L));
925: PetscCallCUSPARSE(cusparseSpSV_bufferSize(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, &fs->spsvBufferSize_L));
927: PetscCallCUSPARSE(cusparseSpSV_createDescr(&fs->spsvDescr_U));
928: PetscCallCUSPARSE(cusparseSpSV_bufferSize(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_U, &fs->spsvBufferSize_U));
930: /* From my experiment with the example at https://github.com/NVIDIA/CUDALibrarySamples/tree/master/cuSPARSE/bicgstab,
931: and discussion at https://github.com/NVIDIA/CUDALibrarySamples/issues/77,
932: spsvBuffer_L/U can not be shared (i.e., the same) for our case, but factBuffer_M can share with either of spsvBuffer_L/U.
933: To save memory, we make factBuffer_M share with the bigger of spsvBuffer_L/U.
934: */
935: if (fs->spsvBufferSize_L > fs->spsvBufferSize_U) {
936: PetscCallCUDA(cudaMalloc((void **)&fs->factBuffer_M, PetscMax(fs->spsvBufferSize_L, (size_t)fs->factBufferSize_M)));
937: fs->spsvBuffer_L = fs->factBuffer_M;
938: PetscCallCUDA(cudaMalloc((void **)&fs->spsvBuffer_U, fs->spsvBufferSize_U));
939: } else {
940: PetscCallCUDA(cudaMalloc((void **)&fs->factBuffer_M, PetscMax(fs->spsvBufferSize_U, (size_t)fs->factBufferSize_M)));
941: fs->spsvBuffer_U = fs->factBuffer_M;
942: PetscCallCUDA(cudaMalloc((void **)&fs->spsvBuffer_L, fs->spsvBufferSize_L));
943: }
945: /* ========================================================================== */
946: /* Perform analysis of ilu0 on M, SpSv on L and U */
947: /* The lower(upper) triangular part of M has the same sparsity pattern as L(U)*/
948: /* ========================================================================== */
949: int structural_zero;
950: cusparseStatus_t status;
952: fs->policy_M = CUSPARSE_SOLVE_POLICY_USE_LEVEL;
953: if (m)
954: PetscCallCUSPARSE(cusparseXcsrilu02_analysis(fs->handle, m, nz, /* cusparseXcsrilu02 errors out with empty matrices (m=0) */
955: fs->matDescr_M, fs->csrVal, fs->csrRowPtr32, fs->csrColIdx32, fs->ilu0Info_M, fs->policy_M, fs->factBuffer_M));
956: if (PetscDefined(USE_DEBUG)) {
957: /* cusparseXcsrilu02_zeroPivot() is a blocking call. It calls cudaDeviceSynchronize() to make sure all previous kernels are done. */
958: status = cusparseXcsrilu02_zeroPivot(fs->handle, fs->ilu0Info_M, &structural_zero);
959: PetscCheck(CUSPARSE_STATUS_ZERO_PIVOT != status, PETSC_COMM_SELF, PETSC_ERR_USER_INPUT, "Structural zero pivot detected in csrilu02: A(%d,%d) is missing", structural_zero, structural_zero);
960: }
962: /* Estimate FLOPs of the numeric factorization */
963: {
964: Mat_SeqAIJ *Aseq = (Mat_SeqAIJ *)A->data;
965: PetscInt *Ai, nzRow, nzLeft;
966: const PetscInt *adiag;
967: PetscLogDouble flops = 0.0;
969: PetscCall(MatGetDiagonalMarkers_SeqAIJ(A, &adiag, NULL));
970: Ai = Aseq->i;
971: for (PetscInt i = 0; i < m; i++) {
972: if (Ai[i] < adiag[i] && adiag[i] < Ai[i + 1]) { /* There are nonzeros left to the diagonal of row i */
973: nzRow = Ai[i + 1] - Ai[i];
974: nzLeft = adiag[i] - Ai[i];
975: /* We want to eliminate nonzeros left to the diagonal one by one. Assume each time, nonzeros right
976: and include the eliminated one will be updated, which incurs a multiplication and an addition.
977: */
978: nzLeft = (nzRow - 1) / 2;
979: flops += nzLeft * (2.0 * nzRow - nzLeft + 1);
980: }
981: }
982: fs->numericFactFlops = flops;
983: }
984: fact->ops->lufactornumeric = MatILUFactorNumeric_SeqAIJCUSPARSE_ILU0;
985: PetscFunctionReturn(PETSC_SUCCESS);
986: }
988: static PetscErrorCode MatSolve_SeqAIJCUSPARSE_ICC0(Mat fact, Vec b, Vec x)
989: {
990: Mat_SeqAIJCUSPARSETriFactors *fs = (Mat_SeqAIJCUSPARSETriFactors *)fact->spptr;
991: Mat_SeqAIJ *aij = (Mat_SeqAIJ *)fact->data;
992: const PetscScalar *barray;
993: PetscScalar *xarray;
995: PetscFunctionBegin;
996: PetscCall(VecCUDAGetArrayWrite(x, &xarray));
997: PetscCall(VecCUDAGetArrayRead(b, &barray));
998: PetscCall(PetscLogGpuTimeBegin());
1000: /* Solve L*y = b */
1001: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, (void *)barray));
1002: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_Y, fs->Y));
1003: PetscCallCUSPARSE(cusparseSpSV_solve(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, /* L Y = X */
1004: fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L));
1006: /* Solve Lt*x = y */
1007: PetscCallCUSPARSE(cusparseDnVecSetValues(fs->dnVecDescr_X, xarray));
1008: PetscCallCUSPARSE(cusparseSpSV_solve(fs->handle, CUSPARSE_OPERATION_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, /* Lt X = Y */
1009: fs->dnVecDescr_Y, fs->dnVecDescr_X, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Lt));
1011: PetscCall(VecCUDARestoreArrayRead(b, &barray));
1012: PetscCall(VecCUDARestoreArrayWrite(x, &xarray));
1014: PetscCall(PetscLogGpuTimeEnd());
1015: PetscCall(PetscLogGpuFlops(2.0 * aij->nz - fact->rmap->n));
1016: PetscFunctionReturn(PETSC_SUCCESS);
1017: }
1019: static PetscErrorCode MatICCFactorNumeric_SeqAIJCUSPARSE_ICC0(Mat fact, Mat A, const MatFactorInfo *)
1020: {
1021: Mat_SeqAIJCUSPARSETriFactors *fs = (Mat_SeqAIJCUSPARSETriFactors *)fact->spptr;
1022: Mat_SeqAIJ *aij = (Mat_SeqAIJ *)fact->data;
1023: Mat_SeqAIJCUSPARSE *Acusp = (Mat_SeqAIJCUSPARSE *)A->spptr;
1024: CsrMatrix *Acsr;
1025: PetscInt m, nz;
1026: PetscBool flg;
1028: PetscFunctionBegin;
1029: if (PetscDefined(USE_DEBUG)) {
1030: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJCUSPARSE, &flg));
1031: PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Expected MATSEQAIJCUSPARSE, but input is %s", ((PetscObject)A)->type_name);
1032: }
1034: /* Copy A's value to fact */
1035: m = fact->rmap->n;
1036: nz = aij->nz;
1037: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
1038: Acsr = (CsrMatrix *)Acusp->mat->mat;
1039: PetscCallCUDA(cudaMemcpyAsync(fs->csrVal, Acsr->values->data().get(), sizeof(PetscScalar) * nz, cudaMemcpyDeviceToDevice, PetscDefaultCudaStream));
1041: /* Factorize fact inplace */
1042: /* https://docs.nvidia.com/cuda/cusparse/index.html#csric02_solve
1043: csric02() only takes the lower triangular part of matrix A to perform factorization.
1044: The matrix type must be CUSPARSE_MATRIX_TYPE_GENERAL, the fill mode and diagonal type are ignored,
1045: and the strictly upper triangular part is ignored and never touched. It does not matter if A is Hermitian or not.
1046: In other words, from the point of view of csric02() A is Hermitian and only the lower triangular part is provided.
1047: */
1048: if (m) PetscCallCUSPARSE(cusparseXcsric02(fs->handle, m, nz, fs->matDescr_M, fs->csrVal, fs->csrRowPtr32, fs->csrColIdx32, fs->ic0Info_M, fs->policy_M, fs->factBuffer_M));
1049: if (PetscDefined(USE_DEBUG)) {
1050: int numerical_zero;
1051: cusparseStatus_t status;
1052: status = cusparseXcsric02_zeroPivot(fs->handle, fs->ic0Info_M, &numerical_zero);
1053: PetscAssert(CUSPARSE_STATUS_ZERO_PIVOT != status, PETSC_COMM_SELF, PETSC_ERR_USER_INPUT, "Numerical zero pivot detected in csric02: A(%d,%d) is zero", numerical_zero, numerical_zero);
1054: }
1056: #if PETSC_PKG_CUDA_VERSION_GE(12, 1, 1)
1057: if (fs->updatedSpSVAnalysis) {
1058: if (fs->csrVal) PetscCallCUSPARSE(cusparseSpSV_updateMatrix(fs->handle, fs->spsvDescr_L, fs->csrVal, CUSPARSE_SPSV_UPDATE_GENERAL));
1059: if (fs->csrVal) PetscCallCUSPARSE(cusparseSpSV_updateMatrix(fs->handle, fs->spsvDescr_Lt, fs->csrVal, CUSPARSE_SPSV_UPDATE_GENERAL));
1060: } else
1061: #endif
1062: {
1063: PetscCallCUSPARSE(cusparseSpSV_analysis(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, fs->spsvBuffer_L));
1065: /* Note that cusparse reports this error if we use double and CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE
1066: ** On entry to cusparseSpSV_analysis(): conjugate transpose (opA) is not supported for matA data type, current -> CUDA_R_64F
1067: */
1068: PetscCallCUSPARSE(cusparseSpSV_analysis(fs->handle, CUSPARSE_OPERATION_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Lt, fs->spsvBuffer_Lt));
1069: fs->updatedSpSVAnalysis = PETSC_TRUE;
1070: }
1072: fact->offloadmask = PETSC_OFFLOAD_GPU;
1073: fact->ops->solve = MatSolve_SeqAIJCUSPARSE_ICC0;
1074: fact->ops->solvetranspose = MatSolve_SeqAIJCUSPARSE_ICC0;
1075: fact->ops->matsolve = NULL;
1076: fact->ops->matsolvetranspose = NULL;
1077: PetscCall(PetscLogGpuFlops(fs->numericFactFlops));
1078: PetscFunctionReturn(PETSC_SUCCESS);
1079: }
1081: static PetscErrorCode MatICCFactorSymbolic_SeqAIJCUSPARSE_ICC0(Mat fact, Mat A, IS, const MatFactorInfo *info)
1082: {
1083: Mat_SeqAIJCUSPARSETriFactors *fs = (Mat_SeqAIJCUSPARSETriFactors *)fact->spptr;
1084: Mat_SeqAIJ *aij = (Mat_SeqAIJ *)fact->data;
1085: PetscInt m, nz;
1087: PetscFunctionBegin;
1088: if (PetscDefined(USE_DEBUG)) {
1089: PetscBool flg, diagDense;
1091: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJCUSPARSE, &flg));
1092: PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Expected MATSEQAIJCUSPARSE, but input is %s", ((PetscObject)A)->type_name);
1093: PetscCheck(A->rmap->n == A->cmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Must be square matrix, rows %" PetscInt_FMT " columns %" PetscInt_FMT, A->rmap->n, A->cmap->n);
1094: PetscCall(MatGetDiagonalMarkers_SeqAIJ(A, NULL, &diagDense));
1095: PetscCheck(diagDense, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix is missing diagonal entries");
1096: }
1098: /* Free the old stale stuff */
1099: PetscCall(MatSeqAIJCUSPARSETriFactors_Reset(&fs));
1101: /* Copy over A's meta data to fact. Note that we also allocated fact's i,j,a on host,
1102: but they will not be used. Allocate them just for easy debugging.
1103: */
1104: PetscCall(MatDuplicateNoCreate_SeqAIJ(fact, A, MAT_DO_NOT_COPY_VALUES, PETSC_TRUE /*malloc*/));
1106: fact->offloadmask = PETSC_OFFLOAD_BOTH;
1107: fact->factortype = MAT_FACTOR_ICC;
1108: fact->info.factor_mallocs = 0;
1109: fact->info.fill_ratio_given = info->fill;
1110: fact->info.fill_ratio_needed = 1.0;
1112: aij->row = NULL;
1113: aij->col = NULL;
1115: /* ====================================================================== */
1116: /* Copy A's i, j to fact and also allocate the value array of fact. */
1117: /* We'll do in-place factorization on fact */
1118: /* ====================================================================== */
1119: const PetscInt *Ai, *Aj;
1121: m = fact->rmap->n;
1122: nz = aij->nz;
1124: PetscCallCUDA(cudaMalloc((void **)&fs->csrRowPtr32, sizeof(*fs->csrRowPtr32) * (m + 1)));
1125: PetscCallCUDA(cudaMalloc((void **)&fs->csrColIdx32, sizeof(*fs->csrColIdx32) * nz));
1126: PetscCallCUDA(cudaMalloc((void **)&fs->csrVal, sizeof(PetscScalar) * nz));
1127: PetscCall(MatSeqAIJCUSPARSEGetIJ(A, PETSC_FALSE, &Ai, &Aj)); // Ai is uncompressed
1129: PetscCheck(nz <= INT_MAX && m <= INT_MAX, PETSC_COMM_SELF, PETSC_ERR_SUP, "nnz %" PetscInt_FMT " and rows %" PetscInt_FMT " overflow C int", nz, m);
1130: PetscCallThrust(thrust::transform(thrust::cuda::par.on(PetscDefaultCudaStream), Ai, Ai + m + 1, fs->csrRowPtr32, PetscIntToCInt()));
1131: PetscCallThrust(thrust::transform(thrust::cuda::par.on(PetscDefaultCudaStream), Aj, Aj + nz, fs->csrColIdx32, PetscIntToCInt()));
1133: /* ====================================================================== */
1134: /* Create mat descriptors for M, L */
1135: /* ====================================================================== */
1136: cusparseFillMode_t fillMode;
1137: cusparseDiagType_t diagType;
1139: PetscCallCUSPARSE(cusparseCreateMatDescr(&fs->matDescr_M));
1140: PetscCallCUSPARSE(cusparseSetMatIndexBase(fs->matDescr_M, CUSPARSE_INDEX_BASE_ZERO));
1141: PetscCallCUSPARSE(cusparseSetMatType(fs->matDescr_M, CUSPARSE_MATRIX_TYPE_GENERAL));
1143: /* https://docs.nvidia.com/cuda/cusparse/index.html#cusparseDiagType_t
1144: cusparseDiagType_t: This type indicates if the matrix diagonal entries are unity. The diagonal elements are always
1145: assumed to be present, but if CUSPARSE_DIAG_TYPE_UNIT is passed to an API routine, then the routine assumes that
1146: all diagonal entries are unity and will not read or modify those entries. Note that in this case the routine
1147: assumes the diagonal entries are equal to one, regardless of what those entries are actually set to in memory.
1148: */
1149: fillMode = CUSPARSE_FILL_MODE_LOWER;
1150: diagType = CUSPARSE_DIAG_TYPE_NON_UNIT;
1151: PetscCallCUSPARSE(cusparseCreateCsr(&fs->spMatDescr_L, m, m, nz, fs->csrRowPtr32, fs->csrColIdx32, fs->csrVal, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
1152: PetscCallCUSPARSE(cusparseSpMatSetAttribute(fs->spMatDescr_L, CUSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
1153: PetscCallCUSPARSE(cusparseSpMatSetAttribute(fs->spMatDescr_L, CUSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));
1155: /* ========================================================================= */
1156: /* Query buffer sizes for csric0, SpSV of L and Lt, and allocate buffers */
1157: /* ========================================================================= */
1158: PetscCallCUSPARSE(cusparseCreateCsric02Info(&fs->ic0Info_M));
1159: if (m) PetscCallCUSPARSE(cusparseXcsric02_bufferSize(fs->handle, m, nz, fs->matDescr_M, fs->csrVal, fs->csrRowPtr32, fs->csrColIdx32, fs->ic0Info_M, &fs->factBufferSize_M));
1161: PetscCallCUDA(cudaMalloc((void **)&fs->X, sizeof(PetscScalar) * m));
1162: PetscCallCUDA(cudaMalloc((void **)&fs->Y, sizeof(PetscScalar) * m));
1164: PetscCallCUSPARSE(cusparseCreateDnVec(&fs->dnVecDescr_X, m, fs->X, cusparse_scalartype));
1165: PetscCallCUSPARSE(cusparseCreateDnVec(&fs->dnVecDescr_Y, m, fs->Y, cusparse_scalartype));
1167: PetscCallCUSPARSE(cusparseSpSV_createDescr(&fs->spsvDescr_L));
1168: PetscCallCUSPARSE(cusparseSpSV_bufferSize(fs->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, &fs->spsvBufferSize_L));
1170: PetscCallCUSPARSE(cusparseSpSV_createDescr(&fs->spsvDescr_Lt));
1171: PetscCallCUSPARSE(cusparseSpSV_bufferSize(fs->handle, CUSPARSE_OPERATION_TRANSPOSE, &PETSC_CUSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, cusparse_scalartype, CUSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Lt, &fs->spsvBufferSize_Lt));
1173: /* To save device memory, we make the factorization buffer share with one of the solver buffer.
1174: See also comments in MatILUFactorSymbolic_SeqAIJCUSPARSE_ILU0().
1175: */
1176: if (fs->spsvBufferSize_L > fs->spsvBufferSize_Lt) {
1177: PetscCallCUDA(cudaMalloc((void **)&fs->factBuffer_M, PetscMax(fs->spsvBufferSize_L, (size_t)fs->factBufferSize_M)));
1178: fs->spsvBuffer_L = fs->factBuffer_M;
1179: PetscCallCUDA(cudaMalloc((void **)&fs->spsvBuffer_Lt, fs->spsvBufferSize_Lt));
1180: } else {
1181: PetscCallCUDA(cudaMalloc((void **)&fs->factBuffer_M, PetscMax(fs->spsvBufferSize_Lt, (size_t)fs->factBufferSize_M)));
1182: fs->spsvBuffer_Lt = fs->factBuffer_M;
1183: PetscCallCUDA(cudaMalloc((void **)&fs->spsvBuffer_L, fs->spsvBufferSize_L));
1184: }
1186: /* ========================================================================== */
1187: /* Perform analysis of ic0 on M */
1188: /* The lower triangular part of M has the same sparsity pattern as L */
1189: /* ========================================================================== */
1190: int structural_zero;
1191: cusparseStatus_t status;
1193: fs->policy_M = CUSPARSE_SOLVE_POLICY_USE_LEVEL;
1194: if (m) PetscCallCUSPARSE(cusparseXcsric02_analysis(fs->handle, m, nz, fs->matDescr_M, fs->csrVal, fs->csrRowPtr32, fs->csrColIdx32, fs->ic0Info_M, fs->policy_M, fs->factBuffer_M));
1195: if (PetscDefined(USE_DEBUG)) {
1196: /* cusparseXcsric02_zeroPivot() is a blocking call. It calls cudaDeviceSynchronize() to make sure all previous kernels are done. */
1197: status = cusparseXcsric02_zeroPivot(fs->handle, fs->ic0Info_M, &structural_zero);
1198: PetscCheck(CUSPARSE_STATUS_ZERO_PIVOT != status, PETSC_COMM_SELF, PETSC_ERR_USER_INPUT, "Structural zero pivot detected in csric02: A(%d,%d) is missing", structural_zero, structural_zero);
1199: }
1201: /* Estimate FLOPs of the numeric factorization */
1202: {
1203: Mat_SeqAIJ *Aseq = (Mat_SeqAIJ *)A->data;
1204: PetscInt *Ai, nzRow, nzLeft;
1205: PetscLogDouble flops = 0.0;
1207: Ai = Aseq->i;
1208: for (PetscInt i = 0; i < m; i++) {
1209: nzRow = Ai[i + 1] - Ai[i];
1210: if (nzRow > 1) {
1211: /* We want to eliminate nonzeros left to the diagonal one by one. Assume each time, nonzeros right
1212: and include the eliminated one will be updated, which incurs a multiplication and an addition.
1213: */
1214: nzLeft = (nzRow - 1) / 2;
1215: flops += nzLeft * (2.0 * nzRow - nzLeft + 1);
1216: }
1217: }
1218: fs->numericFactFlops = flops;
1219: }
1220: fact->ops->choleskyfactornumeric = MatICCFactorNumeric_SeqAIJCUSPARSE_ICC0;
1221: PetscFunctionReturn(PETSC_SUCCESS);
1222: }
1224: static PetscErrorCode MatLUFactorNumeric_SeqAIJCUSPARSE(Mat B, Mat A, const MatFactorInfo *info)
1225: {
1226: // use_cpu_solve is a field in Mat_SeqAIJCUSPARSE. B, a factored matrix, uses Mat_SeqAIJCUSPARSETriFactors.
1227: Mat_SeqAIJCUSPARSE *cusparsestruct = static_cast<Mat_SeqAIJCUSPARSE *>(A->spptr);
1229: PetscFunctionBegin;
1230: PetscCall(MatSeqAIJCUSPARSECopyFromGPU(A));
1231: PetscCall(MatLUFactorNumeric_SeqAIJ(B, A, info));
1232: B->offloadmask = PETSC_OFFLOAD_CPU;
1234: if (!cusparsestruct->use_cpu_solve) {
1235: B->ops->solve = MatSolve_SeqAIJCUSPARSE_LU;
1236: B->ops->solvetranspose = MatSolveTranspose_SeqAIJCUSPARSE_LU;
1237: }
1238: B->ops->matsolve = NULL;
1239: B->ops->matsolvetranspose = NULL;
1241: /* get the triangular factors */
1242: if (!cusparsestruct->use_cpu_solve) PetscCall(MatSeqAIJCUSPARSEILUAnalysisAndCopyToGPU(B));
1243: PetscFunctionReturn(PETSC_SUCCESS);
1244: }
1246: static PetscErrorCode MatLUFactorSymbolic_SeqAIJCUSPARSE(Mat B, Mat A, IS isrow, IS iscol, const MatFactorInfo *info)
1247: {
1248: Mat_SeqAIJCUSPARSETriFactors *cusparseTriFactors = static_cast<Mat_SeqAIJCUSPARSETriFactors *>(B->spptr);
1250: PetscFunctionBegin;
1251: PetscCall(MatSeqAIJCUSPARSETriFactors_Reset(&cusparseTriFactors));
1252: PetscCall(MatLUFactorSymbolic_SeqAIJ(B, A, isrow, iscol, info));
1253: B->ops->lufactornumeric = MatLUFactorNumeric_SeqAIJCUSPARSE;
1254: PetscFunctionReturn(PETSC_SUCCESS);
1255: }
1257: static PetscErrorCode MatILUFactorSymbolic_SeqAIJCUSPARSE(Mat B, Mat A, IS isrow, IS iscol, const MatFactorInfo *info)
1258: {
1259: Mat_SeqAIJCUSPARSETriFactors *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors *)B->spptr;
1261: PetscFunctionBegin;
1262: PetscBool row_identity = PETSC_FALSE, col_identity = PETSC_FALSE;
1263: if (!info->factoronhost) {
1264: PetscCall(ISIdentity(isrow, &row_identity));
1265: PetscCall(ISIdentity(iscol, &col_identity));
1266: }
1267: if (!info->levels && row_identity && col_identity) {
1268: PetscCall(MatILUFactorSymbolic_SeqAIJCUSPARSE_ILU0(B, A, isrow, iscol, info));
1269: } else {
1270: PetscCall(MatSeqAIJCUSPARSETriFactors_Reset(&cusparseTriFactors));
1271: PetscCall(MatILUFactorSymbolic_SeqAIJ(B, A, isrow, iscol, info));
1272: B->ops->lufactornumeric = MatLUFactorNumeric_SeqAIJCUSPARSE;
1273: }
1274: PetscFunctionReturn(PETSC_SUCCESS);
1275: }
1277: static PetscErrorCode MatICCFactorSymbolic_SeqAIJCUSPARSE(Mat B, Mat A, IS perm, const MatFactorInfo *info)
1278: {
1279: Mat_SeqAIJCUSPARSETriFactors *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors *)B->spptr;
1281: PetscFunctionBegin;
1282: PetscBool perm_identity = PETSC_FALSE;
1283: if (!info->factoronhost) PetscCall(ISIdentity(perm, &perm_identity));
1284: if (!info->levels && perm_identity) {
1285: PetscCall(MatICCFactorSymbolic_SeqAIJCUSPARSE_ICC0(B, A, perm, info));
1286: } else {
1287: PetscCall(MatSeqAIJCUSPARSETriFactors_Reset(&cusparseTriFactors));
1288: PetscCall(MatICCFactorSymbolic_SeqAIJ(B, A, perm, info));
1289: B->ops->choleskyfactornumeric = MatCholeskyFactorNumeric_SeqAIJCUSPARSE;
1290: }
1291: PetscFunctionReturn(PETSC_SUCCESS);
1292: }
1294: static PetscErrorCode MatCholeskyFactorSymbolic_SeqAIJCUSPARSE(Mat B, Mat A, IS perm, const MatFactorInfo *info)
1295: {
1296: Mat_SeqAIJCUSPARSETriFactors *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors *)B->spptr;
1298: PetscFunctionBegin;
1299: PetscCall(MatSeqAIJCUSPARSETriFactors_Reset(&cusparseTriFactors));
1300: PetscCall(MatCholeskyFactorSymbolic_SeqAIJ(B, A, perm, info));
1301: B->ops->choleskyfactornumeric = MatCholeskyFactorNumeric_SeqAIJCUSPARSE;
1302: PetscFunctionReturn(PETSC_SUCCESS);
1303: }
1305: static PetscErrorCode MatFactorGetSolverType_seqaij_cusparse(Mat, MatSolverType *type)
1306: {
1307: PetscFunctionBegin;
1308: *type = MATSOLVERCUSPARSE;
1309: PetscFunctionReturn(PETSC_SUCCESS);
1310: }
1312: /*MC
1313: MATSOLVERCUSPARSE = "cusparse" - A matrix type providing triangular solvers for seq matrices
1314: on a single GPU of type, `MATSEQAIJCUSPARSE`. Currently supported
1315: algorithms are ILU(k) and ICC(k). Typically, deeper factorizations (larger k) results in poorer
1316: performance in the triangular solves. Full LU, and Cholesky decompositions can be solved through the
1317: CuSPARSE triangular solve algorithm. However, the performance can be quite poor and thus these
1318: algorithms are not recommended. This class does NOT support direct solver operations.
1320: Level: beginner
1322: .seealso: [](ch_matrices), `Mat`, `MATSEQAIJCUSPARSE`, `PCFactorSetMatSolverType()`, `MatSolverType`, `MatCreateSeqAIJCUSPARSE()`,
1323: `MATAIJCUSPARSE`, `MatCreateAIJCUSPARSE()`, `MatCUSPARSESetFormat()`, `MatCUSPARSEStorageFormat`, `MatCUSPARSEFormatOperation`
1324: M*/
1326: PETSC_EXTERN PetscErrorCode MatGetFactor_seqaijcusparse_cusparse(Mat A, MatFactorType ftype, Mat *B)
1327: {
1328: PetscInt n = A->rmap->n;
1330: PetscFunctionBegin;
1331: PetscCall(MatCreate(PetscObjectComm((PetscObject)A), B));
1332: PetscCall(MatSetSizes(*B, n, n, n, n));
1333: (*B)->factortype = ftype; // factortype makes MatSetType() allocate spptr of type Mat_SeqAIJCUSPARSETriFactors
1334: PetscCall(MatSetType(*B, MATSEQAIJCUSPARSE));
1336: if (A->boundtocpu && A->bindingpropagates) PetscCall(MatBindToCPU(*B, PETSC_TRUE));
1337: if (ftype == MAT_FACTOR_LU || ftype == MAT_FACTOR_ILU || ftype == MAT_FACTOR_ILUDT) {
1338: PetscCall(MatSetBlockSizesFromMats(*B, A, A));
1339: if (!A->boundtocpu) {
1340: (*B)->ops->ilufactorsymbolic = MatILUFactorSymbolic_SeqAIJCUSPARSE;
1341: (*B)->ops->lufactorsymbolic = MatLUFactorSymbolic_SeqAIJCUSPARSE;
1342: } else {
1343: (*B)->ops->ilufactorsymbolic = MatILUFactorSymbolic_SeqAIJ;
1344: (*B)->ops->lufactorsymbolic = MatLUFactorSymbolic_SeqAIJ;
1345: }
1346: PetscCall(PetscStrallocpy(MATORDERINGND, (char **)&(*B)->preferredordering[MAT_FACTOR_LU]));
1347: PetscCall(PetscStrallocpy(MATORDERINGNATURAL, (char **)&(*B)->preferredordering[MAT_FACTOR_ILU]));
1348: PetscCall(PetscStrallocpy(MATORDERINGNATURAL, (char **)&(*B)->preferredordering[MAT_FACTOR_ILUDT]));
1349: } else if (ftype == MAT_FACTOR_CHOLESKY || ftype == MAT_FACTOR_ICC) {
1350: if (!A->boundtocpu) {
1351: (*B)->ops->iccfactorsymbolic = MatICCFactorSymbolic_SeqAIJCUSPARSE;
1352: (*B)->ops->choleskyfactorsymbolic = MatCholeskyFactorSymbolic_SeqAIJCUSPARSE;
1353: } else {
1354: (*B)->ops->iccfactorsymbolic = MatICCFactorSymbolic_SeqAIJ;
1355: (*B)->ops->choleskyfactorsymbolic = MatCholeskyFactorSymbolic_SeqAIJ;
1356: }
1357: PetscCall(PetscStrallocpy(MATORDERINGND, (char **)&(*B)->preferredordering[MAT_FACTOR_CHOLESKY]));
1358: PetscCall(PetscStrallocpy(MATORDERINGNATURAL, (char **)&(*B)->preferredordering[MAT_FACTOR_ICC]));
1359: } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "Factor type not supported for CUSPARSE Matrix Types");
1361: PetscCall(MatSeqAIJSetPreallocation(*B, MAT_SKIP_ALLOCATION, NULL));
1362: (*B)->canuseordering = PETSC_TRUE;
1363: PetscCall(PetscObjectComposeFunction((PetscObject)*B, "MatFactorGetSolverType_C", MatFactorGetSolverType_seqaij_cusparse));
1364: PetscFunctionReturn(PETSC_SUCCESS);
1365: }
1367: static PetscErrorCode MatSeqAIJCUSPARSECopyFromGPU(Mat A)
1368: {
1369: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;
1370: Mat_SeqAIJCUSPARSE *cusp = (Mat_SeqAIJCUSPARSE *)A->spptr;
1371: Mat_SeqAIJCUSPARSETriFactors *fs = (Mat_SeqAIJCUSPARSETriFactors *)A->spptr;
1373: PetscFunctionBegin;
1374: if (A->offloadmask == PETSC_OFFLOAD_GPU) {
1375: PetscCall(PetscLogEventBegin(MAT_CUSPARSECopyFromGPU, A, 0, 0, 0));
1376: if (A->factortype == MAT_FACTOR_NONE) {
1377: CsrMatrix *matrix = (CsrMatrix *)cusp->mat->mat;
1378: PetscCallCUDA(cudaMemcpy(a->a, matrix->values->data().get(), a->nz * sizeof(PetscScalar), cudaMemcpyDeviceToHost));
1379: } else if (fs->csrVal) {
1380: /* We have a factorized matrix on device and are able to copy it to host */
1381: PetscCallCUDA(cudaMemcpy(a->a, fs->csrVal, a->nz * sizeof(PetscScalar), cudaMemcpyDeviceToHost));
1382: } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "No support for copying this type of factorized matrix from device to host");
1383: PetscCall(PetscLogGpuToCpu(a->nz * sizeof(PetscScalar)));
1384: PetscCall(PetscLogEventEnd(MAT_CUSPARSECopyFromGPU, A, 0, 0, 0));
1385: A->offloadmask = PETSC_OFFLOAD_BOTH;
1386: }
1387: PetscFunctionReturn(PETSC_SUCCESS);
1388: }
1390: /* Policy struct for MatSeqAIJCUSPARSE_CUPM shared template (CUDA specialisation) */
1391: struct MatSeqAIJCUSPARSE_Policy {
1392: typedef Mat_SeqAIJCUSPARSE mat_struct_type;
1393: typedef Mat_SeqAIJCUSPARSEMultStruct mult_struct_type;
1395: static int storage_format_csr() { return (int)MAT_CUSPARSE_CSR; }
1396: static int storage_format_ell() { return (int)MAT_CUSPARSE_ELL; }
1397: static int storage_format_hyb() { return (int)MAT_CUSPARSE_HYB; }
1399: static PetscErrorCode CopyToGPU(Mat A) { return MatSeqAIJCUSPARSECopyToGPU(A); }
1400: static PetscErrorCode CopyFromGPU(Mat A) { return MatSeqAIJCUSPARSECopyFromGPU(A); }
1401: static PetscErrorCode InvalidateTranspose(Mat A, PetscBool d) { return MatSeqAIJCUSPARSEInvalidateTranspose(A, d); }
1402: static PetscErrorCode ConvertFromSeqAIJ(Mat B, MatType t, MatReuse r, Mat *C) { return MatConvert_SeqAIJ_SeqAIJCUSPARSE(B, t, r, C); }
1403: static const char *mat_type_name;
1405: static PetscErrorCode Destroy(Mat A) { return MatSeqAIJCUSPARSE_Destroy(A); }
1406: static PetscErrorCode TriFactorsDestroy(void **spptr) { return MatSeqAIJCUSPARSETriFactors_Destroy((Mat_SeqAIJCUSPARSETriFactors **)spptr); }
1407: static const char *set_format_c;
1408: static const char *set_use_cpu_solve_c;
1409: static const char *product_seqdense_device_c;
1410: static const char *product_seqdense_c;
1411: static const char *product_self_c;
1412: static const char *seq_convert_hypre_c;
1414: static PetscErrorCode VecGetArrayRead(Vec v, const PetscScalar **a) { return VecCUDAGetArrayRead(v, a); }
1415: static PetscErrorCode VecRestoreArrayRead(Vec v, const PetscScalar **a) { return VecCUDARestoreArrayRead(v, a); }
1416: static PetscErrorCode VecGetArrayWrite(Vec v, PetscScalar **a) { return VecCUDAGetArrayWrite(v, a); }
1417: static PetscErrorCode VecRestoreArrayWrite(Vec v, PetscScalar **a) { return VecCUDARestoreArrayWrite(v, a); }
1418: };
1419: const char *MatSeqAIJCUSPARSE_Policy::mat_type_name = MATSEQAIJCUSPARSE;
1420: const char *MatSeqAIJCUSPARSE_Policy::set_format_c = "MatCUSPARSESetFormat_C";
1421: const char *MatSeqAIJCUSPARSE_Policy::set_use_cpu_solve_c = "MatCUSPARSESetUseCPUSolve_C";
1422: const char *MatSeqAIJCUSPARSE_Policy::product_seqdense_device_c = "MatProductSetFromOptions_seqaijcusparse_seqdensecuda_C";
1423: const char *MatSeqAIJCUSPARSE_Policy::product_seqdense_c = "MatProductSetFromOptions_seqaijcusparse_seqdense_C";
1424: const char *MatSeqAIJCUSPARSE_Policy::product_self_c = "MatProductSetFromOptions_seqaijcusparse_seqaijcusparse_C";
1425: const char *MatSeqAIJCUSPARSE_Policy::seq_convert_hypre_c = "MatConvert_seqaijcusparse_hypre_C";
1427: using MatSeqAIJCUSPARSE_CUPM_t = Petsc::mat::aij::cupm::impl::MatSeqAIJCUSPARSE_CUPM<Petsc::device::cupm::DeviceType::CUDA, MatSeqAIJCUSPARSE_Policy>;
1429: static PetscErrorCode MatSeqAIJGetArray_SeqAIJCUSPARSE(Mat A, PetscScalar *array[])
1430: {
1431: return MatSeqAIJCUSPARSE_CUPM_t::SeqAIJGetArray(A, array);
1432: }
1434: static PetscErrorCode MatSeqAIJRestoreArray_SeqAIJCUSPARSE(Mat A, PetscScalar *array[])
1435: {
1436: return MatSeqAIJCUSPARSE_CUPM_t::SeqAIJRestoreArray(A, array);
1437: }
1439: static PetscErrorCode MatSeqAIJGetArrayRead_SeqAIJCUSPARSE(Mat A, const PetscScalar *array[])
1440: {
1441: return MatSeqAIJCUSPARSE_CUPM_t::SeqAIJGetArrayRead(A, array);
1442: }
1444: static PetscErrorCode MatSeqAIJRestoreArrayRead_SeqAIJCUSPARSE(Mat A, const PetscScalar *array[])
1445: {
1446: return MatSeqAIJCUSPARSE_CUPM_t::SeqAIJRestoreArrayRead(A, array);
1447: }
1449: static PetscErrorCode MatSeqAIJGetArrayWrite_SeqAIJCUSPARSE(Mat A, PetscScalar *array[])
1450: {
1451: return MatSeqAIJCUSPARSE_CUPM_t::SeqAIJGetArrayWrite(A, array);
1452: }
1454: static PetscErrorCode MatSeqAIJRestoreArrayWrite_SeqAIJCUSPARSE(Mat A, PetscScalar *array[])
1455: {
1456: return MatSeqAIJCUSPARSE_CUPM_t::SeqAIJRestoreArrayWrite(A, array);
1457: }
1459: static PetscErrorCode MatSeqAIJGetCSRAndMemType_SeqAIJCUSPARSE(Mat A, const PetscInt **i, const PetscInt **j, PetscScalar **a, PetscMemType *mtype)
1460: {
1461: Mat_SeqAIJCUSPARSE *cusp;
1462: CsrMatrix *matrix;
1464: PetscFunctionBegin;
1465: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
1466: PetscCheck(A->factortype == MAT_FACTOR_NONE, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Not for factored matrix");
1467: cusp = static_cast<Mat_SeqAIJCUSPARSE *>(A->spptr);
1468: PetscCheck(cusp != NULL, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "cusp is NULL");
1469: matrix = (CsrMatrix *)cusp->mat->mat;
1471: if (i) *i = matrix->row_offsets->data().get();
1472: if (j) *j = matrix->column_indices->data().get();
1473: if (a) *a = matrix->values->data().get();
1474: if (mtype) *mtype = PETSC_MEMTYPE_CUDA;
1475: PetscFunctionReturn(PETSC_SUCCESS);
1476: }
1478: PETSC_INTERN PetscErrorCode MatSeqAIJCUSPARSECopyToGPU(Mat A)
1479: {
1480: Mat_SeqAIJCUSPARSE *cusparsestruct = (Mat_SeqAIJCUSPARSE *)A->spptr;
1481: Mat_SeqAIJCUSPARSEMultStruct *matstruct = cusparsestruct->mat;
1482: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;
1483: PetscInt m = A->rmap->n, *ii, *ridx, tmp;
1484: PetscBool both = PETSC_TRUE;
1486: PetscFunctionBegin;
1487: PetscCheck(!A->boundtocpu, PETSC_COMM_SELF, PETSC_ERR_GPU, "Cannot copy to GPU");
1488: if (A->offloadmask == PETSC_OFFLOAD_UNALLOCATED || A->offloadmask == PETSC_OFFLOAD_CPU) {
1489: if (A->nonzerostate == cusparsestruct->nonzerostate && cusparsestruct->format == MAT_CUSPARSE_CSR) { /* Copy values only */
1490: CsrMatrix *matrix;
1491: matrix = (CsrMatrix *)cusparsestruct->mat->mat;
1493: PetscCheck(!a->nz || a->a, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CSR values");
1494: PetscCall(PetscLogEventBegin(MAT_CUSPARSECopyToGPU, A, 0, 0, 0));
1495: matrix->values->assign(a->a, a->a + a->nz);
1496: PetscCallCUDA(WaitForCUDA());
1497: PetscCall(PetscLogCpuToGpu(a->nz * sizeof(PetscScalar)));
1498: PetscCall(PetscLogEventEnd(MAT_CUSPARSECopyToGPU, A, 0, 0, 0));
1499: PetscCall(MatSeqAIJCUSPARSEInvalidateTranspose(A, PETSC_FALSE));
1500: } else {
1501: PetscInt nnz;
1502: PetscCall(PetscLogEventBegin(MAT_CUSPARSECopyToGPU, A, 0, 0, 0));
1503: PetscCall(MatSeqAIJCUSPARSEMultStruct_Destroy(&cusparsestruct->mat, cusparsestruct->format));
1504: PetscCall(MatSeqAIJCUSPARSEInvalidateTranspose(A, PETSC_TRUE));
1505: delete cusparsestruct->workVector;
1506: delete cusparsestruct->rowoffsets_gpu;
1507: cusparsestruct->workVector = NULL;
1508: cusparsestruct->rowoffsets_gpu = NULL;
1509: try {
1510: if (a->compressedrow.use) {
1511: m = a->compressedrow.nrows;
1512: ii = a->compressedrow.i;
1513: ridx = a->compressedrow.rindex;
1514: } else {
1515: m = A->rmap->n;
1516: ii = a->i;
1517: ridx = NULL;
1518: }
1519: PetscCheck(ii, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CSR row data");
1520: if (!a->a) {
1521: nnz = ii[m];
1522: both = PETSC_FALSE;
1523: } else nnz = a->nz;
1524: PetscCheck(!nnz || a->j, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CSR column data");
1526: /* create cusparse matrix */
1527: cusparsestruct->nrows = m;
1528: matstruct = new Mat_SeqAIJCUSPARSEMultStruct;
1529: PetscCallCUSPARSE(cusparseCreateMatDescr(&matstruct->descr));
1530: PetscCallCUSPARSE(cusparseSetMatIndexBase(matstruct->descr, CUSPARSE_INDEX_BASE_ZERO));
1531: PetscCallCUSPARSE(cusparseSetMatType(matstruct->descr, CUSPARSE_MATRIX_TYPE_GENERAL));
1533: PetscCallCUDA(cudaMalloc((void **)&matstruct->alpha_one, sizeof(PetscScalar)));
1534: PetscCallCUDA(cudaMalloc((void **)&matstruct->beta_zero, sizeof(PetscScalar)));
1535: PetscCallCUDA(cudaMalloc((void **)&matstruct->beta_one, sizeof(PetscScalar)));
1536: PetscCallCUDA(cudaMemcpy(matstruct->alpha_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
1537: PetscCallCUDA(cudaMemcpy(matstruct->beta_zero, &PETSC_CUSPARSE_ZERO, sizeof(PetscScalar), cudaMemcpyHostToDevice));
1538: PetscCallCUDA(cudaMemcpy(matstruct->beta_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
1539: PetscCallCUSPARSE(cusparseSetPointerMode(cusparsestruct->handle, CUSPARSE_POINTER_MODE_DEVICE));
1541: /* Build a hybrid/ellpack matrix if this option is chosen for the storage */
1542: if (cusparsestruct->format == MAT_CUSPARSE_CSR) {
1543: /* set the matrix */
1544: CsrMatrix *mat = new CsrMatrix;
1545: mat->num_rows = m;
1546: mat->num_cols = A->cmap->n;
1547: mat->num_entries = nnz;
1548: PetscCallCXX(mat->row_offsets = new THRUSTINTARRAY(m + 1));
1549: mat->row_offsets->assign(ii, ii + m + 1);
1550: PetscCallCXX(mat->column_indices = new THRUSTINTARRAY(nnz));
1551: mat->column_indices->assign(a->j, a->j + nnz);
1553: PetscCallCXX(mat->values = new THRUSTARRAY(nnz));
1554: if (a->a) mat->values->assign(a->a, a->a + nnz);
1556: /* assign the pointer */
1557: matstruct->mat = mat;
1558: if (mat->num_rows) { /* cusparse errors on empty matrices! */
1559: PetscCallCUSPARSE(cusparseCreateCsr(&matstruct->matDescr, mat->num_rows, mat->num_cols, mat->num_entries, mat->row_offsets->data().get(), mat->column_indices->data().get(), mat->values->data().get(), csrRowOffsetsType, csrColIndType, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
1560: }
1561: } else if (cusparsestruct->format == MAT_CUSPARSE_ELL || cusparsestruct->format == MAT_CUSPARSE_HYB) {
1562: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MAT_CUSPARSE_ELL and MAT_CUSPARSE_HYB are not supported since CUDA-11.0");
1563: }
1565: /* assign the compressed row indices */
1566: if (a->compressedrow.use) {
1567: PetscCallCXX(cusparsestruct->workVector = new THRUSTARRAY(m));
1568: PetscCallCXX(matstruct->cprowIndices = new THRUSTINTARRAY(m));
1569: matstruct->cprowIndices->assign(ridx, ridx + m);
1570: tmp = m;
1571: } else {
1572: cusparsestruct->workVector = NULL;
1573: matstruct->cprowIndices = NULL;
1574: tmp = 0;
1575: }
1576: PetscCall(PetscLogCpuToGpu(((m + 1) + (a->nz)) * sizeof(int) + tmp * sizeof(PetscInt) + (3 + (a->nz)) * sizeof(PetscScalar)));
1578: /* assign the pointer */
1579: cusparsestruct->mat = matstruct;
1580: } catch (char *ex) {
1581: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_LIB, "CUSPARSE error: %s", ex);
1582: }
1583: PetscCallCUDA(WaitForCUDA());
1584: PetscCall(PetscLogEventEnd(MAT_CUSPARSECopyToGPU, A, 0, 0, 0));
1585: cusparsestruct->nonzerostate = A->nonzerostate;
1586: }
1587: if (both) A->offloadmask = PETSC_OFFLOAD_BOTH;
1588: }
1589: PetscFunctionReturn(PETSC_SUCCESS);
1590: }
1592: struct VecCUDAPlusEquals {
1593: template <typename Tuple>
1594: __host__ __device__ void operator()(Tuple t)
1595: {
1596: thrust::get<1>(t) = thrust::get<1>(t) + thrust::get<0>(t);
1597: }
1598: };
1600: struct VecCUDAEquals {
1601: template <typename Tuple>
1602: __host__ __device__ void operator()(Tuple t)
1603: {
1604: thrust::get<1>(t) = thrust::get<0>(t);
1605: }
1606: };
1608: struct VecCUDAEqualsReverse {
1609: template <typename Tuple>
1610: __host__ __device__ void operator()(Tuple t)
1611: {
1612: thrust::get<0>(t) = thrust::get<1>(t);
1613: }
1614: };
1616: struct MatProductCtx_MatMatCusparse {
1617: PetscBool cisdense;
1618: PetscScalar *Bt;
1619: Mat X;
1620: PetscBool reusesym; /* Cusparse does not have split symbolic and numeric phases for sparse matmat operations */
1621: PetscLogDouble flops;
1622: CsrMatrix *Bcsr;
1624: cusparseSpMatDescr_t matSpBDescr;
1625: PetscBool initialized; /* C = alpha op(A) op(B) + beta C */
1626: cusparseDnMatDescr_t matBDescr;
1627: cusparseDnMatDescr_t matCDescr;
1628: PetscInt Blda, Clda; /* Record leading dimensions of B and C here to detect changes*/
1629: void *dBuffer4;
1630: void *dBuffer5;
1631: size_t mmBufferSize;
1632: void *mmBuffer;
1633: void *mmBuffer2; /* SpGEMM WorkEstimation buffer */
1634: cusparseSpGEMMDescr_t spgemmDesc;
1635: };
1637: static PetscErrorCode MatProductCtxDestroy_MatMatCusparse(PetscCtxRt data)
1638: {
1639: MatProductCtx_MatMatCusparse *mmdata = *(MatProductCtx_MatMatCusparse **)data;
1641: PetscFunctionBegin;
1642: PetscCallCUDA(cudaFree(mmdata->Bt));
1643: delete mmdata->Bcsr;
1644: if (mmdata->matSpBDescr) PetscCallCUSPARSE(cusparseDestroySpMat(mmdata->matSpBDescr));
1645: if (mmdata->matBDescr) PetscCallCUSPARSE(cusparseDestroyDnMat(mmdata->matBDescr));
1646: if (mmdata->matCDescr) PetscCallCUSPARSE(cusparseDestroyDnMat(mmdata->matCDescr));
1647: if (mmdata->spgemmDesc) PetscCallCUSPARSE(cusparseSpGEMM_destroyDescr(mmdata->spgemmDesc));
1648: PetscCallCUDA(cudaFree(mmdata->dBuffer4));
1649: PetscCallCUDA(cudaFree(mmdata->dBuffer5));
1650: PetscCallCUDA(cudaFree(mmdata->mmBuffer));
1651: PetscCallCUDA(cudaFree(mmdata->mmBuffer2));
1652: PetscCall(MatDestroy(&mmdata->X));
1653: PetscCall(PetscFree(mmdata));
1654: PetscFunctionReturn(PETSC_SUCCESS);
1655: }
1657: #include <../src/mat/impls/dense/seq/dense.h>
1659: static PetscErrorCode MatProductNumeric_SeqAIJCUSPARSE_SeqDENSECUDA(Mat C)
1660: {
1661: Mat_Product *product = C->product;
1662: Mat A, B;
1663: PetscInt m, n, blda, clda;
1664: PetscBool flg, biscuda, compressed;
1665: Mat_SeqAIJ *a;
1666: Mat_SeqAIJCUSPARSE *cusp;
1667: cusparseOperation_t opA;
1668: const PetscScalar *barray;
1669: PetscScalar *carray;
1670: MatProductCtx_MatMatCusparse *mmdata;
1671: Mat_SeqAIJCUSPARSEMultStruct *mat;
1672: CsrMatrix *csrmat;
1674: PetscFunctionBegin;
1675: MatCheckProduct(C, 1);
1676: PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data empty");
1677: mmdata = (MatProductCtx_MatMatCusparse *)product->data;
1678: A = product->A;
1679: B = product->B;
1680: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJCUSPARSE, &flg));
1681: PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
1682: /* currently CopyToGpu does not copy if the matrix is bound to CPU
1683: Instead of silently accepting the wrong answer, I prefer to raise the error */
1684: PetscCheck(!A->boundtocpu, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Cannot bind to CPU a CUSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
1685: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
1686: a = (Mat_SeqAIJ *)A->data;
1687: cusp = (Mat_SeqAIJCUSPARSE *)A->spptr;
1688: switch (product->type) {
1689: case MATPRODUCT_AB:
1690: case MATPRODUCT_PtAP:
1691: mat = cusp->mat;
1692: opA = CUSPARSE_OPERATION_NON_TRANSPOSE;
1693: m = A->rmap->n;
1694: n = B->cmap->n;
1695: break;
1696: case MATPRODUCT_AtB:
1697: if (!A->form_explicit_transpose) {
1698: mat = cusp->mat;
1699: opA = CUSPARSE_OPERATION_TRANSPOSE;
1700: } else {
1701: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(A));
1702: mat = cusp->matTranspose;
1703: opA = CUSPARSE_OPERATION_NON_TRANSPOSE;
1704: }
1705: m = A->cmap->n;
1706: n = B->cmap->n;
1707: break;
1708: case MATPRODUCT_ABt:
1709: case MATPRODUCT_RARt:
1710: mat = cusp->mat;
1711: opA = CUSPARSE_OPERATION_NON_TRANSPOSE;
1712: m = A->rmap->n;
1713: n = B->rmap->n;
1714: break;
1715: default:
1716: SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
1717: }
1718: PetscCheck(mat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing Mat_SeqAIJCUSPARSEMultStruct");
1719: csrmat = (CsrMatrix *)mat->mat;
1720: /* when the rows of A are compressed on the device, csrmat holds only the nonempty rows of A, so the
1721: SpMM descriptor below must be built with the full row offsets instead of those of csrmat */
1722: compressed = (PetscBool)(mat->cprowIndices != NULL);
1723: /* if the user passed a CPU matrix, copy the data to the GPU */
1724: PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQDENSECUDA, &biscuda));
1725: if (!biscuda) PetscCall(MatConvert(B, MATSEQDENSECUDA, MAT_INPLACE_MATRIX, &B));
1726: PetscCall(MatDenseGetArrayReadAndMemType(B, &barray, nullptr));
1728: PetscCall(MatDenseGetLDA(B, &blda));
1729: if (product->type == MATPRODUCT_RARt || product->type == MATPRODUCT_PtAP) {
1730: PetscCall(MatDenseGetArrayWriteAndMemType(mmdata->X, &carray, nullptr));
1731: PetscCall(MatDenseGetLDA(mmdata->X, &clda));
1732: } else {
1733: PetscCall(MatDenseGetArrayWriteAndMemType(C, &carray, nullptr));
1734: PetscCall(MatDenseGetLDA(C, &clda));
1735: }
1737: PetscCall(PetscLogGpuTimeBegin());
1738: cusparseOperation_t opB = (product->type == MATPRODUCT_ABt || product->type == MATPRODUCT_RARt) ? CUSPARSE_OPERATION_TRANSPOSE : CUSPARSE_OPERATION_NON_TRANSPOSE;
1739: #if PETSC_PKG_CUDA_VERSION_GE(12, 4, 0)
1740: cusparseSpMatDescr_t &matADescr = mat->matDescr_SpMM[opA];
1741: #else
1742: /* mat->matDescr is also used by the SpGEMM code, which relies on its compressed dimensions, so when A is
1743: compressed the SpMM needs a descriptor of its own */
1744: cusparseSpMatDescr_t &matADescr = compressed ? mat->matDescr_SpMM[opA] : mat->matDescr;
1745: #endif
1747: /* (re)allocate mmBuffer if not initialized or LDAs are different */
1748: if (!mmdata->initialized || mmdata->Blda != blda || mmdata->Clda != clda) {
1749: size_t mmBufferSize;
1750: if (mmdata->initialized && mmdata->Blda != blda) {
1751: PetscCallCUSPARSE(cusparseDestroyDnMat(mmdata->matBDescr));
1752: mmdata->matBDescr = NULL;
1753: }
1754: if (!mmdata->matBDescr) {
1755: PetscCallCUSPARSE(cusparseCreateDnMat(&mmdata->matBDescr, B->rmap->n, B->cmap->n, blda, (void *)barray, cusparse_scalartype, CUSPARSE_ORDER_COL));
1756: mmdata->Blda = blda;
1757: }
1759: if (mmdata->initialized && mmdata->Clda != clda) {
1760: PetscCallCUSPARSE(cusparseDestroyDnMat(mmdata->matCDescr));
1761: mmdata->matCDescr = NULL;
1762: }
1763: if (!mmdata->matCDescr) { /* matCDescr is for C or mmdata->X */
1764: PetscCallCUSPARSE(cusparseCreateDnMat(&mmdata->matCDescr, m, n, clda, (void *)carray, cusparse_scalartype, CUSPARSE_ORDER_COL));
1765: mmdata->Clda = clda;
1766: }
1768: #if PETSC_PKG_CUDA_VERSION_GE(12, 4, 0) // tested up to 12.6.0
1769: if (matADescr) {
1770: PetscCallCUSPARSE(cusparseDestroySpMat(matADescr)); // Because I find I could not reuse matADescr. It could be a cusparse bug
1771: matADescr = NULL;
1772: }
1773: #endif
1775: if (!matADescr) {
1776: if (compressed) {
1777: if (!cusp->rowoffsets_gpu) { /* the full row offsets may be absent when we did not construct the transpose with csr2csc */
1778: cusp->rowoffsets_gpu = new THRUSTINTARRAY(A->rmap->n + 1);
1779: cusp->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
1780: PetscCall(PetscLogCpuToGpu((A->rmap->n + 1) * sizeof(PetscInt)));
1781: }
1782: PetscCallCUSPARSE(cusparseCreateCsr(&matADescr, A->rmap->n, csrmat->num_cols, csrmat->num_entries, cusp->rowoffsets_gpu->data().get(), csrmat->column_indices->data().get(), csrmat->values->data().get(), csrRowOffsetsType, csrColIndType, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
1783: } else {
1784: PetscCallCUSPARSE(cusparseCreateCsr(&matADescr, csrmat->num_rows, csrmat->num_cols, csrmat->num_entries, csrmat->row_offsets->data().get(), csrmat->column_indices->data().get(), csrmat->values->data().get(), csrRowOffsetsType, csrColIndType, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
1785: }
1786: }
1788: PetscCallCUSPARSE(cusparseSpMM_bufferSize(cusp->handle, opA, opB, mat->alpha_one, matADescr, mmdata->matBDescr, mat->beta_zero, mmdata->matCDescr, cusparse_scalartype, cusp->spmmAlg, &mmBufferSize));
1790: if ((mmdata->mmBuffer && mmdata->mmBufferSize < mmBufferSize) || !mmdata->mmBuffer) {
1791: PetscCallCUDA(cudaFree(mmdata->mmBuffer));
1792: PetscCallCUDA(cudaMalloc(&mmdata->mmBuffer, mmBufferSize));
1793: mmdata->mmBufferSize = mmBufferSize;
1794: }
1796: #if PETSC_PKG_CUDA_VERSION_GE(12, 4, 0) // the _preprocess was added in 11.2.1, but PETSc worked without it until 12.4.0
1797: PetscCallCUSPARSE(cusparseSpMM_preprocess(cusp->handle, opA, opB, mat->alpha_one, matADescr, mmdata->matBDescr, mat->beta_zero, mmdata->matCDescr, cusparse_scalartype, cusp->spmmAlg, mmdata->mmBuffer));
1798: #endif
1800: mmdata->initialized = PETSC_TRUE;
1801: } else {
1802: /* to be safe, always update pointers of the mats */
1803: PetscCallCUSPARSE(cusparseSpMatSetValues(matADescr, csrmat->values->data().get()));
1804: PetscCallCUSPARSE(cusparseDnMatSetValues(mmdata->matBDescr, (void *)barray));
1805: PetscCallCUSPARSE(cusparseDnMatSetValues(mmdata->matCDescr, (void *)carray));
1806: }
1808: /* do cusparseSpMM, which supports transpose on B */
1809: PetscCallCUSPARSE(cusparseSpMM(cusp->handle, opA, opB, mat->alpha_one, matADescr, mmdata->matBDescr, mat->beta_zero, mmdata->matCDescr, cusparse_scalartype, cusp->spmmAlg, mmdata->mmBuffer));
1811: PetscCall(PetscLogGpuTimeEnd());
1812: PetscCall(PetscLogGpuFlops(n * 2.0 * csrmat->num_entries));
1813: PetscCall(MatDenseRestoreArrayReadAndMemType(B, &barray));
1814: if (product->type == MATPRODUCT_RARt) {
1815: PetscCall(MatDenseRestoreArrayWriteAndMemType(mmdata->X, &carray));
1816: PetscCall(MatMatMultNumeric_SeqDenseCUDA_SeqDenseCUDA_Internal(B, mmdata->X, C, PETSC_FALSE, PETSC_FALSE));
1817: } else if (product->type == MATPRODUCT_PtAP) {
1818: PetscCall(MatDenseRestoreArrayWriteAndMemType(mmdata->X, &carray));
1819: PetscCall(MatMatMultNumeric_SeqDenseCUDA_SeqDenseCUDA_Internal(B, mmdata->X, C, PETSC_TRUE, PETSC_FALSE));
1820: } else {
1821: PetscCall(MatDenseRestoreArrayWriteAndMemType(C, &carray));
1822: }
1823: if (mmdata->cisdense) PetscCall(MatConvert(C, MATSEQDENSE, MAT_INPLACE_MATRIX, &C));
1824: if (!biscuda) PetscCall(MatConvert(B, MATSEQDENSE, MAT_INPLACE_MATRIX, &B));
1825: PetscFunctionReturn(PETSC_SUCCESS);
1826: }
1828: static PetscErrorCode MatProductSymbolic_SeqAIJCUSPARSE_SeqDENSECUDA(Mat C)
1829: {
1830: Mat_Product *product = C->product;
1831: Mat A, B;
1832: PetscInt m, n;
1833: PetscBool cisdense, flg;
1834: MatProductCtx_MatMatCusparse *mmdata;
1835: Mat_SeqAIJCUSPARSE *cusp;
1837: PetscFunctionBegin;
1838: MatCheckProduct(C, 1);
1839: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data not empty");
1840: A = product->A;
1841: B = product->B;
1842: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJCUSPARSE, &flg));
1843: PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
1844: cusp = (Mat_SeqAIJCUSPARSE *)A->spptr;
1845: PetscCheck(cusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
1846: switch (product->type) {
1847: case MATPRODUCT_AB:
1848: m = A->rmap->n;
1849: n = B->cmap->n;
1850: PetscCall(MatSetBlockSizesFromMats(C, A, B));
1851: break;
1852: case MATPRODUCT_AtB:
1853: m = A->cmap->n;
1854: n = B->cmap->n;
1855: if (A->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->rmap, A->cmap->bs));
1856: if (B->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->cmap, B->cmap->bs));
1857: break;
1858: case MATPRODUCT_ABt:
1859: m = A->rmap->n;
1860: n = B->rmap->n;
1861: if (A->rmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->rmap, A->rmap->bs));
1862: if (B->rmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->cmap, B->rmap->bs));
1863: break;
1864: case MATPRODUCT_PtAP:
1865: m = B->cmap->n;
1866: n = B->cmap->n;
1867: if (B->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->rmap, B->cmap->bs));
1868: if (B->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->cmap, B->cmap->bs));
1869: break;
1870: case MATPRODUCT_RARt:
1871: m = B->rmap->n;
1872: n = B->rmap->n;
1873: if (B->rmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->rmap, B->rmap->bs));
1874: if (B->rmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->cmap, B->rmap->bs));
1875: break;
1876: default:
1877: SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
1878: }
1879: PetscCall(MatSetSizes(C, m, n, m, n));
1880: /* if C is of type MATSEQDENSE (CPU), perform the operation on the GPU and then copy on the CPU */
1881: PetscCall(PetscObjectTypeCompare((PetscObject)C, MATSEQDENSE, &cisdense));
1882: PetscCall(MatSetType(C, MATSEQDENSECUDA));
1883: /* keep the VecType of B, e.g. VECKOKKOS from MatCreateDenseFromVecType(), but only if B is a MATSEQDENSECUDA like C, since the VecType of a MATSEQDENSE B, e.g. VECSTANDARD, is not valid for C */
1884: PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQDENSECUDA, &flg));
1885: if (flg) PetscCall(MatSetVecType(C, B->defaultvectype));
1887: /* product data */
1888: PetscCall(PetscNew(&mmdata));
1889: mmdata->cisdense = cisdense;
1890: /* for these products we need intermediate storage */
1891: if (product->type == MATPRODUCT_RARt || product->type == MATPRODUCT_PtAP) {
1892: PetscCall(MatCreate(PetscObjectComm((PetscObject)C), &mmdata->X));
1893: PetscCall(MatSetType(mmdata->X, MATSEQDENSECUDA));
1894: if (product->type == MATPRODUCT_RARt) { /* do not preallocate, since the first call to MatDenseCUDAGetArray will preallocate on the GPU for us */
1895: PetscCall(MatSetSizes(mmdata->X, A->rmap->n, B->rmap->n, A->rmap->n, B->rmap->n));
1896: } else {
1897: PetscCall(MatSetSizes(mmdata->X, A->rmap->n, B->cmap->n, A->rmap->n, B->cmap->n));
1898: }
1899: }
1900: C->product->data = mmdata;
1901: C->product->destroy = MatProductCtxDestroy_MatMatCusparse;
1903: C->ops->productnumeric = MatProductNumeric_SeqAIJCUSPARSE_SeqDENSECUDA;
1904: PetscFunctionReturn(PETSC_SUCCESS);
1905: }
1907: static PetscErrorCode MatProductNumeric_SeqAIJCUSPARSE_SeqAIJCUSPARSE(Mat C)
1908: {
1909: Mat_Product *product = C->product;
1910: Mat A, B;
1911: Mat_SeqAIJCUSPARSE *Acusp, *Bcusp, *Ccusp;
1912: Mat_SeqAIJ *c = (Mat_SeqAIJ *)C->data;
1913: Mat_SeqAIJCUSPARSEMultStruct *Amat, *Bmat, *Cmat;
1914: CsrMatrix *Acsr, *Bcsr, *Ccsr;
1915: PetscBool flg;
1916: MatProductType ptype;
1917: MatProductCtx_MatMatCusparse *mmdata;
1918: cusparseSpMatDescr_t BmatSpDescr;
1919: cusparseOperation_t opA = CUSPARSE_OPERATION_NON_TRANSPOSE, opB = CUSPARSE_OPERATION_NON_TRANSPOSE; /* cuSPARSE spgemm doesn't support transpose yet */
1921: PetscFunctionBegin;
1922: MatCheckProduct(C, 1);
1923: PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data empty");
1924: PetscCall(PetscObjectTypeCompare((PetscObject)C, MATSEQAIJCUSPARSE, &flg));
1925: PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for C of type %s", ((PetscObject)C)->type_name);
1926: mmdata = (MatProductCtx_MatMatCusparse *)C->product->data;
1927: A = product->A;
1928: B = product->B;
1929: if (mmdata->reusesym) { /* this happens when api_user is true, meaning that the matrix values have been already computed in the MatProductSymbolic phase */
1930: mmdata->reusesym = PETSC_FALSE;
1931: Ccusp = (Mat_SeqAIJCUSPARSE *)C->spptr;
1932: PetscCheck(Ccusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
1933: Cmat = Ccusp->mat;
1934: PetscCheck(Cmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C mult struct for product type %s", MatProductTypes[C->product->type]);
1935: Ccsr = (CsrMatrix *)Cmat->mat;
1936: PetscCheck(Ccsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C CSR struct");
1937: goto finalize;
1938: }
1939: if (!c->nz) goto finalize;
1940: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJCUSPARSE, &flg));
1941: PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
1942: PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQAIJCUSPARSE, &flg));
1943: PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for B of type %s", ((PetscObject)B)->type_name);
1944: PetscCheck(!A->boundtocpu, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONG, "Cannot bind to CPU a CUSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
1945: PetscCheck(!B->boundtocpu, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONG, "Cannot bind to CPU a CUSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
1946: Acusp = (Mat_SeqAIJCUSPARSE *)A->spptr;
1947: Bcusp = (Mat_SeqAIJCUSPARSE *)B->spptr;
1948: Ccusp = (Mat_SeqAIJCUSPARSE *)C->spptr;
1949: PetscCheck(Acusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
1950: PetscCheck(Bcusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
1951: PetscCheck(Ccusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
1952: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
1953: PetscCall(MatSeqAIJCUSPARSECopyToGPU(B));
1955: ptype = product->type;
1956: if (A->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_AtB) {
1957: ptype = MATPRODUCT_AB;
1958: PetscCheck(product->symbolic_used_the_fact_A_is_symmetric, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Symbolic should have been built using the fact that A is symmetric");
1959: }
1960: if (B->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_ABt) {
1961: ptype = MATPRODUCT_AB;
1962: PetscCheck(product->symbolic_used_the_fact_B_is_symmetric, PetscObjectComm((PetscObject)C), PETSC_ERR_PLIB, "Symbolic should have been built using the fact that B is symmetric");
1963: }
1964: switch (ptype) {
1965: case MATPRODUCT_AB:
1966: Amat = Acusp->mat;
1967: Bmat = Bcusp->mat;
1968: break;
1969: case MATPRODUCT_AtB:
1970: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(A));
1971: Amat = Acusp->matTranspose;
1972: Bmat = Bcusp->mat;
1973: break;
1974: case MATPRODUCT_ABt:
1975: Amat = Acusp->mat;
1976: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(B));
1977: Bmat = Bcusp->matTranspose;
1978: break;
1979: default:
1980: SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
1981: }
1982: Cmat = Ccusp->mat;
1983: PetscCheck(Amat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A mult struct for product type %s", MatProductTypes[ptype]);
1984: PetscCheck(Bmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B mult struct for product type %s", MatProductTypes[ptype]);
1985: PetscCheck(Cmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C mult struct for product type %s", MatProductTypes[ptype]);
1986: Acsr = (CsrMatrix *)Amat->mat;
1987: Bcsr = mmdata->Bcsr ? mmdata->Bcsr : (CsrMatrix *)Bmat->mat; /* B may be in compressed row storage */
1988: Ccsr = (CsrMatrix *)Cmat->mat;
1989: PetscCheck(Acsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A CSR struct");
1990: PetscCheck(Bcsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B CSR struct");
1991: PetscCheck(Ccsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C CSR struct");
1992: PetscCall(PetscLogGpuTimeBegin());
1993: BmatSpDescr = mmdata->Bcsr ? mmdata->matSpBDescr : Bmat->matDescr; /* B may be in compressed row storage */
1994: PetscCallCUSPARSE(cusparseSetPointerMode(Ccusp->handle, CUSPARSE_POINTER_MODE_DEVICE));
1995: PetscCallCUSPARSE(cusparseSpGEMM_compute(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &mmdata->mmBufferSize, mmdata->mmBuffer));
1996: PetscCallCUSPARSE(cusparseSpGEMM_copy(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc));
1997: PetscCall(PetscLogGpuFlops(mmdata->flops));
1998: PetscCallCUDA(WaitForCUDA());
1999: PetscCall(PetscLogGpuTimeEnd());
2000: C->offloadmask = PETSC_OFFLOAD_GPU;
2001: finalize:
2002: /* shorter version of MatAssemblyEnd_SeqAIJ */
2003: PetscCall(PetscInfo(C, "Matrix size: %" PetscInt_FMT " X %" PetscInt_FMT "; storage space: 0 unneeded, %" PetscInt_FMT " used\n", C->rmap->n, C->cmap->n, c->nz));
2004: PetscCall(PetscInfo(C, "Number of mallocs during MatSetValues() is 0\n"));
2005: PetscCall(PetscInfo(C, "Maximum nonzeros in any row is %" PetscInt_FMT "\n", c->rmax));
2006: c->reallocs = 0;
2007: C->info.mallocs += 0;
2008: C->info.nz_unneeded = 0;
2009: C->assembled = C->was_assembled = PETSC_TRUE;
2010: C->num_ass++;
2011: PetscFunctionReturn(PETSC_SUCCESS);
2012: }
2014: static PetscErrorCode MatProductSymbolic_SeqAIJCUSPARSE_SeqAIJCUSPARSE(Mat C)
2015: {
2016: Mat_Product *product = C->product;
2017: Mat A, B;
2018: Mat_SeqAIJCUSPARSE *Acusp, *Bcusp, *Ccusp;
2019: Mat_SeqAIJ *a, *b, *c;
2020: Mat_SeqAIJCUSPARSEMultStruct *Amat, *Bmat, *Cmat;
2021: CsrMatrix *Acsr, *Bcsr, *Ccsr;
2022: PetscInt i, j, m, n, k;
2023: PetscBool flg;
2024: MatProductType ptype;
2025: MatProductCtx_MatMatCusparse *mmdata;
2026: PetscLogDouble flops;
2027: PetscBool biscompressed, ciscompressed;
2028: int64_t C_num_rows1, C_num_cols1, C_nnz1;
2029: cusparseSpMatDescr_t BmatSpDescr;
2030: cusparseOperation_t opA = CUSPARSE_OPERATION_NON_TRANSPOSE, opB = CUSPARSE_OPERATION_NON_TRANSPOSE; /* cuSPARSE spgemm doesn't support transpose yet */
2031: size_t bufSize2;
2033: PetscFunctionBegin;
2034: MatCheckProduct(C, 1);
2035: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data not empty");
2036: A = product->A;
2037: B = product->B;
2038: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJCUSPARSE, &flg));
2039: PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
2040: PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQAIJCUSPARSE, &flg));
2041: PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for B of type %s", ((PetscObject)B)->type_name);
2042: a = (Mat_SeqAIJ *)A->data;
2043: b = (Mat_SeqAIJ *)B->data;
2044: /* product data */
2045: PetscCall(PetscNew(&mmdata));
2046: C->product->data = mmdata;
2047: C->product->destroy = MatProductCtxDestroy_MatMatCusparse;
2049: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
2050: PetscCall(MatSeqAIJCUSPARSECopyToGPU(B));
2051: Acusp = (Mat_SeqAIJCUSPARSE *)A->spptr; /* Access spptr after MatSeqAIJCUSPARSECopyToGPU, not before */
2052: Bcusp = (Mat_SeqAIJCUSPARSE *)B->spptr;
2053: PetscCheck(Acusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
2054: PetscCheck(Bcusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
2056: ptype = product->type;
2057: if (A->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_AtB) {
2058: ptype = MATPRODUCT_AB;
2059: product->symbolic_used_the_fact_A_is_symmetric = PETSC_TRUE;
2060: }
2061: if (B->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_ABt) {
2062: ptype = MATPRODUCT_AB;
2063: product->symbolic_used_the_fact_B_is_symmetric = PETSC_TRUE;
2064: }
2065: biscompressed = PETSC_FALSE;
2066: ciscompressed = PETSC_FALSE;
2067: switch (ptype) {
2068: case MATPRODUCT_AB:
2069: m = A->rmap->n;
2070: n = B->cmap->n;
2071: k = A->cmap->n;
2072: Amat = Acusp->mat;
2073: Bmat = Bcusp->mat;
2074: if (a->compressedrow.use) ciscompressed = PETSC_TRUE;
2075: if (b->compressedrow.use) biscompressed = PETSC_TRUE;
2076: break;
2077: case MATPRODUCT_AtB:
2078: m = A->cmap->n;
2079: n = B->cmap->n;
2080: k = A->rmap->n;
2081: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(A));
2082: Amat = Acusp->matTranspose;
2083: Bmat = Bcusp->mat;
2084: if (b->compressedrow.use) biscompressed = PETSC_TRUE;
2085: break;
2086: case MATPRODUCT_ABt:
2087: m = A->rmap->n;
2088: n = B->rmap->n;
2089: k = A->cmap->n;
2090: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(B));
2091: Amat = Acusp->mat;
2092: Bmat = Bcusp->matTranspose;
2093: if (a->compressedrow.use) ciscompressed = PETSC_TRUE;
2094: break;
2095: default:
2096: SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
2097: }
2099: /* create cusparse matrix */
2100: PetscCall(MatSetSizes(C, m, n, m, n));
2101: PetscCall(MatSetType(C, MATSEQAIJCUSPARSE));
2102: c = (Mat_SeqAIJ *)C->data;
2103: Ccusp = (Mat_SeqAIJCUSPARSE *)C->spptr;
2104: Cmat = new Mat_SeqAIJCUSPARSEMultStruct;
2105: Ccsr = new CsrMatrix;
2107: c->compressedrow.use = ciscompressed;
2108: if (c->compressedrow.use) { /* if a is in compressed row, than c will be in compressed row format */
2109: c->compressedrow.nrows = a->compressedrow.nrows;
2110: PetscCall(PetscMalloc2(c->compressedrow.nrows + 1, &c->compressedrow.i, c->compressedrow.nrows, &c->compressedrow.rindex));
2111: PetscCall(PetscArraycpy(c->compressedrow.rindex, a->compressedrow.rindex, c->compressedrow.nrows));
2112: Ccusp->workVector = new THRUSTARRAY(c->compressedrow.nrows);
2113: Cmat->cprowIndices = new THRUSTINTARRAY(c->compressedrow.nrows);
2114: Cmat->cprowIndices->assign(c->compressedrow.rindex, c->compressedrow.rindex + c->compressedrow.nrows);
2115: } else {
2116: c->compressedrow.nrows = 0;
2117: c->compressedrow.i = NULL;
2118: c->compressedrow.rindex = NULL;
2119: Ccusp->workVector = NULL;
2120: Cmat->cprowIndices = NULL;
2121: }
2122: Ccusp->nrows = ciscompressed ? c->compressedrow.nrows : m;
2123: Ccusp->mat = Cmat;
2124: Ccusp->mat->mat = Ccsr;
2125: Ccsr->num_rows = Ccusp->nrows;
2126: Ccsr->num_cols = n;
2127: Ccsr->row_offsets = new THRUSTINTARRAY(Ccusp->nrows + 1);
2128: PetscCallCUSPARSE(cusparseCreateMatDescr(&Cmat->descr));
2129: PetscCallCUSPARSE(cusparseSetMatIndexBase(Cmat->descr, CUSPARSE_INDEX_BASE_ZERO));
2130: PetscCallCUSPARSE(cusparseSetMatType(Cmat->descr, CUSPARSE_MATRIX_TYPE_GENERAL));
2131: PetscCallCUDA(cudaMalloc((void **)&Cmat->alpha_one, sizeof(PetscScalar)));
2132: PetscCallCUDA(cudaMalloc((void **)&Cmat->beta_zero, sizeof(PetscScalar)));
2133: PetscCallCUDA(cudaMalloc((void **)&Cmat->beta_one, sizeof(PetscScalar)));
2134: PetscCallCUDA(cudaMemcpy(Cmat->alpha_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
2135: PetscCallCUDA(cudaMemcpy(Cmat->beta_zero, &PETSC_CUSPARSE_ZERO, sizeof(PetscScalar), cudaMemcpyHostToDevice));
2136: PetscCallCUDA(cudaMemcpy(Cmat->beta_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
2137: if (!Ccsr->num_rows || !Ccsr->num_cols || !a->nz || !b->nz) { /* cusparse raise errors in different calls when matrices have zero rows/columns! */
2138: PetscCallThrust(thrust::fill(thrust::device, Ccsr->row_offsets->begin(), Ccsr->row_offsets->end(), 0));
2139: c->nz = 0;
2140: Ccsr->column_indices = new THRUSTINTARRAY(c->nz);
2141: Ccsr->values = new THRUSTARRAY(c->nz);
2142: goto finalizesym;
2143: }
2145: PetscCheck(Amat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A mult struct for product type %s", MatProductTypes[ptype]);
2146: PetscCheck(Bmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B mult struct for product type %s", MatProductTypes[ptype]);
2147: Acsr = (CsrMatrix *)Amat->mat;
2148: if (!biscompressed) {
2149: Bcsr = (CsrMatrix *)Bmat->mat;
2150: BmatSpDescr = Bmat->matDescr;
2151: } else { /* we need to use row offsets for the full matrix */
2152: CsrMatrix *cBcsr = (CsrMatrix *)Bmat->mat;
2153: Bcsr = new CsrMatrix;
2154: Bcsr->num_rows = B->rmap->n;
2155: Bcsr->num_cols = cBcsr->num_cols;
2156: Bcsr->num_entries = cBcsr->num_entries;
2157: Bcsr->column_indices = cBcsr->column_indices;
2158: Bcsr->values = cBcsr->values;
2159: if (!Bcusp->rowoffsets_gpu) {
2160: Bcusp->rowoffsets_gpu = new THRUSTINTARRAY(B->rmap->n + 1);
2161: Bcusp->rowoffsets_gpu->assign(b->i, b->i + B->rmap->n + 1);
2162: PetscCall(PetscLogCpuToGpu((B->rmap->n + 1) * sizeof(PetscInt)));
2163: }
2164: Bcsr->row_offsets = Bcusp->rowoffsets_gpu;
2165: mmdata->Bcsr = Bcsr;
2166: if (Bcsr->num_rows && Bcsr->num_cols) {
2167: PetscCallCUSPARSE(cusparseCreateCsr(&mmdata->matSpBDescr, Bcsr->num_rows, Bcsr->num_cols, Bcsr->num_entries, Bcsr->row_offsets->data().get(), Bcsr->column_indices->data().get(), Bcsr->values->data().get(), csrRowOffsetsType, csrColIndType, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
2168: }
2169: BmatSpDescr = mmdata->matSpBDescr;
2170: }
2171: PetscCheck(Acsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A CSR struct");
2172: PetscCheck(Bcsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B CSR struct");
2173: /* precompute flops count */
2174: if (ptype == MATPRODUCT_AB) {
2175: for (i = 0, flops = 0; i < A->rmap->n; i++) {
2176: const PetscInt st = a->i[i];
2177: const PetscInt en = a->i[i + 1];
2178: for (j = st; j < en; j++) {
2179: const PetscInt brow = a->j[j];
2180: flops += 2. * (b->i[brow + 1] - b->i[brow]);
2181: }
2182: }
2183: } else if (ptype == MATPRODUCT_AtB) {
2184: for (i = 0, flops = 0; i < A->rmap->n; i++) {
2185: const PetscInt anzi = a->i[i + 1] - a->i[i];
2186: const PetscInt bnzi = b->i[i + 1] - b->i[i];
2187: flops += (2. * anzi) * bnzi;
2188: }
2189: } else { /* TODO */
2190: flops = 0.;
2191: }
2193: mmdata->flops = flops;
2194: PetscCall(PetscLogGpuTimeBegin());
2196: PetscCallCUSPARSE(cusparseSetPointerMode(Ccusp->handle, CUSPARSE_POINTER_MODE_DEVICE));
2197: // cuda-12.2 requires non-null csrRowOffsets
2198: PetscCallCUSPARSE(cusparseCreateCsr(&Cmat->matDescr, Ccsr->num_rows, Ccsr->num_cols, 0, Ccsr->row_offsets->data().get(), NULL, NULL, csrRowOffsetsType, csrColIndType, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
2199: PetscCallCUSPARSE(cusparseSpGEMM_createDescr(&mmdata->spgemmDesc));
2200: // Note that cusparseSpGEMMreuse is deprecated in CUDA 13.2.1
2202: PetscCheck(!PetscDefined(USE_64BIT_INDICES) || PETSC_PKG_CUDA_VERSION_GE(13, 0, 0), PETSC_COMM_SELF, PETSC_ERR_SUP_SYS, "cusparseSpGEMM did not support 64-bit indices before CUDA 13.0. Update your CUDA installation.");
2203: /* ask bufferSize bytes for external memory */
2204: PetscCallCUSPARSE(cusparseSpGEMM_workEstimation(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufSize2, NULL));
2205: PetscCallCUDA(cudaMalloc((void **)&mmdata->mmBuffer2, bufSize2));
2206: /* inspect the matrices A and B to understand the memory requirement for the next step */
2207: PetscCallCUSPARSE(cusparseSpGEMM_workEstimation(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufSize2, mmdata->mmBuffer2));
2208: /* ask bufferSize again bytes for external memory */
2209: PetscCallCUSPARSE(cusparseSpGEMM_compute(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &mmdata->mmBufferSize, NULL));
2210: /* The CUSPARSE documentation is not clear, nor the API
2211: We need both buffers to perform the operations properly!
2212: mmdata->mmBuffer2 does not appear anywhere in the compute/copy API
2213: it only appears for the workEstimation stuff, but it seems it is needed in compute, so probably the address
2214: is stored in the descriptor! What a messy API... */
2215: PetscCallCUDA(cudaMalloc((void **)&mmdata->mmBuffer, mmdata->mmBufferSize));
2216: /* compute the intermediate product of A * B */
2217: PetscCallCUSPARSE(cusparseSpGEMM_compute(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &mmdata->mmBufferSize, mmdata->mmBuffer));
2218: /* get matrix C non-zero entries C_nnz1 */
2219: PetscCallCUSPARSE(cusparseSpMatGetSize(Cmat->matDescr, &C_num_rows1, &C_num_cols1, &C_nnz1));
2220: PetscCall(PetscIntCast(C_nnz1, &c->nz));
2221: PetscCall(PetscInfo(C, "Buffer sizes for type %s, result %" PetscInt_FMT " x %" PetscInt_FMT " (k %" PetscInt_FMT ", nzA %" PetscInt_FMT ", nzB %" PetscInt_FMT ", nzC %" PetscInt_FMT ") are: %ldKB %ldKB\n", MatProductTypes[ptype], m, n, k, a->nz, b->nz, c->nz, bufSize2 / 1024,
2222: mmdata->mmBufferSize / 1024));
2223: Ccsr->column_indices = new THRUSTINTARRAY(c->nz);
2224: PetscCallCUDA(cudaPeekAtLastError()); /* catch out of memory errors */
2225: Ccsr->values = new THRUSTARRAY(c->nz);
2226: PetscCallCUDA(cudaPeekAtLastError()); /* catch out of memory errors */
2227: if (c->nz) PetscCallCUSPARSE(cusparseCsrSetPointers(Cmat->matDescr, Ccsr->row_offsets->data().get(), Ccsr->column_indices->data().get(), Ccsr->values->data().get()));
2228: PetscCallCUSPARSE(cusparseSpGEMM_copy(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc));
2229: PetscCall(PetscLogGpuFlops(mmdata->flops));
2230: PetscCall(PetscLogGpuTimeEnd());
2231: finalizesym:
2232: c->free_a = PETSC_TRUE;
2233: PetscCall(PetscShmgetAllocateArray(c->nz, sizeof(PetscInt), (void **)&c->j));
2234: PetscCall(PetscShmgetAllocateArray(m + 1, sizeof(PetscInt), (void **)&c->i));
2235: c->free_ij = PETSC_TRUE;
2237: PetscInt *d_i = c->i;
2238: if (ciscompressed) d_i = c->compressedrow.i;
2239: PetscCallCUDA(cudaMemcpy(d_i, Ccsr->row_offsets->data().get(), Ccsr->row_offsets->size() * sizeof(PetscInt), cudaMemcpyDeviceToHost));
2240: PetscCallCUDA(cudaMemcpy(c->j, Ccsr->column_indices->data().get(), Ccsr->column_indices->size() * sizeof(PetscInt), cudaMemcpyDeviceToHost));
2241: if (ciscompressed) { /* need to expand host row offsets */
2242: PetscInt r = 0;
2243: c->i[0] = 0;
2244: for (k = 0; k < c->compressedrow.nrows; k++) {
2245: const PetscInt next = c->compressedrow.rindex[k];
2246: const PetscInt old = c->compressedrow.i[k];
2247: for (; r < next; r++) c->i[r + 1] = old;
2248: }
2249: for (; r < m; r++) c->i[r + 1] = c->compressedrow.i[c->compressedrow.nrows];
2250: }
2251: PetscCall(PetscLogGpuToCpu((Ccsr->column_indices->size() + Ccsr->row_offsets->size()) * sizeof(PetscInt)));
2252: PetscCall(PetscMalloc1(m, &c->ilen));
2253: PetscCall(PetscMalloc1(m, &c->imax));
2254: c->maxnz = c->nz;
2255: c->nonzerorowcnt = 0;
2256: c->rmax = 0;
2257: for (k = 0; k < m; k++) {
2258: const PetscInt nn = c->i[k + 1] - c->i[k];
2259: c->ilen[k] = c->imax[k] = nn;
2260: c->nonzerorowcnt += (PetscInt)!!nn;
2261: c->rmax = PetscMax(c->rmax, nn);
2262: }
2263: PetscCall(PetscMalloc1(c->nz, &c->a));
2264: Ccsr->num_entries = c->nz;
2266: C->nonzerostate++;
2267: PetscCall(PetscLayoutSetUp(C->rmap));
2268: PetscCall(PetscLayoutSetUp(C->cmap));
2269: Ccusp->nonzerostate = C->nonzerostate;
2270: C->offloadmask = PETSC_OFFLOAD_UNALLOCATED;
2271: C->preallocated = PETSC_TRUE;
2272: C->assembled = PETSC_FALSE;
2273: C->was_assembled = PETSC_FALSE;
2274: if (product->api_user && A->offloadmask == PETSC_OFFLOAD_BOTH && B->offloadmask == PETSC_OFFLOAD_BOTH) { /* flag the matrix C values as computed, so that the numeric phase will only call MatAssembly */
2275: mmdata->reusesym = PETSC_TRUE;
2276: C->offloadmask = PETSC_OFFLOAD_GPU;
2277: }
2278: C->ops->productnumeric = MatProductNumeric_SeqAIJCUSPARSE_SeqAIJCUSPARSE;
2279: PetscFunctionReturn(PETSC_SUCCESS);
2280: }
2282: PETSC_INTERN PetscErrorCode MatProductSetFromOptions_SeqAIJ_SeqDense(Mat);
2284: /* handles sparse or dense B */
2285: static PetscErrorCode MatProductSetFromOptions_SeqAIJCUSPARSE(Mat mat)
2286: {
2287: Mat_Product *product = mat->product;
2288: PetscBool isdense = PETSC_FALSE, Biscusp = PETSC_FALSE, Ciscusp = PETSC_TRUE;
2290: PetscFunctionBegin;
2291: MatCheckProduct(mat, 1);
2292: PetscCall(PetscObjectBaseTypeCompare((PetscObject)product->B, MATSEQDENSE, &isdense));
2293: if (!product->A->boundtocpu && !product->B->boundtocpu) PetscCall(PetscObjectTypeCompare((PetscObject)product->B, MATSEQAIJCUSPARSE, &Biscusp));
2294: if (product->type == MATPRODUCT_ABC) {
2295: Ciscusp = PETSC_FALSE;
2296: if (!product->C->boundtocpu) PetscCall(PetscObjectTypeCompare((PetscObject)product->C, MATSEQAIJCUSPARSE, &Ciscusp));
2297: }
2298: if (Biscusp && Ciscusp) { /* we can always select the CPU backend */
2299: PetscBool usecpu = PETSC_FALSE;
2300: switch (product->type) {
2301: case MATPRODUCT_AB:
2302: if (product->api_user) {
2303: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatMatMult", "Mat");
2304: PetscCall(PetscOptionsBool("-matmatmult_backend_cpu", "Use CPU code", "MatMatMult", usecpu, &usecpu, NULL));
2305: PetscOptionsEnd();
2306: } else {
2307: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_AB", "Mat");
2308: PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatMatMult", usecpu, &usecpu, NULL));
2309: PetscOptionsEnd();
2310: }
2311: break;
2312: case MATPRODUCT_AtB:
2313: if (product->api_user) {
2314: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatTransposeMatMult", "Mat");
2315: PetscCall(PetscOptionsBool("-mattransposematmult_backend_cpu", "Use CPU code", "MatTransposeMatMult", usecpu, &usecpu, NULL));
2316: PetscOptionsEnd();
2317: } else {
2318: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_AtB", "Mat");
2319: PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatTransposeMatMult", usecpu, &usecpu, NULL));
2320: PetscOptionsEnd();
2321: }
2322: break;
2323: case MATPRODUCT_PtAP:
2324: if (product->api_user) {
2325: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatPtAP", "Mat");
2326: PetscCall(PetscOptionsBool("-matptap_backend_cpu", "Use CPU code", "MatPtAP", usecpu, &usecpu, NULL));
2327: PetscOptionsEnd();
2328: } else {
2329: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_PtAP", "Mat");
2330: PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatPtAP", usecpu, &usecpu, NULL));
2331: PetscOptionsEnd();
2332: }
2333: break;
2334: case MATPRODUCT_RARt:
2335: if (product->api_user) {
2336: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatRARt", "Mat");
2337: PetscCall(PetscOptionsBool("-matrart_backend_cpu", "Use CPU code", "MatRARt", usecpu, &usecpu, NULL));
2338: PetscOptionsEnd();
2339: } else {
2340: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_RARt", "Mat");
2341: PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatRARt", usecpu, &usecpu, NULL));
2342: PetscOptionsEnd();
2343: }
2344: break;
2345: case MATPRODUCT_ABC:
2346: if (product->api_user) {
2347: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatMatMatMult", "Mat");
2348: PetscCall(PetscOptionsBool("-matmatmatmult_backend_cpu", "Use CPU code", "MatMatMatMult", usecpu, &usecpu, NULL));
2349: PetscOptionsEnd();
2350: } else {
2351: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_ABC", "Mat");
2352: PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatMatMatMult", usecpu, &usecpu, NULL));
2353: PetscOptionsEnd();
2354: }
2355: break;
2356: default:
2357: break;
2358: }
2359: if (usecpu) Biscusp = Ciscusp = PETSC_FALSE;
2360: }
2361: /* dispatch */
2362: if (isdense) {
2363: switch (product->type) {
2364: case MATPRODUCT_AB:
2365: case MATPRODUCT_AtB:
2366: case MATPRODUCT_ABt:
2367: case MATPRODUCT_PtAP:
2368: case MATPRODUCT_RARt:
2369: if (product->A->boundtocpu) {
2370: PetscCall(MatProductSetFromOptions_SeqAIJ_SeqDense(mat));
2371: } else {
2372: mat->ops->productsymbolic = MatProductSymbolic_SeqAIJCUSPARSE_SeqDENSECUDA;
2373: }
2374: break;
2375: case MATPRODUCT_ABC:
2376: mat->ops->productsymbolic = MatProductSymbolic_ABC_Basic;
2377: break;
2378: default:
2379: break;
2380: }
2381: } else if (Biscusp && Ciscusp) {
2382: switch (product->type) {
2383: case MATPRODUCT_AB:
2384: case MATPRODUCT_AtB:
2385: case MATPRODUCT_ABt:
2386: mat->ops->productsymbolic = MatProductSymbolic_SeqAIJCUSPARSE_SeqAIJCUSPARSE;
2387: break;
2388: case MATPRODUCT_PtAP:
2389: case MATPRODUCT_RARt:
2390: case MATPRODUCT_ABC:
2391: mat->ops->productsymbolic = MatProductSymbolic_ABC_Basic;
2392: break;
2393: default:
2394: break;
2395: }
2396: } else { /* fallback for AIJ */
2397: PetscCall(MatProductSetFromOptions_SeqAIJ(mat));
2398: }
2399: PetscFunctionReturn(PETSC_SUCCESS);
2400: }
2402: static PetscErrorCode MatMult_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy)
2403: {
2404: PetscFunctionBegin;
2405: PetscCall(MatMultAddKernel_SeqAIJCUSPARSE(A, xx, NULL, yy, PETSC_FALSE, PETSC_FALSE));
2406: PetscFunctionReturn(PETSC_SUCCESS);
2407: }
2409: static PetscErrorCode MatMultAdd_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy, Vec zz)
2410: {
2411: PetscFunctionBegin;
2412: PetscCall(MatMultAddKernel_SeqAIJCUSPARSE(A, xx, yy, zz, PETSC_FALSE, PETSC_FALSE));
2413: PetscFunctionReturn(PETSC_SUCCESS);
2414: }
2416: static PetscErrorCode MatMultHermitianTranspose_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy)
2417: {
2418: PetscFunctionBegin;
2419: PetscCall(MatMultAddKernel_SeqAIJCUSPARSE(A, xx, NULL, yy, PETSC_TRUE, PETSC_TRUE));
2420: PetscFunctionReturn(PETSC_SUCCESS);
2421: }
2423: static PetscErrorCode MatMultHermitianTransposeAdd_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy, Vec zz)
2424: {
2425: PetscFunctionBegin;
2426: PetscCall(MatMultAddKernel_SeqAIJCUSPARSE(A, xx, yy, zz, PETSC_TRUE, PETSC_TRUE));
2427: PetscFunctionReturn(PETSC_SUCCESS);
2428: }
2430: static PetscErrorCode MatMultTranspose_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy)
2431: {
2432: PetscFunctionBegin;
2433: PetscCall(MatMultAddKernel_SeqAIJCUSPARSE(A, xx, NULL, yy, PETSC_TRUE, PETSC_FALSE));
2434: PetscFunctionReturn(PETSC_SUCCESS);
2435: }
2437: __global__ static void ScatterAdd(PetscInt n, PetscInt *idx, const PetscScalar *x, PetscScalar *y)
2438: {
2439: int i = blockIdx.x * blockDim.x + threadIdx.x;
2440: if (i < n) y[idx[i]] += x[i];
2441: }
2443: /* z = op(A) x + y. If trans & !herm, op = ^T; if trans & herm, op = ^H; if !trans, op = no-op */
2444: static PetscErrorCode MatMultAddKernel_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy, Vec zz, PetscBool trans, PetscBool herm)
2445: {
2446: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;
2447: Mat_SeqAIJCUSPARSE *cusparsestruct = (Mat_SeqAIJCUSPARSE *)A->spptr;
2448: Mat_SeqAIJCUSPARSEMultStruct *matstruct;
2449: PetscScalar *xarray, *zarray, *dptr, *beta, *xptr;
2450: cusparseOperation_t opA = CUSPARSE_OPERATION_NON_TRANSPOSE;
2451: PetscBool compressed;
2452: PetscInt nx, ny;
2454: PetscFunctionBegin;
2455: PetscCheck(!herm || trans, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Hermitian and not transpose not supported");
2456: if (!a->nz) {
2457: if (yy) PetscCall(VecSeq_CUDA::Copy(yy, zz));
2458: else PetscCall(VecSeq_CUDA::Set(zz, 0));
2459: PetscFunctionReturn(PETSC_SUCCESS);
2460: }
2461: /* The line below is necessary due to the operations that modify the matrix on the CPU (axpy, scale, etc) */
2462: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
2463: if (!trans) {
2464: matstruct = (Mat_SeqAIJCUSPARSEMultStruct *)cusparsestruct->mat;
2465: PetscCheck(matstruct, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "SeqAIJCUSPARSE does not have a 'mat' (need to fix)");
2466: } else {
2467: if (herm || !A->form_explicit_transpose) {
2468: opA = herm ? CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE : CUSPARSE_OPERATION_TRANSPOSE;
2469: matstruct = (Mat_SeqAIJCUSPARSEMultStruct *)cusparsestruct->mat;
2470: } else {
2471: if (!A->transupdated) PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(A));
2472: matstruct = (Mat_SeqAIJCUSPARSEMultStruct *)cusparsestruct->matTranspose;
2473: }
2474: }
2475: /* Does the matrix use compressed rows (i.e., drop zero rows)? */
2476: compressed = matstruct->cprowIndices ? PETSC_TRUE : PETSC_FALSE;
2478: try {
2479: PetscCall(VecCUDAGetArrayRead(xx, (const PetscScalar **)&xarray));
2480: if (yy == zz) PetscCall(VecCUDAGetArray(zz, &zarray)); /* read & write zz, so need to get up-to-date zarray on GPU */
2481: else PetscCall(VecCUDAGetArrayWrite(zz, &zarray)); /* write zz, so no need to init zarray on GPU */
2483: PetscCall(PetscLogGpuTimeBegin());
2484: if (opA == CUSPARSE_OPERATION_NON_TRANSPOSE) {
2485: /* z = A x + beta y.
2486: If A is compressed (with less rows), then Ax is shorter than the full z, so we need a work vector to store Ax.
2487: When A is non-compressed, and z = y, we can set beta=1 to compute y = Ax + y in one call.
2488: */
2489: xptr = xarray;
2490: dptr = compressed ? cusparsestruct->workVector->data().get() : zarray;
2491: beta = (yy == zz && !compressed) ? matstruct->beta_one : matstruct->beta_zero;
2492: /* Get length of x, y for y=Ax. ny might be shorter than the work vector's allocated length, since the work vector is
2493: allocated to accommodate different uses. So we get the length info directly from mat.
2494: */
2495: if (cusparsestruct->format == MAT_CUSPARSE_CSR) {
2496: CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
2497: nx = mat->num_cols; // since y = Ax
2498: ny = mat->num_rows;
2499: }
2500: } else {
2501: /* z = A^T x + beta y
2502: If A is compressed, then we need a work vector as the shorter version of x to compute A^T x.
2503: Note A^Tx is of full length, so we set beta to 1.0 if y exists.
2504: */
2505: xptr = compressed ? cusparsestruct->workVector->data().get() : xarray;
2506: dptr = zarray;
2507: beta = yy ? matstruct->beta_one : matstruct->beta_zero;
2508: if (compressed) { /* Scatter x to work vector */
2509: thrust::device_ptr<PetscScalar> xarr = thrust::device_pointer_cast(xarray);
2511: thrust::for_each(
2512: #if PetscDefined(HAVE_THRUST_ASYNC)
2513: thrust::cuda::par.on(PetscDefaultCudaStream),
2514: #endif
2515: thrust::make_zip_iterator(thrust::make_tuple(cusparsestruct->workVector->begin(), thrust::make_permutation_iterator(xarr, matstruct->cprowIndices->begin()))),
2516: thrust::make_zip_iterator(thrust::make_tuple(cusparsestruct->workVector->begin(), thrust::make_permutation_iterator(xarr, matstruct->cprowIndices->begin()))) + matstruct->cprowIndices->size(), VecCUDAEqualsReverse());
2517: }
2518: if (cusparsestruct->format == MAT_CUSPARSE_CSR) {
2519: CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
2520: nx = mat->num_rows; // since y = A^T x
2521: ny = mat->num_cols;
2522: }
2523: }
2525: /* csr_spmv does y = alpha op(A) x + beta y */
2526: if (cusparsestruct->format == MAT_CUSPARSE_CSR) {
2527: PetscCheck(opA >= 0 && opA <= 2, PETSC_COMM_SELF, PETSC_ERR_SUP, "cuSPARSE ABI on cusparseOperation_t has changed and PETSc has not been updated accordingly");
2529: if (!matstruct->cuSpMV[opA].initialized) { /* built on demand */
2530: CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
2531: PetscCallCUSPARSE(cusparseCreateCsr(&matstruct->cuSpMV[opA].matDescr, mat->num_rows, mat->num_cols, mat->num_entries, mat->row_offsets->data().get(), mat->column_indices->data().get(), mat->values->data().get(), csrRowOffsetsType, csrColIndType, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
2532: PetscCallCUSPARSE(cusparseCreateDnVec(&matstruct->cuSpMV[opA].vecXDescr, nx, xptr, cusparse_scalartype));
2533: PetscCallCUSPARSE(cusparseCreateDnVec(&matstruct->cuSpMV[opA].vecYDescr, ny, dptr, cusparse_scalartype));
2534: PetscCallCUSPARSE(cusparseSpMV_bufferSize(cusparsestruct->handle, opA, matstruct->alpha_one, matstruct->cuSpMV[opA].matDescr, matstruct->cuSpMV[opA].vecXDescr, beta, matstruct->cuSpMV[opA].vecYDescr, cusparse_scalartype, cusparsestruct->spmvAlg,
2535: &matstruct->cuSpMV[opA].spmvBufferSize));
2536: PetscCallCUDA(cudaMalloc(&matstruct->cuSpMV[opA].spmvBuffer, matstruct->cuSpMV[opA].spmvBufferSize));
2537: #if PETSC_PKG_CUDA_VERSION_GE(12, 4, 0) // cusparseSpMV_preprocess is added in 12.4
2538: PetscCallCUSPARSE(cusparseSpMV_preprocess(cusparsestruct->handle, opA, matstruct->alpha_one, matstruct->cuSpMV[opA].matDescr, matstruct->cuSpMV[opA].vecXDescr, beta, matstruct->cuSpMV[opA].vecYDescr, cusparse_scalartype, cusparsestruct->spmvAlg,
2539: matstruct->cuSpMV[opA].spmvBuffer));
2540: #endif
2541: matstruct->cuSpMV[opA].initialized = PETSC_TRUE;
2542: } else {
2543: /* x, y's value pointers might change between calls, but their shape is kept, so we just update pointers */
2544: PetscCallCUSPARSE(cusparseDnVecSetValues(matstruct->cuSpMV[opA].vecXDescr, xptr));
2545: PetscCallCUSPARSE(cusparseDnVecSetValues(matstruct->cuSpMV[opA].vecYDescr, dptr));
2546: }
2548: PetscCallCUSPARSE(
2549: cusparseSpMV(cusparsestruct->handle, opA, matstruct->alpha_one, matstruct->cuSpMV[opA].matDescr, matstruct->cuSpMV[opA].vecXDescr, beta, matstruct->cuSpMV[opA].vecYDescr, cusparse_scalartype, cusparsestruct->spmvAlg, matstruct->cuSpMV[opA].spmvBuffer));
2551: } else {
2552: if (cusparsestruct->nrows) {
2553: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MAT_CUSPARSE_ELL and MAT_CUSPARSE_HYB are not supported since CUDA-11.0");
2554: }
2555: }
2556: PetscCall(PetscLogGpuTimeEnd());
2558: if (opA == CUSPARSE_OPERATION_NON_TRANSPOSE) {
2559: if (yy) { /* MatMultAdd: zz = A*xx + yy */
2560: if (compressed) { /* A is compressed. We first copy yy to zz, then ScatterAdd the work vector to zz */
2561: PetscCall(VecSeq_CUDA::Copy(yy, zz)); /* zz = yy */
2562: } else if (zz != yy) { /* A is not compressed. zz already contains A*xx, and we just need to add yy */
2563: PetscCall(VecSeq_CUDA::AXPY(zz, 1.0, yy)); /* zz += yy */
2564: }
2565: } else if (compressed) { /* MatMult: zz = A*xx. A is compressed, so we zero zz first, then ScatterAdd the work vector to zz */
2566: PetscCall(VecSeq_CUDA::Set(zz, 0));
2567: }
2569: /* ScatterAdd the result from work vector into the full vector when A is compressed */
2570: if (compressed) {
2571: PetscCall(PetscLogGpuTimeBegin());
2572: PetscInt n = (PetscInt)matstruct->cprowIndices->size();
2573: ScatterAdd<<<(int)((n + 255) / 256), 256, 0, PetscDefaultCudaStream>>>(n, matstruct->cprowIndices->data().get(), cusparsestruct->workVector->data().get(), zarray);
2574: PetscCall(PetscLogGpuTimeEnd());
2575: }
2576: } else {
2577: if (yy && yy != zz) PetscCall(VecSeq_CUDA::AXPY(zz, 1.0, yy)); /* zz += yy */
2578: }
2579: PetscCall(VecCUDARestoreArrayRead(xx, (const PetscScalar **)&xarray));
2580: if (yy == zz) PetscCall(VecCUDARestoreArray(zz, &zarray));
2581: else PetscCall(VecCUDARestoreArrayWrite(zz, &zarray));
2582: } catch (char *ex) {
2583: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_LIB, "CUSPARSE error: %s", ex);
2584: }
2585: if (yy) PetscCall(PetscLogGpuFlops(2.0 * a->nz));
2586: else PetscCall(PetscLogGpuFlops(2.0 * a->nz - a->nonzerorowcnt));
2587: PetscFunctionReturn(PETSC_SUCCESS);
2588: }
2590: static PetscErrorCode MatMultTransposeAdd_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy, Vec zz)
2591: {
2592: PetscFunctionBegin;
2593: PetscCall(MatMultAddKernel_SeqAIJCUSPARSE(A, xx, yy, zz, PETSC_TRUE, PETSC_FALSE));
2594: PetscFunctionReturn(PETSC_SUCCESS);
2595: }
2597: static PetscErrorCode MatGetDiagonal_SeqAIJCUSPARSE(Mat A, Vec diag)
2598: {
2599: PetscFunctionBegin;
2600: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::GetDiagonal(A, diag));
2601: PetscFunctionReturn(PETSC_SUCCESS);
2602: }
2604: static PetscErrorCode MatGetRowSumAbs_SeqAIJCUSPARSE(Mat A, Vec v)
2605: {
2606: PetscFunctionBegin;
2607: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::GetRowSumAbs(A, v));
2608: PetscFunctionReturn(PETSC_SUCCESS);
2609: }
2611: static PetscErrorCode MatDiagonalScale_SeqAIJCUSPARSE(Mat A, Vec ll, Vec rr)
2612: {
2613: PetscFunctionBegin;
2614: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::DiagonalScale(A, ll, rr));
2615: PetscFunctionReturn(PETSC_SUCCESS);
2616: }
2618: static PetscErrorCode MatAssemblyEnd_SeqAIJCUSPARSE(Mat A, MatAssemblyType mode)
2619: {
2620: PetscFunctionBegin;
2621: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::AssemblyEnd(A, mode));
2622: PetscFunctionReturn(PETSC_SUCCESS);
2623: }
2625: /*@
2626: MatCreateSeqAIJCUSPARSE - Creates a sparse matrix in `MATAIJCUSPARSE` (compressed row) format for use on NVIDIA GPUs
2628: Collective
2630: Input Parameters:
2631: + comm - MPI communicator, set to `PETSC_COMM_SELF`
2632: . m - number of rows
2633: . n - number of columns
2634: . nz - number of nonzeros per row (same for all rows), ignored if `nnz` is provide
2635: - nnz - array containing the number of nonzeros in the various rows (possibly different for each row) or `NULL`
2637: Output Parameter:
2638: . A - the matrix
2640: Level: intermediate
2642: Notes:
2643: This matrix will ultimately pushed down to NVIDIA GPUs and use the CuSPARSE library for
2644: calculations. For good matrix assembly performance the user should preallocate the matrix
2645: storage by setting the parameter `nz` (or the array `nnz`).
2647: It is recommended that one use the `MatCreate()`, `MatSetType()` and/or `MatSetFromOptions()`,
2648: MatXXXXSetPreallocation() paradgm instead of this routine directly.
2649: [MatXXXXSetPreallocation() is, for example, `MatSeqAIJSetPreallocation()`]
2651: The AIJ format, also called
2652: compressed row storage, is fully compatible with standard Fortran
2653: storage. That is, the stored row and column indices can begin at
2654: either one (as in Fortran) or zero.
2656: Specify the preallocated storage with either nz or nnz (not both).
2657: Set `nz` = `PETSC_DEFAULT` and `nnz` = `NULL` for PETSc to control dynamic memory
2658: allocation.
2660: When working with matrices for GPUs, it is often better to use the `MatSetPreallocationCOO()` and `MatSetValuesCOO()` paradigm rather than using this routine and `MatSetValues()`
2662: .seealso: [](ch_matrices), `Mat`, `MATSEQAIJCUSPARSE`, `MatCreate()`, `MatCreateAIJ()`, `MatSetValues()`, `MatSeqAIJSetColumnIndices()`, `MatCreateSeqAIJWithArrays()`, `MATAIJCUSPARSE`,
2663: `MatSetPreallocationCOO()`, `MatSetValuesCOO()`
2664: @*/
2665: PetscErrorCode MatCreateSeqAIJCUSPARSE(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt nz, const PetscInt nnz[], Mat *A)
2666: {
2667: return MatSeqAIJCUSPARSE_CUPM_t::CreateSeqAIJ(comm, m, n, nz, nnz, A);
2668: }
2670: static PetscErrorCode MatDestroy_SeqAIJCUSPARSE(Mat A)
2671: {
2672: return MatSeqAIJCUSPARSE_CUPM_t::Destroy(A);
2673: }
2675: static PetscErrorCode MatDuplicate_SeqAIJCUSPARSE(Mat A, MatDuplicateOption cpvalues, Mat *B)
2676: {
2677: PetscFunctionBegin;
2678: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::Duplicate(A, cpvalues, B));
2679: PetscFunctionReturn(PETSC_SUCCESS);
2680: }
2682: static PetscErrorCode MatAXPY_SeqAIJCUSPARSE(Mat Y, PetscScalar a, Mat X, MatStructure str)
2683: {
2684: Mat_SeqAIJ *x = (Mat_SeqAIJ *)X->data, *y = (Mat_SeqAIJ *)Y->data;
2685: Mat_SeqAIJCUSPARSE *cy;
2686: Mat_SeqAIJCUSPARSE *cx;
2687: CsrMatrix *csry, *csrx;
2689: PetscFunctionBegin;
2690: cy = (Mat_SeqAIJCUSPARSE *)Y->spptr;
2691: cx = (Mat_SeqAIJCUSPARSE *)X->spptr;
2692: if (X->ops->axpy != Y->ops->axpy) {
2693: PetscCall(MatSeqAIJCUSPARSEInvalidateTranspose(Y, PETSC_FALSE));
2694: PetscCall(MatAXPY_SeqAIJ(Y, a, X, str));
2695: PetscFunctionReturn(PETSC_SUCCESS);
2696: }
2697: /* if we are here, it means both matrices are bound to GPU */
2698: PetscCall(MatSeqAIJCUSPARSECopyToGPU(Y));
2699: PetscCall(MatSeqAIJCUSPARSECopyToGPU(X));
2700: PetscCheck(cy->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)Y), PETSC_ERR_GPU, "only MAT_CUSPARSE_CSR supported");
2701: PetscCheck(cx->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)X), PETSC_ERR_GPU, "only MAT_CUSPARSE_CSR supported");
2702: csry = (CsrMatrix *)cy->mat->mat;
2703: csrx = (CsrMatrix *)cx->mat->mat;
2704: /* see if we can turn this into a cublas axpy */
2705: if (str != SAME_NONZERO_PATTERN && x->nz == y->nz && !x->compressedrow.use && !y->compressedrow.use) {
2706: bool eq = thrust::equal(thrust::device, csry->row_offsets->begin(), csry->row_offsets->end(), csrx->row_offsets->begin());
2707: if (eq) eq = thrust::equal(thrust::device, csry->column_indices->begin(), csry->column_indices->end(), csrx->column_indices->begin());
2708: if (eq) str = SAME_NONZERO_PATTERN;
2709: }
2710: /* spgeam is buggy with one column */
2711: if (Y->cmap->n == 1 && str != SAME_NONZERO_PATTERN) str = DIFFERENT_NONZERO_PATTERN;
2713: #if !PetscDefined(USE_64BIT_INDICES) // cusparseScsrgeam2 etc. do not support 64bit indices
2714: if (str == SUBSET_NONZERO_PATTERN) {
2715: PetscScalar *ay, b = 1.0;
2716: const PetscScalar *ax;
2717: size_t bufferSize;
2718: void *buffer;
2720: PetscCall(MatSeqAIJCUSPARSEGetArrayRead(X, &ax));
2721: PetscCall(MatSeqAIJCUSPARSEGetArray(Y, &ay));
2722: PetscCallCUSPARSE(cusparseSetPointerMode(cy->handle, CUSPARSE_POINTER_MODE_HOST));
2723: PetscCallCUSPARSE(cusparse_csr_spgeam_bufferSize(cy->handle, Y->rmap->n, Y->cmap->n, &a, cx->mat->descr, x->nz, ax, csrx->row_offsets->data().get(), csrx->column_indices->data().get(), &b, cy->mat->descr, y->nz, ay, csry->row_offsets->data().get(),
2724: csry->column_indices->data().get(), cy->mat->descr, ay, csry->row_offsets->data().get(), csry->column_indices->data().get(), &bufferSize));
2725: PetscCallCUDA(cudaMalloc(&buffer, bufferSize));
2726: PetscCall(PetscLogGpuTimeBegin());
2727: PetscCallCUSPARSE(cusparse_csr_spgeam(cy->handle, Y->rmap->n, Y->cmap->n, &a, cx->mat->descr, x->nz, ax, csrx->row_offsets->data().get(), csrx->column_indices->data().get(), &b, cy->mat->descr, y->nz, ay, csry->row_offsets->data().get(),
2728: csry->column_indices->data().get(), cy->mat->descr, ay, csry->row_offsets->data().get(), csry->column_indices->data().get(), buffer));
2729: PetscCall(PetscLogGpuFlops(x->nz + y->nz));
2730: PetscCall(PetscLogGpuTimeEnd());
2731: PetscCallCUDA(cudaFree(buffer));
2733: PetscCallCUSPARSE(cusparseSetPointerMode(cy->handle, CUSPARSE_POINTER_MODE_DEVICE));
2734: PetscCall(MatSeqAIJCUSPARSERestoreArrayRead(X, &ax));
2735: PetscCall(MatSeqAIJCUSPARSERestoreArray(Y, &ay));
2736: } else
2737: #endif
2738: if (str == SAME_NONZERO_PATTERN) {
2739: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::AXPY_SameNZ(Y, a, X));
2740: } else {
2741: PetscCall(MatSeqAIJCUSPARSEInvalidateTranspose(Y, PETSC_FALSE));
2742: PetscCall(MatAXPY_SeqAIJ(Y, a, X, str));
2743: }
2744: PetscFunctionReturn(PETSC_SUCCESS);
2745: }
2747: static PetscErrorCode MatScale_SeqAIJCUSPARSE(Mat Y, PetscScalar a)
2748: {
2749: PetscFunctionBegin;
2750: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::Scale(Y, a));
2751: PetscFunctionReturn(PETSC_SUCCESS);
2752: }
2754: static PetscErrorCode MatZeroEntries_SeqAIJCUSPARSE(Mat A)
2755: {
2756: PetscFunctionBegin;
2757: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::ZeroEntries(A));
2758: PetscFunctionReturn(PETSC_SUCCESS);
2759: }
2761: static PetscErrorCode MatGetCurrentMemType_SeqAIJCUSPARSE(Mat A, PetscMemType *m)
2762: {
2763: PetscFunctionBegin;
2764: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::GetCurrentMemType(A, m));
2765: PetscFunctionReturn(PETSC_SUCCESS);
2766: }
2768: static PetscErrorCode MatBindToCPU_SeqAIJCUSPARSE(Mat A, PetscBool flg)
2769: {
2770: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;
2772: PetscFunctionBegin;
2773: if (A->factortype != MAT_FACTOR_NONE) {
2774: A->boundtocpu = flg;
2775: PetscFunctionReturn(PETSC_SUCCESS);
2776: }
2777: if (flg) {
2778: PetscCall(MatSeqAIJCUSPARSECopyFromGPU(A));
2780: A->ops->scale = MatScale_SeqAIJ;
2781: A->ops->getdiagonal = MatGetDiagonal_SeqAIJ;
2782: A->ops->getrowsumabs = MatGetRowSumAbs_SeqAIJ;
2783: A->ops->diagonalscale = MatDiagonalScale_SeqAIJ;
2784: A->ops->axpy = MatAXPY_SeqAIJ;
2785: A->ops->zeroentries = MatZeroEntries_SeqAIJ;
2786: A->ops->mult = MatMult_SeqAIJ;
2787: A->ops->multadd = MatMultAdd_SeqAIJ;
2788: A->ops->multtranspose = MatMultTranspose_SeqAIJ;
2789: A->ops->multtransposeadd = MatMultTransposeAdd_SeqAIJ;
2790: A->ops->multhermitiantranspose = NULL;
2791: A->ops->multhermitiantransposeadd = NULL;
2792: A->ops->productsetfromoptions = MatProductSetFromOptions_SeqAIJ;
2793: A->ops->getcurrentmemtype = NULL;
2794: PetscCall(PetscMemzero(a->ops, sizeof(Mat_SeqAIJOps)));
2795: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSeqAIJCopySubArray_C", NULL));
2796: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijcusparse_seqdensecuda_C", NULL));
2797: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijcusparse_seqdense_C", NULL));
2798: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetPreallocationCOO_C", NULL));
2799: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetValuesCOO_C", NULL));
2800: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijcusparse_seqaijcusparse_C", NULL));
2801: } else {
2802: A->ops->scale = MatScale_SeqAIJCUSPARSE;
2803: A->ops->getdiagonal = MatGetDiagonal_SeqAIJCUSPARSE;
2804: A->ops->getrowsumabs = MatGetRowSumAbs_SeqAIJCUSPARSE;
2805: A->ops->diagonalscale = MatDiagonalScale_SeqAIJCUSPARSE;
2806: A->ops->axpy = MatAXPY_SeqAIJCUSPARSE;
2807: A->ops->zeroentries = MatZeroEntries_SeqAIJCUSPARSE;
2808: A->ops->mult = MatMult_SeqAIJCUSPARSE;
2809: A->ops->multadd = MatMultAdd_SeqAIJCUSPARSE;
2810: A->ops->multtranspose = MatMultTranspose_SeqAIJCUSPARSE;
2811: A->ops->multtransposeadd = MatMultTransposeAdd_SeqAIJCUSPARSE;
2812: A->ops->multhermitiantranspose = MatMultHermitianTranspose_SeqAIJCUSPARSE;
2813: A->ops->multhermitiantransposeadd = MatMultHermitianTransposeAdd_SeqAIJCUSPARSE;
2814: A->ops->productsetfromoptions = MatProductSetFromOptions_SeqAIJCUSPARSE;
2815: A->ops->getcurrentmemtype = MatGetCurrentMemType_SeqAIJCUSPARSE;
2816: a->ops->getarray = MatSeqAIJGetArray_SeqAIJCUSPARSE;
2817: a->ops->restorearray = MatSeqAIJRestoreArray_SeqAIJCUSPARSE;
2818: a->ops->getarrayread = MatSeqAIJGetArrayRead_SeqAIJCUSPARSE;
2819: a->ops->restorearrayread = MatSeqAIJRestoreArrayRead_SeqAIJCUSPARSE;
2820: a->ops->getarraywrite = MatSeqAIJGetArrayWrite_SeqAIJCUSPARSE;
2821: a->ops->restorearraywrite = MatSeqAIJRestoreArrayWrite_SeqAIJCUSPARSE;
2822: a->ops->getcsrandmemtype = MatSeqAIJGetCSRAndMemType_SeqAIJCUSPARSE;
2824: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSeqAIJCopySubArray_C", MatSeqAIJCopySubArray_SeqAIJCUSPARSE));
2825: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijcusparse_seqdensecuda_C", MatProductSetFromOptions_SeqAIJCUSPARSE));
2826: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijcusparse_seqdense_C", MatProductSetFromOptions_SeqAIJCUSPARSE));
2827: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetPreallocationCOO_C", MatSetPreallocationCOO_SeqAIJCUSPARSE));
2828: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetValuesCOO_C", MatSetValuesCOO_SeqAIJCUSPARSE));
2829: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijcusparse_seqaijcusparse_C", MatProductSetFromOptions_SeqAIJCUSPARSE));
2830: }
2831: A->boundtocpu = flg;
2832: a->inode.use = (flg && a->inode.size_csr) ? PETSC_TRUE : PETSC_FALSE;
2833: PetscFunctionReturn(PETSC_SUCCESS);
2834: }
2836: PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJCUSPARSE(Mat A, MatType, MatReuse reuse, Mat *newmat)
2837: {
2838: Mat B;
2840: PetscFunctionBegin;
2841: PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUDA)); /* first use of CUSPARSE may be via MatConvert */
2842: if (reuse == MAT_INITIAL_MATRIX) {
2843: PetscCall(MatDuplicate(A, MAT_COPY_VALUES, newmat));
2844: } else if (reuse == MAT_REUSE_MATRIX) {
2845: PetscCall(MatCopy(A, *newmat, SAME_NONZERO_PATTERN));
2846: }
2847: B = *newmat;
2849: PetscCall(PetscFree(B->defaultvectype));
2850: PetscCall(PetscStrallocpy(VECCUDA, &B->defaultvectype));
2852: if (reuse != MAT_REUSE_MATRIX && !B->spptr) {
2853: if (B->factortype == MAT_FACTOR_NONE) {
2854: Mat_SeqAIJCUSPARSE *spptr;
2855: PetscCall(PetscNew(&spptr));
2856: PetscCallCUSPARSE(cusparseCreate(&spptr->handle));
2857: PetscCallCUSPARSE(cusparseSetStream(spptr->handle, PetscDefaultCudaStream));
2858: spptr->format = MAT_CUSPARSE_CSR;
2859: spptr->spmvAlg = CUSPARSE_SPMV_CSR_ALG1; /* default, since we only support csr */
2860: spptr->spmmAlg = CUSPARSE_SPMM_CSR_ALG1; /* default, only support column-major dense matrix B */
2861: spptr->csr2cscAlg = CUSPARSE_CSR2CSC_ALG1;
2862: B->spptr = spptr;
2863: } else {
2864: Mat_SeqAIJCUSPARSETriFactors *spptr;
2866: PetscCall(PetscNew(&spptr));
2867: PetscCallCUSPARSE(cusparseCreate(&spptr->handle));
2868: PetscCallCUSPARSE(cusparseSetStream(spptr->handle, PetscDefaultCudaStream));
2869: B->spptr = spptr;
2870: }
2871: B->offloadmask = PETSC_OFFLOAD_UNALLOCATED;
2872: }
2873: B->ops->assemblyend = MatAssemblyEnd_SeqAIJCUSPARSE;
2874: B->ops->destroy = MatDestroy_SeqAIJCUSPARSE;
2875: B->ops->setoption = MatSetOption_SeqAIJCUSPARSE;
2876: B->ops->setfromoptions = MatSetFromOptions_SeqAIJCUSPARSE;
2877: B->ops->bindtocpu = MatBindToCPU_SeqAIJCUSPARSE;
2878: B->ops->duplicate = MatDuplicate_SeqAIJCUSPARSE;
2879: B->ops->getcurrentmemtype = MatGetCurrentMemType_SeqAIJCUSPARSE;
2881: PetscCall(MatBindToCPU_SeqAIJCUSPARSE(B, PETSC_FALSE));
2882: PetscCall(PetscObjectChangeTypeName((PetscObject)B, MATSEQAIJCUSPARSE));
2883: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatCUSPARSESetFormat_C", MatCUSPARSESetFormat_SeqAIJCUSPARSE));
2884: #if PetscDefined(HAVE_HYPRE)
2885: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_seqaijcusparse_hypre_C", MatConvert_AIJ_HYPRE));
2886: #endif
2887: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatCUSPARSESetUseCPUSolve_C", MatCUSPARSESetUseCPUSolve_SeqAIJCUSPARSE));
2888: PetscFunctionReturn(PETSC_SUCCESS);
2889: }
2891: PETSC_EXTERN PetscErrorCode MatCreate_SeqAIJCUSPARSE(Mat B)
2892: {
2893: PetscFunctionBegin;
2894: PetscCall(MatCreate_SeqAIJ(B));
2895: PetscCall(MatConvert_SeqAIJ_SeqAIJCUSPARSE(B, MATSEQAIJCUSPARSE, MAT_INPLACE_MATRIX, &B));
2896: PetscFunctionReturn(PETSC_SUCCESS);
2897: }
2899: /*MC
2900: MATSEQAIJCUSPARSE - MATAIJCUSPARSE = "(seq)aijcusparse" - A matrix type to be used for sparse matrices on NVIDIA GPUs.
2902: Options Database Keys:
2903: + -mat_type aijcusparse - Sets the matrix type to `MATSEQAIJCUSPARSE` during a call to `MatSetFromOptions()`
2904: . -mat_cusparse_storage_format (csr|ell|hyb) - Sets the storage format of matrices (for `MatMult()` and factors in `MatSolve()`).
2905: . -mat_cusparse_mult_storage_format (csr|ell|hyb) - Sets the storage format of matrices (for `MatMult()`).
2906: - -mat_cusparse_use_cpu_solve - Performs the `MatSolve()` on the CPU.
2908: Level: beginner
2910: Notes:
2911: These matrices can be in either CSR, ELL, or HYB format.
2913: All matrix calculations are performed on NVIDIA GPUs using the cuSPARSE library.
2915: Uses 32-bit integers internally. If PETSc is configured `--with-64-bit-indices`, the integer row and column indices are stored on the GPU with `int`. It is unclear what happens
2916: if some integer values passed in do not fit in `int`.
2918: .seealso: [](ch_matrices), `Mat`, `MatCreateSeqAIJCUSPARSE()`, `MatCUSPARSESetUseCPUSolve()`, `MATAIJCUSPARSE`, `MatCreateAIJCUSPARSE()`, `MatCUSPARSESetFormat()`, `MatCUSPARSEStorageFormat`, `MatCUSPARSEFormatOperation`
2919: M*/
2921: PETSC_INTERN PetscErrorCode MatSolverTypeRegister_CUSPARSE(void)
2922: {
2923: PetscFunctionBegin;
2924: PetscCall(MatSolverTypeRegister(MATSOLVERCUSPARSE, MATSEQAIJCUSPARSE, MAT_FACTOR_LU, MatGetFactor_seqaijcusparse_cusparse));
2925: PetscCall(MatSolverTypeRegister(MATSOLVERCUSPARSE, MATSEQAIJCUSPARSE, MAT_FACTOR_CHOLESKY, MatGetFactor_seqaijcusparse_cusparse));
2926: PetscCall(MatSolverTypeRegister(MATSOLVERCUSPARSE, MATSEQAIJCUSPARSE, MAT_FACTOR_ILU, MatGetFactor_seqaijcusparse_cusparse));
2927: PetscCall(MatSolverTypeRegister(MATSOLVERCUSPARSE, MATSEQAIJCUSPARSE, MAT_FACTOR_ICC, MatGetFactor_seqaijcusparse_cusparse));
2928: PetscFunctionReturn(PETSC_SUCCESS);
2929: }
2931: static PetscErrorCode MatSeqAIJCUSPARSE_Destroy(Mat mat)
2932: {
2933: Mat_SeqAIJCUSPARSE *cusp = static_cast<Mat_SeqAIJCUSPARSE *>(mat->spptr);
2935: PetscFunctionBegin;
2936: if (cusp) {
2937: PetscCall(MatSeqAIJCUSPARSEMultStruct_Destroy(&cusp->mat, cusp->format));
2938: PetscCall(MatSeqAIJCUSPARSEMultStruct_Destroy(&cusp->matTranspose, cusp->format));
2939: delete cusp->workVector;
2940: delete cusp->rowoffsets_gpu;
2941: delete cusp->csr2csc_i;
2942: delete cusp->coords;
2943: if (cusp->handle) PetscCallCUSPARSE(cusparseDestroy(cusp->handle));
2944: PetscCall(PetscFree(mat->spptr));
2945: }
2946: PetscFunctionReturn(PETSC_SUCCESS);
2947: }
2949: static PetscErrorCode CsrMatrix_Destroy(CsrMatrix **mat)
2950: {
2951: PetscFunctionBegin;
2952: if (*mat) {
2953: delete (*mat)->values;
2954: delete (*mat)->column_indices;
2955: delete (*mat)->row_offsets;
2956: delete *mat;
2957: *mat = 0;
2958: }
2959: PetscFunctionReturn(PETSC_SUCCESS);
2960: }
2962: static PetscErrorCode MatSeqAIJCUSPARSEMultStruct_Destroy(Mat_SeqAIJCUSPARSEMultStruct **matstruct, MatCUSPARSEStorageFormat format)
2963: {
2964: CsrMatrix *mat;
2966: PetscFunctionBegin;
2967: if (*matstruct) {
2968: if ((*matstruct)->mat) {
2969: if (format == MAT_CUSPARSE_ELL || format == MAT_CUSPARSE_HYB) {
2970: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MAT_CUSPARSE_ELL and MAT_CUSPARSE_HYB are not supported since CUDA-11.0");
2971: } else {
2972: mat = (CsrMatrix *)(*matstruct)->mat;
2973: PetscCall(CsrMatrix_Destroy(&mat));
2974: }
2975: }
2976: if ((*matstruct)->descr) PetscCallCUSPARSE(cusparseDestroyMatDescr((*matstruct)->descr));
2977: delete (*matstruct)->cprowIndices;
2978: PetscCallCUDA(cudaFree((*matstruct)->alpha_one));
2979: PetscCallCUDA(cudaFree((*matstruct)->beta_zero));
2980: PetscCallCUDA(cudaFree((*matstruct)->beta_one));
2982: Mat_SeqAIJCUSPARSEMultStruct *mdata = *matstruct;
2983: if (mdata->matDescr) PetscCallCUSPARSE(cusparseDestroySpMat(mdata->matDescr));
2985: for (int i = 0; i < 3; i++) {
2986: if (mdata->cuSpMV[i].initialized) {
2987: PetscCallCUDA(cudaFree(mdata->cuSpMV[i].spmvBuffer));
2988: PetscCallCUSPARSE(cusparseDestroyDnVec(mdata->cuSpMV[i].vecXDescr));
2989: PetscCallCUSPARSE(cusparseDestroyDnVec(mdata->cuSpMV[i].vecYDescr));
2990: PetscCallCUSPARSE(cusparseDestroySpMat(mdata->cuSpMV[i].matDescr));
2991: }
2992: if (mdata->matDescr_SpMM[i]) PetscCallCUSPARSE(cusparseDestroySpMat(mdata->matDescr_SpMM[i]));
2993: }
2994: delete *matstruct;
2995: *matstruct = NULL;
2996: }
2997: PetscFunctionReturn(PETSC_SUCCESS);
2998: }
3000: PetscErrorCode MatSeqAIJCUSPARSETriFactors_Reset(Mat_SeqAIJCUSPARSETriFactors_p *trifactors)
3001: {
3002: Mat_SeqAIJCUSPARSETriFactors *fs = *trifactors;
3004: PetscFunctionBegin;
3005: if (fs) {
3006: delete fs->rpermIndices;
3007: delete fs->cpermIndices;
3008: fs->rpermIndices = NULL;
3009: fs->cpermIndices = NULL;
3010: fs->init_dev_prop = PETSC_FALSE;
3011: PetscCallCUDA(cudaFree(fs->csrRowPtr));
3012: PetscCallCUDA(cudaFree(fs->csrColIdx));
3013: PetscCallCUDA(cudaFree(fs->csrRowPtr32));
3014: PetscCallCUDA(cudaFree(fs->csrColIdx32));
3015: PetscCallCUDA(cudaFree(fs->csrVal));
3016: PetscCallCUDA(cudaFree(fs->diag));
3017: PetscCallCUDA(cudaFree(fs->X));
3018: PetscCallCUDA(cudaFree(fs->Y));
3019: // PetscCallCUDA(cudaFree(fs->factBuffer_M)); /* No needed since factBuffer_M shares with one of spsvBuffer_L/U */
3020: PetscCallCUDA(cudaFree(fs->spsvBuffer_L));
3021: PetscCallCUDA(cudaFree(fs->spsvBuffer_U));
3022: PetscCallCUDA(cudaFree(fs->spsvBuffer_Lt));
3023: PetscCallCUDA(cudaFree(fs->spsvBuffer_Ut));
3024: PetscCallCUSPARSE(cusparseDestroyMatDescr(fs->matDescr_M));
3025: if (fs->spMatDescr_L) PetscCallCUSPARSE(cusparseDestroySpMat(fs->spMatDescr_L));
3026: if (fs->spMatDescr_U) PetscCallCUSPARSE(cusparseDestroySpMat(fs->spMatDescr_U));
3027: PetscCallCUSPARSE(cusparseSpSV_destroyDescr(fs->spsvDescr_L));
3028: PetscCallCUSPARSE(cusparseSpSV_destroyDescr(fs->spsvDescr_Lt));
3029: PetscCallCUSPARSE(cusparseSpSV_destroyDescr(fs->spsvDescr_U));
3030: PetscCallCUSPARSE(cusparseSpSV_destroyDescr(fs->spsvDescr_Ut));
3031: if (fs->dnVecDescr_X) PetscCallCUSPARSE(cusparseDestroyDnVec(fs->dnVecDescr_X));
3032: if (fs->dnVecDescr_Y) PetscCallCUSPARSE(cusparseDestroyDnVec(fs->dnVecDescr_Y));
3033: PetscCallCUSPARSE(cusparseDestroyCsrilu02Info(fs->ilu0Info_M));
3034: PetscCallCUSPARSE(cusparseDestroyCsric02Info(fs->ic0Info_M));
3035: PetscCall(PetscFree(fs->csrRowPtr_h));
3036: PetscCall(PetscFree(fs->csrVal_h));
3037: PetscCall(PetscFree(fs->diag_h));
3038: fs->createdTransposeSpSVDescr = PETSC_FALSE;
3039: fs->updatedTransposeSpSVAnalysis = PETSC_FALSE;
3040: }
3041: PetscFunctionReturn(PETSC_SUCCESS);
3042: }
3044: static PetscErrorCode MatSeqAIJCUSPARSETriFactors_Destroy(Mat_SeqAIJCUSPARSETriFactors **trifactors)
3045: {
3046: PetscFunctionBegin;
3047: if (*trifactors) {
3048: PetscCall(MatSeqAIJCUSPARSETriFactors_Reset(trifactors));
3049: PetscCallCUSPARSE(cusparseDestroy((*trifactors)->handle));
3050: PetscCall(PetscFree(*trifactors));
3051: }
3052: PetscFunctionReturn(PETSC_SUCCESS);
3053: }
3055: static PetscErrorCode MatSeqAIJCUSPARSEInvalidateTranspose(Mat A, PetscBool destroy)
3056: {
3057: Mat_SeqAIJCUSPARSE *cusp = (Mat_SeqAIJCUSPARSE *)A->spptr;
3059: PetscFunctionBegin;
3060: PetscCheckTypeName(A, MATSEQAIJCUSPARSE);
3061: if (!cusp) PetscFunctionReturn(PETSC_SUCCESS);
3062: if (destroy) {
3063: PetscCall(MatSeqAIJCUSPARSEMultStruct_Destroy(&cusp->matTranspose, cusp->format));
3064: delete cusp->csr2csc_i;
3065: cusp->csr2csc_i = NULL;
3066: }
3067: A->transupdated = PETSC_FALSE;
3068: PetscFunctionReturn(PETSC_SUCCESS);
3069: }
3071: static PetscErrorCode MatSetPreallocationCOO_SeqAIJCUSPARSE(Mat mat, PetscCount coo_n, PetscInt coo_i[], PetscInt coo_j[])
3072: {
3073: PetscFunctionBegin;
3074: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::SetPreallocationCOO(mat, coo_n, coo_i, coo_j));
3075: PetscFunctionReturn(PETSC_SUCCESS);
3076: }
3078: static PetscErrorCode MatSetValuesCOO_SeqAIJCUSPARSE(Mat A, const PetscScalar v[], InsertMode imode)
3079: {
3080: PetscFunctionBegin;
3081: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::SetValuesCOO(A, v, imode));
3082: PetscFunctionReturn(PETSC_SUCCESS);
3083: }
3085: /*@
3086: MatSeqAIJCUSPARSEGetIJ - returns the device row storage `i` and `j` indices for `MATSEQAIJCUSPARSE` matrices.
3088: Not Collective
3090: Input Parameters:
3091: + A - the matrix
3092: - compressed - `PETSC_TRUE` or `PETSC_FALSE` indicating the matrix data structure should be always returned in compressed form
3094: Output Parameters:
3095: + i - the CSR row pointers
3096: - j - the CSR column indices
3098: Level: developer
3100: Note:
3101: When compressed is true, the CSR structure does not contain empty rows
3103: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSERestoreIJ()`, `MatSeqAIJCUSPARSEGetArrayRead()`
3104: @*/
3105: PetscErrorCode MatSeqAIJCUSPARSEGetIJ(Mat A, PetscBool compressed, const PetscInt *i[], const PetscInt *j[])
3106: {
3107: PetscFunctionBegin;
3108: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::GetIJ(A, compressed, i, j));
3109: PetscFunctionReturn(PETSC_SUCCESS);
3110: }
3112: /*@
3113: MatSeqAIJCUSPARSERestoreIJ - restore the device row storage `i` and `j` indices obtained with `MatSeqAIJCUSPARSEGetIJ()`
3115: Not Collective
3117: Input Parameters:
3118: + A - the matrix
3119: . compressed - `PETSC_TRUE` or `PETSC_FALSE` indicating the matrix data structure should be always returned in compressed form
3120: . i - the CSR row pointers
3121: - j - the CSR column indices
3123: Level: developer
3125: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetIJ()`
3126: @*/
3127: PetscErrorCode MatSeqAIJCUSPARSERestoreIJ(Mat A, PetscBool compressed, const PetscInt *i[], const PetscInt *j[])
3128: {
3129: PetscFunctionBegin;
3130: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::RestoreIJ(A, compressed, i, j));
3131: PetscFunctionReturn(PETSC_SUCCESS);
3132: }
3134: /*@
3135: MatSeqAIJCUSPARSEGetArrayRead - gives read-only access to the array where the device data for a `MATSEQAIJCUSPARSE` matrix nonzero entries are stored
3137: Not Collective
3139: Input Parameter:
3140: . A - a `MATSEQAIJCUSPARSE` matrix
3142: Output Parameter:
3143: . a - pointer to the device data
3145: Level: developer
3147: Note:
3148: Will trigger host-to-device copies if the most up-to-date matrix data is on the host
3150: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetArray()`, `MatSeqAIJCUSPARSEGetArrayWrite()`, `MatSeqAIJCUSPARSERestoreArrayRead()`
3151: @*/
3152: PetscErrorCode MatSeqAIJCUSPARSEGetArrayRead(Mat A, const PetscScalar **a)
3153: {
3154: return MatSeqAIJCUSPARSE_CUPM_t::GetArrayRead(A, a);
3155: }
3157: /*@
3158: MatSeqAIJCUSPARSERestoreArrayRead - restore the read-only access array obtained from `MatSeqAIJCUSPARSEGetArrayRead()`
3160: Not Collective
3162: Input Parameters:
3163: + A - a `MATSEQAIJCUSPARSE` matrix
3164: - a - pointer to the device data
3166: Level: developer
3168: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetArrayRead()`
3169: @*/
3170: PetscErrorCode MatSeqAIJCUSPARSERestoreArrayRead(Mat A, const PetscScalar **a)
3171: {
3172: return MatSeqAIJCUSPARSE_CUPM_t::RestoreArrayRead(A, a);
3173: }
3175: /*@
3176: MatSeqAIJCUSPARSEGetArray - gives read-write access to the array where the device data for a `MATSEQAIJCUSPARSE` matrix is stored
3178: Not Collective
3180: Input Parameter:
3181: . A - a `MATSEQAIJCUSPARSE` matrix
3183: Output Parameter:
3184: . a - pointer to the device data
3186: Level: developer
3188: Note:
3189: Will trigger host-to-device copies if the most up-to-date matrix data is on the host
3191: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetArrayRead()`, `MatSeqAIJCUSPARSEGetArrayWrite()`, `MatSeqAIJCUSPARSERestoreArray()`
3192: @*/
3193: PetscErrorCode MatSeqAIJCUSPARSEGetArray(Mat A, PetscScalar **a)
3194: {
3195: return MatSeqAIJCUSPARSE_CUPM_t::GetArray(A, a);
3196: }
3197: /*@
3198: MatSeqAIJCUSPARSERestoreArray - restore the read-write access array obtained from `MatSeqAIJCUSPARSEGetArray()`
3200: Not Collective
3202: Input Parameters:
3203: + A - a `MATSEQAIJCUSPARSE` matrix
3204: - a - pointer to the device data
3206: Level: developer
3208: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetArray()`
3209: @*/
3210: PetscErrorCode MatSeqAIJCUSPARSERestoreArray(Mat A, PetscScalar **a)
3211: {
3212: return MatSeqAIJCUSPARSE_CUPM_t::RestoreArray(A, a);
3213: }
3215: /*@
3216: MatSeqAIJCUSPARSEGetArrayWrite - gives write access to the array where the device data for a `MATSEQAIJCUSPARSE` matrix is stored
3218: Not Collective
3220: Input Parameter:
3221: . A - a `MATSEQAIJCUSPARSE` matrix
3223: Output Parameter:
3224: . a - pointer to the device data
3226: Level: developer
3228: Note:
3229: Does not trigger any host to device copies.
3231: It marks the data GPU valid so users must set all the values in `a` to ensure out-of-date data is not considered current
3233: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetArray()`, `MatSeqAIJCUSPARSEGetArrayRead()`, `MatSeqAIJCUSPARSERestoreArrayWrite()`
3234: @*/
3235: PetscErrorCode MatSeqAIJCUSPARSEGetArrayWrite(Mat A, PetscScalar **a)
3236: {
3237: return MatSeqAIJCUSPARSE_CUPM_t::GetArrayWrite(A, a);
3238: }
3240: /*@
3241: MatSeqAIJCUSPARSERestoreArrayWrite - restore the write-only access array obtained from `MatSeqAIJCUSPARSEGetArrayWrite()`
3243: Not Collective
3245: Input Parameters:
3246: + A - a `MATSEQAIJCUSPARSE` matrix
3247: - a - pointer to the device data
3249: Level: developer
3251: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetArrayWrite()`
3252: @*/
3253: PetscErrorCode MatSeqAIJCUSPARSERestoreArrayWrite(Mat A, PetscScalar **a)
3254: {
3255: return MatSeqAIJCUSPARSE_CUPM_t::RestoreArrayWrite(A, a);
3256: }
3258: struct IJCompare4 {
3259: __host__ __device__ inline bool operator()(const thrust::tuple<PetscInt, PetscInt, PetscScalar, PetscInt> &t1, const thrust::tuple<PetscInt, PetscInt, PetscScalar, PetscInt> &t2)
3260: {
3261: if (thrust::get<0>(t1) < thrust::get<0>(t2)) return true;
3262: if (thrust::get<0>(t1) == thrust::get<0>(t2)) return thrust::get<1>(t1) < thrust::get<1>(t2);
3263: return false;
3264: }
3265: };
3267: struct Shift {
3268: PetscInt _shift;
3270: Shift(PetscInt shift) : _shift(shift) { }
3271: __host__ __device__ inline PetscInt operator()(const PetscInt &c) { return c + _shift; }
3272: };
3274: /* merges two SeqAIJCUSPARSE matrices A, B by concatenating their rows. [A';B']' operation in MATLAB notation */
3275: PetscErrorCode MatSeqAIJCUSPARSEMergeMats(Mat A, Mat B, MatReuse reuse, Mat *C)
3276: {
3277: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data, *b = (Mat_SeqAIJ *)B->data, *c;
3278: Mat_SeqAIJCUSPARSE *Acusp = (Mat_SeqAIJCUSPARSE *)A->spptr, *Bcusp = (Mat_SeqAIJCUSPARSE *)B->spptr, *Ccusp;
3279: Mat_SeqAIJCUSPARSEMultStruct *Cmat;
3280: CsrMatrix *Acsr, *Bcsr, *Ccsr;
3281: PetscInt Annz, Bnnz;
3282: PetscInt i, m, n, zero = 0;
3284: PetscFunctionBegin;
3287: PetscAssertPointer(C, 4);
3288: PetscCheckTypeName(A, MATSEQAIJCUSPARSE);
3289: PetscCheckTypeName(B, MATSEQAIJCUSPARSE);
3290: PetscCheck(A->rmap->n == B->rmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Invalid number or rows %" PetscInt_FMT " != %" PetscInt_FMT, A->rmap->n, B->rmap->n);
3291: PetscCheck(reuse != MAT_INPLACE_MATRIX, PETSC_COMM_SELF, PETSC_ERR_SUP, "MAT_INPLACE_MATRIX not supported");
3292: PetscCheck(Acusp->format != MAT_CUSPARSE_ELL && Acusp->format != MAT_CUSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
3293: PetscCheck(Bcusp->format != MAT_CUSPARSE_ELL && Bcusp->format != MAT_CUSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
3294: if (reuse == MAT_INITIAL_MATRIX) {
3295: m = A->rmap->n;
3296: n = A->cmap->n + B->cmap->n;
3297: PetscCall(MatCreate(PETSC_COMM_SELF, C));
3298: PetscCall(MatSetSizes(*C, m, n, m, n));
3299: PetscCall(MatSetType(*C, MATSEQAIJCUSPARSE));
3300: c = (Mat_SeqAIJ *)(*C)->data;
3301: Ccusp = (Mat_SeqAIJCUSPARSE *)(*C)->spptr;
3302: Cmat = new Mat_SeqAIJCUSPARSEMultStruct;
3303: Ccsr = new CsrMatrix;
3304: Cmat->cprowIndices = NULL;
3305: c->compressedrow.use = PETSC_FALSE;
3306: c->compressedrow.nrows = 0;
3307: c->compressedrow.i = NULL;
3308: c->compressedrow.rindex = NULL;
3309: Ccusp->workVector = NULL;
3310: Ccusp->nrows = m;
3311: Ccusp->mat = Cmat;
3312: Ccusp->mat->mat = Ccsr;
3313: Ccsr->num_rows = m;
3314: Ccsr->num_cols = n;
3315: PetscCallCUSPARSE(cusparseCreateMatDescr(&Cmat->descr));
3316: PetscCallCUSPARSE(cusparseSetMatIndexBase(Cmat->descr, CUSPARSE_INDEX_BASE_ZERO));
3317: PetscCallCUSPARSE(cusparseSetMatType(Cmat->descr, CUSPARSE_MATRIX_TYPE_GENERAL));
3318: PetscCallCUDA(cudaMalloc((void **)&Cmat->alpha_one, sizeof(PetscScalar)));
3319: PetscCallCUDA(cudaMalloc((void **)&Cmat->beta_zero, sizeof(PetscScalar)));
3320: PetscCallCUDA(cudaMalloc((void **)&Cmat->beta_one, sizeof(PetscScalar)));
3321: PetscCallCUDA(cudaMemcpy(Cmat->alpha_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
3322: PetscCallCUDA(cudaMemcpy(Cmat->beta_zero, &PETSC_CUSPARSE_ZERO, sizeof(PetscScalar), cudaMemcpyHostToDevice));
3323: PetscCallCUDA(cudaMemcpy(Cmat->beta_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
3324: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
3325: PetscCall(MatSeqAIJCUSPARSECopyToGPU(B));
3326: PetscCheck(Acusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJCUSPARSEMultStruct");
3327: PetscCheck(Bcusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJCUSPARSEMultStruct");
3329: Acsr = (CsrMatrix *)Acusp->mat->mat;
3330: Bcsr = (CsrMatrix *)Bcusp->mat->mat;
3331: Annz = (PetscInt)Acsr->column_indices->size();
3332: Bnnz = (PetscInt)Bcsr->column_indices->size();
3333: c->nz = Annz + Bnnz;
3334: Ccsr->row_offsets = new THRUSTINTARRAY(m + 1);
3335: Ccsr->column_indices = new THRUSTINTARRAY(c->nz);
3336: Ccsr->values = new THRUSTARRAY(c->nz);
3337: Ccsr->num_entries = c->nz;
3338: Ccusp->coords = new THRUSTINTARRAY(c->nz);
3339: if (c->nz) {
3340: auto Acoo = new THRUSTINTARRAY(Annz); // initialized with zeros
3341: auto Bcoo = new THRUSTINTARRAY(Bnnz);
3342: auto Ccoo = new THRUSTINTARRAY(c->nz);
3343: THRUSTINTARRAY *Aroff, *Broff;
3345: if (a->compressedrow.use) { /* need full row offset */
3346: if (!Acusp->rowoffsets_gpu) {
3347: Acusp->rowoffsets_gpu = new THRUSTINTARRAY(A->rmap->n + 1);
3348: Acusp->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
3349: PetscCall(PetscLogCpuToGpu((A->rmap->n + 1) * sizeof(PetscInt)));
3350: }
3351: Aroff = Acusp->rowoffsets_gpu;
3352: } else Aroff = Acsr->row_offsets;
3353: if (b->compressedrow.use) { /* need full row offset */
3354: if (!Bcusp->rowoffsets_gpu) {
3355: Bcusp->rowoffsets_gpu = new THRUSTINTARRAY(B->rmap->n + 1);
3356: Bcusp->rowoffsets_gpu->assign(b->i, b->i + B->rmap->n + 1);
3357: PetscCall(PetscLogCpuToGpu((B->rmap->n + 1) * sizeof(PetscInt)));
3358: }
3359: Broff = Bcusp->rowoffsets_gpu;
3360: } else Broff = Bcsr->row_offsets;
3361: PetscCall(PetscLogGpuTimeBegin());
3362: // Implement cusparseXcsr2coo() with Thrust, as the former doesn't support 64-bit indices.
3363: PetscCallThrust(thrust::for_each(thrust::device, thrust::counting_iterator<PetscInt>(0), thrust::counting_iterator<PetscInt>(m), Csr2coo(Aroff->data().get(), Acoo->data().get())));
3364: PetscCallThrust(thrust::for_each(thrust::device, thrust::counting_iterator<PetscInt>(0), thrust::counting_iterator<PetscInt>(m), Csr2coo(Broff->data().get(), Bcoo->data().get())));
3366: /* Issues when using bool with large matrices on SUMMIT 10.2.89 */
3367: #if CCCL_VERSION >= 3004000
3368: auto Aperm = cuda::make_constant_iterator(1);
3369: auto Bperm = cuda::make_constant_iterator(0);
3370: #else
3371: auto Aperm = thrust::make_constant_iterator(1);
3372: auto Bperm = thrust::make_constant_iterator(0);
3373: #endif
3374: auto Bcib = thrust::make_transform_iterator(Bcsr->column_indices->begin(), Shift(A->cmap->n));
3375: auto Bcie = thrust::make_transform_iterator(Bcsr->column_indices->end(), Shift(A->cmap->n));
3376: auto wPerm = new THRUSTINTARRAY(Annz + Bnnz);
3377: auto Azb = thrust::make_zip_iterator(thrust::make_tuple(Acoo->begin(), Acsr->column_indices->begin(), Acsr->values->begin(), Aperm));
3378: auto Aze = thrust::make_zip_iterator(thrust::make_tuple(Acoo->end(), Acsr->column_indices->end(), Acsr->values->end(), Aperm));
3379: auto Bzb = thrust::make_zip_iterator(thrust::make_tuple(Bcoo->begin(), Bcib, Bcsr->values->begin(), Bperm)); // Use B column indices shifted by A->cmap->n
3380: auto Bze = thrust::make_zip_iterator(thrust::make_tuple(Bcoo->end(), Bcie, Bcsr->values->end(), Bperm));
3381: auto Czb = thrust::make_zip_iterator(thrust::make_tuple(Ccoo->begin(), Ccsr->column_indices->begin(), Ccsr->values->begin(), wPerm->begin()));
3382: auto p1 = Ccusp->coords->begin();
3383: auto p2 = Ccusp->coords->begin();
3384: #if CCCL_VERSION >= 3001000
3385: cuda::std::advance(p2, Annz);
3386: #else
3387: thrust::advance(p2, Annz);
3388: #endif
3389: PetscCallThrust(thrust::merge(thrust::device, Azb, Aze, Bzb, Bze, Czb, IJCompare4())); // put nonzeros in A and B to C in sorted order (by row and then by column)
3390: auto cci = thrust::make_counting_iterator(zero);
3391: auto cce = thrust::make_counting_iterator(c->nz);
3392: #if THRUST_VERSION < 200800
3393: auto pred = thrust::identity<int>();
3394: #else
3395: auto pred = cuda::std::identity();
3396: #endif
3397: PetscCallThrust(thrust::copy_if(thrust::device, cci, cce, wPerm->begin(), p1, pred));
3398: PetscCallThrust(thrust::remove_copy_if(thrust::device, cci, cce, wPerm->begin(), p2, pred));
3399: // Implement a simplified cusparseXcoo2csr() with Thrust (assuming the row indices are already sorted), as the former doesn't support 64-bit indices.
3400: PetscCallThrust(thrust::lower_bound(thrust::device, Ccoo->begin(), Ccoo->end(), thrust::counting_iterator<PetscInt>(0), thrust::counting_iterator<PetscInt>(m + 1), Ccsr->row_offsets->begin()));
3401: PetscCall(PetscLogGpuTimeEnd());
3402: delete wPerm;
3403: delete Acoo;
3404: delete Bcoo;
3405: delete Ccoo;
3406: PetscCallCUSPARSE(cusparseCreateCsr(&Cmat->matDescr, Ccsr->num_rows, Ccsr->num_cols, Ccsr->num_entries, Ccsr->row_offsets->data().get(), Ccsr->column_indices->data().get(), Ccsr->values->data().get(), csrRowOffsetsType, csrColIndType, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
3407: if (A->form_explicit_transpose && B->form_explicit_transpose) { /* if A and B have the transpose, generate C transpose too */
3408: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(A));
3409: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(B));
3410: PetscBool AT = Acusp->matTranspose ? PETSC_TRUE : PETSC_FALSE, BT = Bcusp->matTranspose ? PETSC_TRUE : PETSC_FALSE;
3411: Mat_SeqAIJCUSPARSEMultStruct *CmatT = new Mat_SeqAIJCUSPARSEMultStruct;
3412: CsrMatrix *CcsrT = new CsrMatrix;
3413: CsrMatrix *AcsrT = AT ? (CsrMatrix *)Acusp->matTranspose->mat : NULL;
3414: CsrMatrix *BcsrT = BT ? (CsrMatrix *)Bcusp->matTranspose->mat : NULL;
3416: (*C)->form_explicit_transpose = PETSC_TRUE;
3417: (*C)->transupdated = PETSC_TRUE;
3418: Ccusp->rowoffsets_gpu = NULL;
3419: CmatT->cprowIndices = NULL;
3420: CmatT->mat = CcsrT;
3421: CcsrT->num_rows = n;
3422: CcsrT->num_cols = m;
3423: CcsrT->num_entries = c->nz;
3425: CcsrT->row_offsets = new THRUSTINTARRAY(n + 1);
3426: CcsrT->column_indices = new THRUSTINTARRAY(c->nz);
3427: CcsrT->values = new THRUSTARRAY(c->nz);
3429: PetscCall(PetscLogGpuTimeBegin());
3430: auto rT = CcsrT->row_offsets->begin();
3431: if (AT) {
3432: rT = thrust::copy(AcsrT->row_offsets->begin(), AcsrT->row_offsets->end(), rT);
3433: #if CCCL_VERSION >= 3001000
3434: cuda::std::advance(rT, -1);
3435: #else
3436: thrust::advance(rT, -1);
3437: #endif
3438: }
3439: if (BT) {
3440: auto titb = thrust::make_transform_iterator(BcsrT->row_offsets->begin(), Shift(a->nz));
3441: auto tite = thrust::make_transform_iterator(BcsrT->row_offsets->end(), Shift(a->nz));
3442: thrust::copy(titb, tite, rT);
3443: }
3444: auto cT = CcsrT->column_indices->begin();
3445: if (AT) cT = thrust::copy(AcsrT->column_indices->begin(), AcsrT->column_indices->end(), cT);
3446: if (BT) thrust::copy(BcsrT->column_indices->begin(), BcsrT->column_indices->end(), cT);
3447: auto vT = CcsrT->values->begin();
3448: if (AT) vT = thrust::copy(AcsrT->values->begin(), AcsrT->values->end(), vT);
3449: if (BT) thrust::copy(BcsrT->values->begin(), BcsrT->values->end(), vT);
3450: PetscCall(PetscLogGpuTimeEnd());
3452: PetscCallCUSPARSE(cusparseCreateMatDescr(&CmatT->descr));
3453: PetscCallCUSPARSE(cusparseSetMatIndexBase(CmatT->descr, CUSPARSE_INDEX_BASE_ZERO));
3454: PetscCallCUSPARSE(cusparseSetMatType(CmatT->descr, CUSPARSE_MATRIX_TYPE_GENERAL));
3455: PetscCallCUDA(cudaMalloc((void **)&CmatT->alpha_one, sizeof(PetscScalar)));
3456: PetscCallCUDA(cudaMalloc((void **)&CmatT->beta_zero, sizeof(PetscScalar)));
3457: PetscCallCUDA(cudaMalloc((void **)&CmatT->beta_one, sizeof(PetscScalar)));
3458: PetscCallCUDA(cudaMemcpy(CmatT->alpha_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
3459: PetscCallCUDA(cudaMemcpy(CmatT->beta_zero, &PETSC_CUSPARSE_ZERO, sizeof(PetscScalar), cudaMemcpyHostToDevice));
3460: PetscCallCUDA(cudaMemcpy(CmatT->beta_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
3461: PetscCallCUSPARSE(cusparseCreateCsr(&CmatT->matDescr, CcsrT->num_rows, CcsrT->num_cols, CcsrT->num_entries, CcsrT->row_offsets->data().get(), CcsrT->column_indices->data().get(), CcsrT->values->data().get(), csrRowOffsetsType, csrColIndType, CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype));
3462: Ccusp->matTranspose = CmatT;
3463: }
3464: }
3466: c->free_a = PETSC_TRUE;
3467: PetscCall(PetscShmgetAllocateArray(c->nz, sizeof(PetscInt), (void **)&c->j));
3468: PetscCall(PetscShmgetAllocateArray(m + 1, sizeof(PetscInt), (void **)&c->i));
3469: c->free_ij = PETSC_TRUE;
3470: PetscCallCUDA(cudaMemcpy(c->i, Ccsr->row_offsets->data().get(), Ccsr->row_offsets->size() * sizeof(PetscInt), cudaMemcpyDeviceToHost));
3471: PetscCallCUDA(cudaMemcpy(c->j, Ccsr->column_indices->data().get(), Ccsr->column_indices->size() * sizeof(PetscInt), cudaMemcpyDeviceToHost));
3472: PetscCall(PetscLogGpuToCpu((Ccsr->column_indices->size() + Ccsr->row_offsets->size()) * sizeof(PetscInt)));
3473: PetscCall(PetscMalloc1(m, &c->ilen));
3474: PetscCall(PetscMalloc1(m, &c->imax));
3475: c->maxnz = c->nz;
3476: c->nonzerorowcnt = 0;
3477: c->rmax = 0;
3478: for (i = 0; i < m; i++) {
3479: const PetscInt nn = c->i[i + 1] - c->i[i];
3480: c->ilen[i] = c->imax[i] = nn;
3481: c->nonzerorowcnt += (PetscInt)!!nn;
3482: c->rmax = PetscMax(c->rmax, nn);
3483: }
3484: PetscCall(PetscMalloc1(c->nz, &c->a));
3485: (*C)->nonzerostate++;
3486: PetscCall(PetscLayoutSetUp((*C)->rmap));
3487: PetscCall(PetscLayoutSetUp((*C)->cmap));
3488: Ccusp->nonzerostate = (*C)->nonzerostate;
3489: (*C)->preallocated = PETSC_TRUE;
3490: } else {
3491: PetscCheck((*C)->rmap->n == B->rmap->n, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Invalid number or rows %" PetscInt_FMT " != %" PetscInt_FMT, (*C)->rmap->n, B->rmap->n);
3492: c = (Mat_SeqAIJ *)(*C)->data;
3493: if (c->nz) {
3494: Ccusp = (Mat_SeqAIJCUSPARSE *)(*C)->spptr;
3495: PetscCheck(Ccusp->coords, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing coords");
3496: PetscCheck(Ccusp->format != MAT_CUSPARSE_ELL && Ccusp->format != MAT_CUSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
3497: PetscCheck(Ccusp->nonzerostate == (*C)->nonzerostate, PETSC_COMM_SELF, PETSC_ERR_COR, "Wrong nonzerostate");
3498: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
3499: PetscCall(MatSeqAIJCUSPARSECopyToGPU(B));
3500: PetscCheck(Acusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJCUSPARSEMultStruct");
3501: PetscCheck(Bcusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJCUSPARSEMultStruct");
3502: Acsr = (CsrMatrix *)Acusp->mat->mat;
3503: Bcsr = (CsrMatrix *)Bcusp->mat->mat;
3504: Ccsr = (CsrMatrix *)Ccusp->mat->mat;
3505: PetscCheck(Acsr->num_entries == (PetscInt)Acsr->values->size(), PETSC_COMM_SELF, PETSC_ERR_COR, "A nnz %" PetscInt_FMT " != %" PetscInt_FMT, Acsr->num_entries, (PetscInt)Acsr->values->size());
3506: PetscCheck(Bcsr->num_entries == (PetscInt)Bcsr->values->size(), PETSC_COMM_SELF, PETSC_ERR_COR, "B nnz %" PetscInt_FMT " != %" PetscInt_FMT, Bcsr->num_entries, (PetscInt)Bcsr->values->size());
3507: PetscCheck(Ccsr->num_entries == (PetscInt)Ccsr->values->size(), PETSC_COMM_SELF, PETSC_ERR_COR, "C nnz %" PetscInt_FMT " != %" PetscInt_FMT, Ccsr->num_entries, (PetscInt)Ccsr->values->size());
3508: PetscCheck(Ccsr->num_entries == Acsr->num_entries + Bcsr->num_entries, PETSC_COMM_SELF, PETSC_ERR_COR, "C nnz %" PetscInt_FMT " != %" PetscInt_FMT " + %" PetscInt_FMT, Ccsr->num_entries, Acsr->num_entries, Bcsr->num_entries);
3509: PetscCheck(Ccusp->coords->size() == Ccsr->values->size(), PETSC_COMM_SELF, PETSC_ERR_COR, "permSize %" PetscInt_FMT " != %" PetscInt_FMT, (PetscInt)Ccusp->coords->size(), (PetscInt)Ccsr->values->size());
3510: auto pmid = Ccusp->coords->begin();
3511: #if CCCL_VERSION >= 3001000
3512: cuda::std::advance(pmid, Acsr->num_entries);
3513: #else
3514: thrust::advance(pmid, Acsr->num_entries);
3515: #endif
3516: PetscCall(PetscLogGpuTimeBegin());
3517: auto zibait = thrust::make_zip_iterator(thrust::make_tuple(Acsr->values->begin(), thrust::make_permutation_iterator(Ccsr->values->begin(), Ccusp->coords->begin())));
3518: auto zieait = thrust::make_zip_iterator(thrust::make_tuple(Acsr->values->end(), thrust::make_permutation_iterator(Ccsr->values->begin(), pmid)));
3519: thrust::for_each(zibait, zieait, VecCUDAEquals());
3520: auto zibbit = thrust::make_zip_iterator(thrust::make_tuple(Bcsr->values->begin(), thrust::make_permutation_iterator(Ccsr->values->begin(), pmid)));
3521: auto ziebit = thrust::make_zip_iterator(thrust::make_tuple(Bcsr->values->end(), thrust::make_permutation_iterator(Ccsr->values->begin(), Ccusp->coords->end())));
3522: thrust::for_each(zibbit, ziebit, VecCUDAEquals());
3523: PetscCall(MatSeqAIJCUSPARSEInvalidateTranspose(*C, PETSC_FALSE));
3524: if (A->form_explicit_transpose && B->form_explicit_transpose && (*C)->form_explicit_transpose) {
3525: PetscCheck(Ccusp->matTranspose, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing transpose Mat_SeqAIJCUSPARSEMultStruct");
3526: PetscBool AT = Acusp->matTranspose ? PETSC_TRUE : PETSC_FALSE, BT = Bcusp->matTranspose ? PETSC_TRUE : PETSC_FALSE;
3527: CsrMatrix *AcsrT = AT ? (CsrMatrix *)Acusp->matTranspose->mat : NULL;
3528: CsrMatrix *BcsrT = BT ? (CsrMatrix *)Bcusp->matTranspose->mat : NULL;
3529: CsrMatrix *CcsrT = (CsrMatrix *)Ccusp->matTranspose->mat;
3530: auto vT = CcsrT->values->begin();
3531: if (AT) vT = thrust::copy(AcsrT->values->begin(), AcsrT->values->end(), vT);
3532: if (BT) thrust::copy(BcsrT->values->begin(), BcsrT->values->end(), vT);
3533: (*C)->transupdated = PETSC_TRUE;
3534: }
3535: PetscCall(PetscLogGpuTimeEnd());
3536: }
3537: }
3538: PetscCall(PetscObjectStateIncrease((PetscObject)*C));
3539: (*C)->assembled = PETSC_TRUE;
3540: (*C)->was_assembled = PETSC_FALSE;
3541: (*C)->offloadmask = PETSC_OFFLOAD_GPU;
3542: PetscFunctionReturn(PETSC_SUCCESS);
3543: }
3545: static PetscErrorCode MatSeqAIJCopySubArray_SeqAIJCUSPARSE(Mat A, PetscInt n, const PetscInt idx[], PetscScalar v[])
3546: {
3547: PetscFunctionBegin;
3548: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::CopySubArray(A, n, idx, v));
3549: PetscFunctionReturn(PETSC_SUCCESS);
3550: }