[SOT] Move Internal Function to External Scope (#71144)

DrRyanHuang · web-flow · commit 4dea9fddc411 · 2025-02-20T09:08:37.000+08:00
diff --git a/python/paddle/jit/api.py b/python/paddle/jit/api.py
@@ -284,7 +284,7 @@ def to_static(
 
     def decorated(python_func):
         """
-        Decorates a python function into a ASTStaticFunction object.
+        Decorates a python function into a ASTStaticFunction or SymbolicStaticFunction object.
         """
 
         nonlocal full_graph
@@ -298,10 +298,9 @@ def decorated(python_func):
             )
             full_graph = True
 
-        StaticClass = {
-            False: SymbolicStaticFunction,
-            True: ASTStaticFunction,
-        }[full_graph]
+        StaticClass = (
+            ASTStaticFunction if full_graph else SymbolicStaticFunction
+        )
 
         # Step 1. unwrap the function if it is already decorated.
         _, python_func = unwrap_decorators(python_func)
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -868,6 +868,28 @@ def gen_load_inputs(
                 self.pycode_gen.gen_load_const("int64")
                 self.pycode_gen.gen_call_function(3)
 
+    @staticmethod
+    def _is_graph_output(
+        var,
+    ) -> TypeGuard[TensorVariable | SymbolicVariable]:
+        return isinstance(
+            var.tracker, (DummyTracker, SymbolicOperationTracker)
+        ) and isinstance(var, (TensorVariable, SymbolicVariable))
+
+    @staticmethod
+    def _collect_related_dummy_tensor(var):
+        if not isinstance(
+            var.tracker, (DummyTracker, SymbolicOperationTracker)
+        ):
+            return []
+        if FunctionGraph._is_graph_output(var):
+            return [var]
+
+        retval = []
+        for inp in var.tracker.inputs:
+            retval.extend(FunctionGraph._collect_related_dummy_tensor(inp))
+        return retval
+
     def _find_tensor_outputs(
         self, outputs: list[VariableBase]
     ) -> OrderedSet[TensorVariable | SymbolicVariable]:
@@ -878,26 +900,6 @@ def _find_tensor_outputs(
             outputs: output variables
         """
 
-        def is_graph_output(
-            var,
-        ) -> TypeGuard[TensorVariable | SymbolicVariable]:
-            return isinstance(
-                var.tracker, (DummyTracker, SymbolicOperationTracker)
-            ) and isinstance(var, (TensorVariable, SymbolicVariable))
-
-        def collect_related_dummy_tensor(var):
-            if isinstance(
-                var.tracker, (DummyTracker, SymbolicOperationTracker)
-            ):
-                if is_graph_output(var):
-                    return [var]
-                else:
-                    retval = []
-                    for inp in var.tracker.inputs:
-                        retval.extend(collect_related_dummy_tensor(inp))
-                    return retval
-            return []
-
         output_tensors: OrderedSet[TensorVariable | SymbolicVariable] = (
             OrderedSet()
         )
@@ -906,11 +908,13 @@ def collect_related_dummy_tensor(var):
             if isinstance(
                 output.tracker, (DummyTracker, SymbolicOperationTracker)
             ):
-                if is_graph_output(output):
+                if FunctionGraph._is_graph_output(output):
                     output_tensors.add(output)
                 else:
                     for inp in output.tracker.inputs:
-                        for _var in collect_related_dummy_tensor(inp):
+                        for _var in FunctionGraph._collect_related_dummy_tensor(
+                            inp
+                        ):
                             output_tensors.add(_var)
                     # Guard output that can not be traced.
                     self.add_global_guarded_variable(output)
@@ -919,7 +923,7 @@ def collect_related_dummy_tensor(var):
             if isinstance(side_effect_var, (ListVariable, DictVariable)):
                 for var in side_effect_var.flatten_inner_vars():
                     if (
-                        is_graph_output(var)
+                        FunctionGraph._is_graph_output(var)
                         and side_effect_var.tracker.is_traceable()
                     ):
                         output_tensors.add(var)
@@ -934,12 +938,12 @@ def collect_related_dummy_tensor(var):
                 for record in proxy_records:
                     if isinstance(record, (MutationSet, MutationNew)):
                         for var in record.value.flatten_inner_vars():
-                            if is_graph_output(var):
+                            if FunctionGraph._is_graph_output(var):
                                 output_tensors.add(var)
         # Find Tensor in print_stmts
         for print_stmt in self._print_variables:
             for var in print_stmt.flatten_inner_vars():
-                if is_graph_output(var):
+                if FunctionGraph._is_graph_output(var):
                     output_tensors.add(var)
 
         # add inplace tensors into output tensors.
diff --git a/python/paddle/jit/sot/utils/utils.py b/python/paddle/jit/sot/utils/utils.py
@@ -338,42 +338,36 @@ def get_unbound_method(obj, name):
 
 class GraphLogger(metaclass=Singleton):
     graph_num: int
-    op_num: int
     graphs: list[Program]
-    ops: list[paddle.base.framework.Operator]
+    num_ops_per_graph: list[int]
 
     def __init__(self):
         self.clear()
 
     def clear(self):
         self.graph_num = 0
-        self.op_num = 0
         self.graphs = []
-        self.ops = []
+        self.num_ops_per_graph = []
 
     def get_graph_num(self):
         return self.graph_num
 
     def get_op_num(self):
-        return self.op_num
+        return sum(self.num_ops_per_graph)
 
     def add_subgraph(self, program: Program):
         self.graph_num += 1
         self.graphs.append(program)
-
-        sub_op_num = 0
-        for op in program.global_block().ops:
-            self.op_num += 1
-            sub_op_num += 1
-        self.ops.append(sub_op_num)
+        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.ops[i]}")
+            strs.append(f"subgraph {i}, OpNum: {self.num_ops_per_graph[i]}")
             strs.append(f"{self.graphs[i]}")
 
     def __str__(self):
