[SOT] Allow specify some sir fallback to dygraph (#70740)

SigureMo · web-flow · commit 044b4269d446 · 2025-01-10T20:44:06.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -24,14 +24,15 @@
 import types
 from dataclasses import dataclass
 from itertools import chain
-from typing import Any, Callable
+from typing import TYPE_CHECKING, Any, Callable
 
 from paddle.jit.utils import OrderedSet
 
 from ...profiler import EventGuard, event_register
 from ...psdb import NO_BREAKGRAPH_CODES
 from ...utils import (
     ENV_MIN_GRAPH_SIZE,
+    ENV_SOT_FORCE_FALLBACK_SIR_IDS,
     BreakGraphError,
     FallbackError,
     InnerError,
@@ -108,6 +109,9 @@
     VariableFactory,
 )
 
+if TYPE_CHECKING:
+    from .function_graph import CompileGraphResult
+
 SUPPORT_COMPARE_OP = {
     ">": operator.gt,
     "<": operator.lt,
@@ -317,6 +321,34 @@ def inner(*args, **kwargs):
     return inner
 
 
+def parse_force_fallback_sir_ids() -> set[int]:
+    ids_string = ENV_SOT_FORCE_FALLBACK_SIR_IDS.get()
+    if not ids_string:
+        return set()
+    ids = set()
+    for comma_sep_part in ids_string.split(","):
+        comma_sep_part = comma_sep_part.strip()
+        range_parts = comma_sep_part.split("-")
+        if len(range_parts) == 1:
+            ids.add(int(comma_sep_part))
+        else:
+            ids |= set(range(int(range_parts[0]), int(range_parts[1]) + 1))
+    return ids
+
+
+def need_fallback(compile_graph_result: CompileGraphResult) -> bool:
+    graph_fn, (statement_ir, _, _) = compile_graph_result
+
+    assert statement_ir.name[:4] == "SIR_"
+    sir_id = int(statement_ir.name[4:])
+    if (
+        graph_fn.graph_size() < ENV_MIN_GRAPH_SIZE.get()
+        or sir_id in parse_force_fallback_sir_ids()
+    ):
+        return True
+    return False
+
+
 class OpcodeExecutorBase:
     """
     Base class for executing opcode instructions.
@@ -2013,8 +2045,7 @@ def RETURN_CONST(self, instr: Instruction):
 
     def compile_return(self, ret_val):
         compile_graph_result = self._graph.compile_graph(ret_val)
-        graph_fn, _ = compile_graph_result
-        if graph_fn.graph_size() < ENV_MIN_GRAPH_SIZE.get():
+        if need_fallback(compile_graph_result):
             self.new_code = None
         else:
             self._graph.compile_function(compile_graph_result, [ret_val])
@@ -2054,9 +2085,8 @@ def get_compute_fn_and_update_changed_vars(
             store_var_info[_var.id].append(name)
 
         compile_graph_result = self._graph.compile_graph(*store_vars)
-        graph_fn, _ = compile_graph_result
 
-        if graph_fn.graph_size() < ENV_MIN_GRAPH_SIZE.get():
+        if need_fallback(compile_graph_result):
             return self._graph._restore_origin_opcode(
                 list(stack), store_var_info, end_idx
             )
diff --git a/python/paddle/jit/sot/utils/__init__.py b/python/paddle/jit/sot/utils/__init__.py
@@ -21,6 +21,7 @@
     ENV_SOT_ENABLE_FASTER_GUARD,
     ENV_SOT_ENABLE_GUARD_TREE,
     ENV_SOT_EXPORT,
+    ENV_SOT_FORCE_FALLBACK_SIR_IDS,
     ENV_SOT_LOG_LEVEL,
     ENV_SOT_WITH_CONTROL_FLOW,
     ENV_STRICT_MODE,
diff --git a/python/paddle/jit/sot/utils/envs.py b/python/paddle/jit/sot/utils/envs.py
@@ -55,6 +55,9 @@
     "SOT_BREAK_GRAPH_ON_GET_SYMBOLIC_VALUE", False
 )
 ENV_SOT_COLLECT_INFO = StringListEnvironmentVariable("SOT_COLLECT_INFO", [])
+ENV_SOT_FORCE_FALLBACK_SIR_IDS = StringEnvironmentVariable(
+    "SOT_FORCE_FALLBACK_SIR_IDS", ""
+)
 
 
 @contextmanager

Original file line number	Diff line number	Diff line change
`@@ -55,6 +55,9 @@`
`55`	`55`	`"SOT_BREAK_GRAPH_ON_GET_SYMBOLIC_VALUE", False`
`56`	`56`	`)`
`57`	`57`	`ENV_SOT_COLLECT_INFO = StringListEnvironmentVariable("SOT_COLLECT_INFO", [])`
	`58`	`+ENV_SOT_FORCE_FALLBACK_SIR_IDS = StringEnvironmentVariable(`
	`59`	`+ "SOT_FORCE_FALLBACK_SIR_IDS", ""`
	`60`	`+)`
`58`	`61`
`59`	`62`
`60`	`63`	`@contextmanager`