[SOT][PIR] Adapt SOT call stack in PIR mode (#62680)

SigureMo · web-flow · commit 17c964fa8023 · 2024-03-13T16:39:29.000+08:00
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -656,11 +656,23 @@ void BindOperation(py::module *m) {
             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);
+            PADDLE_ENFORCE(op_callstack.isa<pir::ArrayAttribute>(),
+                           phi::errors::PreconditionNotMet(
+                               "The callstack of operation `%s` should be an "
+                               "array attribute.",
+                               self.name()));
+            auto op_callstack_array_attr =
+                op_callstack.dyn_cast<pir::ArrayAttribute>();
+            for (size_t i = 0; i < op_callstack_array_attr.size(); ++i) {
+              PADDLE_ENFORCE(
+                  op_callstack_array_attr.at(i).isa<pir::StrAttribute>(),
+                  phi::errors::PreconditionNotMet(
+                      "The callstack info of operation `%s` should be array of "
+                      "string attribute.",
+                      self.name()));
+              callstack_list.append(op_callstack_array_attr.at(i)
+                                        .dyn_cast<pir::StrAttribute>()
+                                        .AsString());
             }
             return callstack_list;
           },
diff --git a/python/paddle/jit/sot/symbolic/interpreter.py b/python/paddle/jit/sot/symbolic/interpreter.py
@@ -19,7 +19,7 @@
 import paddle
 from paddle.utils import to_sequence
 
-from ..utils import InnerError, map_if, map_if_extend
+from ..utils import InnerError, log_do, map_if, map_if_extend
 from .statement_ir import SIRRuntimeCache, Symbol
 
 if TYPE_CHECKING:
@@ -51,16 +51,17 @@ def replace_symbol(
 
 
 def _append_opstack_between(start, end, stack):
-    # NOTE(xiongkun): we don't sync for speed. careful!!
-    # [start, end)
-    if paddle.base.framework.use_pir_api():
-        return
+    # The range is [start, end)
     from paddle.framework import core
 
     op_maker = core.op_proto_and_checker_maker
     callstack_attr_name = op_maker.kOpCreationCallstackAttrName()
     for op in for_each_ops_between(start, end):
-        op._set_attr(callstack_attr_name, stack)
+        if paddle.framework.use_pir_api():
+            op.callstack = stack
+        else:
+            # NOTE(xiongkun): we don't sync for speed. careful!!
+            op._set_attr(callstack_attr_name, stack)
 
 
 def for_each_ops_between(start, end):
@@ -121,8 +122,11 @@ def _set(v, s):
             if len(to_sequence(outs)) != len(to_sequence(stmt.outputs)):
                 raise InnerError("Number output mismatch, some error happen.")
 
-            _append_opstack_between(
-                before_stmt_opnum, opnum_in_program() + 1, stmt.stmt_stack
+            log_do(
+                3,
+                lambda: _append_opstack_between(
+                    before_stmt_opnum, opnum_in_program() + 1, stmt.stmt_stack
+                ),
             )
 
             map_if(
