Date: 2025-03-29T05:58:52+00:00

Actual source code: sfgather.c

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

      4: /* Reuse the type. The difference is some fields (i.e., displs, recvcounts) are not used in Gather, which is not a big deal */
      5: typedef PetscSF_Allgatherv PetscSF_Gather;

      7: static PetscErrorCode PetscSFLinkStartCommunication_Gather(PetscSF sf, PetscSFLink link, PetscSFDirection direction)
      8: {
      9:   MPI_Comm     comm    = MPI_COMM_NULL;
     10:   void        *rootbuf = NULL, *leafbuf = NULL;
     11:   MPI_Request *req = NULL;
     12:   PetscMPIInt  count;
     13:   MPI_Datatype unit = link->unit;

     15:   PetscFunctionBegin;
     16:   if (direction == PETSCSF_ROOT2LEAF) {
     17:     PetscCall(PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host before sending */));
     18:   } else {
     19:     PetscCall(PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(sf, link, PETSC_TRUE /* device2host */));
     20:   }
     21:   PetscCall(PetscObjectGetComm((PetscObject)sf, &comm));
     22:   PetscCall(PetscMPIIntCast(sf->nroots, &count));
     23:   PetscCall(PetscSFLinkGetMPIBuffersAndRequests(sf, link, direction, &rootbuf, &leafbuf, &req, NULL));
     24:   PetscCall(PetscSFLinkSyncStreamBeforeCallMPI(sf, link));

     26:   if (direction == PETSCSF_ROOT2LEAF) {
     27:     PetscCallMPI(MPIU_Igather(rootbuf == leafbuf ? MPI_IN_PLACE : rootbuf, count, unit, leafbuf, count, unit, 0 /*rank 0*/, comm, req));
     28:   } else {
     29:     PetscCallMPI(MPIU_Iscatter(leafbuf, count, unit, rootbuf == leafbuf ? MPI_IN_PLACE : rootbuf, count, unit, 0 /*rank 0*/, comm, req));
     30:   }
     31:   PetscFunctionReturn(PETSC_SUCCESS);
     32: }

     34: static PetscErrorCode PetscSFSetCommunicationOps_Gather(PetscSF sf, PetscSFLink link)
     35: {
     36:   PetscFunctionBegin;
     37:   link->StartCommunication = PetscSFLinkStartCommunication_Gather;
     38:   PetscFunctionReturn(PETSC_SUCCESS);
     39: }

     41: PETSC_INTERN PetscErrorCode PetscSFCreate_Gather(PetscSF sf)
     42: {
     43:   PetscSF_Gather *dat = (PetscSF_Gather *)sf->data;

     45:   PetscFunctionBegin;
     46:   sf->ops->BcastBegin  = PetscSFBcastBegin_Basic;
     47:   sf->ops->BcastEnd    = PetscSFBcastEnd_Basic;
     48:   sf->ops->ReduceBegin = PetscSFReduceBegin_Basic;
     49:   sf->ops->ReduceEnd   = PetscSFReduceEnd_Basic;

     51:   /* Inherit from Allgatherv */
     52:   sf->ops->Reset         = PetscSFReset_Allgatherv;
     53:   sf->ops->Destroy       = PetscSFDestroy_Allgatherv;
     54:   sf->ops->GetGraph      = PetscSFGetGraph_Allgatherv;
     55:   sf->ops->GetRootRanks  = PetscSFGetRootRanks_Allgatherv;
     56:   sf->ops->GetLeafRanks  = PetscSFGetLeafRanks_Allgatherv;
     57:   sf->ops->FetchAndOpEnd = PetscSFFetchAndOpEnd_Allgatherv;
     58:   sf->ops->CreateLocalSF = PetscSFCreateLocalSF_Allgatherv;

     60:   /* Inherit from Allgather */
     61:   sf->ops->SetUp = PetscSFSetUp_Allgather;

     63:   /* Inherit from Gatherv */
     64:   sf->ops->FetchAndOpBegin = PetscSFFetchAndOpBegin_Gatherv;

     66:   sf->ops->SetCommunicationOps = PetscSFSetCommunicationOps_Gather;

     68:   sf->collective = PETSC_TRUE;

     70:   PetscCall(PetscNew(&dat));
     71:   sf->data = (void *)dat;
     72:   PetscFunctionReturn(PETSC_SUCCESS);
     73: }