Actual source code: sfimpl.h

      1: #if !defined(PETSCSFIMPL_H)
      2: #define PETSCSFIMPL_H

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

      9: #if defined(PETSC_HAVE_CUDA)
     10:   #include <cuda_runtime.h>
     11: #endif

     13: #if defined(PETSC_HAVE_HIP)
     14:   #include <hip/hip_runtime.h>
     15: #endif

     17: PETSC_EXTERN PetscLogEvent PETSCSF_SetGraph;
     18: PETSC_EXTERN PetscLogEvent PETSCSF_SetUp;
     19: PETSC_EXTERN PetscLogEvent PETSCSF_BcastBegin;
     20: PETSC_EXTERN PetscLogEvent PETSCSF_BcastEnd;
     21: PETSC_EXTERN PetscLogEvent PETSCSF_BcastBegin;
     22: PETSC_EXTERN PetscLogEvent PETSCSF_BcastEnd;
     23: PETSC_EXTERN PetscLogEvent PETSCSF_ReduceBegin;
     24: PETSC_EXTERN PetscLogEvent PETSCSF_ReduceEnd;
     25: PETSC_EXTERN PetscLogEvent PETSCSF_FetchAndOpBegin;
     26: PETSC_EXTERN PetscLogEvent PETSCSF_FetchAndOpEnd;
     27: PETSC_EXTERN PetscLogEvent PETSCSF_EmbedSF;
     28: PETSC_EXTERN PetscLogEvent PETSCSF_DistSect;
     29: PETSC_EXTERN PetscLogEvent PETSCSF_SectSF;
     30: PETSC_EXTERN PetscLogEvent PETSCSF_RemoteOff;
     31: PETSC_EXTERN PetscLogEvent PETSCSF_Pack;
     32: PETSC_EXTERN PetscLogEvent PETSCSF_Unpack;

     34: typedef enum {PETSCSF_../../..2LEAF=0, PETSCSF_LEAF2../../..} PetscSFDirection;
     35: typedef enum {PETSCSF_BCAST=0, PETSCSF_REDUCE, PETSCSF_FETCH} PetscSFOperation;
     36: /* When doing device-aware MPI, a backend refers to the SF/device interface */
     37: typedef enum {PETSCSF_BACKEND_INVALID=0,PETSCSF_BACKEND_CUDA,PETSCSF_BACKEND_HIP,PETSCSF_BACKEND_KOKKOS} PetscSFBackend;

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

     60:   PetscErrorCode (*Malloc)(PetscMemType,size_t,void**);
     61:   PetscErrorCode (*Free)(PetscMemType,void*);
     62: };

     64: typedef struct _n_PetscSFPackOpt *PetscSFPackOpt;

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

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

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

    103:   PetscInt        nleafreqs;       /* Number of MPI reqests for leaves */
    104:   PetscInt        *rremote;        /* Concatenated array holding remote indices referenced for each remote rank */
    105:   PetscBool       degreeknown;     /* The degree is currently known, do not have to recompute */
    106:   PetscInt        *degree;         /* Degree of each of my root vertices */
    107:   PetscInt        *degreetmp;      /* Temporary local array for computing degree */
    108:   PetscBool       rankorder;       /* Sort ranks for gather and scatter operations */
    109:   MPI_Group       ingroup;         /* Group of processes connected to my roots */
    110:   MPI_Group       outgroup;        /* Group of processes connected to my leaves */
    111:   PetscSF         multi;           /* Internal graph used to implement gather and scatter operations */
    112:   PetscBool       graphset;        /* Flag indicating that the graph has been set, required before calling communication routines */
    113:   PetscBool       setupcalled;     /* Type and communication structures have been set up */
    114:   PetscSFPattern  pattern;         /* Pattern of the graph */
    115:   PetscBool       persistent;      /* Does this SF use MPI persistent requests for communication */
    116:   PetscLayout     map;             /* Layout of leaves over all processes when building a patterned graph */
    117:   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 */
    118:   PetscBool       use_gpu_aware_mpi;   /* If true, SF assumes it can pass GPU pointers to MPI */
    119:   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 */
    120:   PetscInt        maxResidentThreadsPerGPU;
    121:   PetscSFBackend  backend;         /* The device backend (if any) SF will use */
    122:   void *data;                      /* Pointer to implementation */

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

    133:   /* 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! */
    134:   PetscInt        *rootsigdisp;    /* [nRemoteRootRanks]. For my i-th remote root rank, I will access its rootsigdisp[i]-th root signal */
    135:   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 */

    137:   PetscInt        *rootbufdisp_d;
    138:   PetscInt        *rootsigdisp_d;  /* Copy of rootsigdisp[] on device */
    139:   PetscMPIInt     *ranks_d;        /* Copy of the remote part of (root) ranks[] on device */
    140:   PetscInt        *roffset_d;      /* Copy of the remote part of roffset[] on device */
    141:  #endif
    142: };

    144: PETSC_EXTERN PetscBool PetscSFRegisterAllCalled;
    145: PETSC_EXTERN PetscErrorCode PetscSFRegisterAll(void);

    147: PETSC_INTERN PetscErrorCode PetscSFCreateLocalSF_Private(PetscSF,PetscSF*);
    148: PETSC_INTERN PetscErrorCode PetscSFBcastToZero_Private(PetscSF,MPI_Datatype,const void*,void*);

    150: PETSC_EXTERN PetscErrorCode MPIPetsc_Type_unwrap(MPI_Datatype,MPI_Datatype*,PetscBool*);
    151: PETSC_EXTERN PetscErrorCode MPIPetsc_Type_compare(MPI_Datatype,MPI_Datatype,PetscBool*);
    152: PETSC_EXTERN PetscErrorCode MPIPetsc_Type_compare_contig(MPI_Datatype,MPI_Datatype,PetscInt*);

    154: #if defined(PETSC_HAVE_MPI_NONBLOCKING_COLLECTIVES)
    155: #define MPIU_Iscatter(a,b,c,d,e,f,g,h,req)     MPI_Iscatter(a,b,c,d,e,f,g,h,req)
    156: #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)
    157: #define MPIU_Igather(a,b,c,d,e,f,g,h,req)      MPI_Igather(a,b,c,d,e,f,g,h,req)
    158: #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)
    159: #define MPIU_Iallgather(a,b,c,d,e,f,g,req)     MPI_Iallgather(a,b,c,d,e,f,g,req)
    160: #define MPIU_Iallgatherv(a,b,c,d,e,f,g,h,req)  MPI_Iallgatherv(a,b,c,d,e,f,g,h,req)
    161: #define MPIU_Ialltoall(a,b,c,d,e,f,g,req)      MPI_Ialltoall(a,b,c,d,e,f,g,req)
    162: #else
    163: /* Ignore req, the MPI_Request argument, and use MPI blocking collectives. One should initialize req
    164:    to MPI_REQUEST_NULL so that one can do MPI_Wait(req,status) no matter the call is blocking or not.
    165:  */
    166: #define MPIU_Iscatter(a,b,c,d,e,f,g,h,req)     MPI_Scatter(a,b,c,d,e,f,g,h)
    167: #define MPIU_Iscatterv(a,b,c,d,e,f,g,h,i,req)  MPI_Scatterv(a,b,c,d,e,f,g,h,i)
    168: #define MPIU_Igather(a,b,c,d,e,f,g,h,req)      MPI_Gather(a,b,c,d,e,f,g,h)
    169: #define MPIU_Igatherv(a,b,c,d,e,f,g,h,i,req)   MPI_Gatherv(a,b,c,d,e,f,g,h,i)
    170: #define MPIU_Iallgather(a,b,c,d,e,f,g,req)     MPI_Allgather(a,b,c,d,e,f,g)
    171: #define MPIU_Iallgatherv(a,b,c,d,e,f,g,h,req)  MPI_Allgatherv(a,b,c,d,e,f,g,h)
    172: #define MPIU_Ialltoall(a,b,c,d,e,f,g,req)      MPI_Alltoall(a,b,c,d,e,f,g)
    173: #endif

    175: PETSC_EXTERN PetscErrorCode VecScatterGetRemoteCount_Private(VecScatter,PetscBool,PetscInt*,PetscInt*);
    176: PETSC_EXTERN PetscErrorCode VecScatterGetRemote_Private(VecScatter,PetscBool,PetscInt*,const PetscInt**,const PetscInt**,const PetscMPIInt**,PetscInt*);
    177: PETSC_EXTERN PetscErrorCode VecScatterGetRemoteOrdered_Private(VecScatter,PetscBool,PetscInt*,const PetscInt**,const PetscInt**,const PetscMPIInt**,PetscInt*);
    178: PETSC_EXTERN PetscErrorCode VecScatterRestoreRemote_Private(VecScatter,PetscBool,PetscInt*,const PetscInt**,const PetscInt**,const PetscMPIInt**,PetscInt*);
    179: PETSC_EXTERN PetscErrorCode VecScatterRestoreRemoteOrdered_Private(VecScatter,PetscBool,PetscInt*,const PetscInt**,const PetscInt**,const PetscMPIInt**,PetscInt*);

    181: #if defined(PETSC_HAVE_CUDA)
    182: PETSC_EXTERN PetscErrorCode PetscSFMalloc_CUDA(PetscMemType,size_t,void**);
    183: PETSC_EXTERN PetscErrorCode PetscSFFree_CUDA(PetscMemType,void*);
    184: #endif
    185: #if defined(PETSC_HAVE_HIP)
    186: PETSC_EXTERN PetscErrorCode PetscSFMalloc_HIP(PetscMemType,size_t,void**);
    187: PETSC_EXTERN PetscErrorCode PetscSFFree_HIP(PetscMemType,void*);
    188: #endif

    190: #if defined(PETSC_HAVE_KOKKOS)
    191: PETSC_EXTERN PetscErrorCode PetscSFMalloc_Kokkos(PetscMemType,size_t,void**);
    192: PETSC_EXTERN PetscErrorCode PetscSFFree_Kokkos(PetscMemType,void*);
    193: #endif

    196: /* SF only supports CUDA and Kokkos devices. Even VIENNACL is a device, its device pointers are invisible to SF.
    197:    Through VecGetArray(), we copy data of VECVIENNACL from device to host and pass host pointers to SF.
    198:  */
    199: #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_KOKKOS) || defined(PETSC_HAVE_HIP)
    200:   #define PetscSFMalloc(sf,mtype,sz,ptr)  ((*(sf)->ops->Malloc)(mtype,sz,ptr))
    201:   /* Free memory and set ptr to NULL when succeeded */
    202:   #define PetscSFFree(sf,mtype,ptr)       ((ptr) && ((*(sf)->ops->Free)(mtype,ptr) || ((ptr)=NULL,0)))
    203: #else
    204:   /* If pure host code, do with less indirection */
    205:   #define PetscSFMalloc(sf,mtype,sz,ptr)  PetscMalloc(sz,ptr)
    206:   #define PetscSFFree(sf,mtype,ptr)       PetscFree(ptr)
    207: #endif

    209: #endif