[Mlir-commits] [mlir] 310c949 - Re-engineer MLIR python build support.

Stella Laurenzo llvmlistbot at llvm.org
Tue Jul 27 09:03:07 PDT 2021


Author: Stella Laurenzo
Date: 2021-07-27T15:54:58Z
New Revision: 310c9496d80961188e8d8f8ad306cdf44bd7541f

URL: https://github.com/llvm/llvm-project/commit/310c9496d80961188e8d8f8ad306cdf44bd7541f
DIFF: https://github.com/llvm/llvm-project/commit/310c9496d80961188e8d8f8ad306cdf44bd7541f.diff

LOG: Re-engineer MLIR python build support.

* Implements all of the discussed features:
  - Links against common CAPI libraries that are self contained.
  - Stops using the 'python/' directory at the root for everything, opening the namespace up for multiple projects to embed the MLIR python API.
  - Separates declaration of sources (py and C++) needed to build the extension from building, allowing external projects to build custom assemblies from core parts of the API.
  - Makes the core python API relocatable (i.e. it could be embedded as something like 'npcomp.ir', 'npcomp.dialects', etc). Still a bit more to do to make it truly isolated but the main structural reset is done.
  - When building statically, installed python packages are completely self contained, suitable for direct setup and upload to PyPi, et al.
  - Lets external projects assemble their own CAPI common runtime library that all extensions use. No more possibilities for TypeID issues.
  - Begins modularizing the API so that external projects that just include a piece pay only for what they use.
* I also rolled in a re-organization of the native libraries that matches how I was packaging these out of tree and is a better layering (i.e. all libraries go into a nested _mlir_libs package). There is some further cleanup that I resisted since it would have required source changes that I'd rather do in a followup once everything stabilizes.
* Note that I made a somewhat odd choice in choosing to recompile all extensions for each project they are included into (as opposed to compiling once and just linking). While not leveraged yet, this will let us set definitions controlling the namespacing of the extensions so that they can be made to not conflict across projects (with preprocessor definitions).
* This will be a relatively substantial breaking change for downstreams. I will handle the npcomp migration and will coordinate with the circt folks before landing. We should stage this and make sure it isn't causing problems before landing.
* Fixed a couple of absolute imports that were causing issues.

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

Added: 
    mlir/python/mlir/_mlir_libs/__init__.py
    mlir/python/mlir/dialects/PythonTest.td

Modified: 
    mlir/cmake/modules/AddMLIR.cmake
    mlir/cmake/modules/AddMLIRPython.cmake
    mlir/docs/Bindings/Python.md
    mlir/lib/CMakeLists.txt
    mlir/python/CMakeLists.txt
    mlir/python/mlir/_cext_loader.py
    mlir/python/mlir/dialects/_builtin_ops_ext.py
    mlir/python/mlir/dialects/_linalg_ops_ext.py
    mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py
    mlir/test/CMakeLists.txt
    mlir/test/lit.cfg.py
    mlir/test/python/ir/operation.py

Removed: 
    mlir/lib/Bindings/CMakeLists.txt
    mlir/lib/Bindings/Python/CMakeLists.txt
    mlir/lib/Bindings/Python/Conversions/CMakeLists.txt
    mlir/lib/Bindings/Python/Transforms/CMakeLists.txt
    mlir/python/mlir/dialects/CMakeLists.txt
    mlir/test/python/CMakeLists.txt


################################################################################
diff  --git a/mlir/cmake/modules/AddMLIR.cmake b/mlir/cmake/modules/AddMLIR.cmake
index 76f6841f772e9..1c24c91ba2019 100644
--- a/mlir/cmake/modules/AddMLIR.cmake
+++ b/mlir/cmake/modules/AddMLIR.cmake
@@ -54,7 +54,7 @@ endfunction()
 #   with large dependencies.
 function(add_mlir_library name)
   cmake_parse_arguments(ARG
-    "SHARED;INSTALL_WITH_TOOLCHAIN;EXCLUDE_FROM_LIBMLIR"
+    "SHARED;INSTALL_WITH_TOOLCHAIN;EXCLUDE_FROM_LIBMLIR;DISABLE_INSTALL"
     ""
     "ADDITIONAL_HEADERS;DEPENDS;LINK_COMPONENTS;LINK_LIBS"
     ${ARGN})
@@ -131,7 +131,9 @@ function(add_mlir_library name)
 
   if(TARGET ${name})
     target_link_libraries(${name} INTERFACE ${LLVM_COMMON_LIBS})
-    add_mlir_library_install(${name})
+    if(NOT ARG_DISABLE_INSTALL)
+      add_mlir_library_install(${name})
+    endif()
   else()
     # Add empty "phony" target
     add_custom_target(${name})

diff  --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 046acbc085dd7..a552e9b5fc1f4 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -1,10 +1,404 @@
+################################################################################
+# Python modules
+# MLIR's Python modules are both directly used by the core project and are
+# available for use and embedding into external projects (in their own
+# namespace and with their own deps). In order to facilitate this, python
+# artifacts are split between declarations, which make a subset of
+# things available to be built and "add", which in line with the normal LLVM
+# nomenclature, adds libraries.
+################################################################################
+
+# Function: declare_mlir_python_sources
+# Declares pure python sources as part of a named grouping that can be built
+# later.
+# Arguments:
+#   ROOT_DIR: Directory where the python namespace begins (defaults to
+#     CMAKE_CURRENT_SOURCE_DIR). For non-relocatable sources, this will
+#     typically just be the root of the python source tree (current directory).
+#     For relocatable sources, this will point deeper into the directory that
+#     can be relocated. For generated sources, can be relative to
+#     CMAKE_CURRENT_BINARY_DIR. Generated and non generated sources cannot be
+#     mixed.
+#   ADD_TO_PARENT: Adds this source grouping to a previously declared source
+#     grouping. Source groupings form a DAG.
+#   SOURCES: List of specific source files relative to ROOT_DIR to include.
+#   SOURCES_GLOB: List of glob patterns relative to ROOT_DIR to include.
+function(declare_mlir_python_sources name)
+  cmake_parse_arguments(ARG
+    ""
+    "ROOT_DIR;ADD_TO_PARENT"
+    "SOURCES;SOURCES_GLOB"
+    ${ARGN})
+
+  if(NOT ARG_ROOT_DIR)
+    set(ARG_ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
+  endif()
+
+  # Process the glob.
+  set(_glob_sources)
+  if(ARG_SOURCES_GLOB)
+    set(_glob_spec ${ARG_SOURCES_GLOB})
+    list(TRANSFORM _glob_spec PREPEND "${ARG_ROOT_DIR}/")
+    file(GLOB_RECURSE _glob_sources
+      RELATIVE "${ARG_ROOT_DIR}"
+      ${_glob_spec}
+    )
+    list(APPEND ARG_SOURCES ${_glob_sources})
+  endif()
+
+  # We create a custom target to carry properties and dependencies for
+  # generated sources.
+  add_custom_target(${name})
+  set(_file_depends "${ARG_SOURCES}")
+  list(TRANSFORM _file_depends PREPEND "${ARG_ROOT_DIR}/")
+  set_target_properties(${name} PROPERTIES
+    PYTHON_SOURCES_TYPE pure
+    PYTHON_ROOT_DIR "${ARG_ROOT_DIR}"
+    PYTHON_SOURCES "${ARG_SOURCES}"
+    PYTHON_FILE_DEPENDS "${_file_depends}"
+    PYTHON_DEPENDS ""
+  )
+
+  # Add to parent.
+  if(ARG_ADD_TO_PARENT)
+    set_property(TARGET ${ARG_ADD_TO_PARENT} APPEND PROPERTY PYTHON_DEPENDS ${name})
+  endif()
+endfunction()
+
+# Function: declare_mlir_python_extension
+# Declares a buildable python extension from C++ source files. The built
+# module is considered a python source file and included as everything else.
+# Arguments:
+#   MODULE_NAME: Local import name of the module (i.e. "_mlir").
+#   ADD_TO_PARENT: Same as for declare_mlir_python_sources.
+#   SOURCES: C++ sources making up the module.
+#   PRIVATE_LINK_LIBS: List of libraries to link in privately to the module
+#     regardless of how it is included in the project (generally should be
+#     static libraries that can be included with hidden visibility).
+#   EMBED_CAPI_LINK_LIBS: Dependent CAPI libraries that this extension depends
+#     on. These will be collected for all extensions and put into an
+#     aggregate dylib that is linked against.
+function(declare_mlir_python_extension name)
+  cmake_parse_arguments(ARG
+    ""
+    "MODULE_NAME;ADD_TO_PARENT"
+    "SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS"
+    ${ARGN})
+
+  add_custom_target(${name})
+  set_target_properties(${name} PROPERTIES
+    PYTHON_SOURCES_TYPE extension
+    PYTHON_EXTENSION_MODULE_NAME "${ARG_MODULE_NAME}"
+    PYTHON_CPP_SOURCES "${ARG_SOURCES}"
+    PYTHON_PRIVATE_LINK_LIBS "${ARG_PRIVATE_LINK_LIBS}"
+    PYTHON_EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
+    PYTHON_FILE_DEPENDS ""
+    PYTHON_DEPENDS ""
+  )
+
+  # Add to parent.
+  if(ARG_ADD_TO_PARENT)
+    set_property(TARGET ${ARG_ADD_TO_PARENT} APPEND PROPERTY PYTHON_DEPENDS ${name})
+  endif()
+endfunction()
+
+# Function: add_mlir_python_modules
+# Adds python modules to a project, building them from a list of declared
+# source groupings (see declare_mlir_python_sources and
+# declare_mlir_python_extension). One of these must be called for each
+# packaging root in use.
+# Arguments:
+#   ROOT_PREFIX: The directory in the build tree to emit sources. This will
+#     typically be something like ${MY_BINARY_DIR}/python_packages/foobar
+#     for non-relocatable modules or a deeper directory tree for relocatable.
+#   INSTALL_PREFIX: Prefix into the install tree for installing the package.
+#     Typically mirrors the path above but without an absolute path.
+#   DECLARED_SOURCES: List of declared source groups to include. The entire
+#     DAG of source modules is included.
+#   COMMON_CAPI_LINK_LIBS: List of dylibs (typically one) to make every
+#     extension depend on (see mlir_python_add_common_capi_library).
+function(add_mlir_python_modules name)
+  cmake_parse_arguments(ARG
+    ""
+    "ROOT_PREFIX;INSTALL_PREFIX;COMMON_CAPI_LINK_LIBS"
+    "DECLARED_SOURCES"
+    ${ARGN})
+  # Helper to process an individual target.
+  function(_process_target modules_target sources_target)
+    get_target_property(_source_type ${sources_target} PYTHON_SOURCES_TYPE)
+    if(_source_type STREQUAL "pure")
+      # Pure python sources to link into the tree.
+      get_target_property(_python_root_dir ${sources_target} PYTHON_ROOT_DIR)
+      get_target_property(_python_sources ${sources_target} PYTHON_SOURCES)
+      foreach(_source_relative_path ${_python_sources})
+        set(_src_path "${_python_root_dir}/${_source_relative_path}")
+        set(_dest_path "${ARG_ROOT_PREFIX}/${_source_relative_path}")
+
+        get_filename_component(_dest_dir "${_dest_path}" DIRECTORY)
+        get_filename_component(_install_path "${ARG_INSTALL_PREFIX}/${_source_relative_path}" DIRECTORY)
+
+        file(MAKE_DIRECTORY "${_dest_dir}")
+        add_custom_command(
+          TARGET ${modules_target} PRE_BUILD
+          COMMENT "Copying python source ${_src_path} -> ${_dest_path}"
+          DEPENDS "${_src_path}"
+          BYPRODUCTS "${_dest_path}"
+          COMMAND "${CMAKE_COMMAND}" -E create_symlink
+              "${_src_path}" "${_dest_path}"
+        )
+        install(
+          FILES "${_src_path}"
+          DESTINATION "${_install_path}"
+          COMPONENT ${modules_target}
+        )
+      endforeach()
+    elseif(_source_type STREQUAL "extension")
+      # Native CPP extension.
+      get_target_property(_module_name ${sources_target} PYTHON_EXTENSION_MODULE_NAME)
+      get_target_property(_cpp_sources ${sources_target} PYTHON_CPP_SOURCES)
+      get_target_property(_private_link_libs ${sources_target} PYTHON_PRIVATE_LINK_LIBS)
+      set(_extension_target "${name}.extension.${_module_name}.dso")
+      add_mlir_python_extension(${_extension_target} "${_module_name}"
+        INSTALL_COMPONENT ${modules_target}
+        INSTALL_DIR "${ARG_INSTALL_PREFIX}/_mlir_libs"
+        OUTPUT_DIRECTORY "${ARG_ROOT_PREFIX}/_mlir_libs"
+        SOURCES ${_cpp_sources}
+        LINK_LIBS PRIVATE
+          ${_private_link_libs}
+          ${ARG_COMMON_CAPI_LINK_LIBS}
+      )
+      add_dependencies(${name} ${_extension_target})
+      mlir_python_setup_extension_rpath(${_extension_target})
+    else()
+      message(SEND_ERROR "Unrecognized source type '${_source_type}' for python source target ${sources_target}")
+      return()
+    endif()
+  endfunction()
+
+  _flatten_mlir_python_targets(_flat_targets ${ARG_DECLARED_SOURCES})
+  # Collect dependencies.
+  set(_depends)
+  foreach(sources_target ${_flat_targets})
+    get_target_property(_local_depends ${sources_target} PYTHON_FILE_DEPENDS)
+    list(APPEND _depends ${_local_depends})
+  endforeach()
+
+  # Build the modules target.
+  add_custom_target(${name} ALL DEPENDS ${_depends})
+  foreach(sources_target ${_flat_targets})
+    _process_target(${name} ${sources_target})
+  endforeach()
+
+  # Create an install target.
+  if (NOT LLVM_ENABLE_IDE)
+    add_llvm_install_targets(
+      install-${name}
+      DEPENDS ${name}
+      COMPONENT ${name})
+  endif()
+endfunction()
+
+# Function: declare_mlir_dialect_python_bindings
+# Helper to generate source groups for dialects, including both static source
+# files and a TD_FILE to generate wrappers.
+#
+# This will generate a source group named ${ADD_TO_PARENT}.${DIALECT_NAME}.
+#
+# Arguments:
+#   ROOT_DIR: Same as for declare_mlir_python_sources().
+#   ADD_TO_PARENT: Same as for declare_mlir_python_sources(). Unique names
+#     for the subordinate source groups are derived from this.
+#   TD_FILE: Tablegen file to generate source for (relative to ROOT_DIR).
+#   DIALECT_NAME: Python name of the dialect.
+#   SOURCES: Same as declare_mlir_python_sources().
+#   SOURCES_GLOB: Same as declare_mlir_python_sources().
+#   DEPENDS: Additional dependency targets.
+function(declare_mlir_dialect_python_bindings)
+  cmake_parse_arguments(ARG
+    ""
+    "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME"
+    "SOURCES;SOURCES_GLOB;DEPENDS"
+    ${ARGN})
+  # Sources.
+  set(_dialect_target "${ARG_ADD_TO_PARENT}.${ARG_DIALECT_NAME}")
+  declare_mlir_python_sources(${_dialect_target}
+    ROOT_DIR "${ARG_ROOT_DIR}"
+    ADD_TO_PARENT "${ARG_ADD_TO_PARENT}"
+    SOURCES "${ARG_SOURCES}"
+    SOURCES_GLOB "${ARG_SOURCES_GLOB}"
+  )
+
+  # Tablegen
+  if(ARG_TD_FILE)
+    set(tblgen_target "${ARG_ADD_TO}.${ARG_DIALECT_NAME}.tablegen")
+    set(td_file "${ARG_ROOT_DIR}/${ARG_TD_FILE}")
+    get_filename_component(relative_td_directory "${ARG_TD_FILE}" DIRECTORY)
+    set(dialect_filename "${relative_td_directory}/_${ARG_DIALECT_NAME}_ops_gen.py")
+    set(LLVM_TARGET_DEFINITIONS ${td_file})
+    mlir_tablegen("${dialect_filename}" -gen-python-op-bindings
+                  -bind-dialect=${ARG_DIALECT_NAME})
+    add_public_tablegen_target(${tblgen_target})
+    if(ARG_DEPENDS)
+      add_dependencies(${tblgen_target} ${ARG_DEPENDS})
+    endif()
+
+    # Generated.
+    declare_mlir_python_sources("${ARG_ADD_TO_PARENT}.${ARG_DIALECT_NAME}.ops_gen"
+      ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
+      ADD_TO_PARENT "${_dialect_target}"
+      SOURCES "${dialect_filename}"
+    )
+  endif()
+endfunction()
+
+# Function: mlir_python_setup_extension_rpath
+# Sets RPATH properties on a target, assuming that it is being output to
+# an _mlir_libs directory with all other libraries. For static linkage,
+# the RPATH will just be the origin. If linking dynamically, then the LLVM
+# library directory will be added.
+# Arguments:
+#   RELATIVE_INSTALL_ROOT: If building dynamically, an RPATH entry will be
+#     added to the install tree lib/ directory by first traversing this
+#     path relative to the installation location. Typically a number of ".."
+#     entries, one for each level of the install path.
+function(mlir_python_setup_extension_rpath target)
+  cmake_parse_arguments(ARG
+    ""
+    "RELATIVE_INSTALL_ROOT"
+    ""
+    ${ARGN})
+
+  # RPATH handling.
+  # For the build tree, include the LLVM lib directory and the current
+  # directory for RPATH searching. For install, just the current directory
+  # (assumes that needed dependencies have been installed).
+  if(NOT APPLE AND NOT UNIX)
+    return()
+  endif()
+
+  set(_origin_prefix "\$ORIGIN")
+  if(APPLE)
+    set(_origin_prefix "@loader_path")
+  endif()
+  set_target_properties(${target} PROPERTIES
+    BUILD_WITH_INSTALL_RPATH OFF
+    BUILD_RPATH "${_origin_prefix}"
+    INSTALL_RPATH "${_origin_prefix}"
+  )
+
+  # For static builds, that is all that is needed: all dependencies will be in
+  # the one directory. For shared builds, then we also need to add the global
+  # lib directory. This will be absolute for the build tree and relative for
+  # install.
+  # When we have access to CMake >= 3.20, there is a helper to calculate this.
+  if(BUILD_SHARED_LIBS AND ARG_RELATIVE_INSTALL_ROOT)
+    get_filename_component(_real_lib_dir "${LLVM_LIBRARY_OUTPUT_INTDIR}" REALPATH)
+    set_property(TARGET ${target} APPEND PROPERTY
+      BUILD_RPATH "${_real_lib_dir}")
+    set_property(TARGET ${target} APPEND PROPERTY
+      INSTALL_RPATH "${_origin_prefix}/${ARG_RELATIVE_INSTALL_ROOT}/lib${LLVM_LIBDIR_SUFFIX}")
+  endif()
+endfunction()
+
+# Function: add_mlir_python_common_capi_library
+# Adds a shared library which embeds dependent CAPI libraries needed to link
+# all extensions.
+# Arguments:
+#   INSTALL_COMPONENT: Name of the install component. Typically same as the
+#     target name passed to add_mlir_python_modules().
+#   INSTALL_DESTINATION: Prefix into the install tree in which to install the
+#     library.
+#   OUTPUT_DIRECTORY: Full path in the build tree in which to create the
+#     library. Typically, this will be the common _mlir_libs directory where
+#     all extensions are emitted.
+#   RELATIVE_INSTALL_ROOT: See mlir_python_setup_extension_rpath().
+#   DECLARED_SOURCES: Source groups from which to discover dependent
+#     EMBED_CAPI_LINK_LIBS.
+#   EMBED_LIBS: Additional libraries to embed (must be built with OBJECTS and
+#     have an "obj.${name}" object library associated).
+function(add_mlir_python_common_capi_library name)
+  cmake_parse_arguments(ARG
+    ""
+    "INSTALL_COMPONENT;INSTALL_DESTINATION;OUTPUT_DIRECTORY;RELATIVE_INSTALL_ROOT"
+    "DECLARED_SOURCES;EMBED_LIBS"
+    ${ARGN})
+  # TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
+  # once ready.
+
+  # Collect all explicit and transitive embed libs.
+  set(_embed_libs ${ARG_EMBED_LIBS})
+  _flatten_mlir_python_targets(_all_source_targets ${ARG_DECLARED_SOURCES})
+  foreach(t ${_all_source_targets})
+    get_target_property(_local_embed_libs ${t} PYTHON_EMBED_CAPI_LINK_LIBS)
+    if(_local_embed_libs)
+      list(APPEND _embed_libs ${_local_embed_libs})
+    endif()
+  endforeach()
+  list(REMOVE_DUPLICATES _embed_libs)
+
+  foreach(lib ${_embed_libs})
+    if(XCODE)
+      # Xcode doesn't support object libraries, so we have to trick it into
+      # linking the static libraries instead.
+      list(APPEND _deps "-force_load" ${lib})
+    else()
+      list(APPEND _objects $<TARGET_OBJECTS:obj.${lib}>)
+    endif()
+    # Accumulate transitive deps of each exported lib into _DEPS.
+    list(APPEND _deps $<TARGET_PROPERTY:${lib},LINK_LIBRARIES>)
+  endforeach()
+
+  add_mlir_library(${name}
+    PARTIAL_SOURCES_INTENDED
+    SHARED
+    DISABLE_INSTALL
+    ${_objects}
+    EXCLUDE_FROM_LIBMLIR
+    LINK_LIBS
+    ${_deps}
+  )
+  if(MSVC)
+    set_property(TARGET ${name} PROPERTY WINDOWS_EXPORT_ALL_SYMBOLS ON)
+  endif()
+  set_target_properties(${name} PROPERTIES
+    LIBRARY_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
+    BINARY_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
+  )
+  mlir_python_setup_extension_rpath(${name}
+    RELATIVE_INSTALL_ROOT "${ARG_RELATIVE_INSTALL_ROOT}"
+  )
+  install(TARGETS ${name}
+    COMPONENT ${ARG_INSTALL_COMPONENT}
+    LIBRARY DESTINATION "${ARG_INSTALL_DESTINATION}"
+    RUNTIME DESTINATION "${ARG_INSTALL_DESTINATION}"
+  )
+
+endfunction()
+
+function(_flatten_mlir_python_targets output_var)
+  set(_flattened)
+  foreach(t ${ARGN})
+    get_target_property(_source_type ${t} PYTHON_SOURCES_TYPE)
+    get_target_property(_depends ${t} PYTHON_DEPENDS)
+    if(_source_type)
+      list(APPEND _flattened "${t}")
+      if(_depends)
+        _flatten_mlir_python_targets(_local_flattened ${_depends})
+        list(APPEND _flattened ${_local_flattened})
+      endif()
+    endif()
+  endforeach()
+  list(REMOVE_DUPLICATES _flattened)
+  set(${output_var} "${_flattened}" PARENT_SCOPE)
+endfunction()
+
 ################################################################################
 # Build python extension
 ################################################################################
 function(add_mlir_python_extension libname extname)
   cmake_parse_arguments(ARG
   ""
-  "INSTALL_DIR"
+  "INSTALL_COMPONENT;INSTALL_DIR;OUTPUT_DIRECTORY"
   "SOURCES;LINK_LIBS"
   ${ARGN})
   if (ARG_UNPARSED_ARGUMENTS)
@@ -14,6 +408,18 @@ function(add_mlir_python_extension libname extname)
     message(FATAL_ERROR " Missing SOURCES argument to add_mlir_python_extension(${libname}, ...")
   endif()
 
+  # Build-time RPath layouts require to be a directory one up from the
+  # binary root.
+  # TODO: Don't reference the LLVM_BINARY_DIR here: the invariant is that
+  # the output directory must be at the same level of the lib directory
+  # where libMLIR.so is installed. This is presently not optimal from a
+  # project separation perspective and a discussion on how to better
+  # segment MLIR libraries needs to happen.
+  # TODO: Remove this when downstreams are moved off of it.
+  if(NOT ARG_OUTPUT_DIRECTORY)
+    set(ARG_OUTPUT_DIRECTORY ${LLVM_BINARY_DIR}/python)
+  endif()
+
   # Normally on unix-like platforms, extensions are built as "MODULE" libraries
   # and do not explicitly link to the python shared object. This allows for
   # some greater deployment flexibility since the extension will bind to
@@ -67,14 +473,7 @@ function(add_mlir_python_extension libname extname)
   # Configure the output to match python expectations.
   set_target_properties(
     ${libname} PROPERTIES
-    # Build-time RPath layouts require to be a directory one up from the
-    # binary root.
-    # TODO: Don't reference the LLVM_BINARY_DIR here: the invariant is that
-    # the output directory must be at the same level of the lib directory
-    # where libMLIR.so is installed. This is presently not optimal from a
-    # project separation perspective and a discussion on how to better
-    # segment MLIR libraries needs to happen.
-    LIBRARY_OUTPUT_DIRECTORY ${LLVM_BINARY_DIR}/python
+    LIBRARY_OUTPUT_DIRECTORY ${ARG_OUTPUT_DIRECTORY}
     OUTPUT_NAME "${extname}"
     PREFIX "${PYTHON_MODULE_PREFIX}"
     SUFFIX "${PYTHON_MODULE_SUFFIX}${PYTHON_MODULE_EXTENSION}"
@@ -85,7 +484,7 @@ function(add_mlir_python_extension libname extname)
     # control where the .dll gets written.
     set_target_properties(
       ${libname} PROPERTIES
-      RUNTIME_OUTPUT_DIRECTORY ${LLVM_BINARY_DIR}/python
+      RUNTIME_OUTPUT_DIRECTORY ${ARG_OUTPUT_DIRECTORY}
     )
   endif()
 
@@ -111,53 +510,16 @@ function(add_mlir_python_extension libname extname)
       $<$<PLATFORM_ID:Linux>:LINKER:--exclude-libs,ALL>
   )
 
-  llvm_setup_rpath(${libname})
-
   ################################################################################
   # Install
   ################################################################################
   if (ARG_INSTALL_DIR)
     install(TARGETS ${libname}
-      COMPONENT ${libname}
+      COMPONENT ${ARG_INSTALL_COMPONENT}
       LIBRARY DESTINATION ${ARG_INSTALL_DIR}
       ARCHIVE DESTINATION ${ARG_INSTALL_DIR}
       # NOTE: Even on DLL-platforms, extensions go in the lib directory tree.
       RUNTIME DESTINATION ${ARG_INSTALL_DIR}
     )
   endif()
-
-  if (NOT LLVM_ENABLE_IDE)
-    add_llvm_install_targets(
-      install-${libname}
-      DEPENDS ${libname}
-      COMPONENT ${libname})
-  endif()
-
 endfunction()
-
-function(add_mlir_dialect_python_bindings tblgen_target)
-  cmake_parse_arguments(ARG
-    ""
-    "TD_FILE;DIALECT_NAME"
-    "DEPENDS"
-    ${ARGN})
-
-  set(dialect_filename "_${ARG_DIALECT_NAME}_ops_gen.py")
-  set(LLVM_TARGET_DEFINITIONS ${ARG_TD_FILE})
-  mlir_tablegen("${dialect_filename}" -gen-python-op-bindings
-                -bind-dialect=${ARG_DIALECT_NAME})
-  add_public_tablegen_target(
-    ${tblgen_target})
-  if(ARG_DEPENDS)
-    add_dependencies(${tblgen_target} ${ARG_DEPENDS})
-  endif()
-
-  add_custom_command(
-    TARGET ${tblgen_target} POST_BUILD
-    COMMENT "Copying generated python source \"dialects/${dialect_filename}\""
-    BYPRODUCTS "${PROJECT_BINARY_DIR}/python/mlir/dialects/${dialect_filename}"
-    COMMAND "${CMAKE_COMMAND}" -E copy_if_
diff erent
-      "${CMAKE_CURRENT_BINARY_DIR}/${dialect_filename}"
-      "${PROJECT_BINARY_DIR}/python/mlir/dialects/${dialect_filename}")
-endfunction()
-

diff  --git a/mlir/docs/Bindings/Python.md b/mlir/docs/Bindings/Python.md
index d00c7ab616ccc..853d380013a66 100644
--- a/mlir/docs/Bindings/Python.md
+++ b/mlir/docs/Bindings/Python.md
@@ -60,13 +60,19 @@ python -m pip install -r mlir/python/requirements.txt
 # Now run `cmake`, `ninja`, et al.
 ```
 
-For interactive use, it is sufficient to add the `python` directory in your
-`build/` directory to the `PYTHONPATH`. Typically:
+For interactive use, it is sufficient to add the
+`tools/mlir/python_packages/mlir_core/` directory in your `build/` directory to
+the `PYTHONPATH`. Typically:
 
 ```shell
-export PYTHONPATH=$(cd build && pwd)/python
+export PYTHONPATH=$(cd build && pwd)/tools/mlir/python_packages/mlir_core
 ```
 
+Note that if you have installed (i.e. via `ninja install`, et al), then
+python packages for all enabled projects will be in your install tree under
+`python_packages/` (i.e. `python_packages/mlir_core`). Official distributions
+are built with a more specialized setup.
+
 ## Design
 
 ### Use cases

diff  --git a/mlir/lib/Bindings/CMakeLists.txt b/mlir/lib/Bindings/CMakeLists.txt
deleted file mode 100644
index 41227a414a10e..0000000000000
--- a/mlir/lib/Bindings/CMakeLists.txt
+++ /dev/null
@@ -1,3 +0,0 @@
-if(MLIR_ENABLE_BINDINGS_PYTHON)
-  add_subdirectory(Python)
-endif()

diff  --git a/mlir/lib/Bindings/Python/CMakeLists.txt b/mlir/lib/Bindings/Python/CMakeLists.txt
deleted file mode 100644
index 8a112b7f43525..0000000000000
--- a/mlir/lib/Bindings/Python/CMakeLists.txt
+++ /dev/null
@@ -1,137 +0,0 @@
-include(AddMLIRPython)
-add_custom_target(MLIRBindingsPythonExtension)
-
-################################################################################
-# All python extensions must link through one DSO which exports the CAPI, and
-# this must have a globally unique name amongst all embeddors of the python
-# library since it will effectively have global scope.
-#
-# The presence of this aggregate library is part of the long term plan, but its
-# use needs to be made more flexible.
-################################################################################
-
-set(public_api_libs
-  MLIRCAPIConversion
-  MLIRCAPIDebug
-  MLIRCEXECUTIONENGINE
-  MLIRCAPIIR
-  MLIRCAPIRegistration
-  MLIRCAPITransforms
-
-  # Dialects
-  MLIRCAPIAsync
-  MLIRCAPIGPU
-  MLIRCAPILinalg
-  MLIRCAPILLVM
-  MLIRCAPIShape
-  MLIRCAPISparseTensor
-  MLIRCAPIStandard
-  MLIRCAPISCF
-  MLIRCAPITensor
-)
-
-foreach(lib ${public_api_libs})
-  if(XCODE)
-    # Xcode doesn't support object libraries, so we have to trick it into
-    # linking the static libraries instead.
-    list(APPEND _DEPS "-force_load" ${lib})
-  else()
-    list(APPEND _OBJECTS $<TARGET_OBJECTS:obj.${lib}>)
-  endif()
-  # Accumulate transitive deps of each exported lib into _DEPS.
-  list(APPEND _DEPS  $<TARGET_PROPERTY:${lib},LINK_LIBRARIES>)
-endforeach()
-
-add_mlir_library(MLIRPythonCAPI
-  PARTIAL_SOURCES_INTENDED
-  SHARED
-  ${_OBJECTS}
-  EXCLUDE_FROM_LIBMLIR
-  LINK_LIBS
-  ${_DEPS}
-)
-if(MSVC)
-  set_property(TARGET MLIRPythonCAPI PROPERTY WINDOWS_EXPORT_ALL_SYMBOLS ON)
-endif()
-
-################################################################################
-# Build core python extension
-################################################################################
-add_mlir_python_extension(MLIRCoreBindingsPythonExtension _mlir
-  INSTALL_DIR
-    python
-  SOURCES
-    DialectLinalg.cpp
-    DialectSparseTensor.cpp
-    MainModule.cpp
-    IRAffine.cpp
-    IRAttributes.cpp
-    IRCore.cpp
-    IRModule.cpp
-    IRTypes.cpp
-    PybindUtils.cpp
-    Pass.cpp
-    ExecutionEngine.cpp
-  LINK_LIBS PRIVATE
-    LLVMSupport
-    MLIRPythonCAPI
-)
-add_dependencies(MLIRBindingsPythonExtension MLIRCoreBindingsPythonExtension)
-
-add_subdirectory(Transforms)
-add_subdirectory(Conversions)
-
-add_mlir_python_extension(MLIRAllPassesRegistrationBindingsPythonExtension _mlirAllPassesRegistration
-  INSTALL_DIR
-    python
-  SOURCES
-    AllPassesRegistration.cpp
-  LINK_LIBS PRIVATE
-    LLVMSupport
-    MLIRPythonCAPI
-)
-add_dependencies(MLIRBindingsPythonExtension MLIRAllPassesRegistrationBindingsPythonExtension)
-
-add_mlir_python_extension(MLIRAsyncPassesBindingsPythonExtension _mlirAsyncPasses
-  INSTALL_DIR
-    python
-  SOURCES
-    AsyncPasses.cpp
-  LINK_LIBS PRIVATE
-    LLVMSupport
-    MLIRPythonCAPI
-)
-add_dependencies(MLIRBindingsPythonExtension MLIRAsyncPassesBindingsPythonExtension)
-
-add_mlir_python_extension(MLIRSparseTensorPassesBindingsPythonExtension _mlirSparseTensorPasses
-  INSTALL_DIR
-    python
-  SOURCES
-    SparseTensorPasses.cpp
-  LINK_LIBS PRIVATE
-    LLVMSupport
-    MLIRPythonCAPI
-)
-add_dependencies(MLIRBindingsPythonExtension MLIRSparseTensorPassesBindingsPythonExtension)
-
-add_mlir_python_extension(MLIRGPUPassesBindingsPythonExtension _mlirGPUPasses
-  INSTALL_DIR
-    python
-  SOURCES
-    GPUPasses.cpp
-  LINK_LIBS PRIVATE
-    LLVMSupport
-    MLIRPythonCAPI
-)
-add_dependencies(MLIRBindingsPythonExtension MLIRGPUPassesBindingsPythonExtension)
-
-add_mlir_python_extension(MLIRLinalgPassesBindingsPythonExtension _mlirLinalgPasses
-  INSTALL_DIR
-    python
-  SOURCES
-    LinalgPasses.cpp
-  LINK_LIBS PRIVATE
-    LLVMSupport
-    MLIRPythonCAPI
-)
-add_dependencies(MLIRBindingsPythonExtension MLIRLinalgPassesBindingsPythonExtension)

diff  --git a/mlir/lib/Bindings/Python/Conversions/CMakeLists.txt b/mlir/lib/Bindings/Python/Conversions/CMakeLists.txt
deleted file mode 100644
index e39707d0c21de..0000000000000
--- a/mlir/lib/Bindings/Python/Conversions/CMakeLists.txt
+++ /dev/null
@@ -1,12 +0,0 @@
-################################################################################
-# Build python extension
-################################################################################
-
-add_mlir_python_extension(MLIRConversionsBindingsPythonExtension _mlirConversions
-  INSTALL_DIR
-    python
-  SOURCES
-  Conversions.cpp
-  LINK_LIBS PRIVATE
-    MLIRPythonCAPI
-)

diff  --git a/mlir/lib/Bindings/Python/Transforms/CMakeLists.txt b/mlir/lib/Bindings/Python/Transforms/CMakeLists.txt
deleted file mode 100644
index b33d1503b8859..0000000000000
--- a/mlir/lib/Bindings/Python/Transforms/CMakeLists.txt
+++ /dev/null
@@ -1,12 +0,0 @@
-################################################################################
-# Build python extension
-################################################################################
-
-add_mlir_python_extension(MLIRTransformsBindingsPythonExtension _mlirTransforms
-  INSTALL_DIR
-    python
-  SOURCES
-  Transforms.cpp
-  LINK_LIBS PRIVATE
-    MLIRPythonCAPI
-)

diff  --git a/mlir/lib/CMakeLists.txt b/mlir/lib/CMakeLists.txt
index f35c684222ccb..8b9a00e1dc918 100644
--- a/mlir/lib/CMakeLists.txt
+++ b/mlir/lib/CMakeLists.txt
@@ -2,7 +2,6 @@
 add_flag_if_supported("-Werror=global-constructors" WERROR_GLOBAL_CONSTRUCTOR)
 
 add_subdirectory(Analysis)
-add_subdirectory(Bindings)
 add_subdirectory(Conversion)
 add_subdirectory(Dialect)
 add_subdirectory(ExecutionEngine)

diff  --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 1b9705a0318e6..15b1813894a7b 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -1,49 +1,310 @@
+include(AddMLIRPython)
+
 ################################################################################
-# Copy python source tree.
+# Structural groupings.
 ################################################################################
 
-file(GLOB_RECURSE PY_SRC_FILES
-  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
-  "${CMAKE_CURRENT_SOURCE_DIR}/mlir/*.py")
+declare_mlir_python_sources(MLIRPythonSources)
+declare_mlir_python_sources(MLIRPythonSources.Dialects
+  ADD_TO_PARENT MLIRPythonSources)
+
+declare_mlir_python_sources(MLIRPythonTestSources)
+declare_mlir_python_sources(MLIRPythonTestSources.Dialects
+  ADD_TO_PARENT MLIRPythonTestSources)
 
-add_custom_target(MLIRBindingsPythonSources ALL
-  DEPENDS
-    ${PY_SRC_FILES}
+################################################################################
+# Pure python sources and generated code
+################################################################################
+
+declare_mlir_python_sources(MLIRPythonSources.Core
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  ADD_TO_PARENT MLIRPythonSources
+  SOURCES
+    _cext_loader.py
+    _dlloader.py
+    _mlir_libs/__init__.py
+    ir.py
+    passmanager.py
+    dialects/_ods_common.py
 )
 
-foreach(PY_SRC_FILE ${PY_SRC_FILES})
-  set(PY_DEST_FILE "${PROJECT_BINARY_DIR}/python/${PY_SRC_FILE}")
-  get_filename_component(PY_DEST_DIR "${PY_DEST_FILE}" DIRECTORY)
-  file(MAKE_DIRECTORY "${PY_DEST_DIR}")
-  add_custom_command(
-    TARGET MLIRBindingsPythonSources PRE_BUILD
-    COMMENT "Copying python source ${PY_SRC_FILE} -> ${PY_DEST_FILE}"
-    DEPENDS "${PY_SRC_FILE}"
-    BYPRODUCTS "${PY_DEST_FILE}"
-    COMMAND "${CMAKE_COMMAND}" -E create_symlink
-        "${CMAKE_CURRENT_SOURCE_DIR}/${PY_SRC_FILE}" "${PY_DEST_FILE}"
-  )
-endforeach()
+declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  ADD_TO_PARENT MLIRPythonSources
+  SOURCES
+    execution_engine.py
+  SOURCES_GLOB
+    runtime/*.py
+)
+
+declare_mlir_python_sources(MLIRPythonSources.Passes
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  ADD_TO_PARENT MLIRPythonSources
+  SOURCES_GLOB
+    all_passes_registration/*.py
+    conversions/*.py
+    transforms/*.py
+)
+
+################################################################################
+# Dialect bindings
+################################################################################
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/AsyncOps.td
+  SOURCES_GLOB dialects/async_dialect/*.py
+  DIALECT_NAME async_dialect)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/BuiltinOps.td
+  SOURCES
+    dialects/builtin.py
+    dialects/_builtin_ops_ext.py
+  DIALECT_NAME builtin)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/GPUOps.td
+  SOURCES_GLOB dialects/gpu/*.py
+  DIALECT_NAME gpu)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/LinalgOps.td
+  SOURCES
+    dialects/_linalg_ops_ext.py
+  SOURCES_GLOB
+    dialects/linalg/*.py
+  DIALECT_NAME linalg
+  DEPENDS LinalgOdsGen)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/MathOps.td
+  SOURCES dialects/math.py
+  DIALECT_NAME math)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/MemRefOps.td
+  SOURCES dialects/memref.py
+  DIALECT_NAME memref)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonTestSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/PythonTest.td
+  SOURCES dialects/python_test.py
+  DIALECT_NAME python_test)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/ShapeOps.td
+  SOURCES dialects/shape.py
+  DIALECT_NAME shape)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  SOURCES dialects/sparse_tensor.py
+  DIALECT_NAME sparse_tensor)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/StandardOps.td
+  SOURCES dialects/std.py
+  DIALECT_NAME std)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/TensorOps.td
+  SOURCES dialects/tensor.py
+  DIALECT_NAME tensor)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/TosaOps.td
+  SOURCES dialects/tosa.py
+  DIALECT_NAME tosa)
+
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/VectorOps.td
+  SOURCES dialects/vector.py
+  DIALECT_NAME vector)
+
+################################################################################
+# Python extensions.
+# The sources for these are all in lib/Bindings/Python, but since they have to
+# be rebuilt for each package and integrate with the source setup here, we
+# just reference them here instead of having ordered, cross package target
+# dependencies.
+################################################################################
+
+set(PYTHON_SOURCE_DIR "${MLIR_SOURCE_DIR}/lib/Bindings/Python")
+declare_mlir_python_extension(MLIRPythonExtension.Core
+  MODULE_NAME _mlir
+  ADD_TO_PARENT MLIRPythonSources.Core
+  SOURCES
+    ${PYTHON_SOURCE_DIR}/DialectLinalg.cpp  # TODO: Break this out.
+    ${PYTHON_SOURCE_DIR}/DialectSparseTensor.cpp  # TODO: Break this out.
+    ${PYTHON_SOURCE_DIR}/MainModule.cpp
+    ${PYTHON_SOURCE_DIR}/IRAffine.cpp
+    ${PYTHON_SOURCE_DIR}/IRAttributes.cpp
+    ${PYTHON_SOURCE_DIR}/IRCore.cpp
+    ${PYTHON_SOURCE_DIR}/IRModule.cpp
+    ${PYTHON_SOURCE_DIR}/IRTypes.cpp
+    ${PYTHON_SOURCE_DIR}/PybindUtils.cpp
+    ${PYTHON_SOURCE_DIR}/Pass.cpp
+    ${PYTHON_SOURCE_DIR}/ExecutionEngine.cpp # TODO: Break this out.
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIDebug
+    MLIRCAPIIR
+    MLIRCAPIRegistration  # TODO: See about dis-aggregating
+
+    # Dialects
+    MLIRCAPILinalg  # TODO: Remove when above is removed.
+    MLIRCAPISparseTensor  # TODO: Remove when above is removed.
+    MLIRCAPIStandard
+
+    # Execution engine (remove once disaggregated).
+    MLIRCEXECUTIONENGINE
+)
+
+declare_mlir_python_extension(MLIRPythonExtension.AllPassesRegistration
+  MODULE_NAME _mlirAllPassesRegistration
+  SOURCES
+    ${PYTHON_SOURCE_DIR}/AllPassesRegistration.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIConversion
+    MLIRCAPITransforms
+)
+
+declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
+  MODULE_NAME _mlirAsyncPasses
+  ADD_TO_PARENT MLIRPythonSources.Dialects.async_dialect
+  SOURCES
+    ${PYTHON_SOURCE_DIR}/AsyncPasses.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIAsync
+)
+
+declare_mlir_python_extension(MLIRPythonExtension.Conversions
+  MODULE_NAME _mlirConversions
+  ADD_TO_PARENT MLIRPythonSources.Passes
+  SOURCES
+    ${PYTHON_SOURCE_DIR}/Conversions/Conversions.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+  MLIRCAPIConversion
+)
+
+declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
+  MODULE_NAME _mlirGPUPasses
+  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
+  SOURCES
+    ${PYTHON_SOURCE_DIR}/GPUPasses.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIGPU
+)
+
+declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
+  MODULE_NAME _mlirLinalgPasses
+  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
+  SOURCES
+    ${PYTHON_SOURCE_DIR}/LinalgPasses.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPILinalg
+)
 
-# Note that we copy from the source tree just like for headers because
-# it will not be polluted with py_cache runtime artifacts (from testing and
-# such).
-install(
-  DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/mlir
-  DESTINATION python
-  COMPONENT MLIRBindingsPythonSources
-  FILES_MATCHING PATTERN "*.py"
+declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
+  MODULE_NAME _mlirSparseTensorPasses
+  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
+  SOURCES
+    ${PYTHON_SOURCE_DIR}/SparseTensorPasses.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPISparseTensor
 )
 
-if (NOT LLVM_ENABLE_IDE)
-  add_llvm_install_targets(
-    install-MLIRBindingsPythonSources
-    DEPENDS MLIRBindingsPythonSources
-    COMPONENT MLIRBindingsPythonSources)
-endif()
+declare_mlir_python_extension(MLIRPythonExtension.Transforms
+  MODULE_NAME _mlirTransforms
+  ADD_TO_PARENT MLIRPythonSources.Passes
+  SOURCES
+    ${PYTHON_SOURCE_DIR}/Transforms/Transforms.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITransforms
+)
 
 ################################################################################
-# Generated sources.
+# Common CAPI dependency DSO.
+# All python extensions must link through one DSO which exports the CAPI, and
+# this must have a globally unique name amongst all embeddors of the python
+# library since it will effectively have global scope.
+#
+# The presence of this aggregate library is part of the long term plan, but its
+# use needs to be made more flexible.
+#
+# TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
+# once ready.
 ################################################################################
 
-add_subdirectory(mlir/dialects)
+add_mlir_python_common_capi_library(MLIRPythonCAPI
+  INSTALL_COMPONENT MLIRPythonModules
+  INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
+  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
+  RELATIVE_INSTALL_ROOT "../../../.."
+  DECLARED_SOURCES
+    MLIRPythonSources
+    MLIRPythonExtension.AllPassesRegistration
+)
+
+################################################################################
+# The fully assembled package of modules.
+# This must come last.
+################################################################################
+
+add_mlir_python_modules(MLIRPythonModules
+  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
+  INSTALL_PREFIX "python_packages/mlir_core/mlir"
+  DECLARED_SOURCES
+    MLIRPythonSources
+    MLIRPythonExtension.AllPassesRegistration
+  COMMON_CAPI_LINK_LIBS
+    MLIRPythonCAPI
+  )
+
+
+add_mlir_python_modules(MLIRPythonTestModules
+  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_test/mlir"
+  INSTALL_PREFIX "python_packages/mlir_test/mlir"
+  DECLARED_SOURCES
+    MLIRPythonTestSources
+  )

diff  --git a/mlir/python/mlir/_cext_loader.py b/mlir/python/mlir/_cext_loader.py
index c691fccb47375..5f2de7f006960 100644
--- a/mlir/python/mlir/_cext_loader.py
+++ b/mlir/python/mlir/_cext_loader.py
@@ -3,28 +3,27 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 """Common module for looking up and manipulating C-Extensions."""
 
-# 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.
+# The normal layout is to have a nested _mlir_libs package that contains
+# all native libraries and extensions. If that exists, use it, but also fallback
+# to old behavior where extensions were at the top level as loose libraries.
+# TODO: Remove the fallback once downstreams adapt.
 try:
-  import _mlir_libs
+  from ._mlir_libs import *
+  # TODO: Remove these aliases once everything migrates
+  _preload_dependency = preload_dependency
+  _load_extension = load_extension
 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
+  # TODO: Remove _dlloader once all consolidated on the _mlir_libs approach.
+  from ._dlloader import preload_dependency
 
-  def _load_extension(name):
+  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("MLIRPythonCAPI")
+preload_dependency("MLIRPythonCAPI")
 
 # Expose the corresponding C-Extension module with a well-known name at this
 # top-level module. This allows relative imports like the following to
@@ -32,7 +31,7 @@ def _load_extension(name):
 #   from .._cext_loader 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.
-_cext = _load_extension("_mlir")
+_cext = load_extension("_mlir")
 
 
 def _reexport_cext(cext_module_name, target_module_name):

diff  --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py
new file mode 100644
index 0000000000000..54ed5e8ef975a
--- /dev/null
+++ b/mlir/python/mlir/_mlir_libs/__init__.py
@@ -0,0 +1,21 @@
+# Licensed 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 importlib
+import os
+
+__all__ = [
+  "load_extension",
+  "preload_dependency",
+]
+
+_this_dir = os.path.dirname(__file__)
+
+def load_extension(name):
+  return importlib.import_module(f".{name}", __package__)
+
+
+def preload_dependency(public_name):
+  # TODO: Implement this hook to pre-load DLLs with ctypes on Windows.
+  pass

diff  --git a/mlir/python/mlir/dialects/CMakeLists.txt b/mlir/python/mlir/dialects/CMakeLists.txt
deleted file mode 100644
index 3c0434475f304..0000000000000
--- a/mlir/python/mlir/dialects/CMakeLists.txt
+++ /dev/null
@@ -1,86 +0,0 @@
-include(AddMLIRPython)
-
-################################################################################
-# Generate dialect-specific bindings.
-################################################################################
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonAsyncOps
-  TD_FILE AsyncOps.td
-  DIALECT_NAME async_dialect)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonAsyncOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonBuiltinOps
-  TD_FILE BuiltinOps.td
-  DIALECT_NAME builtin)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonBuiltinOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonGPUOps
-  TD_FILE GPUOps.td
-  DIALECT_NAME gpu)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonGPUOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonLinalgOps
-  TD_FILE LinalgOps.td
-  DIALECT_NAME linalg
-  DEPENDS LinalgOdsGen)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonLinalgOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonMathOps
-  TD_FILE MathOps.td
-  DIALECT_NAME math)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonMathOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonMemRefOps
-  TD_FILE MemRefOps.td
-  DIALECT_NAME memref)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonMemRefOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonShapeOps
-  TD_FILE ShapeOps.td
-  DIALECT_NAME shape)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonShapeOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonStandardOps
-  TD_FILE StandardOps.td
-  DIALECT_NAME std)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonStandardOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonTensorOps
-  TD_FILE TensorOps.td
-  DIALECT_NAME tensor)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonTensorOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonTosaOps
-  TD_FILE TosaOps.td
-  DIALECT_NAME tosa)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonTosaOps)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonVectorOps
-  TD_FILE VectorOps.td
-  DIALECT_NAME vector)
-add_dependencies(MLIRBindingsPythonSources MLIRBindingsPythonVectorOps)
-
-################################################################################
-# Installation.
-################################################################################
-
-# Dialect sources are generated. Install separately.
-# Note that __pycache__ directories may have been left by tests and other
-# executions. And __init__.py is handled as a regular source file.
-# TODO: Eliminate this glob install, instead adding INSTALL_COMPONENT to
-# add_mlir_dialect_python_bindings and installing the precise file there.
-install(
-  DIRECTORY ${PROJECT_BINARY_DIR}/python/mlir/dialects
-  DESTINATION python/mlir
-  COMPONENT MLIRBindingsPythonDialects
-  FILES_MATCHING PATTERN "_*_gen.py"
-  PATTERN "__pycache__" EXCLUDE
-  PATTERN "__init__.py" EXCLUDE
-)
-
-if (NOT LLVM_ENABLE_IDE)
-  add_llvm_install_targets(
-    install-MLIRBindingsPythonDialects
-    DEPENDS MLIRBindingsPythonSources
-    COMPONENT MLIRBindingsPythonDialects)
-endif()

diff  --git a/mlir/python/mlir/dialects/PythonTest.td b/mlir/python/mlir/dialects/PythonTest.td
new file mode 100644
index 0000000000000..d3d49395ad45a
--- /dev/null
+++ b/mlir/python/mlir/dialects/PythonTest.td
@@ -0,0 +1,33 @@
+//===-- python_test_ops.td - Python test Op definitions ----*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef PYTHON_TEST_OPS
+#define PYTHON_TEST_OPS
+
+include "mlir/Bindings/Python/Attributes.td"
+include "mlir/IR/OpBase.td"
+
+def Python_Test_Dialect : Dialect {
+  let name = "python_test";
+  let cppNamespace = "PythonTest";
+}
+class TestOp<string mnemonic, list<OpTrait> traits = []>
+    : Op<Python_Test_Dialect, mnemonic, traits>;
+
+def AttributedOp : TestOp<"attributed_op"> {
+  let arguments = (ins I32Attr:$mandatory_i32,
+                   OptionalAttr<I32Attr>:$optional_i32,
+                   UnitAttr:$unit);
+}
+
+def PropertyOp : TestOp<"property_op"> {
+  let arguments = (ins I32Attr:$property,
+                   I32:$idx);
+}
+
+#endif // PYTHON_TEST_OPS

diff  --git a/mlir/python/mlir/dialects/_builtin_ops_ext.py b/mlir/python/mlir/dialects/_builtin_ops_ext.py
index 6598efe3e0828..99783d8333753 100644
--- a/mlir/python/mlir/dialects/_builtin_ops_ext.py
+++ b/mlir/python/mlir/dialects/_builtin_ops_ext.py
@@ -2,11 +2,14 @@
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from typing import Optional, Sequence
+try:
+  from typing import Optional, Sequence
 
-import inspect
+  import inspect
 
-from ..ir import *
+  from ..ir import *
+except ImportError as e:
+  raise RuntimeError("Error loading imports from extension module") from e
 
 
 class ModuleOp:

diff  --git a/mlir/python/mlir/dialects/_linalg_ops_ext.py b/mlir/python/mlir/dialects/_linalg_ops_ext.py
index bce4e08ae3a06..656992cac26d7 100644
--- a/mlir/python/mlir/dialects/_linalg_ops_ext.py
+++ b/mlir/python/mlir/dialects/_linalg_ops_ext.py
@@ -2,12 +2,16 @@
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-from typing import Optional, Sequence, Union
-from ..ir import *
-from ._ods_common import get_default_loc_context
-# TODO: resolve name collision for Linalg functionality that is injected inside
-# the _mlir.dialects.linalg directly via pybind.
-from _mlir.dialects.linalg import fill_builtin_region
+try:
+  from typing import Optional, Sequence, Union
+  from ..ir import *
+  from ._ods_common import get_default_loc_context
+  # TODO: resolve name collision for Linalg functionality that is injected inside
+  # the _mlir.dialects.linalg directly via pybind.
+  from .._cext_loader import _cext
+  fill_builtin_region = _cext.dialects.linalg.fill_builtin_region
+except ImportError as e:
+  raise RuntimeError("Error loading imports from extension module") from e
 
 
 def isa(cls: Type, ty: Type):

diff  --git a/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py b/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py
index 3810df9dff74a..4568298d1be88 100644
--- a/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py
+++ b/mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py
@@ -10,7 +10,8 @@
 from mlir.dialects import math
 # TODO: resolve name collision for Linalg functionality that is injected inside
 # the _mlir.dialects.linalg directly via pybind.
-from _mlir.dialects.linalg import fill_builtin_region
+from ....._cext_loader import _cext
+fill_builtin_region = _cext.dialects.linalg.fill_builtin_region
 
 from .scalar_expr import *
 from .config import *

diff  --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt
index 153e0ab2cf53b..7a5e50c557316 100644
--- a/mlir/test/CMakeLists.txt
+++ b/mlir/test/CMakeLists.txt
@@ -1,10 +1,6 @@
 add_subdirectory(CAPI)
 add_subdirectory(lib)
 
-if(MLIR_ENABLE_BINDINGS_PYTHON)
-  add_subdirectory(python)
-endif()
-
 llvm_canonicalize_cmake_booleans(
   MLIR_ENABLE_BINDINGS_PYTHON
   LLVM_BUILD_EXAMPLES
@@ -124,11 +120,8 @@ endif()
 
 if(MLIR_ENABLE_BINDINGS_PYTHON)
   list(APPEND MLIR_TEST_DEPENDS
-    MLIRBindingsPythonExtension
-    MLIRBindingsPythonSources
-    MLIRBindingsPythonTestOps
-    MLIRTransformsBindingsPythonExtension
-    MLIRConversionsBindingsPythonExtension
+    MLIRPythonModules
+    MLIRPythonTestModules
   )
 endif()
 

diff  --git a/mlir/test/lit.cfg.py b/mlir/test/lit.cfg.py
index a362645bb75c8..a7811faa4be64 100644
--- a/mlir/test/lit.cfg.py
+++ b/mlir/test/lit.cfg.py
@@ -103,13 +103,8 @@
 # by copying/linking sources to build.
 if config.enable_bindings_python:
     llvm_config.with_environment('PYTHONPATH', [
-        # TODO: Don't reference the llvm_obj_root here: the invariant is that
-        # the python/ must be at the same level of the lib directory
-        # where libMLIR.so is installed. This is presently not optimal from a
-        # project separation perspective and a discussion on how to better
-        # segment MLIR libraries needs to happen. See also
-        # lib/Bindings/Python/CMakeLists.txt for where this is set up.
-        os.path.join(config.llvm_obj_root, 'python'),
+        os.path.join(config.mlir_obj_root, 'python_packages', 'mlir_core'),
+        os.path.join(config.mlir_obj_root, 'python_packages', 'mlir_test'),
     ], append_path=True)
 
 if config.enable_assertions:

diff  --git a/mlir/test/python/CMakeLists.txt b/mlir/test/python/CMakeLists.txt
deleted file mode 100644
index 19f5a79bcd75e..0000000000000
--- a/mlir/test/python/CMakeLists.txt
+++ /dev/null
@@ -1,5 +0,0 @@
-include(AddMLIRPython)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonTestOps
-  TD_FILE python_test_ops.td
-  DIALECT_NAME python_test)

diff  --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index 8bd6c9aef6283..c76da46d7a516 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -533,12 +533,12 @@ def testKnownOpView():
 
     # One of the custom ops should resolve to the default OpView.
     custom = module.body.operations[0]
-    # CHECK: <_mlir.ir.OpView object
+    # CHECK: OpView object
     print(repr(custom))
 
     # Check again to make sure negative caching works.
     custom = module.body.operations[0]
-    # CHECK: <_mlir.ir.OpView object
+    # CHECK: OpView object
     print(repr(custom))
 
 run(testKnownOpView)


        


More information about the Mlir-commits mailing list