[Mlir-commits] [mlir] [mlir][python] make loading dialect module (for type and value casting) best effort only (PR #72338)
Maksim Levental
llvmlistbot at llvm.org
Tue Nov 14 23:53:43 PST 2023
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/3] [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 dc6a222ff6100b47aac7d8ca55e0f0a8215f45be 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/3] [mlir][python] loading dialect module is best effort only
---
mlir/lib/Bindings/Python/IRModule.cpp | 9 ++++-----
mlir/python/CMakeLists.txt | 19 ++++++++-----------
2 files changed, 12 insertions(+), 16 deletions(-)
diff --git a/mlir/lib/Bindings/Python/IRModule.cpp b/mlir/lib/Bindings/Python/IRModule.cpp
index 5538924d2481849..6727860c094a2a1 100644
--- a/mlir/lib/Bindings/Python/IRModule.cpp
+++ b/mlir/lib/Bindings/Python/IRModule.cpp
@@ -132,10 +132,8 @@ 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;
-
+ // Try to load dialect module.
+ (void)loadDialectModule(unwrap(mlirDialectGetNamespace(dialect)));
const auto foundIt = typeCasterMap.find(mlirTypeID);
if (foundIt != typeCasterMap.end()) {
assert(foundIt->second && "type caster is defined");
@@ -146,7 +144,8 @@ std::optional<py::function> PyGlobals::lookupTypeCaster(MlirTypeID mlirTypeID,
std::optional<py::function> PyGlobals::lookupValueCaster(MlirTypeID mlirTypeID,
MlirDialect dialect) {
- loadDialectModule(unwrap(mlirDialectGetNamespace(dialect)));
+ // Try to load dialect module.
+ (void)loadDialectModule(unwrap(mlirDialectGetNamespace(dialect)));
const auto foundIt = valueCasterMap.find(mlirTypeID);
if (foundIt != valueCasterMap.end()) {
assert(foundIt->second && "value 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"
>From a03dd807de50469880f23f7cfc8988b5ff109f12 Mon Sep 17 00:00:00 2001
From: max <maksim.levental at gmail.com>
Date: Wed, 15 Nov 2023 01:53:30 -0600
Subject: [PATCH 3/3] remove dialect module loading entirely for value and type
casters
---
mlir/lib/Bindings/Python/IRModule.cpp | 4 ----
mlir/test/python/ir/builtin_types.py | 2 ++
2 files changed, 2 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Bindings/Python/IRModule.cpp b/mlir/lib/Bindings/Python/IRModule.cpp
index 6727860c094a2a1..c13a07df1169cb4 100644
--- a/mlir/lib/Bindings/Python/IRModule.cpp
+++ b/mlir/lib/Bindings/Python/IRModule.cpp
@@ -132,8 +132,6 @@ PyGlobals::lookupAttributeBuilder(const std::string &attributeKind) {
std::optional<py::function> PyGlobals::lookupTypeCaster(MlirTypeID mlirTypeID,
MlirDialect dialect) {
- // Try to load dialect module.
- (void)loadDialectModule(unwrap(mlirDialectGetNamespace(dialect)));
const auto foundIt = typeCasterMap.find(mlirTypeID);
if (foundIt != typeCasterMap.end()) {
assert(foundIt->second && "type caster is defined");
@@ -144,8 +142,6 @@ std::optional<py::function> PyGlobals::lookupTypeCaster(MlirTypeID mlirTypeID,
std::optional<py::function> PyGlobals::lookupValueCaster(MlirTypeID mlirTypeID,
MlirDialect dialect) {
- // Try to load dialect module.
- (void)loadDialectModule(unwrap(mlirDialectGetNamespace(dialect)));
const auto foundIt = valueCasterMap.find(mlirTypeID);
if (foundIt != valueCasterMap.end()) {
assert(foundIt->second && "value caster is defined");
diff --git a/mlir/test/python/ir/builtin_types.py b/mlir/test/python/ir/builtin_types.py
index d4fed86b4f135ee..738e7be4b681d02 100644
--- a/mlir/test/python/ir/builtin_types.py
+++ b/mlir/test/python/ir/builtin_types.py
@@ -766,6 +766,8 @@ def default_builder():
# I.e., we get a transform.OperationType without explicitly importing the transform dialect.
@run
def testCustomTypeTypeCaster():
+ from mlir.dialects import transform
+
with Context() as ctx, Location.unknown():
t = Type.parse('!transform.op<"foo.bar">', Context())
# CHECK: !transform.op<"foo.bar">
More information about the Mlir-commits
mailing list