[SOT] Add dispatch eq, ne, iadd for list and tuple (#67746)

SigureMo · web-flow · commit 02a98601bc59 · 2024-08-28T10:36:36.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py b/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py
@@ -335,6 +335,38 @@ def dispatch_dict_fromkeys(
     ("TupleVariable", "VariableBase"),
     lambda var, value: var.index(value),
 )
+Dispatcher.register(
+    operator.add,
+    ("TupleVariable", "TupleVariable"),
+    lambda var, other: var.concat(other),
+)
+Dispatcher.register(
+    operator.iadd,
+    ("TupleVariable", "TupleVariable"),
+    lambda var, other: var.concat(other),
+)
+
+
+@Dispatcher.register_decorator(operator.eq)
+def dispatch_tuple_eq(lhs: TupleVariable, rhs: TupleVariable):
+    if len(lhs) != len(rhs):
+        return ConstantVariable(False, lhs.graph, DummyTracker([lhs, rhs]))
+    size = len(lhs)
+
+    return ConstantVariable(
+        all(
+            Dispatcher.call(operator.eq, lhs[i], rhs[i]).get_py_value()
+            for i in range(size)
+        ),
+        lhs.graph,
+        DummyTracker([lhs, rhs]),
+    )
+
+
+@Dispatcher.register_decorator(operator.ne)
+def dispatch_tuple_ne(lhs: TupleVariable, rhs: TupleVariable):
+    return Dispatcher.call(operator.eq, lhs, rhs).bool_not()
+
 
 # list
 Dispatcher.register(
@@ -427,16 +459,38 @@ def dispatch_dict_fromkeys(
     lambda var, other: var.concat(other),
 )
 Dispatcher.register(
-    operator.add,
-    ("TupleVariable", "TupleVariable"),
-    lambda var, other: var.concat(other),
+    operator.iadd,
+    ("ListVariable", "ListVariable"),
+    lambda var, other: var.inplace_concat(other),
 )
 Dispatcher.register(
     operator.mul,
     ("ListVariable | TupleVariable", "ConstantVariable"),
     lambda var, other: var.repeat(other),
 )
 
+
+@Dispatcher.register_decorator(operator.eq)
+def dispatch_list_eq(lhs: ListVariable, rhs: ListVariable):
+    if len(lhs) != len(rhs):
+        return ConstantVariable(False, lhs.graph, DummyTracker([lhs, rhs]))
+    size = len(lhs)
+
+    return ConstantVariable(
+        all(
+            Dispatcher.call(operator.eq, lhs[i], rhs[i]).get_py_value()
+            for i in range(size)
+        ),
+        lhs.graph,
+        DummyTracker([lhs, rhs]),
+    )
+
+
+@Dispatcher.register_decorator(operator.ne)
+def dispatch_list_ne(lhs: ListVariable, rhs: ListVariable):
+    return Dispatcher.call(operator.eq, lhs, rhs).bool_not()
+
+
 # getattr
 Dispatcher.register(
     getattr,
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/container.py b/python/paddle/jit/sot/opcode_translator/executor/variables/container.py
@@ -266,6 +266,11 @@ def concat(self, list_):
             DummyTracker([self, list_]),
         )
 
+    def inplace_concat(self, list_):
+        assert isinstance(list_, ListVariable)
+        self.extend(list_)
+        return self
+
     def repeat(self, length):
         assert isinstance(length, ConstantVariable)
         return ListVariable(
diff --git a/test/sot/test_03_tuple.py b/test/sot/test_03_tuple.py
@@ -65,6 +65,31 @@ def tuple_index_tensor(x: paddle.Tensor, y: tuple[paddle.Tensor]):
     return y.index(x)
 
 
+@check_no_breakgraph
+def tuple_compare():
+    # TODO(SigureMo): support gt, ge, lt, le
+    l1 = (1, 2, 3)
+    l2 = (1, 2, 3)
+    l3 = (1, 2, 4)
+    return l1 == l2, l1 == l3, l1 != l2, l1 != l3
+
+
+@check_no_breakgraph
+def tuple_add():
+    l0 = (1, 2, 3)
+    l1 = (4, 5, 6)
+    return l0 + l1
+
+
+@check_no_breakgraph
+def tuple_inplace_add():
+    l0 = (1, 2, 3)
+    l1 = l0
+    l2 = (4, 5, 6)
+    l0 += l2
+    return l0, l1
+
+
 class TestBuildTuple(TestCaseBase):
     def test_build_tuple(self):
         self.assert_results(build_tuple, 1, paddle.to_tensor(2))
@@ -88,6 +113,15 @@ def test_tuple_methods_tensor(self):
         self.assert_results(tuple_count_tensor, a, (a, b, a, b))
         self.assert_results(tuple_index_tensor, b, (b, b, b, a))
 
+    def test_tuple_compare(self):
+        self.assert_results(tuple_compare)
+
+    def test_tuple_add(self):
+        self.assert_results(tuple_add)
+
+    def test_tuple_inplace_add(self):
+        self.assert_results(tuple_inplace_add)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/sot/test_04_list.py b/test/sot/test_04_list.py
@@ -222,6 +222,31 @@ def list_no_arguments():
     return l1[0] + l2[0]
 
 
+@check_no_breakgraph
+def list_compare():
+    # TODO(SigureMo): support gt, ge, lt, le
+    l1 = [1, 2, 3]
+    l2 = [1, 2, 3]
+    l3 = [1, 2, 4]
+    return l1 == l2, l1 == l3, l1 != l2, l1 != l3
+
+
+@check_no_breakgraph
+def list_add():
+    l0 = [1, 2, 3]
+    l1 = [4, 5, 6]
+    return l0 + l1
+
+
+@check_no_breakgraph
+def list_inplace_add():
+    l0 = [1, 2, 3]
+    l1 = l0
+    l2 = [4, 5, 6]
+    l0 += l2
+    return l0, l1
+
+
 class TestListBasic(TestCaseBase):
     def test_list_basic(self):
         self.assert_results(list_getitem_int, 1, paddle.to_tensor(2))
@@ -322,6 +347,15 @@ def test_list_max_min(self):
     def test_list_noargs(self):
         self.assert_results(list_no_arguments)
 
+    def test_list_compare(self):
+        self.assert_results(list_compare)
+
+    def test_list_add(self):
+        self.assert_results(list_add)
+
+    def test_list_inplace_add(self):
+        self.assert_results(list_inplace_add)
+
 
 if __name__ == "__main__":
     unittest.main()
