[SOT] Support builtin function super (#71865)

DrRyanHuang · web-flow · commit 6f98615ec3d9 · 2025-03-28T20:44:35.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -35,7 +35,6 @@
     ENV_MIN_GRAPH_SIZE,
     ENV_SOT_FORCE_FALLBACK_SIR_IDS,
     BreakGraphError,
-    BuiltinFunctionBreak,
     FallbackError,
     InnerError,
     SotUndefinedVar,
@@ -798,10 +797,32 @@ def LOAD_ATTR(self, instr: Instruction):
 
     @call_break_graph_decorator(push_n=1)
     def LOAD_SUPER_ATTR(self, instr: Instruction):
-        # This bytecode is for Python 3.12+, and it will break graph in Python 3.11-.
-        # We align it's behavior with Python 3.11-.
-        raise BreakGraphError(
-            BuiltinFunctionBreak(reason_str="call super is not supported")
+        # Handle LOAD_SUPER_ATTR bytecode (introduced in Python 3.12+) by simulating its execution
+
+        assert isinstance(instr.arg, int)
+
+        name_idx = instr.arg >> 2  # Name index in co_names
+        is_method = bool(instr.arg & 1)  # Method binding flag
+
+        args = self.stack.pop_n(2)
+        super_func = self.stack.pop()
+        self.stack.push(super_func(*args))
+
+        attr_name = self.vframe.code.co_names[name_idx]
+
+        if is_method:
+            # Handle method binding
+            self.load_method(attr_name)
+            return
+
+        # Handle attribute lookup
+        attr_name_var = ConstantVariable.wrap_literal(attr_name, self._graph)
+        obj = self.stack.pop()
+
+        self.stack.push(
+            BuiltinVariable(
+                getattr, graph=self._graph, tracker=DanglingTracker()
+            )(obj, attr_name_var)
         )
 
     def LOAD_CONST(self, instr: Instruction):
@@ -1144,6 +1165,19 @@ def BUILD_MAP_UNPACK_WITH_CALL(self, instr: Instruction):
             )
         )
 
+    def handle_super_init_without_args(self, fn, args, kwargs):
+        if (
+            isinstance(fn, BuiltinVariable)
+            and fn.value is super
+            and len(args) == 0
+        ):
+            self_name = self.vframe.code.co_varnames[0]
+            args = (
+                self.vframe.cells['__class__'].value,
+                self.vframe.locals[self_name],
+            )
+        return fn, args, kwargs
+
     @call_break_graph_decorator(push_n=1)
     def PRECALL__CALL(self, instr: Instruction):
         """
@@ -1193,6 +1227,7 @@ def call_impl_py312_minus(
         if not is_method:
             # pop the NULL variable
             self.stack.pop()
+        fn, args, kwargs = self.handle_super_init_without_args(fn, args, kwargs)
         self.stack.push(fn(*args, **kwargs))
         self._call_shape = None
 
@@ -1262,6 +1297,7 @@ def CALL_FUNCTION(self, instr: Instruction):
         args = self.stack.pop_n(n_args)
         kwargs = {}
         fn = self.stack.pop()
+        fn, args, kwargs = self.handle_super_init_without_args(fn, args, kwargs)
         ret = fn(*args, **kwargs)
         self.stack.push(ret)
 
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py b/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py
@@ -65,6 +65,7 @@
     NumpyVariable,
     RangeVariable,
     SliceVariable,
+    SuperVariable,
     SymbolicVariable,
     TupleVariable,
     VariableBase,
@@ -236,6 +237,19 @@ def inner(*args, **kwargs):
 )
 
 
+# super
+Dispatcher.register(
+    super,
+    ("ClassVariable", "VariableBase"),
+    lambda cls, obj: SuperVariable(
+        cls=cls,
+        obj=obj,
+        graph=Dispatcher.graph,
+        tracker=DummyTracker([cls, obj]),
+    ),
+)
+
+
 @Dispatcher.register_decorator(dict)
 def dispatch_dict_kwargs(**kwargs: VariableBase):
     res_dict = {}
@@ -1461,7 +1475,7 @@ def constant_numpy_equal(left, right):
 
         @Dispatcher.register_decorator(unary_fn)
         def numpy_unary_dispatcher(var: NumpyArrayVariable):
-            raise FallbackError('Numpy operator need fallback to dygraph')
+            raise FallbackError("Numpy operator need fallback to dygraph")
 
 
 Dispatcher.register(
@@ -1476,7 +1490,7 @@ def numpy_unary_dispatcher(var: NumpyArrayVariable):
 
         @Dispatcher.register_decorator(binary_fn)
         def numpy_binary_dispatcher(var: NumpyVariable, other: NumpyVariable):
-            raise FallbackError('Numpy operator need fallback to dygraph')
+            raise FallbackError("Numpy operator need fallback to dygraph")
 
 
 Dispatcher.register(
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/__init__.py b/python/paddle/jit/sot/opcode_translator/executor/variables/__init__.py
@@ -33,6 +33,7 @@
     ObjectVariable,
     ParameterVariable,
     SliceVariable,
+    SuperVariable,
     SymbolicVariable,
     TensorVariable,
 )
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -50,6 +50,7 @@
 )
 from ....utils.envs import ENV_SOT_BREAK_GRAPH_ON_GET_SYMBOLIC_VALUE
 from ....utils.exceptions import (
+    HasNoAttributeError,
     InnerError,
     UnsupportedPaddleAPIBreak,
 )
@@ -79,7 +80,7 @@
 if TYPE_CHECKING:
     from ..function_graph import FunctionGraph
     from ..pycode_generator import PyCodeGen
-    from .callable import FunctionVariable
+    from .callable import ClassVariable, FunctionVariable
 
 
 FP_DTYPE_ABBRS = {
@@ -766,14 +767,14 @@ def getattr(self, name: str, default=None):
 
             return BuiltinVariable(
                 builtin_fn, self.graph, DanglingTracker()
-            ).bind(self, name)
+            ).bind_dangling_fn(self, name)
         elif name in get_tensor_methods():
             from .callable import TensorFunctionVariable
 
             fn_var = TensorFunctionVariable(
                 name, graph=self.graph, tracker=DanglingTracker()
             )
-            return fn_var.bind(self, name)
+            return fn_var.bind_dangling_fn(self, name)
         else:
             raise BreakGraphError(
                 UnsupportedPaddleAPIBreak(fn_name=f"Tensor.{name}")
@@ -1124,6 +1125,84 @@ def get_py_value(self, allow_tensor=False) -> Any:
         return self.value
 
 
+class SuperVariable(VariableBase):
+    """
+    Enhanced support for `super()` calls in Python.
+    The `super()` function facilitates method delegation to parent classes
+    following the method resolution order (MRO).
+
+    Args:
+        obj(Any): The object 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.
+    """
+
+    def __init__(self, cls: ClassVariable, obj: VariableBase, graph, tracker):
+        super().__init__(graph, tracker)
+        self.cls = cls
+        self.obj = obj
+
+    @property
+    def main_info(self) -> dict[str, Any]:
+        if printable(self.obj):
+            return {
+                "value": f"super({self.cls.get_py_value().__name__}, {self.obj})"
+            }
+        return {"value": f"super({self.cls.get_py_value().__name__}, self)"}
+
+    def get_py_value(self, allow_tensor=False) -> Any:
+        cls = self.cls.get_py_value()
+        obj = self.obj.get_py_value()
+        return super(cls, obj)
+
+    @check_guard
+    def make_stringified_guard(self) -> list[StringifiedExpression]:
+        guards = []
+        if self.cls.tracker.need_guard():
+            guards.extend(self.cls.make_stringified_guard())
+        if self.obj.tracker.need_guard():
+            guards.extend(self.obj.make_stringified_guard())
+        return guards
+
+    def getattr(self, name: str, default=None) -> VariableBase:
+        from .callable import FunctionVariable
+
+        mro = VariableFactory.from_value(
+            self.cls.get_py_value().__mro__,
+            self.graph,
+            GetAttrTracker(self.cls, "__mro__"),
+        )
+        # `__mro__` contains currently class, so remove it
+        super_mro = mro.get_wrapped_items()[1:]
+
+        for super_cls in super_mro:
+            if not super_cls.hasattr(name):
+                continue
+            attr = super_cls.getattr(name)
+            if isinstance(attr, FunctionVariable):
+                attr = attr.bind(self.obj, name)
+            return attr
+
+        raise HasNoAttributeError(
+            f"{self.obj.__class__.__name__} {self} has no attribute {name}"
+        )
+
+    @VariableFactory.register_from_value()
+    def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
+        if not isinstance(value, super):
+            return None
+        cls = VariableFactory.from_value(
+            value.__thisclass__, graph, DanglingTracker()
+        )
+        obj = VariableFactory.from_value(
+            value.__self__, graph, DanglingTracker()
+        )
+        super_var = SuperVariable(cls, obj, graph, tracker)
+        cls.tracker = GetAttrTracker(super_var, "__thisclass__")
+        obj.tracker = GetAttrTracker(super_var, "__self__")
+        return super_var
+
+
 class SliceVariable(VariableBase):
     """
     SliceVariable is a subclass of VariableBase used to wrap a Variable of the slice type.
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -156,13 +156,8 @@ def get_code(self) -> VariableBase:
         )
         return code_obj_var
 
-    def bind(self, instance: VariableBase, name: str):
-        method_var = MethodVariable(
-            instance,
-            self,
-            graph=self.graph,
-            tracker=GetAttrTracker(instance, name),
-        )
+    def bind_dangling_fn(self, instance: VariableBase, name: str):
+        method_var = self.bind(instance, name)
         class_var = VariableFactory.from_value(
             instance.get_py_type(),
             graph=self.graph,
@@ -172,6 +167,14 @@ def bind(self, instance: VariableBase, name: str):
         self.tracker = GetAttrTracker(class_var, name)
         return method_var
 
+    def bind(self, instance: VariableBase, name: str, class_var=None):
+        return MethodVariable(
+            instance,
+            self,
+            graph=self.graph,
+            tracker=GetAttrTracker(instance, name),
+        )
+
     make_stringified_guard = object_equal_stringified_guard
 
 
@@ -375,13 +378,8 @@ def call_function(self, /, *args, **kwargs):
             )
         return self.graph.call_tensor_method(self.method_name, *args, **kwargs)
 
-    def bind(self, instance: VariableBase, name: str):
-        method_var = MethodVariable(
-            instance,
-            self,
-            graph=self.graph,
-            tracker=GetAttrTracker(instance, name),
-        )
+    def bind_dangling_fn(self, instance: VariableBase, name: str):
+        method_var = self.bind(instance, name)
         class_var = VariableFactory.from_value(
             instance.get_py_type(),
             graph=self.graph,
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/container.py b/python/paddle/jit/sot/opcode_translator/executor/variables/container.py
@@ -438,7 +438,7 @@ def getattr(self, name: str, default=None):
             builtin_fn = method_name_to_builtin_fn[name]
             return BuiltinVariable(
                 builtin_fn, self.graph, DanglingTracker()
-            ).bind(self, name)
+            ).bind_dangling_fn(self, name)
         else:
             raise FallbackError(f"attribute {name} for list is not implemented")
 
@@ -491,7 +491,7 @@ def getattr(self, name: str, default=None):
             builtin_fn = method_name_to_builtin_fn[name]
             return BuiltinVariable(
                 builtin_fn, self.graph, DanglingTracker()
-            ).bind(self, name)
+            ).bind_dangling_fn(self, name)
         else:
             raise FallbackError(
                 f"attribute {name} for tuple is not implemented"
@@ -1013,7 +1013,7 @@ def getattr(self, name: str, default=None):
             builtin_fn = method_name_to_builtin_fn[name]
             return BuiltinVariable(
                 builtin_fn, self.graph, DanglingTracker()
-            ).bind(self, name)
+            ).bind_dangling_fn(self, name)
         else:
             raise FallbackError(f"attribute {name} for dict is not implemented")
 
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/iter.py b/python/paddle/jit/sot/opcode_translator/executor/variables/iter.py
@@ -34,7 +34,7 @@
     SotErrorBase,
     UnsupportedOperationBreak,
 )
-from ..tracker import ConstTracker, DummyTracker, GetAttrTracker
+from ..tracker import ConstTracker, DanglingTracker, DummyTracker
 from .base import VariableFactory
 from .basic import ConstantVariable
 from .callable import BuiltinVariable
@@ -382,8 +382,8 @@ def getattr(self, name: str, default=None):
             )
         if name == "send":
             return BuiltinVariable(
-                generator_send, self.graph, GetAttrTracker(self, "send")
-            ).bind(self, "send")
+                generator_send, self.graph, DanglingTracker()
+            ).bind_dangling_fn(self, "send")
         unreached()
 
     def get_py_value(self, allow_tensor=False):
diff --git a/test/sot/test_23_super.py b/test/sot/test_23_super.py

Original file line number	Diff line number	Diff line change
`@@ -33,6 +33,7 @@`
`33`	`33`	`ObjectVariable,`
`34`	`34`	`ParameterVariable,`
`35`	`35`	`SliceVariable,`
	`36`	`+ SuperVariable,`
`36`	`37`	`SymbolicVariable,`
`37`	`38`	`TensorVariable,`
`38`	`39`	`)`