undo transform extension stuffusers/makslevental/specific-reg

61 files changed, 215 insertions, 529 deletions

diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake
index ceebc931e60d..d8b6d493f985 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/examples/standalone/pyproject.toml b/mlir/examples/standalone/pyproject.toml
index 5a1e6e86513c..90535030df81 100644
--- a/mlir/examples/standalone/pyproject.toml
+++ b/mlir/examples/standalone/pyproject.toml
@@ -37,6 +37,8 @@ cmake.source-dir = "."
 # This is for installing/distributing the python bindings target and only the python bindings target.
 build.targets = ["StandalonePythonModules"]
 install.components = ["StandalonePythonModules"]
+# The default is true but make it explicit to highlight that this option exists (turn off for debug symbols).
+install.strip = true
 
 [tool.scikit-build.cmake.define]
 # Optional
@@ -51,6 +53,9 @@ LLVM_USE_LINKER = { env = "LLVM_USE_LINKER", default = "" }
 CMAKE_VISIBILITY_INLINES_HIDDEN = "ON"
 CMAKE_C_VISIBILITY_PRESET = "hidden"
 CMAKE_CXX_VISIBILITY_PRESET = "hidden"
+# Disables generation of "version soname" (i.e. libFoo.so.<version>),
+# which causes pure duplication of various shlibs for Python wheels.
+CMAKE_PLATFORM_NO_VERSIONED_SONAME = "ON"
 
 # Non-optional (alternatively you could use CMAKE_PREFIX_PATH here).
 MLIR_DIR = { env = "MLIR_DIR", default = "" }
diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt
index 905c94493975..505add523208 100644
--- a/mlir/examples/standalone/python/CMakeLists.txt
+++ b/mlir/examples/standalone/python/CMakeLists.txt
@@ -138,6 +138,7 @@ set(_declared_sources
   StandalonePythonSources
   MLIRPythonSources.Core
   MLIRPythonSources.Dialects.builtin
+  MLIRPythonSources.Dialects.arith
 )
 # For an external projects build, the MLIRPythonExtension.Core.type_stub_gen
 # target already exists and can just be added to DECLARED_SOURCES.
diff --git a/mlir/include/mlir-c/Dialect/Transform.h b/mlir/include/mlir-c/Dialect/Transform.h
index 1b5232c106ed..02c99b592188 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/RegisterAllExtensions.h b/mlir/include/mlir-c/RegisterAllExtensions.h
index 7238334ed5b8..dc2dcfbb2809 100644
--- a/mlir/include/mlir-c/RegisterAllExtensions.h
+++ b/mlir/include/mlir-c/RegisterAllExtensions.h
@@ -7,8 +7,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_C_REGISTER_Extensions_H
-#define MLIR_C_REGISTER_Extensions_H
+#ifndef MLIR_C_REGISTER_EXTENSIONS_H
+#define MLIR_C_REGISTER_EXTENSIONS_H
 
 #include "mlir-c/IR.h"
 
@@ -23,4 +23,4 @@ MLIR_CAPI_EXPORTED void mlirRegisterAllExtensions(MlirDialectRegistry registry);
 }
 #endif
 
-#endif // MLIR_C_REGISTER_Extensions_H
+#endif // MLIR_C_REGISTER_EXTENSIONS_H
diff --git a/mlir/include/mlir-c/RegisterAllExternalModels.h b/mlir/include/mlir-c/RegisterAllExternalModels.h
new file mode 100644
index 000000000000..baa8a742cd01
--- /dev/null
+++ b/mlir/include/mlir-c/RegisterAllExternalModels.h
@@ -0,0 +1,27 @@
+//===-- mlir-c/RegisterAllExternalModels.h - Register all MLIR model ----*-===//
+//
+// 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_EXTERNAL_MODELS_H
+#define MLIR_C_REGISTER_EXTERNAL_MODELS_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/// Register all compiler External models of MLIR.
+MLIR_CAPI_EXPORTED void
+mlirRegisterAllExternalModels(MlirDialectRegistry registry);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_REGISTER_EXTERNAL_MODELS_H
diff --git a/mlir/include/mlir-c/RegisterEverything.h b/mlir/include/mlir-c/RegisterEverything.h
index e759f6b18877..ea2ea8644972 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/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index c3d28733c2ce..21e60d5550a5 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3139,6 +3139,10 @@ void mlir::python::populateIRCore(nb::module_ &m) {
   //----------------------------------------------------------------------------
 
   nb::class_<MlirDialectHandle>(m, "DialectHandle")
+      .def_prop_ro_static(
+          "_capsule_name",
+          [](nb::handle &) { return MLIR_PYTHON_CAPSULE_DIALECT_HANDLE; },
+          nb::sig("def _capsule_name(/) -> str"))
       .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR,
                   &createMlirDialectHandleFromCapsule);
 
@@ -3147,6 +3151,10 @@ void mlir::python::populateIRCore(nb::module_ &m) {
   //----------------------------------------------------------------------------
   nb::class_<PyDialectRegistry>(m, "DialectRegistry")
       .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyDialectRegistry::getCapsule)
+      .def_prop_ro_static(
+          "_capsule_name",
+          [](nb::handle &) { return MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY; },
+          nb::sig("def _capsule_name(/) -> str"))
       .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR,
                   &PyDialectRegistry::createFromCapsule)
       .def(nb::init<>())
diff --git a/mlir/lib/Bindings/Python/RegisterEverything.cpp b/mlir/lib/Bindings/Python/RegisterEverything.cpp
index d773c15fbfc3..e825bceace87 100644
--- a/mlir/lib/Bindings/Python/RegisterEverything.cpp
+++ b/mlir/lib/Bindings/Python/RegisterEverything.cpp
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir-c/RegisterAllExtensions.h"
+#include "mlir-c/RegisterAllExternalModels.h"
 #include "mlir-c/RegisterAllLLVMTranslations.h"
 #include "mlir-c/RegisterAllPasses.h"
 #include "mlir/Bindings/Python/Nanobind.h"
@@ -16,6 +17,9 @@ NB_MODULE(_mlirRegisterEverything, m) {
   m.doc() =
       "MLIR All Upstream Extensions, Translations and Passes Registration";
 
+  m.def("register_external_models", [](MlirDialectRegistry registry) {
+    mlirRegisterAllExternalModels(registry);
+  });
   m.def("register_extensions", [](MlirDialectRegistry registry) {
     mlirRegisterAllExtensions(registry);
   });
diff --git a/mlir/lib/CAPI/Dialect/CMakeLists.txt b/mlir/lib/CAPI/Dialect/CMakeLists.txt
index 8b1321d0c125..1deafe41a4b7 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 8942d0a00cde..5fd773572bd3 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 6fd7744e8c2c..000000000000
--- 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 f1bd61be1088..000000000000
--- 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 51187a28fa75..000000000000
--- 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 5e62a80011a1..000000000000
--- 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 9b8029574851..000000000000
--- 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 9159dc319aeb..000000000000
--- 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 550bcabee55e..000000000000
--- 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 449129dac99e..000000000000
--- 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 3e5b91099a6f..000000000000
--- 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 fc7a30b5b237..000000000000
--- 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 027d37ea34e7..000000000000
--- 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/lib/CAPI/RegisterEverything/CMakeLists.txt b/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt
index 2ac01bc4c84f..473c3844538c 100644
--- a/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt
+++ b/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt
@@ -43,4 +43,13 @@ add_mlir_upstream_c_api_library(MLIRCAPIRegisterAllExtensions
   LINK_LIBS PUBLIC
   MLIRCAPIIR
   MLIRRegisterAllExtensions
-)
\ No newline at end of file
+)
+
+add_mlir_upstream_c_api_library(MLIRCAPIRegisterAllExternalModels
+  RegisterAllExternalModels.cpp
+  PARTIAL_SOURCES_INTENDED
+
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+  MLIRRegisterAllDialects
+)
diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp
index 9cc9fd394e72..9be46bfff012 100644
--- a/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp
+++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp
@@ -1,5 +1,4 @@
-//===- RegisterAllExtensions.cpp - Register all MLIR entities
-//-----------------===//
+//===- RegisterAllExtensions.cpp - 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.
@@ -11,6 +10,9 @@
 #include "mlir/CAPI/IR.h"
 #include "mlir/InitAllExtensions.h"
 
-void mlirRegisterAllExtensions(MlirDialectRegistry registry) {
-  mlir::registerAllExtensions(*unwrap(registry));
+using namespace mlir;
+
+void mlirRegisterAllExtensions(MlirDialectRegistry mlirRegistry) {
+  mlir::DialectRegistry *registry = unwrap(mlirRegistry);
+  mlir::registerAllExtensions(*registry);
 }
diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllExternalModels.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllExternalModels.cpp
new file mode 100644
index 000000000000..cee3d717ccb3
--- /dev/null
+++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllExternalModels.cpp
@@ -0,0 +1,101 @@
+//===- RegisterAllExtensions.cpp - 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/RegisterAllExternalModels.h"
+#include "mlir/CAPI/IR.h"
+#include "mlir/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/Transforms/BufferDeallocationOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/Transforms/BufferViewFlowOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/Transforms/ShardingInterfaceImpl.h"
+#include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/ControlFlow/Transforms/BufferDeallocationOpInterfaceImpl.h"
+#include "mlir/Dialect/ControlFlow/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/GPU/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/GPU/Transforms/BufferDeallocationOpInterfaceImpl.h"
+#include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
+#include "mlir/Dialect/Linalg/Transforms/AllInterfaces.h"
+#include "mlir/Dialect/Linalg/Transforms/RuntimeOpVerification.h"
+#include "mlir/Dialect/MLProgram/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/MemRef/IR/MemRefMemorySlot.h"
+#include "mlir/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.h"
+#include "mlir/Dialect/MemRef/Transforms/BufferViewFlowOpInterfaceImpl.h"
+#include "mlir/Dialect/MemRef/Transforms/RuntimeOpVerification.h"
+#include "mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/SCF/Transforms/BufferDeallocationOpInterfaceImpl.h"
+#include "mlir/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/SparseTensor/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/TransformOps/TensorTransformOps.h"
+#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Tensor/Transforms/RuntimeOpVerification.h"
+#include "mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h"
+#include "mlir/Dialect/Tosa/IR/ShardingInterfaceImpl.h"
+#include "mlir/Dialect/Vector/IR/ValueBoundsOpInterfaceImpl.h"
+#include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.h"
+#include "mlir/Interfaces/CastInterfaces.h"
+#include "mlir/Target/LLVM/NVVM/Target.h"
+#include "mlir/Target/LLVM/ROCDL/Target.h"
+#include "mlir/Target/LLVM/XeVM/Target.h"
+#include "mlir/Target/SPIRV/Target.h"
+
+using namespace mlir;
+
+void mlirRegisterAllExternalModels(MlirDialectRegistry mlirRegistry) {
+  mlir::DialectRegistry *registry = unwrap(mlirRegistry);
+  // Register all external models.
+  affine::registerValueBoundsOpInterfaceExternalModels(*registry);
+  arith::registerBufferDeallocationOpInterfaceExternalModels(*registry);
+  arith::registerBufferizableOpInterfaceExternalModels(*registry);
+  arith::registerBufferViewFlowOpInterfaceExternalModels(*registry);
+  arith::registerShardingInterfaceExternalModels(*registry);
+  arith::registerValueBoundsOpInterfaceExternalModels(*registry);
+  bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
+      *registry);
+  builtin::registerCastOpInterfaceExternalModels(*registry);
+  cf::registerBufferizableOpInterfaceExternalModels(*registry);
+  cf::registerBufferDeallocationOpInterfaceExternalModels(*registry);
+  gpu::registerBufferDeallocationOpInterfaceExternalModels(*registry);
+  gpu::registerValueBoundsOpInterfaceExternalModels(*registry);
+  LLVM::registerInlinerInterface(*registry);
+  NVVM::registerInlinerInterface(*registry);
+  linalg::registerAllDialectInterfaceImplementations(*registry);
+  linalg::registerRuntimeVerifiableOpInterfaceExternalModels(*registry);
+  memref::registerAllocationOpInterfaceExternalModels(*registry);
+  memref::registerBufferViewFlowOpInterfaceExternalModels(*registry);
+  memref::registerRuntimeVerifiableOpInterfaceExternalModels(*registry);
+  memref::registerValueBoundsOpInterfaceExternalModels(*registry);
+  memref::registerMemorySlotExternalModels(*registry);
+  ml_program::registerBufferizableOpInterfaceExternalModels(*registry);
+  scf::registerBufferDeallocationOpInterfaceExternalModels(*registry);
+  scf::registerBufferizableOpInterfaceExternalModels(*registry);
+  scf::registerValueBoundsOpInterfaceExternalModels(*registry);
+  shape::registerBufferizableOpInterfaceExternalModels(*registry);
+  sparse_tensor::registerBufferizableOpInterfaceExternalModels(*registry);
+  tensor::registerBufferizableOpInterfaceExternalModels(*registry);
+  tensor::registerFindPayloadReplacementOpInterfaceExternalModels(*registry);
+  tensor::registerInferTypeOpInterfaceExternalModels(*registry);
+  tensor::registerRuntimeVerifiableOpInterfaceExternalModels(*registry);
+  tensor::registerSubsetOpInterfaceExternalModels(*registry);
+  tensor::registerTilingInterfaceExternalModels(*registry);
+  tensor::registerValueBoundsOpInterfaceExternalModels(*registry);
+  tosa::registerShardingInterfaceExternalModels(*registry);
+  vector::registerBufferizableOpInterfaceExternalModels(*registry);
+  vector::registerSubsetOpInterfaceExternalModels(*registry);
+  vector::registerValueBoundsOpInterfaceExternalModels(*registry);
+  NVVM::registerNVVMTargetInterfaceExternalModels(*registry);
+  ROCDL::registerROCDLTargetInterfaceExternalModels(*registry);
+  spirv::registerSPIRVTargetInterfaceExternalModels(*registry);
+  xevm::registerXeVMTargetInterfaceExternalModels(*registry);
+}
diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp
index c5dc64ca1783..0eeeaa279aef 100644
--- a/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp
+++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp
@@ -1,4 +1,4 @@
-//===- RegisterAllPasses.cpp - Register all MLIR entities -----------------===//
+//===- RegisterAllPasses.cpp - Register all MLIR 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.
diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp
index 159dff0b86e4..88a189cf7360 100644
--- a/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp
+++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp
@@ -1,4 +1,4 @@
-//===- RegisterAllPasses.cpp - Register all MLIR entities -----------------===//
+//===- RegisterAllPasses.cpp - Register all MLIR passes -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 9a45c5026f60..de736c77ebdd 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -23,6 +23,7 @@ declare_mlir_python_sources(MLIRPythonSources.Core.Python
   ADD_TO_PARENT MLIRPythonSources.Core
   SOURCES
     _mlir_libs/__init__.py
+    _mlir_libs/_capi.py
     _mlir_libs/_mlir/py.typed
     ir.py
     passmanager.py
@@ -204,11 +205,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 +214,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 +270,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 +279,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 +288,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 +297,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 +306,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 +318,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 +327,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 +336,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 +348,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(
@@ -691,6 +651,7 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
     MLIRCAPIRegisterAllPasses
     MLIRCAPIRegisterAllExtensions
     MLIRCAPIRegisterAllLLVMTranslations
+    MLIRCAPIRegisterAllExternalModels
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
@@ -888,7 +849,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
     MLIRCAPIIR
     MLIRCAPISMT
     MLIRCAPIExportSMTLIB
-    MLIRCAPISMTTransformDialectExtension
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
@@ -989,20 +949,8 @@ 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(${MLIR_PYTHON_CAPI_DYLIB_NAME}
+add_mlir_python_common_capi_library(MLIRPythonCAPI
   INSTALL_COMPONENT MLIRPythonModules
   INSTALL_DESTINATION "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs"
   OUTPUT_DIRECTORY "${MLIRPythonModules_ROOT_PREFIX}/_mlir_libs"
@@ -1120,8 +1068,7 @@ endif()
 
 set(_declared_sources
   MLIRPythonSources
-  MLIRPythonExtension.RegisterEverything
-  MLIRPythonCAPICTypesBinding)
+  MLIRPythonExtension.RegisterEverything)
 if(NOT CMAKE_CROSSCOMPILING)
   list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen)
 endif()
diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py
index 2c36397d2cf8..e36ed636b619 100644
--- a/mlir/python/mlir/_mlir_libs/__init__.py
+++ b/mlir/python/mlir/_mlir_libs/__init__.py
@@ -126,6 +126,9 @@ def _site_initialize():
         if hasattr(m, "register_extensions"):
             logger.debug("Registering extensions from initializer %r", m)
             m.register_extensions(get_dialect_registry())
+        if hasattr(m, "register_external_models"):
+            logger.debug("Registering external models from initializer %r", m)
+            m.register_external_models(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)
diff --git a/mlir/python/mlir/_mlir_libs/_capi.py.in b/mlir/python/mlir/_mlir_libs/_capi.py
index 2313ecc4f982..3d7836daa1c8 100644
--- a/mlir/python/mlir/_mlir_libs/_capi.py.in
+++ b/mlir/python/mlir/_mlir_libs/_capi.py
@@ -3,14 +3,33 @@
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 import ctypes
+import platform
 from pathlib import Path
 
-_capi = ctypes.CDLL(
-    str(
-        Path(__file__).parent
-        / "@CMAKE_SHARED_LIBRARY_PREFIX@@MLIR_PYTHON_CAPI_DYLIB_NAME@@CMAKE_SHARED_LIBRARY_SUFFIX@"
-    )
-)
+from . import _mlir, get_dialect_registry as _get_dialect_registry
+
+_get_dialect_registry()
+
+MLIR_PYTHON_CAPSULE_DIALECT_HANDLE = _mlir.ir.DialectHandle._capsule_name.encode()
+
+MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY = _mlir.ir.DialectRegistry._capsule_name.encode()
+
+if platform.system() == "Windows":
+    _ext_suffix = "dll"
+elif platform.system() == "Darwin":
+    _ext_suffix = "dylib"
+else:
+    _ext_suffix = "so"
+
+for fp in Path(__file__).parent.glob(f"*.{_ext_suffix}"):
+    if "CAPI" in fp.name:
+        _capi_dylib = fp
+        break
+else:
+    raise ValueError("Couldn't find CAPI dylib")
+
+
+_capi = ctypes.CDLL(str(Path(__file__).parent / _capi_dylib))
 
 PyCapsule_New = ctypes.pythonapi.PyCapsule_New
 PyCapsule_New.restype = ctypes.py_object
@@ -20,14 +39,6 @@ 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@.ir.DialectHandle._CAPIPtr"
-).encode()
-
-MLIR_PYTHON_CAPSULE_DIALECT_REGISTRY = (
-    "@MLIR_PYTHON_PACKAGE_PREFIX@.ir.DialectRegistry._CAPIPtr"
-).encode()
-
 
 def register_dialect(dialect_handle_capi_name, dialect_registry):
     if not hasattr(_capi, dialect_handle_capi_name):
@@ -37,14 +48,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 228c9dab52b8..485a8a36b630 100644
--- a/mlir/python/mlir/dialects/transform/bufferization.py
+++ b/mlir/python/mlir/dialects/transform/bufferization.py
@@ -15,15 +15,6 @@ except ImportError as e:
 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 353e4f07c883..f7c04268dc03 100644
--- a/mlir/python/mlir/dialects/transform/debug.py
+++ b/mlir/python/mlir/dialects/transform/debug.py
@@ -15,15 +15,6 @@ except ImportError as e:
 
 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 b35004cb4a09..00cf0840eeae 100644
--- a/mlir/python/mlir/dialects/transform/gpu.py
+++ b/mlir/python/mlir/dialects/transform/gpu.py
@@ -15,16 +15,6 @@ except ImportError as e:
 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 3311ae47b4fa..c4770b1c4067 100644
--- a/mlir/python/mlir/dialects/transform/loop.py
+++ b/mlir/python/mlir/dialects/transform/loop.py
@@ -16,15 +16,6 @@ except ImportError as e:
 
 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 f45d8054bfa5..56ea61eb817f 100644
--- a/mlir/python/mlir/dialects/transform/memref.py
+++ b/mlir/python/mlir/dialects/transform/memref.py
@@ -14,15 +14,6 @@ except ImportError as e:
 
 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 3d6427d1b1ed..74ba4c9aeb6c 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 15b0c3783790..bb5fa7ffd306 100644
--- a/mlir/python/mlir/dialects/transform/pdl.py
+++ b/mlir/python/mlir/dialects/transform/pdl.py
@@ -18,16 +18,6 @@ except ImportError as e:
 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 a4e8eb96d388..1f0b7f066118 100644
--- a/mlir/python/mlir/dialects/transform/smt.py
+++ b/mlir/python/mlir/dialects/transform/smt.py
@@ -15,16 +15,6 @@ except ImportError as e:
     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 efaf438e07b8..8b33270dc23a 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 85833d32ba55..bf40cc532065 100644
--- a/mlir/python/mlir/dialects/transform/structured.py
+++ b/mlir/python/mlir/dialects/transform/structured.py
@@ -28,15 +28,6 @@ except ImportError as e:
 
 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 904f57827146..4eb30398f087 100644
--- a/mlir/python/mlir/dialects/transform/tensor.py
+++ b/mlir/python/mlir/dialects/transform/tensor.py
@@ -14,15 +14,6 @@ except ImportError as e:
 
 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 6eadd9b4e9d5..f63f88a38242 100644
--- a/mlir/python/mlir/dialects/transform/tune.py
+++ b/mlir/python/mlir/dialects/transform/tune.py
@@ -25,15 +25,6 @@ except ImportError as e:
 
 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 1d7620188da8..af2435cb26cc 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()
-)
diff --git a/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py b/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
index 96e3f6efcc5d..cda570fceb64 100644
--- a/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
+++ b/mlir/test/python/dialects/gpu/module-to-binary-nvvm.py
@@ -2,7 +2,7 @@
 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
-import mlir.dialects.gpu as gpu
+from mlir.dialects import gpu, llvm, nvvm
 import mlir.dialects.gpu.passes
 from mlir.passmanager import *
 
diff --git a/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py b/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py
index f5ee2856f395..6e8c13b8dea3 100644
--- a/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py
+++ b/mlir/test/python/dialects/gpu/module-to-binary-rocdl.py
@@ -2,7 +2,7 @@
 # RUN: %PYTHON %s | FileCheck %s
 
 from mlir.ir import *
-import mlir.dialects.gpu as gpu
+from mlir.dialects import gpu, llvm, rocdl
 import mlir.dialects.gpu.passes
 from mlir.passmanager import *
 
diff --git a/mlir/test/python/dialects/transform.py b/mlir/test/python/dialects/transform.py
index 3477539b8d2f..6c5e4e5505b1 100644
--- a/mlir/test/python/dialects/transform.py
+++ b/mlir/test/python/dialects/transform.py
@@ -12,7 +12,6 @@ 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 39014735aa8f..03b16c324a40 100644
--- a/mlir/test/python/dialects/transform_bufferization_ext.py
+++ b/mlir/test/python/dialects/transform_bufferization_ext.py
@@ -13,7 +13,6 @@ 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 0826ed14d199..2dfdaed34386 100644
--- a/mlir/test/python/dialects/transform_debug_ext.py
+++ b/mlir/test/python/dialects/transform_debug_ext.py
@@ -19,7 +19,6 @@ 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 3aeac681f8d4..c52d5b41e11b 100644
--- a/mlir/test/python/dialects/transform_extras.py
+++ b/mlir/test/python/dialects/transform_extras.py
@@ -30,7 +30,6 @@ 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 5771ebd13ab7..17f5a6d38f8c 100644
--- a/mlir/test/python/dialects/transform_gpu_ext.py
+++ b/mlir/test/python/dialects/transform_gpu_ext.py
@@ -19,7 +19,6 @@ 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 50c93fe9f46d..430b33fba04c 100644
--- a/mlir/test/python/dialects/transform_loop_ext.py
+++ b/mlir/test/python/dialects/transform_loop_ext.py
@@ -13,7 +13,6 @@ 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 55eceb57f97c..e7d871c9eac8 100644
--- a/mlir/test/python/dialects/transform_memref_ext.py
+++ b/mlir/test/python/dialects/transform_memref_ext.py
@@ -13,7 +13,6 @@ 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 45672ff3b58e..1de8b25bab7a 100644
--- a/mlir/test/python/dialects/transform_nvgpu_ext.py
+++ b/mlir/test/python/dialects/transform_nvgpu_ext.py
@@ -12,7 +12,6 @@ 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 0593421c7713..3692fd92344a 100644
--- a/mlir/test/python/dialects/transform_smt_ext.py
+++ b/mlir/test/python/dialects/transform_smt_ext.py
@@ -19,7 +19,6 @@ 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 b22f27ab6494..e11cc6bf1e07 100644
--- a/mlir/test/python/dialects/transform_sparse_tensor_ext.py
+++ b/mlir/test/python/dialects/transform_sparse_tensor_ext.py
@@ -19,7 +19,6 @@ 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 ee15ac43e3d1..8785d6d36007 100644
--- a/mlir/test/python/dialects/transform_structured_ext.py
+++ b/mlir/test/python/dialects/transform_structured_ext.py
@@ -19,7 +19,6 @@ 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 9dde72f52a8f..a2e7aa242b9d 100644
--- a/mlir/test/python/dialects/transform_tensor_ext.py
+++ b/mlir/test/python/dialects/transform_tensor_ext.py
@@ -19,7 +19,6 @@ 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 f8db151e1070..dfb93594bca5 100644
--- a/mlir/test/python/dialects/transform_tune_ext.py
+++ b/mlir/test/python/dialects/transform_tune_ext.py
@@ -19,7 +19,6 @@ 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 0e841f2895b5..0cd9333dc121 100644
--- a/mlir/test/python/dialects/transform_vector_ext.py
+++ b/mlir/test/python/dialects/transform_vector_ext.py
@@ -21,7 +21,6 @@ def run_apply_patterns(f):
                 transform.YieldOp()
         print("\nTEST:", f.__name__)
         print(module)
-        module.operation.verify()
     return f
 
 
diff --git a/mlir/test/python/execution_engine.py b/mlir/test/python/execution_engine.py
index d569fcef32bf..f93337c5cbb6 100644
--- a/mlir/test/python/execution_engine.py
+++ b/mlir/test/python/execution_engine.py
@@ -5,6 +5,7 @@ from mlir.ir import *
 from mlir.passmanager import *
 from mlir.execution_engine import *
 from mlir.runtime import *
+from mlir.dialects import func, memref, arith, complex, cf, llvm
 
 try:
     from ml_dtypes import bfloat16, float8_e5m2
diff --git a/mlir/test/python/global_constructors.py b/mlir/test/python/global_constructors.py
index 5020c00344a3..a11045bc8526 100644
--- a/mlir/test/python/global_constructors.py
+++ b/mlir/test/python/global_constructors.py
@@ -6,6 +6,7 @@ from mlir.ir import *
 from mlir.passmanager import *
 from mlir.execution_engine import *
 from mlir.runtime import *
+from mlir.dialects import func, llvm
 
 
 # Log everything to stderr and flush so that we have a unified stream to match