xref: /petsc/src/mat/impls/aij/seq/matrart.c (revision 4e8afd12df985612b40e26aef947da2f566aee4e)
1 
2 /*
3   Defines projective product routines where A is a SeqAIJ matrix
4           C = R * A * R^T
5 */
6 
7 #include <../src/mat/impls/aij/seq/aij.h>
8 #include <../src/mat/utils/freespace.h>
9 #include <../src/mat/impls/dense/seq/dense.h> /*I "petscmat.h" I*/
10 
11 PetscErrorCode MatDestroy_SeqAIJ_RARt(Mat A)
12 {
13   PetscErrorCode ierr;
14   Mat_SeqAIJ     *a    = (Mat_SeqAIJ*)A->data;
15   Mat_RARt       *rart = a->rart;
16 
17   PetscFunctionBegin;
18   ierr = MatTransposeColoringDestroy(&rart->matcoloring);CHKERRQ(ierr);
19   ierr = MatDestroy(&rart->Rt);CHKERRQ(ierr);
20   ierr = MatDestroy(&rart->RARt);CHKERRQ(ierr);
21   ierr = MatDestroy(&rart->ARt);CHKERRQ(ierr);
22   ierr = PetscFree(rart->work);CHKERRQ(ierr);
23 
24   A->ops->destroy = rart->destroy;
25   if (A->ops->destroy) {
26     ierr = (*A->ops->destroy)(A);CHKERRQ(ierr);
27   }
28   ierr = PetscFree(rart);CHKERRQ(ierr);
29   PetscFunctionReturn(0);
30 }
31 
32 PetscErrorCode MatRARtSymbolic_SeqAIJ_SeqAIJ_colorrart(Mat A,Mat R,PetscReal fill,Mat C)
33 {
34   PetscErrorCode       ierr;
35   Mat                  P;
36   PetscInt             *rti,*rtj;
37   Mat_RARt             *rart;
38   MatColoring          coloring;
39   MatTransposeColoring matcoloring;
40   ISColoring           iscoloring;
41   Mat                  Rt_dense,RARt_dense;
42   Mat_SeqAIJ           *c;
43 
44   PetscFunctionBegin;
45   /* create symbolic P=Rt */
46   ierr = MatGetSymbolicTranspose_SeqAIJ(R,&rti,&rtj);CHKERRQ(ierr);
47   ierr = MatCreateSeqAIJWithArrays(PETSC_COMM_SELF,R->cmap->n,R->rmap->n,rti,rtj,NULL,&P);CHKERRQ(ierr);
48 
49   /* get symbolic C=Pt*A*P */
50   ierr = MatPtAPSymbolic_SeqAIJ_SeqAIJ_SparseAxpy(A,P,fill,C);CHKERRQ(ierr);
51   ierr = MatSetBlockSizes(C,PetscAbs(R->rmap->bs),PetscAbs(R->rmap->bs));CHKERRQ(ierr);
52   C->ops->rartnumeric = MatRARtNumeric_SeqAIJ_SeqAIJ_colorrart;
53 
54   /* create a supporting struct */
55   ierr    = PetscNew(&rart);CHKERRQ(ierr);
56   c       = (Mat_SeqAIJ*)C->data;
57   c->rart = rart;
58 
59   /* ------ Use coloring ---------- */
60   /* inode causes memory problem */
61   ierr = MatSetOption(C,MAT_USE_INODES,PETSC_FALSE);CHKERRQ(ierr);
62 
63   /* Create MatTransposeColoring from symbolic C=R*A*R^T */
64   ierr = MatColoringCreate(C,&coloring);CHKERRQ(ierr);
65   ierr = MatColoringSetDistance(coloring,2);CHKERRQ(ierr);
66   ierr = MatColoringSetType(coloring,MATCOLORINGSL);CHKERRQ(ierr);
67   ierr = MatColoringSetFromOptions(coloring);CHKERRQ(ierr);
68   ierr = MatColoringApply(coloring,&iscoloring);CHKERRQ(ierr);
69   ierr = MatColoringDestroy(&coloring);CHKERRQ(ierr);
70   ierr = MatTransposeColoringCreate(C,iscoloring,&matcoloring);CHKERRQ(ierr);
71 
72   rart->matcoloring = matcoloring;
73   ierr = ISColoringDestroy(&iscoloring);CHKERRQ(ierr);
74 
75   /* Create Rt_dense */
76   ierr = MatCreate(PETSC_COMM_SELF,&Rt_dense);CHKERRQ(ierr);
77   ierr = MatSetSizes(Rt_dense,A->cmap->n,matcoloring->ncolors,A->cmap->n,matcoloring->ncolors);CHKERRQ(ierr);
78   ierr = MatSetType(Rt_dense,MATSEQDENSE);CHKERRQ(ierr);
79   ierr = MatSeqDenseSetPreallocation(Rt_dense,NULL);CHKERRQ(ierr);
80 
81   Rt_dense->assembled = PETSC_TRUE;
82   rart->Rt            = Rt_dense;
83 
84   /* Create RARt_dense = R*A*Rt_dense */
85   ierr = MatCreate(PETSC_COMM_SELF,&RARt_dense);CHKERRQ(ierr);
86   ierr = MatSetSizes(RARt_dense,C->rmap->n,matcoloring->ncolors,C->rmap->n,matcoloring->ncolors);CHKERRQ(ierr);
87   ierr = MatSetType(RARt_dense,MATSEQDENSE);CHKERRQ(ierr);
88   ierr = MatSeqDenseSetPreallocation(RARt_dense,NULL);CHKERRQ(ierr);
89 
90   rart->RARt = RARt_dense;
91 
92   /* Allocate work array to store columns of A*R^T used in MatMatMatMultNumeric_SeqAIJ_SeqAIJ_SeqDense() */
93   ierr = PetscMalloc1(A->rmap->n*4,&rart->work);CHKERRQ(ierr);
94 
95   rart->destroy       = C->ops->destroy;
96   C->ops->destroy     = MatDestroy_SeqAIJ_RARt;
97 
98   /* clean up */
99   ierr = MatRestoreSymbolicTranspose_SeqAIJ(R,&rti,&rtj);CHKERRQ(ierr);
100   ierr = MatDestroy(&P);CHKERRQ(ierr);
101 
102 #if defined(PETSC_USE_INFO)
103   {
104     PetscReal density= (PetscReal)(c->nz)/(RARt_dense->rmap->n*RARt_dense->cmap->n);
105     ierr = PetscInfo(C,"C=R*(A*Rt) via coloring C - use sparse-dense inner products\n");CHKERRQ(ierr);
106     ierr = PetscInfo6(C,"RARt_den %D %D; Rt %D %D (RARt->nz %D)/(m*ncolors)=%g\n",RARt_dense->rmap->n,RARt_dense->cmap->n,R->cmap->n,R->rmap->n,c->nz,density);CHKERRQ(ierr);
107   }
108 #endif
109   PetscFunctionReturn(0);
110 }
111 
112 /*
113  RAB = R * A * B, R and A in seqaij format, B in dense format;
114 */
115 PetscErrorCode MatMatMatMultNumeric_SeqAIJ_SeqAIJ_SeqDense(Mat R,Mat A,Mat B,Mat RAB,PetscScalar *work)
116 {
117   Mat_SeqAIJ        *a=(Mat_SeqAIJ*)A->data,*r=(Mat_SeqAIJ*)R->data;
118   PetscErrorCode    ierr;
119   PetscScalar       r1,r2,r3,r4;
120   const PetscScalar *b,*b1,*b2,*b3,*b4;
121   MatScalar         *aa,*ra;
122   PetscInt          cn =B->cmap->n,bm=B->rmap->n,col,i,j,n,*ai=a->i,*aj,am=A->rmap->n;
123   PetscInt          am2=2*am,am3=3*am,bm4=4*bm;
124   PetscScalar       *d,*c,*c2,*c3,*c4;
125   PetscInt          *rj,rm=R->rmap->n,dm=RAB->rmap->n,dn=RAB->cmap->n;
126   PetscInt         rm2=2*rm,rm3=3*rm,colrm;
127 
128   PetscFunctionBegin;
129   if (!dm || !dn) PetscFunctionReturn(0);
130   if (bm != A->cmap->n) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Number columns in A %D not equal rows in B %D\n",A->cmap->n,bm);
131   if (am != R->cmap->n) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Number columns in R %D not equal rows in A %D\n",R->cmap->n,am);
132   if (R->rmap->n != RAB->rmap->n) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Number rows in RAB %D not equal rows in R %D\n",RAB->rmap->n,R->rmap->n);
133   if (B->cmap->n != RAB->cmap->n) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Number columns in RAB %D not equal columns in B %D\n",RAB->cmap->n,B->cmap->n);
134 
135   { /*
136      This approach is not as good as original ones (will be removed later), but it reveals that
137      AB_den=A*B takes almost all execution time in R*A*B for src/ksp/ksp/tutorials/ex56.c
138      */
139     PetscBool via_matmatmult=PETSC_FALSE;
140     ierr = PetscOptionsGetBool(NULL,NULL,"-matrart_via_matmatmult",&via_matmatmult,NULL);CHKERRQ(ierr);
141     if (via_matmatmult) {
142       Mat AB_den = NULL;
143       ierr = MatCreate(PetscObjectComm((PetscObject)A),&AB_den);CHKERRQ(ierr);
144       ierr = MatMatMultSymbolic_SeqAIJ_SeqDense(A,B,0.0,AB_den);CHKERRQ(ierr);
145       ierr = MatMatMultNumeric_SeqAIJ_SeqDense(A,B,AB_den);CHKERRQ(ierr);
146       ierr = MatMatMultNumeric_SeqAIJ_SeqDense(R,AB_den,RAB);CHKERRQ(ierr);
147       ierr = MatDestroy(&AB_den);CHKERRQ(ierr);
148       PetscFunctionReturn(0);
149     }
150   }
151 
152   ierr = MatDenseGetArrayRead(B,&b);CHKERRQ(ierr);
153   ierr = MatDenseGetArray(RAB,&d);CHKERRQ(ierr);
154   b1   = b; b2 = b1 + bm; b3 = b2 + bm; b4 = b3 + bm;
155   c    = work; c2 = c + am; c3 = c2 + am; c4 = c3 + am;
156   for (col=0; col<cn-4; col += 4) {  /* over columns of C */
157     for (i=0; i<am; i++) {        /* over rows of A in those columns */
158       r1 = r2 = r3 = r4 = 0.0;
159       n  = ai[i+1] - ai[i];
160       aj = a->j + ai[i];
161       aa = a->a + ai[i];
162       for (j=0; j<n; j++) {
163         r1 += (*aa)*b1[*aj];
164         r2 += (*aa)*b2[*aj];
165         r3 += (*aa)*b3[*aj];
166         r4 += (*aa++)*b4[*aj++];
167       }
168       c[i]       = r1;
169       c[am  + i] = r2;
170       c[am2 + i] = r3;
171       c[am3 + i] = r4;
172     }
173     b1 += bm4;
174     b2 += bm4;
175     b3 += bm4;
176     b4 += bm4;
177 
178     /* RAB[:,col] = R*C[:,col] */
179     colrm = col*rm;
180     for (i=0; i<rm; i++) {        /* over rows of R in those columns */
181       r1 = r2 = r3 = r4 = 0.0;
182       n  = r->i[i+1] - r->i[i];
183       rj = r->j + r->i[i];
184       ra = r->a + r->i[i];
185       for (j=0; j<n; j++) {
186         r1 += (*ra)*c[*rj];
187         r2 += (*ra)*c2[*rj];
188         r3 += (*ra)*c3[*rj];
189         r4 += (*ra++)*c4[*rj++];
190       }
191       d[colrm + i]       = r1;
192       d[colrm + rm + i]  = r2;
193       d[colrm + rm2 + i] = r3;
194       d[colrm + rm3 + i] = r4;
195     }
196   }
197   for (; col<cn; col++) {     /* over extra columns of C */
198     for (i=0; i<am; i++) {  /* over rows of A in those columns */
199       r1 = 0.0;
200       n  = a->i[i+1] - a->i[i];
201       aj = a->j + a->i[i];
202       aa = a->a + a->i[i];
203       for (j=0; j<n; j++) {
204         r1 += (*aa++)*b1[*aj++];
205       }
206       c[i] = r1;
207     }
208     b1 += bm;
209 
210     for (i=0; i<rm; i++) {  /* over rows of R in those columns */
211       r1 = 0.0;
212       n  = r->i[i+1] - r->i[i];
213       rj = r->j + r->i[i];
214       ra = r->a + r->i[i];
215       for (j=0; j<n; j++) {
216         r1 += (*ra++)*c[*rj++];
217       }
218       d[col*rm + i] = r1;
219     }
220   }
221   ierr = PetscLogFlops(cn*2.0*(a->nz + r->nz));CHKERRQ(ierr);
222 
223   ierr = MatDenseRestoreArrayRead(B,&b);CHKERRQ(ierr);
224   ierr = MatDenseRestoreArray(RAB,&d);CHKERRQ(ierr);
225   ierr = MatAssemblyBegin(RAB,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
226   ierr = MatAssemblyEnd(RAB,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
227   PetscFunctionReturn(0);
228 }
229 
230 PetscErrorCode MatRARtNumeric_SeqAIJ_SeqAIJ_colorrart(Mat A,Mat R,Mat C)
231 {
232   PetscErrorCode       ierr;
233   Mat_SeqAIJ           *c = (Mat_SeqAIJ*)C->data;
234   Mat_RARt             *rart=c->rart;
235   MatTransposeColoring matcoloring;
236   Mat                  Rt,RARt;
237 
238   PetscFunctionBegin;
239   /* Get dense Rt by Apply MatTransposeColoring to R */
240   matcoloring = rart->matcoloring;
241   Rt          = rart->Rt;
242   ierr  = MatTransColoringApplySpToDen(matcoloring,R,Rt);CHKERRQ(ierr);
243 
244   /* Get dense RARt = R*A*Rt -- dominates! */
245   RARt = rart->RARt;
246   ierr = MatMatMatMultNumeric_SeqAIJ_SeqAIJ_SeqDense(R,A,Rt,RARt,rart->work);CHKERRQ(ierr);
247 
248   /* Recover C from C_dense */
249   ierr = MatTransColoringApplyDenToSp(matcoloring,RARt,C);CHKERRQ(ierr);
250   PetscFunctionReturn(0);
251 }
252 
253 PetscErrorCode MatRARtSymbolic_SeqAIJ_SeqAIJ_matmattransposemult(Mat A,Mat R,PetscReal fill,Mat C)
254 {
255   PetscErrorCode ierr;
256   Mat            ARt;
257   Mat_SeqAIJ     *c;
258   Mat_RARt       *rart;
259   Mat_Product    *product = C->product;
260   MatProductAlgorithm alg=product->alg;
261 
262   PetscFunctionBegin;
263   /* create symbolic ARt = A*R^T  */
264   ierr = MatProductCreate(A,R,NULL,&ARt);CHKERRQ(ierr);
265   ierr = MatProductSetType(ARt,MATPRODUCT_ABt);CHKERRQ(ierr);
266   ierr = MatProductSetAlgorithm(ARt,"sorted");CHKERRQ(ierr);
267   ierr = MatProductSetFill(ARt,fill);CHKERRQ(ierr);
268   ierr = MatProductSetFromOptions(ARt);CHKERRQ(ierr);
269   ierr = MatProductSymbolic(ARt);CHKERRQ(ierr);
270 
271   /* compute symbolic C = R*ARt */
272   C->product->alg = "sorted"; /* set algorithm for C = R*ARt */
273   ierr = MatMatMultSymbolic_SeqAIJ_SeqAIJ(R,ARt,fill,C);CHKERRQ(ierr);
274   C->product->alg = alg; /* resume original algorithm for C */
275 
276   C->ops->rartnumeric = MatRARtNumeric_SeqAIJ_SeqAIJ_matmattransposemult;
277 
278   ierr = PetscNew(&rart);CHKERRQ(ierr);
279   c         = (Mat_SeqAIJ*)C->data;
280   c->rart   = rart;
281   rart->ARt = ARt;
282   rart->destroy   = C->ops->destroy;
283   C->ops->destroy = MatDestroy_SeqAIJ_RARt;
284 #if defined(PETSC_USE_INFO)
285   ierr = PetscInfo(C,"Use ARt=A*R^T, C=R*ARt via MatMatTransposeMult(). Coloring can be applied to A*R^T.\n");CHKERRQ(ierr);
286 #endif
287   PetscFunctionReturn(0);
288 }
289 
290 PetscErrorCode MatRARtNumeric_SeqAIJ_SeqAIJ_matmattransposemult(Mat A,Mat R,Mat C)
291 {
292   PetscErrorCode  ierr;
293   Mat_SeqAIJ      *c=(Mat_SeqAIJ*)C->data;
294   Mat_RARt        *rart=c->rart;
295   Mat             ARt=rart->ARt;
296 
297   PetscFunctionBegin;
298   ierr = MatMatTransposeMultNumeric_SeqAIJ_SeqAIJ(A,R,ARt);CHKERRQ(ierr); /* dominate! */
299   ierr = MatMatMultNumeric_SeqAIJ_SeqAIJ(R,ARt,C);CHKERRQ(ierr);
300   PetscFunctionReturn(0);
301 }
302 
303 PetscErrorCode MatRARtSymbolic_SeqAIJ_SeqAIJ(Mat A,Mat R,PetscReal fill,Mat C)
304 {
305   PetscErrorCode ierr;
306   Mat            Rt;
307   Mat_SeqAIJ     *c;
308   Mat_RARt       *rart;
309 
310   PetscFunctionBegin;
311   ierr = MatTranspose_SeqAIJ(R,MAT_INITIAL_MATRIX,&Rt);CHKERRQ(ierr);
312   ierr = MatMatMatMultSymbolic_SeqAIJ_SeqAIJ_SeqAIJ(R,A,Rt,fill,C);CHKERRQ(ierr);
313 
314   ierr = PetscNew(&rart);CHKERRQ(ierr);
315   rart->Rt = Rt;
316   c        = (Mat_SeqAIJ*)C->data;
317   c->rart  = rart;
318   rart->destroy       = C->ops->destroy;
319   C->ops->destroy     = MatDestroy_SeqAIJ_RARt;
320   C->ops->rartnumeric = MatRARtNumeric_SeqAIJ_SeqAIJ;
321 #if defined(PETSC_USE_INFO)
322   ierr = PetscInfo(C,"Use Rt=R^T and C=R*A*Rt via MatMatMatMult() to avoid sparse inner products\n");CHKERRQ(ierr);
323 #endif
324   PetscFunctionReturn(0);
325 }
326 
327 PetscErrorCode MatRARtNumeric_SeqAIJ_SeqAIJ(Mat A,Mat R,Mat C)
328 {
329   PetscErrorCode  ierr;
330   Mat_SeqAIJ      *c = (Mat_SeqAIJ*)C->data;
331   Mat_RARt        *rart = c->rart;
332   Mat             Rt = rart->Rt;
333 
334   PetscFunctionBegin;
335   ierr = MatTranspose_SeqAIJ(R,MAT_REUSE_MATRIX,&Rt);CHKERRQ(ierr);
336   ierr = MatMatMatMultNumeric_SeqAIJ_SeqAIJ_SeqAIJ(R,A,Rt,C);CHKERRQ(ierr);
337   PetscFunctionReturn(0);
338 }
339 
340 PetscErrorCode MatRARt_SeqAIJ_SeqAIJ(Mat A,Mat R,MatReuse scall,PetscReal fill,Mat *C)
341 {
342   PetscErrorCode ierr;
343   const char     *algTypes[3] = {"matmatmatmult","matmattransposemult","coloring_rart"};
344   PetscInt       alg=0; /* set default algorithm */
345 
346   PetscFunctionBegin;
347   if (scall == MAT_INITIAL_MATRIX) {
348     ierr = PetscOptionsBegin(PetscObjectComm((PetscObject)A),((PetscObject)A)->prefix,"MatRARt","Mat");CHKERRQ(ierr);
349     ierr = PetscOptionsEList("-matrart_via","Algorithmic approach","MatRARt",algTypes,3,algTypes[0],&alg,NULL);CHKERRQ(ierr);
350     ierr = PetscOptionsEnd();CHKERRQ(ierr);
351 
352     ierr = PetscLogEventBegin(MAT_RARtSymbolic,A,R,0,0);CHKERRQ(ierr);
353     ierr = MatCreate(PETSC_COMM_SELF,C);CHKERRQ(ierr);
354     switch (alg) {
355     case 1:
356       /* via matmattransposemult: ARt=A*R^T, C=R*ARt - matrix coloring can be applied to A*R^T */
357       ierr = MatRARtSymbolic_SeqAIJ_SeqAIJ_matmattransposemult(A,R,fill,*C);CHKERRQ(ierr);
358       break;
359     case 2:
360       /* via coloring_rart: apply coloring C = R*A*R^T                          */
361       ierr = MatRARtSymbolic_SeqAIJ_SeqAIJ_colorrart(A,R,fill,*C);CHKERRQ(ierr);
362       break;
363     default:
364       /* via matmatmatmult: Rt=R^T, C=R*A*Rt - avoid inefficient sparse inner products */
365       ierr = MatRARtSymbolic_SeqAIJ_SeqAIJ(A,R,fill,*C);CHKERRQ(ierr);
366       break;
367     }
368     ierr = PetscLogEventEnd(MAT_RARtSymbolic,A,R,0,0);CHKERRQ(ierr);
369   }
370 
371   ierr = PetscLogEventBegin(MAT_RARtNumeric,A,R,0,0);CHKERRQ(ierr);
372   ierr = ((*C)->ops->rartnumeric)(A,R,*C);CHKERRQ(ierr);
373   ierr = PetscLogEventEnd(MAT_RARtNumeric,A,R,0,0);CHKERRQ(ierr);
374   PetscFunctionReturn(0);
375 }
376 
377 /* ------------------------------------------------------------- */
378 PetscErrorCode MatProductSymbolic_RARt_SeqAIJ_SeqAIJ(Mat C)
379 {
380   PetscErrorCode      ierr;
381   Mat_Product         *product = C->product;
382   Mat                 A=product->A,R=product->B;
383   MatProductAlgorithm alg=product->alg;
384   PetscReal           fill=product->fill;
385   PetscBool           flg;
386 
387   PetscFunctionBegin;
388   ierr = PetscStrcmp(alg,"r*a*rt",&flg);CHKERRQ(ierr);
389   if (flg) {
390     ierr = MatRARtSymbolic_SeqAIJ_SeqAIJ(A,R,fill,C);CHKERRQ(ierr);
391     goto next;
392   }
393 
394   ierr = PetscStrcmp(alg,"r*art",&flg);CHKERRQ(ierr);
395   if (flg) {
396     ierr = MatRARtSymbolic_SeqAIJ_SeqAIJ_matmattransposemult(A,R,fill,C);CHKERRQ(ierr);
397     goto next;
398   }
399 
400   ierr = PetscStrcmp(alg,"coloring_rart",&flg);CHKERRQ(ierr);
401   if (flg) {
402     ierr = MatRARtSymbolic_SeqAIJ_SeqAIJ_colorrart(A,R,fill,C);CHKERRQ(ierr);
403     goto next;
404   }
405 
406   SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"MatProductAlgorithm is not supported");
407 
408 next:
409   C->ops->productnumeric = MatProductNumeric_RARt;
410   PetscFunctionReturn(0);
411 }
412