[SOT][3.13] Support closure (#69753)

gouzil · web-flow · commit 9d49a1160932 · 2024-11-29T01:08:17.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
@@ -101,6 +101,7 @@
     SymbolicVariable,
     TensorVariable,
     TupleVariable,
+    UserCodeVariable,
     UserDefinedFunctionVariable,
     UserDefinedGeneratorFunctionVariable,
     VariableBase,
@@ -1384,14 +1385,28 @@ def MAKE_FUNCTION(self, instr: Instruction):
 
         # the function has no default values in 3.13
         if sys.version_info >= (3, 13):
-            flag = 0
-        else:
-            flag = instr.arg
+            if len(codeobj.get_py_value().co_freevars) > 0:
+                self.stack.push(
+                    UserCodeVariable(
+                        codeobj, self._graph, DummyTracker([codeobj])
+                    )
+                )
+            else:
+                self.push_new_fn_on_stack(
+                    codeobj.get_py_value(),
+                    global_dict,
+                    fn_name.get_py_value(),
+                    (),
+                    (),
+                    related_list,
+                    (),
+                )
+            return
 
+        flag = instr.arg
         closure, related_list, kw_defaults, default_args = (
             self.attach_new_attribute(flag, related_list)
         )
-
         self.push_new_fn_on_stack(
             codeobj.get_py_value(),
             global_dict,
@@ -1404,13 +1419,36 @@ def MAKE_FUNCTION(self, instr: Instruction):
 
     def SET_FUNCTION_ATTRIBUTE(self, instr: Instruction):
         origin_func = self.stack.pop()
+        flag = instr.arg
+
+        if isinstance(origin_func, UserCodeVariable):
+            origin_codeobj = origin_func.codeobj
+            fn_name = ConstantVariable(
+                origin_codeobj.value.co_qualname,
+                self._graph,
+                DummyTracker([origin_codeobj]),
+            )
+            related_list = [fn_name, origin_codeobj]
+            closure, related_list, kw_defaults, default_args = (
+                self.attach_new_attribute(flag, related_list)
+            )
+            self.push_new_fn_on_stack(
+                origin_codeobj.get_py_value(),
+                self._globals.get_value(),
+                fn_name.get_py_value(),
+                default_args,
+                closure,
+                related_list,
+                kw_defaults,
+            )
+            return
+
         # The object we manipulate must be a functionVariable
         assert isinstance(
             origin_func,
             (UserDefinedGeneratorFunctionVariable, UserDefinedFunctionVariable),
         ), f"The object we manipulate must be a function object. But now got {type(origin_func)}"
         origin_func_val = origin_func.get_py_value()
-        flag = instr.arg
         related_list = [origin_func]
         closure, related_list, kw_defaults, default_args = (
             self.attach_new_attribute(flag, related_list)
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/__init__.py b/python/paddle/jit/sot/opcode_translator/executor/variables/__init__.py
@@ -44,6 +44,7 @@
     MethodVariable,
     PaddleApiVariable,
     PaddleLayerVariable,
+    UserCodeVariable,
     UserDefinedFunctionVariable,
     UserDefinedGeneratorFunctionVariable,
     UserDefinedLayerVariable,
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -44,6 +44,7 @@
 from ....utils.exceptions import (
     BreakGraphError,
     FallbackError,
+    InnerError,
     SotErrorBase,
 )
 from ..dispatcher import Dispatcher
@@ -65,7 +66,12 @@
     Tracker,
 )
 from .base import VariableBase, VariableFactory
-from .basic import ConstantVariable, PrintStmtVariable, SliceVariable
+from .basic import (
+    ConstantVariable,
+    ObjectVariable,
+    PrintStmtVariable,
+    SliceVariable,
+)
 
 if TYPE_CHECKING:
     from ..function_graph import FunctionGraph
@@ -230,6 +236,22 @@ def main_info(self) -> dict[str, Any]:
         }
 
 
+class UserCodeVariable(FunctionVariable):
+    """
+    UserCodeVariable is a subclass of Function
+    Variable used to wrap a make function variable.
+    """
+
+    def __init__(
+        self, codeobj: ObjectVariable, graph: FunctionGraph, tracker: Tracker
+    ):
+        super().__init__(codeobj, graph, tracker)
+        self.codeobj = codeobj
+
+    def call_function(self, /, *args, **kwargs):
+        raise InnerError("UserCodeVariable call_function is not implemented.")
+
+
 class PaddleApiVariable(FunctionVariable):
     """
     PaddleApiVariable is a subclass of FunctionVariable used to wrap a paddlepaddle API function.
diff --git a/test/sot/skip_files_py313 b/test/sot/skip_files_py313
diff --git a/test/sot/test_13_make_function.py b/test/sot/test_13_make_function.py
@@ -73,8 +73,7 @@ def test_simple(self):
         self.assert_results(make_fn_default, paddle.to_tensor(1))
         self.assert_results(make_fn_annotation, paddle.to_tensor(1))
         self.assert_results(make_fn_kwdefault, paddle.to_tensor(1))
-        # self.assert_results(make_fn_closure, paddle.to_tensor(1))
-        # we haven't pass this test yet
+        self.assert_results(make_fn_closure, paddle.to_tensor(1))
         self.assert_results(make_fn_mix, paddle.to_tensor(1))
 
 
