xref: /petsc/src/mat/impls/aij/mpi/mpiptap.c (revision f33d1a9a8e1cc388a7592d5011fcd885988d41e4)
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 
742c57489SHong Zhang #include "src/mat/impls/aij/seq/aij.h"   /*I "petscmat.h" I*/
842c57489SHong Zhang #include "src/mat/utils/freespace.h"
942c57489SHong Zhang #include "src/mat/impls/aij/mpi/mpiaij.h"
1042c57489SHong Zhang #include "petscbt.h"
1142c57489SHong Zhang 
1242c57489SHong Zhang #undef __FUNCT__
1342c57489SHong Zhang #define __FUNCT__ "MatPtAP_MPIAIJ_MPIAIJ"
1442c57489SHong Zhang PetscErrorCode MatPtAP_MPIAIJ_MPIAIJ(Mat A,Mat P,MatReuse scall,PetscReal fill,Mat *C)
1542c57489SHong Zhang {
1642c57489SHong Zhang   PetscErrorCode       ierr;
17a6b2eed2SHong Zhang   Mat_MatMatMultMPI    *mult;
1842c57489SHong Zhang   PetscObjectContainer container;
1942c57489SHong Zhang 
2042c57489SHong Zhang   PetscFunctionBegin;
2142c57489SHong Zhang   if (scall == MAT_INITIAL_MATRIX){
2242c57489SHong Zhang     ierr = PetscLogEventBegin(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
23*f33d1a9aSHong Zhang     /* destroy the container 'Mat_MatMatMultMPI' in case that P is attached to it */
24*f33d1a9aSHong Zhang     ierr = PetscObjectContainerDestroy_Mat_MatMatMultMPI((PetscObject)P);CHKERRQ(ierr);
25*f33d1a9aSHong Zhang     /* create the container 'Mat_MatMatMultMPI' and attach it to P */
26a6b2eed2SHong Zhang     ierr = PetscNew(Mat_MatMatMultMPI,&mult);CHKERRQ(ierr);
27*f33d1a9aSHong Zhang     ierr = PetscMemzero(mult,sizeof(Mat_MatMatMultMPI));CHKERRQ(ierr);
28a6b2eed2SHong Zhang     mult->B_loc=PETSC_NULL; mult->B_oth=PETSC_NULL;
29a6b2eed2SHong Zhang     mult->abi=PETSC_NULL; mult->abj=PETSC_NULL;
30a6b2eed2SHong Zhang     mult->abnz_max = 0; /* symbolic A*P is not done yet */
3142c57489SHong Zhang 
3242c57489SHong Zhang     /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */
33a6b2eed2SHong Zhang     ierr = MatGetBrowsOfAoCols(A,P,scall,&mult->startsj,&mult->bufa,&mult->B_oth);CHKERRQ(ierr);
3442c57489SHong Zhang 
3542c57489SHong Zhang     /* get P_loc by taking all local rows of P */
36a6b2eed2SHong Zhang     ierr = MatGetLocalMat(P,scall,&mult->B_loc);CHKERRQ(ierr);
3742c57489SHong Zhang 
38*f33d1a9aSHong Zhang     /* attach the container 'Mat_MatMatMultMPI' to P */
39*f33d1a9aSHong Zhang     P->ops->destroy  = MatDestroy_MPIAIJ_MatMatMult;
4042c57489SHong Zhang     ierr = PetscObjectContainerCreate(PETSC_COMM_SELF,&container);CHKERRQ(ierr);
41a6b2eed2SHong Zhang     ierr = PetscObjectContainerSetPointer(container,mult);CHKERRQ(ierr);
42*f33d1a9aSHong Zhang     ierr = PetscObjectCompose((PetscObject)P,"Mat_MatMatMultMPI",(PetscObject)container);CHKERRQ(ierr);
4342c57489SHong Zhang 
4442c57489SHong Zhang     /* now, compute symbolic local P^T*A*P */
4542c57489SHong Zhang     ierr = MatPtAPSymbolic_MPIAIJ_MPIAIJ(A,P,fill,C);CHKERRQ(ierr);/* numeric product is computed as well */
4642c57489SHong Zhang     ierr = PetscLogEventEnd(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr);
4742c57489SHong Zhang   } else if (scall == MAT_REUSE_MATRIX){
48*f33d1a9aSHong Zhang     ierr = PetscObjectQuery((PetscObject)P,"Mat_MatMatMultMPI",(PetscObject *)&container);CHKERRQ(ierr);
4942c57489SHong Zhang     if (container) {
50a6b2eed2SHong Zhang       ierr  = PetscObjectContainerGetPointer(container,(void **)&mult);CHKERRQ(ierr);
5142c57489SHong Zhang     } else {
5242c57489SHong Zhang       SETERRQ(PETSC_ERR_ARG_WRONGSTATE, "Matrix P does not posses an object container");
5342c57489SHong Zhang     }
5442c57489SHong Zhang     /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */
55a6b2eed2SHong Zhang     ierr = MatGetBrowsOfAoCols(A,P,scall,&mult->startsj,&mult->bufa,&mult->B_oth);CHKERRQ(ierr);
5642c57489SHong Zhang     /* get P_loc by taking all local rows of P */
57a6b2eed2SHong Zhang     ierr = MatGetLocalMat(P,scall,&mult->B_loc);CHKERRQ(ierr);
5842c57489SHong Zhang   } else {
5942c57489SHong Zhang     SETERRQ1(PETSC_ERR_ARG_WRONG,"Invalid MatReuse %d",scall);
6042c57489SHong Zhang   }
6142c57489SHong Zhang   /* now, compute numeric local P^T*A*P */
6242c57489SHong Zhang   ierr = PetscLogEventBegin(MAT_PtAPNumeric,A,P,0,0);CHKERRQ(ierr);
6342c57489SHong Zhang   ierr = MatPtAPNumeric_MPIAIJ_MPIAIJ(A,P,*C);CHKERRQ(ierr);
6442c57489SHong Zhang   ierr = PetscLogEventEnd(MAT_PtAPNumeric,A,P,0,0);CHKERRQ(ierr);
6542c57489SHong Zhang 
6642c57489SHong Zhang   PetscFunctionReturn(0);
6742c57489SHong Zhang }
6842c57489SHong Zhang 
6942c57489SHong Zhang #undef __FUNCT__
7042c57489SHong Zhang #define __FUNCT__ "MatPtAPSymbolic_MPIAIJ_MPIAIJ"
7142c57489SHong Zhang PetscErrorCode MatPtAPSymbolic_MPIAIJ_MPIAIJ(Mat A,Mat P,PetscReal fill,Mat *C)
7242c57489SHong Zhang {
7342c57489SHong Zhang   PetscErrorCode       ierr;
74de0260b3SHong Zhang   Mat                  P_loc,P_oth,B_mpi;
75de0260b3SHong Zhang   Mat_MatMatMultMPI    *ap;
7642c57489SHong Zhang   PetscObjectContainer container;
7742c57489SHong Zhang   FreeSpaceList        free_space=PETSC_NULL,current_space=PETSC_NULL;
7883445d95SHong Zhang   Mat_MPIAIJ           *a=(Mat_MPIAIJ*)A->data,*p=(Mat_MPIAIJ*)P->data;
7942c57489SHong Zhang   Mat_SeqAIJ           *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data;
8042c57489SHong Zhang   Mat_SeqAIJ           *p_loc,*p_oth;
81ded4f561SHong Zhang   PetscInt             *pi_loc,*pj_loc,*pi_oth,*pj_oth,*pdti,*pdtj,*poti,*potj,*ptJ;
8283445d95SHong Zhang   PetscInt             *adi=ad->i,*adj=ad->j,*aoi=ao->i,*aoj=ao->j,nnz;
8313f74950SBarry Smith   PetscInt             nlnk,*lnk,*owners_co,*coi,*coj,i,k,pnz,row;
84a6b2eed2SHong Zhang   PetscInt             am=A->m,pN=P->N,pn=P->n;
8542c57489SHong Zhang   PetscBT              lnkbt;
8642c57489SHong Zhang   MPI_Comm             comm=A->comm;
87ded4f561SHong Zhang   PetscMPIInt          size,rank,tag,*len_si,*len_s,*len_ri;
8842c57489SHong Zhang   PetscInt             **buf_rj,**buf_ri,**buf_ri_k;
8942c57489SHong Zhang   PetscInt             len,proc,*dnz,*onz,*owners;
90ded4f561SHong Zhang   PetscInt             nzi,*bi,*bj;
9142c57489SHong Zhang   PetscInt             nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci;
92ded4f561SHong Zhang   MPI_Request          *swaits,*rwaits;
93ded4f561SHong Zhang   MPI_Status           *sstatus,rstatus;
9442c57489SHong Zhang   Mat_Merge_SeqsToMPI  *merge;
95ded4f561SHong Zhang   PetscInt             *api,*apj,*Jptr,apnz,*prmap=p->garray,pon;
9613f74950SBarry Smith   PetscMPIInt          j;
9742c57489SHong Zhang 
9842c57489SHong Zhang   PetscFunctionBegin;
99*f33d1a9aSHong Zhang   ierr = PetscObjectQuery((PetscObject)P,"Mat_MatMatMultMPI",(PetscObject *)&container);CHKERRQ(ierr);
10042c57489SHong Zhang   if (container) {
101de0260b3SHong Zhang     ierr  = PetscObjectContainerGetPointer(container,(void **)&ap);CHKERRQ(ierr);
10242c57489SHong Zhang   } else {
10342c57489SHong Zhang     SETERRQ(PETSC_ERR_ARG_WRONGSTATE, "Matrix P does not posses an object container");
10442c57489SHong Zhang   }
10583445d95SHong Zhang   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
10683445d95SHong Zhang   ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
10783445d95SHong Zhang 
10842c57489SHong Zhang   /* get data from P_loc and P_oth */
109a6b2eed2SHong Zhang   P_loc=ap->B_loc; P_oth=ap->B_oth;
11042c57489SHong Zhang   p_loc=(Mat_SeqAIJ*)P_loc->data; p_oth=(Mat_SeqAIJ*)P_oth->data;
11142c57489SHong Zhang   pi_loc=p_loc->i; pj_loc=p_loc->j; pi_oth=p_oth->i; pj_oth=p_oth->j;
11242c57489SHong Zhang 
11383445d95SHong Zhang   /* first, compute symbolic AP = A_loc*P */
114de0260b3SHong Zhang   /*--------------------------------------*/
11582412ba7SHong Zhang   ierr  = PetscMalloc((am+2)*sizeof(PetscInt),&api);CHKERRQ(ierr);
11682412ba7SHong Zhang   ap->abi = api;
11782412ba7SHong Zhang   api[0] = 0;
11883445d95SHong Zhang 
11983445d95SHong Zhang   /* create and initialize a linked list */
12083445d95SHong Zhang   nlnk = pN+1;
12183445d95SHong Zhang   ierr = PetscLLCreate(pN,pN,nlnk,lnk,lnkbt);CHKERRQ(ierr);
12283445d95SHong Zhang 
12382412ba7SHong Zhang   /* Initial FreeSpace size is fill*nnz(A) */
124de0260b3SHong Zhang   ierr = GetMoreSpace((PetscInt)(fill*(adi[am]+aoi[am])),&free_space);CHKERRQ(ierr);
125f4a743e1SHong Zhang   current_space = free_space;
126f4a743e1SHong Zhang 
127f4a743e1SHong Zhang   for (i=0;i<am;i++) {
12882412ba7SHong Zhang     apnz = 0;
129f4a743e1SHong Zhang     /* diagonal portion of A */
130ded4f561SHong Zhang     nzi = adi[i+1] - adi[i];
131ded4f561SHong Zhang     for (j=0; j<nzi; j++){
13282412ba7SHong Zhang       row = *adj++;
13382412ba7SHong Zhang       pnz = pi_loc[row+1] - pi_loc[row];
134ded4f561SHong Zhang       Jptr  = pj_loc + pi_loc[row];
13583445d95SHong Zhang       /* add non-zero cols of P into the sorted linked list lnk */
136ded4f561SHong Zhang       ierr = PetscLLAdd(pnz,Jptr,pN,nlnk,lnk,lnkbt);CHKERRQ(ierr);
13782412ba7SHong Zhang       apnz += nlnk;
138f4a743e1SHong Zhang     }
139f4a743e1SHong Zhang     /* off-diagonal portion of A */
140ded4f561SHong Zhang     nzi = aoi[i+1] - aoi[i];
141ded4f561SHong Zhang     for (j=0; j<nzi; j++){
14282412ba7SHong Zhang       row = *aoj++;
14382412ba7SHong Zhang       pnz = pi_oth[row+1] - pi_oth[row];
144ded4f561SHong Zhang       Jptr  = pj_oth + pi_oth[row];
145ded4f561SHong Zhang       ierr = PetscLLAdd(pnz,Jptr,pN,nlnk,lnk,lnkbt);CHKERRQ(ierr);
14682412ba7SHong Zhang       apnz += nlnk;
147f4a743e1SHong Zhang     }
148f4a743e1SHong Zhang 
14982412ba7SHong Zhang     api[i+1] = api[i] + apnz;
15082412ba7SHong Zhang     if (ap->abnz_max < apnz) ap->abnz_max = apnz;
151f4a743e1SHong Zhang 
15283445d95SHong Zhang     /* if free space is not available, double the total space in the list */
15382412ba7SHong Zhang     if (current_space->local_remaining<apnz) {
154f4a743e1SHong Zhang       ierr = GetMoreSpace(current_space->total_array_size,&current_space);CHKERRQ(ierr);
155f4a743e1SHong Zhang     }
156f4a743e1SHong Zhang 
157f4a743e1SHong Zhang     /* Copy data into free space, then initialize lnk */
15882412ba7SHong Zhang     ierr = PetscLLClean(pN,pN,apnz,lnk,current_space->array,lnkbt);CHKERRQ(ierr);
15982412ba7SHong Zhang     current_space->array           += apnz;
16082412ba7SHong Zhang     current_space->local_used      += apnz;
16182412ba7SHong Zhang     current_space->local_remaining -= apnz;
162f4a743e1SHong Zhang   }
16382412ba7SHong Zhang   /* Allocate space for apj, initialize apj, and */
164f4a743e1SHong Zhang   /* destroy list of free space and other temporary array(s) */
16582412ba7SHong Zhang   ierr = PetscMalloc((api[am]+1)*sizeof(PetscInt),&ap->abj);CHKERRQ(ierr);
16682412ba7SHong Zhang   apj = ap->abj;
167de0260b3SHong Zhang   ierr = MakeSpaceContiguous(&free_space,ap->abj);CHKERRQ(ierr);
168f4a743e1SHong Zhang 
169ded4f561SHong Zhang   /* determine symbolic Co=(p->B)^T*AP - send to others */
170ded4f561SHong Zhang   /*----------------------------------------------------*/
171de0260b3SHong Zhang   ierr = MatGetSymbolicTranspose_SeqAIJ(p->B,&poti,&potj);CHKERRQ(ierr);
17242c57489SHong Zhang 
173ded4f561SHong Zhang   /* then, compute symbolic Co = (p->B)^T*AP */
174de0260b3SHong Zhang   pon = (p->B)->n; /* total num of rows to be sent to other processors
17583445d95SHong Zhang                          >= (num of nonzero rows of C_seq) - pn */
176de0260b3SHong Zhang   ierr = PetscMalloc((pon+1)*sizeof(PetscInt),&coi);CHKERRQ(ierr);
177de0260b3SHong Zhang   coi[0] = 0;
17842c57489SHong Zhang 
17982412ba7SHong Zhang   /* set initial free space to be 3*pon*max( nnz(AP) per row) */
18082412ba7SHong Zhang   nnz           = 3*pon*ap->abnz_max + 1;
181de0260b3SHong Zhang   ierr          = GetMoreSpace(nnz,&free_space);
18242c57489SHong Zhang   current_space = free_space;
18342c57489SHong Zhang 
184de0260b3SHong Zhang   for (i=0; i<pon; i++) {
185ded4f561SHong Zhang     nnz  = 0;
18682412ba7SHong Zhang     pnz = poti[i+1] - poti[i];
18782412ba7SHong Zhang     j     = pnz;
188de0260b3SHong Zhang     ptJ   = potj + poti[i+1];
18983445d95SHong Zhang     while (j){/* assume cols are almost in increasing order, starting from its end saves computation */
19083445d95SHong Zhang       j--; ptJ--;
19182412ba7SHong Zhang       row  = *ptJ; /* row of AP == col of Pot */
19282412ba7SHong Zhang       apnz = api[row+1] - api[row];
193ded4f561SHong Zhang       Jptr   = apj + api[row];
19483445d95SHong Zhang       /* add non-zero cols of AP into the sorted linked list lnk */
195ded4f561SHong Zhang       ierr = PetscLLAdd(apnz,Jptr,pN,nlnk,lnk,lnkbt);CHKERRQ(ierr);
196ded4f561SHong Zhang       nnz += nlnk;
19742c57489SHong Zhang     }
19842c57489SHong Zhang 
19983445d95SHong Zhang     /* If free space is not available, double the total space in the list */
200ded4f561SHong Zhang     if (current_space->local_remaining<nnz) {
20142c57489SHong Zhang       ierr = GetMoreSpace(current_space->total_array_size,&current_space);CHKERRQ(ierr);
20242c57489SHong Zhang     }
20342c57489SHong Zhang 
20442c57489SHong Zhang     /* Copy data into free space, and zero out denserows */
205ded4f561SHong Zhang     ierr = PetscLLClean(pN,pN,nnz,lnk,current_space->array,lnkbt);CHKERRQ(ierr);
206ded4f561SHong Zhang     current_space->array           += nnz;
207ded4f561SHong Zhang     current_space->local_used      += nnz;
208ded4f561SHong Zhang     current_space->local_remaining -= nnz;
209ded4f561SHong Zhang     coi[i+1] = coi[i] + nnz;
21042c57489SHong Zhang   }
211de0260b3SHong Zhang   ierr = PetscMalloc((coi[pon]+1)*sizeof(PetscInt),&coj);CHKERRQ(ierr);
212de0260b3SHong Zhang   ierr = MakeSpaceContiguous(&free_space,coj);CHKERRQ(ierr);
213de0260b3SHong Zhang   ierr = MatRestoreSymbolicTranspose_SeqAIJ(p->B,&poti,&potj);CHKERRQ(ierr);
21442c57489SHong Zhang 
215ded4f561SHong Zhang   /* send j-array (coj) of Co to other processors */
216ded4f561SHong Zhang   /*----------------------------------------------*/
21742c57489SHong Zhang   /* determine row ownership */
21883445d95SHong Zhang   ierr = PetscNew(Mat_Merge_SeqsToMPI,&merge);CHKERRQ(ierr);
219*f33d1a9aSHong Zhang   ierr = PetscMemzero(merge,sizeof(Mat_Merge_SeqsToMPI));CHKERRQ(ierr);
22042c57489SHong Zhang   ierr = PetscMapCreate(comm,&merge->rowmap);CHKERRQ(ierr);
22142c57489SHong Zhang   ierr = PetscMapSetLocalSize(merge->rowmap,pn);CHKERRQ(ierr);
22242c57489SHong Zhang   ierr = PetscMapSetType(merge->rowmap,MAP_MPI);CHKERRQ(ierr);
22342c57489SHong Zhang   ierr = PetscMapGetGlobalRange(merge->rowmap,&owners);CHKERRQ(ierr);
22442c57489SHong Zhang 
22542c57489SHong Zhang   /* determine the number of messages to send, their lengths */
22642c57489SHong Zhang   ierr = PetscMalloc(size*sizeof(PetscMPIInt),&len_si);CHKERRQ(ierr);
22783445d95SHong Zhang   ierr = PetscMemzero(len_si,size*sizeof(PetscMPIInt));CHKERRQ(ierr);
22842c57489SHong Zhang   ierr = PetscMalloc(size*sizeof(PetscMPIInt),&merge->len_s);CHKERRQ(ierr);
22942c57489SHong Zhang   len_s = merge->len_s;
23042c57489SHong Zhang   merge->nsend = 0;
23183445d95SHong Zhang 
232de0260b3SHong Zhang   ierr = PetscMalloc((size+2)*sizeof(PetscInt),&owners_co);CHKERRQ(ierr);
233de0260b3SHong Zhang   ierr = PetscMemzero(len_s,size*sizeof(PetscMPIInt));CHKERRQ(ierr);
234de0260b3SHong Zhang 
23583445d95SHong Zhang   proc = 0;
236de0260b3SHong Zhang   for (i=0; i<pon; i++){
237de0260b3SHong Zhang     while (prmap[i] >= owners[proc+1]) proc++;
238de0260b3SHong Zhang     len_si[proc]++;  /* num of rows in Co to be sent to [proc] */
239de0260b3SHong Zhang     len_s[proc] += coi[i+1] - coi[i];
24083445d95SHong Zhang   }
241de0260b3SHong Zhang 
242de0260b3SHong Zhang   len   = 0;  /* max length of buf_si[] */
243de0260b3SHong Zhang   owners_co[0] = 0;
24442c57489SHong Zhang   for (proc=0; proc<size; proc++){
245de0260b3SHong Zhang     owners_co[proc+1] = owners_co[proc] + len_si[proc];
24683445d95SHong Zhang     if (len_si[proc]){
24742c57489SHong Zhang       merge->nsend++;
24883445d95SHong Zhang       len_si[proc] = 2*(len_si[proc] + 1);
24942c57489SHong Zhang       len += len_si[proc];
25042c57489SHong Zhang     }
25142c57489SHong Zhang   }
25242c57489SHong Zhang 
253ded4f561SHong Zhang   /* determine the number and length of messages to receive for coi and coj  */
25442c57489SHong Zhang   ierr = PetscGatherNumberOfMessages(comm,PETSC_NULL,len_s,&merge->nrecv);CHKERRQ(ierr);
25542c57489SHong Zhang   ierr = PetscGatherMessageLengths2(comm,merge->nsend,merge->nrecv,len_s,len_si,&merge->id_r,&merge->len_r,&len_ri);CHKERRQ(ierr);
25642c57489SHong Zhang 
257ded4f561SHong Zhang   /* post the Irecv and Isend of coj */
258ded4f561SHong Zhang   ierr = PetscObjectGetNewTag((PetscObject)merge->rowmap,&tag);CHKERRQ(ierr);
259ded4f561SHong Zhang   ierr = PetscPostIrecvInt(comm,tag,merge->nrecv,merge->id_r,merge->len_r,&buf_rj,&rwaits);CHKERRQ(ierr);
26042c57489SHong Zhang 
261ded4f561SHong Zhang   ierr = PetscMalloc((merge->nsend+1)*sizeof(MPI_Request),&swaits);CHKERRQ(ierr);
26242c57489SHong Zhang 
26342c57489SHong Zhang   for (proc=0, k=0; proc<size; proc++){
26442c57489SHong Zhang     if (!len_s[proc]) continue;
265de0260b3SHong Zhang     i = owners_co[proc];
266ded4f561SHong Zhang     ierr = MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tag,comm,swaits+k);CHKERRQ(ierr);
26742c57489SHong Zhang     k++;
26842c57489SHong Zhang   }
26942c57489SHong Zhang 
270ded4f561SHong Zhang   /* receives and sends of coj are complete */
271ded4f561SHong Zhang   ierr = PetscMalloc(size*sizeof(MPI_Status),&sstatus);CHKERRQ(ierr);
272ded4f561SHong Zhang   i = merge->nrecv;
273ded4f561SHong Zhang   while (i--) {
274ded4f561SHong Zhang     ierr = MPI_Waitany(merge->nrecv,rwaits,&j,&rstatus);CHKERRQ(ierr);
275ded4f561SHong Zhang   }
276ded4f561SHong Zhang   ierr = PetscFree(rwaits);CHKERRQ(ierr);
277ded4f561SHong Zhang   if (merge->nsend){
278ded4f561SHong Zhang     ierr = MPI_Waitall(merge->nsend,swaits,sstatus);CHKERRQ(ierr);
27982412ba7SHong Zhang   }
28082412ba7SHong Zhang 
281ded4f561SHong Zhang   /* send and recv coi */
282ded4f561SHong Zhang   /*-------------------*/
283ded4f561SHong Zhang   ierr = PetscPostIrecvInt(comm,tag,merge->nrecv,merge->id_r,len_ri,&buf_ri,&rwaits);CHKERRQ(ierr);
28442c57489SHong Zhang 
28542c57489SHong Zhang   ierr = PetscMalloc((len+1)*sizeof(PetscInt),&buf_s);CHKERRQ(ierr);
28642c57489SHong Zhang   buf_si = buf_s;  /* points to the beginning of k-th msg to be sent */
28742c57489SHong Zhang   for (proc=0,k=0; proc<size; proc++){
28842c57489SHong Zhang     if (!len_s[proc]) continue;
28942c57489SHong Zhang     /* form outgoing message for i-structure:
29042c57489SHong Zhang          buf_si[0]:                 nrows to be sent
29142c57489SHong Zhang                [1:nrows]:           row index (global)
29242c57489SHong Zhang                [nrows+1:2*nrows+1]: i-structure index
29342c57489SHong Zhang     */
29442c57489SHong Zhang     /*-------------------------------------------*/
29542c57489SHong Zhang     nrows = len_si[proc]/2 - 1;
29642c57489SHong Zhang     buf_si_i    = buf_si + nrows+1;
29742c57489SHong Zhang     buf_si[0]   = nrows;
29842c57489SHong Zhang     buf_si_i[0] = 0;
29942c57489SHong Zhang     nrows = 0;
300de0260b3SHong Zhang     for (i=owners_co[proc]; i<owners_co[proc+1]; i++){
301ded4f561SHong Zhang       nzi = coi[i+1] - coi[i];
302ded4f561SHong Zhang       buf_si_i[nrows+1] = buf_si_i[nrows] + nzi; /* i-structure */
303de0260b3SHong Zhang       buf_si[nrows+1] =prmap[i] -owners[proc]; /* local row index */
30442c57489SHong Zhang       nrows++;
30542c57489SHong Zhang     }
306ded4f561SHong Zhang     ierr = MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tag,comm,swaits+k);CHKERRQ(ierr);
30742c57489SHong Zhang     k++;
30842c57489SHong Zhang     buf_si += len_si[proc];
30942c57489SHong Zhang   }
310ded4f561SHong Zhang   i = merge->nrecv;
311ded4f561SHong Zhang   while (i--) {
312ded4f561SHong Zhang     ierr = MPI_Waitany(merge->nrecv,rwaits,&j,&rstatus);CHKERRQ(ierr);
313ded4f561SHong Zhang   }
314ded4f561SHong Zhang   ierr = PetscFree(rwaits);CHKERRQ(ierr);
315ded4f561SHong Zhang   if (merge->nsend){
316ded4f561SHong Zhang     ierr = MPI_Waitall(merge->nsend,swaits,sstatus);CHKERRQ(ierr);
317ded4f561SHong Zhang   }
31842c57489SHong Zhang 
31942c57489SHong Zhang   ierr = PetscLogInfo((PetscObject)A,"MatMerge_SeqsToMPI: nsend: %d, nrecv: %d\n",merge->nsend,merge->nrecv);CHKERRQ(ierr);
32042c57489SHong Zhang   for (i=0; i<merge->nrecv; i++){
32142c57489SHong Zhang     ierr = PetscLogInfo((PetscObject)A,"MatMerge_SeqsToMPI:   recv len_ri=%d, len_rj=%d from [%d]\n",len_ri[i],merge->len_r[i],merge->id_r[i]);CHKERRQ(ierr);
32242c57489SHong Zhang   }
32342c57489SHong Zhang 
32442c57489SHong Zhang   ierr = PetscFree(len_si);CHKERRQ(ierr);
32542c57489SHong Zhang   ierr = PetscFree(len_ri);CHKERRQ(ierr);
326ded4f561SHong Zhang   ierr = PetscFree(swaits);CHKERRQ(ierr);
327ded4f561SHong Zhang   ierr = PetscFree(sstatus);CHKERRQ(ierr);
32842c57489SHong Zhang   ierr = PetscFree(buf_s);CHKERRQ(ierr);
32942c57489SHong Zhang 
330de0260b3SHong Zhang   /* compute the local portion of C (mpi mat) */
331de0260b3SHong Zhang   /*------------------------------------------*/
332ded4f561SHong Zhang   ierr = MatGetSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj);CHKERRQ(ierr);
333ded4f561SHong Zhang 
33442c57489SHong Zhang   /* allocate bi array and free space for accumulating nonzero column info */
33542c57489SHong Zhang   ierr = PetscMalloc((pn+1)*sizeof(PetscInt),&bi);CHKERRQ(ierr);
33642c57489SHong Zhang   bi[0] = 0;
33742c57489SHong Zhang 
338ded4f561SHong Zhang   /* set initial free space to be 3*pn*max( nnz(AP) per row) */
339ded4f561SHong Zhang   nnz           = 3*pn*ap->abnz_max + 1;
340ded4f561SHong Zhang   ierr          = GetMoreSpace(nnz,&free_space);
34142c57489SHong Zhang   current_space = free_space;
34242c57489SHong Zhang 
34342c57489SHong Zhang   ierr = PetscMalloc((3*merge->nrecv+1)*sizeof(PetscInt**),&buf_ri_k);CHKERRQ(ierr);
34442c57489SHong Zhang   nextrow = buf_ri_k + merge->nrecv;
34542c57489SHong Zhang   nextci  = nextrow + merge->nrecv;
34642c57489SHong Zhang   for (k=0; k<merge->nrecv; k++){
34742c57489SHong Zhang     buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */
34842c57489SHong Zhang     nrows       = *buf_ri_k[k];
34942c57489SHong Zhang     nextrow[k]  = buf_ri_k[k] + 1;  /* next row number of k-th recved i-structure */
35042c57489SHong Zhang     nextci[k]   = buf_ri_k[k] + (nrows + 1);/* poins to the next i-structure of k-th recved i-structure  */
35142c57489SHong Zhang   }
35242c57489SHong Zhang   ierr = MatPreallocateInitialize(comm,pn,pn,dnz,onz);CHKERRQ(ierr);
35342c57489SHong Zhang   for (i=0; i<pn; i++) {
354ded4f561SHong Zhang     /* add pdt[i,:]*AP into lnk */
355ded4f561SHong Zhang     nnz = 0;
356ded4f561SHong Zhang     pnz  = pdti[i+1] - pdti[i];
357ded4f561SHong Zhang     j    = pnz;
358ded4f561SHong Zhang     ptJ  = pdtj + pdti[i+1];
359ded4f561SHong Zhang     while (j){/* assume cols are almost in increasing order, starting from its end saves computation */
360ded4f561SHong Zhang       j--; ptJ--;
361ded4f561SHong Zhang       row  = *ptJ; /* row of AP == col of Pt */
362ded4f561SHong Zhang       apnz = api[row+1] - api[row];
363ded4f561SHong Zhang       Jptr   = apj + api[row];
364ded4f561SHong Zhang       /* add non-zero cols of AP into the sorted linked list lnk */
365ded4f561SHong Zhang       ierr = PetscLLAdd(apnz,Jptr,pN,nlnk,lnk,lnkbt);CHKERRQ(ierr);
366ded4f561SHong Zhang       nnz += nlnk;
367ded4f561SHong Zhang     }
36842c57489SHong Zhang     /* add received col data into lnk */
36942c57489SHong Zhang     for (k=0; k<merge->nrecv; k++){ /* k-th received message */
37042c57489SHong Zhang       if (i == *nextrow[k]) { /* i-th row */
371ded4f561SHong Zhang         nzi = *(nextci[k]+1) - *nextci[k];
372ded4f561SHong Zhang         Jptr  = buf_rj[k] + *nextci[k];
373ded4f561SHong Zhang         ierr = PetscLLAdd(nzi,Jptr,pN,nlnk,lnk,lnkbt);CHKERRQ(ierr);
374ded4f561SHong Zhang         nnz += nlnk;
37542c57489SHong Zhang         nextrow[k]++; nextci[k]++;
37642c57489SHong Zhang       }
37742c57489SHong Zhang     }
37842c57489SHong Zhang 
37942c57489SHong Zhang     /* if free space is not available, make more free space */
380ded4f561SHong Zhang     if (current_space->local_remaining<nnz) {
38142c57489SHong Zhang       ierr = GetMoreSpace(current_space->total_array_size,&current_space);CHKERRQ(ierr);
38242c57489SHong Zhang     }
38342c57489SHong Zhang     /* copy data into free space, then initialize lnk */
384ded4f561SHong Zhang     ierr = PetscLLClean(pN,pN,nnz,lnk,current_space->array,lnkbt);CHKERRQ(ierr);
385ded4f561SHong Zhang     ierr = MatPreallocateSet(i+owners[rank],nnz,current_space->array,dnz,onz);CHKERRQ(ierr);
386ded4f561SHong Zhang     current_space->array           += nnz;
387ded4f561SHong Zhang     current_space->local_used      += nnz;
388ded4f561SHong Zhang     current_space->local_remaining -= nnz;
389ded4f561SHong Zhang     bi[i+1] = bi[i] + nnz;
39042c57489SHong Zhang   }
391ded4f561SHong Zhang   ierr = MatRestoreSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj);CHKERRQ(ierr);
39242c57489SHong Zhang   ierr = PetscFree(buf_ri_k);CHKERRQ(ierr);
39342c57489SHong Zhang 
39442c57489SHong Zhang   ierr = PetscMalloc((bi[pn]+1)*sizeof(PetscInt),&bj);CHKERRQ(ierr);
39542c57489SHong Zhang   ierr = MakeSpaceContiguous(&free_space,bj);CHKERRQ(ierr);
39642c57489SHong Zhang   ierr = PetscLLDestroy(lnk,lnkbt);CHKERRQ(ierr);
39742c57489SHong Zhang 
39842c57489SHong Zhang   /* create symbolic parallel matrix B_mpi */
39942c57489SHong Zhang   /*---------------------------------------*/
40042c57489SHong Zhang   ierr = MatCreate(comm,pn,pn,PETSC_DETERMINE,PETSC_DETERMINE,&B_mpi);CHKERRQ(ierr);
40142c57489SHong Zhang   ierr = MatSetType(B_mpi,MATMPIAIJ);CHKERRQ(ierr);
40242c57489SHong Zhang   ierr = MatMPIAIJSetPreallocation(B_mpi,0,dnz,0,onz);CHKERRQ(ierr);
40342c57489SHong Zhang   ierr = MatPreallocateFinalize(dnz,onz);CHKERRQ(ierr);
40442c57489SHong Zhang 
40582412ba7SHong Zhang   /* B_mpi is not ready for use - assembly will be done by MatPtAPNumeric() */
40642c57489SHong Zhang   B_mpi->assembled     = PETSC_FALSE;
40742c57489SHong Zhang   B_mpi->ops->destroy  = MatDestroy_MPIAIJ_SeqsToMPI;
40842c57489SHong Zhang   merge->bi            = bi;
40942c57489SHong Zhang   merge->bj            = bj;
410de0260b3SHong Zhang   merge->coi           = coi;
411de0260b3SHong Zhang   merge->coj           = coj;
41242c57489SHong Zhang   merge->buf_ri        = buf_ri;
41342c57489SHong Zhang   merge->buf_rj        = buf_rj;
414de0260b3SHong Zhang   merge->owners_co     = owners_co;
41542c57489SHong Zhang 
41642c57489SHong Zhang   /* attach the supporting struct to B_mpi for reuse */
41742c57489SHong Zhang   ierr = PetscObjectContainerCreate(PETSC_COMM_SELF,&container);CHKERRQ(ierr);
41842c57489SHong Zhang   ierr = PetscObjectContainerSetPointer(container,merge);CHKERRQ(ierr);
41942c57489SHong Zhang   ierr = PetscObjectCompose((PetscObject)B_mpi,"MatMergeSeqsToMPI",(PetscObject)container);CHKERRQ(ierr);
42042c57489SHong Zhang   *C = B_mpi;
42142c57489SHong Zhang 
42242c57489SHong Zhang   PetscFunctionReturn(0);
42342c57489SHong Zhang }
42442c57489SHong Zhang 
42542c57489SHong Zhang #undef __FUNCT__
42642c57489SHong Zhang #define __FUNCT__ "MatPtAPNumeric_MPIAIJ_MPIAIJ"
42742c57489SHong Zhang PetscErrorCode MatPtAPNumeric_MPIAIJ_MPIAIJ(Mat A,Mat P,Mat C)
42842c57489SHong Zhang {
42942c57489SHong Zhang   PetscErrorCode       ierr;
43042c57489SHong Zhang   Mat_Merge_SeqsToMPI  *merge;
431de0260b3SHong Zhang   Mat_MatMatMultMPI    *ap;
43242c57489SHong Zhang   PetscObjectContainer cont_merge,cont_ptap;
433de0260b3SHong Zhang   Mat_MPIAIJ     *a=(Mat_MPIAIJ*)A->data,*p=(Mat_MPIAIJ*)P->data;
43442c57489SHong Zhang   Mat_SeqAIJ     *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data;
435de0260b3SHong Zhang   Mat_SeqAIJ     *pd=(Mat_SeqAIJ*)(p->A)->data,*po=(Mat_SeqAIJ*)(p->B)->data;
43642c57489SHong Zhang   Mat_SeqAIJ     *p_loc,*p_oth;
43782412ba7SHong Zhang   PetscInt       *adi=ad->i,*aoi=ao->i,*adj=ad->j,*aoj=ao->j,*apJ,nextp,flops=0;
43882412ba7SHong Zhang   PetscInt       *pi_loc,*pj_loc,*pi_oth,*pj_oth,*pJ,*pj;
43982412ba7SHong Zhang   PetscInt       i,j,k,anz,pnz,apnz,nextap,row,*cj;
44082412ba7SHong Zhang   MatScalar      *ada=ad->a,*aoa=ao->a,*apa,*pa,*ca,*pa_loc,*pa_oth;
44182412ba7SHong Zhang   PetscInt       am=A->m,cm=C->m,pon=(p->B)->n;
44242c57489SHong Zhang   MPI_Comm       comm=C->comm;
44342c57489SHong Zhang   PetscMPIInt    size,rank,taga,*len_s;
44482412ba7SHong Zhang   PetscInt       *owners,proc,nrows,**buf_ri_k,**nextrow,**nextci;
44542c57489SHong Zhang   PetscInt       **buf_ri,**buf_rj;
44650f5bf86SHong Zhang   PetscInt       cnz=0,*bj_i,*bi,*bj,bnz,nextcj; /* bi,bj,ba: local array of C(mpi mat) */
44742c57489SHong Zhang   MPI_Request    *s_waits,*r_waits;
44842c57489SHong Zhang   MPI_Status     *status;
44982412ba7SHong Zhang   MatScalar      **abuf_r,*ba_i,*pA,*coa,*ba;
45082412ba7SHong Zhang   PetscInt       *api,*apj,*coi,*coj;
45182412ba7SHong Zhang   PetscInt       *poJ=po->j,*pdJ=pd->j,pcstart=p->cstart,pcend=p->cend;
45242c57489SHong Zhang 
45342c57489SHong Zhang   PetscFunctionBegin;
45442c57489SHong Zhang   ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr);
45542c57489SHong Zhang   ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr);
45642c57489SHong Zhang 
45742c57489SHong Zhang   ierr = PetscObjectQuery((PetscObject)C,"MatMergeSeqsToMPI",(PetscObject *)&cont_merge);CHKERRQ(ierr);
45842c57489SHong Zhang   if (cont_merge) {
45942c57489SHong Zhang     ierr  = PetscObjectContainerGetPointer(cont_merge,(void **)&merge);CHKERRQ(ierr);
46042c57489SHong Zhang   } else {
46142c57489SHong Zhang     SETERRQ(PETSC_ERR_ARG_WRONGSTATE, "Matrix C does not posses an object container");
46242c57489SHong Zhang   }
463*f33d1a9aSHong Zhang   ierr = PetscObjectQuery((PetscObject)P,"Mat_MatMatMultMPI",(PetscObject *)&cont_ptap);CHKERRQ(ierr);
46442c57489SHong Zhang   if (cont_ptap) {
465de0260b3SHong Zhang     ierr  = PetscObjectContainerGetPointer(cont_ptap,(void **)&ap);CHKERRQ(ierr);
46642c57489SHong Zhang   } else {
46742c57489SHong Zhang     SETERRQ(PETSC_ERR_ARG_WRONGSTATE, "Matrix P does not posses an object container");
46842c57489SHong Zhang   }
46942c57489SHong Zhang 
47042c57489SHong Zhang   /* get data from symbolic products */
471de0260b3SHong Zhang   p_loc = (Mat_SeqAIJ*)(ap->B_loc)->data;
472de0260b3SHong Zhang   p_oth = (Mat_SeqAIJ*)(ap->B_oth)->data;
47382412ba7SHong Zhang   pi_loc=p_loc->i; pj_loc=p_loc->j; pJ=pj_loc; pa_loc=p_loc->a,pA=pa_loc;
47442c57489SHong Zhang   pi_oth=p_oth->i; pj_oth=p_oth->j; pa_oth=p_oth->a;
47542c57489SHong Zhang 
476de0260b3SHong Zhang   coi = merge->coi; coj = merge->coj;
477de0260b3SHong Zhang   ierr = PetscMalloc((coi[pon]+1)*sizeof(MatScalar),&coa);CHKERRQ(ierr);
478de0260b3SHong Zhang   ierr = PetscMemzero(coa,coi[pon]*sizeof(MatScalar));CHKERRQ(ierr);
479de0260b3SHong Zhang 
480de0260b3SHong Zhang   bi = merge->bi; bj = merge->bj;
481de0260b3SHong Zhang   ierr  = PetscMapGetGlobalRange(merge->rowmap,&owners);CHKERRQ(ierr);
482de0260b3SHong Zhang   ierr  = PetscMalloc((bi[cm]+1)*sizeof(MatScalar),&ba);CHKERRQ(ierr);
483de0260b3SHong Zhang   ierr  = PetscMemzero(ba,bi[cm]*sizeof(MatScalar));CHKERRQ(ierr);
48442c57489SHong Zhang 
485f4a743e1SHong Zhang   /* get data from symbolic A*P */
486de0260b3SHong Zhang   ierr = PetscMalloc((ap->abnz_max+1)*sizeof(MatScalar),&apa);CHKERRQ(ierr);
487de0260b3SHong Zhang   ierr = PetscMemzero(apa,ap->abnz_max*sizeof(MatScalar));CHKERRQ(ierr);
48842c57489SHong Zhang 
489f4a743e1SHong Zhang   /* compute numeric C_seq=P_loc^T*A_loc*P */
49082412ba7SHong Zhang   api = ap->abi; apj = ap->abj;
49142c57489SHong Zhang   for (i=0;i<am;i++) {
492f4a743e1SHong Zhang     /* form i-th sparse row of A*P */
49382412ba7SHong Zhang     apnz = api[i+1] - api[i];
49482412ba7SHong Zhang     apJ  = apj + api[i];
49542c57489SHong Zhang     /* diagonal portion of A */
49682412ba7SHong Zhang     anz  = adi[i+1] - adi[i];
49782412ba7SHong Zhang     for (j=0;j<anz;j++) {
49882412ba7SHong Zhang       row = *adj++;
49982412ba7SHong Zhang       pnz = pi_loc[row+1] - pi_loc[row];
50082412ba7SHong Zhang       pj  = pj_loc + pi_loc[row];
50182412ba7SHong Zhang       pa  = pa_loc + pi_loc[row];
502f4a743e1SHong Zhang       nextp = 0;
50382412ba7SHong Zhang       for (k=0; nextp<pnz; k++) {
50482412ba7SHong Zhang         if (apJ[k] == pj[nextp]) { /* col of AP == col of P */
50582412ba7SHong Zhang           apa[k] += (*ada)*pa[nextp++];
50642c57489SHong Zhang         }
50742c57489SHong Zhang       }
50882412ba7SHong Zhang       flops += 2*pnz;
50942c57489SHong Zhang       ada++;
51042c57489SHong Zhang     }
51142c57489SHong Zhang     /* off-diagonal portion of A */
51282412ba7SHong Zhang     anz  = aoi[i+1] - aoi[i];
51382412ba7SHong Zhang     for (j=0; j<anz; j++) {
51482412ba7SHong Zhang       row = *aoj++;
51582412ba7SHong Zhang       pnz = pi_oth[row+1] - pi_oth[row];
51682412ba7SHong Zhang       pj  = pj_oth + pi_oth[row];
51782412ba7SHong Zhang       pa  = pa_oth + pi_oth[row];
518f4a743e1SHong Zhang       nextp = 0;
51982412ba7SHong Zhang       for (k=0; nextp<pnz; k++) {
52082412ba7SHong Zhang         if (apJ[k] == pj[nextp]) { /* col of AP == col of P */
52182412ba7SHong Zhang           apa[k] += (*aoa)*pa[nextp++];
52242c57489SHong Zhang         }
52342c57489SHong Zhang       }
52482412ba7SHong Zhang       flops += 2*pnz;
52542c57489SHong Zhang       aoa++;
52642c57489SHong Zhang     }
52742c57489SHong Zhang 
52882412ba7SHong Zhang     /* Compute P_loc[i,:]^T*AP[i,:] using outer product */
52982412ba7SHong Zhang     pnz = pi_loc[i+1] - pi_loc[i];
53082412ba7SHong Zhang     for (j=0; j<pnz; j++) {
53142c57489SHong Zhang       nextap = 0;
53282412ba7SHong Zhang       row    = *pJ++; /* global index */
53382412ba7SHong Zhang       if (row < pcstart || row >=pcend) { /* put the value into Co */
53482412ba7SHong Zhang         cj  = coj + coi[*poJ];
53582412ba7SHong Zhang         ca  = coa + coi[*poJ++];
53682412ba7SHong Zhang       } else {                            /* put the value into Cd */
53782412ba7SHong Zhang         cj   = bj + bi[*pdJ];
53882412ba7SHong Zhang         ca   = ba + bi[*pdJ++];
53942c57489SHong Zhang       }
54082412ba7SHong Zhang       for (k=0; nextap<apnz; k++) {
54182412ba7SHong Zhang         if (cj[k]==apJ[nextap]) ca[k] += (*pA)*apa[nextap++];
54242c57489SHong Zhang       }
54382412ba7SHong Zhang       flops += 2*apnz;
54482412ba7SHong Zhang       pA++;
545de0260b3SHong Zhang     }
546de0260b3SHong Zhang 
54742c57489SHong Zhang     /* zero the current row info for A*P */
54882412ba7SHong Zhang     ierr = PetscMemzero(apa,apnz*sizeof(MatScalar));CHKERRQ(ierr);
54942c57489SHong Zhang   }
55042c57489SHong Zhang   ierr = PetscFree(apa);CHKERRQ(ierr);
55142c57489SHong Zhang 
552de0260b3SHong Zhang   /* send and recv matrix values */
553de0260b3SHong Zhang   /*-----------------------------*/
55442c57489SHong Zhang   buf_ri = merge->buf_ri;
55542c57489SHong Zhang   buf_rj = merge->buf_rj;
55642c57489SHong Zhang   len_s  = merge->len_s;
55742c57489SHong Zhang   ierr = PetscObjectGetNewTag((PetscObject)merge->rowmap,&taga);CHKERRQ(ierr);
55842c57489SHong Zhang   ierr = PetscPostIrecvScalar(comm,taga,merge->nrecv,merge->id_r,merge->len_r,&abuf_r,&r_waits);CHKERRQ(ierr);
55942c57489SHong Zhang 
56042c57489SHong Zhang   ierr = PetscMalloc((merge->nsend+1)*sizeof(MPI_Request),&s_waits);CHKERRQ(ierr);
56142c57489SHong Zhang   for (proc=0,k=0; proc<size; proc++){
56242c57489SHong Zhang     if (!len_s[proc]) continue;
563de0260b3SHong Zhang     i = merge->owners_co[proc];
564de0260b3SHong Zhang     ierr = MPI_Isend(coa+coi[i],len_s[proc],MPIU_MATSCALAR,proc,taga,comm,s_waits+k);CHKERRQ(ierr);
56542c57489SHong Zhang     k++;
56642c57489SHong Zhang   }
56782412ba7SHong Zhang   ierr = PetscMalloc(size*sizeof(MPI_Status),&status);CHKERRQ(ierr);
56842c57489SHong Zhang   ierr = MPI_Waitall(merge->nrecv,r_waits,status);CHKERRQ(ierr);
56942c57489SHong Zhang   ierr = MPI_Waitall(merge->nsend,s_waits,status);CHKERRQ(ierr);
57042c57489SHong Zhang   ierr = PetscFree(status);CHKERRQ(ierr);
57142c57489SHong Zhang 
57242c57489SHong Zhang   ierr = PetscFree(s_waits);CHKERRQ(ierr);
57342c57489SHong Zhang   ierr = PetscFree(r_waits);CHKERRQ(ierr);
57482412ba7SHong Zhang   ierr = PetscFree(coa);CHKERRQ(ierr);
57542c57489SHong Zhang 
57682412ba7SHong Zhang   /* insert local and received values into C */
57782412ba7SHong Zhang   /*-----------------------------------------*/
57842c57489SHong Zhang   ierr = PetscMalloc((3*merge->nrecv+1)*sizeof(PetscInt**),&buf_ri_k);CHKERRQ(ierr);
57942c57489SHong Zhang   nextrow   = buf_ri_k + merge->nrecv;
58082412ba7SHong Zhang   nextci = nextrow + merge->nrecv;
58142c57489SHong Zhang 
58242c57489SHong Zhang   for (k=0; k<merge->nrecv; k++){
58342c57489SHong Zhang     buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */
58442c57489SHong Zhang     nrows       = *(buf_ri_k[k]);
58542c57489SHong Zhang     nextrow[k]  = buf_ri_k[k]+1;  /* next row number of k-th recved i-structure */
58682412ba7SHong Zhang     nextci[k]   = buf_ri_k[k] + (nrows + 1);/* poins to the next i-structure of k-th recved i-structure  */
58742c57489SHong Zhang   }
58842c57489SHong Zhang 
589de0260b3SHong Zhang   for (i=0; i<cm; i++) {
59082412ba7SHong Zhang     row = owners[rank] + i; /* global row index of C_seq */
59142c57489SHong Zhang     bj_i = bj + bi[i];  /* col indices of the i-th row of C */
592de0260b3SHong Zhang     ba_i = ba + bi[i];
59382412ba7SHong Zhang     bnz  = bi[i+1] - bi[i];
59442c57489SHong Zhang     /* add received vals into ba */
59542c57489SHong Zhang     for (k=0; k<merge->nrecv; k++){ /* k-th received message */
59642c57489SHong Zhang       /* i-th row */
59742c57489SHong Zhang       if (i == *nextrow[k]) {
59882412ba7SHong Zhang         cnz = *(nextci[k]+1) - *nextci[k];
59982412ba7SHong Zhang         cj  = buf_rj[k] + *(nextci[k]);
60082412ba7SHong Zhang         ca  = abuf_r[k] + *(nextci[k]);
60182412ba7SHong Zhang         nextcj = 0;
60282412ba7SHong Zhang         for (j=0; nextcj<cnz; j++){
60382412ba7SHong Zhang           if (bj_i[j] == cj[nextcj]){ /* bcol == ccol */
60482412ba7SHong Zhang             ba_i[j] += ca[nextcj++];
60542c57489SHong Zhang           }
60642c57489SHong Zhang         }
60782412ba7SHong Zhang         nextrow[k]++; nextci[k]++;
60842c57489SHong Zhang       }
60942c57489SHong Zhang     }
61082412ba7SHong Zhang     ierr = MatSetValues(C,1,&row,bnz,bj_i,ba_i,INSERT_VALUES);CHKERRQ(ierr);
61182412ba7SHong Zhang     flops += 2*cnz;
61242c57489SHong Zhang   }
61382412ba7SHong Zhang   ierr = MatSetOption(C,MAT_COLUMNS_SORTED);CHKERRQ(ierr); /* -cause delay? */
61442c57489SHong Zhang   ierr = MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
61542c57489SHong Zhang   ierr = MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
61642c57489SHong Zhang 
617de0260b3SHong Zhang   ierr = PetscFree(ba);CHKERRQ(ierr);
61842c57489SHong Zhang   ierr = PetscFree(abuf_r);CHKERRQ(ierr);
61942c57489SHong Zhang   ierr = PetscFree(buf_ri_k);CHKERRQ(ierr);
62042c57489SHong Zhang   ierr = PetscLogFlops(flops);CHKERRQ(ierr);
62142c57489SHong Zhang 
62242c57489SHong Zhang   PetscFunctionReturn(0);
62342c57489SHong Zhang }
624