Actual source code: aijhipsparse.hip.cxx

  1: /*
  2:   Defines the basic matrix operations for the AIJ (compressed row)
  3:   matrix storage format using the HIPSPARSE library,
  4:   Portions of this code are under:
  5:   Copyright (c) 2022 Advanced Micro Devices, Inc. All rights reserved.
  6: */
  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/mat/impls/dense/seq/dense.h>
 11: #include <../src/vec/vec/impls/dvecimpl.h>
 12: #include <petsc/private/vecimpl.h>
 13: #undef VecType
 14: #include <../src/mat/impls/aij/seq/seqhipsparse/hipsparsematimpl.h>
 15: #include <../src/mat/impls/aij/seq/cupm/aijcupm.hpp>
 16: #include <thrust/adjacent_difference.h>
 17: #include <thrust/iterator/transform_iterator.h>
 18: #if PETSC_CPP_VERSION >= 14
 19:   #define PETSC_HAVE_THRUST_ASYNC 1
 20:   #include <thrust/async/for_each.h>
 21: #endif
 22: #include <thrust/iterator/constant_iterator.h>
 23: #include <thrust/iterator/discard_iterator.h>
 24: #include <thrust/binary_search.h>
 25: #include <thrust/remove.h>
 26: #include <thrust/sort.h>
 27: #include <thrust/unique.h>
 28: #include <thrust/gather.h>
 29: #include <thrust/functional.h>

 31: const char *const MatHIPSPARSEStorageFormats[] = {"CSR", "ELL", "HYB", "MatHIPSPARSEStorageFormat", "MAT_HIPSPARSE_", 0};
 32: const char *const MatHIPSPARSESpMVAlgorithms[] = {"MV_ALG_DEFAULT", "COOMV_ALG", "CSRMV_ALG1", "CSRMV_ALG2", "SPMV_ALG_DEFAULT", "SPMV_COO_ALG1", "SPMV_COO_ALG2", "SPMV_CSR_ALG1", "SPMV_CSR_ALG2", "hipsparseSpMVAlg_t", "HIPSPARSE_", 0};
 33: const char *const MatHIPSPARSESpMMAlgorithms[] = {"ALG_DEFAULT", "COO_ALG1", "COO_ALG2", "COO_ALG3", "CSR_ALG1", "COO_ALG4", "CSR_ALG2", "hipsparseSpMMAlg_t", "HIPSPARSE_SPMM_", 0};
 34: //const char *const MatHIPSPARSECsr2CscAlgorithms[] = {"INVALID"/*HIPSPARSE does not have enum 0! We created one*/, "ALG1", "ALG2", "hipsparseCsr2CscAlg_t", "HIPSPARSE_CSR2CSC_", 0};

 36: static PetscErrorCode MatICCFactorSymbolic_SeqAIJHIPSPARSE(Mat, Mat, IS, const MatFactorInfo *);
 37: static PetscErrorCode MatCholeskyFactorSymbolic_SeqAIJHIPSPARSE(Mat, Mat, IS, const MatFactorInfo *);
 38: static PetscErrorCode MatCholeskyFactorNumeric_SeqAIJHIPSPARSE(Mat, Mat, const MatFactorInfo *);
 39: static PetscErrorCode MatILUFactorSymbolic_SeqAIJHIPSPARSE(Mat, Mat, IS, IS, const MatFactorInfo *);
 40: static PetscErrorCode MatSetFromOptions_SeqAIJHIPSPARSE(Mat, PetscOptionItems PetscOptionsObject);
 41: static PetscErrorCode MatAXPY_SeqAIJHIPSPARSE(Mat, PetscScalar, Mat, MatStructure);
 42: static PetscErrorCode MatScale_SeqAIJHIPSPARSE(Mat, PetscScalar);
 43: static PetscErrorCode MatDiagonalScale_SeqAIJHIPSPARSE(Mat, Vec, Vec);
 44: static PetscErrorCode MatMult_SeqAIJHIPSPARSE(Mat, Vec, Vec);
 45: static PetscErrorCode MatMultAdd_SeqAIJHIPSPARSE(Mat, Vec, Vec, Vec);
 46: static PetscErrorCode MatMultTranspose_SeqAIJHIPSPARSE(Mat, Vec, Vec);
 47: static PetscErrorCode MatMultTransposeAdd_SeqAIJHIPSPARSE(Mat, Vec, Vec, Vec);
 48: static PetscErrorCode MatMultHermitianTranspose_SeqAIJHIPSPARSE(Mat, Vec, Vec);
 49: static PetscErrorCode MatMultHermitianTransposeAdd_SeqAIJHIPSPARSE(Mat, Vec, Vec, Vec);
 50: static PetscErrorCode MatMultAddKernel_SeqAIJHIPSPARSE(Mat, Vec, Vec, Vec, PetscBool, PetscBool);

 52: static PetscErrorCode CsrMatrix_Destroy(CsrMatrix **);
 53: static PetscErrorCode MatSeqAIJHIPSPARSEMultStruct_Destroy(Mat_SeqAIJHIPSPARSEMultStruct **, MatHIPSPARSEStorageFormat);
 54: static PetscErrorCode MatSeqAIJHIPSPARSETriFactors_Destroy(Mat_SeqAIJHIPSPARSETriFactors **);
 55: static PetscErrorCode MatSeqAIJHIPSPARSE_Destroy(Mat);

 57: static PetscErrorCode MatSeqAIJHIPSPARSECopyFromGPU(Mat);
 58: static PetscErrorCode MatSeqAIJHIPSPARSEInvalidateTranspose(Mat, PetscBool);

 60: static PetscErrorCode MatSeqAIJCopySubArray_SeqAIJHIPSPARSE(Mat, PetscInt, const PetscInt[], PetscScalar[]);
 61: static PetscErrorCode MatSetPreallocationCOO_SeqAIJHIPSPARSE(Mat, PetscCount, PetscInt[], PetscInt[]);
 62: static PetscErrorCode MatSetValuesCOO_SeqAIJHIPSPARSE(Mat, const PetscScalar[], InsertMode);

 64: PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJHIPSPARSE(Mat, MatType, MatReuse, Mat *);

 66: // cusparseCreateCsr() separates types for row offsets and column indices in prototype, but requires them to have the same type at runtime!
 67: const hipsparseIndexType_t csrRowOffsetsType = PetscDefined(USE_64BIT_INDICES) ? HIPSPARSE_INDEX_64I : HIPSPARSE_INDEX_32I;
 68: const hipsparseIndexType_t csrColIndType     = PetscDefined(USE_64BIT_INDICES) ? HIPSPARSE_INDEX_64I : HIPSPARSE_INDEX_32I;

 70: using Csr2coo        = Petsc::mat::aij::cupm::impl::Csr2coo;
 71: using PetscIntToCInt = Petsc::mat::aij::cupm::impl::PetscIntToCInt;

 73: /*
 74: PetscErrorCode MatHIPSPARSESetStream(Mat A, const hipStream_t stream)
 75: {
 76:   Mat_SeqAIJHIPSPARSE *hipsparsestruct = (Mat_SeqAIJHIPSPARSE*)A->spptr;

 78:   PetscFunctionBegin;
 79:   PetscCheck(hipsparsestruct, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing spptr");
 80:   hipsparsestruct->stream = stream;
 81:   PetscCallHIPSPARSE(hipsparseSetStream(hipsparsestruct->handle, hipsparsestruct->stream));
 82:   PetscFunctionReturn(PETSC_SUCCESS);
 83: }

 85: PetscErrorCode MatHIPSPARSESetHandle(Mat A, const hipsparseHandle_t handle)
 86: {
 87:   Mat_SeqAIJHIPSPARSE *hipsparsestruct = (Mat_SeqAIJHIPSPARSE*)A->spptr;

 89:   PetscFunctionBegin;
 90:   PetscCheck(hipsparsestruct, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing spptr");
 91:   if (hipsparsestruct->handle != handle) {
 92:     if (hipsparsestruct->handle) PetscCallHIPSPARSE(hipsparseDestroy(hipsparsestruct->handle));
 93:     hipsparsestruct->handle = handle;
 94:   }
 95:   PetscCallHIPSPARSE(hipsparseSetPointerMode(hipsparsestruct->handle, HIPSPARSE_POINTER_MODE_DEVICE));
 96:   PetscFunctionReturn(PETSC_SUCCESS);
 97: }

 99: PetscErrorCode MatHIPSPARSEClearHandle(Mat A)
100: {
101:   Mat_SeqAIJHIPSPARSE *hipsparsestruct = (Mat_SeqAIJHIPSPARSE*)A->spptr;
102:   PetscBool            flg;

104:   PetscFunctionBegin;
105:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
106:   if (!flg || !hipsparsestruct) PetscFunctionReturn(PETSC_SUCCESS);
107:   if (hipsparsestruct->handle) hipsparsestruct->handle = 0;
108:   PetscFunctionReturn(PETSC_SUCCESS);
109: }
110: */

112: PETSC_INTERN PetscErrorCode MatHIPSPARSESetFormat_SeqAIJHIPSPARSE(Mat A, MatHIPSPARSEFormatOperation op, MatHIPSPARSEStorageFormat format)
113: {
114:   Mat_SeqAIJHIPSPARSE *hipsparsestruct = (Mat_SeqAIJHIPSPARSE *)A->spptr;

116:   PetscFunctionBegin;
117:   switch (op) {
118:   case MAT_HIPSPARSE_MULT:
119:     hipsparsestruct->format = format;
120:     break;
121:   case MAT_HIPSPARSE_ALL:
122:     hipsparsestruct->format = format;
123:     break;
124:   default:
125:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "unsupported operation %d for MatHIPSPARSEFormatOperation. MAT_HIPSPARSE_MULT and MAT_HIPSPARSE_ALL are currently supported.", op);
126:   }
127:   PetscFunctionReturn(PETSC_SUCCESS);
128: }

130: /*@
131:   MatHIPSPARSESetFormat - Sets the storage format of `MATSEQHIPSPARSE` matrices for a particular
132:   operation. Only the `MatMult()` operation can use different GPU storage formats

134:   Not Collective

136:   Input Parameters:
137: + A      - Matrix of type `MATSEQAIJHIPSPARSE`
138: . op     - `MatHIPSPARSEFormatOperation`. `MATSEQAIJHIPSPARSE` matrices support `MAT_HIPSPARSE_MULT` and `MAT_HIPSPARSE_ALL`.
139:          `MATMPIAIJHIPSPARSE` matrices support `MAT_HIPSPARSE_MULT_DIAG`, `MAT_HIPSPARSE_MULT_OFFDIAG`, and `MAT_HIPSPARSE_ALL`.
140: - format - `MatHIPSPARSEStorageFormat` (one of `MAT_HIPSPARSE_CSR`, `MAT_HIPSPARSE_ELL`, `MAT_HIPSPARSE_HYB`.)

142:   Level: intermediate

144: .seealso: [](ch_matrices), `Mat`, `MATSEQAIJHIPSPARSE`, `MatHIPSPARSEStorageFormat`, `MatHIPSPARSEFormatOperation`
145: @*/
146: PetscErrorCode MatHIPSPARSESetFormat(Mat A, MatHIPSPARSEFormatOperation op, MatHIPSPARSEStorageFormat format)
147: {
148:   PetscFunctionBegin;
150:   PetscTryMethod(A, "MatHIPSPARSESetFormat_C", (Mat, MatHIPSPARSEFormatOperation, MatHIPSPARSEStorageFormat), (A, op, format));
151:   PetscFunctionReturn(PETSC_SUCCESS);
152: }

154: PETSC_INTERN PetscErrorCode MatHIPSPARSESetUseCPUSolve_SeqAIJHIPSPARSE(Mat A, PetscBool use_cpu)
155: {
156:   Mat_SeqAIJHIPSPARSE *hipsparsestruct = (Mat_SeqAIJHIPSPARSE *)A->spptr;

158:   PetscFunctionBegin;
159:   hipsparsestruct->use_cpu_solve = use_cpu;
160:   PetscFunctionReturn(PETSC_SUCCESS);
161: }

163: /*@
164:   MatHIPSPARSESetUseCPUSolve - Sets use CPU `MatSolve()`.

166:   Input Parameters:
167: + A       - Matrix of type `MATSEQAIJHIPSPARSE`
168: - use_cpu - set flag for using the built-in CPU `MatSolve()`

170:   Level: intermediate

172:   Notes:
173:   The hipSparse LU solver currently computes the factors with the built-in CPU method
174:   and moves the factors to the GPU for the solve. We have observed better performance keeping the data on the CPU and computing the solve there.
175:   This method to specifies if the solve is done on the CPU or GPU (GPU is the default).

177: .seealso: [](ch_matrices), `Mat`, `MatSolve()`, `MATSEQAIJHIPSPARSE`, `MatHIPSPARSEStorageFormat`, `MatHIPSPARSEFormatOperation`
178: @*/
179: PetscErrorCode MatHIPSPARSESetUseCPUSolve(Mat A, PetscBool use_cpu)
180: {
181:   PetscFunctionBegin;
183:   PetscTryMethod(A, "MatHIPSPARSESetUseCPUSolve_C", (Mat, PetscBool), (A, use_cpu));
184:   PetscFunctionReturn(PETSC_SUCCESS);
185: }

187: static PetscErrorCode MatSetOption_SeqAIJHIPSPARSE(Mat A, MatOption op, PetscBool flg)
188: {
189:   PetscFunctionBegin;
190:   switch (op) {
191:   case MAT_FORM_EXPLICIT_TRANSPOSE:
192:     /* need to destroy the transpose matrix if present to prevent from logic errors if flg is set to true later */
193:     if (A->form_explicit_transpose && !flg) PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(A, PETSC_TRUE));
194:     A->form_explicit_transpose = flg;
195:     break;
196:   default:
197:     PetscCall(MatSetOption_SeqAIJ(A, op, flg));
198:     break;
199:   }
200:   PetscFunctionReturn(PETSC_SUCCESS);
201: }

203: static PetscErrorCode MatSetFromOptions_SeqAIJHIPSPARSE(Mat A, PetscOptionItems PetscOptionsObject)
204: {
205:   MatHIPSPARSEStorageFormat format;
206:   PetscBool                 flg;
207:   Mat_SeqAIJHIPSPARSE      *hipsparsestruct = (Mat_SeqAIJHIPSPARSE *)A->spptr;

209:   PetscFunctionBegin;
210:   PetscOptionsHeadBegin(PetscOptionsObject, "SeqAIJHIPSPARSE options");
211:   if (A->factortype == MAT_FACTOR_NONE) {
212:     PetscCall(PetscOptionsEnum("-mat_hipsparse_mult_storage_format", "sets storage format of (seq)aijhipsparse gpu matrices for SpMV", "MatHIPSPARSESetFormat", MatHIPSPARSEStorageFormats, (PetscEnum)hipsparsestruct->format, (PetscEnum *)&format, &flg));
213:     if (flg) PetscCall(MatHIPSPARSESetFormat(A, MAT_HIPSPARSE_MULT, format));
214:     PetscCall(PetscOptionsEnum("-mat_hipsparse_storage_format", "sets storage format of (seq)aijhipsparse gpu matrices for SpMV and TriSolve", "MatHIPSPARSESetFormat", MatHIPSPARSEStorageFormats, (PetscEnum)hipsparsestruct->format, (PetscEnum *)&format, &flg));
215:     if (flg) PetscCall(MatHIPSPARSESetFormat(A, MAT_HIPSPARSE_ALL, format));
216:     PetscCall(PetscOptionsBool("-mat_hipsparse_use_cpu_solve", "Use CPU (I)LU solve", "MatHIPSPARSESetUseCPUSolve", hipsparsestruct->use_cpu_solve, &hipsparsestruct->use_cpu_solve, &flg));
217:     if (flg) PetscCall(MatHIPSPARSESetUseCPUSolve(A, hipsparsestruct->use_cpu_solve));
218:     PetscCall(
219:       PetscOptionsEnum("-mat_hipsparse_spmv_alg", "sets hipSPARSE algorithm used in sparse-mat dense-vector multiplication (SpMV)", "hipsparseSpMVAlg_t", MatHIPSPARSESpMVAlgorithms, (PetscEnum)hipsparsestruct->spmvAlg, (PetscEnum *)&hipsparsestruct->spmvAlg, &flg));
220:     /* If user did use this option, check its consistency with hipSPARSE, since PetscOptionsEnum() sets enum values based on their position in MatHIPSPARSESpMVAlgorithms[] */
221:     PetscCheck(!flg || HIPSPARSE_SPMV_CSR_ALG1 == 2, PETSC_COMM_SELF, PETSC_ERR_SUP, "hipSPARSE enum hipsparseSpMVAlg_t has been changed but PETSc has not been updated accordingly");
222:     PetscCall(
223:       PetscOptionsEnum("-mat_hipsparse_spmm_alg", "sets hipSPARSE algorithm used in sparse-mat dense-mat multiplication (SpMM)", "hipsparseSpMMAlg_t", MatHIPSPARSESpMMAlgorithms, (PetscEnum)hipsparsestruct->spmmAlg, (PetscEnum *)&hipsparsestruct->spmmAlg, &flg));
224:     PetscCheck(!flg || HIPSPARSE_SPMM_CSR_ALG1 == 4, PETSC_COMM_SELF, PETSC_ERR_SUP, "hipSPARSE enum hipsparseSpMMAlg_t has been changed but PETSc has not been updated accordingly");
225:     /*
226:     PetscCall(PetscOptionsEnum("-mat_hipsparse_csr2csc_alg", "sets hipSPARSE algorithm used in converting CSR matrices to CSC matrices", "hipsparseCsr2CscAlg_t", MatHIPSPARSECsr2CscAlgorithms, (PetscEnum)hipsparsestruct->csr2cscAlg, (PetscEnum*)&hipsparsestruct->csr2cscAlg, &flg));
227:     PetscCheck(!flg || HIPSPARSE_CSR2CSC_ALG1 == 1, PETSC_COMM_SELF, PETSC_ERR_SUP, "hipSPARSE enum hipsparseCsr2CscAlg_t has been changed but PETSc has not been updated accordingly");
228:     */
229:   }
230:   PetscOptionsHeadEnd();
231:   PetscFunctionReturn(PETSC_SUCCESS);
232: }

234: static PetscErrorCode MatSeqAIJHIPSPARSEBuildFactoredMatrix_LU(Mat A)
235: {
236:   Mat_SeqAIJ                    *a  = static_cast<Mat_SeqAIJ *>(A->data);
237:   PetscInt                       m  = A->rmap->n;
238:   Mat_SeqAIJHIPSPARSETriFactors *fs = static_cast<Mat_SeqAIJHIPSPARSETriFactors *>(A->spptr);
239:   const PetscInt                *Ai = a->i, *Aj = a->j, *adiag;
240:   const MatScalar               *Aa = a->a;
241:   PetscInt                      *Mi, *Mj, Mnz;
242:   PetscScalar                   *Ma;

244:   PetscFunctionBegin;
245:   PetscCall(MatGetDiagonalMarkers_SeqAIJ(A, &adiag, NULL));
246:   if (A->offloadmask == PETSC_OFFLOAD_CPU) { // A's latest factors are on CPU
247:     if (!fs->csrRowPtr) {                    // Is this the first time we are doing setup? Use csrRowPtr since it is not null even when m=0
248:       // Re-arrange the (skewed) factored matrix and put the result into M, a regular csr matrix on host
249:       Mnz = (Ai[m] - Ai[0]) + (adiag[0] - adiag[m]); // Lnz (without the unit diagonal) + Unz (with the non-unit diagonal)
250:       PetscCall(PetscMalloc1(m + 1, &Mi));
251:       PetscCall(PetscMalloc1(Mnz, &Mj)); // Mj is temp
252:       PetscCall(PetscMalloc1(Mnz, &Ma));
253:       Mi[0] = 0;
254:       for (PetscInt i = 0; i < m; i++) {
255:         PetscInt llen = Ai[i + 1] - Ai[i];
256:         PetscInt ulen = adiag[i] - adiag[i + 1];
257:         PetscCall(PetscArraycpy(Mj + Mi[i], Aj + Ai[i], llen));                           // entries of L
258:         Mj[Mi[i] + llen] = i;                                                             // diagonal entry
259:         PetscCall(PetscArraycpy(Mj + Mi[i] + llen + 1, Aj + adiag[i + 1] + 1, ulen - 1)); // entries of U on the right of the diagonal
260:         Mi[i + 1] = Mi[i] + llen + ulen;
261:       }
262:       // Copy M (L,U) from host to device
263:       PetscCallHIP(hipMalloc(&fs->csrRowPtr, sizeof(*fs->csrRowPtr) * (m + 1)));
264:       PetscCallHIP(hipMalloc(&fs->csrColIdx, sizeof(*fs->csrColIdx) * Mnz));
265:       PetscCallHIP(hipMalloc(&fs->csrVal, sizeof(*fs->csrVal) * Mnz));
266:       PetscCallHIP(hipMemcpy(fs->csrRowPtr, Mi, sizeof(*fs->csrRowPtr) * (m + 1), hipMemcpyHostToDevice));
267:       PetscCallHIP(hipMemcpy(fs->csrColIdx, Mj, sizeof(*fs->csrColIdx) * Mnz, hipMemcpyHostToDevice));

269:       // Create descriptors for L, U. See https://docs.nvidia.com/cuda/cusparse/index.html#cusparseDiagType_t
270:       // cusparseDiagType_t: This type indicates if the matrix diagonal entries are unity. The diagonal elements are always
271:       // assumed to be present, but if CUSPARSE_DIAG_TYPE_UNIT is passed to an API routine, then the routine assumes that
272:       // all diagonal entries are unity and will not read or modify those entries. Note that in this case the routine
273:       // assumes the diagonal entries are equal to one, regardless of what those entries are actually set to in memory.
274:       hipsparseFillMode_t fillMode = HIPSPARSE_FILL_MODE_LOWER;
275:       hipsparseDiagType_t diagType = HIPSPARSE_DIAG_TYPE_UNIT;

277:       PetscCallHIPSPARSE(hipsparseCreateCsr(&fs->spMatDescr_L, m, m, Mnz, fs->csrRowPtr, fs->csrColIdx, fs->csrVal, csrRowOffsetsType, csrColIndType, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
278:       PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_L, HIPSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
279:       PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_L, HIPSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));

281:       fillMode = HIPSPARSE_FILL_MODE_UPPER;
282:       diagType = HIPSPARSE_DIAG_TYPE_NON_UNIT;
283:       PetscCallHIPSPARSE(hipsparseCreateCsr(&fs->spMatDescr_U, m, m, Mnz, fs->csrRowPtr, fs->csrColIdx, fs->csrVal, csrRowOffsetsType, csrColIndType, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
284:       PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_U, HIPSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
285:       PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_U, HIPSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));

287:       // Allocate work vectors in SpSv
288:       PetscCallHIP(hipMalloc((void **)&fs->X, sizeof(*fs->X) * m));
289:       PetscCallHIP(hipMalloc((void **)&fs->Y, sizeof(*fs->Y) * m));

291:       PetscCallHIPSPARSE(hipsparseCreateDnVec(&fs->dnVecDescr_X, m, fs->X, hipsparse_scalartype));
292:       PetscCallHIPSPARSE(hipsparseCreateDnVec(&fs->dnVecDescr_Y, m, fs->Y, hipsparse_scalartype));

294:       // Query buffer sizes for SpSV and then allocate buffers, temporarily assuming opA = HIPSPARSE_OPERATION_NON_TRANSPOSE
295:       PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_L));
296:       PetscCallHIPSPARSE(hipsparseSpSV_bufferSize(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, &fs->spsvBufferSize_L));
297:       PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_U));
298:       PetscCallHIPSPARSE(hipsparseSpSV_bufferSize(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_U, &fs->spsvBufferSize_U));
299:       PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_U, fs->spsvBufferSize_U));
300:       PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_L, fs->spsvBufferSize_L));

302:       // Record for reuse
303:       fs->csrRowPtr_h = Mi;
304:       fs->csrVal_h    = Ma;
305:       PetscCall(PetscFree(Mj));
306:     }
307:     // Copy the value
308:     Mi  = fs->csrRowPtr_h;
309:     Ma  = fs->csrVal_h;
310:     Mnz = Mi[m];
311:     for (PetscInt i = 0; i < m; i++) {
312:       PetscInt llen = Ai[i + 1] - Ai[i];
313:       PetscInt ulen = adiag[i] - adiag[i + 1];
314:       PetscCall(PetscArraycpy(Ma + Mi[i], Aa + Ai[i], llen));                           // entries of L
315:       Ma[Mi[i] + llen] = (MatScalar)1.0 / Aa[adiag[i]];                                 // recover the diagonal entry
316:       PetscCall(PetscArraycpy(Ma + Mi[i] + llen + 1, Aa + adiag[i + 1] + 1, ulen - 1)); // entries of U on the right of the diagonal
317:     }
318:     PetscCallHIP(hipMemcpy(fs->csrVal, Ma, sizeof(*Ma) * Mnz, hipMemcpyHostToDevice));

320:     {
321:       // Do hipsparseSpSV_analysis(), which is numeric and requires valid and up-to-date matrix values
322:       PetscCallHIPSPARSE(hipsparseSpSV_analysis(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, fs->spsvBuffer_L));

324:       PetscCallHIPSPARSE(hipsparseSpSV_analysis(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_U, fs->spsvBuffer_U));
325:       fs->updatedSpSVAnalysis          = PETSC_TRUE;
326:       fs->updatedTransposeSpSVAnalysis = PETSC_FALSE;
327:     }
328:   }
329:   PetscFunctionReturn(PETSC_SUCCESS);
330: }

332: static PetscErrorCode MatSeqAIJHIPSPARSEILUAnalysisAndCopyToGPU(Mat A)
333: {
334:   Mat_SeqAIJ                    *a                   = (Mat_SeqAIJ *)A->data;
335:   Mat_SeqAIJHIPSPARSETriFactors *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;
336:   IS                             isrow = a->row, isicol = a->icol;
337:   PetscBool                      row_identity, col_identity;
338:   PetscInt                       n = A->rmap->n;

340:   PetscFunctionBegin;
341:   PetscCheck(hipsparseTriFactors, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing hipsparseTriFactors");
342:   PetscCall(MatSeqAIJHIPSPARSEBuildFactoredMatrix_LU(A));

344:   hipsparseTriFactors->nnz = a->nz;

346:   A->offloadmask = PETSC_OFFLOAD_BOTH; // factored matrix is sync'ed to GPU
347:   /* lower triangular indices */
348:   PetscCall(ISIdentity(isrow, &row_identity));
349:   if (!row_identity && !hipsparseTriFactors->rpermIndices) {
350:     const PetscInt *r;

352:     PetscCall(ISGetIndices(isrow, &r));
353:     hipsparseTriFactors->rpermIndices = new THRUSTINTARRAY(n);
354:     hipsparseTriFactors->rpermIndices->assign(r, r + n);
355:     PetscCall(ISRestoreIndices(isrow, &r));
356:     PetscCall(PetscLogCpuToGpu(n * sizeof(PetscInt)));
357:   }

359:   /* upper triangular indices */
360:   PetscCall(ISIdentity(isicol, &col_identity));
361:   if (!col_identity && !hipsparseTriFactors->cpermIndices) {
362:     const PetscInt *c;

364:     PetscCall(ISGetIndices(isicol, &c));
365:     hipsparseTriFactors->cpermIndices = new THRUSTINTARRAY(n);
366:     hipsparseTriFactors->cpermIndices->assign(c, c + n);
367:     PetscCall(ISRestoreIndices(isicol, &c));
368:     PetscCall(PetscLogCpuToGpu(n * sizeof(PetscInt)));
369:   }
370:   PetscFunctionReturn(PETSC_SUCCESS);
371: }

373: /*
374:   Affected rocSPARSE versions cache SpSV analysis with a descriptor-wide flag, so analysing an
375:   ordinary solve causes subsequent transpose analysis to be skipped. This occurs in the inspected
376:   ROCm releases from 5.0.0 through 7.2.0. Separate descriptors avoid sharing that flag while sharing
377:   the CSR arrays. Upstream fixes the cache to distinguish operation and fill mode:
378:   https://github.com/ROCm/rocm-libraries/pull/9295
379: */
380: static PetscErrorCode MatSeqAIJHIPSPARSECreateTransposeSolveMat(hipsparseSpMatDescr_t mat, hipsparseSpMatDescr_t *trans)
381: {
382:   int64_t              rows, cols, nnz;
383:   void                *rowptr, *colidx, *values;
384:   hipsparseIndexType_t rowtype, coltype;
385:   hipsparseIndexBase_t base;
386:   hipDataType          datatype;
387:   hipsparseFillMode_t  fill;
388:   hipsparseDiagType_t  diag;

390:   PetscFunctionBegin;
391:   PetscCallHIPSPARSE(hipsparseCsrGet(mat, &rows, &cols, &nnz, &rowptr, &colidx, &values, &rowtype, &coltype, &base, &datatype));
392:   PetscCallHIPSPARSE(hipsparseSpMatGetAttribute(mat, HIPSPARSE_SPMAT_FILL_MODE, &fill, sizeof(fill)));
393:   PetscCallHIPSPARSE(hipsparseSpMatGetAttribute(mat, HIPSPARSE_SPMAT_DIAG_TYPE, &diag, sizeof(diag)));
394:   PetscCallHIPSPARSE(hipsparseCreateCsr(trans, rows, cols, nnz, rowptr, colidx, values, rowtype, coltype, base, datatype));
395:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(*trans, HIPSPARSE_SPMAT_FILL_MODE, &fill, sizeof(fill)));
396:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(*trans, HIPSPARSE_SPMAT_DIAG_TYPE, &diag, sizeof(diag)));
397:   PetscFunctionReturn(PETSC_SUCCESS);
398: }

400: static PetscErrorCode MatSeqAIJHIPSPARSEBuildFactoredMatrix_Cholesky(Mat A)
401: {
402:   Mat_SeqAIJ                    *a  = static_cast<Mat_SeqAIJ *>(A->data);
403:   PetscInt                       m  = A->rmap->n;
404:   Mat_SeqAIJHIPSPARSETriFactors *fs = static_cast<Mat_SeqAIJHIPSPARSETriFactors *>(A->spptr);
405:   const PetscInt                *Ai = a->i, *Aj = a->j, *adiag;
406:   const MatScalar               *Aa = a->a;
407:   PetscInt                      *Mj, Mnz;
408:   PetscScalar                   *Ma, *D;

410:   PetscFunctionBegin;
411:   PetscCall(MatGetDiagonalMarkers_SeqAIJ(A, &adiag, NULL));
412:   if (A->offloadmask == PETSC_OFFLOAD_CPU) { // A's latest factors are on CPU
413:     if (!fs->csrRowPtr) {                    // Is this the first time we are doing setup? Use csrRowPtr since it is not null even m=0
414:       // Re-arrange the (skewed) factored matrix and put the result into M, a regular csr matrix on host.
415:       // See comments at MatICCFactorSymbolic_SeqAIJ() on the layout of the factored matrix (U) on host.
416:       Mnz = Ai[m]; // Unz (with the unit diagonal)
417:       PetscCall(PetscMalloc1(Mnz, &Ma));
418:       PetscCall(PetscMalloc1(Mnz, &Mj)); // Mj[] is temp
419:       PetscCall(PetscMalloc1(m, &D));    // the diagonal
420:       for (PetscInt i = 0; i < m; i++) {
421:         PetscInt ulen = Ai[i + 1] - Ai[i];
422:         Mj[Ai[i]]     = i;                                              // diagonal entry
423:         PetscCall(PetscArraycpy(Mj + Ai[i] + 1, Aj + Ai[i], ulen - 1)); // entries of U on the right of the diagonal
424:       }
425:       // Copy M (U) from host to device
426:       PetscCallHIP(hipMalloc(&fs->csrRowPtr, sizeof(*fs->csrRowPtr) * (m + 1)));
427:       PetscCallHIP(hipMalloc(&fs->csrColIdx, sizeof(*fs->csrColIdx) * Mnz));
428:       PetscCallHIP(hipMalloc(&fs->csrVal, sizeof(*fs->csrVal) * Mnz));
429:       PetscCallHIP(hipMalloc(&fs->diag, sizeof(*fs->diag) * m));
430:       PetscCallHIP(hipMemcpy(fs->csrRowPtr, Ai, sizeof(*Ai) * (m + 1), hipMemcpyHostToDevice));
431:       PetscCallHIP(hipMemcpy(fs->csrColIdx, Mj, sizeof(*Mj) * Mnz, hipMemcpyHostToDevice));

433:       // Create descriptors for L, U. See https://docs.nvidia.com/cuda/cusparse/index.html#cusparseDiagType_t
434:       // cusparseDiagType_t: This type indicates if the matrix diagonal entries are unity. The diagonal elements are always
435:       // assumed to be present, but if CUSPARSE_DIAG_TYPE_UNIT is passed to an API routine, then the routine assumes that
436:       // all diagonal entries are unity and will not read or modify those entries. Note that in this case the routine
437:       // assumes the diagonal entries are equal to one, regardless of what those entries are actually set to in memory.
438:       hipsparseFillMode_t fillMode = HIPSPARSE_FILL_MODE_UPPER;
439:       hipsparseDiagType_t diagType = HIPSPARSE_DIAG_TYPE_UNIT; // U is unit diagonal

441:       PetscCallHIPSPARSE(hipsparseCreateCsr(&fs->spMatDescr_U, m, m, Mnz, fs->csrRowPtr, fs->csrColIdx, fs->csrVal, csrRowOffsetsType, csrColIndType, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
442:       PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_U, HIPSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
443:       PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_U, HIPSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));

445:       // Allocate work vectors in SpSv
446:       PetscCallHIP(hipMalloc((void **)&fs->X, sizeof(*fs->X) * m));
447:       PetscCallHIP(hipMalloc((void **)&fs->Y, sizeof(*fs->Y) * m));

449:       PetscCallHIPSPARSE(hipsparseCreateDnVec(&fs->dnVecDescr_X, m, fs->X, hipsparse_scalartype));
450:       PetscCallHIPSPARSE(hipsparseCreateDnVec(&fs->dnVecDescr_Y, m, fs->Y, hipsparse_scalartype));

452:       // Query buffer sizes for SpSV and then allocate buffers
453:       PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_U));
454:       PetscCallHIPSPARSE(hipsparseSpSV_bufferSize(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_U, &fs->spsvBufferSize_U));
455:       PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_U, fs->spsvBufferSize_U));

457:       PetscCall(MatSeqAIJHIPSPARSECreateTransposeSolveMat(fs->spMatDescr_U, &fs->spMatDescr_Ut));
458:       PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_Ut));
459:       PetscCallHIPSPARSE(hipsparseSpSV_bufferSize(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Ut, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Ut, &fs->spsvBufferSize_Ut));
460:       PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_Ut, fs->spsvBufferSize_Ut));

462:       // Record for reuse
463:       fs->csrVal_h = Ma;
464:       fs->diag_h   = D;
465:       PetscCall(PetscFree(Mj));
466:     }
467:     // Copy the value
468:     Ma  = fs->csrVal_h;
469:     D   = fs->diag_h;
470:     Mnz = Ai[m];
471:     for (PetscInt i = 0; i < m; i++) {
472:       D[i]      = Aa[adiag[i]];   // actually Aa[adiag[i]] is the inverse of the diagonal
473:       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
474:       for (PetscInt k = 0; k < Ai[i + 1] - Ai[i] - 1; k++) Ma[Ai[i] + 1 + k] = -Aa[Ai[i] + k];
475:     }
476:     PetscCallHIP(hipMemcpy(fs->csrVal, Ma, sizeof(*Ma) * Mnz, hipMemcpyHostToDevice));
477:     PetscCallHIP(hipMemcpy(fs->diag, D, sizeof(*D) * m, hipMemcpyHostToDevice));

479:     {
480:       // Do hipsparseSpSV_analysis(), which is numeric and requires valid and up-to-date matrix values
481:       PetscCallHIPSPARSE(hipsparseSpSV_analysis(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_U, fs->spsvBuffer_U));
482:       PetscCallHIPSPARSE(hipsparseSpSV_analysis(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Ut, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Ut, fs->spsvBuffer_Ut));
483:       fs->updatedSpSVAnalysis = PETSC_TRUE;
484:     }
485:   }
486:   PetscFunctionReturn(PETSC_SUCCESS);
487: }

489: // Solve Ut D U x = b
490: static PetscErrorCode MatSolve_SeqAIJHIPSPARSE_Cholesky(Mat A, Vec b, Vec x)
491: {
492:   Mat_SeqAIJHIPSPARSETriFactors        *fs  = static_cast<Mat_SeqAIJHIPSPARSETriFactors *>(A->spptr);
493:   Mat_SeqAIJ                           *aij = static_cast<Mat_SeqAIJ *>(A->data);
494:   const PetscScalar                    *barray;
495:   PetscScalar                          *xarray;
496:   thrust::device_ptr<const PetscScalar> bGPU;
497:   thrust::device_ptr<PetscScalar>       xGPU;
498:   const hipsparseSpSVAlg_t              alg = HIPSPARSE_SPSV_ALG_DEFAULT;
499:   PetscInt                              m   = A->rmap->n;

501:   PetscFunctionBegin;
502:   PetscCall(PetscLogGpuTimeBegin());
503:   PetscCall(VecHIPGetArrayWrite(x, &xarray));
504:   PetscCall(VecHIPGetArrayRead(b, &barray));
505:   xGPU = thrust::device_pointer_cast(xarray);
506:   bGPU = thrust::device_pointer_cast(barray);

508:   // Reorder b with the row permutation if needed, and wrap the result in fs->X
509:   if (fs->rpermIndices)
510:     PetscCallThrust(thrust::copy(thrust::device.on(PetscDefaultHipStream), thrust::make_permutation_iterator(bGPU, fs->rpermIndices->begin()), thrust::make_permutation_iterator(bGPU, fs->rpermIndices->end()), thrust::device_pointer_cast(fs->X)));

512:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, fs->rpermIndices ? fs->X : (void *)barray));

514:   // Solve Ut Y = X
515:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_Y, fs->Y));
516: #if PETSC_PKG_HIP_VERSION_EQ(5, 6, 0) || PETSC_PKG_HIP_VERSION_GE(6, 0, 0)
517:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Ut, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, alg, fs->spsvDescr_Ut));
518: #else
519:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Ut, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, alg, fs->spsvDescr_Ut, fs->spsvBuffer_Ut));
520: #endif

522:   // Solve diag(D) Z = Y. Actually just do Y = Y*D since D is already inverted in MatCholeskyFactorNumeric_SeqAIJ().
523:   // It is basically a vector element-wise multiplication, but cublas does not have it!
524:   auto multiplies = thrust::multiplies<PetscScalar>();
525:   PetscCallThrust(thrust::transform(thrust::device.on(PetscDefaultHipStream), 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));

527:   // Solve U X = Y
528:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, fs->cpermIndices ? fs->X : xarray)); // if need to permute, we need to use the intermediate buffer X

530: #if PETSC_PKG_HIP_VERSION_EQ(5, 6, 0) || PETSC_PKG_HIP_VERSION_GE(6, 0, 0)
531:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_Y, fs->dnVecDescr_X, hipsparse_scalartype, alg, fs->spsvDescr_U));
532: #else
533:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_Y, fs->dnVecDescr_X, hipsparse_scalartype, alg, fs->spsvDescr_U, fs->spsvBuffer_U));
534: #endif

536:   // Reorder X with the column permutation if needed, and put the result back to x
537:   if (fs->cpermIndices)
538:     PetscCallThrust(
539:       thrust::copy(thrust::device.on(PetscDefaultHipStream), thrust::make_permutation_iterator(thrust::device_pointer_cast(fs->X), fs->cpermIndices->begin()), thrust::make_permutation_iterator(thrust::device_pointer_cast(fs->X + m), fs->cpermIndices->end()), xGPU));

541:   PetscCall(VecHIPRestoreArrayRead(b, &barray));
542:   PetscCall(VecHIPRestoreArrayWrite(x, &xarray));
543:   PetscCall(PetscLogGpuTimeEnd());
544:   PetscCall(PetscLogGpuFlops(4.0 * aij->nz - A->rmap->n));
545:   PetscFunctionReturn(PETSC_SUCCESS);
546: }

548: static PetscErrorCode MatSeqAIJHIPSPARSEICCAnalysisAndCopyToGPU(Mat A)
549: {
550:   Mat_SeqAIJ                    *a                   = (Mat_SeqAIJ *)A->data;
551:   Mat_SeqAIJHIPSPARSETriFactors *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;
552:   IS                             ip                  = a->row;
553:   PetscBool                      perm_identity;
554:   PetscInt                       n = A->rmap->n;

556:   PetscFunctionBegin;
557:   PetscCheck(hipsparseTriFactors, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing hipsparseTriFactors");

559:   PetscCall(MatSeqAIJHIPSPARSEBuildFactoredMatrix_Cholesky(A));
560:   hipsparseTriFactors->nnz = (a->nz - n) * 2 + n;

562:   A->offloadmask = PETSC_OFFLOAD_BOTH;

564:   /* lower triangular indices */
565:   PetscCall(ISIdentity(ip, &perm_identity));
566:   if (!perm_identity) {
567:     IS              iip;
568:     const PetscInt *irip, *rip;

570:     PetscCall(ISInvertPermutation(ip, PETSC_DECIDE, &iip));
571:     PetscCall(ISGetIndices(iip, &irip));
572:     PetscCall(ISGetIndices(ip, &rip));
573:     hipsparseTriFactors->rpermIndices = new THRUSTINTARRAY(n);
574:     hipsparseTriFactors->rpermIndices->assign(rip, rip + n);
575:     hipsparseTriFactors->cpermIndices = new THRUSTINTARRAY(n);
576:     hipsparseTriFactors->cpermIndices->assign(irip, irip + n);
577:     PetscCall(ISRestoreIndices(iip, &irip));
578:     PetscCall(ISDestroy(&iip));
579:     PetscCall(ISRestoreIndices(ip, &rip));
580:     PetscCall(PetscLogCpuToGpu(2. * n * sizeof(PetscInt)));
581:   }
582:   PetscFunctionReturn(PETSC_SUCCESS);
583: }

585: static PetscErrorCode MatCholeskyFactorNumeric_SeqAIJHIPSPARSE(Mat B, Mat A, const MatFactorInfo *info)
586: {
587:   PetscFunctionBegin;
588:   PetscCall(MatSeqAIJHIPSPARSECopyFromGPU(A));
589:   PetscCall(MatCholeskyFactorNumeric_SeqAIJ(B, A, info));
590:   B->offloadmask            = PETSC_OFFLOAD_CPU;
591:   B->ops->solve             = MatSolve_SeqAIJHIPSPARSE_Cholesky;
592:   B->ops->solvetranspose    = MatSolve_SeqAIJHIPSPARSE_Cholesky; // since symmetric
593:   B->ops->matsolve          = NULL;
594:   B->ops->matsolvetranspose = NULL;
595:   /* get the triangular factors */
596:   PetscCall(MatSeqAIJHIPSPARSEICCAnalysisAndCopyToGPU(B));
597:   PetscFunctionReturn(PETSC_SUCCESS);
598: }

600: static PetscErrorCode MatSeqAIJHIPSPARSEFormExplicitTranspose(Mat A)
601: {
602:   Mat_SeqAIJHIPSPARSE           *hipsparsestruct = (Mat_SeqAIJHIPSPARSE *)A->spptr;
603:   Mat_SeqAIJHIPSPARSEMultStruct *matstruct, *matstructT;
604:   Mat_SeqAIJ                    *a = (Mat_SeqAIJ *)A->data;
605:   hipsparseIndexBase_t           indexBase;

607:   PetscFunctionBegin;
608:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
609:   matstruct = (Mat_SeqAIJHIPSPARSEMultStruct *)hipsparsestruct->mat;
610:   PetscCheck(matstruct, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing mat struct");
611:   matstructT = (Mat_SeqAIJHIPSPARSEMultStruct *)hipsparsestruct->matTranspose;
612:   PetscCheck(!A->transupdated || matstructT, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing matTranspose struct");
613:   if (A->transupdated) PetscFunctionReturn(PETSC_SUCCESS);
614:   PetscCall(PetscLogEventBegin(MAT_HIPSPARSEGenerateTranspose, A, 0, 0, 0));
615:   PetscCall(PetscLogGpuTimeBegin());
616:   if (hipsparsestruct->format != MAT_HIPSPARSE_CSR) PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(A, PETSC_TRUE));
617:   if (!hipsparsestruct->matTranspose) { /* create hipsparse matrix */
618:     matstructT = new Mat_SeqAIJHIPSPARSEMultStruct;
619:     PetscCallHIPSPARSE(hipsparseCreateMatDescr(&matstructT->descr));
620:     indexBase = hipsparseGetMatIndexBase(matstruct->descr);
621:     PetscCallHIPSPARSE(hipsparseSetMatIndexBase(matstructT->descr, indexBase));
622:     PetscCallHIPSPARSE(hipsparseSetMatType(matstructT->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));

624:     /* set alpha and beta */
625:     PetscCallHIP(hipMalloc((void **)&matstructT->alpha_one, sizeof(PetscScalar)));
626:     PetscCallHIP(hipMalloc((void **)&matstructT->beta_zero, sizeof(PetscScalar)));
627:     PetscCallHIP(hipMalloc((void **)&matstructT->beta_one, sizeof(PetscScalar)));
628:     PetscCallHIP(hipMemcpy(matstructT->alpha_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
629:     PetscCallHIP(hipMemcpy(matstructT->beta_zero, &PETSC_HIPSPARSE_ZERO, sizeof(PetscScalar), hipMemcpyHostToDevice));
630:     PetscCallHIP(hipMemcpy(matstructT->beta_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));

632:     if (hipsparsestruct->format == MAT_HIPSPARSE_CSR) {
633:       CsrMatrix *matrixT      = new CsrMatrix;
634:       matstructT->mat         = matrixT;
635:       matrixT->num_rows       = A->cmap->n;
636:       matrixT->num_cols       = A->rmap->n;
637:       matrixT->num_entries    = a->nz;
638:       matrixT->row_offsets    = new THRUSTINTARRAY(matrixT->num_rows + 1);
639:       matrixT->column_indices = new THRUSTINTARRAY(a->nz);
640:       matrixT->values         = new THRUSTARRAY(a->nz);

642:       if (!hipsparsestruct->rowoffsets_gpu) hipsparsestruct->rowoffsets_gpu = new THRUSTINTARRAY(A->rmap->n + 1);
643:       hipsparsestruct->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
644:       PetscCallHIPSPARSE(hipsparseCreateCsr(&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, hipsparse_scalartype));
645:     } else if (hipsparsestruct->format == MAT_HIPSPARSE_ELL || hipsparsestruct->format == MAT_HIPSPARSE_HYB) {
646:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MAT_HIPSPARSE_ELL and MAT_HIPSPARSE_HYB are not supported");
647:     }
648:   }
649:   if (hipsparsestruct->format == MAT_HIPSPARSE_CSR) { /* transpose mat struct may be already present, update data */
650:     CsrMatrix *matrix  = (CsrMatrix *)matstruct->mat;
651:     CsrMatrix *matrixT = (CsrMatrix *)matstructT->mat;
652:     PetscCheck(matrix, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrix");
653:     PetscCheck(matrix->row_offsets, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrix rows");
654:     PetscCheck(matrix->column_indices, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrix cols");
655:     PetscCheck(matrix->values, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrix values");
656:     PetscCheck(matrixT, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrixT");
657:     PetscCheck(matrixT->row_offsets, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrixT rows");
658:     PetscCheck(matrixT->column_indices, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrixT cols");
659:     PetscCheck(matrixT->values, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrixT values");
660:     if (!hipsparsestruct->rowoffsets_gpu) { /* this may be absent when we did not construct the transpose with csr2csc */
661:       hipsparsestruct->rowoffsets_gpu = new THRUSTINTARRAY(A->rmap->n + 1);
662:       hipsparsestruct->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
663:       PetscCall(PetscLogCpuToGpu((A->rmap->n + 1) * sizeof(PetscInt)));
664:     }
665:     if (!hipsparsestruct->csr2csc_i) { // not using hipsparseCsr2cscEx2() because it requires 32-bit indices
666:       THRUSTINTARRAY row_indices(matrix->num_entries);

668:       // Transpose the matrix via COO, i.e., by putting the row indices in column_indices[] and the column indices in row_indices[]
669:       hipsparsestruct->csr2csc_i = new THRUSTINTARRAY(matrix->num_entries); // will store the matrix to matrixT permutation, i.e., entry matrixT[i] is matrix[csr2csc_i[i]]
670:       PetscCallThrust(thrust::sequence(thrust::device, hipsparsestruct->csr2csc_i->begin(), hipsparsestruct->csr2csc_i->end()));
671:       PetscCallThrust(thrust::for_each(thrust::device, thrust::counting_iterator<PetscInt>(0), thrust::counting_iterator<PetscInt>(A->rmap->n), Csr2coo(hipsparsestruct->rowoffsets_gpu->data().get(), matrixT->column_indices->data().get())));
672:       row_indices = *matrix->column_indices;
673:       // Sort the COO by row then column, and get the permutation csr2csc_i[]
674:       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())),
675:                                           hipsparsestruct->csr2csc_i->begin()));
676:       // Finalize matrixT's row_offsets by looking up row_indices[]
677:       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()));
678:     }
679:     PetscCallThrust(thrust::gather(thrust::device, hipsparsestruct->csr2csc_i->begin(), hipsparsestruct->csr2csc_i->end(), matrix->values->begin(), matrixT->values->begin()));
680:   }
681:   PetscCall(PetscLogGpuTimeEnd());
682:   PetscCall(PetscLogEventEnd(MAT_HIPSPARSEGenerateTranspose, A, 0, 0, 0));
683:   /* the compressed row indices is not used for matTranspose */
684:   matstructT->cprowIndices = NULL;
685:   /* assign the pointer */
686:   ((Mat_SeqAIJHIPSPARSE *)A->spptr)->matTranspose = matstructT;
687:   A->transupdated                                 = PETSC_TRUE;
688:   PetscFunctionReturn(PETSC_SUCCESS);
689: }

691: static PetscErrorCode MatSolve_SeqAIJHIPSPARSE_LU(Mat A, Vec b, Vec x)
692: {
693:   const PetscScalar                    *barray;
694:   PetscScalar                          *xarray;
695:   thrust::device_ptr<const PetscScalar> bGPU;
696:   thrust::device_ptr<PetscScalar>       xGPU;
697:   Mat_SeqAIJHIPSPARSETriFactors        *fs  = static_cast<Mat_SeqAIJHIPSPARSETriFactors *>(A->spptr);
698:   const Mat_SeqAIJ                     *aij = static_cast<Mat_SeqAIJ *>(A->data);
699:   const hipsparseOperation_t            op  = HIPSPARSE_OPERATION_NON_TRANSPOSE;
700:   const hipsparseSpSVAlg_t              alg = HIPSPARSE_SPSV_ALG_DEFAULT;
701:   PetscInt                              m   = A->rmap->n;

703:   PetscFunctionBegin;
704:   PetscCall(PetscLogGpuTimeBegin());
705:   PetscCall(VecHIPGetArrayWrite(x, &xarray));
706:   PetscCall(VecHIPGetArrayRead(b, &barray));
707:   xGPU = thrust::device_pointer_cast(xarray);
708:   bGPU = thrust::device_pointer_cast(barray);

710:   // Reorder b with the row permutation if needed, and wrap the result in fs->X
711:   if (fs->rpermIndices)
712:     PetscCallThrust(thrust::copy(thrust::device.on(PetscDefaultHipStream), thrust::make_permutation_iterator(bGPU, fs->rpermIndices->begin()), thrust::make_permutation_iterator(bGPU, fs->rpermIndices->end()), thrust::device_pointer_cast(fs->X)));

714:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, fs->rpermIndices ? fs->X : (void *)barray));

716:   // Solve L Y = X
717:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_Y, fs->Y));
718:   // Note that hipsparseSpSV_solve() secretly uses the external buffer used in hipsparseSpSV_analysis()!

720: #if PETSC_PKG_HIP_VERSION_EQ(5, 6, 0) || PETSC_PKG_HIP_VERSION_GE(6, 0, 0)
721:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, op, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, alg, fs->spsvDescr_L));
722: #else
723:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, op, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, alg, fs->spsvDescr_L, fs->spsvBuffer_L));
724: #endif

726:   // Solve U X = Y
727:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, fs->cpermIndices ? fs->X : xarray));

729: #if PETSC_PKG_HIP_VERSION_EQ(5, 6, 0) || PETSC_PKG_HIP_VERSION_GE(6, 0, 0)
730:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, op, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_Y, fs->dnVecDescr_X, hipsparse_scalartype, alg, fs->spsvDescr_U));
731: #else
732:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, op, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_Y, fs->dnVecDescr_X, hipsparse_scalartype, alg, fs->spsvDescr_U, fs->spsvBuffer_U));
733: #endif

735:   // Reorder X with the column permutation if needed, and put the result back to x
736:   if (fs->cpermIndices)
737:     PetscCallThrust(
738:       thrust::copy(thrust::device.on(PetscDefaultHipStream), thrust::make_permutation_iterator(thrust::device_pointer_cast(fs->X), fs->cpermIndices->begin()), thrust::make_permutation_iterator(thrust::device_pointer_cast(fs->X + m), fs->cpermIndices->end()), xGPU));
739:   PetscCall(VecHIPRestoreArrayRead(b, &barray));
740:   PetscCall(VecHIPRestoreArrayWrite(x, &xarray));
741:   PetscCall(PetscLogGpuTimeEnd());
742:   PetscCall(PetscLogGpuFlops(2.0 * aij->nz - m));
743:   PetscFunctionReturn(PETSC_SUCCESS);
744: }

746: static PetscErrorCode MatSolveTranspose_SeqAIJHIPSPARSE_LU(Mat A, Vec b, Vec x)
747: {
748:   Mat_SeqAIJHIPSPARSETriFactors        *fs  = static_cast<Mat_SeqAIJHIPSPARSETriFactors *>(A->spptr);
749:   Mat_SeqAIJ                           *aij = static_cast<Mat_SeqAIJ *>(A->data);
750:   const PetscScalar                    *barray;
751:   PetscScalar                          *xarray;
752:   thrust::device_ptr<const PetscScalar> bGPU;
753:   thrust::device_ptr<PetscScalar>       xGPU;
754:   const hipsparseOperation_t            opA = HIPSPARSE_OPERATION_TRANSPOSE;
755:   const hipsparseSpSVAlg_t              alg = HIPSPARSE_SPSV_ALG_DEFAULT;
756:   PetscInt                              m   = A->rmap->n;

758:   PetscFunctionBegin;
759:   PetscCall(PetscLogGpuTimeBegin());
760:   if (!fs->createdTransposeSpSVDescr) { // Call MatSolveTranspose() for the first time
761:     PetscCall(MatSeqAIJHIPSPARSECreateTransposeSolveMat(fs->spMatDescr_L, &fs->spMatDescr_Lt));
762:     PetscCall(MatSeqAIJHIPSPARSECreateTransposeSolveMat(fs->spMatDescr_U, &fs->spMatDescr_Ut));
763:     PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_Lt));
764:     PetscCallHIPSPARSE(hipsparseSpSV_bufferSize(fs->handle, opA, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Lt, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, alg, fs->spsvDescr_Lt, &fs->spsvBufferSize_Lt));

766:     PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_Ut));
767:     PetscCallHIPSPARSE(hipsparseSpSV_bufferSize(fs->handle, opA, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Ut, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, alg, fs->spsvDescr_Ut, &fs->spsvBufferSize_Ut));
768:     PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_Lt, fs->spsvBufferSize_Lt));
769:     PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_Ut, fs->spsvBufferSize_Ut));
770:     fs->createdTransposeSpSVDescr = PETSC_TRUE;
771:   }

773:   if (!fs->updatedTransposeSpSVAnalysis) {
774:     PetscCallHIPSPARSE(hipsparseSpSV_analysis(fs->handle, opA, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Lt, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, alg, fs->spsvDescr_Lt, fs->spsvBuffer_Lt));

776:     PetscCallHIPSPARSE(hipsparseSpSV_analysis(fs->handle, opA, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Ut, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, alg, fs->spsvDescr_Ut, fs->spsvBuffer_Ut));
777:     fs->updatedTransposeSpSVAnalysis = PETSC_TRUE;
778:   }

780:   PetscCall(VecHIPGetArrayWrite(x, &xarray));
781:   PetscCall(VecHIPGetArrayRead(b, &barray));
782:   xGPU = thrust::device_pointer_cast(xarray);
783:   bGPU = thrust::device_pointer_cast(barray);

785:   // Reorder b with the row permutation if needed, and wrap the result in fs->X
786:   if (fs->rpermIndices)
787:     PetscCallThrust(thrust::copy(thrust::device.on(PetscDefaultHipStream), thrust::make_permutation_iterator(bGPU, fs->rpermIndices->begin()), thrust::make_permutation_iterator(bGPU, fs->rpermIndices->end()), thrust::device_pointer_cast(fs->X)));

789:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, fs->rpermIndices ? fs->X : (void *)barray));

791:   // Solve Ut Y = X
792:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_Y, fs->Y));
793: #if PETSC_PKG_HIP_VERSION_EQ(5, 6, 0) || PETSC_PKG_HIP_VERSION_GE(6, 0, 0)
794:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, opA, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Ut, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, alg, fs->spsvDescr_Ut));
795: #else
796:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, opA, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Ut, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, alg, fs->spsvDescr_Ut, fs->spsvBuffer_Ut));
797: #endif

799:   // Solve Lt X = Y
800:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, fs->cpermIndices ? fs->X : xarray)); // if need to permute, we need to use the intermediate buffer X

802: #if PETSC_PKG_HIP_VERSION_EQ(5, 6, 0) || PETSC_PKG_HIP_VERSION_GE(6, 0, 0)
803:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, opA, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Lt, fs->dnVecDescr_Y, fs->dnVecDescr_X, hipsparse_scalartype, alg, fs->spsvDescr_Lt));
804: #else
805:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, opA, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Lt, fs->dnVecDescr_Y, fs->dnVecDescr_X, hipsparse_scalartype, alg, fs->spsvDescr_Lt, fs->spsvBuffer_Lt));
806: #endif

808:   // Reorder X with the column permutation if needed, and put the result back to x
809:   if (fs->cpermIndices)
810:     PetscCallThrust(
811:       thrust::copy(thrust::device.on(PetscDefaultHipStream), thrust::make_permutation_iterator(thrust::device_pointer_cast(fs->X), fs->cpermIndices->begin()), thrust::make_permutation_iterator(thrust::device_pointer_cast(fs->X + m), fs->cpermIndices->end()), xGPU));

813:   PetscCall(VecHIPRestoreArrayRead(b, &barray));
814:   PetscCall(VecHIPRestoreArrayWrite(x, &xarray));
815:   PetscCall(PetscLogGpuTimeEnd());
816:   PetscCall(PetscLogGpuFlops(2.0 * aij->nz - A->rmap->n));
817:   PetscFunctionReturn(PETSC_SUCCESS);
818: }

820: static PetscErrorCode MatILUFactorNumeric_SeqAIJHIPSPARSE_ILU0(Mat fact, Mat A, const MatFactorInfo *info)
821: {
822:   Mat_SeqAIJHIPSPARSETriFactors *fs    = (Mat_SeqAIJHIPSPARSETriFactors *)fact->spptr;
823:   Mat_SeqAIJ                    *aij   = (Mat_SeqAIJ *)fact->data;
824:   Mat_SeqAIJHIPSPARSE           *Acusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
825:   CsrMatrix                     *Acsr;
826:   PetscInt                       m, nz;
827:   PetscBool                      flg;

829:   PetscFunctionBegin;
830:   if (PetscDefined(USE_DEBUG)) {
831:     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
832:     PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Expected MATSEQAIJHIPSPARSE, but input is %s", ((PetscObject)A)->type_name);
833:   }

835:   /* Copy A's value to fact */
836:   m  = fact->rmap->n;
837:   nz = aij->nz;
838:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
839:   Acsr = (CsrMatrix *)Acusp->mat->mat;
840:   PetscCallHIP(hipMemcpyAsync(fs->csrVal, Acsr->values->data().get(), sizeof(PetscScalar) * nz, hipMemcpyDeviceToDevice, PetscDefaultHipStream));

842:   PetscCall(PetscLogGpuTimeBegin());
843:   /* Factorize fact inplace */
844:   if (m)
845:     PetscCallHIPSPARSE(hipsparseXcsrilu02(fs->handle, m, nz, /* hipsparseXcsrilu02 errors out with empty matrices (m=0) */
846:                                           fs->matDescr_M, fs->csrVal, fs->csrRowPtr32, fs->csrColIdx32, fs->ilu0Info_M, fs->policy_M, fs->factBuffer_M));
847:   if (PetscDefined(USE_DEBUG)) {
848:     int               numerical_zero;
849:     hipsparseStatus_t status;
850:     status = hipsparseXcsrilu02_zeroPivot(fs->handle, fs->ilu0Info_M, &numerical_zero);
851:     PetscAssert(HIPSPARSE_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);
852:   }

854:   {
855:     /* hipsparseSpSV_analysis() is numeric, i.e., it requires valid matrix values, therefore, we do it after hipsparseXcsrilu02()
856:      See discussion at https://github.com/NVIDIA/CUDALibrarySamples/issues/78
857:     */
858:     PetscCallHIPSPARSE(hipsparseSpSV_analysis(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, fs->spsvBuffer_L));

860:     PetscCallHIPSPARSE(hipsparseSpSV_analysis(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_U, fs->spsvBuffer_U));

862:     fs->updatedSpSVAnalysis = PETSC_TRUE;
863:     /* L, U values have changed, reset the flag to indicate we need to redo hipsparseSpSV_analysis() for transpose solve */
864:     fs->updatedTransposeSpSVAnalysis = PETSC_FALSE;
865:   }

867:   fact->offloadmask            = PETSC_OFFLOAD_GPU;
868:   fact->ops->solve             = MatSolve_SeqAIJHIPSPARSE_LU; // spMatDescr_L/U uses 32-bit indices, but hipsparseSpSV_solve() supports both 32 and 64. The info is encoded in hipsparseSpMatDescr_t.
869:   fact->ops->solvetranspose    = MatSolveTranspose_SeqAIJHIPSPARSE_LU;
870:   fact->ops->matsolve          = NULL;
871:   fact->ops->matsolvetranspose = NULL;
872:   PetscCall(PetscLogGpuTimeEnd());
873:   PetscCall(PetscLogGpuFlops(fs->numericFactFlops));
874:   PetscFunctionReturn(PETSC_SUCCESS);
875: }

877: static PetscErrorCode MatILUFactorSymbolic_SeqAIJHIPSPARSE_ILU0(Mat fact, Mat A, IS isrow, IS iscol, const MatFactorInfo *info)
878: {
879:   Mat_SeqAIJHIPSPARSETriFactors *fs  = (Mat_SeqAIJHIPSPARSETriFactors *)fact->spptr;
880:   Mat_SeqAIJ                    *aij = (Mat_SeqAIJ *)fact->data;
881:   PetscInt                       m, nz;

883:   PetscFunctionBegin;
884:   if (PetscDefined(USE_DEBUG)) {
885:     PetscBool flg, diagDense;

887:     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
888:     PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Expected MATSEQAIJHIPSPARSE, but input is %s", ((PetscObject)A)->type_name);
889:     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);
890:     PetscCall(MatGetDiagonalMarkers_SeqAIJ(A, NULL, &diagDense));
891:     PetscCheck(diagDense, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix is missing a diagonal entry");
892:   }

894:   /* Free the old stale stuff */
895:   PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(&fs));

897:   /* Copy over A's meta data to fact. Note that we also allocated fact's i,j,a on host,
898:      but they will not be used. Allocate them just for easy debugging.
899:    */
900:   PetscCall(MatDuplicateNoCreate_SeqAIJ(fact, A, MAT_DO_NOT_COPY_VALUES, PETSC_TRUE /*malloc*/));

902:   fact->offloadmask            = PETSC_OFFLOAD_BOTH;
903:   fact->factortype             = MAT_FACTOR_ILU;
904:   fact->info.factor_mallocs    = 0;
905:   fact->info.fill_ratio_given  = info->fill;
906:   fact->info.fill_ratio_needed = 1.0;

908:   aij->row = NULL;
909:   aij->col = NULL;

911:   /* ====================================================================== */
912:   /* Copy A's i, j to fact and also allocate the value array of fact.       */
913:   /* We'll do in-place factorization on fact                                */
914:   /* ====================================================================== */
915:   const PetscInt *Ai, *Aj;

917:   m  = fact->rmap->n;
918:   nz = aij->nz;

920:   PetscCallHIP(hipMalloc((void **)&fs->csrRowPtr32, sizeof(*fs->csrRowPtr32) * (m + 1)));
921:   PetscCallHIP(hipMalloc((void **)&fs->csrColIdx32, sizeof(*fs->csrColIdx32) * nz));
922:   PetscCallHIP(hipMalloc((void **)&fs->csrVal, sizeof(*fs->csrVal) * nz));
923:   PetscCall(MatSeqAIJHIPSPARSEGetIJ(A, PETSC_FALSE, &Ai, &Aj)); // Ai is uncompressed

925:   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);
926:   PetscCallThrust(thrust::transform(thrust::device.on(PetscDefaultHipStream), Ai, Ai + m + 1, fs->csrRowPtr32, PetscIntToCInt()));
927:   PetscCallThrust(thrust::transform(thrust::device.on(PetscDefaultHipStream), Aj, Aj + nz, fs->csrColIdx32, PetscIntToCInt()));

929:   /* ====================================================================== */
930:   /* Create descriptors for M, L, U                                         */
931:   /* ====================================================================== */
932:   hipsparseFillMode_t fillMode;
933:   hipsparseDiagType_t diagType;

935:   PetscCallHIPSPARSE(hipsparseCreateMatDescr(&fs->matDescr_M));
936:   PetscCallHIPSPARSE(hipsparseSetMatIndexBase(fs->matDescr_M, HIPSPARSE_INDEX_BASE_ZERO));
937:   PetscCallHIPSPARSE(hipsparseSetMatType(fs->matDescr_M, HIPSPARSE_MATRIX_TYPE_GENERAL));

939:   /* https://docs.nvidia.com/cuda/cusparse/index.html#cusparseDiagType_t
940:     cusparseDiagType_t: This type indicates if the matrix diagonal entries are unity. The diagonal elements are always
941:     assumed to be present, but if CUSPARSE_DIAG_TYPE_UNIT is passed to an API routine, then the routine assumes that
942:     all diagonal entries are unity and will not read or modify those entries. Note that in this case the routine
943:     assumes the diagonal entries are equal to one, regardless of what those entries are actually set to in memory.
944:   */
945:   fillMode = HIPSPARSE_FILL_MODE_LOWER;
946:   diagType = HIPSPARSE_DIAG_TYPE_UNIT;
947:   PetscCallHIPSPARSE(hipsparseCreateCsr(&fs->spMatDescr_L, m, m, nz, fs->csrRowPtr32, fs->csrColIdx32, fs->csrVal, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
948:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_L, HIPSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
949:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_L, HIPSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));

951:   fillMode = HIPSPARSE_FILL_MODE_UPPER;
952:   diagType = HIPSPARSE_DIAG_TYPE_NON_UNIT;
953:   PetscCallHIPSPARSE(hipsparseCreateCsr(&fs->spMatDescr_U, m, m, nz, fs->csrRowPtr32, fs->csrColIdx32, fs->csrVal, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
954:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_U, HIPSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
955:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_U, HIPSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));

957:   /* ========================================================================= */
958:   /* Query buffer sizes for csrilu0, SpSV and allocate buffers                 */
959:   /* ========================================================================= */
960:   PetscCallHIPSPARSE(hipsparseCreateCsrilu02Info(&fs->ilu0Info_M));
961:   if (m)
962:     PetscCallHIPSPARSE(hipsparseXcsrilu02_bufferSize(fs->handle, m, nz, /* hipsparseXcsrilu02 errors out with empty matrices (m=0) */
963:                                                      fs->matDescr_M, fs->csrVal, fs->csrRowPtr32, fs->csrColIdx32, fs->ilu0Info_M, &fs->factBufferSize_M));

965:   PetscCallHIP(hipMalloc((void **)&fs->X, sizeof(PetscScalar) * m));
966:   PetscCallHIP(hipMalloc((void **)&fs->Y, sizeof(PetscScalar) * m));

968:   PetscCallHIPSPARSE(hipsparseCreateDnVec(&fs->dnVecDescr_X, m, fs->X, hipsparse_scalartype));
969:   PetscCallHIPSPARSE(hipsparseCreateDnVec(&fs->dnVecDescr_Y, m, fs->Y, hipsparse_scalartype));

971:   PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_L));
972:   PetscCallHIPSPARSE(hipsparseSpSV_bufferSize(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, &fs->spsvBufferSize_L));

974:   PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_U));
975:   PetscCallHIPSPARSE(hipsparseSpSV_bufferSize(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_U, &fs->spsvBufferSize_U));

977:   /* From my experiment with the example at https://github.com/NVIDIA/CUDALibrarySamples/tree/master/cuSPARSE/bicgstab,
978:      and discussion at https://github.com/NVIDIA/CUDALibrarySamples/issues/77,
979:      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.
980:      To save memory, we make factBuffer_M share with the bigger of spsvBuffer_L/U.
981:    */
982:   if (fs->spsvBufferSize_L > fs->spsvBufferSize_U) {
983:     PetscCallHIP(hipMalloc((void **)&fs->factBuffer_M, PetscMax(fs->spsvBufferSize_L, (size_t)fs->factBufferSize_M)));
984:     fs->spsvBuffer_L = fs->factBuffer_M;
985:     PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_U, fs->spsvBufferSize_U));
986:   } else {
987:     PetscCallHIP(hipMalloc((void **)&fs->factBuffer_M, PetscMax(fs->spsvBufferSize_U, (size_t)fs->factBufferSize_M)));
988:     fs->spsvBuffer_U = fs->factBuffer_M;
989:     PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_L, fs->spsvBufferSize_L));
990:   }

992:   /* ========================================================================== */
993:   /* Perform analysis of ilu0 on M, SpSv on L and U                             */
994:   /* The lower(upper) triangular part of M has the same sparsity pattern as L(U)*/
995:   /* ========================================================================== */
996:   int               structural_zero;
997:   hipsparseStatus_t status;

999:   fs->policy_M = HIPSPARSE_SOLVE_POLICY_USE_LEVEL;
1000:   if (m)
1001:     PetscCallHIPSPARSE(hipsparseXcsrilu02_analysis(fs->handle, m, nz, /* hipsparseXcsrilu02 errors out with empty matrices (m=0) */
1002:                                                    fs->matDescr_M, fs->csrVal, fs->csrRowPtr32, fs->csrColIdx32, fs->ilu0Info_M, fs->policy_M, fs->factBuffer_M));
1003:   if (PetscDefined(USE_DEBUG)) {
1004:     /* hipsparseXcsrilu02_zeroPivot() is a blocking call. It calls hipDeviceSynchronize() to make sure all previous kernels are done. */
1005:     status = hipsparseXcsrilu02_zeroPivot(fs->handle, fs->ilu0Info_M, &structural_zero);
1006:     PetscCheck(HIPSPARSE_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);
1007:   }

1009:   /* Estimate FLOPs of the numeric factorization */
1010:   {
1011:     Mat_SeqAIJ     *Aseq = (Mat_SeqAIJ *)A->data;
1012:     PetscInt       *Ai, nzRow, nzLeft;
1013:     const PetscInt *adiag;
1014:     PetscLogDouble  flops = 0.0;

1016:     PetscCall(MatGetDiagonalMarkers_SeqAIJ(A, &adiag, NULL));
1017:     Ai = Aseq->i;
1018:     for (PetscInt i = 0; i < m; i++) {
1019:       if (Ai[i] < adiag[i] && adiag[i] < Ai[i + 1]) { /* There are nonzeros left to the diagonal of row i */
1020:         nzRow  = Ai[i + 1] - Ai[i];
1021:         nzLeft = adiag[i] - Ai[i];
1022:         /* We want to eliminate nonzeros left to the diagonal one by one. Assume each time, nonzeros right
1023:           and include the eliminated one will be updated, which incurs a multiplication and an addition.
1024:         */
1025:         nzLeft = (nzRow - 1) / 2;
1026:         flops += nzLeft * (2.0 * nzRow - nzLeft + 1);
1027:       }
1028:     }
1029:     fs->numericFactFlops = flops;
1030:   }
1031:   fact->ops->lufactornumeric = MatILUFactorNumeric_SeqAIJHIPSPARSE_ILU0;
1032:   PetscFunctionReturn(PETSC_SUCCESS);
1033: }

1035: static PetscErrorCode MatSolve_SeqAIJHIPSPARSE_ICC0(Mat fact, Vec b, Vec x)
1036: {
1037:   Mat_SeqAIJHIPSPARSETriFactors *fs  = (Mat_SeqAIJHIPSPARSETriFactors *)fact->spptr;
1038:   Mat_SeqAIJ                    *aij = (Mat_SeqAIJ *)fact->data;
1039:   const PetscScalar             *barray;
1040:   PetscScalar                   *xarray;

1042:   PetscFunctionBegin;
1043:   PetscCall(VecHIPGetArrayWrite(x, &xarray));
1044:   PetscCall(VecHIPGetArrayRead(b, &barray));
1045:   PetscCall(PetscLogGpuTimeBegin());

1047:   /* Solve L*y = b */
1048:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, (void *)barray));
1049:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_Y, fs->Y));
1050: #if PETSC_PKG_HIP_VERSION_EQ(5, 6, 0) || PETSC_PKG_HIP_VERSION_GE(6, 0, 0)
1051:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, /* L Y = X */
1052:                                          fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L));
1053: #else
1054:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, /* L Y = X */
1055:                                          fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, fs->spsvBuffer_L));
1056: #endif
1057:   /* Solve Lt*x = y */
1058:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, xarray));
1059: #if PETSC_PKG_HIP_VERSION_EQ(5, 6, 0) || PETSC_PKG_HIP_VERSION_GE(6, 0, 0)
1060:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Lt, /* Lt X = Y */
1061:                                          fs->dnVecDescr_Y, fs->dnVecDescr_X, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Lt));
1062: #else
1063:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Lt, /* Lt X = Y */
1064:                                          fs->dnVecDescr_Y, fs->dnVecDescr_X, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Lt, fs->spsvBuffer_Lt));
1065: #endif
1066:   PetscCall(VecHIPRestoreArrayRead(b, &barray));
1067:   PetscCall(VecHIPRestoreArrayWrite(x, &xarray));

1069:   PetscCall(PetscLogGpuTimeEnd());
1070:   PetscCall(PetscLogGpuFlops(2.0 * aij->nz - fact->rmap->n));
1071:   PetscFunctionReturn(PETSC_SUCCESS);
1072: }

1074: static PetscErrorCode MatICCFactorNumeric_SeqAIJHIPSPARSE_ICC0(Mat fact, Mat A, const MatFactorInfo *info)
1075: {
1076:   Mat_SeqAIJHIPSPARSETriFactors *fs    = (Mat_SeqAIJHIPSPARSETriFactors *)fact->spptr;
1077:   Mat_SeqAIJ                    *aij   = (Mat_SeqAIJ *)fact->data;
1078:   Mat_SeqAIJHIPSPARSE           *Acusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
1079:   CsrMatrix                     *Acsr;
1080:   PetscInt                       m, nz;
1081:   PetscBool                      flg;

1083:   PetscFunctionBegin;
1084:   if (PetscDefined(USE_DEBUG)) {
1085:     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
1086:     PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Expected MATSEQAIJHIPSPARSE, but input is %s", ((PetscObject)A)->type_name);
1087:   }

1089:   /* Copy A's value to fact */
1090:   m  = fact->rmap->n;
1091:   nz = aij->nz;
1092:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
1093:   Acsr = (CsrMatrix *)Acusp->mat->mat;
1094:   PetscCallHIP(hipMemcpyAsync(fs->csrVal, Acsr->values->data().get(), sizeof(PetscScalar) * nz, hipMemcpyDeviceToDevice, PetscDefaultHipStream));

1096:   /* Factorize fact inplace */
1097:   /* https://docs.nvidia.com/cuda/cusparse/index.html#csric02_solve
1098:      csric02() only takes the lower triangular part of matrix A to perform factorization.
1099:      The matrix type must be CUSPARSE_MATRIX_TYPE_GENERAL, the fill mode and diagonal type are ignored,
1100:      and the strictly upper triangular part is ignored and never touched. It does not matter if A is Hermitian or not.
1101:      In other words, from the point of view of csric02() A is Hermitian and only the lower triangular part is provided.
1102:    */
1103:   if (m) PetscCallHIPSPARSE(hipsparseXcsric02(fs->handle, m, nz, fs->matDescr_M, fs->csrVal, fs->csrRowPtr32, fs->csrColIdx32, fs->ic0Info_M, fs->policy_M, fs->factBuffer_M));
1104:   if (PetscDefined(USE_DEBUG)) {
1105:     int               numerical_zero;
1106:     hipsparseStatus_t status;
1107:     status = hipsparseXcsric02_zeroPivot(fs->handle, fs->ic0Info_M, &numerical_zero);
1108:     PetscAssert(HIPSPARSE_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);
1109:   }

1111:   {
1112:     PetscCallHIPSPARSE(hipsparseSpSV_analysis(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, fs->spsvBuffer_L));

1114:     /* Note that cusparse reports this error if we use double and CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE
1115:     ** On entry to cusparseSpSV_analysis(): conjugate transpose (opA) is not supported for matA data type, current -> CUDA_R_64F
1116:   */
1117:     PetscCallHIPSPARSE(hipsparseSpSV_analysis(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Lt, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Lt, fs->spsvBuffer_Lt));
1118:     fs->updatedSpSVAnalysis = PETSC_TRUE;
1119:   }

1121:   fact->offloadmask            = PETSC_OFFLOAD_GPU;
1122:   fact->ops->solve             = MatSolve_SeqAIJHIPSPARSE_ICC0;
1123:   fact->ops->solvetranspose    = MatSolve_SeqAIJHIPSPARSE_ICC0;
1124:   fact->ops->matsolve          = NULL;
1125:   fact->ops->matsolvetranspose = NULL;
1126:   PetscCall(PetscLogGpuFlops(fs->numericFactFlops));
1127:   PetscFunctionReturn(PETSC_SUCCESS);
1128: }

1130: static PetscErrorCode MatICCFactorSymbolic_SeqAIJHIPSPARSE_ICC0(Mat fact, Mat A, IS perm, const MatFactorInfo *info)
1131: {
1132:   Mat_SeqAIJHIPSPARSETriFactors *fs  = (Mat_SeqAIJHIPSPARSETriFactors *)fact->spptr;
1133:   Mat_SeqAIJ                    *aij = (Mat_SeqAIJ *)fact->data;
1134:   PetscInt                       m, nz;

1136:   PetscFunctionBegin;
1137:   if (PetscDefined(USE_DEBUG)) {
1138:     PetscBool flg, diagDense;

1140:     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
1141:     PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Expected MATSEQAIJHIPSPARSE, but input is %s", ((PetscObject)A)->type_name);
1142:     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);
1143:     PetscCall(MatGetDiagonalMarkers_SeqAIJ(A, NULL, &diagDense));
1144:     PetscCheck(diagDense, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix is missing diagonal entries");
1145:   }

1147:   /* Free the old stale stuff */
1148:   PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(&fs));

1150:   /* Copy over A's meta data to fact. Note that we also allocated fact's i,j,a on host,
1151:      but they will not be used. Allocate them just for easy debugging.
1152:    */
1153:   PetscCall(MatDuplicateNoCreate_SeqAIJ(fact, A, MAT_DO_NOT_COPY_VALUES, PETSC_TRUE /*malloc*/));

1155:   fact->offloadmask            = PETSC_OFFLOAD_BOTH;
1156:   fact->factortype             = MAT_FACTOR_ICC;
1157:   fact->info.factor_mallocs    = 0;
1158:   fact->info.fill_ratio_given  = info->fill;
1159:   fact->info.fill_ratio_needed = 1.0;

1161:   aij->row = NULL;
1162:   aij->col = NULL;

1164:   /* ====================================================================== */
1165:   /* Copy A's i, j to fact and also allocate the value array of fact.       */
1166:   /* We'll do in-place factorization on fact                                */
1167:   /* ====================================================================== */
1168:   const PetscInt *Ai, *Aj;

1170:   m  = fact->rmap->n;
1171:   nz = aij->nz;

1173:   PetscCallHIP(hipMalloc((void **)&fs->csrRowPtr32, sizeof(*fs->csrRowPtr32) * (m + 1)));
1174:   PetscCallHIP(hipMalloc((void **)&fs->csrColIdx32, sizeof(*fs->csrColIdx32) * nz));
1175:   PetscCallHIP(hipMalloc((void **)&fs->csrVal, sizeof(PetscScalar) * nz));
1176:   PetscCall(MatSeqAIJHIPSPARSEGetIJ(A, PETSC_FALSE, &Ai, &Aj)); // Ai is uncompressed

1178:   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);
1179:   PetscCallThrust(thrust::transform(thrust::device.on(PetscDefaultHipStream), Ai, Ai + m + 1, fs->csrRowPtr32, PetscIntToCInt()));
1180:   PetscCallThrust(thrust::transform(thrust::device.on(PetscDefaultHipStream), Aj, Aj + nz, fs->csrColIdx32, PetscIntToCInt()));

1182:   /* ====================================================================== */
1183:   /* Create mat descriptors for M, L                                        */
1184:   /* ====================================================================== */
1185:   hipsparseFillMode_t fillMode;
1186:   hipsparseDiagType_t diagType;

1188:   PetscCallHIPSPARSE(hipsparseCreateMatDescr(&fs->matDescr_M));
1189:   PetscCallHIPSPARSE(hipsparseSetMatIndexBase(fs->matDescr_M, HIPSPARSE_INDEX_BASE_ZERO));
1190:   PetscCallHIPSPARSE(hipsparseSetMatType(fs->matDescr_M, HIPSPARSE_MATRIX_TYPE_GENERAL));

1192:   /* https://docs.nvidia.com/cuda/cusparse/index.html#cusparseDiagType_t
1193:     cusparseDiagType_t: This type indicates if the matrix diagonal entries are unity. The diagonal elements are always
1194:     assumed to be present, but if CUSPARSE_DIAG_TYPE_UNIT is passed to an API routine, then the routine assumes that
1195:     all diagonal entries are unity and will not read or modify those entries. Note that in this case the routine
1196:     assumes the diagonal entries are equal to one, regardless of what those entries are actually set to in memory.
1197:   */
1198:   fillMode = HIPSPARSE_FILL_MODE_LOWER;
1199:   diagType = HIPSPARSE_DIAG_TYPE_NON_UNIT;
1200:   PetscCallHIPSPARSE(hipsparseCreateCsr(&fs->spMatDescr_L, m, m, nz, fs->csrRowPtr32, fs->csrColIdx32, fs->csrVal, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
1201:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_L, HIPSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
1202:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_L, HIPSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));

1204:   /* ========================================================================= */
1205:   /* Query buffer sizes for csric0, SpSV of L and Lt, and allocate buffers     */
1206:   /* ========================================================================= */
1207:   PetscCallHIPSPARSE(hipsparseCreateCsric02Info(&fs->ic0Info_M));
1208:   if (m) PetscCallHIPSPARSE(hipsparseXcsric02_bufferSize(fs->handle, m, nz, fs->matDescr_M, fs->csrVal, fs->csrRowPtr32, fs->csrColIdx32, fs->ic0Info_M, &fs->factBufferSize_M));

1210:   PetscCallHIP(hipMalloc((void **)&fs->X, sizeof(PetscScalar) * m));
1211:   PetscCallHIP(hipMalloc((void **)&fs->Y, sizeof(PetscScalar) * m));

1213:   PetscCallHIPSPARSE(hipsparseCreateDnVec(&fs->dnVecDescr_X, m, fs->X, hipsparse_scalartype));
1214:   PetscCallHIPSPARSE(hipsparseCreateDnVec(&fs->dnVecDescr_Y, m, fs->Y, hipsparse_scalartype));

1216:   PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_L));
1217:   PetscCallHIPSPARSE(hipsparseSpSV_bufferSize(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, &fs->spsvBufferSize_L));

1219:   PetscCall(MatSeqAIJHIPSPARSECreateTransposeSolveMat(fs->spMatDescr_L, &fs->spMatDescr_Lt));
1220:   PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_Lt));
1221:   PetscCallHIPSPARSE(hipsparseSpSV_bufferSize(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_Lt, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Lt, &fs->spsvBufferSize_Lt));

1223:   /* To save device memory, we make the factorization buffer share with one of the solver buffer.
1224:      See also comments in MatILUFactorSymbolic_SeqAIJCUSPARSE_ILU0().
1225:    */
1226:   if (fs->spsvBufferSize_L > fs->spsvBufferSize_Lt) {
1227:     PetscCallHIP(hipMalloc((void **)&fs->factBuffer_M, PetscMax(fs->spsvBufferSize_L, (size_t)fs->factBufferSize_M)));
1228:     fs->spsvBuffer_L = fs->factBuffer_M;
1229:     PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_Lt, fs->spsvBufferSize_Lt));
1230:   } else {
1231:     PetscCallHIP(hipMalloc((void **)&fs->factBuffer_M, PetscMax(fs->spsvBufferSize_Lt, (size_t)fs->factBufferSize_M)));
1232:     fs->spsvBuffer_Lt = fs->factBuffer_M;
1233:     PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_L, fs->spsvBufferSize_L));
1234:   }

1236:   /* ========================================================================== */
1237:   /* Perform analysis of ic0 on M                                               */
1238:   /* The lower triangular part of M has the same sparsity pattern as L          */
1239:   /* ========================================================================== */
1240:   int               structural_zero;
1241:   hipsparseStatus_t status;

1243:   fs->policy_M = HIPSPARSE_SOLVE_POLICY_USE_LEVEL;
1244:   if (m) PetscCallHIPSPARSE(hipsparseXcsric02_analysis(fs->handle, m, nz, fs->matDescr_M, fs->csrVal, fs->csrRowPtr32, fs->csrColIdx32, fs->ic0Info_M, fs->policy_M, fs->factBuffer_M));
1245:   if (PetscDefined(USE_DEBUG)) {
1246:     /* hipsparseXcsric02_zeroPivot() is a blocking call. It calls cudaDeviceSynchronize() to make sure all previous kernels are done. */
1247:     status = hipsparseXcsric02_zeroPivot(fs->handle, fs->ic0Info_M, &structural_zero);
1248:     PetscCheck(HIPSPARSE_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);
1249:   }

1251:   /* Estimate FLOPs of the numeric factorization */
1252:   {
1253:     Mat_SeqAIJ    *Aseq = (Mat_SeqAIJ *)A->data;
1254:     PetscInt      *Ai, nzRow, nzLeft;
1255:     PetscLogDouble flops = 0.0;

1257:     Ai = Aseq->i;
1258:     for (PetscInt i = 0; i < m; i++) {
1259:       nzRow = Ai[i + 1] - Ai[i];
1260:       if (nzRow > 1) {
1261:         /* We want to eliminate nonzeros left to the diagonal one by one. Assume each time, nonzeros right
1262:           and include the eliminated one will be updated, which incurs a multiplication and an addition.
1263:         */
1264:         nzLeft = (nzRow - 1) / 2;
1265:         flops += nzLeft * (2.0 * nzRow - nzLeft + 1);
1266:       }
1267:     }
1268:     fs->numericFactFlops = flops;
1269:   }
1270:   fact->ops->choleskyfactornumeric = MatICCFactorNumeric_SeqAIJHIPSPARSE_ICC0;
1271:   PetscFunctionReturn(PETSC_SUCCESS);
1272: }

1274: static PetscErrorCode MatLUFactorNumeric_SeqAIJHIPSPARSE(Mat B, Mat A, const MatFactorInfo *info)
1275: {
1276:   // use_cpu_solve is a field in Mat_SeqAIJHIPSPARSE. B, a factored matrix, uses Mat_SeqAIJHIPSPARSETriFactors.
1277:   Mat_SeqAIJHIPSPARSE *hipsparsestruct = static_cast<Mat_SeqAIJHIPSPARSE *>(A->spptr);

1279:   PetscFunctionBegin;
1280:   PetscCall(MatSeqAIJHIPSPARSECopyFromGPU(A));
1281:   PetscCall(MatLUFactorNumeric_SeqAIJ(B, A, info));
1282:   B->offloadmask = PETSC_OFFLOAD_CPU;

1284:   if (!hipsparsestruct->use_cpu_solve) {
1285:     B->ops->solve          = MatSolve_SeqAIJHIPSPARSE_LU;
1286:     B->ops->solvetranspose = MatSolveTranspose_SeqAIJHIPSPARSE_LU;
1287:   }
1288:   B->ops->matsolve          = NULL;
1289:   B->ops->matsolvetranspose = NULL;

1291:   /* get the triangular factors */
1292:   if (!hipsparsestruct->use_cpu_solve) PetscCall(MatSeqAIJHIPSPARSEILUAnalysisAndCopyToGPU(B));
1293:   PetscFunctionReturn(PETSC_SUCCESS);
1294: }

1296: static PetscErrorCode MatLUFactorSymbolic_SeqAIJHIPSPARSE(Mat B, Mat A, IS isrow, IS iscol, const MatFactorInfo *info)
1297: {
1298:   Mat_SeqAIJHIPSPARSETriFactors *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)B->spptr;

1300:   PetscFunctionBegin;
1301:   PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(&hipsparseTriFactors));
1302:   PetscCall(MatLUFactorSymbolic_SeqAIJ(B, A, isrow, iscol, info));
1303:   B->ops->lufactornumeric = MatLUFactorNumeric_SeqAIJHIPSPARSE;
1304:   PetscFunctionReturn(PETSC_SUCCESS);
1305: }

1307: static PetscErrorCode MatILUFactorSymbolic_SeqAIJHIPSPARSE(Mat B, Mat A, IS isrow, IS iscol, const MatFactorInfo *info)
1308: {
1309:   Mat_SeqAIJHIPSPARSETriFactors *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)B->spptr;

1311:   PetscFunctionBegin;
1312:   PetscBool row_identity = PETSC_FALSE, col_identity = PETSC_FALSE;
1313:   if (!info->factoronhost) {
1314:     PetscCall(ISIdentity(isrow, &row_identity));
1315:     PetscCall(ISIdentity(iscol, &col_identity));
1316:   }
1317:   if (!info->levels && row_identity && col_identity) PetscCall(MatILUFactorSymbolic_SeqAIJHIPSPARSE_ILU0(B, A, isrow, iscol, info));
1318:   else {
1319:     PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(&hipsparseTriFactors));
1320:     PetscCall(MatILUFactorSymbolic_SeqAIJ(B, A, isrow, iscol, info));
1321:     B->ops->lufactornumeric = MatLUFactorNumeric_SeqAIJHIPSPARSE;
1322:   }
1323:   PetscFunctionReturn(PETSC_SUCCESS);
1324: }

1326: static PetscErrorCode MatICCFactorSymbolic_SeqAIJHIPSPARSE(Mat B, Mat A, IS perm, const MatFactorInfo *info)
1327: {
1328:   Mat_SeqAIJHIPSPARSETriFactors *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)B->spptr;

1330:   PetscFunctionBegin;
1331:   PetscBool perm_identity = PETSC_FALSE;
1332:   if (!info->factoronhost) PetscCall(ISIdentity(perm, &perm_identity));
1333:   if (!info->levels && perm_identity) PetscCall(MatICCFactorSymbolic_SeqAIJHIPSPARSE_ICC0(B, A, perm, info));
1334:   else {
1335:     PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(&hipsparseTriFactors));
1336:     PetscCall(MatICCFactorSymbolic_SeqAIJ(B, A, perm, info));
1337:     B->ops->choleskyfactornumeric = MatCholeskyFactorNumeric_SeqAIJHIPSPARSE;
1338:   }
1339:   PetscFunctionReturn(PETSC_SUCCESS);
1340: }

1342: static PetscErrorCode MatCholeskyFactorSymbolic_SeqAIJHIPSPARSE(Mat B, Mat A, IS perm, const MatFactorInfo *info)
1343: {
1344:   Mat_SeqAIJHIPSPARSETriFactors *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)B->spptr;

1346:   PetscFunctionBegin;
1347:   PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(&hipsparseTriFactors));
1348:   PetscCall(MatCholeskyFactorSymbolic_SeqAIJ(B, A, perm, info));
1349:   B->ops->choleskyfactornumeric = MatCholeskyFactorNumeric_SeqAIJHIPSPARSE;
1350:   PetscFunctionReturn(PETSC_SUCCESS);
1351: }

1353: static PetscErrorCode MatFactorGetSolverType_seqaij_hipsparse(Mat A, MatSolverType *type)
1354: {
1355:   PetscFunctionBegin;
1356:   *type = MATSOLVERHIPSPARSE;
1357:   PetscFunctionReturn(PETSC_SUCCESS);
1358: }

1360: /*MC
1361:   MATSOLVERHIPSPARSE = "hipsparse" - A matrix type providing triangular solvers for sequential matrices
1362:   on a single GPU of type, `MATSEQAIJHIPSPARSE`. Currently supported
1363:   algorithms are ILU(k) and ICC(k). Typically, deeper factorizations (larger k) results in poorer
1364:   performance in the triangular solves. Full LU, and Cholesky decompositions can be solved through the
1365:   HipSPARSE triangular solve algorithm. However, the performance can be quite poor and thus these
1366:   algorithms are not recommended. This class does NOT support direct solver operations.

1368:   Level: beginner

1370: .seealso: [](ch_matrices), `Mat`, `MATSEQAIJHIPSPARSE`, `PCFactorSetMatSolverType()`, `MatSolverType`, `MatCreateSeqAIJHIPSPARSE()`, `MATAIJHIPSPARSE`, `MatCreateAIJHIPSPARSE()`, `MatHIPSPARSESetFormat()`, `MatHIPSPARSEStorageFormat`, `MatHIPSPARSEFormatOperation`
1371: M*/

1373: PETSC_EXTERN PetscErrorCode MatGetFactor_seqaijhipsparse_hipsparse(Mat A, MatFactorType ftype, Mat *B)
1374: {
1375:   PetscInt n = A->rmap->n;

1377:   PetscFunctionBegin;
1378:   PetscCall(MatCreate(PetscObjectComm((PetscObject)A), B));
1379:   PetscCall(MatSetSizes(*B, n, n, n, n));
1380:   (*B)->factortype = ftype; // factortype makes MatSetType() allocate spptr of type Mat_SeqAIJHIPSPARSETriFactors
1381:   PetscCall(MatSetType(*B, MATSEQAIJHIPSPARSE));

1383:   if (A->boundtocpu && A->bindingpropagates) PetscCall(MatBindToCPU(*B, PETSC_TRUE));
1384:   if (ftype == MAT_FACTOR_LU || ftype == MAT_FACTOR_ILU || ftype == MAT_FACTOR_ILUDT) {
1385:     PetscCall(MatSetBlockSizesFromMats(*B, A, A));
1386:     if (!A->boundtocpu) {
1387:       (*B)->ops->ilufactorsymbolic = MatILUFactorSymbolic_SeqAIJHIPSPARSE;
1388:       (*B)->ops->lufactorsymbolic  = MatLUFactorSymbolic_SeqAIJHIPSPARSE;
1389:     } else {
1390:       (*B)->ops->ilufactorsymbolic = MatILUFactorSymbolic_SeqAIJ;
1391:       (*B)->ops->lufactorsymbolic  = MatLUFactorSymbolic_SeqAIJ;
1392:     }
1393:     PetscCall(PetscStrallocpy(MATORDERINGND, (char **)&(*B)->preferredordering[MAT_FACTOR_LU]));
1394:     PetscCall(PetscStrallocpy(MATORDERINGNATURAL, (char **)&(*B)->preferredordering[MAT_FACTOR_ILU]));
1395:     PetscCall(PetscStrallocpy(MATORDERINGNATURAL, (char **)&(*B)->preferredordering[MAT_FACTOR_ILUDT]));
1396:   } else if (ftype == MAT_FACTOR_CHOLESKY || ftype == MAT_FACTOR_ICC) {
1397:     if (!A->boundtocpu) {
1398:       (*B)->ops->iccfactorsymbolic      = MatICCFactorSymbolic_SeqAIJHIPSPARSE;
1399:       (*B)->ops->choleskyfactorsymbolic = MatCholeskyFactorSymbolic_SeqAIJHIPSPARSE;
1400:     } else {
1401:       (*B)->ops->iccfactorsymbolic      = MatICCFactorSymbolic_SeqAIJ;
1402:       (*B)->ops->choleskyfactorsymbolic = MatCholeskyFactorSymbolic_SeqAIJ;
1403:     }
1404:     PetscCall(PetscStrallocpy(MATORDERINGND, (char **)&(*B)->preferredordering[MAT_FACTOR_CHOLESKY]));
1405:     PetscCall(PetscStrallocpy(MATORDERINGNATURAL, (char **)&(*B)->preferredordering[MAT_FACTOR_ICC]));
1406:   } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "Factor type not supported for HIPSPARSE Matrix Types");

1408:   PetscCall(MatSeqAIJSetPreallocation(*B, MAT_SKIP_ALLOCATION, NULL));
1409:   (*B)->canuseordering = PETSC_TRUE;
1410:   PetscCall(PetscObjectComposeFunction((PetscObject)*B, "MatFactorGetSolverType_C", MatFactorGetSolverType_seqaij_hipsparse));
1411:   PetscFunctionReturn(PETSC_SUCCESS);
1412: }

1414: static PetscErrorCode MatSeqAIJHIPSPARSECopyFromGPU(Mat A)
1415: {
1416:   Mat_SeqAIJ                    *a    = (Mat_SeqAIJ *)A->data;
1417:   Mat_SeqAIJHIPSPARSE           *cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
1418:   Mat_SeqAIJHIPSPARSETriFactors *fs   = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;

1420:   PetscFunctionBegin;
1421:   if (A->offloadmask == PETSC_OFFLOAD_GPU) {
1422:     PetscCall(PetscLogEventBegin(MAT_HIPSPARSECopyFromGPU, A, 0, 0, 0));
1423:     if (A->factortype == MAT_FACTOR_NONE) {
1424:       CsrMatrix *matrix = (CsrMatrix *)cusp->mat->mat;
1425:       PetscCallHIP(hipMemcpy(a->a, matrix->values->data().get(), a->nz * sizeof(PetscScalar), hipMemcpyDeviceToHost));
1426:     } else if (fs->csrVal) {
1427:       /* We have a factorized matrix on device and are able to copy it to host */
1428:       PetscCallHIP(hipMemcpy(a->a, fs->csrVal, a->nz * sizeof(PetscScalar), hipMemcpyDeviceToHost));
1429:     } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "No support for copying this type of factorized matrix from device to host");
1430:     PetscCall(PetscLogGpuToCpu(a->nz * sizeof(PetscScalar)));
1431:     PetscCall(PetscLogEventEnd(MAT_HIPSPARSECopyFromGPU, A, 0, 0, 0));
1432:     A->offloadmask = PETSC_OFFLOAD_BOTH;
1433:   }
1434:   PetscFunctionReturn(PETSC_SUCCESS);
1435: }

1437: /* Policy struct for MatSeqAIJCUSPARSE_CUPM shared template (HIP specialisation) */
1438: struct MatSeqAIJHIPSPARSE_Policy {
1439:   typedef Mat_SeqAIJHIPSPARSE           mat_struct_type;
1440:   typedef Mat_SeqAIJHIPSPARSEMultStruct mult_struct_type;

1442:   static int storage_format_csr() { return (int)MAT_HIPSPARSE_CSR; }
1443:   static int storage_format_ell() { return (int)MAT_HIPSPARSE_ELL; }
1444:   static int storage_format_hyb() { return (int)MAT_HIPSPARSE_HYB; }

1446:   static PetscErrorCode CopyToGPU(Mat A) { return MatSeqAIJHIPSPARSECopyToGPU(A); }
1447:   static PetscErrorCode CopyFromGPU(Mat A) { return MatSeqAIJHIPSPARSECopyFromGPU(A); }
1448:   static PetscErrorCode InvalidateTranspose(Mat A, PetscBool d) { return MatSeqAIJHIPSPARSEInvalidateTranspose(A, d); }
1449:   static PetscErrorCode ConvertFromSeqAIJ(Mat B, MatType t, MatReuse r, Mat *C) { return MatConvert_SeqAIJ_SeqAIJHIPSPARSE(B, t, r, C); }
1450:   static const char    *mat_type_name;

1452:   static PetscErrorCode Destroy(Mat A) { return MatSeqAIJHIPSPARSE_Destroy(A); }
1453:   static PetscErrorCode TriFactorsDestroy(void **spptr) { return MatSeqAIJHIPSPARSETriFactors_Destroy((Mat_SeqAIJHIPSPARSETriFactors **)spptr); }
1454:   static const char    *set_format_c;
1455:   static const char    *set_use_cpu_solve_c;
1456:   static const char    *product_seqdense_device_c;
1457:   static const char    *product_seqdense_c;
1458:   static const char    *product_self_c;
1459:   static const char    *seq_convert_hypre_c;

1461:   static PetscErrorCode VecGetArrayRead(Vec v, const PetscScalar **a) { return VecHIPGetArrayRead(v, a); }
1462:   static PetscErrorCode VecRestoreArrayRead(Vec v, const PetscScalar **a) { return VecHIPRestoreArrayRead(v, a); }
1463:   static PetscErrorCode VecGetArrayWrite(Vec v, PetscScalar **a) { return VecHIPGetArrayWrite(v, a); }
1464:   static PetscErrorCode VecRestoreArrayWrite(Vec v, PetscScalar **a) { return VecHIPRestoreArrayWrite(v, a); }
1465: };
1466: const char *MatSeqAIJHIPSPARSE_Policy::mat_type_name             = MATSEQAIJHIPSPARSE;
1467: const char *MatSeqAIJHIPSPARSE_Policy::set_format_c              = "MatHIPSPARSESetFormat_C";
1468: const char *MatSeqAIJHIPSPARSE_Policy::set_use_cpu_solve_c       = "MatHIPSPARSESetUseCPUSolve_C";
1469: const char *MatSeqAIJHIPSPARSE_Policy::product_seqdense_device_c = "MatProductSetFromOptions_seqaijhipsparse_seqdensehip_C";
1470: const char *MatSeqAIJHIPSPARSE_Policy::product_seqdense_c        = "MatProductSetFromOptions_seqaijhipsparse_seqdense_C";
1471: const char *MatSeqAIJHIPSPARSE_Policy::product_self_c            = "MatProductSetFromOptions_seqaijhipsparse_seqaijhipsparse_C";
1472: const char *MatSeqAIJHIPSPARSE_Policy::seq_convert_hypre_c       = "MatConvert_seqaijhipsparse_hypre_C";

1474: using MatSeqAIJHIPSPARSE_CUPM_t = Petsc::mat::aij::cupm::impl::MatSeqAIJCUSPARSE_CUPM<Petsc::device::cupm::DeviceType::HIP, MatSeqAIJHIPSPARSE_Policy>;

1476: static PetscErrorCode MatSeqAIJGetArray_SeqAIJHIPSPARSE(Mat A, PetscScalar *array[])
1477: {
1478:   return MatSeqAIJHIPSPARSE_CUPM_t::SeqAIJGetArray(A, array);
1479: }

1481: static PetscErrorCode MatSeqAIJRestoreArray_SeqAIJHIPSPARSE(Mat A, PetscScalar *array[])
1482: {
1483:   return MatSeqAIJHIPSPARSE_CUPM_t::SeqAIJRestoreArray(A, array);
1484: }

1486: static PetscErrorCode MatSeqAIJGetArrayRead_SeqAIJHIPSPARSE(Mat A, const PetscScalar *array[])
1487: {
1488:   return MatSeqAIJHIPSPARSE_CUPM_t::SeqAIJGetArrayRead(A, array);
1489: }

1491: static PetscErrorCode MatSeqAIJRestoreArrayRead_SeqAIJHIPSPARSE(Mat A, const PetscScalar *array[])
1492: {
1493:   return MatSeqAIJHIPSPARSE_CUPM_t::SeqAIJRestoreArrayRead(A, array);
1494: }

1496: static PetscErrorCode MatSeqAIJGetArrayWrite_SeqAIJHIPSPARSE(Mat A, PetscScalar *array[])
1497: {
1498:   return MatSeqAIJHIPSPARSE_CUPM_t::SeqAIJGetArrayWrite(A, array);
1499: }

1501: static PetscErrorCode MatSeqAIJRestoreArrayWrite_SeqAIJHIPSPARSE(Mat A, PetscScalar *array[])
1502: {
1503:   return MatSeqAIJHIPSPARSE_CUPM_t::SeqAIJRestoreArrayWrite(A, array);
1504: }

1506: static PetscErrorCode MatSeqAIJGetCSRAndMemType_SeqAIJHIPSPARSE(Mat A, const PetscInt **i, const PetscInt **j, PetscScalar **a, PetscMemType *mtype)
1507: {
1508:   Mat_SeqAIJHIPSPARSE *cusp;
1509:   CsrMatrix           *matrix;

1511:   PetscFunctionBegin;
1512:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
1513:   PetscCheck(A->factortype == MAT_FACTOR_NONE, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Not for factored matrix");
1514:   cusp = static_cast<Mat_SeqAIJHIPSPARSE *>(A->spptr);
1515:   PetscCheck(cusp != NULL, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "cusp is NULL");
1516:   matrix = (CsrMatrix *)cusp->mat->mat;

1518:   if (i) *i = matrix->row_offsets->data().get();
1519:   if (j) *j = matrix->column_indices->data().get();
1520:   if (a) *a = matrix->values->data().get();
1521:   if (mtype) *mtype = PETSC_MEMTYPE_HIP;
1522:   PetscFunctionReturn(PETSC_SUCCESS);
1523: }

1525: PETSC_INTERN PetscErrorCode MatSeqAIJHIPSPARSECopyToGPU(Mat A)
1526: {
1527:   Mat_SeqAIJHIPSPARSE           *hipsparsestruct = (Mat_SeqAIJHIPSPARSE *)A->spptr;
1528:   Mat_SeqAIJHIPSPARSEMultStruct *matstruct       = hipsparsestruct->mat;
1529:   Mat_SeqAIJ                    *a               = (Mat_SeqAIJ *)A->data;
1530:   PetscInt                       m               = A->rmap->n, *ii, *ridx, tmp;
1531:   PetscBool                      both            = PETSC_TRUE;

1533:   PetscFunctionBegin;
1534:   PetscCheck(!A->boundtocpu, PETSC_COMM_SELF, PETSC_ERR_GPU, "Cannot copy to GPU");
1535:   if (A->offloadmask == PETSC_OFFLOAD_UNALLOCATED || A->offloadmask == PETSC_OFFLOAD_CPU) {
1536:     if (A->nonzerostate == hipsparsestruct->nonzerostate && hipsparsestruct->format == MAT_HIPSPARSE_CSR) { /* Copy values only */
1537:       CsrMatrix *matrix;
1538:       matrix = (CsrMatrix *)hipsparsestruct->mat->mat;

1540:       PetscCheck(!a->nz || a->a, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CSR values");
1541:       PetscCall(PetscLogEventBegin(MAT_HIPSPARSECopyToGPU, A, 0, 0, 0));
1542:       matrix->values->assign(a->a, a->a + a->nz);
1543:       PetscCallHIP(WaitForHIP());
1544:       PetscCall(PetscLogCpuToGpu(a->nz * sizeof(PetscScalar)));
1545:       PetscCall(PetscLogEventEnd(MAT_HIPSPARSECopyToGPU, A, 0, 0, 0));
1546:       PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(A, PETSC_FALSE));
1547:     } else {
1548:       PetscInt nnz;
1549:       PetscCall(PetscLogEventBegin(MAT_HIPSPARSECopyToGPU, A, 0, 0, 0));
1550:       PetscCall(MatSeqAIJHIPSPARSEMultStruct_Destroy(&hipsparsestruct->mat, hipsparsestruct->format));
1551:       PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(A, PETSC_TRUE));
1552:       delete hipsparsestruct->workVector;
1553:       delete hipsparsestruct->rowoffsets_gpu;
1554:       hipsparsestruct->workVector     = NULL;
1555:       hipsparsestruct->rowoffsets_gpu = NULL;
1556:       try {
1557:         if (a->compressedrow.use) {
1558:           m    = a->compressedrow.nrows;
1559:           ii   = a->compressedrow.i;
1560:           ridx = a->compressedrow.rindex;
1561:         } else {
1562:           m    = A->rmap->n;
1563:           ii   = a->i;
1564:           ridx = NULL;
1565:         }
1566:         PetscCheck(ii, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CSR row data");
1567:         if (!a->a) {
1568:           nnz  = ii[m];
1569:           both = PETSC_FALSE;
1570:         } else nnz = a->nz;
1571:         PetscCheck(!nnz || a->j, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CSR column data");

1573:         /* create cusparse matrix */
1574:         hipsparsestruct->nrows = m;
1575:         matstruct              = new Mat_SeqAIJHIPSPARSEMultStruct;
1576:         PetscCallHIPSPARSE(hipsparseCreateMatDescr(&matstruct->descr));
1577:         PetscCallHIPSPARSE(hipsparseSetMatIndexBase(matstruct->descr, HIPSPARSE_INDEX_BASE_ZERO));
1578:         PetscCallHIPSPARSE(hipsparseSetMatType(matstruct->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));

1580:         PetscCallHIP(hipMalloc((void **)&matstruct->alpha_one, sizeof(PetscScalar)));
1581:         PetscCallHIP(hipMalloc((void **)&matstruct->beta_zero, sizeof(PetscScalar)));
1582:         PetscCallHIP(hipMalloc((void **)&matstruct->beta_one, sizeof(PetscScalar)));
1583:         PetscCallHIP(hipMemcpy(matstruct->alpha_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
1584:         PetscCallHIP(hipMemcpy(matstruct->beta_zero, &PETSC_HIPSPARSE_ZERO, sizeof(PetscScalar), hipMemcpyHostToDevice));
1585:         PetscCallHIP(hipMemcpy(matstruct->beta_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
1586:         PetscCallHIPSPARSE(hipsparseSetPointerMode(hipsparsestruct->handle, HIPSPARSE_POINTER_MODE_DEVICE));

1588:         /* Build a hybrid/ellpack matrix if this option is chosen for the storage */
1589:         if (hipsparsestruct->format == MAT_HIPSPARSE_CSR) {
1590:           /* set the matrix */
1591:           CsrMatrix *mat   = new CsrMatrix;
1592:           mat->num_rows    = m;
1593:           mat->num_cols    = A->cmap->n;
1594:           mat->num_entries = nnz;
1595:           PetscCallCXX(mat->row_offsets = new THRUSTINTARRAY(m + 1));
1596:           mat->row_offsets->assign(ii, ii + m + 1);
1597:           PetscCallCXX(mat->column_indices = new THRUSTINTARRAY(nnz));
1598:           mat->column_indices->assign(a->j, a->j + nnz);

1600:           PetscCallCXX(mat->values = new THRUSTARRAY(nnz));
1601:           if (a->a) mat->values->assign(a->a, a->a + nnz);

1603:           /* assign the pointer */
1604:           matstruct->mat = mat;
1605:           if (mat->num_rows) { /* cusparse errors on empty matrices! */
1606:             PetscCallHIPSPARSE(hipsparseCreateCsr(&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, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
1607:           }
1608:         } else if (hipsparsestruct->format == MAT_HIPSPARSE_ELL || hipsparsestruct->format == MAT_HIPSPARSE_HYB) {
1609:           SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "MAT_HIPSPARSE_ELL and MAT_HIPSPARSE_HYB are not supported");
1610:         }

1612:         /* assign the compressed row indices */
1613:         if (a->compressedrow.use) {
1614:           PetscCallCXX(hipsparsestruct->workVector = new THRUSTARRAY(m));
1615:           PetscCallCXX(matstruct->cprowIndices = new THRUSTINTARRAY(m));
1616:           matstruct->cprowIndices->assign(ridx, ridx + m);
1617:           tmp = m;
1618:         } else {
1619:           hipsparsestruct->workVector = NULL;
1620:           matstruct->cprowIndices     = NULL;
1621:           tmp                         = 0;
1622:         }
1623:         PetscCall(PetscLogCpuToGpu(((m + 1) + (a->nz)) * sizeof(int) + tmp * sizeof(PetscInt) + (3 + (a->nz)) * sizeof(PetscScalar)));

1625:         /* assign the pointer */
1626:         hipsparsestruct->mat = matstruct;
1627:       } catch (char *ex) {
1628:         SETERRQ(PETSC_COMM_SELF, PETSC_ERR_LIB, "HIPSPARSE error: %s", ex);
1629:       }
1630:       PetscCallHIP(WaitForHIP());
1631:       PetscCall(PetscLogEventEnd(MAT_HIPSPARSECopyToGPU, A, 0, 0, 0));
1632:       hipsparsestruct->nonzerostate = A->nonzerostate;
1633:     }
1634:     if (both) A->offloadmask = PETSC_OFFLOAD_BOTH;
1635:   }
1636:   PetscFunctionReturn(PETSC_SUCCESS);
1637: }

1639: struct VecHIPPlusEquals {
1640:   template <typename Tuple>
1641:   __host__ __device__ void operator()(Tuple t)
1642:   {
1643:     thrust::get<1>(t) = thrust::get<1>(t) + thrust::get<0>(t);
1644:   }
1645: };

1647: struct VecHIPEquals {
1648:   template <typename Tuple>
1649:   __host__ __device__ void operator()(Tuple t)
1650:   {
1651:     thrust::get<1>(t) = thrust::get<0>(t);
1652:   }
1653: };

1655: struct VecHIPEqualsReverse {
1656:   template <typename Tuple>
1657:   __host__ __device__ void operator()(Tuple t)
1658:   {
1659:     thrust::get<0>(t) = thrust::get<1>(t);
1660:   }
1661: };

1663: struct MatProductCtx_MatMatHipsparse {
1664:   PetscBool             cisdense;
1665:   PetscScalar          *Bt;
1666:   Mat                   X;
1667:   PetscBool             reusesym; /* Hipsparse does not have split symbolic and numeric phases for sparse matmat operations */
1668:   PetscLogDouble        flops;
1669:   CsrMatrix            *Bcsr;
1670:   hipsparseSpMatDescr_t matSpBDescr;
1671:   PetscBool             initialized; /* C = alpha op(A) op(B) + beta C */
1672:   hipsparseDnMatDescr_t matBDescr;
1673:   hipsparseDnMatDescr_t matCDescr;
1674:   PetscInt              Blda, Clda; /* Record leading dimensions of B and C here to detect changes*/
1675: #if PETSC_PKG_HIP_VERSION_GE(5, 1, 0)
1676:   void *dBuffer4, *dBuffer5;
1677: #endif
1678:   size_t                 mmBufferSize;
1679:   void                  *mmBuffer, *mmBuffer2; /* SpGEMM WorkEstimation buffer */
1680:   hipsparseSpGEMMDescr_t spgemmDesc;
1681: };

1683: static PetscErrorCode MatProductCtxDestroy_MatMatHipsparse(PetscCtxRt data)
1684: {
1685:   MatProductCtx_MatMatHipsparse *mmdata = *(MatProductCtx_MatMatHipsparse **)data;

1687:   PetscFunctionBegin;
1688:   PetscCallHIP(hipFree(mmdata->Bt));
1689:   delete mmdata->Bcsr;
1690:   if (mmdata->matSpBDescr) PetscCallHIPSPARSE(hipsparseDestroySpMat(mmdata->matSpBDescr));
1691:   if (mmdata->matBDescr) PetscCallHIPSPARSE(hipsparseDestroyDnMat(mmdata->matBDescr));
1692:   if (mmdata->matCDescr) PetscCallHIPSPARSE(hipsparseDestroyDnMat(mmdata->matCDescr));
1693:   if (mmdata->spgemmDesc) PetscCallHIPSPARSE(hipsparseSpGEMM_destroyDescr(mmdata->spgemmDesc));
1694: #if PETSC_PKG_HIP_VERSION_GE(5, 1, 0)
1695:   PetscCallHIP(hipFree(mmdata->dBuffer4));
1696:   PetscCallHIP(hipFree(mmdata->dBuffer5));
1697: #endif
1698:   PetscCallHIP(hipFree(mmdata->mmBuffer));
1699:   PetscCallHIP(hipFree(mmdata->mmBuffer2));
1700:   PetscCall(MatDestroy(&mmdata->X));
1701:   PetscCall(PetscFree(*(void **)data));
1702:   PetscFunctionReturn(PETSC_SUCCESS);
1703: }

1705: static PetscErrorCode MatProductNumeric_SeqAIJHIPSPARSE_SeqDENSEHIP(Mat C)
1706: {
1707:   Mat_Product                   *product = C->product;
1708:   Mat                            A, B;
1709:   PetscInt                       m, n, blda, clda;
1710:   PetscBool                      flg, biship, compressed;
1711:   Mat_SeqAIJ                    *a;
1712:   Mat_SeqAIJHIPSPARSE           *cusp;
1713:   hipsparseOperation_t           opA;
1714:   const PetscScalar             *barray;
1715:   PetscScalar                   *carray;
1716:   MatProductCtx_MatMatHipsparse *mmdata;
1717:   Mat_SeqAIJHIPSPARSEMultStruct *mat;
1718:   CsrMatrix                     *csrmat;

1720:   PetscFunctionBegin;
1721:   MatCheckProduct(C, 1);
1722:   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data empty");
1723:   mmdata = (MatProductCtx_MatMatHipsparse *)product->data;
1724:   A      = product->A;
1725:   B      = product->B;
1726:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
1727:   PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
1728:   /* currently CopyToGpu does not copy if the matrix is bound to CPU
1729:      Instead of silently accepting the wrong answer, I prefer to raise the error */
1730:   PetscCheck(!A->boundtocpu, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Cannot bind to CPU a HIPSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
1731:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
1732:   a    = (Mat_SeqAIJ *)A->data;
1733:   cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
1734:   switch (product->type) {
1735:   case MATPRODUCT_AB:
1736:   case MATPRODUCT_PtAP:
1737:     mat = cusp->mat;
1738:     opA = HIPSPARSE_OPERATION_NON_TRANSPOSE;
1739:     m   = A->rmap->n;
1740:     n   = B->cmap->n;
1741:     break;
1742:   case MATPRODUCT_AtB:
1743:     if (!A->form_explicit_transpose) {
1744:       mat = cusp->mat;
1745:       opA = HIPSPARSE_OPERATION_TRANSPOSE;
1746:     } else {
1747:       PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(A));
1748:       mat = cusp->matTranspose;
1749:       opA = HIPSPARSE_OPERATION_NON_TRANSPOSE;
1750:     }
1751:     m = A->cmap->n;
1752:     n = B->cmap->n;
1753:     break;
1754:   case MATPRODUCT_ABt:
1755:   case MATPRODUCT_RARt:
1756:     mat = cusp->mat;
1757:     opA = HIPSPARSE_OPERATION_NON_TRANSPOSE;
1758:     m   = A->rmap->n;
1759:     n   = B->rmap->n;
1760:     break;
1761:   default:
1762:     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
1763:   }
1764:   PetscCheck(mat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing Mat_SeqAIJHIPSPARSEMultStruct");
1765:   csrmat = (CsrMatrix *)mat->mat;
1766:   /* when the rows of A are compressed on the device, csrmat holds only the nonempty rows of A, so the
1767:      SpMM descriptor below must be built with the full row offsets instead of those of csrmat */
1768:   compressed = (PetscBool)(mat->cprowIndices != NULL);
1769:   /* if the user passed a CPU matrix, copy the data to the GPU */
1770:   PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQDENSEHIP, &biship));
1771:   if (!biship) PetscCall(MatConvert(B, MATSEQDENSEHIP, MAT_INPLACE_MATRIX, &B));
1772:   PetscCall(MatDenseGetArrayReadAndMemType(B, &barray, nullptr));
1773:   PetscCall(MatDenseGetLDA(B, &blda));
1774:   if (product->type == MATPRODUCT_RARt || product->type == MATPRODUCT_PtAP) {
1775:     PetscCall(MatDenseGetArrayWriteAndMemType(mmdata->X, &carray, nullptr));
1776:     PetscCall(MatDenseGetLDA(mmdata->X, &clda));
1777:   } else {
1778:     PetscCall(MatDenseGetArrayWriteAndMemType(C, &carray, nullptr));
1779:     PetscCall(MatDenseGetLDA(C, &clda));
1780:   }

1782:   PetscCall(PetscLogGpuTimeBegin());
1783:   hipsparseOperation_t opB = (product->type == MATPRODUCT_ABt || product->type == MATPRODUCT_RARt) ? HIPSPARSE_OPERATION_TRANSPOSE : HIPSPARSE_OPERATION_NON_TRANSPOSE;
1784:   /* mat->matDescr is also used by the SpGEMM code, which relies on its compressed dimensions, so when A is
1785:      compressed the SpMM needs a descriptor of its own */
1786:   hipsparseSpMatDescr_t &matADescr = compressed ? mat->matDescr_SpMM[opA] : mat->matDescr;
1787:   /* (re)allocate mmBuffer if not initialized or LDAs are different */
1788:   if (!mmdata->initialized || mmdata->Blda != blda || mmdata->Clda != clda) {
1789:     size_t mmBufferSize;
1790:     if (mmdata->initialized && mmdata->Blda != blda) {
1791:       PetscCallHIPSPARSE(hipsparseDestroyDnMat(mmdata->matBDescr));
1792:       mmdata->matBDescr = NULL;
1793:     }
1794:     if (!mmdata->matBDescr) {
1795:       PetscCallHIPSPARSE(hipsparseCreateDnMat(&mmdata->matBDescr, B->rmap->n, B->cmap->n, blda, (void *)barray, hipsparse_scalartype, HIPSPARSE_ORDER_COL));
1796:       mmdata->Blda = blda;
1797:     }
1798:     if (mmdata->initialized && mmdata->Clda != clda) {
1799:       PetscCallHIPSPARSE(hipsparseDestroyDnMat(mmdata->matCDescr));
1800:       mmdata->matCDescr = NULL;
1801:     }
1802:     if (!mmdata->matCDescr) { /* matCDescr is for C or mmdata->X */
1803:       PetscCallHIPSPARSE(hipsparseCreateDnMat(&mmdata->matCDescr, m, n, clda, (void *)carray, hipsparse_scalartype, HIPSPARSE_ORDER_COL));
1804:       mmdata->Clda = clda;
1805:     }
1806:     if (!matADescr) {
1807:       if (compressed) {
1808:         if (!cusp->rowoffsets_gpu) { /* the full row offsets may be absent when we did not construct the transpose with csr2csc */
1809:           cusp->rowoffsets_gpu = new THRUSTINTARRAY(A->rmap->n + 1);
1810:           cusp->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
1811:           PetscCall(PetscLogCpuToGpu((A->rmap->n + 1) * sizeof(PetscInt)));
1812:         }
1813:         PetscCallHIPSPARSE(hipsparseCreateCsr(&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, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
1814:       } else {
1815:         PetscCallHIPSPARSE(hipsparseCreateCsr(&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, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
1816:       }
1817:     }
1818:     PetscCallHIPSPARSE(hipsparseSpMM_bufferSize(cusp->handle, opA, opB, mat->alpha_one, matADescr, mmdata->matBDescr, mat->beta_zero, mmdata->matCDescr, hipsparse_scalartype, cusp->spmmAlg, &mmBufferSize));
1819:     if ((mmdata->mmBuffer && mmdata->mmBufferSize < mmBufferSize) || !mmdata->mmBuffer) {
1820:       PetscCallHIP(hipFree(mmdata->mmBuffer));
1821:       PetscCallHIP(hipMalloc(&mmdata->mmBuffer, mmBufferSize));
1822:       mmdata->mmBufferSize = mmBufferSize;
1823:     }
1824:     mmdata->initialized = PETSC_TRUE;
1825:   } else {
1826:     /* to be safe, always update pointers of the mats */
1827:     PetscCallHIPSPARSE(hipsparseSpMatSetValues(matADescr, csrmat->values->data().get()));
1828:     PetscCallHIPSPARSE(hipsparseDnMatSetValues(mmdata->matBDescr, (void *)barray));
1829:     PetscCallHIPSPARSE(hipsparseDnMatSetValues(mmdata->matCDescr, (void *)carray));
1830:   }

1832:   /* do hipsparseSpMM, which supports transpose on B */
1833:   PetscCallHIPSPARSE(hipsparseSpMM(cusp->handle, opA, opB, mat->alpha_one, matADescr, mmdata->matBDescr, mat->beta_zero, mmdata->matCDescr, hipsparse_scalartype, cusp->spmmAlg, mmdata->mmBuffer));

1835:   PetscCall(PetscLogGpuTimeEnd());
1836:   PetscCall(PetscLogGpuFlops(n * 2.0 * csrmat->num_entries));
1837:   PetscCall(MatDenseRestoreArrayReadAndMemType(B, &barray));
1838:   if (product->type == MATPRODUCT_RARt) {
1839:     PetscCall(MatDenseRestoreArrayWriteAndMemType(mmdata->X, &carray));
1840:     PetscCall(MatMatMultNumeric_SeqDenseHIP_SeqDenseHIP_Internal(B, mmdata->X, C, PETSC_FALSE, PETSC_FALSE));
1841:   } else if (product->type == MATPRODUCT_PtAP) {
1842:     PetscCall(MatDenseRestoreArrayWriteAndMemType(mmdata->X, &carray));
1843:     PetscCall(MatMatMultNumeric_SeqDenseHIP_SeqDenseHIP_Internal(B, mmdata->X, C, PETSC_TRUE, PETSC_FALSE));
1844:   } else PetscCall(MatDenseRestoreArrayWriteAndMemType(C, &carray));
1845:   if (mmdata->cisdense) PetscCall(MatConvert(C, MATSEQDENSE, MAT_INPLACE_MATRIX, &C));
1846:   if (!biship) PetscCall(MatConvert(B, MATSEQDENSE, MAT_INPLACE_MATRIX, &B));
1847:   PetscFunctionReturn(PETSC_SUCCESS);
1848: }

1850: static PetscErrorCode MatProductSymbolic_SeqAIJHIPSPARSE_SeqDENSEHIP(Mat C)
1851: {
1852:   Mat_Product                   *product = C->product;
1853:   Mat                            A, B;
1854:   PetscInt                       m, n;
1855:   PetscBool                      cisdense, flg;
1856:   MatProductCtx_MatMatHipsparse *mmdata;
1857:   Mat_SeqAIJHIPSPARSE           *cusp;

1859:   PetscFunctionBegin;
1860:   MatCheckProduct(C, 1);
1861:   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data not empty");
1862:   A = product->A;
1863:   B = product->B;
1864:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
1865:   PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
1866:   cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
1867:   PetscCheck(cusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");
1868:   switch (product->type) {
1869:   case MATPRODUCT_AB:
1870:     m = A->rmap->n;
1871:     n = B->cmap->n;
1872:     break;
1873:   case MATPRODUCT_AtB:
1874:     m = A->cmap->n;
1875:     n = B->cmap->n;
1876:     break;
1877:   case MATPRODUCT_ABt:
1878:     m = A->rmap->n;
1879:     n = B->rmap->n;
1880:     break;
1881:   case MATPRODUCT_PtAP:
1882:     m = B->cmap->n;
1883:     n = B->cmap->n;
1884:     break;
1885:   case MATPRODUCT_RARt:
1886:     m = B->rmap->n;
1887:     n = B->rmap->n;
1888:     break;
1889:   default:
1890:     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
1891:   }
1892:   PetscCall(MatSetSizes(C, m, n, m, n));
1893:   /* if C is of type MATSEQDENSE (CPU), perform the operation on the GPU and then copy on the CPU */
1894:   PetscCall(PetscObjectTypeCompare((PetscObject)C, MATSEQDENSE, &cisdense));
1895:   PetscCall(MatSetType(C, MATSEQDENSEHIP));
1896:   /* keep the VecType of B, e.g. VECKOKKOS from MatCreateDenseFromVecType(), but only if B is a MATSEQDENSEHIP like C, since the VecType of a MATSEQDENSE B, e.g. VECSTANDARD, is not valid for C */
1897:   PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQDENSEHIP, &flg));
1898:   if (flg) PetscCall(MatSetVecType(C, B->defaultvectype));

1900:   /* product data */
1901:   PetscCall(PetscNew(&mmdata));
1902:   mmdata->cisdense = cisdense;
1903:   /* for these products we need intermediate storage */
1904:   if (product->type == MATPRODUCT_RARt || product->type == MATPRODUCT_PtAP) {
1905:     PetscCall(MatCreate(PetscObjectComm((PetscObject)C), &mmdata->X));
1906:     PetscCall(MatSetType(mmdata->X, MATSEQDENSEHIP));
1907:     /* do not preallocate, since the first call to MatDenseHIPGetArray will preallocate on the GPU for us */
1908:     if (product->type == MATPRODUCT_RARt) PetscCall(MatSetSizes(mmdata->X, A->rmap->n, B->rmap->n, A->rmap->n, B->rmap->n));
1909:     else PetscCall(MatSetSizes(mmdata->X, A->rmap->n, B->cmap->n, A->rmap->n, B->cmap->n));
1910:   }
1911:   C->product->data       = mmdata;
1912:   C->product->destroy    = MatProductCtxDestroy_MatMatHipsparse;
1913:   C->ops->productnumeric = MatProductNumeric_SeqAIJHIPSPARSE_SeqDENSEHIP;
1914:   PetscFunctionReturn(PETSC_SUCCESS);
1915: }

1917: static PetscErrorCode MatProductNumeric_SeqAIJHIPSPARSE_SeqAIJHIPSPARSE(Mat C)
1918: {
1919:   Mat_Product                   *product = C->product;
1920:   Mat                            A, B;
1921:   Mat_SeqAIJHIPSPARSE           *Acusp, *Bcusp, *Ccusp;
1922:   Mat_SeqAIJ                    *c = (Mat_SeqAIJ *)C->data;
1923:   Mat_SeqAIJHIPSPARSEMultStruct *Amat, *Bmat, *Cmat;
1924:   CsrMatrix                     *Acsr, *Bcsr, *Ccsr;
1925:   PetscBool                      flg;
1926:   MatProductType                 ptype;
1927:   MatProductCtx_MatMatHipsparse *mmdata;
1928:   hipsparseSpMatDescr_t          BmatSpDescr;
1929:   hipsparseOperation_t           opA = HIPSPARSE_OPERATION_NON_TRANSPOSE, opB = HIPSPARSE_OPERATION_NON_TRANSPOSE; /* hipSPARSE spgemm doesn't support transpose yet */

1931:   PetscFunctionBegin;
1932:   MatCheckProduct(C, 1);
1933:   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data empty");
1934:   PetscCall(PetscObjectTypeCompare((PetscObject)C, MATSEQAIJHIPSPARSE, &flg));
1935:   PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for C of type %s", ((PetscObject)C)->type_name);
1936:   mmdata = (MatProductCtx_MatMatHipsparse *)C->product->data;
1937:   A      = product->A;
1938:   B      = product->B;
1939:   if (mmdata->reusesym) { /* this happens when api_user is true, meaning that the matrix values have been already computed in the MatProductSymbolic phase */
1940:     mmdata->reusesym = PETSC_FALSE;
1941:     Ccusp            = (Mat_SeqAIJHIPSPARSE *)C->spptr;
1942:     PetscCheck(Ccusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");
1943:     Cmat = Ccusp->mat;
1944:     PetscCheck(Cmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C mult struct for product type %s", MatProductTypes[C->product->type]);
1945:     Ccsr = (CsrMatrix *)Cmat->mat;
1946:     PetscCheck(Ccsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C CSR struct");
1947:     goto finalize;
1948:   }
1949:   if (!c->nz) goto finalize;
1950:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
1951:   PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
1952:   PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQAIJHIPSPARSE, &flg));
1953:   PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for B of type %s", ((PetscObject)B)->type_name);
1954:   PetscCheck(!A->boundtocpu, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONG, "Cannot bind to CPU a HIPSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
1955:   PetscCheck(!B->boundtocpu, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONG, "Cannot bind to CPU a HIPSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
1956:   Acusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
1957:   Bcusp = (Mat_SeqAIJHIPSPARSE *)B->spptr;
1958:   Ccusp = (Mat_SeqAIJHIPSPARSE *)C->spptr;
1959:   PetscCheck(Acusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");
1960:   PetscCheck(Bcusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");
1961:   PetscCheck(Ccusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");
1962:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
1963:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(B));

1965:   ptype = product->type;
1966:   if (A->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_AtB) {
1967:     ptype = MATPRODUCT_AB;
1968:     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");
1969:   }
1970:   if (B->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_ABt) {
1971:     ptype = MATPRODUCT_AB;
1972:     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");
1973:   }
1974:   switch (ptype) {
1975:   case MATPRODUCT_AB:
1976:     Amat = Acusp->mat;
1977:     Bmat = Bcusp->mat;
1978:     break;
1979:   case MATPRODUCT_AtB:
1980:     PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(A));
1981:     Amat = Acusp->matTranspose;
1982:     Bmat = Bcusp->mat;
1983:     break;
1984:   case MATPRODUCT_ABt:
1985:     Amat = Acusp->mat;
1986:     PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(B));
1987:     Bmat = Bcusp->matTranspose;
1988:     break;
1989:   default:
1990:     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
1991:   }
1992:   Cmat = Ccusp->mat;
1993:   PetscCheck(Amat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A mult struct for product type %s", MatProductTypes[ptype]);
1994:   PetscCheck(Bmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B mult struct for product type %s", MatProductTypes[ptype]);
1995:   PetscCheck(Cmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C mult struct for product type %s", MatProductTypes[ptype]);
1996:   Acsr = (CsrMatrix *)Amat->mat;
1997:   Bcsr = mmdata->Bcsr ? mmdata->Bcsr : (CsrMatrix *)Bmat->mat; /* B may be in compressed row storage */
1998:   Ccsr = (CsrMatrix *)Cmat->mat;
1999:   PetscCheck(Acsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A CSR struct");
2000:   PetscCheck(Bcsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B CSR struct");
2001:   PetscCheck(Ccsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C CSR struct");
2002:   PetscCall(PetscLogGpuTimeBegin());
2003:   BmatSpDescr = mmdata->Bcsr ? mmdata->matSpBDescr : Bmat->matDescr; /* B may be in compressed row storage */
2004:   PetscCallHIPSPARSE(hipsparseSetPointerMode(Ccusp->handle, HIPSPARSE_POINTER_MODE_DEVICE));
2005:   if (c->nz) PetscCallHIPSPARSE(hipsparseSpGEMMreuse_compute(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, hipsparse_scalartype, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc));
2006:   PetscCall(PetscLogGpuFlops(mmdata->flops));
2007:   PetscCallHIP(WaitForHIP());
2008:   PetscCall(PetscLogGpuTimeEnd());
2009:   C->offloadmask = PETSC_OFFLOAD_GPU;
2010: finalize:
2011:   /* shorter version of MatAssemblyEnd_SeqAIJ */
2012:   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));
2013:   PetscCall(PetscInfo(C, "Number of mallocs during MatSetValues() is 0\n"));
2014:   PetscCall(PetscInfo(C, "Maximum nonzeros in any row is %" PetscInt_FMT "\n", c->rmax));
2015:   c->reallocs = 0;
2016:   C->info.mallocs += 0;
2017:   C->info.nz_unneeded = 0;
2018:   C->assembled = C->was_assembled = PETSC_TRUE;
2019:   C->num_ass++;
2020:   PetscFunctionReturn(PETSC_SUCCESS);
2021: }

2023: static PetscErrorCode MatProductSymbolic_SeqAIJHIPSPARSE_SeqAIJHIPSPARSE(Mat C)
2024: {
2025:   Mat_Product                   *product = C->product;
2026:   Mat                            A, B;
2027:   Mat_SeqAIJHIPSPARSE           *Acusp, *Bcusp, *Ccusp;
2028:   Mat_SeqAIJ                    *a, *b, *c;
2029:   Mat_SeqAIJHIPSPARSEMultStruct *Amat, *Bmat, *Cmat;
2030:   CsrMatrix                     *Acsr, *Bcsr, *Ccsr;
2031:   PetscInt                       i, j, m, n, k;
2032:   PetscBool                      flg;
2033:   MatProductType                 ptype;
2034:   MatProductCtx_MatMatHipsparse *mmdata;
2035:   PetscLogDouble                 flops;
2036:   PetscBool                      biscompressed, ciscompressed;
2037:   int64_t                        C_num_rows1, C_num_cols1, C_nnz1;
2038:   hipsparseSpMatDescr_t          BmatSpDescr;
2039:   hipsparseOperation_t           opA = HIPSPARSE_OPERATION_NON_TRANSPOSE, opB = HIPSPARSE_OPERATION_NON_TRANSPOSE; /* HIPSPARSE spgemm doesn't support transpose yet */

2041:   PetscFunctionBegin;
2042:   MatCheckProduct(C, 1);
2043:   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data not empty");
2044:   A = product->A;
2045:   B = product->B;
2046:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
2047:   PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
2048:   PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQAIJHIPSPARSE, &flg));
2049:   PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for B of type %s", ((PetscObject)B)->type_name);
2050:   a = (Mat_SeqAIJ *)A->data;
2051:   b = (Mat_SeqAIJ *)B->data;
2052:   /* product data */
2053:   PetscCall(PetscNew(&mmdata));
2054:   C->product->data    = mmdata;
2055:   C->product->destroy = MatProductCtxDestroy_MatMatHipsparse;

2057:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
2058:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(B));
2059:   Acusp = (Mat_SeqAIJHIPSPARSE *)A->spptr; /* Access spptr after MatSeqAIJHIPSPARSECopyToGPU, not before */
2060:   Bcusp = (Mat_SeqAIJHIPSPARSE *)B->spptr;
2061:   PetscCheck(Acusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");
2062:   PetscCheck(Bcusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");

2064:   ptype = product->type;
2065:   if (A->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_AtB) {
2066:     ptype                                          = MATPRODUCT_AB;
2067:     product->symbolic_used_the_fact_A_is_symmetric = PETSC_TRUE;
2068:   }
2069:   if (B->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_ABt) {
2070:     ptype                                          = MATPRODUCT_AB;
2071:     product->symbolic_used_the_fact_B_is_symmetric = PETSC_TRUE;
2072:   }
2073:   biscompressed = PETSC_FALSE;
2074:   ciscompressed = PETSC_FALSE;
2075:   switch (ptype) {
2076:   case MATPRODUCT_AB:
2077:     m    = A->rmap->n;
2078:     n    = B->cmap->n;
2079:     k    = A->cmap->n;
2080:     Amat = Acusp->mat;
2081:     Bmat = Bcusp->mat;
2082:     if (a->compressedrow.use) ciscompressed = PETSC_TRUE;
2083:     if (b->compressedrow.use) biscompressed = PETSC_TRUE;
2084:     break;
2085:   case MATPRODUCT_AtB:
2086:     m = A->cmap->n;
2087:     n = B->cmap->n;
2088:     k = A->rmap->n;
2089:     PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(A));
2090:     Amat = Acusp->matTranspose;
2091:     Bmat = Bcusp->mat;
2092:     if (b->compressedrow.use) biscompressed = PETSC_TRUE;
2093:     break;
2094:   case MATPRODUCT_ABt:
2095:     m = A->rmap->n;
2096:     n = B->rmap->n;
2097:     k = A->cmap->n;
2098:     PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(B));
2099:     Amat = Acusp->mat;
2100:     Bmat = Bcusp->matTranspose;
2101:     if (a->compressedrow.use) ciscompressed = PETSC_TRUE;
2102:     break;
2103:   default:
2104:     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
2105:   }

2107:   /* create hipsparse matrix */
2108:   PetscCall(MatSetSizes(C, m, n, m, n));
2109:   PetscCall(MatSetType(C, MATSEQAIJHIPSPARSE));
2110:   c     = (Mat_SeqAIJ *)C->data;
2111:   Ccusp = (Mat_SeqAIJHIPSPARSE *)C->spptr;
2112:   Cmat  = new Mat_SeqAIJHIPSPARSEMultStruct;
2113:   Ccsr  = new CsrMatrix;

2115:   c->compressedrow.use = ciscompressed;
2116:   if (c->compressedrow.use) { /* if a is in compressed row, than c will be in compressed row format */
2117:     c->compressedrow.nrows = a->compressedrow.nrows;
2118:     PetscCall(PetscMalloc2(c->compressedrow.nrows + 1, &c->compressedrow.i, c->compressedrow.nrows, &c->compressedrow.rindex));
2119:     PetscCall(PetscArraycpy(c->compressedrow.rindex, a->compressedrow.rindex, c->compressedrow.nrows));
2120:     Ccusp->workVector  = new THRUSTARRAY(c->compressedrow.nrows);
2121:     Cmat->cprowIndices = new THRUSTINTARRAY(c->compressedrow.nrows);
2122:     Cmat->cprowIndices->assign(c->compressedrow.rindex, c->compressedrow.rindex + c->compressedrow.nrows);
2123:   } else {
2124:     c->compressedrow.nrows  = 0;
2125:     c->compressedrow.i      = NULL;
2126:     c->compressedrow.rindex = NULL;
2127:     Ccusp->workVector       = NULL;
2128:     Cmat->cprowIndices      = NULL;
2129:   }
2130:   Ccusp->nrows      = ciscompressed ? c->compressedrow.nrows : m;
2131:   Ccusp->mat        = Cmat;
2132:   Ccusp->mat->mat   = Ccsr;
2133:   Ccsr->num_rows    = Ccusp->nrows;
2134:   Ccsr->num_cols    = n;
2135:   Ccsr->row_offsets = new THRUSTINTARRAY(Ccusp->nrows + 1);
2136:   PetscCallHIPSPARSE(hipsparseCreateMatDescr(&Cmat->descr));
2137:   PetscCallHIPSPARSE(hipsparseSetMatIndexBase(Cmat->descr, HIPSPARSE_INDEX_BASE_ZERO));
2138:   PetscCallHIPSPARSE(hipsparseSetMatType(Cmat->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));
2139:   PetscCallHIP(hipMalloc((void **)&Cmat->alpha_one, sizeof(PetscScalar)));
2140:   PetscCallHIP(hipMalloc((void **)&Cmat->beta_zero, sizeof(PetscScalar)));
2141:   PetscCallHIP(hipMalloc((void **)&Cmat->beta_one, sizeof(PetscScalar)));
2142:   PetscCallHIP(hipMemcpy(Cmat->alpha_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
2143:   PetscCallHIP(hipMemcpy(Cmat->beta_zero, &PETSC_HIPSPARSE_ZERO, sizeof(PetscScalar), hipMemcpyHostToDevice));
2144:   PetscCallHIP(hipMemcpy(Cmat->beta_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
2145:   if (!Ccsr->num_rows || !Ccsr->num_cols || !a->nz || !b->nz) { /* hipsparse raise errors in different calls when matrices have zero rows/columns! */
2146:     PetscCallThrust(thrust::fill(thrust::device, Ccsr->row_offsets->begin(), Ccsr->row_offsets->end(), 0));
2147:     c->nz                = 0;
2148:     Ccsr->column_indices = new THRUSTINTARRAY(c->nz);
2149:     Ccsr->values         = new THRUSTARRAY(c->nz);
2150:     goto finalizesym;
2151:   }

2153:   PetscCheck(Amat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A mult struct for product type %s", MatProductTypes[ptype]);
2154:   PetscCheck(Bmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B mult struct for product type %s", MatProductTypes[ptype]);
2155:   Acsr = (CsrMatrix *)Amat->mat;
2156:   if (!biscompressed) {
2157:     Bcsr        = (CsrMatrix *)Bmat->mat;
2158:     BmatSpDescr = Bmat->matDescr;
2159:   } else { /* we need to use row offsets for the full matrix */
2160:     CsrMatrix *cBcsr     = (CsrMatrix *)Bmat->mat;
2161:     Bcsr                 = new CsrMatrix;
2162:     Bcsr->num_rows       = B->rmap->n;
2163:     Bcsr->num_cols       = cBcsr->num_cols;
2164:     Bcsr->num_entries    = cBcsr->num_entries;
2165:     Bcsr->column_indices = cBcsr->column_indices;
2166:     Bcsr->values         = cBcsr->values;
2167:     if (!Bcusp->rowoffsets_gpu) {
2168:       Bcusp->rowoffsets_gpu = new THRUSTINTARRAY(B->rmap->n + 1);
2169:       Bcusp->rowoffsets_gpu->assign(b->i, b->i + B->rmap->n + 1);
2170:       PetscCall(PetscLogCpuToGpu((B->rmap->n + 1) * sizeof(PetscInt)));
2171:     }
2172:     Bcsr->row_offsets = Bcusp->rowoffsets_gpu;
2173:     mmdata->Bcsr      = Bcsr;
2174:     if (Bcsr->num_rows && Bcsr->num_cols) {
2175:       PetscCallHIPSPARSE(hipsparseCreateCsr(&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, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
2176:     }
2177:     BmatSpDescr = mmdata->matSpBDescr;
2178:   }
2179:   PetscCheck(Acsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A CSR struct");
2180:   PetscCheck(Bcsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B CSR struct");
2181:   /* precompute flops count */
2182:   if (ptype == MATPRODUCT_AB) {
2183:     for (i = 0, flops = 0; i < A->rmap->n; i++) {
2184:       const PetscInt st = a->i[i];
2185:       const PetscInt en = a->i[i + 1];
2186:       for (j = st; j < en; j++) {
2187:         const PetscInt brow = a->j[j];
2188:         flops += 2. * (b->i[brow + 1] - b->i[brow]);
2189:       }
2190:     }
2191:   } else if (ptype == MATPRODUCT_AtB) {
2192:     for (i = 0, flops = 0; i < A->rmap->n; i++) {
2193:       const PetscInt anzi = a->i[i + 1] - a->i[i];
2194:       const PetscInt bnzi = b->i[i + 1] - b->i[i];
2195:       flops += (2. * anzi) * bnzi;
2196:     }
2197:   } else flops = 0.; /* TODO */

2199:   mmdata->flops = flops;
2200:   PetscCall(PetscLogGpuTimeBegin());

2202:   // cuSparse first had cusparseSpGEMM_compute(), then had cusparseSpGEMMreuse_compute(). But NVIDIA deprecated cusparseSpGEMMreuse_compute()
2203:   // later and went back to cusparseSpGEMM_compute(). AMD followed suit. However, hipsparseSpGEMM_compute() is correct on rocm-5.4.3 but buggy
2204:   // in rocm-6.2 and 7.2. Second call to hipsparseSpGEMM_compute() gives stale result with updated input matrices.
2205:   // See https://github.com/ROCm/rocm-libraries/issues/8878. So we will use the _reuse version until AMD fixes the bug.
2206:   {
2207:     PetscCallHIPSPARSE(hipsparseSetPointerMode(Ccusp->handle, HIPSPARSE_POINTER_MODE_DEVICE));
2208:     PetscCallHIPSPARSE(hipsparseCreateCsr(&Cmat->matDescr, Ccsr->num_rows, Ccsr->num_cols, 0, Ccsr->row_offsets->data().get(), NULL, NULL, csrRowOffsetsType, csrColIndType, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
2209:     PetscCallHIPSPARSE(hipsparseSpGEMM_createDescr(&mmdata->spgemmDesc));

2211:     /* hipsparseSpGEMMreuse has more reasonable APIs than hipsparseSpGEMM, so we prefer to use it.
2212:      We follow the sample code at https://github.com/ROCmSoftwarePlatform/hipSPARSE/blob/develop/clients/include/testing_spgemmreuse_csr.hpp
2213:   */
2214:     void *dBuffer1 = NULL;
2215:     void *dBuffer2 = NULL;
2216:     void *dBuffer3 = NULL;
2217:     /* dBuffer4, dBuffer5 are needed by hipsparseSpGEMMreuse_compute, and therefore are stored in mmdata */
2218:     size_t bufferSize1 = 0;
2219:     size_t bufferSize2 = 0;
2220:     size_t bufferSize3 = 0;
2221:     size_t bufferSize4 = 0;
2222:     size_t bufferSize5 = 0;

2224:     /* ask bufferSize1 bytes for external memory */
2225:     PetscCallHIPSPARSE(hipsparseSpGEMMreuse_workEstimation(Ccusp->handle, opA, opB, Amat->matDescr, BmatSpDescr, Cmat->matDescr, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufferSize1, NULL));
2226:     PetscCallHIP(hipMalloc((void **)&dBuffer1, bufferSize1));
2227:     /* inspect the matrices A and B to understand the memory requirement for the next step */
2228:     PetscCallHIPSPARSE(hipsparseSpGEMMreuse_workEstimation(Ccusp->handle, opA, opB, Amat->matDescr, BmatSpDescr, Cmat->matDescr, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufferSize1, dBuffer1));

2230:     PetscCallHIPSPARSE(hipsparseSpGEMMreuse_nnz(Ccusp->handle, opA, opB, Amat->matDescr, BmatSpDescr, Cmat->matDescr, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufferSize2, NULL, &bufferSize3, NULL, &bufferSize4, NULL));
2231:     PetscCallHIP(hipMalloc((void **)&dBuffer2, bufferSize2));
2232:     PetscCallHIP(hipMalloc((void **)&dBuffer3, bufferSize3));
2233:     PetscCallHIP(hipMalloc((void **)&mmdata->dBuffer4, bufferSize4));
2234:     PetscCallHIPSPARSE(hipsparseSpGEMMreuse_nnz(Ccusp->handle, opA, opB, Amat->matDescr, BmatSpDescr, Cmat->matDescr, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufferSize2, dBuffer2, &bufferSize3, dBuffer3, &bufferSize4, mmdata->dBuffer4));
2235:     PetscCallHIP(hipFree(dBuffer1));
2236:     PetscCallHIP(hipFree(dBuffer2));

2238:     /* get matrix C non-zero entries C_nnz1 */
2239:     PetscCallHIPSPARSE(hipsparseSpMatGetSize(Cmat->matDescr, &C_num_rows1, &C_num_cols1, &C_nnz1));
2240:     PetscCall(PetscIntCast(C_nnz1, &c->nz));
2241:     /* allocate matrix C */
2242:     Ccsr->column_indices = new THRUSTINTARRAY(c->nz);
2243:     PetscCallHIP(hipPeekAtLastError()); /* catch out of memory errors */
2244:     Ccsr->values = new THRUSTARRAY(c->nz);
2245:     PetscCallHIP(hipPeekAtLastError()); /* catch out of memory errors */
2246:     /* update matC with the new pointers */
2247:     if (c->nz) { /* 5.5.1 has a bug with nz = 0, exposed by mat_tests_ex123_2_hypre */
2248:       PetscCallHIPSPARSE(hipsparseCsrSetPointers(Cmat->matDescr, Ccsr->row_offsets->data().get(), Ccsr->column_indices->data().get(), Ccsr->values->data().get()));

2250:       PetscCallHIPSPARSE(hipsparseSpGEMMreuse_copy(Ccusp->handle, opA, opB, Amat->matDescr, BmatSpDescr, Cmat->matDescr, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufferSize5, NULL));
2251:       PetscCallHIP(hipMalloc((void **)&mmdata->dBuffer5, bufferSize5));
2252:       PetscCallHIPSPARSE(hipsparseSpGEMMreuse_copy(Ccusp->handle, opA, opB, Amat->matDescr, BmatSpDescr, Cmat->matDescr, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufferSize5, mmdata->dBuffer5));
2253:       PetscCallHIPSPARSE(hipsparseSpGEMMreuse_compute(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, hipsparse_scalartype, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc));
2254:     }
2255:     PetscCallHIP(hipFree(dBuffer3));
2256:     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, bufferSize4 / 1024, bufferSize5 / 1024));
2257:   }

2259:   PetscCall(PetscLogGpuFlops(mmdata->flops));
2260:   PetscCall(PetscLogGpuTimeEnd());
2261: finalizesym:
2262:   c->free_a = PETSC_TRUE;
2263:   PetscCall(PetscShmgetAllocateArray(c->nz, sizeof(PetscInt), (void **)&c->j));
2264:   PetscCall(PetscShmgetAllocateArray(m + 1, sizeof(PetscInt), (void **)&c->i));
2265:   c->free_ij = PETSC_TRUE;

2267:   PetscInt *d_i = c->i;
2268:   if (ciscompressed) d_i = c->compressedrow.i;
2269:   PetscCallHIP(hipMemcpy(d_i, Ccsr->row_offsets->data().get(), Ccsr->row_offsets->size() * sizeof(PetscInt), hipMemcpyDeviceToHost));
2270:   PetscCallHIP(hipMemcpy(c->j, Ccsr->column_indices->data().get(), Ccsr->column_indices->size() * sizeof(PetscInt), hipMemcpyDeviceToHost));
2271:   if (ciscompressed) { /* need to expand host row offsets */
2272:     PetscInt r = 0;
2273:     c->i[0]    = 0;
2274:     for (k = 0; k < c->compressedrow.nrows; k++) {
2275:       const PetscInt next = c->compressedrow.rindex[k];
2276:       const PetscInt old  = c->compressedrow.i[k];
2277:       for (; r < next; r++) c->i[r + 1] = old;
2278:     }
2279:     for (; r < m; r++) c->i[r + 1] = c->compressedrow.i[c->compressedrow.nrows];
2280:   }
2281:   PetscCall(PetscLogGpuToCpu((Ccsr->column_indices->size() + Ccsr->row_offsets->size()) * sizeof(PetscInt)));
2282:   PetscCall(PetscMalloc1(m, &c->ilen));
2283:   PetscCall(PetscMalloc1(m, &c->imax));
2284:   c->maxnz         = c->nz;
2285:   c->nonzerorowcnt = 0;
2286:   c->rmax          = 0;
2287:   for (k = 0; k < m; k++) {
2288:     const PetscInt nn = c->i[k + 1] - c->i[k];
2289:     c->ilen[k] = c->imax[k] = nn;
2290:     c->nonzerorowcnt += (PetscInt)!!nn;
2291:     c->rmax = PetscMax(c->rmax, nn);
2292:   }
2293:   PetscCall(PetscMalloc1(c->nz, &c->a));
2294:   Ccsr->num_entries = c->nz;

2296:   C->nonzerostate++;
2297:   PetscCall(PetscLayoutSetUp(C->rmap));
2298:   PetscCall(PetscLayoutSetUp(C->cmap));
2299:   Ccusp->nonzerostate = C->nonzerostate;
2300:   C->offloadmask      = PETSC_OFFLOAD_UNALLOCATED;
2301:   C->preallocated     = PETSC_TRUE;
2302:   C->assembled        = PETSC_FALSE;
2303:   C->was_assembled    = PETSC_FALSE;
2304:   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 */
2305:     mmdata->reusesym = PETSC_TRUE;
2306:     C->offloadmask   = PETSC_OFFLOAD_GPU;
2307:   }
2308:   C->ops->productnumeric = MatProductNumeric_SeqAIJHIPSPARSE_SeqAIJHIPSPARSE;
2309:   PetscFunctionReturn(PETSC_SUCCESS);
2310: }

2312: /* handles sparse or dense B */
2313: static PetscErrorCode MatProductSetFromOptions_SeqAIJHIPSPARSE(Mat mat)
2314: {
2315:   Mat_Product *product = mat->product;
2316:   PetscBool    isdense = PETSC_FALSE, Biscusp = PETSC_FALSE, Ciscusp = PETSC_TRUE;

2318:   PetscFunctionBegin;
2319:   MatCheckProduct(mat, 1);
2320:   PetscCall(PetscObjectBaseTypeCompare((PetscObject)product->B, MATSEQDENSE, &isdense));
2321:   if (!product->A->boundtocpu && !product->B->boundtocpu) PetscCall(PetscObjectTypeCompare((PetscObject)product->B, MATSEQAIJHIPSPARSE, &Biscusp));
2322:   if (product->type == MATPRODUCT_ABC) {
2323:     Ciscusp = PETSC_FALSE;
2324:     if (!product->C->boundtocpu) PetscCall(PetscObjectTypeCompare((PetscObject)product->C, MATSEQAIJHIPSPARSE, &Ciscusp));
2325:   }
2326:   if (Biscusp && Ciscusp) { /* we can always select the CPU backend */
2327:     PetscBool usecpu = PETSC_FALSE;
2328:     switch (product->type) {
2329:     case MATPRODUCT_AB:
2330:       if (product->api_user) {
2331:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatMatMult", "Mat");
2332:         PetscCall(PetscOptionsBool("-matmatmult_backend_cpu", "Use CPU code", "MatMatMult", usecpu, &usecpu, NULL));
2333:         PetscOptionsEnd();
2334:       } else {
2335:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_AB", "Mat");
2336:         PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatMatMult", usecpu, &usecpu, NULL));
2337:         PetscOptionsEnd();
2338:       }
2339:       break;
2340:     case MATPRODUCT_AtB:
2341:       if (product->api_user) {
2342:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatTransposeMatMult", "Mat");
2343:         PetscCall(PetscOptionsBool("-mattransposematmult_backend_cpu", "Use CPU code", "MatTransposeMatMult", usecpu, &usecpu, NULL));
2344:         PetscOptionsEnd();
2345:       } else {
2346:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_AtB", "Mat");
2347:         PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatTransposeMatMult", usecpu, &usecpu, NULL));
2348:         PetscOptionsEnd();
2349:       }
2350:       break;
2351:     case MATPRODUCT_PtAP:
2352:       if (product->api_user) {
2353:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatPtAP", "Mat");
2354:         PetscCall(PetscOptionsBool("-matptap_backend_cpu", "Use CPU code", "MatPtAP", usecpu, &usecpu, NULL));
2355:         PetscOptionsEnd();
2356:       } else {
2357:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_PtAP", "Mat");
2358:         PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatPtAP", usecpu, &usecpu, NULL));
2359:         PetscOptionsEnd();
2360:       }
2361:       break;
2362:     case MATPRODUCT_RARt:
2363:       if (product->api_user) {
2364:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatRARt", "Mat");
2365:         PetscCall(PetscOptionsBool("-matrart_backend_cpu", "Use CPU code", "MatRARt", usecpu, &usecpu, NULL));
2366:         PetscOptionsEnd();
2367:       } else {
2368:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_RARt", "Mat");
2369:         PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatRARt", usecpu, &usecpu, NULL));
2370:         PetscOptionsEnd();
2371:       }
2372:       break;
2373:     case MATPRODUCT_ABC:
2374:       if (product->api_user) {
2375:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatMatMatMult", "Mat");
2376:         PetscCall(PetscOptionsBool("-matmatmatmult_backend_cpu", "Use CPU code", "MatMatMatMult", usecpu, &usecpu, NULL));
2377:         PetscOptionsEnd();
2378:       } else {
2379:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_ABC", "Mat");
2380:         PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatMatMatMult", usecpu, &usecpu, NULL));
2381:         PetscOptionsEnd();
2382:       }
2383:       break;
2384:     default:
2385:       break;
2386:     }
2387:     if (usecpu) Biscusp = Ciscusp = PETSC_FALSE;
2388:   }
2389:   /* dispatch */
2390:   if (isdense) {
2391:     switch (product->type) {
2392:     case MATPRODUCT_AB:
2393:     case MATPRODUCT_AtB:
2394:     case MATPRODUCT_ABt:
2395:     case MATPRODUCT_PtAP:
2396:     case MATPRODUCT_RARt:
2397:       if (product->A->boundtocpu) PetscCall(MatProductSetFromOptions_SeqAIJ_SeqDense(mat));
2398:       else mat->ops->productsymbolic = MatProductSymbolic_SeqAIJHIPSPARSE_SeqDENSEHIP;
2399:       break;
2400:     case MATPRODUCT_ABC:
2401:       mat->ops->productsymbolic = MatProductSymbolic_ABC_Basic;
2402:       break;
2403:     default:
2404:       break;
2405:     }
2406:   } else if (Biscusp && Ciscusp) {
2407:     switch (product->type) {
2408:     case MATPRODUCT_AB:
2409:     case MATPRODUCT_AtB:
2410:     case MATPRODUCT_ABt:
2411:       mat->ops->productsymbolic = MatProductSymbolic_SeqAIJHIPSPARSE_SeqAIJHIPSPARSE;
2412:       break;
2413:     case MATPRODUCT_PtAP:
2414:     case MATPRODUCT_RARt:
2415:     case MATPRODUCT_ABC:
2416:       mat->ops->productsymbolic = MatProductSymbolic_ABC_Basic;
2417:       break;
2418:     default:
2419:       break;
2420:     }
2421:   } else PetscCall(MatProductSetFromOptions_SeqAIJ(mat)); /* fallback for AIJ */
2422:   PetscFunctionReturn(PETSC_SUCCESS);
2423: }

2425: static PetscErrorCode MatMult_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy)
2426: {
2427:   PetscFunctionBegin;
2428:   PetscCall(MatMultAddKernel_SeqAIJHIPSPARSE(A, xx, NULL, yy, PETSC_FALSE, PETSC_FALSE));
2429:   PetscFunctionReturn(PETSC_SUCCESS);
2430: }

2432: static PetscErrorCode MatMultAdd_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy, Vec zz)
2433: {
2434:   PetscFunctionBegin;
2435:   PetscCall(MatMultAddKernel_SeqAIJHIPSPARSE(A, xx, yy, zz, PETSC_FALSE, PETSC_FALSE));
2436:   PetscFunctionReturn(PETSC_SUCCESS);
2437: }

2439: static PetscErrorCode MatMultHermitianTranspose_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy)
2440: {
2441:   PetscFunctionBegin;
2442:   PetscCall(MatMultAddKernel_SeqAIJHIPSPARSE(A, xx, NULL, yy, PETSC_TRUE, PETSC_TRUE));
2443:   PetscFunctionReturn(PETSC_SUCCESS);
2444: }

2446: static PetscErrorCode MatMultHermitianTransposeAdd_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy, Vec zz)
2447: {
2448:   PetscFunctionBegin;
2449:   PetscCall(MatMultAddKernel_SeqAIJHIPSPARSE(A, xx, yy, zz, PETSC_TRUE, PETSC_TRUE));
2450:   PetscFunctionReturn(PETSC_SUCCESS);
2451: }

2453: static PetscErrorCode MatMultTranspose_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy)
2454: {
2455:   PetscFunctionBegin;
2456:   PetscCall(MatMultAddKernel_SeqAIJHIPSPARSE(A, xx, NULL, yy, PETSC_TRUE, PETSC_FALSE));
2457:   PetscFunctionReturn(PETSC_SUCCESS);
2458: }

2460: __global__ static void ScatterAdd(PetscInt n, PetscInt *idx, const PetscScalar *x, PetscScalar *y)
2461: {
2462:   int i = blockIdx.x * blockDim.x + threadIdx.x;
2463:   if (i < n) y[idx[i]] += x[i];
2464: }

2466: /* z = op(A) x + y. If trans & !herm, op = ^T; if trans & herm, op = ^H; if !trans, op = no-op */
2467: static PetscErrorCode MatMultAddKernel_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy, Vec zz, PetscBool trans, PetscBool herm)
2468: {
2469:   Mat_SeqAIJ                    *a               = (Mat_SeqAIJ *)A->data;
2470:   Mat_SeqAIJHIPSPARSE           *hipsparsestruct = (Mat_SeqAIJHIPSPARSE *)A->spptr;
2471:   Mat_SeqAIJHIPSPARSEMultStruct *matstruct;
2472:   PetscScalar                   *xarray, *zarray, *dptr, *beta, *xptr;
2473:   hipsparseOperation_t           opA = HIPSPARSE_OPERATION_NON_TRANSPOSE;
2474:   PetscBool                      compressed;
2475:   PetscInt                       nx, ny;

2477:   PetscFunctionBegin;
2478:   PetscCheck(!herm || trans, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Hermitian and not transpose not supported");
2479:   if (!a->nz) {
2480:     if (yy) PetscCall(VecSeq_HIP::Copy(yy, zz));
2481:     else PetscCall(VecSeq_HIP::Set(zz, 0));
2482:     PetscFunctionReturn(PETSC_SUCCESS);
2483:   }
2484:   /* The line below is necessary due to the operations that modify the matrix on the CPU (axpy, scale, etc) */
2485:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
2486:   if (!trans) {
2487:     matstruct = (Mat_SeqAIJHIPSPARSEMultStruct *)hipsparsestruct->mat;
2488:     PetscCheck(matstruct, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "SeqAIJHIPSPARSE does not have a 'mat' (need to fix)");
2489:   } else {
2490:     if (herm || !A->form_explicit_transpose) {
2491:       opA       = herm ? HIPSPARSE_OPERATION_CONJUGATE_TRANSPOSE : HIPSPARSE_OPERATION_TRANSPOSE;
2492:       matstruct = (Mat_SeqAIJHIPSPARSEMultStruct *)hipsparsestruct->mat;
2493:     } else {
2494:       if (!A->transupdated) PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(A));
2495:       matstruct = (Mat_SeqAIJHIPSPARSEMultStruct *)hipsparsestruct->matTranspose;
2496:     }
2497:   }
2498:   /* Does the matrix use compressed rows (i.e., drop zero rows)? */
2499:   compressed = matstruct->cprowIndices ? PETSC_TRUE : PETSC_FALSE;
2500:   try {
2501:     PetscCall(VecHIPGetArrayRead(xx, (const PetscScalar **)&xarray));
2502:     if (yy == zz) PetscCall(VecHIPGetArray(zz, &zarray)); /* read & write zz, so need to get up-to-date zarray on GPU */
2503:     else PetscCall(VecHIPGetArrayWrite(zz, &zarray));     /* write zz, so no need to init zarray on GPU */

2505:     PetscCall(PetscLogGpuTimeBegin());
2506:     if (opA == HIPSPARSE_OPERATION_NON_TRANSPOSE) {
2507:       /* z = A x + beta y.
2508:          If A is compressed (with less rows), then Ax is shorter than the full z, so we need a work vector to store Ax.
2509:          When A is non-compressed, and z = y, we can set beta=1 to compute y = Ax + y in one call.
2510:       */
2511:       xptr = xarray;
2512:       dptr = compressed ? hipsparsestruct->workVector->data().get() : zarray;
2513:       beta = (yy == zz && !compressed) ? matstruct->beta_one : matstruct->beta_zero;
2514:       /* Get length of x, y for y=Ax. ny might be shorter than the work vector's allocated length, since the work vector is
2515:           allocated to accommodate different uses. So we get the length info directly from mat.
2516:        */
2517:       if (hipsparsestruct->format == MAT_HIPSPARSE_CSR) {
2518:         CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
2519:         nx             = mat->num_cols;
2520:         ny             = mat->num_rows;
2521:       }
2522:     } else {
2523:       /* z = A^T x + beta y
2524:          If A is compressed, then we need a work vector as the shorter version of x to compute A^T x.
2525:          Note A^Tx is of full length, so we set beta to 1.0 if y exists.
2526:        */
2527:       xptr = compressed ? hipsparsestruct->workVector->data().get() : xarray;
2528:       dptr = zarray;
2529:       beta = yy ? matstruct->beta_one : matstruct->beta_zero;
2530:       if (compressed) { /* Scatter x to work vector */
2531:         thrust::device_ptr<PetscScalar> xarr = thrust::device_pointer_cast(xarray);
2532:         thrust::for_each(
2533: #if PetscDefined(HAVE_THRUST_ASYNC)
2534:           thrust::device.on(PetscDefaultHipStream),
2535: #endif
2536:           thrust::make_zip_iterator(thrust::make_tuple(hipsparsestruct->workVector->begin(), thrust::make_permutation_iterator(xarr, matstruct->cprowIndices->begin()))),
2537:           thrust::make_zip_iterator(thrust::make_tuple(hipsparsestruct->workVector->begin(), thrust::make_permutation_iterator(xarr, matstruct->cprowIndices->begin()))) + matstruct->cprowIndices->size(), VecHIPEqualsReverse());
2538:       }
2539:       if (hipsparsestruct->format == MAT_HIPSPARSE_CSR) {
2540:         CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
2541:         nx             = mat->num_rows;
2542:         ny             = mat->num_cols;
2543:       }
2544:     }
2545:     /* csr_spmv does y = alpha op(A) x + beta y */
2546:     if (hipsparsestruct->format == MAT_HIPSPARSE_CSR) {
2547:       hipsparseSpMVAlg_t spmvAlg = hipsparsestruct->spmvAlg;

2549: #if PETSC_PKG_HIP_VERSION_GE(6, 4, 0) && PETSC_PKG_HIP_VERSION_LT(7, 3, 0) // ALG_DEFAULT in some versions don't support transpose. See https://github.com/ROCm/rocm-libraries/issues/4803
2550:       // We found rocm-7.2.4 also had this problem. We assume rocm-7.3+ will contain a fix.
2551:       if (opA != HIPSPARSE_OPERATION_NON_TRANSPOSE && spmvAlg == HIPSPARSE_SPMV_ALG_DEFAULT) spmvAlg = HIPSPARSE_SPMV_CSR_ALG1;
2552: #endif
2553:       PetscCheck(opA >= 0 && opA <= 2, PETSC_COMM_SELF, PETSC_ERR_SUP, "hipSPARSE API on hipsparseOperation_t has changed and PETSc has not been updated accordingly");

2555:       if (!matstruct->hipSpMV[opA].initialized) { /* built on demand */
2556:         CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
2557:         PetscCallHIPSPARSE(hipsparseCreateCsr(&matstruct->hipSpMV[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, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
2558:         PetscCallHIPSPARSE(hipsparseCreateDnVec(&matstruct->hipSpMV[opA].vecXDescr, nx, xptr, hipsparse_scalartype));
2559:         PetscCallHIPSPARSE(hipsparseCreateDnVec(&matstruct->hipSpMV[opA].vecYDescr, ny, dptr, hipsparse_scalartype));
2560:         PetscCallHIPSPARSE(hipsparseSpMV_bufferSize(hipsparsestruct->handle, opA, matstruct->alpha_one, matstruct->hipSpMV[opA].matDescr, matstruct->hipSpMV[opA].vecXDescr, beta, matstruct->hipSpMV[opA].vecYDescr, hipsparse_scalartype, spmvAlg,
2561:                                                     &matstruct->hipSpMV[opA].spmvBufferSize));
2562:         PetscCallHIP(hipMalloc(&matstruct->hipSpMV[opA].spmvBuffer, matstruct->hipSpMV[opA].spmvBufferSize));
2563: #if PETSC_PKG_HIP_VERSION_GE(6, 4, 0) // hipsparseSpMV_preprocess is added in rocm-6.4.0
2564:         PetscCallHIPSPARSE(hipsparseSpMV_preprocess(hipsparsestruct->handle, opA, matstruct->alpha_one, matstruct->hipSpMV[opA].matDescr, matstruct->hipSpMV[opA].vecXDescr, beta, matstruct->hipSpMV[opA].vecYDescr, hipsparse_scalartype, spmvAlg,
2565:                                                     matstruct->hipSpMV[opA].spmvBuffer));
2566: #endif
2567:         matstruct->hipSpMV[opA].initialized = PETSC_TRUE;
2568:       } else {
2569:         /* x, y's value pointers might change between calls, but their shape is kept, so we just update pointers */
2570:         PetscCallHIPSPARSE(hipsparseDnVecSetValues(matstruct->hipSpMV[opA].vecXDescr, xptr));
2571:         PetscCallHIPSPARSE(hipsparseDnVecSetValues(matstruct->hipSpMV[opA].vecYDescr, dptr));
2572:       }
2573:       PetscCallHIPSPARSE(
2574:         hipsparseSpMV(hipsparsestruct->handle, opA, matstruct->alpha_one, matstruct->hipSpMV[opA].matDescr, matstruct->hipSpMV[opA].vecXDescr, beta, matstruct->hipSpMV[opA].vecYDescr, hipsparse_scalartype, spmvAlg, matstruct->hipSpMV[opA].spmvBuffer));
2575:     } else PetscCheck(hipsparsestruct->nrows == 0, PETSC_COMM_SELF, PETSC_ERR_SUP, "MAT_HIPSPARSE_ELL and MAT_HIPSPARSE_HYB are not supported");

2577:     PetscCall(PetscLogGpuTimeEnd());

2579:     if (opA == HIPSPARSE_OPERATION_NON_TRANSPOSE) {
2580:       if (yy) {                                     /* MatMultAdd: zz = A*xx + yy */
2581:         if (compressed) {                           /* A is compressed. We first copy yy to zz, then ScatterAdd the work vector to zz */
2582:           PetscCall(VecSeq_HIP::Copy(yy, zz));      /* zz = yy */
2583:         } else if (zz != yy) {                      /* A is not compressed. zz already contains A*xx, and we just need to add yy */
2584:           PetscCall(VecSeq_HIP::AXPY(zz, 1.0, yy)); /* zz += yy */
2585:         }
2586:       } else if (compressed) { /* MatMult: zz = A*xx. A is compressed, so we zero zz first, then ScatterAdd the work vector to zz */
2587:         PetscCall(VecSeq_HIP::Set(zz, 0));
2588:       }

2590:       /* ScatterAdd the result from work vector into the full vector when A is compressed */
2591:       if (compressed) {
2592:         PetscCall(PetscLogGpuTimeBegin());
2593:         /* I wanted to make this for_each asynchronous but failed. thrust::async::for_each() returns an event (internally registered)
2594:            and in the destructor of the scope, it will call hipStreamSynchronize() on this stream. One has to store all events to
2595:            prevent that. So I just add a ScatterAdd kernel.
2596:          */
2597: #if 0
2598:         thrust::device_ptr<PetscScalar> zptr = thrust::device_pointer_cast(zarray);
2599:         thrust::async::for_each(thrust::device.on(hipsparsestruct->stream),
2600:                          thrust::make_zip_iterator(thrust::make_tuple(hipsparsestruct->workVector->begin(), thrust::make_permutation_iterator(zptr, matstruct->cprowIndices->begin()))),
2601:                          thrust::make_zip_iterator(thrust::make_tuple(hipsparsestruct->workVector->begin(), thrust::make_permutation_iterator(zptr, matstruct->cprowIndices->begin()))) + matstruct->cprowIndices->size(),
2602:                          VecHIPPlusEquals());
2603: #else
2604:         PetscInt n = matstruct->cprowIndices->size();
2605:         hipLaunchKernelGGL(ScatterAdd, dim3((n + 255) / 256), dim3(256), 0, PetscDefaultHipStream, n, matstruct->cprowIndices->data().get(), hipsparsestruct->workVector->data().get(), zarray);
2606: #endif
2607:         PetscCall(PetscLogGpuTimeEnd());
2608:       }
2609:     } else {
2610:       if (yy && yy != zz) PetscCall(VecSeq_HIP::AXPY(zz, 1.0, yy)); /* zz += yy */
2611:     }
2612:     PetscCall(VecHIPRestoreArrayRead(xx, (const PetscScalar **)&xarray));
2613:     if (yy == zz) PetscCall(VecHIPRestoreArray(zz, &zarray));
2614:     else PetscCall(VecHIPRestoreArrayWrite(zz, &zarray));
2615:   } catch (char *ex) {
2616:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_LIB, "HIPSPARSE error: %s", ex);
2617:   }
2618:   if (yy) PetscCall(PetscLogGpuFlops(2.0 * a->nz));
2619:   else PetscCall(PetscLogGpuFlops(2.0 * a->nz - a->nonzerorowcnt));
2620:   PetscFunctionReturn(PETSC_SUCCESS);
2621: }

2623: static PetscErrorCode MatMultTransposeAdd_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy, Vec zz)
2624: {
2625:   PetscFunctionBegin;
2626:   PetscCall(MatMultAddKernel_SeqAIJHIPSPARSE(A, xx, yy, zz, PETSC_TRUE, PETSC_FALSE));
2627:   PetscFunctionReturn(PETSC_SUCCESS);
2628: }

2630: static PetscErrorCode MatAssemblyEnd_SeqAIJHIPSPARSE(Mat A, MatAssemblyType mode)
2631: {
2632:   PetscFunctionBegin;
2633:   PetscCall(MatSeqAIJHIPSPARSE_CUPM_t::AssemblyEnd(A, mode));
2634:   PetscFunctionReturn(PETSC_SUCCESS);
2635: }

2637: /*@
2638:   MatCreateSeqAIJHIPSPARSE - Creates a sparse matrix in `MATAIJHIPSPARSE` (compressed row) format.
2639:   This matrix will ultimately pushed down to AMD GPUs and use the HIPSPARSE library for calculations.

2641:   Collective

2643:   Input Parameters:
2644: + comm - MPI communicator, set to `PETSC_COMM_SELF`
2645: . m    - number of rows
2646: . n    - number of columns
2647: . nz   - number of nonzeros per row (same for all rows), ignored if `nnz` is set
2648: - nnz  - array containing the number of nonzeros in the various rows (possibly different for each row) or `NULL`

2650:   Output Parameter:
2651: . A - the matrix

2653:   Level: intermediate

2655:   Notes:
2656:   It is recommended that one use the `MatCreate()`, `MatSetType()` and/or `MatSetFromOptions()`,
2657:   `MatXXXXSetPreallocation()` paradgm instead of this routine directly.
2658:   [MatXXXXSetPreallocation() is, for example, `MatSeqAIJSetPreallocation`]

2660:   The AIJ format (compressed row storage), is fully compatible with standard Fortran
2661:   storage.  That is, the stored row and column indices can begin at
2662:   either one (as in Fortran) or zero.

2664:   Specify the preallocated storage with either `nz` or `nnz` (not both).
2665:   Set `nz` = `PETSC_DEFAULT` and `nnz` = `NULL` for PETSc to control dynamic memory
2666:   allocation.

2668: .seealso: [](ch_matrices), `Mat`, `MatCreate()`, `MatCreateAIJ()`, `MatSetValues()`, `MatSeqAIJSetColumnIndices()`, `MatCreateSeqAIJWithArrays()`, `MATSEQAIJHIPSPARSE`, `MATAIJHIPSPARSE`
2669: @*/
2670: PetscErrorCode MatCreateSeqAIJHIPSPARSE(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt nz, const PetscInt nnz[], Mat *A)
2671: {
2672:   return MatSeqAIJHIPSPARSE_CUPM_t::CreateSeqAIJ(comm, m, n, nz, nnz, A);
2673: }

2675: static PetscErrorCode MatDestroy_SeqAIJHIPSPARSE(Mat A)
2676: {
2677:   return MatSeqAIJHIPSPARSE_CUPM_t::Destroy(A);
2678: }

2680: static PetscErrorCode MatDuplicate_SeqAIJHIPSPARSE(Mat A, MatDuplicateOption cpvalues, Mat *B)
2681: {
2682:   PetscFunctionBegin;
2683:   PetscCall(MatSeqAIJHIPSPARSE_CUPM_t::Duplicate(A, cpvalues, B));
2684:   PetscFunctionReturn(PETSC_SUCCESS);
2685: }

2687: static PetscErrorCode MatAXPY_SeqAIJHIPSPARSE(Mat Y, PetscScalar a, Mat X, MatStructure str)
2688: {
2689:   Mat_SeqAIJ          *x = (Mat_SeqAIJ *)X->data, *y = (Mat_SeqAIJ *)Y->data;
2690:   Mat_SeqAIJHIPSPARSE *cy;
2691:   Mat_SeqAIJHIPSPARSE *cx;
2692:   CsrMatrix           *csry, *csrx;

2694:   PetscFunctionBegin;
2695:   cy = (Mat_SeqAIJHIPSPARSE *)Y->spptr;
2696:   cx = (Mat_SeqAIJHIPSPARSE *)X->spptr;
2697:   if (X->ops->axpy != Y->ops->axpy) {
2698:     PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(Y, PETSC_FALSE));
2699:     PetscCall(MatAXPY_SeqAIJ(Y, a, X, str));
2700:     PetscFunctionReturn(PETSC_SUCCESS);
2701:   }
2702:   /* if we are here, it means both matrices are bound to GPU */
2703:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(Y));
2704:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(X));
2705:   PetscCheck(cy->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)Y), PETSC_ERR_GPU, "only MAT_HIPSPARSE_CSR supported");
2706:   PetscCheck(cx->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)X), PETSC_ERR_GPU, "only MAT_HIPSPARSE_CSR supported");
2707:   csry = (CsrMatrix *)cy->mat->mat;
2708:   csrx = (CsrMatrix *)cx->mat->mat;
2709:   /* see if we can turn this into a cublas axpy */
2710:   if (str != SAME_NONZERO_PATTERN && x->nz == y->nz && !x->compressedrow.use && !y->compressedrow.use) {
2711:     bool eq = thrust::equal(thrust::device, csry->row_offsets->begin(), csry->row_offsets->end(), csrx->row_offsets->begin());
2712:     if (eq) eq = thrust::equal(thrust::device, csry->column_indices->begin(), csry->column_indices->end(), csrx->column_indices->begin());
2713:     if (eq) str = SAME_NONZERO_PATTERN;
2714:   }
2715:   /* spgeam is buggy with one column */
2716:   if (Y->cmap->n == 1 && str != SAME_NONZERO_PATTERN) str = DIFFERENT_NONZERO_PATTERN;

2718: #if !PetscDefined(USE_64BIT_INDICES) // hipsparseScsrgeam2 etc. do not support 64bit indices
2719:   if (str == SUBSET_NONZERO_PATTERN) {
2720:     PetscScalar       *ay, b = 1.0;
2721:     const PetscScalar *ax;
2722:     size_t             bufferSize;
2723:     void              *buffer;

2725:     PetscCall(MatSeqAIJHIPSPARSEGetArrayRead(X, &ax));
2726:     PetscCall(MatSeqAIJHIPSPARSEGetArray(Y, &ay));
2727:     PetscCallHIPSPARSE(hipsparseSetPointerMode(cy->handle, HIPSPARSE_POINTER_MODE_HOST));
2728:     PetscCallHIPSPARSE(hipsparse_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(),
2729:                                                        csry->column_indices->data().get(), cy->mat->descr, ay, csry->row_offsets->data().get(), csry->column_indices->data().get(), &bufferSize));
2730:     PetscCallHIP(hipMalloc(&buffer, bufferSize));
2731:     PetscCall(PetscLogGpuTimeBegin());
2732:     PetscCallHIPSPARSE(hipsparse_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(),
2733:                                             csry->column_indices->data().get(), cy->mat->descr, ay, csry->row_offsets->data().get(), csry->column_indices->data().get(), buffer));
2734:     PetscCall(PetscLogGpuFlops(x->nz + y->nz));
2735:     PetscCall(PetscLogGpuTimeEnd());
2736:     PetscCallHIP(hipFree(buffer));

2738:     PetscCallHIPSPARSE(hipsparseSetPointerMode(cy->handle, HIPSPARSE_POINTER_MODE_DEVICE));
2739:     PetscCall(MatSeqAIJHIPSPARSERestoreArrayRead(X, &ax));
2740:     PetscCall(MatSeqAIJHIPSPARSERestoreArray(Y, &ay));
2741:   } else
2742: #endif
2743:     if (str == SAME_NONZERO_PATTERN) {
2744:     PetscCall(MatSeqAIJHIPSPARSE_CUPM_t::AXPY_SameNZ(Y, a, X));
2745:   } else {
2746:     PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(Y, PETSC_FALSE));
2747:     PetscCall(MatAXPY_SeqAIJ(Y, a, X, str));
2748:   }
2749:   PetscFunctionReturn(PETSC_SUCCESS);
2750: }

2752: static PetscErrorCode MatScale_SeqAIJHIPSPARSE(Mat Y, PetscScalar a)
2753: {
2754:   PetscFunctionBegin;
2755:   PetscCall(MatSeqAIJHIPSPARSE_CUPM_t::Scale(Y, a));
2756:   PetscFunctionReturn(PETSC_SUCCESS);
2757: }

2759: static PetscErrorCode MatGetDiagonal_SeqAIJHIPSPARSE(Mat A, Vec diag)
2760: {
2761:   PetscFunctionBegin;
2762:   PetscCall(MatSeqAIJHIPSPARSE_CUPM_t::GetDiagonal(A, diag));
2763:   PetscFunctionReturn(PETSC_SUCCESS);
2764: }

2766: static PetscErrorCode MatDiagonalScale_SeqAIJHIPSPARSE(Mat A, Vec ll, Vec rr)
2767: {
2768:   PetscFunctionBegin;
2769:   PetscCall(MatSeqAIJHIPSPARSE_CUPM_t::DiagonalScale(A, ll, rr));
2770:   PetscFunctionReturn(PETSC_SUCCESS);
2771: }

2773: static PetscErrorCode MatZeroEntries_SeqAIJHIPSPARSE(Mat A)
2774: {
2775:   PetscFunctionBegin;
2776:   PetscCall(MatSeqAIJHIPSPARSE_CUPM_t::ZeroEntries(A));
2777:   PetscFunctionReturn(PETSC_SUCCESS);
2778: }

2780: static PetscErrorCode MatGetCurrentMemType_SeqAIJHIPSPARSE(Mat A, PetscMemType *m)
2781: {
2782:   PetscFunctionBegin;
2783:   PetscCall(MatSeqAIJHIPSPARSE_CUPM_t::GetCurrentMemType(A, m));
2784:   PetscFunctionReturn(PETSC_SUCCESS);
2785: }

2787: static PetscErrorCode MatBindToCPU_SeqAIJHIPSPARSE(Mat A, PetscBool flg)
2788: {
2789:   Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;

2791:   PetscFunctionBegin;
2792:   if (A->factortype != MAT_FACTOR_NONE) {
2793:     A->boundtocpu = flg;
2794:     PetscFunctionReturn(PETSC_SUCCESS);
2795:   }
2796:   if (flg) {
2797:     PetscCall(MatSeqAIJHIPSPARSECopyFromGPU(A));

2799:     A->ops->scale                     = MatScale_SeqAIJ;
2800:     A->ops->getdiagonal               = MatGetDiagonal_SeqAIJ;
2801:     A->ops->diagonalscale             = MatDiagonalScale_SeqAIJ;
2802:     A->ops->axpy                      = MatAXPY_SeqAIJ;
2803:     A->ops->zeroentries               = MatZeroEntries_SeqAIJ;
2804:     A->ops->mult                      = MatMult_SeqAIJ;
2805:     A->ops->multadd                   = MatMultAdd_SeqAIJ;
2806:     A->ops->multtranspose             = MatMultTranspose_SeqAIJ;
2807:     A->ops->multtransposeadd          = MatMultTransposeAdd_SeqAIJ;
2808:     A->ops->multhermitiantranspose    = NULL;
2809:     A->ops->multhermitiantransposeadd = NULL;
2810:     A->ops->productsetfromoptions     = MatProductSetFromOptions_SeqAIJ;
2811:     A->ops->getcurrentmemtype         = NULL;
2812:     PetscCall(PetscMemzero(a->ops, sizeof(Mat_SeqAIJOps)));
2813:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSeqAIJCopySubArray_C", NULL));
2814:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqdensehip_C", NULL));
2815:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqdense_C", NULL));
2816:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetPreallocationCOO_C", NULL));
2817:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetValuesCOO_C", NULL));
2818:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqaijhipsparse_C", NULL));
2819:   } else {
2820:     A->ops->scale                     = MatScale_SeqAIJHIPSPARSE;
2821:     A->ops->getdiagonal               = MatGetDiagonal_SeqAIJHIPSPARSE;
2822:     A->ops->diagonalscale             = MatDiagonalScale_SeqAIJHIPSPARSE;
2823:     A->ops->axpy                      = MatAXPY_SeqAIJHIPSPARSE;
2824:     A->ops->zeroentries               = MatZeroEntries_SeqAIJHIPSPARSE;
2825:     A->ops->mult                      = MatMult_SeqAIJHIPSPARSE;
2826:     A->ops->multadd                   = MatMultAdd_SeqAIJHIPSPARSE;
2827:     A->ops->multtranspose             = MatMultTranspose_SeqAIJHIPSPARSE;
2828:     A->ops->multtransposeadd          = MatMultTransposeAdd_SeqAIJHIPSPARSE;
2829:     A->ops->multhermitiantranspose    = MatMultHermitianTranspose_SeqAIJHIPSPARSE;
2830:     A->ops->multhermitiantransposeadd = MatMultHermitianTransposeAdd_SeqAIJHIPSPARSE;
2831:     A->ops->productsetfromoptions     = MatProductSetFromOptions_SeqAIJHIPSPARSE;
2832:     A->ops->getcurrentmemtype         = MatGetCurrentMemType_SeqAIJHIPSPARSE;
2833:     a->ops->getarray                  = MatSeqAIJGetArray_SeqAIJHIPSPARSE;
2834:     a->ops->restorearray              = MatSeqAIJRestoreArray_SeqAIJHIPSPARSE;
2835:     a->ops->getarrayread              = MatSeqAIJGetArrayRead_SeqAIJHIPSPARSE;
2836:     a->ops->restorearrayread          = MatSeqAIJRestoreArrayRead_SeqAIJHIPSPARSE;
2837:     a->ops->getarraywrite             = MatSeqAIJGetArrayWrite_SeqAIJHIPSPARSE;
2838:     a->ops->restorearraywrite         = MatSeqAIJRestoreArrayWrite_SeqAIJHIPSPARSE;
2839:     a->ops->getcsrandmemtype          = MatSeqAIJGetCSRAndMemType_SeqAIJHIPSPARSE;

2841:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSeqAIJCopySubArray_C", MatSeqAIJCopySubArray_SeqAIJHIPSPARSE));
2842:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqdensehip_C", MatProductSetFromOptions_SeqAIJHIPSPARSE));
2843:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqdense_C", MatProductSetFromOptions_SeqAIJHIPSPARSE));
2844:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetPreallocationCOO_C", MatSetPreallocationCOO_SeqAIJHIPSPARSE));
2845:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetValuesCOO_C", MatSetValuesCOO_SeqAIJHIPSPARSE));
2846:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqaijhipsparse_C", MatProductSetFromOptions_SeqAIJHIPSPARSE));
2847:   }
2848:   A->boundtocpu = flg;
2849:   a->inode.use  = (flg && a->inode.size_csr) ? PETSC_TRUE : PETSC_FALSE;
2850:   PetscFunctionReturn(PETSC_SUCCESS);
2851: }

2853: PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJHIPSPARSE(Mat A, MatType mtype, MatReuse reuse, Mat *newmat)
2854: {
2855:   Mat B;

2857:   PetscFunctionBegin;
2858:   PetscCall(PetscDeviceInitialize(PETSC_DEVICE_HIP)); /* first use of HIPSPARSE may be via MatConvert */
2859:   if (reuse == MAT_INITIAL_MATRIX) {
2860:     PetscCall(MatDuplicate(A, MAT_COPY_VALUES, newmat));
2861:   } else if (reuse == MAT_REUSE_MATRIX) {
2862:     PetscCall(MatCopy(A, *newmat, SAME_NONZERO_PATTERN));
2863:   }
2864:   B = *newmat;
2865:   PetscCall(PetscFree(B->defaultvectype));
2866:   PetscCall(PetscStrallocpy(VECHIP, &B->defaultvectype));
2867:   if (reuse != MAT_REUSE_MATRIX && !B->spptr) {
2868:     if (B->factortype == MAT_FACTOR_NONE) {
2869:       Mat_SeqAIJHIPSPARSE *spptr;
2870:       PetscCall(PetscNew(&spptr));
2871:       PetscCallHIPSPARSE(hipsparseCreate(&spptr->handle));
2872:       PetscCallHIPSPARSE(hipsparseSetStream(spptr->handle, PetscDefaultHipStream));
2873:       spptr->format  = MAT_HIPSPARSE_CSR;
2874:       spptr->spmvAlg = HIPSPARSE_SPMV_ALG_DEFAULT;
2875:       spptr->spmmAlg = HIPSPARSE_SPMM_CSR_ALG1; /* default, only support column-major dense matrix B */
2876:       //spptr->csr2cscAlg = HIPSPARSE_CSR2CSC_ALG1;

2878:       B->spptr = spptr;
2879:     } else {
2880:       Mat_SeqAIJHIPSPARSETriFactors *spptr;

2882:       PetscCall(PetscNew(&spptr));
2883:       PetscCallHIPSPARSE(hipsparseCreate(&spptr->handle));
2884:       PetscCallHIPSPARSE(hipsparseSetStream(spptr->handle, PetscDefaultHipStream));
2885:       B->spptr = spptr;
2886:     }
2887:     B->offloadmask = PETSC_OFFLOAD_UNALLOCATED;
2888:   }
2889:   B->ops->assemblyend       = MatAssemblyEnd_SeqAIJHIPSPARSE;
2890:   B->ops->destroy           = MatDestroy_SeqAIJHIPSPARSE;
2891:   B->ops->setoption         = MatSetOption_SeqAIJHIPSPARSE;
2892:   B->ops->setfromoptions    = MatSetFromOptions_SeqAIJHIPSPARSE;
2893:   B->ops->bindtocpu         = MatBindToCPU_SeqAIJHIPSPARSE;
2894:   B->ops->duplicate         = MatDuplicate_SeqAIJHIPSPARSE;
2895:   B->ops->getcurrentmemtype = MatGetCurrentMemType_SeqAIJHIPSPARSE;

2897:   PetscCall(MatBindToCPU_SeqAIJHIPSPARSE(B, PETSC_FALSE));
2898:   PetscCall(PetscObjectChangeTypeName((PetscObject)B, MATSEQAIJHIPSPARSE));
2899:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatHIPSPARSESetFormat_C", MatHIPSPARSESetFormat_SeqAIJHIPSPARSE));
2900: #if PetscDefined(HAVE_HYPRE)
2901:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_seqaijhipsparse_hypre_C", MatConvert_AIJ_HYPRE));
2902: #endif
2903:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatHIPSPARSESetUseCPUSolve_C", MatHIPSPARSESetUseCPUSolve_SeqAIJHIPSPARSE));
2904:   PetscFunctionReturn(PETSC_SUCCESS);
2905: }

2907: PETSC_EXTERN PetscErrorCode MatCreate_SeqAIJHIPSPARSE(Mat B)
2908: {
2909:   PetscFunctionBegin;
2910:   PetscCall(MatCreate_SeqAIJ(B));
2911:   PetscCall(MatConvert_SeqAIJ_SeqAIJHIPSPARSE(B, MATSEQAIJHIPSPARSE, MAT_INPLACE_MATRIX, &B));
2912:   PetscFunctionReturn(PETSC_SUCCESS);
2913: }

2915: /*MC
2916:   MATSEQAIJHIPSPARSE - MATAIJHIPSPARSE = "(seq)aijhipsparse" - A matrix type to be used for sparse matrices on AMD GPUs

2918:   A matrix type whose data resides on AMD GPUs. These matrices can be in either CSR, ELL, or Hybrid format.
2919:   All matrix calculations are performed on AMD/NVIDIA GPUs using the HIPSPARSE library.

2921:   Options Database Keys:
2922: + -mat_type aijhipsparse                           - sets the matrix type to `MATSEQAIJHIPSPARSE`
2923: . -mat_hipsparse_storage_format (csr|ell|hyb)      - sets the storage format of matrices (for `MatMult()` and factors in `MatSolve()`).
2924: . -mat_hipsparse_mult_storage_format (csr|ell|hyb) - sets the storage format of matrices (for `MatMult()`)
2925: - -mat_hipsparse_use_cpu_solve                     - do `MatSolve()` on the CPU

2927:   Level: beginner

2929: .seealso: [](ch_matrices), `Mat`, `MatCreateSeqAIJHIPSPARSE()`, `MATAIJHIPSPARSE`, `MatCreateAIJHIPSPARSE()`, `MatHIPSPARSESetFormat()`, `MatHIPSPARSEStorageFormat`, `MatHIPSPARSEFormatOperation`
2930: M*/

2932: PETSC_INTERN PetscErrorCode MatSolverTypeRegister_HIPSPARSE(void)
2933: {
2934:   PetscFunctionBegin;
2935:   PetscCall(MatSolverTypeRegister(MATSOLVERHIPSPARSE, MATSEQAIJHIPSPARSE, MAT_FACTOR_LU, MatGetFactor_seqaijhipsparse_hipsparse));
2936:   PetscCall(MatSolverTypeRegister(MATSOLVERHIPSPARSE, MATSEQAIJHIPSPARSE, MAT_FACTOR_CHOLESKY, MatGetFactor_seqaijhipsparse_hipsparse));
2937:   PetscCall(MatSolverTypeRegister(MATSOLVERHIPSPARSE, MATSEQAIJHIPSPARSE, MAT_FACTOR_ILU, MatGetFactor_seqaijhipsparse_hipsparse));
2938:   PetscCall(MatSolverTypeRegister(MATSOLVERHIPSPARSE, MATSEQAIJHIPSPARSE, MAT_FACTOR_ICC, MatGetFactor_seqaijhipsparse_hipsparse));
2939:   PetscFunctionReturn(PETSC_SUCCESS);
2940: }

2942: static PetscErrorCode MatSeqAIJHIPSPARSE_Destroy(Mat mat)
2943: {
2944:   Mat_SeqAIJHIPSPARSE *cusp = static_cast<Mat_SeqAIJHIPSPARSE *>(mat->spptr);

2946:   PetscFunctionBegin;
2947:   if (cusp) {
2948:     PetscCall(MatSeqAIJHIPSPARSEMultStruct_Destroy(&cusp->mat, cusp->format));
2949:     PetscCall(MatSeqAIJHIPSPARSEMultStruct_Destroy(&cusp->matTranspose, cusp->format));
2950:     delete cusp->workVector;
2951:     delete cusp->rowoffsets_gpu;
2952:     delete cusp->csr2csc_i;
2953:     delete cusp->coords;
2954:     if (cusp->handle) PetscCallHIPSPARSE(hipsparseDestroy(cusp->handle));
2955:     PetscCall(PetscFree(mat->spptr));
2956:   }
2957:   PetscFunctionReturn(PETSC_SUCCESS);
2958: }

2960: static PetscErrorCode CsrMatrix_Destroy(CsrMatrix **mat)
2961: {
2962:   PetscFunctionBegin;
2963:   if (*mat) {
2964:     delete (*mat)->values;
2965:     delete (*mat)->column_indices;
2966:     delete (*mat)->row_offsets;
2967:     delete *mat;
2968:     *mat = 0;
2969:   }
2970:   PetscFunctionReturn(PETSC_SUCCESS);
2971: }

2973: static PetscErrorCode MatSeqAIJHIPSPARSEMultStruct_Destroy(Mat_SeqAIJHIPSPARSEMultStruct **matstruct, MatHIPSPARSEStorageFormat format)
2974: {
2975:   CsrMatrix *mat;

2977:   PetscFunctionBegin;
2978:   if (*matstruct) {
2979:     if ((*matstruct)->mat) {
2980: #if !defined(CUDART_VERSION) || CUDART_VERSION < 11000
2981:       if (format == MAT_HIPSPARSE_ELL || format == MAT_HIPSPARSE_HYB) {
2982:         hipsparseHybMat_t hybMat = (hipsparseHybMat_t)(*matstruct)->mat;
2983:         PetscCallHIPSPARSE(hipsparseDestroyHybMat(hybMat));
2984:       } else
2985: #endif
2986:       {
2987:         mat = (CsrMatrix *)(*matstruct)->mat;
2988:         PetscCall(CsrMatrix_Destroy(&mat));
2989:       }
2990:     }
2991:     if ((*matstruct)->descr) PetscCallHIPSPARSE(hipsparseDestroyMatDescr((*matstruct)->descr));
2992:     delete (*matstruct)->cprowIndices;
2993:     PetscCallHIP(hipFree((*matstruct)->alpha_one));
2994:     PetscCallHIP(hipFree((*matstruct)->beta_zero));
2995:     PetscCallHIP(hipFree((*matstruct)->beta_one));

2997:     Mat_SeqAIJHIPSPARSEMultStruct *mdata = *matstruct;
2998:     if (mdata->matDescr) PetscCallHIPSPARSE(hipsparseDestroySpMat(mdata->matDescr));
2999:     for (int i = 0; i < 3; i++) {
3000:       if (mdata->hipSpMV[i].initialized) {
3001:         PetscCallHIP(hipFree(mdata->hipSpMV[i].spmvBuffer));
3002:         PetscCallHIPSPARSE(hipsparseDestroyDnVec(mdata->hipSpMV[i].vecXDescr));
3003:         PetscCallHIPSPARSE(hipsparseDestroyDnVec(mdata->hipSpMV[i].vecYDescr));
3004:         PetscCallHIPSPARSE(hipsparseDestroySpMat(mdata->hipSpMV[i].matDescr));
3005:       }
3006:       if (mdata->matDescr_SpMM[i]) PetscCallHIPSPARSE(hipsparseDestroySpMat(mdata->matDescr_SpMM[i]));
3007:     }
3008:     delete *matstruct;
3009:     *matstruct = NULL;
3010:   }
3011:   PetscFunctionReturn(PETSC_SUCCESS);
3012: }

3014: PetscErrorCode MatSeqAIJHIPSPARSETriFactors_Reset(Mat_SeqAIJHIPSPARSETriFactors_p *trifactors)
3015: {
3016:   Mat_SeqAIJHIPSPARSETriFactors *fs = *trifactors;

3018:   PetscFunctionBegin;
3019:   if (fs) {
3020:     delete fs->rpermIndices;
3021:     delete fs->cpermIndices;
3022:     fs->rpermIndices  = NULL;
3023:     fs->cpermIndices  = NULL;
3024:     fs->init_dev_prop = PETSC_FALSE;
3025:     PetscCallHIP(hipFree(fs->csrRowPtr));
3026:     PetscCallHIP(hipFree(fs->csrColIdx));
3027:     PetscCallHIP(hipFree(fs->csrRowPtr32));
3028:     PetscCallHIP(hipFree(fs->csrColIdx32));
3029:     PetscCallHIP(hipFree(fs->csrVal));
3030:     PetscCallHIP(hipFree(fs->diag));
3031:     PetscCallHIP(hipFree(fs->X));
3032:     PetscCallHIP(hipFree(fs->Y));
3033:     // PetscCallHIP(hipFree(fs->factBuffer_M)); /* No needed since factBuffer_M shares with one of spsvBuffer_L/U */
3034:     PetscCallHIP(hipFree(fs->spsvBuffer_L));
3035:     PetscCallHIP(hipFree(fs->spsvBuffer_U));
3036:     PetscCallHIP(hipFree(fs->spsvBuffer_Lt));
3037:     PetscCallHIP(hipFree(fs->spsvBuffer_Ut));
3038:     PetscCallHIPSPARSE(hipsparseDestroyMatDescr(fs->matDescr_M));
3039:     if (fs->spMatDescr_L) PetscCallHIPSPARSE(hipsparseDestroySpMat(fs->spMatDescr_L));
3040:     if (fs->spMatDescr_U) PetscCallHIPSPARSE(hipsparseDestroySpMat(fs->spMatDescr_U));
3041:     if (fs->spMatDescr_Lt) PetscCallHIPSPARSE(hipsparseDestroySpMat(fs->spMatDescr_Lt));
3042:     if (fs->spMatDescr_Ut) PetscCallHIPSPARSE(hipsparseDestroySpMat(fs->spMatDescr_Ut));
3043:     fs->spMatDescr_Lt = fs->spMatDescr_Ut = NULL;
3044:     PetscCallHIPSPARSE(hipsparseSpSV_destroyDescr(fs->spsvDescr_L));
3045:     PetscCallHIPSPARSE(hipsparseSpSV_destroyDescr(fs->spsvDescr_Lt));
3046:     PetscCallHIPSPARSE(hipsparseSpSV_destroyDescr(fs->spsvDescr_U));
3047:     PetscCallHIPSPARSE(hipsparseSpSV_destroyDescr(fs->spsvDescr_Ut));
3048:     if (fs->dnVecDescr_X) PetscCallHIPSPARSE(hipsparseDestroyDnVec(fs->dnVecDescr_X));
3049:     if (fs->dnVecDescr_Y) PetscCallHIPSPARSE(hipsparseDestroyDnVec(fs->dnVecDescr_Y));
3050:     PetscCallHIPSPARSE(hipsparseDestroyCsrilu02Info(fs->ilu0Info_M));
3051:     PetscCallHIPSPARSE(hipsparseDestroyCsric02Info(fs->ic0Info_M));
3052:     PetscCall(PetscFree(fs->csrRowPtr_h));
3053:     PetscCall(PetscFree(fs->csrVal_h));
3054:     PetscCall(PetscFree(fs->diag_h));
3055:     fs->createdTransposeSpSVDescr    = PETSC_FALSE;
3056:     fs->updatedTransposeSpSVAnalysis = PETSC_FALSE;
3057:   }
3058:   PetscFunctionReturn(PETSC_SUCCESS);
3059: }

3061: static PetscErrorCode MatSeqAIJHIPSPARSETriFactors_Destroy(Mat_SeqAIJHIPSPARSETriFactors **trifactors)
3062: {
3063:   PetscFunctionBegin;
3064:   if (*trifactors) {
3065:     PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(trifactors));
3066:     PetscCallHIPSPARSE(hipsparseDestroy((*trifactors)->handle));
3067:     PetscCall(PetscFree(*trifactors));
3068:   }
3069:   PetscFunctionReturn(PETSC_SUCCESS);
3070: }

3072: static PetscErrorCode MatSeqAIJHIPSPARSEInvalidateTranspose(Mat A, PetscBool destroy)
3073: {
3074:   Mat_SeqAIJHIPSPARSE *cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;

3076:   PetscFunctionBegin;
3077:   PetscCheckTypeName(A, MATSEQAIJHIPSPARSE);
3078:   if (!cusp) PetscFunctionReturn(PETSC_SUCCESS);
3079:   if (destroy) {
3080:     PetscCall(MatSeqAIJHIPSPARSEMultStruct_Destroy(&cusp->matTranspose, cusp->format));
3081:     delete cusp->csr2csc_i;
3082:     cusp->csr2csc_i = NULL;
3083:   }
3084:   A->transupdated = PETSC_FALSE;
3085:   PetscFunctionReturn(PETSC_SUCCESS);
3086: }

3088: static PetscErrorCode MatSetPreallocationCOO_SeqAIJHIPSPARSE(Mat mat, PetscCount coo_n, PetscInt coo_i[], PetscInt coo_j[])
3089: {
3090:   PetscFunctionBegin;
3091:   PetscCall(MatSeqAIJHIPSPARSE_CUPM_t::SetPreallocationCOO(mat, coo_n, coo_i, coo_j));
3092:   PetscFunctionReturn(PETSC_SUCCESS);
3093: }

3095: static PetscErrorCode MatSetValuesCOO_SeqAIJHIPSPARSE(Mat A, const PetscScalar v[], InsertMode imode)
3096: {
3097:   PetscFunctionBegin;
3098:   PetscCall(MatSeqAIJHIPSPARSE_CUPM_t::SetValuesCOO(A, v, imode));
3099:   PetscFunctionReturn(PETSC_SUCCESS);
3100: }

3102: /*@
3103:   MatSeqAIJHIPSPARSEGetIJ - returns the device row storage `i` and `j` indices for `MATSEQAIJHIPSPARSE` matrices.

3105:   Not Collective

3107:   Input Parameters:
3108: + A          - the matrix
3109: - compressed - `PETSC_TRUE` or `PETSC_FALSE` indicating the matrix data structure should be always returned in compressed form

3111:   Output Parameters:
3112: + i - the CSR row pointers
3113: - j - the CSR column indices

3115:   Level: developer

3117:   Note:
3118:   When compressed is true, the CSR structure does not contain empty rows

3120: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSERestoreIJ()`, `MatSeqAIJHIPSPARSEGetArrayRead()`
3121: @*/
3122: PetscErrorCode MatSeqAIJHIPSPARSEGetIJ(Mat A, PetscBool compressed, const PetscInt *i[], const PetscInt *j[])
3123: {
3124:   PetscFunctionBegin;
3125:   PetscCall(MatSeqAIJHIPSPARSE_CUPM_t::GetIJ(A, compressed, i, j));
3126:   PetscFunctionReturn(PETSC_SUCCESS);
3127: }

3129: /*@
3130:   MatSeqAIJHIPSPARSERestoreIJ - restore the device row storage `i` and `j` indices obtained with `MatSeqAIJHIPSPARSEGetIJ()`

3132:   Not Collective

3134:   Input Parameters:
3135: + A          - the matrix
3136: . compressed - `PETSC_TRUE` or `PETSC_FALSE` indicating the matrix data structure should be always returned in compressed form
3137: . i          - the CSR row pointers
3138: - j          - the CSR column indices

3140:   Level: developer

3142: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetIJ()`
3143: @*/
3144: PetscErrorCode MatSeqAIJHIPSPARSERestoreIJ(Mat A, PetscBool compressed, const PetscInt *i[], const PetscInt *j[])
3145: {
3146:   PetscFunctionBegin;
3147:   PetscCall(MatSeqAIJHIPSPARSE_CUPM_t::RestoreIJ(A, compressed, i, j));
3148:   PetscFunctionReturn(PETSC_SUCCESS);
3149: }

3151: /*@
3152:   MatSeqAIJHIPSPARSEGetArrayRead - gives read-only access to the array where the device data for a `MATSEQAIJHIPSPARSE` matrix is stored

3154:   Not Collective

3156:   Input Parameter:
3157: . A - a `MATSEQAIJHIPSPARSE` matrix

3159:   Output Parameter:
3160: . a - pointer to the device data

3162:   Level: developer

3164:   Note:
3165:   May trigger host-device copies if the up-to-date matrix data is on host

3167: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetArray()`, `MatSeqAIJHIPSPARSEGetArrayWrite()`, `MatSeqAIJHIPSPARSERestoreArrayRead()`
3168: @*/
3169: PetscErrorCode MatSeqAIJHIPSPARSEGetArrayRead(Mat A, const PetscScalar *a[])
3170: {
3171:   return MatSeqAIJHIPSPARSE_CUPM_t::GetArrayRead(A, a);
3172: }

3174: /*@
3175:   MatSeqAIJHIPSPARSERestoreArrayRead - restore the read-only access array obtained from `MatSeqAIJHIPSPARSEGetArrayRead()`

3177:   Not Collective

3179:   Input Parameters:
3180: + A - a `MATSEQAIJHIPSPARSE` matrix
3181: - a - pointer to the device data

3183:   Level: developer

3185: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetArrayRead()`
3186: @*/
3187: PetscErrorCode MatSeqAIJHIPSPARSERestoreArrayRead(Mat A, const PetscScalar *a[])
3188: {
3189:   return MatSeqAIJHIPSPARSE_CUPM_t::RestoreArrayRead(A, a);
3190: }

3192: /*@
3193:   MatSeqAIJHIPSPARSEGetArray - gives read-write access to the array where the device data for a `MATSEQAIJHIPSPARSE` matrix is stored

3195:   Not Collective

3197:   Input Parameter:
3198: . A - a `MATSEQAIJHIPSPARSE` matrix

3200:   Output Parameter:
3201: . a - pointer to the device data

3203:   Level: developer

3205:   Note:
3206:   May trigger host-device copies if up-to-date matrix data is on host

3208: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetArrayRead()`, `MatSeqAIJHIPSPARSEGetArrayWrite()`, `MatSeqAIJHIPSPARSERestoreArray()`
3209: @*/
3210: PetscErrorCode MatSeqAIJHIPSPARSEGetArray(Mat A, PetscScalar *a[])
3211: {
3212:   return MatSeqAIJHIPSPARSE_CUPM_t::GetArray(A, a);
3213: }
3214: /*@
3215:   MatSeqAIJHIPSPARSERestoreArray - restore the read-write access array obtained from `MatSeqAIJHIPSPARSEGetArray()`

3217:   Not Collective

3219:   Input Parameters:
3220: + A - a `MATSEQAIJHIPSPARSE` matrix
3221: - a - pointer to the device data

3223:   Level: developer

3225: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetArray()`
3226: @*/
3227: PetscErrorCode MatSeqAIJHIPSPARSERestoreArray(Mat A, PetscScalar *a[])
3228: {
3229:   return MatSeqAIJHIPSPARSE_CUPM_t::RestoreArray(A, a);
3230: }

3232: /*@
3233:   MatSeqAIJHIPSPARSEGetArrayWrite - gives write access to the array where the device data for a `MATSEQAIJHIPSPARSE` matrix is stored

3235:   Not Collective

3237:   Input Parameter:
3238: . A - a `MATSEQAIJHIPSPARSE` matrix

3240:   Output Parameter:
3241: . a - pointer to the device data

3243:   Level: developer

3245:   Note:
3246:   Does not trigger host-device copies and flags data validity on the GPU

3248: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetArray()`, `MatSeqAIJHIPSPARSEGetArrayRead()`, `MatSeqAIJHIPSPARSERestoreArrayWrite()`
3249: @*/
3250: PetscErrorCode MatSeqAIJHIPSPARSEGetArrayWrite(Mat A, PetscScalar *a[])
3251: {
3252:   return MatSeqAIJHIPSPARSE_CUPM_t::GetArrayWrite(A, a);
3253: }

3255: /*@
3256:   MatSeqAIJHIPSPARSERestoreArrayWrite - restore the write-only access array obtained from `MatSeqAIJHIPSPARSEGetArrayWrite()`

3258:   Not Collective

3260:   Input Parameters:
3261: + A - a `MATSEQAIJHIPSPARSE` matrix
3262: - a - pointer to the device data

3264:   Level: developer

3266: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetArrayWrite()`
3267: @*/
3268: PetscErrorCode MatSeqAIJHIPSPARSERestoreArrayWrite(Mat A, PetscScalar *a[])
3269: {
3270:   return MatSeqAIJHIPSPARSE_CUPM_t::RestoreArrayWrite(A, a);
3271: }

3273: struct IJCompare4 {
3274:   __host__ __device__ inline bool operator()(const thrust::tuple<int, int, PetscScalar, int> &t1, const thrust::tuple<int, int, PetscScalar, int> &t2)
3275:   {
3276:     if (thrust::get<0>(t1) < thrust::get<0>(t2)) return true;
3277:     if (thrust::get<0>(t1) == thrust::get<0>(t2)) return thrust::get<1>(t1) < thrust::get<1>(t2);
3278:     return false;
3279:   }
3280: };

3282: struct Shift {
3283:   int _shift;

3285:   Shift(int shift) : _shift(shift) { }
3286:   __host__ __device__ inline int operator()(const int &c) { return c + _shift; }
3287: };

3289: /* merges two SeqAIJHIPSPARSE matrices A, B by concatenating their rows. [A';B']' operation in MATLAB notation */
3290: PetscErrorCode MatSeqAIJHIPSPARSEMergeMats(Mat A, Mat B, MatReuse reuse, Mat *C)
3291: {
3292:   Mat_SeqAIJ                    *a = (Mat_SeqAIJ *)A->data, *b = (Mat_SeqAIJ *)B->data, *c;
3293:   Mat_SeqAIJHIPSPARSE           *Acusp = (Mat_SeqAIJHIPSPARSE *)A->spptr, *Bcusp = (Mat_SeqAIJHIPSPARSE *)B->spptr, *Ccusp;
3294:   Mat_SeqAIJHIPSPARSEMultStruct *Cmat;
3295:   CsrMatrix                     *Acsr, *Bcsr, *Ccsr;
3296:   PetscInt                       Annz, Bnnz;
3297:   PetscInt                       i, m, n, zero = 0;

3299:   PetscFunctionBegin;
3302:   PetscAssertPointer(C, 4);
3303:   PetscCheckTypeName(A, MATSEQAIJHIPSPARSE);
3304:   PetscCheckTypeName(B, MATSEQAIJHIPSPARSE);
3305:   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);
3306:   PetscCheck(reuse != MAT_INPLACE_MATRIX, PETSC_COMM_SELF, PETSC_ERR_SUP, "MAT_INPLACE_MATRIX not supported");
3307:   PetscCheck(Acusp->format != MAT_HIPSPARSE_ELL && Acusp->format != MAT_HIPSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
3308:   PetscCheck(Bcusp->format != MAT_HIPSPARSE_ELL && Bcusp->format != MAT_HIPSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
3309:   if (reuse == MAT_INITIAL_MATRIX) {
3310:     m = A->rmap->n;
3311:     n = A->cmap->n + B->cmap->n;
3312:     PetscCall(MatCreate(PETSC_COMM_SELF, C));
3313:     PetscCall(MatSetSizes(*C, m, n, m, n));
3314:     PetscCall(MatSetType(*C, MATSEQAIJHIPSPARSE));
3315:     c                       = (Mat_SeqAIJ *)(*C)->data;
3316:     Ccusp                   = (Mat_SeqAIJHIPSPARSE *)(*C)->spptr;
3317:     Cmat                    = new Mat_SeqAIJHIPSPARSEMultStruct;
3318:     Ccsr                    = new CsrMatrix;
3319:     Cmat->cprowIndices      = NULL;
3320:     c->compressedrow.use    = PETSC_FALSE;
3321:     c->compressedrow.nrows  = 0;
3322:     c->compressedrow.i      = NULL;
3323:     c->compressedrow.rindex = NULL;
3324:     Ccusp->workVector       = NULL;
3325:     Ccusp->nrows            = m;
3326:     Ccusp->mat              = Cmat;
3327:     Ccusp->mat->mat         = Ccsr;
3328:     Ccsr->num_rows          = m;
3329:     Ccsr->num_cols          = n;
3330:     PetscCallHIPSPARSE(hipsparseCreateMatDescr(&Cmat->descr));
3331:     PetscCallHIPSPARSE(hipsparseSetMatIndexBase(Cmat->descr, HIPSPARSE_INDEX_BASE_ZERO));
3332:     PetscCallHIPSPARSE(hipsparseSetMatType(Cmat->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));
3333:     PetscCallHIP(hipMalloc((void **)&Cmat->alpha_one, sizeof(PetscScalar)));
3334:     PetscCallHIP(hipMalloc((void **)&Cmat->beta_zero, sizeof(PetscScalar)));
3335:     PetscCallHIP(hipMalloc((void **)&Cmat->beta_one, sizeof(PetscScalar)));
3336:     PetscCallHIP(hipMemcpy(Cmat->alpha_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
3337:     PetscCallHIP(hipMemcpy(Cmat->beta_zero, &PETSC_HIPSPARSE_ZERO, sizeof(PetscScalar), hipMemcpyHostToDevice));
3338:     PetscCallHIP(hipMemcpy(Cmat->beta_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
3339:     PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
3340:     PetscCall(MatSeqAIJHIPSPARSECopyToGPU(B));
3341:     PetscCheck(Acusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJHIPSPARSEMultStruct");
3342:     PetscCheck(Bcusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJHIPSPARSEMultStruct");

3344:     Acsr                 = (CsrMatrix *)Acusp->mat->mat;
3345:     Bcsr                 = (CsrMatrix *)Bcusp->mat->mat;
3346:     Annz                 = (PetscInt)Acsr->column_indices->size();
3347:     Bnnz                 = (PetscInt)Bcsr->column_indices->size();
3348:     c->nz                = Annz + Bnnz;
3349:     Ccsr->row_offsets    = new THRUSTINTARRAY(m + 1);
3350:     Ccsr->column_indices = new THRUSTINTARRAY(c->nz);
3351:     Ccsr->values         = new THRUSTARRAY(c->nz);
3352:     Ccsr->num_entries    = c->nz;
3353:     Ccusp->coords        = new THRUSTINTARRAY(c->nz);
3354:     if (c->nz) {
3355:       auto            Acoo = new THRUSTINTARRAY(Annz); // initialized with zeros
3356:       auto            Bcoo = new THRUSTINTARRAY(Bnnz);
3357:       auto            Ccoo = new THRUSTINTARRAY(c->nz);
3358:       THRUSTINTARRAY *Aroff, *Broff;

3360:       if (a->compressedrow.use) { /* need full row offset */
3361:         if (!Acusp->rowoffsets_gpu) {
3362:           Acusp->rowoffsets_gpu = new THRUSTINTARRAY(A->rmap->n + 1);
3363:           Acusp->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
3364:           PetscCall(PetscLogCpuToGpu((A->rmap->n + 1) * sizeof(PetscInt)));
3365:         }
3366:         Aroff = Acusp->rowoffsets_gpu;
3367:       } else Aroff = Acsr->row_offsets;
3368:       if (b->compressedrow.use) { /* need full row offset */
3369:         if (!Bcusp->rowoffsets_gpu) {
3370:           Bcusp->rowoffsets_gpu = new THRUSTINTARRAY(B->rmap->n + 1);
3371:           Bcusp->rowoffsets_gpu->assign(b->i, b->i + B->rmap->n + 1);
3372:           PetscCall(PetscLogCpuToGpu((B->rmap->n + 1) * sizeof(PetscInt)));
3373:         }
3374:         Broff = Bcusp->rowoffsets_gpu;
3375:       } else Broff = Bcsr->row_offsets;
3376:       PetscCall(PetscLogGpuTimeBegin());
3377:       // Implement cusparseXcsr2coo() with Thrust, as the former doesn't support 64-bit indices.
3378:       PetscCallThrust(thrust::for_each(thrust::device, thrust::counting_iterator<PetscInt>(0), thrust::counting_iterator<PetscInt>(m), Csr2coo(Aroff->data().get(), Acoo->data().get())));
3379:       PetscCallThrust(thrust::for_each(thrust::device, thrust::counting_iterator<PetscInt>(0), thrust::counting_iterator<PetscInt>(m), Csr2coo(Broff->data().get(), Bcoo->data().get())));

3381:       /* Issues when using bool with large matrices on SUMMIT 10.2.89 */
3382:       auto Aperm = thrust::make_constant_iterator(1);
3383:       auto Bperm = thrust::make_constant_iterator(0);
3384:       auto Bcib  = thrust::make_transform_iterator(Bcsr->column_indices->begin(), Shift(A->cmap->n));
3385:       auto Bcie  = thrust::make_transform_iterator(Bcsr->column_indices->end(), Shift(A->cmap->n));
3386:       auto wPerm = new THRUSTINTARRAY(Annz + Bnnz);
3387:       auto Azb   = thrust::make_zip_iterator(thrust::make_tuple(Acoo->begin(), Acsr->column_indices->begin(), Acsr->values->begin(), Aperm));
3388:       auto Aze   = thrust::make_zip_iterator(thrust::make_tuple(Acoo->end(), Acsr->column_indices->end(), Acsr->values->end(), Aperm));
3389:       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
3390:       auto Bze   = thrust::make_zip_iterator(thrust::make_tuple(Bcoo->end(), Bcie, Bcsr->values->end(), Bperm));
3391:       auto Czb   = thrust::make_zip_iterator(thrust::make_tuple(Ccoo->begin(), Ccsr->column_indices->begin(), Ccsr->values->begin(), wPerm->begin()));
3392:       auto p1    = Ccusp->coords->begin();
3393:       auto p2    = Ccusp->coords->begin();
3394:       thrust::advance(p2, Annz);
3395:       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)
3396:       auto cci = thrust::make_counting_iterator(zero);
3397:       auto cce = thrust::make_counting_iterator(c->nz);
3398: #if THRUST_VERSION < 200800
3399:       auto pred = thrust::identity<int>();
3400: #else
3401:       auto pred = ::internal::identity();
3402: #endif
3403:       PetscCallThrust(thrust::copy_if(thrust::device, cci, cce, wPerm->begin(), p1, pred));
3404:       PetscCallThrust(thrust::remove_copy_if(thrust::device, cci, cce, wPerm->begin(), p2, pred));
3405:       // Implement a simplified hipsparseXcoo2csr() with Thrust (assuming the row indices are already sorted), as the former doesn't support 64-bit indices.
3406:       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()));
3407:       PetscCall(PetscLogGpuTimeEnd());
3408:       delete wPerm;
3409:       delete Acoo;
3410:       delete Bcoo;
3411:       delete Ccoo;
3412:       PetscCallHIPSPARSE(hipsparseCreateCsr(&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, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
3413:       if (A->form_explicit_transpose && B->form_explicit_transpose) { /* if A and B have the transpose, generate C transpose too */
3414:         PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(A));
3415:         PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(B));
3416:         PetscBool                      AT = Acusp->matTranspose ? PETSC_TRUE : PETSC_FALSE, BT = Bcusp->matTranspose ? PETSC_TRUE : PETSC_FALSE;
3417:         Mat_SeqAIJHIPSPARSEMultStruct *CmatT = new Mat_SeqAIJHIPSPARSEMultStruct;
3418:         CsrMatrix                     *CcsrT = new CsrMatrix;
3419:         CsrMatrix                     *AcsrT = AT ? (CsrMatrix *)Acusp->matTranspose->mat : NULL;
3420:         CsrMatrix                     *BcsrT = BT ? (CsrMatrix *)Bcusp->matTranspose->mat : NULL;

3422:         (*C)->form_explicit_transpose = PETSC_TRUE;
3423:         (*C)->transupdated            = PETSC_TRUE;
3424:         Ccusp->rowoffsets_gpu         = NULL;
3425:         CmatT->cprowIndices           = NULL;
3426:         CmatT->mat                    = CcsrT;
3427:         CcsrT->num_rows               = n;
3428:         CcsrT->num_cols               = m;
3429:         CcsrT->num_entries            = c->nz;

3431:         CcsrT->row_offsets    = new THRUSTINTARRAY(n + 1);
3432:         CcsrT->column_indices = new THRUSTINTARRAY(c->nz);
3433:         CcsrT->values         = new THRUSTARRAY(c->nz);

3435:         PetscCall(PetscLogGpuTimeBegin());
3436:         auto rT = CcsrT->row_offsets->begin();
3437:         if (AT) {
3438:           rT = thrust::copy(AcsrT->row_offsets->begin(), AcsrT->row_offsets->end(), rT);
3439:           thrust::advance(rT, -1);
3440:         }
3441:         if (BT) {
3442:           auto titb = thrust::make_transform_iterator(BcsrT->row_offsets->begin(), Shift(a->nz));
3443:           auto tite = thrust::make_transform_iterator(BcsrT->row_offsets->end(), Shift(a->nz));
3444:           thrust::copy(titb, tite, rT);
3445:         }
3446:         auto cT = CcsrT->column_indices->begin();
3447:         if (AT) cT = thrust::copy(AcsrT->column_indices->begin(), AcsrT->column_indices->end(), cT);
3448:         if (BT) thrust::copy(BcsrT->column_indices->begin(), BcsrT->column_indices->end(), cT);
3449:         auto vT = CcsrT->values->begin();
3450:         if (AT) vT = thrust::copy(AcsrT->values->begin(), AcsrT->values->end(), vT);
3451:         if (BT) thrust::copy(BcsrT->values->begin(), BcsrT->values->end(), vT);
3452:         PetscCall(PetscLogGpuTimeEnd());

3454:         PetscCallHIPSPARSE(hipsparseCreateMatDescr(&CmatT->descr));
3455:         PetscCallHIPSPARSE(hipsparseSetMatIndexBase(CmatT->descr, HIPSPARSE_INDEX_BASE_ZERO));
3456:         PetscCallHIPSPARSE(hipsparseSetMatType(CmatT->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));
3457:         PetscCallHIP(hipMalloc((void **)&CmatT->alpha_one, sizeof(PetscScalar)));
3458:         PetscCallHIP(hipMalloc((void **)&CmatT->beta_zero, sizeof(PetscScalar)));
3459:         PetscCallHIP(hipMalloc((void **)&CmatT->beta_one, sizeof(PetscScalar)));
3460:         PetscCallHIP(hipMemcpy(CmatT->alpha_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
3461:         PetscCallHIP(hipMemcpy(CmatT->beta_zero, &PETSC_HIPSPARSE_ZERO, sizeof(PetscScalar), hipMemcpyHostToDevice));
3462:         PetscCallHIP(hipMemcpy(CmatT->beta_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
3463:         PetscCallHIPSPARSE(hipsparseCreateCsr(&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, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
3464:         Ccusp->matTranspose = CmatT;
3465:       }
3466:     }

3468:     c->free_a = PETSC_TRUE;
3469:     PetscCall(PetscShmgetAllocateArray(c->nz, sizeof(PetscInt), (void **)&c->j));
3470:     PetscCall(PetscShmgetAllocateArray(m + 1, sizeof(PetscInt), (void **)&c->i));
3471:     c->free_ij = PETSC_TRUE;
3472:     PetscCallHIP(hipMemcpy(c->i, Ccsr->row_offsets->data().get(), Ccsr->row_offsets->size() * sizeof(PetscInt), hipMemcpyDeviceToHost));
3473:     PetscCallHIP(hipMemcpy(c->j, Ccsr->column_indices->data().get(), Ccsr->column_indices->size() * sizeof(PetscInt), hipMemcpyDeviceToHost));
3474:     PetscCall(PetscLogGpuToCpu((Ccsr->column_indices->size() + Ccsr->row_offsets->size()) * sizeof(PetscInt)));
3475:     PetscCall(PetscMalloc1(m, &c->ilen));
3476:     PetscCall(PetscMalloc1(m, &c->imax));
3477:     c->maxnz         = c->nz;
3478:     c->nonzerorowcnt = 0;
3479:     c->rmax          = 0;
3480:     for (i = 0; i < m; i++) {
3481:       const PetscInt nn = c->i[i + 1] - c->i[i];
3482:       c->ilen[i] = c->imax[i] = nn;
3483:       c->nonzerorowcnt += (PetscInt)!!nn;
3484:       c->rmax = PetscMax(c->rmax, nn);
3485:     }
3486:     PetscCall(PetscMalloc1(c->nz, &c->a));
3487:     (*C)->nonzerostate++;
3488:     PetscCall(PetscLayoutSetUp((*C)->rmap));
3489:     PetscCall(PetscLayoutSetUp((*C)->cmap));
3490:     Ccusp->nonzerostate = (*C)->nonzerostate;
3491:     (*C)->preallocated  = PETSC_TRUE;
3492:   } else {
3493:     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);
3494:     c = (Mat_SeqAIJ *)(*C)->data;
3495:     if (c->nz) {
3496:       Ccusp = (Mat_SeqAIJHIPSPARSE *)(*C)->spptr;
3497:       PetscCheck(Ccusp->coords, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing coords");
3498:       PetscCheck(Ccusp->format != MAT_HIPSPARSE_ELL && Ccusp->format != MAT_HIPSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
3499:       PetscCheck(Ccusp->nonzerostate == (*C)->nonzerostate, PETSC_COMM_SELF, PETSC_ERR_COR, "Wrong nonzerostate");
3500:       PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
3501:       PetscCall(MatSeqAIJHIPSPARSECopyToGPU(B));
3502:       PetscCheck(Acusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJHIPSPARSEMultStruct");
3503:       PetscCheck(Bcusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJHIPSPARSEMultStruct");
3504:       Acsr = (CsrMatrix *)Acusp->mat->mat;
3505:       Bcsr = (CsrMatrix *)Bcusp->mat->mat;
3506:       Ccsr = (CsrMatrix *)Ccusp->mat->mat;
3507:       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());
3508:       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());
3509:       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());
3510:       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);
3511:       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());
3512:       auto pmid = Ccusp->coords->begin();
3513:       thrust::advance(pmid, Acsr->num_entries);
3514:       PetscCall(PetscLogGpuTimeBegin());
3515:       auto zibait = thrust::make_zip_iterator(thrust::make_tuple(Acsr->values->begin(), thrust::make_permutation_iterator(Ccsr->values->begin(), Ccusp->coords->begin())));
3516:       auto zieait = thrust::make_zip_iterator(thrust::make_tuple(Acsr->values->end(), thrust::make_permutation_iterator(Ccsr->values->begin(), pmid)));
3517:       thrust::for_each(zibait, zieait, VecHIPEquals());
3518:       auto zibbit = thrust::make_zip_iterator(thrust::make_tuple(Bcsr->values->begin(), thrust::make_permutation_iterator(Ccsr->values->begin(), pmid)));
3519:       auto ziebit = thrust::make_zip_iterator(thrust::make_tuple(Bcsr->values->end(), thrust::make_permutation_iterator(Ccsr->values->begin(), Ccusp->coords->end())));
3520:       thrust::for_each(zibbit, ziebit, VecHIPEquals());
3521:       PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(*C, PETSC_FALSE));
3522:       if (A->form_explicit_transpose && B->form_explicit_transpose && (*C)->form_explicit_transpose) {
3523:         PetscCheck(Ccusp->matTranspose, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing transpose Mat_SeqAIJHIPSPARSEMultStruct");
3524:         PetscBool  AT = Acusp->matTranspose ? PETSC_TRUE : PETSC_FALSE, BT = Bcusp->matTranspose ? PETSC_TRUE : PETSC_FALSE;
3525:         CsrMatrix *AcsrT = AT ? (CsrMatrix *)Acusp->matTranspose->mat : NULL;
3526:         CsrMatrix *BcsrT = BT ? (CsrMatrix *)Bcusp->matTranspose->mat : NULL;
3527:         CsrMatrix *CcsrT = (CsrMatrix *)Ccusp->matTranspose->mat;
3528:         auto       vT    = CcsrT->values->begin();
3529:         if (AT) vT = thrust::copy(AcsrT->values->begin(), AcsrT->values->end(), vT);
3530:         if (BT) thrust::copy(BcsrT->values->begin(), BcsrT->values->end(), vT);
3531:         (*C)->transupdated = PETSC_TRUE;
3532:       }
3533:       PetscCall(PetscLogGpuTimeEnd());
3534:     }
3535:   }
3536:   PetscCall(PetscObjectStateIncrease((PetscObject)*C));
3537:   (*C)->assembled     = PETSC_TRUE;
3538:   (*C)->was_assembled = PETSC_FALSE;
3539:   (*C)->offloadmask   = PETSC_OFFLOAD_GPU;
3540:   PetscFunctionReturn(PETSC_SUCCESS);
3541: }

3543: static PetscErrorCode MatSeqAIJCopySubArray_SeqAIJHIPSPARSE(Mat A, PetscInt n, const PetscInt idx[], PetscScalar v[])
3544: {
3545:   PetscFunctionBegin;
3546:   PetscCall(MatSeqAIJHIPSPARSE_CUPM_t::CopySubArray(A, n, idx, v));
3547:   PetscFunctionReturn(PETSC_SUCCESS);
3548: }