[SOT][dynamic shape] support fallback to static shape (#67876)

zrr1999 · SigureMo · web-flow · commit 7d8c490ceaf4 · 2024-09-04T19:58:16.000+08:00
---------

Co-authored-by: Nyakku Shigure &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py b/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
@@ -84,6 +84,7 @@ def __call__(self, frame: types.FrameType, **kwargs) -> CustomCode:
         if code not in self.cache:
             log(2, f"[Cache]: Firstly call {code}\n")
             new_custom_code, guard_fn = self.translate(frame, **kwargs)
+            assert guard_fn is not None
             self.cache[code] = [(new_custom_code, guard_fn)]
             return new_custom_code
         guarded_fns = self.cache[code]
@@ -101,7 +102,7 @@ def lookup(
             guarded_fns (GuardedFunctions): The list of guarded functions associated with the code object.
 
         Returns:
-            CustomCode | None: The custom code object if a matching guard function is found, otherwise None.
+            CustomCode: The custom code object if a matching guard function is found, otherwise None.
         """
 
         if len(guarded_fns) >= self.MAX_CACHE_SIZE:
@@ -142,7 +143,8 @@ def lookup(
 
         log(2, "[Cache]: all guards missed\n")
         new_custom_code, guard_fn = self.translate(frame, **kwargs)
-        guarded_fns.append((new_custom_code, guard_fn))
+        if guard_fn is not None:
+            guarded_fns.append((new_custom_code, guard_fn))
         return new_custom_code
 
     def before_translate_hook(self, frame: types.FrameType):
@@ -151,7 +153,7 @@ def before_translate_hook(self, frame: types.FrameType):
 
     def translate(
         self, frame: types.FrameType, **kwargs
-    ) -> tuple[CustomCode, Guard]:
+    ) -> tuple[CustomCode, Guard | None]:
         """
         Translates the given frame's code object and returns the cache getter function and a guarded function for the translated code object.
 
@@ -202,22 +204,26 @@ def inner():
 def start_translate(
     frame: types.FrameType,
     **kwargs,
-) -> GuardedFunction:
+) -> tuple[CustomCode, Guard | None]:
     """
     Starts the translation process for the given frame and returns the translated code object and its guard function, or None if translation fails.
 
     Args:
         frame: The frame to be translated.
 
     Returns:
-        GuardedFunction | None: The translated code object and its guard function, or None if translation fails.
+        tuple[CustomCode, Guard | None]: The translated code object and its guard function, or None if translation fails.
     """
     simulator = OpcodeExecutor(frame, **kwargs)
     try:
         simulator.check_code_simulatable()
         new_custom_code, guard_fn = simulator.transform()
+        if not simulator._graph.need_cache:
+            return (
+                CustomCode(None, True),
+                None,
+            )
         return new_custom_code, guard_fn
-    # TODO(zrr1999): InnerError maybe place before (FallbackError, BreakGraphError)
     # TODO(0x45f): handle BreakGraphError to trigger fallback
     except BreakGraphError as e:
         raise RuntimeError(
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -30,6 +30,7 @@
 from paddle.jit.utils import OrderedSet
 from paddle.utils import flatten
 
+from .....utils.layers_utils import NotSupportedTensorArgumentError
 from ...infer_meta import (
     InferMetaCache,
     LayerInferMetaCache,
@@ -40,6 +41,8 @@
 from ...symbolic.statement_ir import Reference, StatementIR, Symbol
 from ...symbolic.symbolic_context import SymbolicTraceContext
 from ...utils import (
+    ENV_SOT_ALLOW_DYNAMIC_SHAPE,
+    BreakGraphError,
     NameGenerator,
     SotUndefinedVar,
     inner_error_default_handler,
@@ -162,6 +165,15 @@ def get_params_and_non_param_symbol(*args, **kwargs):
     return params, non_params
 
 
+def replace_symbolic_var_with_constant_var(inputs):
+    def func(x):
+        if isinstance(x, SymbolicVariable):
+            return x.to_constant()
+        return x
+
+    return map_variables(func, inputs, restore_variable=True)
+
+
 class VariableLoader:
     def __init__(self, store_var_info, pycode_gen):
         self._store_var_info = store_var_info
@@ -194,6 +206,7 @@ class FunctionGraph:
             "side_effects_state",
             "print_variables",
             "inplace_tensors",
+            "need_cache",
         ],
     )
 
@@ -203,6 +216,7 @@ def __init__(self, frame, **kwargs):
         self.input_variables = []  # Store variables required within a function
         self.pycode_gen = PyCodeGen(frame, disable_eval_frame=True)
         self.side_effects = SideEffects()
+        self.need_cache = True
         self._global_guarded_variables: OrderedSet[VariableBase] = OrderedSet()
         self._print_variables = []
         self._inplace_tensors = OrderedSet()
@@ -261,6 +275,7 @@ def save_memo(self) -> FunctionGraph.Memo:
             side_effects_state=self.side_effects.get_state(),
             print_variables=list(self._print_variables),
             inplace_tensors=OrderedSet(self._inplace_tensors),
+            need_cache=self.need_cache,
         )
 
     def restore_memo(self, memo: FunctionGraph.Memo):
@@ -278,6 +293,7 @@ def restore_memo(self, memo: FunctionGraph.Memo):
         self.side_effects.restore_state(memo.side_effects_state)
         self._print_variables = memo.print_variables
         self._inplace_tensors = memo.inplace_tensors
+        self.need_cache = memo.need_cache
 
     def collect_input_variables(self, inputs: list[VariableBase]):
         """
@@ -432,11 +448,9 @@ def compile_function(
         from ..breakpoint import BreakpointManager
 
         BreakpointManager().on_event("compile_function")
-        graph_fn, (
-            statement_ir,
-            symbolic_inputs,
-            symbolic_outputs,
-        ) = compile_graph_result
+        graph_fn, (statement_ir, symbolic_inputs, symbolic_outputs) = (
+            compile_graph_result
+        )
         compiled_fn_name = f"___graph_fn_{statement_ir.name}"
         # prepare function and inputs
         self.pycode_gen.gen_load_object(graph_fn, compiled_fn_name)
@@ -617,13 +631,49 @@ def symbolic_call(
             compute_fn   : function for add stmt to sir, (func, input_symbols, outputs_symbols, stacks) -> None
             func         : the logical function which will be represent as a stmt
         """
+
+        def try_infer_meta_fn(args, kwargs) -> Any:
+            try:
+                metas = convert_to_meta(args)
+                kwmetas = convert_to_meta(kwargs)
+                return args, kwargs, infer_meta_fn(func, *metas, **kwmetas)
+            except NotSupportedTensorArgumentError as e:
+                bound_arguments = inspect.signature(func).bind(*args, **kwargs)
+                bound_arguments.apply_defaults()
+                if e.name not in bound_arguments.arguments:
+                    # TODO(zrr1999): fallback static shape for all symbolic variables
+                    raise BreakGraphError(
+                        f"Can't find {e.name} in bound arguments."
+                    )
+                original_var = bound_arguments.arguments[e.name]
+                flatten_vars = original_var.flatten_items()
+
+                if not any(
+                    isinstance(arg, SymbolicVariable) for arg in flatten_vars
+                ):
+                    raise e
+
+                args, kwargs = map_if(
+                    (args, kwargs),
+                    pred=lambda x: x is original_var,
+                    true_fn=lambda x: replace_symbolic_var_with_constant_var(x),
+                    false_fn=lambda x: x,
+                )
+
+                metas = convert_to_meta(args)
+                kwmetas = convert_to_meta(kwargs)
+                return args, kwargs, infer_meta_fn(func, *metas, **kwmetas)
+
+        if ENV_SOT_ALLOW_DYNAMIC_SHAPE.get():
+            args, kwargs, out_metas = try_infer_meta_fn(args, kwargs)
+        else:
+            metas = convert_to_meta(args)
+            kwmetas = convert_to_meta(kwargs)
+            out_metas = infer_meta_fn(func, *metas, **kwmetas)
+
         self.collect_input_variables(list(args))
         self.collect_input_variables(list(kwargs.values()))
 
-        metas = convert_to_meta(args)
-        kwmetas = convert_to_meta(kwargs)
-
-        out_metas = infer_meta_fn(func, *metas, **kwmetas)
         inputs_symbols = (
             convert_to_symbol(args),
             convert_to_symbol(kwargs),
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/base.py b/python/paddle/jit/sot/opcode_translator/executor/variables/base.py
@@ -90,7 +90,12 @@ def find_traceable_vars(
     return results
 
 
-def map_variables(map_func, variables: list[VariableBase]):
+def map_variables(
+    map_func,
+    variables: list[VariableBase],
+    *,
+    restore_variable=False,
+) -> list[VariableBase]:
     """
     This function maps the given map_func to the given list of variables in a recursive manner.
     Args:
@@ -100,23 +105,42 @@ def map_variables(map_func, variables: list[VariableBase]):
     Returns:
         tuple: The result of applying the map_func to the variables.
     """
+    from .basic import SliceVariable
+    from .container import ContainerVariable
+
+    def _map_container_variable(variable: VariableBase | object):
+        if not isinstance(variable, ContainerVariable):
+            return variable
+        new_container = paddle.utils.map_structure(
+            _map_variable, variable.get_wrapped_items()
+        )
+        if not restore_variable:
+            return new_container
+        return VariableFactory.from_value(
+            new_container,
+            variable.graph,
+            DummyTracker(paddle.utils.flatten(new_container)),
+        )
 
-    def _map_variable(variable: VariableBase | object):
-        from .basic import SliceVariable
-        from .container import ContainerVariable
-
-        if isinstance(variable, ContainerVariable):
-            return paddle.utils.map_structure(
-                _map_variable, variable.get_wrapped_items()
-            )
-
-        if isinstance(variable, SliceVariable):
-            return slice(
-                map_func(variable.getattr("start")),
-                map_func(variable.getattr("stop")),
-                map_func(variable.getattr("step")),
-            )
+    def _map_slice_variable(variable: VariableBase | object):
+        if not isinstance(variable, SliceVariable):
+            return variable
+        new_slice = slice(
+            map_func(variable.getattr("start")),
+            map_func(variable.getattr("stop")),
+            map_func(variable.getattr("step")),
+        )
+        if not restore_variable:
+            return new_slice
+        return VariableFactory.from_value(
+            new_slice,
+            variable.graph,
+            DummyTracker([new_slice.start, new_slice.stop, new_slice.step]),
+        )
 
+    def _map_variable(variable: VariableBase | object):
+        variable = _map_container_variable(variable)
+        variable = _map_slice_variable(variable)
         return map_func(variable)
 
     return paddle.utils.map_structure(_map_variable, variables)
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -689,6 +689,24 @@ def __init__(
         self.graph.side_effects.record_mutable_variable(self)
 
     def to_constant(self):
+        from ..executor_cache import (
+            OpcodeExecutorCache,
+        )
+
+        symbolic_inputs = OpcodeExecutorCache().get_symbolic_inputs(
+            self.graph.pycode_gen._origin_code
+        )
+
+        def disable_symbolic(var: VariableBase):
+            if var.tracker.is_traceable():
+                tracker_expr = var.tracker.trace_value_from_frame().inlined_expr
+                symbolic_inputs[tracker_expr] = None
+                return
+            for input_var in var.tracker.inputs:
+                disable_symbolic(input_var)
+
+        disable_symbolic(self)
+        self.graph.need_cache = False
         return ConstantVariable(
             self.get_py_value(), self.graph, DummyTracker([self])
         )
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -668,19 +668,6 @@ def call_function(self, /, *args, **kwargs):
                     self.graph.pycode_gen._origin_code
                 )
 
-                for var in itertools.chain(args, kwargs.values()):
-                    if isinstance(var, SymbolicVariable):
-                        if var.tracker.is_traceable():
-                            tracker_expr = (
-                                var.tracker.trace_value_from_frame().inlined_expr
-                            )
-                            symbolic_inputs[tracker_expr] = None
-                        else:
-                            for traceable_var in var.get_traceable_inputs():
-                                tracker_expr = (
-                                    traceable_var.tracker.trace_value_from_frame().inlined_expr
-                                )
-                                symbolic_inputs[tracker_expr] = None
                 args, kwargs = map_if(
                     (args, kwargs),
                     pred=lambda x: isinstance(x, SymbolicVariable),
diff --git a/python/paddle/jit/sot/utils/__init__.py b/python/paddle/jit/sot/utils/__init__.py
@@ -32,7 +32,6 @@
 )
 from .exceptions import (  # noqa: F401
     BreakGraphError,
-    DynamicShapeFallbackError,
     ExportError,
     FallbackError,
     InnerError,
diff --git a/python/paddle/jit/sot/utils/exceptions.py b/python/paddle/jit/sot/utils/exceptions.py
@@ -42,10 +42,6 @@ def __init__(self, msg, disable_eval_frame=False):
         self.disable_eval_frame = disable_eval_frame
 
 
-class DynamicShapeFallbackError(SotErrorBase):
-    pass
-
-
 # raise in inline function call strategy.
 class BreakGraphError(SotErrorBase):
     pass
diff --git a/python/paddle/utils/layers_utils.py b/python/paddle/utils/layers_utils.py
@@ -43,6 +43,12 @@
 _U = TypeVar("_U")
 
 
+class NotSupportedTensorArgumentError(TypeError):
+    def __init__(self, msg, name: str):
+        super().__init__(msg)
+        self.name = name
+
+
 def convert_to_list(value, n, name, dtype=int):
     """
     Converts a single numerical type or iterable of numerical
@@ -64,9 +70,7 @@ def convert_to_list(value, n, name, dtype=int):
         passed.
     """
     if isinstance(value, dtype):
-        return [
-            value,
-        ] * n
+        return [value] * n
     else:
         try:
             value_list = list(value)
@@ -79,9 +83,11 @@ def convert_to_list(value, n, name, dtype=int):
                 f"The {name}'s length must be {n}. Received: {value}"
             )
         for single_value in value_list:
-            assert not isinstance(
-                single_value, (Variable, paddle.pir.Value)
-            ), f"Required numerical type with '{dtype}', but received Tensor."
+            if isinstance(single_value, (Variable, paddle.pir.Value)):
+                raise NotSupportedTensorArgumentError(
+                    f"`{name}` required numerical type with `{dtype}`, but received Tensor.",
+                    name,
+                )
             try:
                 dtype(single_value)
             except (ValueError, TypeError):
diff --git a/test/sot/test_sot_dynamic_shape.py b/test/sot/test_sot_dynamic_shape.py
diff --git a/test/sot/test_trace_list_arg.py b/test/sot/test_trace_list_arg.py

Original file line number	Diff line number	Diff line change
`@@ -32,7 +32,6 @@`
`32`	`32`	`)`
`33`	`33`	`from .exceptions import ( # noqa: F401`
`34`	`34`	`BreakGraphError,`
`35`		`- DynamicShapeFallbackError,`
`36`	`35`	`ExportError,`
`37`	`36`	`FallbackError,`
`38`	`37`	`InnerError,`