From 99b1c42fd3ac405471ff56402c270a1056383f82 Mon Sep 17 00:00:00 2001
From: Stella Laurenzo <stellaraccident@gmail.com>
Date: Fri, 6 Nov 2020 09:48:20 -0800
Subject: [PATCH] [mlir][Python] Add Windows DLL loader to get python
 extensions working there.

Differential Revision: https://reviews.llvm.org/D90958
---
 .../modules/AddMLIRPythonExtension.cmake      |  9 +++
 mlir/lib/Bindings/Python/CMakeLists.txt       |  2 +
 mlir/lib/Bindings/Python/mlir/__init__.py     |  5 ++
 mlir/lib/Bindings/Python/mlir/_dlloader.py    | 59 +++++++++++++++++++
 4 files changed, 75 insertions(+)
 create mode 100644 mlir/lib/Bindings/Python/mlir/_dlloader.py

diff --git a/mlir/cmake/modules/AddMLIRPythonExtension.cmake b/mlir/cmake/modules/AddMLIRPythonExtension.cmake
index 8a74675e2543..43ad869a400b 100644
--- a/mlir/cmake/modules/AddMLIRPythonExtension.cmake
+++ b/mlir/cmake/modules/AddMLIRPythonExtension.cmake
@@ -75,6 +75,15 @@ function(add_mlir_python_extension libname extname)
     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
index bf6f1d8d5567..596bff23093e 100644
--- a/mlir/lib/Bindings/Python/CMakeLists.txt
+++ b/mlir/lib/Bindings/Python/CMakeLists.txt
@@ -6,6 +6,8 @@ add_custom_target(MLIRBindingsPythonExtension)
 
 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
index c63c4332be68..8e027f64bdb6 100644
--- a/mlir/lib/Bindings/Python/mlir/__init__.py
+++ b/mlir/lib/Bindings/Python/mlir/__init__.py
@@ -13,6 +13,11 @@ __all__ = [
   "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
index 000000000000..454a7b7f137f
--- /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))