xref: /petsc/src/mat/impls/aij/seq/matmatmult.c (revision f3fe499b4cc4d64bf04aa4f5e4963dcc4eb56541)
1 #define PETSCMAT_DLL
2 
3 /*
4   Defines matrix-matrix product routines for pairs of SeqAIJ matrices
5           C = A * B
6 */
7 
8 #include "../src/mat/impls/aij/seq/aij.h" /*I "petscmat.h" I*/
9 #include "../src/mat/utils/freespace.h"
10 #include "petscbt.h"
11 #include "../src/mat/impls/dense/seq/dense.h" /*I "petscmat.h" I*/
12 
13 EXTERN_C_BEGIN
14 #undef __FUNCT__
15 #define __FUNCT__ "MatMatMult_SeqAIJ_SeqAIJ"
16 PetscErrorCode MatMatMult_SeqAIJ_SeqAIJ(Mat A,Mat B,MatReuse scall,PetscReal fill,Mat *C)
17 {
18   PetscErrorCode ierr;
19 
20   PetscFunctionBegin;
21   if (scall == MAT_INITIAL_MATRIX){
22     ierr = MatMatMultSymbolic_SeqAIJ_SeqAIJ(A,B,fill,C);CHKERRQ(ierr);
23   }
24   ierr = MatMatMultNumeric_SeqAIJ_SeqAIJ(A,B,*C);CHKERRQ(ierr);
25   PetscFunctionReturn(0);
26 }
27 EXTERN_C_END
28 
29 #undef __FUNCT__
30 #define __FUNCT__ "MatMatMultSymbolic_SeqAIJ_SeqAIJ"
31 PetscErrorCode MatMatMultSymbolic_SeqAIJ_SeqAIJ(Mat A,Mat B,PetscReal fill,Mat *C)
32 {
33   PetscErrorCode     ierr;
34   PetscFreeSpaceList free_space=PETSC_NULL,current_space=PETSC_NULL;
35   Mat_SeqAIJ         *a=(Mat_SeqAIJ*)A->data,*b=(Mat_SeqAIJ*)B->data,*c;
36   PetscInt           *ai=a->i,*aj=a->j,*bi=b->i,*bj=b->j,*bjj,*ci,*cj;
37   PetscInt           am=A->rmap->N,bn=B->cmap->N,bm=B->rmap->N;
38   PetscInt           i,j,anzi,brow,bnzj,cnzi,nlnk,*lnk,nspacedouble=0;
39   MatScalar          *ca;
40   PetscBT            lnkbt;
41 
42   PetscFunctionBegin;
43   /* Set up */
44   /* Allocate ci array, arrays for fill computation and */
45   /* free space for accumulating nonzero column info */
46   ierr = PetscMalloc(((am+1)+1)*sizeof(PetscInt),&ci);CHKERRQ(ierr);
47   ci[0] = 0;
48 
49   /* create and initialize a linked list */
50   nlnk = bn+1;
51   ierr = PetscLLCreate(bn,bn,nlnk,lnk,lnkbt);CHKERRQ(ierr);
52 
53   /* Initial FreeSpace size is fill*(nnz(A)+nnz(B)) */
54   ierr = PetscFreeSpaceGet((PetscInt)(fill*(ai[am]+bi[bm])),&free_space);CHKERRQ(ierr);
55   current_space = free_space;
56 
57   /* Determine symbolic info for each row of the product: */
58   for (i=0;i<am;i++) {
59     anzi = ai[i+1] - ai[i];
60     cnzi = 0;
61     j    = anzi;
62     aj   = a->j + ai[i];
63     while (j){/* assume cols are almost in increasing order, starting from its end saves computation */
64       j--;
65       brow = *(aj + j);
66       bnzj = bi[brow+1] - bi[brow];
67       bjj  = bj + bi[brow];
68       /* add non-zero cols of B into the sorted linked list lnk */
69       ierr = PetscLLAdd(bnzj,bjj,bn,nlnk,lnk,lnkbt);CHKERRQ(ierr);
70       cnzi += nlnk;
71     }
72 
73     /* If free space is not available, make more free space */
74     /* Double the amount of total space in the list */
75     if (current_space->local_remaining<cnzi) {
76       ierr = PetscFreeSpaceGet(cnzi+current_space->total_array_size,&current_space);CHKERRQ(ierr);
77       nspacedouble++;
78     }
79 
80     /* Copy data into free space, then initialize lnk */
81     ierr = PetscLLClean(bn,bn,cnzi,lnk,current_space->array,lnkbt);CHKERRQ(ierr);
82     current_space->array           += cnzi;
83     current_space->local_used      += cnzi;
84     current_space->local_remaining -= cnzi;
85 
86     ci[i+1] = ci[i] + cnzi;
87   }
88 
89   /* Column indices are in the list of free space */
90   /* Allocate space for cj, initialize cj, and */
91   /* destroy list of free space and other temporary array(s) */
92   ierr = PetscMalloc((ci[am]+1)*sizeof(PetscInt),&cj);CHKERRQ(ierr);
93   ierr = PetscFreeSpaceContiguous(&free_space,cj);CHKERRQ(ierr);
94   ierr = PetscLLDestroy(lnk,lnkbt);CHKERRQ(ierr);
95 
96   /* Allocate space for ca */
97   ierr = PetscMalloc((ci[am]+1)*sizeof(MatScalar),&ca);CHKERRQ(ierr);
98   ierr = PetscMemzero(ca,(ci[am]+1)*sizeof(MatScalar));CHKERRQ(ierr);
99 
100   /* put together the new symbolic matrix */
101   ierr = MatCreateSeqAIJWithArrays(((PetscObject)A)->comm,am,bn,ci,cj,ca,C);CHKERRQ(ierr);
102 
103   /* MatCreateSeqAIJWithArrays flags matrix so PETSc doesn't free the user's arrays. */
104   /* These are PETSc arrays, so change flags so arrays can be deleted by PETSc */
105   c = (Mat_SeqAIJ *)((*C)->data);
106   c->free_a   = PETSC_TRUE;
107   c->free_ij  = PETSC_TRUE;
108   c->nonew    = 0;
109 
110 #if defined(PETSC_USE_INFO)
111   if (ci[am] != 0) {
112     PetscReal afill = ((PetscReal)ci[am])/(ai[am]+bi[bm]);
113     if (afill < 1.0) afill = 1.0;
114     ierr = PetscInfo3((*C),"Reallocs %D; Fill ratio: given %G needed %G.\n",nspacedouble,fill,afill);CHKERRQ(ierr);
115     ierr = PetscInfo1((*C),"Use MatMatMult(A,B,MatReuse,%G,&C) for best performance.;\n",afill);CHKERRQ(ierr);
116   } else {
117     ierr = PetscInfo((*C),"Empty matrix product\n");CHKERRQ(ierr);
118   }
119 #endif
120   PetscFunctionReturn(0);
121 }
122 
123 
124 #undef __FUNCT__
125 #define __FUNCT__ "MatMatMultNumeric_SeqAIJ_SeqAIJ"
126 PetscErrorCode MatMatMultNumeric_SeqAIJ_SeqAIJ(Mat A,Mat B,Mat C)
127 {
128   PetscErrorCode ierr;
129   PetscLogDouble flops=0.0;
130   Mat_SeqAIJ     *a = (Mat_SeqAIJ *)A->data;
131   Mat_SeqAIJ     *b = (Mat_SeqAIJ *)B->data;
132   Mat_SeqAIJ     *c = (Mat_SeqAIJ *)C->data;
133   PetscInt       *ai=a->i,*aj=a->j,*bi=b->i,*bj=b->j,*bjj,*ci=c->i,*cj=c->j;
134   PetscInt       am=A->rmap->N,cm=C->rmap->N;
135   PetscInt       i,j,k,anzi,bnzi,cnzi,brow,nextb;
136   MatScalar      *aa=a->a,*ba=b->a,*baj,*ca=c->a;
137 
138   PetscFunctionBegin;
139   /* clean old values in C */
140   ierr = PetscMemzero(ca,ci[cm]*sizeof(MatScalar));CHKERRQ(ierr);
141   /* Traverse A row-wise. */
142   /* Build the ith row in C by summing over nonzero columns in A, */
143   /* the rows of B corresponding to nonzeros of A. */
144   for (i=0;i<am;i++) {
145     anzi = ai[i+1] - ai[i];
146     for (j=0;j<anzi;j++) {
147       brow = *aj++;
148       bnzi = bi[brow+1] - bi[brow];
149       bjj  = bj + bi[brow];
150       baj  = ba + bi[brow];
151       nextb = 0;
152       for (k=0; nextb<bnzi; k++) {
153         if (cj[k] == bjj[nextb]){ /* ccol == bcol */
154           ca[k] += (*aa)*baj[nextb++];
155         }
156       }
157       flops += 2*bnzi;
158       aa++;
159     }
160     cnzi = ci[i+1] - ci[i];
161     ca += cnzi;
162     cj += cnzi;
163   }
164   ierr = MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
165   ierr = MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
166 
167   ierr = PetscLogFlops(flops);CHKERRQ(ierr);
168   PetscFunctionReturn(0);
169 }
170 
171 
172 #undef __FUNCT__
173 #define __FUNCT__ "MatMatMultTranspose_SeqAIJ_SeqAIJ"
174 PetscErrorCode MatMatMultTranspose_SeqAIJ_SeqAIJ(Mat A,Mat B,MatReuse scall,PetscReal fill,Mat *C) {
175   PetscErrorCode ierr;
176 
177   PetscFunctionBegin;
178   if (scall == MAT_INITIAL_MATRIX){
179     ierr = MatMatMultTransposeSymbolic_SeqAIJ_SeqAIJ(A,B,fill,C);CHKERRQ(ierr);
180   }
181   ierr = MatMatMultTransposeNumeric_SeqAIJ_SeqAIJ(A,B,*C);CHKERRQ(ierr);
182   PetscFunctionReturn(0);
183 }
184 
185 #undef __FUNCT__
186 #define __FUNCT__ "MatMatMultTransposeSymbolic_SeqAIJ_SeqAIJ"
187 PetscErrorCode MatMatMultTransposeSymbolic_SeqAIJ_SeqAIJ(Mat A,Mat B,PetscReal fill,Mat *C)
188 {
189   PetscErrorCode ierr;
190   Mat            At;
191   PetscInt       *ati,*atj;
192 
193   PetscFunctionBegin;
194   /* create symbolic At */
195   ierr = MatGetSymbolicTranspose_SeqAIJ(A,&ati,&atj);CHKERRQ(ierr);
196   ierr = MatCreateSeqAIJWithArrays(PETSC_COMM_SELF,A->cmap->n,A->rmap->n,ati,atj,PETSC_NULL,&At);CHKERRQ(ierr);
197 
198   /* get symbolic C=At*B */
199   ierr = MatMatMultSymbolic_SeqAIJ_SeqAIJ(At,B,fill,C);CHKERRQ(ierr);
200 
201   /* clean up */
202   ierr = MatDestroy(At);CHKERRQ(ierr);
203   ierr = MatRestoreSymbolicTranspose_SeqAIJ(A,&ati,&atj);CHKERRQ(ierr);
204 
205   PetscFunctionReturn(0);
206 }
207 
208 #undef __FUNCT__
209 #define __FUNCT__ "MatMatMultTransposeNumeric_SeqAIJ_SeqAIJ"
210 PetscErrorCode MatMatMultTransposeNumeric_SeqAIJ_SeqAIJ(Mat A,Mat B,Mat C)
211 {
212   PetscErrorCode ierr;
213   Mat_SeqAIJ     *a=(Mat_SeqAIJ*)A->data,*b=(Mat_SeqAIJ*)B->data,*c=(Mat_SeqAIJ*)C->data;
214   PetscInt       am=A->rmap->n,anzi,*ai=a->i,*aj=a->j,*bi=b->i,*bj,bnzi,nextb;
215   PetscInt       cm=C->rmap->n,*ci=c->i,*cj=c->j,crow,*cjj,i,j,k;
216   PetscLogDouble flops=0.0;
217   MatScalar      *aa=a->a,*ba,*ca=c->a,*caj;
218 
219   PetscFunctionBegin;
220   /* clear old values in C */
221   ierr = PetscMemzero(ca,ci[cm]*sizeof(MatScalar));CHKERRQ(ierr);
222 
223   /* compute A^T*B using outer product (A^T)[:,i]*B[i,:] */
224   for (i=0;i<am;i++) {
225     bj   = b->j + bi[i];
226     ba   = b->a + bi[i];
227     bnzi = bi[i+1] - bi[i];
228     anzi = ai[i+1] - ai[i];
229     for (j=0; j<anzi; j++) {
230       nextb = 0;
231       crow  = *aj++;
232       cjj   = cj + ci[crow];
233       caj   = ca + ci[crow];
234       /* perform sparse axpy operation.  Note cjj includes bj. */
235       for (k=0; nextb<bnzi; k++) {
236         if (cjj[k] == *(bj+nextb)) { /* ccol == bcol */
237           caj[k] += (*aa)*(*(ba+nextb));
238           nextb++;
239         }
240       }
241       flops += 2*bnzi;
242       aa++;
243     }
244   }
245 
246   /* Assemble the final matrix and clean up */
247   ierr = MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
248   ierr = MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
249   ierr = PetscLogFlops(flops);CHKERRQ(ierr);
250   PetscFunctionReturn(0);
251 }
252 
253 EXTERN_C_BEGIN
254 #undef __FUNCT__
255 #define __FUNCT__ "MatMatMult_SeqAIJ_SeqDense"
256 PetscErrorCode MatMatMult_SeqAIJ_SeqDense(Mat A,Mat B,MatReuse scall,PetscReal fill,Mat *C)
257 {
258   PetscErrorCode ierr;
259 
260   PetscFunctionBegin;
261   if (scall == MAT_INITIAL_MATRIX){
262     ierr = MatMatMultSymbolic_SeqAIJ_SeqDense(A,B,fill,C);CHKERRQ(ierr);
263   }
264   ierr = MatMatMultNumeric_SeqAIJ_SeqDense(A,B,*C);CHKERRQ(ierr);
265   PetscFunctionReturn(0);
266 }
267 EXTERN_C_END
268 
269 #undef __FUNCT__
270 #define __FUNCT__ "MatMatMultSymbolic_SeqAIJ_SeqDense"
271 PetscErrorCode MatMatMultSymbolic_SeqAIJ_SeqDense(Mat A,Mat B,PetscReal fill,Mat *C)
272 {
273   PetscErrorCode ierr;
274 
275   PetscFunctionBegin;
276   ierr = MatMatMultSymbolic_SeqDense_SeqDense(A,B,0.0,C);CHKERRQ(ierr);
277   PetscFunctionReturn(0);
278 }
279 
280 #undef __FUNCT__
281 #define __FUNCT__ "MatMatMultNumeric_SeqAIJ_SeqDense"
282 PetscErrorCode MatMatMultNumeric_SeqAIJ_SeqDense(Mat A,Mat B,Mat C)
283 {
284   Mat_SeqAIJ     *a=(Mat_SeqAIJ*)A->data;
285   PetscErrorCode ierr;
286   PetscScalar    *b,*c,r1,r2,r3,r4,*b1,*b2,*b3,*b4;
287   MatScalar      *aa;
288   PetscInt       cm=C->rmap->n, cn=B->cmap->n, bm=B->rmap->n, col, i,j,n,*aj, am = A->rmap->n;
289   PetscInt       am2 = 2*am, am3 = 3*am,  bm4 = 4*bm,colam;
290 
291   PetscFunctionBegin;
292   if (!cm || !cn) PetscFunctionReturn(0);
293   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);
294   if (A->rmap->n != C->rmap->n) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Number rows in C %D not equal rows in A %D\n",C->rmap->n,A->rmap->n);
295   if (B->cmap->n != C->cmap->n) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Number columns in B %D not equal columns in C %D\n",B->cmap->n,C->cmap->n);
296   ierr = MatGetArray(B,&b);CHKERRQ(ierr);
297   ierr = MatGetArray(C,&c);CHKERRQ(ierr);
298   b1 = b; b2 = b1 + bm; b3 = b2 + bm; b4 = b3 + bm;
299   for (col=0; col<cn-4; col += 4){  /* over columns of C */
300     colam = col*am;
301     for (i=0; i<am; i++) {        /* over rows of C in those columns */
302       r1 = r2 = r3 = r4 = 0.0;
303       n   = a->i[i+1] - a->i[i];
304       aj  = a->j + a->i[i];
305       aa  = a->a + a->i[i];
306       for (j=0; j<n; j++) {
307         r1 += (*aa)*b1[*aj];
308         r2 += (*aa)*b2[*aj];
309         r3 += (*aa)*b3[*aj];
310         r4 += (*aa++)*b4[*aj++];
311       }
312       c[colam + i]       = r1;
313       c[colam + am + i]  = r2;
314       c[colam + am2 + i] = r3;
315       c[colam + am3 + i] = r4;
316     }
317     b1 += bm4;
318     b2 += bm4;
319     b3 += bm4;
320     b4 += bm4;
321   }
322   for (;col<cn; col++){     /* over extra columns of C */
323     for (i=0; i<am; i++) {  /* over rows of C in those columns */
324       r1 = 0.0;
325       n   = a->i[i+1] - a->i[i];
326       aj  = a->j + a->i[i];
327       aa  = a->a + a->i[i];
328 
329       for (j=0; j<n; j++) {
330         r1 += (*aa++)*b1[*aj++];
331       }
332       c[col*am + i]     = r1;
333     }
334     b1 += bm;
335   }
336   ierr = PetscLogFlops(cn*(2.0*a->nz));CHKERRQ(ierr);
337   ierr = MatRestoreArray(B,&b);CHKERRQ(ierr);
338   ierr = MatRestoreArray(C,&c);CHKERRQ(ierr);
339   ierr = MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
340   ierr = MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
341   PetscFunctionReturn(0);
342 }
343 
344 /*
345    Note very similar to MatMult_SeqAIJ(), should generate both codes from same base
346 */
347 #undef __FUNCT__
348 #define __FUNCT__ "MatMatMultNumericAdd_SeqAIJ_SeqDense"
349 PetscErrorCode MatMatMultNumericAdd_SeqAIJ_SeqDense(Mat A,Mat B,Mat C)
350 {
351   Mat_SeqAIJ     *a=(Mat_SeqAIJ*)A->data;
352   PetscErrorCode ierr;
353   PetscScalar    *b,*c,r1,r2,r3,r4,*b1,*b2,*b3,*b4;
354   MatScalar      *aa;
355   PetscInt       cm=C->rmap->n, cn=B->cmap->n, bm=B->rmap->n, col, i,j,n,*aj, am = A->rmap->n,*ii,arm;
356   PetscInt       am2 = 2*am, am3 = 3*am,  bm4 = 4*bm,colam,*ridx;
357 
358   PetscFunctionBegin;
359   if (!cm || !cn) PetscFunctionReturn(0);
360   ierr = MatGetArray(B,&b);CHKERRQ(ierr);
361   ierr = MatGetArray(C,&c);CHKERRQ(ierr);
362   b1 = b; b2 = b1 + bm; b3 = b2 + bm; b4 = b3 + bm;
363 
364   if (a->compressedrow.use){ /* use compressed row format */
365     for (col=0; col<cn-4; col += 4){  /* over columns of C */
366       colam = col*am;
367       arm   = a->compressedrow.nrows;
368       ii    = a->compressedrow.i;
369       ridx  = a->compressedrow.rindex;
370       for (i=0; i<arm; i++) {        /* over rows of C in those columns */
371 	r1 = r2 = r3 = r4 = 0.0;
372 	n   = ii[i+1] - ii[i];
373 	aj  = a->j + ii[i];
374 	aa  = a->a + ii[i];
375 	for (j=0; j<n; j++) {
376 	  r1 += (*aa)*b1[*aj];
377 	  r2 += (*aa)*b2[*aj];
378 	  r3 += (*aa)*b3[*aj];
379 	  r4 += (*aa++)*b4[*aj++];
380 	}
381 	c[colam       + ridx[i]] += r1;
382 	c[colam + am  + ridx[i]] += r2;
383 	c[colam + am2 + ridx[i]] += r3;
384 	c[colam + am3 + ridx[i]] += r4;
385       }
386       b1 += bm4;
387       b2 += bm4;
388       b3 += bm4;
389       b4 += bm4;
390     }
391     for (;col<cn; col++){     /* over extra columns of C */
392       colam = col*am;
393       arm   = a->compressedrow.nrows;
394       ii    = a->compressedrow.i;
395       ridx  = a->compressedrow.rindex;
396       for (i=0; i<arm; i++) {  /* over rows of C in those columns */
397 	r1 = 0.0;
398 	n   = ii[i+1] - ii[i];
399 	aj  = a->j + ii[i];
400 	aa  = a->a + ii[i];
401 
402 	for (j=0; j<n; j++) {
403 	  r1 += (*aa++)*b1[*aj++];
404 	}
405 	c[col*am + ridx[i]] += r1;
406       }
407       b1 += bm;
408     }
409   } else {
410     for (col=0; col<cn-4; col += 4){  /* over columns of C */
411       colam = col*am;
412       for (i=0; i<am; i++) {        /* over rows of C in those columns */
413 	r1 = r2 = r3 = r4 = 0.0;
414 	n   = a->i[i+1] - a->i[i];
415 	aj  = a->j + a->i[i];
416 	aa  = a->a + a->i[i];
417 	for (j=0; j<n; j++) {
418 	  r1 += (*aa)*b1[*aj];
419 	  r2 += (*aa)*b2[*aj];
420 	  r3 += (*aa)*b3[*aj];
421 	  r4 += (*aa++)*b4[*aj++];
422 	}
423 	c[colam + i]       += r1;
424 	c[colam + am + i]  += r2;
425 	c[colam + am2 + i] += r3;
426 	c[colam + am3 + i] += r4;
427       }
428       b1 += bm4;
429       b2 += bm4;
430       b3 += bm4;
431       b4 += bm4;
432     }
433     for (;col<cn; col++){     /* over extra columns of C */
434       for (i=0; i<am; i++) {  /* over rows of C in those columns */
435 	r1 = 0.0;
436 	n   = a->i[i+1] - a->i[i];
437 	aj  = a->j + a->i[i];
438 	aa  = a->a + a->i[i];
439 
440 	for (j=0; j<n; j++) {
441 	  r1 += (*aa++)*b1[*aj++];
442 	}
443 	c[col*am + i]     += r1;
444       }
445       b1 += bm;
446     }
447   }
448   ierr = PetscLogFlops(cn*2.0*a->nz);CHKERRQ(ierr);
449   ierr = MatRestoreArray(B,&b);CHKERRQ(ierr);
450   ierr = MatRestoreArray(C,&c);CHKERRQ(ierr);
451   PetscFunctionReturn(0);
452 }
453