[Mlir-commits] [mlir] 5664c5e - [MLIR] [Python] Add `owner` to PyValue and fix its parent reference

[John Demme]

Author: John Demme
Date: 2021-07-14T20:32:43-07:00
New Revision: 5664c5e24ed42f58175428c5aeb53418b4ff76b0

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

LOG: [MLIR] [Python] Add `owner` to PyValue and fix its parent reference

Adds `owner` python call to `mlir.ir.Value`.

Assuming that `PyValue.parentOperation` is intended to be the value's owner, this fixes the construction of it from `PyOpOperandList`.

Reviewed By: stellaraccident

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

Added: 
    

Modified: 
    mlir/lib/Bindings/Python/IRCore.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
index 3f08522b9c1de..b5197d9aa05e5 100644
--- a/mlir/lib/Bindings/Python/IRCore.cpp
+++ b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -1652,7 +1652,17 @@ class PyOpOperandList : public Sliceable<PyOpOperandList, PyValue> {
   }
 
   PyValue getElement(intptr_t pos) {
-    return PyValue(operation, mlirOperationGetOperand(operation->get(), pos));
+    MlirValue operand = mlirOperationGetOperand(operation->get(), pos);
+    MlirOperation owner;
+    if (mlirValueIsAOpResult(operand))
+      owner = mlirOpResultGetOwner(operand);
+    else if (mlirValueIsABlockArgument(operand))
+      owner = mlirBlockGetParentOperation(mlirBlockArgumentGetOwner(operand));
+    else
+      assert(false && "Value must be an block arg or op result.");
+    PyOperationRef pyOwner =
+        PyOperation::forOperation(operation->getContext(), owner);
+    return PyValue(pyOwner, operand);
   }
 
   PyOpOperandList slice(intptr_t startIndex, intptr_t length, intptr_t step) {
@@ -2429,6 +2439,15 @@ void mlir::python::populateIRCore(py::module &m) {
       .def(
           "dump", [](PyValue &self) { mlirValueDump(self.get()); },
           kDumpDocstring)
+      .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();
+          })
       .def("__eq__",
            [](PyValue &self, PyValue &other) {
              return self.get().ptr == other.get().ptr;