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 17 PetscFunctionBegin; 18 ierr = PetscSFPackGet_Allgather(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);CHKERRQ(ierr); 19 ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr); 20 ierr = PetscMPIIntCast(sf->nroots,&sendcount);CHKERRQ(ierr); 21 22 if (op == MPIU_REPLACE) { 23 ierr = MPIU_Iallgather(rootdata,sendcount,unit,leafdata,sendcount,unit,comm,link->rootreqs[PETSCSF_ROOT2LEAF_BCAST][rootmtype]);CHKERRQ(ierr); 24 } else { 25 /* Allgather to the leaf buffer and then add leaf buffer to rootdata */ 26 if (!link->leafbuf[leafmtype]) {ierr = PetscMallocWithMemType(leafmtype,sf->nleaves*link->unitbytes,(void**)&link->leafbuf[leafmtype]);CHKERRQ(ierr);} 27 ierr = MPIU_Iallgather(rootdata,sendcount,unit,link->leafbuf[leafmtype],sendcount,unit,comm,link->rootreqs[PETSCSF_ROOT2LEAF_BCAST][rootmtype]);CHKERRQ(ierr); 28 } 29 PetscFunctionReturn(0); 30 } 31 32 static PetscErrorCode PetscSFBcastToZero_Allgather(PetscSF sf,MPI_Datatype unit,PetscMemType rootmtype,const void *rootdata,PetscMemType leafmtype,void *leafdata) 33 { 34 PetscErrorCode ierr; 35 PetscSFPack link; 36 37 PetscFunctionBegin; 38 ierr = PetscSFBcastAndOpBegin_Gather(sf,unit,rootmtype,rootdata,leafmtype,leafdata,MPIU_REPLACE);CHKERRQ(ierr); 39 /* A simplified PetscSFBcastAndOpEnd_Allgatherv */ 40 ierr = PetscSFPackGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);CHKERRQ(ierr); 41 ierr = MPI_Wait(link->rootreqs[PETSCSF_ROOT2LEAF_BCAST][rootmtype],MPI_STATUS_IGNORE);CHKERRQ(ierr); 42 ierr = PetscSFPackReclaim(sf,&link);CHKERRQ(ierr); 43 PetscFunctionReturn(0); 44 } 45 46 static PetscErrorCode PetscSFReduceBegin_Allgather(PetscSF sf,MPI_Datatype unit,PetscMemType leafmtype,const void *leafdata,PetscMemType rootmtype,void *rootdata,MPI_Op op) 47 { 48 PetscErrorCode ierr; 49 PetscSFPack link; 50 PetscMPIInt rank,count,sendcount; 51 PetscInt rstart; 52 MPI_Comm comm; 53 54 PetscFunctionBegin; 55 ierr = PetscSFPackGet_Allgather(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);CHKERRQ(ierr); 56 ierr = PetscObjectGetComm((PetscObject)sf,&comm);CHKERRQ(ierr); 57 ierr = MPI_Comm_rank(comm,&rank);CHKERRQ(ierr); 58 59 if (op == MPIU_REPLACE) { 60 /* REPLACE is only meaningful when all processes have the same leafdata to reduce. Therefore copy from local leafdata is fine */ 61 ierr = PetscLayoutGetRange(sf->map,&rstart,NULL);CHKERRQ(ierr); 62 ierr = PetscMemcpyWithMemType(rootmtype,leafmtype,rootdata,(const char*)leafdata+(size_t)rstart*link->unitbytes,(size_t)sf->nroots*link->unitbytes);CHKERRQ(ierr); 63 } else { 64 /* Reduce all leafdata on rank 0, then scatter the result to root buffer, then reduce root buffer to leafdata */ 65 if (!rank && !link->leafbuf[leafmtype]) {ierr = PetscMallocWithMemType(leafmtype,sf->nleaves*link->unitbytes,(void**)&link->leafbuf[leafmtype]);CHKERRQ(ierr);} 66 ierr = PetscMPIIntCast(sf->nleaves*link->bs,&count);CHKERRQ(ierr); 67 ierr = PetscMPIIntCast(sf->nroots,&sendcount);CHKERRQ(ierr); 68 ierr = MPI_Reduce(leafdata,link->leafbuf[leafmtype],count,link->basicunit,op,0/*rank 0*/,comm);CHKERRQ(ierr); /* Must do reduce with MPI builltin datatype basicunit */ 69 if (!link->rootbuf[rootmtype]) {ierr = PetscMallocWithMemType(rootmtype,sf->nroots*link->unitbytes,(void**)&link->rootbuf[rootmtype]);CHKERRQ(ierr);} /* Allocate root buffer */ 70 ierr = MPIU_Iscatter(link->leafbuf[leafmtype],sendcount,unit,link->rootbuf[rootmtype],sendcount,unit,0/*rank 0*/,comm,link->rootreqs[PETSCSF_LEAF2ROOT_REDUCE][rootmtype]);CHKERRQ(ierr); 71 } 72 PetscFunctionReturn(0); 73 } 74 75 PETSC_INTERN PetscErrorCode PetscSFCreate_Allgather(PetscSF sf) 76 { 77 PetscErrorCode ierr; 78 PetscSF_Allgather *dat = (PetscSF_Allgather*)sf->data; 79 80 PetscFunctionBegin; 81 82 /* Inherit from Allgatherv */ 83 sf->ops->Reset = PetscSFReset_Allgatherv; 84 sf->ops->Destroy = PetscSFDestroy_Allgatherv; 85 sf->ops->BcastAndOpEnd = PetscSFBcastAndOpEnd_Allgatherv; 86 sf->ops->ReduceEnd = PetscSFReduceEnd_Allgatherv; 87 sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Allgatherv; 88 sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Allgatherv; 89 sf->ops->GetRootRanks = PetscSFGetRootRanks_Allgatherv; 90 sf->ops->CreateLocalSF = PetscSFCreateLocalSF_Allgatherv; 91 sf->ops->GetGraph = PetscSFGetGraph_Allgatherv; 92 sf->ops->GetLeafRanks = PetscSFGetLeafRanks_Allgatherv; 93 94 /* Allgather stuff */ 95 sf->ops->BcastAndOpBegin = PetscSFBcastAndOpBegin_Allgather; 96 sf->ops->ReduceBegin = PetscSFReduceBegin_Allgather; 97 sf->ops->BcastToZero = PetscSFBcastToZero_Allgather; 98 99 ierr = PetscNewLog(sf,&dat);CHKERRQ(ierr); 100 sf->data = (void*)dat; 101 PetscFunctionReturn(0); 102 } 103