Actual source code: sfpack.h

  1: #pragma once

  3: #include <../src/vec/is/sf/impls/basic/sfbasic.h>
  4: #if defined(PETSC_HAVE_CUDA)
  5: #include <petscdevice_cuda.h>
  6: typedef cudaStream_t cupmStream_t;
  7: typedef cudaEvent_t  cupmEvent_t;
  8: #endif

 10: #if defined(PETSC_HAVE_HIP)
 11: #include <petscdevice_hip.h>
 12: typedef hipStream_t cupmStream_t;
 13: typedef hipEvent_t  cupmEvent_t;
 14: #endif

 16: /* In terms of function overloading, long long int is a different type than int64_t, which PetscInt might be defined to.
 17:    We prefer long long int over PetscInt (int64_t), since CUDA atomics are built around (unsigned) long long int.
 18:  */
 19: typedef long long int          llint;
 20: typedef unsigned long long int ullint;

 22: /* We separate SF communications for SFBasic and SFNeighbor in two parts: local (self,intra-rank) and remote (inter-rank) */
 23: typedef enum {
 24:   PETSCSF_LOCAL = 0,
 25:   PETSCSF_REMOTE
 26: } PetscSFScope;

 28: /* Optimizations in packing & unpacking for destination ranks.

 30:   Suppose there are m indices stored in idx[], and two addresses u, p. We want to do packing:
 31:      p[i] = u[idx[i]], for i in [0,m)

 33:   Indices are associated with n ranks and each rank's indices are stored consecutively in idx[].
 34:   We go through indices for each rank and see if they are indices of a 3D submatrix of size [dx,dy,dz] in
 35:   a parent matrix of size [X,Y,Z], with the submatrix's first index being <start>.

 37:   E.g., for indices 1,2,3, 6,7,8, 11,12,13, the submatrix size is [3,3,1] with start=1, and the parent matrix's size
 38:   is [5,3,1]. For simplicity, if any destination rank does not have this pattern, we give up the optimization.

 40:   Note before using this per-rank optimization, one should check leafcontig[], rootcontig[], which say
 41:   indices in whole are contiguous, and therefore much more useful than this one when true.
 42:  */
 43: struct _n_PetscSFPackOpt {
 44:   PetscInt *array;        /* [7*n+2] Memory pool for other fields in this struct. Used to easily copy this struct to GPU */
 45:   PetscInt  n;            /* Number of destination ranks */
 46:   PetscInt *offset;       /* [n+1] Offsets of indices for each rank. offset[0]=0, offset[i+1]=offset[i]+dx[i]*dy[i]*dz[i] */
 47:   PetscInt *start;        /* [n] First index */
 48:   PetscInt *dx, *dy, *dz; /* [n] Lengths of the submatrix in X, Y, Z dimension. */
 49:   PetscInt *X, *Y;        /* [n] Lengths of the outer matrix in X, Y. We do not care Z. */
 50: };

 52: /* An abstract class that defines a communication link, which includes how to pack/unpack data and send/recv buffers
 53:  */
 54: struct _n_PetscSFLink {
 55:   PetscErrorCode (*Memcpy)(PetscSFLink, PetscMemType, void *, PetscMemType, const void *, size_t); /* Async device memcopy might use stream in the link */
 56:   PetscErrorCode (*PrePack)(PetscSF, PetscSFLink, PetscSFDirection);
 57:   PetscErrorCode (*PostUnpack)(PetscSF, PetscSFLink, PetscSFDirection);
 58:   PetscErrorCode (*StartCommunication)(PetscSF, PetscSFLink, PetscSFDirection);
 59:   PetscErrorCode (*FinishCommunication)(PetscSF, PetscSFLink, PetscSFDirection);
 60:   PetscErrorCode (*SyncDevice)(PetscSFLink);
 61:   PetscErrorCode (*SyncStream)(PetscSFLink);
 62:   PetscErrorCode (*Destroy)(PetscSF, PetscSFLink);

 64:   PetscErrorCode (*BuildDependenceBegin)(PetscSF, PetscSFLink, PetscSFDirection);
 65:   PetscErrorCode (*BuildDependenceEnd)(PetscSF, PetscSFLink, PetscSFDirection);

 67:   PetscErrorCode (*h_Pack)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, void *);
 68:   PetscErrorCode (*h_UnpackAndInsert)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
 69:   PetscErrorCode (*h_UnpackAndAdd)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
 70:   PetscErrorCode (*h_UnpackAndMin)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
 71:   PetscErrorCode (*h_UnpackAndMax)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
 72:   PetscErrorCode (*h_UnpackAndMinloc)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
 73:   PetscErrorCode (*h_UnpackAndMaxloc)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
 74:   PetscErrorCode (*h_UnpackAndMult)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
 75:   PetscErrorCode (*h_UnpackAndLAND)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
 76:   PetscErrorCode (*h_UnpackAndBAND)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
 77:   PetscErrorCode (*h_UnpackAndLOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
 78:   PetscErrorCode (*h_UnpackAndBOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
 79:   PetscErrorCode (*h_UnpackAndLXOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
 80:   PetscErrorCode (*h_UnpackAndBXOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
 81:   PetscErrorCode (*h_FetchAndAdd)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, void *);

 83:   PetscErrorCode (*h_ScatterAndInsert)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
 84:   PetscErrorCode (*h_ScatterAndAdd)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
 85:   PetscErrorCode (*h_ScatterAndMin)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
 86:   PetscErrorCode (*h_ScatterAndMax)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
 87:   PetscErrorCode (*h_ScatterAndMinloc)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
 88:   PetscErrorCode (*h_ScatterAndMaxloc)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
 89:   PetscErrorCode (*h_ScatterAndMult)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
 90:   PetscErrorCode (*h_ScatterAndLAND)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
 91:   PetscErrorCode (*h_ScatterAndBAND)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
 92:   PetscErrorCode (*h_ScatterAndLOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
 93:   PetscErrorCode (*h_ScatterAndBOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
 94:   PetscErrorCode (*h_ScatterAndLXOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
 95:   PetscErrorCode (*h_ScatterAndBXOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);

 97:   PetscErrorCode (*h_FetchAndAddLocal)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, void *);

 99:   PetscBool deviceinited; /* Are device related fields initialized? */
100: #if defined(PETSC_HAVE_DEVICE)
101:   /* These fields are lazily initialized in a sense that only when device pointers are passed to an SF, the SF
102:      will set them, otherwise it just leaves them alone. Packing routines using regular ops when there are no data race chances.
103:   */
104:   PetscErrorCode (*d_Pack)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, void *);
105:   PetscErrorCode (*d_UnpackAndInsert)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
106:   PetscErrorCode (*d_UnpackAndAdd)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
107:   PetscErrorCode (*d_UnpackAndMin)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
108:   PetscErrorCode (*d_UnpackAndMax)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
109:   PetscErrorCode (*d_UnpackAndMinloc)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
110:   PetscErrorCode (*d_UnpackAndMaxloc)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
111:   PetscErrorCode (*d_UnpackAndMult)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
112:   PetscErrorCode (*d_UnpackAndLAND)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
113:   PetscErrorCode (*d_UnpackAndBAND)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
114:   PetscErrorCode (*d_UnpackAndLOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
115:   PetscErrorCode (*d_UnpackAndBOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
116:   PetscErrorCode (*d_UnpackAndLXOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
117:   PetscErrorCode (*d_UnpackAndBXOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
118:   PetscErrorCode (*d_FetchAndAdd)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, void *);

120:   PetscErrorCode (*d_ScatterAndInsert)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
121:   PetscErrorCode (*d_ScatterAndAdd)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
122:   PetscErrorCode (*d_ScatterAndMin)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
123:   PetscErrorCode (*d_ScatterAndMax)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
124:   PetscErrorCode (*d_ScatterAndMinloc)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
125:   PetscErrorCode (*d_ScatterAndMaxloc)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
126:   PetscErrorCode (*d_ScatterAndMult)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
127:   PetscErrorCode (*d_ScatterAndLAND)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
128:   PetscErrorCode (*d_ScatterAndBAND)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
129:   PetscErrorCode (*d_ScatterAndLOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
130:   PetscErrorCode (*d_ScatterAndBOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
131:   PetscErrorCode (*d_ScatterAndLXOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
132:   PetscErrorCode (*d_ScatterAndBXOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
133:   PetscErrorCode (*d_FetchAndAddLocal)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, void *);

135:   /* Packing routines using atomics when there are data race chances */
136:   PetscErrorCode (*da_UnpackAndInsert)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
137:   PetscErrorCode (*da_UnpackAndAdd)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
138:   PetscErrorCode (*da_UnpackAndMin)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
139:   PetscErrorCode (*da_UnpackAndMax)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
140:   PetscErrorCode (*da_UnpackAndMinloc)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
141:   PetscErrorCode (*da_UnpackAndMaxloc)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
142:   PetscErrorCode (*da_UnpackAndMult)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
143:   PetscErrorCode (*da_UnpackAndLAND)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
144:   PetscErrorCode (*da_UnpackAndBAND)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
145:   PetscErrorCode (*da_UnpackAndLOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
146:   PetscErrorCode (*da_UnpackAndBOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
147:   PetscErrorCode (*da_UnpackAndLXOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
148:   PetscErrorCode (*da_UnpackAndBXOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *);
149:   PetscErrorCode (*da_FetchAndAdd)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, void *);

151:   PetscErrorCode (*da_ScatterAndInsert)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
152:   PetscErrorCode (*da_ScatterAndAdd)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
153:   PetscErrorCode (*da_ScatterAndMin)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
154:   PetscErrorCode (*da_ScatterAndMax)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
155:   PetscErrorCode (*da_ScatterAndMinloc)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
156:   PetscErrorCode (*da_ScatterAndMaxloc)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
157:   PetscErrorCode (*da_ScatterAndMult)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
158:   PetscErrorCode (*da_ScatterAndLAND)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
159:   PetscErrorCode (*da_ScatterAndBAND)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
160:   PetscErrorCode (*da_ScatterAndLOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
161:   PetscErrorCode (*da_ScatterAndBOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
162:   PetscErrorCode (*da_ScatterAndLXOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
163:   PetscErrorCode (*da_ScatterAndBXOR)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *);
164:   PetscErrorCode (*da_FetchAndAddLocal)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, void *);
165:   #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP)
166:   PetscInt     maxResidentThreadsPerGPU; /* It is a copy from SF for convenience */
167:   cupmStream_t stream;                   /* stream on which input/output root/leafdata is computed on (default is PetscDefaultCudaStream) */
168:   #endif
169: #endif
170:   PetscMPIInt  tag;                  /* Each link has a tag so we can perform multiple SF ops at the same time */
171:   MPI_Datatype unit;                 /* The MPI datatype this PetscSFLink is built for */
172:   MPI_Datatype basicunit;            /* unit is made of MPI builtin dataype basicunit */
173:   PetscBool    isbuiltin;            /* Is unit an MPI/PETSc builtin datatype? If it is true, then bs=1 and basicunit is equivalent to unit */
174:   size_t       unitbytes;            /* Number of bytes in a unit */
175:   PetscInt     bs;                   /* Number of basic units in a unit */
176:   const void  *rootdata, *leafdata;  /* rootdata and leafdata the link is working on. They are used as keys for pending links. */
177:   PetscMemType rootmtype, leafmtype; /* root/leafdata's memory type */

179:   /* For local and remote communication */
180:   PetscMemType rootmtype_mpi, leafmtype_mpi;   /* Mtypes of buffers passed to MPI. If use_gpu_aware_mpi, they are same as root/leafmtype. Otherwise they are PETSC_MEMTYPE_HOST */
181:   PetscBool    rootdirect[2], leafdirect[2];   /* Can root/leafdata be directly passed to SF (i.e., without buffering). In layout of [PETSCSF_LOCAL/REMOTE]. See more in PetscSFLinkCreate() */
182:   PetscInt     rootdirect_mpi, leafdirect_mpi; /* Can root/leafdata for remote be directly passed to MPI? 1: yes, 0: no. See more in PetscSFLinkCreate() */
183:   const void  *rootdatadirect[2][2];           /* The root/leafdata used to init root/leaf requests, in layout of [PETSCSF_DIRECTION][PETSC_MEMTYPE]. */
184:   const void  *leafdatadirect[2][2];           /* ... We need them to look up links when root/leafdirect_mpi are true */
185:   char        *rootbuf[2][2];                  /* Buffers for packed roots, in layout of [PETSCSF_LOCAL/REMOTE][PETSC_MEMTYPE]. PETSCSF_LOCAL does not need MPI, .. */
186:                                                /* .. but in case rootmtype is different from leafmtype, we still need to pack local roots and then copy them to memory of leafmtype */
187:   char        *rootbuf_alloc[2][2];            /* Log memory allocated by petsc. We need it since rootbuf[][] may point to rootdata given by user */
188:   char        *leafbuf[2][2];                  /* Buffers for packed leaves, in layout of [PETSCSF_LOCAL/REMOTE][PETSC_MEMTYPE] */
189:   char        *leafbuf_alloc[2][2];
190:   MPI_Request *rootreqs[2][2][2];       /* Root requests in layout of [PETSCSF_DIRECTION][PETSC_MEMTYPE][rootdirect_mpi] */
191:   MPI_Request *leafreqs[2][2][2];       /* Leaf requests in layout of [PETSCSF_DIRECTION][PETSC_MEMTYPE][leafdirect_mpi] */
192:   PetscBool    rootreqsinited[2][2][2]; /* Are root requests initialized? Also in layout of [PETSCSF_DIRECTION][PETSC_MEMTYPE][rootdirect_mpi]*/
193:   PetscBool    leafreqsinited[2][2][2]; /* Are leaf requests initialized? Also in layout of [PETSCSF_DIRECTION][PETSC_MEMTYPE][leafdirect_mpi]*/
194:   MPI_Request *reqs;                    /* An array of length (nrootreqs+nleafreqs)*8. Pointers in rootreqs[][][] and leafreqs[][][] point here */
195:   PetscSFLink  next;

197:   PetscBool use_nvshmem; /* Does this link use nvshem (vs. MPI) for communication? */
198: #if defined(PETSC_HAVE_NVSHMEM)
199:   cupmEvent_t  dataReady;        /* Events to mark readiness of root/leafdata */
200:   cupmEvent_t  endRemoteComm;    /* Events to mark end of local/remote communication */
201:   cupmStream_t remoteCommStream; /* Streams for remote (i.e., inter-rank) communication */

203:   /* The buffers are allocated in device symmetric heap. Their length is the maximal length over all ranks in the comm, and therefore is the same. */
204:   uint64_t *rootSendSig, *rootRecvSig; /* [max{niranks-ndiranks}], signals used when rootbuf works as send/recv buf */
205:   uint64_t *leafSendSig, *leafRecvSig; /* [max{nranks-ndranks}], signals used when leafbuf works as send/recv buf */
206: #endif
207: };

209: PETSC_INTERN PetscErrorCode PetscSFSetErrorOnUnsupportedOverlap(PetscSF, MPI_Datatype, const void *, const void *);

211: /* Create/setup/retrieve/destroy a link */
212: PETSC_INTERN PetscErrorCode PetscSFLinkCreate(PetscSF, MPI_Datatype, PetscMemType, const void *, PetscMemType, const void *, MPI_Op, PetscSFOperation, PetscSFLink *);
213: PETSC_INTERN PetscErrorCode PetscSFLinkSetUp_Host(PetscSF, PetscSFLink, MPI_Datatype);
214: PETSC_INTERN PetscErrorCode PetscSFLinkGetInUse(PetscSF, MPI_Datatype, const void *, const void *, PetscCopyMode, PetscSFLink *);
215: PETSC_INTERN PetscErrorCode PetscSFLinkReclaim(PetscSF, PetscSFLink *);
216: PETSC_INTERN PetscErrorCode PetscSFLinkDestroy(PetscSF, PetscSFLink);

218: /* Get pack/unpack function pointers from a link */
219: static inline PetscErrorCode PetscSFLinkGetPack(PetscSFLink link, PetscMemType mtype, PetscErrorCode (**Pack)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, void *))
220: {
221:   PetscFunctionBegin;
222:   if (PetscMemTypeHost(mtype)) *Pack = link->h_Pack;
223: #if defined(PETSC_HAVE_DEVICE)
224:   else *Pack = link->d_Pack;
225: #endif
226:   PetscFunctionReturn(PETSC_SUCCESS);
227: }

229: PETSC_INTERN PetscErrorCode PetscSFLinkGetUnpackAndOp(PetscSFLink, PetscMemType, MPI_Op, PetscBool, PetscErrorCode (**UnpackAndOp)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, const void *));
230: PETSC_INTERN PetscErrorCode PetscSFLinkGetFetchAndOp(PetscSFLink, PetscMemType, MPI_Op, PetscBool, PetscErrorCode (**FetchAndOp)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, void *));
231: PETSC_INTERN PetscErrorCode PetscSFLinkGetScatterAndOp(PetscSFLink, PetscMemType, MPI_Op, PetscBool, PetscErrorCode (**ScatterAndOp)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, PetscInt, PetscSFPackOpt, const PetscInt *, void *));
232: PETSC_INTERN PetscErrorCode PetscSFLinkGetFetchAndOpLocal(PetscSFLink, PetscMemType, MPI_Op, PetscBool, PetscErrorCode (**FetchAndOpLocal)(PetscSFLink, PetscInt, PetscInt, PetscSFPackOpt, const PetscInt *, void *, PetscInt, PetscSFPackOpt, const PetscInt *, const void *, void *));
233: PETSC_INTERN PetscErrorCode PetscSFLinkGetMPIBuffersAndRequests(PetscSF, PetscSFLink, PetscSFDirection, void **, void **, MPI_Request **, MPI_Request **);

235: /* Do Pack/Unpack/Fetch/Scatter with the link */
236: PETSC_INTERN PetscErrorCode PetscSFLinkPackRootData(PetscSF, PetscSFLink, PetscSFScope, const void *);
237: PETSC_INTERN PetscErrorCode PetscSFLinkPackLeafData(PetscSF, PetscSFLink, PetscSFScope, const void *);
238: PETSC_INTERN PetscErrorCode PetscSFLinkUnpackRootData(PetscSF, PetscSFLink, PetscSFScope, void *, MPI_Op);
239: PETSC_INTERN PetscErrorCode PetscSFLinkUnpackLeafData(PetscSF, PetscSFLink, PetscSFScope, void *, MPI_Op);
240: PETSC_INTERN PetscErrorCode PetscSFLinkFetchAndOpRemote(PetscSF, PetscSFLink, void *, MPI_Op);

242: PETSC_INTERN PetscErrorCode PetscSFLinkScatterLocal(PetscSF, PetscSFLink, PetscSFDirection, void *, void *, MPI_Op);
243: PETSC_INTERN PetscErrorCode PetscSFLinkFetchAndOpLocal(PetscSF, PetscSFLink, void *, const void *, void *, MPI_Op);

245: PETSC_INTERN PetscErrorCode PetscSFSetUpPackFields(PetscSF);
246: PETSC_INTERN PetscErrorCode PetscSFResetPackFields(PetscSF);
247: PETSC_INTERN PetscErrorCode PetscSFLinkCreate_MPI(PetscSF, MPI_Datatype, PetscMemType, const void *, PetscMemType, const void *, MPI_Op, PetscSFOperation, PetscSFLink *);

249: #if defined(PETSC_HAVE_CUDA)
250: PETSC_INTERN PetscErrorCode PetscSFLinkSetUp_CUDA(PetscSF, PetscSFLink, MPI_Datatype);
251: #endif

253: #if defined(PETSC_HAVE_HIP)
254: PETSC_INTERN PetscErrorCode PetscSFLinkSetUp_HIP(PetscSF, PetscSFLink, MPI_Datatype);
255: #endif

257: #if defined(PETSC_HAVE_KOKKOS)
258: PETSC_INTERN PetscErrorCode PetscSFLinkSetUp_Kokkos(PetscSF, PetscSFLink, MPI_Datatype);
259: #endif

261: #if defined(PETSC_HAVE_NVSHMEM)
262: PETSC_INTERN PetscErrorCode PetscSFLinkCreate_NVSHMEM(PetscSF, MPI_Datatype, PetscMemType, const void *, PetscMemType, const void *, MPI_Op, PetscSFOperation, PetscSFLink *);
263: PETSC_INTERN PetscErrorCode PetscSFLinkNvshmemCheck(PetscSF, PetscMemType, const void *, PetscMemType, const void *, PetscBool *);
264: #endif

266: static inline PetscErrorCode PetscSFLinkStartCommunication(PetscSF sf, PetscSFLink link, PetscSFDirection direction)
267: {
268:   PetscFunctionBegin;
269:   if (link->StartCommunication) PetscCall((*link->StartCommunication)(sf, link, direction));
270:   PetscFunctionReturn(PETSC_SUCCESS);
271: }

273: static inline PetscErrorCode PetscSFLinkFinishCommunication(PetscSF sf, PetscSFLink link, PetscSFDirection direction)
274: {
275:   PetscFunctionBegin;
276:   if (link->FinishCommunication) PetscCall((*link->FinishCommunication)(sf, link, direction));
277:   PetscFunctionReturn(PETSC_SUCCESS);
278: }

280: /* A set of helper routines for Pack/Unpack/Scatter on GPUs */
281: #if defined(PETSC_HAVE_CUDA) || defined(PETSC_HAVE_HIP) || defined(PETSC_HAVE_SYCL)
282: /* PetscSFLinkCopyXxxxBufferInCaseNotUseGpuAwareMPI routines are simple: if not use_gpu_aware_mpi, we need
283:    to copy the buffer from GPU to CPU before MPI calls, and from CPU to GPU after MPI calls.
284: */
285: static inline PetscErrorCode PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(PetscSF sf, PetscSFLink link, PetscBool device2host)
286: {
287:   PetscSF_Basic *bas = (PetscSF_Basic *)sf->data;

289:   PetscFunctionBegin;
290:   /* rootdata is on device but we use regular MPI for communication */
291:   if (PetscMemTypeDevice(link->rootmtype) && PetscMemTypeHost(link->rootmtype_mpi) && bas->rootbuflen[PETSCSF_REMOTE]) {
292:     void  *h_buf = link->rootbuf[PETSCSF_REMOTE][PETSC_MEMTYPE_HOST];
293:     void  *d_buf = link->rootbuf[PETSCSF_REMOTE][PETSC_MEMTYPE_DEVICE];
294:     size_t count = bas->rootbuflen[PETSCSF_REMOTE] * link->unitbytes;
295:     if (device2host) {
296:       PetscCall((*link->Memcpy)(link, PETSC_MEMTYPE_HOST, h_buf, PETSC_MEMTYPE_DEVICE, d_buf, count));
297:       PetscCall(PetscLogGpuToCpu(count));
298:     } else {
299:       PetscCall((*link->Memcpy)(link, PETSC_MEMTYPE_DEVICE, d_buf, PETSC_MEMTYPE_HOST, h_buf, count));
300:       PetscCall(PetscLogCpuToGpu(count));
301:     }
302:   }
303:   PetscFunctionReturn(PETSC_SUCCESS);
304: }

306: static inline PetscErrorCode PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(PetscSF sf, PetscSFLink link, PetscBool device2host)
307: {
308:   PetscFunctionBegin;
309:   if (PetscMemTypeDevice(link->leafmtype) && PetscMemTypeHost(link->leafmtype_mpi) && sf->leafbuflen[PETSCSF_REMOTE]) {
310:     void  *h_buf = link->leafbuf[PETSCSF_REMOTE][PETSC_MEMTYPE_HOST];
311:     void  *d_buf = link->leafbuf[PETSCSF_REMOTE][PETSC_MEMTYPE_DEVICE];
312:     size_t count = sf->leafbuflen[PETSCSF_REMOTE] * link->unitbytes;
313:     if (device2host) {
314:       PetscCall((*link->Memcpy)(link, PETSC_MEMTYPE_HOST, h_buf, PETSC_MEMTYPE_DEVICE, d_buf, count));
315:       PetscCall(PetscLogGpuToCpu(count));
316:     } else {
317:       PetscCall((*link->Memcpy)(link, PETSC_MEMTYPE_DEVICE, d_buf, PETSC_MEMTYPE_HOST, h_buf, count));
318:       PetscCall(PetscLogCpuToGpu(count));
319:     }
320:   }
321:   PetscFunctionReturn(PETSC_SUCCESS);
322: }

324: /* Make sure root/leafbuf for the remote is ready for MPI */
325: static inline PetscErrorCode PetscSFLinkSyncStreamBeforeCallMPI(PetscSF sf, PetscSFLink link, PetscSFDirection direction)
326: {
327:   PetscSF_Basic *bas;
328:   PetscInt       buflen;
329:   PetscMemType   mtype;

331:   PetscFunctionBegin;
332:   if (direction == PETSCSF_ROOT2LEAF) {
333:     bas    = (PetscSF_Basic *)sf->data;
334:     mtype  = link->rootmtype;
335:     buflen = bas->rootbuflen[PETSCSF_REMOTE];
336:   } else {
337:     mtype  = link->leafmtype;
338:     buflen = sf->leafbuflen[PETSCSF_REMOTE];
339:   }

341:   if (PetscMemTypeDevice(mtype) && buflen) PetscCall((*link->SyncStream)(link));
342:   PetscFunctionReturn(PETSC_SUCCESS);
343: }
344: #else /* Host only */
345:   #define PetscSFLinkCopyRootBufferInCaseNotUseGpuAwareMPI(a, b, c) PETSC_SUCCESS
346:   #define PetscSFLinkCopyLeafBufferInCaseNotUseGpuAwareMPI(a, b, c) PETSC_SUCCESS
347:   #define PetscSFLinkSyncStreamBeforeCallMPI(a, b, c)               PETSC_SUCCESS
348: #endif

350: /* Get root indices used for pack/unpack

352: Input arguments:
353:   +sf    - StarForest
354:   .link  - The link, which provides the stream for the async memcpy (In SF, we make all GPU operations asynchronous to avoid unexpected pipeline stalls)
355:   .mtype - In what type of memory? (PETSC_MEMTYPE_DEVICE or PETSC_MEMTYPE_HOST)
356:   -scope - Which part of the indices? (PETSCSF_LOCAL or PETSCSF_REMOTE)

358:  Output arguments:
359:   +count   - Count of indices
360:   .start   - The first index (only useful when indices is NULL)
361:   .opt     - Packing optimizations
362:   -indices - Indices of roots for pack/unpack. NULL means indices are contiguous
363:  */
364: static inline PetscErrorCode PetscSFLinkGetRootPackOptAndIndices(PetscSF sf, PetscSFLink link, PetscMemType mtype, PetscSFScope scope, PetscInt *count, PetscInt *start, PetscSFPackOpt *opt, const PetscInt **indices)
365: {
366:   PetscSF_Basic *bas = (PetscSF_Basic *)sf->data;
367:   PetscInt       offset;

369:   PetscFunctionBegin;
370:   *count   = bas->rootbuflen[scope];
371:   *start   = bas->rootstart[scope];
372:   *opt     = NULL;
373:   *indices = NULL;

375:   /* We have these rules:
376:     1) opt == NULL && indices == NULL ==> indices are contiguous.
377:     2) opt != NULL ==> indices are in 3D but not contiguous. On host, indices != NULL since indices are already available and we do not
378:        want to enforce all operations to use opt; but on device, indices = NULL since we do not want to copy indices to device.
379:   */
380:   if (!bas->rootcontig[scope]) {
381:     offset = (scope == PETSCSF_LOCAL) ? 0 : bas->ioffset[bas->ndiranks];
382:     if (PetscMemTypeHost(mtype)) {
383:       *opt     = bas->rootpackopt[scope];
384:       *indices = bas->irootloc + offset;
385:     } else {
386:       size_t size;
387:       if (bas->rootpackopt[scope]) {
388:         if (!bas->rootpackopt_d[scope]) {
389:           PetscCall(PetscMalloc1(1, &bas->rootpackopt_d[scope]));
390:           PetscCall(PetscArraycpy(bas->rootpackopt_d[scope], bas->rootpackopt[scope], 1)); /* Make pointers in bas->rootpackopt_d[] still work on host */
391:           size = (bas->rootpackopt[scope]->n * 7 + 2) * sizeof(PetscInt);                  /* See comments at struct _n_PetscSFPackOpt*/
392:           PetscCall(PetscSFMalloc(sf, PETSC_MEMTYPE_DEVICE, size, (void **)&bas->rootpackopt_d[scope]->array));
393:           PetscCall((*link->Memcpy)(link, PETSC_MEMTYPE_DEVICE, bas->rootpackopt_d[scope]->array, PETSC_MEMTYPE_HOST, bas->rootpackopt[scope]->array, size));
394:         }
395:         *opt = bas->rootpackopt_d[scope];
396:       } else { /* On device, we only provide indices when there is no optimization. We're reluctant to copy indices to device. */
397:         if (!bas->irootloc_d[scope]) {
398:           size = bas->rootbuflen[scope] * sizeof(PetscInt);
399:           PetscCall(PetscSFMalloc(sf, PETSC_MEMTYPE_DEVICE, size, (void **)&bas->irootloc_d[scope]));
400:           PetscCall((*link->Memcpy)(link, PETSC_MEMTYPE_DEVICE, bas->irootloc_d[scope], PETSC_MEMTYPE_HOST, bas->irootloc + offset, size));
401:         }
402:         *indices = bas->irootloc_d[scope];
403:       }
404:     }
405:   }
406:   PetscFunctionReturn(PETSC_SUCCESS);
407: }

409: /* Get leaf indices used for pack/unpack

411:   See also PetscSFLinkGetRootPackOptAndIndices()
412:  */
413: static inline PetscErrorCode PetscSFLinkGetLeafPackOptAndIndices(PetscSF sf, PetscSFLink link, PetscMemType mtype, PetscSFScope scope, PetscInt *count, PetscInt *start, PetscSFPackOpt *opt, const PetscInt **indices)
414: {
415:   PetscInt offset;

417:   PetscFunctionBegin;
418:   *count   = sf->leafbuflen[scope];
419:   *start   = sf->leafstart[scope];
420:   *opt     = NULL;
421:   *indices = NULL;
422:   if (!sf->leafcontig[scope]) {
423:     offset = (scope == PETSCSF_LOCAL) ? 0 : sf->roffset[sf->ndranks];
424:     if (PetscMemTypeHost(mtype)) {
425:       *opt     = sf->leafpackopt[scope];
426:       *indices = sf->rmine + offset;
427:     } else {
428:       size_t size;
429:       if (sf->leafpackopt[scope]) {
430:         if (!sf->leafpackopt_d[scope]) {
431:           PetscCall(PetscMalloc1(1, &sf->leafpackopt_d[scope]));
432:           PetscCall(PetscArraycpy(sf->leafpackopt_d[scope], sf->leafpackopt[scope], 1));
433:           size = (sf->leafpackopt[scope]->n * 7 + 2) * sizeof(PetscInt);                                       /* See comments at struct _n_PetscSFPackOpt*/
434:           PetscCall(PetscSFMalloc(sf, PETSC_MEMTYPE_DEVICE, size, (void **)&sf->leafpackopt_d[scope]->array)); /* Change ->array to a device pointer */
435:           PetscCall((*link->Memcpy)(link, PETSC_MEMTYPE_DEVICE, sf->leafpackopt_d[scope]->array, PETSC_MEMTYPE_HOST, sf->leafpackopt[scope]->array, size));
436:         }
437:         *opt = sf->leafpackopt_d[scope];
438:       } else {
439:         if (!sf->rmine_d[scope]) {
440:           size = sf->leafbuflen[scope] * sizeof(PetscInt);
441:           PetscCall(PetscSFMalloc(sf, PETSC_MEMTYPE_DEVICE, size, (void **)&sf->rmine_d[scope]));
442:           PetscCall((*link->Memcpy)(link, PETSC_MEMTYPE_DEVICE, sf->rmine_d[scope], PETSC_MEMTYPE_HOST, sf->rmine + offset, size));
443:         }
444:         *indices = sf->rmine_d[scope];
445:       }
446:     }
447:   }
448:   PetscFunctionReturn(PETSC_SUCCESS);
449: }