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)); 1059467f45fSHong Zhang 1062c71b3e2SJacob Faibussowitsch PetscCheckFalse(!ptap->P_oth && size>1,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"AP cannot be reused. Do not call MatProductClear()"); 10780da8df7SHong Zhang 10865e4b4d4SStefano Zampini /* flag CPU mask for C */ 10965e4b4d4SStefano Zampini #if defined(PETSC_HAVE_DEVICE) 11065e4b4d4SStefano Zampini if (C->offloadmask != PETSC_OFFLOAD_UNALLOCATED) C->offloadmask = PETSC_OFFLOAD_CPU; 11165e4b4d4SStefano Zampini if (c->A->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->A->offloadmask = PETSC_OFFLOAD_CPU; 11265e4b4d4SStefano Zampini if (c->B->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->B->offloadmask = PETSC_OFFLOAD_CPU; 11365e4b4d4SStefano Zampini #endif 11465e4b4d4SStefano Zampini 115a1a4e84aSHong Zhang /* 1) get P_oth = ptap->P_oth and P_loc = ptap->P_loc */ 116598bc09dSHong Zhang /*-----------------------------------------------------*/ 117a1a4e84aSHong Zhang /* update numerical values of P_oth and P_loc */ 1189566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_REUSE_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 1199566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_REUSE_MATRIX,&ptap->P_loc)); 120598bc09dSHong Zhang 121598bc09dSHong Zhang /* 2) compute numeric C_loc = A_loc*P = Ad*P_loc + Ao*P_oth */ 122598bc09dSHong Zhang /*----------------------------------------------------------*/ 123598bc09dSHong Zhang /* get data from symbolic products */ 124a1a4e84aSHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 125c12557a1SHong Zhang p_oth = NULL; 1269467f45fSHong Zhang if (size >1) { 1279467f45fSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 1289467f45fSHong Zhang } 129598bc09dSHong Zhang 130598bc09dSHong Zhang /* get apa for storing dense row A[i,:]*P */ 131598bc09dSHong Zhang apa = ptap->apa; 132598bc09dSHong Zhang 13329825010SHong Zhang api = ptap->api; 13429825010SHong Zhang apj = ptap->apj; 135ce496241SStefano Zampini /* trigger copy to CPU */ 1369566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->A,&dummy)); 1379566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->A,&dummy)); 1389566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->B,&dummy)); 1399566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->B,&dummy)); 140598bc09dSHong Zhang for (i=0; i<cm; i++) { 141c12557a1SHong Zhang /* compute apa = A[i,:]*P */ 142e497d3c8SHong Zhang AProw_nonscalable(i,ad,ao,p_loc,p_oth,apa); 143598bc09dSHong Zhang 144598bc09dSHong Zhang /* set values in C */ 145598bc09dSHong Zhang apJ = apj + api[i]; 146598bc09dSHong Zhang cdnz = cd->i[i+1] - cd->i[i]; 147598bc09dSHong Zhang conz = co->i[i+1] - co->i[i]; 148598bc09dSHong Zhang 149910cf402Sprj- /* 1st off-diagonal part of C */ 150598bc09dSHong Zhang ca = coa + co->i[i]; 151598bc09dSHong Zhang k = 0; 152598bc09dSHong Zhang for (k0=0; k0<conz; k0++) { 153598bc09dSHong Zhang if (apJ[k] >= cstart) break; 154598bc09dSHong Zhang ca[k0] = apa[apJ[k]]; 1555cab4f04SHong Zhang apa[apJ[k++]] = 0.0; 156598bc09dSHong Zhang } 157598bc09dSHong Zhang 158598bc09dSHong Zhang /* diagonal part of C */ 159598bc09dSHong Zhang ca = cda + cd->i[i]; 160598bc09dSHong Zhang for (k1=0; k1<cdnz; k1++) { 161598bc09dSHong Zhang ca[k1] = apa[apJ[k]]; 1625cab4f04SHong Zhang apa[apJ[k++]] = 0.0; 163598bc09dSHong Zhang } 164598bc09dSHong Zhang 165910cf402Sprj- /* 2nd off-diagonal part of C */ 166598bc09dSHong Zhang ca = coa + co->i[i]; 167598bc09dSHong Zhang for (; k0<conz; k0++) { 168598bc09dSHong Zhang ca[k0] = apa[apJ[k]]; 1695cab4f04SHong Zhang apa[apJ[k++]] = 0.0; 170598bc09dSHong Zhang } 171598bc09dSHong Zhang } 1729566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 1739566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 174598bc09dSHong Zhang PetscFunctionReturn(0); 175598bc09dSHong Zhang } 176598bc09dSHong Zhang 1774222ddf1SHong Zhang PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(Mat A,Mat P,PetscReal fill,Mat C) 178598bc09dSHong Zhang { 179ce94432eSBarry Smith MPI_Comm comm; 1809467f45fSHong Zhang PetscMPIInt size; 1813cdca5ebSHong Zhang Mat_APMPI *ptap; 1820298fd71SBarry Smith PetscFreeSpaceList free_space=NULL,current_space=NULL; 1836718818eSStefano Zampini Mat_MPIAIJ *a=(Mat_MPIAIJ*)A->data; 184bfcd1a73SHong Zhang Mat_SeqAIJ *ad=(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_loc,*p_oth; 1857dbaa2c6Sandi selinger PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*dnz,*onz; 1864ae0847bSHong Zhang PetscInt *adi=ad->i,*adj=ad->j,*aoi=ao->i,*aoj=ao->j,rstart=A->rmap->rstart; 187d14fa243SHong Zhang PetscInt *lnk,i,pnz,row,*api,*apj,*Jptr,apnz,nspacedouble=0,j,nzi; 1885cab4f04SHong Zhang PetscInt am=A->rmap->n,pN=P->cmap->N,pn=P->cmap->n,pm=P->rmap->n; 189598bc09dSHong Zhang PetscBT lnkbt; 190bfcd1a73SHong Zhang PetscReal afill; 191b92f168fSBarry Smith MatType mtype; 192598bc09dSHong Zhang 193598bc09dSHong Zhang PetscFunctionBegin; 1946718818eSStefano Zampini MatCheckProduct(C,4); 19528b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Extra product struct not empty"); 1969566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 1979566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 1989467f45fSHong Zhang 1993cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 2009566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 201598bc09dSHong Zhang 202598bc09dSHong Zhang /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 2039566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 20419f0e57aSHong Zhang 205598bc09dSHong Zhang /* get P_loc by taking all local rows of P */ 2069566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc)); 207598bc09dSHong Zhang 208a1a4e84aSHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 209598bc09dSHong Zhang pi_loc = p_loc->i; pj_loc = p_loc->j; 2109467f45fSHong Zhang if (size > 1) { 2119467f45fSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 212598bc09dSHong Zhang pi_oth = p_oth->i; pj_oth = p_oth->j; 21320e3dc75SHong Zhang } else { 21420e3dc75SHong Zhang p_oth = NULL; 21520e3dc75SHong Zhang pi_oth = NULL; pj_oth = NULL; 2169467f45fSHong Zhang } 217598bc09dSHong Zhang 218598bc09dSHong Zhang /* first, compute symbolic AP = A_loc*P = A_diag*P_loc + A_off*P_oth */ 219598bc09dSHong Zhang /*-------------------------------------------------------------------*/ 2209566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&api)); 221a1a4e84aSHong Zhang ptap->api = api; 222598bc09dSHong Zhang api[0] = 0; 223598bc09dSHong Zhang 2245cab4f04SHong Zhang /* create and initialize a linked list */ 2259566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate(pN,pN,&lnk,&lnkbt)); 226598bc09dSHong Zhang 227bfcd1a73SHong Zhang /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */ 2289566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(adi[am],PetscIntSumTruncate(aoi[am],pi_loc[pm]))),&free_space)); 229598bc09dSHong Zhang current_space = free_space; 230598bc09dSHong Zhang 231*d0609cedSBarry Smith MatPreallocateBegin(comm,am,pn,dnz,onz); 232598bc09dSHong Zhang for (i=0; i<am; i++) { 233598bc09dSHong Zhang /* diagonal portion of A */ 234598bc09dSHong Zhang nzi = adi[i+1] - adi[i]; 235598bc09dSHong Zhang for (j=0; j<nzi; j++) { 236598bc09dSHong Zhang row = *adj++; 237598bc09dSHong Zhang pnz = pi_loc[row+1] - pi_loc[row]; 238598bc09dSHong Zhang Jptr = pj_loc + pi_loc[row]; 239598bc09dSHong Zhang /* add non-zero cols of P into the sorted linked list lnk */ 2409566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(pnz,Jptr,lnk,lnkbt)); 241598bc09dSHong Zhang } 242598bc09dSHong Zhang /* off-diagonal portion of A */ 243598bc09dSHong Zhang nzi = aoi[i+1] - aoi[i]; 244598bc09dSHong Zhang for (j=0; j<nzi; j++) { 245598bc09dSHong Zhang row = *aoj++; 246598bc09dSHong Zhang pnz = pi_oth[row+1] - pi_oth[row]; 247598bc09dSHong Zhang Jptr = pj_oth + pi_oth[row]; 2489566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(pnz,Jptr,lnk,lnkbt)); 249598bc09dSHong Zhang } 2508c78258cSHong Zhang /* add possible missing diagonal entry */ 2518c78258cSHong Zhang if (C->force_diagonals) { 2528c78258cSHong Zhang j = i + rstart; /* column index */ 2539566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(1,&j,lnk,lnkbt)); 2548c78258cSHong Zhang } 255598bc09dSHong Zhang 256d14fa243SHong Zhang apnz = lnk[0]; 257598bc09dSHong Zhang api[i+1] = api[i] + apnz; 258598bc09dSHong Zhang 259598bc09dSHong Zhang /* if free space is not available, double the total space in the list */ 260598bc09dSHong Zhang if (current_space->local_remaining<apnz) { 2619566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(apnz,current_space->total_array_size),¤t_space)); 262598bc09dSHong Zhang nspacedouble++; 263598bc09dSHong Zhang } 264598bc09dSHong Zhang 265598bc09dSHong Zhang /* Copy data into free space, then initialize lnk */ 2669566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean(pN,apnz,current_space->array,lnk,lnkbt)); 2679566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+rstart,apnz,current_space->array,dnz,onz)); 2682205254eSKarl Rupp 269598bc09dSHong Zhang current_space->array += apnz; 270598bc09dSHong Zhang current_space->local_used += apnz; 271598bc09dSHong Zhang current_space->local_remaining -= apnz; 272598bc09dSHong Zhang } 273598bc09dSHong Zhang 274598bc09dSHong Zhang /* Allocate space for apj, initialize apj, and */ 275598bc09dSHong Zhang /* destroy list of free space and other temporary array(s) */ 2769566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(api[am]+1,&ptap->apj)); 277a1a4e84aSHong Zhang apj = ptap->apj; 2789566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,ptap->apj)); 2799566063dSJacob Faibussowitsch PetscCall(PetscLLDestroy(lnk,lnkbt)); 280598bc09dSHong Zhang 281f84c536bSHong Zhang /* malloc apa to store dense row A[i,:]*P */ 2829566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(pN,&ptap->apa)); 283f84c536bSHong Zhang 2844222ddf1SHong Zhang /* set and assemble symbolic parallel matrix C */ 2854222ddf1SHong Zhang /*---------------------------------------------*/ 2869566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,am,pn,PETSC_DETERMINE,PETSC_DETERMINE)); 2879566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,P)); 288a2f3521dSMark F. Adams 2899566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 2909566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 2919566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 292*d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 29317f75172Sandi selinger 2949566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api)); 2959566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 2969566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 2979566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 298598bc09dSHong Zhang 2994222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable; 3004222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AB; 301598bc09dSHong Zhang 3024222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 3036718818eSStefano Zampini C->product->data = ptap; 3046718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_MatMatMult; 305598bc09dSHong Zhang 306bfcd1a73SHong Zhang /* set MatInfo */ 307118727c9SMark F. Adams afill = (PetscReal)api[am]/(adi[am]+aoi[am]+pi_loc[pm]+1) + 1.e-5; 308bfcd1a73SHong Zhang if (afill < 1.0) afill = 1.0; 3094222ddf1SHong Zhang C->info.mallocs = nspacedouble; 3104222ddf1SHong Zhang C->info.fill_ratio_given = fill; 3114222ddf1SHong Zhang C->info.fill_ratio_needed = afill; 312bfcd1a73SHong Zhang 313bfcd1a73SHong Zhang #if defined(PETSC_USE_INFO) 314bfcd1a73SHong Zhang if (api[am]) { 3159566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 3169566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n",(double)afill)); 317bfcd1a73SHong Zhang } else { 3189566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 319bfcd1a73SHong Zhang } 320bfcd1a73SHong Zhang #endif 321598bc09dSHong Zhang PetscFunctionReturn(0); 322598bc09dSHong Zhang } 323598bc09dSHong Zhang 3244222ddf1SHong Zhang /* ------------------------------------------------------- */ 3256718818eSStefano Zampini static PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIDense(Mat,Mat,PetscReal,Mat); 3266718818eSStefano Zampini static PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIDense(Mat,Mat,Mat); 3276718818eSStefano Zampini 3284222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense_AB(Mat C) 3299123193aSHong Zhang { 3304222ddf1SHong Zhang Mat_Product *product = C->product; 3314222ddf1SHong Zhang Mat A = product->A,B=product->B; 3329123193aSHong Zhang 3339123193aSHong Zhang PetscFunctionBegin; 3344222ddf1SHong Zhang if (A->cmap->rstart != B->rmap->rstart || A->cmap->rend != B->rmap->rend) 33598921bdaSJacob 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); 3364222ddf1SHong Zhang 3374222ddf1SHong Zhang C->ops->matmultsymbolic = MatMatMultSymbolic_MPIAIJ_MPIDense; 3384222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AB; 3399123193aSHong Zhang PetscFunctionReturn(0); 3409123193aSHong Zhang } 3414222ddf1SHong Zhang /* -------------------------------------------------------------------- */ 3424222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense_AtB(Mat C) 3434222ddf1SHong Zhang { 3444222ddf1SHong Zhang Mat_Product *product = C->product; 3454222ddf1SHong Zhang Mat A = product->A,B=product->B; 3464222ddf1SHong Zhang 3474222ddf1SHong Zhang PetscFunctionBegin; 3484222ddf1SHong Zhang if (A->rmap->rstart != B->rmap->rstart || A->rmap->rend != B->rmap->rend) 34998921bdaSJacob 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); 3504222ddf1SHong Zhang 3514222ddf1SHong Zhang C->ops->transposematmultsymbolic = MatTransposeMatMultSymbolic_MPIAIJ_MPIDense; 3524222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AtB; 3534222ddf1SHong Zhang PetscFunctionReturn(0); 3544222ddf1SHong Zhang } 3554222ddf1SHong Zhang 3564222ddf1SHong Zhang /* --------------------------------------------------------------------- */ 3574222ddf1SHong Zhang PETSC_INTERN PetscErrorCode MatProductSetFromOptions_MPIAIJ_MPIDense(Mat C) 3584222ddf1SHong Zhang { 3594222ddf1SHong Zhang Mat_Product *product = C->product; 3604222ddf1SHong Zhang 3614222ddf1SHong Zhang PetscFunctionBegin; 3624222ddf1SHong Zhang switch (product->type) { 3634222ddf1SHong Zhang case MATPRODUCT_AB: 3649566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_MPIDense_AB(C)); 3654222ddf1SHong Zhang break; 3664222ddf1SHong Zhang case MATPRODUCT_AtB: 3679566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_MPIDense_AtB(C)); 3684222ddf1SHong Zhang break; 3694222ddf1SHong Zhang default: 3706718818eSStefano Zampini break; 3714222ddf1SHong Zhang } 3724222ddf1SHong Zhang PetscFunctionReturn(0); 3734222ddf1SHong Zhang } 3744222ddf1SHong Zhang /* ------------------------------------------------------- */ 3759123193aSHong Zhang 3764324174eSBarry Smith typedef struct { 3775ea7661aSPierre Jolivet Mat workB,workB1; 3784324174eSBarry Smith MPI_Request *rwaits,*swaits; 37987753ddeSHong Zhang PetscInt nsends,nrecvs; 38087753ddeSHong Zhang MPI_Datatype *stype,*rtype; 3816718818eSStefano Zampini PetscInt blda; 3824324174eSBarry Smith } MPIAIJ_MPIDense; 3834324174eSBarry Smith 38496b95a6bSBarry Smith PetscErrorCode MatMPIAIJ_MPIDenseDestroy(void *ctx) 3854324174eSBarry Smith { 3864324174eSBarry Smith MPIAIJ_MPIDense *contents = (MPIAIJ_MPIDense*)ctx; 38787753ddeSHong Zhang PetscInt i; 3884324174eSBarry Smith 3894324174eSBarry Smith PetscFunctionBegin; 3909566063dSJacob Faibussowitsch PetscCall(MatDestroy(&contents->workB)); 3919566063dSJacob Faibussowitsch PetscCall(MatDestroy(&contents->workB1)); 39287753ddeSHong Zhang for (i=0; i<contents->nsends; i++) { 3939566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&contents->stype[i])); 39487753ddeSHong Zhang } 39587753ddeSHong Zhang for (i=0; i<contents->nrecvs; i++) { 3969566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&contents->rtype[i])); 39787753ddeSHong Zhang } 3989566063dSJacob Faibussowitsch PetscCall(PetscFree4(contents->stype,contents->rtype,contents->rwaits,contents->swaits)); 3999566063dSJacob Faibussowitsch PetscCall(PetscFree(contents)); 4004324174eSBarry Smith PetscFunctionReturn(0); 4014324174eSBarry Smith } 4024324174eSBarry Smith 4036718818eSStefano Zampini static PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIDense(Mat A,Mat B,PetscReal fill,Mat C) 4049123193aSHong Zhang { 405f32d5d43SBarry Smith Mat_MPIAIJ *aij=(Mat_MPIAIJ*)A->data; 406b45e3bf4SStefano Zampini PetscInt nz=aij->B->cmap->n,nsends,nrecvs,i,nrows_to,j,blda,m,M,n,N; 4074324174eSBarry Smith MPIAIJ_MPIDense *contents; 4084324174eSBarry Smith VecScatter ctx=aij->Mvctx; 4095ea7661aSPierre Jolivet PetscInt Am=A->rmap->n,Bm=B->rmap->n,BN=B->cmap->N,Bbn,Bbn1,bs,nrows_from,numBb; 41087753ddeSHong Zhang MPI_Comm comm; 41187753ddeSHong Zhang MPI_Datatype type1,*stype,*rtype; 41287753ddeSHong Zhang const PetscInt *sindices,*sstarts,*rstarts; 41387753ddeSHong Zhang PetscMPIInt *disp; 4146718818eSStefano Zampini PetscBool cisdense; 4159123193aSHong Zhang 4169123193aSHong Zhang PetscFunctionBegin; 4176718818eSStefano Zampini MatCheckProduct(C,4); 41828b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data not empty"); 4199566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 4209566063dSJacob Faibussowitsch PetscCall(PetscObjectBaseTypeCompare((PetscObject)C,MATMPIDENSE,&cisdense)); 4216718818eSStefano Zampini if (!cisdense) { 4229566063dSJacob Faibussowitsch PetscCall(MatSetType(C,((PetscObject)B)->type_name)); 4236718818eSStefano Zampini } 4249566063dSJacob Faibussowitsch PetscCall(MatGetLocalSize(C,&m,&n)); 4259566063dSJacob Faibussowitsch PetscCall(MatGetSize(C,&M,&N)); 426825bfe5fSPierre Jolivet if (m == PETSC_DECIDE || n == PETSC_DECIDE || M == PETSC_DECIDE || N == PETSC_DECIDE) { 4279566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,Am,B->cmap->n,A->rmap->N,BN)); 428825bfe5fSPierre Jolivet } 4299566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,B)); 4309566063dSJacob Faibussowitsch PetscCall(MatSetUp(C)); 4319566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(B,&blda)); 4329566063dSJacob Faibussowitsch PetscCall(PetscNew(&contents)); 4334324174eSBarry Smith 4349566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,NULL,NULL)); 4359566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,NULL,NULL)); 43687753ddeSHong Zhang 43787753ddeSHong Zhang /* Create column block of B and C for memory scalability when BN is too large */ 43887753ddeSHong Zhang /* Estimate Bbn, column size of Bb */ 43987753ddeSHong Zhang if (nz) { 44087753ddeSHong Zhang Bbn1 = 2*Am*BN/nz; 441da7cbbe7SHong Zhang if (!Bbn1) Bbn1 = 1; 44287753ddeSHong Zhang } else Bbn1 = BN; 44387753ddeSHong Zhang 44487753ddeSHong Zhang bs = PetscAbs(B->cmap->bs); 44587753ddeSHong Zhang Bbn1 = Bbn1/bs *bs; /* Bbn1 is a multiple of bs */ 44687753ddeSHong Zhang if (Bbn1 > BN) Bbn1 = BN; 4479566063dSJacob Faibussowitsch PetscCallMPI(MPI_Allreduce(&Bbn1,&Bbn,1,MPIU_INT,MPI_MAX,comm)); 44887753ddeSHong Zhang 44987753ddeSHong Zhang /* Enable runtime option for Bbn */ 450*d0609cedSBarry Smith PetscOptionsBegin(comm,((PetscObject)C)->prefix,"MatMatMult","Mat"); 4519566063dSJacob Faibussowitsch PetscCall(PetscOptionsInt("-matmatmult_Bbn","Number of columns in Bb","MatMatMult",Bbn,&Bbn,NULL)); 452*d0609cedSBarry Smith PetscOptionsEnd(); 4536718818eSStefano Zampini Bbn = PetscMin(Bbn,BN); 45487753ddeSHong Zhang 4556718818eSStefano Zampini if (Bbn > 0 && Bbn < BN) { 4565ea7661aSPierre Jolivet numBb = BN/Bbn; 4575ea7661aSPierre Jolivet Bbn1 = BN - numBb*Bbn; 4585ea7661aSPierre Jolivet } else numBb = 0; 45987753ddeSHong Zhang 4605ea7661aSPierre Jolivet if (numBb) { 4619566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"use Bb, BN=%" PetscInt_FMT ", Bbn=%" PetscInt_FMT "; numBb=%" PetscInt_FMT "\n",BN,Bbn,numBb)); 4625ea7661aSPierre Jolivet if (Bbn1) { /* Create workB1 for the remaining columns */ 4639566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"use Bb1, BN=%" PetscInt_FMT ", Bbn1=%" PetscInt_FMT "\n",BN,Bbn1)); 46487753ddeSHong Zhang /* Create work matrix used to store off processor rows of B needed for local product */ 4659566063dSJacob Faibussowitsch PetscCall(MatCreateSeqDense(PETSC_COMM_SELF,nz,Bbn1,NULL,&contents->workB1)); 4665ea7661aSPierre Jolivet } else contents->workB1 = NULL; 46787753ddeSHong Zhang } 46887753ddeSHong Zhang 46987753ddeSHong Zhang /* Create work matrix used to store off processor rows of B needed for local product */ 4709566063dSJacob Faibussowitsch PetscCall(MatCreateSeqDense(PETSC_COMM_SELF,nz,Bbn,NULL,&contents->workB)); 47187753ddeSHong Zhang 47287753ddeSHong Zhang /* Use MPI derived data type to reduce memory required by the send/recv buffers */ 4739566063dSJacob Faibussowitsch PetscCall(PetscMalloc4(nsends,&stype,nrecvs,&rtype,nrecvs,&contents->rwaits,nsends,&contents->swaits)); 47487753ddeSHong Zhang contents->stype = stype; 47587753ddeSHong Zhang contents->nsends = nsends; 47687753ddeSHong Zhang 47787753ddeSHong Zhang contents->rtype = rtype; 47887753ddeSHong Zhang contents->nrecvs = nrecvs; 4796718818eSStefano Zampini contents->blda = blda; 48087753ddeSHong Zhang 4819566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(Bm+1,&disp)); 48287753ddeSHong Zhang for (i=0; i<nsends; i++) { 48387753ddeSHong Zhang nrows_to = sstarts[i+1]-sstarts[i]; 484b45e3bf4SStefano Zampini for (j=0; j<nrows_to; j++) disp[j] = sindices[sstarts[i]+j]; /* rowB to be sent */ 4859566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_indexed_block(nrows_to,1,disp,MPIU_SCALAR,&type1)); 4869566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_resized(type1,0,blda*sizeof(PetscScalar),&stype[i])); 4879566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_commit(&stype[i])); 4889566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&type1)); 48987753ddeSHong Zhang } 49087753ddeSHong Zhang 49187753ddeSHong Zhang for (i=0; i<nrecvs; i++) { 49287753ddeSHong Zhang /* received values from a process form a (nrows_from x Bbn) row block in workB (column-wise) */ 49387753ddeSHong Zhang nrows_from = rstarts[i+1]-rstarts[i]; 49487753ddeSHong Zhang disp[0] = 0; 4959566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_indexed_block(1,nrows_from,disp,MPIU_SCALAR,&type1)); 4969566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_create_resized(type1,0,nz*sizeof(PetscScalar),&rtype[i])); 4979566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_commit(&rtype[i])); 4989566063dSJacob Faibussowitsch PetscCallMPI(MPI_Type_free(&type1)); 49987753ddeSHong Zhang } 50087753ddeSHong Zhang 5019566063dSJacob Faibussowitsch PetscCall(PetscFree(disp)); 5029566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,NULL,NULL)); 5039566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,NULL,NULL)); 5049566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NO_OFF_PROC_ENTRIES,PETSC_TRUE)); 5059566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 5069566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 5079566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 5086718818eSStefano Zampini 5096718818eSStefano Zampini C->product->data = contents; 5106718818eSStefano Zampini C->product->destroy = MatMPIAIJ_MPIDenseDestroy; 5114222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIDense; 5129123193aSHong Zhang PetscFunctionReturn(0); 5139123193aSHong Zhang } 5149123193aSHong Zhang 51593aa15f2SStefano Zampini PETSC_INTERN PetscErrorCode MatMatMultNumericAdd_SeqAIJ_SeqDense(Mat,Mat,Mat,const PetscBool); 516f32d5d43SBarry Smith /* 5172636ff26SBarry Smith Performs an efficient scatter on the rows of B needed by this process; this is 5182636ff26SBarry Smith a modification of the VecScatterBegin_() routines. 51987753ddeSHong Zhang 52087753ddeSHong Zhang Input: Bbidx = 0: B = Bb 52187753ddeSHong Zhang = 1: B = Bb1, see MatMatMultSymbolic_MPIAIJ_MPIDense() 522f32d5d43SBarry Smith */ 52387753ddeSHong Zhang PetscErrorCode MatMPIDenseScatter(Mat A,Mat B,PetscInt Bbidx,Mat C,Mat *outworkB) 524f32d5d43SBarry Smith { 525f32d5d43SBarry Smith Mat_MPIAIJ *aij = (Mat_MPIAIJ*)A->data; 5261683a169SBarry Smith const PetscScalar *b; 52787753ddeSHong Zhang PetscScalar *rvalues; 528f32d5d43SBarry Smith VecScatter ctx = aij->Mvctx; 52987753ddeSHong Zhang const PetscInt *sindices,*sstarts,*rstarts; 5303515ee7fSJunchao Zhang const PetscMPIInt *sprocs,*rprocs; 5316718818eSStefano Zampini PetscInt i,nsends,nrecvs; 532f32d5d43SBarry Smith MPI_Request *swaits,*rwaits; 533ce94432eSBarry Smith MPI_Comm comm; 5346718818eSStefano Zampini PetscMPIInt tag=((PetscObject)ctx)->tag,ncols=B->cmap->N,nrows=aij->B->cmap->n,nsends_mpi,nrecvs_mpi; 5354324174eSBarry Smith MPIAIJ_MPIDense *contents; 5364324174eSBarry Smith Mat workB; 53787753ddeSHong Zhang MPI_Datatype *stype,*rtype; 5386718818eSStefano Zampini PetscInt blda; 539f32d5d43SBarry Smith 540f32d5d43SBarry Smith PetscFunctionBegin; 5416718818eSStefano Zampini MatCheckProduct(C,4); 54228b400f6SJacob Faibussowitsch PetscCheck(C->product->data,PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data empty"); 5436718818eSStefano Zampini contents = (MPIAIJ_MPIDense*)C->product->data; 5449566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,&sprocs,NULL/*bs*/)); 5459566063dSJacob Faibussowitsch PetscCall(VecScatterGetRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,&rprocs,NULL/*bs*/)); 5469566063dSJacob Faibussowitsch PetscCall(PetscMPIIntCast(nsends,&nsends_mpi)); 5479566063dSJacob Faibussowitsch PetscCall(PetscMPIIntCast(nrecvs,&nrecvs_mpi)); 548b45e3bf4SStefano Zampini if (Bbidx == 0) workB = *outworkB = contents->workB; 549b45e3bf4SStefano Zampini else workB = *outworkB = contents->workB1; 55008401ef6SPierre 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); 5513515ee7fSJunchao Zhang swaits = contents->swaits; 5523515ee7fSJunchao Zhang rwaits = contents->rwaits; 553f32d5d43SBarry Smith 5549566063dSJacob Faibussowitsch PetscCall(MatDenseGetArrayRead(B,&b)); 5559566063dSJacob Faibussowitsch PetscCall(MatDenseGetLDA(B,&blda)); 55608401ef6SPierre 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); 5579566063dSJacob Faibussowitsch PetscCall(MatDenseGetArray(workB,&rvalues)); 558f32d5d43SBarry Smith 55987753ddeSHong Zhang /* Post recv, use MPI derived data type to save memory */ 5609566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 56187753ddeSHong Zhang rtype = contents->rtype; 5623515ee7fSJunchao Zhang for (i=0; i<nrecvs; i++) { 5639566063dSJacob Faibussowitsch PetscCallMPI(MPI_Irecv(rvalues+(rstarts[i]-rstarts[0]),ncols,rtype[i],rprocs[i],tag,comm,rwaits+i)); 564f32d5d43SBarry Smith } 565f32d5d43SBarry Smith 56687753ddeSHong Zhang stype = contents->stype; 5673515ee7fSJunchao Zhang for (i=0; i<nsends; i++) { 5689566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(b,ncols,stype[i],sprocs[i],tag,comm,swaits+i)); 569f32d5d43SBarry Smith } 5702636ff26SBarry Smith 5719566063dSJacob Faibussowitsch if (nrecvs) PetscCallMPI(MPI_Waitall(nrecvs_mpi,rwaits,MPI_STATUSES_IGNORE)); 5729566063dSJacob Faibussowitsch if (nsends) PetscCallMPI(MPI_Waitall(nsends_mpi,swaits,MPI_STATUSES_IGNORE)); 573f32d5d43SBarry Smith 5749566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemote_Private(ctx,PETSC_TRUE/*send*/,&nsends,&sstarts,&sindices,&sprocs,NULL)); 5759566063dSJacob Faibussowitsch PetscCall(VecScatterRestoreRemoteOrdered_Private(ctx,PETSC_FALSE/*recv*/,&nrecvs,&rstarts,NULL,&rprocs,NULL)); 5769566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArrayRead(B,&b)); 5779566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreArray(workB,&rvalues)); 57887753ddeSHong Zhang PetscFunctionReturn(0); 57987753ddeSHong Zhang } 580f32d5d43SBarry Smith 5816718818eSStefano Zampini static PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIDense(Mat A,Mat B,Mat C) 5829123193aSHong Zhang { 583f32d5d43SBarry Smith Mat_MPIAIJ *aij = (Mat_MPIAIJ*)A->data; 584f32d5d43SBarry Smith Mat_MPIDense *bdense = (Mat_MPIDense*)B->data; 585f32d5d43SBarry Smith Mat_MPIDense *cdense = (Mat_MPIDense*)C->data; 586f32d5d43SBarry Smith Mat workB; 58787753ddeSHong Zhang MPIAIJ_MPIDense *contents; 5889123193aSHong Zhang 5899123193aSHong Zhang PetscFunctionBegin; 5906718818eSStefano Zampini MatCheckProduct(C,3); 59128b400f6SJacob Faibussowitsch PetscCheck(C->product->data,PetscObjectComm((PetscObject)C),PETSC_ERR_PLIB,"Product data empty"); 5926718818eSStefano Zampini contents = (MPIAIJ_MPIDense*)C->product->data; 593f32d5d43SBarry Smith /* diagonal block of A times all local rows of B */ 5945ea7661aSPierre Jolivet /* TODO: this calls a symbolic multiplication every time, which could be avoided */ 5959566063dSJacob Faibussowitsch PetscCall(MatMatMult(aij->A,bdense->A,MAT_REUSE_MATRIX,PETSC_DEFAULT,&cdense->A)); 5965ea7661aSPierre Jolivet if (contents->workB->cmap->n == B->cmap->N) { 59787753ddeSHong Zhang /* get off processor parts of B needed to complete C=A*B */ 5989566063dSJacob Faibussowitsch PetscCall(MatMPIDenseScatter(A,B,0,C,&workB)); 599f32d5d43SBarry Smith 600f32d5d43SBarry Smith /* off-diagonal block of A times nonlocal rows of B */ 6019566063dSJacob Faibussowitsch PetscCall(MatMatMultNumericAdd_SeqAIJ_SeqDense(aij->B,workB,cdense->A,PETSC_TRUE)); 60287753ddeSHong Zhang } else { 6035ea7661aSPierre Jolivet Mat Bb,Cb; 6045ea7661aSPierre Jolivet PetscInt BN=B->cmap->N,n=contents->workB->cmap->n,i; 605664c1e85SStefano Zampini PetscBool ccpu; 60687753ddeSHong Zhang 60708401ef6SPierre Jolivet PetscCheck(n > 0,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"Column block size %" PetscInt_FMT " must be positive",n); 608664c1e85SStefano Zampini /* Prevent from unneeded copies back and forth from the GPU 609664c1e85SStefano Zampini when getting and restoring the submatrix 610664c1e85SStefano Zampini We need a proper GPU code for AIJ * dense in parallel */ 6119566063dSJacob Faibussowitsch PetscCall(MatBoundToCPU(C,&ccpu)); 6129566063dSJacob Faibussowitsch PetscCall(MatBindToCPU(C,PETSC_TRUE)); 6135ea7661aSPierre Jolivet for (i=0; i<BN; i+=n) { 6149566063dSJacob Faibussowitsch PetscCall(MatDenseGetSubMatrix(B,i,PetscMin(i+n,BN),&Bb)); 6159566063dSJacob Faibussowitsch PetscCall(MatDenseGetSubMatrix(C,i,PetscMin(i+n,BN),&Cb)); 61687753ddeSHong Zhang 6175ea7661aSPierre Jolivet /* get off processor parts of B needed to complete C=A*B */ 6189566063dSJacob Faibussowitsch PetscCall(MatMPIDenseScatter(A,Bb,(i+n)>BN,C,&workB)); 6195ea7661aSPierre Jolivet 6205ea7661aSPierre Jolivet /* off-diagonal block of A times nonlocal rows of B */ 6215ea7661aSPierre Jolivet cdense = (Mat_MPIDense*)Cb->data; 6229566063dSJacob Faibussowitsch PetscCall(MatMatMultNumericAdd_SeqAIJ_SeqDense(aij->B,workB,cdense->A,PETSC_TRUE)); 6239566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreSubMatrix(B,&Bb)); 6249566063dSJacob Faibussowitsch PetscCall(MatDenseRestoreSubMatrix(C,&Cb)); 62587753ddeSHong Zhang } 6269566063dSJacob Faibussowitsch PetscCall(MatBindToCPU(C,ccpu)); 62787753ddeSHong Zhang } 6289123193aSHong Zhang PetscFunctionReturn(0); 6299123193aSHong Zhang } 630cf1a0672SHong Zhang 631f8efd71cSHong Zhang PetscErrorCode MatMatMultNumeric_MPIAIJ_MPIAIJ(Mat A,Mat P,Mat C) 6321634b032SHong Zhang { 633cf1a0672SHong Zhang Mat_MPIAIJ *a = (Mat_MPIAIJ*)A->data,*c=(Mat_MPIAIJ*)C->data; 634cf1a0672SHong Zhang Mat_SeqAIJ *ad = (Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data; 635cf1a0672SHong Zhang Mat_SeqAIJ *cd = (Mat_SeqAIJ*)(c->A)->data,*co=(Mat_SeqAIJ*)(c->B)->data; 636cf1a0672SHong Zhang PetscInt *adi = ad->i,*adj,*aoi=ao->i,*aoj; 637cf1a0672SHong Zhang PetscScalar *ada,*aoa,*cda=cd->a,*coa=co->a; 638cf1a0672SHong Zhang Mat_SeqAIJ *p_loc,*p_oth; 639cf1a0672SHong Zhang PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*pj; 64029825010SHong Zhang PetscScalar *pa_loc,*pa_oth,*pa,valtmp,*ca; 64129825010SHong Zhang PetscInt cm = C->rmap->n,anz,pnz; 6426718818eSStefano Zampini Mat_APMPI *ptap; 64380da8df7SHong Zhang PetscScalar *apa_sparse; 644ce496241SStefano Zampini const PetscScalar *dummy; 645cf1a0672SHong Zhang PetscInt *api,*apj,*apJ,i,j,k,row; 64629825010SHong Zhang PetscInt cstart = C->cmap->rstart; 64729825010SHong Zhang PetscInt cdnz,conz,k0,k1,nextp; 648a7c7454dSHong Zhang MPI_Comm comm; 649a7c7454dSHong Zhang PetscMPIInt size; 6501634b032SHong Zhang 6511634b032SHong Zhang PetscFunctionBegin; 6526718818eSStefano Zampini MatCheckProduct(C,3); 6536718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 65428b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 6559566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 6569566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 6572c71b3e2SJacob Faibussowitsch PetscCheckFalse(!ptap->P_oth && size>1,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"AP cannot be reused. Do not call MatProductClear()"); 658a7c7454dSHong Zhang 65965e4b4d4SStefano Zampini /* flag CPU mask for C */ 66065e4b4d4SStefano Zampini #if defined(PETSC_HAVE_DEVICE) 66165e4b4d4SStefano Zampini if (C->offloadmask != PETSC_OFFLOAD_UNALLOCATED) C->offloadmask = PETSC_OFFLOAD_CPU; 66265e4b4d4SStefano Zampini if (c->A->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->A->offloadmask = PETSC_OFFLOAD_CPU; 66365e4b4d4SStefano Zampini if (c->B->offloadmask != PETSC_OFFLOAD_UNALLOCATED) c->B->offloadmask = PETSC_OFFLOAD_CPU; 66465e4b4d4SStefano Zampini #endif 66580da8df7SHong Zhang apa_sparse = ptap->apa; 66680da8df7SHong Zhang 667cf1a0672SHong Zhang /* 1) get P_oth = ptap->P_oth and P_loc = ptap->P_loc */ 668cf1a0672SHong Zhang /*-----------------------------------------------------*/ 669cf1a0672SHong Zhang /* update numerical values of P_oth and P_loc */ 6709566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_REUSE_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 6719566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_REUSE_MATRIX,&ptap->P_loc)); 672cf1a0672SHong Zhang 673cf1a0672SHong Zhang /* 2) compute numeric C_loc = A_loc*P = Ad*P_loc + Ao*P_oth */ 674cf1a0672SHong Zhang /*----------------------------------------------------------*/ 675cf1a0672SHong Zhang /* get data from symbolic products */ 676cf1a0672SHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 677cf1a0672SHong Zhang pi_loc = p_loc->i; pj_loc = p_loc->j; pa_loc = p_loc->a; 678a7c7454dSHong Zhang if (size >1) { 679a7c7454dSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 680cf1a0672SHong Zhang pi_oth = p_oth->i; pj_oth = p_oth->j; pa_oth = p_oth->a; 68120e3dc75SHong Zhang } else { 68220e3dc75SHong Zhang p_oth = NULL; pi_oth = NULL; pj_oth = NULL; pa_oth = NULL; 683a7c7454dSHong Zhang } 684cf1a0672SHong Zhang 685ce496241SStefano Zampini /* trigger copy to CPU */ 6869566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->A,&dummy)); 6879566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->A,&dummy)); 6889566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(a->B,&dummy)); 6899566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(a->B,&dummy)); 69029825010SHong Zhang api = ptap->api; 69129825010SHong Zhang apj = ptap->apj; 692cf1a0672SHong Zhang for (i=0; i<cm; i++) { 69329825010SHong Zhang apJ = apj + api[i]; 69429825010SHong Zhang 695cf1a0672SHong Zhang /* diagonal portion of A */ 696cf1a0672SHong Zhang anz = adi[i+1] - adi[i]; 697cf1a0672SHong Zhang adj = ad->j + adi[i]; 698cf1a0672SHong Zhang ada = ad->a + adi[i]; 699cf1a0672SHong Zhang for (j=0; j<anz; j++) { 700cf1a0672SHong Zhang row = adj[j]; 701cf1a0672SHong Zhang pnz = pi_loc[row+1] - pi_loc[row]; 702cf1a0672SHong Zhang pj = pj_loc + pi_loc[row]; 703cf1a0672SHong Zhang pa = pa_loc + pi_loc[row]; 70429825010SHong Zhang /* perform sparse axpy */ 705cf1a0672SHong Zhang valtmp = ada[j]; 70629825010SHong Zhang nextp = 0; 70729825010SHong Zhang for (k=0; nextp<pnz; k++) { 70829825010SHong Zhang if (apJ[k] == pj[nextp]) { /* column of AP == column of P */ 70929825010SHong Zhang apa_sparse[k] += valtmp*pa[nextp++]; 71029825010SHong Zhang } 7111634b032SHong Zhang } 7129566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*pnz)); 713cf1a0672SHong Zhang } 7141634b032SHong Zhang 715cf1a0672SHong Zhang /* off-diagonal portion of A */ 716cf1a0672SHong Zhang anz = aoi[i+1] - aoi[i]; 717cf1a0672SHong Zhang aoj = ao->j + aoi[i]; 718cf1a0672SHong Zhang aoa = ao->a + aoi[i]; 719cf1a0672SHong Zhang for (j=0; j<anz; j++) { 720cf1a0672SHong Zhang row = aoj[j]; 721cf1a0672SHong Zhang pnz = pi_oth[row+1] - pi_oth[row]; 722cf1a0672SHong Zhang pj = pj_oth + pi_oth[row]; 723cf1a0672SHong Zhang pa = pa_oth + pi_oth[row]; 72429825010SHong Zhang /* perform sparse axpy */ 725cf1a0672SHong Zhang valtmp = aoa[j]; 72629825010SHong Zhang nextp = 0; 72729825010SHong Zhang for (k=0; nextp<pnz; k++) { 72829825010SHong Zhang if (apJ[k] == pj[nextp]) { /* column of AP == column of P */ 72929825010SHong Zhang apa_sparse[k] += valtmp*pa[nextp++]; 73029825010SHong Zhang } 731cf1a0672SHong Zhang } 7329566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*pnz)); 733cf1a0672SHong Zhang } 7341634b032SHong Zhang 735cf1a0672SHong Zhang /* set values in C */ 736cf1a0672SHong Zhang cdnz = cd->i[i+1] - cd->i[i]; 737cf1a0672SHong Zhang conz = co->i[i+1] - co->i[i]; 7381634b032SHong Zhang 739910cf402Sprj- /* 1st off-diagonal part of C */ 740cf1a0672SHong Zhang ca = coa + co->i[i]; 741cf1a0672SHong Zhang k = 0; 742cf1a0672SHong Zhang for (k0=0; k0<conz; k0++) { 743cf1a0672SHong Zhang if (apJ[k] >= cstart) break; 74429825010SHong Zhang ca[k0] = apa_sparse[k]; 74529825010SHong Zhang apa_sparse[k] = 0.0; 746cf1a0672SHong Zhang k++; 747cf1a0672SHong Zhang } 7481634b032SHong Zhang 749cf1a0672SHong Zhang /* diagonal part of C */ 750cf1a0672SHong Zhang ca = cda + cd->i[i]; 751cf1a0672SHong Zhang for (k1=0; k1<cdnz; k1++) { 75229825010SHong Zhang ca[k1] = apa_sparse[k]; 75329825010SHong Zhang apa_sparse[k] = 0.0; 754cf1a0672SHong Zhang k++; 755cf1a0672SHong Zhang } 756cf1a0672SHong Zhang 757910cf402Sprj- /* 2nd off-diagonal part of C */ 758cf1a0672SHong Zhang ca = coa + co->i[i]; 759cf1a0672SHong Zhang for (; k0<conz; k0++) { 76029825010SHong Zhang ca[k0] = apa_sparse[k]; 76129825010SHong Zhang apa_sparse[k] = 0.0; 762cf1a0672SHong Zhang k++; 763cf1a0672SHong Zhang } 764cf1a0672SHong Zhang } 7659566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 7669566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 767cf1a0672SHong Zhang PetscFunctionReturn(0); 768cf1a0672SHong Zhang } 769cf1a0672SHong Zhang 7700fc8cf34SHong Zhang /* same as MatMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(), except using LLCondensed to avoid O(BN) memory requirement */ 7714222ddf1SHong Zhang PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ(Mat A,Mat P,PetscReal fill,Mat C) 772cf1a0672SHong Zhang { 773ce94432eSBarry Smith MPI_Comm comm; 774a7c7454dSHong Zhang PetscMPIInt size; 7753cdca5ebSHong Zhang Mat_APMPI *ptap; 7760298fd71SBarry Smith PetscFreeSpaceList free_space = NULL,current_space=NULL; 7776718818eSStefano Zampini Mat_MPIAIJ *a = (Mat_MPIAIJ*)A->data; 778cf1a0672SHong Zhang Mat_SeqAIJ *ad = (Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_loc,*p_oth; 7797dbaa2c6Sandi selinger PetscInt *pi_loc,*pj_loc,*pi_oth,*pj_oth,*dnz,*onz; 780cf1a0672SHong Zhang PetscInt *adi=ad->i,*adj=ad->j,*aoi=ao->i,*aoj=ao->j,rstart=A->rmap->rstart; 7818c78258cSHong Zhang PetscInt i,pnz,row,*api,*apj,*Jptr,apnz,nspacedouble=0,j,nzi,*lnk,apnz_max=1; 782e5ab8c0aSKarl Rupp PetscInt am=A->rmap->n,pn=P->cmap->n,pm=P->rmap->n,lsize=pn+20; 783cf1a0672SHong Zhang PetscReal afill; 784b92f168fSBarry Smith MatType mtype; 785cf1a0672SHong Zhang 786cf1a0672SHong Zhang PetscFunctionBegin; 7876718818eSStefano Zampini MatCheckProduct(C,4); 78828b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Extra product struct not empty"); 7899566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 7909566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 791a7c7454dSHong Zhang 7923cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 7939566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 794cf1a0672SHong Zhang 795cf1a0672SHong Zhang /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 7969566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 79719f0e57aSHong Zhang 798cf1a0672SHong Zhang /* get P_loc by taking all local rows of P */ 7999566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc)); 800cf1a0672SHong Zhang 801cf1a0672SHong Zhang p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 802cf1a0672SHong Zhang pi_loc = p_loc->i; pj_loc = p_loc->j; 803a7c7454dSHong Zhang if (size > 1) { 804a7c7454dSHong Zhang p_oth = (Mat_SeqAIJ*)(ptap->P_oth)->data; 80520e3dc75SHong Zhang pi_oth = p_oth->i; pj_oth = p_oth->j; 806968382fdSHong Zhang } else { 80720e3dc75SHong Zhang p_oth = NULL; 80820e3dc75SHong Zhang pi_oth = NULL; pj_oth = NULL; 809a7c7454dSHong Zhang } 810cf1a0672SHong Zhang 811cf1a0672SHong Zhang /* first, compute symbolic AP = A_loc*P = A_diag*P_loc + A_off*P_oth */ 812cf1a0672SHong Zhang /*-------------------------------------------------------------------*/ 8139566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&api)); 814cf1a0672SHong Zhang ptap->api = api; 815cf1a0672SHong Zhang api[0] = 0; 816cf1a0672SHong Zhang 8179566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(lsize,&lnk)); 818cf1a0672SHong Zhang 819cf1a0672SHong Zhang /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */ 8209566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(adi[am],PetscIntSumTruncate(aoi[am],pi_loc[pm]))),&free_space)); 821cf1a0672SHong Zhang current_space = free_space; 822*d0609cedSBarry Smith MatPreallocateBegin(comm,am,pn,dnz,onz); 823cf1a0672SHong Zhang for (i=0; i<am; i++) { 824cf1a0672SHong Zhang /* diagonal portion of A */ 825cf1a0672SHong Zhang nzi = adi[i+1] - adi[i]; 826cf1a0672SHong Zhang for (j=0; j<nzi; j++) { 827cf1a0672SHong Zhang row = *adj++; 828cf1a0672SHong Zhang pnz = pi_loc[row+1] - pi_loc[row]; 829cf1a0672SHong Zhang Jptr = pj_loc + pi_loc[row]; 8308a210135Sandi selinger /* Expand list if it is not long enough */ 8318328dd82Sandi selinger if (pnz+apnz_max > lsize) { 8328328dd82Sandi selinger lsize = pnz+apnz_max; 8339566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedExpand_Scalable(lsize, &lnk)); 8348a210135Sandi selinger } 835cf1a0672SHong Zhang /* add non-zero cols of P into the sorted linked list lnk */ 8369566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(pnz,Jptr,lnk)); 8378a210135Sandi selinger apnz = *lnk; /* The first element in the list is the number of items in the list */ 8388a210135Sandi selinger api[i+1] = api[i] + apnz; 8398c78258cSHong Zhang if (apnz > apnz_max) apnz_max = apnz + 1; /* '1' for diagonal entry */ 840cf1a0672SHong Zhang } 841cf1a0672SHong Zhang /* off-diagonal portion of A */ 842cf1a0672SHong Zhang nzi = aoi[i+1] - aoi[i]; 843cf1a0672SHong Zhang for (j=0; j<nzi; j++) { 844cf1a0672SHong Zhang row = *aoj++; 845cf1a0672SHong Zhang pnz = pi_oth[row+1] - pi_oth[row]; 846cf1a0672SHong Zhang Jptr = pj_oth + pi_oth[row]; 8478a210135Sandi selinger /* Expand list if it is not long enough */ 8488328dd82Sandi selinger if (pnz+apnz_max > lsize) { 8498328dd82Sandi selinger lsize = pnz + apnz_max; 8509566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedExpand_Scalable(lsize, &lnk)); 8518a210135Sandi selinger } 8528a210135Sandi selinger /* add non-zero cols of P into the sorted linked list lnk */ 8539566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(pnz,Jptr,lnk)); 8548a210135Sandi selinger apnz = *lnk; /* The first element in the list is the number of items in the list */ 8558a210135Sandi selinger api[i+1] = api[i] + apnz; 8568c78258cSHong Zhang if (apnz > apnz_max) apnz_max = apnz + 1; /* '1' for diagonal entry */ 857cf1a0672SHong Zhang } 8588c78258cSHong Zhang 8598c78258cSHong Zhang /* add missing diagonal entry */ 8608c78258cSHong Zhang if (C->force_diagonals) { 8618c78258cSHong Zhang j = i + rstart; /* column index */ 8629566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(1,&j,lnk)); 8638c78258cSHong Zhang } 8648c78258cSHong Zhang 865f84c536bSHong Zhang apnz = *lnk; 866cf1a0672SHong Zhang api[i+1] = api[i] + apnz; 8678a210135Sandi selinger if (apnz > apnz_max) apnz_max = apnz; 868cf1a0672SHong Zhang 869cf1a0672SHong Zhang /* if free space is not available, double the total space in the list */ 870cf1a0672SHong Zhang if (current_space->local_remaining<apnz) { 8719566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(apnz,current_space->total_array_size),¤t_space)); 872cf1a0672SHong Zhang nspacedouble++; 873cf1a0672SHong Zhang } 874cf1a0672SHong Zhang 875cf1a0672SHong Zhang /* Copy data into free space, then initialize lnk */ 8769566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(apnz,current_space->array,lnk)); 8779566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+rstart,apnz,current_space->array,dnz,onz)); 8782205254eSKarl Rupp 879cf1a0672SHong Zhang current_space->array += apnz; 880cf1a0672SHong Zhang current_space->local_used += apnz; 881cf1a0672SHong Zhang current_space->local_remaining -= apnz; 882cf1a0672SHong Zhang } 883cf1a0672SHong Zhang 884cf1a0672SHong Zhang /* Allocate space for apj, initialize apj, and */ 885cf1a0672SHong Zhang /* destroy list of free space and other temporary array(s) */ 8869566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(api[am]+1,&ptap->apj)); 887cf1a0672SHong Zhang apj = ptap->apj; 8889566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,ptap->apj)); 8899566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); 890cf1a0672SHong Zhang 8914222ddf1SHong Zhang /* create and assemble symbolic parallel matrix C */ 892cf1a0672SHong Zhang /*----------------------------------------------------*/ 8939566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,am,pn,PETSC_DETERMINE,PETSC_DETERMINE)); 8949566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,P)); 8959566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 8969566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 8979566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 898*d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 899cf1a0672SHong Zhang 9004222ddf1SHong Zhang /* malloc apa for assembly C */ 9019566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(apnz_max,&ptap->apa)); 90217f75172Sandi selinger 9039566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api)); 9049566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 9059566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 9069566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 907cf1a0672SHong Zhang 9084222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ; 9094222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AB; 910cf1a0672SHong Zhang 9114222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 9126718818eSStefano Zampini C->product->data = ptap; 9136718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_MatMatMult; 914cf1a0672SHong Zhang 915cf1a0672SHong Zhang /* set MatInfo */ 916118727c9SMark F. Adams afill = (PetscReal)api[am]/(adi[am]+aoi[am]+pi_loc[pm]+1) + 1.e-5; 917cf1a0672SHong Zhang if (afill < 1.0) afill = 1.0; 9184222ddf1SHong Zhang C->info.mallocs = nspacedouble; 9194222ddf1SHong Zhang C->info.fill_ratio_given = fill; 9204222ddf1SHong Zhang C->info.fill_ratio_needed = afill; 921cf1a0672SHong Zhang 922cf1a0672SHong Zhang #if defined(PETSC_USE_INFO) 923cf1a0672SHong Zhang if (api[am]) { 9249566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 9259566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n",(double)afill)); 926cf1a0672SHong Zhang } else { 9279566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 928cf1a0672SHong Zhang } 929cf1a0672SHong Zhang #endif 9301634b032SHong Zhang PetscFunctionReturn(0); 9311634b032SHong Zhang } 932f96d37fcSHong Zhang 933a9d7e0ccSandi selinger /* This function is needed for the seqMPI matrix-matrix multiplication. */ 934a9d7e0ccSandi selinger /* Three input arrays are merged to one output array. The size of the */ 935a9d7e0ccSandi selinger /* output array is also output. Duplicate entries only show up once. */ 936a9d7e0ccSandi selinger static void Merge3SortedArrays(PetscInt size1, PetscInt *in1, 937a9d7e0ccSandi selinger PetscInt size2, PetscInt *in2, 938a9d7e0ccSandi selinger PetscInt size3, PetscInt *in3, 939a9d7e0ccSandi selinger PetscInt *size4, PetscInt *out) 940a9d7e0ccSandi selinger { 941a9d7e0ccSandi selinger int i = 0, j = 0, k = 0, l = 0; 942a9d7e0ccSandi selinger 943a9d7e0ccSandi selinger /* Traverse all three arrays */ 944a9d7e0ccSandi selinger while (i<size1 && j<size2 && k<size3) { 945a9d7e0ccSandi selinger if (in1[i] < in2[j] && in1[i] < in3[k]) { 946a9d7e0ccSandi selinger out[l++] = in1[i++]; 947a9d7e0ccSandi selinger } 948a9d7e0ccSandi selinger else if (in2[j] < in1[i] && in2[j] < in3[k]) { 949a9d7e0ccSandi selinger out[l++] = in2[j++]; 950a9d7e0ccSandi selinger } 951a9d7e0ccSandi selinger else if (in3[k] < in1[i] && in3[k] < in2[j]) { 952a9d7e0ccSandi selinger out[l++] = in3[k++]; 953a9d7e0ccSandi selinger } 954a9d7e0ccSandi selinger else if (in1[i] == in2[j] && in1[i] < in3[k]) { 955a9d7e0ccSandi selinger out[l++] = in1[i]; 956a9d7e0ccSandi selinger i++, j++; 957a9d7e0ccSandi selinger } 958a9d7e0ccSandi selinger else if (in1[i] == in3[k] && in1[i] < in2[j]) { 959a9d7e0ccSandi selinger out[l++] = in1[i]; 960a9d7e0ccSandi selinger i++, k++; 961a9d7e0ccSandi selinger } 962a9d7e0ccSandi selinger else if (in3[k] == in2[j] && in2[j] < in1[i]) { 963a9d7e0ccSandi selinger out[l++] = in2[j]; 964a9d7e0ccSandi selinger k++, j++; 965a9d7e0ccSandi selinger } 966a9d7e0ccSandi selinger else if (in1[i] == in2[j] && in1[i] == in3[k]) { 967a9d7e0ccSandi selinger out[l++] = in1[i]; 968a9d7e0ccSandi selinger i++, j++, k++; 969a9d7e0ccSandi selinger } 970a9d7e0ccSandi selinger } 971a9d7e0ccSandi selinger 972a9d7e0ccSandi selinger /* Traverse two remaining arrays */ 973a9d7e0ccSandi selinger while (i<size1 && j<size2) { 974a9d7e0ccSandi selinger if (in1[i] < in2[j]) { 975a9d7e0ccSandi selinger out[l++] = in1[i++]; 976a9d7e0ccSandi selinger } 977a9d7e0ccSandi selinger else if (in1[i] > in2[j]) { 978a9d7e0ccSandi selinger out[l++] = in2[j++]; 979a9d7e0ccSandi selinger } 980a9d7e0ccSandi selinger else { 981a9d7e0ccSandi selinger out[l++] = in1[i]; 982a9d7e0ccSandi selinger i++, j++; 983a9d7e0ccSandi selinger } 984a9d7e0ccSandi selinger } 985a9d7e0ccSandi selinger 986a9d7e0ccSandi selinger while (i<size1 && k<size3) { 987a9d7e0ccSandi selinger if (in1[i] < in3[k]) { 988a9d7e0ccSandi selinger out[l++] = in1[i++]; 989a9d7e0ccSandi selinger } 990a9d7e0ccSandi selinger else if (in1[i] > in3[k]) { 991a9d7e0ccSandi selinger out[l++] = in3[k++]; 992a9d7e0ccSandi selinger } 993a9d7e0ccSandi selinger else { 994a9d7e0ccSandi selinger out[l++] = in1[i]; 995a9d7e0ccSandi selinger i++, k++; 996a9d7e0ccSandi selinger } 997a9d7e0ccSandi selinger } 998a9d7e0ccSandi selinger 999a9d7e0ccSandi selinger while (k<size3 && j<size2) { 1000a9d7e0ccSandi selinger if (in3[k] < in2[j]) { 1001a9d7e0ccSandi selinger out[l++] = in3[k++]; 1002a9d7e0ccSandi selinger } 1003a9d7e0ccSandi selinger else if (in3[k] > in2[j]) { 1004a9d7e0ccSandi selinger out[l++] = in2[j++]; 1005a9d7e0ccSandi selinger } 1006a9d7e0ccSandi selinger else { 1007a9d7e0ccSandi selinger out[l++] = in3[k]; 1008a9d7e0ccSandi selinger k++, j++; 1009a9d7e0ccSandi selinger } 1010a9d7e0ccSandi selinger } 1011a9d7e0ccSandi selinger 1012a9d7e0ccSandi selinger /* Traverse one remaining array */ 1013a9d7e0ccSandi selinger while (i<size1) out[l++] = in1[i++]; 1014a9d7e0ccSandi selinger while (j<size2) out[l++] = in2[j++]; 1015a9d7e0ccSandi selinger while (k<size3) out[l++] = in3[k++]; 1016a9d7e0ccSandi selinger 1017a9d7e0ccSandi selinger *size4 = l; 1018a9d7e0ccSandi selinger } 1019a9d7e0ccSandi selinger 102073b67375Sandi selinger /* This matrix-matrix multiplication algorithm divides the multiplication into three multiplications and */ 102173b67375Sandi selinger /* adds up the products. Two of these three multiplications are performed with existing (sequential) */ 102273b67375Sandi selinger /* matrix-matrix multiplications. */ 10234222ddf1SHong Zhang PetscErrorCode MatMatMultSymbolic_MPIAIJ_MPIAIJ_seqMPI(Mat A, Mat P, PetscReal fill, Mat C) 1024a9d7e0ccSandi selinger { 1025a9d7e0ccSandi selinger MPI_Comm comm; 1026a9d7e0ccSandi selinger PetscMPIInt size; 10273cdca5ebSHong Zhang Mat_APMPI *ptap; 1028a9d7e0ccSandi selinger PetscFreeSpaceList free_space_diag=NULL, current_space=NULL; 1029a9d7e0ccSandi selinger Mat_MPIAIJ *a =(Mat_MPIAIJ*)A->data; 1030a9d7e0ccSandi selinger Mat_SeqAIJ *ad =(Mat_SeqAIJ*)(a->A)->data,*ao=(Mat_SeqAIJ*)(a->B)->data,*p_loc; 1031a9d7e0ccSandi selinger Mat_MPIAIJ *p =(Mat_MPIAIJ*)P->data; 1032a9d7e0ccSandi selinger Mat_SeqAIJ *adpd_seq, *p_off, *aopoth_seq; 1033a9d7e0ccSandi selinger PetscInt adponz, adpdnz; 10347dbaa2c6Sandi selinger PetscInt *pi_loc,*dnz,*onz; 1035a9d7e0ccSandi selinger PetscInt *adi=ad->i,*adj=ad->j,*aoi=ao->i,rstart=A->rmap->rstart; 1036a9d7e0ccSandi selinger PetscInt *lnk,i, i1=0,pnz,row,*adpoi,*adpoj, *api, *adpoJ, *aopJ, *apJ,*Jptr, aopnz, nspacedouble=0,j,nzi, 1037a9d7e0ccSandi selinger *apj,apnz, *adpdi, *adpdj, *adpdJ, *poff_i, *poff_j, *j_temp, *aopothi, *aopothj; 1038a016958eSandi selinger PetscInt am=A->rmap->n,pN=P->cmap->N,pn=P->cmap->n,pm=P->rmap->n, p_colstart, p_colend; 1039a9d7e0ccSandi selinger PetscBT lnkbt; 1040a9d7e0ccSandi selinger PetscReal afill; 1041a9d7e0ccSandi selinger PetscMPIInt rank; 1042a9d7e0ccSandi selinger Mat adpd, aopoth; 1043b92f168fSBarry Smith MatType mtype; 1044e83e5f86SFande Kong const char *prefix; 1045a9d7e0ccSandi selinger 1046a9d7e0ccSandi selinger PetscFunctionBegin; 10476718818eSStefano Zampini MatCheckProduct(C,4); 104828b400f6SJacob Faibussowitsch PetscCheck(!C->product->data,PETSC_COMM_SELF,PETSC_ERR_PLIB,"Extra product struct not empty"); 10499566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 10509566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 10519566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm, &rank)); 10529566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRangeColumn(P, &p_colstart, &p_colend)); 1053a9d7e0ccSandi selinger 10543cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 10559566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 1056a9d7e0ccSandi selinger 1057a9d7e0ccSandi selinger /* get P_oth by taking rows of P (= non-zero cols of local A) from other processors */ 10589566063dSJacob Faibussowitsch PetscCall(MatGetBrowsOfAoCols_MPIAIJ(A,P,MAT_INITIAL_MATRIX,&ptap->startsj_s,&ptap->startsj_r,&ptap->bufa,&ptap->P_oth)); 1059a9d7e0ccSandi selinger 1060a9d7e0ccSandi selinger /* get P_loc by taking all local rows of P */ 10619566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(P,MAT_INITIAL_MATRIX,&ptap->P_loc)); 1062a9d7e0ccSandi selinger 1063a9d7e0ccSandi selinger p_loc = (Mat_SeqAIJ*)(ptap->P_loc)->data; 1064a9d7e0ccSandi selinger pi_loc = p_loc->i; 1065a9d7e0ccSandi selinger 1066a9d7e0ccSandi selinger /* Allocate memory for the i arrays of the matrices A*P, A_diag*P_off and A_offd * P */ 10679566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&api)); 10689566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(am+2,&adpoi)); 1069a9d7e0ccSandi selinger 1070a9d7e0ccSandi selinger adpoi[0] = 0; 1071a9d7e0ccSandi selinger ptap->api = api; 1072a9d7e0ccSandi selinger api[0] = 0; 1073a9d7e0ccSandi selinger 1074a9d7e0ccSandi selinger /* create and initialize a linked list, will be used for both A_diag * P_loc_off and A_offd * P_oth */ 10759566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate(pN,pN,&lnk,&lnkbt)); 1076*d0609cedSBarry Smith MatPreallocateBegin(comm,am,pn,dnz,onz); 1077a9d7e0ccSandi selinger 1078a9d7e0ccSandi selinger /* Symbolic calc of A_loc_diag * P_loc_diag */ 10799566063dSJacob Faibussowitsch PetscCall(MatGetOptionsPrefix(A,&prefix)); 10809566063dSJacob Faibussowitsch PetscCall(MatProductCreate(a->A,p->A,NULL,&adpd)); 10819566063dSJacob Faibussowitsch PetscCall(MatGetOptionsPrefix(A,&prefix)); 10829566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(adpd,prefix)); 10839566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(adpd,"inner_diag_")); 10844222ddf1SHong Zhang 10859566063dSJacob Faibussowitsch PetscCall(MatProductSetType(adpd,MATPRODUCT_AB)); 10869566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(adpd,"sorted")); 10879566063dSJacob Faibussowitsch PetscCall(MatProductSetFill(adpd,fill)); 10889566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions(adpd)); 10898c78258cSHong Zhang 10908c78258cSHong Zhang adpd->force_diagonals = C->force_diagonals; 10919566063dSJacob Faibussowitsch PetscCall(MatProductSymbolic(adpd)); 10924222ddf1SHong Zhang 1093a9d7e0ccSandi selinger adpd_seq = (Mat_SeqAIJ*)((adpd)->data); 1094a9d7e0ccSandi selinger adpdi = adpd_seq->i; adpdj = adpd_seq->j; 1095a9d7e0ccSandi selinger p_off = (Mat_SeqAIJ*)((p->B)->data); 1096a9d7e0ccSandi selinger poff_i = p_off->i; poff_j = p_off->j; 1097a9d7e0ccSandi selinger 109873b67375Sandi selinger /* j_temp stores indices of a result row before they are added to the linked list */ 10999566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pN+2,&j_temp)); 1100a9d7e0ccSandi selinger 1101a9d7e0ccSandi selinger /* Symbolic calc of the A_diag * p_loc_off */ 1102a9d7e0ccSandi selinger /* Initial FreeSpace size is fill*(nnz(A)+nnz(P)) */ 11039566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscRealIntMultTruncate(fill,PetscIntSumTruncate(adi[am],PetscIntSumTruncate(aoi[am],pi_loc[pm]))),&free_space_diag)); 1104a9d7e0ccSandi selinger current_space = free_space_diag; 1105a9d7e0ccSandi selinger 1106a9d7e0ccSandi selinger for (i=0; i<am; i++) { 1107a9d7e0ccSandi selinger /* A_diag * P_loc_off */ 1108a9d7e0ccSandi selinger nzi = adi[i+1] - adi[i]; 1109a9d7e0ccSandi selinger for (j=0; j<nzi; j++) { 1110a9d7e0ccSandi selinger row = *adj++; 1111a9d7e0ccSandi selinger pnz = poff_i[row+1] - poff_i[row]; 1112a9d7e0ccSandi selinger Jptr = poff_j + poff_i[row]; 1113a9d7e0ccSandi selinger for (i1 = 0; i1 < pnz; i1++) { 1114a9d7e0ccSandi selinger j_temp[i1] = p->garray[Jptr[i1]]; 1115a9d7e0ccSandi selinger } 1116a9d7e0ccSandi selinger /* add non-zero cols of P into the sorted linked list lnk */ 11179566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(pnz,j_temp,lnk,lnkbt)); 1118a9d7e0ccSandi selinger } 1119a9d7e0ccSandi selinger 1120a9d7e0ccSandi selinger adponz = lnk[0]; 1121a9d7e0ccSandi selinger adpoi[i+1] = adpoi[i] + adponz; 1122a9d7e0ccSandi selinger 1123a9d7e0ccSandi selinger /* if free space is not available, double the total space in the list */ 1124a9d7e0ccSandi selinger if (current_space->local_remaining<adponz) { 11259566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(adponz,current_space->total_array_size),¤t_space)); 1126a9d7e0ccSandi selinger nspacedouble++; 1127a9d7e0ccSandi selinger } 1128a9d7e0ccSandi selinger 1129a9d7e0ccSandi selinger /* Copy data into free space, then initialize lnk */ 11309566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean(pN,adponz,current_space->array,lnk,lnkbt)); 1131a9d7e0ccSandi selinger 1132a9d7e0ccSandi selinger current_space->array += adponz; 1133a9d7e0ccSandi selinger current_space->local_used += adponz; 1134a9d7e0ccSandi selinger current_space->local_remaining -= adponz; 1135a9d7e0ccSandi selinger } 1136a9d7e0ccSandi selinger 1137a9d7e0ccSandi selinger /* Symbolic calc of A_off * P_oth */ 11389566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(a->B,prefix)); 11399566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(a->B,"inner_offdiag_")); 11409566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF,&aopoth)); 11419566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(a->B, ptap->P_oth, fill, aopoth)); 1142a9d7e0ccSandi selinger aopoth_seq = (Mat_SeqAIJ*)((aopoth)->data); 1143a9d7e0ccSandi selinger aopothi = aopoth_seq->i; aopothj = aopoth_seq->j; 1144a9d7e0ccSandi selinger 1145a9d7e0ccSandi selinger /* Allocate space for apj, adpj, aopj, ... */ 1146a9d7e0ccSandi selinger /* destroy lists of free space and other temporary array(s) */ 1147a9d7e0ccSandi selinger 11489566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(aopothi[am] + adpoi[am] + adpdi[am]+2, &ptap->apj)); 11499566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(adpoi[am]+2, &adpoj)); 1150a9d7e0ccSandi selinger 1151a9d7e0ccSandi selinger /* Copy from linked list to j-array */ 11529566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space_diag,adpoj)); 11539566063dSJacob Faibussowitsch PetscCall(PetscLLDestroy(lnk,lnkbt)); 1154a9d7e0ccSandi selinger 1155a9d7e0ccSandi selinger adpoJ = adpoj; 1156a9d7e0ccSandi selinger adpdJ = adpdj; 1157a9d7e0ccSandi selinger aopJ = aopothj; 1158a9d7e0ccSandi selinger apj = ptap->apj; 1159a9d7e0ccSandi selinger apJ = apj; /* still empty */ 1160a9d7e0ccSandi selinger 1161a9d7e0ccSandi selinger /* Merge j-arrays of A_off * P, A_diag * P_loc_off, and */ 1162a9d7e0ccSandi selinger /* A_diag * P_loc_diag to get A*P */ 1163a9d7e0ccSandi selinger for (i = 0; i < am; i++) { 1164a9d7e0ccSandi selinger aopnz = aopothi[i+1] - aopothi[i]; 1165a9d7e0ccSandi selinger adponz = adpoi[i+1] - adpoi[i]; 1166a9d7e0ccSandi selinger adpdnz = adpdi[i+1] - adpdi[i]; 1167a9d7e0ccSandi selinger 1168a9d7e0ccSandi selinger /* Correct indices from A_diag*P_diag */ 1169a9d7e0ccSandi selinger for (i1 = 0; i1 < adpdnz; i1++) { 1170a016958eSandi selinger adpdJ[i1] += p_colstart; 1171a9d7e0ccSandi selinger } 1172a9d7e0ccSandi selinger /* Merge j-arrays of A_diag * P_loc_off and A_diag * P_loc_diag and A_off * P_oth */ 1173a9d7e0ccSandi selinger Merge3SortedArrays(adponz, adpoJ, adpdnz, adpdJ, aopnz, aopJ, &apnz, apJ); 11749566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+rstart, apnz, apJ, dnz, onz)); 1175a9d7e0ccSandi selinger 1176a9d7e0ccSandi selinger aopJ += aopnz; 1177a9d7e0ccSandi selinger adpoJ += adponz; 1178a9d7e0ccSandi selinger adpdJ += adpdnz; 1179a9d7e0ccSandi selinger apJ += apnz; 1180a9d7e0ccSandi selinger api[i+1] = api[i] + apnz; 1181a9d7e0ccSandi selinger } 1182a9d7e0ccSandi selinger 1183a9d7e0ccSandi selinger /* malloc apa to store dense row A[i,:]*P */ 11849566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(pN+2,&ptap->apa)); 1185a9d7e0ccSandi selinger 11864222ddf1SHong Zhang /* create and assemble symbolic parallel matrix C */ 11879566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,am,pn,PETSC_DETERMINE,PETSC_DETERMINE)); 11889566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizesFromMats(C,A,P)); 11899566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 11909566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 11919566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 1192*d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 119317f75172Sandi selinger 11949566063dSJacob Faibussowitsch PetscCall(MatSetValues_MPIAIJ_CopyFromCSRFormat_Symbolic(C, apj, api)); 11959566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 11969566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 11979566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 119817f75172Sandi selinger 11994222ddf1SHong Zhang C->ops->matmultnumeric = MatMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable; 12004222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AB; 1201a9d7e0ccSandi selinger 12024222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 12036718818eSStefano Zampini C->product->data = ptap; 12046718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_MatMatMult; 1205a9d7e0ccSandi selinger 1206a9d7e0ccSandi selinger /* set MatInfo */ 120773b67375Sandi selinger afill = (PetscReal)api[am]/(adi[am]+aoi[am]+pi_loc[pm]+1) + 1.e-5; 120873b67375Sandi selinger if (afill < 1.0) afill = 1.0; 12094222ddf1SHong Zhang C->info.mallocs = nspacedouble; 12104222ddf1SHong Zhang C->info.fill_ratio_given = fill; 12114222ddf1SHong Zhang C->info.fill_ratio_needed = afill; 1212a9d7e0ccSandi selinger 1213a9d7e0ccSandi selinger #if defined(PETSC_USE_INFO) 1214a9d7e0ccSandi selinger if (api[am]) { 12159566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 12169566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatMatMult(A,B,MatReuse,%g,&C) for best performance.;\n",(double)afill)); 1217a9d7e0ccSandi selinger } else { 12189566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 1219a9d7e0ccSandi selinger } 1220a9d7e0ccSandi selinger #endif 1221a9d7e0ccSandi selinger 12229566063dSJacob Faibussowitsch PetscCall(MatDestroy(&aopoth)); 12239566063dSJacob Faibussowitsch PetscCall(MatDestroy(&adpd)); 12249566063dSJacob Faibussowitsch PetscCall(PetscFree(j_temp)); 12259566063dSJacob Faibussowitsch PetscCall(PetscFree(adpoj)); 12269566063dSJacob Faibussowitsch PetscCall(PetscFree(adpoi)); 1227a9d7e0ccSandi selinger PetscFunctionReturn(0); 1228a9d7e0ccSandi selinger } 1229a9d7e0ccSandi selinger 1230f96d37fcSHong Zhang /*-------------------------------------------------------------------------*/ 1231c216dbf3SHong Zhang /* This routine only works when scall=MAT_REUSE_MATRIX! */ 1232c216dbf3SHong Zhang PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_matmatmult(Mat P,Mat A,Mat C) 1233c216dbf3SHong Zhang { 12346718818eSStefano Zampini Mat_APMPI *ptap; 123560552ceaSHong Zhang Mat Pt; 1236c216dbf3SHong Zhang 1237c216dbf3SHong Zhang PetscFunctionBegin; 12386718818eSStefano Zampini MatCheckProduct(C,3); 12396718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 124028b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 124128b400f6SJacob Faibussowitsch PetscCheck(ptap->Pt,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be reused. Do not call MatProductClear()"); 124260552ceaSHong Zhang 124360552ceaSHong Zhang Pt = ptap->Pt; 12449566063dSJacob Faibussowitsch PetscCall(MatTranspose(P,MAT_REUSE_MATRIX,&Pt)); 12459566063dSJacob Faibussowitsch PetscCall(MatMatMultNumeric_MPIAIJ_MPIAIJ(Pt,A,C)); 1246f96d37fcSHong Zhang PetscFunctionReturn(0); 1247f96d37fcSHong Zhang } 1248f96d37fcSHong Zhang 1249c5af039cSHong Zhang /* This routine is modified from MatPtAPSymbolic_MPIAIJ_MPIAIJ() */ 12504222ddf1SHong Zhang PetscErrorCode MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(Mat P,Mat A,PetscReal fill,Mat C) 1251912837bbSHong Zhang { 12523cdca5ebSHong Zhang Mat_APMPI *ptap; 12536718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data; 1254912837bbSHong Zhang MPI_Comm comm; 1255912837bbSHong Zhang PetscMPIInt size,rank; 1256912837bbSHong Zhang PetscFreeSpaceList free_space=NULL,current_space=NULL; 12574ede91e3SHong Zhang PetscInt pn=P->cmap->n,aN=A->cmap->N,an=A->cmap->n; 125869be949aSStefano Zampini PetscInt *lnk,i,k,nsend,rstart; 1259912837bbSHong Zhang PetscBT lnkbt; 1260ec4bef21SJose E. Roman PetscMPIInt tagi,tagj,*len_si,*len_s,*len_ri,nrecv; 1261ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted=0; 12628c78258cSHong Zhang PetscInt **buf_rj,**buf_ri,**buf_ri_k,row,ncols,*cols; 12637f3e1f30SHong Zhang PetscInt len,proc,*dnz,*onz,*owners,nzi; 1264912837bbSHong Zhang PetscInt nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci; 1265912837bbSHong Zhang MPI_Request *swaits,*rwaits; 1266912837bbSHong Zhang MPI_Status *sstatus,rstatus; 1267912837bbSHong Zhang PetscLayout rowmap; 1268912837bbSHong Zhang PetscInt *owners_co,*coi,*coj; /* i and j array of (p->B)^T*A*P - used in the communication */ 1269912837bbSHong Zhang PetscMPIInt *len_r,*id_r; /* array of length of comm->size, store send/recv matrix values */ 12707f3e1f30SHong Zhang PetscInt *Jptr,*prmap=p->garray,con,j,Crmax; 12714ede91e3SHong Zhang Mat_SeqAIJ *a_loc,*c_loc,*c_oth; 1272912837bbSHong Zhang PetscTable ta; 1273b92f168fSBarry Smith MatType mtype; 1274e83e5f86SFande Kong const char *prefix; 1275912837bbSHong Zhang 1276912837bbSHong Zhang PetscFunctionBegin; 12779566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 12789566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 12799566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1280912837bbSHong Zhang 12814222ddf1SHong Zhang /* create symbolic parallel matrix C */ 12829566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 12839566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 1284912837bbSHong Zhang 12854222ddf1SHong Zhang C->ops->transposematmultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable; 1286912837bbSHong Zhang 12873cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 12889566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 1289912837bbSHong Zhang ptap->reuse = MAT_INITIAL_MATRIX; 1290912837bbSHong Zhang 1291912837bbSHong Zhang /* (0) compute Rd = Pd^T, Ro = Po^T */ 1292912837bbSHong Zhang /* --------------------------------- */ 12939566063dSJacob Faibussowitsch PetscCall(MatTranspose_SeqAIJ(p->A,MAT_INITIAL_MATRIX,&ptap->Rd)); 12949566063dSJacob Faibussowitsch PetscCall(MatTranspose_SeqAIJ(p->B,MAT_INITIAL_MATRIX,&ptap->Ro)); 1295912837bbSHong Zhang 1296669e9c86SHong Zhang /* (1) compute symbolic A_loc */ 1297669e9c86SHong Zhang /* ---------------------------*/ 12989566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_INITIAL_MATRIX,&ptap->A_loc)); 1299912837bbSHong Zhang 13004ede91e3SHong Zhang /* (2-1) compute symbolic C_oth = Ro*A_loc */ 1301912837bbSHong Zhang /* ------------------------------------ */ 13029566063dSJacob Faibussowitsch PetscCall(MatGetOptionsPrefix(A,&prefix)); 13039566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(ptap->Ro,prefix)); 13049566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(ptap->Ro,"inner_offdiag_")); 13059566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF,&ptap->C_oth)); 13069566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Ro,ptap->A_loc,fill,ptap->C_oth)); 1307912837bbSHong Zhang 1308912837bbSHong Zhang /* (3) send coj of C_oth to other processors */ 1309912837bbSHong Zhang /* ------------------------------------------ */ 1310912837bbSHong Zhang /* determine row ownership */ 13119566063dSJacob Faibussowitsch PetscCall(PetscLayoutCreate(comm,&rowmap)); 1312912837bbSHong Zhang rowmap->n = pn; 1313912837bbSHong Zhang rowmap->bs = 1; 13149566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(rowmap)); 1315912837bbSHong Zhang owners = rowmap->range; 1316912837bbSHong Zhang 1317912837bbSHong Zhang /* determine the number of messages to send, their lengths */ 13189566063dSJacob Faibussowitsch PetscCall(PetscMalloc4(size,&len_s,size,&len_si,size,&sstatus,size+2,&owners_co)); 13199566063dSJacob Faibussowitsch PetscCall(PetscArrayzero(len_s,size)); 13209566063dSJacob Faibussowitsch PetscCall(PetscArrayzero(len_si,size)); 1321912837bbSHong Zhang 1322912837bbSHong Zhang c_oth = (Mat_SeqAIJ*)ptap->C_oth->data; 1323912837bbSHong Zhang coi = c_oth->i; coj = c_oth->j; 1324912837bbSHong Zhang con = ptap->C_oth->rmap->n; 1325912837bbSHong Zhang proc = 0; 1326912837bbSHong Zhang for (i=0; i<con; i++) { 1327912837bbSHong Zhang while (prmap[i] >= owners[proc+1]) proc++; 13287f3e1f30SHong Zhang len_si[proc]++; /* num of rows in Co(=Pt*A) to be sent to [proc] */ 1329912837bbSHong Zhang len_s[proc] += coi[i+1] - coi[i]; /* num of nonzeros in Co to be sent to [proc] */ 1330912837bbSHong Zhang } 1331912837bbSHong Zhang 1332912837bbSHong Zhang len = 0; /* max length of buf_si[], see (4) */ 1333912837bbSHong Zhang owners_co[0] = 0; 1334912837bbSHong Zhang nsend = 0; 1335912837bbSHong Zhang for (proc=0; proc<size; proc++) { 1336912837bbSHong Zhang owners_co[proc+1] = owners_co[proc] + len_si[proc]; 1337912837bbSHong Zhang if (len_s[proc]) { 1338912837bbSHong Zhang nsend++; 1339912837bbSHong Zhang len_si[proc] = 2*(len_si[proc] + 1); /* length of buf_si to be sent to [proc] */ 1340912837bbSHong Zhang len += len_si[proc]; 1341912837bbSHong Zhang } 1342912837bbSHong Zhang } 1343912837bbSHong Zhang 1344912837bbSHong Zhang /* determine the number and length of messages to receive for coi and coj */ 13459566063dSJacob Faibussowitsch PetscCall(PetscGatherNumberOfMessages(comm,NULL,len_s,&nrecv)); 13469566063dSJacob Faibussowitsch PetscCall(PetscGatherMessageLengths2(comm,nsend,nrecv,len_s,len_si,&id_r,&len_r,&len_ri)); 1347912837bbSHong Zhang 1348912837bbSHong Zhang /* post the Irecv and Isend of coj */ 13499566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagj)); 13509566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagj,nrecv,id_r,len_r,&buf_rj,&rwaits)); 13519566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(nsend+1,&swaits)); 1352912837bbSHong Zhang for (proc=0, k=0; proc<size; proc++) { 1353912837bbSHong Zhang if (!len_s[proc]) continue; 1354912837bbSHong Zhang i = owners_co[proc]; 13559566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tagj,comm,swaits+k)); 1356912837bbSHong Zhang k++; 1357912837bbSHong Zhang } 1358912837bbSHong Zhang 13597f3e1f30SHong Zhang /* (2-2) compute symbolic C_loc = Rd*A_loc */ 1360912837bbSHong Zhang /* ---------------------------------------- */ 13619566063dSJacob Faibussowitsch PetscCall(MatSetOptionsPrefix(ptap->Rd,prefix)); 13629566063dSJacob Faibussowitsch PetscCall(MatAppendOptionsPrefix(ptap->Rd,"inner_diag_")); 13639566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_SELF,&ptap->C_loc)); 13649566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_SeqAIJ_SeqAIJ(ptap->Rd,ptap->A_loc,fill,ptap->C_loc)); 1365912837bbSHong Zhang c_loc = (Mat_SeqAIJ*)ptap->C_loc->data; 1366912837bbSHong Zhang 1367912837bbSHong Zhang /* receives coj are complete */ 1368912837bbSHong Zhang for (i=0; i<nrecv; i++) { 13699566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(nrecv,rwaits,&icompleted,&rstatus)); 1370912837bbSHong Zhang } 13719566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 13729566063dSJacob Faibussowitsch if (nsend) PetscCallMPI(MPI_Waitall(nsend,swaits,sstatus)); 1373912837bbSHong Zhang 1374912837bbSHong Zhang /* add received column indices into ta to update Crmax */ 13754ede91e3SHong Zhang a_loc = (Mat_SeqAIJ*)(ptap->A_loc)->data; 1376669e9c86SHong Zhang 1377669e9c86SHong Zhang /* create and initialize a linked list */ 13789566063dSJacob Faibussowitsch PetscCall(PetscTableCreate(an,aN,&ta)); /* for compute Crmax */ 13794ede91e3SHong Zhang MatRowMergeMax_SeqAIJ(a_loc,ptap->A_loc->rmap->N,ta); 1380669e9c86SHong Zhang 1381912837bbSHong Zhang for (k=0; k<nrecv; k++) {/* k-th received message */ 1382912837bbSHong Zhang Jptr = buf_rj[k]; 1383912837bbSHong Zhang for (j=0; j<len_r[k]; j++) { 13849566063dSJacob Faibussowitsch PetscCall(PetscTableAdd(ta,*(Jptr+j)+1,1,INSERT_VALUES)); 1385912837bbSHong Zhang } 1386912837bbSHong Zhang } 13879566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Crmax)); 13889566063dSJacob Faibussowitsch PetscCall(PetscTableDestroy(&ta)); 1389912837bbSHong Zhang 1390912837bbSHong Zhang /* (4) send and recv coi */ 1391912837bbSHong Zhang /*-----------------------*/ 13929566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagi)); 13939566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagi,nrecv,id_r,len_ri,&buf_ri,&rwaits)); 13949566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(len+1,&buf_s)); 1395912837bbSHong Zhang buf_si = buf_s; /* points to the beginning of k-th msg to be sent */ 1396912837bbSHong Zhang for (proc=0,k=0; proc<size; proc++) { 1397912837bbSHong Zhang if (!len_s[proc]) continue; 1398912837bbSHong Zhang /* form outgoing message for i-structure: 1399912837bbSHong Zhang buf_si[0]: nrows to be sent 1400912837bbSHong Zhang [1:nrows]: row index (global) 1401912837bbSHong Zhang [nrows+1:2*nrows+1]: i-structure index 1402912837bbSHong Zhang */ 1403912837bbSHong Zhang /*-------------------------------------------*/ 1404912837bbSHong Zhang nrows = len_si[proc]/2 - 1; /* num of rows in Co to be sent to [proc] */ 1405912837bbSHong Zhang buf_si_i = buf_si + nrows+1; 1406912837bbSHong Zhang buf_si[0] = nrows; 1407912837bbSHong Zhang buf_si_i[0] = 0; 1408912837bbSHong Zhang nrows = 0; 1409912837bbSHong Zhang for (i=owners_co[proc]; i<owners_co[proc+1]; i++) { 1410912837bbSHong Zhang nzi = coi[i+1] - coi[i]; 1411912837bbSHong Zhang buf_si_i[nrows+1] = buf_si_i[nrows] + nzi; /* i-structure */ 1412912837bbSHong Zhang buf_si[nrows+1] = prmap[i] -owners[proc]; /* local row index */ 1413912837bbSHong Zhang nrows++; 1414912837bbSHong Zhang } 14159566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tagi,comm,swaits+k)); 1416912837bbSHong Zhang k++; 1417912837bbSHong Zhang buf_si += len_si[proc]; 1418912837bbSHong Zhang } 1419912837bbSHong Zhang for (i=0; i<nrecv; i++) { 14209566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(nrecv,rwaits,&icompleted,&rstatus)); 1421912837bbSHong Zhang } 14229566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 14239566063dSJacob Faibussowitsch if (nsend) PetscCallMPI(MPI_Waitall(nsend,swaits,sstatus)); 1424912837bbSHong Zhang 14259566063dSJacob Faibussowitsch PetscCall(PetscFree4(len_s,len_si,sstatus,owners_co)); 14269566063dSJacob Faibussowitsch PetscCall(PetscFree(len_ri)); 14279566063dSJacob Faibussowitsch PetscCall(PetscFree(swaits)); 14289566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_s)); 1429912837bbSHong Zhang 14304222ddf1SHong Zhang /* (5) compute the local portion of C */ 1431912837bbSHong Zhang /* ------------------------------------------ */ 14324222ddf1SHong Zhang /* set initial free space to be Crmax, sufficient for holding nozeros in each row of C */ 14339566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(Crmax,&free_space)); 1434912837bbSHong Zhang current_space = free_space; 1435912837bbSHong Zhang 14369566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(nrecv,&buf_ri_k,nrecv,&nextrow,nrecv,&nextci)); 1437912837bbSHong Zhang for (k=0; k<nrecv; k++) { 1438912837bbSHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 1439912837bbSHong Zhang nrows = *buf_ri_k[k]; 1440912837bbSHong Zhang nextrow[k] = buf_ri_k[k] + 1; /* next row number of k-th recved i-structure */ 1441a5b23f4aSJose E. Roman nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th recved i-structure */ 1442912837bbSHong Zhang } 1443912837bbSHong Zhang 1444*d0609cedSBarry Smith MatPreallocateBegin(comm,pn,an,dnz,onz); 14459566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate(Crmax,aN,&lnk,&lnkbt)); 14468c78258cSHong Zhang for (i=0; i<pn; i++) { /* for each local row of C */ 14474222ddf1SHong Zhang /* add C_loc into C */ 1448912837bbSHong Zhang nzi = c_loc->i[i+1] - c_loc->i[i]; 1449912837bbSHong Zhang Jptr = c_loc->j + c_loc->i[i]; 14509566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(nzi,Jptr,lnk,lnkbt)); 1451912837bbSHong Zhang 1452912837bbSHong Zhang /* add received col data into lnk */ 1453912837bbSHong Zhang for (k=0; k<nrecv; k++) { /* k-th received message */ 1454912837bbSHong Zhang if (i == *nextrow[k]) { /* i-th row */ 1455912837bbSHong Zhang nzi = *(nextci[k]+1) - *nextci[k]; 1456912837bbSHong Zhang Jptr = buf_rj[k] + *nextci[k]; 14579566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(nzi,Jptr,lnk,lnkbt)); 1458912837bbSHong Zhang nextrow[k]++; nextci[k]++; 1459912837bbSHong Zhang } 1460912837bbSHong Zhang } 14618c78258cSHong Zhang 14628c78258cSHong Zhang /* add missing diagonal entry */ 14638c78258cSHong Zhang if (C->force_diagonals) { 14648c78258cSHong Zhang k = i + owners[rank]; /* column index */ 14659566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted(1,&k,lnk,lnkbt)); 14668c78258cSHong Zhang } 14678c78258cSHong Zhang 1468912837bbSHong Zhang nzi = lnk[0]; 1469912837bbSHong Zhang 1470912837bbSHong Zhang /* copy data into free space, then initialize lnk */ 14719566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean(aN,nzi,current_space->array,lnk,lnkbt)); 14729566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+owners[rank],nzi,current_space->array,dnz,onz)); 1473912837bbSHong Zhang } 14749566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 14759566063dSJacob Faibussowitsch PetscCall(PetscLLDestroy(lnk,lnkbt)); 14769566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceDestroy(free_space)); 1477912837bbSHong Zhang 1478912837bbSHong Zhang /* local sizes and preallocation */ 14799566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,pn,an,PETSC_DETERMINE,PETSC_DETERMINE)); 14809566063dSJacob Faibussowitsch if (P->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->rmap,P->cmap->bs)); 14819566063dSJacob Faibussowitsch if (A->cmap->bs > 0) PetscCall(PetscLayoutSetBlockSize(C->cmap,A->cmap->bs)); 14829566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 1483*d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 1484912837bbSHong Zhang 148569be949aSStefano Zampini /* add C_loc and C_oth to C */ 14869566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRange(C,&rstart,NULL)); 148769be949aSStefano Zampini for (i=0; i<pn; i++) { 14888c78258cSHong Zhang ncols = c_loc->i[i+1] - c_loc->i[i]; 14898c78258cSHong Zhang cols = c_loc->j + c_loc->i[i]; 14908c78258cSHong Zhang row = rstart + i; 14919566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,(const PetscInt*)&row,ncols,(const PetscInt*)cols,NULL,INSERT_VALUES)); 14928c78258cSHong Zhang 14938c78258cSHong Zhang if (C->force_diagonals) { 14949566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,(const PetscInt*)&row,1,(const PetscInt*)&row,NULL,INSERT_VALUES)); 14958c78258cSHong Zhang } 149669be949aSStefano Zampini } 149769be949aSStefano Zampini for (i=0; i<con; i++) { 14988c78258cSHong Zhang ncols = c_oth->i[i+1] - c_oth->i[i]; 14998c78258cSHong Zhang cols = c_oth->j + c_oth->i[i]; 15008c78258cSHong Zhang row = prmap[i]; 15019566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,(const PetscInt*)&row,ncols,(const PetscInt*)cols,NULL,INSERT_VALUES)); 150269be949aSStefano Zampini } 15039566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 15049566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 15059566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 150669be949aSStefano Zampini 1507912837bbSHong Zhang /* members in merge */ 15089566063dSJacob Faibussowitsch PetscCall(PetscFree(id_r)); 15099566063dSJacob Faibussowitsch PetscCall(PetscFree(len_r)); 15109566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_ri[0])); 15119566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_ri)); 15129566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_rj[0])); 15139566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_rj)); 15149566063dSJacob Faibussowitsch PetscCall(PetscLayoutDestroy(&rowmap)); 1515912837bbSHong Zhang 15164222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 15176718818eSStefano Zampini C->product->data = ptap; 15186718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 1519912837bbSHong Zhang PetscFunctionReturn(0); 1520912837bbSHong Zhang } 1521912837bbSHong Zhang 1522912837bbSHong Zhang PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_nonscalable(Mat P,Mat A,Mat C) 1523912837bbSHong Zhang { 15246718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data; 1525669e9c86SHong Zhang Mat_SeqAIJ *c_seq; 15266718818eSStefano Zampini Mat_APMPI *ptap; 1527669e9c86SHong Zhang Mat A_loc,C_loc,C_oth; 1528912837bbSHong Zhang PetscInt i,rstart,rend,cm,ncols,row; 1529912837bbSHong Zhang const PetscInt *cols; 1530912837bbSHong Zhang const PetscScalar *vals; 1531912837bbSHong Zhang 1532912837bbSHong Zhang PetscFunctionBegin; 15336718818eSStefano Zampini MatCheckProduct(C,3); 15346718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 153528b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 153628b400f6SJacob Faibussowitsch PetscCheck(ptap->A_loc,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be reused. Do not call MatProductClear()"); 15379566063dSJacob Faibussowitsch PetscCall(MatZeroEntries(C)); 1538912837bbSHong Zhang 1539912837bbSHong Zhang if (ptap->reuse == MAT_REUSE_MATRIX) { 1540669e9c86SHong Zhang /* These matrices are obtained in MatTransposeMatMultSymbolic() */ 1541669e9c86SHong Zhang /* 1) get R = Pd^T, Ro = Po^T */ 1542669e9c86SHong Zhang /*----------------------------*/ 15439566063dSJacob Faibussowitsch PetscCall(MatTranspose_SeqAIJ(p->A,MAT_REUSE_MATRIX,&ptap->Rd)); 15449566063dSJacob Faibussowitsch PetscCall(MatTranspose_SeqAIJ(p->B,MAT_REUSE_MATRIX,&ptap->Ro)); 1545669e9c86SHong Zhang 1546669e9c86SHong Zhang /* 2) compute numeric A_loc */ 1547669e9c86SHong Zhang /*--------------------------*/ 15489566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_REUSE_MATRIX,&ptap->A_loc)); 1549912837bbSHong Zhang } 1550912837bbSHong Zhang 1551669e9c86SHong Zhang /* 3) C_loc = Rd*A_loc, C_oth = Ro*A_loc */ 15524ede91e3SHong Zhang A_loc = ptap->A_loc; 15539566063dSJacob Faibussowitsch PetscCall(((ptap->C_loc)->ops->matmultnumeric)(ptap->Rd,A_loc,ptap->C_loc)); 15549566063dSJacob Faibussowitsch PetscCall(((ptap->C_oth)->ops->matmultnumeric)(ptap->Ro,A_loc,ptap->C_oth)); 1555912837bbSHong Zhang C_loc = ptap->C_loc; 1556912837bbSHong Zhang C_oth = ptap->C_oth; 1557912837bbSHong Zhang 155869be949aSStefano Zampini /* add C_loc and C_oth to C */ 15599566063dSJacob Faibussowitsch PetscCall(MatGetOwnershipRange(C,&rstart,&rend)); 1560912837bbSHong Zhang 1561912837bbSHong Zhang /* C_loc -> C */ 1562912837bbSHong Zhang cm = C_loc->rmap->N; 1563912837bbSHong Zhang c_seq = (Mat_SeqAIJ*)C_loc->data; 1564912837bbSHong Zhang cols = c_seq->j; 1565912837bbSHong Zhang vals = c_seq->a; 1566912837bbSHong Zhang for (i=0; i<cm; i++) { 1567912837bbSHong Zhang ncols = c_seq->i[i+1] - c_seq->i[i]; 1568912837bbSHong Zhang row = rstart + i; 15699566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,ncols,cols,vals,ADD_VALUES)); 1570912837bbSHong Zhang cols += ncols; vals += ncols; 1571912837bbSHong Zhang } 1572912837bbSHong Zhang 1573912837bbSHong Zhang /* Co -> C, off-processor part */ 1574912837bbSHong Zhang cm = C_oth->rmap->N; 1575912837bbSHong Zhang c_seq = (Mat_SeqAIJ*)C_oth->data; 1576912837bbSHong Zhang cols = c_seq->j; 1577912837bbSHong Zhang vals = c_seq->a; 1578912837bbSHong Zhang for (i=0; i<cm; i++) { 1579912837bbSHong Zhang ncols = c_seq->i[i+1] - c_seq->i[i]; 1580912837bbSHong Zhang row = p->garray[i]; 15819566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,ncols,cols,vals,ADD_VALUES)); 1582912837bbSHong Zhang cols += ncols; vals += ncols; 1583912837bbSHong Zhang } 15849566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 15859566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 15869566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 1587912837bbSHong Zhang 1588912837bbSHong Zhang ptap->reuse = MAT_REUSE_MATRIX; 1589912837bbSHong Zhang PetscFunctionReturn(0); 1590912837bbSHong Zhang } 1591912837bbSHong Zhang 15926da69ca6SHong Zhang PetscErrorCode MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ(Mat P,Mat A,Mat C) 1593d6ab1dc0SHong Zhang { 1594d6ab1dc0SHong Zhang Mat_Merge_SeqsToMPI *merge; 15956718818eSStefano Zampini Mat_MPIAIJ *p =(Mat_MPIAIJ*)P->data; 1596d6ab1dc0SHong Zhang Mat_SeqAIJ *pd=(Mat_SeqAIJ*)(p->A)->data,*po=(Mat_SeqAIJ*)(p->B)->data; 15973cdca5ebSHong Zhang Mat_APMPI *ptap; 1598e745005bSHong Zhang PetscInt *adj; 1599e745005bSHong Zhang PetscInt i,j,k,anz,pnz,row,*cj,nexta; 1600e745005bSHong Zhang MatScalar *ada,*ca,valtmp; 1601d6ab1dc0SHong Zhang PetscInt am=A->rmap->n,cm=C->rmap->n,pon=(p->B)->cmap->n; 1602ce94432eSBarry Smith MPI_Comm comm; 1603d6ab1dc0SHong Zhang PetscMPIInt size,rank,taga,*len_s; 1604d6ab1dc0SHong Zhang PetscInt *owners,proc,nrows,**buf_ri_k,**nextrow,**nextci; 1605d6ab1dc0SHong Zhang PetscInt **buf_ri,**buf_rj; 1606d6ab1dc0SHong Zhang PetscInt cnz=0,*bj_i,*bi,*bj,bnz,nextcj; /* bi,bj,ba: local array of C(mpi mat) */ 1607d6ab1dc0SHong Zhang MPI_Request *s_waits,*r_waits; 1608d6ab1dc0SHong Zhang MPI_Status *status; 1609d6ab1dc0SHong Zhang MatScalar **abuf_r,*ba_i,*pA,*coa,*ba; 1610ce496241SStefano Zampini const PetscScalar *dummy; 16114e201bd7SHong Zhang PetscInt *ai,*aj,*coi,*coj,*poJ,*pdJ; 1612d6ab1dc0SHong Zhang Mat A_loc; 1613d6ab1dc0SHong Zhang Mat_SeqAIJ *a_loc; 1614d6ab1dc0SHong Zhang 1615d6ab1dc0SHong Zhang PetscFunctionBegin; 16166718818eSStefano Zampini MatCheckProduct(C,3); 16176718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 161828b400f6SJacob Faibussowitsch PetscCheck(ptap,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtAP cannot be computed. Missing data"); 161928b400f6SJacob Faibussowitsch PetscCheck(ptap->A_loc,PetscObjectComm((PetscObject)C),PETSC_ERR_ARG_WRONGSTATE,"PtA cannot be reused. Do not call MatProductClear()"); 16209566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 16219566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 16229566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1623d6ab1dc0SHong Zhang 1624d6ab1dc0SHong Zhang merge = ptap->merge; 1625d6ab1dc0SHong Zhang 1626e745005bSHong Zhang /* 2) compute numeric C_seq = P_loc^T*A_loc */ 1627e745005bSHong Zhang /*------------------------------------------*/ 1628d6ab1dc0SHong Zhang /* get data from symbolic products */ 1629d6ab1dc0SHong Zhang coi = merge->coi; coj = merge->coj; 16309566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(coi[pon]+1,&coa)); 1631d6ab1dc0SHong Zhang bi = merge->bi; bj = merge->bj; 1632d6ab1dc0SHong Zhang owners = merge->rowmap->range; 16339566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(bi[cm]+1,&ba)); 1634d6ab1dc0SHong Zhang 1635d6ab1dc0SHong Zhang /* get A_loc by taking all local rows of A */ 1636d6ab1dc0SHong Zhang A_loc = ptap->A_loc; 16379566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_REUSE_MATRIX,&A_loc)); 1638d6ab1dc0SHong Zhang a_loc = (Mat_SeqAIJ*)(A_loc)->data; 1639d6ab1dc0SHong Zhang ai = a_loc->i; 1640d6ab1dc0SHong Zhang aj = a_loc->j; 1641d6ab1dc0SHong Zhang 1642ce496241SStefano Zampini /* trigger copy to CPU */ 16439566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(p->A,&dummy)); 16449566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(p->A,&dummy)); 16459566063dSJacob Faibussowitsch PetscCall(MatSeqAIJGetArrayRead(p->B,&dummy)); 16469566063dSJacob Faibussowitsch PetscCall(MatSeqAIJRestoreArrayRead(p->B,&dummy)); 1647d6ab1dc0SHong Zhang for (i=0; i<am; i++) { 1648d6ab1dc0SHong Zhang anz = ai[i+1] - ai[i]; 1649d6ab1dc0SHong Zhang adj = aj + ai[i]; 1650d6ab1dc0SHong Zhang ada = a_loc->a + ai[i]; 1651d6ab1dc0SHong Zhang 1652d6ab1dc0SHong Zhang /* 2-b) Compute Cseq = P_loc[i,:]^T*A[i,:] using outer product */ 1653e745005bSHong Zhang /*-------------------------------------------------------------*/ 1654d6ab1dc0SHong Zhang /* put the value into Co=(p->B)^T*A (off-diagonal part, send to others) */ 1655d6ab1dc0SHong Zhang pnz = po->i[i+1] - po->i[i]; 1656d6ab1dc0SHong Zhang poJ = po->j + po->i[i]; 1657d6ab1dc0SHong Zhang pA = po->a + po->i[i]; 1658d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1659d6ab1dc0SHong Zhang row = poJ[j]; 1660d6ab1dc0SHong Zhang cj = coj + coi[row]; 1661d6ab1dc0SHong Zhang ca = coa + coi[row]; 1662e745005bSHong Zhang /* perform sparse axpy */ 1663e745005bSHong Zhang nexta = 0; 1664d6ab1dc0SHong Zhang valtmp = pA[j]; 1665e745005bSHong Zhang for (k=0; nexta<anz; k++) { 1666e745005bSHong Zhang if (cj[k] == adj[nexta]) { 1667e745005bSHong Zhang ca[k] += valtmp*ada[nexta]; 1668e745005bSHong Zhang nexta++; 1669d6ab1dc0SHong Zhang } 1670e745005bSHong Zhang } 16719566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*anz)); 1672d6ab1dc0SHong Zhang } 1673d6ab1dc0SHong Zhang 1674d6ab1dc0SHong Zhang /* put the value into Cd (diagonal part) */ 1675d6ab1dc0SHong Zhang pnz = pd->i[i+1] - pd->i[i]; 1676d6ab1dc0SHong Zhang pdJ = pd->j + pd->i[i]; 1677d6ab1dc0SHong Zhang pA = pd->a + pd->i[i]; 1678d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1679d6ab1dc0SHong Zhang row = pdJ[j]; 1680d6ab1dc0SHong Zhang cj = bj + bi[row]; 1681d6ab1dc0SHong Zhang ca = ba + bi[row]; 1682e745005bSHong Zhang /* perform sparse axpy */ 1683e745005bSHong Zhang nexta = 0; 1684d6ab1dc0SHong Zhang valtmp = pA[j]; 1685e745005bSHong Zhang for (k=0; nexta<anz; k++) { 1686e745005bSHong Zhang if (cj[k] == adj[nexta]) { 1687e745005bSHong Zhang ca[k] += valtmp*ada[nexta]; 1688e745005bSHong Zhang nexta++; 1689d6ab1dc0SHong Zhang } 1690e745005bSHong Zhang } 16919566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*anz)); 1692d6ab1dc0SHong Zhang } 1693d6ab1dc0SHong Zhang } 1694d6ab1dc0SHong Zhang 1695d6ab1dc0SHong Zhang /* 3) send and recv matrix values coa */ 1696d6ab1dc0SHong Zhang /*------------------------------------*/ 1697d6ab1dc0SHong Zhang buf_ri = merge->buf_ri; 1698d6ab1dc0SHong Zhang buf_rj = merge->buf_rj; 1699d6ab1dc0SHong Zhang len_s = merge->len_s; 17009566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&taga)); 17019566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvScalar(comm,taga,merge->nrecv,merge->id_r,merge->len_r,&abuf_r,&r_waits)); 1702d6ab1dc0SHong Zhang 17039566063dSJacob Faibussowitsch PetscCall(PetscMalloc2(merge->nsend+1,&s_waits,size,&status)); 1704d6ab1dc0SHong Zhang for (proc=0,k=0; proc<size; proc++) { 1705d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1706d6ab1dc0SHong Zhang i = merge->owners_co[proc]; 17079566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coa+coi[i],len_s[proc],MPIU_MATSCALAR,proc,taga,comm,s_waits+k)); 1708d6ab1dc0SHong Zhang k++; 1709d6ab1dc0SHong Zhang } 17109566063dSJacob Faibussowitsch if (merge->nrecv) PetscCallMPI(MPI_Waitall(merge->nrecv,r_waits,status)); 17119566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,s_waits,status)); 1712d6ab1dc0SHong Zhang 17139566063dSJacob Faibussowitsch PetscCall(PetscFree2(s_waits,status)); 17149566063dSJacob Faibussowitsch PetscCall(PetscFree(r_waits)); 17159566063dSJacob Faibussowitsch PetscCall(PetscFree(coa)); 1716d6ab1dc0SHong Zhang 1717d6ab1dc0SHong Zhang /* 4) insert local Cseq and received values into Cmpi */ 1718d6ab1dc0SHong Zhang /*----------------------------------------------------*/ 17199566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(merge->nrecv,&buf_ri_k,merge->nrecv,&nextrow,merge->nrecv,&nextci)); 1720d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { 1721e745005bSHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 1722d6ab1dc0SHong Zhang nrows = *(buf_ri_k[k]); 1723d6ab1dc0SHong Zhang nextrow[k] = buf_ri_k[k]+1; /* next row number of k-th recved i-structure */ 1724a5b23f4aSJose E. Roman nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th recved i-structure */ 1725d6ab1dc0SHong Zhang } 1726d6ab1dc0SHong Zhang 1727d6ab1dc0SHong Zhang for (i=0; i<cm; i++) { 1728d6ab1dc0SHong Zhang row = owners[rank] + i; /* global row index of C_seq */ 1729d6ab1dc0SHong Zhang bj_i = bj + bi[i]; /* col indices of the i-th row of C */ 1730d6ab1dc0SHong Zhang ba_i = ba + bi[i]; 1731d6ab1dc0SHong Zhang bnz = bi[i+1] - bi[i]; 1732d6ab1dc0SHong Zhang /* add received vals into ba */ 1733d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { /* k-th received message */ 1734d6ab1dc0SHong Zhang /* i-th row */ 1735d6ab1dc0SHong Zhang if (i == *nextrow[k]) { 1736d6ab1dc0SHong Zhang cnz = *(nextci[k]+1) - *nextci[k]; 1737d6ab1dc0SHong Zhang cj = buf_rj[k] + *(nextci[k]); 1738d6ab1dc0SHong Zhang ca = abuf_r[k] + *(nextci[k]); 1739d6ab1dc0SHong Zhang nextcj = 0; 1740d6ab1dc0SHong Zhang for (j=0; nextcj<cnz; j++) { 1741d6ab1dc0SHong Zhang if (bj_i[j] == cj[nextcj]) { /* bcol == ccol */ 1742d6ab1dc0SHong Zhang ba_i[j] += ca[nextcj++]; 1743d6ab1dc0SHong Zhang } 1744d6ab1dc0SHong Zhang } 1745d6ab1dc0SHong Zhang nextrow[k]++; nextci[k]++; 17469566063dSJacob Faibussowitsch PetscCall(PetscLogFlops(2.0*cnz)); 1747d6ab1dc0SHong Zhang } 1748d6ab1dc0SHong Zhang } 17499566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,bnz,bj_i,ba_i,INSERT_VALUES)); 1750d6ab1dc0SHong Zhang } 17519566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 17529566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 1753d6ab1dc0SHong Zhang 17549566063dSJacob Faibussowitsch PetscCall(PetscFree(ba)); 17559566063dSJacob Faibussowitsch PetscCall(PetscFree(abuf_r[0])); 17569566063dSJacob Faibussowitsch PetscCall(PetscFree(abuf_r)); 17579566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 1758d6ab1dc0SHong Zhang PetscFunctionReturn(0); 1759d6ab1dc0SHong Zhang } 1760d6ab1dc0SHong Zhang 17614222ddf1SHong Zhang PetscErrorCode MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ(Mat P,Mat A,PetscReal fill,Mat C) 1762d6ab1dc0SHong Zhang { 1763b419842dSStefano Zampini Mat A_loc; 17643cdca5ebSHong Zhang Mat_APMPI *ptap; 17650298fd71SBarry Smith PetscFreeSpaceList free_space=NULL,current_space=NULL; 17666718818eSStefano Zampini Mat_MPIAIJ *p=(Mat_MPIAIJ*)P->data,*a=(Mat_MPIAIJ*)A->data; 1767d6ab1dc0SHong Zhang PetscInt *pdti,*pdtj,*poti,*potj,*ptJ; 1768d6ab1dc0SHong Zhang PetscInt nnz; 1769d6ab1dc0SHong Zhang PetscInt *lnk,*owners_co,*coi,*coj,i,k,pnz,row; 1770d6ab1dc0SHong Zhang PetscInt am =A->rmap->n,pn=P->cmap->n; 1771ce94432eSBarry Smith MPI_Comm comm; 1772d6ab1dc0SHong Zhang PetscMPIInt size,rank,tagi,tagj,*len_si,*len_s,*len_ri; 1773d6ab1dc0SHong Zhang PetscInt **buf_rj,**buf_ri,**buf_ri_k; 1774d6ab1dc0SHong Zhang PetscInt len,proc,*dnz,*onz,*owners; 1775d6ab1dc0SHong Zhang PetscInt nzi,*bi,*bj; 1776d6ab1dc0SHong Zhang PetscInt nrows,*buf_s,*buf_si,*buf_si_i,**nextrow,**nextci; 1777d6ab1dc0SHong Zhang MPI_Request *swaits,*rwaits; 1778d6ab1dc0SHong Zhang MPI_Status *sstatus,rstatus; 1779d6ab1dc0SHong Zhang Mat_Merge_SeqsToMPI *merge; 1780d6ab1dc0SHong Zhang PetscInt *ai,*aj,*Jptr,anz,*prmap=p->garray,pon,nspacedouble=0,j; 1781d6ab1dc0SHong Zhang PetscReal afill =1.0,afill_tmp; 17824b38b95cSHong Zhang PetscInt rstart = P->cmap->rstart,rmax,aN=A->cmap->N,Armax; 1783b419842dSStefano Zampini Mat_SeqAIJ *a_loc; 17840acc65b4SHong Zhang PetscTable ta; 1785b92f168fSBarry Smith MatType mtype; 1786d6ab1dc0SHong Zhang 1787d6ab1dc0SHong Zhang PetscFunctionBegin; 17889566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)A,&comm)); 1789d6ab1dc0SHong Zhang /* check if matrix local sizes are compatible */ 17902c71b3e2SJacob Faibussowitsch PetscCheckFalse(A->rmap->rstart != P->rmap->rstart || A->rmap->rend != P->rmap->rend,comm,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, A (%" PetscInt_FMT ", %" PetscInt_FMT ") != P (%" PetscInt_FMT ",%" PetscInt_FMT ")",A->rmap->rstart,A->rmap->rend,P->rmap->rstart,P->rmap->rend); 1791d6ab1dc0SHong Zhang 17929566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm,&size)); 17939566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(comm,&rank)); 1794d6ab1dc0SHong Zhang 17953cdca5ebSHong Zhang /* create struct Mat_APMPI and attached it to C later */ 17969566063dSJacob Faibussowitsch PetscCall(PetscNew(&ptap)); 1797d6ab1dc0SHong Zhang 1798d6ab1dc0SHong Zhang /* get A_loc by taking all local rows of A */ 17999566063dSJacob Faibussowitsch PetscCall(MatMPIAIJGetLocalMat(A,MAT_INITIAL_MATRIX,&A_loc)); 18002205254eSKarl Rupp 1801d6ab1dc0SHong Zhang ptap->A_loc = A_loc; 1802d6ab1dc0SHong Zhang a_loc = (Mat_SeqAIJ*)(A_loc)->data; 1803d6ab1dc0SHong Zhang ai = a_loc->i; 1804d6ab1dc0SHong Zhang aj = a_loc->j; 1805d6ab1dc0SHong Zhang 1806d6ab1dc0SHong Zhang /* determine symbolic Co=(p->B)^T*A - send to others */ 1807d6ab1dc0SHong Zhang /*----------------------------------------------------*/ 18089566063dSJacob Faibussowitsch PetscCall(MatGetSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj)); 18099566063dSJacob Faibussowitsch PetscCall(MatGetSymbolicTranspose_SeqAIJ(p->B,&poti,&potj)); 1810d6ab1dc0SHong Zhang pon = (p->B)->cmap->n; /* total num of rows to be sent to other processors 1811d6ab1dc0SHong Zhang >= (num of nonzero rows of C_seq) - pn */ 18129566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pon+1,&coi)); 1813d6ab1dc0SHong Zhang coi[0] = 0; 1814d6ab1dc0SHong Zhang 1815d6ab1dc0SHong Zhang /* set initial free space to be fill*(nnz(p->B) + nnz(A)) */ 1816f91af8c7SBarry Smith nnz = PetscRealIntMultTruncate(fill,PetscIntSumTruncate(poti[pon],ai[am])); 18179566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(nnz,&free_space)); 1818d6ab1dc0SHong Zhang current_space = free_space; 181919f0e57aSHong Zhang 1820d6ab1dc0SHong Zhang /* create and initialize a linked list */ 18219566063dSJacob Faibussowitsch PetscCall(PetscTableCreate(A->cmap->n + a->B->cmap->N,aN,&ta)); 18224b38b95cSHong Zhang MatRowMergeMax_SeqAIJ(a_loc,am,ta); 18239566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Armax)); 182433ba5e2fSHong Zhang 18259566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(Armax,&lnk)); 1826d6ab1dc0SHong Zhang 1827d6ab1dc0SHong Zhang for (i=0; i<pon; i++) { 1828d6ab1dc0SHong Zhang pnz = poti[i+1] - poti[i]; 1829d6ab1dc0SHong Zhang ptJ = potj + poti[i]; 1830d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 1831d6ab1dc0SHong Zhang row = ptJ[j]; /* row of A_loc == col of Pot */ 1832d6ab1dc0SHong Zhang anz = ai[row+1] - ai[row]; 1833d6ab1dc0SHong Zhang Jptr = aj + ai[row]; 1834d6ab1dc0SHong Zhang /* add non-zero cols of AP into the sorted linked list lnk */ 18359566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(anz,Jptr,lnk)); 1836d6ab1dc0SHong Zhang } 1837d6ab1dc0SHong Zhang nnz = lnk[0]; 1838d6ab1dc0SHong Zhang 1839d6ab1dc0SHong Zhang /* If free space is not available, double the total space in the list */ 1840d6ab1dc0SHong Zhang if (current_space->local_remaining<nnz) { 18419566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(nnz,current_space->total_array_size),¤t_space)); 1842d6ab1dc0SHong Zhang nspacedouble++; 1843d6ab1dc0SHong Zhang } 1844d6ab1dc0SHong Zhang 1845d6ab1dc0SHong Zhang /* Copy data into free space, and zero out denserows */ 18469566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(nnz,current_space->array,lnk)); 18472205254eSKarl Rupp 1848d6ab1dc0SHong Zhang current_space->array += nnz; 1849d6ab1dc0SHong Zhang current_space->local_used += nnz; 1850d6ab1dc0SHong Zhang current_space->local_remaining -= nnz; 18512205254eSKarl Rupp 1852d6ab1dc0SHong Zhang coi[i+1] = coi[i] + nnz; 1853d6ab1dc0SHong Zhang } 1854d6ab1dc0SHong Zhang 18559566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(coi[pon]+1,&coj)); 18569566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,coj)); 18579566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); /* must destroy to get a new one for C */ 18582205254eSKarl Rupp 1859118727c9SMark F. Adams afill_tmp = (PetscReal)coi[pon]/(poti[pon] + ai[am]+1); 1860d6ab1dc0SHong Zhang if (afill_tmp > afill) afill = afill_tmp; 1861d6ab1dc0SHong Zhang 1862d6ab1dc0SHong Zhang /* send j-array (coj) of Co to other processors */ 1863d6ab1dc0SHong Zhang /*----------------------------------------------*/ 1864d6ab1dc0SHong Zhang /* determine row ownership */ 18659566063dSJacob Faibussowitsch PetscCall(PetscNew(&merge)); 18669566063dSJacob Faibussowitsch PetscCall(PetscLayoutCreate(comm,&merge->rowmap)); 18672205254eSKarl Rupp 1868d6ab1dc0SHong Zhang merge->rowmap->n = pn; 1869d6ab1dc0SHong Zhang merge->rowmap->bs = 1; 18702205254eSKarl Rupp 18719566063dSJacob Faibussowitsch PetscCall(PetscLayoutSetUp(merge->rowmap)); 1872d6ab1dc0SHong Zhang owners = merge->rowmap->range; 1873d6ab1dc0SHong Zhang 1874d6ab1dc0SHong Zhang /* determine the number of messages to send, their lengths */ 18759566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(size,&len_si)); 18769566063dSJacob Faibussowitsch PetscCall(PetscCalloc1(size,&merge->len_s)); 18772205254eSKarl Rupp 1878d6ab1dc0SHong Zhang len_s = merge->len_s; 1879d6ab1dc0SHong Zhang merge->nsend = 0; 1880d6ab1dc0SHong Zhang 18819566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size+2,&owners_co)); 1882d6ab1dc0SHong Zhang 1883d6ab1dc0SHong Zhang proc = 0; 1884d6ab1dc0SHong Zhang for (i=0; i<pon; i++) { 1885d6ab1dc0SHong Zhang while (prmap[i] >= owners[proc+1]) proc++; 1886d6ab1dc0SHong Zhang len_si[proc]++; /* num of rows in Co to be sent to [proc] */ 1887d6ab1dc0SHong Zhang len_s[proc] += coi[i+1] - coi[i]; 1888d6ab1dc0SHong Zhang } 1889d6ab1dc0SHong Zhang 1890d6ab1dc0SHong Zhang len = 0; /* max length of buf_si[] */ 1891d6ab1dc0SHong Zhang owners_co[0] = 0; 1892d6ab1dc0SHong Zhang for (proc=0; proc<size; proc++) { 1893d6ab1dc0SHong Zhang owners_co[proc+1] = owners_co[proc] + len_si[proc]; 1894d6ab1dc0SHong Zhang if (len_si[proc]) { 1895d6ab1dc0SHong Zhang merge->nsend++; 1896d6ab1dc0SHong Zhang len_si[proc] = 2*(len_si[proc] + 1); 1897d6ab1dc0SHong Zhang len += len_si[proc]; 1898d6ab1dc0SHong Zhang } 1899d6ab1dc0SHong Zhang } 1900d6ab1dc0SHong Zhang 1901d6ab1dc0SHong Zhang /* determine the number and length of messages to receive for coi and coj */ 19029566063dSJacob Faibussowitsch PetscCall(PetscGatherNumberOfMessages(comm,NULL,len_s,&merge->nrecv)); 19039566063dSJacob Faibussowitsch PetscCall(PetscGatherMessageLengths2(comm,merge->nsend,merge->nrecv,len_s,len_si,&merge->id_r,&merge->len_r,&len_ri)); 1904d6ab1dc0SHong Zhang 1905d6ab1dc0SHong Zhang /* post the Irecv and Isend of coj */ 19069566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagj)); 19079566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagj,merge->nrecv,merge->id_r,merge->len_r,&buf_rj,&rwaits)); 19089566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(merge->nsend+1,&swaits)); 1909d6ab1dc0SHong Zhang for (proc=0, k=0; proc<size; proc++) { 1910d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1911d6ab1dc0SHong Zhang i = owners_co[proc]; 19129566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(coj+coi[i],len_s[proc],MPIU_INT,proc,tagj,comm,swaits+k)); 1913d6ab1dc0SHong Zhang k++; 1914d6ab1dc0SHong Zhang } 1915d6ab1dc0SHong Zhang 1916d6ab1dc0SHong Zhang /* receives and sends of coj are complete */ 19179566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(size,&sstatus)); 1918d6ab1dc0SHong Zhang for (i=0; i<merge->nrecv; i++) { 1919ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted; 19209566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(merge->nrecv,rwaits,&icompleted,&rstatus)); 1921d6ab1dc0SHong Zhang } 19229566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 19239566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,swaits,sstatus)); 1924d6ab1dc0SHong Zhang 19250acc65b4SHong Zhang /* add received column indices into table to update Armax */ 1926c4762a1bSJed Brown /* Armax can be as large as aN if a P[row,:] is dense, see src/ksp/ksp/tutorials/ex56.c! */ 19270acc65b4SHong Zhang for (k=0; k<merge->nrecv; k++) {/* k-th received message */ 19280acc65b4SHong Zhang Jptr = buf_rj[k]; 19290acc65b4SHong Zhang for (j=0; j<merge->len_r[k]; j++) { 19309566063dSJacob Faibussowitsch PetscCall(PetscTableAdd(ta,*(Jptr+j)+1,1,INSERT_VALUES)); 19310acc65b4SHong Zhang } 19320acc65b4SHong Zhang } 19339566063dSJacob Faibussowitsch PetscCall(PetscTableGetCount(ta,&Armax)); 19340acc65b4SHong Zhang 1935d6ab1dc0SHong Zhang /* send and recv coi */ 1936d6ab1dc0SHong Zhang /*-------------------*/ 19379566063dSJacob Faibussowitsch PetscCall(PetscCommGetNewTag(comm,&tagi)); 19389566063dSJacob Faibussowitsch PetscCall(PetscPostIrecvInt(comm,tagi,merge->nrecv,merge->id_r,len_ri,&buf_ri,&rwaits)); 19399566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(len+1,&buf_s)); 1940d6ab1dc0SHong Zhang buf_si = buf_s; /* points to the beginning of k-th msg to be sent */ 1941d6ab1dc0SHong Zhang for (proc=0,k=0; proc<size; proc++) { 1942d6ab1dc0SHong Zhang if (!len_s[proc]) continue; 1943d6ab1dc0SHong Zhang /* form outgoing message for i-structure: 1944d6ab1dc0SHong Zhang buf_si[0]: nrows to be sent 1945d6ab1dc0SHong Zhang [1:nrows]: row index (global) 1946d6ab1dc0SHong Zhang [nrows+1:2*nrows+1]: i-structure index 1947d6ab1dc0SHong Zhang */ 1948d6ab1dc0SHong Zhang /*-------------------------------------------*/ 1949d6ab1dc0SHong Zhang nrows = len_si[proc]/2 - 1; 1950d6ab1dc0SHong Zhang buf_si_i = buf_si + nrows+1; 1951d6ab1dc0SHong Zhang buf_si[0] = nrows; 1952d6ab1dc0SHong Zhang buf_si_i[0] = 0; 1953d6ab1dc0SHong Zhang nrows = 0; 1954d6ab1dc0SHong Zhang for (i=owners_co[proc]; i<owners_co[proc+1]; i++) { 1955d6ab1dc0SHong Zhang nzi = coi[i+1] - coi[i]; 1956d6ab1dc0SHong Zhang buf_si_i[nrows+1] = buf_si_i[nrows] + nzi; /* i-structure */ 1957d6ab1dc0SHong Zhang buf_si[nrows+1] = prmap[i] -owners[proc]; /* local row index */ 1958d6ab1dc0SHong Zhang nrows++; 1959d6ab1dc0SHong Zhang } 19609566063dSJacob Faibussowitsch PetscCallMPI(MPI_Isend(buf_si,len_si[proc],MPIU_INT,proc,tagi,comm,swaits+k)); 1961d6ab1dc0SHong Zhang k++; 1962d6ab1dc0SHong Zhang buf_si += len_si[proc]; 1963d6ab1dc0SHong Zhang } 1964d6ab1dc0SHong Zhang i = merge->nrecv; 1965d6ab1dc0SHong Zhang while (i--) { 1966ec4bef21SJose E. Roman PETSC_UNUSED PetscMPIInt icompleted; 19679566063dSJacob Faibussowitsch PetscCallMPI(MPI_Waitany(merge->nrecv,rwaits,&icompleted,&rstatus)); 1968d6ab1dc0SHong Zhang } 19699566063dSJacob Faibussowitsch PetscCall(PetscFree(rwaits)); 19709566063dSJacob Faibussowitsch if (merge->nsend) PetscCallMPI(MPI_Waitall(merge->nsend,swaits,sstatus)); 19719566063dSJacob Faibussowitsch PetscCall(PetscFree(len_si)); 19729566063dSJacob Faibussowitsch PetscCall(PetscFree(len_ri)); 19739566063dSJacob Faibussowitsch PetscCall(PetscFree(swaits)); 19749566063dSJacob Faibussowitsch PetscCall(PetscFree(sstatus)); 19759566063dSJacob Faibussowitsch PetscCall(PetscFree(buf_s)); 1976d6ab1dc0SHong Zhang 1977d6ab1dc0SHong Zhang /* compute the local portion of C (mpi mat) */ 1978d6ab1dc0SHong Zhang /*------------------------------------------*/ 1979d6ab1dc0SHong Zhang /* allocate bi array and free space for accumulating nonzero column info */ 19809566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(pn+1,&bi)); 1981d6ab1dc0SHong Zhang bi[0] = 0; 1982d6ab1dc0SHong Zhang 1983d6ab1dc0SHong Zhang /* set initial free space to be fill*(nnz(P) + nnz(AP)) */ 1984f91af8c7SBarry Smith nnz = PetscRealIntMultTruncate(fill,PetscIntSumTruncate(pdti[pn],PetscIntSumTruncate(poti[pon],ai[am]))); 19859566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(nnz,&free_space)); 1986d6ab1dc0SHong Zhang current_space = free_space; 1987d6ab1dc0SHong Zhang 19889566063dSJacob Faibussowitsch PetscCall(PetscMalloc3(merge->nrecv,&buf_ri_k,merge->nrecv,&nextrow,merge->nrecv,&nextci)); 1989d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { 1990d6ab1dc0SHong Zhang buf_ri_k[k] = buf_ri[k]; /* beginning of k-th recved i-structure */ 1991d6ab1dc0SHong Zhang nrows = *buf_ri_k[k]; 1992d6ab1dc0SHong Zhang nextrow[k] = buf_ri_k[k] + 1; /* next row number of k-th recved i-structure */ 19932d4ee042Sprj- nextci[k] = buf_ri_k[k] + (nrows + 1); /* points to the next i-structure of k-th received i-structure */ 1994d6ab1dc0SHong Zhang } 1995d6ab1dc0SHong Zhang 19969566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedCreate_Scalable(Armax,&lnk)); 1997*d0609cedSBarry Smith MatPreallocateBegin(comm,pn,A->cmap->n,dnz,onz); 1998d6ab1dc0SHong Zhang rmax = 0; 1999d6ab1dc0SHong Zhang for (i=0; i<pn; i++) { 2000d6ab1dc0SHong Zhang /* add pdt[i,:]*AP into lnk */ 2001d6ab1dc0SHong Zhang pnz = pdti[i+1] - pdti[i]; 2002d6ab1dc0SHong Zhang ptJ = pdtj + pdti[i]; 2003d6ab1dc0SHong Zhang for (j=0; j<pnz; j++) { 2004d6ab1dc0SHong Zhang row = ptJ[j]; /* row of AP == col of Pt */ 2005d6ab1dc0SHong Zhang anz = ai[row+1] - ai[row]; 2006d6ab1dc0SHong Zhang Jptr = aj + ai[row]; 2007d6ab1dc0SHong Zhang /* add non-zero cols of AP into the sorted linked list lnk */ 20089566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(anz,Jptr,lnk)); 2009d6ab1dc0SHong Zhang } 2010d6ab1dc0SHong Zhang 2011d6ab1dc0SHong Zhang /* add received col data into lnk */ 2012d6ab1dc0SHong Zhang for (k=0; k<merge->nrecv; k++) { /* k-th received message */ 2013d6ab1dc0SHong Zhang if (i == *nextrow[k]) { /* i-th row */ 2014d6ab1dc0SHong Zhang nzi = *(nextci[k]+1) - *nextci[k]; 2015d6ab1dc0SHong Zhang Jptr = buf_rj[k] + *nextci[k]; 20169566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(nzi,Jptr,lnk)); 2017d6ab1dc0SHong Zhang nextrow[k]++; nextci[k]++; 2018d6ab1dc0SHong Zhang } 2019d6ab1dc0SHong Zhang } 20208c78258cSHong Zhang 20218c78258cSHong Zhang /* add missing diagonal entry */ 20228c78258cSHong Zhang if (C->force_diagonals) { 20238c78258cSHong Zhang k = i + owners[rank]; /* column index */ 20249566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedAddSorted_Scalable(1,&k,lnk)); 20258c78258cSHong Zhang } 20268c78258cSHong Zhang 2027d6ab1dc0SHong Zhang nnz = lnk[0]; 2028d6ab1dc0SHong Zhang 2029d6ab1dc0SHong Zhang /* if free space is not available, make more free space */ 2030d6ab1dc0SHong Zhang if (current_space->local_remaining<nnz) { 20319566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceGet(PetscIntSumTruncate(nnz,current_space->total_array_size),¤t_space)); 2032d6ab1dc0SHong Zhang nspacedouble++; 2033d6ab1dc0SHong Zhang } 2034d6ab1dc0SHong Zhang /* copy data into free space, then initialize lnk */ 20359566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedClean_Scalable(nnz,current_space->array,lnk)); 20369566063dSJacob Faibussowitsch PetscCall(MatPreallocateSet(i+owners[rank],nnz,current_space->array,dnz,onz)); 20372205254eSKarl Rupp 2038d6ab1dc0SHong Zhang current_space->array += nnz; 2039d6ab1dc0SHong Zhang current_space->local_used += nnz; 2040d6ab1dc0SHong Zhang current_space->local_remaining -= nnz; 20412205254eSKarl Rupp 2042d6ab1dc0SHong Zhang bi[i+1] = bi[i] + nnz; 2043d6ab1dc0SHong Zhang if (nnz > rmax) rmax = nnz; 2044d6ab1dc0SHong Zhang } 20459566063dSJacob Faibussowitsch PetscCall(PetscFree3(buf_ri_k,nextrow,nextci)); 2046d6ab1dc0SHong Zhang 20479566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(bi[pn]+1,&bj)); 20489566063dSJacob Faibussowitsch PetscCall(PetscFreeSpaceContiguous(&free_space,bj)); 2049118727c9SMark F. Adams afill_tmp = (PetscReal)bi[pn]/(pdti[pn] + poti[pon] + ai[am]+1); 2050d6ab1dc0SHong Zhang if (afill_tmp > afill) afill = afill_tmp; 20519566063dSJacob Faibussowitsch PetscCall(PetscLLCondensedDestroy_Scalable(lnk)); 20529566063dSJacob Faibussowitsch PetscCall(PetscTableDestroy(&ta)); 20539566063dSJacob Faibussowitsch PetscCall(MatRestoreSymbolicTranspose_SeqAIJ(p->A,&pdti,&pdtj)); 20549566063dSJacob Faibussowitsch PetscCall(MatRestoreSymbolicTranspose_SeqAIJ(p->B,&poti,&potj)); 2055d6ab1dc0SHong Zhang 20564222ddf1SHong Zhang /* create symbolic parallel matrix C - why cannot be assembled in Numeric part */ 20574222ddf1SHong Zhang /*-------------------------------------------------------------------------------*/ 20589566063dSJacob Faibussowitsch PetscCall(MatSetSizes(C,pn,A->cmap->n,PETSC_DETERMINE,PETSC_DETERMINE)); 20599566063dSJacob Faibussowitsch PetscCall(MatSetBlockSizes(C,PetscAbs(P->cmap->bs),PetscAbs(A->cmap->bs))); 20609566063dSJacob Faibussowitsch PetscCall(MatGetType(A,&mtype)); 20619566063dSJacob Faibussowitsch PetscCall(MatSetType(C,mtype)); 20629566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocation(C,0,dnz,0,onz)); 2063*d0609cedSBarry Smith MatPreallocateEnd(dnz,onz); 20649566063dSJacob Faibussowitsch PetscCall(MatSetBlockSize(C,1)); 20659566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NO_OFF_PROC_ENTRIES,PETSC_TRUE)); 2066d6ab1dc0SHong Zhang for (i=0; i<pn; i++) { 2067d6ab1dc0SHong Zhang row = i + rstart; 2068d6ab1dc0SHong Zhang nnz = bi[i+1] - bi[i]; 2069d6ab1dc0SHong Zhang Jptr = bj + bi[i]; 20709566063dSJacob Faibussowitsch PetscCall(MatSetValues(C,1,&row,nnz,Jptr,NULL,INSERT_VALUES)); 2071d6ab1dc0SHong Zhang } 20729566063dSJacob Faibussowitsch PetscCall(MatAssemblyBegin(C,MAT_FINAL_ASSEMBLY)); 20739566063dSJacob Faibussowitsch PetscCall(MatAssemblyEnd(C,MAT_FINAL_ASSEMBLY)); 20749566063dSJacob Faibussowitsch PetscCall(MatSetOption(C,MAT_NEW_NONZERO_LOCATION_ERR,PETSC_TRUE)); 2075d6ab1dc0SHong Zhang merge->bi = bi; 2076d6ab1dc0SHong Zhang merge->bj = bj; 2077d6ab1dc0SHong Zhang merge->coi = coi; 2078d6ab1dc0SHong Zhang merge->coj = coj; 2079d6ab1dc0SHong Zhang merge->buf_ri = buf_ri; 2080d6ab1dc0SHong Zhang merge->buf_rj = buf_rj; 2081d6ab1dc0SHong Zhang merge->owners_co = owners_co; 2082d6ab1dc0SHong Zhang 20834222ddf1SHong Zhang /* attach the supporting struct to C for reuse */ 20846718818eSStefano Zampini C->product->data = ptap; 20856718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 2086d6ab1dc0SHong Zhang ptap->merge = merge; 2087a560ef98SHong Zhang 20884222ddf1SHong Zhang C->ops->mattransposemultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ; 2089f96d37fcSHong Zhang 2090f96d37fcSHong Zhang #if defined(PETSC_USE_INFO) 2091f96d37fcSHong Zhang if (bi[pn] != 0) { 20929566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Reallocs %" PetscInt_FMT "; Fill ratio: given %g needed %g.\n",nspacedouble,(double)fill,(double)afill)); 20939566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Use MatTransposeMatMult(A,B,MatReuse,%g,&C) for best performance.\n",(double)afill)); 2094f96d37fcSHong Zhang } else { 20959566063dSJacob Faibussowitsch PetscCall(PetscInfo(C,"Empty matrix product\n")); 2096f96d37fcSHong Zhang } 2097f96d37fcSHong Zhang #endif 2098f96d37fcSHong Zhang PetscFunctionReturn(0); 2099f96d37fcSHong Zhang } 21004222ddf1SHong Zhang 21014222ddf1SHong Zhang /* ---------------------------------------------------------------- */ 21024222ddf1SHong Zhang static PetscErrorCode MatProductSymbolic_AtB_MPIAIJ_MPIAIJ(Mat C) 21034222ddf1SHong Zhang { 21044222ddf1SHong Zhang Mat_Product *product = C->product; 21054222ddf1SHong Zhang Mat A=product->A,B=product->B; 21064222ddf1SHong Zhang PetscReal fill=product->fill; 21074222ddf1SHong Zhang PetscBool flg; 21084222ddf1SHong Zhang 21094222ddf1SHong Zhang PetscFunctionBegin; 21104222ddf1SHong Zhang /* scalable */ 21119566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"scalable",&flg)); 21124222ddf1SHong Zhang if (flg) { 21139566063dSJacob Faibussowitsch PetscCall(MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ(A,B,fill,C)); 21144222ddf1SHong Zhang goto next; 21154222ddf1SHong Zhang } 21164222ddf1SHong Zhang 21174222ddf1SHong Zhang /* nonscalable */ 21189566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"nonscalable",&flg)); 21194222ddf1SHong Zhang if (flg) { 21209566063dSJacob Faibussowitsch PetscCall(MatTransposeMatMultSymbolic_MPIAIJ_MPIAIJ_nonscalable(A,B,fill,C)); 21214222ddf1SHong Zhang goto next; 21224222ddf1SHong Zhang } 21234222ddf1SHong Zhang 21244222ddf1SHong Zhang /* matmatmult */ 21259566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"at*b",&flg)); 21264222ddf1SHong Zhang if (flg) { 21274222ddf1SHong Zhang Mat At; 21284222ddf1SHong Zhang Mat_APMPI *ptap; 21294222ddf1SHong Zhang 21309566063dSJacob Faibussowitsch PetscCall(MatTranspose(A,MAT_INITIAL_MATRIX,&At)); 21319566063dSJacob Faibussowitsch PetscCall(MatMatMultSymbolic_MPIAIJ_MPIAIJ(At,B,fill,C)); 21326718818eSStefano Zampini ptap = (Mat_APMPI*)C->product->data; 21334222ddf1SHong Zhang if (ptap) { 21344222ddf1SHong Zhang ptap->Pt = At; 21356718818eSStefano Zampini C->product->destroy = MatDestroy_MPIAIJ_PtAP; 21364222ddf1SHong Zhang } 21374222ddf1SHong Zhang C->ops->transposematmultnumeric = MatTransposeMatMultNumeric_MPIAIJ_MPIAIJ_matmatmult; 21384222ddf1SHong Zhang goto next; 21394222ddf1SHong Zhang } 21404222ddf1SHong Zhang 21414e84afc0SStefano Zampini /* backend general code */ 21429566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(product->alg,"backend",&flg)); 21434e84afc0SStefano Zampini if (flg) { 21449566063dSJacob Faibussowitsch PetscCall(MatProductSymbolic_MPIAIJBACKEND(C)); 21454e84afc0SStefano Zampini PetscFunctionReturn(0); 21464e84afc0SStefano Zampini } 21474e84afc0SStefano Zampini 21484222ddf1SHong Zhang SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"MatProduct type is not supported"); 21494222ddf1SHong Zhang 21504222ddf1SHong Zhang next: 21514222ddf1SHong Zhang C->ops->productnumeric = MatProductNumeric_AtB; 21524222ddf1SHong Zhang PetscFunctionReturn(0); 21534222ddf1SHong Zhang } 21544222ddf1SHong Zhang 21554222ddf1SHong Zhang /* ---------------------------------------------------------------- */ 21564222ddf1SHong Zhang /* Set options for MatMatMultxxx_MPIAIJ_MPIAIJ */ 21574222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_AB(Mat C) 21584222ddf1SHong Zhang { 21594222ddf1SHong Zhang Mat_Product *product = C->product; 21604222ddf1SHong Zhang Mat A=product->A,B=product->B; 21614222ddf1SHong Zhang #if defined(PETSC_HAVE_HYPRE) 21624e84afc0SStefano Zampini const char *algTypes[5] = {"scalable","nonscalable","seqmpi","backend","hypre"}; 21634e84afc0SStefano Zampini PetscInt nalg = 5; 21644222ddf1SHong Zhang #else 21654e84afc0SStefano Zampini const char *algTypes[4] = {"scalable","nonscalable","seqmpi","backend",}; 21664e84afc0SStefano Zampini PetscInt nalg = 4; 21674222ddf1SHong Zhang #endif 21684222ddf1SHong Zhang PetscInt alg = 1; /* set nonscalable algorithm as default */ 21694222ddf1SHong Zhang PetscBool flg; 21704222ddf1SHong Zhang MPI_Comm comm; 21714222ddf1SHong Zhang 21724222ddf1SHong Zhang PetscFunctionBegin; 21734222ddf1SHong Zhang /* Check matrix local sizes */ 21749566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 21752c71b3e2SJacob Faibussowitsch PetscCheckFalse(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); 21764222ddf1SHong Zhang 21774222ddf1SHong Zhang /* Set "nonscalable" as default algorithm */ 21789566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 21794222ddf1SHong Zhang if (flg) { 21809566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 21814222ddf1SHong Zhang 21824222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 21834222ddf1SHong Zhang if (B->cmap->N > 100000) { /* may switch to scalable algorithm as default */ 21844222ddf1SHong Zhang MatInfo Ainfo,Binfo; 21854222ddf1SHong Zhang PetscInt nz_local; 21864222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 21874222ddf1SHong Zhang 21889566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 21899566063dSJacob Faibussowitsch PetscCall(MatGetInfo(B,MAT_LOCAL,&Binfo)); 21904222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Binfo.nz_allocated); 21914222ddf1SHong Zhang 21924222ddf1SHong Zhang if (B->cmap->N > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 21931c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 21944222ddf1SHong Zhang 21954222ddf1SHong Zhang if (alg_scalable) { 21964222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 21979566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 21989566063dSJacob Faibussowitsch PetscCall(PetscInfo(B,"Use scalable algorithm, BN %" PetscInt_FMT ", fill*nz_allocated %g\n",B->cmap->N,(double)(product->fill*nz_local))); 21994222ddf1SHong Zhang } 22004222ddf1SHong Zhang } 22014222ddf1SHong Zhang } 22024222ddf1SHong Zhang 22034222ddf1SHong Zhang /* Get runtime option */ 22044222ddf1SHong Zhang if (product->api_user) { 2205*d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatMatMult","Mat"); 22069566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matmatmult_via","Algorithmic approach","MatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2207*d0609cedSBarry Smith PetscOptionsEnd(); 22084222ddf1SHong Zhang } else { 2209*d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_AB","Mat"); 22109566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2211*d0609cedSBarry Smith PetscOptionsEnd(); 22124222ddf1SHong Zhang } 22134222ddf1SHong Zhang if (flg) { 22149566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22154222ddf1SHong Zhang } 22164222ddf1SHong Zhang 22174222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AB_MPIAIJ_MPIAIJ; 22184222ddf1SHong Zhang PetscFunctionReturn(0); 22194222ddf1SHong Zhang } 22204222ddf1SHong Zhang 22214222ddf1SHong Zhang /* Set options for MatTransposeMatMultXXX_MPIAIJ_MPIAIJ */ 22224222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_AtB(Mat C) 22234222ddf1SHong Zhang { 22244222ddf1SHong Zhang Mat_Product *product = C->product; 22254222ddf1SHong Zhang Mat A=product->A,B=product->B; 22264e84afc0SStefano Zampini const char *algTypes[4] = {"scalable","nonscalable","at*b","backend"}; 22274e84afc0SStefano Zampini PetscInt nalg = 4; 22284222ddf1SHong Zhang PetscInt alg = 1; /* set default algorithm */ 22294222ddf1SHong Zhang PetscBool flg; 22304222ddf1SHong Zhang MPI_Comm comm; 22314222ddf1SHong Zhang 22324222ddf1SHong Zhang PetscFunctionBegin; 22334222ddf1SHong Zhang /* Check matrix local sizes */ 22349566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 22352c71b3e2SJacob Faibussowitsch PetscCheckFalse(A->rmap->rstart != B->rmap->rstart || A->rmap->rend != B->rmap->rend,PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, A (%" PetscInt_FMT ", %" PetscInt_FMT ") != B (%" PetscInt_FMT ",%" PetscInt_FMT ")",A->rmap->rstart,A->rmap->rend,B->rmap->rstart,B->rmap->rend); 22364222ddf1SHong Zhang 22374222ddf1SHong Zhang /* Set default algorithm */ 22389566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 22394222ddf1SHong Zhang if (flg) { 22409566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22414222ddf1SHong Zhang } 22424222ddf1SHong Zhang 22434222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 22444222ddf1SHong Zhang if (alg && B->cmap->N > 100000) { /* may switch to scalable algorithm as default */ 22454222ddf1SHong Zhang MatInfo Ainfo,Binfo; 22464222ddf1SHong Zhang PetscInt nz_local; 22474222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 22484222ddf1SHong Zhang 22499566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 22509566063dSJacob Faibussowitsch PetscCall(MatGetInfo(B,MAT_LOCAL,&Binfo)); 22514222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Binfo.nz_allocated); 22524222ddf1SHong Zhang 22534222ddf1SHong Zhang if (B->cmap->N > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 22541c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 22554222ddf1SHong Zhang 22564222ddf1SHong Zhang if (alg_scalable) { 22574222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 22589566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22599566063dSJacob Faibussowitsch PetscCall(PetscInfo(B,"Use scalable algorithm, BN %" PetscInt_FMT ", fill*nz_allocated %g\n",B->cmap->N,(double)(product->fill*nz_local))); 22604222ddf1SHong Zhang } 22614222ddf1SHong Zhang } 22624222ddf1SHong Zhang 22634222ddf1SHong Zhang /* Get runtime option */ 22644222ddf1SHong Zhang if (product->api_user) { 2265*d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatTransposeMatMult","Mat"); 22669566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mattransposematmult_via","Algorithmic approach","MatTransposeMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2267*d0609cedSBarry Smith PetscOptionsEnd(); 22684222ddf1SHong Zhang } else { 2269*d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_AtB","Mat"); 22709566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatTransposeMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2271*d0609cedSBarry Smith PetscOptionsEnd(); 22724222ddf1SHong Zhang } 22734222ddf1SHong Zhang if (flg) { 22749566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 22754222ddf1SHong Zhang } 22764222ddf1SHong Zhang 22774222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_AtB_MPIAIJ_MPIAIJ; 22784222ddf1SHong Zhang PetscFunctionReturn(0); 22794222ddf1SHong Zhang } 22804222ddf1SHong Zhang 22814222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_PtAP(Mat C) 22824222ddf1SHong Zhang { 22834222ddf1SHong Zhang Mat_Product *product = C->product; 22844222ddf1SHong Zhang Mat A=product->A,P=product->B; 22854222ddf1SHong Zhang MPI_Comm comm; 22864222ddf1SHong Zhang PetscBool flg; 22874222ddf1SHong Zhang PetscInt alg=1; /* set default algorithm */ 22884222ddf1SHong Zhang #if !defined(PETSC_HAVE_HYPRE) 22894e84afc0SStefano Zampini const char *algTypes[5] = {"scalable","nonscalable","allatonce","allatonce_merged","backend"}; 22904222ddf1SHong Zhang PetscInt nalg=5; 22914e84afc0SStefano Zampini #else 22924e84afc0SStefano Zampini const char *algTypes[6] = {"scalable","nonscalable","allatonce","allatonce_merged","backend","hypre"}; 22934e84afc0SStefano Zampini PetscInt nalg=6; 22944222ddf1SHong Zhang #endif 22954222ddf1SHong Zhang PetscInt pN=P->cmap->N; 22964222ddf1SHong Zhang 22974222ddf1SHong Zhang PetscFunctionBegin; 22984222ddf1SHong Zhang /* Check matrix local sizes */ 22999566063dSJacob Faibussowitsch PetscCall(PetscObjectGetComm((PetscObject)C,&comm)); 23002c71b3e2SJacob Faibussowitsch PetscCheckFalse(A->rmap->rstart != P->rmap->rstart || A->rmap->rend != P->rmap->rend,PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, Arow (%" PetscInt_FMT ", %" PetscInt_FMT ") != Prow (%" PetscInt_FMT ",%" PetscInt_FMT ")",A->rmap->rstart,A->rmap->rend,P->rmap->rstart,P->rmap->rend); 23012c71b3e2SJacob Faibussowitsch PetscCheckFalse(A->cmap->rstart != P->rmap->rstart || A->cmap->rend != P->rmap->rend,PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, Acol (%" PetscInt_FMT ", %" PetscInt_FMT ") != Prow (%" PetscInt_FMT ",%" PetscInt_FMT ")",A->cmap->rstart,A->cmap->rend,P->rmap->rstart,P->rmap->rend); 23024222ddf1SHong Zhang 23034222ddf1SHong Zhang /* Set "nonscalable" as default algorithm */ 23049566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 23054222ddf1SHong Zhang if (flg) { 23069566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23074222ddf1SHong Zhang 23084222ddf1SHong Zhang /* Set "scalable" as default if BN and local nonzeros of A and B are large */ 23094222ddf1SHong Zhang if (pN > 100000) { 23104222ddf1SHong Zhang MatInfo Ainfo,Pinfo; 23114222ddf1SHong Zhang PetscInt nz_local; 23124222ddf1SHong Zhang PetscBool alg_scalable_loc=PETSC_FALSE,alg_scalable; 23134222ddf1SHong Zhang 23149566063dSJacob Faibussowitsch PetscCall(MatGetInfo(A,MAT_LOCAL,&Ainfo)); 23159566063dSJacob Faibussowitsch PetscCall(MatGetInfo(P,MAT_LOCAL,&Pinfo)); 23164222ddf1SHong Zhang nz_local = (PetscInt)(Ainfo.nz_allocated + Pinfo.nz_allocated); 23174222ddf1SHong Zhang 23184222ddf1SHong Zhang if (pN > product->fill*nz_local) alg_scalable_loc = PETSC_TRUE; 23191c2dc1cbSBarry Smith PetscCall(MPIU_Allreduce(&alg_scalable_loc,&alg_scalable,1,MPIU_BOOL,MPI_LOR,comm)); 23204222ddf1SHong Zhang 23214222ddf1SHong Zhang if (alg_scalable) { 23224222ddf1SHong Zhang alg = 0; /* scalable algorithm would 50% slower than nonscalable algorithm */ 23239566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23244222ddf1SHong Zhang } 23254222ddf1SHong Zhang } 23264222ddf1SHong Zhang } 23274222ddf1SHong Zhang 23284222ddf1SHong Zhang /* Get runtime option */ 23294222ddf1SHong Zhang if (product->api_user) { 2330*d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatPtAP","Mat"); 23319566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matptap_via","Algorithmic approach","MatPtAP",algTypes,nalg,algTypes[alg],&alg,&flg)); 2332*d0609cedSBarry Smith PetscOptionsEnd(); 23334222ddf1SHong Zhang } else { 2334*d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_PtAP","Mat"); 23359566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatPtAP",algTypes,nalg,algTypes[alg],&alg,&flg)); 2336*d0609cedSBarry Smith PetscOptionsEnd(); 23374222ddf1SHong Zhang } 23384222ddf1SHong Zhang if (flg) { 23399566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23404222ddf1SHong Zhang } 23414222ddf1SHong Zhang 23424222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_PtAP_MPIAIJ_MPIAIJ; 23434222ddf1SHong Zhang PetscFunctionReturn(0); 23444222ddf1SHong Zhang } 23454222ddf1SHong Zhang 23464222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_RARt(Mat C) 23474222ddf1SHong Zhang { 23484222ddf1SHong Zhang Mat_Product *product = C->product; 23494222ddf1SHong Zhang Mat A = product->A,R=product->B; 23504222ddf1SHong Zhang 23514222ddf1SHong Zhang PetscFunctionBegin; 23524222ddf1SHong Zhang /* Check matrix local sizes */ 23532c71b3e2SJacob Faibussowitsch PetscCheckFalse(A->cmap->n != R->cmap->n || A->rmap->n != R->cmap->n,PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ,"Matrix local dimensions are incompatible, A local (%" PetscInt_FMT ", %" PetscInt_FMT "), R local (%" PetscInt_FMT ",%" PetscInt_FMT ")",A->rmap->n,A->rmap->n,R->rmap->n,R->cmap->n); 23544222ddf1SHong Zhang 23554222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_RARt_MPIAIJ_MPIAIJ; 23564222ddf1SHong Zhang PetscFunctionReturn(0); 23574222ddf1SHong Zhang } 23584222ddf1SHong Zhang 23594222ddf1SHong Zhang /* 23604222ddf1SHong Zhang Set options for ABC = A*B*C = A*(B*C); ABC's algorithm must be chosen from AB's algorithm 23614222ddf1SHong Zhang */ 23624222ddf1SHong Zhang static PetscErrorCode MatProductSetFromOptions_MPIAIJ_ABC(Mat C) 23634222ddf1SHong Zhang { 23644222ddf1SHong Zhang Mat_Product *product = C->product; 23654222ddf1SHong Zhang PetscBool flg = PETSC_FALSE; 23664222ddf1SHong Zhang PetscInt alg = 1; /* default algorithm */ 23674222ddf1SHong Zhang const char *algTypes[3] = {"scalable","nonscalable","seqmpi"}; 23684222ddf1SHong Zhang PetscInt nalg = 3; 23694222ddf1SHong Zhang 23704222ddf1SHong Zhang PetscFunctionBegin; 23714222ddf1SHong Zhang /* Set default algorithm */ 23729566063dSJacob Faibussowitsch PetscCall(PetscStrcmp(C->product->alg,"default",&flg)); 23734222ddf1SHong Zhang if (flg) { 23749566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23754222ddf1SHong Zhang } 23764222ddf1SHong Zhang 23774222ddf1SHong Zhang /* Get runtime option */ 23784222ddf1SHong Zhang if (product->api_user) { 2379*d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatMatMatMult","Mat"); 23809566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-matmatmatmult_via","Algorithmic approach","MatMatMatMult",algTypes,nalg,algTypes[alg],&alg,&flg)); 2381*d0609cedSBarry Smith PetscOptionsEnd(); 23824222ddf1SHong Zhang } else { 2383*d0609cedSBarry Smith PetscOptionsBegin(PetscObjectComm((PetscObject)C),((PetscObject)C)->prefix,"MatProduct_ABC","Mat"); 23849566063dSJacob Faibussowitsch PetscCall(PetscOptionsEList("-mat_product_algorithm","Algorithmic approach","MatProduct_ABC",algTypes,nalg,algTypes[alg],&alg,&flg)); 2385*d0609cedSBarry Smith PetscOptionsEnd(); 23864222ddf1SHong Zhang } 23874222ddf1SHong Zhang if (flg) { 23889566063dSJacob Faibussowitsch PetscCall(MatProductSetAlgorithm(C,(MatProductAlgorithm)algTypes[alg])); 23894222ddf1SHong Zhang } 23904222ddf1SHong Zhang 23914222ddf1SHong Zhang C->ops->matmatmultsymbolic = MatMatMatMultSymbolic_MPIAIJ_MPIAIJ_MPIAIJ; 23924222ddf1SHong Zhang C->ops->productsymbolic = MatProductSymbolic_ABC; 23934222ddf1SHong Zhang PetscFunctionReturn(0); 23944222ddf1SHong Zhang } 23954222ddf1SHong Zhang 23964222ddf1SHong Zhang PETSC_INTERN PetscErrorCode MatProductSetFromOptions_MPIAIJ(Mat C) 23974222ddf1SHong Zhang { 23984222ddf1SHong Zhang Mat_Product *product = C->product; 23994222ddf1SHong Zhang 24004222ddf1SHong Zhang PetscFunctionBegin; 24014222ddf1SHong Zhang switch (product->type) { 24024222ddf1SHong Zhang case MATPRODUCT_AB: 24039566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_AB(C)); 24044222ddf1SHong Zhang break; 24054222ddf1SHong Zhang case MATPRODUCT_AtB: 24069566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_AtB(C)); 24074222ddf1SHong Zhang break; 24084222ddf1SHong Zhang case MATPRODUCT_PtAP: 24099566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_PtAP(C)); 24104222ddf1SHong Zhang break; 24114222ddf1SHong Zhang case MATPRODUCT_RARt: 24129566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_RARt(C)); 24134222ddf1SHong Zhang break; 24144222ddf1SHong Zhang case MATPRODUCT_ABC: 24159566063dSJacob Faibussowitsch PetscCall(MatProductSetFromOptions_MPIAIJ_ABC(C)); 24164222ddf1SHong Zhang break; 24176718818eSStefano Zampini default: 24186718818eSStefano Zampini break; 24194222ddf1SHong Zhang } 24204222ddf1SHong Zhang PetscFunctionReturn(0); 24214222ddf1SHong Zhang } 2422