[SOT] Add new conditional fallback mechanism and stop specialize 1 (#72417)

SigureMo · web-flow · commit 221ebf82c8b7 · 2025-04-24T19:39:11.000+08:00
* [SOT] Add new conditional fallback mechanism and stop specialize `1`

* update test_sot_dynamic_shape
diff --git a/python/paddle/jit/sot/infer_meta.py b/python/paddle/jit/sot/infer_meta.py
@@ -592,4 +592,4 @@ def __init__(self, dynamic_axes: list[int], *args, **kwargs):
         )  # (idx of dim, min, max)
         super().__init__(*args, **kwargs)
         for i in dynamic_axes:
-            self.ranges.append((i, 2, None))
+            self.ranges.append((i, 1, None))
diff --git a/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py b/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
@@ -29,6 +29,7 @@
     ENV_SOT_ENABLE_STRICT_GUARD_CHECK,
     BreakGraphError,
     CompileCountInfo,
+    ConditionalFallbackError,
     FallbackError,
     InfoCollector,
     InnerError,
@@ -307,10 +308,11 @@ def start_translate(
     Returns:
         tuple[CustomCode, Guard | None, GuardChain | None]: The translated code object, its guard function and its guard tree node, or None if translation fails.
     """
+    simulator = None
     graph = FunctionGraph(frame.f_code, frame.f_globals, **kwargs)
-    vframe = VirtualFrame.from_real_frame(frame, graph)
-    simulator = OpcodeExecutor(vframe, graph)
     try:
+        vframe = VirtualFrame.from_real_frame(frame, graph)
+        simulator = OpcodeExecutor(vframe, graph)
         simulator.check_code_simulatable()
         InfoCollector().attach(CompileCountInfo, frame.f_code)
         with sot_simulation_mode_guard(True):
@@ -332,9 +334,9 @@ def start_translate(
             f"Found BreakGraphError raised, it should not be catch at start_translate!\n{e}"
         )
     except FallbackError as e:
-        if simulator.vframe.code in NO_FALLBACK_CODES:
+        if frame.f_code in NO_FALLBACK_CODES:
             raise InnerError(
-                f"{simulator.vframe.code.co_name} should not fallback, but got '{e}'"
+                f"{frame.f_code.co_name} should not fallback, but got '{e}'"
             )
         if is_strict_mode():
             raise
@@ -343,22 +345,30 @@ def start_translate(
             f"Unsupported Frame is {frame.f_code}, error message is: \n"
             + "".join(traceback.format_exception(type(e), e, e.__traceback__)),
         )
-
         dummy_guard_chain = [
             # TODO(zrr1999): GuardNode should support zero-expr constructor
             paddle.framework.core.GuardNode(
                 paddle.framework.core.DummyGuard(),
                 [paddle.framework.core.ConstantExprNode(True)],
             )
         ]
+        guard, guard_chain = dummy_guard, dummy_guard_chain
+
+        if isinstance(e, ConditionalFallbackError):
+            # Guard global variables only
+            graph.input_variables.clear()
+            guard = graph.guard_fn
+            guard_chain = graph.guard_chain
+
         return (
             CustomCode(None, e.disable_eval_frame),
-            dummy_guard,
-            dummy_guard_chain,
+            guard,
+            guard_chain,
         )
     except Exception as e:
         raise InnerError(OpcodeExecutorBase.error_message_summary(e)) from e
     finally:
-        simulator.cleanup()
+        if simulator is not None:
+            simulator.cleanup()
         del simulator
         gc.collect()
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -72,6 +72,7 @@
     ENV_SOT_ALLOW_DYNAMIC_SHAPE,
     BreakGraphError,
     BuiltinFunctionBreak,
+    ConditionalFallbackError,
     ConstTypes,
     DataDependencyDynamicShapeBreak,
     DataDependencyOperationBreak,
@@ -112,6 +113,7 @@
     GetItemTracker,
     GetIterTracker,
     GlobalTracker,
+    LocalTracker,
     SymbolicOperationTracker,
     Tracker,
 )
@@ -570,7 +572,7 @@ def make_stringified_guard(self) -> list[StringifiedExpression]:
                     )
                     if not isinstance(meta.shape[i], SymbolicInt)
                     else StringifiedExpression(
-                        f"{{}}.shape[{i}] >= 2",
+                        f"{{}}.shape[{i}] >= 1",
                         [frame_value_tracer],
                         union_free_vars(frame_value_tracer.free_vars),
                     )
@@ -951,7 +953,7 @@ def __init__(
                 (
                     GreaterEqualConstraintNode(
                         SymbolicConstraintNode(self.var_name),
-                        ConstantConstraintNode(2),
+                        ConstantConstraintNode(1),
                     ),
                     {self.var_name: self},
                 )
@@ -1258,7 +1260,7 @@ def should_create_symbolic_variable(
         value: Any,
         tracker: Tracker,
         symbolic_inputs: dict[str, dict[int, int] | None],
-    ):
+    ) -> bool | None:
         tracker_expr = tracker.trace_value_from_frame().inlined_expr
         symbolic_inputs.setdefault(tracker_expr, {})
         if tracker_expr in symbolic_inputs:
@@ -1267,7 +1269,9 @@ def should_create_symbolic_variable(
                 return False
             symbolic_input.setdefault(value, 0)
             symbolic_input[value] += 1
-            if value < 2:  # Specialize 0 and 1
+            if value == 0:
+                return None  # Fallback to dygraph
+            if value < 1:  # Specialize 0 and 1
                 return False
             if len(symbolic_input.keys()) > 1:
                 return True
@@ -1292,6 +1296,15 @@ def find_tensor_shape_source(dim_tracker: Tracker):
         shape_idx = dim_tracker.key
         return tensor_var, shape_idx
 
+    @staticmethod
+    def has_local_leaf(tracker: Tracker) -> bool:
+        if isinstance(tracker, LocalTracker):
+            return True
+        return any(
+            SymbolicVariable.has_local_leaf(input.tracker)
+            for input in tracker.inputs
+        )
+
     @VariableFactory.register_from_value(successor="ConstantVariable")
     def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
         if not ENV_SOT_ALLOW_DYNAMIC_SHAPE.get():
@@ -1322,16 +1335,30 @@ def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
             return None
         if not tracker.is_traceable():
             return None
+        if not SymbolicVariable.has_local_leaf(tracker):
+            return None
 
         from ..executor_cache import OpcodeExecutorCache
 
         symbolic_inputs = OpcodeExecutorCache().get_symbolic_inputs(
             graph.pycode_gen._origin_code
         )
 
-        if SymbolicVariable.should_create_symbolic_variable(
+        should_create_sym = SymbolicVariable.should_create_symbolic_variable(
             value, tracker, symbolic_inputs
+        )
+        if (
+            should_create_sym is None
+            and SymbolicVariable.find_tensor_shape_source(tracker) is not None
         ):
+            graph.add_global_guarded_variable(
+                ConstantVariable(value, graph, tracker)
+            )
+            raise ConditionalFallbackError(
+                "Fallback graph since input has 0 size Tensor",
+                disable_eval_frame=True,
+            )
+        elif should_create_sym:
             return SymbolicVariable(SymbolicInt(value), graph, tracker)
         return None
 
diff --git a/python/paddle/jit/sot/utils/__init__.py b/python/paddle/jit/sot/utils/__init__.py
@@ -41,6 +41,7 @@
     BreakGraphError,
     BreakGraphReasonBase,
     BuiltinFunctionBreak,
+    ConditionalFallbackError,
     DataDependencyControlFlowBreak,
     DataDependencyDynamicShapeBreak,
     DataDependencyOperationBreak,
diff --git a/python/paddle/jit/sot/utils/exceptions.py b/python/paddle/jit/sot/utils/exceptions.py
@@ -195,6 +195,9 @@ def __init__(self, msg, disable_eval_frame=False):
         self.disable_eval_frame = disable_eval_frame
 
 
+class ConditionalFallbackError(FallbackError): ...
+
+
 # raise in inline function call strategy.
 class BreakGraphError(SotErrorBase):
     def __init__(self, reason: BreakGraphReasonBase = None):
diff --git a/test/sot/test_sot_dynamic_shape.py b/test/sot/test_sot_dynamic_shape.py
@@ -25,6 +25,7 @@
 import paddle
 from paddle.jit.sot.psdb import check_no_breakgraph
 from paddle.jit.sot.utils import (
+    ConditionalFallbackError,
     allow_dynamic_shape_guard,
 )
 
@@ -138,11 +139,11 @@ def test_dynamic_int_input_cache_hit_case3(self):
         ), test_instruction_translator_cache_context() as ctx:
             translate_count_map = {
                 0: 1,
-                1: 2,  # 0, 1 is specialize to static dim
-                2: 3,  # 2 is dynamic dim
-                3: 3,  # 3 hit cache, no recompile
-                4: 4,  # 4 is dynamic dim, but it not hit cache
-                5: 4,  # 5 hit cache, no recompile
+                1: 2,  # 1 is dynamic dim
+                2: 2,  # 2 hit cache, no recompile
+                3: 2,  # 3 hit cache, no recompile
+                4: 3,  # 4 is dynamic dim, but it not hit cache
+                5: 3,  # 5 hit cache, no recompile
             }
             for i in range(0, 6):
                 self.assert_results(
@@ -251,7 +252,7 @@ def test_dynamic_shape_constraint(self):
         ), test_instruction_translator_cache_context() as ctx:
             const_dim = 6
             self.assert_results(
-                dynamic_shape_constraint, paddle.randn([0, 1, const_dim])
+                dynamic_shape_constraint, paddle.randn([1, 1, const_dim])
             )
             self.assertEqual(ctx.translate_count, 1)
 
@@ -326,6 +327,11 @@ def test_dynamic_shape_constraint(self):
                 ctx.translate_count, 5  # hit 2 * (s0 + s1 - 2) <= 30
             )
 
+            with self.assertRaises(ConditionalFallbackError):
+                self.assert_results(
+                    dynamic_shape_constraint, paddle.randn([0, 1, const_dim])
+                )
+
     def test_mixed_dynamic_and_static(self):
         with allow_dynamic_shape_guard(
             True
@@ -352,7 +358,7 @@ def test_mixed_static_after_dynamic(self):
             self.assert_results(dynamic_int_input_func1, a, 0)
             self.assertEqual(ctx.translate_count, 3)
             self.assert_results(dynamic_int_input_func1, a, 1)
-            self.assertEqual(ctx.translate_count, 4)
+            self.assertEqual(ctx.translate_count, 3)
 
     def test_dynamic_shape_with_constraints(self):
         with allow_dynamic_shape_guard(
