[llvm-branch-commits] [mlir] 42c57dc - [mlir][python] Tweaks to make python extensions packagable/distributable.

[Stella Laurenzo via llvm-branch-commits]

Author: Stella Laurenzo
Date: 2020-12-30T23:35:46-08:00
New Revision: 42c57dcc35ea9d6a09469c61473b17c95edcaa2e

URL: https://github.com/llvm/llvm-project/commit/42c57dcc35ea9d6a09469c61473b17c95edcaa2e
DIFF: https://github.com/llvm/llvm-project/commit/42c57dcc35ea9d6a09469c61473b17c95edcaa2e.diff

LOG: [mlir][python] Tweaks to make python extensions packagable/distributable.

* Works in tandem with prototype packaging scripts here: https://github.com/stellaraccident/mlir-py-release
* The `mlir` top-level now differentiates between in-tree builds where all packages are co-located and distribution mode where all native components are under a top-level `_mlir_libs` package.
* Also fixes the generated dialect python installation again. Hopefully the last tweak.
* With this, I am able to install and generate archives with the above setup script on Linux. Archive size=31M with just host codegen and headers/shared-libraries. Will need more linker tweaks when wiring up the next dependent project.

Differential Revision: https://reviews.llvm.org/D93936

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/CMakeLists.txt
    mlir/lib/Bindings/Python/mlir/__init__.py
    mlir/lib/Bindings/Python/mlir/transforms/__init__.py

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/CMakeLists.txt b/mlir/lib/Bindings/Python/CMakeLists.txt
index 83e978a6e046..0c34f5b55415 100644
--- a/mlir/lib/Bindings/Python/CMakeLists.txt
+++ b/mlir/lib/Bindings/Python/CMakeLists.txt
@@ -1,15 +1,6 @@
 include(AddMLIRPythonExtension)
 add_custom_target(MLIRBindingsPythonExtension)
 
-################################################################################
-# Generate dialect-specific bindings.
-################################################################################
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonStandardOps
-  StandardOps.td
-  std)
-add_dependencies(MLIRBindingsPythonExtension MLIRBindingsPythonStandardOps)
-
 ################################################################################
 # Copy python source tree.
 ################################################################################
@@ -40,6 +31,15 @@ foreach(PY_SRC_FILE ${PY_SRC_FILES})
   )
 endforeach()
 
+################################################################################
+# Generate dialect-specific bindings.
+################################################################################
+
+add_mlir_dialect_python_bindings(MLIRBindingsPythonStandardOps
+  StandardOps.td
+  std)
+add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonStandardOps)
+
 ################################################################################
 # Build core python extension
 ################################################################################

diff  --git a/mlir/lib/Bindings/Python/mlir/__init__.py b/mlir/lib/Bindings/Python/mlir/__init__.py
index 8e027f64bdb6..5ae8151e18f1 100644
--- a/mlir/lib/Bindings/Python/mlir/__init__.py
+++ b/mlir/lib/Bindings/Python/mlir/__init__.py
@@ -13,18 +13,35 @@
   "passmanager",
 ]
 
-# The _dlloader takes care of platform specific setup before we try to
-# load a shared library.
-from . import _dlloader
-_dlloader.preload_dependency("MLIRPublicAPI")
+# Packaged installs have a top-level _mlir_libs package with symbols:
+#   load_extension(name): Loads a named extension module
+#   preload_dependency(public_name): Loads a shared-library/DLL into the
+#     namespace. TODO: Remove this in favor of a more robust mechanism.
+# Conditionally switch based on whether we are in a package context.
+try:
+  import _mlir_libs
+except ModuleNotFoundError:
+  # Assume that we are in-tree.
+  # The _dlloader takes care of platform specific setup before we try to
+  # load a shared library.
+  from ._dlloader import preload_dependency as _preload_dependency
 
+  def _load_extension(name):
+    import importlib
+    return importlib.import_module(name)  # i.e. '_mlir' at the top level
+else:
+  # Packaged distribution.
+  _load_extension = _mlir_libs.load_extension
+  _preload_dependency = _mlir_libs.preload_dependency
+
+_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
+_cext = _load_extension("_mlir")
 
 def _reexport_cext(cext_module_name, target_module_name):
   """Re-exports a named sub-module of the C-Extension into another module.

diff  --git a/mlir/lib/Bindings/Python/mlir/transforms/__init__.py b/mlir/lib/Bindings/Python/mlir/transforms/__init__.py
index d6172521295a..1155c6392eb8 100644
--- a/mlir/lib/Bindings/Python/mlir/transforms/__init__.py
+++ b/mlir/lib/Bindings/Python/mlir/transforms/__init__.py
@@ -4,5 +4,5 @@
 
 # Expose the corresponding C-Extension module with a well-known name at this
 # level.
-import _mlirTransforms as _cextTransforms
-
+from .. import _load_extension
+_cextTransforms = _load_extension("_mlirTransforms")