[SOT] Fallback when breakgraph on comprehensive (#64126)

SigureMo · web-flow · commit f085c6c9e32c · 2024-05-09T09:58:54.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
@@ -38,6 +38,7 @@
     InnerError,
     SotUndefinedVar,
     get_static_function,
+    is_comprehensive_name,
     log,
     log_do,
 )
@@ -1714,6 +1715,12 @@ def FOR_ITER(self, instr):
                 iterator.idx = backup_iter_idx
             self._graph.remove_global_guarded_variable(iterator)
             self.stack.push(iterator)
+            if is_comprehensive_name(self._code.co_name):
+                # NOTE(SigureMo): The loop body of comprehensive will access the
+                # value out of the loop, so we simply fallback it now.
+                raise FallbackError(
+                    "Comprehensive for loop break graph is not supported."
+                )
             self._break_graph_when_for_loop(iterator, instr)
             return Stop(state="BreakGraph")
 
diff --git a/python/paddle/jit/sot/utils/__init__.py b/python/paddle/jit/sot/utils/__init__.py
@@ -64,6 +64,7 @@
     is_break_graph_api,
     is_builtin_fn,
     is_clean_code,
+    is_comprehensive_name,
     is_paddle_api,
     is_strict_mode,
     list_contain_by_id,
diff --git a/python/paddle/jit/sot/utils/utils.py b/python/paddle/jit/sot/utils/utils.py
@@ -149,6 +149,10 @@ def no_eval_frame_func(*args, **kwargs):
     return no_eval_frame_func
 
 
+def is_comprehensive_name(name):
+    return name in ["<listcomp>", "<dictcomp>", "<setcomp>", "<genexpr>"]
+
+
 def is_paddle_api(func):
     if isinstance(func, paddle.nn.Layer):  # ignore all the classes
         return False
diff --git a/test/sot/test_12_for_loop.py b/test/sot/test_12_for_loop.py
@@ -295,5 +295,18 @@ def test_undefined_var_case_1(self):
         self.assert_results(undefined_var_case_1)
 
 
+def comp_with_fallback(x):
+    paddle.jit.sot.psdb.fallback()
+    y = [len(t) for t in x]
+    return paddle.to_tensor(y)
+
+
+class TestListCompWithFallback(TestCaseBase):
+    @strict_mode_guard(False)
+    def test_list_comp_with_fallback(self):
+        x = [paddle.randn([4, 6])]
+        self.assert_results(comp_with_fallback, x)
+
+
 if __name__ == "__main__":
     unittest.main()
