[SOT][dynamic shape] add SymbolicTracker to replace DummyTracker (#64500)

zrr1999 · web-flow · commit ded7e0007d55 · 2024-05-29T17:13:27.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -64,7 +64,7 @@
     SideEffectRestorer,
     SideEffects,
 )
-from .tracker import BuiltinTracker, DummyTracker
+from .tracker import BuiltinTracker, DummyTracker, SymbolicOperationTracker
 from .variables import (
     ConstantVariable,
     DictVariable,
@@ -580,13 +580,20 @@ def message_handler(*args, **kwargs):
         try:
             return inner_error_default_handler(
                 self.symbolic_call, message_handler
-            )(ast_infer_meta, compute_fn, static_function, *args, **kwargs)
+            )(
+                ast_infer_meta,
+                compute_fn,
+                static_function,
+                False,
+                *args,
+                **kwargs,
+            )
         except Exception as e:
             log(3, f"[call AST] {e}")
             return None
 
     def symbolic_call(
-        self, infer_meta_fn, compute_fn, func, is_symbolic_int, *args, **kwargs
+        self, infer_meta_fn, compute_fn, func, is_symbolic_var, *args, **kwargs
     ):
         """
         Using infer_meta_fn and compute_fn convert func to symbolic function.
@@ -612,14 +619,21 @@ def symbolic_call(
 
         log(3, f"         inputs : {inputs_symbols}", "\n")
 
-        var_cls = SymbolicVariable if is_symbolic_int else TensorVariable
+        if is_symbolic_var:
+            var_cls = SymbolicVariable
+            tracker = SymbolicOperationTracker(
+                list(args) + list(kwargs.values()), func
+            )
+        else:
+            var_cls = TensorVariable
+            tracker = DummyTracker(list(args) + list(kwargs.values()))
         outputs = map_if(
             out_metas,
             pred=lambda x: isinstance(x, MetaInfo),
             true_fn=lambda x: var_cls(
                 x,
                 self,
-                tracker=DummyTracker(list(args) + list(kwargs.values())),
+                tracker=tracker,
             ),
             false_fn=lambda x: x,
         )
@@ -649,9 +663,15 @@ def symbolic_call(
                     stmt_stacks,
                 )  # symbolic only contain symbols.
                 self._put_inner(outputs)
-            return VariableFactory.from_value(
-                outputs, self, DummyTracker(list(args) + list(kwargs.values()))
-            )
+            if is_symbolic_var:
+                # compute_fn should be call_method
+                tracker = SymbolicOperationTracker(
+                    list(args) + list(kwargs.values()), func
+                )
+            else:
+                tracker = DummyTracker(list(args) + list(kwargs.values()))
+
+            return VariableFactory.from_value(outputs, self, tracker)
         else:
             return ConstantVariable.wrap_literal(None, self)
 
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -88,6 +88,7 @@
     NullVariable,
     SequenceIterVariable,
     SliceVariable,
+    SymbolicVariable,
     TensorVariable,
     TupleVariable,
     UserDefinedFunctionVariable,
@@ -204,7 +205,7 @@ def inner(self: OpcodeExecutorBase, instr: Instruction):
                     fn, graph=self._graph, tracker=DanglingTracker()
                 )(res)
 
-            assert isinstance(res, ConstantVariable)
+            assert isinstance(res, (ConstantVariable, SymbolicVariable))
             is_jump = res.get_py_value()
             assert isinstance(is_jump, bool)
             if is_jump:
diff --git a/python/paddle/jit/sot/opcode_translator/executor/tracker.py b/python/paddle/jit/sot/opcode_translator/executor/tracker.py
@@ -127,6 +127,31 @@ def need_guard(self) -> bool:
         return False
 
 
+class SymbolicOperationTracker(DummyTracker):
+    """
+    SymbolicOperationTracker is a subclass of Tracker that specifically tracks variables cannot be reproduced from the frame.
+    It is mostly generated by complex operations of symbolic variables.
+
+    Args:
+        inputs (list[VariableBase]): The input variables associated with the generated variables.
+    """
+
+    def __init__(self, inputs: Sequence[VariableBase], method_name: str):
+        super().__init__(inputs)
+        self.method_name = method_name
+
+    def gen_instructions(self, codegen: PyCodeGen):
+        raise InnerError("SymbolicOperationTracker has no instructions")
+
+    def trace_value_from_frame(self):
+        raise InnerError(
+            "SymbolicOperationTracker can't trace value from frame"
+        )
+
+    def __repr__(self) -> str:
+        return f"SymbolicOperationTracker(num_inputs={len(self.inputs)})"
+
+
 class DanglingTracker(Tracker):
     """
     DanglingTracker is a subclass of Tracker that specifically tracks variables that are not in the frame.
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py b/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py
@@ -576,12 +576,12 @@ def dispatch_reversed(var: ContainerVariable):
 # bool
 Dispatcher.register(
     bool,
-    ("ContainerVariable",),
+    ("ContainerVariable | SymbolicVariable",),
     lambda var: var.bool(),
 )
 Dispatcher.register(
     operator.truth,
-    ("ConstantVariable",),
+    ("ConstantVariable | SymbolicVariable",),
     lambda var: var.bool(),
 )
 
@@ -941,19 +941,6 @@ def is_not_func(var: VariableBase, other: VariableBase):
                     magic_method.name,
                 ),
             )
-            Dispatcher.register(
-                binary_fn,
-                (
-                    "SymbolicVariable",
-                    "ConstantVariable | SymbolicVariable",
-                ),
-                partial(
-                    lambda magic_name, var, other: var.graph.call_symbolic_method(
-                        magic_name, var, other
-                    ),
-                    magic_method.name,
-                ),
-            )
         else:
             # skip __mod__ for str and TensorVariable
             if magic_method.name == "__rmod__":
@@ -983,17 +970,39 @@ def tensor_mod_dispatcher(
                         magic_method.name,
                     ),
                 )
+# Symbolic
+for binary_fn in BINARY_OPS:
+    for magic_method in magic_method_builtin_dispatch(binary_fn):
+        # skip all inplace magic method name, we will dispatch it to non-inplace
+        # magic methods
+        if magic_method.is_inplace:
+            continue
 
-                Dispatcher.register(
-                    binary_fn,
-                    ("ConstantVariable", "SymbolicVariable"),
-                    partial(
-                        lambda magic_name, var, other: var.graph.call_symbolic_method(
-                            magic_name, var, other
-                        ),
-                        magic_method.name,
+        if not magic_method.is_reverse:
+            Dispatcher.register(
+                binary_fn,
+                (
+                    "SymbolicVariable",
+                    "ConstantVariable | SymbolicVariable",
+                ),
+                partial(
+                    lambda magic_name, var, other: var.graph.call_symbolic_method(
+                        magic_name, var, other
                     ),
-                )
+                    magic_method.name,
+                ),
+            )
+        else:
+            Dispatcher.register(
+                binary_fn,
+                ("ConstantVariable", "SymbolicVariable"),
+                partial(
+                    lambda magic_name, var, other: var.graph.call_symbolic_method(
+                        magic_name, var, other
+                    ),
+                    magic_method.name,
+                ),
+            )
 
 # Register dispatch for NumpyVariable: fallback !
 for unary_fn in UNARY_OPS:
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -53,6 +53,7 @@
     GetAttrTracker,
     GetIterTracker,
     GlobalTracker,
+    SymbolicOperationTracker,
     Tracker,
 )
 from .base import VariableBase, VariableFactory
@@ -618,16 +619,35 @@ def __init__(
             self.meta = MetaInfo(
                 [], paddle.int64, True, self.var_name, False, None, None
             )
+        self.need_guard_value = False
 
     def get_py_value(self, allow_tensor=False):
+        self.need_guard_value = True
+        if self.value is None:
+            assert isinstance(
+                self.tracker, SymbolicOperationTracker
+            ), f"self.value is None, but tracker is not SymbolicOperationTracker. tracker: {self.tracker}"
+            inputs = self.tracker.inputs
+            assert len(inputs) >= 1
+            other_inputs_value = [x.get_py_value() for x in inputs[1:]]
+            self.value = getattr(
+                inputs[0].get_py_value(), self.tracker.method_name
+            )(*other_inputs_value)
         return self.value
 
     def get_py_type(self):
-        return int
+        # TODO(zrr1999): not need to use value to get type
+        return super().get_py_type()
 
     def get_symbol(self) -> Symbol:
         return Symbol(self.var_name)
 
+    def __bool__(self) -> bool:
+        return bool(self.get_py_value())
+
+    def bool(self):
+        return ConstantVariable(bool(self), self.graph, DummyTracker([self]))
+
     @property
     def out_var_name(self):
         return f"{self.graph.OUT_VAR_PREFIX}{self.var_name}"
@@ -651,10 +671,11 @@ def make_stringify_guard(self) -> list[StringifyExpression]:
         symbolic_input = symbolic_inputs[frame_value_tracer.inlined_expr]
         symbolic_input.setdefault(self.value, 0)
         symbolic_input[self.value] += 1
-
+        if self.need_guard_value:
+            return super().make_stringify_guard()
         return [
             StringifyExpression(
-                "isinstance({}, int)",
+                f"id(type({{}})) == {id(self.get_py_type())}",
                 [frame_value_tracer],
                 union_free_vars(frame_value_tracer.free_vars),
             )
diff --git a/test/sot/test_sot_dynamic_shape.py b/test/sot/test_sot_dynamic_shape.py
@@ -39,6 +39,13 @@ def dynamic_int_input_func2(x, n):
     return x + n[1]
 
 
+def dynamic_int_input_func3(x, n):
+    if n < 4:
+        return 1
+    x = paddle.reshape(x, [n, -1])
+    return (x + n) * 2 - 1, (-n + 1) * 2 - 1
+
+
 class TestOpcodeExecutorDynamicShapeCache(TestCaseBase):
     def test_dynamic_int_input_cache_hit_case1(self):
         with with_allow_dynamic_shape_guard(
@@ -68,6 +75,16 @@ def test_dynamic_int_input_cache_hit_case2(self):
                 )
                 self.assertEqual(ctx.translate_count, 2)
 
+    def test_dynamic_int_input_cache_hit_case3(self):
+        with with_allow_dynamic_shape_guard(
+            True
+        ), test_instruction_translator_cache_context() as ctx:
+            for i in range(0, 6):
+                self.assert_results(
+                    dynamic_int_input_func3, paddle.randn([3, 4, 5]), i
+                )
+                self.assertEqual(ctx.translate_count, i + 1)
+
 
 if __name__ == '__main__':
     unittest.main()
