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