Actual source code: aijhipsparse.hip.cpp

  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 <thrust/adjacent_difference.h>
 16: #include <thrust/iterator/transform_iterator.h>
 17: #if PETSC_CPP_VERSION >= 14
 18:   #define PETSC_HAVE_THRUST_ASYNC 1
 19:   #include <thrust/async/for_each.h>
 20: #endif
 21: #include <thrust/iterator/constant_iterator.h>
 22: #include <thrust/iterator/discard_iterator.h>
 23: #include <thrust/binary_search.h>
 24: #include <thrust/remove.h>
 25: #include <thrust/sort.h>
 26: #include <thrust/unique.h>

 28: const char *const MatHIPSPARSEStorageFormats[] = {"CSR", "ELL", "HYB", "MatHIPSPARSEStorageFormat", "MAT_HIPSPARSE_", 0};
 29: 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};
 30: const char *const MatHIPSPARSESpMMAlgorithms[] = {"ALG_DEFAULT", "COO_ALG1", "COO_ALG2", "COO_ALG3", "CSR_ALG1", "COO_ALG4", "CSR_ALG2", "hipsparseSpMMAlg_t", "HIPSPARSE_SPMM_", 0};
 31: //const char *const MatHIPSPARSECsr2CscAlgorithms[] = {"INVALID"/*HIPSPARSE does not have enum 0! We created one*/, "ALG1", "ALG2", "hipsparseCsr2CscAlg_t", "HIPSPARSE_CSR2CSC_", 0};

 33: static PetscErrorCode MatICCFactorSymbolic_SeqAIJHIPSPARSE(Mat, Mat, IS, const MatFactorInfo *);
 34: static PetscErrorCode MatCholeskyFactorSymbolic_SeqAIJHIPSPARSE(Mat, Mat, IS, const MatFactorInfo *);
 35: static PetscErrorCode MatCholeskyFactorNumeric_SeqAIJHIPSPARSE(Mat, Mat, const MatFactorInfo *);
 36: static PetscErrorCode MatILUFactorSymbolic_SeqAIJHIPSPARSE(Mat, Mat, IS, IS, const MatFactorInfo *);
 37: static PetscErrorCode MatLUFactorSymbolic_SeqAIJHIPSPARSE(Mat, Mat, IS, IS, const MatFactorInfo *);
 38: static PetscErrorCode MatLUFactorNumeric_SeqAIJHIPSPARSE(Mat, Mat, const MatFactorInfo *);
 39: static PetscErrorCode MatSolve_SeqAIJHIPSPARSE(Mat, Vec, Vec);
 40: static PetscErrorCode MatSolve_SeqAIJHIPSPARSE_NaturalOrdering(Mat, Vec, Vec);
 41: static PetscErrorCode MatSolveTranspose_SeqAIJHIPSPARSE(Mat, Vec, Vec);
 42: static PetscErrorCode MatSolveTranspose_SeqAIJHIPSPARSE_NaturalOrdering(Mat, Vec, Vec);
 43: static PetscErrorCode MatSetFromOptions_SeqAIJHIPSPARSE(Mat, PetscOptionItems *PetscOptionsObject);
 44: static PetscErrorCode MatAXPY_SeqAIJHIPSPARSE(Mat, PetscScalar, Mat, MatStructure);
 45: static PetscErrorCode MatScale_SeqAIJHIPSPARSE(Mat, PetscScalar);
 46: static PetscErrorCode MatMult_SeqAIJHIPSPARSE(Mat, Vec, Vec);
 47: static PetscErrorCode MatMultAdd_SeqAIJHIPSPARSE(Mat, Vec, Vec, Vec);
 48: static PetscErrorCode MatMultTranspose_SeqAIJHIPSPARSE(Mat, Vec, Vec);
 49: static PetscErrorCode MatMultTransposeAdd_SeqAIJHIPSPARSE(Mat, Vec, Vec, Vec);
 50: static PetscErrorCode MatMultHermitianTranspose_SeqAIJHIPSPARSE(Mat, Vec, Vec);
 51: static PetscErrorCode MatMultHermitianTransposeAdd_SeqAIJHIPSPARSE(Mat, Vec, Vec, Vec);
 52: static PetscErrorCode MatMultAddKernel_SeqAIJHIPSPARSE(Mat, Vec, Vec, Vec, PetscBool, PetscBool);
 53: static PetscErrorCode CsrMatrix_Destroy(CsrMatrix **);
 54: static PetscErrorCode MatSeqAIJHIPSPARSEMultStruct_Destroy(Mat_SeqAIJHIPSPARSETriFactorStruct **);
 55: static PetscErrorCode MatSeqAIJHIPSPARSEMultStruct_Destroy(Mat_SeqAIJHIPSPARSEMultStruct **, MatHIPSPARSEStorageFormat);
 56: static PetscErrorCode MatSeqAIJHIPSPARSETriFactors_Destroy(Mat_SeqAIJHIPSPARSETriFactors **);
 57: static PetscErrorCode MatSeqAIJHIPSPARSE_Destroy(Mat_SeqAIJHIPSPARSE **);
 58: static PetscErrorCode MatSeqAIJHIPSPARSECopyFromGPU(Mat);
 59: static PetscErrorCode MatSeqAIJHIPSPARSEILUAnalysisAndCopyToGPU(Mat);
 60: static PetscErrorCode MatSeqAIJHIPSPARSEInvalidateTranspose(Mat, PetscBool);
 61: static PetscErrorCode MatSeqAIJCopySubArray_SeqAIJHIPSPARSE(Mat, PetscInt, const PetscInt[], PetscScalar[]);
 62: static PetscErrorCode MatBindToCPU_SeqAIJHIPSPARSE(Mat, PetscBool);
 63: static PetscErrorCode MatSetPreallocationCOO_SeqAIJHIPSPARSE(Mat, PetscCount, PetscInt[], PetscInt[]);
 64: static PetscErrorCode MatSetValuesCOO_SeqAIJHIPSPARSE(Mat, const PetscScalar[], InsertMode);

 66: PETSC_INTERN PetscErrorCode MatProductSetFromOptions_SeqAIJ_SeqDense(Mat);
 67: PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJHIPSPARSE(Mat, MatType, MatReuse, Mat *);
 68: PETSC_EXTERN PetscErrorCode MatGetFactor_seqaijhipsparse_hipsparse_band(Mat, MatFactorType, Mat *);

 70: /*
 71: PetscErrorCode MatHIPSPARSESetStream(Mat A, const hipStream_t stream)
 72: {
 73:   Mat_SeqAIJHIPSPARSE *hipsparsestruct = (Mat_SeqAIJHIPSPARSE*)A->spptr;

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

 82: PetscErrorCode MatHIPSPARSESetHandle(Mat A, const hipsparseHandle_t handle)
 83: {
 84:   Mat_SeqAIJHIPSPARSE *hipsparsestruct = (Mat_SeqAIJHIPSPARSE*)A->spptr;

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

 96: PetscErrorCode MatHIPSPARSEClearHandle(Mat A)
 97: {
 98:   Mat_SeqAIJHIPSPARSE *hipsparsestruct = (Mat_SeqAIJHIPSPARSE*)A->spptr;
 99:   PetscBool            flg;

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

109: PETSC_INTERN PetscErrorCode MatHIPSPARSESetFormat_SeqAIJHIPSPARSE(Mat A, MatHIPSPARSEFormatOperation op, MatHIPSPARSEStorageFormat format)
110: {
111:   Mat_SeqAIJHIPSPARSE *hipsparsestruct = (Mat_SeqAIJHIPSPARSE *)A->spptr;

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

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

131:    Not Collective

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

139:    Level: intermediate

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

151: PETSC_INTERN PetscErrorCode MatHIPSPARSESetUseCPUSolve_SeqAIJHIPSPARSE(Mat A, PetscBool use_cpu)
152: {
153:   Mat_SeqAIJHIPSPARSE *hipsparsestruct = (Mat_SeqAIJHIPSPARSE *)A->spptr;

155:   PetscFunctionBegin;
156:   hipsparsestruct->use_cpu_solve = use_cpu;
157:   PetscFunctionReturn(PETSC_SUCCESS);
158: }

160: /*@
161:    MatHIPSPARSESetUseCPUSolve - Sets use CPU `MatSolve()`.

163:    Input Parameters:
164: +  A - Matrix of type `MATSEQAIJHIPSPARSE`
165: -  use_cpu - set flag for using the built-in CPU `MatSolve()`

167:    Level: intermediate

169:    Notes:
170:    The hipSparse LU solver currently computes the factors with the built-in CPU method
171:    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.
172:    This method to specifies if the solve is done on the CPU or GPU (GPU is the default).

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

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

200: static PetscErrorCode MatLUFactorNumeric_SeqAIJHIPSPARSE(Mat B, Mat A, const MatFactorInfo *info)
201: {
202:   PetscBool            row_identity, col_identity;
203:   Mat_SeqAIJ          *b     = (Mat_SeqAIJ *)B->data;
204:   IS                   isrow = b->row, iscol = b->col;
205:   Mat_SeqAIJHIPSPARSE *hipsparsestruct = (Mat_SeqAIJHIPSPARSE *)B->spptr;

207:   PetscFunctionBegin;
208:   PetscCall(MatSeqAIJHIPSPARSECopyFromGPU(A));
209:   PetscCall(MatLUFactorNumeric_SeqAIJ(B, A, info));
210:   B->offloadmask = PETSC_OFFLOAD_CPU;
211:   /* determine which version of MatSolve needs to be used. */
212:   PetscCall(ISIdentity(isrow, &row_identity));
213:   PetscCall(ISIdentity(iscol, &col_identity));
214:   if (!hipsparsestruct->use_cpu_solve) {
215:     if (row_identity && col_identity) {
216:       B->ops->solve          = MatSolve_SeqAIJHIPSPARSE_NaturalOrdering;
217:       B->ops->solvetranspose = MatSolveTranspose_SeqAIJHIPSPARSE_NaturalOrdering;
218:     } else {
219:       B->ops->solve          = MatSolve_SeqAIJHIPSPARSE;
220:       B->ops->solvetranspose = MatSolveTranspose_SeqAIJHIPSPARSE;
221:     }
222:   }
223:   B->ops->matsolve          = NULL;
224:   B->ops->matsolvetranspose = NULL;

226:   /* get the triangular factors */
227:   if (!hipsparsestruct->use_cpu_solve) { PetscCall(MatSeqAIJHIPSPARSEILUAnalysisAndCopyToGPU(B)); }
228:   PetscFunctionReturn(PETSC_SUCCESS);
229: }

231: static PetscErrorCode MatSetFromOptions_SeqAIJHIPSPARSE(Mat A, PetscOptionItems *PetscOptionsObject)
232: {
233:   MatHIPSPARSEStorageFormat format;
234:   PetscBool                 flg;
235:   Mat_SeqAIJHIPSPARSE      *hipsparsestruct = (Mat_SeqAIJHIPSPARSE *)A->spptr;

237:   PetscFunctionBegin;
238:   PetscOptionsHeadBegin(PetscOptionsObject, "SeqAIJHIPSPARSE options");
239:   if (A->factortype == MAT_FACTOR_NONE) {
240:     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));
241:     if (flg) PetscCall(MatHIPSPARSESetFormat(A, MAT_HIPSPARSE_MULT, format));
242:     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));
243:     if (flg) PetscCall(MatHIPSPARSESetFormat(A, MAT_HIPSPARSE_ALL, format));
244:     PetscCall(PetscOptionsBool("-mat_hipsparse_use_cpu_solve", "Use CPU (I)LU solve", "MatHIPSPARSESetUseCPUSolve", hipsparsestruct->use_cpu_solve, &hipsparsestruct->use_cpu_solve, &flg));
245:     if (flg) PetscCall(MatHIPSPARSESetUseCPUSolve(A, hipsparsestruct->use_cpu_solve));
246:     PetscCall(
247:       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));
248:     /* If user did use this option, check its consistency with hipSPARSE, since PetscOptionsEnum() sets enum values based on their position in MatHIPSPARSESpMVAlgorithms[] */
249:     PetscCheck(!flg || HIPSPARSE_CSRMV_ALG1 == 2, PETSC_COMM_SELF, PETSC_ERR_SUP, "hipSPARSE enum hipsparseSpMVAlg_t has been changed but PETSc has not been updated accordingly");
250:     PetscCall(
251:       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));
252:     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");
253:     /*
254:     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));
255:     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");
256:     */
257:   }
258:   PetscOptionsHeadEnd();
259:   PetscFunctionReturn(PETSC_SUCCESS);
260: }

262: static PetscErrorCode MatSeqAIJHIPSPARSEBuildILULowerTriMatrix(Mat A)
263: {
264:   Mat_SeqAIJ                         *a                   = (Mat_SeqAIJ *)A->data;
265:   PetscInt                            n                   = A->rmap->n;
266:   Mat_SeqAIJHIPSPARSETriFactors      *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;
267:   Mat_SeqAIJHIPSPARSETriFactorStruct *loTriFactor         = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->loTriFactorPtr;
268:   const PetscInt                     *ai = a->i, *aj = a->j, *vi;
269:   const MatScalar                    *aa = a->a, *v;
270:   PetscInt                           *AiLo, *AjLo;
271:   PetscInt                            i, nz, nzLower, offset, rowOffset;

273:   PetscFunctionBegin;
274:   if (!n) PetscFunctionReturn(PETSC_SUCCESS);
275:   if (A->offloadmask == PETSC_OFFLOAD_UNALLOCATED || A->offloadmask == PETSC_OFFLOAD_CPU) {
276:     try {
277:       /* first figure out the number of nonzeros in the lower triangular matrix including 1's on the diagonal. */
278:       nzLower = n + ai[n] - ai[1];
279:       if (!loTriFactor) {
280:         PetscScalar *AALo;
281:         PetscCallHIP(hipHostMalloc((void **)&AALo, nzLower * sizeof(PetscScalar)));

283:         /* Allocate Space for the lower triangular matrix */
284:         PetscCallHIP(hipHostMalloc((void **)&AiLo, (n + 1) * sizeof(PetscInt)));
285:         PetscCallHIP(hipHostMalloc((void **)&AjLo, nzLower * sizeof(PetscInt)));

287:         /* Fill the lower triangular matrix */
288:         AiLo[0]   = (PetscInt)0;
289:         AiLo[n]   = nzLower;
290:         AjLo[0]   = (PetscInt)0;
291:         AALo[0]   = (MatScalar)1.0;
292:         v         = aa;
293:         vi        = aj;
294:         offset    = 1;
295:         rowOffset = 1;
296:         for (i = 1; i < n; i++) {
297:           nz = ai[i + 1] - ai[i];
298:           /* additional 1 for the term on the diagonal */
299:           AiLo[i] = rowOffset;
300:           rowOffset += nz + 1;

302:           PetscCall(PetscArraycpy(&(AjLo[offset]), vi, nz));
303:           PetscCall(PetscArraycpy(&(AALo[offset]), v, nz));
304:           offset += nz;
305:           AjLo[offset] = (PetscInt)i;
306:           AALo[offset] = (MatScalar)1.0;
307:           offset += 1;
308:           v += nz;
309:           vi += nz;
310:         }

312:         /* allocate space for the triangular factor information */
313:         PetscCall(PetscNew(&loTriFactor));
314:         loTriFactor->solvePolicy = HIPSPARSE_SOLVE_POLICY_USE_LEVEL;
315:         /* Create the matrix description */
316:         PetscCallHIPSPARSE(hipsparseCreateMatDescr(&loTriFactor->descr));
317:         PetscCallHIPSPARSE(hipsparseSetMatIndexBase(loTriFactor->descr, HIPSPARSE_INDEX_BASE_ZERO));
318:         PetscCallHIPSPARSE(hipsparseSetMatType(loTriFactor->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));
319:         PetscCallHIPSPARSE(hipsparseSetMatFillMode(loTriFactor->descr, HIPSPARSE_FILL_MODE_LOWER));
320:         PetscCallHIPSPARSE(hipsparseSetMatDiagType(loTriFactor->descr, HIPSPARSE_DIAG_TYPE_UNIT));

322:         /* set the operation */
323:         loTriFactor->solveOp = HIPSPARSE_OPERATION_NON_TRANSPOSE;

325:         /* set the matrix */
326:         loTriFactor->csrMat                 = new CsrMatrix;
327:         loTriFactor->csrMat->num_rows       = n;
328:         loTriFactor->csrMat->num_cols       = n;
329:         loTriFactor->csrMat->num_entries    = nzLower;
330:         loTriFactor->csrMat->row_offsets    = new THRUSTINTARRAY32(n + 1);
331:         loTriFactor->csrMat->column_indices = new THRUSTINTARRAY32(nzLower);
332:         loTriFactor->csrMat->values         = new THRUSTARRAY(nzLower);

334:         loTriFactor->csrMat->row_offsets->assign(AiLo, AiLo + n + 1);
335:         loTriFactor->csrMat->column_indices->assign(AjLo, AjLo + nzLower);
336:         loTriFactor->csrMat->values->assign(AALo, AALo + nzLower);

338:         /* Create the solve analysis information */
339:         PetscCall(PetscLogEventBegin(MAT_HIPSPARSESolveAnalysis, A, 0, 0, 0));
340:         PetscCallHIPSPARSE(hipsparseCreateCsrsvInfo(&loTriFactor->solveInfo));
341:         PetscCallHIPSPARSE(hipsparseXcsrsv_buffsize(hipsparseTriFactors->handle, loTriFactor->solveOp, loTriFactor->csrMat->num_rows, loTriFactor->csrMat->num_entries, loTriFactor->descr, loTriFactor->csrMat->values->data().get(),
342:                                                     loTriFactor->csrMat->row_offsets->data().get(), loTriFactor->csrMat->column_indices->data().get(), loTriFactor->solveInfo, &loTriFactor->solveBufferSize));
343:         PetscCallHIP(hipMalloc(&loTriFactor->solveBuffer, loTriFactor->solveBufferSize));

345:         /* perform the solve analysis */
346:         PetscCallHIPSPARSE(hipsparseXcsrsv_analysis(hipsparseTriFactors->handle, loTriFactor->solveOp, loTriFactor->csrMat->num_rows, loTriFactor->csrMat->num_entries, loTriFactor->descr, loTriFactor->csrMat->values->data().get(),
347:                                                     loTriFactor->csrMat->row_offsets->data().get(), loTriFactor->csrMat->column_indices->data().get(), loTriFactor->solveInfo, loTriFactor->solvePolicy, loTriFactor->solveBuffer));

349:         PetscCallHIP(WaitForHIP());
350:         PetscCall(PetscLogEventEnd(MAT_HIPSPARSESolveAnalysis, A, 0, 0, 0));

352:         /* assign the pointer */
353:         ((Mat_SeqAIJHIPSPARSETriFactors *)A->spptr)->loTriFactorPtr = loTriFactor;
354:         loTriFactor->AA_h                                           = AALo;
355:         PetscCallHIP(hipHostFree(AiLo));
356:         PetscCallHIP(hipHostFree(AjLo));
357:         PetscCall(PetscLogCpuToGpu((n + 1 + nzLower) * sizeof(int) + nzLower * sizeof(PetscScalar)));
358:       } else { /* update values only */
359:         if (!loTriFactor->AA_h) PetscCallHIP(hipHostMalloc((void **)&loTriFactor->AA_h, nzLower * sizeof(PetscScalar)));
360:         /* Fill the lower triangular matrix */
361:         loTriFactor->AA_h[0] = 1.0;
362:         v                    = aa;
363:         vi                   = aj;
364:         offset               = 1;
365:         for (i = 1; i < n; i++) {
366:           nz = ai[i + 1] - ai[i];
367:           PetscCall(PetscArraycpy(&(loTriFactor->AA_h[offset]), v, nz));
368:           offset += nz;
369:           loTriFactor->AA_h[offset] = 1.0;
370:           offset += 1;
371:           v += nz;
372:         }
373:         loTriFactor->csrMat->values->assign(loTriFactor->AA_h, loTriFactor->AA_h + nzLower);
374:         PetscCall(PetscLogCpuToGpu(nzLower * sizeof(PetscScalar)));
375:       }
376:     } catch (char *ex) {
377:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_LIB, "HIPSPARSE error: %s", ex);
378:     }
379:   }
380:   PetscFunctionReturn(PETSC_SUCCESS);
381: }

383: static PetscErrorCode MatSeqAIJHIPSPARSEBuildILUUpperTriMatrix(Mat A)
384: {
385:   Mat_SeqAIJ                         *a                   = (Mat_SeqAIJ *)A->data;
386:   PetscInt                            n                   = A->rmap->n;
387:   Mat_SeqAIJHIPSPARSETriFactors      *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;
388:   Mat_SeqAIJHIPSPARSETriFactorStruct *upTriFactor         = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->upTriFactorPtr;
389:   const PetscInt                     *aj = a->j, *adiag = a->diag, *vi;
390:   const MatScalar                    *aa = a->a, *v;
391:   PetscInt                           *AiUp, *AjUp;
392:   PetscInt                            i, nz, nzUpper, offset;

394:   PetscFunctionBegin;
395:   if (!n) PetscFunctionReturn(PETSC_SUCCESS);
396:   if (A->offloadmask == PETSC_OFFLOAD_UNALLOCATED || A->offloadmask == PETSC_OFFLOAD_CPU) {
397:     try {
398:       /* next, figure out the number of nonzeros in the upper triangular matrix. */
399:       nzUpper = adiag[0] - adiag[n];
400:       if (!upTriFactor) {
401:         PetscScalar *AAUp;
402:         PetscCallHIP(hipHostMalloc((void **)&AAUp, nzUpper * sizeof(PetscScalar)));

404:         /* Allocate Space for the upper triangular matrix */
405:         PetscCallHIP(hipHostMalloc((void **)&AiUp, (n + 1) * sizeof(PetscInt)));
406:         PetscCallHIP(hipHostMalloc((void **)&AjUp, nzUpper * sizeof(PetscInt)));

408:         /* Fill the upper triangular matrix */
409:         AiUp[0] = (PetscInt)0;
410:         AiUp[n] = nzUpper;
411:         offset  = nzUpper;
412:         for (i = n - 1; i >= 0; i--) {
413:           v  = aa + adiag[i + 1] + 1;
414:           vi = aj + adiag[i + 1] + 1;
415:           nz = adiag[i] - adiag[i + 1] - 1; /* number of elements NOT on the diagonal */
416:           offset -= (nz + 1);               /* decrement the offset */

418:           /* first, set the diagonal elements */
419:           AjUp[offset] = (PetscInt)i;
420:           AAUp[offset] = (MatScalar)1. / v[nz];
421:           AiUp[i]      = AiUp[i + 1] - (nz + 1);

423:           PetscCall(PetscArraycpy(&(AjUp[offset + 1]), vi, nz));
424:           PetscCall(PetscArraycpy(&(AAUp[offset + 1]), v, nz));
425:         }

427:         /* allocate space for the triangular factor information */
428:         PetscCall(PetscNew(&upTriFactor));
429:         upTriFactor->solvePolicy = HIPSPARSE_SOLVE_POLICY_USE_LEVEL;

431:         /* Create the matrix description */
432:         PetscCallHIPSPARSE(hipsparseCreateMatDescr(&upTriFactor->descr));
433:         PetscCallHIPSPARSE(hipsparseSetMatIndexBase(upTriFactor->descr, HIPSPARSE_INDEX_BASE_ZERO));
434:         PetscCallHIPSPARSE(hipsparseSetMatType(upTriFactor->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));
435:         PetscCallHIPSPARSE(hipsparseSetMatFillMode(upTriFactor->descr, HIPSPARSE_FILL_MODE_UPPER));
436:         PetscCallHIPSPARSE(hipsparseSetMatDiagType(upTriFactor->descr, HIPSPARSE_DIAG_TYPE_NON_UNIT));

438:         /* set the operation */
439:         upTriFactor->solveOp = HIPSPARSE_OPERATION_NON_TRANSPOSE;

441:         /* set the matrix */
442:         upTriFactor->csrMat                 = new CsrMatrix;
443:         upTriFactor->csrMat->num_rows       = n;
444:         upTriFactor->csrMat->num_cols       = n;
445:         upTriFactor->csrMat->num_entries    = nzUpper;
446:         upTriFactor->csrMat->row_offsets    = new THRUSTINTARRAY32(n + 1);
447:         upTriFactor->csrMat->column_indices = new THRUSTINTARRAY32(nzUpper);
448:         upTriFactor->csrMat->values         = new THRUSTARRAY(nzUpper);
449:         upTriFactor->csrMat->row_offsets->assign(AiUp, AiUp + n + 1);
450:         upTriFactor->csrMat->column_indices->assign(AjUp, AjUp + nzUpper);
451:         upTriFactor->csrMat->values->assign(AAUp, AAUp + nzUpper);

453:         /* Create the solve analysis information */
454:         PetscCall(PetscLogEventBegin(MAT_HIPSPARSESolveAnalysis, A, 0, 0, 0));
455:         PetscCallHIPSPARSE(hipsparseCreateCsrsvInfo(&upTriFactor->solveInfo));
456:         PetscCallHIPSPARSE(hipsparseXcsrsv_buffsize(hipsparseTriFactors->handle, upTriFactor->solveOp, upTriFactor->csrMat->num_rows, upTriFactor->csrMat->num_entries, upTriFactor->descr, upTriFactor->csrMat->values->data().get(),
457:                                                     upTriFactor->csrMat->row_offsets->data().get(), upTriFactor->csrMat->column_indices->data().get(), upTriFactor->solveInfo, &upTriFactor->solveBufferSize));
458:         PetscCallHIP(hipMalloc(&upTriFactor->solveBuffer, upTriFactor->solveBufferSize));

460:         /* perform the solve analysis */
461:         PetscCallHIPSPARSE(hipsparseXcsrsv_analysis(hipsparseTriFactors->handle, upTriFactor->solveOp, upTriFactor->csrMat->num_rows, upTriFactor->csrMat->num_entries, upTriFactor->descr, upTriFactor->csrMat->values->data().get(),
462:                                                     upTriFactor->csrMat->row_offsets->data().get(), upTriFactor->csrMat->column_indices->data().get(), upTriFactor->solveInfo, upTriFactor->solvePolicy, upTriFactor->solveBuffer));

464:         PetscCallHIP(WaitForHIP());
465:         PetscCall(PetscLogEventEnd(MAT_HIPSPARSESolveAnalysis, A, 0, 0, 0));

467:         /* assign the pointer */
468:         ((Mat_SeqAIJHIPSPARSETriFactors *)A->spptr)->upTriFactorPtr = upTriFactor;
469:         upTriFactor->AA_h                                           = AAUp;
470:         PetscCallHIP(hipHostFree(AiUp));
471:         PetscCallHIP(hipHostFree(AjUp));
472:         PetscCall(PetscLogCpuToGpu((n + 1 + nzUpper) * sizeof(int) + nzUpper * sizeof(PetscScalar)));
473:       } else {
474:         if (!upTriFactor->AA_h) PetscCallHIP(hipHostMalloc((void **)&upTriFactor->AA_h, nzUpper * sizeof(PetscScalar)));
475:         /* Fill the upper triangular matrix */
476:         offset = nzUpper;
477:         for (i = n - 1; i >= 0; i--) {
478:           v  = aa + adiag[i + 1] + 1;
479:           nz = adiag[i] - adiag[i + 1] - 1; /* number of elements NOT on the diagonal */
480:           offset -= (nz + 1);               /* decrement the offset */

482:           /* first, set the diagonal elements */
483:           upTriFactor->AA_h[offset] = 1. / v[nz];
484:           PetscCall(PetscArraycpy(&(upTriFactor->AA_h[offset + 1]), v, nz));
485:         }
486:         upTriFactor->csrMat->values->assign(upTriFactor->AA_h, upTriFactor->AA_h + nzUpper);
487:         PetscCall(PetscLogCpuToGpu(nzUpper * sizeof(PetscScalar)));
488:       }
489:     } catch (char *ex) {
490:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_LIB, "HIPSPARSE error: %s", ex);
491:     }
492:   }
493:   PetscFunctionReturn(PETSC_SUCCESS);
494: }

496: static PetscErrorCode MatSeqAIJHIPSPARSEILUAnalysisAndCopyToGPU(Mat A)
497: {
498:   PetscBool                      row_identity, col_identity;
499:   Mat_SeqAIJ                    *a                   = (Mat_SeqAIJ *)A->data;
500:   Mat_SeqAIJHIPSPARSETriFactors *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;
501:   IS                             isrow = a->row, iscol = a->icol;
502:   PetscInt                       n = A->rmap->n;

504:   PetscFunctionBegin;
505:   PetscCheck(hipsparseTriFactors, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing hipsparseTriFactors");
506:   PetscCall(MatSeqAIJHIPSPARSEBuildILULowerTriMatrix(A));
507:   PetscCall(MatSeqAIJHIPSPARSEBuildILUUpperTriMatrix(A));

509:   if (!hipsparseTriFactors->workVector) hipsparseTriFactors->workVector = new THRUSTARRAY(n);
510:   hipsparseTriFactors->nnz = a->nz;

512:   A->offloadmask = PETSC_OFFLOAD_BOTH;
513:   /* lower triangular indices */
514:   PetscCall(ISIdentity(isrow, &row_identity));
515:   if (!row_identity && !hipsparseTriFactors->rpermIndices) {
516:     const PetscInt *r;

518:     PetscCall(ISGetIndices(isrow, &r));
519:     hipsparseTriFactors->rpermIndices = new THRUSTINTARRAY(n);
520:     hipsparseTriFactors->rpermIndices->assign(r, r + n);
521:     PetscCall(ISRestoreIndices(isrow, &r));
522:     PetscCall(PetscLogCpuToGpu(n * sizeof(PetscInt)));
523:   }
524:   /* upper triangular indices */
525:   PetscCall(ISIdentity(iscol, &col_identity));
526:   if (!col_identity && !hipsparseTriFactors->cpermIndices) {
527:     const PetscInt *c;

529:     PetscCall(ISGetIndices(iscol, &c));
530:     hipsparseTriFactors->cpermIndices = new THRUSTINTARRAY(n);
531:     hipsparseTriFactors->cpermIndices->assign(c, c + n);
532:     PetscCall(ISRestoreIndices(iscol, &c));
533:     PetscCall(PetscLogCpuToGpu(n * sizeof(PetscInt)));
534:   }
535:   PetscFunctionReturn(PETSC_SUCCESS);
536: }

538: static PetscErrorCode MatSeqAIJHIPSPARSEBuildICCTriMatrices(Mat A)
539: {
540:   Mat_SeqAIJ                         *a                   = (Mat_SeqAIJ *)A->data;
541:   Mat_SeqAIJHIPSPARSETriFactors      *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;
542:   Mat_SeqAIJHIPSPARSETriFactorStruct *loTriFactor         = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->loTriFactorPtr;
543:   Mat_SeqAIJHIPSPARSETriFactorStruct *upTriFactor         = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->upTriFactorPtr;
544:   PetscInt                           *AiUp, *AjUp;
545:   PetscScalar                        *AAUp;
546:   PetscScalar                        *AALo;
547:   PetscInt                            nzUpper = a->nz, n = A->rmap->n, i, offset, nz, j;
548:   Mat_SeqSBAIJ                       *b  = (Mat_SeqSBAIJ *)A->data;
549:   const PetscInt                     *ai = b->i, *aj = b->j, *vj;
550:   const MatScalar                    *aa = b->a, *v;

552:   PetscFunctionBegin;
553:   if (!n) PetscFunctionReturn(PETSC_SUCCESS);
554:   if (A->offloadmask == PETSC_OFFLOAD_UNALLOCATED || A->offloadmask == PETSC_OFFLOAD_CPU) {
555:     try {
556:       PetscCallHIP(hipHostMalloc((void **)&AAUp, nzUpper * sizeof(PetscScalar)));
557:       PetscCallHIP(hipHostMalloc((void **)&AALo, nzUpper * sizeof(PetscScalar)));
558:       if (!upTriFactor && !loTriFactor) {
559:         /* Allocate Space for the upper triangular matrix */
560:         PetscCallHIP(hipHostMalloc((void **)&AiUp, (n + 1) * sizeof(PetscInt)));
561:         PetscCallHIP(hipHostMalloc((void **)&AjUp, nzUpper * sizeof(PetscInt)));

563:         /* Fill the upper triangular matrix */
564:         AiUp[0] = (PetscInt)0;
565:         AiUp[n] = nzUpper;
566:         offset  = 0;
567:         for (i = 0; i < n; i++) {
568:           /* set the pointers */
569:           v  = aa + ai[i];
570:           vj = aj + ai[i];
571:           nz = ai[i + 1] - ai[i] - 1; /* exclude diag[i] */

573:           /* first, set the diagonal elements */
574:           AjUp[offset] = (PetscInt)i;
575:           AAUp[offset] = (MatScalar)1.0 / v[nz];
576:           AiUp[i]      = offset;
577:           AALo[offset] = (MatScalar)1.0 / v[nz];

579:           offset += 1;
580:           if (nz > 0) {
581:             PetscCall(PetscArraycpy(&(AjUp[offset]), vj, nz));
582:             PetscCall(PetscArraycpy(&(AAUp[offset]), v, nz));
583:             for (j = offset; j < offset + nz; j++) {
584:               AAUp[j] = -AAUp[j];
585:               AALo[j] = AAUp[j] / v[nz];
586:             }
587:             offset += nz;
588:           }
589:         }

591:         /* allocate space for the triangular factor information */
592:         PetscCall(PetscNew(&upTriFactor));
593:         upTriFactor->solvePolicy = HIPSPARSE_SOLVE_POLICY_USE_LEVEL;

595:         /* Create the matrix description */
596:         PetscCallHIPSPARSE(hipsparseCreateMatDescr(&upTriFactor->descr));
597:         PetscCallHIPSPARSE(hipsparseSetMatIndexBase(upTriFactor->descr, HIPSPARSE_INDEX_BASE_ZERO));
598:         PetscCallHIPSPARSE(hipsparseSetMatType(upTriFactor->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));
599:         PetscCallHIPSPARSE(hipsparseSetMatFillMode(upTriFactor->descr, HIPSPARSE_FILL_MODE_UPPER));
600:         PetscCallHIPSPARSE(hipsparseSetMatDiagType(upTriFactor->descr, HIPSPARSE_DIAG_TYPE_UNIT));

602:         /* set the matrix */
603:         upTriFactor->csrMat                 = new CsrMatrix;
604:         upTriFactor->csrMat->num_rows       = A->rmap->n;
605:         upTriFactor->csrMat->num_cols       = A->cmap->n;
606:         upTriFactor->csrMat->num_entries    = a->nz;
607:         upTriFactor->csrMat->row_offsets    = new THRUSTINTARRAY32(A->rmap->n + 1);
608:         upTriFactor->csrMat->column_indices = new THRUSTINTARRAY32(a->nz);
609:         upTriFactor->csrMat->values         = new THRUSTARRAY(a->nz);
610:         upTriFactor->csrMat->row_offsets->assign(AiUp, AiUp + A->rmap->n + 1);
611:         upTriFactor->csrMat->column_indices->assign(AjUp, AjUp + a->nz);
612:         upTriFactor->csrMat->values->assign(AAUp, AAUp + a->nz);

614:         /* set the operation */
615:         upTriFactor->solveOp = HIPSPARSE_OPERATION_NON_TRANSPOSE;

617:         /* Create the solve analysis information */
618:         PetscCall(PetscLogEventBegin(MAT_HIPSPARSESolveAnalysis, A, 0, 0, 0));
619:         PetscCallHIPSPARSE(hipsparseCreateCsrsvInfo(&upTriFactor->solveInfo));
620:         PetscCallHIPSPARSE(hipsparseXcsrsv_buffsize(hipsparseTriFactors->handle, upTriFactor->solveOp, upTriFactor->csrMat->num_rows, upTriFactor->csrMat->num_entries, upTriFactor->descr, upTriFactor->csrMat->values->data().get(),
621:                                                     upTriFactor->csrMat->row_offsets->data().get(), upTriFactor->csrMat->column_indices->data().get(), upTriFactor->solveInfo, &upTriFactor->solveBufferSize));
622:         PetscCallHIP(hipMalloc(&upTriFactor->solveBuffer, upTriFactor->solveBufferSize));

624:         /* perform the solve analysis */
625:         PetscCallHIPSPARSE(hipsparseXcsrsv_analysis(hipsparseTriFactors->handle, upTriFactor->solveOp, upTriFactor->csrMat->num_rows, upTriFactor->csrMat->num_entries, upTriFactor->descr, upTriFactor->csrMat->values->data().get(),
626:                                                     upTriFactor->csrMat->row_offsets->data().get(), upTriFactor->csrMat->column_indices->data().get(), upTriFactor->solveInfo, upTriFactor->solvePolicy, upTriFactor->solveBuffer));

628:         PetscCallHIP(WaitForHIP());
629:         PetscCall(PetscLogEventEnd(MAT_HIPSPARSESolveAnalysis, A, 0, 0, 0));

631:         /* assign the pointer */
632:         ((Mat_SeqAIJHIPSPARSETriFactors *)A->spptr)->upTriFactorPtr = upTriFactor;

634:         /* allocate space for the triangular factor information */
635:         PetscCall(PetscNew(&loTriFactor));
636:         loTriFactor->solvePolicy = HIPSPARSE_SOLVE_POLICY_USE_LEVEL;

638:         /* Create the matrix description */
639:         PetscCallHIPSPARSE(hipsparseCreateMatDescr(&loTriFactor->descr));
640:         PetscCallHIPSPARSE(hipsparseSetMatIndexBase(loTriFactor->descr, HIPSPARSE_INDEX_BASE_ZERO));
641:         PetscCallHIPSPARSE(hipsparseSetMatType(loTriFactor->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));
642:         PetscCallHIPSPARSE(hipsparseSetMatFillMode(loTriFactor->descr, HIPSPARSE_FILL_MODE_UPPER));
643:         PetscCallHIPSPARSE(hipsparseSetMatDiagType(loTriFactor->descr, HIPSPARSE_DIAG_TYPE_NON_UNIT));

645:         /* set the operation */
646:         loTriFactor->solveOp = HIPSPARSE_OPERATION_TRANSPOSE;

648:         /* set the matrix */
649:         loTriFactor->csrMat                 = new CsrMatrix;
650:         loTriFactor->csrMat->num_rows       = A->rmap->n;
651:         loTriFactor->csrMat->num_cols       = A->cmap->n;
652:         loTriFactor->csrMat->num_entries    = a->nz;
653:         loTriFactor->csrMat->row_offsets    = new THRUSTINTARRAY32(A->rmap->n + 1);
654:         loTriFactor->csrMat->column_indices = new THRUSTINTARRAY32(a->nz);
655:         loTriFactor->csrMat->values         = new THRUSTARRAY(a->nz);
656:         loTriFactor->csrMat->row_offsets->assign(AiUp, AiUp + A->rmap->n + 1);
657:         loTriFactor->csrMat->column_indices->assign(AjUp, AjUp + a->nz);
658:         loTriFactor->csrMat->values->assign(AALo, AALo + a->nz);

660:         /* Create the solve analysis information */
661:         PetscCall(PetscLogEventBegin(MAT_HIPSPARSESolveAnalysis, A, 0, 0, 0));
662:         PetscCallHIPSPARSE(hipsparseCreateCsrsvInfo(&loTriFactor->solveInfo));
663:         PetscCallHIPSPARSE(hipsparseXcsrsv_buffsize(hipsparseTriFactors->handle, loTriFactor->solveOp, loTriFactor->csrMat->num_rows, loTriFactor->csrMat->num_entries, loTriFactor->descr, loTriFactor->csrMat->values->data().get(),
664:                                                     loTriFactor->csrMat->row_offsets->data().get(), loTriFactor->csrMat->column_indices->data().get(), loTriFactor->solveInfo, &loTriFactor->solveBufferSize));
665:         PetscCallHIP(hipMalloc(&loTriFactor->solveBuffer, loTriFactor->solveBufferSize));

667:         /* perform the solve analysis */
668:         PetscCallHIPSPARSE(hipsparseXcsrsv_analysis(hipsparseTriFactors->handle, loTriFactor->solveOp, loTriFactor->csrMat->num_rows, loTriFactor->csrMat->num_entries, loTriFactor->descr, loTriFactor->csrMat->values->data().get(),
669:                                                     loTriFactor->csrMat->row_offsets->data().get(), loTriFactor->csrMat->column_indices->data().get(), loTriFactor->solveInfo, loTriFactor->solvePolicy, loTriFactor->solveBuffer));

671:         PetscCallHIP(WaitForHIP());
672:         PetscCall(PetscLogEventEnd(MAT_HIPSPARSESolveAnalysis, A, 0, 0, 0));

674:         /* assign the pointer */
675:         ((Mat_SeqAIJHIPSPARSETriFactors *)A->spptr)->loTriFactorPtr = loTriFactor;

677:         PetscCall(PetscLogCpuToGpu(2 * (((A->rmap->n + 1) + (a->nz)) * sizeof(int) + (a->nz) * sizeof(PetscScalar))));
678:         PetscCallHIP(hipHostFree(AiUp));
679:         PetscCallHIP(hipHostFree(AjUp));
680:       } else {
681:         /* Fill the upper triangular matrix */
682:         offset = 0;
683:         for (i = 0; i < n; i++) {
684:           /* set the pointers */
685:           v  = aa + ai[i];
686:           nz = ai[i + 1] - ai[i] - 1; /* exclude diag[i] */

688:           /* first, set the diagonal elements */
689:           AAUp[offset] = 1.0 / v[nz];
690:           AALo[offset] = 1.0 / v[nz];

692:           offset += 1;
693:           if (nz > 0) {
694:             PetscCall(PetscArraycpy(&(AAUp[offset]), v, nz));
695:             for (j = offset; j < offset + nz; j++) {
696:               AAUp[j] = -AAUp[j];
697:               AALo[j] = AAUp[j] / v[nz];
698:             }
699:             offset += nz;
700:           }
701:         }
702:         PetscCheck(upTriFactor, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing hipsparseTriFactors");
703:         PetscCheck(loTriFactor, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing hipsparseTriFactors");
704:         upTriFactor->csrMat->values->assign(AAUp, AAUp + a->nz);
705:         loTriFactor->csrMat->values->assign(AALo, AALo + a->nz);
706:         PetscCall(PetscLogCpuToGpu(2 * (a->nz) * sizeof(PetscScalar)));
707:       }
708:       PetscCallHIP(hipHostFree(AAUp));
709:       PetscCallHIP(hipHostFree(AALo));
710:     } catch (char *ex) {
711:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_LIB, "HIPSPARSE error: %s", ex);
712:     }
713:   }
714:   PetscFunctionReturn(PETSC_SUCCESS);
715: }

717: static PetscErrorCode MatSeqAIJHIPSPARSEICCAnalysisAndCopyToGPU(Mat A)
718: {
719:   PetscBool                      perm_identity;
720:   Mat_SeqAIJ                    *a                   = (Mat_SeqAIJ *)A->data;
721:   Mat_SeqAIJHIPSPARSETriFactors *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;
722:   IS                             ip                  = a->row;
723:   PetscInt                       n                   = A->rmap->n;

725:   PetscFunctionBegin;
726:   PetscCheck(hipsparseTriFactors, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing hipsparseTriFactors");
727:   PetscCall(MatSeqAIJHIPSPARSEBuildICCTriMatrices(A));
728:   if (!hipsparseTriFactors->workVector) hipsparseTriFactors->workVector = new THRUSTARRAY(n);
729:   hipsparseTriFactors->nnz = (a->nz - n) * 2 + n;

731:   A->offloadmask = PETSC_OFFLOAD_BOTH;
732:   /* lower triangular indices */
733:   PetscCall(ISIdentity(ip, &perm_identity));
734:   if (!perm_identity) {
735:     IS              iip;
736:     const PetscInt *irip, *rip;

738:     PetscCall(ISInvertPermutation(ip, PETSC_DECIDE, &iip));
739:     PetscCall(ISGetIndices(iip, &irip));
740:     PetscCall(ISGetIndices(ip, &rip));
741:     hipsparseTriFactors->rpermIndices = new THRUSTINTARRAY(n);
742:     hipsparseTriFactors->cpermIndices = new THRUSTINTARRAY(n);
743:     hipsparseTriFactors->rpermIndices->assign(rip, rip + n);
744:     hipsparseTriFactors->cpermIndices->assign(irip, irip + n);
745:     PetscCall(ISRestoreIndices(iip, &irip));
746:     PetscCall(ISDestroy(&iip));
747:     PetscCall(ISRestoreIndices(ip, &rip));
748:     PetscCall(PetscLogCpuToGpu(2. * n * sizeof(PetscInt)));
749:   }
750:   PetscFunctionReturn(PETSC_SUCCESS);
751: }

753: static PetscErrorCode MatCholeskyFactorNumeric_SeqAIJHIPSPARSE(Mat B, Mat A, const MatFactorInfo *info)
754: {
755:   PetscBool   perm_identity;
756:   Mat_SeqAIJ *b  = (Mat_SeqAIJ *)B->data;
757:   IS          ip = b->row;

759:   PetscFunctionBegin;
760:   PetscCall(MatSeqAIJHIPSPARSECopyFromGPU(A));
761:   PetscCall(MatCholeskyFactorNumeric_SeqAIJ(B, A, info));
762:   B->offloadmask = PETSC_OFFLOAD_CPU;
763:   /* determine which version of MatSolve needs to be used. */
764:   PetscCall(ISIdentity(ip, &perm_identity));
765:   if (perm_identity) {
766:     B->ops->solve             = MatSolve_SeqAIJHIPSPARSE_NaturalOrdering;
767:     B->ops->solvetranspose    = MatSolveTranspose_SeqAIJHIPSPARSE_NaturalOrdering;
768:     B->ops->matsolve          = NULL;
769:     B->ops->matsolvetranspose = NULL;
770:   } else {
771:     B->ops->solve             = MatSolve_SeqAIJHIPSPARSE;
772:     B->ops->solvetranspose    = MatSolveTranspose_SeqAIJHIPSPARSE;
773:     B->ops->matsolve          = NULL;
774:     B->ops->matsolvetranspose = NULL;
775:   }

777:   /* get the triangular factors */
778:   PetscCall(MatSeqAIJHIPSPARSEICCAnalysisAndCopyToGPU(B));
779:   PetscFunctionReturn(PETSC_SUCCESS);
780: }

782: static PetscErrorCode MatSeqAIJHIPSPARSEAnalyzeTransposeForSolve(Mat A)
783: {
784:   Mat_SeqAIJHIPSPARSETriFactors      *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;
785:   Mat_SeqAIJHIPSPARSETriFactorStruct *loTriFactor         = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->loTriFactorPtr;
786:   Mat_SeqAIJHIPSPARSETriFactorStruct *upTriFactor         = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->upTriFactorPtr;
787:   Mat_SeqAIJHIPSPARSETriFactorStruct *loTriFactorT;
788:   Mat_SeqAIJHIPSPARSETriFactorStruct *upTriFactorT;
789:   hipsparseIndexBase_t                indexBase;
790:   hipsparseMatrixType_t               matrixType;
791:   hipsparseFillMode_t                 fillMode;
792:   hipsparseDiagType_t                 diagType;

794:   PetscFunctionBegin;
795:   /* allocate space for the transpose of the lower triangular factor */
796:   PetscCall(PetscNew(&loTriFactorT));
797:   loTriFactorT->solvePolicy = HIPSPARSE_SOLVE_POLICY_USE_LEVEL;

799:   /* set the matrix descriptors of the lower triangular factor */
800:   matrixType = hipsparseGetMatType(loTriFactor->descr);
801:   indexBase  = hipsparseGetMatIndexBase(loTriFactor->descr);
802:   fillMode   = hipsparseGetMatFillMode(loTriFactor->descr) == HIPSPARSE_FILL_MODE_UPPER ? HIPSPARSE_FILL_MODE_LOWER : HIPSPARSE_FILL_MODE_UPPER;
803:   diagType   = hipsparseGetMatDiagType(loTriFactor->descr);

805:   /* Create the matrix description */
806:   PetscCallHIPSPARSE(hipsparseCreateMatDescr(&loTriFactorT->descr));
807:   PetscCallHIPSPARSE(hipsparseSetMatIndexBase(loTriFactorT->descr, indexBase));
808:   PetscCallHIPSPARSE(hipsparseSetMatType(loTriFactorT->descr, matrixType));
809:   PetscCallHIPSPARSE(hipsparseSetMatFillMode(loTriFactorT->descr, fillMode));
810:   PetscCallHIPSPARSE(hipsparseSetMatDiagType(loTriFactorT->descr, diagType));

812:   /* set the operation */
813:   loTriFactorT->solveOp = HIPSPARSE_OPERATION_NON_TRANSPOSE;

815:   /* allocate GPU space for the CSC of the lower triangular factor*/
816:   loTriFactorT->csrMat                 = new CsrMatrix;
817:   loTriFactorT->csrMat->num_rows       = loTriFactor->csrMat->num_cols;
818:   loTriFactorT->csrMat->num_cols       = loTriFactor->csrMat->num_rows;
819:   loTriFactorT->csrMat->num_entries    = loTriFactor->csrMat->num_entries;
820:   loTriFactorT->csrMat->row_offsets    = new THRUSTINTARRAY32(loTriFactorT->csrMat->num_rows + 1);
821:   loTriFactorT->csrMat->column_indices = new THRUSTINTARRAY32(loTriFactorT->csrMat->num_entries);
822:   loTriFactorT->csrMat->values         = new THRUSTARRAY(loTriFactorT->csrMat->num_entries);

824:   /* compute the transpose of the lower triangular factor, i.e. the CSC */
825:   /* Csr2cscEx2 is not implemented in ROCm-5.2.0 and is planned for implementation in hipsparse with future releases of ROCm
826: #if PETSC_PKG_HIP_VERSION_GE(5, 2, 0)
827:   PetscCallHIPSPARSE(hipsparseCsr2cscEx2_bufferSize(hipsparseTriFactors->handle, loTriFactor->csrMat->num_rows, loTriFactor->csrMat->num_cols, loTriFactor->csrMat->num_entries, loTriFactor->csrMat->values->data().get(),
828:                                                   loTriFactor->csrMat->row_offsets->data().get(), loTriFactor->csrMat->column_indices->data().get(), loTriFactorT->csrMat->values->data().get(), loTriFactorT->csrMat->row_offsets->data().get(),
829:                                                   loTriFactorT->csrMat->column_indices->data().get(), hipsparse_scalartype, HIPSPARSE_ACTION_NUMERIC, indexBase, HIPSPARSE_CSR2CSC_ALG1, &loTriFactor->csr2cscBufferSize));
830:   PetscCallHIP(hipMalloc(&loTriFactor->csr2cscBuffer, loTriFactor->csr2cscBufferSize));
831: #endif
832: */
833:   PetscCall(PetscLogEventBegin(MAT_HIPSPARSEGenerateTranspose, A, 0, 0, 0));

835:   PetscCallHIPSPARSE(hipsparse_csr2csc(hipsparseTriFactors->handle, loTriFactor->csrMat->num_rows, loTriFactor->csrMat->num_cols, loTriFactor->csrMat->num_entries, loTriFactor->csrMat->values->data().get(), loTriFactor->csrMat->row_offsets->data().get(),
836:                                        loTriFactor->csrMat->column_indices->data().get(), loTriFactorT->csrMat->values->data().get(),
837: #if 0 /* when Csr2cscEx2 is implemented in hipSparse PETSC_PKG_HIP_VERSION_GE(5, 2, 0)*/
838:                           loTriFactorT->csrMat->row_offsets->data().get(), loTriFactorT->csrMat->column_indices->data().get(),
839:                           hipsparse_scalartype, HIPSPARSE_ACTION_NUMERIC, indexBase, HIPSPARSE_CSR2CSC_ALG1, loTriFactor->csr2cscBuffer));
840: #else
841:                                        loTriFactorT->csrMat->column_indices->data().get(), loTriFactorT->csrMat->row_offsets->data().get(), HIPSPARSE_ACTION_NUMERIC, indexBase));
842: #endif

844:   PetscCallHIP(WaitForHIP());
845:   PetscCall(PetscLogEventBegin(MAT_HIPSPARSEGenerateTranspose, A, 0, 0, 0));

847:   /* Create the solve analysis information */
848:   PetscCall(PetscLogEventBegin(MAT_HIPSPARSESolveAnalysis, A, 0, 0, 0));
849:   PetscCallHIPSPARSE(hipsparseCreateCsrsvInfo(&loTriFactorT->solveInfo));
850:   PetscCallHIPSPARSE(hipsparseXcsrsv_buffsize(hipsparseTriFactors->handle, loTriFactorT->solveOp, loTriFactorT->csrMat->num_rows, loTriFactorT->csrMat->num_entries, loTriFactorT->descr, loTriFactorT->csrMat->values->data().get(),
851:                                               loTriFactorT->csrMat->row_offsets->data().get(), loTriFactorT->csrMat->column_indices->data().get(), loTriFactorT->solveInfo, &loTriFactorT->solveBufferSize));
852:   PetscCallHIP(hipMalloc(&loTriFactorT->solveBuffer, loTriFactorT->solveBufferSize));

854:   /* perform the solve analysis */
855:   PetscCallHIPSPARSE(hipsparseXcsrsv_analysis(hipsparseTriFactors->handle, loTriFactorT->solveOp, loTriFactorT->csrMat->num_rows, loTriFactorT->csrMat->num_entries, loTriFactorT->descr, loTriFactorT->csrMat->values->data().get(),
856:                                               loTriFactorT->csrMat->row_offsets->data().get(), loTriFactorT->csrMat->column_indices->data().get(), loTriFactorT->solveInfo, loTriFactorT->solvePolicy, loTriFactorT->solveBuffer));

858:   PetscCallHIP(WaitForHIP());
859:   PetscCall(PetscLogEventEnd(MAT_HIPSPARSESolveAnalysis, A, 0, 0, 0));

861:   /* assign the pointer */
862:   ((Mat_SeqAIJHIPSPARSETriFactors *)A->spptr)->loTriFactorPtrTranspose = loTriFactorT;

864:   /*********************************************/
865:   /* Now the Transpose of the Upper Tri Factor */
866:   /*********************************************/

868:   /* allocate space for the transpose of the upper triangular factor */
869:   PetscCall(PetscNew(&upTriFactorT));
870:   upTriFactorT->solvePolicy = HIPSPARSE_SOLVE_POLICY_USE_LEVEL;

872:   /* set the matrix descriptors of the upper triangular factor */
873:   matrixType = hipsparseGetMatType(upTriFactor->descr);
874:   indexBase  = hipsparseGetMatIndexBase(upTriFactor->descr);
875:   fillMode   = hipsparseGetMatFillMode(upTriFactor->descr) == HIPSPARSE_FILL_MODE_UPPER ? HIPSPARSE_FILL_MODE_LOWER : HIPSPARSE_FILL_MODE_UPPER;
876:   diagType   = hipsparseGetMatDiagType(upTriFactor->descr);

878:   /* Create the matrix description */
879:   PetscCallHIPSPARSE(hipsparseCreateMatDescr(&upTriFactorT->descr));
880:   PetscCallHIPSPARSE(hipsparseSetMatIndexBase(upTriFactorT->descr, indexBase));
881:   PetscCallHIPSPARSE(hipsparseSetMatType(upTriFactorT->descr, matrixType));
882:   PetscCallHIPSPARSE(hipsparseSetMatFillMode(upTriFactorT->descr, fillMode));
883:   PetscCallHIPSPARSE(hipsparseSetMatDiagType(upTriFactorT->descr, diagType));

885:   /* set the operation */
886:   upTriFactorT->solveOp = HIPSPARSE_OPERATION_NON_TRANSPOSE;

888:   /* allocate GPU space for the CSC of the upper triangular factor*/
889:   upTriFactorT->csrMat                 = new CsrMatrix;
890:   upTriFactorT->csrMat->num_rows       = upTriFactor->csrMat->num_cols;
891:   upTriFactorT->csrMat->num_cols       = upTriFactor->csrMat->num_rows;
892:   upTriFactorT->csrMat->num_entries    = upTriFactor->csrMat->num_entries;
893:   upTriFactorT->csrMat->row_offsets    = new THRUSTINTARRAY32(upTriFactorT->csrMat->num_rows + 1);
894:   upTriFactorT->csrMat->column_indices = new THRUSTINTARRAY32(upTriFactorT->csrMat->num_entries);
895:   upTriFactorT->csrMat->values         = new THRUSTARRAY(upTriFactorT->csrMat->num_entries);

897:   /* compute the transpose of the upper triangular factor, i.e. the CSC */
898:   /* Csr2cscEx2 is not implemented in ROCm-5.2.0 and is planned for implementation in hipsparse with future releases of ROCm
899: #if PETSC_PKG_HIP_VERSION_GE(5, 2, 0)
900:   PetscCallHIPSPARSE(hipsparseCsr2cscEx2_bufferSize(hipsparseTriFactors->handle, upTriFactor->csrMat->num_rows, upTriFactor->csrMat->num_cols, upTriFactor->csrMat->num_entries, upTriFactor->csrMat->values->data().get(),
901:                                                   upTriFactor->csrMat->row_offsets->data().get(), upTriFactor->csrMat->column_indices->data().get(), upTriFactorT->csrMat->values->data().get(), upTriFactorT->csrMat->row_offsets->data().get(),
902:                                                   upTriFactorT->csrMat->column_indices->data().get(), hipsparse_scalartype, HIPSPARSE_ACTION_NUMERIC, indexBase, HIPSPARSE_CSR2CSC_ALG1, &upTriFactor->csr2cscBufferSize));
903:   PetscCallHIP(hipMalloc(&upTriFactor->csr2cscBuffer, upTriFactor->csr2cscBufferSize));
904: #endif
905: */
906:   PetscCall(PetscLogEventBegin(MAT_HIPSPARSEGenerateTranspose, A, 0, 0, 0));
907:   PetscCallHIPSPARSE(hipsparse_csr2csc(hipsparseTriFactors->handle, upTriFactor->csrMat->num_rows, upTriFactor->csrMat->num_cols, upTriFactor->csrMat->num_entries, upTriFactor->csrMat->values->data().get(), upTriFactor->csrMat->row_offsets->data().get(),
908:                                        upTriFactor->csrMat->column_indices->data().get(), upTriFactorT->csrMat->values->data().get(),
909: #if 0 /* when Csr2cscEx2 is implemented in hipSparse PETSC_PKG_HIP_VERSION_GE(5, 2, 0)*/
910:                           upTriFactorT->csrMat->row_offsets->data().get(), upTriFactorT->csrMat->column_indices->data().get(),
911:                           hipsparse_scalartype, HIPSPARSE_ACTION_NUMERIC, indexBase, HIPSPARSE_CSR2CSC_ALG1, upTriFactor->csr2cscBuffer));
912: #else
913:                                        upTriFactorT->csrMat->column_indices->data().get(), upTriFactorT->csrMat->row_offsets->data().get(), HIPSPARSE_ACTION_NUMERIC, indexBase));
914: #endif

916:   PetscCallHIP(WaitForHIP());
917:   PetscCall(PetscLogEventBegin(MAT_HIPSPARSEGenerateTranspose, A, 0, 0, 0));

919:   /* Create the solve analysis information */
920:   PetscCall(PetscLogEventBegin(MAT_HIPSPARSESolveAnalysis, A, 0, 0, 0));
921:   PetscCallHIPSPARSE(hipsparseCreateCsrsvInfo(&upTriFactorT->solveInfo));
922:   PetscCallHIPSPARSE(hipsparseXcsrsv_buffsize(hipsparseTriFactors->handle, upTriFactorT->solveOp, upTriFactorT->csrMat->num_rows, upTriFactorT->csrMat->num_entries, upTriFactorT->descr, upTriFactorT->csrMat->values->data().get(),
923:                                               upTriFactorT->csrMat->row_offsets->data().get(), upTriFactorT->csrMat->column_indices->data().get(), upTriFactorT->solveInfo, &upTriFactorT->solveBufferSize));
924:   PetscCallHIP(hipMalloc(&upTriFactorT->solveBuffer, upTriFactorT->solveBufferSize));

926:   /* perform the solve analysis */
927:   PetscCallHIPSPARSE(hipsparseXcsrsv_analysis(hipsparseTriFactors->handle, upTriFactorT->solveOp, upTriFactorT->csrMat->num_rows, upTriFactorT->csrMat->num_entries, upTriFactorT->descr, upTriFactorT->csrMat->values->data().get(),
928:                                               upTriFactorT->csrMat->row_offsets->data().get(), upTriFactorT->csrMat->column_indices->data().get(), upTriFactorT->solveInfo, upTriFactorT->solvePolicy, upTriFactorT->solveBuffer));

930:   PetscCallHIP(WaitForHIP());
931:   PetscCall(PetscLogEventEnd(MAT_HIPSPARSESolveAnalysis, A, 0, 0, 0));

933:   /* assign the pointer */
934:   ((Mat_SeqAIJHIPSPARSETriFactors *)A->spptr)->upTriFactorPtrTranspose = upTriFactorT;
935:   PetscFunctionReturn(PETSC_SUCCESS);
936: }

938: struct PetscScalarToPetscInt {
939:   __host__ __device__ PetscInt operator()(PetscScalar s) { return (PetscInt)PetscRealPart(s); }
940: };

942: static PetscErrorCode MatSeqAIJHIPSPARSEFormExplicitTranspose(Mat A)
943: {
944:   Mat_SeqAIJHIPSPARSE           *hipsparsestruct = (Mat_SeqAIJHIPSPARSE *)A->spptr;
945:   Mat_SeqAIJHIPSPARSEMultStruct *matstruct, *matstructT;
946:   Mat_SeqAIJ                    *a = (Mat_SeqAIJ *)A->data;
947:   hipsparseIndexBase_t           indexBase;

949:   PetscFunctionBegin;
950:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
951:   matstruct = (Mat_SeqAIJHIPSPARSEMultStruct *)hipsparsestruct->mat;
952:   PetscCheck(matstruct, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing mat struct");
953:   matstructT = (Mat_SeqAIJHIPSPARSEMultStruct *)hipsparsestruct->matTranspose;
954:   PetscCheck(!A->transupdated || matstructT, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing matTranspose struct");
955:   if (A->transupdated) PetscFunctionReturn(PETSC_SUCCESS);
956:   PetscCall(PetscLogEventBegin(MAT_HIPSPARSEGenerateTranspose, A, 0, 0, 0));
957:   PetscCall(PetscLogGpuTimeBegin());
958:   if (hipsparsestruct->format != MAT_HIPSPARSE_CSR) PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(A, PETSC_TRUE));
959:   if (!hipsparsestruct->matTranspose) { /* create hipsparse matrix */
960:     matstructT = new Mat_SeqAIJHIPSPARSEMultStruct;
961:     PetscCallHIPSPARSE(hipsparseCreateMatDescr(&matstructT->descr));
962:     indexBase = hipsparseGetMatIndexBase(matstruct->descr);
963:     PetscCallHIPSPARSE(hipsparseSetMatIndexBase(matstructT->descr, indexBase));
964:     PetscCallHIPSPARSE(hipsparseSetMatType(matstructT->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));

966:     /* set alpha and beta */
967:     PetscCallHIP(hipMalloc((void **)&(matstructT->alpha_one), sizeof(PetscScalar)));
968:     PetscCallHIP(hipMalloc((void **)&(matstructT->beta_zero), sizeof(PetscScalar)));
969:     PetscCallHIP(hipMalloc((void **)&(matstructT->beta_one), sizeof(PetscScalar)));
970:     PetscCallHIP(hipMemcpy(matstructT->alpha_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
971:     PetscCallHIP(hipMemcpy(matstructT->beta_zero, &PETSC_HIPSPARSE_ZERO, sizeof(PetscScalar), hipMemcpyHostToDevice));
972:     PetscCallHIP(hipMemcpy(matstructT->beta_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));

974:     if (hipsparsestruct->format == MAT_HIPSPARSE_CSR) {
975:       CsrMatrix *matrixT      = new CsrMatrix;
976:       matstructT->mat         = matrixT;
977:       matrixT->num_rows       = A->cmap->n;
978:       matrixT->num_cols       = A->rmap->n;
979:       matrixT->num_entries    = a->nz;
980:       matrixT->row_offsets    = new THRUSTINTARRAY32(matrixT->num_rows + 1);
981:       matrixT->column_indices = new THRUSTINTARRAY32(a->nz);
982:       matrixT->values         = new THRUSTARRAY(a->nz);

984:       if (!hipsparsestruct->rowoffsets_gpu) hipsparsestruct->rowoffsets_gpu = new THRUSTINTARRAY32(A->rmap->n + 1);
985:       hipsparsestruct->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);

987:       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(), HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, /* row offset, col idx type due to THRUSTINTARRAY32 */
988:                                             indexBase, hipsparse_scalartype));
989:     } else if (hipsparsestruct->format == MAT_HIPSPARSE_ELL || hipsparsestruct->format == MAT_HIPSPARSE_HYB) {
990:       CsrMatrix *temp  = new CsrMatrix;
991:       CsrMatrix *tempT = new CsrMatrix;
992:       /* First convert HYB to CSR */
993:       temp->num_rows       = A->rmap->n;
994:       temp->num_cols       = A->cmap->n;
995:       temp->num_entries    = a->nz;
996:       temp->row_offsets    = new THRUSTINTARRAY32(A->rmap->n + 1);
997:       temp->column_indices = new THRUSTINTARRAY32(a->nz);
998:       temp->values         = new THRUSTARRAY(a->nz);

1000:       PetscCallHIPSPARSE(hipsparse_hyb2csr(hipsparsestruct->handle, matstruct->descr, (hipsparseHybMat_t)matstruct->mat, temp->values->data().get(), temp->row_offsets->data().get(), temp->column_indices->data().get()));

1002:       /* Next, convert CSR to CSC (i.e. the matrix transpose) */
1003:       tempT->num_rows       = A->rmap->n;
1004:       tempT->num_cols       = A->cmap->n;
1005:       tempT->num_entries    = a->nz;
1006:       tempT->row_offsets    = new THRUSTINTARRAY32(A->rmap->n + 1);
1007:       tempT->column_indices = new THRUSTINTARRAY32(a->nz);
1008:       tempT->values         = new THRUSTARRAY(a->nz);

1010:       PetscCallHIPSPARSE(hipsparse_csr2csc(hipsparsestruct->handle, temp->num_rows, temp->num_cols, temp->num_entries, temp->values->data().get(), temp->row_offsets->data().get(), temp->column_indices->data().get(), tempT->values->data().get(),
1011:                                            tempT->column_indices->data().get(), tempT->row_offsets->data().get(), HIPSPARSE_ACTION_NUMERIC, indexBase));

1013:       /* Last, convert CSC to HYB */
1014:       hipsparseHybMat_t hybMat;
1015:       PetscCallHIPSPARSE(hipsparseCreateHybMat(&hybMat));
1016:       hipsparseHybPartition_t partition = hipsparsestruct->format == MAT_HIPSPARSE_ELL ? HIPSPARSE_HYB_PARTITION_MAX : HIPSPARSE_HYB_PARTITION_AUTO;
1017:       PetscCallHIPSPARSE(hipsparse_csr2hyb(hipsparsestruct->handle, A->rmap->n, A->cmap->n, matstructT->descr, tempT->values->data().get(), tempT->row_offsets->data().get(), tempT->column_indices->data().get(), hybMat, 0, partition));

1019:       /* assign the pointer */
1020:       matstructT->mat = hybMat;
1021:       A->transupdated = PETSC_TRUE;
1022:       /* delete temporaries */
1023:       if (tempT) {
1024:         if (tempT->values) delete (THRUSTARRAY *)tempT->values;
1025:         if (tempT->column_indices) delete (THRUSTINTARRAY32 *)tempT->column_indices;
1026:         if (tempT->row_offsets) delete (THRUSTINTARRAY32 *)tempT->row_offsets;
1027:         delete (CsrMatrix *)tempT;
1028:       }
1029:       if (temp) {
1030:         if (temp->values) delete (THRUSTARRAY *)temp->values;
1031:         if (temp->column_indices) delete (THRUSTINTARRAY32 *)temp->column_indices;
1032:         if (temp->row_offsets) delete (THRUSTINTARRAY32 *)temp->row_offsets;
1033:         delete (CsrMatrix *)temp;
1034:       }
1035:     }
1036:   }
1037:   if (hipsparsestruct->format == MAT_HIPSPARSE_CSR) { /* transpose mat struct may be already present, update data */
1038:     CsrMatrix *matrix  = (CsrMatrix *)matstruct->mat;
1039:     CsrMatrix *matrixT = (CsrMatrix *)matstructT->mat;
1040:     PetscCheck(matrix, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrix");
1041:     PetscCheck(matrix->row_offsets, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrix rows");
1042:     PetscCheck(matrix->column_indices, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrix cols");
1043:     PetscCheck(matrix->values, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrix values");
1044:     PetscCheck(matrixT, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrixT");
1045:     PetscCheck(matrixT->row_offsets, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrixT rows");
1046:     PetscCheck(matrixT->column_indices, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrixT cols");
1047:     PetscCheck(matrixT->values, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CsrMatrixT values");
1048:     if (!hipsparsestruct->rowoffsets_gpu) { /* this may be absent when we did not construct the transpose with csr2csc */
1049:       hipsparsestruct->rowoffsets_gpu = new THRUSTINTARRAY32(A->rmap->n + 1);
1050:       hipsparsestruct->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
1051:       PetscCall(PetscLogCpuToGpu((A->rmap->n + 1) * sizeof(PetscInt)));
1052:     }
1053:     if (!hipsparsestruct->csr2csc_i) {
1054:       THRUSTARRAY csr2csc_a(matrix->num_entries);
1055:       PetscCallThrust(thrust::sequence(thrust::device, csr2csc_a.begin(), csr2csc_a.end(), 0.0));

1057:       indexBase = hipsparseGetMatIndexBase(matstruct->descr);
1058:       if (matrix->num_entries) {
1059:         /* This routine is known to give errors with CUDA-11, but works fine with CUDA-10
1060:            Need to verify this for ROCm.
1061:         */
1062:         PetscCallHIPSPARSE(hipsparse_csr2csc(hipsparsestruct->handle, A->rmap->n, A->cmap->n, matrix->num_entries, csr2csc_a.data().get(), hipsparsestruct->rowoffsets_gpu->data().get(), matrix->column_indices->data().get(), matrixT->values->data().get(),
1063:                                              matrixT->column_indices->data().get(), matrixT->row_offsets->data().get(), HIPSPARSE_ACTION_NUMERIC, indexBase));
1064:       } else {
1065:         matrixT->row_offsets->assign(matrixT->row_offsets->size(), indexBase);
1066:       }

1068:       hipsparsestruct->csr2csc_i = new THRUSTINTARRAY(matrix->num_entries);
1069:       PetscCallThrust(thrust::transform(thrust::device, matrixT->values->begin(), matrixT->values->end(), hipsparsestruct->csr2csc_i->begin(), PetscScalarToPetscInt()));
1070:     }
1071:     PetscCallThrust(
1072:       thrust::copy(thrust::device, thrust::make_permutation_iterator(matrix->values->begin(), hipsparsestruct->csr2csc_i->begin()), thrust::make_permutation_iterator(matrix->values->begin(), hipsparsestruct->csr2csc_i->end()), matrixT->values->begin()));
1073:   }
1074:   PetscCall(PetscLogGpuTimeEnd());
1075:   PetscCall(PetscLogEventEnd(MAT_HIPSPARSEGenerateTranspose, A, 0, 0, 0));
1076:   /* the compressed row indices is not used for matTranspose */
1077:   matstructT->cprowIndices = NULL;
1078:   /* assign the pointer */
1079:   ((Mat_SeqAIJHIPSPARSE *)A->spptr)->matTranspose = matstructT;
1080:   A->transupdated                                 = PETSC_TRUE;
1081:   PetscFunctionReturn(PETSC_SUCCESS);
1082: }

1084: /* Why do we need to analyze the transposed matrix again? Can't we just use op(A) = HIPSPARSE_OPERATION_TRANSPOSE in MatSolve_SeqAIJHIPSPARSE? */
1085: static PetscErrorCode MatSolveTranspose_SeqAIJHIPSPARSE(Mat A, Vec bb, Vec xx)
1086: {
1087:   PetscInt                              n = xx->map->n;
1088:   const PetscScalar                    *barray;
1089:   PetscScalar                          *xarray;
1090:   thrust::device_ptr<const PetscScalar> bGPU;
1091:   thrust::device_ptr<PetscScalar>       xGPU;
1092:   Mat_SeqAIJHIPSPARSETriFactors        *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;
1093:   Mat_SeqAIJHIPSPARSETriFactorStruct   *loTriFactorT        = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->loTriFactorPtrTranspose;
1094:   Mat_SeqAIJHIPSPARSETriFactorStruct   *upTriFactorT        = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->upTriFactorPtrTranspose;
1095:   THRUSTARRAY                          *tempGPU             = (THRUSTARRAY *)hipsparseTriFactors->workVector;

1097:   PetscFunctionBegin;
1098:   /* Analyze the matrix and create the transpose ... on the fly */
1099:   if (!loTriFactorT && !upTriFactorT) {
1100:     PetscCall(MatSeqAIJHIPSPARSEAnalyzeTransposeForSolve(A));
1101:     loTriFactorT = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->loTriFactorPtrTranspose;
1102:     upTriFactorT = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->upTriFactorPtrTranspose;
1103:   }

1105:   /* Get the GPU pointers */
1106:   PetscCall(VecHIPGetArrayWrite(xx, &xarray));
1107:   PetscCall(VecHIPGetArrayRead(bb, &barray));
1108:   xGPU = thrust::device_pointer_cast(xarray);
1109:   bGPU = thrust::device_pointer_cast(barray);

1111:   PetscCall(PetscLogGpuTimeBegin());
1112:   /* First, reorder with the row permutation */
1113:   thrust::copy(thrust::hip::par.on(PetscDefaultHipStream), thrust::make_permutation_iterator(bGPU, hipsparseTriFactors->rpermIndices->begin()), thrust::make_permutation_iterator(bGPU + n, hipsparseTriFactors->rpermIndices->end()), xGPU);

1115:   /* First, solve U */
1116:   PetscCallHIPSPARSE(hipsparseXcsrsv_solve(hipsparseTriFactors->handle, upTriFactorT->solveOp, upTriFactorT->csrMat->num_rows, upTriFactorT->csrMat->num_entries, &PETSC_HIPSPARSE_ONE, upTriFactorT->descr, upTriFactorT->csrMat->values->data().get(),
1117:                                            upTriFactorT->csrMat->row_offsets->data().get(), upTriFactorT->csrMat->column_indices->data().get(), upTriFactorT->solveInfo, xarray, tempGPU->data().get(), upTriFactorT->solvePolicy, upTriFactorT->solveBuffer));

1119:   /* Then, solve L */
1120:   PetscCallHIPSPARSE(hipsparseXcsrsv_solve(hipsparseTriFactors->handle, loTriFactorT->solveOp, loTriFactorT->csrMat->num_rows, loTriFactorT->csrMat->num_entries, &PETSC_HIPSPARSE_ONE, loTriFactorT->descr, loTriFactorT->csrMat->values->data().get(),
1121:                                            loTriFactorT->csrMat->row_offsets->data().get(), loTriFactorT->csrMat->column_indices->data().get(), loTriFactorT->solveInfo, tempGPU->data().get(), xarray, loTriFactorT->solvePolicy, loTriFactorT->solveBuffer));

1123:   /* Last, copy the solution, xGPU, into a temporary with the column permutation ... can't be done in place. */
1124:   thrust::copy(thrust::hip::par.on(PetscDefaultHipStream), thrust::make_permutation_iterator(xGPU, hipsparseTriFactors->cpermIndices->begin()), thrust::make_permutation_iterator(xGPU + n, hipsparseTriFactors->cpermIndices->end()), tempGPU->begin());

1126:   /* Copy the temporary to the full solution. */
1127:   thrust::copy(thrust::hip::par.on(PetscDefaultHipStream), tempGPU->begin(), tempGPU->end(), xGPU);

1129:   /* restore */
1130:   PetscCall(VecHIPRestoreArrayRead(bb, &barray));
1131:   PetscCall(VecHIPRestoreArrayWrite(xx, &xarray));
1132:   PetscCall(PetscLogGpuTimeEnd());
1133:   PetscCall(PetscLogGpuFlops(2.0 * hipsparseTriFactors->nnz - A->cmap->n));
1134:   PetscFunctionReturn(PETSC_SUCCESS);
1135: }

1137: static PetscErrorCode MatSolveTranspose_SeqAIJHIPSPARSE_NaturalOrdering(Mat A, Vec bb, Vec xx)
1138: {
1139:   const PetscScalar                  *barray;
1140:   PetscScalar                        *xarray;
1141:   Mat_SeqAIJHIPSPARSETriFactors      *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;
1142:   Mat_SeqAIJHIPSPARSETriFactorStruct *loTriFactorT        = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->loTriFactorPtrTranspose;
1143:   Mat_SeqAIJHIPSPARSETriFactorStruct *upTriFactorT        = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->upTriFactorPtrTranspose;
1144:   THRUSTARRAY                        *tempGPU             = (THRUSTARRAY *)hipsparseTriFactors->workVector;

1146:   PetscFunctionBegin;
1147:   /* Analyze the matrix and create the transpose ... on the fly */
1148:   if (!loTriFactorT && !upTriFactorT) {
1149:     PetscCall(MatSeqAIJHIPSPARSEAnalyzeTransposeForSolve(A));
1150:     loTriFactorT = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->loTriFactorPtrTranspose;
1151:     upTriFactorT = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->upTriFactorPtrTranspose;
1152:   }

1154:   /* Get the GPU pointers */
1155:   PetscCall(VecHIPGetArrayWrite(xx, &xarray));
1156:   PetscCall(VecHIPGetArrayRead(bb, &barray));

1158:   PetscCall(PetscLogGpuTimeBegin());
1159:   /* First, solve U */
1160:   PetscCallHIPSPARSE(hipsparseXcsrsv_solve(hipsparseTriFactors->handle, upTriFactorT->solveOp, upTriFactorT->csrMat->num_rows, upTriFactorT->csrMat->num_entries, &PETSC_HIPSPARSE_ONE, upTriFactorT->descr, upTriFactorT->csrMat->values->data().get(),
1161:                                            upTriFactorT->csrMat->row_offsets->data().get(), upTriFactorT->csrMat->column_indices->data().get(), upTriFactorT->solveInfo, barray, tempGPU->data().get(), upTriFactorT->solvePolicy, upTriFactorT->solveBuffer));

1163:   /* Then, solve L */
1164:   PetscCallHIPSPARSE(hipsparseXcsrsv_solve(hipsparseTriFactors->handle, loTriFactorT->solveOp, loTriFactorT->csrMat->num_rows, loTriFactorT->csrMat->num_entries, &PETSC_HIPSPARSE_ONE, loTriFactorT->descr, loTriFactorT->csrMat->values->data().get(),
1165:                                            loTriFactorT->csrMat->row_offsets->data().get(), loTriFactorT->csrMat->column_indices->data().get(), loTriFactorT->solveInfo, tempGPU->data().get(), xarray, loTriFactorT->solvePolicy, loTriFactorT->solveBuffer));

1167:   /* restore */
1168:   PetscCall(VecHIPRestoreArrayRead(bb, &barray));
1169:   PetscCall(VecHIPRestoreArrayWrite(xx, &xarray));
1170:   PetscCall(PetscLogGpuTimeEnd());
1171:   PetscCall(PetscLogGpuFlops(2.0 * hipsparseTriFactors->nnz - A->cmap->n));
1172:   PetscFunctionReturn(PETSC_SUCCESS);
1173: }

1175: static PetscErrorCode MatSolve_SeqAIJHIPSPARSE(Mat A, Vec bb, Vec xx)
1176: {
1177:   const PetscScalar                    *barray;
1178:   PetscScalar                          *xarray;
1179:   thrust::device_ptr<const PetscScalar> bGPU;
1180:   thrust::device_ptr<PetscScalar>       xGPU;
1181:   Mat_SeqAIJHIPSPARSETriFactors        *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;
1182:   Mat_SeqAIJHIPSPARSETriFactorStruct   *loTriFactor         = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->loTriFactorPtr;
1183:   Mat_SeqAIJHIPSPARSETriFactorStruct   *upTriFactor         = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->upTriFactorPtr;
1184:   THRUSTARRAY                          *tempGPU             = (THRUSTARRAY *)hipsparseTriFactors->workVector;

1186:   PetscFunctionBegin;
1187:   /* Get the GPU pointers */
1188:   PetscCall(VecHIPGetArrayWrite(xx, &xarray));
1189:   PetscCall(VecHIPGetArrayRead(bb, &barray));
1190:   xGPU = thrust::device_pointer_cast(xarray);
1191:   bGPU = thrust::device_pointer_cast(barray);

1193:   PetscCall(PetscLogGpuTimeBegin());
1194:   /* First, reorder with the row permutation */
1195:   thrust::copy(thrust::hip::par.on(PetscDefaultHipStream), thrust::make_permutation_iterator(bGPU, hipsparseTriFactors->rpermIndices->begin()), thrust::make_permutation_iterator(bGPU, hipsparseTriFactors->rpermIndices->end()), tempGPU->begin());

1197:   /* Next, solve L */
1198:   PetscCallHIPSPARSE(hipsparseXcsrsv_solve(hipsparseTriFactors->handle, loTriFactor->solveOp, loTriFactor->csrMat->num_rows, loTriFactor->csrMat->num_entries, &PETSC_HIPSPARSE_ONE, loTriFactor->descr, loTriFactor->csrMat->values->data().get(),
1199:                                            loTriFactor->csrMat->row_offsets->data().get(), loTriFactor->csrMat->column_indices->data().get(), loTriFactor->solveInfo, tempGPU->data().get(), xarray, loTriFactor->solvePolicy, loTriFactor->solveBuffer));

1201:   /* Then, solve U */
1202:   PetscCallHIPSPARSE(hipsparseXcsrsv_solve(hipsparseTriFactors->handle, upTriFactor->solveOp, upTriFactor->csrMat->num_rows, upTriFactor->csrMat->num_entries, &PETSC_HIPSPARSE_ONE, upTriFactor->descr, upTriFactor->csrMat->values->data().get(),
1203:                                            upTriFactor->csrMat->row_offsets->data().get(), upTriFactor->csrMat->column_indices->data().get(), upTriFactor->solveInfo, xarray, tempGPU->data().get(), upTriFactor->solvePolicy, upTriFactor->solveBuffer));

1205:   /* Last, reorder with the column permutation */
1206:   thrust::copy(thrust::hip::par.on(PetscDefaultHipStream), thrust::make_permutation_iterator(tempGPU->begin(), hipsparseTriFactors->cpermIndices->begin()), thrust::make_permutation_iterator(tempGPU->begin(), hipsparseTriFactors->cpermIndices->end()), xGPU);

1208:   PetscCall(VecHIPRestoreArrayRead(bb, &barray));
1209:   PetscCall(VecHIPRestoreArrayWrite(xx, &xarray));
1210:   PetscCall(PetscLogGpuTimeEnd());
1211:   PetscCall(PetscLogGpuFlops(2.0 * hipsparseTriFactors->nnz - A->cmap->n));
1212:   PetscFunctionReturn(PETSC_SUCCESS);
1213: }

1215: static PetscErrorCode MatSolve_SeqAIJHIPSPARSE_NaturalOrdering(Mat A, Vec bb, Vec xx)
1216: {
1217:   const PetscScalar                  *barray;
1218:   PetscScalar                        *xarray;
1219:   Mat_SeqAIJHIPSPARSETriFactors      *hipsparseTriFactors = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;
1220:   Mat_SeqAIJHIPSPARSETriFactorStruct *loTriFactor         = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->loTriFactorPtr;
1221:   Mat_SeqAIJHIPSPARSETriFactorStruct *upTriFactor         = (Mat_SeqAIJHIPSPARSETriFactorStruct *)hipsparseTriFactors->upTriFactorPtr;
1222:   THRUSTARRAY                        *tempGPU             = (THRUSTARRAY *)hipsparseTriFactors->workVector;

1224:   PetscFunctionBegin;
1225:   /* Get the GPU pointers */
1226:   PetscCall(VecHIPGetArrayWrite(xx, &xarray));
1227:   PetscCall(VecHIPGetArrayRead(bb, &barray));

1229:   PetscCall(PetscLogGpuTimeBegin());
1230:   /* First, solve L */
1231:   PetscCallHIPSPARSE(hipsparseXcsrsv_solve(hipsparseTriFactors->handle, loTriFactor->solveOp, loTriFactor->csrMat->num_rows, loTriFactor->csrMat->num_entries, &PETSC_HIPSPARSE_ONE, loTriFactor->descr, loTriFactor->csrMat->values->data().get(),
1232:                                            loTriFactor->csrMat->row_offsets->data().get(), loTriFactor->csrMat->column_indices->data().get(), loTriFactor->solveInfo, barray, tempGPU->data().get(), loTriFactor->solvePolicy, loTriFactor->solveBuffer));

1234:   /* Next, solve U */
1235:   PetscCallHIPSPARSE(hipsparseXcsrsv_solve(hipsparseTriFactors->handle, upTriFactor->solveOp, upTriFactor->csrMat->num_rows, upTriFactor->csrMat->num_entries, &PETSC_HIPSPARSE_ONE, upTriFactor->descr, upTriFactor->csrMat->values->data().get(),
1236:                                            upTriFactor->csrMat->row_offsets->data().get(), upTriFactor->csrMat->column_indices->data().get(), upTriFactor->solveInfo, tempGPU->data().get(), xarray, upTriFactor->solvePolicy, upTriFactor->solveBuffer));

1238:   PetscCall(VecHIPRestoreArrayRead(bb, &barray));
1239:   PetscCall(VecHIPRestoreArrayWrite(xx, &xarray));
1240:   PetscCall(PetscLogGpuTimeEnd());
1241:   PetscCall(PetscLogGpuFlops(2.0 * hipsparseTriFactors->nnz - A->cmap->n));
1242:   PetscFunctionReturn(PETSC_SUCCESS);
1243: }

1245: #if PETSC_PKG_HIP_VERSION_GE(4, 5, 0)
1246: /* hipsparseSpSV_solve() and related functions first appeared in ROCm-4.5.0*/
1247: static PetscErrorCode MatSolve_SeqAIJHIPSPARSE_ILU0(Mat fact, Vec b, Vec x)
1248: {
1249:   Mat_SeqAIJHIPSPARSETriFactors *fs  = (Mat_SeqAIJHIPSPARSETriFactors *)fact->spptr;
1250:   Mat_SeqAIJ                    *aij = (Mat_SeqAIJ *)fact->data;
1251:   const PetscScalar             *barray;
1252:   PetscScalar                   *xarray;

1254:   PetscFunctionBegin;
1255:   PetscCall(VecHIPGetArrayWrite(x, &xarray));
1256:   PetscCall(VecHIPGetArrayRead(b, &barray));
1257:   PetscCall(PetscLogGpuTimeBegin());

1259:   /* Solve L*y = b */
1260:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, (void *)barray));
1261:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_Y, fs->Y));
1262:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L,                                     /* L Y = X */
1263:                                          fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, fs->spsvBuffer_L)); // hipsparseSpSV_solve() secretely uses the external buffer used in hipsparseSpSV_analysis()!

1265:   /* Solve U*x = y */
1266:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, xarray));
1267:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, /* U X = Y */
1268:                                          fs->dnVecDescr_Y, fs->dnVecDescr_X, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_U, fs->spsvBuffer_U));

1270:   PetscCall(VecHIPRestoreArrayRead(b, &barray));
1271:   PetscCall(VecHIPRestoreArrayWrite(x, &xarray));

1273:   PetscCall(PetscLogGpuTimeEnd());
1274:   PetscCall(PetscLogGpuFlops(2.0 * aij->nz - fact->rmap->n));
1275:   PetscFunctionReturn(PETSC_SUCCESS);
1276: }

1278: static PetscErrorCode MatSolveTranspose_SeqAIJHIPSPARSE_ILU0(Mat fact, Vec b, Vec x)
1279: {
1280:   Mat_SeqAIJHIPSPARSETriFactors *fs  = (Mat_SeqAIJHIPSPARSETriFactors *)fact->spptr;
1281:   Mat_SeqAIJ                    *aij = (Mat_SeqAIJ *)fact->data;
1282:   const PetscScalar             *barray;
1283:   PetscScalar                   *xarray;

1285:   PetscFunctionBegin;
1286:   if (!fs->createdTransposeSpSVDescr) { /* Call MatSolveTranspose() for the first time */
1287:     PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_Lt));
1288:     PetscCallHIPSPARSE(hipsparseSpSV_bufferSize(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, /* The matrix is still L. We only do transpose solve with it */
1289:                                                 fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Lt, &fs->spsvBufferSize_Lt));

1291:     PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_Ut));
1292:     PetscCallHIPSPARSE(hipsparseSpSV_bufferSize(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Ut, &fs->spsvBufferSize_Ut));
1293:     PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_Lt, fs->spsvBufferSize_Lt));
1294:     PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_Ut, fs->spsvBufferSize_Ut));
1295:     fs->createdTransposeSpSVDescr = PETSC_TRUE;
1296:   }

1298:   if (!fs->updatedTransposeSpSVAnalysis) {
1299:     PetscCallHIPSPARSE(hipsparseSpSV_analysis(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Lt, fs->spsvBuffer_Lt));

1301:     PetscCallHIPSPARSE(hipsparseSpSV_analysis(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Ut, fs->spsvBuffer_Ut));
1302:     fs->updatedTransposeSpSVAnalysis = PETSC_TRUE;
1303:   }

1305:   PetscCall(VecHIPGetArrayWrite(x, &xarray));
1306:   PetscCall(VecHIPGetArrayRead(b, &barray));
1307:   PetscCall(PetscLogGpuTimeBegin());

1309:   /* Solve Ut*y = b */
1310:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, (void *)barray));
1311:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_Y, fs->Y));
1312:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_U, /* Ut Y = X */
1313:                                          fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Ut, fs->spsvBuffer_Ut));

1315:   /* Solve Lt*x = y */
1316:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, xarray));
1317:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, /* Lt X = Y */
1318:                                          fs->dnVecDescr_Y, fs->dnVecDescr_X, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Lt, fs->spsvBuffer_Lt));

1320:   PetscCall(VecHIPRestoreArrayRead(b, &barray));
1321:   PetscCall(VecHIPRestoreArrayWrite(x, &xarray));
1322:   PetscCall(PetscLogGpuTimeEnd());
1323:   PetscCall(PetscLogGpuFlops(2.0 * aij->nz - fact->rmap->n));
1324:   PetscFunctionReturn(PETSC_SUCCESS);
1325: }

1327: static PetscErrorCode MatILUFactorNumeric_SeqAIJHIPSPARSE_ILU0(Mat fact, Mat A, const MatFactorInfo *info)
1328: {
1329:   Mat_SeqAIJHIPSPARSETriFactors *fs    = (Mat_SeqAIJHIPSPARSETriFactors *)fact->spptr;
1330:   Mat_SeqAIJ                    *aij   = (Mat_SeqAIJ *)fact->data;
1331:   Mat_SeqAIJHIPSPARSE           *Acusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
1332:   CsrMatrix                     *Acsr;
1333:   PetscInt                       m, nz;
1334:   PetscBool                      flg;

1336:   PetscFunctionBegin;
1337:   if (PetscDefined(USE_DEBUG)) {
1338:     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
1339:     PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Expected MATSEQAIJHIPSPARSE, but input is %s", ((PetscObject)A)->type_name);
1340:   }

1342:   /* Copy A's value to fact */
1343:   m  = fact->rmap->n;
1344:   nz = aij->nz;
1345:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
1346:   Acsr = (CsrMatrix *)Acusp->mat->mat;
1347:   PetscCallHIP(hipMemcpyAsync(fs->csrVal, Acsr->values->data().get(), sizeof(PetscScalar) * nz, hipMemcpyDeviceToDevice, PetscDefaultHipStream));

1349:   /* Factorize fact inplace */
1350:   if (m)
1351:     PetscCallHIPSPARSE(hipsparseXcsrilu02(fs->handle, m, nz, /* hipsparseXcsrilu02 errors out with empty matrices (m=0) */
1352:                                           fs->matDescr_M, fs->csrVal, fs->csrRowPtr, fs->csrColIdx, fs->ilu0Info_M, fs->policy_M, fs->factBuffer_M));
1353:   if (PetscDefined(USE_DEBUG)) {
1354:     int               numerical_zero;
1355:     hipsparseStatus_t status;
1356:     status = hipsparseXcsrilu02_zeroPivot(fs->handle, fs->ilu0Info_M, &numerical_zero);
1357:     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);
1358:   }

1360:   /* hipsparseSpSV_analysis() is numeric, i.e., it requires valid matrix values, therefore, we do it after hipsparseXcsrilu02() */
1361:   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));

1363:   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));

1365:   /* L, U values have changed, reset the flag to indicate we need to redo hipsparseSpSV_analysis() for transpose solve */
1366:   fs->updatedTransposeSpSVAnalysis = PETSC_FALSE;

1368:   fact->offloadmask            = PETSC_OFFLOAD_GPU;
1369:   fact->ops->solve             = MatSolve_SeqAIJHIPSPARSE_ILU0;
1370:   fact->ops->solvetranspose    = MatSolveTranspose_SeqAIJHIPSPARSE_ILU0;
1371:   fact->ops->matsolve          = NULL;
1372:   fact->ops->matsolvetranspose = NULL;
1373:   PetscCall(PetscLogGpuFlops(fs->numericFactFlops));
1374:   PetscFunctionReturn(PETSC_SUCCESS);
1375: }

1377: static PetscErrorCode MatILUFactorSymbolic_SeqAIJHIPSPARSE_ILU0(Mat fact, Mat A, IS isrow, IS iscol, const MatFactorInfo *info)
1378: {
1379:   Mat_SeqAIJHIPSPARSETriFactors *fs  = (Mat_SeqAIJHIPSPARSETriFactors *)fact->spptr;
1380:   Mat_SeqAIJ                    *aij = (Mat_SeqAIJ *)fact->data;
1381:   PetscInt                       m, nz;

1383:   PetscFunctionBegin;
1384:   if (PetscDefined(USE_DEBUG)) {
1385:     PetscInt  i;
1386:     PetscBool flg, missing;

1388:     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
1389:     PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Expected MATSEQAIJHIPSPARSE, but input is %s", ((PetscObject)A)->type_name);
1390:     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);
1391:     PetscCall(MatMissingDiagonal(A, &missing, &i));
1392:     PetscCheck(!missing, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix is missing diagonal entry %" PetscInt_FMT, i);
1393:   }

1395:   /* Free the old stale stuff */
1396:   PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(&fs));

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

1403:   fact->offloadmask            = PETSC_OFFLOAD_BOTH;
1404:   fact->factortype             = MAT_FACTOR_ILU;
1405:   fact->info.factor_mallocs    = 0;
1406:   fact->info.fill_ratio_given  = info->fill;
1407:   fact->info.fill_ratio_needed = 1.0;

1409:   aij->row = NULL;
1410:   aij->col = NULL;

1412:   /* ====================================================================== */
1413:   /* Copy A's i, j to fact and also allocate the value array of fact.       */
1414:   /* We'll do in-place factorization on fact                                */
1415:   /* ====================================================================== */
1416:   const int *Ai, *Aj;

1418:   m  = fact->rmap->n;
1419:   nz = aij->nz;

1421:   PetscCallHIP(hipMalloc((void **)&fs->csrRowPtr, sizeof(int) * (m + 1)));
1422:   PetscCallHIP(hipMalloc((void **)&fs->csrColIdx, sizeof(int) * nz));
1423:   PetscCallHIP(hipMalloc((void **)&fs->csrVal, sizeof(PetscScalar) * nz));
1424:   PetscCall(MatSeqAIJHIPSPARSEGetIJ(A, PETSC_FALSE, &Ai, &Aj)); /* Do not use compressed Ai */
1425:   PetscCallHIP(hipMemcpyAsync(fs->csrRowPtr, Ai, sizeof(int) * (m + 1), hipMemcpyDeviceToDevice, PetscDefaultHipStream));
1426:   PetscCallHIP(hipMemcpyAsync(fs->csrColIdx, Aj, sizeof(int) * nz, hipMemcpyDeviceToDevice, PetscDefaultHipStream));

1428:   /* ====================================================================== */
1429:   /* Create descriptors for M, L, U                                         */
1430:   /* ====================================================================== */
1431:   hipsparseFillMode_t fillMode;
1432:   hipsparseDiagType_t diagType;

1434:   PetscCallHIPSPARSE(hipsparseCreateMatDescr(&fs->matDescr_M));
1435:   PetscCallHIPSPARSE(hipsparseSetMatIndexBase(fs->matDescr_M, HIPSPARSE_INDEX_BASE_ZERO));
1436:   PetscCallHIPSPARSE(hipsparseSetMatType(fs->matDescr_M, HIPSPARSE_MATRIX_TYPE_GENERAL));

1438:   /* https://docs.amd.com/bundle/hipSPARSE-Documentation---hipSPARSE-documentation/page/usermanual.html/#hipsparse_8h_1a79e036b6c0680cb37e2aa53d3542a054
1439:     hipsparseDiagType_t: This type indicates if the matrix diagonal entries are unity. The diagonal elements are always
1440:     assumed to be present, but if HIPSPARSE_DIAG_TYPE_UNIT is passed to an API routine, then the routine assumes that
1441:     all diagonal entries are unity and will not read or modify those entries. Note that in this case the routine
1442:     assumes the diagonal entries are equal to one, regardless of what those entries are actually set to in memory.
1443:   */
1444:   fillMode = HIPSPARSE_FILL_MODE_LOWER;
1445:   diagType = HIPSPARSE_DIAG_TYPE_UNIT;
1446:   PetscCallHIPSPARSE(hipsparseCreateCsr(&fs->spMatDescr_L, m, m, nz, fs->csrRowPtr, fs->csrColIdx, fs->csrVal, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
1447:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_L, HIPSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
1448:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_L, HIPSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));

1450:   fillMode = HIPSPARSE_FILL_MODE_UPPER;
1451:   diagType = HIPSPARSE_DIAG_TYPE_NON_UNIT;
1452:   PetscCallHIPSPARSE(hipsparseCreateCsr(&fs->spMatDescr_U, m, m, nz, fs->csrRowPtr, fs->csrColIdx, fs->csrVal, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
1453:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_U, HIPSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
1454:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_U, HIPSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));

1456:   /* ========================================================================= */
1457:   /* Query buffer sizes for csrilu0, SpSV and allocate buffers                 */
1458:   /* ========================================================================= */
1459:   PetscCallHIPSPARSE(hipsparseCreateCsrilu02Info(&fs->ilu0Info_M));
1460:   if (m)
1461:     PetscCallHIPSPARSE(hipsparseXcsrilu02_bufferSize(fs->handle, m, nz, /* hipsparseXcsrilu02 errors out with empty matrices (m=0) */
1462:                                                      fs->matDescr_M, fs->csrVal, fs->csrRowPtr, fs->csrColIdx, fs->ilu0Info_M, &fs->factBufferSize_M));

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

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

1470:   PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_L));
1471:   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));

1473:   PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_U));
1474:   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));

1476:   /* It appears 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.
1477:      To save memory, we make factBuffer_M share with the bigger of spsvBuffer_L/U.
1478:    */
1479:   if (fs->spsvBufferSize_L > fs->spsvBufferSize_U) {
1480:     PetscCallHIP(hipMalloc((void **)&fs->factBuffer_M, PetscMax(fs->spsvBufferSize_L, (size_t)fs->factBufferSize_M)));
1481:     fs->spsvBuffer_L = fs->factBuffer_M;
1482:     PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_U, fs->spsvBufferSize_U));
1483:   } else {
1484:     PetscCallHIP(hipMalloc((void **)&fs->factBuffer_M, PetscMax(fs->spsvBufferSize_U, (size_t)fs->factBufferSize_M)));
1485:     fs->spsvBuffer_U = fs->factBuffer_M;
1486:     PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_L, fs->spsvBufferSize_L));
1487:   }

1489:   /* ========================================================================== */
1490:   /* Perform analysis of ilu0 on M, SpSv on L and U                             */
1491:   /* The lower(upper) triangular part of M has the same sparsity pattern as L(U)*/
1492:   /* ========================================================================== */
1493:   int structural_zero;

1495:   fs->policy_M = HIPSPARSE_SOLVE_POLICY_USE_LEVEL;
1496:   if (m)
1497:     PetscCallHIPSPARSE(hipsparseXcsrilu02_analysis(fs->handle, m, nz, /* hipsparseXcsrilu02 errors out with empty matrices (m=0) */
1498:                                                    fs->matDescr_M, fs->csrVal, fs->csrRowPtr, fs->csrColIdx, fs->ilu0Info_M, fs->policy_M, fs->factBuffer_M));
1499:   if (PetscDefined(USE_DEBUG)) {
1500:     /* Function hipsparseXcsrilu02_zeroPivot() is a blocking call. It calls hipDeviceSynchronize() to make sure all previous kernels are done. */
1501:     hipsparseStatus_t status;
1502:     status = hipsparseXcsrilu02_zeroPivot(fs->handle, fs->ilu0Info_M, &structural_zero);
1503:     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);
1504:   }

1506:   /* Estimate FLOPs of the numeric factorization */
1507:   {
1508:     Mat_SeqAIJ    *Aseq = (Mat_SeqAIJ *)A->data;
1509:     PetscInt      *Ai, *Adiag, nzRow, nzLeft;
1510:     PetscLogDouble flops = 0.0;

1512:     PetscCall(MatMarkDiagonal_SeqAIJ(A));
1513:     Ai    = Aseq->i;
1514:     Adiag = Aseq->diag;
1515:     for (PetscInt i = 0; i < m; i++) {
1516:       if (Ai[i] < Adiag[i] && Adiag[i] < Ai[i + 1]) { /* There are nonzeros left to the diagonal of row i */
1517:         nzRow  = Ai[i + 1] - Ai[i];
1518:         nzLeft = Adiag[i] - Ai[i];
1519:         /* We want to eliminate nonzeros left to the diagonal one by one. Assume each time, nonzeros right
1520:           and include the eliminated one will be updated, which incurs a multiplication and an addition.
1521:         */
1522:         nzLeft = (nzRow - 1) / 2;
1523:         flops += nzLeft * (2.0 * nzRow - nzLeft + 1);
1524:       }
1525:     }
1526:     fs->numericFactFlops = flops;
1527:   }
1528:   fact->ops->lufactornumeric = MatILUFactorNumeric_SeqAIJHIPSPARSE_ILU0;
1529:   PetscFunctionReturn(PETSC_SUCCESS);
1530: }

1532: static PetscErrorCode MatSolve_SeqAIJHIPSPARSE_ICC0(Mat fact, Vec b, Vec x)
1533: {
1534:   Mat_SeqAIJHIPSPARSETriFactors *fs  = (Mat_SeqAIJHIPSPARSETriFactors *)fact->spptr;
1535:   Mat_SeqAIJ                    *aij = (Mat_SeqAIJ *)fact->data;
1536:   const PetscScalar             *barray;
1537:   PetscScalar                   *xarray;

1539:   PetscFunctionBegin;
1540:   PetscCall(VecHIPGetArrayWrite(x, &xarray));
1541:   PetscCall(VecHIPGetArrayRead(b, &barray));
1542:   PetscCall(PetscLogGpuTimeBegin());

1544:   /* Solve L*y = b */
1545:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, (void *)barray));
1546:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_Y, fs->Y));
1547:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_NON_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, /* L Y = X */
1548:                                          fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_L, fs->spsvBuffer_L));

1550:   /* Solve Lt*x = y */
1551:   PetscCallHIPSPARSE(hipsparseDnVecSetValues(fs->dnVecDescr_X, xarray));
1552:   PetscCallHIPSPARSE(hipsparseSpSV_solve(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, /* Lt X = Y */
1553:                                          fs->dnVecDescr_Y, fs->dnVecDescr_X, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Lt, fs->spsvBuffer_Lt));

1555:   PetscCall(VecHIPRestoreArrayRead(b, &barray));
1556:   PetscCall(VecHIPRestoreArrayWrite(x, &xarray));

1558:   PetscCall(PetscLogGpuTimeEnd());
1559:   PetscCall(PetscLogGpuFlops(2.0 * aij->nz - fact->rmap->n));
1560:   PetscFunctionReturn(PETSC_SUCCESS);
1561: }

1563: static PetscErrorCode MatICCFactorNumeric_SeqAIJHIPSPARSE_ICC0(Mat fact, Mat A, const MatFactorInfo *info)
1564: {
1565:   Mat_SeqAIJHIPSPARSETriFactors *fs    = (Mat_SeqAIJHIPSPARSETriFactors *)fact->spptr;
1566:   Mat_SeqAIJ                    *aij   = (Mat_SeqAIJ *)fact->data;
1567:   Mat_SeqAIJHIPSPARSE           *Acusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
1568:   CsrMatrix                     *Acsr;
1569:   PetscInt                       m, nz;
1570:   PetscBool                      flg;

1572:   PetscFunctionBegin;
1573:   if (PetscDefined(USE_DEBUG)) {
1574:     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
1575:     PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Expected MATSEQAIJHIPSPARSE, but input is %s", ((PetscObject)A)->type_name);
1576:   }

1578:   /* Copy A's value to fact */
1579:   m  = fact->rmap->n;
1580:   nz = aij->nz;
1581:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
1582:   Acsr = (CsrMatrix *)Acusp->mat->mat;
1583:   PetscCallHIP(hipMemcpyAsync(fs->csrVal, Acsr->values->data().get(), sizeof(PetscScalar) * nz, hipMemcpyDeviceToDevice, PetscDefaultHipStream));

1585:   /* Factorize fact inplace */
1586:   /* Function csric02() only takes the lower triangular part of matrix A to perform factorization.
1587:      The matrix type must be HIPSPARSE_MATRIX_TYPE_GENERAL, the fill mode and diagonal type are ignored,
1588:      and the strictly upper triangular part is ignored and never touched. It does not matter if A is Hermitian or not.
1589:      In other words, from the point of view of csric02() A is Hermitian and only the lower triangular part is provided.
1590:    */
1591:   if (m) PetscCallHIPSPARSE(hipsparseXcsric02(fs->handle, m, nz, fs->matDescr_M, fs->csrVal, fs->csrRowPtr, fs->csrColIdx, fs->ic0Info_M, fs->policy_M, fs->factBuffer_M));
1592:   if (PetscDefined(USE_DEBUG)) {
1593:     int               numerical_zero;
1594:     hipsparseStatus_t status;
1595:     status = hipsparseXcsric02_zeroPivot(fs->handle, fs->ic0Info_M, &numerical_zero);
1596:     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);
1597:   }

1599:   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));

1601:   /* Note that hipsparse reports this error if we use double and HIPSPARSE_OPERATION_CONJUGATE_TRANSPOSE
1602:     ** On entry to hipsparseSpSV_analysis(): conjugate transpose (opA) is not supported for matA data type, current -> CUDA_R_64F
1603:   */
1604:   PetscCallHIPSPARSE(hipsparseSpSV_analysis(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Lt, fs->spsvBuffer_Lt));

1606:   fact->offloadmask            = PETSC_OFFLOAD_GPU;
1607:   fact->ops->solve             = MatSolve_SeqAIJHIPSPARSE_ICC0;
1608:   fact->ops->solvetranspose    = MatSolve_SeqAIJHIPSPARSE_ICC0;
1609:   fact->ops->matsolve          = NULL;
1610:   fact->ops->matsolvetranspose = NULL;
1611:   PetscCall(PetscLogGpuFlops(fs->numericFactFlops));
1612:   PetscFunctionReturn(PETSC_SUCCESS);
1613: }

1615: static PetscErrorCode MatICCFactorSymbolic_SeqAIJHIPSPARSE_ICC0(Mat fact, Mat A, IS perm, const MatFactorInfo *info)
1616: {
1617:   Mat_SeqAIJHIPSPARSETriFactors *fs  = (Mat_SeqAIJHIPSPARSETriFactors *)fact->spptr;
1618:   Mat_SeqAIJ                    *aij = (Mat_SeqAIJ *)fact->data;
1619:   PetscInt                       m, nz;

1621:   PetscFunctionBegin;
1622:   if (PetscDefined(USE_DEBUG)) {
1623:     PetscInt  i;
1624:     PetscBool flg, missing;

1626:     PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
1627:     PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Expected MATSEQAIJHIPSPARSE, but input is %s", ((PetscObject)A)->type_name);
1628:     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);
1629:     PetscCall(MatMissingDiagonal(A, &missing, &i));
1630:     PetscCheck(!missing, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "Matrix is missing diagonal entry %" PetscInt_FMT, i);
1631:   }

1633:   /* Free the old stale stuff */
1634:   PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(&fs));

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

1641:   fact->offloadmask            = PETSC_OFFLOAD_BOTH;
1642:   fact->factortype             = MAT_FACTOR_ICC;
1643:   fact->info.factor_mallocs    = 0;
1644:   fact->info.fill_ratio_given  = info->fill;
1645:   fact->info.fill_ratio_needed = 1.0;

1647:   aij->row = NULL;
1648:   aij->col = NULL;

1650:   /* ====================================================================== */
1651:   /* Copy A's i, j to fact and also allocate the value array of fact.       */
1652:   /* We'll do in-place factorization on fact                                */
1653:   /* ====================================================================== */
1654:   const int *Ai, *Aj;

1656:   m  = fact->rmap->n;
1657:   nz = aij->nz;

1659:   PetscCallHIP(hipMalloc((void **)&fs->csrRowPtr, sizeof(int) * (m + 1)));
1660:   PetscCallHIP(hipMalloc((void **)&fs->csrColIdx, sizeof(int) * nz));
1661:   PetscCallHIP(hipMalloc((void **)&fs->csrVal, sizeof(PetscScalar) * nz));
1662:   PetscCall(MatSeqAIJHIPSPARSEGetIJ(A, PETSC_FALSE, &Ai, &Aj)); /* Do not use compressed Ai */
1663:   PetscCallHIP(hipMemcpyAsync(fs->csrRowPtr, Ai, sizeof(int) * (m + 1), hipMemcpyDeviceToDevice, PetscDefaultHipStream));
1664:   PetscCallHIP(hipMemcpyAsync(fs->csrColIdx, Aj, sizeof(int) * nz, hipMemcpyDeviceToDevice, PetscDefaultHipStream));

1666:   /* ====================================================================== */
1667:   /* Create mat descriptors for M, L                                        */
1668:   /* ====================================================================== */
1669:   hipsparseFillMode_t fillMode;
1670:   hipsparseDiagType_t diagType;

1672:   PetscCallHIPSPARSE(hipsparseCreateMatDescr(&fs->matDescr_M));
1673:   PetscCallHIPSPARSE(hipsparseSetMatIndexBase(fs->matDescr_M, HIPSPARSE_INDEX_BASE_ZERO));
1674:   PetscCallHIPSPARSE(hipsparseSetMatType(fs->matDescr_M, HIPSPARSE_MATRIX_TYPE_GENERAL));

1676:   /* https://docs.amd.com/bundle/hipSPARSE-Documentation---hipSPARSE-documentation/page/usermanual.html/#hipsparse_8h_1a79e036b6c0680cb37e2aa53d3542a054
1677:     hipsparseDiagType_t: This type indicates if the matrix diagonal entries are unity. The diagonal elements are always
1678:     assumed to be present, but if HIPSPARSE_DIAG_TYPE_UNIT is passed to an API routine, then the routine assumes that
1679:     all diagonal entries are unity and will not read or modify those entries. Note that in this case the routine
1680:     assumes the diagonal entries are equal to one, regardless of what those entries are actually set to in memory.
1681:   */
1682:   fillMode = HIPSPARSE_FILL_MODE_LOWER;
1683:   diagType = HIPSPARSE_DIAG_TYPE_NON_UNIT;
1684:   PetscCallHIPSPARSE(hipsparseCreateCsr(&fs->spMatDescr_L, m, m, nz, fs->csrRowPtr, fs->csrColIdx, fs->csrVal, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
1685:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_L, HIPSPARSE_SPMAT_FILL_MODE, &fillMode, sizeof(fillMode)));
1686:   PetscCallHIPSPARSE(hipsparseSpMatSetAttribute(fs->spMatDescr_L, HIPSPARSE_SPMAT_DIAG_TYPE, &diagType, sizeof(diagType)));

1688:   /* ========================================================================= */
1689:   /* Query buffer sizes for csric0, SpSV of L and Lt, and allocate buffers     */
1690:   /* ========================================================================= */
1691:   PetscCallHIPSPARSE(hipsparseCreateCsric02Info(&fs->ic0Info_M));
1692:   if (m) PetscCallHIPSPARSE(hipsparseXcsric02_bufferSize(fs->handle, m, nz, fs->matDescr_M, fs->csrVal, fs->csrRowPtr, fs->csrColIdx, fs->ic0Info_M, &fs->factBufferSize_M));

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

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

1700:   PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_L));
1701:   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));

1703:   PetscCallHIPSPARSE(hipsparseSpSV_createDescr(&fs->spsvDescr_Lt));
1704:   PetscCallHIPSPARSE(hipsparseSpSV_bufferSize(fs->handle, HIPSPARSE_OPERATION_TRANSPOSE, &PETSC_HIPSPARSE_ONE, fs->spMatDescr_L, fs->dnVecDescr_X, fs->dnVecDescr_Y, hipsparse_scalartype, HIPSPARSE_SPSV_ALG_DEFAULT, fs->spsvDescr_Lt, &fs->spsvBufferSize_Lt));

1706:   /* To save device memory, we make the factorization buffer share with one of the solver buffer.
1707:      See also comments in `MatILUFactorSymbolic_SeqAIJHIPSPARSE_ILU0()`.
1708:    */
1709:   if (fs->spsvBufferSize_L > fs->spsvBufferSize_Lt) {
1710:     PetscCallHIP(hipMalloc((void **)&fs->factBuffer_M, PetscMax(fs->spsvBufferSize_L, (size_t)fs->factBufferSize_M)));
1711:     fs->spsvBuffer_L = fs->factBuffer_M;
1712:     PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_Lt, fs->spsvBufferSize_Lt));
1713:   } else {
1714:     PetscCallHIP(hipMalloc((void **)&fs->factBuffer_M, PetscMax(fs->spsvBufferSize_Lt, (size_t)fs->factBufferSize_M)));
1715:     fs->spsvBuffer_Lt = fs->factBuffer_M;
1716:     PetscCallHIP(hipMalloc((void **)&fs->spsvBuffer_L, fs->spsvBufferSize_L));
1717:   }

1719:   /* ========================================================================== */
1720:   /* Perform analysis of ic0 on M                                               */
1721:   /* The lower triangular part of M has the same sparsity pattern as L          */
1722:   /* ========================================================================== */
1723:   int structural_zero;

1725:   fs->policy_M = HIPSPARSE_SOLVE_POLICY_USE_LEVEL;
1726:   if (m) PetscCallHIPSPARSE(hipsparseXcsric02_analysis(fs->handle, m, nz, fs->matDescr_M, fs->csrVal, fs->csrRowPtr, fs->csrColIdx, fs->ic0Info_M, fs->policy_M, fs->factBuffer_M));
1727:   if (PetscDefined(USE_DEBUG)) {
1728:     hipsparseStatus_t status;
1729:     /* Function hipsparseXcsric02_zeroPivot() is a blocking call. It calls hipDeviceSynchronize() to make sure all previous kernels are done. */
1730:     status = hipsparseXcsric02_zeroPivot(fs->handle, fs->ic0Info_M, &structural_zero);
1731:     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);
1732:   }

1734:   /* Estimate FLOPs of the numeric factorization */
1735:   {
1736:     Mat_SeqAIJ    *Aseq = (Mat_SeqAIJ *)A->data;
1737:     PetscInt      *Ai, nzRow, nzLeft;
1738:     PetscLogDouble flops = 0.0;

1740:     Ai = Aseq->i;
1741:     for (PetscInt i = 0; i < m; i++) {
1742:       nzRow = Ai[i + 1] - Ai[i];
1743:       if (nzRow > 1) {
1744:         /* We want to eliminate nonzeros left to the diagonal one by one. Assume each time, nonzeros right
1745:           and include the eliminated one will be updated, which incurs a multiplication and an addition.
1746:         */
1747:         nzLeft = (nzRow - 1) / 2;
1748:         flops += nzLeft * (2.0 * nzRow - nzLeft + 1);
1749:       }
1750:     }
1751:     fs->numericFactFlops = flops;
1752:   }
1753:   fact->ops->choleskyfactornumeric = MatICCFactorNumeric_SeqAIJHIPSPARSE_ICC0;
1754:   PetscFunctionReturn(PETSC_SUCCESS);
1755: }
1756: #endif

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

1762:   PetscFunctionBegin;
1763: #if PETSC_PKG_HIP_VERSION_GE(4, 5, 0)
1764:   PetscBool row_identity = PETSC_FALSE, col_identity = PETSC_FALSE;
1765:   if (hipsparseTriFactors->factorizeOnDevice) {
1766:     PetscCall(ISIdentity(isrow, &row_identity));
1767:     PetscCall(ISIdentity(iscol, &col_identity));
1768:   }
1769:   if (!info->levels && row_identity && col_identity) PetscCall(MatILUFactorSymbolic_SeqAIJHIPSPARSE_ILU0(B, A, isrow, iscol, info));
1770:   else
1771: #endif
1772:   {
1773:     PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(&hipsparseTriFactors));
1774:     PetscCall(MatILUFactorSymbolic_SeqAIJ(B, A, isrow, iscol, info));
1775:     B->ops->lufactornumeric = MatLUFactorNumeric_SeqAIJHIPSPARSE;
1776:   }
1777:   PetscFunctionReturn(PETSC_SUCCESS);
1778: }

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

1784:   PetscFunctionBegin;
1785:   PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(&hipsparseTriFactors));
1786:   PetscCall(MatLUFactorSymbolic_SeqAIJ(B, A, isrow, iscol, info));
1787:   B->ops->lufactornumeric = MatLUFactorNumeric_SeqAIJHIPSPARSE;
1788:   PetscFunctionReturn(PETSC_SUCCESS);
1789: }

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

1795:   PetscFunctionBegin;
1796: #if PETSC_PKG_HIP_VERSION_GE(4, 5, 0)
1797:   PetscBool perm_identity = PETSC_FALSE;
1798:   if (hipsparseTriFactors->factorizeOnDevice) PetscCall(ISIdentity(perm, &perm_identity));
1799:   if (!info->levels && perm_identity) PetscCall(MatICCFactorSymbolic_SeqAIJHIPSPARSE_ICC0(B, A, perm, info));
1800:   else
1801: #endif
1802:   {
1803:     PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(&hipsparseTriFactors));
1804:     PetscCall(MatICCFactorSymbolic_SeqAIJ(B, A, perm, info));
1805:     B->ops->choleskyfactornumeric = MatCholeskyFactorNumeric_SeqAIJHIPSPARSE;
1806:   }
1807:   PetscFunctionReturn(PETSC_SUCCESS);
1808: }

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

1814:   PetscFunctionBegin;
1815:   PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(&hipsparseTriFactors));
1816:   PetscCall(MatCholeskyFactorSymbolic_SeqAIJ(B, A, perm, info));
1817:   B->ops->choleskyfactornumeric = MatCholeskyFactorNumeric_SeqAIJHIPSPARSE;
1818:   PetscFunctionReturn(PETSC_SUCCESS);
1819: }

1821: PetscErrorCode MatFactorGetSolverType_seqaij_hipsparse(Mat A, MatSolverType *type)
1822: {
1823:   PetscFunctionBegin;
1824:   *type = MATSOLVERHIPSPARSE;
1825:   PetscFunctionReturn(PETSC_SUCCESS);
1826: }

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

1836:   Level: beginner

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

1841: PETSC_EXTERN PetscErrorCode MatGetFactor_seqaijhipsparse_hipsparse(Mat A, MatFactorType ftype, Mat *B)
1842: {
1843:   PetscInt  n = A->rmap->n;
1844:   PetscBool factOnDevice, factOnHost;
1845:   char     *prefix;
1846:   char      factPlace[32] = "device"; /* the default */

1848:   PetscFunctionBegin;
1849:   PetscCall(MatCreate(PetscObjectComm((PetscObject)A), B));
1850:   PetscCall(MatSetSizes(*B, n, n, n, n));
1851:   (*B)->factortype = ftype;
1852:   PetscCall(MatSetType(*B, MATSEQAIJHIPSPARSE));

1854:   prefix = (*B)->factorprefix ? (*B)->factorprefix : ((PetscObject)A)->prefix;
1855:   PetscOptionsBegin(PetscObjectComm((PetscObject)(*B)), prefix, "MatGetFactor", "Mat");
1856:   PetscCall(PetscOptionsString("-mat_factor_bind_factorization", "Do matrix factorization on host or device when possible", "MatGetFactor", NULL, factPlace, sizeof(factPlace), NULL));
1857:   PetscOptionsEnd();
1858:   PetscCall(PetscStrcasecmp("device", factPlace, &factOnDevice));
1859:   PetscCall(PetscStrcasecmp("host", factPlace, &factOnHost));
1860:   PetscCheck(factOnDevice || factOnHost, PetscObjectComm((PetscObject)(*B)), PETSC_ERR_ARG_OUTOFRANGE, "Wrong option %s to -mat_factor_bind_factorization <string>. Only host and device are allowed", factPlace);
1861:   ((Mat_SeqAIJHIPSPARSETriFactors *)(*B)->spptr)->factorizeOnDevice = factOnDevice;

1863:   if (A->boundtocpu && A->bindingpropagates) PetscCall(MatBindToCPU(*B, PETSC_TRUE));
1864:   if (ftype == MAT_FACTOR_LU || ftype == MAT_FACTOR_ILU || ftype == MAT_FACTOR_ILUDT) {
1865:     PetscCall(MatSetBlockSizesFromMats(*B, A, A));
1866:     if (!A->boundtocpu) {
1867:       (*B)->ops->ilufactorsymbolic = MatILUFactorSymbolic_SeqAIJHIPSPARSE;
1868:       (*B)->ops->lufactorsymbolic  = MatLUFactorSymbolic_SeqAIJHIPSPARSE;
1869:     } else {
1870:       (*B)->ops->ilufactorsymbolic = MatILUFactorSymbolic_SeqAIJ;
1871:       (*B)->ops->lufactorsymbolic  = MatLUFactorSymbolic_SeqAIJ;
1872:     }
1873:     PetscCall(PetscStrallocpy(MATORDERINGND, (char **)&(*B)->preferredordering[MAT_FACTOR_LU]));
1874:     PetscCall(PetscStrallocpy(MATORDERINGNATURAL, (char **)&(*B)->preferredordering[MAT_FACTOR_ILU]));
1875:     PetscCall(PetscStrallocpy(MATORDERINGNATURAL, (char **)&(*B)->preferredordering[MAT_FACTOR_ILUDT]));
1876:   } else if (ftype == MAT_FACTOR_CHOLESKY || ftype == MAT_FACTOR_ICC) {
1877:     if (!A->boundtocpu) {
1878:       (*B)->ops->iccfactorsymbolic      = MatICCFactorSymbolic_SeqAIJHIPSPARSE;
1879:       (*B)->ops->choleskyfactorsymbolic = MatCholeskyFactorSymbolic_SeqAIJHIPSPARSE;
1880:     } else {
1881:       (*B)->ops->iccfactorsymbolic      = MatICCFactorSymbolic_SeqAIJ;
1882:       (*B)->ops->choleskyfactorsymbolic = MatCholeskyFactorSymbolic_SeqAIJ;
1883:     }
1884:     PetscCall(PetscStrallocpy(MATORDERINGND, (char **)&(*B)->preferredordering[MAT_FACTOR_CHOLESKY]));
1885:     PetscCall(PetscStrallocpy(MATORDERINGNATURAL, (char **)&(*B)->preferredordering[MAT_FACTOR_ICC]));
1886:   } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "Factor type not supported for HIPSPARSE Matrix Types");

1888:   PetscCall(MatSeqAIJSetPreallocation(*B, MAT_SKIP_ALLOCATION, NULL));
1889:   (*B)->canuseordering = PETSC_TRUE;
1890:   PetscCall(PetscObjectComposeFunction((PetscObject)(*B), "MatFactorGetSolverType_C", MatFactorGetSolverType_seqaij_hipsparse));
1891:   PetscFunctionReturn(PETSC_SUCCESS);
1892: }

1894: static PetscErrorCode MatSeqAIJHIPSPARSECopyFromGPU(Mat A)
1895: {
1896:   Mat_SeqAIJ          *a    = (Mat_SeqAIJ *)A->data;
1897:   Mat_SeqAIJHIPSPARSE *cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
1898: #if PETSC_PKG_HIP_VERSION_GE(4, 5, 0)
1899:   Mat_SeqAIJHIPSPARSETriFactors *fs = (Mat_SeqAIJHIPSPARSETriFactors *)A->spptr;
1900: #endif

1902:   PetscFunctionBegin;
1903:   if (A->offloadmask == PETSC_OFFLOAD_GPU) {
1904:     PetscCall(PetscLogEventBegin(MAT_HIPSPARSECopyFromGPU, A, 0, 0, 0));
1905:     if (A->factortype == MAT_FACTOR_NONE) {
1906:       CsrMatrix *matrix = (CsrMatrix *)cusp->mat->mat;
1907:       PetscCallHIP(hipMemcpy(a->a, matrix->values->data().get(), a->nz * sizeof(PetscScalar), hipMemcpyDeviceToHost));
1908:     }
1909: #if PETSC_PKG_HIP_VERSION_GE(4, 5, 0)
1910:     else if (fs->csrVal) {
1911:       /* We have a factorized matrix on device and are able to copy it to host */
1912:       PetscCallHIP(hipMemcpy(a->a, fs->csrVal, a->nz * sizeof(PetscScalar), hipMemcpyDeviceToHost));
1913:     }
1914: #endif
1915:     else
1916:       SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "No support for copying this type of factorized matrix from device to host");
1917:     PetscCall(PetscLogGpuToCpu(a->nz * sizeof(PetscScalar)));
1918:     PetscCall(PetscLogEventEnd(MAT_HIPSPARSECopyFromGPU, A, 0, 0, 0));
1919:     A->offloadmask = PETSC_OFFLOAD_BOTH;
1920:   }
1921:   PetscFunctionReturn(PETSC_SUCCESS);
1922: }

1924: static PetscErrorCode MatSeqAIJGetArray_SeqAIJHIPSPARSE(Mat A, PetscScalar *array[])
1925: {
1926:   PetscFunctionBegin;
1927:   PetscCall(MatSeqAIJHIPSPARSECopyFromGPU(A));
1928:   *array = ((Mat_SeqAIJ *)A->data)->a;
1929:   PetscFunctionReturn(PETSC_SUCCESS);
1930: }

1932: static PetscErrorCode MatSeqAIJRestoreArray_SeqAIJHIPSPARSE(Mat A, PetscScalar *array[])
1933: {
1934:   PetscFunctionBegin;
1935:   A->offloadmask = PETSC_OFFLOAD_CPU;
1936:   *array         = NULL;
1937:   PetscFunctionReturn(PETSC_SUCCESS);
1938: }

1940: static PetscErrorCode MatSeqAIJGetArrayRead_SeqAIJHIPSPARSE(Mat A, const PetscScalar *array[])
1941: {
1942:   PetscFunctionBegin;
1943:   PetscCall(MatSeqAIJHIPSPARSECopyFromGPU(A));
1944:   *array = ((Mat_SeqAIJ *)A->data)->a;
1945:   PetscFunctionReturn(PETSC_SUCCESS);
1946: }

1948: static PetscErrorCode MatSeqAIJRestoreArrayRead_SeqAIJHIPSPARSE(Mat A, const PetscScalar *array[])
1949: {
1950:   PetscFunctionBegin;
1951:   *array = NULL;
1952:   PetscFunctionReturn(PETSC_SUCCESS);
1953: }

1955: static PetscErrorCode MatSeqAIJGetArrayWrite_SeqAIJHIPSPARSE(Mat A, PetscScalar *array[])
1956: {
1957:   PetscFunctionBegin;
1958:   *array = ((Mat_SeqAIJ *)A->data)->a;
1959:   PetscFunctionReturn(PETSC_SUCCESS);
1960: }

1962: static PetscErrorCode MatSeqAIJRestoreArrayWrite_SeqAIJHIPSPARSE(Mat A, PetscScalar *array[])
1963: {
1964:   PetscFunctionBegin;
1965:   A->offloadmask = PETSC_OFFLOAD_CPU;
1966:   *array         = NULL;
1967:   PetscFunctionReturn(PETSC_SUCCESS);
1968: }

1970: static PetscErrorCode MatSeqAIJGetCSRAndMemType_SeqAIJHIPSPARSE(Mat A, const PetscInt **i, const PetscInt **j, PetscScalar **a, PetscMemType *mtype)
1971: {
1972:   Mat_SeqAIJHIPSPARSE *cusp;
1973:   CsrMatrix           *matrix;

1975:   PetscFunctionBegin;
1976:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
1977:   PetscCheck(A->factortype == MAT_FACTOR_NONE, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "Not for factored matrix");
1978:   cusp = static_cast<Mat_SeqAIJHIPSPARSE *>(A->spptr);
1979:   PetscCheck(cusp != NULL, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONGSTATE, "cusp is NULL");
1980:   matrix = (CsrMatrix *)cusp->mat->mat;

1982:   if (i) {
1983: #if !defined(PETSC_USE_64BIT_INDICES)
1984:     *i = matrix->row_offsets->data().get();
1985: #else
1986:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "hipSparse does not supported 64-bit indices");
1987: #endif
1988:   }
1989:   if (j) {
1990: #if !defined(PETSC_USE_64BIT_INDICES)
1991:     *j = matrix->column_indices->data().get();
1992: #else
1993:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "hipSparse does not supported 64-bit indices");
1994: #endif
1995:   }
1996:   if (a) *a = matrix->values->data().get();
1997:   if (mtype) *mtype = PETSC_MEMTYPE_HIP;
1998:   PetscFunctionReturn(PETSC_SUCCESS);
1999: }

2001: PETSC_INTERN PetscErrorCode MatSeqAIJHIPSPARSECopyToGPU(Mat A)
2002: {
2003:   Mat_SeqAIJHIPSPARSE           *hipsparsestruct = (Mat_SeqAIJHIPSPARSE *)A->spptr;
2004:   Mat_SeqAIJHIPSPARSEMultStruct *matstruct       = hipsparsestruct->mat;
2005:   Mat_SeqAIJ                    *a               = (Mat_SeqAIJ *)A->data;
2006:   PetscBool                      both            = PETSC_TRUE;
2007:   PetscInt                       m               = A->rmap->n, *ii, *ridx, tmp;

2009:   PetscFunctionBegin;
2010:   PetscCheck(!A->boundtocpu, PETSC_COMM_SELF, PETSC_ERR_GPU, "Cannot copy to GPU");
2011:   if (A->offloadmask == PETSC_OFFLOAD_UNALLOCATED || A->offloadmask == PETSC_OFFLOAD_CPU) {
2012:     if (A->nonzerostate == hipsparsestruct->nonzerostate && hipsparsestruct->format == MAT_HIPSPARSE_CSR) { /* Copy values only */
2013:       CsrMatrix *matrix;
2014:       matrix = (CsrMatrix *)hipsparsestruct->mat->mat;

2016:       PetscCheck(!a->nz || a->a, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CSR values");
2017:       PetscCall(PetscLogEventBegin(MAT_HIPSPARSECopyToGPU, A, 0, 0, 0));
2018:       matrix->values->assign(a->a, a->a + a->nz);
2019:       PetscCallHIP(WaitForHIP());
2020:       PetscCall(PetscLogCpuToGpu((a->nz) * sizeof(PetscScalar)));
2021:       PetscCall(PetscLogEventEnd(MAT_HIPSPARSECopyToGPU, A, 0, 0, 0));
2022:       PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(A, PETSC_FALSE));
2023:     } else {
2024:       PetscInt nnz;
2025:       PetscCall(PetscLogEventBegin(MAT_HIPSPARSECopyToGPU, A, 0, 0, 0));
2026:       PetscCall(MatSeqAIJHIPSPARSEMultStruct_Destroy(&hipsparsestruct->mat, hipsparsestruct->format));
2027:       PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(A, PETSC_TRUE));
2028:       delete hipsparsestruct->workVector;
2029:       delete hipsparsestruct->rowoffsets_gpu;
2030:       hipsparsestruct->workVector     = NULL;
2031:       hipsparsestruct->rowoffsets_gpu = NULL;
2032:       try {
2033:         if (a->compressedrow.use) {
2034:           m    = a->compressedrow.nrows;
2035:           ii   = a->compressedrow.i;
2036:           ridx = a->compressedrow.rindex;
2037:         } else {
2038:           m    = A->rmap->n;
2039:           ii   = a->i;
2040:           ridx = NULL;
2041:         }
2042:         PetscCheck(ii, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CSR row data");
2043:         if (!a->a) {
2044:           nnz  = ii[m];
2045:           both = PETSC_FALSE;
2046:         } else nnz = a->nz;
2047:         PetscCheck(!nnz || a->j, PETSC_COMM_SELF, PETSC_ERR_GPU, "Missing CSR column data");

2049:         /* create hipsparse matrix */
2050:         hipsparsestruct->nrows = m;
2051:         matstruct              = new Mat_SeqAIJHIPSPARSEMultStruct;
2052:         PetscCallHIPSPARSE(hipsparseCreateMatDescr(&matstruct->descr));
2053:         PetscCallHIPSPARSE(hipsparseSetMatIndexBase(matstruct->descr, HIPSPARSE_INDEX_BASE_ZERO));
2054:         PetscCallHIPSPARSE(hipsparseSetMatType(matstruct->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));

2056:         PetscCallHIP(hipMalloc((void **)&(matstruct->alpha_one), sizeof(PetscScalar)));
2057:         PetscCallHIP(hipMalloc((void **)&(matstruct->beta_zero), sizeof(PetscScalar)));
2058:         PetscCallHIP(hipMalloc((void **)&(matstruct->beta_one), sizeof(PetscScalar)));
2059:         PetscCallHIP(hipMemcpy(matstruct->alpha_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
2060:         PetscCallHIP(hipMemcpy(matstruct->beta_zero, &PETSC_HIPSPARSE_ZERO, sizeof(PetscScalar), hipMemcpyHostToDevice));
2061:         PetscCallHIP(hipMemcpy(matstruct->beta_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
2062:         PetscCallHIPSPARSE(hipsparseSetPointerMode(hipsparsestruct->handle, HIPSPARSE_POINTER_MODE_DEVICE));

2064:         /* Build a hybrid/ellpack matrix if this option is chosen for the storage */
2065:         if (hipsparsestruct->format == MAT_HIPSPARSE_CSR) {
2066:           /* set the matrix */
2067:           CsrMatrix *mat      = new CsrMatrix;
2068:           mat->num_rows       = m;
2069:           mat->num_cols       = A->cmap->n;
2070:           mat->num_entries    = nnz;
2071:           mat->row_offsets    = new THRUSTINTARRAY32(m + 1);
2072:           mat->column_indices = new THRUSTINTARRAY32(nnz);
2073:           mat->values         = new THRUSTARRAY(nnz);
2074:           mat->row_offsets->assign(ii, ii + m + 1);
2075:           mat->column_indices->assign(a->j, a->j + nnz);
2076:           if (a->a) mat->values->assign(a->a, a->a + nnz);

2078:           /* assign the pointer */
2079:           matstruct->mat = mat;
2080:           if (mat->num_rows) { /* hipsparse errors on empty matrices! */
2081:             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(), HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, /* row offset, col idx types due to THRUSTINTARRAY32 */
2082:                                                   HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
2083:           }
2084:         } else if (hipsparsestruct->format == MAT_HIPSPARSE_ELL || hipsparsestruct->format == MAT_HIPSPARSE_HYB) {
2085:           CsrMatrix *mat      = new CsrMatrix;
2086:           mat->num_rows       = m;
2087:           mat->num_cols       = A->cmap->n;
2088:           mat->num_entries    = nnz;
2089:           mat->row_offsets    = new THRUSTINTARRAY32(m + 1);
2090:           mat->column_indices = new THRUSTINTARRAY32(nnz);
2091:           mat->values         = new THRUSTARRAY(nnz);
2092:           mat->row_offsets->assign(ii, ii + m + 1);
2093:           mat->column_indices->assign(a->j, a->j + nnz);
2094:           if (a->a) mat->values->assign(a->a, a->a + nnz);

2096:           hipsparseHybMat_t hybMat;
2097:           PetscCallHIPSPARSE(hipsparseCreateHybMat(&hybMat));
2098:           hipsparseHybPartition_t partition = hipsparsestruct->format == MAT_HIPSPARSE_ELL ? HIPSPARSE_HYB_PARTITION_MAX : HIPSPARSE_HYB_PARTITION_AUTO;
2099:           PetscCallHIPSPARSE(hipsparse_csr2hyb(hipsparsestruct->handle, mat->num_rows, mat->num_cols, matstruct->descr, mat->values->data().get(), mat->row_offsets->data().get(), mat->column_indices->data().get(), hybMat, 0, partition));
2100:           /* assign the pointer */
2101:           matstruct->mat = hybMat;

2103:           if (mat) {
2104:             if (mat->values) delete (THRUSTARRAY *)mat->values;
2105:             if (mat->column_indices) delete (THRUSTINTARRAY32 *)mat->column_indices;
2106:             if (mat->row_offsets) delete (THRUSTINTARRAY32 *)mat->row_offsets;
2107:             delete (CsrMatrix *)mat;
2108:           }
2109:         }

2111:         /* assign the compressed row indices */
2112:         if (a->compressedrow.use) {
2113:           hipsparsestruct->workVector = new THRUSTARRAY(m);
2114:           matstruct->cprowIndices     = new THRUSTINTARRAY(m);
2115:           matstruct->cprowIndices->assign(ridx, ridx + m);
2116:           tmp = m;
2117:         } else {
2118:           hipsparsestruct->workVector = NULL;
2119:           matstruct->cprowIndices     = NULL;
2120:           tmp                         = 0;
2121:         }
2122:         PetscCall(PetscLogCpuToGpu(((m + 1) + (a->nz)) * sizeof(int) + tmp * sizeof(PetscInt) + (3 + (a->nz)) * sizeof(PetscScalar)));

2124:         /* assign the pointer */
2125:         hipsparsestruct->mat = matstruct;
2126:       } catch (char *ex) {
2127:         SETERRQ(PETSC_COMM_SELF, PETSC_ERR_LIB, "HIPSPARSE error: %s", ex);
2128:       }
2129:       PetscCallHIP(WaitForHIP());
2130:       PetscCall(PetscLogEventEnd(MAT_HIPSPARSECopyToGPU, A, 0, 0, 0));
2131:       hipsparsestruct->nonzerostate = A->nonzerostate;
2132:     }
2133:     if (both) A->offloadmask = PETSC_OFFLOAD_BOTH;
2134:   }
2135:   PetscFunctionReturn(PETSC_SUCCESS);
2136: }

2138: struct VecHIPPlusEquals {
2139:   template <typename Tuple>
2140:   __host__ __device__ void operator()(Tuple t)
2141:   {
2142:     thrust::get<1>(t) = thrust::get<1>(t) + thrust::get<0>(t);
2143:   }
2144: };

2146: struct VecHIPEquals {
2147:   template <typename Tuple>
2148:   __host__ __device__ void operator()(Tuple t)
2149:   {
2150:     thrust::get<1>(t) = thrust::get<0>(t);
2151:   }
2152: };

2154: struct VecHIPEqualsReverse {
2155:   template <typename Tuple>
2156:   __host__ __device__ void operator()(Tuple t)
2157:   {
2158:     thrust::get<0>(t) = thrust::get<1>(t);
2159:   }
2160: };

2162: struct MatMatHipsparse {
2163:   PetscBool             cisdense;
2164:   PetscScalar          *Bt;
2165:   Mat                   X;
2166:   PetscBool             reusesym; /* Hipsparse does not have split symbolic and numeric phases for sparse matmat operations */
2167:   PetscLogDouble        flops;
2168:   CsrMatrix            *Bcsr;
2169:   hipsparseSpMatDescr_t matSpBDescr;
2170:   PetscBool             initialized; /* C = alpha op(A) op(B) + beta C */
2171:   hipsparseDnMatDescr_t matBDescr;
2172:   hipsparseDnMatDescr_t matCDescr;
2173:   PetscInt              Blda, Clda; /* Record leading dimensions of B and C here to detect changes*/
2174: #if PETSC_PKG_HIP_VERSION_GE(5, 1, 0)
2175:   void *dBuffer4, *dBuffer5;
2176: #endif
2177:   size_t                 mmBufferSize;
2178:   void                  *mmBuffer, *mmBuffer2; /* SpGEMM WorkEstimation buffer */
2179:   hipsparseSpGEMMDescr_t spgemmDesc;
2180: };

2182: static PetscErrorCode MatDestroy_MatMatHipsparse(void *data)
2183: {
2184:   MatMatHipsparse *mmdata = (MatMatHipsparse *)data;

2186:   PetscFunctionBegin;
2187:   PetscCallHIP(hipFree(mmdata->Bt));
2188:   delete mmdata->Bcsr;
2189:   if (mmdata->matSpBDescr) PetscCallHIPSPARSE(hipsparseDestroySpMat(mmdata->matSpBDescr));
2190:   if (mmdata->matBDescr) PetscCallHIPSPARSE(hipsparseDestroyDnMat(mmdata->matBDescr));
2191:   if (mmdata->matCDescr) PetscCallHIPSPARSE(hipsparseDestroyDnMat(mmdata->matCDescr));
2192:   if (mmdata->spgemmDesc) PetscCallHIPSPARSE(hipsparseSpGEMM_destroyDescr(mmdata->spgemmDesc));
2193: #if PETSC_PKG_HIP_VERSION_GE(5, 1, 0)
2194:   if (mmdata->dBuffer4) PetscCallHIP(hipFree(mmdata->dBuffer4));
2195:   if (mmdata->dBuffer5) PetscCallHIP(hipFree(mmdata->dBuffer5));
2196: #endif
2197:   if (mmdata->mmBuffer) PetscCallHIP(hipFree(mmdata->mmBuffer));
2198:   if (mmdata->mmBuffer2) PetscCallHIP(hipFree(mmdata->mmBuffer2));
2199:   PetscCall(MatDestroy(&mmdata->X));
2200:   PetscCall(PetscFree(data));
2201:   PetscFunctionReturn(PETSC_SUCCESS);
2202: }

2204: static PetscErrorCode MatProductNumeric_SeqAIJHIPSPARSE_SeqDENSEHIP(Mat C)
2205: {
2206:   Mat_Product                   *product = C->product;
2207:   Mat                            A, B;
2208:   PetscInt                       m, n, blda, clda;
2209:   PetscBool                      flg, biship;
2210:   Mat_SeqAIJHIPSPARSE           *cusp;
2211:   hipsparseOperation_t           opA;
2212:   const PetscScalar             *barray;
2213:   PetscScalar                   *carray;
2214:   MatMatHipsparse               *mmdata;
2215:   Mat_SeqAIJHIPSPARSEMultStruct *mat;
2216:   CsrMatrix                     *csrmat;

2218:   PetscFunctionBegin;
2219:   MatCheckProduct(C, 1);
2220:   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data empty");
2221:   mmdata = (MatMatHipsparse *)product->data;
2222:   A      = product->A;
2223:   B      = product->B;
2224:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
2225:   PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
2226:   /* currently CopyToGpu does not copy if the matrix is bound to CPU
2227:      Instead of silently accepting the wrong answer, I prefer to raise the error */
2228:   PetscCheck(!A->boundtocpu, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Cannot bind to CPU a HIPSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
2229:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
2230:   cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
2231:   switch (product->type) {
2232:   case MATPRODUCT_AB:
2233:   case MATPRODUCT_PtAP:
2234:     mat = cusp->mat;
2235:     opA = HIPSPARSE_OPERATION_NON_TRANSPOSE;
2236:     m   = A->rmap->n;
2237:     n   = B->cmap->n;
2238:     break;
2239:   case MATPRODUCT_AtB:
2240:     if (!A->form_explicit_transpose) {
2241:       mat = cusp->mat;
2242:       opA = HIPSPARSE_OPERATION_TRANSPOSE;
2243:     } else {
2244:       PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(A));
2245:       mat = cusp->matTranspose;
2246:       opA = HIPSPARSE_OPERATION_NON_TRANSPOSE;
2247:     }
2248:     m = A->cmap->n;
2249:     n = B->cmap->n;
2250:     break;
2251:   case MATPRODUCT_ABt:
2252:   case MATPRODUCT_RARt:
2253:     mat = cusp->mat;
2254:     opA = HIPSPARSE_OPERATION_NON_TRANSPOSE;
2255:     m   = A->rmap->n;
2256:     n   = B->rmap->n;
2257:     break;
2258:   default:
2259:     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
2260:   }
2261:   PetscCheck(mat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing Mat_SeqAIJHIPSPARSEMultStruct");
2262:   csrmat = (CsrMatrix *)mat->mat;
2263:   /* if the user passed a CPU matrix, copy the data to the GPU */
2264:   PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQDENSEHIP, &biship));
2265:   if (!biship) { PetscCall(MatConvert(B, MATSEQDENSEHIP, MAT_INPLACE_MATRIX, &B)); }
2266:   PetscCall(MatDenseGetArrayReadAndMemType(B, &barray, nullptr));
2267:   PetscCall(MatDenseGetLDA(B, &blda));
2268:   if (product->type == MATPRODUCT_RARt || product->type == MATPRODUCT_PtAP) {
2269:     PetscCall(MatDenseGetArrayWriteAndMemType(mmdata->X, &carray, nullptr));
2270:     PetscCall(MatDenseGetLDA(mmdata->X, &clda));
2271:   } else {
2272:     PetscCall(MatDenseGetArrayWriteAndMemType(C, &carray, nullptr));
2273:     PetscCall(MatDenseGetLDA(C, &clda));
2274:   }

2276:   PetscCall(PetscLogGpuTimeBegin());
2277:   hipsparseOperation_t opB = (product->type == MATPRODUCT_ABt || product->type == MATPRODUCT_RARt) ? HIPSPARSE_OPERATION_TRANSPOSE : HIPSPARSE_OPERATION_NON_TRANSPOSE;
2278:   /* (re)allocate mmBuffer if not initialized or LDAs are different */
2279:   if (!mmdata->initialized || mmdata->Blda != blda || mmdata->Clda != clda) {
2280:     size_t mmBufferSize;
2281:     if (mmdata->initialized && mmdata->Blda != blda) {
2282:       PetscCallHIPSPARSE(hipsparseDestroyDnMat(mmdata->matBDescr));
2283:       mmdata->matBDescr = NULL;
2284:     }
2285:     if (!mmdata->matBDescr) {
2286:       PetscCallHIPSPARSE(hipsparseCreateDnMat(&mmdata->matBDescr, B->rmap->n, B->cmap->n, blda, (void *)barray, hipsparse_scalartype, HIPSPARSE_ORDER_COL));
2287:       mmdata->Blda = blda;
2288:     }
2289:     if (mmdata->initialized && mmdata->Clda != clda) {
2290:       PetscCallHIPSPARSE(hipsparseDestroyDnMat(mmdata->matCDescr));
2291:       mmdata->matCDescr = NULL;
2292:     }
2293:     if (!mmdata->matCDescr) { /* matCDescr is for C or mmdata->X */
2294:       PetscCallHIPSPARSE(hipsparseCreateDnMat(&mmdata->matCDescr, m, n, clda, (void *)carray, hipsparse_scalartype, HIPSPARSE_ORDER_COL));
2295:       mmdata->Clda = clda;
2296:     }
2297:     if (!mat->matDescr) {
2298:       PetscCallHIPSPARSE(hipsparseCreateCsr(&mat->matDescr, csrmat->num_rows, csrmat->num_cols, csrmat->num_entries, csrmat->row_offsets->data().get(), csrmat->column_indices->data().get(), csrmat->values->data().get(), HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, /* row offset, col idx types due to THRUSTINTARRAY32 */
2299:                                             HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
2300:     }
2301:     PetscCallHIPSPARSE(hipsparseSpMM_bufferSize(cusp->handle, opA, opB, mat->alpha_one, mat->matDescr, mmdata->matBDescr, mat->beta_zero, mmdata->matCDescr, hipsparse_scalartype, cusp->spmmAlg, &mmBufferSize));
2302:     if ((mmdata->mmBuffer && mmdata->mmBufferSize < mmBufferSize) || !mmdata->mmBuffer) {
2303:       PetscCallHIP(hipFree(mmdata->mmBuffer));
2304:       PetscCallHIP(hipMalloc(&mmdata->mmBuffer, mmBufferSize));
2305:       mmdata->mmBufferSize = mmBufferSize;
2306:     }
2307:     mmdata->initialized = PETSC_TRUE;
2308:   } else {
2309:     /* to be safe, always update pointers of the mats */
2310:     PetscCallHIPSPARSE(hipsparseSpMatSetValues(mat->matDescr, csrmat->values->data().get()));
2311:     PetscCallHIPSPARSE(hipsparseDnMatSetValues(mmdata->matBDescr, (void *)barray));
2312:     PetscCallHIPSPARSE(hipsparseDnMatSetValues(mmdata->matCDescr, (void *)carray));
2313:   }

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

2318:   PetscCall(PetscLogGpuTimeEnd());
2319:   PetscCall(PetscLogGpuFlops(n * 2.0 * csrmat->num_entries));
2320:   PetscCall(MatDenseRestoreArrayReadAndMemType(B, &barray));
2321:   if (product->type == MATPRODUCT_RARt) {
2322:     PetscCall(MatDenseRestoreArrayWriteAndMemType(mmdata->X, &carray));
2323:     PetscCall(MatMatMultNumeric_SeqDenseHIP_SeqDenseHIP_Internal(B, mmdata->X, C, PETSC_FALSE, PETSC_FALSE));
2324:   } else if (product->type == MATPRODUCT_PtAP) {
2325:     PetscCall(MatDenseRestoreArrayWriteAndMemType(mmdata->X, &carray));
2326:     PetscCall(MatMatMultNumeric_SeqDenseHIP_SeqDenseHIP_Internal(B, mmdata->X, C, PETSC_TRUE, PETSC_FALSE));
2327:   } else PetscCall(MatDenseRestoreArrayWriteAndMemType(C, &carray));
2328:   if (mmdata->cisdense) PetscCall(MatConvert(C, MATSEQDENSE, MAT_INPLACE_MATRIX, &C));
2329:   if (!biship) PetscCall(MatConvert(B, MATSEQDENSE, MAT_INPLACE_MATRIX, &B));
2330:   PetscFunctionReturn(PETSC_SUCCESS);
2331: }

2333: static PetscErrorCode MatProductSymbolic_SeqAIJHIPSPARSE_SeqDENSEHIP(Mat C)
2334: {
2335:   Mat_Product         *product = C->product;
2336:   Mat                  A, B;
2337:   PetscInt             m, n;
2338:   PetscBool            cisdense, flg;
2339:   MatMatHipsparse     *mmdata;
2340:   Mat_SeqAIJHIPSPARSE *cusp;

2342:   PetscFunctionBegin;
2343:   MatCheckProduct(C, 1);
2344:   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data not empty");
2345:   A = product->A;
2346:   B = product->B;
2347:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
2348:   PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
2349:   cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
2350:   PetscCheck(cusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");
2351:   switch (product->type) {
2352:   case MATPRODUCT_AB:
2353:     m = A->rmap->n;
2354:     n = B->cmap->n;
2355:     break;
2356:   case MATPRODUCT_AtB:
2357:     m = A->cmap->n;
2358:     n = B->cmap->n;
2359:     break;
2360:   case MATPRODUCT_ABt:
2361:     m = A->rmap->n;
2362:     n = B->rmap->n;
2363:     break;
2364:   case MATPRODUCT_PtAP:
2365:     m = B->cmap->n;
2366:     n = B->cmap->n;
2367:     break;
2368:   case MATPRODUCT_RARt:
2369:     m = B->rmap->n;
2370:     n = B->rmap->n;
2371:     break;
2372:   default:
2373:     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
2374:   }
2375:   PetscCall(MatSetSizes(C, m, n, m, n));
2376:   /* if C is of type MATSEQDENSE (CPU), perform the operation on the GPU and then copy on the CPU */
2377:   PetscCall(PetscObjectTypeCompare((PetscObject)C, MATSEQDENSE, &cisdense));
2378:   PetscCall(MatSetType(C, MATSEQDENSEHIP));

2380:   /* product data */
2381:   PetscCall(PetscNew(&mmdata));
2382:   mmdata->cisdense = cisdense;
2383:   /* for these products we need intermediate storage */
2384:   if (product->type == MATPRODUCT_RARt || product->type == MATPRODUCT_PtAP) {
2385:     PetscCall(MatCreate(PetscObjectComm((PetscObject)C), &mmdata->X));
2386:     PetscCall(MatSetType(mmdata->X, MATSEQDENSEHIP));
2387:     /* do not preallocate, since the first call to MatDenseHIPGetArray will preallocate on the GPU for us */
2388:     if (product->type == MATPRODUCT_RARt) PetscCall(MatSetSizes(mmdata->X, A->rmap->n, B->rmap->n, A->rmap->n, B->rmap->n));
2389:     else PetscCall(MatSetSizes(mmdata->X, A->rmap->n, B->cmap->n, A->rmap->n, B->cmap->n));
2390:   }
2391:   C->product->data       = mmdata;
2392:   C->product->destroy    = MatDestroy_MatMatHipsparse;
2393:   C->ops->productnumeric = MatProductNumeric_SeqAIJHIPSPARSE_SeqDENSEHIP;
2394:   PetscFunctionReturn(PETSC_SUCCESS);
2395: }

2397: static PetscErrorCode MatProductNumeric_SeqAIJHIPSPARSE_SeqAIJHIPSPARSE(Mat C)
2398: {
2399:   Mat_Product                   *product = C->product;
2400:   Mat                            A, B;
2401:   Mat_SeqAIJHIPSPARSE           *Acusp, *Bcusp, *Ccusp;
2402:   Mat_SeqAIJ                    *c = (Mat_SeqAIJ *)C->data;
2403:   Mat_SeqAIJHIPSPARSEMultStruct *Amat, *Bmat, *Cmat;
2404:   CsrMatrix                     *Acsr, *Bcsr, *Ccsr;
2405:   PetscBool                      flg;
2406:   MatProductType                 ptype;
2407:   MatMatHipsparse               *mmdata;
2408:   hipsparseSpMatDescr_t          BmatSpDescr;
2409:   hipsparseOperation_t           opA = HIPSPARSE_OPERATION_NON_TRANSPOSE, opB = HIPSPARSE_OPERATION_NON_TRANSPOSE; /* hipSPARSE spgemm doesn't support transpose yet */

2411:   PetscFunctionBegin;
2412:   MatCheckProduct(C, 1);
2413:   PetscCheck(C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data empty");
2414:   PetscCall(PetscObjectTypeCompare((PetscObject)C, MATSEQAIJHIPSPARSE, &flg));
2415:   PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for C of type %s", ((PetscObject)C)->type_name);
2416:   mmdata = (MatMatHipsparse *)C->product->data;
2417:   A      = product->A;
2418:   B      = product->B;
2419:   if (mmdata->reusesym) { /* this happens when api_user is true, meaning that the matrix values have been already computed in the MatProductSymbolic phase */
2420:     mmdata->reusesym = PETSC_FALSE;
2421:     Ccusp            = (Mat_SeqAIJHIPSPARSE *)C->spptr;
2422:     PetscCheck(Ccusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");
2423:     Cmat = Ccusp->mat;
2424:     PetscCheck(Cmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C mult struct for product type %s", MatProductTypes[C->product->type]);
2425:     Ccsr = (CsrMatrix *)Cmat->mat;
2426:     PetscCheck(Ccsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C CSR struct");
2427:     goto finalize;
2428:   }
2429:   if (!c->nz) goto finalize;
2430:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
2431:   PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
2432:   PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQAIJHIPSPARSE, &flg));
2433:   PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for B of type %s", ((PetscObject)B)->type_name);
2434:   PetscCheck(!A->boundtocpu, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONG, "Cannot bind to CPU a HIPSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
2435:   PetscCheck(!B->boundtocpu, PetscObjectComm((PetscObject)C), PETSC_ERR_ARG_WRONG, "Cannot bind to CPU a HIPSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
2436:   Acusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
2437:   Bcusp = (Mat_SeqAIJHIPSPARSE *)B->spptr;
2438:   Ccusp = (Mat_SeqAIJHIPSPARSE *)C->spptr;
2439:   PetscCheck(Acusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");
2440:   PetscCheck(Bcusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");
2441:   PetscCheck(Ccusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");
2442:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
2443:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(B));

2445:   ptype = product->type;
2446:   if (A->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_AtB) {
2447:     ptype = MATPRODUCT_AB;
2448:     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");
2449:   }
2450:   if (B->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_ABt) {
2451:     ptype = MATPRODUCT_AB;
2452:     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");
2453:   }
2454:   switch (ptype) {
2455:   case MATPRODUCT_AB:
2456:     Amat = Acusp->mat;
2457:     Bmat = Bcusp->mat;
2458:     break;
2459:   case MATPRODUCT_AtB:
2460:     Amat = Acusp->matTranspose;
2461:     Bmat = Bcusp->mat;
2462:     break;
2463:   case MATPRODUCT_ABt:
2464:     Amat = Acusp->mat;
2465:     Bmat = Bcusp->matTranspose;
2466:     break;
2467:   default:
2468:     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
2469:   }
2470:   Cmat = Ccusp->mat;
2471:   PetscCheck(Amat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A mult struct for product type %s", MatProductTypes[ptype]);
2472:   PetscCheck(Bmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B mult struct for product type %s", MatProductTypes[ptype]);
2473:   PetscCheck(Cmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C mult struct for product type %s", MatProductTypes[ptype]);
2474:   Acsr = (CsrMatrix *)Amat->mat;
2475:   Bcsr = mmdata->Bcsr ? mmdata->Bcsr : (CsrMatrix *)Bmat->mat; /* B may be in compressed row storage */
2476:   Ccsr = (CsrMatrix *)Cmat->mat;
2477:   PetscCheck(Acsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A CSR struct");
2478:   PetscCheck(Bcsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B CSR struct");
2479:   PetscCheck(Ccsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing C CSR struct");
2480:   PetscCall(PetscLogGpuTimeBegin());
2481: #if PETSC_PKG_HIP_VERSION_GE(5, 0, 0)
2482:   BmatSpDescr = mmdata->Bcsr ? mmdata->matSpBDescr : Bmat->matDescr; /* B may be in compressed row storage */
2483:   PetscCallHIPSPARSE(hipsparseSetPointerMode(Ccusp->handle, HIPSPARSE_POINTER_MODE_DEVICE));
2484:   #if PETSC_PKG_HIP_VERSION_GE(5, 1, 0)
2485:   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));
2486:   #else
2487:   PetscCallHIPSPARSE(hipsparseSpGEMM_compute(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, hipsparse_scalartype, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &mmdata->mmBufferSize, mmdata->mmBuffer));
2488:   PetscCallHIPSPARSE(hipsparseSpGEMM_copy(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, hipsparse_scalartype, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc));
2489:   #endif
2490: #else
2491:   PetscCallHIPSPARSE(hipsparse_csr_spgemm(Ccusp->handle, opA, opB, Acsr->num_rows, Bcsr->num_cols, Acsr->num_cols, Amat->descr, Acsr->num_entries, Acsr->values->data().get(), Acsr->row_offsets->data().get(), Acsr->column_indices->data().get(), Bmat->descr,
2492:                                           Bcsr->num_entries, Bcsr->values->data().get(), Bcsr->row_offsets->data().get(), Bcsr->column_indices->data().get(), Cmat->descr, Ccsr->values->data().get(), Ccsr->row_offsets->data().get(),
2493:                                           Ccsr->column_indices->data().get()));
2494: #endif
2495:   PetscCall(PetscLogGpuFlops(mmdata->flops));
2496:   PetscCallHIP(WaitForHIP());
2497:   PetscCall(PetscLogGpuTimeEnd());
2498:   C->offloadmask = PETSC_OFFLOAD_GPU;
2499: finalize:
2500:   /* shorter version of MatAssemblyEnd_SeqAIJ */
2501:   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));
2502:   PetscCall(PetscInfo(C, "Number of mallocs during MatSetValues() is 0\n"));
2503:   PetscCall(PetscInfo(C, "Maximum nonzeros in any row is %" PetscInt_FMT "\n", c->rmax));
2504:   c->reallocs = 0;
2505:   C->info.mallocs += 0;
2506:   C->info.nz_unneeded = 0;
2507:   C->assembled = C->was_assembled = PETSC_TRUE;
2508:   C->num_ass++;
2509:   PetscFunctionReturn(PETSC_SUCCESS);
2510: }

2512: static PetscErrorCode MatProductSymbolic_SeqAIJHIPSPARSE_SeqAIJHIPSPARSE(Mat C)
2513: {
2514:   Mat_Product                   *product = C->product;
2515:   Mat                            A, B;
2516:   Mat_SeqAIJHIPSPARSE           *Acusp, *Bcusp, *Ccusp;
2517:   Mat_SeqAIJ                    *a, *b, *c;
2518:   Mat_SeqAIJHIPSPARSEMultStruct *Amat, *Bmat, *Cmat;
2519:   CsrMatrix                     *Acsr, *Bcsr, *Ccsr;
2520:   PetscInt                       i, j, m, n, k;
2521:   PetscBool                      flg;
2522:   MatProductType                 ptype;
2523:   MatMatHipsparse               *mmdata;
2524:   PetscLogDouble                 flops;
2525:   PetscBool                      biscompressed, ciscompressed;
2526: #if PETSC_PKG_HIP_VERSION_GE(5, 0, 0)
2527:   int64_t               C_num_rows1, C_num_cols1, C_nnz1;
2528:   hipsparseSpMatDescr_t BmatSpDescr;
2529: #else
2530:   int cnz;
2531: #endif
2532:   hipsparseOperation_t opA = HIPSPARSE_OPERATION_NON_TRANSPOSE, opB = HIPSPARSE_OPERATION_NON_TRANSPOSE; /* hipSPARSE spgemm doesn't support transpose yet */

2534:   PetscFunctionBegin;
2535:   MatCheckProduct(C, 1);
2536:   PetscCheck(!C->product->data, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Product data not empty");
2537:   A = product->A;
2538:   B = product->B;
2539:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATSEQAIJHIPSPARSE, &flg));
2540:   PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for type %s", ((PetscObject)A)->type_name);
2541:   PetscCall(PetscObjectTypeCompare((PetscObject)B, MATSEQAIJHIPSPARSE, &flg));
2542:   PetscCheck(flg, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Not for B of type %s", ((PetscObject)B)->type_name);
2543:   a = (Mat_SeqAIJ *)A->data;
2544:   b = (Mat_SeqAIJ *)B->data;
2545:   /* product data */
2546:   PetscCall(PetscNew(&mmdata));
2547:   C->product->data    = mmdata;
2548:   C->product->destroy = MatDestroy_MatMatHipsparse;

2550:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
2551:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(B));
2552:   Acusp = (Mat_SeqAIJHIPSPARSE *)A->spptr; /* Access spptr after MatSeqAIJHIPSPARSECopyToGPU, not before */
2553:   Bcusp = (Mat_SeqAIJHIPSPARSE *)B->spptr;
2554:   PetscCheck(Acusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");
2555:   PetscCheck(Bcusp->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Only for MAT_HIPSPARSE_CSR format");

2557:   ptype = product->type;
2558:   if (A->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_AtB) {
2559:     ptype                                          = MATPRODUCT_AB;
2560:     product->symbolic_used_the_fact_A_is_symmetric = PETSC_TRUE;
2561:   }
2562:   if (B->symmetric == PETSC_BOOL3_TRUE && ptype == MATPRODUCT_ABt) {
2563:     ptype                                          = MATPRODUCT_AB;
2564:     product->symbolic_used_the_fact_B_is_symmetric = PETSC_TRUE;
2565:   }
2566:   biscompressed = PETSC_FALSE;
2567:   ciscompressed = PETSC_FALSE;
2568:   switch (ptype) {
2569:   case MATPRODUCT_AB:
2570:     m    = A->rmap->n;
2571:     n    = B->cmap->n;
2572:     k    = A->cmap->n;
2573:     Amat = Acusp->mat;
2574:     Bmat = Bcusp->mat;
2575:     if (a->compressedrow.use) ciscompressed = PETSC_TRUE;
2576:     if (b->compressedrow.use) biscompressed = PETSC_TRUE;
2577:     break;
2578:   case MATPRODUCT_AtB:
2579:     m = A->cmap->n;
2580:     n = B->cmap->n;
2581:     k = A->rmap->n;
2582:     PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(A));
2583:     Amat = Acusp->matTranspose;
2584:     Bmat = Bcusp->mat;
2585:     if (b->compressedrow.use) biscompressed = PETSC_TRUE;
2586:     break;
2587:   case MATPRODUCT_ABt:
2588:     m = A->rmap->n;
2589:     n = B->rmap->n;
2590:     k = A->cmap->n;
2591:     PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(B));
2592:     Amat = Acusp->mat;
2593:     Bmat = Bcusp->matTranspose;
2594:     if (a->compressedrow.use) ciscompressed = PETSC_TRUE;
2595:     break;
2596:   default:
2597:     SETERRQ(PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Unsupported product type %s", MatProductTypes[product->type]);
2598:   }

2600:   /* create hipsparse matrix */
2601:   PetscCall(MatSetSizes(C, m, n, m, n));
2602:   PetscCall(MatSetType(C, MATSEQAIJHIPSPARSE));
2603:   c     = (Mat_SeqAIJ *)C->data;
2604:   Ccusp = (Mat_SeqAIJHIPSPARSE *)C->spptr;
2605:   Cmat  = new Mat_SeqAIJHIPSPARSEMultStruct;
2606:   Ccsr  = new CsrMatrix;

2608:   c->compressedrow.use = ciscompressed;
2609:   if (c->compressedrow.use) { /* if a is in compressed row, than c will be in compressed row format */
2610:     c->compressedrow.nrows = a->compressedrow.nrows;
2611:     PetscCall(PetscMalloc2(c->compressedrow.nrows + 1, &c->compressedrow.i, c->compressedrow.nrows, &c->compressedrow.rindex));
2612:     PetscCall(PetscArraycpy(c->compressedrow.rindex, a->compressedrow.rindex, c->compressedrow.nrows));
2613:     Ccusp->workVector  = new THRUSTARRAY(c->compressedrow.nrows);
2614:     Cmat->cprowIndices = new THRUSTINTARRAY(c->compressedrow.nrows);
2615:     Cmat->cprowIndices->assign(c->compressedrow.rindex, c->compressedrow.rindex + c->compressedrow.nrows);
2616:   } else {
2617:     c->compressedrow.nrows  = 0;
2618:     c->compressedrow.i      = NULL;
2619:     c->compressedrow.rindex = NULL;
2620:     Ccusp->workVector       = NULL;
2621:     Cmat->cprowIndices      = NULL;
2622:   }
2623:   Ccusp->nrows      = ciscompressed ? c->compressedrow.nrows : m;
2624:   Ccusp->mat        = Cmat;
2625:   Ccusp->mat->mat   = Ccsr;
2626:   Ccsr->num_rows    = Ccusp->nrows;
2627:   Ccsr->num_cols    = n;
2628:   Ccsr->row_offsets = new THRUSTINTARRAY32(Ccusp->nrows + 1);
2629:   PetscCallHIPSPARSE(hipsparseCreateMatDescr(&Cmat->descr));
2630:   PetscCallHIPSPARSE(hipsparseSetMatIndexBase(Cmat->descr, HIPSPARSE_INDEX_BASE_ZERO));
2631:   PetscCallHIPSPARSE(hipsparseSetMatType(Cmat->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));
2632:   PetscCallHIP(hipMalloc((void **)&(Cmat->alpha_one), sizeof(PetscScalar)));
2633:   PetscCallHIP(hipMalloc((void **)&(Cmat->beta_zero), sizeof(PetscScalar)));
2634:   PetscCallHIP(hipMalloc((void **)&(Cmat->beta_one), sizeof(PetscScalar)));
2635:   PetscCallHIP(hipMemcpy(Cmat->alpha_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
2636:   PetscCallHIP(hipMemcpy(Cmat->beta_zero, &PETSC_HIPSPARSE_ZERO, sizeof(PetscScalar), hipMemcpyHostToDevice));
2637:   PetscCallHIP(hipMemcpy(Cmat->beta_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
2638:   if (!Ccsr->num_rows || !Ccsr->num_cols || !a->nz || !b->nz) { /* hipsparse raise errors in different calls when matrices have zero rows/columns! */
2639:     thrust::fill(thrust::device, Ccsr->row_offsets->begin(), Ccsr->row_offsets->end(), 0);
2640:     c->nz                = 0;
2641:     Ccsr->column_indices = new THRUSTINTARRAY32(c->nz);
2642:     Ccsr->values         = new THRUSTARRAY(c->nz);
2643:     goto finalizesym;
2644:   }

2646:   PetscCheck(Amat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A mult struct for product type %s", MatProductTypes[ptype]);
2647:   PetscCheck(Bmat, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B mult struct for product type %s", MatProductTypes[ptype]);
2648:   Acsr = (CsrMatrix *)Amat->mat;
2649:   if (!biscompressed) {
2650:     Bcsr        = (CsrMatrix *)Bmat->mat;
2651:     BmatSpDescr = Bmat->matDescr;
2652:   } else { /* we need to use row offsets for the full matrix */
2653:     CsrMatrix *cBcsr     = (CsrMatrix *)Bmat->mat;
2654:     Bcsr                 = new CsrMatrix;
2655:     Bcsr->num_rows       = B->rmap->n;
2656:     Bcsr->num_cols       = cBcsr->num_cols;
2657:     Bcsr->num_entries    = cBcsr->num_entries;
2658:     Bcsr->column_indices = cBcsr->column_indices;
2659:     Bcsr->values         = cBcsr->values;
2660:     if (!Bcusp->rowoffsets_gpu) {
2661:       Bcusp->rowoffsets_gpu = new THRUSTINTARRAY32(B->rmap->n + 1);
2662:       Bcusp->rowoffsets_gpu->assign(b->i, b->i + B->rmap->n + 1);
2663:       PetscCall(PetscLogCpuToGpu((B->rmap->n + 1) * sizeof(PetscInt)));
2664:     }
2665:     Bcsr->row_offsets = Bcusp->rowoffsets_gpu;
2666:     mmdata->Bcsr      = Bcsr;
2667:     if (Bcsr->num_rows && Bcsr->num_cols) {
2668:       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(), HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
2669:     }
2670:     BmatSpDescr = mmdata->matSpBDescr;
2671:   }
2672:   PetscCheck(Acsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing A CSR struct");
2673:   PetscCheck(Bcsr, PetscObjectComm((PetscObject)C), PETSC_ERR_GPU, "Missing B CSR struct");
2674:   /* precompute flops count */
2675:   if (ptype == MATPRODUCT_AB) {
2676:     for (i = 0, flops = 0; i < A->rmap->n; i++) {
2677:       const PetscInt st = a->i[i];
2678:       const PetscInt en = a->i[i + 1];
2679:       for (j = st; j < en; j++) {
2680:         const PetscInt brow = a->j[j];
2681:         flops += 2. * (b->i[brow + 1] - b->i[brow]);
2682:       }
2683:     }
2684:   } else if (ptype == MATPRODUCT_AtB) {
2685:     for (i = 0, flops = 0; i < A->rmap->n; i++) {
2686:       const PetscInt anzi = a->i[i + 1] - a->i[i];
2687:       const PetscInt bnzi = b->i[i + 1] - b->i[i];
2688:       flops += (2. * anzi) * bnzi;
2689:     }
2690:   } else flops = 0.; /* TODO */

2692:   mmdata->flops = flops;
2693:   PetscCall(PetscLogGpuTimeBegin());
2694: #if PETSC_PKG_HIP_VERSION_GE(5, 0, 0)
2695:   PetscCallHIPSPARSE(hipsparseSetPointerMode(Ccusp->handle, HIPSPARSE_POINTER_MODE_DEVICE));
2696:   PetscCallHIPSPARSE(hipsparseCreateCsr(&Cmat->matDescr, Ccsr->num_rows, Ccsr->num_cols, 0, Ccsr->row_offsets->data().get(), NULL, NULL, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
2697:   PetscCallHIPSPARSE(hipsparseSpGEMM_createDescr(&mmdata->spgemmDesc));
2698:   #if PETSC_PKG_HIP_VERSION_GE(5, 1, 0)
2699:   {
2700:     /* hipsparseSpGEMMreuse has more reasonable APIs than hipsparseSpGEMM, so we prefer to use it.
2701:      We follow the sample code at https://github.com/ROCmSoftwarePlatform/hipSPARSE/blob/develop/clients/include/testing_spgemmreuse_csr.hpp
2702:   */
2703:     void *dBuffer1 = NULL;
2704:     void *dBuffer2 = NULL;
2705:     void *dBuffer3 = NULL;
2706:     /* dBuffer4, dBuffer5 are needed by hipsparseSpGEMMreuse_compute, and therefore are stored in mmdata */
2707:     size_t bufferSize1 = 0;
2708:     size_t bufferSize2 = 0;
2709:     size_t bufferSize3 = 0;
2710:     size_t bufferSize4 = 0;
2711:     size_t bufferSize5 = 0;

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

2719:     PetscCallHIPSPARSE(hipsparseSpGEMMreuse_nnz(Ccusp->handle, opA, opB, Amat->matDescr, BmatSpDescr, Cmat->matDescr, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufferSize2, NULL, &bufferSize3, NULL, &bufferSize4, NULL));
2720:     PetscCallHIP(hipMalloc((void **)&dBuffer2, bufferSize2));
2721:     PetscCallHIP(hipMalloc((void **)&dBuffer3, bufferSize3));
2722:     PetscCallHIP(hipMalloc((void **)&mmdata->dBuffer4, bufferSize4));
2723:     PetscCallHIPSPARSE(hipsparseSpGEMMreuse_nnz(Ccusp->handle, opA, opB, Amat->matDescr, BmatSpDescr, Cmat->matDescr, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufferSize2, dBuffer2, &bufferSize3, dBuffer3, &bufferSize4, mmdata->dBuffer4));
2724:     PetscCallHIP(hipFree(dBuffer1));
2725:     PetscCallHIP(hipFree(dBuffer2));

2727:     /* get matrix C non-zero entries C_nnz1 */
2728:     PetscCallHIPSPARSE(hipsparseSpMatGetSize(Cmat->matDescr, &C_num_rows1, &C_num_cols1, &C_nnz1));
2729:     c->nz = (PetscInt)C_nnz1;
2730:     /* allocate matrix C */
2731:     Ccsr->column_indices = new THRUSTINTARRAY32(c->nz);
2732:     PetscCallHIP(hipPeekAtLastError()); /* catch out of memory errors */
2733:     Ccsr->values = new THRUSTARRAY(c->nz);
2734:     PetscCallHIP(hipPeekAtLastError()); /* catch out of memory errors */
2735:     /* update matC with the new pointers */
2736:     PetscCallHIPSPARSE(hipsparseCsrSetPointers(Cmat->matDescr, Ccsr->row_offsets->data().get(), Ccsr->column_indices->data().get(), Ccsr->values->data().get()));

2738:     PetscCallHIPSPARSE(hipsparseSpGEMMreuse_copy(Ccusp->handle, opA, opB, Amat->matDescr, BmatSpDescr, Cmat->matDescr, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufferSize5, NULL));
2739:     PetscCallHIP(hipMalloc((void **)&mmdata->dBuffer5, bufferSize5));
2740:     PetscCallHIPSPARSE(hipsparseSpGEMMreuse_copy(Ccusp->handle, opA, opB, Amat->matDescr, BmatSpDescr, Cmat->matDescr, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufferSize5, mmdata->dBuffer5));
2741:     PetscCallHIP(hipFree(dBuffer3));
2742:     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));
2743:     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));
2744:   }
2745:   #else
2746:   size_t bufSize2;
2747:   /* ask bufferSize bytes for external memory */
2748:   PetscCallHIPSPARSE(hipsparseSpGEMM_workEstimation(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, hipsparse_scalartype, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufSize2, NULL));
2749:   PetscCallHIP(hipMalloc((void **)&mmdata->mmBuffer2, bufSize2));
2750:   /* inspect the matrices A and B to understand the memory requirement for the next step */
2751:   PetscCallHIPSPARSE(hipsparseSpGEMM_workEstimation(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, hipsparse_scalartype, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &bufSize2, mmdata->mmBuffer2));
2752:   /* ask bufferSize again bytes for external memory */
2753:   PetscCallHIPSPARSE(hipsparseSpGEMM_compute(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, hipsparse_scalartype, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &mmdata->mmBufferSize, NULL));
2754:   /* Similar to CUSPARSE, we need both buffers to perform the operations properly!
2755:      mmdata->mmBuffer2 does not appear anywhere in the compute/copy API
2756:      it only appears for the workEstimation stuff, but it seems it is needed in compute, so probably the address
2757:      is stored in the descriptor! What a messy API... */
2758:   PetscCallHIP(hipMalloc((void **)&mmdata->mmBuffer, mmdata->mmBufferSize));
2759:   /* compute the intermediate product of A * B */
2760:   PetscCallHIPSPARSE(hipsparseSpGEMM_compute(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, hipsparse_scalartype, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc, &mmdata->mmBufferSize, mmdata->mmBuffer));
2761:   /* get matrix C non-zero entries C_nnz1 */
2762:   PetscCallHIPSPARSE(hipsparseSpMatGetSize(Cmat->matDescr, &C_num_rows1, &C_num_cols1, &C_nnz1));
2763:   c->nz = (PetscInt)C_nnz1;
2764:   PetscCall(PetscInfo(C, "Buffer sizes for type %s, result %" PetscInt_FMT " x %" PetscInt_FMT " (k %" PetscInt_FMT ", nzA %" PetscInt_FMT ", nzB %" PetscInt_FMT ", nzC %" PetscInt_FMT ") are: %ldKB %ldKB\n", MatProductTypes[ptype], m, n, k, a->nz, b->nz, c->nz, bufSize2 / 1024,
2765:                       mmdata->mmBufferSize / 1024));
2766:   Ccsr->column_indices = new THRUSTINTARRAY32(c->nz);
2767:   PetscCallHIP(hipPeekAtLastError()); /* catch out of memory errors */
2768:   Ccsr->values = new THRUSTARRAY(c->nz);
2769:   PetscCallHIP(hipPeekAtLastError()); /* catch out of memory errors */
2770:   PetscCallHIPSPARSE(hipsparseCsrSetPointers(Cmat->matDescr, Ccsr->row_offsets->data().get(), Ccsr->column_indices->data().get(), Ccsr->values->data().get()));
2771:   PetscCallHIPSPARSE(hipsparseSpGEMM_copy(Ccusp->handle, opA, opB, Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr, hipsparse_scalartype, HIPSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc));
2772:   #endif
2773: #else
2774:   PetscCallHIPSPARSE(hipsparseSetPointerMode(Ccusp->handle, HIPSPARSE_POINTER_MODE_HOST));
2775:   PetscCallHIPSPARSE(hipsparseXcsrgemmNnz(Ccusp->handle, opA, opB, Acsr->num_rows, Bcsr->num_cols, Acsr->num_cols, Amat->descr, Acsr->num_entries, Acsr->row_offsets->data().get(), Acsr->column_indices->data().get(), Bmat->descr, Bcsr->num_entries,
2776:                                           Bcsr->row_offsets->data().get(), Bcsr->column_indices->data().get(), Cmat->descr, Ccsr->row_offsets->data().get(), &cnz));
2777:   c->nz                = cnz;
2778:   Ccsr->column_indices = new THRUSTINTARRAY32(c->nz);
2779:   PetscCallHIP(hipPeekAtLastError()); /* catch out of memory errors */
2780:   Ccsr->values = new THRUSTARRAY(c->nz);
2781:   PetscCallHIP(hipPeekAtLastError()); /* catch out of memory errors */

2783:   PetscCallHIPSPARSE(hipsparseSetPointerMode(Ccusp->handle, HIPSPARSE_POINTER_MODE_DEVICE));
2784:   /* with the old gemm interface (removed from 11.0 on) we cannot compute the symbolic factorization only.
2785:       I have tried using the gemm2 interface (alpha * A * B + beta * D), which allows to do symbolic by passing NULL for values, but it seems quite buggy when
2786:       D is NULL, despite the fact that CUSPARSE documentation claims it is supported! */
2787:   PetscCallHIPSPARSE(hipsparse_csr_spgemm(Ccusp->handle, opA, opB, Acsr->num_rows, Bcsr->num_cols, Acsr->num_cols, Amat->descr, Acsr->num_entries, Acsr->values->data().get(), Acsr->row_offsets->data().get(), Acsr->column_indices->data().get(), Bmat->descr,
2788:                                           Bcsr->num_entries, Bcsr->values->data().get(), Bcsr->row_offsets->data().get(), Bcsr->column_indices->data().get(), Cmat->descr, Ccsr->values->data().get(), Ccsr->row_offsets->data().get(),
2789:                                           Ccsr->column_indices->data().get()));
2790: #endif
2791:   PetscCall(PetscLogGpuFlops(mmdata->flops));
2792:   PetscCall(PetscLogGpuTimeEnd());
2793: finalizesym:
2794:   c->singlemalloc = PETSC_FALSE;
2795:   c->free_a       = PETSC_TRUE;
2796:   c->free_ij      = PETSC_TRUE;
2797:   PetscCall(PetscMalloc1(m + 1, &c->i));
2798:   PetscCall(PetscMalloc1(c->nz, &c->j));
2799:   if (PetscDefined(USE_64BIT_INDICES)) { /* 32 to 64-bit conversion on the GPU and then copy to host (lazy) */
2800:     PetscInt      *d_i = c->i;
2801:     THRUSTINTARRAY ii(Ccsr->row_offsets->size());
2802:     THRUSTINTARRAY jj(Ccsr->column_indices->size());
2803:     ii = *Ccsr->row_offsets;
2804:     jj = *Ccsr->column_indices;
2805:     if (ciscompressed) d_i = c->compressedrow.i;
2806:     PetscCallHIP(hipMemcpy(d_i, ii.data().get(), Ccsr->row_offsets->size() * sizeof(PetscInt), hipMemcpyDeviceToHost));
2807:     PetscCallHIP(hipMemcpy(c->j, jj.data().get(), Ccsr->column_indices->size() * sizeof(PetscInt), hipMemcpyDeviceToHost));
2808:   } else {
2809:     PetscInt *d_i = c->i;
2810:     if (ciscompressed) d_i = c->compressedrow.i;
2811:     PetscCallHIP(hipMemcpy(d_i, Ccsr->row_offsets->data().get(), Ccsr->row_offsets->size() * sizeof(PetscInt), hipMemcpyDeviceToHost));
2812:     PetscCallHIP(hipMemcpy(c->j, Ccsr->column_indices->data().get(), Ccsr->column_indices->size() * sizeof(PetscInt), hipMemcpyDeviceToHost));
2813:   }
2814:   if (ciscompressed) { /* need to expand host row offsets */
2815:     PetscInt r = 0;
2816:     c->i[0]    = 0;
2817:     for (k = 0; k < c->compressedrow.nrows; k++) {
2818:       const PetscInt next = c->compressedrow.rindex[k];
2819:       const PetscInt old  = c->compressedrow.i[k];
2820:       for (; r < next; r++) c->i[r + 1] = old;
2821:     }
2822:     for (; r < m; r++) c->i[r + 1] = c->compressedrow.i[c->compressedrow.nrows];
2823:   }
2824:   PetscCall(PetscLogGpuToCpu((Ccsr->column_indices->size() + Ccsr->row_offsets->size()) * sizeof(PetscInt)));
2825:   PetscCall(PetscMalloc1(m, &c->ilen));
2826:   PetscCall(PetscMalloc1(m, &c->imax));
2827:   c->maxnz         = c->nz;
2828:   c->nonzerorowcnt = 0;
2829:   c->rmax          = 0;
2830:   for (k = 0; k < m; k++) {
2831:     const PetscInt nn = c->i[k + 1] - c->i[k];
2832:     c->ilen[k] = c->imax[k] = nn;
2833:     c->nonzerorowcnt += (PetscInt) !!nn;
2834:     c->rmax = PetscMax(c->rmax, nn);
2835:   }
2836:   PetscCall(MatMarkDiagonal_SeqAIJ(C));
2837:   PetscCall(PetscMalloc1(c->nz, &c->a));
2838:   Ccsr->num_entries = c->nz;

2840:   C->nonzerostate++;
2841:   PetscCall(PetscLayoutSetUp(C->rmap));
2842:   PetscCall(PetscLayoutSetUp(C->cmap));
2843:   Ccusp->nonzerostate = C->nonzerostate;
2844:   C->offloadmask      = PETSC_OFFLOAD_UNALLOCATED;
2845:   C->preallocated     = PETSC_TRUE;
2846:   C->assembled        = PETSC_FALSE;
2847:   C->was_assembled    = PETSC_FALSE;
2848:   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 */
2849:     mmdata->reusesym = PETSC_TRUE;
2850:     C->offloadmask   = PETSC_OFFLOAD_GPU;
2851:   }
2852:   C->ops->productnumeric = MatProductNumeric_SeqAIJHIPSPARSE_SeqAIJHIPSPARSE;
2853:   PetscFunctionReturn(PETSC_SUCCESS);
2854: }

2856: /* handles sparse or dense B */
2857: static PetscErrorCode MatProductSetFromOptions_SeqAIJHIPSPARSE(Mat mat)
2858: {
2859:   Mat_Product *product = mat->product;
2860:   PetscBool    isdense = PETSC_FALSE, Biscusp = PETSC_FALSE, Ciscusp = PETSC_TRUE;

2862:   PetscFunctionBegin;
2863:   MatCheckProduct(mat, 1);
2864:   PetscCall(PetscObjectBaseTypeCompare((PetscObject)product->B, MATSEQDENSE, &isdense));
2865:   if (!product->A->boundtocpu && !product->B->boundtocpu) PetscCall(PetscObjectTypeCompare((PetscObject)product->B, MATSEQAIJHIPSPARSE, &Biscusp));
2866:   if (product->type == MATPRODUCT_ABC) {
2867:     Ciscusp = PETSC_FALSE;
2868:     if (!product->C->boundtocpu) PetscCall(PetscObjectTypeCompare((PetscObject)product->C, MATSEQAIJHIPSPARSE, &Ciscusp));
2869:   }
2870:   if (Biscusp && Ciscusp) { /* we can always select the CPU backend */
2871:     PetscBool usecpu = PETSC_FALSE;
2872:     switch (product->type) {
2873:     case MATPRODUCT_AB:
2874:       if (product->api_user) {
2875:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatMatMult", "Mat");
2876:         PetscCall(PetscOptionsBool("-matmatmult_backend_cpu", "Use CPU code", "MatMatMult", usecpu, &usecpu, NULL));
2877:         PetscOptionsEnd();
2878:       } else {
2879:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_AB", "Mat");
2880:         PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatMatMult", usecpu, &usecpu, NULL));
2881:         PetscOptionsEnd();
2882:       }
2883:       break;
2884:     case MATPRODUCT_AtB:
2885:       if (product->api_user) {
2886:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatTransposeMatMult", "Mat");
2887:         PetscCall(PetscOptionsBool("-mattransposematmult_backend_cpu", "Use CPU code", "MatTransposeMatMult", usecpu, &usecpu, NULL));
2888:         PetscOptionsEnd();
2889:       } else {
2890:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_AtB", "Mat");
2891:         PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatTransposeMatMult", usecpu, &usecpu, NULL));
2892:         PetscOptionsEnd();
2893:       }
2894:       break;
2895:     case MATPRODUCT_PtAP:
2896:       if (product->api_user) {
2897:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatPtAP", "Mat");
2898:         PetscCall(PetscOptionsBool("-matptap_backend_cpu", "Use CPU code", "MatPtAP", usecpu, &usecpu, NULL));
2899:         PetscOptionsEnd();
2900:       } else {
2901:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_PtAP", "Mat");
2902:         PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatPtAP", usecpu, &usecpu, NULL));
2903:         PetscOptionsEnd();
2904:       }
2905:       break;
2906:     case MATPRODUCT_RARt:
2907:       if (product->api_user) {
2908:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatRARt", "Mat");
2909:         PetscCall(PetscOptionsBool("-matrart_backend_cpu", "Use CPU code", "MatRARt", usecpu, &usecpu, NULL));
2910:         PetscOptionsEnd();
2911:       } else {
2912:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_RARt", "Mat");
2913:         PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatRARt", usecpu, &usecpu, NULL));
2914:         PetscOptionsEnd();
2915:       }
2916:       break;
2917:     case MATPRODUCT_ABC:
2918:       if (product->api_user) {
2919:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatMatMatMult", "Mat");
2920:         PetscCall(PetscOptionsBool("-matmatmatmult_backend_cpu", "Use CPU code", "MatMatMatMult", usecpu, &usecpu, NULL));
2921:         PetscOptionsEnd();
2922:       } else {
2923:         PetscOptionsBegin(PetscObjectComm((PetscObject)mat), ((PetscObject)mat)->prefix, "MatProduct_ABC", "Mat");
2924:         PetscCall(PetscOptionsBool("-mat_product_algorithm_backend_cpu", "Use CPU code", "MatMatMatMult", usecpu, &usecpu, NULL));
2925:         PetscOptionsEnd();
2926:       }
2927:       break;
2928:     default:
2929:       break;
2930:     }
2931:     if (usecpu) Biscusp = Ciscusp = PETSC_FALSE;
2932:   }
2933:   /* dispatch */
2934:   if (isdense) {
2935:     switch (product->type) {
2936:     case MATPRODUCT_AB:
2937:     case MATPRODUCT_AtB:
2938:     case MATPRODUCT_ABt:
2939:     case MATPRODUCT_PtAP:
2940:     case MATPRODUCT_RARt:
2941:       if (product->A->boundtocpu) PetscCall(MatProductSetFromOptions_SeqAIJ_SeqDense(mat));
2942:       else mat->ops->productsymbolic = MatProductSymbolic_SeqAIJHIPSPARSE_SeqDENSEHIP;
2943:       break;
2944:     case MATPRODUCT_ABC:
2945:       mat->ops->productsymbolic = MatProductSymbolic_ABC_Basic;
2946:       break;
2947:     default:
2948:       break;
2949:     }
2950:   } else if (Biscusp && Ciscusp) {
2951:     switch (product->type) {
2952:     case MATPRODUCT_AB:
2953:     case MATPRODUCT_AtB:
2954:     case MATPRODUCT_ABt:
2955:       mat->ops->productsymbolic = MatProductSymbolic_SeqAIJHIPSPARSE_SeqAIJHIPSPARSE;
2956:       break;
2957:     case MATPRODUCT_PtAP:
2958:     case MATPRODUCT_RARt:
2959:     case MATPRODUCT_ABC:
2960:       mat->ops->productsymbolic = MatProductSymbolic_ABC_Basic;
2961:       break;
2962:     default:
2963:       break;
2964:     }
2965:   } else PetscCall(MatProductSetFromOptions_SeqAIJ(mat)); /* fallback for AIJ */
2966:   PetscFunctionReturn(PETSC_SUCCESS);
2967: }

2969: static PetscErrorCode MatMult_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy)
2970: {
2971:   PetscFunctionBegin;
2972:   PetscCall(MatMultAddKernel_SeqAIJHIPSPARSE(A, xx, NULL, yy, PETSC_FALSE, PETSC_FALSE));
2973:   PetscFunctionReturn(PETSC_SUCCESS);
2974: }

2976: static PetscErrorCode MatMultAdd_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy, Vec zz)
2977: {
2978:   PetscFunctionBegin;
2979:   PetscCall(MatMultAddKernel_SeqAIJHIPSPARSE(A, xx, yy, zz, PETSC_FALSE, PETSC_FALSE));
2980:   PetscFunctionReturn(PETSC_SUCCESS);
2981: }

2983: static PetscErrorCode MatMultHermitianTranspose_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy)
2984: {
2985:   PetscFunctionBegin;
2986:   PetscCall(MatMultAddKernel_SeqAIJHIPSPARSE(A, xx, NULL, yy, PETSC_TRUE, PETSC_TRUE));
2987:   PetscFunctionReturn(PETSC_SUCCESS);
2988: }

2990: static PetscErrorCode MatMultHermitianTransposeAdd_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy, Vec zz)
2991: {
2992:   PetscFunctionBegin;
2993:   PetscCall(MatMultAddKernel_SeqAIJHIPSPARSE(A, xx, yy, zz, PETSC_TRUE, PETSC_TRUE));
2994:   PetscFunctionReturn(PETSC_SUCCESS);
2995: }

2997: static PetscErrorCode MatMultTranspose_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy)
2998: {
2999:   PetscFunctionBegin;
3000:   PetscCall(MatMultAddKernel_SeqAIJHIPSPARSE(A, xx, NULL, yy, PETSC_TRUE, PETSC_FALSE));
3001:   PetscFunctionReturn(PETSC_SUCCESS);
3002: }

3004: __global__ static void ScatterAdd(PetscInt n, PetscInt *idx, const PetscScalar *x, PetscScalar *y)
3005: {
3006:   int i = blockIdx.x * blockDim.x + threadIdx.x;
3007:   if (i < n) y[idx[i]] += x[i];
3008: }

3010: /* z = op(A) x + y. If trans & !herm, op = ^T; if trans & herm, op = ^H; if !trans, op = no-op */
3011: static PetscErrorCode MatMultAddKernel_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy, Vec zz, PetscBool trans, PetscBool herm)
3012: {
3013:   Mat_SeqAIJ                    *a               = (Mat_SeqAIJ *)A->data;
3014:   Mat_SeqAIJHIPSPARSE           *hipsparsestruct = (Mat_SeqAIJHIPSPARSE *)A->spptr;
3015:   Mat_SeqAIJHIPSPARSEMultStruct *matstruct;
3016:   PetscScalar                   *xarray, *zarray, *dptr, *beta, *xptr;
3017:   hipsparseOperation_t           opA = HIPSPARSE_OPERATION_NON_TRANSPOSE;
3018:   PetscBool                      compressed;
3019:   PetscInt                       nx, ny;

3021:   PetscFunctionBegin;
3022:   PetscCheck(!herm || trans, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "Hermitian and not transpose not supported");
3023:   if (!a->nz) {
3024:     if (yy) PetscCall(VecSeq_HIP::Copy(yy, zz));
3025:     else PetscCall(VecSeq_HIP::Set(zz, 0));
3026:     PetscFunctionReturn(PETSC_SUCCESS);
3027:   }
3028:   /* The line below is necessary due to the operations that modify the matrix on the CPU (axpy, scale, etc) */
3029:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
3030:   if (!trans) {
3031:     matstruct = (Mat_SeqAIJHIPSPARSEMultStruct *)hipsparsestruct->mat;
3032:     PetscCheck(matstruct, PetscObjectComm((PetscObject)A), PETSC_ERR_GPU, "SeqAIJHIPSPARSE does not have a 'mat' (need to fix)");
3033:   } else {
3034:     if (herm || !A->form_explicit_transpose) {
3035:       opA       = herm ? HIPSPARSE_OPERATION_CONJUGATE_TRANSPOSE : HIPSPARSE_OPERATION_TRANSPOSE;
3036:       matstruct = (Mat_SeqAIJHIPSPARSEMultStruct *)hipsparsestruct->mat;
3037:     } else {
3038:       if (!hipsparsestruct->matTranspose) PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(A));
3039:       matstruct = (Mat_SeqAIJHIPSPARSEMultStruct *)hipsparsestruct->matTranspose;
3040:     }
3041:   }
3042:   /* Does the matrix use compressed rows (i.e., drop zero rows)? */
3043:   compressed = matstruct->cprowIndices ? PETSC_TRUE : PETSC_FALSE;
3044:   try {
3045:     PetscCall(VecHIPGetArrayRead(xx, (const PetscScalar **)&xarray));
3046:     if (yy == zz) PetscCall(VecHIPGetArray(zz, &zarray)); /* read & write zz, so need to get up-to-date zarray on GPU */
3047:     else PetscCall(VecHIPGetArrayWrite(zz, &zarray));     /* write zz, so no need to init zarray on GPU */

3049:     PetscCall(PetscLogGpuTimeBegin());
3050:     if (opA == HIPSPARSE_OPERATION_NON_TRANSPOSE) {
3051:       /* z = A x + beta y.
3052:          If A is compressed (with less rows), then Ax is shorter than the full z, so we need a work vector to store Ax.
3053:          When A is non-compressed, and z = y, we can set beta=1 to compute y = Ax + y in one call.
3054:       */
3055:       xptr = xarray;
3056:       dptr = compressed ? hipsparsestruct->workVector->data().get() : zarray;
3057:       beta = (yy == zz && !compressed) ? matstruct->beta_one : matstruct->beta_zero;
3058:       /* Get length of x, y for y=Ax. ny might be shorter than the work vector's allocated length, since the work vector is
3059:           allocated to accommodate different uses. So we get the length info directly from mat.
3060:        */
3061:       if (hipsparsestruct->format == MAT_HIPSPARSE_CSR) {
3062:         CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
3063:         nx             = mat->num_cols;
3064:         ny             = mat->num_rows;
3065:       }
3066:     } else {
3067:       /* z = A^T x + beta y
3068:          If A is compressed, then we need a work vector as the shorter version of x to compute A^T x.
3069:          Note A^Tx is of full length, so we set beta to 1.0 if y exists.
3070:        */
3071:       xptr = compressed ? hipsparsestruct->workVector->data().get() : xarray;
3072:       dptr = zarray;
3073:       beta = yy ? matstruct->beta_one : matstruct->beta_zero;
3074:       if (compressed) { /* Scatter x to work vector */
3075:         thrust::device_ptr<PetscScalar> xarr = thrust::device_pointer_cast(xarray);
3076:         thrust::for_each(
3077: #if PetscDefined(HAVE_THRUST_ASYNC)
3078:           thrust::hip::par.on(PetscDefaultHipStream),
3079: #endif
3080:           thrust::make_zip_iterator(thrust::make_tuple(hipsparsestruct->workVector->begin(), thrust::make_permutation_iterator(xarr, matstruct->cprowIndices->begin()))),
3081:           thrust::make_zip_iterator(thrust::make_tuple(hipsparsestruct->workVector->begin(), thrust::make_permutation_iterator(xarr, matstruct->cprowIndices->begin()))) + matstruct->cprowIndices->size(), VecHIPEqualsReverse());
3082:       }
3083:       if (hipsparsestruct->format == MAT_HIPSPARSE_CSR) {
3084:         CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
3085:         nx             = mat->num_rows;
3086:         ny             = mat->num_cols;
3087:       }
3088:     }
3089:     /* csr_spmv does y = alpha op(A) x + beta y */
3090:     if (hipsparsestruct->format == MAT_HIPSPARSE_CSR) {
3091: #if PETSC_PKG_HIP_VERSION_GE(5, 1, 0)
3092:       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");
3093:       if (!matstruct->hipSpMV[opA].initialized) { /* built on demand */
3094:         PetscCallHIPSPARSE(hipsparseCreateDnVec(&matstruct->hipSpMV[opA].vecXDescr, nx, xptr, hipsparse_scalartype));
3095:         PetscCallHIPSPARSE(hipsparseCreateDnVec(&matstruct->hipSpMV[opA].vecYDescr, ny, dptr, hipsparse_scalartype));
3096:         PetscCallHIPSPARSE(hipsparseSpMV_bufferSize(hipsparsestruct->handle, opA, matstruct->alpha_one, matstruct->matDescr, matstruct->hipSpMV[opA].vecXDescr, beta, matstruct->hipSpMV[opA].vecYDescr, hipsparse_scalartype, hipsparsestruct->spmvAlg,
3097:                                                     &matstruct->hipSpMV[opA].spmvBufferSize));
3098:         PetscCallHIP(hipMalloc(&matstruct->hipSpMV[opA].spmvBuffer, matstruct->hipSpMV[opA].spmvBufferSize));
3099:         matstruct->hipSpMV[opA].initialized = PETSC_TRUE;
3100:       } else {
3101:         /* x, y's value pointers might change between calls, but their shape is kept, so we just update pointers */
3102:         PetscCallHIPSPARSE(hipsparseDnVecSetValues(matstruct->hipSpMV[opA].vecXDescr, xptr));
3103:         PetscCallHIPSPARSE(hipsparseDnVecSetValues(matstruct->hipSpMV[opA].vecYDescr, dptr));
3104:       }
3105:       PetscCallHIPSPARSE(hipsparseSpMV(hipsparsestruct->handle, opA, matstruct->alpha_one, matstruct->matDescr, /* built in MatSeqAIJHIPSPARSECopyToGPU() or MatSeqAIJHIPSPARSEFormExplicitTranspose() */
3106:                                        matstruct->hipSpMV[opA].vecXDescr, beta, matstruct->hipSpMV[opA].vecYDescr, hipsparse_scalartype, hipsparsestruct->spmvAlg, matstruct->hipSpMV[opA].spmvBuffer));
3107: #else
3108:       CsrMatrix *mat = (CsrMatrix *)matstruct->mat;
3109:       PetscCallHIPSPARSE(hipsparse_csr_spmv(hipsparsestruct->handle, opA, mat->num_rows, mat->num_cols, mat->num_entries, matstruct->alpha_one, matstruct->descr, mat->values->data().get(), mat->row_offsets->data().get(), mat->column_indices->data().get(), xptr, beta, dptr));
3110: #endif
3111:     } else {
3112:       if (hipsparsestruct->nrows) {
3113:         hipsparseHybMat_t hybMat = (hipsparseHybMat_t)matstruct->mat;
3114:         PetscCallHIPSPARSE(hipsparse_hyb_spmv(hipsparsestruct->handle, opA, matstruct->alpha_one, matstruct->descr, hybMat, xptr, beta, dptr));
3115:       }
3116:     }
3117:     PetscCall(PetscLogGpuTimeEnd());

3119:     if (opA == HIPSPARSE_OPERATION_NON_TRANSPOSE) {
3120:       if (yy) {                                     /* MatMultAdd: zz = A*xx + yy */
3121:         if (compressed) {                           /* A is compressed. We first copy yy to zz, then ScatterAdd the work vector to zz */
3122:           PetscCall(VecSeq_HIP::Copy(yy, zz));      /* zz = yy */
3123:         } else if (zz != yy) {                      /* A is not compressed. zz already contains A*xx, and we just need to add yy */
3124:           PetscCall(VecSeq_HIP::AXPY(zz, 1.0, yy)); /* zz += yy */
3125:         }
3126:       } else if (compressed) { /* MatMult: zz = A*xx. A is compressed, so we zero zz first, then ScatterAdd the work vector to zz */
3127:         PetscCall(VecSeq_HIP::Set(zz, 0));
3128:       }

3130:       /* ScatterAdd the result from work vector into the full vector when A is compressed */
3131:       if (compressed) {
3132:         PetscCall(PetscLogGpuTimeBegin());
3133:         /* I wanted to make this for_each asynchronous but failed. thrust::async::for_each() returns an event (internally registered)
3134:            and in the destructor of the scope, it will call hipStreamSynchronize() on this stream. One has to store all events to
3135:            prevent that. So I just add a ScatterAdd kernel.
3136:          */
3137: #if 0
3138:         thrust::device_ptr<PetscScalar> zptr = thrust::device_pointer_cast(zarray);
3139:         thrust::async::for_each(thrust::hip::par.on(hipsparsestruct->stream),
3140:                          thrust::make_zip_iterator(thrust::make_tuple(hipsparsestruct->workVector->begin(), thrust::make_permutation_iterator(zptr, matstruct->cprowIndices->begin()))),
3141:                          thrust::make_zip_iterator(thrust::make_tuple(hipsparsestruct->workVector->begin(), thrust::make_permutation_iterator(zptr, matstruct->cprowIndices->begin()))) + matstruct->cprowIndices->size(),
3142:                          VecHIPPlusEquals());
3143: #else
3144:         PetscInt n = matstruct->cprowIndices->size();
3145:         hipLaunchKernelGGL(ScatterAdd, dim3((n + 255) / 256), dim3(256), 0, PetscDefaultHipStream, n, matstruct->cprowIndices->data().get(), hipsparsestruct->workVector->data().get(), zarray);
3146: #endif
3147:         PetscCall(PetscLogGpuTimeEnd());
3148:       }
3149:     } else {
3150:       if (yy && yy != zz) PetscCall(VecSeq_HIP::AXPY(zz, 1.0, yy)); /* zz += yy */
3151:     }
3152:     PetscCall(VecHIPRestoreArrayRead(xx, (const PetscScalar **)&xarray));
3153:     if (yy == zz) PetscCall(VecHIPRestoreArray(zz, &zarray));
3154:     else PetscCall(VecHIPRestoreArrayWrite(zz, &zarray));
3155:   } catch (char *ex) {
3156:     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_LIB, "HIPSPARSE error: %s", ex);
3157:   }
3158:   if (yy) PetscCall(PetscLogGpuFlops(2.0 * a->nz));
3159:   else PetscCall(PetscLogGpuFlops(2.0 * a->nz - a->nonzerorowcnt));
3160:   PetscFunctionReturn(PETSC_SUCCESS);
3161: }

3163: static PetscErrorCode MatMultTransposeAdd_SeqAIJHIPSPARSE(Mat A, Vec xx, Vec yy, Vec zz)
3164: {
3165:   PetscFunctionBegin;
3166:   PetscCall(MatMultAddKernel_SeqAIJHIPSPARSE(A, xx, yy, zz, PETSC_TRUE, PETSC_FALSE));
3167:   PetscFunctionReturn(PETSC_SUCCESS);
3168: }

3170: static PetscErrorCode MatAssemblyEnd_SeqAIJHIPSPARSE(Mat A, MatAssemblyType mode)
3171: {
3172:   PetscObjectState     onnz = A->nonzerostate;
3173:   Mat_SeqAIJHIPSPARSE *cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;

3175:   PetscFunctionBegin;
3176:   PetscCall(MatAssemblyEnd_SeqAIJ(A, mode));
3177:   if (onnz != A->nonzerostate && cusp->deviceMat) {
3178:     PetscCall(PetscInfo(A, "Destroy device mat since nonzerostate changed\n"));
3179:     PetscCallHIP(hipFree(cusp->deviceMat));
3180:     cusp->deviceMat = NULL;
3181:   }
3182:   PetscFunctionReturn(PETSC_SUCCESS);
3183: }

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

3189:    Collective

3191:    Input Parameters:
3192: +  comm - MPI communicator, set to `PETSC_COMM_SELF`
3193: .  m - number of rows
3194: .  n - number of columns
3195: .  nz - number of nonzeros per row (same for all rows), ignored if `nnz` is set
3196: -  nnz - array containing the number of nonzeros in the various rows (possibly different for each row) or `NULL`

3198:    Output Parameter:
3199: .  A - the matrix

3201:    Level: intermediate

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

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

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

3216: .seealso: [](ch_matrices), `Mat`, `MatCreate()`, `MatCreateAIJ()`, `MatSetValues()`, `MatSeqAIJSetColumnIndices()`, `MatCreateSeqAIJWithArrays()`, `MatCreateAIJ()`, `MATSEQAIJHIPSPARSE`, `MATAIJHIPSPARSE`
3217: @*/
3218: PetscErrorCode MatCreateSeqAIJHIPSPARSE(MPI_Comm comm, PetscInt m, PetscInt n, PetscInt nz, const PetscInt nnz[], Mat *A)
3219: {
3220:   PetscFunctionBegin;
3221:   PetscCall(MatCreate(comm, A));
3222:   PetscCall(MatSetSizes(*A, m, n, m, n));
3223:   PetscCall(MatSetType(*A, MATSEQAIJHIPSPARSE));
3224:   PetscCall(MatSeqAIJSetPreallocation_SeqAIJ(*A, nz, (PetscInt *)nnz));
3225:   PetscFunctionReturn(PETSC_SUCCESS);
3226: }

3228: static PetscErrorCode MatDestroy_SeqAIJHIPSPARSE(Mat A)
3229: {
3230:   PetscFunctionBegin;
3231:   if (A->factortype == MAT_FACTOR_NONE) PetscCall(MatSeqAIJHIPSPARSE_Destroy((Mat_SeqAIJHIPSPARSE **)&A->spptr));
3232:   else PetscCall(MatSeqAIJHIPSPARSETriFactors_Destroy((Mat_SeqAIJHIPSPARSETriFactors **)&A->spptr));
3233:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSeqAIJCopySubArray_C", NULL));
3234:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatHIPSPARSESetFormat_C", NULL));
3235:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatHIPSPARSESetUseCPUSolve_C", NULL));
3236:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqdensehip_C", NULL));
3237:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqdense_C", NULL));
3238:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqaijhipsparse_C", NULL));
3239:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatFactorGetSolverType_C", NULL));
3240:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetPreallocationCOO_C", NULL));
3241:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetValuesCOO_C", NULL));
3242:   PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatConvert_seqaijhipsparse_hypre_C", NULL));
3243:   PetscCall(MatDestroy_SeqAIJ(A));
3244:   PetscFunctionReturn(PETSC_SUCCESS);
3245: }

3247: static PetscErrorCode MatDuplicate_SeqAIJHIPSPARSE(Mat A, MatDuplicateOption cpvalues, Mat *B)
3248: {
3249:   PetscFunctionBegin;
3250:   PetscCall(MatDuplicate_SeqAIJ(A, cpvalues, B));
3251:   PetscCall(MatConvert_SeqAIJ_SeqAIJHIPSPARSE(*B, MATSEQAIJHIPSPARSE, MAT_INPLACE_MATRIX, B));
3252:   PetscFunctionReturn(PETSC_SUCCESS);
3253: }

3255: static PetscErrorCode MatAXPY_SeqAIJHIPSPARSE(Mat Y, PetscScalar a, Mat X, MatStructure str)
3256: {
3257:   Mat_SeqAIJ          *x = (Mat_SeqAIJ *)X->data, *y = (Mat_SeqAIJ *)Y->data;
3258:   Mat_SeqAIJHIPSPARSE *cy;
3259:   Mat_SeqAIJHIPSPARSE *cx;
3260:   PetscScalar         *ay;
3261:   const PetscScalar   *ax;
3262:   CsrMatrix           *csry, *csrx;

3264:   PetscFunctionBegin;
3265:   cy = (Mat_SeqAIJHIPSPARSE *)Y->spptr;
3266:   cx = (Mat_SeqAIJHIPSPARSE *)X->spptr;
3267:   if (X->ops->axpy != Y->ops->axpy) {
3268:     PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(Y, PETSC_FALSE));
3269:     PetscCall(MatAXPY_SeqAIJ(Y, a, X, str));
3270:     PetscFunctionReturn(PETSC_SUCCESS);
3271:   }
3272:   /* if we are here, it means both matrices are bound to GPU */
3273:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(Y));
3274:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(X));
3275:   PetscCheck(cy->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)Y), PETSC_ERR_GPU, "only MAT_HIPSPARSE_CSR supported");
3276:   PetscCheck(cx->format == MAT_HIPSPARSE_CSR, PetscObjectComm((PetscObject)X), PETSC_ERR_GPU, "only MAT_HIPSPARSE_CSR supported");
3277:   csry = (CsrMatrix *)cy->mat->mat;
3278:   csrx = (CsrMatrix *)cx->mat->mat;
3279:   /* see if we can turn this into a hipblas axpy */
3280:   if (str != SAME_NONZERO_PATTERN && x->nz == y->nz && !x->compressedrow.use && !y->compressedrow.use) {
3281:     bool eq = thrust::equal(thrust::device, csry->row_offsets->begin(), csry->row_offsets->end(), csrx->row_offsets->begin());
3282:     if (eq) eq = thrust::equal(thrust::device, csry->column_indices->begin(), csry->column_indices->end(), csrx->column_indices->begin());
3283:     if (eq) str = SAME_NONZERO_PATTERN;
3284:   }
3285:   /* spgeam is buggy with one column */
3286:   if (Y->cmap->n == 1 && str != SAME_NONZERO_PATTERN) str = DIFFERENT_NONZERO_PATTERN;
3287:   if (str == SUBSET_NONZERO_PATTERN) {
3288:     PetscScalar b = 1.0;
3289: #if PETSC_PKG_HIP_VERSION_GE(4, 5, 0)
3290:     size_t bufferSize;
3291:     void  *buffer;
3292: #endif

3294:     PetscCall(MatSeqAIJHIPSPARSEGetArrayRead(X, &ax));
3295:     PetscCall(MatSeqAIJHIPSPARSEGetArray(Y, &ay));
3296:     PetscCallHIPSPARSE(hipsparseSetPointerMode(cy->handle, HIPSPARSE_POINTER_MODE_HOST));
3297: #if PETSC_PKG_HIP_VERSION_GE(4, 5, 0)
3298:     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(),
3299:                                                        csry->column_indices->data().get(), cy->mat->descr, ay, csry->row_offsets->data().get(), csry->column_indices->data().get(), &bufferSize));
3300:     PetscCallHIP(hipMalloc(&buffer, bufferSize));
3301:     PetscCall(PetscLogGpuTimeBegin());
3302:     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(),
3303:                                             csry->column_indices->data().get(), cy->mat->descr, ay, csry->row_offsets->data().get(), csry->column_indices->data().get(), buffer));
3304:     PetscCall(PetscLogGpuFlops(x->nz + y->nz));
3305:     PetscCall(PetscLogGpuTimeEnd());
3306:     PetscCallHIP(hipFree(buffer));
3307: #else
3308:     PetscCall(PetscLogGpuTimeBegin());
3309:     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(),
3310:                                             csry->column_indices->data().get(), cy->mat->descr, ay, csry->row_offsets->data().get(), csry->column_indices->data().get()));
3311:     PetscCall(PetscLogGpuFlops(x->nz + y->nz));
3312:     PetscCall(PetscLogGpuTimeEnd());
3313: #endif
3314:     PetscCallHIPSPARSE(hipsparseSetPointerMode(cy->handle, HIPSPARSE_POINTER_MODE_DEVICE));
3315:     PetscCall(MatSeqAIJHIPSPARSERestoreArrayRead(X, &ax));
3316:     PetscCall(MatSeqAIJHIPSPARSERestoreArray(Y, &ay));
3317:     PetscCall(MatSeqAIJInvalidateDiagonal(Y));
3318:   } else if (str == SAME_NONZERO_PATTERN) {
3319:     hipblasHandle_t hipblasv2handle;
3320:     PetscBLASInt    one = 1, bnz = 1;

3322:     PetscCall(MatSeqAIJHIPSPARSEGetArrayRead(X, &ax));
3323:     PetscCall(MatSeqAIJHIPSPARSEGetArray(Y, &ay));
3324:     PetscCall(PetscHIPBLASGetHandle(&hipblasv2handle));
3325:     PetscCall(PetscBLASIntCast(x->nz, &bnz));
3326:     PetscCall(PetscLogGpuTimeBegin());
3327:     PetscCallHIPBLAS(hipblasXaxpy(hipblasv2handle, bnz, &a, ax, one, ay, one));
3328:     PetscCall(PetscLogGpuFlops(2.0 * bnz));
3329:     PetscCall(PetscLogGpuTimeEnd());
3330:     PetscCall(MatSeqAIJHIPSPARSERestoreArrayRead(X, &ax));
3331:     PetscCall(MatSeqAIJHIPSPARSERestoreArray(Y, &ay));
3332:     PetscCall(MatSeqAIJInvalidateDiagonal(Y));
3333:   } else {
3334:     PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(Y, PETSC_FALSE));
3335:     PetscCall(MatAXPY_SeqAIJ(Y, a, X, str));
3336:   }
3337:   PetscFunctionReturn(PETSC_SUCCESS);
3338: }

3340: static PetscErrorCode MatScale_SeqAIJHIPSPARSE(Mat Y, PetscScalar a)
3341: {
3342:   Mat_SeqAIJ     *y = (Mat_SeqAIJ *)Y->data;
3343:   PetscScalar    *ay;
3344:   hipblasHandle_t hipblasv2handle;
3345:   PetscBLASInt    one = 1, bnz = 1;

3347:   PetscFunctionBegin;
3348:   PetscCall(MatSeqAIJHIPSPARSEGetArray(Y, &ay));
3349:   PetscCall(PetscHIPBLASGetHandle(&hipblasv2handle));
3350:   PetscCall(PetscBLASIntCast(y->nz, &bnz));
3351:   PetscCall(PetscLogGpuTimeBegin());
3352:   PetscCallHIPBLAS(hipblasXscal(hipblasv2handle, bnz, &a, ay, one));
3353:   PetscCall(PetscLogGpuFlops(bnz));
3354:   PetscCall(PetscLogGpuTimeEnd());
3355:   PetscCall(MatSeqAIJHIPSPARSERestoreArray(Y, &ay));
3356:   PetscCall(MatSeqAIJInvalidateDiagonal(Y));
3357:   PetscFunctionReturn(PETSC_SUCCESS);
3358: }

3360: static PetscErrorCode MatZeroEntries_SeqAIJHIPSPARSE(Mat A)
3361: {
3362:   PetscBool   both = PETSC_FALSE;
3363:   Mat_SeqAIJ *a    = (Mat_SeqAIJ *)A->data;

3365:   PetscFunctionBegin;
3366:   if (A->factortype == MAT_FACTOR_NONE) {
3367:     Mat_SeqAIJHIPSPARSE *spptr = (Mat_SeqAIJHIPSPARSE *)A->spptr;
3368:     if (spptr->mat) {
3369:       CsrMatrix *matrix = (CsrMatrix *)spptr->mat->mat;
3370:       if (matrix->values) {
3371:         both = PETSC_TRUE;
3372:         thrust::fill(thrust::device, matrix->values->begin(), matrix->values->end(), 0.);
3373:       }
3374:     }
3375:     if (spptr->matTranspose) {
3376:       CsrMatrix *matrix = (CsrMatrix *)spptr->matTranspose->mat;
3377:       if (matrix->values) { thrust::fill(thrust::device, matrix->values->begin(), matrix->values->end(), 0.); }
3378:     }
3379:   }
3380:   //PetscCall(MatZeroEntries_SeqAIJ(A));
3381:   PetscCall(PetscArrayzero(a->a, a->i[A->rmap->n]));
3382:   PetscCall(MatSeqAIJInvalidateDiagonal(A));
3383:   if (both) A->offloadmask = PETSC_OFFLOAD_BOTH;
3384:   else A->offloadmask = PETSC_OFFLOAD_CPU;
3385:   PetscFunctionReturn(PETSC_SUCCESS);
3386: }

3388: static PetscErrorCode MatBindToCPU_SeqAIJHIPSPARSE(Mat A, PetscBool flg)
3389: {
3390:   Mat_SeqAIJ *a = (Mat_SeqAIJ *)A->data;

3392:   PetscFunctionBegin;
3393:   if (A->factortype != MAT_FACTOR_NONE) {
3394:     A->boundtocpu = flg;
3395:     PetscFunctionReturn(PETSC_SUCCESS);
3396:   }
3397:   if (flg) {
3398:     PetscCall(MatSeqAIJHIPSPARSECopyFromGPU(A));

3400:     A->ops->scale                     = MatScale_SeqAIJ;
3401:     A->ops->axpy                      = MatAXPY_SeqAIJ;
3402:     A->ops->zeroentries               = MatZeroEntries_SeqAIJ;
3403:     A->ops->mult                      = MatMult_SeqAIJ;
3404:     A->ops->multadd                   = MatMultAdd_SeqAIJ;
3405:     A->ops->multtranspose             = MatMultTranspose_SeqAIJ;
3406:     A->ops->multtransposeadd          = MatMultTransposeAdd_SeqAIJ;
3407:     A->ops->multhermitiantranspose    = NULL;
3408:     A->ops->multhermitiantransposeadd = NULL;
3409:     A->ops->productsetfromoptions     = MatProductSetFromOptions_SeqAIJ;
3410:     PetscCall(PetscMemzero(a->ops, sizeof(Mat_SeqAIJOps)));
3411:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSeqAIJCopySubArray_C", NULL));
3412:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqdensehip_C", NULL));
3413:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqdense_C", NULL));
3414:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetPreallocationCOO_C", NULL));
3415:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetValuesCOO_C", NULL));
3416:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqaijhipsparse_C", NULL));
3417:   } else {
3418:     A->ops->scale                     = MatScale_SeqAIJHIPSPARSE;
3419:     A->ops->axpy                      = MatAXPY_SeqAIJHIPSPARSE;
3420:     A->ops->zeroentries               = MatZeroEntries_SeqAIJHIPSPARSE;
3421:     A->ops->mult                      = MatMult_SeqAIJHIPSPARSE;
3422:     A->ops->multadd                   = MatMultAdd_SeqAIJHIPSPARSE;
3423:     A->ops->multtranspose             = MatMultTranspose_SeqAIJHIPSPARSE;
3424:     A->ops->multtransposeadd          = MatMultTransposeAdd_SeqAIJHIPSPARSE;
3425:     A->ops->multhermitiantranspose    = MatMultHermitianTranspose_SeqAIJHIPSPARSE;
3426:     A->ops->multhermitiantransposeadd = MatMultHermitianTransposeAdd_SeqAIJHIPSPARSE;
3427:     A->ops->productsetfromoptions     = MatProductSetFromOptions_SeqAIJHIPSPARSE;
3428:     a->ops->getarray                  = MatSeqAIJGetArray_SeqAIJHIPSPARSE;
3429:     a->ops->restorearray              = MatSeqAIJRestoreArray_SeqAIJHIPSPARSE;
3430:     a->ops->getarrayread              = MatSeqAIJGetArrayRead_SeqAIJHIPSPARSE;
3431:     a->ops->restorearrayread          = MatSeqAIJRestoreArrayRead_SeqAIJHIPSPARSE;
3432:     a->ops->getarraywrite             = MatSeqAIJGetArrayWrite_SeqAIJHIPSPARSE;
3433:     a->ops->restorearraywrite         = MatSeqAIJRestoreArrayWrite_SeqAIJHIPSPARSE;
3434:     a->ops->getcsrandmemtype          = MatSeqAIJGetCSRAndMemType_SeqAIJHIPSPARSE;
3435:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSeqAIJCopySubArray_C", MatSeqAIJCopySubArray_SeqAIJHIPSPARSE));
3436:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqdensehip_C", MatProductSetFromOptions_SeqAIJHIPSPARSE));
3437:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqdense_C", MatProductSetFromOptions_SeqAIJHIPSPARSE));
3438:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetPreallocationCOO_C", MatSetPreallocationCOO_SeqAIJHIPSPARSE));
3439:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatSetValuesCOO_C", MatSetValuesCOO_SeqAIJHIPSPARSE));
3440:     PetscCall(PetscObjectComposeFunction((PetscObject)A, "MatProductSetFromOptions_seqaijhipsparse_seqaijhipsparse_C", MatProductSetFromOptions_SeqAIJHIPSPARSE));
3441:   }
3442:   A->boundtocpu = flg;
3443:   if (flg && a->inode.size) a->inode.use = PETSC_TRUE;
3444:   else a->inode.use = PETSC_FALSE;

3446:   PetscFunctionReturn(PETSC_SUCCESS);
3447: }

3449: PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJHIPSPARSE(Mat A, MatType mtype, MatReuse reuse, Mat *newmat)
3450: {
3451:   Mat B;

3453:   PetscFunctionBegin;
3454:   PetscCall(PetscDeviceInitialize(PETSC_DEVICE_HIP)); /* first use of HIPSPARSE may be via MatConvert */
3455:   if (reuse == MAT_INITIAL_MATRIX) {
3456:     PetscCall(MatDuplicate(A, MAT_COPY_VALUES, newmat));
3457:   } else if (reuse == MAT_REUSE_MATRIX) {
3458:     PetscCall(MatCopy(A, *newmat, SAME_NONZERO_PATTERN));
3459:   }
3460:   B = *newmat;
3461:   PetscCall(PetscFree(B->defaultvectype));
3462:   PetscCall(PetscStrallocpy(VECHIP, &B->defaultvectype));
3463:   if (reuse != MAT_REUSE_MATRIX && !B->spptr) {
3464:     if (B->factortype == MAT_FACTOR_NONE) {
3465:       Mat_SeqAIJHIPSPARSE *spptr;
3466:       PetscCall(PetscNew(&spptr));
3467:       PetscCallHIPSPARSE(hipsparseCreate(&spptr->handle));
3468:       PetscCallHIPSPARSE(hipsparseSetStream(spptr->handle, PetscDefaultHipStream));
3469:       spptr->format = MAT_HIPSPARSE_CSR;
3470: #if PETSC_PKG_HIP_VERSION_GE(4, 5, 0)
3471:       spptr->spmvAlg = HIPSPARSE_SPMV_CSR_ALG1;
3472: #else
3473:       spptr->spmvAlg = HIPSPARSE_CSRMV_ALG1; /* default, since we only support csr */
3474: #endif
3475:       spptr->spmmAlg = HIPSPARSE_SPMM_CSR_ALG1; /* default, only support column-major dense matrix B */
3476:       //spptr->csr2cscAlg = HIPSPARSE_CSR2CSC_ALG1;

3478:       B->spptr = spptr;
3479:     } else {
3480:       Mat_SeqAIJHIPSPARSETriFactors *spptr;

3482:       PetscCall(PetscNew(&spptr));
3483:       PetscCallHIPSPARSE(hipsparseCreate(&spptr->handle));
3484:       PetscCallHIPSPARSE(hipsparseSetStream(spptr->handle, PetscDefaultHipStream));
3485:       B->spptr = spptr;
3486:     }
3487:     B->offloadmask = PETSC_OFFLOAD_UNALLOCATED;
3488:   }
3489:   B->ops->assemblyend    = MatAssemblyEnd_SeqAIJHIPSPARSE;
3490:   B->ops->destroy        = MatDestroy_SeqAIJHIPSPARSE;
3491:   B->ops->setoption      = MatSetOption_SeqAIJHIPSPARSE;
3492:   B->ops->setfromoptions = MatSetFromOptions_SeqAIJHIPSPARSE;
3493:   B->ops->bindtocpu      = MatBindToCPU_SeqAIJHIPSPARSE;
3494:   B->ops->duplicate      = MatDuplicate_SeqAIJHIPSPARSE;

3496:   PetscCall(MatBindToCPU_SeqAIJHIPSPARSE(B, PETSC_FALSE));
3497:   PetscCall(PetscObjectChangeTypeName((PetscObject)B, MATSEQAIJHIPSPARSE));
3498:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatHIPSPARSESetFormat_C", MatHIPSPARSESetFormat_SeqAIJHIPSPARSE));
3499: #if defined(PETSC_HAVE_HYPRE)
3500:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatConvert_seqaijhipsparse_hypre_C", MatConvert_AIJ_HYPRE));
3501: #endif
3502:   PetscCall(PetscObjectComposeFunction((PetscObject)B, "MatHIPSPARSESetUseCPUSolve_C", MatHIPSPARSESetUseCPUSolve_SeqAIJHIPSPARSE));
3503:   PetscFunctionReturn(PETSC_SUCCESS);
3504: }

3506: PETSC_EXTERN PetscErrorCode MatCreate_SeqAIJHIPSPARSE(Mat B)
3507: {
3508:   PetscFunctionBegin;
3509:   PetscCall(MatCreate_SeqAIJ(B));
3510:   PetscCall(MatConvert_SeqAIJ_SeqAIJHIPSPARSE(B, MATSEQAIJHIPSPARSE, MAT_INPLACE_MATRIX, &B));
3511:   PetscFunctionReturn(PETSC_SUCCESS);
3512: }

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

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

3521:    Options Database Keys:
3522: +  -mat_type aijhipsparse - sets the matrix type to `MATSEQAIJHIPSPARSE`
3523: .  -mat_hipsparse_storage_format csr - sets the storage format of matrices (for `MatMult()` and factors in `MatSolve()`).
3524:                                        Other options include ell (ellpack) or hyb (hybrid).
3525: . -mat_hipsparse_mult_storage_format csr - sets the storage format of matrices (for `MatMult()`). Other options include ell (ellpack) or hyb (hybrid).
3526: -  -mat_hipsparse_use_cpu_solve - Do `MatSolve()` on the CPU

3528:   Level: beginner

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

3533: PETSC_EXTERN PetscErrorCode MatSolverTypeRegister_HIPSPARSE(void)
3534: {
3535:   PetscFunctionBegin;
3536:   PetscCall(MatSolverTypeRegister(MATSOLVERHIPSPARSEBAND, MATSEQAIJ, MAT_FACTOR_LU, MatGetFactor_seqaijhipsparse_hipsparse_band));
3537:   PetscCall(MatSolverTypeRegister(MATSOLVERHIPSPARSE, MATSEQAIJHIPSPARSE, MAT_FACTOR_LU, MatGetFactor_seqaijhipsparse_hipsparse));
3538:   PetscCall(MatSolverTypeRegister(MATSOLVERHIPSPARSE, MATSEQAIJHIPSPARSE, MAT_FACTOR_CHOLESKY, MatGetFactor_seqaijhipsparse_hipsparse));
3539:   PetscCall(MatSolverTypeRegister(MATSOLVERHIPSPARSE, MATSEQAIJHIPSPARSE, MAT_FACTOR_ILU, MatGetFactor_seqaijhipsparse_hipsparse));
3540:   PetscCall(MatSolverTypeRegister(MATSOLVERHIPSPARSE, MATSEQAIJHIPSPARSE, MAT_FACTOR_ICC, MatGetFactor_seqaijhipsparse_hipsparse));

3542:   PetscFunctionReturn(PETSC_SUCCESS);
3543: }

3545: static PetscErrorCode MatResetPreallocationCOO_SeqAIJHIPSPARSE(Mat mat)
3546: {
3547:   Mat_SeqAIJHIPSPARSE *cusp = (Mat_SeqAIJHIPSPARSE *)mat->spptr;

3549:   PetscFunctionBegin;
3550:   if (!cusp) PetscFunctionReturn(PETSC_SUCCESS);
3551:   delete cusp->cooPerm;
3552:   delete cusp->cooPerm_a;
3553:   cusp->cooPerm   = NULL;
3554:   cusp->cooPerm_a = NULL;
3555:   if (cusp->use_extended_coo) {
3556:     PetscCallHIP(hipFree(cusp->jmap_d));
3557:     PetscCallHIP(hipFree(cusp->perm_d));
3558:   }
3559:   cusp->use_extended_coo = PETSC_FALSE;
3560:   PetscFunctionReturn(PETSC_SUCCESS);
3561: }

3563: static PetscErrorCode MatSeqAIJHIPSPARSE_Destroy(Mat_SeqAIJHIPSPARSE **hipsparsestruct)
3564: {
3565:   PetscFunctionBegin;
3566:   if (*hipsparsestruct) {
3567:     PetscCall(MatSeqAIJHIPSPARSEMultStruct_Destroy(&(*hipsparsestruct)->mat, (*hipsparsestruct)->format));
3568:     PetscCall(MatSeqAIJHIPSPARSEMultStruct_Destroy(&(*hipsparsestruct)->matTranspose, (*hipsparsestruct)->format));
3569:     delete (*hipsparsestruct)->workVector;
3570:     delete (*hipsparsestruct)->rowoffsets_gpu;
3571:     delete (*hipsparsestruct)->cooPerm;
3572:     delete (*hipsparsestruct)->cooPerm_a;
3573:     delete (*hipsparsestruct)->csr2csc_i;
3574:     if ((*hipsparsestruct)->handle) PetscCallHIPSPARSE(hipsparseDestroy((*hipsparsestruct)->handle));
3575:     if ((*hipsparsestruct)->jmap_d) PetscCallHIP(hipFree((*hipsparsestruct)->jmap_d));
3576:     if ((*hipsparsestruct)->perm_d) PetscCallHIP(hipFree((*hipsparsestruct)->perm_d));
3577:     PetscCall(PetscFree(*hipsparsestruct));
3578:   }
3579:   PetscFunctionReturn(PETSC_SUCCESS);
3580: }

3582: static PetscErrorCode CsrMatrix_Destroy(CsrMatrix **mat)
3583: {
3584:   PetscFunctionBegin;
3585:   if (*mat) {
3586:     delete (*mat)->values;
3587:     delete (*mat)->column_indices;
3588:     delete (*mat)->row_offsets;
3589:     delete *mat;
3590:     *mat = 0;
3591:   }
3592:   PetscFunctionReturn(PETSC_SUCCESS);
3593: }

3595: static PetscErrorCode MatSeqAIJHIPSPARSEMultStruct_Destroy(Mat_SeqAIJHIPSPARSETriFactorStruct **trifactor)
3596: {
3597:   PetscFunctionBegin;
3598:   if (*trifactor) {
3599:     if ((*trifactor)->descr) PetscCallHIPSPARSE(hipsparseDestroyMatDescr((*trifactor)->descr));
3600:     if ((*trifactor)->solveInfo) PetscCallHIPSPARSE(hipsparseDestroyCsrsvInfo((*trifactor)->solveInfo));
3601:     PetscCall(CsrMatrix_Destroy(&(*trifactor)->csrMat));
3602:     if ((*trifactor)->solveBuffer) PetscCallHIP(hipFree((*trifactor)->solveBuffer));
3603:     if ((*trifactor)->AA_h) PetscCallHIP(hipHostFree((*trifactor)->AA_h));
3604:     if ((*trifactor)->csr2cscBuffer) PetscCallHIP(hipFree((*trifactor)->csr2cscBuffer));
3605:     PetscCall(PetscFree(*trifactor));
3606:   }
3607:   PetscFunctionReturn(PETSC_SUCCESS);
3608: }

3610: static PetscErrorCode MatSeqAIJHIPSPARSEMultStruct_Destroy(Mat_SeqAIJHIPSPARSEMultStruct **matstruct, MatHIPSPARSEStorageFormat format)
3611: {
3612:   CsrMatrix *mat;

3614:   PetscFunctionBegin;
3615:   if (*matstruct) {
3616:     if ((*matstruct)->mat) {
3617:       if (format == MAT_HIPSPARSE_ELL || format == MAT_HIPSPARSE_HYB) {
3618:         hipsparseHybMat_t hybMat = (hipsparseHybMat_t)(*matstruct)->mat;
3619:         PetscCallHIPSPARSE(hipsparseDestroyHybMat(hybMat));
3620:       } else {
3621:         mat = (CsrMatrix *)(*matstruct)->mat;
3622:         PetscCall(CsrMatrix_Destroy(&mat));
3623:       }
3624:     }
3625:     if ((*matstruct)->descr) PetscCallHIPSPARSE(hipsparseDestroyMatDescr((*matstruct)->descr));
3626:     delete (*matstruct)->cprowIndices;
3627:     if ((*matstruct)->alpha_one) PetscCallHIP(hipFree((*matstruct)->alpha_one));
3628:     if ((*matstruct)->beta_zero) PetscCallHIP(hipFree((*matstruct)->beta_zero));
3629:     if ((*matstruct)->beta_one) PetscCallHIP(hipFree((*matstruct)->beta_one));

3631:     Mat_SeqAIJHIPSPARSEMultStruct *mdata = *matstruct;
3632:     if (mdata->matDescr) PetscCallHIPSPARSE(hipsparseDestroySpMat(mdata->matDescr));
3633:     for (int i = 0; i < 3; i++) {
3634:       if (mdata->hipSpMV[i].initialized) {
3635:         PetscCallHIP(hipFree(mdata->hipSpMV[i].spmvBuffer));
3636:         PetscCallHIPSPARSE(hipsparseDestroyDnVec(mdata->hipSpMV[i].vecXDescr));
3637:         PetscCallHIPSPARSE(hipsparseDestroyDnVec(mdata->hipSpMV[i].vecYDescr));
3638:       }
3639:     }
3640:     delete *matstruct;
3641:     *matstruct = NULL;
3642:   }
3643:   PetscFunctionReturn(PETSC_SUCCESS);
3644: }

3646: PetscErrorCode MatSeqAIJHIPSPARSETriFactors_Reset(Mat_SeqAIJHIPSPARSETriFactors_p *trifactors)
3647: {
3648:   Mat_SeqAIJHIPSPARSETriFactors *fs = *trifactors;

3650:   PetscFunctionBegin;
3651:   if (fs) {
3652:     PetscCall(MatSeqAIJHIPSPARSEMultStruct_Destroy(&fs->loTriFactorPtr));
3653:     PetscCall(MatSeqAIJHIPSPARSEMultStruct_Destroy(&fs->upTriFactorPtr));
3654:     PetscCall(MatSeqAIJHIPSPARSEMultStruct_Destroy(&fs->loTriFactorPtrTranspose));
3655:     PetscCall(MatSeqAIJHIPSPARSEMultStruct_Destroy(&fs->upTriFactorPtrTranspose));
3656:     delete fs->rpermIndices;
3657:     delete fs->cpermIndices;
3658:     delete fs->workVector;
3659:     fs->rpermIndices = NULL;
3660:     fs->cpermIndices = NULL;
3661:     fs->workVector   = NULL;
3662:     if (fs->a_band_d) PetscCallHIP(hipFree(fs->a_band_d));
3663:     if (fs->i_band_d) PetscCallHIP(hipFree(fs->i_band_d));
3664:     fs->init_dev_prop = PETSC_FALSE;
3665: #if PETSC_PKG_HIP_VERSION_GE(4, 5, 0)
3666:     PetscCallHIP(hipFree(fs->csrRowPtr));
3667:     PetscCallHIP(hipFree(fs->csrColIdx));
3668:     PetscCallHIP(hipFree(fs->csrVal));
3669:     PetscCallHIP(hipFree(fs->X));
3670:     PetscCallHIP(hipFree(fs->Y));
3671:     // PetscCallHIP(hipFree(fs->factBuffer_M)); /* No needed since factBuffer_M shares with one of spsvBuffer_L/U */
3672:     PetscCallHIP(hipFree(fs->spsvBuffer_L));
3673:     PetscCallHIP(hipFree(fs->spsvBuffer_U));
3674:     PetscCallHIP(hipFree(fs->spsvBuffer_Lt));
3675:     PetscCallHIP(hipFree(fs->spsvBuffer_Ut));
3676:     PetscCallHIPSPARSE(hipsparseDestroyMatDescr(fs->matDescr_M));
3677:     if (fs->spMatDescr_L) PetscCallHIPSPARSE(hipsparseDestroySpMat(fs->spMatDescr_L));
3678:     if (fs->spMatDescr_U) PetscCallHIPSPARSE(hipsparseDestroySpMat(fs->spMatDescr_U));
3679:     PetscCallHIPSPARSE(hipsparseSpSV_destroyDescr(fs->spsvDescr_L));
3680:     PetscCallHIPSPARSE(hipsparseSpSV_destroyDescr(fs->spsvDescr_Lt));
3681:     PetscCallHIPSPARSE(hipsparseSpSV_destroyDescr(fs->spsvDescr_U));
3682:     PetscCallHIPSPARSE(hipsparseSpSV_destroyDescr(fs->spsvDescr_Ut));
3683:     if (fs->dnVecDescr_X) PetscCallHIPSPARSE(hipsparseDestroyDnVec(fs->dnVecDescr_X));
3684:     if (fs->dnVecDescr_Y) PetscCallHIPSPARSE(hipsparseDestroyDnVec(fs->dnVecDescr_Y));
3685:     PetscCallHIPSPARSE(hipsparseDestroyCsrilu02Info(fs->ilu0Info_M));
3686:     PetscCallHIPSPARSE(hipsparseDestroyCsric02Info(fs->ic0Info_M));

3688:     fs->createdTransposeSpSVDescr    = PETSC_FALSE;
3689:     fs->updatedTransposeSpSVAnalysis = PETSC_FALSE;
3690: #endif
3691:   }
3692:   PetscFunctionReturn(PETSC_SUCCESS);
3693: }

3695: static PetscErrorCode MatSeqAIJHIPSPARSETriFactors_Destroy(Mat_SeqAIJHIPSPARSETriFactors **trifactors)
3696: {
3697:   hipsparseHandle_t handle;

3699:   PetscFunctionBegin;
3700:   if (*trifactors) {
3701:     PetscCall(MatSeqAIJHIPSPARSETriFactors_Reset(trifactors));
3702:     if ((handle = (*trifactors)->handle)) PetscCallHIPSPARSE(hipsparseDestroy(handle));
3703:     PetscCall(PetscFree(*trifactors));
3704:   }
3705:   PetscFunctionReturn(PETSC_SUCCESS);
3706: }

3708: struct IJCompare {
3709:   __host__ __device__ inline bool operator()(const thrust::tuple<PetscInt, PetscInt> &t1, const thrust::tuple<PetscInt, PetscInt> &t2)
3710:   {
3711:     if (t1.get<0>() < t2.get<0>()) return true;
3712:     if (t1.get<0>() == t2.get<0>()) return t1.get<1>() < t2.get<1>();
3713:     return false;
3714:   }
3715: };

3717: struct IJEqual {
3718:   __host__ __device__ inline bool operator()(const thrust::tuple<PetscInt, PetscInt> &t1, const thrust::tuple<PetscInt, PetscInt> &t2)
3719:   {
3720:     if (t1.get<0>() != t2.get<0>() || t1.get<1>() != t2.get<1>()) return false;
3721:     return true;
3722:   }
3723: };

3725: struct IJDiff {
3726:   __host__ __device__ inline PetscInt operator()(const PetscInt &t1, const PetscInt &t2) { return t1 == t2 ? 0 : 1; }
3727: };

3729: struct IJSum {
3730:   __host__ __device__ inline PetscInt operator()(const PetscInt &t1, const PetscInt &t2) { return t1 || t2; }
3731: };

3733: PetscErrorCode MatSetValuesCOO_SeqAIJHIPSPARSE_Basic(Mat A, const PetscScalar v[], InsertMode imode)
3734: {
3735:   Mat_SeqAIJHIPSPARSE                  *cusp      = (Mat_SeqAIJHIPSPARSE *)A->spptr;
3736:   Mat_SeqAIJ                           *a         = (Mat_SeqAIJ *)A->data;
3737:   THRUSTARRAY                          *cooPerm_v = NULL;
3738:   thrust::device_ptr<const PetscScalar> d_v;
3739:   CsrMatrix                            *matrix;
3740:   PetscInt                              n;

3742:   PetscFunctionBegin;
3743:   PetscCheck(cusp, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing HIPSPARSE struct");
3744:   PetscCheck(cusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing HIPSPARSE CsrMatrix");
3745:   if (!cusp->cooPerm) {
3746:     PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
3747:     PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
3748:     PetscFunctionReturn(PETSC_SUCCESS);
3749:   }
3750:   matrix = (CsrMatrix *)cusp->mat->mat;
3751:   PetscCheck(matrix->values, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing HIP memory");
3752:   if (!v) {
3753:     if (imode == INSERT_VALUES) thrust::fill(thrust::device, matrix->values->begin(), matrix->values->end(), 0.);
3754:     goto finalize;
3755:   }
3756:   n = cusp->cooPerm->size();
3757:   if (isHipMem(v)) d_v = thrust::device_pointer_cast(v);
3758:   else {
3759:     cooPerm_v = new THRUSTARRAY(n);
3760:     cooPerm_v->assign(v, v + n);
3761:     d_v = cooPerm_v->data();
3762:     PetscCall(PetscLogCpuToGpu(n * sizeof(PetscScalar)));
3763:   }
3764:   PetscCall(PetscLogGpuTimeBegin());
3765:   if (imode == ADD_VALUES) { /* ADD VALUES means add to existing ones */
3766:     if (cusp->cooPerm_a) {   /* there are repeated entries in d_v[], and we need to add these them */
3767:       THRUSTARRAY *cooPerm_w = new THRUSTARRAY(matrix->values->size());
3768:       auto         vbit      = thrust::make_permutation_iterator(d_v, cusp->cooPerm->begin());
3769:       /* thrust::reduce_by_key(keys_first,keys_last,values_first,keys_output,values_output)
3770:         cooPerm_a = [0,0,1,2,3,4]. The length is n, number of nonozeros in d_v[].
3771:         cooPerm_a is ordered. d_v[i] is the cooPerm_a[i]-th unique nonzero.
3772:       */
3773:       thrust::reduce_by_key(cusp->cooPerm_a->begin(), cusp->cooPerm_a->end(), vbit, thrust::make_discard_iterator(), cooPerm_w->begin(), thrust::equal_to<PetscInt>(), thrust::plus<PetscScalar>());
3774:       thrust::transform(cooPerm_w->begin(), cooPerm_w->end(), matrix->values->begin(), matrix->values->begin(), thrust::plus<PetscScalar>());
3775:       delete cooPerm_w;
3776:     } else {
3777:       /* all nonzeros in d_v[] are unique entries */
3778:       auto zibit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(d_v, cusp->cooPerm->begin()), matrix->values->begin()));
3779:       auto zieit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(d_v, cusp->cooPerm->end()), matrix->values->end()));
3780:       thrust::for_each(zibit, zieit, VecHIPPlusEquals()); /* values[i] += d_v[cooPerm[i]]  */
3781:     }
3782:   } else {
3783:     if (cusp->cooPerm_a) { /* repeated entries in COO, with INSERT_VALUES -> reduce */
3784:       auto vbit = thrust::make_permutation_iterator(d_v, cusp->cooPerm->begin());
3785:       thrust::reduce_by_key(cusp->cooPerm_a->begin(), cusp->cooPerm_a->end(), vbit, thrust::make_discard_iterator(), matrix->values->begin(), thrust::equal_to<PetscInt>(), thrust::plus<PetscScalar>());
3786:     } else {
3787:       auto zibit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(d_v, cusp->cooPerm->begin()), matrix->values->begin()));
3788:       auto zieit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(d_v, cusp->cooPerm->end()), matrix->values->end()));
3789:       thrust::for_each(zibit, zieit, VecHIPEquals());
3790:     }
3791:   }
3792:   PetscCall(PetscLogGpuTimeEnd());
3793: finalize:
3794:   delete cooPerm_v;
3795:   A->offloadmask = PETSC_OFFLOAD_GPU;
3796:   PetscCall(PetscObjectStateIncrease((PetscObject)A));
3797:   /* shorter version of MatAssemblyEnd_SeqAIJ */
3798:   PetscCall(PetscInfo(A, "Matrix size: %" PetscInt_FMT " X %" PetscInt_FMT "; storage space: 0 unneeded,%" PetscInt_FMT " used\n", A->rmap->n, A->cmap->n, a->nz));
3799:   PetscCall(PetscInfo(A, "Number of mallocs during MatSetValues() is 0\n"));
3800:   PetscCall(PetscInfo(A, "Maximum nonzeros in any row is %" PetscInt_FMT "\n", a->rmax));
3801:   a->reallocs = 0;
3802:   A->info.mallocs += 0;
3803:   A->info.nz_unneeded = 0;
3804:   A->assembled = A->was_assembled = PETSC_TRUE;
3805:   A->num_ass++;
3806:   PetscFunctionReturn(PETSC_SUCCESS);
3807: }

3809: PetscErrorCode MatSeqAIJHIPSPARSEInvalidateTranspose(Mat A, PetscBool destroy)
3810: {
3811:   Mat_SeqAIJHIPSPARSE *cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;

3813:   PetscFunctionBegin;
3814:   PetscCheckTypeName(A, MATSEQAIJHIPSPARSE);
3815:   if (!cusp) PetscFunctionReturn(PETSC_SUCCESS);
3816:   if (destroy) {
3817:     PetscCall(MatSeqAIJHIPSPARSEMultStruct_Destroy(&cusp->matTranspose, cusp->format));
3818:     delete cusp->csr2csc_i;
3819:     cusp->csr2csc_i = NULL;
3820:   }
3821:   A->transupdated = PETSC_FALSE;
3822:   PetscFunctionReturn(PETSC_SUCCESS);
3823: }

3825: PetscErrorCode MatSetPreallocationCOO_SeqAIJHIPSPARSE_Basic(Mat A, PetscCount n, PetscInt coo_i[], PetscInt coo_j[])
3826: {
3827:   Mat_SeqAIJHIPSPARSE *cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
3828:   Mat_SeqAIJ          *a    = (Mat_SeqAIJ *)A->data;
3829:   PetscInt             cooPerm_n, nzr = 0;

3831:   PetscFunctionBegin;
3832:   PetscCall(PetscLayoutSetUp(A->rmap));
3833:   PetscCall(PetscLayoutSetUp(A->cmap));
3834:   cooPerm_n = cusp->cooPerm ? cusp->cooPerm->size() : 0;
3835:   if (n != cooPerm_n) {
3836:     delete cusp->cooPerm;
3837:     delete cusp->cooPerm_a;
3838:     cusp->cooPerm   = NULL;
3839:     cusp->cooPerm_a = NULL;
3840:   }
3841:   if (n) {
3842:     thrust::device_ptr<PetscInt> d_i, d_j;
3843:     PetscInt                    *d_raw_i, *d_raw_j;
3844:     PetscBool                    free_raw_i = PETSC_FALSE, free_raw_j = PETSC_FALSE;
3845:     PetscMemType                 imtype, jmtype;

3847:     PetscCall(PetscGetMemType(coo_i, &imtype));
3848:     if (PetscMemTypeHost(imtype)) {
3849:       PetscCallHIP(hipMalloc(&d_raw_i, sizeof(PetscInt) * n));
3850:       PetscCallHIP(hipMemcpy(d_raw_i, coo_i, sizeof(PetscInt) * n, hipMemcpyHostToDevice));
3851:       d_i        = thrust::device_pointer_cast(d_raw_i);
3852:       free_raw_i = PETSC_TRUE;
3853:       PetscCall(PetscLogCpuToGpu(1. * n * sizeof(PetscInt)));
3854:     } else {
3855:       d_i = thrust::device_pointer_cast(coo_i);
3856:     }

3858:     PetscCall(PetscGetMemType(coo_j, &jmtype));
3859:     if (PetscMemTypeHost(jmtype)) { // MatSetPreallocationCOO_MPIAIJHIPSPARSE_Basic() passes device coo_i[] and host coo_j[]!
3860:       PetscCallHIP(hipMalloc(&d_raw_j, sizeof(PetscInt) * n));
3861:       PetscCallHIP(hipMemcpy(d_raw_j, coo_j, sizeof(PetscInt) * n, hipMemcpyHostToDevice));
3862:       d_j        = thrust::device_pointer_cast(d_raw_j);
3863:       free_raw_j = PETSC_TRUE;
3864:       PetscCall(PetscLogCpuToGpu(1. * n * sizeof(PetscInt)));
3865:     } else {
3866:       d_j = thrust::device_pointer_cast(coo_j);
3867:     }

3869:     THRUSTINTARRAY ii(A->rmap->n);

3871:     if (!cusp->cooPerm) cusp->cooPerm = new THRUSTINTARRAY(n);
3872:     if (!cusp->cooPerm_a) cusp->cooPerm_a = new THRUSTINTARRAY(n);
3873:     /* Ex.
3874:       n = 6
3875:       coo_i = [3,3,1,4,1,4]
3876:       coo_j = [3,2,2,5,2,6]
3877:     */
3878:     auto fkey = thrust::make_zip_iterator(thrust::make_tuple(d_i, d_j));
3879:     auto ekey = thrust::make_zip_iterator(thrust::make_tuple(d_i + n, d_j + n));

3881:     PetscCall(PetscLogGpuTimeBegin());
3882:     thrust::sequence(thrust::device, cusp->cooPerm->begin(), cusp->cooPerm->end(), 0);
3883:     thrust::sort_by_key(fkey, ekey, cusp->cooPerm->begin(), IJCompare()); /* sort by row, then by col */
3884:     (*cusp->cooPerm_a).assign(d_i, d_i + n);                              /* copy the sorted array */
3885:     THRUSTINTARRAY w(d_j, d_j + n);
3886:     /*
3887:       d_i     = [1,1,3,3,4,4]
3888:       d_j     = [2,2,2,3,5,6]
3889:       cooPerm = [2,4,1,0,3,5]
3890:     */
3891:     auto nekey = thrust::unique(fkey, ekey, IJEqual()); /* unique (d_i, d_j) */
3892:     /*
3893:       d_i     = [1,3,3,4,4,x]
3894:                             ^ekey
3895:       d_j     = [2,2,3,5,6,x]
3896:                            ^nekye
3897:     */
3898:     if (nekey == ekey) { /* all entries are unique */
3899:       delete cusp->cooPerm_a;
3900:       cusp->cooPerm_a = NULL;
3901:     } else { /* Stefano: I couldn't come up with a more elegant algorithm */
3902:       /* idea: any change in i or j in the (i,j) sequence implies a new nonzero */
3903:       adjacent_difference(cusp->cooPerm_a->begin(), cusp->cooPerm_a->end(), cusp->cooPerm_a->begin(), IJDiff()); /* cooPerm_a: [1,1,3,3,4,4] => [1,0,1,0,1,0]*/
3904:       adjacent_difference(w.begin(), w.end(), w.begin(), IJDiff());                                              /* w:         [2,2,2,3,5,6] => [2,0,0,1,1,1]*/
3905:       (*cusp->cooPerm_a)[0] = 0;                                                                                 /* clear the first entry, though accessing an entry on device implies a hipMemcpy */
3906:       w[0]                  = 0;
3907:       thrust::transform(cusp->cooPerm_a->begin(), cusp->cooPerm_a->end(), w.begin(), cusp->cooPerm_a->begin(), IJSum());            /* cooPerm_a =          [0,0,1,1,1,1]*/
3908:       thrust::inclusive_scan(cusp->cooPerm_a->begin(), cusp->cooPerm_a->end(), cusp->cooPerm_a->begin(), thrust::plus<PetscInt>()); /*cooPerm_a=[0,0,1,2,3,4]*/
3909:     }
3910:     thrust::counting_iterator<PetscInt> search_begin(0);
3911:     thrust::upper_bound(d_i, nekey.get_iterator_tuple().get<0>(), /* binary search entries of [0,1,2,3,4,5,6) in ordered array d_i = [1,3,3,4,4], supposing A->rmap->n = 6. */
3912:                         search_begin, search_begin + A->rmap->n,  /* return in ii[] the index of last position in d_i[] where value could be inserted without violating the ordering */
3913:                         ii.begin());                              /* ii = [0,1,1,3,5,5]. A leading 0 will be added later */
3914:     PetscCall(PetscLogGpuTimeEnd());

3916:     PetscCall(MatSeqXAIJFreeAIJ(A, &a->a, &a->j, &a->i));
3917:     a->singlemalloc = PETSC_FALSE;
3918:     a->free_a       = PETSC_TRUE;
3919:     a->free_ij      = PETSC_TRUE;
3920:     PetscCall(PetscMalloc1(A->rmap->n + 1, &a->i));
3921:     a->i[0] = 0; /* a->i = [0,0,1,1,3,5,5] */
3922:     PetscCallHIP(hipMemcpy(a->i + 1, ii.data().get(), A->rmap->n * sizeof(PetscInt), hipMemcpyDeviceToHost));
3923:     a->nz = a->maxnz = a->i[A->rmap->n];
3924:     a->rmax          = 0;
3925:     PetscCall(PetscMalloc1(a->nz, &a->a));
3926:     PetscCall(PetscMalloc1(a->nz, &a->j));
3927:     PetscCallHIP(hipMemcpy(a->j, thrust::raw_pointer_cast(d_j), a->nz * sizeof(PetscInt), hipMemcpyDeviceToHost));
3928:     if (!a->ilen) PetscCall(PetscMalloc1(A->rmap->n, &a->ilen));
3929:     if (!a->imax) PetscCall(PetscMalloc1(A->rmap->n, &a->imax));
3930:     for (PetscInt i = 0; i < A->rmap->n; i++) {
3931:       const PetscInt nnzr = a->i[i + 1] - a->i[i];
3932:       nzr += (PetscInt) !!(nnzr);
3933:       a->ilen[i] = a->imax[i] = nnzr;
3934:       a->rmax                 = PetscMax(a->rmax, nnzr);
3935:     }
3936:     a->nonzerorowcnt = nzr;
3937:     A->preallocated  = PETSC_TRUE;
3938:     PetscCall(PetscLogGpuToCpu((A->rmap->n + a->nz) * sizeof(PetscInt)));
3939:     PetscCall(MatMarkDiagonal_SeqAIJ(A));
3940:     if (free_raw_i) PetscCallHIP(hipFree(d_raw_i));
3941:     if (free_raw_j) PetscCallHIP(hipFree(d_raw_j));
3942:   } else PetscCall(MatSeqAIJSetPreallocation(A, 0, NULL));
3943:   PetscCall(MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE));
3944:   /* We want to allocate the HIPSPARSE struct for matvec now.
3945:      The code is so convoluted now that I prefer to copy zeros */
3946:   PetscCall(PetscArrayzero(a->a, a->nz));
3947:   PetscCall(MatCheckCompressedRow(A, nzr, &a->compressedrow, a->i, A->rmap->n, 0.6));
3948:   A->offloadmask = PETSC_OFFLOAD_CPU;
3949:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
3950:   PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(A, PETSC_TRUE));
3951:   PetscFunctionReturn(PETSC_SUCCESS);
3952: }

3954: PetscErrorCode MatSetPreallocationCOO_SeqAIJHIPSPARSE(Mat mat, PetscCount coo_n, PetscInt coo_i[], PetscInt coo_j[])
3955: {
3956:   Mat_SeqAIJ          *seq;
3957:   Mat_SeqAIJHIPSPARSE *dev;
3958:   PetscBool            coo_basic = PETSC_TRUE;
3959:   PetscMemType         mtype     = PETSC_MEMTYPE_DEVICE;

3961:   PetscFunctionBegin;
3962:   PetscCall(MatResetPreallocationCOO_SeqAIJ(mat));
3963:   PetscCall(MatResetPreallocationCOO_SeqAIJHIPSPARSE(mat));
3964:   if (coo_i) {
3965:     PetscCall(PetscGetMemType(coo_i, &mtype));
3966:     if (PetscMemTypeHost(mtype)) {
3967:       for (PetscCount k = 0; k < coo_n; k++) {
3968:         if (coo_i[k] < 0 || coo_j[k] < 0) {
3969:           coo_basic = PETSC_FALSE;
3970:           break;
3971:         }
3972:       }
3973:     }
3974:   }

3976:   if (coo_basic) { /* i,j are on device or do not contain negative indices */
3977:     PetscCall(MatSetPreallocationCOO_SeqAIJHIPSPARSE_Basic(mat, coo_n, coo_i, coo_j));
3978:   } else {
3979:     PetscCall(MatSetPreallocationCOO_SeqAIJ(mat, coo_n, coo_i, coo_j));
3980:     mat->offloadmask = PETSC_OFFLOAD_CPU;
3981:     PetscCall(MatSeqAIJHIPSPARSECopyToGPU(mat));
3982:     seq = static_cast<Mat_SeqAIJ *>(mat->data);
3983:     dev = static_cast<Mat_SeqAIJHIPSPARSE *>(mat->spptr);
3984:     PetscCallHIP(hipMalloc((void **)&dev->jmap_d, (seq->nz + 1) * sizeof(PetscCount)));
3985:     PetscCallHIP(hipMemcpy(dev->jmap_d, seq->jmap, (seq->nz + 1) * sizeof(PetscCount), hipMemcpyHostToDevice));
3986:     PetscCallHIP(hipMalloc((void **)&dev->perm_d, seq->Atot * sizeof(PetscCount)));
3987:     PetscCallHIP(hipMemcpy(dev->perm_d, seq->perm, seq->Atot * sizeof(PetscCount), hipMemcpyHostToDevice));
3988:     dev->use_extended_coo = PETSC_TRUE;
3989:   }
3990:   PetscFunctionReturn(PETSC_SUCCESS);
3991: }

3993: __global__ static void MatAddCOOValues(const PetscScalar kv[], PetscCount nnz, const PetscCount jmap[], const PetscCount perm[], InsertMode imode, PetscScalar a[])
3994: {
3995:   PetscCount       i         = blockIdx.x * blockDim.x + threadIdx.x;
3996:   const PetscCount grid_size = gridDim.x * blockDim.x;
3997:   for (; i < nnz; i += grid_size) {
3998:     PetscScalar sum = 0.0;
3999:     for (PetscCount k = jmap[i]; k < jmap[i + 1]; k++) sum += kv[perm[k]];
4000:     a[i] = (imode == INSERT_VALUES ? 0.0 : a[i]) + sum;
4001:   }
4002: }

4004: PetscErrorCode MatSetValuesCOO_SeqAIJHIPSPARSE(Mat A, const PetscScalar v[], InsertMode imode)
4005: {
4006:   Mat_SeqAIJ          *seq  = (Mat_SeqAIJ *)A->data;
4007:   Mat_SeqAIJHIPSPARSE *dev  = (Mat_SeqAIJHIPSPARSE *)A->spptr;
4008:   PetscCount           Annz = seq->nz;
4009:   PetscMemType         memtype;
4010:   const PetscScalar   *v1 = v;
4011:   PetscScalar         *Aa;

4013:   PetscFunctionBegin;
4014:   if (dev->use_extended_coo) {
4015:     PetscCall(PetscGetMemType(v, &memtype));
4016:     if (PetscMemTypeHost(memtype)) { /* If user gave v[] in host, we might need to copy it to device if any */
4017:       PetscCallHIP(hipMalloc((void **)&v1, seq->coo_n * sizeof(PetscScalar)));
4018:       PetscCallHIP(hipMemcpy((void *)v1, v, seq->coo_n * sizeof(PetscScalar), hipMemcpyHostToDevice));
4019:     }

4021:     if (imode == INSERT_VALUES) PetscCall(MatSeqAIJHIPSPARSEGetArrayWrite(A, &Aa));
4022:     else PetscCall(MatSeqAIJHIPSPARSEGetArray(A, &Aa));

4024:     if (Annz) {
4025:       hipLaunchKernelGGL(HIP_KERNEL_NAME(MatAddCOOValues), dim3((Annz + 255) / 256), dim3(256), 0, PetscDefaultHipStream, v1, Annz, dev->jmap_d, dev->perm_d, imode, Aa);
4026:       PetscCallHIP(hipPeekAtLastError());
4027:     }

4029:     if (imode == INSERT_VALUES) PetscCall(MatSeqAIJHIPSPARSERestoreArrayWrite(A, &Aa));
4030:     else PetscCall(MatSeqAIJHIPSPARSERestoreArray(A, &Aa));

4032:     if (PetscMemTypeHost(memtype)) PetscCallHIP(hipFree((void *)v1));
4033:   } else {
4034:     PetscCall(MatSetValuesCOO_SeqAIJHIPSPARSE_Basic(A, v, imode));
4035:   }
4036:   PetscFunctionReturn(PETSC_SUCCESS);
4037: }

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

4042:     Not Collective

4044:     Input Parameters:
4045: +   A - the matrix
4046: -   compressed - `PETSC_TRUE` or `PETSC_FALSE` indicating the matrix data structure should be always returned in compressed form

4048:     Output Parameters:
4049: +   i - the CSR row pointers
4050: -   j - the CSR column indices

4052:     Level: developer

4054:     Note:
4055:       When compressed is true, the CSR structure does not contain empty rows

4057: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSERestoreIJ()`, `MatSeqAIJHIPSPARSEGetArrayRead()`
4058: @*/
4059: PetscErrorCode MatSeqAIJHIPSPARSEGetIJ(Mat A, PetscBool compressed, const int **i, const int **j)
4060: {
4061:   Mat_SeqAIJHIPSPARSE *cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
4062:   Mat_SeqAIJ          *a    = (Mat_SeqAIJ *)A->data;
4063:   CsrMatrix           *csr;

4065:   PetscFunctionBegin;
4067:   if (!i || !j) PetscFunctionReturn(PETSC_SUCCESS);
4068:   PetscCheckTypeName(A, MATSEQAIJHIPSPARSE);
4069:   PetscCheck(cusp->format != MAT_HIPSPARSE_ELL && cusp->format != MAT_HIPSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
4070:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
4071:   PetscCheck(cusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJHIPSPARSEMultStruct");
4072:   csr = (CsrMatrix *)cusp->mat->mat;
4073:   if (i) {
4074:     if (!compressed && a->compressedrow.use) { /* need full row offset */
4075:       if (!cusp->rowoffsets_gpu) {
4076:         cusp->rowoffsets_gpu = new THRUSTINTARRAY32(A->rmap->n + 1);
4077:         cusp->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
4078:         PetscCall(PetscLogCpuToGpu((A->rmap->n + 1) * sizeof(PetscInt)));
4079:       }
4080:       *i = cusp->rowoffsets_gpu->data().get();
4081:     } else *i = csr->row_offsets->data().get();
4082:   }
4083:   if (j) *j = csr->column_indices->data().get();
4084:   PetscFunctionReturn(PETSC_SUCCESS);
4085: }

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

4090:     Not Collective

4092:     Input Parameters:
4093: +   A - the matrix
4094: .   compressed - `PETSC_TRUE` or `PETSC_FALSE` indicating the matrix data structure should be always returned in compressed form
4095: .   i - the CSR row pointers
4096: -   j - the CSR column indices

4098:     Level: developer

4100: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetIJ()`
4101: @*/
4102: PetscErrorCode MatSeqAIJHIPSPARSERestoreIJ(Mat A, PetscBool compressed, const int **i, const int **j)
4103: {
4104:   PetscFunctionBegin;
4106:   PetscCheckTypeName(A, MATSEQAIJHIPSPARSE);
4107:   if (i) *i = NULL;
4108:   if (j) *j = NULL;
4109:   PetscFunctionReturn(PETSC_SUCCESS);
4110: }

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

4115:    Not Collective

4117:    Input Parameter:
4118: .   A - a `MATSEQAIJHIPSPARSE` matrix

4120:    Output Parameter:
4121: .   a - pointer to the device data

4123:    Level: developer

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

4128: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetArray()`, `MatSeqAIJHIPSPARSEGetArrayWrite()`, `MatSeqAIJHIPSPARSERestoreArrayRead()`
4129: @*/
4130: PetscErrorCode MatSeqAIJHIPSPARSEGetArrayRead(Mat A, const PetscScalar **a)
4131: {
4132:   Mat_SeqAIJHIPSPARSE *cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
4133:   CsrMatrix           *csr;

4135:   PetscFunctionBegin;
4138:   PetscCheckTypeName(A, MATSEQAIJHIPSPARSE);
4139:   PetscCheck(cusp->format != MAT_HIPSPARSE_ELL && cusp->format != MAT_HIPSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
4140:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
4141:   PetscCheck(cusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJHIPSPARSEMultStruct");
4142:   csr = (CsrMatrix *)cusp->mat->mat;
4143:   PetscCheck(csr->values, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing HIP memory");
4144:   *a = csr->values->data().get();
4145:   PetscFunctionReturn(PETSC_SUCCESS);
4146: }

4148: /*@C
4149:    MatSeqAIJHIPSPARSERestoreArrayRead - restore the read-only access array obtained from `MatSeqAIJHIPSPARSEGetArrayRead()`

4151:    Not Collective

4153:    Input Parameters:
4154: +   A - a `MATSEQAIJHIPSPARSE` matrix
4155: -   a - pointer to the device data

4157:    Level: developer

4159: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetArrayRead()`
4160: @*/
4161: PetscErrorCode MatSeqAIJHIPSPARSERestoreArrayRead(Mat A, const PetscScalar **a)
4162: {
4163:   PetscFunctionBegin;
4166:   PetscCheckTypeName(A, MATSEQAIJHIPSPARSE);
4167:   *a = NULL;
4168:   PetscFunctionReturn(PETSC_SUCCESS);
4169: }

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

4174:    Not Collective

4176:    Input Parameter:
4177: .   A - a `MATSEQAIJHIPSPARSE` matrix

4179:    Output Parameter:
4180: .   a - pointer to the device data

4182:    Level: developer

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

4187: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetArrayRead()`, `MatSeqAIJHIPSPARSEGetArrayWrite()`, `MatSeqAIJHIPSPARSERestoreArray()`
4188: @*/
4189: PetscErrorCode MatSeqAIJHIPSPARSEGetArray(Mat A, PetscScalar **a)
4190: {
4191:   Mat_SeqAIJHIPSPARSE *cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
4192:   CsrMatrix           *csr;

4194:   PetscFunctionBegin;
4197:   PetscCheckTypeName(A, MATSEQAIJHIPSPARSE);
4198:   PetscCheck(cusp->format != MAT_HIPSPARSE_ELL && cusp->format != MAT_HIPSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
4199:   PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
4200:   PetscCheck(cusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJHIPSPARSEMultStruct");
4201:   csr = (CsrMatrix *)cusp->mat->mat;
4202:   PetscCheck(csr->values, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing HIP memory");
4203:   *a             = csr->values->data().get();
4204:   A->offloadmask = PETSC_OFFLOAD_GPU;
4205:   PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(A, PETSC_FALSE));
4206:   PetscFunctionReturn(PETSC_SUCCESS);
4207: }
4208: /*@C
4209:    MatSeqAIJHIPSPARSERestoreArray - restore the read-write access array obtained from `MatSeqAIJHIPSPARSEGetArray()`

4211:    Not Collective

4213:    Input Parameters:
4214: +   A - a `MATSEQAIJHIPSPARSE` matrix
4215: -   a - pointer to the device data

4217:    Level: developer

4219: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetArray()`
4220: @*/
4221: PetscErrorCode MatSeqAIJHIPSPARSERestoreArray(Mat A, PetscScalar **a)
4222: {
4223:   PetscFunctionBegin;
4226:   PetscCheckTypeName(A, MATSEQAIJHIPSPARSE);
4227:   PetscCall(MatSeqAIJInvalidateDiagonal(A));
4228:   PetscCall(PetscObjectStateIncrease((PetscObject)A));
4229:   *a = NULL;
4230:   PetscFunctionReturn(PETSC_SUCCESS);
4231: }

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

4236:    Not Collective

4238:    Input Parameter:
4239: .   A - a `MATSEQAIJHIPSPARSE` matrix

4241:    Output Parameter:
4242: .   a - pointer to the device data

4244:    Level: developer

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

4249: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetArray()`, `MatSeqAIJHIPSPARSEGetArrayRead()`, `MatSeqAIJHIPSPARSERestoreArrayWrite()`
4250: @*/
4251: PetscErrorCode MatSeqAIJHIPSPARSEGetArrayWrite(Mat A, PetscScalar **a)
4252: {
4253:   Mat_SeqAIJHIPSPARSE *cusp = (Mat_SeqAIJHIPSPARSE *)A->spptr;
4254:   CsrMatrix           *csr;

4256:   PetscFunctionBegin;
4259:   PetscCheckTypeName(A, MATSEQAIJHIPSPARSE);
4260:   PetscCheck(cusp->format != MAT_HIPSPARSE_ELL && cusp->format != MAT_HIPSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
4261:   PetscCheck(cusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJHIPSPARSEMultStruct");
4262:   csr = (CsrMatrix *)cusp->mat->mat;
4263:   PetscCheck(csr->values, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing HIP memory");
4264:   *a             = csr->values->data().get();
4265:   A->offloadmask = PETSC_OFFLOAD_GPU;
4266:   PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(A, PETSC_FALSE));
4267:   PetscFunctionReturn(PETSC_SUCCESS);
4268: }

4270: /*@C
4271:    MatSeqAIJHIPSPARSERestoreArrayWrite - restore the write-only access array obtained from `MatSeqAIJHIPSPARSEGetArrayWrite()`

4273:    Not Collective

4275:    Input Parameters:
4276: +   A - a `MATSEQAIJHIPSPARSE` matrix
4277: -   a - pointer to the device data

4279:    Level: developer

4281: .seealso: [](ch_matrices), `Mat`, `MatSeqAIJHIPSPARSEGetArrayWrite()`
4282: @*/
4283: PetscErrorCode MatSeqAIJHIPSPARSERestoreArrayWrite(Mat A, PetscScalar **a)
4284: {
4285:   PetscFunctionBegin;
4288:   PetscCheckTypeName(A, MATSEQAIJHIPSPARSE);
4289:   PetscCall(MatSeqAIJInvalidateDiagonal(A));
4290:   PetscCall(PetscObjectStateIncrease((PetscObject)A));
4291:   *a = NULL;
4292:   PetscFunctionReturn(PETSC_SUCCESS);
4293: }

4295: struct IJCompare4 {
4296:   __host__ __device__ inline bool operator()(const thrust::tuple<int, int, PetscScalar, int> &t1, const thrust::tuple<int, int, PetscScalar, int> &t2)
4297:   {
4298:     if (t1.get<0>() < t2.get<0>()) return true;
4299:     if (t1.get<0>() == t2.get<0>()) return t1.get<1>() < t2.get<1>();
4300:     return false;
4301:   }
4302: };

4304: struct Shift {
4305:   int _shift;

4307:   Shift(int shift) : _shift(shift) { }
4308:   __host__ __device__ inline int operator()(const int &c) { return c + _shift; }
4309: };

4311: /* merges two SeqAIJHIPSPARSE matrices A, B by concatenating their rows. [A';B']' operation in matlab notation */
4312: PetscErrorCode MatSeqAIJHIPSPARSEMergeMats(Mat A, Mat B, MatReuse reuse, Mat *C)
4313: {
4314:   Mat_SeqAIJ                    *a = (Mat_SeqAIJ *)A->data, *b = (Mat_SeqAIJ *)B->data, *c;
4315:   Mat_SeqAIJHIPSPARSE           *Acusp = (Mat_SeqAIJHIPSPARSE *)A->spptr, *Bcusp = (Mat_SeqAIJHIPSPARSE *)B->spptr, *Ccusp;
4316:   Mat_SeqAIJHIPSPARSEMultStruct *Cmat;
4317:   CsrMatrix                     *Acsr, *Bcsr, *Ccsr;
4318:   PetscInt                       Annz, Bnnz;
4319:   PetscInt                       i, m, n, zero = 0;

4321:   PetscFunctionBegin;
4325:   PetscCheckTypeName(A, MATSEQAIJHIPSPARSE);
4326:   PetscCheckTypeName(B, MATSEQAIJHIPSPARSE);
4327:   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);
4328:   PetscCheck(reuse != MAT_INPLACE_MATRIX, PETSC_COMM_SELF, PETSC_ERR_SUP, "MAT_INPLACE_MATRIX not supported");
4329:   PetscCheck(Acusp->format != MAT_HIPSPARSE_ELL && Acusp->format != MAT_HIPSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
4330:   PetscCheck(Bcusp->format != MAT_HIPSPARSE_ELL && Bcusp->format != MAT_HIPSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
4331:   if (reuse == MAT_INITIAL_MATRIX) {
4332:     m = A->rmap->n;
4333:     n = A->cmap->n + B->cmap->n;
4334:     PetscCall(MatCreate(PETSC_COMM_SELF, C));
4335:     PetscCall(MatSetSizes(*C, m, n, m, n));
4336:     PetscCall(MatSetType(*C, MATSEQAIJHIPSPARSE));
4337:     c                       = (Mat_SeqAIJ *)(*C)->data;
4338:     Ccusp                   = (Mat_SeqAIJHIPSPARSE *)(*C)->spptr;
4339:     Cmat                    = new Mat_SeqAIJHIPSPARSEMultStruct;
4340:     Ccsr                    = new CsrMatrix;
4341:     Cmat->cprowIndices      = NULL;
4342:     c->compressedrow.use    = PETSC_FALSE;
4343:     c->compressedrow.nrows  = 0;
4344:     c->compressedrow.i      = NULL;
4345:     c->compressedrow.rindex = NULL;
4346:     Ccusp->workVector       = NULL;
4347:     Ccusp->nrows            = m;
4348:     Ccusp->mat              = Cmat;
4349:     Ccusp->mat->mat         = Ccsr;
4350:     Ccsr->num_rows          = m;
4351:     Ccsr->num_cols          = n;
4352:     PetscCallHIPSPARSE(hipsparseCreateMatDescr(&Cmat->descr));
4353:     PetscCallHIPSPARSE(hipsparseSetMatIndexBase(Cmat->descr, HIPSPARSE_INDEX_BASE_ZERO));
4354:     PetscCallHIPSPARSE(hipsparseSetMatType(Cmat->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));
4355:     PetscCallHIP(hipMalloc((void **)&(Cmat->alpha_one), sizeof(PetscScalar)));
4356:     PetscCallHIP(hipMalloc((void **)&(Cmat->beta_zero), sizeof(PetscScalar)));
4357:     PetscCallHIP(hipMalloc((void **)&(Cmat->beta_one), sizeof(PetscScalar)));
4358:     PetscCallHIP(hipMemcpy(Cmat->alpha_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
4359:     PetscCallHIP(hipMemcpy(Cmat->beta_zero, &PETSC_HIPSPARSE_ZERO, sizeof(PetscScalar), hipMemcpyHostToDevice));
4360:     PetscCallHIP(hipMemcpy(Cmat->beta_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
4361:     PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
4362:     PetscCall(MatSeqAIJHIPSPARSECopyToGPU(B));
4363:     PetscCheck(Acusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJHIPSPARSEMultStruct");
4364:     PetscCheck(Bcusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJHIPSPARSEMultStruct");

4366:     Acsr                 = (CsrMatrix *)Acusp->mat->mat;
4367:     Bcsr                 = (CsrMatrix *)Bcusp->mat->mat;
4368:     Annz                 = (PetscInt)Acsr->column_indices->size();
4369:     Bnnz                 = (PetscInt)Bcsr->column_indices->size();
4370:     c->nz                = Annz + Bnnz;
4371:     Ccsr->row_offsets    = new THRUSTINTARRAY32(m + 1);
4372:     Ccsr->column_indices = new THRUSTINTARRAY32(c->nz);
4373:     Ccsr->values         = new THRUSTARRAY(c->nz);
4374:     Ccsr->num_entries    = c->nz;
4375:     Ccusp->cooPerm       = new THRUSTINTARRAY(c->nz);
4376:     if (c->nz) {
4377:       auto              Acoo = new THRUSTINTARRAY32(Annz);
4378:       auto              Bcoo = new THRUSTINTARRAY32(Bnnz);
4379:       auto              Ccoo = new THRUSTINTARRAY32(c->nz);
4380:       THRUSTINTARRAY32 *Aroff, *Broff;

4382:       if (a->compressedrow.use) { /* need full row offset */
4383:         if (!Acusp->rowoffsets_gpu) {
4384:           Acusp->rowoffsets_gpu = new THRUSTINTARRAY32(A->rmap->n + 1);
4385:           Acusp->rowoffsets_gpu->assign(a->i, a->i + A->rmap->n + 1);
4386:           PetscCall(PetscLogCpuToGpu((A->rmap->n + 1) * sizeof(PetscInt)));
4387:         }
4388:         Aroff = Acusp->rowoffsets_gpu;
4389:       } else Aroff = Acsr->row_offsets;
4390:       if (b->compressedrow.use) { /* need full row offset */
4391:         if (!Bcusp->rowoffsets_gpu) {
4392:           Bcusp->rowoffsets_gpu = new THRUSTINTARRAY32(B->rmap->n + 1);
4393:           Bcusp->rowoffsets_gpu->assign(b->i, b->i + B->rmap->n + 1);
4394:           PetscCall(PetscLogCpuToGpu((B->rmap->n + 1) * sizeof(PetscInt)));
4395:         }
4396:         Broff = Bcusp->rowoffsets_gpu;
4397:       } else Broff = Bcsr->row_offsets;
4398:       PetscCall(PetscLogGpuTimeBegin());
4399:       PetscCallHIPSPARSE(hipsparseXcsr2coo(Acusp->handle, Aroff->data().get(), Annz, m, Acoo->data().get(), HIPSPARSE_INDEX_BASE_ZERO));
4400:       PetscCallHIPSPARSE(hipsparseXcsr2coo(Bcusp->handle, Broff->data().get(), Bnnz, m, Bcoo->data().get(), HIPSPARSE_INDEX_BASE_ZERO));
4401:       /* Issues when using bool with large matrices on SUMMIT 10.2.89 */
4402:       auto Aperm = thrust::make_constant_iterator(1);
4403:       auto Bperm = thrust::make_constant_iterator(0);
4404:       auto Bcib  = thrust::make_transform_iterator(Bcsr->column_indices->begin(), Shift(A->cmap->n));
4405:       auto Bcie  = thrust::make_transform_iterator(Bcsr->column_indices->end(), Shift(A->cmap->n));
4406:       auto wPerm = new THRUSTINTARRAY32(Annz + Bnnz);
4407:       auto Azb   = thrust::make_zip_iterator(thrust::make_tuple(Acoo->begin(), Acsr->column_indices->begin(), Acsr->values->begin(), Aperm));
4408:       auto Aze   = thrust::make_zip_iterator(thrust::make_tuple(Acoo->end(), Acsr->column_indices->end(), Acsr->values->end(), Aperm));
4409:       auto Bzb   = thrust::make_zip_iterator(thrust::make_tuple(Bcoo->begin(), Bcib, Bcsr->values->begin(), Bperm));
4410:       auto Bze   = thrust::make_zip_iterator(thrust::make_tuple(Bcoo->end(), Bcie, Bcsr->values->end(), Bperm));
4411:       auto Czb   = thrust::make_zip_iterator(thrust::make_tuple(Ccoo->begin(), Ccsr->column_indices->begin(), Ccsr->values->begin(), wPerm->begin()));
4412:       auto p1    = Ccusp->cooPerm->begin();
4413:       auto p2    = Ccusp->cooPerm->begin();
4414:       thrust::advance(p2, Annz);
4415:       PetscCallThrust(thrust::merge(thrust::device, Azb, Aze, Bzb, Bze, Czb, IJCompare4()));
4416:       auto cci = thrust::make_counting_iterator(zero);
4417:       auto cce = thrust::make_counting_iterator(c->nz);
4418: #if 0 //Errors on SUMMIT cuda 11.1.0
4419:       PetscCallThrust(thrust::partition_copy(thrust::device, cci, cce, wPerm->begin(), p1, p2, thrust::identity<int>()));
4420: #else
4421:       auto pred = thrust::identity<int>();
4422:       PetscCallThrust(thrust::copy_if(thrust::device, cci, cce, wPerm->begin(), p1, pred));
4423:       PetscCallThrust(thrust::remove_copy_if(thrust::device, cci, cce, wPerm->begin(), p2, pred));
4424: #endif
4425:       PetscCallHIPSPARSE(hipsparseXcoo2csr(Ccusp->handle, Ccoo->data().get(), c->nz, m, Ccsr->row_offsets->data().get(), HIPSPARSE_INDEX_BASE_ZERO));
4426:       PetscCall(PetscLogGpuTimeEnd());
4427:       delete wPerm;
4428:       delete Acoo;
4429:       delete Bcoo;
4430:       delete Ccoo;
4431:       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(), HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));

4433:       if (A->form_explicit_transpose && B->form_explicit_transpose) { /* if A and B have the transpose, generate C transpose too */
4434:         PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(A));
4435:         PetscCall(MatSeqAIJHIPSPARSEFormExplicitTranspose(B));
4436:         PetscBool                      AT = Acusp->matTranspose ? PETSC_TRUE : PETSC_FALSE, BT = Bcusp->matTranspose ? PETSC_TRUE : PETSC_FALSE;
4437:         Mat_SeqAIJHIPSPARSEMultStruct *CmatT = new Mat_SeqAIJHIPSPARSEMultStruct;
4438:         CsrMatrix                     *CcsrT = new CsrMatrix;
4439:         CsrMatrix                     *AcsrT = AT ? (CsrMatrix *)Acusp->matTranspose->mat : NULL;
4440:         CsrMatrix                     *BcsrT = BT ? (CsrMatrix *)Bcusp->matTranspose->mat : NULL;

4442:         (*C)->form_explicit_transpose = PETSC_TRUE;
4443:         (*C)->transupdated            = PETSC_TRUE;
4444:         Ccusp->rowoffsets_gpu         = NULL;
4445:         CmatT->cprowIndices           = NULL;
4446:         CmatT->mat                    = CcsrT;
4447:         CcsrT->num_rows               = n;
4448:         CcsrT->num_cols               = m;
4449:         CcsrT->num_entries            = c->nz;
4450:         CcsrT->row_offsets            = new THRUSTINTARRAY32(n + 1);
4451:         CcsrT->column_indices         = new THRUSTINTARRAY32(c->nz);
4452:         CcsrT->values                 = new THRUSTARRAY(c->nz);

4454:         PetscCall(PetscLogGpuTimeBegin());
4455:         auto rT = CcsrT->row_offsets->begin();
4456:         if (AT) {
4457:           rT = thrust::copy(AcsrT->row_offsets->begin(), AcsrT->row_offsets->end(), rT);
4458:           thrust::advance(rT, -1);
4459:         }
4460:         if (BT) {
4461:           auto titb = thrust::make_transform_iterator(BcsrT->row_offsets->begin(), Shift(a->nz));
4462:           auto tite = thrust::make_transform_iterator(BcsrT->row_offsets->end(), Shift(a->nz));
4463:           thrust::copy(titb, tite, rT);
4464:         }
4465:         auto cT = CcsrT->column_indices->begin();
4466:         if (AT) cT = thrust::copy(AcsrT->column_indices->begin(), AcsrT->column_indices->end(), cT);
4467:         if (BT) thrust::copy(BcsrT->column_indices->begin(), BcsrT->column_indices->end(), cT);
4468:         auto vT = CcsrT->values->begin();
4469:         if (AT) vT = thrust::copy(AcsrT->values->begin(), AcsrT->values->end(), vT);
4470:         if (BT) thrust::copy(BcsrT->values->begin(), BcsrT->values->end(), vT);
4471:         PetscCall(PetscLogGpuTimeEnd());

4473:         PetscCallHIPSPARSE(hipsparseCreateMatDescr(&CmatT->descr));
4474:         PetscCallHIPSPARSE(hipsparseSetMatIndexBase(CmatT->descr, HIPSPARSE_INDEX_BASE_ZERO));
4475:         PetscCallHIPSPARSE(hipsparseSetMatType(CmatT->descr, HIPSPARSE_MATRIX_TYPE_GENERAL));
4476:         PetscCallHIP(hipMalloc((void **)&(CmatT->alpha_one), sizeof(PetscScalar)));
4477:         PetscCallHIP(hipMalloc((void **)&(CmatT->beta_zero), sizeof(PetscScalar)));
4478:         PetscCallHIP(hipMalloc((void **)&(CmatT->beta_one), sizeof(PetscScalar)));
4479:         PetscCallHIP(hipMemcpy(CmatT->alpha_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));
4480:         PetscCallHIP(hipMemcpy(CmatT->beta_zero, &PETSC_HIPSPARSE_ZERO, sizeof(PetscScalar), hipMemcpyHostToDevice));
4481:         PetscCallHIP(hipMemcpy(CmatT->beta_one, &PETSC_HIPSPARSE_ONE, sizeof(PetscScalar), hipMemcpyHostToDevice));

4483:         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(), HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_32I, HIPSPARSE_INDEX_BASE_ZERO, hipsparse_scalartype));
4484:         Ccusp->matTranspose = CmatT;
4485:       }
4486:     }

4488:     c->singlemalloc = PETSC_FALSE;
4489:     c->free_a       = PETSC_TRUE;
4490:     c->free_ij      = PETSC_TRUE;
4491:     PetscCall(PetscMalloc1(m + 1, &c->i));
4492:     PetscCall(PetscMalloc1(c->nz, &c->j));
4493:     if (PetscDefined(USE_64BIT_INDICES)) { /* 32 to 64-bit conversion on the GPU and then copy to host (lazy) */
4494:       THRUSTINTARRAY ii(Ccsr->row_offsets->size());
4495:       THRUSTINTARRAY jj(Ccsr->column_indices->size());
4496:       ii = *Ccsr->row_offsets;
4497:       jj = *Ccsr->column_indices;
4498:       PetscCallHIP(hipMemcpy(c->i, ii.data().get(), Ccsr->row_offsets->size() * sizeof(PetscInt), hipMemcpyDeviceToHost));
4499:       PetscCallHIP(hipMemcpy(c->j, jj.data().get(), Ccsr->column_indices->size() * sizeof(PetscInt), hipMemcpyDeviceToHost));
4500:     } else {
4501:       PetscCallHIP(hipMemcpy(c->i, Ccsr->row_offsets->data().get(), Ccsr->row_offsets->size() * sizeof(PetscInt), hipMemcpyDeviceToHost));
4502:       PetscCallHIP(hipMemcpy(c->j, Ccsr->column_indices->data().get(), Ccsr->column_indices->size() * sizeof(PetscInt), hipMemcpyDeviceToHost));
4503:     }
4504:     PetscCall(PetscLogGpuToCpu((Ccsr->column_indices->size() + Ccsr->row_offsets->size()) * sizeof(PetscInt)));
4505:     PetscCall(PetscMalloc1(m, &c->ilen));
4506:     PetscCall(PetscMalloc1(m, &c->imax));
4507:     c->maxnz         = c->nz;
4508:     c->nonzerorowcnt = 0;
4509:     c->rmax          = 0;
4510:     for (i = 0; i < m; i++) {
4511:       const PetscInt nn = c->i[i + 1] - c->i[i];
4512:       c->ilen[i] = c->imax[i] = nn;
4513:       c->nonzerorowcnt += (PetscInt) !!nn;
4514:       c->rmax = PetscMax(c->rmax, nn);
4515:     }
4516:     PetscCall(MatMarkDiagonal_SeqAIJ(*C));
4517:     PetscCall(PetscMalloc1(c->nz, &c->a));
4518:     (*C)->nonzerostate++;
4519:     PetscCall(PetscLayoutSetUp((*C)->rmap));
4520:     PetscCall(PetscLayoutSetUp((*C)->cmap));
4521:     Ccusp->nonzerostate = (*C)->nonzerostate;
4522:     (*C)->preallocated  = PETSC_TRUE;
4523:   } else {
4524:     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);
4525:     c = (Mat_SeqAIJ *)(*C)->data;
4526:     if (c->nz) {
4527:       Ccusp = (Mat_SeqAIJHIPSPARSE *)(*C)->spptr;
4528:       PetscCheck(Ccusp->cooPerm, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing cooPerm");
4529:       PetscCheck(Ccusp->format != MAT_HIPSPARSE_ELL && Ccusp->format != MAT_HIPSPARSE_HYB, PETSC_COMM_SELF, PETSC_ERR_SUP, "Not implemented");
4530:       PetscCheck(Ccusp->nonzerostate == (*C)->nonzerostate, PETSC_COMM_SELF, PETSC_ERR_COR, "Wrong nonzerostate");
4531:       PetscCall(MatSeqAIJHIPSPARSECopyToGPU(A));
4532:       PetscCall(MatSeqAIJHIPSPARSECopyToGPU(B));
4533:       PetscCheck(Acusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJHIPSPARSEMultStruct");
4534:       PetscCheck(Bcusp->mat, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing Mat_SeqAIJHIPSPARSEMultStruct");
4535:       Acsr = (CsrMatrix *)Acusp->mat->mat;
4536:       Bcsr = (CsrMatrix *)Bcusp->mat->mat;
4537:       Ccsr = (CsrMatrix *)Ccusp->mat->mat;
4538:       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());
4539:       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());
4540:       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());
4541:       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);
4542:       PetscCheck(Ccusp->cooPerm->size() == Ccsr->values->size(), PETSC_COMM_SELF, PETSC_ERR_COR, "permSize %" PetscInt_FMT " != %" PetscInt_FMT, (PetscInt)Ccusp->cooPerm->size(), (PetscInt)Ccsr->values->size());
4543:       auto pmid = Ccusp->cooPerm->begin();
4544:       thrust::advance(pmid, Acsr->num_entries);
4545:       PetscCall(PetscLogGpuTimeBegin());
4546:       auto zibait = thrust::make_zip_iterator(thrust::make_tuple(Acsr->values->begin(), thrust::make_permutation_iterator(Ccsr->values->begin(), Ccusp->cooPerm->begin())));
4547:       auto zieait = thrust::make_zip_iterator(thrust::make_tuple(Acsr->values->end(), thrust::make_permutation_iterator(Ccsr->values->begin(), pmid)));
4548:       thrust::for_each(zibait, zieait, VecHIPEquals());
4549:       auto zibbit = thrust::make_zip_iterator(thrust::make_tuple(Bcsr->values->begin(), thrust::make_permutation_iterator(Ccsr->values->begin(), pmid)));
4550:       auto ziebit = thrust::make_zip_iterator(thrust::make_tuple(Bcsr->values->end(), thrust::make_permutation_iterator(Ccsr->values->begin(), Ccusp->cooPerm->end())));
4551:       thrust::for_each(zibbit, ziebit, VecHIPEquals());
4552:       PetscCall(MatSeqAIJHIPSPARSEInvalidateTranspose(*C, PETSC_FALSE));
4553:       if (A->form_explicit_transpose && B->form_explicit_transpose && (*C)->form_explicit_transpose) {
4554:         PetscCheck(Ccusp->matTranspose, PETSC_COMM_SELF, PETSC_ERR_COR, "Missing transpose Mat_SeqAIJHIPSPARSEMultStruct");
4555:         PetscBool  AT = Acusp->matTranspose ? PETSC_TRUE : PETSC_FALSE, BT = Bcusp->matTranspose ? PETSC_TRUE : PETSC_FALSE;
4556:         CsrMatrix *AcsrT = AT ? (CsrMatrix *)Acusp->matTranspose->mat : NULL;
4557:         CsrMatrix *BcsrT = BT ? (CsrMatrix *)Bcusp->matTranspose->mat : NULL;
4558:         CsrMatrix *CcsrT = (CsrMatrix *)Ccusp->matTranspose->mat;
4559:         auto       vT    = CcsrT->values->begin();
4560:         if (AT) vT = thrust::copy(AcsrT->values->begin(), AcsrT->values->end(), vT);
4561:         if (BT) thrust::copy(BcsrT->values->begin(), BcsrT->values->end(), vT);
4562:         (*C)->transupdated = PETSC_TRUE;
4563:       }
4564:       PetscCall(PetscLogGpuTimeEnd());
4565:     }
4566:   }
4567:   PetscCall(PetscObjectStateIncrease((PetscObject)*C));
4568:   (*C)->assembled     = PETSC_TRUE;
4569:   (*C)->was_assembled = PETSC_FALSE;
4570:   (*C)->offloadmask   = PETSC_OFFLOAD_GPU;
4571:   PetscFunctionReturn(PETSC_SUCCESS);
4572: }

4574: static PetscErrorCode MatSeqAIJCopySubArray_SeqAIJHIPSPARSE(Mat A, PetscInt n, const PetscInt idx[], PetscScalar v[])
4575: {
4576:   bool               dmem;
4577:   const PetscScalar *av;

4579:   PetscFunctionBegin;
4580:   dmem = isHipMem(v);
4581:   PetscCall(MatSeqAIJHIPSPARSEGetArrayRead(A, &av));
4582:   if (n && idx) {
4583:     THRUSTINTARRAY widx(n);
4584:     widx.assign(idx, idx + n);
4585:     PetscCall(PetscLogCpuToGpu(n * sizeof(PetscInt)));

4587:     THRUSTARRAY                    *w = NULL;
4588:     thrust::device_ptr<PetscScalar> dv;
4589:     if (dmem) dv = thrust::device_pointer_cast(v);
4590:     else {
4591:       w  = new THRUSTARRAY(n);
4592:       dv = w->data();
4593:     }
4594:     thrust::device_ptr<const PetscScalar> dav = thrust::device_pointer_cast(av);

4596:     auto zibit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(dav, widx.begin()), dv));
4597:     auto zieit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(dav, widx.end()), dv + n));
4598:     thrust::for_each(zibit, zieit, VecHIPEquals());
4599:     if (w) PetscCallHIP(hipMemcpy(v, w->data().get(), n * sizeof(PetscScalar), hipMemcpyDeviceToHost));
4600:     delete w;
4601:   } else PetscCallHIP(hipMemcpy(v, av, n * sizeof(PetscScalar), dmem ? hipMemcpyDeviceToDevice : hipMemcpyDeviceToHost));

4603:   if (!dmem) PetscCall(PetscLogCpuToGpu(n * sizeof(PetscScalar)));
4604:   PetscCall(MatSeqAIJHIPSPARSERestoreArrayRead(A, &av));
4605:   PetscFunctionReturn(PETSC_SUCCESS);
4606: }