[Mlir-commits] [mlir] d747a17 - [MLIR] [Python] Fix `Value.owner` to handle BlockArgs

[John Demme]

Author: John Demme
Date: 2022-08-09T19:37:04-07:00
New Revision: d747a170a47dce64aefb906211e8aaed0c6bd6f6

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

LOG: [MLIR] [Python] Fix `Value.owner` to handle BlockArgs

Previously, calling `Value.owner()` would C++ assert in debug builds if
`Value` was a block argument. Additionally, the behavior was just wrong
in release builds. This patch adds support for BlockArg Values.

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/IRCore.cpp
    mlir/test/python/ir/value.py

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index beb0c6cb9f72b..db199b38b9158 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -3118,11 +3118,22 @@ void mlir::python::populateIRCore(py::module &m) {
       .def_property_readonly(
           "owner",
           [](PyValue &self) {
-            assert(mlirOperationEqual(self.getParentOperation()->get(),
-                                      mlirOpResultGetOwner(self.get())) &&
-                   "expected the owner of the value in Python to match that in "
-                   "the IR");
-            return self.getParentOperation().getObject();
+            MlirValue v = self.get();
+            if (mlirValueIsAOpResult(v)) {
+              assert(
+                  mlirOperationEqual(self.getParentOperation()->get(),
+                                     mlirOpResultGetOwner(self.get())) &&
+                  "expected the owner of the value in Python to match that in "
+                  "the IR");
+              return self.getParentOperation().getObject();
+            }
+
+            if (mlirValueIsABlockArgument(v)) {
+              MlirBlock block = mlirBlockArgumentGetOwner(self.get());
+              return py::cast(PyBlock(self.getParentOperation(), block));
+            }
+
+            assert(false && "Value must be a block argument or an op result");
           })
       .def("__eq__",
            [](PyValue &self, PyValue &other) {

diff  --git a/mlir/test/python/ir/value.py b/mlir/test/python/ir/value.py
index 4eebc53ab0c93..262896ec317f9 100644
--- a/mlir/test/python/ir/value.py
+++ b/mlir/test/python/ir/value.py
@@ -37,6 +37,21 @@ def testOpResultOwner():
     assert op.result.owner == op
 
 
+# CHECK-LABEL: TEST: testBlockArgOwner
+ at run
+def testBlockArgOwner():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  module = Module.parse(
+      r"""
+    func.func @foo(%arg0: f32) {
+      return
+    }""", ctx)
+  func = module.body.operations[0]
+  block = func.regions[0].blocks[0]
+  assert block.arguments[0].owner == block
+
+
 # CHECK-LABEL: TEST: testValueIsInstance
 @run
 def testValueIsInstance():