1632d0f97SHong Zhang 2632d0f97SHong Zhang /* 3632d0f97SHong Zhang Routines to compute overlapping regions of a parallel MPI matrix. 4632d0f97SHong Zhang Used for finding submatrices that were shared across processors. 5632d0f97SHong Zhang */ 6c6db04a5SJed Brown #include <../src/mat/impls/sbaij/mpi/mpisbaij.h> 7c6db04a5SJed Brown #include <petscbt.h> 8632d0f97SHong Zhang 913f74950SBarry Smith static PetscErrorCode MatIncreaseOverlap_MPISBAIJ_Once(Mat, PetscInt, IS *); 1013f74950SBarry Smith static PetscErrorCode MatIncreaseOverlap_MPISBAIJ_Local(Mat, PetscInt *, PetscInt, PetscInt *, PetscBT *); 11632d0f97SHong Zhang 12d71ae5a4SJacob Faibussowitsch PetscErrorCode MatIncreaseOverlap_MPISBAIJ(Mat C, PetscInt is_max, IS is[], PetscInt ov) 13d71ae5a4SJacob Faibussowitsch { 14b11de2a9SHong Zhang PetscInt i, N = C->cmap->N, bs = C->rmap->bs, M = C->rmap->N, Mbs = M / bs, *nidx, isz, iov; 15b11de2a9SHong Zhang IS *is_new, *is_row; 16b11de2a9SHong Zhang Mat *submats; 17b11de2a9SHong Zhang Mat_MPISBAIJ *c = (Mat_MPISBAIJ *)C->data; 18b11de2a9SHong Zhang Mat_SeqSBAIJ *asub_i; 19b11de2a9SHong Zhang PetscBT table; 20b11de2a9SHong Zhang PetscInt *ai, brow, nz, nis, l, nmax, nstages_local, nstages, max_no, pos; 21b11de2a9SHong Zhang const PetscInt *idx; 22b829d874SHong Zhang PetscBool flg; 23632d0f97SHong Zhang 24632d0f97SHong Zhang PetscFunctionBegin; 259566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(is_max, &is_new)); 26632d0f97SHong Zhang /* Convert the indices into block format */ 279566063dSJacob Faibussowitsch PetscCall(ISCompressIndicesGeneral(N, C->rmap->n, bs, is_max, is, is_new)); 2808401ef6SPierre Jolivet PetscCheck(ov >= 0, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Negative overlap specified"); 29db41cccfSHong Zhang 30*2ef1f0ffSBarry Smith /* previous non-scalable implementation */ 31a0769a71SSatish Balay flg = PETSC_FALSE; 329566063dSJacob Faibussowitsch PetscCall(PetscOptionsHasName(NULL, NULL, "-IncreaseOverlap_old", &flg)); 337a868f3eSHong Zhang if (flg) { /* previous non-scalable implementation */ 347a868f3eSHong Zhang printf("use previous non-scalable implementation...\n"); 3548a46eb9SPierre Jolivet for (i = 0; i < ov; ++i) PetscCall(MatIncreaseOverlap_MPISBAIJ_Once(C, is_max, is_new)); 361c4982ebSHong Zhang } else { /* implementation using modified BAIJ routines */ 371c4982ebSHong Zhang 389566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(Mbs + 1, &nidx)); 399566063dSJacob Faibussowitsch PetscCall(PetscBTCreate(Mbs, &table)); /* for column search */ 40db41cccfSHong Zhang 41db41cccfSHong Zhang /* Create is_row */ 429566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(is_max, &is_row)); 439566063dSJacob Faibussowitsch PetscCall(ISCreateStride(PETSC_COMM_SELF, Mbs, 0, 1, &is_row[0])); 4426fbe8dcSKarl Rupp 459371c9d4SSatish Balay for (i = 1; i < is_max; i++) { is_row[i] = is_row[0]; /* reuse is_row[0] */ } 46db41cccfSHong Zhang 477dae84e0SHong Zhang /* Allocate memory to hold all the submatrices - Modified from MatCreateSubMatrices_MPIBAIJ() */ 489566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(is_max + 1, &submats)); 49b11de2a9SHong Zhang 50db41cccfSHong Zhang /* Determine the number of stages through which submatrices are done */ 51db41cccfSHong Zhang nmax = 20 * 1000000 / (c->Nbs * sizeof(PetscInt)); 52db41cccfSHong Zhang if (!nmax) nmax = 1; 53db41cccfSHong Zhang nstages_local = is_max / nmax + ((is_max % nmax) ? 1 : 0); 54db41cccfSHong Zhang 55db41cccfSHong Zhang /* Make sure every processor loops through the nstages */ 561c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&nstages_local, &nstages, 1, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)C))); 57b11de2a9SHong Zhang 58b11de2a9SHong Zhang for (iov = 0; iov < ov; ++iov) { 59b11de2a9SHong Zhang /* 1) Get submats for column search */ 60db41cccfSHong Zhang for (i = 0, pos = 0; i < nstages; i++) { 61db41cccfSHong Zhang if (pos + nmax <= is_max) max_no = nmax; 62db41cccfSHong Zhang else if (pos == is_max) max_no = 0; 63db41cccfSHong Zhang else max_no = is_max - pos; 64b829d874SHong Zhang c->ijonly = PETSC_TRUE; /* only matrix data structures are requested */ 65476417e5SBarry Smith /* The resulting submatrices should be BAIJ, not SBAIJ, hence we change this value to trigger that */ 669566063dSJacob Faibussowitsch PetscCall(PetscStrcpy(((PetscObject)c->A)->type_name, MATSEQBAIJ)); 679566063dSJacob Faibussowitsch PetscCall(MatCreateSubMatrices_MPIBAIJ_local(C, max_no, is_row + pos, is_new + pos, MAT_INITIAL_MATRIX, submats + pos)); 689566063dSJacob Faibussowitsch PetscCall(PetscStrcpy(((PetscObject)c->A)->type_name, MATSEQSBAIJ)); 69db41cccfSHong Zhang pos += max_no; 70db41cccfSHong Zhang } 71db41cccfSHong Zhang 72b11de2a9SHong Zhang /* 2) Row search */ 739566063dSJacob Faibussowitsch PetscCall(MatIncreaseOverlap_MPIBAIJ_Once(C, is_max, is_new)); 74db41cccfSHong Zhang 75b11de2a9SHong Zhang /* 3) Column search */ 76db41cccfSHong Zhang for (i = 0; i < is_max; i++) { 77db41cccfSHong Zhang asub_i = (Mat_SeqSBAIJ *)submats[i]->data; 78feb237baSPierre Jolivet ai = asub_i->i; 79db41cccfSHong Zhang 80db41cccfSHong Zhang /* put is_new obtained from MatIncreaseOverlap_MPIBAIJ() to table */ 819566063dSJacob Faibussowitsch PetscCall(PetscBTMemzero(Mbs, table)); 82db41cccfSHong Zhang 839566063dSJacob Faibussowitsch PetscCall(ISGetIndices(is_new[i], &idx)); 849566063dSJacob Faibussowitsch PetscCall(ISGetLocalSize(is_new[i], &nis)); 85db41cccfSHong Zhang for (l = 0; l < nis; l++) { 869566063dSJacob Faibussowitsch PetscCall(PetscBTSet(table, idx[l])); 87db41cccfSHong Zhang nidx[l] = idx[l]; 88db41cccfSHong Zhang } 89db41cccfSHong Zhang isz = nis; 90db41cccfSHong Zhang 91b11de2a9SHong Zhang /* add column entries to table */ 92db41cccfSHong Zhang for (brow = 0; brow < Mbs; brow++) { 93db41cccfSHong Zhang nz = ai[brow + 1] - ai[brow]; 94db41cccfSHong Zhang if (nz) { 95db41cccfSHong Zhang if (!PetscBTLookupSet(table, brow)) nidx[isz++] = brow; 96db41cccfSHong Zhang } 97db41cccfSHong Zhang } 989566063dSJacob Faibussowitsch PetscCall(ISRestoreIndices(is_new[i], &idx)); 999566063dSJacob Faibussowitsch PetscCall(ISDestroy(&is_new[i])); 100db41cccfSHong Zhang 101db41cccfSHong Zhang /* create updated is_new */ 1029566063dSJacob Faibussowitsch PetscCall(ISCreateGeneral(PETSC_COMM_SELF, isz, nidx, PETSC_COPY_VALUES, is_new + i)); 103db41cccfSHong Zhang } 104db41cccfSHong Zhang 105db41cccfSHong Zhang /* Free tmp spaces */ 10648a46eb9SPierre Jolivet for (i = 0; i < is_max; i++) PetscCall(MatDestroy(&submats[i])); 107db41cccfSHong Zhang } 108b829d874SHong Zhang 1099566063dSJacob Faibussowitsch PetscCall(PetscBTDestroy(&table)); 1109566063dSJacob Faibussowitsch PetscCall(PetscFree(submats)); 1119566063dSJacob Faibussowitsch PetscCall(ISDestroy(&is_row[0])); 1129566063dSJacob Faibussowitsch PetscCall(PetscFree(is_row)); 1139566063dSJacob Faibussowitsch PetscCall(PetscFree(nidx)); 114db41cccfSHong Zhang } 115db41cccfSHong Zhang 1169566063dSJacob Faibussowitsch for (i = 0; i < is_max; i++) PetscCall(ISDestroy(&is[i])); 1179566063dSJacob Faibussowitsch PetscCall(ISExpandIndicesGeneral(N, N, bs, is_max, is_new, is)); 118db41cccfSHong Zhang 1199566063dSJacob Faibussowitsch for (i = 0; i < is_max; i++) PetscCall(ISDestroy(&is_new[i])); 1209566063dSJacob Faibussowitsch PetscCall(PetscFree(is_new)); 1213ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 122632d0f97SHong Zhang } 123632d0f97SHong Zhang 1249371c9d4SSatish Balay typedef enum { 1259371c9d4SSatish Balay MINE, 1269371c9d4SSatish Balay OTHER 1279371c9d4SSatish Balay } WhoseOwner; 1280472cc68SHong Zhang /* data1, odata1 and odata2 are packed in the format (for communication): 129a2a9f22aSHong Zhang data[0] = is_max, no of is 130a2a9f22aSHong Zhang data[1] = size of is[0] 131a2a9f22aSHong Zhang ... 132a2a9f22aSHong Zhang data[is_max] = size of is[is_max-1] 133a2a9f22aSHong Zhang data[is_max + 1] = data(is[0]) 134a2a9f22aSHong Zhang ... 135a2a9f22aSHong Zhang data[is_max+1+sum(size of is[k]), k=0,...,i-1] = data(is[i]) 136a2a9f22aSHong Zhang ... 1370472cc68SHong Zhang data2 is packed in the format (for creating output is[]): 1380472cc68SHong Zhang data[0] = is_max, no of is 1390472cc68SHong Zhang data[1] = size of is[0] 1400472cc68SHong Zhang ... 1410472cc68SHong Zhang data[is_max] = size of is[is_max-1] 1420472cc68SHong Zhang data[is_max + 1] = data(is[0]) 1430472cc68SHong Zhang ... 1440472cc68SHong Zhang data[is_max + 1 + Mbs*i) = data(is[i]) 1450472cc68SHong Zhang ... 146a2a9f22aSHong Zhang */ 147d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatIncreaseOverlap_MPISBAIJ_Once(Mat C, PetscInt is_max, IS is[]) 148d71ae5a4SJacob Faibussowitsch { 149632d0f97SHong Zhang Mat_MPISBAIJ *c = (Mat_MPISBAIJ *)C->data; 150adcec1e5SJed Brown PetscMPIInt size, rank, tag1, tag2, *len_s, nrqr, nrqs, *id_r1, *len_r1, flag, len, *iwork; 1515d0c19d7SBarry Smith const PetscInt *idx_i; 15226fbe8dcSKarl Rupp PetscInt idx, isz, col, *n, *data1, **data1_start, *data2, *data2_i, *data, *data_i; 15326fbe8dcSKarl Rupp PetscInt Mbs, i, j, k, *odata1, *odata2; 154f4259b30SLisandro Dalcin PetscInt proc_id, **odata2_ptr, *ctable = NULL, *btable, len_max, len_est; 155adcec1e5SJed Brown PetscInt proc_end = 0, len_unused, nodata2; 15613f74950SBarry Smith PetscInt ois_max; /* max no of is[] in each of processor */ 157bfc6803cSHong Zhang char *t_p; 158632d0f97SHong Zhang MPI_Comm comm; 159e8527bf2SHong Zhang MPI_Request *s_waits1, *s_waits2, r_req; 160632d0f97SHong Zhang MPI_Status *s_status, r_status; 16145f43ab7SHong Zhang PetscBT *table; /* mark indices of this processor's is[] */ 162fc70d14eSHong Zhang PetscBT table_i; 163bfc6803cSHong Zhang PetscBT otable; /* mark indices of other processors' is[] */ 164d0f46423SBarry Smith PetscInt bs = C->rmap->bs, Bn = c->B->cmap->n, Bnbs = Bn / bs, *Bowners; 16510eea884SHong Zhang IS garray_local, garray_gl; 1665483b11dSHong Zhang 167632d0f97SHong Zhang PetscFunctionBegin; 1689566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C, &comm)); 169632d0f97SHong Zhang size = c->size; 170632d0f97SHong Zhang rank = c->rank; 171632d0f97SHong Zhang Mbs = c->Mbs; 172632d0f97SHong Zhang 1739566063dSJacob Faibussowitsch PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag1)); 1749566063dSJacob Faibussowitsch PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag2)); 175c910923dSHong Zhang 176430a0127SHong Zhang /* create tables used in 177430a0127SHong Zhang step 1: table[i] - mark c->garray of proc [i] 17845f43ab7SHong Zhang step 3: table[i] - mark indices of is[i] when whose=MINE 179430a0127SHong Zhang table[0] - mark incideces of is[] when whose=OTHER */ 1802f613bf5SBarry Smith len = PetscMax(is_max, size); 1819566063dSJacob Faibussowitsch PetscCall(PetscMalloc2(len, &table, (Mbs / PETSC_BITS_PER_BYTE + 1) * len, &t_p)); 182ad540459SPierre Jolivet for (i = 0; i < len; i++) table[i] = t_p + (Mbs / PETSC_BITS_PER_BYTE + 1) * i; 183430a0127SHong Zhang 1841c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&is_max, &ois_max, 1, MPIU_INT, MPI_MAX, comm)); 18510eea884SHong Zhang 1865483b11dSHong Zhang /* 1. Send this processor's is[] to other processors */ 187e8527bf2SHong Zhang /* allocate spaces */ 1889566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(is_max, &n)); 1895483b11dSHong Zhang len = 0; 190c910923dSHong Zhang for (i = 0; i < is_max; i++) { 1919566063dSJacob Faibussowitsch PetscCall(ISGetLocalSize(is[i], &n[i])); 192632d0f97SHong Zhang len += n[i]; 193632d0f97SHong Zhang } 194958c9bccSBarry Smith if (!len) { 1955483b11dSHong Zhang is_max = 0; 1965483b11dSHong Zhang } else { 1975483b11dSHong Zhang len += 1 + is_max; /* max length of data1 for one processor */ 1985483b11dSHong Zhang } 1995483b11dSHong Zhang 2009566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size * len + 1, &data1)); 2019566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size, &data1_start)); 2025483b11dSHong Zhang for (i = 0; i < size; i++) data1_start[i] = data1 + i * len; 2035483b11dSHong Zhang 2049566063dSJacob Faibussowitsch PetscCall(PetscMalloc4(size, &len_s, size, &btable, size, &iwork, size + 1, &Bowners)); 205e8527bf2SHong Zhang 20676f244e2SHong Zhang /* gather c->garray from all processors */ 2079566063dSJacob Faibussowitsch PetscCall(ISCreateGeneral(comm, Bnbs, c->garray, PETSC_COPY_VALUES, &garray_local)); 2089566063dSJacob Faibussowitsch PetscCall(ISAllGather(garray_local, &garray_gl)); 2099566063dSJacob Faibussowitsch PetscCall(ISDestroy(&garray_local)); 2109566063dSJacob Faibussowitsch PetscCallMPI(MPI_Allgather(&Bnbs, 1, MPIU_INT, Bowners + 1, 1, MPIU_INT, comm)); 21126fbe8dcSKarl Rupp 21276f244e2SHong Zhang Bowners[0] = 0; 21376f244e2SHong Zhang for (i = 0; i < size; i++) Bowners[i + 1] += Bowners[i]; 21476f244e2SHong Zhang 2155483b11dSHong Zhang if (is_max) { 216430a0127SHong Zhang /* hash table ctable which maps c->row to proc_id) */ 2179566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(Mbs, &ctable)); 2185483b11dSHong Zhang for (proc_id = 0, j = 0; proc_id < size; proc_id++) { 21926fbe8dcSKarl Rupp for (; j < C->rmap->range[proc_id + 1] / bs; j++) ctable[j] = proc_id; 2205483b11dSHong Zhang } 2215483b11dSHong Zhang 222430a0127SHong Zhang /* hash tables marking c->garray */ 2239566063dSJacob Faibussowitsch PetscCall(ISGetIndices(garray_gl, &idx_i)); 2245483b11dSHong Zhang for (i = 0; i < size; i++) { 225430a0127SHong Zhang table_i = table[i]; 2269566063dSJacob Faibussowitsch PetscCall(PetscBTMemzero(Mbs, table_i)); 227430a0127SHong Zhang for (j = Bowners[i]; j < Bowners[i + 1]; j++) { /* go through B cols of proc[i]*/ 2289566063dSJacob Faibussowitsch PetscCall(PetscBTSet(table_i, idx_i[j])); 2295483b11dSHong Zhang } 2305483b11dSHong Zhang } 2319566063dSJacob Faibussowitsch PetscCall(ISRestoreIndices(garray_gl, &idx_i)); 2325483b11dSHong Zhang } /* if (is_max) */ 2339566063dSJacob Faibussowitsch PetscCall(ISDestroy(&garray_gl)); 2345483b11dSHong Zhang 2355483b11dSHong Zhang /* evaluate communication - mesg to who, length, and buffer space */ 236e8527bf2SHong Zhang for (i = 0; i < size; i++) len_s[i] = 0; 2375483b11dSHong Zhang 2385483b11dSHong Zhang /* header of data1 */ 2395483b11dSHong Zhang for (proc_id = 0; proc_id < size; proc_id++) { 2405483b11dSHong Zhang iwork[proc_id] = 0; 2415483b11dSHong Zhang *data1_start[proc_id] = is_max; 2425483b11dSHong Zhang data1_start[proc_id]++; 2435483b11dSHong Zhang for (j = 0; j < is_max; j++) { 2445483b11dSHong Zhang if (proc_id == rank) { 2455483b11dSHong Zhang *data1_start[proc_id] = n[j]; 2465483b11dSHong Zhang } else { 2475483b11dSHong Zhang *data1_start[proc_id] = 0; 2485483b11dSHong Zhang } 2495483b11dSHong Zhang data1_start[proc_id]++; 2505483b11dSHong Zhang } 2515483b11dSHong Zhang } 2525483b11dSHong Zhang 2535483b11dSHong Zhang for (i = 0; i < is_max; i++) { 2549566063dSJacob Faibussowitsch PetscCall(ISGetIndices(is[i], &idx_i)); 2555483b11dSHong Zhang for (j = 0; j < n[i]; j++) { 2565483b11dSHong Zhang idx = idx_i[j]; 2579371c9d4SSatish Balay *data1_start[rank] = idx; 25835cb6cd3SPierre Jolivet data1_start[rank]++; /* for local processing */ 2595483b11dSHong Zhang proc_end = ctable[idx]; 2605483b11dSHong Zhang for (proc_id = 0; proc_id <= proc_end; proc_id++) { /* for others to process */ 261e8527bf2SHong Zhang if (proc_id == rank) continue; /* done before this loop */ 26226fbe8dcSKarl Rupp if (proc_id < proc_end && !PetscBTLookup(table[proc_id], idx)) continue; /* no need for sending idx to [proc_id] */ 2639371c9d4SSatish Balay *data1_start[proc_id] = idx; 2649371c9d4SSatish Balay data1_start[proc_id]++; 2655483b11dSHong Zhang len_s[proc_id]++; 2665483b11dSHong Zhang } 2675483b11dSHong Zhang } 2685483b11dSHong Zhang /* update header data */ 2692cfbe0a4SHong Zhang for (proc_id = 0; proc_id < size; proc_id++) { 2705483b11dSHong Zhang if (proc_id == rank) continue; 2715483b11dSHong Zhang *(data1 + proc_id * len + 1 + i) = len_s[proc_id] - iwork[proc_id]; 2725483b11dSHong Zhang iwork[proc_id] = len_s[proc_id]; 2735483b11dSHong Zhang } 2749566063dSJacob Faibussowitsch PetscCall(ISRestoreIndices(is[i], &idx_i)); 275e8527bf2SHong Zhang } 2765483b11dSHong Zhang 2779371c9d4SSatish Balay nrqs = 0; 2789371c9d4SSatish Balay nrqr = 0; 2795483b11dSHong Zhang for (i = 0; i < size; i++) { 2805483b11dSHong Zhang data1_start[i] = data1 + i * len; 2815483b11dSHong Zhang if (len_s[i]) { 2825483b11dSHong Zhang nrqs++; 2835483b11dSHong Zhang len_s[i] += 1 + is_max; /* add no. of header msg */ 2845483b11dSHong Zhang } 2855483b11dSHong Zhang } 2865483b11dSHong Zhang 28748a46eb9SPierre Jolivet for (i = 0; i < is_max; i++) PetscCall(ISDestroy(&is[i])); 2889566063dSJacob Faibussowitsch PetscCall(PetscFree(n)); 2899566063dSJacob Faibussowitsch PetscCall(PetscFree(ctable)); 2905483b11dSHong Zhang 2915483b11dSHong Zhang /* Determine the number of messages to expect, their lengths, from from-ids */ 2929566063dSJacob Faibussowitsch PetscCall(PetscGatherNumberOfMessages(comm, NULL, len_s, &nrqr)); 2939566063dSJacob Faibussowitsch PetscCall(PetscGatherMessageLengths(comm, nrqs, nrqr, len_s, &id_r1, &len_r1)); 294632d0f97SHong Zhang 295632d0f97SHong Zhang /* Now post the sends */ 2969566063dSJacob Faibussowitsch PetscCall(PetscMalloc2(size, &s_waits1, size, &s_waits2)); 297632d0f97SHong Zhang k = 0; 2985483b11dSHong Zhang for (proc_id = 0; proc_id < size; proc_id++) { /* send data1 to processor [proc_id] */ 2995483b11dSHong Zhang if (len_s[proc_id]) { 3009566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(data1_start[proc_id], len_s[proc_id], MPIU_INT, proc_id, tag1, comm, s_waits1 + k)); 301632d0f97SHong Zhang k++; 302632d0f97SHong Zhang } 303632d0f97SHong Zhang } 304632d0f97SHong Zhang 30545f43ab7SHong Zhang /* 2. Receive other's is[] and process. Then send back */ 3065483b11dSHong Zhang len = 0; 3075483b11dSHong Zhang for (i = 0; i < nrqr; i++) { 3085483b11dSHong Zhang if (len_r1[i] > len) len = len_r1[i]; 3095483b11dSHong Zhang } 3109566063dSJacob Faibussowitsch PetscCall(PetscFree(len_r1)); 3119566063dSJacob Faibussowitsch PetscCall(PetscFree(id_r1)); 31245f43ab7SHong Zhang 31326fbe8dcSKarl Rupp for (proc_id = 0; proc_id < size; proc_id++) len_s[proc_id] = iwork[proc_id] = 0; 31445f43ab7SHong Zhang 3159566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(len + 1, &odata1)); 3169566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size, &odata2_ptr)); 3179566063dSJacob Faibussowitsch PetscCall(PetscBTCreate(Mbs, &otable)); 31810eea884SHong Zhang 31910eea884SHong Zhang len_max = ois_max * (Mbs + 1); /* max space storing all is[] for each receive */ 320240e5896SHong Zhang len_est = 2 * len_max; /* estimated space of storing is[] for all receiving messages */ 3219566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(len_est + 1, &odata2)); 32210eea884SHong Zhang nodata2 = 0; /* nodata2+1: num of PetscMalloc(,&odata2_ptr[]) called */ 32326fbe8dcSKarl Rupp 324240e5896SHong Zhang odata2_ptr[nodata2] = odata2; 32526fbe8dcSKarl Rupp 32610eea884SHong Zhang len_unused = len_est; /* unused space in the array odata2_ptr[nodata2]-- needs to be >= len_max */ 32710eea884SHong Zhang 328632d0f97SHong Zhang k = 0; 3295483b11dSHong Zhang while (k < nrqr) { 330632d0f97SHong Zhang /* Receive messages */ 3319566063dSJacob Faibussowitsch PetscCallMPI(MPI_Iprobe(MPI_ANY_SOURCE, tag1, comm, &flag, &r_status)); 332632d0f97SHong Zhang if (flag) { 3339566063dSJacob Faibussowitsch PetscCallMPI(MPI_Get_count(&r_status, MPIU_INT, &len)); 334632d0f97SHong Zhang proc_id = r_status.MPI_SOURCE; 3359566063dSJacob Faibussowitsch PetscCallMPI(MPI_Irecv(odata1, len, MPIU_INT, proc_id, r_status.MPI_TAG, comm, &r_req)); 3369566063dSJacob Faibussowitsch PetscCallMPI(MPI_Wait(&r_req, &r_status)); 337632d0f97SHong Zhang 338fc70d14eSHong Zhang /* Process messages */ 339240e5896SHong Zhang /* make sure there is enough unused space in odata2 array */ 34010eea884SHong Zhang if (len_unused < len_max) { /* allocate more space for odata2 */ 3419566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(len_est + 1, &odata2)); 34226fbe8dcSKarl Rupp 343240e5896SHong Zhang odata2_ptr[++nodata2] = odata2; 34426fbe8dcSKarl Rupp 34510eea884SHong Zhang len_unused = len_est; 34610eea884SHong Zhang } 34710eea884SHong Zhang 3489566063dSJacob Faibussowitsch PetscCall(MatIncreaseOverlap_MPISBAIJ_Local(C, odata1, OTHER, odata2, &otable)); 349a2a9f22aSHong Zhang len = 1 + odata2[0]; 35026fbe8dcSKarl Rupp for (i = 0; i < odata2[0]; i++) len += odata2[1 + i]; 351632d0f97SHong Zhang 352632d0f97SHong Zhang /* Send messages back */ 3539566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(odata2, len, MPIU_INT, proc_id, tag2, comm, s_waits2 + k)); 354fc70d14eSHong Zhang k++; 35510eea884SHong Zhang odata2 += len; 35610eea884SHong Zhang len_unused -= len; 35745f43ab7SHong Zhang len_s[proc_id] = len; /* num of messages sending back to [proc_id] by this proc */ 358632d0f97SHong Zhang } 3595483b11dSHong Zhang } 3609566063dSJacob Faibussowitsch PetscCall(PetscFree(odata1)); 3619566063dSJacob Faibussowitsch PetscCall(PetscBTDestroy(&otable)); 362632d0f97SHong Zhang 36345f43ab7SHong Zhang /* 3. Do local work on this processor's is[] */ 36445f43ab7SHong Zhang /* make sure there is enough unused space in odata2(=data) array */ 36545f43ab7SHong Zhang len_max = is_max * (Mbs + 1); /* max space storing all is[] for this processor */ 36610eea884SHong Zhang if (len_unused < len_max) { /* allocate more space for odata2 */ 3679566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(len_est + 1, &odata2)); 36826fbe8dcSKarl Rupp 369240e5896SHong Zhang odata2_ptr[++nodata2] = odata2; 37010eea884SHong Zhang } 371bfc6803cSHong Zhang 372240e5896SHong Zhang data = odata2; 3739566063dSJacob Faibussowitsch PetscCall(MatIncreaseOverlap_MPISBAIJ_Local(C, data1_start[rank], MINE, data, table)); 3749566063dSJacob Faibussowitsch PetscCall(PetscFree(data1_start)); 37545f43ab7SHong Zhang 37645f43ab7SHong Zhang /* 4. Receive work done on other processors, then merge */ 37745f43ab7SHong Zhang /* get max number of messages that this processor expects to recv */ 3781c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(len_s, iwork, size, MPI_INT, MPI_MAX, comm)); 3799566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(iwork[rank] + 1, &data2)); 3809566063dSJacob Faibussowitsch PetscCall(PetscFree4(len_s, btable, iwork, Bowners)); 38145f43ab7SHong Zhang 382632d0f97SHong Zhang k = 0; 3835483b11dSHong Zhang while (k < nrqs) { 384632d0f97SHong Zhang /* Receive messages */ 3859566063dSJacob Faibussowitsch PetscCallMPI(MPI_Iprobe(MPI_ANY_SOURCE, tag2, comm, &flag, &r_status)); 386632d0f97SHong Zhang if (flag) { 3879566063dSJacob Faibussowitsch PetscCallMPI(MPI_Get_count(&r_status, MPIU_INT, &len)); 38826fbe8dcSKarl Rupp 389632d0f97SHong Zhang proc_id = r_status.MPI_SOURCE; 39026fbe8dcSKarl Rupp 3919566063dSJacob Faibussowitsch PetscCallMPI(MPI_Irecv(data2, len, MPIU_INT, proc_id, r_status.MPI_TAG, comm, &r_req)); 3929566063dSJacob Faibussowitsch PetscCallMPI(MPI_Wait(&r_req, &r_status)); 3935483b11dSHong Zhang if (len > 1 + is_max) { /* Add data2 into data */ 3940472cc68SHong Zhang data2_i = data2 + 1 + is_max; 395fc70d14eSHong Zhang for (i = 0; i < is_max; i++) { 396fc70d14eSHong Zhang table_i = table[i]; 397bfc6803cSHong Zhang data_i = data + 1 + is_max + Mbs * i; 398bfc6803cSHong Zhang isz = data[1 + i]; 3990472cc68SHong Zhang for (j = 0; j < data2[1 + i]; j++) { 4000472cc68SHong Zhang col = data2_i[j]; 40126fbe8dcSKarl Rupp if (!PetscBTLookupSet(table_i, col)) data_i[isz++] = col; 402fc70d14eSHong Zhang } 403bfc6803cSHong Zhang data[1 + i] = isz; 4040472cc68SHong Zhang if (i < is_max - 1) data2_i += data2[1 + i]; 405fc70d14eSHong Zhang } 4065483b11dSHong Zhang } 407632d0f97SHong Zhang k++; 408632d0f97SHong Zhang } 4095483b11dSHong Zhang } 4109566063dSJacob Faibussowitsch PetscCall(PetscFree(data2)); 4119566063dSJacob Faibussowitsch PetscCall(PetscFree2(table, t_p)); 4125483b11dSHong Zhang 4135483b11dSHong Zhang /* phase 1 sends are complete */ 4149566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size, &s_status)); 4159566063dSJacob Faibussowitsch if (nrqs) PetscCallMPI(MPI_Waitall(nrqs, s_waits1, s_status)); 4169566063dSJacob Faibussowitsch PetscCall(PetscFree(data1)); 4175483b11dSHong Zhang 418240e5896SHong Zhang /* phase 2 sends are complete */ 4199566063dSJacob Faibussowitsch if (nrqr) PetscCallMPI(MPI_Waitall(nrqr, s_waits2, s_status)); 4209566063dSJacob Faibussowitsch PetscCall(PetscFree2(s_waits1, s_waits2)); 4219566063dSJacob Faibussowitsch PetscCall(PetscFree(s_status)); 422632d0f97SHong Zhang 423c910923dSHong Zhang /* 5. Create new is[] */ 424c910923dSHong Zhang for (i = 0; i < is_max; i++) { 425bfc6803cSHong Zhang data_i = data + 1 + is_max + Mbs * i; 4269566063dSJacob Faibussowitsch PetscCall(ISCreateGeneral(PETSC_COMM_SELF, data[1 + i], data_i, PETSC_COPY_VALUES, is + i)); 427632d0f97SHong Zhang } 42848a46eb9SPierre Jolivet for (k = 0; k <= nodata2; k++) PetscCall(PetscFree(odata2_ptr[k])); 4299566063dSJacob Faibussowitsch PetscCall(PetscFree(odata2_ptr)); 4303ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 431632d0f97SHong Zhang } 432632d0f97SHong Zhang 433632d0f97SHong Zhang /* 434dc008846SHong Zhang MatIncreaseOverlap_MPISBAIJ_Local - Called by MatIncreaseOverlap, to do 435632d0f97SHong Zhang the work on the local processor. 436632d0f97SHong Zhang 437632d0f97SHong Zhang Inputs: 438632d0f97SHong Zhang C - MAT_MPISBAIJ; 439bfc6803cSHong Zhang data - holds is[]. See MatIncreaseOverlap_MPISBAIJ_Once() for the format. 440bfc6803cSHong Zhang whose - whose is[] to be processed, 441bfc6803cSHong Zhang MINE: this processor's is[] 442bfc6803cSHong Zhang OTHER: other processor's is[] 443632d0f97SHong Zhang Output: 44410eea884SHong Zhang nidx - whose = MINE: 4450472cc68SHong Zhang holds input and newly found indices in the same format as data 4460472cc68SHong Zhang whose = OTHER: 4470472cc68SHong Zhang only holds the newly found indices 4480472cc68SHong Zhang table - table[i]: mark the indices of is[i], i=0,...,is_max. Used only in the case 'whose=MINE'. 449632d0f97SHong Zhang */ 45076f244e2SHong Zhang /* Would computation be reduced by swapping the loop 'for each is' and 'for each row'? */ 451d71ae5a4SJacob Faibussowitsch static PetscErrorCode MatIncreaseOverlap_MPISBAIJ_Local(Mat C, PetscInt *data, PetscInt whose, PetscInt *nidx, PetscBT *table) 452d71ae5a4SJacob Faibussowitsch { 453632d0f97SHong Zhang Mat_MPISBAIJ *c = (Mat_MPISBAIJ *)C->data; 454dc008846SHong Zhang Mat_SeqSBAIJ *a = (Mat_SeqSBAIJ *)(c->A)->data; 455dc008846SHong Zhang Mat_SeqBAIJ *b = (Mat_SeqBAIJ *)(c->B)->data; 45613f74950SBarry Smith PetscInt row, mbs, Mbs, *nidx_i, col, col_max, isz, isz0, *ai, *aj, *bi, *bj, *garray, rstart, l; 45713f74950SBarry Smith PetscInt a_start, a_end, b_start, b_end, i, j, k, is_max, *idx_i, n; 458bfc6803cSHong Zhang PetscBT table0; /* mark the indices of input is[] for look up */ 459da81f932SPierre Jolivet PetscBT table_i; /* points to i-th table. When whose=OTHER, a single table is used for all is[] */ 460632d0f97SHong Zhang 461632d0f97SHong Zhang PetscFunctionBegin; 4629371c9d4SSatish Balay Mbs = c->Mbs; 4639371c9d4SSatish Balay mbs = a->mbs; 4649371c9d4SSatish Balay ai = a->i; 4659371c9d4SSatish Balay aj = a->j; 4669371c9d4SSatish Balay bi = b->i; 4679371c9d4SSatish Balay bj = b->j; 468632d0f97SHong Zhang garray = c->garray; 469899cda47SBarry Smith rstart = c->rstartbs; 470dc008846SHong Zhang is_max = data[0]; 471632d0f97SHong Zhang 4729566063dSJacob Faibussowitsch PetscCall(PetscBTCreate(Mbs, &table0)); 473fc70d14eSHong Zhang 474fc70d14eSHong Zhang nidx[0] = is_max; 475fc70d14eSHong Zhang idx_i = data + is_max + 1; /* ptr to input is[0] array */ 476bfc6803cSHong Zhang nidx_i = nidx + is_max + 1; /* ptr to output is[0] array */ 477dc008846SHong Zhang for (i = 0; i < is_max; i++) { /* for each is */ 478dc008846SHong Zhang isz = 0; 479fc70d14eSHong Zhang n = data[1 + i]; /* size of input is[i] */ 480dc008846SHong Zhang 48176f244e2SHong Zhang /* initialize and set table_i(mark idx and nidx) and table0(only mark idx) */ 482bfc6803cSHong Zhang if (whose == MINE) { /* process this processor's is[] */ 483bfc6803cSHong Zhang table_i = table[i]; 4840472cc68SHong Zhang nidx_i = nidx + 1 + is_max + Mbs * i; 485bfc6803cSHong Zhang } else { /* process other processor's is[] - only use one temp table */ 486430a0127SHong Zhang table_i = table[0]; 487bfc6803cSHong Zhang } 4889566063dSJacob Faibussowitsch PetscCall(PetscBTMemzero(Mbs, table_i)); 4899566063dSJacob Faibussowitsch PetscCall(PetscBTMemzero(Mbs, table0)); 49076f244e2SHong Zhang if (n == 0) { 49176f244e2SHong Zhang nidx[1 + i] = 0; /* size of new is[i] */ 49276f244e2SHong Zhang continue; 49376f244e2SHong Zhang } 49476f244e2SHong Zhang 4959371c9d4SSatish Balay isz0 = 0; 4969371c9d4SSatish Balay col_max = 0; 497dc008846SHong Zhang for (j = 0; j < n; j++) { 498dc008846SHong Zhang col = idx_i[j]; 49908401ef6SPierre Jolivet PetscCheck(col < Mbs, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "index col %" PetscInt_FMT " >= Mbs %" PetscInt_FMT, col, Mbs); 500bfc6803cSHong Zhang if (!PetscBTLookupSet(table_i, col)) { 5019566063dSJacob Faibussowitsch PetscCall(PetscBTSet(table0, col)); 50226fbe8dcSKarl Rupp if (whose == MINE) nidx_i[isz0] = col; 503dbe03f88SHong Zhang if (col_max < col) col_max = col; 504bfc6803cSHong Zhang isz0++; 505bfc6803cSHong Zhang } 506632d0f97SHong Zhang } 507dc008846SHong Zhang 50826fbe8dcSKarl Rupp if (whose == MINE) isz = isz0; 509fc70d14eSHong Zhang k = 0; /* no. of indices from input is[i] that have been examined */ 510dc008846SHong Zhang for (row = 0; row < mbs; row++) { 5119371c9d4SSatish Balay a_start = ai[row]; 5129371c9d4SSatish Balay a_end = ai[row + 1]; 5139371c9d4SSatish Balay b_start = bi[row]; 5149371c9d4SSatish Balay b_end = bi[row + 1]; 5150472cc68SHong Zhang if (PetscBTLookup(table0, row + rstart)) { /* row is on input is[i]: 5160472cc68SHong Zhang do row search: collect all col in this row */ 517dc008846SHong Zhang for (l = a_start; l < a_end; l++) { /* Amat */ 518dc008846SHong Zhang col = aj[l] + rstart; 51926fbe8dcSKarl Rupp if (!PetscBTLookupSet(table_i, col)) nidx_i[isz++] = col; 520dc008846SHong Zhang } 521dc008846SHong Zhang for (l = b_start; l < b_end; l++) { /* Bmat */ 522dc008846SHong Zhang col = garray[bj[l]]; 52326fbe8dcSKarl Rupp if (!PetscBTLookupSet(table_i, col)) nidx_i[isz++] = col; 524dc008846SHong Zhang } 525dc008846SHong Zhang k++; 526dc008846SHong Zhang if (k >= isz0) break; /* for (row=0; row<mbs; row++) */ 5270472cc68SHong Zhang } else { /* row is not on input is[i]: 528da81f932SPierre Jolivet do col search: add row onto nidx_i if there is a col in nidx_i */ 529dc008846SHong Zhang for (l = a_start; l < a_end; l++) { /* Amat */ 530dc008846SHong Zhang col = aj[l] + rstart; 53176f244e2SHong Zhang if (col > col_max) break; 532dc008846SHong Zhang if (PetscBTLookup(table0, col)) { 53326fbe8dcSKarl Rupp if (!PetscBTLookupSet(table_i, row + rstart)) nidx_i[isz++] = row + rstart; 534dc008846SHong Zhang break; /* for l = start; l<end ; l++) */ 535632d0f97SHong Zhang } 536632d0f97SHong Zhang } 537dc008846SHong Zhang for (l = b_start; l < b_end; l++) { /* Bmat */ 538dc008846SHong Zhang col = garray[bj[l]]; 53976f244e2SHong Zhang if (col > col_max) break; 540dc008846SHong Zhang if (PetscBTLookup(table0, col)) { 54126fbe8dcSKarl Rupp if (!PetscBTLookupSet(table_i, row + rstart)) nidx_i[isz++] = row + rstart; 542dc008846SHong Zhang break; /* for l = start; l<end ; l++) */ 543632d0f97SHong Zhang } 544dc008846SHong Zhang } 545dc008846SHong Zhang } 546bfc6803cSHong Zhang } 547dc008846SHong Zhang 548dc008846SHong Zhang if (i < is_max - 1) { 549fc70d14eSHong Zhang idx_i += n; /* ptr to input is[i+1] array */ 550bfc6803cSHong Zhang nidx_i += isz; /* ptr to output is[i+1] array */ 551dc008846SHong Zhang } 552fc70d14eSHong Zhang nidx[1 + i] = isz; /* size of new is[i] */ 5531ab97a2aSSatish Balay } /* for each is */ 5549566063dSJacob Faibussowitsch PetscCall(PetscBTDestroy(&table0)); 5553ba16761SJacob Faibussowitsch PetscFunctionReturn(PETSC_SUCCESS); 556632d0f97SHong Zhang } 557