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

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

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

 17:   PetscFunctionBegin;
 18:   PetscCall(PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_BCAST, &link));
 19:   PetscCall(PetscSFLinkPackRootData(sf, link, PETSCSF_REMOTE, rootdata));
 20:   PetscCall(PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */));
 21:   PetscCall(PetscObjectGetComm((PetscObject)sf, &comm));
 22:   PetscCall(PetscMPIIntCast(sf->nroots, &sendcount));
 23:   PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_ROOT2LEAF, &rootbuf, &leafbuf, &req, NULL));
 24:   PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_ROOT2LEAF));
 25:   PetscCallMPI(MPIU_Igatherv(rootbuf, sendcount, unit, leafbuf, dat->recvcounts, dat->displs, unit, 0 /*rank 0*/, comm, req));
 26:   PetscFunctionReturn(PETSC_SUCCESS);
 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:   PetscFunctionBegin;
 39:   PetscCall(PetscSFLinkCreate(sf, unit, rootmtype, rootdata, leafmtype, leafdata, op, PETSCSF_REDUCE, &link));
 40:   PetscCall(PetscSFLinkPackLeafData(sf, link, PETSCSF_REMOTE, leafdata));
 41:   PetscCall(PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */));
 42:   PetscCall(PetscObjectGetComm((PetscObject)sf, &comm));
 43:   PetscCall(PetscMPIIntCast(sf->nroots, &recvcount));
 44:   PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, PETSCSF_LEAF2ROOT, &rootbuf, &leafbuf, &req, NULL));
 45:   PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link, PETSCSF_LEAF2ROOT));
 46:   PetscCallMPI(MPIU_Iscatterv(leafbuf, dat->recvcounts, dat->displs, unit, rootbuf, recvcount, unit, 0, comm, req));
 47:   PetscFunctionReturn(PETSC_SUCCESS);
 48: }

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

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

 64:   PetscFunctionBegin;
 65:   sf->ops->BcastEnd  = PetscSFBcastEnd_Basic;
 66:   sf->ops->ReduceEnd = PetscSFReduceEnd_Basic;

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

 78:   /* Gatherv stuff */
 79:   sf->ops->BcastBegin      = PetscSFBcastBegin_Gatherv;
 80:   sf->ops->ReduceBegin     = PetscSFReduceBegin_Gatherv;
 81:   sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Gatherv;

 83:   PetscCall(PetscNew(&dat));
 84:   sf->data = (void *)dat;
 85:   PetscFunctionReturn(PETSC_SUCCESS);
 86: }