1be1d678aSKris Buschelman 22499ec78SHong Zhang /* 32499ec78SHong Zhang Defines matrix-matrix product routines for pairs of MPIAIJ matrices 42499ec78SHong Zhang C = A * B 52499ec78SHong Zhang */ 6c6db04a5SJed Brown #include <../src/mat/impls/aij/seq/aij.h> /*I "petscmat.h" I*/ 7c6db04a5SJed Brown #include <../src/mat/utils/freespace.h> 8c6db04a5SJed Brown #include <../src/mat/impls/aij/mpi/mpiaij.h> 9c6db04a5SJed Brown #include <petscbt.h> 10c6db04a5SJed Brown #include <../src/mat/impls/dense/mpi/mpidense.h> 11af0996ceSBarry Smith #include <petsc/private/vecimpl.h> 1297929ea7SJunchao Zhang #include <petsc/private/sfimpl.h> 131634b032SHong Zhang 145e5acdf2Sstefano_zampini #if defined(PETSC_HAVE_HYPRE) 154222ddf1SHong Zhang PETSC_INTERN PetscErrorCode MatMatMultSymbolic_AIJ_AIJ_wHYPRE(Mat,Mat,PetscReal,Mat); 165e5acdf2Sstefano_zampini #endif 175e5acdf2Sstefano_zampini 18*8a9c020eSBarry Smith PETSC_INTERN PetscErrorCode MatProductSymbolic_ABt_MPIAIJ_MPIAIJ(Mat C) 19*8a9c020eSBarry Smith { 20*8a9c020eSBarry Smith Mat_Product *product = C->product; 21*8a9c020eSBarry Smith Mat B=product->B; 22*8a9c020eSBarry Smith 23*8a9c020eSBarry Smith PetscFunctionBegin; 24*8a9c020eSBarry Smith PetscCall(MatTranspose(B,MAT_INITIAL_MATRIX,&product->B)); 25*8a9c020eSBarry Smith PetscCall(MatDestroy(&B)); 26*8a9c020eSBarry Smith PetscCall(MatProductSymbolic_AB_MPIAIJ_MPIAIJ(C)); 27*8a9c020eSBarry Smith PetscFunctionReturn(0); 28*8a9c020eSBarry Smith } 29*8a9c020eSBarry Smith 304222ddf1SHong Zhang PETSC_INTERN PetscErrorCode MatProductSymbolic_AB_MPIAIJ_MPIAIJ(Mat C) 312499ec78SHong Zhang { 324222ddf1SHong Zhang Mat_Product *product = C->product; 334222ddf1SHong Zhang Mat A=product->A,B=product->B; 344222ddf1SHong Zhang MatProductAlgorithm alg=product->alg; 354222ddf1SHong Zhang PetscReal fill=product->fill; 365e68f438SHong Zhang PetscBool flg; 372499ec78SHong Zhang 382499ec78SHong Zhang PetscFunctionBegin; 394222ddf1SHong Zhang /* scalable */ 409566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(alg,"scalable",&flg)); 414222ddf1SHong Zhang if (flg) { 429566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ(A,B,fill,C)); 436718818eSStefano Zampini PetscFunctionReturn(0); 440fc8cf34SHong Zhang } 457d0a89b7SHong Zhang 464222ddf1SHong Zhang /* nonscalable */ 479566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(alg,"nonscalable",&flg)); 484222ddf1SHong Zhang if (flg) { 499566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(A,B,fill,C)); 506718818eSStefano Zampini PetscFunctionReturn(0); 514222ddf1SHong Zhang } 524222ddf1SHong Zhang 534222ddf1SHong Zhang /* seqmpi */ 549566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(alg,"seqmpi",&flg)); 554222ddf1SHong Zhang if (flg) { 569566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ_seqMPI(A,B,fill,C)); 576718818eSStefano Zampini PetscFunctionReturn(0); 584222ddf1SHong Zhang } 594222ddf1SHong Zhang 604e84afc0SStefano Zampini /* backend general code */ 619566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(alg,"backend",&flg)); 624e84afc0SStefano Zampini if (flg) { 639566063dSJacob Faibussowitsch PetscCall(MatProductSymbolic_MPIAIJBACKEND(C)); 644e84afc0SStefano Zampini PetscFunctionReturn(0); 654e84afc0SStefano Zampini } 664e84afc0SStefano Zampini 674222ddf1SHong Zhang #if defined(PETSC_HAVE_HYPRE) 689566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(alg,"hypre",&flg)); 694222ddf1SHong Zhang if (flg) { 709566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_AIJ_AIJ_wHYPRE(A,B,fill,C)); 714222ddf1SHong Zhang PetscFunctionReturn(0); 724222ddf1SHong Zhang } 734222ddf1SHong Zhang #endif 744222ddf1SHong Zhang SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_SUP,"Mat Product Algorithm is not supported"); 752499ec78SHong Zhang } 762499ec78SHong Zhang 776718818eSStefano Zampini PetscErrorCode MatDestroy_MPIAIJ_MatMatMult(void *data) 78598bc09dSHong Zhang { 796718818eSStefano Zampini Mat_APMPI *ptap = (Mat_APMPI*)data; 80598bc09dSHong Zhang 81598bc09dSHong Zhang PetscFunctionBegin; 829566063dSJacob Faibussowitsch PetscCall(PetscFree2(ptap->startsj_s,ptap->startsj_r)); 839566063dSJacob Faibussowitsch PetscCall(PetscFree(ptap->bufa)); 849566063dSJacob Faibussowitsch PetscCall(MatDestroy(&ptap->P_loc)); 859566063dSJacob Faibussowitsch PetscCall(MatDestroy(&ptap->P_oth)); 869566063dSJacob Faibussowitsch PetscCall(MatDestroy(&ptap->Pt)); 879566063dSJacob Faibussowitsch PetscCall(PetscFree(ptap->api)); 889566063dSJacob Faibussowitsch PetscCall(PetscFree(ptap->apj)); 899566063dSJacob Faibussowitsch PetscCall(PetscFree(ptap->apa)); 909566063dSJacob Faibussowitsch PetscCall(PetscFree(ptap)); 91598bc09dSHong Zhang PetscFunctionReturn(0); 92598bc09dSHong Zhang } 93598bc09dSHong Zhang 94f8efd71cSHong Zhang PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable(Mat A,Mat P,Mat C) 95598bc09dSHong Zhang { 964ae0847bSHong Zhang Mat_MPIAIJ *a =(Mat_MPIAIJ*)A->data,*c=(Mat_MPIAIJ*)C->data; 97598bc09dSHong Zhang Mat_SeqAIJ *ad =(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data; 98598bc09dSHong Zhang Mat_SeqAIJ *cd =(Mat_SeqAIJ*)(c->A)->data,*co=(Mat_SeqAIJ*)(c->B)->data; 99c12557a1SHong Zhang PetscScalar *cda=cd->a,*coa=co->a; 100598bc09dSHong Zhang Mat_SeqAIJ *p_loc,*p_oth; 1019ce11a7cSHong Zhang PetscScalar *apa,*ca; 102c12557a1SHong Zhang PetscInt cm =C->rmap->n; 1036718818eSStefano Zampini Mat_APMPI *ptap; 104c12557a1SHong Zhang PetscInt *api,*apj,*apJ,i,k; 10529825010SHong Zhang PetscInt cstart=C->cmap->rstart; 106598bc09dSHong Zhang PetscInt cdnz,conz,k0,k1; 107ce496241SStefano Zampini const PetscScalar *dummy; 1089467f45fSHong Zhang MPI_Comm comm; 1099467f45fSHong Zhang PetscMPIInt size; 110598bc09dSHong Zhang 111598bc09dSHong Zhang PetscFunctionBegin; 1126718818eSStefano Zampini MatCheckProduct(C,3); 1136718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 11428b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 1159566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 1169566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 1179467f45fSHong Zhang 1182c71b3e2SJacob Faibussowitsch PetscCheckFalse(!ptap->P_oth && size>1,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"AP cannot be reused. Do not call MatProductClear()"); 11980da8df7SHong Zhang 12065e4b4d4SStefano Zampini /* flag CPU mask for C */ 12165e4b4d4SStefano Zampini #if defined(PETSC_HAVE_DEVICE) 12265e4b4d4SStefano Zampini if (C->offloadmask != PETSC_OFFLOAD_UNALLOCATED) C->offloadmask = PETSC_OFFLOAD_CPU; 12365e4b4d4SStefano Zampini if (c->A->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->A->offloadmask = PETSC_OFFLOAD_CPU; 12465e4b4d4SStefano Zampini if (c->B->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->B->offloadmask = PETSC_OFFLOAD_CPU; 12565e4b4d4SStefano Zampini #endif 12665e4b4d4SStefano Zampini 127a1a4e84aSHong Zhang /* 1) get P_oth = ptap->P_oth and P_loc = ptap->P_loc */ 128598bc09dSHong Zhang /*-----------------------------------------------------*/ 129a1a4e84aSHong Zhang /* update numerical values of P_oth and P_loc */ 1309566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_REUSE_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 1319566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_REUSE_MATRIX,&ptap->P_loc)); 132598bc09dSHong Zhang 133598bc09dSHong Zhang /* 2) compute numeric C_loc = A_loc*P = Ad*P_loc + Ao*P_oth */ 134598bc09dSHong Zhang /*----------------------------------------------------------*/ 135598bc09dSHong Zhang /* get data from symbolic products */ 136a1a4e84aSHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 137c12557a1SHong Zhang p_oth = NULL; 1389467f45fSHong Zhang if (size >1) { 1399467f45fSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 1409467f45fSHong Zhang } 141598bc09dSHong Zhang 142598bc09dSHong Zhang /* get apa for storing dense row A[i,:]*P */ 143598bc09dSHong Zhang apa = ptap->apa; 144598bc09dSHong Zhang 14529825010SHong Zhang api = ptap->api; 14629825010SHong Zhang apj = ptap->apj; 147ce496241SStefano Zampini /* trigger copy to CPU */ 1489566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->A,&dummy)); 1499566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->A,&dummy)); 1509566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->B,&dummy)); 1519566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->B,&dummy)); 152598bc09dSHong Zhang for (i=0; i<cm; i++) { 153c12557a1SHong Zhang /* compute apa = A[i,:]*P */ 154e497d3c8SHong Zhang AProw_nonscalable(i,ad,ao,p_loc,p_oth,apa); 155598bc09dSHong Zhang 156598bc09dSHong Zhang /* set values in C */ 157598bc09dSHong Zhang apJ = apj + api[i]; 158598bc09dSHong Zhang cdnz = cd->i[i+1] - cd->i[i]; 159598bc09dSHong Zhang conz = co->i[i+1] - co->i[i]; 160598bc09dSHong Zhang 161910cf402Sprj- /* 1st off-diagonal part of C */ 162598bc09dSHong Zhang ca = coa + co->i[i]; 163598bc09dSHong Zhang k = 0; 164598bc09dSHong Zhang for (k0=0; k0<conz; k0++) { 165598bc09dSHong Zhang if (apJ[k] >= cstart) break; 166598bc09dSHong Zhang ca[k0] = apa[apJ[k]]; 1675cab4f04SHong Zhang apa[apJ[k++]] = 0.0; 168598bc09dSHong Zhang } 169598bc09dSHong Zhang 170598bc09dSHong Zhang /* diagonal part of C */ 171598bc09dSHong Zhang ca = cda + cd->i[i]; 172598bc09dSHong Zhang for (k1=0; k1<cdnz; k1++) { 173598bc09dSHong Zhang ca[k1] = apa[apJ[k]]; 1745cab4f04SHong Zhang apa[apJ[k++]] = 0.0; 175598bc09dSHong Zhang } 176598bc09dSHong Zhang 177910cf402Sprj- /* 2nd off-diagonal part of C */ 178598bc09dSHong Zhang ca = coa + co->i[i]; 179598bc09dSHong Zhang for (; k0<conz; k0++) { 180598bc09dSHong Zhang ca[k0] = apa[apJ[k]]; 1815cab4f04SHong Zhang apa[apJ[k++]] = 0.0; 182598bc09dSHong Zhang } 183598bc09dSHong Zhang } 1849566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 1859566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 186598bc09dSHong Zhang PetscFunctionReturn(0); 187598bc09dSHong Zhang } 188598bc09dSHong Zhang 1894222ddf1SHong Zhang PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(Mat A,Mat P,PetscReal fill,Mat C) 190598bc09dSHong Zhang { 191ce94432eSBarry Smith MPI_Comm comm; 1929467f45fSHong Zhang PetscMPIInt size; 1933cdca5ebSHong Zhang Mat_APMPI *ptap; 1940298fd71SBarry Smith PetscFreeSpaceList free_space=NULL,current_space=NULL; 1956718818eSStefano Zampini Mat_MPIAIJ *a=(Mat_MPIAIJ*)A->data; 196bfcd1a73SHong Zhang Mat_SeqAIJ *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_loc,*p_oth; 1977dbaa2c6Sandi selinger PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*dnz,*onz; 1984ae0847bSHong Zhang PetscInt *adi=ad->i,*adj=ad->j,*aoi=ao->i,*aoj=ao->j,rstart=A->rmap->rstart; 199d14fa243SHong Zhang PetscInt *lnk,i,pnz,row,*api,*apj,*Jptr,apnz,nspacedouble=0,j,nzi; 2005cab4f04SHong Zhang PetscInt am=A->rmap->n,pN=P->cmap->N,pn=P->cmap->n,pm=P->rmap->n; 201598bc09dSHong Zhang PetscBT lnkbt; 202bfcd1a73SHong Zhang PetscReal afill; 203b92f168fSBarry Smith MatType mtype; 204598bc09dSHong Zhang 205598bc09dSHong Zhang PetscFunctionBegin; 2066718818eSStefano Zampini MatCheckProduct(C,4); 20728b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Extra product struct not empty"); 2089566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 2099566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 2109467f45fSHong Zhang 2113cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 2129566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 213598bc09dSHong Zhang 214598bc09dSHong Zhang /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 2159566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 21619f0e57aSHong Zhang 217598bc09dSHong Zhang /* get P_loc by taking all local rows of P */ 2189566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc)); 219598bc09dSHong Zhang 220a1a4e84aSHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 221598bc09dSHong Zhang pi_loc = p_loc->i; pj_loc = p_loc->j; 2229467f45fSHong Zhang if (size > 1) { 2239467f45fSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 224598bc09dSHong Zhang pi_oth = p_oth->i; pj_oth = p_oth->j; 22520e3dc75SHong Zhang } else { 22620e3dc75SHong Zhang p_oth = NULL; 22720e3dc75SHong Zhang pi_oth = NULL; pj_oth = NULL; 2289467f45fSHong Zhang } 229598bc09dSHong Zhang 230598bc09dSHong Zhang /* first, compute symbolic AP = A_loc*P = A_diag*P_loc + A_off*P_oth */ 231598bc09dSHong Zhang /*-------------------------------------------------------------------*/ 2329566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&api)); 233a1a4e84aSHong Zhang ptap->api = api; 234598bc09dSHong Zhang api[0] = 0; 235598bc09dSHong Zhang 2365cab4f04SHong Zhang /* create and initialize a linked list */ 2379566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate(pN,pN,&lnk,&lnkbt)); 238598bc09dSHong Zhang 239bfcd1a73SHong Zhang /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */ 2409566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(adi[am],PetscIntSumTruncate(aoi[am],pi_loc[pm]))),&free_space)); 241598bc09dSHong Zhang current_space = free_space; 242598bc09dSHong Zhang 243d0609cedSBarry Smith MatPreallocateBegin(comm,am,pn,dnz,onz); 244598bc09dSHong Zhang for (i=0; i<am; i++) { 245598bc09dSHong Zhang /* diagonal portion of A */ 246598bc09dSHong Zhang nzi = adi[i+1] - adi[i]; 247598bc09dSHong Zhang for (j=0; j<nzi; j++) { 248598bc09dSHong Zhang row = *adj++; 249598bc09dSHong Zhang pnz = pi_loc[row+1] - pi_loc[row]; 250598bc09dSHong Zhang Jptr = pj_loc + pi_loc[row]; 251598bc09dSHong Zhang /* add non-zero cols of P into the sorted linked list lnk */ 2529566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(pnz,Jptr,lnk,lnkbt)); 253598bc09dSHong Zhang } 254598bc09dSHong Zhang /* off-diagonal portion of A */ 255598bc09dSHong Zhang nzi = aoi[i+1] - aoi[i]; 256598bc09dSHong Zhang for (j=0; j<nzi; j++) { 257598bc09dSHong Zhang row = *aoj++; 258598bc09dSHong Zhang pnz = pi_oth[row+1] - pi_oth[row]; 259598bc09dSHong Zhang Jptr = pj_oth + pi_oth[row]; 2609566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(pnz,Jptr,lnk,lnkbt)); 261598bc09dSHong Zhang } 2628c78258cSHong Zhang /* add possible missing diagonal entry */ 2638c78258cSHong Zhang if (C->force_diagonals) { 2648c78258cSHong Zhang j = i + rstart; /* column index */ 2659566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(1,&j,lnk,lnkbt)); 2668c78258cSHong Zhang } 267598bc09dSHong Zhang 268d14fa243SHong Zhang apnz = lnk[0]; 269598bc09dSHong Zhang api[i+1] = api[i] + apnz; 270598bc09dSHong Zhang 271598bc09dSHong Zhang /* if free space is not available, double the total space in the list */ 272598bc09dSHong Zhang if (current_space->local_remaining<apnz) { 2739566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(apnz,current_space->total_array_size),¤t_space)); 274598bc09dSHong Zhang nspacedouble++; 275598bc09dSHong Zhang } 276598bc09dSHong Zhang 277598bc09dSHong Zhang /* Copy data into free space, then initialize lnk */ 2789566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean(pN,apnz,current_space->array,lnk,lnkbt)); 2799566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+rstart,apnz,current_space->array,dnz,onz)); 2802205254eSKarl Rupp 281598bc09dSHong Zhang current_space->array += apnz; 282598bc09dSHong Zhang current_space->local_used += apnz; 283598bc09dSHong Zhang current_space->local_remaining -= apnz; 284598bc09dSHong Zhang } 285598bc09dSHong Zhang 286598bc09dSHong Zhang /* Allocate space for apj, initialize apj, and */ 287598bc09dSHong Zhang /* destroy list of free space and other temporary array(s) */ 2889566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(api[am]+1,&ptap->apj)); 289a1a4e84aSHong Zhang apj = ptap->apj; 2909566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,ptap->apj)); 2919566063dSJacob Faibussowitsch PetscCall(PetscLLDestroy(lnk,lnkbt)); 292598bc09dSHong Zhang 293f84c536bSHong Zhang /* malloc apa to store dense row A[i,:]*P */ 2949566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(pN,&ptap->apa)); 295f84c536bSHong Zhang 2964222ddf1SHong Zhang /* set and assemble symbolic parallel matrix C */ 2974222ddf1SHong Zhang /*---------------------------------------------*/ 2989566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,am,pn,PETSC_DETERMINE,PETSC_DETERMINE)); 2999566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,P)); 300a2f3521dSMark F. Adams 3019566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 3029566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 3039566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 304d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 30517f75172Sandi selinger 3069566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api)); 3079566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 3089566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 3099566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 310598bc09dSHong Zhang 3114222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable; 3124222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AB; 313598bc09dSHong Zhang 3144222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 3156718818eSStefano Zampini C->product->data = ptap; 3166718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_MatMatMult; 317598bc09dSHong Zhang 318bfcd1a73SHong Zhang /* set MatInfo */ 319118727c9SMark F. Adams afill = (PetscReal)api[am]/(adi[am]+aoi[am]+pi_loc[pm]+1) + 1.e-5; 320bfcd1a73SHong Zhang if (afill < 1.0) afill = 1.0; 3214222ddf1SHong Zhang C->info.mallocs = nspacedouble; 3224222ddf1SHong Zhang C->info.fill_ratio_given = fill; 3234222ddf1SHong Zhang C->info.fill_ratio_needed = afill; 324bfcd1a73SHong Zhang 325bfcd1a73SHong Zhang #if defined(PETSC_USE_INFO) 326bfcd1a73SHong Zhang if (api[am]) { 3279566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 3289566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n",(double)afill)); 329bfcd1a73SHong Zhang } else { 3309566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 331bfcd1a73SHong Zhang } 332bfcd1a73SHong Zhang #endif 333598bc09dSHong Zhang PetscFunctionReturn(0); 334598bc09dSHong Zhang } 335598bc09dSHong Zhang 3364222ddf1SHong Zhang /* ------------------------------------------------------- */ 3376718818eSStefano Zampini static PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIDense(Mat,Mat,PetscReal,Mat); 3386718818eSStefano Zampini static PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIDense(Mat,Mat,Mat); 3396718818eSStefano Zampini 3404222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense_AB(Mat C) 3419123193aSHong Zhang { 3424222ddf1SHong Zhang Mat_Product *product = C->product; 3434222ddf1SHong Zhang Mat A = product->A,B=product->B; 3449123193aSHong Zhang 3459123193aSHong Zhang PetscFunctionBegin; 3464222ddf1SHong Zhang if (A->cmap->rstart != B->rmap->rstart || A->cmap->rend != B->rmap->rend) 34798921bdaSJacob Faibussowitsch SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, (%" PetscInt_FMT ", %" PetscInt_FMT ") != (%" PetscInt_FMT ",%" PetscInt_FMT ")",A->cmap->rstart,A->cmap->rend,B->rmap->rstart,B->rmap->rend); 3484222ddf1SHong Zhang 3494222ddf1SHong Zhang C->ops->matmultsymbolic = MatMatMultSymbolic_MPIAIJ_MPIDense; 3504222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AB; 3519123193aSHong Zhang PetscFunctionReturn(0); 3529123193aSHong Zhang } 3534222ddf1SHong Zhang /* -------------------------------------------------------------------- */ 3544222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense_AtB(Mat C) 3554222ddf1SHong Zhang { 3564222ddf1SHong Zhang Mat_Product *product = C->product; 3574222ddf1SHong Zhang Mat A = product->A,B=product->B; 3584222ddf1SHong Zhang 3594222ddf1SHong Zhang PetscFunctionBegin; 3604222ddf1SHong Zhang if (A->rmap->rstart != B->rmap->rstart || A->rmap->rend != B->rmap->rend) 36198921bdaSJacob Faibussowitsch SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, (%" PetscInt_FMT ", %" PetscInt_FMT ") != (%" PetscInt_FMT ",%" PetscInt_FMT ")",A->rmap->rstart,A->rmap->rend,B->rmap->rstart,B->rmap->rend); 3624222ddf1SHong Zhang 3634222ddf1SHong Zhang C->ops->transposematmultsymbolic = MatTransposeMatMultSymbolic_MPIAIJ_MPIDense; 3644222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AtB; 3654222ddf1SHong Zhang PetscFunctionReturn(0); 3664222ddf1SHong Zhang } 3674222ddf1SHong Zhang 3684222ddf1SHong Zhang /* --------------------------------------------------------------------- */ 3694222ddf1SHong Zhang PETSC_INTERN PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense(Mat C) 3704222ddf1SHong Zhang { 3714222ddf1SHong Zhang Mat_Product *product = C->product; 3724222ddf1SHong Zhang 3734222ddf1SHong Zhang PetscFunctionBegin; 3744222ddf1SHong Zhang switch (product->type) { 3754222ddf1SHong Zhang case MATPRODUCT_AB: 3769566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_MPIDense_AB(C)); 3774222ddf1SHong Zhang break; 3784222ddf1SHong Zhang case MATPRODUCT_AtB: 3799566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_MPIDense_AtB(C)); 3804222ddf1SHong Zhang break; 3814222ddf1SHong Zhang default: 3826718818eSStefano Zampini break; 3834222ddf1SHong Zhang } 3844222ddf1SHong Zhang PetscFunctionReturn(0); 3854222ddf1SHong Zhang } 3864222ddf1SHong Zhang /* ------------------------------------------------------- */ 3879123193aSHong Zhang 3884324174eSBarry Smith typedef struct { 3895ea7661aSPierre Jolivet Mat workB,workB1; 3904324174eSBarry Smith MPI_Request *rwaits,*swaits; 39187753ddeSHong Zhang PetscInt nsends,nrecvs; 39287753ddeSHong Zhang MPI_Datatype *stype,*rtype; 3936718818eSStefano Zampini PetscInt blda; 3944324174eSBarry Smith } MPIAIJ_MPIDense; 3954324174eSBarry Smith 39696b95a6bSBarry Smith PetscErrorCode MatMPIAIJ_MPIDenseDestroy(void *ctx) 3974324174eSBarry Smith { 3984324174eSBarry Smith MPIAIJ_MPIDense *contents = (MPIAIJ_MPIDense*)ctx; 39987753ddeSHong Zhang PetscInt i; 4004324174eSBarry Smith 4014324174eSBarry Smith PetscFunctionBegin; 4029566063dSJacob Faibussowitsch PetscCall(MatDestroy(&contents->workB)); 4039566063dSJacob Faibussowitsch PetscCall(MatDestroy(&contents->workB1)); 40487753ddeSHong Zhang for (i=0; i<contents->nsends; i++) { 4059566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&contents->stype[i])); 40687753ddeSHong Zhang } 40787753ddeSHong Zhang for (i=0; i<contents->nrecvs; i++) { 4089566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&contents->rtype[i])); 40987753ddeSHong Zhang } 4109566063dSJacob Faibussowitsch PetscCall(PetscFree4(contents->stype,contents->rtype,contents->rwaits,contents->swaits)); 4119566063dSJacob Faibussowitsch PetscCall(PetscFree(contents)); 4124324174eSBarry Smith PetscFunctionReturn(0); 4134324174eSBarry Smith } 4144324174eSBarry Smith 4156718818eSStefano Zampini static PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIDense(Mat A,Mat B,PetscReal fill,Mat C) 4169123193aSHong Zhang { 417f32d5d43SBarry Smith Mat_MPIAIJ *aij=(Mat_MPIAIJ*)A->data; 418b45e3bf4SStefano Zampini PetscInt nz=aij->B->cmap->n,nsends,nrecvs,i,nrows_to,j,blda,m,M,n,N; 4194324174eSBarry Smith MPIAIJ_MPIDense *contents; 4204324174eSBarry Smith VecScatter ctx=aij->Mvctx; 4215ea7661aSPierre Jolivet PetscInt Am=A->rmap->n,Bm=B->rmap->n,BN=B->cmap->N,Bbn,Bbn1,bs,nrows_from,numBb; 42287753ddeSHong Zhang MPI_Comm comm; 42387753ddeSHong Zhang MPI_Datatype type1,*stype,*rtype; 42487753ddeSHong Zhang const PetscInt *sindices,*sstarts,*rstarts; 42587753ddeSHong Zhang PetscMPIInt *disp; 4266718818eSStefano Zampini PetscBool cisdense; 4279123193aSHong Zhang 4289123193aSHong Zhang PetscFunctionBegin; 4296718818eSStefano Zampini MatCheckProduct(C,4); 43028b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data not empty"); 4319566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 4329566063dSJacob Faibussowitsch PetscCall(PetscObjectBaseTypeCompare((PetscObject)C,MATMPIDENSE,&cisdense)); 4336718818eSStefano Zampini if (!cisdense) { 4349566063dSJacob Faibussowitsch PetscCall(MatSetType(C,((PetscObject)B)->type_name)); 4356718818eSStefano Zampini } 4369566063dSJacob Faibussowitsch PetscCall(MatGetLocalSize(C,&m,&n)); 4379566063dSJacob Faibussowitsch PetscCall(MatGetSize(C,&M,&N)); 438825bfe5fSPierre Jolivet if (m == PETSC_DECIDE || n == PETSC_DECIDE || M == PETSC_DECIDE || N == PETSC_DECIDE) { 4399566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,Am,B->cmap->n,A->rmap->N,BN)); 440825bfe5fSPierre Jolivet } 4419566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,B)); 4429566063dSJacob Faibussowitsch PetscCall(MatSetUp(C)); 4439566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(B,&blda)); 4449566063dSJacob Faibussowitsch PetscCall(PetscNew(&contents)); 4454324174eSBarry Smith 4469566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,NULL,NULL)); 4479566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,NULL,NULL)); 44887753ddeSHong Zhang 44987753ddeSHong Zhang /* Create column block of B and C for memory scalability when BN is too large */ 45087753ddeSHong Zhang /* Estimate Bbn, column size of Bb */ 45187753ddeSHong Zhang if (nz) { 45287753ddeSHong Zhang Bbn1 = 2*Am*BN/nz; 453da7cbbe7SHong Zhang if (!Bbn1) Bbn1 = 1; 45487753ddeSHong Zhang } else Bbn1 = BN; 45587753ddeSHong Zhang 45687753ddeSHong Zhang bs = PetscAbs(B->cmap->bs); 45787753ddeSHong Zhang Bbn1 = Bbn1/bs *bs; /* Bbn1 is a multiple of bs */ 45887753ddeSHong Zhang if (Bbn1 > BN) Bbn1 = BN; 4599566063dSJacob Faibussowitsch PetscCallMPI(MPI_Allreduce(&Bbn1,&Bbn,1,MPIU_INT,MPI_MAX,comm)); 46087753ddeSHong Zhang 46187753ddeSHong Zhang /* Enable runtime option for Bbn */ 462d0609cedSBarry Smith PetscOptionsBegin(comm,((PetscObject)C)->prefix,"MatMatMult","Mat"); 4639566063dSJacob Faibussowitsch PetscCall(PetscOptionsInt("-matmatmult_Bbn","Number of columns in Bb","MatMatMult",Bbn,&Bbn,NULL)); 464d0609cedSBarry Smith PetscOptionsEnd(); 4656718818eSStefano Zampini Bbn = PetscMin(Bbn,BN); 46687753ddeSHong Zhang 4676718818eSStefano Zampini if (Bbn > 0 && Bbn < BN) { 4685ea7661aSPierre Jolivet numBb = BN/Bbn; 4695ea7661aSPierre Jolivet Bbn1 = BN - numBb*Bbn; 4705ea7661aSPierre Jolivet } else numBb = 0; 47187753ddeSHong Zhang 4725ea7661aSPierre Jolivet if (numBb) { 4739566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"use Bb, BN=%" PetscInt_FMT ", Bbn=%" PetscInt_FMT "; numBb=%" PetscInt_FMT "\n",BN,Bbn,numBb)); 4745ea7661aSPierre Jolivet if (Bbn1) { /* Create workB1 for the remaining columns */ 4759566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"use Bb1, BN=%" PetscInt_FMT ", Bbn1=%" PetscInt_FMT "\n",BN,Bbn1)); 47687753ddeSHong Zhang /* Create work matrix used to store off processor rows of B needed for local product */ 4779566063dSJacob Faibussowitsch PetscCall(MatCreateSeqDense(PETSC_COMM_SELF,nz,Bbn1,NULL,&contents->workB1)); 4785ea7661aSPierre Jolivet } else contents->workB1 = NULL; 47987753ddeSHong Zhang } 48087753ddeSHong Zhang 48187753ddeSHong Zhang /* Create work matrix used to store off processor rows of B needed for local product */ 4829566063dSJacob Faibussowitsch PetscCall(MatCreateSeqDense(PETSC_COMM_SELF,nz,Bbn,NULL,&contents->workB)); 48387753ddeSHong Zhang 48487753ddeSHong Zhang /* Use MPI derived data type to reduce memory required by the send/recv buffers */ 4859566063dSJacob Faibussowitsch PetscCall(PetscMalloc4(nsends,&stype,nrecvs,&rtype,nrecvs,&contents->rwaits,nsends,&contents->swaits)); 48687753ddeSHong Zhang contents->stype = stype; 48787753ddeSHong Zhang contents->nsends = nsends; 48887753ddeSHong Zhang 48987753ddeSHong Zhang contents->rtype = rtype; 49087753ddeSHong Zhang contents->nrecvs = nrecvs; 4916718818eSStefano Zampini contents->blda = blda; 49287753ddeSHong Zhang 4939566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(Bm+1,&disp)); 49487753ddeSHong Zhang for (i=0; i<nsends; i++) { 49587753ddeSHong Zhang nrows_to = sstarts[i+1]-sstarts[i]; 496b45e3bf4SStefano Zampini for (j=0; j<nrows_to; j++) disp[j] = sindices[sstarts[i]+j]; /* rowB to be sent */ 4979566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_indexed_block(nrows_to,1,disp,MPIU_SCALAR,&type1)); 4989566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_resized(type1,0,blda*sizeof(PetscScalar),&stype[i])); 4999566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_commit(&stype[i])); 5009566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&type1)); 50187753ddeSHong Zhang } 50287753ddeSHong Zhang 50387753ddeSHong Zhang for (i=0; i<nrecvs; i++) { 50487753ddeSHong Zhang /* received values from a process form a (nrows_from x Bbn) row block in workB (column-wise) */ 50587753ddeSHong Zhang nrows_from = rstarts[i+1]-rstarts[i]; 50687753ddeSHong Zhang disp[0] = 0; 5079566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_indexed_block(1,nrows_from,disp,MPIU_SCALAR,&type1)); 5089566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_resized(type1,0,nz*sizeof(PetscScalar),&rtype[i])); 5099566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_commit(&rtype[i])); 5109566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&type1)); 51187753ddeSHong Zhang } 51287753ddeSHong Zhang 5139566063dSJacob Faibussowitsch PetscCall(PetscFree(disp)); 5149566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,NULL,NULL)); 5159566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,NULL,NULL)); 5169566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NO_OFF_PROC_ENTRIES,PETSC_TRUE)); 5179566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 5189566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 5199566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 5206718818eSStefano Zampini 5216718818eSStefano Zampini C->product->data = contents; 5226718818eSStefano Zampini C->product->destroy = MatMPIAIJ_MPIDenseDestroy; 5234222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIDense; 5249123193aSHong Zhang PetscFunctionReturn(0); 5259123193aSHong Zhang } 5269123193aSHong Zhang 52793aa15f2SStefano Zampini PETSC_INTERN PetscErrorCode MatMatMultNumericAdd_SeqAIJ_SeqDense(Mat,Mat,Mat,const PetscBool); 528f32d5d43SBarry Smith /* 5292636ff26SBarry Smith Performs an efficient scatter on the rows of B needed by this process; this is 5302636ff26SBarry Smith a modification of the VecScatterBegin_() routines. 53187753ddeSHong Zhang 53287753ddeSHong Zhang Input: Bbidx = 0: B = Bb 53387753ddeSHong Zhang = 1: B = Bb1, see MatMatMultSymbolic_MPIAIJ_MPIDense() 534f32d5d43SBarry Smith */ 53587753ddeSHong Zhang PetscErrorCode MatMPIDenseScatter(Mat A,Mat B,PetscInt Bbidx,Mat C,Mat *outworkB) 536f32d5d43SBarry Smith { 537f32d5d43SBarry Smith Mat_MPIAIJ *aij = (Mat_MPIAIJ*)A->data; 5381683a169SBarry Smith const PetscScalar *b; 53987753ddeSHong Zhang PetscScalar *rvalues; 540f32d5d43SBarry Smith VecScatter ctx = aij->Mvctx; 54187753ddeSHong Zhang const PetscInt *sindices,*sstarts,*rstarts; 5423515ee7fSJunchao Zhang const PetscMPIInt *sprocs,*rprocs; 5436718818eSStefano Zampini PetscInt i,nsends,nrecvs; 544f32d5d43SBarry Smith MPI_Request *swaits,*rwaits; 545ce94432eSBarry Smith MPI_Comm comm; 5466718818eSStefano Zampini PetscMPIInt tag=((PetscObject)ctx)->tag,ncols=B->cmap->N,nrows=aij->B->cmap->n,nsends_mpi,nrecvs_mpi; 5474324174eSBarry Smith MPIAIJ_MPIDense *contents; 5484324174eSBarry Smith Mat workB; 54987753ddeSHong Zhang MPI_Datatype *stype,*rtype; 5506718818eSStefano Zampini PetscInt blda; 551f32d5d43SBarry Smith 552f32d5d43SBarry Smith PetscFunctionBegin; 5536718818eSStefano Zampini MatCheckProduct(C,4); 55428b400f6SJacob Faibussowitsch PetscCheck(C->product->data,PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data empty"); 5556718818eSStefano Zampini contents = (MPIAIJ_MPIDense*)C->product->data; 5569566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,&sprocs,NULL/*bs*/)); 5579566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,&rprocs,NULL/*bs*/)); 5589566063dSJacob Faibussowitsch PetscCall(PetscMPIIntCast(nsends,&nsends_mpi)); 5599566063dSJacob Faibussowitsch PetscCall(PetscMPIIntCast(nrecvs,&nrecvs_mpi)); 560b45e3bf4SStefano Zampini if (Bbidx == 0) workB = *outworkB = contents->workB; 561b45e3bf4SStefano Zampini else workB = *outworkB = contents->workB1; 56208401ef6SPierre Jolivet PetscCheck(nrows == workB->rmap->n,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Number of rows of workB %" PetscInt_FMT " not equal to columns of aij->B %d",workB->cmap->n,nrows); 5633515ee7fSJunchao Zhang swaits = contents->swaits; 5643515ee7fSJunchao Zhang rwaits = contents->rwaits; 565f32d5d43SBarry Smith 5669566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(B,&b)); 5679566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(B,&blda)); 56808401ef6SPierre Jolivet PetscCheck(blda == contents->blda,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"Cannot reuse an input matrix with lda %" PetscInt_FMT " != %" PetscInt_FMT,blda,contents->blda); 5699566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(workB,&rvalues)); 570f32d5d43SBarry Smith 57187753ddeSHong Zhang /* Post recv, use MPI derived data type to save memory */ 5729566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 57387753ddeSHong Zhang rtype = contents->rtype; 5743515ee7fSJunchao Zhang for (i=0; i<nrecvs; i++) { 5759566063dSJacob Faibussowitsch PetscCallMPI(MPI_Irecv(rvalues+(rstarts[i]-rstarts[0]),ncols,rtype[i],rprocs[i],tag,comm,rwaits+i)); 576f32d5d43SBarry Smith } 577f32d5d43SBarry Smith 57887753ddeSHong Zhang stype = contents->stype; 5793515ee7fSJunchao Zhang for (i=0; i<nsends; i++) { 5809566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(b,ncols,stype[i],sprocs[i],tag,comm,swaits+i)); 581f32d5d43SBarry Smith } 5822636ff26SBarry Smith 5839566063dSJacob Faibussowitsch if (nrecvs) PetscCallMPI(MPI_Waitall(nrecvs_mpi,rwaits,MPI_STATUSES_IGNORE)); 5849566063dSJacob Faibussowitsch if (nsends) PetscCallMPI(MPI_Waitall(nsends_mpi,swaits,MPI_STATUSES_IGNORE)); 585f32d5d43SBarry Smith 5869566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,&sprocs,NULL)); 5879566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,&rprocs,NULL)); 5889566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(B,&b)); 5899566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(workB,&rvalues)); 59087753ddeSHong Zhang PetscFunctionReturn(0); 59187753ddeSHong Zhang } 592f32d5d43SBarry Smith 5936718818eSStefano Zampini static PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIDense(Mat A,Mat B,Mat C) 5949123193aSHong Zhang { 595f32d5d43SBarry Smith Mat_MPIAIJ *aij = (Mat_MPIAIJ*)A->data; 596f32d5d43SBarry Smith Mat_MPIDense *bdense = (Mat_MPIDense*)B->data; 597f32d5d43SBarry Smith Mat_MPIDense *cdense = (Mat_MPIDense*)C->data; 598f32d5d43SBarry Smith Mat workB; 59987753ddeSHong Zhang MPIAIJ_MPIDense *contents; 6009123193aSHong Zhang 6019123193aSHong Zhang PetscFunctionBegin; 6026718818eSStefano Zampini MatCheckProduct(C,3); 60328b400f6SJacob Faibussowitsch PetscCheck(C->product->data,PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data empty"); 6046718818eSStefano Zampini contents = (MPIAIJ_MPIDense*)C->product->data; 605f32d5d43SBarry Smith /* diagonal block of A times all local rows of B */ 6065ea7661aSPierre Jolivet /* TODO: this calls a symbolic multiplication every time, which could be avoided */ 6079566063dSJacob Faibussowitsch PetscCall(MatMatMult(aij->A,bdense->A,MAT_REUSE_MATRIX,PETSC_DEFAULT,&cdense->A)); 6085ea7661aSPierre Jolivet if (contents->workB->cmap->n == B->cmap->N) { 60987753ddeSHong Zhang /* get off processor parts of B needed to complete C=A*B */ 6109566063dSJacob Faibussowitsch PetscCall(MatMPIDenseScatter(A,B,0,C,&workB)); 611f32d5d43SBarry Smith 612f32d5d43SBarry Smith /* off-diagonal block of A times nonlocal rows of B */ 6139566063dSJacob Faibussowitsch PetscCall(MatMatMultNumericAdd_SeqAIJ_SeqDense(aij->B,workB,cdense->A,PETSC_TRUE)); 61487753ddeSHong Zhang } else { 6155ea7661aSPierre Jolivet Mat Bb,Cb; 6165ea7661aSPierre Jolivet PetscInt BN=B->cmap->N,n=contents->workB->cmap->n,i; 617664c1e85SStefano Zampini PetscBool ccpu; 61887753ddeSHong Zhang 61908401ef6SPierre Jolivet PetscCheck(n > 0,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"Column block size %" PetscInt_FMT " must be positive",n); 620664c1e85SStefano Zampini /* Prevent from unneeded copies back and forth from the GPU 621664c1e85SStefano Zampini when getting and restoring the submatrix 622664c1e85SStefano Zampini We need a proper GPU code for AIJ * dense in parallel */ 6239566063dSJacob Faibussowitsch PetscCall(MatBoundToCPU(C,&ccpu)); 6249566063dSJacob Faibussowitsch PetscCall(MatBindToCPU(C,PETSC_TRUE)); 6255ea7661aSPierre Jolivet for (i=0; i<BN; i+=n) { 6269566063dSJacob Faibussowitsch PetscCall(MatDenseGetSubMatrix(B,i,PetscMin(i+n,BN),&Bb)); 6279566063dSJacob Faibussowitsch PetscCall(MatDenseGetSubMatrix(C,i,PetscMin(i+n,BN),&Cb)); 62887753ddeSHong Zhang 6295ea7661aSPierre Jolivet /* get off processor parts of B needed to complete C=A*B */ 6309566063dSJacob Faibussowitsch PetscCall(MatMPIDenseScatter(A,Bb,(i+n)>BN,C,&workB)); 6315ea7661aSPierre Jolivet 6325ea7661aSPierre Jolivet /* off-diagonal block of A times nonlocal rows of B */ 6335ea7661aSPierre Jolivet cdense = (Mat_MPIDense*)Cb->data; 6349566063dSJacob Faibussowitsch PetscCall(MatMatMultNumericAdd_SeqAIJ_SeqDense(aij->B,workB,cdense->A,PETSC_TRUE)); 6359566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreSubMatrix(B,&Bb)); 6369566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreSubMatrix(C,&Cb)); 63787753ddeSHong Zhang } 6389566063dSJacob Faibussowitsch PetscCall(MatBindToCPU(C,ccpu)); 63987753ddeSHong Zhang } 6409123193aSHong Zhang PetscFunctionReturn(0); 6419123193aSHong Zhang } 642cf1a0672SHong Zhang 643f8efd71cSHong Zhang PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIAIJ(Mat A,Mat P,Mat C) 6441634b032SHong Zhang { 645cf1a0672SHong Zhang Mat_MPIAIJ *a = (Mat_MPIAIJ*)A->data,*c=(Mat_MPIAIJ*)C->data; 646cf1a0672SHong Zhang Mat_SeqAIJ *ad = (Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data; 647cf1a0672SHong Zhang Mat_SeqAIJ *cd = (Mat_SeqAIJ*)(c->A)->data,*co=(Mat_SeqAIJ*)(c->B)->data; 648cf1a0672SHong Zhang PetscInt *adi = ad->i,*adj,*aoi=ao->i,*aoj; 649cf1a0672SHong Zhang PetscScalar *ada,*aoa,*cda=cd->a,*coa=co->a; 650cf1a0672SHong Zhang Mat_SeqAIJ *p_loc,*p_oth; 651cf1a0672SHong Zhang PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*pj; 65229825010SHong Zhang PetscScalar *pa_loc,*pa_oth,*pa,valtmp,*ca; 65329825010SHong Zhang PetscInt cm = C->rmap->n,anz,pnz; 6546718818eSStefano Zampini Mat_APMPI *ptap; 65580da8df7SHong Zhang PetscScalar *apa_sparse; 656ce496241SStefano Zampini const PetscScalar *dummy; 657cf1a0672SHong Zhang PetscInt *api,*apj,*apJ,i,j,k,row; 65829825010SHong Zhang PetscInt cstart = C->cmap->rstart; 65929825010SHong Zhang PetscInt cdnz,conz,k0,k1,nextp; 660a7c7454dSHong Zhang MPI_Comm comm; 661a7c7454dSHong Zhang PetscMPIInt size; 6621634b032SHong Zhang 6631634b032SHong Zhang PetscFunctionBegin; 6646718818eSStefano Zampini MatCheckProduct(C,3); 6656718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 66628b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 6679566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 6689566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 6692c71b3e2SJacob Faibussowitsch PetscCheckFalse(!ptap->P_oth && size>1,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"AP cannot be reused. Do not call MatProductClear()"); 670a7c7454dSHong Zhang 67165e4b4d4SStefano Zampini /* flag CPU mask for C */ 67265e4b4d4SStefano Zampini #if defined(PETSC_HAVE_DEVICE) 67365e4b4d4SStefano Zampini if (C->offloadmask != PETSC_OFFLOAD_UNALLOCATED) C->offloadmask = PETSC_OFFLOAD_CPU; 67465e4b4d4SStefano Zampini if (c->A->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->A->offloadmask = PETSC_OFFLOAD_CPU; 67565e4b4d4SStefano Zampini if (c->B->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->B->offloadmask = PETSC_OFFLOAD_CPU; 67665e4b4d4SStefano Zampini #endif 67780da8df7SHong Zhang apa_sparse = ptap->apa; 67880da8df7SHong Zhang 679cf1a0672SHong Zhang /* 1) get P_oth = ptap->P_oth and P_loc = ptap->P_loc */ 680cf1a0672SHong Zhang /*-----------------------------------------------------*/ 681cf1a0672SHong Zhang /* update numerical values of P_oth and P_loc */ 6829566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_REUSE_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 6839566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_REUSE_MATRIX,&ptap->P_loc)); 684cf1a0672SHong Zhang 685cf1a0672SHong Zhang /* 2) compute numeric C_loc = A_loc*P = Ad*P_loc + Ao*P_oth */ 686cf1a0672SHong Zhang /*----------------------------------------------------------*/ 687cf1a0672SHong Zhang /* get data from symbolic products */ 688cf1a0672SHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 689cf1a0672SHong Zhang pi_loc = p_loc->i; pj_loc = p_loc->j; pa_loc = p_loc->a; 690a7c7454dSHong Zhang if (size >1) { 691a7c7454dSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 692cf1a0672SHong Zhang pi_oth = p_oth->i; pj_oth = p_oth->j; pa_oth = p_oth->a; 69320e3dc75SHong Zhang } else { 69420e3dc75SHong Zhang p_oth = NULL; pi_oth = NULL; pj_oth = NULL; pa_oth = NULL; 695a7c7454dSHong Zhang } 696cf1a0672SHong Zhang 697ce496241SStefano Zampini /* trigger copy to CPU */ 6989566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->A,&dummy)); 6999566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->A,&dummy)); 7009566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->B,&dummy)); 7019566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->B,&dummy)); 70229825010SHong Zhang api = ptap->api; 70329825010SHong Zhang apj = ptap->apj; 704cf1a0672SHong Zhang for (i=0; i<cm; i++) { 70529825010SHong Zhang apJ = apj + api[i]; 70629825010SHong Zhang 707cf1a0672SHong Zhang /* diagonal portion of A */ 708cf1a0672SHong Zhang anz = adi[i+1] - adi[i]; 709cf1a0672SHong Zhang adj = ad->j + adi[i]; 710cf1a0672SHong Zhang ada = ad->a + adi[i]; 711cf1a0672SHong Zhang for (j=0; j<anz; j++) { 712cf1a0672SHong Zhang row = adj[j]; 713cf1a0672SHong Zhang pnz = pi_loc[row+1] - pi_loc[row]; 714cf1a0672SHong Zhang pj = pj_loc + pi_loc[row]; 715cf1a0672SHong Zhang pa = pa_loc + pi_loc[row]; 71629825010SHong Zhang /* perform sparse axpy */ 717cf1a0672SHong Zhang valtmp = ada[j]; 71829825010SHong Zhang nextp = 0; 71929825010SHong Zhang for (k=0; nextp<pnz; k++) { 72029825010SHong Zhang if (apJ[k] == pj[nextp]) { /* column of AP == column of P */ 72129825010SHong Zhang apa_sparse[k] += valtmp*pa[nextp++]; 72229825010SHong Zhang } 7231634b032SHong Zhang } 7249566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*pnz)); 725cf1a0672SHong Zhang } 7261634b032SHong Zhang 727cf1a0672SHong Zhang /* off-diagonal portion of A */ 728cf1a0672SHong Zhang anz = aoi[i+1] - aoi[i]; 729cf1a0672SHong Zhang aoj = ao->j + aoi[i]; 730cf1a0672SHong Zhang aoa = ao->a + aoi[i]; 731cf1a0672SHong Zhang for (j=0; j<anz; j++) { 732cf1a0672SHong Zhang row = aoj[j]; 733cf1a0672SHong Zhang pnz = pi_oth[row+1] - pi_oth[row]; 734cf1a0672SHong Zhang pj = pj_oth + pi_oth[row]; 735cf1a0672SHong Zhang pa = pa_oth + pi_oth[row]; 73629825010SHong Zhang /* perform sparse axpy */ 737cf1a0672SHong Zhang valtmp = aoa[j]; 73829825010SHong Zhang nextp = 0; 73929825010SHong Zhang for (k=0; nextp<pnz; k++) { 74029825010SHong Zhang if (apJ[k] == pj[nextp]) { /* column of AP == column of P */ 74129825010SHong Zhang apa_sparse[k] += valtmp*pa[nextp++]; 74229825010SHong Zhang } 743cf1a0672SHong Zhang } 7449566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*pnz)); 745cf1a0672SHong Zhang } 7461634b032SHong Zhang 747cf1a0672SHong Zhang /* set values in C */ 748cf1a0672SHong Zhang cdnz = cd->i[i+1] - cd->i[i]; 749cf1a0672SHong Zhang conz = co->i[i+1] - co->i[i]; 7501634b032SHong Zhang 751910cf402Sprj- /* 1st off-diagonal part of C */ 752cf1a0672SHong Zhang ca = coa + co->i[i]; 753cf1a0672SHong Zhang k = 0; 754cf1a0672SHong Zhang for (k0=0; k0<conz; k0++) { 755cf1a0672SHong Zhang if (apJ[k] >= cstart) break; 75629825010SHong Zhang ca[k0] = apa_sparse[k]; 75729825010SHong Zhang apa_sparse[k] = 0.0; 758cf1a0672SHong Zhang k++; 759cf1a0672SHong Zhang } 7601634b032SHong Zhang 761cf1a0672SHong Zhang /* diagonal part of C */ 762cf1a0672SHong Zhang ca = cda + cd->i[i]; 763cf1a0672SHong Zhang for (k1=0; k1<cdnz; k1++) { 76429825010SHong Zhang ca[k1] = apa_sparse[k]; 76529825010SHong Zhang apa_sparse[k] = 0.0; 766cf1a0672SHong Zhang k++; 767cf1a0672SHong Zhang } 768cf1a0672SHong Zhang 769910cf402Sprj- /* 2nd off-diagonal part of C */ 770cf1a0672SHong Zhang ca = coa + co->i[i]; 771cf1a0672SHong Zhang for (; k0<conz; k0++) { 77229825010SHong Zhang ca[k0] = apa_sparse[k]; 77329825010SHong Zhang apa_sparse[k] = 0.0; 774cf1a0672SHong Zhang k++; 775cf1a0672SHong Zhang } 776cf1a0672SHong Zhang } 7779566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 7789566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 779cf1a0672SHong Zhang PetscFunctionReturn(0); 780cf1a0672SHong Zhang } 781cf1a0672SHong Zhang 7820fc8cf34SHong Zhang /* same as MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(), except using LLCondensed to avoid O(BN) memory requirement */ 7834222ddf1SHong Zhang PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ(Mat A,Mat P,PetscReal fill,Mat C) 784cf1a0672SHong Zhang { 785ce94432eSBarry Smith MPI_Comm comm; 786a7c7454dSHong Zhang PetscMPIInt size; 7873cdca5ebSHong Zhang Mat_APMPI *ptap; 7880298fd71SBarry Smith PetscFreeSpaceList free_space = NULL,current_space=NULL; 7896718818eSStefano Zampini Mat_MPIAIJ *a = (Mat_MPIAIJ*)A->data; 790cf1a0672SHong Zhang Mat_SeqAIJ *ad = (Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_loc,*p_oth; 7917dbaa2c6Sandi selinger PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*dnz,*onz; 792cf1a0672SHong Zhang PetscInt *adi=ad->i,*adj=ad->j,*aoi=ao->i,*aoj=ao->j,rstart=A->rmap->rstart; 7938c78258cSHong Zhang PetscInt i,pnz,row,*api,*apj,*Jptr,apnz,nspacedouble=0,j,nzi,*lnk,apnz_max=1; 794e5ab8c0aSKarl Rupp PetscInt am=A->rmap->n,pn=P->cmap->n,pm=P->rmap->n,lsize=pn+20; 795cf1a0672SHong Zhang PetscReal afill; 796b92f168fSBarry Smith MatType mtype; 797cf1a0672SHong Zhang 798cf1a0672SHong Zhang PetscFunctionBegin; 7996718818eSStefano Zampini MatCheckProduct(C,4); 80028b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Extra product struct not empty"); 8019566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 8029566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 803a7c7454dSHong Zhang 8043cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 8059566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 806cf1a0672SHong Zhang 807cf1a0672SHong Zhang /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 8089566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 80919f0e57aSHong Zhang 810cf1a0672SHong Zhang /* get P_loc by taking all local rows of P */ 8119566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc)); 812cf1a0672SHong Zhang 813cf1a0672SHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 814cf1a0672SHong Zhang pi_loc = p_loc->i; pj_loc = p_loc->j; 815a7c7454dSHong Zhang if (size > 1) { 816a7c7454dSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 81720e3dc75SHong Zhang pi_oth = p_oth->i; pj_oth = p_oth->j; 818968382fdSHong Zhang } else { 81920e3dc75SHong Zhang p_oth = NULL; 82020e3dc75SHong Zhang pi_oth = NULL; pj_oth = NULL; 821a7c7454dSHong Zhang } 822cf1a0672SHong Zhang 823cf1a0672SHong Zhang /* first, compute symbolic AP = A_loc*P = A_diag*P_loc + A_off*P_oth */ 824cf1a0672SHong Zhang /*-------------------------------------------------------------------*/ 8259566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&api)); 826cf1a0672SHong Zhang ptap->api = api; 827cf1a0672SHong Zhang api[0] = 0; 828cf1a0672SHong Zhang 8299566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(lsize,&lnk)); 830cf1a0672SHong Zhang 831cf1a0672SHong Zhang /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */ 8329566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(adi[am],PetscIntSumTruncate(aoi[am],pi_loc[pm]))),&free_space)); 833cf1a0672SHong Zhang current_space = free_space; 834d0609cedSBarry Smith MatPreallocateBegin(comm,am,pn,dnz,onz); 835cf1a0672SHong Zhang for (i=0; i<am; i++) { 836cf1a0672SHong Zhang /* diagonal portion of A */ 837cf1a0672SHong Zhang nzi = adi[i+1] - adi[i]; 838cf1a0672SHong Zhang for (j=0; j<nzi; j++) { 839cf1a0672SHong Zhang row = *adj++; 840cf1a0672SHong Zhang pnz = pi_loc[row+1] - pi_loc[row]; 841cf1a0672SHong Zhang Jptr = pj_loc + pi_loc[row]; 8428a210135Sandi selinger /* Expand list if it is not long enough */ 8438328dd82Sandi selinger if (pnz+apnz_max > lsize) { 8448328dd82Sandi selinger lsize = pnz+apnz_max; 8459566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedExpand_Scalable(lsize, &lnk)); 8468a210135Sandi selinger } 847cf1a0672SHong Zhang /* add non-zero cols of P into the sorted linked list lnk */ 8489566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(pnz,Jptr,lnk)); 8498a210135Sandi selinger apnz = *lnk; /* The first element in the list is the number of items in the list */ 8508a210135Sandi selinger api[i+1] = api[i] + apnz; 8518c78258cSHong Zhang if (apnz > apnz_max) apnz_max = apnz + 1; /* '1' for diagonal entry */ 852cf1a0672SHong Zhang } 853cf1a0672SHong Zhang /* off-diagonal portion of A */ 854cf1a0672SHong Zhang nzi = aoi[i+1] - aoi[i]; 855cf1a0672SHong Zhang for (j=0; j<nzi; j++) { 856cf1a0672SHong Zhang row = *aoj++; 857cf1a0672SHong Zhang pnz = pi_oth[row+1] - pi_oth[row]; 858cf1a0672SHong Zhang Jptr = pj_oth + pi_oth[row]; 8598a210135Sandi selinger /* Expand list if it is not long enough */ 8608328dd82Sandi selinger if (pnz+apnz_max > lsize) { 8618328dd82Sandi selinger lsize = pnz + apnz_max; 8629566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedExpand_Scalable(lsize, &lnk)); 8638a210135Sandi selinger } 8648a210135Sandi selinger /* add non-zero cols of P into the sorted linked list lnk */ 8659566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(pnz,Jptr,lnk)); 8668a210135Sandi selinger apnz = *lnk; /* The first element in the list is the number of items in the list */ 8678a210135Sandi selinger api[i+1] = api[i] + apnz; 8688c78258cSHong Zhang if (apnz > apnz_max) apnz_max = apnz + 1; /* '1' for diagonal entry */ 869cf1a0672SHong Zhang } 8708c78258cSHong Zhang 8718c78258cSHong Zhang /* add missing diagonal entry */ 8728c78258cSHong Zhang if (C->force_diagonals) { 8738c78258cSHong Zhang j = i + rstart; /* column index */ 8749566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(1,&j,lnk)); 8758c78258cSHong Zhang } 8768c78258cSHong Zhang 877f84c536bSHong Zhang apnz = *lnk; 878cf1a0672SHong Zhang api[i+1] = api[i] + apnz; 8798a210135Sandi selinger if (apnz > apnz_max) apnz_max = apnz; 880cf1a0672SHong Zhang 881cf1a0672SHong Zhang /* if free space is not available, double the total space in the list */ 882cf1a0672SHong Zhang if (current_space->local_remaining<apnz) { 8839566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(apnz,current_space->total_array_size),¤t_space)); 884cf1a0672SHong Zhang nspacedouble++; 885cf1a0672SHong Zhang } 886cf1a0672SHong Zhang 887cf1a0672SHong Zhang /* Copy data into free space, then initialize lnk */ 8889566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(apnz,current_space->array,lnk)); 8899566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+rstart,apnz,current_space->array,dnz,onz)); 8902205254eSKarl Rupp 891cf1a0672SHong Zhang current_space->array += apnz; 892cf1a0672SHong Zhang current_space->local_used += apnz; 893cf1a0672SHong Zhang current_space->local_remaining -= apnz; 894cf1a0672SHong Zhang } 895cf1a0672SHong Zhang 896cf1a0672SHong Zhang /* Allocate space for apj, initialize apj, and */ 897cf1a0672SHong Zhang /* destroy list of free space and other temporary array(s) */ 8989566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(api[am]+1,&ptap->apj)); 899cf1a0672SHong Zhang apj = ptap->apj; 9009566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,ptap->apj)); 9019566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); 902cf1a0672SHong Zhang 9034222ddf1SHong Zhang /* create and assemble symbolic parallel matrix C */ 904cf1a0672SHong Zhang /*----------------------------------------------------*/ 9059566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,am,pn,PETSC_DETERMINE,PETSC_DETERMINE)); 9069566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,P)); 9079566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 9089566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 9099566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 910d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 911cf1a0672SHong Zhang 9124222ddf1SHong Zhang /* malloc apa for assembly C */ 9139566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(apnz_max,&ptap->apa)); 91417f75172Sandi selinger 9159566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api)); 9169566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 9179566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 9189566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 919cf1a0672SHong Zhang 9204222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ; 9214222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AB; 922cf1a0672SHong Zhang 9234222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 9246718818eSStefano Zampini C->product->data = ptap; 9256718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_MatMatMult; 926cf1a0672SHong Zhang 927cf1a0672SHong Zhang /* set MatInfo */ 928118727c9SMark F. Adams afill = (PetscReal)api[am]/(adi[am]+aoi[am]+pi_loc[pm]+1) + 1.e-5; 929cf1a0672SHong Zhang if (afill < 1.0) afill = 1.0; 9304222ddf1SHong Zhang C->info.mallocs = nspacedouble; 9314222ddf1SHong Zhang C->info.fill_ratio_given = fill; 9324222ddf1SHong Zhang C->info.fill_ratio_needed = afill; 933cf1a0672SHong Zhang 934cf1a0672SHong Zhang #if defined(PETSC_USE_INFO) 935cf1a0672SHong Zhang if (api[am]) { 9369566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 9379566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n",(double)afill)); 938cf1a0672SHong Zhang } else { 9399566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 940cf1a0672SHong Zhang } 941cf1a0672SHong Zhang #endif 9421634b032SHong Zhang PetscFunctionReturn(0); 9431634b032SHong Zhang } 944f96d37fcSHong Zhang 945a9d7e0ccSandi selinger /* This function is needed for the seqMPI matrix-matrix multiplication. */ 946a9d7e0ccSandi selinger /* Three input arrays are merged to one output array. The size of the */ 947a9d7e0ccSandi selinger /* output array is also output. Duplicate entries only show up once. */ 948a9d7e0ccSandi selinger static void Merge3SortedArrays(PetscInt size1, PetscInt *in1, 949a9d7e0ccSandi selinger PetscInt size2, PetscInt *in2, 950a9d7e0ccSandi selinger PetscInt size3, PetscInt *in3, 951a9d7e0ccSandi selinger PetscInt *size4, PetscInt *out) 952a9d7e0ccSandi selinger { 953a9d7e0ccSandi selinger int i = 0, j = 0, k = 0, l = 0; 954a9d7e0ccSandi selinger 955a9d7e0ccSandi selinger /* Traverse all three arrays */ 956a9d7e0ccSandi selinger while (i<size1 && j<size2 && k<size3) { 957a9d7e0ccSandi selinger if (in1[i] < in2[j] && in1[i] < in3[k]) { 958a9d7e0ccSandi selinger out[l++] = in1[i++]; 959a9d7e0ccSandi selinger } 960a9d7e0ccSandi selinger else if (in2[j] < in1[i] && in2[j] < in3[k]) { 961a9d7e0ccSandi selinger out[l++] = in2[j++]; 962a9d7e0ccSandi selinger } 963a9d7e0ccSandi selinger else if (in3[k] < in1[i] && in3[k] < in2[j]) { 964a9d7e0ccSandi selinger out[l++] = in3[k++]; 965a9d7e0ccSandi selinger } 966a9d7e0ccSandi selinger else if (in1[i] == in2[j] && in1[i] < in3[k]) { 967a9d7e0ccSandi selinger out[l++] = in1[i]; 968a9d7e0ccSandi selinger i++, j++; 969a9d7e0ccSandi selinger } 970a9d7e0ccSandi selinger else if (in1[i] == in3[k] && in1[i] < in2[j]) { 971a9d7e0ccSandi selinger out[l++] = in1[i]; 972a9d7e0ccSandi selinger i++, k++; 973a9d7e0ccSandi selinger } 974a9d7e0ccSandi selinger else if (in3[k] == in2[j] && in2[j] < in1[i]) { 975a9d7e0ccSandi selinger out[l++] = in2[j]; 976a9d7e0ccSandi selinger k++, j++; 977a9d7e0ccSandi selinger } 978a9d7e0ccSandi selinger else if (in1[i] == in2[j] && in1[i] == in3[k]) { 979a9d7e0ccSandi selinger out[l++] = in1[i]; 980a9d7e0ccSandi selinger i++, j++, k++; 981a9d7e0ccSandi selinger } 982a9d7e0ccSandi selinger } 983a9d7e0ccSandi selinger 984a9d7e0ccSandi selinger /* Traverse two remaining arrays */ 985a9d7e0ccSandi selinger while (i<size1 && j<size2) { 986a9d7e0ccSandi selinger if (in1[i] < in2[j]) { 987a9d7e0ccSandi selinger out[l++] = in1[i++]; 988a9d7e0ccSandi selinger } 989a9d7e0ccSandi selinger else if (in1[i] > in2[j]) { 990a9d7e0ccSandi selinger out[l++] = in2[j++]; 991a9d7e0ccSandi selinger } 992a9d7e0ccSandi selinger else { 993a9d7e0ccSandi selinger out[l++] = in1[i]; 994a9d7e0ccSandi selinger i++, j++; 995a9d7e0ccSandi selinger } 996a9d7e0ccSandi selinger } 997a9d7e0ccSandi selinger 998a9d7e0ccSandi selinger while (i<size1 && k<size3) { 999a9d7e0ccSandi selinger if (in1[i] < in3[k]) { 1000a9d7e0ccSandi selinger out[l++] = in1[i++]; 1001a9d7e0ccSandi selinger } 1002a9d7e0ccSandi selinger else if (in1[i] > in3[k]) { 1003a9d7e0ccSandi selinger out[l++] = in3[k++]; 1004a9d7e0ccSandi selinger } 1005a9d7e0ccSandi selinger else { 1006a9d7e0ccSandi selinger out[l++] = in1[i]; 1007a9d7e0ccSandi selinger i++, k++; 1008a9d7e0ccSandi selinger } 1009a9d7e0ccSandi selinger } 1010a9d7e0ccSandi selinger 1011a9d7e0ccSandi selinger while (k<size3 && j<size2) { 1012a9d7e0ccSandi selinger if (in3[k] < in2[j]) { 1013a9d7e0ccSandi selinger out[l++] = in3[k++]; 1014a9d7e0ccSandi selinger } 1015a9d7e0ccSandi selinger else if (in3[k] > in2[j]) { 1016a9d7e0ccSandi selinger out[l++] = in2[j++]; 1017a9d7e0ccSandi selinger } 1018a9d7e0ccSandi selinger else { 1019a9d7e0ccSandi selinger out[l++] = in3[k]; 1020a9d7e0ccSandi selinger k++, j++; 1021a9d7e0ccSandi selinger } 1022a9d7e0ccSandi selinger } 1023a9d7e0ccSandi selinger 1024a9d7e0ccSandi selinger /* Traverse one remaining array */ 1025a9d7e0ccSandi selinger while (i<size1) out[l++] = in1[i++]; 1026a9d7e0ccSandi selinger while (j<size2) out[l++] = in2[j++]; 1027a9d7e0ccSandi selinger while (k<size3) out[l++] = in3[k++]; 1028a9d7e0ccSandi selinger 1029a9d7e0ccSandi selinger *size4 = l; 1030a9d7e0ccSandi selinger } 1031a9d7e0ccSandi selinger 103273b67375Sandi selinger /* This matrix-matrix multiplication algorithm divides the multiplication into three multiplications and */ 103373b67375Sandi selinger /* adds up the products. Two of these three multiplications are performed with existing (sequential) */ 103473b67375Sandi selinger /* matrix-matrix multiplications. */ 10354222ddf1SHong Zhang PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ_seqMPI(Mat A, Mat P, PetscReal fill, Mat C) 1036a9d7e0ccSandi selinger { 1037a9d7e0ccSandi selinger MPI_Comm comm; 1038a9d7e0ccSandi selinger PetscMPIInt size; 10393cdca5ebSHong Zhang Mat_APMPI *ptap; 1040a9d7e0ccSandi selinger PetscFreeSpaceList free_space_diag=NULL, current_space=NULL; 1041a9d7e0ccSandi selinger Mat_MPIAIJ *a =(Mat_MPIAIJ*)A->data; 1042a9d7e0ccSandi selinger Mat_SeqAIJ *ad =(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_loc; 1043a9d7e0ccSandi selinger Mat_MPIAIJ *p =(Mat_MPIAIJ*)P->data; 1044a9d7e0ccSandi selinger Mat_SeqAIJ *adpd_seq, *p_off, *aopoth_seq; 1045a9d7e0ccSandi selinger PetscInt adponz, adpdnz; 10467dbaa2c6Sandi selinger PetscInt *pi_loc,*dnz,*onz; 1047a9d7e0ccSandi selinger PetscInt *adi=ad->i,*adj=ad->j,*aoi=ao->i,rstart=A->rmap->rstart; 1048a9d7e0ccSandi selinger PetscInt *lnk,i, i1=0,pnz,row,*adpoi,*adpoj, *api, *adpoJ, *aopJ, *apJ,*Jptr, aopnz, nspacedouble=0,j,nzi, 1049a9d7e0ccSandi selinger *apj,apnz, *adpdi, *adpdj, *adpdJ, *poff_i, *poff_j, *j_temp, *aopothi, *aopothj; 1050a016958eSandi selinger PetscInt am=A->rmap->n,pN=P->cmap->N,pn=P->cmap->n,pm=P->rmap->n, p_colstart, p_colend; 1051a9d7e0ccSandi selinger PetscBT lnkbt; 1052a9d7e0ccSandi selinger PetscReal afill; 1053a9d7e0ccSandi selinger PetscMPIInt rank; 1054a9d7e0ccSandi selinger Mat adpd, aopoth; 1055b92f168fSBarry Smith MatType mtype; 1056e83e5f86SFande Kong const char *prefix; 1057a9d7e0ccSandi selinger 1058a9d7e0ccSandi selinger PetscFunctionBegin; 10596718818eSStefano Zampini MatCheckProduct(C,4); 106028b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Extra product struct not empty"); 10619566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 10629566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 10639566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm, &rank)); 10649566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRangeColumn(P, &p_colstart, &p_colend)); 1065a9d7e0ccSandi selinger 10663cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 10679566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 1068a9d7e0ccSandi selinger 1069a9d7e0ccSandi selinger /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 10709566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 1071a9d7e0ccSandi selinger 1072a9d7e0ccSandi selinger /* get P_loc by taking all local rows of P */ 10739566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc)); 1074a9d7e0ccSandi selinger 1075a9d7e0ccSandi selinger p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 1076a9d7e0ccSandi selinger pi_loc = p_loc->i; 1077a9d7e0ccSandi selinger 1078a9d7e0ccSandi selinger /* Allocate memory for the i arrays of the matrices A*P, A_diag*P_off and A_offd * P */ 10799566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&api)); 10809566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&adpoi)); 1081a9d7e0ccSandi selinger 1082a9d7e0ccSandi selinger adpoi[0] = 0; 1083a9d7e0ccSandi selinger ptap->api = api; 1084a9d7e0ccSandi selinger api[0] = 0; 1085a9d7e0ccSandi selinger 1086a9d7e0ccSandi selinger /* create and initialize a linked list, will be used for both A_diag * P_loc_off and A_offd * P_oth */ 10879566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate(pN,pN,&lnk,&lnkbt)); 1088d0609cedSBarry Smith MatPreallocateBegin(comm,am,pn,dnz,onz); 1089a9d7e0ccSandi selinger 1090a9d7e0ccSandi selinger /* Symbolic calc of A_loc_diag * P_loc_diag */ 10919566063dSJacob Faibussowitsch PetscCall(MatGetOptionsPrefix(A,&prefix)); 10929566063dSJacob Faibussowitsch PetscCall(MatProductCreate(a->A,p->A,NULL,&adpd)); 10939566063dSJacob Faibussowitsch PetscCall(MatGetOptionsPrefix(A,&prefix)); 10949566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(adpd,prefix)); 10959566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(adpd,"inner_diag_")); 10964222ddf1SHong Zhang 10979566063dSJacob Faibussowitsch PetscCall(MatProductSetType(adpd,MATPRODUCT_AB)); 10989566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(adpd,"sorted")); 10999566063dSJacob Faibussowitsch PetscCall(MatProductSetFill(adpd,fill)); 11009566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions(adpd)); 11018c78258cSHong Zhang 11028c78258cSHong Zhang adpd->force_diagonals = C->force_diagonals; 11039566063dSJacob Faibussowitsch PetscCall(MatProductSymbolic(adpd)); 11044222ddf1SHong Zhang 1105a9d7e0ccSandi selinger adpd_seq = (Mat_SeqAIJ*)((adpd)->data); 1106a9d7e0ccSandi selinger adpdi = adpd_seq->i; adpdj = adpd_seq->j; 1107a9d7e0ccSandi selinger p_off = (Mat_SeqAIJ*)((p->B)->data); 1108a9d7e0ccSandi selinger poff_i = p_off->i; poff_j = p_off->j; 1109a9d7e0ccSandi selinger 111073b67375Sandi selinger /* j_temp stores indices of a result row before they are added to the linked list */ 11119566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pN+2,&j_temp)); 1112a9d7e0ccSandi selinger 1113a9d7e0ccSandi selinger /* Symbolic calc of the A_diag * p_loc_off */ 1114a9d7e0ccSandi selinger /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */ 11159566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(adi[am],PetscIntSumTruncate(aoi[am],pi_loc[pm]))),&free_space_diag)); 1116a9d7e0ccSandi selinger current_space = free_space_diag; 1117a9d7e0ccSandi selinger 1118a9d7e0ccSandi selinger for (i=0; i<am; i++) { 1119a9d7e0ccSandi selinger /* A_diag * P_loc_off */ 1120a9d7e0ccSandi selinger nzi = adi[i+1] - adi[i]; 1121a9d7e0ccSandi selinger for (j=0; j<nzi; j++) { 1122a9d7e0ccSandi selinger row = *adj++; 1123a9d7e0ccSandi selinger pnz = poff_i[row+1] - poff_i[row]; 1124a9d7e0ccSandi selinger Jptr = poff_j + poff_i[row]; 1125a9d7e0ccSandi selinger for (i1 = 0; i1 < pnz; i1++) { 1126a9d7e0ccSandi selinger j_temp[i1] = p->garray[Jptr[i1]]; 1127a9d7e0ccSandi selinger } 1128a9d7e0ccSandi selinger /* add non-zero cols of P into the sorted linked list lnk */ 11299566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(pnz,j_temp,lnk,lnkbt)); 1130a9d7e0ccSandi selinger } 1131a9d7e0ccSandi selinger 1132a9d7e0ccSandi selinger adponz = lnk[0]; 1133a9d7e0ccSandi selinger adpoi[i+1] = adpoi[i] + adponz; 1134a9d7e0ccSandi selinger 1135a9d7e0ccSandi selinger /* if free space is not available, double the total space in the list */ 1136a9d7e0ccSandi selinger if (current_space->local_remaining<adponz) { 11379566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(adponz,current_space->total_array_size),¤t_space)); 1138a9d7e0ccSandi selinger nspacedouble++; 1139a9d7e0ccSandi selinger } 1140a9d7e0ccSandi selinger 1141a9d7e0ccSandi selinger /* Copy data into free space, then initialize lnk */ 11429566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean(pN,adponz,current_space->array,lnk,lnkbt)); 1143a9d7e0ccSandi selinger 1144a9d7e0ccSandi selinger current_space->array += adponz; 1145a9d7e0ccSandi selinger current_space->local_used += adponz; 1146a9d7e0ccSandi selinger current_space->local_remaining -= adponz; 1147a9d7e0ccSandi selinger } 1148a9d7e0ccSandi selinger 1149a9d7e0ccSandi selinger /* Symbolic calc of A_off * P_oth */ 11509566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(a->B,prefix)); 11519566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(a->B,"inner_offdiag_")); 11529566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF,&aopoth)); 11539566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(a->B, ptap->P_oth, fill, aopoth)); 1154a9d7e0ccSandi selinger aopoth_seq = (Mat_SeqAIJ*)((aopoth)->data); 1155a9d7e0ccSandi selinger aopothi = aopoth_seq->i; aopothj = aopoth_seq->j; 1156a9d7e0ccSandi selinger 1157a9d7e0ccSandi selinger /* Allocate space for apj, adpj, aopj, ... */ 1158a9d7e0ccSandi selinger /* destroy lists of free space and other temporary array(s) */ 1159a9d7e0ccSandi selinger 11609566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(aopothi[am] + adpoi[am] + adpdi[am]+2, &ptap->apj)); 11619566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(adpoi[am]+2, &adpoj)); 1162a9d7e0ccSandi selinger 1163a9d7e0ccSandi selinger /* Copy from linked list to j-array */ 11649566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space_diag,adpoj)); 11659566063dSJacob Faibussowitsch PetscCall(PetscLLDestroy(lnk,lnkbt)); 1166a9d7e0ccSandi selinger 1167a9d7e0ccSandi selinger adpoJ = adpoj; 1168a9d7e0ccSandi selinger adpdJ = adpdj; 1169a9d7e0ccSandi selinger aopJ = aopothj; 1170a9d7e0ccSandi selinger apj = ptap->apj; 1171a9d7e0ccSandi selinger apJ = apj; /* still empty */ 1172a9d7e0ccSandi selinger 1173a9d7e0ccSandi selinger /* Merge j-arrays of A_off * P, A_diag * P_loc_off, and */ 1174a9d7e0ccSandi selinger /* A_diag * P_loc_diag to get A*P */ 1175a9d7e0ccSandi selinger for (i = 0; i < am; i++) { 1176a9d7e0ccSandi selinger aopnz = aopothi[i+1] - aopothi[i]; 1177a9d7e0ccSandi selinger adponz = adpoi[i+1] - adpoi[i]; 1178a9d7e0ccSandi selinger adpdnz = adpdi[i+1] - adpdi[i]; 1179a9d7e0ccSandi selinger 1180a9d7e0ccSandi selinger /* Correct indices from A_diag*P_diag */ 1181a9d7e0ccSandi selinger for (i1 = 0; i1 < adpdnz; i1++) { 1182a016958eSandi selinger adpdJ[i1] += p_colstart; 1183a9d7e0ccSandi selinger } 1184a9d7e0ccSandi selinger /* Merge j-arrays of A_diag * P_loc_off and A_diag * P_loc_diag and A_off * P_oth */ 1185a9d7e0ccSandi selinger Merge3SortedArrays(adponz, adpoJ, adpdnz, adpdJ, aopnz, aopJ, &apnz, apJ); 11869566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+rstart, apnz, apJ, dnz, onz)); 1187a9d7e0ccSandi selinger 1188a9d7e0ccSandi selinger aopJ += aopnz; 1189a9d7e0ccSandi selinger adpoJ += adponz; 1190a9d7e0ccSandi selinger adpdJ += adpdnz; 1191a9d7e0ccSandi selinger apJ += apnz; 1192a9d7e0ccSandi selinger api[i+1] = api[i] + apnz; 1193a9d7e0ccSandi selinger } 1194a9d7e0ccSandi selinger 1195a9d7e0ccSandi selinger /* malloc apa to store dense row A[i,:]*P */ 11969566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(pN+2,&ptap->apa)); 1197a9d7e0ccSandi selinger 11984222ddf1SHong Zhang /* create and assemble symbolic parallel matrix C */ 11999566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,am,pn,PETSC_DETERMINE,PETSC_DETERMINE)); 12009566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,P)); 12019566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 12029566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 12039566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 1204d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 120517f75172Sandi selinger 12069566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api)); 12079566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 12089566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 12099566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 121017f75172Sandi selinger 12114222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable; 12124222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AB; 1213a9d7e0ccSandi selinger 12144222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 12156718818eSStefano Zampini C->product->data = ptap; 12166718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_MatMatMult; 1217a9d7e0ccSandi selinger 1218a9d7e0ccSandi selinger /* set MatInfo */ 121973b67375Sandi selinger afill = (PetscReal)api[am]/(adi[am]+aoi[am]+pi_loc[pm]+1) + 1.e-5; 122073b67375Sandi selinger if (afill < 1.0) afill = 1.0; 12214222ddf1SHong Zhang C->info.mallocs = nspacedouble; 12224222ddf1SHong Zhang C->info.fill_ratio_given = fill; 12234222ddf1SHong Zhang C->info.fill_ratio_needed = afill; 1224a9d7e0ccSandi selinger 1225a9d7e0ccSandi selinger #if defined(PETSC_USE_INFO) 1226a9d7e0ccSandi selinger if (api[am]) { 12279566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 12289566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n",(double)afill)); 1229a9d7e0ccSandi selinger } else { 12309566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 1231a9d7e0ccSandi selinger } 1232a9d7e0ccSandi selinger #endif 1233a9d7e0ccSandi selinger 12349566063dSJacob Faibussowitsch PetscCall(MatDestroy(&aopoth)); 12359566063dSJacob Faibussowitsch PetscCall(MatDestroy(&adpd)); 12369566063dSJacob Faibussowitsch PetscCall(PetscFree(j_temp)); 12379566063dSJacob Faibussowitsch PetscCall(PetscFree(adpoj)); 12389566063dSJacob Faibussowitsch PetscCall(PetscFree(adpoi)); 1239a9d7e0ccSandi selinger PetscFunctionReturn(0); 1240a9d7e0ccSandi selinger } 1241a9d7e0ccSandi selinger 1242f96d37fcSHong Zhang /*-------------------------------------------------------------------------*/ 1243c216dbf3SHong Zhang /* This routine only works when scall=MAT_REUSE_MATRIX! */ 1244c216dbf3SHong Zhang PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_matmatmult(Mat P,Mat A,Mat C) 1245c216dbf3SHong Zhang { 12466718818eSStefano Zampini Mat_APMPI *ptap; 124760552ceaSHong Zhang Mat Pt; 1248c216dbf3SHong Zhang 1249c216dbf3SHong Zhang PetscFunctionBegin; 12506718818eSStefano Zampini MatCheckProduct(C,3); 12516718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 125228b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 125328b400f6SJacob Faibussowitsch PetscCheck(ptap->Pt,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be reused. Do not call MatProductClear()"); 125460552ceaSHong Zhang 125560552ceaSHong Zhang Pt = ptap->Pt; 12569566063dSJacob Faibussowitsch PetscCall(MatTranspose(P,MAT_REUSE_MATRIX,&Pt)); 12579566063dSJacob Faibussowitsch PetscCall(MatMatMultNumeric_MPIAIJ_MPIAIJ(Pt,A,C)); 1258f96d37fcSHong Zhang PetscFunctionReturn(0); 1259f96d37fcSHong Zhang } 1260f96d37fcSHong Zhang 1261c5af039cSHong Zhang /* This routine is modified from MatPtAPSymbolic_MPIAIJ_MPIAIJ() */ 12624222ddf1SHong Zhang PetscErrorCode MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(Mat P,Mat A,PetscReal fill,Mat C) 1263912837bbSHong Zhang { 12643cdca5ebSHong Zhang Mat_APMPI *ptap; 12656718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data; 1266912837bbSHong Zhang MPI_Comm comm; 1267912837bbSHong Zhang PetscMPIInt size,rank; 1268912837bbSHong Zhang PetscFreeSpaceList free_space=NULL,current_space=NULL; 12694ede91e3SHong Zhang PetscInt pn=P->cmap->n,aN=A->cmap->N,an=A->cmap->n; 127069be949aSStefano Zampini PetscInt *lnk,i,k,nsend,rstart; 1271912837bbSHong Zhang PetscBT lnkbt; 1272ec4bef21SJose E. Roman PetscMPIInt tagi,tagj,*len_si,*len_s,*len_ri,nrecv; 1273ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted=0; 12748c78258cSHong Zhang PetscInt **buf_rj,**buf_ri,**buf_ri_k,row,ncols,*cols; 12757f3e1f30SHong Zhang PetscInt len,proc,*dnz,*onz,*owners,nzi; 1276912837bbSHong Zhang PetscInt nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci; 1277912837bbSHong Zhang MPI_Request *swaits,*rwaits; 1278912837bbSHong Zhang MPI_Status *sstatus,rstatus; 1279912837bbSHong Zhang PetscLayout rowmap; 1280912837bbSHong Zhang PetscInt *owners_co,*coi,*coj; /* i and j array of (p->B)^T*A*P - used in the communication */ 1281912837bbSHong Zhang PetscMPIInt *len_r,*id_r; /* array of length of comm->size, store send/recv matrix values */ 12827f3e1f30SHong Zhang PetscInt *Jptr,*prmap=p->garray,con,j,Crmax; 12834ede91e3SHong Zhang Mat_SeqAIJ *a_loc,*c_loc,*c_oth; 1284912837bbSHong Zhang PetscTable ta; 1285b92f168fSBarry Smith MatType mtype; 1286e83e5f86SFande Kong const char *prefix; 1287912837bbSHong Zhang 1288912837bbSHong Zhang PetscFunctionBegin; 12899566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 12909566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 12919566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1292912837bbSHong Zhang 12934222ddf1SHong Zhang /* create symbolic parallel matrix C */ 12949566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 12959566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 1296912837bbSHong Zhang 12974222ddf1SHong Zhang C->ops->transposematmultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable; 1298912837bbSHong Zhang 12993cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 13009566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 1301912837bbSHong Zhang ptap->reuse = MAT_INITIAL_MATRIX; 1302912837bbSHong Zhang 1303912837bbSHong Zhang /* (0) compute Rd = Pd^T, Ro = Po^T */ 1304912837bbSHong Zhang /* --------------------------------- */ 13059566063dSJacob Faibussowitsch PetscCall(MatTranspose_SeqAIJ(p->A,MAT_INITIAL_MATRIX,&ptap->Rd)); 13069566063dSJacob Faibussowitsch PetscCall(MatTranspose_SeqAIJ(p->B,MAT_INITIAL_MATRIX,&ptap->Ro)); 1307912837bbSHong Zhang 1308669e9c86SHong Zhang /* (1) compute symbolic A_loc */ 1309669e9c86SHong Zhang /* ---------------------------*/ 13109566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_INITIAL_MATRIX,&ptap->A_loc)); 1311912837bbSHong Zhang 13124ede91e3SHong Zhang /* (2-1) compute symbolic C_oth = Ro*A_loc */ 1313912837bbSHong Zhang /* ------------------------------------ */ 13149566063dSJacob Faibussowitsch PetscCall(MatGetOptionsPrefix(A,&prefix)); 13159566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(ptap->Ro,prefix)); 13169566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(ptap->Ro,"inner_offdiag_")); 13179566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF,&ptap->C_oth)); 13189566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Ro,ptap->A_loc,fill,ptap->C_oth)); 1319912837bbSHong Zhang 1320912837bbSHong Zhang /* (3) send coj of C_oth to other processors */ 1321912837bbSHong Zhang /* ------------------------------------------ */ 1322912837bbSHong Zhang /* determine row ownership */ 13239566063dSJacob Faibussowitsch PetscCall(PetscLayoutCreate(comm,&rowmap)); 1324912837bbSHong Zhang rowmap->n = pn; 1325912837bbSHong Zhang rowmap->bs = 1; 13269566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(rowmap)); 1327912837bbSHong Zhang owners = rowmap->range; 1328912837bbSHong Zhang 1329912837bbSHong Zhang /* determine the number of messages to send, their lengths */ 13309566063dSJacob Faibussowitsch PetscCall(PetscMalloc4(size,&len_s,size,&len_si,size,&sstatus,size+2,&owners_co)); 13319566063dSJacob Faibussowitsch PetscCall(PetscArrayzero(len_s,size)); 13329566063dSJacob Faibussowitsch PetscCall(PetscArrayzero(len_si,size)); 1333912837bbSHong Zhang 1334912837bbSHong Zhang c_oth = (Mat_SeqAIJ*)ptap->C_oth->data; 1335912837bbSHong Zhang coi = c_oth->i; coj = c_oth->j; 1336912837bbSHong Zhang con = ptap->C_oth->rmap->n; 1337912837bbSHong Zhang proc = 0; 1338912837bbSHong Zhang for (i=0; i<con; i++) { 1339912837bbSHong Zhang while (prmap[i] >= owners[proc+1]) proc++; 13407f3e1f30SHong Zhang len_si[proc]++; /* num of rows in Co(=Pt*A) to be sent to [proc] */ 1341912837bbSHong Zhang len_s[proc] += coi[i+1] - coi[i]; /* num of nonzeros in Co to be sent to [proc] */ 1342912837bbSHong Zhang } 1343912837bbSHong Zhang 1344912837bbSHong Zhang len = 0; /* max length of buf_si[], see (4) */ 1345912837bbSHong Zhang owners_co[0] = 0; 1346912837bbSHong Zhang nsend = 0; 1347912837bbSHong Zhang for (proc=0; proc<size; proc++) { 1348912837bbSHong Zhang owners_co[proc+1] = owners_co[proc] + len_si[proc]; 1349912837bbSHong Zhang if (len_s[proc]) { 1350912837bbSHong Zhang nsend++; 1351912837bbSHong Zhang len_si[proc] = 2*(len_si[proc] + 1); /* length of buf_si to be sent to [proc] */ 1352912837bbSHong Zhang len += len_si[proc]; 1353912837bbSHong Zhang } 1354912837bbSHong Zhang } 1355912837bbSHong Zhang 1356912837bbSHong Zhang /* determine the number and length of messages to receive for coi and coj */ 13579566063dSJacob Faibussowitsch PetscCall(PetscGatherNumberOfMessages(comm,NULL,len_s,&nrecv)); 13589566063dSJacob Faibussowitsch PetscCall(PetscGatherMessageLengths2(comm,nsend,nrecv,len_s,len_si,&id_r,&len_r,&len_ri)); 1359912837bbSHong Zhang 1360912837bbSHong Zhang /* post the Irecv and Isend of coj */ 13619566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagj)); 13629566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagj,nrecv,id_r,len_r,&buf_rj,&rwaits)); 13639566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(nsend+1,&swaits)); 1364912837bbSHong Zhang for (proc=0, k=0; proc<size; proc++) { 1365912837bbSHong Zhang if (!len_s[proc]) continue; 1366912837bbSHong Zhang i = owners_co[proc]; 13679566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tagj,comm,swaits+k)); 1368912837bbSHong Zhang k++; 1369912837bbSHong Zhang } 1370912837bbSHong Zhang 13717f3e1f30SHong Zhang /* (2-2) compute symbolic C_loc = Rd*A_loc */ 1372912837bbSHong Zhang /* ---------------------------------------- */ 13739566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(ptap->Rd,prefix)); 13749566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(ptap->Rd,"inner_diag_")); 13759566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF,&ptap->C_loc)); 13769566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Rd,ptap->A_loc,fill,ptap->C_loc)); 1377912837bbSHong Zhang c_loc = (Mat_SeqAIJ*)ptap->C_loc->data; 1378912837bbSHong Zhang 1379912837bbSHong Zhang /* receives coj are complete */ 1380912837bbSHong Zhang for (i=0; i<nrecv; i++) { 13819566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(nrecv,rwaits,&icompleted,&rstatus)); 1382912837bbSHong Zhang } 13839566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 13849566063dSJacob Faibussowitsch if (nsend) PetscCallMPI(MPI_Waitall(nsend,swaits,sstatus)); 1385912837bbSHong Zhang 1386912837bbSHong Zhang /* add received column indices into ta to update Crmax */ 13874ede91e3SHong Zhang a_loc = (Mat_SeqAIJ*)(ptap->A_loc)->data; 1388669e9c86SHong Zhang 1389669e9c86SHong Zhang /* create and initialize a linked list */ 13909566063dSJacob Faibussowitsch PetscCall(PetscTableCreate(an,aN,&ta)); /* for compute Crmax */ 13914ede91e3SHong Zhang MatRowMergeMax_SeqAIJ(a_loc,ptap->A_loc->rmap->N,ta); 1392669e9c86SHong Zhang 1393912837bbSHong Zhang for (k=0; k<nrecv; k++) {/* k-th received message */ 1394912837bbSHong Zhang Jptr = buf_rj[k]; 1395912837bbSHong Zhang for (j=0; j<len_r[k]; j++) { 13969566063dSJacob Faibussowitsch PetscCall(PetscTableAdd(ta,*(Jptr+j)+1,1,INSERT_VALUES)); 1397912837bbSHong Zhang } 1398912837bbSHong Zhang } 13999566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Crmax)); 14009566063dSJacob Faibussowitsch PetscCall(PetscTableDestroy(&ta)); 1401912837bbSHong Zhang 1402912837bbSHong Zhang /* (4) send and recv coi */ 1403912837bbSHong Zhang /*-----------------------*/ 14049566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagi)); 14059566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagi,nrecv,id_r,len_ri,&buf_ri,&rwaits)); 14069566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(len+1,&buf_s)); 1407912837bbSHong Zhang buf_si = buf_s; /* points to the beginning of k-th msg to be sent */ 1408912837bbSHong Zhang for (proc=0,k=0; proc<size; proc++) { 1409912837bbSHong Zhang if (!len_s[proc]) continue; 1410912837bbSHong Zhang /* form outgoing message for i-structure: 1411912837bbSHong Zhang buf_si[0]: nrows to be sent 1412912837bbSHong Zhang [1:nrows]: row index (global) 1413912837bbSHong Zhang [nrows+1:2*nrows+1]: i-structure index 1414912837bbSHong Zhang */ 1415912837bbSHong Zhang /*-------------------------------------------*/ 1416912837bbSHong Zhang nrows = len_si[proc]/2 - 1; /* num of rows in Co to be sent to [proc] */ 1417912837bbSHong Zhang buf_si_i = buf_si + nrows+1; 1418912837bbSHong Zhang buf_si[0] = nrows; 1419912837bbSHong Zhang buf_si_i[0] = 0; 1420912837bbSHong Zhang nrows = 0; 1421912837bbSHong Zhang for (i=owners_co[proc]; i<owners_co[proc+1]; i++) { 1422912837bbSHong Zhang nzi = coi[i+1] - coi[i]; 1423912837bbSHong Zhang buf_si_i[nrows+1] = buf_si_i[nrows] + nzi; /* i-structure */ 1424912837bbSHong Zhang buf_si[nrows+1] = prmap[i] -owners[proc]; /* local row index */ 1425912837bbSHong Zhang nrows++; 1426912837bbSHong Zhang } 14279566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tagi,comm,swaits+k)); 1428912837bbSHong Zhang k++; 1429912837bbSHong Zhang buf_si += len_si[proc]; 1430912837bbSHong Zhang } 1431912837bbSHong Zhang for (i=0; i<nrecv; i++) { 14329566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(nrecv,rwaits,&icompleted,&rstatus)); 1433912837bbSHong Zhang } 14349566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 14359566063dSJacob Faibussowitsch if (nsend) PetscCallMPI(MPI_Waitall(nsend,swaits,sstatus)); 1436912837bbSHong Zhang 14379566063dSJacob Faibussowitsch PetscCall(PetscFree4(len_s,len_si,sstatus,owners_co)); 14389566063dSJacob Faibussowitsch PetscCall(PetscFree(len_ri)); 14399566063dSJacob Faibussowitsch PetscCall(PetscFree(swaits)); 14409566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_s)); 1441912837bbSHong Zhang 14424222ddf1SHong Zhang /* (5) compute the local portion of C */ 1443912837bbSHong Zhang /* ------------------------------------------ */ 14444222ddf1SHong Zhang /* set initial free space to be Crmax, sufficient for holding nozeros in each row of C */ 14459566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(Crmax,&free_space)); 1446912837bbSHong Zhang current_space = free_space; 1447912837bbSHong Zhang 14489566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(nrecv,&buf_ri_k,nrecv,&nextrow,nrecv,&nextci)); 1449912837bbSHong Zhang for (k=0; k<nrecv; k++) { 1450912837bbSHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 1451912837bbSHong Zhang nrows = *buf_ri_k[k]; 1452912837bbSHong Zhang nextrow[k] = buf_ri_k[k] + 1; /* next row number of k-th recved i-structure */ 1453a5b23f4aSJose E. Roman nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th recved i-structure */ 1454912837bbSHong Zhang } 1455912837bbSHong Zhang 1456d0609cedSBarry Smith MatPreallocateBegin(comm,pn,an,dnz,onz); 14579566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate(Crmax,aN,&lnk,&lnkbt)); 14588c78258cSHong Zhang for (i=0; i<pn; i++) { /* for each local row of C */ 14594222ddf1SHong Zhang /* add C_loc into C */ 1460912837bbSHong Zhang nzi = c_loc->i[i+1] - c_loc->i[i]; 1461912837bbSHong Zhang Jptr = c_loc->j + c_loc->i[i]; 14629566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(nzi,Jptr,lnk,lnkbt)); 1463912837bbSHong Zhang 1464912837bbSHong Zhang /* add received col data into lnk */ 1465912837bbSHong Zhang for (k=0; k<nrecv; k++) { /* k-th received message */ 1466912837bbSHong Zhang if (i == *nextrow[k]) { /* i-th row */ 1467912837bbSHong Zhang nzi = *(nextci[k]+1) - *nextci[k]; 1468912837bbSHong Zhang Jptr = buf_rj[k] + *nextci[k]; 14699566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(nzi,Jptr,lnk,lnkbt)); 1470912837bbSHong Zhang nextrow[k]++; nextci[k]++; 1471912837bbSHong Zhang } 1472912837bbSHong Zhang } 14738c78258cSHong Zhang 14748c78258cSHong Zhang /* add missing diagonal entry */ 14758c78258cSHong Zhang if (C->force_diagonals) { 14768c78258cSHong Zhang k = i + owners[rank]; /* column index */ 14779566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(1,&k,lnk,lnkbt)); 14788c78258cSHong Zhang } 14798c78258cSHong Zhang 1480912837bbSHong Zhang nzi = lnk[0]; 1481912837bbSHong Zhang 1482912837bbSHong Zhang /* copy data into free space, then initialize lnk */ 14839566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean(aN,nzi,current_space->array,lnk,lnkbt)); 14849566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+owners[rank],nzi,current_space->array,dnz,onz)); 1485912837bbSHong Zhang } 14869566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 14879566063dSJacob Faibussowitsch PetscCall(PetscLLDestroy(lnk,lnkbt)); 14889566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceDestroy(free_space)); 1489912837bbSHong Zhang 1490912837bbSHong Zhang /* local sizes and preallocation */ 14919566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,pn,an,PETSC_DETERMINE,PETSC_DETERMINE)); 14929566063dSJacob Faibussowitsch if (P->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->rmap,P->cmap->bs)); 14939566063dSJacob Faibussowitsch if (A->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->cmap,A->cmap->bs)); 14949566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 1495d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 1496912837bbSHong Zhang 149769be949aSStefano Zampini /* add C_loc and C_oth to C */ 14989566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRange(C,&rstart,NULL)); 149969be949aSStefano Zampini for (i=0; i<pn; i++) { 15008c78258cSHong Zhang ncols = c_loc->i[i+1] - c_loc->i[i]; 15018c78258cSHong Zhang cols = c_loc->j + c_loc->i[i]; 15028c78258cSHong Zhang row = rstart + i; 15039566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,(const PetscInt*)&row,ncols,(const PetscInt*)cols,NULL,INSERT_VALUES)); 15048c78258cSHong Zhang 15058c78258cSHong Zhang if (C->force_diagonals) { 15069566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,(const PetscInt*)&row,1,(const PetscInt*)&row,NULL,INSERT_VALUES)); 15078c78258cSHong Zhang } 150869be949aSStefano Zampini } 150969be949aSStefano Zampini for (i=0; i<con; i++) { 15108c78258cSHong Zhang ncols = c_oth->i[i+1] - c_oth->i[i]; 15118c78258cSHong Zhang cols = c_oth->j + c_oth->i[i]; 15128c78258cSHong Zhang row = prmap[i]; 15139566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,(const PetscInt*)&row,ncols,(const PetscInt*)cols,NULL,INSERT_VALUES)); 151469be949aSStefano Zampini } 15159566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 15169566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 15179566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 151869be949aSStefano Zampini 1519912837bbSHong Zhang /* members in merge */ 15209566063dSJacob Faibussowitsch PetscCall(PetscFree(id_r)); 15219566063dSJacob Faibussowitsch PetscCall(PetscFree(len_r)); 15229566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_ri[0])); 15239566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_ri)); 15249566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_rj[0])); 15259566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_rj)); 15269566063dSJacob Faibussowitsch PetscCall(PetscLayoutDestroy(&rowmap)); 1527912837bbSHong Zhang 15284222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 15296718818eSStefano Zampini C->product->data = ptap; 15306718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 1531912837bbSHong Zhang PetscFunctionReturn(0); 1532912837bbSHong Zhang } 1533912837bbSHong Zhang 1534912837bbSHong Zhang PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable(Mat P,Mat A,Mat C) 1535912837bbSHong Zhang { 15366718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data; 1537669e9c86SHong Zhang Mat_SeqAIJ *c_seq; 15386718818eSStefano Zampini Mat_APMPI *ptap; 1539669e9c86SHong Zhang Mat A_loc,C_loc,C_oth; 1540912837bbSHong Zhang PetscInt i,rstart,rend,cm,ncols,row; 1541912837bbSHong Zhang const PetscInt *cols; 1542912837bbSHong Zhang const PetscScalar *vals; 1543912837bbSHong Zhang 1544912837bbSHong Zhang PetscFunctionBegin; 15456718818eSStefano Zampini MatCheckProduct(C,3); 15466718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 154728b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 154828b400f6SJacob Faibussowitsch PetscCheck(ptap->A_loc,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be reused. Do not call MatProductClear()"); 15499566063dSJacob Faibussowitsch PetscCall(MatZeroEntries(C)); 1550912837bbSHong Zhang 1551912837bbSHong Zhang if (ptap->reuse == MAT_REUSE_MATRIX) { 1552669e9c86SHong Zhang /* These matrices are obtained in MatTransposeMatMultSymbolic() */ 1553669e9c86SHong Zhang /* 1) get R = Pd^T, Ro = Po^T */ 1554669e9c86SHong Zhang /*----------------------------*/ 15559566063dSJacob Faibussowitsch PetscCall(MatTranspose_SeqAIJ(p->A,MAT_REUSE_MATRIX,&ptap->Rd)); 15569566063dSJacob Faibussowitsch PetscCall(MatTranspose_SeqAIJ(p->B,MAT_REUSE_MATRIX,&ptap->Ro)); 1557669e9c86SHong Zhang 1558669e9c86SHong Zhang /* 2) compute numeric A_loc */ 1559669e9c86SHong Zhang /*--------------------------*/ 15609566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_REUSE_MATRIX,&ptap->A_loc)); 1561912837bbSHong Zhang } 1562912837bbSHong Zhang 1563669e9c86SHong Zhang /* 3) C_loc = Rd*A_loc, C_oth = Ro*A_loc */ 15644ede91e3SHong Zhang A_loc = ptap->A_loc; 15659566063dSJacob Faibussowitsch PetscCall(((ptap->C_loc)->ops->matmultnumeric)(ptap->Rd,A_loc,ptap->C_loc)); 15669566063dSJacob Faibussowitsch PetscCall(((ptap->C_oth)->ops->matmultnumeric)(ptap->Ro,A_loc,ptap->C_oth)); 1567912837bbSHong Zhang C_loc = ptap->C_loc; 1568912837bbSHong Zhang C_oth = ptap->C_oth; 1569912837bbSHong Zhang 157069be949aSStefano Zampini /* add C_loc and C_oth to C */ 15719566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRange(C,&rstart,&rend)); 1572912837bbSHong Zhang 1573912837bbSHong Zhang /* C_loc -> C */ 1574912837bbSHong Zhang cm = C_loc->rmap->N; 1575912837bbSHong Zhang c_seq = (Mat_SeqAIJ*)C_loc->data; 1576912837bbSHong Zhang cols = c_seq->j; 1577912837bbSHong Zhang vals = c_seq->a; 1578912837bbSHong Zhang for (i=0; i<cm; i++) { 1579912837bbSHong Zhang ncols = c_seq->i[i+1] - c_seq->i[i]; 1580912837bbSHong Zhang row = rstart + i; 15819566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,ncols,cols,vals,ADD_VALUES)); 1582912837bbSHong Zhang cols += ncols; vals += ncols; 1583912837bbSHong Zhang } 1584912837bbSHong Zhang 1585912837bbSHong Zhang /* Co -> C, off-processor part */ 1586912837bbSHong Zhang cm = C_oth->rmap->N; 1587912837bbSHong Zhang c_seq = (Mat_SeqAIJ*)C_oth->data; 1588912837bbSHong Zhang cols = c_seq->j; 1589912837bbSHong Zhang vals = c_seq->a; 1590912837bbSHong Zhang for (i=0; i<cm; i++) { 1591912837bbSHong Zhang ncols = c_seq->i[i+1] - c_seq->i[i]; 1592912837bbSHong Zhang row = p->garray[i]; 15939566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,ncols,cols,vals,ADD_VALUES)); 1594912837bbSHong Zhang cols += ncols; vals += ncols; 1595912837bbSHong Zhang } 15969566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 15979566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 15989566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 1599912837bbSHong Zhang 1600912837bbSHong Zhang ptap->reuse = MAT_REUSE_MATRIX; 1601912837bbSHong Zhang PetscFunctionReturn(0); 1602912837bbSHong Zhang } 1603912837bbSHong Zhang 16046da69ca6SHong Zhang PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ(Mat P,Mat A,Mat C) 1605d6ab1dc0SHong Zhang { 1606d6ab1dc0SHong Zhang Mat_Merge_SeqsToMPI *merge; 16076718818eSStefano Zampini Mat_MPIAIJ *p =(Mat_MPIAIJ*)P->data; 1608d6ab1dc0SHong Zhang Mat_SeqAIJ *pd=(Mat_SeqAIJ*)(p->A)->data,*po=(Mat_SeqAIJ*)(p->B)->data; 16093cdca5ebSHong Zhang Mat_APMPI *ptap; 1610e745005bSHong Zhang PetscInt *adj; 1611e745005bSHong Zhang PetscInt i,j,k,anz,pnz,row,*cj,nexta; 1612e745005bSHong Zhang MatScalar *ada,*ca,valtmp; 1613d6ab1dc0SHong Zhang PetscInt am=A->rmap->n,cm=C->rmap->n,pon=(p->B)->cmap->n; 1614ce94432eSBarry Smith MPI_Comm comm; 1615d6ab1dc0SHong Zhang PetscMPIInt size,rank,taga,*len_s; 1616d6ab1dc0SHong Zhang PetscInt *owners,proc,nrows,**buf_ri_k,**nextrow,**nextci; 1617d6ab1dc0SHong Zhang PetscInt **buf_ri,**buf_rj; 1618d6ab1dc0SHong Zhang PetscInt cnz=0,*bj_i,*bi,*bj,bnz,nextcj; /* bi,bj,ba: local array of C(mpi mat) */ 1619d6ab1dc0SHong Zhang MPI_Request *s_waits,*r_waits; 1620d6ab1dc0SHong Zhang MPI_Status *status; 1621d6ab1dc0SHong Zhang MatScalar **abuf_r,*ba_i,*pA,*coa,*ba; 1622ce496241SStefano Zampini const PetscScalar *dummy; 16234e201bd7SHong Zhang PetscInt *ai,*aj,*coi,*coj,*poJ,*pdJ; 1624d6ab1dc0SHong Zhang Mat A_loc; 1625d6ab1dc0SHong Zhang Mat_SeqAIJ *a_loc; 1626d6ab1dc0SHong Zhang 1627d6ab1dc0SHong Zhang PetscFunctionBegin; 16286718818eSStefano Zampini MatCheckProduct(C,3); 16296718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 163028b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 163128b400f6SJacob Faibussowitsch PetscCheck(ptap->A_loc,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be reused. Do not call MatProductClear()"); 16329566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 16339566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 16349566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1635d6ab1dc0SHong Zhang 1636d6ab1dc0SHong Zhang merge = ptap->merge; 1637d6ab1dc0SHong Zhang 1638e745005bSHong Zhang /* 2) compute numeric C_seq = P_loc^T*A_loc */ 1639e745005bSHong Zhang /*------------------------------------------*/ 1640d6ab1dc0SHong Zhang /* get data from symbolic products */ 1641d6ab1dc0SHong Zhang coi = merge->coi; coj = merge->coj; 16429566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(coi[pon]+1,&coa)); 1643d6ab1dc0SHong Zhang bi = merge->bi; bj = merge->bj; 1644d6ab1dc0SHong Zhang owners = merge->rowmap->range; 16459566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(bi[cm]+1,&ba)); 1646d6ab1dc0SHong Zhang 1647d6ab1dc0SHong Zhang /* get A_loc by taking all local rows of A */ 1648d6ab1dc0SHong Zhang A_loc = ptap->A_loc; 16499566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_REUSE_MATRIX,&A_loc)); 1650d6ab1dc0SHong Zhang a_loc = (Mat_SeqAIJ*)(A_loc)->data; 1651d6ab1dc0SHong Zhang ai = a_loc->i; 1652d6ab1dc0SHong Zhang aj = a_loc->j; 1653d6ab1dc0SHong Zhang 1654ce496241SStefano Zampini /* trigger copy to CPU */ 16559566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(p->A,&dummy)); 16569566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(p->A,&dummy)); 16579566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(p->B,&dummy)); 16589566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(p->B,&dummy)); 1659d6ab1dc0SHong Zhang for (i=0; i<am; i++) { 1660d6ab1dc0SHong Zhang anz = ai[i+1] - ai[i]; 1661d6ab1dc0SHong Zhang adj = aj + ai[i]; 1662d6ab1dc0SHong Zhang ada = a_loc->a + ai[i]; 1663d6ab1dc0SHong Zhang 1664d6ab1dc0SHong Zhang /* 2-b) Compute Cseq = P_loc[i,:]^T*A[i,:] using outer product */ 1665e745005bSHong Zhang /*-------------------------------------------------------------*/ 1666d6ab1dc0SHong Zhang /* put the value into Co=(p->B)^T*A (off-diagonal part, send to others) */ 1667d6ab1dc0SHong Zhang pnz = po->i[i+1] - po->i[i]; 1668d6ab1dc0SHong Zhang poJ = po->j + po->i[i]; 1669d6ab1dc0SHong Zhang pA = po->a + po->i[i]; 1670d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1671d6ab1dc0SHong Zhang row = poJ[j]; 1672d6ab1dc0SHong Zhang cj = coj + coi[row]; 1673d6ab1dc0SHong Zhang ca = coa + coi[row]; 1674e745005bSHong Zhang /* perform sparse axpy */ 1675e745005bSHong Zhang nexta = 0; 1676d6ab1dc0SHong Zhang valtmp = pA[j]; 1677e745005bSHong Zhang for (k=0; nexta<anz; k++) { 1678e745005bSHong Zhang if (cj[k] == adj[nexta]) { 1679e745005bSHong Zhang ca[k] += valtmp*ada[nexta]; 1680e745005bSHong Zhang nexta++; 1681d6ab1dc0SHong Zhang } 1682e745005bSHong Zhang } 16839566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*anz)); 1684d6ab1dc0SHong Zhang } 1685d6ab1dc0SHong Zhang 1686d6ab1dc0SHong Zhang /* put the value into Cd (diagonal part) */ 1687d6ab1dc0SHong Zhang pnz = pd->i[i+1] - pd->i[i]; 1688d6ab1dc0SHong Zhang pdJ = pd->j + pd->i[i]; 1689d6ab1dc0SHong Zhang pA = pd->a + pd->i[i]; 1690d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1691d6ab1dc0SHong Zhang row = pdJ[j]; 1692d6ab1dc0SHong Zhang cj = bj + bi[row]; 1693d6ab1dc0SHong Zhang ca = ba + bi[row]; 1694e745005bSHong Zhang /* perform sparse axpy */ 1695e745005bSHong Zhang nexta = 0; 1696d6ab1dc0SHong Zhang valtmp = pA[j]; 1697e745005bSHong Zhang for (k=0; nexta<anz; k++) { 1698e745005bSHong Zhang if (cj[k] == adj[nexta]) { 1699e745005bSHong Zhang ca[k] += valtmp*ada[nexta]; 1700e745005bSHong Zhang nexta++; 1701d6ab1dc0SHong Zhang } 1702e745005bSHong Zhang } 17039566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*anz)); 1704d6ab1dc0SHong Zhang } 1705d6ab1dc0SHong Zhang } 1706d6ab1dc0SHong Zhang 1707d6ab1dc0SHong Zhang /* 3) send and recv matrix values coa */ 1708d6ab1dc0SHong Zhang /*------------------------------------*/ 1709d6ab1dc0SHong Zhang buf_ri = merge->buf_ri; 1710d6ab1dc0SHong Zhang buf_rj = merge->buf_rj; 1711d6ab1dc0SHong Zhang len_s = merge->len_s; 17129566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&taga)); 17139566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvScalar(comm,taga,merge->nrecv,merge->id_r,merge->len_r,&abuf_r,&r_waits)); 1714d6ab1dc0SHong Zhang 17159566063dSJacob Faibussowitsch PetscCall(PetscMalloc2(merge->nsend+1,&s_waits,size,&status)); 1716d6ab1dc0SHong Zhang for (proc=0,k=0; proc<size; proc++) { 1717d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1718d6ab1dc0SHong Zhang i = merge->owners_co[proc]; 17199566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coa+coi[i],len_s[proc],MPIU_MATSCALAR,proc,taga,comm,s_waits+k)); 1720d6ab1dc0SHong Zhang k++; 1721d6ab1dc0SHong Zhang } 17229566063dSJacob Faibussowitsch if (merge->nrecv) PetscCallMPI(MPI_Waitall(merge->nrecv,r_waits,status)); 17239566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,s_waits,status)); 1724d6ab1dc0SHong Zhang 17259566063dSJacob Faibussowitsch PetscCall(PetscFree2(s_waits,status)); 17269566063dSJacob Faibussowitsch PetscCall(PetscFree(r_waits)); 17279566063dSJacob Faibussowitsch PetscCall(PetscFree(coa)); 1728d6ab1dc0SHong Zhang 1729d6ab1dc0SHong Zhang /* 4) insert local Cseq and received values into Cmpi */ 1730d6ab1dc0SHong Zhang /*----------------------------------------------------*/ 17319566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(merge->nrecv,&buf_ri_k,merge->nrecv,&nextrow,merge->nrecv,&nextci)); 1732d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { 1733e745005bSHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 1734d6ab1dc0SHong Zhang nrows = *(buf_ri_k[k]); 1735d6ab1dc0SHong Zhang nextrow[k] = buf_ri_k[k]+1; /* next row number of k-th recved i-structure */ 1736a5b23f4aSJose E. Roman nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th recved i-structure */ 1737d6ab1dc0SHong Zhang } 1738d6ab1dc0SHong Zhang 1739d6ab1dc0SHong Zhang for (i=0; i<cm; i++) { 1740d6ab1dc0SHong Zhang row = owners[rank] + i; /* global row index of C_seq */ 1741d6ab1dc0SHong Zhang bj_i = bj + bi[i]; /* col indices of the i-th row of C */ 1742d6ab1dc0SHong Zhang ba_i = ba + bi[i]; 1743d6ab1dc0SHong Zhang bnz = bi[i+1] - bi[i]; 1744d6ab1dc0SHong Zhang /* add received vals into ba */ 1745d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { /* k-th received message */ 1746d6ab1dc0SHong Zhang /* i-th row */ 1747d6ab1dc0SHong Zhang if (i == *nextrow[k]) { 1748d6ab1dc0SHong Zhang cnz = *(nextci[k]+1) - *nextci[k]; 1749d6ab1dc0SHong Zhang cj = buf_rj[k] + *(nextci[k]); 1750d6ab1dc0SHong Zhang ca = abuf_r[k] + *(nextci[k]); 1751d6ab1dc0SHong Zhang nextcj = 0; 1752d6ab1dc0SHong Zhang for (j=0; nextcj<cnz; j++) { 1753d6ab1dc0SHong Zhang if (bj_i[j] == cj[nextcj]) { /* bcol == ccol */ 1754d6ab1dc0SHong Zhang ba_i[j] += ca[nextcj++]; 1755d6ab1dc0SHong Zhang } 1756d6ab1dc0SHong Zhang } 1757d6ab1dc0SHong Zhang nextrow[k]++; nextci[k]++; 17589566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*cnz)); 1759d6ab1dc0SHong Zhang } 1760d6ab1dc0SHong Zhang } 17619566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,bnz,bj_i,ba_i,INSERT_VALUES)); 1762d6ab1dc0SHong Zhang } 17639566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 17649566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 1765d6ab1dc0SHong Zhang 17669566063dSJacob Faibussowitsch PetscCall(PetscFree(ba)); 17679566063dSJacob Faibussowitsch PetscCall(PetscFree(abuf_r[0])); 17689566063dSJacob Faibussowitsch PetscCall(PetscFree(abuf_r)); 17699566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 1770d6ab1dc0SHong Zhang PetscFunctionReturn(0); 1771d6ab1dc0SHong Zhang } 1772d6ab1dc0SHong Zhang 17734222ddf1SHong Zhang PetscErrorCode MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ(Mat P,Mat A,PetscReal fill,Mat C) 1774d6ab1dc0SHong Zhang { 1775b419842dSStefano Zampini Mat A_loc; 17763cdca5ebSHong Zhang Mat_APMPI *ptap; 17770298fd71SBarry Smith PetscFreeSpaceList free_space=NULL,current_space=NULL; 17786718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data,*a=(Mat_MPIAIJ*)A->data; 1779d6ab1dc0SHong Zhang PetscInt *pdti,*pdtj,*poti,*potj,*ptJ; 1780d6ab1dc0SHong Zhang PetscInt nnz; 1781d6ab1dc0SHong Zhang PetscInt *lnk,*owners_co,*coi,*coj,i,k,pnz,row; 1782d6ab1dc0SHong Zhang PetscInt am =A->rmap->n,pn=P->cmap->n; 1783ce94432eSBarry Smith MPI_Comm comm; 1784d6ab1dc0SHong Zhang PetscMPIInt size,rank,tagi,tagj,*len_si,*len_s,*len_ri; 1785d6ab1dc0SHong Zhang PetscInt **buf_rj,**buf_ri,**buf_ri_k; 1786d6ab1dc0SHong Zhang PetscInt len,proc,*dnz,*onz,*owners; 1787d6ab1dc0SHong Zhang PetscInt nzi,*bi,*bj; 1788d6ab1dc0SHong Zhang PetscInt nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci; 1789d6ab1dc0SHong Zhang MPI_Request *swaits,*rwaits; 1790d6ab1dc0SHong Zhang MPI_Status *sstatus,rstatus; 1791d6ab1dc0SHong Zhang Mat_Merge_SeqsToMPI *merge; 1792d6ab1dc0SHong Zhang PetscInt *ai,*aj,*Jptr,anz,*prmap=p->garray,pon,nspacedouble=0,j; 1793d6ab1dc0SHong Zhang PetscReal afill =1.0,afill_tmp; 17944b38b95cSHong Zhang PetscInt rstart = P->cmap->rstart,rmax,aN=A->cmap->N,Armax; 1795b419842dSStefano Zampini Mat_SeqAIJ *a_loc; 17960acc65b4SHong Zhang PetscTable ta; 1797b92f168fSBarry Smith MatType mtype; 1798d6ab1dc0SHong Zhang 1799d6ab1dc0SHong Zhang PetscFunctionBegin; 18009566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 1801d6ab1dc0SHong Zhang /* check if matrix local sizes are compatible */ 18022c71b3e2SJacob Faibussowitsch PetscCheckFalse(A->rmap->rstart != P->rmap->rstart || A->rmap->rend != P->rmap->rend,comm,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, A (%" PetscInt_FMT ", %" PetscInt_FMT ") != P (%" PetscInt_FMT ",%" PetscInt_FMT ")",A->rmap->rstart,A->rmap->rend,P->rmap->rstart,P->rmap->rend); 1803d6ab1dc0SHong Zhang 18049566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 18059566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1806d6ab1dc0SHong Zhang 18073cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 18089566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 1809d6ab1dc0SHong Zhang 1810d6ab1dc0SHong Zhang /* get A_loc by taking all local rows of A */ 18119566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_INITIAL_MATRIX,&A_loc)); 18122205254eSKarl Rupp 1813d6ab1dc0SHong Zhang ptap->A_loc = A_loc; 1814d6ab1dc0SHong Zhang a_loc = (Mat_SeqAIJ*)(A_loc)->data; 1815d6ab1dc0SHong Zhang ai = a_loc->i; 1816d6ab1dc0SHong Zhang aj = a_loc->j; 1817d6ab1dc0SHong Zhang 1818d6ab1dc0SHong Zhang /* determine symbolic Co=(p->B)^T*A - send to others */ 1819d6ab1dc0SHong Zhang /*----------------------------------------------------*/ 18209566063dSJacob Faibussowitsch PetscCall(MatGetSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj)); 18219566063dSJacob Faibussowitsch PetscCall(MatGetSymbolicTranspose_SeqAIJ(p->B,&poti,&potj)); 1822d6ab1dc0SHong Zhang pon = (p->B)->cmap->n; /* total num of rows to be sent to other processors 1823d6ab1dc0SHong Zhang >= (num of nonzero rows of C_seq) - pn */ 18249566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pon+1,&coi)); 1825d6ab1dc0SHong Zhang coi[0] = 0; 1826d6ab1dc0SHong Zhang 1827d6ab1dc0SHong Zhang /* set initial free space to be fill*(nnz(p->B) + nnz(A)) */ 1828f91af8c7SBarry Smith nnz = PetscRealIntMultTruncate(fill,PetscIntSumTruncate(poti[pon],ai[am])); 18299566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(nnz,&free_space)); 1830d6ab1dc0SHong Zhang current_space = free_space; 183119f0e57aSHong Zhang 1832d6ab1dc0SHong Zhang /* create and initialize a linked list */ 18339566063dSJacob Faibussowitsch PetscCall(PetscTableCreate(A->cmap->n + a->B->cmap->N,aN,&ta)); 18344b38b95cSHong Zhang MatRowMergeMax_SeqAIJ(a_loc,am,ta); 18359566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Armax)); 183633ba5e2fSHong Zhang 18379566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(Armax,&lnk)); 1838d6ab1dc0SHong Zhang 1839d6ab1dc0SHong Zhang for (i=0; i<pon; i++) { 1840d6ab1dc0SHong Zhang pnz = poti[i+1] - poti[i]; 1841d6ab1dc0SHong Zhang ptJ = potj + poti[i]; 1842d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1843d6ab1dc0SHong Zhang row = ptJ[j]; /* row of A_loc == col of Pot */ 1844d6ab1dc0SHong Zhang anz = ai[row+1] - ai[row]; 1845d6ab1dc0SHong Zhang Jptr = aj + ai[row]; 1846d6ab1dc0SHong Zhang /* add non-zero cols of AP into the sorted linked list lnk */ 18479566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(anz,Jptr,lnk)); 1848d6ab1dc0SHong Zhang } 1849d6ab1dc0SHong Zhang nnz = lnk[0]; 1850d6ab1dc0SHong Zhang 1851d6ab1dc0SHong Zhang /* If free space is not available, double the total space in the list */ 1852d6ab1dc0SHong Zhang if (current_space->local_remaining<nnz) { 18539566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(nnz,current_space->total_array_size),¤t_space)); 1854d6ab1dc0SHong Zhang nspacedouble++; 1855d6ab1dc0SHong Zhang } 1856d6ab1dc0SHong Zhang 1857d6ab1dc0SHong Zhang /* Copy data into free space, and zero out denserows */ 18589566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(nnz,current_space->array,lnk)); 18592205254eSKarl Rupp 1860d6ab1dc0SHong Zhang current_space->array += nnz; 1861d6ab1dc0SHong Zhang current_space->local_used += nnz; 1862d6ab1dc0SHong Zhang current_space->local_remaining -= nnz; 18632205254eSKarl Rupp 1864d6ab1dc0SHong Zhang coi[i+1] = coi[i] + nnz; 1865d6ab1dc0SHong Zhang } 1866d6ab1dc0SHong Zhang 18679566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(coi[pon]+1,&coj)); 18689566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,coj)); 18699566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); /* must destroy to get a new one for C */ 18702205254eSKarl Rupp 1871118727c9SMark F. Adams afill_tmp = (PetscReal)coi[pon]/(poti[pon] + ai[am]+1); 1872d6ab1dc0SHong Zhang if (afill_tmp > afill) afill = afill_tmp; 1873d6ab1dc0SHong Zhang 1874d6ab1dc0SHong Zhang /* send j-array (coj) of Co to other processors */ 1875d6ab1dc0SHong Zhang /*----------------------------------------------*/ 1876d6ab1dc0SHong Zhang /* determine row ownership */ 18779566063dSJacob Faibussowitsch PetscCall(PetscNew(&merge)); 18789566063dSJacob Faibussowitsch PetscCall(PetscLayoutCreate(comm,&merge->rowmap)); 18792205254eSKarl Rupp 1880d6ab1dc0SHong Zhang merge->rowmap->n = pn; 1881d6ab1dc0SHong Zhang merge->rowmap->bs = 1; 18822205254eSKarl Rupp 18839566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(merge->rowmap)); 1884d6ab1dc0SHong Zhang owners = merge->rowmap->range; 1885d6ab1dc0SHong Zhang 1886d6ab1dc0SHong Zhang /* determine the number of messages to send, their lengths */ 18879566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(size,&len_si)); 18889566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(size,&merge->len_s)); 18892205254eSKarl Rupp 1890d6ab1dc0SHong Zhang len_s = merge->len_s; 1891d6ab1dc0SHong Zhang merge->nsend = 0; 1892d6ab1dc0SHong Zhang 18939566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size+2,&owners_co)); 1894d6ab1dc0SHong Zhang 1895d6ab1dc0SHong Zhang proc = 0; 1896d6ab1dc0SHong Zhang for (i=0; i<pon; i++) { 1897d6ab1dc0SHong Zhang while (prmap[i] >= owners[proc+1]) proc++; 1898d6ab1dc0SHong Zhang len_si[proc]++; /* num of rows in Co to be sent to [proc] */ 1899d6ab1dc0SHong Zhang len_s[proc] += coi[i+1] - coi[i]; 1900d6ab1dc0SHong Zhang } 1901d6ab1dc0SHong Zhang 1902d6ab1dc0SHong Zhang len = 0; /* max length of buf_si[] */ 1903d6ab1dc0SHong Zhang owners_co[0] = 0; 1904d6ab1dc0SHong Zhang for (proc=0; proc<size; proc++) { 1905d6ab1dc0SHong Zhang owners_co[proc+1] = owners_co[proc] + len_si[proc]; 1906d6ab1dc0SHong Zhang if (len_si[proc]) { 1907d6ab1dc0SHong Zhang merge->nsend++; 1908d6ab1dc0SHong Zhang len_si[proc] = 2*(len_si[proc] + 1); 1909d6ab1dc0SHong Zhang len += len_si[proc]; 1910d6ab1dc0SHong Zhang } 1911d6ab1dc0SHong Zhang } 1912d6ab1dc0SHong Zhang 1913d6ab1dc0SHong Zhang /* determine the number and length of messages to receive for coi and coj */ 19149566063dSJacob Faibussowitsch PetscCall(PetscGatherNumberOfMessages(comm,NULL,len_s,&merge->nrecv)); 19159566063dSJacob Faibussowitsch PetscCall(PetscGatherMessageLengths2(comm,merge->nsend,merge->nrecv,len_s,len_si,&merge->id_r,&merge->len_r,&len_ri)); 1916d6ab1dc0SHong Zhang 1917d6ab1dc0SHong Zhang /* post the Irecv and Isend of coj */ 19189566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagj)); 19199566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagj,merge->nrecv,merge->id_r,merge->len_r,&buf_rj,&rwaits)); 19209566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(merge->nsend+1,&swaits)); 1921d6ab1dc0SHong Zhang for (proc=0, k=0; proc<size; proc++) { 1922d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1923d6ab1dc0SHong Zhang i = owners_co[proc]; 19249566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tagj,comm,swaits+k)); 1925d6ab1dc0SHong Zhang k++; 1926d6ab1dc0SHong Zhang } 1927d6ab1dc0SHong Zhang 1928d6ab1dc0SHong Zhang /* receives and sends of coj are complete */ 19299566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size,&sstatus)); 1930d6ab1dc0SHong Zhang for (i=0; i<merge->nrecv; i++) { 1931ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted; 19329566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(merge->nrecv,rwaits,&icompleted,&rstatus)); 1933d6ab1dc0SHong Zhang } 19349566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 19359566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,swaits,sstatus)); 1936d6ab1dc0SHong Zhang 19370acc65b4SHong Zhang /* add received column indices into table to update Armax */ 1938c4762a1bSJed Brown /* Armax can be as large as aN if a P[row,:] is dense, see src/ksp/ksp/tutorials/ex56.c! */ 19390acc65b4SHong Zhang for (k=0; k<merge->nrecv; k++) {/* k-th received message */ 19400acc65b4SHong Zhang Jptr = buf_rj[k]; 19410acc65b4SHong Zhang for (j=0; j<merge->len_r[k]; j++) { 19429566063dSJacob Faibussowitsch PetscCall(PetscTableAdd(ta,*(Jptr+j)+1,1,INSERT_VALUES)); 19430acc65b4SHong Zhang } 19440acc65b4SHong Zhang } 19459566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Armax)); 19460acc65b4SHong Zhang 1947d6ab1dc0SHong Zhang /* send and recv coi */ 1948d6ab1dc0SHong Zhang /*-------------------*/ 19499566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagi)); 19509566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagi,merge->nrecv,merge->id_r,len_ri,&buf_ri,&rwaits)); 19519566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(len+1,&buf_s)); 1952d6ab1dc0SHong Zhang buf_si = buf_s; /* points to the beginning of k-th msg to be sent */ 1953d6ab1dc0SHong Zhang for (proc=0,k=0; proc<size; proc++) { 1954d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1955d6ab1dc0SHong Zhang /* form outgoing message for i-structure: 1956d6ab1dc0SHong Zhang buf_si[0]: nrows to be sent 1957d6ab1dc0SHong Zhang [1:nrows]: row index (global) 1958d6ab1dc0SHong Zhang [nrows+1:2*nrows+1]: i-structure index 1959d6ab1dc0SHong Zhang */ 1960d6ab1dc0SHong Zhang /*-------------------------------------------*/ 1961d6ab1dc0SHong Zhang nrows = len_si[proc]/2 - 1; 1962d6ab1dc0SHong Zhang buf_si_i = buf_si + nrows+1; 1963d6ab1dc0SHong Zhang buf_si[0] = nrows; 1964d6ab1dc0SHong Zhang buf_si_i[0] = 0; 1965d6ab1dc0SHong Zhang nrows = 0; 1966d6ab1dc0SHong Zhang for (i=owners_co[proc]; i<owners_co[proc+1]; i++) { 1967d6ab1dc0SHong Zhang nzi = coi[i+1] - coi[i]; 1968d6ab1dc0SHong Zhang buf_si_i[nrows+1] = buf_si_i[nrows] + nzi; /* i-structure */ 1969d6ab1dc0SHong Zhang buf_si[nrows+1] = prmap[i] -owners[proc]; /* local row index */ 1970d6ab1dc0SHong Zhang nrows++; 1971d6ab1dc0SHong Zhang } 19729566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tagi,comm,swaits+k)); 1973d6ab1dc0SHong Zhang k++; 1974d6ab1dc0SHong Zhang buf_si += len_si[proc]; 1975d6ab1dc0SHong Zhang } 1976d6ab1dc0SHong Zhang i = merge->nrecv; 1977d6ab1dc0SHong Zhang while (i--) { 1978ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted; 19799566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(merge->nrecv,rwaits,&icompleted,&rstatus)); 1980d6ab1dc0SHong Zhang } 19819566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 19829566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,swaits,sstatus)); 19839566063dSJacob Faibussowitsch PetscCall(PetscFree(len_si)); 19849566063dSJacob Faibussowitsch PetscCall(PetscFree(len_ri)); 19859566063dSJacob Faibussowitsch PetscCall(PetscFree(swaits)); 19869566063dSJacob Faibussowitsch PetscCall(PetscFree(sstatus)); 19879566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_s)); 1988d6ab1dc0SHong Zhang 1989d6ab1dc0SHong Zhang /* compute the local portion of C (mpi mat) */ 1990d6ab1dc0SHong Zhang /*------------------------------------------*/ 1991d6ab1dc0SHong Zhang /* allocate bi array and free space for accumulating nonzero column info */ 19929566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pn+1,&bi)); 1993d6ab1dc0SHong Zhang bi[0] = 0; 1994d6ab1dc0SHong Zhang 1995d6ab1dc0SHong Zhang /* set initial free space to be fill*(nnz(P) + nnz(AP)) */ 1996f91af8c7SBarry Smith nnz = PetscRealIntMultTruncate(fill,PetscIntSumTruncate(pdti[pn],PetscIntSumTruncate(poti[pon],ai[am]))); 19979566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(nnz,&free_space)); 1998d6ab1dc0SHong Zhang current_space = free_space; 1999d6ab1dc0SHong Zhang 20009566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(merge->nrecv,&buf_ri_k,merge->nrecv,&nextrow,merge->nrecv,&nextci)); 2001d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { 2002d6ab1dc0SHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 2003d6ab1dc0SHong Zhang nrows = *buf_ri_k[k]; 2004d6ab1dc0SHong Zhang nextrow[k] = buf_ri_k[k] + 1; /* next row number of k-th recved i-structure */ 20052d4ee042Sprj- nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th received i-structure */ 2006d6ab1dc0SHong Zhang } 2007d6ab1dc0SHong Zhang 20089566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(Armax,&lnk)); 2009d0609cedSBarry Smith MatPreallocateBegin(comm,pn,A->cmap->n,dnz,onz); 2010d6ab1dc0SHong Zhang rmax = 0; 2011d6ab1dc0SHong Zhang for (i=0; i<pn; i++) { 2012d6ab1dc0SHong Zhang /* add pdt[i,:]*AP into lnk */ 2013d6ab1dc0SHong Zhang pnz = pdti[i+1] - pdti[i]; 2014d6ab1dc0SHong Zhang ptJ = pdtj + pdti[i]; 2015d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 2016d6ab1dc0SHong Zhang row = ptJ[j]; /* row of AP == col of Pt */ 2017d6ab1dc0SHong Zhang anz = ai[row+1] - ai[row]; 2018d6ab1dc0SHong Zhang Jptr = aj + ai[row]; 2019d6ab1dc0SHong Zhang /* add non-zero cols of AP into the sorted linked list lnk */ 20209566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(anz,Jptr,lnk)); 2021d6ab1dc0SHong Zhang } 2022d6ab1dc0SHong Zhang 2023d6ab1dc0SHong Zhang /* add received col data into lnk */ 2024d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { /* k-th received message */ 2025d6ab1dc0SHong Zhang if (i == *nextrow[k]) { /* i-th row */ 2026d6ab1dc0SHong Zhang nzi = *(nextci[k]+1) - *nextci[k]; 2027d6ab1dc0SHong Zhang Jptr = buf_rj[k] + *nextci[k]; 20289566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(nzi,Jptr,lnk)); 2029d6ab1dc0SHong Zhang nextrow[k]++; nextci[k]++; 2030d6ab1dc0SHong Zhang } 2031d6ab1dc0SHong Zhang } 20328c78258cSHong Zhang 20338c78258cSHong Zhang /* add missing diagonal entry */ 20348c78258cSHong Zhang if (C->force_diagonals) { 20358c78258cSHong Zhang k = i + owners[rank]; /* column index */ 20369566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(1,&k,lnk)); 20378c78258cSHong Zhang } 20388c78258cSHong Zhang 2039d6ab1dc0SHong Zhang nnz = lnk[0]; 2040d6ab1dc0SHong Zhang 2041d6ab1dc0SHong Zhang /* if free space is not available, make more free space */ 2042d6ab1dc0SHong Zhang if (current_space->local_remaining<nnz) { 20439566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(nnz,current_space->total_array_size),¤t_space)); 2044d6ab1dc0SHong Zhang nspacedouble++; 2045d6ab1dc0SHong Zhang } 2046d6ab1dc0SHong Zhang /* copy data into free space, then initialize lnk */ 20479566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(nnz,current_space->array,lnk)); 20489566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+owners[rank],nnz,current_space->array,dnz,onz)); 20492205254eSKarl Rupp 2050d6ab1dc0SHong Zhang current_space->array += nnz; 2051d6ab1dc0SHong Zhang current_space->local_used += nnz; 2052d6ab1dc0SHong Zhang current_space->local_remaining -= nnz; 20532205254eSKarl Rupp 2054d6ab1dc0SHong Zhang bi[i+1] = bi[i] + nnz; 2055d6ab1dc0SHong Zhang if (nnz > rmax) rmax = nnz; 2056d6ab1dc0SHong Zhang } 20579566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 2058d6ab1dc0SHong Zhang 20599566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(bi[pn]+1,&bj)); 20609566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,bj)); 2061118727c9SMark F. Adams afill_tmp = (PetscReal)bi[pn]/(pdti[pn] + poti[pon] + ai[am]+1); 2062d6ab1dc0SHong Zhang if (afill_tmp > afill) afill = afill_tmp; 20639566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); 20649566063dSJacob Faibussowitsch PetscCall(PetscTableDestroy(&ta)); 20659566063dSJacob Faibussowitsch PetscCall(MatRestoreSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj)); 20669566063dSJacob Faibussowitsch PetscCall(MatRestoreSymbolicTranspose_SeqAIJ(p->B,&poti,&potj)); 2067d6ab1dc0SHong Zhang 20684222ddf1SHong Zhang /* create symbolic parallel matrix C - why cannot be assembled in Numeric part */ 20694222ddf1SHong Zhang /*-------------------------------------------------------------------------------*/ 20709566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,pn,A->cmap->n,PETSC_DETERMINE,PETSC_DETERMINE)); 20719566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizes(C,PetscAbs(P->cmap->bs),PetscAbs(A->cmap->bs))); 20729566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 20739566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 20749566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 2075d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 20769566063dSJacob Faibussowitsch PetscCall(MatSetBlockSize(C,1)); 20779566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NO_OFF_PROC_ENTRIES,PETSC_TRUE)); 2078d6ab1dc0SHong Zhang for (i=0; i<pn; i++) { 2079d6ab1dc0SHong Zhang row = i + rstart; 2080d6ab1dc0SHong Zhang nnz = bi[i+1] - bi[i]; 2081d6ab1dc0SHong Zhang Jptr = bj + bi[i]; 20829566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,nnz,Jptr,NULL,INSERT_VALUES)); 2083d6ab1dc0SHong Zhang } 20849566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 20859566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 20869566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 2087d6ab1dc0SHong Zhang merge->bi = bi; 2088d6ab1dc0SHong Zhang merge->bj = bj; 2089d6ab1dc0SHong Zhang merge->coi = coi; 2090d6ab1dc0SHong Zhang merge->coj = coj; 2091d6ab1dc0SHong Zhang merge->buf_ri = buf_ri; 2092d6ab1dc0SHong Zhang merge->buf_rj = buf_rj; 2093d6ab1dc0SHong Zhang merge->owners_co = owners_co; 2094d6ab1dc0SHong Zhang 20954222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 20966718818eSStefano Zampini C->product->data = ptap; 20976718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 2098d6ab1dc0SHong Zhang ptap->merge = merge; 2099a560ef98SHong Zhang 21004222ddf1SHong Zhang C->ops->mattransposemultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ; 2101f96d37fcSHong Zhang 2102f96d37fcSHong Zhang #if defined(PETSC_USE_INFO) 2103f96d37fcSHong Zhang if (bi[pn] != 0) { 21049566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 21059566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatTransposeMatMult(A,B,MatReuse,%g,&C) for best performance.\n",(double)afill)); 2106f96d37fcSHong Zhang } else { 21079566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 2108f96d37fcSHong Zhang } 2109f96d37fcSHong Zhang #endif 2110f96d37fcSHong Zhang PetscFunctionReturn(0); 2111f96d37fcSHong Zhang } 21124222ddf1SHong Zhang 21134222ddf1SHong Zhang /* ---------------------------------------------------------------- */ 21144222ddf1SHong Zhang static PetscErrorCode MatProductSymbolic_AtB_MPIAIJ_MPIAIJ(Mat C) 21154222ddf1SHong Zhang { 21164222ddf1SHong Zhang Mat_Product *product = C->product; 21174222ddf1SHong Zhang Mat A=product->A,B=product->B; 21184222ddf1SHong Zhang PetscReal fill=product->fill; 21194222ddf1SHong Zhang PetscBool flg; 21204222ddf1SHong Zhang 21214222ddf1SHong Zhang PetscFunctionBegin; 21224222ddf1SHong Zhang /* scalable */ 21239566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"scalable",&flg)); 21244222ddf1SHong Zhang if (flg) { 21259566063dSJacob Faibussowitsch PetscCall(MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ(A,B,fill,C)); 21264222ddf1SHong Zhang goto next; 21274222ddf1SHong Zhang } 21284222ddf1SHong Zhang 21294222ddf1SHong Zhang /* nonscalable */ 21309566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"nonscalable",&flg)); 21314222ddf1SHong Zhang if (flg) { 21329566063dSJacob Faibussowitsch PetscCall(MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(A,B,fill,C)); 21334222ddf1SHong Zhang goto next; 21344222ddf1SHong Zhang } 21354222ddf1SHong Zhang 21364222ddf1SHong Zhang /* matmatmult */ 21379566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"at*b",&flg)); 21384222ddf1SHong Zhang if (flg) { 21394222ddf1SHong Zhang Mat At; 21404222ddf1SHong Zhang Mat_APMPI *ptap; 21414222ddf1SHong Zhang 21429566063dSJacob Faibussowitsch PetscCall(MatTranspose(A,MAT_INITIAL_MATRIX,&At)); 21439566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ(At,B,fill,C)); 21446718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 21454222ddf1SHong Zhang if (ptap) { 21464222ddf1SHong Zhang ptap->Pt = At; 21476718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 21484222ddf1SHong Zhang } 21494222ddf1SHong Zhang C->ops->transposematmultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_matmatmult; 21504222ddf1SHong Zhang goto next; 21514222ddf1SHong Zhang } 21524222ddf1SHong Zhang 21534e84afc0SStefano Zampini /* backend general code */ 21549566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"backend",&flg)); 21554e84afc0SStefano Zampini if (flg) { 21569566063dSJacob Faibussowitsch PetscCall(MatProductSymbolic_MPIAIJBACKEND(C)); 21574e84afc0SStefano Zampini PetscFunctionReturn(0); 21584e84afc0SStefano Zampini } 21594e84afc0SStefano Zampini 21604222ddf1SHong Zhang SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"MatProduct type is not supported"); 21614222ddf1SHong Zhang 21624222ddf1SHong Zhang next: 21634222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AtB; 21644222ddf1SHong Zhang PetscFunctionReturn(0); 21654222ddf1SHong Zhang } 21664222ddf1SHong Zhang 21674222ddf1SHong Zhang /* ---------------------------------------------------------------- */ 21684222ddf1SHong Zhang /* Set options for MatMatMultxxx_MPIAIJ_MPIAIJ */ 21694222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_AB(Mat C) 21704222ddf1SHong Zhang { 21714222ddf1SHong Zhang Mat_Product *product = C->product; 21724222ddf1SHong Zhang Mat A=product->A,B=product->B; 21734222ddf1SHong Zhang #if defined(PETSC_HAVE_HYPRE) 21744e84afc0SStefano Zampini const char *algTypes[5] = {"scalable","nonscalable","seqmpi","backend","hypre"}; 21754e84afc0SStefano Zampini PetscInt nalg = 5; 21764222ddf1SHong Zhang #else 21774e84afc0SStefano Zampini const char *algTypes[4] = {"scalable","nonscalable","seqmpi","backend",}; 21784e84afc0SStefano Zampini PetscInt nalg = 4; 21794222ddf1SHong Zhang #endif 21804222ddf1SHong Zhang PetscInt alg = 1; /* set nonscalable algorithm as default */ 21814222ddf1SHong Zhang PetscBool flg; 21824222ddf1SHong Zhang MPI_Comm comm; 21834222ddf1SHong Zhang 21844222ddf1SHong Zhang PetscFunctionBegin; 21854222ddf1SHong Zhang /* Check matrix local sizes */ 21869566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 21874222ddf1SHong Zhang 21884222ddf1SHong Zhang /* Set "nonscalable" as default algorithm */ 21899566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 21904222ddf1SHong Zhang if (flg) { 21919566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 21924222ddf1SHong Zhang 21934222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 21944222ddf1SHong Zhang if (B->cmap->N > 100000) { /* may switch to scalable algorithm as default */ 21954222ddf1SHong Zhang MatInfo Ainfo,Binfo; 21964222ddf1SHong Zhang PetscInt nz_local; 21974222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 21984222ddf1SHong Zhang 21999566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 22009566063dSJacob Faibussowitsch PetscCall(MatGetInfo(B,MAT_LOCAL,&Binfo)); 22014222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Binfo.nz_allocated); 22024222ddf1SHong Zhang 22034222ddf1SHong Zhang if (B->cmap->N > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 22041c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 22054222ddf1SHong Zhang 22064222ddf1SHong Zhang if (alg_scalable) { 22074222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 22089566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22099566063dSJacob Faibussowitsch PetscCall(PetscInfo(B,"Use scalable algorithm, BN %" PetscInt_FMT ", fill*nz_allocated %g\n",B->cmap->N,(double)(product->fill*nz_local))); 22104222ddf1SHong Zhang } 22114222ddf1SHong Zhang } 22124222ddf1SHong Zhang } 22134222ddf1SHong Zhang 22144222ddf1SHong Zhang /* Get runtime option */ 22154222ddf1SHong Zhang if (product->api_user) { 2216d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatMatMult","Mat"); 22179566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matmatmult_via","Algorithmic approach","MatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2218d0609cedSBarry Smith PetscOptionsEnd(); 22194222ddf1SHong Zhang } else { 2220d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_AB","Mat"); 22219566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2222d0609cedSBarry Smith PetscOptionsEnd(); 22234222ddf1SHong Zhang } 22244222ddf1SHong Zhang if (flg) { 22259566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22264222ddf1SHong Zhang } 22274222ddf1SHong Zhang 22284222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AB_MPIAIJ_MPIAIJ; 22294222ddf1SHong Zhang PetscFunctionReturn(0); 22304222ddf1SHong Zhang } 22314222ddf1SHong Zhang 2232*8a9c020eSBarry Smith static PetscErrorCode MatProductSetFromOptions_MPIAIJ_ABt(Mat C) 2233*8a9c020eSBarry Smith { 2234*8a9c020eSBarry Smith PetscFunctionBegin; 2235*8a9c020eSBarry Smith PetscCall(MatProductSetFromOptions_MPIAIJ_AB(C)); 2236*8a9c020eSBarry Smith C->ops->productsymbolic = MatProductSymbolic_ABt_MPIAIJ_MPIAIJ; 2237*8a9c020eSBarry Smith PetscFunctionReturn(0); 2238*8a9c020eSBarry Smith } 2239*8a9c020eSBarry Smith 22404222ddf1SHong Zhang /* Set options for MatTransposeMatMultXXX_MPIAIJ_MPIAIJ */ 22414222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_AtB(Mat C) 22424222ddf1SHong Zhang { 22434222ddf1SHong Zhang Mat_Product *product = C->product; 22444222ddf1SHong Zhang Mat A=product->A,B=product->B; 22454e84afc0SStefano Zampini const char *algTypes[4] = {"scalable","nonscalable","at*b","backend"}; 22464e84afc0SStefano Zampini PetscInt nalg = 4; 22474222ddf1SHong Zhang PetscInt alg = 1; /* set default algorithm */ 22484222ddf1SHong Zhang PetscBool flg; 22494222ddf1SHong Zhang MPI_Comm comm; 22504222ddf1SHong Zhang 22514222ddf1SHong Zhang PetscFunctionBegin; 22524222ddf1SHong Zhang /* Check matrix local sizes */ 22539566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 22542c71b3e2SJacob Faibussowitsch PetscCheckFalse(A->rmap->rstart != B->rmap->rstart || A->rmap->rend != B->rmap->rend,PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, A (%" PetscInt_FMT ", %" PetscInt_FMT ") != B (%" PetscInt_FMT ",%" PetscInt_FMT ")",A->rmap->rstart,A->rmap->rend,B->rmap->rstart,B->rmap->rend); 22554222ddf1SHong Zhang 22564222ddf1SHong Zhang /* Set default algorithm */ 22579566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 22584222ddf1SHong Zhang if (flg) { 22599566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22604222ddf1SHong Zhang } 22614222ddf1SHong Zhang 22624222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 22634222ddf1SHong Zhang if (alg && B->cmap->N > 100000) { /* may switch to scalable algorithm as default */ 22644222ddf1SHong Zhang MatInfo Ainfo,Binfo; 22654222ddf1SHong Zhang PetscInt nz_local; 22664222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 22674222ddf1SHong Zhang 22689566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 22699566063dSJacob Faibussowitsch PetscCall(MatGetInfo(B,MAT_LOCAL,&Binfo)); 22704222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Binfo.nz_allocated); 22714222ddf1SHong Zhang 22724222ddf1SHong Zhang if (B->cmap->N > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 22731c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 22744222ddf1SHong Zhang 22754222ddf1SHong Zhang if (alg_scalable) { 22764222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 22779566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22789566063dSJacob Faibussowitsch PetscCall(PetscInfo(B,"Use scalable algorithm, BN %" PetscInt_FMT ", fill*nz_allocated %g\n",B->cmap->N,(double)(product->fill*nz_local))); 22794222ddf1SHong Zhang } 22804222ddf1SHong Zhang } 22814222ddf1SHong Zhang 22824222ddf1SHong Zhang /* Get runtime option */ 22834222ddf1SHong Zhang if (product->api_user) { 2284d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatTransposeMatMult","Mat"); 22859566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mattransposematmult_via","Algorithmic approach","MatTransposeMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2286d0609cedSBarry Smith PetscOptionsEnd(); 22874222ddf1SHong Zhang } else { 2288d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_AtB","Mat"); 22899566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatTransposeMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2290d0609cedSBarry Smith PetscOptionsEnd(); 22914222ddf1SHong Zhang } 22924222ddf1SHong Zhang if (flg) { 22939566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22944222ddf1SHong Zhang } 22954222ddf1SHong Zhang 22964222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AtB_MPIAIJ_MPIAIJ; 22974222ddf1SHong Zhang PetscFunctionReturn(0); 22984222ddf1SHong Zhang } 22994222ddf1SHong Zhang 23004222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_PtAP(Mat C) 23014222ddf1SHong Zhang { 23024222ddf1SHong Zhang Mat_Product *product = C->product; 23034222ddf1SHong Zhang Mat A=product->A,P=product->B; 23044222ddf1SHong Zhang MPI_Comm comm; 23054222ddf1SHong Zhang PetscBool flg; 23064222ddf1SHong Zhang PetscInt alg=1; /* set default algorithm */ 23074222ddf1SHong Zhang #if !defined(PETSC_HAVE_HYPRE) 23084e84afc0SStefano Zampini const char *algTypes[5] = {"scalable","nonscalable","allatonce","allatonce_merged","backend"}; 23094222ddf1SHong Zhang PetscInt nalg=5; 23104e84afc0SStefano Zampini #else 23114e84afc0SStefano Zampini const char *algTypes[6] = {"scalable","nonscalable","allatonce","allatonce_merged","backend","hypre"}; 23124e84afc0SStefano Zampini PetscInt nalg=6; 23134222ddf1SHong Zhang #endif 23144222ddf1SHong Zhang PetscInt pN=P->cmap->N; 23154222ddf1SHong Zhang 23164222ddf1SHong Zhang PetscFunctionBegin; 23174222ddf1SHong Zhang /* Check matrix local sizes */ 23189566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 23192c71b3e2SJacob Faibussowitsch PetscCheckFalse(A->rmap->rstart != P->rmap->rstart || A->rmap->rend != P->rmap->rend,PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, Arow (%" PetscInt_FMT ", %" PetscInt_FMT ") != Prow (%" PetscInt_FMT ",%" PetscInt_FMT ")",A->rmap->rstart,A->rmap->rend,P->rmap->rstart,P->rmap->rend); 23202c71b3e2SJacob Faibussowitsch PetscCheckFalse(A->cmap->rstart != P->rmap->rstart || A->cmap->rend != P->rmap->rend,PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, Acol (%" PetscInt_FMT ", %" PetscInt_FMT ") != Prow (%" PetscInt_FMT ",%" PetscInt_FMT ")",A->cmap->rstart,A->cmap->rend,P->rmap->rstart,P->rmap->rend); 23214222ddf1SHong Zhang 23224222ddf1SHong Zhang /* Set "nonscalable" as default algorithm */ 23239566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 23244222ddf1SHong Zhang if (flg) { 23259566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23264222ddf1SHong Zhang 23274222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 23284222ddf1SHong Zhang if (pN > 100000) { 23294222ddf1SHong Zhang MatInfo Ainfo,Pinfo; 23304222ddf1SHong Zhang PetscInt nz_local; 23314222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 23324222ddf1SHong Zhang 23339566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 23349566063dSJacob Faibussowitsch PetscCall(MatGetInfo(P,MAT_LOCAL,&Pinfo)); 23354222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Pinfo.nz_allocated); 23364222ddf1SHong Zhang 23374222ddf1SHong Zhang if (pN > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 23381c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 23394222ddf1SHong Zhang 23404222ddf1SHong Zhang if (alg_scalable) { 23414222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 23429566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23434222ddf1SHong Zhang } 23444222ddf1SHong Zhang } 23454222ddf1SHong Zhang } 23464222ddf1SHong Zhang 23474222ddf1SHong Zhang /* Get runtime option */ 23484222ddf1SHong Zhang if (product->api_user) { 2349d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatPtAP","Mat"); 23509566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matptap_via","Algorithmic approach","MatPtAP",algTypes,nalg,algTypes[alg],&alg,&flg)); 2351d0609cedSBarry Smith PetscOptionsEnd(); 23524222ddf1SHong Zhang } else { 2353d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_PtAP","Mat"); 23549566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatPtAP",algTypes,nalg,algTypes[alg],&alg,&flg)); 2355d0609cedSBarry Smith PetscOptionsEnd(); 23564222ddf1SHong Zhang } 23574222ddf1SHong Zhang if (flg) { 23589566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23594222ddf1SHong Zhang } 23604222ddf1SHong Zhang 23614222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_PtAP_MPIAIJ_MPIAIJ; 23624222ddf1SHong Zhang PetscFunctionReturn(0); 23634222ddf1SHong Zhang } 23644222ddf1SHong Zhang 23654222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_RARt(Mat C) 23664222ddf1SHong Zhang { 23674222ddf1SHong Zhang Mat_Product *product = C->product; 23684222ddf1SHong Zhang Mat A = product->A,R=product->B; 23694222ddf1SHong Zhang 23704222ddf1SHong Zhang PetscFunctionBegin; 23714222ddf1SHong Zhang /* Check matrix local sizes */ 23722c71b3e2SJacob Faibussowitsch PetscCheckFalse(A->cmap->n != R->cmap->n || A->rmap->n != R->cmap->n,PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, A local (%" PetscInt_FMT ", %" PetscInt_FMT "), R local (%" PetscInt_FMT ",%" PetscInt_FMT ")",A->rmap->n,A->rmap->n,R->rmap->n,R->cmap->n); 23734222ddf1SHong Zhang 23744222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_RARt_MPIAIJ_MPIAIJ; 23754222ddf1SHong Zhang PetscFunctionReturn(0); 23764222ddf1SHong Zhang } 23774222ddf1SHong Zhang 23784222ddf1SHong Zhang /* 23794222ddf1SHong Zhang Set options for ABC = A*B*C = A*(B*C); ABC's algorithm must be chosen from AB's algorithm 23804222ddf1SHong Zhang */ 23814222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_ABC(Mat C) 23824222ddf1SHong Zhang { 23834222ddf1SHong Zhang Mat_Product *product = C->product; 23844222ddf1SHong Zhang PetscBool flg = PETSC_FALSE; 23854222ddf1SHong Zhang PetscInt alg = 1; /* default algorithm */ 23864222ddf1SHong Zhang const char *algTypes[3] = {"scalable","nonscalable","seqmpi"}; 23874222ddf1SHong Zhang PetscInt nalg = 3; 23884222ddf1SHong Zhang 23894222ddf1SHong Zhang PetscFunctionBegin; 23904222ddf1SHong Zhang /* Set default algorithm */ 23919566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 23924222ddf1SHong Zhang if (flg) { 23939566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23944222ddf1SHong Zhang } 23954222ddf1SHong Zhang 23964222ddf1SHong Zhang /* Get runtime option */ 23974222ddf1SHong Zhang if (product->api_user) { 2398d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatMatMatMult","Mat"); 23999566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matmatmatmult_via","Algorithmic approach","MatMatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2400d0609cedSBarry Smith PetscOptionsEnd(); 24014222ddf1SHong Zhang } else { 2402d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_ABC","Mat"); 24039566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatProduct_ABC",algTypes,nalg,algTypes[alg],&alg,&flg)); 2404d0609cedSBarry Smith PetscOptionsEnd(); 24054222ddf1SHong Zhang } 24064222ddf1SHong Zhang if (flg) { 24079566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 24084222ddf1SHong Zhang } 24094222ddf1SHong Zhang 24104222ddf1SHong Zhang C->ops->matmatmultsymbolic = MatMatMatMultSymbolic_MPIAIJ_MPIAIJ_MPIAIJ; 24114222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_ABC; 24124222ddf1SHong Zhang PetscFunctionReturn(0); 24134222ddf1SHong Zhang } 24144222ddf1SHong Zhang 24154222ddf1SHong Zhang PETSC_INTERN PetscErrorCode MatProductSetFromOptions_MPIAIJ(Mat C) 24164222ddf1SHong Zhang { 24174222ddf1SHong Zhang Mat_Product *product = C->product; 24184222ddf1SHong Zhang 24194222ddf1SHong Zhang PetscFunctionBegin; 24204222ddf1SHong Zhang switch (product->type) { 24214222ddf1SHong Zhang case MATPRODUCT_AB: 24229566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_AB(C)); 24234222ddf1SHong Zhang break; 2424*8a9c020eSBarry Smith case MATPRODUCT_ABt: 2425*8a9c020eSBarry Smith PetscCall(MatProductSetFromOptions_MPIAIJ_ABt(C)); 2426*8a9c020eSBarry Smith break; 24274222ddf1SHong Zhang case MATPRODUCT_AtB: 24289566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_AtB(C)); 24294222ddf1SHong Zhang break; 24304222ddf1SHong Zhang case MATPRODUCT_PtAP: 24319566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_PtAP(C)); 24324222ddf1SHong Zhang break; 24334222ddf1SHong Zhang case MATPRODUCT_RARt: 24349566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_RARt(C)); 24354222ddf1SHong Zhang break; 24364222ddf1SHong Zhang case MATPRODUCT_ABC: 24379566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_ABC(C)); 24384222ddf1SHong Zhang break; 24396718818eSStefano Zampini default: 24406718818eSStefano Zampini break; 24414222ddf1SHong Zhang } 24424222ddf1SHong Zhang PetscFunctionReturn(0); 24434222ddf1SHong Zhang } 2444