1:  #include <../src/vec/is/sf/impls/basic/allgatherv/sfallgatherv.h>

  3: #define PetscSFPackGet_Allgather PetscSFPackGet_Allgatherv

  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;

  8: PETSC_INTERN PetscErrorCode PetscSFBcastAndOpBegin_Gather(PetscSF,MPI_Datatype,PetscMemType,const void*,PetscMemType,void*,MPI_Op);

 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;

 18:   PetscSFPackGet_Allgather(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);
 19:   PetscObjectGetComm((PetscObject)sf,&comm);
 20:   PetscMPIIntCast(sf->nroots,&sendcount);

 22:   if (op == MPIU_REPLACE) {
 23:     MPIU_Iallgather(rootdata,sendcount,unit,leafdata,sendcount,unit,comm,link->rootreqs[PETSCSF_../../../../../..2LEAF_BCAST][rootmtype]);
 24:   } else {
 25:     /* Allgather to the leaf buffer and then add leaf buffer to rootdata */
 26:     if (!link->leafbuf[leafmtype]) {PetscMallocWithMemType(leafmtype,sf->nleaves*link->unitbytes,(void**)&link->leafbuf[leafmtype]);}
 27:     MPIU_Iallgather(rootdata,sendcount,unit,link->leafbuf[leafmtype],sendcount,unit,comm,link->rootreqs[PETSCSF_../../../../../..2LEAF_BCAST][rootmtype]);
 28:   }
 29:   return(0);
 30: }

 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;

 38:   PetscSFBcastAndOpBegin_Gather(sf,unit,rootmtype,rootdata,leafmtype,leafdata,MPIU_REPLACE);
 39:   /* A simplified PetscSFBcastAndOpEnd_Allgatherv */
 40:   PetscSFPackGetInUse(sf,unit,rootdata,leafdata,PETSC_OWN_POINTER,&link);
 41:   MPI_Wait(link->rootreqs[PETSCSF_../../../../../..2LEAF_BCAST][rootmtype],MPI_STATUS_IGNORE);
 42:   PetscSFPackReclaim(sf,&link);
 43:   return(0);
 44: }

 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;

 55:   PetscSFPackGet_Allgather(sf,unit,rootmtype,rootdata,leafmtype,leafdata,&link);
 56:   PetscObjectGetComm((PetscObject)sf,&comm);
 57:   MPI_Comm_rank(comm,&rank);

 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:     PetscLayoutGetRange(sf->map,&rstart,NULL);
 62:     PetscMemcpyWithMemType(rootmtype,leafmtype,rootdata,(const char*)leafdata+(size_t)rstart*link->unitbytes,(size_t)sf->nroots*link->unitbytes);
 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]) {PetscMallocWithMemType(leafmtype,sf->nleaves*link->unitbytes,(void**)&link->leafbuf[leafmtype]);}
 66:     PetscMPIIntCast(sf->nleaves*link->bs,&count);
 67:     PetscMPIIntCast(sf->nroots,&sendcount);
 68:     MPI_Reduce(leafdata,link->leafbuf[leafmtype],count,link->basicunit,op,0/*rank 0*/,comm); /* Must do reduce with MPI builltin datatype basicunit */
 69:     if (!link->rootbuf[rootmtype]) {PetscMallocWithMemType(rootmtype,sf->nroots*link->unitbytes,(void**)&link->rootbuf[rootmtype]);} /* Allocate root buffer */
 70:     MPIU_Iscatter(link->leafbuf[leafmtype],sendcount,unit,link->rootbuf[rootmtype],sendcount,unit,0/*rank 0*/,comm,link->rootreqs[PETSCSF_LEAF2../../../../../.._REDUCE][rootmtype]);
 71:   }
 72:   return(0);
 73: }

 75: PETSC_INTERN PetscErrorCode PetscSFCreate_Allgather(PetscSF sf)
 76: {
 77:   PetscErrorCode    ierr;
 78:   PetscSF_Allgather *dat = (PetscSF_Allgather*)sf->data;


 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;

 94:   /* Allgather stuff */
 95:   sf->ops->BcastAndOpBegin = PetscSFBcastAndOpBegin_Allgather;
 96:   sf->ops->ReduceBegin     = PetscSFReduceBegin_Allgather;
 97:   sf->ops->BcastToZero     = PetscSFBcastToZero_Allgather;

 99:   PetscNewLog(sf,&dat);
100:   sf->data = (void*)dat;
101:   return(0);
102: }