[SOT] GraphLogger->SubGraphInfo to collect graph info (#71412)

DrRyanHuang · web-flow · commit 6fb1e000f3ed · 2025-03-05T12:13:41.000+08:00
diff --git a/python/paddle/jit/sot/symbolic/compile_cache.py b/python/paddle/jit/sot/symbolic/compile_cache.py
@@ -27,11 +27,11 @@
 from ..utils import (
     ENV_SOT_EXPORT,
     Cache,
-    GraphLogger,
     InfoCollector,
     NewSymbolHitRateInfo,
     Singleton,
     StepInfoManager,
+    SubGraphInfo,
     SubGraphRelationInfo,
     log,
     log_do,
@@ -41,6 +41,7 @@
 from .interpreter import compile_sir
 
 if TYPE_CHECKING:
+    from paddle.framework import Program
     from paddle.static import InputSpec
 
     from .symbolic_context import SymbolicTraceContext
@@ -217,6 +218,18 @@ def collect_subgraph_relation(self, inputs, outputs, partial_program_layer):
             self.graph_size(),
         )
 
+    def collect_subgraph_info(self, program: Program):
+        if not InfoCollector().need_collect(SubGraphInfo):
+            return
+
+        InfoCollector().attach(
+            SubGraphInfo,
+            program,
+            self.graph_size(),
+            self.SIR.name,
+            self.is_first_call,
+        )
+
     def __call__(self, *args, **kwargs):
         with EventGuard(f"FallbackWrapper: {self.SIR.name}"):
             if StepInfoManager().need_back_trace:
@@ -247,18 +260,13 @@ def __call__(self, *args, **kwargs):
                 outputs = self.partial_program.sot_call(*args, **kwargs)
 
             clear_eager_tensor_name(outputs)
-            log_do(
-                1,
-                lambda: GraphLogger().add_subgraph(
-                    self.concrete_program.main_program
-                ),
-            )
             log_do(
                 4,
                 lambda: print("[CompileCache] run sir forward success."),
             )
             self.collect_new_symbol_hit_rate(args, outputs)
             self.collect_subgraph_relation(args, outputs, self.partial_program)
+            self.collect_subgraph_info(self.concrete_program.main_program)
             if ENV_SOT_EXPORT.get() != "" and not self.exported:
                 export(self.SIR, ENV_SOT_EXPORT.get())
                 self.exported = True
diff --git a/python/paddle/jit/sot/translate.py b/python/paddle/jit/sot/translate.py
@@ -23,10 +23,8 @@
 from .opcode_translator import eval_frame_callback
 from .profiler import SotStepProfilerGuard
 from .utils import (
-    GraphLogger,
     InfoCollector,
     StepInfoManager,
-    log_do,
 )
 
 P = ParamSpec("P")
@@ -100,7 +98,6 @@ def impl(*args: P.args, **kwargs: P.kwargs) -> R:
             assert hasattr(
                 fn, "__code__"
             ), "Target function doesn't have code for simulating."
-            GraphLogger().clear()
             InfoCollector().clear_step_info()
             paddle.framework.core.set_eval_frame(callback)
             try:
@@ -110,7 +107,6 @@ def impl(*args: P.args, **kwargs: P.kwargs) -> R:
             finally:
                 paddle.framework.core.set_eval_frame(None)
 
-            log_do(1, lambda: GraphLogger().print_info())
             InfoCollector().print_step_report()
             return outs
 
diff --git a/python/paddle/jit/sot/utils/__init__.py b/python/paddle/jit/sot/utils/__init__.py
@@ -52,6 +52,7 @@
     CompileCountInfo,
     InfoCollector,
     NewSymbolHitRateInfo,
+    SubGraphInfo,
     SubGraphRelationInfo,
 )
 from .magic_methods import magic_method_builtin_dispatch  # noqa: F401
@@ -65,7 +66,6 @@
 from .utils import (  # noqa: F401
     Cache,
     ConstTypes,
-    GraphLogger,
     NameGenerator,
     ResumeFnNameFactory,
     Singleton,
diff --git a/python/paddle/jit/sot/utils/info_collector.py b/python/paddle/jit/sot/utils/info_collector.py
@@ -26,13 +26,11 @@
 from .envs import ENV_SOT_COLLECT_INFO
 from .utils import Singleton
 
-if TYPE_CHECKING:
-    from .exceptions import BreakGraphReasonBase
-
-
 if TYPE_CHECKING:
     import types
 
+    from .exceptions import BreakGraphReasonBase
+
 
 def try_import_graphviz():
     try:
@@ -220,9 +218,9 @@ def to_tensor_node_name(
             dot.node(
                 subgraph_id,
                 f"Subgraph {i} ({info.subgraph_name}, size={info.graph_size})",
-                shape='oval',
-                fillcolor='cyan' if info.is_first_call else None,
-                style='filled' if info.is_first_call else None,
+                shape="oval",
+                fillcolor="cyan" if info.is_first_call else None,
+                style="filled" if info.is_first_call else None,
             )
             for shape_info in info.input_shape_infos:
                 dot.edge(
@@ -281,7 +279,6 @@ def __init__(self, reason: BreakGraphReasonBase):
 
     @classmethod
     def summary(cls, history: list[Self]) -> str:
-
         reason_dict = {}
 
         for info in history:
@@ -306,3 +303,30 @@ def collect_break_graph_reason(reason: BreakGraphReasonBase):
             return
 
         InfoCollector().attach(BreakGraphReasonInfo, reason)
+
+
+class SubGraphInfo(InfoBase):
+    SHORT_NAME = "subgraph_info"
+    TYPE = InfoType.E2E_INFO
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.clear()
+
+        self.graph, self.op_num, *_ = args
+
+    def clear(self):
+        self.graph = None
+        self.op_num = 0
+
+    def __str__(self):
+        return f"OpNum: {self.op_num}\n{self.graph}"
+
+    @classmethod
+    def summary(cls, history: list[Self]) -> str:
+        return "\n".join(
+            [
+                f"SubGraphIdx: {idx} {info}"
+                for idx, info in enumerate(map(str, history))
+            ]
+        )
diff --git a/python/paddle/jit/sot/utils/utils.py b/python/paddle/jit/sot/utils/utils.py
@@ -42,7 +42,6 @@
 
 if TYPE_CHECKING:
     from paddle._typing import NestedStructure
-    from paddle.framework import Program
 
 T = TypeVar("T")
 T1 = TypeVar("T1")
@@ -329,59 +328,6 @@ def get_unbound_method(obj, name):
     return getattr(obj.__class__, name)
 
 
-class GraphLogger(metaclass=Singleton):
-    graph_num: int
-    graphs: list[Program]
-    num_ops_per_graph: list[int]
-
-    def __init__(self):
-        self.clear()
-
-    def clear(self):
-        self.graph_num = 0
-        self.graphs = []
-        self.num_ops_per_graph = []
-
-    def get_graph_num(self):
-        return self.graph_num
-
-    def get_op_num(self):
-        return sum(self.num_ops_per_graph)
-
-    def add_subgraph(self, program: Program):
-        self.graph_num += 1
-        self.graphs.append(program)
-        sub_graph_op_num = len(program.global_block().ops)
-        self.num_ops_per_graph.append(sub_graph_op_num)
-
-    def add_subgraph_info(self, strs):
-        for i in range(len(self.graphs)):
-            strs.append(
-                "------------------------------------------------------"
-            )
-
-            strs.append(f"subgraph {i}, OpNum: {self.num_ops_per_graph[i]}")
-            strs.append(f"{self.graphs[i]}")
-
-    def __str__(self):
-        strs = []
-        strs.append("---------------- PaddleSOT graph info ----------------")
-        strs.append(f"SubgraphNum: {self.get_graph_num()}")
-        strs.append(f"OpNum: {self.get_op_num()}")
-
-        # We can display every subgraph info
-        log_do(5, lambda: self.add_subgraph_info(strs))
-
-        strs.append("---------------- PaddleSOT graph info ----------------")
-        return "\n".join(strs)
-
-    def __repr__(self):
-        return self.__str__()
-
-    def print_info(self):
-        print(self)
-
-
 class SotUndefinedVar(metaclass=Singleton):
     pass
 
