1 #include <../src/vec/is/sf/impls/basic/allgatherv/sfallgatherv.h> 2 3 #define PetscSFPackGet_Allgather PetscSFPackGet_Allgatherv 4 5 /* Reuse the type. The difference is some fields (i.e., displs, recvcounts) are not used in Allgather on rank != 0, which is not a big deal */ 6 typedef PetscSF_Allgatherv PetscSF_Allgather; 7 8 PETSC_INTERN PetscErrorCode PetscSFBcastAndOpBegin_Gather(PetscSF,MPI_Datatype,PetscMemType,const void*,PetscMemType,void*,MPI_Op); 9 10 static PetscErrorCode PetscSFBcastAndOpBegin_Allgather(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata,MPI_Op op) 11 { 12 PetscErrorCode ierr; 13 PetscSFPack link; 14 PetscMPIInt sendcount; 15 MPI_Comm comm; 16 const void *rootbuf_mpi; /* buffer used by MPI */ 17 void *leafbuf_mpi; 18 PetscMemType rootmtype_mpi,leafmtype_mpi; 19 20 PetscFunctionBegin; 21 ierr = PetscSFPackGet_Allgather(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);CHKERRQ(ierr); 22 ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr); 23 ierr = PetscMPIIntCast(sf->nroots,&sendcount);CHKERRQ(ierr); 24 ierr = PetscSFBcastPrepareMPIBuffers_Allgatherv(sf,link,op,&rootmtype_mpi,&rootbuf_mpi,&leafmtype_mpi,&leafbuf_mpi);CHKERRQ(ierr); 25 ierr = MPIU_Iallgather(rootbuf_mpi,sendcount,unit,leafbuf_mpi,sendcount,unit,comm,link->rootreqs[PETSCSF_ROOT2LEAF_BCAST][rootmtype_mpi]);CHKERRQ(ierr); 26 PetscFunctionReturn(0); 27 } 28 29 static PetscErrorCode PetscSFReduceBegin_Allgather(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op) 30 { 31 PetscErrorCode ierr; 32 PetscSFPack link; 33 PetscMPIInt sendcount; 34 PetscInt rstart; 35 MPI_Comm comm; 36 const void *leafbuf_mpi; 37 void *rootbuf_mpi; 38 PetscMemType leafmtype_mpi,rootmtype_mpi; 39 40 PetscFunctionBegin; 41 ierr = PetscSFPackGet_Allgather(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);CHKERRQ(ierr); 42 ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr); 43 44 if (op == MPIU_REPLACE) { 45 /* REPLACE is only meaningful when all processes have the same leafdata to reduce. Therefore copy from local leafdata is fine */ 46 ierr = PetscLayoutGetRange(sf->map,&rstart,NULL);CHKERRQ(ierr); 47 ierr = PetscMemcpyWithMemType(rootmtype,leafmtype,rootdata,(const char*)leafdata+(size_t)rstart*link->unitbytes,(size_t)sf->nroots*link->unitbytes);CHKERRQ(ierr); 48 } else { 49 ierr = PetscMPIIntCast(sf->nroots,&sendcount);CHKERRQ(ierr); 50 ierr = PetscSFReducePrepareMPIBuffers_Allgatherv(sf,link,op,&rootmtype_mpi,&rootbuf_mpi,&leafmtype_mpi,&leafbuf_mpi);CHKERRQ(ierr); 51 ierr = MPIU_Iscatter(leafbuf_mpi,sendcount,unit,rootbuf_mpi,sendcount,unit,0/*rank 0*/,comm,link->rootreqs[PETSCSF_LEAF2ROOT_REDUCE][rootmtype_mpi]);CHKERRQ(ierr); 52 } 53 PetscFunctionReturn(0); 54 } 55 56 static PetscErrorCode PetscSFBcastToZero_Allgather(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata) 57 { 58 PetscErrorCode ierr; 59 PetscSFPack link; 60 PetscMPIInt rank; 61 62 PetscFunctionBegin; 63 ierr = PetscSFBcastAndOpBegin_Gather(sf,unit,rootmtype,rootdata,leafmtype,leafdata,MPIU_REPLACE);CHKERRQ(ierr); 64 ierr = PetscSFPackGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr); 65 ierr = PetscSFPackWaitall(link,PETSCSF_ROOT2LEAF_BCAST);CHKERRQ(ierr); 66 ierr = MPI_Comm_rank(PetscObjectComm((PetscObject)sf),&rank);CHKERRQ(ierr); 67 if (!rank && leafmtype == PETSC_MEMTYPE_DEVICE && !use_gpu_aware_mpi) { 68 ierr = PetscMemcpyWithMemType(PETSC_MEMTYPE_DEVICE,PETSC_MEMTYPE_HOST,leafdata,link->leafbuf[PETSC_MEMTYPE_HOST],link->leafbuflen*link->unitbytes);CHKERRQ(ierr); 69 } 70 ierr = PetscSFPackReclaim(sf,&link);CHKERRQ(ierr); 71 PetscFunctionReturn(0); 72 } 73 74 PETSC_INTERN PetscErrorCode PetscSFCreate_Allgather(PetscSF sf) 75 { 76 PetscErrorCode ierr; 77 PetscSF_Allgather *dat = (PetscSF_Allgather*)sf->data; 78 79 PetscFunctionBegin; 80 81 /* Inherit from Allgatherv */ 82 sf->ops->Reset = PetscSFReset_Allgatherv; 83 sf->ops->Destroy = PetscSFDestroy_Allgatherv; 84 sf->ops->BcastAndOpEnd = PetscSFBcastAndOpEnd_Allgatherv; 85 sf->ops->ReduceEnd = PetscSFReduceEnd_Allgatherv; 86 sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Allgatherv; 87 sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Allgatherv; 88 sf->ops->GetRootRanks = PetscSFGetRootRanks_Allgatherv; 89 sf->ops->CreateLocalSF = PetscSFCreateLocalSF_Allgatherv; 90 sf->ops->GetGraph = PetscSFGetGraph_Allgatherv; 91 sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Allgatherv; 92 93 /* Allgather stuff */ 94 sf->ops->BcastAndOpBegin = PetscSFBcastAndOpBegin_Allgather; 95 sf->ops->ReduceBegin = PetscSFReduceBegin_Allgather; 96 sf->ops->BcastToZero = PetscSFBcastToZero_Allgather; 97 98 ierr = PetscNewLog(sf,&dat);CHKERRQ(ierr); 99 sf->data = (void*)dat; 100 PetscFunctionReturn(0); 101 } 102