[Mlir-commits] [mlir] [MLIR][Python] enable precise registration (PR #160742)

Maksim Levental llvmlistbot at llvm.org
Sat Sep 27 01:07:48 PDT 2025


https://github.com/makslevental updated https://github.com/llvm/llvm-project/pull/160742

>From 17d76f802f323ce0b909064580921021724e007d Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Thu, 25 Sep 2025 09:40:53 -0700
Subject: [PATCH 1/4] [MLIR][Python] enable precise registration

---
 mlir/cmake/modules/AddMLIRPython.cmake  | 14 +++++++++++---
 mlir/python/CMakeLists.txt              | 24 +++++++++++++++++++++---
 mlir/python/mlir/_mlir_libs/_capi.py.in |  8 ++++++++
 mlir/test/python/ir/capi.py             |  6 ++++++
 4 files changed, 46 insertions(+), 6 deletions(-)
 create mode 100644 mlir/python/mlir/_mlir_libs/_capi.py.in
 create mode 100644 mlir/test/python/ir/capi.py

diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index 208cbdd1dd535..d8b6d493f985c 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -23,11 +23,14 @@
 #     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.
+#   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_sources name)
   cmake_parse_arguments(ARG
     ""
     "ROOT_DIR;ADD_TO_PARENT"
-    "SOURCES;SOURCES_GLOB"
+    "SOURCES;SOURCES_GLOB;EMBED_CAPI_LINK_LIBS"
     ${ARGN})
 
   if(NOT ARG_ROOT_DIR)
@@ -53,9 +56,10 @@ function(declare_mlir_python_sources name)
   set_target_properties(${name} PROPERTIES
     # Yes: Leading-lowercase property names are load bearing and the recommended
     # way to do this: https://gitlab.kitware.com/cmake/cmake/-/issues/19261
-    EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_DEPENDS"
+    EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_DEPENDS;mlir_python_EMBED_CAPI_LINK_LIBS"
     mlir_python_SOURCES_TYPE pure
     mlir_python_DEPENDS ""
+    mlir_python_EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
   )
 
   # Use the interface include directories and sources on the target to carry the
@@ -374,6 +378,9 @@ endfunction()
 #     This file is where the *EnumAttrs are defined, not where the *Enums are defined.
 #     **WARNING**: This arg will shortly be removed when the just-below TODO is satisfied. Use at your
 #     risk.
+#   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.
 #
 # TODO: Right now `TD_FILE` can't be the actual dialect tablegen file, since we
 #       use its path to determine where to place the generated python file. If
@@ -383,7 +390,7 @@ function(declare_mlir_dialect_python_bindings)
   cmake_parse_arguments(ARG
     "GEN_ENUM_BINDINGS"
     "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME"
-    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE"
+    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE;EMBED_CAPI_LINK_LIBS"
     ${ARGN})
   # Sources.
   set(_dialect_target "${ARG_ADD_TO_PARENT}.${ARG_DIALECT_NAME}")
@@ -424,6 +431,7 @@ function(declare_mlir_dialect_python_bindings)
       ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
       ADD_TO_PARENT "${_dialect_target}"
       SOURCES ${_sources}
+      EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
     )
   endif()
 endfunction()
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 9f5246de6bda0..368f344841ed1 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -326,7 +326,10 @@ declare_mlir_dialect_python_bindings(
   SOURCES
     dialects/arith.py
   DIALECT_NAME arith
-  GEN_ENUM_BINDINGS)
+  GEN_ENUM_BINDINGS
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIArith
+  )
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -846,8 +849,20 @@ endif()
 # once ready.
 ################################################################################
 
+set(MLIR_PYTHON_CAPI_DYLIB_NAME MLIRPythonCAPI)
+configure_file(
+  "${CMAKE_CURRENT_LIST_DIR}/mlir/_mlir_libs/_capi.py.in"
+  "${CMAKE_CURRENT_BINARY_DIR}/_mlir_libs/_capi.py"
+  @ONLY
+)
+declare_mlir_python_sources(
+  MLIRPythonCAPICTypesBinding
+  ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
+  SOURCES _mlir_libs/_capi.py
+)
+
 set(MLIRPythonModules_ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}")
-add_mlir_python_common_capi_library(MLIRPythonCAPI
+add_mlir_python_common_capi_library(${MLIR_PYTHON_CAPI_DYLIB_NAME}
   INSTALL_COMPONENT MLIRPythonModules
   INSTALL_DESTINATION "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs"
   OUTPUT_DIRECTORY "${MLIRPythonModules_ROOT_PREFIX}/_mlir_libs"
@@ -963,7 +978,10 @@ endif()
 # This must come last.
 ################################################################################
 
-set(_declared_sources MLIRPythonSources MLIRPythonExtension.RegisterEverything)
+set(_declared_sources
+  MLIRPythonSources
+  MLIRPythonExtension.RegisterEverything
+  MLIRPythonCAPICTypesBinding)
 if(NOT CMAKE_CROSSCOMPILING)
   list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen)
 endif()
diff --git a/mlir/python/mlir/_mlir_libs/_capi.py.in b/mlir/python/mlir/_mlir_libs/_capi.py.in
new file mode 100644
index 0000000000000..9568845e67de9
--- /dev/null
+++ b/mlir/python/mlir/_mlir_libs/_capi.py.in
@@ -0,0 +1,8 @@
+#  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
+
+import ctypes
+from pathlib import Path
+
+_capi = ctypes.CDLL(str(Path(__file__).parent / "@CMAKE_SHARED_LIBRARY_PREFIX@@MLIR_PYTHON_CAPI_DYLIB_NAME@@CMAKE_SHARED_LIBRARY_SUFFIX@"))
\ No newline at end of file
diff --git a/mlir/test/python/ir/capi.py b/mlir/test/python/ir/capi.py
new file mode 100644
index 0000000000000..d60fbd820f91e
--- /dev/null
+++ b/mlir/test/python/ir/capi.py
@@ -0,0 +1,6 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir._mlir_libs._capi import _capi
+
+print("success")
+# CHECK: success
\ No newline at end of file

>From b02964f047538eca9f1b1a10ee3e6ad6cf20fc2b Mon Sep 17 00:00:00 2001
From: makslevental <maksim.levental at gmail.com>
Date: Thu, 25 Sep 2025 14:59:58 -0700
Subject: [PATCH 2/4] [MLIR][Python] enable precise registration

---
 mlir/cmake/modules/AddMLIRPython.cmake        |   6 +-
 mlir/include/mlir-c/Bindings/Python/Interop.h |   9 +
 mlir/include/mlir-c/Dialect/Affine.h          |  36 +++
 mlir/include/mlir-c/Dialect/Bufferization.h   |  36 +++
 mlir/include/mlir-c/Dialect/Builtin.h         |  33 +++
 mlir/include/mlir-c/Dialect/Complex.h         |  33 +++
 mlir/include/mlir-c/Dialect/GPU.h             |   3 +
 mlir/include/mlir-c/Dialect/Linalg.h          |   3 +
 mlir/include/mlir-c/Dialect/MemRef.h          |   3 +
 mlir/include/mlir-c/Dialect/NVGPU.h           |   3 +
 mlir/include/mlir-c/Dialect/PDL.h             |   3 +
 mlir/include/mlir-c/Dialect/SMT.h             |   3 +
 mlir/include/mlir-c/Dialect/SparseTensor.h    |   3 +
 mlir/include/mlir-c/Dialect/Tensor.h          |   3 +
 mlir/include/mlir-c/Dialect/Tosa.h            |  33 +++
 mlir/include/mlir-c/Dialect/Transform.h       |  13 +
 mlir/include/mlir-c/Dialect/UB.h              |  33 +++
 mlir/include/mlir-c/Dialect/Vector.h          |   3 +
 mlir/include/mlir-c/IR.h                      |  18 +-
 mlir/lib/Bindings/Python/IRCore.cpp           |  34 ++-
 mlir/lib/CAPI/Dialect/Affine.cpp              |  14 +
 mlir/lib/CAPI/Dialect/Bufferization.cpp       |  14 +
 mlir/lib/CAPI/Dialect/Builtin.cpp             |  13 +
 mlir/lib/CAPI/Dialect/CMakeLists.txt          | 144 ++++++++++
 mlir/lib/CAPI/Dialect/Complex.cpp             |  14 +
 mlir/lib/CAPI/Dialect/Tosa.cpp                |  13 +
 mlir/lib/CAPI/Dialect/Transform.cpp           |  15 +
 mlir/lib/CAPI/Dialect/TransformAffine.cpp     |  16 ++
 .../CAPI/Dialect/TransformBufferization.cpp   |  16 ++
 mlir/lib/CAPI/Dialect/TransformGPU.cpp        |  15 +
 .../lib/CAPI/Dialect/TransformInterpreter.cpp |   3 -
 mlir/lib/CAPI/Dialect/TransformLinalg.cpp     |  15 +
 mlir/lib/CAPI/Dialect/TransformMemRef.cpp     |  15 +
 mlir/lib/CAPI/Dialect/TransformNVGPU.cpp      |  15 +
 mlir/lib/CAPI/Dialect/TransformPDL.cpp        |  15 +
 mlir/lib/CAPI/Dialect/TransformSMT.cpp        |  15 +
 .../CAPI/Dialect/TransformSparseTensor.cpp    |  16 ++
 mlir/lib/CAPI/Dialect/TransformTensor.cpp     |  15 +
 mlir/lib/CAPI/Dialect/TransformVector.cpp     |  15 +
 mlir/lib/CAPI/Dialect/UB.cpp                  |  13 +
 mlir/lib/CAPI/IR/IR.cpp                       |  12 +
 .../lib/Dialect/Transform/IR/TransformOps.cpp |   3 +-
 mlir/python/CMakeLists.txt                    | 265 +++++++++++++-----
 mlir/python/mlir/_mlir_libs/_capi.py.in       |  44 ++-
 mlir/python/mlir/dialects/QuantOps.td         |  14 +
 mlir/python/mlir/dialects/quant.py            |   1 +
 .../mlir/dialects/transform/bufferization.py  |   9 +
 mlir/python/mlir/dialects/transform/debug.py  |   9 +
 mlir/python/mlir/dialects/transform/gpu.py    |  10 +
 mlir/python/mlir/dialects/transform/loop.py   |   9 +
 mlir/python/mlir/dialects/transform/memref.py |   9 +
 mlir/python/mlir/dialects/transform/nvgpu.py  |   9 +
 mlir/python/mlir/dialects/transform/pdl.py    |  10 +
 mlir/python/mlir/dialects/transform/smt.py    |  10 +
 .../mlir/dialects/transform/sparse_tensor.py  |   8 +
 .../mlir/dialects/transform/structured.py     |   9 +
 mlir/python/mlir/dialects/transform/tensor.py |   9 +
 mlir/python/mlir/dialects/transform/tune.py   |   9 +
 mlir/python/mlir/dialects/transform/vector.py |   8 +
 mlir/test/python/dialects/gpu/dialect.py      |   1 +
 mlir/test/python/dialects/irdl.py             |   1 +
 mlir/test/python/dialects/memref.py           |   1 +
 mlir/test/python/dialects/transform.py        |   1 +
 .../dialects/transform_bufferization_ext.py   |   1 +
 .../python/dialects/transform_debug_ext.py    |   1 +
 mlir/test/python/dialects/transform_extras.py |   3 +-
 .../test/python/dialects/transform_gpu_ext.py |   1 +
 .../python/dialects/transform_loop_ext.py     |   1 +
 .../python/dialects/transform_memref_ext.py   |   1 +
 .../python/dialects/transform_nvgpu_ext.py    |   1 +
 .../test/python/dialects/transform_smt_ext.py |   1 +
 .../dialects/transform_sparse_tensor_ext.py   |   1 +
 .../dialects/transform_structured_ext.py      |   1 +
 .../python/dialects/transform_tensor_ext.py   |   1 +
 .../python/dialects/transform_tune_ext.py     |   1 +
 .../python/dialects/transform_vector_ext.py   |   1 +
 mlir/test/python/integration/dialects/pdl.py  |   2 +-
 .../python/integration/dialects/transform.py  |   1 +
 mlir/test/python/ir/builtin_types.py          |   2 +-
 mlir/test/python/ir/capi.py                   |  39 ++-
 mlir/test/python/ir/diagnostic_handler.py     |   1 +
 mlir/test/python/ir/dialects.py               |   1 +
 mlir/test/python/ir/exception.py              |   1 +
 mlir/test/python/ir/insertion_point.py        |   1 +
 mlir/test/python/ir/module.py                 |   1 +
 mlir/test/python/ir/operation.py              |   2 +-
 mlir/test/python/ir/symbol_table.py           |   1 +
 mlir/test/python/multithreaded_tests.py       |   2 +-
 mlir/test/python/pass_manager.py              |   1 +
 mlir/test/python/python_pass.py               |   2 +-
 mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp |   9 +-
 91 files changed, 1168 insertions(+), 96 deletions(-)
 create mode 100644 mlir/include/mlir-c/Dialect/Affine.h
 create mode 100644 mlir/include/mlir-c/Dialect/Bufferization.h
 create mode 100644 mlir/include/mlir-c/Dialect/Builtin.h
 create mode 100644 mlir/include/mlir-c/Dialect/Complex.h
 create mode 100644 mlir/include/mlir-c/Dialect/Tosa.h
 create mode 100644 mlir/include/mlir-c/Dialect/UB.h
 create mode 100644 mlir/lib/CAPI/Dialect/Affine.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/Bufferization.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/Builtin.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/Complex.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/Tosa.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/TransformAffine.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/TransformBufferization.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/TransformGPU.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/TransformLinalg.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/TransformMemRef.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/TransformNVGPU.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/TransformPDL.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/TransformSMT.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/TransformTensor.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/TransformVector.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/UB.cpp
 create mode 100644 mlir/python/mlir/dialects/QuantOps.td

diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index d8b6d493f985c..ceebc931e60dc 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -457,11 +457,14 @@ endfunction()
 #     This file is where the *Attrs are defined, not where the *Enums are defined.
 #     **WARNING**: This arg will shortly be removed when the TODO for
 #     declare_mlir_dialect_python_bindings is satisfied. Use at your risk.
+#   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_dialect_extension_python_bindings)
   cmake_parse_arguments(ARG
     "GEN_ENUM_BINDINGS"
     "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME;EXTENSION_NAME"
-    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE"
+    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE;EMBED_CAPI_LINK_LIBS"
     ${ARGN})
   # Source files.
   set(_extension_target "${ARG_ADD_TO_PARENT}.${ARG_EXTENSION_NAME}")
@@ -503,6 +506,7 @@ function(declare_mlir_dialect_extension_python_bindings)
       ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
       ADD_TO_PARENT "${_extension_target}"
       SOURCES ${_sources}
+      EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
     )
   endif()
 endfunction()
diff --git a/mlir/include/mlir-c/Bindings/Python/Interop.h b/mlir/include/mlir-c/Bindings/Python/Interop.h
index a33190c380d37..89559da689017 100644
--- a/mlir/include/mlir-c/Bindings/Python/Interop.h
+++ b/mlir/include/mlir-c/Bindings/Python/Interop.h
@@ -84,6 +84,8 @@
 #define MLIR_PYTHON_CAPSULE_VALUE MAKE_MLIR_PYTHON_QUALNAME("ir.Value._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_TYPEID                                             \
   MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID._CAPIPtr")
+#define MLIR_PYTHON_CAPSULE_DIALECT_HANDLE                                     \
+  MAKE_MLIR_PYTHON_QUALNAME("ir.DialectHandle._CAPIPtr")
 
 /** Attribute on MLIR Python objects that expose their C-API pointer.
  * This will be a type-specific capsule created as per one of the helpers
@@ -457,6 +459,13 @@ static inline MlirValue mlirPythonCapsuleToValue(PyObject *capsule) {
   return value;
 }
 
+static inline MlirDialectHandle
+mlirPythonCapsuleToDialectHandle(PyObject *capsule) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_DIALECT_HANDLE);
+  MlirDialectHandle handle = {ptr};
+  return handle;
+}
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/include/mlir-c/Dialect/Affine.h b/mlir/include/mlir-c/Dialect/Affine.h
new file mode 100644
index 0000000000000..b2bf5aad44de9
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Affine.h
@@ -0,0 +1,36 @@
+//===-- mlir-c/Dialect/Affine.h - C API for Affine dialect --------*- C -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Affine dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_AFFINE_H
+#define MLIR_C_DIALECT_AFFINE_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Affine, affine);
+
+MLIR_CAPI_EXPORTED void
+mlirAffineRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_AFFINE_H
diff --git a/mlir/include/mlir-c/Dialect/Bufferization.h b/mlir/include/mlir-c/Dialect/Bufferization.h
new file mode 100644
index 0000000000000..41af7a294eb5c
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Bufferization.h
@@ -0,0 +1,36 @@
+//===-- mlir-c/Dialect/Bufferization.h - C API for Bufferization dialect --===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Bufferization dialect. A dialect should be registered with a context to make
+// it available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_BUFFERIZATION_H
+#define MLIR_C_DIALECT_BUFFERIZATION_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Bufferization, bufferization);
+
+MLIR_CAPI_EXPORTED void mlirBufferizationRegisterTransformDialectExtension(
+    MlirDialectRegistry registry);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_BUFFERIZATION_H
diff --git a/mlir/include/mlir-c/Dialect/Builtin.h b/mlir/include/mlir-c/Dialect/Builtin.h
new file mode 100644
index 0000000000000..c5d958249b36f
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Builtin.h
@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/Builtin.h - C API for Builtin dialect ------*- C -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Builtin dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_BUILTIN_H
+#define MLIR_C_DIALECT_BUILTIN_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Builtin, builtin);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_BUILTIN_H
diff --git a/mlir/include/mlir-c/Dialect/Complex.h b/mlir/include/mlir-c/Dialect/Complex.h
new file mode 100644
index 0000000000000..e51a67346a6ee
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Complex.h
@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/Complex.h - C API for Complex dialect ------*- C -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Complex dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_COMPLEX_H
+#define MLIR_C_DIALECT_COMPLEX_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Complex, complex);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_COMPLEX_H
diff --git a/mlir/include/mlir-c/Dialect/GPU.h b/mlir/include/mlir-c/Dialect/GPU.h
index 321c1122c3370..44d6308c270b2 100644
--- a/mlir/include/mlir-c/Dialect/GPU.h
+++ b/mlir/include/mlir-c/Dialect/GPU.h
@@ -63,6 +63,9 @@ mlirGPUObjectAttrHasKernels(MlirAttribute mlirObjectAttr);
 MLIR_CAPI_EXPORTED MlirAttribute
 mlirGPUObjectAttrGetKernels(MlirAttribute mlirObjectAttr);
 
+MLIR_CAPI_EXPORTED void
+mlirGPURegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/include/mlir-c/Dialect/Linalg.h b/mlir/include/mlir-c/Dialect/Linalg.h
index 339e63d667c5e..a3fa3a93a70c7 100644
--- a/mlir/include/mlir-c/Dialect/Linalg.h
+++ b/mlir/include/mlir-c/Dialect/Linalg.h
@@ -55,6 +55,9 @@ mlirLinalgGetIndexingMapsAttribute(MlirOperation op);
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Linalg, linalg);
 
+MLIR_CAPI_EXPORTED void
+mlirLinalgRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/include/mlir-c/Dialect/MemRef.h b/mlir/include/mlir-c/Dialect/MemRef.h
index 087a4b3f85b1d..38184404d6853 100644
--- a/mlir/include/mlir-c/Dialect/MemRef.h
+++ b/mlir/include/mlir-c/Dialect/MemRef.h
@@ -26,6 +26,9 @@ extern "C" {
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(MemRef, memref);
 
+MLIR_CAPI_EXPORTED void
+mlirMemRefRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/include/mlir-c/Dialect/NVGPU.h b/mlir/include/mlir-c/Dialect/NVGPU.h
index e58015a4a3421..89783d3758f9a 100644
--- a/mlir/include/mlir-c/Dialect/NVGPU.h
+++ b/mlir/include/mlir-c/Dialect/NVGPU.h
@@ -29,6 +29,9 @@ MLIR_CAPI_EXPORTED MlirType mlirNVGPUTensorMapDescriptorTypeGet(
     MlirContext ctx, MlirType tensorMemrefType, int swizzle, int l2promo,
     int oobFill, int interleave);
 
+MLIR_CAPI_EXPORTED void
+mlirNVGPURegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/include/mlir-c/Dialect/PDL.h b/mlir/include/mlir-c/Dialect/PDL.h
index 6ad2e2da62d87..eb611964e278b 100644
--- a/mlir/include/mlir-c/Dialect/PDL.h
+++ b/mlir/include/mlir-c/Dialect/PDL.h
@@ -66,6 +66,9 @@ MLIR_CAPI_EXPORTED bool mlirTypeIsAPDLValueType(MlirType type);
 
 MLIR_CAPI_EXPORTED MlirType mlirPDLValueTypeGet(MlirContext ctx);
 
+MLIR_CAPI_EXPORTED void
+mlirPDLRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/include/mlir-c/Dialect/SMT.h b/mlir/include/mlir-c/Dialect/SMT.h
index 0ad64746f148b..e6d8e14efed6e 100644
--- a/mlir/include/mlir-c/Dialect/SMT.h
+++ b/mlir/include/mlir-c/Dialect/SMT.h
@@ -104,6 +104,9 @@ mlirSMTAttrGetBVCmpPredicate(MlirContext ctx, MlirStringRef str);
 MLIR_CAPI_EXPORTED MlirAttribute mlirSMTAttrGetIntPredicate(MlirContext ctx,
                                                             MlirStringRef str);
 
+MLIR_CAPI_EXPORTED void
+mlirSMTRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/include/mlir-c/Dialect/SparseTensor.h b/mlir/include/mlir-c/Dialect/SparseTensor.h
index c816c1b58690e..0701211375a48 100644
--- a/mlir/include/mlir-c/Dialect/SparseTensor.h
+++ b/mlir/include/mlir-c/Dialect/SparseTensor.h
@@ -107,6 +107,9 @@ mlirSparseTensorEncodingAttrBuildLvlType(
     const enum MlirSparseTensorLevelPropertyNondefault *properties,
     unsigned propSize, unsigned n, unsigned m);
 
+MLIR_CAPI_EXPORTED void
+mlirSparseTensorRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/include/mlir-c/Dialect/Tensor.h b/mlir/include/mlir-c/Dialect/Tensor.h
index 74cbc5a6f0fe1..5891ade3bf3fd 100644
--- a/mlir/include/mlir-c/Dialect/Tensor.h
+++ b/mlir/include/mlir-c/Dialect/Tensor.h
@@ -18,6 +18,9 @@ extern "C" {
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Tensor, tensor);
 
+MLIR_CAPI_EXPORTED void
+mlirTensorRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/include/mlir-c/Dialect/Tosa.h b/mlir/include/mlir-c/Dialect/Tosa.h
new file mode 100644
index 0000000000000..ed55577996604
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Tosa.h
@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/Tosa.h - C API for Tosa dialect ----------*- C ---*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Tosa dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_TOSA_H
+#define MLIR_C_DIALECT_TOSA_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Tosa, tosa);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_TOSA_H
diff --git a/mlir/include/mlir-c/Dialect/Transform.h b/mlir/include/mlir-c/Dialect/Transform.h
index 02c99b5921882..1b5232c106edd 100644
--- a/mlir/include/mlir-c/Dialect/Transform.h
+++ b/mlir/include/mlir-c/Dialect/Transform.h
@@ -76,6 +76,19 @@ MLIR_CAPI_EXPORTED MlirType mlirTransformParamTypeGet(MlirContext ctx,
 
 MLIR_CAPI_EXPORTED MlirType mlirTransformParamTypeGetType(MlirType type);
 
+//===---------------------------------------------------------------------===//
+// Extension Registration
+//===---------------------------------------------------------------------===//
+
+MLIR_CAPI_EXPORTED void
+mlirDebugRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
+MLIR_CAPI_EXPORTED void
+mlirLoopRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
+MLIR_CAPI_EXPORTED void
+mlirTuneRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/include/mlir-c/Dialect/UB.h b/mlir/include/mlir-c/Dialect/UB.h
new file mode 100644
index 0000000000000..74159f0c705de
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/UB.h
@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/UB.h - C API for UB dialect ----------------*- C -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// UB dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_UB_H
+#define MLIR_C_DIALECT_UB_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(UB, ub);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_UB_H
diff --git a/mlir/include/mlir-c/Dialect/Vector.h b/mlir/include/mlir-c/Dialect/Vector.h
index 6256c82d172ce..b7248b7985057 100644
--- a/mlir/include/mlir-c/Dialect/Vector.h
+++ b/mlir/include/mlir-c/Dialect/Vector.h
@@ -26,6 +26,9 @@ extern "C" {
 
 MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Vector, vector);
 
+MLIR_CAPI_EXPORTED void
+mlirVectorRegisterTransformDialectExtension(MlirDialectRegistry registry);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/include/mlir-c/IR.h b/mlir/include/mlir-c/IR.h
index 061d7620ba077..55cc86accb8a0 100644
--- a/mlir/include/mlir-c/IR.h
+++ b/mlir/include/mlir-c/IR.h
@@ -66,6 +66,7 @@ DEFINE_C_API_STRUCT(MlirLocation, const void);
 DEFINE_C_API_STRUCT(MlirModule, const void);
 DEFINE_C_API_STRUCT(MlirType, const void);
 DEFINE_C_API_STRUCT(MlirValue, const void);
+DEFINE_C_API_STRUCT(MlirDialectHandle, const void);
 
 #undef DEFINE_C_API_STRUCT
 
@@ -207,11 +208,6 @@ MLIR_CAPI_EXPORTED MlirStringRef mlirDialectGetNamespace(MlirDialect dialect);
 // registration schemes.
 //===----------------------------------------------------------------------===//
 
-struct MlirDialectHandle {
-  const void *ptr;
-};
-typedef struct MlirDialectHandle MlirDialectHandle;
-
 #define MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Name, Namespace)                \
   MLIR_CAPI_EXPORTED MlirDialectHandle mlirGetDialectHandle__##Namespace##__(  \
       void)
@@ -233,6 +229,11 @@ MLIR_CAPI_EXPORTED void mlirDialectHandleRegisterDialect(MlirDialectHandle,
 MLIR_CAPI_EXPORTED MlirDialect mlirDialectHandleLoadDialect(MlirDialectHandle,
                                                             MlirContext);
 
+/// Checks if the dialect handle is null.
+static inline bool mlirDialectHandleIsNull(MlirDialectHandle handle) {
+  return !handle.ptr;
+}
+
 //===----------------------------------------------------------------------===//
 // DialectRegistry API.
 //===----------------------------------------------------------------------===//
@@ -249,6 +250,13 @@ static inline bool mlirDialectRegistryIsNull(MlirDialectRegistry registry) {
 MLIR_CAPI_EXPORTED void
 mlirDialectRegistryDestroy(MlirDialectRegistry registry);
 
+MLIR_CAPI_EXPORTED int64_t
+mlirDialectRegistryGetNumDialectNames(MlirDialectRegistry registry);
+
+MLIR_CAPI_EXPORTED void
+mlirDialectRegistryGetDialectNames(MlirDialectRegistry registry,
+                                   MlirStringRef *dialectNames);
+
 //===----------------------------------------------------------------------===//
 // Location API.
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 83a8757bb72c7..c3d28733c2cee 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2897,6 +2897,14 @@ maybeGetTracebackLocation(const std::optional<PyLocation> &location) {
 // Populates the core exports of the 'ir' submodule.
 //------------------------------------------------------------------------------
 
+MlirDialectHandle createMlirDialectHandleFromCapsule(nb::object capsule) {
+  MlirDialectHandle rawRegistry =
+      mlirPythonCapsuleToDialectHandle(capsule.ptr());
+  if (mlirDialectHandleIsNull(rawRegistry))
+    throw nb::python_error();
+  return rawRegistry;
+}
+
 void mlir::python::populateIRCore(nb::module_ &m) {
   // disable leak warnings which tend to be false positives.
   nb::set_leak_warnings(false);
@@ -3126,6 +3134,14 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           },
           nb::sig("def __repr__(self) -> str"));
 
+  //----------------------------------------------------------------------------
+  // Mapping of MlirDialectHandle
+  //----------------------------------------------------------------------------
+
+  nb::class_<MlirDialectHandle>(m, "DialectHandle")
+      .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR,
+                  &createMlirDialectHandleFromCapsule);
+
   //----------------------------------------------------------------------------
   // Mapping of PyDialectRegistry
   //----------------------------------------------------------------------------
@@ -3133,7 +3149,23 @@ void mlir::python::populateIRCore(nb::module_ &m) {
       .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyDialectRegistry::getCapsule)
       .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR,
                   &PyDialectRegistry::createFromCapsule)
-      .def(nb::init<>());
+      .def(nb::init<>())
+      .def("insert_dialect",
+           [](PyDialectRegistry &self, MlirDialectHandle handle) {
+             mlirDialectHandleInsertDialect(handle, self.get());
+           })
+      .def("insert_dialect",
+           [](PyDialectRegistry &self, intptr_t ptr) {
+             mlirDialectHandleInsertDialect(
+                 {reinterpret_cast<const void *>(ptr)}, self.get());
+           })
+      .def_prop_ro("dialect_names", [](PyDialectRegistry &self) {
+        int64_t numDialectNames =
+            mlirDialectRegistryGetNumDialectNames(self.get());
+        std::vector<MlirStringRef> dialectNames(numDialectNames);
+        mlirDialectRegistryGetDialectNames(self.get(), dialectNames.data());
+        return dialectNames;
+      });
 
   //----------------------------------------------------------------------------
   // Mapping of Location
diff --git a/mlir/lib/CAPI/Dialect/Affine.cpp b/mlir/lib/CAPI/Dialect/Affine.cpp
new file mode 100644
index 0000000000000..b796523390a29
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/Affine.cpp
@@ -0,0 +1,14 @@
+//===- Affine.cpp - C Interface for Affine dialect ------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Affine.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Affine, affine,
+                                      mlir::affine::AffineDialect)
diff --git a/mlir/lib/CAPI/Dialect/Bufferization.cpp b/mlir/lib/CAPI/Dialect/Bufferization.cpp
new file mode 100644
index 0000000000000..da50dfaab44b6
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/Bufferization.cpp
@@ -0,0 +1,14 @@
+//===- Bufferization.cpp - C Interface for Bufferization dialect ----------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Bufferization.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Bufferization, bufferization,
+                                      mlir::bufferization::BufferizationDialect)
diff --git a/mlir/lib/CAPI/Dialect/Builtin.cpp b/mlir/lib/CAPI/Dialect/Builtin.cpp
new file mode 100644
index 0000000000000..d095daa294c56
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/Builtin.cpp
@@ -0,0 +1,13 @@
+//===- Builtin.cpp - C Interface for Builtin dialect ----------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Builtin.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/IR/BuiltinDialect.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Builtin, builtin, mlir::BuiltinDialect)
diff --git a/mlir/lib/CAPI/Dialect/CMakeLists.txt b/mlir/lib/CAPI/Dialect/CMakeLists.txt
index bb1fdf8be3c8f..8b1321d0c125a 100644
--- a/mlir/lib/CAPI/Dialect/CMakeLists.txt
+++ b/mlir/lib/CAPI/Dialect/CMakeLists.txt
@@ -7,6 +7,15 @@ add_mlir_upstream_c_api_library(MLIRCAPIAMDGPU
   MLIRAMDGPUDialect
 )
 
+add_mlir_upstream_c_api_library(MLIRCAPIAffine
+  Affine.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRAffineDialect
+)
+
 add_mlir_upstream_c_api_library(MLIRCAPIArith
   Arith.cpp
 
@@ -16,6 +25,23 @@ add_mlir_upstream_c_api_library(MLIRCAPIArith
   MLIRArithDialect
 )
 
+add_mlir_upstream_c_api_library(MLIRCAPIBufferization
+  Bufferization.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRBufferizationDialect
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIBuiltin
+  Builtin.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+)
+
 add_mlir_upstream_c_api_library(MLIRCAPIAsync
   Async.cpp
   AsyncPasses.cpp
@@ -31,6 +57,15 @@ add_mlir_upstream_c_api_library(MLIRCAPIAsync
   MLIRPass
 )
 
+add_mlir_upstream_c_api_library(MLIRCAPIComplex
+  Complex.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRComplexDialect
+)
+
 add_mlir_upstream_c_api_library(MLIRCAPIControlFlow
   ControlFlow.cpp
 
@@ -223,6 +258,9 @@ add_mlir_upstream_c_api_library(MLIRCAPITransformDialect
   LINK_LIBS PUBLIC
   MLIRCAPIIR
   MLIRTransformDialect
+  MLIRTransformDebugExtension
+  MLIRTransformLoopExtension
+  MLIRTransformTuneExtension
 )
 
 add_mlir_upstream_c_api_library(MLIRCAPITransformDialectTransforms
@@ -234,6 +272,94 @@ add_mlir_upstream_c_api_library(MLIRCAPITransformDialectTransforms
   MLIRTransformDialectTransforms
 )
 
+add_mlir_upstream_c_api_library(MLIRCAPIAffineTransformDialectExtension
+  TransformAffine.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRAffineTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIBufferizationTransformDialectExtension
+  TransformBufferization.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRBufferizationTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIGPUTransformDialectExtension
+  TransformGPU.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRGPUTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPILinalgTransformDialectExtension
+  TransformLinalg.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRLinalgTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIMemRefTransformDialectExtension
+  TransformMemRef.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRMemRefTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPINVGPUTransformDialectExtension
+  TransformNVGPU.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRNVGPUTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIPDLTransformDialectExtension
+  TransformPDL.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRTransformPDLExtension
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPISMTTransformDialectExtension
+  TransformSMT.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRTransformSMTExtension
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPISparseTensorTransformDialectExtension
+  TransformSparseTensor.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRSparseTensorTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPITensorTransformDialectExtension
+  TransformTensor.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRTensorTransformOps
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIVectorTransformDialectExtension
+  TransformVector.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRVectorTransformOps
+)
+
 add_mlir_upstream_c_api_library(MLIRCAPIQuant
   Quant.cpp
 
@@ -278,3 +404,21 @@ add_mlir_upstream_c_api_library(MLIRCAPISMT
   MLIRCAPIIR
   MLIRSMT
 )
+
+add_mlir_upstream_c_api_library(MLIRCAPITosa
+  Tosa.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRTosaDialect
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIUB
+  UB.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRUBDialect
+)
diff --git a/mlir/lib/CAPI/Dialect/Complex.cpp b/mlir/lib/CAPI/Dialect/Complex.cpp
new file mode 100644
index 0000000000000..7063028e5d640
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/Complex.cpp
@@ -0,0 +1,14 @@
+//===- Complex.cpp - C Interface for Complex dialect ----------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Complex.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Complex, complex,
+                                      mlir::complex::ComplexDialect)
diff --git a/mlir/lib/CAPI/Dialect/Tosa.cpp b/mlir/lib/CAPI/Dialect/Tosa.cpp
new file mode 100644
index 0000000000000..357717a2ed5a5
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/Tosa.cpp
@@ -0,0 +1,13 @@
+//===- Tosa.cpp - C Interface for Tosa dialect ----------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Tosa.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Tosa, tosa, mlir::tosa::TosaDialect)
diff --git a/mlir/lib/CAPI/Dialect/Transform.cpp b/mlir/lib/CAPI/Dialect/Transform.cpp
index 5fd773572bd3c..8942d0a00cde3 100644
--- a/mlir/lib/CAPI/Dialect/Transform.cpp
+++ b/mlir/lib/CAPI/Dialect/Transform.cpp
@@ -9,8 +9,11 @@
 #include "mlir-c/Dialect/Transform.h"
 #include "mlir-c/Support.h"
 #include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Transform/DebugExtension/DebugExtension.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
+#include "mlir/Dialect/Transform/LoopExtension/LoopExtension.h"
+#include "mlir/Dialect/Transform/TuneExtension/TuneExtension.h"
 
 using namespace mlir;
 
@@ -106,3 +109,15 @@ MlirType mlirTransformParamTypeGet(MlirContext ctx, MlirType type) {
 MlirType mlirTransformParamTypeGetType(MlirType type) {
   return wrap(cast<transform::ParamType>(unwrap(type)).getType());
 }
+
+void mlirDebugRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+  mlir::transform::registerDebugExtension(*unwrap(registry));
+}
+
+void mlirLoopRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+  mlir::transform::registerLoopExtension(*unwrap(registry));
+}
+
+void mlirTuneRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+  mlir::transform::registerTuneExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformAffine.cpp b/mlir/lib/CAPI/Dialect/TransformAffine.cpp
new file mode 100644
index 0000000000000..6fd7744e8c2c9
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformAffine.cpp
@@ -0,0 +1,16 @@
+//===- TransformAffine.cpp - C Interface for Transform affine extension ---===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Affine.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h"
+
+MLIR_CAPI_EXPORTED void
+mlirAffineRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+  mlir::affine::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformBufferization.cpp b/mlir/lib/CAPI/Dialect/TransformBufferization.cpp
new file mode 100644
index 0000000000000..f1bd61be10880
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformBufferization.cpp
@@ -0,0 +1,16 @@
+//==- TransformBufferization.cpp - C Interface for bufferization extension -==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Bufferization.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
+
+void mlirBufferizationRegisterTransformDialectExtension(
+    MlirDialectRegistry registry) {
+  mlir::bufferization::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformGPU.cpp b/mlir/lib/CAPI/Dialect/TransformGPU.cpp
new file mode 100644
index 0000000000000..51187a28fa752
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformGPU.cpp
@@ -0,0 +1,15 @@
+//===- TransformGPU.cpp - C Interface for Transform GPU extension ---------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/GPU.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
+
+void mlirGPURegisterTransformDialectExtension(MlirDialectRegistry registry) {
+  mlir::gpu::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformInterpreter.cpp b/mlir/lib/CAPI/Dialect/TransformInterpreter.cpp
index 145455e1c1b3d..05648de959e7a 100644
--- a/mlir/lib/CAPI/Dialect/TransformInterpreter.cpp
+++ b/mlir/lib/CAPI/Dialect/TransformInterpreter.cpp
@@ -23,8 +23,6 @@ using namespace mlir;
 
 DEFINE_C_API_PTR_METHODS(MlirTransformOptions, transform::TransformOptions)
 
-extern "C" {
-
 MlirTransformOptions mlirTransformOptionsCreate() {
   return wrap(new transform::TransformOptions);
 }
@@ -80,4 +78,3 @@ MlirLogicalResult mlirMergeSymbolsIntoFromClone(MlirOperation target,
       unwrap(target), std::move(otherOwning));
   return wrap(result);
 }
-}
diff --git a/mlir/lib/CAPI/Dialect/TransformLinalg.cpp b/mlir/lib/CAPI/Dialect/TransformLinalg.cpp
new file mode 100644
index 0000000000000..5e62a80011a1b
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformLinalg.cpp
@@ -0,0 +1,15 @@
+//===- TransformLinalg.cpp - C Interface for Transform Linalg extension ---===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Linalg.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h"
+
+void mlirLinalgRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+  mlir::linalg::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformMemRef.cpp b/mlir/lib/CAPI/Dialect/TransformMemRef.cpp
new file mode 100644
index 0000000000000..9b80295748514
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformMemRef.cpp
@@ -0,0 +1,15 @@
+//===- TransformMemRef.cpp - C Interface for Transform MemRef extension ---===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/MemRef.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
+
+void mlirMemRefRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+  mlir::memref::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformNVGPU.cpp b/mlir/lib/CAPI/Dialect/TransformNVGPU.cpp
new file mode 100644
index 0000000000000..9159dc319aeb1
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformNVGPU.cpp
@@ -0,0 +1,15 @@
+//===- TransformNVGPU.cpp - C Interface for Transform NVGPU extension -----===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/NVGPU.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/NVGPU/TransformOps/NVGPUTransformOps.h"
+
+void mlirNVGPURegisterTransformDialectExtension(MlirDialectRegistry registry) {
+  mlir::nvgpu::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformPDL.cpp b/mlir/lib/CAPI/Dialect/TransformPDL.cpp
new file mode 100644
index 0000000000000..550bcabee55ec
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformPDL.cpp
@@ -0,0 +1,15 @@
+//===- TransformSMT.cpp - C Interface for Transform SMT extension ---------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/PDL.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Transform/PDLExtension/PDLExtension.h"
+
+void mlirPDLRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+  mlir::transform::registerPDLExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformSMT.cpp b/mlir/lib/CAPI/Dialect/TransformSMT.cpp
new file mode 100644
index 0000000000000..449129dac99e5
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformSMT.cpp
@@ -0,0 +1,15 @@
+//===- TransformSMT.cpp - C Interface for Transform SMT extension ---------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/SMT.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Transform/SMTExtension/SMTExtension.h"
+
+void mlirSMTRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+  mlir::transform::registerSMTExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp b/mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp
new file mode 100644
index 0000000000000..3e5b91099a6fb
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp
@@ -0,0 +1,16 @@
+//===- TransformSparseTensor.cpp - C Interface for SparseTensor extension -===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/SparseTensor.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.h"
+
+void mlirSparseTensorRegisterTransformDialectExtension(
+    MlirDialectRegistry registry) {
+  mlir::sparse_tensor::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformTensor.cpp b/mlir/lib/CAPI/Dialect/TransformTensor.cpp
new file mode 100644
index 0000000000000..fc7a30b5b237e
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformTensor.cpp
@@ -0,0 +1,15 @@
+//===- TransformTensor.cpp - C Interface for Transform tensor extensio ----===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Tensor.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
+
+void mlirTensorRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+  mlir::tensor::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/TransformVector.cpp b/mlir/lib/CAPI/Dialect/TransformVector.cpp
new file mode 100644
index 0000000000000..027d37ea34e79
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/TransformVector.cpp
@@ -0,0 +1,15 @@
+//===- TransformVector.cpp - C Interface for Transform Vector extension ---===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Vector.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
+
+void mlirVectorRegisterTransformDialectExtension(MlirDialectRegistry registry) {
+  mlir::vector::registerTransformDialectExtension(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/Dialect/UB.cpp b/mlir/lib/CAPI/Dialect/UB.cpp
new file mode 100644
index 0000000000000..de989237159c4
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/UB.cpp
@@ -0,0 +1,13 @@
+//===- Ub.cpp - C Interface for UB dialect --------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/UB.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(UB, ub, mlir::ub::UBDialect)
diff --git a/mlir/lib/CAPI/IR/IR.cpp b/mlir/lib/CAPI/IR/IR.cpp
index e9844a7cc1909..a81e2a14e5255 100644
--- a/mlir/lib/CAPI/IR/IR.cpp
+++ b/mlir/lib/CAPI/IR/IR.cpp
@@ -150,6 +150,18 @@ void mlirDialectRegistryDestroy(MlirDialectRegistry registry) {
   delete unwrap(registry);
 }
 
+int64_t mlirDialectRegistryGetNumDialectNames(MlirDialectRegistry registry) {
+  auto dialectNames = unwrap(registry)->getDialectNames();
+  return std::distance(dialectNames.begin(), dialectNames.end());
+}
+
+void mlirDialectRegistryGetDialectNames(MlirDialectRegistry registry,
+                                        MlirStringRef *dialectNames) {
+  for (auto [i, location] :
+       llvm::enumerate(unwrap(registry)->getDialectNames()))
+    dialectNames[i] = wrap(location);
+}
+
 //===----------------------------------------------------------------------===//
 // AsmState API.
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
index 132ed815c354e..d8e54bf060bd6 100644
--- a/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
+++ b/mlir/lib/Dialect/Transform/IR/TransformOps.cpp
@@ -2928,7 +2928,8 @@ LogicalResult transform::SequenceOp::verify() {
       InFlightDiagnostic diag =
           emitOpError()
           << "expected children ops to implement TransformOpInterface";
-      diag.attachNote(child.getLoc()) << "op without interface";
+      diag.attachNote(child.getLoc())
+          << "op without interface: " << child.getName();
       return diag;
     }
 
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 368f344841ed1..082e4ebac43ee 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -64,7 +64,10 @@ declare_mlir_dialect_python_bindings(
   SOURCES
     dialects/affine.py
   DIALECT_NAME affine
-  GEN_ENUM_BINDINGS)
+  GEN_ENUM_BINDINGS
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIAffine
+  )
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -73,14 +76,20 @@ declare_mlir_dialect_python_bindings(
   SOURCES
     dialects/amdgpu.py
   DIALECT_NAME amdgpu
-  GEN_ENUM_BINDINGS)
+  GEN_ENUM_BINDINGS
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIAMDGPU
+  )
 
 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_NAME async
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIAsync
+  )
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -91,6 +100,8 @@ declare_mlir_dialect_python_bindings(
   DIALECT_NAME bufferization
   GEN_ENUM_BINDINGS_TD_FILE
     "../../include/mlir/Dialect/Bufferization/IR/BufferizationEnums.td"
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIBufferization
 )
 
 declare_mlir_dialect_python_bindings(
@@ -99,7 +110,10 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/BuiltinOps.td
   SOURCES
     dialects/builtin.py
-  DIALECT_NAME builtin)
+  DIALECT_NAME builtin
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIBuiltin
+  )
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -107,7 +121,10 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/ComplexOps.td
   SOURCES
     dialects/complex.py
-  DIALECT_NAME complex)
+  DIALECT_NAME complex
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIComplex
+  )
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -116,7 +133,10 @@ declare_mlir_dialect_python_bindings(
   SOURCES
     dialects/index.py
   DIALECT_NAME index
-  GEN_ENUM_BINDINGS)
+  GEN_ENUM_BINDINGS
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIIndex
+  )
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -124,7 +144,10 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/ControlFlowOps.td
   SOURCES
     dialects/cf.py
-  DIALECT_NAME cf)
+  DIALECT_NAME cf
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIControlFlow
+  )
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -132,7 +155,10 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/FuncOps.td
   SOURCES
     dialects/func.py
-  DIALECT_NAME func)
+  DIALECT_NAME func
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIFunc
+  )
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -140,7 +166,10 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/GPUOps.td
   SOURCES_GLOB dialects/gpu/*.py
   DIALECT_NAME gpu
-  GEN_ENUM_BINDINGS)
+  GEN_ENUM_BINDINGS
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIGPU
+  )
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -151,7 +180,10 @@ declare_mlir_dialect_python_bindings(
     dialects/linalg/*.py
   DIALECT_NAME linalg
   DEPENDS LinalgOdsGen
-  GEN_ENUM_BINDINGS)
+  GEN_ENUM_BINDINGS
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPILinalg
+  )
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -160,29 +192,40 @@ declare_mlir_dialect_python_bindings(
   SOURCES
     dialects/llvm.py
   DIALECT_NAME llvm
-  GEN_ENUM_BINDINGS)
+  GEN_ENUM_BINDINGS
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPILLVM
+  )
 
 declare_mlir_dialect_extension_python_bindings(
-ADD_TO_PARENT MLIRPythonSources.Dialects
-ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/TransformPDLExtensionOps.td
   SOURCES
     dialects/transform/pdl.py
   DIALECT_NAME transform
-  EXTENSION_NAME transform_pdl_extension)
+  EXTENSION_NAME transform_pdl_extension
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITransformDialect
+    MLIRCAPIPDLTransformDialectExtension
+)
 
 declare_mlir_dialect_extension_python_bindings(
-ADD_TO_PARENT MLIRPythonSources.Dialects
-ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/TransformSMTExtensionOps.td
   SOURCES
     dialects/transform/smt.py
   DIALECT_NAME transform
-  EXTENSION_NAME transform_smt_extension)
+  EXTENSION_NAME transform_smt_extension
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITransformDialect
+    MLIRCAPISMTTransformDialectExtension
+)
 
 declare_mlir_dialect_extension_python_bindings(
-ADD_TO_PARENT MLIRPythonSources.Dialects
-ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/TransformDebugExtensionOps.td
   SOURCES
     dialects/transform/debug.py
@@ -190,8 +233,8 @@ ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   EXTENSION_NAME transform_debug_extension)
 
 declare_mlir_dialect_extension_python_bindings(
-ADD_TO_PARENT MLIRPythonSources.Dialects
-ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/TransformTuneExtensionOps.td
   SOURCES
     dialects/transform/tune.py
@@ -208,7 +251,9 @@ declare_mlir_dialect_python_bindings(
   DIALECT_NAME transform
   GEN_ENUM_BINDINGS_TD_FILE
     "../../include/mlir/Dialect/Transform/IR/TransformAttrs.td"
-)
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITransformDialect
+  )
 
 declare_mlir_python_sources(
   MLIRPythonSources.Dialects.transform.extras
@@ -232,7 +277,11 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/bufferization.py
   DIALECT_NAME transform
-  EXTENSION_NAME bufferization_transform)
+  EXTENSION_NAME bufferization_transform
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITransformDialect
+    MLIRCAPIBufferizationTransformDialectExtension
+  )
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -241,7 +290,11 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/gpu.py
   DIALECT_NAME transform
-  EXTENSION_NAME gpu_transform)
+  EXTENSION_NAME gpu_transform
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITransformDialect
+    MLIRCAPIGPUTransformDialectExtension
+)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -250,7 +303,10 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/loop.py
   DIALECT_NAME transform
-  EXTENSION_NAME loop_transform)
+  EXTENSION_NAME loop_transform
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITransformDialect
+)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -259,7 +315,11 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/memref.py
   DIALECT_NAME transform
-  EXTENSION_NAME memref_transform)
+  EXTENSION_NAME memref_transform
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITransformDialect
+    MLIRCAPIMemRefTransformDialectExtension
+)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -268,7 +328,11 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/nvgpu.py
   DIALECT_NAME transform
-  EXTENSION_NAME nvgpu_transform)
+  EXTENSION_NAME nvgpu_transform
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITransformDialect
+    MLIRCAPINVGPUTransformDialectExtension
+)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -280,6 +344,9 @@ declare_mlir_dialect_extension_python_bindings(
   EXTENSION_NAME structured_transform
   GEN_ENUM_BINDINGS_TD_FILE
     "../../include/mlir/Dialect/Linalg/TransformOps/LinalgTransformEnums.td"
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITransformDialect
+    MLIRCAPILinalgTransformDialectExtension
 )
 
 declare_mlir_dialect_extension_python_bindings(
@@ -289,7 +356,11 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/sparse_tensor.py
   DIALECT_NAME transform
-  EXTENSION_NAME sparse_tensor_transform)
+  EXTENSION_NAME sparse_tensor_transform
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITransformDialect
+    MLIRCAPISparseTensorTransformDialectExtension
+)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -298,7 +369,11 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/tensor.py
   DIALECT_NAME transform
-  EXTENSION_NAME tensor_transform)
+  EXTENSION_NAME tensor_transform
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITransformDialect
+    MLIRCAPITensorTransformDialectExtension
+)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -310,6 +385,9 @@ declare_mlir_dialect_extension_python_bindings(
   EXTENSION_NAME vector_transform
   GEN_ENUM_BINDINGS_TD_FILE
     "../../include/mlir/Dialect/Vector/Transforms/VectorTransformsBase.td"
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITransformDialect
+    MLIRCAPIVectorTransformDialectExtension
 )
 
 declare_mlir_dialect_python_bindings(
@@ -317,7 +395,10 @@ declare_mlir_dialect_python_bindings(
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/MathOps.td
   SOURCES dialects/math.py
-  DIALECT_NAME math)
+  DIALECT_NAME math
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIMath
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -329,7 +410,7 @@ declare_mlir_dialect_python_bindings(
   GEN_ENUM_BINDINGS
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIArith
-  )
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -337,7 +418,10 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/MemRefOps.td
   SOURCES
     dialects/memref.py
-  DIALECT_NAME memref)
+  DIALECT_NAME memref
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIMemRef
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -345,7 +429,10 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/MLProgramOps.td
   SOURCES
     dialects/ml_program.py
-  DIALECT_NAME ml_program)
+  DIALECT_NAME ml_program
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIMLProgram
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -354,7 +441,10 @@ declare_mlir_dialect_python_bindings(
   SOURCES
     dialects/nvgpu.py
   DIALECT_NAME nvgpu
-  GEN_ENUM_BINDINGS)
+  GEN_ENUM_BINDINGS
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPINVGPU
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -363,7 +453,10 @@ declare_mlir_dialect_python_bindings(
   SOURCES
     dialects/nvvm.py
   DIALECT_NAME nvvm
-  GEN_ENUM_BINDINGS)
+  GEN_ENUM_BINDINGS
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPINVVM
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -371,16 +464,23 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/ROCDLOps.td
   SOURCES
     dialects/rocdl.py
-  DIALECT_NAME rocdl)
+  DIALECT_NAME rocdl
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIROCDL
+)
 
-declare_mlir_python_sources(
-  MLIRPythonSources.Dialects.quant
+declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/QuantOps.td
   GEN_ENUM_BINDINGS
   SOURCES
     dialects/quant.py
-    _mlir_libs/_mlir/dialects/quant.pyi)
+    _mlir_libs/_mlir/dialects/quant.pyi
+  DIALECT_NAME quant
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIQuant
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -388,7 +488,10 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/EmitC.td
   SOURCES
     dialects/emitc.py
-  DIALECT_NAME emitc)
+  DIALECT_NAME emitc
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIEmitC
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -397,7 +500,10 @@ declare_mlir_dialect_python_bindings(
   SOURCES
     dialects/pdl.py
     _mlir_libs/_mlir/dialects/pdl.pyi
-  DIALECT_NAME pdl)
+  DIALECT_NAME pdl
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIPDL
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -406,7 +512,10 @@ declare_mlir_dialect_python_bindings(
   SOURCES
     dialects/openmp.py
   DIALECT_NAME omp
-  DEPENDS omp_common_td)
+  DEPENDS omp_common_td
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIOpenMP
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -414,14 +523,20 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/SCFOps.td
   SOURCES
     dialects/scf.py
-  DIALECT_NAME scf)
+  DIALECT_NAME scf
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPISCF
+)
 
 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)
+  DIALECT_NAME shape
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIShape
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -431,6 +546,8 @@ declare_mlir_dialect_python_bindings(
   DIALECT_NAME sparse_tensor
   GEN_ENUM_BINDINGS_TD_FILE
     "../../include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td"
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPISparseTensor
 )
 
 declare_mlir_dialect_python_bindings(
@@ -440,14 +557,20 @@ declare_mlir_dialect_python_bindings(
   GEN_ENUM_BINDINGS
   SOURCES
     dialects/smt.py
-  DIALECT_NAME smt)
+  DIALECT_NAME smt
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPISMT
+)
 
 declare_mlir_dialect_python_bindings(
-    ADD_TO_PARENT MLIRPythonSources.Dialects
-    ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
-    TD_FILE dialects/SPIRVOps.td
-    SOURCES dialects/spirv.py
-    DIALECT_NAME spirv)
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/SPIRVOps.td
+  SOURCES dialects/spirv.py
+  DIALECT_NAME spirv
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPISPIRV
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -455,7 +578,10 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/TensorOps.td
   SOURCES
     dialects/tensor.py
-  DIALECT_NAME tensor)
+  DIALECT_NAME tensor
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITensor
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -463,6 +589,8 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/TosaOps.td
   SOURCES dialects/tosa.py
   DIALECT_NAME tosa
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPITosa
 )
 
 declare_mlir_dialect_python_bindings(
@@ -471,6 +599,8 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/UBOps.td
   SOURCES dialects/ub.py
   DIALECT_NAME ub
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIUB
 )
 
 declare_mlir_dialect_python_bindings(
@@ -480,7 +610,10 @@ declare_mlir_dialect_python_bindings(
   SOURCES dialects/vector.py
   DIALECT_NAME vector
   GEN_ENUM_BINDINGS_TD_FILE
-    "dialects/VectorAttributes.td")
+    "dialects/VectorAttributes.td"
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIVector
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -489,6 +622,8 @@ declare_mlir_dialect_python_bindings(
   SOURCES dialects/irdl.py
   DIALECT_NAME irdl
   GEN_ENUM_BINDINGS
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIIRDL
 )
 
 ################################################################################
@@ -528,32 +663,13 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
     MLIRCAPIDebug
     MLIRCAPIIR
     MLIRCAPIInterfaces
+    MLIRCAPITransforms
+    MLIRCAPIBuiltin
 
     # Dialects
     MLIRCAPIFunc
 )
 
-# This extension exposes an API to register all dialects, extensions, and passes
-# packaged in upstream MLIR and it is used for the upstream "mlir" Python
-# package. Downstreams will likely want to provide their own and not depend
-# on this one, since it links in the world.
-# Note that this is not added to any top-level source target for transitive
-# inclusion: It must be included explicitly by downstreams if desired. Note that
-# this has a very large impact on what gets built/packaged.
-declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
-  MODULE_NAME _mlirRegisterEverything
-  ROOT_DIR "${PYTHON_SOURCE_DIR}"
-  PYTHON_BINDINGS_LIBRARY nanobind
-  SOURCES
-    RegisterEverything.cpp
-  PRIVATE_LINK_LIBS
-    LLVMSupport
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPIConversion
-    MLIRCAPITransforms
-    MLIRCAPIRegisterEverything
-)
-
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
   MODULE_NAME _mlirDialectsLinalg
   ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
@@ -749,6 +865,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
     MLIRCAPIIR
     MLIRCAPISMT
     MLIRCAPIExportSMTLIB
+    MLIRCAPISMTTransformDialectExtension
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
@@ -871,7 +988,6 @@ add_mlir_python_common_capi_library(${MLIR_PYTHON_CAPI_DYLIB_NAME}
     MLIRPythonCAPI.HeaderSources
   DECLARED_SOURCES
     MLIRPythonSources
-    MLIRPythonExtension.RegisterEverything
     ${_ADDL_TEST_SOURCES}
 )
 
@@ -980,7 +1096,6 @@ endif()
 
 set(_declared_sources
   MLIRPythonSources
-  MLIRPythonExtension.RegisterEverything
   MLIRPythonCAPICTypesBinding)
 if(NOT CMAKE_CROSSCOMPILING)
   list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen)
diff --git a/mlir/python/mlir/_mlir_libs/_capi.py.in b/mlir/python/mlir/_mlir_libs/_capi.py.in
index 9568845e67de9..2313ecc4f9823 100644
--- a/mlir/python/mlir/_mlir_libs/_capi.py.in
+++ b/mlir/python/mlir/_mlir_libs/_capi.py.in
@@ -5,4 +5,46 @@
 import ctypes
 from pathlib import Path
 
-_capi = ctypes.CDLL(str(Path(__file__).parent / "@CMAKE_SHARED_LIBRARY_PREFIX@@MLIR_PYTHON_CAPI_DYLIB_NAME@@CMAKE_SHARED_LIBRARY_SUFFIX@"))
\ No newline at end of file
+_capi = ctypes.CDLL(
+    str(
+        Path(__file__).parent
+        / "@CMAKE_SHARED_LIBRARY_PREFIX@@MLIR_PYTHON_CAPI_DYLIB_NAME@@CMAKE_SHARED_LIBRARY_SUFFIX@"
+    )
+)
+
+PyCapsule_New = ctypes.pythonapi.PyCapsule_New
+PyCapsule_New.restype = ctypes.py_object
+PyCapsule_New.argtypes = ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p
+
+PyCapsule_GetPointer = ctypes.pythonapi.PyCapsule_GetPointer
+PyCapsule_GetPointer.argtypes = [ctypes.py_object, ctypes.c_char_p]
+PyCapsule_GetPointer.restype = ctypes.c_void_p
+
+MLIR_PYTHON_CAPSULE_DIALECT_HANDLE = (
+    "@MLIR_PYTHON_PACKAGE_PREFIX at .ir.DialectHandle._CAPIPtr"
+).encode()
+
+MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY = (
+    "@MLIR_PYTHON_PACKAGE_PREFIX at .ir.DialectRegistry._CAPIPtr"
+).encode()
+
+
+def register_dialect(dialect_handle_capi_name, dialect_registry):
+    if not hasattr(_capi, dialect_handle_capi_name):
+        raise RuntimeError(f"missing {dialect_handle_capi_name} API")
+    dialect_handle_capi = getattr(_capi, dialect_handle_capi_name)
+    dialect_handle_capi.argtypes = []
+    dialect_handle_capi.restype = ctypes.c_void_p
+    handle = dialect_handle_capi()
+    dialect_registry.insert_dialect(handle)
+
+
+def register_transform_dialect_extension(registration_capi_name, dialect_registry):
+    if not hasattr(_capi, registration_capi_name):
+        raise RuntimeError(f"missing {registration_capi_name} API")
+    registration_capi = getattr(_capi, registration_capi_name)
+    registration_capi.argtypes = [ctypes.c_void_p]
+    dialect_registry_ptr = PyCapsule_GetPointer(
+        dialect_registry._CAPIPtr, MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY
+    )
+    registration_capi(dialect_registry_ptr)
diff --git a/mlir/python/mlir/dialects/QuantOps.td b/mlir/python/mlir/dialects/QuantOps.td
new file mode 100644
index 0000000000000..46385fd00ac0c
--- /dev/null
+++ b/mlir/python/mlir/dialects/QuantOps.td
@@ -0,0 +1,14 @@
+//===-- QuantOps.td - Entry point for QuantOps bind --------*- 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_BINDINGS_QUANT_OPS
+#define PYTHON_BINDINGS_QUANT_OPS
+
+include "mlir/Dialect/Quant/IR/QuantOps.td"
+
+#endif
diff --git a/mlir/python/mlir/dialects/quant.py b/mlir/python/mlir/dialects/quant.py
index bf1fc5f2de378..7a7273d8e26be 100644
--- a/mlir/python/mlir/dialects/quant.py
+++ b/mlir/python/mlir/dialects/quant.py
@@ -2,4 +2,5 @@
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+from ._quant_ops_gen import *
 from .._mlir_libs._mlirDialectsQuant import *
diff --git a/mlir/python/mlir/dialects/transform/bufferization.py b/mlir/python/mlir/dialects/transform/bufferization.py
index 485a8a36b6305..228c9dab52b8d 100644
--- a/mlir/python/mlir/dialects/transform/bufferization.py
+++ b/mlir/python/mlir/dialects/transform/bufferization.py
@@ -15,6 +15,15 @@
 from enum import Enum
 from typing import Optional, overload, Union
 
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+    register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+    "mlirBufferizationRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class EmptyTensorToAllocTensorOp(EmptyTensorToAllocTensorOp):
diff --git a/mlir/python/mlir/dialects/transform/debug.py b/mlir/python/mlir/dialects/transform/debug.py
index f7c04268dc03d..353e4f07c883a 100644
--- a/mlir/python/mlir/dialects/transform/debug.py
+++ b/mlir/python/mlir/dialects/transform/debug.py
@@ -15,6 +15,15 @@
 
 from typing import Union
 
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+    register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+    "mlirDebugRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class EmitParamAsRemarkOp(EmitParamAsRemarkOp):
diff --git a/mlir/python/mlir/dialects/transform/gpu.py b/mlir/python/mlir/dialects/transform/gpu.py
index 00cf0840eeae9..b35004cb4a09a 100644
--- a/mlir/python/mlir/dialects/transform/gpu.py
+++ b/mlir/python/mlir/dialects/transform/gpu.py
@@ -15,6 +15,16 @@
 from typing import Optional, Sequence, Union, overload
 
 
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+    register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+    "mlirGPURegisterTransformDialectExtension", _get_dialect_registry()
+)
+
+
 @_ods_cext.register_operation(_Dialect, replace=True)
 class MapForallToBlocks(MapForallToBlocks):
     """Specialization for MapForallToBlocks class."""
diff --git a/mlir/python/mlir/dialects/transform/loop.py b/mlir/python/mlir/dialects/transform/loop.py
index c4770b1c4067e..3311ae47b4fa9 100644
--- a/mlir/python/mlir/dialects/transform/loop.py
+++ b/mlir/python/mlir/dialects/transform/loop.py
@@ -16,6 +16,15 @@
 
 from typing import Optional, Union
 
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+    register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+    "mlirLoopRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class LoopOutlineOp(LoopOutlineOp):
diff --git a/mlir/python/mlir/dialects/transform/memref.py b/mlir/python/mlir/dialects/transform/memref.py
index 56ea61eb817f8..f45d8054bfa58 100644
--- a/mlir/python/mlir/dialects/transform/memref.py
+++ b/mlir/python/mlir/dialects/transform/memref.py
@@ -14,6 +14,15 @@
 
 from typing import Optional, overload, Union
 
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+    register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+    "mlirMemRefRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class MemRefAllocaToGlobalOp(MemRefAllocaToGlobalOp):
diff --git a/mlir/python/mlir/dialects/transform/nvgpu.py b/mlir/python/mlir/dialects/transform/nvgpu.py
index 74ba4c9aeb6c1..3d6427d1b1ed9 100644
--- a/mlir/python/mlir/dialects/transform/nvgpu.py
+++ b/mlir/python/mlir/dialects/transform/nvgpu.py
@@ -3,3 +3,12 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from .._nvgpu_transform_ops_gen import *
+
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+    register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+    "mlirNVGPURegisterTransformDialectExtension", _get_dialect_registry()
+)
diff --git a/mlir/python/mlir/dialects/transform/pdl.py b/mlir/python/mlir/dialects/transform/pdl.py
index bb5fa7ffd3065..15b0c37837902 100644
--- a/mlir/python/mlir/dialects/transform/pdl.py
+++ b/mlir/python/mlir/dialects/transform/pdl.py
@@ -18,6 +18,16 @@
 from typing import Union
 
 
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+    register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+    "mlirPDLRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
+
 @_ods_cext.register_operation(_Dialect, replace=True)
 class PDLMatchOp(PDLMatchOp):
     def __init__(
diff --git a/mlir/python/mlir/dialects/transform/smt.py b/mlir/python/mlir/dialects/transform/smt.py
index 1f0b7f066118c..a4e8eb96d388c 100644
--- a/mlir/python/mlir/dialects/transform/smt.py
+++ b/mlir/python/mlir/dialects/transform/smt.py
@@ -15,6 +15,16 @@
     raise RuntimeError("Error loading imports from extension module") from e
 
 
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+    register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+    "mlirSMTRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
+
 @_ods_cext.register_operation(_Dialect, replace=True)
 class ConstrainParamsOp(ConstrainParamsOp):
     def __init__(
diff --git a/mlir/python/mlir/dialects/transform/sparse_tensor.py b/mlir/python/mlir/dialects/transform/sparse_tensor.py
index 8b33270dc23a1..efaf438e07b83 100644
--- a/mlir/python/mlir/dialects/transform/sparse_tensor.py
+++ b/mlir/python/mlir/dialects/transform/sparse_tensor.py
@@ -3,3 +3,11 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from .._sparse_tensor_transform_ops_gen import *
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+    register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+    "mlirSparseTensorRegisterTransformDialectExtension", _get_dialect_registry()
+)
diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index bf40cc532065d..85833d32ba55e 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -28,6 +28,15 @@
 
 from typing import List, Optional, Sequence, Union, overload
 
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+    register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+    "mlirLinalgRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class BufferizeToAllocationOp(BufferizeToAllocationOp):
diff --git a/mlir/python/mlir/dialects/transform/tensor.py b/mlir/python/mlir/dialects/transform/tensor.py
index 4eb30398f0872..904f57827146d 100644
--- a/mlir/python/mlir/dialects/transform/tensor.py
+++ b/mlir/python/mlir/dialects/transform/tensor.py
@@ -14,6 +14,15 @@
 
 from typing import Optional, overload, Union
 
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+    register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+    "mlirTensorRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class MakeLoopIndependentOp(MakeLoopIndependentOp):
diff --git a/mlir/python/mlir/dialects/transform/tune.py b/mlir/python/mlir/dialects/transform/tune.py
index f63f88a382422..6eadd9b4e9d5d 100644
--- a/mlir/python/mlir/dialects/transform/tune.py
+++ b/mlir/python/mlir/dialects/transform/tune.py
@@ -25,6 +25,15 @@
 
 from typing import Union
 
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+    register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+    "mlirTuneRegisterTransformDialectExtension", _get_dialect_registry()
+)
+
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class KnobOp(KnobOp):
diff --git a/mlir/python/mlir/dialects/transform/vector.py b/mlir/python/mlir/dialects/transform/vector.py
index af2435cb26cc4..1d7620188da83 100644
--- a/mlir/python/mlir/dialects/transform/vector.py
+++ b/mlir/python/mlir/dialects/transform/vector.py
@@ -4,3 +4,11 @@
 
 from .._vector_transform_enum_gen import *
 from .._vector_transform_ops_gen import *
+from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
+from ..._mlir_libs._capi import (
+    register_transform_dialect_extension as _register_transform_dialect_extension,
+)
+
+_register_transform_dialect_extension(
+    "mlirVectorRegisterTransformDialectExtension", _get_dialect_registry()
+)
diff --git a/mlir/test/python/dialects/gpu/dialect.py b/mlir/test/python/dialects/gpu/dialect.py
index 26ee9f34cb332..3517b6519b93e 100644
--- a/mlir/test/python/dialects/gpu/dialect.py
+++ b/mlir/test/python/dialects/gpu/dialect.py
@@ -2,6 +2,7 @@
 
 from mlir.ir import *
 import mlir.dialects.gpu as gpu
+import mlir.dialects.nvvm as nvvm
 import mlir.dialects.gpu.passes
 from mlir.passmanager import *
 
diff --git a/mlir/test/python/dialects/irdl.py b/mlir/test/python/dialects/irdl.py
index ed62db9b69968..4dbcdd6d0cef2 100644
--- a/mlir/test/python/dialects/irdl.py
+++ b/mlir/test/python/dialects/irdl.py
@@ -2,6 +2,7 @@
 
 from mlir.ir import *
 from mlir.dialects.irdl import *
+import mlir.dialects.arith
 import sys
 
 
diff --git a/mlir/test/python/dialects/memref.py b/mlir/test/python/dialects/memref.py
index b91fdc367cf30..0967e7e5d070c 100644
--- a/mlir/test/python/dialects/memref.py
+++ b/mlir/test/python/dialects/memref.py
@@ -2,6 +2,7 @@
 
 import mlir.dialects.arith as arith
 import mlir.dialects.memref as memref
+import mlir.dialects.func
 import mlir.extras.types as T
 from mlir.dialects.memref import _infer_memref_subview_result_type
 from mlir.ir import *
diff --git a/mlir/test/python/dialects/transform.py b/mlir/test/python/dialects/transform.py
index 6c5e4e5505b1c..3477539b8d2f0 100644
--- a/mlir/test/python/dialects/transform.py
+++ b/mlir/test/python/dialects/transform.py
@@ -12,6 +12,7 @@ def run(f):
             print("\nTEST:", f.__name__)
             f(module)
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/dialects/transform_bufferization_ext.py b/mlir/test/python/dialects/transform_bufferization_ext.py
index 03b16c324a407..39014735aa8fb 100644
--- a/mlir/test/python/dialects/transform_bufferization_ext.py
+++ b/mlir/test/python/dialects/transform_bufferization_ext.py
@@ -13,6 +13,7 @@ def run(f):
             print("\nTEST:", f.__name__)
             f()
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/dialects/transform_debug_ext.py b/mlir/test/python/dialects/transform_debug_ext.py
index 2dfdaed343865..0826ed14d199b 100644
--- a/mlir/test/python/dialects/transform_debug_ext.py
+++ b/mlir/test/python/dialects/transform_debug_ext.py
@@ -19,6 +19,7 @@ def run(f):
                 f(sequence.bodyTarget)
                 transform.YieldOp()
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/dialects/transform_extras.py b/mlir/test/python/dialects/transform_extras.py
index ea47f170cb632..3aeac681f8d49 100644
--- a/mlir/test/python/dialects/transform_extras.py
+++ b/mlir/test/python/dialects/transform_extras.py
@@ -2,7 +2,7 @@
 
 from typing import Callable
 from mlir import ir
-from mlir.dialects import scf, pdl
+from mlir.dialects import scf, pdl, arith
 from mlir.dialects.transform import (
     structured,
     get_parent_op,
@@ -30,6 +30,7 @@ def construct_and_print_in_module(f):
         with ir.InsertionPoint(module.body):
             f()
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/dialects/transform_gpu_ext.py b/mlir/test/python/dialects/transform_gpu_ext.py
index 17f5a6d38f8c1..5771ebd13ab77 100644
--- a/mlir/test/python/dialects/transform_gpu_ext.py
+++ b/mlir/test/python/dialects/transform_gpu_ext.py
@@ -19,6 +19,7 @@ def run(f):
                 transform.YieldOp()
         print("\nTEST:", f.__name__)
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/dialects/transform_loop_ext.py b/mlir/test/python/dialects/transform_loop_ext.py
index 430b33fba04c7..50c93fe9f46d7 100644
--- a/mlir/test/python/dialects/transform_loop_ext.py
+++ b/mlir/test/python/dialects/transform_loop_ext.py
@@ -13,6 +13,7 @@ def run(f):
             print("\nTEST:", f.__name__)
             f()
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/dialects/transform_memref_ext.py b/mlir/test/python/dialects/transform_memref_ext.py
index e7d871c9eac8c..55eceb57f97c8 100644
--- a/mlir/test/python/dialects/transform_memref_ext.py
+++ b/mlir/test/python/dialects/transform_memref_ext.py
@@ -13,6 +13,7 @@ def run(f):
             print("\nTEST:", f.__name__)
             f()
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/dialects/transform_nvgpu_ext.py b/mlir/test/python/dialects/transform_nvgpu_ext.py
index 1de8b25bab7a5..45672ff3b58e4 100644
--- a/mlir/test/python/dialects/transform_nvgpu_ext.py
+++ b/mlir/test/python/dialects/transform_nvgpu_ext.py
@@ -12,6 +12,7 @@ def run(f):
             print("\nTEST:", f.__name__)
             f()
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/dialects/transform_smt_ext.py b/mlir/test/python/dialects/transform_smt_ext.py
index 3692fd92344a6..0593421c7713d 100644
--- a/mlir/test/python/dialects/transform_smt_ext.py
+++ b/mlir/test/python/dialects/transform_smt_ext.py
@@ -19,6 +19,7 @@ def run(f):
                 f(sequence.bodyTarget)
                 transform.YieldOp()
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/dialects/transform_sparse_tensor_ext.py b/mlir/test/python/dialects/transform_sparse_tensor_ext.py
index e11cc6bf1e074..b22f27ab6494b 100644
--- a/mlir/test/python/dialects/transform_sparse_tensor_ext.py
+++ b/mlir/test/python/dialects/transform_sparse_tensor_ext.py
@@ -19,6 +19,7 @@ def run(f):
                 transform.YieldOp()
         print("\nTEST:", f.__name__)
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py
index 8785d6d360074..ee15ac43e3d1a 100644
--- a/mlir/test/python/dialects/transform_structured_ext.py
+++ b/mlir/test/python/dialects/transform_structured_ext.py
@@ -19,6 +19,7 @@ def run(f):
             f()
         module.operation.verify()
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/dialects/transform_tensor_ext.py b/mlir/test/python/dialects/transform_tensor_ext.py
index a2e7aa242b9da..9dde72f52a8fc 100644
--- a/mlir/test/python/dialects/transform_tensor_ext.py
+++ b/mlir/test/python/dialects/transform_tensor_ext.py
@@ -19,6 +19,7 @@ def run(f):
                 f(sequence.bodyTarget)
                 transform.YieldOp()
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/dialects/transform_tune_ext.py b/mlir/test/python/dialects/transform_tune_ext.py
index dfb93594bca52..f8db151e10703 100644
--- a/mlir/test/python/dialects/transform_tune_ext.py
+++ b/mlir/test/python/dialects/transform_tune_ext.py
@@ -19,6 +19,7 @@ def run(f):
                 f(sequence.bodyTarget)
                 transform.YieldOp()
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/dialects/transform_vector_ext.py b/mlir/test/python/dialects/transform_vector_ext.py
index 0cd9333dc1218..0e841f2895b55 100644
--- a/mlir/test/python/dialects/transform_vector_ext.py
+++ b/mlir/test/python/dialects/transform_vector_ext.py
@@ -21,6 +21,7 @@ def run_apply_patterns(f):
                 transform.YieldOp()
         print("\nTEST:", f.__name__)
         print(module)
+        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/integration/dialects/pdl.py b/mlir/test/python/integration/dialects/pdl.py
index c8e6197e03842..5876ecfa41a5a 100644
--- a/mlir/test/python/integration/dialects/pdl.py
+++ b/mlir/test/python/integration/dialects/pdl.py
@@ -1,6 +1,6 @@
 # RUN: %PYTHON %s 2>&1 | FileCheck %s
 
-from mlir.dialects import arith, func, pdl
+from mlir.dialects import arith, func, pdl, irdl
 from mlir.dialects.builtin import module
 from mlir.ir import *
 from mlir.rewrite import *
diff --git a/mlir/test/python/integration/dialects/transform.py b/mlir/test/python/integration/dialects/transform.py
index 303274a8f8828..f57eee7d81819 100644
--- a/mlir/test/python/integration/dialects/transform.py
+++ b/mlir/test/python/integration/dialects/transform.py
@@ -24,6 +24,7 @@ def construct_and_print_in_module(f):
             module = f(module)
         if module is not None:
             print(module)
+            module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/ir/builtin_types.py b/mlir/test/python/ir/builtin_types.py
index b42bfd9bc6587..22a68b5d76b97 100644
--- a/mlir/test/python/ir/builtin_types.py
+++ b/mlir/test/python/ir/builtin_types.py
@@ -2,7 +2,7 @@
 
 import gc
 from mlir.ir import *
-from mlir.dialects import arith, tensor, func, memref
+from mlir.dialects import arith, tensor, func, memref, transform, scf
 import mlir.extras.types as T
 
 
diff --git a/mlir/test/python/ir/capi.py b/mlir/test/python/ir/capi.py
index d60fbd820f91e..934fab91f05a0 100644
--- a/mlir/test/python/ir/capi.py
+++ b/mlir/test/python/ir/capi.py
@@ -1,6 +1,41 @@
 # RUN: %PYTHON %s | FileCheck %s
 
-from mlir._mlir_libs._capi import _capi
+import ctypes
+
+from mlir._mlir_libs import get_dialect_registry
+from mlir._mlir_libs._capi import (
+    _capi,
+    PyCapsule_New,
+    MLIR_PYTHON_CAPSULE_DIALECT_HANDLE,
+)
+from mlir.ir import DialectHandle
 
 print("success")
-# CHECK: success
\ No newline at end of file
+# CHECK: success
+
+
+if not hasattr(_capi, "mlirGetDialectHandle__arith__"):
+    raise Exception("missing API")
+_capi.mlirGetDialectHandle__arith__.argtypes = []
+_capi.mlirGetDialectHandle__arith__.restype = ctypes.c_void_p
+
+if not hasattr(_capi, "mlirGetDialectHandle__quant__"):
+    raise Exception("missing API")
+_capi.mlirGetDialectHandle__quant__.argtypes = []
+_capi.mlirGetDialectHandle__quant__.restype = ctypes.c_void_p
+
+dialect_registry = get_dialect_registry()
+# CHECK: ['builtin']
+print(dialect_registry.dialect_names)
+
+arith_handle = _capi.mlirGetDialectHandle__arith__()
+dialect_registry.insert_dialect(arith_handle)
+# CHECK: ['arith', 'builtin']
+print(dialect_registry.dialect_names)
+
+quant_handle = _capi.mlirGetDialectHandle__quant__()
+capsule = PyCapsule_New(quant_handle, MLIR_PYTHON_CAPSULE_DIALECT_HANDLE, None)
+dialect_handle = DialectHandle._CAPICreate(capsule)
+dialect_registry.insert_dialect(dialect_handle)
+# CHECK: ['arith', 'builtin', 'quant']
+print(dialect_registry.dialect_names)
diff --git a/mlir/test/python/ir/diagnostic_handler.py b/mlir/test/python/ir/diagnostic_handler.py
index 6d273e5092e42..f06090d982476 100644
--- a/mlir/test/python/ir/diagnostic_handler.py
+++ b/mlir/test/python/ir/diagnostic_handler.py
@@ -5,6 +5,7 @@
 from mlir._mlir_libs._mlirPythonTestNanobind import (
     test_diagnostics_with_errors_and_notes,
 )
+import mlir.dialects.arith
 
 
 def run(f):
diff --git a/mlir/test/python/ir/dialects.py b/mlir/test/python/ir/dialects.py
index 5a2ed684d298b..5ac077830cf02 100644
--- a/mlir/test/python/ir/dialects.py
+++ b/mlir/test/python/ir/dialects.py
@@ -4,6 +4,7 @@
 import sys
 from mlir.ir import *
 from mlir.dialects._ods_common import _cext
+from mlir.dialects import func, arith, scf, cf
 
 
 def run(f):
diff --git a/mlir/test/python/ir/exception.py b/mlir/test/python/ir/exception.py
index 74085cd349643..bcf95bf43acc3 100644
--- a/mlir/test/python/ir/exception.py
+++ b/mlir/test/python/ir/exception.py
@@ -2,6 +2,7 @@
 
 import gc
 from mlir.ir import *
+import mlir.dialects.func
 
 
 def run(f):
diff --git a/mlir/test/python/ir/insertion_point.py b/mlir/test/python/ir/insertion_point.py
index f48beb25f04b2..74223ce18599e 100644
--- a/mlir/test/python/ir/insertion_point.py
+++ b/mlir/test/python/ir/insertion_point.py
@@ -2,6 +2,7 @@
 
 import gc
 from mlir.ir import *
+import mlir.dialects.func
 
 
 def run(f):
diff --git a/mlir/test/python/ir/module.py b/mlir/test/python/ir/module.py
index 33959bea9ffb6..c9bb8fc309cc2 100644
--- a/mlir/test/python/ir/module.py
+++ b/mlir/test/python/ir/module.py
@@ -3,6 +3,7 @@
 import gc
 from tempfile import NamedTemporaryFile
 from mlir.ir import *
+import mlir.dialects.func
 
 
 def run(f):
diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
index 4a3625c953d52..2b37f79251af5 100644
--- a/mlir/test/python/ir/operation.py
+++ b/mlir/test/python/ir/operation.py
@@ -6,7 +6,7 @@
 from tempfile import NamedTemporaryFile
 from mlir.ir import *
 from mlir.dialects.builtin import ModuleOp
-from mlir.dialects import arith
+from mlir.dialects import arith, func
 from mlir.dialects._ods_common import _cext
 
 
diff --git a/mlir/test/python/ir/symbol_table.py b/mlir/test/python/ir/symbol_table.py
index 99d5fadfea10a..81efe3725f911 100644
--- a/mlir/test/python/ir/symbol_table.py
+++ b/mlir/test/python/ir/symbol_table.py
@@ -4,6 +4,7 @@
 import io
 import itertools
 from mlir.ir import *
+import mlir.dialects.func
 
 
 def run(f):
diff --git a/mlir/test/python/multithreaded_tests.py b/mlir/test/python/multithreaded_tests.py
index 6e1a668346872..513ec90fddfeb 100644
--- a/mlir/test/python/multithreaded_tests.py
+++ b/mlir/test/python/multithreaded_tests.py
@@ -50,7 +50,7 @@
 from typing import Optional, List
 
 import mlir.dialects.arith as arith
-from mlir.dialects import transform
+from mlir.dialects import transform, func, memref
 from mlir.ir import Context, Location, Module, IntegerType, InsertionPoint
 
 
diff --git a/mlir/test/python/pass_manager.py b/mlir/test/python/pass_manager.py
index 5f92f5b52a09a..db2915246d30f 100644
--- a/mlir/test/python/pass_manager.py
+++ b/mlir/test/python/pass_manager.py
@@ -5,6 +5,7 @@
 from mlir.passmanager import *
 from mlir.dialects.func import FuncOp
 from mlir.dialects.builtin import ModuleOp
+import mlir.dialects.arith
 
 
 # Log everything to stderr and flush so that we have a unified stream to match
diff --git a/mlir/test/python/python_pass.py b/mlir/test/python/python_pass.py
index 50c42102f66d3..7d47cf3684f5c 100644
--- a/mlir/test/python/python_pass.py
+++ b/mlir/test/python/python_pass.py
@@ -4,7 +4,7 @@
 from mlir.ir import *
 from mlir.passmanager import *
 from mlir.dialects.builtin import ModuleOp
-from mlir.dialects import pdl
+from mlir.dialects import pdl, func, arith
 from mlir.rewrite import *
 
 
diff --git a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
index 0172b3fa38a6b..6c98d3ed8d0a9 100644
--- a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
@@ -44,6 +44,13 @@ from ._ods_common import (
 _ods_ir = _ods_cext.ir
 _ods_cext.globals.register_traceback_file_exclusion(__file__)
 
+from .._mlir_libs import get_dialect_registry as _get_dialect_registry
+from .._mlir_libs._capi import register_dialect as _register_dialect
+
+_dialect_registry = _get_dialect_registry()
+if "{0}" not in _dialect_registry.dialect_names:
+    _register_dialect("mlirGetDialectHandle__{0}__", _dialect_registry)
+
 import builtins
 from typing import Sequence as _Sequence, Union as _Union, Optional as _Optional
 
@@ -1191,7 +1198,7 @@ static bool emitAllOps(const RecordKeeper &records, raw_ostream &os) {
   if (clDialectName.empty())
     llvm::PrintFatalError("dialect name not provided");
 
-  os << fileHeader;
+  os << formatv(fileHeader, clDialectName.getValue());
   if (!clDialectExtensionName.empty())
     os << formatv(dialectExtensionTemplate, clDialectName.getValue());
   else

>From e85d9c80c4dbc77b22a98f6d3ace5af36c35aacc Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental at gmail.com>
Date: Sat, 27 Sep 2025 03:58:37 -0400
Subject: [PATCH 3/4] factor registereverything

---
 mlir/include/mlir-c/RegisterAllExtensions.h   | 26 ++++++++++++++++
 .../mlir-c/RegisterAllLLVMTranslations.h      | 26 ++++++++++++++++
 mlir/include/mlir-c/RegisterAllPasses.h       | 26 ++++++++++++++++
 mlir/include/mlir-c/RegisterEverything.h      |  4 +--
 .../Bindings/Python/RegisterEverything.cpp    | 11 ++++---
 .../CAPI/RegisterEverything/CMakeLists.txt    | 30 +++++++++++++++++++
 .../RegisterAllExtensions.cpp                 | 16 ++++++++++
 .../RegisterAllLLVMTranslations.cpp           | 20 +++++++++++++
 .../RegisterEverything/RegisterAllPasses.cpp  | 12 ++++++++
 mlir/python/CMakeLists.txt                    | 25 ++++++++++++++++
 mlir/python/mlir/_mlir_libs/__init__.py       |  3 ++
 11 files changed, 193 insertions(+), 6 deletions(-)
 create mode 100644 mlir/include/mlir-c/RegisterAllExtensions.h
 create mode 100644 mlir/include/mlir-c/RegisterAllLLVMTranslations.h
 create mode 100644 mlir/include/mlir-c/RegisterAllPasses.h
 create mode 100644 mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp
 create mode 100644 mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp
 create mode 100644 mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp

diff --git a/mlir/include/mlir-c/RegisterAllExtensions.h b/mlir/include/mlir-c/RegisterAllExtensions.h
new file mode 100644
index 0000000000000..7238334ed5b8f
--- /dev/null
+++ b/mlir/include/mlir-c/RegisterAllExtensions.h
@@ -0,0 +1,26 @@
+//===-- mlir-c/RegisterAllExtensions.h - Register all MLIR Extensions ---*-===//
+//
+// 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 MLIR_C_REGISTER_Extensions_H
+#define MLIR_C_REGISTER_Extensions_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/// Register all compiler Extensions of MLIR.
+MLIR_CAPI_EXPORTED void mlirRegisterAllExtensions(MlirDialectRegistry registry);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_REGISTER_Extensions_H
diff --git a/mlir/include/mlir-c/RegisterAllLLVMTranslations.h b/mlir/include/mlir-c/RegisterAllLLVMTranslations.h
new file mode 100644
index 0000000000000..b2be041e8a93a
--- /dev/null
+++ b/mlir/include/mlir-c/RegisterAllLLVMTranslations.h
@@ -0,0 +1,26 @@
+//===-- mlir-c/RegisterAllLLVMTranslations.h - Register all LLVM translations //
+//
+// 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 MLIR_C_REGISTER_LLVM_TRANSLATIONS_H
+#define MLIR_C_REGISTER_LLVM_TRANSLATIONS_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/// Register all LLVM translations of MLIR.
+MLIR_CAPI_EXPORTED void mlirRegisterAllLLVMTranslations(MlirContext context);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_REGISTER_LLVM_TRANSLATIONS_H
diff --git a/mlir/include/mlir-c/RegisterAllPasses.h b/mlir/include/mlir-c/RegisterAllPasses.h
new file mode 100644
index 0000000000000..797b865885682
--- /dev/null
+++ b/mlir/include/mlir-c/RegisterAllPasses.h
@@ -0,0 +1,26 @@
+//===-- mlir-c/RegisterAllPasses.h - Register all MLIR Pass --*- C ------*-===//
+//
+// 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 MLIR_C_REGISTER_PASSES_H
+#define MLIR_C_REGISTER_PASSES_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/// Register all compiler passes of MLIR.
+MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_REGISTER_PASSES_H
diff --git a/mlir/include/mlir-c/RegisterEverything.h b/mlir/include/mlir-c/RegisterEverything.h
index ea2ea86449727..e759f6b188778 100644
--- a/mlir/include/mlir-c/RegisterEverything.h
+++ b/mlir/include/mlir-c/RegisterEverything.h
@@ -28,8 +28,8 @@ MLIR_CAPI_EXPORTED void mlirRegisterAllDialects(MlirDialectRegistry registry);
 /// Register all translations to LLVM IR for dialects that can support it.
 MLIR_CAPI_EXPORTED void mlirRegisterAllLLVMTranslations(MlirContext context);
 
-/// Register all compiler passes of MLIR.
-MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void);
+// /// Register all compiler passes of MLIR.
+// MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void);
 
 #ifdef __cplusplus
 }
diff --git a/mlir/lib/Bindings/Python/RegisterEverything.cpp b/mlir/lib/Bindings/Python/RegisterEverything.cpp
index 3edcb099c0a24..d773c15fbfc30 100644
--- a/mlir/lib/Bindings/Python/RegisterEverything.cpp
+++ b/mlir/lib/Bindings/Python/RegisterEverything.cpp
@@ -6,15 +6,18 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir-c/RegisterEverything.h"
+#include "mlir-c/RegisterAllExtensions.h"
+#include "mlir-c/RegisterAllLLVMTranslations.h"
+#include "mlir-c/RegisterAllPasses.h"
 #include "mlir/Bindings/Python/Nanobind.h"
 #include "mlir/Bindings/Python/NanobindAdaptors.h"
 
 NB_MODULE(_mlirRegisterEverything, m) {
-  m.doc() = "MLIR All Upstream Dialects, Translations and Passes Registration";
+  m.doc() =
+      "MLIR All Upstream Extensions, Translations and Passes Registration";
 
-  m.def("register_dialects", [](MlirDialectRegistry registry) {
-    mlirRegisterAllDialects(registry);
+  m.def("register_extensions", [](MlirDialectRegistry registry) {
+    mlirRegisterAllExtensions(registry);
   });
   m.def("register_llvm_translations",
         [](MlirContext context) { mlirRegisterAllLLVMTranslations(context); });
diff --git a/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt b/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt
index ccda668ece3d4..2ac01bc4c84f3 100644
--- a/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt
+++ b/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt
@@ -2,6 +2,7 @@
 get_property(translation_libs GLOBAL PROPERTY MLIR_TRANSLATION_LIBS)
 add_mlir_upstream_c_api_library(MLIRCAPIRegisterEverything
   RegisterEverything.cpp
+  PARTIAL_SOURCES_INTENDED
 
   LINK_LIBS PUBLIC
   ${translation_libs}
@@ -14,3 +15,32 @@ add_mlir_upstream_c_api_library(MLIRCAPIRegisterEverything
   MLIRRegisterAllExtensions
   MLIRRegisterAllPasses
 )
+
+add_mlir_upstream_c_api_library(MLIRCAPIRegisterAllPasses
+  RegisterAllPasses.cpp
+  PARTIAL_SOURCES_INTENDED
+
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRRegisterAllPasses
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIRegisterAllLLVMTranslations
+  RegisterAllLLVMTranslations.cpp
+  PARTIAL_SOURCES_INTENDED
+
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  ${translation_libs}
+  MLIRBuiltinToLLVMIRTranslation
+  MLIRLLVMToLLVMIRTranslation
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIRegisterAllExtensions
+  RegisterAllExtensions.cpp
+  PARTIAL_SOURCES_INTENDED
+
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRRegisterAllExtensions
+)
\ No newline at end of file
diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp
new file mode 100644
index 0000000000000..9cc9fd394e728
--- /dev/null
+++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp
@@ -0,0 +1,16 @@
+//===- RegisterAllExtensions.cpp - Register all MLIR entities
+//-----------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/RegisterAllExtensions.h"
+#include "mlir/CAPI/IR.h"
+#include "mlir/InitAllExtensions.h"
+
+void mlirRegisterAllExtensions(MlirDialectRegistry registry) {
+  mlir::registerAllExtensions(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp
new file mode 100644
index 0000000000000..c5dc64ca17834
--- /dev/null
+++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp
@@ -0,0 +1,20 @@
+//===- RegisterAllPasses.cpp - Register all MLIR entities -----------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/RegisterAllLLVMTranslations.h"
+#include "mlir/CAPI/IR.h"
+#include "mlir/Target/LLVMIR/Dialect/All.h"
+#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
+
+void mlirRegisterAllLLVMTranslations(MlirContext context) {
+  auto &ctx = *unwrap(context);
+  mlir::DialectRegistry registry;
+  mlir::registerAllToLLVMIRTranslations(registry);
+  ctx.appendDialectRegistry(registry);
+}
diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp
new file mode 100644
index 0000000000000..159dff0b86e44
--- /dev/null
+++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp
@@ -0,0 +1,12 @@
+//===- RegisterAllPasses.cpp - Register all MLIR entities -----------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/RegisterAllPasses.h"
+#include "mlir/InitAllPasses.h"
+
+void mlirRegisterAllPasses() { mlir::registerAllPasses(); }
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 082e4ebac43ee..9a45c5026f60e 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -670,6 +670,29 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
     MLIRCAPIFunc
 )
 
+# This extension exposes an API to register all dialects, extensions, and passes
+# packaged in upstream MLIR and it is used for the upstream "mlir" Python
+# package. Downstreams will likely want to provide their own and not depend
+# on this one, since it links in the world.
+# Note that this is not added to any top-level source target for transitive
+# inclusion: It must be included explicitly by downstreams if desired. Note that
+# this has a very large impact on what gets built/packaged.
+declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
+  MODULE_NAME _mlirRegisterEverything
+  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
+  SOURCES
+    RegisterEverything.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIConversion
+    MLIRCAPITransforms
+    MLIRCAPIRegisterAllPasses
+    MLIRCAPIRegisterAllExtensions
+    MLIRCAPIRegisterAllLLVMTranslations
+)
+
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
   MODULE_NAME _mlirDialectsLinalg
   ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
@@ -988,6 +1011,7 @@ add_mlir_python_common_capi_library(${MLIR_PYTHON_CAPI_DYLIB_NAME}
     MLIRPythonCAPI.HeaderSources
   DECLARED_SOURCES
     MLIRPythonSources
+    MLIRPythonExtension.RegisterEverything
     ${_ADDL_TEST_SOURCES}
 )
 
@@ -1096,6 +1120,7 @@ endif()
 
 set(_declared_sources
   MLIRPythonSources
+  MLIRPythonExtension.RegisterEverything
   MLIRPythonCAPICTypesBinding)
 if(NOT CMAKE_CROSSCOMPILING)
   list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen)
diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py
index 63244212ba42c..2c36397d2cf8b 100644
--- a/mlir/python/mlir/_mlir_libs/__init__.py
+++ b/mlir/python/mlir/_mlir_libs/__init__.py
@@ -123,6 +123,9 @@ def process_initializer_module(module_name):
         if hasattr(m, "register_dialects"):
             logger.debug("Registering dialects from initializer %r", m)
             m.register_dialects(get_dialect_registry())
+        if hasattr(m, "register_extensions"):
+            logger.debug("Registering extensions from initializer %r", m)
+            m.register_extensions(get_dialect_registry())
         if hasattr(m, "context_init_hook"):
             logger.debug("Adding context init hook from %r", m)
             post_init_hooks.append(m.context_init_hook)

>From 72ff4c7c6ed5007dc3291300aaf9901780b678fe Mon Sep 17 00:00:00 2001
From: Maksim Levental <maksim.levental at gmail.com>
Date: Sat, 27 Sep 2025 04:06:35 -0400
Subject: [PATCH 4/4] undo transform extension stuff

---
 mlir/cmake/modules/AddMLIRPython.cmake        |  6 +-
 mlir/include/mlir-c/Dialect/Transform.h       | 13 ---
 mlir/include/mlir-c/RegisterEverything.h      |  4 +-
 mlir/lib/CAPI/Dialect/CMakeLists.txt          | 91 -------------------
 mlir/lib/CAPI/Dialect/Transform.cpp           | 15 ---
 mlir/lib/CAPI/Dialect/TransformAffine.cpp     | 16 ----
 .../CAPI/Dialect/TransformBufferization.cpp   | 16 ----
 mlir/lib/CAPI/Dialect/TransformGPU.cpp        | 15 ---
 mlir/lib/CAPI/Dialect/TransformLinalg.cpp     | 15 ---
 mlir/lib/CAPI/Dialect/TransformMemRef.cpp     | 15 ---
 mlir/lib/CAPI/Dialect/TransformNVGPU.cpp      | 15 ---
 mlir/lib/CAPI/Dialect/TransformPDL.cpp        | 15 ---
 mlir/lib/CAPI/Dialect/TransformSMT.cpp        | 15 ---
 .../CAPI/Dialect/TransformSparseTensor.cpp    | 16 ----
 mlir/lib/CAPI/Dialect/TransformTensor.cpp     | 15 ---
 mlir/lib/CAPI/Dialect/TransformVector.cpp     | 15 ---
 mlir/python/CMakeLists.txt                    | 60 ++----------
 mlir/python/mlir/_mlir_libs/_capi.py.in       | 11 ---
 .../mlir/dialects/transform/bufferization.py  |  9 --
 mlir/python/mlir/dialects/transform/debug.py  |  9 --
 mlir/python/mlir/dialects/transform/gpu.py    | 10 --
 mlir/python/mlir/dialects/transform/loop.py   |  9 --
 mlir/python/mlir/dialects/transform/memref.py |  9 --
 mlir/python/mlir/dialects/transform/nvgpu.py  |  9 --
 mlir/python/mlir/dialects/transform/pdl.py    | 10 --
 mlir/python/mlir/dialects/transform/smt.py    | 10 --
 .../mlir/dialects/transform/sparse_tensor.py  |  8 --
 .../mlir/dialects/transform/structured.py     |  9 --
 mlir/python/mlir/dialects/transform/tensor.py |  9 --
 mlir/python/mlir/dialects/transform/tune.py   |  9 --
 mlir/python/mlir/dialects/transform/vector.py |  8 --
 31 files changed, 12 insertions(+), 474 deletions(-)
 delete mode 100644 mlir/lib/CAPI/Dialect/TransformAffine.cpp
 delete mode 100644 mlir/lib/CAPI/Dialect/TransformBufferization.cpp
 delete mode 100644 mlir/lib/CAPI/Dialect/TransformGPU.cpp
 delete mode 100644 mlir/lib/CAPI/Dialect/TransformLinalg.cpp
 delete mode 100644 mlir/lib/CAPI/Dialect/TransformMemRef.cpp
 delete mode 100644 mlir/lib/CAPI/Dialect/TransformNVGPU.cpp
 delete mode 100644 mlir/lib/CAPI/Dialect/TransformPDL.cpp
 delete mode 100644 mlir/lib/CAPI/Dialect/TransformSMT.cpp
 delete mode 100644 mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp
 delete mode 100644 mlir/lib/CAPI/Dialect/TransformTensor.cpp
 delete mode 100644 mlir/lib/CAPI/Dialect/TransformVector.cpp

diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index ceebc931e60dc..d8b6d493f985c 100644
--- a/mlir/cmake/modules/AddMLIRPython.cmake
+++ b/mlir/cmake/modules/AddMLIRPython.cmake
@@ -457,14 +457,11 @@ endfunction()
 #     This file is where the *Attrs are defined, not where the *Enums are defined.
 #     **WARNING**: This arg will shortly be removed when the TODO for
 #     declare_mlir_dialect_python_bindings is satisfied. Use at your risk.
-#   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_dialect_extension_python_bindings)
   cmake_parse_arguments(ARG
     "GEN_ENUM_BINDINGS"
     "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME;EXTENSION_NAME"
-    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE;EMBED_CAPI_LINK_LIBS"
+    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE"
     ${ARGN})
   # Source files.
   set(_extension_target "${ARG_ADD_TO_PARENT}.${ARG_EXTENSION_NAME}")
@@ -506,7 +503,6 @@ function(declare_mlir_dialect_extension_python_bindings)
       ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
       ADD_TO_PARENT "${_extension_target}"
       SOURCES ${_sources}
-      EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
     )
   endif()
 endfunction()
diff --git a/mlir/include/mlir-c/Dialect/Transform.h b/mlir/include/mlir-c/Dialect/Transform.h
index 1b5232c106edd..02c99b5921882 100644
--- a/mlir/include/mlir-c/Dialect/Transform.h
+++ b/mlir/include/mlir-c/Dialect/Transform.h
@@ -76,19 +76,6 @@ MLIR_CAPI_EXPORTED MlirType mlirTransformParamTypeGet(MlirContext ctx,
 
 MLIR_CAPI_EXPORTED MlirType mlirTransformParamTypeGetType(MlirType type);
 
-//===---------------------------------------------------------------------===//
-// Extension Registration
-//===---------------------------------------------------------------------===//
-
-MLIR_CAPI_EXPORTED void
-mlirDebugRegisterTransformDialectExtension(MlirDialectRegistry registry);
-
-MLIR_CAPI_EXPORTED void
-mlirLoopRegisterTransformDialectExtension(MlirDialectRegistry registry);
-
-MLIR_CAPI_EXPORTED void
-mlirTuneRegisterTransformDialectExtension(MlirDialectRegistry registry);
-
 #ifdef __cplusplus
 }
 #endif
diff --git a/mlir/include/mlir-c/RegisterEverything.h b/mlir/include/mlir-c/RegisterEverything.h
index e759f6b188778..ea2ea86449727 100644
--- a/mlir/include/mlir-c/RegisterEverything.h
+++ b/mlir/include/mlir-c/RegisterEverything.h
@@ -28,8 +28,8 @@ MLIR_CAPI_EXPORTED void mlirRegisterAllDialects(MlirDialectRegistry registry);
 /// Register all translations to LLVM IR for dialects that can support it.
 MLIR_CAPI_EXPORTED void mlirRegisterAllLLVMTranslations(MlirContext context);
 
-// /// Register all compiler passes of MLIR.
-// MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void);
+/// Register all compiler passes of MLIR.
+MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void);
 
 #ifdef __cplusplus
 }
diff --git a/mlir/lib/CAPI/Dialect/CMakeLists.txt b/mlir/lib/CAPI/Dialect/CMakeLists.txt
index 8b1321d0c125a..1deafe41a4b7a 100644
--- a/mlir/lib/CAPI/Dialect/CMakeLists.txt
+++ b/mlir/lib/CAPI/Dialect/CMakeLists.txt
@@ -258,9 +258,6 @@ add_mlir_upstream_c_api_library(MLIRCAPITransformDialect
   LINK_LIBS PUBLIC
   MLIRCAPIIR
   MLIRTransformDialect
-  MLIRTransformDebugExtension
-  MLIRTransformLoopExtension
-  MLIRTransformTuneExtension
 )
 
 add_mlir_upstream_c_api_library(MLIRCAPITransformDialectTransforms
@@ -272,94 +269,6 @@ add_mlir_upstream_c_api_library(MLIRCAPITransformDialectTransforms
   MLIRTransformDialectTransforms
 )
 
-add_mlir_upstream_c_api_library(MLIRCAPIAffineTransformDialectExtension
-  TransformAffine.cpp
-
-  PARTIAL_SOURCES_INTENDED
-  LINK_LIBS PUBLIC
-  MLIRAffineTransformOps
-)
-
-add_mlir_upstream_c_api_library(MLIRCAPIBufferizationTransformDialectExtension
-  TransformBufferization.cpp
-
-  PARTIAL_SOURCES_INTENDED
-  LINK_LIBS PUBLIC
-  MLIRBufferizationTransformOps
-)
-
-add_mlir_upstream_c_api_library(MLIRCAPIGPUTransformDialectExtension
-  TransformGPU.cpp
-
-  PARTIAL_SOURCES_INTENDED
-  LINK_LIBS PUBLIC
-  MLIRGPUTransformOps
-)
-
-add_mlir_upstream_c_api_library(MLIRCAPILinalgTransformDialectExtension
-  TransformLinalg.cpp
-
-  PARTIAL_SOURCES_INTENDED
-  LINK_LIBS PUBLIC
-  MLIRLinalgTransformOps
-)
-
-add_mlir_upstream_c_api_library(MLIRCAPIMemRefTransformDialectExtension
-  TransformMemRef.cpp
-
-  PARTIAL_SOURCES_INTENDED
-  LINK_LIBS PUBLIC
-  MLIRMemRefTransformOps
-)
-
-add_mlir_upstream_c_api_library(MLIRCAPINVGPUTransformDialectExtension
-  TransformNVGPU.cpp
-
-  PARTIAL_SOURCES_INTENDED
-  LINK_LIBS PUBLIC
-  MLIRNVGPUTransformOps
-)
-
-add_mlir_upstream_c_api_library(MLIRCAPIPDLTransformDialectExtension
-  TransformPDL.cpp
-
-  PARTIAL_SOURCES_INTENDED
-  LINK_LIBS PUBLIC
-  MLIRTransformPDLExtension
-)
-
-add_mlir_upstream_c_api_library(MLIRCAPISMTTransformDialectExtension
-  TransformSMT.cpp
-
-  PARTIAL_SOURCES_INTENDED
-  LINK_LIBS PUBLIC
-  MLIRTransformSMTExtension
-)
-
-add_mlir_upstream_c_api_library(MLIRCAPISparseTensorTransformDialectExtension
-  TransformSparseTensor.cpp
-
-  PARTIAL_SOURCES_INTENDED
-  LINK_LIBS PUBLIC
-  MLIRSparseTensorTransformOps
-)
-
-add_mlir_upstream_c_api_library(MLIRCAPITensorTransformDialectExtension
-  TransformTensor.cpp
-
-  PARTIAL_SOURCES_INTENDED
-  LINK_LIBS PUBLIC
-  MLIRTensorTransformOps
-)
-
-add_mlir_upstream_c_api_library(MLIRCAPIVectorTransformDialectExtension
-  TransformVector.cpp
-
-  PARTIAL_SOURCES_INTENDED
-  LINK_LIBS PUBLIC
-  MLIRVectorTransformOps
-)
-
 add_mlir_upstream_c_api_library(MLIRCAPIQuant
   Quant.cpp
 
diff --git a/mlir/lib/CAPI/Dialect/Transform.cpp b/mlir/lib/CAPI/Dialect/Transform.cpp
index 8942d0a00cde3..5fd773572bd3c 100644
--- a/mlir/lib/CAPI/Dialect/Transform.cpp
+++ b/mlir/lib/CAPI/Dialect/Transform.cpp
@@ -9,11 +9,8 @@
 #include "mlir-c/Dialect/Transform.h"
 #include "mlir-c/Support.h"
 #include "mlir/CAPI/Registration.h"
-#include "mlir/Dialect/Transform/DebugExtension/DebugExtension.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
-#include "mlir/Dialect/Transform/LoopExtension/LoopExtension.h"
-#include "mlir/Dialect/Transform/TuneExtension/TuneExtension.h"
 
 using namespace mlir;
 
@@ -109,15 +106,3 @@ MlirType mlirTransformParamTypeGet(MlirContext ctx, MlirType type) {
 MlirType mlirTransformParamTypeGetType(MlirType type) {
   return wrap(cast<transform::ParamType>(unwrap(type)).getType());
 }
-
-void mlirDebugRegisterTransformDialectExtension(MlirDialectRegistry registry) {
-  mlir::transform::registerDebugExtension(*unwrap(registry));
-}
-
-void mlirLoopRegisterTransformDialectExtension(MlirDialectRegistry registry) {
-  mlir::transform::registerLoopExtension(*unwrap(registry));
-}
-
-void mlirTuneRegisterTransformDialectExtension(MlirDialectRegistry registry) {
-  mlir::transform::registerTuneExtension(*unwrap(registry));
-}
diff --git a/mlir/lib/CAPI/Dialect/TransformAffine.cpp b/mlir/lib/CAPI/Dialect/TransformAffine.cpp
deleted file mode 100644
index 6fd7744e8c2c9..0000000000000
--- a/mlir/lib/CAPI/Dialect/TransformAffine.cpp
+++ /dev/null
@@ -1,16 +0,0 @@
-//===- TransformAffine.cpp - C Interface for Transform affine extension ---===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir-c/Dialect/Affine.h"
-#include "mlir/CAPI/Registration.h"
-#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h"
-
-MLIR_CAPI_EXPORTED void
-mlirAffineRegisterTransformDialectExtension(MlirDialectRegistry registry) {
-  mlir::affine::registerTransformDialectExtension(*unwrap(registry));
-}
diff --git a/mlir/lib/CAPI/Dialect/TransformBufferization.cpp b/mlir/lib/CAPI/Dialect/TransformBufferization.cpp
deleted file mode 100644
index f1bd61be10880..0000000000000
--- a/mlir/lib/CAPI/Dialect/TransformBufferization.cpp
+++ /dev/null
@@ -1,16 +0,0 @@
-//==- TransformBufferization.cpp - C Interface for bufferization extension -==//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir-c/Dialect/Bufferization.h"
-#include "mlir/CAPI/Registration.h"
-#include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
-
-void mlirBufferizationRegisterTransformDialectExtension(
-    MlirDialectRegistry registry) {
-  mlir::bufferization::registerTransformDialectExtension(*unwrap(registry));
-}
diff --git a/mlir/lib/CAPI/Dialect/TransformGPU.cpp b/mlir/lib/CAPI/Dialect/TransformGPU.cpp
deleted file mode 100644
index 51187a28fa752..0000000000000
--- a/mlir/lib/CAPI/Dialect/TransformGPU.cpp
+++ /dev/null
@@ -1,15 +0,0 @@
-//===- TransformGPU.cpp - C Interface for Transform GPU extension ---------===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir-c/Dialect/GPU.h"
-#include "mlir/CAPI/Registration.h"
-#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
-
-void mlirGPURegisterTransformDialectExtension(MlirDialectRegistry registry) {
-  mlir::gpu::registerTransformDialectExtension(*unwrap(registry));
-}
diff --git a/mlir/lib/CAPI/Dialect/TransformLinalg.cpp b/mlir/lib/CAPI/Dialect/TransformLinalg.cpp
deleted file mode 100644
index 5e62a80011a1b..0000000000000
--- a/mlir/lib/CAPI/Dialect/TransformLinalg.cpp
+++ /dev/null
@@ -1,15 +0,0 @@
-//===- TransformLinalg.cpp - C Interface for Transform Linalg extension ---===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir-c/Dialect/Linalg.h"
-#include "mlir/CAPI/Registration.h"
-#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h"
-
-void mlirLinalgRegisterTransformDialectExtension(MlirDialectRegistry registry) {
-  mlir::linalg::registerTransformDialectExtension(*unwrap(registry));
-}
diff --git a/mlir/lib/CAPI/Dialect/TransformMemRef.cpp b/mlir/lib/CAPI/Dialect/TransformMemRef.cpp
deleted file mode 100644
index 9b80295748514..0000000000000
--- a/mlir/lib/CAPI/Dialect/TransformMemRef.cpp
+++ /dev/null
@@ -1,15 +0,0 @@
-//===- TransformMemRef.cpp - C Interface for Transform MemRef extension ---===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir-c/Dialect/MemRef.h"
-#include "mlir/CAPI/Registration.h"
-#include "mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.h"
-
-void mlirMemRefRegisterTransformDialectExtension(MlirDialectRegistry registry) {
-  mlir::memref::registerTransformDialectExtension(*unwrap(registry));
-}
diff --git a/mlir/lib/CAPI/Dialect/TransformNVGPU.cpp b/mlir/lib/CAPI/Dialect/TransformNVGPU.cpp
deleted file mode 100644
index 9159dc319aeb1..0000000000000
--- a/mlir/lib/CAPI/Dialect/TransformNVGPU.cpp
+++ /dev/null
@@ -1,15 +0,0 @@
-//===- TransformNVGPU.cpp - C Interface for Transform NVGPU extension -----===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir-c/Dialect/NVGPU.h"
-#include "mlir/CAPI/Registration.h"
-#include "mlir/Dialect/NVGPU/TransformOps/NVGPUTransformOps.h"
-
-void mlirNVGPURegisterTransformDialectExtension(MlirDialectRegistry registry) {
-  mlir::nvgpu::registerTransformDialectExtension(*unwrap(registry));
-}
diff --git a/mlir/lib/CAPI/Dialect/TransformPDL.cpp b/mlir/lib/CAPI/Dialect/TransformPDL.cpp
deleted file mode 100644
index 550bcabee55ec..0000000000000
--- a/mlir/lib/CAPI/Dialect/TransformPDL.cpp
+++ /dev/null
@@ -1,15 +0,0 @@
-//===- TransformSMT.cpp - C Interface for Transform SMT extension ---------===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir-c/Dialect/PDL.h"
-#include "mlir/CAPI/Registration.h"
-#include "mlir/Dialect/Transform/PDLExtension/PDLExtension.h"
-
-void mlirPDLRegisterTransformDialectExtension(MlirDialectRegistry registry) {
-  mlir::transform::registerPDLExtension(*unwrap(registry));
-}
diff --git a/mlir/lib/CAPI/Dialect/TransformSMT.cpp b/mlir/lib/CAPI/Dialect/TransformSMT.cpp
deleted file mode 100644
index 449129dac99e5..0000000000000
--- a/mlir/lib/CAPI/Dialect/TransformSMT.cpp
+++ /dev/null
@@ -1,15 +0,0 @@
-//===- TransformSMT.cpp - C Interface for Transform SMT extension ---------===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir-c/Dialect/SMT.h"
-#include "mlir/CAPI/Registration.h"
-#include "mlir/Dialect/Transform/SMTExtension/SMTExtension.h"
-
-void mlirSMTRegisterTransformDialectExtension(MlirDialectRegistry registry) {
-  mlir::transform::registerSMTExtension(*unwrap(registry));
-}
diff --git a/mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp b/mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp
deleted file mode 100644
index 3e5b91099a6fb..0000000000000
--- a/mlir/lib/CAPI/Dialect/TransformSparseTensor.cpp
+++ /dev/null
@@ -1,16 +0,0 @@
-//===- TransformSparseTensor.cpp - C Interface for SparseTensor extension -===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir-c/Dialect/SparseTensor.h"
-#include "mlir/CAPI/Registration.h"
-#include "mlir/Dialect/SparseTensor/TransformOps/SparseTensorTransformOps.h"
-
-void mlirSparseTensorRegisterTransformDialectExtension(
-    MlirDialectRegistry registry) {
-  mlir::sparse_tensor::registerTransformDialectExtension(*unwrap(registry));
-}
diff --git a/mlir/lib/CAPI/Dialect/TransformTensor.cpp b/mlir/lib/CAPI/Dialect/TransformTensor.cpp
deleted file mode 100644
index fc7a30b5b237e..0000000000000
--- a/mlir/lib/CAPI/Dialect/TransformTensor.cpp
+++ /dev/null
@@ -1,15 +0,0 @@
-//===- TransformTensor.cpp - C Interface for Transform tensor extensio ----===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir-c/Dialect/Tensor.h"
-#include "mlir/CAPI/Registration.h"
-#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
-
-void mlirTensorRegisterTransformDialectExtension(MlirDialectRegistry registry) {
-  mlir::tensor::registerTransformDialectExtension(*unwrap(registry));
-}
diff --git a/mlir/lib/CAPI/Dialect/TransformVector.cpp b/mlir/lib/CAPI/Dialect/TransformVector.cpp
deleted file mode 100644
index 027d37ea34e79..0000000000000
--- a/mlir/lib/CAPI/Dialect/TransformVector.cpp
+++ /dev/null
@@ -1,15 +0,0 @@
-//===- TransformVector.cpp - C Interface for Transform Vector extension ---===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-
-#include "mlir-c/Dialect/Vector.h"
-#include "mlir/CAPI/Registration.h"
-#include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
-
-void mlirVectorRegisterTransformDialectExtension(MlirDialectRegistry registry) {
-  mlir::vector::registerTransformDialectExtension(*unwrap(registry));
-}
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 9a45c5026f60e..a29709ee68740 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -204,11 +204,7 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/pdl.py
   DIALECT_NAME transform
-  EXTENSION_NAME transform_pdl_extension
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPITransformDialect
-    MLIRCAPIPDLTransformDialectExtension
-)
+  EXTENSION_NAME transform_pdl_extension)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -217,11 +213,7 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/smt.py
   DIALECT_NAME transform
-  EXTENSION_NAME transform_smt_extension
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPITransformDialect
-    MLIRCAPISMTTransformDialectExtension
-)
+  EXTENSION_NAME transform_smt_extension)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -277,11 +269,7 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/bufferization.py
   DIALECT_NAME transform
-  EXTENSION_NAME bufferization_transform
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPITransformDialect
-    MLIRCAPIBufferizationTransformDialectExtension
-  )
+  EXTENSION_NAME bufferization_transform)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -290,11 +278,7 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/gpu.py
   DIALECT_NAME transform
-  EXTENSION_NAME gpu_transform
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPITransformDialect
-    MLIRCAPIGPUTransformDialectExtension
-)
+  EXTENSION_NAME gpu_transform)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -303,10 +287,7 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/loop.py
   DIALECT_NAME transform
-  EXTENSION_NAME loop_transform
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPITransformDialect
-)
+  EXTENSION_NAME loop_transform)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -315,11 +296,7 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/memref.py
   DIALECT_NAME transform
-  EXTENSION_NAME memref_transform
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPITransformDialect
-    MLIRCAPIMemRefTransformDialectExtension
-)
+  EXTENSION_NAME memref_transform)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -328,11 +305,7 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/nvgpu.py
   DIALECT_NAME transform
-  EXTENSION_NAME nvgpu_transform
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPITransformDialect
-    MLIRCAPINVGPUTransformDialectExtension
-)
+  EXTENSION_NAME nvgpu_transform)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -344,9 +317,6 @@ declare_mlir_dialect_extension_python_bindings(
   EXTENSION_NAME structured_transform
   GEN_ENUM_BINDINGS_TD_FILE
     "../../include/mlir/Dialect/Linalg/TransformOps/LinalgTransformEnums.td"
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPITransformDialect
-    MLIRCAPILinalgTransformDialectExtension
 )
 
 declare_mlir_dialect_extension_python_bindings(
@@ -356,11 +326,7 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/sparse_tensor.py
   DIALECT_NAME transform
-  EXTENSION_NAME sparse_tensor_transform
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPITransformDialect
-    MLIRCAPISparseTensorTransformDialectExtension
-)
+  EXTENSION_NAME sparse_tensor_transform)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -369,11 +335,7 @@ declare_mlir_dialect_extension_python_bindings(
   SOURCES
     dialects/transform/tensor.py
   DIALECT_NAME transform
-  EXTENSION_NAME tensor_transform
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPITransformDialect
-    MLIRCAPITensorTransformDialectExtension
-)
+  EXTENSION_NAME tensor_transform)
 
 declare_mlir_dialect_extension_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -385,9 +347,6 @@ declare_mlir_dialect_extension_python_bindings(
   EXTENSION_NAME vector_transform
   GEN_ENUM_BINDINGS_TD_FILE
     "../../include/mlir/Dialect/Vector/Transforms/VectorTransformsBase.td"
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPITransformDialect
-    MLIRCAPIVectorTransformDialectExtension
 )
 
 declare_mlir_dialect_python_bindings(
@@ -888,7 +847,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
     MLIRCAPIIR
     MLIRCAPISMT
     MLIRCAPIExportSMTLIB
-    MLIRCAPISMTTransformDialectExtension
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
diff --git a/mlir/python/mlir/_mlir_libs/_capi.py.in b/mlir/python/mlir/_mlir_libs/_capi.py.in
index 2313ecc4f9823..f822c00c06038 100644
--- a/mlir/python/mlir/_mlir_libs/_capi.py.in
+++ b/mlir/python/mlir/_mlir_libs/_capi.py.in
@@ -37,14 +37,3 @@ def register_dialect(dialect_handle_capi_name, dialect_registry):
     dialect_handle_capi.restype = ctypes.c_void_p
     handle = dialect_handle_capi()
     dialect_registry.insert_dialect(handle)
-
-
-def register_transform_dialect_extension(registration_capi_name, dialect_registry):
-    if not hasattr(_capi, registration_capi_name):
-        raise RuntimeError(f"missing {registration_capi_name} API")
-    registration_capi = getattr(_capi, registration_capi_name)
-    registration_capi.argtypes = [ctypes.c_void_p]
-    dialect_registry_ptr = PyCapsule_GetPointer(
-        dialect_registry._CAPIPtr, MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY
-    )
-    registration_capi(dialect_registry_ptr)
diff --git a/mlir/python/mlir/dialects/transform/bufferization.py b/mlir/python/mlir/dialects/transform/bufferization.py
index 228c9dab52b8d..485a8a36b6305 100644
--- a/mlir/python/mlir/dialects/transform/bufferization.py
+++ b/mlir/python/mlir/dialects/transform/bufferization.py
@@ -15,15 +15,6 @@
 from enum import Enum
 from typing import Optional, overload, Union
 
-from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
-from ..._mlir_libs._capi import (
-    register_transform_dialect_extension as _register_transform_dialect_extension,
-)
-
-_register_transform_dialect_extension(
-    "mlirBufferizationRegisterTransformDialectExtension", _get_dialect_registry()
-)
-
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class EmptyTensorToAllocTensorOp(EmptyTensorToAllocTensorOp):
diff --git a/mlir/python/mlir/dialects/transform/debug.py b/mlir/python/mlir/dialects/transform/debug.py
index 353e4f07c883a..f7c04268dc03d 100644
--- a/mlir/python/mlir/dialects/transform/debug.py
+++ b/mlir/python/mlir/dialects/transform/debug.py
@@ -15,15 +15,6 @@
 
 from typing import Union
 
-from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
-from ..._mlir_libs._capi import (
-    register_transform_dialect_extension as _register_transform_dialect_extension,
-)
-
-_register_transform_dialect_extension(
-    "mlirDebugRegisterTransformDialectExtension", _get_dialect_registry()
-)
-
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class EmitParamAsRemarkOp(EmitParamAsRemarkOp):
diff --git a/mlir/python/mlir/dialects/transform/gpu.py b/mlir/python/mlir/dialects/transform/gpu.py
index b35004cb4a09a..00cf0840eeae9 100644
--- a/mlir/python/mlir/dialects/transform/gpu.py
+++ b/mlir/python/mlir/dialects/transform/gpu.py
@@ -15,16 +15,6 @@
 from typing import Optional, Sequence, Union, overload
 
 
-from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
-from ..._mlir_libs._capi import (
-    register_transform_dialect_extension as _register_transform_dialect_extension,
-)
-
-_register_transform_dialect_extension(
-    "mlirGPURegisterTransformDialectExtension", _get_dialect_registry()
-)
-
-
 @_ods_cext.register_operation(_Dialect, replace=True)
 class MapForallToBlocks(MapForallToBlocks):
     """Specialization for MapForallToBlocks class."""
diff --git a/mlir/python/mlir/dialects/transform/loop.py b/mlir/python/mlir/dialects/transform/loop.py
index 3311ae47b4fa9..c4770b1c4067e 100644
--- a/mlir/python/mlir/dialects/transform/loop.py
+++ b/mlir/python/mlir/dialects/transform/loop.py
@@ -16,15 +16,6 @@
 
 from typing import Optional, Union
 
-from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
-from ..._mlir_libs._capi import (
-    register_transform_dialect_extension as _register_transform_dialect_extension,
-)
-
-_register_transform_dialect_extension(
-    "mlirLoopRegisterTransformDialectExtension", _get_dialect_registry()
-)
-
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class LoopOutlineOp(LoopOutlineOp):
diff --git a/mlir/python/mlir/dialects/transform/memref.py b/mlir/python/mlir/dialects/transform/memref.py
index f45d8054bfa58..56ea61eb817f8 100644
--- a/mlir/python/mlir/dialects/transform/memref.py
+++ b/mlir/python/mlir/dialects/transform/memref.py
@@ -14,15 +14,6 @@
 
 from typing import Optional, overload, Union
 
-from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
-from ..._mlir_libs._capi import (
-    register_transform_dialect_extension as _register_transform_dialect_extension,
-)
-
-_register_transform_dialect_extension(
-    "mlirMemRefRegisterTransformDialectExtension", _get_dialect_registry()
-)
-
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class MemRefAllocaToGlobalOp(MemRefAllocaToGlobalOp):
diff --git a/mlir/python/mlir/dialects/transform/nvgpu.py b/mlir/python/mlir/dialects/transform/nvgpu.py
index 3d6427d1b1ed9..74ba4c9aeb6c1 100644
--- a/mlir/python/mlir/dialects/transform/nvgpu.py
+++ b/mlir/python/mlir/dialects/transform/nvgpu.py
@@ -3,12 +3,3 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from .._nvgpu_transform_ops_gen import *
-
-from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
-from ..._mlir_libs._capi import (
-    register_transform_dialect_extension as _register_transform_dialect_extension,
-)
-
-_register_transform_dialect_extension(
-    "mlirNVGPURegisterTransformDialectExtension", _get_dialect_registry()
-)
diff --git a/mlir/python/mlir/dialects/transform/pdl.py b/mlir/python/mlir/dialects/transform/pdl.py
index 15b0c37837902..bb5fa7ffd3065 100644
--- a/mlir/python/mlir/dialects/transform/pdl.py
+++ b/mlir/python/mlir/dialects/transform/pdl.py
@@ -18,16 +18,6 @@
 from typing import Union
 
 
-from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
-from ..._mlir_libs._capi import (
-    register_transform_dialect_extension as _register_transform_dialect_extension,
-)
-
-_register_transform_dialect_extension(
-    "mlirPDLRegisterTransformDialectExtension", _get_dialect_registry()
-)
-
-
 @_ods_cext.register_operation(_Dialect, replace=True)
 class PDLMatchOp(PDLMatchOp):
     def __init__(
diff --git a/mlir/python/mlir/dialects/transform/smt.py b/mlir/python/mlir/dialects/transform/smt.py
index a4e8eb96d388c..1f0b7f066118c 100644
--- a/mlir/python/mlir/dialects/transform/smt.py
+++ b/mlir/python/mlir/dialects/transform/smt.py
@@ -15,16 +15,6 @@
     raise RuntimeError("Error loading imports from extension module") from e
 
 
-from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
-from ..._mlir_libs._capi import (
-    register_transform_dialect_extension as _register_transform_dialect_extension,
-)
-
-_register_transform_dialect_extension(
-    "mlirSMTRegisterTransformDialectExtension", _get_dialect_registry()
-)
-
-
 @_ods_cext.register_operation(_Dialect, replace=True)
 class ConstrainParamsOp(ConstrainParamsOp):
     def __init__(
diff --git a/mlir/python/mlir/dialects/transform/sparse_tensor.py b/mlir/python/mlir/dialects/transform/sparse_tensor.py
index efaf438e07b83..8b33270dc23a1 100644
--- a/mlir/python/mlir/dialects/transform/sparse_tensor.py
+++ b/mlir/python/mlir/dialects/transform/sparse_tensor.py
@@ -3,11 +3,3 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from .._sparse_tensor_transform_ops_gen import *
-from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
-from ..._mlir_libs._capi import (
-    register_transform_dialect_extension as _register_transform_dialect_extension,
-)
-
-_register_transform_dialect_extension(
-    "mlirSparseTensorRegisterTransformDialectExtension", _get_dialect_registry()
-)
diff --git a/mlir/python/mlir/dialects/transform/structured.py b/mlir/python/mlir/dialects/transform/structured.py
index 85833d32ba55e..bf40cc532065d 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -28,15 +28,6 @@
 
 from typing import List, Optional, Sequence, Union, overload
 
-from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
-from ..._mlir_libs._capi import (
-    register_transform_dialect_extension as _register_transform_dialect_extension,
-)
-
-_register_transform_dialect_extension(
-    "mlirLinalgRegisterTransformDialectExtension", _get_dialect_registry()
-)
-
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class BufferizeToAllocationOp(BufferizeToAllocationOp):
diff --git a/mlir/python/mlir/dialects/transform/tensor.py b/mlir/python/mlir/dialects/transform/tensor.py
index 904f57827146d..4eb30398f0872 100644
--- a/mlir/python/mlir/dialects/transform/tensor.py
+++ b/mlir/python/mlir/dialects/transform/tensor.py
@@ -14,15 +14,6 @@
 
 from typing import Optional, overload, Union
 
-from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
-from ..._mlir_libs._capi import (
-    register_transform_dialect_extension as _register_transform_dialect_extension,
-)
-
-_register_transform_dialect_extension(
-    "mlirTensorRegisterTransformDialectExtension", _get_dialect_registry()
-)
-
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class MakeLoopIndependentOp(MakeLoopIndependentOp):
diff --git a/mlir/python/mlir/dialects/transform/tune.py b/mlir/python/mlir/dialects/transform/tune.py
index 6eadd9b4e9d5d..f63f88a382422 100644
--- a/mlir/python/mlir/dialects/transform/tune.py
+++ b/mlir/python/mlir/dialects/transform/tune.py
@@ -25,15 +25,6 @@
 
 from typing import Union
 
-from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
-from ..._mlir_libs._capi import (
-    register_transform_dialect_extension as _register_transform_dialect_extension,
-)
-
-_register_transform_dialect_extension(
-    "mlirTuneRegisterTransformDialectExtension", _get_dialect_registry()
-)
-
 
 @_ods_cext.register_operation(_Dialect, replace=True)
 class KnobOp(KnobOp):
diff --git a/mlir/python/mlir/dialects/transform/vector.py b/mlir/python/mlir/dialects/transform/vector.py
index 1d7620188da83..af2435cb26cc4 100644
--- a/mlir/python/mlir/dialects/transform/vector.py
+++ b/mlir/python/mlir/dialects/transform/vector.py
@@ -4,11 +4,3 @@
 
 from .._vector_transform_enum_gen import *
 from .._vector_transform_ops_gen import *
-from ..._mlir_libs import get_dialect_registry as _get_dialect_registry
-from ..._mlir_libs._capi import (
-    register_transform_dialect_extension as _register_transform_dialect_extension,
-)
-
-_register_transform_dialect_extension(
-    "mlirVectorRegisterTransformDialectExtension", _get_dialect_registry()
-)



More information about the Mlir-commits mailing list