factor registereverything

makslevental · makslevental · commit e85d9c80c4db · 2025-09-27T03:58:37.000-04:00
diff --git a/mlir/include/mlir-c/RegisterAllExtensions.h 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/RegisterAllLLVMTranslations.h 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
@@ -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/include/mlir-c/RegisterEverything.h b/mlir/include/mlir-c/RegisterEverything.h
@@ -28,8 +28,8 @@ MLIR_CAPI_EXPORTED void mlirRegisterAllDialects(MlirDialectRegistry registry);
 /// Register all translations to LLVM IR for dialects that can support it.
 MLIR_CAPI_EXPORTED void mlirRegisterAllLLVMTranslations(MlirContext context);
 
-/// Register all compiler passes of MLIR.
-MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void);
+// /// Register all compiler passes of MLIR.
+// MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void);
 
 #ifdef __cplusplus
 }
diff --git a/mlir/lib/Bindings/Python/RegisterEverything.cpp b/mlir/lib/Bindings/Python/RegisterEverything.cpp
@@ -6,15 +6,18 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir-c/RegisterEverything.h"
+#include "mlir-c/RegisterAllExtensions.h"
+#include "mlir-c/RegisterAllLLVMTranslations.h"
+#include "mlir-c/RegisterAllPasses.h"
 #include "mlir/Bindings/Python/Nanobind.h"
 #include "mlir/Bindings/Python/NanobindAdaptors.h"
 
 NB_MODULE(_mlirRegisterEverything, m) {
-  m.doc() = "MLIR All Upstream Dialects, Translations and Passes Registration";
+  m.doc() =
+      "MLIR All Upstream Extensions, Translations and Passes Registration";
 
-  m.def("register_dialects", [](MlirDialectRegistry registry) {
-    mlirRegisterAllDialects(registry);
+  m.def("register_extensions", [](MlirDialectRegistry registry) {
+    mlirRegisterAllExtensions(registry);
   });
   m.def("register_llvm_translations",
         [](MlirContext context) { mlirRegisterAllLLVMTranslations(context); });
diff --git 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,32 @@ 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
+)
diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllExtensions.cpp
@@ -0,0 +1,16 @@
+//===- RegisterAllExtensions.cpp - Register all MLIR entities
+//-----------------===//
+//
+// 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"
+
+void mlirRegisterAllExtensions(MlirDialectRegistry registry) {
+  mlir::registerAllExtensions(*unwrap(registry));
+}
diff --git a/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp b/mlir/lib/CAPI/RegisterEverything/RegisterAllLLVMTranslations.cpp
@@ -0,0 +1,20 @@
+//===- RegisterAllPasses.cpp - Register all MLIR entities -----------------===//
+//
+// 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
@@ -0,0 +1,12 @@
+//===- RegisterAllPasses.cpp - Register all MLIR entities -----------------===//
+//
+// 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/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
@@ -670,6 +670,29 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
     MLIRCAPIFunc
 )
 
+# This extension exposes an API to register all dialects, extensions, and passes
+# packaged in upstream MLIR and it is used for the upstream "mlir" Python
+# package. Downstreams will likely want to provide their own and not depend
+# on this one, since it links in the world.
+# Note that this is not added to any top-level source target for transitive
+# inclusion: It must be included explicitly by downstreams if desired. Note that
+# this has a very large impact on what gets built/packaged.
+declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
+  MODULE_NAME _mlirRegisterEverything
+  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+  PYTHON_BINDINGS_LIBRARY nanobind
+  SOURCES
+    RegisterEverything.cpp
+  PRIVATE_LINK_LIBS
+    LLVMSupport
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIConversion
+    MLIRCAPITransforms
+    MLIRCAPIRegisterAllPasses
+    MLIRCAPIRegisterAllExtensions
+    MLIRCAPIRegisterAllLLVMTranslations
+)
+
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
   MODULE_NAME _mlirDialectsLinalg
   ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
@@ -988,6 +1011,7 @@ add_mlir_python_common_capi_library(${MLIR_PYTHON_CAPI_DYLIB_NAME}
     MLIRPythonCAPI.HeaderSources
   DECLARED_SOURCES
     MLIRPythonSources
+    MLIRPythonExtension.RegisterEverything
     ${_ADDL_TEST_SOURCES}
 )
 
@@ -1096,6 +1120,7 @@ endif()
 
 set(_declared_sources
   MLIRPythonSources
+  MLIRPythonExtension.RegisterEverything
   MLIRPythonCAPICTypesBinding)
 if(NOT CMAKE_CROSSCOMPILING)
   list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen)
diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py
@@ -123,6 +123,9 @@ def process_initializer_module(module_name):
         if hasattr(m, "register_dialects"):
             logger.debug("Registering dialects from initializer %r", m)
             m.register_dialects(get_dialect_registry())
+        if hasattr(m, "register_extensions"):
+            logger.debug("Registering extensions from initializer %r", m)
+            m.register_extensions(get_dialect_registry())
         if hasattr(m, "context_init_hook"):
             logger.debug("Adding context init hook from %r", m)
             post_init_hooks.append(m.context_init_hook)

Original file line number	Diff line number	Diff line change
`@@ -28,8 +28,8 @@ MLIR_CAPI_EXPORTED void mlirRegisterAllDialects(MlirDialectRegistry registry);`
`28`	`28`	`/// Register all translations to LLVM IR for dialects that can support it.`
`29`	`29`	`MLIR_CAPI_EXPORTED void mlirRegisterAllLLVMTranslations(MlirContext context);`
`30`	`30`
`31`		`-/// Register all compiler passes of MLIR.`
`32`		`-MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void);`
	`31`	`+// /// Register all compiler passes of MLIR.`
	`32`	`+// MLIR_CAPI_EXPORTED void mlirRegisterAllPasses(void);`
`33`	`33`
`34`	`34`	`#ifdef __cplusplus`
`35`	`35`	`}`