1 #include <petsc/private/sfimpl.h>
2
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
12
MPIPetsc_Type_free(MPI_Datatype * a)13 static PetscErrorCode MPIPetsc_Type_free(MPI_Datatype *a)
14 {
15 MPIU_Count nints, naddrs, ncounts, ntypes;
16 PetscMPIInt combiner;
17
18 PetscFunctionBegin;
19 PetscCallMPI(MPIPetsc_Type_get_envelope(*a, &nints, &naddrs, &ncounts, &ntypes, &combiner));
20
21 if (combiner != MPI_COMBINER_NAMED) PetscCallMPI(MPI_Type_free(a));
22
23 *a = MPI_DATATYPE_NULL;
24 PetscFunctionReturn(PETSC_SUCCESS);
25 }
26
27 // PETSc wrapper for MPI_Type_get_envelope_c using MPIU_Count arguments; works even when MPI large count is not available
MPIPetsc_Type_get_envelope(MPI_Datatype datatype,MPIU_Count * nints,MPIU_Count * naddrs,MPIU_Count * ncounts,MPIU_Count * ntypes,PetscMPIInt * combiner)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 }
45
46 // PETSc wrapper for MPI_Type_get_contents_c using MPIU_Count arguments; works even when MPI large count is not available
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[])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 }
58
59 /*
60 Unwrap an MPI datatype recursively in case it is dupped or MPI_Type_contiguous(1,...)'ed from another type.
61
62 Input Parameter:
63 . a - the datatype to be unwrapped
64
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 */
MPIPetsc_Type_unwrap(MPI_Datatype a,MPI_Datatype * atype,PetscBool * flg)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};
75
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 }
107
MPIPetsc_Type_compare(MPI_Datatype a,MPI_Datatype b,PetscBool * match)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;
115
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;
137
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 }
174
175 /* Check whether a was created via MPI_Type_contiguous from b
176 *
177 */
MPIPetsc_Type_compare_contig(MPI_Datatype a,MPI_Datatype b,PetscInt * n)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;
184
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 }
217
218 if (freeatype) PetscCall(MPIPetsc_Type_free(&atype));
219 if (freebtype) PetscCall(MPIPetsc_Type_free(&btype));
220 PetscFunctionReturn(PETSC_SUCCESS);
221 }
222