[Mlir-commits] [mlir] [mlir][python] automatically bundle builtin dialect with core sources (PR #72338)

[Maksim Levental]

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

>From 8ae7d2d868c8dbc827f1d996722c564b12b4eaf0 Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Tue, 14 Nov 2023 21:33:39 -0600
Subject: [PATCH 1/2] [mlir][python] automatically bundle builtin dialect with
 core

---
 mlir/python/CMakeLists.txt | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 971ad2dd214a15f..6be84ec311e23ca 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -44,6 +44,17 @@ declare_mlir_python_sources(MLIRPythonCAPI.HeaderSources
   SOURCES_GLOB "mlir-c/*.h"
 )
 
+# The builtin dialect is special (e.g., type casters and such expect it to be loaded
+# in order to work) so we add it to Core instead of Dialects so that anyone depending on
+# Core will get it automatically.
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Core
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/BuiltinOps.td
+  SOURCES
+    dialects/builtin.py
+  DIALECT_NAME builtin)
+
 ################################################################################
 # Dialect bindings
 ################################################################################
@@ -84,14 +95,6 @@ declare_mlir_dialect_python_bindings(
     "../../include/mlir/Dialect/Bufferization/IR/BufferizationEnums.td"
 )
 
-declare_mlir_dialect_python_bindings(
-  ADD_TO_PARENT MLIRPythonSources.Dialects
-  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
-  TD_FILE dialects/BuiltinOps.td
-  SOURCES
-    dialects/builtin.py
-  DIALECT_NAME builtin)
-
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"

>From 140223dd50cd938c71b5c6bb74975aa896ca8d0d Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Tue, 14 Nov 2023 22:53:12 -0600
Subject: [PATCH 2/2] [mlir][python] loading dialect module is best effort only

---
 mlir/lib/Bindings/Python/IRModule.cpp |  4 +---
 mlir/python/CMakeLists.txt            | 19 ++++++++-----------
 2 files changed, 9 insertions(+), 14 deletions(-)

diff --git a/mlir/lib/Bindings/Python/IRModule.cpp b/mlir/lib/Bindings/Python/IRModule.cpp
index 5538924d2481849..328bdc46ca1182d 100644
--- a/mlir/lib/Bindings/Python/IRModule.cpp
+++ b/mlir/lib/Bindings/Python/IRModule.cpp
@@ -133,9 +133,7 @@ PyGlobals::lookupAttributeBuilder(const std::string &attributeKind) {
 std::optional<py::function> PyGlobals::lookupTypeCaster(MlirTypeID mlirTypeID,
                                                         MlirDialect dialect) {
   // Make sure dialect module is loaded.
-  if (!loadDialectModule(unwrap(mlirDialectGetNamespace(dialect))))
-    return std::nullopt;
-
+  loadDialectModule(unwrap(mlirDialectGetNamespace(dialect)));
   const auto foundIt = typeCasterMap.find(mlirTypeID);
   if (foundIt != typeCasterMap.end()) {
     assert(foundIt->second && "type caster is defined");
diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt
index 6be84ec311e23ca..971ad2dd214a15f 100644
--- a/mlir/python/CMakeLists.txt
+++ b/mlir/python/CMakeLists.txt
@@ -44,17 +44,6 @@ declare_mlir_python_sources(MLIRPythonCAPI.HeaderSources
   SOURCES_GLOB "mlir-c/*.h"
 )
 
-# The builtin dialect is special (e.g., type casters and such expect it to be loaded
-# in order to work) so we add it to Core instead of Dialects so that anyone depending on
-# Core will get it automatically.
-declare_mlir_dialect_python_bindings(
-  ADD_TO_PARENT MLIRPythonSources.Core
-  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
-  TD_FILE dialects/BuiltinOps.td
-  SOURCES
-    dialects/builtin.py
-  DIALECT_NAME builtin)
-
 ################################################################################
 # Dialect bindings
 ################################################################################
@@ -95,6 +84,14 @@ declare_mlir_dialect_python_bindings(
     "../../include/mlir/Dialect/Bufferization/IR/BufferizationEnums.td"
 )
 
+declare_mlir_dialect_python_bindings(
+  ADD_TO_PARENT MLIRPythonSources.Dialects
+  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
+  TD_FILE dialects/BuiltinOps.td
+  SOURCES
+    dialects/builtin.py
+  DIALECT_NAME builtin)
+
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"