1 /* 2 Support for the parallel dense matrix vector multiply 3 */ 4 #include <../src/mat/impls/dense/mpi/mpidense.h> 5 #include <petscblaslapack.h> 6 7 PetscErrorCode MatSetUpMultiply_MPIDense(Mat mat) 8 { 9 Mat_MPIDense *mdn = (Mat_MPIDense *)mat->data; 10 11 PetscFunctionBegin; 12 if (!mdn->Mvctx) { 13 /* Create local vector that is used to scatter into */ 14 PetscCall(VecDestroy(&mdn->lvec)); 15 if (mdn->A) PetscCall(MatCreateVecs(mdn->A, &mdn->lvec, NULL)); 16 PetscCall(PetscLayoutSetUp(mat->cmap)); 17 PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)mat), &mdn->Mvctx)); 18 PetscCall(PetscSFSetGraphWithPattern(mdn->Mvctx, mat->cmap, PETSCSF_PATTERN_ALLGATHER)); 19 } 20 PetscFunctionReturn(PETSC_SUCCESS); 21 } 22 23 static PetscErrorCode MatCreateSubMatrices_MPIDense_Local(Mat, PetscInt, const IS[], const IS[], MatReuse, Mat *); 24 25 PetscErrorCode MatCreateSubMatrices_MPIDense(Mat C, PetscInt ismax, const IS isrow[], const IS iscol[], MatReuse scall, Mat *submat[]) 26 { 27 PetscInt nmax, nstages_local, nstages, i, pos, max_no; 28 29 PetscFunctionBegin; 30 /* Allocate memory to hold all the submatrices */ 31 if (scall != MAT_REUSE_MATRIX) PetscCall(PetscCalloc1(ismax + 1, submat)); 32 /* Determine the number of stages through which submatrices are done */ 33 nmax = 20 * 1000000 / (C->cmap->N * sizeof(PetscInt)); 34 if (!nmax) nmax = 1; 35 nstages_local = ismax / nmax + ((ismax % nmax) ? 1 : 0); 36 37 /* Make sure every processor loops through the nstages */ 38 PetscCallMPI(MPIU_Allreduce(&nstages_local, &nstages, 1, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)C))); 39 40 for (i = 0, pos = 0; i < nstages; i++) { 41 if (pos + nmax <= ismax) max_no = nmax; 42 else if (pos == ismax) max_no = 0; 43 else max_no = ismax - pos; 44 PetscCall(MatCreateSubMatrices_MPIDense_Local(C, max_no, isrow + pos, iscol + pos, scall, *submat + pos)); 45 pos += max_no; 46 } 47 PetscFunctionReturn(PETSC_SUCCESS); 48 } 49 50 static PetscErrorCode MatCreateSubMatrices_MPIDense_Local(Mat C, PetscInt ismax, const IS isrow[], const IS iscol[], MatReuse scall, Mat *submats) 51 { 52 Mat_MPIDense *c = (Mat_MPIDense *)C->data; 53 Mat A = c->A; 54 Mat_SeqDense *a = (Mat_SeqDense *)A->data, *mat; 55 PetscMPIInt rank, size, tag0, tag1, idex, end, i, proc, nrqs, *rtable, *pa, nrqr; 56 PetscInt N = C->cmap->N, rstart = C->rmap->rstart, count; 57 const PetscInt **irow, **icol, *irow_i; 58 PetscInt *nrow, *ncol, *w1, *w3, *w4, start, inrqr; 59 PetscInt **sbuf1, m, ct1, **rbuf1, row; 60 PetscInt msz, **ptr, *ctr, *tmp, bsz; 61 PetscInt is_no, jmax, **rmap, *rmap_i; 62 PetscInt ctr_j, *sbuf1_j, *rbuf1_i; 63 MPI_Request *s_waits1, *r_waits1, *s_waits2, *r_waits2; 64 MPI_Status *r_status1, *r_status2, *s_status1, *s_status2; 65 MPI_Comm comm; 66 PetscScalar **rbuf2, **sbuf2; 67 PetscBool sorted; 68 69 PetscFunctionBegin; 70 PetscCall(PetscObjectGetComm((PetscObject)C, &comm)); 71 tag0 = ((PetscObject)C)->tag; 72 PetscCallMPI(MPI_Comm_rank(comm, &rank)); 73 PetscCallMPI(MPI_Comm_size(comm, &size)); 74 m = C->rmap->N; 75 76 /* Get some new tags to keep the communication clean */ 77 PetscCall(PetscObjectGetNewTag((PetscObject)C, &tag1)); 78 79 /* Check if the col indices are sorted */ 80 for (PetscInt i = 0; i < ismax; i++) { 81 PetscCall(ISSorted(isrow[i], &sorted)); 82 PetscCheck(sorted, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "ISrow is not sorted"); 83 PetscCall(ISSorted(iscol[i], &sorted)); 84 PetscCheck(sorted, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "IScol is not sorted"); 85 } 86 87 PetscCall(PetscMalloc5(ismax, (PetscInt ***)&irow, ismax, (PetscInt ***)&icol, ismax, &nrow, ismax, &ncol, m, &rtable)); 88 for (PetscInt i = 0; i < ismax; i++) { 89 PetscCall(ISGetIndices(isrow[i], &irow[i])); 90 PetscCall(ISGetIndices(iscol[i], &icol[i])); 91 PetscCall(ISGetLocalSize(isrow[i], &nrow[i])); 92 PetscCall(ISGetLocalSize(iscol[i], &ncol[i])); 93 } 94 95 /* Create hash table for the mapping :row -> proc*/ 96 for (PetscMPIInt i = 0, j = 0; i < size; i++) { 97 jmax = C->rmap->range[i + 1]; 98 for (; j < jmax; j++) rtable[j] = i; 99 } 100 101 /* evaluate communication - mesg to who,length of mesg, and buffer space 102 required. Based on this, buffers are allocated, and data copied into them*/ 103 PetscCall(PetscMalloc3(2 * size, &w1, size, &w3, size, &w4)); 104 PetscCall(PetscArrayzero(w1, size * 2)); /* initialize work vector*/ 105 PetscCall(PetscArrayzero(w3, size)); /* initialize work vector*/ 106 for (PetscInt i = 0; i < ismax; i++) { 107 PetscCall(PetscArrayzero(w4, size)); /* initialize work vector*/ 108 jmax = nrow[i]; 109 irow_i = irow[i]; 110 for (PetscInt j = 0; j < jmax; j++) { 111 row = irow_i[j]; 112 proc = rtable[row]; 113 w4[proc]++; 114 } 115 for (PetscMPIInt j = 0; j < size; j++) { 116 if (w4[j]) { 117 w1[2 * j] += w4[j]; 118 w3[j]++; 119 } 120 } 121 } 122 123 nrqs = 0; /* no of outgoing messages */ 124 msz = 0; /* total mesg length (for all procs) */ 125 w1[2 * rank] = 0; /* no mesg sent to self */ 126 w3[rank] = 0; 127 for (PetscMPIInt i = 0; i < size; i++) { 128 if (w1[2 * i]) { 129 w1[2 * i + 1] = 1; 130 nrqs++; 131 } /* there exists a message to proc i */ 132 } 133 PetscCall(PetscMalloc1(nrqs + 1, &pa)); /*(proc -array)*/ 134 for (PetscMPIInt i = 0, j = 0; i < size; i++) { 135 if (w1[2 * i]) { 136 pa[j] = i; 137 j++; 138 } 139 } 140 141 /* Each message would have a header = 1 + 2*(no of IS) + data */ 142 for (PetscMPIInt i = 0; i < nrqs; i++) { 143 PetscMPIInt j = pa[i]; 144 w1[2 * j] += w1[2 * j + 1] + 2 * w3[j]; 145 msz += w1[2 * j]; 146 } 147 /* Do a global reduction to determine how many messages to expect*/ 148 PetscCall(PetscMaxSum(comm, w1, &bsz, &inrqr)); 149 PetscCall(PetscMPIIntCast(inrqr, &nrqr)); 150 151 /* Allocate memory for recv buffers . Make sure rbuf1[0] exists by adding 1 to the buffer length */ 152 PetscCall(PetscMalloc1(nrqr + 1, &rbuf1)); 153 PetscCall(PetscMalloc1(nrqr * bsz, &rbuf1[0])); 154 for (PetscInt i = 1; i < nrqr; ++i) rbuf1[i] = rbuf1[i - 1] + bsz; 155 156 /* Post the receives */ 157 PetscCall(PetscMalloc1(nrqr + 1, &r_waits1)); 158 for (PetscInt i = 0; i < nrqr; ++i) PetscCallMPI(MPIU_Irecv(rbuf1[i], bsz, MPIU_INT, MPI_ANY_SOURCE, tag0, comm, r_waits1 + i)); 159 160 /* Allocate Memory for outgoing messages */ 161 PetscCall(PetscMalloc4(size, &sbuf1, size, &ptr, 2 * msz, &tmp, size, &ctr)); 162 PetscCall(PetscArrayzero(sbuf1, size)); 163 PetscCall(PetscArrayzero(ptr, size)); 164 { 165 PetscInt *iptr = tmp, ict = 0; 166 for (PetscMPIInt i = 0; i < nrqs; i++) { 167 PetscMPIInt j = pa[i]; 168 iptr += ict; 169 sbuf1[j] = iptr; 170 ict = w1[2 * j]; 171 } 172 } 173 174 /* Form the outgoing messages */ 175 /* Initialize the header space */ 176 for (PetscMPIInt i = 0; i < nrqs; i++) { 177 PetscInt j = pa[i]; 178 sbuf1[j][0] = 0; 179 PetscCall(PetscArrayzero(sbuf1[j] + 1, 2 * w3[j])); 180 ptr[j] = sbuf1[j] + 2 * w3[j] + 1; 181 } 182 183 /* Parse the isrow and copy data into outbuf */ 184 for (PetscInt i = 0; i < ismax; i++) { 185 PetscCall(PetscArrayzero(ctr, size)); 186 irow_i = irow[i]; 187 jmax = nrow[i]; 188 for (PetscInt j = 0; j < jmax; j++) { /* parse the indices of each IS */ 189 row = irow_i[j]; 190 proc = rtable[row]; 191 if (proc != rank) { /* copy to the outgoing buf*/ 192 ctr[proc]++; 193 *ptr[proc] = row; 194 ptr[proc]++; 195 } 196 } 197 /* Update the headers for the current IS */ 198 for (PetscMPIInt j = 0; j < size; j++) { /* Can Optimise this loop too */ 199 if ((ctr_j = ctr[j])) { 200 PetscInt k; 201 sbuf1_j = sbuf1[j]; 202 k = ++sbuf1_j[0]; 203 sbuf1_j[2 * k] = ctr_j; 204 sbuf1_j[2 * k - 1] = i; 205 } 206 } 207 } 208 209 /* Now post the sends */ 210 PetscCall(PetscMalloc1(nrqs + 1, &s_waits1)); 211 for (PetscMPIInt i = 0; i < nrqs; ++i) { 212 PetscMPIInt j = pa[i]; 213 PetscCallMPI(MPIU_Isend(sbuf1[j], w1[2 * j], MPIU_INT, j, tag0, comm, s_waits1 + i)); 214 } 215 216 /* Post receives to capture the row_data from other procs */ 217 PetscCall(PetscMalloc1(nrqs + 1, &r_waits2)); 218 PetscCall(PetscMalloc1(nrqs + 1, &rbuf2)); 219 for (PetscMPIInt i = 0; i < nrqs; i++) { 220 PetscMPIInt j = pa[i]; 221 count = (w1[2 * j] - (2 * sbuf1[j][0] + 1)) * N; 222 PetscCall(PetscMalloc1(count + 1, &rbuf2[i])); 223 PetscCallMPI(MPIU_Irecv(rbuf2[i], count, MPIU_SCALAR, j, tag1, comm, r_waits2 + i)); 224 } 225 226 /* Receive messages(row_nos) and then, pack and send off the rowvalues 227 to the correct processors */ 228 229 PetscCall(PetscMalloc1(nrqr + 1, &s_waits2)); 230 PetscCall(PetscMalloc1(nrqr + 1, &r_status1)); 231 PetscCall(PetscMalloc1(nrqr + 1, &sbuf2)); 232 233 { 234 PetscScalar *sbuf2_i; 235 const PetscScalar *v_start, *v; 236 PetscMPIInt s_proc; 237 238 PetscCall(MatDenseGetArrayRead(A, &v)); 239 for (PetscMPIInt i = 0; i < nrqr; ++i) { 240 PetscCallMPI(MPI_Waitany(nrqr, r_waits1, &idex, r_status1 + i)); 241 s_proc = r_status1[i].MPI_SOURCE; /* send processor */ 242 rbuf1_i = rbuf1[idex]; /* Actual message from s_proc */ 243 /* no of rows = end - start; since start is array idex[], 0idex, whel end 244 is length of the buffer - which is 1idex */ 245 start = 2 * rbuf1_i[0] + 1; 246 PetscCallMPI(MPI_Get_count(r_status1 + i, MPIU_INT, &end)); 247 /* allocate memory sufficinet to hold all the row values */ 248 PetscCall(PetscMalloc1((end - start) * N, &sbuf2[idex])); 249 sbuf2_i = sbuf2[idex]; 250 /* Now pack the data */ 251 for (PetscInt j = start; j < end; j++) { 252 row = rbuf1_i[j] - rstart; 253 v_start = v + row; 254 for (PetscInt k = 0; k < N; k++) { 255 sbuf2_i[0] = v_start[0]; 256 sbuf2_i++; 257 v_start += a->lda; 258 } 259 } 260 /* Now send off the data */ 261 PetscCallMPI(MPIU_Isend(sbuf2[idex], (end - start) * N, MPIU_SCALAR, s_proc, tag1, comm, s_waits2 + i)); 262 } 263 PetscCall(MatDenseRestoreArrayRead(A, &v)); 264 } 265 /* End Send-Recv of IS + row_numbers */ 266 PetscCall(PetscFree(r_status1)); 267 PetscCall(PetscFree(r_waits1)); 268 PetscCall(PetscMalloc1(nrqs + 1, &s_status1)); 269 if (nrqs) PetscCallMPI(MPI_Waitall(nrqs, s_waits1, s_status1)); 270 PetscCall(PetscFree(s_status1)); 271 PetscCall(PetscFree(s_waits1)); 272 273 /* Create the submatrices */ 274 if (scall == MAT_REUSE_MATRIX) { 275 for (PetscInt i = 0; i < ismax; i++) { 276 mat = (Mat_SeqDense *)submats[i]->data; 277 PetscCheck(!(submats[i]->rmap->n != nrow[i]) && !(submats[i]->cmap->n != ncol[i]), PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Cannot reuse matrix. wrong size"); 278 PetscCall(PetscArrayzero(mat->v, submats[i]->rmap->n * submats[i]->cmap->n)); 279 280 submats[i]->factortype = C->factortype; 281 } 282 } else { 283 for (PetscInt i = 0; i < ismax; i++) { 284 PetscCall(MatCreate(PETSC_COMM_SELF, submats + i)); 285 PetscCall(MatSetSizes(submats[i], nrow[i], ncol[i], nrow[i], ncol[i])); 286 PetscCall(MatSetType(submats[i], ((PetscObject)A)->type_name)); 287 PetscCall(MatSeqDenseSetPreallocation(submats[i], NULL)); 288 } 289 } 290 291 /* Assemble the matrices */ 292 { 293 PetscInt col; 294 PetscScalar *imat_v, *mat_v, *imat_vi, *mat_vi; 295 296 for (PetscInt i = 0; i < ismax; i++) { 297 mat = (Mat_SeqDense *)submats[i]->data; 298 mat_v = a->v; 299 imat_v = mat->v; 300 irow_i = irow[i]; 301 m = nrow[i]; 302 for (PetscInt j = 0; j < m; j++) { 303 row = irow_i[j]; 304 proc = rtable[row]; 305 if (proc == rank) { 306 row = row - rstart; 307 mat_vi = mat_v + row; 308 imat_vi = imat_v + j; 309 for (PetscInt k = 0; k < ncol[i]; k++) { 310 col = icol[i][k]; 311 imat_vi[k * m] = mat_vi[col * a->lda]; 312 } 313 } 314 } 315 } 316 } 317 318 /* Create row map-> This maps c->row to submat->row for each submat*/ 319 /* this is a very expensive operation wrt memory usage */ 320 PetscCall(PetscMalloc1(ismax, &rmap)); 321 PetscCall(PetscCalloc1(ismax * C->rmap->N, &rmap[0])); 322 for (PetscInt i = 1; i < ismax; i++) rmap[i] = rmap[i - 1] + C->rmap->N; 323 for (PetscInt i = 0; i < ismax; i++) { 324 rmap_i = rmap[i]; 325 irow_i = irow[i]; 326 jmax = nrow[i]; 327 for (PetscInt j = 0; j < jmax; j++) rmap_i[irow_i[j]] = j; 328 } 329 330 /* Now Receive the row_values and assemble the rest of the matrix */ 331 PetscCall(PetscMalloc1(nrqs + 1, &r_status2)); 332 { 333 PetscInt is_max, tmp1, col, *sbuf1_i, is_sz; 334 PetscScalar *rbuf2_i, *imat_v, *imat_vi; 335 336 for (tmp1 = 0; tmp1 < nrqs; tmp1++) { /* For each message */ 337 PetscCallMPI(MPI_Waitany(nrqs, r_waits2, &i, r_status2 + tmp1)); 338 /* Now dig out the corresponding sbuf1, which contains the IS data_structure */ 339 sbuf1_i = sbuf1[pa[i]]; 340 is_max = sbuf1_i[0]; 341 ct1 = 2 * is_max + 1; 342 rbuf2_i = rbuf2[i]; 343 for (PetscInt j = 1; j <= is_max; j++) { /* For each IS belonging to the message */ 344 is_no = sbuf1_i[2 * j - 1]; 345 is_sz = sbuf1_i[2 * j]; 346 mat = (Mat_SeqDense *)submats[is_no]->data; 347 imat_v = mat->v; 348 rmap_i = rmap[is_no]; 349 m = nrow[is_no]; 350 for (PetscInt k = 0; k < is_sz; k++, rbuf2_i += N) { /* For each row */ 351 row = sbuf1_i[ct1]; 352 ct1++; 353 row = rmap_i[row]; 354 imat_vi = imat_v + row; 355 for (PetscInt l = 0; l < ncol[is_no]; l++) { /* For each col */ 356 col = icol[is_no][l]; 357 imat_vi[l * m] = rbuf2_i[col]; 358 } 359 } 360 } 361 } 362 } 363 /* End Send-Recv of row_values */ 364 PetscCall(PetscFree(r_status2)); 365 PetscCall(PetscFree(r_waits2)); 366 PetscCall(PetscMalloc1(nrqr + 1, &s_status2)); 367 if (nrqr) PetscCallMPI(MPI_Waitall(nrqr, s_waits2, s_status2)); 368 PetscCall(PetscFree(s_status2)); 369 PetscCall(PetscFree(s_waits2)); 370 371 /* Restore the indices */ 372 for (PetscMPIInt i = 0; i < ismax; i++) { 373 PetscCall(ISRestoreIndices(isrow[i], irow + i)); 374 PetscCall(ISRestoreIndices(iscol[i], icol + i)); 375 } 376 377 PetscCall(PetscFree5(*(PetscInt ***)&irow, *(PetscInt ***)&icol, nrow, ncol, rtable)); 378 PetscCall(PetscFree3(w1, w3, w4)); 379 PetscCall(PetscFree(pa)); 380 381 for (PetscMPIInt i = 0; i < nrqs; ++i) PetscCall(PetscFree(rbuf2[i])); 382 PetscCall(PetscFree(rbuf2)); 383 PetscCall(PetscFree4(sbuf1, ptr, tmp, ctr)); 384 PetscCall(PetscFree(rbuf1[0])); 385 PetscCall(PetscFree(rbuf1)); 386 387 for (PetscMPIInt i = 0; i < nrqr; ++i) PetscCall(PetscFree(sbuf2[i])); 388 389 PetscCall(PetscFree(sbuf2)); 390 PetscCall(PetscFree(rmap[0])); 391 PetscCall(PetscFree(rmap)); 392 393 for (PetscInt i = 0; i < ismax; i++) { 394 PetscCall(MatAssemblyBegin(submats[i], MAT_FINAL_ASSEMBLY)); 395 PetscCall(MatAssemblyEnd(submats[i], MAT_FINAL_ASSEMBLY)); 396 } 397 PetscFunctionReturn(PETSC_SUCCESS); 398 } 399 400 PETSC_INTERN PetscErrorCode MatScale_MPIDense(Mat inA, PetscScalar alpha) 401 { 402 Mat_MPIDense *A = (Mat_MPIDense *)inA->data; 403 404 PetscFunctionBegin; 405 PetscCall(MatScale(A->A, alpha)); 406 PetscFunctionReturn(PETSC_SUCCESS); 407 } 408