[SOT][CINN] Add flag to disable SOT fallback when compile time exceed compile time limit (PaddlePaddle#76386)

Author: SigureMo

python/paddle/jit/sot/opcode_translator/executor/executor_cache.py

@@ -25,6 +25,7 @@
 from ...psdb import NO_FALLBACK_CODES
 from ...utils import (
     ENV_SOT_ALLOW_DYNAMIC_SHAPE,
+    ENV_SOT_ENABLE_COMPILE_TIME_LIMIT,
     ENV_SOT_ENABLE_GUARD_TREE,
     ENV_SOT_ENABLE_STRICT_GUARD_CHECK,
     ENV_SOT_UNSAFE_CACHE_FASTPATH,
@@ -183,6 +184,7 @@ def lookup(
         enable_strict_guard = ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get()
         enable_guard_tree = ENV_SOT_ENABLE_GUARD_TREE.get()
         enable_unsafe_cache_fastpath = ENV_SOT_UNSAFE_CACHE_FASTPATH.get()
+        enable_compile_time_limit = ENV_SOT_ENABLE_COMPILE_TIME_LIMIT.get()
 
         if enable_unsafe_cache_fastpath and (
             self.is_fastpath_threshold_reached(code)
@@ -205,13 +207,19 @@ def lookup(
                 return guarded_fns[cache_index][0]
             else:
                 log(2, "[Cache] all guards missed (guard tree mode)\n")
-                if compile_time_for_code >= self.MAX_COMPILE_TIME_PER_CODE:
+                if (
+                    enable_compile_time_limit
+                    and compile_time_for_code >= self.MAX_COMPILE_TIME_PER_CODE
+                ):
                     log(
                         2,
                         "[Cache] Exceed max compile time per code, skip it\n",
                     )
                     return CustomCode(None, False)
-                if compile_time_total >= self.MAX_COMPILE_TIME_TOTAL:
+                if (
+                    enable_compile_time_limit
+                    and compile_time_total >= self.MAX_COMPILE_TIME_TOTAL
+                ):
                     log_once(
                         f"[SOT] Current total compile time is {compile_time_total}, exceed max compile time total {self.MAX_COMPILE_TIME_TOTAL}, fallback new function to dygraph"
                     )
@@ -307,10 +315,16 @@ def lookup(
                     )
 
         log(2, "[Cache] all guards missed\n")
-        if compile_time_for_code >= self.MAX_COMPILE_TIME_PER_CODE:
+        if (
+            enable_compile_time_limit
+            and compile_time_for_code >= self.MAX_COMPILE_TIME_PER_CODE
+        ):
             log(2, "[Cache] Exceed max compile time per code, skip it\n")
             return CustomCode(None, False)
-        if compile_time_total >= self.MAX_COMPILE_TIME_TOTAL:
+        if (
+            enable_compile_time_limit
+            and compile_time_total >= self.MAX_COMPILE_TIME_TOTAL
+        ):
             log_once(
                 f"[SOT] Current compile time total is {compile_time_total}, exceed max compile time total {self.MAX_COMPILE_TIME_TOTAL}, fallback new function to dygraph"
             )

python/paddle/jit/sot/utils/__init__.py

@@ -19,6 +19,7 @@
     ENV_SOT_CE_DEBUG_MODE,
     ENV_SOT_COLLECT_INFO,
     ENV_SOT_ENABLE_0_SIZE_FALLBACK,
+    ENV_SOT_ENABLE_COMPILE_TIME_LIMIT,
     ENV_SOT_ENABLE_FASTER_GUARD,
     ENV_SOT_ENABLE_GUARD_TREE,
     ENV_SOT_ENABLE_STRICT_GUARD_CHECK,

python/paddle/jit/sot/utils/envs.py

@@ -157,6 +157,9 @@ def parse_parameterized_key(input_str: str) -> dict[str, list[str]]:
 ENV_SOT_SPECIALIZED_DIM_NUMBERS = StringEnvironmentVariable(
     "SOT_SPECIALIZED_DIM_NUMBERS", "0"
 )
+ENV_SOT_ENABLE_COMPILE_TIME_LIMIT = BooleanEnvironmentVariable(
+    "SOT_ENABLE_COMPILE_TIME_LIMIT", True
+)
 
 
 def update_ce_flags():