Skip to content

Conversation

makslevental
Copy link
Contributor

@makslevental makslevental commented Sep 29, 2025

TL;DR: This PR enables registering dialects piecemeal (in the Python bindings) without using _site_initialize_0 or anything like that.

The PR is large but it's mostly boilerplate. The way the new registration mechanism works is by exposing a ctypes-based CAPI for mlirGetDialectHandle__DIALECT__ and adding DialectRegistry.insert_dialect. The three changes to focus on are

  1. python/mlir/_mlir_libs/_capi.py
  2. test/python/ir/capi.py
  3. examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py

The core of it looks like this:

from .._mlir_libs import get_dialect_registry
from .._mlir_libs._capi import register_dialect

if "quant" not in dialect_registry.dialect_names:
    register_dialect("quant", get_dialect_registry())

I.e., if someone chooses to forgo using MLIRPythonExtension.RegisterEverything, they can drop these registration calls somewhere (see the changes to the standalone example which now registers the quant dialect using this path) and everything will work.

The rest of the PR simply adds the boilerplate necessary for users to actually make these calls for all the dialects.

Some notes:

  1. If you do not link MLIRPythonExtension.RegisterEverything you will not transitively link your chosen dialects. Thus you will need to add MLIRCAPI<DIALECT> to either your declare_mlir_python_extension invocation or your add_mlir_python_common_capi_library invocation (see the changes to examples/standalone/python/CMakeLists.txt);
  2. The core bindings are not actually migrated to use this new mechanism because that would be a breaking change. But the change was implemented and tested for core as well (see this other version of the PR [MLIR][Python] enable precise registration #160742);
  3. This does not automagically clean up the build graph associated with MLIRPythonExtension.Core (i.e., MLIRPythonExtension.Core still depends on buildings and linking libs for all the dialects). This is because in addition to registering dialects, MLIRPythonExtension.RegisterEverything registers all passes, all dialect extensions, all interfaces, and all LLVMIR translations. This is tedious but doable (cf the other PR);
  4. This PR also demonstrates how to interop with our nanobind types using bare ctypes - the added test shows how to create PyCapsules from ctypes.c_void_p.

@makslevental makslevental force-pushed the users/makslevental/specific-reg-v3 branch 5 times, most recently from e19a7b4 to 16aca5e Compare September 29, 2025 04:48
@makslevental makslevental changed the title [MLIR][Python] enable precise registration [MLIR][Python] enable narrow registration (avoid registering all dialects) Sep 29, 2025
@makslevental makslevental force-pushed the users/makslevental/specific-reg-v3 branch 3 times, most recently from 99683dd to a3272cb Compare September 29, 2025 05:23
@makslevental makslevental changed the title [MLIR][Python] enable narrow registration (avoid registering all dialects) [MLIR][Python] enable narrow registration directly from Python (avoid registering all dialects) Sep 29, 2025
@makslevental makslevental force-pushed the users/makslevental/specific-reg-v3 branch 2 times, most recently from 4cb401a to 72f15f3 Compare September 29, 2025 05:37
@makslevental
Copy link
Contributor Author

cc @nsmithtt @Ahajha

@makslevental makslevental marked this pull request as ready for review September 29, 2025 05:42
@llvmbot llvmbot added mlir:python MLIR Python bindings mlir labels Sep 29, 2025
@llvmbot
Copy link
Member

llvmbot commented Sep 29, 2025

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

TL;DR: This PR enables registering dialects piecemeal (in the Python bindings) without using _site_initialize_0 or anything like that.

The PR is large but it's mostly boilerplate. The way the new registration mechanism works is by exposing a ctypes-based CAPI for mlirGetDialectHandle__DIALECT__ and adding DialectRegistry.insert_dialect. The three changes to focus on are

  1. python/mlir/_mlir_libs/_capi.py
  2. test/python/ir/capi.py
  3. examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py

The core of it looks like this:

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)

I.e., if someone chooses to forgo using MLIRPythonExtension.RegisterEverything, they can drop these registration calls somewhere (see the changes to the standalone example which now registers the quant dialect using this path) and everything will work.

The rest of the PR simply adds the boilerplate necessary for users to actually make these calls for all the dialects.

Some notes:

  1. If you do not link MLIRPythonExtension.RegisterEverything you will not transitively link your chosen dialects. Thus you will need to add MLIRCAPI&lt;DIALECT&gt; to either your declare_mlir_python_extension invocation or your add_mlir_python_common_capi_library invocation (see the changes to examples/standalone/python/CMakeLists.txt);
  2. The core bindings are not actually migrated to use this new mechanism because that would be a breaking change. But the change was implemented and tested for core as well (see this other version of the PR [MLIR][Python] enable precise registration #160742);
  3. This does not automagically clean up the build graph associated with MLIRPythonExtension.Core (i.e., that MLIRPythonExtension.Core still depends on buildings and linking libs for all the dialects). This is because in addition to registering dialects, MLIRPythonExtension.RegisterEverything registers all passes, all dialect extensions, all interfaces, and all LLVMIR translations. This is tedious but doable (cf the other PR);
  4. This PR also demonstrates how to interop with our nanobind types using bare ctypes - the added test shows how to create PyCapsules from ctypes.c_void_p.

Patch is 37.37 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/161118.diff

38 Files Affected:

  • (modified) mlir/examples/standalone/pyproject.toml (+5)
  • (modified) mlir/examples/standalone/python/CMakeLists.txt (+4-1)
  • (modified) mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp (+1-1)
  • (modified) mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp (+1-1)
  • (modified) mlir/examples/standalone/python/mlir_standalone/dialects/standalone_nanobind.py (+7)
  • (modified) mlir/examples/standalone/python/mlir_standalone/dialects/standalone_pybind11.py (+7)
  • (modified) mlir/examples/standalone/test/python/smoketest.py (+19-8)
  • (modified) mlir/include/mlir-c/Bindings/Python/Interop.h (+9)
  • (added) mlir/include/mlir-c/Dialect/Affine.h (+36)
  • (added) mlir/include/mlir-c/Dialect/Bufferization.h (+36)
  • (added) mlir/include/mlir-c/Dialect/Complex.h (+33)
  • (modified) mlir/include/mlir-c/Dialect/GPU.h (+3)
  • (modified) mlir/include/mlir-c/Dialect/Linalg.h (+3)
  • (modified) mlir/include/mlir-c/Dialect/MemRef.h (+3)
  • (modified) mlir/include/mlir-c/Dialect/NVGPU.h (+3)
  • (modified) mlir/include/mlir-c/Dialect/PDL.h (+3)
  • (modified) mlir/include/mlir-c/Dialect/SMT.h (+3)
  • (modified) mlir/include/mlir-c/Dialect/SparseTensor.h (+3)
  • (modified) mlir/include/mlir-c/Dialect/Tensor.h (+3)
  • (added) mlir/include/mlir-c/Dialect/Tosa.h (+33)
  • (added) mlir/include/mlir-c/Dialect/UB.h (+33)
  • (modified) mlir/include/mlir-c/Dialect/Vector.h (+3)
  • (modified) mlir/include/mlir-c/IR.h (+13-5)
  • (modified) mlir/lib/Bindings/Python/IRCore.cpp (+41-1)
  • (added) mlir/lib/CAPI/Dialect/Affine.cpp (+14)
  • (added) mlir/lib/CAPI/Dialect/Bufferization.cpp (+14)
  • (modified) mlir/lib/CAPI/Dialect/CMakeLists.txt (+45)
  • (added) mlir/lib/CAPI/Dialect/Complex.cpp (+14)
  • (added) mlir/lib/CAPI/Dialect/Tosa.cpp (+13)
  • (modified) mlir/lib/CAPI/Dialect/TransformInterpreter.cpp (-3)
  • (added) mlir/lib/CAPI/Dialect/UB.cpp (+13)
  • (modified) mlir/lib/CAPI/IR/IR.cpp (+12)
  • (modified) mlir/lib/Dialect/Transform/IR/TransformOps.cpp (+2-1)
  • (modified) mlir/python/CMakeLists.txt (+8-3)
  • (added) mlir/python/mlir/_mlir_libs/_capi.py (+50)
  • (added) mlir/python/mlir/dialects/QuantOps.td (+14)
  • (modified) mlir/python/mlir/dialects/quant.py (+1)
  • (added) mlir/test/python/ir/capi.py (+41)
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..6d86d0e4ce3a8 100644
--- a/mlir/examples/standalone/python/CMakeLists.txt
+++ b/mlir/examples/standalone/python/CMakeLists.txt
@@ -65,7 +65,8 @@ add_mlir_python_common_capi_library(StandalonePythonCAPI
   DECLARED_SOURCES
     StandalonePythonSources
     MLIRPythonSources.Core
-    MLIRPythonSources.Dialects.builtin
+  EMBED_LIBS
+    MLIRCAPIQuant
 )
 
 ################################################################################
@@ -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/StandaloneExtensionNanobind.cpp b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
index 0ec6cdfa7994b..e28744b32ccfe 100644
--- a/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
+++ b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
@@ -32,7 +32,7 @@ NB_MODULE(_standaloneDialectsNanobind, m) {
         mlirDialectHandleRegisterDialect(standaloneHandle, context);
         if (load) {
           mlirDialectHandleLoadDialect(arithHandle, context);
-          mlirDialectHandleRegisterDialect(standaloneHandle, context);
+          mlirDialectHandleLoadDialect(standaloneHandle, context);
         }
       },
       nb::arg("context").none() = nb::none(), nb::arg("load") = true,
diff --git a/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
index da8c2167dc36b..db7645bcf5331 100644
--- a/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
+++ b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
@@ -31,7 +31,7 @@ PYBIND11_MODULE(_standaloneDialectsPybind11, m) {
         mlirDialectHandleRegisterDialect(standaloneHandle, context);
         if (load) {
           mlirDialectHandleLoadDialect(arithHandle, context);
-          mlirDialectHandleRegisterDialect(standaloneHandle, context);
+          mlirDialectHandleLoadDialect(standaloneHandle, context);
         }
       },
       py::arg("context") = py::none(), py::arg("load") = true);
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..3628522360c74 100644
--- a/mlir/examples/standalone/test/python/smoketest.py
+++ b/mlir/examples/standalone/test/python/smoketest.py
@@ -11,15 +11,26 @@
 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)
+    i32 = IntegerType.get_signless(32)
+    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):
+        two_i32 = arith.constant(i32, 2)
+        standalone_d.foo(two_i32)
+        two_f32 = arith.constant(f32, 2.0)
+        quant.qcast(uniform, two_f32)
+    # CHECK: %[[TWOI32:.*]] = arith.constant 2 : i32
+    # CHECK: standalone.foo %[[TWOI32]] : i32
+    # CHECK: %[[TWOF32:.*]] = arith.constant 2.000000e+00 : f32
+    # CHECK: quant.qcast %[[TWOF32]] : 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 dialec...
[truncated]

@makslevental
Copy link
Contributor Author

For context: here is a previous PR that attempted to do the same thing #74252.

@makslevental makslevental force-pushed the users/makslevental/specific-reg-v3 branch from 72f15f3 to 2dd5d63 Compare September 29, 2025 05:44
@makslevental makslevental changed the title [MLIR][Python] enable narrow registration directly from Python (avoid registering all dialects) [MLIR][Python] enable piecemeal dialect registration directly from Python (avoid registering all dialects) Sep 29, 2025
@makslevental makslevental changed the title [MLIR][Python] enable piecemeal dialect registration directly from Python (avoid registering all dialects) [MLIR][Python] enable piecemeal dialect registration directly from Python Sep 29, 2025
@makslevental makslevental force-pushed the users/makslevental/specific-reg-v3 branch 3 times, most recently from 969af55 to 4a54377 Compare September 29, 2025 05:51
@makslevental makslevental force-pushed the users/makslevental/specific-reg-v3 branch 3 times, most recently from 54f76f4 to b703e53 Compare September 29, 2025 06:28
@makslevental makslevental force-pushed the users/makslevental/specific-reg-v3 branch from b703e53 to f8e5414 Compare September 29, 2025 06:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
mlir:python MLIR Python bindings mlir
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants