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;
1664: Mat_SeqAIJCUSPARSE *cusp;
1665: cusparseOperation_t opA;
1666: const PetscScalar *barray;
1667: PetscScalar *carray;
1668: MatProductCtx_MatMatCusparse *mmdata;
1669: Mat_SeqAIJCUSPARSEMultStruct *mat;
1670: CsrMatrix *csrmat;
1672: PetscFunctionBegin;
1673: MatCheckProduct(C, 1);
1674: PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data empty");
1675: mmdata = (MatProductCtx_MatMatCusparse *)product->data;
1676: A = product->A;
1677: B = product->B;
1678: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJCUSPARSE, &flg));
1679: PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
1680: /* currently CopyToGpu does not copy if the matrix is bound to CPU
1681: Instead of silently accepting the wrong answer, I prefer to raise the error */
1682: PetscCheck(!A->boundtocpu, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Cannot bind to CPU a CUSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
1683: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
1684: cusp = (Mat_SeqAIJCUSPARSE *)A->spptr;
1685: switch (product->type) {
1686: case MATPRODUCT_AB:
1687: case MATPRODUCT_PtAP:
1688: mat = cusp->mat;
1689: opA = CUSPARSE_OPERATION_NON_TRANSPOSE;
1690: m = A->rmap->n;
1691: n = B->cmap->n;
1692: break;
1693: case MATPRODUCT_AtB:
1694: if (!A->form_explicit_transpose) {
1695: mat = cusp->mat;
1696: opA = CUSPARSE_OPERATION_TRANSPOSE;
1697: } else {
1698: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(A));
1699: mat = cusp->matTranspose;
1700: opA = CUSPARSE_OPERATION_NON_TRANSPOSE;
1701: }
1702: m = A->cmap->n;
1703: n = B->cmap->n;
1704: break;
1705: case MATPRODUCT_ABt:
1706: case MATPRODUCT_RARt:
1707: mat = cusp->mat;
1708: opA = CUSPARSE_OPERATION_NON_TRANSPOSE;
1709: m = A->rmap->n;
1710: n = B->rmap->n;
1711: break;
1712: default:
1713: SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
1714: }
1715: PetscCheck(mat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing Mat_SeqAIJCUSPARSEMultStruct");
1716: csrmat = (CsrMatrix *)mat->mat;
1717: /* if the user passed a CPU matrix, copy the data to the GPU */
1718: PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQDENSECUDA, &biscuda));
1719: if (!biscuda) PetscCall(MatConvert(B, MATSEQDENSECUDA, MAT_INPLACE_MATRIX, &B));
1720: PetscCall(MatDenseGetArrayReadAndMemType(B, &barray, nullptr));
1722: PetscCall(MatDenseGetLDA(B, &blda));
1723: if (product->type == MATPRODUCT_RARt || product->type == MATPRODUCT_PtAP) {
1724: PetscCall(MatDenseGetArrayWriteAndMemType(mmdata->X, &carray, nullptr));
1725: PetscCall(MatDenseGetLDA(mmdata->X, &clda));
1726: } else {
1727: PetscCall(MatDenseGetArrayWriteAndMemType(C, &carray, nullptr));
1728: PetscCall(MatDenseGetLDA(C, &clda));
1729: }
1731: PetscCall(PetscLogGpuTimeBegin());
1732: cusparseOperation_t opB = (product->type == MATPRODUCT_ABt || product->type == MATPRODUCT_RARt) ? CUSPARSE_OPERATION_TRANSPOSE : CUSPARSE_OPERATION_NON_TRANSPOSE;
1733: #if PETSC_PKG_CUDA_VERSION_GE(12, 4, 0)
1734: cusparseSpMatDescr_t &matADescr = mat->matDescr_SpMM[opA];
1735: #else
1736: cusparseSpMatDescr_t &matADescr = mat->matDescr;
1737: #endif
1739: /* (re)allocate mmBuffer if not initialized or LDAs are different */
1740: if (!mmdata->initialized || mmdata->Blda != blda || mmdata->Clda != clda) {
1741: size_t mmBufferSize;
1742: if (mmdata->initialized && mmdata->Blda != blda) {
1743: PetscCallCUSPARSE(cusparseDestroyDnMat(mmdata->matBDescr));
1744: mmdata->matBDescr = NULL;
1745: }
1746: if (!mmdata->matBDescr) {
1747: PetscCallCUSPARSE(cusparseCreateDnMat(&mmdata->matBDescr, B->rmap->n, B->cmap->n, blda, (void *)barray, cusparse_scalartype, CUSPARSE_ORDER_COL));
1748: mmdata->Blda = blda;
1749: }
1751: if (mmdata->initialized && mmdata->Clda != clda) {
1752: PetscCallCUSPARSE(cusparseDestroyDnMat(mmdata->matCDescr));
1753: mmdata->matCDescr = NULL;
1754: }
1755: if (!mmdata->matCDescr) { /* matCDescr is for C or mmdata->X */
1756: PetscCallCUSPARSE(cusparseCreateDnMat(&mmdata->matCDescr, m, n, clda, (void *)carray, cusparse_scalartype, CUSPARSE_ORDER_COL));
1757: mmdata->Clda = clda;
1758: }
1760: #if PETSC_PKG_CUDA_VERSION_GE(12, 4, 0) // tested up to 12.6.0
1761: if (matADescr) {
1762: PetscCallCUSPARSE(cusparseDestroySpMat(matADescr)); // Because I find I could not reuse matADescr. It could be a cusparse bug
1763: matADescr = NULL;
1764: }
1765: #endif
1767: if (!matADescr) {
1768: 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));
1769: }
1771: PetscCallCUSPARSE(cusparseSpMM_bufferSize(cusp->handle, opA, opB, mat->alpha_one, matADescr, mmdata->matBDescr, mat->beta_zero, mmdata->matCDescr, cusparse_scalartype, cusp->spmmAlg, &mmBufferSize));
1773: if ((mmdata->mmBuffer && mmdata->mmBufferSize < mmBufferSize) || !mmdata->mmBuffer) {
1774: PetscCallCUDA(cudaFree(mmdata->mmBuffer));
1775: PetscCallCUDA(cudaMalloc(&mmdata->mmBuffer, mmBufferSize));
1776: mmdata->mmBufferSize = mmBufferSize;
1777: }
1779: #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
1780: PetscCallCUSPARSE(cusparseSpMM_preprocess(cusp->handle, opA, opB, mat->alpha_one, matADescr, mmdata->matBDescr, mat->beta_zero, mmdata->matCDescr, cusparse_scalartype, cusp->spmmAlg, mmdata->mmBuffer));
1781: #endif
1783: mmdata->initialized = PETSC_TRUE;
1784: } else {
1785: /* to be safe, always update pointers of the mats */
1786: PetscCallCUSPARSE(cusparseSpMatSetValues(matADescr, csrmat->values->data().get()));
1787: PetscCallCUSPARSE(cusparseDnMatSetValues(mmdata->matBDescr, (void *)barray));
1788: PetscCallCUSPARSE(cusparseDnMatSetValues(mmdata->matCDescr, (void *)carray));
1789: }
1791: /* do cusparseSpMM, which supports transpose on B */
1792: PetscCallCUSPARSE(cusparseSpMM(cusp->handle, opA, opB, mat->alpha_one, matADescr, mmdata->matBDescr, mat->beta_zero, mmdata->matCDescr, cusparse_scalartype, cusp->spmmAlg, mmdata->mmBuffer));
1794: PetscCall(PetscLogGpuTimeEnd());
1795: PetscCall(PetscLogGpuFlops(n * 2.0 * csrmat->num_entries));
1796: PetscCall(MatDenseRestoreArrayReadAndMemType(B, &barray));
1797: if (product->type == MATPRODUCT_RARt) {
1798: PetscCall(MatDenseRestoreArrayWriteAndMemType(mmdata->X, &carray));
1799: PetscCall(MatMatMultNumeric_SeqDenseCUDA_SeqDenseCUDA_Internal(B, mmdata->X, C, PETSC_FALSE, PETSC_FALSE));
1800: } else if (product->type == MATPRODUCT_PtAP) {
1801: PetscCall(MatDenseRestoreArrayWriteAndMemType(mmdata->X, &carray));
1802: PetscCall(MatMatMultNumeric_SeqDenseCUDA_SeqDenseCUDA_Internal(B, mmdata->X, C, PETSC_TRUE, PETSC_FALSE));
1803: } else {
1804: PetscCall(MatDenseRestoreArrayWriteAndMemType(C, &carray));
1805: }
1806: if (mmdata->cisdense) PetscCall(MatConvert(C, MATSEQDENSE, MAT_INPLACE_MATRIX, &C));
1807: if (!biscuda) PetscCall(MatConvert(B, MATSEQDENSE, MAT_INPLACE_MATRIX, &B));
1808: PetscFunctionReturn(PETSC_SUCCESS);
1809: }
1811: static PetscErrorCode MatProductSymbolic_SeqAIJCUSPARSE_SeqDENSECUDA(Mat C)
1812: {
1813: Mat_Product *product = C->product;
1814: Mat A, B;
1815: PetscInt m, n;
1816: PetscBool cisdense, flg;
1817: MatProductCtx_MatMatCusparse *mmdata;
1818: Mat_SeqAIJCUSPARSE *cusp;
1820: PetscFunctionBegin;
1821: MatCheckProduct(C, 1);
1822: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data not empty");
1823: A = product->A;
1824: B = product->B;
1825: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJCUSPARSE, &flg));
1826: PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
1827: cusp = (Mat_SeqAIJCUSPARSE *)A->spptr;
1828: PetscCheck(cusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
1829: switch (product->type) {
1830: case MATPRODUCT_AB:
1831: m = A->rmap->n;
1832: n = B->cmap->n;
1833: PetscCall(MatSetBlockSizesFromMats(C, A, B));
1834: break;
1835: case MATPRODUCT_AtB:
1836: m = A->cmap->n;
1837: n = B->cmap->n;
1838: if (A->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->rmap, A->cmap->bs));
1839: if (B->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->cmap, B->cmap->bs));
1840: break;
1841: case MATPRODUCT_ABt:
1842: m = A->rmap->n;
1843: n = B->rmap->n;
1844: if (A->rmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->rmap, A->rmap->bs));
1845: if (B->rmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->cmap, B->rmap->bs));
1846: break;
1847: case MATPRODUCT_PtAP:
1848: m = B->cmap->n;
1849: n = B->cmap->n;
1850: if (B->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->rmap, B->cmap->bs));
1851: if (B->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->cmap, B->cmap->bs));
1852: break;
1853: case MATPRODUCT_RARt:
1854: m = B->rmap->n;
1855: n = B->rmap->n;
1856: if (B->rmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->rmap, B->rmap->bs));
1857: if (B->rmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->cmap, B->rmap->bs));
1858: break;
1859: default:
1860: SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
1861: }
1862: PetscCall(MatSetSizes(C, m, n, m, n));
1863: /* if C is of type MATSEQDENSE (CPU), perform the operation on the GPU and then copy on the CPU */
1864: PetscCall(PetscObjectTypeCompare((PetscObject)C, MATSEQDENSE, &cisdense));
1865: PetscCall(MatSetType(C, MATSEQDENSECUDA));
1867: /* product data */
1868: PetscCall(PetscNew(&mmdata));
1869: mmdata->cisdense = cisdense;
1870: /* for these products we need intermediate storage */
1871: if (product->type == MATPRODUCT_RARt || product->type == MATPRODUCT_PtAP) {
1872: PetscCall(MatCreate(PetscObjectComm((PetscObject)C), &mmdata->X));
1873: PetscCall(MatSetType(mmdata->X, MATSEQDENSECUDA));
1874: if (product->type == MATPRODUCT_RARt) { /* do not preallocate, since the first call to MatDenseCUDAGetArray will preallocate on the GPU for us */
1875: PetscCall(MatSetSizes(mmdata->X, A->rmap->n, B->rmap->n, A->rmap->n, B->rmap->n));
1876: } else {
1877: PetscCall(MatSetSizes(mmdata->X, A->rmap->n, B->cmap->n, A->rmap->n, B->cmap->n));
1878: }
1879: }
1880: C->product->data = mmdata;
1881: C->product->destroy = MatProductCtxDestroy_MatMatCusparse;
1883: C->ops->productnumeric = MatProductNumeric_SeqAIJCUSPARSE_SeqDENSECUDA;
1884: PetscFunctionReturn(PETSC_SUCCESS);
1885: }
1887: static PetscErrorCode MatProductNumeric_SeqAIJCUSPARSE_SeqAIJCUSPARSE(Mat C)
1888: {
1889: Mat_Product *product = C->product;
1890: Mat A, B;
1891: Mat_SeqAIJCUSPARSE *Acusp, *Bcusp, *Ccusp;
1892: Mat_SeqAIJ *c = (Mat_SeqAIJ *)C->data;
1893: Mat_SeqAIJCUSPARSEMultStruct *Amat, *Bmat, *Cmat;
1894: CsrMatrix *Acsr, *Bcsr, *Ccsr;
1895: PetscBool flg;
1896: MatProductType ptype;
1897: MatProductCtx_MatMatCusparse *mmdata;
1898: cusparseSpMatDescr_t BmatSpDescr;
1899: cusparseOperation_t opA = CUSPARSE_OPERATION_NON_TRANSPOSE, opB = CUSPARSE_OPERATION_NON_TRANSPOSE; /* cuSPARSE spgemm doesn't support transpose yet */
1901: PetscFunctionBegin;
1902: MatCheckProduct(C, 1);
1903: PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data empty");
1904: PetscCall(PetscObjectTypeCompare((PetscObject)C, MATSEQAIJCUSPARSE, &flg));
1905: PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for C of type %s", ((PetscObject)C)->type_name);
1906: mmdata = (MatProductCtx_MatMatCusparse *)C->product->data;
1907: A = product->A;
1908: B = product->B;
1909: if (mmdata->reusesym) { /* this happens when api_user is true, meaning that the matrix values have been already computed in the MatProductSymbolic phase */
1910: mmdata->reusesym = PETSC_FALSE;
1911: Ccusp = (Mat_SeqAIJCUSPARSE *)C->spptr;
1912: PetscCheck(Ccusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
1913: Cmat = Ccusp->mat;
1914: PetscCheck(Cmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C mult struct for product type %s", MatProductTypes[C->product->type]);
1915: Ccsr = (CsrMatrix *)Cmat->mat;
1916: PetscCheck(Ccsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C CSR struct");
1917: goto finalize;
1918: }
1919: if (!c->nz) goto finalize;
1920: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJCUSPARSE, &flg));
1921: PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
1922: PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQAIJCUSPARSE, &flg));
1923: PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for B of type %s", ((PetscObject)B)->type_name);
1924: PetscCheck(!A->boundtocpu, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONG, "Cannot bind to CPU a CUSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
1925: PetscCheck(!B->boundtocpu, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONG, "Cannot bind to CPU a CUSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
1926: Acusp = (Mat_SeqAIJCUSPARSE *)A->spptr;
1927: Bcusp = (Mat_SeqAIJCUSPARSE *)B->spptr;
1928: Ccusp = (Mat_SeqAIJCUSPARSE *)C->spptr;
1929: PetscCheck(Acusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
1930: PetscCheck(Bcusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
1931: PetscCheck(Ccusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
1932: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
1933: PetscCall(MatSeqAIJCUSPARSECopyToGPU(B));
1935: ptype = product->type;
1936: if (A->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_AtB) {
1937: ptype = MATPRODUCT_AB;
1938: 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");
1939: }
1940: if (B->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_ABt) {
1941: ptype = MATPRODUCT_AB;
1942: 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");
1943: }
1944: switch (ptype) {
1945: case MATPRODUCT_AB:
1946: Amat = Acusp->mat;
1947: Bmat = Bcusp->mat;
1948: break;
1949: case MATPRODUCT_AtB:
1950: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(A));
1951: Amat = Acusp->matTranspose;
1952: Bmat = Bcusp->mat;
1953: break;
1954: case MATPRODUCT_ABt:
1955: Amat = Acusp->mat;
1956: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(B));
1957: Bmat = Bcusp->matTranspose;
1958: break;
1959: default:
1960: SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
1961: }
1962: Cmat = Ccusp->mat;
1963: PetscCheck(Amat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A mult struct for product type %s", MatProductTypes[ptype]);
1964: PetscCheck(Bmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B mult struct for product type %s", MatProductTypes[ptype]);
1965: PetscCheck(Cmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C mult struct for product type %s", MatProductTypes[ptype]);
1966: Acsr = (CsrMatrix *)Amat->mat;
1967: Bcsr = mmdata->Bcsr ? mmdata->Bcsr : (CsrMatrix *)Bmat->mat; /* B may be in compressed row storage */
1968: Ccsr = (CsrMatrix *)Cmat->mat;
1969: PetscCheck(Acsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A CSR struct");
1970: PetscCheck(Bcsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B CSR struct");
1971: PetscCheck(Ccsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C CSR struct");
1972: PetscCall(PetscLogGpuTimeBegin());
1973: BmatSpDescr = mmdata->Bcsr ? mmdata->matSpBDescr : Bmat->matDescr; /* B may be in compressed row storage */
1974: PetscCallCUSPARSE(cusparseSetPointerMode(Ccusp->handle, CUSPARSE_POINTER_MODE_DEVICE));
1975: 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));
1976: 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));
1977: PetscCall(PetscLogGpuFlops(mmdata->flops));
1978: PetscCallCUDA(WaitForCUDA());
1979: PetscCall(PetscLogGpuTimeEnd());
1980: C->offloadmask = PETSC_OFFLOAD_GPU;
1981: finalize:
1982: /* shorter version of MatAssemblyEnd_SeqAIJ */
1983: 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));
1984: PetscCall(PetscInfo(C, "Number of mallocs during MatSetValues() is 0\n"));
1985: PetscCall(PetscInfo(C, "Maximum nonzeros in any row is %" PetscInt_FMT "\n", c->rmax));
1986: c->reallocs = 0;
1987: C->info.mallocs += 0;
1988: C->info.nz_unneeded = 0;
1989: C->assembled = C->was_assembled = PETSC_TRUE;
1990: C->num_ass++;
1991: PetscFunctionReturn(PETSC_SUCCESS);
1992: }
1994: static PetscErrorCode MatProductSymbolic_SeqAIJCUSPARSE_SeqAIJCUSPARSE(Mat C)
1995: {
1996: Mat_Product *product = C->product;
1997: Mat A, B;
1998: Mat_SeqAIJCUSPARSE *Acusp, *Bcusp, *Ccusp;
1999: Mat_SeqAIJ *a, *b, *c;
2000: Mat_SeqAIJCUSPARSEMultStruct *Amat, *Bmat, *Cmat;
2001: CsrMatrix *Acsr, *Bcsr, *Ccsr;
2002: PetscInt i, j, m, n, k;
2003: PetscBool flg;
2004: MatProductType ptype;
2005: MatProductCtx_MatMatCusparse *mmdata;
2006: PetscLogDouble flops;
2007: PetscBool biscompressed, ciscompressed;
2008: int64_t C_num_rows1, C_num_cols1, C_nnz1;
2009: cusparseSpMatDescr_t BmatSpDescr;
2010: cusparseOperation_t opA = CUSPARSE_OPERATION_NON_TRANSPOSE, opB = CUSPARSE_OPERATION_NON_TRANSPOSE; /* cuSPARSE spgemm doesn't support transpose yet */
2011: size_t bufSize2;
2013: PetscFunctionBegin;
2014: MatCheckProduct(C, 1);
2015: PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data not empty");
2016: A = product->A;
2017: B = product->B;
2018: PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJCUSPARSE, &flg));
2019: PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
2020: PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQAIJCUSPARSE, &flg));
2021: PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for B of type %s", ((PetscObject)B)->type_name);
2022: a = (Mat_SeqAIJ *)A->data;
2023: b = (Mat_SeqAIJ *)B->data;
2024: /* product data */
2025: PetscCall(PetscNew(&mmdata));
2026: C->product->data = mmdata;
2027: C->product->destroy = MatProductCtxDestroy_MatMatCusparse;
2029: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
2030: PetscCall(MatSeqAIJCUSPARSECopyToGPU(B));
2031: Acusp = (Mat_SeqAIJCUSPARSE *)A->spptr; /* Access spptr after MatSeqAIJCUSPARSECopyToGPU, not before */
2032: Bcusp = (Mat_SeqAIJCUSPARSE *)B->spptr;
2033: PetscCheck(Acusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
2034: PetscCheck(Bcusp->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_CUSPARSE_CSR format");
2036: ptype = product->type;
2037: if (A->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_AtB) {
2038: ptype = MATPRODUCT_AB;
2039: product->symbolic_used_the_fact_A_is_symmetric = PETSC_TRUE;
2040: }
2041: if (B->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_ABt) {
2042: ptype = MATPRODUCT_AB;
2043: product->symbolic_used_the_fact_B_is_symmetric = PETSC_TRUE;
2044: }
2045: biscompressed = PETSC_FALSE;
2046: ciscompressed = PETSC_FALSE;
2047: switch (ptype) {
2048: case MATPRODUCT_AB:
2049: m = A->rmap->n;
2050: n = B->cmap->n;
2051: k = A->cmap->n;
2052: Amat = Acusp->mat;
2053: Bmat = Bcusp->mat;
2054: if (a->compressedrow.use) ciscompressed = PETSC_TRUE;
2055: if (b->compressedrow.use) biscompressed = PETSC_TRUE;
2056: break;
2057: case MATPRODUCT_AtB:
2058: m = A->cmap->n;
2059: n = B->cmap->n;
2060: k = A->rmap->n;
2061: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(A));
2062: Amat = Acusp->matTranspose;
2063: Bmat = Bcusp->mat;
2064: if (b->compressedrow.use) biscompressed = PETSC_TRUE;
2065: break;
2066: case MATPRODUCT_ABt:
2067: m = A->rmap->n;
2068: n = B->rmap->n;
2069: k = A->cmap->n;
2070: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(B));
2071: Amat = Acusp->mat;
2072: Bmat = Bcusp->matTranspose;
2073: if (a->compressedrow.use) ciscompressed = PETSC_TRUE;
2074: break;
2075: default:
2076: SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
2077: }
2079: /* create cusparse matrix */
2080: PetscCall(MatSetSizes(C, m, n, m, n));
2081: PetscCall(MatSetType(C, MATSEQAIJCUSPARSE));
2082: c = (Mat_SeqAIJ *)C->data;
2083: Ccusp = (Mat_SeqAIJCUSPARSE *)C->spptr;
2084: Cmat = new Mat_SeqAIJCUSPARSEMultStruct;
2085: Ccsr = new CsrMatrix;
2087: c->compressedrow.use = ciscompressed;
2088: if (c->compressedrow.use) { /* if a is in compressed row, than c will be in compressed row format */
2089: c->compressedrow.nrows = a->compressedrow.nrows;
2090: PetscCall(PetscMalloc2(c->compressedrow.nrows + 1, &c->compressedrow.i, c->compressedrow.nrows, &c->compressedrow.rindex));
2091: PetscCall(PetscArraycpy(c->compressedrow.rindex, a->compressedrow.rindex, c->compressedrow.nrows));
2092: Ccusp->workVector = new THRUSTARRAY(c->compressedrow.nrows);
2093: Cmat->cprowIndices = new THRUSTINTARRAY(c->compressedrow.nrows);
2094: Cmat->cprowIndices->assign(c->compressedrow.rindex, c->compressedrow.rindex + c->compressedrow.nrows);
2095: } else {
2096: c->compressedrow.nrows = 0;
2097: c->compressedrow.i = NULL;
2098: c->compressedrow.rindex = NULL;
2099: Ccusp->workVector = NULL;
2100: Cmat->cprowIndices = NULL;
2101: }
2102: Ccusp->nrows = ciscompressed ? c->compressedrow.nrows : m;
2103: Ccusp->mat = Cmat;
2104: Ccusp->mat->mat = Ccsr;
2105: Ccsr->num_rows = Ccusp->nrows;
2106: Ccsr->num_cols = n;
2107: Ccsr->row_offsets = new THRUSTINTARRAY(Ccusp->nrows + 1);
2108: PetscCallCUSPARSE(cusparseCreateMatDescr(&Cmat->descr));
2109: PetscCallCUSPARSE(cusparseSetMatIndexBase(Cmat->descr, CUSPARSE_INDEX_BASE_ZERO));
2110: PetscCallCUSPARSE(cusparseSetMatType(Cmat->descr, CUSPARSE_MATRIX_TYPE_GENERAL));
2111: PetscCallCUDA(cudaMalloc((void **)&Cmat->alpha_one, sizeof(PetscScalar)));
2112: PetscCallCUDA(cudaMalloc((void **)&Cmat->beta_zero, sizeof(PetscScalar)));
2113: PetscCallCUDA(cudaMalloc((void **)&Cmat->beta_one, sizeof(PetscScalar)));
2114: PetscCallCUDA(cudaMemcpy(Cmat->alpha_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
2115: PetscCallCUDA(cudaMemcpy(Cmat->beta_zero, &PETSC_CUSPARSE_ZERO, sizeof(PetscScalar), cudaMemcpyHostToDevice));
2116: PetscCallCUDA(cudaMemcpy(Cmat->beta_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
2117: if (!Ccsr->num_rows || !Ccsr->num_cols || !a->nz || !b->nz) { /* cusparse raise errors in different calls when matrices have zero rows/columns! */
2118: PetscCallThrust(thrust::fill(thrust::device, Ccsr->row_offsets->begin(), Ccsr->row_offsets->end(), 0));
2119: c->nz = 0;
2120: Ccsr->column_indices = new THRUSTINTARRAY(c->nz);
2121: Ccsr->values = new THRUSTARRAY(c->nz);
2122: goto finalizesym;
2123: }
2125: PetscCheck(Amat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A mult struct for product type %s", MatProductTypes[ptype]);
2126: PetscCheck(Bmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B mult struct for product type %s", MatProductTypes[ptype]);
2127: Acsr = (CsrMatrix *)Amat->mat;
2128: if (!biscompressed) {
2129: Bcsr = (CsrMatrix *)Bmat->mat;
2130: BmatSpDescr = Bmat->matDescr;
2131: } else { /* we need to use row offsets for the full matrix */
2132: CsrMatrix *cBcsr = (CsrMatrix *)Bmat->mat;
2133: Bcsr = new CsrMatrix;
2134: Bcsr->num_rows = B->rmap->n;
2135: Bcsr->num_cols = cBcsr->num_cols;
2136: Bcsr->num_entries = cBcsr->num_entries;
2137: Bcsr->column_indices = cBcsr->column_indices;
2138: Bcsr->values = cBcsr->values;
2139: if (!Bcusp->rowoffsets_gpu) {
2140: Bcusp->rowoffsets_gpu = new THRUSTINTARRAY(B->rmap->n + 1);
2141: Bcusp->rowoffsets_gpu->assign(b->i, b->i + B->rmap->n + 1);
2142: PetscCall(PetscLogCpuToGpu((B->rmap->n + 1) * sizeof(PetscInt)));
2143: }
2144: Bcsr->row_offsets = Bcusp->rowoffsets_gpu;
2145: mmdata->Bcsr = Bcsr;
2146: if (Bcsr->num_rows && Bcsr->num_cols) {
2147: 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));
2148: }
2149: BmatSpDescr = mmdata->matSpBDescr;
2150: }
2151: PetscCheck(Acsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A CSR struct");
2152: PetscCheck(Bcsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B CSR struct");
2153: /* precompute flops count */
2154: if (ptype == MATPRODUCT_AB) {
2155: for (i = 0, flops = 0; i < A->rmap->n; i++) {
2156: const PetscInt st = a->i[i];
2157: const PetscInt en = a->i[i + 1];
2158: for (j = st; j < en; j++) {
2159: const PetscInt brow = a->j[j];
2160: flops += 2. * (b->i[brow + 1] - b->i[brow]);
2161: }
2162: }
2163: } else if (ptype == MATPRODUCT_AtB) {
2164: for (i = 0, flops = 0; i < A->rmap->n; i++) {
2165: const PetscInt anzi = a->i[i + 1] - a->i[i];
2166: const PetscInt bnzi = b->i[i + 1] - b->i[i];
2167: flops += (2. * anzi) * bnzi;
2168: }
2169: } else { /* TODO */
2170: flops = 0.;
2171: }
2173: mmdata->flops = flops;
2174: PetscCall(PetscLogGpuTimeBegin());
2176: PetscCallCUSPARSE(cusparseSetPointerMode(Ccusp->handle, CUSPARSE_POINTER_MODE_DEVICE));
2177: // cuda-12.2 requires non-null csrRowOffsets
2178: 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));
2179: PetscCallCUSPARSE(cusparseSpGEMM_createDescr(&mmdata->spgemmDesc));
2180: // Note that cusparseSpGEMMreuse is deprecated in CUDA 13.2.1
2182: 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.");
2183: /* ask bufferSize bytes for external memory */
2184: 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));
2185: PetscCallCUDA(cudaMalloc((void **)&mmdata->mmBuffer2, bufSize2));
2186: /* inspect the matrices A and B to understand the memory requirement for the next step */
2187: 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));
2188: /* ask bufferSize again bytes for external memory */
2189: 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));
2190: /* The CUSPARSE documentation is not clear, nor the API
2191: We need both buffers to perform the operations properly!
2192: mmdata->mmBuffer2 does not appear anywhere in the compute/copy API
2193: it only appears for the workEstimation stuff, but it seems it is needed in compute, so probably the address
2194: is stored in the descriptor! What a messy API... */
2195: PetscCallCUDA(cudaMalloc((void **)&mmdata->mmBuffer, mmdata->mmBufferSize));
2196: /* compute the intermediate product of A * B */
2197: 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));
2198: /* get matrix C non-zero entries C_nnz1 */
2199: PetscCallCUSPARSE(cusparseSpMatGetSize(Cmat->matDescr, &C_num_rows1, &C_num_cols1, &C_nnz1));
2200: PetscCall(PetscIntCast(C_nnz1, &c->nz));
2201: 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,
2202: mmdata->mmBufferSize / 1024));
2203: Ccsr->column_indices = new THRUSTINTARRAY(c->nz);
2204: PetscCallCUDA(cudaPeekAtLastError()); /* catch out of memory errors */
2205: Ccsr->values = new THRUSTARRAY(c->nz);
2206: PetscCallCUDA(cudaPeekAtLastError()); /* catch out of memory errors */
2207: if (c->nz) PetscCallCUSPARSE(cusparseCsrSetPointers(Cmat->matDescr, Ccsr->row_offsets->data().get(), Ccsr->column_indices->data().get(), Ccsr->values->data().get()));
2208: 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));
2209: PetscCall(PetscLogGpuFlops(mmdata->flops));
2210: PetscCall(PetscLogGpuTimeEnd());
2211: finalizesym:
2212: c->free_a = PETSC_TRUE;
2213: PetscCall(PetscShmgetAllocateArray(c->nz, sizeof(PetscInt), (void **)&c->j));
2214: PetscCall(PetscShmgetAllocateArray(m + 1, sizeof(PetscInt), (void **)&c->i));
2215: c->free_ij = PETSC_TRUE;
2217: PetscInt *d_i = c->i;
2218: if (ciscompressed) d_i = c->compressedrow.i;
2219: PetscCallCUDA(cudaMemcpy(d_i, Ccsr->row_offsets->data().get(), Ccsr->row_offsets->size() * sizeof(PetscInt), cudaMemcpyDeviceToHost));
2220: PetscCallCUDA(cudaMemcpy(c->j, Ccsr->column_indices->data().get(), Ccsr->column_indices->size() * sizeof(PetscInt), cudaMemcpyDeviceToHost));
2221: if (ciscompressed) { /* need to expand host row offsets */
2222: PetscInt r = 0;
2223: c->i[0] = 0;
2224: for (k = 0; k < c->compressedrow.nrows; k++) {
2225: const PetscInt next = c->compressedrow.rindex[k];
2226: const PetscInt old = c->compressedrow.i[k];
2227: for (; r < next; r++) c->i[r + 1] = old;
2228: }
2229: for (; r < m; r++) c->i[r + 1] = c->compressedrow.i[c->compressedrow.nrows];
2230: }
2231: PetscCall(PetscLogGpuToCpu((Ccsr->column_indices->size() + Ccsr->row_offsets->size()) * sizeof(PetscInt)));
2232: PetscCall(PetscMalloc1(m, &c->ilen));
2233: PetscCall(PetscMalloc1(m, &c->imax));
2234: c->maxnz = c->nz;
2235: c->nonzerorowcnt = 0;
2236: c->rmax = 0;
2237: for (k = 0; k < m; k++) {
2238: const PetscInt nn = c->i[k + 1] - c->i[k];
2239: c->ilen[k] = c->imax[k] = nn;
2240: c->nonzerorowcnt += (PetscInt)!!nn;
2241: c->rmax = PetscMax(c->rmax, nn);
2242: }
2243: PetscCall(PetscMalloc1(c->nz, &c->a));
2244: Ccsr->num_entries = c->nz;
2246: C->nonzerostate++;
2247: PetscCall(PetscLayoutSetUp(C->rmap));
2248: PetscCall(PetscLayoutSetUp(C->cmap));
2249: Ccusp->nonzerostate = C->nonzerostate;
2250: C->offloadmask = PETSC_OFFLOAD_UNALLOCATED;
2251: C->preallocated = PETSC_TRUE;
2252: C->assembled = PETSC_FALSE;
2253: C->was_assembled = PETSC_FALSE;
2254: 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 */
2255: mmdata->reusesym = PETSC_TRUE;
2256: C->offloadmask = PETSC_OFFLOAD_GPU;
2257: }
2258: C->ops->productnumeric = MatProductNumeric_SeqAIJCUSPARSE_SeqAIJCUSPARSE;
2259: PetscFunctionReturn(PETSC_SUCCESS);
2260: }
2262: PETSC_INTERN PetscErrorCode MatProductSetFromOptions_SeqAIJ_SeqDense(Mat);
2264: /* handles sparse or dense B */
2265: static PetscErrorCode MatProductSetFromOptions_SeqAIJCUSPARSE(Mat mat)
2266: {
2267: Mat_Product *product = mat->product;
2268: PetscBool isdense = PETSC_FALSE, Biscusp = PETSC_FALSE, Ciscusp = PETSC_TRUE;
2270: PetscFunctionBegin;
2271: MatCheckProduct(mat, 1);
2272: PetscCall(PetscObjectBaseTypeCompare((PetscObject)product->B, MATSEQDENSE, &isdense));
2273: if (!product->A->boundtocpu && !product->B->boundtocpu) PetscCall(PetscObjectTypeCompare((PetscObject)product->B, MATSEQAIJCUSPARSE, &Biscusp));
2274: if (product->type == MATPRODUCT_ABC) {
2275: Ciscusp = PETSC_FALSE;
2276: if (!product->C->boundtocpu) PetscCall(PetscObjectTypeCompare((PetscObject)product->C, MATSEQAIJCUSPARSE, &Ciscusp));
2277: }
2278: if (Biscusp && Ciscusp) { /* we can always select the CPU backend */
2279: PetscBool usecpu = PETSC_FALSE;
2280: switch (product->type) {
2281: case MATPRODUCT_AB:
2282: if (product->api_user) {
2283: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatMatMult", "Mat");
2284: PetscCall(PetscOptionsBool("-matmatmult_backend_cpu", "Use CPU code", "MatMatMult", usecpu, &usecpu, NULL));
2285: PetscOptionsEnd();
2286: } else {
2287: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_AB", "Mat");
2288: PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatMatMult", usecpu, &usecpu, NULL));
2289: PetscOptionsEnd();
2290: }
2291: break;
2292: case MATPRODUCT_AtB:
2293: if (product->api_user) {
2294: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatTransposeMatMult", "Mat");
2295: PetscCall(PetscOptionsBool("-mattransposematmult_backend_cpu", "Use CPU code", "MatTransposeMatMult", usecpu, &usecpu, NULL));
2296: PetscOptionsEnd();
2297: } else {
2298: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_AtB", "Mat");
2299: PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatTransposeMatMult", usecpu, &usecpu, NULL));
2300: PetscOptionsEnd();
2301: }
2302: break;
2303: case MATPRODUCT_PtAP:
2304: if (product->api_user) {
2305: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatPtAP", "Mat");
2306: PetscCall(PetscOptionsBool("-matptap_backend_cpu", "Use CPU code", "MatPtAP", usecpu, &usecpu, NULL));
2307: PetscOptionsEnd();
2308: } else {
2309: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_PtAP", "Mat");
2310: PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatPtAP", usecpu, &usecpu, NULL));
2311: PetscOptionsEnd();
2312: }
2313: break;
2314: case MATPRODUCT_RARt:
2315: if (product->api_user) {
2316: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatRARt", "Mat");
2317: PetscCall(PetscOptionsBool("-matrart_backend_cpu", "Use CPU code", "MatRARt", usecpu, &usecpu, NULL));
2318: PetscOptionsEnd();
2319: } else {
2320: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_RARt", "Mat");
2321: PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatRARt", usecpu, &usecpu, NULL));
2322: PetscOptionsEnd();
2323: }
2324: break;
2325: case MATPRODUCT_ABC:
2326: if (product->api_user) {
2327: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatMatMatMult", "Mat");
2328: PetscCall(PetscOptionsBool("-matmatmatmult_backend_cpu", "Use CPU code", "MatMatMatMult", usecpu, &usecpu, NULL));
2329: PetscOptionsEnd();
2330: } else {
2331: PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_ABC", "Mat");
2332: PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatMatMatMult", usecpu, &usecpu, NULL));
2333: PetscOptionsEnd();
2334: }
2335: break;
2336: default:
2337: break;
2338: }
2339: if (usecpu) Biscusp = Ciscusp = PETSC_FALSE;
2340: }
2341: /* dispatch */
2342: if (isdense) {
2343: switch (product->type) {
2344: case MATPRODUCT_AB:
2345: case MATPRODUCT_AtB:
2346: case MATPRODUCT_ABt:
2347: case MATPRODUCT_PtAP:
2348: case MATPRODUCT_RARt:
2349: if (product->A->boundtocpu) {
2350: PetscCall(MatProductSetFromOptions_SeqAIJ_SeqDense(mat));
2351: } else {
2352: mat->ops->productsymbolic = MatProductSymbolic_SeqAIJCUSPARSE_SeqDENSECUDA;
2353: }
2354: break;
2355: case MATPRODUCT_ABC:
2356: mat->ops->productsymbolic = MatProductSymbolic_ABC_Basic;
2357: break;
2358: default:
2359: break;
2360: }
2361: } else if (Biscusp && Ciscusp) {
2362: switch (product->type) {
2363: case MATPRODUCT_AB:
2364: case MATPRODUCT_AtB:
2365: case MATPRODUCT_ABt:
2366: mat->ops->productsymbolic = MatProductSymbolic_SeqAIJCUSPARSE_SeqAIJCUSPARSE;
2367: break;
2368: case MATPRODUCT_PtAP:
2369: case MATPRODUCT_RARt:
2370: case MATPRODUCT_ABC:
2371: mat->ops->productsymbolic = MatProductSymbolic_ABC_Basic;
2372: break;
2373: default:
2374: break;
2375: }
2376: } else { /* fallback for AIJ */
2377: PetscCall(MatProductSetFromOptions_SeqAIJ(mat));
2378: }
2379: PetscFunctionReturn(PETSC_SUCCESS);
2380: }
2382: static PetscErrorCode MatMult_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy)
2383: {
2384: PetscFunctionBegin;
2385: PetscCall(MatMultAddKernel_SeqAIJCUSPARSE(A, xx, NULL, yy, PETSC_FALSE, PETSC_FALSE));
2386: PetscFunctionReturn(PETSC_SUCCESS);
2387: }
2389: static PetscErrorCode MatMultAdd_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy, Vec zz)
2390: {
2391: PetscFunctionBegin;
2392: PetscCall(MatMultAddKernel_SeqAIJCUSPARSE(A, xx, yy, zz, PETSC_FALSE, PETSC_FALSE));
2393: PetscFunctionReturn(PETSC_SUCCESS);
2394: }
2396: static PetscErrorCode MatMultHermitianTranspose_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy)
2397: {
2398: PetscFunctionBegin;
2399: PetscCall(MatMultAddKernel_SeqAIJCUSPARSE(A, xx, NULL, yy, PETSC_TRUE, PETSC_TRUE));
2400: PetscFunctionReturn(PETSC_SUCCESS);
2401: }
2403: static PetscErrorCode MatMultHermitianTransposeAdd_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy, Vec zz)
2404: {
2405: PetscFunctionBegin;
2406: PetscCall(MatMultAddKernel_SeqAIJCUSPARSE(A, xx, yy, zz, PETSC_TRUE, PETSC_TRUE));
2407: PetscFunctionReturn(PETSC_SUCCESS);
2408: }
2410: static PetscErrorCode MatMultTranspose_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy)
2411: {
2412: PetscFunctionBegin;
2413: PetscCall(MatMultAddKernel_SeqAIJCUSPARSE(A, xx, NULL, yy, PETSC_TRUE, PETSC_FALSE));
2414: PetscFunctionReturn(PETSC_SUCCESS);
2415: }
2417: __global__ static void ScatterAdd(PetscInt n, PetscInt *idx, const PetscScalar *x, PetscScalar *y)
2418: {
2419: int i = blockIdx.x * blockDim.x + threadIdx.x;
2420: if (i < n) y[idx[i]] += x[i];
2421: }
2423: /* z = op(A) x + y. If trans & !herm, op = ^T; if trans & herm, op = ^H; if !trans, op = no-op */
2424: static PetscErrorCode MatMultAddKernel_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy, Vec zz, PetscBool trans, PetscBool herm)
2425: {
2426: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;
2427: Mat_SeqAIJCUSPARSE *cusparsestruct = (Mat_SeqAIJCUSPARSE *)A->spptr;
2428: Mat_SeqAIJCUSPARSEMultStruct *matstruct;
2429: PetscScalar *xarray, *zarray, *dptr, *beta, *xptr;
2430: cusparseOperation_t opA = CUSPARSE_OPERATION_NON_TRANSPOSE;
2431: PetscBool compressed;
2432: PetscInt nx, ny;
2434: PetscFunctionBegin;
2435: PetscCheck(!herm || trans, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Hermitian and not transpose not supported");
2436: if (!a->nz) {
2437: if (yy) PetscCall(VecSeq_CUDA::Copy(yy, zz));
2438: else PetscCall(VecSeq_CUDA::Set(zz, 0));
2439: PetscFunctionReturn(PETSC_SUCCESS);
2440: }
2441: /* The line below is necessary due to the operations that modify the matrix on the CPU (axpy, scale, etc) */
2442: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
2443: if (!trans) {
2444: matstruct = (Mat_SeqAIJCUSPARSEMultStruct *)cusparsestruct->mat;
2445: PetscCheck(matstruct, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "SeqAIJCUSPARSE does not have a 'mat' (need to fix)");
2446: } else {
2447: if (herm || !A->form_explicit_transpose) {
2448: opA = herm ? CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE : CUSPARSE_OPERATION_TRANSPOSE;
2449: matstruct = (Mat_SeqAIJCUSPARSEMultStruct *)cusparsestruct->mat;
2450: } else {
2451: if (!cusparsestruct->matTranspose) PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(A));
2452: matstruct = (Mat_SeqAIJCUSPARSEMultStruct *)cusparsestruct->matTranspose;
2453: }
2454: }
2455: /* Does the matrix use compressed rows (i.e., drop zero rows)? */
2456: compressed = matstruct->cprowIndices ? PETSC_TRUE : PETSC_FALSE;
2458: try {
2459: PetscCall(VecCUDAGetArrayRead(xx, (const PetscScalar **)&xarray));
2460: if (yy == zz) PetscCall(VecCUDAGetArray(zz, &zarray)); /* read & write zz, so need to get up-to-date zarray on GPU */
2461: else PetscCall(VecCUDAGetArrayWrite(zz, &zarray)); /* write zz, so no need to init zarray on GPU */
2463: PetscCall(PetscLogGpuTimeBegin());
2464: if (opA == CUSPARSE_OPERATION_NON_TRANSPOSE) {
2465: /* z = A x + beta y.
2466: If A is compressed (with less rows), then Ax is shorter than the full z, so we need a work vector to store Ax.
2467: When A is non-compressed, and z = y, we can set beta=1 to compute y = Ax + y in one call.
2468: */
2469: xptr = xarray;
2470: dptr = compressed ? cusparsestruct->workVector->data().get() : zarray;
2471: beta = (yy == zz && !compressed) ? matstruct->beta_one : matstruct->beta_zero;
2472: /* Get length of x, y for y=Ax. ny might be shorter than the work vector's allocated length, since the work vector is
2473: allocated to accommodate different uses. So we get the length info directly from mat.
2474: */
2475: if (cusparsestruct->format == MAT_CUSPARSE_CSR) {
2476: CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
2477: nx = mat->num_cols; // since y = Ax
2478: ny = mat->num_rows;
2479: }
2480: } else {
2481: /* z = A^T x + beta y
2482: If A is compressed, then we need a work vector as the shorter version of x to compute A^T x.
2483: Note A^Tx is of full length, so we set beta to 1.0 if y exists.
2484: */
2485: xptr = compressed ? cusparsestruct->workVector->data().get() : xarray;
2486: dptr = zarray;
2487: beta = yy ? matstruct->beta_one : matstruct->beta_zero;
2488: if (compressed) { /* Scatter x to work vector */
2489: thrust::device_ptr<PetscScalar> xarr = thrust::device_pointer_cast(xarray);
2491: thrust::for_each(
2492: #if PetscDefined(HAVE_THRUST_ASYNC)
2493: thrust::cuda::par.on(PetscDefaultCudaStream),
2494: #endif
2495: thrust::make_zip_iterator(thrust::make_tuple(cusparsestruct->workVector->begin(), thrust::make_permutation_iterator(xarr, matstruct->cprowIndices->begin()))),
2496: thrust::make_zip_iterator(thrust::make_tuple(cusparsestruct->workVector->begin(), thrust::make_permutation_iterator(xarr, matstruct->cprowIndices->begin()))) + matstruct->cprowIndices->size(), VecCUDAEqualsReverse());
2497: }
2498: if (cusparsestruct->format == MAT_CUSPARSE_CSR) {
2499: CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
2500: nx = mat->num_rows; // since y = A^T x
2501: ny = mat->num_cols;
2502: }
2503: }
2505: /* csr_spmv does y = alpha op(A) x + beta y */
2506: if (cusparsestruct->format == MAT_CUSPARSE_CSR) {
2507: 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");
2509: if (!matstruct->cuSpMV[opA].initialized) { /* built on demand */
2510: CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
2511: 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));
2512: PetscCallCUSPARSE(cusparseCreateDnVec(&matstruct->cuSpMV[opA].vecXDescr, nx, xptr, cusparse_scalartype));
2513: PetscCallCUSPARSE(cusparseCreateDnVec(&matstruct->cuSpMV[opA].vecYDescr, ny, dptr, cusparse_scalartype));
2514: 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,
2515: &matstruct->cuSpMV[opA].spmvBufferSize));
2516: PetscCallCUDA(cudaMalloc(&matstruct->cuSpMV[opA].spmvBuffer, matstruct->cuSpMV[opA].spmvBufferSize));
2517: #if PETSC_PKG_CUDA_VERSION_GE(12, 4, 0) // cusparseSpMV_preprocess is added in 12.4
2518: 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,
2519: matstruct->cuSpMV[opA].spmvBuffer));
2520: #endif
2521: matstruct->cuSpMV[opA].initialized = PETSC_TRUE;
2522: } else {
2523: /* x, y's value pointers might change between calls, but their shape is kept, so we just update pointers */
2524: PetscCallCUSPARSE(cusparseDnVecSetValues(matstruct->cuSpMV[opA].vecXDescr, xptr));
2525: PetscCallCUSPARSE(cusparseDnVecSetValues(matstruct->cuSpMV[opA].vecYDescr, dptr));
2526: }
2528: PetscCallCUSPARSE(
2529: 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));
2531: } else {
2532: if (cusparsestruct->nrows) {
2533: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MAT_CUSPARSE_ELL and MAT_CUSPARSE_HYB are not supported since CUDA-11.0");
2534: }
2535: }
2536: PetscCall(PetscLogGpuTimeEnd());
2538: if (opA == CUSPARSE_OPERATION_NON_TRANSPOSE) {
2539: if (yy) { /* MatMultAdd: zz = A*xx + yy */
2540: if (compressed) { /* A is compressed. We first copy yy to zz, then ScatterAdd the work vector to zz */
2541: PetscCall(VecSeq_CUDA::Copy(yy, zz)); /* zz = yy */
2542: } else if (zz != yy) { /* A is not compressed. zz already contains A*xx, and we just need to add yy */
2543: PetscCall(VecSeq_CUDA::AXPY(zz, 1.0, yy)); /* zz += yy */
2544: }
2545: } else if (compressed) { /* MatMult: zz = A*xx. A is compressed, so we zero zz first, then ScatterAdd the work vector to zz */
2546: PetscCall(VecSeq_CUDA::Set(zz, 0));
2547: }
2549: /* ScatterAdd the result from work vector into the full vector when A is compressed */
2550: if (compressed) {
2551: PetscCall(PetscLogGpuTimeBegin());
2552: PetscInt n = (PetscInt)matstruct->cprowIndices->size();
2553: ScatterAdd<<<(int)((n + 255) / 256), 256, 0, PetscDefaultCudaStream>>>(n, matstruct->cprowIndices->data().get(), cusparsestruct->workVector->data().get(), zarray);
2554: PetscCall(PetscLogGpuTimeEnd());
2555: }
2556: } else {
2557: if (yy && yy != zz) PetscCall(VecSeq_CUDA::AXPY(zz, 1.0, yy)); /* zz += yy */
2558: }
2559: PetscCall(VecCUDARestoreArrayRead(xx, (const PetscScalar **)&xarray));
2560: if (yy == zz) PetscCall(VecCUDARestoreArray(zz, &zarray));
2561: else PetscCall(VecCUDARestoreArrayWrite(zz, &zarray));
2562: } catch (char *ex) {
2563: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_LIB, "CUSPARSE error: %s", ex);
2564: }
2565: if (yy) PetscCall(PetscLogGpuFlops(2.0 * a->nz));
2566: else PetscCall(PetscLogGpuFlops(2.0 * a->nz - a->nonzerorowcnt));
2567: PetscFunctionReturn(PETSC_SUCCESS);
2568: }
2570: static PetscErrorCode MatMultTransposeAdd_SeqAIJCUSPARSE(Mat A, Vec xx, Vec yy, Vec zz)
2571: {
2572: PetscFunctionBegin;
2573: PetscCall(MatMultAddKernel_SeqAIJCUSPARSE(A, xx, yy, zz, PETSC_TRUE, PETSC_FALSE));
2574: PetscFunctionReturn(PETSC_SUCCESS);
2575: }
2577: static PetscErrorCode MatGetDiagonal_SeqAIJCUSPARSE(Mat A, Vec diag)
2578: {
2579: PetscFunctionBegin;
2580: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::GetDiagonal(A, diag));
2581: PetscFunctionReturn(PETSC_SUCCESS);
2582: }
2584: static PetscErrorCode MatDiagonalScale_SeqAIJCUSPARSE(Mat A, Vec ll, Vec rr)
2585: {
2586: PetscFunctionBegin;
2587: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::DiagonalScale(A, ll, rr));
2588: PetscFunctionReturn(PETSC_SUCCESS);
2589: }
2591: static PetscErrorCode MatAssemblyEnd_SeqAIJCUSPARSE(Mat A, MatAssemblyType mode)
2592: {
2593: PetscFunctionBegin;
2594: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::AssemblyEnd(A, mode));
2595: PetscFunctionReturn(PETSC_SUCCESS);
2596: }
2598: /*@
2599: MatCreateSeqAIJCUSPARSE - Creates a sparse matrix in `MATAIJCUSPARSE` (compressed row) format for use on NVIDIA GPUs
2601: Collective
2603: Input Parameters:
2604: + comm - MPI communicator, set to `PETSC_COMM_SELF`
2605: . m - number of rows
2606: . n - number of columns
2607: . nz - number of nonzeros per row (same for all rows), ignored if `nnz` is provide
2608: - nnz - array containing the number of nonzeros in the various rows (possibly different for each row) or `NULL`
2610: Output Parameter:
2611: . A - the matrix
2613: Level: intermediate
2615: Notes:
2616: This matrix will ultimately pushed down to NVIDIA GPUs and use the CuSPARSE library for
2617: calculations. For good matrix assembly performance the user should preallocate the matrix
2618: storage by setting the parameter `nz` (or the array `nnz`).
2620: It is recommended that one use the `MatCreate()`, `MatSetType()` and/or `MatSetFromOptions()`,
2621: MatXXXXSetPreallocation() paradgm instead of this routine directly.
2622: [MatXXXXSetPreallocation() is, for example, `MatSeqAIJSetPreallocation()`]
2624: The AIJ format, also called
2625: compressed row storage, is fully compatible with standard Fortran
2626: storage. That is, the stored row and column indices can begin at
2627: either one (as in Fortran) or zero.
2629: Specify the preallocated storage with either nz or nnz (not both).
2630: Set `nz` = `PETSC_DEFAULT` and `nnz` = `NULL` for PETSc to control dynamic memory
2631: allocation.
2633: When working with matrices for GPUs, it is often better to use the `MatSetPreallocationCOO()` and `MatSetValuesCOO()` paradigm rather than using this routine and `MatSetValues()`
2635: .seealso: [](ch_matrices), `Mat`, `MATSEQAIJCUSPARSE`, `MatCreate()`, `MatCreateAIJ()`, `MatSetValues()`, `MatSeqAIJSetColumnIndices()`, `MatCreateSeqAIJWithArrays()`, `MATAIJCUSPARSE`,
2636: `MatSetPreallocationCOO()`, `MatSetValuesCOO()`
2637: @*/
2638: PetscErrorCode MatCreateSeqAIJCUSPARSE(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt nz, const PetscInt nnz[], Mat *A)
2639: {
2640: return MatSeqAIJCUSPARSE_CUPM_t::CreateSeqAIJ(comm, m, n, nz, nnz, A);
2641: }
2643: static PetscErrorCode MatDestroy_SeqAIJCUSPARSE(Mat A)
2644: {
2645: return MatSeqAIJCUSPARSE_CUPM_t::Destroy(A);
2646: }
2648: static PetscErrorCode MatDuplicate_SeqAIJCUSPARSE(Mat A, MatDuplicateOption cpvalues, Mat *B)
2649: {
2650: PetscFunctionBegin;
2651: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::Duplicate(A, cpvalues, B));
2652: PetscFunctionReturn(PETSC_SUCCESS);
2653: }
2655: static PetscErrorCode MatAXPY_SeqAIJCUSPARSE(Mat Y, PetscScalar a, Mat X, MatStructure str)
2656: {
2657: Mat_SeqAIJ *x = (Mat_SeqAIJ *)X->data, *y = (Mat_SeqAIJ *)Y->data;
2658: Mat_SeqAIJCUSPARSE *cy;
2659: Mat_SeqAIJCUSPARSE *cx;
2660: CsrMatrix *csry, *csrx;
2662: PetscFunctionBegin;
2663: cy = (Mat_SeqAIJCUSPARSE *)Y->spptr;
2664: cx = (Mat_SeqAIJCUSPARSE *)X->spptr;
2665: if (X->ops->axpy != Y->ops->axpy) {
2666: PetscCall(MatSeqAIJCUSPARSEInvalidateTranspose(Y, PETSC_FALSE));
2667: PetscCall(MatAXPY_SeqAIJ(Y, a, X, str));
2668: PetscFunctionReturn(PETSC_SUCCESS);
2669: }
2670: /* if we are here, it means both matrices are bound to GPU */
2671: PetscCall(MatSeqAIJCUSPARSECopyToGPU(Y));
2672: PetscCall(MatSeqAIJCUSPARSECopyToGPU(X));
2673: PetscCheck(cy->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)Y), PETSC_ERR_GPU, "only MAT_CUSPARSE_CSR supported");
2674: PetscCheck(cx->format == MAT_CUSPARSE_CSR, PetscObjectComm((PetscObject)X), PETSC_ERR_GPU, "only MAT_CUSPARSE_CSR supported");
2675: csry = (CsrMatrix *)cy->mat->mat;
2676: csrx = (CsrMatrix *)cx->mat->mat;
2677: /* see if we can turn this into a cublas axpy */
2678: if (str != SAME_NONZERO_PATTERN && x->nz == y->nz && !x->compressedrow.use && !y->compressedrow.use) {
2679: bool eq = thrust::equal(thrust::device, csry->row_offsets->begin(), csry->row_offsets->end(), csrx->row_offsets->begin());
2680: if (eq) eq = thrust::equal(thrust::device, csry->column_indices->begin(), csry->column_indices->end(), csrx->column_indices->begin());
2681: if (eq) str = SAME_NONZERO_PATTERN;
2682: }
2683: /* spgeam is buggy with one column */
2684: if (Y->cmap->n == 1 && str != SAME_NONZERO_PATTERN) str = DIFFERENT_NONZERO_PATTERN;
2686: #if !PetscDefined(USE_64BIT_INDICES) // cusparseScsrgeam2 etc. do not support 64bit indices
2687: if (str == SUBSET_NONZERO_PATTERN) {
2688: PetscScalar *ay, b = 1.0;
2689: const PetscScalar *ax;
2690: size_t bufferSize;
2691: void *buffer;
2693: PetscCall(MatSeqAIJCUSPARSEGetArrayRead(X, &ax));
2694: PetscCall(MatSeqAIJCUSPARSEGetArray(Y, &ay));
2695: PetscCallCUSPARSE(cusparseSetPointerMode(cy->handle, CUSPARSE_POINTER_MODE_HOST));
2696: 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(),
2697: csry->column_indices->data().get(), cy->mat->descr, ay, csry->row_offsets->data().get(), csry->column_indices->data().get(), &bufferSize));
2698: PetscCallCUDA(cudaMalloc(&buffer, bufferSize));
2699: PetscCall(PetscLogGpuTimeBegin());
2700: 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(),
2701: csry->column_indices->data().get(), cy->mat->descr, ay, csry->row_offsets->data().get(), csry->column_indices->data().get(), buffer));
2702: PetscCall(PetscLogGpuFlops(x->nz + y->nz));
2703: PetscCall(PetscLogGpuTimeEnd());
2704: PetscCallCUDA(cudaFree(buffer));
2706: PetscCallCUSPARSE(cusparseSetPointerMode(cy->handle, CUSPARSE_POINTER_MODE_DEVICE));
2707: PetscCall(MatSeqAIJCUSPARSERestoreArrayRead(X, &ax));
2708: PetscCall(MatSeqAIJCUSPARSERestoreArray(Y, &ay));
2709: } else
2710: #endif
2711: if (str == SAME_NONZERO_PATTERN) {
2712: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::AXPY_SameNZ(Y, a, X));
2713: } else {
2714: PetscCall(MatSeqAIJCUSPARSEInvalidateTranspose(Y, PETSC_FALSE));
2715: PetscCall(MatAXPY_SeqAIJ(Y, a, X, str));
2716: }
2717: PetscFunctionReturn(PETSC_SUCCESS);
2718: }
2720: static PetscErrorCode MatScale_SeqAIJCUSPARSE(Mat Y, PetscScalar a)
2721: {
2722: PetscFunctionBegin;
2723: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::Scale(Y, a));
2724: PetscFunctionReturn(PETSC_SUCCESS);
2725: }
2727: static PetscErrorCode MatZeroEntries_SeqAIJCUSPARSE(Mat A)
2728: {
2729: PetscFunctionBegin;
2730: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::ZeroEntries(A));
2731: PetscFunctionReturn(PETSC_SUCCESS);
2732: }
2734: static PetscErrorCode MatGetCurrentMemType_SeqAIJCUSPARSE(Mat A, PetscMemType *m)
2735: {
2736: PetscFunctionBegin;
2737: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::GetCurrentMemType(A, m));
2738: PetscFunctionReturn(PETSC_SUCCESS);
2739: }
2741: static PetscErrorCode MatBindToCPU_SeqAIJCUSPARSE(Mat A, PetscBool flg)
2742: {
2743: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;
2745: PetscFunctionBegin;
2746: if (A->factortype != MAT_FACTOR_NONE) {
2747: A->boundtocpu = flg;
2748: PetscFunctionReturn(PETSC_SUCCESS);
2749: }
2750: if (flg) {
2751: PetscCall(MatSeqAIJCUSPARSECopyFromGPU(A));
2753: A->ops->scale = MatScale_SeqAIJ;
2754: A->ops->getdiagonal = MatGetDiagonal_SeqAIJ;
2755: A->ops->diagonalscale = MatDiagonalScale_SeqAIJ;
2756: A->ops->axpy = MatAXPY_SeqAIJ;
2757: A->ops->zeroentries = MatZeroEntries_SeqAIJ;
2758: A->ops->mult = MatMult_SeqAIJ;
2759: A->ops->multadd = MatMultAdd_SeqAIJ;
2760: A->ops->multtranspose = MatMultTranspose_SeqAIJ;
2761: A->ops->multtransposeadd = MatMultTransposeAdd_SeqAIJ;
2762: A->ops->multhermitiantranspose = NULL;
2763: A->ops->multhermitiantransposeadd = NULL;
2764: A->ops->productsetfromoptions = MatProductSetFromOptions_SeqAIJ;
2765: A->ops->getcurrentmemtype = NULL;
2766: PetscCall(PetscMemzero(a->ops, sizeof(Mat_SeqAIJOps)));
2767: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSeqAIJCopySubArray_C", NULL));
2768: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijcusparse_seqdensecuda_C", NULL));
2769: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijcusparse_seqdense_C", NULL));
2770: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetPreallocationCOO_C", NULL));
2771: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetValuesCOO_C", NULL));
2772: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijcusparse_seqaijcusparse_C", NULL));
2773: } else {
2774: A->ops->scale = MatScale_SeqAIJCUSPARSE;
2775: A->ops->getdiagonal = MatGetDiagonal_SeqAIJCUSPARSE;
2776: A->ops->diagonalscale = MatDiagonalScale_SeqAIJCUSPARSE;
2777: A->ops->axpy = MatAXPY_SeqAIJCUSPARSE;
2778: A->ops->zeroentries = MatZeroEntries_SeqAIJCUSPARSE;
2779: A->ops->mult = MatMult_SeqAIJCUSPARSE;
2780: A->ops->multadd = MatMultAdd_SeqAIJCUSPARSE;
2781: A->ops->multtranspose = MatMultTranspose_SeqAIJCUSPARSE;
2782: A->ops->multtransposeadd = MatMultTransposeAdd_SeqAIJCUSPARSE;
2783: A->ops->multhermitiantranspose = MatMultHermitianTranspose_SeqAIJCUSPARSE;
2784: A->ops->multhermitiantransposeadd = MatMultHermitianTransposeAdd_SeqAIJCUSPARSE;
2785: A->ops->productsetfromoptions = MatProductSetFromOptions_SeqAIJCUSPARSE;
2786: A->ops->getcurrentmemtype = MatGetCurrentMemType_SeqAIJCUSPARSE;
2787: a->ops->getarray = MatSeqAIJGetArray_SeqAIJCUSPARSE;
2788: a->ops->restorearray = MatSeqAIJRestoreArray_SeqAIJCUSPARSE;
2789: a->ops->getarrayread = MatSeqAIJGetArrayRead_SeqAIJCUSPARSE;
2790: a->ops->restorearrayread = MatSeqAIJRestoreArrayRead_SeqAIJCUSPARSE;
2791: a->ops->getarraywrite = MatSeqAIJGetArrayWrite_SeqAIJCUSPARSE;
2792: a->ops->restorearraywrite = MatSeqAIJRestoreArrayWrite_SeqAIJCUSPARSE;
2793: a->ops->getcsrandmemtype = MatSeqAIJGetCSRAndMemType_SeqAIJCUSPARSE;
2795: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSeqAIJCopySubArray_C", MatSeqAIJCopySubArray_SeqAIJCUSPARSE));
2796: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijcusparse_seqdensecuda_C", MatProductSetFromOptions_SeqAIJCUSPARSE));
2797: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijcusparse_seqdense_C", MatProductSetFromOptions_SeqAIJCUSPARSE));
2798: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetPreallocationCOO_C", MatSetPreallocationCOO_SeqAIJCUSPARSE));
2799: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetValuesCOO_C", MatSetValuesCOO_SeqAIJCUSPARSE));
2800: PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijcusparse_seqaijcusparse_C", MatProductSetFromOptions_SeqAIJCUSPARSE));
2801: }
2802: A->boundtocpu = flg;
2803: a->inode.use = (flg && a->inode.size_csr) ? PETSC_TRUE : PETSC_FALSE;
2804: PetscFunctionReturn(PETSC_SUCCESS);
2805: }
2807: PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJCUSPARSE(Mat A, MatType, MatReuse reuse, Mat *newmat)
2808: {
2809: Mat B;
2811: PetscFunctionBegin;
2812: PetscCall(PetscDeviceInitialize(PETSC_DEVICE_CUDA)); /* first use of CUSPARSE may be via MatConvert */
2813: if (reuse == MAT_INITIAL_MATRIX) {
2814: PetscCall(MatDuplicate(A, MAT_COPY_VALUES, newmat));
2815: } else if (reuse == MAT_REUSE_MATRIX) {
2816: PetscCall(MatCopy(A, *newmat, SAME_NONZERO_PATTERN));
2817: }
2818: B = *newmat;
2820: PetscCall(PetscFree(B->defaultvectype));
2821: PetscCall(PetscStrallocpy(VECCUDA, &B->defaultvectype));
2823: if (reuse != MAT_REUSE_MATRIX && !B->spptr) {
2824: if (B->factortype == MAT_FACTOR_NONE) {
2825: Mat_SeqAIJCUSPARSE *spptr;
2826: PetscCall(PetscNew(&spptr));
2827: PetscCallCUSPARSE(cusparseCreate(&spptr->handle));
2828: PetscCallCUSPARSE(cusparseSetStream(spptr->handle, PetscDefaultCudaStream));
2829: spptr->format = MAT_CUSPARSE_CSR;
2830: spptr->spmvAlg = CUSPARSE_SPMV_CSR_ALG1; /* default, since we only support csr */
2831: spptr->spmmAlg = CUSPARSE_SPMM_CSR_ALG1; /* default, only support column-major dense matrix B */
2832: spptr->csr2cscAlg = CUSPARSE_CSR2CSC_ALG1;
2833: B->spptr = spptr;
2834: } else {
2835: Mat_SeqAIJCUSPARSETriFactors *spptr;
2837: PetscCall(PetscNew(&spptr));
2838: PetscCallCUSPARSE(cusparseCreate(&spptr->handle));
2839: PetscCallCUSPARSE(cusparseSetStream(spptr->handle, PetscDefaultCudaStream));
2840: B->spptr = spptr;
2841: }
2842: B->offloadmask = PETSC_OFFLOAD_UNALLOCATED;
2843: }
2844: B->ops->assemblyend = MatAssemblyEnd_SeqAIJCUSPARSE;
2845: B->ops->destroy = MatDestroy_SeqAIJCUSPARSE;
2846: B->ops->setoption = MatSetOption_SeqAIJCUSPARSE;
2847: B->ops->setfromoptions = MatSetFromOptions_SeqAIJCUSPARSE;
2848: B->ops->bindtocpu = MatBindToCPU_SeqAIJCUSPARSE;
2849: B->ops->duplicate = MatDuplicate_SeqAIJCUSPARSE;
2850: B->ops->getcurrentmemtype = MatGetCurrentMemType_SeqAIJCUSPARSE;
2852: PetscCall(MatBindToCPU_SeqAIJCUSPARSE(B, PETSC_FALSE));
2853: PetscCall(PetscObjectChangeTypeName((PetscObject)B, MATSEQAIJCUSPARSE));
2854: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatCUSPARSESetFormat_C", MatCUSPARSESetFormat_SeqAIJCUSPARSE));
2855: #if PetscDefined(HAVE_HYPRE)
2856: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_seqaijcusparse_hypre_C", MatConvert_AIJ_HYPRE));
2857: #endif
2858: PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatCUSPARSESetUseCPUSolve_C", MatCUSPARSESetUseCPUSolve_SeqAIJCUSPARSE));
2859: PetscFunctionReturn(PETSC_SUCCESS);
2860: }
2862: PETSC_EXTERN PetscErrorCode MatCreate_SeqAIJCUSPARSE(Mat B)
2863: {
2864: PetscFunctionBegin;
2865: PetscCall(MatCreate_SeqAIJ(B));
2866: PetscCall(MatConvert_SeqAIJ_SeqAIJCUSPARSE(B, MATSEQAIJCUSPARSE, MAT_INPLACE_MATRIX, &B));
2867: PetscFunctionReturn(PETSC_SUCCESS);
2868: }
2870: /*MC
2871: MATSEQAIJCUSPARSE - MATAIJCUSPARSE = "(seq)aijcusparse" - A matrix type to be used for sparse matrices on NVIDIA GPUs.
2873: Options Database Keys:
2874: + -mat_type aijcusparse - Sets the matrix type to "seqaijcusparse" during a call to `MatSetFromOptions()`
2875: . -mat_cusparse_storage_format csr - Sets the storage format of matrices (for `MatMult()` and factors in `MatSolve()`).
2876: Other options include ell (ellpack) or hyb (hybrid).
2877: . -mat_cusparse_mult_storage_format csr - Sets the storage format of matrices (for `MatMult()`). Other options include ell (ellpack) or hyb (hybrid).
2878: - -mat_cusparse_use_cpu_solve - Performs the `MatSolve()` on the CPU
2880: Level: beginner
2882: Notes:
2883: These matrices can be in either CSR, ELL, or HYB format.
2885: All matrix calculations are performed on NVIDIA GPUs using the cuSPARSE library.
2887: 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
2888: if some integer values passed in do not fit in `int`.
2890: .seealso: [](ch_matrices), `Mat`, `MatCreateSeqAIJCUSPARSE()`, `MatCUSPARSESetUseCPUSolve()`, `MATAIJCUSPARSE`, `MatCreateAIJCUSPARSE()`, `MatCUSPARSESetFormat()`, `MatCUSPARSEStorageFormat`, `MatCUSPARSEFormatOperation`
2891: M*/
2893: PETSC_INTERN PetscErrorCode MatSolverTypeRegister_CUSPARSE(void)
2894: {
2895: PetscFunctionBegin;
2896: PetscCall(MatSolverTypeRegister(MATSOLVERCUSPARSE, MATSEQAIJCUSPARSE, MAT_FACTOR_LU, MatGetFactor_seqaijcusparse_cusparse));
2897: PetscCall(MatSolverTypeRegister(MATSOLVERCUSPARSE, MATSEQAIJCUSPARSE, MAT_FACTOR_CHOLESKY, MatGetFactor_seqaijcusparse_cusparse));
2898: PetscCall(MatSolverTypeRegister(MATSOLVERCUSPARSE, MATSEQAIJCUSPARSE, MAT_FACTOR_ILU, MatGetFactor_seqaijcusparse_cusparse));
2899: PetscCall(MatSolverTypeRegister(MATSOLVERCUSPARSE, MATSEQAIJCUSPARSE, MAT_FACTOR_ICC, MatGetFactor_seqaijcusparse_cusparse));
2900: PetscFunctionReturn(PETSC_SUCCESS);
2901: }
2903: static PetscErrorCode MatSeqAIJCUSPARSE_Destroy(Mat mat)
2904: {
2905: Mat_SeqAIJCUSPARSE *cusp = static_cast<Mat_SeqAIJCUSPARSE *>(mat->spptr);
2907: PetscFunctionBegin;
2908: if (cusp) {
2909: PetscCall(MatSeqAIJCUSPARSEMultStruct_Destroy(&cusp->mat, cusp->format));
2910: PetscCall(MatSeqAIJCUSPARSEMultStruct_Destroy(&cusp->matTranspose, cusp->format));
2911: delete cusp->workVector;
2912: delete cusp->rowoffsets_gpu;
2913: delete cusp->csr2csc_i;
2914: delete cusp->coords;
2915: if (cusp->handle) PetscCallCUSPARSE(cusparseDestroy(cusp->handle));
2916: PetscCall(PetscFree(mat->spptr));
2917: }
2918: PetscFunctionReturn(PETSC_SUCCESS);
2919: }
2921: static PetscErrorCode CsrMatrix_Destroy(CsrMatrix **mat)
2922: {
2923: PetscFunctionBegin;
2924: if (*mat) {
2925: delete (*mat)->values;
2926: delete (*mat)->column_indices;
2927: delete (*mat)->row_offsets;
2928: delete *mat;
2929: *mat = 0;
2930: }
2931: PetscFunctionReturn(PETSC_SUCCESS);
2932: }
2934: static PetscErrorCode MatSeqAIJCUSPARSEMultStruct_Destroy(Mat_SeqAIJCUSPARSEMultStruct **matstruct, MatCUSPARSEStorageFormat format)
2935: {
2936: CsrMatrix *mat;
2938: PetscFunctionBegin;
2939: if (*matstruct) {
2940: if ((*matstruct)->mat) {
2941: if (format == MAT_CUSPARSE_ELL || format == MAT_CUSPARSE_HYB) {
2942: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MAT_CUSPARSE_ELL and MAT_CUSPARSE_HYB are not supported since CUDA-11.0");
2943: } else {
2944: mat = (CsrMatrix *)(*matstruct)->mat;
2945: PetscCall(CsrMatrix_Destroy(&mat));
2946: }
2947: }
2948: if ((*matstruct)->descr) PetscCallCUSPARSE(cusparseDestroyMatDescr((*matstruct)->descr));
2949: delete (*matstruct)->cprowIndices;
2950: PetscCallCUDA(cudaFree((*matstruct)->alpha_one));
2951: PetscCallCUDA(cudaFree((*matstruct)->beta_zero));
2952: PetscCallCUDA(cudaFree((*matstruct)->beta_one));
2954: Mat_SeqAIJCUSPARSEMultStruct *mdata = *matstruct;
2955: if (mdata->matDescr) PetscCallCUSPARSE(cusparseDestroySpMat(mdata->matDescr));
2957: for (int i = 0; i < 3; i++) {
2958: if (mdata->cuSpMV[i].initialized) {
2959: PetscCallCUDA(cudaFree(mdata->cuSpMV[i].spmvBuffer));
2960: PetscCallCUSPARSE(cusparseDestroyDnVec(mdata->cuSpMV[i].vecXDescr));
2961: PetscCallCUSPARSE(cusparseDestroyDnVec(mdata->cuSpMV[i].vecYDescr));
2962: PetscCallCUSPARSE(cusparseDestroySpMat(mdata->cuSpMV[i].matDescr));
2963: #if PETSC_PKG_CUDA_VERSION_GE(12, 4, 0)
2964: if (mdata->matDescr_SpMM[i]) PetscCallCUSPARSE(cusparseDestroySpMat(mdata->matDescr_SpMM[i]));
2965: #endif
2966: }
2967: }
2968: delete *matstruct;
2969: *matstruct = NULL;
2970: }
2971: PetscFunctionReturn(PETSC_SUCCESS);
2972: }
2974: PetscErrorCode MatSeqAIJCUSPARSETriFactors_Reset(Mat_SeqAIJCUSPARSETriFactors_p *trifactors)
2975: {
2976: Mat_SeqAIJCUSPARSETriFactors *fs = *trifactors;
2978: PetscFunctionBegin;
2979: if (fs) {
2980: delete fs->rpermIndices;
2981: delete fs->cpermIndices;
2982: fs->rpermIndices = NULL;
2983: fs->cpermIndices = NULL;
2984: fs->init_dev_prop = PETSC_FALSE;
2985: PetscCallCUDA(cudaFree(fs->csrRowPtr));
2986: PetscCallCUDA(cudaFree(fs->csrColIdx));
2987: PetscCallCUDA(cudaFree(fs->csrRowPtr32));
2988: PetscCallCUDA(cudaFree(fs->csrColIdx32));
2989: PetscCallCUDA(cudaFree(fs->csrVal));
2990: PetscCallCUDA(cudaFree(fs->diag));
2991: PetscCallCUDA(cudaFree(fs->X));
2992: PetscCallCUDA(cudaFree(fs->Y));
2993: // PetscCallCUDA(cudaFree(fs->factBuffer_M)); /* No needed since factBuffer_M shares with one of spsvBuffer_L/U */
2994: PetscCallCUDA(cudaFree(fs->spsvBuffer_L));
2995: PetscCallCUDA(cudaFree(fs->spsvBuffer_U));
2996: PetscCallCUDA(cudaFree(fs->spsvBuffer_Lt));
2997: PetscCallCUDA(cudaFree(fs->spsvBuffer_Ut));
2998: PetscCallCUSPARSE(cusparseDestroyMatDescr(fs->matDescr_M));
2999: if (fs->spMatDescr_L) PetscCallCUSPARSE(cusparseDestroySpMat(fs->spMatDescr_L));
3000: if (fs->spMatDescr_U) PetscCallCUSPARSE(cusparseDestroySpMat(fs->spMatDescr_U));
3001: PetscCallCUSPARSE(cusparseSpSV_destroyDescr(fs->spsvDescr_L));
3002: PetscCallCUSPARSE(cusparseSpSV_destroyDescr(fs->spsvDescr_Lt));
3003: PetscCallCUSPARSE(cusparseSpSV_destroyDescr(fs->spsvDescr_U));
3004: PetscCallCUSPARSE(cusparseSpSV_destroyDescr(fs->spsvDescr_Ut));
3005: if (fs->dnVecDescr_X) PetscCallCUSPARSE(cusparseDestroyDnVec(fs->dnVecDescr_X));
3006: if (fs->dnVecDescr_Y) PetscCallCUSPARSE(cusparseDestroyDnVec(fs->dnVecDescr_Y));
3007: PetscCallCUSPARSE(cusparseDestroyCsrilu02Info(fs->ilu0Info_M));
3008: PetscCallCUSPARSE(cusparseDestroyCsric02Info(fs->ic0Info_M));
3009: PetscCall(PetscFree(fs->csrRowPtr_h));
3010: PetscCall(PetscFree(fs->csrVal_h));
3011: PetscCall(PetscFree(fs->diag_h));
3012: fs->createdTransposeSpSVDescr = PETSC_FALSE;
3013: fs->updatedTransposeSpSVAnalysis = PETSC_FALSE;
3014: }
3015: PetscFunctionReturn(PETSC_SUCCESS);
3016: }
3018: static PetscErrorCode MatSeqAIJCUSPARSETriFactors_Destroy(Mat_SeqAIJCUSPARSETriFactors **trifactors)
3019: {
3020: PetscFunctionBegin;
3021: if (*trifactors) {
3022: PetscCall(MatSeqAIJCUSPARSETriFactors_Reset(trifactors));
3023: PetscCallCUSPARSE(cusparseDestroy((*trifactors)->handle));
3024: PetscCall(PetscFree(*trifactors));
3025: }
3026: PetscFunctionReturn(PETSC_SUCCESS);
3027: }
3029: static PetscErrorCode MatSeqAIJCUSPARSEInvalidateTranspose(Mat A, PetscBool destroy)
3030: {
3031: Mat_SeqAIJCUSPARSE *cusp = (Mat_SeqAIJCUSPARSE *)A->spptr;
3033: PetscFunctionBegin;
3034: PetscCheckTypeName(A, MATSEQAIJCUSPARSE);
3035: if (!cusp) PetscFunctionReturn(PETSC_SUCCESS);
3036: if (destroy) {
3037: PetscCall(MatSeqAIJCUSPARSEMultStruct_Destroy(&cusp->matTranspose, cusp->format));
3038: delete cusp->csr2csc_i;
3039: cusp->csr2csc_i = NULL;
3040: }
3041: A->transupdated = PETSC_FALSE;
3042: PetscFunctionReturn(PETSC_SUCCESS);
3043: }
3045: static PetscErrorCode MatSetPreallocationCOO_SeqAIJCUSPARSE(Mat mat, PetscCount coo_n, PetscInt coo_i[], PetscInt coo_j[])
3046: {
3047: PetscFunctionBegin;
3048: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::SetPreallocationCOO(mat, coo_n, coo_i, coo_j));
3049: PetscFunctionReturn(PETSC_SUCCESS);
3050: }
3052: static PetscErrorCode MatSetValuesCOO_SeqAIJCUSPARSE(Mat A, const PetscScalar v[], InsertMode imode)
3053: {
3054: PetscFunctionBegin;
3055: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::SetValuesCOO(A, v, imode));
3056: PetscFunctionReturn(PETSC_SUCCESS);
3057: }
3059: /*@
3060: MatSeqAIJCUSPARSEGetIJ - returns the device row storage `i` and `j` indices for `MATSEQAIJCUSPARSE` matrices.
3062: Not Collective
3064: Input Parameters:
3065: + A - the matrix
3066: - compressed - `PETSC_TRUE` or `PETSC_FALSE` indicating the matrix data structure should be always returned in compressed form
3068: Output Parameters:
3069: + i - the CSR row pointers
3070: - j - the CSR column indices
3072: Level: developer
3074: Note:
3075: When compressed is true, the CSR structure does not contain empty rows
3077: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSERestoreIJ()`, `MatSeqAIJCUSPARSEGetArrayRead()`
3078: @*/
3079: PetscErrorCode MatSeqAIJCUSPARSEGetIJ(Mat A, PetscBool compressed, const PetscInt *i[], const PetscInt *j[])
3080: {
3081: PetscFunctionBegin;
3082: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::GetIJ(A, compressed, i, j));
3083: PetscFunctionReturn(PETSC_SUCCESS);
3084: }
3086: /*@
3087: MatSeqAIJCUSPARSERestoreIJ - restore the device row storage `i` and `j` indices obtained with `MatSeqAIJCUSPARSEGetIJ()`
3089: Not Collective
3091: Input Parameters:
3092: + A - the matrix
3093: . compressed - `PETSC_TRUE` or `PETSC_FALSE` indicating the matrix data structure should be always returned in compressed form
3094: . i - the CSR row pointers
3095: - j - the CSR column indices
3097: Level: developer
3099: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetIJ()`
3100: @*/
3101: PetscErrorCode MatSeqAIJCUSPARSERestoreIJ(Mat A, PetscBool compressed, const PetscInt *i[], const PetscInt *j[])
3102: {
3103: PetscFunctionBegin;
3104: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::RestoreIJ(A, compressed, i, j));
3105: PetscFunctionReturn(PETSC_SUCCESS);
3106: }
3108: /*@
3109: MatSeqAIJCUSPARSEGetArrayRead - gives read-only access to the array where the device data for a `MATSEQAIJCUSPARSE` matrix nonzero entries are stored
3111: Not Collective
3113: Input Parameter:
3114: . A - a `MATSEQAIJCUSPARSE` matrix
3116: Output Parameter:
3117: . a - pointer to the device data
3119: Level: developer
3121: Note:
3122: Will trigger host-to-device copies if the most up-to-date matrix data is on the host
3124: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetArray()`, `MatSeqAIJCUSPARSEGetArrayWrite()`, `MatSeqAIJCUSPARSERestoreArrayRead()`
3125: @*/
3126: PetscErrorCode MatSeqAIJCUSPARSEGetArrayRead(Mat A, const PetscScalar **a)
3127: {
3128: return MatSeqAIJCUSPARSE_CUPM_t::GetArrayRead(A, a);
3129: }
3131: /*@
3132: MatSeqAIJCUSPARSERestoreArrayRead - restore the read-only access array obtained from `MatSeqAIJCUSPARSEGetArrayRead()`
3134: Not Collective
3136: Input Parameters:
3137: + A - a `MATSEQAIJCUSPARSE` matrix
3138: - a - pointer to the device data
3140: Level: developer
3142: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetArrayRead()`
3143: @*/
3144: PetscErrorCode MatSeqAIJCUSPARSERestoreArrayRead(Mat A, const PetscScalar **a)
3145: {
3146: return MatSeqAIJCUSPARSE_CUPM_t::RestoreArrayRead(A, a);
3147: }
3149: /*@
3150: MatSeqAIJCUSPARSEGetArray - gives read-write access to the array where the device data for a `MATSEQAIJCUSPARSE` matrix is stored
3152: Not Collective
3154: Input Parameter:
3155: . A - a `MATSEQAIJCUSPARSE` matrix
3157: Output Parameter:
3158: . a - pointer to the device data
3160: Level: developer
3162: Note:
3163: Will trigger host-to-device copies if the most up-to-date matrix data is on the host
3165: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetArrayRead()`, `MatSeqAIJCUSPARSEGetArrayWrite()`, `MatSeqAIJCUSPARSERestoreArray()`
3166: @*/
3167: PetscErrorCode MatSeqAIJCUSPARSEGetArray(Mat A, PetscScalar **a)
3168: {
3169: return MatSeqAIJCUSPARSE_CUPM_t::GetArray(A, a);
3170: }
3171: /*@
3172: MatSeqAIJCUSPARSERestoreArray - restore the read-write access array obtained from `MatSeqAIJCUSPARSEGetArray()`
3174: Not Collective
3176: Input Parameters:
3177: + A - a `MATSEQAIJCUSPARSE` matrix
3178: - a - pointer to the device data
3180: Level: developer
3182: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetArray()`
3183: @*/
3184: PetscErrorCode MatSeqAIJCUSPARSERestoreArray(Mat A, PetscScalar **a)
3185: {
3186: return MatSeqAIJCUSPARSE_CUPM_t::RestoreArray(A, a);
3187: }
3189: /*@
3190: MatSeqAIJCUSPARSEGetArrayWrite - gives write access to the array where the device data for a `MATSEQAIJCUSPARSE` matrix is stored
3192: Not Collective
3194: Input Parameter:
3195: . A - a `MATSEQAIJCUSPARSE` matrix
3197: Output Parameter:
3198: . a - pointer to the device data
3200: Level: developer
3202: Note:
3203: Does not trigger any host to device copies.
3205: 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
3207: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetArray()`, `MatSeqAIJCUSPARSEGetArrayRead()`, `MatSeqAIJCUSPARSERestoreArrayWrite()`
3208: @*/
3209: PetscErrorCode MatSeqAIJCUSPARSEGetArrayWrite(Mat A, PetscScalar **a)
3210: {
3211: return MatSeqAIJCUSPARSE_CUPM_t::GetArrayWrite(A, a);
3212: }
3214: /*@
3215: MatSeqAIJCUSPARSERestoreArrayWrite - restore the write-only access array obtained from `MatSeqAIJCUSPARSEGetArrayWrite()`
3217: Not Collective
3219: Input Parameters:
3220: + A - a `MATSEQAIJCUSPARSE` matrix
3221: - a - pointer to the device data
3223: Level: developer
3225: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJCUSPARSEGetArrayWrite()`
3226: @*/
3227: PetscErrorCode MatSeqAIJCUSPARSERestoreArrayWrite(Mat A, PetscScalar **a)
3228: {
3229: return MatSeqAIJCUSPARSE_CUPM_t::RestoreArrayWrite(A, a);
3230: }
3232: struct IJCompare4 {
3233: __host__ __device__ inline bool operator()(const thrust::tuple<PetscInt, PetscInt, PetscScalar, PetscInt> &t1, const thrust::tuple<PetscInt, PetscInt, PetscScalar, PetscInt> &t2)
3234: {
3235: if (thrust::get<0>(t1) < thrust::get<0>(t2)) return true;
3236: if (thrust::get<0>(t1) == thrust::get<0>(t2)) return thrust::get<1>(t1) < thrust::get<1>(t2);
3237: return false;
3238: }
3239: };
3241: struct Shift {
3242: PetscInt _shift;
3244: Shift(PetscInt shift) : _shift(shift) { }
3245: __host__ __device__ inline PetscInt operator()(const PetscInt &c) { return c + _shift; }
3246: };
3248: /* merges two SeqAIJCUSPARSE matrices A, B by concatenating their rows. [A';B']' operation in MATLAB notation */
3249: PetscErrorCode MatSeqAIJCUSPARSEMergeMats(Mat A, Mat B, MatReuse reuse, Mat *C)
3250: {
3251: Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data, *b = (Mat_SeqAIJ *)B->data, *c;
3252: Mat_SeqAIJCUSPARSE *Acusp = (Mat_SeqAIJCUSPARSE *)A->spptr, *Bcusp = (Mat_SeqAIJCUSPARSE *)B->spptr, *Ccusp;
3253: Mat_SeqAIJCUSPARSEMultStruct *Cmat;
3254: CsrMatrix *Acsr, *Bcsr, *Ccsr;
3255: PetscInt Annz, Bnnz;
3256: PetscInt i, m, n, zero = 0;
3258: PetscFunctionBegin;
3261: PetscAssertPointer(C, 4);
3262: PetscCheckTypeName(A, MATSEQAIJCUSPARSE);
3263: PetscCheckTypeName(B, MATSEQAIJCUSPARSE);
3264: 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);
3265: PetscCheck(reuse != MAT_INPLACE_MATRIX, PETSC_COMM_SELF, PETSC_ERR_SUP, "MAT_INPLACE_MATRIX not supported");
3266: PetscCheck(Acusp->format != MAT_CUSPARSE_ELL && Acusp->format != MAT_CUSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
3267: PetscCheck(Bcusp->format != MAT_CUSPARSE_ELL && Bcusp->format != MAT_CUSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
3268: if (reuse == MAT_INITIAL_MATRIX) {
3269: m = A->rmap->n;
3270: n = A->cmap->n + B->cmap->n;
3271: PetscCall(MatCreate(PETSC_COMM_SELF, C));
3272: PetscCall(MatSetSizes(*C, m, n, m, n));
3273: PetscCall(MatSetType(*C, MATSEQAIJCUSPARSE));
3274: c = (Mat_SeqAIJ *)(*C)->data;
3275: Ccusp = (Mat_SeqAIJCUSPARSE *)(*C)->spptr;
3276: Cmat = new Mat_SeqAIJCUSPARSEMultStruct;
3277: Ccsr = new CsrMatrix;
3278: Cmat->cprowIndices = NULL;
3279: c->compressedrow.use = PETSC_FALSE;
3280: c->compressedrow.nrows = 0;
3281: c->compressedrow.i = NULL;
3282: c->compressedrow.rindex = NULL;
3283: Ccusp->workVector = NULL;
3284: Ccusp->nrows = m;
3285: Ccusp->mat = Cmat;
3286: Ccusp->mat->mat = Ccsr;
3287: Ccsr->num_rows = m;
3288: Ccsr->num_cols = n;
3289: PetscCallCUSPARSE(cusparseCreateMatDescr(&Cmat->descr));
3290: PetscCallCUSPARSE(cusparseSetMatIndexBase(Cmat->descr, CUSPARSE_INDEX_BASE_ZERO));
3291: PetscCallCUSPARSE(cusparseSetMatType(Cmat->descr, CUSPARSE_MATRIX_TYPE_GENERAL));
3292: PetscCallCUDA(cudaMalloc((void **)&Cmat->alpha_one, sizeof(PetscScalar)));
3293: PetscCallCUDA(cudaMalloc((void **)&Cmat->beta_zero, sizeof(PetscScalar)));
3294: PetscCallCUDA(cudaMalloc((void **)&Cmat->beta_one, sizeof(PetscScalar)));
3295: PetscCallCUDA(cudaMemcpy(Cmat->alpha_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
3296: PetscCallCUDA(cudaMemcpy(Cmat->beta_zero, &PETSC_CUSPARSE_ZERO, sizeof(PetscScalar), cudaMemcpyHostToDevice));
3297: PetscCallCUDA(cudaMemcpy(Cmat->beta_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
3298: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
3299: PetscCall(MatSeqAIJCUSPARSECopyToGPU(B));
3300: PetscCheck(Acusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJCUSPARSEMultStruct");
3301: PetscCheck(Bcusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJCUSPARSEMultStruct");
3303: Acsr = (CsrMatrix *)Acusp->mat->mat;
3304: Bcsr = (CsrMatrix *)Bcusp->mat->mat;
3305: Annz = (PetscInt)Acsr->column_indices->size();
3306: Bnnz = (PetscInt)Bcsr->column_indices->size();
3307: c->nz = Annz + Bnnz;
3308: Ccsr->row_offsets = new THRUSTINTARRAY(m + 1);
3309: Ccsr->column_indices = new THRUSTINTARRAY(c->nz);
3310: Ccsr->values = new THRUSTARRAY(c->nz);
3311: Ccsr->num_entries = c->nz;
3312: Ccusp->coords = new THRUSTINTARRAY(c->nz);
3313: if (c->nz) {
3314: auto Acoo = new THRUSTINTARRAY(Annz); // initialized with zeros
3315: auto Bcoo = new THRUSTINTARRAY(Bnnz);
3316: auto Ccoo = new THRUSTINTARRAY(c->nz);
3317: THRUSTINTARRAY *Aroff, *Broff;
3319: if (a->compressedrow.use) { /* need full row offset */
3320: if (!Acusp->rowoffsets_gpu) {
3321: Acusp->rowoffsets_gpu = new THRUSTINTARRAY(A->rmap->n + 1);
3322: Acusp->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
3323: PetscCall(PetscLogCpuToGpu((A->rmap->n + 1) * sizeof(PetscInt)));
3324: }
3325: Aroff = Acusp->rowoffsets_gpu;
3326: } else Aroff = Acsr->row_offsets;
3327: if (b->compressedrow.use) { /* need full row offset */
3328: if (!Bcusp->rowoffsets_gpu) {
3329: Bcusp->rowoffsets_gpu = new THRUSTINTARRAY(B->rmap->n + 1);
3330: Bcusp->rowoffsets_gpu->assign(b->i, b->i + B->rmap->n + 1);
3331: PetscCall(PetscLogCpuToGpu((B->rmap->n + 1) * sizeof(PetscInt)));
3332: }
3333: Broff = Bcusp->rowoffsets_gpu;
3334: } else Broff = Bcsr->row_offsets;
3335: PetscCall(PetscLogGpuTimeBegin());
3336: // Implement cusparseXcsr2coo() with Thrust, as the former doesn't support 64-bit indices.
3337: PetscCallThrust(thrust::for_each(thrust::device, thrust::counting_iterator<PetscInt>(0), thrust::counting_iterator<PetscInt>(m), Csr2coo(Aroff->data().get(), Acoo->data().get())));
3338: PetscCallThrust(thrust::for_each(thrust::device, thrust::counting_iterator<PetscInt>(0), thrust::counting_iterator<PetscInt>(m), Csr2coo(Broff->data().get(), Bcoo->data().get())));
3340: /* Issues when using bool with large matrices on SUMMIT 10.2.89 */
3341: #if CCCL_VERSION >= 3004000
3342: auto Aperm = cuda::make_constant_iterator(1);
3343: auto Bperm = cuda::make_constant_iterator(0);
3344: #else
3345: auto Aperm = thrust::make_constant_iterator(1);
3346: auto Bperm = thrust::make_constant_iterator(0);
3347: #endif
3348: auto Bcib = thrust::make_transform_iterator(Bcsr->column_indices->begin(), Shift(A->cmap->n));
3349: auto Bcie = thrust::make_transform_iterator(Bcsr->column_indices->end(), Shift(A->cmap->n));
3350: auto wPerm = new THRUSTINTARRAY(Annz + Bnnz);
3351: auto Azb = thrust::make_zip_iterator(thrust::make_tuple(Acoo->begin(), Acsr->column_indices->begin(), Acsr->values->begin(), Aperm));
3352: auto Aze = thrust::make_zip_iterator(thrust::make_tuple(Acoo->end(), Acsr->column_indices->end(), Acsr->values->end(), Aperm));
3353: 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
3354: auto Bze = thrust::make_zip_iterator(thrust::make_tuple(Bcoo->end(), Bcie, Bcsr->values->end(), Bperm));
3355: auto Czb = thrust::make_zip_iterator(thrust::make_tuple(Ccoo->begin(), Ccsr->column_indices->begin(), Ccsr->values->begin(), wPerm->begin()));
3356: auto p1 = Ccusp->coords->begin();
3357: auto p2 = Ccusp->coords->begin();
3358: #if CCCL_VERSION >= 3001000
3359: cuda::std::advance(p2, Annz);
3360: #else
3361: thrust::advance(p2, Annz);
3362: #endif
3363: 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)
3364: auto cci = thrust::make_counting_iterator(zero);
3365: auto cce = thrust::make_counting_iterator(c->nz);
3366: #if PETSC_PKG_CUDA_VERSION_LT(12, 9, 0) || PetscDefined(HAVE_THRUST)
3367: auto pred = thrust::identity<int>();
3368: #else
3369: auto pred = cuda::std::identity();
3370: #endif
3371: PetscCallThrust(thrust::copy_if(thrust::device, cci, cce, wPerm->begin(), p1, pred));
3372: PetscCallThrust(thrust::remove_copy_if(thrust::device, cci, cce, wPerm->begin(), p2, pred));
3373: // Implement a simplified cusparseXcoo2csr() with Thrust (assuming the row indices are already sorted), as the former doesn't support 64-bit indices.
3374: 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()));
3375: PetscCall(PetscLogGpuTimeEnd());
3376: delete wPerm;
3377: delete Acoo;
3378: delete Bcoo;
3379: delete Ccoo;
3380: 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));
3381: if (A->form_explicit_transpose && B->form_explicit_transpose) { /* if A and B have the transpose, generate C transpose too */
3382: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(A));
3383: PetscCall(MatSeqAIJCUSPARSEFormExplicitTranspose(B));
3384: PetscBool AT = Acusp->matTranspose ? PETSC_TRUE : PETSC_FALSE, BT = Bcusp->matTranspose ? PETSC_TRUE : PETSC_FALSE;
3385: Mat_SeqAIJCUSPARSEMultStruct *CmatT = new Mat_SeqAIJCUSPARSEMultStruct;
3386: CsrMatrix *CcsrT = new CsrMatrix;
3387: CsrMatrix *AcsrT = AT ? (CsrMatrix *)Acusp->matTranspose->mat : NULL;
3388: CsrMatrix *BcsrT = BT ? (CsrMatrix *)Bcusp->matTranspose->mat : NULL;
3390: (*C)->form_explicit_transpose = PETSC_TRUE;
3391: (*C)->transupdated = PETSC_TRUE;
3392: Ccusp->rowoffsets_gpu = NULL;
3393: CmatT->cprowIndices = NULL;
3394: CmatT->mat = CcsrT;
3395: CcsrT->num_rows = n;
3396: CcsrT->num_cols = m;
3397: CcsrT->num_entries = c->nz;
3399: CcsrT->row_offsets = new THRUSTINTARRAY(n + 1);
3400: CcsrT->column_indices = new THRUSTINTARRAY(c->nz);
3401: CcsrT->values = new THRUSTARRAY(c->nz);
3403: PetscCall(PetscLogGpuTimeBegin());
3404: auto rT = CcsrT->row_offsets->begin();
3405: if (AT) {
3406: rT = thrust::copy(AcsrT->row_offsets->begin(), AcsrT->row_offsets->end(), rT);
3407: #if CCCL_VERSION >= 3001000
3408: cuda::std::advance(rT, -1);
3409: #else
3410: thrust::advance(rT, -1);
3411: #endif
3412: }
3413: if (BT) {
3414: auto titb = thrust::make_transform_iterator(BcsrT->row_offsets->begin(), Shift(a->nz));
3415: auto tite = thrust::make_transform_iterator(BcsrT->row_offsets->end(), Shift(a->nz));
3416: thrust::copy(titb, tite, rT);
3417: }
3418: auto cT = CcsrT->column_indices->begin();
3419: if (AT) cT = thrust::copy(AcsrT->column_indices->begin(), AcsrT->column_indices->end(), cT);
3420: if (BT) thrust::copy(BcsrT->column_indices->begin(), BcsrT->column_indices->end(), cT);
3421: auto vT = CcsrT->values->begin();
3422: if (AT) vT = thrust::copy(AcsrT->values->begin(), AcsrT->values->end(), vT);
3423: if (BT) thrust::copy(BcsrT->values->begin(), BcsrT->values->end(), vT);
3424: PetscCall(PetscLogGpuTimeEnd());
3426: PetscCallCUSPARSE(cusparseCreateMatDescr(&CmatT->descr));
3427: PetscCallCUSPARSE(cusparseSetMatIndexBase(CmatT->descr, CUSPARSE_INDEX_BASE_ZERO));
3428: PetscCallCUSPARSE(cusparseSetMatType(CmatT->descr, CUSPARSE_MATRIX_TYPE_GENERAL));
3429: PetscCallCUDA(cudaMalloc((void **)&CmatT->alpha_one, sizeof(PetscScalar)));
3430: PetscCallCUDA(cudaMalloc((void **)&CmatT->beta_zero, sizeof(PetscScalar)));
3431: PetscCallCUDA(cudaMalloc((void **)&CmatT->beta_one, sizeof(PetscScalar)));
3432: PetscCallCUDA(cudaMemcpy(CmatT->alpha_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
3433: PetscCallCUDA(cudaMemcpy(CmatT->beta_zero, &PETSC_CUSPARSE_ZERO, sizeof(PetscScalar), cudaMemcpyHostToDevice));
3434: PetscCallCUDA(cudaMemcpy(CmatT->beta_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar), cudaMemcpyHostToDevice));
3435: 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));
3436: Ccusp->matTranspose = CmatT;
3437: }
3438: }
3440: c->free_a = PETSC_TRUE;
3441: PetscCall(PetscShmgetAllocateArray(c->nz, sizeof(PetscInt), (void **)&c->j));
3442: PetscCall(PetscShmgetAllocateArray(m + 1, sizeof(PetscInt), (void **)&c->i));
3443: c->free_ij = PETSC_TRUE;
3444: PetscCallCUDA(cudaMemcpy(c->i, Ccsr->row_offsets->data().get(), Ccsr->row_offsets->size() * sizeof(PetscInt), cudaMemcpyDeviceToHost));
3445: PetscCallCUDA(cudaMemcpy(c->j, Ccsr->column_indices->data().get(), Ccsr->column_indices->size() * sizeof(PetscInt), cudaMemcpyDeviceToHost));
3446: PetscCall(PetscLogGpuToCpu((Ccsr->column_indices->size() + Ccsr->row_offsets->size()) * sizeof(PetscInt)));
3447: PetscCall(PetscMalloc1(m, &c->ilen));
3448: PetscCall(PetscMalloc1(m, &c->imax));
3449: c->maxnz = c->nz;
3450: c->nonzerorowcnt = 0;
3451: c->rmax = 0;
3452: for (i = 0; i < m; i++) {
3453: const PetscInt nn = c->i[i + 1] - c->i[i];
3454: c->ilen[i] = c->imax[i] = nn;
3455: c->nonzerorowcnt += (PetscInt)!!nn;
3456: c->rmax = PetscMax(c->rmax, nn);
3457: }
3458: PetscCall(PetscMalloc1(c->nz, &c->a));
3459: (*C)->nonzerostate++;
3460: PetscCall(PetscLayoutSetUp((*C)->rmap));
3461: PetscCall(PetscLayoutSetUp((*C)->cmap));
3462: Ccusp->nonzerostate = (*C)->nonzerostate;
3463: (*C)->preallocated = PETSC_TRUE;
3464: } else {
3465: 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);
3466: c = (Mat_SeqAIJ *)(*C)->data;
3467: if (c->nz) {
3468: Ccusp = (Mat_SeqAIJCUSPARSE *)(*C)->spptr;
3469: PetscCheck(Ccusp->coords, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing coords");
3470: PetscCheck(Ccusp->format != MAT_CUSPARSE_ELL && Ccusp->format != MAT_CUSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
3471: PetscCheck(Ccusp->nonzerostate == (*C)->nonzerostate, PETSC_COMM_SELF, PETSC_ERR_COR, "Wrong nonzerostate");
3472: PetscCall(MatSeqAIJCUSPARSECopyToGPU(A));
3473: PetscCall(MatSeqAIJCUSPARSECopyToGPU(B));
3474: PetscCheck(Acusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJCUSPARSEMultStruct");
3475: PetscCheck(Bcusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJCUSPARSEMultStruct");
3476: Acsr = (CsrMatrix *)Acusp->mat->mat;
3477: Bcsr = (CsrMatrix *)Bcusp->mat->mat;
3478: Ccsr = (CsrMatrix *)Ccusp->mat->mat;
3479: 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());
3480: 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());
3481: 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());
3482: 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);
3483: 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());
3484: auto pmid = Ccusp->coords->begin();
3485: #if CCCL_VERSION >= 3001000
3486: cuda::std::advance(pmid, Acsr->num_entries);
3487: #else
3488: thrust::advance(pmid, Acsr->num_entries);
3489: #endif
3490: PetscCall(PetscLogGpuTimeBegin());
3491: auto zibait = thrust::make_zip_iterator(thrust::make_tuple(Acsr->values->begin(), thrust::make_permutation_iterator(Ccsr->values->begin(), Ccusp->coords->begin())));
3492: auto zieait = thrust::make_zip_iterator(thrust::make_tuple(Acsr->values->end(), thrust::make_permutation_iterator(Ccsr->values->begin(), pmid)));
3493: thrust::for_each(zibait, zieait, VecCUDAEquals());
3494: auto zibbit = thrust::make_zip_iterator(thrust::make_tuple(Bcsr->values->begin(), thrust::make_permutation_iterator(Ccsr->values->begin(), pmid)));
3495: auto ziebit = thrust::make_zip_iterator(thrust::make_tuple(Bcsr->values->end(), thrust::make_permutation_iterator(Ccsr->values->begin(), Ccusp->coords->end())));
3496: thrust::for_each(zibbit, ziebit, VecCUDAEquals());
3497: PetscCall(MatSeqAIJCUSPARSEInvalidateTranspose(*C, PETSC_FALSE));
3498: if (A->form_explicit_transpose && B->form_explicit_transpose && (*C)->form_explicit_transpose) {
3499: PetscCheck(Ccusp->matTranspose, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing transpose Mat_SeqAIJCUSPARSEMultStruct");
3500: PetscBool AT = Acusp->matTranspose ? PETSC_TRUE : PETSC_FALSE, BT = Bcusp->matTranspose ? PETSC_TRUE : PETSC_FALSE;
3501: CsrMatrix *AcsrT = AT ? (CsrMatrix *)Acusp->matTranspose->mat : NULL;
3502: CsrMatrix *BcsrT = BT ? (CsrMatrix *)Bcusp->matTranspose->mat : NULL;
3503: CsrMatrix *CcsrT = (CsrMatrix *)Ccusp->matTranspose->mat;
3504: auto vT = CcsrT->values->begin();
3505: if (AT) vT = thrust::copy(AcsrT->values->begin(), AcsrT->values->end(), vT);
3506: if (BT) thrust::copy(BcsrT->values->begin(), BcsrT->values->end(), vT);
3507: (*C)->transupdated = PETSC_TRUE;
3508: }
3509: PetscCall(PetscLogGpuTimeEnd());
3510: }
3511: }
3512: PetscCall(PetscObjectStateIncrease((PetscObject)*C));
3513: (*C)->assembled = PETSC_TRUE;
3514: (*C)->was_assembled = PETSC_FALSE;
3515: (*C)->offloadmask = PETSC_OFFLOAD_GPU;
3516: PetscFunctionReturn(PETSC_SUCCESS);
3517: }
3519: static PetscErrorCode MatSeqAIJCopySubArray_SeqAIJCUSPARSE(Mat A, PetscInt n, const PetscInt idx[], PetscScalar v[])
3520: {
3521: PetscFunctionBegin;
3522: PetscCall(MatSeqAIJCUSPARSE_CUPM_t::CopySubArray(A, n, idx, v));
3523: PetscFunctionReturn(PETSC_SUCCESS);
3524: }