xref: /petsc/src/mat/impls/aij/mpi/strumpack/strumpack.c (revision 7d6ea485c49fcf8e9bc72cfc84041f9531d1ac38)
1 #include <../src/mat/impls/aij/seq/aij.h>
2 #include <../src/mat/impls/aij/mpi/mpiaij.h>
3 #include <StrumpackSparseSolver.h>
4 
5 /*
6   These are only relevant for MATMPIAIJ, not for MATSEQAIJ.
7     REPLICATED  - STRUMPACK expects the entire sparse matrix and right-hand side on every process.
8     DISTRIBUTED - STRUMPACK expects the sparse matrix and right-hand side to be distributed across the entire MPI communicator.
9 */
10 typedef enum {REPLICATED, DISTRIBUTED} STRUMPACK_MatInputMode;
11 const char *STRUMPACK_MatInputModes[] = {"REPLICATED","DISTRIBUTED","STRUMPACK_MatInputMode","PETSC_",0};
12 
13 typedef struct {
14   STRUMPACK_SparseSolver S;
15   STRUMPACK_MatInputMode MatInputMode;
16 } STRUMPACK_data;
17 
18 extern PetscErrorCode MatFactorInfo_STRUMPACK(Mat,PetscViewer);
19 extern PetscErrorCode MatLUFactorNumeric_STRUMPACK(Mat,Mat,const MatFactorInfo*);
20 extern PetscErrorCode MatDestroy_STRUMPACK(Mat);
21 extern PetscErrorCode MatView_STRUMPACK(Mat,PetscViewer);
22 extern PetscErrorCode MatSolve_STRUMPACK(Mat,Vec,Vec);
23 extern PetscErrorCode MatLUFactorSymbolic_STRUMPACK(Mat,Mat,IS,IS,const MatFactorInfo*);
24 extern PetscErrorCode MatDestroy_MPIAIJ(Mat);
25 
26 #undef __FUNCT__
27 #define __FUNCT__ "MatGetDiagonal_STRUMPACK"
28 PetscErrorCode MatGetDiagonal_STRUMPACK(Mat A,Vec v)
29 {
30   PetscFunctionBegin;
31   SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_SUP,"Mat type: STRUMPACK factor");
32   PetscFunctionReturn(0);
33 }
34 
35 #undef __FUNCT__
36 #define __FUNCT__ "MatDestroy_STRUMPACK"
37 PetscErrorCode MatDestroy_STRUMPACK(Mat A)
38 {
39   STRUMPACK_data *sp = (STRUMPACK_data*)A->spptr;
40   PetscErrorCode ierr;
41   PetscBool      flg;
42 
43   PetscFunctionBegin;
44   /* Deallocate STRUMPACK storage */
45   PetscStackCall("STRUMPACK_destroy",STRUMPACK_destroy(&(sp->S)));
46   ierr = PetscFree(A->spptr);CHKERRQ(ierr);
47   ierr = PetscObjectTypeCompare((PetscObject)A,MATSEQAIJ,&flg);CHKERRQ(ierr);
48   if (flg) {
49     ierr = MatDestroy_SeqAIJ(A);CHKERRQ(ierr);
50   } else {
51     ierr = MatDestroy_MPIAIJ(A);CHKERRQ(ierr);
52   }
53   PetscFunctionReturn(0);
54 }
55 
56 #undef __FUNCT__
57 #define __FUNCT__ "MatSolve_STRUMPACK"
58 PetscErrorCode MatSolve_STRUMPACK(Mat A,Vec b_mpi,Vec x)
59 {
60   STRUMPACK_data        *sp = (STRUMPACK_data*)A->spptr;
61   STRUMPACK_RETURN_CODE sp_err;
62   PetscErrorCode        ierr;
63   PetscMPIInt           size;
64   PetscInt              N=A->cmap->N;
65   const PetscScalar     *bptr;
66   PetscScalar           *xptr;
67   Vec                   x_seq,b_seq;
68   IS                    iden;
69   VecScatter            scat;
70 
71   PetscFunctionBegin;
72   ierr = MPI_Comm_size(PetscObjectComm((PetscObject)A),&size);CHKERRQ(ierr);
73   if (size > 1 && sp->MatInputMode == REPLICATED) {
74     ierr = VecCreateSeq(PETSC_COMM_SELF,N,&x_seq);CHKERRQ(ierr);
75     ierr = VecGetArray(x_seq,&xptr);CHKERRQ(ierr);
76     /* global mat input, convert b to b_seq */
77     ierr = VecCreateSeq(PETSC_COMM_SELF,N,&b_seq);CHKERRQ(ierr);
78     ierr = ISCreateStride(PETSC_COMM_SELF,N,0,1,&iden);CHKERRQ(ierr);
79     ierr = VecScatterCreate(b_mpi,iden,b_seq,iden,&scat);CHKERRQ(ierr);
80     ierr = ISDestroy(&iden);CHKERRQ(ierr);
81     ierr = VecScatterBegin(scat,b_mpi,b_seq,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
82     ierr = VecScatterEnd(scat,b_mpi,b_seq,INSERT_VALUES,SCATTER_FORWARD);CHKERRQ(ierr);
83     ierr = VecGetArrayRead(b_seq,&bptr);CHKERRQ(ierr);
84   } else { /* size==1 || distributed mat input */
85     ierr = VecGetArray(x,&xptr);CHKERRQ(ierr);
86     ierr = VecGetArrayRead(b_mpi,&bptr);CHKERRQ(ierr);
87   }
88 
89   PetscStackCall("STRUMPACK_solve",sp_err = STRUMPACK_solve(sp->S,(PetscScalar*)bptr,xptr,0));
90 
91   if (sp_err != STRUMPACK_SUCCESS) {
92     if (sp_err == STRUMPACK_MATRIX_NOT_SET)        SETERRQ(PETSC_COMM_SELF,PETSC_ERR_LIB,"STRUMPACK error: matrix was not set");
93     else if (sp_err == STRUMPACK_REORDERING_ERROR) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_LIB,"STRUMPACK error: matrix reordering failed");
94     else                                           SETERRQ(PETSC_COMM_SELF,PETSC_ERR_LIB,"STRUMPACK error: solve failed");
95   }
96 
97   if (size > 1 && sp->MatInputMode == REPLICATED) {
98     ierr = VecRestoreArrayRead(b_seq,&bptr);CHKERRQ(ierr);
99     ierr = VecDestroy(&b_seq);CHKERRQ(ierr);
100     /* convert seq x to mpi x */
101     ierr = VecRestoreArray(x_seq,&xptr);CHKERRQ(ierr);
102     ierr = VecScatterBegin(scat,x_seq,x,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
103     ierr = VecScatterEnd(scat,x_seq,x,INSERT_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
104     ierr = VecScatterDestroy(&scat);CHKERRQ(ierr);
105     ierr = VecDestroy(&x_seq);CHKERRQ(ierr);
106   } else {
107     ierr = VecRestoreArray(x,&xptr);CHKERRQ(ierr);
108     ierr = VecRestoreArrayRead(b_mpi,&bptr);CHKERRQ(ierr);
109   }
110 
111   PetscFunctionReturn(0);
112 }
113 
114 #undef __FUNCT__
115 #define __FUNCT__ "MatMatSolve_STRUMPACK"
116 PetscErrorCode MatMatSolve_STRUMPACK(Mat A,Mat B_mpi,Mat X)
117 {
118   PetscErrorCode   ierr;
119   PetscBool        flg;
120 
121   PetscFunctionBegin;
122   ierr = PetscObjectTypeCompareAny((PetscObject)B_mpi,&flg,MATSEQDENSE,MATMPIDENSE,NULL);CHKERRQ(ierr);
123   if (!flg) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_ARG_WRONG,"Matrix B must be MATDENSE matrix");
124   ierr = PetscObjectTypeCompareAny((PetscObject)X,&flg,MATSEQDENSE,MATMPIDENSE,NULL);CHKERRQ(ierr);
125   if (!flg) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_ARG_WRONG,"Matrix X must be MATDENSE matrix");
126   SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"MatMatSolve_STRUMPACK() is not implemented yet");
127   PetscFunctionReturn(0);
128 }
129 
130 
131 #undef __FUNCT__
132 #define __FUNCT__ "MatLUFactorNumeric_STRUMPACK"
133 PetscErrorCode MatLUFactorNumeric_STRUMPACK(Mat F,Mat A,const MatFactorInfo *info)
134 {
135   STRUMPACK_data        *sp = (STRUMPACK_data*)F->spptr;
136   STRUMPACK_RETURN_CODE sp_err;
137   Mat                   *tseq,A_seq = NULL;
138   Mat_SeqAIJ            *A_d,*A_o;
139   Mat_MPIAIJ            *mat;
140   PetscErrorCode        ierr;
141   PetscInt              M=A->rmap->N,m=A->rmap->n,N=A->cmap->N;
142   PetscMPIInt           size;
143   IS                    isrow;
144   PetscBool             flg;
145 
146   PetscFunctionBegin;
147   ierr = MPI_Comm_size(PetscObjectComm((PetscObject)A),&size);CHKERRQ(ierr);
148 
149   ierr = PetscObjectTypeCompare((PetscObject)A,MATMPIAIJ,&flg);CHKERRQ(ierr);
150   if (flg) { /* A is MATMPIAIJ */
151     if (sp->MatInputMode == REPLICATED) {
152       if (size > 1) { /* convert mpi A to seq mat A */
153         ierr = ISCreateStride(PETSC_COMM_SELF,M,0,1,&isrow);CHKERRQ(ierr);
154         ierr = MatGetSubMatrices(A,1,&isrow,&isrow,MAT_INITIAL_MATRIX,&tseq);CHKERRQ(ierr);
155         ierr = ISDestroy(&isrow);CHKERRQ(ierr);
156         A_seq = *tseq;
157         ierr  = PetscFree(tseq);CHKERRQ(ierr);
158         A_d   = (Mat_SeqAIJ*)A_seq->data;
159       } else { /* size == 1 */
160         mat = (Mat_MPIAIJ*)A->data;
161         A_d = (Mat_SeqAIJ*)(mat->A)->data;
162       }
163       PetscStackCall("STRUMPACK_set_csr_matrix",STRUMPACK_set_csr_matrix(sp->S,&N,A_d->i,A_d->j,A_d->a,0));
164     } else { /* sp->MatInputMode == DISTRIBUTED */
165       mat = (Mat_MPIAIJ*)A->data;
166       A_d = (Mat_SeqAIJ*)(mat->A)->data;
167       A_o = (Mat_SeqAIJ*)(mat->B)->data;
168       PetscStackCall("STRUMPACK_set_MPIAIJ_matrix",STRUMPACK_set_MPIAIJ_matrix(sp->S,&m,A_d->i,A_d->j,A_d->a,A_o->i,A_o->j,A_o->a,mat->garray));
169     }
170   } else { /* A is MATSEQAIJ */
171     A_d = (Mat_SeqAIJ*)A->data;
172     PetscStackCall("STRUMPACK_set_csr_matrix",STRUMPACK_set_csr_matrix(sp->S,&N,A_d->i,A_d->j,A_d->a,0));
173   }
174 
175   /* Reorder and Factor the matrix. */
176   /* TODO figure out how to avoid reorder if the matrix values changed, but the pattern remains the same. */
177   PetscStackCall("STRUMPACK_reorder",sp_err = STRUMPACK_reorder(sp->S));
178   PetscStackCall("STRUMPACK_factor",sp_err = STRUMPACK_factor(sp->S));
179   if (sp_err != STRUMPACK_SUCCESS) {
180     if (sp_err == STRUMPACK_MATRIX_NOT_SET)        SETERRQ(PETSC_COMM_SELF,PETSC_ERR_LIB,"STRUMPACK error: matrix was not set");
181     else if (sp_err == STRUMPACK_REORDERING_ERROR) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_LIB,"STRUMPACK error: matrix reordering failed");
182     else                                           SETERRQ(PETSC_COMM_SELF,PETSC_ERR_LIB,"STRUMPACK error: factorization failed");
183   }
184   if (flg && sp->MatInputMode == REPLICATED && size > 1) {
185     ierr = MatDestroy(&A_seq);CHKERRQ(ierr);
186   }
187   PetscFunctionReturn(0);
188 }
189 
190 #undef __FUNCT__
191 #define __FUNCT__ "MatLUFactorSymbolic_STRUMPACK"
192 PetscErrorCode MatLUFactorSymbolic_STRUMPACK(Mat F,Mat A,IS r,IS c,const MatFactorInfo *info)
193 {
194   PetscFunctionBegin;
195   F->ops->lufactornumeric = MatLUFactorNumeric_STRUMPACK;
196   F->ops->solve           = MatSolve_STRUMPACK;
197   F->ops->matsolve        = MatMatSolve_STRUMPACK;
198   PetscFunctionReturn(0);
199 }
200 
201 #undef __FUNCT__
202 #define __FUNCT__ "MatFactorGetSolverPackage_aij_strumpack"
203 PetscErrorCode MatFactorGetSolverPackage_aij_strumpack(Mat A,const MatSolverPackage *type)
204 {
205   PetscFunctionBegin;
206   *type = MATSOLVERSTRUMPACK;
207   PetscFunctionReturn(0);
208 }
209 
210 #undef __FUNCT__
211 #define __FUNCT__ "MatGetFactor_aij_strumpack"
212 PETSC_EXTERN PetscErrorCode MatGetFactor_aij_strumpack(Mat A,MatFactorType ftype,Mat *F)
213 {
214   Mat                 B;
215   STRUMPACK_data      *sp;
216   PetscErrorCode      ierr;
217   PetscInt            M=A->rmap->N,N=A->cmap->N;
218   PetscMPIInt         size;
219   STRUMPACK_INTERFACE iface;
220   PetscBool           verb,flg;
221   int                 argc;
222   char                **args;
223   const STRUMPACK_PRECISION table[2][2][2] = {{{STRUMPACK_FLOATCOMPLEX_64,STRUMPACK_DOUBLECOMPLEX_64},
224                                                {STRUMPACK_FLOAT_64,STRUMPACK_DOUBLE_64}},
225                                               {{STRUMPACK_FLOATCOMPLEX,STRUMPACK_DOUBLECOMPLEX},
226                                                {STRUMPACK_FLOAT,STRUMPACK_DOUBLE}}};
227   const STRUMPACK_PRECISION prec = table[(sizeof(PetscInt)==8)?0:1][(PETSC_SCALAR==PETSC_COMPLEX)?0:1][(PETSC_REAL==PETSC_FLOAT)?0:1];
228 
229   PetscFunctionBegin;
230   ierr = MPI_Comm_size(PetscObjectComm((PetscObject)A),&size);CHKERRQ(ierr);
231   /* Create the factorization matrix */
232   ierr = MatCreate(PetscObjectComm((PetscObject)A),&B);CHKERRQ(ierr);
233   ierr = MatSetSizes(B,A->rmap->n,A->cmap->n,M,N);CHKERRQ(ierr);
234   ierr = MatSetType(B,((PetscObject)A)->type_name);CHKERRQ(ierr);
235   ierr = MatSeqAIJSetPreallocation(B,0,NULL);
236   ierr = MatMPIAIJSetPreallocation(B,0,NULL,0,NULL);CHKERRQ(ierr);
237   B->ops->lufactorsymbolic = MatLUFactorSymbolic_STRUMPACK;
238   B->ops->view             = MatView_STRUMPACK;
239   B->ops->destroy          = MatDestroy_STRUMPACK;
240   B->ops->getdiagonal      = MatGetDiagonal_STRUMPACK;
241   ierr = PetscObjectComposeFunction((PetscObject)B,"MatFactorGetSolverPackage_C",MatFactorGetSolverPackage_aij_strumpack);CHKERRQ(ierr);
242   B->factortype = MAT_FACTOR_LU;
243   ierr     = PetscNewLog(B,&sp);CHKERRQ(ierr);
244   B->spptr = sp;
245 
246   ierr = PetscOptionsBegin(PetscObjectComm((PetscObject)A),((PetscObject)A)->prefix,"STRUMPACK Options","Mat");CHKERRQ(ierr);
247   sp->MatInputMode = DISTRIBUTED;
248   iface = STRUMPACK_MPI_DIST;
249   ierr = PetscOptionsEnum("-mat_strumpack_matinput","Matrix input mode (global or distributed)","None",STRUMPACK_MatInputModes,
250                           (PetscEnum)sp->MatInputMode,(PetscEnum*)&sp->MatInputMode,NULL);CHKERRQ(ierr);
251   if (sp->MatInputMode == DISTRIBUTED && size == 1) sp->MatInputMode = REPLICATED;
252   if (sp->MatInputMode == DISTRIBUTED)     iface = STRUMPACK_MPI_DIST;
253   else if (sp->MatInputMode == REPLICATED) iface = STRUMPACK_MPI;
254 
255   ierr = PetscObjectTypeCompare((PetscObject)A,MATSEQAIJ,&flg);
256   if (flg) iface = STRUMPACK_MT;
257 
258   if (PetscLogPrintInfo) verb = PETSC_TRUE;
259   else verb = PETSC_FALSE;
260   ierr = PetscOptionsBool("-mat_strumpack_verbose","Print STRUMPACK information","None",verb,&verb,NULL);CHKERRQ(ierr);
261   PetscOptionsEnd();
262 
263   ierr = PetscGetArgs(&argc,&args);CHKERRQ(ierr);
264   PetscStackCall("STRUMPACK_init",STRUMPACK_init(&(sp->S),PetscObjectComm((PetscObject)A),prec,iface,argc,args,verb));
265   PetscStackCall("STRUMPACK_set_from_options",STRUMPACK_set_from_options(sp->S));
266 
267   *F = B;
268   PetscFunctionReturn(0);
269 }
270 
271 #undef __FUNCT__
272 #define __FUNCT__ "MatSolverPackageRegister_STRUMPACK"
273 PETSC_EXTERN PetscErrorCode MatSolverPackageRegister_STRUMPACK(void)
274 {
275   PetscErrorCode ierr;
276 
277   PetscFunctionBegin;
278   ierr = MatSolverPackageRegister(MATSOLVERSTRUMPACK,MATMPIAIJ,MAT_FACTOR_LU,MatGetFactor_aij_strumpack);CHKERRQ(ierr);
279   ierr = MatSolverPackageRegister(MATSOLVERSTRUMPACK,MATSEQAIJ,MAT_FACTOR_LU,MatGetFactor_aij_strumpack);CHKERRQ(ierr);
280   PetscFunctionReturn(0);
281 }
282 
283 #undef __FUNCT__
284 #define __FUNCT__ "MatFactorInfo_STRUMPACK"
285 PetscErrorCode MatFactorInfo_STRUMPACK(Mat A,PetscViewer viewer)
286 {
287   PetscErrorCode  ierr;
288 
289   PetscFunctionBegin;
290   /* check if matrix is strumpack type */
291   if (A->ops->solve != MatSolve_STRUMPACK) PetscFunctionReturn(0);
292   ierr = PetscViewerASCIIPrintf(viewer,"STRUMPACK sparse solver!\n");CHKERRQ(ierr);
293   PetscFunctionReturn(0);
294 }
295 
296 #undef __FUNCT__
297 #define __FUNCT__ "MatView_STRUMPACK"
298 PetscErrorCode MatView_STRUMPACK(Mat A,PetscViewer viewer)
299 {
300   PetscErrorCode    ierr;
301   PetscBool         iascii;
302   PetscViewerFormat format;
303 
304   PetscFunctionBegin;
305   ierr = PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);CHKERRQ(ierr);
306   if (iascii) {
307     ierr = PetscViewerGetFormat(viewer,&format);CHKERRQ(ierr);
308     if (format == PETSC_VIEWER_ASCII_INFO) {
309       ierr = MatFactorInfo_STRUMPACK(A,viewer);CHKERRQ(ierr);
310     }
311   }
312   PetscFunctionReturn(0);
313 }
314 
315 
316 /*MC
317   MATSOLVERSSTRUMPACK - Parallel direct solver package for LU factorization
318 
319   Use ./configure --download-strumpack to have PETSc installed with STRUMPACK
320 
321   Use -pc_type lu -pc_factor_mat_solver_package strumpack to us this direct solver
322 
323    Works with AIJ matrices
324 
325 .seealso: PCLU
326 
327 .seealso: PCFactorSetMatSolverPackage(), MatSolverPackage
328 
329 M*/
330 
331