[PIR]get value from op id (#69909)

risemeup1 · web-flow · commit 39cab51ba6e2 · 2024-12-05T10:52:42.000+08:00
* get value from op id

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -3520,7 +3520,9 @@ void ProcessBlock(
 std::unique_ptr<pir::Program> PdOpLowerToKernelPass(pir::Program* prog,
                                                     phi::Place place) {
   auto program = std::make_unique<pir::Program>(pir::IrContext::Instance());
-
+  if (FLAGS_print_ir) {
+    std::cout << "IR before lowering = " << *prog << std::endl;
+  }
   auto block = prog->block();
 
   pir::IrContext* ctx = pir::IrContext::Instance();
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -513,6 +513,40 @@ void BindProgram(py::module *m) {
             return op_list;
           },
           return_value_policy::reference)
+      .def(
+          "get_value_by_op_id",
+          [](Program &self, py::object op_ids) {
+            std::vector<int> op_ids_list;
+            if (py::isinstance<py::int_>(op_ids)) {
+              op_ids_list.push_back(op_ids.cast<int>());
+            } else if (py::isinstance<py::list>(op_ids)) {
+              for (auto item : op_ids) {
+                op_ids_list.push_back(item.cast<int>());
+              }
+            } else {
+              PADDLE_THROW(
+                  "Invalid op_ids format. Please provide either a single "
+                  "integer or a list of integers.");
+            }
+
+            std::list<Operation *> all_ops = self.block()->get_recursive_ops();
+            std::vector<pir::Value> value_list;
+
+            for (auto op : all_ops) {
+              if (std::find(op_ids_list.begin(), op_ids_list.end(), op->id()) !=
+                  op_ids_list.end()) {
+                for (auto value : op->results()) {
+                  value_list.push_back(value);
+                }
+              }
+            }
+
+            if (value_list.empty()) {
+              PADDLE_THROW(
+                  "Can't find the corresponding opresult from the op ids");
+            }
+            return value_list;
+          })
       .def("get_output_value_by_name",
            [](Program &self, const std::string &name) {
              return name_analysis::GetOutputValueByName(self, name);
diff --git a/paddle/pir/include/core/block.h b/paddle/pir/include/core/block.h
@@ -176,6 +176,12 @@ class IR_API Block {
     return num;
   }
 
+  OpListType get_recursive_ops() {
+    OpListType ops;
+    Walk([&ops](Operation *op) { ops.push_back(op); });
+    return ops;
+  }
+
  private:
   Block(Block &) = delete;
   Block &operator=(const Block &) = delete;
diff --git a/paddle/pir/src/dialect/shape/utils/shape_analysis.cc b/paddle/pir/src/dialect/shape/utils/shape_analysis.cc
@@ -804,7 +804,6 @@ pir::PrintHooks ShapeConstraintIRAnalysis::PrintHook() {
       }
     }
     printer.os << " }";
-    printer.os << "\t(op_" << op.id() << ")";
   };
   return print_hook;
 }
diff --git a/test/ir/pir/test_build_op.py b/test/ir/pir/test_build_op.py
@@ -175,5 +175,40 @@ def test_build_tensorrt_engine_op(self):
             )
 
 
+class TestGetValueByOpId(unittest.TestCase):
+    def test_get_value_by_op_id(self):
+        def true_func():
+            return paddle.tensor.fill_constant(
+                shape=[2, 3], dtype='int32', value=2
+            )
+
+        def false_func():
+            return paddle.tensor.fill_constant(
+                shape=[3, 2], dtype='int32', value=-1
+            )
+
+        main_program = paddle.static.Program()
+        startup_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program, startup_program):
+            x = paddle.tensor.fill_constant(
+                shape=[1], dtype='float32', value=0.1
+            )
+            y = paddle.tensor.fill_constant(
+                shape=[1], dtype='float32', value=0.23
+            )
+            pred = paddle.less_than(y, x)
+            out = paddle.static.nn.cond(pred, true_func, false_func)
+            value1 = main_program.get_value_by_op_id(65)
+            self.assertEqual(
+                out.get_defining_op().id(),
+                value1[0].get_defining_op().id(),
+            )
+            value2 = main_program.get_value_by_op_id([58, 65])
+            self.assertEqual(
+                58,
+                value2[0].get_defining_op().id(),
+            )
+
+
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -804,7 +804,6 @@ pir::PrintHooks ShapeConstraintIRAnalysis::PrintHook() {`
`804`	`804`	`}`
`805`	`805`	`}`
`806`	`806`	`printer.os << " }";`
`807`		`- printer.os << "\t(op_" << op.id() << ")";`
`808`	`807`	`};`
`809`	`808`	`return print_hook;`
`810`	`809`	`}`