xref: /petsc/src/mat/impls/baij/mpi/baijov.c (revision 6849ba73f22fecb8f92ef896a42e4e8bd4cd6965)
1 /*
2    Routines to compute overlapping regions of a parallel MPI matrix
3   and to find submatrices that were shared across processors.
4 */
5 #include "src/mat/impls/baij/mpi/mpibaij.h"
6 #include "petscbt.h"
7 
8 static PetscErrorCode MatIncreaseOverlap_MPIBAIJ_Once(Mat,int,IS *);
9 static PetscErrorCode MatIncreaseOverlap_MPIBAIJ_Local(Mat,int,char **,int*,int**);
10 static PetscErrorCode MatIncreaseOverlap_MPIBAIJ_Receive(Mat,int,int **,int**,int*);
11 EXTERN PetscErrorCode MatGetRow_MPIBAIJ(Mat,int,int*,int**,PetscScalar**);
12 EXTERN PetscErrorCode MatRestoreRow_MPIBAIJ(Mat,int,int*,int**,PetscScalar**);
13 
14 #undef __FUNCT__
15 #define __FUNCT__ "MatIncreaseOverlap_MPIBAIJ"
16 PetscErrorCode MatIncreaseOverlap_MPIBAIJ(Mat C,int imax,IS is[],int ov)
17 {
18   Mat_MPIBAIJ  *c = (Mat_MPIBAIJ*)C->data;
19   PetscErrorCode ierr;
20   int          i,N=C->N, bs=c->bs;
21   IS           *is_new;
22 
23   PetscFunctionBegin;
24   ierr = PetscMalloc(imax*sizeof(IS),&is_new);CHKERRQ(ierr);
25   /* Convert the indices into block format */
26   ierr = ISCompressIndicesGeneral(N,bs,imax,is,is_new);CHKERRQ(ierr);
27   if (ov < 0){ SETERRQ(PETSC_ERR_ARG_OUTOFRANGE,"Negative overlap specified\n");}
28   for (i=0; i<ov; ++i) {
29     ierr = MatIncreaseOverlap_MPIBAIJ_Once(C,imax,is_new);CHKERRQ(ierr);
30   }
31   for (i=0; i<imax; i++) {ierr = ISDestroy(is[i]);CHKERRQ(ierr);}
32   ierr = ISExpandIndicesGeneral(N,bs,imax,is_new,is);CHKERRQ(ierr);
33   for (i=0; i<imax; i++) {ierr = ISDestroy(is_new[i]);CHKERRQ(ierr);}
34   ierr = PetscFree(is_new);CHKERRQ(ierr);
35   PetscFunctionReturn(0);
36 }
37 
38 /*
39   Sample message format:
40   If a processor A wants processor B to process some elements corresponding
41   to index sets is[1], is[5]
42   mesg [0] = 2   (no of index sets in the mesg)
43   -----------
44   mesg [1] = 1 => is[1]
45   mesg [2] = sizeof(is[1]);
46   -----------
47   mesg [5] = 5  => is[5]
48   mesg [6] = sizeof(is[5]);
49   -----------
50   mesg [7]
51   mesg [n]  data(is[1])
52   -----------
53   mesg[n+1]
54   mesg[m]  data(is[5])
55   -----------
56 
57   Notes:
58   nrqs - no of requests sent (or to be sent out)
59   nrqr - no of requests recieved (which have to be or which have been processed
60 */
61 #undef __FUNCT__
62 #define __FUNCT__ "MatIncreaseOverlap_MPIBAIJ_Once"
63 static PetscErrorCode MatIncreaseOverlap_MPIBAIJ_Once(Mat C,int imax,IS is[])
64 {
65   Mat_MPIBAIJ  *c = (Mat_MPIBAIJ*)C->data;
66   int         **idx,*n,*w1,*w2,*w3,*w4,*rtable,**data,len,*idx_i;
67   PetscErrorCode ierr;
68   int         size,rank,Mbs,i,j,k,**rbuf,row,proc,nrqs,msz,**outdat,**ptr;
69   int         *ctr,*pa,*tmp,nrqr,*isz,*isz1,**xdata,**rbuf2;
70   int         *onodes1,*olengths1,tag1,tag2,*onodes2,*olengths2;
71   PetscBT     *table;
72   MPI_Comm    comm;
73   MPI_Request *s_waits1,*r_waits1,*s_waits2,*r_waits2;
74   MPI_Status  *s_status,*recv_status;
75 
76   PetscFunctionBegin;
77   comm   = C->comm;
78   size   = c->size;
79   rank   = c->rank;
80   Mbs    = c->Mbs;
81 
82   ierr = PetscObjectGetNewTag((PetscObject)C,&tag1);CHKERRQ(ierr);
83   ierr = PetscObjectGetNewTag((PetscObject)C,&tag2);CHKERRQ(ierr);
84 
85   len    = (imax+1)*sizeof(int*)+ (imax + Mbs)*sizeof(int);
86   ierr   = PetscMalloc(len,&idx);CHKERRQ(ierr);
87   n      = (int*)(idx + imax);
88   rtable = n + imax;
89 
90   for (i=0; i<imax; i++) {
91     ierr = ISGetIndices(is[i],&idx[i]);CHKERRQ(ierr);
92     ierr = ISGetLocalSize(is[i],&n[i]);CHKERRQ(ierr);
93   }
94 
95   /* Create hash table for the mapping :row -> proc*/
96   for (i=0,j=0; i<size; i++) {
97     len = c->rowners[i+1];
98     for (; j<len; j++) {
99       rtable[j] = i;
100     }
101   }
102 
103   /* evaluate communication - mesg to who,length of mesg, and buffer space
104      required. Based on this, buffers are allocated, and data copied into them*/
105   ierr = PetscMalloc(size*4*sizeof(int),&w1);CHKERRQ(ierr);/*  mesg size */
106   w2   = w1 + size;       /* if w2[i] marked, then a message to proc i*/
107   w3   = w2 + size;       /* no of IS that needs to be sent to proc i */
108   w4   = w3 + size;       /* temp work space used in determining w1, w2, w3 */
109   ierr = PetscMemzero(w1,size*3*sizeof(int));CHKERRQ(ierr); /* initialise work vector*/
110   for (i=0; i<imax; i++) {
111     ierr  = PetscMemzero(w4,size*sizeof(int));CHKERRQ(ierr); /* initialise work vector*/
112     idx_i = idx[i];
113     len   = n[i];
114     for (j=0; j<len; j++) {
115       row  = idx_i[j];
116       if (row < 0) {
117         SETERRQ(1,"Index set cannot have negative entries");
118       }
119       proc = rtable[row];
120       w4[proc]++;
121     }
122     for (j=0; j<size; j++){
123       if (w4[j]) { w1[j] += w4[j]; w3[j]++;}
124     }
125   }
126 
127   nrqs     = 0;              /* no of outgoing messages */
128   msz      = 0;              /* total mesg length (for all proc */
129   w1[rank] = 0;              /* no mesg sent to itself */
130   w3[rank] = 0;
131   for (i=0; i<size; i++) {
132     if (w1[i])  {w2[i] = 1; nrqs++;} /* there exists a message to proc i */
133   }
134   /* pa - is list of processors to communicate with */
135   ierr = PetscMalloc((nrqs+1)*sizeof(int),&pa);CHKERRQ(ierr);
136   for (i=0,j=0; i<size; i++) {
137     if (w1[i]) {pa[j] = i; j++;}
138   }
139 
140   /* Each message would have a header = 1 + 2*(no of IS) + data */
141   for (i=0; i<nrqs; i++) {
142     j      = pa[i];
143     w1[j] += w2[j] + 2*w3[j];
144     msz   += w1[j];
145   }
146 
147   /* Determine the number of messages to expect, their lengths, from from-ids */
148   ierr = PetscGatherNumberOfMessages(comm,w2,w1,&nrqr);CHKERRQ(ierr);
149   ierr = PetscGatherMessageLengths(comm,nrqs,nrqr,w1,&onodes1,&olengths1);CHKERRQ(ierr);
150 
151   /* Now post the Irecvs corresponding to these messages */
152   ierr = PetscPostIrecvInt(comm,tag1,nrqr,onodes1,olengths1,&rbuf,&r_waits1);CHKERRQ(ierr);
153 
154   /* Allocate Memory for outgoing messages */
155   len    = 2*size*sizeof(int*) + (size+msz)*sizeof(int);
156   ierr = PetscMalloc(len,&outdat);CHKERRQ(ierr);
157   ptr    = outdat + size;     /* Pointers to the data in outgoing buffers */
158   ierr   = PetscMemzero(outdat,2*size*sizeof(int*));CHKERRQ(ierr);
159   tmp    = (int*)(outdat + 2*size);
160   ctr    = tmp + msz;
161 
162   {
163     int *iptr = tmp,ict  = 0;
164     for (i=0; i<nrqs; i++) {
165       j         = pa[i];
166       iptr     +=  ict;
167       outdat[j] = iptr;
168       ict       = w1[j];
169     }
170   }
171 
172   /* Form the outgoing messages */
173   /*plug in the headers*/
174   for (i=0; i<nrqs; i++) {
175     j            = pa[i];
176     outdat[j][0] = 0;
177     ierr = PetscMemzero(outdat[j]+1,2*w3[j]*sizeof(int));CHKERRQ(ierr);
178     ptr[j]       = outdat[j] + 2*w3[j] + 1;
179   }
180 
181   /* Memory for doing local proc's work*/
182   {
183     int  *d_p;
184     char *t_p;
185 
186     len  = (imax)*(sizeof(PetscBT) + sizeof(int*)+ sizeof(int)) +
187       (Mbs)*imax*sizeof(int)  + (Mbs/PETSC_BITS_PER_BYTE+1)*imax*sizeof(char) + 1;
188     ierr = PetscMalloc(len,&table);CHKERRQ(ierr);
189     ierr = PetscMemzero(table,len);CHKERRQ(ierr);
190     data = (int **)(table + imax);
191     isz  = (int  *)(data  + imax);
192     d_p  = (int  *)(isz   + imax);
193     t_p  = (char *)(d_p   + Mbs*imax);
194     for (i=0; i<imax; i++) {
195       table[i] = t_p + (Mbs/PETSC_BITS_PER_BYTE+1)*i;
196       data[i]  = d_p + (Mbs)*i;
197     }
198   }
199 
200   /* Parse the IS and update local tables and the outgoing buf with the data*/
201   {
202     int     n_i,*data_i,isz_i,*outdat_j,ctr_j;
203     PetscBT table_i;
204 
205     for (i=0; i<imax; i++) {
206       ierr    = PetscMemzero(ctr,size*sizeof(int));CHKERRQ(ierr);
207       n_i     = n[i];
208       table_i = table[i];
209       idx_i   = idx[i];
210       data_i  = data[i];
211       isz_i   = isz[i];
212       for (j=0;  j<n_i; j++) {  /* parse the indices of each IS */
213         row  = idx_i[j];
214         proc = rtable[row];
215         if (proc != rank) { /* copy to the outgoing buffer */
216           ctr[proc]++;
217           *ptr[proc] = row;
218           ptr[proc]++;
219         }
220         else { /* Update the local table */
221           if (!PetscBTLookupSet(table_i,row)) { data_i[isz_i++] = row;}
222         }
223       }
224       /* Update the headers for the current IS */
225       for (j=0; j<size; j++) { /* Can Optimise this loop by using pa[] */
226         if ((ctr_j = ctr[j])) {
227           outdat_j        = outdat[j];
228           k               = ++outdat_j[0];
229           outdat_j[2*k]   = ctr_j;
230           outdat_j[2*k-1] = i;
231         }
232       }
233       isz[i] = isz_i;
234     }
235   }
236 
237   /*  Now  post the sends */
238   ierr = PetscMalloc((nrqs+1)*sizeof(MPI_Request),&s_waits1);CHKERRQ(ierr);
239   for (i=0; i<nrqs; ++i) {
240     j    = pa[i];
241     ierr = MPI_Isend(outdat[j],w1[j],MPI_INT,j,tag1,comm,s_waits1+i);CHKERRQ(ierr);
242   }
243 
244   /* No longer need the original indices*/
245   for (i=0; i<imax; ++i) {
246     ierr = ISRestoreIndices(is[i],idx+i);CHKERRQ(ierr);
247   }
248   ierr = PetscFree(idx);CHKERRQ(ierr);
249 
250   for (i=0; i<imax; ++i) {
251     ierr = ISDestroy(is[i]);CHKERRQ(ierr);
252   }
253 
254   /* Do Local work*/
255   ierr = MatIncreaseOverlap_MPIBAIJ_Local(C,imax,table,isz,data);CHKERRQ(ierr);
256 
257   /* Receive messages*/
258   ierr = PetscMalloc((nrqr+1)*sizeof(MPI_Status),&recv_status);CHKERRQ(ierr);
259   ierr = MPI_Waitall(nrqr,r_waits1,recv_status);CHKERRQ(ierr);
260 
261   ierr = PetscMalloc((nrqs+1)*sizeof(MPI_Status),&s_status);CHKERRQ(ierr);
262   ierr = MPI_Waitall(nrqs,s_waits1,s_status);CHKERRQ(ierr);
263 
264   /* Phase 1 sends are complete - deallocate buffers */
265   ierr = PetscFree(outdat);CHKERRQ(ierr);
266   ierr = PetscFree(w1);CHKERRQ(ierr);
267 
268   ierr = PetscMalloc((nrqr+1)*sizeof(int*),&xdata);CHKERRQ(ierr);
269   ierr = PetscMalloc((nrqr+1)*sizeof(int),&isz1);CHKERRQ(ierr);
270   ierr = MatIncreaseOverlap_MPIBAIJ_Receive(C,nrqr,rbuf,xdata,isz1);CHKERRQ(ierr);
271   ierr = PetscFree(rbuf);CHKERRQ(ierr);
272 
273   /* Send the data back*/
274   /* Do a global reduction to know the buffer space req for incoming messages*/
275   {
276     int *rw1;
277 
278     ierr = PetscMalloc(size*sizeof(int),&rw1);CHKERRQ(ierr);
279     ierr = PetscMemzero(rw1,size*sizeof(int));CHKERRQ(ierr);
280 
281     for (i=0; i<nrqr; ++i) {
282       proc      = recv_status[i].MPI_SOURCE;
283       if (proc != onodes1[i]) SETERRQ(1,"MPI_SOURCE mismatch");
284       rw1[proc] = isz1[i];
285     }
286 
287     ierr = PetscFree(onodes1);CHKERRQ(ierr);
288     ierr = PetscFree(olengths1);CHKERRQ(ierr);
289 
290     /* Determine the number of messages to expect, their lengths, from from-ids */
291     ierr = PetscGatherMessageLengths(comm,nrqr,nrqs,rw1,&onodes2,&olengths2);CHKERRQ(ierr);
292     PetscFree(rw1);
293   }
294   /* Now post the Irecvs corresponding to these messages */
295   ierr = PetscPostIrecvInt(comm,tag2,nrqs,onodes2,olengths2,&rbuf2,&r_waits2);CHKERRQ(ierr);
296 
297   /*  Now  post the sends */
298   ierr = PetscMalloc((nrqr+1)*sizeof(MPI_Request),&s_waits2);CHKERRQ(ierr);
299   for (i=0; i<nrqr; ++i) {
300     j    = recv_status[i].MPI_SOURCE;
301     ierr = MPI_Isend(xdata[i],isz1[i],MPI_INT,j,tag2,comm,s_waits2+i);CHKERRQ(ierr);
302   }
303 
304   /* receive work done on other processors*/
305   {
306     int         idex,is_no,ct1,max,*rbuf2_i,isz_i,*data_i,jmax;
307     PetscBT     table_i;
308     MPI_Status  *status2;
309 
310     ierr = PetscMalloc((PetscMax(nrqr,nrqs)+1)*sizeof(MPI_Status),&status2);CHKERRQ(ierr);
311 
312     for (i=0; i<nrqs; ++i) {
313       ierr = MPI_Waitany(nrqs,r_waits2,&idex,status2+i);CHKERRQ(ierr);
314       /* Process the message*/
315       rbuf2_i = rbuf2[idex];
316       ct1     = 2*rbuf2_i[0]+1;
317       jmax    = rbuf2[idex][0];
318       for (j=1; j<=jmax; j++) {
319         max     = rbuf2_i[2*j];
320         is_no   = rbuf2_i[2*j-1];
321         isz_i   = isz[is_no];
322         data_i  = data[is_no];
323         table_i = table[is_no];
324         for (k=0; k<max; k++,ct1++) {
325           row = rbuf2_i[ct1];
326           if (!PetscBTLookupSet(table_i,row)) { data_i[isz_i++] = row;}
327         }
328         isz[is_no] = isz_i;
329       }
330     }
331     ierr = MPI_Waitall(nrqr,s_waits2,status2);CHKERRQ(ierr);
332     ierr = PetscFree(status2);CHKERRQ(ierr);
333   }
334 
335   for (i=0; i<imax; ++i) {
336     ierr = ISCreateGeneral(PETSC_COMM_SELF,isz[i],data[i],is+i);CHKERRQ(ierr);
337   }
338 
339 
340   ierr = PetscFree(onodes2);CHKERRQ(ierr);
341   ierr = PetscFree(olengths2);CHKERRQ(ierr);
342 
343   ierr = PetscFree(pa);CHKERRQ(ierr);
344   ierr = PetscFree(rbuf2);CHKERRQ(ierr);
345   ierr = PetscFree(s_waits1);CHKERRQ(ierr);
346   ierr = PetscFree(r_waits1);CHKERRQ(ierr);
347   ierr = PetscFree(s_waits2);CHKERRQ(ierr);
348   ierr = PetscFree(r_waits2);CHKERRQ(ierr);
349   ierr = PetscFree(table);CHKERRQ(ierr);
350   ierr = PetscFree(s_status);CHKERRQ(ierr);
351   ierr = PetscFree(recv_status);CHKERRQ(ierr);
352   ierr = PetscFree(xdata[0]);CHKERRQ(ierr);
353   ierr = PetscFree(xdata);CHKERRQ(ierr);
354   ierr = PetscFree(isz1);CHKERRQ(ierr);
355   PetscFunctionReturn(0);
356 }
357 
358 #undef __FUNCT__
359 #define __FUNCT__ "MatIncreaseOverlap_MPIBAIJ_Local"
360 /*
361    MatIncreaseOverlap_MPIBAIJ_Local - Called by MatincreaseOverlap, to do
362        the work on the local processor.
363 
364      Inputs:
365       C      - MAT_MPIBAIJ;
366       imax - total no of index sets processed at a time;
367       table  - an array of char - size = Mbs bits.
368 
369      Output:
370       isz    - array containing the count of the solution elements corresponding
371                to each index set;
372       data   - pointer to the solutions
373 */
374 static PetscErrorCode MatIncreaseOverlap_MPIBAIJ_Local(Mat C,int imax,PetscBT *table,int *isz,int **data)
375 {
376   Mat_MPIBAIJ *c = (Mat_MPIBAIJ*)C->data;
377   Mat         A = c->A,B = c->B;
378   Mat_SeqBAIJ *a = (Mat_SeqBAIJ*)A->data,*b = (Mat_SeqBAIJ*)B->data;
379   int         start,end,val,max,rstart,cstart,*ai,*aj;
380   int         *bi,*bj,*garray,i,j,k,row,*data_i,isz_i;
381   PetscBT     table_i;
382 
383   PetscFunctionBegin;
384   rstart = c->rstart;
385   cstart = c->cstart;
386   ai     = a->i;
387   aj     = a->j;
388   bi     = b->i;
389   bj     = b->j;
390   garray = c->garray;
391 
392 
393   for (i=0; i<imax; i++) {
394     data_i  = data[i];
395     table_i = table[i];
396     isz_i   = isz[i];
397     for (j=0,max=isz[i]; j<max; j++) {
398       row   = data_i[j] - rstart;
399       start = ai[row];
400       end   = ai[row+1];
401       for (k=start; k<end; k++) { /* Amat */
402         val = aj[k] + cstart;
403         if (!PetscBTLookupSet(table_i,val)) { data_i[isz_i++] = val;}
404       }
405       start = bi[row];
406       end   = bi[row+1];
407       for (k=start; k<end; k++) { /* Bmat */
408         val = garray[bj[k]];
409         if (!PetscBTLookupSet(table_i,val)) { data_i[isz_i++] = val;}
410       }
411     }
412     isz[i] = isz_i;
413   }
414   PetscFunctionReturn(0);
415 }
416 #undef __FUNCT__
417 #define __FUNCT__ "MatIncreaseOverlap_MPIBAIJ_Receive"
418 /*
419       MatIncreaseOverlap_MPIBAIJ_Receive - Process the recieved messages,
420          and return the output
421 
422          Input:
423            C    - the matrix
424            nrqr - no of messages being processed.
425            rbuf - an array of pointers to the recieved requests
426 
427          Output:
428            xdata - array of messages to be sent back
429            isz1  - size of each message
430 
431   For better efficiency perhaps we should malloc seperately each xdata[i],
432 then if a remalloc is required we need only copy the data for that one row
433 rather than all previous rows as it is now where a single large chunck of
434 memory is used.
435 
436 */
437 static PetscErrorCode MatIncreaseOverlap_MPIBAIJ_Receive(Mat C,int nrqr,int **rbuf,int **xdata,int * isz1)
438 {
439   Mat_MPIBAIJ *c = (Mat_MPIBAIJ*)C->data;
440   Mat         A = c->A,B = c->B;
441   Mat_SeqBAIJ *a = (Mat_SeqBAIJ*)A->data,*b = (Mat_SeqBAIJ*)B->data;
442   PetscErrorCode ierr;
443   int         rstart,cstart,*ai,*aj,*bi,*bj,*garray,i,j,k;
444   int         row,total_sz,ct,ct1,ct2,ct3,mem_estimate,oct2,l,start,end;
445   int         val,max1,max2,rank,Mbs,no_malloc =0,*tmp,new_estimate,ctr;
446   int         *rbuf_i,kmax,rbuf_0;
447   PetscBT     xtable;
448 
449   PetscFunctionBegin;
450   rank   = c->rank;
451   Mbs    = c->Mbs;
452   rstart = c->rstart;
453   cstart = c->cstart;
454   ai     = a->i;
455   aj     = a->j;
456   bi     = b->i;
457   bj     = b->j;
458   garray = c->garray;
459 
460 
461   for (i=0,ct=0,total_sz=0; i<nrqr; ++i) {
462     rbuf_i  =  rbuf[i];
463     rbuf_0  =  rbuf_i[0];
464     ct     += rbuf_0;
465     for (j=1; j<=rbuf_0; j++) { total_sz += rbuf_i[2*j]; }
466   }
467 
468   if (c->Mbs) max1 = ct*(a->nz +b->nz)/c->Mbs;
469   else        max1 = 1;
470   mem_estimate = 3*((total_sz > max1 ? total_sz : max1)+1);
471   ierr         = PetscMalloc(mem_estimate*sizeof(int),&xdata[0]);CHKERRQ(ierr);
472   ++no_malloc;
473   ierr         = PetscBTCreate(Mbs,xtable);CHKERRQ(ierr);
474   ierr         = PetscMemzero(isz1,nrqr*sizeof(int));CHKERRQ(ierr);
475 
476   ct3 = 0;
477   for (i=0; i<nrqr; i++) { /* for easch mesg from proc i */
478     rbuf_i =  rbuf[i];
479     rbuf_0 =  rbuf_i[0];
480     ct1    =  2*rbuf_0+1;
481     ct2    =  ct1;
482     ct3    += ct1;
483     for (j=1; j<=rbuf_0; j++) { /* for each IS from proc i*/
484       ierr = PetscBTMemzero(Mbs,xtable);CHKERRQ(ierr);
485       oct2 = ct2;
486       kmax = rbuf_i[2*j];
487       for (k=0; k<kmax; k++,ct1++) {
488         row = rbuf_i[ct1];
489         if (!PetscBTLookupSet(xtable,row)) {
490           if (!(ct3 < mem_estimate)) {
491             new_estimate = (int)(1.5*mem_estimate)+1;
492             ierr = PetscMalloc(new_estimate * sizeof(int),&tmp);CHKERRQ(ierr);
493             ierr = PetscMemcpy(tmp,xdata[0],mem_estimate*sizeof(int));CHKERRQ(ierr);
494             ierr = PetscFree(xdata[0]);CHKERRQ(ierr);
495             xdata[0]     = tmp;
496             mem_estimate = new_estimate; ++no_malloc;
497             for (ctr=1; ctr<=i; ctr++) { xdata[ctr] = xdata[ctr-1] + isz1[ctr-1];}
498           }
499           xdata[i][ct2++] = row;
500           ct3++;
501         }
502       }
503       for (k=oct2,max2=ct2; k<max2; k++)  {
504         row   = xdata[i][k] - rstart;
505         start = ai[row];
506         end   = ai[row+1];
507         for (l=start; l<end; l++) {
508           val = aj[l] + cstart;
509           if (!PetscBTLookupSet(xtable,val)) {
510             if (!(ct3 < mem_estimate)) {
511               new_estimate = (int)(1.5*mem_estimate)+1;
512               ierr = PetscMalloc(new_estimate * sizeof(int),&tmp);CHKERRQ(ierr);
513               ierr = PetscMemcpy(tmp,xdata[0],mem_estimate*sizeof(int));CHKERRQ(ierr);
514               ierr = PetscFree(xdata[0]);CHKERRQ(ierr);
515               xdata[0]     = tmp;
516               mem_estimate = new_estimate; ++no_malloc;
517               for (ctr=1; ctr<=i; ctr++) { xdata[ctr] = xdata[ctr-1] + isz1[ctr-1];}
518             }
519             xdata[i][ct2++] = val;
520             ct3++;
521           }
522         }
523         start = bi[row];
524         end   = bi[row+1];
525         for (l=start; l<end; l++) {
526           val = garray[bj[l]];
527           if (!PetscBTLookupSet(xtable,val)) {
528             if (!(ct3 < mem_estimate)) {
529               new_estimate = (int)(1.5*mem_estimate)+1;
530               ierr = PetscMalloc(new_estimate * sizeof(int),&tmp);CHKERRQ(ierr);
531               ierr = PetscMemcpy(tmp,xdata[0],mem_estimate*sizeof(int));CHKERRQ(ierr);
532               ierr = PetscFree(xdata[0]);CHKERRQ(ierr);
533               xdata[0]     = tmp;
534               mem_estimate = new_estimate; ++no_malloc;
535               for (ctr =1; ctr <=i; ctr++) { xdata[ctr] = xdata[ctr-1] + isz1[ctr-1];}
536             }
537             xdata[i][ct2++] = val;
538             ct3++;
539           }
540         }
541       }
542       /* Update the header*/
543       xdata[i][2*j]   = ct2 - oct2; /* Undo the vector isz1 and use only a var*/
544       xdata[i][2*j-1] = rbuf_i[2*j-1];
545     }
546     xdata[i][0] = rbuf_0;
547     xdata[i+1]  = xdata[i] + ct2;
548     isz1[i]     = ct2; /* size of each message */
549   }
550   ierr = PetscBTDestroy(xtable);CHKERRQ(ierr);
551   PetscLogInfo(0,"MatIncreaseOverlap_MPIBAIJ:[%d] Allocated %d bytes, required %d, no of mallocs = %d\n",rank,mem_estimate,ct3,no_malloc);
552   PetscFunctionReturn(0);
553 }
554 
555 static PetscErrorCode MatGetSubMatrices_MPIBAIJ_local(Mat,int,const IS[],const IS[],MatReuse,Mat *);
556 
557 #undef __FUNCT__
558 #define __FUNCT__ "MatGetSubMatrices_MPIBAIJ"
559 PetscErrorCode MatGetSubMatrices_MPIBAIJ(Mat C,int ismax,const IS isrow[],const IS iscol[],MatReuse scall,Mat *submat[])
560 {
561   IS          *isrow_new,*iscol_new;
562   Mat_MPIBAIJ *c = (Mat_MPIBAIJ*)C->data;
563   PetscErrorCode ierr;
564   int         nmax,nstages_local,nstages,i,pos,max_no,N=C->N,bs=c->bs;
565 
566   PetscFunctionBegin;
567   /* The compression and expansion should be avoided. Does'nt point
568      out errors might change the indices hence buggey */
569 
570   ierr = PetscMalloc(2*(ismax+1)*sizeof(IS),&isrow_new);CHKERRQ(ierr);
571   iscol_new = isrow_new + ismax;
572   ierr = ISCompressIndicesSorted(N,bs,ismax,isrow,isrow_new);CHKERRQ(ierr);
573   ierr = ISCompressIndicesSorted(N,bs,ismax,iscol,iscol_new);CHKERRQ(ierr);
574 
575   /* Allocate memory to hold all the submatrices */
576   if (scall != MAT_REUSE_MATRIX) {
577     ierr = PetscMalloc((ismax+1)*sizeof(Mat),submat);CHKERRQ(ierr);
578   }
579   /* Determine the number of stages through which submatrices are done */
580   nmax          = 20*1000000 / (c->Nbs * sizeof(int));
581   if (!nmax) nmax = 1;
582   nstages_local = ismax/nmax + ((ismax % nmax)?1:0);
583 
584   /* Make sure every processor loops through the nstages */
585   ierr = MPI_Allreduce(&nstages_local,&nstages,1,MPI_INT,MPI_MAX,C->comm);CHKERRQ(ierr);
586   for (i=0,pos=0; i<nstages; i++) {
587     if (pos+nmax <= ismax) max_no = nmax;
588     else if (pos == ismax) max_no = 0;
589     else                   max_no = ismax-pos;
590     ierr = MatGetSubMatrices_MPIBAIJ_local(C,max_no,isrow_new+pos,iscol_new+pos,scall,*submat+pos);CHKERRQ(ierr);
591     pos += max_no;
592   }
593 
594   for (i=0; i<ismax; i++) {
595     ierr = ISDestroy(isrow_new[i]);CHKERRQ(ierr);
596     ierr = ISDestroy(iscol_new[i]);CHKERRQ(ierr);
597   }
598   ierr = PetscFree(isrow_new);CHKERRQ(ierr);
599   PetscFunctionReturn(0);
600 }
601 
602 #if defined (PETSC_USE_CTABLE)
603 #undef __FUNCT__
604 #define __FUNCT__ "PetscGetProc"
605 PetscErrorCode PetscGetProc(const int gid, const int size, const int proc_gnode[], int *proc)
606 {
607   int nGlobalNd = proc_gnode[size];
608   int fproc = (int) ((float)gid * (float)size / (float)nGlobalNd + 0.5);
609 
610   PetscFunctionBegin;
611   /* if(fproc < 0) SETERRQ(1,"fproc < 0");*/
612   if (fproc > size) fproc = size;
613   while (gid < proc_gnode[fproc] || gid >= proc_gnode[fproc+1]) {
614     if (gid < proc_gnode[fproc]) fproc--;
615     else                         fproc++;
616   }
617   /* if(fproc<0 || fproc>=size) { SETERRQ(1,"fproc < 0 || fproc >= size"); }*/
618   *proc = fproc;
619   PetscFunctionReturn(0);
620 }
621 #endif
622 
623 /* -------------------------------------------------------------------------*/
624 #undef __FUNCT__
625 #define __FUNCT__ "MatGetSubMatrices_MPIBAIJ_local"
626 static PetscErrorCode MatGetSubMatrices_MPIBAIJ_local(Mat C,int ismax,const IS isrow[],const IS iscol[],MatReuse scall,Mat *submats)
627 {
628   Mat_MPIBAIJ *c = (Mat_MPIBAIJ*)C->data;
629   Mat         A = c->A;
630   Mat_SeqBAIJ *a = (Mat_SeqBAIJ*)A->data,*b = (Mat_SeqBAIJ*)c->B->data,*mat;
631   int         **irow,**icol,*nrow,*ncol,*w1,*w2,*w3,*w4,start,end,size;
632   PetscErrorCode ierr;
633   int         **sbuf1,**sbuf2,rank,i,j,k,l,ct1,ct2,**rbuf1,row,proc;
634   int         nrqs,msz,**ptr,idex,*req_size,*ctr,*pa,*tmp,tcol,nrqr;
635   int         **rbuf3,*req_source,**sbuf_aj,**rbuf2,max1,max2;
636   int         **lens,is_no,ncols,*cols,mat_i,*mat_j,tmp2,jmax,*irow_i;
637   int         len,ctr_j,*sbuf1_j,*sbuf_aj_i,*rbuf1_i,kmax,*lens_i;
638   int         bs=c->bs,bs2=c->bs2,*a_j=a->j,*b_j=b->j,*cworkA,*cworkB;
639   int         cstart = c->cstart,nzA,nzB,*a_i=a->i,*b_i=b->i,imark;
640   int         *bmap = c->garray,ctmp,rstart=c->rstart,tag0,tag1,tag2,tag3;
641   MPI_Request *s_waits1,*r_waits1,*s_waits2,*r_waits2,*r_waits3;
642   MPI_Request *r_waits4,*s_waits3,*s_waits4;
643   MPI_Status  *r_status1,*r_status2,*s_status1,*s_status3,*s_status2;
644   MPI_Status  *r_status3,*r_status4,*s_status4;
645   MPI_Comm    comm;
646   MatScalar   **rbuf4,**sbuf_aa,*vals,*mat_a,*sbuf_aa_i,*vworkA,*vworkB;
647   MatScalar   *a_a=a->a,*b_a=b->a;
648   PetscTruth  flag;
649   int         *onodes1,*olengths1;
650 
651 #if defined (PETSC_USE_CTABLE)
652   int tt;
653   PetscTable  *rowmaps,*colmaps,lrow1_grow1,lcol1_gcol1;
654 #else
655   int         **cmap,*cmap_i,*rtable,*rmap_i,**rmap, Mbs = c->Mbs;
656 #endif
657 
658   PetscFunctionBegin;
659   comm   = C->comm;
660   tag0   = C->tag;
661   size   = c->size;
662   rank   = c->rank;
663 
664   /* Get some new tags to keep the communication clean */
665   ierr = PetscObjectGetNewTag((PetscObject)C,&tag1);CHKERRQ(ierr);
666   ierr = PetscObjectGetNewTag((PetscObject)C,&tag2);CHKERRQ(ierr);
667   ierr = PetscObjectGetNewTag((PetscObject)C,&tag3);CHKERRQ(ierr);
668 
669   /* Check if the col indices are sorted */
670   for (i=0; i<ismax; i++) {
671     ierr = ISSorted(iscol[i],(PetscTruth*)&j);CHKERRQ(ierr);
672     if (!j) SETERRQ(PETSC_ERR_ARG_WRONGSTATE,"IS is not sorted");
673   }
674 
675   len    = (2*ismax+1)*(sizeof(int*)+ sizeof(int));
676 #if !defined (PETSC_USE_CTABLE)
677   len    += (Mbs+1)*sizeof(int);
678 #endif
679   ierr = PetscMalloc(len,&irow);CHKERRQ(ierr);
680   icol = irow + ismax;
681   nrow = (int*)(icol + ismax);
682   ncol = nrow + ismax;
683 #if !defined (PETSC_USE_CTABLE)
684   rtable = ncol + ismax;
685   /* Create hash table for the mapping :row -> proc*/
686   for (i=0,j=0; i<size; i++) {
687     jmax = c->rowners[i+1];
688     for (; j<jmax; j++) {
689       rtable[j] = i;
690     }
691   }
692 #endif
693 
694   for (i=0; i<ismax; i++) {
695     ierr = ISGetIndices(isrow[i],&irow[i]);CHKERRQ(ierr);
696     ierr = ISGetIndices(iscol[i],&icol[i]);CHKERRQ(ierr);
697     ierr = ISGetLocalSize(isrow[i],&nrow[i]);CHKERRQ(ierr);
698     ierr = ISGetLocalSize(iscol[i],&ncol[i]);CHKERRQ(ierr);
699   }
700 
701   /* evaluate communication - mesg to who,length of mesg,and buffer space
702      required. Based on this, buffers are allocated, and data copied into them*/
703   ierr = PetscMalloc(size*4*sizeof(int),&w1);CHKERRQ(ierr); /* mesg size */
704   w2   = w1 + size;      /* if w2[i] marked, then a message to proc i*/
705   w3   = w2 + size;      /* no of IS that needs to be sent to proc i */
706   w4   = w3 + size;      /* temp work space used in determining w1, w2, w3 */
707   ierr = PetscMemzero(w1,size*3*sizeof(int));CHKERRQ(ierr); /* initialise work vector*/
708   for (i=0; i<ismax; i++) {
709     ierr   = PetscMemzero(w4,size*sizeof(int));CHKERRQ(ierr); /* initialise work vector*/
710     jmax   = nrow[i];
711     irow_i = irow[i];
712     for (j=0; j<jmax; j++) {
713       row  = irow_i[j];
714 #if defined (PETSC_USE_CTABLE)
715       ierr = PetscGetProc(row,size,c->rowners,&proc);CHKERRQ(ierr);
716 #else
717       proc = rtable[row];
718 #endif
719       w4[proc]++;
720     }
721     for (j=0; j<size; j++) {
722       if (w4[j]) { w1[j] += w4[j];  w3[j]++;}
723     }
724   }
725 
726   nrqs     = 0;              /* no of outgoing messages */
727   msz      = 0;              /* total mesg length for all proc */
728   w1[rank] = 0;              /* no mesg sent to intself */
729   w3[rank] = 0;
730   for (i=0; i<size; i++) {
731     if (w1[i])  { w2[i] = 1; nrqs++;} /* there exists a message to proc i */
732   }
733   ierr = PetscMalloc((nrqs+1)*sizeof(int),&pa);CHKERRQ(ierr); /*(proc -array)*/
734   for (i=0,j=0; i<size; i++) {
735     if (w1[i]) { pa[j] = i; j++; }
736   }
737 
738   /* Each message would have a header = 1 + 2*(no of IS) + data */
739   for (i=0; i<nrqs; i++) {
740     j     = pa[i];
741     w1[j] += w2[j] + 2* w3[j];
742     msz   += w1[j];
743   }
744 
745   /* Determine the number of messages to expect, their lengths, from from-ids */
746   ierr = PetscGatherNumberOfMessages(comm,w2,w1,&nrqr);CHKERRQ(ierr);
747   ierr = PetscGatherMessageLengths(comm,nrqs,nrqr,w1,&onodes1,&olengths1);CHKERRQ(ierr);
748 
749   /* Now post the Irecvs corresponding to these messages */
750   ierr = PetscPostIrecvInt(comm,tag0,nrqr,onodes1,olengths1,&rbuf1,&r_waits1);CHKERRQ(ierr);
751 
752   ierr = PetscFree(onodes1);CHKERRQ(ierr);
753   ierr = PetscFree(olengths1);CHKERRQ(ierr);
754 
755   /* Allocate Memory for outgoing messages */
756   len  = 2*size*sizeof(int*) + 2*msz*sizeof(int) + size*sizeof(int);
757   ierr = PetscMalloc(len,&sbuf1);CHKERRQ(ierr);
758   ptr  = sbuf1 + size;   /* Pointers to the data in outgoing buffers */
759   ierr = PetscMemzero(sbuf1,2*size*sizeof(int*));CHKERRQ(ierr);
760   /* allocate memory for outgoing data + buf to receive the first reply */
761   tmp  = (int*)(ptr + size);
762   ctr  = tmp + 2*msz;
763 
764   {
765     int *iptr = tmp,ict = 0;
766     for (i=0; i<nrqs; i++) {
767       j         = pa[i];
768       iptr     += ict;
769       sbuf1[j]  = iptr;
770       ict       = w1[j];
771     }
772   }
773 
774   /* Form the outgoing messages */
775   /* Initialise the header space */
776   for (i=0; i<nrqs; i++) {
777     j           = pa[i];
778     sbuf1[j][0] = 0;
779     ierr        = PetscMemzero(sbuf1[j]+1,2*w3[j]*sizeof(int));CHKERRQ(ierr);
780     ptr[j]      = sbuf1[j] + 2*w3[j] + 1;
781   }
782 
783   /* Parse the isrow and copy data into outbuf */
784   for (i=0; i<ismax; i++) {
785     ierr   = PetscMemzero(ctr,size*sizeof(int));CHKERRQ(ierr);
786     irow_i = irow[i];
787     jmax   = nrow[i];
788     for (j=0; j<jmax; j++) {  /* parse the indices of each IS */
789       row  = irow_i[j];
790 #if defined (PETSC_USE_CTABLE)
791       ierr = PetscGetProc(row,size,c->rowners,&proc);CHKERRQ(ierr);
792 #else
793       proc = rtable[row];
794 #endif
795       if (proc != rank) { /* copy to the outgoing buf*/
796         ctr[proc]++;
797         *ptr[proc] = row;
798         ptr[proc]++;
799       }
800     }
801     /* Update the headers for the current IS */
802     for (j=0; j<size; j++) { /* Can Optimise this loop too */
803       if ((ctr_j = ctr[j])) {
804         sbuf1_j        = sbuf1[j];
805         k              = ++sbuf1_j[0];
806         sbuf1_j[2*k]   = ctr_j;
807         sbuf1_j[2*k-1] = i;
808       }
809     }
810   }
811 
812   /*  Now  post the sends */
813   ierr = PetscMalloc((nrqs+1)*sizeof(MPI_Request),&s_waits1);CHKERRQ(ierr);
814   for (i=0; i<nrqs; ++i) {
815     j = pa[i];
816     ierr = MPI_Isend(sbuf1[j],w1[j],MPI_INT,j,tag0,comm,s_waits1+i);CHKERRQ(ierr);
817   }
818 
819   /* Post Recieves to capture the buffer size */
820   ierr     = PetscMalloc((nrqs+1)*sizeof(MPI_Request),&r_waits2);CHKERRQ(ierr);
821   ierr     = PetscMalloc((nrqs+1)*sizeof(int*),&rbuf2);CHKERRQ(ierr);
822   rbuf2[0] = tmp + msz;
823   for (i=1; i<nrqs; ++i) {
824     j        = pa[i];
825     rbuf2[i] = rbuf2[i-1]+w1[pa[i-1]];
826   }
827   for (i=0; i<nrqs; ++i) {
828     j    = pa[i];
829     ierr = MPI_Irecv(rbuf2[i],w1[j],MPI_INT,j,tag1,comm,r_waits2+i);CHKERRQ(ierr);
830   }
831 
832   /* Send to other procs the buf size they should allocate */
833 
834   /* Receive messages*/
835   ierr       = PetscMalloc((nrqr+1)*sizeof(MPI_Request),&s_waits2);CHKERRQ(ierr);
836   ierr       = PetscMalloc((nrqr+1)*sizeof(MPI_Status),&r_status1);CHKERRQ(ierr);
837   len        = 2*nrqr*sizeof(int) + (nrqr+1)*sizeof(int*);
838   ierr       = PetscMalloc(len,&sbuf2);CHKERRQ(ierr);
839   req_size   = (int*)(sbuf2 + nrqr);
840   req_source = req_size + nrqr;
841 
842   {
843     Mat_SeqBAIJ *sA = (Mat_SeqBAIJ*)c->A->data,*sB = (Mat_SeqBAIJ*)c->B->data;
844     int        *sAi = sA->i,*sBi = sB->i,id,*sbuf2_i;
845 
846     for (i=0; i<nrqr; ++i) {
847       ierr = MPI_Waitany(nrqr,r_waits1,&idex,r_status1+i);CHKERRQ(ierr);
848       req_size[idex] = 0;
849       rbuf1_i         = rbuf1[idex];
850       start           = 2*rbuf1_i[0] + 1;
851       ierr            = MPI_Get_count(r_status1+i,MPI_INT,&end);CHKERRQ(ierr);
852       ierr            = PetscMalloc(end*sizeof(int),&sbuf2[idex]);CHKERRQ(ierr);
853       sbuf2_i         = sbuf2[idex];
854       for (j=start; j<end; j++) {
855         id               = rbuf1_i[j] - rstart;
856         ncols            = sAi[id+1] - sAi[id] + sBi[id+1] - sBi[id];
857         sbuf2_i[j]       = ncols;
858         req_size[idex] += ncols;
859       }
860       req_source[idex] = r_status1[i].MPI_SOURCE;
861       /* form the header */
862       sbuf2_i[0]   = req_size[idex];
863       for (j=1; j<start; j++) { sbuf2_i[j] = rbuf1_i[j]; }
864       ierr = MPI_Isend(sbuf2_i,end,MPI_INT,req_source[idex],tag1,comm,s_waits2+i);CHKERRQ(ierr);
865     }
866   }
867   ierr = PetscFree(r_status1);CHKERRQ(ierr);
868   ierr = PetscFree(r_waits1);CHKERRQ(ierr);
869 
870   /*  recv buffer sizes */
871   /* Receive messages*/
872 
873   ierr = PetscMalloc((nrqs+1)*sizeof(int*),&rbuf3);CHKERRQ(ierr);
874   ierr = PetscMalloc((nrqs+1)*sizeof(MatScalar*),&rbuf4);CHKERRQ(ierr);
875   ierr = PetscMalloc((nrqs+1)*sizeof(MPI_Request),&r_waits3);CHKERRQ(ierr);
876   ierr = PetscMalloc((nrqs+1)*sizeof(MPI_Request),&r_waits4);CHKERRQ(ierr);
877   ierr = PetscMalloc((nrqs+1)*sizeof(MPI_Status),&r_status2);CHKERRQ(ierr);
878 
879   for (i=0; i<nrqs; ++i) {
880     ierr = MPI_Waitany(nrqs,r_waits2,&idex,r_status2+i);CHKERRQ(ierr);
881     ierr = PetscMalloc(rbuf2[idex][0]*sizeof(int),&rbuf3[idex]);CHKERRQ(ierr);
882     ierr = PetscMalloc(rbuf2[idex][0]*bs2*sizeof(MatScalar),&rbuf4[idex]);CHKERRQ(ierr);
883     ierr = MPI_Irecv(rbuf3[idex],rbuf2[idex][0],MPI_INT,
884                      r_status2[i].MPI_SOURCE,tag2,comm,r_waits3+idex);CHKERRQ(ierr);
885     ierr = MPI_Irecv(rbuf4[idex],rbuf2[idex][0]*bs2,MPIU_MATSCALAR,
886                      r_status2[i].MPI_SOURCE,tag3,comm,r_waits4+idex);CHKERRQ(ierr);
887   }
888   ierr = PetscFree(r_status2);CHKERRQ(ierr);
889   ierr = PetscFree(r_waits2);CHKERRQ(ierr);
890 
891   /* Wait on sends1 and sends2 */
892   ierr = PetscMalloc((nrqs+1)*sizeof(MPI_Status),&s_status1);CHKERRQ(ierr);
893   ierr = PetscMalloc((nrqr+1)*sizeof(MPI_Status),&s_status2);CHKERRQ(ierr);
894 
895   ierr = MPI_Waitall(nrqs,s_waits1,s_status1);CHKERRQ(ierr);
896   ierr = MPI_Waitall(nrqr,s_waits2,s_status2);CHKERRQ(ierr);
897   ierr = PetscFree(s_status1);CHKERRQ(ierr);
898   ierr = PetscFree(s_status2);CHKERRQ(ierr);
899   ierr = PetscFree(s_waits1);CHKERRQ(ierr);
900   ierr = PetscFree(s_waits2);CHKERRQ(ierr);
901 
902   /* Now allocate buffers for a->j, and send them off */
903   ierr = PetscMalloc((nrqr+1)*sizeof(int*),&sbuf_aj);CHKERRQ(ierr);
904   for (i=0,j=0; i<nrqr; i++) j += req_size[i];
905   ierr = PetscMalloc((j+1)*sizeof(int),&sbuf_aj[0]);CHKERRQ(ierr);
906   for (i=1; i<nrqr; i++)  sbuf_aj[i] = sbuf_aj[i-1] + req_size[i-1];
907 
908   ierr = PetscMalloc((nrqr+1)*sizeof(MPI_Request),&s_waits3);CHKERRQ(ierr);
909   {
910      for (i=0; i<nrqr; i++) {
911       rbuf1_i   = rbuf1[i];
912       sbuf_aj_i = sbuf_aj[i];
913       ct1       = 2*rbuf1_i[0] + 1;
914       ct2       = 0;
915       for (j=1,max1=rbuf1_i[0]; j<=max1; j++) {
916         kmax = rbuf1[i][2*j];
917         for (k=0; k<kmax; k++,ct1++) {
918           row    = rbuf1_i[ct1] - rstart;
919           nzA    = a_i[row+1] - a_i[row];     nzB = b_i[row+1] - b_i[row];
920           ncols  = nzA + nzB;
921           cworkA = a_j + a_i[row]; cworkB = b_j + b_i[row];
922 
923           /* load the column indices for this row into cols*/
924           cols  = sbuf_aj_i + ct2;
925           for (l=0; l<nzB; l++) {
926             if ((ctmp = bmap[cworkB[l]]) < cstart)  cols[l] = ctmp;
927             else break;
928           }
929           imark = l;
930           for (l=0; l<nzA; l++)   cols[imark+l] = cstart + cworkA[l];
931           for (l=imark; l<nzB; l++) cols[nzA+l] = bmap[cworkB[l]];
932           ct2 += ncols;
933         }
934       }
935       ierr = MPI_Isend(sbuf_aj_i,req_size[i],MPI_INT,req_source[i],tag2,comm,s_waits3+i);CHKERRQ(ierr);
936     }
937   }
938   ierr = PetscMalloc((nrqs+1)*sizeof(MPI_Status),&r_status3);CHKERRQ(ierr);
939   ierr = PetscMalloc((nrqr+1)*sizeof(MPI_Status),&s_status3);CHKERRQ(ierr);
940 
941   /* Allocate buffers for a->a, and send them off */
942   ierr = PetscMalloc((nrqr+1)*sizeof(MatScalar *),&sbuf_aa);CHKERRQ(ierr);
943   for (i=0,j=0; i<nrqr; i++) j += req_size[i];
944   ierr = PetscMalloc((j+1)*bs2*sizeof(MatScalar),&sbuf_aa[0]);CHKERRQ(ierr);
945   for (i=1; i<nrqr; i++)  sbuf_aa[i] = sbuf_aa[i-1] + req_size[i-1]*bs2;
946 
947   ierr = PetscMalloc((nrqr+1)*sizeof(MPI_Request),&s_waits4);CHKERRQ(ierr);
948   {
949     for (i=0; i<nrqr; i++) {
950       rbuf1_i   = rbuf1[i];
951       sbuf_aa_i = sbuf_aa[i];
952       ct1       = 2*rbuf1_i[0]+1;
953       ct2       = 0;
954       for (j=1,max1=rbuf1_i[0]; j<=max1; j++) {
955         kmax = rbuf1_i[2*j];
956         for (k=0; k<kmax; k++,ct1++) {
957           row    = rbuf1_i[ct1] - rstart;
958           nzA    = a_i[row+1] - a_i[row];     nzB = b_i[row+1] - b_i[row];
959           ncols  = nzA + nzB;
960           cworkA = a_j + a_i[row];     cworkB = b_j + b_i[row];
961           vworkA = a_a + a_i[row]*bs2; vworkB = b_a + b_i[row]*bs2;
962 
963           /* load the column values for this row into vals*/
964           vals  = sbuf_aa_i+ct2*bs2;
965           for (l=0; l<nzB; l++) {
966             if ((bmap[cworkB[l]]) < cstart) {
967               ierr = PetscMemcpy(vals+l*bs2,vworkB+l*bs2,bs2*sizeof(MatScalar));CHKERRQ(ierr);
968             }
969             else break;
970           }
971           imark = l;
972           for (l=0; l<nzA; l++) {
973             ierr = PetscMemcpy(vals+(imark+l)*bs2,vworkA+l*bs2,bs2*sizeof(MatScalar));CHKERRQ(ierr);
974           }
975           for (l=imark; l<nzB; l++) {
976             ierr = PetscMemcpy(vals+(nzA+l)*bs2,vworkB+l*bs2,bs2*sizeof(MatScalar));CHKERRQ(ierr);
977           }
978           ct2 += ncols;
979         }
980       }
981       ierr = MPI_Isend(sbuf_aa_i,req_size[i]*bs2,MPIU_MATSCALAR,req_source[i],tag3,comm,s_waits4+i);CHKERRQ(ierr);
982     }
983   }
984   ierr = PetscMalloc((nrqs+1)*sizeof(MPI_Status),&r_status4);CHKERRQ(ierr);
985   ierr = PetscMalloc((nrqr+1)*sizeof(MPI_Status),&s_status4);CHKERRQ(ierr);
986   ierr = PetscFree(rbuf1);CHKERRQ(ierr);
987 
988   /* Form the matrix */
989   /* create col map */
990   {
991     int *icol_i;
992 #if defined (PETSC_USE_CTABLE)
993     /* Create row map*/
994     ierr = PetscMalloc((1+ismax)*sizeof(PetscTable),&colmaps);CHKERRQ(ierr);
995     for (i=0; i<ismax; i++) {
996       ierr = PetscTableCreate(ncol[i]+1,&colmaps[i]);CHKERRQ(ierr);
997     }
998 #else
999     len     = (1+ismax)*sizeof(int*)+ ismax*c->Nbs*sizeof(int);
1000     ierr    = PetscMalloc(len,&cmap);CHKERRQ(ierr);
1001     cmap[0] = (int*)(cmap + ismax);
1002     ierr    = PetscMemzero(cmap[0],(1+ismax*c->Nbs)*sizeof(int));CHKERRQ(ierr);
1003     for (i=1; i<ismax; i++) { cmap[i] = cmap[i-1] + c->Nbs; }
1004 #endif
1005     for (i=0; i<ismax; i++) {
1006       jmax   = ncol[i];
1007       icol_i = icol[i];
1008 #if defined (PETSC_USE_CTABLE)
1009       lcol1_gcol1 = colmaps[i];
1010       for (j=0; j<jmax; j++) {
1011 	ierr = PetscTableAdd(lcol1_gcol1,icol_i[j]+1,j+1);CHKERRQ(ierr);
1012       }
1013 #else
1014       cmap_i = cmap[i];
1015       for (j=0; j<jmax; j++) {
1016         cmap_i[icol_i[j]] = j+1;
1017       }
1018 #endif
1019     }
1020   }
1021 
1022   /* Create lens which is required for MatCreate... */
1023   for (i=0,j=0; i<ismax; i++) { j += nrow[i]; }
1024   len     = (1+ismax)*sizeof(int*)+ j*sizeof(int);
1025   ierr    = PetscMalloc(len,&lens);CHKERRQ(ierr);
1026   lens[0] = (int*)(lens + ismax);
1027   ierr    = PetscMemzero(lens[0],j*sizeof(int));CHKERRQ(ierr);
1028   for (i=1; i<ismax; i++) { lens[i] = lens[i-1] + nrow[i-1]; }
1029 
1030   /* Update lens from local data */
1031   for (i=0; i<ismax; i++) {
1032     jmax   = nrow[i];
1033 #if defined (PETSC_USE_CTABLE)
1034     lcol1_gcol1 = colmaps[i];
1035 #else
1036     cmap_i = cmap[i];
1037 #endif
1038     irow_i = irow[i];
1039     lens_i = lens[i];
1040     for (j=0; j<jmax; j++) {
1041       row  = irow_i[j];
1042 #if defined (PETSC_USE_CTABLE)
1043       ierr = PetscGetProc(row,size,c->rowners,&proc);CHKERRQ(ierr);
1044 #else
1045       proc = rtable[row];
1046 #endif
1047       if (proc == rank) {
1048         /* Get indices from matA and then from matB */
1049         row    = row - rstart;
1050         nzA    = a_i[row+1] - a_i[row];     nzB = b_i[row+1] - b_i[row];
1051         cworkA =  a_j + a_i[row]; cworkB = b_j + b_i[row];
1052 #if defined (PETSC_USE_CTABLE)
1053        for (k=0; k<nzA; k++) {
1054 	 ierr = PetscTableFind(lcol1_gcol1,cstart+cworkA[k]+1,&tt);CHKERRQ(ierr);
1055           if (tt) { lens_i[j]++; }
1056         }
1057         for (k=0; k<nzB; k++) {
1058 	  ierr = PetscTableFind(lcol1_gcol1,bmap[cworkB[k]]+1,&tt);CHKERRQ(ierr);
1059           if (tt) { lens_i[j]++; }
1060         }
1061 #else
1062         for (k=0; k<nzA; k++) {
1063           if (cmap_i[cstart + cworkA[k]]) { lens_i[j]++; }
1064         }
1065         for (k=0; k<nzB; k++) {
1066           if (cmap_i[bmap[cworkB[k]]]) { lens_i[j]++; }
1067         }
1068 #endif
1069       }
1070     }
1071   }
1072 #if defined (PETSC_USE_CTABLE)
1073   /* Create row map*/
1074   ierr = PetscMalloc((1+ismax)*sizeof(PetscTable),&rowmaps);CHKERRQ(ierr);
1075   for (i=0; i<ismax; i++){
1076     ierr = PetscTableCreate(nrow[i]+1,&rowmaps[i]);CHKERRQ(ierr);
1077   }
1078 #else
1079   /* Create row map*/
1080   len     = (1+ismax)*sizeof(int*)+ ismax*Mbs*sizeof(int);
1081   ierr    = PetscMalloc(len,&rmap);CHKERRQ(ierr);
1082   rmap[0] = (int*)(rmap + ismax);
1083   ierr    = PetscMemzero(rmap[0],ismax*Mbs*sizeof(int));CHKERRQ(ierr);
1084   for (i=1; i<ismax; i++) { rmap[i] = rmap[i-1] + Mbs;}
1085 #endif
1086   for (i=0; i<ismax; i++) {
1087     irow_i = irow[i];
1088     jmax   = nrow[i];
1089 #if defined (PETSC_USE_CTABLE)
1090     lrow1_grow1 = rowmaps[i];
1091     for (j=0; j<jmax; j++) {
1092       ierr = PetscTableAdd(lrow1_grow1,irow_i[j]+1,j+1);CHKERRQ(ierr);
1093     }
1094 #else
1095     rmap_i = rmap[i];
1096     for (j=0; j<jmax; j++) {
1097       rmap_i[irow_i[j]] = j;
1098     }
1099 #endif
1100   }
1101 
1102   /* Update lens from offproc data */
1103   {
1104     int *rbuf2_i,*rbuf3_i,*sbuf1_i;
1105 
1106     for (tmp2=0; tmp2<nrqs; tmp2++) {
1107       ierr    = MPI_Waitany(nrqs,r_waits3,&i,r_status3+tmp2);CHKERRQ(ierr);
1108       idex   = pa[i];
1109       sbuf1_i = sbuf1[idex];
1110       jmax    = sbuf1_i[0];
1111       ct1     = 2*jmax+1;
1112       ct2     = 0;
1113       rbuf2_i = rbuf2[i];
1114       rbuf3_i = rbuf3[i];
1115       for (j=1; j<=jmax; j++) {
1116         is_no   = sbuf1_i[2*j-1];
1117         max1    = sbuf1_i[2*j];
1118         lens_i  = lens[is_no];
1119 #if defined (PETSC_USE_CTABLE)
1120 	lcol1_gcol1 = colmaps[is_no];
1121 	lrow1_grow1 = rowmaps[is_no];
1122 #else
1123         cmap_i  = cmap[is_no];
1124 	rmap_i  = rmap[is_no];
1125 #endif
1126         for (k=0; k<max1; k++,ct1++) {
1127 #if defined (PETSC_USE_CTABLE)
1128 	  ierr = PetscTableFind(lrow1_grow1,sbuf1_i[ct1]+1,&row);CHKERRQ(ierr);
1129           row--;
1130           if(row < 0) { SETERRQ(1,"row not found in table"); }
1131 #else
1132           row  = rmap_i[sbuf1_i[ct1]]; /* the val in the new matrix to be */
1133 #endif
1134           max2 = rbuf2_i[ct1];
1135           for (l=0; l<max2; l++,ct2++) {
1136 #if defined (PETSC_USE_CTABLE)
1137 	    ierr = PetscTableFind(lcol1_gcol1,rbuf3_i[ct2]+1,&tt);CHKERRQ(ierr);
1138 	    if (tt) {
1139               lens_i[row]++;
1140             }
1141 #else
1142             if (cmap_i[rbuf3_i[ct2]]) {
1143               lens_i[row]++;
1144             }
1145 #endif
1146           }
1147         }
1148       }
1149     }
1150   }
1151   ierr = PetscFree(r_status3);CHKERRQ(ierr);
1152   ierr = PetscFree(r_waits3);CHKERRQ(ierr);
1153   ierr = MPI_Waitall(nrqr,s_waits3,s_status3);CHKERRQ(ierr);
1154   ierr = PetscFree(s_status3);CHKERRQ(ierr);
1155   ierr = PetscFree(s_waits3);CHKERRQ(ierr);
1156 
1157   /* Create the submatrices */
1158   if (scall == MAT_REUSE_MATRIX) {
1159     /*
1160         Assumes new rows are same length as the old rows, hence bug!
1161     */
1162     for (i=0; i<ismax; i++) {
1163       mat = (Mat_SeqBAIJ *)(submats[i]->data);
1164       if ((mat->mbs != nrow[i]) || (mat->nbs != ncol[i] || mat->bs != bs)) {
1165         SETERRQ(PETSC_ERR_ARG_SIZ,"Cannot reuse matrix. wrong size");
1166       }
1167       ierr = PetscMemcmp(mat->ilen,lens[i],mat->mbs *sizeof(int),&flag);CHKERRQ(ierr);
1168       if (flag == PETSC_FALSE) {
1169         SETERRQ(PETSC_ERR_ARG_INCOMP,"Cannot reuse matrix. wrong no of nonzeros");
1170       }
1171       /* Initial matrix as if empty */
1172       ierr = PetscMemzero(mat->ilen,mat->mbs*sizeof(int));CHKERRQ(ierr);
1173       submats[i]->factor = C->factor;
1174     }
1175   } else {
1176     for (i=0; i<ismax; i++) {
1177       ierr = MatCreate(PETSC_COMM_SELF,nrow[i]*bs,ncol[i]*bs,nrow[i]*bs,ncol[i]*bs,submats+i);CHKERRQ(ierr);
1178       ierr = MatSetType(submats[i],A->type_name);CHKERRQ(ierr);
1179       ierr = MatSeqBAIJSetPreallocation(submats[i],a->bs,0,lens[i]);CHKERRQ(ierr);
1180       ierr = MatSeqSBAIJSetPreallocation(submats[i],a->bs,0,lens[i]);CHKERRQ(ierr);
1181     }
1182   }
1183 
1184   /* Assemble the matrices */
1185   /* First assemble the local rows */
1186   {
1187     int       ilen_row,*imat_ilen,*imat_j,*imat_i;
1188     MatScalar *imat_a;
1189 
1190     for (i=0; i<ismax; i++) {
1191       mat       = (Mat_SeqBAIJ*)submats[i]->data;
1192       imat_ilen = mat->ilen;
1193       imat_j    = mat->j;
1194       imat_i    = mat->i;
1195       imat_a    = mat->a;
1196 
1197 #if defined (PETSC_USE_CTABLE)
1198       lcol1_gcol1 = colmaps[i];
1199       lrow1_grow1 = rowmaps[i];
1200 #else
1201       cmap_i    = cmap[i];
1202       rmap_i    = rmap[i];
1203 #endif
1204       irow_i    = irow[i];
1205       jmax      = nrow[i];
1206       for (j=0; j<jmax; j++) {
1207         row      = irow_i[j];
1208 #if defined (PETSC_USE_CTABLE)
1209 	ierr = PetscGetProc(row,size,c->rowners,&proc);CHKERRQ(ierr);
1210 #else
1211 	proc = rtable[row];
1212 #endif
1213         if (proc == rank) {
1214           row      = row - rstart;
1215           nzA      = a_i[row+1] - a_i[row];
1216           nzB      = b_i[row+1] - b_i[row];
1217           cworkA   = a_j + a_i[row];
1218           cworkB   = b_j + b_i[row];
1219           vworkA   = a_a + a_i[row]*bs2;
1220           vworkB   = b_a + b_i[row]*bs2;
1221 #if defined (PETSC_USE_CTABLE)
1222 	  ierr = PetscTableFind(lrow1_grow1,row+rstart+1,&row);CHKERRQ(ierr);
1223           row--;
1224           if (row < 0) { SETERRQ(1,"row not found in table"); }
1225 #else
1226           row      = rmap_i[row + rstart];
1227 #endif
1228           mat_i    = imat_i[row];
1229           mat_a    = imat_a + mat_i*bs2;
1230           mat_j    = imat_j + mat_i;
1231           ilen_row = imat_ilen[row];
1232 
1233           /* load the column indices for this row into cols*/
1234           for (l=0; l<nzB; l++) {
1235 	    if ((ctmp = bmap[cworkB[l]]) < cstart) {
1236 #if defined (PETSC_USE_CTABLE)
1237 	      ierr = PetscTableFind(lcol1_gcol1,ctmp+1,&tcol);CHKERRQ(ierr);
1238 	      if (tcol) {
1239 #else
1240               if ((tcol = cmap_i[ctmp])) {
1241 #endif
1242                 *mat_j++ = tcol - 1;
1243                 ierr     = PetscMemcpy(mat_a,vworkB+l*bs2,bs2*sizeof(MatScalar));CHKERRQ(ierr);
1244                 mat_a   += bs2;
1245                 ilen_row++;
1246               }
1247             } else break;
1248           }
1249           imark = l;
1250           for (l=0; l<nzA; l++) {
1251 #if defined (PETSC_USE_CTABLE)
1252 	    ierr = PetscTableFind(lcol1_gcol1,cstart+cworkA[l]+1,&tcol);CHKERRQ(ierr);
1253 	    if (tcol) {
1254 #else
1255 	    if ((tcol = cmap_i[cstart + cworkA[l]])) {
1256 #endif
1257               *mat_j++ = tcol - 1;
1258               ierr     = PetscMemcpy(mat_a,vworkA+l*bs2,bs2*sizeof(MatScalar));CHKERRQ(ierr);
1259               mat_a   += bs2;
1260               ilen_row++;
1261             }
1262           }
1263           for (l=imark; l<nzB; l++) {
1264 #if defined (PETSC_USE_CTABLE)
1265 	    ierr = PetscTableFind(lcol1_gcol1,bmap[cworkB[l]]+1,&tcol);CHKERRQ(ierr);
1266 	    if (tcol) {
1267 #else
1268             if ((tcol = cmap_i[bmap[cworkB[l]]])) {
1269 #endif
1270               *mat_j++ = tcol - 1;
1271               ierr     = PetscMemcpy(mat_a,vworkB+l*bs2,bs2*sizeof(MatScalar));CHKERRQ(ierr);
1272               mat_a   += bs2;
1273               ilen_row++;
1274             }
1275           }
1276           imat_ilen[row] = ilen_row;
1277         }
1278       }
1279 
1280     }
1281   }
1282 
1283   /*   Now assemble the off proc rows*/
1284   {
1285     int       *sbuf1_i,*rbuf2_i,*rbuf3_i,*imat_ilen,ilen;
1286     int       *imat_j,*imat_i;
1287     MatScalar *imat_a,*rbuf4_i;
1288 
1289     for (tmp2=0; tmp2<nrqs; tmp2++) {
1290       ierr    = MPI_Waitany(nrqs,r_waits4,&i,r_status4+tmp2);CHKERRQ(ierr);
1291       idex   = pa[i];
1292       sbuf1_i = sbuf1[idex];
1293       jmax    = sbuf1_i[0];
1294       ct1     = 2*jmax + 1;
1295       ct2     = 0;
1296       rbuf2_i = rbuf2[i];
1297       rbuf3_i = rbuf3[i];
1298       rbuf4_i = rbuf4[i];
1299       for (j=1; j<=jmax; j++) {
1300         is_no     = sbuf1_i[2*j-1];
1301 #if defined (PETSC_USE_CTABLE)
1302 	lrow1_grow1 = rowmaps[is_no];
1303 	lcol1_gcol1 = colmaps[is_no];
1304 #else
1305         rmap_i    = rmap[is_no];
1306         cmap_i    = cmap[is_no];
1307 #endif
1308         mat       = (Mat_SeqBAIJ*)submats[is_no]->data;
1309         imat_ilen = mat->ilen;
1310         imat_j    = mat->j;
1311         imat_i    = mat->i;
1312         imat_a    = mat->a;
1313         max1      = sbuf1_i[2*j];
1314         for (k=0; k<max1; k++,ct1++) {
1315           row   = sbuf1_i[ct1];
1316 #if defined (PETSC_USE_CTABLE)
1317 	  ierr = PetscTableFind(lrow1_grow1,row+1,&row);CHKERRQ(ierr);
1318           row--;
1319           if(row < 0) { SETERRQ(1,"row not found in table"); }
1320 #else
1321           row   = rmap_i[row];
1322 #endif
1323           ilen  = imat_ilen[row];
1324           mat_i = imat_i[row];
1325           mat_a = imat_a + mat_i*bs2;
1326           mat_j = imat_j + mat_i;
1327           max2 = rbuf2_i[ct1];
1328           for (l=0; l<max2; l++,ct2++) {
1329 #if defined (PETSC_USE_CTABLE)
1330 	    ierr = PetscTableFind(lcol1_gcol1,rbuf3_i[ct2]+1,&tcol);CHKERRQ(ierr);
1331 	    if (tcol) {
1332 #else
1333 	    if ((tcol = cmap_i[rbuf3_i[ct2]])) {
1334 #endif
1335               *mat_j++    = tcol - 1;
1336               /* *mat_a++= rbuf4_i[ct2]; */
1337               ierr        = PetscMemcpy(mat_a,rbuf4_i+ct2*bs2,bs2*sizeof(MatScalar));CHKERRQ(ierr);
1338               mat_a      += bs2;
1339               ilen++;
1340             }
1341           }
1342           imat_ilen[row] = ilen;
1343         }
1344       }
1345     }
1346   }
1347   ierr = PetscFree(r_status4);CHKERRQ(ierr);
1348   ierr = PetscFree(r_waits4);CHKERRQ(ierr);
1349   ierr = MPI_Waitall(nrqr,s_waits4,s_status4);CHKERRQ(ierr);
1350   ierr = PetscFree(s_waits4);CHKERRQ(ierr);
1351   ierr = PetscFree(s_status4);CHKERRQ(ierr);
1352 
1353   /* Restore the indices */
1354   for (i=0; i<ismax; i++) {
1355     ierr = ISRestoreIndices(isrow[i],irow+i);CHKERRQ(ierr);
1356     ierr = ISRestoreIndices(iscol[i],icol+i);CHKERRQ(ierr);
1357   }
1358 
1359   /* Destroy allocated memory */
1360   ierr = PetscFree(irow);CHKERRQ(ierr);
1361   ierr = PetscFree(w1);CHKERRQ(ierr);
1362   ierr = PetscFree(pa);CHKERRQ(ierr);
1363 
1364   ierr = PetscFree(sbuf1);CHKERRQ(ierr);
1365   ierr = PetscFree(rbuf2);CHKERRQ(ierr);
1366   for (i=0; i<nrqr; ++i) {
1367     ierr = PetscFree(sbuf2[i]);CHKERRQ(ierr);
1368   }
1369   for (i=0; i<nrqs; ++i) {
1370     ierr = PetscFree(rbuf3[i]);CHKERRQ(ierr);
1371     ierr = PetscFree(rbuf4[i]);CHKERRQ(ierr);
1372   }
1373 
1374   ierr = PetscFree(sbuf2);CHKERRQ(ierr);
1375   ierr = PetscFree(rbuf3);CHKERRQ(ierr);
1376   ierr = PetscFree(rbuf4);CHKERRQ(ierr);
1377   ierr = PetscFree(sbuf_aj[0]);CHKERRQ(ierr);
1378   ierr = PetscFree(sbuf_aj);CHKERRQ(ierr);
1379   ierr = PetscFree(sbuf_aa[0]);CHKERRQ(ierr);
1380   ierr = PetscFree(sbuf_aa);CHKERRQ(ierr);
1381 
1382 #if defined (PETSC_USE_CTABLE)
1383   for (i=0; i<ismax; i++){
1384     ierr = PetscTableDelete(rowmaps[i]);CHKERRQ(ierr);
1385     ierr = PetscTableDelete(colmaps[i]);CHKERRQ(ierr);
1386   }
1387   ierr = PetscFree(colmaps);CHKERRQ(ierr);
1388   ierr = PetscFree(rowmaps);CHKERRQ(ierr);
1389 #else
1390   ierr = PetscFree(rmap);CHKERRQ(ierr);
1391   ierr = PetscFree(cmap);CHKERRQ(ierr);
1392 #endif
1393   ierr = PetscFree(lens);CHKERRQ(ierr);
1394 
1395   for (i=0; i<ismax; i++) {
1396     ierr = MatAssemblyBegin(submats[i],MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
1397     ierr = MatAssemblyEnd(submats[i],MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
1398   }
1399 
1400   PetscFunctionReturn(0);
1401 }
1402 
1403