import itertools import os import pathlib import sys from mpi4py import MPI from mpi4py.util.dtlib import ( from_numpy_dtype as fromnumpy, to_numpy_dtype as tonumpy, ) try: import mpitestutil as testutil import mpiunittest as unittest except ImportError: sys.path.append(os.fspath(pathlib.Path(__file__).resolve().parent)) import mpitestutil as testutil import mpiunittest as unittest try: import numpy np_dtype = numpy.dtype np_version = tuple(map(int, numpy.__version__.split(".", 2)[:2])) except ImportError: numpy = None np_dtype = None np_version = None typecodes = list("?cbhilqpBHILQPefdgFDG") typecodes += [f"b{n:d}" for n in (1,)] typecodes += [f"i{n:d}" for n in (1, 2, 4, 8)] typecodes += [f"u{n:d}" for n in (1, 2, 4, 8)] typecodes += [f"f{n:d}" for n in (2, 4, 8)] if os.environ.get("COVERAGE_RUN") == "true": typecodes = [*list("cif"), "b1", "i8", "f8"] if np_version and np_version < (1, 17): for tc in "LFDG": if tc in typecodes: typecodes.remove(tc) if ( MPI.FLOAT16_T == MPI.DATATYPE_NULL or unittest.is_mpi("mpich(<5.0.0)") or unittest.is_mpi("openmpi(<5.0.0)") or unittest.is_mpi("impi") ): for tc in ("e", "f2"): if tc in typecodes: typecodes.remove(tc) if unittest.is_mpi("mpich(<4.0.0)"): for tc in "FDG": if tc in typecodes: typecodes.remove(tc) if unittest.is_mpi("impi(>=2021.12.0)") and os.name == "nt": for tc in (*"lLg", "i4", "u4"): if tc in typecodes: typecodes.remove(tc) datatypes_c = [MPI.Datatype.fromcode(t) for t in typecodes] + [ MPI.BYTE, MPI.AINT, MPI.OFFSET, MPI.COUNT, ] mpipairtypes = [ MPI.SHORT_INT, MPI.INT_INT, MPI.LONG_INT, MPI.FLOAT_INT, MPI.DOUBLE_INT, MPI.LONG_DOUBLE_INT, ] mpif77types = [ MPI.CHARACTER, MPI.LOGICAL, MPI.INTEGER, MPI.REAL, MPI.DOUBLE_PRECISION, MPI.COMPLEX, MPI.DOUBLE_COMPLEX, ] mpif90types = [ MPI.LOGICAL1, MPI.LOGICAL2, MPI.LOGICAL4, MPI.LOGICAL8, MPI.LOGICAL16, MPI.INTEGER1, MPI.INTEGER2, MPI.INTEGER4, MPI.INTEGER8, MPI.INTEGER16, MPI.REAL2, MPI.REAL4, MPI.REAL8, MPI.REAL16, MPI.COMPLEX4, MPI.COMPLEX8, MPI.COMPLEX16, MPI.COMPLEX32, ] for typelist in [ datatypes_c, mpif77types, mpif90types, ]: typelist[:] = [t for t in typelist if testutil.has_datatype(t)] del typelist datatypes = [] datatypes += datatypes_c datatypes += mpif77types datatypes += mpif90types def try_dtype(*args): for spec in args: if isinstance(spec, MPI.Datatype): spec = spec.typestr if np_dtype is not None: try: np_dtype(spec) except TypeError: return False return True class TestUtilDTLib(unittest.TestCase): # def check(self, arg, *args): if numpy is None: if isinstance(arg, MPI.Datatype): mt1 = arg.Dup() dt1 = tonumpy(mt1) mt1.Free() return if not try_dtype(arg): return if isinstance(arg, MPI.Datatype): mt1 = arg.Dup() dt1 = tonumpy(mt1) else: dt1 = np_dtype(arg, *args) mt1 = fromnumpy(dt1) dt2 = tonumpy(mt1) mt2 = fromnumpy(dt2) dt3 = tonumpy(mt2) mt3 = fromnumpy(dt3) try: self.assertEqual(dt1, dt2) self.assertEqual(dt2, dt3) if isinstance(arg, MPI.Datatype): if arg.combiner not in ( MPI.COMBINER_INDEXED, MPI.COMBINER_HINDEXED, MPI.COMBINER_INDEXED_BLOCK, MPI.COMBINER_HINDEXED_BLOCK, ): self.assertEqual(dt1.itemsize, mt1.extent) self.assertEqual(dt2.itemsize, mt2.extent) self.assertEqual(dt3.itemsize, mt3.extent) finally: mt1.Free() mt2.Free() mt3.Free() def testBasic(self): for spec in typecodes: with self.subTest(spec=spec): self.check(spec) for mpit in datatypes: with self.subTest(name=mpit.name): self.check(mpit) def testSubarray1(self): shapes = [ (1,), (1, 1), (1, 1, 1), (3,), (3, 4), (2, 3, 4), ] for dt, shape in itertools.product(typecodes, shapes): spec = f"{shape}{dt}" with self.subTest(spec=spec): self.check(spec) def testSubarray2(self): shapes = [ (1,), (1, 1), (1, 1, 1), (3,), (3, 4), (2, 3, 4), ] orders = [MPI.ORDER_C, MPI.ORDER_FORTRAN] for mt, shape, order in itertools.product(datatypes, shapes, orders): if not try_dtype(mt): continue with self.subTest(name=mt.name, shape=shape, order=order): starts = (0,) * len(shape) mt1 = mt.Create_subarray(shape, shape, starts, order) self.check(mt1) mt1.Free() @unittest.skipMPI("msmpi") def testStruct1(self): shapes = [ (), (1,), (3,), (3, 5), ] iter1 = itertools.product(shapes, typecodes) iter2 = itertools.product(shapes, typecodes) iterN = itertools.product(iter1, iter2) iterA = iter([False, True]) for nt, align in itertools.product(iterN, iterA): s1, t1, s2, t2 = sum(nt, ()) spec = f"{s1}{t1},{s2}{t2}" with self.subTest(spec=spec, align=align): self.check(spec, align) @unittest.skipMPI("msmpi") def testStruct2(self): iter1 = iter(typecodes) iter2 = iter(typecodes) iter3 = iter(typecodes) iterN = itertools.product(iter1, iter2, iter3) iterA = iter([False, True]) for tp, align in itertools.product(iterN, iterA): t1, t2, t3 = tp spec = f"{t1},{t2},{t3}" with self.subTest(spec=spec, align=align): self.check(spec, align) @unittest.skipMPI("msmpi") def testStruct3(self): blens = [1, 2, 3] disps = [1, 27, 71] types = [MPI.INT, MPI.DOUBLE, MPI.INT] mt1 = MPI.Datatype.Create_struct(blens, disps, types) mt2 = MPI.Datatype.Create_struct([1], [0], [mt1]) self.check(mt1) self.check(mt2) mt1.Free() mt2.Free() def makeStruct(self, dt, mt): dt = numpy.dtype(dt).str stp = numpy.dtype(f"B,{dt},B", align=True) def off(i): return stp.fields[stp.names[i]][1] blens = [1, 1, 1] disps = [0, off(1), off(2)] types = [MPI.BYTE, mt, MPI.BYTE] mtp = MPI.Datatype.Create_struct(blens, disps, types) return stp, mtp @unittest.skipMPI("msmpi") @unittest.skipIf(numpy is None, "numpy") def testStruct4(self): for t in typecodes: with self.subTest(typecode=t): dt0 = np_dtype(t) mt0 = fromnumpy(dt0) stp, mt1 = self.makeStruct(t, mt0) ex1 = stp.itemsize for n, mt in ( (1, mt1), (1, mt1.Dup()), (1, mt1.Create_resized(0, 1 * ex1)), (3, mt1.Create_resized(0, 3 * ex1)), (3, mt1.Create_contiguous(3)), (5, mt1.Create_subarray([5], [5], [0])), (7, MPI.Datatype.Create_struct([7], [0], [mt1])), ): dt = tonumpy(mt) self.assertEqual(mt.extent, n * ex1) self.assertEqual(dt.itemsize, n * ex1) self.assertTrue(dt.isalignedstruct) self.check(mt) self.check(dt) if mt != mt1: mt.Free() mt0.Free() mt1.Free() @unittest.skipMPI("msmpi") @unittest.skipIf(numpy is None, "numpy") def testStruct5(self): for t1, t2 in itertools.product(*[typecodes] * 2): with self.subTest(t1=t1, t2=t2): dtlist = [] dt = np_dtype(f"c,{t1},{t2},c", align=True) dtlist.append(dt) for _ in range(3): dt = np_dtype([("", dt)] * 2, align=True) dtlist.append(dt) for dt in dtlist: mt = fromnumpy(dt) dt2 = tonumpy(mt) mt.Free() self.assertEqual(dt, dt2) def testVector(self): for mt in datatypes: if not try_dtype(mt): continue with self.subTest(name=mt.name): mt1 = mt.Create_vector(3, 4, 6) mt2 = mt.Create_hvector(3, 4, 6 * mt.extent) self.check(mt1) self.check(mt2) dt1 = tonumpy(mt1) dt2 = tonumpy(mt2) self.check(dt1) self.check(dt2) self.assertEqual(dt1, dt2) mt3 = mt1.Create_vector(2, 3, 4) mt4 = mt2.Create_hvector(2, 3, 4 * mt2.extent) self.check(mt3) self.check(mt4) dt3 = tonumpy(mt3) dt4 = tonumpy(mt4) self.check(dt3) self.check(dt4) self.assertEqual(dt3, dt4) mt3.Free() mt4.Free() mt1.Free() mt2.Free() def testHVector(self): for mt in datatypes: if not try_dtype(mt): continue with self.subTest(name=mt.name): mt1 = mt.Create_hvector(3, 4, 6 * mt.extent + 1) mt2 = mt1.Dup() self.check(mt1) self.check(mt2) dt1 = tonumpy(mt1) dt2 = tonumpy(mt2) self.check(dt1) self.check(dt2) self.assertEqual(dt1, dt2) mt3 = mt1.Create_hvector(2, 3, 4 * mt1.extent + 1) mt4 = mt2.Create_hvector(2, 3, 4 * mt2.extent + 1) self.check(mt3) self.check(mt4) dt3 = tonumpy(mt3) dt4 = tonumpy(mt4) self.check(dt3) self.check(dt4) self.assertEqual(dt3, dt4) mt3.Free() mt4.Free() mt1.Free() mt2.Free() def testIndexed(self): disps = [1, 6, 12] for mt in datatypes: if not try_dtype(mt): continue with self.subTest(name=mt.name): mt1 = mt.Create_indexed([4] * 3, disps) mt2 = mt.Create_indexed_block(4, disps) self.check(mt1) self.check(mt2) dt1 = tonumpy(mt1) dt2 = tonumpy(mt2) self.check(dt1) self.check(dt2) self.assertEqual(dt1, dt2) mt3 = mt1.Create_indexed([1], [0]) mt4 = mt2.Create_indexed_block(1, [0]) self.check(mt3) self.check(mt4) dt3 = tonumpy(mt3) dt4 = tonumpy(mt4) self.check(dt3) self.check(dt4) self.assertEqual(dt3, dt4) mt3.Free() mt4.Free() mt1.Free() mt2.Free() def testHIndexed(self): disps = [0, 6, 12] for mt in datatypes: if not try_dtype(mt): continue with self.subTest(name=mt.name): mt1 = mt.Create_hindexed( [4] * 3, [d * mt.extent + 1 for d in disps] ) mt2 = mt.Create_hindexed_block( 4, [d * mt.extent + 1 for d in disps] ) self.check(mt1) self.check(mt2) dt1 = tonumpy(mt1) dt2 = tonumpy(mt2) self.check(dt1) self.check(dt2) self.assertEqual(dt1, dt2) mt3 = mt1.Create_hindexed([1], [0]) mt4 = mt2.Create_hindexed_block(1, [0]) self.check(mt3) self.check(mt4) mt3.Free() mt4.Free() mt1.Free() mt2.Free() @unittest.skipMPI("msmpi") def testF77(self): for mt in mpif77types: dt = tonumpy(mt) if np_dtype is not None: self.assertEqual(dt.itemsize, mt.extent) @unittest.skipMPI("msmpi") def testF90(self): for mt in mpif90types: if np_dtype is not None: typestr = mt.typestr try: np_dtype(typestr) except TypeError: continue dt = tonumpy(mt) if np_dtype is not None: self.assertEqual(dt.itemsize, mt.extent) @unittest.skipMPI("msmpi") def testF90Integer(self): try: mt = MPI.Datatype.Create_f90_integer(1) except NotImplementedError: mt = MPI.DATATYPE_NULL if not testutil.has_datatype(mt): self.skipTest("mpi-type-create-f90-integer") for r in range(1, 19): with self.subTest(r=r): mt = MPI.Datatype.Create_f90_integer(r) dt = tonumpy(mt) if np_dtype is not None: self.assertEqual(dt.kind, "i") self.assertEqual(dt.itemsize, mt.extent) size = mt.Get_size() tstr = f"i{size}" stp, mtp = self.makeStruct(tstr, mt) self.assertEqual(stp.itemsize, mtp.extent) self.check(mtp) mtp.Free() @unittest.skipMPI("msmpi") def testF90Real(self): try: mt = MPI.Datatype.Create_f90_real(7, MPI.UNDEFINED) except NotImplementedError: mt = MPI.DATATYPE_NULL if not testutil.has_datatype(mt): self.skipTest("mpi-type-create-f90-real") for p in (6, 7, 14, 15): with self.subTest(p=p): mt = MPI.Datatype.Create_f90_real(p, MPI.UNDEFINED) dt = tonumpy(mt) if np_dtype is not None: self.assertEqual(dt.kind, "f") self.assertEqual(dt.itemsize, mt.extent) size = mt.Get_size() tstr = f"i{size}" stp, mtp = self.makeStruct(tstr, mt) self.assertEqual(stp.itemsize, mtp.extent) self.check(mtp) mtp.Free() @unittest.skipMPI("msmpi") def testF90Complex(self): try: mt = MPI.Datatype.Create_f90_complex(7, MPI.UNDEFINED) except NotImplementedError: mt = MPI.DATATYPE_NULL if not testutil.has_datatype(mt): self.skipTest("mpi-type-create-f90-complex") for p in (6, 7, 14, 15): with self.subTest(p=p): mt = MPI.Datatype.Create_f90_complex(p, MPI.UNDEFINED) dt = tonumpy(mt) if np_dtype is not None: self.assertEqual(dt.kind, "c") self.assertEqual(dt.itemsize, mt.extent) def testPair(self): for mt in mpipairtypes: with self.subTest(datatype=mt.name): dt = tonumpy(mt) if np_dtype is not None: self.assertTrue(dt.isalignedstruct) self.assertEqual(dt.itemsize, mt.extent) integral = "bhilqpBHILQP" floating = "fdg" vtypes = integral + floating itypes = integral for vcode, icode in itertools.product(vtypes, itypes): value = MPI.Datatype.fromcode(vcode) index = MPI.Datatype.fromcode(icode) pair = MPI.Datatype.Get_value_index(value, index) if pair == MPI.DATATYPE_NULL: continue vt, it, pt = map(tonumpy, (value, index, pair)) dt = (f"{vt},{it}", {"align": True}) if np_dtype is not None: dt = np_dtype(dt[0], **dt[1]) self.assertEqual(pt, dt) def testPairStruct(self): cases = [mpipairtypes] * 3 + [[False, True]] for mt1, mt2, mt3, dup in itertools.product(*cases): with self.subTest(mt1=mt1.name, mt2=mt2.name, mt3=mt3.name): if dup: mt1 = mt1.Dup() mt2 = mt2.Dup() mt3 = mt3.Dup() align = max(mt.extent for mt in (mt1, mt2, mt3)) structtype = MPI.Datatype.Create_struct( [1, 1, 1], [0, align, align * 2], [mt1, mt2, mt3], ) if dup: mt1.Free() mt2.Free() mt3.Free() dt = tonumpy(structtype) structtype.Free() if np_dtype is not None: self.assertTrue(dt.isalignedstruct) def testAlignmentComplex(self): complexcodes = list("FDG") complexcodes += [f"c{n}" for n in (8, 16)] for t in typecodes + complexcodes: with self.subTest(typecode=t): datatype = MPI.Datatype.fromcode(t) alignment1 = MPI._typealign(datatype) if np_dtype is not None: alignment2 = np_dtype(t).alignment self.assertEqual(alignment1, alignment2) def testAlignmentPair(self): for pairtype in mpipairtypes: alignment1 = MPI._typealign(pairtype) self.assertIn(alignment1, (2, 4, 8, 16)) if np_dtype is not None: alignment2 = tonumpy(pairtype).alignment self.assertEqual(alignment1, alignment2) def testAlignmentStruct(self): off = MPI.DOUBLE.extent structtype = MPI.Datatype.Create_struct( [1, 1], [0, off], [MPI.INT, MPI.DOUBLE], ) alignment = MPI._typealign(structtype) self.assertIsNone(alignment) structtype.Free() def testMissingNumPy(self): from mpi4py.util import dtlib np_dtype = dtlib._np_dtype dtlib._np_dtype = None try: for t in typecodes: with self.subTest(typecode=t): mt = MPI.Datatype.fromcode(t) dt = tonumpy(mt) code = mt.tocode() self.assertEqual(dt, code) arraytype = mt.Create_contiguous(7) dt = tonumpy(arraytype) arraytype.Free() self.assertIsInstance(dt, tuple) self.assertEqual(dt[0], code) self.assertEqual(dt[1], (7,)) structtype = MPI.Datatype.Create_struct( [1, 1], [0, mt.extent], [mt, mt], ) dt = tonumpy(structtype) structtype.Free() self.assertIsInstance(dt, dict) self.assertEqual(dt["formats"], [code] * 2) self.assertEqual(dt["offsets"], [0, mt.extent]) self.assertEqual(dt["itemsize"], mt.extent * 2) self.assertTrue(dt["aligned"]) with self.assertRaises(RuntimeError): fromnumpy(None) finally: dtlib._np_dtype = np_dtype @unittest.skipIf(numpy is None, "numpy") def testFailures(self): endian = ">" if np_dtype("