diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp --- a/mlir/lib/Bindings/Python/IRCore.cpp +++ b/mlir/lib/Bindings/Python/IRCore.cpp @@ -1652,7 +1652,17 @@ } 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 @@ .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;