[SOT] make sure parameter holds in IterVariable.__init__() is a list (#71437)

GoldenStain · web-flow · commit a21bafa8b621 · 2025-03-06T19:19:50.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/iter.py b/python/paddle/jit/sot/opcode_translator/executor/variables/iter.py
@@ -40,13 +40,18 @@ class IterVariable(VariableBase):
     """
 
     def __init__(
-        self, obj: VariableBase, graph: FunctionGraph, tracker: Tracker
+        self, holds: list[VariableBase], graph: FunctionGraph, tracker: Tracker
     ):
+
         super().__init__(graph, tracker)
-        self.hold = obj
+        self.hold = holds
 
     def make_stringified_guard(self):
-        return self.hold.make_stringified_guard()
+        return [
+            result
+            for holds in self.hold
+            for result in holds.make_stringified_guard()
+        ]
 
     def next(self):
         raise NotImplementedError(f"Can not simulate `next` for {type(self)}")
@@ -56,8 +61,6 @@ def get_iter(self):
 
     def flatten_inner_vars(self) -> list[VariableBase]:
         holds = self.hold
-        if not isinstance(holds, (list, tuple)):
-            holds = [holds]
         return [
             inner_var
             for hold in holds
@@ -79,15 +82,23 @@ class SequenceIterVariable(IterVariable):
 
     mutable_attrs = ["idx"]
 
-    def __init__(self, obj, graph: FunctionGraph, tracker: Tracker):
-        super().__init__(obj, graph, tracker)
+    def __init__(
+        self,
+        holds: VariableBase | list[VariableBase],
+        graph: FunctionGraph,
+        tracker: Tracker,
+    ):
+        if not isinstance(holds, list):
+            holds = [holds]
+        super().__init__(holds, graph, tracker)
         self.idx = 0
         self.graph.side_effects.record_mutable_variable(self)
 
     def next(self):
+        holds = self.hold[0]
         # TODO: self.hold should have a __len__ method
-        if self.idx < len(self.hold):
-            val = self.hold[self.idx]
+        if self.idx < len(holds):
+            val = holds[self.idx]
             self.idx += 1
             return val
         else:
@@ -96,10 +107,11 @@ def next(self):
     def to_list(self) -> list:
         if self.has_side_effect():
             raise FallbackError("Can not convert an used iterator into list")
-        self.idx = len(self.hold)
+        holds = self.hold[0]
+        self.idx = len(holds)
         retval = []
-        for i in range(len(self.hold)):
-            retval.append(self.hold[i])
+        for i in range(len(holds)):
+            retval.append(holds[i])
         return retval
 
     def has_side_effect(self) -> bool:
@@ -109,7 +121,7 @@ def _reconstruct(self, codegen: PyCodeGen):
         if self.has_side_effect():
             super()._reconstruct(codegen)
         else:
-            self.hold.reconstruct(codegen)
+            self.hold[0].reconstruct(codegen)
             codegen.gen_get_iter()
 
     @property
@@ -124,34 +136,36 @@ class EnumerateVariable(SequenceIterVariable):
     EnumerateVariable holds a SequenceIterVariable and return additional index
     """
 
-    def __init__(self, val_iterator, graph, tracker):
+    def __init__(
+        self, val_iterator: IterVariable, graph: FunctionGraph, tracker: Tracker
+    ):
         super().__init__(val_iterator, graph, tracker)
 
     def next(self):
-        val = self.hold.next()
+        val = self.hold[0].next()
         idx_var = ConstantVariable(self.idx, self.graph, ConstTracker(self.idx))
         self.idx += 1
         return TupleVariable(
             (idx_var, val), self.graph, DummyTracker([idx_var, val])
         )
 
     def to_list(self):
-        values = self.hold.to_list()
+        values = self.hold[0].to_list()
         idx = [
             ConstantVariable(i, self.graph, ConstTracker(i))
             for i in range(len(values))
         ]
         return list(zip(idx, values))
 
     def has_side_effect(self) -> bool:
-        return self.hold.has_side_effect()
+        return self.hold[0].has_side_effect()
 
     def _reconstruct(self, codegen: PyCodeGen):
         if self.has_side_effect():
             super()._reconstruct(codegen)
         else:
             codegen.gen_load_global("enumerate", push_null=True)
-            self.hold.reconstruct(codegen)
+            self.hold[0].reconstruct(codegen)
             codegen.gen_call_function(1)
 
     @staticmethod
@@ -168,7 +182,9 @@ class ZipVariable(SequenceIterVariable):
     ZipVariable holds a list of SequenceIterVariable
     """
 
-    def __init__(self, iters, graph, tracker):
+    def __init__(
+        self, iters: list[IterVariable], graph: FunctionGraph, tracker: Tracker
+    ):
         super().__init__(iters, graph, tracker)
 
     def next(self):
@@ -232,11 +248,13 @@ class MapVariable(SequenceIterVariable):
     MapVariable holds a SequenceIterVariable and return a Iterable Variable after map function
     """
 
-    def __init__(self, fn, iters: TupleVariable, graph, tracker):
+    def __init__(self, fn, iters: list[IterVariable], graph, tracker):
+
         super().__init__(iters, graph, tracker)
         self.fn = fn
 
     def next(self):
+
         return self.fn(*[iter_var.next() for iter_var in self.hold])
 
     def to_list(self) -> list:
@@ -256,8 +274,9 @@ def _reconstruct(self, codegen: PyCodeGen):
         else:
             codegen.gen_load_global("map", push_null=True)
             self.fn.reconstruct(codegen)
-            self.hold.reconstruct(codegen)
-            codegen.gen_call_function(2)
+            for iter_var in self.hold:
+                iter_var.reconstruct(codegen)
+            codegen.gen_call_function(len(self.hold) + 1)
 
     @staticmethod
     def from_iterator(
@@ -279,8 +298,15 @@ def from_iterator(
 
 # what UserDefinedIterVariable holds doesn't matter, because use user defined iterator will trigger break graph
 class UserDefinedIterVariable(IterVariable):
-    def __init__(self, obj, graph, tracker):
-        super().__init__(obj, graph, tracker)
+    def __init__(
+        self,
+        holds: VariableBase | list[VariableBase],
+        graph: FunctionGraph,
+        tracker: Tracker,
+    ):
+        if not isinstance(holds, list):
+            holds = [holds]
+        super().__init__(holds, graph, tracker)
 
     def next(self):
         raise BreakGraphError(
