Actual source code: sfgatherv.c


  2: #include <../src/vec/is/sf/impls/basic/gatherv/sfgatherv.h>

  4: /* Reuse the type. The difference is some fields (displs, recvcounts) are only significant
  5:    on rank 0 in Gatherv. On other ranks they are harmless NULL.
  6:  */
  7: typedef PetscSF_Allgatherv PetscSF_Gatherv;

  9: PETSC_INTERN PetscErrorCode PetscSFBcastBegin_Gatherv(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, const void *rootdata, PetscMemType leafmtype, void *leafdata, MPI_Op op)
 10: {
 11:   PetscSFLink      link;
 12:   PetscMPIInt      sendcount;
 13:   MPI_Comm         comm;
 14:   PetscSF_Gatherv *dat     = (PetscSF_Gatherv *)sf->data;
 15:   void            *rootbuf = NULL, *leafbuf = NULL; /* buffer seen by MPI */
 16:   MPI_Request     *req;

 18:   PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_BCAST, &link);
 19:   PetscSFLinkPackRootData(sf, link, PETSCSF_REMOTE, rootdata);
 20:   PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */);
 21:   PetscObjectGetComm((PetscObject)sf, &comm);
 22:   PetscMPIIntCast(sf->nroots, &sendcount);
 23:   PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_../../../../../..2LEAF, &rootbuf, &leafbuf, &req, NULL);
 24:   PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_../../../../../..2LEAF);
 25:   MPIU_Igatherv(rootbuf, sendcount, unit, leafbuf, dat->recvcounts, dat->displs, unit, 0 /*rank 0*/, comm, req);
 26:   return 0;
 27: }

 29: static PetscErrorCode PetscSFReduceBegin_Gatherv(PetscSF sf, MPI_Datatype unit, PetscMemType leafmtype, const void *leafdata, PetscMemType rootmtype, void *rootdata, MPI_Op op)
 30: {
 31:   PetscSFLink      link;
 32:   PetscMPIInt      recvcount;
 33:   MPI_Comm         comm;
 34:   PetscSF_Gatherv *dat     = (PetscSF_Gatherv *)sf->data;
 35:   void            *rootbuf = NULL, *leafbuf = NULL; /* buffer seen by MPI */
 36:   MPI_Request     *req;

 38:   PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_REDUCE, &link);
 39:   PetscSFLinkPackLeafData(sf, link, PETSCSF_REMOTE, leafdata);
 40:   PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */);
 41:   PetscObjectGetComm((PetscObject)sf, &comm);
 42:   PetscMPIIntCast(sf->nroots, &recvcount);
 43:   PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_LEAF2../../../../../.., &rootbuf, &leafbuf, &req, NULL);
 44:   PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_LEAF2../../../../../..);
 45:   MPIU_Iscatterv(leafbuf, dat->recvcounts, dat->displs, unit, rootbuf, recvcount, unit, 0, comm, req);
 46:   return 0;
 47: }

 49: PETSC_INTERN PetscErrorCode PetscSFFetchAndOpBegin_Gatherv(PetscSF sf, MPI_Datatype unit, PetscMemType rootmtype, void *rootdata, PetscMemType leafmtype, const void *leafdata, void *leafupdate, MPI_Op op)
 50: {
 51:   /* In Gatherv, each root only has one leaf. So we just need to bcast rootdata to leafupdate and then reduce leafdata to rootdata */
 52:   PetscSFBcastBegin(sf, unit, rootdata, leafupdate, MPI_REPLACE);
 53:   PetscSFBcastEnd(sf, unit, rootdata, leafupdate, MPI_REPLACE);
 54:   PetscSFReduceBegin(sf, unit, leafdata, rootdata, op);
 55:   return 0;
 56: }

 58: PETSC_INTERN PetscErrorCode PetscSFCreate_Gatherv(PetscSF sf)
 59: {
 60:   PetscSF_Gatherv *dat = (PetscSF_Gatherv *)sf->data;

 62:   sf->ops->BcastEnd  = PetscSFBcastEnd_Basic;
 63:   sf->ops->ReduceEnd = PetscSFReduceEnd_Basic;

 65:   /* Inherit from Allgatherv */
 66:   sf->ops->SetUp         = PetscSFSetUp_Allgatherv;
 67:   sf->ops->Reset         = PetscSFReset_Allgatherv;
 68:   sf->ops->Destroy       = PetscSFDestroy_Allgatherv;
 69:   sf->ops->GetGraph      = PetscSFGetGraph_Allgatherv;
 70:   sf->ops->GetLeafRanks  = PetscSFGetLeafRanks_Allgatherv;
 71:   sf->ops->GetRootRanks  = PetscSFGetRootRanks_Allgatherv;
 72:   sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Allgatherv;
 73:   sf->ops->CreateLocalSF = PetscSFCreateLocalSF_Allgatherv;

 75:   /* Gatherv stuff */
 76:   sf->ops->BcastBegin      = PetscSFBcastBegin_Gatherv;
 77:   sf->ops->ReduceBegin     = PetscSFReduceBegin_Gatherv;
 78:   sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Gatherv;

 80:   PetscNew(&dat);
 81:   sf->data = (void *)dat;
 82:   return 0;
 83: }