[SOT][dynamic shape] Use guard value instead of raise_break_graph_fn to call int/float (#66346)

zrr1999 · web-flow · commit 4d898d919e57 · 2024-08-28T19:34:57.000+08:00
* use guard value instead of raise_break_graph_fn

* add test

* fix

* add

* improve

* improve

* fix dispatch

* fix
diff --git a/python/paddle/jit/sot/infer_meta.py b/python/paddle/jit/sot/infer_meta.py
@@ -41,21 +41,23 @@ def get_static_type(self) -> type:
 
 
 class SymbolicBool(SymbolicValue):
-    def get_static_type(self) -> type:
+    def get_static_type(self) -> type[bool]:
         return bool
 
 
 class SymbolicInt(SymbolicValue):
-    def get_static_type(self) -> type:
+    def get_static_type(self) -> type[int]:
         return int
 
 
 class SymbolicFloat(SymbolicValue):
-    def get_static_type(self) -> type:
+    def get_static_type(self) -> type[float]:
         return float
 
 
 class MetaInfo:
+    shape: list[int | SymbolicInt]
+
     def __init__(
         self,
         shape,
@@ -73,15 +75,15 @@ def __init__(
         self.persistable = persistable
         self.type = type
         self.place = place
-        self.shape: list[int | SymbolicInt] = shape
+        self.shape = shape
         self.dtype = dtype
         self.stop_gradient = stop_gradient
 
     def shape_with_special_symbol(
         self, dynamic_symbol: DynamicSymbolT = -1
     ) -> list[int | DynamicSymbolT]:
         return [
-            dim if not isinstance(dim, SymbolicInt) else dynamic_symbol
+            dynamic_symbol if isinstance(dim, SymbolicInt) else dim
             for dim in self.shape
         ]
 
diff --git a/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py b/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
@@ -58,7 +58,7 @@ class OpcodeExecutorCache(metaclass=Singleton):
     MAX_CACHE_SIZE = 20
     cache: dict[types.CodeType, GuardedFunctions]
     translate_count: int
-    code_symbolic_inputs: dict[types.CodeType, dict[str, dict[int, int]]]
+    code_symbolic_inputs: dict[types.CodeType, dict[str, None | dict[int, int]]]
 
     def __init__(self):
         self.cache = {}
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py b/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py
@@ -947,11 +947,10 @@ def is_not_func(var: VariableBase, other: VariableBase):
 Dispatcher.register(tensor_numel, ("TensorVariable",), lambda x: x.numel())
 
 for unary_fn in UNARY_OPS:
-    # TODO(zrr1999): SymbolicVariable should have special dispatch for fallback_tensor_unary_method
     if unary_fn in fallback_tensor_unary_method:
         Dispatcher.register(
             unary_fn,
-            ("TensorVariable | SymbolicVariable",),
+            ("TensorVariable",),
             raise_break_graph_fn,
         )
         continue
@@ -1036,6 +1035,20 @@ def tensor_mod_dispatcher(
                     ),
                 )
 # Symbolic
+for unary_fn in fallback_tensor_unary_method:
+    Dispatcher.register(
+        unary_fn,
+        ("SymbolicVariable",),
+        partial(
+            lambda fn, var: VariableFactory.from_value(
+                fn(var.get_py_value()),
+                var.graph,
+                tracker=DummyTracker([var]),
+            ),
+            unary_fn,
+        ),
+    )
+
 for binary_fn in BINARY_OPS:
     for magic_method in magic_method_builtin_dispatch(binary_fn):
         if magic_method.name not in get_tensor_methods():
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -686,7 +686,7 @@ def __init__(
             )
         self.need_guard_value = False
 
-    def get_py_value(self, allow_tensor=False):
+    def get_py_value(self, allow_tensor: bool = False) -> bool | int | float:
         self.need_guard_value = True
         if isinstance(self.value, SymbolicValue):
             assert isinstance(
@@ -698,6 +698,7 @@ def get_py_value(self, allow_tensor=False):
             self.value = getattr(
                 inputs[0].get_py_value(), self.tracker.method_name
             )(*other_inputs_value)
+            assert isinstance(self.value, (bool, int, float))
         return self.value
 
     def get_py_type(self):
@@ -714,6 +715,18 @@ def __bool__(self) -> bool:
     def bool(self):
         return ConstantVariable(bool(self), self.graph, DummyTracker([self]))
 
+    def __int__(self) -> int:
+        return int(self.get_py_value())
+
+    def int(self):
+        return ConstantVariable(int(self), self.graph, DummyTracker([self]))
+
+    def __float__(self) -> float:
+        return float(self.get_py_value())
+
+    def float(self):
+        return ConstantVariable(float(self), self.graph, DummyTracker([self]))
+
     @property
     def out_var_name(self):
         return f"{self.graph.OUT_VAR_PREFIX}{self.var_name}"
diff --git a/test/sot/test_sot_dynamic_shape.py b/test/sot/test_sot_dynamic_shape.py
@@ -22,7 +22,10 @@
 )
 
 import paddle
-from paddle.jit.sot.utils import with_allow_dynamic_shape_guard
+from paddle.jit.sot.psdb import check_no_breakgraph
+from paddle.jit.sot.utils import (
+    with_allow_dynamic_shape_guard,
+)
 
 
 def dynamic_shape_input_func1(x):
@@ -123,23 +126,34 @@ def test_dynamic_shape_input_cache_hit_case2(self):
                 )
                 self.assertEqual(ctx.translate_count, 2)
 
-    # def test_dynamic_shape_in_list(self):
-    #     with with_allow_dynamic_shape_guard(
-    #         True
-    #     ), test_instruction_translator_cache_context() as ctx:
-    #         self.assert_results(
-    #             dynamic_shape_in_list,
-    #             paddle.randn([1, 4, 5]),
-    #             [4, 5],
-    #         )
-    #         self.assertEqual(ctx.translate_count, 1)
-    #         for i in range(2, 6):
-    #             self.assert_results(
-    #                 dynamic_shape_in_list,
-    #                 paddle.randn([i, 4, 5]),
-    #                 [i * 4, 5],
-    #             )
-    #             self.assertEqual(ctx.translate_count, 2)
+    def test_dynamic_shape_cast(self):
+        with with_allow_dynamic_shape_guard(
+            True
+        ), test_instruction_translator_cache_context() as ctx:
+            func1 = check_no_breakgraph(lambda n: bool(n))
+            func2 = check_no_breakgraph(lambda n: int(n))
+            func3 = check_no_breakgraph(lambda n: float(n))
+            for func in [func1, func2, func3]:
+                self.assert_results(func, 1)
+                self.assert_results(func, 2)
+
+    def test_dynamic_shape_in_list(self):
+        with with_allow_dynamic_shape_guard(
+            True
+        ), test_instruction_translator_cache_context() as ctx:
+            self.assert_results(
+                dynamic_shape_in_list,
+                paddle.randn([1, 4, 5]),
+                [4, 5],
+            )
+            self.assertEqual(ctx.translate_count, 1)
+            for i in range(2, 6):
+                self.assert_results(
+                    dynamic_shape_in_list,
+                    paddle.randn([i, 4, 5]),
+                    [i * 4, 5],
+                )
+                self.assertEqual(ctx.translate_count, 2)
 
 
 if __name__ == '__main__':
