Actual source code: sfimpl.h

      1: #pragma once

      3: #include <petscvec.h>
      4: #include <petscsf.h>
      5: #include <petsc/private/deviceimpl.h>
      6: #include <petsc/private/petscimpl.h>

      8: PETSC_EXTERN PetscLogEvent PETSCSF_SetGraph;
      9: PETSC_EXTERN PetscLogEvent PETSCSF_SetUp;
     10: PETSC_EXTERN PetscLogEvent PETSCSF_BcastBegin;
     11: PETSC_EXTERN PetscLogEvent PETSCSF_BcastEnd;
     12: PETSC_EXTERN PetscLogEvent PETSCSF_BcastBegin;
     13: PETSC_EXTERN PetscLogEvent PETSCSF_BcastEnd;
     14: PETSC_EXTERN PetscLogEvent PETSCSF_ReduceBegin;
     15: PETSC_EXTERN PetscLogEvent PETSCSF_ReduceEnd;
     16: PETSC_EXTERN PetscLogEvent PETSCSF_FetchAndOpBegin;
     17: PETSC_EXTERN PetscLogEvent PETSCSF_FetchAndOpEnd;
     18: PETSC_EXTERN PetscLogEvent PETSCSF_EmbedSF;
     19: PETSC_EXTERN PetscLogEvent PETSCSF_DistSect;
     20: PETSC_EXTERN PetscLogEvent PETSCSF_SectSF;
     21: PETSC_EXTERN PetscLogEvent PETSCSF_RemoteOff;
     22: PETSC_EXTERN PetscLogEvent PETSCSF_Pack;
     23: PETSC_EXTERN PetscLogEvent PETSCSF_Unpack;

     25: typedef enum {
     26:   PETSCSF_ROOT2LEAF = 0,
     27:   PETSCSF_LEAF2ROOT
     28: } PetscSFDirection;
     29: typedef enum {
     30:   PETSCSF_BCAST = 0,
     31:   PETSCSF_REDUCE,
     32:   PETSCSF_FETCH
     33: } PetscSFOperation;
     34: /* When doing device-aware MPI, a backend refers to the SF/device interface */
     35: typedef enum {
     36:   PETSCSF_BACKEND_INVALID = 0,
     37:   PETSCSF_BACKEND_CUDA,
     38:   PETSCSF_BACKEND_HIP,
     39:   PETSCSF_BACKEND_KOKKOS
     40: } PetscSFBackend;

     42: typedef struct _n_PetscSFLink *PetscSFLink;

     44: struct _PetscSFOps {
     45:   PetscErrorCode (*Reset)(PetscSF);
     46:   PetscErrorCode (*Destroy)(PetscSF);
     47:   PetscErrorCode (*SetUp)(PetscSF);
     48:   PetscErrorCode (*SetFromOptions)(PetscSF, PetscOptionItems *);
     49:   PetscErrorCode (*View)(PetscSF, PetscViewer);
     50:   PetscErrorCode (*Duplicate)(PetscSF, PetscSFDuplicateOption, PetscSF);
     51:   PetscErrorCode (*BcastBegin)(PetscSF, MPI_Datatype, PetscMemType, const void *, PetscMemType, void *, MPI_Op);
     52:   PetscErrorCode (*BcastEnd)(PetscSF, MPI_Datatype, const void *, void *, MPI_Op);
     53:   PetscErrorCode (*ReduceBegin)(PetscSF, MPI_Datatype, PetscMemType, const void *, PetscMemType, void *, MPI_Op);
     54:   PetscErrorCode (*ReduceEnd)(PetscSF, MPI_Datatype, const void *, void *, MPI_Op);
     55:   PetscErrorCode (*FetchAndOpBegin)(PetscSF, MPI_Datatype, PetscMemType, void *, PetscMemType, const void *, void *, MPI_Op);
     56:   PetscErrorCode (*FetchAndOpEnd)(PetscSF, MPI_Datatype, void *, const void *, void *, MPI_Op);
     57:   PetscErrorCode (*BcastToZero)(PetscSF, MPI_Datatype, PetscMemType, const void *, PetscMemType, void *); /* For internal use only */
     58:   PetscErrorCode (*GetRootRanks)(PetscSF, PetscMPIInt *, const PetscMPIInt **, const PetscInt **, const PetscInt **, const PetscInt **);
     59:   PetscErrorCode (*GetLeafRanks)(PetscSF, PetscMPIInt *, const PetscMPIInt **, const PetscInt **, const PetscInt **);
     60:   PetscErrorCode (*CreateLocalSF)(PetscSF, PetscSF *);
     61:   PetscErrorCode (*GetGraph)(PetscSF, PetscInt *, PetscInt *, const PetscInt **, const PetscSFNode **);
     62:   PetscErrorCode (*CreateEmbeddedRootSF)(PetscSF, PetscInt, const PetscInt *, PetscSF *);
     63:   PetscErrorCode (*CreateEmbeddedLeafSF)(PetscSF, PetscInt, const PetscInt *, PetscSF *);
     64:   PetscErrorCode (*SetCommunicationOps)(PetscSF, PetscSFLink);

     66:   PetscErrorCode (*Malloc)(PetscMemType, size_t, void **);
     67:   PetscErrorCode (*Free)(PetscMemType, void *);
     68: };

     70: typedef struct _n_PetscSFPackOpt *PetscSFPackOpt;

     72: struct _p_PetscSF {
     73:   PETSCHEADER(struct _PetscSFOps);
     74:   struct {                                  /* Fields needed to implement VecScatter behavior */
     75:     PetscInt           from_n, to_n;        /* Recorded local sizes of the input from/to vectors in VecScatterCreate(). Used subsequently for error checking. */
     76:     PetscBool          beginandendtogether; /* Indicates that the scatter begin and end  function are called together, VecScatterEnd() is then treated as a nop */
     77:     const PetscScalar *xdata;               /* Vector data to read from */
     78:     PetscScalar       *ydata;               /* Vector data to write to. The two pointers are recorded in VecScatterBegin. Memory is not managed by SF. */
     79:     PetscSF            lsf;                 /* The local part of the scatter, used in SCATTER_LOCAL. Built on demand. */
     80:     PetscInt           bs;                  /* Block size, determined by IS passed to VecScatterCreate */
     81:     MPI_Datatype       unit;                /* one unit = bs PetscScalars */
     82:     PetscBool          logging;             /* Indicate if vscat log events are happening. If yes, avoid duplicated SF logging to have clear -log_view */
     83:   } vscat;

     85:   /* Fields for generic PetscSF functionality */
     86:   PetscInt     nroots;  /* Number of root vertices on current process (candidates for incoming edges) */
     87:   PetscInt     nleaves; /* Number of leaf vertices on current process (this process specifies a root for each leaf) */
     88:   PetscInt    *mine;    /* Location of leaves in leafdata arrays provided to the communication routines */
     89:   PetscInt    *mine_alloc;
     90:   PetscInt     minleaf, maxleaf;
     91:   PetscSFNode *remote; /* Remote references to roots for each local leaf */
     92:   PetscSFNode *remote_alloc;
     93:   PetscMPIInt  nranks;     /* Number of ranks owning roots connected to my leaves */
     94:   PetscMPIInt  ndranks;    /* Number of ranks in distinguished group holding roots connected to my leaves */
     95:   PetscMPIInt *ranks;      /* List of ranks referenced by "remote" */
     96:   PetscInt    *roffset;    /* Array of length nranks+1, offset in rmine/rremote for each rank */
     97:   PetscInt    *rmine;      /* Concatenated array holding local indices referencing each remote rank */
     98:   PetscInt    *rmine_d[2]; /* A copy of rmine[local/remote] in device memory if needed */

    100:   /* Some results useful in packing by analyzing rmine[] */
    101:   PetscInt       leafbuflen[2];    /* Length (in unit) of leaf buffers, in layout of [PETSCSF_LOCAL/REMOTE] */
    102:   PetscBool      leafcontig[2];    /* True means indices in rmine[self part] or rmine[remote part] are contiguous, and they start from ... */
    103:   PetscInt       leafstart[2];     /* ... leafstart[0] and leafstart[1] respectively */
    104:   PetscSFPackOpt leafpackopt[2];   /* Optimization plans to (un)pack leaves connected to remote roots, based on index patterns in rmine[]. NULL for no optimization */
    105:   PetscSFPackOpt leafpackopt_d[2]; /* Copy of leafpackopt_d[] on device if needed */
    106:   PetscBool      leafdups[2];      /* Indices in rmine[] for self(0)/remote(1) communication have dups respectively? TRUE implies threads working on them in parallel may have data race. */

    108:   PetscMPIInt    nleafreqs;            /* Number of MPI requests for leaves */
    109:   PetscInt      *rremote;              /* Concatenated array holding remote indices referenced for each remote rank */
    110:   PetscBool      degreeknown;          /* The degree is currently known, do not have to recompute */
    111:   PetscInt      *degree;               /* Degree of each of my root vertices */
    112:   PetscInt      *degreetmp;            /* Temporary local array for computing degree */
    113:   PetscBool      rankorder;            /* Sort ranks for gather and scatter operations */
    114:   MPI_Group      ingroup;              /* Group of processes connected to my roots */
    115:   MPI_Group      outgroup;             /* Group of processes connected to my leaves */
    116:   PetscSF        multi;                /* Internal graph used to implement gather and scatter operations */
    117:   PetscSF        rankssf;              /* Internal graph used to implement communications with root ranks */
    118:   PetscBool      graphset;             /* Flag indicating that the graph has been set, required before calling communication routines */
    119:   PetscBool      setupcalled;          /* Type and communication structures have been set up */
    120:   PetscSFPattern pattern;              /* Pattern of the graph */
    121:   PetscBool      persistent;           /* Does this SF use MPI persistent requests for communication */
    122:   PetscBool      collective;           /* Is this SF collective? Currently only SFBASIC/SFWINDOW are not collective */
    123:   PetscLayout    map;                  /* Layout of leaves over all processes when building a patterned graph */
    124:   PetscBool      unknown_input_stream; /* If true, SF does not know which streams root/leafdata is on. Default is false, since we only use petsc default stream */
    125:   PetscBool      use_gpu_aware_mpi;    /* If true, SF assumes it can pass GPU pointers to MPI */
    126:   PetscBool      use_stream_aware_mpi; /* If true, SF assumes the underlying MPI is cuda-stream aware and we won't sync streams for send/recv buffers passed to MPI */
    127:   PetscInt       maxResidentThreadsPerGPU;
    128:   PetscBool      allow_multi_leaves;
    129:   PetscSFBackend backend; /* The device backend (if any) SF will use */
    130:   void          *data;    /* Pointer to implementation */

    132: #if defined(PETSC_HAVE_NVSHMEM)
    133:   PetscBool use_nvshmem;                 /* TRY to use nvshmem on cuda devices with this SF when possible */
    134:   PetscBool use_nvshmem_get;             /* If true, use nvshmem_get based protocol, otherwise, use nvshmem_put based protocol */
    135:   PetscBool checked_nvshmem_eligibility; /* Have we checked eligibility of using NVSHMEM on this sf? */
    136:   PetscBool setup_nvshmem;               /* Have we already set up NVSHMEM related fields below? These fields are built on-demand */
    137:   PetscInt  leafbuflen_rmax;             /* max leafbuflen[REMOTE] over comm */
    138:   PetscInt  nRemoteRootRanks;            /* nranks - ndranks */
    139:   PetscInt  nRemoteRootRanksMax;         /* max nranks-ndranks over comm */

    141:   /* The following two fields look confusing but actually make sense: They are offsets of buffers at the remote side. We're doing one-sided communication! */
    142:   PetscInt *rootsigdisp; /* [nRemoteRootRanks]. For my i-th remote root rank, I will access its rootsigdisp[i]-th root signal */
    143:   PetscInt *rootbufdisp; /* [nRemoteRootRanks]. For my i-th remote root rank, I will access its root buf at offset rootbufdisp[i], in <unit> to be set */

    145:   PetscInt    *rootbufdisp_d;
    146:   PetscInt    *rootsigdisp_d; /* Copy of rootsigdisp[] on device */
    147:   PetscMPIInt *ranks_d;       /* Copy of the remote part of (root) ranks[] on device */
    148:   PetscInt    *roffset_d;     /* Copy of the remote part of roffset[] on device */
    149: #endif
    150: #if defined(PETSC_HAVE_MPIX_STREAM)
    151:   MPIX_Stream mpi_stream;
    152:   MPI_Comm    stream_comm; /* gpu stream aware MPI communicator */
    153: #endif
    154: };

    156: PETSC_EXTERN PetscBool      PetscSFRegisterAllCalled;
    157: PETSC_EXTERN PetscErrorCode PetscSFRegisterAll(void);

    159: PETSC_INTERN PetscErrorCode PetscSFGetDatatypeSize_Internal(MPI_Comm, MPI_Datatype, MPI_Aint *);

    161: PETSC_INTERN PetscErrorCode PetscSFCreateLocalSF_Private(PetscSF, PetscSF *);
    162: PETSC_INTERN PetscErrorCode PetscSFBcastToZero_Private(PetscSF, MPI_Datatype, const void *, void *) PETSC_ATTRIBUTE_MPI_POINTER_WITH_TYPE(3, 2) PETSC_ATTRIBUTE_MPI_POINTER_WITH_TYPE(4, 2);

    164: PETSC_INTERN PetscErrorCode MPIPetsc_Type_unwrap(MPI_Datatype, MPI_Datatype *, PetscBool *);
    165: PETSC_INTERN PetscErrorCode MPIPetsc_Type_compare(MPI_Datatype, MPI_Datatype, PetscBool *);
    166: PETSC_INTERN PetscErrorCode MPIPetsc_Type_compare_contig(MPI_Datatype, MPI_Datatype, PetscInt *);

    168: PETSC_INTERN PetscErrorCode MPIPetsc_Type_get_envelope(MPI_Datatype, MPIU_Count *, MPIU_Count *, MPIU_Count *, MPIU_Count *, PetscMPIInt *);
    169: PETSC_INTERN PetscErrorCode MPIPetsc_Type_get_contents(MPI_Datatype, MPIU_Count, MPIU_Count, MPIU_Count, MPIU_Count, int *, MPI_Aint *, MPIU_Count *, MPI_Datatype *);

    171: #if defined(PETSC_HAVE_MPI_NONBLOCKING_COLLECTIVES)
    172:   #define MPIU_Ibcast(a, b, c, d, e, req)                MPI_Ibcast(a, b, c, d, e, req)
    173:   #define MPIU_Ireduce(a, b, c, d, e, f, g, req)         MPI_Ireduce(a, b, (PetscMPIInt)c, d, e, f, g, req)
    174:   #define MPIU_Iscatter(a, b, c, d, e, f, g, h, req)     MPI_Iscatter(a, b, c, d, e, f, g, h, req)
    175:   #define MPIU_Iscatterv(a, b, c, d, e, f, g, h, i, req) MPI_Iscatterv(a, b, c, d, e, f, g, h, i, req)
    176:   #define MPIU_Igather(a, b, c, d, e, f, g, h, req)      MPI_Igather(a, b, c, d, e, f, g, h, req)
    177:   #define MPIU_Igatherv(a, b, c, d, e, f, g, h, i, req)  MPI_Igatherv(a, b, c, d, e, f, g, h, i, req)
    178:   #define MPIU_Iallgather(a, b, c, d, e, f, g, req)      MPI_Iallgather(a, b, c, d, e, f, g, req)
    179:   #define MPIU_Iallgatherv(a, b, c, d, e, f, g, h, req)  MPI_Iallgatherv(a, b, c, d, e, f, g, h, req)
    180:   #define MPIU_Ialltoall(a, b, c, d, e, f, g, req)       MPI_Ialltoall(a, b, c, d, e, f, g, req)
    181: #else
    182:   /* Ignore req, the MPI_Request argument, and use MPI blocking collectives. One should initialize req
    183:    to MPI_REQUEST_NULL so that one can do MPI_Wait(req,status) no matter the call is blocking or not.
    184:  */
    185:   #define MPIU_Ibcast(a, b, c, d, e, req)                MPI_Bcast(a, b, c, d, e)
    186:   #define MPIU_Ireduce(a, b, c, d, e, f, g, req)         MPI_Reduce(a, b, (PetscMPIInt)c, d, e, f, g)
    187:   #define MPIU_Iscatter(a, b, c, d, e, f, g, h, req)     MPI_Scatter(a, b, c, d, e, f, g, h)
    188:   #define MPIU_Iscatterv(a, b, c, d, e, f, g, h, i, req) MPI_Scatterv(a, b, c, d, e, f, g, h, i)
    189:   #define MPIU_Igather(a, b, c, d, e, f, g, h, req)      MPI_Gather(a, b, c, d, e, f, g, h)
    190:   #define MPIU_Igatherv(a, b, c, d, e, f, g, h, i, req)  MPI_Gatherv(a, b, c, d, e, f, g, h, i)
    191:   #define MPIU_Iallgather(a, b, c, d, e, f, g, req)      MPI_Allgather(a, b, c, d, e, f, g)
    192:   #define MPIU_Iallgatherv(a, b, c, d, e, f, g, h, req)  MPI_Allgatherv(a, b, c, d, e, f, g, h)
    193:   #define MPIU_Ialltoall(a, b, c, d, e, f, g, req)       MPI_Alltoall(a, b, c, d, e, f, g)
    194: #endif

    196: PETSC_EXTERN PetscErrorCode VecScatterGetRemoteCount_Private(VecScatter, PetscBool, PetscInt *, PetscInt *);
    197: PETSC_EXTERN PetscErrorCode VecScatterGetRemote_Private(VecScatter, PetscBool, PetscMPIInt *, const PetscInt **, const PetscInt **, const PetscMPIInt **, PetscInt *);
    198: PETSC_EXTERN PetscErrorCode VecScatterGetRemoteOrdered_Private(VecScatter, PetscBool, PetscMPIInt *, const PetscInt **, const PetscInt **, const PetscMPIInt **, PetscInt *);
    199: PETSC_EXTERN PetscErrorCode VecScatterRestoreRemote_Private(VecScatter, PetscBool, PetscMPIInt *, const PetscInt **, const PetscInt **, const PetscMPIInt **, PetscInt *);
    200: PETSC_EXTERN PetscErrorCode VecScatterRestoreRemoteOrdered_Private(VecScatter, PetscBool, PetscMPIInt *, const PetscInt **, const PetscInt **, const PetscMPIInt **, PetscInt *);

    202: #if defined(PETSC_HAVE_CUDA)
    203: PETSC_EXTERN PetscErrorCode PetscSFMalloc_CUDA(PetscMemType, size_t, void **);
    204: PETSC_EXTERN PetscErrorCode PetscSFFree_CUDA(PetscMemType, void *);
    205: #endif
    206: #if defined(PETSC_HAVE_HIP)
    207: PETSC_EXTERN PetscErrorCode PetscSFMalloc_HIP(PetscMemType, size_t, void **);
    208: PETSC_EXTERN PetscErrorCode PetscSFFree_HIP(PetscMemType, void *);
    209: #endif
    210: #if defined(PETSC_HAVE_KOKKOS)
    211: PETSC_EXTERN PetscErrorCode PetscSFMalloc_Kokkos(PetscMemType, size_t, void **);
    212: PETSC_EXTERN PetscErrorCode PetscSFFree_Kokkos(PetscMemType, void *);
    213: #endif

    215: /* SF only supports CUDA and Kokkos devices. Even VIENNACL is a device, its device pointers are invisible to SF.
    216:    Through VecGetArray(), we copy data of VECVIENNACL from device to host and pass host pointers to SF.
    217:  */
    218: #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_KOKKOS) || defined(PETSC_HAVE_HIP)
    219:   #define PetscSFMalloc(sf, mtype, sz, ptr) ((*(sf)->ops->Malloc)(mtype, sz, ptr))
    220:   /* Free memory and set ptr to NULL when succeeded */
    221:   #define PetscSFFree(sf, mtype, ptr) ((PetscErrorCode)((ptr) && ((*(sf)->ops->Free)(mtype, ptr) || ((ptr) = NULL, PETSC_SUCCESS))))
    222: #else
    223:   /* If pure host code, do with less indirection */
    224:   #define PetscSFMalloc(sf, mtype, sz, ptr) PetscMalloc(sz, ptr)
    225:   #define PetscSFFree(sf, mtype, ptr)       PetscFree(ptr)
    226: #endif