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 184222ddf1SHong Zhang PETSC_INTERN PetscErrorCode MatProductSymbolic_AB_MPIAIJ_MPIAIJ(Mat C) 192499ec78SHong Zhang { 204222ddf1SHong Zhang Mat_Product *product = C->product; 214222ddf1SHong Zhang Mat A=product->A,B=product->B; 224222ddf1SHong Zhang MatProductAlgorithm alg=product->alg; 234222ddf1SHong Zhang PetscReal fill=product->fill; 245e68f438SHong Zhang PetscBool flg; 252499ec78SHong Zhang 262499ec78SHong Zhang PetscFunctionBegin; 274222ddf1SHong Zhang /* scalable */ 289566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(alg,"scalable",&flg)); 294222ddf1SHong Zhang if (flg) { 309566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ(A,B,fill,C)); 316718818eSStefano Zampini PetscFunctionReturn(0); 320fc8cf34SHong Zhang } 337d0a89b7SHong Zhang 344222ddf1SHong Zhang /* nonscalable */ 359566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(alg,"nonscalable",&flg)); 364222ddf1SHong Zhang if (flg) { 379566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(A,B,fill,C)); 386718818eSStefano Zampini PetscFunctionReturn(0); 394222ddf1SHong Zhang } 404222ddf1SHong Zhang 414222ddf1SHong Zhang /* seqmpi */ 429566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(alg,"seqmpi",&flg)); 434222ddf1SHong Zhang if (flg) { 449566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ_seqMPI(A,B,fill,C)); 456718818eSStefano Zampini PetscFunctionReturn(0); 464222ddf1SHong Zhang } 474222ddf1SHong Zhang 484e84afc0SStefano Zampini /* backend general code */ 499566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(alg,"backend",&flg)); 504e84afc0SStefano Zampini if (flg) { 519566063dSJacob Faibussowitsch PetscCall(MatProductSymbolic_MPIAIJBACKEND(C)); 524e84afc0SStefano Zampini PetscFunctionReturn(0); 534e84afc0SStefano Zampini } 544e84afc0SStefano Zampini 554222ddf1SHong Zhang #if defined(PETSC_HAVE_HYPRE) 569566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(alg,"hypre",&flg)); 574222ddf1SHong Zhang if (flg) { 589566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_AIJ_AIJ_wHYPRE(A,B,fill,C)); 594222ddf1SHong Zhang PetscFunctionReturn(0); 604222ddf1SHong Zhang } 614222ddf1SHong Zhang #endif 624222ddf1SHong Zhang SETERRQ(PetscObjectComm((PetscObject)C),PETSC_ERR_SUP,"Mat Product Algorithm is not supported"); 632499ec78SHong Zhang } 642499ec78SHong Zhang 656718818eSStefano Zampini PetscErrorCode MatDestroy_MPIAIJ_MatMatMult(void *data) 66598bc09dSHong Zhang { 676718818eSStefano Zampini Mat_APMPI *ptap = (Mat_APMPI*)data; 68598bc09dSHong Zhang 69598bc09dSHong Zhang PetscFunctionBegin; 709566063dSJacob Faibussowitsch PetscCall(PetscFree2(ptap->startsj_s,ptap->startsj_r)); 719566063dSJacob Faibussowitsch PetscCall(PetscFree(ptap->bufa)); 729566063dSJacob Faibussowitsch PetscCall(MatDestroy(&ptap->P_loc)); 739566063dSJacob Faibussowitsch PetscCall(MatDestroy(&ptap->P_oth)); 749566063dSJacob Faibussowitsch PetscCall(MatDestroy(&ptap->Pt)); 759566063dSJacob Faibussowitsch PetscCall(PetscFree(ptap->api)); 769566063dSJacob Faibussowitsch PetscCall(PetscFree(ptap->apj)); 779566063dSJacob Faibussowitsch PetscCall(PetscFree(ptap->apa)); 789566063dSJacob Faibussowitsch PetscCall(PetscFree(ptap)); 79598bc09dSHong Zhang PetscFunctionReturn(0); 80598bc09dSHong Zhang } 81598bc09dSHong Zhang 82f8efd71cSHong Zhang PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable(Mat A,Mat P,Mat C) 83598bc09dSHong Zhang { 844ae0847bSHong Zhang Mat_MPIAIJ *a =(Mat_MPIAIJ*)A->data,*c=(Mat_MPIAIJ*)C->data; 85598bc09dSHong Zhang Mat_SeqAIJ *ad =(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data; 86598bc09dSHong Zhang Mat_SeqAIJ *cd =(Mat_SeqAIJ*)(c->A)->data,*co=(Mat_SeqAIJ*)(c->B)->data; 87c12557a1SHong Zhang PetscScalar *cda=cd->a,*coa=co->a; 88598bc09dSHong Zhang Mat_SeqAIJ *p_loc,*p_oth; 899ce11a7cSHong Zhang PetscScalar *apa,*ca; 90c12557a1SHong Zhang PetscInt cm =C->rmap->n; 916718818eSStefano Zampini Mat_APMPI *ptap; 92c12557a1SHong Zhang PetscInt *api,*apj,*apJ,i,k; 9329825010SHong Zhang PetscInt cstart=C->cmap->rstart; 94598bc09dSHong Zhang PetscInt cdnz,conz,k0,k1; 95ce496241SStefano Zampini const PetscScalar *dummy; 969467f45fSHong Zhang MPI_Comm comm; 979467f45fSHong Zhang PetscMPIInt size; 98598bc09dSHong Zhang 99598bc09dSHong Zhang PetscFunctionBegin; 1006718818eSStefano Zampini MatCheckProduct(C,3); 1016718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 10228b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 1039566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 1049566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 105*aed4548fSBarry Smith PetscCheck(ptap->P_oth || size <= 1,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"AP cannot be reused. Do not call MatProductClear()"); 10680da8df7SHong Zhang 10765e4b4d4SStefano Zampini /* flag CPU mask for C */ 10865e4b4d4SStefano Zampini #if defined(PETSC_HAVE_DEVICE) 10965e4b4d4SStefano Zampini if (C->offloadmask != PETSC_OFFLOAD_UNALLOCATED) C->offloadmask = PETSC_OFFLOAD_CPU; 11065e4b4d4SStefano Zampini if (c->A->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->A->offloadmask = PETSC_OFFLOAD_CPU; 11165e4b4d4SStefano Zampini if (c->B->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->B->offloadmask = PETSC_OFFLOAD_CPU; 11265e4b4d4SStefano Zampini #endif 11365e4b4d4SStefano Zampini 114a1a4e84aSHong Zhang /* 1) get P_oth = ptap->P_oth and P_loc = ptap->P_loc */ 115598bc09dSHong Zhang /*-----------------------------------------------------*/ 116a1a4e84aSHong Zhang /* update numerical values of P_oth and P_loc */ 1179566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_REUSE_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 1189566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_REUSE_MATRIX,&ptap->P_loc)); 119598bc09dSHong Zhang 120598bc09dSHong Zhang /* 2) compute numeric C_loc = A_loc*P = Ad*P_loc + Ao*P_oth */ 121598bc09dSHong Zhang /*----------------------------------------------------------*/ 122598bc09dSHong Zhang /* get data from symbolic products */ 123a1a4e84aSHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 124c12557a1SHong Zhang p_oth = NULL; 1259467f45fSHong Zhang if (size >1) { 1269467f45fSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 1279467f45fSHong Zhang } 128598bc09dSHong Zhang 129598bc09dSHong Zhang /* get apa for storing dense row A[i,:]*P */ 130598bc09dSHong Zhang apa = ptap->apa; 131598bc09dSHong Zhang 13229825010SHong Zhang api = ptap->api; 13329825010SHong Zhang apj = ptap->apj; 134ce496241SStefano Zampini /* trigger copy to CPU */ 1359566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->A,&dummy)); 1369566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->A,&dummy)); 1379566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->B,&dummy)); 1389566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->B,&dummy)); 139598bc09dSHong Zhang for (i=0; i<cm; i++) { 140c12557a1SHong Zhang /* compute apa = A[i,:]*P */ 141e497d3c8SHong Zhang AProw_nonscalable(i,ad,ao,p_loc,p_oth,apa); 142598bc09dSHong Zhang 143598bc09dSHong Zhang /* set values in C */ 144598bc09dSHong Zhang apJ = apj + api[i]; 145598bc09dSHong Zhang cdnz = cd->i[i+1] - cd->i[i]; 146598bc09dSHong Zhang conz = co->i[i+1] - co->i[i]; 147598bc09dSHong Zhang 148910cf402Sprj- /* 1st off-diagonal part of C */ 149598bc09dSHong Zhang ca = coa + co->i[i]; 150598bc09dSHong Zhang k = 0; 151598bc09dSHong Zhang for (k0=0; k0<conz; k0++) { 152598bc09dSHong Zhang if (apJ[k] >= cstart) break; 153598bc09dSHong Zhang ca[k0] = apa[apJ[k]]; 1545cab4f04SHong Zhang apa[apJ[k++]] = 0.0; 155598bc09dSHong Zhang } 156598bc09dSHong Zhang 157598bc09dSHong Zhang /* diagonal part of C */ 158598bc09dSHong Zhang ca = cda + cd->i[i]; 159598bc09dSHong Zhang for (k1=0; k1<cdnz; k1++) { 160598bc09dSHong Zhang ca[k1] = apa[apJ[k]]; 1615cab4f04SHong Zhang apa[apJ[k++]] = 0.0; 162598bc09dSHong Zhang } 163598bc09dSHong Zhang 164910cf402Sprj- /* 2nd off-diagonal part of C */ 165598bc09dSHong Zhang ca = coa + co->i[i]; 166598bc09dSHong Zhang for (; k0<conz; k0++) { 167598bc09dSHong Zhang ca[k0] = apa[apJ[k]]; 1685cab4f04SHong Zhang apa[apJ[k++]] = 0.0; 169598bc09dSHong Zhang } 170598bc09dSHong Zhang } 1719566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 1729566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 173598bc09dSHong Zhang PetscFunctionReturn(0); 174598bc09dSHong Zhang } 175598bc09dSHong Zhang 1764222ddf1SHong Zhang PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(Mat A,Mat P,PetscReal fill,Mat C) 177598bc09dSHong Zhang { 178ce94432eSBarry Smith MPI_Comm comm; 1799467f45fSHong Zhang PetscMPIInt size; 1803cdca5ebSHong Zhang Mat_APMPI *ptap; 1810298fd71SBarry Smith PetscFreeSpaceList free_space=NULL,current_space=NULL; 1826718818eSStefano Zampini Mat_MPIAIJ *a=(Mat_MPIAIJ*)A->data; 183bfcd1a73SHong Zhang Mat_SeqAIJ *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_loc,*p_oth; 1847dbaa2c6Sandi selinger PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*dnz,*onz; 1854ae0847bSHong Zhang PetscInt *adi=ad->i,*adj=ad->j,*aoi=ao->i,*aoj=ao->j,rstart=A->rmap->rstart; 186d14fa243SHong Zhang PetscInt *lnk,i,pnz,row,*api,*apj,*Jptr,apnz,nspacedouble=0,j,nzi; 1875cab4f04SHong Zhang PetscInt am=A->rmap->n,pN=P->cmap->N,pn=P->cmap->n,pm=P->rmap->n; 188598bc09dSHong Zhang PetscBT lnkbt; 189bfcd1a73SHong Zhang PetscReal afill; 190b92f168fSBarry Smith MatType mtype; 191598bc09dSHong Zhang 192598bc09dSHong Zhang PetscFunctionBegin; 1936718818eSStefano Zampini MatCheckProduct(C,4); 19428b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Extra product struct not empty"); 1959566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 1969566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 1979467f45fSHong Zhang 1983cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 1999566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 200598bc09dSHong Zhang 201598bc09dSHong Zhang /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 2029566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 20319f0e57aSHong Zhang 204598bc09dSHong Zhang /* get P_loc by taking all local rows of P */ 2059566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc)); 206598bc09dSHong Zhang 207a1a4e84aSHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 208598bc09dSHong Zhang pi_loc = p_loc->i; pj_loc = p_loc->j; 2099467f45fSHong Zhang if (size > 1) { 2109467f45fSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 211598bc09dSHong Zhang pi_oth = p_oth->i; pj_oth = p_oth->j; 21220e3dc75SHong Zhang } else { 21320e3dc75SHong Zhang p_oth = NULL; 21420e3dc75SHong Zhang pi_oth = NULL; pj_oth = NULL; 2159467f45fSHong Zhang } 216598bc09dSHong Zhang 217598bc09dSHong Zhang /* first, compute symbolic AP = A_loc*P = A_diag*P_loc + A_off*P_oth */ 218598bc09dSHong Zhang /*-------------------------------------------------------------------*/ 2199566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&api)); 220a1a4e84aSHong Zhang ptap->api = api; 221598bc09dSHong Zhang api[0] = 0; 222598bc09dSHong Zhang 2235cab4f04SHong Zhang /* create and initialize a linked list */ 2249566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate(pN,pN,&lnk,&lnkbt)); 225598bc09dSHong Zhang 226bfcd1a73SHong Zhang /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */ 2279566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(adi[am],PetscIntSumTruncate(aoi[am],pi_loc[pm]))),&free_space)); 228598bc09dSHong Zhang current_space = free_space; 229598bc09dSHong Zhang 230d0609cedSBarry Smith MatPreallocateBegin(comm,am,pn,dnz,onz); 231598bc09dSHong Zhang for (i=0; i<am; i++) { 232598bc09dSHong Zhang /* diagonal portion of A */ 233598bc09dSHong Zhang nzi = adi[i+1] - adi[i]; 234598bc09dSHong Zhang for (j=0; j<nzi; j++) { 235598bc09dSHong Zhang row = *adj++; 236598bc09dSHong Zhang pnz = pi_loc[row+1] - pi_loc[row]; 237598bc09dSHong Zhang Jptr = pj_loc + pi_loc[row]; 238598bc09dSHong Zhang /* add non-zero cols of P into the sorted linked list lnk */ 2399566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(pnz,Jptr,lnk,lnkbt)); 240598bc09dSHong Zhang } 241598bc09dSHong Zhang /* off-diagonal portion of A */ 242598bc09dSHong Zhang nzi = aoi[i+1] - aoi[i]; 243598bc09dSHong Zhang for (j=0; j<nzi; j++) { 244598bc09dSHong Zhang row = *aoj++; 245598bc09dSHong Zhang pnz = pi_oth[row+1] - pi_oth[row]; 246598bc09dSHong Zhang Jptr = pj_oth + pi_oth[row]; 2479566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(pnz,Jptr,lnk,lnkbt)); 248598bc09dSHong Zhang } 2498c78258cSHong Zhang /* add possible missing diagonal entry */ 2508c78258cSHong Zhang if (C->force_diagonals) { 2518c78258cSHong Zhang j = i + rstart; /* column index */ 2529566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(1,&j,lnk,lnkbt)); 2538c78258cSHong Zhang } 254598bc09dSHong Zhang 255d14fa243SHong Zhang apnz = lnk[0]; 256598bc09dSHong Zhang api[i+1] = api[i] + apnz; 257598bc09dSHong Zhang 258598bc09dSHong Zhang /* if free space is not available, double the total space in the list */ 259598bc09dSHong Zhang if (current_space->local_remaining<apnz) { 2609566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(apnz,current_space->total_array_size),¤t_space)); 261598bc09dSHong Zhang nspacedouble++; 262598bc09dSHong Zhang } 263598bc09dSHong Zhang 264598bc09dSHong Zhang /* Copy data into free space, then initialize lnk */ 2659566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean(pN,apnz,current_space->array,lnk,lnkbt)); 2669566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+rstart,apnz,current_space->array,dnz,onz)); 2672205254eSKarl Rupp 268598bc09dSHong Zhang current_space->array += apnz; 269598bc09dSHong Zhang current_space->local_used += apnz; 270598bc09dSHong Zhang current_space->local_remaining -= apnz; 271598bc09dSHong Zhang } 272598bc09dSHong Zhang 273598bc09dSHong Zhang /* Allocate space for apj, initialize apj, and */ 274598bc09dSHong Zhang /* destroy list of free space and other temporary array(s) */ 2759566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(api[am]+1,&ptap->apj)); 276a1a4e84aSHong Zhang apj = ptap->apj; 2779566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,ptap->apj)); 2789566063dSJacob Faibussowitsch PetscCall(PetscLLDestroy(lnk,lnkbt)); 279598bc09dSHong Zhang 280f84c536bSHong Zhang /* malloc apa to store dense row A[i,:]*P */ 2819566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(pN,&ptap->apa)); 282f84c536bSHong Zhang 2834222ddf1SHong Zhang /* set and assemble symbolic parallel matrix C */ 2844222ddf1SHong Zhang /*---------------------------------------------*/ 2859566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,am,pn,PETSC_DETERMINE,PETSC_DETERMINE)); 2869566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,P)); 287a2f3521dSMark F. Adams 2889566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 2899566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 2909566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 291d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 29217f75172Sandi selinger 2939566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api)); 2949566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 2959566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 2969566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 297598bc09dSHong Zhang 2984222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable; 2994222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AB; 300598bc09dSHong Zhang 3014222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 3026718818eSStefano Zampini C->product->data = ptap; 3036718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_MatMatMult; 304598bc09dSHong Zhang 305bfcd1a73SHong Zhang /* set MatInfo */ 306118727c9SMark F. Adams afill = (PetscReal)api[am]/(adi[am]+aoi[am]+pi_loc[pm]+1) + 1.e-5; 307bfcd1a73SHong Zhang if (afill < 1.0) afill = 1.0; 3084222ddf1SHong Zhang C->info.mallocs = nspacedouble; 3094222ddf1SHong Zhang C->info.fill_ratio_given = fill; 3104222ddf1SHong Zhang C->info.fill_ratio_needed = afill; 311bfcd1a73SHong Zhang 312bfcd1a73SHong Zhang #if defined(PETSC_USE_INFO) 313bfcd1a73SHong Zhang if (api[am]) { 3149566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 3159566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n",(double)afill)); 316bfcd1a73SHong Zhang } else { 3179566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 318bfcd1a73SHong Zhang } 319bfcd1a73SHong Zhang #endif 320598bc09dSHong Zhang PetscFunctionReturn(0); 321598bc09dSHong Zhang } 322598bc09dSHong Zhang 3234222ddf1SHong Zhang /* ------------------------------------------------------- */ 3246718818eSStefano Zampini static PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIDense(Mat,Mat,PetscReal,Mat); 3256718818eSStefano Zampini static PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIDense(Mat,Mat,Mat); 3266718818eSStefano Zampini 3274222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense_AB(Mat C) 3289123193aSHong Zhang { 3294222ddf1SHong Zhang Mat_Product *product = C->product; 3304222ddf1SHong Zhang Mat A = product->A,B=product->B; 3319123193aSHong Zhang 3329123193aSHong Zhang PetscFunctionBegin; 3334222ddf1SHong Zhang if (A->cmap->rstart != B->rmap->rstart || A->cmap->rend != B->rmap->rend) 33498921bdaSJacob 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); 3354222ddf1SHong Zhang 3364222ddf1SHong Zhang C->ops->matmultsymbolic = MatMatMultSymbolic_MPIAIJ_MPIDense; 3374222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AB; 3389123193aSHong Zhang PetscFunctionReturn(0); 3399123193aSHong Zhang } 3404222ddf1SHong Zhang /* -------------------------------------------------------------------- */ 3414222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense_AtB(Mat C) 3424222ddf1SHong Zhang { 3434222ddf1SHong Zhang Mat_Product *product = C->product; 3444222ddf1SHong Zhang Mat A = product->A,B=product->B; 3454222ddf1SHong Zhang 3464222ddf1SHong Zhang PetscFunctionBegin; 3474222ddf1SHong Zhang if (A->rmap->rstart != B->rmap->rstart || A->rmap->rend != B->rmap->rend) 34898921bdaSJacob 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); 3494222ddf1SHong Zhang 3504222ddf1SHong Zhang C->ops->transposematmultsymbolic = MatTransposeMatMultSymbolic_MPIAIJ_MPIDense; 3514222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AtB; 3524222ddf1SHong Zhang PetscFunctionReturn(0); 3534222ddf1SHong Zhang } 3544222ddf1SHong Zhang 3554222ddf1SHong Zhang /* --------------------------------------------------------------------- */ 3564222ddf1SHong Zhang PETSC_INTERN PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense(Mat C) 3574222ddf1SHong Zhang { 3584222ddf1SHong Zhang Mat_Product *product = C->product; 3594222ddf1SHong Zhang 3604222ddf1SHong Zhang PetscFunctionBegin; 3614222ddf1SHong Zhang switch (product->type) { 3624222ddf1SHong Zhang case MATPRODUCT_AB: 3639566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_MPIDense_AB(C)); 3644222ddf1SHong Zhang break; 3654222ddf1SHong Zhang case MATPRODUCT_AtB: 3669566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_MPIDense_AtB(C)); 3674222ddf1SHong Zhang break; 3684222ddf1SHong Zhang default: 3696718818eSStefano Zampini break; 3704222ddf1SHong Zhang } 3714222ddf1SHong Zhang PetscFunctionReturn(0); 3724222ddf1SHong Zhang } 3734222ddf1SHong Zhang /* ------------------------------------------------------- */ 3749123193aSHong Zhang 3754324174eSBarry Smith typedef struct { 3765ea7661aSPierre Jolivet Mat workB,workB1; 3774324174eSBarry Smith MPI_Request *rwaits,*swaits; 37887753ddeSHong Zhang PetscInt nsends,nrecvs; 37987753ddeSHong Zhang MPI_Datatype *stype,*rtype; 3806718818eSStefano Zampini PetscInt blda; 3814324174eSBarry Smith } MPIAIJ_MPIDense; 3824324174eSBarry Smith 38396b95a6bSBarry Smith PetscErrorCode MatMPIAIJ_MPIDenseDestroy(void *ctx) 3844324174eSBarry Smith { 3854324174eSBarry Smith MPIAIJ_MPIDense *contents = (MPIAIJ_MPIDense*)ctx; 38687753ddeSHong Zhang PetscInt i; 3874324174eSBarry Smith 3884324174eSBarry Smith PetscFunctionBegin; 3899566063dSJacob Faibussowitsch PetscCall(MatDestroy(&contents->workB)); 3909566063dSJacob Faibussowitsch PetscCall(MatDestroy(&contents->workB1)); 39187753ddeSHong Zhang for (i=0; i<contents->nsends; i++) { 3929566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&contents->stype[i])); 39387753ddeSHong Zhang } 39487753ddeSHong Zhang for (i=0; i<contents->nrecvs; i++) { 3959566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&contents->rtype[i])); 39687753ddeSHong Zhang } 3979566063dSJacob Faibussowitsch PetscCall(PetscFree4(contents->stype,contents->rtype,contents->rwaits,contents->swaits)); 3989566063dSJacob Faibussowitsch PetscCall(PetscFree(contents)); 3994324174eSBarry Smith PetscFunctionReturn(0); 4004324174eSBarry Smith } 4014324174eSBarry Smith 4026718818eSStefano Zampini static PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIDense(Mat A,Mat B,PetscReal fill,Mat C) 4039123193aSHong Zhang { 404f32d5d43SBarry Smith Mat_MPIAIJ *aij=(Mat_MPIAIJ*)A->data; 405b45e3bf4SStefano Zampini PetscInt nz=aij->B->cmap->n,nsends,nrecvs,i,nrows_to,j,blda,m,M,n,N; 4064324174eSBarry Smith MPIAIJ_MPIDense *contents; 4074324174eSBarry Smith VecScatter ctx=aij->Mvctx; 4085ea7661aSPierre Jolivet PetscInt Am=A->rmap->n,Bm=B->rmap->n,BN=B->cmap->N,Bbn,Bbn1,bs,nrows_from,numBb; 40987753ddeSHong Zhang MPI_Comm comm; 41087753ddeSHong Zhang MPI_Datatype type1,*stype,*rtype; 41187753ddeSHong Zhang const PetscInt *sindices,*sstarts,*rstarts; 41287753ddeSHong Zhang PetscMPIInt *disp; 4136718818eSStefano Zampini PetscBool cisdense; 4149123193aSHong Zhang 4159123193aSHong Zhang PetscFunctionBegin; 4166718818eSStefano Zampini MatCheckProduct(C,4); 41728b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data not empty"); 4189566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 4199566063dSJacob Faibussowitsch PetscCall(PetscObjectBaseTypeCompare((PetscObject)C,MATMPIDENSE,&cisdense)); 4206718818eSStefano Zampini if (!cisdense) { 4219566063dSJacob Faibussowitsch PetscCall(MatSetType(C,((PetscObject)B)->type_name)); 4226718818eSStefano Zampini } 4239566063dSJacob Faibussowitsch PetscCall(MatGetLocalSize(C,&m,&n)); 4249566063dSJacob Faibussowitsch PetscCall(MatGetSize(C,&M,&N)); 425825bfe5fSPierre Jolivet if (m == PETSC_DECIDE || n == PETSC_DECIDE || M == PETSC_DECIDE || N == PETSC_DECIDE) { 4269566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,Am,B->cmap->n,A->rmap->N,BN)); 427825bfe5fSPierre Jolivet } 4289566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,B)); 4299566063dSJacob Faibussowitsch PetscCall(MatSetUp(C)); 4309566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(B,&blda)); 4319566063dSJacob Faibussowitsch PetscCall(PetscNew(&contents)); 4324324174eSBarry Smith 4339566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,NULL,NULL)); 4349566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,NULL,NULL)); 43587753ddeSHong Zhang 43687753ddeSHong Zhang /* Create column block of B and C for memory scalability when BN is too large */ 43787753ddeSHong Zhang /* Estimate Bbn, column size of Bb */ 43887753ddeSHong Zhang if (nz) { 43987753ddeSHong Zhang Bbn1 = 2*Am*BN/nz; 440da7cbbe7SHong Zhang if (!Bbn1) Bbn1 = 1; 44187753ddeSHong Zhang } else Bbn1 = BN; 44287753ddeSHong Zhang 44387753ddeSHong Zhang bs = PetscAbs(B->cmap->bs); 44487753ddeSHong Zhang Bbn1 = Bbn1/bs *bs; /* Bbn1 is a multiple of bs */ 44587753ddeSHong Zhang if (Bbn1 > BN) Bbn1 = BN; 4469566063dSJacob Faibussowitsch PetscCallMPI(MPI_Allreduce(&Bbn1,&Bbn,1,MPIU_INT,MPI_MAX,comm)); 44787753ddeSHong Zhang 44887753ddeSHong Zhang /* Enable runtime option for Bbn */ 449d0609cedSBarry Smith PetscOptionsBegin(comm,((PetscObject)C)->prefix,"MatMatMult","Mat"); 4509566063dSJacob Faibussowitsch PetscCall(PetscOptionsInt("-matmatmult_Bbn","Number of columns in Bb","MatMatMult",Bbn,&Bbn,NULL)); 451d0609cedSBarry Smith PetscOptionsEnd(); 4526718818eSStefano Zampini Bbn = PetscMin(Bbn,BN); 45387753ddeSHong Zhang 4546718818eSStefano Zampini if (Bbn > 0 && Bbn < BN) { 4555ea7661aSPierre Jolivet numBb = BN/Bbn; 4565ea7661aSPierre Jolivet Bbn1 = BN - numBb*Bbn; 4575ea7661aSPierre Jolivet } else numBb = 0; 45887753ddeSHong Zhang 4595ea7661aSPierre Jolivet if (numBb) { 4609566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"use Bb, BN=%" PetscInt_FMT ", Bbn=%" PetscInt_FMT "; numBb=%" PetscInt_FMT "\n",BN,Bbn,numBb)); 4615ea7661aSPierre Jolivet if (Bbn1) { /* Create workB1 for the remaining columns */ 4629566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"use Bb1, BN=%" PetscInt_FMT ", Bbn1=%" PetscInt_FMT "\n",BN,Bbn1)); 46387753ddeSHong Zhang /* Create work matrix used to store off processor rows of B needed for local product */ 4649566063dSJacob Faibussowitsch PetscCall(MatCreateSeqDense(PETSC_COMM_SELF,nz,Bbn1,NULL,&contents->workB1)); 4655ea7661aSPierre Jolivet } else contents->workB1 = NULL; 46687753ddeSHong Zhang } 46787753ddeSHong Zhang 46887753ddeSHong Zhang /* Create work matrix used to store off processor rows of B needed for local product */ 4699566063dSJacob Faibussowitsch PetscCall(MatCreateSeqDense(PETSC_COMM_SELF,nz,Bbn,NULL,&contents->workB)); 47087753ddeSHong Zhang 47187753ddeSHong Zhang /* Use MPI derived data type to reduce memory required by the send/recv buffers */ 4729566063dSJacob Faibussowitsch PetscCall(PetscMalloc4(nsends,&stype,nrecvs,&rtype,nrecvs,&contents->rwaits,nsends,&contents->swaits)); 47387753ddeSHong Zhang contents->stype = stype; 47487753ddeSHong Zhang contents->nsends = nsends; 47587753ddeSHong Zhang 47687753ddeSHong Zhang contents->rtype = rtype; 47787753ddeSHong Zhang contents->nrecvs = nrecvs; 4786718818eSStefano Zampini contents->blda = blda; 47987753ddeSHong Zhang 4809566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(Bm+1,&disp)); 48187753ddeSHong Zhang for (i=0; i<nsends; i++) { 48287753ddeSHong Zhang nrows_to = sstarts[i+1]-sstarts[i]; 483b45e3bf4SStefano Zampini for (j=0; j<nrows_to; j++) disp[j] = sindices[sstarts[i]+j]; /* rowB to be sent */ 4849566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_indexed_block(nrows_to,1,disp,MPIU_SCALAR,&type1)); 4859566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_resized(type1,0,blda*sizeof(PetscScalar),&stype[i])); 4869566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_commit(&stype[i])); 4879566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&type1)); 48887753ddeSHong Zhang } 48987753ddeSHong Zhang 49087753ddeSHong Zhang for (i=0; i<nrecvs; i++) { 49187753ddeSHong Zhang /* received values from a process form a (nrows_from x Bbn) row block in workB (column-wise) */ 49287753ddeSHong Zhang nrows_from = rstarts[i+1]-rstarts[i]; 49387753ddeSHong Zhang disp[0] = 0; 4949566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_indexed_block(1,nrows_from,disp,MPIU_SCALAR,&type1)); 4959566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_resized(type1,0,nz*sizeof(PetscScalar),&rtype[i])); 4969566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_commit(&rtype[i])); 4979566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&type1)); 49887753ddeSHong Zhang } 49987753ddeSHong Zhang 5009566063dSJacob Faibussowitsch PetscCall(PetscFree(disp)); 5019566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,NULL,NULL)); 5029566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,NULL,NULL)); 5039566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NO_OFF_PROC_ENTRIES,PETSC_TRUE)); 5049566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 5059566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 5069566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 5076718818eSStefano Zampini 5086718818eSStefano Zampini C->product->data = contents; 5096718818eSStefano Zampini C->product->destroy = MatMPIAIJ_MPIDenseDestroy; 5104222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIDense; 5119123193aSHong Zhang PetscFunctionReturn(0); 5129123193aSHong Zhang } 5139123193aSHong Zhang 51493aa15f2SStefano Zampini PETSC_INTERN PetscErrorCode MatMatMultNumericAdd_SeqAIJ_SeqDense(Mat,Mat,Mat,const PetscBool); 515f32d5d43SBarry Smith /* 5162636ff26SBarry Smith Performs an efficient scatter on the rows of B needed by this process; this is 5172636ff26SBarry Smith a modification of the VecScatterBegin_() routines. 51887753ddeSHong Zhang 51987753ddeSHong Zhang Input: Bbidx = 0: B = Bb 52087753ddeSHong Zhang = 1: B = Bb1, see MatMatMultSymbolic_MPIAIJ_MPIDense() 521f32d5d43SBarry Smith */ 52287753ddeSHong Zhang PetscErrorCode MatMPIDenseScatter(Mat A,Mat B,PetscInt Bbidx,Mat C,Mat *outworkB) 523f32d5d43SBarry Smith { 524f32d5d43SBarry Smith Mat_MPIAIJ *aij = (Mat_MPIAIJ*)A->data; 5251683a169SBarry Smith const PetscScalar *b; 52687753ddeSHong Zhang PetscScalar *rvalues; 527f32d5d43SBarry Smith VecScatter ctx = aij->Mvctx; 52887753ddeSHong Zhang const PetscInt *sindices,*sstarts,*rstarts; 5293515ee7fSJunchao Zhang const PetscMPIInt *sprocs,*rprocs; 5306718818eSStefano Zampini PetscInt i,nsends,nrecvs; 531f32d5d43SBarry Smith MPI_Request *swaits,*rwaits; 532ce94432eSBarry Smith MPI_Comm comm; 5336718818eSStefano Zampini PetscMPIInt tag=((PetscObject)ctx)->tag,ncols=B->cmap->N,nrows=aij->B->cmap->n,nsends_mpi,nrecvs_mpi; 5344324174eSBarry Smith MPIAIJ_MPIDense *contents; 5354324174eSBarry Smith Mat workB; 53687753ddeSHong Zhang MPI_Datatype *stype,*rtype; 5376718818eSStefano Zampini PetscInt blda; 538f32d5d43SBarry Smith 539f32d5d43SBarry Smith PetscFunctionBegin; 5406718818eSStefano Zampini MatCheckProduct(C,4); 54128b400f6SJacob Faibussowitsch PetscCheck(C->product->data,PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data empty"); 5426718818eSStefano Zampini contents = (MPIAIJ_MPIDense*)C->product->data; 5439566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,&sprocs,NULL/*bs*/)); 5449566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,&rprocs,NULL/*bs*/)); 5459566063dSJacob Faibussowitsch PetscCall(PetscMPIIntCast(nsends,&nsends_mpi)); 5469566063dSJacob Faibussowitsch PetscCall(PetscMPIIntCast(nrecvs,&nrecvs_mpi)); 547b45e3bf4SStefano Zampini if (Bbidx == 0) workB = *outworkB = contents->workB; 548b45e3bf4SStefano Zampini else workB = *outworkB = contents->workB1; 54908401ef6SPierre 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); 5503515ee7fSJunchao Zhang swaits = contents->swaits; 5513515ee7fSJunchao Zhang rwaits = contents->rwaits; 552f32d5d43SBarry Smith 5539566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(B,&b)); 5549566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(B,&blda)); 55508401ef6SPierre 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); 5569566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(workB,&rvalues)); 557f32d5d43SBarry Smith 55887753ddeSHong Zhang /* Post recv, use MPI derived data type to save memory */ 5599566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 56087753ddeSHong Zhang rtype = contents->rtype; 5613515ee7fSJunchao Zhang for (i=0; i<nrecvs; i++) { 5629566063dSJacob Faibussowitsch PetscCallMPI(MPI_Irecv(rvalues+(rstarts[i]-rstarts[0]),ncols,rtype[i],rprocs[i],tag,comm,rwaits+i)); 563f32d5d43SBarry Smith } 564f32d5d43SBarry Smith 56587753ddeSHong Zhang stype = contents->stype; 5663515ee7fSJunchao Zhang for (i=0; i<nsends; i++) { 5679566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(b,ncols,stype[i],sprocs[i],tag,comm,swaits+i)); 568f32d5d43SBarry Smith } 5692636ff26SBarry Smith 5709566063dSJacob Faibussowitsch if (nrecvs) PetscCallMPI(MPI_Waitall(nrecvs_mpi,rwaits,MPI_STATUSES_IGNORE)); 5719566063dSJacob Faibussowitsch if (nsends) PetscCallMPI(MPI_Waitall(nsends_mpi,swaits,MPI_STATUSES_IGNORE)); 572f32d5d43SBarry Smith 5739566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,&sprocs,NULL)); 5749566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,&rprocs,NULL)); 5759566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(B,&b)); 5769566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(workB,&rvalues)); 57787753ddeSHong Zhang PetscFunctionReturn(0); 57887753ddeSHong Zhang } 579f32d5d43SBarry Smith 5806718818eSStefano Zampini static PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIDense(Mat A,Mat B,Mat C) 5819123193aSHong Zhang { 582f32d5d43SBarry Smith Mat_MPIAIJ *aij = (Mat_MPIAIJ*)A->data; 583f32d5d43SBarry Smith Mat_MPIDense *bdense = (Mat_MPIDense*)B->data; 584f32d5d43SBarry Smith Mat_MPIDense *cdense = (Mat_MPIDense*)C->data; 585f32d5d43SBarry Smith Mat workB; 58687753ddeSHong Zhang MPIAIJ_MPIDense *contents; 5879123193aSHong Zhang 5889123193aSHong Zhang PetscFunctionBegin; 5896718818eSStefano Zampini MatCheckProduct(C,3); 59028b400f6SJacob Faibussowitsch PetscCheck(C->product->data,PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data empty"); 5916718818eSStefano Zampini contents = (MPIAIJ_MPIDense*)C->product->data; 592f32d5d43SBarry Smith /* diagonal block of A times all local rows of B */ 5935ea7661aSPierre Jolivet /* TODO: this calls a symbolic multiplication every time, which could be avoided */ 5949566063dSJacob Faibussowitsch PetscCall(MatMatMult(aij->A,bdense->A,MAT_REUSE_MATRIX,PETSC_DEFAULT,&cdense->A)); 5955ea7661aSPierre Jolivet if (contents->workB->cmap->n == B->cmap->N) { 59687753ddeSHong Zhang /* get off processor parts of B needed to complete C=A*B */ 5979566063dSJacob Faibussowitsch PetscCall(MatMPIDenseScatter(A,B,0,C,&workB)); 598f32d5d43SBarry Smith 599f32d5d43SBarry Smith /* off-diagonal block of A times nonlocal rows of B */ 6009566063dSJacob Faibussowitsch PetscCall(MatMatMultNumericAdd_SeqAIJ_SeqDense(aij->B,workB,cdense->A,PETSC_TRUE)); 60187753ddeSHong Zhang } else { 6025ea7661aSPierre Jolivet Mat Bb,Cb; 6035ea7661aSPierre Jolivet PetscInt BN=B->cmap->N,n=contents->workB->cmap->n,i; 604664c1e85SStefano Zampini PetscBool ccpu; 60587753ddeSHong Zhang 60608401ef6SPierre Jolivet PetscCheck(n > 0,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"Column block size %" PetscInt_FMT " must be positive",n); 607664c1e85SStefano Zampini /* Prevent from unneeded copies back and forth from the GPU 608664c1e85SStefano Zampini when getting and restoring the submatrix 609664c1e85SStefano Zampini We need a proper GPU code for AIJ * dense in parallel */ 6109566063dSJacob Faibussowitsch PetscCall(MatBoundToCPU(C,&ccpu)); 6119566063dSJacob Faibussowitsch PetscCall(MatBindToCPU(C,PETSC_TRUE)); 6125ea7661aSPierre Jolivet for (i=0; i<BN; i+=n) { 6139566063dSJacob Faibussowitsch PetscCall(MatDenseGetSubMatrix(B,i,PetscMin(i+n,BN),&Bb)); 6149566063dSJacob Faibussowitsch PetscCall(MatDenseGetSubMatrix(C,i,PetscMin(i+n,BN),&Cb)); 61587753ddeSHong Zhang 6165ea7661aSPierre Jolivet /* get off processor parts of B needed to complete C=A*B */ 6179566063dSJacob Faibussowitsch PetscCall(MatMPIDenseScatter(A,Bb,(i+n)>BN,C,&workB)); 6185ea7661aSPierre Jolivet 6195ea7661aSPierre Jolivet /* off-diagonal block of A times nonlocal rows of B */ 6205ea7661aSPierre Jolivet cdense = (Mat_MPIDense*)Cb->data; 6219566063dSJacob Faibussowitsch PetscCall(MatMatMultNumericAdd_SeqAIJ_SeqDense(aij->B,workB,cdense->A,PETSC_TRUE)); 6229566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreSubMatrix(B,&Bb)); 6239566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreSubMatrix(C,&Cb)); 62487753ddeSHong Zhang } 6259566063dSJacob Faibussowitsch PetscCall(MatBindToCPU(C,ccpu)); 62687753ddeSHong Zhang } 6279123193aSHong Zhang PetscFunctionReturn(0); 6289123193aSHong Zhang } 629cf1a0672SHong Zhang 630f8efd71cSHong Zhang PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIAIJ(Mat A,Mat P,Mat C) 6311634b032SHong Zhang { 632cf1a0672SHong Zhang Mat_MPIAIJ *a = (Mat_MPIAIJ*)A->data,*c=(Mat_MPIAIJ*)C->data; 633cf1a0672SHong Zhang Mat_SeqAIJ *ad = (Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data; 634cf1a0672SHong Zhang Mat_SeqAIJ *cd = (Mat_SeqAIJ*)(c->A)->data,*co=(Mat_SeqAIJ*)(c->B)->data; 635cf1a0672SHong Zhang PetscInt *adi = ad->i,*adj,*aoi=ao->i,*aoj; 636cf1a0672SHong Zhang PetscScalar *ada,*aoa,*cda=cd->a,*coa=co->a; 637cf1a0672SHong Zhang Mat_SeqAIJ *p_loc,*p_oth; 638cf1a0672SHong Zhang PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*pj; 63929825010SHong Zhang PetscScalar *pa_loc,*pa_oth,*pa,valtmp,*ca; 64029825010SHong Zhang PetscInt cm = C->rmap->n,anz,pnz; 6416718818eSStefano Zampini Mat_APMPI *ptap; 64280da8df7SHong Zhang PetscScalar *apa_sparse; 643ce496241SStefano Zampini const PetscScalar *dummy; 644cf1a0672SHong Zhang PetscInt *api,*apj,*apJ,i,j,k,row; 64529825010SHong Zhang PetscInt cstart = C->cmap->rstart; 64629825010SHong Zhang PetscInt cdnz,conz,k0,k1,nextp; 647a7c7454dSHong Zhang MPI_Comm comm; 648a7c7454dSHong Zhang PetscMPIInt size; 6491634b032SHong Zhang 6501634b032SHong Zhang PetscFunctionBegin; 6516718818eSStefano Zampini MatCheckProduct(C,3); 6526718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 65328b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 6549566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 6559566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 656*aed4548fSBarry Smith PetscCheck(ptap->P_oth || size <= 1,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"AP cannot be reused. Do not call MatProductClear()"); 657a7c7454dSHong Zhang 65865e4b4d4SStefano Zampini /* flag CPU mask for C */ 65965e4b4d4SStefano Zampini #if defined(PETSC_HAVE_DEVICE) 66065e4b4d4SStefano Zampini if (C->offloadmask != PETSC_OFFLOAD_UNALLOCATED) C->offloadmask = PETSC_OFFLOAD_CPU; 66165e4b4d4SStefano Zampini if (c->A->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->A->offloadmask = PETSC_OFFLOAD_CPU; 66265e4b4d4SStefano Zampini if (c->B->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->B->offloadmask = PETSC_OFFLOAD_CPU; 66365e4b4d4SStefano Zampini #endif 66480da8df7SHong Zhang apa_sparse = ptap->apa; 66580da8df7SHong Zhang 666cf1a0672SHong Zhang /* 1) get P_oth = ptap->P_oth and P_loc = ptap->P_loc */ 667cf1a0672SHong Zhang /*-----------------------------------------------------*/ 668cf1a0672SHong Zhang /* update numerical values of P_oth and P_loc */ 6699566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_REUSE_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 6709566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_REUSE_MATRIX,&ptap->P_loc)); 671cf1a0672SHong Zhang 672cf1a0672SHong Zhang /* 2) compute numeric C_loc = A_loc*P = Ad*P_loc + Ao*P_oth */ 673cf1a0672SHong Zhang /*----------------------------------------------------------*/ 674cf1a0672SHong Zhang /* get data from symbolic products */ 675cf1a0672SHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 676cf1a0672SHong Zhang pi_loc = p_loc->i; pj_loc = p_loc->j; pa_loc = p_loc->a; 677a7c7454dSHong Zhang if (size >1) { 678a7c7454dSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 679cf1a0672SHong Zhang pi_oth = p_oth->i; pj_oth = p_oth->j; pa_oth = p_oth->a; 68020e3dc75SHong Zhang } else { 68120e3dc75SHong Zhang p_oth = NULL; pi_oth = NULL; pj_oth = NULL; pa_oth = NULL; 682a7c7454dSHong Zhang } 683cf1a0672SHong Zhang 684ce496241SStefano Zampini /* trigger copy to CPU */ 6859566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->A,&dummy)); 6869566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->A,&dummy)); 6879566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->B,&dummy)); 6889566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->B,&dummy)); 68929825010SHong Zhang api = ptap->api; 69029825010SHong Zhang apj = ptap->apj; 691cf1a0672SHong Zhang for (i=0; i<cm; i++) { 69229825010SHong Zhang apJ = apj + api[i]; 69329825010SHong Zhang 694cf1a0672SHong Zhang /* diagonal portion of A */ 695cf1a0672SHong Zhang anz = adi[i+1] - adi[i]; 696cf1a0672SHong Zhang adj = ad->j + adi[i]; 697cf1a0672SHong Zhang ada = ad->a + adi[i]; 698cf1a0672SHong Zhang for (j=0; j<anz; j++) { 699cf1a0672SHong Zhang row = adj[j]; 700cf1a0672SHong Zhang pnz = pi_loc[row+1] - pi_loc[row]; 701cf1a0672SHong Zhang pj = pj_loc + pi_loc[row]; 702cf1a0672SHong Zhang pa = pa_loc + pi_loc[row]; 70329825010SHong Zhang /* perform sparse axpy */ 704cf1a0672SHong Zhang valtmp = ada[j]; 70529825010SHong Zhang nextp = 0; 70629825010SHong Zhang for (k=0; nextp<pnz; k++) { 70729825010SHong Zhang if (apJ[k] == pj[nextp]) { /* column of AP == column of P */ 70829825010SHong Zhang apa_sparse[k] += valtmp*pa[nextp++]; 70929825010SHong Zhang } 7101634b032SHong Zhang } 7119566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*pnz)); 712cf1a0672SHong Zhang } 7131634b032SHong Zhang 714cf1a0672SHong Zhang /* off-diagonal portion of A */ 715cf1a0672SHong Zhang anz = aoi[i+1] - aoi[i]; 716cf1a0672SHong Zhang aoj = ao->j + aoi[i]; 717cf1a0672SHong Zhang aoa = ao->a + aoi[i]; 718cf1a0672SHong Zhang for (j=0; j<anz; j++) { 719cf1a0672SHong Zhang row = aoj[j]; 720cf1a0672SHong Zhang pnz = pi_oth[row+1] - pi_oth[row]; 721cf1a0672SHong Zhang pj = pj_oth + pi_oth[row]; 722cf1a0672SHong Zhang pa = pa_oth + pi_oth[row]; 72329825010SHong Zhang /* perform sparse axpy */ 724cf1a0672SHong Zhang valtmp = aoa[j]; 72529825010SHong Zhang nextp = 0; 72629825010SHong Zhang for (k=0; nextp<pnz; k++) { 72729825010SHong Zhang if (apJ[k] == pj[nextp]) { /* column of AP == column of P */ 72829825010SHong Zhang apa_sparse[k] += valtmp*pa[nextp++]; 72929825010SHong Zhang } 730cf1a0672SHong Zhang } 7319566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*pnz)); 732cf1a0672SHong Zhang } 7331634b032SHong Zhang 734cf1a0672SHong Zhang /* set values in C */ 735cf1a0672SHong Zhang cdnz = cd->i[i+1] - cd->i[i]; 736cf1a0672SHong Zhang conz = co->i[i+1] - co->i[i]; 7371634b032SHong Zhang 738910cf402Sprj- /* 1st off-diagonal part of C */ 739cf1a0672SHong Zhang ca = coa + co->i[i]; 740cf1a0672SHong Zhang k = 0; 741cf1a0672SHong Zhang for (k0=0; k0<conz; k0++) { 742cf1a0672SHong Zhang if (apJ[k] >= cstart) break; 74329825010SHong Zhang ca[k0] = apa_sparse[k]; 74429825010SHong Zhang apa_sparse[k] = 0.0; 745cf1a0672SHong Zhang k++; 746cf1a0672SHong Zhang } 7471634b032SHong Zhang 748cf1a0672SHong Zhang /* diagonal part of C */ 749cf1a0672SHong Zhang ca = cda + cd->i[i]; 750cf1a0672SHong Zhang for (k1=0; k1<cdnz; k1++) { 75129825010SHong Zhang ca[k1] = apa_sparse[k]; 75229825010SHong Zhang apa_sparse[k] = 0.0; 753cf1a0672SHong Zhang k++; 754cf1a0672SHong Zhang } 755cf1a0672SHong Zhang 756910cf402Sprj- /* 2nd off-diagonal part of C */ 757cf1a0672SHong Zhang ca = coa + co->i[i]; 758cf1a0672SHong Zhang for (; k0<conz; k0++) { 75929825010SHong Zhang ca[k0] = apa_sparse[k]; 76029825010SHong Zhang apa_sparse[k] = 0.0; 761cf1a0672SHong Zhang k++; 762cf1a0672SHong Zhang } 763cf1a0672SHong Zhang } 7649566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 7659566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 766cf1a0672SHong Zhang PetscFunctionReturn(0); 767cf1a0672SHong Zhang } 768cf1a0672SHong Zhang 7690fc8cf34SHong Zhang /* same as MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(), except using LLCondensed to avoid O(BN) memory requirement */ 7704222ddf1SHong Zhang PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ(Mat A,Mat P,PetscReal fill,Mat C) 771cf1a0672SHong Zhang { 772ce94432eSBarry Smith MPI_Comm comm; 773a7c7454dSHong Zhang PetscMPIInt size; 7743cdca5ebSHong Zhang Mat_APMPI *ptap; 7750298fd71SBarry Smith PetscFreeSpaceList free_space = NULL,current_space=NULL; 7766718818eSStefano Zampini Mat_MPIAIJ *a = (Mat_MPIAIJ*)A->data; 777cf1a0672SHong Zhang Mat_SeqAIJ *ad = (Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_loc,*p_oth; 7787dbaa2c6Sandi selinger PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*dnz,*onz; 779cf1a0672SHong Zhang PetscInt *adi=ad->i,*adj=ad->j,*aoi=ao->i,*aoj=ao->j,rstart=A->rmap->rstart; 7808c78258cSHong Zhang PetscInt i,pnz,row,*api,*apj,*Jptr,apnz,nspacedouble=0,j,nzi,*lnk,apnz_max=1; 781e5ab8c0aSKarl Rupp PetscInt am=A->rmap->n,pn=P->cmap->n,pm=P->rmap->n,lsize=pn+20; 782cf1a0672SHong Zhang PetscReal afill; 783b92f168fSBarry Smith MatType mtype; 784cf1a0672SHong Zhang 785cf1a0672SHong Zhang PetscFunctionBegin; 7866718818eSStefano Zampini MatCheckProduct(C,4); 78728b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Extra product struct not empty"); 7889566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 7899566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 790a7c7454dSHong Zhang 7913cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 7929566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 793cf1a0672SHong Zhang 794cf1a0672SHong Zhang /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 7959566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 79619f0e57aSHong Zhang 797cf1a0672SHong Zhang /* get P_loc by taking all local rows of P */ 7989566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc)); 799cf1a0672SHong Zhang 800cf1a0672SHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 801cf1a0672SHong Zhang pi_loc = p_loc->i; pj_loc = p_loc->j; 802a7c7454dSHong Zhang if (size > 1) { 803a7c7454dSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 80420e3dc75SHong Zhang pi_oth = p_oth->i; pj_oth = p_oth->j; 805968382fdSHong Zhang } else { 80620e3dc75SHong Zhang p_oth = NULL; 80720e3dc75SHong Zhang pi_oth = NULL; pj_oth = NULL; 808a7c7454dSHong Zhang } 809cf1a0672SHong Zhang 810cf1a0672SHong Zhang /* first, compute symbolic AP = A_loc*P = A_diag*P_loc + A_off*P_oth */ 811cf1a0672SHong Zhang /*-------------------------------------------------------------------*/ 8129566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&api)); 813cf1a0672SHong Zhang ptap->api = api; 814cf1a0672SHong Zhang api[0] = 0; 815cf1a0672SHong Zhang 8169566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(lsize,&lnk)); 817cf1a0672SHong Zhang 818cf1a0672SHong Zhang /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */ 8199566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(adi[am],PetscIntSumTruncate(aoi[am],pi_loc[pm]))),&free_space)); 820cf1a0672SHong Zhang current_space = free_space; 821d0609cedSBarry Smith MatPreallocateBegin(comm,am,pn,dnz,onz); 822cf1a0672SHong Zhang for (i=0; i<am; i++) { 823cf1a0672SHong Zhang /* diagonal portion of A */ 824cf1a0672SHong Zhang nzi = adi[i+1] - adi[i]; 825cf1a0672SHong Zhang for (j=0; j<nzi; j++) { 826cf1a0672SHong Zhang row = *adj++; 827cf1a0672SHong Zhang pnz = pi_loc[row+1] - pi_loc[row]; 828cf1a0672SHong Zhang Jptr = pj_loc + pi_loc[row]; 8298a210135Sandi selinger /* Expand list if it is not long enough */ 8308328dd82Sandi selinger if (pnz+apnz_max > lsize) { 8318328dd82Sandi selinger lsize = pnz+apnz_max; 8329566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedExpand_Scalable(lsize, &lnk)); 8338a210135Sandi selinger } 834cf1a0672SHong Zhang /* add non-zero cols of P into the sorted linked list lnk */ 8359566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(pnz,Jptr,lnk)); 8368a210135Sandi selinger apnz = *lnk; /* The first element in the list is the number of items in the list */ 8378a210135Sandi selinger api[i+1] = api[i] + apnz; 8388c78258cSHong Zhang if (apnz > apnz_max) apnz_max = apnz + 1; /* '1' for diagonal entry */ 839cf1a0672SHong Zhang } 840cf1a0672SHong Zhang /* off-diagonal portion of A */ 841cf1a0672SHong Zhang nzi = aoi[i+1] - aoi[i]; 842cf1a0672SHong Zhang for (j=0; j<nzi; j++) { 843cf1a0672SHong Zhang row = *aoj++; 844cf1a0672SHong Zhang pnz = pi_oth[row+1] - pi_oth[row]; 845cf1a0672SHong Zhang Jptr = pj_oth + pi_oth[row]; 8468a210135Sandi selinger /* Expand list if it is not long enough */ 8478328dd82Sandi selinger if (pnz+apnz_max > lsize) { 8488328dd82Sandi selinger lsize = pnz + apnz_max; 8499566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedExpand_Scalable(lsize, &lnk)); 8508a210135Sandi selinger } 8518a210135Sandi selinger /* add non-zero cols of P into the sorted linked list lnk */ 8529566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(pnz,Jptr,lnk)); 8538a210135Sandi selinger apnz = *lnk; /* The first element in the list is the number of items in the list */ 8548a210135Sandi selinger api[i+1] = api[i] + apnz; 8558c78258cSHong Zhang if (apnz > apnz_max) apnz_max = apnz + 1; /* '1' for diagonal entry */ 856cf1a0672SHong Zhang } 8578c78258cSHong Zhang 8588c78258cSHong Zhang /* add missing diagonal entry */ 8598c78258cSHong Zhang if (C->force_diagonals) { 8608c78258cSHong Zhang j = i + rstart; /* column index */ 8619566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(1,&j,lnk)); 8628c78258cSHong Zhang } 8638c78258cSHong Zhang 864f84c536bSHong Zhang apnz = *lnk; 865cf1a0672SHong Zhang api[i+1] = api[i] + apnz; 8668a210135Sandi selinger if (apnz > apnz_max) apnz_max = apnz; 867cf1a0672SHong Zhang 868cf1a0672SHong Zhang /* if free space is not available, double the total space in the list */ 869cf1a0672SHong Zhang if (current_space->local_remaining<apnz) { 8709566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(apnz,current_space->total_array_size),¤t_space)); 871cf1a0672SHong Zhang nspacedouble++; 872cf1a0672SHong Zhang } 873cf1a0672SHong Zhang 874cf1a0672SHong Zhang /* Copy data into free space, then initialize lnk */ 8759566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(apnz,current_space->array,lnk)); 8769566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+rstart,apnz,current_space->array,dnz,onz)); 8772205254eSKarl Rupp 878cf1a0672SHong Zhang current_space->array += apnz; 879cf1a0672SHong Zhang current_space->local_used += apnz; 880cf1a0672SHong Zhang current_space->local_remaining -= apnz; 881cf1a0672SHong Zhang } 882cf1a0672SHong Zhang 883cf1a0672SHong Zhang /* Allocate space for apj, initialize apj, and */ 884cf1a0672SHong Zhang /* destroy list of free space and other temporary array(s) */ 8859566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(api[am]+1,&ptap->apj)); 886cf1a0672SHong Zhang apj = ptap->apj; 8879566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,ptap->apj)); 8889566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); 889cf1a0672SHong Zhang 8904222ddf1SHong Zhang /* create and assemble symbolic parallel matrix C */ 891cf1a0672SHong Zhang /*----------------------------------------------------*/ 8929566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,am,pn,PETSC_DETERMINE,PETSC_DETERMINE)); 8939566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,P)); 8949566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 8959566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 8969566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 897d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 898cf1a0672SHong Zhang 8994222ddf1SHong Zhang /* malloc apa for assembly C */ 9009566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(apnz_max,&ptap->apa)); 90117f75172Sandi selinger 9029566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api)); 9039566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 9049566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 9059566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 906cf1a0672SHong Zhang 9074222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ; 9084222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AB; 909cf1a0672SHong Zhang 9104222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 9116718818eSStefano Zampini C->product->data = ptap; 9126718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_MatMatMult; 913cf1a0672SHong Zhang 914cf1a0672SHong Zhang /* set MatInfo */ 915118727c9SMark F. Adams afill = (PetscReal)api[am]/(adi[am]+aoi[am]+pi_loc[pm]+1) + 1.e-5; 916cf1a0672SHong Zhang if (afill < 1.0) afill = 1.0; 9174222ddf1SHong Zhang C->info.mallocs = nspacedouble; 9184222ddf1SHong Zhang C->info.fill_ratio_given = fill; 9194222ddf1SHong Zhang C->info.fill_ratio_needed = afill; 920cf1a0672SHong Zhang 921cf1a0672SHong Zhang #if defined(PETSC_USE_INFO) 922cf1a0672SHong Zhang if (api[am]) { 9239566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 9249566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n",(double)afill)); 925cf1a0672SHong Zhang } else { 9269566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 927cf1a0672SHong Zhang } 928cf1a0672SHong Zhang #endif 9291634b032SHong Zhang PetscFunctionReturn(0); 9301634b032SHong Zhang } 931f96d37fcSHong Zhang 932a9d7e0ccSandi selinger /* This function is needed for the seqMPI matrix-matrix multiplication. */ 933a9d7e0ccSandi selinger /* Three input arrays are merged to one output array. The size of the */ 934a9d7e0ccSandi selinger /* output array is also output. Duplicate entries only show up once. */ 935a9d7e0ccSandi selinger static void Merge3SortedArrays(PetscInt size1, PetscInt *in1, 936a9d7e0ccSandi selinger PetscInt size2, PetscInt *in2, 937a9d7e0ccSandi selinger PetscInt size3, PetscInt *in3, 938a9d7e0ccSandi selinger PetscInt *size4, PetscInt *out) 939a9d7e0ccSandi selinger { 940a9d7e0ccSandi selinger int i = 0, j = 0, k = 0, l = 0; 941a9d7e0ccSandi selinger 942a9d7e0ccSandi selinger /* Traverse all three arrays */ 943a9d7e0ccSandi selinger while (i<size1 && j<size2 && k<size3) { 944a9d7e0ccSandi selinger if (in1[i] < in2[j] && in1[i] < in3[k]) { 945a9d7e0ccSandi selinger out[l++] = in1[i++]; 946a9d7e0ccSandi selinger } 947a9d7e0ccSandi selinger else if (in2[j] < in1[i] && in2[j] < in3[k]) { 948a9d7e0ccSandi selinger out[l++] = in2[j++]; 949a9d7e0ccSandi selinger } 950a9d7e0ccSandi selinger else if (in3[k] < in1[i] && in3[k] < in2[j]) { 951a9d7e0ccSandi selinger out[l++] = in3[k++]; 952a9d7e0ccSandi selinger } 953a9d7e0ccSandi selinger else if (in1[i] == in2[j] && in1[i] < in3[k]) { 954a9d7e0ccSandi selinger out[l++] = in1[i]; 955a9d7e0ccSandi selinger i++, j++; 956a9d7e0ccSandi selinger } 957a9d7e0ccSandi selinger else if (in1[i] == in3[k] && in1[i] < in2[j]) { 958a9d7e0ccSandi selinger out[l++] = in1[i]; 959a9d7e0ccSandi selinger i++, k++; 960a9d7e0ccSandi selinger } 961a9d7e0ccSandi selinger else if (in3[k] == in2[j] && in2[j] < in1[i]) { 962a9d7e0ccSandi selinger out[l++] = in2[j]; 963a9d7e0ccSandi selinger k++, j++; 964a9d7e0ccSandi selinger } 965a9d7e0ccSandi selinger else if (in1[i] == in2[j] && in1[i] == in3[k]) { 966a9d7e0ccSandi selinger out[l++] = in1[i]; 967a9d7e0ccSandi selinger i++, j++, k++; 968a9d7e0ccSandi selinger } 969a9d7e0ccSandi selinger } 970a9d7e0ccSandi selinger 971a9d7e0ccSandi selinger /* Traverse two remaining arrays */ 972a9d7e0ccSandi selinger while (i<size1 && j<size2) { 973a9d7e0ccSandi selinger if (in1[i] < in2[j]) { 974a9d7e0ccSandi selinger out[l++] = in1[i++]; 975a9d7e0ccSandi selinger } 976a9d7e0ccSandi selinger else if (in1[i] > in2[j]) { 977a9d7e0ccSandi selinger out[l++] = in2[j++]; 978a9d7e0ccSandi selinger } 979a9d7e0ccSandi selinger else { 980a9d7e0ccSandi selinger out[l++] = in1[i]; 981a9d7e0ccSandi selinger i++, j++; 982a9d7e0ccSandi selinger } 983a9d7e0ccSandi selinger } 984a9d7e0ccSandi selinger 985a9d7e0ccSandi selinger while (i<size1 && k<size3) { 986a9d7e0ccSandi selinger if (in1[i] < in3[k]) { 987a9d7e0ccSandi selinger out[l++] = in1[i++]; 988a9d7e0ccSandi selinger } 989a9d7e0ccSandi selinger else if (in1[i] > in3[k]) { 990a9d7e0ccSandi selinger out[l++] = in3[k++]; 991a9d7e0ccSandi selinger } 992a9d7e0ccSandi selinger else { 993a9d7e0ccSandi selinger out[l++] = in1[i]; 994a9d7e0ccSandi selinger i++, k++; 995a9d7e0ccSandi selinger } 996a9d7e0ccSandi selinger } 997a9d7e0ccSandi selinger 998a9d7e0ccSandi selinger while (k<size3 && j<size2) { 999a9d7e0ccSandi selinger if (in3[k] < in2[j]) { 1000a9d7e0ccSandi selinger out[l++] = in3[k++]; 1001a9d7e0ccSandi selinger } 1002a9d7e0ccSandi selinger else if (in3[k] > in2[j]) { 1003a9d7e0ccSandi selinger out[l++] = in2[j++]; 1004a9d7e0ccSandi selinger } 1005a9d7e0ccSandi selinger else { 1006a9d7e0ccSandi selinger out[l++] = in3[k]; 1007a9d7e0ccSandi selinger k++, j++; 1008a9d7e0ccSandi selinger } 1009a9d7e0ccSandi selinger } 1010a9d7e0ccSandi selinger 1011a9d7e0ccSandi selinger /* Traverse one remaining array */ 1012a9d7e0ccSandi selinger while (i<size1) out[l++] = in1[i++]; 1013a9d7e0ccSandi selinger while (j<size2) out[l++] = in2[j++]; 1014a9d7e0ccSandi selinger while (k<size3) out[l++] = in3[k++]; 1015a9d7e0ccSandi selinger 1016a9d7e0ccSandi selinger *size4 = l; 1017a9d7e0ccSandi selinger } 1018a9d7e0ccSandi selinger 101973b67375Sandi selinger /* This matrix-matrix multiplication algorithm divides the multiplication into three multiplications and */ 102073b67375Sandi selinger /* adds up the products. Two of these three multiplications are performed with existing (sequential) */ 102173b67375Sandi selinger /* matrix-matrix multiplications. */ 10224222ddf1SHong Zhang PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ_seqMPI(Mat A, Mat P, PetscReal fill, Mat C) 1023a9d7e0ccSandi selinger { 1024a9d7e0ccSandi selinger MPI_Comm comm; 1025a9d7e0ccSandi selinger PetscMPIInt size; 10263cdca5ebSHong Zhang Mat_APMPI *ptap; 1027a9d7e0ccSandi selinger PetscFreeSpaceList free_space_diag=NULL, current_space=NULL; 1028a9d7e0ccSandi selinger Mat_MPIAIJ *a =(Mat_MPIAIJ*)A->data; 1029a9d7e0ccSandi selinger Mat_SeqAIJ *ad =(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_loc; 1030a9d7e0ccSandi selinger Mat_MPIAIJ *p =(Mat_MPIAIJ*)P->data; 1031a9d7e0ccSandi selinger Mat_SeqAIJ *adpd_seq, *p_off, *aopoth_seq; 1032a9d7e0ccSandi selinger PetscInt adponz, adpdnz; 10337dbaa2c6Sandi selinger PetscInt *pi_loc,*dnz,*onz; 1034a9d7e0ccSandi selinger PetscInt *adi=ad->i,*adj=ad->j,*aoi=ao->i,rstart=A->rmap->rstart; 1035a9d7e0ccSandi selinger PetscInt *lnk,i, i1=0,pnz,row,*adpoi,*adpoj, *api, *adpoJ, *aopJ, *apJ,*Jptr, aopnz, nspacedouble=0,j,nzi, 1036a9d7e0ccSandi selinger *apj,apnz, *adpdi, *adpdj, *adpdJ, *poff_i, *poff_j, *j_temp, *aopothi, *aopothj; 1037a016958eSandi selinger PetscInt am=A->rmap->n,pN=P->cmap->N,pn=P->cmap->n,pm=P->rmap->n, p_colstart, p_colend; 1038a9d7e0ccSandi selinger PetscBT lnkbt; 1039a9d7e0ccSandi selinger PetscReal afill; 1040a9d7e0ccSandi selinger PetscMPIInt rank; 1041a9d7e0ccSandi selinger Mat adpd, aopoth; 1042b92f168fSBarry Smith MatType mtype; 1043e83e5f86SFande Kong const char *prefix; 1044a9d7e0ccSandi selinger 1045a9d7e0ccSandi selinger PetscFunctionBegin; 10466718818eSStefano Zampini MatCheckProduct(C,4); 104728b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Extra product struct not empty"); 10489566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 10499566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 10509566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm, &rank)); 10519566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRangeColumn(P, &p_colstart, &p_colend)); 1052a9d7e0ccSandi selinger 10533cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 10549566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 1055a9d7e0ccSandi selinger 1056a9d7e0ccSandi selinger /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 10579566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 1058a9d7e0ccSandi selinger 1059a9d7e0ccSandi selinger /* get P_loc by taking all local rows of P */ 10609566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc)); 1061a9d7e0ccSandi selinger 1062a9d7e0ccSandi selinger p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 1063a9d7e0ccSandi selinger pi_loc = p_loc->i; 1064a9d7e0ccSandi selinger 1065a9d7e0ccSandi selinger /* Allocate memory for the i arrays of the matrices A*P, A_diag*P_off and A_offd * P */ 10669566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&api)); 10679566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&adpoi)); 1068a9d7e0ccSandi selinger 1069a9d7e0ccSandi selinger adpoi[0] = 0; 1070a9d7e0ccSandi selinger ptap->api = api; 1071a9d7e0ccSandi selinger api[0] = 0; 1072a9d7e0ccSandi selinger 1073a9d7e0ccSandi selinger /* create and initialize a linked list, will be used for both A_diag * P_loc_off and A_offd * P_oth */ 10749566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate(pN,pN,&lnk,&lnkbt)); 1075d0609cedSBarry Smith MatPreallocateBegin(comm,am,pn,dnz,onz); 1076a9d7e0ccSandi selinger 1077a9d7e0ccSandi selinger /* Symbolic calc of A_loc_diag * P_loc_diag */ 10789566063dSJacob Faibussowitsch PetscCall(MatGetOptionsPrefix(A,&prefix)); 10799566063dSJacob Faibussowitsch PetscCall(MatProductCreate(a->A,p->A,NULL,&adpd)); 10809566063dSJacob Faibussowitsch PetscCall(MatGetOptionsPrefix(A,&prefix)); 10819566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(adpd,prefix)); 10829566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(adpd,"inner_diag_")); 10834222ddf1SHong Zhang 10849566063dSJacob Faibussowitsch PetscCall(MatProductSetType(adpd,MATPRODUCT_AB)); 10859566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(adpd,"sorted")); 10869566063dSJacob Faibussowitsch PetscCall(MatProductSetFill(adpd,fill)); 10879566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions(adpd)); 10888c78258cSHong Zhang 10898c78258cSHong Zhang adpd->force_diagonals = C->force_diagonals; 10909566063dSJacob Faibussowitsch PetscCall(MatProductSymbolic(adpd)); 10914222ddf1SHong Zhang 1092a9d7e0ccSandi selinger adpd_seq = (Mat_SeqAIJ*)((adpd)->data); 1093a9d7e0ccSandi selinger adpdi = adpd_seq->i; adpdj = adpd_seq->j; 1094a9d7e0ccSandi selinger p_off = (Mat_SeqAIJ*)((p->B)->data); 1095a9d7e0ccSandi selinger poff_i = p_off->i; poff_j = p_off->j; 1096a9d7e0ccSandi selinger 109773b67375Sandi selinger /* j_temp stores indices of a result row before they are added to the linked list */ 10989566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pN+2,&j_temp)); 1099a9d7e0ccSandi selinger 1100a9d7e0ccSandi selinger /* Symbolic calc of the A_diag * p_loc_off */ 1101a9d7e0ccSandi selinger /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */ 11029566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(adi[am],PetscIntSumTruncate(aoi[am],pi_loc[pm]))),&free_space_diag)); 1103a9d7e0ccSandi selinger current_space = free_space_diag; 1104a9d7e0ccSandi selinger 1105a9d7e0ccSandi selinger for (i=0; i<am; i++) { 1106a9d7e0ccSandi selinger /* A_diag * P_loc_off */ 1107a9d7e0ccSandi selinger nzi = adi[i+1] - adi[i]; 1108a9d7e0ccSandi selinger for (j=0; j<nzi; j++) { 1109a9d7e0ccSandi selinger row = *adj++; 1110a9d7e0ccSandi selinger pnz = poff_i[row+1] - poff_i[row]; 1111a9d7e0ccSandi selinger Jptr = poff_j + poff_i[row]; 1112a9d7e0ccSandi selinger for (i1 = 0; i1 < pnz; i1++) { 1113a9d7e0ccSandi selinger j_temp[i1] = p->garray[Jptr[i1]]; 1114a9d7e0ccSandi selinger } 1115a9d7e0ccSandi selinger /* add non-zero cols of P into the sorted linked list lnk */ 11169566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(pnz,j_temp,lnk,lnkbt)); 1117a9d7e0ccSandi selinger } 1118a9d7e0ccSandi selinger 1119a9d7e0ccSandi selinger adponz = lnk[0]; 1120a9d7e0ccSandi selinger adpoi[i+1] = adpoi[i] + adponz; 1121a9d7e0ccSandi selinger 1122a9d7e0ccSandi selinger /* if free space is not available, double the total space in the list */ 1123a9d7e0ccSandi selinger if (current_space->local_remaining<adponz) { 11249566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(adponz,current_space->total_array_size),¤t_space)); 1125a9d7e0ccSandi selinger nspacedouble++; 1126a9d7e0ccSandi selinger } 1127a9d7e0ccSandi selinger 1128a9d7e0ccSandi selinger /* Copy data into free space, then initialize lnk */ 11299566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean(pN,adponz,current_space->array,lnk,lnkbt)); 1130a9d7e0ccSandi selinger 1131a9d7e0ccSandi selinger current_space->array += adponz; 1132a9d7e0ccSandi selinger current_space->local_used += adponz; 1133a9d7e0ccSandi selinger current_space->local_remaining -= adponz; 1134a9d7e0ccSandi selinger } 1135a9d7e0ccSandi selinger 1136a9d7e0ccSandi selinger /* Symbolic calc of A_off * P_oth */ 11379566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(a->B,prefix)); 11389566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(a->B,"inner_offdiag_")); 11399566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF,&aopoth)); 11409566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(a->B, ptap->P_oth, fill, aopoth)); 1141a9d7e0ccSandi selinger aopoth_seq = (Mat_SeqAIJ*)((aopoth)->data); 1142a9d7e0ccSandi selinger aopothi = aopoth_seq->i; aopothj = aopoth_seq->j; 1143a9d7e0ccSandi selinger 1144a9d7e0ccSandi selinger /* Allocate space for apj, adpj, aopj, ... */ 1145a9d7e0ccSandi selinger /* destroy lists of free space and other temporary array(s) */ 1146a9d7e0ccSandi selinger 11479566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(aopothi[am] + adpoi[am] + adpdi[am]+2, &ptap->apj)); 11489566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(adpoi[am]+2, &adpoj)); 1149a9d7e0ccSandi selinger 1150a9d7e0ccSandi selinger /* Copy from linked list to j-array */ 11519566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space_diag,adpoj)); 11529566063dSJacob Faibussowitsch PetscCall(PetscLLDestroy(lnk,lnkbt)); 1153a9d7e0ccSandi selinger 1154a9d7e0ccSandi selinger adpoJ = adpoj; 1155a9d7e0ccSandi selinger adpdJ = adpdj; 1156a9d7e0ccSandi selinger aopJ = aopothj; 1157a9d7e0ccSandi selinger apj = ptap->apj; 1158a9d7e0ccSandi selinger apJ = apj; /* still empty */ 1159a9d7e0ccSandi selinger 1160a9d7e0ccSandi selinger /* Merge j-arrays of A_off * P, A_diag * P_loc_off, and */ 1161a9d7e0ccSandi selinger /* A_diag * P_loc_diag to get A*P */ 1162a9d7e0ccSandi selinger for (i = 0; i < am; i++) { 1163a9d7e0ccSandi selinger aopnz = aopothi[i+1] - aopothi[i]; 1164a9d7e0ccSandi selinger adponz = adpoi[i+1] - adpoi[i]; 1165a9d7e0ccSandi selinger adpdnz = adpdi[i+1] - adpdi[i]; 1166a9d7e0ccSandi selinger 1167a9d7e0ccSandi selinger /* Correct indices from A_diag*P_diag */ 1168a9d7e0ccSandi selinger for (i1 = 0; i1 < adpdnz; i1++) { 1169a016958eSandi selinger adpdJ[i1] += p_colstart; 1170a9d7e0ccSandi selinger } 1171a9d7e0ccSandi selinger /* Merge j-arrays of A_diag * P_loc_off and A_diag * P_loc_diag and A_off * P_oth */ 1172a9d7e0ccSandi selinger Merge3SortedArrays(adponz, adpoJ, adpdnz, adpdJ, aopnz, aopJ, &apnz, apJ); 11739566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+rstart, apnz, apJ, dnz, onz)); 1174a9d7e0ccSandi selinger 1175a9d7e0ccSandi selinger aopJ += aopnz; 1176a9d7e0ccSandi selinger adpoJ += adponz; 1177a9d7e0ccSandi selinger adpdJ += adpdnz; 1178a9d7e0ccSandi selinger apJ += apnz; 1179a9d7e0ccSandi selinger api[i+1] = api[i] + apnz; 1180a9d7e0ccSandi selinger } 1181a9d7e0ccSandi selinger 1182a9d7e0ccSandi selinger /* malloc apa to store dense row A[i,:]*P */ 11839566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(pN+2,&ptap->apa)); 1184a9d7e0ccSandi selinger 11854222ddf1SHong Zhang /* create and assemble symbolic parallel matrix C */ 11869566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,am,pn,PETSC_DETERMINE,PETSC_DETERMINE)); 11879566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,P)); 11889566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 11899566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 11909566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 1191d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 119217f75172Sandi selinger 11939566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api)); 11949566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 11959566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 11969566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 119717f75172Sandi selinger 11984222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable; 11994222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AB; 1200a9d7e0ccSandi selinger 12014222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 12026718818eSStefano Zampini C->product->data = ptap; 12036718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_MatMatMult; 1204a9d7e0ccSandi selinger 1205a9d7e0ccSandi selinger /* set MatInfo */ 120673b67375Sandi selinger afill = (PetscReal)api[am]/(adi[am]+aoi[am]+pi_loc[pm]+1) + 1.e-5; 120773b67375Sandi selinger if (afill < 1.0) afill = 1.0; 12084222ddf1SHong Zhang C->info.mallocs = nspacedouble; 12094222ddf1SHong Zhang C->info.fill_ratio_given = fill; 12104222ddf1SHong Zhang C->info.fill_ratio_needed = afill; 1211a9d7e0ccSandi selinger 1212a9d7e0ccSandi selinger #if defined(PETSC_USE_INFO) 1213a9d7e0ccSandi selinger if (api[am]) { 12149566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 12159566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n",(double)afill)); 1216a9d7e0ccSandi selinger } else { 12179566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 1218a9d7e0ccSandi selinger } 1219a9d7e0ccSandi selinger #endif 1220a9d7e0ccSandi selinger 12219566063dSJacob Faibussowitsch PetscCall(MatDestroy(&aopoth)); 12229566063dSJacob Faibussowitsch PetscCall(MatDestroy(&adpd)); 12239566063dSJacob Faibussowitsch PetscCall(PetscFree(j_temp)); 12249566063dSJacob Faibussowitsch PetscCall(PetscFree(adpoj)); 12259566063dSJacob Faibussowitsch PetscCall(PetscFree(adpoi)); 1226a9d7e0ccSandi selinger PetscFunctionReturn(0); 1227a9d7e0ccSandi selinger } 1228a9d7e0ccSandi selinger 1229f96d37fcSHong Zhang /*-------------------------------------------------------------------------*/ 1230c216dbf3SHong Zhang /* This routine only works when scall=MAT_REUSE_MATRIX! */ 1231c216dbf3SHong Zhang PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_matmatmult(Mat P,Mat A,Mat C) 1232c216dbf3SHong Zhang { 12336718818eSStefano Zampini Mat_APMPI *ptap; 123460552ceaSHong Zhang Mat Pt; 1235c216dbf3SHong Zhang 1236c216dbf3SHong Zhang PetscFunctionBegin; 12376718818eSStefano Zampini MatCheckProduct(C,3); 12386718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 123928b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 124028b400f6SJacob Faibussowitsch PetscCheck(ptap->Pt,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be reused. Do not call MatProductClear()"); 124160552ceaSHong Zhang 124260552ceaSHong Zhang Pt = ptap->Pt; 12439566063dSJacob Faibussowitsch PetscCall(MatTranspose(P,MAT_REUSE_MATRIX,&Pt)); 12449566063dSJacob Faibussowitsch PetscCall(MatMatMultNumeric_MPIAIJ_MPIAIJ(Pt,A,C)); 1245f96d37fcSHong Zhang PetscFunctionReturn(0); 1246f96d37fcSHong Zhang } 1247f96d37fcSHong Zhang 1248c5af039cSHong Zhang /* This routine is modified from MatPtAPSymbolic_MPIAIJ_MPIAIJ() */ 12494222ddf1SHong Zhang PetscErrorCode MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(Mat P,Mat A,PetscReal fill,Mat C) 1250912837bbSHong Zhang { 12513cdca5ebSHong Zhang Mat_APMPI *ptap; 12526718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data; 1253912837bbSHong Zhang MPI_Comm comm; 1254912837bbSHong Zhang PetscMPIInt size,rank; 1255912837bbSHong Zhang PetscFreeSpaceList free_space=NULL,current_space=NULL; 12564ede91e3SHong Zhang PetscInt pn=P->cmap->n,aN=A->cmap->N,an=A->cmap->n; 125769be949aSStefano Zampini PetscInt *lnk,i,k,nsend,rstart; 1258912837bbSHong Zhang PetscBT lnkbt; 1259ec4bef21SJose E. Roman PetscMPIInt tagi,tagj,*len_si,*len_s,*len_ri,nrecv; 1260ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted=0; 12618c78258cSHong Zhang PetscInt **buf_rj,**buf_ri,**buf_ri_k,row,ncols,*cols; 12627f3e1f30SHong Zhang PetscInt len,proc,*dnz,*onz,*owners,nzi; 1263912837bbSHong Zhang PetscInt nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci; 1264912837bbSHong Zhang MPI_Request *swaits,*rwaits; 1265912837bbSHong Zhang MPI_Status *sstatus,rstatus; 1266912837bbSHong Zhang PetscLayout rowmap; 1267912837bbSHong Zhang PetscInt *owners_co,*coi,*coj; /* i and j array of (p->B)^T*A*P - used in the communication */ 1268912837bbSHong Zhang PetscMPIInt *len_r,*id_r; /* array of length of comm->size, store send/recv matrix values */ 12697f3e1f30SHong Zhang PetscInt *Jptr,*prmap=p->garray,con,j,Crmax; 12704ede91e3SHong Zhang Mat_SeqAIJ *a_loc,*c_loc,*c_oth; 1271912837bbSHong Zhang PetscTable ta; 1272b92f168fSBarry Smith MatType mtype; 1273e83e5f86SFande Kong const char *prefix; 1274912837bbSHong Zhang 1275912837bbSHong Zhang PetscFunctionBegin; 12769566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 12779566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 12789566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1279912837bbSHong Zhang 12804222ddf1SHong Zhang /* create symbolic parallel matrix C */ 12819566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 12829566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 1283912837bbSHong Zhang 12844222ddf1SHong Zhang C->ops->transposematmultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable; 1285912837bbSHong Zhang 12863cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 12879566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 1288912837bbSHong Zhang ptap->reuse = MAT_INITIAL_MATRIX; 1289912837bbSHong Zhang 1290912837bbSHong Zhang /* (0) compute Rd = Pd^T, Ro = Po^T */ 1291912837bbSHong Zhang /* --------------------------------- */ 12929566063dSJacob Faibussowitsch PetscCall(MatTranspose_SeqAIJ(p->A,MAT_INITIAL_MATRIX,&ptap->Rd)); 12939566063dSJacob Faibussowitsch PetscCall(MatTranspose_SeqAIJ(p->B,MAT_INITIAL_MATRIX,&ptap->Ro)); 1294912837bbSHong Zhang 1295669e9c86SHong Zhang /* (1) compute symbolic A_loc */ 1296669e9c86SHong Zhang /* ---------------------------*/ 12979566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_INITIAL_MATRIX,&ptap->A_loc)); 1298912837bbSHong Zhang 12994ede91e3SHong Zhang /* (2-1) compute symbolic C_oth = Ro*A_loc */ 1300912837bbSHong Zhang /* ------------------------------------ */ 13019566063dSJacob Faibussowitsch PetscCall(MatGetOptionsPrefix(A,&prefix)); 13029566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(ptap->Ro,prefix)); 13039566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(ptap->Ro,"inner_offdiag_")); 13049566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF,&ptap->C_oth)); 13059566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Ro,ptap->A_loc,fill,ptap->C_oth)); 1306912837bbSHong Zhang 1307912837bbSHong Zhang /* (3) send coj of C_oth to other processors */ 1308912837bbSHong Zhang /* ------------------------------------------ */ 1309912837bbSHong Zhang /* determine row ownership */ 13109566063dSJacob Faibussowitsch PetscCall(PetscLayoutCreate(comm,&rowmap)); 1311912837bbSHong Zhang rowmap->n = pn; 1312912837bbSHong Zhang rowmap->bs = 1; 13139566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(rowmap)); 1314912837bbSHong Zhang owners = rowmap->range; 1315912837bbSHong Zhang 1316912837bbSHong Zhang /* determine the number of messages to send, their lengths */ 13179566063dSJacob Faibussowitsch PetscCall(PetscMalloc4(size,&len_s,size,&len_si,size,&sstatus,size+2,&owners_co)); 13189566063dSJacob Faibussowitsch PetscCall(PetscArrayzero(len_s,size)); 13199566063dSJacob Faibussowitsch PetscCall(PetscArrayzero(len_si,size)); 1320912837bbSHong Zhang 1321912837bbSHong Zhang c_oth = (Mat_SeqAIJ*)ptap->C_oth->data; 1322912837bbSHong Zhang coi = c_oth->i; coj = c_oth->j; 1323912837bbSHong Zhang con = ptap->C_oth->rmap->n; 1324912837bbSHong Zhang proc = 0; 1325912837bbSHong Zhang for (i=0; i<con; i++) { 1326912837bbSHong Zhang while (prmap[i] >= owners[proc+1]) proc++; 13277f3e1f30SHong Zhang len_si[proc]++; /* num of rows in Co(=Pt*A) to be sent to [proc] */ 1328912837bbSHong Zhang len_s[proc] += coi[i+1] - coi[i]; /* num of nonzeros in Co to be sent to [proc] */ 1329912837bbSHong Zhang } 1330912837bbSHong Zhang 1331912837bbSHong Zhang len = 0; /* max length of buf_si[], see (4) */ 1332912837bbSHong Zhang owners_co[0] = 0; 1333912837bbSHong Zhang nsend = 0; 1334912837bbSHong Zhang for (proc=0; proc<size; proc++) { 1335912837bbSHong Zhang owners_co[proc+1] = owners_co[proc] + len_si[proc]; 1336912837bbSHong Zhang if (len_s[proc]) { 1337912837bbSHong Zhang nsend++; 1338912837bbSHong Zhang len_si[proc] = 2*(len_si[proc] + 1); /* length of buf_si to be sent to [proc] */ 1339912837bbSHong Zhang len += len_si[proc]; 1340912837bbSHong Zhang } 1341912837bbSHong Zhang } 1342912837bbSHong Zhang 1343912837bbSHong Zhang /* determine the number and length of messages to receive for coi and coj */ 13449566063dSJacob Faibussowitsch PetscCall(PetscGatherNumberOfMessages(comm,NULL,len_s,&nrecv)); 13459566063dSJacob Faibussowitsch PetscCall(PetscGatherMessageLengths2(comm,nsend,nrecv,len_s,len_si,&id_r,&len_r,&len_ri)); 1346912837bbSHong Zhang 1347912837bbSHong Zhang /* post the Irecv and Isend of coj */ 13489566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagj)); 13499566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagj,nrecv,id_r,len_r,&buf_rj,&rwaits)); 13509566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(nsend+1,&swaits)); 1351912837bbSHong Zhang for (proc=0, k=0; proc<size; proc++) { 1352912837bbSHong Zhang if (!len_s[proc]) continue; 1353912837bbSHong Zhang i = owners_co[proc]; 13549566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tagj,comm,swaits+k)); 1355912837bbSHong Zhang k++; 1356912837bbSHong Zhang } 1357912837bbSHong Zhang 13587f3e1f30SHong Zhang /* (2-2) compute symbolic C_loc = Rd*A_loc */ 1359912837bbSHong Zhang /* ---------------------------------------- */ 13609566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(ptap->Rd,prefix)); 13619566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(ptap->Rd,"inner_diag_")); 13629566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF,&ptap->C_loc)); 13639566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Rd,ptap->A_loc,fill,ptap->C_loc)); 1364912837bbSHong Zhang c_loc = (Mat_SeqAIJ*)ptap->C_loc->data; 1365912837bbSHong Zhang 1366912837bbSHong Zhang /* receives coj are complete */ 1367912837bbSHong Zhang for (i=0; i<nrecv; i++) { 13689566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(nrecv,rwaits,&icompleted,&rstatus)); 1369912837bbSHong Zhang } 13709566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 13719566063dSJacob Faibussowitsch if (nsend) PetscCallMPI(MPI_Waitall(nsend,swaits,sstatus)); 1372912837bbSHong Zhang 1373912837bbSHong Zhang /* add received column indices into ta to update Crmax */ 13744ede91e3SHong Zhang a_loc = (Mat_SeqAIJ*)(ptap->A_loc)->data; 1375669e9c86SHong Zhang 1376669e9c86SHong Zhang /* create and initialize a linked list */ 13779566063dSJacob Faibussowitsch PetscCall(PetscTableCreate(an,aN,&ta)); /* for compute Crmax */ 13784ede91e3SHong Zhang MatRowMergeMax_SeqAIJ(a_loc,ptap->A_loc->rmap->N,ta); 1379669e9c86SHong Zhang 1380912837bbSHong Zhang for (k=0; k<nrecv; k++) {/* k-th received message */ 1381912837bbSHong Zhang Jptr = buf_rj[k]; 1382912837bbSHong Zhang for (j=0; j<len_r[k]; j++) { 13839566063dSJacob Faibussowitsch PetscCall(PetscTableAdd(ta,*(Jptr+j)+1,1,INSERT_VALUES)); 1384912837bbSHong Zhang } 1385912837bbSHong Zhang } 13869566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Crmax)); 13879566063dSJacob Faibussowitsch PetscCall(PetscTableDestroy(&ta)); 1388912837bbSHong Zhang 1389912837bbSHong Zhang /* (4) send and recv coi */ 1390912837bbSHong Zhang /*-----------------------*/ 13919566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagi)); 13929566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagi,nrecv,id_r,len_ri,&buf_ri,&rwaits)); 13939566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(len+1,&buf_s)); 1394912837bbSHong Zhang buf_si = buf_s; /* points to the beginning of k-th msg to be sent */ 1395912837bbSHong Zhang for (proc=0,k=0; proc<size; proc++) { 1396912837bbSHong Zhang if (!len_s[proc]) continue; 1397912837bbSHong Zhang /* form outgoing message for i-structure: 1398912837bbSHong Zhang buf_si[0]: nrows to be sent 1399912837bbSHong Zhang [1:nrows]: row index (global) 1400912837bbSHong Zhang [nrows+1:2*nrows+1]: i-structure index 1401912837bbSHong Zhang */ 1402912837bbSHong Zhang /*-------------------------------------------*/ 1403912837bbSHong Zhang nrows = len_si[proc]/2 - 1; /* num of rows in Co to be sent to [proc] */ 1404912837bbSHong Zhang buf_si_i = buf_si + nrows+1; 1405912837bbSHong Zhang buf_si[0] = nrows; 1406912837bbSHong Zhang buf_si_i[0] = 0; 1407912837bbSHong Zhang nrows = 0; 1408912837bbSHong Zhang for (i=owners_co[proc]; i<owners_co[proc+1]; i++) { 1409912837bbSHong Zhang nzi = coi[i+1] - coi[i]; 1410912837bbSHong Zhang buf_si_i[nrows+1] = buf_si_i[nrows] + nzi; /* i-structure */ 1411912837bbSHong Zhang buf_si[nrows+1] = prmap[i] -owners[proc]; /* local row index */ 1412912837bbSHong Zhang nrows++; 1413912837bbSHong Zhang } 14149566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tagi,comm,swaits+k)); 1415912837bbSHong Zhang k++; 1416912837bbSHong Zhang buf_si += len_si[proc]; 1417912837bbSHong Zhang } 1418912837bbSHong Zhang for (i=0; i<nrecv; i++) { 14199566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(nrecv,rwaits,&icompleted,&rstatus)); 1420912837bbSHong Zhang } 14219566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 14229566063dSJacob Faibussowitsch if (nsend) PetscCallMPI(MPI_Waitall(nsend,swaits,sstatus)); 1423912837bbSHong Zhang 14249566063dSJacob Faibussowitsch PetscCall(PetscFree4(len_s,len_si,sstatus,owners_co)); 14259566063dSJacob Faibussowitsch PetscCall(PetscFree(len_ri)); 14269566063dSJacob Faibussowitsch PetscCall(PetscFree(swaits)); 14279566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_s)); 1428912837bbSHong Zhang 14294222ddf1SHong Zhang /* (5) compute the local portion of C */ 1430912837bbSHong Zhang /* ------------------------------------------ */ 14314222ddf1SHong Zhang /* set initial free space to be Crmax, sufficient for holding nozeros in each row of C */ 14329566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(Crmax,&free_space)); 1433912837bbSHong Zhang current_space = free_space; 1434912837bbSHong Zhang 14359566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(nrecv,&buf_ri_k,nrecv,&nextrow,nrecv,&nextci)); 1436912837bbSHong Zhang for (k=0; k<nrecv; k++) { 1437912837bbSHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 1438912837bbSHong Zhang nrows = *buf_ri_k[k]; 1439912837bbSHong Zhang nextrow[k] = buf_ri_k[k] + 1; /* next row number of k-th recved i-structure */ 1440a5b23f4aSJose E. Roman nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th recved i-structure */ 1441912837bbSHong Zhang } 1442912837bbSHong Zhang 1443d0609cedSBarry Smith MatPreallocateBegin(comm,pn,an,dnz,onz); 14449566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate(Crmax,aN,&lnk,&lnkbt)); 14458c78258cSHong Zhang for (i=0; i<pn; i++) { /* for each local row of C */ 14464222ddf1SHong Zhang /* add C_loc into C */ 1447912837bbSHong Zhang nzi = c_loc->i[i+1] - c_loc->i[i]; 1448912837bbSHong Zhang Jptr = c_loc->j + c_loc->i[i]; 14499566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(nzi,Jptr,lnk,lnkbt)); 1450912837bbSHong Zhang 1451912837bbSHong Zhang /* add received col data into lnk */ 1452912837bbSHong Zhang for (k=0; k<nrecv; k++) { /* k-th received message */ 1453912837bbSHong Zhang if (i == *nextrow[k]) { /* i-th row */ 1454912837bbSHong Zhang nzi = *(nextci[k]+1) - *nextci[k]; 1455912837bbSHong Zhang Jptr = buf_rj[k] + *nextci[k]; 14569566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(nzi,Jptr,lnk,lnkbt)); 1457912837bbSHong Zhang nextrow[k]++; nextci[k]++; 1458912837bbSHong Zhang } 1459912837bbSHong Zhang } 14608c78258cSHong Zhang 14618c78258cSHong Zhang /* add missing diagonal entry */ 14628c78258cSHong Zhang if (C->force_diagonals) { 14638c78258cSHong Zhang k = i + owners[rank]; /* column index */ 14649566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(1,&k,lnk,lnkbt)); 14658c78258cSHong Zhang } 14668c78258cSHong Zhang 1467912837bbSHong Zhang nzi = lnk[0]; 1468912837bbSHong Zhang 1469912837bbSHong Zhang /* copy data into free space, then initialize lnk */ 14709566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean(aN,nzi,current_space->array,lnk,lnkbt)); 14719566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+owners[rank],nzi,current_space->array,dnz,onz)); 1472912837bbSHong Zhang } 14739566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 14749566063dSJacob Faibussowitsch PetscCall(PetscLLDestroy(lnk,lnkbt)); 14759566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceDestroy(free_space)); 1476912837bbSHong Zhang 1477912837bbSHong Zhang /* local sizes and preallocation */ 14789566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,pn,an,PETSC_DETERMINE,PETSC_DETERMINE)); 14799566063dSJacob Faibussowitsch if (P->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->rmap,P->cmap->bs)); 14809566063dSJacob Faibussowitsch if (A->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->cmap,A->cmap->bs)); 14819566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 1482d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 1483912837bbSHong Zhang 148469be949aSStefano Zampini /* add C_loc and C_oth to C */ 14859566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRange(C,&rstart,NULL)); 148669be949aSStefano Zampini for (i=0; i<pn; i++) { 14878c78258cSHong Zhang ncols = c_loc->i[i+1] - c_loc->i[i]; 14888c78258cSHong Zhang cols = c_loc->j + c_loc->i[i]; 14898c78258cSHong Zhang row = rstart + i; 14909566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,(const PetscInt*)&row,ncols,(const PetscInt*)cols,NULL,INSERT_VALUES)); 14918c78258cSHong Zhang 14928c78258cSHong Zhang if (C->force_diagonals) { 14939566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,(const PetscInt*)&row,1,(const PetscInt*)&row,NULL,INSERT_VALUES)); 14948c78258cSHong Zhang } 149569be949aSStefano Zampini } 149669be949aSStefano Zampini for (i=0; i<con; i++) { 14978c78258cSHong Zhang ncols = c_oth->i[i+1] - c_oth->i[i]; 14988c78258cSHong Zhang cols = c_oth->j + c_oth->i[i]; 14998c78258cSHong Zhang row = prmap[i]; 15009566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,(const PetscInt*)&row,ncols,(const PetscInt*)cols,NULL,INSERT_VALUES)); 150169be949aSStefano Zampini } 15029566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 15039566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 15049566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 150569be949aSStefano Zampini 1506912837bbSHong Zhang /* members in merge */ 15079566063dSJacob Faibussowitsch PetscCall(PetscFree(id_r)); 15089566063dSJacob Faibussowitsch PetscCall(PetscFree(len_r)); 15099566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_ri[0])); 15109566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_ri)); 15119566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_rj[0])); 15129566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_rj)); 15139566063dSJacob Faibussowitsch PetscCall(PetscLayoutDestroy(&rowmap)); 1514912837bbSHong Zhang 15154222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 15166718818eSStefano Zampini C->product->data = ptap; 15176718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 1518912837bbSHong Zhang PetscFunctionReturn(0); 1519912837bbSHong Zhang } 1520912837bbSHong Zhang 1521912837bbSHong Zhang PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable(Mat P,Mat A,Mat C) 1522912837bbSHong Zhang { 15236718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data; 1524669e9c86SHong Zhang Mat_SeqAIJ *c_seq; 15256718818eSStefano Zampini Mat_APMPI *ptap; 1526669e9c86SHong Zhang Mat A_loc,C_loc,C_oth; 1527912837bbSHong Zhang PetscInt i,rstart,rend,cm,ncols,row; 1528912837bbSHong Zhang const PetscInt *cols; 1529912837bbSHong Zhang const PetscScalar *vals; 1530912837bbSHong Zhang 1531912837bbSHong Zhang PetscFunctionBegin; 15326718818eSStefano Zampini MatCheckProduct(C,3); 15336718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 153428b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 153528b400f6SJacob Faibussowitsch PetscCheck(ptap->A_loc,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be reused. Do not call MatProductClear()"); 15369566063dSJacob Faibussowitsch PetscCall(MatZeroEntries(C)); 1537912837bbSHong Zhang 1538912837bbSHong Zhang if (ptap->reuse == MAT_REUSE_MATRIX) { 1539669e9c86SHong Zhang /* These matrices are obtained in MatTransposeMatMultSymbolic() */ 1540669e9c86SHong Zhang /* 1) get R = Pd^T, Ro = Po^T */ 1541669e9c86SHong Zhang /*----------------------------*/ 15429566063dSJacob Faibussowitsch PetscCall(MatTranspose_SeqAIJ(p->A,MAT_REUSE_MATRIX,&ptap->Rd)); 15439566063dSJacob Faibussowitsch PetscCall(MatTranspose_SeqAIJ(p->B,MAT_REUSE_MATRIX,&ptap->Ro)); 1544669e9c86SHong Zhang 1545669e9c86SHong Zhang /* 2) compute numeric A_loc */ 1546669e9c86SHong Zhang /*--------------------------*/ 15479566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_REUSE_MATRIX,&ptap->A_loc)); 1548912837bbSHong Zhang } 1549912837bbSHong Zhang 1550669e9c86SHong Zhang /* 3) C_loc = Rd*A_loc, C_oth = Ro*A_loc */ 15514ede91e3SHong Zhang A_loc = ptap->A_loc; 15529566063dSJacob Faibussowitsch PetscCall(((ptap->C_loc)->ops->matmultnumeric)(ptap->Rd,A_loc,ptap->C_loc)); 15539566063dSJacob Faibussowitsch PetscCall(((ptap->C_oth)->ops->matmultnumeric)(ptap->Ro,A_loc,ptap->C_oth)); 1554912837bbSHong Zhang C_loc = ptap->C_loc; 1555912837bbSHong Zhang C_oth = ptap->C_oth; 1556912837bbSHong Zhang 155769be949aSStefano Zampini /* add C_loc and C_oth to C */ 15589566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRange(C,&rstart,&rend)); 1559912837bbSHong Zhang 1560912837bbSHong Zhang /* C_loc -> C */ 1561912837bbSHong Zhang cm = C_loc->rmap->N; 1562912837bbSHong Zhang c_seq = (Mat_SeqAIJ*)C_loc->data; 1563912837bbSHong Zhang cols = c_seq->j; 1564912837bbSHong Zhang vals = c_seq->a; 1565912837bbSHong Zhang for (i=0; i<cm; i++) { 1566912837bbSHong Zhang ncols = c_seq->i[i+1] - c_seq->i[i]; 1567912837bbSHong Zhang row = rstart + i; 15689566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,ncols,cols,vals,ADD_VALUES)); 1569912837bbSHong Zhang cols += ncols; vals += ncols; 1570912837bbSHong Zhang } 1571912837bbSHong Zhang 1572912837bbSHong Zhang /* Co -> C, off-processor part */ 1573912837bbSHong Zhang cm = C_oth->rmap->N; 1574912837bbSHong Zhang c_seq = (Mat_SeqAIJ*)C_oth->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 = p->garray[i]; 15809566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,ncols,cols,vals,ADD_VALUES)); 1581912837bbSHong Zhang cols += ncols; vals += ncols; 1582912837bbSHong Zhang } 15839566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 15849566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 15859566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 1586912837bbSHong Zhang 1587912837bbSHong Zhang ptap->reuse = MAT_REUSE_MATRIX; 1588912837bbSHong Zhang PetscFunctionReturn(0); 1589912837bbSHong Zhang } 1590912837bbSHong Zhang 15916da69ca6SHong Zhang PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ(Mat P,Mat A,Mat C) 1592d6ab1dc0SHong Zhang { 1593d6ab1dc0SHong Zhang Mat_Merge_SeqsToMPI *merge; 15946718818eSStefano Zampini Mat_MPIAIJ *p =(Mat_MPIAIJ*)P->data; 1595d6ab1dc0SHong Zhang Mat_SeqAIJ *pd=(Mat_SeqAIJ*)(p->A)->data,*po=(Mat_SeqAIJ*)(p->B)->data; 15963cdca5ebSHong Zhang Mat_APMPI *ptap; 1597e745005bSHong Zhang PetscInt *adj; 1598e745005bSHong Zhang PetscInt i,j,k,anz,pnz,row,*cj,nexta; 1599e745005bSHong Zhang MatScalar *ada,*ca,valtmp; 1600d6ab1dc0SHong Zhang PetscInt am=A->rmap->n,cm=C->rmap->n,pon=(p->B)->cmap->n; 1601ce94432eSBarry Smith MPI_Comm comm; 1602d6ab1dc0SHong Zhang PetscMPIInt size,rank,taga,*len_s; 1603d6ab1dc0SHong Zhang PetscInt *owners,proc,nrows,**buf_ri_k,**nextrow,**nextci; 1604d6ab1dc0SHong Zhang PetscInt **buf_ri,**buf_rj; 1605d6ab1dc0SHong Zhang PetscInt cnz=0,*bj_i,*bi,*bj,bnz,nextcj; /* bi,bj,ba: local array of C(mpi mat) */ 1606d6ab1dc0SHong Zhang MPI_Request *s_waits,*r_waits; 1607d6ab1dc0SHong Zhang MPI_Status *status; 1608d6ab1dc0SHong Zhang MatScalar **abuf_r,*ba_i,*pA,*coa,*ba; 1609ce496241SStefano Zampini const PetscScalar *dummy; 16104e201bd7SHong Zhang PetscInt *ai,*aj,*coi,*coj,*poJ,*pdJ; 1611d6ab1dc0SHong Zhang Mat A_loc; 1612d6ab1dc0SHong Zhang Mat_SeqAIJ *a_loc; 1613d6ab1dc0SHong Zhang 1614d6ab1dc0SHong Zhang PetscFunctionBegin; 16156718818eSStefano Zampini MatCheckProduct(C,3); 16166718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 161728b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 161828b400f6SJacob Faibussowitsch PetscCheck(ptap->A_loc,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be reused. Do not call MatProductClear()"); 16199566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 16209566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 16219566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1622d6ab1dc0SHong Zhang 1623d6ab1dc0SHong Zhang merge = ptap->merge; 1624d6ab1dc0SHong Zhang 1625e745005bSHong Zhang /* 2) compute numeric C_seq = P_loc^T*A_loc */ 1626e745005bSHong Zhang /*------------------------------------------*/ 1627d6ab1dc0SHong Zhang /* get data from symbolic products */ 1628d6ab1dc0SHong Zhang coi = merge->coi; coj = merge->coj; 16299566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(coi[pon]+1,&coa)); 1630d6ab1dc0SHong Zhang bi = merge->bi; bj = merge->bj; 1631d6ab1dc0SHong Zhang owners = merge->rowmap->range; 16329566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(bi[cm]+1,&ba)); 1633d6ab1dc0SHong Zhang 1634d6ab1dc0SHong Zhang /* get A_loc by taking all local rows of A */ 1635d6ab1dc0SHong Zhang A_loc = ptap->A_loc; 16369566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_REUSE_MATRIX,&A_loc)); 1637d6ab1dc0SHong Zhang a_loc = (Mat_SeqAIJ*)(A_loc)->data; 1638d6ab1dc0SHong Zhang ai = a_loc->i; 1639d6ab1dc0SHong Zhang aj = a_loc->j; 1640d6ab1dc0SHong Zhang 1641ce496241SStefano Zampini /* trigger copy to CPU */ 16429566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(p->A,&dummy)); 16439566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(p->A,&dummy)); 16449566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(p->B,&dummy)); 16459566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(p->B,&dummy)); 1646d6ab1dc0SHong Zhang for (i=0; i<am; i++) { 1647d6ab1dc0SHong Zhang anz = ai[i+1] - ai[i]; 1648d6ab1dc0SHong Zhang adj = aj + ai[i]; 1649d6ab1dc0SHong Zhang ada = a_loc->a + ai[i]; 1650d6ab1dc0SHong Zhang 1651d6ab1dc0SHong Zhang /* 2-b) Compute Cseq = P_loc[i,:]^T*A[i,:] using outer product */ 1652e745005bSHong Zhang /*-------------------------------------------------------------*/ 1653d6ab1dc0SHong Zhang /* put the value into Co=(p->B)^T*A (off-diagonal part, send to others) */ 1654d6ab1dc0SHong Zhang pnz = po->i[i+1] - po->i[i]; 1655d6ab1dc0SHong Zhang poJ = po->j + po->i[i]; 1656d6ab1dc0SHong Zhang pA = po->a + po->i[i]; 1657d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1658d6ab1dc0SHong Zhang row = poJ[j]; 1659d6ab1dc0SHong Zhang cj = coj + coi[row]; 1660d6ab1dc0SHong Zhang ca = coa + coi[row]; 1661e745005bSHong Zhang /* perform sparse axpy */ 1662e745005bSHong Zhang nexta = 0; 1663d6ab1dc0SHong Zhang valtmp = pA[j]; 1664e745005bSHong Zhang for (k=0; nexta<anz; k++) { 1665e745005bSHong Zhang if (cj[k] == adj[nexta]) { 1666e745005bSHong Zhang ca[k] += valtmp*ada[nexta]; 1667e745005bSHong Zhang nexta++; 1668d6ab1dc0SHong Zhang } 1669e745005bSHong Zhang } 16709566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*anz)); 1671d6ab1dc0SHong Zhang } 1672d6ab1dc0SHong Zhang 1673d6ab1dc0SHong Zhang /* put the value into Cd (diagonal part) */ 1674d6ab1dc0SHong Zhang pnz = pd->i[i+1] - pd->i[i]; 1675d6ab1dc0SHong Zhang pdJ = pd->j + pd->i[i]; 1676d6ab1dc0SHong Zhang pA = pd->a + pd->i[i]; 1677d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1678d6ab1dc0SHong Zhang row = pdJ[j]; 1679d6ab1dc0SHong Zhang cj = bj + bi[row]; 1680d6ab1dc0SHong Zhang ca = ba + bi[row]; 1681e745005bSHong Zhang /* perform sparse axpy */ 1682e745005bSHong Zhang nexta = 0; 1683d6ab1dc0SHong Zhang valtmp = pA[j]; 1684e745005bSHong Zhang for (k=0; nexta<anz; k++) { 1685e745005bSHong Zhang if (cj[k] == adj[nexta]) { 1686e745005bSHong Zhang ca[k] += valtmp*ada[nexta]; 1687e745005bSHong Zhang nexta++; 1688d6ab1dc0SHong Zhang } 1689e745005bSHong Zhang } 16909566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*anz)); 1691d6ab1dc0SHong Zhang } 1692d6ab1dc0SHong Zhang } 1693d6ab1dc0SHong Zhang 1694d6ab1dc0SHong Zhang /* 3) send and recv matrix values coa */ 1695d6ab1dc0SHong Zhang /*------------------------------------*/ 1696d6ab1dc0SHong Zhang buf_ri = merge->buf_ri; 1697d6ab1dc0SHong Zhang buf_rj = merge->buf_rj; 1698d6ab1dc0SHong Zhang len_s = merge->len_s; 16999566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&taga)); 17009566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvScalar(comm,taga,merge->nrecv,merge->id_r,merge->len_r,&abuf_r,&r_waits)); 1701d6ab1dc0SHong Zhang 17029566063dSJacob Faibussowitsch PetscCall(PetscMalloc2(merge->nsend+1,&s_waits,size,&status)); 1703d6ab1dc0SHong Zhang for (proc=0,k=0; proc<size; proc++) { 1704d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1705d6ab1dc0SHong Zhang i = merge->owners_co[proc]; 17069566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coa+coi[i],len_s[proc],MPIU_MATSCALAR,proc,taga,comm,s_waits+k)); 1707d6ab1dc0SHong Zhang k++; 1708d6ab1dc0SHong Zhang } 17099566063dSJacob Faibussowitsch if (merge->nrecv) PetscCallMPI(MPI_Waitall(merge->nrecv,r_waits,status)); 17109566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,s_waits,status)); 1711d6ab1dc0SHong Zhang 17129566063dSJacob Faibussowitsch PetscCall(PetscFree2(s_waits,status)); 17139566063dSJacob Faibussowitsch PetscCall(PetscFree(r_waits)); 17149566063dSJacob Faibussowitsch PetscCall(PetscFree(coa)); 1715d6ab1dc0SHong Zhang 1716d6ab1dc0SHong Zhang /* 4) insert local Cseq and received values into Cmpi */ 1717d6ab1dc0SHong Zhang /*----------------------------------------------------*/ 17189566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(merge->nrecv,&buf_ri_k,merge->nrecv,&nextrow,merge->nrecv,&nextci)); 1719d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { 1720e745005bSHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 1721d6ab1dc0SHong Zhang nrows = *(buf_ri_k[k]); 1722d6ab1dc0SHong Zhang nextrow[k] = buf_ri_k[k]+1; /* next row number of k-th recved i-structure */ 1723a5b23f4aSJose E. Roman nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th recved i-structure */ 1724d6ab1dc0SHong Zhang } 1725d6ab1dc0SHong Zhang 1726d6ab1dc0SHong Zhang for (i=0; i<cm; i++) { 1727d6ab1dc0SHong Zhang row = owners[rank] + i; /* global row index of C_seq */ 1728d6ab1dc0SHong Zhang bj_i = bj + bi[i]; /* col indices of the i-th row of C */ 1729d6ab1dc0SHong Zhang ba_i = ba + bi[i]; 1730d6ab1dc0SHong Zhang bnz = bi[i+1] - bi[i]; 1731d6ab1dc0SHong Zhang /* add received vals into ba */ 1732d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { /* k-th received message */ 1733d6ab1dc0SHong Zhang /* i-th row */ 1734d6ab1dc0SHong Zhang if (i == *nextrow[k]) { 1735d6ab1dc0SHong Zhang cnz = *(nextci[k]+1) - *nextci[k]; 1736d6ab1dc0SHong Zhang cj = buf_rj[k] + *(nextci[k]); 1737d6ab1dc0SHong Zhang ca = abuf_r[k] + *(nextci[k]); 1738d6ab1dc0SHong Zhang nextcj = 0; 1739d6ab1dc0SHong Zhang for (j=0; nextcj<cnz; j++) { 1740d6ab1dc0SHong Zhang if (bj_i[j] == cj[nextcj]) { /* bcol == ccol */ 1741d6ab1dc0SHong Zhang ba_i[j] += ca[nextcj++]; 1742d6ab1dc0SHong Zhang } 1743d6ab1dc0SHong Zhang } 1744d6ab1dc0SHong Zhang nextrow[k]++; nextci[k]++; 17459566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*cnz)); 1746d6ab1dc0SHong Zhang } 1747d6ab1dc0SHong Zhang } 17489566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,bnz,bj_i,ba_i,INSERT_VALUES)); 1749d6ab1dc0SHong Zhang } 17509566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 17519566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 1752d6ab1dc0SHong Zhang 17539566063dSJacob Faibussowitsch PetscCall(PetscFree(ba)); 17549566063dSJacob Faibussowitsch PetscCall(PetscFree(abuf_r[0])); 17559566063dSJacob Faibussowitsch PetscCall(PetscFree(abuf_r)); 17569566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 1757d6ab1dc0SHong Zhang PetscFunctionReturn(0); 1758d6ab1dc0SHong Zhang } 1759d6ab1dc0SHong Zhang 17604222ddf1SHong Zhang PetscErrorCode MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ(Mat P,Mat A,PetscReal fill,Mat C) 1761d6ab1dc0SHong Zhang { 1762b419842dSStefano Zampini Mat A_loc; 17633cdca5ebSHong Zhang Mat_APMPI *ptap; 17640298fd71SBarry Smith PetscFreeSpaceList free_space=NULL,current_space=NULL; 17656718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data,*a=(Mat_MPIAIJ*)A->data; 1766d6ab1dc0SHong Zhang PetscInt *pdti,*pdtj,*poti,*potj,*ptJ; 1767d6ab1dc0SHong Zhang PetscInt nnz; 1768d6ab1dc0SHong Zhang PetscInt *lnk,*owners_co,*coi,*coj,i,k,pnz,row; 1769d6ab1dc0SHong Zhang PetscInt am =A->rmap->n,pn=P->cmap->n; 1770ce94432eSBarry Smith MPI_Comm comm; 1771d6ab1dc0SHong Zhang PetscMPIInt size,rank,tagi,tagj,*len_si,*len_s,*len_ri; 1772d6ab1dc0SHong Zhang PetscInt **buf_rj,**buf_ri,**buf_ri_k; 1773d6ab1dc0SHong Zhang PetscInt len,proc,*dnz,*onz,*owners; 1774d6ab1dc0SHong Zhang PetscInt nzi,*bi,*bj; 1775d6ab1dc0SHong Zhang PetscInt nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci; 1776d6ab1dc0SHong Zhang MPI_Request *swaits,*rwaits; 1777d6ab1dc0SHong Zhang MPI_Status *sstatus,rstatus; 1778d6ab1dc0SHong Zhang Mat_Merge_SeqsToMPI *merge; 1779d6ab1dc0SHong Zhang PetscInt *ai,*aj,*Jptr,anz,*prmap=p->garray,pon,nspacedouble=0,j; 1780d6ab1dc0SHong Zhang PetscReal afill =1.0,afill_tmp; 17814b38b95cSHong Zhang PetscInt rstart = P->cmap->rstart,rmax,aN=A->cmap->N,Armax; 1782b419842dSStefano Zampini Mat_SeqAIJ *a_loc; 17830acc65b4SHong Zhang PetscTable ta; 1784b92f168fSBarry Smith MatType mtype; 1785d6ab1dc0SHong Zhang 1786d6ab1dc0SHong Zhang PetscFunctionBegin; 17879566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 1788d6ab1dc0SHong Zhang /* check if matrix local sizes are compatible */ 1789*aed4548fSBarry 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); 1790d6ab1dc0SHong Zhang 17919566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 17929566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1793d6ab1dc0SHong Zhang 17943cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 17959566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 1796d6ab1dc0SHong Zhang 1797d6ab1dc0SHong Zhang /* get A_loc by taking all local rows of A */ 17989566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_INITIAL_MATRIX,&A_loc)); 17992205254eSKarl Rupp 1800d6ab1dc0SHong Zhang ptap->A_loc = A_loc; 1801d6ab1dc0SHong Zhang a_loc = (Mat_SeqAIJ*)(A_loc)->data; 1802d6ab1dc0SHong Zhang ai = a_loc->i; 1803d6ab1dc0SHong Zhang aj = a_loc->j; 1804d6ab1dc0SHong Zhang 1805d6ab1dc0SHong Zhang /* determine symbolic Co=(p->B)^T*A - send to others */ 1806d6ab1dc0SHong Zhang /*----------------------------------------------------*/ 18079566063dSJacob Faibussowitsch PetscCall(MatGetSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj)); 18089566063dSJacob Faibussowitsch PetscCall(MatGetSymbolicTranspose_SeqAIJ(p->B,&poti,&potj)); 1809d6ab1dc0SHong Zhang pon = (p->B)->cmap->n; /* total num of rows to be sent to other processors 1810d6ab1dc0SHong Zhang >= (num of nonzero rows of C_seq) - pn */ 18119566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pon+1,&coi)); 1812d6ab1dc0SHong Zhang coi[0] = 0; 1813d6ab1dc0SHong Zhang 1814d6ab1dc0SHong Zhang /* set initial free space to be fill*(nnz(p->B) + nnz(A)) */ 1815f91af8c7SBarry Smith nnz = PetscRealIntMultTruncate(fill,PetscIntSumTruncate(poti[pon],ai[am])); 18169566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(nnz,&free_space)); 1817d6ab1dc0SHong Zhang current_space = free_space; 181819f0e57aSHong Zhang 1819d6ab1dc0SHong Zhang /* create and initialize a linked list */ 18209566063dSJacob Faibussowitsch PetscCall(PetscTableCreate(A->cmap->n + a->B->cmap->N,aN,&ta)); 18214b38b95cSHong Zhang MatRowMergeMax_SeqAIJ(a_loc,am,ta); 18229566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Armax)); 182333ba5e2fSHong Zhang 18249566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(Armax,&lnk)); 1825d6ab1dc0SHong Zhang 1826d6ab1dc0SHong Zhang for (i=0; i<pon; i++) { 1827d6ab1dc0SHong Zhang pnz = poti[i+1] - poti[i]; 1828d6ab1dc0SHong Zhang ptJ = potj + poti[i]; 1829d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1830d6ab1dc0SHong Zhang row = ptJ[j]; /* row of A_loc == col of Pot */ 1831d6ab1dc0SHong Zhang anz = ai[row+1] - ai[row]; 1832d6ab1dc0SHong Zhang Jptr = aj + ai[row]; 1833d6ab1dc0SHong Zhang /* add non-zero cols of AP into the sorted linked list lnk */ 18349566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(anz,Jptr,lnk)); 1835d6ab1dc0SHong Zhang } 1836d6ab1dc0SHong Zhang nnz = lnk[0]; 1837d6ab1dc0SHong Zhang 1838d6ab1dc0SHong Zhang /* If free space is not available, double the total space in the list */ 1839d6ab1dc0SHong Zhang if (current_space->local_remaining<nnz) { 18409566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(nnz,current_space->total_array_size),¤t_space)); 1841d6ab1dc0SHong Zhang nspacedouble++; 1842d6ab1dc0SHong Zhang } 1843d6ab1dc0SHong Zhang 1844d6ab1dc0SHong Zhang /* Copy data into free space, and zero out denserows */ 18459566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(nnz,current_space->array,lnk)); 18462205254eSKarl Rupp 1847d6ab1dc0SHong Zhang current_space->array += nnz; 1848d6ab1dc0SHong Zhang current_space->local_used += nnz; 1849d6ab1dc0SHong Zhang current_space->local_remaining -= nnz; 18502205254eSKarl Rupp 1851d6ab1dc0SHong Zhang coi[i+1] = coi[i] + nnz; 1852d6ab1dc0SHong Zhang } 1853d6ab1dc0SHong Zhang 18549566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(coi[pon]+1,&coj)); 18559566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,coj)); 18569566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); /* must destroy to get a new one for C */ 18572205254eSKarl Rupp 1858118727c9SMark F. Adams afill_tmp = (PetscReal)coi[pon]/(poti[pon] + ai[am]+1); 1859d6ab1dc0SHong Zhang if (afill_tmp > afill) afill = afill_tmp; 1860d6ab1dc0SHong Zhang 1861d6ab1dc0SHong Zhang /* send j-array (coj) of Co to other processors */ 1862d6ab1dc0SHong Zhang /*----------------------------------------------*/ 1863d6ab1dc0SHong Zhang /* determine row ownership */ 18649566063dSJacob Faibussowitsch PetscCall(PetscNew(&merge)); 18659566063dSJacob Faibussowitsch PetscCall(PetscLayoutCreate(comm,&merge->rowmap)); 18662205254eSKarl Rupp 1867d6ab1dc0SHong Zhang merge->rowmap->n = pn; 1868d6ab1dc0SHong Zhang merge->rowmap->bs = 1; 18692205254eSKarl Rupp 18709566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(merge->rowmap)); 1871d6ab1dc0SHong Zhang owners = merge->rowmap->range; 1872d6ab1dc0SHong Zhang 1873d6ab1dc0SHong Zhang /* determine the number of messages to send, their lengths */ 18749566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(size,&len_si)); 18759566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(size,&merge->len_s)); 18762205254eSKarl Rupp 1877d6ab1dc0SHong Zhang len_s = merge->len_s; 1878d6ab1dc0SHong Zhang merge->nsend = 0; 1879d6ab1dc0SHong Zhang 18809566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size+2,&owners_co)); 1881d6ab1dc0SHong Zhang 1882d6ab1dc0SHong Zhang proc = 0; 1883d6ab1dc0SHong Zhang for (i=0; i<pon; i++) { 1884d6ab1dc0SHong Zhang while (prmap[i] >= owners[proc+1]) proc++; 1885d6ab1dc0SHong Zhang len_si[proc]++; /* num of rows in Co to be sent to [proc] */ 1886d6ab1dc0SHong Zhang len_s[proc] += coi[i+1] - coi[i]; 1887d6ab1dc0SHong Zhang } 1888d6ab1dc0SHong Zhang 1889d6ab1dc0SHong Zhang len = 0; /* max length of buf_si[] */ 1890d6ab1dc0SHong Zhang owners_co[0] = 0; 1891d6ab1dc0SHong Zhang for (proc=0; proc<size; proc++) { 1892d6ab1dc0SHong Zhang owners_co[proc+1] = owners_co[proc] + len_si[proc]; 1893d6ab1dc0SHong Zhang if (len_si[proc]) { 1894d6ab1dc0SHong Zhang merge->nsend++; 1895d6ab1dc0SHong Zhang len_si[proc] = 2*(len_si[proc] + 1); 1896d6ab1dc0SHong Zhang len += len_si[proc]; 1897d6ab1dc0SHong Zhang } 1898d6ab1dc0SHong Zhang } 1899d6ab1dc0SHong Zhang 1900d6ab1dc0SHong Zhang /* determine the number and length of messages to receive for coi and coj */ 19019566063dSJacob Faibussowitsch PetscCall(PetscGatherNumberOfMessages(comm,NULL,len_s,&merge->nrecv)); 19029566063dSJacob Faibussowitsch PetscCall(PetscGatherMessageLengths2(comm,merge->nsend,merge->nrecv,len_s,len_si,&merge->id_r,&merge->len_r,&len_ri)); 1903d6ab1dc0SHong Zhang 1904d6ab1dc0SHong Zhang /* post the Irecv and Isend of coj */ 19059566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagj)); 19069566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagj,merge->nrecv,merge->id_r,merge->len_r,&buf_rj,&rwaits)); 19079566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(merge->nsend+1,&swaits)); 1908d6ab1dc0SHong Zhang for (proc=0, k=0; proc<size; proc++) { 1909d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1910d6ab1dc0SHong Zhang i = owners_co[proc]; 19119566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tagj,comm,swaits+k)); 1912d6ab1dc0SHong Zhang k++; 1913d6ab1dc0SHong Zhang } 1914d6ab1dc0SHong Zhang 1915d6ab1dc0SHong Zhang /* receives and sends of coj are complete */ 19169566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size,&sstatus)); 1917d6ab1dc0SHong Zhang for (i=0; i<merge->nrecv; i++) { 1918ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted; 19199566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(merge->nrecv,rwaits,&icompleted,&rstatus)); 1920d6ab1dc0SHong Zhang } 19219566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 19229566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,swaits,sstatus)); 1923d6ab1dc0SHong Zhang 19240acc65b4SHong Zhang /* add received column indices into table to update Armax */ 1925c4762a1bSJed Brown /* Armax can be as large as aN if a P[row,:] is dense, see src/ksp/ksp/tutorials/ex56.c! */ 19260acc65b4SHong Zhang for (k=0; k<merge->nrecv; k++) {/* k-th received message */ 19270acc65b4SHong Zhang Jptr = buf_rj[k]; 19280acc65b4SHong Zhang for (j=0; j<merge->len_r[k]; j++) { 19299566063dSJacob Faibussowitsch PetscCall(PetscTableAdd(ta,*(Jptr+j)+1,1,INSERT_VALUES)); 19300acc65b4SHong Zhang } 19310acc65b4SHong Zhang } 19329566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Armax)); 19330acc65b4SHong Zhang 1934d6ab1dc0SHong Zhang /* send and recv coi */ 1935d6ab1dc0SHong Zhang /*-------------------*/ 19369566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagi)); 19379566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagi,merge->nrecv,merge->id_r,len_ri,&buf_ri,&rwaits)); 19389566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(len+1,&buf_s)); 1939d6ab1dc0SHong Zhang buf_si = buf_s; /* points to the beginning of k-th msg to be sent */ 1940d6ab1dc0SHong Zhang for (proc=0,k=0; proc<size; proc++) { 1941d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1942d6ab1dc0SHong Zhang /* form outgoing message for i-structure: 1943d6ab1dc0SHong Zhang buf_si[0]: nrows to be sent 1944d6ab1dc0SHong Zhang [1:nrows]: row index (global) 1945d6ab1dc0SHong Zhang [nrows+1:2*nrows+1]: i-structure index 1946d6ab1dc0SHong Zhang */ 1947d6ab1dc0SHong Zhang /*-------------------------------------------*/ 1948d6ab1dc0SHong Zhang nrows = len_si[proc]/2 - 1; 1949d6ab1dc0SHong Zhang buf_si_i = buf_si + nrows+1; 1950d6ab1dc0SHong Zhang buf_si[0] = nrows; 1951d6ab1dc0SHong Zhang buf_si_i[0] = 0; 1952d6ab1dc0SHong Zhang nrows = 0; 1953d6ab1dc0SHong Zhang for (i=owners_co[proc]; i<owners_co[proc+1]; i++) { 1954d6ab1dc0SHong Zhang nzi = coi[i+1] - coi[i]; 1955d6ab1dc0SHong Zhang buf_si_i[nrows+1] = buf_si_i[nrows] + nzi; /* i-structure */ 1956d6ab1dc0SHong Zhang buf_si[nrows+1] = prmap[i] -owners[proc]; /* local row index */ 1957d6ab1dc0SHong Zhang nrows++; 1958d6ab1dc0SHong Zhang } 19599566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tagi,comm,swaits+k)); 1960d6ab1dc0SHong Zhang k++; 1961d6ab1dc0SHong Zhang buf_si += len_si[proc]; 1962d6ab1dc0SHong Zhang } 1963d6ab1dc0SHong Zhang i = merge->nrecv; 1964d6ab1dc0SHong Zhang while (i--) { 1965ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted; 19669566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(merge->nrecv,rwaits,&icompleted,&rstatus)); 1967d6ab1dc0SHong Zhang } 19689566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 19699566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,swaits,sstatus)); 19709566063dSJacob Faibussowitsch PetscCall(PetscFree(len_si)); 19719566063dSJacob Faibussowitsch PetscCall(PetscFree(len_ri)); 19729566063dSJacob Faibussowitsch PetscCall(PetscFree(swaits)); 19739566063dSJacob Faibussowitsch PetscCall(PetscFree(sstatus)); 19749566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_s)); 1975d6ab1dc0SHong Zhang 1976d6ab1dc0SHong Zhang /* compute the local portion of C (mpi mat) */ 1977d6ab1dc0SHong Zhang /*------------------------------------------*/ 1978d6ab1dc0SHong Zhang /* allocate bi array and free space for accumulating nonzero column info */ 19799566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pn+1,&bi)); 1980d6ab1dc0SHong Zhang bi[0] = 0; 1981d6ab1dc0SHong Zhang 1982d6ab1dc0SHong Zhang /* set initial free space to be fill*(nnz(P) + nnz(AP)) */ 1983f91af8c7SBarry Smith nnz = PetscRealIntMultTruncate(fill,PetscIntSumTruncate(pdti[pn],PetscIntSumTruncate(poti[pon],ai[am]))); 19849566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(nnz,&free_space)); 1985d6ab1dc0SHong Zhang current_space = free_space; 1986d6ab1dc0SHong Zhang 19879566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(merge->nrecv,&buf_ri_k,merge->nrecv,&nextrow,merge->nrecv,&nextci)); 1988d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { 1989d6ab1dc0SHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 1990d6ab1dc0SHong Zhang nrows = *buf_ri_k[k]; 1991d6ab1dc0SHong Zhang nextrow[k] = buf_ri_k[k] + 1; /* next row number of k-th recved i-structure */ 19922d4ee042Sprj- nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th received i-structure */ 1993d6ab1dc0SHong Zhang } 1994d6ab1dc0SHong Zhang 19959566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(Armax,&lnk)); 1996d0609cedSBarry Smith MatPreallocateBegin(comm,pn,A->cmap->n,dnz,onz); 1997d6ab1dc0SHong Zhang rmax = 0; 1998d6ab1dc0SHong Zhang for (i=0; i<pn; i++) { 1999d6ab1dc0SHong Zhang /* add pdt[i,:]*AP into lnk */ 2000d6ab1dc0SHong Zhang pnz = pdti[i+1] - pdti[i]; 2001d6ab1dc0SHong Zhang ptJ = pdtj + pdti[i]; 2002d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 2003d6ab1dc0SHong Zhang row = ptJ[j]; /* row of AP == col of Pt */ 2004d6ab1dc0SHong Zhang anz = ai[row+1] - ai[row]; 2005d6ab1dc0SHong Zhang Jptr = aj + ai[row]; 2006d6ab1dc0SHong Zhang /* add non-zero cols of AP into the sorted linked list lnk */ 20079566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(anz,Jptr,lnk)); 2008d6ab1dc0SHong Zhang } 2009d6ab1dc0SHong Zhang 2010d6ab1dc0SHong Zhang /* add received col data into lnk */ 2011d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { /* k-th received message */ 2012d6ab1dc0SHong Zhang if (i == *nextrow[k]) { /* i-th row */ 2013d6ab1dc0SHong Zhang nzi = *(nextci[k]+1) - *nextci[k]; 2014d6ab1dc0SHong Zhang Jptr = buf_rj[k] + *nextci[k]; 20159566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(nzi,Jptr,lnk)); 2016d6ab1dc0SHong Zhang nextrow[k]++; nextci[k]++; 2017d6ab1dc0SHong Zhang } 2018d6ab1dc0SHong Zhang } 20198c78258cSHong Zhang 20208c78258cSHong Zhang /* add missing diagonal entry */ 20218c78258cSHong Zhang if (C->force_diagonals) { 20228c78258cSHong Zhang k = i + owners[rank]; /* column index */ 20239566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(1,&k,lnk)); 20248c78258cSHong Zhang } 20258c78258cSHong Zhang 2026d6ab1dc0SHong Zhang nnz = lnk[0]; 2027d6ab1dc0SHong Zhang 2028d6ab1dc0SHong Zhang /* if free space is not available, make more free space */ 2029d6ab1dc0SHong Zhang if (current_space->local_remaining<nnz) { 20309566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(nnz,current_space->total_array_size),¤t_space)); 2031d6ab1dc0SHong Zhang nspacedouble++; 2032d6ab1dc0SHong Zhang } 2033d6ab1dc0SHong Zhang /* copy data into free space, then initialize lnk */ 20349566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(nnz,current_space->array,lnk)); 20359566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+owners[rank],nnz,current_space->array,dnz,onz)); 20362205254eSKarl Rupp 2037d6ab1dc0SHong Zhang current_space->array += nnz; 2038d6ab1dc0SHong Zhang current_space->local_used += nnz; 2039d6ab1dc0SHong Zhang current_space->local_remaining -= nnz; 20402205254eSKarl Rupp 2041d6ab1dc0SHong Zhang bi[i+1] = bi[i] + nnz; 2042d6ab1dc0SHong Zhang if (nnz > rmax) rmax = nnz; 2043d6ab1dc0SHong Zhang } 20449566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 2045d6ab1dc0SHong Zhang 20469566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(bi[pn]+1,&bj)); 20479566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,bj)); 2048118727c9SMark F. Adams afill_tmp = (PetscReal)bi[pn]/(pdti[pn] + poti[pon] + ai[am]+1); 2049d6ab1dc0SHong Zhang if (afill_tmp > afill) afill = afill_tmp; 20509566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); 20519566063dSJacob Faibussowitsch PetscCall(PetscTableDestroy(&ta)); 20529566063dSJacob Faibussowitsch PetscCall(MatRestoreSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj)); 20539566063dSJacob Faibussowitsch PetscCall(MatRestoreSymbolicTranspose_SeqAIJ(p->B,&poti,&potj)); 2054d6ab1dc0SHong Zhang 20554222ddf1SHong Zhang /* create symbolic parallel matrix C - why cannot be assembled in Numeric part */ 20564222ddf1SHong Zhang /*-------------------------------------------------------------------------------*/ 20579566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,pn,A->cmap->n,PETSC_DETERMINE,PETSC_DETERMINE)); 20589566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizes(C,PetscAbs(P->cmap->bs),PetscAbs(A->cmap->bs))); 20599566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 20609566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 20619566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 2062d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 20639566063dSJacob Faibussowitsch PetscCall(MatSetBlockSize(C,1)); 20649566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NO_OFF_PROC_ENTRIES,PETSC_TRUE)); 2065d6ab1dc0SHong Zhang for (i=0; i<pn; i++) { 2066d6ab1dc0SHong Zhang row = i + rstart; 2067d6ab1dc0SHong Zhang nnz = bi[i+1] - bi[i]; 2068d6ab1dc0SHong Zhang Jptr = bj + bi[i]; 20699566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,nnz,Jptr,NULL,INSERT_VALUES)); 2070d6ab1dc0SHong Zhang } 20719566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 20729566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 20739566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 2074d6ab1dc0SHong Zhang merge->bi = bi; 2075d6ab1dc0SHong Zhang merge->bj = bj; 2076d6ab1dc0SHong Zhang merge->coi = coi; 2077d6ab1dc0SHong Zhang merge->coj = coj; 2078d6ab1dc0SHong Zhang merge->buf_ri = buf_ri; 2079d6ab1dc0SHong Zhang merge->buf_rj = buf_rj; 2080d6ab1dc0SHong Zhang merge->owners_co = owners_co; 2081d6ab1dc0SHong Zhang 20824222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 20836718818eSStefano Zampini C->product->data = ptap; 20846718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 2085d6ab1dc0SHong Zhang ptap->merge = merge; 2086a560ef98SHong Zhang 20874222ddf1SHong Zhang C->ops->mattransposemultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ; 2088f96d37fcSHong Zhang 2089f96d37fcSHong Zhang #if defined(PETSC_USE_INFO) 2090f96d37fcSHong Zhang if (bi[pn] != 0) { 20919566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 20929566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatTransposeMatMult(A,B,MatReuse,%g,&C) for best performance.\n",(double)afill)); 2093f96d37fcSHong Zhang } else { 20949566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 2095f96d37fcSHong Zhang } 2096f96d37fcSHong Zhang #endif 2097f96d37fcSHong Zhang PetscFunctionReturn(0); 2098f96d37fcSHong Zhang } 20994222ddf1SHong Zhang 21004222ddf1SHong Zhang /* ---------------------------------------------------------------- */ 21014222ddf1SHong Zhang static PetscErrorCode MatProductSymbolic_AtB_MPIAIJ_MPIAIJ(Mat C) 21024222ddf1SHong Zhang { 21034222ddf1SHong Zhang Mat_Product *product = C->product; 21044222ddf1SHong Zhang Mat A=product->A,B=product->B; 21054222ddf1SHong Zhang PetscReal fill=product->fill; 21064222ddf1SHong Zhang PetscBool flg; 21074222ddf1SHong Zhang 21084222ddf1SHong Zhang PetscFunctionBegin; 21094222ddf1SHong Zhang /* scalable */ 21109566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"scalable",&flg)); 21114222ddf1SHong Zhang if (flg) { 21129566063dSJacob Faibussowitsch PetscCall(MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ(A,B,fill,C)); 21134222ddf1SHong Zhang goto next; 21144222ddf1SHong Zhang } 21154222ddf1SHong Zhang 21164222ddf1SHong Zhang /* nonscalable */ 21179566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"nonscalable",&flg)); 21184222ddf1SHong Zhang if (flg) { 21199566063dSJacob Faibussowitsch PetscCall(MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(A,B,fill,C)); 21204222ddf1SHong Zhang goto next; 21214222ddf1SHong Zhang } 21224222ddf1SHong Zhang 21234222ddf1SHong Zhang /* matmatmult */ 21249566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"at*b",&flg)); 21254222ddf1SHong Zhang if (flg) { 21264222ddf1SHong Zhang Mat At; 21274222ddf1SHong Zhang Mat_APMPI *ptap; 21284222ddf1SHong Zhang 21299566063dSJacob Faibussowitsch PetscCall(MatTranspose(A,MAT_INITIAL_MATRIX,&At)); 21309566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ(At,B,fill,C)); 21316718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 21324222ddf1SHong Zhang if (ptap) { 21334222ddf1SHong Zhang ptap->Pt = At; 21346718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 21354222ddf1SHong Zhang } 21364222ddf1SHong Zhang C->ops->transposematmultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_matmatmult; 21374222ddf1SHong Zhang goto next; 21384222ddf1SHong Zhang } 21394222ddf1SHong Zhang 21404e84afc0SStefano Zampini /* backend general code */ 21419566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"backend",&flg)); 21424e84afc0SStefano Zampini if (flg) { 21439566063dSJacob Faibussowitsch PetscCall(MatProductSymbolic_MPIAIJBACKEND(C)); 21444e84afc0SStefano Zampini PetscFunctionReturn(0); 21454e84afc0SStefano Zampini } 21464e84afc0SStefano Zampini 21474222ddf1SHong Zhang SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"MatProduct type is not supported"); 21484222ddf1SHong Zhang 21494222ddf1SHong Zhang next: 21504222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AtB; 21514222ddf1SHong Zhang PetscFunctionReturn(0); 21524222ddf1SHong Zhang } 21534222ddf1SHong Zhang 21544222ddf1SHong Zhang /* ---------------------------------------------------------------- */ 21554222ddf1SHong Zhang /* Set options for MatMatMultxxx_MPIAIJ_MPIAIJ */ 21564222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_AB(Mat C) 21574222ddf1SHong Zhang { 21584222ddf1SHong Zhang Mat_Product *product = C->product; 21594222ddf1SHong Zhang Mat A=product->A,B=product->B; 21604222ddf1SHong Zhang #if defined(PETSC_HAVE_HYPRE) 21614e84afc0SStefano Zampini const char *algTypes[5] = {"scalable","nonscalable","seqmpi","backend","hypre"}; 21624e84afc0SStefano Zampini PetscInt nalg = 5; 21634222ddf1SHong Zhang #else 21644e84afc0SStefano Zampini const char *algTypes[4] = {"scalable","nonscalable","seqmpi","backend",}; 21654e84afc0SStefano Zampini PetscInt nalg = 4; 21664222ddf1SHong Zhang #endif 21674222ddf1SHong Zhang PetscInt alg = 1; /* set nonscalable algorithm as default */ 21684222ddf1SHong Zhang PetscBool flg; 21694222ddf1SHong Zhang MPI_Comm comm; 21704222ddf1SHong Zhang 21714222ddf1SHong Zhang PetscFunctionBegin; 21724222ddf1SHong Zhang /* Check matrix local sizes */ 21739566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 2174*aed4548fSBarry Smith PetscCheck(A->cmap->rstart == B->rmap->rstart && A->cmap->rend == B->rmap->rend,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); 21754222ddf1SHong Zhang 21764222ddf1SHong Zhang /* Set "nonscalable" as default algorithm */ 21779566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 21784222ddf1SHong Zhang if (flg) { 21799566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 21804222ddf1SHong Zhang 21814222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 21824222ddf1SHong Zhang if (B->cmap->N > 100000) { /* may switch to scalable algorithm as default */ 21834222ddf1SHong Zhang MatInfo Ainfo,Binfo; 21844222ddf1SHong Zhang PetscInt nz_local; 21854222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 21864222ddf1SHong Zhang 21879566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 21889566063dSJacob Faibussowitsch PetscCall(MatGetInfo(B,MAT_LOCAL,&Binfo)); 21894222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Binfo.nz_allocated); 21904222ddf1SHong Zhang 21914222ddf1SHong Zhang if (B->cmap->N > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 21921c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 21934222ddf1SHong Zhang 21944222ddf1SHong Zhang if (alg_scalable) { 21954222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 21969566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 21979566063dSJacob Faibussowitsch PetscCall(PetscInfo(B,"Use scalable algorithm, BN %" PetscInt_FMT ", fill*nz_allocated %g\n",B->cmap->N,(double)(product->fill*nz_local))); 21984222ddf1SHong Zhang } 21994222ddf1SHong Zhang } 22004222ddf1SHong Zhang } 22014222ddf1SHong Zhang 22024222ddf1SHong Zhang /* Get runtime option */ 22034222ddf1SHong Zhang if (product->api_user) { 2204d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatMatMult","Mat"); 22059566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matmatmult_via","Algorithmic approach","MatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2206d0609cedSBarry Smith PetscOptionsEnd(); 22074222ddf1SHong Zhang } else { 2208d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_AB","Mat"); 22099566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2210d0609cedSBarry Smith PetscOptionsEnd(); 22114222ddf1SHong Zhang } 22124222ddf1SHong Zhang if (flg) { 22139566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22144222ddf1SHong Zhang } 22154222ddf1SHong Zhang 22164222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AB_MPIAIJ_MPIAIJ; 22174222ddf1SHong Zhang PetscFunctionReturn(0); 22184222ddf1SHong Zhang } 22194222ddf1SHong Zhang 22204222ddf1SHong Zhang /* Set options for MatTransposeMatMultXXX_MPIAIJ_MPIAIJ */ 22214222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_AtB(Mat C) 22224222ddf1SHong Zhang { 22234222ddf1SHong Zhang Mat_Product *product = C->product; 22244222ddf1SHong Zhang Mat A=product->A,B=product->B; 22254e84afc0SStefano Zampini const char *algTypes[4] = {"scalable","nonscalable","at*b","backend"}; 22264e84afc0SStefano Zampini PetscInt nalg = 4; 22274222ddf1SHong Zhang PetscInt alg = 1; /* set default algorithm */ 22284222ddf1SHong Zhang PetscBool flg; 22294222ddf1SHong Zhang MPI_Comm comm; 22304222ddf1SHong Zhang 22314222ddf1SHong Zhang PetscFunctionBegin; 22324222ddf1SHong Zhang /* Check matrix local sizes */ 22339566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 2234*aed4548fSBarry 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); 22354222ddf1SHong Zhang 22364222ddf1SHong Zhang /* Set default algorithm */ 22379566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 22384222ddf1SHong Zhang if (flg) { 22399566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22404222ddf1SHong Zhang } 22414222ddf1SHong Zhang 22424222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 22434222ddf1SHong Zhang if (alg && B->cmap->N > 100000) { /* may switch to scalable algorithm as default */ 22444222ddf1SHong Zhang MatInfo Ainfo,Binfo; 22454222ddf1SHong Zhang PetscInt nz_local; 22464222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 22474222ddf1SHong Zhang 22489566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 22499566063dSJacob Faibussowitsch PetscCall(MatGetInfo(B,MAT_LOCAL,&Binfo)); 22504222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Binfo.nz_allocated); 22514222ddf1SHong Zhang 22524222ddf1SHong Zhang if (B->cmap->N > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 22531c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 22544222ddf1SHong Zhang 22554222ddf1SHong Zhang if (alg_scalable) { 22564222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 22579566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22589566063dSJacob Faibussowitsch PetscCall(PetscInfo(B,"Use scalable algorithm, BN %" PetscInt_FMT ", fill*nz_allocated %g\n",B->cmap->N,(double)(product->fill*nz_local))); 22594222ddf1SHong Zhang } 22604222ddf1SHong Zhang } 22614222ddf1SHong Zhang 22624222ddf1SHong Zhang /* Get runtime option */ 22634222ddf1SHong Zhang if (product->api_user) { 2264d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatTransposeMatMult","Mat"); 22659566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mattransposematmult_via","Algorithmic approach","MatTransposeMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2266d0609cedSBarry Smith PetscOptionsEnd(); 22674222ddf1SHong Zhang } else { 2268d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_AtB","Mat"); 22699566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatTransposeMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2270d0609cedSBarry Smith PetscOptionsEnd(); 22714222ddf1SHong Zhang } 22724222ddf1SHong Zhang if (flg) { 22739566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22744222ddf1SHong Zhang } 22754222ddf1SHong Zhang 22764222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AtB_MPIAIJ_MPIAIJ; 22774222ddf1SHong Zhang PetscFunctionReturn(0); 22784222ddf1SHong Zhang } 22794222ddf1SHong Zhang 22804222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_PtAP(Mat C) 22814222ddf1SHong Zhang { 22824222ddf1SHong Zhang Mat_Product *product = C->product; 22834222ddf1SHong Zhang Mat A=product->A,P=product->B; 22844222ddf1SHong Zhang MPI_Comm comm; 22854222ddf1SHong Zhang PetscBool flg; 22864222ddf1SHong Zhang PetscInt alg=1; /* set default algorithm */ 22874222ddf1SHong Zhang #if !defined(PETSC_HAVE_HYPRE) 22884e84afc0SStefano Zampini const char *algTypes[5] = {"scalable","nonscalable","allatonce","allatonce_merged","backend"}; 22894222ddf1SHong Zhang PetscInt nalg=5; 22904e84afc0SStefano Zampini #else 22914e84afc0SStefano Zampini const char *algTypes[6] = {"scalable","nonscalable","allatonce","allatonce_merged","backend","hypre"}; 22924e84afc0SStefano Zampini PetscInt nalg=6; 22934222ddf1SHong Zhang #endif 22944222ddf1SHong Zhang PetscInt pN=P->cmap->N; 22954222ddf1SHong Zhang 22964222ddf1SHong Zhang PetscFunctionBegin; 22974222ddf1SHong Zhang /* Check matrix local sizes */ 22989566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 2299*aed4548fSBarry 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); 2300*aed4548fSBarry 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); 23014222ddf1SHong Zhang 23024222ddf1SHong Zhang /* Set "nonscalable" as default algorithm */ 23039566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 23044222ddf1SHong Zhang if (flg) { 23059566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23064222ddf1SHong Zhang 23074222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 23084222ddf1SHong Zhang if (pN > 100000) { 23094222ddf1SHong Zhang MatInfo Ainfo,Pinfo; 23104222ddf1SHong Zhang PetscInt nz_local; 23114222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 23124222ddf1SHong Zhang 23139566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 23149566063dSJacob Faibussowitsch PetscCall(MatGetInfo(P,MAT_LOCAL,&Pinfo)); 23154222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Pinfo.nz_allocated); 23164222ddf1SHong Zhang 23174222ddf1SHong Zhang if (pN > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 23181c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 23194222ddf1SHong Zhang 23204222ddf1SHong Zhang if (alg_scalable) { 23214222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 23229566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23234222ddf1SHong Zhang } 23244222ddf1SHong Zhang } 23254222ddf1SHong Zhang } 23264222ddf1SHong Zhang 23274222ddf1SHong Zhang /* Get runtime option */ 23284222ddf1SHong Zhang if (product->api_user) { 2329d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatPtAP","Mat"); 23309566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matptap_via","Algorithmic approach","MatPtAP",algTypes,nalg,algTypes[alg],&alg,&flg)); 2331d0609cedSBarry Smith PetscOptionsEnd(); 23324222ddf1SHong Zhang } else { 2333d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_PtAP","Mat"); 23349566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatPtAP",algTypes,nalg,algTypes[alg],&alg,&flg)); 2335d0609cedSBarry Smith PetscOptionsEnd(); 23364222ddf1SHong Zhang } 23374222ddf1SHong Zhang if (flg) { 23389566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23394222ddf1SHong Zhang } 23404222ddf1SHong Zhang 23414222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_PtAP_MPIAIJ_MPIAIJ; 23424222ddf1SHong Zhang PetscFunctionReturn(0); 23434222ddf1SHong Zhang } 23444222ddf1SHong Zhang 23454222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_RARt(Mat C) 23464222ddf1SHong Zhang { 23474222ddf1SHong Zhang Mat_Product *product = C->product; 23484222ddf1SHong Zhang Mat A = product->A,R=product->B; 23494222ddf1SHong Zhang 23504222ddf1SHong Zhang PetscFunctionBegin; 23514222ddf1SHong Zhang /* Check matrix local sizes */ 2352*aed4548fSBarry 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); 23534222ddf1SHong Zhang 23544222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_RARt_MPIAIJ_MPIAIJ; 23554222ddf1SHong Zhang PetscFunctionReturn(0); 23564222ddf1SHong Zhang } 23574222ddf1SHong Zhang 23584222ddf1SHong Zhang /* 23594222ddf1SHong Zhang Set options for ABC = A*B*C = A*(B*C); ABC's algorithm must be chosen from AB's algorithm 23604222ddf1SHong Zhang */ 23614222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_ABC(Mat C) 23624222ddf1SHong Zhang { 23634222ddf1SHong Zhang Mat_Product *product = C->product; 23644222ddf1SHong Zhang PetscBool flg = PETSC_FALSE; 23654222ddf1SHong Zhang PetscInt alg = 1; /* default algorithm */ 23664222ddf1SHong Zhang const char *algTypes[3] = {"scalable","nonscalable","seqmpi"}; 23674222ddf1SHong Zhang PetscInt nalg = 3; 23684222ddf1SHong Zhang 23694222ddf1SHong Zhang PetscFunctionBegin; 23704222ddf1SHong Zhang /* Set default algorithm */ 23719566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 23724222ddf1SHong Zhang if (flg) { 23739566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23744222ddf1SHong Zhang } 23754222ddf1SHong Zhang 23764222ddf1SHong Zhang /* Get runtime option */ 23774222ddf1SHong Zhang if (product->api_user) { 2378d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatMatMatMult","Mat"); 23799566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matmatmatmult_via","Algorithmic approach","MatMatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2380d0609cedSBarry Smith PetscOptionsEnd(); 23814222ddf1SHong Zhang } else { 2382d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_ABC","Mat"); 23839566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatProduct_ABC",algTypes,nalg,algTypes[alg],&alg,&flg)); 2384d0609cedSBarry Smith PetscOptionsEnd(); 23854222ddf1SHong Zhang } 23864222ddf1SHong Zhang if (flg) { 23879566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23884222ddf1SHong Zhang } 23894222ddf1SHong Zhang 23904222ddf1SHong Zhang C->ops->matmatmultsymbolic = MatMatMatMultSymbolic_MPIAIJ_MPIAIJ_MPIAIJ; 23914222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_ABC; 23924222ddf1SHong Zhang PetscFunctionReturn(0); 23934222ddf1SHong Zhang } 23944222ddf1SHong Zhang 23954222ddf1SHong Zhang PETSC_INTERN PetscErrorCode MatProductSetFromOptions_MPIAIJ(Mat C) 23964222ddf1SHong Zhang { 23974222ddf1SHong Zhang Mat_Product *product = C->product; 23984222ddf1SHong Zhang 23994222ddf1SHong Zhang PetscFunctionBegin; 24004222ddf1SHong Zhang switch (product->type) { 24014222ddf1SHong Zhang case MATPRODUCT_AB: 24029566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_AB(C)); 24034222ddf1SHong Zhang break; 24044222ddf1SHong Zhang case MATPRODUCT_AtB: 24059566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_AtB(C)); 24064222ddf1SHong Zhang break; 24074222ddf1SHong Zhang case MATPRODUCT_PtAP: 24089566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_PtAP(C)); 24094222ddf1SHong Zhang break; 24104222ddf1SHong Zhang case MATPRODUCT_RARt: 24119566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_RARt(C)); 24124222ddf1SHong Zhang break; 24134222ddf1SHong Zhang case MATPRODUCT_ABC: 24149566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_ABC(C)); 24154222ddf1SHong Zhang break; 24166718818eSStefano Zampini default: 24176718818eSStefano Zampini break; 24184222ddf1SHong Zhang } 24194222ddf1SHong Zhang PetscFunctionReturn(0); 24204222ddf1SHong Zhang } 2421