xref: /petsc/src/mat/impls/aij/seq/seqcusparse/aijcusparse.cu (revision bfcc362789b558bcb4d7977f28a2c9d2d6233b89)
19ae82921SPaul Mullowney /*
29ae82921SPaul Mullowney   Defines the basic matrix operations for the AIJ (compressed row)
3fd7c363cSSatish Balay   matrix storage format using the CUSPARSE library,
49ae82921SPaul Mullowney */
5dced61a5SBarry Smith #define PETSC_SKIP_SPINLOCK
653800007SKarl Rupp #define PETSC_SKIP_CXX_COMPLEX_FIX
799acd6aaSStefano Zampini #define PETSC_SKIP_IMMINTRIN_H_CUDAWORKAROUND 1
89ae82921SPaul Mullowney 
93d13b8fdSMatthew G. Knepley #include <petscconf.h>
103d13b8fdSMatthew G. Knepley #include <../src/mat/impls/aij/seq/aij.h>          /*I "petscmat.h" I*/
11087f3262SPaul Mullowney #include <../src/mat/impls/sbaij/seq/sbaij.h>
123d13b8fdSMatthew G. Knepley #include <../src/vec/vec/impls/dvecimpl.h>
13af0996ceSBarry Smith #include <petsc/private/vecimpl.h>
149ae82921SPaul Mullowney #undef VecType
153d13b8fdSMatthew G. Knepley #include <../src/mat/impls/aij/seq/seqcusparse/cusparsematimpl.h>
16bc3f50f2SPaul Mullowney 
17e057df02SPaul Mullowney const char *const MatCUSPARSEStorageFormats[]    = {"CSR","ELL","HYB","MatCUSPARSEStorageFormat","MAT_CUSPARSE_",0};
18afb2bd1cSJunchao Zhang #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
19afb2bd1cSJunchao Zhang   /* The following are copied from cusparse.h in CUDA-11.0. In MatCUSPARSESpMVAlgorithms[] etc, we copy them in
20afb2bd1cSJunchao Zhang     0-based integer value order, since we want to use PetscOptionsEnum() to parse user command line options for them.
21afb2bd1cSJunchao Zhang 
22afb2bd1cSJunchao Zhang   typedef enum {
23afb2bd1cSJunchao Zhang       CUSPARSE_MV_ALG_DEFAULT = 0,
24afb2bd1cSJunchao Zhang       CUSPARSE_COOMV_ALG      = 1,
25afb2bd1cSJunchao Zhang       CUSPARSE_CSRMV_ALG1     = 2,
26afb2bd1cSJunchao Zhang       CUSPARSE_CSRMV_ALG2     = 3
27afb2bd1cSJunchao Zhang   } cusparseSpMVAlg_t;
28afb2bd1cSJunchao Zhang 
29afb2bd1cSJunchao Zhang   typedef enum {
30afb2bd1cSJunchao Zhang       CUSPARSE_MM_ALG_DEFAULT     CUSPARSE_DEPRECATED_ENUM(CUSPARSE_SPMM_ALG_DEFAULT) = 0,
31afb2bd1cSJunchao Zhang       CUSPARSE_COOMM_ALG1         CUSPARSE_DEPRECATED_ENUM(CUSPARSE_SPMM_COO_ALG1)    = 1,
32afb2bd1cSJunchao Zhang       CUSPARSE_COOMM_ALG2         CUSPARSE_DEPRECATED_ENUM(CUSPARSE_SPMM_COO_ALG2)    = 2,
33afb2bd1cSJunchao Zhang       CUSPARSE_COOMM_ALG3         CUSPARSE_DEPRECATED_ENUM(CUSPARSE_SPMM_COO_ALG3)    = 3,
34afb2bd1cSJunchao Zhang       CUSPARSE_CSRMM_ALG1         CUSPARSE_DEPRECATED_ENUM(CUSPARSE_SPMM_CSR_ALG1)    = 4,
35afb2bd1cSJunchao Zhang       CUSPARSE_SPMM_ALG_DEFAULT = 0,
36afb2bd1cSJunchao Zhang       CUSPARSE_SPMM_COO_ALG1    = 1,
37afb2bd1cSJunchao Zhang       CUSPARSE_SPMM_COO_ALG2    = 2,
38afb2bd1cSJunchao Zhang       CUSPARSE_SPMM_COO_ALG3    = 3,
39afb2bd1cSJunchao Zhang       CUSPARSE_SPMM_COO_ALG4    = 5,
40afb2bd1cSJunchao Zhang       CUSPARSE_SPMM_CSR_ALG1    = 4,
41afb2bd1cSJunchao Zhang       CUSPARSE_SPMM_CSR_ALG2    = 6,
42afb2bd1cSJunchao Zhang   } cusparseSpMMAlg_t;
43afb2bd1cSJunchao Zhang 
44afb2bd1cSJunchao Zhang   typedef enum {
45afb2bd1cSJunchao Zhang       CUSPARSE_CSR2CSC_ALG1 = 1, // faster than V2 (in general), deterministc
46afb2bd1cSJunchao Zhang       CUSPARSE_CSR2CSC_ALG2 = 2  // low memory requirement, non-deterministc
47afb2bd1cSJunchao Zhang   } cusparseCsr2CscAlg_t;
48afb2bd1cSJunchao Zhang   */
49afb2bd1cSJunchao Zhang   const char *const MatCUSPARSESpMVAlgorithms[]    = {"MV_ALG_DEFAULT","COOMV_ALG", "CSRMV_ALG1","CSRMV_ALG2", "cusparseSpMVAlg_t","CUSPARSE_",0};
50afb2bd1cSJunchao Zhang   const char *const MatCUSPARSESpMMAlgorithms[]    = {"ALG_DEFAULT","COO_ALG1","COO_ALG2","COO_ALG3","CSR_ALG1","COO_ALG4","CSR_ALG2","cusparseSpMMAlg_t","CUSPARSE_SPMM_",0};
51afb2bd1cSJunchao Zhang   const char *const MatCUSPARSECsr2CscAlgorithms[] = {"INVALID"/*cusparse does not have enum 0! We created one*/,"ALG1","ALG2","cusparseCsr2CscAlg_t","CUSPARSE_CSR2CSC_",0};
52afb2bd1cSJunchao Zhang #endif
539ae82921SPaul Mullowney 
54087f3262SPaul Mullowney static PetscErrorCode MatICCFactorSymbolic_SeqAIJCUSPARSE(Mat,Mat,IS,const MatFactorInfo*);
55087f3262SPaul Mullowney static PetscErrorCode MatCholeskyFactorSymbolic_SeqAIJCUSPARSE(Mat,Mat,IS,const MatFactorInfo*);
56087f3262SPaul Mullowney static PetscErrorCode MatCholeskyFactorNumeric_SeqAIJCUSPARSE(Mat,Mat,const MatFactorInfo*);
57087f3262SPaul Mullowney 
586fa9248bSJed Brown static PetscErrorCode MatILUFactorSymbolic_SeqAIJCUSPARSE(Mat,Mat,IS,IS,const MatFactorInfo*);
596fa9248bSJed Brown static PetscErrorCode MatLUFactorSymbolic_SeqAIJCUSPARSE(Mat,Mat,IS,IS,const MatFactorInfo*);
606fa9248bSJed Brown static PetscErrorCode MatLUFactorNumeric_SeqAIJCUSPARSE(Mat,Mat,const MatFactorInfo*);
61087f3262SPaul Mullowney 
626fa9248bSJed Brown static PetscErrorCode MatSolve_SeqAIJCUSPARSE(Mat,Vec,Vec);
636fa9248bSJed Brown static PetscErrorCode MatSolve_SeqAIJCUSPARSE_NaturalOrdering(Mat,Vec,Vec);
646fa9248bSJed Brown static PetscErrorCode MatSolveTranspose_SeqAIJCUSPARSE(Mat,Vec,Vec);
656fa9248bSJed Brown static PetscErrorCode MatSolveTranspose_SeqAIJCUSPARSE_NaturalOrdering(Mat,Vec,Vec);
664416b707SBarry Smith static PetscErrorCode MatSetFromOptions_SeqAIJCUSPARSE(PetscOptionItems *PetscOptionsObject,Mat);
67a587d139SMark static PetscErrorCode MatAXPY_SeqAIJCUSPARSE(Mat,PetscScalar,Mat,MatStructure);
686fa9248bSJed Brown static PetscErrorCode MatMult_SeqAIJCUSPARSE(Mat,Vec,Vec);
696fa9248bSJed Brown static PetscErrorCode MatMultAdd_SeqAIJCUSPARSE(Mat,Vec,Vec,Vec);
706fa9248bSJed Brown static PetscErrorCode MatMultTranspose_SeqAIJCUSPARSE(Mat,Vec,Vec);
716fa9248bSJed Brown static PetscErrorCode MatMultTransposeAdd_SeqAIJCUSPARSE(Mat,Vec,Vec,Vec);
72e6e9a74fSStefano Zampini static PetscErrorCode MatMultHermitianTranspose_SeqAIJCUSPARSE(Mat,Vec,Vec);
73e6e9a74fSStefano Zampini static PetscErrorCode MatMultHermitianTransposeAdd_SeqAIJCUSPARSE(Mat,Vec,Vec,Vec);
74e6e9a74fSStefano Zampini static PetscErrorCode MatMultAddKernel_SeqAIJCUSPARSE(Mat,Vec,Vec,Vec,PetscBool,PetscBool);
759ae82921SPaul Mullowney 
767f756511SDominic Meiser static PetscErrorCode CsrMatrix_Destroy(CsrMatrix**);
77470880abSPatrick Sanan static PetscErrorCode MatSeqAIJCUSPARSEMultStruct_Destroy(Mat_SeqAIJCUSPARSETriFactorStruct**);
78470880abSPatrick Sanan static PetscErrorCode MatSeqAIJCUSPARSEMultStruct_Destroy(Mat_SeqAIJCUSPARSEMultStruct**,MatCUSPARSEStorageFormat);
79ccdfe979SStefano Zampini static PetscErrorCode MatSeqAIJCUSPARSETriFactors_Reset(Mat_SeqAIJCUSPARSETriFactors**);
80470880abSPatrick Sanan static PetscErrorCode MatSeqAIJCUSPARSETriFactors_Destroy(Mat_SeqAIJCUSPARSETriFactors**);
81470880abSPatrick Sanan static PetscErrorCode MatSeqAIJCUSPARSE_Destroy(Mat_SeqAIJCUSPARSE**);
827f756511SDominic Meiser 
8357181aedSStefano Zampini static PetscErrorCode MatSeqAIJCUSPARSECopyToGPU(Mat);
8457181aedSStefano Zampini static PetscErrorCode MatSeqAIJCUSPARSECopyFromGPU(Mat);
8557181aedSStefano Zampini 
867e8381f9SStefano Zampini PETSC_INTERN PetscErrorCode MatSetPreallocationCOO_SeqAIJCUSPARSE(Mat,PetscInt,const PetscInt[],const PetscInt[]);
877e8381f9SStefano Zampini PETSC_INTERN PetscErrorCode MatSetValuesCOO_SeqAIJCUSPARSE(Mat,const PetscScalar[],InsertMode);
887e8381f9SStefano Zampini 
89b06137fdSPaul Mullowney PetscErrorCode MatCUSPARSESetStream(Mat A,const cudaStream_t stream)
90b06137fdSPaul Mullowney {
91b06137fdSPaul Mullowney   cusparseStatus_t   stat;
92b06137fdSPaul Mullowney   Mat_SeqAIJCUSPARSE *cusparsestruct = (Mat_SeqAIJCUSPARSE*)A->spptr;
93b06137fdSPaul Mullowney 
94b06137fdSPaul Mullowney   PetscFunctionBegin;
95d98d7c49SStefano Zampini   if (!cusparsestruct) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing spptr");
96b06137fdSPaul Mullowney   cusparsestruct->stream = stream;
9757d48284SJunchao Zhang   stat = cusparseSetStream(cusparsestruct->handle,cusparsestruct->stream);CHKERRCUSPARSE(stat);
98b06137fdSPaul Mullowney   PetscFunctionReturn(0);
99b06137fdSPaul Mullowney }
100b06137fdSPaul Mullowney 
101b06137fdSPaul Mullowney PetscErrorCode MatCUSPARSESetHandle(Mat A,const cusparseHandle_t handle)
102b06137fdSPaul Mullowney {
103b06137fdSPaul Mullowney   cusparseStatus_t   stat;
104b06137fdSPaul Mullowney   Mat_SeqAIJCUSPARSE *cusparsestruct = (Mat_SeqAIJCUSPARSE*)A->spptr;
105b06137fdSPaul Mullowney 
106b06137fdSPaul Mullowney   PetscFunctionBegin;
107d98d7c49SStefano Zampini   if (!cusparsestruct) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing spptr");
1086b1cf21dSAlejandro Lamas Daviña   if (cusparsestruct->handle != handle) {
10916a2e217SAlejandro Lamas Daviña     if (cusparsestruct->handle) {
11057d48284SJunchao Zhang       stat = cusparseDestroy(cusparsestruct->handle);CHKERRCUSPARSE(stat);
11116a2e217SAlejandro Lamas Daviña     }
112b06137fdSPaul Mullowney     cusparsestruct->handle = handle;
1136b1cf21dSAlejandro Lamas Daviña   }
11457d48284SJunchao Zhang   stat = cusparseSetPointerMode(cusparsestruct->handle, CUSPARSE_POINTER_MODE_DEVICE);CHKERRCUSPARSE(stat);
115b06137fdSPaul Mullowney   PetscFunctionReturn(0);
116b06137fdSPaul Mullowney }
117b06137fdSPaul Mullowney 
118b06137fdSPaul Mullowney PetscErrorCode MatCUSPARSEClearHandle(Mat A)
119b06137fdSPaul Mullowney {
120b06137fdSPaul Mullowney   Mat_SeqAIJCUSPARSE *cusparsestruct = (Mat_SeqAIJCUSPARSE*)A->spptr;
1217e8381f9SStefano Zampini   PetscBool          flg;
1227e8381f9SStefano Zampini   PetscErrorCode     ierr;
123ccdfe979SStefano Zampini 
124b06137fdSPaul Mullowney   PetscFunctionBegin;
1257e8381f9SStefano Zampini   ierr = PetscObjectTypeCompare((PetscObject)A,MATSEQAIJCUSPARSE,&flg);CHKERRQ(ierr);
1267e8381f9SStefano Zampini   if (!flg || !cusparsestruct) PetscFunctionReturn(0);
127ccdfe979SStefano Zampini   if (cusparsestruct->handle) cusparsestruct->handle = 0;
128b06137fdSPaul Mullowney   PetscFunctionReturn(0);
129b06137fdSPaul Mullowney }
130b06137fdSPaul Mullowney 
131ea799195SBarry Smith PetscErrorCode MatFactorGetSolverType_seqaij_cusparse(Mat A,MatSolverType *type)
1329ae82921SPaul Mullowney {
1339ae82921SPaul Mullowney   PetscFunctionBegin;
1349ae82921SPaul Mullowney   *type = MATSOLVERCUSPARSE;
1359ae82921SPaul Mullowney   PetscFunctionReturn(0);
1369ae82921SPaul Mullowney }
1379ae82921SPaul Mullowney 
138c708e6cdSJed Brown /*MC
139087f3262SPaul Mullowney   MATSOLVERCUSPARSE = "cusparse" - A matrix type providing triangular solvers for seq matrices
140087f3262SPaul Mullowney   on a single GPU of type, seqaijcusparse, aijcusparse, or seqaijcusp, aijcusp. Currently supported
141087f3262SPaul Mullowney   algorithms are ILU(k) and ICC(k). Typically, deeper factorizations (larger k) results in poorer
142087f3262SPaul Mullowney   performance in the triangular solves. Full LU, and Cholesky decompositions can be solved through the
143087f3262SPaul Mullowney   CUSPARSE triangular solve algorithm. However, the performance can be quite poor and thus these
144087f3262SPaul Mullowney   algorithms are not recommended. This class does NOT support direct solver operations.
145c708e6cdSJed Brown 
1469ae82921SPaul Mullowney   Level: beginner
147c708e6cdSJed Brown 
1483ca39a21SBarry Smith .seealso: PCFactorSetMatSolverType(), MatSolverType, MatCreateSeqAIJCUSPARSE(), MATAIJCUSPARSE, MatCreateAIJCUSPARSE(), MatCUSPARSESetFormat(), MatCUSPARSEStorageFormat, MatCUSPARSEFormatOperation
149c708e6cdSJed Brown M*/
1509ae82921SPaul Mullowney 
15142c9c57cSBarry Smith PETSC_EXTERN PetscErrorCode MatGetFactor_seqaijcusparse_cusparse(Mat A,MatFactorType ftype,Mat *B)
1529ae82921SPaul Mullowney {
1539ae82921SPaul Mullowney   PetscErrorCode ierr;
154bc3f50f2SPaul Mullowney   PetscInt       n = A->rmap->n;
1559ae82921SPaul Mullowney 
1569ae82921SPaul Mullowney   PetscFunctionBegin;
157bc3f50f2SPaul Mullowney   ierr = MatCreate(PetscObjectComm((PetscObject)A),B);CHKERRQ(ierr);
158bc3f50f2SPaul Mullowney   ierr = MatSetSizes(*B,n,n,n,n);CHKERRQ(ierr);
1592c7c0729SBarry Smith   (*B)->factortype = ftype;
1602c7c0729SBarry Smith   (*B)->useordering = PETSC_TRUE;
1619ae82921SPaul Mullowney   ierr = MatSetType(*B,MATSEQAIJCUSPARSE);CHKERRQ(ierr);
1622205254eSKarl Rupp 
163087f3262SPaul Mullowney   if (ftype == MAT_FACTOR_LU || ftype == MAT_FACTOR_ILU || ftype == MAT_FACTOR_ILUDT) {
16433d57670SJed Brown     ierr = MatSetBlockSizesFromMats(*B,A,A);CHKERRQ(ierr);
1659ae82921SPaul Mullowney     (*B)->ops->ilufactorsymbolic = MatILUFactorSymbolic_SeqAIJCUSPARSE;
1669ae82921SPaul Mullowney     (*B)->ops->lufactorsymbolic  = MatLUFactorSymbolic_SeqAIJCUSPARSE;
167087f3262SPaul Mullowney   } else if (ftype == MAT_FACTOR_CHOLESKY || ftype == MAT_FACTOR_ICC) {
168087f3262SPaul Mullowney     (*B)->ops->iccfactorsymbolic      = MatICCFactorSymbolic_SeqAIJCUSPARSE;
169087f3262SPaul Mullowney     (*B)->ops->choleskyfactorsymbolic = MatCholeskyFactorSymbolic_SeqAIJCUSPARSE;
1709ae82921SPaul Mullowney   } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Factor type not supported for CUSPARSE Matrix Types");
171bc3f50f2SPaul Mullowney 
172fa03d054SJed Brown   ierr = MatSeqAIJSetPreallocation(*B,MAT_SKIP_ALLOCATION,NULL);CHKERRQ(ierr);
1733ca39a21SBarry Smith   ierr = PetscObjectComposeFunction((PetscObject)(*B),"MatFactorGetSolverType_C",MatFactorGetSolverType_seqaij_cusparse);CHKERRQ(ierr);
1749ae82921SPaul Mullowney   PetscFunctionReturn(0);
1759ae82921SPaul Mullowney }
1769ae82921SPaul Mullowney 
177bc3f50f2SPaul Mullowney PETSC_INTERN PetscErrorCode MatCUSPARSESetFormat_SeqAIJCUSPARSE(Mat A,MatCUSPARSEFormatOperation op,MatCUSPARSEStorageFormat format)
178ca45077fSPaul Mullowney {
179aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSE *cusparsestruct = (Mat_SeqAIJCUSPARSE*)A->spptr;
1806e111a19SKarl Rupp 
181ca45077fSPaul Mullowney   PetscFunctionBegin;
182ca45077fSPaul Mullowney   switch (op) {
183e057df02SPaul Mullowney   case MAT_CUSPARSE_MULT:
184aa372e3fSPaul Mullowney     cusparsestruct->format = format;
185ca45077fSPaul Mullowney     break;
186e057df02SPaul Mullowney   case MAT_CUSPARSE_ALL:
187aa372e3fSPaul Mullowney     cusparsestruct->format = format;
188ca45077fSPaul Mullowney     break;
189ca45077fSPaul Mullowney   default:
19036d62e41SPaul Mullowney     SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_SUP,"unsupported operation %d for MatCUSPARSEFormatOperation. MAT_CUSPARSE_MULT and MAT_CUSPARSE_ALL are currently supported.",op);
191ca45077fSPaul Mullowney   }
192ca45077fSPaul Mullowney   PetscFunctionReturn(0);
193ca45077fSPaul Mullowney }
1949ae82921SPaul Mullowney 
195e057df02SPaul Mullowney /*@
196e057df02SPaul Mullowney    MatCUSPARSESetFormat - Sets the storage format of CUSPARSE matrices for a particular
197e057df02SPaul Mullowney    operation. Only the MatMult operation can use different GPU storage formats
198aa372e3fSPaul Mullowney    for MPIAIJCUSPARSE matrices.
199e057df02SPaul Mullowney    Not Collective
200e057df02SPaul Mullowney 
201e057df02SPaul Mullowney    Input Parameters:
2028468deeeSKarl Rupp +  A - Matrix of type SEQAIJCUSPARSE
20336d62e41SPaul Mullowney .  op - MatCUSPARSEFormatOperation. SEQAIJCUSPARSE matrices support MAT_CUSPARSE_MULT and MAT_CUSPARSE_ALL. MPIAIJCUSPARSE matrices support MAT_CUSPARSE_MULT_DIAG, MAT_CUSPARSE_MULT_OFFDIAG, and MAT_CUSPARSE_ALL.
2042692e278SPaul Mullowney -  format - MatCUSPARSEStorageFormat (one of MAT_CUSPARSE_CSR, MAT_CUSPARSE_ELL, MAT_CUSPARSE_HYB. The latter two require CUDA 4.2)
205e057df02SPaul Mullowney 
206e057df02SPaul Mullowney    Output Parameter:
207e057df02SPaul Mullowney 
208e057df02SPaul Mullowney    Level: intermediate
209e057df02SPaul Mullowney 
2108468deeeSKarl Rupp .seealso: MatCUSPARSEStorageFormat, MatCUSPARSEFormatOperation
211e057df02SPaul Mullowney @*/
212e057df02SPaul Mullowney PetscErrorCode MatCUSPARSESetFormat(Mat A,MatCUSPARSEFormatOperation op,MatCUSPARSEStorageFormat format)
213e057df02SPaul Mullowney {
214e057df02SPaul Mullowney   PetscErrorCode ierr;
2156e111a19SKarl Rupp 
216e057df02SPaul Mullowney   PetscFunctionBegin;
217e057df02SPaul Mullowney   PetscValidHeaderSpecific(A, MAT_CLASSID,1);
218e057df02SPaul Mullowney   ierr = PetscTryMethod(A,"MatCUSPARSESetFormat_C",(Mat,MatCUSPARSEFormatOperation,MatCUSPARSEStorageFormat),(A,op,format));CHKERRQ(ierr);
219e057df02SPaul Mullowney   PetscFunctionReturn(0);
220e057df02SPaul Mullowney }
221e057df02SPaul Mullowney 
222e6e9a74fSStefano Zampini /*@
223e589036eSStefano Zampini    MatSeqAIJCUSPARSESetGenerateTranspose - Sets the flag to explicitly generate the transpose matrix before calling MatMultTranspose
224e6e9a74fSStefano Zampini 
225e6e9a74fSStefano Zampini    Collective on mat
226e6e9a74fSStefano Zampini 
227e6e9a74fSStefano Zampini    Input Parameters:
228e6e9a74fSStefano Zampini +  A - Matrix of type SEQAIJCUSPARSE
229e6e9a74fSStefano Zampini -  transgen - the boolean flag
230e6e9a74fSStefano Zampini 
231e6e9a74fSStefano Zampini    Level: intermediate
232e6e9a74fSStefano Zampini 
233e589036eSStefano Zampini .seealso: MATSEQAIJCUSPARSE, MatAIJCUSPARSESetGenerateTranspose()
234e6e9a74fSStefano Zampini @*/
235e6e9a74fSStefano Zampini PetscErrorCode MatSeqAIJCUSPARSESetGenerateTranspose(Mat A,PetscBool transgen)
236e6e9a74fSStefano Zampini {
237e6e9a74fSStefano Zampini   PetscErrorCode ierr;
238e6e9a74fSStefano Zampini   PetscBool      flg;
239e6e9a74fSStefano Zampini 
240e6e9a74fSStefano Zampini   PetscFunctionBegin;
241e6e9a74fSStefano Zampini   PetscValidHeaderSpecific(A,MAT_CLASSID,1);
242e6e9a74fSStefano Zampini   ierr = PetscObjectTypeCompare(((PetscObject)A),MATSEQAIJCUSPARSE,&flg);CHKERRQ(ierr);
243e6e9a74fSStefano Zampini   if (flg) {
244e6e9a74fSStefano Zampini     Mat_SeqAIJCUSPARSE *cusp = (Mat_SeqAIJCUSPARSE*)A->spptr;
24554da937aSStefano Zampini 
246e6e9a74fSStefano Zampini     if (A->factortype) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_ARG_WRONGSTATE,"Not for factored matrix");
247e6e9a74fSStefano Zampini     cusp->transgen = transgen;
24854da937aSStefano Zampini     if (!transgen) { /* need to destroy the transpose matrix if present to prevent from logic errors if transgen is set to true later */
24954da937aSStefano Zampini       ierr = MatSeqAIJCUSPARSEMultStruct_Destroy(&cusp->matTranspose,cusp->format);CHKERRQ(ierr);
25054da937aSStefano Zampini     }
251e6e9a74fSStefano Zampini   }
252e6e9a74fSStefano Zampini   PetscFunctionReturn(0);
253e6e9a74fSStefano Zampini }
254e6e9a74fSStefano Zampini 
2554416b707SBarry Smith static PetscErrorCode MatSetFromOptions_SeqAIJCUSPARSE(PetscOptionItems *PetscOptionsObject,Mat A)
2569ae82921SPaul Mullowney {
2579ae82921SPaul Mullowney   PetscErrorCode           ierr;
258e057df02SPaul Mullowney   MatCUSPARSEStorageFormat format;
2599ae82921SPaul Mullowney   PetscBool                flg;
260a183c035SDominic Meiser   Mat_SeqAIJCUSPARSE       *cusparsestruct = (Mat_SeqAIJCUSPARSE*)A->spptr;
2616e111a19SKarl Rupp 
2629ae82921SPaul Mullowney   PetscFunctionBegin;
263e55864a3SBarry Smith   ierr = PetscOptionsHead(PetscOptionsObject,"SeqAIJCUSPARSE options");CHKERRQ(ierr);
2649ae82921SPaul Mullowney   if (A->factortype == MAT_FACTOR_NONE) {
26554da937aSStefano Zampini     PetscBool transgen = cusparsestruct->transgen;
26654da937aSStefano Zampini 
26754da937aSStefano Zampini     ierr = PetscOptionsBool("-mat_cusparse_transgen","Generate explicit transpose for MatMultTranspose","MatSeqAIJCUSPARSESetGenerateTranspose",transgen,&transgen,&flg);CHKERRQ(ierr);
268afb2bd1cSJunchao Zhang     if (flg) {ierr = MatSeqAIJCUSPARSESetGenerateTranspose(A,transgen);CHKERRQ(ierr);}
269afb2bd1cSJunchao Zhang 
270e057df02SPaul Mullowney     ierr = PetscOptionsEnum("-mat_cusparse_mult_storage_format","sets storage format of (seq)aijcusparse gpu matrices for SpMV",
271a183c035SDominic Meiser                             "MatCUSPARSESetFormat",MatCUSPARSEStorageFormats,(PetscEnum)cusparsestruct->format,(PetscEnum*)&format,&flg);CHKERRQ(ierr);
272afb2bd1cSJunchao Zhang     if (flg) {ierr = MatCUSPARSESetFormat(A,MAT_CUSPARSE_MULT,format);CHKERRQ(ierr);}
273afb2bd1cSJunchao Zhang 
2744c87dfd4SPaul Mullowney     ierr = PetscOptionsEnum("-mat_cusparse_storage_format","sets storage format of (seq)aijcusparse gpu matrices for SpMV and TriSolve",
275a183c035SDominic Meiser                             "MatCUSPARSESetFormat",MatCUSPARSEStorageFormats,(PetscEnum)cusparsestruct->format,(PetscEnum*)&format,&flg);CHKERRQ(ierr);
276afb2bd1cSJunchao Zhang     if (flg) {ierr = MatCUSPARSESetFormat(A,MAT_CUSPARSE_ALL,format);CHKERRQ(ierr);}
277afb2bd1cSJunchao Zhang    #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
278afb2bd1cSJunchao Zhang     cusparsestruct->spmvAlg = CUSPARSE_CSRMV_ALG1; /* default, since we only support csr */
279afb2bd1cSJunchao Zhang     ierr = PetscOptionsEnum("-mat_cusparse_spmv_alg","sets cuSPARSE algorithm used in sparse-mat dense-vector multiplication (SpMV)",
280afb2bd1cSJunchao Zhang                             "cusparseSpMVAlg_t",MatCUSPARSESpMVAlgorithms,(PetscEnum)cusparsestruct->spmvAlg,(PetscEnum*)&cusparsestruct->spmvAlg,&flg);CHKERRQ(ierr);
281afb2bd1cSJunchao Zhang     /* If user did use this option, check its consistency with cuSPARSE, since PetscOptionsEnum() sets enum values based on their position in MatCUSPARSESpMVAlgorithms[] */
282afb2bd1cSJunchao Zhang     if (flg && CUSPARSE_CSRMV_ALG1 != 2) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"cuSPARSE enum cusparseSpMVAlg_t has been changed but PETSc has not been updated accordingly");
283afb2bd1cSJunchao Zhang 
284afb2bd1cSJunchao Zhang     cusparsestruct->spmmAlg = CUSPARSE_SPMM_CSR_ALG1; /* default, only support column-major dense matrix B */
285afb2bd1cSJunchao Zhang     ierr = PetscOptionsEnum("-mat_cusparse_spmm_alg","sets cuSPARSE algorithm used in sparse-mat dense-mat multiplication (SpMM)",
286afb2bd1cSJunchao Zhang                             "cusparseSpMMAlg_t",MatCUSPARSESpMMAlgorithms,(PetscEnum)cusparsestruct->spmmAlg,(PetscEnum*)&cusparsestruct->spmmAlg,&flg);CHKERRQ(ierr);
287afb2bd1cSJunchao Zhang     if (flg && CUSPARSE_SPMM_CSR_ALG1 != 4) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"cuSPARSE enum cusparseSpMMAlg_t has been changed but PETSc has not been updated accordingly");
288afb2bd1cSJunchao Zhang 
289afb2bd1cSJunchao Zhang     cusparsestruct->csr2cscAlg = CUSPARSE_CSR2CSC_ALG1;
290afb2bd1cSJunchao Zhang     ierr = PetscOptionsEnum("-mat_cusparse_csr2csc_alg","sets cuSPARSE algorithm used in converting CSR matrices to CSC matrices",
291afb2bd1cSJunchao Zhang                             "cusparseCsr2CscAlg_t",MatCUSPARSECsr2CscAlgorithms,(PetscEnum)cusparsestruct->csr2cscAlg,(PetscEnum*)&cusparsestruct->csr2cscAlg,&flg);CHKERRQ(ierr);
292afb2bd1cSJunchao Zhang     if (flg && CUSPARSE_CSR2CSC_ALG1 != 1) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"cuSPARSE enum cusparseCsr2CscAlg_t has been changed but PETSc has not been updated accordingly");
293afb2bd1cSJunchao Zhang    #endif
2944c87dfd4SPaul Mullowney   }
2950af67c1bSStefano Zampini   ierr = PetscOptionsTail();CHKERRQ(ierr);
2969ae82921SPaul Mullowney   PetscFunctionReturn(0);
2979ae82921SPaul Mullowney }
2989ae82921SPaul Mullowney 
2996fa9248bSJed Brown static PetscErrorCode MatILUFactorSymbolic_SeqAIJCUSPARSE(Mat B,Mat A,IS isrow,IS iscol,const MatFactorInfo *info)
3009ae82921SPaul Mullowney {
301da79fbbcSStefano Zampini   Mat_SeqAIJCUSPARSETriFactors *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)B->spptr;
3029ae82921SPaul Mullowney   PetscErrorCode               ierr;
3039ae82921SPaul Mullowney 
3049ae82921SPaul Mullowney   PetscFunctionBegin;
305da79fbbcSStefano Zampini   ierr = MatSeqAIJCUSPARSETriFactors_Reset(&cusparseTriFactors);CHKERRQ(ierr);
3069ae82921SPaul Mullowney   ierr = MatILUFactorSymbolic_SeqAIJ(B,A,isrow,iscol,info);CHKERRQ(ierr);
3079ae82921SPaul Mullowney   B->ops->lufactornumeric = MatLUFactorNumeric_SeqAIJCUSPARSE;
3089ae82921SPaul Mullowney   PetscFunctionReturn(0);
3099ae82921SPaul Mullowney }
3109ae82921SPaul Mullowney 
3116fa9248bSJed Brown static PetscErrorCode MatLUFactorSymbolic_SeqAIJCUSPARSE(Mat B,Mat A,IS isrow,IS iscol,const MatFactorInfo *info)
3129ae82921SPaul Mullowney {
313da79fbbcSStefano Zampini   Mat_SeqAIJCUSPARSETriFactors *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)B->spptr;
3149ae82921SPaul Mullowney   PetscErrorCode               ierr;
3159ae82921SPaul Mullowney 
3169ae82921SPaul Mullowney   PetscFunctionBegin;
317da79fbbcSStefano Zampini   ierr = MatSeqAIJCUSPARSETriFactors_Reset(&cusparseTriFactors);CHKERRQ(ierr);
3189ae82921SPaul Mullowney   ierr = MatLUFactorSymbolic_SeqAIJ(B,A,isrow,iscol,info);CHKERRQ(ierr);
3199ae82921SPaul Mullowney   B->ops->lufactornumeric = MatLUFactorNumeric_SeqAIJCUSPARSE;
3209ae82921SPaul Mullowney   PetscFunctionReturn(0);
3219ae82921SPaul Mullowney }
3229ae82921SPaul Mullowney 
323087f3262SPaul Mullowney static PetscErrorCode MatICCFactorSymbolic_SeqAIJCUSPARSE(Mat B,Mat A,IS perm,const MatFactorInfo *info)
324087f3262SPaul Mullowney {
325da79fbbcSStefano Zampini   Mat_SeqAIJCUSPARSETriFactors *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)B->spptr;
326087f3262SPaul Mullowney   PetscErrorCode               ierr;
327087f3262SPaul Mullowney 
328087f3262SPaul Mullowney   PetscFunctionBegin;
329da79fbbcSStefano Zampini   ierr = MatSeqAIJCUSPARSETriFactors_Reset(&cusparseTriFactors);CHKERRQ(ierr);
330087f3262SPaul Mullowney   ierr = MatICCFactorSymbolic_SeqAIJ(B,A,perm,info);CHKERRQ(ierr);
331087f3262SPaul Mullowney   B->ops->choleskyfactornumeric = MatCholeskyFactorNumeric_SeqAIJCUSPARSE;
332087f3262SPaul Mullowney   PetscFunctionReturn(0);
333087f3262SPaul Mullowney }
334087f3262SPaul Mullowney 
335087f3262SPaul Mullowney static PetscErrorCode MatCholeskyFactorSymbolic_SeqAIJCUSPARSE(Mat B,Mat A,IS perm,const MatFactorInfo *info)
336087f3262SPaul Mullowney {
337da79fbbcSStefano Zampini   Mat_SeqAIJCUSPARSETriFactors *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)B->spptr;
338087f3262SPaul Mullowney   PetscErrorCode               ierr;
339087f3262SPaul Mullowney 
340087f3262SPaul Mullowney   PetscFunctionBegin;
341da79fbbcSStefano Zampini   ierr = MatSeqAIJCUSPARSETriFactors_Reset(&cusparseTriFactors);CHKERRQ(ierr);
342087f3262SPaul Mullowney   ierr = MatCholeskyFactorSymbolic_SeqAIJ(B,A,perm,info);CHKERRQ(ierr);
343087f3262SPaul Mullowney   B->ops->choleskyfactornumeric = MatCholeskyFactorNumeric_SeqAIJCUSPARSE;
344087f3262SPaul Mullowney   PetscFunctionReturn(0);
345087f3262SPaul Mullowney }
346087f3262SPaul Mullowney 
347087f3262SPaul Mullowney static PetscErrorCode MatSeqAIJCUSPARSEBuildILULowerTriMatrix(Mat A)
3489ae82921SPaul Mullowney {
3499ae82921SPaul Mullowney   Mat_SeqAIJ                        *a = (Mat_SeqAIJ*)A->data;
3509ae82921SPaul Mullowney   PetscInt                          n = A->rmap->n;
3519ae82921SPaul Mullowney   Mat_SeqAIJCUSPARSETriFactors      *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)A->spptr;
352aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct *loTriFactor = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->loTriFactorPtr;
3539ae82921SPaul Mullowney   cusparseStatus_t                  stat;
3549ae82921SPaul Mullowney   const PetscInt                    *ai = a->i,*aj = a->j,*vi;
3559ae82921SPaul Mullowney   const MatScalar                   *aa = a->a,*v;
3569ae82921SPaul Mullowney   PetscInt                          *AiLo, *AjLo;
3579ae82921SPaul Mullowney   PetscInt                          i,nz, nzLower, offset, rowOffset;
358b175d8bbSPaul Mullowney   PetscErrorCode                    ierr;
35957d48284SJunchao Zhang   cudaError_t                       cerr;
3609ae82921SPaul Mullowney 
3619ae82921SPaul Mullowney   PetscFunctionBegin;
362cf00fe3bSKarl Rupp   if (!n) PetscFunctionReturn(0);
363c70f7ee4SJunchao Zhang   if (A->offloadmask == PETSC_OFFLOAD_UNALLOCATED || A->offloadmask == PETSC_OFFLOAD_CPU) {
3649ae82921SPaul Mullowney     try {
3659ae82921SPaul Mullowney       /* first figure out the number of nonzeros in the lower triangular matrix including 1's on the diagonal. */
3669ae82921SPaul Mullowney       nzLower=n+ai[n]-ai[1];
367da79fbbcSStefano Zampini       if (!loTriFactor) {
3682cbc15d9SMark         PetscScalar                       *AALo;
3692cbc15d9SMark 
3702cbc15d9SMark         cerr = cudaMallocHost((void**) &AALo, nzLower*sizeof(PetscScalar));CHKERRCUDA(cerr);
3719ae82921SPaul Mullowney 
3729ae82921SPaul Mullowney         /* Allocate Space for the lower triangular matrix */
37357d48284SJunchao Zhang         cerr = cudaMallocHost((void**) &AiLo, (n+1)*sizeof(PetscInt));CHKERRCUDA(cerr);
37457d48284SJunchao Zhang         cerr = cudaMallocHost((void**) &AjLo, nzLower*sizeof(PetscInt));CHKERRCUDA(cerr);
3759ae82921SPaul Mullowney 
3769ae82921SPaul Mullowney         /* Fill the lower triangular matrix */
3779ae82921SPaul Mullowney         AiLo[0]  = (PetscInt) 0;
3789ae82921SPaul Mullowney         AiLo[n]  = nzLower;
3799ae82921SPaul Mullowney         AjLo[0]  = (PetscInt) 0;
3809ae82921SPaul Mullowney         AALo[0]  = (MatScalar) 1.0;
3819ae82921SPaul Mullowney         v        = aa;
3829ae82921SPaul Mullowney         vi       = aj;
3839ae82921SPaul Mullowney         offset   = 1;
3849ae82921SPaul Mullowney         rowOffset= 1;
3859ae82921SPaul Mullowney         for (i=1; i<n; i++) {
3869ae82921SPaul Mullowney           nz = ai[i+1] - ai[i];
387e057df02SPaul Mullowney           /* additional 1 for the term on the diagonal */
3889ae82921SPaul Mullowney           AiLo[i]    = rowOffset;
3899ae82921SPaul Mullowney           rowOffset += nz+1;
3909ae82921SPaul Mullowney 
391580bdb30SBarry Smith           ierr = PetscArraycpy(&(AjLo[offset]), vi, nz);CHKERRQ(ierr);
392580bdb30SBarry Smith           ierr = PetscArraycpy(&(AALo[offset]), v, nz);CHKERRQ(ierr);
3939ae82921SPaul Mullowney 
3949ae82921SPaul Mullowney           offset      += nz;
3959ae82921SPaul Mullowney           AjLo[offset] = (PetscInt) i;
3969ae82921SPaul Mullowney           AALo[offset] = (MatScalar) 1.0;
3979ae82921SPaul Mullowney           offset      += 1;
3989ae82921SPaul Mullowney 
3999ae82921SPaul Mullowney           v  += nz;
4009ae82921SPaul Mullowney           vi += nz;
4019ae82921SPaul Mullowney         }
4022205254eSKarl Rupp 
403aa372e3fSPaul Mullowney         /* allocate space for the triangular factor information */
404da79fbbcSStefano Zampini         ierr = PetscNew(&loTriFactor);CHKERRQ(ierr);
405da79fbbcSStefano Zampini         loTriFactor->solvePolicy = CUSPARSE_SOLVE_POLICY_USE_LEVEL;
406aa372e3fSPaul Mullowney         /* Create the matrix description */
40757d48284SJunchao Zhang         stat = cusparseCreateMatDescr(&loTriFactor->descr);CHKERRCUSPARSE(stat);
40857d48284SJunchao Zhang         stat = cusparseSetMatIndexBase(loTriFactor->descr, CUSPARSE_INDEX_BASE_ZERO);CHKERRCUSPARSE(stat);
4091b0a6780SStefano Zampini        #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
410afb2bd1cSJunchao Zhang         stat = cusparseSetMatType(loTriFactor->descr, CUSPARSE_MATRIX_TYPE_GENERAL);CHKERRCUSPARSE(stat);
411afb2bd1cSJunchao Zhang        #else
41257d48284SJunchao Zhang         stat = cusparseSetMatType(loTriFactor->descr, CUSPARSE_MATRIX_TYPE_TRIANGULAR);CHKERRCUSPARSE(stat);
413afb2bd1cSJunchao Zhang        #endif
41457d48284SJunchao Zhang         stat = cusparseSetMatFillMode(loTriFactor->descr, CUSPARSE_FILL_MODE_LOWER);CHKERRCUSPARSE(stat);
41557d48284SJunchao Zhang         stat = cusparseSetMatDiagType(loTriFactor->descr, CUSPARSE_DIAG_TYPE_UNIT);CHKERRCUSPARSE(stat);
416aa372e3fSPaul Mullowney 
417aa372e3fSPaul Mullowney         /* set the operation */
418aa372e3fSPaul Mullowney         loTriFactor->solveOp = CUSPARSE_OPERATION_NON_TRANSPOSE;
419aa372e3fSPaul Mullowney 
420aa372e3fSPaul Mullowney         /* set the matrix */
421aa372e3fSPaul Mullowney         loTriFactor->csrMat = new CsrMatrix;
422aa372e3fSPaul Mullowney         loTriFactor->csrMat->num_rows = n;
423aa372e3fSPaul Mullowney         loTriFactor->csrMat->num_cols = n;
424aa372e3fSPaul Mullowney         loTriFactor->csrMat->num_entries = nzLower;
425aa372e3fSPaul Mullowney 
426aa372e3fSPaul Mullowney         loTriFactor->csrMat->row_offsets = new THRUSTINTARRAY32(n+1);
427aa372e3fSPaul Mullowney         loTriFactor->csrMat->row_offsets->assign(AiLo, AiLo+n+1);
428aa372e3fSPaul Mullowney 
429aa372e3fSPaul Mullowney         loTriFactor->csrMat->column_indices = new THRUSTINTARRAY32(nzLower);
430aa372e3fSPaul Mullowney         loTriFactor->csrMat->column_indices->assign(AjLo, AjLo+nzLower);
431aa372e3fSPaul Mullowney 
432aa372e3fSPaul Mullowney         loTriFactor->csrMat->values = new THRUSTARRAY(nzLower);
433aa372e3fSPaul Mullowney         loTriFactor->csrMat->values->assign(AALo, AALo+nzLower);
434aa372e3fSPaul Mullowney 
435afb2bd1cSJunchao Zhang         /* Create the solve analysis information */
436da79fbbcSStefano Zampini         ierr = PetscLogEventBegin(MAT_CUSPARSESolveAnalysis,A,0,0,0);CHKERRQ(ierr);
437afb2bd1cSJunchao Zhang         stat = cusparse_create_analysis_info(&loTriFactor->solveInfo);CHKERRCUSPARSE(stat);
4381b0a6780SStefano Zampini       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
439afb2bd1cSJunchao Zhang         stat = cusparse_get_svbuffsize(cusparseTriFactors->handle, loTriFactor->solveOp,
440afb2bd1cSJunchao Zhang                                        loTriFactor->csrMat->num_rows, loTriFactor->csrMat->num_entries, loTriFactor->descr,
441afb2bd1cSJunchao Zhang                                        loTriFactor->csrMat->values->data().get(), loTriFactor->csrMat->row_offsets->data().get(),
442afb2bd1cSJunchao Zhang                                        loTriFactor->csrMat->column_indices->data().get(), loTriFactor->solveInfo,
443afb2bd1cSJunchao Zhang                                        &loTriFactor->solveBufferSize);CHKERRCUSPARSE(stat);
444afb2bd1cSJunchao Zhang         cerr = cudaMalloc(&loTriFactor->solveBuffer,loTriFactor->solveBufferSize);CHKERRCUDA(cerr);
445afb2bd1cSJunchao Zhang       #endif
446afb2bd1cSJunchao Zhang 
447aa372e3fSPaul Mullowney         /* perform the solve analysis */
448aa372e3fSPaul Mullowney         stat = cusparse_analysis(cusparseTriFactors->handle, loTriFactor->solveOp,
449aa372e3fSPaul Mullowney                                  loTriFactor->csrMat->num_rows, loTriFactor->csrMat->num_entries, loTriFactor->descr,
450aa372e3fSPaul Mullowney                                  loTriFactor->csrMat->values->data().get(), loTriFactor->csrMat->row_offsets->data().get(),
451afb2bd1cSJunchao Zhang                                  loTriFactor->csrMat->column_indices->data().get(), loTriFactor->solveInfo
4521b0a6780SStefano Zampini                                #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
453afb2bd1cSJunchao Zhang                                  ,loTriFactor->solvePolicy, loTriFactor->solveBuffer
454afb2bd1cSJunchao Zhang                                #endif
455afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
456da79fbbcSStefano Zampini         cerr = WaitForCUDA();CHKERRCUDA(cerr);
457da79fbbcSStefano Zampini         ierr = PetscLogEventEnd(MAT_CUSPARSESolveAnalysis,A,0,0,0);CHKERRQ(ierr);
458aa372e3fSPaul Mullowney 
459da79fbbcSStefano Zampini         /* assign the pointer */
460aa372e3fSPaul Mullowney         ((Mat_SeqAIJCUSPARSETriFactors*)A->spptr)->loTriFactorPtr = loTriFactor;
4612cbc15d9SMark         loTriFactor->AA_h = AALo;
46257d48284SJunchao Zhang         cerr = cudaFreeHost(AiLo);CHKERRCUDA(cerr);
46357d48284SJunchao Zhang         cerr = cudaFreeHost(AjLo);CHKERRCUDA(cerr);
4644863603aSSatish Balay         ierr = PetscLogCpuToGpu((n+1+nzLower)*sizeof(int)+nzLower*sizeof(PetscScalar));CHKERRQ(ierr);
465da79fbbcSStefano Zampini       } else { /* update values only */
4662cbc15d9SMark         if (!loTriFactor->AA_h) {
4672cbc15d9SMark           cerr = cudaMallocHost((void**) &loTriFactor->AA_h, nzLower*sizeof(PetscScalar));CHKERRCUDA(cerr);
4682cbc15d9SMark         }
469da79fbbcSStefano Zampini         /* Fill the lower triangular matrix */
4702cbc15d9SMark         loTriFactor->AA_h[0]  = 1.0;
471da79fbbcSStefano Zampini         v        = aa;
472da79fbbcSStefano Zampini         vi       = aj;
473da79fbbcSStefano Zampini         offset   = 1;
474da79fbbcSStefano Zampini         for (i=1; i<n; i++) {
475da79fbbcSStefano Zampini           nz = ai[i+1] - ai[i];
4762cbc15d9SMark           ierr = PetscArraycpy(&(loTriFactor->AA_h[offset]), v, nz);CHKERRQ(ierr);
477da79fbbcSStefano Zampini           offset      += nz;
4782cbc15d9SMark           loTriFactor->AA_h[offset] = 1.0;
479da79fbbcSStefano Zampini           offset      += 1;
480da79fbbcSStefano Zampini           v  += nz;
481da79fbbcSStefano Zampini         }
4822cbc15d9SMark         loTriFactor->csrMat->values->assign(loTriFactor->AA_h, loTriFactor->AA_h+nzLower);
483da79fbbcSStefano Zampini         ierr = PetscLogCpuToGpu(nzLower*sizeof(PetscScalar));CHKERRQ(ierr);
484da79fbbcSStefano Zampini       }
4859ae82921SPaul Mullowney     } catch(char *ex) {
4869ae82921SPaul Mullowney       SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_LIB,"CUSPARSE error: %s", ex);
4879ae82921SPaul Mullowney     }
4889ae82921SPaul Mullowney   }
4899ae82921SPaul Mullowney   PetscFunctionReturn(0);
4909ae82921SPaul Mullowney }
4919ae82921SPaul Mullowney 
492087f3262SPaul Mullowney static PetscErrorCode MatSeqAIJCUSPARSEBuildILUUpperTriMatrix(Mat A)
4939ae82921SPaul Mullowney {
4949ae82921SPaul Mullowney   Mat_SeqAIJ                        *a = (Mat_SeqAIJ*)A->data;
4959ae82921SPaul Mullowney   PetscInt                          n = A->rmap->n;
4969ae82921SPaul Mullowney   Mat_SeqAIJCUSPARSETriFactors      *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)A->spptr;
497aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct *upTriFactor = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->upTriFactorPtr;
4989ae82921SPaul Mullowney   cusparseStatus_t                  stat;
4999ae82921SPaul Mullowney   const PetscInt                    *aj = a->j,*adiag = a->diag,*vi;
5009ae82921SPaul Mullowney   const MatScalar                   *aa = a->a,*v;
5019ae82921SPaul Mullowney   PetscInt                          *AiUp, *AjUp;
5029ae82921SPaul Mullowney   PetscInt                          i,nz, nzUpper, offset;
5039ae82921SPaul Mullowney   PetscErrorCode                    ierr;
50457d48284SJunchao Zhang   cudaError_t                       cerr;
5059ae82921SPaul Mullowney 
5069ae82921SPaul Mullowney   PetscFunctionBegin;
507cf00fe3bSKarl Rupp   if (!n) PetscFunctionReturn(0);
508c70f7ee4SJunchao Zhang   if (A->offloadmask == PETSC_OFFLOAD_UNALLOCATED || A->offloadmask == PETSC_OFFLOAD_CPU) {
5099ae82921SPaul Mullowney     try {
5109ae82921SPaul Mullowney       /* next, figure out the number of nonzeros in the upper triangular matrix. */
5119ae82921SPaul Mullowney       nzUpper = adiag[0]-adiag[n];
512da79fbbcSStefano Zampini       if (!upTriFactor) {
5132cbc15d9SMark         PetscScalar *AAUp;
5142cbc15d9SMark 
5152cbc15d9SMark         cerr = cudaMallocHost((void**) &AAUp, nzUpper*sizeof(PetscScalar));CHKERRCUDA(cerr);
5162cbc15d9SMark 
5179ae82921SPaul Mullowney         /* Allocate Space for the upper triangular matrix */
51857d48284SJunchao Zhang         cerr = cudaMallocHost((void**) &AiUp, (n+1)*sizeof(PetscInt));CHKERRCUDA(cerr);
51957d48284SJunchao Zhang         cerr = cudaMallocHost((void**) &AjUp, nzUpper*sizeof(PetscInt));CHKERRCUDA(cerr);
5209ae82921SPaul Mullowney 
5219ae82921SPaul Mullowney         /* Fill the upper triangular matrix */
5229ae82921SPaul Mullowney         AiUp[0]=(PetscInt) 0;
5239ae82921SPaul Mullowney         AiUp[n]=nzUpper;
5249ae82921SPaul Mullowney         offset = nzUpper;
5259ae82921SPaul Mullowney         for (i=n-1; i>=0; i--) {
5269ae82921SPaul Mullowney           v  = aa + adiag[i+1] + 1;
5279ae82921SPaul Mullowney           vi = aj + adiag[i+1] + 1;
5289ae82921SPaul Mullowney 
529e057df02SPaul Mullowney           /* number of elements NOT on the diagonal */
5309ae82921SPaul Mullowney           nz = adiag[i] - adiag[i+1]-1;
5319ae82921SPaul Mullowney 
532e057df02SPaul Mullowney           /* decrement the offset */
5339ae82921SPaul Mullowney           offset -= (nz+1);
5349ae82921SPaul Mullowney 
535e057df02SPaul Mullowney           /* first, set the diagonal elements */
5369ae82921SPaul Mullowney           AjUp[offset] = (PetscInt) i;
53709f51544SAlejandro Lamas Daviña           AAUp[offset] = (MatScalar)1./v[nz];
5389ae82921SPaul Mullowney           AiUp[i]      = AiUp[i+1] - (nz+1);
5399ae82921SPaul Mullowney 
540580bdb30SBarry Smith           ierr = PetscArraycpy(&(AjUp[offset+1]), vi, nz);CHKERRQ(ierr);
541580bdb30SBarry Smith           ierr = PetscArraycpy(&(AAUp[offset+1]), v, nz);CHKERRQ(ierr);
5429ae82921SPaul Mullowney         }
5432205254eSKarl Rupp 
544aa372e3fSPaul Mullowney         /* allocate space for the triangular factor information */
545da79fbbcSStefano Zampini         ierr = PetscNew(&upTriFactor);CHKERRQ(ierr);
546da79fbbcSStefano Zampini         upTriFactor->solvePolicy = CUSPARSE_SOLVE_POLICY_USE_LEVEL;
5472205254eSKarl Rupp 
548aa372e3fSPaul Mullowney         /* Create the matrix description */
54957d48284SJunchao Zhang         stat = cusparseCreateMatDescr(&upTriFactor->descr);CHKERRCUSPARSE(stat);
55057d48284SJunchao Zhang         stat = cusparseSetMatIndexBase(upTriFactor->descr, CUSPARSE_INDEX_BASE_ZERO);CHKERRCUSPARSE(stat);
5511b0a6780SStefano Zampini        #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
552afb2bd1cSJunchao Zhang         stat = cusparseSetMatType(upTriFactor->descr, CUSPARSE_MATRIX_TYPE_GENERAL);CHKERRCUSPARSE(stat);
553afb2bd1cSJunchao Zhang        #else
55457d48284SJunchao Zhang         stat = cusparseSetMatType(upTriFactor->descr, CUSPARSE_MATRIX_TYPE_TRIANGULAR);CHKERRCUSPARSE(stat);
555afb2bd1cSJunchao Zhang        #endif
55657d48284SJunchao Zhang         stat = cusparseSetMatFillMode(upTriFactor->descr, CUSPARSE_FILL_MODE_UPPER);CHKERRCUSPARSE(stat);
55757d48284SJunchao Zhang         stat = cusparseSetMatDiagType(upTriFactor->descr, CUSPARSE_DIAG_TYPE_NON_UNIT);CHKERRCUSPARSE(stat);
558aa372e3fSPaul Mullowney 
559aa372e3fSPaul Mullowney         /* set the operation */
560aa372e3fSPaul Mullowney         upTriFactor->solveOp = CUSPARSE_OPERATION_NON_TRANSPOSE;
561aa372e3fSPaul Mullowney 
562aa372e3fSPaul Mullowney         /* set the matrix */
563aa372e3fSPaul Mullowney         upTriFactor->csrMat = new CsrMatrix;
564aa372e3fSPaul Mullowney         upTriFactor->csrMat->num_rows = n;
565aa372e3fSPaul Mullowney         upTriFactor->csrMat->num_cols = n;
566aa372e3fSPaul Mullowney         upTriFactor->csrMat->num_entries = nzUpper;
567aa372e3fSPaul Mullowney 
568aa372e3fSPaul Mullowney         upTriFactor->csrMat->row_offsets = new THRUSTINTARRAY32(n+1);
569aa372e3fSPaul Mullowney         upTriFactor->csrMat->row_offsets->assign(AiUp, AiUp+n+1);
570aa372e3fSPaul Mullowney 
571aa372e3fSPaul Mullowney         upTriFactor->csrMat->column_indices = new THRUSTINTARRAY32(nzUpper);
572aa372e3fSPaul Mullowney         upTriFactor->csrMat->column_indices->assign(AjUp, AjUp+nzUpper);
573aa372e3fSPaul Mullowney 
574aa372e3fSPaul Mullowney         upTriFactor->csrMat->values = new THRUSTARRAY(nzUpper);
575aa372e3fSPaul Mullowney         upTriFactor->csrMat->values->assign(AAUp, AAUp+nzUpper);
576aa372e3fSPaul Mullowney 
577afb2bd1cSJunchao Zhang         /* Create the solve analysis information */
578da79fbbcSStefano Zampini         ierr = PetscLogEventBegin(MAT_CUSPARSESolveAnalysis,A,0,0,0);CHKERRQ(ierr);
579afb2bd1cSJunchao Zhang         stat = cusparse_create_analysis_info(&upTriFactor->solveInfo);CHKERRCUSPARSE(stat);
5801b0a6780SStefano Zampini       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
581afb2bd1cSJunchao Zhang         stat = cusparse_get_svbuffsize(cusparseTriFactors->handle, upTriFactor->solveOp,
582afb2bd1cSJunchao Zhang                                      upTriFactor->csrMat->num_rows, upTriFactor->csrMat->num_entries, upTriFactor->descr,
583afb2bd1cSJunchao Zhang                                      upTriFactor->csrMat->values->data().get(), upTriFactor->csrMat->row_offsets->data().get(),
584afb2bd1cSJunchao Zhang                                      upTriFactor->csrMat->column_indices->data().get(), upTriFactor->solveInfo,
585afb2bd1cSJunchao Zhang                                      &upTriFactor->solveBufferSize);CHKERRCUSPARSE(stat);
586afb2bd1cSJunchao Zhang         cerr = cudaMalloc(&upTriFactor->solveBuffer,upTriFactor->solveBufferSize);CHKERRCUDA(cerr);
587afb2bd1cSJunchao Zhang       #endif
588afb2bd1cSJunchao Zhang 
589aa372e3fSPaul Mullowney         /* perform the solve analysis */
590aa372e3fSPaul Mullowney         stat = cusparse_analysis(cusparseTriFactors->handle, upTriFactor->solveOp,
591aa372e3fSPaul Mullowney                                  upTriFactor->csrMat->num_rows, upTriFactor->csrMat->num_entries, upTriFactor->descr,
592aa372e3fSPaul Mullowney                                  upTriFactor->csrMat->values->data().get(), upTriFactor->csrMat->row_offsets->data().get(),
593afb2bd1cSJunchao Zhang                                  upTriFactor->csrMat->column_indices->data().get(), upTriFactor->solveInfo
5941b0a6780SStefano Zampini                                #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
595afb2bd1cSJunchao Zhang                                  ,upTriFactor->solvePolicy, upTriFactor->solveBuffer
596afb2bd1cSJunchao Zhang                                #endif
597afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
598da79fbbcSStefano Zampini         cerr = WaitForCUDA();CHKERRCUDA(cerr);
599da79fbbcSStefano Zampini         ierr = PetscLogEventEnd(MAT_CUSPARSESolveAnalysis,A,0,0,0);CHKERRQ(ierr);
600aa372e3fSPaul Mullowney 
601da79fbbcSStefano Zampini         /* assign the pointer */
602aa372e3fSPaul Mullowney         ((Mat_SeqAIJCUSPARSETriFactors*)A->spptr)->upTriFactorPtr = upTriFactor;
6032cbc15d9SMark         upTriFactor->AA_h = AAUp;
60457d48284SJunchao Zhang         cerr = cudaFreeHost(AiUp);CHKERRCUDA(cerr);
60557d48284SJunchao Zhang         cerr = cudaFreeHost(AjUp);CHKERRCUDA(cerr);
6064863603aSSatish Balay         ierr = PetscLogCpuToGpu((n+1+nzUpper)*sizeof(int)+nzUpper*sizeof(PetscScalar));CHKERRQ(ierr);
607da79fbbcSStefano Zampini       } else {
6082cbc15d9SMark         if (!upTriFactor->AA_h) {
6092cbc15d9SMark           cerr = cudaMallocHost((void**) &upTriFactor->AA_h, nzUpper*sizeof(PetscScalar));CHKERRCUDA(cerr);
6102cbc15d9SMark         }
611da79fbbcSStefano Zampini         /* Fill the upper triangular matrix */
612da79fbbcSStefano Zampini         offset = nzUpper;
613da79fbbcSStefano Zampini         for (i=n-1; i>=0; i--) {
614da79fbbcSStefano Zampini           v  = aa + adiag[i+1] + 1;
615da79fbbcSStefano Zampini 
616da79fbbcSStefano Zampini           /* number of elements NOT on the diagonal */
617da79fbbcSStefano Zampini           nz = adiag[i] - adiag[i+1]-1;
618da79fbbcSStefano Zampini 
619da79fbbcSStefano Zampini           /* decrement the offset */
620da79fbbcSStefano Zampini           offset -= (nz+1);
621da79fbbcSStefano Zampini 
622da79fbbcSStefano Zampini           /* first, set the diagonal elements */
6232cbc15d9SMark           upTriFactor->AA_h[offset] = 1./v[nz];
6242cbc15d9SMark           ierr = PetscArraycpy(&(upTriFactor->AA_h[offset+1]), v, nz);CHKERRQ(ierr);
625da79fbbcSStefano Zampini         }
6262cbc15d9SMark         upTriFactor->csrMat->values->assign(upTriFactor->AA_h, upTriFactor->AA_h+nzUpper);
627da79fbbcSStefano Zampini         ierr = PetscLogCpuToGpu(nzUpper*sizeof(PetscScalar));CHKERRQ(ierr);
628da79fbbcSStefano Zampini       }
6299ae82921SPaul Mullowney     } catch(char *ex) {
6309ae82921SPaul Mullowney       SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_LIB,"CUSPARSE error: %s", ex);
6319ae82921SPaul Mullowney     }
6329ae82921SPaul Mullowney   }
6339ae82921SPaul Mullowney   PetscFunctionReturn(0);
6349ae82921SPaul Mullowney }
6359ae82921SPaul Mullowney 
636087f3262SPaul Mullowney static PetscErrorCode MatSeqAIJCUSPARSEILUAnalysisAndCopyToGPU(Mat A)
6379ae82921SPaul Mullowney {
6389ae82921SPaul Mullowney   PetscErrorCode               ierr;
6399ae82921SPaul Mullowney   Mat_SeqAIJ                   *a                  = (Mat_SeqAIJ*)A->data;
6409ae82921SPaul Mullowney   Mat_SeqAIJCUSPARSETriFactors *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)A->spptr;
6419ae82921SPaul Mullowney   IS                           isrow = a->row,iscol = a->icol;
6429ae82921SPaul Mullowney   PetscBool                    row_identity,col_identity;
6439ae82921SPaul Mullowney   PetscInt                     n = A->rmap->n;
6449ae82921SPaul Mullowney 
6459ae82921SPaul Mullowney   PetscFunctionBegin;
646da79fbbcSStefano Zampini   if (!cusparseTriFactors) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing cusparseTriFactors");
647087f3262SPaul Mullowney   ierr = MatSeqAIJCUSPARSEBuildILULowerTriMatrix(A);CHKERRQ(ierr);
648087f3262SPaul Mullowney   ierr = MatSeqAIJCUSPARSEBuildILUUpperTriMatrix(A);CHKERRQ(ierr);
6492205254eSKarl Rupp 
650da79fbbcSStefano Zampini   if (!cusparseTriFactors->workVector) { cusparseTriFactors->workVector = new THRUSTARRAY(n); }
651aa372e3fSPaul Mullowney   cusparseTriFactors->nnz=a->nz;
6529ae82921SPaul Mullowney 
653c70f7ee4SJunchao Zhang   A->offloadmask = PETSC_OFFLOAD_BOTH;
654e057df02SPaul Mullowney   /* lower triangular indices */
6559ae82921SPaul Mullowney   ierr = ISIdentity(isrow,&row_identity);CHKERRQ(ierr);
656da79fbbcSStefano Zampini   if (!row_identity && !cusparseTriFactors->rpermIndices) {
657da79fbbcSStefano Zampini     const PetscInt *r;
658da79fbbcSStefano Zampini 
659da79fbbcSStefano Zampini     ierr = ISGetIndices(isrow,&r);CHKERRQ(ierr);
660aa372e3fSPaul Mullowney     cusparseTriFactors->rpermIndices = new THRUSTINTARRAY(n);
661aa372e3fSPaul Mullowney     cusparseTriFactors->rpermIndices->assign(r, r+n);
6629ae82921SPaul Mullowney     ierr = ISRestoreIndices(isrow,&r);CHKERRQ(ierr);
663da79fbbcSStefano Zampini     ierr = PetscLogCpuToGpu(n*sizeof(PetscInt));CHKERRQ(ierr);
664da79fbbcSStefano Zampini   }
6659ae82921SPaul Mullowney 
666e057df02SPaul Mullowney   /* upper triangular indices */
6679ae82921SPaul Mullowney   ierr = ISIdentity(iscol,&col_identity);CHKERRQ(ierr);
668da79fbbcSStefano Zampini   if (!col_identity && !cusparseTriFactors->cpermIndices) {
669da79fbbcSStefano Zampini     const PetscInt *c;
670da79fbbcSStefano Zampini 
671da79fbbcSStefano Zampini     ierr = ISGetIndices(iscol,&c);CHKERRQ(ierr);
672aa372e3fSPaul Mullowney     cusparseTriFactors->cpermIndices = new THRUSTINTARRAY(n);
673aa372e3fSPaul Mullowney     cusparseTriFactors->cpermIndices->assign(c, c+n);
6749ae82921SPaul Mullowney     ierr = ISRestoreIndices(iscol,&c);CHKERRQ(ierr);
675da79fbbcSStefano Zampini     ierr = PetscLogCpuToGpu(n*sizeof(PetscInt));CHKERRQ(ierr);
676da79fbbcSStefano Zampini   }
6779ae82921SPaul Mullowney   PetscFunctionReturn(0);
6789ae82921SPaul Mullowney }
6799ae82921SPaul Mullowney 
680087f3262SPaul Mullowney static PetscErrorCode MatSeqAIJCUSPARSEBuildICCTriMatrices(Mat A)
681087f3262SPaul Mullowney {
682087f3262SPaul Mullowney   Mat_SeqAIJ                        *a = (Mat_SeqAIJ*)A->data;
683087f3262SPaul Mullowney   Mat_SeqAIJCUSPARSETriFactors      *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)A->spptr;
684aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct *loTriFactor = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->loTriFactorPtr;
685aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct *upTriFactor = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->upTriFactorPtr;
686087f3262SPaul Mullowney   cusparseStatus_t                  stat;
687087f3262SPaul Mullowney   PetscErrorCode                    ierr;
68857d48284SJunchao Zhang   cudaError_t                       cerr;
689087f3262SPaul Mullowney   PetscInt                          *AiUp, *AjUp;
690087f3262SPaul Mullowney   PetscScalar                       *AAUp;
691087f3262SPaul Mullowney   PetscScalar                       *AALo;
692087f3262SPaul Mullowney   PetscInt                          nzUpper = a->nz,n = A->rmap->n,i,offset,nz,j;
693087f3262SPaul Mullowney   Mat_SeqSBAIJ                      *b = (Mat_SeqSBAIJ*)A->data;
694087f3262SPaul Mullowney   const PetscInt                    *ai = b->i,*aj = b->j,*vj;
695087f3262SPaul Mullowney   const MatScalar                   *aa = b->a,*v;
696087f3262SPaul Mullowney 
697087f3262SPaul Mullowney   PetscFunctionBegin;
698cf00fe3bSKarl Rupp   if (!n) PetscFunctionReturn(0);
699c70f7ee4SJunchao Zhang   if (A->offloadmask == PETSC_OFFLOAD_UNALLOCATED || A->offloadmask == PETSC_OFFLOAD_CPU) {
700087f3262SPaul Mullowney     try {
701da79fbbcSStefano Zampini       cerr = cudaMallocHost((void**) &AAUp, nzUpper*sizeof(PetscScalar));CHKERRCUDA(cerr);
702da79fbbcSStefano Zampini       cerr = cudaMallocHost((void**) &AALo, nzUpper*sizeof(PetscScalar));CHKERRCUDA(cerr);
703da79fbbcSStefano Zampini       if (!upTriFactor && !loTriFactor) {
704087f3262SPaul Mullowney         /* Allocate Space for the upper triangular matrix */
70557d48284SJunchao Zhang         cerr = cudaMallocHost((void**) &AiUp, (n+1)*sizeof(PetscInt));CHKERRCUDA(cerr);
70657d48284SJunchao Zhang         cerr = cudaMallocHost((void**) &AjUp, nzUpper*sizeof(PetscInt));CHKERRCUDA(cerr);
707087f3262SPaul Mullowney 
708087f3262SPaul Mullowney         /* Fill the upper triangular matrix */
709087f3262SPaul Mullowney         AiUp[0]=(PetscInt) 0;
710087f3262SPaul Mullowney         AiUp[n]=nzUpper;
711087f3262SPaul Mullowney         offset = 0;
712087f3262SPaul Mullowney         for (i=0; i<n; i++) {
713087f3262SPaul Mullowney           /* set the pointers */
714087f3262SPaul Mullowney           v  = aa + ai[i];
715087f3262SPaul Mullowney           vj = aj + ai[i];
716087f3262SPaul Mullowney           nz = ai[i+1] - ai[i] - 1; /* exclude diag[i] */
717087f3262SPaul Mullowney 
718087f3262SPaul Mullowney           /* first, set the diagonal elements */
719087f3262SPaul Mullowney           AjUp[offset] = (PetscInt) i;
72009f51544SAlejandro Lamas Daviña           AAUp[offset] = (MatScalar)1.0/v[nz];
721087f3262SPaul Mullowney           AiUp[i]      = offset;
72209f51544SAlejandro Lamas Daviña           AALo[offset] = (MatScalar)1.0/v[nz];
723087f3262SPaul Mullowney 
724087f3262SPaul Mullowney           offset+=1;
725087f3262SPaul Mullowney           if (nz>0) {
726f22e0265SBarry Smith             ierr = PetscArraycpy(&(AjUp[offset]), vj, nz);CHKERRQ(ierr);
727580bdb30SBarry Smith             ierr = PetscArraycpy(&(AAUp[offset]), v, nz);CHKERRQ(ierr);
728087f3262SPaul Mullowney             for (j=offset; j<offset+nz; j++) {
729087f3262SPaul Mullowney               AAUp[j] = -AAUp[j];
730087f3262SPaul Mullowney               AALo[j] = AAUp[j]/v[nz];
731087f3262SPaul Mullowney             }
732087f3262SPaul Mullowney             offset+=nz;
733087f3262SPaul Mullowney           }
734087f3262SPaul Mullowney         }
735087f3262SPaul Mullowney 
736aa372e3fSPaul Mullowney         /* allocate space for the triangular factor information */
737da79fbbcSStefano Zampini         ierr = PetscNew(&upTriFactor);CHKERRQ(ierr);
738da79fbbcSStefano Zampini         upTriFactor->solvePolicy = CUSPARSE_SOLVE_POLICY_USE_LEVEL;
739087f3262SPaul Mullowney 
740aa372e3fSPaul Mullowney         /* Create the matrix description */
74157d48284SJunchao Zhang         stat = cusparseCreateMatDescr(&upTriFactor->descr);CHKERRCUSPARSE(stat);
74257d48284SJunchao Zhang         stat = cusparseSetMatIndexBase(upTriFactor->descr, CUSPARSE_INDEX_BASE_ZERO);CHKERRCUSPARSE(stat);
7431b0a6780SStefano Zampini        #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
744afb2bd1cSJunchao Zhang         stat = cusparseSetMatType(upTriFactor->descr, CUSPARSE_MATRIX_TYPE_GENERAL);CHKERRCUSPARSE(stat);
745afb2bd1cSJunchao Zhang        #else
74657d48284SJunchao Zhang         stat = cusparseSetMatType(upTriFactor->descr, CUSPARSE_MATRIX_TYPE_TRIANGULAR);CHKERRCUSPARSE(stat);
747afb2bd1cSJunchao Zhang        #endif
74857d48284SJunchao Zhang         stat = cusparseSetMatFillMode(upTriFactor->descr, CUSPARSE_FILL_MODE_UPPER);CHKERRCUSPARSE(stat);
74957d48284SJunchao Zhang         stat = cusparseSetMatDiagType(upTriFactor->descr, CUSPARSE_DIAG_TYPE_UNIT);CHKERRCUSPARSE(stat);
750087f3262SPaul Mullowney 
751aa372e3fSPaul Mullowney         /* set the matrix */
752aa372e3fSPaul Mullowney         upTriFactor->csrMat = new CsrMatrix;
753aa372e3fSPaul Mullowney         upTriFactor->csrMat->num_rows = A->rmap->n;
754aa372e3fSPaul Mullowney         upTriFactor->csrMat->num_cols = A->cmap->n;
755aa372e3fSPaul Mullowney         upTriFactor->csrMat->num_entries = a->nz;
756aa372e3fSPaul Mullowney 
757aa372e3fSPaul Mullowney         upTriFactor->csrMat->row_offsets = new THRUSTINTARRAY32(A->rmap->n+1);
758aa372e3fSPaul Mullowney         upTriFactor->csrMat->row_offsets->assign(AiUp, AiUp+A->rmap->n+1);
759aa372e3fSPaul Mullowney 
760aa372e3fSPaul Mullowney         upTriFactor->csrMat->column_indices = new THRUSTINTARRAY32(a->nz);
761aa372e3fSPaul Mullowney         upTriFactor->csrMat->column_indices->assign(AjUp, AjUp+a->nz);
762aa372e3fSPaul Mullowney 
763aa372e3fSPaul Mullowney         upTriFactor->csrMat->values = new THRUSTARRAY(a->nz);
764aa372e3fSPaul Mullowney         upTriFactor->csrMat->values->assign(AAUp, AAUp+a->nz);
765aa372e3fSPaul Mullowney 
766afb2bd1cSJunchao Zhang         /* set the operation */
767afb2bd1cSJunchao Zhang         upTriFactor->solveOp = CUSPARSE_OPERATION_NON_TRANSPOSE;
768afb2bd1cSJunchao Zhang 
769afb2bd1cSJunchao Zhang         /* Create the solve analysis information */
770da79fbbcSStefano Zampini         ierr = PetscLogEventBegin(MAT_CUSPARSESolveAnalysis,A,0,0,0);CHKERRQ(ierr);
771afb2bd1cSJunchao Zhang         stat = cusparse_create_analysis_info(&upTriFactor->solveInfo);CHKERRCUSPARSE(stat);
7721b0a6780SStefano Zampini       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
773afb2bd1cSJunchao Zhang         stat = cusparse_get_svbuffsize(cusparseTriFactors->handle, upTriFactor->solveOp,
774afb2bd1cSJunchao Zhang                                        upTriFactor->csrMat->num_rows, upTriFactor->csrMat->num_entries, upTriFactor->descr,
775afb2bd1cSJunchao Zhang                                        upTriFactor->csrMat->values->data().get(), upTriFactor->csrMat->row_offsets->data().get(),
776afb2bd1cSJunchao Zhang                                        upTriFactor->csrMat->column_indices->data().get(), upTriFactor->solveInfo,
777afb2bd1cSJunchao Zhang                                        &upTriFactor->solveBufferSize);CHKERRCUSPARSE(stat);
778afb2bd1cSJunchao Zhang         cerr = cudaMalloc(&upTriFactor->solveBuffer,upTriFactor->solveBufferSize);CHKERRCUDA(cerr);
779afb2bd1cSJunchao Zhang       #endif
780afb2bd1cSJunchao Zhang 
781aa372e3fSPaul Mullowney         /* perform the solve analysis */
782aa372e3fSPaul Mullowney         stat = cusparse_analysis(cusparseTriFactors->handle, upTriFactor->solveOp,
783aa372e3fSPaul Mullowney                                  upTriFactor->csrMat->num_rows, upTriFactor->csrMat->num_entries, upTriFactor->descr,
784aa372e3fSPaul Mullowney                                  upTriFactor->csrMat->values->data().get(), upTriFactor->csrMat->row_offsets->data().get(),
785afb2bd1cSJunchao Zhang                                  upTriFactor->csrMat->column_indices->data().get(), upTriFactor->solveInfo
7861b0a6780SStefano Zampini                                 #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
787afb2bd1cSJunchao Zhang                                  ,upTriFactor->solvePolicy, upTriFactor->solveBuffer
788afb2bd1cSJunchao Zhang                                 #endif
789afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
790da79fbbcSStefano Zampini         cerr = WaitForCUDA();CHKERRCUDA(cerr);
791da79fbbcSStefano Zampini         ierr = PetscLogEventEnd(MAT_CUSPARSESolveAnalysis,A,0,0,0);CHKERRQ(ierr);
792aa372e3fSPaul Mullowney 
793da79fbbcSStefano Zampini         /* assign the pointer */
794aa372e3fSPaul Mullowney         ((Mat_SeqAIJCUSPARSETriFactors*)A->spptr)->upTriFactorPtr = upTriFactor;
795aa372e3fSPaul Mullowney 
796aa372e3fSPaul Mullowney         /* allocate space for the triangular factor information */
797da79fbbcSStefano Zampini         ierr = PetscNew(&loTriFactor);CHKERRQ(ierr);
798da79fbbcSStefano Zampini         loTriFactor->solvePolicy = CUSPARSE_SOLVE_POLICY_USE_LEVEL;
799aa372e3fSPaul Mullowney 
800aa372e3fSPaul Mullowney         /* Create the matrix description */
80157d48284SJunchao Zhang         stat = cusparseCreateMatDescr(&loTriFactor->descr);CHKERRCUSPARSE(stat);
80257d48284SJunchao Zhang         stat = cusparseSetMatIndexBase(loTriFactor->descr, CUSPARSE_INDEX_BASE_ZERO);CHKERRCUSPARSE(stat);
8031b0a6780SStefano Zampini        #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
804afb2bd1cSJunchao Zhang         stat = cusparseSetMatType(loTriFactor->descr, CUSPARSE_MATRIX_TYPE_GENERAL);CHKERRCUSPARSE(stat);
805afb2bd1cSJunchao Zhang        #else
80657d48284SJunchao Zhang         stat = cusparseSetMatType(loTriFactor->descr, CUSPARSE_MATRIX_TYPE_TRIANGULAR);CHKERRCUSPARSE(stat);
807afb2bd1cSJunchao Zhang        #endif
80857d48284SJunchao Zhang         stat = cusparseSetMatFillMode(loTriFactor->descr, CUSPARSE_FILL_MODE_UPPER);CHKERRCUSPARSE(stat);
80957d48284SJunchao Zhang         stat = cusparseSetMatDiagType(loTriFactor->descr, CUSPARSE_DIAG_TYPE_NON_UNIT);CHKERRCUSPARSE(stat);
810aa372e3fSPaul Mullowney 
811aa372e3fSPaul Mullowney         /* set the operation */
812aa372e3fSPaul Mullowney         loTriFactor->solveOp = CUSPARSE_OPERATION_TRANSPOSE;
813aa372e3fSPaul Mullowney 
814aa372e3fSPaul Mullowney         /* set the matrix */
815aa372e3fSPaul Mullowney         loTriFactor->csrMat = new CsrMatrix;
816aa372e3fSPaul Mullowney         loTriFactor->csrMat->num_rows = A->rmap->n;
817aa372e3fSPaul Mullowney         loTriFactor->csrMat->num_cols = A->cmap->n;
818aa372e3fSPaul Mullowney         loTriFactor->csrMat->num_entries = a->nz;
819aa372e3fSPaul Mullowney 
820aa372e3fSPaul Mullowney         loTriFactor->csrMat->row_offsets = new THRUSTINTARRAY32(A->rmap->n+1);
821aa372e3fSPaul Mullowney         loTriFactor->csrMat->row_offsets->assign(AiUp, AiUp+A->rmap->n+1);
822aa372e3fSPaul Mullowney 
823aa372e3fSPaul Mullowney         loTriFactor->csrMat->column_indices = new THRUSTINTARRAY32(a->nz);
824aa372e3fSPaul Mullowney         loTriFactor->csrMat->column_indices->assign(AjUp, AjUp+a->nz);
825aa372e3fSPaul Mullowney 
826aa372e3fSPaul Mullowney         loTriFactor->csrMat->values = new THRUSTARRAY(a->nz);
827aa372e3fSPaul Mullowney         loTriFactor->csrMat->values->assign(AALo, AALo+a->nz);
828aa372e3fSPaul Mullowney 
829afb2bd1cSJunchao Zhang         /* Create the solve analysis information */
830da79fbbcSStefano Zampini         ierr = PetscLogEventBegin(MAT_CUSPARSESolveAnalysis,A,0,0,0);CHKERRQ(ierr);
831afb2bd1cSJunchao Zhang         stat = cusparse_create_analysis_info(&loTriFactor->solveInfo);CHKERRCUSPARSE(stat);
8321b0a6780SStefano Zampini       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
833afb2bd1cSJunchao Zhang         stat = cusparse_get_svbuffsize(cusparseTriFactors->handle, loTriFactor->solveOp,
834afb2bd1cSJunchao Zhang                                        loTriFactor->csrMat->num_rows, loTriFactor->csrMat->num_entries, loTriFactor->descr,
835afb2bd1cSJunchao Zhang                                        loTriFactor->csrMat->values->data().get(), loTriFactor->csrMat->row_offsets->data().get(),
836afb2bd1cSJunchao Zhang                                        loTriFactor->csrMat->column_indices->data().get(), loTriFactor->solveInfo,
837afb2bd1cSJunchao Zhang                                        &loTriFactor->solveBufferSize);CHKERRCUSPARSE(stat);
838afb2bd1cSJunchao Zhang         cerr = cudaMalloc(&loTriFactor->solveBuffer,loTriFactor->solveBufferSize);CHKERRCUDA(cerr);
839afb2bd1cSJunchao Zhang       #endif
840afb2bd1cSJunchao Zhang 
841aa372e3fSPaul Mullowney         /* perform the solve analysis */
842aa372e3fSPaul Mullowney         stat = cusparse_analysis(cusparseTriFactors->handle, loTriFactor->solveOp,
843aa372e3fSPaul Mullowney                                  loTriFactor->csrMat->num_rows, loTriFactor->csrMat->num_entries, loTriFactor->descr,
844aa372e3fSPaul Mullowney                                  loTriFactor->csrMat->values->data().get(), loTriFactor->csrMat->row_offsets->data().get(),
845afb2bd1cSJunchao Zhang                                  loTriFactor->csrMat->column_indices->data().get(), loTriFactor->solveInfo
8461b0a6780SStefano Zampini                                 #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
847afb2bd1cSJunchao Zhang                                  ,loTriFactor->solvePolicy, loTriFactor->solveBuffer
848afb2bd1cSJunchao Zhang                                 #endif
849afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
850da79fbbcSStefano Zampini         cerr = WaitForCUDA();CHKERRCUDA(cerr);
851da79fbbcSStefano Zampini         ierr = PetscLogEventEnd(MAT_CUSPARSESolveAnalysis,A,0,0,0);CHKERRQ(ierr);
852aa372e3fSPaul Mullowney 
853da79fbbcSStefano Zampini         /* assign the pointer */
854aa372e3fSPaul Mullowney         ((Mat_SeqAIJCUSPARSETriFactors*)A->spptr)->loTriFactorPtr = loTriFactor;
855087f3262SPaul Mullowney 
856da79fbbcSStefano Zampini         ierr = PetscLogCpuToGpu(2*(((A->rmap->n+1)+(a->nz))*sizeof(int)+(a->nz)*sizeof(PetscScalar)));CHKERRQ(ierr);
85757d48284SJunchao Zhang         cerr = cudaFreeHost(AiUp);CHKERRCUDA(cerr);
85857d48284SJunchao Zhang         cerr = cudaFreeHost(AjUp);CHKERRCUDA(cerr);
859da79fbbcSStefano Zampini       } else {
860da79fbbcSStefano Zampini         /* Fill the upper triangular matrix */
861da79fbbcSStefano Zampini         offset = 0;
862da79fbbcSStefano Zampini         for (i=0; i<n; i++) {
863da79fbbcSStefano Zampini           /* set the pointers */
864da79fbbcSStefano Zampini           v  = aa + ai[i];
865da79fbbcSStefano Zampini           nz = ai[i+1] - ai[i] - 1; /* exclude diag[i] */
866da79fbbcSStefano Zampini 
867da79fbbcSStefano Zampini           /* first, set the diagonal elements */
868da79fbbcSStefano Zampini           AAUp[offset] = 1.0/v[nz];
869da79fbbcSStefano Zampini           AALo[offset] = 1.0/v[nz];
870da79fbbcSStefano Zampini 
871da79fbbcSStefano Zampini           offset+=1;
872da79fbbcSStefano Zampini           if (nz>0) {
873da79fbbcSStefano Zampini             ierr = PetscArraycpy(&(AAUp[offset]), v, nz);CHKERRQ(ierr);
874da79fbbcSStefano Zampini             for (j=offset; j<offset+nz; j++) {
875da79fbbcSStefano Zampini               AAUp[j] = -AAUp[j];
876da79fbbcSStefano Zampini               AALo[j] = AAUp[j]/v[nz];
877da79fbbcSStefano Zampini             }
878da79fbbcSStefano Zampini             offset+=nz;
879da79fbbcSStefano Zampini           }
880da79fbbcSStefano Zampini         }
881da79fbbcSStefano Zampini         if (!upTriFactor) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing cusparseTriFactors");
882da79fbbcSStefano Zampini         if (!loTriFactor) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing cusparseTriFactors");
883da79fbbcSStefano Zampini         upTriFactor->csrMat->values->assign(AAUp, AAUp+a->nz);
884da79fbbcSStefano Zampini         loTriFactor->csrMat->values->assign(AALo, AALo+a->nz);
885da79fbbcSStefano Zampini         ierr = PetscLogCpuToGpu(2*(a->nz)*sizeof(PetscScalar));CHKERRQ(ierr);
886da79fbbcSStefano Zampini       }
88757d48284SJunchao Zhang       cerr = cudaFreeHost(AAUp);CHKERRCUDA(cerr);
88857d48284SJunchao Zhang       cerr = cudaFreeHost(AALo);CHKERRCUDA(cerr);
889087f3262SPaul Mullowney     } catch(char *ex) {
890087f3262SPaul Mullowney       SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_LIB,"CUSPARSE error: %s", ex);
891087f3262SPaul Mullowney     }
892087f3262SPaul Mullowney   }
893087f3262SPaul Mullowney   PetscFunctionReturn(0);
894087f3262SPaul Mullowney }
895087f3262SPaul Mullowney 
896087f3262SPaul Mullowney static PetscErrorCode MatSeqAIJCUSPARSEICCAnalysisAndCopyToGPU(Mat A)
8979ae82921SPaul Mullowney {
8989ae82921SPaul Mullowney   PetscErrorCode               ierr;
899087f3262SPaul Mullowney   Mat_SeqAIJ                   *a                  = (Mat_SeqAIJ*)A->data;
900087f3262SPaul Mullowney   Mat_SeqAIJCUSPARSETriFactors *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)A->spptr;
901087f3262SPaul Mullowney   IS                           ip = a->row;
902087f3262SPaul Mullowney   PetscBool                    perm_identity;
903087f3262SPaul Mullowney   PetscInt                     n = A->rmap->n;
904087f3262SPaul Mullowney 
905087f3262SPaul Mullowney   PetscFunctionBegin;
906da79fbbcSStefano Zampini   if (!cusparseTriFactors) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing cusparseTriFactors");
907087f3262SPaul Mullowney   ierr = MatSeqAIJCUSPARSEBuildICCTriMatrices(A);CHKERRQ(ierr);
908da79fbbcSStefano Zampini   if (!cusparseTriFactors->workVector) { cusparseTriFactors->workVector = new THRUSTARRAY(n); }
909aa372e3fSPaul Mullowney   cusparseTriFactors->nnz=(a->nz-n)*2 + n;
910aa372e3fSPaul Mullowney 
911da79fbbcSStefano Zampini   A->offloadmask = PETSC_OFFLOAD_BOTH;
912da79fbbcSStefano Zampini 
913087f3262SPaul Mullowney   /* lower triangular indices */
914087f3262SPaul Mullowney   ierr = ISIdentity(ip,&perm_identity);CHKERRQ(ierr);
915087f3262SPaul Mullowney   if (!perm_identity) {
9164e4bbfaaSStefano Zampini     IS             iip;
917da79fbbcSStefano Zampini     const PetscInt *irip,*rip;
9184e4bbfaaSStefano Zampini 
9194e4bbfaaSStefano Zampini     ierr = ISInvertPermutation(ip,PETSC_DECIDE,&iip);CHKERRQ(ierr);
9204e4bbfaaSStefano Zampini     ierr = ISGetIndices(iip,&irip);CHKERRQ(ierr);
921da79fbbcSStefano Zampini     ierr = ISGetIndices(ip,&rip);CHKERRQ(ierr);
922aa372e3fSPaul Mullowney     cusparseTriFactors->rpermIndices = new THRUSTINTARRAY(n);
923aa372e3fSPaul Mullowney     cusparseTriFactors->rpermIndices->assign(rip, rip+n);
924aa372e3fSPaul Mullowney     cusparseTriFactors->cpermIndices = new THRUSTINTARRAY(n);
9254e4bbfaaSStefano Zampini     cusparseTriFactors->cpermIndices->assign(irip, irip+n);
9264e4bbfaaSStefano Zampini     ierr = ISRestoreIndices(iip,&irip);CHKERRQ(ierr);
9274e4bbfaaSStefano Zampini     ierr = ISDestroy(&iip);CHKERRQ(ierr);
928087f3262SPaul Mullowney     ierr = ISRestoreIndices(ip,&rip);CHKERRQ(ierr);
929da79fbbcSStefano Zampini     ierr = PetscLogCpuToGpu(2.*n*sizeof(PetscInt));CHKERRQ(ierr);
930da79fbbcSStefano Zampini   }
931087f3262SPaul Mullowney   PetscFunctionReturn(0);
932087f3262SPaul Mullowney }
933087f3262SPaul Mullowney 
9346fa9248bSJed Brown static PetscErrorCode MatLUFactorNumeric_SeqAIJCUSPARSE(Mat B,Mat A,const MatFactorInfo *info)
9359ae82921SPaul Mullowney {
9369ae82921SPaul Mullowney   Mat_SeqAIJ     *b = (Mat_SeqAIJ*)B->data;
9379ae82921SPaul Mullowney   IS             isrow = b->row,iscol = b->col;
9389ae82921SPaul Mullowney   PetscBool      row_identity,col_identity;
939b175d8bbSPaul Mullowney   PetscErrorCode ierr;
9409ae82921SPaul Mullowney 
9419ae82921SPaul Mullowney   PetscFunctionBegin;
94257181aedSStefano Zampini   ierr = MatSeqAIJCUSPARSECopyFromGPU(A);CHKERRQ(ierr);
9439ae82921SPaul Mullowney   ierr = MatLUFactorNumeric_SeqAIJ(B,A,info);CHKERRQ(ierr);
944ccdfe979SStefano Zampini   B->offloadmask = PETSC_OFFLOAD_CPU;
945e057df02SPaul Mullowney   /* determine which version of MatSolve needs to be used. */
9469ae82921SPaul Mullowney   ierr = ISIdentity(isrow,&row_identity);CHKERRQ(ierr);
9479ae82921SPaul Mullowney   ierr = ISIdentity(iscol,&col_identity);CHKERRQ(ierr);
948bda325fcSPaul Mullowney   if (row_identity && col_identity) {
949bda325fcSPaul Mullowney     B->ops->solve = MatSolve_SeqAIJCUSPARSE_NaturalOrdering;
950bda325fcSPaul Mullowney     B->ops->solvetranspose = MatSolveTranspose_SeqAIJCUSPARSE_NaturalOrdering;
9514e4bbfaaSStefano Zampini     B->ops->matsolve = NULL;
9524e4bbfaaSStefano Zampini     B->ops->matsolvetranspose = NULL;
953bda325fcSPaul Mullowney   } else {
954bda325fcSPaul Mullowney     B->ops->solve = MatSolve_SeqAIJCUSPARSE;
955bda325fcSPaul Mullowney     B->ops->solvetranspose = MatSolveTranspose_SeqAIJCUSPARSE;
9564e4bbfaaSStefano Zampini     B->ops->matsolve = NULL;
9574e4bbfaaSStefano Zampini     B->ops->matsolvetranspose = NULL;
958bda325fcSPaul Mullowney   }
9598dc1d2a3SPaul Mullowney 
960e057df02SPaul Mullowney   /* get the triangular factors */
961087f3262SPaul Mullowney   ierr = MatSeqAIJCUSPARSEILUAnalysisAndCopyToGPU(B);CHKERRQ(ierr);
9629ae82921SPaul Mullowney   PetscFunctionReturn(0);
9639ae82921SPaul Mullowney }
9649ae82921SPaul Mullowney 
965087f3262SPaul Mullowney static PetscErrorCode MatCholeskyFactorNumeric_SeqAIJCUSPARSE(Mat B,Mat A,const MatFactorInfo *info)
966087f3262SPaul Mullowney {
967087f3262SPaul Mullowney   Mat_SeqAIJ     *b = (Mat_SeqAIJ*)B->data;
968087f3262SPaul Mullowney   IS             ip = b->row;
969087f3262SPaul Mullowney   PetscBool      perm_identity;
970b175d8bbSPaul Mullowney   PetscErrorCode ierr;
971087f3262SPaul Mullowney 
972087f3262SPaul Mullowney   PetscFunctionBegin;
97357181aedSStefano Zampini   ierr = MatSeqAIJCUSPARSECopyFromGPU(A);CHKERRQ(ierr);
974087f3262SPaul Mullowney   ierr = MatCholeskyFactorNumeric_SeqAIJ(B,A,info);CHKERRQ(ierr);
975ccdfe979SStefano Zampini   B->offloadmask = PETSC_OFFLOAD_CPU;
976087f3262SPaul Mullowney   /* determine which version of MatSolve needs to be used. */
977087f3262SPaul Mullowney   ierr = ISIdentity(ip,&perm_identity);CHKERRQ(ierr);
978087f3262SPaul Mullowney   if (perm_identity) {
979087f3262SPaul Mullowney     B->ops->solve = MatSolve_SeqAIJCUSPARSE_NaturalOrdering;
980087f3262SPaul Mullowney     B->ops->solvetranspose = MatSolveTranspose_SeqAIJCUSPARSE_NaturalOrdering;
9814e4bbfaaSStefano Zampini     B->ops->matsolve = NULL;
9824e4bbfaaSStefano Zampini     B->ops->matsolvetranspose = NULL;
983087f3262SPaul Mullowney   } else {
984087f3262SPaul Mullowney     B->ops->solve = MatSolve_SeqAIJCUSPARSE;
985087f3262SPaul Mullowney     B->ops->solvetranspose = MatSolveTranspose_SeqAIJCUSPARSE;
9864e4bbfaaSStefano Zampini     B->ops->matsolve = NULL;
9874e4bbfaaSStefano Zampini     B->ops->matsolvetranspose = NULL;
988087f3262SPaul Mullowney   }
989087f3262SPaul Mullowney 
990087f3262SPaul Mullowney   /* get the triangular factors */
991087f3262SPaul Mullowney   ierr = MatSeqAIJCUSPARSEICCAnalysisAndCopyToGPU(B);CHKERRQ(ierr);
992087f3262SPaul Mullowney   PetscFunctionReturn(0);
993087f3262SPaul Mullowney }
9949ae82921SPaul Mullowney 
995b175d8bbSPaul Mullowney static PetscErrorCode MatSeqAIJCUSPARSEAnalyzeTransposeForSolve(Mat A)
996bda325fcSPaul Mullowney {
997bda325fcSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactors      *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)A->spptr;
998aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct *loTriFactor = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->loTriFactorPtr;
999aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct *upTriFactor = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->upTriFactorPtr;
1000da79fbbcSStefano Zampini   Mat_SeqAIJCUSPARSETriFactorStruct *loTriFactorT;
1001da79fbbcSStefano Zampini   Mat_SeqAIJCUSPARSETriFactorStruct *upTriFactorT;
1002bda325fcSPaul Mullowney   cusparseStatus_t                  stat;
1003aa372e3fSPaul Mullowney   cusparseIndexBase_t               indexBase;
1004aa372e3fSPaul Mullowney   cusparseMatrixType_t              matrixType;
1005aa372e3fSPaul Mullowney   cusparseFillMode_t                fillMode;
1006aa372e3fSPaul Mullowney   cusparseDiagType_t                diagType;
10071b0a6780SStefano Zampini   cudaError_t                       cerr;
1008da79fbbcSStefano Zampini   PetscErrorCode                    ierr;
1009b175d8bbSPaul Mullowney 
1010bda325fcSPaul Mullowney   PetscFunctionBegin;
1011aa372e3fSPaul Mullowney   /* allocate space for the transpose of the lower triangular factor */
1012da79fbbcSStefano Zampini   ierr = PetscNew(&loTriFactorT);CHKERRQ(ierr);
1013da79fbbcSStefano Zampini   loTriFactorT->solvePolicy = CUSPARSE_SOLVE_POLICY_USE_LEVEL;
1014aa372e3fSPaul Mullowney 
1015aa372e3fSPaul Mullowney   /* set the matrix descriptors of the lower triangular factor */
1016aa372e3fSPaul Mullowney   matrixType = cusparseGetMatType(loTriFactor->descr);
1017aa372e3fSPaul Mullowney   indexBase = cusparseGetMatIndexBase(loTriFactor->descr);
1018aa372e3fSPaul Mullowney   fillMode = cusparseGetMatFillMode(loTriFactor->descr)==CUSPARSE_FILL_MODE_UPPER ?
1019aa372e3fSPaul Mullowney     CUSPARSE_FILL_MODE_LOWER : CUSPARSE_FILL_MODE_UPPER;
1020aa372e3fSPaul Mullowney   diagType = cusparseGetMatDiagType(loTriFactor->descr);
1021aa372e3fSPaul Mullowney 
1022aa372e3fSPaul Mullowney   /* Create the matrix description */
102357d48284SJunchao Zhang   stat = cusparseCreateMatDescr(&loTriFactorT->descr);CHKERRCUSPARSE(stat);
102457d48284SJunchao Zhang   stat = cusparseSetMatIndexBase(loTriFactorT->descr, indexBase);CHKERRCUSPARSE(stat);
102557d48284SJunchao Zhang   stat = cusparseSetMatType(loTriFactorT->descr, matrixType);CHKERRCUSPARSE(stat);
102657d48284SJunchao Zhang   stat = cusparseSetMatFillMode(loTriFactorT->descr, fillMode);CHKERRCUSPARSE(stat);
102757d48284SJunchao Zhang   stat = cusparseSetMatDiagType(loTriFactorT->descr, diagType);CHKERRCUSPARSE(stat);
1028aa372e3fSPaul Mullowney 
1029aa372e3fSPaul Mullowney   /* set the operation */
1030aa372e3fSPaul Mullowney   loTriFactorT->solveOp = CUSPARSE_OPERATION_NON_TRANSPOSE;
1031aa372e3fSPaul Mullowney 
1032aa372e3fSPaul Mullowney   /* allocate GPU space for the CSC of the lower triangular factor*/
1033aa372e3fSPaul Mullowney   loTriFactorT->csrMat = new CsrMatrix;
1034afb2bd1cSJunchao Zhang   loTriFactorT->csrMat->num_rows       = loTriFactor->csrMat->num_cols;
1035afb2bd1cSJunchao Zhang   loTriFactorT->csrMat->num_cols       = loTriFactor->csrMat->num_rows;
1036aa372e3fSPaul Mullowney   loTriFactorT->csrMat->num_entries    = loTriFactor->csrMat->num_entries;
1037afb2bd1cSJunchao Zhang   loTriFactorT->csrMat->row_offsets    = new THRUSTINTARRAY32(loTriFactorT->csrMat->num_rows+1);
1038afb2bd1cSJunchao Zhang   loTriFactorT->csrMat->column_indices = new THRUSTINTARRAY32(loTriFactorT->csrMat->num_entries);
1039afb2bd1cSJunchao Zhang   loTriFactorT->csrMat->values         = new THRUSTARRAY(loTriFactorT->csrMat->num_entries);
1040aa372e3fSPaul Mullowney 
1041aa372e3fSPaul Mullowney   /* compute the transpose of the lower triangular factor, i.e. the CSC */
1042afb2bd1cSJunchao Zhang #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1043afb2bd1cSJunchao Zhang   stat = cusparseCsr2cscEx2_bufferSize(cusparseTriFactors->handle, loTriFactor->csrMat->num_rows,
1044afb2bd1cSJunchao Zhang                                        loTriFactor->csrMat->num_cols, loTriFactor->csrMat->num_entries,
1045afb2bd1cSJunchao Zhang                                        loTriFactor->csrMat->values->data().get(),
1046afb2bd1cSJunchao Zhang                                        loTriFactor->csrMat->row_offsets->data().get(),
1047afb2bd1cSJunchao Zhang                                        loTriFactor->csrMat->column_indices->data().get(),
1048afb2bd1cSJunchao Zhang                                        loTriFactorT->csrMat->values->data().get(),
1049afb2bd1cSJunchao Zhang                                        loTriFactorT->csrMat->row_offsets->data().get(), loTriFactorT->csrMat->column_indices->data().get(), cusparse_scalartype,
1050afb2bd1cSJunchao Zhang                                        CUSPARSE_ACTION_NUMERIC,indexBase,
1051afb2bd1cSJunchao Zhang                                        CUSPARSE_CSR2CSC_ALG1, &loTriFactor->csr2cscBufferSize);CHKERRCUSPARSE(stat);
10521b0a6780SStefano Zampini   cerr = cudaMalloc(&loTriFactor->csr2cscBuffer,loTriFactor->csr2cscBufferSize);CHKERRCUDA(cerr);
1053afb2bd1cSJunchao Zhang #endif
1054afb2bd1cSJunchao Zhang 
1055da79fbbcSStefano Zampini   ierr = PetscLogEventBegin(MAT_CUSPARSEGenerateTranspose,A,0,0,0);CHKERRQ(ierr);
1056aa372e3fSPaul Mullowney   stat = cusparse_csr2csc(cusparseTriFactors->handle, loTriFactor->csrMat->num_rows,
1057aa372e3fSPaul Mullowney                           loTriFactor->csrMat->num_cols, loTriFactor->csrMat->num_entries,
1058aa372e3fSPaul Mullowney                           loTriFactor->csrMat->values->data().get(),
1059aa372e3fSPaul Mullowney                           loTriFactor->csrMat->row_offsets->data().get(),
1060aa372e3fSPaul Mullowney                           loTriFactor->csrMat->column_indices->data().get(),
1061aa372e3fSPaul Mullowney                           loTriFactorT->csrMat->values->data().get(),
1062afb2bd1cSJunchao Zhang                         #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1063afb2bd1cSJunchao Zhang                           loTriFactorT->csrMat->row_offsets->data().get(), loTriFactorT->csrMat->column_indices->data().get(), cusparse_scalartype,
1064afb2bd1cSJunchao Zhang                           CUSPARSE_ACTION_NUMERIC, indexBase,
1065afb2bd1cSJunchao Zhang                           CUSPARSE_CSR2CSC_ALG1, loTriFactor->csr2cscBuffer
1066afb2bd1cSJunchao Zhang                         #else
1067afb2bd1cSJunchao Zhang                           loTriFactorT->csrMat->column_indices->data().get(), loTriFactorT->csrMat->row_offsets->data().get(),
1068afb2bd1cSJunchao Zhang                           CUSPARSE_ACTION_NUMERIC, indexBase
1069afb2bd1cSJunchao Zhang                         #endif
1070afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
1071da79fbbcSStefano Zampini   cerr = WaitForCUDA();CHKERRCUDA(cerr);
1072da79fbbcSStefano Zampini   ierr = PetscLogEventBegin(MAT_CUSPARSEGenerateTranspose,A,0,0,0);CHKERRQ(ierr);
1073aa372e3fSPaul Mullowney 
1074afb2bd1cSJunchao Zhang   /* Create the solve analysis information */
1075da79fbbcSStefano Zampini   ierr = PetscLogEventBegin(MAT_CUSPARSESolveAnalysis,A,0,0,0);CHKERRQ(ierr);
1076afb2bd1cSJunchao Zhang   stat = cusparse_create_analysis_info(&loTriFactorT->solveInfo);CHKERRCUSPARSE(stat);
10771b0a6780SStefano Zampini #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1078afb2bd1cSJunchao Zhang   stat = cusparse_get_svbuffsize(cusparseTriFactors->handle, loTriFactorT->solveOp,
1079afb2bd1cSJunchao Zhang                                 loTriFactorT->csrMat->num_rows, loTriFactorT->csrMat->num_entries, loTriFactorT->descr,
1080afb2bd1cSJunchao Zhang                                 loTriFactorT->csrMat->values->data().get(), loTriFactorT->csrMat->row_offsets->data().get(),
1081afb2bd1cSJunchao Zhang                                 loTriFactorT->csrMat->column_indices->data().get(), loTriFactorT->solveInfo,
1082afb2bd1cSJunchao Zhang                                 &loTriFactorT->solveBufferSize);CHKERRCUSPARSE(stat);
1083afb2bd1cSJunchao Zhang   cerr = cudaMalloc(&loTriFactorT->solveBuffer,loTriFactorT->solveBufferSize);CHKERRCUDA(cerr);
1084afb2bd1cSJunchao Zhang #endif
1085afb2bd1cSJunchao Zhang 
1086afb2bd1cSJunchao Zhang   /* perform the solve analysis */
1087aa372e3fSPaul Mullowney   stat = cusparse_analysis(cusparseTriFactors->handle, loTriFactorT->solveOp,
1088afb2bd1cSJunchao Zhang                            loTriFactorT->csrMat->num_rows, loTriFactorT->csrMat->num_entries, loTriFactorT->descr,
1089afb2bd1cSJunchao Zhang                            loTriFactorT->csrMat->values->data().get(), loTriFactorT->csrMat->row_offsets->data().get(),
1090afb2bd1cSJunchao Zhang                            loTriFactorT->csrMat->column_indices->data().get(), loTriFactorT->solveInfo
10911b0a6780SStefano Zampini                           #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1092afb2bd1cSJunchao Zhang                            ,loTriFactorT->solvePolicy, loTriFactorT->solveBuffer
1093afb2bd1cSJunchao Zhang                           #endif
1094afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
1095da79fbbcSStefano Zampini   cerr = WaitForCUDA();CHKERRCUDA(cerr);
1096da79fbbcSStefano Zampini   ierr = PetscLogEventEnd(MAT_CUSPARSESolveAnalysis,A,0,0,0);CHKERRQ(ierr);
1097aa372e3fSPaul Mullowney 
1098da79fbbcSStefano Zampini   /* assign the pointer */
1099aa372e3fSPaul Mullowney   ((Mat_SeqAIJCUSPARSETriFactors*)A->spptr)->loTriFactorPtrTranspose = loTriFactorT;
1100aa372e3fSPaul Mullowney 
1101aa372e3fSPaul Mullowney   /*********************************************/
1102aa372e3fSPaul Mullowney   /* Now the Transpose of the Upper Tri Factor */
1103aa372e3fSPaul Mullowney   /*********************************************/
1104aa372e3fSPaul Mullowney 
1105aa372e3fSPaul Mullowney   /* allocate space for the transpose of the upper triangular factor */
1106da79fbbcSStefano Zampini   ierr = PetscNew(&upTriFactorT);CHKERRQ(ierr);
1107da79fbbcSStefano Zampini   upTriFactorT->solvePolicy = CUSPARSE_SOLVE_POLICY_USE_LEVEL;
1108aa372e3fSPaul Mullowney 
1109aa372e3fSPaul Mullowney   /* set the matrix descriptors of the upper triangular factor */
1110aa372e3fSPaul Mullowney   matrixType = cusparseGetMatType(upTriFactor->descr);
1111aa372e3fSPaul Mullowney   indexBase = cusparseGetMatIndexBase(upTriFactor->descr);
1112aa372e3fSPaul Mullowney   fillMode = cusparseGetMatFillMode(upTriFactor->descr)==CUSPARSE_FILL_MODE_UPPER ?
1113aa372e3fSPaul Mullowney     CUSPARSE_FILL_MODE_LOWER : CUSPARSE_FILL_MODE_UPPER;
1114aa372e3fSPaul Mullowney   diagType = cusparseGetMatDiagType(upTriFactor->descr);
1115aa372e3fSPaul Mullowney 
1116aa372e3fSPaul Mullowney   /* Create the matrix description */
111757d48284SJunchao Zhang   stat = cusparseCreateMatDescr(&upTriFactorT->descr);CHKERRCUSPARSE(stat);
111857d48284SJunchao Zhang   stat = cusparseSetMatIndexBase(upTriFactorT->descr, indexBase);CHKERRCUSPARSE(stat);
111957d48284SJunchao Zhang   stat = cusparseSetMatType(upTriFactorT->descr, matrixType);CHKERRCUSPARSE(stat);
112057d48284SJunchao Zhang   stat = cusparseSetMatFillMode(upTriFactorT->descr, fillMode);CHKERRCUSPARSE(stat);
112157d48284SJunchao Zhang   stat = cusparseSetMatDiagType(upTriFactorT->descr, diagType);CHKERRCUSPARSE(stat);
1122aa372e3fSPaul Mullowney 
1123aa372e3fSPaul Mullowney   /* set the operation */
1124aa372e3fSPaul Mullowney   upTriFactorT->solveOp = CUSPARSE_OPERATION_NON_TRANSPOSE;
1125aa372e3fSPaul Mullowney 
1126aa372e3fSPaul Mullowney   /* allocate GPU space for the CSC of the upper triangular factor*/
1127aa372e3fSPaul Mullowney   upTriFactorT->csrMat = new CsrMatrix;
1128afb2bd1cSJunchao Zhang   upTriFactorT->csrMat->num_rows       = upTriFactor->csrMat->num_cols;
1129afb2bd1cSJunchao Zhang   upTriFactorT->csrMat->num_cols       = upTriFactor->csrMat->num_rows;
1130aa372e3fSPaul Mullowney   upTriFactorT->csrMat->num_entries    = upTriFactor->csrMat->num_entries;
1131afb2bd1cSJunchao Zhang   upTriFactorT->csrMat->row_offsets    = new THRUSTINTARRAY32(upTriFactorT->csrMat->num_rows+1);
1132afb2bd1cSJunchao Zhang   upTriFactorT->csrMat->column_indices = new THRUSTINTARRAY32(upTriFactorT->csrMat->num_entries);
1133afb2bd1cSJunchao Zhang   upTriFactorT->csrMat->values         = new THRUSTARRAY(upTriFactorT->csrMat->num_entries);
1134aa372e3fSPaul Mullowney 
1135aa372e3fSPaul Mullowney   /* compute the transpose of the upper triangular factor, i.e. the CSC */
1136afb2bd1cSJunchao Zhang #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1137afb2bd1cSJunchao Zhang   stat = cusparseCsr2cscEx2_bufferSize(cusparseTriFactors->handle,upTriFactor->csrMat->num_rows,
1138afb2bd1cSJunchao Zhang                                 upTriFactor->csrMat->num_cols, upTriFactor->csrMat->num_entries,
1139afb2bd1cSJunchao Zhang                                 upTriFactor->csrMat->values->data().get(),
1140afb2bd1cSJunchao Zhang                                 upTriFactor->csrMat->row_offsets->data().get(),
1141afb2bd1cSJunchao Zhang                                 upTriFactor->csrMat->column_indices->data().get(),
1142afb2bd1cSJunchao Zhang                                 upTriFactorT->csrMat->values->data().get(),
1143afb2bd1cSJunchao Zhang                                 upTriFactorT->csrMat->row_offsets->data().get(), upTriFactorT->csrMat->column_indices->data().get(), cusparse_scalartype,
1144afb2bd1cSJunchao Zhang                                 CUSPARSE_ACTION_NUMERIC,indexBase,
1145afb2bd1cSJunchao Zhang                                 CUSPARSE_CSR2CSC_ALG1, &upTriFactor->csr2cscBufferSize);CHKERRCUSPARSE(stat);
1146afb2bd1cSJunchao Zhang   cerr = cudaMalloc(&upTriFactor->csr2cscBuffer,upTriFactor->csr2cscBufferSize);CHKERRCUDA(cerr);
1147afb2bd1cSJunchao Zhang #endif
1148afb2bd1cSJunchao Zhang 
1149da79fbbcSStefano Zampini   ierr = PetscLogEventBegin(MAT_CUSPARSEGenerateTranspose,A,0,0,0);CHKERRQ(ierr);
1150aa372e3fSPaul Mullowney   stat = cusparse_csr2csc(cusparseTriFactors->handle, upTriFactor->csrMat->num_rows,
1151aa372e3fSPaul Mullowney                           upTriFactor->csrMat->num_cols, upTriFactor->csrMat->num_entries,
1152aa372e3fSPaul Mullowney                           upTriFactor->csrMat->values->data().get(),
1153aa372e3fSPaul Mullowney                           upTriFactor->csrMat->row_offsets->data().get(),
1154aa372e3fSPaul Mullowney                           upTriFactor->csrMat->column_indices->data().get(),
1155aa372e3fSPaul Mullowney                           upTriFactorT->csrMat->values->data().get(),
1156afb2bd1cSJunchao Zhang                         #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1157afb2bd1cSJunchao Zhang                           upTriFactorT->csrMat->row_offsets->data().get(), upTriFactorT->csrMat->column_indices->data().get(), cusparse_scalartype,
1158afb2bd1cSJunchao Zhang                           CUSPARSE_ACTION_NUMERIC, indexBase,
1159afb2bd1cSJunchao Zhang                           CUSPARSE_CSR2CSC_ALG1, upTriFactor->csr2cscBuffer
1160afb2bd1cSJunchao Zhang                         #else
1161afb2bd1cSJunchao Zhang                           upTriFactorT->csrMat->column_indices->data().get(), upTriFactorT->csrMat->row_offsets->data().get(),
1162afb2bd1cSJunchao Zhang                           CUSPARSE_ACTION_NUMERIC, indexBase
1163afb2bd1cSJunchao Zhang                         #endif
1164afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
1165da79fbbcSStefano Zampini   cerr = WaitForCUDA();CHKERRCUDA(cerr);
1166da79fbbcSStefano Zampini   ierr = PetscLogEventBegin(MAT_CUSPARSEGenerateTranspose,A,0,0,0);CHKERRQ(ierr);
1167aa372e3fSPaul Mullowney 
1168afb2bd1cSJunchao Zhang   /* Create the solve analysis information */
1169da79fbbcSStefano Zampini   ierr = PetscLogEventBegin(MAT_CUSPARSESolveAnalysis,A,0,0,0);CHKERRQ(ierr);
1170afb2bd1cSJunchao Zhang   stat = cusparse_create_analysis_info(&upTriFactorT->solveInfo);CHKERRCUSPARSE(stat);
11711b0a6780SStefano Zampini   #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1172afb2bd1cSJunchao Zhang   stat = cusparse_get_svbuffsize(cusparseTriFactors->handle, upTriFactorT->solveOp,
1173afb2bd1cSJunchao Zhang                                  upTriFactorT->csrMat->num_rows, upTriFactorT->csrMat->num_entries, upTriFactorT->descr,
1174afb2bd1cSJunchao Zhang                                  upTriFactorT->csrMat->values->data().get(), upTriFactorT->csrMat->row_offsets->data().get(),
1175afb2bd1cSJunchao Zhang                                  upTriFactorT->csrMat->column_indices->data().get(), upTriFactorT->solveInfo,
1176afb2bd1cSJunchao Zhang                                  &upTriFactorT->solveBufferSize);CHKERRCUSPARSE(stat);
1177afb2bd1cSJunchao Zhang   cerr = cudaMalloc(&upTriFactorT->solveBuffer,upTriFactorT->solveBufferSize);CHKERRCUDA(cerr);
1178afb2bd1cSJunchao Zhang   #endif
1179afb2bd1cSJunchao Zhang 
1180afb2bd1cSJunchao Zhang   /* perform the solve analysis */
1181aa372e3fSPaul Mullowney   stat = cusparse_analysis(cusparseTriFactors->handle, upTriFactorT->solveOp,
1182afb2bd1cSJunchao Zhang                            upTriFactorT->csrMat->num_rows, upTriFactorT->csrMat->num_entries, upTriFactorT->descr,
1183afb2bd1cSJunchao Zhang                            upTriFactorT->csrMat->values->data().get(), upTriFactorT->csrMat->row_offsets->data().get(),
1184afb2bd1cSJunchao Zhang                            upTriFactorT->csrMat->column_indices->data().get(), upTriFactorT->solveInfo
11851b0a6780SStefano Zampini                           #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1186afb2bd1cSJunchao Zhang                            ,upTriFactorT->solvePolicy, upTriFactorT->solveBuffer
1187afb2bd1cSJunchao Zhang                           #endif
1188afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
1189da79fbbcSStefano Zampini   cerr = WaitForCUDA();CHKERRCUDA(cerr);
1190da79fbbcSStefano Zampini   ierr = PetscLogEventEnd(MAT_CUSPARSESolveAnalysis,A,0,0,0);CHKERRQ(ierr);
1191aa372e3fSPaul Mullowney 
1192da79fbbcSStefano Zampini   /* assign the pointer */
1193aa372e3fSPaul Mullowney   ((Mat_SeqAIJCUSPARSETriFactors*)A->spptr)->upTriFactorPtrTranspose = upTriFactorT;
1194bda325fcSPaul Mullowney   PetscFunctionReturn(0);
1195bda325fcSPaul Mullowney }
1196bda325fcSPaul Mullowney 
1197b175d8bbSPaul Mullowney static PetscErrorCode MatSeqAIJCUSPARSEGenerateTransposeForMult(Mat A)
1198bda325fcSPaul Mullowney {
1199aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSE           *cusparsestruct = (Mat_SeqAIJCUSPARSE*)A->spptr;
1200aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSEMultStruct *matstruct = (Mat_SeqAIJCUSPARSEMultStruct*)cusparsestruct->mat;
1201aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSEMultStruct *matstructT = (Mat_SeqAIJCUSPARSEMultStruct*)cusparsestruct->matTranspose;
1202bda325fcSPaul Mullowney   Mat_SeqAIJ                   *a = (Mat_SeqAIJ*)A->data;
1203bda325fcSPaul Mullowney   cusparseStatus_t             stat;
1204aa372e3fSPaul Mullowney   cusparseIndexBase_t          indexBase;
1205b06137fdSPaul Mullowney   cudaError_t                  err;
120685ba7357SStefano Zampini   PetscErrorCode               ierr;
1207b175d8bbSPaul Mullowney 
1208bda325fcSPaul Mullowney   PetscFunctionBegin;
1209fcdce8c4SStefano Zampini   if (!cusparsestruct->transgen || cusparsestruct->matTranspose || !A->rmap->n || !A->cmap->n) PetscFunctionReturn(0);
121085ba7357SStefano Zampini   ierr = PetscLogEventBegin(MAT_CUSPARSEGenerateTranspose,A,0,0,0);CHKERRQ(ierr);
121185ba7357SStefano Zampini   ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
121285ba7357SStefano Zampini   /* create cusparse matrix */
1213aa372e3fSPaul Mullowney   matstructT = new Mat_SeqAIJCUSPARSEMultStruct;
121457d48284SJunchao Zhang   stat = cusparseCreateMatDescr(&matstructT->descr);CHKERRCUSPARSE(stat);
1215aa372e3fSPaul Mullowney   indexBase = cusparseGetMatIndexBase(matstruct->descr);
121657d48284SJunchao Zhang   stat = cusparseSetMatIndexBase(matstructT->descr, indexBase);CHKERRCUSPARSE(stat);
121757d48284SJunchao Zhang   stat = cusparseSetMatType(matstructT->descr, CUSPARSE_MATRIX_TYPE_GENERAL);CHKERRCUSPARSE(stat);
1218aa372e3fSPaul Mullowney 
1219b06137fdSPaul Mullowney   /* set alpha and beta */
1220afb2bd1cSJunchao Zhang   err = cudaMalloc((void **)&(matstructT->alpha_one),sizeof(PetscScalar));CHKERRCUDA(err);
12217656d835SStefano Zampini   err = cudaMalloc((void **)&(matstructT->beta_zero),sizeof(PetscScalar));CHKERRCUDA(err);
12227656d835SStefano Zampini   err = cudaMalloc((void **)&(matstructT->beta_one), sizeof(PetscScalar));CHKERRCUDA(err);
1223afb2bd1cSJunchao Zhang   err = cudaMemcpy(matstructT->alpha_one,&PETSC_CUSPARSE_ONE, sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(err);
12247656d835SStefano Zampini   err = cudaMemcpy(matstructT->beta_zero,&PETSC_CUSPARSE_ZERO,sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(err);
12257656d835SStefano Zampini   err = cudaMemcpy(matstructT->beta_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(err);
122657d48284SJunchao Zhang   stat = cusparseSetPointerMode(cusparsestruct->handle, CUSPARSE_POINTER_MODE_DEVICE);CHKERRCUSPARSE(stat);
1227b06137fdSPaul Mullowney 
1228aa372e3fSPaul Mullowney   if (cusparsestruct->format==MAT_CUSPARSE_CSR) {
1229aa372e3fSPaul Mullowney     CsrMatrix *matrix = (CsrMatrix*)matstruct->mat;
1230aa372e3fSPaul Mullowney     CsrMatrix *matrixT= new CsrMatrix;
1231554b8892SKarl Rupp     matrixT->num_rows = A->cmap->n;
1232554b8892SKarl Rupp     matrixT->num_cols = A->rmap->n;
1233aa372e3fSPaul Mullowney     matrixT->num_entries = a->nz;
1234a8bd5306SMark Adams     matrixT->row_offsets = new THRUSTINTARRAY32(matrixT->num_rows+1);
1235aa372e3fSPaul Mullowney     matrixT->column_indices = new THRUSTINTARRAY32(a->nz);
1236aa372e3fSPaul Mullowney     matrixT->values = new THRUSTARRAY(a->nz);
1237a3fdcf43SKarl Rupp 
123881902715SJunchao Zhang     cusparsestruct->rowoffsets_gpu = new THRUSTINTARRAY32(A->rmap->n+1);
123981902715SJunchao Zhang     cusparsestruct->rowoffsets_gpu->assign(a->i,a->i+A->rmap->n+1);
1240afb2bd1cSJunchao Zhang 
124181902715SJunchao Zhang     /* compute the transpose, i.e. the CSC */
1242afb2bd1cSJunchao Zhang    #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1243afb2bd1cSJunchao Zhang     stat = cusparseCsr2cscEx2_bufferSize(cusparsestruct->handle, A->rmap->n,
1244afb2bd1cSJunchao Zhang                                   A->cmap->n, matrix->num_entries,
1245afb2bd1cSJunchao Zhang                                   matrix->values->data().get(),
1246afb2bd1cSJunchao Zhang                                   cusparsestruct->rowoffsets_gpu->data().get(),
1247afb2bd1cSJunchao Zhang                                   matrix->column_indices->data().get(),
1248afb2bd1cSJunchao Zhang                                   matrixT->values->data().get(),
1249afb2bd1cSJunchao Zhang                                   matrixT->row_offsets->data().get(), matrixT->column_indices->data().get(), cusparse_scalartype,
1250afb2bd1cSJunchao Zhang                                   CUSPARSE_ACTION_NUMERIC,indexBase,
1251afb2bd1cSJunchao Zhang                                   cusparsestruct->csr2cscAlg, &cusparsestruct->csr2cscBufferSize);CHKERRCUSPARSE(stat);
1252afb2bd1cSJunchao Zhang     err = cudaMalloc(&cusparsestruct->csr2cscBuffer,cusparsestruct->csr2cscBufferSize);CHKERRCUDA(err);
1253afb2bd1cSJunchao Zhang    #endif
1254afb2bd1cSJunchao Zhang 
1255a3fdcf43SKarl Rupp     stat = cusparse_csr2csc(cusparsestruct->handle, A->rmap->n,
1256a3fdcf43SKarl Rupp                             A->cmap->n, matrix->num_entries,
1257aa372e3fSPaul Mullowney                             matrix->values->data().get(),
125881902715SJunchao Zhang                             cusparsestruct->rowoffsets_gpu->data().get(),
1259aa372e3fSPaul Mullowney                             matrix->column_indices->data().get(),
1260aa372e3fSPaul Mullowney                             matrixT->values->data().get(),
1261afb2bd1cSJunchao Zhang                           #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1262afb2bd1cSJunchao Zhang                             matrixT->row_offsets->data().get(), matrixT->column_indices->data().get(), cusparse_scalartype,
1263afb2bd1cSJunchao Zhang                             CUSPARSE_ACTION_NUMERIC,indexBase,
1264afb2bd1cSJunchao Zhang                             cusparsestruct->csr2cscAlg, cusparsestruct->csr2cscBuffer
1265afb2bd1cSJunchao Zhang                           #else
1266afb2bd1cSJunchao Zhang                             matrixT->column_indices->data().get(), matrixT->row_offsets->data().get(),
1267afb2bd1cSJunchao Zhang                             CUSPARSE_ACTION_NUMERIC, indexBase
1268afb2bd1cSJunchao Zhang                           #endif
1269afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
1270aa372e3fSPaul Mullowney     matstructT->mat = matrixT;
1271afb2bd1cSJunchao Zhang 
1272afb2bd1cSJunchao Zhang    #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1273afb2bd1cSJunchao Zhang     stat = cusparseCreateCsr(&matstructT->matDescr,
1274afb2bd1cSJunchao Zhang                              matrixT->num_rows, matrixT->num_cols, matrixT->num_entries,
1275afb2bd1cSJunchao Zhang                              matrixT->row_offsets->data().get(), matrixT->column_indices->data().get(),
1276afb2bd1cSJunchao Zhang                              matrixT->values->data().get(),
1277afb2bd1cSJunchao Zhang                              CUSPARSE_INDEX_32I,CUSPARSE_INDEX_32I, /* row offset, col idx type due to THRUSTINTARRAY32 */
1278afb2bd1cSJunchao Zhang                              indexBase,cusparse_scalartype);CHKERRCUSPARSE(stat);
1279afb2bd1cSJunchao Zhang    #endif
1280aa372e3fSPaul Mullowney   } else if (cusparsestruct->format==MAT_CUSPARSE_ELL || cusparsestruct->format==MAT_CUSPARSE_HYB) {
1281afb2bd1cSJunchao Zhang    #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1282afb2bd1cSJunchao Zhang     SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"MAT_CUSPARSE_ELL and MAT_CUSPARSE_HYB are not supported since CUDA-11.0");
1283afb2bd1cSJunchao Zhang    #else
1284aa372e3fSPaul Mullowney     CsrMatrix *temp  = new CsrMatrix;
128551c6d536SStefano Zampini     CsrMatrix *tempT = new CsrMatrix;
128651c6d536SStefano Zampini     /* First convert HYB to CSR */
1287aa372e3fSPaul Mullowney     temp->num_rows = A->rmap->n;
1288aa372e3fSPaul Mullowney     temp->num_cols = A->cmap->n;
1289aa372e3fSPaul Mullowney     temp->num_entries = a->nz;
1290aa372e3fSPaul Mullowney     temp->row_offsets = new THRUSTINTARRAY32(A->rmap->n+1);
1291aa372e3fSPaul Mullowney     temp->column_indices = new THRUSTINTARRAY32(a->nz);
1292aa372e3fSPaul Mullowney     temp->values = new THRUSTARRAY(a->nz);
1293aa372e3fSPaul Mullowney 
1294aa372e3fSPaul Mullowney     stat = cusparse_hyb2csr(cusparsestruct->handle,
1295aa372e3fSPaul Mullowney                             matstruct->descr, (cusparseHybMat_t)matstruct->mat,
1296aa372e3fSPaul Mullowney                             temp->values->data().get(),
1297aa372e3fSPaul Mullowney                             temp->row_offsets->data().get(),
129857d48284SJunchao Zhang                             temp->column_indices->data().get());CHKERRCUSPARSE(stat);
1299aa372e3fSPaul Mullowney 
1300aa372e3fSPaul Mullowney     /* Next, convert CSR to CSC (i.e. the matrix transpose) */
1301aa372e3fSPaul Mullowney     tempT->num_rows = A->rmap->n;
1302aa372e3fSPaul Mullowney     tempT->num_cols = A->cmap->n;
1303aa372e3fSPaul Mullowney     tempT->num_entries = a->nz;
1304aa372e3fSPaul Mullowney     tempT->row_offsets = new THRUSTINTARRAY32(A->rmap->n+1);
1305aa372e3fSPaul Mullowney     tempT->column_indices = new THRUSTINTARRAY32(a->nz);
1306aa372e3fSPaul Mullowney     tempT->values = new THRUSTARRAY(a->nz);
1307aa372e3fSPaul Mullowney 
1308aa372e3fSPaul Mullowney     stat = cusparse_csr2csc(cusparsestruct->handle, temp->num_rows,
1309aa372e3fSPaul Mullowney                             temp->num_cols, temp->num_entries,
1310aa372e3fSPaul Mullowney                             temp->values->data().get(),
1311aa372e3fSPaul Mullowney                             temp->row_offsets->data().get(),
1312aa372e3fSPaul Mullowney                             temp->column_indices->data().get(),
1313aa372e3fSPaul Mullowney                             tempT->values->data().get(),
1314aa372e3fSPaul Mullowney                             tempT->column_indices->data().get(),
1315aa372e3fSPaul Mullowney                             tempT->row_offsets->data().get(),
131657d48284SJunchao Zhang                             CUSPARSE_ACTION_NUMERIC, indexBase);CHKERRCUSPARSE(stat);
1317aa372e3fSPaul Mullowney 
1318aa372e3fSPaul Mullowney     /* Last, convert CSC to HYB */
1319aa372e3fSPaul Mullowney     cusparseHybMat_t hybMat;
132057d48284SJunchao Zhang     stat = cusparseCreateHybMat(&hybMat);CHKERRCUSPARSE(stat);
1321aa372e3fSPaul Mullowney     cusparseHybPartition_t partition = cusparsestruct->format==MAT_CUSPARSE_ELL ?
1322aa372e3fSPaul Mullowney       CUSPARSE_HYB_PARTITION_MAX : CUSPARSE_HYB_PARTITION_AUTO;
1323aa372e3fSPaul Mullowney     stat = cusparse_csr2hyb(cusparsestruct->handle, A->rmap->n, A->cmap->n,
1324aa372e3fSPaul Mullowney                             matstructT->descr, tempT->values->data().get(),
1325aa372e3fSPaul Mullowney                             tempT->row_offsets->data().get(),
1326aa372e3fSPaul Mullowney                             tempT->column_indices->data().get(),
132757d48284SJunchao Zhang                             hybMat, 0, partition);CHKERRCUSPARSE(stat);
1328aa372e3fSPaul Mullowney 
1329aa372e3fSPaul Mullowney     /* assign the pointer */
1330aa372e3fSPaul Mullowney     matstructT->mat = hybMat;
1331aa372e3fSPaul Mullowney     /* delete temporaries */
1332aa372e3fSPaul Mullowney     if (tempT) {
1333aa372e3fSPaul Mullowney       if (tempT->values) delete (THRUSTARRAY*) tempT->values;
1334aa372e3fSPaul Mullowney       if (tempT->column_indices) delete (THRUSTINTARRAY32*) tempT->column_indices;
1335aa372e3fSPaul Mullowney       if (tempT->row_offsets) delete (THRUSTINTARRAY32*) tempT->row_offsets;
1336aa372e3fSPaul Mullowney       delete (CsrMatrix*) tempT;
1337087f3262SPaul Mullowney     }
1338aa372e3fSPaul Mullowney     if (temp) {
1339aa372e3fSPaul Mullowney       if (temp->values) delete (THRUSTARRAY*) temp->values;
1340aa372e3fSPaul Mullowney       if (temp->column_indices) delete (THRUSTINTARRAY32*) temp->column_indices;
1341aa372e3fSPaul Mullowney       if (temp->row_offsets) delete (THRUSTINTARRAY32*) temp->row_offsets;
1342aa372e3fSPaul Mullowney       delete (CsrMatrix*) temp;
1343aa372e3fSPaul Mullowney     }
1344afb2bd1cSJunchao Zhang    #endif
1345aa372e3fSPaul Mullowney   }
134605035670SJunchao Zhang   err  = WaitForCUDA();CHKERRCUDA(err);
134785ba7357SStefano Zampini   ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
134885ba7357SStefano Zampini   ierr = PetscLogEventEnd(MAT_CUSPARSEGenerateTranspose,A,0,0,0);CHKERRQ(ierr);
1349213423ffSJunchao Zhang   /* the compressed row indices is not used for matTranspose */
1350213423ffSJunchao Zhang   matstructT->cprowIndices = NULL;
1351aa372e3fSPaul Mullowney   /* assign the pointer */
1352aa372e3fSPaul Mullowney   ((Mat_SeqAIJCUSPARSE*)A->spptr)->matTranspose = matstructT;
1353bda325fcSPaul Mullowney   PetscFunctionReturn(0);
1354bda325fcSPaul Mullowney }
1355bda325fcSPaul Mullowney 
13564e4bbfaaSStefano Zampini /* Why do we need to analyze the tranposed matrix again? Can't we just use op(A) = CUSPARSE_OPERATION_TRANSPOSE in MatSolve_SeqAIJCUSPARSE? */
13576fa9248bSJed Brown static PetscErrorCode MatSolveTranspose_SeqAIJCUSPARSE(Mat A,Vec bb,Vec xx)
1358bda325fcSPaul Mullowney {
1359c41cb2e2SAlejandro Lamas Daviña   PetscInt                              n = xx->map->n;
1360465f34aeSAlejandro Lamas Daviña   const PetscScalar                     *barray;
1361465f34aeSAlejandro Lamas Daviña   PetscScalar                           *xarray;
1362465f34aeSAlejandro Lamas Daviña   thrust::device_ptr<const PetscScalar> bGPU;
1363465f34aeSAlejandro Lamas Daviña   thrust::device_ptr<PetscScalar>       xGPU;
1364bda325fcSPaul Mullowney   cusparseStatus_t                      stat;
1365bda325fcSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactors          *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)A->spptr;
1366aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct     *loTriFactorT = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->loTriFactorPtrTranspose;
1367aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct     *upTriFactorT = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->upTriFactorPtrTranspose;
1368aa372e3fSPaul Mullowney   THRUSTARRAY                           *tempGPU = (THRUSTARRAY*)cusparseTriFactors->workVector;
1369b175d8bbSPaul Mullowney   PetscErrorCode                        ierr;
137057d48284SJunchao Zhang   cudaError_t                           cerr;
1371bda325fcSPaul Mullowney 
1372bda325fcSPaul Mullowney   PetscFunctionBegin;
1373aa372e3fSPaul Mullowney   /* Analyze the matrix and create the transpose ... on the fly */
1374aa372e3fSPaul Mullowney   if (!loTriFactorT && !upTriFactorT) {
1375bda325fcSPaul Mullowney     ierr = MatSeqAIJCUSPARSEAnalyzeTransposeForSolve(A);CHKERRQ(ierr);
1376aa372e3fSPaul Mullowney     loTriFactorT = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->loTriFactorPtrTranspose;
1377aa372e3fSPaul Mullowney     upTriFactorT = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->upTriFactorPtrTranspose;
1378bda325fcSPaul Mullowney   }
1379bda325fcSPaul Mullowney 
1380bda325fcSPaul Mullowney   /* Get the GPU pointers */
1381c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDAGetArrayWrite(xx,&xarray);CHKERRQ(ierr);
1382c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDAGetArrayRead(bb,&barray);CHKERRQ(ierr);
1383c41cb2e2SAlejandro Lamas Daviña   xGPU = thrust::device_pointer_cast(xarray);
1384c41cb2e2SAlejandro Lamas Daviña   bGPU = thrust::device_pointer_cast(barray);
1385bda325fcSPaul Mullowney 
13867a052e47Shannah_mairs   ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
1387aa372e3fSPaul Mullowney   /* First, reorder with the row permutation */
1388c41cb2e2SAlejandro Lamas Daviña   thrust::copy(thrust::make_permutation_iterator(bGPU, cusparseTriFactors->rpermIndices->begin()),
1389c41cb2e2SAlejandro Lamas Daviña                thrust::make_permutation_iterator(bGPU+n, cusparseTriFactors->rpermIndices->end()),
1390c41cb2e2SAlejandro Lamas Daviña                xGPU);
1391aa372e3fSPaul Mullowney 
1392aa372e3fSPaul Mullowney   /* First, solve U */
1393aa372e3fSPaul Mullowney   stat = cusparse_solve(cusparseTriFactors->handle, upTriFactorT->solveOp,
1394afb2bd1cSJunchao Zhang                         upTriFactorT->csrMat->num_rows,
13951b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1396afb2bd1cSJunchao Zhang                         upTriFactorT->csrMat->num_entries,
1397afb2bd1cSJunchao Zhang                       #endif
1398afb2bd1cSJunchao Zhang                         &PETSC_CUSPARSE_ONE, upTriFactorT->descr,
1399aa372e3fSPaul Mullowney                         upTriFactorT->csrMat->values->data().get(),
1400aa372e3fSPaul Mullowney                         upTriFactorT->csrMat->row_offsets->data().get(),
1401aa372e3fSPaul Mullowney                         upTriFactorT->csrMat->column_indices->data().get(),
1402aa372e3fSPaul Mullowney                         upTriFactorT->solveInfo,
1403afb2bd1cSJunchao Zhang                         xarray, tempGPU->data().get()
14041b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1405afb2bd1cSJunchao Zhang                         ,upTriFactorT->solvePolicy, upTriFactorT->solveBuffer
1406afb2bd1cSJunchao Zhang                       #endif
1407afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
1408aa372e3fSPaul Mullowney 
1409aa372e3fSPaul Mullowney   /* Then, solve L */
1410aa372e3fSPaul Mullowney   stat = cusparse_solve(cusparseTriFactors->handle, loTriFactorT->solveOp,
1411afb2bd1cSJunchao Zhang                         loTriFactorT->csrMat->num_rows,
14121b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1413afb2bd1cSJunchao Zhang                         loTriFactorT->csrMat->num_entries,
1414afb2bd1cSJunchao Zhang                       #endif
1415afb2bd1cSJunchao Zhang                         &PETSC_CUSPARSE_ONE, loTriFactorT->descr,
1416aa372e3fSPaul Mullowney                         loTriFactorT->csrMat->values->data().get(),
1417aa372e3fSPaul Mullowney                         loTriFactorT->csrMat->row_offsets->data().get(),
1418aa372e3fSPaul Mullowney                         loTriFactorT->csrMat->column_indices->data().get(),
1419aa372e3fSPaul Mullowney                         loTriFactorT->solveInfo,
1420afb2bd1cSJunchao Zhang                         tempGPU->data().get(), xarray
14211b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1422afb2bd1cSJunchao Zhang                         ,loTriFactorT->solvePolicy, loTriFactorT->solveBuffer
1423afb2bd1cSJunchao Zhang                       #endif
1424afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
1425aa372e3fSPaul Mullowney 
1426aa372e3fSPaul Mullowney   /* Last, copy the solution, xGPU, into a temporary with the column permutation ... can't be done in place. */
1427c41cb2e2SAlejandro Lamas Daviña   thrust::copy(thrust::make_permutation_iterator(xGPU, cusparseTriFactors->cpermIndices->begin()),
1428c41cb2e2SAlejandro Lamas Daviña                thrust::make_permutation_iterator(xGPU+n, cusparseTriFactors->cpermIndices->end()),
1429aa372e3fSPaul Mullowney                tempGPU->begin());
1430aa372e3fSPaul Mullowney 
1431aa372e3fSPaul Mullowney   /* Copy the temporary to the full solution. */
1432c41cb2e2SAlejandro Lamas Daviña   thrust::copy(tempGPU->begin(), tempGPU->end(), xGPU);
1433bda325fcSPaul Mullowney 
1434bda325fcSPaul Mullowney   /* restore */
1435c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDARestoreArrayRead(bb,&barray);CHKERRQ(ierr);
1436c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDARestoreArrayWrite(xx,&xarray);CHKERRQ(ierr);
143705035670SJunchao Zhang   cerr = WaitForCUDA();CHKERRCUDA(cerr);
1438661c2d29Shannah_mairs   ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
1439958c4211Shannah_mairs   ierr = PetscLogGpuFlops(2.0*cusparseTriFactors->nnz - A->cmap->n);CHKERRQ(ierr);
1440bda325fcSPaul Mullowney   PetscFunctionReturn(0);
1441bda325fcSPaul Mullowney }
1442bda325fcSPaul Mullowney 
14436fa9248bSJed Brown static PetscErrorCode MatSolveTranspose_SeqAIJCUSPARSE_NaturalOrdering(Mat A,Vec bb,Vec xx)
1444bda325fcSPaul Mullowney {
1445465f34aeSAlejandro Lamas Daviña   const PetscScalar                 *barray;
1446465f34aeSAlejandro Lamas Daviña   PetscScalar                       *xarray;
1447bda325fcSPaul Mullowney   cusparseStatus_t                  stat;
1448bda325fcSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactors      *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)A->spptr;
1449aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct *loTriFactorT = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->loTriFactorPtrTranspose;
1450aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct *upTriFactorT = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->upTriFactorPtrTranspose;
1451aa372e3fSPaul Mullowney   THRUSTARRAY                       *tempGPU = (THRUSTARRAY*)cusparseTriFactors->workVector;
1452b175d8bbSPaul Mullowney   PetscErrorCode                    ierr;
145357d48284SJunchao Zhang   cudaError_t                       cerr;
1454bda325fcSPaul Mullowney 
1455bda325fcSPaul Mullowney   PetscFunctionBegin;
1456aa372e3fSPaul Mullowney   /* Analyze the matrix and create the transpose ... on the fly */
1457aa372e3fSPaul Mullowney   if (!loTriFactorT && !upTriFactorT) {
1458bda325fcSPaul Mullowney     ierr = MatSeqAIJCUSPARSEAnalyzeTransposeForSolve(A);CHKERRQ(ierr);
1459aa372e3fSPaul Mullowney     loTriFactorT       = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->loTriFactorPtrTranspose;
1460aa372e3fSPaul Mullowney     upTriFactorT       = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->upTriFactorPtrTranspose;
1461bda325fcSPaul Mullowney   }
1462bda325fcSPaul Mullowney 
1463bda325fcSPaul Mullowney   /* Get the GPU pointers */
1464c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDAGetArrayWrite(xx,&xarray);CHKERRQ(ierr);
1465c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDAGetArrayRead(bb,&barray);CHKERRQ(ierr);
1466bda325fcSPaul Mullowney 
14677a052e47Shannah_mairs   ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
1468aa372e3fSPaul Mullowney   /* First, solve U */
1469aa372e3fSPaul Mullowney   stat = cusparse_solve(cusparseTriFactors->handle, upTriFactorT->solveOp,
1470afb2bd1cSJunchao Zhang                         upTriFactorT->csrMat->num_rows,
14711b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1472afb2bd1cSJunchao Zhang                         upTriFactorT->csrMat->num_entries,
1473afb2bd1cSJunchao Zhang                       #endif
1474afb2bd1cSJunchao Zhang                         &PETSC_CUSPARSE_ONE, upTriFactorT->descr,
1475aa372e3fSPaul Mullowney                         upTriFactorT->csrMat->values->data().get(),
1476aa372e3fSPaul Mullowney                         upTriFactorT->csrMat->row_offsets->data().get(),
1477aa372e3fSPaul Mullowney                         upTriFactorT->csrMat->column_indices->data().get(),
1478aa372e3fSPaul Mullowney                         upTriFactorT->solveInfo,
1479afb2bd1cSJunchao Zhang                         barray, tempGPU->data().get()
14801b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1481afb2bd1cSJunchao Zhang                         ,upTriFactorT->solvePolicy, upTriFactorT->solveBuffer
1482afb2bd1cSJunchao Zhang                       #endif
1483afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
1484aa372e3fSPaul Mullowney 
1485aa372e3fSPaul Mullowney   /* Then, solve L */
1486aa372e3fSPaul Mullowney   stat = cusparse_solve(cusparseTriFactors->handle, loTriFactorT->solveOp,
1487afb2bd1cSJunchao Zhang                         loTriFactorT->csrMat->num_rows,
14881b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1489afb2bd1cSJunchao Zhang                         loTriFactorT->csrMat->num_entries,
1490afb2bd1cSJunchao Zhang                       #endif
1491afb2bd1cSJunchao Zhang                         &PETSC_CUSPARSE_ONE, loTriFactorT->descr,
1492aa372e3fSPaul Mullowney                         loTriFactorT->csrMat->values->data().get(),
1493aa372e3fSPaul Mullowney                         loTriFactorT->csrMat->row_offsets->data().get(),
1494aa372e3fSPaul Mullowney                         loTriFactorT->csrMat->column_indices->data().get(),
1495aa372e3fSPaul Mullowney                         loTriFactorT->solveInfo,
1496afb2bd1cSJunchao Zhang                         tempGPU->data().get(), xarray
14971b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1498afb2bd1cSJunchao Zhang                         ,loTriFactorT->solvePolicy, loTriFactorT->solveBuffer
1499afb2bd1cSJunchao Zhang                       #endif
1500afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
1501bda325fcSPaul Mullowney 
1502bda325fcSPaul Mullowney   /* restore */
1503c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDARestoreArrayRead(bb,&barray);CHKERRQ(ierr);
1504c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDARestoreArrayWrite(xx,&xarray);CHKERRQ(ierr);
150505035670SJunchao Zhang   cerr = WaitForCUDA();CHKERRCUDA(cerr);
1506661c2d29Shannah_mairs   ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
1507958c4211Shannah_mairs   ierr = PetscLogGpuFlops(2.0*cusparseTriFactors->nnz - A->cmap->n);CHKERRQ(ierr);
1508bda325fcSPaul Mullowney   PetscFunctionReturn(0);
1509bda325fcSPaul Mullowney }
1510bda325fcSPaul Mullowney 
15116fa9248bSJed Brown static PetscErrorCode MatSolve_SeqAIJCUSPARSE(Mat A,Vec bb,Vec xx)
15129ae82921SPaul Mullowney {
1513465f34aeSAlejandro Lamas Daviña   const PetscScalar                     *barray;
1514465f34aeSAlejandro Lamas Daviña   PetscScalar                           *xarray;
1515465f34aeSAlejandro Lamas Daviña   thrust::device_ptr<const PetscScalar> bGPU;
1516465f34aeSAlejandro Lamas Daviña   thrust::device_ptr<PetscScalar>       xGPU;
15179ae82921SPaul Mullowney   cusparseStatus_t                      stat;
15189ae82921SPaul Mullowney   Mat_SeqAIJCUSPARSETriFactors          *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)A->spptr;
1519aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct     *loTriFactor = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->loTriFactorPtr;
1520aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct     *upTriFactor = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->upTriFactorPtr;
1521aa372e3fSPaul Mullowney   THRUSTARRAY                           *tempGPU = (THRUSTARRAY*)cusparseTriFactors->workVector;
1522b175d8bbSPaul Mullowney   PetscErrorCode                        ierr;
152357d48284SJunchao Zhang   cudaError_t                           cerr;
15249ae82921SPaul Mullowney 
15259ae82921SPaul Mullowney   PetscFunctionBegin;
1526ebc8f436SDominic Meiser 
1527e057df02SPaul Mullowney   /* Get the GPU pointers */
1528c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDAGetArrayWrite(xx,&xarray);CHKERRQ(ierr);
1529c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDAGetArrayRead(bb,&barray);CHKERRQ(ierr);
1530c41cb2e2SAlejandro Lamas Daviña   xGPU = thrust::device_pointer_cast(xarray);
1531c41cb2e2SAlejandro Lamas Daviña   bGPU = thrust::device_pointer_cast(barray);
15329ae82921SPaul Mullowney 
15337a052e47Shannah_mairs   ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
1534aa372e3fSPaul Mullowney   /* First, reorder with the row permutation */
1535c41cb2e2SAlejandro Lamas Daviña   thrust::copy(thrust::make_permutation_iterator(bGPU, cusparseTriFactors->rpermIndices->begin()),
1536c41cb2e2SAlejandro Lamas Daviña                thrust::make_permutation_iterator(bGPU, cusparseTriFactors->rpermIndices->end()),
15374e4bbfaaSStefano Zampini                tempGPU->begin());
1538aa372e3fSPaul Mullowney 
1539aa372e3fSPaul Mullowney   /* Next, solve L */
1540aa372e3fSPaul Mullowney   stat = cusparse_solve(cusparseTriFactors->handle, loTriFactor->solveOp,
1541afb2bd1cSJunchao Zhang                         loTriFactor->csrMat->num_rows,
15421b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1543afb2bd1cSJunchao Zhang                         loTriFactor->csrMat->num_entries,
1544afb2bd1cSJunchao Zhang                       #endif
1545afb2bd1cSJunchao Zhang                         &PETSC_CUSPARSE_ONE, loTriFactor->descr,
1546aa372e3fSPaul Mullowney                         loTriFactor->csrMat->values->data().get(),
1547aa372e3fSPaul Mullowney                         loTriFactor->csrMat->row_offsets->data().get(),
1548aa372e3fSPaul Mullowney                         loTriFactor->csrMat->column_indices->data().get(),
1549aa372e3fSPaul Mullowney                         loTriFactor->solveInfo,
1550afb2bd1cSJunchao Zhang                         tempGPU->data().get(), xarray
15511b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1552afb2bd1cSJunchao Zhang                         ,loTriFactor->solvePolicy, loTriFactor->solveBuffer
1553afb2bd1cSJunchao Zhang                       #endif
1554afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
1555aa372e3fSPaul Mullowney 
1556aa372e3fSPaul Mullowney   /* Then, solve U */
1557aa372e3fSPaul Mullowney   stat = cusparse_solve(cusparseTriFactors->handle, upTriFactor->solveOp,
1558afb2bd1cSJunchao Zhang                         upTriFactor->csrMat->num_rows,
15591b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1560afb2bd1cSJunchao Zhang                         upTriFactor->csrMat->num_entries,
1561afb2bd1cSJunchao Zhang                       #endif
1562afb2bd1cSJunchao Zhang                         &PETSC_CUSPARSE_ONE, upTriFactor->descr,
1563aa372e3fSPaul Mullowney                         upTriFactor->csrMat->values->data().get(),
1564aa372e3fSPaul Mullowney                         upTriFactor->csrMat->row_offsets->data().get(),
1565aa372e3fSPaul Mullowney                         upTriFactor->csrMat->column_indices->data().get(),
1566aa372e3fSPaul Mullowney                         upTriFactor->solveInfo,
1567afb2bd1cSJunchao Zhang                         xarray, tempGPU->data().get()
15681b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1569afb2bd1cSJunchao Zhang                         ,upTriFactor->solvePolicy, upTriFactor->solveBuffer
1570afb2bd1cSJunchao Zhang                       #endif
1571afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
1572aa372e3fSPaul Mullowney 
15734e4bbfaaSStefano Zampini   /* Last, reorder with the column permutation */
15744e4bbfaaSStefano Zampini   thrust::copy(thrust::make_permutation_iterator(tempGPU->begin(), cusparseTriFactors->cpermIndices->begin()),
15754e4bbfaaSStefano Zampini                thrust::make_permutation_iterator(tempGPU->begin(), cusparseTriFactors->cpermIndices->end()),
15764e4bbfaaSStefano Zampini                xGPU);
15779ae82921SPaul Mullowney 
1578c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDARestoreArrayRead(bb,&barray);CHKERRQ(ierr);
1579c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDARestoreArrayWrite(xx,&xarray);CHKERRQ(ierr);
158005035670SJunchao Zhang   cerr = WaitForCUDA();CHKERRCUDA(cerr);
1581661c2d29Shannah_mairs   ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
1582958c4211Shannah_mairs   ierr = PetscLogGpuFlops(2.0*cusparseTriFactors->nnz - A->cmap->n);CHKERRQ(ierr);
15839ae82921SPaul Mullowney   PetscFunctionReturn(0);
15849ae82921SPaul Mullowney }
15859ae82921SPaul Mullowney 
15866fa9248bSJed Brown static PetscErrorCode MatSolve_SeqAIJCUSPARSE_NaturalOrdering(Mat A,Vec bb,Vec xx)
15879ae82921SPaul Mullowney {
1588465f34aeSAlejandro Lamas Daviña   const PetscScalar                 *barray;
1589465f34aeSAlejandro Lamas Daviña   PetscScalar                       *xarray;
15909ae82921SPaul Mullowney   cusparseStatus_t                  stat;
15919ae82921SPaul Mullowney   Mat_SeqAIJCUSPARSETriFactors      *cusparseTriFactors = (Mat_SeqAIJCUSPARSETriFactors*)A->spptr;
1592aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct *loTriFactor = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->loTriFactorPtr;
1593aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSETriFactorStruct *upTriFactor = (Mat_SeqAIJCUSPARSETriFactorStruct*)cusparseTriFactors->upTriFactorPtr;
1594aa372e3fSPaul Mullowney   THRUSTARRAY                       *tempGPU = (THRUSTARRAY*)cusparseTriFactors->workVector;
1595b175d8bbSPaul Mullowney   PetscErrorCode                    ierr;
159657d48284SJunchao Zhang   cudaError_t                       cerr;
15979ae82921SPaul Mullowney 
15989ae82921SPaul Mullowney   PetscFunctionBegin;
1599e057df02SPaul Mullowney   /* Get the GPU pointers */
1600c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDAGetArrayWrite(xx,&xarray);CHKERRQ(ierr);
1601c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDAGetArrayRead(bb,&barray);CHKERRQ(ierr);
16029ae82921SPaul Mullowney 
16037a052e47Shannah_mairs   ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
1604aa372e3fSPaul Mullowney   /* First, solve L */
1605aa372e3fSPaul Mullowney   stat = cusparse_solve(cusparseTriFactors->handle, loTriFactor->solveOp,
1606afb2bd1cSJunchao Zhang                         loTriFactor->csrMat->num_rows,
16071b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1608afb2bd1cSJunchao Zhang                         loTriFactor->csrMat->num_entries,
1609afb2bd1cSJunchao Zhang                       #endif
1610afb2bd1cSJunchao Zhang                         &PETSC_CUSPARSE_ONE, loTriFactor->descr,
1611aa372e3fSPaul Mullowney                         loTriFactor->csrMat->values->data().get(),
1612aa372e3fSPaul Mullowney                         loTriFactor->csrMat->row_offsets->data().get(),
1613aa372e3fSPaul Mullowney                         loTriFactor->csrMat->column_indices->data().get(),
1614aa372e3fSPaul Mullowney                         loTriFactor->solveInfo,
1615afb2bd1cSJunchao Zhang                         barray, tempGPU->data().get()
16161b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1617afb2bd1cSJunchao Zhang                         ,loTriFactor->solvePolicy, loTriFactor->solveBuffer
1618afb2bd1cSJunchao Zhang                       #endif
1619afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
1620aa372e3fSPaul Mullowney 
1621aa372e3fSPaul Mullowney   /* Next, solve U */
1622aa372e3fSPaul Mullowney   stat = cusparse_solve(cusparseTriFactors->handle, upTriFactor->solveOp,
1623afb2bd1cSJunchao Zhang                         upTriFactor->csrMat->num_rows,
16241b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1625afb2bd1cSJunchao Zhang                         upTriFactor->csrMat->num_entries,
1626afb2bd1cSJunchao Zhang                       #endif
1627afb2bd1cSJunchao Zhang                         &PETSC_CUSPARSE_ONE, upTriFactor->descr,
1628aa372e3fSPaul Mullowney                         upTriFactor->csrMat->values->data().get(),
1629aa372e3fSPaul Mullowney                         upTriFactor->csrMat->row_offsets->data().get(),
1630aa372e3fSPaul Mullowney                         upTriFactor->csrMat->column_indices->data().get(),
1631aa372e3fSPaul Mullowney                         upTriFactor->solveInfo,
1632afb2bd1cSJunchao Zhang                         tempGPU->data().get(), xarray
16331b0a6780SStefano Zampini                       #if PETSC_PKG_CUDA_VERSION_GE(9,0,0)
1634afb2bd1cSJunchao Zhang                         ,upTriFactor->solvePolicy, upTriFactor->solveBuffer
1635afb2bd1cSJunchao Zhang                       #endif
1636afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
16379ae82921SPaul Mullowney 
1638c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDARestoreArrayRead(bb,&barray);CHKERRQ(ierr);
1639c41cb2e2SAlejandro Lamas Daviña   ierr = VecCUDARestoreArrayWrite(xx,&xarray);CHKERRQ(ierr);
164005035670SJunchao Zhang   cerr = WaitForCUDA();CHKERRCUDA(cerr);
1641661c2d29Shannah_mairs   ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
1642958c4211Shannah_mairs   ierr = PetscLogGpuFlops(2.0*cusparseTriFactors->nnz - A->cmap->n);CHKERRQ(ierr);
16439ae82921SPaul Mullowney   PetscFunctionReturn(0);
16449ae82921SPaul Mullowney }
16459ae82921SPaul Mullowney 
16467e8381f9SStefano Zampini static PetscErrorCode MatSeqAIJCUSPARSECopyFromGPU(Mat A)
16477e8381f9SStefano Zampini {
16487e8381f9SStefano Zampini   Mat_SeqAIJ         *a = (Mat_SeqAIJ*)A->data;
16497e8381f9SStefano Zampini   Mat_SeqAIJCUSPARSE *cusp = (Mat_SeqAIJCUSPARSE*)A->spptr;
16507e8381f9SStefano Zampini   cudaError_t        cerr;
16517e8381f9SStefano Zampini   PetscErrorCode     ierr;
16527e8381f9SStefano Zampini 
16537e8381f9SStefano Zampini   PetscFunctionBegin;
16547e8381f9SStefano Zampini   if (A->offloadmask == PETSC_OFFLOAD_GPU) {
16557e8381f9SStefano Zampini     CsrMatrix *matrix = (CsrMatrix*)cusp->mat->mat;
16567e8381f9SStefano Zampini 
16577e8381f9SStefano Zampini     ierr = PetscLogEventBegin(MAT_CUSPARSECopyFromGPU,A,0,0,0);CHKERRQ(ierr);
16587e8381f9SStefano Zampini     cerr = cudaMemcpy(a->a, matrix->values->data().get(), a->nz*sizeof(PetscScalar), cudaMemcpyDeviceToHost);CHKERRCUDA(cerr);
16597e8381f9SStefano Zampini     cerr = WaitForCUDA();CHKERRCUDA(cerr);
16607e8381f9SStefano Zampini     ierr = PetscLogGpuToCpu(a->nz*sizeof(PetscScalar));CHKERRQ(ierr);
16617e8381f9SStefano Zampini     ierr = PetscLogEventEnd(MAT_CUSPARSECopyFromGPU,A,0,0,0);CHKERRQ(ierr);
16627e8381f9SStefano Zampini     A->offloadmask = PETSC_OFFLOAD_BOTH;
16637e8381f9SStefano Zampini   }
16647e8381f9SStefano Zampini   PetscFunctionReturn(0);
16657e8381f9SStefano Zampini }
16667e8381f9SStefano Zampini 
16677e8381f9SStefano Zampini static PetscErrorCode MatSeqAIJGetArray_SeqAIJCUSPARSE(Mat A,PetscScalar *array[])
16687e8381f9SStefano Zampini {
16697e8381f9SStefano Zampini   Mat_SeqAIJ     *a = (Mat_SeqAIJ*)A->data;
16707e8381f9SStefano Zampini   PetscErrorCode ierr;
16717e8381f9SStefano Zampini 
16727e8381f9SStefano Zampini   PetscFunctionBegin;
16737e8381f9SStefano Zampini   ierr = MatSeqAIJCUSPARSECopyFromGPU(A);CHKERRQ(ierr);
16747e8381f9SStefano Zampini   *array = a->a;
16757e8381f9SStefano Zampini   A->offloadmask = PETSC_OFFLOAD_CPU;
16767e8381f9SStefano Zampini   PetscFunctionReturn(0);
16777e8381f9SStefano Zampini }
16787e8381f9SStefano Zampini 
16796fa9248bSJed Brown static PetscErrorCode MatSeqAIJCUSPARSECopyToGPU(Mat A)
16809ae82921SPaul Mullowney {
1681aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSE           *cusparsestruct = (Mat_SeqAIJCUSPARSE*)A->spptr;
16827c700b8dSJunchao Zhang   Mat_SeqAIJCUSPARSEMultStruct *matstruct = cusparsestruct->mat;
16839ae82921SPaul Mullowney   Mat_SeqAIJ                   *a = (Mat_SeqAIJ*)A->data;
1684213423ffSJunchao Zhang   PetscInt                     m = A->rmap->n,*ii,*ridx,tmp;
16859ae82921SPaul Mullowney   PetscErrorCode               ierr;
1686aa372e3fSPaul Mullowney   cusparseStatus_t             stat;
1687b06137fdSPaul Mullowney   cudaError_t                  err;
16889ae82921SPaul Mullowney 
16899ae82921SPaul Mullowney   PetscFunctionBegin;
1690fcdce8c4SStefano Zampini   if (A->boundtocpu) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"Cannot copy to GPU");
1691c70f7ee4SJunchao Zhang   if (A->offloadmask == PETSC_OFFLOAD_UNALLOCATED || A->offloadmask == PETSC_OFFLOAD_CPU) {
1692fcdce8c4SStefano Zampini     if (A->nonzerostate == cusparsestruct->nonzerostate && cusparsestruct->format == MAT_CUSPARSE_CSR) {
169381902715SJunchao Zhang       /* Copy values only */
1694afb2bd1cSJunchao Zhang       CsrMatrix *matrix,*matrixT;
1695afb2bd1cSJunchao Zhang       matrix = (CsrMatrix*)cusparsestruct->mat->mat;
169685ba7357SStefano Zampini 
169785ba7357SStefano Zampini       ierr = PetscLogEventBegin(MAT_CUSPARSECopyToGPU,A,0,0,0);CHKERRQ(ierr);
1698afb2bd1cSJunchao Zhang       matrix->values->assign(a->a, a->a+a->nz);
169905035670SJunchao Zhang       err  = WaitForCUDA();CHKERRCUDA(err);
17004863603aSSatish Balay       ierr = PetscLogCpuToGpu((a->nz)*sizeof(PetscScalar));CHKERRQ(ierr);
170185ba7357SStefano Zampini       ierr = PetscLogEventEnd(MAT_CUSPARSECopyToGPU,A,0,0,0);CHKERRQ(ierr);
170281902715SJunchao Zhang 
170381902715SJunchao Zhang       /* Update matT when it was built before */
170481902715SJunchao Zhang       if (cusparsestruct->matTranspose) {
170581902715SJunchao Zhang         cusparseIndexBase_t indexBase = cusparseGetMatIndexBase(cusparsestruct->mat->descr);
1706afb2bd1cSJunchao Zhang         matrixT = (CsrMatrix*)cusparsestruct->matTranspose->mat;
170785ba7357SStefano Zampini         ierr = PetscLogEventBegin(MAT_CUSPARSEGenerateTranspose,A,0,0,0);CHKERRQ(ierr);
170881902715SJunchao Zhang         stat = cusparse_csr2csc(cusparsestruct->handle, A->rmap->n,
1709afb2bd1cSJunchao Zhang                             A->cmap->n, matrix->num_entries,
1710afb2bd1cSJunchao Zhang                             matrix->values->data().get(),
171181902715SJunchao Zhang                             cusparsestruct->rowoffsets_gpu->data().get(),
1712afb2bd1cSJunchao Zhang                             matrix->column_indices->data().get(),
1713afb2bd1cSJunchao Zhang                             matrixT->values->data().get(),
1714afb2bd1cSJunchao Zhang                           #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1715afb2bd1cSJunchao Zhang                             matrixT->row_offsets->data().get(), matrixT->column_indices->data().get(), cusparse_scalartype,
1716afb2bd1cSJunchao Zhang                             CUSPARSE_ACTION_NUMERIC,indexBase,
1717afb2bd1cSJunchao Zhang                             cusparsestruct->csr2cscAlg, cusparsestruct->csr2cscBuffer
1718afb2bd1cSJunchao Zhang                           #else
1719afb2bd1cSJunchao Zhang                             matrixT->column_indices->data().get(), matrixT->row_offsets->data().get(),
1720afb2bd1cSJunchao Zhang                             CUSPARSE_ACTION_NUMERIC, indexBase
1721afb2bd1cSJunchao Zhang                           #endif
1722afb2bd1cSJunchao Zhang );CHKERRCUSPARSE(stat);
172305035670SJunchao Zhang         err  = WaitForCUDA();CHKERRCUDA(err);
172485ba7357SStefano Zampini         ierr = PetscLogEventEnd(MAT_CUSPARSEGenerateTranspose,A,0,0,0);CHKERRQ(ierr);
172581902715SJunchao Zhang       }
172634d6c7a5SJose E. Roman     } else {
172785ba7357SStefano Zampini       ierr = PetscLogEventBegin(MAT_CUSPARSECopyToGPU,A,0,0,0);CHKERRQ(ierr);
17287c700b8dSJunchao Zhang       ierr = MatSeqAIJCUSPARSEMultStruct_Destroy(&cusparsestruct->mat,cusparsestruct->format);CHKERRQ(ierr);
17297c700b8dSJunchao Zhang       ierr = MatSeqAIJCUSPARSEMultStruct_Destroy(&cusparsestruct->matTranspose,cusparsestruct->format);CHKERRQ(ierr);
17307c700b8dSJunchao Zhang       delete cusparsestruct->workVector;
173181902715SJunchao Zhang       delete cusparsestruct->rowoffsets_gpu;
17329ae82921SPaul Mullowney       try {
17339ae82921SPaul Mullowney         if (a->compressedrow.use) {
17349ae82921SPaul Mullowney           m    = a->compressedrow.nrows;
17359ae82921SPaul Mullowney           ii   = a->compressedrow.i;
17369ae82921SPaul Mullowney           ridx = a->compressedrow.rindex;
17379ae82921SPaul Mullowney         } else {
1738213423ffSJunchao Zhang           m    = A->rmap->n;
1739213423ffSJunchao Zhang           ii   = a->i;
1740e6e9a74fSStefano Zampini           ridx = NULL;
17419ae82921SPaul Mullowney         }
1742213423ffSJunchao Zhang         cusparsestruct->nrows = m;
17439ae82921SPaul Mullowney 
174485ba7357SStefano Zampini         /* create cusparse matrix */
1745aa372e3fSPaul Mullowney         matstruct = new Mat_SeqAIJCUSPARSEMultStruct;
174657d48284SJunchao Zhang         stat = cusparseCreateMatDescr(&matstruct->descr);CHKERRCUSPARSE(stat);
174757d48284SJunchao Zhang         stat = cusparseSetMatIndexBase(matstruct->descr, CUSPARSE_INDEX_BASE_ZERO);CHKERRCUSPARSE(stat);
174857d48284SJunchao Zhang         stat = cusparseSetMatType(matstruct->descr, CUSPARSE_MATRIX_TYPE_GENERAL);CHKERRCUSPARSE(stat);
17499ae82921SPaul Mullowney 
1750afb2bd1cSJunchao Zhang         err = cudaMalloc((void **)&(matstruct->alpha_one),sizeof(PetscScalar));CHKERRCUDA(err);
17517656d835SStefano Zampini         err = cudaMalloc((void **)&(matstruct->beta_zero),sizeof(PetscScalar));CHKERRCUDA(err);
17527656d835SStefano Zampini         err = cudaMalloc((void **)&(matstruct->beta_one), sizeof(PetscScalar));CHKERRCUDA(err);
1753afb2bd1cSJunchao Zhang         err = cudaMemcpy(matstruct->alpha_one,&PETSC_CUSPARSE_ONE, sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(err);
17547656d835SStefano Zampini         err = cudaMemcpy(matstruct->beta_zero,&PETSC_CUSPARSE_ZERO,sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(err);
17557656d835SStefano Zampini         err = cudaMemcpy(matstruct->beta_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(err);
175657d48284SJunchao Zhang         stat = cusparseSetPointerMode(cusparsestruct->handle, CUSPARSE_POINTER_MODE_DEVICE);CHKERRCUSPARSE(stat);
1757b06137fdSPaul Mullowney 
1758aa372e3fSPaul Mullowney         /* Build a hybrid/ellpack matrix if this option is chosen for the storage */
1759aa372e3fSPaul Mullowney         if (cusparsestruct->format==MAT_CUSPARSE_CSR) {
1760aa372e3fSPaul Mullowney           /* set the matrix */
1761afb2bd1cSJunchao Zhang           CsrMatrix *mat= new CsrMatrix;
1762afb2bd1cSJunchao Zhang           mat->num_rows = m;
1763afb2bd1cSJunchao Zhang           mat->num_cols = A->cmap->n;
1764afb2bd1cSJunchao Zhang           mat->num_entries = a->nz;
1765afb2bd1cSJunchao Zhang           mat->row_offsets = new THRUSTINTARRAY32(m+1);
1766afb2bd1cSJunchao Zhang           mat->row_offsets->assign(ii, ii + m+1);
17679ae82921SPaul Mullowney 
1768afb2bd1cSJunchao Zhang           mat->column_indices = new THRUSTINTARRAY32(a->nz);
1769afb2bd1cSJunchao Zhang           mat->column_indices->assign(a->j, a->j+a->nz);
1770aa372e3fSPaul Mullowney 
1771afb2bd1cSJunchao Zhang           mat->values = new THRUSTARRAY(a->nz);
1772afb2bd1cSJunchao Zhang           mat->values->assign(a->a, a->a+a->nz);
1773aa372e3fSPaul Mullowney 
1774aa372e3fSPaul Mullowney           /* assign the pointer */
1775afb2bd1cSJunchao Zhang           matstruct->mat = mat;
1776afb2bd1cSJunchao Zhang          #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1777afb2bd1cSJunchao Zhang           if (mat->num_rows) { /* cusparse errors on empty matrices! */
1778afb2bd1cSJunchao Zhang             stat = cusparseCreateCsr(&matstruct->matDescr,
1779afb2bd1cSJunchao Zhang                                     mat->num_rows, mat->num_cols, mat->num_entries,
1780afb2bd1cSJunchao Zhang                                     mat->row_offsets->data().get(), mat->column_indices->data().get(),
1781afb2bd1cSJunchao Zhang                                     mat->values->data().get(),
1782afb2bd1cSJunchao Zhang                                     CUSPARSE_INDEX_32I,CUSPARSE_INDEX_32I, /* row offset, col idx types due to THRUSTINTARRAY32 */
1783afb2bd1cSJunchao Zhang                                     CUSPARSE_INDEX_BASE_ZERO,cusparse_scalartype);CHKERRCUSPARSE(stat);
1784afb2bd1cSJunchao Zhang           }
1785afb2bd1cSJunchao Zhang          #endif
1786aa372e3fSPaul Mullowney         } else if (cusparsestruct->format==MAT_CUSPARSE_ELL || cusparsestruct->format==MAT_CUSPARSE_HYB) {
1787afb2bd1cSJunchao Zhang          #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1788afb2bd1cSJunchao Zhang           SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"MAT_CUSPARSE_ELL and MAT_CUSPARSE_HYB are not supported since CUDA-11.0");
1789afb2bd1cSJunchao Zhang          #else
1790afb2bd1cSJunchao Zhang           CsrMatrix *mat= new CsrMatrix;
1791afb2bd1cSJunchao Zhang           mat->num_rows = m;
1792afb2bd1cSJunchao Zhang           mat->num_cols = A->cmap->n;
1793afb2bd1cSJunchao Zhang           mat->num_entries = a->nz;
1794afb2bd1cSJunchao Zhang           mat->row_offsets = new THRUSTINTARRAY32(m+1);
1795afb2bd1cSJunchao Zhang           mat->row_offsets->assign(ii, ii + m+1);
1796aa372e3fSPaul Mullowney 
1797afb2bd1cSJunchao Zhang           mat->column_indices = new THRUSTINTARRAY32(a->nz);
1798afb2bd1cSJunchao Zhang           mat->column_indices->assign(a->j, a->j+a->nz);
1799aa372e3fSPaul Mullowney 
1800afb2bd1cSJunchao Zhang           mat->values = new THRUSTARRAY(a->nz);
1801afb2bd1cSJunchao Zhang           mat->values->assign(a->a, a->a+a->nz);
1802aa372e3fSPaul Mullowney 
1803aa372e3fSPaul Mullowney           cusparseHybMat_t hybMat;
180457d48284SJunchao Zhang           stat = cusparseCreateHybMat(&hybMat);CHKERRCUSPARSE(stat);
1805aa372e3fSPaul Mullowney           cusparseHybPartition_t partition = cusparsestruct->format==MAT_CUSPARSE_ELL ?
1806aa372e3fSPaul Mullowney             CUSPARSE_HYB_PARTITION_MAX : CUSPARSE_HYB_PARTITION_AUTO;
1807afb2bd1cSJunchao Zhang           stat = cusparse_csr2hyb(cusparsestruct->handle, mat->num_rows, mat->num_cols,
1808afb2bd1cSJunchao Zhang               matstruct->descr, mat->values->data().get(),
1809afb2bd1cSJunchao Zhang               mat->row_offsets->data().get(),
1810afb2bd1cSJunchao Zhang               mat->column_indices->data().get(),
181157d48284SJunchao Zhang               hybMat, 0, partition);CHKERRCUSPARSE(stat);
1812aa372e3fSPaul Mullowney           /* assign the pointer */
1813aa372e3fSPaul Mullowney           matstruct->mat = hybMat;
1814aa372e3fSPaul Mullowney 
1815afb2bd1cSJunchao Zhang           if (mat) {
1816afb2bd1cSJunchao Zhang             if (mat->values) delete (THRUSTARRAY*)mat->values;
1817afb2bd1cSJunchao Zhang             if (mat->column_indices) delete (THRUSTINTARRAY32*)mat->column_indices;
1818afb2bd1cSJunchao Zhang             if (mat->row_offsets) delete (THRUSTINTARRAY32*)mat->row_offsets;
1819afb2bd1cSJunchao Zhang             delete (CsrMatrix*)mat;
1820087f3262SPaul Mullowney           }
1821afb2bd1cSJunchao Zhang          #endif
1822087f3262SPaul Mullowney         }
1823ca45077fSPaul Mullowney 
1824aa372e3fSPaul Mullowney         /* assign the compressed row indices */
1825213423ffSJunchao Zhang         if (a->compressedrow.use) {
1826213423ffSJunchao Zhang           cusparsestruct->workVector = new THRUSTARRAY(m);
1827aa372e3fSPaul Mullowney           matstruct->cprowIndices    = new THRUSTINTARRAY(m);
1828aa372e3fSPaul Mullowney           matstruct->cprowIndices->assign(ridx,ridx+m);
1829213423ffSJunchao Zhang           tmp = m;
1830213423ffSJunchao Zhang         } else {
1831213423ffSJunchao Zhang           cusparsestruct->workVector = NULL;
1832213423ffSJunchao Zhang           matstruct->cprowIndices    = NULL;
1833213423ffSJunchao Zhang           tmp = 0;
1834213423ffSJunchao Zhang         }
1835213423ffSJunchao Zhang         ierr = PetscLogCpuToGpu(((m+1)+(a->nz))*sizeof(int)+tmp*sizeof(PetscInt)+(3+(a->nz))*sizeof(PetscScalar));CHKERRQ(ierr);
1836aa372e3fSPaul Mullowney 
1837aa372e3fSPaul Mullowney         /* assign the pointer */
1838aa372e3fSPaul Mullowney         cusparsestruct->mat = matstruct;
18399ae82921SPaul Mullowney       } catch(char *ex) {
18409ae82921SPaul Mullowney         SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_LIB,"CUSPARSE error: %s", ex);
18419ae82921SPaul Mullowney       }
184205035670SJunchao Zhang       err  = WaitForCUDA();CHKERRCUDA(err);
184385ba7357SStefano Zampini       ierr = PetscLogEventEnd(MAT_CUSPARSECopyToGPU,A,0,0,0);CHKERRQ(ierr);
184434d6c7a5SJose E. Roman       cusparsestruct->nonzerostate = A->nonzerostate;
184534d6c7a5SJose E. Roman     }
1846c70f7ee4SJunchao Zhang     A->offloadmask = PETSC_OFFLOAD_BOTH;
18479ae82921SPaul Mullowney   }
18489ae82921SPaul Mullowney   PetscFunctionReturn(0);
18499ae82921SPaul Mullowney }
18509ae82921SPaul Mullowney 
1851c41cb2e2SAlejandro Lamas Daviña struct VecCUDAPlusEquals
1852aa372e3fSPaul Mullowney {
1853aa372e3fSPaul Mullowney   template <typename Tuple>
1854aa372e3fSPaul Mullowney   __host__ __device__
1855aa372e3fSPaul Mullowney   void operator()(Tuple t)
1856aa372e3fSPaul Mullowney   {
1857aa372e3fSPaul Mullowney     thrust::get<1>(t) = thrust::get<1>(t) + thrust::get<0>(t);
1858aa372e3fSPaul Mullowney   }
1859aa372e3fSPaul Mullowney };
1860aa372e3fSPaul Mullowney 
18617e8381f9SStefano Zampini struct VecCUDAEquals
18627e8381f9SStefano Zampini {
18637e8381f9SStefano Zampini   template <typename Tuple>
18647e8381f9SStefano Zampini   __host__ __device__
18657e8381f9SStefano Zampini   void operator()(Tuple t)
18667e8381f9SStefano Zampini   {
18677e8381f9SStefano Zampini     thrust::get<1>(t) = thrust::get<0>(t);
18687e8381f9SStefano Zampini   }
18697e8381f9SStefano Zampini };
18707e8381f9SStefano Zampini 
1871e6e9a74fSStefano Zampini struct VecCUDAEqualsReverse
1872e6e9a74fSStefano Zampini {
1873e6e9a74fSStefano Zampini   template <typename Tuple>
1874e6e9a74fSStefano Zampini   __host__ __device__
1875e6e9a74fSStefano Zampini   void operator()(Tuple t)
1876e6e9a74fSStefano Zampini   {
1877e6e9a74fSStefano Zampini     thrust::get<0>(t) = thrust::get<1>(t);
1878e6e9a74fSStefano Zampini   }
1879e6e9a74fSStefano Zampini };
1880e6e9a74fSStefano Zampini 
1881afb2bd1cSJunchao Zhang struct MatMatCusparse {
1882ccdfe979SStefano Zampini   PetscBool             cisdense;
1883ccdfe979SStefano Zampini   PetscScalar           *Bt;
1884ccdfe979SStefano Zampini   Mat                   X;
1885fcdce8c4SStefano Zampini   PetscBool             reusesym; /* Cusparse does not have split symbolic and numeric phases for sparse matmat operations */
1886fcdce8c4SStefano Zampini   PetscLogDouble        flops;
1887fcdce8c4SStefano Zampini   CsrMatrix             *Bcsr;
1888afb2bd1cSJunchao Zhang #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1889fcdce8c4SStefano Zampini   cusparseSpMatDescr_t  matSpBDescr;
1890afb2bd1cSJunchao Zhang   PetscBool             initialized;   /* C = alpha op(A) op(B) + beta C */
1891afb2bd1cSJunchao Zhang   cusparseDnMatDescr_t  matBDescr;
1892afb2bd1cSJunchao Zhang   cusparseDnMatDescr_t  matCDescr;
1893afb2bd1cSJunchao Zhang   PetscInt              Blda,Clda; /* Record leading dimensions of B and C here to detect changes*/
1894fcdce8c4SStefano Zampini   size_t                mmBufferSize;
1895fcdce8c4SStefano Zampini   void                  *mmBuffer;
1896fcdce8c4SStefano Zampini   void                  *mmBuffer2; /* SpGEMM WorkEstimation buffer */
1897fcdce8c4SStefano Zampini   cusparseSpGEMMDescr_t spgemmDesc;
1898afb2bd1cSJunchao Zhang #endif
1899afb2bd1cSJunchao Zhang };
1900ccdfe979SStefano Zampini 
1901ccdfe979SStefano Zampini static PetscErrorCode MatDestroy_MatMatCusparse(void *data)
1902ccdfe979SStefano Zampini {
1903ccdfe979SStefano Zampini   PetscErrorCode   ierr;
1904ccdfe979SStefano Zampini   MatMatCusparse   *mmdata = (MatMatCusparse *)data;
1905ccdfe979SStefano Zampini   cudaError_t      cerr;
1906fcdce8c4SStefano Zampini  #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1907fcdce8c4SStefano Zampini   cusparseStatus_t stat;
1908fcdce8c4SStefano Zampini  #endif
1909ccdfe979SStefano Zampini 
1910ccdfe979SStefano Zampini   PetscFunctionBegin;
1911ccdfe979SStefano Zampini   cerr = cudaFree(mmdata->Bt);CHKERRCUDA(cerr);
1912fcdce8c4SStefano Zampini   delete mmdata->Bcsr;
1913afb2bd1cSJunchao Zhang  #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
1914fcdce8c4SStefano Zampini   if (mmdata->matSpBDescr) { stat = cusparseDestroySpMat(mmdata->matSpBDescr);CHKERRCUSPARSE(stat); }
1915fcdce8c4SStefano Zampini   if (mmdata->mmBuffer)    { cerr = cudaFree(mmdata->mmBuffer);CHKERRCUDA(cerr); }
1916fcdce8c4SStefano Zampini   if (mmdata->mmBuffer2)   { cerr = cudaFree(mmdata->mmBuffer2);CHKERRCUDA(cerr); }
1917afb2bd1cSJunchao Zhang   if (mmdata->matBDescr)   { stat = cusparseDestroyDnMat(mmdata->matBDescr);CHKERRCUSPARSE(stat); }
1918afb2bd1cSJunchao Zhang   if (mmdata->matCDescr)   { stat = cusparseDestroyDnMat(mmdata->matCDescr);CHKERRCUSPARSE(stat); }
1919fcdce8c4SStefano Zampini   if (mmdata->spgemmDesc)  { stat = cusparseSpGEMM_destroyDescr(mmdata->spgemmDesc);CHKERRCUSPARSE(stat); }
1920afb2bd1cSJunchao Zhang  #endif
1921ccdfe979SStefano Zampini   ierr = MatDestroy(&mmdata->X);CHKERRQ(ierr);
1922ccdfe979SStefano Zampini   ierr = PetscFree(data);CHKERRQ(ierr);
1923ccdfe979SStefano Zampini   PetscFunctionReturn(0);
1924ccdfe979SStefano Zampini }
1925ccdfe979SStefano Zampini 
1926ccdfe979SStefano Zampini PETSC_INTERN PetscErrorCode MatMatMultNumeric_SeqDenseCUDA_SeqDenseCUDA_Private(Mat,Mat,Mat,PetscBool,PetscBool);
1927ccdfe979SStefano Zampini 
1928ccdfe979SStefano Zampini static PetscErrorCode MatProductNumeric_SeqAIJCUSPARSE_SeqDENSECUDA(Mat C)
1929ccdfe979SStefano Zampini {
1930ccdfe979SStefano Zampini   Mat_Product                  *product = C->product;
1931ccdfe979SStefano Zampini   Mat                          A,B;
1932afb2bd1cSJunchao Zhang   PetscInt                     m,n,blda,clda;
1933ccdfe979SStefano Zampini   PetscBool                    flg,biscuda;
1934ccdfe979SStefano Zampini   Mat_SeqAIJCUSPARSE           *cusp;
1935ccdfe979SStefano Zampini   cusparseStatus_t             stat;
1936ccdfe979SStefano Zampini   cusparseOperation_t          opA;
1937ccdfe979SStefano Zampini   const PetscScalar            *barray;
1938ccdfe979SStefano Zampini   PetscScalar                  *carray;
1939ccdfe979SStefano Zampini   PetscErrorCode               ierr;
1940ccdfe979SStefano Zampini   MatMatCusparse               *mmdata;
1941ccdfe979SStefano Zampini   Mat_SeqAIJCUSPARSEMultStruct *mat;
1942ccdfe979SStefano Zampini   CsrMatrix                    *csrmat;
1943afb2bd1cSJunchao Zhang   cudaError_t                  cerr;
1944ccdfe979SStefano Zampini 
1945ccdfe979SStefano Zampini   PetscFunctionBegin;
1946ccdfe979SStefano Zampini   MatCheckProduct(C,1);
1947ccdfe979SStefano Zampini   if (!C->product->data) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data empty");
1948ccdfe979SStefano Zampini   mmdata = (MatMatCusparse*)product->data;
1949ccdfe979SStefano Zampini   A    = product->A;
1950ccdfe979SStefano Zampini   B    = product->B;
1951ccdfe979SStefano Zampini   ierr = PetscObjectTypeCompare((PetscObject)A,MATSEQAIJCUSPARSE,&flg);CHKERRQ(ierr);
1952ccdfe979SStefano Zampini   if (!flg) SETERRQ1(PetscObjectComm((PetscObject)A),PETSC_ERR_PLIB,"Not for type %s",((PetscObject)A)->type_name);
1953ccdfe979SStefano Zampini   /* currently CopyToGpu does not copy if the matrix is bound to CPU
1954ccdfe979SStefano Zampini      Instead of silently accepting the wrong answer, I prefer to raise the error */
1955ccdfe979SStefano Zampini   if (A->boundtocpu) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_ARG_WRONG,"Cannot bind to CPU a CUSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
1956ccdfe979SStefano Zampini   ierr = MatSeqAIJCUSPARSECopyToGPU(A);CHKERRQ(ierr);
1957ccdfe979SStefano Zampini   cusp = (Mat_SeqAIJCUSPARSE*)A->spptr;
1958ccdfe979SStefano Zampini   switch (product->type) {
1959ccdfe979SStefano Zampini   case MATPRODUCT_AB:
1960ccdfe979SStefano Zampini   case MATPRODUCT_PtAP:
1961ccdfe979SStefano Zampini     mat = cusp->mat;
1962ccdfe979SStefano Zampini     opA = CUSPARSE_OPERATION_NON_TRANSPOSE;
1963ccdfe979SStefano Zampini     m   = A->rmap->n;
1964ccdfe979SStefano Zampini     n   = B->cmap->n;
1965ccdfe979SStefano Zampini     break;
1966ccdfe979SStefano Zampini   case MATPRODUCT_AtB:
1967e6e9a74fSStefano Zampini     if (!cusp->transgen) {
1968e6e9a74fSStefano Zampini       mat = cusp->mat;
1969e6e9a74fSStefano Zampini       opA = CUSPARSE_OPERATION_TRANSPOSE;
1970e6e9a74fSStefano Zampini     } else {
1971ccdfe979SStefano Zampini       ierr = MatSeqAIJCUSPARSEGenerateTransposeForMult(A);CHKERRQ(ierr);
1972ccdfe979SStefano Zampini       mat  = cusp->matTranspose;
1973ccdfe979SStefano Zampini       opA  = CUSPARSE_OPERATION_NON_TRANSPOSE;
1974e6e9a74fSStefano Zampini     }
1975ccdfe979SStefano Zampini     m = A->cmap->n;
1976ccdfe979SStefano Zampini     n = B->cmap->n;
1977ccdfe979SStefano Zampini     break;
1978ccdfe979SStefano Zampini   case MATPRODUCT_ABt:
1979ccdfe979SStefano Zampini   case MATPRODUCT_RARt:
1980ccdfe979SStefano Zampini     mat = cusp->mat;
1981ccdfe979SStefano Zampini     opA = CUSPARSE_OPERATION_NON_TRANSPOSE;
1982ccdfe979SStefano Zampini     m   = A->rmap->n;
1983ccdfe979SStefano Zampini     n   = B->rmap->n;
1984ccdfe979SStefano Zampini     break;
1985ccdfe979SStefano Zampini   default:
1986ccdfe979SStefano Zampini     SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Unsupported product type %s",MatProductTypes[product->type]);
1987ccdfe979SStefano Zampini   }
1988ccdfe979SStefano Zampini   if (!mat) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Missing Mat_SeqAIJCUSPARSEMultStruct");
1989ccdfe979SStefano Zampini   csrmat = (CsrMatrix*)mat->mat;
1990ccdfe979SStefano Zampini   /* if the user passed a CPU matrix, copy the data to the GPU */
1991ccdfe979SStefano Zampini   ierr = PetscObjectTypeCompare((PetscObject)B,MATSEQDENSECUDA,&biscuda);CHKERRQ(ierr);
1992afb2bd1cSJunchao Zhang   if (!biscuda) {ierr = MatConvert(B,MATSEQDENSECUDA,MAT_INPLACE_MATRIX,&B);CHKERRQ(ierr);}
1993ccdfe979SStefano Zampini   ierr = MatDenseCUDAGetArrayRead(B,&barray);CHKERRQ(ierr);
1994afb2bd1cSJunchao Zhang 
1995ccdfe979SStefano Zampini   ierr = MatDenseGetLDA(B,&blda);CHKERRQ(ierr);
1996c8378d12SStefano Zampini   if (product->type == MATPRODUCT_RARt || product->type == MATPRODUCT_PtAP) {
1997c8378d12SStefano Zampini     ierr = MatDenseCUDAGetArrayWrite(mmdata->X,&carray);CHKERRQ(ierr);
1998c8378d12SStefano Zampini     ierr = MatDenseGetLDA(mmdata->X,&clda);CHKERRQ(ierr);
1999c8378d12SStefano Zampini   } else {
2000c8378d12SStefano Zampini     ierr = MatDenseCUDAGetArrayWrite(C,&carray);CHKERRQ(ierr);
2001c8378d12SStefano Zampini     ierr = MatDenseGetLDA(C,&clda);CHKERRQ(ierr);
2002c8378d12SStefano Zampini   }
2003c8378d12SStefano Zampini 
2004c8378d12SStefano Zampini   ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
2005afb2bd1cSJunchao Zhang  #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
2006afb2bd1cSJunchao Zhang   cusparseOperation_t opB = (product->type == MATPRODUCT_ABt || product->type == MATPRODUCT_RARt) ? CUSPARSE_OPERATION_TRANSPOSE : CUSPARSE_OPERATION_NON_TRANSPOSE;
2007fcdce8c4SStefano Zampini   /* (re)allcoate mmBuffer if not initialized or LDAs are different */
2008afb2bd1cSJunchao Zhang   if (!mmdata->initialized || mmdata->Blda != blda || mmdata->Clda != clda) {
2009fcdce8c4SStefano Zampini     size_t mmBufferSize;
2010afb2bd1cSJunchao Zhang     if (mmdata->initialized && mmdata->Blda != blda) {stat = cusparseDestroyDnMat(mmdata->matBDescr);CHKERRCUSPARSE(stat); mmdata->matBDescr = NULL;}
2011afb2bd1cSJunchao Zhang     if (!mmdata->matBDescr) {
2012afb2bd1cSJunchao Zhang       stat         = cusparseCreateDnMat(&mmdata->matBDescr,B->rmap->n,B->cmap->n,blda,(void*)barray,cusparse_scalartype,CUSPARSE_ORDER_COL);CHKERRCUSPARSE(stat);
2013afb2bd1cSJunchao Zhang       mmdata->Blda = blda;
2014afb2bd1cSJunchao Zhang     }
2015c8378d12SStefano Zampini 
2016afb2bd1cSJunchao Zhang     if (mmdata->initialized && mmdata->Clda != clda) {stat = cusparseDestroyDnMat(mmdata->matCDescr);CHKERRCUSPARSE(stat); mmdata->matCDescr = NULL;}
2017afb2bd1cSJunchao Zhang     if (!mmdata->matCDescr) { /* matCDescr is for C or mmdata->X */
2018afb2bd1cSJunchao Zhang       stat         = cusparseCreateDnMat(&mmdata->matCDescr,m,n,clda,(void*)carray,cusparse_scalartype,CUSPARSE_ORDER_COL);CHKERRCUSPARSE(stat);
2019afb2bd1cSJunchao Zhang       mmdata->Clda = clda;
2020afb2bd1cSJunchao Zhang     }
2021afb2bd1cSJunchao Zhang 
2022afb2bd1cSJunchao Zhang     if (!mat->matDescr) {
2023afb2bd1cSJunchao Zhang       stat = cusparseCreateCsr(&mat->matDescr,
2024afb2bd1cSJunchao Zhang                                csrmat->num_rows, csrmat->num_cols, csrmat->num_entries,
2025afb2bd1cSJunchao Zhang                                csrmat->row_offsets->data().get(), csrmat->column_indices->data().get(),
2026afb2bd1cSJunchao Zhang                                csrmat->values->data().get(),
2027afb2bd1cSJunchao Zhang                                CUSPARSE_INDEX_32I,CUSPARSE_INDEX_32I, /* row offset, col idx types due to THRUSTINTARRAY32 */
2028afb2bd1cSJunchao Zhang                                CUSPARSE_INDEX_BASE_ZERO,cusparse_scalartype);CHKERRCUSPARSE(stat);
2029afb2bd1cSJunchao Zhang     }
2030afb2bd1cSJunchao Zhang     stat = cusparseSpMM_bufferSize(cusp->handle,opA,opB,mat->alpha_one,
2031afb2bd1cSJunchao Zhang                                    mat->matDescr,mmdata->matBDescr,mat->beta_zero,
2032afb2bd1cSJunchao Zhang                                    mmdata->matCDescr,cusparse_scalartype,
2033fcdce8c4SStefano Zampini                                    cusp->spmmAlg,&mmBufferSize);CHKERRCUSPARSE(stat);
2034fcdce8c4SStefano Zampini     if ((mmdata->mmBuffer && mmdata->mmBufferSize < mmBufferSize) || !mmdata->mmBuffer) {
2035fcdce8c4SStefano Zampini       cerr = cudaFree(mmdata->mmBuffer);CHKERRCUDA(cerr);
2036fcdce8c4SStefano Zampini       cerr = cudaMalloc(&mmdata->mmBuffer,mmBufferSize);CHKERRCUDA(cerr);
2037fcdce8c4SStefano Zampini       mmdata->mmBufferSize = mmBufferSize;
2038fcdce8c4SStefano Zampini     }
2039afb2bd1cSJunchao Zhang     mmdata->initialized = PETSC_TRUE;
2040afb2bd1cSJunchao Zhang   } else {
2041afb2bd1cSJunchao Zhang     /* to be safe, always update pointers of the mats */
2042afb2bd1cSJunchao Zhang     stat = cusparseSpMatSetValues(mat->matDescr,csrmat->values->data().get());CHKERRCUSPARSE(stat);
2043afb2bd1cSJunchao Zhang     stat = cusparseDnMatSetValues(mmdata->matBDescr,(void*)barray);CHKERRCUSPARSE(stat);
2044afb2bd1cSJunchao Zhang     stat = cusparseDnMatSetValues(mmdata->matCDescr,(void*)carray);CHKERRCUSPARSE(stat);
2045afb2bd1cSJunchao Zhang   }
2046afb2bd1cSJunchao Zhang 
2047afb2bd1cSJunchao Zhang   /* do cusparseSpMM, which supports transpose on B */
2048afb2bd1cSJunchao Zhang   stat = cusparseSpMM(cusp->handle,opA,opB,mat->alpha_one,
2049afb2bd1cSJunchao Zhang                       mat->matDescr,mmdata->matBDescr,mat->beta_zero,
2050afb2bd1cSJunchao Zhang                       mmdata->matCDescr,cusparse_scalartype,
2051fcdce8c4SStefano Zampini                       cusp->spmmAlg,mmdata->mmBuffer);CHKERRCUSPARSE(stat);
2052afb2bd1cSJunchao Zhang  #else
2053afb2bd1cSJunchao Zhang   PetscInt k;
2054afb2bd1cSJunchao Zhang   /* cusparseXcsrmm does not support transpose on B */
2055ccdfe979SStefano Zampini   if (product->type == MATPRODUCT_ABt || product->type == MATPRODUCT_RARt) {
2056ccdfe979SStefano Zampini     cublasHandle_t cublasv2handle;
2057ccdfe979SStefano Zampini     cublasStatus_t cerr;
2058ccdfe979SStefano Zampini 
2059ccdfe979SStefano Zampini     ierr = PetscCUBLASGetHandle(&cublasv2handle);CHKERRQ(ierr);
2060ccdfe979SStefano Zampini     cerr = cublasXgeam(cublasv2handle,CUBLAS_OP_T,CUBLAS_OP_T,
2061ccdfe979SStefano Zampini                        B->cmap->n,B->rmap->n,
2062ccdfe979SStefano Zampini                        &PETSC_CUSPARSE_ONE ,barray,blda,
2063ccdfe979SStefano Zampini                        &PETSC_CUSPARSE_ZERO,barray,blda,
2064ccdfe979SStefano Zampini                        mmdata->Bt,B->cmap->n);CHKERRCUBLAS(cerr);
2065ccdfe979SStefano Zampini     blda = B->cmap->n;
2066afb2bd1cSJunchao Zhang     k    = B->cmap->n;
2067afb2bd1cSJunchao Zhang   } else {
2068afb2bd1cSJunchao Zhang     k    = B->rmap->n;
2069ccdfe979SStefano Zampini   }
2070ccdfe979SStefano Zampini 
2071afb2bd1cSJunchao Zhang   /* perform the MatMat operation, op(A) is m x k, op(B) is k x n */
2072ccdfe979SStefano Zampini   stat = cusparse_csr_spmm(cusp->handle,opA,m,n,k,
2073afb2bd1cSJunchao Zhang                            csrmat->num_entries,mat->alpha_one,mat->descr,
2074ccdfe979SStefano Zampini                            csrmat->values->data().get(),
2075ccdfe979SStefano Zampini                            csrmat->row_offsets->data().get(),
2076ccdfe979SStefano Zampini                            csrmat->column_indices->data().get(),
2077ccdfe979SStefano Zampini                            mmdata->Bt ? mmdata->Bt : barray,blda,mat->beta_zero,
2078ccdfe979SStefano Zampini                            carray,clda);CHKERRCUSPARSE(stat);
2079afb2bd1cSJunchao Zhang  #endif
2080afb2bd1cSJunchao Zhang   cerr = WaitForCUDA();CHKERRCUDA(cerr);
2081c8378d12SStefano Zampini   ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
2082c8378d12SStefano Zampini   ierr = PetscLogGpuFlops(n*2.0*csrmat->num_entries);CHKERRQ(ierr);
2083ccdfe979SStefano Zampini   ierr = MatDenseCUDARestoreArrayRead(B,&barray);CHKERRQ(ierr);
2084ccdfe979SStefano Zampini   if (product->type == MATPRODUCT_RARt) {
2085ccdfe979SStefano Zampini     ierr = MatDenseCUDARestoreArrayWrite(mmdata->X,&carray);CHKERRQ(ierr);
2086ccdfe979SStefano Zampini     ierr = MatMatMultNumeric_SeqDenseCUDA_SeqDenseCUDA_Private(B,mmdata->X,C,PETSC_FALSE,PETSC_FALSE);CHKERRQ(ierr);
2087ccdfe979SStefano Zampini   } else if (product->type == MATPRODUCT_PtAP) {
2088ccdfe979SStefano Zampini     ierr = MatDenseCUDARestoreArrayWrite(mmdata->X,&carray);CHKERRQ(ierr);
2089ccdfe979SStefano Zampini     ierr = MatMatMultNumeric_SeqDenseCUDA_SeqDenseCUDA_Private(B,mmdata->X,C,PETSC_TRUE,PETSC_FALSE);CHKERRQ(ierr);
2090ccdfe979SStefano Zampini   } else {
2091ccdfe979SStefano Zampini     ierr = MatDenseCUDARestoreArrayWrite(C,&carray);CHKERRQ(ierr);
2092ccdfe979SStefano Zampini   }
2093ccdfe979SStefano Zampini   if (mmdata->cisdense) {
2094ccdfe979SStefano Zampini     ierr = MatConvert(C,MATSEQDENSE,MAT_INPLACE_MATRIX,&C);CHKERRQ(ierr);
2095ccdfe979SStefano Zampini   }
2096ccdfe979SStefano Zampini   if (!biscuda) {
2097ccdfe979SStefano Zampini     ierr = MatConvert(B,MATSEQDENSE,MAT_INPLACE_MATRIX,&B);CHKERRQ(ierr);
2098ccdfe979SStefano Zampini   }
2099ccdfe979SStefano Zampini   PetscFunctionReturn(0);
2100ccdfe979SStefano Zampini }
2101ccdfe979SStefano Zampini 
2102ccdfe979SStefano Zampini static PetscErrorCode MatProductSymbolic_SeqAIJCUSPARSE_SeqDENSECUDA(Mat C)
2103ccdfe979SStefano Zampini {
2104ccdfe979SStefano Zampini   Mat_Product        *product = C->product;
2105ccdfe979SStefano Zampini   Mat                A,B;
2106ccdfe979SStefano Zampini   PetscInt           m,n;
2107ccdfe979SStefano Zampini   PetscBool          cisdense,flg;
2108ccdfe979SStefano Zampini   PetscErrorCode     ierr;
2109ccdfe979SStefano Zampini   MatMatCusparse     *mmdata;
2110ccdfe979SStefano Zampini   Mat_SeqAIJCUSPARSE *cusp;
2111ccdfe979SStefano Zampini 
2112ccdfe979SStefano Zampini   PetscFunctionBegin;
2113ccdfe979SStefano Zampini   MatCheckProduct(C,1);
2114ccdfe979SStefano Zampini   if (C->product->data) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data not empty");
2115ccdfe979SStefano Zampini   A    = product->A;
2116ccdfe979SStefano Zampini   B    = product->B;
2117ccdfe979SStefano Zampini   ierr = PetscObjectTypeCompare((PetscObject)A,MATSEQAIJCUSPARSE,&flg);CHKERRQ(ierr);
2118ccdfe979SStefano Zampini   if (!flg) SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Not for type %s",((PetscObject)A)->type_name);
2119ccdfe979SStefano Zampini   cusp = (Mat_SeqAIJCUSPARSE*)A->spptr;
2120ccdfe979SStefano Zampini   if (cusp->format != MAT_CUSPARSE_CSR) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Only for MAT_CUSPARSE_CSR format");
2121ccdfe979SStefano Zampini   switch (product->type) {
2122ccdfe979SStefano Zampini   case MATPRODUCT_AB:
2123ccdfe979SStefano Zampini     m = A->rmap->n;
2124ccdfe979SStefano Zampini     n = B->cmap->n;
2125ccdfe979SStefano Zampini     break;
2126ccdfe979SStefano Zampini   case MATPRODUCT_AtB:
2127ccdfe979SStefano Zampini     m = A->cmap->n;
2128ccdfe979SStefano Zampini     n = B->cmap->n;
2129ccdfe979SStefano Zampini     break;
2130ccdfe979SStefano Zampini   case MATPRODUCT_ABt:
2131ccdfe979SStefano Zampini     m = A->rmap->n;
2132ccdfe979SStefano Zampini     n = B->rmap->n;
2133ccdfe979SStefano Zampini     break;
2134ccdfe979SStefano Zampini   case MATPRODUCT_PtAP:
2135ccdfe979SStefano Zampini     m = B->cmap->n;
2136ccdfe979SStefano Zampini     n = B->cmap->n;
2137ccdfe979SStefano Zampini     break;
2138ccdfe979SStefano Zampini   case MATPRODUCT_RARt:
2139ccdfe979SStefano Zampini     m = B->rmap->n;
2140ccdfe979SStefano Zampini     n = B->rmap->n;
2141ccdfe979SStefano Zampini     break;
2142ccdfe979SStefano Zampini   default:
2143ccdfe979SStefano Zampini     SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Unsupported product type %s",MatProductTypes[product->type]);
2144ccdfe979SStefano Zampini   }
2145ccdfe979SStefano Zampini   ierr = MatSetSizes(C,m,n,m,n);CHKERRQ(ierr);
2146ccdfe979SStefano Zampini   /* if C is of type MATSEQDENSE (CPU), perform the operation on the GPU and then copy on the CPU */
2147ccdfe979SStefano Zampini   ierr = PetscObjectTypeCompare((PetscObject)C,MATSEQDENSE,&cisdense);CHKERRQ(ierr);
2148ccdfe979SStefano Zampini   ierr = MatSetType(C,MATSEQDENSECUDA);CHKERRQ(ierr);
2149ccdfe979SStefano Zampini 
2150ccdfe979SStefano Zampini   /* product data */
2151ccdfe979SStefano Zampini   ierr = PetscNew(&mmdata);CHKERRQ(ierr);
2152ccdfe979SStefano Zampini   mmdata->cisdense = cisdense;
2153afb2bd1cSJunchao Zhang  #if PETSC_PKG_CUDA_VERSION_LT(11,0,0)
2154afb2bd1cSJunchao Zhang   /* cusparseXcsrmm does not support transpose on B, so we allocate buffer to store B^T */
2155ccdfe979SStefano Zampini   if (product->type == MATPRODUCT_ABt || product->type == MATPRODUCT_RARt) {
2156afb2bd1cSJunchao Zhang     cudaError_t cerr = cudaMalloc((void**)&mmdata->Bt,(size_t)B->rmap->n*(size_t)B->cmap->n*sizeof(PetscScalar));CHKERRCUDA(cerr);
2157ccdfe979SStefano Zampini   }
2158afb2bd1cSJunchao Zhang  #endif
2159ccdfe979SStefano Zampini   /* for these products we need intermediate storage */
2160ccdfe979SStefano Zampini   if (product->type == MATPRODUCT_RARt || product->type == MATPRODUCT_PtAP) {
2161ccdfe979SStefano Zampini     ierr = MatCreate(PetscObjectComm((PetscObject)C),&mmdata->X);CHKERRQ(ierr);
2162ccdfe979SStefano Zampini     ierr = MatSetType(mmdata->X,MATSEQDENSECUDA);CHKERRQ(ierr);
2163ccdfe979SStefano Zampini     if (product->type == MATPRODUCT_RARt) { /* do not preallocate, since the first call to MatDenseCUDAGetArray will preallocate on the GPU for us */
2164ccdfe979SStefano Zampini       ierr = MatSetSizes(mmdata->X,A->rmap->n,B->rmap->n,A->rmap->n,B->rmap->n);CHKERRQ(ierr);
2165ccdfe979SStefano Zampini     } else {
2166ccdfe979SStefano Zampini       ierr = MatSetSizes(mmdata->X,A->rmap->n,B->cmap->n,A->rmap->n,B->cmap->n);CHKERRQ(ierr);
2167ccdfe979SStefano Zampini     }
2168ccdfe979SStefano Zampini   }
2169ccdfe979SStefano Zampini   C->product->data    = mmdata;
2170ccdfe979SStefano Zampini   C->product->destroy = MatDestroy_MatMatCusparse;
2171ccdfe979SStefano Zampini 
2172ccdfe979SStefano Zampini   C->ops->productnumeric = MatProductNumeric_SeqAIJCUSPARSE_SeqDENSECUDA;
2173ccdfe979SStefano Zampini   PetscFunctionReturn(0);
2174ccdfe979SStefano Zampini }
2175ccdfe979SStefano Zampini 
2176fcdce8c4SStefano Zampini static PetscErrorCode MatProductNumeric_SeqAIJCUSPARSE_SeqAIJCUSPARSE(Mat C)
2177ccdfe979SStefano Zampini {
2178ccdfe979SStefano Zampini   Mat_Product                  *product = C->product;
2179fcdce8c4SStefano Zampini   Mat                          A,B;
2180fcdce8c4SStefano Zampini   Mat_SeqAIJCUSPARSE           *Acusp,*Bcusp,*Ccusp;
2181fcdce8c4SStefano Zampini   Mat_SeqAIJ                   *c = (Mat_SeqAIJ*)C->data;
2182fcdce8c4SStefano Zampini   Mat_SeqAIJCUSPARSEMultStruct *Amat,*Bmat,*Cmat;
2183fcdce8c4SStefano Zampini   CsrMatrix                    *Acsr,*Bcsr,*Ccsr;
2184fcdce8c4SStefano Zampini   PetscBool                    flg;
2185ccdfe979SStefano Zampini   PetscErrorCode               ierr;
2186fcdce8c4SStefano Zampini   cusparseStatus_t             stat;
2187fcdce8c4SStefano Zampini   cudaError_t                  cerr;
2188fcdce8c4SStefano Zampini   MatProductType               ptype;
2189fcdce8c4SStefano Zampini   MatMatCusparse               *mmdata;
2190fcdce8c4SStefano Zampini #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
2191fcdce8c4SStefano Zampini   cusparseSpMatDescr_t         BmatSpDescr;
2192fcdce8c4SStefano Zampini #endif
2193ccdfe979SStefano Zampini 
2194ccdfe979SStefano Zampini   PetscFunctionBegin;
2195ccdfe979SStefano Zampini   MatCheckProduct(C,1);
2196fcdce8c4SStefano Zampini   if (!C->product->data) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data empty");
2197fcdce8c4SStefano Zampini   ierr = PetscObjectTypeCompare((PetscObject)C,MATSEQAIJCUSPARSE,&flg);CHKERRQ(ierr);
2198fcdce8c4SStefano Zampini   if (!flg) SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Not for C of type %s",((PetscObject)C)->type_name);
2199fcdce8c4SStefano Zampini   mmdata = (MatMatCusparse*)C->product->data;
2200fcdce8c4SStefano Zampini   A = product->A;
2201fcdce8c4SStefano Zampini   B = product->B;
2202fcdce8c4SStefano Zampini   if (mmdata->reusesym) { /* this happens when api_user is true, meaning that the matrix values have been already computed in the MatProductSymbolic phase */
2203fcdce8c4SStefano Zampini     mmdata->reusesym = PETSC_FALSE;
2204fcdce8c4SStefano Zampini     Ccusp = (Mat_SeqAIJCUSPARSE*)C->spptr;
2205fcdce8c4SStefano Zampini     if (Ccusp->format != MAT_CUSPARSE_CSR) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Only for MAT_CUSPARSE_CSR format");
2206fcdce8c4SStefano Zampini     Cmat = Ccusp->mat;
2207fcdce8c4SStefano Zampini     if (!Cmat) SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Missing C mult struct for product type %s",MatProductTypes[C->product->type]);
2208fcdce8c4SStefano Zampini     Ccsr = (CsrMatrix*)Cmat->mat;
2209fcdce8c4SStefano Zampini     if (!Ccsr) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Missing C CSR struct");
2210fcdce8c4SStefano Zampini     goto finalize;
2211fcdce8c4SStefano Zampini   }
2212fcdce8c4SStefano Zampini   if (!c->nz) goto finalize;
2213fcdce8c4SStefano Zampini   ierr = PetscObjectTypeCompare((PetscObject)A,MATSEQAIJCUSPARSE,&flg);CHKERRQ(ierr);
2214fcdce8c4SStefano Zampini   if (!flg) SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Not for type %s",((PetscObject)A)->type_name);
2215fcdce8c4SStefano Zampini   ierr = PetscObjectTypeCompare((PetscObject)B,MATSEQAIJCUSPARSE,&flg);CHKERRQ(ierr);
2216fcdce8c4SStefano Zampini   if (!flg) SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Not for B of type %s",((PetscObject)B)->type_name);
2217fcdce8c4SStefano Zampini   if (A->boundtocpu) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONG,"Cannot bind to CPU a CUSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
2218fcdce8c4SStefano Zampini   if (B->boundtocpu) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONG,"Cannot bind to CPU a CUSPARSE matrix between MatProductSymbolic and MatProductNumeric phases");
2219fcdce8c4SStefano Zampini   Acusp = (Mat_SeqAIJCUSPARSE*)A->spptr;
2220fcdce8c4SStefano Zampini   Bcusp = (Mat_SeqAIJCUSPARSE*)B->spptr;
2221fcdce8c4SStefano Zampini   Ccusp = (Mat_SeqAIJCUSPARSE*)C->spptr;
2222fcdce8c4SStefano Zampini   if (Acusp->format != MAT_CUSPARSE_CSR) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Only for MAT_CUSPARSE_CSR format");
2223fcdce8c4SStefano Zampini   if (Bcusp->format != MAT_CUSPARSE_CSR) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Only for MAT_CUSPARSE_CSR format");
2224fcdce8c4SStefano Zampini   if (Ccusp->format != MAT_CUSPARSE_CSR) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Only for MAT_CUSPARSE_CSR format");
2225fcdce8c4SStefano Zampini   ierr = MatSeqAIJCUSPARSECopyToGPU(A);CHKERRQ(ierr);
2226fcdce8c4SStefano Zampini   ierr = MatSeqAIJCUSPARSECopyToGPU(B);CHKERRQ(ierr);
2227fcdce8c4SStefano Zampini 
2228fcdce8c4SStefano Zampini   ptype = product->type;
2229fcdce8c4SStefano Zampini   if (A->symmetric && ptype == MATPRODUCT_AtB) ptype = MATPRODUCT_AB;
2230fcdce8c4SStefano Zampini   if (B->symmetric && ptype == MATPRODUCT_ABt) ptype = MATPRODUCT_AB;
2231fcdce8c4SStefano Zampini   switch (ptype) {
2232fcdce8c4SStefano Zampini   case MATPRODUCT_AB:
2233fcdce8c4SStefano Zampini     Amat = Acusp->mat;
2234fcdce8c4SStefano Zampini     Bmat = Bcusp->mat;
2235fcdce8c4SStefano Zampini     break;
2236fcdce8c4SStefano Zampini   case MATPRODUCT_AtB:
2237fcdce8c4SStefano Zampini     Amat = Acusp->matTranspose;
2238fcdce8c4SStefano Zampini     Bmat = Bcusp->mat;
2239fcdce8c4SStefano Zampini     break;
2240fcdce8c4SStefano Zampini   case MATPRODUCT_ABt:
2241fcdce8c4SStefano Zampini     Amat = Acusp->mat;
2242fcdce8c4SStefano Zampini     Bmat = Bcusp->matTranspose;
2243fcdce8c4SStefano Zampini     break;
2244fcdce8c4SStefano Zampini   default:
2245fcdce8c4SStefano Zampini     SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Unsupported product type %s",MatProductTypes[product->type]);
2246fcdce8c4SStefano Zampini   }
2247fcdce8c4SStefano Zampini   Cmat = Ccusp->mat;
2248fcdce8c4SStefano Zampini   if (!Amat) SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Missing A mult struct for product type %s",MatProductTypes[ptype]);
2249fcdce8c4SStefano Zampini   if (!Bmat) SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Missing B mult struct for product type %s",MatProductTypes[ptype]);
2250fcdce8c4SStefano Zampini   if (!Cmat) SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Missing C mult struct for product type %s",MatProductTypes[ptype]);
2251fcdce8c4SStefano Zampini   Acsr = (CsrMatrix*)Amat->mat;
2252fcdce8c4SStefano Zampini   Bcsr = mmdata->Bcsr ? mmdata->Bcsr : (CsrMatrix*)Bmat->mat; /* B may be in compressed row storage */
2253fcdce8c4SStefano Zampini   Ccsr = (CsrMatrix*)Cmat->mat;
2254fcdce8c4SStefano Zampini   if (!Acsr) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Missing A CSR struct");
2255fcdce8c4SStefano Zampini   if (!Bcsr) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Missing B CSR struct");
2256fcdce8c4SStefano Zampini   if (!Ccsr) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Missing C CSR struct");
2257fcdce8c4SStefano Zampini   ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
2258fcdce8c4SStefano Zampini #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
2259fcdce8c4SStefano Zampini   BmatSpDescr = mmdata->Bcsr ? mmdata->matSpBDescr : Bmat->matDescr; /* B may be in compressed row storage */
2260fcdce8c4SStefano Zampini   stat = cusparseSpGEMM_compute(Ccusp->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, CUSPARSE_OPERATION_NON_TRANSPOSE,
2261fcdce8c4SStefano Zampini                                 Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr,
2262fcdce8c4SStefano Zampini                                 cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT,
2263fcdce8c4SStefano Zampini                                 mmdata->spgemmDesc, &mmdata->mmBufferSize, mmdata->mmBuffer);CHKERRCUSPARSE(stat);
2264fcdce8c4SStefano Zampini   stat = cusparseSpGEMM_copy(Ccusp->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, CUSPARSE_OPERATION_NON_TRANSPOSE,
2265fcdce8c4SStefano Zampini                              Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr,
2266fcdce8c4SStefano Zampini                              cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc);CHKERRCUSPARSE(stat);
2267fcdce8c4SStefano Zampini #else
2268fcdce8c4SStefano Zampini   stat = cusparse_csr_spgemm(Ccusp->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, CUSPARSE_OPERATION_NON_TRANSPOSE,
2269fcdce8c4SStefano Zampini                              Acsr->num_rows, Bcsr->num_cols, Acsr->num_cols,
2270fcdce8c4SStefano Zampini                              Amat->descr, Acsr->num_entries, Acsr->values->data().get(), Acsr->row_offsets->data().get(), Acsr->column_indices->data().get(),
2271fcdce8c4SStefano Zampini                              Bmat->descr, Bcsr->num_entries, Bcsr->values->data().get(), Bcsr->row_offsets->data().get(), Bcsr->column_indices->data().get(),
2272fcdce8c4SStefano Zampini                              Cmat->descr, Ccsr->values->data().get(), Ccsr->row_offsets->data().get(), Ccsr->column_indices->data().get());CHKERRCUSPARSE(stat);
2273fcdce8c4SStefano Zampini #endif
2274fcdce8c4SStefano Zampini   ierr = PetscLogGpuFlops(mmdata->flops);CHKERRQ(ierr);
2275fcdce8c4SStefano Zampini   cerr = WaitForCUDA();CHKERRCUDA(cerr);
2276fcdce8c4SStefano Zampini   ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
2277fcdce8c4SStefano Zampini   C->offloadmask = PETSC_OFFLOAD_GPU;
2278fcdce8c4SStefano Zampini finalize:
2279fcdce8c4SStefano Zampini   /* shorter version of MatAssemblyEnd_SeqAIJ */
2280fcdce8c4SStefano Zampini   ierr = PetscInfo3(C,"Matrix size: %D X %D; storage space: 0 unneeded,%D used\n",C->rmap->n,C->cmap->n,c->nz);CHKERRQ(ierr);
2281fcdce8c4SStefano Zampini   ierr = PetscInfo(C,"Number of mallocs during MatSetValues() is 0\n");CHKERRQ(ierr);
2282fcdce8c4SStefano Zampini   ierr = PetscInfo1(C,"Maximum nonzeros in any row is %D\n",c->rmax);CHKERRQ(ierr);
2283fcdce8c4SStefano Zampini   c->reallocs         = 0;
2284fcdce8c4SStefano Zampini   C->info.mallocs    += 0;
2285fcdce8c4SStefano Zampini   C->info.nz_unneeded = 0;
2286fcdce8c4SStefano Zampini   C->assembled = C->was_assembled = PETSC_TRUE;
2287fcdce8c4SStefano Zampini   C->num_ass++;
2288ccdfe979SStefano Zampini   PetscFunctionReturn(0);
2289ccdfe979SStefano Zampini }
2290fcdce8c4SStefano Zampini 
2291fcdce8c4SStefano Zampini static PetscErrorCode MatProductSymbolic_SeqAIJCUSPARSE_SeqAIJCUSPARSE(Mat C)
2292fcdce8c4SStefano Zampini {
2293fcdce8c4SStefano Zampini   Mat_Product                  *product = C->product;
2294fcdce8c4SStefano Zampini   Mat                          A,B;
2295fcdce8c4SStefano Zampini   Mat_SeqAIJCUSPARSE           *Acusp,*Bcusp,*Ccusp;
2296fcdce8c4SStefano Zampini   Mat_SeqAIJ                   *a,*b,*c;
2297fcdce8c4SStefano Zampini   Mat_SeqAIJCUSPARSEMultStruct *Amat,*Bmat,*Cmat;
2298fcdce8c4SStefano Zampini   CsrMatrix                    *Acsr,*Bcsr,*Ccsr;
2299fcdce8c4SStefano Zampini   PetscInt                     i,j,m,n,k;
2300fcdce8c4SStefano Zampini   PetscBool                    flg;
2301fcdce8c4SStefano Zampini   PetscErrorCode               ierr;
2302fcdce8c4SStefano Zampini   cusparseStatus_t             stat;
2303fcdce8c4SStefano Zampini   cudaError_t                  cerr;
2304fcdce8c4SStefano Zampini   MatProductType               ptype;
2305fcdce8c4SStefano Zampini   MatMatCusparse               *mmdata;
2306fcdce8c4SStefano Zampini   PetscLogDouble               flops;
2307fcdce8c4SStefano Zampini   PetscBool                    biscompressed,ciscompressed;
2308fcdce8c4SStefano Zampini #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
2309fcdce8c4SStefano Zampini   int64_t                      C_num_rows1, C_num_cols1, C_nnz1;
2310fcdce8c4SStefano Zampini   size_t                       bufSize2;
2311fcdce8c4SStefano Zampini   cusparseSpMatDescr_t         BmatSpDescr;
2312fcdce8c4SStefano Zampini #else
2313fcdce8c4SStefano Zampini   int                          cnz;
2314fcdce8c4SStefano Zampini #endif
2315fcdce8c4SStefano Zampini 
2316fcdce8c4SStefano Zampini   PetscFunctionBegin;
2317fcdce8c4SStefano Zampini   MatCheckProduct(C,1);
2318fcdce8c4SStefano Zampini   if (C->product->data) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data not empty");
2319fcdce8c4SStefano Zampini   A    = product->A;
2320fcdce8c4SStefano Zampini   B    = product->B;
2321fcdce8c4SStefano Zampini   ierr = PetscObjectTypeCompare((PetscObject)A,MATSEQAIJCUSPARSE,&flg);CHKERRQ(ierr);
2322fcdce8c4SStefano Zampini   if (!flg) SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Not for type %s",((PetscObject)A)->type_name);
2323fcdce8c4SStefano Zampini   ierr = PetscObjectTypeCompare((PetscObject)B,MATSEQAIJCUSPARSE,&flg);CHKERRQ(ierr);
2324fcdce8c4SStefano Zampini   if (!flg) SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Not for B of type %s",((PetscObject)B)->type_name);
2325fcdce8c4SStefano Zampini   a = (Mat_SeqAIJ*)A->data;
2326fcdce8c4SStefano Zampini   b = (Mat_SeqAIJ*)B->data;
2327fcdce8c4SStefano Zampini   Acusp = (Mat_SeqAIJCUSPARSE*)A->spptr;
2328fcdce8c4SStefano Zampini   Bcusp = (Mat_SeqAIJCUSPARSE*)B->spptr;
2329fcdce8c4SStefano Zampini   if (Acusp->format != MAT_CUSPARSE_CSR) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Only for MAT_CUSPARSE_CSR format");
2330fcdce8c4SStefano Zampini   if (Bcusp->format != MAT_CUSPARSE_CSR) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Only for MAT_CUSPARSE_CSR format");
2331fcdce8c4SStefano Zampini 
2332fcdce8c4SStefano Zampini   /* product data */
2333fcdce8c4SStefano Zampini   ierr = PetscNew(&mmdata);CHKERRQ(ierr);
2334fcdce8c4SStefano Zampini   C->product->data    = mmdata;
2335fcdce8c4SStefano Zampini   C->product->destroy = MatDestroy_MatMatCusparse;
2336fcdce8c4SStefano Zampini 
2337fcdce8c4SStefano Zampini   ierr = MatSeqAIJCUSPARSECopyToGPU(A);CHKERRQ(ierr);
2338fcdce8c4SStefano Zampini   ierr = MatSeqAIJCUSPARSECopyToGPU(B);CHKERRQ(ierr);
2339fcdce8c4SStefano Zampini   ptype = product->type;
2340fcdce8c4SStefano Zampini   if (A->symmetric && ptype == MATPRODUCT_AtB) ptype = MATPRODUCT_AB;
2341fcdce8c4SStefano Zampini   if (B->symmetric && ptype == MATPRODUCT_ABt) ptype = MATPRODUCT_AB;
2342fcdce8c4SStefano Zampini   biscompressed = PETSC_FALSE;
2343fcdce8c4SStefano Zampini   ciscompressed = PETSC_FALSE;
2344fcdce8c4SStefano Zampini   switch (ptype) {
2345fcdce8c4SStefano Zampini   case MATPRODUCT_AB:
2346fcdce8c4SStefano Zampini     m = A->rmap->n;
2347fcdce8c4SStefano Zampini     n = B->cmap->n;
2348fcdce8c4SStefano Zampini     k = A->cmap->n;
2349fcdce8c4SStefano Zampini     Amat = Acusp->mat;
2350fcdce8c4SStefano Zampini     Bmat = Bcusp->mat;
2351fcdce8c4SStefano Zampini     if (a->compressedrow.use) ciscompressed = PETSC_TRUE;
2352fcdce8c4SStefano Zampini     if (b->compressedrow.use) biscompressed = PETSC_TRUE;
2353fcdce8c4SStefano Zampini     break;
2354fcdce8c4SStefano Zampini   case MATPRODUCT_AtB:
2355fcdce8c4SStefano Zampini     m = A->cmap->n;
2356fcdce8c4SStefano Zampini     n = B->cmap->n;
2357fcdce8c4SStefano Zampini     k = A->rmap->n;
2358fcdce8c4SStefano Zampini     ierr = MatSeqAIJCUSPARSEGenerateTransposeForMult(A);CHKERRQ(ierr);
2359fcdce8c4SStefano Zampini     Amat = Acusp->matTranspose;
2360fcdce8c4SStefano Zampini     Bmat = Bcusp->mat;
2361fcdce8c4SStefano Zampini     if (b->compressedrow.use) biscompressed = PETSC_TRUE;
2362fcdce8c4SStefano Zampini     break;
2363fcdce8c4SStefano Zampini   case MATPRODUCT_ABt:
2364fcdce8c4SStefano Zampini     m = A->rmap->n;
2365fcdce8c4SStefano Zampini     n = B->rmap->n;
2366fcdce8c4SStefano Zampini     k = A->cmap->n;
2367fcdce8c4SStefano Zampini     ierr = MatSeqAIJCUSPARSEGenerateTransposeForMult(B);CHKERRQ(ierr);
2368fcdce8c4SStefano Zampini     Amat = Acusp->mat;
2369fcdce8c4SStefano Zampini     Bmat = Bcusp->matTranspose;
2370fcdce8c4SStefano Zampini     if (a->compressedrow.use) ciscompressed = PETSC_TRUE;
2371fcdce8c4SStefano Zampini     break;
2372fcdce8c4SStefano Zampini   default:
2373fcdce8c4SStefano Zampini     SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Unsupported product type %s",MatProductTypes[product->type]);
2374fcdce8c4SStefano Zampini   }
2375fcdce8c4SStefano Zampini 
2376fcdce8c4SStefano Zampini   /* create cusparse matrix */
2377fcdce8c4SStefano Zampini   ierr  = MatSetSizes(C,m,n,m,n);CHKERRQ(ierr);
2378fcdce8c4SStefano Zampini   ierr  = MatSetType(C,MATSEQAIJCUSPARSE);CHKERRQ(ierr);
2379fcdce8c4SStefano Zampini   c     = (Mat_SeqAIJ*)C->data;
2380fcdce8c4SStefano Zampini   Ccusp = (Mat_SeqAIJCUSPARSE*)C->spptr;
2381fcdce8c4SStefano Zampini   Cmat  = new Mat_SeqAIJCUSPARSEMultStruct;
2382fcdce8c4SStefano Zampini   Ccsr  = new CsrMatrix;
2383fcdce8c4SStefano Zampini 
2384fcdce8c4SStefano Zampini   c->compressedrow.use = ciscompressed;
2385fcdce8c4SStefano Zampini   if (c->compressedrow.use) { /* if a is in compressed row, than c will be in compressed row format */
2386fcdce8c4SStefano Zampini     c->compressedrow.nrows = a->compressedrow.nrows;
2387fcdce8c4SStefano Zampini     ierr = PetscMalloc2(c->compressedrow.nrows+1,&c->compressedrow.i,c->compressedrow.nrows,&c->compressedrow.rindex);CHKERRQ(ierr);
2388fcdce8c4SStefano Zampini     ierr = PetscArraycpy(c->compressedrow.rindex,a->compressedrow.rindex,c->compressedrow.nrows);CHKERRQ(ierr);
2389fcdce8c4SStefano Zampini     Ccusp->workVector  = new THRUSTARRAY(c->compressedrow.nrows);
2390fcdce8c4SStefano Zampini     Cmat->cprowIndices = new THRUSTINTARRAY(c->compressedrow.nrows);
2391fcdce8c4SStefano Zampini     Cmat->cprowIndices->assign(c->compressedrow.rindex,c->compressedrow.rindex + c->compressedrow.nrows);
2392fcdce8c4SStefano Zampini   } else {
2393fcdce8c4SStefano Zampini     c->compressedrow.nrows  = 0;
2394fcdce8c4SStefano Zampini     c->compressedrow.i      = NULL;
2395fcdce8c4SStefano Zampini     c->compressedrow.rindex = NULL;
2396fcdce8c4SStefano Zampini     Ccusp->workVector       = NULL;
2397fcdce8c4SStefano Zampini     Cmat->cprowIndices      = NULL;
2398fcdce8c4SStefano Zampini   }
2399fcdce8c4SStefano Zampini   Ccusp->nrows    = ciscompressed ? c->compressedrow.nrows : m;
2400fcdce8c4SStefano Zampini   Ccusp->mat      = Cmat;
2401fcdce8c4SStefano Zampini   Ccusp->mat->mat = Ccsr;
2402fcdce8c4SStefano Zampini   Ccsr->num_rows    = Ccusp->nrows;
2403fcdce8c4SStefano Zampini   Ccsr->num_cols    = n;
2404fcdce8c4SStefano Zampini   Ccsr->row_offsets = new THRUSTINTARRAY32(Ccusp->nrows+1);
2405fcdce8c4SStefano Zampini   stat = cusparseCreateMatDescr(&Cmat->descr);CHKERRCUSPARSE(stat);
2406fcdce8c4SStefano Zampini   stat = cusparseSetMatIndexBase(Cmat->descr, CUSPARSE_INDEX_BASE_ZERO);CHKERRCUSPARSE(stat);
2407fcdce8c4SStefano Zampini   stat = cusparseSetMatType(Cmat->descr, CUSPARSE_MATRIX_TYPE_GENERAL);CHKERRCUSPARSE(stat);
2408fcdce8c4SStefano Zampini   cerr = cudaMalloc((void **)&(Cmat->alpha_one),sizeof(PetscScalar));CHKERRCUDA(cerr);
2409fcdce8c4SStefano Zampini   cerr = cudaMalloc((void **)&(Cmat->beta_zero),sizeof(PetscScalar));CHKERRCUDA(cerr);
2410fcdce8c4SStefano Zampini   cerr = cudaMalloc((void **)&(Cmat->beta_one), sizeof(PetscScalar));CHKERRCUDA(cerr);
2411fcdce8c4SStefano Zampini   cerr = cudaMemcpy(Cmat->alpha_one,&PETSC_CUSPARSE_ONE, sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(cerr);
2412fcdce8c4SStefano Zampini   cerr = cudaMemcpy(Cmat->beta_zero,&PETSC_CUSPARSE_ZERO,sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(cerr);
2413fcdce8c4SStefano Zampini   cerr = cudaMemcpy(Cmat->beta_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(cerr);
2414fcdce8c4SStefano Zampini   if (!Ccsr->num_rows || !Ccsr->num_cols || !a->nz || !b->nz) { /* cusparse raise errors in different calls when matrices have zero rows/columns! */
2415fcdce8c4SStefano Zampini     thrust::fill(thrust::device,Ccsr->row_offsets->begin(),Ccsr->row_offsets->end(),0);
2416fcdce8c4SStefano Zampini     c->nz = 0;
2417fcdce8c4SStefano Zampini     Ccsr->column_indices = new THRUSTINTARRAY32(c->nz);
2418fcdce8c4SStefano Zampini     Ccsr->values = new THRUSTARRAY(c->nz);
2419fcdce8c4SStefano Zampini     goto finalizesym;
2420fcdce8c4SStefano Zampini   }
2421fcdce8c4SStefano Zampini 
2422fcdce8c4SStefano Zampini   if (!Amat) SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Missing A mult struct for product type %s",MatProductTypes[ptype]);
2423fcdce8c4SStefano Zampini   if (!Bmat) SETERRQ1(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Missing B mult struct for product type %s",MatProductTypes[ptype]);
2424fcdce8c4SStefano Zampini   Acsr = (CsrMatrix*)Amat->mat;
2425fcdce8c4SStefano Zampini   if (!biscompressed) {
2426fcdce8c4SStefano Zampini     Bcsr = (CsrMatrix*)Bmat->mat;
2427fcdce8c4SStefano Zampini #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
2428fcdce8c4SStefano Zampini     BmatSpDescr = Bmat->matDescr;
2429fcdce8c4SStefano Zampini #endif
2430fcdce8c4SStefano Zampini   } else { /* we need to use row offsets for the full matrix */
2431fcdce8c4SStefano Zampini     CsrMatrix *cBcsr = (CsrMatrix*)Bmat->mat;
2432fcdce8c4SStefano Zampini     Bcsr = new CsrMatrix;
2433fcdce8c4SStefano Zampini     Bcsr->num_rows       = B->rmap->n;
2434fcdce8c4SStefano Zampini     Bcsr->num_cols       = cBcsr->num_cols;
2435fcdce8c4SStefano Zampini     Bcsr->num_entries    = cBcsr->num_entries;
2436fcdce8c4SStefano Zampini     Bcsr->column_indices = cBcsr->column_indices;
2437fcdce8c4SStefano Zampini     Bcsr->values         = cBcsr->values;
2438fcdce8c4SStefano Zampini     if (!Bcusp->rowoffsets_gpu) {
2439fcdce8c4SStefano Zampini       Bcusp->rowoffsets_gpu  = new THRUSTINTARRAY32(B->rmap->n + 1);
2440fcdce8c4SStefano Zampini       Bcusp->rowoffsets_gpu->assign(b->i,b->i + B->rmap->n + 1);
2441fcdce8c4SStefano Zampini       ierr = PetscLogCpuToGpu((B->rmap->n + 1)*sizeof(PetscInt));CHKERRQ(ierr);
2442fcdce8c4SStefano Zampini     }
2443fcdce8c4SStefano Zampini     Bcsr->row_offsets = Bcusp->rowoffsets_gpu;
2444fcdce8c4SStefano Zampini     mmdata->Bcsr = Bcsr;
2445fcdce8c4SStefano Zampini #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
2446fcdce8c4SStefano Zampini     if (Bcsr->num_rows && Bcsr->num_cols) {
2447fcdce8c4SStefano Zampini       stat = cusparseCreateCsr(&mmdata->matSpBDescr, Bcsr->num_rows, Bcsr->num_cols, Bcsr->num_entries,
2448fcdce8c4SStefano Zampini                                Bcsr->row_offsets->data().get(), Bcsr->column_indices->data().get(),
2449fcdce8c4SStefano Zampini                                Bcsr->values->data().get(),
2450fcdce8c4SStefano Zampini                                CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I,
2451fcdce8c4SStefano Zampini                                CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype);CHKERRCUSPARSE(stat);
2452fcdce8c4SStefano Zampini     }
2453fcdce8c4SStefano Zampini     BmatSpDescr = mmdata->matSpBDescr;
2454fcdce8c4SStefano Zampini #endif
2455fcdce8c4SStefano Zampini   }
2456fcdce8c4SStefano Zampini   if (!Acsr) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Missing A CSR struct");
2457fcdce8c4SStefano Zampini   if (!Bcsr) SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Missing B CSR struct");
2458fcdce8c4SStefano Zampini   /* precompute flops count */
2459fcdce8c4SStefano Zampini   if (ptype == MATPRODUCT_AB) {
2460fcdce8c4SStefano Zampini     for (i=0, flops = 0; i<A->rmap->n; i++) {
2461fcdce8c4SStefano Zampini       const PetscInt st = a->i[i];
2462fcdce8c4SStefano Zampini       const PetscInt en = a->i[i+1];
2463fcdce8c4SStefano Zampini       for (j=st; j<en; j++) {
2464fcdce8c4SStefano Zampini         const PetscInt brow = a->j[j];
2465fcdce8c4SStefano Zampini         flops += 2.*(b->i[brow+1] - b->i[brow]);
2466fcdce8c4SStefano Zampini       }
2467fcdce8c4SStefano Zampini     }
2468fcdce8c4SStefano Zampini   } else if (ptype == MATPRODUCT_AtB) {
2469fcdce8c4SStefano Zampini     for (i=0, flops = 0; i<A->rmap->n; i++) {
2470fcdce8c4SStefano Zampini       const PetscInt anzi = a->i[i+1] - a->i[i];
2471fcdce8c4SStefano Zampini       const PetscInt bnzi = b->i[i+1] - b->i[i];
2472fcdce8c4SStefano Zampini       flops += (2.*anzi)*bnzi;
2473fcdce8c4SStefano Zampini     }
2474fcdce8c4SStefano Zampini   } else { /* TODO */
2475fcdce8c4SStefano Zampini     flops = 0.;
2476fcdce8c4SStefano Zampini   }
2477fcdce8c4SStefano Zampini 
2478fcdce8c4SStefano Zampini   mmdata->flops = flops;
2479fcdce8c4SStefano Zampini   ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
2480fcdce8c4SStefano Zampini #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
2481fcdce8c4SStefano Zampini   stat = cusparseSetPointerMode(Ccusp->handle, CUSPARSE_POINTER_MODE_DEVICE);CHKERRCUSPARSE(stat);
2482fcdce8c4SStefano Zampini   stat = cusparseCreateCsr(&Cmat->matDescr, Ccsr->num_rows, Ccsr->num_cols, 0,
2483fcdce8c4SStefano Zampini                            NULL, NULL, NULL,
2484fcdce8c4SStefano Zampini                            CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I,
2485fcdce8c4SStefano Zampini                            CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype);CHKERRCUSPARSE(stat);
2486fcdce8c4SStefano Zampini   stat = cusparseSpGEMM_createDescr(&mmdata->spgemmDesc);CHKERRCUSPARSE(stat);
2487fcdce8c4SStefano Zampini   /* ask bufferSize bytes for external memory */
2488fcdce8c4SStefano Zampini   stat = cusparseSpGEMM_workEstimation(Ccusp->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, CUSPARSE_OPERATION_NON_TRANSPOSE,
2489fcdce8c4SStefano Zampini                                        Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr,
2490fcdce8c4SStefano Zampini                                        cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT,
2491fcdce8c4SStefano Zampini                                        mmdata->spgemmDesc, &bufSize2, NULL);CHKERRCUSPARSE(stat);
2492*bfcc3627SStefano Zampini   cerr = cudaMalloc((void**) &mmdata->mmBuffer2, bufSize2);CHKERRCUDA(cerr);
2493fcdce8c4SStefano Zampini   /* inspect the matrices A and B to understand the memory requirement for the next step */
2494fcdce8c4SStefano Zampini   stat = cusparseSpGEMM_workEstimation(Ccusp->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, CUSPARSE_OPERATION_NON_TRANSPOSE,
2495fcdce8c4SStefano Zampini                                        Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr,
2496fcdce8c4SStefano Zampini                                        cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT,
2497fcdce8c4SStefano Zampini                                        mmdata->spgemmDesc, &bufSize2, mmdata->mmBuffer2);CHKERRCUSPARSE(stat);
2498fcdce8c4SStefano Zampini   /* ask bufferSize again bytes for external memory */
2499fcdce8c4SStefano Zampini   stat = cusparseSpGEMM_compute(Ccusp->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, CUSPARSE_OPERATION_NON_TRANSPOSE,
2500fcdce8c4SStefano Zampini                                 Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr,
2501fcdce8c4SStefano Zampini                                 cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT,
2502fcdce8c4SStefano Zampini                                 mmdata->spgemmDesc, &mmdata->mmBufferSize, NULL);CHKERRCUSPARSE(stat);
2503fcdce8c4SStefano Zampini   /* The CUSPARSE documentation is not clear, nor the API
2504fcdce8c4SStefano Zampini      We need both buffers to perform the operations properly!
2505fcdce8c4SStefano Zampini      mmdata->mmBuffer2 does not appear anywhere in the compute/copy API
2506fcdce8c4SStefano Zampini      it only appears for the workEstimation stuff, but it seems it is needed in compute, so probably the address
2507fcdce8c4SStefano Zampini      is stored in the descriptor! What a messy API... */
2508*bfcc3627SStefano Zampini   cerr = cudaMalloc((void**) &mmdata->mmBuffer, mmdata->mmBufferSize);CHKERRCUDA(cerr);
2509fcdce8c4SStefano Zampini   /* compute the intermediate product of A * B */
2510fcdce8c4SStefano Zampini   stat = cusparseSpGEMM_compute(Ccusp->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, CUSPARSE_OPERATION_NON_TRANSPOSE,
2511fcdce8c4SStefano Zampini                                 Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr,
2512fcdce8c4SStefano Zampini                                 cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT,
2513fcdce8c4SStefano Zampini                                 mmdata->spgemmDesc, &mmdata->mmBufferSize, mmdata->mmBuffer);CHKERRCUSPARSE(stat);
2514fcdce8c4SStefano Zampini   /* get matrix C non-zero entries C_nnz1 */
2515fcdce8c4SStefano Zampini   stat = cusparseSpMatGetSize(Cmat->matDescr, &C_num_rows1, &C_num_cols1, &C_nnz1);CHKERRCUSPARSE(stat);
2516fcdce8c4SStefano Zampini   c->nz = (PetscInt) C_nnz1;
2517fcdce8c4SStefano Zampini   Ccsr->column_indices = new THRUSTINTARRAY32(c->nz);
2518fcdce8c4SStefano Zampini   CHKERRCUDA(cudaPeekAtLastError()); /* catch out of memory errors */
2519fcdce8c4SStefano Zampini   Ccsr->values = new THRUSTARRAY(c->nz);
2520fcdce8c4SStefano Zampini   CHKERRCUDA(cudaPeekAtLastError()); /* catch out of memory errors */
2521fcdce8c4SStefano Zampini   stat = cusparseCsrSetPointers(Cmat->matDescr, Ccsr->row_offsets->data().get(), Ccsr->column_indices->data().get(),
2522fcdce8c4SStefano Zampini                                 Ccsr->values->data().get());CHKERRCUSPARSE(stat);
2523fcdce8c4SStefano Zampini   stat = cusparseSpGEMM_copy(Ccusp->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, CUSPARSE_OPERATION_NON_TRANSPOSE,
2524fcdce8c4SStefano Zampini                              Cmat->alpha_one, Amat->matDescr, BmatSpDescr, Cmat->beta_zero, Cmat->matDescr,
2525fcdce8c4SStefano Zampini                              cusparse_scalartype, CUSPARSE_SPGEMM_DEFAULT, mmdata->spgemmDesc);CHKERRCUSPARSE(stat);
2526fcdce8c4SStefano Zampini #else
2527fcdce8c4SStefano Zampini   stat = cusparseSetPointerMode(Ccusp->handle, CUSPARSE_POINTER_MODE_HOST);CHKERRCUSPARSE(stat);
2528fcdce8c4SStefano Zampini   stat = cusparseXcsrgemmNnz(Ccusp->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, CUSPARSE_OPERATION_NON_TRANSPOSE,
2529fcdce8c4SStefano Zampini                              Acsr->num_rows, Bcsr->num_cols, Acsr->num_cols,
2530fcdce8c4SStefano Zampini                              Amat->descr, Acsr->num_entries, Acsr->row_offsets->data().get(), Acsr->column_indices->data().get(),
2531fcdce8c4SStefano Zampini                              Bmat->descr, Bcsr->num_entries, Bcsr->row_offsets->data().get(), Bcsr->column_indices->data().get(),
2532fcdce8c4SStefano Zampini                              Cmat->descr, Ccsr->row_offsets->data().get(), &cnz);CHKERRCUSPARSE(stat);
2533fcdce8c4SStefano Zampini   c->nz = cnz;
2534fcdce8c4SStefano Zampini   Ccsr->column_indices = new THRUSTINTARRAY32(c->nz);
2535fcdce8c4SStefano Zampini   CHKERRCUDA(cudaPeekAtLastError()); /* catch out of memory errors */
2536fcdce8c4SStefano Zampini   Ccsr->values = new THRUSTARRAY(c->nz);
2537fcdce8c4SStefano Zampini   CHKERRCUDA(cudaPeekAtLastError()); /* catch out of memory errors */
2538fcdce8c4SStefano Zampini 
2539fcdce8c4SStefano Zampini   stat = cusparseSetPointerMode(Ccusp->handle, CUSPARSE_POINTER_MODE_DEVICE);CHKERRCUSPARSE(stat);
2540fcdce8c4SStefano Zampini   /* with the old gemm interface (removed from 11.0 on) we cannot compute the symbolic factorization only.
2541fcdce8c4SStefano Zampini      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
2542fcdce8c4SStefano Zampini      D is NULL, despite the fact that CUSPARSE documentation claims it is supported! */
2543fcdce8c4SStefano Zampini   stat = cusparse_csr_spgemm(Ccusp->handle, CUSPARSE_OPERATION_NON_TRANSPOSE, CUSPARSE_OPERATION_NON_TRANSPOSE,
2544fcdce8c4SStefano Zampini                              Acsr->num_rows, Bcsr->num_cols, Acsr->num_cols,
2545fcdce8c4SStefano Zampini                              Amat->descr, Acsr->num_entries, Acsr->values->data().get(), Acsr->row_offsets->data().get(), Acsr->column_indices->data().get(),
2546fcdce8c4SStefano Zampini                              Bmat->descr, Bcsr->num_entries, Bcsr->values->data().get(), Bcsr->row_offsets->data().get(), Bcsr->column_indices->data().get(),
2547fcdce8c4SStefano Zampini                              Cmat->descr, Ccsr->values->data().get(), Ccsr->row_offsets->data().get(), Ccsr->column_indices->data().get());CHKERRCUSPARSE(stat);
2548fcdce8c4SStefano Zampini #endif
2549fcdce8c4SStefano Zampini   cerr = WaitForCUDA();CHKERRCUDA(cerr);
2550fcdce8c4SStefano Zampini   ierr = PetscLogGpuFlops(mmdata->flops);CHKERRQ(ierr);
2551fcdce8c4SStefano Zampini   ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
2552fcdce8c4SStefano Zampini finalizesym:
2553fcdce8c4SStefano Zampini   c->singlemalloc = PETSC_FALSE;
2554fcdce8c4SStefano Zampini   c->free_a       = PETSC_TRUE;
2555fcdce8c4SStefano Zampini   c->free_ij      = PETSC_TRUE;
2556fcdce8c4SStefano Zampini   ierr = PetscMalloc1(m+1,&c->i);CHKERRQ(ierr);
2557fcdce8c4SStefano Zampini   ierr = PetscMalloc1(c->nz,&c->j);CHKERRQ(ierr);
2558fcdce8c4SStefano Zampini   if (PetscDefined(USE_64BIT_INDICES)) { /* 32 to 64 bit conversion on the GPU and then copy to host (lazy) */
2559fcdce8c4SStefano Zampini     PetscInt *d_i = c->i;
2560fcdce8c4SStefano Zampini     THRUSTINTARRAY ii(Ccsr->row_offsets->size());
2561fcdce8c4SStefano Zampini     THRUSTINTARRAY jj(Ccsr->column_indices->size());
2562fcdce8c4SStefano Zampini     ii   = *Ccsr->row_offsets;
2563fcdce8c4SStefano Zampini     jj   = *Ccsr->column_indices;
2564fcdce8c4SStefano Zampini     if (ciscompressed) d_i = c->compressedrow.i;
2565fcdce8c4SStefano Zampini     cerr = cudaMemcpy(d_i,ii.data().get(),Ccsr->row_offsets->size()*sizeof(PetscInt),cudaMemcpyDeviceToHost);CHKERRCUDA(cerr);
2566fcdce8c4SStefano Zampini     cerr = cudaMemcpy(c->j,jj.data().get(),Ccsr->column_indices->size()*sizeof(PetscInt),cudaMemcpyDeviceToHost);CHKERRCUDA(cerr);
2567fcdce8c4SStefano Zampini   } else {
2568fcdce8c4SStefano Zampini     PetscInt *d_i = c->i;
2569fcdce8c4SStefano Zampini     if (ciscompressed) d_i = c->compressedrow.i;
2570fcdce8c4SStefano Zampini     cerr = cudaMemcpy(d_i,Ccsr->row_offsets->data().get(),Ccsr->row_offsets->size()*sizeof(PetscInt),cudaMemcpyDeviceToHost);CHKERRCUDA(cerr);
2571fcdce8c4SStefano Zampini     cerr = cudaMemcpy(c->j,Ccsr->column_indices->data().get(),Ccsr->column_indices->size()*sizeof(PetscInt),cudaMemcpyDeviceToHost);CHKERRCUDA(cerr);
2572fcdce8c4SStefano Zampini   }
2573fcdce8c4SStefano Zampini   if (ciscompressed) { /* need to expand host row offsets */
2574fcdce8c4SStefano Zampini     PetscInt r = 0;
2575fcdce8c4SStefano Zampini     c->i[0] = 0;
2576fcdce8c4SStefano Zampini     for (k = 0; k < c->compressedrow.nrows; k++) {
2577fcdce8c4SStefano Zampini       const PetscInt next = c->compressedrow.rindex[k];
2578fcdce8c4SStefano Zampini       const PetscInt old = c->compressedrow.i[k];
2579fcdce8c4SStefano Zampini       for (; r < next; r++) c->i[r+1] = old;
2580fcdce8c4SStefano Zampini     }
2581fcdce8c4SStefano Zampini     for (; r < m; r++) c->i[r+1] = c->compressedrow.i[c->compressedrow.nrows];
2582fcdce8c4SStefano Zampini   }
2583fcdce8c4SStefano Zampini   ierr = PetscLogGpuToCpu((Ccsr->column_indices->size() + Ccsr->row_offsets->size())*sizeof(PetscInt));CHKERRQ(ierr);
2584fcdce8c4SStefano Zampini   ierr = PetscMalloc1(m,&c->ilen);CHKERRQ(ierr);
2585fcdce8c4SStefano Zampini   ierr = PetscMalloc1(m,&c->imax);CHKERRQ(ierr);
2586fcdce8c4SStefano Zampini   c->maxnz = c->nz;
2587fcdce8c4SStefano Zampini   c->nonzerorowcnt = 0;
2588fcdce8c4SStefano Zampini   c->rmax = 0;
2589fcdce8c4SStefano Zampini   for (k = 0; k < m; k++) {
2590fcdce8c4SStefano Zampini     const PetscInt nn = c->i[k+1] - c->i[k];
2591fcdce8c4SStefano Zampini     c->ilen[k] = c->imax[k] = nn;
2592fcdce8c4SStefano Zampini     c->nonzerorowcnt += (PetscInt)!!nn;
2593fcdce8c4SStefano Zampini     c->rmax = PetscMax(c->rmax,nn);
2594fcdce8c4SStefano Zampini   }
2595fcdce8c4SStefano Zampini   ierr = MatMarkDiagonal_SeqAIJ(C);CHKERRQ(ierr);
2596fcdce8c4SStefano Zampini   ierr = PetscMalloc1(c->nz,&c->a);CHKERRQ(ierr);
2597fcdce8c4SStefano Zampini   Ccsr->num_entries = c->nz;
2598fcdce8c4SStefano Zampini 
2599fcdce8c4SStefano Zampini   C->nonzerostate++;
2600fcdce8c4SStefano Zampini   ierr = PetscLayoutSetUp(C->rmap);CHKERRQ(ierr);
2601fcdce8c4SStefano Zampini   ierr = PetscLayoutSetUp(C->cmap);CHKERRQ(ierr);
2602fcdce8c4SStefano Zampini   Ccusp->nonzerostate = C->nonzerostate;
2603fcdce8c4SStefano Zampini   C->offloadmask   = PETSC_OFFLOAD_UNALLOCATED;
2604fcdce8c4SStefano Zampini   C->preallocated  = PETSC_TRUE;
2605fcdce8c4SStefano Zampini   C->assembled     = PETSC_FALSE;
2606fcdce8c4SStefano Zampini   C->was_assembled = PETSC_FALSE;
2607fcdce8c4SStefano Zampini   if (product->api_user) { /* flag the matrix C values as computed, so that the numeric phase will only call MatAssembly */
2608fcdce8c4SStefano Zampini     mmdata->reusesym = PETSC_TRUE;
2609fcdce8c4SStefano Zampini     C->offloadmask   = PETSC_OFFLOAD_GPU;
2610fcdce8c4SStefano Zampini   }
2611fcdce8c4SStefano Zampini   C->ops->productnumeric = MatProductNumeric_SeqAIJCUSPARSE_SeqAIJCUSPARSE;
2612fcdce8c4SStefano Zampini   PetscFunctionReturn(0);
2613fcdce8c4SStefano Zampini }
2614fcdce8c4SStefano Zampini 
2615fcdce8c4SStefano Zampini PETSC_INTERN PetscErrorCode MatProductSetFromOptions_SeqAIJ_SeqDense(Mat);
2616fcdce8c4SStefano Zampini 
2617fcdce8c4SStefano Zampini /* handles sparse or dense B */
2618fcdce8c4SStefano Zampini static PetscErrorCode MatProductSetFromOptions_SeqAIJCUSPARSE(Mat mat)
2619fcdce8c4SStefano Zampini {
2620fcdce8c4SStefano Zampini   Mat_Product    *product = mat->product;
2621fcdce8c4SStefano Zampini   PetscErrorCode ierr;
2622fcdce8c4SStefano Zampini   PetscBool      isdense = PETSC_FALSE,Biscusp = PETSC_FALSE,Ciscusp = PETSC_TRUE;
2623fcdce8c4SStefano Zampini 
2624fcdce8c4SStefano Zampini   PetscFunctionBegin;
2625fcdce8c4SStefano Zampini   MatCheckProduct(mat,1);
2626fcdce8c4SStefano Zampini   ierr = PetscObjectBaseTypeCompare((PetscObject)product->B,MATSEQDENSE,&isdense);CHKERRQ(ierr);
2627fcdce8c4SStefano Zampini   if (!product->B->boundtocpu) {
2628fcdce8c4SStefano Zampini     ierr = PetscObjectTypeCompare((PetscObject)product->B,MATSEQAIJCUSPARSE,&Biscusp);CHKERRQ(ierr);
2629fcdce8c4SStefano Zampini   }
2630fcdce8c4SStefano Zampini   if (product->type == MATPRODUCT_ABC) {
2631fcdce8c4SStefano Zampini     Ciscusp = PETSC_FALSE;
2632fcdce8c4SStefano Zampini     if (!product->C->boundtocpu) {
2633fcdce8c4SStefano Zampini       ierr = PetscObjectTypeCompare((PetscObject)product->C,MATSEQAIJCUSPARSE,&Ciscusp);CHKERRQ(ierr);
2634fcdce8c4SStefano Zampini     }
2635fcdce8c4SStefano Zampini   }
2636fcdce8c4SStefano Zampini   if (isdense) {
2637ccdfe979SStefano Zampini     switch (product->type) {
2638ccdfe979SStefano Zampini     case MATPRODUCT_AB:
2639ccdfe979SStefano Zampini     case MATPRODUCT_AtB:
2640ccdfe979SStefano Zampini     case MATPRODUCT_ABt:
2641ccdfe979SStefano Zampini     case MATPRODUCT_PtAP:
2642ccdfe979SStefano Zampini     case MATPRODUCT_RARt:
2643fcdce8c4SStefano Zampini      if (product->A->boundtocpu) {
2644fcdce8c4SStefano Zampini         ierr = MatProductSetFromOptions_SeqAIJ_SeqDense(mat);CHKERRQ(ierr);
2645fcdce8c4SStefano Zampini       } else {
2646fcdce8c4SStefano Zampini         mat->ops->productsymbolic = MatProductSymbolic_SeqAIJCUSPARSE_SeqDENSECUDA;
2647fcdce8c4SStefano Zampini       }
2648fcdce8c4SStefano Zampini       break;
2649fcdce8c4SStefano Zampini     case MATPRODUCT_ABC:
2650fcdce8c4SStefano Zampini       mat->ops->productsymbolic = MatProductSymbolic_ABC_Basic;
2651fcdce8c4SStefano Zampini       break;
2652ccdfe979SStefano Zampini     default:
2653ccdfe979SStefano Zampini       break;
2654ccdfe979SStefano Zampini     }
2655fcdce8c4SStefano Zampini   } else if (Biscusp && Ciscusp) {
2656fcdce8c4SStefano Zampini     switch (product->type) {
2657fcdce8c4SStefano Zampini     case MATPRODUCT_AB:
2658fcdce8c4SStefano Zampini     case MATPRODUCT_AtB:
2659fcdce8c4SStefano Zampini     case MATPRODUCT_ABt:
2660fcdce8c4SStefano Zampini       mat->ops->productsymbolic = MatProductSymbolic_SeqAIJCUSPARSE_SeqAIJCUSPARSE;
2661fcdce8c4SStefano Zampini       break;
2662fcdce8c4SStefano Zampini     case MATPRODUCT_PtAP:
2663fcdce8c4SStefano Zampini     case MATPRODUCT_RARt:
2664fcdce8c4SStefano Zampini     case MATPRODUCT_ABC:
2665fcdce8c4SStefano Zampini       mat->ops->productsymbolic = MatProductSymbolic_ABC_Basic;
2666fcdce8c4SStefano Zampini       break;
2667fcdce8c4SStefano Zampini     default:
2668fcdce8c4SStefano Zampini       break;
2669fcdce8c4SStefano Zampini     }
2670fcdce8c4SStefano Zampini   } else { /* fallback for AIJ */
2671fcdce8c4SStefano Zampini     ierr = MatProductSetFromOptions_SeqAIJ(mat);CHKERRQ(ierr);
2672fcdce8c4SStefano Zampini   }
2673ccdfe979SStefano Zampini   PetscFunctionReturn(0);
2674ccdfe979SStefano Zampini }
2675ccdfe979SStefano Zampini 
26766fa9248bSJed Brown static PetscErrorCode MatMult_SeqAIJCUSPARSE(Mat A,Vec xx,Vec yy)
26779ae82921SPaul Mullowney {
2678b175d8bbSPaul Mullowney   PetscErrorCode ierr;
26799ae82921SPaul Mullowney 
26809ae82921SPaul Mullowney   PetscFunctionBegin;
2681e6e9a74fSStefano Zampini   ierr = MatMultAddKernel_SeqAIJCUSPARSE(A,xx,NULL,yy,PETSC_FALSE,PETSC_FALSE);CHKERRQ(ierr);
2682e6e9a74fSStefano Zampini   PetscFunctionReturn(0);
2683e6e9a74fSStefano Zampini }
2684e6e9a74fSStefano Zampini 
2685e6e9a74fSStefano Zampini static PetscErrorCode MatMultAdd_SeqAIJCUSPARSE(Mat A,Vec xx,Vec yy, Vec zz)
2686e6e9a74fSStefano Zampini {
2687e6e9a74fSStefano Zampini   PetscErrorCode ierr;
2688e6e9a74fSStefano Zampini 
2689e6e9a74fSStefano Zampini   PetscFunctionBegin;
2690e6e9a74fSStefano Zampini   ierr = MatMultAddKernel_SeqAIJCUSPARSE(A,xx,yy,zz,PETSC_FALSE,PETSC_FALSE);CHKERRQ(ierr);
2691e6e9a74fSStefano Zampini   PetscFunctionReturn(0);
2692e6e9a74fSStefano Zampini }
2693e6e9a74fSStefano Zampini 
2694e6e9a74fSStefano Zampini static PetscErrorCode MatMultHermitianTranspose_SeqAIJCUSPARSE(Mat A,Vec xx,Vec yy)
2695e6e9a74fSStefano Zampini {
2696e6e9a74fSStefano Zampini   PetscErrorCode ierr;
2697e6e9a74fSStefano Zampini 
2698e6e9a74fSStefano Zampini   PetscFunctionBegin;
2699e6e9a74fSStefano Zampini   ierr = MatMultAddKernel_SeqAIJCUSPARSE(A,xx,NULL,yy,PETSC_TRUE,PETSC_TRUE);CHKERRQ(ierr);
2700e6e9a74fSStefano Zampini   PetscFunctionReturn(0);
2701e6e9a74fSStefano Zampini }
2702e6e9a74fSStefano Zampini 
2703e6e9a74fSStefano Zampini static PetscErrorCode MatMultHermitianTransposeAdd_SeqAIJCUSPARSE(Mat A,Vec xx,Vec yy,Vec zz)
2704e6e9a74fSStefano Zampini {
2705e6e9a74fSStefano Zampini   PetscErrorCode ierr;
2706e6e9a74fSStefano Zampini 
2707e6e9a74fSStefano Zampini   PetscFunctionBegin;
2708e6e9a74fSStefano Zampini   ierr = MatMultAddKernel_SeqAIJCUSPARSE(A,xx,yy,zz,PETSC_TRUE,PETSC_TRUE);CHKERRQ(ierr);
27099ae82921SPaul Mullowney   PetscFunctionReturn(0);
27109ae82921SPaul Mullowney }
27119ae82921SPaul Mullowney 
27126fa9248bSJed Brown static PetscErrorCode MatMultTranspose_SeqAIJCUSPARSE(Mat A,Vec xx,Vec yy)
2713ca45077fSPaul Mullowney {
2714b175d8bbSPaul Mullowney   PetscErrorCode ierr;
2715ca45077fSPaul Mullowney 
2716ca45077fSPaul Mullowney   PetscFunctionBegin;
2717e6e9a74fSStefano Zampini   ierr = MatMultAddKernel_SeqAIJCUSPARSE(A,xx,NULL,yy,PETSC_TRUE,PETSC_FALSE);CHKERRQ(ierr);
2718ca45077fSPaul Mullowney   PetscFunctionReturn(0);
2719ca45077fSPaul Mullowney }
2720ca45077fSPaul Mullowney 
2721afb2bd1cSJunchao Zhang /* z = op(A) x + y. If trans & !herm, op = ^T; if trans & herm, op = ^H; if !trans, op = no-op */
2722e6e9a74fSStefano Zampini static PetscErrorCode MatMultAddKernel_SeqAIJCUSPARSE(Mat A,Vec xx,Vec yy,Vec zz,PetscBool trans,PetscBool herm)
27239ae82921SPaul Mullowney {
27249ae82921SPaul Mullowney   Mat_SeqAIJ                   *a = (Mat_SeqAIJ*)A->data;
2725aa372e3fSPaul Mullowney   Mat_SeqAIJCUSPARSE           *cusparsestruct = (Mat_SeqAIJCUSPARSE*)A->spptr;
27269ff858a8SKarl Rupp   Mat_SeqAIJCUSPARSEMultStruct *matstruct;
2727e6e9a74fSStefano Zampini   PetscScalar                  *xarray,*zarray,*dptr,*beta,*xptr;
2728b175d8bbSPaul Mullowney   PetscErrorCode               ierr;
272957d48284SJunchao Zhang   cudaError_t                  cerr;
2730aa372e3fSPaul Mullowney   cusparseStatus_t             stat;
2731e6e9a74fSStefano Zampini   cusparseOperation_t          opA = CUSPARSE_OPERATION_NON_TRANSPOSE;
2732e6e9a74fSStefano Zampini   PetscBool                    compressed;
2733afb2bd1cSJunchao Zhang #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
2734afb2bd1cSJunchao Zhang   PetscInt                     nx,ny;
2735afb2bd1cSJunchao Zhang #endif
27366e111a19SKarl Rupp 
27379ae82921SPaul Mullowney   PetscFunctionBegin;
2738e6e9a74fSStefano Zampini   if (herm && !trans) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_PLIB,"Hermitian and not transpose not supported");
2739e6e9a74fSStefano Zampini   if (!a->nonzerorowcnt) {
2740afb2bd1cSJunchao Zhang     if (!yy) {ierr = VecSet_SeqCUDA(zz,0);CHKERRQ(ierr);}
2741d38a13f6SStefano Zampini     else {ierr = VecCopy_SeqCUDA(yy,zz);CHKERRQ(ierr);}
2742e6e9a74fSStefano Zampini     PetscFunctionReturn(0);
2743e6e9a74fSStefano Zampini   }
274434d6c7a5SJose E. Roman   /* The line below is necessary due to the operations that modify the matrix on the CPU (axpy, scale, etc) */
274534d6c7a5SJose E. Roman   ierr = MatSeqAIJCUSPARSECopyToGPU(A);CHKERRQ(ierr);
2746e6e9a74fSStefano Zampini   if (!trans) {
27479ff858a8SKarl Rupp     matstruct = (Mat_SeqAIJCUSPARSEMultStruct*)cusparsestruct->mat;
2748c9567895SMark     if (!matstruct) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_PLIB,"SeqAIJCUSPARSE does not have a 'mat' (need to fix)");
2749e6e9a74fSStefano Zampini   } else {
2750e6e9a74fSStefano Zampini     if (herm || !cusparsestruct->transgen) {
2751e6e9a74fSStefano Zampini       opA = herm ? CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE : CUSPARSE_OPERATION_TRANSPOSE;
2752e6e9a74fSStefano Zampini       matstruct = (Mat_SeqAIJCUSPARSEMultStruct*)cusparsestruct->mat;
2753e6e9a74fSStefano Zampini     } else {
2754afb2bd1cSJunchao Zhang       if (!cusparsestruct->matTranspose) {ierr = MatSeqAIJCUSPARSEGenerateTransposeForMult(A);CHKERRQ(ierr);}
2755e6e9a74fSStefano Zampini       matstruct = (Mat_SeqAIJCUSPARSEMultStruct*)cusparsestruct->matTranspose;
2756e6e9a74fSStefano Zampini     }
2757e6e9a74fSStefano Zampini   }
2758e6e9a74fSStefano Zampini   /* Does the matrix use compressed rows (i.e., drop zero rows)? */
2759e6e9a74fSStefano Zampini   compressed = matstruct->cprowIndices ? PETSC_TRUE : PETSC_FALSE;
2760213423ffSJunchao Zhang 
2761e6e9a74fSStefano Zampini   try {
2762e6e9a74fSStefano Zampini     ierr = VecCUDAGetArrayRead(xx,(const PetscScalar**)&xarray);CHKERRQ(ierr);
2763213423ffSJunchao Zhang     if (yy == zz) {ierr = VecCUDAGetArray(zz,&zarray);CHKERRQ(ierr);} /* read & write zz, so need to get uptodate zarray on GPU */
2764213423ffSJunchao Zhang     else {ierr = VecCUDAGetArrayWrite(zz,&zarray);CHKERRQ(ierr);} /* write zz, so no need to init zarray on GPU */
2765afb2bd1cSJunchao Zhang 
276685ba7357SStefano Zampini     ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
2767e6e9a74fSStefano Zampini     if (opA == CUSPARSE_OPERATION_NON_TRANSPOSE) {
2768afb2bd1cSJunchao Zhang       /* z = A x + beta y.
2769afb2bd1cSJunchao Zhang          If A is compressed (with less rows), then Ax is shorter than the full z, so we need a work vector to store Ax.
2770afb2bd1cSJunchao Zhang          When A is non-compressed, and z = y, we can set beta=1 to compute y = Ax + y in one call.
2771afb2bd1cSJunchao Zhang       */
2772e6e9a74fSStefano Zampini       xptr = xarray;
2773afb2bd1cSJunchao Zhang       dptr = compressed ? cusparsestruct->workVector->data().get() : zarray;
2774213423ffSJunchao Zhang       beta = (yy == zz && !compressed) ? matstruct->beta_one : matstruct->beta_zero;
2775afb2bd1cSJunchao Zhang      #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
2776afb2bd1cSJunchao Zhang       /* Get length of x, y for y=Ax. ny might be shorter than the work vector's allocated length, since the work vector is
2777afb2bd1cSJunchao Zhang           allocated to accommodate different uses. So we get the length info directly from mat.
2778afb2bd1cSJunchao Zhang        */
2779afb2bd1cSJunchao Zhang       if (cusparsestruct->format == MAT_CUSPARSE_CSR) {
2780afb2bd1cSJunchao Zhang         CsrMatrix *mat = (CsrMatrix*)matstruct->mat;
2781afb2bd1cSJunchao Zhang         nx = mat->num_cols;
2782afb2bd1cSJunchao Zhang         ny = mat->num_rows;
2783afb2bd1cSJunchao Zhang       }
2784afb2bd1cSJunchao Zhang      #endif
2785e6e9a74fSStefano Zampini     } else {
2786afb2bd1cSJunchao Zhang       /* z = A^T x + beta y
2787afb2bd1cSJunchao Zhang          If A is compressed, then we need a work vector as the shorter version of x to compute A^T x.
2788afb2bd1cSJunchao Zhang          Note A^Tx is of full length, so we set beta to 1.0 if y exists.
2789afb2bd1cSJunchao Zhang        */
2790afb2bd1cSJunchao Zhang       xptr = compressed ? cusparsestruct->workVector->data().get() : xarray;
2791e6e9a74fSStefano Zampini       dptr = zarray;
2792e6e9a74fSStefano Zampini       beta = yy ? matstruct->beta_one : matstruct->beta_zero;
2793afb2bd1cSJunchao Zhang       if (compressed) { /* Scatter x to work vector */
2794e6e9a74fSStefano Zampini         thrust::device_ptr<PetscScalar> xarr = thrust::device_pointer_cast(xarray);
2795e6e9a74fSStefano Zampini         thrust::for_each(thrust::make_zip_iterator(thrust::make_tuple(cusparsestruct->workVector->begin(), thrust::make_permutation_iterator(xarr, matstruct->cprowIndices->begin()))),
2796e6e9a74fSStefano Zampini                          thrust::make_zip_iterator(thrust::make_tuple(cusparsestruct->workVector->begin(), thrust::make_permutation_iterator(xarr, matstruct->cprowIndices->begin()))) + matstruct->cprowIndices->size(),
2797e6e9a74fSStefano Zampini                          VecCUDAEqualsReverse());
2798e6e9a74fSStefano Zampini       }
2799afb2bd1cSJunchao Zhang      #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
2800afb2bd1cSJunchao Zhang       if (cusparsestruct->format == MAT_CUSPARSE_CSR) {
2801afb2bd1cSJunchao Zhang         CsrMatrix *mat = (CsrMatrix*)matstruct->mat;
2802afb2bd1cSJunchao Zhang         nx = mat->num_rows;
2803afb2bd1cSJunchao Zhang         ny = mat->num_cols;
2804afb2bd1cSJunchao Zhang       }
2805afb2bd1cSJunchao Zhang      #endif
2806e6e9a74fSStefano Zampini     }
28079ae82921SPaul Mullowney 
2808afb2bd1cSJunchao Zhang     /* csr_spmv does y = alpha op(A) x + beta y */
2809aa372e3fSPaul Mullowney     if (cusparsestruct->format == MAT_CUSPARSE_CSR) {
2810afb2bd1cSJunchao Zhang      #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
2811afb2bd1cSJunchao Zhang       if (opA < 0 || opA > 2) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"cuSPARSE ABI on cusparseOperation_t has changed and PETSc has not been updated accordingly");
2812afb2bd1cSJunchao Zhang       if (!matstruct->cuSpMV[opA].initialized) { /* built on demand */
2813afb2bd1cSJunchao Zhang         stat = cusparseCreateDnVec(&matstruct->cuSpMV[opA].vecXDescr,nx,xptr,cusparse_scalartype);CHKERRCUSPARSE(stat);
2814afb2bd1cSJunchao Zhang         stat = cusparseCreateDnVec(&matstruct->cuSpMV[opA].vecYDescr,ny,dptr,cusparse_scalartype);CHKERRCUSPARSE(stat);
2815afb2bd1cSJunchao Zhang         stat = cusparseSpMV_bufferSize(cusparsestruct->handle, opA, matstruct->alpha_one,
2816afb2bd1cSJunchao Zhang                                 matstruct->matDescr,
2817afb2bd1cSJunchao Zhang                                 matstruct->cuSpMV[opA].vecXDescr, beta,
2818afb2bd1cSJunchao Zhang                                 matstruct->cuSpMV[opA].vecYDescr,
2819afb2bd1cSJunchao Zhang                                 cusparse_scalartype,
2820afb2bd1cSJunchao Zhang                                 cusparsestruct->spmvAlg,
2821afb2bd1cSJunchao Zhang                                 &matstruct->cuSpMV[opA].spmvBufferSize);CHKERRCUSPARSE(stat);
2822afb2bd1cSJunchao Zhang         cerr = cudaMalloc(&matstruct->cuSpMV[opA].spmvBuffer,matstruct->cuSpMV[opA].spmvBufferSize);CHKERRCUDA(cerr);
2823afb2bd1cSJunchao Zhang 
2824afb2bd1cSJunchao Zhang         matstruct->cuSpMV[opA].initialized = PETSC_TRUE;
2825afb2bd1cSJunchao Zhang       } else {
2826afb2bd1cSJunchao Zhang         /* x, y's value pointers might change between calls, but their shape is kept, so we just update pointers */
2827afb2bd1cSJunchao Zhang         stat = cusparseDnVecSetValues(matstruct->cuSpMV[opA].vecXDescr,xptr);CHKERRCUSPARSE(stat);
2828afb2bd1cSJunchao Zhang         stat = cusparseDnVecSetValues(matstruct->cuSpMV[opA].vecYDescr,dptr);CHKERRCUSPARSE(stat);
2829afb2bd1cSJunchao Zhang       }
2830afb2bd1cSJunchao Zhang 
2831afb2bd1cSJunchao Zhang       stat = cusparseSpMV(cusparsestruct->handle, opA,
2832afb2bd1cSJunchao Zhang                                matstruct->alpha_one,
2833afb2bd1cSJunchao Zhang                                matstruct->matDescr, /* built in MatSeqAIJCUSPARSECopyToGPU() or MatSeqAIJCUSPARSEGenerateTransposeForMult() */
2834afb2bd1cSJunchao Zhang                                matstruct->cuSpMV[opA].vecXDescr,
2835afb2bd1cSJunchao Zhang                                beta,
2836afb2bd1cSJunchao Zhang                                matstruct->cuSpMV[opA].vecYDescr,
2837afb2bd1cSJunchao Zhang                                cusparse_scalartype,
2838afb2bd1cSJunchao Zhang                                cusparsestruct->spmvAlg,
2839afb2bd1cSJunchao Zhang                                matstruct->cuSpMV[opA].spmvBuffer);CHKERRCUSPARSE(stat);
2840afb2bd1cSJunchao Zhang      #else
28417656d835SStefano Zampini       CsrMatrix *mat = (CsrMatrix*)matstruct->mat;
2842e6e9a74fSStefano Zampini       stat = cusparse_csr_spmv(cusparsestruct->handle, opA,
2843a65300a6SPaul Mullowney                                mat->num_rows, mat->num_cols,
2844afb2bd1cSJunchao Zhang                                mat->num_entries, matstruct->alpha_one, matstruct->descr,
2845aa372e3fSPaul Mullowney                                mat->values->data().get(), mat->row_offsets->data().get(),
2846e6e9a74fSStefano Zampini                                mat->column_indices->data().get(), xptr, beta,
284757d48284SJunchao Zhang                                dptr);CHKERRCUSPARSE(stat);
2848afb2bd1cSJunchao Zhang      #endif
2849aa372e3fSPaul Mullowney     } else {
2850213423ffSJunchao Zhang       if (cusparsestruct->nrows) {
2851afb2bd1cSJunchao Zhang        #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
2852afb2bd1cSJunchao Zhang         SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"MAT_CUSPARSE_ELL and MAT_CUSPARSE_HYB are not supported since CUDA-11.0");
2853afb2bd1cSJunchao Zhang        #else
2854301298b4SMark Adams         cusparseHybMat_t hybMat = (cusparseHybMat_t)matstruct->mat;
2855e6e9a74fSStefano Zampini         stat = cusparse_hyb_spmv(cusparsestruct->handle, opA,
2856afb2bd1cSJunchao Zhang                                  matstruct->alpha_one, matstruct->descr, hybMat,
2857e6e9a74fSStefano Zampini                                  xptr, beta,
285857d48284SJunchao Zhang                                  dptr);CHKERRCUSPARSE(stat);
2859afb2bd1cSJunchao Zhang        #endif
2860a65300a6SPaul Mullowney       }
2861aa372e3fSPaul Mullowney     }
286205035670SJunchao Zhang     cerr = WaitForCUDA();CHKERRCUDA(cerr);
2863958c4211Shannah_mairs     ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
2864aa372e3fSPaul Mullowney 
2865e6e9a74fSStefano Zampini     if (opA == CUSPARSE_OPERATION_NON_TRANSPOSE) {
2866213423ffSJunchao Zhang       if (yy) { /* MatMultAdd: zz = A*xx + yy */
2867213423ffSJunchao Zhang         if (compressed) { /* A is compressed. We first copy yy to zz, then ScatterAdd the work vector to zz */
2868213423ffSJunchao Zhang           ierr = VecCopy_SeqCUDA(yy,zz);CHKERRQ(ierr); /* zz = yy */
2869e6e9a74fSStefano Zampini         } else if (zz != yy) { /* A is not compressed. zz already contains A*xx, and we just need to add yy */
2870213423ffSJunchao Zhang           ierr = VecAXPY_SeqCUDA(zz,1.0,yy);CHKERRQ(ierr); /* zz += yy */
28717656d835SStefano Zampini         }
2872213423ffSJunchao Zhang       } else if (compressed) { /* MatMult: zz = A*xx. A is compressed, so we zero zz first, then ScatterAdd the work vector to zz */
2873c1fb3f03SStefano Zampini         ierr = VecSet_SeqCUDA(zz,0);CHKERRQ(ierr);
28747656d835SStefano Zampini       }
28757656d835SStefano Zampini 
2876213423ffSJunchao Zhang       /* ScatterAdd the result from work vector into the full vector when A is compressed */
2877213423ffSJunchao Zhang       if (compressed) {
2878213423ffSJunchao Zhang         thrust::device_ptr<PetscScalar> zptr = thrust::device_pointer_cast(zarray);
2879e6e9a74fSStefano Zampini         ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
2880c41cb2e2SAlejandro Lamas Daviña         thrust::for_each(thrust::make_zip_iterator(thrust::make_tuple(cusparsestruct->workVector->begin(), thrust::make_permutation_iterator(zptr, matstruct->cprowIndices->begin()))),
2881e6e9a74fSStefano Zampini                          thrust::make_zip_iterator(thrust::make_tuple(cusparsestruct->workVector->begin(), thrust::make_permutation_iterator(zptr, matstruct->cprowIndices->begin()))) + matstruct->cprowIndices->size(),
2882c41cb2e2SAlejandro Lamas Daviña                          VecCUDAPlusEquals());
288305035670SJunchao Zhang         cerr = WaitForCUDA();CHKERRCUDA(cerr);
2884958c4211Shannah_mairs         ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
2885e6e9a74fSStefano Zampini       }
2886e6e9a74fSStefano Zampini     } else {
2887e6e9a74fSStefano Zampini       if (yy && yy != zz) {
2888e6e9a74fSStefano Zampini         ierr = VecAXPY_SeqCUDA(zz,1.0,yy);CHKERRQ(ierr); /* zz += yy */
2889e6e9a74fSStefano Zampini       }
2890e6e9a74fSStefano Zampini     }
2891e6e9a74fSStefano Zampini     ierr = VecCUDARestoreArrayRead(xx,(const PetscScalar**)&xarray);CHKERRQ(ierr);
2892213423ffSJunchao Zhang     if (yy == zz) {ierr = VecCUDARestoreArray(zz,&zarray);CHKERRQ(ierr);}
2893213423ffSJunchao Zhang     else {ierr = VecCUDARestoreArrayWrite(zz,&zarray);CHKERRQ(ierr);}
28949ae82921SPaul Mullowney   } catch(char *ex) {
28959ae82921SPaul Mullowney     SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_LIB,"CUSPARSE error: %s", ex);
28969ae82921SPaul Mullowney   }
2897e6e9a74fSStefano Zampini   if (yy) {
2898958c4211Shannah_mairs     ierr = PetscLogGpuFlops(2.0*a->nz);CHKERRQ(ierr);
2899e6e9a74fSStefano Zampini   } else {
2900e6e9a74fSStefano Zampini     ierr = PetscLogGpuFlops(2.0*a->nz-a->nonzerorowcnt);CHKERRQ(ierr);
2901e6e9a74fSStefano Zampini   }
29029ae82921SPaul Mullowney   PetscFunctionReturn(0);
29039ae82921SPaul Mullowney }
29049ae82921SPaul Mullowney 
29056fa9248bSJed Brown static PetscErrorCode MatMultTransposeAdd_SeqAIJCUSPARSE(Mat A,Vec xx,Vec yy,Vec zz)
2906ca45077fSPaul Mullowney {
2907b175d8bbSPaul Mullowney   PetscErrorCode ierr;
29086e111a19SKarl Rupp 
2909ca45077fSPaul Mullowney   PetscFunctionBegin;
2910e6e9a74fSStefano Zampini   ierr = MatMultAddKernel_SeqAIJCUSPARSE(A,xx,yy,zz,PETSC_TRUE,PETSC_FALSE);CHKERRQ(ierr);
2911ca45077fSPaul Mullowney   PetscFunctionReturn(0);
2912ca45077fSPaul Mullowney }
2913ca45077fSPaul Mullowney 
29146fa9248bSJed Brown static PetscErrorCode MatAssemblyEnd_SeqAIJCUSPARSE(Mat A,MatAssemblyType mode)
29159ae82921SPaul Mullowney {
29169ae82921SPaul Mullowney   PetscErrorCode              ierr;
2917a587d139SMark   PetscSplitCSRDataStructure  *d_mat = NULL;
29189ae82921SPaul Mullowney   PetscFunctionBegin;
2919bc3f50f2SPaul Mullowney   if (A->factortype == MAT_FACTOR_NONE) {
29203fa6b06aSMark Adams     d_mat = ((Mat_SeqAIJCUSPARSE*)A->spptr)->deviceMat;
2921bc3f50f2SPaul Mullowney   }
29223fa6b06aSMark Adams   ierr = MatAssemblyEnd_SeqAIJ(A,mode);CHKERRQ(ierr); // this does very little if assembled on GPU - call it?
29233fa6b06aSMark Adams   if (mode == MAT_FLUSH_ASSEMBLY || A->boundtocpu) PetscFunctionReturn(0);
2924a587d139SMark   if (d_mat) {
29253fa6b06aSMark Adams     A->offloadmask = PETSC_OFFLOAD_GPU;
29263fa6b06aSMark Adams   }
29273fa6b06aSMark Adams 
29289ae82921SPaul Mullowney   PetscFunctionReturn(0);
29299ae82921SPaul Mullowney }
29309ae82921SPaul Mullowney 
29319ae82921SPaul Mullowney /* --------------------------------------------------------------------------------*/
2932e057df02SPaul Mullowney /*@
29339ae82921SPaul Mullowney    MatCreateSeqAIJCUSPARSE - Creates a sparse matrix in AIJ (compressed row) format
2934e057df02SPaul Mullowney    (the default parallel PETSc format). This matrix will ultimately pushed down
2935e057df02SPaul Mullowney    to NVidia GPUs and use the CUSPARSE library for calculations. For good matrix
2936e057df02SPaul Mullowney    assembly performance the user should preallocate the matrix storage by setting
2937e057df02SPaul Mullowney    the parameter nz (or the array nnz).  By setting these parameters accurately,
2938e057df02SPaul Mullowney    performance during matrix assembly can be increased by more than a factor of 50.
29399ae82921SPaul Mullowney 
2940d083f849SBarry Smith    Collective
29419ae82921SPaul Mullowney 
29429ae82921SPaul Mullowney    Input Parameters:
29439ae82921SPaul Mullowney +  comm - MPI communicator, set to PETSC_COMM_SELF
29449ae82921SPaul Mullowney .  m - number of rows
29459ae82921SPaul Mullowney .  n - number of columns
29469ae82921SPaul Mullowney .  nz - number of nonzeros per row (same for all rows)
29479ae82921SPaul Mullowney -  nnz - array containing the number of nonzeros in the various rows
29480298fd71SBarry Smith          (possibly different for each row) or NULL
29499ae82921SPaul Mullowney 
29509ae82921SPaul Mullowney    Output Parameter:
29519ae82921SPaul Mullowney .  A - the matrix
29529ae82921SPaul Mullowney 
29539ae82921SPaul Mullowney    It is recommended that one use the MatCreate(), MatSetType() and/or MatSetFromOptions(),
29549ae82921SPaul Mullowney    MatXXXXSetPreallocation() paradgm instead of this routine directly.
29559ae82921SPaul Mullowney    [MatXXXXSetPreallocation() is, for example, MatSeqAIJSetPreallocation]
29569ae82921SPaul Mullowney 
29579ae82921SPaul Mullowney    Notes:
29589ae82921SPaul Mullowney    If nnz is given then nz is ignored
29599ae82921SPaul Mullowney 
29609ae82921SPaul Mullowney    The AIJ format (also called the Yale sparse matrix format or
29619ae82921SPaul Mullowney    compressed row storage), is fully compatible with standard Fortran 77
29629ae82921SPaul Mullowney    storage.  That is, the stored row and column indices can begin at
29639ae82921SPaul Mullowney    either one (as in Fortran) or zero.  See the users' manual for details.
29649ae82921SPaul Mullowney 
29659ae82921SPaul Mullowney    Specify the preallocated storage with either nz or nnz (not both).
29660298fd71SBarry Smith    Set nz=PETSC_DEFAULT and nnz=NULL for PETSc to control dynamic memory
29679ae82921SPaul Mullowney    allocation.  For large problems you MUST preallocate memory or you
29689ae82921SPaul Mullowney    will get TERRIBLE performance, see the users' manual chapter on matrices.
29699ae82921SPaul Mullowney 
29709ae82921SPaul Mullowney    By default, this format uses inodes (identical nodes) when possible, to
29719ae82921SPaul Mullowney    improve numerical efficiency of matrix-vector products and solves. We
29729ae82921SPaul Mullowney    search for consecutive rows with the same nonzero structure, thereby
29739ae82921SPaul Mullowney    reusing matrix information to achieve increased efficiency.
29749ae82921SPaul Mullowney 
29759ae82921SPaul Mullowney    Level: intermediate
29769ae82921SPaul Mullowney 
2977e057df02SPaul Mullowney .seealso: MatCreate(), MatCreateAIJ(), MatSetValues(), MatSeqAIJSetColumnIndices(), MatCreateSeqAIJWithArrays(), MatCreateAIJ(), MATSEQAIJCUSPARSE, MATAIJCUSPARSE
29789ae82921SPaul Mullowney @*/
29799ae82921SPaul Mullowney PetscErrorCode  MatCreateSeqAIJCUSPARSE(MPI_Comm comm,PetscInt m,PetscInt n,PetscInt nz,const PetscInt nnz[],Mat *A)
29809ae82921SPaul Mullowney {
29819ae82921SPaul Mullowney   PetscErrorCode ierr;
29829ae82921SPaul Mullowney 
29839ae82921SPaul Mullowney   PetscFunctionBegin;
29849ae82921SPaul Mullowney   ierr = MatCreate(comm,A);CHKERRQ(ierr);
29859ae82921SPaul Mullowney   ierr = MatSetSizes(*A,m,n,m,n);CHKERRQ(ierr);
29869ae82921SPaul Mullowney   ierr = MatSetType(*A,MATSEQAIJCUSPARSE);CHKERRQ(ierr);
29879ae82921SPaul Mullowney   ierr = MatSeqAIJSetPreallocation_SeqAIJ(*A,nz,(PetscInt*)nnz);CHKERRQ(ierr);
29889ae82921SPaul Mullowney   PetscFunctionReturn(0);
29899ae82921SPaul Mullowney }
29909ae82921SPaul Mullowney 
29916fa9248bSJed Brown static PetscErrorCode MatDestroy_SeqAIJCUSPARSE(Mat A)
29929ae82921SPaul Mullowney {
29939ae82921SPaul Mullowney   PetscErrorCode              ierr;
29943fa6b06aSMark Adams   PetscSplitCSRDataStructure  *d_mat = NULL;
2995ab25e6cbSDominic Meiser 
29969ae82921SPaul Mullowney   PetscFunctionBegin;
29979ae82921SPaul Mullowney   if (A->factortype == MAT_FACTOR_NONE) {
29983fa6b06aSMark Adams     d_mat = ((Mat_SeqAIJCUSPARSE*)A->spptr)->deviceMat;
29993fa6b06aSMark Adams     ((Mat_SeqAIJCUSPARSE*)A->spptr)->deviceMat = NULL;
3000470880abSPatrick Sanan     ierr = MatSeqAIJCUSPARSE_Destroy((Mat_SeqAIJCUSPARSE**)&A->spptr);CHKERRQ(ierr);
30019ae82921SPaul Mullowney   } else {
3002470880abSPatrick Sanan     ierr = MatSeqAIJCUSPARSETriFactors_Destroy((Mat_SeqAIJCUSPARSETriFactors**)&A->spptr);CHKERRQ(ierr);
3003aa372e3fSPaul Mullowney   }
30043fa6b06aSMark Adams   if (d_mat) {
30053fa6b06aSMark Adams     Mat_SeqAIJ                 *a = (Mat_SeqAIJ*)A->data;
30063fa6b06aSMark Adams     cudaError_t                err;
30073fa6b06aSMark Adams     PetscSplitCSRDataStructure h_mat;
30083fa6b06aSMark Adams     ierr = PetscInfo(A,"Have device matrix\n");CHKERRQ(ierr);
30093fa6b06aSMark Adams     err = cudaMemcpy( &h_mat, d_mat, sizeof(PetscSplitCSRDataStructure), cudaMemcpyDeviceToHost);CHKERRCUDA(err);
30103fa6b06aSMark Adams     if (a->compressedrow.use) {
30113fa6b06aSMark Adams       err = cudaFree(h_mat.diag.i);CHKERRCUDA(err);
30123fa6b06aSMark Adams     }
30133fa6b06aSMark Adams     err = cudaFree(d_mat);CHKERRCUDA(err);
30143fa6b06aSMark Adams   }
3015ccdfe979SStefano Zampini   ierr = PetscObjectComposeFunction((PetscObject)A,"MatCUSPARSESetFormat_C",NULL);CHKERRQ(ierr);
3016ccdfe979SStefano Zampini   ierr = PetscObjectComposeFunction((PetscObject)A,"MatProductSetFromOptions_seqaijcusparse_seqdensecuda_C",NULL);CHKERRQ(ierr);
3017ccdfe979SStefano Zampini   ierr = PetscObjectComposeFunction((PetscObject)A,"MatProductSetFromOptions_seqaijcusparse_seqdense_C",NULL);CHKERRQ(ierr);
3018fcdce8c4SStefano Zampini   ierr = PetscObjectComposeFunction((PetscObject)A,"MatProductSetFromOptions_seqaijcusparse_seqaijcusparse_C",NULL);CHKERRQ(ierr);
3019ccdfe979SStefano Zampini   ierr = PetscObjectComposeFunction((PetscObject)A,"MatFactorGetSolverType_C",NULL);CHKERRQ(ierr);
30207e8381f9SStefano Zampini   ierr = PetscObjectComposeFunction((PetscObject)A,"MatSetPreallocationCOO_C",NULL);CHKERRQ(ierr);
30217e8381f9SStefano Zampini   ierr = PetscObjectComposeFunction((PetscObject)A,"MatSetValuesCOO_C",NULL);CHKERRQ(ierr);
30229ae82921SPaul Mullowney   ierr = MatDestroy_SeqAIJ(A);CHKERRQ(ierr);
30239ae82921SPaul Mullowney   PetscFunctionReturn(0);
30249ae82921SPaul Mullowney }
30259ae82921SPaul Mullowney 
3026ccdfe979SStefano Zampini PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJCUSPARSE(Mat,MatType,MatReuse,Mat*);
302795639643SRichard Tran Mills static PetscErrorCode MatBindToCPU_SeqAIJCUSPARSE(Mat,PetscBool);
30289ff858a8SKarl Rupp static PetscErrorCode MatDuplicate_SeqAIJCUSPARSE(Mat A,MatDuplicateOption cpvalues,Mat *B)
30299ff858a8SKarl Rupp {
30309ff858a8SKarl Rupp   PetscErrorCode ierr;
30319ff858a8SKarl Rupp 
30329ff858a8SKarl Rupp   PetscFunctionBegin;
30339ff858a8SKarl Rupp   ierr = MatDuplicate_SeqAIJ(A,cpvalues,B);CHKERRQ(ierr);
3034ccdfe979SStefano Zampini   ierr = MatConvert_SeqAIJ_SeqAIJCUSPARSE(*B,MATSEQAIJCUSPARSE,MAT_INPLACE_MATRIX,B);CHKERRQ(ierr);
30359ff858a8SKarl Rupp   PetscFunctionReturn(0);
30369ff858a8SKarl Rupp }
30379ff858a8SKarl Rupp 
3038a587d139SMark static PetscErrorCode MatAXPY_SeqAIJCUSPARSE(Mat Y,PetscScalar a,Mat X,MatStructure str) // put axpy in aijcusparse, etc.
303995639643SRichard Tran Mills {
3040e6e9a74fSStefano Zampini   PetscErrorCode ierr;
3041a587d139SMark   Mat_SeqAIJ     *x = (Mat_SeqAIJ*)X->data,*y = (Mat_SeqAIJ*)Y->data;
3042a587d139SMark   PetscBool      flgx,flgy;
3043e6e9a74fSStefano Zampini 
304495639643SRichard Tran Mills   PetscFunctionBegin;
3045a587d139SMark   if (a == (PetscScalar)0.0) PetscFunctionReturn(0);
3046a587d139SMark   PetscValidHeaderSpecific(Y,MAT_CLASSID,1);
3047a587d139SMark   PetscValidHeaderSpecific(X,MAT_CLASSID,3);
3048a587d139SMark   ierr = PetscObjectTypeCompare((PetscObject)Y,MATSEQAIJCUSPARSE,&flgy);CHKERRQ(ierr);
3049a587d139SMark   ierr = PetscObjectTypeCompare((PetscObject)X,MATSEQAIJCUSPARSE,&flgx);CHKERRQ(ierr);
3050a587d139SMark   if (!flgx || !flgy) {
3051a587d139SMark     ierr = MatAXPY_SeqAIJ( Y, a, X, str);CHKERRQ(ierr);
3052a587d139SMark     PetscFunctionReturn(0);
305395639643SRichard Tran Mills   }
3054a587d139SMark   if (Y->factortype != MAT_FACTOR_NONE || X->factortype != MAT_FACTOR_NONE) SETERRQ(PetscObjectComm((PetscObject)X),PETSC_ERR_PLIB,"both matrices must be MAT_FACTOR_NONE");
3055a587d139SMark   if (str == DIFFERENT_NONZERO_PATTERN) {
3056a587d139SMark     if (x->nz == y->nz) {
3057a587d139SMark       PetscBool e;
3058a587d139SMark       ierr = PetscArraycmp(x->i,y->i,Y->rmap->n+1,&e);CHKERRQ(ierr);
3059a587d139SMark       if (e) {
3060a587d139SMark         ierr = PetscArraycmp(x->j,y->j,y->nz,&e);CHKERRQ(ierr);
3061a587d139SMark         if (e) {
3062a587d139SMark           str = SAME_NONZERO_PATTERN;
3063a587d139SMark         }
3064a587d139SMark       }
3065a587d139SMark     }
3066a587d139SMark   }
3067a587d139SMark   if (str != SAME_NONZERO_PATTERN) {
3068a587d139SMark     ierr = MatAXPY_SeqAIJ( Y, a, X, str);CHKERRQ(ierr);
3069a587d139SMark     PetscFunctionReturn(0);
3070a587d139SMark   } else {
3071a587d139SMark     Mat_SeqAIJCUSPARSE           *cusparsestruct_y = (Mat_SeqAIJCUSPARSE*)Y->spptr;
3072a587d139SMark     Mat_SeqAIJCUSPARSE           *cusparsestruct_x = (Mat_SeqAIJCUSPARSE*)X->spptr;
3073a587d139SMark     if (cusparsestruct_y->format!=MAT_CUSPARSE_CSR) SETERRQ(PetscObjectComm((PetscObject)X),PETSC_ERR_PLIB,"only MAT_CUSPARSE_CSR supported");
3074a587d139SMark     if (cusparsestruct_x->format!=MAT_CUSPARSE_CSR) SETERRQ(PetscObjectComm((PetscObject)X),PETSC_ERR_PLIB,"only MAT_CUSPARSE_CSR supported");
3075a587d139SMark     if (!cusparsestruct_y->mat || !cusparsestruct_x->mat) {
3076a587d139SMark       if (Y->offloadmask == PETSC_OFFLOAD_UNALLOCATED || Y->offloadmask == PETSC_OFFLOAD_GPU) {
3077a587d139SMark         ierr = MatSeqAIJCUSPARSECopyFromGPU(Y);CHKERRQ(ierr);
3078a587d139SMark       }
3079a587d139SMark       if (X->offloadmask == PETSC_OFFLOAD_UNALLOCATED || X->offloadmask == PETSC_OFFLOAD_GPU) {
3080a587d139SMark         ierr = MatSeqAIJCUSPARSECopyFromGPU(X);CHKERRQ(ierr);
3081a587d139SMark       }
3082a587d139SMark       ierr = MatAXPY_SeqAIJ(Y,a,X,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
3083a587d139SMark       ierr = MatSeqAIJCUSPARSECopyToGPU(Y);CHKERRQ(ierr);
3084a587d139SMark       ierr = MatSeqAIJCUSPARSECopyToGPU(X);CHKERRQ(ierr);
3085a587d139SMark     } else {
3086a587d139SMark       cublasHandle_t cublasv2handle;
3087a587d139SMark       cublasStatus_t cberr;
3088a587d139SMark       cudaError_t    err;
3089a587d139SMark       PetscScalar    alpha = a;
3090a587d139SMark       PetscBLASInt   one = 1, bnz = 1;
3091a587d139SMark       CsrMatrix      *matrix_y = (CsrMatrix*)cusparsestruct_y->mat->mat;
3092a587d139SMark       CsrMatrix      *matrix_x = (CsrMatrix*)cusparsestruct_x->mat->mat;
3093a587d139SMark       PetscScalar    *aa_y, *aa_x;
3094a587d139SMark       aa_y = matrix_y->values->data().get();
3095a587d139SMark       aa_x = matrix_x->values->data().get();
3096a587d139SMark       ierr = PetscCUBLASGetHandle(&cublasv2handle);CHKERRQ(ierr);
3097a587d139SMark       ierr = PetscBLASIntCast(x->nz,&bnz);CHKERRQ(ierr);
3098a587d139SMark       ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
3099a587d139SMark       cberr = cublasXaxpy(cublasv2handle,bnz,&alpha,aa_x,one,aa_y,one);CHKERRCUBLAS(cberr);
3100a587d139SMark       err  = WaitForCUDA();CHKERRCUDA(err);
3101a587d139SMark       ierr = PetscLogGpuFlops(2.0*bnz);CHKERRQ(ierr);
3102a587d139SMark       ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
3103a587d139SMark       ierr = MatSeqAIJInvalidateDiagonal(Y);CHKERRQ(ierr);
3104a587d139SMark       ierr = PetscObjectStateIncrease((PetscObject)Y);CHKERRQ(ierr);
3105a587d139SMark       if (Y->offloadmask == PETSC_OFFLOAD_BOTH) Y->offloadmask = PETSC_OFFLOAD_GPU;
3106a587d139SMark       else if (Y->offloadmask != PETSC_OFFLOAD_GPU) SETERRQ(PetscObjectComm((PetscObject)X),PETSC_ERR_PLIB,"wrong state");
3107a587d139SMark       ierr = MatSeqAIJCUSPARSECopyFromGPU(Y);CHKERRQ(ierr);
3108a587d139SMark     }
3109a587d139SMark   }
311095639643SRichard Tran Mills   PetscFunctionReturn(0);
311195639643SRichard Tran Mills }
311295639643SRichard Tran Mills 
31133fa6b06aSMark Adams static PetscErrorCode MatZeroEntries_SeqAIJCUSPARSE(Mat A)
31143fa6b06aSMark Adams {
31153fa6b06aSMark Adams   PetscErrorCode             ierr;
31167e8381f9SStefano Zampini   PetscBool                  both = PETSC_FALSE;
3117a587d139SMark   Mat_SeqAIJ                 *a = (Mat_SeqAIJ*)A->data;
31187e8381f9SStefano Zampini 
31193fa6b06aSMark Adams   PetscFunctionBegin;
31203fa6b06aSMark Adams   if (A->factortype == MAT_FACTOR_NONE) {
31213fa6b06aSMark Adams     Mat_SeqAIJCUSPARSE *spptr = (Mat_SeqAIJCUSPARSE*)A->spptr;
31227e8381f9SStefano Zampini     if (spptr->mat) {
31237e8381f9SStefano Zampini       CsrMatrix* matrix = (CsrMatrix*)spptr->mat->mat;
31247e8381f9SStefano Zampini       if (matrix->values) {
31257e8381f9SStefano Zampini         both = PETSC_TRUE;
31267e8381f9SStefano Zampini         thrust::fill(thrust::device,matrix->values->begin(),matrix->values->end(),0.);
31277e8381f9SStefano Zampini       }
31287e8381f9SStefano Zampini     }
31297e8381f9SStefano Zampini     if (spptr->matTranspose) {
31307e8381f9SStefano Zampini       CsrMatrix* matrix = (CsrMatrix*)spptr->matTranspose->mat;
31317e8381f9SStefano Zampini       if (matrix->values) {
31327e8381f9SStefano Zampini         thrust::fill(thrust::device,matrix->values->begin(),matrix->values->end(),0.);
31337e8381f9SStefano Zampini       }
31347e8381f9SStefano Zampini     }
31353fa6b06aSMark Adams   }
3136a587d139SMark   //ierr = MatZeroEntries_SeqAIJ(A);CHKERRQ(ierr);
3137a587d139SMark   ierr = PetscArrayzero(a->a,a->i[A->rmap->n]);CHKERRQ(ierr);
3138a587d139SMark   ierr = MatSeqAIJInvalidateDiagonal(A);CHKERRQ(ierr);
31397e8381f9SStefano Zampini   if (both) A->offloadmask = PETSC_OFFLOAD_BOTH;
3140a587d139SMark   else A->offloadmask = PETSC_OFFLOAD_CPU;
31413fa6b06aSMark Adams 
31423fa6b06aSMark Adams   PetscFunctionReturn(0);
31433fa6b06aSMark Adams }
31443fa6b06aSMark Adams 
3145a587d139SMark static PetscErrorCode MatBindToCPU_SeqAIJCUSPARSE(Mat A,PetscBool flg)
3146a587d139SMark {
3147a587d139SMark   Mat_SeqAIJ     *a = (Mat_SeqAIJ*)A->data;
3148a587d139SMark   PetscErrorCode ierr;
3149a587d139SMark 
3150a587d139SMark   PetscFunctionBegin;
3151a587d139SMark   if (A->factortype != MAT_FACTOR_NONE) PetscFunctionReturn(0);
3152a587d139SMark   if (flg) {
3153a587d139SMark     ierr = MatSeqAIJCUSPARSECopyFromGPU(A);CHKERRQ(ierr);
3154a587d139SMark 
3155a587d139SMark     A->ops->axpy                      = MatAXPY_SeqAIJ;
3156a587d139SMark     A->ops->zeroentries               = MatZeroEntries_SeqAIJ;
3157a587d139SMark     A->ops->mult                      = MatMult_SeqAIJ;
3158a587d139SMark     A->ops->multadd                   = MatMultAdd_SeqAIJ;
3159a587d139SMark     A->ops->multtranspose             = MatMultTranspose_SeqAIJ;
3160a587d139SMark     A->ops->multtransposeadd          = MatMultTransposeAdd_SeqAIJ;
3161a587d139SMark     A->ops->multhermitiantranspose    = NULL;
3162a587d139SMark     A->ops->multhermitiantransposeadd = NULL;
3163fcdce8c4SStefano Zampini     A->ops->productsetfromoptions     = MatProductSetFromOptions_SeqAIJ;
3164a587d139SMark     ierr = PetscObjectComposeFunction((PetscObject)A,"MatProductSetFromOptions_seqaijcusparse_seqdensecuda_C",NULL);CHKERRQ(ierr);
3165a587d139SMark     ierr = PetscObjectComposeFunction((PetscObject)A,"MatProductSetFromOptions_seqaijcusparse_seqdense_C",NULL);CHKERRQ(ierr);
3166a587d139SMark     ierr = PetscObjectComposeFunction((PetscObject)A,"MatSetPreallocationCOO_C",NULL);CHKERRQ(ierr);
3167a587d139SMark     ierr = PetscObjectComposeFunction((PetscObject)A,"MatSetValuesCOO_C",NULL);CHKERRQ(ierr);
3168a587d139SMark     ierr = PetscObjectComposeFunction((PetscObject)A,"MatSeqAIJGetArray_C",MatSeqAIJGetArray_SeqAIJ);CHKERRQ(ierr);
3169fcdce8c4SStefano Zampini     ierr = PetscObjectComposeFunction((PetscObject)A,"MatProductSetFromOptions_seqaijcusparse_seqaijcusparse_C",NULL);CHKERRQ(ierr);
3170a587d139SMark   } else {
3171a587d139SMark     A->ops->axpy                      = MatAXPY_SeqAIJCUSPARSE;
3172a587d139SMark     A->ops->zeroentries               = MatZeroEntries_SeqAIJCUSPARSE;
3173a587d139SMark     A->ops->mult                      = MatMult_SeqAIJCUSPARSE;
3174a587d139SMark     A->ops->multadd                   = MatMultAdd_SeqAIJCUSPARSE;
3175a587d139SMark     A->ops->multtranspose             = MatMultTranspose_SeqAIJCUSPARSE;
3176a587d139SMark     A->ops->multtransposeadd          = MatMultTransposeAdd_SeqAIJCUSPARSE;
3177a587d139SMark     A->ops->multhermitiantranspose    = MatMultHermitianTranspose_SeqAIJCUSPARSE;
3178a587d139SMark     A->ops->multhermitiantransposeadd = MatMultHermitianTransposeAdd_SeqAIJCUSPARSE;
3179fcdce8c4SStefano Zampini     A->ops->productsetfromoptions     = MatProductSetFromOptions_SeqAIJCUSPARSE;
3180a587d139SMark     ierr = PetscObjectComposeFunction((PetscObject)A,"MatProductSetFromOptions_seqaijcusparse_seqdensecuda_C",MatProductSetFromOptions_SeqAIJCUSPARSE);CHKERRQ(ierr);
3181a587d139SMark     ierr = PetscObjectComposeFunction((PetscObject)A,"MatProductSetFromOptions_seqaijcusparse_seqdense_C",MatProductSetFromOptions_SeqAIJCUSPARSE);CHKERRQ(ierr);
3182a587d139SMark     ierr = PetscObjectComposeFunction((PetscObject)A,"MatSetPreallocationCOO_C",MatSetPreallocationCOO_SeqAIJCUSPARSE);CHKERRQ(ierr);
3183a587d139SMark     ierr = PetscObjectComposeFunction((PetscObject)A,"MatSetValuesCOO_C",MatSetValuesCOO_SeqAIJCUSPARSE);CHKERRQ(ierr);
3184a587d139SMark     ierr = PetscObjectComposeFunction((PetscObject)A,"MatSeqAIJGetArray_C",MatSeqAIJGetArray_SeqAIJCUSPARSE);CHKERRQ(ierr);
3185fcdce8c4SStefano Zampini     ierr = PetscObjectComposeFunction((PetscObject)A,"MatProductSetFromOptions_seqaijcusparse_seqaijcusparse_C",MatProductSetFromOptions_SeqAIJCUSPARSE);CHKERRQ(ierr);
3186a587d139SMark   }
3187a587d139SMark   A->boundtocpu = flg;
3188a587d139SMark   a->inode.use = flg;
3189a587d139SMark   PetscFunctionReturn(0);
3190a587d139SMark }
3191a587d139SMark 
319249735bf3SStefano Zampini PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJCUSPARSE(Mat A, MatType mtype, MatReuse reuse, Mat* newmat)
31939ae82921SPaul Mullowney {
31949ae82921SPaul Mullowney   PetscErrorCode   ierr;
3195aa372e3fSPaul Mullowney   cusparseStatus_t stat;
319649735bf3SStefano Zampini   Mat              B;
31979ae82921SPaul Mullowney 
31989ae82921SPaul Mullowney   PetscFunctionBegin;
3199832b2c02SStefano Zampini   ierr = PetscCUDAInitializeCheck();CHKERRQ(ierr); /* first use of CUSPARSE may be via MatConvert */
320049735bf3SStefano Zampini   if (reuse == MAT_INITIAL_MATRIX) {
320149735bf3SStefano Zampini     ierr = MatDuplicate(A,MAT_COPY_VALUES,newmat);CHKERRQ(ierr);
320249735bf3SStefano Zampini   } else if (reuse == MAT_REUSE_MATRIX) {
320349735bf3SStefano Zampini     ierr = MatCopy(A,*newmat,SAME_NONZERO_PATTERN);CHKERRQ(ierr);
320449735bf3SStefano Zampini   }
320549735bf3SStefano Zampini   B = *newmat;
320649735bf3SStefano Zampini 
320734136279SStefano Zampini   ierr = PetscFree(B->defaultvectype);CHKERRQ(ierr);
320834136279SStefano Zampini   ierr = PetscStrallocpy(VECCUDA,&B->defaultvectype);CHKERRQ(ierr);
320934136279SStefano Zampini 
321049735bf3SStefano Zampini   if (reuse != MAT_REUSE_MATRIX && !B->spptr) {
32119ae82921SPaul Mullowney     if (B->factortype == MAT_FACTOR_NONE) {
3212e6e9a74fSStefano Zampini       Mat_SeqAIJCUSPARSE *spptr;
3213e6e9a74fSStefano Zampini 
3214e6e9a74fSStefano Zampini       ierr = PetscNew(&spptr);CHKERRQ(ierr);
3215e6e9a74fSStefano Zampini       spptr->format = MAT_CUSPARSE_CSR;
3216e6e9a74fSStefano Zampini       stat = cusparseCreate(&spptr->handle);CHKERRCUSPARSE(stat);
3217e6e9a74fSStefano Zampini       B->spptr = spptr;
32183fa6b06aSMark Adams       spptr->deviceMat = NULL;
32199ae82921SPaul Mullowney     } else {
3220e6e9a74fSStefano Zampini       Mat_SeqAIJCUSPARSETriFactors *spptr;
3221e6e9a74fSStefano Zampini 
3222e6e9a74fSStefano Zampini       ierr = PetscNew(&spptr);CHKERRQ(ierr);
3223e6e9a74fSStefano Zampini       stat = cusparseCreate(&spptr->handle);CHKERRCUSPARSE(stat);
3224e6e9a74fSStefano Zampini       B->spptr = spptr;
32259ae82921SPaul Mullowney     }
3226e6e9a74fSStefano Zampini     B->offloadmask = PETSC_OFFLOAD_UNALLOCATED;
322749735bf3SStefano Zampini   }
3228693b0035SStefano Zampini   B->ops->assemblyend    = MatAssemblyEnd_SeqAIJCUSPARSE;
32299ae82921SPaul Mullowney   B->ops->destroy        = MatDestroy_SeqAIJCUSPARSE;
32309ae82921SPaul Mullowney   B->ops->setfromoptions = MatSetFromOptions_SeqAIJCUSPARSE;
323195639643SRichard Tran Mills   B->ops->bindtocpu      = MatBindToCPU_SeqAIJCUSPARSE;
3232693b0035SStefano Zampini   B->ops->duplicate      = MatDuplicate_SeqAIJCUSPARSE;
32332205254eSKarl Rupp 
3234e6e9a74fSStefano Zampini   ierr = MatBindToCPU_SeqAIJCUSPARSE(B,PETSC_FALSE);CHKERRQ(ierr);
32359ae82921SPaul Mullowney   ierr = PetscObjectChangeTypeName((PetscObject)B,MATSEQAIJCUSPARSE);CHKERRQ(ierr);
3236bdf89e91SBarry Smith   ierr = PetscObjectComposeFunction((PetscObject)B,"MatCUSPARSESetFormat_C",MatCUSPARSESetFormat_SeqAIJCUSPARSE);CHKERRQ(ierr);
32379ae82921SPaul Mullowney   PetscFunctionReturn(0);
32389ae82921SPaul Mullowney }
32399ae82921SPaul Mullowney 
324002fe1965SBarry Smith PETSC_EXTERN PetscErrorCode MatCreate_SeqAIJCUSPARSE(Mat B)
324102fe1965SBarry Smith {
324202fe1965SBarry Smith   PetscErrorCode ierr;
324302fe1965SBarry Smith 
324402fe1965SBarry Smith   PetscFunctionBegin;
324502fe1965SBarry Smith   ierr = MatCreate_SeqAIJ(B);CHKERRQ(ierr);
32460ce8acdeSStefano Zampini   ierr = MatConvert_SeqAIJ_SeqAIJCUSPARSE(B,MATSEQAIJCUSPARSE,MAT_INPLACE_MATRIX,&B);CHKERRQ(ierr);
3247afb2bd1cSJunchao Zhang   ierr = PetscObjectOptionsBegin((PetscObject)B);CHKERRQ(ierr);
3248afb2bd1cSJunchao Zhang   ierr = MatSetFromOptions_SeqAIJCUSPARSE(PetscOptionsObject,B);CHKERRQ(ierr);
3249afb2bd1cSJunchao Zhang   ierr = PetscOptionsEnd();CHKERRQ(ierr);
325002fe1965SBarry Smith   PetscFunctionReturn(0);
325102fe1965SBarry Smith }
325202fe1965SBarry Smith 
32533ca39a21SBarry Smith /*MC
3254e057df02SPaul Mullowney    MATSEQAIJCUSPARSE - MATAIJCUSPARSE = "(seq)aijcusparse" - A matrix type to be used for sparse matrices.
3255e057df02SPaul Mullowney 
3256e057df02SPaul Mullowney    A matrix type type whose data resides on Nvidia GPUs. These matrices can be in either
32572692e278SPaul Mullowney    CSR, ELL, or Hybrid format. The ELL and HYB formats require CUDA 4.2 or later.
32582692e278SPaul Mullowney    All matrix calculations are performed on Nvidia GPUs using the CUSPARSE library.
3259e057df02SPaul Mullowney 
3260e057df02SPaul Mullowney    Options Database Keys:
3261e057df02SPaul Mullowney +  -mat_type aijcusparse - sets the matrix type to "seqaijcusparse" during a call to MatSetFromOptions()
3262aa372e3fSPaul Mullowney .  -mat_cusparse_storage_format csr - sets the storage format of matrices (for MatMult and factors in MatSolve) during a call to MatSetFromOptions(). Other options include ell (ellpack) or hyb (hybrid).
3263a2b725a8SWilliam Gropp -  -mat_cusparse_mult_storage_format csr - sets the storage format of matrices (for MatMult) during a call to MatSetFromOptions(). Other options include ell (ellpack) or hyb (hybrid).
3264e057df02SPaul Mullowney 
3265e057df02SPaul Mullowney   Level: beginner
3266e057df02SPaul Mullowney 
32678468deeeSKarl Rupp .seealso: MatCreateSeqAIJCUSPARSE(), MATAIJCUSPARSE, MatCreateAIJCUSPARSE(), MatCUSPARSESetFormat(), MatCUSPARSEStorageFormat, MatCUSPARSEFormatOperation
3268e057df02SPaul Mullowney M*/
32697f756511SDominic Meiser 
327042c9c57cSBarry Smith PETSC_EXTERN PetscErrorCode MatGetFactor_seqaijcusparse_cusparse(Mat,MatFactorType,Mat*);
327142c9c57cSBarry Smith 
32720f39cd5aSBarry Smith 
32733ca39a21SBarry Smith PETSC_EXTERN PetscErrorCode MatSolverTypeRegister_CUSPARSE(void)
327442c9c57cSBarry Smith {
327542c9c57cSBarry Smith   PetscErrorCode ierr;
327642c9c57cSBarry Smith 
327742c9c57cSBarry Smith   PetscFunctionBegin;
32783ca39a21SBarry Smith   ierr = MatSolverTypeRegister(MATSOLVERCUSPARSE,MATSEQAIJCUSPARSE,MAT_FACTOR_LU,MatGetFactor_seqaijcusparse_cusparse);CHKERRQ(ierr);
32793ca39a21SBarry Smith   ierr = MatSolverTypeRegister(MATSOLVERCUSPARSE,MATSEQAIJCUSPARSE,MAT_FACTOR_CHOLESKY,MatGetFactor_seqaijcusparse_cusparse);CHKERRQ(ierr);
32803ca39a21SBarry Smith   ierr = MatSolverTypeRegister(MATSOLVERCUSPARSE,MATSEQAIJCUSPARSE,MAT_FACTOR_ILU,MatGetFactor_seqaijcusparse_cusparse);CHKERRQ(ierr);
32813ca39a21SBarry Smith   ierr = MatSolverTypeRegister(MATSOLVERCUSPARSE,MATSEQAIJCUSPARSE,MAT_FACTOR_ICC,MatGetFactor_seqaijcusparse_cusparse);CHKERRQ(ierr);
328242c9c57cSBarry Smith   PetscFunctionReturn(0);
328342c9c57cSBarry Smith }
328429b38603SBarry Smith 
3285470880abSPatrick Sanan static PetscErrorCode MatSeqAIJCUSPARSE_Destroy(Mat_SeqAIJCUSPARSE **cusparsestruct)
32867f756511SDominic Meiser {
3287e6e9a74fSStefano Zampini   PetscErrorCode   ierr;
32887f756511SDominic Meiser   cusparseStatus_t stat;
32897f756511SDominic Meiser 
32907f756511SDominic Meiser   PetscFunctionBegin;
32917f756511SDominic Meiser   if (*cusparsestruct) {
3292e6e9a74fSStefano Zampini     ierr = MatSeqAIJCUSPARSEMultStruct_Destroy(&(*cusparsestruct)->mat,(*cusparsestruct)->format);CHKERRQ(ierr);
3293e6e9a74fSStefano Zampini     ierr = MatSeqAIJCUSPARSEMultStruct_Destroy(&(*cusparsestruct)->matTranspose,(*cusparsestruct)->format);CHKERRQ(ierr);
32947f756511SDominic Meiser     delete (*cusparsestruct)->workVector;
329581902715SJunchao Zhang     delete (*cusparsestruct)->rowoffsets_gpu;
32967e8381f9SStefano Zampini     delete (*cusparsestruct)->cooPerm;
32977e8381f9SStefano Zampini     delete (*cusparsestruct)->cooPerm_a;
32987e8381f9SStefano Zampini     if ((*cusparsestruct)->handle) {stat = cusparseDestroy((*cusparsestruct)->handle);CHKERRCUSPARSE(stat);}
3299afb2bd1cSJunchao Zhang    #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
3300afb2bd1cSJunchao Zhang     cudaError_t cerr = cudaFree((*cusparsestruct)->csr2cscBuffer);CHKERRCUDA(cerr);
3301afb2bd1cSJunchao Zhang    #endif
3302e6e9a74fSStefano Zampini     ierr = PetscFree(*cusparsestruct);CHKERRQ(ierr);
33037f756511SDominic Meiser   }
33047f756511SDominic Meiser   PetscFunctionReturn(0);
33057f756511SDominic Meiser }
33067f756511SDominic Meiser 
33077f756511SDominic Meiser static PetscErrorCode CsrMatrix_Destroy(CsrMatrix **mat)
33087f756511SDominic Meiser {
33097f756511SDominic Meiser   PetscFunctionBegin;
33107f756511SDominic Meiser   if (*mat) {
33117f756511SDominic Meiser     delete (*mat)->values;
33127f756511SDominic Meiser     delete (*mat)->column_indices;
33137f756511SDominic Meiser     delete (*mat)->row_offsets;
33147f756511SDominic Meiser     delete *mat;
33157f756511SDominic Meiser     *mat = 0;
33167f756511SDominic Meiser   }
33177f756511SDominic Meiser   PetscFunctionReturn(0);
33187f756511SDominic Meiser }
33197f756511SDominic Meiser 
3320470880abSPatrick Sanan static PetscErrorCode MatSeqAIJCUSPARSEMultStruct_Destroy(Mat_SeqAIJCUSPARSETriFactorStruct **trifactor)
33217f756511SDominic Meiser {
33227f756511SDominic Meiser   cusparseStatus_t stat;
33237f756511SDominic Meiser   PetscErrorCode   ierr;
33247f756511SDominic Meiser 
33257f756511SDominic Meiser   PetscFunctionBegin;
33267f756511SDominic Meiser   if (*trifactor) {
332757d48284SJunchao Zhang     if ((*trifactor)->descr) { stat = cusparseDestroyMatDescr((*trifactor)->descr);CHKERRCUSPARSE(stat); }
3328afb2bd1cSJunchao Zhang     if ((*trifactor)->solveInfo) { stat = cusparse_destroy_analysis_info((*trifactor)->solveInfo);CHKERRCUSPARSE(stat); }
33297f756511SDominic Meiser     ierr = CsrMatrix_Destroy(&(*trifactor)->csrMat);CHKERRQ(ierr);
33301b0a6780SStefano Zampini     if ((*trifactor)->solveBuffer)   {cudaError_t cerr = cudaFree((*trifactor)->solveBuffer);CHKERRCUDA(cerr);}
33312cbc15d9SMark     if ((*trifactor)->AA_h)   {cudaError_t cerr = cudaFreeHost((*trifactor)->AA_h);CHKERRCUDA(cerr);}
3332afb2bd1cSJunchao Zhang    #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
33331b0a6780SStefano Zampini     if ((*trifactor)->csr2cscBuffer) {cudaError_t cerr = cudaFree((*trifactor)->csr2cscBuffer);CHKERRCUDA(cerr);}
3334afb2bd1cSJunchao Zhang    #endif
3335da79fbbcSStefano Zampini     ierr = PetscFree(*trifactor);CHKERRQ(ierr);
33367f756511SDominic Meiser   }
33377f756511SDominic Meiser   PetscFunctionReturn(0);
33387f756511SDominic Meiser }
33397f756511SDominic Meiser 
3340470880abSPatrick Sanan static PetscErrorCode MatSeqAIJCUSPARSEMultStruct_Destroy(Mat_SeqAIJCUSPARSEMultStruct **matstruct,MatCUSPARSEStorageFormat format)
33417f756511SDominic Meiser {
33427f756511SDominic Meiser   CsrMatrix        *mat;
33437f756511SDominic Meiser   cusparseStatus_t stat;
33447f756511SDominic Meiser   cudaError_t      err;
33457f756511SDominic Meiser 
33467f756511SDominic Meiser   PetscFunctionBegin;
33477f756511SDominic Meiser   if (*matstruct) {
33487f756511SDominic Meiser     if ((*matstruct)->mat) {
33497f756511SDominic Meiser       if (format==MAT_CUSPARSE_ELL || format==MAT_CUSPARSE_HYB) {
3350afb2bd1cSJunchao Zhang        #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
3351afb2bd1cSJunchao Zhang         SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"MAT_CUSPARSE_ELL and MAT_CUSPARSE_HYB are not supported since CUDA-11.0");
3352afb2bd1cSJunchao Zhang        #else
33537f756511SDominic Meiser         cusparseHybMat_t hybMat = (cusparseHybMat_t)(*matstruct)->mat;
335457d48284SJunchao Zhang         stat = cusparseDestroyHybMat(hybMat);CHKERRCUSPARSE(stat);
3355afb2bd1cSJunchao Zhang        #endif
33567f756511SDominic Meiser       } else {
33577f756511SDominic Meiser         mat = (CsrMatrix*)(*matstruct)->mat;
33587f756511SDominic Meiser         CsrMatrix_Destroy(&mat);
33597f756511SDominic Meiser       }
33607f756511SDominic Meiser     }
336157d48284SJunchao Zhang     if ((*matstruct)->descr) { stat = cusparseDestroyMatDescr((*matstruct)->descr);CHKERRCUSPARSE(stat); }
33627f756511SDominic Meiser     delete (*matstruct)->cprowIndices;
3363afb2bd1cSJunchao Zhang     if ((*matstruct)->alpha_one) { err=cudaFree((*matstruct)->alpha_one);CHKERRCUDA(err); }
33647656d835SStefano Zampini     if ((*matstruct)->beta_zero) { err=cudaFree((*matstruct)->beta_zero);CHKERRCUDA(err); }
33657656d835SStefano Zampini     if ((*matstruct)->beta_one)  { err=cudaFree((*matstruct)->beta_one);CHKERRCUDA(err); }
3366afb2bd1cSJunchao Zhang 
3367afb2bd1cSJunchao Zhang    #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
3368afb2bd1cSJunchao Zhang     Mat_SeqAIJCUSPARSEMultStruct *mdata = *matstruct;
3369afb2bd1cSJunchao Zhang     if (mdata->matDescr) {stat = cusparseDestroySpMat(mdata->matDescr);CHKERRCUSPARSE(stat);}
3370afb2bd1cSJunchao Zhang     for (int i=0; i<3; i++) {
3371afb2bd1cSJunchao Zhang       if (mdata->cuSpMV[i].initialized) {
3372afb2bd1cSJunchao Zhang         err  = cudaFree(mdata->cuSpMV[i].spmvBuffer);CHKERRCUDA(err);
3373afb2bd1cSJunchao Zhang         stat = cusparseDestroyDnVec(mdata->cuSpMV[i].vecXDescr);CHKERRCUSPARSE(stat);
3374afb2bd1cSJunchao Zhang         stat = cusparseDestroyDnVec(mdata->cuSpMV[i].vecYDescr);CHKERRCUSPARSE(stat);
3375afb2bd1cSJunchao Zhang       }
3376afb2bd1cSJunchao Zhang     }
3377afb2bd1cSJunchao Zhang    #endif
33787f756511SDominic Meiser     delete *matstruct;
33797e8381f9SStefano Zampini     *matstruct = NULL;
33807f756511SDominic Meiser   }
33817f756511SDominic Meiser   PetscFunctionReturn(0);
33827f756511SDominic Meiser }
33837f756511SDominic Meiser 
3384ccdfe979SStefano Zampini static PetscErrorCode MatSeqAIJCUSPARSETriFactors_Reset(Mat_SeqAIJCUSPARSETriFactors** trifactors)
33857f756511SDominic Meiser {
3386e6e9a74fSStefano Zampini   PetscErrorCode ierr;
3387e6e9a74fSStefano Zampini 
33887f756511SDominic Meiser   PetscFunctionBegin;
33897f756511SDominic Meiser   if (*trifactors) {
3390e6e9a74fSStefano Zampini     ierr = MatSeqAIJCUSPARSEMultStruct_Destroy(&(*trifactors)->loTriFactorPtr);CHKERRQ(ierr);
3391e6e9a74fSStefano Zampini     ierr = MatSeqAIJCUSPARSEMultStruct_Destroy(&(*trifactors)->upTriFactorPtr);CHKERRQ(ierr);
3392e6e9a74fSStefano Zampini     ierr = MatSeqAIJCUSPARSEMultStruct_Destroy(&(*trifactors)->loTriFactorPtrTranspose);CHKERRQ(ierr);
3393e6e9a74fSStefano Zampini     ierr = MatSeqAIJCUSPARSEMultStruct_Destroy(&(*trifactors)->upTriFactorPtrTranspose);CHKERRQ(ierr);
33947f756511SDominic Meiser     delete (*trifactors)->rpermIndices;
33957f756511SDominic Meiser     delete (*trifactors)->cpermIndices;
33967f756511SDominic Meiser     delete (*trifactors)->workVector;
33977e8381f9SStefano Zampini     (*trifactors)->rpermIndices = NULL;
33987e8381f9SStefano Zampini     (*trifactors)->cpermIndices = NULL;
33997e8381f9SStefano Zampini     (*trifactors)->workVector = NULL;
3400ccdfe979SStefano Zampini   }
3401ccdfe979SStefano Zampini   PetscFunctionReturn(0);
3402ccdfe979SStefano Zampini }
3403ccdfe979SStefano Zampini 
3404ccdfe979SStefano Zampini static PetscErrorCode MatSeqAIJCUSPARSETriFactors_Destroy(Mat_SeqAIJCUSPARSETriFactors** trifactors)
3405ccdfe979SStefano Zampini {
3406e6e9a74fSStefano Zampini   PetscErrorCode   ierr;
3407ccdfe979SStefano Zampini   cusparseHandle_t handle;
3408ccdfe979SStefano Zampini   cusparseStatus_t stat;
3409ccdfe979SStefano Zampini 
3410ccdfe979SStefano Zampini   PetscFunctionBegin;
3411ccdfe979SStefano Zampini   if (*trifactors) {
3412e6e9a74fSStefano Zampini     ierr = MatSeqAIJCUSPARSETriFactors_Reset(trifactors);CHKERRQ(ierr);
34137f756511SDominic Meiser     if (handle = (*trifactors)->handle) {
341457d48284SJunchao Zhang       stat = cusparseDestroy(handle);CHKERRCUSPARSE(stat);
34157f756511SDominic Meiser     }
3416e6e9a74fSStefano Zampini     ierr = PetscFree(*trifactors);CHKERRQ(ierr);
34177f756511SDominic Meiser   }
34187f756511SDominic Meiser   PetscFunctionReturn(0);
34197f756511SDominic Meiser }
34207e8381f9SStefano Zampini 
34217e8381f9SStefano Zampini struct IJCompare
34227e8381f9SStefano Zampini {
34237e8381f9SStefano Zampini   __host__ __device__
34247e8381f9SStefano Zampini   inline bool operator() (const thrust::tuple<PetscInt, PetscInt> &t1, const thrust::tuple<PetscInt, PetscInt> &t2)
34257e8381f9SStefano Zampini   {
34267e8381f9SStefano Zampini     if (t1.get<0>() < t2.get<0>()) return true;
34277e8381f9SStefano Zampini     if (t1.get<0>() == t2.get<0>()) return t1.get<1>() < t2.get<1>();
34287e8381f9SStefano Zampini     return false;
34297e8381f9SStefano Zampini   }
34307e8381f9SStefano Zampini };
34317e8381f9SStefano Zampini 
34327e8381f9SStefano Zampini struct IJEqual
34337e8381f9SStefano Zampini {
34347e8381f9SStefano Zampini   __host__ __device__
34357e8381f9SStefano Zampini   inline bool operator() (const thrust::tuple<PetscInt, PetscInt> &t1, const thrust::tuple<PetscInt, PetscInt> &t2)
34367e8381f9SStefano Zampini   {
34377e8381f9SStefano Zampini     if (t1.get<0>() != t2.get<0>() || t1.get<1>() != t2.get<1>()) return false;
34387e8381f9SStefano Zampini     return true;
34397e8381f9SStefano Zampini   }
34407e8381f9SStefano Zampini };
34417e8381f9SStefano Zampini 
34427e8381f9SStefano Zampini struct IJDiff
34437e8381f9SStefano Zampini {
34447e8381f9SStefano Zampini   __host__ __device__
34457e8381f9SStefano Zampini   inline PetscInt operator() (const PetscInt &t1, const PetscInt &t2)
34467e8381f9SStefano Zampini   {
34477e8381f9SStefano Zampini     return t1 == t2 ? 0 : 1;
34487e8381f9SStefano Zampini   }
34497e8381f9SStefano Zampini };
34507e8381f9SStefano Zampini 
34517e8381f9SStefano Zampini struct IJSum
34527e8381f9SStefano Zampini {
34537e8381f9SStefano Zampini   __host__ __device__
34547e8381f9SStefano Zampini   inline PetscInt operator() (const PetscInt &t1, const PetscInt &t2)
34557e8381f9SStefano Zampini   {
34567e8381f9SStefano Zampini     return t1||t2;
34577e8381f9SStefano Zampini   }
34587e8381f9SStefano Zampini };
34597e8381f9SStefano Zampini 
34607e8381f9SStefano Zampini #include <thrust/iterator/discard_iterator.h>
3461e61fc153SStefano Zampini PetscErrorCode MatSetValuesCOO_SeqAIJCUSPARSE(Mat A, const PetscScalar v[], InsertMode imode)
34627e8381f9SStefano Zampini {
34637e8381f9SStefano Zampini   Mat_SeqAIJCUSPARSE                    *cusp = (Mat_SeqAIJCUSPARSE*)A->spptr;
3464fcdce8c4SStefano Zampini   Mat_SeqAIJ                            *a = (Mat_SeqAIJ*)A->data;
3465*bfcc3627SStefano Zampini   THRUSTARRAY                           *cooPerm_v = NULL;
346608391a17SStefano Zampini   thrust::device_ptr<const PetscScalar> d_v;
34677e8381f9SStefano Zampini   CsrMatrix                             *matrix;
34687e8381f9SStefano Zampini   PetscErrorCode                        ierr;
34697e8381f9SStefano Zampini   cudaError_t                           cerr;
34707e8381f9SStefano Zampini   PetscInt                              n;
34717e8381f9SStefano Zampini 
34727e8381f9SStefano Zampini   PetscFunctionBegin;
34737e8381f9SStefano Zampini   if (!cusp) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing CUSPARSE struct");
34747e8381f9SStefano Zampini   if (!cusp->mat) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing CUSPARSE CsrMatrix");
34757e8381f9SStefano Zampini   if (!cusp->cooPerm) {
34767e8381f9SStefano Zampini     ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
34777e8381f9SStefano Zampini     ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
34787e8381f9SStefano Zampini     PetscFunctionReturn(0);
34797e8381f9SStefano Zampini   }
34807e8381f9SStefano Zampini   matrix = (CsrMatrix*)cusp->mat->mat;
34817e8381f9SStefano Zampini   if (!matrix->values) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing CUDA memory");
3482e61fc153SStefano Zampini   if (!v) {
3483e61fc153SStefano Zampini     if (imode == INSERT_VALUES) thrust::fill(thrust::device,matrix->values->begin(),matrix->values->end(),0.);
3484e61fc153SStefano Zampini     goto finalize;
34857e8381f9SStefano Zampini   }
3486e61fc153SStefano Zampini   n = cusp->cooPerm->size();
348708391a17SStefano Zampini   if (isCudaMem(v)) {
348808391a17SStefano Zampini     d_v = thrust::device_pointer_cast(v);
348908391a17SStefano Zampini   } else {
3490e61fc153SStefano Zampini     cooPerm_v = new THRUSTARRAY(n);
3491e61fc153SStefano Zampini     cooPerm_v->assign(v,v+n);
349208391a17SStefano Zampini     d_v = cooPerm_v->data();
3493e61fc153SStefano Zampini     ierr = PetscLogCpuToGpu(n*sizeof(PetscScalar));CHKERRQ(ierr);
349408391a17SStefano Zampini   }
3495*bfcc3627SStefano Zampini   ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
3496e61fc153SStefano Zampini   if (imode == ADD_VALUES) { /* ADD VALUES means add to existing ones */
34977e8381f9SStefano Zampini     if (cusp->cooPerm_a) {
3498*bfcc3627SStefano Zampini       THRUSTARRAY *cooPerm_w = new THRUSTARRAY(matrix->values->size());
349908391a17SStefano Zampini       auto vbit = thrust::make_permutation_iterator(d_v,cusp->cooPerm->begin());
3500e61fc153SStefano Zampini       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>());
3501e61fc153SStefano Zampini       thrust::transform(cooPerm_w->begin(),cooPerm_w->end(),matrix->values->begin(),matrix->values->begin(),thrust::plus<PetscScalar>());
3502e61fc153SStefano Zampini       delete cooPerm_w;
35037e8381f9SStefano Zampini     } else {
350408391a17SStefano Zampini       auto zibit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(d_v,cusp->cooPerm->begin()),
35057e8381f9SStefano Zampini                                                                 matrix->values->begin()));
350608391a17SStefano Zampini       auto zieit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(d_v,cusp->cooPerm->end()),
35077e8381f9SStefano Zampini                                                                 matrix->values->end()));
35087e8381f9SStefano Zampini       thrust::for_each(zibit,zieit,VecCUDAPlusEquals());
35097e8381f9SStefano Zampini     }
35107e8381f9SStefano Zampini   } else {
3511e61fc153SStefano Zampini     if (cusp->cooPerm_a) { /* repeated entries in COO, with INSERT_VALUES -> reduce */
351208391a17SStefano Zampini       auto vbit = thrust::make_permutation_iterator(d_v,cusp->cooPerm->begin());
3513e61fc153SStefano Zampini       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>());
35147e8381f9SStefano Zampini     } else {
351508391a17SStefano Zampini       auto zibit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(d_v,cusp->cooPerm->begin()),
35167e8381f9SStefano Zampini                                                                 matrix->values->begin()));
351708391a17SStefano Zampini       auto zieit = thrust::make_zip_iterator(thrust::make_tuple(thrust::make_permutation_iterator(d_v,cusp->cooPerm->end()),
35187e8381f9SStefano Zampini                                                                 matrix->values->end()));
35197e8381f9SStefano Zampini       thrust::for_each(zibit,zieit,VecCUDAEquals());
35207e8381f9SStefano Zampini     }
35217e8381f9SStefano Zampini   }
35227e8381f9SStefano Zampini   cerr = WaitForCUDA();CHKERRCUDA(cerr);
3523*bfcc3627SStefano Zampini   ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
3524e61fc153SStefano Zampini finalize:
3525e61fc153SStefano Zampini   delete cooPerm_v;
35267e8381f9SStefano Zampini   A->offloadmask = PETSC_OFFLOAD_GPU;
3527e61fc153SStefano Zampini   ierr = PetscObjectStateIncrease((PetscObject)A);CHKERRQ(ierr);
3528fcdce8c4SStefano Zampini   /* shorter version of MatAssemblyEnd_SeqAIJ */
3529fcdce8c4SStefano Zampini   ierr = PetscInfo3(A,"Matrix size: %D X %D; storage space: 0 unneeded,%D used\n",A->rmap->n,A->cmap->n,a->nz);CHKERRQ(ierr);
3530fcdce8c4SStefano Zampini   ierr = PetscInfo(A,"Number of mallocs during MatSetValues() is 0\n");CHKERRQ(ierr);
3531fcdce8c4SStefano Zampini   ierr = PetscInfo1(A,"Maximum nonzeros in any row is %D\n",a->rmax);CHKERRQ(ierr);
3532fcdce8c4SStefano Zampini   a->reallocs         = 0;
3533fcdce8c4SStefano Zampini   A->info.mallocs    += 0;
3534fcdce8c4SStefano Zampini   A->info.nz_unneeded = 0;
3535fcdce8c4SStefano Zampini   A->assembled = A->was_assembled = PETSC_TRUE;
3536fcdce8c4SStefano Zampini   A->num_ass++;
35377e8381f9SStefano Zampini   PetscFunctionReturn(0);
35387e8381f9SStefano Zampini }
35397e8381f9SStefano Zampini 
35407e8381f9SStefano Zampini #include <thrust/binary_search.h>
3541e61fc153SStefano Zampini PetscErrorCode MatSetPreallocationCOO_SeqAIJCUSPARSE(Mat A, PetscInt n, const PetscInt coo_i[], const PetscInt coo_j[])
35427e8381f9SStefano Zampini {
35437e8381f9SStefano Zampini   PetscErrorCode     ierr;
35447e8381f9SStefano Zampini   Mat_SeqAIJCUSPARSE *cusp = (Mat_SeqAIJCUSPARSE*)A->spptr;
35457e8381f9SStefano Zampini   Mat_SeqAIJ         *a = (Mat_SeqAIJ*)A->data;
35467e8381f9SStefano Zampini   PetscInt           cooPerm_n, nzr = 0;
35477e8381f9SStefano Zampini   cudaError_t        cerr;
35487e8381f9SStefano Zampini 
35497e8381f9SStefano Zampini   PetscFunctionBegin;
35507e8381f9SStefano Zampini   ierr = PetscLayoutSetUp(A->rmap);CHKERRQ(ierr);
35517e8381f9SStefano Zampini   ierr = PetscLayoutSetUp(A->cmap);CHKERRQ(ierr);
35527e8381f9SStefano Zampini   cooPerm_n = cusp->cooPerm ? cusp->cooPerm->size() : 0;
35537e8381f9SStefano Zampini   if (n != cooPerm_n) {
35547e8381f9SStefano Zampini     delete cusp->cooPerm;
35557e8381f9SStefano Zampini     delete cusp->cooPerm_a;
35567e8381f9SStefano Zampini     cusp->cooPerm = NULL;
35577e8381f9SStefano Zampini     cusp->cooPerm_a = NULL;
35587e8381f9SStefano Zampini   }
35597e8381f9SStefano Zampini   if (n) {
35607e8381f9SStefano Zampini     THRUSTINTARRAY d_i(n);
35617e8381f9SStefano Zampini     THRUSTINTARRAY d_j(n);
35627e8381f9SStefano Zampini     THRUSTINTARRAY ii(A->rmap->n);
35637e8381f9SStefano Zampini 
35647e8381f9SStefano Zampini     if (!cusp->cooPerm)   { cusp->cooPerm   = new THRUSTINTARRAY(n); }
35657e8381f9SStefano Zampini     if (!cusp->cooPerm_a) { cusp->cooPerm_a = new THRUSTINTARRAY(n); }
35667e8381f9SStefano Zampini 
35677e8381f9SStefano Zampini     ierr = PetscLogCpuToGpu(2.*n*sizeof(PetscInt));CHKERRQ(ierr);
35687e8381f9SStefano Zampini     d_i.assign(coo_i,coo_i+n);
35697e8381f9SStefano Zampini     d_j.assign(coo_j,coo_j+n);
35707e8381f9SStefano Zampini     auto fkey = thrust::make_zip_iterator(thrust::make_tuple(d_i.begin(),d_j.begin()));
35717e8381f9SStefano Zampini     auto ekey = thrust::make_zip_iterator(thrust::make_tuple(d_i.end(),d_j.end()));
35727e8381f9SStefano Zampini 
357308391a17SStefano Zampini     ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
35747e8381f9SStefano Zampini     thrust::sequence(thrust::device, cusp->cooPerm->begin(), cusp->cooPerm->end(), 0);
35757e8381f9SStefano Zampini     thrust::sort_by_key(fkey, ekey, cusp->cooPerm->begin(), IJCompare());
35767e8381f9SStefano Zampini     *cusp->cooPerm_a = d_i;
35777e8381f9SStefano Zampini     THRUSTINTARRAY w = d_j;
35787e8381f9SStefano Zampini 
35797e8381f9SStefano Zampini     auto nekey = thrust::unique(fkey, ekey, IJEqual());
35807e8381f9SStefano Zampini     if (nekey == ekey) { /* all entries are unique */
35817e8381f9SStefano Zampini       delete cusp->cooPerm_a;
35827e8381f9SStefano Zampini       cusp->cooPerm_a = NULL;
35837e8381f9SStefano Zampini     } else { /* I couldn't come up with a more elegant algorithm */
35847e8381f9SStefano Zampini       adjacent_difference(cusp->cooPerm_a->begin(),cusp->cooPerm_a->end(),cusp->cooPerm_a->begin(),IJDiff());
35857e8381f9SStefano Zampini       adjacent_difference(w.begin(),w.end(),w.begin(),IJDiff());
35867e8381f9SStefano Zampini       (*cusp->cooPerm_a)[0] = 0;
35877e8381f9SStefano Zampini       w[0] = 0;
35887e8381f9SStefano Zampini       thrust::transform(cusp->cooPerm_a->begin(),cusp->cooPerm_a->end(),w.begin(),cusp->cooPerm_a->begin(),IJSum());
35897e8381f9SStefano Zampini       thrust::inclusive_scan(cusp->cooPerm_a->begin(),cusp->cooPerm_a->end(),cusp->cooPerm_a->begin(),thrust::plus<PetscInt>());
35907e8381f9SStefano Zampini     }
35917e8381f9SStefano Zampini     thrust::counting_iterator<PetscInt> search_begin(0);
35927e8381f9SStefano Zampini     thrust::upper_bound(d_i.begin(), nekey.get_iterator_tuple().get<0>(),
35937e8381f9SStefano Zampini                         search_begin, search_begin + A->rmap->n,
35947e8381f9SStefano Zampini                         ii.begin());
359508391a17SStefano Zampini     cerr = WaitForCUDA();CHKERRCUDA(cerr);
359608391a17SStefano Zampini     ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
35977e8381f9SStefano Zampini 
35987e8381f9SStefano Zampini     ierr = MatSeqXAIJFreeAIJ(A,&a->a,&a->j,&a->i);CHKERRQ(ierr);
35997e8381f9SStefano Zampini     a->singlemalloc = PETSC_FALSE;
36007e8381f9SStefano Zampini     a->free_a       = PETSC_TRUE;
36017e8381f9SStefano Zampini     a->free_ij      = PETSC_TRUE;
36027e8381f9SStefano Zampini     ierr = PetscMalloc1(A->rmap->n+1,&a->i);CHKERRQ(ierr);
36037e8381f9SStefano Zampini     a->i[0] = 0;
36047e8381f9SStefano Zampini     cerr = cudaMemcpy(a->i+1,ii.data().get(),A->rmap->n*sizeof(PetscInt),cudaMemcpyDeviceToHost);CHKERRCUDA(cerr);
36057e8381f9SStefano Zampini     a->nz = a->maxnz = a->i[A->rmap->n];
3606fcdce8c4SStefano Zampini     a->rmax = 0;
36077e8381f9SStefano Zampini     ierr = PetscMalloc1(a->nz,&a->a);CHKERRQ(ierr);
36087e8381f9SStefano Zampini     ierr = PetscMalloc1(a->nz,&a->j);CHKERRQ(ierr);
36097e8381f9SStefano Zampini     cerr = cudaMemcpy(a->j,d_j.data().get(),a->nz*sizeof(PetscInt),cudaMemcpyDeviceToHost);CHKERRCUDA(cerr);
36107e8381f9SStefano Zampini     if (!a->ilen) { ierr = PetscMalloc1(A->rmap->n,&a->ilen);CHKERRQ(ierr); }
36117e8381f9SStefano Zampini     if (!a->imax) { ierr = PetscMalloc1(A->rmap->n,&a->imax);CHKERRQ(ierr); }
36127e8381f9SStefano Zampini     for (PetscInt i = 0; i < A->rmap->n; i++) {
36137e8381f9SStefano Zampini       const PetscInt nnzr = a->i[i+1] - a->i[i];
36147e8381f9SStefano Zampini       nzr += (PetscInt)!!(nnzr);
36157e8381f9SStefano Zampini       a->ilen[i] = a->imax[i] = nnzr;
3616fcdce8c4SStefano Zampini       a->rmax = PetscMax(a->rmax,nnzr);
36177e8381f9SStefano Zampini     }
3618fcdce8c4SStefano Zampini     a->nonzerorowcnt = nzr;
36197e8381f9SStefano Zampini     A->preallocated = PETSC_TRUE;
36207e8381f9SStefano Zampini     ierr = PetscLogGpuToCpu((A->rmap->n+a->nz)*sizeof(PetscInt));CHKERRQ(ierr);
3621fcdce8c4SStefano Zampini     ierr = MatMarkDiagonal_SeqAIJ(A);CHKERRQ(ierr);
36227e8381f9SStefano Zampini   } else {
36237e8381f9SStefano Zampini     ierr = MatSeqAIJSetPreallocation(A,0,NULL);CHKERRQ(ierr);
36247e8381f9SStefano Zampini   }
3625e61fc153SStefano Zampini   ierr = MatSetOption(A,MAT_NEW_NONZERO_ALLOCATION_ERR,PETSC_TRUE);CHKERRQ(ierr);
36267e8381f9SStefano Zampini 
36277e8381f9SStefano Zampini   /* We want to allocate the CUSPARSE struct for matvec now.
3628e61fc153SStefano Zampini      The code is so convoluted now that I prefer to copy zeros */
3629e61fc153SStefano Zampini   ierr = PetscArrayzero(a->a,a->nz);CHKERRQ(ierr);
36307e8381f9SStefano Zampini   ierr = MatCheckCompressedRow(A,nzr,&a->compressedrow,a->i,A->rmap->n,0.6);CHKERRQ(ierr);
36317e8381f9SStefano Zampini   A->offloadmask = PETSC_OFFLOAD_CPU;
36327e8381f9SStefano Zampini   A->nonzerostate++;
36337e8381f9SStefano Zampini   ierr = MatSeqAIJCUSPARSECopyToGPU(A);CHKERRQ(ierr);
36347e8381f9SStefano Zampini   ierr = MatSeqAIJCUSPARSEMultStruct_Destroy(&cusp->matTranspose,cusp->format);CHKERRQ(ierr);
36357e8381f9SStefano Zampini 
36367e8381f9SStefano Zampini   A->assembled = PETSC_FALSE;
36377e8381f9SStefano Zampini   A->was_assembled = PETSC_FALSE;
36387e8381f9SStefano Zampini   PetscFunctionReturn(0);
36397e8381f9SStefano Zampini }
3640ed502f03SStefano Zampini 
3641ed502f03SStefano Zampini PetscErrorCode MatSeqAIJCUSPARSEGetArrayRead(Mat A, const PetscScalar** a)
3642ed502f03SStefano Zampini {
3643ed502f03SStefano Zampini   Mat_SeqAIJCUSPARSE *cusp = (Mat_SeqAIJCUSPARSE*)A->spptr;
3644ed502f03SStefano Zampini   CsrMatrix          *csr;
3645ed502f03SStefano Zampini   PetscErrorCode     ierr;
3646ed502f03SStefano Zampini 
3647ed502f03SStefano Zampini   PetscFunctionBegin;
3648ed502f03SStefano Zampini   PetscValidHeaderSpecific(A,MAT_CLASSID,1);
3649ed502f03SStefano Zampini   PetscValidPointer(a,2);
3650ed502f03SStefano Zampini   PetscCheckTypeName(A,MATSEQAIJCUSPARSE);
3651ed502f03SStefano Zampini   if (!cusp->mat) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing Mat_SeqAIJCUSPARSEMultStruct");
3652ed502f03SStefano Zampini   if (cusp->format == MAT_CUSPARSE_ELL || cusp->format == MAT_CUSPARSE_HYB) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not implemented");
3653ed502f03SStefano Zampini   ierr = MatSeqAIJCUSPARSECopyToGPU(A);CHKERRQ(ierr);
3654ed502f03SStefano Zampini   csr = (CsrMatrix*)cusp->mat->mat;
3655ed502f03SStefano Zampini   if (!csr->values) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing CUDA memory");
3656ed502f03SStefano Zampini   *a = csr->values->data().get();
3657ed502f03SStefano Zampini   PetscFunctionReturn(0);
3658ed502f03SStefano Zampini }
3659ed502f03SStefano Zampini 
3660ed502f03SStefano Zampini PetscErrorCode MatSeqAIJCUSPARSERestoreArrayRead(Mat A, const PetscScalar** a)
3661ed502f03SStefano Zampini {
3662ed502f03SStefano Zampini   PetscFunctionBegin;
3663ed502f03SStefano Zampini   PetscValidHeaderSpecific(A,MAT_CLASSID,1);
3664ed502f03SStefano Zampini   PetscValidPointer(a,2);
3665ed502f03SStefano Zampini   PetscCheckTypeName(A,MATSEQAIJCUSPARSE);
3666ed502f03SStefano Zampini   *a = NULL;
3667ed502f03SStefano Zampini   PetscFunctionReturn(0);
3668ed502f03SStefano Zampini }
3669ed502f03SStefano Zampini 
3670ed502f03SStefano Zampini PetscErrorCode MatSeqAIJCUSPARSEGetArrayWrite(Mat A, PetscScalar** a)
3671ed502f03SStefano Zampini {
3672ed502f03SStefano Zampini   Mat_SeqAIJCUSPARSE *cusp = (Mat_SeqAIJCUSPARSE*)A->spptr;
3673ed502f03SStefano Zampini   CsrMatrix          *csr;
3674ed502f03SStefano Zampini 
3675ed502f03SStefano Zampini   PetscFunctionBegin;
3676ed502f03SStefano Zampini   PetscValidHeaderSpecific(A,MAT_CLASSID,1);
3677ed502f03SStefano Zampini   PetscValidPointer(a,2);
3678ed502f03SStefano Zampini   PetscCheckTypeName(A,MATSEQAIJCUSPARSE);
3679ed502f03SStefano Zampini   if (!cusp->mat) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing Mat_SeqAIJCUSPARSEMultStruct");
3680ed502f03SStefano Zampini   if (cusp->format == MAT_CUSPARSE_ELL || cusp->format == MAT_CUSPARSE_HYB) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not implemented");
3681ed502f03SStefano Zampini   csr = (CsrMatrix*)cusp->mat->mat;
3682ed502f03SStefano Zampini   if (!csr->values) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing CUDA memory");
3683ed502f03SStefano Zampini   *a = csr->values->data().get();
3684ed502f03SStefano Zampini   PetscFunctionReturn(0);
3685ed502f03SStefano Zampini }
3686ed502f03SStefano Zampini 
3687ed502f03SStefano Zampini PetscErrorCode MatSeqAIJCUSPARSERestoreArrayWrite(Mat A, PetscScalar** a)
3688ed502f03SStefano Zampini {
3689ed502f03SStefano Zampini   PetscErrorCode ierr;
3690ed502f03SStefano Zampini 
3691ed502f03SStefano Zampini   PetscFunctionBegin;
3692ed502f03SStefano Zampini   PetscValidHeaderSpecific(A,MAT_CLASSID,1);
3693ed502f03SStefano Zampini   PetscValidPointer(a,2);
3694ed502f03SStefano Zampini   PetscCheckTypeName(A,MATSEQAIJCUSPARSE);
3695ed502f03SStefano Zampini   ierr = PetscObjectStateIncrease((PetscObject)A);CHKERRQ(ierr);
3696ed502f03SStefano Zampini   A->offloadmask = PETSC_OFFLOAD_GPU;
3697ed502f03SStefano Zampini   *a = NULL;
3698ed502f03SStefano Zampini   PetscFunctionReturn(0);
3699ed502f03SStefano Zampini }
3700ed502f03SStefano Zampini 
3701ed502f03SStefano Zampini struct IJCompare4
3702ed502f03SStefano Zampini {
3703ed502f03SStefano Zampini   __host__ __device__
3704ed502f03SStefano Zampini   inline bool operator() (const thrust::tuple<int, int, PetscScalar, bool> &t1, const thrust::tuple<int, int, PetscScalar, bool> &t2)
3705ed502f03SStefano Zampini   {
3706ed502f03SStefano Zampini     if (t1.get<0>() < t2.get<0>()) return true;
3707ed502f03SStefano Zampini     if (t1.get<0>() == t2.get<0>()) return t1.get<1>() < t2.get<1>();
3708ed502f03SStefano Zampini     return false;
3709ed502f03SStefano Zampini   }
3710ed502f03SStefano Zampini };
3711ed502f03SStefano Zampini 
37128909a122SStefano Zampini struct Shift
37138909a122SStefano Zampini {
3714ed502f03SStefano Zampini   int _shift;
3715ed502f03SStefano Zampini 
3716ed502f03SStefano Zampini   Shift(int shift) : _shift(shift) {}
3717ed502f03SStefano Zampini   __host__ __device__
3718ed502f03SStefano Zampini   inline int operator() (const int &c)
3719ed502f03SStefano Zampini   {
3720ed502f03SStefano Zampini     return c + _shift;
3721ed502f03SStefano Zampini   }
3722ed502f03SStefano Zampini };
3723ed502f03SStefano Zampini 
3724ed502f03SStefano Zampini /* merges to SeqAIJCUSPARSE matrices, [A';B']' operation in matlab notation */
3725ed502f03SStefano Zampini PetscErrorCode MatSeqAIJCUSPARSEMergeMats(Mat A,Mat B,MatReuse reuse,Mat* C)
3726ed502f03SStefano Zampini {
3727ed502f03SStefano Zampini   PetscErrorCode               ierr;
3728ed502f03SStefano Zampini   Mat_SeqAIJ                   *a = (Mat_SeqAIJ*)A->data, *b = (Mat_SeqAIJ*)B->data, *c;
3729ed502f03SStefano Zampini   Mat_SeqAIJCUSPARSE           *Acusp = (Mat_SeqAIJCUSPARSE*)A->spptr, *Bcusp = (Mat_SeqAIJCUSPARSE*)B->spptr, *Ccusp;
3730ed502f03SStefano Zampini   Mat_SeqAIJCUSPARSEMultStruct *Cmat;
3731ed502f03SStefano Zampini   CsrMatrix                    *Acsr,*Bcsr,*Ccsr;
3732ed502f03SStefano Zampini   PetscInt                     Annz,Bnnz;
3733ed502f03SStefano Zampini   cusparseStatus_t             stat;
3734ed502f03SStefano Zampini   PetscInt                     i,m,n,zero = 0;
3735ed502f03SStefano Zampini   cudaError_t                  cerr;
3736ed502f03SStefano Zampini 
3737ed502f03SStefano Zampini   PetscFunctionBegin;
3738ed502f03SStefano Zampini   PetscValidHeaderSpecific(A,MAT_CLASSID,1);
3739ed502f03SStefano Zampini   PetscValidHeaderSpecific(B,MAT_CLASSID,2);
3740ed502f03SStefano Zampini   PetscValidPointer(C,4);
3741ed502f03SStefano Zampini   PetscCheckTypeName(A,MATSEQAIJCUSPARSE);
3742ed502f03SStefano Zampini   PetscCheckTypeName(B,MATSEQAIJCUSPARSE);
3743ed502f03SStefano Zampini   if (A->rmap->n != B->rmap->n) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Invalid number or rows %D != %D",A->rmap->n,B->rmap->n);
3744ed502f03SStefano Zampini   if (reuse == MAT_INPLACE_MATRIX) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"MAT_INPLACE_MATRIX not supported");
3745ed502f03SStefano Zampini   if (Acusp->format == MAT_CUSPARSE_ELL || Acusp->format == MAT_CUSPARSE_HYB) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not implemented");
3746ed502f03SStefano Zampini   if (Bcusp->format == MAT_CUSPARSE_ELL || Bcusp->format == MAT_CUSPARSE_HYB) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not implemented");
3747ed502f03SStefano Zampini   if (reuse == MAT_INITIAL_MATRIX) {
3748ed502f03SStefano Zampini     m     = A->rmap->n;
3749ed502f03SStefano Zampini     n     = A->cmap->n + B->cmap->n;
3750ed502f03SStefano Zampini     ierr  = MatCreate(PETSC_COMM_SELF,C);CHKERRQ(ierr);
3751ed502f03SStefano Zampini     ierr  = MatSetSizes(*C,m,n,m,n);CHKERRQ(ierr);
3752ed502f03SStefano Zampini     ierr  = MatSetType(*C,MATSEQAIJCUSPARSE);CHKERRQ(ierr);
3753ed502f03SStefano Zampini     c     = (Mat_SeqAIJ*)(*C)->data;
3754ed502f03SStefano Zampini     Ccusp = (Mat_SeqAIJCUSPARSE*)(*C)->spptr;
3755ed502f03SStefano Zampini     Cmat  = new Mat_SeqAIJCUSPARSEMultStruct;
3756ed502f03SStefano Zampini     Ccsr  = new CsrMatrix;
3757ed502f03SStefano Zampini     Cmat->cprowIndices      = NULL;
3758ed502f03SStefano Zampini     c->compressedrow.use    = PETSC_FALSE;
3759ed502f03SStefano Zampini     c->compressedrow.nrows  = 0;
3760ed502f03SStefano Zampini     c->compressedrow.i      = NULL;
3761ed502f03SStefano Zampini     c->compressedrow.rindex = NULL;
3762ed502f03SStefano Zampini     Ccusp->workVector       = NULL;
3763ed502f03SStefano Zampini     Ccusp->nrows    = m;
3764ed502f03SStefano Zampini     Ccusp->mat      = Cmat;
3765ed502f03SStefano Zampini     Ccusp->mat->mat = Ccsr;
3766ed502f03SStefano Zampini     Ccsr->num_rows  = m;
3767ed502f03SStefano Zampini     Ccsr->num_cols  = n;
3768ed502f03SStefano Zampini     stat = cusparseCreateMatDescr(&Cmat->descr);CHKERRCUSPARSE(stat);
3769ed502f03SStefano Zampini     stat = cusparseSetMatIndexBase(Cmat->descr, CUSPARSE_INDEX_BASE_ZERO);CHKERRCUSPARSE(stat);
3770ed502f03SStefano Zampini     stat = cusparseSetMatType(Cmat->descr, CUSPARSE_MATRIX_TYPE_GENERAL);CHKERRCUSPARSE(stat);
3771ed502f03SStefano Zampini     cerr = cudaMalloc((void **)&(Cmat->alpha_one),sizeof(PetscScalar));CHKERRCUDA(cerr);
3772ed502f03SStefano Zampini     cerr = cudaMalloc((void **)&(Cmat->beta_zero),sizeof(PetscScalar));CHKERRCUDA(cerr);
3773ed502f03SStefano Zampini     cerr = cudaMalloc((void **)&(Cmat->beta_one), sizeof(PetscScalar));CHKERRCUDA(cerr);
3774ed502f03SStefano Zampini     cerr = cudaMemcpy(Cmat->alpha_one,&PETSC_CUSPARSE_ONE, sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(cerr);
3775ed502f03SStefano Zampini     cerr = cudaMemcpy(Cmat->beta_zero,&PETSC_CUSPARSE_ZERO,sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(cerr);
3776ed502f03SStefano Zampini     cerr = cudaMemcpy(Cmat->beta_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(cerr);
3777ed502f03SStefano Zampini     ierr = MatSeqAIJCUSPARSECopyToGPU(A);CHKERRQ(ierr);
3778ed502f03SStefano Zampini     ierr = MatSeqAIJCUSPARSECopyToGPU(B);CHKERRQ(ierr);
3779ed502f03SStefano Zampini     ierr = MatSeqAIJCUSPARSEGenerateTransposeForMult(A);CHKERRQ(ierr);
3780ed502f03SStefano Zampini     ierr = MatSeqAIJCUSPARSEGenerateTransposeForMult(B);CHKERRQ(ierr);
3781ed502f03SStefano Zampini     if (!Acusp->mat) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing Mat_SeqAIJCUSPARSEMultStruct");
3782ed502f03SStefano Zampini     if (!Bcusp->mat) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing Mat_SeqAIJCUSPARSEMultStruct");
3783ed502f03SStefano Zampini 
3784ed502f03SStefano Zampini     Acsr = (CsrMatrix*)Acusp->mat->mat;
3785ed502f03SStefano Zampini     Bcsr = (CsrMatrix*)Bcusp->mat->mat;
3786ed502f03SStefano Zampini     Annz = (PetscInt)Acsr->column_indices->size();
3787ed502f03SStefano Zampini     Bnnz = (PetscInt)Bcsr->column_indices->size();
3788ed502f03SStefano Zampini     c->nz = Annz + Bnnz;
3789ed502f03SStefano Zampini     Ccsr->row_offsets = new THRUSTINTARRAY32(m+1);
3790ed502f03SStefano Zampini     Ccsr->column_indices = new THRUSTINTARRAY32(c->nz);
3791ed502f03SStefano Zampini     Ccsr->values = new THRUSTARRAY(c->nz);
3792ed502f03SStefano Zampini     Ccsr->num_entries = c->nz;
3793ed502f03SStefano Zampini     Ccusp->cooPerm = new THRUSTINTARRAY(c->nz);
3794ed502f03SStefano Zampini 
3795ed502f03SStefano Zampini     if (c->nz) {
3796ed502f03SStefano Zampini       THRUSTINTARRAY32 Acoo(Annz);
3797ed502f03SStefano Zampini       THRUSTINTARRAY32 Bcoo(Bnnz);
3798ed502f03SStefano Zampini       THRUSTINTARRAY32 *roff;
3799ed502f03SStefano Zampini       if (a->compressedrow.use) { /* need full row offset */
3800ed502f03SStefano Zampini         if (!Acusp->rowoffsets_gpu) {
3801ed502f03SStefano Zampini           Acusp->rowoffsets_gpu  = new THRUSTINTARRAY32(A->rmap->n + 1);
3802ed502f03SStefano Zampini           Acusp->rowoffsets_gpu->assign(a->i,a->i + A->rmap->n + 1);
3803ed502f03SStefano Zampini           ierr = PetscLogCpuToGpu((A->rmap->n + 1)*sizeof(PetscInt));CHKERRQ(ierr);
3804ed502f03SStefano Zampini         }
3805ed502f03SStefano Zampini         roff = Acusp->rowoffsets_gpu;
3806ed502f03SStefano Zampini       } else roff = Acsr->row_offsets;
3807ed502f03SStefano Zampini       ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
3808ed502f03SStefano Zampini       stat = cusparseXcsr2coo(Acusp->handle,
3809ed502f03SStefano Zampini                               roff->data().get(),
3810ed502f03SStefano Zampini                               Annz,
3811ed502f03SStefano Zampini                               m,
3812ed502f03SStefano Zampini                               Acoo.data().get(),
3813ed502f03SStefano Zampini                               CUSPARSE_INDEX_BASE_ZERO);CHKERRCUSPARSE(stat);
3814ed502f03SStefano Zampini       cerr = WaitForCUDA();CHKERRCUDA(cerr);
3815ed502f03SStefano Zampini       ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
3816ed502f03SStefano Zampini       if (b->compressedrow.use) { /* need full row offset */
3817ed502f03SStefano Zampini         if (!Bcusp->rowoffsets_gpu) {
3818ed502f03SStefano Zampini           Bcusp->rowoffsets_gpu  = new THRUSTINTARRAY32(B->rmap->n + 1);
3819ed502f03SStefano Zampini           Bcusp->rowoffsets_gpu->assign(b->i,b->i + B->rmap->n + 1);
3820ed502f03SStefano Zampini           ierr = PetscLogCpuToGpu((B->rmap->n + 1)*sizeof(PetscInt));CHKERRQ(ierr);
3821ed502f03SStefano Zampini         }
3822ed502f03SStefano Zampini         roff = Bcusp->rowoffsets_gpu;
3823ed502f03SStefano Zampini       } else roff = Bcsr->row_offsets;
3824ed502f03SStefano Zampini       ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
3825ed502f03SStefano Zampini       stat = cusparseXcsr2coo(Bcusp->handle,
3826ed502f03SStefano Zampini                               roff->data().get(),
3827ed502f03SStefano Zampini                               Bnnz,
3828ed502f03SStefano Zampini                               m,
3829ed502f03SStefano Zampini                               Bcoo.data().get(),
3830ed502f03SStefano Zampini                               CUSPARSE_INDEX_BASE_ZERO);CHKERRCUSPARSE(stat);
3831ed502f03SStefano Zampini       cerr = WaitForCUDA();CHKERRCUDA(cerr);
3832ed502f03SStefano Zampini       ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
3833ed502f03SStefano Zampini       THRUSTINTARRAY32 Ccoo(c->nz);
3834ed502f03SStefano Zampini       auto Aperm = thrust::make_constant_iterator(true);
3835ed502f03SStefano Zampini       auto Bperm = thrust::make_constant_iterator(false);
38368909a122SStefano Zampini #if PETSC_PKG_CUDA_VERSION_GE(10,0,0)
3837ed502f03SStefano Zampini       auto Bcib = thrust::make_transform_iterator(Bcsr->column_indices->begin(),Shift(A->cmap->n));
3838ed502f03SStefano Zampini       auto Bcie = thrust::make_transform_iterator(Bcsr->column_indices->end(),Shift(A->cmap->n));
38398909a122SStefano Zampini #else
38408909a122SStefano Zampini       /* there are issues instantiating the merge operation using a transform iterator for the columns of B */
38418909a122SStefano Zampini       auto Bcib = Bcsr->column_indices->begin();
38428909a122SStefano Zampini       auto Bcie = Bcsr->column_indices->end();
38438909a122SStefano Zampini       thrust::transform(Bcib,Bcie,Bcib,Shift(A->cmap->n));
38448909a122SStefano Zampini #endif
3845ed502f03SStefano Zampini       thrust::device_vector<bool> wPerm(Annz+Bnnz);
3846ed502f03SStefano Zampini       auto Azb = thrust::make_zip_iterator(thrust::make_tuple(Acoo.begin(),Acsr->column_indices->begin(),Acsr->values->begin(),Aperm));
3847ed502f03SStefano Zampini       auto Aze = thrust::make_zip_iterator(thrust::make_tuple(Acoo.end(),Acsr->column_indices->end(),Acsr->values->end(),Aperm));
3848ed502f03SStefano Zampini       auto Bzb = thrust::make_zip_iterator(thrust::make_tuple(Bcoo.begin(),Bcib,Bcsr->values->begin(),Bperm));
3849ed502f03SStefano Zampini       auto Bze = thrust::make_zip_iterator(thrust::make_tuple(Bcoo.end(),Bcie,Bcsr->values->end(),Bperm));
3850ed502f03SStefano Zampini       auto Czb = thrust::make_zip_iterator(thrust::make_tuple(Ccoo.begin(),Ccsr->column_indices->begin(),Ccsr->values->begin(),wPerm.begin()));
3851ed502f03SStefano Zampini       auto p1 = Ccusp->cooPerm->begin();
3852ed502f03SStefano Zampini       auto p2 = Ccusp->cooPerm->begin();
3853ed502f03SStefano Zampini       thrust::advance(p2,Annz);
3854ed502f03SStefano Zampini       ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
3855ed502f03SStefano Zampini       thrust::merge(Azb,Aze,Bzb,Bze,Czb,IJCompare4());
38568909a122SStefano Zampini #if PETSC_PKG_CUDA_VERSION_LT(10,0,0)
38578909a122SStefano Zampini       thrust::transform(Bcib,Bcie,Bcib,Shift(-A->cmap->n));
38588909a122SStefano Zampini #endif
3859ed502f03SStefano Zampini       thrust::partition_copy(thrust::make_counting_iterator(zero),thrust::make_counting_iterator(c->nz),wPerm.begin(),p1,p2,thrust::identity<bool>());
3860ed502f03SStefano Zampini       stat = cusparseXcoo2csr(Ccusp->handle,
3861ed502f03SStefano Zampini                               Ccoo.data().get(),
3862ed502f03SStefano Zampini                               c->nz,
3863ed502f03SStefano Zampini                               m,
3864ed502f03SStefano Zampini                               Ccsr->row_offsets->data().get(),
3865ed502f03SStefano Zampini                               CUSPARSE_INDEX_BASE_ZERO);CHKERRCUSPARSE(stat);
3866ed502f03SStefano Zampini       cerr = WaitForCUDA();CHKERRCUDA(cerr);
3867ed502f03SStefano Zampini       ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
3868ed502f03SStefano Zampini #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
3869ed502f03SStefano Zampini       stat = cusparseCreateCsr(&Cmat->matDescr, Ccsr->num_rows, Ccsr->num_cols, Ccsr->num_entries,
3870ed502f03SStefano Zampini                                Ccsr->row_offsets->data().get(), Ccsr->column_indices->data().get(), Ccsr->values->data().get(),
3871ed502f03SStefano Zampini                                CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I,
3872ed502f03SStefano Zampini                                CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype);CHKERRCUSPARSE(stat);
3873ed502f03SStefano Zampini #endif
3874ed502f03SStefano Zampini       if (Acusp->transgen && Bcusp->transgen) { /* if A and B have the transpose, generate C transpose too */
3875ed502f03SStefano Zampini         PetscBool AT = Acusp->matTranspose ? PETSC_TRUE : PETSC_FALSE, BT = Bcusp->matTranspose ? PETSC_TRUE : PETSC_FALSE;
3876ed502f03SStefano Zampini         Mat_SeqAIJCUSPARSEMultStruct *CmatT = new Mat_SeqAIJCUSPARSEMultStruct;
3877ed502f03SStefano Zampini         CsrMatrix *CcsrT = new CsrMatrix;
3878ed502f03SStefano Zampini         CsrMatrix *AcsrT = AT ? (CsrMatrix*)Acusp->matTranspose->mat : NULL;
3879ed502f03SStefano Zampini         CsrMatrix *BcsrT = BT ? (CsrMatrix*)Bcusp->matTranspose->mat : NULL;
3880ed502f03SStefano Zampini 
3881ed502f03SStefano Zampini         Ccusp->transgen = PETSC_TRUE;
3882ed502f03SStefano Zampini         CmatT->cprowIndices  = NULL;
3883ed502f03SStefano Zampini         CmatT->mat = CcsrT;
3884ed502f03SStefano Zampini         CcsrT->num_rows = n;
3885ed502f03SStefano Zampini         CcsrT->num_cols = m;
3886ed502f03SStefano Zampini         CcsrT->num_entries = c->nz;
3887ed502f03SStefano Zampini 
3888ed502f03SStefano Zampini         CcsrT->row_offsets = new THRUSTINTARRAY32(n+1);
3889ed502f03SStefano Zampini         CcsrT->column_indices = new THRUSTINTARRAY32(c->nz);
3890ed502f03SStefano Zampini         CcsrT->values = new THRUSTARRAY(c->nz);
3891ed502f03SStefano Zampini 
3892ed502f03SStefano Zampini         ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
3893ed502f03SStefano Zampini         auto rT = CcsrT->row_offsets->begin();
3894ed502f03SStefano Zampini         if (AT) {
3895ed502f03SStefano Zampini           rT = thrust::copy(AcsrT->row_offsets->begin(),AcsrT->row_offsets->end(),rT);
3896ed502f03SStefano Zampini           thrust::advance(rT,-1);
3897ed502f03SStefano Zampini         }
3898ed502f03SStefano Zampini         if (BT) {
3899ed502f03SStefano Zampini           auto titb = thrust::make_transform_iterator(BcsrT->row_offsets->begin(),Shift(a->nz));
3900ed502f03SStefano Zampini           auto tite = thrust::make_transform_iterator(BcsrT->row_offsets->end(),Shift(a->nz));
3901ed502f03SStefano Zampini           thrust::copy(titb,tite,rT);
3902ed502f03SStefano Zampini         }
3903ed502f03SStefano Zampini         auto cT = CcsrT->column_indices->begin();
3904ed502f03SStefano Zampini         if (AT) cT = thrust::copy(AcsrT->column_indices->begin(),AcsrT->column_indices->end(),cT);
3905ed502f03SStefano Zampini         if (BT) thrust::copy(BcsrT->column_indices->begin(),BcsrT->column_indices->end(),cT);
3906ed502f03SStefano Zampini         auto vT = CcsrT->values->begin();
3907ed502f03SStefano Zampini         if (AT) vT = thrust::copy(AcsrT->values->begin(),AcsrT->values->end(),vT);
3908ed502f03SStefano Zampini         if (BT) thrust::copy(BcsrT->values->begin(),BcsrT->values->end(),vT);
3909ed502f03SStefano Zampini         cerr = WaitForCUDA();CHKERRCUDA(cerr);
3910ed502f03SStefano Zampini         ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
3911ed502f03SStefano Zampini 
3912ed502f03SStefano Zampini         stat = cusparseCreateMatDescr(&CmatT->descr);CHKERRCUSPARSE(stat);
3913ed502f03SStefano Zampini         stat = cusparseSetMatIndexBase(CmatT->descr, CUSPARSE_INDEX_BASE_ZERO);CHKERRCUSPARSE(stat);
3914ed502f03SStefano Zampini         stat = cusparseSetMatType(CmatT->descr, CUSPARSE_MATRIX_TYPE_GENERAL);CHKERRCUSPARSE(stat);
3915ed502f03SStefano Zampini         cerr = cudaMalloc((void **)&(CmatT->alpha_one),sizeof(PetscScalar));CHKERRCUDA(cerr);
3916ed502f03SStefano Zampini         cerr = cudaMalloc((void **)&(CmatT->beta_zero),sizeof(PetscScalar));CHKERRCUDA(cerr);
3917ed502f03SStefano Zampini         cerr = cudaMalloc((void **)&(CmatT->beta_one), sizeof(PetscScalar));CHKERRCUDA(cerr);
3918ed502f03SStefano Zampini         cerr = cudaMemcpy(CmatT->alpha_one,&PETSC_CUSPARSE_ONE, sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(cerr);
3919ed502f03SStefano Zampini         cerr = cudaMemcpy(CmatT->beta_zero,&PETSC_CUSPARSE_ZERO,sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(cerr);
3920ed502f03SStefano Zampini         cerr = cudaMemcpy(CmatT->beta_one, &PETSC_CUSPARSE_ONE, sizeof(PetscScalar),cudaMemcpyHostToDevice);CHKERRCUDA(cerr);
3921ed502f03SStefano Zampini #if PETSC_PKG_CUDA_VERSION_GE(11,0,0)
3922ed502f03SStefano Zampini         stat = cusparseCreateCsr(&CmatT->matDescr, CcsrT->num_rows, CcsrT->num_cols, CcsrT->num_entries,
3923ed502f03SStefano Zampini                                  CcsrT->row_offsets->data().get(), CcsrT->column_indices->data().get(), CcsrT->values->data().get(),
3924ed502f03SStefano Zampini                                  CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I,
3925ed502f03SStefano Zampini                                  CUSPARSE_INDEX_BASE_ZERO, cusparse_scalartype);CHKERRCUSPARSE(stat);
3926ed502f03SStefano Zampini #endif
3927ed502f03SStefano Zampini         Ccusp->matTranspose = CmatT;
3928ed502f03SStefano Zampini       }
3929ed502f03SStefano Zampini     }
3930ed502f03SStefano Zampini 
3931ed502f03SStefano Zampini     c->singlemalloc = PETSC_FALSE;
3932ed502f03SStefano Zampini     c->free_a       = PETSC_TRUE;
3933ed502f03SStefano Zampini     c->free_ij      = PETSC_TRUE;
3934ed502f03SStefano Zampini     ierr = PetscMalloc1(m+1,&c->i);CHKERRQ(ierr);
3935ed502f03SStefano Zampini     ierr = PetscMalloc1(c->nz,&c->j);CHKERRQ(ierr);
3936ed502f03SStefano Zampini     if (PetscDefined(USE_64BIT_INDICES)) { /* 32 to 64 bit conversion on the GPU and then copy to host (lazy) */
3937ed502f03SStefano Zampini       THRUSTINTARRAY ii(Ccsr->row_offsets->size());
3938ed502f03SStefano Zampini       THRUSTINTARRAY jj(Ccsr->column_indices->size());
3939ed502f03SStefano Zampini       ii   = *Ccsr->row_offsets;
3940ed502f03SStefano Zampini       jj   = *Ccsr->column_indices;
3941ed502f03SStefano Zampini       cerr = cudaMemcpy(c->i,ii.data().get(),Ccsr->row_offsets->size()*sizeof(PetscInt),cudaMemcpyDeviceToHost);CHKERRCUDA(cerr);
3942ed502f03SStefano Zampini       cerr = cudaMemcpy(c->j,jj.data().get(),Ccsr->column_indices->size()*sizeof(PetscInt),cudaMemcpyDeviceToHost);CHKERRCUDA(cerr);
3943ed502f03SStefano Zampini     } else {
3944ed502f03SStefano Zampini       cerr = cudaMemcpy(c->i,Ccsr->row_offsets->data().get(),Ccsr->row_offsets->size()*sizeof(PetscInt),cudaMemcpyDeviceToHost);CHKERRCUDA(cerr);
3945ed502f03SStefano Zampini       cerr = cudaMemcpy(c->j,Ccsr->column_indices->data().get(),Ccsr->column_indices->size()*sizeof(PetscInt),cudaMemcpyDeviceToHost);CHKERRCUDA(cerr);
3946ed502f03SStefano Zampini     }
3947ed502f03SStefano Zampini     ierr = PetscLogGpuToCpu((Ccsr->column_indices->size() + Ccsr->row_offsets->size())*sizeof(PetscInt));CHKERRQ(ierr);
3948ed502f03SStefano Zampini     ierr = PetscMalloc1(m,&c->ilen);CHKERRQ(ierr);
3949ed502f03SStefano Zampini     ierr = PetscMalloc1(m,&c->imax);CHKERRQ(ierr);
3950ed502f03SStefano Zampini     c->maxnz = c->nz;
3951ed502f03SStefano Zampini     c->nonzerorowcnt = 0;
3952ed502f03SStefano Zampini     c->rmax = 0;
3953ed502f03SStefano Zampini     for (i = 0; i < m; i++) {
3954ed502f03SStefano Zampini       const PetscInt nn = c->i[i+1] - c->i[i];
3955ed502f03SStefano Zampini       c->ilen[i] = c->imax[i] = nn;
3956ed502f03SStefano Zampini       c->nonzerorowcnt += (PetscInt)!!nn;
3957ed502f03SStefano Zampini       c->rmax = PetscMax(c->rmax,nn);
3958ed502f03SStefano Zampini     }
3959ed502f03SStefano Zampini     ierr = MatMarkDiagonal_SeqAIJ(*C);CHKERRQ(ierr);
3960ed502f03SStefano Zampini     ierr = PetscMalloc1(c->nz,&c->a);CHKERRQ(ierr);
3961ed502f03SStefano Zampini     (*C)->nonzerostate++;
3962ed502f03SStefano Zampini     ierr = PetscLayoutSetUp((*C)->rmap);CHKERRQ(ierr);
3963ed502f03SStefano Zampini     ierr = PetscLayoutSetUp((*C)->cmap);CHKERRQ(ierr);
3964ed502f03SStefano Zampini     Ccusp->nonzerostate = (*C)->nonzerostate;
3965ed502f03SStefano Zampini     (*C)->preallocated  = PETSC_TRUE;
3966ed502f03SStefano Zampini   } else {
3967ed502f03SStefano Zampini     if ((*C)->rmap->n != B->rmap->n) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Invalid number or rows %D != %D",(*C)->rmap->n,B->rmap->n);
3968ed502f03SStefano Zampini     c = (Mat_SeqAIJ*)(*C)->data;
3969ed502f03SStefano Zampini     if (c->nz) {
3970ed502f03SStefano Zampini       Ccusp = (Mat_SeqAIJCUSPARSE*)(*C)->spptr;
3971ed502f03SStefano Zampini       if (!Ccusp->cooPerm) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing cooPerm");
3972ed502f03SStefano Zampini       if (Ccusp->format == MAT_CUSPARSE_ELL || Ccusp->format == MAT_CUSPARSE_HYB) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"Not implemented");
3973ed502f03SStefano Zampini       if (Ccusp->nonzerostate != (*C)->nonzerostate) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Wrong nonzerostate");
3974ed502f03SStefano Zampini       ierr = MatSeqAIJCUSPARSECopyToGPU(A);CHKERRQ(ierr);
3975ed502f03SStefano Zampini       ierr = MatSeqAIJCUSPARSECopyToGPU(B);CHKERRQ(ierr);
3976ed502f03SStefano Zampini       if (!Acusp->mat) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing Mat_SeqAIJCUSPARSEMultStruct");
3977ed502f03SStefano Zampini       if (!Bcusp->mat) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing Mat_SeqAIJCUSPARSEMultStruct");
3978ed502f03SStefano Zampini       Acsr = (CsrMatrix*)Acusp->mat->mat;
3979ed502f03SStefano Zampini       Bcsr = (CsrMatrix*)Bcusp->mat->mat;
3980ed502f03SStefano Zampini       Ccsr = (CsrMatrix*)Ccusp->mat->mat;
3981ed502f03SStefano Zampini       if (Acsr->num_entries != (PetscInt)Acsr->values->size()) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_COR,"A nnz %D != %D",Acsr->num_entries,(PetscInt)Acsr->values->size());
3982ed502f03SStefano Zampini       if (Bcsr->num_entries != (PetscInt)Bcsr->values->size()) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_COR,"B nnz %D != %D",Bcsr->num_entries,(PetscInt)Bcsr->values->size());
3983ed502f03SStefano Zampini       if (Ccsr->num_entries != (PetscInt)Ccsr->values->size()) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_COR,"C nnz %D != %D",Ccsr->num_entries,(PetscInt)Ccsr->values->size());
3984ed502f03SStefano Zampini       if (Ccsr->num_entries != Acsr->num_entries + Bcsr->num_entries) SETERRQ3(PETSC_COMM_SELF,PETSC_ERR_COR,"C nnz %D != %D + %D",Ccsr->num_entries,Acsr->num_entries,Bcsr->num_entries);
3985ed502f03SStefano Zampini       if (Ccusp->cooPerm->size() != Ccsr->values->size()) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_COR,"permSize %D != %D",(PetscInt)Ccusp->cooPerm->size(),(PetscInt)Ccsr->values->size());
3986ed502f03SStefano Zampini       auto pmid = Ccusp->cooPerm->begin();
3987ed502f03SStefano Zampini       thrust::advance(pmid,Acsr->num_entries);
3988ed502f03SStefano Zampini       ierr = PetscLogGpuTimeBegin();CHKERRQ(ierr);
3989ed502f03SStefano Zampini       auto zibait = thrust::make_zip_iterator(thrust::make_tuple(Acsr->values->begin(),
3990ed502f03SStefano Zampini                                                                  thrust::make_permutation_iterator(Ccsr->values->begin(),Ccusp->cooPerm->begin())));
3991ed502f03SStefano Zampini       auto zieait = thrust::make_zip_iterator(thrust::make_tuple(Acsr->values->end(),
3992ed502f03SStefano Zampini                                                                  thrust::make_permutation_iterator(Ccsr->values->begin(),pmid)));
3993ed502f03SStefano Zampini       thrust::for_each(zibait,zieait,VecCUDAEquals());
3994ed502f03SStefano Zampini       auto zibbit = thrust::make_zip_iterator(thrust::make_tuple(Bcsr->values->begin(),
3995ed502f03SStefano Zampini                                                                  thrust::make_permutation_iterator(Ccsr->values->begin(),pmid)));
3996ed502f03SStefano Zampini       auto ziebit = thrust::make_zip_iterator(thrust::make_tuple(Bcsr->values->end(),
3997ed502f03SStefano Zampini                                                                  thrust::make_permutation_iterator(Ccsr->values->begin(),Ccusp->cooPerm->end())));
3998ed502f03SStefano Zampini       thrust::for_each(zibbit,ziebit,VecCUDAEquals());
3999ed502f03SStefano Zampini       if (Acusp->transgen && Bcusp->transgen && Ccusp->transgen) {
4000ed502f03SStefano Zampini         if (!Ccusp->matTranspose) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_COR,"Missing transpose Mat_SeqAIJCUSPARSEMultStruct");
4001ed502f03SStefano Zampini         PetscBool AT = Acusp->matTranspose ? PETSC_TRUE : PETSC_FALSE, BT = Bcusp->matTranspose ? PETSC_TRUE : PETSC_FALSE;
4002ed502f03SStefano Zampini         CsrMatrix *AcsrT = AT ? (CsrMatrix*)Acusp->matTranspose->mat : NULL;
4003ed502f03SStefano Zampini         CsrMatrix *BcsrT = BT ? (CsrMatrix*)Bcusp->matTranspose->mat : NULL;
4004ed502f03SStefano Zampini         CsrMatrix *CcsrT = (CsrMatrix*)Ccusp->matTranspose->mat;
4005ed502f03SStefano Zampini         auto vT = CcsrT->values->begin();
4006ed502f03SStefano Zampini         if (AT) vT = thrust::copy(AcsrT->values->begin(),AcsrT->values->end(),vT);
4007ed502f03SStefano Zampini         if (BT) thrust::copy(BcsrT->values->begin(),BcsrT->values->end(),vT);
4008ed502f03SStefano Zampini       }
4009ed502f03SStefano Zampini       cerr = WaitForCUDA();CHKERRCUDA(cerr);
4010ed502f03SStefano Zampini       ierr = PetscLogGpuTimeEnd();CHKERRQ(ierr);
4011ed502f03SStefano Zampini     }
4012ed502f03SStefano Zampini   }
4013ed502f03SStefano Zampini   ierr = PetscObjectStateIncrease((PetscObject)*C);CHKERRQ(ierr);
4014ed502f03SStefano Zampini   (*C)->assembled     = PETSC_TRUE;
4015ed502f03SStefano Zampini   (*C)->was_assembled = PETSC_FALSE;
4016ed502f03SStefano Zampini   (*C)->offloadmask   = PETSC_OFFLOAD_GPU;
4017ed502f03SStefano Zampini   PetscFunctionReturn(0);
4018ed502f03SStefano Zampini }
4019