[SOT] Add dispatch for check eq between builtin types (list, tuple, dict, const) (#70151)

SigureMo · web-flow · commit cf9918ff5d26 · 2024-12-12T04:41:22.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
@@ -504,6 +504,35 @@ def dispatch_list_ne(lhs: ListVariable, rhs: ListVariable):
     return Dispatcher.call(operator.eq, lhs, rhs).bool_not()
 
 
+BUILTIN_EQ_DISPATCH_TYPES = [
+    "ListVariable",
+    "TupleVariable",
+    "DictVariable",
+    "ConstantVariable",
+]
+
+for i in range(len(BUILTIN_EQ_DISPATCH_TYPES)):
+    current_type = BUILTIN_EQ_DISPATCH_TYPES[i]
+    other_types = (
+        BUILTIN_EQ_DISPATCH_TYPES[:i] + BUILTIN_EQ_DISPATCH_TYPES[i + 1 :]
+    )
+    Dispatcher.register(
+        operator.eq,
+        (current_type, " | ".join(other_types)),
+        lambda var, other: ConstantVariable(
+            False, var.graph, DummyTracker([var, other])
+        ),
+    )
+
+    Dispatcher.register(
+        operator.ne,
+        (current_type, " | ".join(other_types)),
+        lambda var, other: ConstantVariable(
+            True, var.graph, DummyTracker([var, other])
+        ),
+    )
+
+
 # getattr
 Dispatcher.register(
     getattr,
diff --git a/test/sot/test_builtin_dispatch.py b/test/sot/test_builtin_dispatch.py
@@ -135,6 +135,25 @@ def test_log(x: int):
     return math.log(x)
 
 
+@check_no_breakgraph
+def test_builtin_type_check_eq():
+    a = 1
+    b = []
+    c = ()
+    d = {}
+    eq_results = (
+        a == b, a == c, a == d,
+        b == a, b == c, b == d,
+        c == a, c == b, c == d,
+    )  # fmt: skip
+    ne_results = (
+        a != b, a != c, a != d,
+        b != a, b != c, b != d,
+        c != a, c != b, c != d,
+    )  # fmt: skip
+    return eq_results, ne_results
+
+
 class TestBuiltinDispatch(TestCaseBase):
     @test_with_faster_guard
     def test_dispatch_len(self):
@@ -279,6 +298,9 @@ def test_dispatch_min(self):
     def test_dispatch_max(self):
         self.assert_results(test_max)
 
+    def test_dispatch_builtin_type_check_eq(self):
+        self.assert_results(test_builtin_type_check_eq)
+
 
 def run_getattr(x: paddle.Tensor):
     attr = 'dtype'
