[SOT][DynamicShape] Register inplace dispatch for SymbolicVariable (#72159)

SigureMo · web-flow · commit 146fa42aab68 · 2025-04-10T10:41:31.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
@@ -44,6 +44,7 @@
     BINARY_OPS,
     UNARY_OPS,
     magic_method_builtin_dispatch,
+    non_inplace_op_to_inplace_op,
 )
 from .dispatch_functions import (
     create_raise_break_graph_handler,
@@ -1150,22 +1151,32 @@ def tensor_mod_dispatcher(
         ),
     )
 for binary_fn in SYMBOLIC_BINARY_OPS:
-    Dispatcher.register(
-        binary_fn,
-        ("SymbolicVariable", "SymbolicVariable | ConstantVariable"),
-        partial(
-            lambda fn, var, other: var.graph.call_symbolic_api(fn, var, other),
-            binary_fn,
-        ),
-    )
-    Dispatcher.register(
-        binary_fn,
-        ("ConstantVariable", "SymbolicVariable"),
-        partial(
-            lambda fn, var, other: var.graph.call_symbolic_api(fn, var, other),
-            binary_fn,
-        ),
-    )
+    register_fns = [binary_fn]
+    if (
+        inplace_binary_fn := non_inplace_op_to_inplace_op(binary_fn)
+    ) is not None:
+        register_fns.append(inplace_binary_fn)
+    for register_fn in register_fns:
+        Dispatcher.register(
+            register_fn,
+            ("SymbolicVariable", "SymbolicVariable | ConstantVariable"),
+            partial(
+                lambda fn, var, other: var.graph.call_symbolic_api(
+                    fn, var, other
+                ),
+                binary_fn,
+            ),
+        )
+        Dispatcher.register(
+            register_fn,
+            ("ConstantVariable", "SymbolicVariable"),
+            partial(
+                lambda fn, var, other: var.graph.call_symbolic_api(
+                    fn, var, other
+                ),
+                binary_fn,
+            ),
+        )
 
 
 @Dispatcher.register_decorator(bool)
diff --git a/python/paddle/jit/sot/utils/magic_methods.py b/python/paddle/jit/sot/utils/magic_methods.py
@@ -127,3 +127,15 @@ def magic_method_builtin_dispatch(fn: BinaryOp | UnaryOp) -> list[MagicMethod]:
         magic_name = UNARY_OPS_TO_MAGIC_NAMES[fn]
         return [MagicMethod(magic_name)]
     return []
+
+
+def non_inplace_op_to_inplace_op(
+    fn: BinaryOp,
+) -> BinaryOp | None:
+    for inplace_op, (
+        _,
+        non_inplace_op,
+    ) in INPLACE_BINARY_OPS_TO_MAGIC_NAMES.items():
+        if fn is non_inplace_op:
+            return inplace_op
+    return None
diff --git a/test/sot/test_sot_dynamic_shape.py b/test/sot/test_sot_dynamic_shape.py
@@ -271,5 +271,126 @@ def test_dynamic_shape_with_constraints(self):
                 self.assertEqual(ctx.translate_count, 2)
 
 
+@check_no_breakgraph
+def dynamic_shape_non_break_non_inplace_ops(x):
+    s0 = x.shape[0]
+    s1 = s0 + 1
+    s2 = 1 + s0
+    s3 = s1 + s2
+    s4 = s1 * s2
+    s5 = s1 - s2
+    s6 = s1 / s2
+    s7 = s1 // s2
+    s8 = s1 % s2
+    s9 = s1**s2
+    s10 = s1 & s2
+    s11 = s1 | s2
+    s12 = s1 ^ s2
+    s13 = s1 << s2
+    s14 = s1 >> s2
+    s15 = s1 == s2
+    s16 = s1 != s2
+    s17 = s1 < s2
+    s18 = s1 <= s2
+    s19 = s1 > s2
+    s20 = s1 >= s2
+    s21 = bool(s1)
+    s22 = not s2
+    return (
+        s0,
+        s1,
+        s2,
+        s3,
+        s4,
+        s5,
+        s6,
+        s7,
+        s8,
+        s9,
+        s10,
+        s11,
+        s12,
+        s13,
+        s14,
+        s15,
+        s16,
+        s17,
+        s18,
+        s19,
+        s20,
+        s21,
+        s22,
+    )
+
+
+@check_no_breakgraph
+def dynamic_shape_non_break_inplace_ops(x):
+    s0 = x.shape[0]
+    s1 = s0 + 1
+    s2 = s3 = s4 = s5 = s6 = s7 = s8 = s9 = s10 = s11 = s12 = s13 = s0
+    s2 += s1
+    s3 *= s1
+    s4 -= s1
+    # TODO(SigureMo): Open this case, currently the compute result between Python and C++ (Paddle Kernel)
+    # has a small difference (0.8333333134651184 and 0.8333333333333334)
+    # s5 /= s1
+    s6 //= s1
+    s7 %= s1
+    s8 **= s1
+    s9 &= s1
+    s10 |= s1
+    s11 ^= s1
+    s12 <<= s1
+    s13 >>= s1
+    return (
+        s0,
+        s1,
+        s2,
+        s3,
+        s4,
+        # s5,
+        s6,
+        s7,
+        s8,
+        s9,
+        s10,
+        s11,
+        s12,
+        s13,
+    )
+
+
+class TestDynamicShapeNonBreakOps(TestCaseBase):
+    def test_dynamic_shape_non_break_non_inplace_ops(self):
+        with allow_dynamic_shape_guard(
+            True
+        ), test_instruction_translator_cache_context() as ctx:
+            self.assert_results(
+                dynamic_shape_non_break_non_inplace_ops, paddle.randn([4, 5, 6])
+            )
+            self.assertEqual(ctx.translate_count, 1)
+            for i in range(5, 9):
+                self.assert_results(
+                    dynamic_shape_non_break_non_inplace_ops,
+                    paddle.randn([i, 5, 6]),
+                )
+                self.assertEqual(ctx.translate_count, 2)
+
+    def test_dynamic_shape_non_break_inplace_ops(self):
+        with allow_dynamic_shape_guard(
+            True
+        ), test_instruction_translator_cache_context() as ctx:
+            self.assert_results(
+                dynamic_shape_non_break_inplace_ops, paddle.randn([4, 5, 6])
+            )
+            self.assertEqual(ctx.translate_count, 1)
+            for i in range(5, 9):
+                self.assert_results(
+                    dynamic_shape_non_break_inplace_ops,
+                    paddle.randn([i, 5, 6]),
+                )
+                self.assertEqual(ctx.translate_count, 2)
+
+
 if __name__ == '__main__':
     unittest.main()
