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