Provide opoperands(), opresults() methods for the Operation module (#55903)

jjyaoao · web-flow · commit a5ba0b6534d8 · 2023-08-08T14:01:44.000+08:00
* Provide opoperands(), opresults() methods for the Operation module

Signed-off-by: jjyaoao &lt;jjyaoao@126.com&gt;

* Update test_ir_pybind.py

---------

Signed-off-by: jjyaoao &lt;jjyaoao@126.com&gt;
diff --git a/paddle/fluid/pybind/ir.cc b/paddle/fluid/pybind/ir.cc
@@ -110,22 +110,8 @@ void BindOperation(py::module *m) {
       .def("operand", &Operation::operand)
       .def("result", &Operation::result)
       .def("operand_source", &Operation::operand_source)
-      .def("operands",
-           [](Operation &self) -> py::list {
-             py::list op_list;
-             for (uint32_t i = 0; i < self.num_operands(); i++) {
-               op_list.append(self.operand(i));
-             }
-             return op_list;
-           })
-      .def("results",
-           [](Operation &self) -> py::list {
-             py::list op_list;
-             for (uint32_t i = 0; i < self.num_results(); i++) {
-               op_list.append(self.result(i));
-             }
-             return op_list;
-           })
+      .def("operands", &Operation::operands)
+      .def("results", &Operation::results)
       .def("attrs",
            [](Operation &self) -> py::dict {
              py::dict attrs_dict;
diff --git a/paddle/ir/core/operation.cc b/paddle/ir/core/operation.cc
@@ -264,4 +264,20 @@ void Operation::Verify() {
   }
 }
 
+std::vector<OpOperand> Operation::operands() const {
+  std::vector<OpOperand> res;
+  for (uint32_t i = 0; i < num_operands(); ++i) {
+    res.push_back(operand(i));
+  }
+  return res;
+}
+
+std::vector<OpResult> Operation::results() const {
+  std::vector<OpResult> res;
+  for (uint32_t i = 0; i < num_results(); ++i) {
+    res.push_back(result(i));
+  }
+  return res;
+}
+
 }  // namespace ir
diff --git a/paddle/ir/core/operation.h b/paddle/ir/core/operation.h
@@ -142,6 +142,10 @@ class IR_API alignas(8) Operation final {
 
   void Verify();
 
+  std::vector<OpOperand> operands() const;
+
+  std::vector<OpResult> results() const;
+
  private:
   DISABLE_COPY_AND_ASSIGN(Operation);
   Operation(const AttributeMap &attribute,
diff --git a/test/ir/new_ir/test_ir_pybind.py b/test/ir/new_ir/test_ir_pybind.py
@@ -164,6 +164,18 @@ def test_attr(self):
         self.assertEqual(full_attr["dtype"], paddle.fluid.core.DataType.FLOAT32)
         self.assertTrue(isinstance(full_attr["place"], paddle.fluid.core.Place))
 
+    def test_operands(self):
+        newir_program = get_ir_program()
+        matmul_op = newir_program.block().get_ops()[1]
+        operands = matmul_op.operands()
+        self.assertEqual(len(operands), 2)
+
+    def test_results(self):
+        newir_program = get_ir_program()
+        matmul_op = newir_program.block().get_ops()[1]
+        results = matmul_op.results()
+        self.assertEqual(len(results), 1)
+
 
 if __name__ == "__main__":
     unittest.main()
