refactor: MPI standard/legacy ABI support

Author: dalcinl

.cibw/install-mpi.sh

@@ -1,43 +1,28 @@
 #!/bin/bash
 set -euo pipefail
 
-MPI_ABI=${1:-mpi31-mpich}
+MPI_ABI=${1:-mpich}
 MACHINE=${PROCESSOR_ARCHITECTURE:-$(uname -m)}
 MPIARCH=${2:-$MACHINE}
 MPIARCH=${MPIARCH/native/$MACHINE}
 
 MPI_CHANNEL=conda-forge
-MPI_PACKAGE=${MPI_ABI#*-}
-MPI_VERSION="*.*"
-case "$(uname)"-"$MPIARCH"-"$MPI_ABI" in
-# Linux x86_64/aarch64
-Linux-*-mpi41-mpich)         MPI_VERSION=4.2;; # >= 4.2
-Linux-*-mpi40-mpich)         MPI_VERSION=4.1;; # >= 4.0
-Linux-*-mpi31-mpich)         MPI_VERSION=3.4;; # >= 3.2
-Linux-*-mpi31-openmpi)       MPI_VERSION=4.1;; # >= 3.1
-# Darwin x86_64
-Darwin-x86_64-mpi41-mpich)   MPI_VERSION=4.2;;
-Darwin-x86_64-mpi40-mpich)   MPI_VERSION=4.0;;
-Darwin-x86_64-mpi31-mpich)   MPI_VERSION=3.2;;
-Darwin-x86_64-mpi31-openmpi) MPI_VERSION=3.1;;
-# Darwin arm64
-Darwin-arm64-mpi41-mpich)    MPI_VERSION=4.2;;
-Darwin-arm64-mpi40-mpich)    MPI_VERSION=4.0;;
-Darwin-arm64-mpi31-mpich)    MPI_VERSION=3.4;;
-Darwin-arm64-mpi31-openmpi)  MPI_VERSION=4.0;;
-# Windows AMD64
-*NT*-AMD64-mpi20-msmpi)
-    MPI_CHANNEL=conda-forge
-    MPI_PACKAGE=msmpi
-    MPI_VERSION=10.1.1
-    MPI_ROOT=${MPI_ROOT:-~/mpi}
-    ;;
-*NT*-AMD64-mpi31-impi)
-    MPI_CHANNEL=conda-forge
-    MPI_PACKAGE=impi-devel
-    MPI_VERSION=2021.14.0
-    MPI_ROOT=${MPI_ROOT:-~/mpi}
-    ;;
+MPI_PACKAGE=${MPI_ABI}
+case "$MPI_ABI" in
+    mpich)   MPI_VERSION=4 ;;
+    openmpi) MPI_VERSION=5 ;;
+    msmpi)   MPI_VERSION=10.1.1 ;;
+    impi)    MPI_VERSION=2021.14.0 MPI_PACKAGE=impi-devel ;;
+esac
+case "$(uname)" in
+    Linux|Darwin)
+        MPI_ROOT=${MPI_ROOT:-/usr/local}
+        sudo() { [ "$(id -u)" -eq 0 ] || set -- command sudo "$@"; "$@"; }
+        ;;
+    *NT*)
+        MPI_ROOT=${MPI_ROOT:-~/mpi}
+        sudo() { "$@"; }
+        ;;
 esac
 
 echo "Install Micromamba"
@@ -56,7 +41,7 @@ micromamba create --yes --always-copy \
            --prefix "$envdir" \
            --relocate-prefix "$MPI_ROOT" \
            "$MPI_PACKAGE"="$MPI_VERSION"
-test "$(uname)-$MPIARCH-$MPI_ABI" = "Linux-x86_64-mpi41-mpich" && \
+test "$(micromamba list --json --prefix "$envdir" attr)" != "[]" && \
 micromamba remove --yes --force --prefix "$envdir" attr
 micromamba list --prefix "$envdir"
 
@@ -73,16 +58,15 @@ if [ "$MPI_PACKAGE" == openmpi ]; then
     files=("$envdir"/share/openmpi/mpi{cc,c++,fort}-wrapper-data.txt)
     sed -i.orig -E 's/(compiler)=(.*)-(.*)/\1=\3/' "${files[@]}"
     sed -i.orig "s%-Wl,-rpath,$MPI_ROOT/lib%%g" "${files[@]}"
+    sed -i.orig "s%-Wl,-allow-shlib-undefined%%g" "${files[@]}"
+    sed -i.orig "s%-Wl,-rpath -Wl,\${libdir}%%g" "${files[@]}"
 fi
 
-echo
-# shellcheck disable=SC2015
-SUDO=$(test "$(id -u)" -ne 0 && command -v sudo || true)
 echo "Copying MPI to $MPI_ROOT"
-$SUDO mkdir -p "$MPI_ROOT"
-$SUDO cp -RP "$envdir"/. "$MPI_ROOT"
+sudo mkdir -p "$MPI_ROOT"
+sudo cp -RP "$envdir"/. "$MPI_ROOT"
 echo "Rebuild dynamic linker cache"
-$SUDO "$(command -v ldconfig || echo true)"
+sudo "$(command -v ldconfig || echo true)"
 
 echo "Display MPI information"
 if [ "$MPI_PACKAGE" == mpich   ]; then mpichversion; fi
@@ -107,7 +91,7 @@ if [ "$(uname)" == Darwin ] && [ "$MPIARCH" != "$MACHINE" ]; then
     libs=$(find "$envdir2/lib" -type f -name 'lib*.dylib')
     for lib in $libs; do
         lib=$(basename "$lib")
-        $SUDO lipo -create \
+        sudo lipo -create \
               "$envdir1/lib/$lib" \
               "$envdir2/lib/$lib" \
               -output "$MPI_ROOT/lib/$lib"

.cibw/mpi4py_mpiabi.py

@@ -4,7 +4,6 @@
 import importlib.machinery
 import importlib.util
 import os
-import re
 import sys
 import warnings
 
@@ -196,88 +195,68 @@ def _get_mpiabi_from_libmpi(libmpi=None):
     # pylint: disable=import-outside-toplevel
     import ctypes as ct
     lib = _dlopen_libmpi(libmpi)
-    lib.MPI_Get_version.restype = ct.c_int
-    lib.MPI_Get_version.argtypes = [ct.POINTER(ct.c_int)] * 2
-    vmajor, vminor = ct.c_int(0), ct.c_int(0)
-    ierr = lib.MPI_Get_version(ct.byref(vmajor), ct.byref(vminor))
-    if ierr:  # pragma: no cover
-        message = f"MPI_Get_version [ierr={ierr}]"
-        raise RuntimeError(message)
-    vmajor, vminor = vmajor.value, vminor.value
+    abi_get_version = getattr(lib, "MPI_Abi_get_version", None)
+    if abi_get_version:
+        abi_get_version.restype = ct.c_int
+        abi_get_version.argtypes = [ct.POINTER(ct.c_int)] * 2
+        abi_major, abi_minor = ct.c_int(0), ct.c_int(0)
+        ierr = abi_get_version(ct.byref(abi_major), ct.byref(abi_minor))
+        if ierr:  # pragma: no cover
+            message = f"MPI_Abi_get_version [ierr={ierr}]"
+            raise RuntimeError(message)
+        if abi_major.value > 0:
+            return "mpiabi"
     if os.name == "posix":
         openmpi = hasattr(lib, "ompi_mpi_comm_self")
-        family = "openmpi" if openmpi else "mpich"
+        mpiabi = "openmpi" if openmpi else "mpich"
     else:
         msmpi = hasattr(lib, "MSMPI_Get_version")
-        family = "msmpi" if msmpi else "impi"
-    return (vmajor, vminor), family
-
-
-_pattern = re.compile(
-    r"""
-    \.? (
-    (mpi)? (\W|_)* (
-    (?P<vmajor>\d+) \.?
-    (?P<vminor>\d)
-    ) )? (\W|_)*
-    (?P<family>\w+)?
-    """,
-    re.VERBOSE | re.IGNORECASE,
-)
-_pattern_strict = re.compile(
-    r"mpi(?P<vmajor>\d+)(?P<vminor>\d)(-(?P<family>\w+))?",
-    re.VERBOSE | re.IGNORECASE,
-)
-
-
-def _get_mpiabi_from_string(string, strict=False):
-    pattern = _pattern_strict if strict else _pattern
-    match = pattern.match(string)
-    if match is None:
-        message = f"invalid MPI ABI string {string!r}"
-        raise RuntimeError(message)
-    vmajor = match.group("vmajor") or "4"
-    vminor = match.group("vminor") or "0"
-    family = match.group("family") or ""
-    return (int(vmajor), int(vminor)), family.lower() or None
+        mpiabi = "msmpi" if msmpi else "impi"
+    return mpiabi
+
+
+def _get_mpiabi_from_string(string):
+    table = {ord(c): "" for c in " -_"}
+    mpiabi = string.translate(table).lower()
+    if os.name == "posix":
+        if mpiabi == "impi":
+            mpiabi = "mpich"
+    else:
+        if mpiabi == "mpich":
+            mpiabi = "impi"
+    return mpiabi
 
 
 def _get_mpiabi():
-    version = getattr(_get_mpiabi, "version", None)
-    family = getattr(_get_mpiabi, "family", None)
-    if version is None:
-        string = os.environ.get("MPI4PY_MPIABI") or None
+    mpiabi = getattr(_get_mpiabi, "mpiabi", None)
+    if mpiabi is None:
+        mpiabi = os.environ.get("MPI4PY_MPIABI") or None
         libmpi = os.environ.get("MPI4PY_LIBMPI") or None
-        if string is not None:
-            version, family = _get_mpiabi_from_string(string)
+        if mpiabi is not None:
+            mpiabi = _get_mpiabi_from_string(mpiabi)
         else:
-            version, family = _get_mpiabi_from_libmpi(libmpi)
-        _get_mpiabi.version = version  # pyright: ignore
-        _get_mpiabi.family = family  # pyright: ignore
-    return version, family
+            mpiabi = _get_mpiabi_from_libmpi(libmpi)
+        _get_mpiabi.mpiabi = mpiabi  # pyright: ignore
+    return mpiabi
 
 
-_registry = {}  # type: dict[str, dict[str, list[tuple[int, int]]]]
+_registry = {}  # type: dict[str, list[str]]
 
 
 def _register(module, mpiabi):
-    version, family = _get_mpiabi_from_string(mpiabi, strict=True)
-    versions = _registry.setdefault(module, {}).setdefault(family, [])
-    versions.append(version)
-    versions.sort()
+    mpiabi = _get_mpiabi_from_string(mpiabi)
+    registered = _registry.setdefault(module, [])
+    if mpiabi not in registered:
+        registered.append(mpiabi)
 
 
 def _get_mpiabi_suffix(module):
     if module not in _registry:
         return None
-    version, family = _get_mpiabi()
-    versions = _registry[module].get(family)
-    if versions:
-        vmin, vmax = versions[0], versions[-1]
-        version = max(vmin, min(version, vmax))
-    vmajor, vminor = version
-    family_tag = f"-{family}" if family else ""
-    return f".mpi{vmajor}{vminor}{family_tag}"
+    mpiabi = _get_mpiabi()
+    if mpiabi not in _registry[module]:
+        return None
+    return f".{mpiabi}" if mpiabi else ""
 
 
 class _Finder:

.cibw/run-tests-conda.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+set -euo pipefail
+
+mpich=("4.3" "4.1" "3.4")
+openmpi=("5.0" "4.1")
+impi=("2021.14.1" "2021.10.0")
+msmpi=("10.1.1")
+
+mpi="$1"
+mpipackage="$mpi"
+mpiversion="${mpi}[@]"
+test "$mpi" = impi && mpipackage=impi_rt
+
+CONDA=$(command -v micromamba || command -v mamba || command -v conda)
+scriptdir=$(dirname "${BASH_SOURCE[0]}")
+for version in "${!mpiversion}"; do
+    echo "::group::$mpipackage=$version"
+    "$CONDA" install -qy "$mpipackage=$version"
+    "$CONDA" list
+    "$scriptdir"/run-tests-mpi.sh
+    echo "::endgroup::"
+done

.cibw/run-tests-mpi.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+set -euo pipefail
+
+: "${MPIEXEC=mpiexec}"
+: "${PYTHON=python}"
+
+{ set -x; } 2>/dev/null
+"$PYTHON" -m mpi4py --prefix
+"$PYTHON" -m mpi4py --version
+"$PYTHON" -m mpi4py --mpi-library
+"$PYTHON" -m mpi4py --mpi-std-version
+"$PYTHON" -m mpi4py --mpi-lib-version | { head -n 1; } 2>/dev/null
+"$MPIEXEC" -n 2 "$PYTHON" -m mpi4py.bench ringtest
+"$MPIEXEC" -n 2 "$PYTHON" -m mpi4py.bench helloworld
+{ set +x; } 2>/dev/null

.cibw/run-tests.sh

@@ -20,5 +20,6 @@ set -x
 python --version
 python -m mpi4py --prefix
 python -m mpi4py --version
+python -m mpi4py --mpi-library
 python -m mpi4py --mpi-std-version
 python -m mpi4py --mpi-lib-version

.cibw/setup-build.py renamed to .cibw/setup-matrix.py

@@ -37,20 +37,12 @@ def pp3(y_min=8, y_max=10):
 }
 
 MPI_ABI_POSIX = [
-    f"mpi{std}-{mpi}"
-    for std in (31, 40, 41)
-    for mpi in ("mpich", "openmpi")
-    if (std, mpi) not in (
-        (40, "openmpi"),
-        (41, "openmpi"),
-    )
+    "mpich",
+    "openmpi",
 ]
 MPI_ABI_WINNT = [
-    f"mpi{std}-{mpi}"
-    for std, mpi in (
-        (20, "msmpi"),
-        (31, "impi"),
-    )
+    "impi",
+    "msmpi",
 ]
 MPI_ABI = {
     "Linux": MPI_ABI_POSIX[:],
@@ -65,12 +57,12 @@ def pp3(y_min=8, y_max=10):
         None: "ubuntu-latest"
     },
     "macOS": {
-        "arm64": "macos-14",
+        "arm64": "macos-15",
         "x86_64": "macos-13",
         None: "macos-latest"
     },
     "Windows": {
-        "AMD64": "windows-2019",
+        "AMD64": "windows-2022",
         None: "windows-latest"
     },
 }
@@ -107,28 +99,52 @@ def pp3(y_min=8, y_max=10):
 
 matrix_build = [
     {
-        "os": os, "arch": arch,
-        "py": py, "mpi-abi": mpi_abi,
+        "os": os,
+        "arch": arch,
+        "py": py,
+        "mpi-abi": mpi_abi,
         "runner": GHA_RUNNER[os][arch],
     }
     for os in os_arch_py
     for arch in os_arch_py[os]
     for py in os_arch_py[os][arch]
     for mpi_abi in MPI_ABI[os]
 ]
+
 matrix_merge = [
     {
-        "os": os, "arch": arch,
+        "os": os,
+        "arch": arch,
         "runner": GHA_RUNNER[os][None],
     }
     for os in os_arch_py
     for arch in os_arch_py[os]
 ]
-os_arch_list = [
-    "{os}-{arch}".format(**row)
-    for row in matrix_merge
-]
+
+matrix_test = []
+for build in matrix_build:
+    os = build["os"]
+    arch = build["arch"]
+    pytag = build["py"]
+    mpi_abi = build["mpi-abi"]
+    runner = GHA_RUNNER[os][arch]
+    if pytag.startswith("pp"):
+        continue
+    pyver = pytag[2:3] + "." + pytag[3:]
+    mpilist = [mpi_abi]
+    if (os, arch, mpi_abi) == ("Linux", "x86_64", "mpich"):
+        mpilist.insert(0, "impi")
+    matrix_test += [
+        {
+            "mpi": mpi,
+            "py": pyver,
+            "os": os,
+            "arch": arch,
+            "runner": runner,
+        }
+        for mpi in mpilist
+    ]
 
 print(f"matrix-build={json.dumps(matrix_build)}")
 print(f"matrix-merge={json.dumps(matrix_merge)}")
-print(f"os-arch-list={json.dumps(os_arch_list)}")
+print(f"matrix-test={json.dumps(matrix_test)}")