[SOT] Catch BreakGraphError when extend ListVariable with RangeVariable (#67591)

gouzil · SigureMo · web-flow · commit 93b23a2d48e7 · 2024-08-29T14:50:22.000+08:00
---------

Co-authored-by: Nyakku Shigure &lt;sigure.qaq@gmail.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
@@ -1056,7 +1056,8 @@ def build_seq_unpack(self, instr: Instruction):
 
         retval = []
         for item in unpack_values:
-            assert isinstance(item, (TupleVariable, ListVariable))
+            if not isinstance(item, (TupleVariable, ListVariable)):
+                raise BreakGraphError(f"{type(item)} not support unpack")
             retval.extend(item.get_wrapped_items())
 
         if instr.opname in {
@@ -1071,12 +1072,15 @@ def build_seq_unpack(self, instr: Instruction):
             )
         )
 
+    @call_break_graph_decorator(push_n=1)
     def BUILD_TUPLE_UNPACK_WITH_CALL(self, instr: Instruction):
         self.build_seq_unpack(instr)
 
+    @call_break_graph_decorator(push_n=1)
     def BUILD_TUPLE_UNPACK(self, instr: Instruction):
         self.build_seq_unpack(instr)
 
+    @call_break_graph_decorator(push_n=1)
     def BUILD_LIST_UNPACK(self, instr: Instruction):
         self.build_seq_unpack(instr)
 
@@ -1561,6 +1565,7 @@ def MAP_ADD(self, instr: Instruction):
             self.stack.peek[instr.arg], key, value
         )
 
+    @call_break_graph_decorator(push_n=0)
     def LIST_EXTEND(self, instr: Instruction):
         list_value = self.stack.pop()
         assert isinstance(instr.arg, int)
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py b/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py
@@ -400,7 +400,10 @@ def dispatch_tuple_ne(lhs: TupleVariable, rhs: TupleVariable):
 )
 Dispatcher.register(
     list.extend,
-    ("ListVariable", "ListVariable | TupleVariable"),
+    (
+        "ListVariable",
+        "ListVariable | TupleVariable",
+    ),
     lambda var, other: var.extend(other),
 )
 Dispatcher.register(
diff --git a/test/sot/test_04_list.py b/test/sot/test_04_list.py
@@ -247,6 +247,16 @@ def list_inplace_add():
     return l0, l1
 
 
+def list_extend_range(x):
+    return [1, *range(0, len(x.shape))]
+
+
+def list_extend_dict():
+    l1 = []
+    l1.extend({1: 2, 2: 3, 3: 4})
+    return l1
+
+
 class TestListBasic(TestCaseBase):
     def test_list_basic(self):
         self.assert_results(list_getitem_int, 1, paddle.to_tensor(2))
@@ -356,6 +366,12 @@ def test_list_add(self):
     def test_list_inplace_add(self):
         self.assert_results(list_inplace_add)
 
+    def test_list_extend_range(self):
+        self.assert_results(list_extend_range, paddle.to_tensor([1, 2]))
+
+    # def test_list_extend_dict(self):
+    #     self.assert_results(list_extend_dict)
+
 
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -400,7 +400,10 @@ def dispatch_tuple_ne(lhs: TupleVariable, rhs: TupleVariable):`
`400`	`400`	`)`
`401`	`401`	`Dispatcher.register(`
`402`	`402`	`list.extend,`
`403`		`- ("ListVariable", "ListVariable \| TupleVariable"),`
	`403`	`+ (`
	`404`	`+ "ListVariable",`
	`405`	`+ "ListVariable \| TupleVariable",`
	`406`	`+ ),`
`404`	`407`	`lambda var, other: var.extend(other),`
`405`	`408`	`)`
`406`	`409`	`Dispatcher.register(`