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 @@ -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 @@ ) 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 --- 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 --- 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")