[MLIR][Standalone] use narrow registration instead of RegisterEverything #160469

makslevental · 2025-09-24T08:37:55Z

This PR cleans up a long-standing TODO by avoiding MLIRPythonExtension.RegisterEverything in the Standalone example and registering the necessary dialects explicitly instead.

llvmbot · 2025-09-24T09:39:45Z

@llvm/pr-subscribers-mlir

Author: Maksim Levental (makslevental)

Changes

This PR cleans up a long-standing TODO by exposing builtin dialect registration via C thus enabling downstream bindings users to opt for narrower registration i.e., avoiding linking/building MLIRPythonExtension.RegisterEverything.

Full diff: https://github.com/llvm/llvm-project/pull/160469.diff

10 Files Affected:

(modified) mlir/examples/standalone/python/CMakeLists.txt (+8-6)
(modified) mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp (+14-5)
(modified) mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp (+13-4)
(modified) mlir/examples/standalone/test/CAPI/CMakeLists.txt (+2-3)
(modified) mlir/examples/standalone/test/CAPI/standalone-capi-test.c (+4-11)
(modified) mlir/examples/standalone/test/python/smoketest.py (+1-2)
(added) mlir/include/mlir-c/Dialect/Builtin.h (+33)
(added) mlir/lib/CAPI/Dialect/Builtin.cpp (+14)
(modified) mlir/lib/CAPI/Dialect/CMakeLists.txt (+8)
(modified) mlir/test/CMakeLists.txt (+8-1)

diff --git a/mlir/examples/standalone/python/CMakeLists.txt b/mlir/examples/standalone/python/CMakeLists.txt
index 5404128e5fa27..8ac8f501470ba 100644
--- a/mlir/examples/standalone/python/CMakeLists.txt
+++ b/mlir/examples/standalone/python/CMakeLists.txt
@@ -30,6 +30,10 @@ declare_mlir_python_extension(StandalonePythonSources.Pybind11Extension
   PRIVATE_LINK_LIBS
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
+    MLIRCAPIIR
+    MLIRCAPIArith
+    MLIRCAPIBuiltin
+    MLIRCAPITransforms
     StandaloneCAPI
   PYTHON_BINDINGS_LIBRARY pybind11
 )
@@ -42,6 +46,10 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
   PRIVATE_LINK_LIBS
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
+    MLIRCAPIIR
+    MLIRCAPIArith
+    MLIRCAPIBuiltin
+    MLIRCAPITransforms
     StandaloneCAPI
   PYTHON_BINDINGS_LIBRARY nanobind
 )
@@ -58,9 +66,6 @@ add_mlir_python_common_capi_library(StandalonePythonCAPI
   RELATIVE_INSTALL_ROOT "../../../.."
   DECLARED_SOURCES
     StandalonePythonSources
-    # TODO: Remove this in favor of showing fine grained registration once
-    # available.
-    MLIRPythonExtension.RegisterEverything
     MLIRPythonSources.Core
     MLIRPythonSources.Dialects.builtin
 )
@@ -130,9 +135,6 @@ declare_mlir_python_sources(
 )
 set(_declared_sources
   StandalonePythonSources
-  # TODO: Remove this in favor of showing fine grained registration once
-  # available.
-  MLIRPythonExtension.RegisterEverything
   MLIRPythonSources.Core
   MLIRPythonSources.Dialects.builtin
 )
diff --git a/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
index e06ec3b6472b8..d93bfcc91065d 100644
--- a/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
+++ b/mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp
@@ -10,6 +10,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "Standalone-c/Dialects.h"
+#include "mlir-c/Dialect/Arith.h"
+#include "mlir-c/Dialect/Builtin.h"
 #include "mlir/Bindings/Python/Nanobind.h"
 #include "mlir/Bindings/Python/NanobindAdaptors.h"
 
@@ -22,17 +24,24 @@ NB_MODULE(_standaloneDialectsNanobind, m) {
   auto standaloneM = m.def_submodule("standalone");
 
   standaloneM.def(
-      "register_dialect",
+      "register_dialects",
       [](MlirContext context, bool load) {
-        MlirDialectHandle handle = mlirGetDialectHandle__standalone__();
-        mlirDialectHandleRegisterDialect(handle, context);
+        MlirDialectHandle arithHandle = mlirGetDialectHandle__arith__();
+        MlirDialectHandle builtinHandle = mlirGetDialectHandle__builtin__();
+        MlirDialectHandle standaloneHandle =
+            mlirGetDialectHandle__standalone__();
+        mlirDialectHandleRegisterDialect(arithHandle, context);
+        mlirDialectHandleRegisterDialect(builtinHandle, context);
+        mlirDialectHandleRegisterDialect(standaloneHandle, context);
         if (load) {
-          mlirDialectHandleLoadDialect(handle, context);
+          mlirDialectHandleLoadDialect(arithHandle, context);
+          mlirDialectHandleRegisterDialect(builtinHandle, context);
+          mlirDialectHandleRegisterDialect(standaloneHandle, context);
         }
       },
       nb::arg("context").none() = nb::none(), nb::arg("load") = true,
       // clang-format off
-      nb::sig("def register_dialect(context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") ", load: bool = True) -> None")
+      nb::sig("def register_dialects(context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") ", load: bool = True) -> None")
       // clang-format on
   );
 }
diff --git a/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
index 397db4c20e743..530ce316049b6 100644
--- a/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
+++ b/mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp
@@ -10,6 +10,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "Standalone-c/Dialects.h"
+#include "mlir-c/Dialect/Arith.h"
+#include "mlir-c/Dialect/Builtin.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 
 using namespace mlir::python::adaptors;
@@ -21,12 +23,19 @@ PYBIND11_MODULE(_standaloneDialectsPybind11, m) {
   auto standaloneM = m.def_submodule("standalone");
 
   standaloneM.def(
-      "register_dialect",
+      "register_dialects",
       [](MlirContext context, bool load) {
-        MlirDialectHandle handle = mlirGetDialectHandle__standalone__();
-        mlirDialectHandleRegisterDialect(handle, context);
+        MlirDialectHandle arithHandle = mlirGetDialectHandle__arith__();
+        MlirDialectHandle builtinHandle = mlirGetDialectHandle__builtin__();
+        MlirDialectHandle standaloneHandle =
+            mlirGetDialectHandle__standalone__();
+        mlirDialectHandleRegisterDialect(arithHandle, context);
+        mlirDialectHandleRegisterDialect(builtinHandle, context);
+        mlirDialectHandleRegisterDialect(standaloneHandle, context);
         if (load) {
-          mlirDialectHandleLoadDialect(handle, context);
+          mlirDialectHandleLoadDialect(arithHandle, context);
+          mlirDialectHandleRegisterDialect(builtinHandle, context);
+          mlirDialectHandleRegisterDialect(standaloneHandle, context);
         }
       },
       py::arg("context") = py::none(), py::arg("load") = true);
diff --git a/mlir/examples/standalone/test/CAPI/CMakeLists.txt b/mlir/examples/standalone/test/CAPI/CMakeLists.txt
index eaa6cfc102c73..7f4b780ac6efe 100644
--- a/mlir/examples/standalone/test/CAPI/CMakeLists.txt
+++ b/mlir/examples/standalone/test/CAPI/CMakeLists.txt
@@ -6,9 +6,8 @@ add_mlir_aggregate(StandaloneCAPITestLib
   SHARED
   EMBED_LIBS
   MLIRCAPIIR
-  # TODO: Remove this in favor of showing fine grained dialect registration
-  # (once available).
-  MLIRCAPIRegisterEverything
+  MLIRCAPIArith
+  MLIRCAPIBuiltin
   StandaloneCAPI
 )
 
diff --git a/mlir/examples/standalone/test/CAPI/standalone-capi-test.c b/mlir/examples/standalone/test/CAPI/standalone-capi-test.c
index 54f3ca7f7ff14..20e86d39cdc82 100644
--- a/mlir/examples/standalone/test/CAPI/standalone-capi-test.c
+++ b/mlir/examples/standalone/test/CAPI/standalone-capi-test.c
@@ -12,21 +12,14 @@
 #include <stdio.h>
 
 #include "Standalone-c/Dialects.h"
+#include "mlir-c/Dialect/Arith.h"
+#include "mlir-c/Dialect/Builtin.h"
 #include "mlir-c/IR.h"
-#include "mlir-c/RegisterEverything.h"
-
-static void registerAllUpstreamDialects(MlirContext ctx) {
-  MlirDialectRegistry registry = mlirDialectRegistryCreate();
-  mlirRegisterAllDialects(registry);
-  mlirContextAppendDialectRegistry(ctx, registry);
-  mlirDialectRegistryDestroy(registry);
-}
 
 int main(int argc, char **argv) {
   MlirContext ctx = mlirContextCreate();
-  // TODO: Create the dialect handles for the builtin dialects and avoid this.
-  // This adds dozens of MB of binary size over just the standalone dialect.
-  registerAllUpstreamDialects(ctx);
+  mlirDialectHandleRegisterDialect(mlirGetDialectHandle__arith__(), ctx);
+  mlirDialectHandleRegisterDialect(mlirGetDialectHandle__builtin__(), ctx);
   mlirDialectHandleRegisterDialect(mlirGetDialectHandle__standalone__(), ctx);
 
   MlirModule module = mlirModuleCreateParse(
diff --git a/mlir/examples/standalone/test/python/smoketest.py b/mlir/examples/standalone/test/python/smoketest.py
index bd40c65d16164..26d84fd63e947 100644
--- a/mlir/examples/standalone/test/python/smoketest.py
+++ b/mlir/examples/standalone/test/python/smoketest.py
@@ -3,7 +3,6 @@
 
 import sys
 from mlir_standalone.ir import *
-from mlir_standalone.dialects import builtin as builtin_d
 
 if sys.argv[1] == "pybind11":
     from mlir_standalone.dialects import standalone_pybind11 as standalone_d
@@ -14,7 +13,7 @@
 
 
 with Context():
-    standalone_d.register_dialect()
+    standalone_d.register_dialects()
     module = Module.parse(
         """
     %0 = arith.constant 2 : i32
diff --git a/mlir/include/mlir-c/Dialect/Builtin.h b/mlir/include/mlir-c/Dialect/Builtin.h
new file mode 100644
index 0000000000000..c5d958249b36f
--- /dev/null
+++ b/mlir/include/mlir-c/Dialect/Builtin.h
@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/Builtin.h - C API for Builtin dialect ------*- C -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Builtin dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_BUILTIN_H
+#define MLIR_C_DIALECT_BUILTIN_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Builtin, builtin);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_BUILTIN_H
diff --git a/mlir/lib/CAPI/Dialect/Builtin.cpp b/mlir/lib/CAPI/Dialect/Builtin.cpp
new file mode 100644
index 0000000000000..dd44c8db84a14
--- /dev/null
+++ b/mlir/lib/CAPI/Dialect/Builtin.cpp
@@ -0,0 +1,14 @@
+//===- Builtin.cpp - C Interface for Builtin dialect
+//--------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir-c/Dialect/Builtin.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/IR/BuiltinDialect.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Builtin, builtin, mlir::BuiltinDialect)
diff --git a/mlir/lib/CAPI/Dialect/CMakeLists.txt b/mlir/lib/CAPI/Dialect/CMakeLists.txt
index bb1fdf8be3c8f..0b8a542b8826d 100644
--- a/mlir/lib/CAPI/Dialect/CMakeLists.txt
+++ b/mlir/lib/CAPI/Dialect/CMakeLists.txt
@@ -31,6 +31,14 @@ add_mlir_upstream_c_api_library(MLIRCAPIAsync
   MLIRPass
 )
 
+add_mlir_upstream_c_api_library(MLIRCAPIBuiltin
+  Builtin.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+)
+
 add_mlir_upstream_c_api_library(MLIRCAPIControlFlow
   ControlFlow.cpp
 
diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt
index 628adcfb6e285..2b846a9d3dc9f 100644
--- a/mlir/test/CMakeLists.txt
+++ b/mlir/test/CMakeLists.txt
@@ -123,7 +123,14 @@ set(MLIR_TEST_DEPENDS
   tblgen-to-irdl
   )
 if(NOT MLIR_STANDALONE_BUILD)
-  list(APPEND MLIR_TEST_DEPENDS FileCheck count not split-file yaml2obj)
+  list(APPEND
+    MLIR_TEST_DEPENDS
+    FileCheck count not
+    split-file yaml2obj
+    MLIRCAPIIR
+    MLIRCAPIArith
+    MLIRCAPIBuiltin
+  )
 endif()
 
 set(MLIR_TEST_DEPENDS ${MLIR_TEST_DEPENDS}

mlir/examples/standalone/test/CAPI/standalone-capi-test.c

mlir/test/CMakeLists.txt

joker-eph · 2025-09-24T19:18:24Z

mlir/test/CMakeLists.txt

 endif()
+# Examples/standalone/test.toy (vis-a-vis the standalone example) depends on these.
+if(LLVM_INCLUDE_EXAMPLES)
+  list(APPEND MLIR_TEST_DEPENDS MLIRCAPIIR MLIRCAPIArith)


Another question: these aren't already dependencies of the various mlir-capi-*-test listed above?

Yea they should be right? I can try it again but I added them of course because running LIT_FILTER=test.toy ninja check-mlir would fail with errors about missing the corresponding object files. Possibly something about CMake task graph? But I'll try again.

Ah it's just MLIRCAPIArith which is needed here (since there are no acutal Arith C API tests that depend on that lib): https://github.com/llvm/llvm-project/actions/runs/17987615036/job/51169652084#step:3:7493

joker-eph

LGTM modulo the one question I had left.

…ing (llvm#160469) This PR cleans up a long-standing TODO by avoiding `MLIRPythonExtension.RegisterEverything` in the Standalone example and registering the necessary dialects explicitly instead.

makslevental force-pushed the users/makslevental/reduce-standalone branch 5 times, most recently from eb6cf81 to cf4b971 Compare September 24, 2025 08:58

makslevental requested review from Groverkss, jpienaar, rengolin and rolfmorel September 24, 2025 09:02

makslevental force-pushed the users/makslevental/reduce-standalone branch 5 times, most recently from 859ae97 to 370b09d Compare September 24, 2025 09:28

makslevental requested a review from ftynse September 24, 2025 09:37

makslevental marked this pull request as ready for review September 24, 2025 09:39

llvmbot added the mlir label Sep 24, 2025

[MLIR][Standalone] use narrow registration instead of RegisterEverything

5187ae7

makslevental force-pushed the users/makslevental/reduce-standalone branch from 370b09d to 5187ae7 Compare September 24, 2025 09:44

joker-eph reviewed Sep 24, 2025

View reviewed changes

mlir/examples/standalone/test/CAPI/standalone-capi-test.c Outdated Show resolved Hide resolved

no need for builtin

0b3372c

makslevental requested a review from joker-eph September 24, 2025 15:07

joker-eph reviewed Sep 24, 2025

View reviewed changes

mlir/test/CMakeLists.txt Outdated Show resolved Hide resolved

gate MLIRCAPIIR MLIRCAPIArith correctly

2ca77da

makslevental requested a review from joker-eph September 24, 2025 19:17

joker-eph reviewed Sep 24, 2025

View reviewed changes

joker-eph approved these changes Sep 24, 2025

View reviewed changes

disable dep on MLIRCAPIIR MLIRCAPIArith

b2827e3

makslevental force-pushed the users/makslevental/reduce-standalone branch from 0429f84 to b2827e3 Compare September 24, 2025 20:03

makslevental merged commit f963d17 into main Sep 24, 2025
9 checks passed

makslevental deleted the users/makslevental/reduce-standalone branch September 24, 2025 20:13

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[MLIR][Standalone] use narrow registration instead of RegisterEverything #160469

[MLIR][Standalone] use narrow registration instead of RegisterEverything #160469

Uh oh!

makslevental commented Sep 24, 2025 •

edited

Loading

Uh oh!

llvmbot commented Sep 24, 2025

Uh oh!

Uh oh!

Uh oh!

joker-eph Sep 24, 2025

Uh oh!

makslevental Sep 24, 2025

Uh oh!

makslevental Sep 24, 2025

Uh oh!

joker-eph left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

[MLIR][Standalone] use narrow registration instead of RegisterEverything #160469

[MLIR][Standalone] use narrow registration instead of RegisterEverything #160469

Uh oh!

Conversation

makslevental commented Sep 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Sep 24, 2025

Uh oh!

Uh oh!

Uh oh!

joker-eph Sep 24, 2025

Choose a reason for hiding this comment

Uh oh!

makslevental Sep 24, 2025

Choose a reason for hiding this comment

Uh oh!

makslevental Sep 24, 2025

Choose a reason for hiding this comment

Uh oh!

joker-eph left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

makslevental commented Sep 24, 2025 •

edited

Loading