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 188a9c020eSBarry Smith PETSC_INTERN PetscErrorCode MatProductSymbolic_ABt_MPIAIJ_MPIAIJ(Mat C) 198a9c020eSBarry Smith { 208a9c020eSBarry Smith Mat_Product *product = C->product; 218a9c020eSBarry Smith Mat B=product->B; 228a9c020eSBarry Smith 238a9c020eSBarry Smith PetscFunctionBegin; 248a9c020eSBarry Smith PetscCall(MatTranspose(B,MAT_INITIAL_MATRIX,&product->B)); 258a9c020eSBarry Smith PetscCall(MatDestroy(&B)); 268a9c020eSBarry Smith PetscCall(MatProductSymbolic_AB_MPIAIJ_MPIAIJ(C)); 278a9c020eSBarry Smith PetscFunctionReturn(0); 288a9c020eSBarry Smith } 298a9c020eSBarry 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)); 117aed4548fSBarry Smith PetscCheck(ptap->P_oth || size <= 1,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"AP cannot be reused. Do not call MatProductClear()"); 11880da8df7SHong Zhang 11965e4b4d4SStefano Zampini /* flag CPU mask for C */ 12065e4b4d4SStefano Zampini #if defined(PETSC_HAVE_DEVICE) 12165e4b4d4SStefano Zampini if (C->offloadmask != PETSC_OFFLOAD_UNALLOCATED) C->offloadmask = PETSC_OFFLOAD_CPU; 12265e4b4d4SStefano Zampini if (c->A->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->A->offloadmask = PETSC_OFFLOAD_CPU; 12365e4b4d4SStefano Zampini if (c->B->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->B->offloadmask = PETSC_OFFLOAD_CPU; 12465e4b4d4SStefano Zampini #endif 12565e4b4d4SStefano Zampini 126a1a4e84aSHong Zhang /* 1) get P_oth = ptap->P_oth and P_loc = ptap->P_loc */ 127598bc09dSHong Zhang /*-----------------------------------------------------*/ 128a1a4e84aSHong Zhang /* update numerical values of P_oth and P_loc */ 1299566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_REUSE_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 1309566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_REUSE_MATRIX,&ptap->P_loc)); 131598bc09dSHong Zhang 132598bc09dSHong Zhang /* 2) compute numeric C_loc = A_loc*P = Ad*P_loc + Ao*P_oth */ 133598bc09dSHong Zhang /*----------------------------------------------------------*/ 134598bc09dSHong Zhang /* get data from symbolic products */ 135a1a4e84aSHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 136c12557a1SHong Zhang p_oth = NULL; 1379467f45fSHong Zhang if (size >1) { 1389467f45fSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 1399467f45fSHong Zhang } 140598bc09dSHong Zhang 141598bc09dSHong Zhang /* get apa for storing dense row A[i,:]*P */ 142598bc09dSHong Zhang apa = ptap->apa; 143598bc09dSHong Zhang 14429825010SHong Zhang api = ptap->api; 14529825010SHong Zhang apj = ptap->apj; 146ce496241SStefano Zampini /* trigger copy to CPU */ 1479566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->A,&dummy)); 1489566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->A,&dummy)); 1499566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->B,&dummy)); 1509566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->B,&dummy)); 151598bc09dSHong Zhang for (i=0; i<cm; i++) { 152c12557a1SHong Zhang /* compute apa = A[i,:]*P */ 153e497d3c8SHong Zhang AProw_nonscalable(i,ad,ao,p_loc,p_oth,apa); 154598bc09dSHong Zhang 155598bc09dSHong Zhang /* set values in C */ 156598bc09dSHong Zhang apJ = apj + api[i]; 157598bc09dSHong Zhang cdnz = cd->i[i+1] - cd->i[i]; 158598bc09dSHong Zhang conz = co->i[i+1] - co->i[i]; 159598bc09dSHong Zhang 160910cf402Sprj- /* 1st off-diagonal part of C */ 161598bc09dSHong Zhang ca = coa + co->i[i]; 162598bc09dSHong Zhang k = 0; 163598bc09dSHong Zhang for (k0=0; k0<conz; k0++) { 164598bc09dSHong Zhang if (apJ[k] >= cstart) break; 165598bc09dSHong Zhang ca[k0] = apa[apJ[k]]; 1665cab4f04SHong Zhang apa[apJ[k++]] = 0.0; 167598bc09dSHong Zhang } 168598bc09dSHong Zhang 169598bc09dSHong Zhang /* diagonal part of C */ 170598bc09dSHong Zhang ca = cda + cd->i[i]; 171598bc09dSHong Zhang for (k1=0; k1<cdnz; k1++) { 172598bc09dSHong Zhang ca[k1] = apa[apJ[k]]; 1735cab4f04SHong Zhang apa[apJ[k++]] = 0.0; 174598bc09dSHong Zhang } 175598bc09dSHong Zhang 176910cf402Sprj- /* 2nd off-diagonal part of C */ 177598bc09dSHong Zhang ca = coa + co->i[i]; 178598bc09dSHong Zhang for (; k0<conz; k0++) { 179598bc09dSHong Zhang ca[k0] = apa[apJ[k]]; 1805cab4f04SHong Zhang apa[apJ[k++]] = 0.0; 181598bc09dSHong Zhang } 182598bc09dSHong Zhang } 1839566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 1849566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 185598bc09dSHong Zhang PetscFunctionReturn(0); 186598bc09dSHong Zhang } 187598bc09dSHong Zhang 1884222ddf1SHong Zhang PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(Mat A,Mat P,PetscReal fill,Mat C) 189598bc09dSHong Zhang { 190ce94432eSBarry Smith MPI_Comm comm; 1919467f45fSHong Zhang PetscMPIInt size; 1923cdca5ebSHong Zhang Mat_APMPI *ptap; 1930298fd71SBarry Smith PetscFreeSpaceList free_space=NULL,current_space=NULL; 1946718818eSStefano Zampini Mat_MPIAIJ *a=(Mat_MPIAIJ*)A->data; 195bfcd1a73SHong Zhang Mat_SeqAIJ *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_loc,*p_oth; 1967dbaa2c6Sandi selinger PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*dnz,*onz; 1974ae0847bSHong Zhang PetscInt *adi=ad->i,*adj=ad->j,*aoi=ao->i,*aoj=ao->j,rstart=A->rmap->rstart; 198d14fa243SHong Zhang PetscInt *lnk,i,pnz,row,*api,*apj,*Jptr,apnz,nspacedouble=0,j,nzi; 1995cab4f04SHong Zhang PetscInt am=A->rmap->n,pN=P->cmap->N,pn=P->cmap->n,pm=P->rmap->n; 200598bc09dSHong Zhang PetscBT lnkbt; 201bfcd1a73SHong Zhang PetscReal afill; 202b92f168fSBarry Smith MatType mtype; 203598bc09dSHong Zhang 204598bc09dSHong Zhang PetscFunctionBegin; 2056718818eSStefano Zampini MatCheckProduct(C,4); 20628b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Extra product struct not empty"); 2079566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 2089566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 2099467f45fSHong Zhang 2103cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 2119566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 212598bc09dSHong Zhang 213598bc09dSHong Zhang /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 2149566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 21519f0e57aSHong Zhang 216598bc09dSHong Zhang /* get P_loc by taking all local rows of P */ 2179566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc)); 218598bc09dSHong Zhang 219a1a4e84aSHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 220598bc09dSHong Zhang pi_loc = p_loc->i; pj_loc = p_loc->j; 2219467f45fSHong Zhang if (size > 1) { 2229467f45fSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 223598bc09dSHong Zhang pi_oth = p_oth->i; pj_oth = p_oth->j; 22420e3dc75SHong Zhang } else { 22520e3dc75SHong Zhang p_oth = NULL; 22620e3dc75SHong Zhang pi_oth = NULL; pj_oth = NULL; 2279467f45fSHong Zhang } 228598bc09dSHong Zhang 229598bc09dSHong Zhang /* first, compute symbolic AP = A_loc*P = A_diag*P_loc + A_off*P_oth */ 230598bc09dSHong Zhang /*-------------------------------------------------------------------*/ 2319566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&api)); 232a1a4e84aSHong Zhang ptap->api = api; 233598bc09dSHong Zhang api[0] = 0; 234598bc09dSHong Zhang 2355cab4f04SHong Zhang /* create and initialize a linked list */ 2369566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate(pN,pN,&lnk,&lnkbt)); 237598bc09dSHong Zhang 238bfcd1a73SHong Zhang /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */ 2399566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(adi[am],PetscIntSumTruncate(aoi[am],pi_loc[pm]))),&free_space)); 240598bc09dSHong Zhang current_space = free_space; 241598bc09dSHong Zhang 242d0609cedSBarry Smith MatPreallocateBegin(comm,am,pn,dnz,onz); 243598bc09dSHong Zhang for (i=0; i<am; i++) { 244598bc09dSHong Zhang /* diagonal portion of A */ 245598bc09dSHong Zhang nzi = adi[i+1] - adi[i]; 246598bc09dSHong Zhang for (j=0; j<nzi; j++) { 247598bc09dSHong Zhang row = *adj++; 248598bc09dSHong Zhang pnz = pi_loc[row+1] - pi_loc[row]; 249598bc09dSHong Zhang Jptr = pj_loc + pi_loc[row]; 250598bc09dSHong Zhang /* add non-zero cols of P into the sorted linked list lnk */ 2519566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(pnz,Jptr,lnk,lnkbt)); 252598bc09dSHong Zhang } 253598bc09dSHong Zhang /* off-diagonal portion of A */ 254598bc09dSHong Zhang nzi = aoi[i+1] - aoi[i]; 255598bc09dSHong Zhang for (j=0; j<nzi; j++) { 256598bc09dSHong Zhang row = *aoj++; 257598bc09dSHong Zhang pnz = pi_oth[row+1] - pi_oth[row]; 258598bc09dSHong Zhang Jptr = pj_oth + pi_oth[row]; 2599566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(pnz,Jptr,lnk,lnkbt)); 260598bc09dSHong Zhang } 2618c78258cSHong Zhang /* add possible missing diagonal entry */ 2628c78258cSHong Zhang if (C->force_diagonals) { 2638c78258cSHong Zhang j = i + rstart; /* column index */ 2649566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(1,&j,lnk,lnkbt)); 2658c78258cSHong Zhang } 266598bc09dSHong Zhang 267d14fa243SHong Zhang apnz = lnk[0]; 268598bc09dSHong Zhang api[i+1] = api[i] + apnz; 269598bc09dSHong Zhang 270598bc09dSHong Zhang /* if free space is not available, double the total space in the list */ 271598bc09dSHong Zhang if (current_space->local_remaining<apnz) { 2729566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(apnz,current_space->total_array_size),¤t_space)); 273598bc09dSHong Zhang nspacedouble++; 274598bc09dSHong Zhang } 275598bc09dSHong Zhang 276598bc09dSHong Zhang /* Copy data into free space, then initialize lnk */ 2779566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean(pN,apnz,current_space->array,lnk,lnkbt)); 2789566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+rstart,apnz,current_space->array,dnz,onz)); 2792205254eSKarl Rupp 280598bc09dSHong Zhang current_space->array += apnz; 281598bc09dSHong Zhang current_space->local_used += apnz; 282598bc09dSHong Zhang current_space->local_remaining -= apnz; 283598bc09dSHong Zhang } 284598bc09dSHong Zhang 285598bc09dSHong Zhang /* Allocate space for apj, initialize apj, and */ 286598bc09dSHong Zhang /* destroy list of free space and other temporary array(s) */ 2879566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(api[am]+1,&ptap->apj)); 288a1a4e84aSHong Zhang apj = ptap->apj; 2899566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,ptap->apj)); 2909566063dSJacob Faibussowitsch PetscCall(PetscLLDestroy(lnk,lnkbt)); 291598bc09dSHong Zhang 292f84c536bSHong Zhang /* malloc apa to store dense row A[i,:]*P */ 2939566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(pN,&ptap->apa)); 294f84c536bSHong Zhang 2954222ddf1SHong Zhang /* set and assemble symbolic parallel matrix C */ 2964222ddf1SHong Zhang /*---------------------------------------------*/ 2979566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,am,pn,PETSC_DETERMINE,PETSC_DETERMINE)); 2989566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,P)); 299a2f3521dSMark F. Adams 3009566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 3019566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 3029566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 303d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 30417f75172Sandi selinger 3059566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api)); 3069566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 3079566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 3089566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 309598bc09dSHong Zhang 3104222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable; 3114222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AB; 312598bc09dSHong Zhang 3134222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 3146718818eSStefano Zampini C->product->data = ptap; 3156718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_MatMatMult; 316598bc09dSHong Zhang 317bfcd1a73SHong Zhang /* set MatInfo */ 318118727c9SMark F. Adams afill = (PetscReal)api[am]/(adi[am]+aoi[am]+pi_loc[pm]+1) + 1.e-5; 319bfcd1a73SHong Zhang if (afill < 1.0) afill = 1.0; 3204222ddf1SHong Zhang C->info.mallocs = nspacedouble; 3214222ddf1SHong Zhang C->info.fill_ratio_given = fill; 3224222ddf1SHong Zhang C->info.fill_ratio_needed = afill; 323bfcd1a73SHong Zhang 324bfcd1a73SHong Zhang #if defined(PETSC_USE_INFO) 325bfcd1a73SHong Zhang if (api[am]) { 3269566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 3279566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n",(double)afill)); 328bfcd1a73SHong Zhang } else { 3299566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 330bfcd1a73SHong Zhang } 331bfcd1a73SHong Zhang #endif 332598bc09dSHong Zhang PetscFunctionReturn(0); 333598bc09dSHong Zhang } 334598bc09dSHong Zhang 3354222ddf1SHong Zhang /* ------------------------------------------------------- */ 3366718818eSStefano Zampini static PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIDense(Mat,Mat,PetscReal,Mat); 3376718818eSStefano Zampini static PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIDense(Mat,Mat,Mat); 3386718818eSStefano Zampini 3394222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense_AB(Mat C) 3409123193aSHong Zhang { 3414222ddf1SHong Zhang Mat_Product *product = C->product; 3424222ddf1SHong Zhang Mat A = product->A,B=product->B; 3439123193aSHong Zhang 3449123193aSHong Zhang PetscFunctionBegin; 3454222ddf1SHong Zhang if (A->cmap->rstart != B->rmap->rstart || A->cmap->rend != B->rmap->rend) 34698921bdaSJacob 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); 3474222ddf1SHong Zhang 3484222ddf1SHong Zhang C->ops->matmultsymbolic = MatMatMultSymbolic_MPIAIJ_MPIDense; 3494222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AB; 3509123193aSHong Zhang PetscFunctionReturn(0); 3519123193aSHong Zhang } 3524222ddf1SHong Zhang /* -------------------------------------------------------------------- */ 3534222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense_AtB(Mat C) 3544222ddf1SHong Zhang { 3554222ddf1SHong Zhang Mat_Product *product = C->product; 3564222ddf1SHong Zhang Mat A = product->A,B=product->B; 3574222ddf1SHong Zhang 3584222ddf1SHong Zhang PetscFunctionBegin; 3594222ddf1SHong Zhang if (A->rmap->rstart != B->rmap->rstart || A->rmap->rend != B->rmap->rend) 36098921bdaSJacob 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); 3614222ddf1SHong Zhang 3624222ddf1SHong Zhang C->ops->transposematmultsymbolic = MatTransposeMatMultSymbolic_MPIAIJ_MPIDense; 3634222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AtB; 3644222ddf1SHong Zhang PetscFunctionReturn(0); 3654222ddf1SHong Zhang } 3664222ddf1SHong Zhang 3674222ddf1SHong Zhang /* --------------------------------------------------------------------- */ 3684222ddf1SHong Zhang PETSC_INTERN PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense(Mat C) 3694222ddf1SHong Zhang { 3704222ddf1SHong Zhang Mat_Product *product = C->product; 3714222ddf1SHong Zhang 3724222ddf1SHong Zhang PetscFunctionBegin; 3734222ddf1SHong Zhang switch (product->type) { 3744222ddf1SHong Zhang case MATPRODUCT_AB: 3759566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_MPIDense_AB(C)); 3764222ddf1SHong Zhang break; 3774222ddf1SHong Zhang case MATPRODUCT_AtB: 3789566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_MPIDense_AtB(C)); 3794222ddf1SHong Zhang break; 3804222ddf1SHong Zhang default: 3816718818eSStefano Zampini break; 3824222ddf1SHong Zhang } 3834222ddf1SHong Zhang PetscFunctionReturn(0); 3844222ddf1SHong Zhang } 3854222ddf1SHong Zhang /* ------------------------------------------------------- */ 3869123193aSHong Zhang 3874324174eSBarry Smith typedef struct { 3885ea7661aSPierre Jolivet Mat workB,workB1; 3894324174eSBarry Smith MPI_Request *rwaits,*swaits; 39087753ddeSHong Zhang PetscInt nsends,nrecvs; 39187753ddeSHong Zhang MPI_Datatype *stype,*rtype; 3926718818eSStefano Zampini PetscInt blda; 3934324174eSBarry Smith } MPIAIJ_MPIDense; 3944324174eSBarry Smith 39596b95a6bSBarry Smith PetscErrorCode MatMPIAIJ_MPIDenseDestroy(void *ctx) 3964324174eSBarry Smith { 3974324174eSBarry Smith MPIAIJ_MPIDense *contents = (MPIAIJ_MPIDense*)ctx; 39887753ddeSHong Zhang PetscInt i; 3994324174eSBarry Smith 4004324174eSBarry Smith PetscFunctionBegin; 4019566063dSJacob Faibussowitsch PetscCall(MatDestroy(&contents->workB)); 4029566063dSJacob Faibussowitsch PetscCall(MatDestroy(&contents->workB1)); 40387753ddeSHong Zhang for (i=0; i<contents->nsends; i++) { 4049566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&contents->stype[i])); 40587753ddeSHong Zhang } 40687753ddeSHong Zhang for (i=0; i<contents->nrecvs; i++) { 4079566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&contents->rtype[i])); 40887753ddeSHong Zhang } 4099566063dSJacob Faibussowitsch PetscCall(PetscFree4(contents->stype,contents->rtype,contents->rwaits,contents->swaits)); 4109566063dSJacob Faibussowitsch PetscCall(PetscFree(contents)); 4114324174eSBarry Smith PetscFunctionReturn(0); 4124324174eSBarry Smith } 4134324174eSBarry Smith 4146718818eSStefano Zampini static PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIDense(Mat A,Mat B,PetscReal fill,Mat C) 4159123193aSHong Zhang { 416f32d5d43SBarry Smith Mat_MPIAIJ *aij=(Mat_MPIAIJ*)A->data; 417b45e3bf4SStefano Zampini PetscInt nz=aij->B->cmap->n,nsends,nrecvs,i,nrows_to,j,blda,m,M,n,N; 4184324174eSBarry Smith MPIAIJ_MPIDense *contents; 4194324174eSBarry Smith VecScatter ctx=aij->Mvctx; 4205ea7661aSPierre Jolivet PetscInt Am=A->rmap->n,Bm=B->rmap->n,BN=B->cmap->N,Bbn,Bbn1,bs,nrows_from,numBb; 42187753ddeSHong Zhang MPI_Comm comm; 42287753ddeSHong Zhang MPI_Datatype type1,*stype,*rtype; 42387753ddeSHong Zhang const PetscInt *sindices,*sstarts,*rstarts; 42487753ddeSHong Zhang PetscMPIInt *disp; 4256718818eSStefano Zampini PetscBool cisdense; 4269123193aSHong Zhang 4279123193aSHong Zhang PetscFunctionBegin; 4286718818eSStefano Zampini MatCheckProduct(C,4); 42928b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data not empty"); 4309566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 4319566063dSJacob Faibussowitsch PetscCall(PetscObjectBaseTypeCompare((PetscObject)C,MATMPIDENSE,&cisdense)); 4326718818eSStefano Zampini if (!cisdense) { 4339566063dSJacob Faibussowitsch PetscCall(MatSetType(C,((PetscObject)B)->type_name)); 4346718818eSStefano Zampini } 4359566063dSJacob Faibussowitsch PetscCall(MatGetLocalSize(C,&m,&n)); 4369566063dSJacob Faibussowitsch PetscCall(MatGetSize(C,&M,&N)); 437825bfe5fSPierre Jolivet if (m == PETSC_DECIDE || n == PETSC_DECIDE || M == PETSC_DECIDE || N == PETSC_DECIDE) { 4389566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,Am,B->cmap->n,A->rmap->N,BN)); 439825bfe5fSPierre Jolivet } 4409566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,B)); 4419566063dSJacob Faibussowitsch PetscCall(MatSetUp(C)); 4429566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(B,&blda)); 4439566063dSJacob Faibussowitsch PetscCall(PetscNew(&contents)); 4444324174eSBarry Smith 4459566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,NULL,NULL)); 4469566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,NULL,NULL)); 44787753ddeSHong Zhang 44887753ddeSHong Zhang /* Create column block of B and C for memory scalability when BN is too large */ 44987753ddeSHong Zhang /* Estimate Bbn, column size of Bb */ 45087753ddeSHong Zhang if (nz) { 45187753ddeSHong Zhang Bbn1 = 2*Am*BN/nz; 452da7cbbe7SHong Zhang if (!Bbn1) Bbn1 = 1; 45387753ddeSHong Zhang } else Bbn1 = BN; 45487753ddeSHong Zhang 45587753ddeSHong Zhang bs = PetscAbs(B->cmap->bs); 45687753ddeSHong Zhang Bbn1 = Bbn1/bs *bs; /* Bbn1 is a multiple of bs */ 45787753ddeSHong Zhang if (Bbn1 > BN) Bbn1 = BN; 4589566063dSJacob Faibussowitsch PetscCallMPI(MPI_Allreduce(&Bbn1,&Bbn,1,MPIU_INT,MPI_MAX,comm)); 45987753ddeSHong Zhang 46087753ddeSHong Zhang /* Enable runtime option for Bbn */ 461d0609cedSBarry Smith PetscOptionsBegin(comm,((PetscObject)C)->prefix,"MatMatMult","Mat"); 4629566063dSJacob Faibussowitsch PetscCall(PetscOptionsInt("-matmatmult_Bbn","Number of columns in Bb","MatMatMult",Bbn,&Bbn,NULL)); 463d0609cedSBarry Smith PetscOptionsEnd(); 4646718818eSStefano Zampini Bbn = PetscMin(Bbn,BN); 46587753ddeSHong Zhang 4666718818eSStefano Zampini if (Bbn > 0 && Bbn < BN) { 4675ea7661aSPierre Jolivet numBb = BN/Bbn; 4685ea7661aSPierre Jolivet Bbn1 = BN - numBb*Bbn; 4695ea7661aSPierre Jolivet } else numBb = 0; 47087753ddeSHong Zhang 4715ea7661aSPierre Jolivet if (numBb) { 4729566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"use Bb, BN=%" PetscInt_FMT ", Bbn=%" PetscInt_FMT "; numBb=%" PetscInt_FMT "\n",BN,Bbn,numBb)); 4735ea7661aSPierre Jolivet if (Bbn1) { /* Create workB1 for the remaining columns */ 4749566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"use Bb1, BN=%" PetscInt_FMT ", Bbn1=%" PetscInt_FMT "\n",BN,Bbn1)); 47587753ddeSHong Zhang /* Create work matrix used to store off processor rows of B needed for local product */ 4769566063dSJacob Faibussowitsch PetscCall(MatCreateSeqDense(PETSC_COMM_SELF,nz,Bbn1,NULL,&contents->workB1)); 4775ea7661aSPierre Jolivet } else contents->workB1 = NULL; 47887753ddeSHong Zhang } 47987753ddeSHong Zhang 48087753ddeSHong Zhang /* Create work matrix used to store off processor rows of B needed for local product */ 4819566063dSJacob Faibussowitsch PetscCall(MatCreateSeqDense(PETSC_COMM_SELF,nz,Bbn,NULL,&contents->workB)); 48287753ddeSHong Zhang 48387753ddeSHong Zhang /* Use MPI derived data type to reduce memory required by the send/recv buffers */ 4849566063dSJacob Faibussowitsch PetscCall(PetscMalloc4(nsends,&stype,nrecvs,&rtype,nrecvs,&contents->rwaits,nsends,&contents->swaits)); 48587753ddeSHong Zhang contents->stype = stype; 48687753ddeSHong Zhang contents->nsends = nsends; 48787753ddeSHong Zhang 48887753ddeSHong Zhang contents->rtype = rtype; 48987753ddeSHong Zhang contents->nrecvs = nrecvs; 4906718818eSStefano Zampini contents->blda = blda; 49187753ddeSHong Zhang 4929566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(Bm+1,&disp)); 49387753ddeSHong Zhang for (i=0; i<nsends; i++) { 49487753ddeSHong Zhang nrows_to = sstarts[i+1]-sstarts[i]; 495b45e3bf4SStefano Zampini for (j=0; j<nrows_to; j++) disp[j] = sindices[sstarts[i]+j]; /* rowB to be sent */ 4969566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_indexed_block(nrows_to,1,disp,MPIU_SCALAR,&type1)); 4979566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_resized(type1,0,blda*sizeof(PetscScalar),&stype[i])); 4989566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_commit(&stype[i])); 4999566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&type1)); 50087753ddeSHong Zhang } 50187753ddeSHong Zhang 50287753ddeSHong Zhang for (i=0; i<nrecvs; i++) { 50387753ddeSHong Zhang /* received values from a process form a (nrows_from x Bbn) row block in workB (column-wise) */ 50487753ddeSHong Zhang nrows_from = rstarts[i+1]-rstarts[i]; 50587753ddeSHong Zhang disp[0] = 0; 5069566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_indexed_block(1,nrows_from,disp,MPIU_SCALAR,&type1)); 5079566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_resized(type1,0,nz*sizeof(PetscScalar),&rtype[i])); 5089566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_commit(&rtype[i])); 5099566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&type1)); 51087753ddeSHong Zhang } 51187753ddeSHong Zhang 5129566063dSJacob Faibussowitsch PetscCall(PetscFree(disp)); 5139566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,NULL,NULL)); 5149566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,NULL,NULL)); 5159566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NO_OFF_PROC_ENTRIES,PETSC_TRUE)); 5169566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 5179566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 5189566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 5196718818eSStefano Zampini 5206718818eSStefano Zampini C->product->data = contents; 5216718818eSStefano Zampini C->product->destroy = MatMPIAIJ_MPIDenseDestroy; 5224222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIDense; 5239123193aSHong Zhang PetscFunctionReturn(0); 5249123193aSHong Zhang } 5259123193aSHong Zhang 52693aa15f2SStefano Zampini PETSC_INTERN PetscErrorCode MatMatMultNumericAdd_SeqAIJ_SeqDense(Mat,Mat,Mat,const PetscBool); 527f32d5d43SBarry Smith /* 5282636ff26SBarry Smith Performs an efficient scatter on the rows of B needed by this process; this is 5292636ff26SBarry Smith a modification of the VecScatterBegin_() routines. 53087753ddeSHong Zhang 53187753ddeSHong Zhang Input: Bbidx = 0: B = Bb 53287753ddeSHong Zhang = 1: B = Bb1, see MatMatMultSymbolic_MPIAIJ_MPIDense() 533f32d5d43SBarry Smith */ 53487753ddeSHong Zhang PetscErrorCode MatMPIDenseScatter(Mat A,Mat B,PetscInt Bbidx,Mat C,Mat *outworkB) 535f32d5d43SBarry Smith { 536f32d5d43SBarry Smith Mat_MPIAIJ *aij = (Mat_MPIAIJ*)A->data; 5371683a169SBarry Smith const PetscScalar *b; 53887753ddeSHong Zhang PetscScalar *rvalues; 539f32d5d43SBarry Smith VecScatter ctx = aij->Mvctx; 54087753ddeSHong Zhang const PetscInt *sindices,*sstarts,*rstarts; 5413515ee7fSJunchao Zhang const PetscMPIInt *sprocs,*rprocs; 5426718818eSStefano Zampini PetscInt i,nsends,nrecvs; 543f32d5d43SBarry Smith MPI_Request *swaits,*rwaits; 544ce94432eSBarry Smith MPI_Comm comm; 5456718818eSStefano Zampini PetscMPIInt tag=((PetscObject)ctx)->tag,ncols=B->cmap->N,nrows=aij->B->cmap->n,nsends_mpi,nrecvs_mpi; 5464324174eSBarry Smith MPIAIJ_MPIDense *contents; 5474324174eSBarry Smith Mat workB; 54887753ddeSHong Zhang MPI_Datatype *stype,*rtype; 5496718818eSStefano Zampini PetscInt blda; 550f32d5d43SBarry Smith 551f32d5d43SBarry Smith PetscFunctionBegin; 5526718818eSStefano Zampini MatCheckProduct(C,4); 55328b400f6SJacob Faibussowitsch PetscCheck(C->product->data,PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data empty"); 5546718818eSStefano Zampini contents = (MPIAIJ_MPIDense*)C->product->data; 5559566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,&sprocs,NULL/*bs*/)); 5569566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,&rprocs,NULL/*bs*/)); 5579566063dSJacob Faibussowitsch PetscCall(PetscMPIIntCast(nsends,&nsends_mpi)); 5589566063dSJacob Faibussowitsch PetscCall(PetscMPIIntCast(nrecvs,&nrecvs_mpi)); 559b45e3bf4SStefano Zampini if (Bbidx == 0) workB = *outworkB = contents->workB; 560b45e3bf4SStefano Zampini else workB = *outworkB = contents->workB1; 56108401ef6SPierre 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); 5623515ee7fSJunchao Zhang swaits = contents->swaits; 5633515ee7fSJunchao Zhang rwaits = contents->rwaits; 564f32d5d43SBarry Smith 5659566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(B,&b)); 5669566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(B,&blda)); 56708401ef6SPierre 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); 5689566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(workB,&rvalues)); 569f32d5d43SBarry Smith 57087753ddeSHong Zhang /* Post recv, use MPI derived data type to save memory */ 5719566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 57287753ddeSHong Zhang rtype = contents->rtype; 5733515ee7fSJunchao Zhang for (i=0; i<nrecvs; i++) { 5749566063dSJacob Faibussowitsch PetscCallMPI(MPI_Irecv(rvalues+(rstarts[i]-rstarts[0]),ncols,rtype[i],rprocs[i],tag,comm,rwaits+i)); 575f32d5d43SBarry Smith } 576f32d5d43SBarry Smith 57787753ddeSHong Zhang stype = contents->stype; 5783515ee7fSJunchao Zhang for (i=0; i<nsends; i++) { 5799566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(b,ncols,stype[i],sprocs[i],tag,comm,swaits+i)); 580f32d5d43SBarry Smith } 5812636ff26SBarry Smith 5829566063dSJacob Faibussowitsch if (nrecvs) PetscCallMPI(MPI_Waitall(nrecvs_mpi,rwaits,MPI_STATUSES_IGNORE)); 5839566063dSJacob Faibussowitsch if (nsends) PetscCallMPI(MPI_Waitall(nsends_mpi,swaits,MPI_STATUSES_IGNORE)); 584f32d5d43SBarry Smith 5859566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,&sprocs,NULL)); 5869566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,&rprocs,NULL)); 5879566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(B,&b)); 5889566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(workB,&rvalues)); 58987753ddeSHong Zhang PetscFunctionReturn(0); 59087753ddeSHong Zhang } 591f32d5d43SBarry Smith 5926718818eSStefano Zampini static PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIDense(Mat A,Mat B,Mat C) 5939123193aSHong Zhang { 594f32d5d43SBarry Smith Mat_MPIAIJ *aij = (Mat_MPIAIJ*)A->data; 595f32d5d43SBarry Smith Mat_MPIDense *bdense = (Mat_MPIDense*)B->data; 596f32d5d43SBarry Smith Mat_MPIDense *cdense = (Mat_MPIDense*)C->data; 597f32d5d43SBarry Smith Mat workB; 59887753ddeSHong Zhang MPIAIJ_MPIDense *contents; 5999123193aSHong Zhang 6009123193aSHong Zhang PetscFunctionBegin; 6016718818eSStefano Zampini MatCheckProduct(C,3); 60228b400f6SJacob Faibussowitsch PetscCheck(C->product->data,PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data empty"); 6036718818eSStefano Zampini contents = (MPIAIJ_MPIDense*)C->product->data; 604f32d5d43SBarry Smith /* diagonal block of A times all local rows of B */ 6055ea7661aSPierre Jolivet /* TODO: this calls a symbolic multiplication every time, which could be avoided */ 6069566063dSJacob Faibussowitsch PetscCall(MatMatMult(aij->A,bdense->A,MAT_REUSE_MATRIX,PETSC_DEFAULT,&cdense->A)); 6075ea7661aSPierre Jolivet if (contents->workB->cmap->n == B->cmap->N) { 60887753ddeSHong Zhang /* get off processor parts of B needed to complete C=A*B */ 6099566063dSJacob Faibussowitsch PetscCall(MatMPIDenseScatter(A,B,0,C,&workB)); 610f32d5d43SBarry Smith 611f32d5d43SBarry Smith /* off-diagonal block of A times nonlocal rows of B */ 6129566063dSJacob Faibussowitsch PetscCall(MatMatMultNumericAdd_SeqAIJ_SeqDense(aij->B,workB,cdense->A,PETSC_TRUE)); 61387753ddeSHong Zhang } else { 6145ea7661aSPierre Jolivet Mat Bb,Cb; 6155ea7661aSPierre Jolivet PetscInt BN=B->cmap->N,n=contents->workB->cmap->n,i; 616664c1e85SStefano Zampini PetscBool ccpu; 61787753ddeSHong Zhang 61808401ef6SPierre Jolivet PetscCheck(n > 0,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"Column block size %" PetscInt_FMT " must be positive",n); 619664c1e85SStefano Zampini /* Prevent from unneeded copies back and forth from the GPU 620664c1e85SStefano Zampini when getting and restoring the submatrix 621664c1e85SStefano Zampini We need a proper GPU code for AIJ * dense in parallel */ 6229566063dSJacob Faibussowitsch PetscCall(MatBoundToCPU(C,&ccpu)); 6239566063dSJacob Faibussowitsch PetscCall(MatBindToCPU(C,PETSC_TRUE)); 6245ea7661aSPierre Jolivet for (i=0; i<BN; i+=n) { 625a2748737SPierre Jolivet PetscCall(MatDenseGetSubMatrix(B,PETSC_DECIDE,PETSC_DECIDE,i,PetscMin(i+n,BN),&Bb)); 626a2748737SPierre Jolivet PetscCall(MatDenseGetSubMatrix(C,PETSC_DECIDE,PETSC_DECIDE,i,PetscMin(i+n,BN),&Cb)); 62787753ddeSHong Zhang 6285ea7661aSPierre Jolivet /* get off processor parts of B needed to complete C=A*B */ 6299566063dSJacob Faibussowitsch PetscCall(MatMPIDenseScatter(A,Bb,(i+n)>BN,C,&workB)); 6305ea7661aSPierre Jolivet 6315ea7661aSPierre Jolivet /* off-diagonal block of A times nonlocal rows of B */ 6325ea7661aSPierre Jolivet cdense = (Mat_MPIDense*)Cb->data; 6339566063dSJacob Faibussowitsch PetscCall(MatMatMultNumericAdd_SeqAIJ_SeqDense(aij->B,workB,cdense->A,PETSC_TRUE)); 6349566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreSubMatrix(B,&Bb)); 6359566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreSubMatrix(C,&Cb)); 63687753ddeSHong Zhang } 6379566063dSJacob Faibussowitsch PetscCall(MatBindToCPU(C,ccpu)); 63887753ddeSHong Zhang } 6399123193aSHong Zhang PetscFunctionReturn(0); 6409123193aSHong Zhang } 641cf1a0672SHong Zhang 642f8efd71cSHong Zhang PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIAIJ(Mat A,Mat P,Mat C) 6431634b032SHong Zhang { 644cf1a0672SHong Zhang Mat_MPIAIJ *a = (Mat_MPIAIJ*)A->data,*c=(Mat_MPIAIJ*)C->data; 645cf1a0672SHong Zhang Mat_SeqAIJ *ad = (Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data; 646cf1a0672SHong Zhang Mat_SeqAIJ *cd = (Mat_SeqAIJ*)(c->A)->data,*co=(Mat_SeqAIJ*)(c->B)->data; 647cf1a0672SHong Zhang PetscInt *adi = ad->i,*adj,*aoi=ao->i,*aoj; 648cf1a0672SHong Zhang PetscScalar *ada,*aoa,*cda=cd->a,*coa=co->a; 649cf1a0672SHong Zhang Mat_SeqAIJ *p_loc,*p_oth; 650cf1a0672SHong Zhang PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*pj; 65129825010SHong Zhang PetscScalar *pa_loc,*pa_oth,*pa,valtmp,*ca; 65229825010SHong Zhang PetscInt cm = C->rmap->n,anz,pnz; 6536718818eSStefano Zampini Mat_APMPI *ptap; 65480da8df7SHong Zhang PetscScalar *apa_sparse; 655ce496241SStefano Zampini const PetscScalar *dummy; 656cf1a0672SHong Zhang PetscInt *api,*apj,*apJ,i,j,k,row; 65729825010SHong Zhang PetscInt cstart = C->cmap->rstart; 65829825010SHong Zhang PetscInt cdnz,conz,k0,k1,nextp; 659a7c7454dSHong Zhang MPI_Comm comm; 660a7c7454dSHong Zhang PetscMPIInt size; 6611634b032SHong Zhang 6621634b032SHong Zhang PetscFunctionBegin; 6636718818eSStefano Zampini MatCheckProduct(C,3); 6646718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 66528b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 6669566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 6679566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 668aed4548fSBarry Smith PetscCheck(ptap->P_oth || size <= 1,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"AP cannot be reused. Do not call MatProductClear()"); 669a7c7454dSHong Zhang 67065e4b4d4SStefano Zampini /* flag CPU mask for C */ 67165e4b4d4SStefano Zampini #if defined(PETSC_HAVE_DEVICE) 67265e4b4d4SStefano Zampini if (C->offloadmask != PETSC_OFFLOAD_UNALLOCATED) C->offloadmask = PETSC_OFFLOAD_CPU; 67365e4b4d4SStefano Zampini if (c->A->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->A->offloadmask = PETSC_OFFLOAD_CPU; 67465e4b4d4SStefano Zampini if (c->B->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->B->offloadmask = PETSC_OFFLOAD_CPU; 67565e4b4d4SStefano Zampini #endif 67680da8df7SHong Zhang apa_sparse = ptap->apa; 67780da8df7SHong Zhang 678cf1a0672SHong Zhang /* 1) get P_oth = ptap->P_oth and P_loc = ptap->P_loc */ 679cf1a0672SHong Zhang /*-----------------------------------------------------*/ 680cf1a0672SHong Zhang /* update numerical values of P_oth and P_loc */ 6819566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_REUSE_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 6829566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_REUSE_MATRIX,&ptap->P_loc)); 683cf1a0672SHong Zhang 684cf1a0672SHong Zhang /* 2) compute numeric C_loc = A_loc*P = Ad*P_loc + Ao*P_oth */ 685cf1a0672SHong Zhang /*----------------------------------------------------------*/ 686cf1a0672SHong Zhang /* get data from symbolic products */ 687cf1a0672SHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 688cf1a0672SHong Zhang pi_loc = p_loc->i; pj_loc = p_loc->j; pa_loc = p_loc->a; 689a7c7454dSHong Zhang if (size >1) { 690a7c7454dSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 691cf1a0672SHong Zhang pi_oth = p_oth->i; pj_oth = p_oth->j; pa_oth = p_oth->a; 69220e3dc75SHong Zhang } else { 69320e3dc75SHong Zhang p_oth = NULL; pi_oth = NULL; pj_oth = NULL; pa_oth = NULL; 694a7c7454dSHong Zhang } 695cf1a0672SHong Zhang 696ce496241SStefano Zampini /* trigger copy to CPU */ 6979566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->A,&dummy)); 6989566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->A,&dummy)); 6999566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->B,&dummy)); 7009566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->B,&dummy)); 70129825010SHong Zhang api = ptap->api; 70229825010SHong Zhang apj = ptap->apj; 703cf1a0672SHong Zhang for (i=0; i<cm; i++) { 70429825010SHong Zhang apJ = apj + api[i]; 70529825010SHong Zhang 706cf1a0672SHong Zhang /* diagonal portion of A */ 707cf1a0672SHong Zhang anz = adi[i+1] - adi[i]; 708cf1a0672SHong Zhang adj = ad->j + adi[i]; 709cf1a0672SHong Zhang ada = ad->a + adi[i]; 710cf1a0672SHong Zhang for (j=0; j<anz; j++) { 711cf1a0672SHong Zhang row = adj[j]; 712cf1a0672SHong Zhang pnz = pi_loc[row+1] - pi_loc[row]; 713cf1a0672SHong Zhang pj = pj_loc + pi_loc[row]; 714cf1a0672SHong Zhang pa = pa_loc + pi_loc[row]; 71529825010SHong Zhang /* perform sparse axpy */ 716cf1a0672SHong Zhang valtmp = ada[j]; 71729825010SHong Zhang nextp = 0; 71829825010SHong Zhang for (k=0; nextp<pnz; k++) { 71929825010SHong Zhang if (apJ[k] == pj[nextp]) { /* column of AP == column of P */ 72029825010SHong Zhang apa_sparse[k] += valtmp*pa[nextp++]; 72129825010SHong Zhang } 7221634b032SHong Zhang } 7239566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*pnz)); 724cf1a0672SHong Zhang } 7251634b032SHong Zhang 726cf1a0672SHong Zhang /* off-diagonal portion of A */ 727cf1a0672SHong Zhang anz = aoi[i+1] - aoi[i]; 728cf1a0672SHong Zhang aoj = ao->j + aoi[i]; 729cf1a0672SHong Zhang aoa = ao->a + aoi[i]; 730cf1a0672SHong Zhang for (j=0; j<anz; j++) { 731cf1a0672SHong Zhang row = aoj[j]; 732cf1a0672SHong Zhang pnz = pi_oth[row+1] - pi_oth[row]; 733cf1a0672SHong Zhang pj = pj_oth + pi_oth[row]; 734cf1a0672SHong Zhang pa = pa_oth + pi_oth[row]; 73529825010SHong Zhang /* perform sparse axpy */ 736cf1a0672SHong Zhang valtmp = aoa[j]; 73729825010SHong Zhang nextp = 0; 73829825010SHong Zhang for (k=0; nextp<pnz; k++) { 73929825010SHong Zhang if (apJ[k] == pj[nextp]) { /* column of AP == column of P */ 74029825010SHong Zhang apa_sparse[k] += valtmp*pa[nextp++]; 74129825010SHong Zhang } 742cf1a0672SHong Zhang } 7439566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*pnz)); 744cf1a0672SHong Zhang } 7451634b032SHong Zhang 746cf1a0672SHong Zhang /* set values in C */ 747cf1a0672SHong Zhang cdnz = cd->i[i+1] - cd->i[i]; 748cf1a0672SHong Zhang conz = co->i[i+1] - co->i[i]; 7491634b032SHong Zhang 750910cf402Sprj- /* 1st off-diagonal part of C */ 751cf1a0672SHong Zhang ca = coa + co->i[i]; 752cf1a0672SHong Zhang k = 0; 753cf1a0672SHong Zhang for (k0=0; k0<conz; k0++) { 754cf1a0672SHong Zhang if (apJ[k] >= cstart) break; 75529825010SHong Zhang ca[k0] = apa_sparse[k]; 75629825010SHong Zhang apa_sparse[k] = 0.0; 757cf1a0672SHong Zhang k++; 758cf1a0672SHong Zhang } 7591634b032SHong Zhang 760cf1a0672SHong Zhang /* diagonal part of C */ 761cf1a0672SHong Zhang ca = cda + cd->i[i]; 762cf1a0672SHong Zhang for (k1=0; k1<cdnz; k1++) { 76329825010SHong Zhang ca[k1] = apa_sparse[k]; 76429825010SHong Zhang apa_sparse[k] = 0.0; 765cf1a0672SHong Zhang k++; 766cf1a0672SHong Zhang } 767cf1a0672SHong Zhang 768910cf402Sprj- /* 2nd off-diagonal part of C */ 769cf1a0672SHong Zhang ca = coa + co->i[i]; 770cf1a0672SHong Zhang for (; k0<conz; k0++) { 77129825010SHong Zhang ca[k0] = apa_sparse[k]; 77229825010SHong Zhang apa_sparse[k] = 0.0; 773cf1a0672SHong Zhang k++; 774cf1a0672SHong Zhang } 775cf1a0672SHong Zhang } 7769566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 7779566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 778cf1a0672SHong Zhang PetscFunctionReturn(0); 779cf1a0672SHong Zhang } 780cf1a0672SHong Zhang 7810fc8cf34SHong Zhang /* same as MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(), except using LLCondensed to avoid O(BN) memory requirement */ 7824222ddf1SHong Zhang PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ(Mat A,Mat P,PetscReal fill,Mat C) 783cf1a0672SHong Zhang { 784ce94432eSBarry Smith MPI_Comm comm; 785a7c7454dSHong Zhang PetscMPIInt size; 7863cdca5ebSHong Zhang Mat_APMPI *ptap; 7870298fd71SBarry Smith PetscFreeSpaceList free_space = NULL,current_space=NULL; 7886718818eSStefano Zampini Mat_MPIAIJ *a = (Mat_MPIAIJ*)A->data; 789cf1a0672SHong Zhang Mat_SeqAIJ *ad = (Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_loc,*p_oth; 7907dbaa2c6Sandi selinger PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*dnz,*onz; 791cf1a0672SHong Zhang PetscInt *adi=ad->i,*adj=ad->j,*aoi=ao->i,*aoj=ao->j,rstart=A->rmap->rstart; 7928c78258cSHong Zhang PetscInt i,pnz,row,*api,*apj,*Jptr,apnz,nspacedouble=0,j,nzi,*lnk,apnz_max=1; 793e5ab8c0aSKarl Rupp PetscInt am=A->rmap->n,pn=P->cmap->n,pm=P->rmap->n,lsize=pn+20; 794cf1a0672SHong Zhang PetscReal afill; 795b92f168fSBarry Smith MatType mtype; 796cf1a0672SHong Zhang 797cf1a0672SHong Zhang PetscFunctionBegin; 7986718818eSStefano Zampini MatCheckProduct(C,4); 79928b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Extra product struct not empty"); 8009566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 8019566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 802a7c7454dSHong Zhang 8033cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 8049566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 805cf1a0672SHong Zhang 806cf1a0672SHong Zhang /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 8079566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 80819f0e57aSHong Zhang 809cf1a0672SHong Zhang /* get P_loc by taking all local rows of P */ 8109566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc)); 811cf1a0672SHong Zhang 812cf1a0672SHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 813cf1a0672SHong Zhang pi_loc = p_loc->i; pj_loc = p_loc->j; 814a7c7454dSHong Zhang if (size > 1) { 815a7c7454dSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 81620e3dc75SHong Zhang pi_oth = p_oth->i; pj_oth = p_oth->j; 817968382fdSHong Zhang } else { 81820e3dc75SHong Zhang p_oth = NULL; 81920e3dc75SHong Zhang pi_oth = NULL; pj_oth = NULL; 820a7c7454dSHong Zhang } 821cf1a0672SHong Zhang 822cf1a0672SHong Zhang /* first, compute symbolic AP = A_loc*P = A_diag*P_loc + A_off*P_oth */ 823cf1a0672SHong Zhang /*-------------------------------------------------------------------*/ 8249566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&api)); 825cf1a0672SHong Zhang ptap->api = api; 826cf1a0672SHong Zhang api[0] = 0; 827cf1a0672SHong Zhang 8289566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(lsize,&lnk)); 829cf1a0672SHong Zhang 830cf1a0672SHong Zhang /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */ 8319566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(adi[am],PetscIntSumTruncate(aoi[am],pi_loc[pm]))),&free_space)); 832cf1a0672SHong Zhang current_space = free_space; 833d0609cedSBarry Smith MatPreallocateBegin(comm,am,pn,dnz,onz); 834cf1a0672SHong Zhang for (i=0; i<am; i++) { 835cf1a0672SHong Zhang /* diagonal portion of A */ 836cf1a0672SHong Zhang nzi = adi[i+1] - adi[i]; 837cf1a0672SHong Zhang for (j=0; j<nzi; j++) { 838cf1a0672SHong Zhang row = *adj++; 839cf1a0672SHong Zhang pnz = pi_loc[row+1] - pi_loc[row]; 840cf1a0672SHong Zhang Jptr = pj_loc + pi_loc[row]; 8418a210135Sandi selinger /* Expand list if it is not long enough */ 8428328dd82Sandi selinger if (pnz+apnz_max > lsize) { 8438328dd82Sandi selinger lsize = pnz+apnz_max; 8449566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedExpand_Scalable(lsize, &lnk)); 8458a210135Sandi selinger } 846cf1a0672SHong Zhang /* add non-zero cols of P into the sorted linked list lnk */ 8479566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(pnz,Jptr,lnk)); 8488a210135Sandi selinger apnz = *lnk; /* The first element in the list is the number of items in the list */ 8498a210135Sandi selinger api[i+1] = api[i] + apnz; 8508c78258cSHong Zhang if (apnz > apnz_max) apnz_max = apnz + 1; /* '1' for diagonal entry */ 851cf1a0672SHong Zhang } 852cf1a0672SHong Zhang /* off-diagonal portion of A */ 853cf1a0672SHong Zhang nzi = aoi[i+1] - aoi[i]; 854cf1a0672SHong Zhang for (j=0; j<nzi; j++) { 855cf1a0672SHong Zhang row = *aoj++; 856cf1a0672SHong Zhang pnz = pi_oth[row+1] - pi_oth[row]; 857cf1a0672SHong Zhang Jptr = pj_oth + pi_oth[row]; 8588a210135Sandi selinger /* Expand list if it is not long enough */ 8598328dd82Sandi selinger if (pnz+apnz_max > lsize) { 8608328dd82Sandi selinger lsize = pnz + apnz_max; 8619566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedExpand_Scalable(lsize, &lnk)); 8628a210135Sandi selinger } 8638a210135Sandi selinger /* add non-zero cols of P into the sorted linked list lnk */ 8649566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(pnz,Jptr,lnk)); 8658a210135Sandi selinger apnz = *lnk; /* The first element in the list is the number of items in the list */ 8668a210135Sandi selinger api[i+1] = api[i] + apnz; 8678c78258cSHong Zhang if (apnz > apnz_max) apnz_max = apnz + 1; /* '1' for diagonal entry */ 868cf1a0672SHong Zhang } 8698c78258cSHong Zhang 8708c78258cSHong Zhang /* add missing diagonal entry */ 8718c78258cSHong Zhang if (C->force_diagonals) { 8728c78258cSHong Zhang j = i + rstart; /* column index */ 8739566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(1,&j,lnk)); 8748c78258cSHong Zhang } 8758c78258cSHong Zhang 876f84c536bSHong Zhang apnz = *lnk; 877cf1a0672SHong Zhang api[i+1] = api[i] + apnz; 8788a210135Sandi selinger if (apnz > apnz_max) apnz_max = apnz; 879cf1a0672SHong Zhang 880cf1a0672SHong Zhang /* if free space is not available, double the total space in the list */ 881cf1a0672SHong Zhang if (current_space->local_remaining<apnz) { 8829566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(apnz,current_space->total_array_size),¤t_space)); 883cf1a0672SHong Zhang nspacedouble++; 884cf1a0672SHong Zhang } 885cf1a0672SHong Zhang 886cf1a0672SHong Zhang /* Copy data into free space, then initialize lnk */ 8879566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(apnz,current_space->array,lnk)); 8889566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+rstart,apnz,current_space->array,dnz,onz)); 8892205254eSKarl Rupp 890cf1a0672SHong Zhang current_space->array += apnz; 891cf1a0672SHong Zhang current_space->local_used += apnz; 892cf1a0672SHong Zhang current_space->local_remaining -= apnz; 893cf1a0672SHong Zhang } 894cf1a0672SHong Zhang 895cf1a0672SHong Zhang /* Allocate space for apj, initialize apj, and */ 896cf1a0672SHong Zhang /* destroy list of free space and other temporary array(s) */ 8979566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(api[am]+1,&ptap->apj)); 898cf1a0672SHong Zhang apj = ptap->apj; 8999566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,ptap->apj)); 9009566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); 901cf1a0672SHong Zhang 9024222ddf1SHong Zhang /* create and assemble symbolic parallel matrix C */ 903cf1a0672SHong Zhang /*----------------------------------------------------*/ 9049566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,am,pn,PETSC_DETERMINE,PETSC_DETERMINE)); 9059566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,P)); 9069566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 9079566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 9089566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 909d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 910cf1a0672SHong Zhang 9114222ddf1SHong Zhang /* malloc apa for assembly C */ 9129566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(apnz_max,&ptap->apa)); 91317f75172Sandi selinger 9149566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api)); 9159566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 9169566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 9179566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 918cf1a0672SHong Zhang 9194222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ; 9204222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AB; 921cf1a0672SHong Zhang 9224222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 9236718818eSStefano Zampini C->product->data = ptap; 9246718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_MatMatMult; 925cf1a0672SHong Zhang 926cf1a0672SHong Zhang /* set MatInfo */ 927118727c9SMark F. Adams afill = (PetscReal)api[am]/(adi[am]+aoi[am]+pi_loc[pm]+1) + 1.e-5; 928cf1a0672SHong Zhang if (afill < 1.0) afill = 1.0; 9294222ddf1SHong Zhang C->info.mallocs = nspacedouble; 9304222ddf1SHong Zhang C->info.fill_ratio_given = fill; 9314222ddf1SHong Zhang C->info.fill_ratio_needed = afill; 932cf1a0672SHong Zhang 933cf1a0672SHong Zhang #if defined(PETSC_USE_INFO) 934cf1a0672SHong Zhang if (api[am]) { 9359566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 9369566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n",(double)afill)); 937cf1a0672SHong Zhang } else { 9389566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 939cf1a0672SHong Zhang } 940cf1a0672SHong Zhang #endif 9411634b032SHong Zhang PetscFunctionReturn(0); 9421634b032SHong Zhang } 943f96d37fcSHong Zhang 944a9d7e0ccSandi selinger /* This function is needed for the seqMPI matrix-matrix multiplication. */ 945a9d7e0ccSandi selinger /* Three input arrays are merged to one output array. The size of the */ 946a9d7e0ccSandi selinger /* output array is also output. Duplicate entries only show up once. */ 947a9d7e0ccSandi selinger static void Merge3SortedArrays(PetscInt size1, PetscInt *in1, 948a9d7e0ccSandi selinger PetscInt size2, PetscInt *in2, 949a9d7e0ccSandi selinger PetscInt size3, PetscInt *in3, 950a9d7e0ccSandi selinger PetscInt *size4, PetscInt *out) 951a9d7e0ccSandi selinger { 952a9d7e0ccSandi selinger int i = 0, j = 0, k = 0, l = 0; 953a9d7e0ccSandi selinger 954a9d7e0ccSandi selinger /* Traverse all three arrays */ 955a9d7e0ccSandi selinger while (i<size1 && j<size2 && k<size3) { 956a9d7e0ccSandi selinger if (in1[i] < in2[j] && in1[i] < in3[k]) { 957a9d7e0ccSandi selinger out[l++] = in1[i++]; 958a9d7e0ccSandi selinger } 959a9d7e0ccSandi selinger else if (in2[j] < in1[i] && in2[j] < in3[k]) { 960a9d7e0ccSandi selinger out[l++] = in2[j++]; 961a9d7e0ccSandi selinger } 962a9d7e0ccSandi selinger else if (in3[k] < in1[i] && in3[k] < in2[j]) { 963a9d7e0ccSandi selinger out[l++] = in3[k++]; 964a9d7e0ccSandi selinger } 965a9d7e0ccSandi selinger else if (in1[i] == in2[j] && in1[i] < in3[k]) { 966a9d7e0ccSandi selinger out[l++] = in1[i]; 967a9d7e0ccSandi selinger i++, j++; 968a9d7e0ccSandi selinger } 969a9d7e0ccSandi selinger else if (in1[i] == in3[k] && in1[i] < in2[j]) { 970a9d7e0ccSandi selinger out[l++] = in1[i]; 971a9d7e0ccSandi selinger i++, k++; 972a9d7e0ccSandi selinger } 973a9d7e0ccSandi selinger else if (in3[k] == in2[j] && in2[j] < in1[i]) { 974a9d7e0ccSandi selinger out[l++] = in2[j]; 975a9d7e0ccSandi selinger k++, j++; 976a9d7e0ccSandi selinger } 977a9d7e0ccSandi selinger else if (in1[i] == in2[j] && in1[i] == in3[k]) { 978a9d7e0ccSandi selinger out[l++] = in1[i]; 979a9d7e0ccSandi selinger i++, j++, k++; 980a9d7e0ccSandi selinger } 981a9d7e0ccSandi selinger } 982a9d7e0ccSandi selinger 983a9d7e0ccSandi selinger /* Traverse two remaining arrays */ 984a9d7e0ccSandi selinger while (i<size1 && j<size2) { 985a9d7e0ccSandi selinger if (in1[i] < in2[j]) { 986a9d7e0ccSandi selinger out[l++] = in1[i++]; 987a9d7e0ccSandi selinger } 988a9d7e0ccSandi selinger else if (in1[i] > in2[j]) { 989a9d7e0ccSandi selinger out[l++] = in2[j++]; 990a9d7e0ccSandi selinger } 991a9d7e0ccSandi selinger else { 992a9d7e0ccSandi selinger out[l++] = in1[i]; 993a9d7e0ccSandi selinger i++, j++; 994a9d7e0ccSandi selinger } 995a9d7e0ccSandi selinger } 996a9d7e0ccSandi selinger 997a9d7e0ccSandi selinger while (i<size1 && k<size3) { 998a9d7e0ccSandi selinger if (in1[i] < in3[k]) { 999a9d7e0ccSandi selinger out[l++] = in1[i++]; 1000a9d7e0ccSandi selinger } 1001a9d7e0ccSandi selinger else if (in1[i] > in3[k]) { 1002a9d7e0ccSandi selinger out[l++] = in3[k++]; 1003a9d7e0ccSandi selinger } 1004a9d7e0ccSandi selinger else { 1005a9d7e0ccSandi selinger out[l++] = in1[i]; 1006a9d7e0ccSandi selinger i++, k++; 1007a9d7e0ccSandi selinger } 1008a9d7e0ccSandi selinger } 1009a9d7e0ccSandi selinger 1010a9d7e0ccSandi selinger while (k<size3 && j<size2) { 1011a9d7e0ccSandi selinger if (in3[k] < in2[j]) { 1012a9d7e0ccSandi selinger out[l++] = in3[k++]; 1013a9d7e0ccSandi selinger } 1014a9d7e0ccSandi selinger else if (in3[k] > in2[j]) { 1015a9d7e0ccSandi selinger out[l++] = in2[j++]; 1016a9d7e0ccSandi selinger } 1017a9d7e0ccSandi selinger else { 1018a9d7e0ccSandi selinger out[l++] = in3[k]; 1019a9d7e0ccSandi selinger k++, j++; 1020a9d7e0ccSandi selinger } 1021a9d7e0ccSandi selinger } 1022a9d7e0ccSandi selinger 1023a9d7e0ccSandi selinger /* Traverse one remaining array */ 1024a9d7e0ccSandi selinger while (i<size1) out[l++] = in1[i++]; 1025a9d7e0ccSandi selinger while (j<size2) out[l++] = in2[j++]; 1026a9d7e0ccSandi selinger while (k<size3) out[l++] = in3[k++]; 1027a9d7e0ccSandi selinger 1028a9d7e0ccSandi selinger *size4 = l; 1029a9d7e0ccSandi selinger } 1030a9d7e0ccSandi selinger 103173b67375Sandi selinger /* This matrix-matrix multiplication algorithm divides the multiplication into three multiplications and */ 103273b67375Sandi selinger /* adds up the products. Two of these three multiplications are performed with existing (sequential) */ 103373b67375Sandi selinger /* matrix-matrix multiplications. */ 10344222ddf1SHong Zhang PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ_seqMPI(Mat A, Mat P, PetscReal fill, Mat C) 1035a9d7e0ccSandi selinger { 1036a9d7e0ccSandi selinger MPI_Comm comm; 1037a9d7e0ccSandi selinger PetscMPIInt size; 10383cdca5ebSHong Zhang Mat_APMPI *ptap; 1039a9d7e0ccSandi selinger PetscFreeSpaceList free_space_diag=NULL, current_space=NULL; 1040a9d7e0ccSandi selinger Mat_MPIAIJ *a =(Mat_MPIAIJ*)A->data; 1041a9d7e0ccSandi selinger Mat_SeqAIJ *ad =(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_loc; 1042a9d7e0ccSandi selinger Mat_MPIAIJ *p =(Mat_MPIAIJ*)P->data; 1043a9d7e0ccSandi selinger Mat_SeqAIJ *adpd_seq, *p_off, *aopoth_seq; 1044a9d7e0ccSandi selinger PetscInt adponz, adpdnz; 10457dbaa2c6Sandi selinger PetscInt *pi_loc,*dnz,*onz; 1046a9d7e0ccSandi selinger PetscInt *adi=ad->i,*adj=ad->j,*aoi=ao->i,rstart=A->rmap->rstart; 1047a9d7e0ccSandi selinger PetscInt *lnk,i, i1=0,pnz,row,*adpoi,*adpoj, *api, *adpoJ, *aopJ, *apJ,*Jptr, aopnz, nspacedouble=0,j,nzi, 1048a9d7e0ccSandi selinger *apj,apnz, *adpdi, *adpdj, *adpdJ, *poff_i, *poff_j, *j_temp, *aopothi, *aopothj; 1049a016958eSandi selinger PetscInt am=A->rmap->n,pN=P->cmap->N,pn=P->cmap->n,pm=P->rmap->n, p_colstart, p_colend; 1050a9d7e0ccSandi selinger PetscBT lnkbt; 1051a9d7e0ccSandi selinger PetscReal afill; 1052a9d7e0ccSandi selinger PetscMPIInt rank; 1053a9d7e0ccSandi selinger Mat adpd, aopoth; 1054b92f168fSBarry Smith MatType mtype; 1055e83e5f86SFande Kong const char *prefix; 1056a9d7e0ccSandi selinger 1057a9d7e0ccSandi selinger PetscFunctionBegin; 10586718818eSStefano Zampini MatCheckProduct(C,4); 105928b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Extra product struct not empty"); 10609566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 10619566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 10629566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm, &rank)); 10639566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRangeColumn(P, &p_colstart, &p_colend)); 1064a9d7e0ccSandi selinger 10653cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 10669566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 1067a9d7e0ccSandi selinger 1068a9d7e0ccSandi selinger /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 10699566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 1070a9d7e0ccSandi selinger 1071a9d7e0ccSandi selinger /* get P_loc by taking all local rows of P */ 10729566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc)); 1073a9d7e0ccSandi selinger 1074a9d7e0ccSandi selinger p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 1075a9d7e0ccSandi selinger pi_loc = p_loc->i; 1076a9d7e0ccSandi selinger 1077a9d7e0ccSandi selinger /* Allocate memory for the i arrays of the matrices A*P, A_diag*P_off and A_offd * P */ 10789566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&api)); 10799566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&adpoi)); 1080a9d7e0ccSandi selinger 1081a9d7e0ccSandi selinger adpoi[0] = 0; 1082a9d7e0ccSandi selinger ptap->api = api; 1083a9d7e0ccSandi selinger api[0] = 0; 1084a9d7e0ccSandi selinger 1085a9d7e0ccSandi selinger /* create and initialize a linked list, will be used for both A_diag * P_loc_off and A_offd * P_oth */ 10869566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate(pN,pN,&lnk,&lnkbt)); 1087d0609cedSBarry Smith MatPreallocateBegin(comm,am,pn,dnz,onz); 1088a9d7e0ccSandi selinger 1089a9d7e0ccSandi selinger /* Symbolic calc of A_loc_diag * P_loc_diag */ 10909566063dSJacob Faibussowitsch PetscCall(MatGetOptionsPrefix(A,&prefix)); 10919566063dSJacob Faibussowitsch PetscCall(MatProductCreate(a->A,p->A,NULL,&adpd)); 10929566063dSJacob Faibussowitsch PetscCall(MatGetOptionsPrefix(A,&prefix)); 10939566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(adpd,prefix)); 10949566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(adpd,"inner_diag_")); 10954222ddf1SHong Zhang 10969566063dSJacob Faibussowitsch PetscCall(MatProductSetType(adpd,MATPRODUCT_AB)); 10979566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(adpd,"sorted")); 10989566063dSJacob Faibussowitsch PetscCall(MatProductSetFill(adpd,fill)); 10999566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions(adpd)); 11008c78258cSHong Zhang 11018c78258cSHong Zhang adpd->force_diagonals = C->force_diagonals; 11029566063dSJacob Faibussowitsch PetscCall(MatProductSymbolic(adpd)); 11034222ddf1SHong Zhang 1104a9d7e0ccSandi selinger adpd_seq = (Mat_SeqAIJ*)((adpd)->data); 1105a9d7e0ccSandi selinger adpdi = adpd_seq->i; adpdj = adpd_seq->j; 1106a9d7e0ccSandi selinger p_off = (Mat_SeqAIJ*)((p->B)->data); 1107a9d7e0ccSandi selinger poff_i = p_off->i; poff_j = p_off->j; 1108a9d7e0ccSandi selinger 110973b67375Sandi selinger /* j_temp stores indices of a result row before they are added to the linked list */ 11109566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pN+2,&j_temp)); 1111a9d7e0ccSandi selinger 1112a9d7e0ccSandi selinger /* Symbolic calc of the A_diag * p_loc_off */ 1113a9d7e0ccSandi selinger /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */ 11149566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(adi[am],PetscIntSumTruncate(aoi[am],pi_loc[pm]))),&free_space_diag)); 1115a9d7e0ccSandi selinger current_space = free_space_diag; 1116a9d7e0ccSandi selinger 1117a9d7e0ccSandi selinger for (i=0; i<am; i++) { 1118a9d7e0ccSandi selinger /* A_diag * P_loc_off */ 1119a9d7e0ccSandi selinger nzi = adi[i+1] - adi[i]; 1120a9d7e0ccSandi selinger for (j=0; j<nzi; j++) { 1121a9d7e0ccSandi selinger row = *adj++; 1122a9d7e0ccSandi selinger pnz = poff_i[row+1] - poff_i[row]; 1123a9d7e0ccSandi selinger Jptr = poff_j + poff_i[row]; 1124a9d7e0ccSandi selinger for (i1 = 0; i1 < pnz; i1++) { 1125a9d7e0ccSandi selinger j_temp[i1] = p->garray[Jptr[i1]]; 1126a9d7e0ccSandi selinger } 1127a9d7e0ccSandi selinger /* add non-zero cols of P into the sorted linked list lnk */ 11289566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(pnz,j_temp,lnk,lnkbt)); 1129a9d7e0ccSandi selinger } 1130a9d7e0ccSandi selinger 1131a9d7e0ccSandi selinger adponz = lnk[0]; 1132a9d7e0ccSandi selinger adpoi[i+1] = adpoi[i] + adponz; 1133a9d7e0ccSandi selinger 1134a9d7e0ccSandi selinger /* if free space is not available, double the total space in the list */ 1135a9d7e0ccSandi selinger if (current_space->local_remaining<adponz) { 11369566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(adponz,current_space->total_array_size),¤t_space)); 1137a9d7e0ccSandi selinger nspacedouble++; 1138a9d7e0ccSandi selinger } 1139a9d7e0ccSandi selinger 1140a9d7e0ccSandi selinger /* Copy data into free space, then initialize lnk */ 11419566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean(pN,adponz,current_space->array,lnk,lnkbt)); 1142a9d7e0ccSandi selinger 1143a9d7e0ccSandi selinger current_space->array += adponz; 1144a9d7e0ccSandi selinger current_space->local_used += adponz; 1145a9d7e0ccSandi selinger current_space->local_remaining -= adponz; 1146a9d7e0ccSandi selinger } 1147a9d7e0ccSandi selinger 1148a9d7e0ccSandi selinger /* Symbolic calc of A_off * P_oth */ 11499566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(a->B,prefix)); 11509566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(a->B,"inner_offdiag_")); 11519566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF,&aopoth)); 11529566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(a->B, ptap->P_oth, fill, aopoth)); 1153a9d7e0ccSandi selinger aopoth_seq = (Mat_SeqAIJ*)((aopoth)->data); 1154a9d7e0ccSandi selinger aopothi = aopoth_seq->i; aopothj = aopoth_seq->j; 1155a9d7e0ccSandi selinger 1156a9d7e0ccSandi selinger /* Allocate space for apj, adpj, aopj, ... */ 1157a9d7e0ccSandi selinger /* destroy lists of free space and other temporary array(s) */ 1158a9d7e0ccSandi selinger 11599566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(aopothi[am] + adpoi[am] + adpdi[am]+2, &ptap->apj)); 11609566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(adpoi[am]+2, &adpoj)); 1161a9d7e0ccSandi selinger 1162a9d7e0ccSandi selinger /* Copy from linked list to j-array */ 11639566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space_diag,adpoj)); 11649566063dSJacob Faibussowitsch PetscCall(PetscLLDestroy(lnk,lnkbt)); 1165a9d7e0ccSandi selinger 1166a9d7e0ccSandi selinger adpoJ = adpoj; 1167a9d7e0ccSandi selinger adpdJ = adpdj; 1168a9d7e0ccSandi selinger aopJ = aopothj; 1169a9d7e0ccSandi selinger apj = ptap->apj; 1170a9d7e0ccSandi selinger apJ = apj; /* still empty */ 1171a9d7e0ccSandi selinger 1172a9d7e0ccSandi selinger /* Merge j-arrays of A_off * P, A_diag * P_loc_off, and */ 1173a9d7e0ccSandi selinger /* A_diag * P_loc_diag to get A*P */ 1174a9d7e0ccSandi selinger for (i = 0; i < am; i++) { 1175a9d7e0ccSandi selinger aopnz = aopothi[i+1] - aopothi[i]; 1176a9d7e0ccSandi selinger adponz = adpoi[i+1] - adpoi[i]; 1177a9d7e0ccSandi selinger adpdnz = adpdi[i+1] - adpdi[i]; 1178a9d7e0ccSandi selinger 1179a9d7e0ccSandi selinger /* Correct indices from A_diag*P_diag */ 1180a9d7e0ccSandi selinger for (i1 = 0; i1 < adpdnz; i1++) { 1181a016958eSandi selinger adpdJ[i1] += p_colstart; 1182a9d7e0ccSandi selinger } 1183a9d7e0ccSandi selinger /* Merge j-arrays of A_diag * P_loc_off and A_diag * P_loc_diag and A_off * P_oth */ 1184a9d7e0ccSandi selinger Merge3SortedArrays(adponz, adpoJ, adpdnz, adpdJ, aopnz, aopJ, &apnz, apJ); 11859566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+rstart, apnz, apJ, dnz, onz)); 1186a9d7e0ccSandi selinger 1187a9d7e0ccSandi selinger aopJ += aopnz; 1188a9d7e0ccSandi selinger adpoJ += adponz; 1189a9d7e0ccSandi selinger adpdJ += adpdnz; 1190a9d7e0ccSandi selinger apJ += apnz; 1191a9d7e0ccSandi selinger api[i+1] = api[i] + apnz; 1192a9d7e0ccSandi selinger } 1193a9d7e0ccSandi selinger 1194a9d7e0ccSandi selinger /* malloc apa to store dense row A[i,:]*P */ 11959566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(pN+2,&ptap->apa)); 1196a9d7e0ccSandi selinger 11974222ddf1SHong Zhang /* create and assemble symbolic parallel matrix C */ 11989566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,am,pn,PETSC_DETERMINE,PETSC_DETERMINE)); 11999566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,P)); 12009566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 12019566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 12029566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 1203d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 120417f75172Sandi selinger 12059566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api)); 12069566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 12079566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 12089566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 120917f75172Sandi selinger 12104222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable; 12114222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AB; 1212a9d7e0ccSandi selinger 12134222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 12146718818eSStefano Zampini C->product->data = ptap; 12156718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_MatMatMult; 1216a9d7e0ccSandi selinger 1217a9d7e0ccSandi selinger /* set MatInfo */ 121873b67375Sandi selinger afill = (PetscReal)api[am]/(adi[am]+aoi[am]+pi_loc[pm]+1) + 1.e-5; 121973b67375Sandi selinger if (afill < 1.0) afill = 1.0; 12204222ddf1SHong Zhang C->info.mallocs = nspacedouble; 12214222ddf1SHong Zhang C->info.fill_ratio_given = fill; 12224222ddf1SHong Zhang C->info.fill_ratio_needed = afill; 1223a9d7e0ccSandi selinger 1224a9d7e0ccSandi selinger #if defined(PETSC_USE_INFO) 1225a9d7e0ccSandi selinger if (api[am]) { 12269566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 12279566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n",(double)afill)); 1228a9d7e0ccSandi selinger } else { 12299566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 1230a9d7e0ccSandi selinger } 1231a9d7e0ccSandi selinger #endif 1232a9d7e0ccSandi selinger 12339566063dSJacob Faibussowitsch PetscCall(MatDestroy(&aopoth)); 12349566063dSJacob Faibussowitsch PetscCall(MatDestroy(&adpd)); 12359566063dSJacob Faibussowitsch PetscCall(PetscFree(j_temp)); 12369566063dSJacob Faibussowitsch PetscCall(PetscFree(adpoj)); 12379566063dSJacob Faibussowitsch PetscCall(PetscFree(adpoi)); 1238a9d7e0ccSandi selinger PetscFunctionReturn(0); 1239a9d7e0ccSandi selinger } 1240a9d7e0ccSandi selinger 1241f96d37fcSHong Zhang /*-------------------------------------------------------------------------*/ 1242c216dbf3SHong Zhang /* This routine only works when scall=MAT_REUSE_MATRIX! */ 1243c216dbf3SHong Zhang PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_matmatmult(Mat P,Mat A,Mat C) 1244c216dbf3SHong Zhang { 12456718818eSStefano Zampini Mat_APMPI *ptap; 124660552ceaSHong Zhang Mat Pt; 1247c216dbf3SHong Zhang 1248c216dbf3SHong Zhang PetscFunctionBegin; 12496718818eSStefano Zampini MatCheckProduct(C,3); 12506718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 125128b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 125228b400f6SJacob Faibussowitsch PetscCheck(ptap->Pt,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be reused. Do not call MatProductClear()"); 125360552ceaSHong Zhang 125460552ceaSHong Zhang Pt = ptap->Pt; 12559566063dSJacob Faibussowitsch PetscCall(MatTranspose(P,MAT_REUSE_MATRIX,&Pt)); 12569566063dSJacob Faibussowitsch PetscCall(MatMatMultNumeric_MPIAIJ_MPIAIJ(Pt,A,C)); 1257f96d37fcSHong Zhang PetscFunctionReturn(0); 1258f96d37fcSHong Zhang } 1259f96d37fcSHong Zhang 1260c5af039cSHong Zhang /* This routine is modified from MatPtAPSymbolic_MPIAIJ_MPIAIJ() */ 12614222ddf1SHong Zhang PetscErrorCode MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(Mat P,Mat A,PetscReal fill,Mat C) 1262912837bbSHong Zhang { 12633cdca5ebSHong Zhang Mat_APMPI *ptap; 12646718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data; 1265912837bbSHong Zhang MPI_Comm comm; 1266912837bbSHong Zhang PetscMPIInt size,rank; 1267912837bbSHong Zhang PetscFreeSpaceList free_space=NULL,current_space=NULL; 12684ede91e3SHong Zhang PetscInt pn=P->cmap->n,aN=A->cmap->N,an=A->cmap->n; 126969be949aSStefano Zampini PetscInt *lnk,i,k,nsend,rstart; 1270912837bbSHong Zhang PetscBT lnkbt; 1271ec4bef21SJose E. Roman PetscMPIInt tagi,tagj,*len_si,*len_s,*len_ri,nrecv; 1272ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted=0; 12738c78258cSHong Zhang PetscInt **buf_rj,**buf_ri,**buf_ri_k,row,ncols,*cols; 12747f3e1f30SHong Zhang PetscInt len,proc,*dnz,*onz,*owners,nzi; 1275912837bbSHong Zhang PetscInt nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci; 1276912837bbSHong Zhang MPI_Request *swaits,*rwaits; 1277912837bbSHong Zhang MPI_Status *sstatus,rstatus; 1278912837bbSHong Zhang PetscLayout rowmap; 1279912837bbSHong Zhang PetscInt *owners_co,*coi,*coj; /* i and j array of (p->B)^T*A*P - used in the communication */ 1280912837bbSHong Zhang PetscMPIInt *len_r,*id_r; /* array of length of comm->size, store send/recv matrix values */ 12817f3e1f30SHong Zhang PetscInt *Jptr,*prmap=p->garray,con,j,Crmax; 12824ede91e3SHong Zhang Mat_SeqAIJ *a_loc,*c_loc,*c_oth; 1283912837bbSHong Zhang PetscTable ta; 1284b92f168fSBarry Smith MatType mtype; 1285e83e5f86SFande Kong const char *prefix; 1286912837bbSHong Zhang 1287912837bbSHong Zhang PetscFunctionBegin; 12889566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 12899566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 12909566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1291912837bbSHong Zhang 12924222ddf1SHong Zhang /* create symbolic parallel matrix C */ 12939566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 12949566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 1295912837bbSHong Zhang 12964222ddf1SHong Zhang C->ops->transposematmultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable; 1297912837bbSHong Zhang 12983cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 12999566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 1300912837bbSHong Zhang ptap->reuse = MAT_INITIAL_MATRIX; 1301912837bbSHong Zhang 1302912837bbSHong Zhang /* (0) compute Rd = Pd^T, Ro = Po^T */ 1303912837bbSHong Zhang /* --------------------------------- */ 1304*acd337a6SBarry Smith PetscCall(MatTranspose(p->A,MAT_INITIAL_MATRIX,&ptap->Rd)); 1305*acd337a6SBarry Smith PetscCall(MatTranspose(p->B,MAT_INITIAL_MATRIX,&ptap->Ro)); 1306912837bbSHong Zhang 1307669e9c86SHong Zhang /* (1) compute symbolic A_loc */ 1308669e9c86SHong Zhang /* ---------------------------*/ 13099566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_INITIAL_MATRIX,&ptap->A_loc)); 1310912837bbSHong Zhang 13114ede91e3SHong Zhang /* (2-1) compute symbolic C_oth = Ro*A_loc */ 1312912837bbSHong Zhang /* ------------------------------------ */ 13139566063dSJacob Faibussowitsch PetscCall(MatGetOptionsPrefix(A,&prefix)); 13149566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(ptap->Ro,prefix)); 13159566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(ptap->Ro,"inner_offdiag_")); 13169566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF,&ptap->C_oth)); 13179566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Ro,ptap->A_loc,fill,ptap->C_oth)); 1318912837bbSHong Zhang 1319912837bbSHong Zhang /* (3) send coj of C_oth to other processors */ 1320912837bbSHong Zhang /* ------------------------------------------ */ 1321912837bbSHong Zhang /* determine row ownership */ 13229566063dSJacob Faibussowitsch PetscCall(PetscLayoutCreate(comm,&rowmap)); 1323912837bbSHong Zhang rowmap->n = pn; 1324912837bbSHong Zhang rowmap->bs = 1; 13259566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(rowmap)); 1326912837bbSHong Zhang owners = rowmap->range; 1327912837bbSHong Zhang 1328912837bbSHong Zhang /* determine the number of messages to send, their lengths */ 13299566063dSJacob Faibussowitsch PetscCall(PetscMalloc4(size,&len_s,size,&len_si,size,&sstatus,size+2,&owners_co)); 13309566063dSJacob Faibussowitsch PetscCall(PetscArrayzero(len_s,size)); 13319566063dSJacob Faibussowitsch PetscCall(PetscArrayzero(len_si,size)); 1332912837bbSHong Zhang 1333912837bbSHong Zhang c_oth = (Mat_SeqAIJ*)ptap->C_oth->data; 1334912837bbSHong Zhang coi = c_oth->i; coj = c_oth->j; 1335912837bbSHong Zhang con = ptap->C_oth->rmap->n; 1336912837bbSHong Zhang proc = 0; 1337912837bbSHong Zhang for (i=0; i<con; i++) { 1338912837bbSHong Zhang while (prmap[i] >= owners[proc+1]) proc++; 13397f3e1f30SHong Zhang len_si[proc]++; /* num of rows in Co(=Pt*A) to be sent to [proc] */ 1340912837bbSHong Zhang len_s[proc] += coi[i+1] - coi[i]; /* num of nonzeros in Co to be sent to [proc] */ 1341912837bbSHong Zhang } 1342912837bbSHong Zhang 1343912837bbSHong Zhang len = 0; /* max length of buf_si[], see (4) */ 1344912837bbSHong Zhang owners_co[0] = 0; 1345912837bbSHong Zhang nsend = 0; 1346912837bbSHong Zhang for (proc=0; proc<size; proc++) { 1347912837bbSHong Zhang owners_co[proc+1] = owners_co[proc] + len_si[proc]; 1348912837bbSHong Zhang if (len_s[proc]) { 1349912837bbSHong Zhang nsend++; 1350912837bbSHong Zhang len_si[proc] = 2*(len_si[proc] + 1); /* length of buf_si to be sent to [proc] */ 1351912837bbSHong Zhang len += len_si[proc]; 1352912837bbSHong Zhang } 1353912837bbSHong Zhang } 1354912837bbSHong Zhang 1355912837bbSHong Zhang /* determine the number and length of messages to receive for coi and coj */ 13569566063dSJacob Faibussowitsch PetscCall(PetscGatherNumberOfMessages(comm,NULL,len_s,&nrecv)); 13579566063dSJacob Faibussowitsch PetscCall(PetscGatherMessageLengths2(comm,nsend,nrecv,len_s,len_si,&id_r,&len_r,&len_ri)); 1358912837bbSHong Zhang 1359912837bbSHong Zhang /* post the Irecv and Isend of coj */ 13609566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagj)); 13619566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagj,nrecv,id_r,len_r,&buf_rj,&rwaits)); 13629566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(nsend+1,&swaits)); 1363912837bbSHong Zhang for (proc=0, k=0; proc<size; proc++) { 1364912837bbSHong Zhang if (!len_s[proc]) continue; 1365912837bbSHong Zhang i = owners_co[proc]; 13669566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tagj,comm,swaits+k)); 1367912837bbSHong Zhang k++; 1368912837bbSHong Zhang } 1369912837bbSHong Zhang 13707f3e1f30SHong Zhang /* (2-2) compute symbolic C_loc = Rd*A_loc */ 1371912837bbSHong Zhang /* ---------------------------------------- */ 13729566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(ptap->Rd,prefix)); 13739566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(ptap->Rd,"inner_diag_")); 13749566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF,&ptap->C_loc)); 13759566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Rd,ptap->A_loc,fill,ptap->C_loc)); 1376912837bbSHong Zhang c_loc = (Mat_SeqAIJ*)ptap->C_loc->data; 1377912837bbSHong Zhang 1378912837bbSHong Zhang /* receives coj are complete */ 1379912837bbSHong Zhang for (i=0; i<nrecv; i++) { 13809566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(nrecv,rwaits,&icompleted,&rstatus)); 1381912837bbSHong Zhang } 13829566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 13839566063dSJacob Faibussowitsch if (nsend) PetscCallMPI(MPI_Waitall(nsend,swaits,sstatus)); 1384912837bbSHong Zhang 1385912837bbSHong Zhang /* add received column indices into ta to update Crmax */ 13864ede91e3SHong Zhang a_loc = (Mat_SeqAIJ*)(ptap->A_loc)->data; 1387669e9c86SHong Zhang 1388669e9c86SHong Zhang /* create and initialize a linked list */ 13899566063dSJacob Faibussowitsch PetscCall(PetscTableCreate(an,aN,&ta)); /* for compute Crmax */ 13904ede91e3SHong Zhang MatRowMergeMax_SeqAIJ(a_loc,ptap->A_loc->rmap->N,ta); 1391669e9c86SHong Zhang 1392912837bbSHong Zhang for (k=0; k<nrecv; k++) {/* k-th received message */ 1393912837bbSHong Zhang Jptr = buf_rj[k]; 1394912837bbSHong Zhang for (j=0; j<len_r[k]; j++) { 13959566063dSJacob Faibussowitsch PetscCall(PetscTableAdd(ta,*(Jptr+j)+1,1,INSERT_VALUES)); 1396912837bbSHong Zhang } 1397912837bbSHong Zhang } 13989566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Crmax)); 13999566063dSJacob Faibussowitsch PetscCall(PetscTableDestroy(&ta)); 1400912837bbSHong Zhang 1401912837bbSHong Zhang /* (4) send and recv coi */ 1402912837bbSHong Zhang /*-----------------------*/ 14039566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagi)); 14049566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagi,nrecv,id_r,len_ri,&buf_ri,&rwaits)); 14059566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(len+1,&buf_s)); 1406912837bbSHong Zhang buf_si = buf_s; /* points to the beginning of k-th msg to be sent */ 1407912837bbSHong Zhang for (proc=0,k=0; proc<size; proc++) { 1408912837bbSHong Zhang if (!len_s[proc]) continue; 1409912837bbSHong Zhang /* form outgoing message for i-structure: 1410912837bbSHong Zhang buf_si[0]: nrows to be sent 1411912837bbSHong Zhang [1:nrows]: row index (global) 1412912837bbSHong Zhang [nrows+1:2*nrows+1]: i-structure index 1413912837bbSHong Zhang */ 1414912837bbSHong Zhang /*-------------------------------------------*/ 1415912837bbSHong Zhang nrows = len_si[proc]/2 - 1; /* num of rows in Co to be sent to [proc] */ 1416912837bbSHong Zhang buf_si_i = buf_si + nrows+1; 1417912837bbSHong Zhang buf_si[0] = nrows; 1418912837bbSHong Zhang buf_si_i[0] = 0; 1419912837bbSHong Zhang nrows = 0; 1420912837bbSHong Zhang for (i=owners_co[proc]; i<owners_co[proc+1]; i++) { 1421912837bbSHong Zhang nzi = coi[i+1] - coi[i]; 1422912837bbSHong Zhang buf_si_i[nrows+1] = buf_si_i[nrows] + nzi; /* i-structure */ 1423912837bbSHong Zhang buf_si[nrows+1] = prmap[i] -owners[proc]; /* local row index */ 1424912837bbSHong Zhang nrows++; 1425912837bbSHong Zhang } 14269566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tagi,comm,swaits+k)); 1427912837bbSHong Zhang k++; 1428912837bbSHong Zhang buf_si += len_si[proc]; 1429912837bbSHong Zhang } 1430912837bbSHong Zhang for (i=0; i<nrecv; i++) { 14319566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(nrecv,rwaits,&icompleted,&rstatus)); 1432912837bbSHong Zhang } 14339566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 14349566063dSJacob Faibussowitsch if (nsend) PetscCallMPI(MPI_Waitall(nsend,swaits,sstatus)); 1435912837bbSHong Zhang 14369566063dSJacob Faibussowitsch PetscCall(PetscFree4(len_s,len_si,sstatus,owners_co)); 14379566063dSJacob Faibussowitsch PetscCall(PetscFree(len_ri)); 14389566063dSJacob Faibussowitsch PetscCall(PetscFree(swaits)); 14399566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_s)); 1440912837bbSHong Zhang 14414222ddf1SHong Zhang /* (5) compute the local portion of C */ 1442912837bbSHong Zhang /* ------------------------------------------ */ 14434222ddf1SHong Zhang /* set initial free space to be Crmax, sufficient for holding nozeros in each row of C */ 14449566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(Crmax,&free_space)); 1445912837bbSHong Zhang current_space = free_space; 1446912837bbSHong Zhang 14479566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(nrecv,&buf_ri_k,nrecv,&nextrow,nrecv,&nextci)); 1448912837bbSHong Zhang for (k=0; k<nrecv; k++) { 1449912837bbSHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 1450912837bbSHong Zhang nrows = *buf_ri_k[k]; 1451912837bbSHong Zhang nextrow[k] = buf_ri_k[k] + 1; /* next row number of k-th recved i-structure */ 1452a5b23f4aSJose E. Roman nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th recved i-structure */ 1453912837bbSHong Zhang } 1454912837bbSHong Zhang 1455d0609cedSBarry Smith MatPreallocateBegin(comm,pn,an,dnz,onz); 14569566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate(Crmax,aN,&lnk,&lnkbt)); 14578c78258cSHong Zhang for (i=0; i<pn; i++) { /* for each local row of C */ 14584222ddf1SHong Zhang /* add C_loc into C */ 1459912837bbSHong Zhang nzi = c_loc->i[i+1] - c_loc->i[i]; 1460912837bbSHong Zhang Jptr = c_loc->j + c_loc->i[i]; 14619566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(nzi,Jptr,lnk,lnkbt)); 1462912837bbSHong Zhang 1463912837bbSHong Zhang /* add received col data into lnk */ 1464912837bbSHong Zhang for (k=0; k<nrecv; k++) { /* k-th received message */ 1465912837bbSHong Zhang if (i == *nextrow[k]) { /* i-th row */ 1466912837bbSHong Zhang nzi = *(nextci[k]+1) - *nextci[k]; 1467912837bbSHong Zhang Jptr = buf_rj[k] + *nextci[k]; 14689566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(nzi,Jptr,lnk,lnkbt)); 1469912837bbSHong Zhang nextrow[k]++; nextci[k]++; 1470912837bbSHong Zhang } 1471912837bbSHong Zhang } 14728c78258cSHong Zhang 14738c78258cSHong Zhang /* add missing diagonal entry */ 14748c78258cSHong Zhang if (C->force_diagonals) { 14758c78258cSHong Zhang k = i + owners[rank]; /* column index */ 14769566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(1,&k,lnk,lnkbt)); 14778c78258cSHong Zhang } 14788c78258cSHong Zhang 1479912837bbSHong Zhang nzi = lnk[0]; 1480912837bbSHong Zhang 1481912837bbSHong Zhang /* copy data into free space, then initialize lnk */ 14829566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean(aN,nzi,current_space->array,lnk,lnkbt)); 14839566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+owners[rank],nzi,current_space->array,dnz,onz)); 1484912837bbSHong Zhang } 14859566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 14869566063dSJacob Faibussowitsch PetscCall(PetscLLDestroy(lnk,lnkbt)); 14879566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceDestroy(free_space)); 1488912837bbSHong Zhang 1489912837bbSHong Zhang /* local sizes and preallocation */ 14909566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,pn,an,PETSC_DETERMINE,PETSC_DETERMINE)); 14919566063dSJacob Faibussowitsch if (P->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->rmap,P->cmap->bs)); 14929566063dSJacob Faibussowitsch if (A->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->cmap,A->cmap->bs)); 14939566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 1494d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 1495912837bbSHong Zhang 149669be949aSStefano Zampini /* add C_loc and C_oth to C */ 14979566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRange(C,&rstart,NULL)); 149869be949aSStefano Zampini for (i=0; i<pn; i++) { 14998c78258cSHong Zhang ncols = c_loc->i[i+1] - c_loc->i[i]; 15008c78258cSHong Zhang cols = c_loc->j + c_loc->i[i]; 15018c78258cSHong Zhang row = rstart + i; 15029566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,(const PetscInt*)&row,ncols,(const PetscInt*)cols,NULL,INSERT_VALUES)); 15038c78258cSHong Zhang 15048c78258cSHong Zhang if (C->force_diagonals) { 15059566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,(const PetscInt*)&row,1,(const PetscInt*)&row,NULL,INSERT_VALUES)); 15068c78258cSHong Zhang } 150769be949aSStefano Zampini } 150869be949aSStefano Zampini for (i=0; i<con; i++) { 15098c78258cSHong Zhang ncols = c_oth->i[i+1] - c_oth->i[i]; 15108c78258cSHong Zhang cols = c_oth->j + c_oth->i[i]; 15118c78258cSHong Zhang row = prmap[i]; 15129566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,(const PetscInt*)&row,ncols,(const PetscInt*)cols,NULL,INSERT_VALUES)); 151369be949aSStefano Zampini } 15149566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 15159566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 15169566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 151769be949aSStefano Zampini 1518912837bbSHong Zhang /* members in merge */ 15199566063dSJacob Faibussowitsch PetscCall(PetscFree(id_r)); 15209566063dSJacob Faibussowitsch PetscCall(PetscFree(len_r)); 15219566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_ri[0])); 15229566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_ri)); 15239566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_rj[0])); 15249566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_rj)); 15259566063dSJacob Faibussowitsch PetscCall(PetscLayoutDestroy(&rowmap)); 1526912837bbSHong Zhang 15274222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 15286718818eSStefano Zampini C->product->data = ptap; 15296718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 1530912837bbSHong Zhang PetscFunctionReturn(0); 1531912837bbSHong Zhang } 1532912837bbSHong Zhang 1533912837bbSHong Zhang PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable(Mat P,Mat A,Mat C) 1534912837bbSHong Zhang { 15356718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data; 1536669e9c86SHong Zhang Mat_SeqAIJ *c_seq; 15376718818eSStefano Zampini Mat_APMPI *ptap; 1538669e9c86SHong Zhang Mat A_loc,C_loc,C_oth; 1539912837bbSHong Zhang PetscInt i,rstart,rend,cm,ncols,row; 1540912837bbSHong Zhang const PetscInt *cols; 1541912837bbSHong Zhang const PetscScalar *vals; 1542912837bbSHong Zhang 1543912837bbSHong Zhang PetscFunctionBegin; 15446718818eSStefano Zampini MatCheckProduct(C,3); 15456718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 154628b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 154728b400f6SJacob Faibussowitsch PetscCheck(ptap->A_loc,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be reused. Do not call MatProductClear()"); 15489566063dSJacob Faibussowitsch PetscCall(MatZeroEntries(C)); 1549912837bbSHong Zhang 1550912837bbSHong Zhang if (ptap->reuse == MAT_REUSE_MATRIX) { 1551669e9c86SHong Zhang /* These matrices are obtained in MatTransposeMatMultSymbolic() */ 1552669e9c86SHong Zhang /* 1) get R = Pd^T, Ro = Po^T */ 1553669e9c86SHong Zhang /*----------------------------*/ 1554*acd337a6SBarry Smith PetscCall(MatTranspose(p->A,MAT_REUSE_MATRIX,&ptap->Rd)); 1555*acd337a6SBarry Smith PetscCall(MatTranspose(p->B,MAT_REUSE_MATRIX,&ptap->Ro)); 1556669e9c86SHong Zhang 1557669e9c86SHong Zhang /* 2) compute numeric A_loc */ 1558669e9c86SHong Zhang /*--------------------------*/ 15599566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_REUSE_MATRIX,&ptap->A_loc)); 1560912837bbSHong Zhang } 1561912837bbSHong Zhang 1562669e9c86SHong Zhang /* 3) C_loc = Rd*A_loc, C_oth = Ro*A_loc */ 15634ede91e3SHong Zhang A_loc = ptap->A_loc; 15649566063dSJacob Faibussowitsch PetscCall(((ptap->C_loc)->ops->matmultnumeric)(ptap->Rd,A_loc,ptap->C_loc)); 15659566063dSJacob Faibussowitsch PetscCall(((ptap->C_oth)->ops->matmultnumeric)(ptap->Ro,A_loc,ptap->C_oth)); 1566912837bbSHong Zhang C_loc = ptap->C_loc; 1567912837bbSHong Zhang C_oth = ptap->C_oth; 1568912837bbSHong Zhang 156969be949aSStefano Zampini /* add C_loc and C_oth to C */ 15709566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRange(C,&rstart,&rend)); 1571912837bbSHong Zhang 1572912837bbSHong Zhang /* C_loc -> C */ 1573912837bbSHong Zhang cm = C_loc->rmap->N; 1574912837bbSHong Zhang c_seq = (Mat_SeqAIJ*)C_loc->data; 1575912837bbSHong Zhang cols = c_seq->j; 1576912837bbSHong Zhang vals = c_seq->a; 1577912837bbSHong Zhang for (i=0; i<cm; i++) { 1578912837bbSHong Zhang ncols = c_seq->i[i+1] - c_seq->i[i]; 1579912837bbSHong Zhang row = rstart + i; 15809566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,ncols,cols,vals,ADD_VALUES)); 1581912837bbSHong Zhang cols += ncols; vals += ncols; 1582912837bbSHong Zhang } 1583912837bbSHong Zhang 1584912837bbSHong Zhang /* Co -> C, off-processor part */ 1585912837bbSHong Zhang cm = C_oth->rmap->N; 1586912837bbSHong Zhang c_seq = (Mat_SeqAIJ*)C_oth->data; 1587912837bbSHong Zhang cols = c_seq->j; 1588912837bbSHong Zhang vals = c_seq->a; 1589912837bbSHong Zhang for (i=0; i<cm; i++) { 1590912837bbSHong Zhang ncols = c_seq->i[i+1] - c_seq->i[i]; 1591912837bbSHong Zhang row = p->garray[i]; 15929566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,ncols,cols,vals,ADD_VALUES)); 1593912837bbSHong Zhang cols += ncols; vals += ncols; 1594912837bbSHong Zhang } 15959566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 15969566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 15979566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 1598912837bbSHong Zhang 1599912837bbSHong Zhang ptap->reuse = MAT_REUSE_MATRIX; 1600912837bbSHong Zhang PetscFunctionReturn(0); 1601912837bbSHong Zhang } 1602912837bbSHong Zhang 16036da69ca6SHong Zhang PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ(Mat P,Mat A,Mat C) 1604d6ab1dc0SHong Zhang { 1605d6ab1dc0SHong Zhang Mat_Merge_SeqsToMPI *merge; 16066718818eSStefano Zampini Mat_MPIAIJ *p =(Mat_MPIAIJ*)P->data; 1607d6ab1dc0SHong Zhang Mat_SeqAIJ *pd=(Mat_SeqAIJ*)(p->A)->data,*po=(Mat_SeqAIJ*)(p->B)->data; 16083cdca5ebSHong Zhang Mat_APMPI *ptap; 1609e745005bSHong Zhang PetscInt *adj; 1610e745005bSHong Zhang PetscInt i,j,k,anz,pnz,row,*cj,nexta; 1611e745005bSHong Zhang MatScalar *ada,*ca,valtmp; 1612d6ab1dc0SHong Zhang PetscInt am=A->rmap->n,cm=C->rmap->n,pon=(p->B)->cmap->n; 1613ce94432eSBarry Smith MPI_Comm comm; 1614d6ab1dc0SHong Zhang PetscMPIInt size,rank,taga,*len_s; 1615d6ab1dc0SHong Zhang PetscInt *owners,proc,nrows,**buf_ri_k,**nextrow,**nextci; 1616d6ab1dc0SHong Zhang PetscInt **buf_ri,**buf_rj; 1617d6ab1dc0SHong Zhang PetscInt cnz=0,*bj_i,*bi,*bj,bnz,nextcj; /* bi,bj,ba: local array of C(mpi mat) */ 1618d6ab1dc0SHong Zhang MPI_Request *s_waits,*r_waits; 1619d6ab1dc0SHong Zhang MPI_Status *status; 1620d6ab1dc0SHong Zhang MatScalar **abuf_r,*ba_i,*pA,*coa,*ba; 1621ce496241SStefano Zampini const PetscScalar *dummy; 16224e201bd7SHong Zhang PetscInt *ai,*aj,*coi,*coj,*poJ,*pdJ; 1623d6ab1dc0SHong Zhang Mat A_loc; 1624d6ab1dc0SHong Zhang Mat_SeqAIJ *a_loc; 1625d6ab1dc0SHong Zhang 1626d6ab1dc0SHong Zhang PetscFunctionBegin; 16276718818eSStefano Zampini MatCheckProduct(C,3); 16286718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 162928b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 163028b400f6SJacob Faibussowitsch PetscCheck(ptap->A_loc,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be reused. Do not call MatProductClear()"); 16319566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 16329566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 16339566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1634d6ab1dc0SHong Zhang 1635d6ab1dc0SHong Zhang merge = ptap->merge; 1636d6ab1dc0SHong Zhang 1637e745005bSHong Zhang /* 2) compute numeric C_seq = P_loc^T*A_loc */ 1638e745005bSHong Zhang /*------------------------------------------*/ 1639d6ab1dc0SHong Zhang /* get data from symbolic products */ 1640d6ab1dc0SHong Zhang coi = merge->coi; coj = merge->coj; 16419566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(coi[pon]+1,&coa)); 1642d6ab1dc0SHong Zhang bi = merge->bi; bj = merge->bj; 1643d6ab1dc0SHong Zhang owners = merge->rowmap->range; 16449566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(bi[cm]+1,&ba)); 1645d6ab1dc0SHong Zhang 1646d6ab1dc0SHong Zhang /* get A_loc by taking all local rows of A */ 1647d6ab1dc0SHong Zhang A_loc = ptap->A_loc; 16489566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_REUSE_MATRIX,&A_loc)); 1649d6ab1dc0SHong Zhang a_loc = (Mat_SeqAIJ*)(A_loc)->data; 1650d6ab1dc0SHong Zhang ai = a_loc->i; 1651d6ab1dc0SHong Zhang aj = a_loc->j; 1652d6ab1dc0SHong Zhang 1653ce496241SStefano Zampini /* trigger copy to CPU */ 16549566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(p->A,&dummy)); 16559566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(p->A,&dummy)); 16569566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(p->B,&dummy)); 16579566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(p->B,&dummy)); 1658d6ab1dc0SHong Zhang for (i=0; i<am; i++) { 1659d6ab1dc0SHong Zhang anz = ai[i+1] - ai[i]; 1660d6ab1dc0SHong Zhang adj = aj + ai[i]; 1661d6ab1dc0SHong Zhang ada = a_loc->a + ai[i]; 1662d6ab1dc0SHong Zhang 1663d6ab1dc0SHong Zhang /* 2-b) Compute Cseq = P_loc[i,:]^T*A[i,:] using outer product */ 1664e745005bSHong Zhang /*-------------------------------------------------------------*/ 1665d6ab1dc0SHong Zhang /* put the value into Co=(p->B)^T*A (off-diagonal part, send to others) */ 1666d6ab1dc0SHong Zhang pnz = po->i[i+1] - po->i[i]; 1667d6ab1dc0SHong Zhang poJ = po->j + po->i[i]; 1668d6ab1dc0SHong Zhang pA = po->a + po->i[i]; 1669d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1670d6ab1dc0SHong Zhang row = poJ[j]; 1671d6ab1dc0SHong Zhang cj = coj + coi[row]; 1672d6ab1dc0SHong Zhang ca = coa + coi[row]; 1673e745005bSHong Zhang /* perform sparse axpy */ 1674e745005bSHong Zhang nexta = 0; 1675d6ab1dc0SHong Zhang valtmp = pA[j]; 1676e745005bSHong Zhang for (k=0; nexta<anz; k++) { 1677e745005bSHong Zhang if (cj[k] == adj[nexta]) { 1678e745005bSHong Zhang ca[k] += valtmp*ada[nexta]; 1679e745005bSHong Zhang nexta++; 1680d6ab1dc0SHong Zhang } 1681e745005bSHong Zhang } 16829566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*anz)); 1683d6ab1dc0SHong Zhang } 1684d6ab1dc0SHong Zhang 1685d6ab1dc0SHong Zhang /* put the value into Cd (diagonal part) */ 1686d6ab1dc0SHong Zhang pnz = pd->i[i+1] - pd->i[i]; 1687d6ab1dc0SHong Zhang pdJ = pd->j + pd->i[i]; 1688d6ab1dc0SHong Zhang pA = pd->a + pd->i[i]; 1689d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1690d6ab1dc0SHong Zhang row = pdJ[j]; 1691d6ab1dc0SHong Zhang cj = bj + bi[row]; 1692d6ab1dc0SHong Zhang ca = ba + bi[row]; 1693e745005bSHong Zhang /* perform sparse axpy */ 1694e745005bSHong Zhang nexta = 0; 1695d6ab1dc0SHong Zhang valtmp = pA[j]; 1696e745005bSHong Zhang for (k=0; nexta<anz; k++) { 1697e745005bSHong Zhang if (cj[k] == adj[nexta]) { 1698e745005bSHong Zhang ca[k] += valtmp*ada[nexta]; 1699e745005bSHong Zhang nexta++; 1700d6ab1dc0SHong Zhang } 1701e745005bSHong Zhang } 17029566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*anz)); 1703d6ab1dc0SHong Zhang } 1704d6ab1dc0SHong Zhang } 1705d6ab1dc0SHong Zhang 1706d6ab1dc0SHong Zhang /* 3) send and recv matrix values coa */ 1707d6ab1dc0SHong Zhang /*------------------------------------*/ 1708d6ab1dc0SHong Zhang buf_ri = merge->buf_ri; 1709d6ab1dc0SHong Zhang buf_rj = merge->buf_rj; 1710d6ab1dc0SHong Zhang len_s = merge->len_s; 17119566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&taga)); 17129566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvScalar(comm,taga,merge->nrecv,merge->id_r,merge->len_r,&abuf_r,&r_waits)); 1713d6ab1dc0SHong Zhang 17149566063dSJacob Faibussowitsch PetscCall(PetscMalloc2(merge->nsend+1,&s_waits,size,&status)); 1715d6ab1dc0SHong Zhang for (proc=0,k=0; proc<size; proc++) { 1716d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1717d6ab1dc0SHong Zhang i = merge->owners_co[proc]; 17189566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coa+coi[i],len_s[proc],MPIU_MATSCALAR,proc,taga,comm,s_waits+k)); 1719d6ab1dc0SHong Zhang k++; 1720d6ab1dc0SHong Zhang } 17219566063dSJacob Faibussowitsch if (merge->nrecv) PetscCallMPI(MPI_Waitall(merge->nrecv,r_waits,status)); 17229566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,s_waits,status)); 1723d6ab1dc0SHong Zhang 17249566063dSJacob Faibussowitsch PetscCall(PetscFree2(s_waits,status)); 17259566063dSJacob Faibussowitsch PetscCall(PetscFree(r_waits)); 17269566063dSJacob Faibussowitsch PetscCall(PetscFree(coa)); 1727d6ab1dc0SHong Zhang 1728d6ab1dc0SHong Zhang /* 4) insert local Cseq and received values into Cmpi */ 1729d6ab1dc0SHong Zhang /*----------------------------------------------------*/ 17309566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(merge->nrecv,&buf_ri_k,merge->nrecv,&nextrow,merge->nrecv,&nextci)); 1731d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { 1732e745005bSHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 1733d6ab1dc0SHong Zhang nrows = *(buf_ri_k[k]); 1734d6ab1dc0SHong Zhang nextrow[k] = buf_ri_k[k]+1; /* next row number of k-th recved i-structure */ 1735a5b23f4aSJose E. Roman nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th recved i-structure */ 1736d6ab1dc0SHong Zhang } 1737d6ab1dc0SHong Zhang 1738d6ab1dc0SHong Zhang for (i=0; i<cm; i++) { 1739d6ab1dc0SHong Zhang row = owners[rank] + i; /* global row index of C_seq */ 1740d6ab1dc0SHong Zhang bj_i = bj + bi[i]; /* col indices of the i-th row of C */ 1741d6ab1dc0SHong Zhang ba_i = ba + bi[i]; 1742d6ab1dc0SHong Zhang bnz = bi[i+1] - bi[i]; 1743d6ab1dc0SHong Zhang /* add received vals into ba */ 1744d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { /* k-th received message */ 1745d6ab1dc0SHong Zhang /* i-th row */ 1746d6ab1dc0SHong Zhang if (i == *nextrow[k]) { 1747d6ab1dc0SHong Zhang cnz = *(nextci[k]+1) - *nextci[k]; 1748d6ab1dc0SHong Zhang cj = buf_rj[k] + *(nextci[k]); 1749d6ab1dc0SHong Zhang ca = abuf_r[k] + *(nextci[k]); 1750d6ab1dc0SHong Zhang nextcj = 0; 1751d6ab1dc0SHong Zhang for (j=0; nextcj<cnz; j++) { 1752d6ab1dc0SHong Zhang if (bj_i[j] == cj[nextcj]) { /* bcol == ccol */ 1753d6ab1dc0SHong Zhang ba_i[j] += ca[nextcj++]; 1754d6ab1dc0SHong Zhang } 1755d6ab1dc0SHong Zhang } 1756d6ab1dc0SHong Zhang nextrow[k]++; nextci[k]++; 17579566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*cnz)); 1758d6ab1dc0SHong Zhang } 1759d6ab1dc0SHong Zhang } 17609566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,bnz,bj_i,ba_i,INSERT_VALUES)); 1761d6ab1dc0SHong Zhang } 17629566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 17639566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 1764d6ab1dc0SHong Zhang 17659566063dSJacob Faibussowitsch PetscCall(PetscFree(ba)); 17669566063dSJacob Faibussowitsch PetscCall(PetscFree(abuf_r[0])); 17679566063dSJacob Faibussowitsch PetscCall(PetscFree(abuf_r)); 17689566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 1769d6ab1dc0SHong Zhang PetscFunctionReturn(0); 1770d6ab1dc0SHong Zhang } 1771d6ab1dc0SHong Zhang 17724222ddf1SHong Zhang PetscErrorCode MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ(Mat P,Mat A,PetscReal fill,Mat C) 1773d6ab1dc0SHong Zhang { 1774b419842dSStefano Zampini Mat A_loc; 17753cdca5ebSHong Zhang Mat_APMPI *ptap; 17760298fd71SBarry Smith PetscFreeSpaceList free_space=NULL,current_space=NULL; 17776718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data,*a=(Mat_MPIAIJ*)A->data; 1778d6ab1dc0SHong Zhang PetscInt *pdti,*pdtj,*poti,*potj,*ptJ; 1779d6ab1dc0SHong Zhang PetscInt nnz; 1780d6ab1dc0SHong Zhang PetscInt *lnk,*owners_co,*coi,*coj,i,k,pnz,row; 1781d6ab1dc0SHong Zhang PetscInt am =A->rmap->n,pn=P->cmap->n; 1782ce94432eSBarry Smith MPI_Comm comm; 1783d6ab1dc0SHong Zhang PetscMPIInt size,rank,tagi,tagj,*len_si,*len_s,*len_ri; 1784d6ab1dc0SHong Zhang PetscInt **buf_rj,**buf_ri,**buf_ri_k; 1785d6ab1dc0SHong Zhang PetscInt len,proc,*dnz,*onz,*owners; 1786d6ab1dc0SHong Zhang PetscInt nzi,*bi,*bj; 1787d6ab1dc0SHong Zhang PetscInt nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci; 1788d6ab1dc0SHong Zhang MPI_Request *swaits,*rwaits; 1789d6ab1dc0SHong Zhang MPI_Status *sstatus,rstatus; 1790d6ab1dc0SHong Zhang Mat_Merge_SeqsToMPI *merge; 1791d6ab1dc0SHong Zhang PetscInt *ai,*aj,*Jptr,anz,*prmap=p->garray,pon,nspacedouble=0,j; 1792d6ab1dc0SHong Zhang PetscReal afill =1.0,afill_tmp; 17934b38b95cSHong Zhang PetscInt rstart = P->cmap->rstart,rmax,aN=A->cmap->N,Armax; 1794b419842dSStefano Zampini Mat_SeqAIJ *a_loc; 17950acc65b4SHong Zhang PetscTable ta; 1796b92f168fSBarry Smith MatType mtype; 1797d6ab1dc0SHong Zhang 1798d6ab1dc0SHong Zhang PetscFunctionBegin; 17999566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 1800d6ab1dc0SHong Zhang /* check if matrix local sizes are compatible */ 1801aed4548fSBarry Smith PetscCheck(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); 1802d6ab1dc0SHong Zhang 18039566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 18049566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1805d6ab1dc0SHong Zhang 18063cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 18079566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 1808d6ab1dc0SHong Zhang 1809d6ab1dc0SHong Zhang /* get A_loc by taking all local rows of A */ 18109566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_INITIAL_MATRIX,&A_loc)); 18112205254eSKarl Rupp 1812d6ab1dc0SHong Zhang ptap->A_loc = A_loc; 1813d6ab1dc0SHong Zhang a_loc = (Mat_SeqAIJ*)(A_loc)->data; 1814d6ab1dc0SHong Zhang ai = a_loc->i; 1815d6ab1dc0SHong Zhang aj = a_loc->j; 1816d6ab1dc0SHong Zhang 1817d6ab1dc0SHong Zhang /* determine symbolic Co=(p->B)^T*A - send to others */ 1818d6ab1dc0SHong Zhang /*----------------------------------------------------*/ 18199566063dSJacob Faibussowitsch PetscCall(MatGetSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj)); 18209566063dSJacob Faibussowitsch PetscCall(MatGetSymbolicTranspose_SeqAIJ(p->B,&poti,&potj)); 1821d6ab1dc0SHong Zhang pon = (p->B)->cmap->n; /* total num of rows to be sent to other processors 1822d6ab1dc0SHong Zhang >= (num of nonzero rows of C_seq) - pn */ 18239566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pon+1,&coi)); 1824d6ab1dc0SHong Zhang coi[0] = 0; 1825d6ab1dc0SHong Zhang 1826d6ab1dc0SHong Zhang /* set initial free space to be fill*(nnz(p->B) + nnz(A)) */ 1827f91af8c7SBarry Smith nnz = PetscRealIntMultTruncate(fill,PetscIntSumTruncate(poti[pon],ai[am])); 18289566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(nnz,&free_space)); 1829d6ab1dc0SHong Zhang current_space = free_space; 183019f0e57aSHong Zhang 1831d6ab1dc0SHong Zhang /* create and initialize a linked list */ 18329566063dSJacob Faibussowitsch PetscCall(PetscTableCreate(A->cmap->n + a->B->cmap->N,aN,&ta)); 18334b38b95cSHong Zhang MatRowMergeMax_SeqAIJ(a_loc,am,ta); 18349566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Armax)); 183533ba5e2fSHong Zhang 18369566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(Armax,&lnk)); 1837d6ab1dc0SHong Zhang 1838d6ab1dc0SHong Zhang for (i=0; i<pon; i++) { 1839d6ab1dc0SHong Zhang pnz = poti[i+1] - poti[i]; 1840d6ab1dc0SHong Zhang ptJ = potj + poti[i]; 1841d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1842d6ab1dc0SHong Zhang row = ptJ[j]; /* row of A_loc == col of Pot */ 1843d6ab1dc0SHong Zhang anz = ai[row+1] - ai[row]; 1844d6ab1dc0SHong Zhang Jptr = aj + ai[row]; 1845d6ab1dc0SHong Zhang /* add non-zero cols of AP into the sorted linked list lnk */ 18469566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(anz,Jptr,lnk)); 1847d6ab1dc0SHong Zhang } 1848d6ab1dc0SHong Zhang nnz = lnk[0]; 1849d6ab1dc0SHong Zhang 1850d6ab1dc0SHong Zhang /* If free space is not available, double the total space in the list */ 1851d6ab1dc0SHong Zhang if (current_space->local_remaining<nnz) { 18529566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(nnz,current_space->total_array_size),¤t_space)); 1853d6ab1dc0SHong Zhang nspacedouble++; 1854d6ab1dc0SHong Zhang } 1855d6ab1dc0SHong Zhang 1856d6ab1dc0SHong Zhang /* Copy data into free space, and zero out denserows */ 18579566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(nnz,current_space->array,lnk)); 18582205254eSKarl Rupp 1859d6ab1dc0SHong Zhang current_space->array += nnz; 1860d6ab1dc0SHong Zhang current_space->local_used += nnz; 1861d6ab1dc0SHong Zhang current_space->local_remaining -= nnz; 18622205254eSKarl Rupp 1863d6ab1dc0SHong Zhang coi[i+1] = coi[i] + nnz; 1864d6ab1dc0SHong Zhang } 1865d6ab1dc0SHong Zhang 18669566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(coi[pon]+1,&coj)); 18679566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,coj)); 18689566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); /* must destroy to get a new one for C */ 18692205254eSKarl Rupp 1870118727c9SMark F. Adams afill_tmp = (PetscReal)coi[pon]/(poti[pon] + ai[am]+1); 1871d6ab1dc0SHong Zhang if (afill_tmp > afill) afill = afill_tmp; 1872d6ab1dc0SHong Zhang 1873d6ab1dc0SHong Zhang /* send j-array (coj) of Co to other processors */ 1874d6ab1dc0SHong Zhang /*----------------------------------------------*/ 1875d6ab1dc0SHong Zhang /* determine row ownership */ 18769566063dSJacob Faibussowitsch PetscCall(PetscNew(&merge)); 18779566063dSJacob Faibussowitsch PetscCall(PetscLayoutCreate(comm,&merge->rowmap)); 18782205254eSKarl Rupp 1879d6ab1dc0SHong Zhang merge->rowmap->n = pn; 1880d6ab1dc0SHong Zhang merge->rowmap->bs = 1; 18812205254eSKarl Rupp 18829566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(merge->rowmap)); 1883d6ab1dc0SHong Zhang owners = merge->rowmap->range; 1884d6ab1dc0SHong Zhang 1885d6ab1dc0SHong Zhang /* determine the number of messages to send, their lengths */ 18869566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(size,&len_si)); 18879566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(size,&merge->len_s)); 18882205254eSKarl Rupp 1889d6ab1dc0SHong Zhang len_s = merge->len_s; 1890d6ab1dc0SHong Zhang merge->nsend = 0; 1891d6ab1dc0SHong Zhang 18929566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size+2,&owners_co)); 1893d6ab1dc0SHong Zhang 1894d6ab1dc0SHong Zhang proc = 0; 1895d6ab1dc0SHong Zhang for (i=0; i<pon; i++) { 1896d6ab1dc0SHong Zhang while (prmap[i] >= owners[proc+1]) proc++; 1897d6ab1dc0SHong Zhang len_si[proc]++; /* num of rows in Co to be sent to [proc] */ 1898d6ab1dc0SHong Zhang len_s[proc] += coi[i+1] - coi[i]; 1899d6ab1dc0SHong Zhang } 1900d6ab1dc0SHong Zhang 1901d6ab1dc0SHong Zhang len = 0; /* max length of buf_si[] */ 1902d6ab1dc0SHong Zhang owners_co[0] = 0; 1903d6ab1dc0SHong Zhang for (proc=0; proc<size; proc++) { 1904d6ab1dc0SHong Zhang owners_co[proc+1] = owners_co[proc] + len_si[proc]; 1905d6ab1dc0SHong Zhang if (len_si[proc]) { 1906d6ab1dc0SHong Zhang merge->nsend++; 1907d6ab1dc0SHong Zhang len_si[proc] = 2*(len_si[proc] + 1); 1908d6ab1dc0SHong Zhang len += len_si[proc]; 1909d6ab1dc0SHong Zhang } 1910d6ab1dc0SHong Zhang } 1911d6ab1dc0SHong Zhang 1912d6ab1dc0SHong Zhang /* determine the number and length of messages to receive for coi and coj */ 19139566063dSJacob Faibussowitsch PetscCall(PetscGatherNumberOfMessages(comm,NULL,len_s,&merge->nrecv)); 19149566063dSJacob Faibussowitsch PetscCall(PetscGatherMessageLengths2(comm,merge->nsend,merge->nrecv,len_s,len_si,&merge->id_r,&merge->len_r,&len_ri)); 1915d6ab1dc0SHong Zhang 1916d6ab1dc0SHong Zhang /* post the Irecv and Isend of coj */ 19179566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagj)); 19189566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagj,merge->nrecv,merge->id_r,merge->len_r,&buf_rj,&rwaits)); 19199566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(merge->nsend+1,&swaits)); 1920d6ab1dc0SHong Zhang for (proc=0, k=0; proc<size; proc++) { 1921d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1922d6ab1dc0SHong Zhang i = owners_co[proc]; 19239566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tagj,comm,swaits+k)); 1924d6ab1dc0SHong Zhang k++; 1925d6ab1dc0SHong Zhang } 1926d6ab1dc0SHong Zhang 1927d6ab1dc0SHong Zhang /* receives and sends of coj are complete */ 19289566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size,&sstatus)); 1929d6ab1dc0SHong Zhang for (i=0; i<merge->nrecv; i++) { 1930ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted; 19319566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(merge->nrecv,rwaits,&icompleted,&rstatus)); 1932d6ab1dc0SHong Zhang } 19339566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 19349566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,swaits,sstatus)); 1935d6ab1dc0SHong Zhang 19360acc65b4SHong Zhang /* add received column indices into table to update Armax */ 1937c4762a1bSJed Brown /* Armax can be as large as aN if a P[row,:] is dense, see src/ksp/ksp/tutorials/ex56.c! */ 19380acc65b4SHong Zhang for (k=0; k<merge->nrecv; k++) {/* k-th received message */ 19390acc65b4SHong Zhang Jptr = buf_rj[k]; 19400acc65b4SHong Zhang for (j=0; j<merge->len_r[k]; j++) { 19419566063dSJacob Faibussowitsch PetscCall(PetscTableAdd(ta,*(Jptr+j)+1,1,INSERT_VALUES)); 19420acc65b4SHong Zhang } 19430acc65b4SHong Zhang } 19449566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Armax)); 19450acc65b4SHong Zhang 1946d6ab1dc0SHong Zhang /* send and recv coi */ 1947d6ab1dc0SHong Zhang /*-------------------*/ 19489566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagi)); 19499566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagi,merge->nrecv,merge->id_r,len_ri,&buf_ri,&rwaits)); 19509566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(len+1,&buf_s)); 1951d6ab1dc0SHong Zhang buf_si = buf_s; /* points to the beginning of k-th msg to be sent */ 1952d6ab1dc0SHong Zhang for (proc=0,k=0; proc<size; proc++) { 1953d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1954d6ab1dc0SHong Zhang /* form outgoing message for i-structure: 1955d6ab1dc0SHong Zhang buf_si[0]: nrows to be sent 1956d6ab1dc0SHong Zhang [1:nrows]: row index (global) 1957d6ab1dc0SHong Zhang [nrows+1:2*nrows+1]: i-structure index 1958d6ab1dc0SHong Zhang */ 1959d6ab1dc0SHong Zhang /*-------------------------------------------*/ 1960d6ab1dc0SHong Zhang nrows = len_si[proc]/2 - 1; 1961d6ab1dc0SHong Zhang buf_si_i = buf_si + nrows+1; 1962d6ab1dc0SHong Zhang buf_si[0] = nrows; 1963d6ab1dc0SHong Zhang buf_si_i[0] = 0; 1964d6ab1dc0SHong Zhang nrows = 0; 1965d6ab1dc0SHong Zhang for (i=owners_co[proc]; i<owners_co[proc+1]; i++) { 1966d6ab1dc0SHong Zhang nzi = coi[i+1] - coi[i]; 1967d6ab1dc0SHong Zhang buf_si_i[nrows+1] = buf_si_i[nrows] + nzi; /* i-structure */ 1968d6ab1dc0SHong Zhang buf_si[nrows+1] = prmap[i] -owners[proc]; /* local row index */ 1969d6ab1dc0SHong Zhang nrows++; 1970d6ab1dc0SHong Zhang } 19719566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tagi,comm,swaits+k)); 1972d6ab1dc0SHong Zhang k++; 1973d6ab1dc0SHong Zhang buf_si += len_si[proc]; 1974d6ab1dc0SHong Zhang } 1975d6ab1dc0SHong Zhang i = merge->nrecv; 1976d6ab1dc0SHong Zhang while (i--) { 1977ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted; 19789566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(merge->nrecv,rwaits,&icompleted,&rstatus)); 1979d6ab1dc0SHong Zhang } 19809566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 19819566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,swaits,sstatus)); 19829566063dSJacob Faibussowitsch PetscCall(PetscFree(len_si)); 19839566063dSJacob Faibussowitsch PetscCall(PetscFree(len_ri)); 19849566063dSJacob Faibussowitsch PetscCall(PetscFree(swaits)); 19859566063dSJacob Faibussowitsch PetscCall(PetscFree(sstatus)); 19869566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_s)); 1987d6ab1dc0SHong Zhang 1988d6ab1dc0SHong Zhang /* compute the local portion of C (mpi mat) */ 1989d6ab1dc0SHong Zhang /*------------------------------------------*/ 1990d6ab1dc0SHong Zhang /* allocate bi array and free space for accumulating nonzero column info */ 19919566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pn+1,&bi)); 1992d6ab1dc0SHong Zhang bi[0] = 0; 1993d6ab1dc0SHong Zhang 1994d6ab1dc0SHong Zhang /* set initial free space to be fill*(nnz(P) + nnz(AP)) */ 1995f91af8c7SBarry Smith nnz = PetscRealIntMultTruncate(fill,PetscIntSumTruncate(pdti[pn],PetscIntSumTruncate(poti[pon],ai[am]))); 19969566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(nnz,&free_space)); 1997d6ab1dc0SHong Zhang current_space = free_space; 1998d6ab1dc0SHong Zhang 19999566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(merge->nrecv,&buf_ri_k,merge->nrecv,&nextrow,merge->nrecv,&nextci)); 2000d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { 2001d6ab1dc0SHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 2002d6ab1dc0SHong Zhang nrows = *buf_ri_k[k]; 2003d6ab1dc0SHong Zhang nextrow[k] = buf_ri_k[k] + 1; /* next row number of k-th recved i-structure */ 20042d4ee042Sprj- nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th received i-structure */ 2005d6ab1dc0SHong Zhang } 2006d6ab1dc0SHong Zhang 20079566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(Armax,&lnk)); 2008d0609cedSBarry Smith MatPreallocateBegin(comm,pn,A->cmap->n,dnz,onz); 2009d6ab1dc0SHong Zhang rmax = 0; 2010d6ab1dc0SHong Zhang for (i=0; i<pn; i++) { 2011d6ab1dc0SHong Zhang /* add pdt[i,:]*AP into lnk */ 2012d6ab1dc0SHong Zhang pnz = pdti[i+1] - pdti[i]; 2013d6ab1dc0SHong Zhang ptJ = pdtj + pdti[i]; 2014d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 2015d6ab1dc0SHong Zhang row = ptJ[j]; /* row of AP == col of Pt */ 2016d6ab1dc0SHong Zhang anz = ai[row+1] - ai[row]; 2017d6ab1dc0SHong Zhang Jptr = aj + ai[row]; 2018d6ab1dc0SHong Zhang /* add non-zero cols of AP into the sorted linked list lnk */ 20199566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(anz,Jptr,lnk)); 2020d6ab1dc0SHong Zhang } 2021d6ab1dc0SHong Zhang 2022d6ab1dc0SHong Zhang /* add received col data into lnk */ 2023d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { /* k-th received message */ 2024d6ab1dc0SHong Zhang if (i == *nextrow[k]) { /* i-th row */ 2025d6ab1dc0SHong Zhang nzi = *(nextci[k]+1) - *nextci[k]; 2026d6ab1dc0SHong Zhang Jptr = buf_rj[k] + *nextci[k]; 20279566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(nzi,Jptr,lnk)); 2028d6ab1dc0SHong Zhang nextrow[k]++; nextci[k]++; 2029d6ab1dc0SHong Zhang } 2030d6ab1dc0SHong Zhang } 20318c78258cSHong Zhang 20328c78258cSHong Zhang /* add missing diagonal entry */ 20338c78258cSHong Zhang if (C->force_diagonals) { 20348c78258cSHong Zhang k = i + owners[rank]; /* column index */ 20359566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(1,&k,lnk)); 20368c78258cSHong Zhang } 20378c78258cSHong Zhang 2038d6ab1dc0SHong Zhang nnz = lnk[0]; 2039d6ab1dc0SHong Zhang 2040d6ab1dc0SHong Zhang /* if free space is not available, make more free space */ 2041d6ab1dc0SHong Zhang if (current_space->local_remaining<nnz) { 20429566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(nnz,current_space->total_array_size),¤t_space)); 2043d6ab1dc0SHong Zhang nspacedouble++; 2044d6ab1dc0SHong Zhang } 2045d6ab1dc0SHong Zhang /* copy data into free space, then initialize lnk */ 20469566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(nnz,current_space->array,lnk)); 20479566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+owners[rank],nnz,current_space->array,dnz,onz)); 20482205254eSKarl Rupp 2049d6ab1dc0SHong Zhang current_space->array += nnz; 2050d6ab1dc0SHong Zhang current_space->local_used += nnz; 2051d6ab1dc0SHong Zhang current_space->local_remaining -= nnz; 20522205254eSKarl Rupp 2053d6ab1dc0SHong Zhang bi[i+1] = bi[i] + nnz; 2054d6ab1dc0SHong Zhang if (nnz > rmax) rmax = nnz; 2055d6ab1dc0SHong Zhang } 20569566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 2057d6ab1dc0SHong Zhang 20589566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(bi[pn]+1,&bj)); 20599566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,bj)); 2060118727c9SMark F. Adams afill_tmp = (PetscReal)bi[pn]/(pdti[pn] + poti[pon] + ai[am]+1); 2061d6ab1dc0SHong Zhang if (afill_tmp > afill) afill = afill_tmp; 20629566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); 20639566063dSJacob Faibussowitsch PetscCall(PetscTableDestroy(&ta)); 20649566063dSJacob Faibussowitsch PetscCall(MatRestoreSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj)); 20659566063dSJacob Faibussowitsch PetscCall(MatRestoreSymbolicTranspose_SeqAIJ(p->B,&poti,&potj)); 2066d6ab1dc0SHong Zhang 20674222ddf1SHong Zhang /* create symbolic parallel matrix C - why cannot be assembled in Numeric part */ 20684222ddf1SHong Zhang /*-------------------------------------------------------------------------------*/ 20699566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,pn,A->cmap->n,PETSC_DETERMINE,PETSC_DETERMINE)); 20709566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizes(C,PetscAbs(P->cmap->bs),PetscAbs(A->cmap->bs))); 20719566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 20729566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 20739566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 2074d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 20759566063dSJacob Faibussowitsch PetscCall(MatSetBlockSize(C,1)); 20769566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NO_OFF_PROC_ENTRIES,PETSC_TRUE)); 2077d6ab1dc0SHong Zhang for (i=0; i<pn; i++) { 2078d6ab1dc0SHong Zhang row = i + rstart; 2079d6ab1dc0SHong Zhang nnz = bi[i+1] - bi[i]; 2080d6ab1dc0SHong Zhang Jptr = bj + bi[i]; 20819566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,nnz,Jptr,NULL,INSERT_VALUES)); 2082d6ab1dc0SHong Zhang } 20839566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 20849566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 20859566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 2086d6ab1dc0SHong Zhang merge->bi = bi; 2087d6ab1dc0SHong Zhang merge->bj = bj; 2088d6ab1dc0SHong Zhang merge->coi = coi; 2089d6ab1dc0SHong Zhang merge->coj = coj; 2090d6ab1dc0SHong Zhang merge->buf_ri = buf_ri; 2091d6ab1dc0SHong Zhang merge->buf_rj = buf_rj; 2092d6ab1dc0SHong Zhang merge->owners_co = owners_co; 2093d6ab1dc0SHong Zhang 20944222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 20956718818eSStefano Zampini C->product->data = ptap; 20966718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 2097d6ab1dc0SHong Zhang ptap->merge = merge; 2098a560ef98SHong Zhang 20994222ddf1SHong Zhang C->ops->mattransposemultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ; 2100f96d37fcSHong Zhang 2101f96d37fcSHong Zhang #if defined(PETSC_USE_INFO) 2102f96d37fcSHong Zhang if (bi[pn] != 0) { 21039566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 21049566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatTransposeMatMult(A,B,MatReuse,%g,&C) for best performance.\n",(double)afill)); 2105f96d37fcSHong Zhang } else { 21069566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 2107f96d37fcSHong Zhang } 2108f96d37fcSHong Zhang #endif 2109f96d37fcSHong Zhang PetscFunctionReturn(0); 2110f96d37fcSHong Zhang } 21114222ddf1SHong Zhang 21124222ddf1SHong Zhang /* ---------------------------------------------------------------- */ 21134222ddf1SHong Zhang static PetscErrorCode MatProductSymbolic_AtB_MPIAIJ_MPIAIJ(Mat C) 21144222ddf1SHong Zhang { 21154222ddf1SHong Zhang Mat_Product *product = C->product; 21164222ddf1SHong Zhang Mat A=product->A,B=product->B; 21174222ddf1SHong Zhang PetscReal fill=product->fill; 21184222ddf1SHong Zhang PetscBool flg; 21194222ddf1SHong Zhang 21204222ddf1SHong Zhang PetscFunctionBegin; 21214222ddf1SHong Zhang /* scalable */ 21229566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"scalable",&flg)); 21234222ddf1SHong Zhang if (flg) { 21249566063dSJacob Faibussowitsch PetscCall(MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ(A,B,fill,C)); 21254222ddf1SHong Zhang goto next; 21264222ddf1SHong Zhang } 21274222ddf1SHong Zhang 21284222ddf1SHong Zhang /* nonscalable */ 21299566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"nonscalable",&flg)); 21304222ddf1SHong Zhang if (flg) { 21319566063dSJacob Faibussowitsch PetscCall(MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(A,B,fill,C)); 21324222ddf1SHong Zhang goto next; 21334222ddf1SHong Zhang } 21344222ddf1SHong Zhang 21354222ddf1SHong Zhang /* matmatmult */ 21369566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"at*b",&flg)); 21374222ddf1SHong Zhang if (flg) { 21384222ddf1SHong Zhang Mat At; 21394222ddf1SHong Zhang Mat_APMPI *ptap; 21404222ddf1SHong Zhang 21419566063dSJacob Faibussowitsch PetscCall(MatTranspose(A,MAT_INITIAL_MATRIX,&At)); 21429566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ(At,B,fill,C)); 21436718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 21444222ddf1SHong Zhang if (ptap) { 21454222ddf1SHong Zhang ptap->Pt = At; 21466718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 21474222ddf1SHong Zhang } 21484222ddf1SHong Zhang C->ops->transposematmultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_matmatmult; 21494222ddf1SHong Zhang goto next; 21504222ddf1SHong Zhang } 21514222ddf1SHong Zhang 21524e84afc0SStefano Zampini /* backend general code */ 21539566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"backend",&flg)); 21544e84afc0SStefano Zampini if (flg) { 21559566063dSJacob Faibussowitsch PetscCall(MatProductSymbolic_MPIAIJBACKEND(C)); 21564e84afc0SStefano Zampini PetscFunctionReturn(0); 21574e84afc0SStefano Zampini } 21584e84afc0SStefano Zampini 21594222ddf1SHong Zhang SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"MatProduct type is not supported"); 21604222ddf1SHong Zhang 21614222ddf1SHong Zhang next: 21624222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AtB; 21634222ddf1SHong Zhang PetscFunctionReturn(0); 21644222ddf1SHong Zhang } 21654222ddf1SHong Zhang 21664222ddf1SHong Zhang /* ---------------------------------------------------------------- */ 21674222ddf1SHong Zhang /* Set options for MatMatMultxxx_MPIAIJ_MPIAIJ */ 21684222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_AB(Mat C) 21694222ddf1SHong Zhang { 21704222ddf1SHong Zhang Mat_Product *product = C->product; 21714222ddf1SHong Zhang Mat A=product->A,B=product->B; 21724222ddf1SHong Zhang #if defined(PETSC_HAVE_HYPRE) 21734e84afc0SStefano Zampini const char *algTypes[5] = {"scalable","nonscalable","seqmpi","backend","hypre"}; 21744e84afc0SStefano Zampini PetscInt nalg = 5; 21754222ddf1SHong Zhang #else 21764e84afc0SStefano Zampini const char *algTypes[4] = {"scalable","nonscalable","seqmpi","backend",}; 21774e84afc0SStefano Zampini PetscInt nalg = 4; 21784222ddf1SHong Zhang #endif 21794222ddf1SHong Zhang PetscInt alg = 1; /* set nonscalable algorithm as default */ 21804222ddf1SHong Zhang PetscBool flg; 21814222ddf1SHong Zhang MPI_Comm comm; 21824222ddf1SHong Zhang 21834222ddf1SHong Zhang PetscFunctionBegin; 21849566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 21854222ddf1SHong Zhang 21864222ddf1SHong Zhang /* Set "nonscalable" as default algorithm */ 21879566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 21884222ddf1SHong Zhang if (flg) { 21899566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 21904222ddf1SHong Zhang 21914222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 21924222ddf1SHong Zhang if (B->cmap->N > 100000) { /* may switch to scalable algorithm as default */ 21934222ddf1SHong Zhang MatInfo Ainfo,Binfo; 21944222ddf1SHong Zhang PetscInt nz_local; 21954222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 21964222ddf1SHong Zhang 21979566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 21989566063dSJacob Faibussowitsch PetscCall(MatGetInfo(B,MAT_LOCAL,&Binfo)); 21994222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Binfo.nz_allocated); 22004222ddf1SHong Zhang 22014222ddf1SHong Zhang if (B->cmap->N > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 22021c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 22034222ddf1SHong Zhang 22044222ddf1SHong Zhang if (alg_scalable) { 22054222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 22069566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22079566063dSJacob Faibussowitsch PetscCall(PetscInfo(B,"Use scalable algorithm, BN %" PetscInt_FMT ", fill*nz_allocated %g\n",B->cmap->N,(double)(product->fill*nz_local))); 22084222ddf1SHong Zhang } 22094222ddf1SHong Zhang } 22104222ddf1SHong Zhang } 22114222ddf1SHong Zhang 22124222ddf1SHong Zhang /* Get runtime option */ 22134222ddf1SHong Zhang if (product->api_user) { 2214d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatMatMult","Mat"); 22159566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matmatmult_via","Algorithmic approach","MatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2216d0609cedSBarry Smith PetscOptionsEnd(); 22174222ddf1SHong Zhang } else { 2218d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_AB","Mat"); 22199566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2220d0609cedSBarry Smith PetscOptionsEnd(); 22214222ddf1SHong Zhang } 22224222ddf1SHong Zhang if (flg) { 22239566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22244222ddf1SHong Zhang } 22254222ddf1SHong Zhang 22264222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AB_MPIAIJ_MPIAIJ; 22274222ddf1SHong Zhang PetscFunctionReturn(0); 22284222ddf1SHong Zhang } 22294222ddf1SHong Zhang 22308a9c020eSBarry Smith static PetscErrorCode MatProductSetFromOptions_MPIAIJ_ABt(Mat C) 22318a9c020eSBarry Smith { 22328a9c020eSBarry Smith PetscFunctionBegin; 22338a9c020eSBarry Smith PetscCall(MatProductSetFromOptions_MPIAIJ_AB(C)); 22348a9c020eSBarry Smith C->ops->productsymbolic = MatProductSymbolic_ABt_MPIAIJ_MPIAIJ; 22358a9c020eSBarry Smith PetscFunctionReturn(0); 22368a9c020eSBarry Smith } 22378a9c020eSBarry Smith 22384222ddf1SHong Zhang /* Set options for MatTransposeMatMultXXX_MPIAIJ_MPIAIJ */ 22394222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_AtB(Mat C) 22404222ddf1SHong Zhang { 22414222ddf1SHong Zhang Mat_Product *product = C->product; 22424222ddf1SHong Zhang Mat A=product->A,B=product->B; 22434e84afc0SStefano Zampini const char *algTypes[4] = {"scalable","nonscalable","at*b","backend"}; 22444e84afc0SStefano Zampini PetscInt nalg = 4; 22454222ddf1SHong Zhang PetscInt alg = 1; /* set default algorithm */ 22464222ddf1SHong Zhang PetscBool flg; 22474222ddf1SHong Zhang MPI_Comm comm; 22484222ddf1SHong Zhang 22494222ddf1SHong Zhang PetscFunctionBegin; 22504222ddf1SHong Zhang /* Check matrix local sizes */ 22519566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 2252aed4548fSBarry Smith PetscCheck(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); 22534222ddf1SHong Zhang 22544222ddf1SHong Zhang /* Set default algorithm */ 22559566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 22564222ddf1SHong Zhang if (flg) { 22579566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22584222ddf1SHong Zhang } 22594222ddf1SHong Zhang 22604222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 22614222ddf1SHong Zhang if (alg && B->cmap->N > 100000) { /* may switch to scalable algorithm as default */ 22624222ddf1SHong Zhang MatInfo Ainfo,Binfo; 22634222ddf1SHong Zhang PetscInt nz_local; 22644222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 22654222ddf1SHong Zhang 22669566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 22679566063dSJacob Faibussowitsch PetscCall(MatGetInfo(B,MAT_LOCAL,&Binfo)); 22684222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Binfo.nz_allocated); 22694222ddf1SHong Zhang 22704222ddf1SHong Zhang if (B->cmap->N > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 22711c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 22724222ddf1SHong Zhang 22734222ddf1SHong Zhang if (alg_scalable) { 22744222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 22759566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22769566063dSJacob Faibussowitsch PetscCall(PetscInfo(B,"Use scalable algorithm, BN %" PetscInt_FMT ", fill*nz_allocated %g\n",B->cmap->N,(double)(product->fill*nz_local))); 22774222ddf1SHong Zhang } 22784222ddf1SHong Zhang } 22794222ddf1SHong Zhang 22804222ddf1SHong Zhang /* Get runtime option */ 22814222ddf1SHong Zhang if (product->api_user) { 2282d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatTransposeMatMult","Mat"); 22839566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mattransposematmult_via","Algorithmic approach","MatTransposeMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2284d0609cedSBarry Smith PetscOptionsEnd(); 22854222ddf1SHong Zhang } else { 2286d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_AtB","Mat"); 22879566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatTransposeMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2288d0609cedSBarry Smith PetscOptionsEnd(); 22894222ddf1SHong Zhang } 22904222ddf1SHong Zhang if (flg) { 22919566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22924222ddf1SHong Zhang } 22934222ddf1SHong Zhang 22944222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AtB_MPIAIJ_MPIAIJ; 22954222ddf1SHong Zhang PetscFunctionReturn(0); 22964222ddf1SHong Zhang } 22974222ddf1SHong Zhang 22984222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_PtAP(Mat C) 22994222ddf1SHong Zhang { 23004222ddf1SHong Zhang Mat_Product *product = C->product; 23014222ddf1SHong Zhang Mat A=product->A,P=product->B; 23024222ddf1SHong Zhang MPI_Comm comm; 23034222ddf1SHong Zhang PetscBool flg; 23044222ddf1SHong Zhang PetscInt alg=1; /* set default algorithm */ 23054222ddf1SHong Zhang #if !defined(PETSC_HAVE_HYPRE) 23064e84afc0SStefano Zampini const char *algTypes[5] = {"scalable","nonscalable","allatonce","allatonce_merged","backend"}; 23074222ddf1SHong Zhang PetscInt nalg=5; 23084e84afc0SStefano Zampini #else 23094e84afc0SStefano Zampini const char *algTypes[6] = {"scalable","nonscalable","allatonce","allatonce_merged","backend","hypre"}; 23104e84afc0SStefano Zampini PetscInt nalg=6; 23114222ddf1SHong Zhang #endif 23124222ddf1SHong Zhang PetscInt pN=P->cmap->N; 23134222ddf1SHong Zhang 23144222ddf1SHong Zhang PetscFunctionBegin; 23154222ddf1SHong Zhang /* Check matrix local sizes */ 23169566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 2317aed4548fSBarry Smith PetscCheck(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); 2318aed4548fSBarry Smith PetscCheck(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); 23194222ddf1SHong Zhang 23204222ddf1SHong Zhang /* Set "nonscalable" as default algorithm */ 23219566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 23224222ddf1SHong Zhang if (flg) { 23239566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23244222ddf1SHong Zhang 23254222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 23264222ddf1SHong Zhang if (pN > 100000) { 23274222ddf1SHong Zhang MatInfo Ainfo,Pinfo; 23284222ddf1SHong Zhang PetscInt nz_local; 23294222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 23304222ddf1SHong Zhang 23319566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 23329566063dSJacob Faibussowitsch PetscCall(MatGetInfo(P,MAT_LOCAL,&Pinfo)); 23334222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Pinfo.nz_allocated); 23344222ddf1SHong Zhang 23354222ddf1SHong Zhang if (pN > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 23361c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 23374222ddf1SHong Zhang 23384222ddf1SHong Zhang if (alg_scalable) { 23394222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 23409566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23414222ddf1SHong Zhang } 23424222ddf1SHong Zhang } 23434222ddf1SHong Zhang } 23444222ddf1SHong Zhang 23454222ddf1SHong Zhang /* Get runtime option */ 23464222ddf1SHong Zhang if (product->api_user) { 2347d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatPtAP","Mat"); 23489566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matptap_via","Algorithmic approach","MatPtAP",algTypes,nalg,algTypes[alg],&alg,&flg)); 2349d0609cedSBarry Smith PetscOptionsEnd(); 23504222ddf1SHong Zhang } else { 2351d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_PtAP","Mat"); 23529566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatPtAP",algTypes,nalg,algTypes[alg],&alg,&flg)); 2353d0609cedSBarry Smith PetscOptionsEnd(); 23544222ddf1SHong Zhang } 23554222ddf1SHong Zhang if (flg) { 23569566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23574222ddf1SHong Zhang } 23584222ddf1SHong Zhang 23594222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_PtAP_MPIAIJ_MPIAIJ; 23604222ddf1SHong Zhang PetscFunctionReturn(0); 23614222ddf1SHong Zhang } 23624222ddf1SHong Zhang 23634222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_RARt(Mat C) 23644222ddf1SHong Zhang { 23654222ddf1SHong Zhang Mat_Product *product = C->product; 23664222ddf1SHong Zhang Mat A = product->A,R=product->B; 23674222ddf1SHong Zhang 23684222ddf1SHong Zhang PetscFunctionBegin; 23694222ddf1SHong Zhang /* Check matrix local sizes */ 2370aed4548fSBarry Smith PetscCheck(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); 23714222ddf1SHong Zhang 23724222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_RARt_MPIAIJ_MPIAIJ; 23734222ddf1SHong Zhang PetscFunctionReturn(0); 23744222ddf1SHong Zhang } 23754222ddf1SHong Zhang 23764222ddf1SHong Zhang /* 23774222ddf1SHong Zhang Set options for ABC = A*B*C = A*(B*C); ABC's algorithm must be chosen from AB's algorithm 23784222ddf1SHong Zhang */ 23794222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_ABC(Mat C) 23804222ddf1SHong Zhang { 23814222ddf1SHong Zhang Mat_Product *product = C->product; 23824222ddf1SHong Zhang PetscBool flg = PETSC_FALSE; 23834222ddf1SHong Zhang PetscInt alg = 1; /* default algorithm */ 23844222ddf1SHong Zhang const char *algTypes[3] = {"scalable","nonscalable","seqmpi"}; 23854222ddf1SHong Zhang PetscInt nalg = 3; 23864222ddf1SHong Zhang 23874222ddf1SHong Zhang PetscFunctionBegin; 23884222ddf1SHong Zhang /* Set default algorithm */ 23899566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 23904222ddf1SHong Zhang if (flg) { 23919566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23924222ddf1SHong Zhang } 23934222ddf1SHong Zhang 23944222ddf1SHong Zhang /* Get runtime option */ 23954222ddf1SHong Zhang if (product->api_user) { 2396d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatMatMatMult","Mat"); 23979566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matmatmatmult_via","Algorithmic approach","MatMatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2398d0609cedSBarry Smith PetscOptionsEnd(); 23994222ddf1SHong Zhang } else { 2400d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_ABC","Mat"); 24019566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatProduct_ABC",algTypes,nalg,algTypes[alg],&alg,&flg)); 2402d0609cedSBarry Smith PetscOptionsEnd(); 24034222ddf1SHong Zhang } 24044222ddf1SHong Zhang if (flg) { 24059566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 24064222ddf1SHong Zhang } 24074222ddf1SHong Zhang 24084222ddf1SHong Zhang C->ops->matmatmultsymbolic = MatMatMatMultSymbolic_MPIAIJ_MPIAIJ_MPIAIJ; 24094222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_ABC; 24104222ddf1SHong Zhang PetscFunctionReturn(0); 24114222ddf1SHong Zhang } 24124222ddf1SHong Zhang 24134222ddf1SHong Zhang PETSC_INTERN PetscErrorCode MatProductSetFromOptions_MPIAIJ(Mat C) 24144222ddf1SHong Zhang { 24154222ddf1SHong Zhang Mat_Product *product = C->product; 24164222ddf1SHong Zhang 24174222ddf1SHong Zhang PetscFunctionBegin; 24184222ddf1SHong Zhang switch (product->type) { 24194222ddf1SHong Zhang case MATPRODUCT_AB: 24209566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_AB(C)); 24214222ddf1SHong Zhang break; 24228a9c020eSBarry Smith case MATPRODUCT_ABt: 24238a9c020eSBarry Smith PetscCall(MatProductSetFromOptions_MPIAIJ_ABt(C)); 24248a9c020eSBarry Smith break; 24254222ddf1SHong Zhang case MATPRODUCT_AtB: 24269566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_AtB(C)); 24274222ddf1SHong Zhang break; 24284222ddf1SHong Zhang case MATPRODUCT_PtAP: 24299566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_PtAP(C)); 24304222ddf1SHong Zhang break; 24314222ddf1SHong Zhang case MATPRODUCT_RARt: 24329566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_RARt(C)); 24334222ddf1SHong Zhang break; 24344222ddf1SHong Zhang case MATPRODUCT_ABC: 24359566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_ABC(C)); 24364222ddf1SHong Zhang break; 24376718818eSStefano Zampini default: 24386718818eSStefano Zampini break; 24394222ddf1SHong Zhang } 24404222ddf1SHong Zhang PetscFunctionReturn(0); 24414222ddf1SHong Zhang } 2442