[Mlir-commits] [mlir] d2b9664 - [mlir][py] Fix negative cached value in attribute builder

[Jacques Pienaar]

Author: Jacques Pienaar
Date: 2022-12-27T21:56:58-08:00
New Revision: d2b96647595da2c1f0fea5e0555466d3db93815c

URL: https://github.com/llvm/llvm-project/commit/d2b96647595da2c1f0fea5e0555466d3db93815c
DIFF: https://github.com/llvm/llvm-project/commit/d2b96647595da2c1f0fea5e0555466d3db93815c.diff

LOG: [mlir][py] Fix negative cached value in attribute builder

Previously this was incorrectly assigning py::none to where function was
expected which resulted in failure if one used a non-attribute for
attribute without registered builder.

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/Globals.h
    mlir/lib/Bindings/Python/IRModule.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/Globals.h b/mlir/lib/Bindings/Python/Globals.h
index ba6cfb545b71b..f0bf3f556357f 100644
--- a/mlir/lib/Bindings/Python/Globals.h
+++ b/mlir/lib/Bindings/Python/Globals.h
@@ -103,7 +103,7 @@ class PyGlobals {
   /// the OpView base class (bypassing the user class constructor).
   llvm::StringMap<pybind11::object> rawOpViewClassMap;
   /// Map of attribute ODS name to custom builder.
-  llvm::StringMap<pybind11::function> attributeBuilderMap;
+  llvm::StringMap<pybind11::object> attributeBuilderMap;
 
   /// Set of dialect namespaces that we have attempted to import implementation
   /// modules for.

diff  --git a/mlir/lib/Bindings/Python/IRModule.cpp b/mlir/lib/Bindings/Python/IRModule.cpp
index be6de5fd2f1ba..1cdd7e4412668 100644
--- a/mlir/lib/Bindings/Python/IRModule.cpp
+++ b/mlir/lib/Bindings/Python/IRModule.cpp
@@ -62,7 +62,7 @@ void PyGlobals::loadDialectModule(llvm::StringRef dialectNamespace) {
 
 void PyGlobals::registerAttributeBuilder(const std::string &attributeKind,
                                          py::function pyFunc) {
-  py::function &found = attributeBuilderMap[attributeKind];
+  py::object &found = attributeBuilderMap[attributeKind];
   if (found) {
     throw std::runtime_error((llvm::Twine("Attribute builder for '") +
                               attributeKind + "' is already registered")