[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

diff --git a/mlir/lib/Bindings/Python/CMakeLists.txt b/mlir/lib/Bindings/Python/CMakeLists.txt
index 83e978a..0c34f5b 100644 (file)
--- a/mlir/lib/Bindings/Python/CMakeLists.txt
+++ b/mlir/lib/Bindings/Python/CMakeLists.txt
@@ -2,15 +2,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.
 ################################################################################
 
@@ -41,6 +32,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
 ################################################################################
 add_mlir_python_extension(MLIRCoreBindingsPythonExtension _mlir

diff --git a/mlir/lib/Bindings/Python/mlir/__init__.py b/mlir/lib/Bindings/Python/mlir/__init__.py
index 8e027f6..5ae8151 100644 (file)
--- a/mlir/lib/Bindings/Python/mlir/__init__.py
+++ b/mlir/lib/Bindings/Python/mlir/__init__.py
@@ -13,18 +13,35 @@ __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")
+# 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 d617252..1155c63 100644 (file)
--- 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")