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: }