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