[mlir][python] Update `PyOpResult.owner` to get the parent object.

Previously, this API would return the PyObjectRef, rather than the
underlying PyOperation.

Reviewed By: stellaraccident

Differential Revision: https://reviews.llvm.org/D101416

2 files changed, 14 insertions, 1 deletions

diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 0945753f9fc9..781e9aed66e9 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -1566,7 +1566,7 @@ public:
           mlirOperationEqual(self.getParentOperation()->get(),
                              mlirOpResultGetOwner(self.get())) &&
           "expected the owner of the value in Python to match that in the IR");
-      return self.getParentOperation();
+      return self.getParentOperation().getObject();
     });
     c.def_property_readonly("result_number", [](PyOpResult &self) {
       return mlirOpResultGetResultNumber(self.get());
diff --git a/mlir/test/Bindings/Python/ir_value.py b/mlir/test/Bindings/Python/ir_value.py
index 3b88fee375a0..1db9f33a65cd 100644
--- a/mlir/test/Bindings/Python/ir_value.py
+++ b/mlir/test/Bindings/Python/ir_value.py
@@ -25,3 +25,16 @@ def testCapsuleConversions():
 
 
 run(testCapsuleConversions)
+
+
+# CHECK-LABEL: TEST: testOpResultOwner
+def testOpResultOwner():
+  ctx = Context()
+  ctx.allow_unregistered_dialects = True
+  with Location.unknown(ctx):
+    i32 = IntegerType.get_signless(32)
+    op = Operation.create("custom.op1", results=[i32])
+    assert op.result.owner == op
+
+
+run(testOpResultOwner)