xref: /petsc/src/mat/impls/aij/mpi/mpiptap.c (revision d724dfffc20f5834ebb4b97bb1e8ef89c8c2f0ed)
1 
2 /*
3   Defines projective product routines where A is a MPIAIJ matrix
4           C = P^T * A * P
5 */
6 
7 #include <../src/mat/impls/aij/seq/aij.h>   /*I "petscmat.h" I*/
8 #include <../src/mat/utils/freespace.h>
9 #include <../src/mat/impls/aij/mpi/mpiaij.h>
10 #include <petscbt.h>
11 
12 /* #define PTAP_PROFILE */
13 
14 extern PetscErrorCode MatDestroy_MPIAIJ(Mat);
15 #undef __FUNCT__
16 #define __FUNCT__ "MatDestroy_MPIAIJ_PtAP"
17 PetscErrorCode MatDestroy_MPIAIJ_PtAP(Mat A)
18 {
19   PetscErrorCode       ierr;
20   Mat_MPIAIJ           *a=(Mat_MPIAIJ*)A->data;
21   Mat_PtAPMPI          *ptap=a->ptap;
22 
23   PetscFunctionBegin;
24   if (ptap){
25     Mat_Merge_SeqsToMPI  *merge=ptap->merge;
26     ierr = PetscFree2(ptap->startsj_s,ptap->startsj_r);CHKERRQ(ierr);
27     ierr = PetscFree(ptap->bufa);CHKERRQ(ierr);
28     ierr = MatDestroy(&ptap->P_loc);CHKERRQ(ierr);
29     ierr = MatDestroy(&ptap->P_oth);CHKERRQ(ierr);
30     ierr = MatDestroy(&ptap->A_loc);CHKERRQ(ierr); /* used by MatTransposeMatMult() */
31     if (ptap->api){ierr = PetscFree(ptap->api);CHKERRQ(ierr);}
32     if (ptap->apj){ierr = PetscFree(ptap->apj);CHKERRQ(ierr);}
33     if (ptap->apa){ierr = PetscFree(ptap->apa);CHKERRQ(ierr);}
34     if (merge) {
35       ierr = PetscFree(merge->id_r);CHKERRQ(ierr);
36       ierr = PetscFree(merge->len_s);CHKERRQ(ierr);
37       ierr = PetscFree(merge->len_r);CHKERRQ(ierr);
38       ierr = PetscFree(merge->bi);CHKERRQ(ierr);
39       ierr = PetscFree(merge->bj);CHKERRQ(ierr);
40       ierr = PetscFree(merge->buf_ri[0]);CHKERRQ(ierr);
41       ierr = PetscFree(merge->buf_ri);CHKERRQ(ierr);
42       ierr = PetscFree(merge->buf_rj[0]);CHKERRQ(ierr);
43       ierr = PetscFree(merge->buf_rj);CHKERRQ(ierr);
44       ierr = PetscFree(merge->coi);CHKERRQ(ierr);
45       ierr = PetscFree(merge->coj);CHKERRQ(ierr);
46       ierr = PetscFree(merge->owners_co);CHKERRQ(ierr);
47       ierr = PetscLayoutDestroy(&merge->rowmap);CHKERRQ(ierr);
48       ierr = merge->destroy(A);CHKERRQ(ierr);
49       ierr = PetscFree(ptap->merge);CHKERRQ(ierr);
50     }
51     ierr = PetscFree(ptap);CHKERRQ(ierr);
52   }
53   PetscFunctionReturn(0);
54 }
55 
56 #undef __FUNCT__
57 #define __FUNCT__ "MatDuplicate_MPIAIJ_MatPtAP"
58 PetscErrorCode MatDuplicate_MPIAIJ_MatPtAP(Mat A, MatDuplicateOption op, Mat *M)
59 {
60   PetscErrorCode       ierr;
61   Mat_MPIAIJ           *a=(Mat_MPIAIJ*)A->data;
62   Mat_PtAPMPI          *ptap = a->ptap;
63   Mat_Merge_SeqsToMPI  *merge = ptap->merge;
64 
65   PetscFunctionBegin;
66   ierr = (*merge->duplicate)(A,op,M);CHKERRQ(ierr);
67   (*M)->ops->destroy   = merge->destroy;
68   (*M)->ops->duplicate = merge->duplicate;
69   PetscFunctionReturn(0);
70 }
71 
72 #undef __FUNCT__
73 #define __FUNCT__ "MatPtAP_MPIAIJ_MPIAIJ"
74 PetscErrorCode MatPtAP_MPIAIJ_MPIAIJ(Mat A,Mat P,MatReuse scall,PetscReal fill,Mat *C)
75 {
76   PetscErrorCode ierr;
77 
78   PetscFunctionBegin;
79   if (scall == MAT_INITIAL_MATRIX){
80     ierr = PetscLogEventBegin(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
81     ierr = MatPtAPSymbolic_MPIAIJ_MPIAIJ(A,P,fill,C);CHKERRQ(ierr);
82     ierr = PetscLogEventEnd(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
83   }
84   ierr = PetscLogEventBegin(MAT_PtAPNumeric,A,P,0,0);CHKERRQ(ierr);
85   ierr = MatPtAPNumeric_MPIAIJ_MPIAIJ(A,P,*C);CHKERRQ(ierr);
86   ierr = PetscLogEventEnd(MAT_PtAPNumeric,A,P,0,0);CHKERRQ(ierr);
87   PetscFunctionReturn(0);
88 }
89 
90 #undef __FUNCT__
91 #define __FUNCT__ "MatPtAPSymbolic_MPIAIJ_MPIAIJ"
92 PetscErrorCode MatPtAPSymbolic_MPIAIJ_MPIAIJ(Mat A,Mat P,PetscReal fill,Mat *C)
93 {
94   PetscErrorCode       ierr;
95   Mat                  Cmpi;
96   Mat_PtAPMPI          *ptap;
97   PetscFreeSpaceList   free_space=PETSC_NULL,current_space=PETSC_NULL;
98   Mat_MPIAIJ           *a=(Mat_MPIAIJ*)A->data,*p=(Mat_MPIAIJ*)P->data,*c;
99   Mat_SeqAIJ           *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data;
100   Mat_SeqAIJ           *p_loc,*p_oth;
101   PetscInt             *pi_loc,*pj_loc,*pi_oth,*pj_oth,*pdti,*pdtj,*poti,*potj,*ptJ;
102   PetscInt             *adi=ad->i,*aj,*aoi=ao->i,nnz;
103   PetscInt             *lnk,*owners_co,*coi,*coj,i,k,pnz,row;
104   PetscInt             am=A->rmap->n,pN=P->cmap->N,pm=P->rmap->n,pn=P->cmap->n;
105   PetscBT              lnkbt;
106   MPI_Comm             comm=((PetscObject)A)->comm;
107   PetscMPIInt          size,rank,tagi,tagj,*len_si,*len_s,*len_ri,icompleted=0;
108   PetscInt             **buf_rj,**buf_ri,**buf_ri_k;
109   PetscInt             len,proc,*dnz,*onz,*owners;
110   PetscInt             nzi,*pti,*ptj;
111   PetscInt             nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci;
112   MPI_Request          *swaits,*rwaits;
113   MPI_Status           *sstatus,rstatus;
114   Mat_Merge_SeqsToMPI  *merge;
115   PetscInt             *api,*apj,*Jptr,apnz,*prmap=p->garray,pon,nspacedouble=0,j,ap_rmax=0;
116   PetscReal            afill=1.0,afill_tmp;
117   PetscInt             rmax;
118 #if defined(PTAP_PROFILE)
119   PetscLogDouble       t0,t1,t2,t3,t4;
120 #endif
121 
122   PetscFunctionBegin;
123 #if defined(PTAP_PROFILE)
124   ierr = PetscGetTime(&t0);CHKERRQ(ierr);
125 #endif
126 
127   /* check if matrix local sizes are compatible */
128   if (A->rmap->rstart != P->rmap->rstart || A->rmap->rend != P->rmap->rend){
129     SETERRQ4(comm,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, Arow (%D, %D) != Prow (%D,%D)",A->rmap->rstart,A->rmap->rend,P->rmap->rstart,P->rmap->rend);
130   }
131   if (A->cmap->rstart != P->rmap->rstart || A->cmap->rend != P->rmap->rend){
132     SETERRQ4(comm,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, Acol (%D, %D) != Prow (%D,%D)",A->cmap->rstart,A->cmap->rend,P->rmap->rstart,P->rmap->rend);
133   }
134 
135   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
136   ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
137 
138   /* create struct Mat_PtAPMPI and attached it to C later */
139   ierr = PetscNew(Mat_PtAPMPI,&ptap);CHKERRQ(ierr);
140   ptap->reuse = MAT_INITIAL_MATRIX;
141 
142   /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */
143   ierr = MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth);CHKERRQ(ierr);
144 
145   /* get P_loc by taking all local rows of P */
146   ierr = MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc);CHKERRQ(ierr);
147 
148   p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data;
149   p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data;
150   pi_loc = p_loc->i; pj_loc = p_loc->j;
151   pi_oth = p_oth->i; pj_oth = p_oth->j;
152 #if defined(PTAP_PROFILE)
153   ierr = PetscGetTime(&t1);CHKERRQ(ierr);
154 #endif
155 
156   /* first, compute symbolic AP = A_loc*P = A_diag*P_loc + A_off*P_oth */
157   /*-------------------------------------------------------------------*/
158   ierr = PetscMalloc((am+1)*sizeof(PetscInt),&api);CHKERRQ(ierr);
159   api[0] = 0;
160 
161   /* create and initialize a linked list */
162   ierr = PetscLLCondensedCreate(pN,pN,&lnk,&lnkbt);CHKERRQ(ierr);
163 
164   /* Initial FreeSpace size is fill*(nnz(A) + nnz(P)) -OOM for ex56, np=8k on Intrepid! */
165   ierr = PetscFreeSpaceGet((PetscInt)(fill*(adi[am]+aoi[am]+pi_loc[pm])),&free_space);CHKERRQ(ierr);
166   current_space = free_space;
167 
168   for (i=0; i<am; i++) {
169     /* diagonal portion of A */
170     nzi = adi[i+1] - adi[i];
171     aj  = ad->j + adi[i];
172     for (j=0; j<nzi; j++){
173       row  = aj[j];
174       pnz  = pi_loc[row+1] - pi_loc[row];
175       Jptr = pj_loc + pi_loc[row];
176       /* add non-zero cols of P into the sorted linked list lnk */
177       ierr = PetscLLCondensedAddSorted(pnz,Jptr,lnk,lnkbt);CHKERRQ(ierr);
178     }
179     /* off-diagonal portion of A */
180     nzi = aoi[i+1] - aoi[i];
181     aj  = ao->j + aoi[i];
182     for (j=0; j<nzi; j++){
183       row = aj[j];
184       pnz = pi_oth[row+1] - pi_oth[row];
185       Jptr = pj_oth + pi_oth[row];
186       ierr = PetscLLCondensedAddSorted(pnz,Jptr,lnk,lnkbt);CHKERRQ(ierr);
187     }
188     apnz = lnk[0];
189     api[i+1] = api[i] + apnz;
190     if (ap_rmax < apnz) ap_rmax = apnz;
191 
192     /* if free space is not available, double the total space in the list */
193     if (current_space->local_remaining<apnz) {
194       ierr = PetscFreeSpaceGet(apnz+current_space->total_array_size,&current_space);CHKERRQ(ierr);
195       nspacedouble++;
196     }
197 
198     /* Copy data into free space, then initialize lnk */
199     ierr = PetscLLCondensedClean(pN,apnz,current_space->array,lnk,lnkbt);CHKERRQ(ierr);
200     current_space->array           += apnz;
201     current_space->local_used      += apnz;
202     current_space->local_remaining -= apnz;
203   }
204 
205   /* Allocate space for apj, initialize apj, and */
206   /* destroy list of free space and other temporary array(s) */
207   ierr = PetscMalloc((api[am]+1)*sizeof(PetscInt),&apj);CHKERRQ(ierr);
208   ierr = PetscFreeSpaceContiguous(&free_space,apj);CHKERRQ(ierr);
209   afill_tmp = (PetscReal)api[am]/(adi[am]+aoi[am]+pi_loc[pm]+1);
210   if (afill_tmp > afill) afill = afill_tmp;
211 
212 #if defined(PTAP_PROFILE)
213   ierr = PetscGetTime(&t2);CHKERRQ(ierr);
214 #endif
215 
216   /* determine symbolic Co=(p->B)^T*AP - send to others */
217   /*----------------------------------------------------*/
218   ierr = MatGetSymbolicTranspose_SeqAIJ(p->B,&poti,&potj);CHKERRQ(ierr);
219 
220   /* then, compute symbolic Co = (p->B)^T*AP */
221   pon = (p->B)->cmap->n; /* total num of rows to be sent to other processors
222                          >= (num of nonzero rows of C_seq) - pn */
223   ierr = PetscMalloc((pon+1)*sizeof(PetscInt),&coi);CHKERRQ(ierr);
224   coi[0] = 0;
225 
226   /* set initial free space to be fill*(nnz(p->B) + nnz(AP)) */
227   nnz           = fill*(poti[pon] + api[am]);
228   ierr          = PetscFreeSpaceGet(nnz,&free_space);
229   current_space = free_space;
230 
231   for (i=0; i<pon; i++) {
232     pnz = poti[i+1] - poti[i];
233     ptJ = potj + poti[i];
234     for (j=0; j<pnz; j++){
235       row  = ptJ[j]; /* row of AP == col of Pot */
236       apnz = api[row+1] - api[row];
237       Jptr = apj + api[row];
238       /* add non-zero cols of AP into the sorted linked list lnk */
239       ierr = PetscLLCondensedAddSorted(apnz,Jptr,lnk,lnkbt);CHKERRQ(ierr);
240     }
241     nnz = lnk[0];
242 
243     /* If free space is not available, double the total space in the list */
244     if (current_space->local_remaining<nnz) {
245       ierr = PetscFreeSpaceGet(nnz+current_space->total_array_size,&current_space);CHKERRQ(ierr);
246       nspacedouble++;
247     }
248 
249     /* Copy data into free space, and zero out denserows */
250     ierr = PetscLLCondensedClean(pN,nnz,current_space->array,lnk,lnkbt);CHKERRQ(ierr);
251     current_space->array           += nnz;
252     current_space->local_used      += nnz;
253     current_space->local_remaining -= nnz;
254     coi[i+1] = coi[i] + nnz;
255   }
256   ierr = PetscMalloc((coi[pon]+1)*sizeof(PetscInt),&coj);CHKERRQ(ierr);
257   ierr = PetscFreeSpaceContiguous(&free_space,coj);CHKERRQ(ierr);
258   afill_tmp = (PetscReal)coi[pon]/(poti[pon] + api[am]+1);
259   if (afill_tmp > afill) afill = afill_tmp;
260   ierr = MatRestoreSymbolicTranspose_SeqAIJ(p->B,&poti,&potj);CHKERRQ(ierr);
261 
262   /* send j-array (coj) of Co to other processors */
263   /*----------------------------------------------*/
264   /* determine row ownership */
265   ierr = PetscNew(Mat_Merge_SeqsToMPI,&merge);CHKERRQ(ierr);
266   ierr = PetscLayoutCreate(comm,&merge->rowmap);CHKERRQ(ierr);
267   merge->rowmap->n  = pn;
268   merge->rowmap->bs = 1;
269   ierr = PetscLayoutSetUp(merge->rowmap);CHKERRQ(ierr);
270   owners = merge->rowmap->range;
271 
272   /* determine the number of messages to send, their lengths */
273   ierr = PetscMalloc2(size,PetscMPIInt,&len_si,size,MPI_Status,&sstatus);CHKERRQ(ierr);
274   ierr = PetscMemzero(len_si,size*sizeof(PetscMPIInt));CHKERRQ(ierr);
275   ierr = PetscMalloc(size*sizeof(PetscMPIInt),&merge->len_s);CHKERRQ(ierr);
276   len_s = merge->len_s;
277   merge->nsend = 0;
278 
279   ierr = PetscMalloc((size+2)*sizeof(PetscInt),&owners_co);CHKERRQ(ierr);
280   ierr = PetscMemzero(len_s,size*sizeof(PetscMPIInt));CHKERRQ(ierr);
281 
282   proc = 0;
283   for (i=0; i<pon; i++){
284     while (prmap[i] >= owners[proc+1]) proc++;
285     len_si[proc]++;  /* num of rows in Co to be sent to [proc] */
286     len_s[proc] += coi[i+1] - coi[i];
287   }
288 
289   len   = 0;  /* max length of buf_si[] */
290   owners_co[0] = 0;
291   for (proc=0; proc<size; proc++){
292     owners_co[proc+1] = owners_co[proc] + len_si[proc];
293     if (len_si[proc]){
294       merge->nsend++;
295       len_si[proc] = 2*(len_si[proc] + 1);
296       len += len_si[proc];
297     }
298   }
299 
300   /* determine the number and length of messages to receive for coi and coj  */
301   ierr = PetscGatherNumberOfMessages(comm,PETSC_NULL,len_s,&merge->nrecv);CHKERRQ(ierr);
302   ierr = PetscGatherMessageLengths2(comm,merge->nsend,merge->nrecv,len_s,len_si,&merge->id_r,&merge->len_r,&len_ri);CHKERRQ(ierr);
303 
304   /* post the Irecv and Isend of coj */
305   ierr = PetscCommGetNewTag(comm,&tagj);CHKERRQ(ierr);
306   ierr = PetscPostIrecvInt(comm,tagj,merge->nrecv,merge->id_r,merge->len_r,&buf_rj,&rwaits);CHKERRQ(ierr);
307   ierr = PetscMalloc((merge->nsend+1)*sizeof(MPI_Request),&swaits);CHKERRQ(ierr);
308   for (proc=0, k=0; proc<size; proc++){
309     if (!len_s[proc]) continue;
310     i = owners_co[proc];
311     ierr = MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tagj,comm,swaits+k);CHKERRQ(ierr);
312     k++;
313   }
314 
315   /* receives and sends of coj are complete */
316   for (i=0; i<merge->nrecv; i++){
317     ierr = MPI_Waitany(merge->nrecv,rwaits,&icompleted,&rstatus);CHKERRQ(ierr);
318   }
319   ierr = PetscFree(rwaits);CHKERRQ(ierr);
320   if (merge->nsend) {ierr = MPI_Waitall(merge->nsend,swaits,sstatus);CHKERRQ(ierr);}
321 
322   /* send and recv coi */
323   /*-------------------*/
324   ierr = PetscCommGetNewTag(comm,&tagi);CHKERRQ(ierr);
325   ierr = PetscPostIrecvInt(comm,tagi,merge->nrecv,merge->id_r,len_ri,&buf_ri,&rwaits);CHKERRQ(ierr);
326   ierr = PetscMalloc((len+1)*sizeof(PetscInt),&buf_s);CHKERRQ(ierr);
327   buf_si = buf_s;  /* points to the beginning of k-th msg to be sent */
328   for (proc=0,k=0; proc<size; proc++){
329     if (!len_s[proc]) continue;
330     /* form outgoing message for i-structure:
331          buf_si[0]:                 nrows to be sent
332                [1:nrows]:           row index (global)
333                [nrows+1:2*nrows+1]: i-structure index
334     */
335     /*-------------------------------------------*/
336     nrows = len_si[proc]/2 - 1;
337     buf_si_i    = buf_si + nrows+1;
338     buf_si[0]   = nrows;
339     buf_si_i[0] = 0;
340     nrows = 0;
341     for (i=owners_co[proc]; i<owners_co[proc+1]; i++){
342       nzi = coi[i+1] - coi[i];
343       buf_si_i[nrows+1] = buf_si_i[nrows] + nzi; /* i-structure */
344       buf_si[nrows+1] =prmap[i] -owners[proc]; /* local row index */
345       nrows++;
346     }
347     ierr = MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tagi,comm,swaits+k);CHKERRQ(ierr);
348     k++;
349     buf_si += len_si[proc];
350   }
351   i = merge->nrecv;
352   while (i--) {
353     ierr = MPI_Waitany(merge->nrecv,rwaits,&icompleted,&rstatus);CHKERRQ(ierr);
354   }
355   ierr = PetscFree(rwaits);CHKERRQ(ierr);
356   if (merge->nsend) {ierr = MPI_Waitall(merge->nsend,swaits,sstatus);CHKERRQ(ierr);}
357 
358   ierr = PetscFree2(len_si,sstatus);CHKERRQ(ierr);
359   ierr = PetscFree(len_ri);CHKERRQ(ierr);
360   ierr = PetscFree(swaits);CHKERRQ(ierr);
361   ierr = PetscFree(buf_s);CHKERRQ(ierr);
362 
363 #if defined(PTAP_PROFILE)
364   ierr = PetscGetTime(&t3);CHKERRQ(ierr);
365 #endif
366 
367   /* compute the local portion of C (mpi mat) */
368   /*------------------------------------------*/
369   ierr = MatGetSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj);CHKERRQ(ierr);
370 
371   /* allocate pti array and free space for accumulating nonzero column info */
372   ierr = PetscMalloc((pn+1)*sizeof(PetscInt),&pti);CHKERRQ(ierr);
373   pti[0] = 0;
374 
375   /* set initial free space to be fill*(nnz(P) + nnz(AP)) */
376   nnz           = fill*(pi_loc[pm] + api[am]);
377   ierr          = PetscFreeSpaceGet(nnz,&free_space);
378   current_space = free_space;
379 
380   ierr = PetscMalloc3(merge->nrecv,PetscInt**,&buf_ri_k,merge->nrecv,PetscInt*,&nextrow,merge->nrecv,PetscInt*,&nextci);CHKERRQ(ierr);
381   for (k=0; k<merge->nrecv; k++){
382     buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */
383     nrows       = *buf_ri_k[k];
384     nextrow[k]  = buf_ri_k[k] + 1;  /* next row number of k-th recved i-structure */
385     nextci[k]   = buf_ri_k[k] + (nrows + 1);/* poins to the next i-structure of k-th recved i-structure  */
386   }
387   ierr = MatPreallocateInitialize(comm,pn,pn,dnz,onz);CHKERRQ(ierr);
388   rmax = 0;
389   for (i=0; i<pn; i++) {
390     /* add pdt[i,:]*AP into lnk */
391     pnz = pdti[i+1] - pdti[i];
392     ptJ = pdtj + pdti[i];
393     for (j=0; j<pnz; j++){
394       row  = ptJ[j];  /* row of AP == col of Pt */
395       apnz = api[row+1] - api[row];
396       Jptr = apj + api[row];
397       /* add non-zero cols of AP into the sorted linked list lnk */
398       ierr = PetscLLCondensedAddSorted(apnz,Jptr,lnk,lnkbt);CHKERRQ(ierr);
399     }
400 
401     /* add received col data into lnk */
402     for (k=0; k<merge->nrecv; k++){ /* k-th received message */
403       if (i == *nextrow[k]) { /* i-th row */
404         nzi = *(nextci[k]+1) - *nextci[k];
405         Jptr  = buf_rj[k] + *nextci[k];
406         ierr = PetscLLCondensedAddSorted(nzi,Jptr,lnk,lnkbt);CHKERRQ(ierr);
407         nextrow[k]++; nextci[k]++;
408       }
409     }
410     nnz = lnk[0];
411 
412     /* if free space is not available, make more free space */
413     if (current_space->local_remaining<nnz) {
414       ierr = PetscFreeSpaceGet(nnz+current_space->total_array_size,&current_space);CHKERRQ(ierr);
415       nspacedouble++;
416     }
417     /* copy data into free space, then initialize lnk */
418     ierr = PetscLLCondensedClean(pN,nnz,current_space->array,lnk,lnkbt);CHKERRQ(ierr);
419     ierr = MatPreallocateSet(i+owners[rank],nnz,current_space->array,dnz,onz);CHKERRQ(ierr);
420     current_space->array           += nnz;
421     current_space->local_used      += nnz;
422     current_space->local_remaining -= nnz;
423     pti[i+1] = pti[i] + nnz;
424     if (nnz > rmax) rmax = nnz;
425   }
426   ierr = MatRestoreSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj);CHKERRQ(ierr);
427   ierr = PetscFree3(buf_ri_k,nextrow,nextci);CHKERRQ(ierr);
428 
429   ierr = PetscMalloc((pti[pn]+1)*sizeof(PetscInt),&ptj);CHKERRQ(ierr);
430   ierr = PetscFreeSpaceContiguous(&free_space,ptj);CHKERRQ(ierr);
431   afill_tmp = (PetscReal)pti[pn]/(pi_loc[pm] + api[am]+1);
432   if (afill_tmp > afill) afill = afill_tmp;
433   ierr = PetscLLDestroy(lnk,lnkbt);CHKERRQ(ierr);
434 
435   /* create symbolic parallel matrix Cmpi */
436   /*--------------------------------------*/
437   ierr = MatCreate(comm,&Cmpi);CHKERRQ(ierr);
438   ierr = MatSetSizes(Cmpi,pn,pn,PETSC_DETERMINE,PETSC_DETERMINE);CHKERRQ(ierr);
439   ierr = MatSetBlockSizes(Cmpi,P->cmap->bs,P->cmap->bs);CHKERRQ(ierr);
440   ierr = MatSetType(Cmpi,MATMPIAIJ);CHKERRQ(ierr);
441   ierr = MatMPIAIJSetPreallocation(Cmpi,0,dnz,0,onz);CHKERRQ(ierr);
442   ierr = MatPreallocateFinalize(dnz,onz);CHKERRQ(ierr);
443 
444   merge->bi            = pti;
445   merge->bj            = ptj;
446   merge->coi           = coi;
447   merge->coj           = coj;
448   merge->buf_ri        = buf_ri;
449   merge->buf_rj        = buf_rj;
450   merge->owners_co     = owners_co;
451   merge->destroy       = Cmpi->ops->destroy;
452   merge->duplicate     = Cmpi->ops->duplicate;
453 
454   /* Cmpi is not ready for use - assembly will be done by MatPtAPNumeric() */
455   Cmpi->assembled      = PETSC_FALSE;
456   Cmpi->ops->destroy   = MatDestroy_MPIAIJ_PtAP;
457   Cmpi->ops->duplicate = MatDuplicate_MPIAIJ_MatPtAP;
458 
459   /* attach the supporting struct to Cmpi for reuse */
460   c = (Mat_MPIAIJ*)Cmpi->data;
461   c->ptap        = ptap;
462   ptap->api      = api;
463   ptap->apj      = apj;
464   ptap->merge    = merge;
465   ptap->rmax     = ap_rmax;
466   *C             = Cmpi;
467 
468   /* flag 'scalable' determines which implementations to be used:
469        0: do dense axpy in MatPtAPNumeric() - fast, but requires storage of a nonscalable dense array apa;
470        1: do sparse axpy in MatPtAPNumeric() - might slow, uses a sparse array apa */
471   /* set default scalable */
472   ptap->scalable = PETSC_TRUE;
473   ierr = PetscOptionsGetBool(((PetscObject)Cmpi)->prefix,"-matptap_scalable",&ptap->scalable,PETSC_NULL);CHKERRQ(ierr);
474   if (!ptap->scalable){  /* Do dense axpy */
475     ierr = PetscMalloc(pN*sizeof(PetscScalar),&ptap->apa);CHKERRQ(ierr);
476     ierr = PetscMemzero(ptap->apa,pN*sizeof(PetscScalar));CHKERRQ(ierr);
477   } else {
478     ierr = PetscMalloc((ap_rmax+1)*sizeof(PetscScalar),&ptap->apa);CHKERRQ(ierr);
479     ierr = PetscMemzero(ptap->apa,ap_rmax*sizeof(PetscScalar));CHKERRQ(ierr);
480   }
481 
482 #if defined(PTAP_PROFILE)
483   ierr = PetscGetTime(&t4);CHKERRQ(ierr);
484   if (rank==1) PetscPrintf(MPI_COMM_SELF,"  [%d] PtAPSymbolic %g/P + %g/AP + %g/comm + %g/PtAP = %g\n",rank,t1-t0,t2-t1,t3-t2,t4-t3,t4-t0);CHKERRQ(ierr);
485 #endif
486 
487 #if defined(PETSC_USE_INFO)
488   if (pti[pn] != 0) {
489     ierr = PetscInfo3(Cmpi,"Reallocs %D; Fill ratio: given %G needed %G.\n",nspacedouble,fill,afill);CHKERRQ(ierr);
490     ierr = PetscInfo1(Cmpi,"Use MatPtAP(A,P,MatReuse,%G,&C) for best performance.\n",afill);CHKERRQ(ierr);
491   } else {
492     ierr = PetscInfo(Cmpi,"Empty matrix product\n");CHKERRQ(ierr);
493   }
494 #endif
495   PetscFunctionReturn(0);
496 }
497 
498 #undef __FUNCT__
499 #define __FUNCT__ "MatPtAPNumeric_MPIAIJ_MPIAIJ"
500 PetscErrorCode MatPtAPNumeric_MPIAIJ_MPIAIJ(Mat A,Mat P,Mat C)
501 {
502   PetscErrorCode       ierr;
503   Mat_Merge_SeqsToMPI  *merge;
504   Mat_MPIAIJ           *a=(Mat_MPIAIJ*)A->data,*p=(Mat_MPIAIJ*)P->data,*c=(Mat_MPIAIJ*)C->data;
505   Mat_SeqAIJ           *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data;
506   Mat_SeqAIJ           *pd=(Mat_SeqAIJ*)(p->A)->data,*po=(Mat_SeqAIJ*)(p->B)->data;
507   Mat_SeqAIJ           *p_loc,*p_oth;
508   Mat_PtAPMPI          *ptap;
509   PetscInt             *adi=ad->i,*aoi=ao->i,*adj,*aoj,*apJ,nextp;
510   PetscInt             *pi_loc,*pj_loc,*pi_oth,*pj_oth,*pJ,*pj;
511   PetscInt             i,j,k,anz,pnz,apnz,nextap,row,*cj;
512   MatScalar            *ada,*aoa,*apa,*pa,*ca,*pa_loc,*pa_oth,valtmp;
513   PetscInt             am=A->rmap->n,cm=C->rmap->n,pon=(p->B)->cmap->n;
514   MPI_Comm             comm=((PetscObject)C)->comm;
515   PetscMPIInt          size,rank,taga,*len_s;
516   PetscInt             *owners,proc,nrows,**buf_ri_k,**nextrow,**nextci;
517   PetscInt             **buf_ri,**buf_rj;
518   PetscInt             cnz=0,*bj_i,*bi,*bj,bnz,nextcj; /* bi,bj,ba: local array of C(mpi mat) */
519   MPI_Request          *s_waits,*r_waits;
520   MPI_Status           *status;
521   MatScalar            **abuf_r,*ba_i,*pA,*coa,*ba;
522   PetscInt             *api,*apj,*coi,*coj;
523   PetscInt             *poJ=po->j,*pdJ=pd->j,pcstart=P->cmap->rstart,pcend=P->cmap->rend;
524   PetscBool            scalable;
525 #if defined(PTAP_PROFILE)
526   PetscLogDouble       t0,t1,t2,t3,t4;
527 #endif
528 
529   PetscFunctionBegin;
530 #if defined(PTAP_PROFILE)
531   ierr = PetscGetTime(&t0);CHKERRQ(ierr);
532 #endif
533   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
534   ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
535 
536   ptap = c->ptap;
537   if (!ptap) SETERRQ(((PetscObject)C)->comm,PETSC_ERR_ARG_INCOMP,"MatPtAP() has not been called to create matrix C yet, cannot use MAT_REUSE_MATRIX");
538   merge    = ptap->merge;
539   apa      = ptap->apa;
540   scalable = ptap->scalable;
541 
542   /* 1) get P_oth = ptap->P_oth  and P_loc = ptap->P_loc */
543   /*--------------------------------------------------*/
544   if (ptap->reuse == MAT_INITIAL_MATRIX){
545     ptap->reuse = MAT_REUSE_MATRIX;
546   } else { /* update numerical values of P_oth and P_loc */
547     ierr = MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_REUSE_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth);CHKERRQ(ierr);
548     ierr = MatMPIAIJGetLocalMat(P,MAT_REUSE_MATRIX,&ptap->P_loc);CHKERRQ(ierr);
549   }
550 #if defined(PTAP_PROFILE)
551   ierr = PetscGetTime(&t1);CHKERRQ(ierr);
552 #endif
553 
554   /* 2) compute numeric C_seq = P_loc^T*A_loc*P - dominating part */
555   /*--------------------------------------------------------------*/
556   /* get data from symbolic products */
557   p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data;
558   p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data;
559   pi_loc=p_loc->i; pj_loc=p_loc->j; pJ=pj_loc; pa_loc=p_loc->a;
560   pi_oth=p_oth->i; pj_oth=p_oth->j; pa_oth=p_oth->a;
561 
562   coi = merge->coi; coj = merge->coj;
563   ierr = PetscMalloc((coi[pon]+1)*sizeof(MatScalar),&coa);CHKERRQ(ierr);
564   ierr = PetscMemzero(coa,coi[pon]*sizeof(MatScalar));CHKERRQ(ierr);
565 
566   bi     = merge->bi; bj = merge->bj;
567   owners = merge->rowmap->range;
568   ierr   = PetscMalloc((bi[cm]+1)*sizeof(MatScalar),&ba);CHKERRQ(ierr);
569   ierr   = PetscMemzero(ba,bi[cm]*sizeof(MatScalar));CHKERRQ(ierr);
570 
571   api = ptap->api; apj = ptap->apj;
572 
573   if (!scalable){ /* Do dense axpy on apa (length of pN, stores A[i,:]*P) - nonscalable, but fast */
574     /*-----------------------------------------------------------------------------------------------------*/
575     for (i=0; i<am; i++) {
576       /* 2-a) form i-th sparse row of A_loc*P = Ad*P_loc + Ao*P_oth */
577       /*------------------------------------------------------------*/
578       apJ = apj + api[i];
579 
580       /* diagonal portion of A */
581       anz = adi[i+1] - adi[i];
582       adj = ad->j + adi[i];
583       ada = ad->a + adi[i];
584       for (j=0; j<anz; j++) {
585         row = adj[j];
586         pnz = pi_loc[row+1] - pi_loc[row];
587         pj  = pj_loc + pi_loc[row];
588         pa  = pa_loc + pi_loc[row];
589 
590         /* perform dense axpy */
591         valtmp = ada[j];
592         for (k=0; k<pnz; k++){
593           apa[pj[k]] += valtmp*pa[k];
594         }
595         ierr = PetscLogFlops(2.0*pnz);CHKERRQ(ierr);
596       }
597 
598       /* off-diagonal portion of A */
599       anz = aoi[i+1] - aoi[i];
600       aoj = ao->j + aoi[i];
601       aoa = ao->a + aoi[i];
602       for (j=0; j<anz; j++) {
603         row = aoj[j];
604         pnz = pi_oth[row+1] - pi_oth[row];
605         pj  = pj_oth + pi_oth[row];
606         pa  = pa_oth + pi_oth[row];
607 
608         /* perform dense axpy */
609         valtmp = aoa[j];
610         for (k=0; k<pnz; k++){
611           apa[pj[k]] += valtmp*pa[k];
612         }
613         ierr = PetscLogFlops(2.0*pnz);CHKERRQ(ierr);
614       }
615 
616       /* 2-b) Compute Cseq = P_loc[i,:]^T*AP[i,:] using outer product */
617       /*--------------------------------------------------------------*/
618       apnz = api[i+1] - api[i];
619       /* put the value into Co=(p->B)^T*AP (off-diagonal part, send to others) */
620       pnz = po->i[i+1] - po->i[i];
621       poJ = po->j + po->i[i];
622       pA  = po->a + po->i[i];
623       for (j=0; j<pnz; j++){
624         row = poJ[j];
625         cnz = coi[row+1] - coi[row];
626         cj  = coj + coi[row];
627         ca  = coa + coi[row];
628         /* perform dense axpy */
629         valtmp = pA[j];
630         for (k=0; k<cnz; k++) {
631           ca[k] += valtmp*apa[cj[k]];
632         }
633         ierr = PetscLogFlops(2.0*cnz);CHKERRQ(ierr);
634       }
635 
636       /* put the value into Cd (diagonal part) */
637       pnz = pd->i[i+1] - pd->i[i];
638       pdJ = pd->j + pd->i[i];
639       pA  = pd->a + pd->i[i];
640       for (j=0; j<pnz; j++){
641         row = pdJ[j];
642         cnz = bi[row+1] - bi[row];
643         cj  = bj + bi[row];
644         ca  = ba + bi[row];
645         /* perform dense axpy */
646         valtmp = pA[j];
647         for (k=0; k<cnz; k++) {
648           ca[k] += valtmp*apa[cj[k]];
649         }
650         ierr = PetscLogFlops(2.0*cnz);CHKERRQ(ierr);
651       }
652 
653       /* zero the current row of A*P */
654       for (k=0; k<apnz; k++) apa[apJ[k]] = 0.0;
655     }
656   } else {/* Do sparse axpy on apa (length of ap_rmax, stores A[i,:]*P) - scalable, but slower */
657     /*-----------------------------------------------------------------------------------------*/
658     pA=pa_loc;
659     for (i=0; i<am; i++) {
660       /* form i-th sparse row of A*P */
661       apnz = api[i+1] - api[i];
662       apJ  = apj + api[i];
663       /* diagonal portion of A */
664       anz = adi[i+1] - adi[i];
665       adj = ad->j + adi[i];
666       ada = ad->a + adi[i];
667       for (j=0; j<anz; j++) {
668         row = adj[j];
669         pnz = pi_loc[row+1] - pi_loc[row];
670         pj  = pj_loc + pi_loc[row];
671         pa  = pa_loc + pi_loc[row];
672         valtmp = ada[j];
673         nextp  = 0;
674         for (k=0; nextp<pnz; k++) {
675           if (apJ[k] == pj[nextp]) { /* col of AP == col of P */
676             apa[k] += valtmp*pa[nextp++];
677           }
678         }
679         ierr = PetscLogFlops(2.0*pnz);CHKERRQ(ierr);
680       }
681       /* off-diagonal portion of A */
682       anz = aoi[i+1] - aoi[i];
683       aoj = ao->j + aoi[i];
684       aoa = ao->a + aoi[i];
685       for (j=0; j<anz; j++) {
686         row = aoj[j];
687         pnz = pi_oth[row+1] - pi_oth[row];
688         pj  = pj_oth + pi_oth[row];
689         pa  = pa_oth + pi_oth[row];
690         valtmp = aoa[j];
691         nextp  = 0;
692         for (k=0; nextp<pnz; k++) {
693           if (apJ[k] == pj[nextp]) { /* col of AP == col of P */
694             apa[k] += valtmp*pa[nextp++];
695           }
696         }
697         ierr = PetscLogFlops(2.0*pnz);CHKERRQ(ierr);
698       }
699 
700       /* 2-b) Compute Cseq = P_loc[i,:]^T*AP[i,:] using outer product */
701       /*--------------------------------------------------------------*/
702       pnz = pi_loc[i+1] - pi_loc[i];
703       pJ  = pj_loc + pi_loc[i];
704       for (j=0; j<pnz; j++) {
705         nextap = 0;
706         row    = pJ[j]; /* global index */
707         if (row < pcstart || row >=pcend) { /* put the value into Co */
708           row = *poJ;
709           cj  = coj + coi[row];
710           ca  = coa + coi[row]; poJ++;
711         } else {                            /* put the value into Cd */
712           row  = *pdJ;
713           cj   = bj + bi[row];
714           ca   = ba + bi[row]; pdJ++;
715         }
716         valtmp = pA[j];
717         for (k=0; nextap<apnz; k++) {
718           if (cj[k]==apJ[nextap]) ca[k] += valtmp*apa[nextap++];
719         }
720         ierr = PetscLogFlops(2.0*apnz);CHKERRQ(ierr);
721       }
722       pA += pnz;
723       /* zero the current row info for A*P */
724       ierr = PetscMemzero(apa,apnz*sizeof(MatScalar));CHKERRQ(ierr);
725     }
726   }
727 #if defined(PTAP_PROFILE)
728   ierr = PetscGetTime(&t2);CHKERRQ(ierr);
729 #endif
730 
731   /* 3) send and recv matrix values coa */
732   /*------------------------------------*/
733   buf_ri = merge->buf_ri;
734   buf_rj = merge->buf_rj;
735   len_s  = merge->len_s;
736   ierr = PetscCommGetNewTag(comm,&taga);CHKERRQ(ierr);
737   ierr = PetscPostIrecvScalar(comm,taga,merge->nrecv,merge->id_r,merge->len_r,&abuf_r,&r_waits);CHKERRQ(ierr);
738 
739   ierr = PetscMalloc2(merge->nsend+1,MPI_Request,&s_waits,size,MPI_Status,&status);CHKERRQ(ierr);
740   for (proc=0,k=0; proc<size; proc++){
741     if (!len_s[proc]) continue;
742     i = merge->owners_co[proc];
743     ierr = MPI_Isend(coa+coi[i],len_s[proc],MPIU_MATSCALAR,proc,taga,comm,s_waits+k);CHKERRQ(ierr);
744     k++;
745   }
746   if (merge->nrecv) {ierr = MPI_Waitall(merge->nrecv,r_waits,status);CHKERRQ(ierr);}
747   if (merge->nsend) {ierr = MPI_Waitall(merge->nsend,s_waits,status);CHKERRQ(ierr);}
748 
749   ierr = PetscFree2(s_waits,status);CHKERRQ(ierr);
750   ierr = PetscFree(r_waits);CHKERRQ(ierr);
751   ierr = PetscFree(coa);CHKERRQ(ierr);
752 #if defined(PTAP_PROFILE)
753   ierr = PetscGetTime(&t3);CHKERRQ(ierr);
754 #endif
755 
756   /* 4) insert local Cseq and received values into Cmpi */
757   /*------------------------------------------------------*/
758   ierr = PetscMalloc3(merge->nrecv,PetscInt**,&buf_ri_k,merge->nrecv,PetscInt*,&nextrow,merge->nrecv,PetscInt*,&nextci);CHKERRQ(ierr);
759   for (k=0; k<merge->nrecv; k++){
760     buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */
761     nrows       = *(buf_ri_k[k]);
762     nextrow[k]  = buf_ri_k[k]+1;  /* next row number of k-th recved i-structure */
763     nextci[k]   = buf_ri_k[k] + (nrows + 1);/* poins to the next i-structure of k-th recved i-structure  */
764   }
765 
766   for (i=0; i<cm; i++) {
767     row = owners[rank] + i; /* global row index of C_seq */
768     bj_i = bj + bi[i];  /* col indices of the i-th row of C */
769     ba_i = ba + bi[i];
770     bnz  = bi[i+1] - bi[i];
771     /* add received vals into ba */
772     for (k=0; k<merge->nrecv; k++){ /* k-th received message */
773       /* i-th row */
774       if (i == *nextrow[k]) {
775         cnz = *(nextci[k]+1) - *nextci[k];
776         cj  = buf_rj[k] + *(nextci[k]);
777         ca  = abuf_r[k] + *(nextci[k]);
778         nextcj = 0;
779         for (j=0; nextcj<cnz; j++){
780           if (bj_i[j] == cj[nextcj]){ /* bcol == ccol */
781             ba_i[j] += ca[nextcj++];
782           }
783         }
784         nextrow[k]++; nextci[k]++;
785         ierr = PetscLogFlops(2.0*cnz);CHKERRQ(ierr);
786       }
787     }
788     ierr = MatSetValues(C,1,&row,bnz,bj_i,ba_i,INSERT_VALUES);CHKERRQ(ierr);
789   }
790   ierr = MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
791   ierr = MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
792 
793   ierr = PetscFree(ba);CHKERRQ(ierr);
794   ierr = PetscFree(abuf_r[0]);CHKERRQ(ierr);
795   ierr = PetscFree(abuf_r);CHKERRQ(ierr);
796   ierr = PetscFree3(buf_ri_k,nextrow,nextci);CHKERRQ(ierr);
797 #if defined(PTAP_PROFILE)
798   ierr = PetscGetTime(&t4);CHKERRQ(ierr);
799   if (rank==1) PetscPrintf(MPI_COMM_SELF,"  [%d] PtAPNum %g/P + %g/PtAP + %g/comm + %g/Cloc = %g\n\n",rank,t1-t0,t2-t1,t3-t2,t4-t3,t4-t0);CHKERRQ(ierr);
800 #endif
801 
802 
803 #if defined(PETSC_USE_INFO)
804   if (scalable){
805     ierr = PetscInfo(C,"Use scalable sparse axpy\n");CHKERRQ(ierr);
806   } else {
807     ierr = PetscInfo(C,"Use non-scalable dense axpy\n");CHKERRQ(ierr);
808   }
809 #endif
810   PetscFunctionReturn(0);
811 }
812