xref: /petsc/src/mat/impls/aij/mpi/mpiptap.c (revision ef7a94f2f2432c6d7ea17e36da41b31cbe9f5832)
182412ba7SHong Zhang 
242c57489SHong Zhang /*
342c57489SHong Zhang   Defines projective product routines where A is a MPIAIJ matrix
442c57489SHong Zhang           C = P^T * A * P
542c57489SHong Zhang */
642c57489SHong Zhang 
7c6db04a5SJed Brown #include <../src/mat/impls/aij/seq/aij.h>   /*I "petscmat.h" I*/
8c6db04a5SJed Brown #include <../src/mat/utils/freespace.h>
9c6db04a5SJed Brown #include <../src/mat/impls/aij/mpi/mpiaij.h>
10c6db04a5SJed Brown #include <petscbt.h>
118563dfccSBarry Smith #include <petsctime.h>
124a56b808SFande Kong #include <petsc/private/hashmapiv.h>
134a56b808SFande Kong #include <petsc/private/hashseti.h>
144a56b808SFande Kong #include <petscsf.h>
1542c57489SHong Zhang 
16f671be37SHong Zhang /* #define PTAP_PROFILE */
1724ecddacSHong Zhang 
18cf97cf3cSHong Zhang PetscErrorCode MatView_MPIAIJ_PtAP(Mat A,PetscViewer viewer)
19cf97cf3cSHong Zhang {
20cf97cf3cSHong Zhang   PetscErrorCode    ierr;
21cf97cf3cSHong Zhang   Mat_MPIAIJ        *a=(Mat_MPIAIJ*)A->data;
223cdca5ebSHong Zhang   Mat_APMPI         *ptap=a->ap;
23cf97cf3cSHong Zhang   PetscBool         iascii;
24cf97cf3cSHong Zhang   PetscViewerFormat format;
25cf97cf3cSHong Zhang 
26cf97cf3cSHong Zhang   PetscFunctionBegin;
27baa1ba78SHong Zhang   if (!ptap) {
28baa1ba78SHong Zhang     /* hack: MatDuplicate() sets oldmat->ops->view to newmat which is a base mat class with null ptpa! */
29baa1ba78SHong Zhang     A->ops->view = MatView_MPIAIJ;
30baa1ba78SHong Zhang     ierr = (A->ops->view)(A,viewer);CHKERRQ(ierr);
31baa1ba78SHong Zhang     PetscFunctionReturn(0);
32baa1ba78SHong Zhang   }
33baa1ba78SHong Zhang 
34cf97cf3cSHong Zhang   ierr = PetscObjectTypeCompare((PetscObject)viewer,PETSCVIEWERASCII,&iascii);CHKERRQ(ierr);
35cf97cf3cSHong Zhang   if (iascii) {
36cf97cf3cSHong Zhang     ierr = PetscViewerGetFormat(viewer,&format);CHKERRQ(ierr);
37cf97cf3cSHong Zhang     if (format == PETSC_VIEWER_ASCII_INFO || format == PETSC_VIEWER_ASCII_INFO_DETAIL) {
38cf97cf3cSHong Zhang       if (ptap->algType == 0) {
39cf97cf3cSHong Zhang         ierr = PetscViewerASCIIPrintf(viewer,"using scalable MatPtAP() implementation\n");CHKERRQ(ierr);
40cf97cf3cSHong Zhang       } else if (ptap->algType == 1) {
41cf97cf3cSHong Zhang         ierr = PetscViewerASCIIPrintf(viewer,"using nonscalable MatPtAP() implementation\n");CHKERRQ(ierr);
424a56b808SFande Kong       } else if (ptap->algType == 2) {
434a56b808SFande Kong         ierr = PetscViewerASCIIPrintf(viewer,"using allatonce MatPtAP() implementation\n");CHKERRQ(ierr);
444a56b808SFande Kong       } else if (ptap->algType == 3) {
454a56b808SFande Kong         ierr = PetscViewerASCIIPrintf(viewer,"using merged allatonce MatPtAP() implementation\n");CHKERRQ(ierr);
46cf97cf3cSHong Zhang       }
47cf97cf3cSHong Zhang     }
48cf97cf3cSHong Zhang   }
49cf97cf3cSHong Zhang   ierr = (ptap->view)(A,viewer);CHKERRQ(ierr);
50cf97cf3cSHong Zhang   PetscFunctionReturn(0);
51cf97cf3cSHong Zhang }
52cf97cf3cSHong Zhang 
534624976aSHong Zhang PetscErrorCode MatFreeIntermediateDataStructures_MPIAIJ_AP(Mat A)
54cc6cb787SHong Zhang {
55cc6cb787SHong Zhang   PetscErrorCode      ierr;
56f8487c73SHong Zhang   Mat_MPIAIJ          *a=(Mat_MPIAIJ*)A->data;
573cdca5ebSHong Zhang   Mat_APMPI           *ptap=a->ap;
5860552ceaSHong Zhang   Mat_Merge_SeqsToMPI *merge;
59cc6cb787SHong Zhang 
60cc6cb787SHong Zhang   PetscFunctionBegin;
61dd4011a9SHong Zhang   if (!ptap) PetscFunctionReturn(0);
62dd4011a9SHong Zhang 
63b7f45c76SHong Zhang   ierr = PetscFree2(ptap->startsj_s,ptap->startsj_r);CHKERRQ(ierr);
64f8487c73SHong Zhang   ierr = PetscFree(ptap->bufa);CHKERRQ(ierr);
65a1a4e84aSHong Zhang   ierr = MatDestroy(&ptap->P_loc);CHKERRQ(ierr);
66a1a4e84aSHong Zhang   ierr = MatDestroy(&ptap->P_oth);CHKERRQ(ierr);
67c5af039cSHong Zhang   ierr = MatDestroy(&ptap->A_loc);CHKERRQ(ierr); /* used by MatTransposeMatMult() */
6805d62848SHong Zhang   ierr = MatDestroy(&ptap->Rd);CHKERRQ(ierr);
6905d62848SHong Zhang   ierr = MatDestroy(&ptap->Ro);CHKERRQ(ierr);
708403a639SHong Zhang   if (ptap->AP_loc) { /* used by alg_rap */
71681d504bSHong Zhang     Mat_SeqAIJ *ap = (Mat_SeqAIJ*)(ptap->AP_loc)->data;
72681d504bSHong Zhang     ierr = PetscFree(ap->i);CHKERRQ(ierr);
73681d504bSHong Zhang     ierr = PetscFree2(ap->j,ap->a);CHKERRQ(ierr);
740d3441aeSHong Zhang     ierr = MatDestroy(&ptap->AP_loc);CHKERRQ(ierr);
758403a639SHong Zhang   } else { /* used by alg_ptap */
768403a639SHong Zhang     ierr = PetscFree(ptap->api);CHKERRQ(ierr);
778403a639SHong Zhang     ierr = PetscFree(ptap->apj);CHKERRQ(ierr);
78681d504bSHong Zhang   }
792259aa2eSHong Zhang   ierr = MatDestroy(&ptap->C_loc);CHKERRQ(ierr);
802259aa2eSHong Zhang   ierr = MatDestroy(&ptap->C_oth);CHKERRQ(ierr);
81414894bdSHong Zhang   if (ptap->apa) {ierr = PetscFree(ptap->apa);CHKERRQ(ierr);}
82a560ef98SHong Zhang 
8360552ceaSHong Zhang   ierr = MatDestroy(&ptap->Pt);CHKERRQ(ierr);
8460552ceaSHong Zhang 
8560552ceaSHong Zhang   merge=ptap->merge;
868403a639SHong Zhang   if (merge) { /* used by alg_ptap */
87cc6cb787SHong Zhang     ierr = PetscFree(merge->id_r);CHKERRQ(ierr);
88cc6cb787SHong Zhang     ierr = PetscFree(merge->len_s);CHKERRQ(ierr);
89cc6cb787SHong Zhang     ierr = PetscFree(merge->len_r);CHKERRQ(ierr);
90cc6cb787SHong Zhang     ierr = PetscFree(merge->bi);CHKERRQ(ierr);
91cc6cb787SHong Zhang     ierr = PetscFree(merge->bj);CHKERRQ(ierr);
92c05d87d6SBarry Smith     ierr = PetscFree(merge->buf_ri[0]);CHKERRQ(ierr);
93cc6cb787SHong Zhang     ierr = PetscFree(merge->buf_ri);CHKERRQ(ierr);
94c05d87d6SBarry Smith     ierr = PetscFree(merge->buf_rj[0]);CHKERRQ(ierr);
95cc6cb787SHong Zhang     ierr = PetscFree(merge->buf_rj);CHKERRQ(ierr);
96445158ffSHong Zhang     ierr = PetscFree(merge->coi);CHKERRQ(ierr);
97445158ffSHong Zhang     ierr = PetscFree(merge->coj);CHKERRQ(ierr);
9805b42c5fSBarry Smith     ierr = PetscFree(merge->owners_co);CHKERRQ(ierr);
996bf464f9SBarry Smith     ierr = PetscLayoutDestroy(&merge->rowmap);CHKERRQ(ierr);
100f8487c73SHong Zhang     ierr = PetscFree(ptap->merge);CHKERRQ(ierr);
101bf0cc555SLisandro Dalcin   }
102a3bb6f32SFande Kong   ierr = ISLocalToGlobalMappingDestroy(&ptap->ltog);CHKERRQ(ierr);
1034a56b808SFande Kong 
1044a56b808SFande Kong   ierr = PetscSFDestroy(&ptap->sf);CHKERRQ(ierr);
1054a56b808SFande Kong   ierr = PetscFree(ptap->c_othi);CHKERRQ(ierr);
1064a56b808SFande Kong   ierr = PetscFree(ptap->c_rmti);CHKERRQ(ierr);
1073697aca6SHong Zhang   PetscFunctionReturn(0);
108c8b0795cSMark F. Adams }
109dce485f0SBarry Smith 
1103697aca6SHong Zhang PetscErrorCode MatDestroy_MPIAIJ_PtAP(Mat A)
1113697aca6SHong Zhang {
1123697aca6SHong Zhang   PetscErrorCode ierr;
1139b3d8a68SHong Zhang   Mat_MPIAIJ     *a=(Mat_MPIAIJ*)A->data;
1149b3d8a68SHong Zhang   Mat_APMPI      *ptap=a->ap;
1153697aca6SHong Zhang 
1163697aca6SHong Zhang   PetscFunctionBegin;
117dd4011a9SHong Zhang   ierr = (*A->ops->freeintermediatedatastructures)(A);CHKERRQ(ierr);
1189b3d8a68SHong Zhang   ierr = (*ptap->destroy)(A);CHKERRQ(ierr); /* MatDestroy_MPIAIJ(A) */
1199b3d8a68SHong Zhang   ierr = PetscFree(ptap);CHKERRQ(ierr);
120cc6cb787SHong Zhang   PetscFunctionReturn(0);
121cc6cb787SHong Zhang }
122cc6cb787SHong Zhang 
123150d2497SBarry Smith PETSC_INTERN PetscErrorCode MatPtAP_MPIAIJ_MPIAIJ(Mat A,Mat P,MatReuse scall,PetscReal fill,Mat *C)
12442c57489SHong Zhang {
12542c57489SHong Zhang   PetscErrorCode ierr;
126a4ffb656SHong Zhang   PetscBool      flg;
12767a12041SHong Zhang   MPI_Comm       comm;
128a4ffb656SHong Zhang #if !defined(PETSC_HAVE_HYPRE)
1294a56b808SFande Kong   const char          *algTypes[4] = {"scalable","nonscalable","allatonce","allatonce_merged"};
1304a56b808SFande Kong   PetscInt            nalg=4;
131a4ffb656SHong Zhang #else
1324a56b808SFande Kong   const char          *algTypes[5] = {"scalable","nonscalable","allatonce","allatonce_merged","hypre"};
1334a56b808SFande Kong   PetscInt            nalg=5;
134a4ffb656SHong Zhang #endif
135a4ffb656SHong Zhang   PetscInt            pN=P->cmap->N,alg=1; /* set default algorithm */
13642c57489SHong Zhang 
13742c57489SHong Zhang   PetscFunctionBegin;
13867a12041SHong Zhang   /* check if matrix local sizes are compatible */
13967a12041SHong Zhang   ierr = PetscObjectGetComm((PetscObject)A,&comm);CHKERRQ(ierr);
1406c4ed002SBarry Smith   if (A->rmap->rstart != P->rmap->rstart || A->rmap->rend != P->rmap->rend) 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);
1416c4ed002SBarry Smith   if (A->cmap->rstart != P->rmap->rstart || A->cmap->rend != P->rmap->rend) 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);
14267a12041SHong Zhang 
143cf3ca8ceSHong Zhang   if (scall == MAT_INITIAL_MATRIX) {
144a4ffb656SHong Zhang     /* pick an algorithm */
145715a5346SHong Zhang     ierr = PetscOptionsBegin(PetscObjectComm((PetscObject)A),((PetscObject)A)->prefix,"MatPtAP","Mat");CHKERRQ(ierr);
146a4ffb656SHong Zhang     ierr = PetscOptionsEList("-matptap_via","Algorithmic approach","MatPtAP",algTypes,nalg,algTypes[alg],&alg,&flg);CHKERRQ(ierr);
147a4ffb656SHong Zhang     ierr = PetscOptionsEnd();CHKERRQ(ierr);
148a4ffb656SHong Zhang 
149a4ffb656SHong Zhang     if (!flg && pN > 100000) { /* may switch to scalable algorithm as default */
150a4ffb656SHong Zhang       MatInfo     Ainfo,Pinfo;
151a4ffb656SHong Zhang       PetscInt    nz_local;
152a4ffb656SHong Zhang       PetscBool   alg_scalable_loc=PETSC_FALSE,alg_scalable;
153a4ffb656SHong Zhang 
154a4ffb656SHong Zhang       ierr = MatGetInfo(A,MAT_LOCAL,&Ainfo);CHKERRQ(ierr);
155a4ffb656SHong Zhang       ierr = MatGetInfo(P,MAT_LOCAL,&Pinfo);CHKERRQ(ierr);
156a4ffb656SHong Zhang       nz_local = (PetscInt)(Ainfo.nz_allocated + Pinfo.nz_allocated);
157a4ffb656SHong Zhang 
158a4ffb656SHong Zhang       if (pN > fill*nz_local) alg_scalable_loc = PETSC_TRUE;
159a4ffb656SHong Zhang       ierr = MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm);CHKERRQ(ierr);
160a4ffb656SHong Zhang 
161a4ffb656SHong Zhang       if (alg_scalable) {
162a4ffb656SHong Zhang         alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */
1630d3441aeSHong Zhang       }
164a4ffb656SHong Zhang     }
165a4ffb656SHong Zhang 
166a4ffb656SHong Zhang     switch (alg) {
167a4ffb656SHong Zhang     case 1:
16880da8df7SHong Zhang       /* do R=P^T locally, then C=R*A*P -- nonscalable */
169a4ffb656SHong Zhang       ierr = PetscLogEventBegin(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
170a4ffb656SHong Zhang       ierr = MatPtAPSymbolic_MPIAIJ_MPIAIJ(A,P,fill,C);CHKERRQ(ierr);
1713ff4c91cSHong Zhang       ierr = PetscLogEventEnd(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
172a4ffb656SHong Zhang       break;
173a4ffb656SHong Zhang     case 2:
1744a56b808SFande Kong       /* compute C=P^T*A*P allatonce */
1754a56b808SFande Kong       ierr = PetscLogEventBegin(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
1764a56b808SFande Kong       ierr = MatPtAPSymbolic_MPIAIJ_MPIAIJ_allatonce(A,P,fill,C);CHKERRQ(ierr);
1774a56b808SFande Kong       ierr = PetscLogEventEnd(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
1784a56b808SFande Kong       break;
1794a56b808SFande Kong     case 3:
1804a56b808SFande Kong       /* compute C=P^T*A*P allatonce */
1814a56b808SFande Kong       ierr = PetscLogEventBegin(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
1824a56b808SFande Kong       ierr = MatPtAPSymbolic_MPIAIJ_MPIAIJ_allatonce_merged(A,P,fill,C);CHKERRQ(ierr);
1834a56b808SFande Kong       ierr = PetscLogEventEnd(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
1844a56b808SFande Kong       break;
1854a56b808SFande Kong #if defined(PETSC_HAVE_HYPRE)
1864a56b808SFande Kong     case 4:
187a4ffb656SHong Zhang       /* Use boomerAMGBuildCoarseOperator */
1880abfe76eSFande Kong       ierr = PetscLogEventBegin(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
189a4ffb656SHong Zhang       ierr = MatPtAPSymbolic_AIJ_AIJ_wHYPRE(A,P,fill,C);CHKERRQ(ierr);
1900abfe76eSFande Kong       ierr = PetscLogEventEnd(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
191a4ffb656SHong Zhang       break;
192a4ffb656SHong Zhang #endif
193a4ffb656SHong Zhang     default:
194a4ffb656SHong Zhang       /* do R=P^T locally, then C=R*A*P */
195a4ffb656SHong Zhang       ierr = PetscLogEventBegin(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
196a4ffb656SHong Zhang       ierr = MatPtAPSymbolic_MPIAIJ_MPIAIJ_scalable(A,P,fill,C);CHKERRQ(ierr);
197a4ffb656SHong Zhang       ierr = PetscLogEventEnd(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
198a4ffb656SHong Zhang       break;
199a4ffb656SHong Zhang     }
2007d0a89b7SHong Zhang 
2014a56b808SFande Kong     if (alg == 0 || alg == 1 || alg == 2 || alg == 3) {
2027d0a89b7SHong Zhang       Mat_MPIAIJ *c  = (Mat_MPIAIJ*)(*C)->data;
2037d0a89b7SHong Zhang       Mat_APMPI  *ap = c->ap;
2047d0a89b7SHong Zhang       ierr = PetscOptionsBegin(PetscObjectComm((PetscObject)(*C)),((PetscObject)(*C))->prefix,"MatFreeIntermediateDataStructures","Mat");CHKERRQ(ierr);
2057d0a89b7SHong Zhang       ap->freestruct = PETSC_FALSE;
2067d0a89b7SHong Zhang       ierr = PetscOptionsBool("-mat_freeintermediatedatastructures","Free intermediate data structures", "MatFreeIntermediateDataStructures",ap->freestruct,&ap->freestruct, NULL);CHKERRQ(ierr);
2077d0a89b7SHong Zhang       ierr = PetscOptionsEnd();CHKERRQ(ierr);
2087a7894deSKris Buschelman     }
2097d0a89b7SHong Zhang   }
2107d0a89b7SHong Zhang 
2113ff4c91cSHong Zhang   ierr = PetscLogEventBegin(MAT_PtAPNumeric,A,P,0,0);CHKERRQ(ierr);
2128403a639SHong Zhang   ierr = (*(*C)->ops->ptapnumeric)(A,P,*C);CHKERRQ(ierr);
2133ff4c91cSHong Zhang   ierr = PetscLogEventEnd(MAT_PtAPNumeric,A,P,0,0);CHKERRQ(ierr);
21442c57489SHong Zhang   PetscFunctionReturn(0);
21542c57489SHong Zhang }
21642c57489SHong Zhang 
217cf742fccSHong Zhang PetscErrorCode MatPtAPNumeric_MPIAIJ_MPIAIJ_scalable(Mat A,Mat P,Mat C)
218cf742fccSHong Zhang {
219cf742fccSHong Zhang   PetscErrorCode    ierr;
220cf742fccSHong Zhang   Mat_MPIAIJ        *a=(Mat_MPIAIJ*)A->data,*p=(Mat_MPIAIJ*)P->data,*c=(Mat_MPIAIJ*)C->data;
221cf742fccSHong Zhang   Mat_SeqAIJ        *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data;
222cf742fccSHong Zhang   Mat_SeqAIJ        *ap,*p_loc,*p_oth,*c_seq;
2233cdca5ebSHong Zhang   Mat_APMPI         *ptap = c->ap;
224cf742fccSHong Zhang   Mat               AP_loc,C_loc,C_oth;
225a3bb6f32SFande Kong   PetscInt          i,rstart,rend,cm,ncols,row,*api,*apj,am = A->rmap->n,apnz,nout;
226cf742fccSHong Zhang   PetscScalar       *apa;
227cf742fccSHong Zhang   const PetscInt    *cols;
228cf742fccSHong Zhang   const PetscScalar *vals;
229cf742fccSHong Zhang 
230cf742fccSHong Zhang   PetscFunctionBegin;
23180da8df7SHong Zhang   if (!ptap) {
23280da8df7SHong Zhang     MPI_Comm comm;
23380da8df7SHong Zhang     ierr = PetscObjectGetComm((PetscObject)C,&comm);CHKERRQ(ierr);
23480da8df7SHong Zhang     SETERRQ(comm,PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be reused. Do not call MatFreeIntermediateDataStructures() or use '-mat_freeintermediatedatastructures'");
23580da8df7SHong Zhang   }
23680da8df7SHong Zhang 
237cf742fccSHong Zhang   ierr = MatZeroEntries(C);CHKERRQ(ierr);
238cf742fccSHong Zhang 
239cf742fccSHong Zhang   /* 1) get R = Pd^T,Ro = Po^T */
240cf742fccSHong Zhang   if (ptap->reuse == MAT_REUSE_MATRIX) {
241419ecdd9Sandi selinger     ierr = MatTranspose(p->A,MAT_REUSE_MATRIX,&ptap->Rd);CHKERRQ(ierr);
242419ecdd9Sandi selinger     ierr = MatTranspose(p->B,MAT_REUSE_MATRIX,&ptap->Ro);CHKERRQ(ierr);
243cf742fccSHong Zhang   }
244cf742fccSHong Zhang 
245cf742fccSHong Zhang   /* 2) get AP_loc */
246cf742fccSHong Zhang   AP_loc = ptap->AP_loc;
247cf742fccSHong Zhang   ap = (Mat_SeqAIJ*)AP_loc->data;
248cf742fccSHong Zhang 
249cf742fccSHong Zhang   /* 2-1) get P_oth = ptap->P_oth  and P_loc = ptap->P_loc */
250cf742fccSHong Zhang   /*-----------------------------------------------------*/
251cf742fccSHong Zhang   if (ptap->reuse == MAT_REUSE_MATRIX) {
252cf742fccSHong Zhang     /* P_oth and P_loc are obtained in MatPtASymbolic() when reuse == MAT_INITIAL_MATRIX */
253cf742fccSHong Zhang     ierr = MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_REUSE_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth);CHKERRQ(ierr);
254cf742fccSHong Zhang     ierr = MatMPIAIJGetLocalMat(P,MAT_REUSE_MATRIX,&ptap->P_loc);CHKERRQ(ierr);
255cf742fccSHong Zhang   }
256cf742fccSHong Zhang 
257cf742fccSHong Zhang   /* 2-2) compute numeric A_loc*P - dominating part */
258cf742fccSHong Zhang   /* ---------------------------------------------- */
259cf742fccSHong Zhang   /* get data from symbolic products */
260cf742fccSHong Zhang   p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data;
26152f7967eSHong Zhang   if (ptap->P_oth) p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data;
262c072d3e2SSatish Balay   else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"ptap->P_oth is NULL. Cannot proceed!");
26352f7967eSHong Zhang 
264cf742fccSHong Zhang   api   = ap->i;
265cf742fccSHong Zhang   apj   = ap->j;
266a3bb6f32SFande Kong   ierr = ISLocalToGlobalMappingApply(ptap->ltog,api[AP_loc->rmap->n],apj,apj);CHKERRQ(ierr);
267cf742fccSHong Zhang   for (i=0; i<am; i++) {
268cf742fccSHong Zhang     /* AP[i,:] = A[i,:]*P = Ad*P_loc Ao*P_oth */
269cf742fccSHong Zhang     apnz = api[i+1] - api[i];
270b4e8d1b6SHong Zhang     apa = ap->a + api[i];
271b4e8d1b6SHong Zhang     ierr = PetscMemzero(apa,sizeof(PetscScalar)*apnz);CHKERRQ(ierr);
272b4e8d1b6SHong Zhang     AProw_scalable(i,ad,ao,p_loc,p_oth,api,apj,apa);
273cf742fccSHong Zhang     ierr = PetscLogFlops(2.0*apnz);CHKERRQ(ierr);
274cf742fccSHong Zhang   }
275a3bb6f32SFande Kong   ierr = ISGlobalToLocalMappingApply(ptap->ltog,IS_GTOLM_DROP,api[AP_loc->rmap->n],apj,&nout,apj);CHKERRQ(ierr);
276a3bb6f32SFande Kong   if (api[AP_loc->rmap->n] != nout) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Incorrect mapping %D != %D\n",api[AP_loc->rmap->n],nout);
277cf742fccSHong Zhang 
278cf742fccSHong Zhang   /* 3) C_loc = Rd*AP_loc, C_oth = Ro*AP_loc */
279154d0d78SFande Kong   /* Always use scalable version since we are in the MPI scalable version */
280cf742fccSHong Zhang   ierr = MatMatMultNumeric_SeqAIJ_SeqAIJ_Scalable(ptap->Rd,AP_loc,ptap->C_loc);CHKERRQ(ierr);
281cf742fccSHong Zhang   ierr = MatMatMultNumeric_SeqAIJ_SeqAIJ_Scalable(ptap->Ro,AP_loc,ptap->C_oth);CHKERRQ(ierr);
282cf742fccSHong Zhang 
283cf742fccSHong Zhang   C_loc = ptap->C_loc;
284cf742fccSHong Zhang   C_oth = ptap->C_oth;
285cf742fccSHong Zhang 
286cf742fccSHong Zhang   /* add C_loc and Co to to C */
287cf742fccSHong Zhang   ierr = MatGetOwnershipRange(C,&rstart,&rend);CHKERRQ(ierr);
288cf742fccSHong Zhang 
289cf742fccSHong Zhang   /* C_loc -> C */
290cf742fccSHong Zhang   cm    = C_loc->rmap->N;
291cf742fccSHong Zhang   c_seq = (Mat_SeqAIJ*)C_loc->data;
292cf742fccSHong Zhang   cols = c_seq->j;
293cf742fccSHong Zhang   vals = c_seq->a;
294a3bb6f32SFande Kong   ierr = ISLocalToGlobalMappingApply(ptap->ltog,c_seq->i[C_loc->rmap->n],c_seq->j,c_seq->j);CHKERRQ(ierr);
295edeac6deSandi selinger 
296e9ede7d0Sandi selinger   /* The (fast) MatSetValues_MPIAIJ_CopyFromCSRFormat function can only be used when C->was_assembled is PETSC_FALSE and */
297edeac6deSandi selinger   /* when there are no off-processor parts.  */
2981de21080Sandi selinger   /* If was_assembled is true, then the statement aj[rowstart_diag+dnz_row] = mat_j[col] - cstart; in MatSetValues_MPIAIJ_CopyFromCSRFormat */
2991de21080Sandi selinger   /* is no longer true. Then the more complex function MatSetValues_MPIAIJ() has to be used, where the column index is looked up from */
3001de21080Sandi selinger   /* a table, and other, more complex stuff has to be done. */
301edeac6deSandi selinger   if (C->assembled) {
302edeac6deSandi selinger     C->was_assembled = PETSC_TRUE;
303edeac6deSandi selinger     C->assembled     = PETSC_FALSE;
304edeac6deSandi selinger   }
305edeac6deSandi selinger   if (C->was_assembled) {
306cf742fccSHong Zhang     for (i=0; i<cm; i++) {
307cf742fccSHong Zhang       ncols = c_seq->i[i+1] - c_seq->i[i];
308cf742fccSHong Zhang       row = rstart + i;
309edeac6deSandi selinger       ierr = MatSetValues_MPIAIJ(C,1,&row,ncols,cols,vals,ADD_VALUES);CHKERRQ(ierr);
310cf742fccSHong Zhang       cols += ncols; vals += ncols;
311cf742fccSHong Zhang     }
312edeac6deSandi selinger   } else {
313e9ede7d0Sandi selinger     ierr = MatSetValues_MPIAIJ_CopyFromCSRFormat(C,c_seq->j,c_seq->i,c_seq->a);CHKERRQ(ierr);
314edeac6deSandi selinger   }
315a3bb6f32SFande Kong   ierr = ISGlobalToLocalMappingApply(ptap->ltog,IS_GTOLM_DROP,c_seq->i[C_loc->rmap->n],c_seq->j,&nout,c_seq->j);CHKERRQ(ierr);
316a3bb6f32SFande Kong   if (c_seq->i[C_loc->rmap->n] != nout) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Incorrect mapping %D != %D\n",c_seq->i[C_loc->rmap->n],nout);
317cf742fccSHong Zhang 
318cf742fccSHong Zhang   /* Co -> C, off-processor part */
319cf742fccSHong Zhang   cm = C_oth->rmap->N;
320cf742fccSHong Zhang   c_seq = (Mat_SeqAIJ*)C_oth->data;
321cf742fccSHong Zhang   cols = c_seq->j;
322cf742fccSHong Zhang   vals = c_seq->a;
323a3bb6f32SFande Kong   ierr = ISLocalToGlobalMappingApply(ptap->ltog,c_seq->i[C_oth->rmap->n],c_seq->j,c_seq->j);CHKERRQ(ierr);
324cf742fccSHong Zhang   for (i=0; i<cm; i++) {
325cf742fccSHong Zhang     ncols = c_seq->i[i+1] - c_seq->i[i];
326cf742fccSHong Zhang     row = p->garray[i];
327cf742fccSHong Zhang     ierr = MatSetValues(C,1,&row,ncols,cols,vals,ADD_VALUES);CHKERRQ(ierr);
328cf742fccSHong Zhang     cols += ncols; vals += ncols;
329cf742fccSHong Zhang   }
330cf742fccSHong Zhang   ierr = MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
331cf742fccSHong Zhang   ierr = MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
332cf742fccSHong Zhang 
333cf742fccSHong Zhang   ptap->reuse = MAT_REUSE_MATRIX;
33480da8df7SHong Zhang 
335a3bb6f32SFande Kong   ierr = ISGlobalToLocalMappingApply(ptap->ltog,IS_GTOLM_DROP,c_seq->i[C_oth->rmap->n],c_seq->j,&nout,c_seq->j);CHKERRQ(ierr);
336a3bb6f32SFande Kong   if (c_seq->i[C_oth->rmap->n] != nout) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Incorrect mapping %D != %D\n",c_seq->i[C_loc->rmap->n],nout);
337a3bb6f32SFande Kong 
33880da8df7SHong Zhang   /* supporting struct ptap consumes almost same amount of memory as C=PtAP, release it if C will not be updated by A and P */
3397d0a89b7SHong Zhang   if (ptap->freestruct) {
34080da8df7SHong Zhang     ierr = MatFreeIntermediateDataStructures(C);CHKERRQ(ierr);
34180da8df7SHong Zhang   }
342cf742fccSHong Zhang   PetscFunctionReturn(0);
343cf742fccSHong Zhang }
344cf742fccSHong Zhang 
345e953e47cSHong Zhang PetscErrorCode MatPtAPSymbolic_MPIAIJ_MPIAIJ_scalable(Mat A,Mat P,PetscReal fill,Mat *C)
3460d3441aeSHong Zhang {
3470d3441aeSHong Zhang   PetscErrorCode      ierr;
3483cdca5ebSHong Zhang   Mat_APMPI           *ptap;
349681d504bSHong Zhang   Mat_MPIAIJ          *a=(Mat_MPIAIJ*)A->data,*p=(Mat_MPIAIJ*)P->data,*c;
3502259aa2eSHong Zhang   MPI_Comm            comm;
3512259aa2eSHong Zhang   PetscMPIInt         size,rank;
35276db11f6SHong Zhang   Mat                 Cmpi,P_loc,P_oth;
35315a3b8e2SHong Zhang   PetscFreeSpaceList  free_space=NULL,current_space=NULL;
354d0e9a2f7SHong Zhang   PetscInt            am=A->rmap->n,pm=P->rmap->n,pN=P->cmap->N,pn=P->cmap->n;
355d0e9a2f7SHong Zhang   PetscInt            *lnk,i,k,pnz,row,nsend;
356f671be37SHong Zhang   PetscMPIInt         tagi,tagj,*len_si,*len_s,*len_ri,icompleted=0,nrecv;
35715a3b8e2SHong Zhang   PetscInt            **buf_rj,**buf_ri,**buf_ri_k;
358d0e9a2f7SHong Zhang   PetscInt            len,proc,*dnz,*onz,*owners,nzi,nspacedouble;
35915a3b8e2SHong Zhang   PetscInt            nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci;
36015a3b8e2SHong Zhang   MPI_Request         *swaits,*rwaits;
36115a3b8e2SHong Zhang   MPI_Status          *sstatus,rstatus;
362445158ffSHong Zhang   PetscLayout         rowmap;
363445158ffSHong Zhang   PetscInt            *owners_co,*coi,*coj;    /* i and j array of (p->B)^T*A*P - used in the communication */
364445158ffSHong Zhang   PetscMPIInt         *len_r,*id_r;    /* array of length of comm->size, store send/recv matrix values */
365a3bb6f32SFande Kong   PetscInt            *api,*apj,*Jptr,apnz,*prmap=p->garray,con,j,Crmax,*aj,*ai,*pi,nout;
36652f7967eSHong Zhang   Mat_SeqAIJ          *p_loc,*p_oth=NULL,*ad=(Mat_SeqAIJ*)(a->A)->data,*ao=NULL,*c_loc,*c_oth;
36767a12041SHong Zhang   PetscScalar         *apv;
36878d30b94SHong Zhang   PetscTable          ta;
369b92f168fSBarry Smith   MatType             mtype;
370e83e5f86SFande Kong   const char          *prefix;
371aa690a28SHong Zhang #if defined(PETSC_USE_INFO)
3728cb82516SHong Zhang   PetscReal           apfill;
373aa690a28SHong Zhang #endif
37467a12041SHong Zhang 
37567a12041SHong Zhang   PetscFunctionBegin;
37667a12041SHong Zhang   ierr = PetscObjectGetComm((PetscObject)A,&comm);CHKERRQ(ierr);
37767a12041SHong Zhang   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
37867a12041SHong Zhang   ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
379ae5f0867Sstefano_zampini 
38052f7967eSHong Zhang   if (size > 1) ao = (Mat_SeqAIJ*)(a->B)->data;
38152f7967eSHong Zhang 
382ae5f0867Sstefano_zampini   /* create symbolic parallel matrix Cmpi */
383ae5f0867Sstefano_zampini   ierr = MatCreate(comm,&Cmpi);CHKERRQ(ierr);
384b92f168fSBarry Smith   ierr = MatGetType(A,&mtype);CHKERRQ(ierr);
385b92f168fSBarry Smith   ierr = MatSetType(Cmpi,mtype);CHKERRQ(ierr);
386ae5f0867Sstefano_zampini 
3873cdca5ebSHong Zhang   /* create struct Mat_APMPI and attached it to C later */
388e953e47cSHong Zhang   ierr        = PetscNew(&ptap);CHKERRQ(ierr);
389e953e47cSHong Zhang   ptap->reuse = MAT_INITIAL_MATRIX;
390cf97cf3cSHong Zhang   ptap->algType = 0;
391e953e47cSHong Zhang 
392e953e47cSHong Zhang   /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */
39376db11f6SHong Zhang   ierr = MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&P_oth);CHKERRQ(ierr);
394e953e47cSHong Zhang   /* get P_loc by taking all local rows of P */
39576db11f6SHong Zhang   ierr = MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&P_loc);CHKERRQ(ierr);
39676db11f6SHong Zhang 
39776db11f6SHong Zhang   ptap->P_loc = P_loc;
39876db11f6SHong Zhang   ptap->P_oth = P_oth;
399e953e47cSHong Zhang 
400e953e47cSHong Zhang   /* (0) compute Rd = Pd^T, Ro = Po^T  */
401e953e47cSHong Zhang   /* --------------------------------- */
402419ecdd9Sandi selinger   ierr = MatTranspose(p->A,MAT_INITIAL_MATRIX,&ptap->Rd);CHKERRQ(ierr);
403419ecdd9Sandi selinger   ierr = MatTranspose(p->B,MAT_INITIAL_MATRIX,&ptap->Ro);CHKERRQ(ierr);
404e953e47cSHong Zhang 
405e953e47cSHong Zhang   /* (1) compute symbolic AP = A_loc*P = Ad*P_loc + Ao*P_oth (api,apj) */
406e953e47cSHong Zhang   /* ----------------------------------------------------------------- */
40776db11f6SHong Zhang   p_loc  = (Mat_SeqAIJ*)P_loc->data;
40852f7967eSHong Zhang   if (P_oth) p_oth = (Mat_SeqAIJ*)P_oth->data;
409e953e47cSHong Zhang 
410e953e47cSHong Zhang   /* create and initialize a linked list */
411e953e47cSHong Zhang   ierr = PetscTableCreate(pn,pN,&ta);CHKERRQ(ierr); /* for compute AP_loc and Cmpi */
41276db11f6SHong Zhang   MatRowMergeMax_SeqAIJ(p_loc,P_loc->rmap->N,ta);
41376db11f6SHong Zhang   MatRowMergeMax_SeqAIJ(p_oth,P_oth->rmap->N,ta);
414e953e47cSHong Zhang   ierr = PetscTableGetCount(ta,&Crmax);CHKERRQ(ierr); /* Crmax = nnz(sum of Prows) */
415e953e47cSHong Zhang 
41676db11f6SHong Zhang   ierr = PetscLLCondensedCreate_Scalable(Crmax,&lnk);CHKERRQ(ierr);
417e953e47cSHong Zhang 
418e953e47cSHong Zhang   /* Initial FreeSpace size is fill*(nnz(A) + nnz(P)) */
41952f7967eSHong Zhang   if (ao) {
420e953e47cSHong Zhang     ierr = PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(ad->i[am],PetscIntSumTruncate(ao->i[am],p_loc->i[pm]))),&free_space);CHKERRQ(ierr);
42152f7967eSHong Zhang   } else {
42252f7967eSHong Zhang     ierr = PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(ad->i[am],p_loc->i[pm])),&free_space);CHKERRQ(ierr);
42352f7967eSHong Zhang   }
424e953e47cSHong Zhang   current_space = free_space;
425e953e47cSHong Zhang   nspacedouble  = 0;
426e953e47cSHong Zhang 
427e953e47cSHong Zhang   ierr   = PetscMalloc1(am+1,&api);CHKERRQ(ierr);
428e953e47cSHong Zhang   api[0] = 0;
429e953e47cSHong Zhang   for (i=0; i<am; i++) {
430e953e47cSHong Zhang     /* diagonal portion: Ad[i,:]*P */
431e953e47cSHong Zhang     ai = ad->i; pi = p_loc->i;
432e953e47cSHong Zhang     nzi = ai[i+1] - ai[i];
433e953e47cSHong Zhang     aj  = ad->j + ai[i];
434e953e47cSHong Zhang     for (j=0; j<nzi; j++) {
435e953e47cSHong Zhang       row  = aj[j];
436e953e47cSHong Zhang       pnz  = pi[row+1] - pi[row];
437e953e47cSHong Zhang       Jptr = p_loc->j + pi[row];
438e953e47cSHong Zhang       /* add non-zero cols of P into the sorted linked list lnk */
43976db11f6SHong Zhang       ierr = PetscLLCondensedAddSorted_Scalable(pnz,Jptr,lnk);CHKERRQ(ierr);
440e953e47cSHong Zhang     }
441e953e47cSHong Zhang     /* off-diagonal portion: Ao[i,:]*P */
44252f7967eSHong Zhang     if (ao) {
443e953e47cSHong Zhang       ai = ao->i; pi = p_oth->i;
444e953e47cSHong Zhang       nzi = ai[i+1] - ai[i];
445e953e47cSHong Zhang       aj  = ao->j + ai[i];
446e953e47cSHong Zhang       for (j=0; j<nzi; j++) {
447e953e47cSHong Zhang         row  = aj[j];
448e953e47cSHong Zhang         pnz  = pi[row+1] - pi[row];
449e953e47cSHong Zhang         Jptr = p_oth->j + pi[row];
45076db11f6SHong Zhang         ierr = PetscLLCondensedAddSorted_Scalable(pnz,Jptr,lnk);CHKERRQ(ierr);
451e953e47cSHong Zhang       }
45252f7967eSHong Zhang     }
453e953e47cSHong Zhang     apnz     = lnk[0];
454e953e47cSHong Zhang     api[i+1] = api[i] + apnz;
455e953e47cSHong Zhang 
456e953e47cSHong Zhang     /* if free space is not available, double the total space in the list */
457e953e47cSHong Zhang     if (current_space->local_remaining<apnz) {
458e953e47cSHong Zhang       ierr = PetscFreeSpaceGet(PetscIntSumTruncate(apnz,current_space->total_array_size),&current_space);CHKERRQ(ierr);
459e953e47cSHong Zhang       nspacedouble++;
460e953e47cSHong Zhang     }
461e953e47cSHong Zhang 
462e953e47cSHong Zhang     /* Copy data into free space, then initialize lnk */
46376db11f6SHong Zhang     ierr = PetscLLCondensedClean_Scalable(apnz,current_space->array,lnk);CHKERRQ(ierr);
464e953e47cSHong Zhang 
465e953e47cSHong Zhang     current_space->array           += apnz;
466e953e47cSHong Zhang     current_space->local_used      += apnz;
467e953e47cSHong Zhang     current_space->local_remaining -= apnz;
468e953e47cSHong Zhang   }
469e953e47cSHong Zhang   /* Allocate space for apj and apv, initialize apj, and */
470e953e47cSHong Zhang   /* destroy list of free space and other temporary array(s) */
471a3bb6f32SFande Kong   ierr = PetscCalloc2(api[am],&apj,api[am],&apv);CHKERRQ(ierr);
472e953e47cSHong Zhang   ierr = PetscFreeSpaceContiguous(&free_space,apj);CHKERRQ(ierr);
47376db11f6SHong Zhang   ierr = PetscLLCondensedDestroy_Scalable(lnk);CHKERRQ(ierr);
474e953e47cSHong Zhang 
475e953e47cSHong Zhang   /* Create AP_loc for reuse */
476e953e47cSHong Zhang   ierr = MatCreateSeqAIJWithArrays(PETSC_COMM_SELF,am,pN,api,apj,apv,&ptap->AP_loc);CHKERRQ(ierr);
477a3bb6f32SFande Kong   ierr = MatSeqAIJCompactOutExtraColumns_SeqAIJ(ptap->AP_loc, &ptap->ltog);CHKERRQ(ierr);
478e953e47cSHong Zhang 
479e953e47cSHong Zhang #if defined(PETSC_USE_INFO)
48052f7967eSHong Zhang   if (ao) {
481e953e47cSHong Zhang     apfill = (PetscReal)api[am]/(ad->i[am]+ao->i[am]+p_loc->i[pm]+1);
48252f7967eSHong Zhang   } else {
48352f7967eSHong Zhang     apfill = (PetscReal)api[am]/(ad->i[am]+p_loc->i[pm]+1);
48452f7967eSHong Zhang   }
485e953e47cSHong Zhang   ptap->AP_loc->info.mallocs           = nspacedouble;
486e953e47cSHong Zhang   ptap->AP_loc->info.fill_ratio_given  = fill;
487e953e47cSHong Zhang   ptap->AP_loc->info.fill_ratio_needed = apfill;
488e953e47cSHong Zhang 
489e953e47cSHong Zhang   if (api[am]) {
490b11c1ec8SHong Zhang     ierr = PetscInfo3(ptap->AP_loc,"Scalable algorithm, AP_loc reallocs %D; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)apfill);CHKERRQ(ierr);
491e953e47cSHong Zhang     ierr = PetscInfo1(ptap->AP_loc,"Use MatPtAP(A,B,MatReuse,%g,&C) for best AP_loc performance.;\n",(double)apfill);CHKERRQ(ierr);
492e953e47cSHong Zhang   } else {
493b11c1ec8SHong Zhang     ierr = PetscInfo(ptap->AP_loc,"Scalable algorithm, AP_loc is empty \n");CHKERRQ(ierr);
494e953e47cSHong Zhang   }
495e953e47cSHong Zhang #endif
496e953e47cSHong Zhang 
497e953e47cSHong Zhang   /* (2-1) compute symbolic Co = Ro*AP_loc  */
498e953e47cSHong Zhang   /* ------------------------------------ */
499e83e5f86SFande Kong   ierr = MatGetOptionsPrefix(A,&prefix);CHKERRQ(ierr);
500e83e5f86SFande Kong   ierr = MatSetOptionsPrefix(ptap->Ro,prefix);CHKERRQ(ierr);
501e83e5f86SFande Kong   ierr = MatAppendOptionsPrefix(ptap->Ro,"inner_offdiag_");CHKERRQ(ierr);
502e953e47cSHong Zhang   ierr = MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Ro,ptap->AP_loc,fill,&ptap->C_oth);CHKERRQ(ierr);
503e953e47cSHong Zhang 
504e953e47cSHong Zhang   /* (3) send coj of C_oth to other processors  */
505e953e47cSHong Zhang   /* ------------------------------------------ */
506e953e47cSHong Zhang   /* determine row ownership */
507e953e47cSHong Zhang   ierr = PetscLayoutCreate(comm,&rowmap);CHKERRQ(ierr);
508e953e47cSHong Zhang   rowmap->n  = pn;
509e953e47cSHong Zhang   rowmap->bs = 1;
510e953e47cSHong Zhang   ierr   = PetscLayoutSetUp(rowmap);CHKERRQ(ierr);
511e953e47cSHong Zhang   owners = rowmap->range;
512e953e47cSHong Zhang 
513e953e47cSHong Zhang   /* determine the number of messages to send, their lengths */
514e953e47cSHong Zhang   ierr = PetscMalloc4(size,&len_s,size,&len_si,size,&sstatus,size+2,&owners_co);CHKERRQ(ierr);
515e953e47cSHong Zhang   ierr = PetscMemzero(len_s,size*sizeof(PetscMPIInt));CHKERRQ(ierr);
516e953e47cSHong Zhang   ierr = PetscMemzero(len_si,size*sizeof(PetscMPIInt));CHKERRQ(ierr);
517e953e47cSHong Zhang 
518e953e47cSHong Zhang   c_oth = (Mat_SeqAIJ*)ptap->C_oth->data;
519e953e47cSHong Zhang   coi   = c_oth->i; coj = c_oth->j;
520e953e47cSHong Zhang   con   = ptap->C_oth->rmap->n;
521e953e47cSHong Zhang   proc  = 0;
522a3bb6f32SFande Kong   ierr = ISLocalToGlobalMappingApply(ptap->ltog,coi[con],coj,coj);CHKERRQ(ierr);
523e953e47cSHong Zhang   for (i=0; i<con; i++) {
524e953e47cSHong Zhang     while (prmap[i] >= owners[proc+1]) proc++;
525e953e47cSHong Zhang     len_si[proc]++;               /* num of rows in Co(=Pt*AP) to be sent to [proc] */
526e953e47cSHong Zhang     len_s[proc] += coi[i+1] - coi[i]; /* num of nonzeros in Co to be sent to [proc] */
527e953e47cSHong Zhang   }
528e953e47cSHong Zhang 
529e953e47cSHong Zhang   len          = 0; /* max length of buf_si[], see (4) */
530e953e47cSHong Zhang   owners_co[0] = 0;
531e953e47cSHong Zhang   nsend        = 0;
532e953e47cSHong Zhang   for (proc=0; proc<size; proc++) {
533e953e47cSHong Zhang     owners_co[proc+1] = owners_co[proc] + len_si[proc];
534e953e47cSHong Zhang     if (len_s[proc]) {
535e953e47cSHong Zhang       nsend++;
536e953e47cSHong Zhang       len_si[proc] = 2*(len_si[proc] + 1); /* length of buf_si to be sent to [proc] */
537e953e47cSHong Zhang       len         += len_si[proc];
538e953e47cSHong Zhang     }
539e953e47cSHong Zhang   }
540e953e47cSHong Zhang 
541e953e47cSHong Zhang   /* determine the number and length of messages to receive for coi and coj  */
542e953e47cSHong Zhang   ierr = PetscGatherNumberOfMessages(comm,NULL,len_s,&nrecv);CHKERRQ(ierr);
543e953e47cSHong Zhang   ierr = PetscGatherMessageLengths2(comm,nsend,nrecv,len_s,len_si,&id_r,&len_r,&len_ri);CHKERRQ(ierr);
544e953e47cSHong Zhang 
545e953e47cSHong Zhang   /* post the Irecv and Isend of coj */
546e953e47cSHong Zhang   ierr = PetscCommGetNewTag(comm,&tagj);CHKERRQ(ierr);
547e953e47cSHong Zhang   ierr = PetscPostIrecvInt(comm,tagj,nrecv,id_r,len_r,&buf_rj,&rwaits);CHKERRQ(ierr);
548e953e47cSHong Zhang   ierr = PetscMalloc1(nsend+1,&swaits);CHKERRQ(ierr);
549e953e47cSHong Zhang   for (proc=0, k=0; proc<size; proc++) {
550e953e47cSHong Zhang     if (!len_s[proc]) continue;
551e953e47cSHong Zhang     i    = owners_co[proc];
552e953e47cSHong Zhang     ierr = MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tagj,comm,swaits+k);CHKERRQ(ierr);
553e953e47cSHong Zhang     k++;
554e953e47cSHong Zhang   }
555e953e47cSHong Zhang 
556e953e47cSHong Zhang   /* (2-2) compute symbolic C_loc = Rd*AP_loc */
557e953e47cSHong Zhang   /* ---------------------------------------- */
558e83e5f86SFande Kong   ierr = MatSetOptionsPrefix(ptap->Rd,prefix);CHKERRQ(ierr);
559e83e5f86SFande Kong   ierr = MatAppendOptionsPrefix(ptap->Rd,"inner_diag_");CHKERRQ(ierr);
560e953e47cSHong Zhang   ierr = MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Rd,ptap->AP_loc,fill,&ptap->C_loc);CHKERRQ(ierr);
561e953e47cSHong Zhang   c_loc = (Mat_SeqAIJ*)ptap->C_loc->data;
562a3bb6f32SFande Kong   ierr = ISLocalToGlobalMappingApply(ptap->ltog,c_loc->i[ptap->C_loc->rmap->n],c_loc->j,c_loc->j);CHKERRQ(ierr);
563e953e47cSHong Zhang 
564e953e47cSHong Zhang   /* receives coj are complete */
565e953e47cSHong Zhang   for (i=0; i<nrecv; i++) {
566e953e47cSHong Zhang     ierr = MPI_Waitany(nrecv,rwaits,&icompleted,&rstatus);CHKERRQ(ierr);
567e953e47cSHong Zhang   }
568e953e47cSHong Zhang   ierr = PetscFree(rwaits);CHKERRQ(ierr);
569e953e47cSHong Zhang   if (nsend) {ierr = MPI_Waitall(nsend,swaits,sstatus);CHKERRQ(ierr);}
570e953e47cSHong Zhang 
571e953e47cSHong Zhang   /* add received column indices into ta to update Crmax */
572e953e47cSHong Zhang   for (k=0; k<nrecv; k++) {/* k-th received message */
573e953e47cSHong Zhang     Jptr = buf_rj[k];
574e953e47cSHong Zhang     for (j=0; j<len_r[k]; j++) {
575e953e47cSHong Zhang       ierr = PetscTableAdd(ta,*(Jptr+j)+1,1,INSERT_VALUES);CHKERRQ(ierr);
576e953e47cSHong Zhang     }
577e953e47cSHong Zhang   }
578e953e47cSHong Zhang   ierr = PetscTableGetCount(ta,&Crmax);CHKERRQ(ierr);
579e953e47cSHong Zhang   ierr = PetscTableDestroy(&ta);CHKERRQ(ierr);
580e953e47cSHong Zhang 
581e953e47cSHong Zhang   /* (4) send and recv coi */
582e953e47cSHong Zhang   /*-----------------------*/
583e953e47cSHong Zhang   ierr   = PetscCommGetNewTag(comm,&tagi);CHKERRQ(ierr);
584e953e47cSHong Zhang   ierr   = PetscPostIrecvInt(comm,tagi,nrecv,id_r,len_ri,&buf_ri,&rwaits);CHKERRQ(ierr);
585e953e47cSHong Zhang   ierr   = PetscMalloc1(len+1,&buf_s);CHKERRQ(ierr);
586e953e47cSHong Zhang   buf_si = buf_s;  /* points to the beginning of k-th msg to be sent */
587e953e47cSHong Zhang   for (proc=0,k=0; proc<size; proc++) {
588e953e47cSHong Zhang     if (!len_s[proc]) continue;
589e953e47cSHong Zhang     /* form outgoing message for i-structure:
590e953e47cSHong Zhang          buf_si[0]:                 nrows to be sent
591e953e47cSHong Zhang                [1:nrows]:           row index (global)
592e953e47cSHong Zhang                [nrows+1:2*nrows+1]: i-structure index
593e953e47cSHong Zhang     */
594e953e47cSHong Zhang     /*-------------------------------------------*/
595e953e47cSHong Zhang     nrows       = len_si[proc]/2 - 1; /* num of rows in Co to be sent to [proc] */
596e953e47cSHong Zhang     buf_si_i    = buf_si + nrows+1;
597e953e47cSHong Zhang     buf_si[0]   = nrows;
598e953e47cSHong Zhang     buf_si_i[0] = 0;
599e953e47cSHong Zhang     nrows       = 0;
600e953e47cSHong Zhang     for (i=owners_co[proc]; i<owners_co[proc+1]; i++) {
601e953e47cSHong Zhang       nzi = coi[i+1] - coi[i];
602e953e47cSHong Zhang       buf_si_i[nrows+1] = buf_si_i[nrows] + nzi;  /* i-structure */
603e953e47cSHong Zhang       buf_si[nrows+1]   = prmap[i] -owners[proc]; /* local row index */
604e953e47cSHong Zhang       nrows++;
605e953e47cSHong Zhang     }
606e953e47cSHong Zhang     ierr = MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tagi,comm,swaits+k);CHKERRQ(ierr);
607e953e47cSHong Zhang     k++;
608e953e47cSHong Zhang     buf_si += len_si[proc];
609e953e47cSHong Zhang   }
610e953e47cSHong Zhang   for (i=0; i<nrecv; i++) {
611e953e47cSHong Zhang     ierr = MPI_Waitany(nrecv,rwaits,&icompleted,&rstatus);CHKERRQ(ierr);
612e953e47cSHong Zhang   }
613e953e47cSHong Zhang   ierr = PetscFree(rwaits);CHKERRQ(ierr);
614e953e47cSHong Zhang   if (nsend) {ierr = MPI_Waitall(nsend,swaits,sstatus);CHKERRQ(ierr);}
615e953e47cSHong Zhang 
616e953e47cSHong Zhang   ierr = PetscFree4(len_s,len_si,sstatus,owners_co);CHKERRQ(ierr);
617e953e47cSHong Zhang   ierr = PetscFree(len_ri);CHKERRQ(ierr);
618e953e47cSHong Zhang   ierr = PetscFree(swaits);CHKERRQ(ierr);
619e953e47cSHong Zhang   ierr = PetscFree(buf_s);CHKERRQ(ierr);
620b4e8d1b6SHong Zhang 
621e953e47cSHong Zhang   /* (5) compute the local portion of Cmpi      */
622e953e47cSHong Zhang   /* ------------------------------------------ */
623e953e47cSHong Zhang   /* set initial free space to be Crmax, sufficient for holding nozeros in each row of Cmpi */
624e953e47cSHong Zhang   ierr          = PetscFreeSpaceGet(Crmax,&free_space);CHKERRQ(ierr);
625e953e47cSHong Zhang   current_space = free_space;
626e953e47cSHong Zhang 
627e953e47cSHong Zhang   ierr = PetscMalloc3(nrecv,&buf_ri_k,nrecv,&nextrow,nrecv,&nextci);CHKERRQ(ierr);
628e953e47cSHong Zhang   for (k=0; k<nrecv; k++) {
629e953e47cSHong Zhang     buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */
630e953e47cSHong Zhang     nrows       = *buf_ri_k[k];
631e953e47cSHong Zhang     nextrow[k]  = buf_ri_k[k] + 1;  /* next row number of k-th recved i-structure */
632e953e47cSHong Zhang     nextci[k]   = buf_ri_k[k] + (nrows + 1); /* poins to the next i-structure of k-th recved i-structure  */
633e953e47cSHong Zhang   }
634e953e47cSHong Zhang 
635e953e47cSHong Zhang   ierr = MatPreallocateInitialize(comm,pn,pn,dnz,onz);CHKERRQ(ierr);
636cf742fccSHong Zhang   ierr = PetscLLCondensedCreate_Scalable(Crmax,&lnk);CHKERRQ(ierr);
637e953e47cSHong Zhang   for (i=0; i<pn; i++) {
638e953e47cSHong Zhang     /* add C_loc into Cmpi */
639e953e47cSHong Zhang     nzi  = c_loc->i[i+1] - c_loc->i[i];
640e953e47cSHong Zhang     Jptr = c_loc->j + c_loc->i[i];
641cf742fccSHong Zhang     ierr = PetscLLCondensedAddSorted_Scalable(nzi,Jptr,lnk);CHKERRQ(ierr);
642e953e47cSHong Zhang 
643e953e47cSHong Zhang     /* add received col data into lnk */
644e953e47cSHong Zhang     for (k=0; k<nrecv; k++) { /* k-th received message */
645e953e47cSHong Zhang       if (i == *nextrow[k]) { /* i-th row */
646e953e47cSHong Zhang         nzi  = *(nextci[k]+1) - *nextci[k];
647e953e47cSHong Zhang         Jptr = buf_rj[k] + *nextci[k];
648cf742fccSHong Zhang         ierr = PetscLLCondensedAddSorted_Scalable(nzi,Jptr,lnk);CHKERRQ(ierr);
649e953e47cSHong Zhang         nextrow[k]++; nextci[k]++;
650e953e47cSHong Zhang       }
651e953e47cSHong Zhang     }
652e953e47cSHong Zhang     nzi = lnk[0];
653e953e47cSHong Zhang 
654e953e47cSHong Zhang     /* copy data into free space, then initialize lnk */
655cf742fccSHong Zhang     ierr = PetscLLCondensedClean_Scalable(nzi,current_space->array,lnk);CHKERRQ(ierr);
656e953e47cSHong Zhang     ierr = MatPreallocateSet(i+owners[rank],nzi,current_space->array,dnz,onz);CHKERRQ(ierr);
657e953e47cSHong Zhang   }
658e953e47cSHong Zhang   ierr = PetscFree3(buf_ri_k,nextrow,nextci);CHKERRQ(ierr);
659cf742fccSHong Zhang   ierr = PetscLLCondensedDestroy_Scalable(lnk);CHKERRQ(ierr);
660e953e47cSHong Zhang   ierr = PetscFreeSpaceDestroy(free_space);CHKERRQ(ierr);
661e953e47cSHong Zhang 
662e953e47cSHong Zhang   /* local sizes and preallocation */
663e953e47cSHong Zhang   ierr = MatSetSizes(Cmpi,pn,pn,PETSC_DETERMINE,PETSC_DETERMINE);CHKERRQ(ierr);
664e953e47cSHong Zhang   ierr = MatSetBlockSizes(Cmpi,PetscAbs(P->cmap->bs),PetscAbs(P->cmap->bs));CHKERRQ(ierr);
665e953e47cSHong Zhang   ierr = MatMPIAIJSetPreallocation(Cmpi,0,dnz,0,onz);CHKERRQ(ierr);
666e953e47cSHong Zhang   ierr = MatPreallocateFinalize(dnz,onz);CHKERRQ(ierr);
667e953e47cSHong Zhang 
668e953e47cSHong Zhang   /* members in merge */
669e953e47cSHong Zhang   ierr = PetscFree(id_r);CHKERRQ(ierr);
670e953e47cSHong Zhang   ierr = PetscFree(len_r);CHKERRQ(ierr);
671e953e47cSHong Zhang   ierr = PetscFree(buf_ri[0]);CHKERRQ(ierr);
672e953e47cSHong Zhang   ierr = PetscFree(buf_ri);CHKERRQ(ierr);
673e953e47cSHong Zhang   ierr = PetscFree(buf_rj[0]);CHKERRQ(ierr);
674e953e47cSHong Zhang   ierr = PetscFree(buf_rj);CHKERRQ(ierr);
675e953e47cSHong Zhang   ierr = PetscLayoutDestroy(&rowmap);CHKERRQ(ierr);
676e953e47cSHong Zhang 
677e953e47cSHong Zhang   /* attach the supporting struct to Cmpi for reuse */
678e953e47cSHong Zhang   c = (Mat_MPIAIJ*)Cmpi->data;
6793cdca5ebSHong Zhang   c->ap           = ptap;
680e953e47cSHong Zhang   ptap->duplicate = Cmpi->ops->duplicate;
681e953e47cSHong Zhang   ptap->destroy   = Cmpi->ops->destroy;
682cf97cf3cSHong Zhang   ptap->view      = Cmpi->ops->view;
683e953e47cSHong Zhang 
684e953e47cSHong Zhang   /* Cmpi is not ready for use - assembly will be done by MatPtAPNumeric() */
685e953e47cSHong Zhang   Cmpi->assembled        = PETSC_FALSE;
686a4ffb656SHong Zhang   Cmpi->ops->ptapnumeric = MatPtAPNumeric_MPIAIJ_MPIAIJ_scalable;
687e953e47cSHong Zhang   Cmpi->ops->destroy     = MatDestroy_MPIAIJ_PtAP;
688cf97cf3cSHong Zhang   Cmpi->ops->view        = MatView_MPIAIJ_PtAP;
6894624976aSHong Zhang   Cmpi->ops->freeintermediatedatastructures = MatFreeIntermediateDataStructures_MPIAIJ_AP;
690e953e47cSHong Zhang   *C                     = Cmpi;
691a3bb6f32SFande Kong 
692a3bb6f32SFande Kong    nout = 0;
693a3bb6f32SFande Kong    ierr = ISGlobalToLocalMappingApply(ptap->ltog,IS_GTOLM_DROP,c_oth->i[ptap->C_oth->rmap->n],c_oth->j,&nout,c_oth->j);CHKERRQ(ierr);
694a3bb6f32SFande Kong    if (c_oth->i[ptap->C_oth->rmap->n] != nout) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Incorrect mapping %D != %D\n",c_oth->i[ptap->C_oth->rmap->n],nout);
695a3bb6f32SFande Kong    ierr = ISGlobalToLocalMappingApply(ptap->ltog,IS_GTOLM_DROP,c_loc->i[ptap->C_loc->rmap->n],c_loc->j,&nout,c_loc->j);CHKERRQ(ierr);
696a3bb6f32SFande Kong    if (c_loc->i[ptap->C_loc->rmap->n] != nout) SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_INCOMP,"Incorrect mapping %D != %D\n",c_loc->i[ptap->C_loc->rmap->n],nout);
697a3bb6f32SFande Kong 
698e953e47cSHong Zhang   PetscFunctionReturn(0);
699e953e47cSHong Zhang }
700e953e47cSHong Zhang 
7014a56b808SFande Kong PETSC_STATIC_INLINE PetscErrorCode MatPtAPSymbolicComputeOneRowOfAP_private(Mat A,Mat P,Mat P_oth,PetscInt i,PetscHSetI dht,PetscHSetI oht)
7024a56b808SFande Kong {
7034a56b808SFande Kong   Mat_MPIAIJ           *a=(Mat_MPIAIJ*)A->data,*p=(Mat_MPIAIJ*)P->data;
7044a56b808SFande Kong   Mat_SeqAIJ           *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_oth=(Mat_SeqAIJ*)P_oth->data,*pd=(Mat_SeqAIJ*)p->A->data,*po=(Mat_SeqAIJ*)p->B->data;
7054a56b808SFande Kong   PetscInt             *ai,nzi,j,*aj,row,col,*pi,*pj,pnz,nzpi,*p_othcols,k;
7064a56b808SFande Kong   PetscInt             pcstart,pcend,column;
7074a56b808SFande Kong   PetscErrorCode       ierr;
7084a56b808SFande Kong 
7094a56b808SFande Kong   PetscFunctionBegin;
7104a56b808SFande Kong 
7114a56b808SFande Kong   pcstart = P->cmap->rstart;
7124a56b808SFande Kong   pcend   = P->cmap->rend;
7134a56b808SFande Kong   /* diagonal portion: Ad[i,:]*P */
7144a56b808SFande Kong   ai = ad->i;
7154a56b808SFande Kong   nzi = ai[i+1] - ai[i];
7164a56b808SFande Kong   aj  = ad->j + ai[i];
7174a56b808SFande Kong   for (j=0; j<nzi; j++) {
7184a56b808SFande Kong     row  = aj[j];
7194a56b808SFande Kong     nzpi = pd->i[row+1] - pd->i[row];
7204a56b808SFande Kong     pj  = pd->j + pd->i[row];
7214a56b808SFande Kong     for (k=0; k<nzpi; k++) {
7224a56b808SFande Kong       ierr = PetscHSetIAdd(dht,pj[k]+pcstart);CHKERRQ(ierr);
7234a56b808SFande Kong     }
7244a56b808SFande Kong   }
7254a56b808SFande Kong   for (j=0; j<nzi; j++) {
7264a56b808SFande Kong     row  = aj[j];
7274a56b808SFande Kong     nzpi = po->i[row+1] - po->i[row];
7284a56b808SFande Kong     pj  = po->j + po->i[row];
7294a56b808SFande Kong     for (k=0; k<nzpi; k++) {
7304a56b808SFande Kong       ierr = PetscHSetIAdd(oht,p->garray[pj[k]]);CHKERRQ(ierr);
7314a56b808SFande Kong     }
7324a56b808SFande Kong   }
7334a56b808SFande Kong 
7344a56b808SFande Kong   /* off diagonal part: Ao[i, :]*P_oth */
7354a56b808SFande Kong   if (ao) {
7364a56b808SFande Kong     ai = ao->i;
7374a56b808SFande Kong     pi = p_oth->i;
7384a56b808SFande Kong     nzi = ai[i+1] - ai[i];
7394a56b808SFande Kong     aj  = ao->j + ai[i];
7404a56b808SFande Kong     for (j=0; j<nzi; j++) {
7414a56b808SFande Kong       row  = aj[j];
7424a56b808SFande Kong       pnz  = pi[row+1] - pi[row];
7434a56b808SFande Kong       p_othcols = p_oth->j + pi[row];
7444a56b808SFande Kong       for (col=0; col<pnz; col++) {
7454a56b808SFande Kong         column = p_othcols[col];
7464a56b808SFande Kong         if (column>=pcstart && column<pcend) {
7474a56b808SFande Kong           ierr = PetscHSetIAdd(dht,column);CHKERRQ(ierr);
7484a56b808SFande Kong         } else {
7494a56b808SFande Kong           ierr = PetscHSetIAdd(oht,column);CHKERRQ(ierr);
7504a56b808SFande Kong         }
7514a56b808SFande Kong       }
7524a56b808SFande Kong     }
7534a56b808SFande Kong   } /* end if (ao) */
7544a56b808SFande Kong   PetscFunctionReturn(0);
7554a56b808SFande Kong }
7564a56b808SFande Kong 
7574a56b808SFande Kong PETSC_STATIC_INLINE PetscErrorCode MatPtAPNumericComputeOneRowOfAP_private(Mat A,Mat P,Mat P_oth,PetscInt i,PetscHMapIV hmap)
7584a56b808SFande Kong {
7594a56b808SFande Kong   Mat_MPIAIJ           *a=(Mat_MPIAIJ*)A->data,*p=(Mat_MPIAIJ*)P->data;
7604a56b808SFande Kong   Mat_SeqAIJ           *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_oth=(Mat_SeqAIJ*)P_oth->data,*pd=(Mat_SeqAIJ*)p->A->data,*po=(Mat_SeqAIJ*)p->B->data;
7614a56b808SFande Kong   PetscInt             *ai,nzi,j,*aj,row,col,*pi,pnz,*p_othcols,pcstart,*pj,k,nzpi;
7624a56b808SFande Kong   PetscScalar          ra,*aa,*pa;
7634a56b808SFande Kong   PetscErrorCode       ierr;
7644a56b808SFande Kong 
7654a56b808SFande Kong   PetscFunctionBegin;
7664a56b808SFande Kong   pcstart = P->cmap->rstart;
7674a56b808SFande Kong   /* diagonal portion: Ad[i,:]*P */
7684a56b808SFande Kong   ai  = ad->i;
7694a56b808SFande Kong   nzi = ai[i+1] - ai[i];
7704a56b808SFande Kong   aj  = ad->j + ai[i];
7714a56b808SFande Kong   aa  = ad->a + ai[i];
7724a56b808SFande Kong 
7734a56b808SFande Kong   for (j=0; j<nzi; j++) {
7744a56b808SFande Kong     ra   = aa[j];
7754a56b808SFande Kong     row  = aj[j];
7764a56b808SFande Kong     nzpi = pd->i[row+1] - pd->i[row];
7774a56b808SFande Kong     pj = pd->j + pd->i[row];
7784a56b808SFande Kong     pa = pd->a + pd->i[row];
7794a56b808SFande Kong     for (k=0; k<nzpi; k++) {
7804a56b808SFande Kong       ierr = PetscHMapIVAdd(hmap,pj[k]+pcstart,ra*pa[k]);CHKERRQ(ierr);
7814a56b808SFande Kong     }
7824a56b808SFande Kong   }
7834a56b808SFande Kong   for (j=0; j<nzi; j++) {
7844a56b808SFande Kong     ra   = aa[j];
7854a56b808SFande Kong     row  = aj[j];
7864a56b808SFande Kong     nzpi = po->i[row+1] - po->i[row];
7874a56b808SFande Kong     pj = po->j + po->i[row];
7884a56b808SFande Kong     pa = po->a + po->i[row];
7894a56b808SFande Kong     for (k=0; k<nzpi; k++) {
7904a56b808SFande Kong       ierr = PetscHMapIVAdd(hmap,p->garray[pj[k]],ra*pa[k]);CHKERRQ(ierr);
7914a56b808SFande Kong     }
7924a56b808SFande Kong   }
7934a56b808SFande Kong 
7944a56b808SFande Kong 
7954a56b808SFande Kong   /* off diagonal part: Ao[i, :]*P_oth */
7964a56b808SFande Kong   if (ao) {
7974a56b808SFande Kong     ai = ao->i;
7984a56b808SFande Kong     pi = p_oth->i;
7994a56b808SFande Kong     nzi = ai[i+1] - ai[i];
8004a56b808SFande Kong     aj  = ao->j + ai[i];
8014a56b808SFande Kong     aa  = ao->a + ai[i];
8024a56b808SFande Kong     for (j=0; j<nzi; j++) {
8034a56b808SFande Kong       row  = aj[j];
8044a56b808SFande Kong       ra   = aa[j];
8054a56b808SFande Kong       pnz  = pi[row+1] - pi[row];
8064a56b808SFande Kong       p_othcols = p_oth->j + pi[row];
8074a56b808SFande Kong       pa   = p_oth->a + pi[row];
8084a56b808SFande Kong       for (col=0; col<pnz; col++) {
8094a56b808SFande Kong         ierr = PetscHMapIVAdd(hmap,p_othcols[col],ra*pa[col]);CHKERRQ(ierr);
8104a56b808SFande Kong       }
8114a56b808SFande Kong     }
8124a56b808SFande Kong   } /* end if (ao) */
8134a56b808SFande Kong   PetscFunctionReturn(0);
8144a56b808SFande Kong }
8154a56b808SFande Kong 
8164a56b808SFande Kong PetscErrorCode MatPtAPNumeric_MPIAIJ_MPIAIJ_allatonce(Mat A,Mat P,Mat C)
8174a56b808SFande Kong {
8184a56b808SFande Kong   PetscErrorCode    ierr;
8194a56b808SFande Kong   Mat_MPIAIJ        *p=(Mat_MPIAIJ*)P->data,*c=(Mat_MPIAIJ*)C->data;
8204a56b808SFande Kong   Mat_SeqAIJ        *cd,*co,*po,*pd;
8214a56b808SFande Kong   Mat_APMPI         *ptap = c->ap;
8224a56b808SFande Kong   PetscHMapIV       hmap;
8234a56b808SFande Kong   PetscInt          i,j,jj,kk,nzi,*c_rmtj,voff,*c_othj,pn,pon,pcstart,pcend,ccstart,ccend,row,am,*poj,*pdj,*apindices,cmaxr,*c_rmtc,*c_rmtjj,*dcc,*occ,loc;
8244a56b808SFande Kong   PetscScalar       *c_rmta,*c_otha,*poa,*pda,*apvalues,*apvaluestmp,*c_rmtaa;
8254a56b808SFande Kong   MPI_Comm          comm;
8264a56b808SFande Kong 
8274a56b808SFande Kong   PetscFunctionBegin;
8284a56b808SFande Kong   ierr = PetscObjectGetComm((PetscObject)A,&comm);CHKERRQ(ierr);
8294a56b808SFande Kong   if (!ptap) SETERRQ(comm,PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be reused. Do not call MatFreeIntermediateDataStructures() or use '-mat_freeintermediatedatastructures'");
8304a56b808SFande Kong 
8314a56b808SFande Kong   ierr = MatZeroEntries(C);CHKERRQ(ierr);
8324a56b808SFande Kong 
8334a56b808SFande Kong   /* Get P_oth = ptap->P_oth  and P_loc = ptap->P_loc */
8344a56b808SFande Kong   /*-----------------------------------------------------*/
8354a56b808SFande Kong   if (ptap->reuse == MAT_REUSE_MATRIX) {
8364a56b808SFande Kong     /* P_oth and P_loc are obtained in MatPtASymbolic() when reuse == MAT_INITIAL_MATRIX */
8374a56b808SFande Kong     ierr = MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_REUSE_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth);CHKERRQ(ierr);
8384a56b808SFande Kong   }
8394a56b808SFande Kong 
8404a56b808SFande Kong   po = (Mat_SeqAIJ*) p->B->data;
8414a56b808SFande Kong   pd = (Mat_SeqAIJ*) p->A->data;
8424a56b808SFande Kong 
8434a56b808SFande Kong   ierr = MatGetLocalSize(p->B,NULL,&pon);CHKERRQ(ierr);
8444a56b808SFande Kong 
8454a56b808SFande Kong   ierr = PetscCalloc2(ptap->c_rmti[pon],&c_rmtj,ptap->c_rmti[pon],&c_rmta);CHKERRQ(ierr);
8464a56b808SFande Kong   ierr = MatGetLocalSize(A,&am,NULL);CHKERRQ(ierr);
8474a56b808SFande Kong   cmaxr = 0;
8484a56b808SFande Kong   for (i=0; i<pon; i++) {
8494a56b808SFande Kong     cmaxr = PetscMax(cmaxr,ptap->c_rmti[i+1]-ptap->c_rmti[i]);
8504a56b808SFande Kong   }
8514a56b808SFande Kong   ierr = PetscCalloc4(cmaxr,&apindices,cmaxr,&apvalues,cmaxr,&apvaluestmp,pon,&c_rmtc);CHKERRQ(ierr);
8524a56b808SFande Kong   ierr = PetscHMapIVCreate(&hmap);CHKERRQ(ierr);
8534a56b808SFande Kong   ierr = PetscHMapIVResize(hmap,cmaxr);CHKERRQ(ierr);
8544a56b808SFande Kong   for (i=0; i<am && pon; i++) {
8554a56b808SFande Kong     ierr = PetscHMapIVClear(hmap);CHKERRQ(ierr);
8564a56b808SFande Kong     nzi = po->i[i+1] - po->i[i];
8574a56b808SFande Kong     if (!nzi) continue;
8584a56b808SFande Kong     ierr = MatPtAPNumericComputeOneRowOfAP_private(A,P,ptap->P_oth,i,hmap);CHKERRQ(ierr);
8594a56b808SFande Kong     voff = 0;
8604a56b808SFande Kong     ierr = PetscHMapIVGetPairs(hmap,&voff,apindices,apvalues);CHKERRQ(ierr);
8614a56b808SFande Kong     if (!voff) continue;
8624a56b808SFande Kong     /*ierr = PetscMemzero(c_rmtc,sizeof(PetscInt)*pon);CHKERRQ(ierr);*/
8634a56b808SFande Kong 
8644a56b808SFande Kong     /* Form C(ii, :) */
8654a56b808SFande Kong     poj = po->j + po->i[i];
8664a56b808SFande Kong     poa = po->a + po->i[i];
8674a56b808SFande Kong     for (j=0; j<nzi; j++) {
8684a56b808SFande Kong       c_rmtjj = c_rmtj + ptap->c_rmti[poj[j]];
8694a56b808SFande Kong       c_rmtaa = c_rmta + ptap->c_rmti[poj[j]];
8704a56b808SFande Kong       for (jj=0; jj<voff; jj++) {
8714a56b808SFande Kong         apvaluestmp[jj] = apvalues[jj]*poa[j];
8724a56b808SFande Kong         /*If the row is empty */
8734a56b808SFande Kong         if (!c_rmtc[poj[j]]) {
8744a56b808SFande Kong           c_rmtjj[jj] = apindices[jj];
8754a56b808SFande Kong           c_rmtaa[jj] = apvaluestmp[jj];
8764a56b808SFande Kong           c_rmtc[poj[j]]++;
8774a56b808SFande Kong         } else {
8784a56b808SFande Kong           ierr = PetscFindInt(apindices[jj],c_rmtc[poj[j]],c_rmtjj,&loc);CHKERRQ(ierr);
8794a56b808SFande Kong           if (loc>=0){ /* hit */
8804a56b808SFande Kong             c_rmtaa[loc] += apvaluestmp[jj];
8814a56b808SFande Kong           } else { /* new element */
8824a56b808SFande Kong             loc = -(loc+1);
8834a56b808SFande Kong             /* Move data backward */
8844a56b808SFande Kong             for (kk=c_rmtc[poj[j]]; kk>loc; kk--) {
8854a56b808SFande Kong               c_rmtjj[kk] = c_rmtjj[kk-1];
8864a56b808SFande Kong               c_rmtaa[kk] = c_rmtaa[kk-1];
8874a56b808SFande Kong             }/* End kk */
8884a56b808SFande Kong             c_rmtjj[loc] = apindices[jj];
8894a56b808SFande Kong             c_rmtaa[loc] = apvaluestmp[jj];
8904a56b808SFande Kong             c_rmtc[poj[j]]++;
8914a56b808SFande Kong           }
8924a56b808SFande Kong         }
8934a56b808SFande Kong       } /* End jj */
8944a56b808SFande Kong     } /* End j */
8954a56b808SFande Kong   } /* End i */
8964a56b808SFande Kong 
8974a56b808SFande Kong   ierr = PetscFree4(apindices,apvalues,apvaluestmp,c_rmtc);CHKERRQ(ierr);
8984a56b808SFande Kong 
8994a56b808SFande Kong   ierr = MatGetLocalSize(P,NULL,&pn);CHKERRQ(ierr);
9004a56b808SFande Kong   ierr = PetscCalloc2(ptap->c_othi[pn],&c_othj,ptap->c_othi[pn],&c_otha);CHKERRQ(ierr);
9014a56b808SFande Kong 
9024a56b808SFande Kong   ierr = PetscSFReduceBegin(ptap->sf,MPIU_INT,c_rmtj,c_othj,MPIU_REPLACE);CHKERRQ(ierr);
9034a56b808SFande Kong   ierr = PetscSFReduceBegin(ptap->sf,MPIU_SCALAR,c_rmta,c_otha,MPIU_REPLACE);CHKERRQ(ierr);
9044a56b808SFande Kong   ierr = MatGetOwnershipRangeColumn(P,&pcstart,&pcend);CHKERRQ(ierr);
9054a56b808SFande Kong   cd = (Mat_SeqAIJ*)(c->A)->data;
9064a56b808SFande Kong   co = (Mat_SeqAIJ*)(c->B)->data;
9074a56b808SFande Kong 
9084a56b808SFande Kong   cmaxr = 0;
9094a56b808SFande Kong   for (i=0; i<pn; i++) {
9104a56b808SFande Kong     cmaxr = PetscMax(cmaxr,(cd->i[i+1]-cd->i[i])+(co->i[i+1]-co->i[i]));
9114a56b808SFande Kong   }
9124a56b808SFande Kong   ierr = PetscCalloc5(cmaxr,&apindices,cmaxr,&apvalues,cmaxr,&apvaluestmp,pn,&dcc,pn,&occ);CHKERRQ(ierr);
9134a56b808SFande Kong   ierr = PetscHMapIVCreate(&hmap);CHKERRQ(ierr);
9144a56b808SFande Kong   ierr = PetscHMapIVResize(hmap,cmaxr);CHKERRQ(ierr);
9154a56b808SFande Kong   for (i=0; i<am && pn; i++) {
9164a56b808SFande Kong     ierr = PetscHMapIVClear(hmap);CHKERRQ(ierr);
9174a56b808SFande Kong     nzi = pd->i[i+1] - pd->i[i];
9184a56b808SFande Kong     if (!nzi) continue;
9194a56b808SFande Kong     ierr = MatPtAPNumericComputeOneRowOfAP_private(A,P,ptap->P_oth,i,hmap);CHKERRQ(ierr);
9204a56b808SFande Kong     voff = 0;
9214a56b808SFande Kong     ierr = PetscHMapIVGetPairs(hmap,&voff,apindices,apvalues);CHKERRQ(ierr);
9224a56b808SFande Kong     if (!voff) continue;
9234a56b808SFande Kong     /* Form C(ii, :) */
9244a56b808SFande Kong     pdj = pd->j + pd->i[i];
9254a56b808SFande Kong     pda = pd->a + pd->i[i];
9264a56b808SFande Kong     for (j=0; j<nzi; j++) {
9274a56b808SFande Kong       row = pcstart + pdj[j];
9284a56b808SFande Kong       for (jj=0; jj<voff; jj++) {
9294a56b808SFande Kong         apvaluestmp[jj] = apvalues[jj]*pda[j];
9304a56b808SFande Kong       }
9314a56b808SFande Kong       ierr = MatSetValues(C,1,&row,voff,apindices,apvaluestmp,ADD_VALUES);CHKERRQ(ierr);
9324a56b808SFande Kong     }
9334a56b808SFande Kong   }
9344a56b808SFande Kong 
9354a56b808SFande Kong   ierr = MatGetOwnershipRangeColumn(C,&ccstart,&ccend);CHKERRQ(ierr);
9364a56b808SFande Kong   ierr = PetscFree5(apindices,apvalues,apvaluestmp,dcc,occ);CHKERRQ(ierr);
9374a56b808SFande Kong   ierr = PetscHMapIVDestroy(&hmap);CHKERRQ(ierr);
9384a56b808SFande Kong   ierr = PetscSFReduceEnd(ptap->sf,MPIU_INT,c_rmtj,c_othj,MPI_SUM);CHKERRQ(ierr);
9394a56b808SFande Kong   ierr = PetscSFReduceEnd(ptap->sf,MPIU_SCALAR,c_rmta,c_otha,MPIU_REPLACE);CHKERRQ(ierr);
9404a56b808SFande Kong   ierr = PetscFree2(c_rmtj,c_rmta);CHKERRQ(ierr);
9414a56b808SFande Kong 
9424a56b808SFande Kong   /* Add contributions from remote */
9434a56b808SFande Kong   for (i = 0; i < pn; i++) {
9444a56b808SFande Kong     row = i + pcstart;
9454a56b808SFande Kong     ierr = MatSetValues(C,1,&row,ptap->c_othi[i+1]-ptap->c_othi[i],c_othj+ptap->c_othi[i],c_otha+ptap->c_othi[i],ADD_VALUES);CHKERRQ(ierr);
9464a56b808SFande Kong   }
9474a56b808SFande Kong   ierr = PetscFree2(c_othj,c_otha);CHKERRQ(ierr);
9484a56b808SFande Kong 
9494a56b808SFande Kong   ierr = MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
9504a56b808SFande Kong   ierr = MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
9514a56b808SFande Kong 
9524a56b808SFande Kong   ptap->reuse = MAT_REUSE_MATRIX;
9534a56b808SFande Kong 
9544a56b808SFande Kong   /* supporting struct ptap consumes almost same amount of memory as C=PtAP, release it if C will not be updated by A and P */
9554a56b808SFande Kong   if (ptap->freestruct) {
9564a56b808SFande Kong     ierr = MatFreeIntermediateDataStructures(C);CHKERRQ(ierr);
9574a56b808SFande Kong   }
9584a56b808SFande Kong   PetscFunctionReturn(0);
9594a56b808SFande Kong }
9604a56b808SFande Kong 
9614a56b808SFande Kong PetscErrorCode MatPtAPNumeric_MPIAIJ_MPIAIJ_allatonce_merged(Mat A,Mat P,Mat C)
9624a56b808SFande Kong {
9634a56b808SFande Kong   PetscErrorCode    ierr;
9644a56b808SFande Kong   Mat_MPIAIJ        *p=(Mat_MPIAIJ*)P->data,*c=(Mat_MPIAIJ*)C->data;
9654a56b808SFande Kong   Mat_SeqAIJ        *cd,*co,*po,*pd;
9664a56b808SFande Kong   Mat_APMPI         *ptap = c->ap;
9674a56b808SFande Kong   PetscHMapIV       hmap;
9684a56b808SFande Kong   PetscInt          i,j,jj,kk,nzi,dnzi,*c_rmtj,voff,*c_othj,pn,pon,pcstart,pcend,row,am,*poj,*pdj,*apindices,cmaxr,*c_rmtc,*c_rmtjj,loc;
9694a56b808SFande Kong   PetscScalar       *c_rmta,*c_otha,*poa,*pda,*apvalues,*apvaluestmp,*c_rmtaa;
9704a56b808SFande Kong   MPI_Comm          comm;
9714a56b808SFande Kong 
9724a56b808SFande Kong   PetscFunctionBegin;
9734a56b808SFande Kong   ierr = PetscObjectGetComm((PetscObject)A,&comm);CHKERRQ(ierr);
9744a56b808SFande Kong   if (!ptap) SETERRQ(comm,PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be reused. Do not call MatFreeIntermediateDataStructures() or use '-mat_freeintermediatedatastructures'");
9754a56b808SFande Kong 
9764a56b808SFande Kong   ierr = MatZeroEntries(C);CHKERRQ(ierr);
9774a56b808SFande Kong 
9784a56b808SFande Kong   /* Get P_oth = ptap->P_oth  and P_loc = ptap->P_loc */
9794a56b808SFande Kong   /*-----------------------------------------------------*/
9804a56b808SFande Kong   if (ptap->reuse == MAT_REUSE_MATRIX) {
9814a56b808SFande Kong     /* P_oth and P_loc are obtained in MatPtASymbolic() when reuse == MAT_INITIAL_MATRIX */
9824a56b808SFande Kong     ierr = MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_REUSE_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth);CHKERRQ(ierr);
9834a56b808SFande Kong   }
9844a56b808SFande Kong 
9854a56b808SFande Kong   po = (Mat_SeqAIJ*) p->B->data;
9864a56b808SFande Kong   pd = (Mat_SeqAIJ*) p->A->data;
9874a56b808SFande Kong 
9884a56b808SFande Kong   ierr = MatGetLocalSize(p->B,NULL,&pon);CHKERRQ(ierr);
9894a56b808SFande Kong   ierr = MatGetLocalSize(P,NULL,&pn);CHKERRQ(ierr);
9904a56b808SFande Kong 
9914a56b808SFande Kong   ierr = PetscCalloc2(ptap->c_rmti[pon],&c_rmtj,ptap->c_rmti[pon],&c_rmta);CHKERRQ(ierr);
9924a56b808SFande Kong   ierr = MatGetLocalSize(A,&am,NULL);CHKERRQ(ierr);
9934a56b808SFande Kong   ierr = MatGetOwnershipRangeColumn(P,&pcstart,&pcend);CHKERRQ(ierr);
9944a56b808SFande Kong   cmaxr = 0;
9954a56b808SFande Kong   for (i=0; i<pon; i++) {
9964a56b808SFande Kong     cmaxr = PetscMax(cmaxr,ptap->c_rmti[i+1]-ptap->c_rmti[i]);
9974a56b808SFande Kong   }
9984a56b808SFande Kong   cd = (Mat_SeqAIJ*)(c->A)->data;
9994a56b808SFande Kong   co = (Mat_SeqAIJ*)(c->B)->data;
10004a56b808SFande Kong   for (i=0; i<pn; i++) {
10014a56b808SFande Kong     cmaxr = PetscMax(cmaxr,(cd->i[i+1]-cd->i[i])+(co->i[i+1]-co->i[i]));
10024a56b808SFande Kong   }
10034a56b808SFande Kong   ierr = PetscCalloc4(cmaxr,&apindices,cmaxr,&apvalues,cmaxr,&apvaluestmp,pon,&c_rmtc);CHKERRQ(ierr);
10044a56b808SFande Kong   ierr = PetscHMapIVCreate(&hmap);CHKERRQ(ierr);
10054a56b808SFande Kong   ierr = PetscHMapIVResize(hmap,cmaxr);CHKERRQ(ierr);
10064a56b808SFande Kong   for (i=0; i<am && (pon || pn); i++) {
10074a56b808SFande Kong     ierr = PetscHMapIVClear(hmap);CHKERRQ(ierr);
10084a56b808SFande Kong     nzi  = po->i[i+1] - po->i[i];
10094a56b808SFande Kong     dnzi = pd->i[i+1] - pd->i[i];
10104a56b808SFande Kong     if (!nzi && !dnzi) continue;
10114a56b808SFande Kong     ierr = MatPtAPNumericComputeOneRowOfAP_private(A,P,ptap->P_oth,i,hmap);CHKERRQ(ierr);
10124a56b808SFande Kong     voff = 0;
10134a56b808SFande Kong     ierr = PetscHMapIVGetPairs(hmap,&voff,apindices,apvalues);CHKERRQ(ierr);
10144a56b808SFande Kong     if (!voff) continue;
10154a56b808SFande Kong 
10164a56b808SFande Kong     /* Form remote C(ii, :) */
10174a56b808SFande Kong     poj = po->j + po->i[i];
10184a56b808SFande Kong     poa = po->a + po->i[i];
10194a56b808SFande Kong     for (j=0; j<nzi; j++) {
10204a56b808SFande Kong       c_rmtjj = c_rmtj + ptap->c_rmti[poj[j]];
10214a56b808SFande Kong       c_rmtaa = c_rmta + ptap->c_rmti[poj[j]];
10224a56b808SFande Kong       for (jj=0; jj<voff; jj++) {
10234a56b808SFande Kong         apvaluestmp[jj] = apvalues[jj]*poa[j];
10244a56b808SFande Kong         /*If the row is empty */
10254a56b808SFande Kong         if (!c_rmtc[poj[j]]) {
10264a56b808SFande Kong           c_rmtjj[jj] = apindices[jj];
10274a56b808SFande Kong           c_rmtaa[jj] = apvaluestmp[jj];
10284a56b808SFande Kong           c_rmtc[poj[j]]++;
10294a56b808SFande Kong         } else {
10304a56b808SFande Kong           ierr = PetscFindInt(apindices[jj],c_rmtc[poj[j]],c_rmtjj,&loc);CHKERRQ(ierr);
10314a56b808SFande Kong           if (loc>=0){ /* hit */
10324a56b808SFande Kong             c_rmtaa[loc] += apvaluestmp[jj];
10334a56b808SFande Kong           } else { /* new element */
10344a56b808SFande Kong             loc = -(loc+1);
10354a56b808SFande Kong             /* Move data backward */
10364a56b808SFande Kong             for (kk=c_rmtc[poj[j]]; kk>loc; kk--) {
10374a56b808SFande Kong               c_rmtjj[kk] = c_rmtjj[kk-1];
10384a56b808SFande Kong               c_rmtaa[kk] = c_rmtaa[kk-1];
10394a56b808SFande Kong             }/* End kk */
10404a56b808SFande Kong             c_rmtjj[loc] = apindices[jj];
10414a56b808SFande Kong             c_rmtaa[loc] = apvaluestmp[jj];
10424a56b808SFande Kong             c_rmtc[poj[j]]++;
10434a56b808SFande Kong           }
10444a56b808SFande Kong         }
10454a56b808SFande Kong       } /* End jj */
10464a56b808SFande Kong     } /* End j */
10474a56b808SFande Kong 
10484a56b808SFande Kong     /* Form local C(ii, :) */
10494a56b808SFande Kong     pdj = pd->j + pd->i[i];
10504a56b808SFande Kong     pda = pd->a + pd->i[i];
10514a56b808SFande Kong     for (j=0; j<dnzi; j++) {
10524a56b808SFande Kong       row = pcstart + pdj[j];
10534a56b808SFande Kong       for (jj=0; jj<voff; jj++) {
10544a56b808SFande Kong         apvaluestmp[jj] = apvalues[jj]*pda[j];
10554a56b808SFande Kong       }/* End kk */
10564a56b808SFande Kong       ierr = MatSetValues(C,1,&row,voff,apindices,apvaluestmp,ADD_VALUES);CHKERRQ(ierr);
10574a56b808SFande Kong     }/* End j */
10584a56b808SFande Kong   } /* End i */
10594a56b808SFande Kong 
10604a56b808SFande Kong   ierr = PetscFree4(apindices,apvalues,apvaluestmp,c_rmtc);CHKERRQ(ierr);
10614a56b808SFande Kong   ierr = PetscHMapIVDestroy(&hmap);CHKERRQ(ierr);
10624a56b808SFande Kong   ierr = PetscCalloc2(ptap->c_othi[pn],&c_othj,ptap->c_othi[pn],&c_otha);CHKERRQ(ierr);
10634a56b808SFande Kong 
10644a56b808SFande Kong   ierr = PetscSFReduceBegin(ptap->sf,MPIU_INT,c_rmtj,c_othj,MPIU_REPLACE);CHKERRQ(ierr);
10654a56b808SFande Kong   ierr = PetscSFReduceBegin(ptap->sf,MPIU_SCALAR,c_rmta,c_otha,MPIU_REPLACE);CHKERRQ(ierr);
10664a56b808SFande Kong   ierr = PetscSFReduceEnd(ptap->sf,MPIU_INT,c_rmtj,c_othj,MPI_SUM);CHKERRQ(ierr);
10674a56b808SFande Kong   ierr = PetscSFReduceEnd(ptap->sf,MPIU_SCALAR,c_rmta,c_otha,MPIU_REPLACE);CHKERRQ(ierr);
10684a56b808SFande Kong   ierr = PetscFree2(c_rmtj,c_rmta);CHKERRQ(ierr);
10694a56b808SFande Kong 
10704a56b808SFande Kong   /* Add contributions from remote */
10714a56b808SFande Kong   for (i = 0; i < pn; i++) {
10724a56b808SFande Kong     row = i + pcstart;
10734a56b808SFande Kong     ierr = MatSetValues(C,1,&row,ptap->c_othi[i+1]-ptap->c_othi[i],c_othj+ptap->c_othi[i],c_otha+ptap->c_othi[i],ADD_VALUES);CHKERRQ(ierr);
10744a56b808SFande Kong   }
10754a56b808SFande Kong   ierr = PetscFree2(c_othj,c_otha);CHKERRQ(ierr);
10764a56b808SFande Kong 
10774a56b808SFande Kong   ierr = MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
10784a56b808SFande Kong   ierr = MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
10794a56b808SFande Kong 
10804a56b808SFande Kong   ptap->reuse = MAT_REUSE_MATRIX;
10814a56b808SFande Kong 
10824a56b808SFande Kong   /* supporting struct ptap consumes almost same amount of memory as C=PtAP, release it if C will not be updated by A and P */
10834a56b808SFande Kong   if (ptap->freestruct) {
10844a56b808SFande Kong     ierr = MatFreeIntermediateDataStructures(C);CHKERRQ(ierr);
10854a56b808SFande Kong   }
10864a56b808SFande Kong   PetscFunctionReturn(0);
10874a56b808SFande Kong }
10884a56b808SFande Kong 
10894a56b808SFande Kong PetscErrorCode MatPtAPSymbolic_MPIAIJ_MPIAIJ_allatonce(Mat A,Mat P,PetscReal fill,Mat *C)
10904a56b808SFande Kong {
10914a56b808SFande Kong   Mat_APMPI           *ptap;
10924a56b808SFande Kong   Mat_MPIAIJ          *p=(Mat_MPIAIJ*)P->data,*c;
10934a56b808SFande Kong   MPI_Comm            comm;
10944a56b808SFande Kong   Mat                 Cmpi;
10954a56b808SFande Kong   Mat_SeqAIJ          *pd,*po;
10964a56b808SFande Kong   MatType             mtype;
10974a56b808SFande Kong   PetscSF             sf;
10984a56b808SFande Kong   PetscSFNode         *iremote;
10994a56b808SFande Kong   PetscInt            rootspacesize,*rootspace,*rootspaceoffsets,nleaves;
11004a56b808SFande Kong   const PetscInt      *rootdegrees;
11014a56b808SFande Kong   PetscHSetI          ht,oht,*hta,*hto;
11024a56b808SFande Kong   PetscInt            pn,pon,*c_rmtc,i,j,nzi,htsize,htosize,*c_rmtj,off,*c_othj,rcvncols,sendncols,*c_rmtoffsets;
11034a56b808SFande Kong   PetscInt            owner,lidx,*rdj,col,pcstart,pcend,*dnz,*onz,am,arstart,arend,*poj,*pdj;
11044a56b808SFande Kong   PetscInt            nalg=2,alg=0;
11054a56b808SFande Kong   PetscBool           flg;
11064a56b808SFande Kong   const char          *algTypes[2] = {"overlapping","merged"};
11074a56b808SFande Kong   PetscErrorCode      ierr;
11084a56b808SFande Kong 
11094a56b808SFande Kong   PetscFunctionBegin;
11104a56b808SFande Kong   ierr = PetscObjectGetComm((PetscObject)A,&comm);CHKERRQ(ierr);
11114a56b808SFande Kong 
11124a56b808SFande Kong   /* Create symbolic parallel matrix Cmpi */
11134a56b808SFande Kong   ierr = MatGetLocalSize(P,NULL,&pn);CHKERRQ(ierr);
11144a56b808SFande Kong   ierr = MatCreate(comm,&Cmpi);CHKERRQ(ierr);
11154a56b808SFande Kong   ierr = MatGetType(A,&mtype);CHKERRQ(ierr);
11164a56b808SFande Kong   ierr = MatSetType(Cmpi,mtype);CHKERRQ(ierr);
11174a56b808SFande Kong   ierr = MatSetSizes(Cmpi,pn,pn,PETSC_DETERMINE,PETSC_DETERMINE);CHKERRQ(ierr);
11184a56b808SFande Kong 
11194a56b808SFande Kong   ierr        = PetscNew(&ptap);CHKERRQ(ierr);
11204a56b808SFande Kong   ptap->reuse = MAT_INITIAL_MATRIX;
11214a56b808SFande Kong   ptap->algType = 2;
11224a56b808SFande Kong 
11234a56b808SFande Kong   /* Get P_oth by taking rows of P (= non-zero cols of local A) from other processors */
11244a56b808SFande Kong   ierr = MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth);CHKERRQ(ierr);
11254a56b808SFande Kong 
11264a56b808SFande Kong   po = (Mat_SeqAIJ*)p->B->data;
11274a56b808SFande Kong   pd = (Mat_SeqAIJ*)p->A->data;
11284a56b808SFande Kong 
11294a56b808SFande Kong   /* This equals to the number of offdiag columns in P */
11304a56b808SFande Kong   ierr = MatGetLocalSize(p->B,NULL,&pon);CHKERRQ(ierr);
11314a56b808SFande Kong   /* offsets */
11324a56b808SFande Kong   ierr = PetscMalloc1(pon+1,&ptap->c_rmti);CHKERRQ(ierr);
11334a56b808SFande Kong   /* The number of columns we will send to remote ranks */
11344a56b808SFande Kong   ierr = PetscMalloc1(pon,&c_rmtc);CHKERRQ(ierr);
11354a56b808SFande Kong   ierr = PetscMalloc1(pon,&hta);CHKERRQ(ierr);
11364a56b808SFande Kong   for (i=0; i<pon; i++) {
11374a56b808SFande Kong     ierr = PetscHSetICreate(&hta[i]);CHKERRQ(ierr);
11384a56b808SFande Kong   }
11394a56b808SFande Kong   ierr = MatGetLocalSize(A,&am,NULL);CHKERRQ(ierr);
11404a56b808SFande Kong   ierr = MatGetOwnershipRange(A,&arstart,&arend);CHKERRQ(ierr);
11414a56b808SFande Kong   /* Create hash table to merge all columns for C(i, :) */
11424a56b808SFande Kong   ierr = PetscHSetICreate(&ht);CHKERRQ(ierr);
11434a56b808SFande Kong 
11444a56b808SFande Kong   ptap->c_rmti[0] = 0;
11454a56b808SFande Kong   /* 2) Pass 1: calculate the size for C_rmt (a matrix need to be sent to other processors)  */
11464a56b808SFande Kong   for (i=0; i<am && pon; i++) {
11474a56b808SFande Kong     /* Form one row of AP */
11484a56b808SFande Kong     ierr = PetscHSetIClear(ht);CHKERRQ(ierr);
11494a56b808SFande Kong     /* If the off diag is empty, we should not do any calculation */
11504a56b808SFande Kong     nzi = po->i[i+1] - po->i[i];
11514a56b808SFande Kong     if (!nzi) continue;
11524a56b808SFande Kong 
11534a56b808SFande Kong     ierr = MatPtAPSymbolicComputeOneRowOfAP_private(A,P,ptap->P_oth,i,ht,ht);CHKERRQ(ierr);
11544a56b808SFande Kong     ierr = PetscHSetIGetSize(ht,&htsize);CHKERRQ(ierr);
11554a56b808SFande Kong     /* If AP is empty, just continue */
11564a56b808SFande Kong     if (!htsize) continue;
11574a56b808SFande Kong     /* Form C(ii, :) */
11584a56b808SFande Kong     poj = po->j + po->i[i];
11594a56b808SFande Kong     for (j=0; j<nzi; j++) {
11604a56b808SFande Kong       ierr = PetscHSetIAppend(hta[poj[j]],ht);CHKERRQ(ierr);
11614a56b808SFande Kong     }
11624a56b808SFande Kong   }
11634a56b808SFande Kong 
11644a56b808SFande Kong   for (i=0; i<pon; i++) {
11654a56b808SFande Kong     ierr = PetscHSetIGetSize(hta[i],&htsize);CHKERRQ(ierr);
11664a56b808SFande Kong     ptap->c_rmti[i+1] = ptap->c_rmti[i] + htsize;
11674a56b808SFande Kong     c_rmtc[i] = htsize;
11684a56b808SFande Kong   }
11694a56b808SFande Kong 
11704a56b808SFande Kong   ierr = PetscMalloc1(ptap->c_rmti[pon],&c_rmtj);CHKERRQ(ierr);
11714a56b808SFande Kong 
11724a56b808SFande Kong   for (i=0; i<pon; i++) {
11734a56b808SFande Kong     off = 0;
11744a56b808SFande Kong     ierr = PetscHSetIGetElems(hta[i],&off,c_rmtj+ptap->c_rmti[i]);CHKERRQ(ierr);
11754a56b808SFande Kong     ierr = PetscHSetIDestroy(&hta[i]);CHKERRQ(ierr);
11764a56b808SFande Kong   }
11774a56b808SFande Kong   ierr = PetscFree(hta);CHKERRQ(ierr);
11784a56b808SFande Kong 
11794a56b808SFande Kong   ierr = PetscCalloc1(pon,&iremote);CHKERRQ(ierr);
11804a56b808SFande Kong   for (i=0; i<pon; i++) {
1181*ef7a94f2SFande Kong     owner = 0; lidx = 0;
11824a56b808SFande Kong     ierr = PetscLayoutFindOwnerIndex(P->cmap,p->garray[i],&owner,&lidx);CHKERRQ(ierr);
11834a56b808SFande Kong     iremote[i].index = lidx;
11844a56b808SFande Kong     iremote[i].rank  = owner;
11854a56b808SFande Kong   }
11864a56b808SFande Kong 
11874a56b808SFande Kong   ierr = PetscSFCreate(comm,&sf);CHKERRQ(ierr);
11884a56b808SFande Kong   ierr = PetscSFSetGraph(sf,pn,pon,NULL,PETSC_OWN_POINTER,iremote,PETSC_OWN_POINTER);CHKERRQ(ierr);
11894a56b808SFande Kong   /* Reorder ranks properly so that the data handled by gather and scatter have the same order */
11904a56b808SFande Kong   ierr = PetscSFSetRankOrder(sf,PETSC_TRUE);CHKERRQ(ierr);
11914a56b808SFande Kong   ierr = PetscSFSetFromOptions(sf);CHKERRQ(ierr);
11924a56b808SFande Kong   ierr = PetscSFSetUp(sf);CHKERRQ(ierr);
11934a56b808SFande Kong   /* How many neighbors have contributions to my rows? */
11944a56b808SFande Kong   ierr = PetscSFComputeDegreeBegin(sf,&rootdegrees);CHKERRQ(ierr);
11954a56b808SFande Kong   ierr = PetscSFComputeDegreeEnd(sf,&rootdegrees);CHKERRQ(ierr);
11964a56b808SFande Kong   rootspacesize = 0;
11974a56b808SFande Kong   for (i = 0; i < pn; i++) {
11984a56b808SFande Kong     rootspacesize += rootdegrees[i];
11994a56b808SFande Kong   }
12004a56b808SFande Kong   ierr = PetscCalloc1(rootspacesize,&rootspace);CHKERRQ(ierr);
12014a56b808SFande Kong   ierr = PetscCalloc1(rootspacesize+1,&rootspaceoffsets);CHKERRQ(ierr);
12024a56b808SFande Kong   /* Get information from leaves
12034a56b808SFande Kong    * Number of columns other people contribute to my rows
12044a56b808SFande Kong    * */
12054a56b808SFande Kong   ierr = PetscSFGatherBegin(sf,MPIU_INT,c_rmtc,rootspace);CHKERRQ(ierr);
12064a56b808SFande Kong   ierr = PetscSFGatherEnd(sf,MPIU_INT,c_rmtc,rootspace);CHKERRQ(ierr);
12074a56b808SFande Kong   ierr = PetscFree(c_rmtc);CHKERRQ(ierr);
12084a56b808SFande Kong   ierr = PetscCalloc1(pn+1,&ptap->c_othi);CHKERRQ(ierr);
12094a56b808SFande Kong   /* The number of columns is received for each row */
12104a56b808SFande Kong   ptap->c_othi[0] = 0;
12114a56b808SFande Kong   rootspacesize = 0;
12124a56b808SFande Kong   rootspaceoffsets[0] = 0;
12134a56b808SFande Kong   for (i = 0; i < pn; i++) {
12144a56b808SFande Kong     rcvncols = 0;
12154a56b808SFande Kong     for (j = 0; j<rootdegrees[i]; j++) {
12164a56b808SFande Kong       rcvncols += rootspace[rootspacesize];
12174a56b808SFande Kong       rootspaceoffsets[rootspacesize+1] = rootspaceoffsets[rootspacesize] + rootspace[rootspacesize];
12184a56b808SFande Kong       rootspacesize++;
12194a56b808SFande Kong     }
12204a56b808SFande Kong     ptap->c_othi[i+1] = ptap->c_othi[i] + rcvncols;
12214a56b808SFande Kong   }
12224a56b808SFande Kong   ierr = PetscFree(rootspace);CHKERRQ(ierr);
12234a56b808SFande Kong 
12244a56b808SFande Kong   ierr = PetscCalloc1(pon,&c_rmtoffsets);CHKERRQ(ierr);
12254a56b808SFande Kong   ierr = PetscSFScatterBegin(sf,MPIU_INT,rootspaceoffsets,c_rmtoffsets);CHKERRQ(ierr);
12264a56b808SFande Kong   ierr = PetscSFScatterEnd(sf,MPIU_INT,rootspaceoffsets,c_rmtoffsets);CHKERRQ(ierr);
12274a56b808SFande Kong   ierr = PetscSFDestroy(&sf);CHKERRQ(ierr);
12284a56b808SFande Kong   ierr = PetscFree(rootspaceoffsets);CHKERRQ(ierr);
12294a56b808SFande Kong 
12304a56b808SFande Kong   ierr = PetscCalloc1(ptap->c_rmti[pon],&iremote);CHKERRQ(ierr);
12314a56b808SFande Kong   nleaves = 0;
12324a56b808SFande Kong   for (i = 0; i<pon; i++) {
1233*ef7a94f2SFande Kong     owner = 0; lidx = 0;
12344a56b808SFande Kong     ierr = PetscLayoutFindOwnerIndex(P->cmap,p->garray[i],&owner,&lidx);CHKERRQ(ierr);
12354a56b808SFande Kong     sendncols = ptap->c_rmti[i+1] - ptap->c_rmti[i];
12364a56b808SFande Kong     for (j=0; j<sendncols; j++) {
12374a56b808SFande Kong       iremote[nleaves].rank = owner;
12384a56b808SFande Kong       iremote[nleaves++].index = c_rmtoffsets[i] + j;
12394a56b808SFande Kong     }
12404a56b808SFande Kong   }
12414a56b808SFande Kong   ierr = PetscFree(c_rmtoffsets);CHKERRQ(ierr);
12424a56b808SFande Kong   ierr = PetscCalloc1(ptap->c_othi[pn],&c_othj);CHKERRQ(ierr);
12434a56b808SFande Kong 
12444a56b808SFande Kong   ierr = PetscSFCreate(comm,&ptap->sf);CHKERRQ(ierr);
12454a56b808SFande Kong   ierr = PetscSFSetGraph(ptap->sf,ptap->c_othi[pn],nleaves,NULL,PETSC_OWN_POINTER,iremote,PETSC_OWN_POINTER);CHKERRQ(ierr);
12464a56b808SFande Kong   ierr = PetscSFSetFromOptions(ptap->sf);CHKERRQ(ierr);
12474a56b808SFande Kong   /* One to one map */
12484a56b808SFande Kong   ierr = PetscSFReduceBegin(ptap->sf,MPIU_INT,c_rmtj,c_othj,MPIU_REPLACE);CHKERRQ(ierr);
12494a56b808SFande Kong 
12504a56b808SFande Kong   ierr = PetscCalloc2(pn,&dnz,pn,&onz);CHKERRQ(ierr);
12514a56b808SFande Kong   ierr = PetscHSetICreate(&oht);CHKERRQ(ierr);
12524a56b808SFande Kong   ierr = MatGetOwnershipRangeColumn(P,&pcstart,&pcend);CHKERRQ(ierr);
12534a56b808SFande Kong   ierr = PetscMalloc2(pn,&hta,pn,&hto);CHKERRQ(ierr);
12544a56b808SFande Kong   for (i=0; i<pn; i++) {
12554a56b808SFande Kong     ierr = PetscHSetICreate(&hta[i]);CHKERRQ(ierr);
12564a56b808SFande Kong     ierr = PetscHSetICreate(&hto[i]);CHKERRQ(ierr);
12574a56b808SFande Kong   }
12584a56b808SFande Kong   /* Work on local part */
12594a56b808SFande Kong   /* 4) Pass 1: Estimate memory for C_loc */
12604a56b808SFande Kong   for (i=0; i<am && pn; i++) {
12614a56b808SFande Kong     ierr = PetscHSetIClear(ht);CHKERRQ(ierr);
12624a56b808SFande Kong     ierr = PetscHSetIClear(oht);CHKERRQ(ierr);
12634a56b808SFande Kong     nzi = pd->i[i+1] - pd->i[i];
12644a56b808SFande Kong     if (!nzi) continue;
12654a56b808SFande Kong 
12664a56b808SFande Kong     ierr = MatPtAPSymbolicComputeOneRowOfAP_private(A,P,ptap->P_oth,i,ht,oht);CHKERRQ(ierr);
12674a56b808SFande Kong     ierr = PetscHSetIGetSize(ht,&htsize);CHKERRQ(ierr);
12684a56b808SFande Kong     ierr = PetscHSetIGetSize(oht,&htosize);CHKERRQ(ierr);
12694a56b808SFande Kong     if (!(htsize+htosize)) continue;
12704a56b808SFande Kong     /* Form C(ii, :) */
12714a56b808SFande Kong     pdj = pd->j + pd->i[i];
12724a56b808SFande Kong     for (j=0; j<nzi; j++) {
12734a56b808SFande Kong       ierr = PetscHSetIAppend(hta[pdj[j]],ht);CHKERRQ(ierr);
12744a56b808SFande Kong       ierr = PetscHSetIAppend(hto[pdj[j]],oht);CHKERRQ(ierr);
12754a56b808SFande Kong     }
12764a56b808SFande Kong   }
12774a56b808SFande Kong 
12784a56b808SFande Kong   ierr = PetscHSetIDestroy(&ht);CHKERRQ(ierr);
12794a56b808SFande Kong   ierr = PetscHSetIDestroy(&oht);CHKERRQ(ierr);
12804a56b808SFande Kong 
12814a56b808SFande Kong   /* Get remote data */
12824a56b808SFande Kong   ierr = PetscSFReduceEnd(ptap->sf,MPIU_INT,c_rmtj,c_othj,MPIU_REPLACE);CHKERRQ(ierr);
12834a56b808SFande Kong   ierr = PetscFree(c_rmtj);CHKERRQ(ierr);
12844a56b808SFande Kong 
12854a56b808SFande Kong   for (i = 0; i < pn; i++) {
12864a56b808SFande Kong     nzi = ptap->c_othi[i+1] - ptap->c_othi[i];
12874a56b808SFande Kong     rdj = c_othj + ptap->c_othi[i];
12884a56b808SFande Kong     for (j = 0; j < nzi; j++) {
12894a56b808SFande Kong       col =  rdj[j];
12904a56b808SFande Kong       /* diag part */
12914a56b808SFande Kong       if (col>=pcstart && col<pcend) {
12924a56b808SFande Kong         ierr = PetscHSetIAdd(hta[i],col);CHKERRQ(ierr);
12934a56b808SFande Kong       } else { /* off diag */
12944a56b808SFande Kong         ierr = PetscHSetIAdd(hto[i],col);CHKERRQ(ierr);
12954a56b808SFande Kong       }
12964a56b808SFande Kong     }
12974a56b808SFande Kong     ierr = PetscHSetIGetSize(hta[i],&htsize);CHKERRQ(ierr);
12984a56b808SFande Kong     dnz[i] = htsize;
12994a56b808SFande Kong     ierr = PetscHSetIDestroy(&hta[i]);CHKERRQ(ierr);
13004a56b808SFande Kong     ierr = PetscHSetIGetSize(hto[i],&htsize);CHKERRQ(ierr);
13014a56b808SFande Kong     onz[i] = htsize;
13024a56b808SFande Kong     ierr = PetscHSetIDestroy(&hto[i]);CHKERRQ(ierr);
13034a56b808SFande Kong   }
13044a56b808SFande Kong 
13054a56b808SFande Kong   ierr = PetscFree2(hta,hto);CHKERRQ(ierr);
13064a56b808SFande Kong   ierr = PetscFree(c_othj);CHKERRQ(ierr);
13074a56b808SFande Kong 
13084a56b808SFande Kong   /* local sizes and preallocation */
13094a56b808SFande Kong   ierr = MatSetSizes(Cmpi,pn,pn,PETSC_DETERMINE,PETSC_DETERMINE);CHKERRQ(ierr);
13104a56b808SFande Kong   ierr = MatSetBlockSizes(Cmpi,PetscAbs(P->cmap->bs),PetscAbs(P->cmap->bs));CHKERRQ(ierr);
13114a56b808SFande Kong   ierr = MatMPIAIJSetPreallocation(Cmpi,0,dnz,0,onz);CHKERRQ(ierr);
13124a56b808SFande Kong   ierr = PetscFree2(dnz,onz);CHKERRQ(ierr);
13134a56b808SFande Kong 
13144a56b808SFande Kong   /* attach the supporting struct to Cmpi for reuse */
13154a56b808SFande Kong   c = (Mat_MPIAIJ*)Cmpi->data;
13164a56b808SFande Kong   c->ap           = ptap;
13174a56b808SFande Kong   ptap->duplicate = Cmpi->ops->duplicate;
13184a56b808SFande Kong   ptap->destroy   = Cmpi->ops->destroy;
13194a56b808SFande Kong   ptap->view      = Cmpi->ops->view;
13204a56b808SFande Kong 
13214a56b808SFande Kong   /* Cmpi is not ready for use - assembly will be done by MatPtAPNumeric() */
13224a56b808SFande Kong   Cmpi->assembled        = PETSC_FALSE;
13234a56b808SFande Kong   /* pick an algorithm */
13244a56b808SFande Kong   ierr = PetscOptionsBegin(PetscObjectComm((PetscObject)A),((PetscObject)A)->prefix,"MatPtAP","Mat");CHKERRQ(ierr);
13254a56b808SFande Kong   alg = 0;
13264a56b808SFande Kong   ierr = PetscOptionsEList("-matptap_allatonce_via","PtAP allatonce numeric approach","MatPtAP",algTypes,nalg,algTypes[alg],&alg,&flg);CHKERRQ(ierr);
13274a56b808SFande Kong   ierr = PetscOptionsEnd();CHKERRQ(ierr);
13284a56b808SFande Kong   switch (alg) {
13294a56b808SFande Kong     case 0:
13304a56b808SFande Kong       Cmpi->ops->ptapnumeric = MatPtAPNumeric_MPIAIJ_MPIAIJ_allatonce;
13314a56b808SFande Kong       break;
13324a56b808SFande Kong     case 1:
13334a56b808SFande Kong       Cmpi->ops->ptapnumeric = MatPtAPNumeric_MPIAIJ_MPIAIJ_allatonce_merged;
13344a56b808SFande Kong       break;
13354a56b808SFande Kong     default:
13364a56b808SFande Kong       SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG," Unsupported allatonce numerical algorithm \n");
13374a56b808SFande Kong   }
13384a56b808SFande Kong   Cmpi->ops->destroy     = MatDestroy_MPIAIJ_PtAP;
13394a56b808SFande Kong   Cmpi->ops->view        = MatView_MPIAIJ_PtAP;
13404a56b808SFande Kong   Cmpi->ops->freeintermediatedatastructures = MatFreeIntermediateDataStructures_MPIAIJ_AP;
13414a56b808SFande Kong   *C                     = Cmpi;
13424a56b808SFande Kong   PetscFunctionReturn(0);
13434a56b808SFande Kong }
13444a56b808SFande Kong 
13454a56b808SFande Kong PetscErrorCode MatPtAPSymbolic_MPIAIJ_MPIAIJ_allatonce_merged(Mat A,Mat P,PetscReal fill,Mat *C)
13464a56b808SFande Kong {
13474a56b808SFande Kong   Mat_APMPI           *ptap;
13484a56b808SFande Kong   Mat_MPIAIJ          *p=(Mat_MPIAIJ*)P->data,*c;
13494a56b808SFande Kong   MPI_Comm            comm;
13504a56b808SFande Kong   Mat                 Cmpi;
13514a56b808SFande Kong   Mat_SeqAIJ          *pd,*po;
13524a56b808SFande Kong   MatType             mtype;
13534a56b808SFande Kong   PetscSF             sf;
13544a56b808SFande Kong   PetscSFNode         *iremote;
13554a56b808SFande Kong   PetscInt            rootspacesize,*rootspace,*rootspaceoffsets,nleaves;
13564a56b808SFande Kong   const PetscInt      *rootdegrees;
13574a56b808SFande Kong   PetscHSetI          ht,oht,*hta,*hto,*htd;
13584a56b808SFande Kong   PetscInt            pn,pon,*c_rmtc,i,j,nzi,dnzi,htsize,htosize,*c_rmtj,off,*c_othj,rcvncols,sendncols,*c_rmtoffsets;
13594a56b808SFande Kong   PetscInt            owner,lidx,*rdj,col,pcstart,pcend,*dnz,*onz,am,arstart,arend,*poj,*pdj;
13604a56b808SFande Kong   PetscInt            nalg=2,alg=0;
13614a56b808SFande Kong   PetscBool           flg;
13624a56b808SFande Kong   const char          *algTypes[2] = {"merged","overlapping"};
13634a56b808SFande Kong   PetscErrorCode      ierr;
13644a56b808SFande Kong 
13654a56b808SFande Kong   PetscFunctionBegin;
13664a56b808SFande Kong   ierr = PetscObjectGetComm((PetscObject)A,&comm);CHKERRQ(ierr);
13674a56b808SFande Kong 
13684a56b808SFande Kong   /* Create symbolic parallel matrix Cmpi */
13694a56b808SFande Kong   ierr = MatGetLocalSize(P,NULL,&pn);CHKERRQ(ierr);
13704a56b808SFande Kong   ierr = MatCreate(comm,&Cmpi);CHKERRQ(ierr);
13714a56b808SFande Kong   ierr = MatGetType(A,&mtype);CHKERRQ(ierr);
13724a56b808SFande Kong   ierr = MatSetType(Cmpi,mtype);CHKERRQ(ierr);
13734a56b808SFande Kong   ierr = MatSetSizes(Cmpi,pn,pn,PETSC_DETERMINE,PETSC_DETERMINE);CHKERRQ(ierr);
13744a56b808SFande Kong 
13754a56b808SFande Kong   ierr        = PetscNew(&ptap);CHKERRQ(ierr);
13764a56b808SFande Kong   ptap->reuse = MAT_INITIAL_MATRIX;
13774a56b808SFande Kong   ptap->algType = 3;
13784a56b808SFande Kong 
13794a56b808SFande Kong   /* 0) Get P_oth by taking rows of P (= non-zero cols of local A) from other processors */
13804a56b808SFande Kong   ierr = MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth);CHKERRQ(ierr);
13814a56b808SFande Kong 
13824a56b808SFande Kong   po = (Mat_SeqAIJ*)p->B->data;
13834a56b808SFande Kong   pd = (Mat_SeqAIJ*)p->A->data;
13844a56b808SFande Kong 
13854a56b808SFande Kong   /* This equals to the number of offdiag columns in P */
13864a56b808SFande Kong   ierr = MatGetLocalSize(p->B,NULL,&pon);CHKERRQ(ierr);
13874a56b808SFande Kong   /* offsets */
13884a56b808SFande Kong   ierr = PetscMalloc1(pon+1,&ptap->c_rmti);CHKERRQ(ierr);
13894a56b808SFande Kong   /* The number of columns we will send to remote ranks */
13904a56b808SFande Kong   ierr = PetscMalloc1(pon,&c_rmtc);CHKERRQ(ierr);
13914a56b808SFande Kong   ierr = PetscMalloc1(pon,&hta);CHKERRQ(ierr);
13924a56b808SFande Kong   for (i=0; i<pon; i++) {
13934a56b808SFande Kong     ierr = PetscHSetICreate(&hta[i]);CHKERRQ(ierr);
13944a56b808SFande Kong   }
13954a56b808SFande Kong   ierr = MatGetLocalSize(A,&am,NULL);CHKERRQ(ierr);
13964a56b808SFande Kong   ierr = MatGetOwnershipRange(A,&arstart,&arend);CHKERRQ(ierr);
13974a56b808SFande Kong   /* Create hash table to merge all columns for C(i, :) */
13984a56b808SFande Kong   ierr = PetscHSetICreate(&ht);CHKERRQ(ierr);
13994a56b808SFande Kong   ierr = PetscHSetICreate(&oht);CHKERRQ(ierr);
14004a56b808SFande Kong   ierr = PetscMalloc2(pn,&htd,pn,&hto);CHKERRQ(ierr);
14014a56b808SFande Kong   for (i=0; i<pn; i++) {
14024a56b808SFande Kong     ierr = PetscHSetICreate(&htd[i]);CHKERRQ(ierr);
14034a56b808SFande Kong     ierr = PetscHSetICreate(&hto[i]);CHKERRQ(ierr);
14044a56b808SFande Kong   }
14054a56b808SFande Kong   ptap->c_rmti[0] = 0;
14064a56b808SFande Kong   /* 2) Pass 1: calculate the size for C_rmt (a matrix need to be sent to other processors)  */
14074a56b808SFande Kong   for (i=0; i<am && (pon || pn); i++) {
14084a56b808SFande Kong     /* Form one row of AP */
14094a56b808SFande Kong     ierr = PetscHSetIClear(ht);CHKERRQ(ierr);
14104a56b808SFande Kong     ierr = PetscHSetIClear(oht);CHKERRQ(ierr);
14114a56b808SFande Kong     /* If the off diag is empty, we should not do any calculation */
14124a56b808SFande Kong     nzi = po->i[i+1] - po->i[i];
14134a56b808SFande Kong     dnzi = pd->i[i+1] - pd->i[i];
14144a56b808SFande Kong     if (!nzi && !dnzi) continue;
14154a56b808SFande Kong 
14164a56b808SFande Kong     ierr = MatPtAPSymbolicComputeOneRowOfAP_private(A,P,ptap->P_oth,i,ht,oht);CHKERRQ(ierr);
14174a56b808SFande Kong     ierr = PetscHSetIGetSize(ht,&htsize);CHKERRQ(ierr);
14184a56b808SFande Kong     ierr = PetscHSetIGetSize(oht,&htosize);CHKERRQ(ierr);
14194a56b808SFande Kong     /* If AP is empty, just continue */
14204a56b808SFande Kong     if (!(htsize+htosize)) continue;
14214a56b808SFande Kong 
14224a56b808SFande Kong     /* Form remote C(ii, :) */
14234a56b808SFande Kong     poj = po->j + po->i[i];
14244a56b808SFande Kong     for (j=0; j<nzi; j++) {
14254a56b808SFande Kong       ierr = PetscHSetIAppend(hta[poj[j]],ht);CHKERRQ(ierr);
14264a56b808SFande Kong       ierr = PetscHSetIAppend(hta[poj[j]],oht);CHKERRQ(ierr);
14274a56b808SFande Kong     }
14284a56b808SFande Kong 
14294a56b808SFande Kong     /* Form local C(ii, :) */
14304a56b808SFande Kong     pdj = pd->j + pd->i[i];
14314a56b808SFande Kong     for (j=0; j<dnzi; j++) {
14324a56b808SFande Kong       ierr = PetscHSetIAppend(htd[pdj[j]],ht);CHKERRQ(ierr);
14334a56b808SFande Kong       ierr = PetscHSetIAppend(hto[pdj[j]],oht);CHKERRQ(ierr);
14344a56b808SFande Kong     }
14354a56b808SFande Kong   }
14364a56b808SFande Kong 
14374a56b808SFande Kong   ierr = PetscHSetIDestroy(&ht);CHKERRQ(ierr);
14384a56b808SFande Kong   ierr = PetscHSetIDestroy(&oht);CHKERRQ(ierr);
14394a56b808SFande Kong 
14404a56b808SFande Kong   for (i=0; i<pon; i++) {
14414a56b808SFande Kong     ierr = PetscHSetIGetSize(hta[i],&htsize);CHKERRQ(ierr);
14424a56b808SFande Kong     ptap->c_rmti[i+1] = ptap->c_rmti[i] + htsize;
14434a56b808SFande Kong     c_rmtc[i] = htsize;
14444a56b808SFande Kong   }
14454a56b808SFande Kong 
14464a56b808SFande Kong   ierr = PetscMalloc1(ptap->c_rmti[pon],&c_rmtj);CHKERRQ(ierr);
14474a56b808SFande Kong 
14484a56b808SFande Kong   for (i=0; i<pon; i++) {
14494a56b808SFande Kong     off = 0;
14504a56b808SFande Kong     ierr = PetscHSetIGetElems(hta[i],&off,c_rmtj+ptap->c_rmti[i]);CHKERRQ(ierr);
14514a56b808SFande Kong     ierr = PetscHSetIDestroy(&hta[i]);CHKERRQ(ierr);
14524a56b808SFande Kong   }
14534a56b808SFande Kong   ierr = PetscFree(hta);CHKERRQ(ierr);
14544a56b808SFande Kong 
14554a56b808SFande Kong   ierr = PetscCalloc1(pon,&iremote);CHKERRQ(ierr);
14564a56b808SFande Kong   for (i=0; i<pon; i++) {
1457*ef7a94f2SFande Kong     owner = 0; lidx = 0;
14584a56b808SFande Kong     ierr = PetscLayoutFindOwnerIndex(P->cmap,p->garray[i],&owner,&lidx);CHKERRQ(ierr);
14594a56b808SFande Kong     iremote[i].index = lidx;
14604a56b808SFande Kong     iremote[i].rank  = owner;
14614a56b808SFande Kong   }
14624a56b808SFande Kong 
14634a56b808SFande Kong   ierr = PetscSFCreate(comm,&sf);CHKERRQ(ierr);
14644a56b808SFande Kong   ierr = PetscSFSetGraph(sf,pn,pon,NULL,PETSC_OWN_POINTER,iremote,PETSC_OWN_POINTER);CHKERRQ(ierr);
14654a56b808SFande Kong   /* Reorder ranks properly so that the data handled by gather and scatter have the same order */
14664a56b808SFande Kong   ierr = PetscSFSetRankOrder(sf,PETSC_TRUE);CHKERRQ(ierr);
14674a56b808SFande Kong   ierr = PetscSFSetFromOptions(sf);CHKERRQ(ierr);
14684a56b808SFande Kong   ierr = PetscSFSetUp(sf);CHKERRQ(ierr);
14694a56b808SFande Kong   /* How many neighbors have contributions to my rows? */
14704a56b808SFande Kong   ierr = PetscSFComputeDegreeBegin(sf,&rootdegrees);CHKERRQ(ierr);
14714a56b808SFande Kong   ierr = PetscSFComputeDegreeEnd(sf,&rootdegrees);CHKERRQ(ierr);
14724a56b808SFande Kong   rootspacesize = 0;
14734a56b808SFande Kong   for (i = 0; i < pn; i++) {
14744a56b808SFande Kong     rootspacesize += rootdegrees[i];
14754a56b808SFande Kong   }
14764a56b808SFande Kong   ierr = PetscCalloc1(rootspacesize,&rootspace);CHKERRQ(ierr);
14774a56b808SFande Kong   ierr = PetscCalloc1(rootspacesize+1,&rootspaceoffsets);CHKERRQ(ierr);
14784a56b808SFande Kong   /* Get information from leaves
14794a56b808SFande Kong    * Number of columns other people contribute to my rows
14804a56b808SFande Kong    * */
14814a56b808SFande Kong   ierr = PetscSFGatherBegin(sf,MPIU_INT,c_rmtc,rootspace);CHKERRQ(ierr);
14824a56b808SFande Kong   ierr = PetscSFGatherEnd(sf,MPIU_INT,c_rmtc,rootspace);CHKERRQ(ierr);
14834a56b808SFande Kong   ierr = PetscFree(c_rmtc);CHKERRQ(ierr);
14844a56b808SFande Kong   ierr = PetscCalloc1(pn+1,&ptap->c_othi);CHKERRQ(ierr);
14854a56b808SFande Kong   /* The number of columns is received for each row */
14864a56b808SFande Kong   ptap->c_othi[0] = 0;
14874a56b808SFande Kong   rootspacesize = 0;
14884a56b808SFande Kong   rootspaceoffsets[0] = 0;
14894a56b808SFande Kong   for (i = 0; i < pn; i++) {
14904a56b808SFande Kong     rcvncols = 0;
14914a56b808SFande Kong     for (j = 0; j<rootdegrees[i]; j++) {
14924a56b808SFande Kong       rcvncols += rootspace[rootspacesize];
14934a56b808SFande Kong       rootspaceoffsets[rootspacesize+1] = rootspaceoffsets[rootspacesize] + rootspace[rootspacesize];
14944a56b808SFande Kong       rootspacesize++;
14954a56b808SFande Kong     }
14964a56b808SFande Kong     ptap->c_othi[i+1] = ptap->c_othi[i] + rcvncols;
14974a56b808SFande Kong   }
14984a56b808SFande Kong   ierr = PetscFree(rootspace);CHKERRQ(ierr);
14994a56b808SFande Kong 
15004a56b808SFande Kong   ierr = PetscCalloc1(pon,&c_rmtoffsets);CHKERRQ(ierr);
15014a56b808SFande Kong   ierr = PetscSFScatterBegin(sf,MPIU_INT,rootspaceoffsets,c_rmtoffsets);CHKERRQ(ierr);
15024a56b808SFande Kong   ierr = PetscSFScatterEnd(sf,MPIU_INT,rootspaceoffsets,c_rmtoffsets);CHKERRQ(ierr);
15034a56b808SFande Kong   ierr = PetscSFDestroy(&sf);CHKERRQ(ierr);
15044a56b808SFande Kong   ierr = PetscFree(rootspaceoffsets);CHKERRQ(ierr);
15054a56b808SFande Kong 
15064a56b808SFande Kong   ierr = PetscCalloc1(ptap->c_rmti[pon],&iremote);CHKERRQ(ierr);
15074a56b808SFande Kong   nleaves = 0;
15084a56b808SFande Kong   for (i = 0; i<pon; i++) {
1509*ef7a94f2SFande Kong     owner = 0; lidx = 0;
15104a56b808SFande Kong     ierr = PetscLayoutFindOwnerIndex(P->cmap,p->garray[i],&owner,&lidx);CHKERRQ(ierr);
15114a56b808SFande Kong     sendncols = ptap->c_rmti[i+1] - ptap->c_rmti[i];
15124a56b808SFande Kong     for (j=0; j<sendncols; j++) {
15134a56b808SFande Kong       iremote[nleaves].rank = owner;
15144a56b808SFande Kong       iremote[nleaves++].index = c_rmtoffsets[i] + j;
15154a56b808SFande Kong     }
15164a56b808SFande Kong   }
15174a56b808SFande Kong   ierr = PetscFree(c_rmtoffsets);CHKERRQ(ierr);
15184a56b808SFande Kong   ierr = PetscCalloc1(ptap->c_othi[pn],&c_othj);CHKERRQ(ierr);
15194a56b808SFande Kong 
15204a56b808SFande Kong   ierr = PetscSFCreate(comm,&ptap->sf);CHKERRQ(ierr);
15214a56b808SFande Kong   ierr = PetscSFSetGraph(ptap->sf,ptap->c_othi[pn],nleaves,NULL,PETSC_OWN_POINTER,iremote,PETSC_OWN_POINTER);CHKERRQ(ierr);
15224a56b808SFande Kong   ierr = PetscSFSetFromOptions(ptap->sf);CHKERRQ(ierr);
15234a56b808SFande Kong   /* One to one map */
15244a56b808SFande Kong   ierr = PetscSFReduceBegin(ptap->sf,MPIU_INT,c_rmtj,c_othj,MPIU_REPLACE);CHKERRQ(ierr);
15254a56b808SFande Kong   /* Get remote data */
15264a56b808SFande Kong   ierr = PetscSFReduceEnd(ptap->sf,MPIU_INT,c_rmtj,c_othj,MPIU_REPLACE);CHKERRQ(ierr);
15274a56b808SFande Kong   ierr = PetscFree(c_rmtj);CHKERRQ(ierr);
15284a56b808SFande Kong   ierr = PetscCalloc2(pn,&dnz,pn,&onz);CHKERRQ(ierr);
15294a56b808SFande Kong   ierr = MatGetOwnershipRangeColumn(P,&pcstart,&pcend);CHKERRQ(ierr);
15304a56b808SFande Kong 
15314a56b808SFande Kong   for (i = 0; i < pn; i++) {
15324a56b808SFande Kong     nzi = ptap->c_othi[i+1] - ptap->c_othi[i];
15334a56b808SFande Kong     rdj = c_othj + ptap->c_othi[i];
15344a56b808SFande Kong     for (j = 0; j < nzi; j++) {
15354a56b808SFande Kong       col =  rdj[j];
15364a56b808SFande Kong       /* diag part */
15374a56b808SFande Kong       if (col>=pcstart && col<pcend) {
15384a56b808SFande Kong         ierr = PetscHSetIAdd(htd[i],col);CHKERRQ(ierr);
15394a56b808SFande Kong       } else { /* off diag */
15404a56b808SFande Kong         ierr = PetscHSetIAdd(hto[i],col);CHKERRQ(ierr);
15414a56b808SFande Kong       }
15424a56b808SFande Kong     }
15434a56b808SFande Kong     ierr = PetscHSetIGetSize(htd[i],&htsize);CHKERRQ(ierr);
15444a56b808SFande Kong     dnz[i] = htsize;
15454a56b808SFande Kong     ierr = PetscHSetIDestroy(&htd[i]);CHKERRQ(ierr);
15464a56b808SFande Kong     ierr = PetscHSetIGetSize(hto[i],&htsize);CHKERRQ(ierr);
15474a56b808SFande Kong     onz[i] = htsize;
15484a56b808SFande Kong     ierr = PetscHSetIDestroy(&hto[i]);CHKERRQ(ierr);
15494a56b808SFande Kong   }
15504a56b808SFande Kong 
15514a56b808SFande Kong   ierr = PetscFree2(htd,hto);CHKERRQ(ierr);
15524a56b808SFande Kong   ierr = PetscFree(c_othj);CHKERRQ(ierr);
15534a56b808SFande Kong 
15544a56b808SFande Kong   /* local sizes and preallocation */
15554a56b808SFande Kong   ierr = MatSetSizes(Cmpi,pn,pn,PETSC_DETERMINE,PETSC_DETERMINE);CHKERRQ(ierr);
15564a56b808SFande Kong   ierr = MatSetBlockSizes(Cmpi,PetscAbs(P->cmap->bs),PetscAbs(P->cmap->bs));CHKERRQ(ierr);
15574a56b808SFande Kong   ierr = MatMPIAIJSetPreallocation(Cmpi,0,dnz,0,onz);CHKERRQ(ierr);
15584a56b808SFande Kong   ierr = PetscFree2(dnz,onz);CHKERRQ(ierr);
15594a56b808SFande Kong 
15604a56b808SFande Kong   /* attach the supporting struct to Cmpi for reuse */
15614a56b808SFande Kong   c = (Mat_MPIAIJ*)Cmpi->data;
15624a56b808SFande Kong   c->ap           = ptap;
15634a56b808SFande Kong   ptap->duplicate = Cmpi->ops->duplicate;
15644a56b808SFande Kong   ptap->destroy   = Cmpi->ops->destroy;
15654a56b808SFande Kong   ptap->view      = Cmpi->ops->view;
15664a56b808SFande Kong 
15674a56b808SFande Kong   /* Cmpi is not ready for use - assembly will be done by MatPtAPNumeric() */
15684a56b808SFande Kong   Cmpi->assembled        = PETSC_FALSE;
15694a56b808SFande Kong   /* pick an algorithm */
15704a56b808SFande Kong   ierr = PetscOptionsBegin(PetscObjectComm((PetscObject)A),((PetscObject)A)->prefix,"MatPtAP","Mat");CHKERRQ(ierr);
15714a56b808SFande Kong   alg = 0;
15724a56b808SFande Kong   ierr = PetscOptionsEList("-matptap_allatonce_via","PtAP allatonce numeric approach","MatPtAP",algTypes,nalg,algTypes[alg],&alg,&flg);CHKERRQ(ierr);
15734a56b808SFande Kong   ierr = PetscOptionsEnd();CHKERRQ(ierr);
15744a56b808SFande Kong   switch (alg) {
15754a56b808SFande Kong     case 0:
15764a56b808SFande Kong       Cmpi->ops->ptapnumeric = MatPtAPNumeric_MPIAIJ_MPIAIJ_allatonce_merged;
15774a56b808SFande Kong       break;
15784a56b808SFande Kong     case 1:
15794a56b808SFande Kong       Cmpi->ops->ptapnumeric = MatPtAPNumeric_MPIAIJ_MPIAIJ_allatonce;
15804a56b808SFande Kong       break;
15814a56b808SFande Kong     default:
15824a56b808SFande Kong       SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG," Unsupported allatonce numerical algorithm \n");
15834a56b808SFande Kong   }
15844a56b808SFande Kong   Cmpi->ops->destroy     = MatDestroy_MPIAIJ_PtAP;
15854a56b808SFande Kong   Cmpi->ops->view        = MatView_MPIAIJ_PtAP;
15864a56b808SFande Kong   Cmpi->ops->freeintermediatedatastructures = MatFreeIntermediateDataStructures_MPIAIJ_AP;
15874a56b808SFande Kong   *C                     = Cmpi;
15884a56b808SFande Kong   PetscFunctionReturn(0);
15894a56b808SFande Kong }
15904a56b808SFande Kong 
1591e953e47cSHong Zhang PetscErrorCode MatPtAPSymbolic_MPIAIJ_MPIAIJ(Mat A,Mat P,PetscReal fill,Mat *C)
1592e953e47cSHong Zhang {
1593e953e47cSHong Zhang   PetscErrorCode      ierr;
15943cdca5ebSHong Zhang   Mat_APMPI           *ptap;
1595e953e47cSHong Zhang   Mat_MPIAIJ          *a=(Mat_MPIAIJ*)A->data,*p=(Mat_MPIAIJ*)P->data,*c;
1596e953e47cSHong Zhang   MPI_Comm            comm;
1597e953e47cSHong Zhang   PetscMPIInt         size,rank;
1598e953e47cSHong Zhang   Mat                 Cmpi;
1599e953e47cSHong Zhang   PetscFreeSpaceList  free_space=NULL,current_space=NULL;
1600e953e47cSHong Zhang   PetscInt            am=A->rmap->n,pm=P->rmap->n,pN=P->cmap->N,pn=P->cmap->n;
1601e953e47cSHong Zhang   PetscInt            *lnk,i,k,pnz,row,nsend;
1602e953e47cSHong Zhang   PetscBT             lnkbt;
1603e953e47cSHong Zhang   PetscMPIInt         tagi,tagj,*len_si,*len_s,*len_ri,icompleted=0,nrecv;
1604e953e47cSHong Zhang   PetscInt            **buf_rj,**buf_ri,**buf_ri_k;
1605e953e47cSHong Zhang   PetscInt            len,proc,*dnz,*onz,*owners,nzi,nspacedouble;
1606e953e47cSHong Zhang   PetscInt            nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci;
1607e953e47cSHong Zhang   MPI_Request         *swaits,*rwaits;
1608e953e47cSHong Zhang   MPI_Status          *sstatus,rstatus;
1609e953e47cSHong Zhang   PetscLayout         rowmap;
1610e953e47cSHong Zhang   PetscInt            *owners_co,*coi,*coj;    /* i and j array of (p->B)^T*A*P - used in the communication */
1611e953e47cSHong Zhang   PetscMPIInt         *len_r,*id_r;    /* array of length of comm->size, store send/recv matrix values */
1612e953e47cSHong Zhang   PetscInt            *api,*apj,*Jptr,apnz,*prmap=p->garray,con,j,ap_rmax=0,Crmax,*aj,*ai,*pi;
1613ec07b8f8SHong Zhang   Mat_SeqAIJ          *p_loc,*p_oth=NULL,*ad=(Mat_SeqAIJ*)(a->A)->data,*ao=NULL,*c_loc,*c_oth;
1614e953e47cSHong Zhang   PetscScalar         *apv;
1615e953e47cSHong Zhang   PetscTable          ta;
1616b92f168fSBarry Smith   MatType             mtype;
1617e83e5f86SFande Kong   const char          *prefix;
1618e953e47cSHong Zhang #if defined(PETSC_USE_INFO)
1619e953e47cSHong Zhang   PetscReal           apfill;
1620e953e47cSHong Zhang #endif
1621e953e47cSHong Zhang 
1622e953e47cSHong Zhang   PetscFunctionBegin;
1623b4e8d1b6SHong Zhang   ierr = PetscObjectGetComm((PetscObject)A,&comm);CHKERRQ(ierr);
1624e953e47cSHong Zhang   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
1625e953e47cSHong Zhang   ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
1626e953e47cSHong Zhang 
162752f7967eSHong Zhang   if (size > 1) ao = (Mat_SeqAIJ*)(a->B)->data;
1628ec07b8f8SHong Zhang 
1629e953e47cSHong Zhang   /* create symbolic parallel matrix Cmpi */
1630e953e47cSHong Zhang   ierr = MatCreate(comm,&Cmpi);CHKERRQ(ierr);
1631b92f168fSBarry Smith   ierr = MatGetType(A,&mtype);CHKERRQ(ierr);
1632b92f168fSBarry Smith   ierr = MatSetType(Cmpi,mtype);CHKERRQ(ierr);
1633e953e47cSHong Zhang 
1634e953e47cSHong Zhang   /* Do dense axpy in MatPtAPNumeric_MPIAIJ_MPIAIJ() */
1635e953e47cSHong Zhang   Cmpi->ops->ptapnumeric = MatPtAPNumeric_MPIAIJ_MPIAIJ;
1636e953e47cSHong Zhang 
16373cdca5ebSHong Zhang   /* create struct Mat_APMPI and attached it to C later */
163815a3b8e2SHong Zhang   ierr        = PetscNew(&ptap);CHKERRQ(ierr);
163915a3b8e2SHong Zhang   ptap->reuse = MAT_INITIAL_MATRIX;
1640a4ffb656SHong Zhang   ptap->algType = 1;
164115a3b8e2SHong Zhang 
164215a3b8e2SHong Zhang   /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */
164315a3b8e2SHong Zhang   ierr = MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth);CHKERRQ(ierr);
164415a3b8e2SHong Zhang   /* get P_loc by taking all local rows of P */
164515a3b8e2SHong Zhang   ierr = MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc);CHKERRQ(ierr);
164615a3b8e2SHong Zhang 
164767a12041SHong Zhang   /* (0) compute Rd = Pd^T, Ro = Po^T  */
164867a12041SHong Zhang   /* --------------------------------- */
1649419ecdd9Sandi selinger   ierr = MatTranspose(p->A,MAT_INITIAL_MATRIX,&ptap->Rd);CHKERRQ(ierr);
1650419ecdd9Sandi selinger   ierr = MatTranspose(p->B,MAT_INITIAL_MATRIX,&ptap->Ro);CHKERRQ(ierr);
165115a3b8e2SHong Zhang 
165267a12041SHong Zhang   /* (1) compute symbolic AP = A_loc*P = Ad*P_loc + Ao*P_oth (api,apj) */
165367a12041SHong Zhang   /* ----------------------------------------------------------------- */
165467a12041SHong Zhang   p_loc  = (Mat_SeqAIJ*)(ptap->P_loc)->data;
165552f7967eSHong Zhang   if (ptap->P_oth) p_oth  = (Mat_SeqAIJ*)(ptap->P_oth)->data;
1656445158ffSHong Zhang 
16579a6dcab7SHong Zhang   /* create and initialize a linked list */
165845d00d1dSHong Zhang   ierr = PetscTableCreate(pn,pN,&ta);CHKERRQ(ierr); /* for compute AP_loc and Cmpi */
16594b38b95cSHong Zhang   MatRowMergeMax_SeqAIJ(p_loc,ptap->P_loc->rmap->N,ta);
16604b38b95cSHong Zhang   MatRowMergeMax_SeqAIJ(p_oth,ptap->P_oth->rmap->N,ta);
166178d30b94SHong Zhang   ierr = PetscTableGetCount(ta,&Crmax);CHKERRQ(ierr); /* Crmax = nnz(sum of Prows) */
1662d0e9a2f7SHong Zhang   /* printf("[%d] est %d, Crmax %d; pN %d\n",rank,5*(p_loc->rmax+p_oth->rmax + (PetscInt)(1.e-2*pN)),Crmax,pN); */
166378d30b94SHong Zhang 
166478d30b94SHong Zhang   ierr = PetscLLCondensedCreate(Crmax,pN,&lnk,&lnkbt);CHKERRQ(ierr);
1665445158ffSHong Zhang 
16668cb82516SHong Zhang   /* Initial FreeSpace size is fill*(nnz(A) + nnz(P)) */
1667ec07b8f8SHong Zhang   if (ao) {
1668f91af8c7SBarry Smith     ierr = PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(ad->i[am],PetscIntSumTruncate(ao->i[am],p_loc->i[pm]))),&free_space);CHKERRQ(ierr);
1669ec07b8f8SHong Zhang   } else {
1670ec07b8f8SHong Zhang     ierr = PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(ad->i[am],p_loc->i[pm])),&free_space);CHKERRQ(ierr);
1671ec07b8f8SHong Zhang   }
1672445158ffSHong Zhang   current_space = free_space;
167367a12041SHong Zhang   nspacedouble  = 0;
167467a12041SHong Zhang 
167567a12041SHong Zhang   ierr   = PetscMalloc1(am+1,&api);CHKERRQ(ierr);
167667a12041SHong Zhang   api[0] = 0;
1677445158ffSHong Zhang   for (i=0; i<am; i++) {
167867a12041SHong Zhang     /* diagonal portion: Ad[i,:]*P */
167967a12041SHong Zhang     ai = ad->i; pi = p_loc->i;
168067a12041SHong Zhang     nzi = ai[i+1] - ai[i];
168167a12041SHong Zhang     aj  = ad->j + ai[i];
1682445158ffSHong Zhang     for (j=0; j<nzi; j++) {
1683445158ffSHong Zhang       row  = aj[j];
168467a12041SHong Zhang       pnz  = pi[row+1] - pi[row];
168567a12041SHong Zhang       Jptr = p_loc->j + pi[row];
1686445158ffSHong Zhang       /* add non-zero cols of P into the sorted linked list lnk */
1687445158ffSHong Zhang       ierr = PetscLLCondensedAddSorted(pnz,Jptr,lnk,lnkbt);CHKERRQ(ierr);
1688445158ffSHong Zhang     }
168967a12041SHong Zhang     /* off-diagonal portion: Ao[i,:]*P */
1690ec07b8f8SHong Zhang     if (ao) {
169167a12041SHong Zhang       ai = ao->i; pi = p_oth->i;
169267a12041SHong Zhang       nzi = ai[i+1] - ai[i];
169367a12041SHong Zhang       aj  = ao->j + ai[i];
1694445158ffSHong Zhang       for (j=0; j<nzi; j++) {
1695445158ffSHong Zhang         row  = aj[j];
169667a12041SHong Zhang         pnz  = pi[row+1] - pi[row];
169767a12041SHong Zhang         Jptr = p_oth->j + pi[row];
1698445158ffSHong Zhang         ierr = PetscLLCondensedAddSorted(pnz,Jptr,lnk,lnkbt);CHKERRQ(ierr);
1699445158ffSHong Zhang       }
1700ec07b8f8SHong Zhang     }
1701445158ffSHong Zhang     apnz     = lnk[0];
1702445158ffSHong Zhang     api[i+1] = api[i] + apnz;
1703445158ffSHong Zhang     if (ap_rmax < apnz) ap_rmax = apnz;
1704445158ffSHong Zhang 
1705445158ffSHong Zhang     /* if free space is not available, double the total space in the list */
1706445158ffSHong Zhang     if (current_space->local_remaining<apnz) {
1707f91af8c7SBarry Smith       ierr = PetscFreeSpaceGet(PetscIntSumTruncate(apnz,current_space->total_array_size),&current_space);CHKERRQ(ierr);
1708445158ffSHong Zhang       nspacedouble++;
1709445158ffSHong Zhang     }
1710445158ffSHong Zhang 
1711445158ffSHong Zhang     /* Copy data into free space, then initialize lnk */
1712445158ffSHong Zhang     ierr = PetscLLCondensedClean(pN,apnz,current_space->array,lnk,lnkbt);CHKERRQ(ierr);
1713445158ffSHong Zhang 
1714445158ffSHong Zhang     current_space->array           += apnz;
1715445158ffSHong Zhang     current_space->local_used      += apnz;
1716445158ffSHong Zhang     current_space->local_remaining -= apnz;
1717445158ffSHong Zhang   }
1718681d504bSHong Zhang   /* Allocate space for apj and apv, initialize apj, and */
1719445158ffSHong Zhang   /* destroy list of free space and other temporary array(s) */
1720445158ffSHong Zhang   ierr   = PetscMalloc2(api[am],&apj,api[am],&apv);CHKERRQ(ierr);
1721445158ffSHong Zhang   ierr   = PetscFreeSpaceContiguous(&free_space,apj);CHKERRQ(ierr);
17229a6dcab7SHong Zhang   ierr = PetscLLDestroy(lnk,lnkbt);CHKERRQ(ierr);
17239a6dcab7SHong Zhang 
1724aa690a28SHong Zhang   /* Create AP_loc for reuse */
1725445158ffSHong Zhang   ierr = MatCreateSeqAIJWithArrays(PETSC_COMM_SELF,am,pN,api,apj,apv,&ptap->AP_loc);CHKERRQ(ierr);
1726aa690a28SHong Zhang 
1727aa690a28SHong Zhang #if defined(PETSC_USE_INFO)
1728ec07b8f8SHong Zhang   if (ao) {
1729aa690a28SHong Zhang     apfill = (PetscReal)api[am]/(ad->i[am]+ao->i[am]+p_loc->i[pm]+1);
1730ec07b8f8SHong Zhang   } else {
1731ec07b8f8SHong Zhang     apfill = (PetscReal)api[am]/(ad->i[am]+p_loc->i[pm]+1);
1732ec07b8f8SHong Zhang   }
1733aa690a28SHong Zhang   ptap->AP_loc->info.mallocs           = nspacedouble;
1734aa690a28SHong Zhang   ptap->AP_loc->info.fill_ratio_given  = fill;
1735aa690a28SHong Zhang   ptap->AP_loc->info.fill_ratio_needed = apfill;
1736aa690a28SHong Zhang 
1737aa690a28SHong Zhang   if (api[am]) {
1738b11c1ec8SHong Zhang     ierr = PetscInfo3(ptap->AP_loc,"Nonscalable algorithm, AP_loc reallocs %D; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)apfill);CHKERRQ(ierr);
1739aa690a28SHong Zhang     ierr = PetscInfo1(ptap->AP_loc,"Use MatPtAP(A,B,MatReuse,%g,&C) for best AP_loc performance.;\n",(double)apfill);CHKERRQ(ierr);
1740aa690a28SHong Zhang   } else {
1741b11c1ec8SHong Zhang     ierr = PetscInfo(ptap->AP_loc,"Nonscalable algorithm, AP_loc is empty \n");CHKERRQ(ierr);
1742aa690a28SHong Zhang   }
1743aa690a28SHong Zhang #endif
1744aa690a28SHong Zhang 
1745681d504bSHong Zhang   /* (2-1) compute symbolic Co = Ro*AP_loc  */
1746681d504bSHong Zhang   /* ------------------------------------ */
1747e83e5f86SFande Kong   ierr = MatGetOptionsPrefix(A,&prefix);CHKERRQ(ierr);
1748e83e5f86SFande Kong   ierr = MatSetOptionsPrefix(ptap->Ro,prefix);CHKERRQ(ierr);
1749e83e5f86SFande Kong   ierr = MatAppendOptionsPrefix(ptap->Ro,"inner_offdiag_");CHKERRQ(ierr);
175067a12041SHong Zhang   ierr = MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Ro,ptap->AP_loc,fill,&ptap->C_oth);CHKERRQ(ierr);
175115a3b8e2SHong Zhang 
175267a12041SHong Zhang   /* (3) send coj of C_oth to other processors  */
175367a12041SHong Zhang   /* ------------------------------------------ */
175415a3b8e2SHong Zhang   /* determine row ownership */
1755445158ffSHong Zhang   ierr = PetscLayoutCreate(comm,&rowmap);CHKERRQ(ierr);
1756445158ffSHong Zhang   rowmap->n  = pn;
1757445158ffSHong Zhang   rowmap->bs = 1;
1758445158ffSHong Zhang   ierr   = PetscLayoutSetUp(rowmap);CHKERRQ(ierr);
1759445158ffSHong Zhang   owners = rowmap->range;
176015a3b8e2SHong Zhang 
176115a3b8e2SHong Zhang   /* determine the number of messages to send, their lengths */
17628cb82516SHong Zhang   ierr = PetscMalloc4(size,&len_s,size,&len_si,size,&sstatus,size+2,&owners_co);CHKERRQ(ierr);
17638cb82516SHong Zhang   ierr = PetscMemzero(len_s,size*sizeof(PetscMPIInt));CHKERRQ(ierr);
176415a3b8e2SHong Zhang   ierr = PetscMemzero(len_si,size*sizeof(PetscMPIInt));CHKERRQ(ierr);
176515a3b8e2SHong Zhang 
176667a12041SHong Zhang   c_oth = (Mat_SeqAIJ*)ptap->C_oth->data;
176767a12041SHong Zhang   coi   = c_oth->i; coj = c_oth->j;
176867a12041SHong Zhang   con   = ptap->C_oth->rmap->n;
176915a3b8e2SHong Zhang   proc  = 0;
177067a12041SHong Zhang   for (i=0; i<con; i++) {
177115a3b8e2SHong Zhang     while (prmap[i] >= owners[proc+1]) proc++;
177215a3b8e2SHong Zhang     len_si[proc]++;               /* num of rows in Co(=Pt*AP) to be sent to [proc] */
177315a3b8e2SHong Zhang     len_s[proc] += coi[i+1] - coi[i]; /* num of nonzeros in Co to be sent to [proc] */
177415a3b8e2SHong Zhang   }
177515a3b8e2SHong Zhang 
177615a3b8e2SHong Zhang   len          = 0; /* max length of buf_si[], see (4) */
177715a3b8e2SHong Zhang   owners_co[0] = 0;
177867a12041SHong Zhang   nsend        = 0;
177915a3b8e2SHong Zhang   for (proc=0; proc<size; proc++) {
178015a3b8e2SHong Zhang     owners_co[proc+1] = owners_co[proc] + len_si[proc];
178115a3b8e2SHong Zhang     if (len_s[proc]) {
1782445158ffSHong Zhang       nsend++;
178315a3b8e2SHong Zhang       len_si[proc] = 2*(len_si[proc] + 1); /* length of buf_si to be sent to [proc] */
178415a3b8e2SHong Zhang       len         += len_si[proc];
178515a3b8e2SHong Zhang     }
178615a3b8e2SHong Zhang   }
178715a3b8e2SHong Zhang 
178815a3b8e2SHong Zhang   /* determine the number and length of messages to receive for coi and coj  */
1789445158ffSHong Zhang   ierr = PetscGatherNumberOfMessages(comm,NULL,len_s,&nrecv);CHKERRQ(ierr);
1790445158ffSHong Zhang   ierr = PetscGatherMessageLengths2(comm,nsend,nrecv,len_s,len_si,&id_r,&len_r,&len_ri);CHKERRQ(ierr);
179115a3b8e2SHong Zhang 
179215a3b8e2SHong Zhang   /* post the Irecv and Isend of coj */
179315a3b8e2SHong Zhang   ierr = PetscCommGetNewTag(comm,&tagj);CHKERRQ(ierr);
1794445158ffSHong Zhang   ierr = PetscPostIrecvInt(comm,tagj,nrecv,id_r,len_r,&buf_rj,&rwaits);CHKERRQ(ierr);
1795445158ffSHong Zhang   ierr = PetscMalloc1(nsend+1,&swaits);CHKERRQ(ierr);
179615a3b8e2SHong Zhang   for (proc=0, k=0; proc<size; proc++) {
179715a3b8e2SHong Zhang     if (!len_s[proc]) continue;
179815a3b8e2SHong Zhang     i    = owners_co[proc];
179915a3b8e2SHong Zhang     ierr = MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tagj,comm,swaits+k);CHKERRQ(ierr);
180015a3b8e2SHong Zhang     k++;
180115a3b8e2SHong Zhang   }
180215a3b8e2SHong Zhang 
1803681d504bSHong Zhang   /* (2-2) compute symbolic C_loc = Rd*AP_loc */
1804681d504bSHong Zhang   /* ---------------------------------------- */
1805e83e5f86SFande Kong   ierr = MatSetOptionsPrefix(ptap->Rd,prefix);CHKERRQ(ierr);
1806e83e5f86SFande Kong   ierr = MatAppendOptionsPrefix(ptap->Rd,"inner_diag_");CHKERRQ(ierr);
1807681d504bSHong Zhang   ierr = MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Rd,ptap->AP_loc,fill,&ptap->C_loc);CHKERRQ(ierr);
1808681d504bSHong Zhang   c_loc = (Mat_SeqAIJ*)ptap->C_loc->data;
1809681d504bSHong Zhang 
1810681d504bSHong Zhang   /* receives coj are complete */
1811445158ffSHong Zhang   for (i=0; i<nrecv; i++) {
1812445158ffSHong Zhang     ierr = MPI_Waitany(nrecv,rwaits,&icompleted,&rstatus);CHKERRQ(ierr);
181315a3b8e2SHong Zhang   }
181415a3b8e2SHong Zhang   ierr = PetscFree(rwaits);CHKERRQ(ierr);
1815445158ffSHong Zhang   if (nsend) {ierr = MPI_Waitall(nsend,swaits,sstatus);CHKERRQ(ierr);}
181615a3b8e2SHong Zhang 
181778d30b94SHong Zhang   /* add received column indices into ta to update Crmax */
181878d30b94SHong Zhang   for (k=0; k<nrecv; k++) {/* k-th received message */
181978d30b94SHong Zhang     Jptr = buf_rj[k];
182078d30b94SHong Zhang     for (j=0; j<len_r[k]; j++) {
182178d30b94SHong Zhang       ierr = PetscTableAdd(ta,*(Jptr+j)+1,1,INSERT_VALUES);CHKERRQ(ierr);
182278d30b94SHong Zhang     }
182378d30b94SHong Zhang   }
182478d30b94SHong Zhang   ierr = PetscTableGetCount(ta,&Crmax);CHKERRQ(ierr);
182578d30b94SHong Zhang   ierr = PetscTableDestroy(&ta);CHKERRQ(ierr);
18269a6dcab7SHong Zhang 
182715a3b8e2SHong Zhang   /* (4) send and recv coi */
182815a3b8e2SHong Zhang   /*-----------------------*/
182915a3b8e2SHong Zhang   ierr   = PetscCommGetNewTag(comm,&tagi);CHKERRQ(ierr);
1830445158ffSHong Zhang   ierr   = PetscPostIrecvInt(comm,tagi,nrecv,id_r,len_ri,&buf_ri,&rwaits);CHKERRQ(ierr);
183115a3b8e2SHong Zhang   ierr   = PetscMalloc1(len+1,&buf_s);CHKERRQ(ierr);
183215a3b8e2SHong Zhang   buf_si = buf_s;  /* points to the beginning of k-th msg to be sent */
183315a3b8e2SHong Zhang   for (proc=0,k=0; proc<size; proc++) {
183415a3b8e2SHong Zhang     if (!len_s[proc]) continue;
183515a3b8e2SHong Zhang     /* form outgoing message for i-structure:
183615a3b8e2SHong Zhang          buf_si[0]:                 nrows to be sent
183715a3b8e2SHong Zhang                [1:nrows]:           row index (global)
183815a3b8e2SHong Zhang                [nrows+1:2*nrows+1]: i-structure index
183915a3b8e2SHong Zhang     */
184015a3b8e2SHong Zhang     /*-------------------------------------------*/
184115a3b8e2SHong Zhang     nrows       = len_si[proc]/2 - 1; /* num of rows in Co to be sent to [proc] */
184215a3b8e2SHong Zhang     buf_si_i    = buf_si + nrows+1;
184315a3b8e2SHong Zhang     buf_si[0]   = nrows;
184415a3b8e2SHong Zhang     buf_si_i[0] = 0;
184515a3b8e2SHong Zhang     nrows       = 0;
184615a3b8e2SHong Zhang     for (i=owners_co[proc]; i<owners_co[proc+1]; i++) {
184715a3b8e2SHong Zhang       nzi = coi[i+1] - coi[i];
184815a3b8e2SHong Zhang       buf_si_i[nrows+1] = buf_si_i[nrows] + nzi;  /* i-structure */
184915a3b8e2SHong Zhang       buf_si[nrows+1]   = prmap[i] -owners[proc]; /* local row index */
185015a3b8e2SHong Zhang       nrows++;
185115a3b8e2SHong Zhang     }
185215a3b8e2SHong Zhang     ierr = MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tagi,comm,swaits+k);CHKERRQ(ierr);
185315a3b8e2SHong Zhang     k++;
185415a3b8e2SHong Zhang     buf_si += len_si[proc];
185515a3b8e2SHong Zhang   }
1856681d504bSHong Zhang   for (i=0; i<nrecv; i++) {
1857445158ffSHong Zhang     ierr = MPI_Waitany(nrecv,rwaits,&icompleted,&rstatus);CHKERRQ(ierr);
185815a3b8e2SHong Zhang   }
185915a3b8e2SHong Zhang   ierr = PetscFree(rwaits);CHKERRQ(ierr);
1860445158ffSHong Zhang   if (nsend) {ierr = MPI_Waitall(nsend,swaits,sstatus);CHKERRQ(ierr);}
186115a3b8e2SHong Zhang 
18628cb82516SHong Zhang   ierr = PetscFree4(len_s,len_si,sstatus,owners_co);CHKERRQ(ierr);
186315a3b8e2SHong Zhang   ierr = PetscFree(len_ri);CHKERRQ(ierr);
186415a3b8e2SHong Zhang   ierr = PetscFree(swaits);CHKERRQ(ierr);
186515a3b8e2SHong Zhang   ierr = PetscFree(buf_s);CHKERRQ(ierr);
1866a4ffb656SHong Zhang 
186767a12041SHong Zhang   /* (5) compute the local portion of Cmpi      */
186867a12041SHong Zhang   /* ------------------------------------------ */
186978d30b94SHong Zhang   /* set initial free space to be Crmax, sufficient for holding nozeros in each row of Cmpi */
187078d30b94SHong Zhang   ierr          = PetscFreeSpaceGet(Crmax,&free_space);CHKERRQ(ierr);
187115a3b8e2SHong Zhang   current_space = free_space;
187215a3b8e2SHong Zhang 
1873445158ffSHong Zhang   ierr = PetscMalloc3(nrecv,&buf_ri_k,nrecv,&nextrow,nrecv,&nextci);CHKERRQ(ierr);
1874445158ffSHong Zhang   for (k=0; k<nrecv; k++) {
187515a3b8e2SHong Zhang     buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */
187615a3b8e2SHong Zhang     nrows       = *buf_ri_k[k];
187715a3b8e2SHong Zhang     nextrow[k]  = buf_ri_k[k] + 1;  /* next row number of k-th recved i-structure */
187815a3b8e2SHong Zhang     nextci[k]   = buf_ri_k[k] + (nrows + 1); /* poins to the next i-structure of k-th recved i-structure  */
187915a3b8e2SHong Zhang   }
1880445158ffSHong Zhang 
188178d30b94SHong Zhang   ierr = MatPreallocateInitialize(comm,pn,pn,dnz,onz);CHKERRQ(ierr);
188278d30b94SHong Zhang   ierr = PetscLLCondensedCreate(Crmax,pN,&lnk,&lnkbt);CHKERRQ(ierr);
188315a3b8e2SHong Zhang   for (i=0; i<pn; i++) {
188467a12041SHong Zhang     /* add C_loc into Cmpi */
188567a12041SHong Zhang     nzi  = c_loc->i[i+1] - c_loc->i[i];
188667a12041SHong Zhang     Jptr = c_loc->j + c_loc->i[i];
188767a12041SHong Zhang     ierr = PetscLLCondensedAddSorted(nzi,Jptr,lnk,lnkbt);CHKERRQ(ierr);
188815a3b8e2SHong Zhang 
188915a3b8e2SHong Zhang     /* add received col data into lnk */
1890445158ffSHong Zhang     for (k=0; k<nrecv; k++) { /* k-th received message */
189115a3b8e2SHong Zhang       if (i == *nextrow[k]) { /* i-th row */
189215a3b8e2SHong Zhang         nzi  = *(nextci[k]+1) - *nextci[k];
189315a3b8e2SHong Zhang         Jptr = buf_rj[k] + *nextci[k];
189415a3b8e2SHong Zhang         ierr = PetscLLCondensedAddSorted(nzi,Jptr,lnk,lnkbt);CHKERRQ(ierr);
189515a3b8e2SHong Zhang         nextrow[k]++; nextci[k]++;
189615a3b8e2SHong Zhang       }
189715a3b8e2SHong Zhang     }
1898d0e9a2f7SHong Zhang     nzi = lnk[0];
18998cb82516SHong Zhang 
190015a3b8e2SHong Zhang     /* copy data into free space, then initialize lnk */
1901d0e9a2f7SHong Zhang     ierr = PetscLLCondensedClean(pN,nzi,current_space->array,lnk,lnkbt);CHKERRQ(ierr);
1902d0e9a2f7SHong Zhang     ierr = MatPreallocateSet(i+owners[rank],nzi,current_space->array,dnz,onz);CHKERRQ(ierr);
190315a3b8e2SHong Zhang   }
190415a3b8e2SHong Zhang   ierr = PetscFree3(buf_ri_k,nextrow,nextci);CHKERRQ(ierr);
190515a3b8e2SHong Zhang   ierr = PetscLLDestroy(lnk,lnkbt);CHKERRQ(ierr);
1906445158ffSHong Zhang   ierr = PetscFreeSpaceDestroy(free_space);CHKERRQ(ierr);
190780bb4639SHong Zhang 
1908ae5f0867Sstefano_zampini   /* local sizes and preallocation */
190915a3b8e2SHong Zhang   ierr = MatSetSizes(Cmpi,pn,pn,PETSC_DETERMINE,PETSC_DETERMINE);CHKERRQ(ierr);
191015a3b8e2SHong Zhang   ierr = MatSetBlockSizes(Cmpi,PetscAbs(P->cmap->bs),PetscAbs(P->cmap->bs));CHKERRQ(ierr);
191115a3b8e2SHong Zhang   ierr = MatMPIAIJSetPreallocation(Cmpi,0,dnz,0,onz);CHKERRQ(ierr);
191215a3b8e2SHong Zhang   ierr = MatPreallocateFinalize(dnz,onz);CHKERRQ(ierr);
191315a3b8e2SHong Zhang 
1914445158ffSHong Zhang   /* members in merge */
1915445158ffSHong Zhang   ierr = PetscFree(id_r);CHKERRQ(ierr);
1916445158ffSHong Zhang   ierr = PetscFree(len_r);CHKERRQ(ierr);
1917445158ffSHong Zhang   ierr = PetscFree(buf_ri[0]);CHKERRQ(ierr);
1918445158ffSHong Zhang   ierr = PetscFree(buf_ri);CHKERRQ(ierr);
1919445158ffSHong Zhang   ierr = PetscFree(buf_rj[0]);CHKERRQ(ierr);
1920445158ffSHong Zhang   ierr = PetscFree(buf_rj);CHKERRQ(ierr);
1921445158ffSHong Zhang   ierr = PetscLayoutDestroy(&rowmap);CHKERRQ(ierr);
192215a3b8e2SHong Zhang 
192315a3b8e2SHong Zhang   /* attach the supporting struct to Cmpi for reuse */
192415a3b8e2SHong Zhang   c = (Mat_MPIAIJ*)Cmpi->data;
19253cdca5ebSHong Zhang   c->ap           = ptap;
19261a47ec54SHong Zhang   ptap->duplicate = Cmpi->ops->duplicate;
19271a47ec54SHong Zhang   ptap->destroy   = Cmpi->ops->destroy;
1928cf97cf3cSHong Zhang   ptap->view      = Cmpi->ops->view;
19298cb82516SHong Zhang   ierr = PetscCalloc1(pN,&ptap->apa);CHKERRQ(ierr);
19302259aa2eSHong Zhang 
19311a47ec54SHong Zhang   /* Cmpi is not ready for use - assembly will be done by MatPtAPNumeric() */
19321a47ec54SHong Zhang   Cmpi->assembled        = PETSC_FALSE;
19331a47ec54SHong Zhang   Cmpi->ops->destroy     = MatDestroy_MPIAIJ_PtAP;
1934cf97cf3cSHong Zhang   Cmpi->ops->view        = MatView_MPIAIJ_PtAP;
19354624976aSHong Zhang   Cmpi->ops->freeintermediatedatastructures = MatFreeIntermediateDataStructures_MPIAIJ_AP;
19361a47ec54SHong Zhang   *C                     = Cmpi;
19370d3441aeSHong Zhang   PetscFunctionReturn(0);
19380d3441aeSHong Zhang }
19390d3441aeSHong Zhang 
1940aa690a28SHong Zhang PetscErrorCode MatPtAPNumeric_MPIAIJ_MPIAIJ(Mat A,Mat P,Mat C)
19410d3441aeSHong Zhang {
19420d3441aeSHong Zhang   PetscErrorCode    ierr;
19430d3441aeSHong Zhang   Mat_MPIAIJ        *a=(Mat_MPIAIJ*)A->data,*p=(Mat_MPIAIJ*)P->data,*c=(Mat_MPIAIJ*)C->data;
19440d3441aeSHong Zhang   Mat_SeqAIJ        *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data;
19452dd9e643SHong Zhang   Mat_SeqAIJ        *ap,*p_loc,*p_oth=NULL,*c_seq;
19463cdca5ebSHong Zhang   Mat_APMPI         *ptap = c->ap;
19479ce11a7cSHong Zhang   Mat               AP_loc,C_loc,C_oth;
19480d3441aeSHong Zhang   PetscInt          i,rstart,rend,cm,ncols,row;
19498cb82516SHong Zhang   PetscInt          *api,*apj,am = A->rmap->n,j,col,apnz;
19508cb82516SHong Zhang   PetscScalar       *apa;
19510d3441aeSHong Zhang   const PetscInt    *cols;
19520d3441aeSHong Zhang   const PetscScalar *vals;
19530d3441aeSHong Zhang 
19540d3441aeSHong Zhang   PetscFunctionBegin;
1955dd4011a9SHong Zhang   if (!ptap) {
1956a9899c97SHong Zhang     MPI_Comm comm;
1957a9899c97SHong Zhang     ierr = PetscObjectGetComm((PetscObject)C,&comm);CHKERRQ(ierr);
1958dd4011a9SHong Zhang     SETERRQ(comm,PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be reused. Do not call MatFreeIntermediateDataStructures() or use '-mat_freeintermediatedatastructures'");
1959a9899c97SHong Zhang   }
1960a9899c97SHong Zhang 
19610d3441aeSHong Zhang   ierr = MatZeroEntries(C);CHKERRQ(ierr);
1962e497d3c8SHong Zhang   /* 1) get R = Pd^T,Ro = Po^T */
1963748c7196SHong Zhang   if (ptap->reuse == MAT_REUSE_MATRIX) {
1964419ecdd9Sandi selinger     ierr = MatTranspose(p->A,MAT_REUSE_MATRIX,&ptap->Rd);CHKERRQ(ierr);
1965419ecdd9Sandi selinger     ierr = MatTranspose(p->B,MAT_REUSE_MATRIX,&ptap->Ro);CHKERRQ(ierr);
1966748c7196SHong Zhang   }
19670d3441aeSHong Zhang 
1968e497d3c8SHong Zhang   /* 2) get AP_loc */
19690d3441aeSHong Zhang   AP_loc = ptap->AP_loc;
19708cb82516SHong Zhang   ap = (Mat_SeqAIJ*)AP_loc->data;
19710d3441aeSHong Zhang 
1972e497d3c8SHong Zhang   /* 2-1) get P_oth = ptap->P_oth  and P_loc = ptap->P_loc */
19730d3441aeSHong Zhang   /*-----------------------------------------------------*/
1974748c7196SHong Zhang   if (ptap->reuse == MAT_REUSE_MATRIX) {
1975748c7196SHong Zhang     /* P_oth and P_loc are obtained in MatPtASymbolic() when reuse == MAT_INITIAL_MATRIX */
19760d3441aeSHong Zhang     ierr = MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_REUSE_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth);CHKERRQ(ierr);
19770d3441aeSHong Zhang     ierr = MatMPIAIJGetLocalMat(P,MAT_REUSE_MATRIX,&ptap->P_loc);CHKERRQ(ierr);
1978e497d3c8SHong Zhang   }
19790d3441aeSHong Zhang 
19808cb82516SHong Zhang   /* 2-2) compute numeric A_loc*P - dominating part */
19818cb82516SHong Zhang   /* ---------------------------------------------- */
19820d3441aeSHong Zhang   /* get data from symbolic products */
19838cb82516SHong Zhang   p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data;
19842dd9e643SHong Zhang   if (ptap->P_oth) {
19858cb82516SHong Zhang     p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data;
19862dd9e643SHong Zhang   }
19878cb82516SHong Zhang   apa   = ptap->apa;
1988681d504bSHong Zhang   api   = ap->i;
1989681d504bSHong Zhang   apj   = ap->j;
1990e497d3c8SHong Zhang   for (i=0; i<am; i++) {
1991445158ffSHong Zhang     /* AP[i,:] = A[i,:]*P = Ad*P_loc Ao*P_oth */
1992e497d3c8SHong Zhang     AProw_nonscalable(i,ad,ao,p_loc,p_oth,apa);
1993e497d3c8SHong Zhang     apnz = api[i+1] - api[i];
1994e497d3c8SHong Zhang     for (j=0; j<apnz; j++) {
1995e497d3c8SHong Zhang       col = apj[j+api[i]];
1996e497d3c8SHong Zhang       ap->a[j+ap->i[i]] = apa[col];
1997e497d3c8SHong Zhang       apa[col] = 0.0;
19980d3441aeSHong Zhang     }
1999aa690a28SHong Zhang     ierr = PetscLogFlops(2.0*apnz);CHKERRQ(ierr);
20000d3441aeSHong Zhang   }
20010d3441aeSHong Zhang 
20028cb82516SHong Zhang   /* 3) C_loc = Rd*AP_loc, C_oth = Ro*AP_loc */
2003445158ffSHong Zhang   ierr = ((ptap->C_loc)->ops->matmultnumeric)(ptap->Rd,AP_loc,ptap->C_loc);CHKERRQ(ierr);
2004445158ffSHong Zhang   ierr = ((ptap->C_oth)->ops->matmultnumeric)(ptap->Ro,AP_loc,ptap->C_oth);CHKERRQ(ierr);
20059ce11a7cSHong Zhang   C_loc = ptap->C_loc;
20069ce11a7cSHong Zhang   C_oth = ptap->C_oth;
20070d3441aeSHong Zhang 
20080d3441aeSHong Zhang   /* add C_loc and Co to to C */
20090d3441aeSHong Zhang   ierr = MatGetOwnershipRange(C,&rstart,&rend);CHKERRQ(ierr);
20100d3441aeSHong Zhang 
20110d3441aeSHong Zhang   /* C_loc -> C */
20120d3441aeSHong Zhang   cm    = C_loc->rmap->N;
20139ce11a7cSHong Zhang   c_seq = (Mat_SeqAIJ*)C_loc->data;
20148cb82516SHong Zhang   cols = c_seq->j;
20158cb82516SHong Zhang   vals = c_seq->a;
2016904d1e70Sandi selinger 
2017904d1e70Sandi selinger 
2018e9ede7d0Sandi selinger   /* The (fast) MatSetValues_MPIAIJ_CopyFromCSRFormat function can only be used when C->was_assembled is PETSC_FALSE and */
2019904d1e70Sandi selinger   /* when there are no off-processor parts.  */
20201de21080Sandi selinger   /* If was_assembled is true, then the statement aj[rowstart_diag+dnz_row] = mat_j[col] - cstart; in MatSetValues_MPIAIJ_CopyFromCSRFormat */
20211de21080Sandi selinger   /* is no longer true. Then the more complex function MatSetValues_MPIAIJ() has to be used, where the column index is looked up from */
20221de21080Sandi selinger   /* a table, and other, more complex stuff has to be done. */
2023904d1e70Sandi selinger   if (C->assembled) {
2024904d1e70Sandi selinger     C->was_assembled = PETSC_TRUE;
2025904d1e70Sandi selinger     C->assembled     = PETSC_FALSE;
2026904d1e70Sandi selinger   }
2027904d1e70Sandi selinger   if (C->was_assembled) {
20280d3441aeSHong Zhang     for (i=0; i<cm; i++) {
20299ce11a7cSHong Zhang       ncols = c_seq->i[i+1] - c_seq->i[i];
20300d3441aeSHong Zhang       row = rstart + i;
2031904d1e70Sandi selinger       ierr = MatSetValues_MPIAIJ(C,1,&row,ncols,cols,vals,ADD_VALUES);CHKERRQ(ierr);
20328cb82516SHong Zhang       cols += ncols; vals += ncols;
20330d3441aeSHong Zhang     }
2034904d1e70Sandi selinger   } else {
2035e9ede7d0Sandi selinger     ierr = MatSetValues_MPIAIJ_CopyFromCSRFormat(C,c_seq->j,c_seq->i,c_seq->a);CHKERRQ(ierr);
2036904d1e70Sandi selinger   }
20370d3441aeSHong Zhang 
20380d3441aeSHong Zhang   /* Co -> C, off-processor part */
20399ce11a7cSHong Zhang   cm = C_oth->rmap->N;
20409ce11a7cSHong Zhang   c_seq = (Mat_SeqAIJ*)C_oth->data;
20418cb82516SHong Zhang   cols = c_seq->j;
20428cb82516SHong Zhang   vals = c_seq->a;
20439ce11a7cSHong Zhang   for (i=0; i<cm; i++) {
20449ce11a7cSHong Zhang     ncols = c_seq->i[i+1] - c_seq->i[i];
20450d3441aeSHong Zhang     row = p->garray[i];
20460d3441aeSHong Zhang     ierr = MatSetValues(C,1,&row,ncols,cols,vals,ADD_VALUES);CHKERRQ(ierr);
20478cb82516SHong Zhang     cols += ncols; vals += ncols;
20480d3441aeSHong Zhang   }
2049904d1e70Sandi selinger 
20500d3441aeSHong Zhang   ierr = MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
20510d3441aeSHong Zhang   ierr = MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
20520d3441aeSHong Zhang 
2053748c7196SHong Zhang   ptap->reuse = MAT_REUSE_MATRIX;
20543697aca6SHong Zhang 
2055dd4011a9SHong Zhang   /* supporting struct ptap consumes almost same amount of memory as C=PtAP, release it if C will not be updated by A and P */
20567d0a89b7SHong Zhang   if (ptap->freestruct) {
2057dd4011a9SHong Zhang     ierr = MatFreeIntermediateDataStructures(C);CHKERRQ(ierr);
20583697aca6SHong Zhang   }
20590d3441aeSHong Zhang   PetscFunctionReturn(0);
20600d3441aeSHong Zhang }
2061