[PIR][Dy2St] Adapt opcallstack for D2S

paddle/fluid/pybind/pir.cc

@@ -595,13 +595,29 @@ 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(";
-        self.Print(print_stream);
-        print_stream << ")";
-        return print_stream.str();
-      });
+      .def("__repr__",
+           [](Operation &self) {
+             std::ostringstream print_stream;
+             print_stream << "Operation(";
+             self.Print(print_stream);
+             print_stream << ")";
+             return print_stream.str();
+           })
+      .def("set_opcallstack",
+           [](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.

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,36 @@ 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():
+                get_all_pir_block_ops(sub_block, 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, ops)
+        for op in ops:
             if op.has_attr(callstack_var_name):
-                callstack = op.attr(callstack_var_name)
-
+                callstack = op.attrs()[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.set_opcallstack(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

python/paddle/jit/dy2static/program_translator.py

@@ -1259,7 +1259,7 @@ def pir_from_func_spec(
                         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,