diff --git a/mlir/cmake/modules/AddMLIRPythonExtension.cmake b/mlir/cmake/modules/AddMLIRPythonExtension.cmake --- a/mlir/cmake/modules/AddMLIRPythonExtension.cmake +++ b/mlir/cmake/modules/AddMLIRPythonExtension.cmake @@ -75,6 +75,15 @@ 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 diff --git a/mlir/lib/Bindings/Python/CMakeLists.txt b/mlir/lib/Bindings/Python/CMakeLists.txt --- a/mlir/lib/Bindings/Python/CMakeLists.txt +++ b/mlir/lib/Bindings/Python/CMakeLists.txt @@ -6,6 +6,8 @@ set(PY_SRC_FILES mlir/__init__.py + mlir/_dlloader.py + mlir/ir.py mlir/dialects/__init__.py mlir/ir.py mlir/passmanager.py diff --git a/mlir/lib/Bindings/Python/mlir/__init__.py b/mlir/lib/Bindings/Python/mlir/__init__.py --- a/mlir/lib/Bindings/Python/mlir/__init__.py +++ b/mlir/lib/Bindings/Python/mlir/__init__.py @@ -13,6 +13,11 @@ "passmanager", ] +# 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: diff --git a/mlir/lib/Bindings/Python/mlir/_dlloader.py b/mlir/lib/Bindings/Python/mlir/_dlloader.py new file mode 100644 --- /dev/null +++ b/mlir/lib/Bindings/Python/mlir/_dlloader.py @@ -0,0 +1,59 @@ +# 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))