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

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

Author: makslevental

mlir/examples/standalone/python/CMakeLists.txt

@@ -30,6 +30,9 @@ declare_mlir_python_extension(StandalonePythonSources.Pybind11Extension
   PRIVATE_LINK_LIBS
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
+    MLIRCAPIIR
+    MLIRCAPIArith
+    MLIRCAPITransforms
     StandaloneCAPI
   PYTHON_BINDINGS_LIBRARY pybind11
 )
@@ -42,6 +45,9 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
   PRIVATE_LINK_LIBS
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
+    MLIRCAPIIR
+    MLIRCAPIArith
+    MLIRCAPITransforms
     StandaloneCAPI
   PYTHON_BINDINGS_LIBRARY nanobind
 )
@@ -58,9 +64,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 +133,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
 )

mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp

@@ -10,6 +10,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "Standalone-c/Dialects.h"
+#include "mlir-c/Dialect/Arith.h"
 #include "mlir/Bindings/Python/Nanobind.h"
 #include "mlir/Bindings/Python/NanobindAdaptors.h"
 
@@ -22,17 +23,21 @@ 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 standaloneHandle =
+            mlirGetDialectHandle__standalone__();
+        mlirDialectHandleRegisterDialect(arithHandle, context);
+        mlirDialectHandleRegisterDialect(standaloneHandle, context);
         if (load) {
-          mlirDialectHandleLoadDialect(handle, context);
+          mlirDialectHandleLoadDialect(arithHandle, 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
   );
 }

mlir/examples/standalone/python/StandaloneExtensionPybind11.cpp

@@ -10,6 +10,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "Standalone-c/Dialects.h"
+#include "mlir-c/Dialect/Arith.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 
 using namespace mlir::python::adaptors;
@@ -21,12 +22,16 @@ 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 standaloneHandle =
+            mlirGetDialectHandle__standalone__();
+        mlirDialectHandleRegisterDialect(arithHandle, context);
+        mlirDialectHandleRegisterDialect(standaloneHandle, context);
         if (load) {
-          mlirDialectHandleLoadDialect(handle, context);
+          mlirDialectHandleLoadDialect(arithHandle, context);
+          mlirDialectHandleRegisterDialect(standaloneHandle, context);
         }
       },
       py::arg("context") = py::none(), py::arg("load") = true);

mlir/examples/standalone/test/CAPI/CMakeLists.txt

@@ -6,9 +6,7 @@ add_mlir_aggregate(StandaloneCAPITestLib
   SHARED
   EMBED_LIBS
   MLIRCAPIIR
-  # TODO: Remove this in favor of showing fine grained dialect registration
-  # (once available).
-  MLIRCAPIRegisterEverything
+  MLIRCAPIArith
   StandaloneCAPI
 )
 

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

@@ -12,21 +12,12 @@
 #include <stdio.h>
 
 #include "Standalone-c/Dialects.h"
+#include "mlir-c/Dialect/Arith.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__standalone__(), ctx);
 
   MlirModule module = mlirModuleCreateParse(

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

mlir/test/CMakeLists.txt

@@ -126,6 +126,10 @@ set(MLIR_TEST_DEPENDS
 if(NOT MLIR_STANDALONE_BUILD)
   list(APPEND MLIR_TEST_DEPENDS FileCheck count not split-file yaml2obj)
 endif()
+# Examples/standalone/test.toy (vis-a-vis the standalone example) depends on these.
+if(LLVM_INCLUDE_EXAMPLES)
+  list(APPEND MLIR_TEST_DEPENDS MLIRCAPIArith)
+endif()
 
 set(MLIR_TEST_DEPENDS ${MLIR_TEST_DEPENDS}
   mlir-capi-pdl-test