[MLIR] [Python] Add Operation.parent

Attribute to get the parent operation of an operation. Reviewed By: stellaraccident Differential Revision: https://reviews.llvm.org/D102981
llvm · May 24, 2021 · c65bb76 · c65bb76
1 parent 57501e5
commit c65bb76
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/mlir/lib/Bindings/Python/IRCore.cpp b/mlir/lib/Bindings/Python/IRCore.cpp
@@ -2113,6 +2113,10 @@ void mlir::python::populateIRCore(py::module &m) {
                   py::arg("successors") = py::none(), py::arg("regions") = 0,
                   py::arg("loc") = py::none(), py::arg("ip") = py::none(),
                   kOperationCreateDocstring)
+      .def_property_readonly("parent",
+                             [](PyOperation &self) {
+                               return self.getParentOperation().getObject();
+                             })
       .def("erase", &PyOperation::erase)
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyOperation::getCapsule)

diff --git a/mlir/test/python/ir/operation.py b/mlir/test/python/ir/operation.py
@@ -91,15 +91,19 @@ def walk_operations(indent, op):
         for k in range(len(block.operations)):
           child_op = block.operations[k]
           print(f"{indent}    OP {k}: {child_op}")
+          print(f"{indent}    OP {k}: parent {child_op.operation.parent.name}")
           walk_operations(indent + "      ", child_op)
 
   # CHECK: REGION 0:
   # CHECK:   BLOCK 0:
   # CHECK:     OP 0: func
+  # CHECK:     OP 0: parent module
   # CHECK:       REGION 0:
   # CHECK:         BLOCK 0:
   # CHECK:           OP 0: %0 = "custom.addi"
+  # CHECK:           OP 0: parent func
   # CHECK:           OP 1: return
+  # CHECK:           OP 1: parent func
   walk_operations("", module.operation)
 
 run(testTraverseOpRegionBlockIndices)