[SOT] add simulation support for user-defined iterable objects

python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py

@@ -788,8 +788,8 @@ def LOAD_ATTR(self, instr: Instruction):
                 return
         else:
             attr_name = self._code.co_names[instr.arg]
-        attr_name_var = ConstantVariable.wrap_literal(attr_name, self._graph)
         obj = self.stack.pop()
+        attr_name_var = ConstantVariable.wrap_literal(attr_name, self._graph)
         self.stack.push(
             BuiltinVariable(
                 getattr, graph=self._graph, tracker=DanglingTracker()
@@ -860,11 +860,20 @@ def LOAD_GLOBAL(self, instr: Instruction):
         if push_null and CALL_METHOD_LAYOUT_NULL_AFTER_VALUE:
             self.stack.push(NullVariable())
 
+    def load_sequence(self, obj):
+        self.stack.push(obj.get_iter())
+        # skip call
+        while self._instructions[self._lasti].opname != "RETURN_VALUE":
+            self._lasti += 1
+
     def load_method(self, method_name):
+        obj = self.stack.pop()
+        if isinstance(obj, ContainerVariable) and method_name == "__iter__":
+            self.load_sequence(obj)
+            return
         method_name_var = ConstantVariable.wrap_literal(
             method_name, self._graph
         )
-        obj = self.stack.pop()
 
         method = BuiltinVariable(
             getattr, graph=self._graph, tracker=DanglingTracker()

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

@@ -1049,6 +1049,35 @@ def main_info(self) -> dict[str, Any]:
     def get_py_value(self, allow_tensor=False) -> Any:
         return self.value
 
+    def get_iter(self):
+        """
+        To simplify the problem, we only support the case where the __iter__ method returns a list.
+        """
+        from . import (
+            BuiltinVariable,
+            ConstantVariable,
+            SequenceIterVariable,
+            UserDefinedFunctionVariable,
+        )
+
+        if not hasattr(self.value, "__iter__"):
+            return super().get_iter()
+        iter_name_var = ConstantVariable.wrap_literal("__iter__", self.graph)
+        iter_method = BuiltinVariable(
+            getattr, graph=self.graph, tracker=DanglingTracker()
+        )(self, iter_name_var)
+        # If the target object is a builtin object like list_iterator, the iter_method's fn will be a ObjectVariable instead of UserDefinedFunctionVariable.
+        if not isinstance(iter_method.fn, UserDefinedFunctionVariable):
+            return super().get_iter()
+        iter_result = iter_method()
+
+        if iter_result is None or not isinstance(
+            iter_result, SequenceIterVariable
+        ):
+            return super().get_iter()
+
+        return iter_result
+
 
 class SliceVariable(VariableBase):
     """

test/sot/test_04_list.py

@@ -259,6 +259,22 @@ def list_extend_dict():
     return l1
 
 
+class IterableWithList:
+    def __init__(self):
+        self._list = [1, 2, 3]
+
+    def __iter__(self):
+        return self._list.__iter__()
+
+
+@check_no_breakgraph
+def list_within_class(x: paddle.Tensor):
+    my_iterable = IterableWithList()
+    for i in my_iterable:
+        x += i
+    return x
+
+
 class TestListBasic(TestCaseBase):
     def test_list_basic(self):
         self.assert_results(list_getitem_int, 1, paddle.to_tensor(2))
@@ -375,6 +391,9 @@ def test_list_extend_range(self):
     def test_list_extend_dict(self):
         self.assert_results(list_extend_dict)
 
+    def test_list_within_class(self):
+        self.assert_results(list_within_class, paddle.to_tensor(1))
+
 
 if __name__ == "__main__":
     unittest.main()