[mlir][Python] Add Windows DLL loader to get python extensions working there.

[mehdi_amini]

mlir/lib/Bindings/Python/mlir/_dlloader.py

13

Isn't the .dll under lib?

[stellaraccident]

mlir/lib/Bindings/Python/mlir/_dlloader.py

13

No, on windows, dlls go under the runtime directory (bin/). Static libraries only go under lib/. I think it is because on that platform the exe and dll search path is shared.

There are special cases for this everywhere.

[stellaraccident]

Rebase and test on Windows.

mlir/cmake/modules/AddMLIRPythonExtension.cmake

set_target_properties(
     # project separation perspective and a discussion on how to better
     # segment MLIR libraries needs to happen.
     LIBRARY_OUTPUT_DIRECTORY ${LLVM_BINARY_DIR}/python
     OUTPUT_NAME "${extname}"
     PREFIX "${PYTHON_MODULE_PREFIX}"
     SUFFIX "${PYTHON_MODULE_SUFFIX}${PYTHON_MODULE_EXTENSION}"
   )
 
+  if(WIN32)
+    # Need to also set the RUNTIME_OUTPUT_DIRECTORY on Windows in order to
+    # control where the .dll gets written.
+    set_target_properties(
+      ${libname} PROPERTIES
+      RUNTIME_OUTPUT_DIRECTORY ${LLVM_BINARY_DIR}/python
+    )
+  endif()
+
   # pybind11 requires binding code to be compiled with -fvisibility=hidden
   # For static linkage, better code can be generated if the entire project
   # compiles that way, but that is not enforced here. Instead, include a linker
   # script that explicitly hides anything but the PyInit_* symbols, allowing gc
   # to take place.
   set_target_properties(${libname} PROPERTIES CXX_VISIBILITY_PRESET "hidden")
 
   # Python extensions depends *only* on the public API and LLVMSupport unless

mlir/lib/Bindings/Python/CMakeLists.txt

 include(AddMLIRPythonExtension)
 add_custom_target(MLIRBindingsPythonExtension)
 ################################################################################
 # Copy python source tree.
 ################################################################################
 
 set(PY_SRC_FILES
   mlir/__init__.py
+  mlir/_dlloader.py
   mlir/ir.py
   mlir/dialects/__init__.py
   mlir/dialects/std.py
 )
 
 add_custom_target(MLIRBindingsPythonSources ALL
   DEPENDS ${PY_SRC_FILES}
 )

mlir/lib/Bindings/Python/mlir/__init__.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 # Note that the only function of this module is currently to load the
 # native module and re-export its symbols. In the future, this file is
 # reserved as a trampoline to handle environment specific loading needs
 # and arbitrate any one-time initialization needed in various shared-library
 # scenarios.
 
 __all__ = [
   "ir",
 ]
 
+# The _dlloader takes care of platform specific setup before we try to
+# load a shared library.
+from . import _dlloader
+_dlloader.preload_dependency("MLIRPublicAPI")
+
 # Expose the corresponding C-Extension module with a well-known name at this
 # top-level module. This allows relative imports like the following to
 # function:
 #   from .. import _cext
 # This reduces coupling, allowing embedding of the python sources into another
 # project that can just vary based on this top-level loader module.
 import _mlir as _cext
 

mlir/lib/Bindings/Python/mlir/_dlloader.py

+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import os
+import platform
+
+_is_windows = platform.system() == "Windows"
+_this_directory = os.path.dirname(__file__)
+
+# The standard LLVM build/install tree for Windows is laid out as:
+#   bin/
+#     MLIRPublicAPI.dll
+#   python/
+#     _mlir.*.pyd (dll extension)
+#     mlir/
+#       _dlloader.py (this file)
+# First check the python/ directory level for DLLs co-located with the pyd
+# file, and then fall back to searching the bin/ directory.
+# TODO: This should be configurable at some point.
+_dll_search_path = [
+  os.path.join(_this_directory, ".."),
+  os.path.join(_this_directory, "..", "..", "bin"),
+]
+
+# Stash loaded DLLs to keep them alive.
+_loaded_dlls = []
+
+def preload_dependency(public_name):
+  """Preloads a dylib by its soname or DLL name.
+
+  On Windows and Linux, doing this prior to loading a dependency will populate
+  the library in the flat namespace so that a subsequent library that depend
+  on it will resolve to this preloaded version.
+
+  On OSX, resolution is completely path based so this facility no-ops. On
+  Linux, as long as RPATHs are setup properly, resolution is path based but
+  this facility can still act as an escape hatch for relocatable distributions.
+  """
+  if _is_windows:
+    _preload_dependency_windows(public_name)
+
+
+def _preload_dependency_windows(public_name):
+  dll_basename = public_name + ".dll"
+  found_path = None
+  for search_dir in _dll_search_path:
+    candidate_path = os.path.join(search_dir, dll_basename)
+    if os.path.exists(candidate_path):
+      found_path = candidate_path
+      break
+
+  if found_path is None:
+    raise RuntimeError(
+      f"Unable to find dependency DLL {dll_basename} in search "
+      f"path {_dll_search_path}")
+
+  import ctypes
+  _loaded_dlls.append(ctypes.CDLL(found_path))