[SOT] Support use iterable for call with varargs (#71377)

SigureMo · Copilot · web-flow · commit 7984d738edbd · 2025-03-04T14:30:41.000+08:00
---------

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -1306,9 +1306,17 @@ def CALL_FUNCTION_EX(self, instr: Instruction):
             kwargs = {}
 
         args_variable = self.stack.pop()
-        assert isinstance(args_variable, (TupleVariable, ListVariable))
-        args = args_variable.get_wrapped_items()
-
+        args_iter = args_variable.get_iter()
+        assert isinstance(
+            args_iter, IterVariable
+        ), f"args_iter should be IterVariable, but got {args_iter}"
+        if not isinstance(args_iter, SequenceIterVariable):
+            raise BreakGraphError(
+                UnsupportedOperationBreak(
+                    reason_str="CALL_FUNCTION_EX only supports SequenceIterVariable for varargs"
+                )
+            )
+        args = args_iter.to_list()
         if sys.version_info >= (3, 11) and CALL_METHOD_LAYOUT_NULL_AFTER_VALUE:
             null = self.stack.pop()
             assert isinstance(null, NullVariable)
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/iter.py b/python/paddle/jit/sot/opcode_translator/executor/variables/iter.py
@@ -56,7 +56,7 @@ def get_iter(self):
 
     def flatten_inner_vars(self) -> list[VariableBase]:
         holds = self.hold
-        if not isinstance(holds, list):
+        if not isinstance(holds, (list, tuple)):
             holds = [holds]
         return [
             inner_var
@@ -188,7 +188,13 @@ def to_list(self):
         min_len = min(len(l) for l in lists)
         result = []
         for i in range(min_len):
-            result.append(tuple(l[i] for l in lists))
+            result.append(
+                VariableFactory.from_value(
+                    tuple(l[i] for l in lists),
+                    self.graph,
+                    DummyTracker(list(self.hold)),
+                )
+            )
         return result
 
     def has_side_effect(self) -> bool:
@@ -284,6 +290,6 @@ def __init__(self, obj, graph, tracker):
     def next(self):
         raise BreakGraphError(
             UnsupportedOperationBreak(
-                reason_str="Break graph when using user defined iterator"
+                reason_str="Break graph when iterating user defined iterator"
             )
         )
diff --git a/test/sot/test_builtin_map.py b/test/sot/test_builtin_map.py
@@ -105,6 +105,14 @@ def test_map_multi_input(func, tensor_, tuple_):
     return x
 
 
+@check_no_breakgraph
+def test_map_with_zip_and_call_fn_ex(x: list[tuple[int, int]]):
+    fn = lambda x: (0,) + x  # noqa: RUF005
+    map_results = map(fn, x)
+    res = tuple(map(list, zip(*map_results)))
+    return res
+
+
 class TestMap(TestCaseBase):
     @test_with_faster_guard
     def test_map(self):
@@ -140,6 +148,12 @@ def test_map_unpack(self):
     def test_map_for_loop(self):
         self.assert_results(test_map_for_loop, [7, 8, 9, 10])
 
+    @test_with_faster_guard
+    def test_map_with_zip_and_call_fn_ex(self):
+        self.assert_results(
+            test_map_with_zip_and_call_fn_ex, [(1, 2), (3, 4), (5, 6)]
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
