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

[stellaraccident]

Summary

• 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.

[stellaraccident]

I'm going to go ahead and submit this and get your feedback in post-review. It is blocking the ability to make functional distributable binaries.

mlir/lib/Bindings/Python/CMakeLists.txt

 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.
 ################################################################################
 
 set(PY_SRC_FILES
   mlir/__init__.py
   mlir/_dlloader.py
   mlir/ir.py
   mlir/dialects/__init__.py
add_custom_command(
     COMMENT "Copying python source ${PY_SRC_FILE} -> ${PY_DEST_FILE}"
     DEPENDS "${PY_SRC_FILE}"
     COMMAND "${CMAKE_COMMAND}" -E copy_if_different
         "${CMAKE_CURRENT_SOURCE_DIR}/${PY_SRC_FILE}" "${PY_DEST_FILE}"
   )
 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
   INSTALL_DIR
     python
   SOURCES
     MainModule.cpp
     IRModules.cpp

mlir/lib/Bindings/Python/mlir/__init__.py

 #  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
 
 # Note that the only function of this module is currently to load the
 # native module and re-export its symbols. In the future, this file is
 # reserved as a trampoline to handle environment specific loading needs
 # and arbitrate any one-time initialization needed in various shared-library
 # scenarios.
 
 __all__ = [
   "ir",
   "passmanager",
 ]
 
+# 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 . import _dlloader
-_dlloader.preload_dependency("MLIRPublicAPI")
+  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.
 
   Typically:
     from . import _reexport_cext
     _reexport_cext("ir", __name__)
     del _reexport_cext

mlir/lib/Bindings/Python/mlir/transforms/__init__.py

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