[SOT] Add partial support (#72956)

Author: DrRyanHuang

python/paddle/jit/sot/opcode_translator/executor/variables/__init__.py

@@ -50,6 +50,7 @@
     NumPyApiVariable,
     PaddleApiVariable,
     PaddleLayerVariable,
+    PartialVariable,
     UserCodeVariable,
     UserDefinedFunctionVariable,
     UserDefinedGeneratorFunctionVariable,

python/paddle/jit/sot/opcode_translator/executor/variables/callable.py

@@ -21,7 +21,7 @@
 import operator
 import sys
 import types
-from functools import reduce
+from functools import partial, reduce
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -1118,7 +1118,9 @@ def call_function(self, /, *args, **kwargs):
         # do not have init function
         if self.value.__init__ is object.__init__:
             return VariableFactory.from_value(
-                new_object, self.graph, DummyTracker([self])
+                new_object,
+                self.graph,
+                DummyTracker([self, *args, *kwargs.values()]),
             )
 
         if not hasattr(self.value.__init__, "__code__"):
@@ -1138,7 +1140,7 @@ def call_function(self, /, *args, **kwargs):
         new_object_variable = VariableFactory.from_value(
             new_object,
             self.graph,
-            DummyTracker([self, *list(args), *list(kwargs.values())]),
+            DummyTracker([self, *args, *kwargs.values()]),
         )
         fn_var(new_object_variable, *args, **kwargs)
         return new_object_variable
@@ -1234,3 +1236,48 @@ def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
         if is_namedtuple_class(value):
             return NamedTupleClassVariable(value, graph, tracker)
         return None
+
+
+class PartialVariable(CallableVariable):
+    def __init__(
+        self,
+        value: partial,
+        graph: FunctionGraph,
+        tracker: Tracker,
+    ):
+        super().__init__(graph, tracker)
+        self.value = value
+
+    def get_py_value(self, allow_tensor=False):
+        return self.value
+
+    def get_py_type(self):
+        return partial
+
+    def call_function(self, /, *call_args, **call_kwargs):
+        func_variable = VariableFactory.from_value(
+            self.value.func, self.graph, GetAttrTracker(self, "func")
+        )
+        partial_args = VariableFactory.from_value(
+            self.value.args, self.graph, GetAttrTracker(self, "args")
+        )
+        partial_keywords = VariableFactory.from_value(
+            self.value.keywords, self.graph, GetAttrTracker(self, "keywords")
+        )
+        assert isinstance(func_variable, CallableVariable)
+
+        partial_keywords.get_wrapped_items().update(call_kwargs)
+
+        out = func_variable(
+            *partial_args.get_wrapped_items(),
+            *call_args,
+            **(partial_keywords.get_wrapped_items() | call_kwargs),
+        )
+
+        return out
+
+    @VariableFactory.register_from_value()
+    def from_value(value: partial, graph: FunctionGraph, tracker: Tracker):
+        if isinstance(value, partial):
+            return PartialVariable(value, graph, tracker)
+        return None

test/sot/test_functools.py

@@ -16,7 +16,9 @@
 import operator
 import unittest
 
-from test_case_base import TestCaseBase
+from test_case_base import (
+    TestCaseBase,
+)
 
 import paddle
 from paddle.jit.sot.psdb import check_no_breakgraph
@@ -54,6 +56,54 @@ def try_reduce_iter(fn, var, init=None):
     return ans
 
 
+def add(a, b, c, d=2, e=3, f=4):
+    return a, b, c, d, e, f
+
+
+@check_no_breakgraph
+def simple_partial(x=1):
+    partial_func = functools.partial(add, x)
+    out = partial_func(2, 3)
+    return out
+
+
+@check_no_breakgraph
+def simple_partial_with_two_args(x=1):
+    partial_func = functools.partial(add, x, 2)
+    out = partial_func(3)
+    return out
+
+
+@check_no_breakgraph
+def simple_partial_with_n_args(x=1):
+    partial_func = functools.partial(add, x, 2, 3, 4, 5, 6)
+    out = partial_func()
+    return out
+
+
+@check_no_breakgraph
+def simple_partial_with_n_args_kwargs():
+    partial_func = functools.partial(add, 1, 2, d=2)
+    out = partial_func(paddle.to_tensor(3), e=7)
+    return out
+
+
+@check_no_breakgraph
+def simple_partial_with_n_args_same_kwargs():
+    partial_func = functools.partial(add, 1, 2, e=2)
+    out = partial_func(3, e=7)
+    return out
+
+
+# NOTE(DrRyanHuang): Currently, SOT does not support tensor bind yet
+# because this case is not common enough.
+# We could create a PartialClassVariable to prevent breakgraph.
+def simple_partial_with_tensor_bind():
+    partial_func = functools.partial(add, 1, paddle.to_tensor(2.0), e=2)
+    out = partial_func(3, e=7)
+    return out
+
+
 @check_no_breakgraph
 def try_reduce_iter_failed(fn, var):
     it = iter(var)
@@ -97,6 +147,25 @@ def test_reduce_with_init_value(self):
     def test_reduce_with_builtin_fn(self):
         self.assert_results(try_reduce, operator.add, [2, 5, 8])
 
+    def test_simple_partial(self):
+        self.assert_results(simple_partial, 1)
+        self.assert_results(simple_partial, 2)
+
+    def test_simple_partial_with_two_argsn(self):
+        self.assert_results(simple_partial_with_two_args)
+
+    def test_simple_partial_with_n_args(self):
+        self.assert_results(simple_partial_with_n_args)
+
+    def test_simple_partial_with_n_args_kwargs(self):
+        self.assert_results(simple_partial_with_n_args_kwargs)
+
+    def test_simple_partial_with_n_args_same_kwargs(self):
+        self.assert_results(simple_partial_with_n_args_same_kwargs)
+
+    def test_simple_partial_with_tensor_bind(self):
+        self.assert_results(simple_partial_with_tensor_bind)
+
 
 if __name__ == "__main__":
     unittest.main()