Re-engineer MLIR python build support.

[stellaraccident]

Summary

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

[stephenneuendorffer]

if not installing into python/ anymore, can you document where things will get installed? I didn't quite parse it out of the code.

[stellaraccident]

Update docs.

[stellaraccident]

In D106520#2895534, @stephenneuendorffer wrote:

if not installing into python/ anymore, can you document where things will get installed? I didn't quite parse it out of the code.

Updated the Python.md documentation:

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)/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.

Side note: I really wish that MLIR's binary directory was not tools/mlir (and I wish it wasn't gated on the LLVM_BUILD_TOOLS option.

[mikeurbach]

I haven't reviewed yet, but it's great to see this. I'll dig in today, and I'll help integrate the changes into CIRCT. The CI builds of CIRCT use the github.com/circt/llvm repo, and we can stage this there for testing as needed.

[mikeurbach]

I've reviewed, and I think I'm understanding what's going on. I don't think I have enough CMake chops to provide serious suggestions/alternatives on the approach, so I'm just going to try it out with CIRCT and see how it goes. Assuming it works and we don't have to work around TypeID issues and whatnot, I'm satisfied!

[stellaraccident]

In D106520#2898313, @mikeurbach wrote:

I've reviewed, and I think I'm understanding what's going on. I don't think I have enough CMake chops to provide serious suggestions/alternatives on the approach, so I'm just going to try it out with CIRCT and see how it goes. Assuming it works and we don't have to work around TypeID issues and whatnot, I'm satisfied!

Thanks for the honesty. It is a wise path to not have such CMake chops (otherwise, you end up "me") :)

I have this working with one of our internal projects, modulo a couple of minor tweaks I'll push momentarily. If you'd like to wait until tomorrow, I'll have an npcomp branch staged with changes and you should be able to draft onto those. Probably easier than being the first to apply in a new project.

[stellaraccident]

A couple of minor fixes found while integrating a third party project.

[mikeurbach]

In D106520#2898374, @stellaraccident wrote:

If you'd like to wait until tomorrow, I'll have an npcomp branch staged with changes and you should be able to draft onto those. Probably easier than being the first to apply in a new project.

Sounds good, I will apply the latest patch to my local MLIR but I won't dig into the circt changes just yet.

[stellaraccident]

Add sources custom target to ALL.

[stellaraccident]

We have pre-patched npcomp (https://github.com/llvm/mlir-npcomp/pull/251) and the iree-llvm-sandbox (https://github.com/google/iree-llvm-sandbox) successfully based on this patch. Mike says that he has a pending patch for CIRCT. So I think this is ready to land. I'll plan to do so tomorrow.

[mikeurbach]

This has been working out of the box for the CIRCT project. No more TypeID suprises, etc. Thanks Stella, this is super nice!

mlir/cmake/modules/AddMLIR.cmake

 # In addition to everything that llvm_add_librar accepts, this
 # also has the following option:
 # EXCLUDE_FROM_LIBMLIR
 #   Don't include this library in libMLIR.so.  This option should be used
 #   for test libraries, executable-specific libraries, or rarely used libraries
 #   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})
   set(srcs)
   if(MSVC_IDE OR XCODE)
     # Add public headers
     file(RELATIVE_PATH lib_path
       ${MLIR_SOURCE_DIR}/lib/
foreach(lib ${ARG_LINK_LIBS})
     endif()
   endforeach()
 
   list(APPEND ARG_DEPENDS mlir-generic-headers)
   llvm_add_library(${name} ${LIBTYPE} ${ARG_UNPARSED_ARGUMENTS} ${srcs} DEPENDS ${ARG_DEPENDS} LINK_COMPONENTS ${ARG_LINK_COMPONENTS} LINK_LIBS ${ARG_LINK_LIBS})
 
   if(TARGET ${name})
     target_link_libraries(${name} INTERFACE ${LLVM_COMMON_LIBS})
+    if(NOT ARG_DISABLE_INSTALL)
       add_mlir_library_install(${name})
+    endif()
   else()
     # Add empty "phony" target
     add_custom_target(${name})
   endif()
   set_target_properties(${name} PROPERTIES FOLDER "MLIR libraries")
 endfunction(add_mlir_library)
 
 # Adds an MLIR library target for installation.

mlir/cmake/modules/AddMLIRPython.cmake

 ################################################################################
+# 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} 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)
     message(FATAL_ERROR " Unhandled arguments to add_mlir_python_extension(${libname}, ... : ${ARG_UNPARSED_ARGUMENTS}")
   endif()
   if ("${ARG_SOURCES}" STREQUAL "")
     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
   # symbols in the python interpreter on load. However, it also keeps the
   # linker from erroring on undefined symbols, leaving this to (usually obtuse)
   # runtime errors. Building in "SHARED" mode with an explicit link to the
   # python libraries allows us to build with the expectation of no undefined
   # symbols, which is better for development. Note that not all python
target_compile_options(${libname} PRIVATE
     $<$<CXX_COMPILER_ID:MSVC>:
       # Enable RTTI and exceptions.
       /EHsc /GR>
   )
 
   # 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}"
   )
 
   if(WIN32)
     # Need to also set the RUNTIME_OUTPUT_DIRECTORY on Windows in order to
     # 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()
 
   # pybind11 requires binding code to be compiled with -fvisibility=hidden
   # For static linkage, better code can be generated if the entire project
   # compiles that way, but that is not enforced here. Instead, include a linker
   # script that explicitly hides anything but the PyInit_* symbols, allowing gc
   # to take place.
function(add_mlir_python_extension libname extname)
 
   target_link_options(${libname}
     PRIVATE
       # On Linux, disable re-export of any static linked libraries that
       # came through.
       $<$<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_different
-      "${CMAKE_CURRENT_BINARY_DIR}/${dialect_filename}"
-      "${PROJECT_BINARY_DIR}/python/mlir/dialects/${dialect_filename}")
 endfunction()
-

mlir/lib/Bindings/CMakeLists.txt

-if(MLIR_ENABLE_BINDINGS_PYTHON)
-  add_subdirectory(Python)
-endif()

mlir/lib/Bindings/Python/CMakeLists.txt

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

mlir/lib/Bindings/Python/Conversions/CMakeLists.txt

-################################################################################
-# Build python extension
-################################################################################
-
-add_mlir_python_extension(MLIRConversionsBindingsPythonExtension _mlirConversions
-  INSTALL_DIR
-    python
-  SOURCES
-  Conversions.cpp
-  LINK_LIBS PRIVATE
-    MLIRPythonCAPI
-)

mlir/lib/Bindings/Python/Transforms/CMakeLists.txt

-################################################################################
-# Build python extension
-################################################################################
-
-add_mlir_python_extension(MLIRTransformsBindingsPythonExtension _mlirTransforms
-  INSTALL_DIR
-    python
-  SOURCES
-  Transforms.cpp
-  LINK_LIBS PRIVATE
-    MLIRPythonCAPI
-)

mlir/lib/CMakeLists.txt

 # Enable errors for any global constructors.
 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)
 add_subdirectory(IR)
 add_subdirectory(Interfaces)
 add_subdirectory(Parser)
 add_subdirectory(Pass)
 add_subdirectory(Reducer)
 add_subdirectory(Rewrite)
 add_subdirectory(Support)
 add_subdirectory(TableGen)
 add_subdirectory(Target)
 add_subdirectory(Tools)
 add_subdirectory(Transforms)
 add_subdirectory(Translation)

mlir/python/CMakeLists.txt

+include(AddMLIRPython)
+
+################################################################################
+# Structural groupings.
+################################################################################
+
+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)
+
+################################################################################
+# 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
+)
+
+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
+    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)
+
 ################################################################################
-# Copy python source tree.
+# 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.
 ################################################################################
 
-file(GLOB_RECURSE PY_SRC_FILES
-  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
-  "${CMAKE_CURRENT_SOURCE_DIR}/mlir/*.py")
+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
 
-add_custom_target(MLIRBindingsPythonSources ALL
-  DEPENDS
-    ${PY_SRC_FILES}
+    # Execution engine (remove once disaggregated).
+    MLIRCEXECUTIONENGINE
 )
 
-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_extension(MLIRPythonExtension.AllPassesRegistration
+  MODULE_NAME _mlirAllPassesRegistration
+  SOURCES
+    ${PYTHON_SOURCE_DIR}/AllPassesRegistration.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIConversion
+    MLIRCAPITransforms
 
-# 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.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
 )
 
-if (NOT LLVM_ENABLE_IDE)
-  add_llvm_install_targets(
-    install-MLIRBindingsPythonSources
-    DEPENDS MLIRBindingsPythonSources
-    COMPONENT MLIRBindingsPythonSources)
-endif()
+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
+)
+
+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
+)
+
+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
+)
+
+################################################################################
+# 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_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
+)
 
 ################################################################################
-# Generated sources.
+# The fully assembled package of modules.
+# This must come last.
 ################################################################################
 
-add_subdirectory(mlir/dialects)
+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
+  )

mlir/python/mlir/_cext_loader.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
 """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
 # function:
 #   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):
   """Re-exports a named sub-module of the C-Extension into another module.
 
   Typically:
     from ._cext_loader import _reexport_cext
     _reexport_cext("ir", __name__)

mlir/python/mlir/_mlir_libs/__init__.py

+# 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

mlir/python/mlir/dialects/CMakeLists.txt

-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()

mlir/python/mlir/dialects/PythonTest.td

+//===-- 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

mlir/python/mlir/dialects/_builtin_ops_ext.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
 
+try:
   from typing import Optional, Sequence
 
   import inspect
 
   from ..ir import *
+except ImportError as e:
+  raise RuntimeError("Error loading imports from extension module") from e
 
 
 class ModuleOp:
   """Specialization for the module op class."""
 
   def __init__(self, *, loc=None, ip=None):
     super().__init__(self.build_generic(results=[], operands=[], loc=loc,
                                         ip=ip))

mlir/python/mlir/dialects/_linalg_ops_ext.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
 
+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 _mlir.dialects.linalg import fill_builtin_region
+  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):
   try:
     cls(ty)
     return True
   except ValueError:
     return False

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.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
 
 from typing import Dict, Sequence
 
 from mlir.ir import *
 from mlir.dialects import linalg
 from mlir.dialects import std
 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 *
 import numpy as np
 
 __all__ = [
     "emit_generic_structured_op",
     "emit_named_structured_op",

mlir/test/CMakeLists.txt

 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
   MLIR_ENABLE_CUDA_CONVERSIONS
   MLIR_ENABLE_CUDA_RUNNER
   MLIR_ENABLE_ROCM_CONVERSIONS
   MLIR_ENABLE_ROCM_RUNNER
   MLIR_ENABLE_SPIRV_CPU_RUNNER

 if(MLIR_ENABLE_VULKAN_RUNNER)
   list(APPEND MLIR_TEST_DEPENDS
     mlir-vulkan-runner
   )
 endif()
 
 if(MLIR_ENABLE_BINDINGS_PYTHON)
   list(APPEND MLIR_TEST_DEPENDS
-    MLIRBindingsPythonExtension
-    MLIRBindingsPythonSources
-    MLIRBindingsPythonTestOps
-    MLIRTransformsBindingsPythonExtension
-    MLIRConversionsBindingsPythonExtension
+    MLIRPythonModules
+    MLIRPythonTestModules
   )
 endif()
 
 add_lit_testsuite(check-mlir "Running the MLIR regression tests"
   ${CMAKE_CURRENT_BINARY_DIR}
   DEPENDS ${MLIR_TEST_DEPENDS}
   )
 set_target_properties(check-mlir PROPERTIES FOLDER "Tests")
 
 add_lit_testsuites(MLIR ${CMAKE_CURRENT_SOURCE_DIR}
   DEPENDS ${MLIR_TEST_DEPENDS}
 )

mlir/test/lit.cfg.py

if config.target_triple:
     config.available_features.add('default_triple')
 
 # Add the python path for both the source and binary tree.
 # Note that presently, the python sources come from the source tree and the
 # binaries come from the build tree. This should be unified to the build tree
 # 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:
     config.available_features.add('asserts')
 else:
     config.available_features.add('noasserts')

mlir/test/python/CMakeLists.txt

-include(AddMLIRPython)
-
-add_mlir_dialect_python_bindings(MLIRBindingsPythonTestOps
-  TD_FILE python_test_ops.td
-  DIALECT_NAME python_test)

mlir/test/python/ir/operation.py

with Context(), Location.unknown():
     addf = module.body.operations[2]
     # CHECK: <mlir.dialects._std_ops_gen._AddFOp object
     print(repr(addf))
     # CHECK: "custom.f32"()
     print(addf.lhs)
 
     # 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)
 
 
 # CHECK-LABEL: TEST: testSingleResultProperty
 def testSingleResultProperty():
   with Context(), Location.unknown():