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

Maksim Levental llvmlistbot at llvm.org
Sun Sep 28 21:40:11 PDT 2025


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

>From e19a7b4b04d7a0e7d9377820113848675d8c7be5 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] [MLIR][Python] enable precise registration

---
 mlir/examples/standalone/pyproject.toml       |   5 +
 .../examples/standalone/python/CMakeLists.txt |   5 +-
 .../dialects/standalone_nanobind.py           |   7 ++
 .../dialects/standalone_pybind11.py           |   7 ++
 .../standalone/test/python/smoketest.py       |  24 +++--
 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/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/UB.h              |  33 ++++++
 mlir/include/mlir-c/Dialect/Vector.h          |   3 +
 mlir/include/mlir-c/IR.h                      |  18 +++-
 mlir/include/mlir-c/RegisterAllExtensions.h   |  26 +++++
 .../mlir-c/RegisterAllExternalModels.h        |  27 +++++
 .../mlir-c/RegisterAllLLVMTranslations.h      |  26 +++++
 mlir/include/mlir-c/RegisterAllPasses.h       |  26 +++++
 mlir/lib/Bindings/Python/IRCore.cpp           |  42 +++++++-
 mlir/lib/CAPI/Dialect/Affine.cpp              |  14 +++
 mlir/lib/CAPI/Dialect/Bufferization.cpp       |  14 +++
 mlir/lib/CAPI/Dialect/CMakeLists.txt          |  45 ++++++++
 mlir/lib/CAPI/Dialect/Complex.cpp             |  14 +++
 mlir/lib/CAPI/Dialect/Tosa.cpp                |  13 +++
 .../lib/CAPI/Dialect/TransformInterpreter.cpp |   3 -
 mlir/lib/CAPI/Dialect/UB.cpp                  |  13 +++
 mlir/lib/CAPI/IR/IR.cpp                       |  12 +++
 .../CAPI/RegisterEverything/CMakeLists.txt    |  39 +++++++
 .../RegisterAllExtensions.cpp                 |  18 ++++
 .../RegisterAllExternalModels.cpp             | 101 ++++++++++++++++++
 .../RegisterAllLLVMTranslations.cpp           |  20 ++++
 .../RegisterEverything/RegisterAllPasses.cpp  |  12 +++
 .../lib/Dialect/Transform/IR/TransformOps.cpp |   3 +-
 mlir/python/CMakeLists.txt                    |  11 +-
 mlir/python/mlir/_mlir_libs/_capi.py          |  50 +++++++++
 mlir/python/mlir/dialects/QuantOps.td         |  14 +++
 mlir/python/mlir/dialects/quant.py            |   1 +
 mlir/test/python/ir/capi.py                   |  41 +++++++
 45 files changed, 836 insertions(+), 22 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/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/include/mlir-c/RegisterAllExtensions.h
 create mode 100644 mlir/include/mlir-c/RegisterAllExternalModels.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/Dialect/Affine.cpp
 create mode 100644 mlir/lib/CAPI/Dialect/Bufferization.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/UB.cpp
 create mode 100644 mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp
 create mode 100644 mlir/lib/CAPI/RegisterEverything/RegisterAllExternalModels.cpp
 create mode 100644 mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp
 create mode 100644 mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp
 create mode 100644 mlir/python/mlir/_mlir_libs/_capi.py
 create mode 100644 mlir/python/mlir/dialects/QuantOps.td
 create mode 100644 mlir/test/python/ir/capi.py

diff --git a/mlir/examples/standalone/pyproject.toml b/mlir/examples/standalone/pyproject.toml
index 5a1e6e86513c3..90535030df817 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 905c944939756..58ac00339a96b 100644
--- a/mlir/examples/standalone/python/CMakeLists.txt
+++ b/mlir/examples/standalone/python/CMakeLists.txt
@@ -32,6 +32,7 @@ declare_mlir_python_extension(StandalonePythonSources.Pybind11Extension
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPIArith
+    MLIRCAPIQuant
     MLIRCAPITransforms
     StandaloneCAPI
   PYTHON_BINDINGS_LIBRARY pybind11
@@ -47,6 +48,7 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPIArith
+    MLIRCAPIQuant
     MLIRCAPITransforms
     StandaloneCAPI
   PYTHON_BINDINGS_LIBRARY nanobind
@@ -65,7 +67,6 @@ add_mlir_python_common_capi_library(StandalonePythonCAPI
   DECLARED_SOURCES
     StandalonePythonSources
     MLIRPythonSources.Core
-    MLIRPythonSources.Dialects.builtin
 )
 
 ################################################################################
@@ -138,6 +139,8 @@ set(_declared_sources
   StandalonePythonSources
   MLIRPythonSources.Core
   MLIRPythonSources.Dialects.builtin
+  MLIRPythonSources.Dialects.arith
+  MLIRPythonSources.Dialects.quant
 )
 # 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/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py b/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py
index 6218720951c82..3eb011ff4e87e 100644
--- a/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py
+++ b/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py
@@ -4,3 +4,10 @@
 
 from ._standalone_ops_gen import *
 from .._mlir_libs._standaloneDialectsNanobind.standalone import *
+
+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 "mlirGetDialectHandle__quant__" not in _dialect_registry.dialect_names:
+    _register_dialect("mlirGetDialectHandle__quant__", _dialect_registry)
diff --git a/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py b/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py
index bfb98e404e13f..2a259b71b33be 100644
--- a/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py
+++ b/mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py
@@ -4,3 +4,10 @@
 
 from ._standalone_ops_gen import *
 from .._mlir_libs._standaloneDialectsPybind11.standalone import *
+
+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 "mlirGetDialectHandle__quant__" not in _dialect_registry.dialect_names:
+    _register_dialect("mlirGetDialectHandle__quant__", _dialect_registry)
diff --git a/mlir/examples/standalone/test/python/smoketest.py b/mlir/examples/standalone/test/python/smoketest.py
index 26d84fd63e947..149bcc9d57d46 100644
--- a/mlir/examples/standalone/test/python/smoketest.py
+++ b/mlir/examples/standalone/test/python/smoketest.py
@@ -11,15 +11,23 @@
 else:
     raise ValueError("Expected either pybind11 or nanobind as arguments")
 
+from mlir_standalone.dialects import arith, quant
 
-with Context():
+
+with Context(), Location.unknown():
     standalone_d.register_dialects()
-    module = Module.parse(
-        """
-    %0 = arith.constant 2 : i32
-    %1 = standalone.foo %0 : i32
-    """
+    f32 = F32Type.get()
+    i8 = IntegerType.get_signless(8)
+    f32 = F32Type.get()
+    uniform = quant.UniformQuantizedType.get(
+        quant.UniformQuantizedType.FLAG_SIGNED, i8, f32, 0.99872, 127, -8, 7
     )
-    # CHECK: %[[C:.*]] = arith.constant 2 : i32
-    # CHECK: standalone.foo %[[C]] : i32
+    module = Module.create()
+    with InsertionPoint(module.body):
+        v0 = arith.constant(f32, 2.0)
+        v1 = standalone_d.foo(v0)
+        v2 = quant.qcast(uniform, v0)
+    # CHECK: %[[C:.*]] = arith.constant 2.0 : f32
+    # CHECK: standalone.foo %[[C]] : f32
+    # CHECK: quant.qcast %[[C]] : f32 to !quant.uniform<i8<-8:7>:f32, 9.987200e-01:127>
     print(str(module))
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/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/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/include/mlir-c/RegisterAllExtensions.h b/mlir/include/mlir-c/RegisterAllExtensions.h
new file mode 100644
index 0000000000000..dc2dcfbb28098
--- /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/RegisterAllExternalModels.h b/mlir/include/mlir-c/RegisterAllExternalModels.h
new file mode 100644
index 0000000000000..baa8a742cd019
--- /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/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/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 83a8757bb72c7..21e60d5550a51 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,14 +3134,46 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           },
           nb::sig("def __repr__(self) -> str"));
 
+  //----------------------------------------------------------------------------
+  // Mapping of MlirDialectHandle
+  //----------------------------------------------------------------------------
+
+  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);
+
   //----------------------------------------------------------------------------
   // Mapping of PyDialectRegistry
   //----------------------------------------------------------------------------
   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<>());
+      .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/CMakeLists.txt b/mlir/lib/CAPI/Dialect/CMakeLists.txt
index bb1fdf8be3c8f..c667ca354b266 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,15 @@ 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(MLIRCAPIAsync
   Async.cpp
   AsyncPasses.cpp
@@ -31,6 +49,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
 
@@ -278,3 +305,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/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/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/CAPI/RegisterEverything/CMakeLists.txt b/mlir/lib/CAPI/RegisterEverything/CMakeLists.txt
index ccda668ece3d4..473c3844538c5 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,41 @@ 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
+)
+
+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
new file mode 100644
index 0000000000000..9be46bfff012f
--- /dev/null
+++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp
@@ -0,0 +1,18 @@
+//===- 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/RegisterAllExtensions.h"
+#include "mlir/CAPI/IR.h"
+#include "mlir/InitAllExtensions.h"
+
+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 0000000000000..cee3d717ccb33
--- /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
new file mode 100644
index 0000000000000..0eeeaa279aef7
--- /dev/null
+++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp
@@ -0,0 +1,20 @@
+//===- 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.
+// 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..88a189cf73605
--- /dev/null
+++ b/mlir/lib/CAPI/RegisterEverything/RegisterAllPasses.cpp
@@ -0,0 +1,12 @@
+//===- 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.
+// 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/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 9f5246de6bda0..7969e9ca4750b 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
@@ -370,14 +371,16 @@ declare_mlir_dialect_python_bindings(
     dialects/rocdl.py
   DIALECT_NAME rocdl)
 
-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
+)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -528,6 +531,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
 
     # Dialects
     MLIRCAPIFunc
+    MLIRCAPIArith
+    MLIRCAPIQuant
 )
 
 # This extension exposes an API to register all dialects, extensions, and passes
diff --git a/mlir/python/mlir/_mlir_libs/_capi.py b/mlir/python/mlir/_mlir_libs/_capi.py
new file mode 100644
index 0000000000000..3d7836daa1c8b
--- /dev/null
+++ b/mlir/python/mlir/_mlir_libs/_capi.py
@@ -0,0 +1,50 @@
+#  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
+import platform
+from pathlib import Path
+
+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
+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
+
+
+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)
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/test/python/ir/capi.py b/mlir/test/python/ir/capi.py
new file mode 100644
index 0000000000000..d3d4556dff954
--- /dev/null
+++ b/mlir/test/python/ir/capi.py
@@ -0,0 +1,41 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+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, DialectRegistry
+
+print("success")
+# 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 = DialectRegistry()
+# 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)



More information about the Mlir-commits mailing list