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