1 2 /* 3 Defines projective product routines where A is a MPIAIJ matrix 4 C = P^T * A * P 5 */ 6 7 #include "src/mat/impls/aij/seq/aij.h" /*I "petscmat.h" I*/ 8 #include "src/mat/utils/freespace.h" 9 #include "src/mat/impls/aij/mpi/mpiaij.h" 10 #include "petscbt.h" 11 12 #undef __FUNCT__ 13 #define __FUNCT__ "MatPtAP_MPIAIJ_MPIAIJ" 14 PetscErrorCode MatPtAP_MPIAIJ_MPIAIJ(Mat A,Mat P,MatReuse scall,PetscReal fill,Mat *C) 15 { 16 PetscErrorCode ierr; 17 Mat_MatMatMultMPI *mult=PETSC_NULL; 18 PetscObjectContainer container; 19 20 PetscFunctionBegin; 21 if (scall == MAT_INITIAL_MATRIX){ 22 ierr = PetscLogEventBegin(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr); 23 /* destroy the container 'Mat_MatMatMultMPI' in case that P is attached to it */ 24 ierr = PetscObjectQuery((PetscObject)P,"Mat_MatMatMultMPI",(PetscObject *)&container);CHKERRQ(ierr); 25 if (container) { 26 ierr = PetscObjectCompose((PetscObject)P,"Mat_MatMatMultMPI",0);CHKERRQ(ierr); 27 ierr = PetscObjectContainerDestroy(container);CHKERRQ(ierr); 28 } 29 30 /* create the container 'Mat_MatMatMultMPI' and attach it to P */ 31 ierr = PetscNew(Mat_MatMatMultMPI,&mult);CHKERRQ(ierr); 32 mult->B_loc=PETSC_NULL; mult->B_oth=PETSC_NULL; 33 mult->abi=PETSC_NULL; mult->abj=PETSC_NULL; 34 mult->abnz_max = 0; /* symbolic A*P is not done yet */ 35 36 /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 37 ierr = MatGetBrowsOfAoCols(A,P,scall,&mult->startsj,&mult->bufa,&mult->B_oth);CHKERRQ(ierr); 38 39 /* get P_loc by taking all local rows of P */ 40 ierr = MatGetLocalMat(P,scall,&mult->B_loc);CHKERRQ(ierr); 41 42 /* attach the container 'Mat_MatMatMultMPI' to P */ 43 ierr = PetscObjectContainerCreate(PETSC_COMM_SELF,&container);CHKERRQ(ierr); 44 ierr = PetscObjectContainerSetPointer(container,mult);CHKERRQ(ierr); 45 ierr = PetscObjectCompose((PetscObject)P,"Mat_MatMatMultMPI",(PetscObject)container);CHKERRQ(ierr); 46 mult->MatDestroy = P->ops->destroy; 47 P->ops->destroy = MatDestroy_MPIAIJ_MatMatMult; 48 ierr = PetscObjectContainerSetUserDestroy(container,PetscObjectContainerDestroy_Mat_MatMatMultMPI);CHKERRQ(ierr); 49 50 /* now, compute symbolic local P^T*A*P */ 51 ierr = MatPtAPSymbolic_MPIAIJ_MPIAIJ(A,P,fill,C);CHKERRQ(ierr);/* numeric product is computed as well */ 52 ierr = PetscLogEventEnd(MAT_PtAPSymbolic,A,P,0,0);CHKERRQ(ierr); 53 } else if (scall == MAT_REUSE_MATRIX){ 54 ierr = PetscObjectQuery((PetscObject)P,"Mat_MatMatMultMPI",(PetscObject *)&container);CHKERRQ(ierr); 55 if (container) { 56 ierr = PetscObjectContainerGetPointer(container,(void **)&mult);CHKERRQ(ierr); 57 } else { 58 SETERRQ(PETSC_ERR_ARG_WRONGSTATE, "Matrix P does not posses an object container"); 59 } 60 /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 61 ierr = MatGetBrowsOfAoCols(A,P,scall,&mult->startsj,&mult->bufa,&mult->B_oth);CHKERRQ(ierr); 62 /* get P_loc by taking all local rows of P */ 63 ierr = MatGetLocalMat(P,scall,&mult->B_loc);CHKERRQ(ierr); 64 } else { 65 SETERRQ1(PETSC_ERR_ARG_WRONG,"Invalid MatReuse %d",scall); 66 } 67 /* now, compute numeric local P^T*A*P */ 68 ierr = PetscLogEventBegin(MAT_PtAPNumeric,A,P,0,0);CHKERRQ(ierr); 69 ierr = MatPtAPNumeric_MPIAIJ_MPIAIJ(A,P,*C);CHKERRQ(ierr); 70 ierr = PetscLogEventEnd(MAT_PtAPNumeric,A,P,0,0);CHKERRQ(ierr); 71 72 PetscFunctionReturn(0); 73 } 74 75 #undef __FUNCT__ 76 #define __FUNCT__ "MatPtAPSymbolic_MPIAIJ_MPIAIJ" 77 PetscErrorCode MatPtAPSymbolic_MPIAIJ_MPIAIJ(Mat A,Mat P,PetscReal fill,Mat *C) 78 { 79 PetscErrorCode ierr; 80 Mat P_loc,P_oth,B_mpi; 81 Mat_MatMatMultMPI *ap; 82 PetscObjectContainer container; 83 FreeSpaceList free_space=PETSC_NULL,current_space=PETSC_NULL; 84 Mat_MPIAIJ *a=(Mat_MPIAIJ*)A->data,*p=(Mat_MPIAIJ*)P->data; 85 Mat_SeqAIJ *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data; 86 Mat_SeqAIJ *p_loc,*p_oth; 87 PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*pdti,*pdtj,*poti,*potj,*ptJ; 88 PetscInt *adi=ad->i,*adj=ad->j,*aoi=ao->i,*aoj=ao->j,nnz; 89 PetscInt nlnk,*lnk,*owners_co,*coi,*coj,i,k,pnz,row; 90 PetscInt am=A->m,pN=P->N,pn=P->n; 91 PetscBT lnkbt; 92 MPI_Comm comm=A->comm; 93 PetscMPIInt size,rank,tag,*len_si,*len_s,*len_ri; 94 PetscInt **buf_rj,**buf_ri,**buf_ri_k; 95 PetscInt len,proc,*dnz,*onz,*owners; 96 PetscInt nzi,*bi,*bj; 97 PetscInt nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci; 98 MPI_Request *swaits,*rwaits; 99 MPI_Status *sstatus,rstatus; 100 Mat_Merge_SeqsToMPI *merge; 101 PetscInt *api,*apj,*Jptr,apnz,*prmap=p->garray,pon; 102 PetscMPIInt j; 103 104 PetscFunctionBegin; 105 ierr = PetscObjectQuery((PetscObject)P,"Mat_MatMatMultMPI",(PetscObject *)&container);CHKERRQ(ierr); 106 if (container) { 107 ierr = PetscObjectContainerGetPointer(container,(void **)&ap);CHKERRQ(ierr); 108 } else { 109 SETERRQ(PETSC_ERR_ARG_WRONGSTATE, "Matrix P does not posses an object container"); 110 } 111 ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr); 112 ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr); 113 114 /* get data from P_loc and P_oth */ 115 P_loc=ap->B_loc; P_oth=ap->B_oth; 116 p_loc=(Mat_SeqAIJ*)P_loc->data; p_oth=(Mat_SeqAIJ*)P_oth->data; 117 pi_loc=p_loc->i; pj_loc=p_loc->j; pi_oth=p_oth->i; pj_oth=p_oth->j; 118 119 /* first, compute symbolic AP = A_loc*P */ 120 /*--------------------------------------*/ 121 ierr = PetscMalloc((am+2)*sizeof(PetscInt),&api);CHKERRQ(ierr); 122 ap->abi = api; 123 api[0] = 0; 124 125 /* create and initialize a linked list */ 126 nlnk = pN+1; 127 ierr = PetscLLCreate(pN,pN,nlnk,lnk,lnkbt);CHKERRQ(ierr); 128 129 /* Initial FreeSpace size is fill*nnz(A) */ 130 ierr = GetMoreSpace((PetscInt)(fill*(adi[am]+aoi[am])),&free_space);CHKERRQ(ierr); 131 current_space = free_space; 132 133 for (i=0;i<am;i++) { 134 apnz = 0; 135 /* diagonal portion of A */ 136 nzi = adi[i+1] - adi[i]; 137 for (j=0; j<nzi; j++){ 138 row = *adj++; 139 pnz = pi_loc[row+1] - pi_loc[row]; 140 Jptr = pj_loc + pi_loc[row]; 141 /* add non-zero cols of P into the sorted linked list lnk */ 142 ierr = PetscLLAdd(pnz,Jptr,pN,nlnk,lnk,lnkbt);CHKERRQ(ierr); 143 apnz += nlnk; 144 } 145 /* off-diagonal portion of A */ 146 nzi = aoi[i+1] - aoi[i]; 147 for (j=0; j<nzi; j++){ 148 row = *aoj++; 149 pnz = pi_oth[row+1] - pi_oth[row]; 150 Jptr = pj_oth + pi_oth[row]; 151 ierr = PetscLLAdd(pnz,Jptr,pN,nlnk,lnk,lnkbt);CHKERRQ(ierr); 152 apnz += nlnk; 153 } 154 155 api[i+1] = api[i] + apnz; 156 if (ap->abnz_max < apnz) ap->abnz_max = apnz; 157 158 /* if free space is not available, double the total space in the list */ 159 if (current_space->local_remaining<apnz) { 160 ierr = GetMoreSpace(current_space->total_array_size,¤t_space);CHKERRQ(ierr); 161 } 162 163 /* Copy data into free space, then initialize lnk */ 164 ierr = PetscLLClean(pN,pN,apnz,lnk,current_space->array,lnkbt);CHKERRQ(ierr); 165 current_space->array += apnz; 166 current_space->local_used += apnz; 167 current_space->local_remaining -= apnz; 168 } 169 /* Allocate space for apj, initialize apj, and */ 170 /* destroy list of free space and other temporary array(s) */ 171 ierr = PetscMalloc((api[am]+1)*sizeof(PetscInt),&ap->abj);CHKERRQ(ierr); 172 apj = ap->abj; 173 ierr = MakeSpaceContiguous(&free_space,ap->abj);CHKERRQ(ierr); 174 175 /* determine symbolic Co=(p->B)^T*AP - send to others */ 176 /*----------------------------------------------------*/ 177 ierr = MatGetSymbolicTranspose_SeqAIJ(p->B,&poti,&potj);CHKERRQ(ierr); 178 179 /* then, compute symbolic Co = (p->B)^T*AP */ 180 pon = (p->B)->n; /* total num of rows to be sent to other processors 181 >= (num of nonzero rows of C_seq) - pn */ 182 ierr = PetscMalloc((pon+1)*sizeof(PetscInt),&coi);CHKERRQ(ierr); 183 coi[0] = 0; 184 185 /* set initial free space to be 3*pon*max( nnz(AP) per row) */ 186 nnz = 3*pon*ap->abnz_max + 1; 187 ierr = GetMoreSpace(nnz,&free_space); 188 current_space = free_space; 189 190 for (i=0; i<pon; i++) { 191 nnz = 0; 192 pnz = poti[i+1] - poti[i]; 193 j = pnz; 194 ptJ = potj + poti[i+1]; 195 while (j){/* assume cols are almost in increasing order, starting from its end saves computation */ 196 j--; ptJ--; 197 row = *ptJ; /* row of AP == col of Pot */ 198 apnz = api[row+1] - api[row]; 199 Jptr = apj + api[row]; 200 /* add non-zero cols of AP into the sorted linked list lnk */ 201 ierr = PetscLLAdd(apnz,Jptr,pN,nlnk,lnk,lnkbt);CHKERRQ(ierr); 202 nnz += nlnk; 203 } 204 205 /* If free space is not available, double the total space in the list */ 206 if (current_space->local_remaining<nnz) { 207 ierr = GetMoreSpace(current_space->total_array_size,¤t_space);CHKERRQ(ierr); 208 } 209 210 /* Copy data into free space, and zero out denserows */ 211 ierr = PetscLLClean(pN,pN,nnz,lnk,current_space->array,lnkbt);CHKERRQ(ierr); 212 current_space->array += nnz; 213 current_space->local_used += nnz; 214 current_space->local_remaining -= nnz; 215 coi[i+1] = coi[i] + nnz; 216 } 217 ierr = PetscMalloc((coi[pon]+1)*sizeof(PetscInt),&coj);CHKERRQ(ierr); 218 ierr = MakeSpaceContiguous(&free_space,coj);CHKERRQ(ierr); 219 ierr = MatRestoreSymbolicTranspose_SeqAIJ(p->B,&poti,&potj);CHKERRQ(ierr); 220 221 /* send j-array (coj) of Co to other processors */ 222 /*----------------------------------------------*/ 223 /* determine row ownership */ 224 ierr = PetscNew(Mat_Merge_SeqsToMPI,&merge);CHKERRQ(ierr); 225 ierr = PetscMapCreate(comm,&merge->rowmap);CHKERRQ(ierr); 226 ierr = PetscMapSetLocalSize(merge->rowmap,pn);CHKERRQ(ierr); 227 ierr = PetscMapSetType(merge->rowmap,MAP_MPI);CHKERRQ(ierr); 228 ierr = PetscMapGetGlobalRange(merge->rowmap,&owners);CHKERRQ(ierr); 229 230 /* determine the number of messages to send, their lengths */ 231 ierr = PetscMalloc(size*sizeof(PetscMPIInt),&len_si);CHKERRQ(ierr); 232 ierr = PetscMemzero(len_si,size*sizeof(PetscMPIInt));CHKERRQ(ierr); 233 ierr = PetscMalloc(size*sizeof(PetscMPIInt),&merge->len_s);CHKERRQ(ierr); 234 len_s = merge->len_s; 235 merge->nsend = 0; 236 237 ierr = PetscMalloc((size+2)*sizeof(PetscInt),&owners_co);CHKERRQ(ierr); 238 ierr = PetscMemzero(len_s,size*sizeof(PetscMPIInt));CHKERRQ(ierr); 239 240 proc = 0; 241 for (i=0; i<pon; i++){ 242 while (prmap[i] >= owners[proc+1]) proc++; 243 len_si[proc]++; /* num of rows in Co to be sent to [proc] */ 244 len_s[proc] += coi[i+1] - coi[i]; 245 } 246 247 len = 0; /* max length of buf_si[] */ 248 owners_co[0] = 0; 249 for (proc=0; proc<size; proc++){ 250 owners_co[proc+1] = owners_co[proc] + len_si[proc]; 251 if (len_si[proc]){ 252 merge->nsend++; 253 len_si[proc] = 2*(len_si[proc] + 1); 254 len += len_si[proc]; 255 } 256 } 257 258 /* determine the number and length of messages to receive for coi and coj */ 259 ierr = PetscGatherNumberOfMessages(comm,PETSC_NULL,len_s,&merge->nrecv);CHKERRQ(ierr); 260 ierr = PetscGatherMessageLengths2(comm,merge->nsend,merge->nrecv,len_s,len_si,&merge->id_r,&merge->len_r,&len_ri);CHKERRQ(ierr); 261 262 /* post the Irecv and Isend of coj */ 263 ierr = PetscObjectGetNewTag((PetscObject)merge->rowmap,&tag);CHKERRQ(ierr); 264 ierr = PetscPostIrecvInt(comm,tag,merge->nrecv,merge->id_r,merge->len_r,&buf_rj,&rwaits);CHKERRQ(ierr); 265 266 ierr = PetscMalloc((merge->nsend+1)*sizeof(MPI_Request),&swaits);CHKERRQ(ierr); 267 268 for (proc=0, k=0; proc<size; proc++){ 269 if (!len_s[proc]) continue; 270 i = owners_co[proc]; 271 ierr = MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tag,comm,swaits+k);CHKERRQ(ierr); 272 k++; 273 } 274 275 /* receives and sends of coj are complete */ 276 ierr = PetscMalloc(size*sizeof(MPI_Status),&sstatus);CHKERRQ(ierr); 277 i = merge->nrecv; 278 while (i--) { 279 ierr = MPI_Waitany(merge->nrecv,rwaits,&j,&rstatus);CHKERRQ(ierr); 280 } 281 ierr = PetscFree(rwaits);CHKERRQ(ierr); 282 if (merge->nsend){ 283 ierr = MPI_Waitall(merge->nsend,swaits,sstatus);CHKERRQ(ierr); 284 } 285 286 /* send and recv coi */ 287 /*-------------------*/ 288 ierr = PetscPostIrecvInt(comm,tag,merge->nrecv,merge->id_r,len_ri,&buf_ri,&rwaits);CHKERRQ(ierr); 289 290 ierr = PetscMalloc((len+1)*sizeof(PetscInt),&buf_s);CHKERRQ(ierr); 291 buf_si = buf_s; /* points to the beginning of k-th msg to be sent */ 292 for (proc=0,k=0; proc<size; proc++){ 293 if (!len_s[proc]) continue; 294 /* form outgoing message for i-structure: 295 buf_si[0]: nrows to be sent 296 [1:nrows]: row index (global) 297 [nrows+1:2*nrows+1]: i-structure index 298 */ 299 /*-------------------------------------------*/ 300 nrows = len_si[proc]/2 - 1; 301 buf_si_i = buf_si + nrows+1; 302 buf_si[0] = nrows; 303 buf_si_i[0] = 0; 304 nrows = 0; 305 for (i=owners_co[proc]; i<owners_co[proc+1]; i++){ 306 nzi = coi[i+1] - coi[i]; 307 buf_si_i[nrows+1] = buf_si_i[nrows] + nzi; /* i-structure */ 308 buf_si[nrows+1] =prmap[i] -owners[proc]; /* local row index */ 309 nrows++; 310 } 311 ierr = MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tag,comm,swaits+k);CHKERRQ(ierr); 312 k++; 313 buf_si += len_si[proc]; 314 } 315 i = merge->nrecv; 316 while (i--) { 317 ierr = MPI_Waitany(merge->nrecv,rwaits,&j,&rstatus);CHKERRQ(ierr); 318 } 319 ierr = PetscFree(rwaits);CHKERRQ(ierr); 320 if (merge->nsend){ 321 ierr = MPI_Waitall(merge->nsend,swaits,sstatus);CHKERRQ(ierr); 322 } 323 324 ierr = PetscLogInfo((PetscObject)A,"MatMerge_SeqsToMPI: nsend: %d, nrecv: %d\n",merge->nsend,merge->nrecv);CHKERRQ(ierr); 325 for (i=0; i<merge->nrecv; i++){ 326 ierr = PetscLogInfo((PetscObject)A,"MatMerge_SeqsToMPI: recv len_ri=%d, len_rj=%d from [%d]\n",len_ri[i],merge->len_r[i],merge->id_r[i]);CHKERRQ(ierr); 327 } 328 329 ierr = PetscFree(len_si);CHKERRQ(ierr); 330 ierr = PetscFree(len_ri);CHKERRQ(ierr); 331 ierr = PetscFree(swaits);CHKERRQ(ierr); 332 ierr = PetscFree(sstatus);CHKERRQ(ierr); 333 ierr = PetscFree(buf_s);CHKERRQ(ierr); 334 335 /* compute the local portion of C (mpi mat) */ 336 /*------------------------------------------*/ 337 ierr = MatGetSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj);CHKERRQ(ierr); 338 339 /* allocate bi array and free space for accumulating nonzero column info */ 340 ierr = PetscMalloc((pn+1)*sizeof(PetscInt),&bi);CHKERRQ(ierr); 341 bi[0] = 0; 342 343 /* set initial free space to be 3*pn*max( nnz(AP) per row) */ 344 nnz = 3*pn*ap->abnz_max + 1; 345 ierr = GetMoreSpace(nnz,&free_space); 346 current_space = free_space; 347 348 ierr = PetscMalloc((3*merge->nrecv+1)*sizeof(PetscInt**),&buf_ri_k);CHKERRQ(ierr); 349 nextrow = buf_ri_k + merge->nrecv; 350 nextci = nextrow + merge->nrecv; 351 for (k=0; k<merge->nrecv; k++){ 352 buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 353 nrows = *buf_ri_k[k]; 354 nextrow[k] = buf_ri_k[k] + 1; /* next row number of k-th recved i-structure */ 355 nextci[k] = buf_ri_k[k] + (nrows + 1);/* poins to the next i-structure of k-th recved i-structure */ 356 } 357 ierr = MatPreallocateInitialize(comm,pn,pn,dnz,onz);CHKERRQ(ierr); 358 for (i=0; i<pn; i++) { 359 /* add pdt[i,:]*AP into lnk */ 360 nnz = 0; 361 pnz = pdti[i+1] - pdti[i]; 362 j = pnz; 363 ptJ = pdtj + pdti[i+1]; 364 while (j){/* assume cols are almost in increasing order, starting from its end saves computation */ 365 j--; ptJ--; 366 row = *ptJ; /* row of AP == col of Pt */ 367 apnz = api[row+1] - api[row]; 368 Jptr = apj + api[row]; 369 /* add non-zero cols of AP into the sorted linked list lnk */ 370 ierr = PetscLLAdd(apnz,Jptr,pN,nlnk,lnk,lnkbt);CHKERRQ(ierr); 371 nnz += nlnk; 372 } 373 /* add received col data into lnk */ 374 for (k=0; k<merge->nrecv; k++){ /* k-th received message */ 375 if (i == *nextrow[k]) { /* i-th row */ 376 nzi = *(nextci[k]+1) - *nextci[k]; 377 Jptr = buf_rj[k] + *nextci[k]; 378 ierr = PetscLLAdd(nzi,Jptr,pN,nlnk,lnk,lnkbt);CHKERRQ(ierr); 379 nnz += nlnk; 380 nextrow[k]++; nextci[k]++; 381 } 382 } 383 384 /* if free space is not available, make more free space */ 385 if (current_space->local_remaining<nnz) { 386 ierr = GetMoreSpace(current_space->total_array_size,¤t_space);CHKERRQ(ierr); 387 } 388 /* copy data into free space, then initialize lnk */ 389 ierr = PetscLLClean(pN,pN,nnz,lnk,current_space->array,lnkbt);CHKERRQ(ierr); 390 ierr = MatPreallocateSet(i+owners[rank],nnz,current_space->array,dnz,onz);CHKERRQ(ierr); 391 current_space->array += nnz; 392 current_space->local_used += nnz; 393 current_space->local_remaining -= nnz; 394 bi[i+1] = bi[i] + nnz; 395 } 396 ierr = MatRestoreSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj);CHKERRQ(ierr); 397 ierr = PetscFree(buf_ri_k);CHKERRQ(ierr); 398 399 ierr = PetscMalloc((bi[pn]+1)*sizeof(PetscInt),&bj);CHKERRQ(ierr); 400 ierr = MakeSpaceContiguous(&free_space,bj);CHKERRQ(ierr); 401 ierr = PetscLLDestroy(lnk,lnkbt);CHKERRQ(ierr); 402 403 /* create symbolic parallel matrix B_mpi */ 404 /*---------------------------------------*/ 405 ierr = MatCreate(comm,pn,pn,PETSC_DETERMINE,PETSC_DETERMINE,&B_mpi);CHKERRQ(ierr); 406 ierr = MatSetType(B_mpi,MATMPIAIJ);CHKERRQ(ierr); 407 ierr = MatMPIAIJSetPreallocation(B_mpi,0,dnz,0,onz);CHKERRQ(ierr); 408 ierr = MatPreallocateFinalize(dnz,onz);CHKERRQ(ierr); 409 410 /* B_mpi is not ready for use - assembly will be done by MatPtAPNumeric() */ 411 B_mpi->assembled = PETSC_FALSE; 412 B_mpi->ops->destroy = MatDestroy_MPIAIJ_SeqsToMPI; 413 merge->bi = bi; 414 merge->bj = bj; 415 merge->coi = coi; 416 merge->coj = coj; 417 merge->buf_ri = buf_ri; 418 merge->buf_rj = buf_rj; 419 merge->owners_co = owners_co; 420 421 /* attach the supporting struct to B_mpi for reuse */ 422 ierr = PetscObjectContainerCreate(PETSC_COMM_SELF,&container);CHKERRQ(ierr); 423 ierr = PetscObjectContainerSetPointer(container,merge);CHKERRQ(ierr); 424 ierr = PetscObjectCompose((PetscObject)B_mpi,"MatMergeSeqsToMPI",(PetscObject)container);CHKERRQ(ierr); 425 *C = B_mpi; 426 427 PetscFunctionReturn(0); 428 } 429 430 #undef __FUNCT__ 431 #define __FUNCT__ "MatPtAPNumeric_MPIAIJ_MPIAIJ" 432 PetscErrorCode MatPtAPNumeric_MPIAIJ_MPIAIJ(Mat A,Mat P,Mat C) 433 { 434 PetscErrorCode ierr; 435 Mat_Merge_SeqsToMPI *merge; 436 Mat_MatMatMultMPI *ap; 437 PetscObjectContainer cont_merge,cont_ptap; 438 Mat_MPIAIJ *a=(Mat_MPIAIJ*)A->data,*p=(Mat_MPIAIJ*)P->data; 439 Mat_SeqAIJ *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data; 440 Mat_SeqAIJ *pd=(Mat_SeqAIJ*)(p->A)->data,*po=(Mat_SeqAIJ*)(p->B)->data; 441 Mat_SeqAIJ *p_loc,*p_oth; 442 PetscInt *adi=ad->i,*aoi=ao->i,*adj=ad->j,*aoj=ao->j,*apJ,nextp,flops=0; 443 PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*pJ,*pj; 444 PetscInt i,j,k,anz,pnz,apnz,nextap,row,*cj; 445 MatScalar *ada=ad->a,*aoa=ao->a,*apa,*pa,*ca,*pa_loc,*pa_oth; 446 PetscInt am=A->m,cm=C->m,pon=(p->B)->n; 447 MPI_Comm comm=C->comm; 448 PetscMPIInt size,rank,taga,*len_s; 449 PetscInt *owners,proc,nrows,**buf_ri_k,**nextrow,**nextci; 450 PetscInt **buf_ri,**buf_rj; 451 PetscInt cnz=0,*bj_i,*bi,*bj,bnz,nextcj; /* bi,bj,ba: local array of C(mpi mat) */ 452 MPI_Request *s_waits,*r_waits; 453 MPI_Status *status; 454 MatScalar **abuf_r,*ba_i,*pA,*coa,*ba; 455 PetscInt *api,*apj,*coi,*coj; 456 PetscInt *poJ=po->j,*pdJ=pd->j,pcstart=p->cstart,pcend=p->cend; 457 458 PetscFunctionBegin; 459 ierr = MPI_Comm_size(comm,&size);CHKERRQ(ierr); 460 ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr); 461 462 ierr = PetscObjectQuery((PetscObject)C,"MatMergeSeqsToMPI",(PetscObject *)&cont_merge);CHKERRQ(ierr); 463 if (cont_merge) { 464 ierr = PetscObjectContainerGetPointer(cont_merge,(void **)&merge);CHKERRQ(ierr); 465 } else { 466 SETERRQ(PETSC_ERR_ARG_WRONGSTATE, "Matrix C does not posses an object container"); 467 } 468 ierr = PetscObjectQuery((PetscObject)P,"Mat_MatMatMultMPI",(PetscObject *)&cont_ptap);CHKERRQ(ierr); 469 if (cont_ptap) { 470 ierr = PetscObjectContainerGetPointer(cont_ptap,(void **)&ap);CHKERRQ(ierr); 471 } else { 472 SETERRQ(PETSC_ERR_ARG_WRONGSTATE, "Matrix P does not posses an object container"); 473 } 474 475 /* get data from symbolic products */ 476 p_loc = (Mat_SeqAIJ*)(ap->B_loc)->data; 477 p_oth = (Mat_SeqAIJ*)(ap->B_oth)->data; 478 pi_loc=p_loc->i; pj_loc=p_loc->j; pJ=pj_loc; pa_loc=p_loc->a,pA=pa_loc; 479 pi_oth=p_oth->i; pj_oth=p_oth->j; pa_oth=p_oth->a; 480 481 coi = merge->coi; coj = merge->coj; 482 ierr = PetscMalloc((coi[pon]+1)*sizeof(MatScalar),&coa);CHKERRQ(ierr); 483 ierr = PetscMemzero(coa,coi[pon]*sizeof(MatScalar));CHKERRQ(ierr); 484 485 bi = merge->bi; bj = merge->bj; 486 ierr = PetscMapGetGlobalRange(merge->rowmap,&owners);CHKERRQ(ierr); 487 ierr = PetscMalloc((bi[cm]+1)*sizeof(MatScalar),&ba);CHKERRQ(ierr); 488 ierr = PetscMemzero(ba,bi[cm]*sizeof(MatScalar));CHKERRQ(ierr); 489 490 /* get data from symbolic A*P */ 491 ierr = PetscMalloc((ap->abnz_max+1)*sizeof(MatScalar),&apa);CHKERRQ(ierr); 492 ierr = PetscMemzero(apa,ap->abnz_max*sizeof(MatScalar));CHKERRQ(ierr); 493 494 /* compute numeric C_seq=P_loc^T*A_loc*P */ 495 api = ap->abi; apj = ap->abj; 496 for (i=0;i<am;i++) { 497 /* form i-th sparse row of A*P */ 498 apnz = api[i+1] - api[i]; 499 apJ = apj + api[i]; 500 /* diagonal portion of A */ 501 anz = adi[i+1] - adi[i]; 502 for (j=0;j<anz;j++) { 503 row = *adj++; 504 pnz = pi_loc[row+1] - pi_loc[row]; 505 pj = pj_loc + pi_loc[row]; 506 pa = pa_loc + pi_loc[row]; 507 nextp = 0; 508 for (k=0; nextp<pnz; k++) { 509 if (apJ[k] == pj[nextp]) { /* col of AP == col of P */ 510 apa[k] += (*ada)*pa[nextp++]; 511 } 512 } 513 flops += 2*pnz; 514 ada++; 515 } 516 /* off-diagonal portion of A */ 517 anz = aoi[i+1] - aoi[i]; 518 for (j=0; j<anz; j++) { 519 row = *aoj++; 520 pnz = pi_oth[row+1] - pi_oth[row]; 521 pj = pj_oth + pi_oth[row]; 522 pa = pa_oth + pi_oth[row]; 523 nextp = 0; 524 for (k=0; nextp<pnz; k++) { 525 if (apJ[k] == pj[nextp]) { /* col of AP == col of P */ 526 apa[k] += (*aoa)*pa[nextp++]; 527 } 528 } 529 flops += 2*pnz; 530 aoa++; 531 } 532 533 /* Compute P_loc[i,:]^T*AP[i,:] using outer product */ 534 pnz = pi_loc[i+1] - pi_loc[i]; 535 for (j=0; j<pnz; j++) { 536 nextap = 0; 537 row = *pJ++; /* global index */ 538 if (row < pcstart || row >=pcend) { /* put the value into Co */ 539 cj = coj + coi[*poJ]; 540 ca = coa + coi[*poJ++]; 541 } else { /* put the value into Cd */ 542 cj = bj + bi[*pdJ]; 543 ca = ba + bi[*pdJ++]; 544 } 545 for (k=0; nextap<apnz; k++) { 546 if (cj[k]==apJ[nextap]) ca[k] += (*pA)*apa[nextap++]; 547 } 548 flops += 2*apnz; 549 pA++; 550 } 551 552 /* zero the current row info for A*P */ 553 ierr = PetscMemzero(apa,apnz*sizeof(MatScalar));CHKERRQ(ierr); 554 } 555 ierr = PetscFree(apa);CHKERRQ(ierr); 556 557 /* send and recv matrix values */ 558 /*-----------------------------*/ 559 buf_ri = merge->buf_ri; 560 buf_rj = merge->buf_rj; 561 len_s = merge->len_s; 562 ierr = PetscObjectGetNewTag((PetscObject)merge->rowmap,&taga);CHKERRQ(ierr); 563 ierr = PetscPostIrecvScalar(comm,taga,merge->nrecv,merge->id_r,merge->len_r,&abuf_r,&r_waits);CHKERRQ(ierr); 564 565 ierr = PetscMalloc((merge->nsend+1)*sizeof(MPI_Request),&s_waits);CHKERRQ(ierr); 566 for (proc=0,k=0; proc<size; proc++){ 567 if (!len_s[proc]) continue; 568 i = merge->owners_co[proc]; 569 ierr = MPI_Isend(coa+coi[i],len_s[proc],MPIU_MATSCALAR,proc,taga,comm,s_waits+k);CHKERRQ(ierr); 570 k++; 571 } 572 ierr = PetscMalloc(size*sizeof(MPI_Status),&status);CHKERRQ(ierr); 573 ierr = MPI_Waitall(merge->nrecv,r_waits,status);CHKERRQ(ierr); 574 ierr = MPI_Waitall(merge->nsend,s_waits,status);CHKERRQ(ierr); 575 ierr = PetscFree(status);CHKERRQ(ierr); 576 577 ierr = PetscFree(s_waits);CHKERRQ(ierr); 578 ierr = PetscFree(r_waits);CHKERRQ(ierr); 579 ierr = PetscFree(coa);CHKERRQ(ierr); 580 581 /* insert local and received values into C */ 582 /*-----------------------------------------*/ 583 ierr = PetscMalloc((3*merge->nrecv+1)*sizeof(PetscInt**),&buf_ri_k);CHKERRQ(ierr); 584 nextrow = buf_ri_k + merge->nrecv; 585 nextci = nextrow + merge->nrecv; 586 587 for (k=0; k<merge->nrecv; k++){ 588 buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 589 nrows = *(buf_ri_k[k]); 590 nextrow[k] = buf_ri_k[k]+1; /* next row number of k-th recved i-structure */ 591 nextci[k] = buf_ri_k[k] + (nrows + 1);/* poins to the next i-structure of k-th recved i-structure */ 592 } 593 594 for (i=0; i<cm; i++) { 595 row = owners[rank] + i; /* global row index of C_seq */ 596 bj_i = bj + bi[i]; /* col indices of the i-th row of C */ 597 ba_i = ba + bi[i]; 598 bnz = bi[i+1] - bi[i]; 599 /* add received vals into ba */ 600 for (k=0; k<merge->nrecv; k++){ /* k-th received message */ 601 /* i-th row */ 602 if (i == *nextrow[k]) { 603 cnz = *(nextci[k]+1) - *nextci[k]; 604 cj = buf_rj[k] + *(nextci[k]); 605 ca = abuf_r[k] + *(nextci[k]); 606 nextcj = 0; 607 for (j=0; nextcj<cnz; j++){ 608 if (bj_i[j] == cj[nextcj]){ /* bcol == ccol */ 609 ba_i[j] += ca[nextcj++]; 610 } 611 } 612 nextrow[k]++; nextci[k]++; 613 } 614 } 615 ierr = MatSetValues(C,1,&row,bnz,bj_i,ba_i,INSERT_VALUES);CHKERRQ(ierr); 616 flops += 2*cnz; 617 } 618 ierr = MatSetOption(C,MAT_COLUMNS_SORTED);CHKERRQ(ierr); /* -cause delay? */ 619 ierr = MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 620 ierr = MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr); 621 622 ierr = PetscFree(ba);CHKERRQ(ierr); 623 ierr = PetscFree(abuf_r);CHKERRQ(ierr); 624 ierr = PetscFree(buf_ri_k);CHKERRQ(ierr); 625 ierr = PetscLogFlops(flops);CHKERRQ(ierr); 626 627 PetscFunctionReturn(0); 628 } 629