[SOT][3.13] split certain LOAD_ATTR in 3.13+ & replace get_items and flatten_items withflatten_inner_vars,while replacing special cases with inheritance & remove method_name attribute in MethodVariable (#70927)

GoldenStain · web-flow · commit dcb9419abcca · 2025-01-23T10:28:23.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -419,7 +419,7 @@ def compile_graph(self, *ret_vars: VariableBase) -> CompileGraphResult:
         ret_items = [
             ret_item
             for ret_var in ret_vars
-            for ret_item in ret_var.flatten_items()
+            for ret_item in ret_var.flatten_inner_vars()
         ]
 
         symbolic_outputs = self._find_tensor_outputs(ret_items)
@@ -658,7 +658,7 @@ def try_infer_meta_fn(args, kwargs) -> Any:
                     and e.name in bound_arguments.arguments
                 ):
                     original_var = bound_arguments.arguments[e.name]
-                    flatten_vars = original_var.flatten_items()
+                    flatten_vars = original_var.flatten_inner_vars()
                     if not any(
                         isinstance(arg, SymbolicVariable)
                         for arg in flatten_vars
@@ -679,7 +679,7 @@ def try_infer_meta_fn(args, kwargs) -> Any:
                 else:
                     flatten_vars = reduce(
                         lambda x, y: (
-                            x + y.flatten_items()
+                            x + y.flatten_inner_vars()
                             if isinstance(y, VariableBase)
                             else x
                         ),
@@ -917,7 +917,7 @@ def collect_related_dummy_tensor(var):
         # Find Tensor Variables from side effects Variables.
         for side_effect_var in self.side_effects.proxy_variables:
             if isinstance(side_effect_var, (ListVariable, DictVariable)):
-                for var in side_effect_var.flatten_items():
+                for var in side_effect_var.flatten_inner_vars():
                     if (
                         is_graph_output(var)
                         and side_effect_var.tracker.is_traceable()
@@ -933,12 +933,12 @@ def collect_related_dummy_tensor(var):
                     continue
                 for record in proxy_records:
                     if isinstance(record, (MutationSet, MutationNew)):
-                        for var in record.value.flatten_items():
+                        for var in record.value.flatten_inner_vars():
                             if 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_items():
+            for var in print_stmt.flatten_inner_vars():
                 if is_graph_output(var):
                     output_tensors.add(var)
 
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -1077,7 +1077,8 @@ def BUILD_CONST_KEY_MAP(self, instr: Instruction):
         assert map_size + 1 <= len(
             self.stack
         ), f"OpExecutor want BUILD_CONST_KEY_MAP with size {map_size} + 1, but current stack do not have enough elems."
-        keys = self.stack.pop().get_items()
+        keys = self.stack.pop().get_wrapped_items()
+        keys = list(keys) if isinstance(keys, tuple) else keys
         assert len(keys) == map_size
         values = self.stack.pop_n(map_size)
         self.stack.push(self.build_map(keys, values))
@@ -2483,7 +2484,7 @@ def create_after_loop_fn():
         extra_store_vars = (
             [
                 item
-                for item in iterator.flatten_items()
+                for item in iterator.flatten_inner_vars()
                 if isinstance(item, (TensorVariable, SymbolicVariable))
             ]
             if isinstance(iterator, IterVariable)
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/base.py b/python/paddle/jit/sot/opcode_translator/executor/variables/base.py
@@ -422,22 +422,14 @@ def _reconstruct(self, codegen: PyCodeGen) -> None:
             f'{self.__class__.__name__} does not implement "_reconstruct" method'
         )
 
-    def flatten_items(self) -> list[VariableBase]:
+    def flatten_inner_vars(self) -> list[VariableBase]:
         """
         Recursively flatten the items in this container variable to a list of Variable objects.
 
         Returns:
             list[VariableBase]: Flattened items of a container variable.
         """
-        from .container import ContainerVariable
-        from .iter import IterVariable
-
-        if not isinstance(self, (ContainerVariable, IterVariable)):
-            return [self]
-        flattened_items = []
-        for item in self.get_items():
-            flattened_items.extend(item.flatten_items())
-        return flattened_items
+        return [self]
 
     def get_inputs(self) -> list[VariableBase]:
         """
@@ -501,7 +493,6 @@ def attr_proxy_getter(self, proxy: MutableDictLikeData, name: str):
                 fn=fn,
                 graph=self.graph,
                 tracker=GetAttrTracker(self, name),
-                method_name=name,
             )
 
         return VariableFactory.from_value(
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -253,8 +253,8 @@ def _reconstruct(self, codegen: PyCodeGen):
         codegen.gen_call_function(len(self.args))
         codegen.gen_pop_top()
 
-    def flatten_items(self):
-        return self.args
+    def flatten_inner_vars(self):
+        return self.args.flatten_inner_vars()
 
 
 IMPLEMENTED_TENSOR_PROPERTIES = set()
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -26,6 +26,9 @@
 )
 
 import paddle
+from paddle.jit.sot.opcode_translator.executor.variables.base import (
+    VariableBase,
+)
 
 from .... import psdb
 from ....profiler import EventGuard
@@ -66,7 +69,7 @@
     GetIterTracker,
     Tracker,
 )
-from .base import VariableBase, VariableFactory
+from .base import VariableFactory
 from .basic import (
     ConstantVariable,
     ObjectVariable,
@@ -361,7 +364,6 @@ class MethodVariable(CallableVariable):
         fn (VariableBase): The method to be wrapped.
         graph(FunctionGraph): The FunctionGraph object that this variable is associated with.
         tracker(Tracker): The Tracker object that tracks the information of this variable.
-        method_name (str): The name of the method to be wrapped.
     """
 
     def __init__(
@@ -370,13 +372,10 @@ def __init__(
         fn: VariableBase,
         graph: FunctionGraph,
         tracker: Tracker,
-        *,
-        method_name: str | None = None,
     ):
         super().__init__(graph, tracker)
         self.bound_instance = bound_instance
         self.fn = fn
-        self.method_name = method_name
 
     def get_py_value(self, allow_tensor=False):
         return self.fn.get_py_value().__get__(
@@ -385,13 +384,19 @@ def get_py_value(self, allow_tensor=False):
         )
 
     def _reconstruct(self, pycode_gen):
-        assert self.method_name is not None
-        self.tensor.reconstruct(pycode_gen)
-        pycode_gen.gen_load_attr(self.method_name)
+        # We bind the method to the instance before calling the method
+        self.fn.reconstruct(pycode_gen)
+        pycode_gen.gen_load_method("__get__")
+        self.bound_instance.reconstruct(pycode_gen)
+        pycode_gen.gen_call_function(1)
 
     def call_function(self, /, *args, **kwargs):
         return self.fn(*(self.bound_instance, *args), **kwargs)
 
+    def flatten_inner_vars(self) -> list[VariableBase]:
+        # The method's inner_vars is from its bound_instance
+        return self.bound_instance.flatten_inner_vars()
+
     @staticmethod
     def wrap_method(
         value: types.MethodType,
@@ -400,7 +405,6 @@ def wrap_method(
         tracker: Tracker,
         instance: VariableBase | None = None,
         fn: VariableBase | None = None,
-        method_name: str | None = None,
     ):
         # NOTE(SigureMo): Since the method_self need method_var as the obj
         # of the tracker, we need to temporarily set the tracker of method_self
@@ -420,7 +424,6 @@ def wrap_method(
         method_var = MethodVariable(
             instance_var,
             fn_var,
-            method_name=method_name,
             graph=graph,
             tracker=tracker,
         )
@@ -441,7 +444,8 @@ def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
     @property
     def main_info(self) -> dict[str, Any]:
         return {
-            "method": self.method_name,
+            "function": self.fn,
+            "instance": self.bound_instance,
         }
 
 
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/container.py b/python/paddle/jit/sot/opcode_translator/executor/variables/container.py
@@ -20,6 +20,9 @@
 from typing import TYPE_CHECKING, Any
 
 import paddle
+from paddle.jit.sot.opcode_translator.executor.variables.base import (
+    VariableBase,
+)
 
 from ....utils import ConstTypes
 from ....utils.exceptions import FallbackError, InnerError
@@ -39,7 +42,7 @@
     GetIterTracker,
     Tracker,
 )
-from .base import VariableBase, VariableFactory
+from .base import VariableFactory
 from .basic import ConstantVariable
 from .callable import BuiltinVariable, UserDefinedFunctionVariable
 
@@ -57,9 +60,6 @@ class ContainerVariable(VariableBase):
     def init_value(self):
         return self.value
 
-    def get_items(self) -> list[VariableBase]:
-        raise FallbackError('ContainerVariable.get_items do not implement')
-
     def get_wrapped_items(self):
         raise FallbackError(
             "ContainerVariable.get_wrapped_items do not implement"
@@ -77,6 +77,12 @@ def __bool__(self) -> bool:
     def bool(self):
         return ConstantVariable(bool(self), self.graph, DummyTracker([self]))
 
+    def flatten_inner_vars(self) -> list[VariableBase]:
+        items = []
+        for item in self.get_wrapped_items():
+            items.extend(item.flatten_inner_vars())
+        return items
+
     @check_guard
     def make_stringified_guard(self) -> list[StringifiedExpression]:
         frame_value_tracer = self.tracker.trace_value_from_frame()
@@ -168,15 +174,12 @@ def _reconstruct(self, codegen: PyCodeGen):
             Dispatcher.call(operator.getitem, self, idx).reconstruct(codegen)
         codegen.gen_build_list(size)
 
-    def get_items(self):
+    def get_wrapped_items(self):
         size = len(self)
         return [
             Dispatcher.call(operator.getitem, self, idx) for idx in range(size)
         ]
 
-    def get_wrapped_items(self):
-        return self.get_items()
-
     def get_iter(self):
         from .iter import SequenceIterVariable
 
@@ -549,14 +552,11 @@ def _reconstruct(self, codegen: PyCodeGen):
             Dispatcher.call(operator.getitem, self, idx).reconstruct(codegen)
         codegen.gen_build_tuple(size)
 
-    def get_items(self):
+    def get_wrapped_items(self):
         size = len(self)
-        return [
+        return tuple(
             Dispatcher.call(operator.getitem, self, idx) for idx in range(size)
-        ]
-
-    def get_wrapped_items(self):
-        return tuple(self.get_items())
+        )
 
     def get_iter(self):
         from .iter import SequenceIterVariable
@@ -835,19 +835,22 @@ def _reconstruct(self, codegen: PyCodeGen):
             value_var.reconstruct(codegen)
         codegen.gen_build_map(size)
 
-    def get_items(self):
-        items = []
-        for key in self.proxy.get_all().keys():
-            if not isinstance(key, ConstTypes):
-                raise InnerError(
-                    f"[{self.__class__.__name__}]: received {key} as key."
-                )
-            key_var = VariableFactory.from_value(
-                key, self.graph, tracker=ConstTracker(key)
-            )
-            value_var = self[key]
-            items.extend([key_var, value_var])
-        return items
+    def flatten_inner_vars(self):
+        items = self.get_wrapped_items()
+        return reduce(
+            operator.add,
+            (
+                key_var.flatten_inner_vars()
+                + self[key_var].flatten_inner_vars()
+                for key in items.keys()
+                for key_var in [
+                    VariableFactory.from_value(
+                        key, self.graph, tracker=ConstTracker(key)
+                    )
+                ]
+            ),
+            [],
+        )
 
     def get_wrapped_items(self):
         items = {}
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/iter.py b/python/paddle/jit/sot/opcode_translator/executor/variables/iter.py
@@ -16,9 +16,13 @@
 
 from typing import TYPE_CHECKING, Any
 
+from paddle.jit.sot.opcode_translator.executor.variables.base import (
+    VariableBase,
+)
+
 from ....utils import BreakGraphError, FallbackError
 from ..tracker import ConstTracker, DummyTracker
-from .base import VariableBase, VariableFactory
+from .base import VariableFactory
 from .basic import ConstantVariable
 from .container import ContainerVariable, TupleVariable
 
@@ -50,10 +54,19 @@ def next(self):
     def get_iter(self):
         return self
 
-    def get_items(self):
-        if isinstance(self.hold, (ContainerVariable, IterVariable)):
-            return self.hold.get_items()
-        return [self.hold]
+    def flatten_inner_vars(self) -> list[VariableBase]:
+        return [
+            var
+            for item_list in (
+                self.hold.get_wrapped_items()
+                if isinstance(self.hold, (ContainerVariable, IterVariable))
+                else [self.hold]
+            )
+            for item in (
+                item_list if isinstance(item_list, list) else [item_list]
+            )
+            for var in item.flatten_inner_vars()
+        ]
 
 
 class SequenceIterVariable(IterVariable):
@@ -162,15 +175,6 @@ class ZipVariable(SequenceIterVariable):
     def __init__(self, iters, graph, tracker):
         super().__init__(iters, graph, tracker)
 
-    def get_items(self):
-        items = []
-        for hold in self.hold:
-            if isinstance(hold, (ContainerVariable, IterVariable)):
-                items.extend(hold.get_items())
-            else:
-                items.append(hold)
-        return items
-
     def next(self):
         # can not use <listcomp> here, because it will raise a RuntimeError("StopIteration")
         # but we want a StopIteration Exception
diff --git a/python/paddle/jit/sot/opcode_translator/instruction_utils/instruction_utils.py b/python/paddle/jit/sot/opcode_translator/instruction_utils/instruction_utils.py
diff --git a/test/sot/test_min_graph_size.py b/test/sot/test_min_graph_size.py
