Actual source code: sftype.c

  1: #include <petsc/private/sfimpl.h>

  3: #if !defined(PETSC_HAVE_MPI_COMBINER_DUP) && !defined(MPI_COMBINER_DUP) /* We have no way to interpret output of MPI_Type_get_envelope without this. */
  4:   #define MPI_COMBINER_DUP 0
  5: #endif
  6: #if !defined(PETSC_HAVE_MPI_COMBINER_NAMED) && !defined(MPI_COMBINER_NAMED)
  7:   #define MPI_COMBINER_NAMED -2
  8: #endif
  9: #if !defined(PETSC_HAVE_MPI_COMBINER_CONTIGUOUS) && !defined(MPI_COMBINER_CONTIGUOUS) && MPI_VERSION < 2
 10:   #define MPI_COMBINER_CONTIGUOUS -1
 11: #endif

 13: static PetscErrorCode MPIPetsc_Type_free(MPI_Datatype *a)
 14: {
 15:   MPIU_Count  nints, naddrs, ncounts, ntypes;
 16:   PetscMPIInt combiner;

 18:   PetscFunctionBegin;
 19:   PetscCallMPI(MPIPetsc_Type_get_envelope(*a, &nints, &naddrs, &ncounts, &ntypes, &combiner));

 21:   if (combiner != MPI_COMBINER_NAMED) PetscCallMPI(MPI_Type_free(a));

 23:   *a = MPI_DATATYPE_NULL;
 24:   PetscFunctionReturn(PETSC_SUCCESS);
 25: }

 27: // petsc wrapper for MPI_Type_get_envelope_c using MPIU_Count arguments; works even when MPI large count is not available
 28: PetscErrorCode MPIPetsc_Type_get_envelope(MPI_Datatype datatype, MPIU_Count *nints, MPIU_Count *naddrs, MPIU_Count *ncounts, MPIU_Count *ntypes, PetscMPIInt *combiner)
 29: {
 30:   PetscFunctionBegin;
 31: #if defined(PETSC_HAVE_MPI_LARGE_COUNT) && !defined(PETSC_HAVE_MPIUNI) // MPIUNI does not really support large counts in datatype creation
 32:   PetscCallMPI(MPI_Type_get_envelope_c(datatype, nints, naddrs, ncounts, ntypes, combiner));
 33: #else
 34:   PetscMPIInt mints, maddrs, mtypes;
 35:   // As of 2024/09/12, MPI Forum has yet to decide whether it is legal to call MPI_Type_get_envelope() on types created by, e.g.,
 36:   // MPI_Type_contiguous_c(4, MPI_DOUBLE, &newtype). We just let the MPI being used play out (i.e., return error or not)
 37:   PetscCallMPI(MPI_Type_get_envelope(datatype, &mints, &maddrs, &mtypes, combiner));
 38:   *nints   = mints;
 39:   *naddrs  = maddrs;
 40:   *ncounts = 0;
 41:   *ntypes  = mtypes;
 42: #endif
 43:   PetscFunctionReturn(PETSC_SUCCESS);
 44: }

 46: // petsc wrapper for MPI_Type_get_contents_c using MPIU_Count arguments; works even when MPI large count is not available
 47: PetscErrorCode MPIPetsc_Type_get_contents(MPI_Datatype datatype, MPIU_Count nints, MPIU_Count naddrs, MPIU_Count ncounts, MPIU_Count ntypes, int intarray[], MPI_Aint addrarray[], MPIU_Count countarray[], MPI_Datatype typearray[])
 48: {
 49:   PetscFunctionBegin;
 50: #if defined(PETSC_HAVE_MPI_LARGE_COUNT) && !defined(PETSC_HAVE_MPIUNI) // MPI-4.0, so MPIU_Count is MPI_Count
 51:   PetscCallMPI(MPI_Type_get_contents_c(datatype, nints, naddrs, ncounts, ntypes, intarray, addrarray, countarray, typearray));
 52: #else
 53:   PetscCheck(nints <= PETSC_MPI_INT_MAX && naddrs <= PETSC_MPI_INT_MAX && ntypes <= PETSC_MPI_INT_MAX && ncounts == 0, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "The input derived MPI datatype is created with large counts, but petsc is configured with an MPI without the large count support");
 54:   PetscCallMPI(MPI_Type_get_contents(datatype, (PetscMPIInt)nints, (PetscMPIInt)naddrs, (PetscMPIInt)ntypes, intarray, addrarray, typearray));
 55: #endif
 56:   PetscFunctionReturn(PETSC_SUCCESS);
 57: }

 59: /*
 60:   Unwrap an MPI datatype recursively in case it is dupped or MPI_Type_contiguous(1,...)'ed from another type.

 62:    Input Parameter:
 63: .  a  - the datatype to be unwrapped

 65:    Output Parameters:
 66: + atype - the unwrapped datatype, which is either equal(=) to a or equivalent to a.
 67: - flg   - true if atype != a, which implies caller should MPIPetsc_Type_free(atype) after use. Note atype might be MPI builtin.
 68: */
 69: PetscErrorCode MPIPetsc_Type_unwrap(MPI_Datatype a, MPI_Datatype *atype, PetscBool *flg)
 70: {
 71:   MPIU_Count   nints = 0, naddrs = 0, ncounts = 0, ntypes = 0, counts[1] = {0};
 72:   PetscMPIInt  combiner, ints[1] = {0};
 73:   MPI_Aint     addrs[1] = {0};
 74:   MPI_Datatype types[1] = {MPI_INT};

 76:   PetscFunctionBegin;
 77:   *flg   = PETSC_FALSE;
 78:   *atype = a;
 79:   if (a == MPIU_INT || a == MPIU_REAL || a == MPIU_SCALAR) PetscFunctionReturn(PETSC_SUCCESS);
 80:   PetscCall(MPIPetsc_Type_get_envelope(a, &nints, &naddrs, &ncounts, &ntypes, &combiner));
 81:   if (combiner == MPI_COMBINER_DUP) {
 82:     PetscCheck(nints == 0 && naddrs == 0 && ncounts == 0 && ntypes == 1, PETSC_COMM_SELF, PETSC_ERR_LIB, "Unexpected returns from MPI_Type_get_envelope()");
 83:     PetscCallMPI(MPIPetsc_Type_get_contents(a, nints, naddrs, ncounts, ntypes, ints, addrs, counts, types));
 84:     /* Recursively unwrap dupped types. */
 85:     PetscCall(MPIPetsc_Type_unwrap(types[0], atype, flg));
 86:     if (*flg) {
 87:       /* If the recursive call returns a new type, then that means that atype[0] != types[0] and we're on the hook to
 88:        * free types[0].  Note that this case occurs if combiner(types[0]) is MPI_COMBINER_DUP, so we're safe to
 89:        * directly call MPI_Type_free rather than MPIPetsc_Type_free here. */
 90:       PetscCallMPI(MPI_Type_free(&types[0]));
 91:     }
 92:     /* In any case, it's up to the caller to free the returned type in this case. */
 93:     *flg = PETSC_TRUE;
 94:   } else if (combiner == MPI_COMBINER_CONTIGUOUS) {
 95:     PetscCheck((nints + ncounts == 1) && naddrs == 0 && ntypes == 1, PETSC_COMM_SELF, PETSC_ERR_LIB, "Unexpected returns from MPI_Type_get_envelope()");
 96:     PetscCallMPI(MPIPetsc_Type_get_contents(a, nints, naddrs, ncounts, ntypes, ints, addrs, counts, types));
 97:     if ((nints == 1 && ints[0] == 1) || (ncounts == 1 && counts[0] == 1)) { /* If a is created by MPI_Type_contiguous/_c(1,..) */
 98:       PetscCall(MPIPetsc_Type_unwrap(types[0], atype, flg));
 99:       if (*flg) PetscCall(MPIPetsc_Type_free(&types[0]));
100:       *flg = PETSC_TRUE;
101:     } else {
102:       PetscCall(MPIPetsc_Type_free(&types[0]));
103:     }
104:   }
105:   PetscFunctionReturn(PETSC_SUCCESS);
106: }

108: PetscErrorCode MPIPetsc_Type_compare(MPI_Datatype a, MPI_Datatype b, PetscBool *match)
109: {
110:   MPI_Datatype atype, btype;
111:   MPIU_Count   aintcount, aaddrcount, acountcount, atypecount;
112:   MPIU_Count   bintcount, baddrcount, bcountcount, btypecount;
113:   PetscMPIInt  acombiner, bcombiner;
114:   PetscBool    freeatype, freebtype;

116:   PetscFunctionBegin;
117:   if (a == b) { /* this is common when using MPI builtin datatypes */
118:     *match = PETSC_TRUE;
119:     PetscFunctionReturn(PETSC_SUCCESS);
120:   }
121:   PetscCall(MPIPetsc_Type_unwrap(a, &atype, &freeatype));
122:   PetscCall(MPIPetsc_Type_unwrap(b, &btype, &freebtype));
123:   *match = PETSC_FALSE;
124:   if (atype == btype) {
125:     *match = PETSC_TRUE;
126:     goto free_types;
127:   }
128:   PetscCall(MPIPetsc_Type_get_envelope(atype, &aintcount, &aaddrcount, &acountcount, &atypecount, &acombiner));
129:   PetscCall(MPIPetsc_Type_get_envelope(btype, &bintcount, &baddrcount, &bcountcount, &btypecount, &bcombiner));
130:   if (acombiner == bcombiner && aintcount == bintcount && aaddrcount == baddrcount && acountcount == bcountcount && atypecount == btypecount && (aintcount > 0 || aaddrcount > 0 || acountcount > 0 || atypecount > 0)) {
131:     PetscMPIInt  *aints, *bints;
132:     MPI_Aint     *aaddrs, *baddrs;
133:     MPIU_Count   *acounts, *bcounts;
134:     MPI_Datatype *atypes, *btypes;
135:     PetscInt      i;
136:     PetscBool     same;

138:     PetscCall(PetscMalloc4(aintcount, &aints, aaddrcount, &aaddrs, acountcount, &acounts, atypecount, &atypes));
139:     PetscCall(PetscMalloc4(bintcount, &bints, baddrcount, &baddrs, bcountcount, &bcounts, btypecount, &btypes));
140:     PetscCall(MPIPetsc_Type_get_contents(atype, aintcount, aaddrcount, acountcount, atypecount, aints, aaddrs, acounts, atypes));
141:     PetscCall(MPIPetsc_Type_get_contents(btype, bintcount, baddrcount, bcountcount, btypecount, bints, baddrs, bcounts, btypes));
142:     PetscCall(PetscArraycmp(aints, bints, aintcount, &same));
143:     if (same) {
144:       PetscCall(PetscArraycmp(aaddrs, baddrs, aaddrcount, &same));
145:       if (same) {
146:         PetscCall(PetscArraycmp(acounts, bcounts, acountcount, &same));
147:         if (same) {
148:           /* Check for identity first */
149:           PetscCall(PetscArraycmp(atypes, btypes, atypecount, &same));
150:           if (!same) {
151:             /* If the atype or btype were not predefined data types, then the types returned from MPI_Type_get_contents
152:            * will merely be equivalent to the types used in the construction, so we must recursively compare. */
153:             for (i = 0; i < atypecount; i++) {
154:               PetscCall(MPIPetsc_Type_compare(atypes[i], btypes[i], &same));
155:               if (!same) break;
156:             }
157:           }
158:         }
159:       }
160:     }
161:     for (i = 0; i < atypecount; i++) {
162:       PetscCall(MPIPetsc_Type_free(&atypes[i]));
163:       PetscCall(MPIPetsc_Type_free(&btypes[i]));
164:     }
165:     PetscCall(PetscFree4(aints, aaddrs, acounts, atypes));
166:     PetscCall(PetscFree4(bints, baddrs, bcounts, btypes));
167:     if (same) *match = PETSC_TRUE;
168:   }
169: free_types:
170:   if (freeatype) PetscCall(MPIPetsc_Type_free(&atype));
171:   if (freebtype) PetscCall(MPIPetsc_Type_free(&btype));
172:   PetscFunctionReturn(PETSC_SUCCESS);
173: }

175: /* Check whether a was created via MPI_Type_contiguous from b
176:  *
177:  */
178: PetscErrorCode MPIPetsc_Type_compare_contig(MPI_Datatype a, MPI_Datatype b, PetscInt *n)
179: {
180:   MPI_Datatype atype, btype;
181:   MPIU_Count   aintcount, aaddrcount, acountcount, atypecount;
182:   PetscMPIInt  acombiner;
183:   PetscBool    freeatype, freebtype;

185:   PetscFunctionBegin;
186:   if (a == b) {
187:     *n = 1;
188:     PetscFunctionReturn(PETSC_SUCCESS);
189:   }
190:   *n = 0;
191:   PetscCall(MPIPetsc_Type_unwrap(a, &atype, &freeatype));
192:   PetscCall(MPIPetsc_Type_unwrap(b, &btype, &freebtype));
193:   PetscCall(MPIPetsc_Type_get_envelope(atype, &aintcount, &aaddrcount, &acountcount, &atypecount, &acombiner));
194:   if (acombiner == MPI_COMBINER_CONTIGUOUS && (aintcount >= 1 || acountcount >= 1)) {
195:     PetscMPIInt  *aints;
196:     MPI_Aint     *aaddrs;
197:     MPIU_Count   *acounts;
198:     MPI_Datatype *atypes;
199:     PetscBool     same;
200:     PetscCall(PetscMalloc4(aintcount, &aints, aaddrcount, &aaddrs, acountcount, &acounts, atypecount, &atypes));
201:     PetscCall(MPIPetsc_Type_get_contents(atype, aintcount, aaddrcount, acountcount, atypecount, aints, aaddrs, acounts, atypes));
202:     /* Check for identity first. */
203:     if (atypes[0] == btype) {
204:       if (aintcount) *n = aints[0];
205:       else PetscCall(PetscIntCast(acounts[0], n)); // Yet to support real big count values
206:     } else {
207:       /* atypes[0] merely has to be equivalent to the type used to create atype. */
208:       PetscCall(MPIPetsc_Type_compare(atypes[0], btype, &same));
209:       if (same) {
210:         if (aintcount) *n = aints[0];
211:         else PetscCall(PetscIntCast(acounts[0], n)); // Yet to support real big count values
212:       }
213:     }
214:     for (MPIU_Count i = 0; i < atypecount; i++) PetscCall(MPIPetsc_Type_free(&atypes[i]));
215:     PetscCall(PetscFree4(aints, aaddrs, acounts, atypes));
216:   }

218:   if (freeatype) PetscCall(MPIPetsc_Type_free(&atype));
219:   if (freebtype) PetscCall(MPIPetsc_Type_free(&btype));
220:   PetscFunctionReturn(PETSC_SUCCESS);
221: }