Refactor DLL path handling on Windows

mpi4py · May 9, 2024 · 55572a3 · 55572a3
1 parent 8228cba
commit 55572a3
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 14 deletions.
diff --git a/.cibw/merge-wheels.py b/.cibw/merge-wheels.py
@@ -158,11 +158,11 @@ def wheel_tagline(tags: list[str]) -> str:
         for variant in variant_registry:
             for module in extensions:
                 fh.write(f"_mpiabi._register({module!r}, {variant!r})\n")
-        if tags[-1].startswith("win"):
-            fh.write(textwrap.dedent("""\
-            # Set Windows DLL search path
-            _mpiabi._set_windows_dll_path()
-            """))
+        # if tags[-1].startswith("win"):
+        #     fh.write(textwrap.dedent("""\
+        #     # Set Windows DLL search path
+        #     _mpiabi._set_windows_dll_path()
+        #     """))
 
     output_dir.mkdir(parents=True, exist_ok=True)
     wheel_pack.pack(root_dir, output_dir, None)

diff --git a/.cibw/mpi4py_mpiabi.py b/.cibw/mpi4py_mpiabi.py
@@ -54,6 +54,20 @@ def add_rpath_prefix(prefix):
 
     add_rpath("")
 
+    if os.name == "nt":
+        impi_root = os.environ.get("I_MPI_ROOT")
+        if impi_root:
+            library_kind = (
+                os.environ.get("I_MPI_LIBRARY_KIND") or
+                os.environ.get("library_kind") or
+                "release"
+            )
+            add_rpath(impi_root, "bin", library_kind)
+            add_rpath(impi_root, "bin")
+        msmpi_bin = os.environ.get("MSMPI_BIN")
+        if msmpi_bin:
+            add_rpath(msmpi_bin)
+
     if sys.platform == "darwin":
         add_rpath("/usr/local/lib")
         add_rpath("/opt/homebrew/lib")
@@ -80,23 +94,52 @@ def dlopen(name):
             if hasattr(lib, "I_MPI_Check_image_status"):
                 if os.path.basename(name) != name:
                     dlopen_impi_libfabric(os.path.dirname(name))
+        if name is not None and os.name == "nt":
+            if os.path.basename(name).lower() == "impi.dll":
+                if os.path.basename(name) != name:
+                    dlopen_impi_libfabric(os.path.dirname(name))
         return lib
 
+    def search_impi_libfabric(rootdir):
+        if sys.platform == "linux":
+            libdir = "lib"
+            suffix = ".so.1"
+        else:
+            libdir = "bin"
+            suffix = ".dll"
+        for subdir in (
+            ("opt", "mpi", "libfabric", libdir),
+            ("libfabric"),
+            ("libfabric", libdir),
+            (libdir, "libfabric"),
+        ):
+            ofi_libdir = os.path.join(rootdir, *subdir)
+            ofi_filename = os.path.join(ofi_libdir, f"libfabric{suffix}")
+            if os.path.isfile(ofi_filename):
+                return ofi_filename
+        return None
+
     def dlopen_impi_libfabric(libdir):
         ofi_internal = (
             os.environ.get("I_MPI_OFI_LIBRARY_INTERNAL", "").lower()
             in ("", "1", "y", "on", "yes", "true", "enable")
         )
         ofi_required = os.environ.get("I_MPI_FABRICS") != "shm"
-        ofi_library_path = os.path.join(libdir, "libfabric")
-        ofi_provider_path = os.path.join(ofi_library_path, "prov")
-        ofi_filename = os.path.join(ofi_library_path, "libfabric.so.1")
-        if ofi_internal and ofi_required and os.path.isfile(ofi_filename):
-            if "FI_PROVIDER_PATH" not in os.environ:
-                if os.path.isdir(ofi_provider_path):
-                    os.environ["FI_PROVIDER_PATH"] = ofi_provider_path
-            ct.CDLL(ofi_filename, mode)
-            _verbose_info(f"OFI library from {ofi_filename!r}")
+        if not (ofi_internal and ofi_required):
+            return
+        rootdir = os.path.dirname(libdir)
+        if os.path.basename(rootdir).lower() in ("release", "debug"):
+            rootdir = os.path.dirname(rootdir)
+        ofi_filename = search_impi_libfabric(rootdir)
+        if ofi_filename is None:
+            return
+        if "FI_PROVIDER_PATH" not in os.environ:
+            ofi_libdir = os.path.dirname(ofi_filename)
+            ofi_provider_path = os.path.join(ofi_libdir, "prov")
+            if os.path.isdir(ofi_provider_path):
+                os.environ["FI_PROVIDER_PATH"] = ofi_provider_path
+        ct.CDLL(ofi_filename, mode)
+        _verbose_info(f"OFI library from {ofi_filename!r}")
 
     def libmpi_names():
         if os.name == "posix":