[PIR][Dy2St] Adapt opcallstack for D2S (#62536)

xingmingyyj · SigureMo · web-flow · commit b2e563e8fcac · 2024-03-12T10:19:27.000+08:00
---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/.flake8 b/.flake8
@@ -22,7 +22,7 @@ ignore =
     E741
 per-file-ignores =
     # These files need tabs for testing.
-    test/dygraph_to_static/test_legacy_error.py:E101
+    test/dygraph_to_static/test_error.py:E101
 
     # Ignore compare with True in sot unittest
     test/sot/test_dup_top.py:E712
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -36,7 +36,7 @@ repos:
         # Exclude some unit test files that require tabs.
         exclude: |
             (?x)^(
-                test/dygraph_to_static/test_legacy_error.py
+                test/dygraph_to_static/test_error.py
             )$
 -   repo: local
     hooks:
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -630,6 +630,7 @@ void BindOperation(py::module *m) {
       .def("as_while_op",
            [](Operation &self) { return PyWhileOp(self.dyn_cast<WhileOp>()); })
       .def("__repr__",
+
            [](Operation &self) {
              std::ostringstream print_stream;
              print_stream << "Operation(";
@@ -644,9 +645,39 @@ void BindOperation(py::module *m) {
             return ApiBuilder::Instance().GetBuilder()->Insert(op);
           },
           return_value_policy::reference)
-      .def("move_before", [](Operation &self, Operation &other) {
-        self.MoveTo(other.GetParent(), Block::Iterator{other});
-      });
+      .def("move_before",
+           [](Operation &self, Operation &other) {
+             self.MoveTo(other.GetParent(), Block::Iterator{other});
+           })
+      .def_property(
+          "callstack",
+          [](Operation &self) -> py::list {
+            py::list callstack_list;
+            pir::Attribute op_callstack = self.attribute<pir::Attribute>(
+                paddle::framework::OpProtoAndCheckerMaker::
+                    OpCreationCallstackAttrName());
+            auto op_callstack_infos = PADDLE_GET_CONST(
+                std::vector<std::string>,
+                paddle::dialect::GetAttributeData(op_callstack));
+            for (auto &op_callstack_info : op_callstack_infos) {
+              callstack_list.append(op_callstack_info);
+            }
+            return callstack_list;
+          },
+          [](Operation &self,
+             const std::vector<std::string> &callstack) -> void {
+            std::vector<pir::Attribute> op_callstack_infos;
+            for (auto str : callstack) {
+              op_callstack_infos.push_back(
+                  pir::StrAttribute::get(pir::IrContext::Instance(), str));
+            }
+
+            self.set_attribute(
+                paddle::framework::OpProtoAndCheckerMaker::
+                    OpCreationCallstackAttrName(),
+                pir::ArrayAttribute::get(pir::IrContext::Instance(),
+                                         op_callstack_infos));
+          });
   py::class_<Operation::BlockContainer> block_container(
       *m, "Operation_BlockContainer", R"DOC(
     The Operation_BlockContainer only use to walk all blocks in the operation.
diff --git a/pyproject.toml b/pyproject.toml
@@ -113,7 +113,7 @@ known-first-party = ["paddle"]
 
 [tool.ruff.lint.per-file-ignores]
 # These files need tabs for testing.
-"test/dygraph_to_static/test_legacy_error.py" = ["E101", "W191"]
+"test/dygraph_to_static/test_error.py" = ["E101", "W191"]
 # Ignore compare with True in sot unittest
 "test/sot/test_dup_top.py" = ["E712"]
 # Ignore undefined variables in CMake config and some dygraph_to_static tests
diff --git a/python/paddle/jit/dy2static/origin_info.py b/python/paddle/jit/dy2static/origin_info.py
@@ -16,7 +16,7 @@
 from collections.abc import Sequence
 
 from paddle.base import core
-from paddle.base.framework import Program
+from paddle.framework import use_pir_api
 from paddle.utils import gast
 
 from .utils import ORIGIN_INFO
@@ -269,8 +269,6 @@ def update_op_callstack_with_origin_info(program):
     Replaces op callstack information about transformed static code with original dygraph code.
     """
 
-    assert isinstance(program, Program)
-
     def get_new_op_callstack(callstack):
         """
         An example of callstack:
@@ -306,21 +304,35 @@ def get_new_op_callstack(callstack):
 
         return callstack
 
+    def get_all_pir_block_ops(block):
+        ops = []
+        for op in block.ops:
+            ops.append(op)
+            for sub_block in op.blocks():
+                ops += get_all_pir_block_ops(sub_block)
+        return ops
+
     op_maker = core.op_proto_and_checker_maker
     callstack_var_name = op_maker.kOpCreationCallstackAttrName()
 
-    for block in program.blocks:
-        for i, op in enumerate(block.ops):
+    if use_pir_api():
+        global_block = program.global_block()
+        ops = get_all_pir_block_ops(global_block)
+        for op in ops:
             if op.has_attr(callstack_var_name):
-                callstack = op.attr(callstack_var_name)
-
-                callstack = get_new_op_callstack(callstack)
-
-                try:
-                    # (@xiongkun) In 2-order derivative for paddle science, there may exists `pow_grad`
-                    # which has op_proto == nullptr and causes _set_attr failed. so we add a try...except.
-                    op._set_attr(callstack_var_name, callstack)
-                except:
-                    pass
-
+                op.callstack = get_new_op_callstack(op.callstack)
+    else:
+        for block in program.blocks:
+            for i, op in enumerate(block.ops):
+                if op.has_attr(callstack_var_name):
+                    callstack = op.attr(callstack_var_name)
+
+                    callstack = get_new_op_callstack(callstack)
+
+                    try:
+                        # (@xiongkun) In 2-order derivative for paddle science, there may exists `pow_grad`
+                        # which has op_proto == nullptr and causes _set_attr failed. so we add a try...except.
+                        op._set_attr(callstack_var_name, callstack)
+                    except:
+                        pass
     return program
diff --git a/python/paddle/jit/dy2static/program_translator.py b/python/paddle/jit/dy2static/program_translator.py
@@ -1258,8 +1258,7 @@ def pir_from_func_spec(
                     if need_wrap_into_list:
                         outputs = [outputs]
 
-        # TODO(@xiongkun): support op call stack in new ir?
-        # main_program = update_op_callstack_with_origin_info(main_program)
+        main_program = update_op_callstack_with_origin_info(main_program)
 
         return ConcreteProgram(
             inputs=static_inputs,
diff --git a/test/dygraph_to_static/CMakeLists.txt b/test/dygraph_to_static/CMakeLists.txt
@@ -80,3 +80,12 @@ foreach(ITEST ${LEGACY_ONLY_TEST_FILES})
       STATUS "PT Disabled OpTest: not found ${ITEST} in dygraph_to_static")
   endif()
 endforeach()
+
+# PIR only tests for dygraph_to_static
+set(PIR_ONLY_TEST_FILES test_error)
+foreach(ITEST ${PIR_ONLY_TEST_FILES})
+  if(TEST ${ITEST})
+    set_tests_properties(${ITEST} PROPERTIES ENVIRONMENT
+                                             "FLAGS_enable_pir_api=True")
+  endif()
+endforeach()
diff --git a/test/dygraph_to_static/check_approval.py b/test/dygraph_to_static/check_approval.py
@@ -127,6 +127,7 @@ def __init__(self, start: Location, end: Location):
     UseToStaticAsDecoratorDiagnostic: [
         "test_rollback.py",
         "test_legacy_error.py",
+        "test_error.py",
         "test_op_attr.py",
         "test_se_resnet.py",
         "test_lac.py",
@@ -142,6 +143,7 @@ def __init__(self, start: Location, end: Location):
         "test_eval_frame.py",
         "test_ignore_module.py",
         "test_legacy_error.py",
+        "test_error.py",
         "test_local_cast.py",
         "test_ordered_set.py",
         "test_origin_info.py",
diff --git a/test/dygraph_to_static/test_error.py b/test/dygraph_to_static/test_error.py
diff --git a/test/dygraph_to_static/test_legacy_error.py b/test/dygraph_to_static/test_legacy_error.py
