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; 1255*1cdffd5eSHong Zhang PetscCall(MatTransposeSetPrecursor(P,Pt)); 12569566063dSJacob Faibussowitsch PetscCall(MatTranspose(P,MAT_REUSE_MATRIX,&Pt)); 12579566063dSJacob Faibussowitsch PetscCall(MatMatMultNumeric_MPIAIJ_MPIAIJ(Pt,A,C)); 1258f96d37fcSHong Zhang PetscFunctionReturn(0); 1259f96d37fcSHong Zhang } 1260f96d37fcSHong Zhang 1261c5af039cSHong Zhang /* This routine is modified from MatPtAPSymbolic_MPIAIJ_MPIAIJ() */ 12624222ddf1SHong Zhang PetscErrorCode MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(Mat P,Mat A,PetscReal fill,Mat C) 1263912837bbSHong Zhang { 12643cdca5ebSHong Zhang Mat_APMPI *ptap; 12656718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data; 1266912837bbSHong Zhang MPI_Comm comm; 1267912837bbSHong Zhang PetscMPIInt size,rank; 1268912837bbSHong Zhang PetscFreeSpaceList free_space=NULL,current_space=NULL; 12694ede91e3SHong Zhang PetscInt pn=P->cmap->n,aN=A->cmap->N,an=A->cmap->n; 127069be949aSStefano Zampini PetscInt *lnk,i,k,nsend,rstart; 1271912837bbSHong Zhang PetscBT lnkbt; 1272ec4bef21SJose E. Roman PetscMPIInt tagi,tagj,*len_si,*len_s,*len_ri,nrecv; 1273ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted=0; 12748c78258cSHong Zhang PetscInt **buf_rj,**buf_ri,**buf_ri_k,row,ncols,*cols; 12757f3e1f30SHong Zhang PetscInt len,proc,*dnz,*onz,*owners,nzi; 1276912837bbSHong Zhang PetscInt nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci; 1277912837bbSHong Zhang MPI_Request *swaits,*rwaits; 1278912837bbSHong Zhang MPI_Status *sstatus,rstatus; 1279912837bbSHong Zhang PetscLayout rowmap; 1280912837bbSHong Zhang PetscInt *owners_co,*coi,*coj; /* i and j array of (p->B)^T*A*P - used in the communication */ 1281912837bbSHong Zhang PetscMPIInt *len_r,*id_r; /* array of length of comm->size, store send/recv matrix values */ 12827f3e1f30SHong Zhang PetscInt *Jptr,*prmap=p->garray,con,j,Crmax; 12834ede91e3SHong Zhang Mat_SeqAIJ *a_loc,*c_loc,*c_oth; 1284912837bbSHong Zhang PetscTable ta; 1285b92f168fSBarry Smith MatType mtype; 1286e83e5f86SFande Kong const char *prefix; 1287912837bbSHong Zhang 1288912837bbSHong Zhang PetscFunctionBegin; 12899566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 12909566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 12919566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1292912837bbSHong Zhang 12934222ddf1SHong Zhang /* create symbolic parallel matrix C */ 12949566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 12959566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 1296912837bbSHong Zhang 12974222ddf1SHong Zhang C->ops->transposematmultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable; 1298912837bbSHong Zhang 12993cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 13009566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 1301912837bbSHong Zhang 1302912837bbSHong Zhang /* (0) compute Rd = Pd^T, Ro = Po^T */ 1303912837bbSHong Zhang /* --------------------------------- */ 1304acd337a6SBarry Smith PetscCall(MatTranspose(p->A,MAT_INITIAL_MATRIX,&ptap->Rd)); 1305acd337a6SBarry 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 1550669e9c86SHong Zhang /* These matrices are obtained in MatTransposeMatMultSymbolic() */ 1551669e9c86SHong Zhang /* 1) get R = Pd^T, Ro = Po^T */ 1552669e9c86SHong Zhang /*----------------------------*/ 1553*1cdffd5eSHong Zhang PetscCall(MatTransposeSetPrecursor(p->A,ptap->Rd)); 1554acd337a6SBarry Smith PetscCall(MatTranspose(p->A,MAT_REUSE_MATRIX,&ptap->Rd)); 1555*1cdffd5eSHong Zhang PetscCall(MatTransposeSetPrecursor(p->B,ptap->Ro)); 1556acd337a6SBarry Smith PetscCall(MatTranspose(p->B,MAT_REUSE_MATRIX,&ptap->Ro)); 1557669e9c86SHong Zhang 1558669e9c86SHong Zhang /* 2) compute numeric A_loc */ 1559669e9c86SHong Zhang /*--------------------------*/ 15609566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_REUSE_MATRIX,&ptap->A_loc)); 1561912837bbSHong Zhang 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 PetscFunctionReturn(0); 1599912837bbSHong Zhang } 1600912837bbSHong Zhang 16016da69ca6SHong Zhang PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ(Mat P,Mat A,Mat C) 1602d6ab1dc0SHong Zhang { 1603d6ab1dc0SHong Zhang Mat_Merge_SeqsToMPI *merge; 16046718818eSStefano Zampini Mat_MPIAIJ *p =(Mat_MPIAIJ*)P->data; 1605d6ab1dc0SHong Zhang Mat_SeqAIJ *pd=(Mat_SeqAIJ*)(p->A)->data,*po=(Mat_SeqAIJ*)(p->B)->data; 1606*1cdffd5eSHong Zhang Mat_APMPI *ap; 1607e745005bSHong Zhang PetscInt *adj; 1608e745005bSHong Zhang PetscInt i,j,k,anz,pnz,row,*cj,nexta; 1609e745005bSHong Zhang MatScalar *ada,*ca,valtmp; 1610d6ab1dc0SHong Zhang PetscInt am=A->rmap->n,cm=C->rmap->n,pon=(p->B)->cmap->n; 1611ce94432eSBarry Smith MPI_Comm comm; 1612d6ab1dc0SHong Zhang PetscMPIInt size,rank,taga,*len_s; 1613d6ab1dc0SHong Zhang PetscInt *owners,proc,nrows,**buf_ri_k,**nextrow,**nextci; 1614d6ab1dc0SHong Zhang PetscInt **buf_ri,**buf_rj; 1615d6ab1dc0SHong Zhang PetscInt cnz=0,*bj_i,*bi,*bj,bnz,nextcj; /* bi,bj,ba: local array of C(mpi mat) */ 1616d6ab1dc0SHong Zhang MPI_Request *s_waits,*r_waits; 1617d6ab1dc0SHong Zhang MPI_Status *status; 1618d6ab1dc0SHong Zhang MatScalar **abuf_r,*ba_i,*pA,*coa,*ba; 1619ce496241SStefano Zampini const PetscScalar *dummy; 16204e201bd7SHong Zhang PetscInt *ai,*aj,*coi,*coj,*poJ,*pdJ; 1621d6ab1dc0SHong Zhang Mat A_loc; 1622d6ab1dc0SHong Zhang Mat_SeqAIJ *a_loc; 1623d6ab1dc0SHong Zhang 1624d6ab1dc0SHong Zhang PetscFunctionBegin; 16256718818eSStefano Zampini MatCheckProduct(C,3); 1626*1cdffd5eSHong Zhang ap = (Mat_APMPI*)C->product->data; 1627*1cdffd5eSHong Zhang PetscCheck(ap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be computed. Missing data"); 1628*1cdffd5eSHong Zhang PetscCheck(ap->A_loc,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be reused. Do not call MatProductClear()"); 16299566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 16309566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 16319566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1632d6ab1dc0SHong Zhang 1633*1cdffd5eSHong Zhang merge = ap->merge; 1634d6ab1dc0SHong Zhang 1635e745005bSHong Zhang /* 2) compute numeric C_seq = P_loc^T*A_loc */ 1636e745005bSHong Zhang /*------------------------------------------*/ 1637d6ab1dc0SHong Zhang /* get data from symbolic products */ 1638d6ab1dc0SHong Zhang coi = merge->coi; coj = merge->coj; 16399566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(coi[pon]+1,&coa)); 1640d6ab1dc0SHong Zhang bi = merge->bi; bj = merge->bj; 1641d6ab1dc0SHong Zhang owners = merge->rowmap->range; 16429566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(bi[cm]+1,&ba)); 1643d6ab1dc0SHong Zhang 1644d6ab1dc0SHong Zhang /* get A_loc by taking all local rows of A */ 1645*1cdffd5eSHong Zhang A_loc = ap->A_loc; 16469566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_REUSE_MATRIX,&A_loc)); 1647d6ab1dc0SHong Zhang a_loc = (Mat_SeqAIJ*)(A_loc)->data; 1648d6ab1dc0SHong Zhang ai = a_loc->i; 1649d6ab1dc0SHong Zhang aj = a_loc->j; 1650d6ab1dc0SHong Zhang 1651ce496241SStefano Zampini /* trigger copy to CPU */ 16529566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(p->A,&dummy)); 16539566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(p->A,&dummy)); 16549566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(p->B,&dummy)); 16559566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(p->B,&dummy)); 1656d6ab1dc0SHong Zhang for (i=0; i<am; i++) { 1657d6ab1dc0SHong Zhang anz = ai[i+1] - ai[i]; 1658d6ab1dc0SHong Zhang adj = aj + ai[i]; 1659d6ab1dc0SHong Zhang ada = a_loc->a + ai[i]; 1660d6ab1dc0SHong Zhang 1661d6ab1dc0SHong Zhang /* 2-b) Compute Cseq = P_loc[i,:]^T*A[i,:] using outer product */ 1662e745005bSHong Zhang /*-------------------------------------------------------------*/ 1663d6ab1dc0SHong Zhang /* put the value into Co=(p->B)^T*A (off-diagonal part, send to others) */ 1664d6ab1dc0SHong Zhang pnz = po->i[i+1] - po->i[i]; 1665d6ab1dc0SHong Zhang poJ = po->j + po->i[i]; 1666d6ab1dc0SHong Zhang pA = po->a + po->i[i]; 1667d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1668d6ab1dc0SHong Zhang row = poJ[j]; 1669d6ab1dc0SHong Zhang cj = coj + coi[row]; 1670d6ab1dc0SHong Zhang ca = coa + coi[row]; 1671e745005bSHong Zhang /* perform sparse axpy */ 1672e745005bSHong Zhang nexta = 0; 1673d6ab1dc0SHong Zhang valtmp = pA[j]; 1674e745005bSHong Zhang for (k=0; nexta<anz; k++) { 1675e745005bSHong Zhang if (cj[k] == adj[nexta]) { 1676e745005bSHong Zhang ca[k] += valtmp*ada[nexta]; 1677e745005bSHong Zhang nexta++; 1678d6ab1dc0SHong Zhang } 1679e745005bSHong Zhang } 16809566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*anz)); 1681d6ab1dc0SHong Zhang } 1682d6ab1dc0SHong Zhang 1683d6ab1dc0SHong Zhang /* put the value into Cd (diagonal part) */ 1684d6ab1dc0SHong Zhang pnz = pd->i[i+1] - pd->i[i]; 1685d6ab1dc0SHong Zhang pdJ = pd->j + pd->i[i]; 1686d6ab1dc0SHong Zhang pA = pd->a + pd->i[i]; 1687d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1688d6ab1dc0SHong Zhang row = pdJ[j]; 1689d6ab1dc0SHong Zhang cj = bj + bi[row]; 1690d6ab1dc0SHong Zhang ca = ba + bi[row]; 1691e745005bSHong Zhang /* perform sparse axpy */ 1692e745005bSHong Zhang nexta = 0; 1693d6ab1dc0SHong Zhang valtmp = pA[j]; 1694e745005bSHong Zhang for (k=0; nexta<anz; k++) { 1695e745005bSHong Zhang if (cj[k] == adj[nexta]) { 1696e745005bSHong Zhang ca[k] += valtmp*ada[nexta]; 1697e745005bSHong Zhang nexta++; 1698d6ab1dc0SHong Zhang } 1699e745005bSHong Zhang } 17009566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*anz)); 1701d6ab1dc0SHong Zhang } 1702d6ab1dc0SHong Zhang } 1703d6ab1dc0SHong Zhang 1704d6ab1dc0SHong Zhang /* 3) send and recv matrix values coa */ 1705d6ab1dc0SHong Zhang /*------------------------------------*/ 1706d6ab1dc0SHong Zhang buf_ri = merge->buf_ri; 1707d6ab1dc0SHong Zhang buf_rj = merge->buf_rj; 1708d6ab1dc0SHong Zhang len_s = merge->len_s; 17099566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&taga)); 17109566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvScalar(comm,taga,merge->nrecv,merge->id_r,merge->len_r,&abuf_r,&r_waits)); 1711d6ab1dc0SHong Zhang 17129566063dSJacob Faibussowitsch PetscCall(PetscMalloc2(merge->nsend+1,&s_waits,size,&status)); 1713d6ab1dc0SHong Zhang for (proc=0,k=0; proc<size; proc++) { 1714d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1715d6ab1dc0SHong Zhang i = merge->owners_co[proc]; 17169566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coa+coi[i],len_s[proc],MPIU_MATSCALAR,proc,taga,comm,s_waits+k)); 1717d6ab1dc0SHong Zhang k++; 1718d6ab1dc0SHong Zhang } 17199566063dSJacob Faibussowitsch if (merge->nrecv) PetscCallMPI(MPI_Waitall(merge->nrecv,r_waits,status)); 17209566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,s_waits,status)); 1721d6ab1dc0SHong Zhang 17229566063dSJacob Faibussowitsch PetscCall(PetscFree2(s_waits,status)); 17239566063dSJacob Faibussowitsch PetscCall(PetscFree(r_waits)); 17249566063dSJacob Faibussowitsch PetscCall(PetscFree(coa)); 1725d6ab1dc0SHong Zhang 1726d6ab1dc0SHong Zhang /* 4) insert local Cseq and received values into Cmpi */ 1727d6ab1dc0SHong Zhang /*----------------------------------------------------*/ 17289566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(merge->nrecv,&buf_ri_k,merge->nrecv,&nextrow,merge->nrecv,&nextci)); 1729d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { 1730e745005bSHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 1731d6ab1dc0SHong Zhang nrows = *(buf_ri_k[k]); 1732d6ab1dc0SHong Zhang nextrow[k] = buf_ri_k[k]+1; /* next row number of k-th recved i-structure */ 1733a5b23f4aSJose E. Roman nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th recved i-structure */ 1734d6ab1dc0SHong Zhang } 1735d6ab1dc0SHong Zhang 1736d6ab1dc0SHong Zhang for (i=0; i<cm; i++) { 1737d6ab1dc0SHong Zhang row = owners[rank] + i; /* global row index of C_seq */ 1738d6ab1dc0SHong Zhang bj_i = bj + bi[i]; /* col indices of the i-th row of C */ 1739d6ab1dc0SHong Zhang ba_i = ba + bi[i]; 1740d6ab1dc0SHong Zhang bnz = bi[i+1] - bi[i]; 1741d6ab1dc0SHong Zhang /* add received vals into ba */ 1742d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { /* k-th received message */ 1743d6ab1dc0SHong Zhang /* i-th row */ 1744d6ab1dc0SHong Zhang if (i == *nextrow[k]) { 1745d6ab1dc0SHong Zhang cnz = *(nextci[k]+1) - *nextci[k]; 1746d6ab1dc0SHong Zhang cj = buf_rj[k] + *(nextci[k]); 1747d6ab1dc0SHong Zhang ca = abuf_r[k] + *(nextci[k]); 1748d6ab1dc0SHong Zhang nextcj = 0; 1749d6ab1dc0SHong Zhang for (j=0; nextcj<cnz; j++) { 1750d6ab1dc0SHong Zhang if (bj_i[j] == cj[nextcj]) { /* bcol == ccol */ 1751d6ab1dc0SHong Zhang ba_i[j] += ca[nextcj++]; 1752d6ab1dc0SHong Zhang } 1753d6ab1dc0SHong Zhang } 1754d6ab1dc0SHong Zhang nextrow[k]++; nextci[k]++; 17559566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*cnz)); 1756d6ab1dc0SHong Zhang } 1757d6ab1dc0SHong Zhang } 17589566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,bnz,bj_i,ba_i,INSERT_VALUES)); 1759d6ab1dc0SHong Zhang } 17609566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 17619566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 1762d6ab1dc0SHong Zhang 17639566063dSJacob Faibussowitsch PetscCall(PetscFree(ba)); 17649566063dSJacob Faibussowitsch PetscCall(PetscFree(abuf_r[0])); 17659566063dSJacob Faibussowitsch PetscCall(PetscFree(abuf_r)); 17669566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 1767d6ab1dc0SHong Zhang PetscFunctionReturn(0); 1768d6ab1dc0SHong Zhang } 1769d6ab1dc0SHong Zhang 17704222ddf1SHong Zhang PetscErrorCode MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ(Mat P,Mat A,PetscReal fill,Mat C) 1771d6ab1dc0SHong Zhang { 1772b419842dSStefano Zampini Mat A_loc; 1773*1cdffd5eSHong Zhang Mat_APMPI *ap; 17740298fd71SBarry Smith PetscFreeSpaceList free_space=NULL,current_space=NULL; 17756718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data,*a=(Mat_MPIAIJ*)A->data; 1776d6ab1dc0SHong Zhang PetscInt *pdti,*pdtj,*poti,*potj,*ptJ; 1777d6ab1dc0SHong Zhang PetscInt nnz; 1778d6ab1dc0SHong Zhang PetscInt *lnk,*owners_co,*coi,*coj,i,k,pnz,row; 1779d6ab1dc0SHong Zhang PetscInt am =A->rmap->n,pn=P->cmap->n; 1780ce94432eSBarry Smith MPI_Comm comm; 1781d6ab1dc0SHong Zhang PetscMPIInt size,rank,tagi,tagj,*len_si,*len_s,*len_ri; 1782d6ab1dc0SHong Zhang PetscInt **buf_rj,**buf_ri,**buf_ri_k; 1783d6ab1dc0SHong Zhang PetscInt len,proc,*dnz,*onz,*owners; 1784d6ab1dc0SHong Zhang PetscInt nzi,*bi,*bj; 1785d6ab1dc0SHong Zhang PetscInt nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci; 1786d6ab1dc0SHong Zhang MPI_Request *swaits,*rwaits; 1787d6ab1dc0SHong Zhang MPI_Status *sstatus,rstatus; 1788d6ab1dc0SHong Zhang Mat_Merge_SeqsToMPI *merge; 1789d6ab1dc0SHong Zhang PetscInt *ai,*aj,*Jptr,anz,*prmap=p->garray,pon,nspacedouble=0,j; 1790d6ab1dc0SHong Zhang PetscReal afill =1.0,afill_tmp; 17914b38b95cSHong Zhang PetscInt rstart = P->cmap->rstart,rmax,aN=A->cmap->N,Armax; 1792b419842dSStefano Zampini Mat_SeqAIJ *a_loc; 17930acc65b4SHong Zhang PetscTable ta; 1794b92f168fSBarry Smith MatType mtype; 1795d6ab1dc0SHong Zhang 1796d6ab1dc0SHong Zhang PetscFunctionBegin; 17979566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 1798d6ab1dc0SHong Zhang /* check if matrix local sizes are compatible */ 1799aed4548fSBarry 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); 1800d6ab1dc0SHong Zhang 18019566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 18029566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1803d6ab1dc0SHong Zhang 18043cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 1805*1cdffd5eSHong Zhang PetscCall(PetscNew(&ap)); 1806d6ab1dc0SHong Zhang 1807d6ab1dc0SHong Zhang /* get A_loc by taking all local rows of A */ 18089566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_INITIAL_MATRIX,&A_loc)); 18092205254eSKarl Rupp 1810*1cdffd5eSHong Zhang ap->A_loc = A_loc; 1811d6ab1dc0SHong Zhang a_loc = (Mat_SeqAIJ*)(A_loc)->data; 1812d6ab1dc0SHong Zhang ai = a_loc->i; 1813d6ab1dc0SHong Zhang aj = a_loc->j; 1814d6ab1dc0SHong Zhang 1815d6ab1dc0SHong Zhang /* determine symbolic Co=(p->B)^T*A - send to others */ 1816d6ab1dc0SHong Zhang /*----------------------------------------------------*/ 18179566063dSJacob Faibussowitsch PetscCall(MatGetSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj)); 18189566063dSJacob Faibussowitsch PetscCall(MatGetSymbolicTranspose_SeqAIJ(p->B,&poti,&potj)); 1819d6ab1dc0SHong Zhang pon = (p->B)->cmap->n; /* total num of rows to be sent to other processors 1820d6ab1dc0SHong Zhang >= (num of nonzero rows of C_seq) - pn */ 18219566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pon+1,&coi)); 1822d6ab1dc0SHong Zhang coi[0] = 0; 1823d6ab1dc0SHong Zhang 1824d6ab1dc0SHong Zhang /* set initial free space to be fill*(nnz(p->B) + nnz(A)) */ 1825f91af8c7SBarry Smith nnz = PetscRealIntMultTruncate(fill,PetscIntSumTruncate(poti[pon],ai[am])); 18269566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(nnz,&free_space)); 1827d6ab1dc0SHong Zhang current_space = free_space; 182819f0e57aSHong Zhang 1829d6ab1dc0SHong Zhang /* create and initialize a linked list */ 18309566063dSJacob Faibussowitsch PetscCall(PetscTableCreate(A->cmap->n + a->B->cmap->N,aN,&ta)); 18314b38b95cSHong Zhang MatRowMergeMax_SeqAIJ(a_loc,am,ta); 18329566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Armax)); 183333ba5e2fSHong Zhang 18349566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(Armax,&lnk)); 1835d6ab1dc0SHong Zhang 1836d6ab1dc0SHong Zhang for (i=0; i<pon; i++) { 1837d6ab1dc0SHong Zhang pnz = poti[i+1] - poti[i]; 1838d6ab1dc0SHong Zhang ptJ = potj + poti[i]; 1839d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1840d6ab1dc0SHong Zhang row = ptJ[j]; /* row of A_loc == col of Pot */ 1841d6ab1dc0SHong Zhang anz = ai[row+1] - ai[row]; 1842d6ab1dc0SHong Zhang Jptr = aj + ai[row]; 1843d6ab1dc0SHong Zhang /* add non-zero cols of AP into the sorted linked list lnk */ 18449566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(anz,Jptr,lnk)); 1845d6ab1dc0SHong Zhang } 1846d6ab1dc0SHong Zhang nnz = lnk[0]; 1847d6ab1dc0SHong Zhang 1848d6ab1dc0SHong Zhang /* If free space is not available, double the total space in the list */ 1849d6ab1dc0SHong Zhang if (current_space->local_remaining<nnz) { 18509566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(nnz,current_space->total_array_size),¤t_space)); 1851d6ab1dc0SHong Zhang nspacedouble++; 1852d6ab1dc0SHong Zhang } 1853d6ab1dc0SHong Zhang 1854d6ab1dc0SHong Zhang /* Copy data into free space, and zero out denserows */ 18559566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(nnz,current_space->array,lnk)); 18562205254eSKarl Rupp 1857d6ab1dc0SHong Zhang current_space->array += nnz; 1858d6ab1dc0SHong Zhang current_space->local_used += nnz; 1859d6ab1dc0SHong Zhang current_space->local_remaining -= nnz; 18602205254eSKarl Rupp 1861d6ab1dc0SHong Zhang coi[i+1] = coi[i] + nnz; 1862d6ab1dc0SHong Zhang } 1863d6ab1dc0SHong Zhang 18649566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(coi[pon]+1,&coj)); 18659566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,coj)); 18669566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); /* must destroy to get a new one for C */ 18672205254eSKarl Rupp 1868118727c9SMark F. Adams afill_tmp = (PetscReal)coi[pon]/(poti[pon] + ai[am]+1); 1869d6ab1dc0SHong Zhang if (afill_tmp > afill) afill = afill_tmp; 1870d6ab1dc0SHong Zhang 1871d6ab1dc0SHong Zhang /* send j-array (coj) of Co to other processors */ 1872d6ab1dc0SHong Zhang /*----------------------------------------------*/ 1873d6ab1dc0SHong Zhang /* determine row ownership */ 18749566063dSJacob Faibussowitsch PetscCall(PetscNew(&merge)); 18759566063dSJacob Faibussowitsch PetscCall(PetscLayoutCreate(comm,&merge->rowmap)); 18762205254eSKarl Rupp 1877d6ab1dc0SHong Zhang merge->rowmap->n = pn; 1878d6ab1dc0SHong Zhang merge->rowmap->bs = 1; 18792205254eSKarl Rupp 18809566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(merge->rowmap)); 1881d6ab1dc0SHong Zhang owners = merge->rowmap->range; 1882d6ab1dc0SHong Zhang 1883d6ab1dc0SHong Zhang /* determine the number of messages to send, their lengths */ 18849566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(size,&len_si)); 18859566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(size,&merge->len_s)); 18862205254eSKarl Rupp 1887d6ab1dc0SHong Zhang len_s = merge->len_s; 1888d6ab1dc0SHong Zhang merge->nsend = 0; 1889d6ab1dc0SHong Zhang 18909566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size+2,&owners_co)); 1891d6ab1dc0SHong Zhang 1892d6ab1dc0SHong Zhang proc = 0; 1893d6ab1dc0SHong Zhang for (i=0; i<pon; i++) { 1894d6ab1dc0SHong Zhang while (prmap[i] >= owners[proc+1]) proc++; 1895d6ab1dc0SHong Zhang len_si[proc]++; /* num of rows in Co to be sent to [proc] */ 1896d6ab1dc0SHong Zhang len_s[proc] += coi[i+1] - coi[i]; 1897d6ab1dc0SHong Zhang } 1898d6ab1dc0SHong Zhang 1899d6ab1dc0SHong Zhang len = 0; /* max length of buf_si[] */ 1900d6ab1dc0SHong Zhang owners_co[0] = 0; 1901d6ab1dc0SHong Zhang for (proc=0; proc<size; proc++) { 1902d6ab1dc0SHong Zhang owners_co[proc+1] = owners_co[proc] + len_si[proc]; 1903d6ab1dc0SHong Zhang if (len_si[proc]) { 1904d6ab1dc0SHong Zhang merge->nsend++; 1905d6ab1dc0SHong Zhang len_si[proc] = 2*(len_si[proc] + 1); 1906d6ab1dc0SHong Zhang len += len_si[proc]; 1907d6ab1dc0SHong Zhang } 1908d6ab1dc0SHong Zhang } 1909d6ab1dc0SHong Zhang 1910d6ab1dc0SHong Zhang /* determine the number and length of messages to receive for coi and coj */ 19119566063dSJacob Faibussowitsch PetscCall(PetscGatherNumberOfMessages(comm,NULL,len_s,&merge->nrecv)); 19129566063dSJacob Faibussowitsch PetscCall(PetscGatherMessageLengths2(comm,merge->nsend,merge->nrecv,len_s,len_si,&merge->id_r,&merge->len_r,&len_ri)); 1913d6ab1dc0SHong Zhang 1914d6ab1dc0SHong Zhang /* post the Irecv and Isend of coj */ 19159566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagj)); 19169566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagj,merge->nrecv,merge->id_r,merge->len_r,&buf_rj,&rwaits)); 19179566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(merge->nsend+1,&swaits)); 1918d6ab1dc0SHong Zhang for (proc=0, k=0; proc<size; proc++) { 1919d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1920d6ab1dc0SHong Zhang i = owners_co[proc]; 19219566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tagj,comm,swaits+k)); 1922d6ab1dc0SHong Zhang k++; 1923d6ab1dc0SHong Zhang } 1924d6ab1dc0SHong Zhang 1925d6ab1dc0SHong Zhang /* receives and sends of coj are complete */ 19269566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size,&sstatus)); 1927d6ab1dc0SHong Zhang for (i=0; i<merge->nrecv; i++) { 1928ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted; 19299566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(merge->nrecv,rwaits,&icompleted,&rstatus)); 1930d6ab1dc0SHong Zhang } 19319566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 19329566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,swaits,sstatus)); 1933d6ab1dc0SHong Zhang 19340acc65b4SHong Zhang /* add received column indices into table to update Armax */ 1935c4762a1bSJed Brown /* Armax can be as large as aN if a P[row,:] is dense, see src/ksp/ksp/tutorials/ex56.c! */ 19360acc65b4SHong Zhang for (k=0; k<merge->nrecv; k++) {/* k-th received message */ 19370acc65b4SHong Zhang Jptr = buf_rj[k]; 19380acc65b4SHong Zhang for (j=0; j<merge->len_r[k]; j++) { 19399566063dSJacob Faibussowitsch PetscCall(PetscTableAdd(ta,*(Jptr+j)+1,1,INSERT_VALUES)); 19400acc65b4SHong Zhang } 19410acc65b4SHong Zhang } 19429566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Armax)); 19430acc65b4SHong Zhang 1944d6ab1dc0SHong Zhang /* send and recv coi */ 1945d6ab1dc0SHong Zhang /*-------------------*/ 19469566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagi)); 19479566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagi,merge->nrecv,merge->id_r,len_ri,&buf_ri,&rwaits)); 19489566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(len+1,&buf_s)); 1949d6ab1dc0SHong Zhang buf_si = buf_s; /* points to the beginning of k-th msg to be sent */ 1950d6ab1dc0SHong Zhang for (proc=0,k=0; proc<size; proc++) { 1951d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1952d6ab1dc0SHong Zhang /* form outgoing message for i-structure: 1953d6ab1dc0SHong Zhang buf_si[0]: nrows to be sent 1954d6ab1dc0SHong Zhang [1:nrows]: row index (global) 1955d6ab1dc0SHong Zhang [nrows+1:2*nrows+1]: i-structure index 1956d6ab1dc0SHong Zhang */ 1957d6ab1dc0SHong Zhang /*-------------------------------------------*/ 1958d6ab1dc0SHong Zhang nrows = len_si[proc]/2 - 1; 1959d6ab1dc0SHong Zhang buf_si_i = buf_si + nrows+1; 1960d6ab1dc0SHong Zhang buf_si[0] = nrows; 1961d6ab1dc0SHong Zhang buf_si_i[0] = 0; 1962d6ab1dc0SHong Zhang nrows = 0; 1963d6ab1dc0SHong Zhang for (i=owners_co[proc]; i<owners_co[proc+1]; i++) { 1964d6ab1dc0SHong Zhang nzi = coi[i+1] - coi[i]; 1965d6ab1dc0SHong Zhang buf_si_i[nrows+1] = buf_si_i[nrows] + nzi; /* i-structure */ 1966d6ab1dc0SHong Zhang buf_si[nrows+1] = prmap[i] -owners[proc]; /* local row index */ 1967d6ab1dc0SHong Zhang nrows++; 1968d6ab1dc0SHong Zhang } 19699566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tagi,comm,swaits+k)); 1970d6ab1dc0SHong Zhang k++; 1971d6ab1dc0SHong Zhang buf_si += len_si[proc]; 1972d6ab1dc0SHong Zhang } 1973d6ab1dc0SHong Zhang i = merge->nrecv; 1974d6ab1dc0SHong Zhang while (i--) { 1975ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted; 19769566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(merge->nrecv,rwaits,&icompleted,&rstatus)); 1977d6ab1dc0SHong Zhang } 19789566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 19799566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,swaits,sstatus)); 19809566063dSJacob Faibussowitsch PetscCall(PetscFree(len_si)); 19819566063dSJacob Faibussowitsch PetscCall(PetscFree(len_ri)); 19829566063dSJacob Faibussowitsch PetscCall(PetscFree(swaits)); 19839566063dSJacob Faibussowitsch PetscCall(PetscFree(sstatus)); 19849566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_s)); 1985d6ab1dc0SHong Zhang 1986d6ab1dc0SHong Zhang /* compute the local portion of C (mpi mat) */ 1987d6ab1dc0SHong Zhang /*------------------------------------------*/ 1988d6ab1dc0SHong Zhang /* allocate bi array and free space for accumulating nonzero column info */ 19899566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pn+1,&bi)); 1990d6ab1dc0SHong Zhang bi[0] = 0; 1991d6ab1dc0SHong Zhang 1992d6ab1dc0SHong Zhang /* set initial free space to be fill*(nnz(P) + nnz(AP)) */ 1993f91af8c7SBarry Smith nnz = PetscRealIntMultTruncate(fill,PetscIntSumTruncate(pdti[pn],PetscIntSumTruncate(poti[pon],ai[am]))); 19949566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(nnz,&free_space)); 1995d6ab1dc0SHong Zhang current_space = free_space; 1996d6ab1dc0SHong Zhang 19979566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(merge->nrecv,&buf_ri_k,merge->nrecv,&nextrow,merge->nrecv,&nextci)); 1998d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { 1999d6ab1dc0SHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 2000d6ab1dc0SHong Zhang nrows = *buf_ri_k[k]; 2001d6ab1dc0SHong Zhang nextrow[k] = buf_ri_k[k] + 1; /* next row number of k-th recved i-structure */ 20022d4ee042Sprj- nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th received i-structure */ 2003d6ab1dc0SHong Zhang } 2004d6ab1dc0SHong Zhang 20059566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(Armax,&lnk)); 2006d0609cedSBarry Smith MatPreallocateBegin(comm,pn,A->cmap->n,dnz,onz); 2007d6ab1dc0SHong Zhang rmax = 0; 2008d6ab1dc0SHong Zhang for (i=0; i<pn; i++) { 2009d6ab1dc0SHong Zhang /* add pdt[i,:]*AP into lnk */ 2010d6ab1dc0SHong Zhang pnz = pdti[i+1] - pdti[i]; 2011d6ab1dc0SHong Zhang ptJ = pdtj + pdti[i]; 2012d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 2013d6ab1dc0SHong Zhang row = ptJ[j]; /* row of AP == col of Pt */ 2014d6ab1dc0SHong Zhang anz = ai[row+1] - ai[row]; 2015d6ab1dc0SHong Zhang Jptr = aj + ai[row]; 2016d6ab1dc0SHong Zhang /* add non-zero cols of AP into the sorted linked list lnk */ 20179566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(anz,Jptr,lnk)); 2018d6ab1dc0SHong Zhang } 2019d6ab1dc0SHong Zhang 2020d6ab1dc0SHong Zhang /* add received col data into lnk */ 2021d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { /* k-th received message */ 2022d6ab1dc0SHong Zhang if (i == *nextrow[k]) { /* i-th row */ 2023d6ab1dc0SHong Zhang nzi = *(nextci[k]+1) - *nextci[k]; 2024d6ab1dc0SHong Zhang Jptr = buf_rj[k] + *nextci[k]; 20259566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(nzi,Jptr,lnk)); 2026d6ab1dc0SHong Zhang nextrow[k]++; nextci[k]++; 2027d6ab1dc0SHong Zhang } 2028d6ab1dc0SHong Zhang } 20298c78258cSHong Zhang 20308c78258cSHong Zhang /* add missing diagonal entry */ 20318c78258cSHong Zhang if (C->force_diagonals) { 20328c78258cSHong Zhang k = i + owners[rank]; /* column index */ 20339566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(1,&k,lnk)); 20348c78258cSHong Zhang } 20358c78258cSHong Zhang 2036d6ab1dc0SHong Zhang nnz = lnk[0]; 2037d6ab1dc0SHong Zhang 2038d6ab1dc0SHong Zhang /* if free space is not available, make more free space */ 2039d6ab1dc0SHong Zhang if (current_space->local_remaining<nnz) { 20409566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(nnz,current_space->total_array_size),¤t_space)); 2041d6ab1dc0SHong Zhang nspacedouble++; 2042d6ab1dc0SHong Zhang } 2043d6ab1dc0SHong Zhang /* copy data into free space, then initialize lnk */ 20449566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(nnz,current_space->array,lnk)); 20459566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+owners[rank],nnz,current_space->array,dnz,onz)); 20462205254eSKarl Rupp 2047d6ab1dc0SHong Zhang current_space->array += nnz; 2048d6ab1dc0SHong Zhang current_space->local_used += nnz; 2049d6ab1dc0SHong Zhang current_space->local_remaining -= nnz; 20502205254eSKarl Rupp 2051d6ab1dc0SHong Zhang bi[i+1] = bi[i] + nnz; 2052d6ab1dc0SHong Zhang if (nnz > rmax) rmax = nnz; 2053d6ab1dc0SHong Zhang } 20549566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 2055d6ab1dc0SHong Zhang 20569566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(bi[pn]+1,&bj)); 20579566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,bj)); 2058118727c9SMark F. Adams afill_tmp = (PetscReal)bi[pn]/(pdti[pn] + poti[pon] + ai[am]+1); 2059d6ab1dc0SHong Zhang if (afill_tmp > afill) afill = afill_tmp; 20609566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); 20619566063dSJacob Faibussowitsch PetscCall(PetscTableDestroy(&ta)); 20629566063dSJacob Faibussowitsch PetscCall(MatRestoreSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj)); 20639566063dSJacob Faibussowitsch PetscCall(MatRestoreSymbolicTranspose_SeqAIJ(p->B,&poti,&potj)); 2064d6ab1dc0SHong Zhang 20654222ddf1SHong Zhang /* create symbolic parallel matrix C - why cannot be assembled in Numeric part */ 20664222ddf1SHong Zhang /*-------------------------------------------------------------------------------*/ 20679566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,pn,A->cmap->n,PETSC_DETERMINE,PETSC_DETERMINE)); 20689566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizes(C,PetscAbs(P->cmap->bs),PetscAbs(A->cmap->bs))); 20699566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 20709566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 20719566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 2072d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 20739566063dSJacob Faibussowitsch PetscCall(MatSetBlockSize(C,1)); 20749566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NO_OFF_PROC_ENTRIES,PETSC_TRUE)); 2075d6ab1dc0SHong Zhang for (i=0; i<pn; i++) { 2076d6ab1dc0SHong Zhang row = i + rstart; 2077d6ab1dc0SHong Zhang nnz = bi[i+1] - bi[i]; 2078d6ab1dc0SHong Zhang Jptr = bj + bi[i]; 20799566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,nnz,Jptr,NULL,INSERT_VALUES)); 2080d6ab1dc0SHong Zhang } 20819566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 20829566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 20839566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 2084d6ab1dc0SHong Zhang merge->bi = bi; 2085d6ab1dc0SHong Zhang merge->bj = bj; 2086d6ab1dc0SHong Zhang merge->coi = coi; 2087d6ab1dc0SHong Zhang merge->coj = coj; 2088d6ab1dc0SHong Zhang merge->buf_ri = buf_ri; 2089d6ab1dc0SHong Zhang merge->buf_rj = buf_rj; 2090d6ab1dc0SHong Zhang merge->owners_co = owners_co; 2091d6ab1dc0SHong Zhang 20924222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 2093*1cdffd5eSHong Zhang C->product->data = ap; 20946718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 2095*1cdffd5eSHong Zhang ap->merge = merge; 2096a560ef98SHong Zhang 20974222ddf1SHong Zhang C->ops->mattransposemultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ; 2098f96d37fcSHong Zhang 2099f96d37fcSHong Zhang #if defined(PETSC_USE_INFO) 2100f96d37fcSHong Zhang if (bi[pn] != 0) { 21019566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 21029566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatTransposeMatMult(A,B,MatReuse,%g,&C) for best performance.\n",(double)afill)); 2103f96d37fcSHong Zhang } else { 21049566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 2105f96d37fcSHong Zhang } 2106f96d37fcSHong Zhang #endif 2107f96d37fcSHong Zhang PetscFunctionReturn(0); 2108f96d37fcSHong Zhang } 21094222ddf1SHong Zhang 21104222ddf1SHong Zhang /* ---------------------------------------------------------------- */ 21114222ddf1SHong Zhang static PetscErrorCode MatProductSymbolic_AtB_MPIAIJ_MPIAIJ(Mat C) 21124222ddf1SHong Zhang { 21134222ddf1SHong Zhang Mat_Product *product = C->product; 21144222ddf1SHong Zhang Mat A=product->A,B=product->B; 21154222ddf1SHong Zhang PetscReal fill=product->fill; 21164222ddf1SHong Zhang PetscBool flg; 21174222ddf1SHong Zhang 21184222ddf1SHong Zhang PetscFunctionBegin; 21194222ddf1SHong Zhang /* scalable */ 21209566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"scalable",&flg)); 21214222ddf1SHong Zhang if (flg) { 21229566063dSJacob Faibussowitsch PetscCall(MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ(A,B,fill,C)); 21234222ddf1SHong Zhang goto next; 21244222ddf1SHong Zhang } 21254222ddf1SHong Zhang 21264222ddf1SHong Zhang /* nonscalable */ 21279566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"nonscalable",&flg)); 21284222ddf1SHong Zhang if (flg) { 21299566063dSJacob Faibussowitsch PetscCall(MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(A,B,fill,C)); 21304222ddf1SHong Zhang goto next; 21314222ddf1SHong Zhang } 21324222ddf1SHong Zhang 21334222ddf1SHong Zhang /* matmatmult */ 21349566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"at*b",&flg)); 21354222ddf1SHong Zhang if (flg) { 21364222ddf1SHong Zhang Mat At; 21374222ddf1SHong Zhang Mat_APMPI *ptap; 21384222ddf1SHong Zhang 21399566063dSJacob Faibussowitsch PetscCall(MatTranspose(A,MAT_INITIAL_MATRIX,&At)); 21409566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ(At,B,fill,C)); 21416718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 21424222ddf1SHong Zhang if (ptap) { 21434222ddf1SHong Zhang ptap->Pt = At; 21446718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 21454222ddf1SHong Zhang } 21464222ddf1SHong Zhang C->ops->transposematmultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_matmatmult; 21474222ddf1SHong Zhang goto next; 21484222ddf1SHong Zhang } 21494222ddf1SHong Zhang 21504e84afc0SStefano Zampini /* backend general code */ 21519566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"backend",&flg)); 21524e84afc0SStefano Zampini if (flg) { 21539566063dSJacob Faibussowitsch PetscCall(MatProductSymbolic_MPIAIJBACKEND(C)); 21544e84afc0SStefano Zampini PetscFunctionReturn(0); 21554e84afc0SStefano Zampini } 21564e84afc0SStefano Zampini 21574222ddf1SHong Zhang SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"MatProduct type is not supported"); 21584222ddf1SHong Zhang 21594222ddf1SHong Zhang next: 21604222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AtB; 21614222ddf1SHong Zhang PetscFunctionReturn(0); 21624222ddf1SHong Zhang } 21634222ddf1SHong Zhang 21644222ddf1SHong Zhang /* ---------------------------------------------------------------- */ 21654222ddf1SHong Zhang /* Set options for MatMatMultxxx_MPIAIJ_MPIAIJ */ 21664222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_AB(Mat C) 21674222ddf1SHong Zhang { 21684222ddf1SHong Zhang Mat_Product *product = C->product; 21694222ddf1SHong Zhang Mat A=product->A,B=product->B; 21704222ddf1SHong Zhang #if defined(PETSC_HAVE_HYPRE) 21714e84afc0SStefano Zampini const char *algTypes[5] = {"scalable","nonscalable","seqmpi","backend","hypre"}; 21724e84afc0SStefano Zampini PetscInt nalg = 5; 21734222ddf1SHong Zhang #else 21744e84afc0SStefano Zampini const char *algTypes[4] = {"scalable","nonscalable","seqmpi","backend",}; 21754e84afc0SStefano Zampini PetscInt nalg = 4; 21764222ddf1SHong Zhang #endif 21774222ddf1SHong Zhang PetscInt alg = 1; /* set nonscalable algorithm as default */ 21784222ddf1SHong Zhang PetscBool flg; 21794222ddf1SHong Zhang MPI_Comm comm; 21804222ddf1SHong Zhang 21814222ddf1SHong Zhang PetscFunctionBegin; 21829566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 21834222ddf1SHong Zhang 21844222ddf1SHong Zhang /* Set "nonscalable" as default algorithm */ 21859566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 21864222ddf1SHong Zhang if (flg) { 21879566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 21884222ddf1SHong Zhang 21894222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 21904222ddf1SHong Zhang if (B->cmap->N > 100000) { /* may switch to scalable algorithm as default */ 21914222ddf1SHong Zhang MatInfo Ainfo,Binfo; 21924222ddf1SHong Zhang PetscInt nz_local; 21934222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 21944222ddf1SHong Zhang 21959566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 21969566063dSJacob Faibussowitsch PetscCall(MatGetInfo(B,MAT_LOCAL,&Binfo)); 21974222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Binfo.nz_allocated); 21984222ddf1SHong Zhang 21994222ddf1SHong Zhang if (B->cmap->N > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 22001c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 22014222ddf1SHong Zhang 22024222ddf1SHong Zhang if (alg_scalable) { 22034222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 22049566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22059566063dSJacob Faibussowitsch PetscCall(PetscInfo(B,"Use scalable algorithm, BN %" PetscInt_FMT ", fill*nz_allocated %g\n",B->cmap->N,(double)(product->fill*nz_local))); 22064222ddf1SHong Zhang } 22074222ddf1SHong Zhang } 22084222ddf1SHong Zhang } 22094222ddf1SHong Zhang 22104222ddf1SHong Zhang /* Get runtime option */ 22114222ddf1SHong Zhang if (product->api_user) { 2212d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatMatMult","Mat"); 22139566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matmatmult_via","Algorithmic approach","MatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2214d0609cedSBarry Smith PetscOptionsEnd(); 22154222ddf1SHong Zhang } else { 2216d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_AB","Mat"); 22179566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2218d0609cedSBarry Smith PetscOptionsEnd(); 22194222ddf1SHong Zhang } 22204222ddf1SHong Zhang if (flg) { 22219566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22224222ddf1SHong Zhang } 22234222ddf1SHong Zhang 22244222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AB_MPIAIJ_MPIAIJ; 22254222ddf1SHong Zhang PetscFunctionReturn(0); 22264222ddf1SHong Zhang } 22274222ddf1SHong Zhang 22288a9c020eSBarry Smith static PetscErrorCode MatProductSetFromOptions_MPIAIJ_ABt(Mat C) 22298a9c020eSBarry Smith { 22308a9c020eSBarry Smith PetscFunctionBegin; 22318a9c020eSBarry Smith PetscCall(MatProductSetFromOptions_MPIAIJ_AB(C)); 22328a9c020eSBarry Smith C->ops->productsymbolic = MatProductSymbolic_ABt_MPIAIJ_MPIAIJ; 22338a9c020eSBarry Smith PetscFunctionReturn(0); 22348a9c020eSBarry Smith } 22358a9c020eSBarry Smith 22364222ddf1SHong Zhang /* Set options for MatTransposeMatMultXXX_MPIAIJ_MPIAIJ */ 22374222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_AtB(Mat C) 22384222ddf1SHong Zhang { 22394222ddf1SHong Zhang Mat_Product *product = C->product; 22404222ddf1SHong Zhang Mat A=product->A,B=product->B; 22414e84afc0SStefano Zampini const char *algTypes[4] = {"scalable","nonscalable","at*b","backend"}; 22424e84afc0SStefano Zampini PetscInt nalg = 4; 22434222ddf1SHong Zhang PetscInt alg = 1; /* set default algorithm */ 22444222ddf1SHong Zhang PetscBool flg; 22454222ddf1SHong Zhang MPI_Comm comm; 22464222ddf1SHong Zhang 22474222ddf1SHong Zhang PetscFunctionBegin; 22484222ddf1SHong Zhang /* Check matrix local sizes */ 22499566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 2250aed4548fSBarry 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); 22514222ddf1SHong Zhang 22524222ddf1SHong Zhang /* Set default algorithm */ 22539566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 22544222ddf1SHong Zhang if (flg) { 22559566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22564222ddf1SHong Zhang } 22574222ddf1SHong Zhang 22584222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 22594222ddf1SHong Zhang if (alg && B->cmap->N > 100000) { /* may switch to scalable algorithm as default */ 22604222ddf1SHong Zhang MatInfo Ainfo,Binfo; 22614222ddf1SHong Zhang PetscInt nz_local; 22624222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 22634222ddf1SHong Zhang 22649566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 22659566063dSJacob Faibussowitsch PetscCall(MatGetInfo(B,MAT_LOCAL,&Binfo)); 22664222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Binfo.nz_allocated); 22674222ddf1SHong Zhang 22684222ddf1SHong Zhang if (B->cmap->N > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 22691c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 22704222ddf1SHong Zhang 22714222ddf1SHong Zhang if (alg_scalable) { 22724222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 22739566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22749566063dSJacob Faibussowitsch PetscCall(PetscInfo(B,"Use scalable algorithm, BN %" PetscInt_FMT ", fill*nz_allocated %g\n",B->cmap->N,(double)(product->fill*nz_local))); 22754222ddf1SHong Zhang } 22764222ddf1SHong Zhang } 22774222ddf1SHong Zhang 22784222ddf1SHong Zhang /* Get runtime option */ 22794222ddf1SHong Zhang if (product->api_user) { 2280d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatTransposeMatMult","Mat"); 22819566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mattransposematmult_via","Algorithmic approach","MatTransposeMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2282d0609cedSBarry Smith PetscOptionsEnd(); 22834222ddf1SHong Zhang } else { 2284d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_AtB","Mat"); 22859566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatTransposeMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2286d0609cedSBarry Smith PetscOptionsEnd(); 22874222ddf1SHong Zhang } 22884222ddf1SHong Zhang if (flg) { 22899566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22904222ddf1SHong Zhang } 22914222ddf1SHong Zhang 22924222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AtB_MPIAIJ_MPIAIJ; 22934222ddf1SHong Zhang PetscFunctionReturn(0); 22944222ddf1SHong Zhang } 22954222ddf1SHong Zhang 22964222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_PtAP(Mat C) 22974222ddf1SHong Zhang { 22984222ddf1SHong Zhang Mat_Product *product = C->product; 22994222ddf1SHong Zhang Mat A=product->A,P=product->B; 23004222ddf1SHong Zhang MPI_Comm comm; 23014222ddf1SHong Zhang PetscBool flg; 23024222ddf1SHong Zhang PetscInt alg=1; /* set default algorithm */ 23034222ddf1SHong Zhang #if !defined(PETSC_HAVE_HYPRE) 23044e84afc0SStefano Zampini const char *algTypes[5] = {"scalable","nonscalable","allatonce","allatonce_merged","backend"}; 23054222ddf1SHong Zhang PetscInt nalg=5; 23064e84afc0SStefano Zampini #else 23074e84afc0SStefano Zampini const char *algTypes[6] = {"scalable","nonscalable","allatonce","allatonce_merged","backend","hypre"}; 23084e84afc0SStefano Zampini PetscInt nalg=6; 23094222ddf1SHong Zhang #endif 23104222ddf1SHong Zhang PetscInt pN=P->cmap->N; 23114222ddf1SHong Zhang 23124222ddf1SHong Zhang PetscFunctionBegin; 23134222ddf1SHong Zhang /* Check matrix local sizes */ 23149566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 2315aed4548fSBarry 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); 2316aed4548fSBarry 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); 23174222ddf1SHong Zhang 23184222ddf1SHong Zhang /* Set "nonscalable" as default algorithm */ 23199566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 23204222ddf1SHong Zhang if (flg) { 23219566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23224222ddf1SHong Zhang 23234222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 23244222ddf1SHong Zhang if (pN > 100000) { 23254222ddf1SHong Zhang MatInfo Ainfo,Pinfo; 23264222ddf1SHong Zhang PetscInt nz_local; 23274222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 23284222ddf1SHong Zhang 23299566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 23309566063dSJacob Faibussowitsch PetscCall(MatGetInfo(P,MAT_LOCAL,&Pinfo)); 23314222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Pinfo.nz_allocated); 23324222ddf1SHong Zhang 23334222ddf1SHong Zhang if (pN > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 23341c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 23354222ddf1SHong Zhang 23364222ddf1SHong Zhang if (alg_scalable) { 23374222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 23389566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23394222ddf1SHong Zhang } 23404222ddf1SHong Zhang } 23414222ddf1SHong Zhang } 23424222ddf1SHong Zhang 23434222ddf1SHong Zhang /* Get runtime option */ 23444222ddf1SHong Zhang if (product->api_user) { 2345d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatPtAP","Mat"); 23469566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matptap_via","Algorithmic approach","MatPtAP",algTypes,nalg,algTypes[alg],&alg,&flg)); 2347d0609cedSBarry Smith PetscOptionsEnd(); 23484222ddf1SHong Zhang } else { 2349d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_PtAP","Mat"); 23509566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatPtAP",algTypes,nalg,algTypes[alg],&alg,&flg)); 2351d0609cedSBarry Smith PetscOptionsEnd(); 23524222ddf1SHong Zhang } 23534222ddf1SHong Zhang if (flg) { 23549566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23554222ddf1SHong Zhang } 23564222ddf1SHong Zhang 23574222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_PtAP_MPIAIJ_MPIAIJ; 23584222ddf1SHong Zhang PetscFunctionReturn(0); 23594222ddf1SHong Zhang } 23604222ddf1SHong Zhang 23614222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_RARt(Mat C) 23624222ddf1SHong Zhang { 23634222ddf1SHong Zhang Mat_Product *product = C->product; 23644222ddf1SHong Zhang Mat A = product->A,R=product->B; 23654222ddf1SHong Zhang 23664222ddf1SHong Zhang PetscFunctionBegin; 23674222ddf1SHong Zhang /* Check matrix local sizes */ 2368aed4548fSBarry 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); 23694222ddf1SHong Zhang 23704222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_RARt_MPIAIJ_MPIAIJ; 23714222ddf1SHong Zhang PetscFunctionReturn(0); 23724222ddf1SHong Zhang } 23734222ddf1SHong Zhang 23744222ddf1SHong Zhang /* 23754222ddf1SHong Zhang Set options for ABC = A*B*C = A*(B*C); ABC's algorithm must be chosen from AB's algorithm 23764222ddf1SHong Zhang */ 23774222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_ABC(Mat C) 23784222ddf1SHong Zhang { 23794222ddf1SHong Zhang Mat_Product *product = C->product; 23804222ddf1SHong Zhang PetscBool flg = PETSC_FALSE; 23814222ddf1SHong Zhang PetscInt alg = 1; /* default algorithm */ 23824222ddf1SHong Zhang const char *algTypes[3] = {"scalable","nonscalable","seqmpi"}; 23834222ddf1SHong Zhang PetscInt nalg = 3; 23844222ddf1SHong Zhang 23854222ddf1SHong Zhang PetscFunctionBegin; 23864222ddf1SHong Zhang /* Set default algorithm */ 23879566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 23884222ddf1SHong Zhang if (flg) { 23899566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23904222ddf1SHong Zhang } 23914222ddf1SHong Zhang 23924222ddf1SHong Zhang /* Get runtime option */ 23934222ddf1SHong Zhang if (product->api_user) { 2394d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatMatMatMult","Mat"); 23959566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matmatmatmult_via","Algorithmic approach","MatMatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2396d0609cedSBarry Smith PetscOptionsEnd(); 23974222ddf1SHong Zhang } else { 2398d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_ABC","Mat"); 23999566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatProduct_ABC",algTypes,nalg,algTypes[alg],&alg,&flg)); 2400d0609cedSBarry Smith PetscOptionsEnd(); 24014222ddf1SHong Zhang } 24024222ddf1SHong Zhang if (flg) { 24039566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 24044222ddf1SHong Zhang } 24054222ddf1SHong Zhang 24064222ddf1SHong Zhang C->ops->matmatmultsymbolic = MatMatMatMultSymbolic_MPIAIJ_MPIAIJ_MPIAIJ; 24074222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_ABC; 24084222ddf1SHong Zhang PetscFunctionReturn(0); 24094222ddf1SHong Zhang } 24104222ddf1SHong Zhang 24114222ddf1SHong Zhang PETSC_INTERN PetscErrorCode MatProductSetFromOptions_MPIAIJ(Mat C) 24124222ddf1SHong Zhang { 24134222ddf1SHong Zhang Mat_Product *product = C->product; 24144222ddf1SHong Zhang 24154222ddf1SHong Zhang PetscFunctionBegin; 24164222ddf1SHong Zhang switch (product->type) { 24174222ddf1SHong Zhang case MATPRODUCT_AB: 24189566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_AB(C)); 24194222ddf1SHong Zhang break; 24208a9c020eSBarry Smith case MATPRODUCT_ABt: 24218a9c020eSBarry Smith PetscCall(MatProductSetFromOptions_MPIAIJ_ABt(C)); 24228a9c020eSBarry Smith break; 24234222ddf1SHong Zhang case MATPRODUCT_AtB: 24249566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_AtB(C)); 24254222ddf1SHong Zhang break; 24264222ddf1SHong Zhang case MATPRODUCT_PtAP: 24279566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_PtAP(C)); 24284222ddf1SHong Zhang break; 24294222ddf1SHong Zhang case MATPRODUCT_RARt: 24309566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_RARt(C)); 24314222ddf1SHong Zhang break; 24324222ddf1SHong Zhang case MATPRODUCT_ABC: 24339566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_ABC(C)); 24344222ddf1SHong Zhang break; 24356718818eSStefano Zampini default: 24366718818eSStefano Zampini break; 24374222ddf1SHong Zhang } 24384222ddf1SHong Zhang PetscFunctionReturn(0); 24394222ddf1SHong Zhang } 2440