[aotinductor] Only autotune at compile time when enabled via config (#…

…129335) Summary: Pull Request resolved: #129335 --- Test Plan: ci Reviewed By: jingsh, desertfire, houseroad Differential Revision: D58931057
pytorch · Jun 24, 2024 · defde8d · defde8d
1 parent 1a54bb0
commit defde8d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/torch/_inductor/compile_fx.py b/torch/_inductor/compile_fx.py
@@ -1300,7 +1300,7 @@ def compile_fx(
  with config.patch(
  {
  "cpp_wrapper": False,
- "triton.autotune_at_compile_time": True,
+ "triton.autotune_at_compile_time": config.triton.autotune_at_compile_time,
  "triton.autotune_cublasLt": False,
  "triton.cudagraphs": False,
  "triton.store_cubin": True,

diff --git a/torch/_inductor/config.py b/torch/_inductor/config.py
@@ -697,7 +697,10 @@ class triton:
 
  # Tune the generated Triton kernels at compile time instead of first time they run
  autotune_at_compile_time = (
- os.environ.get("TORCHINDUCTOR_TRITON_AUTOTUNE_AT_COMPILE_TIME", "0") == "1"
+ os.environ.get(
+ "TORCHINDUCTOR_TRITON_AUTOTUNE_AT_COMPILE_TIME", "0" if is_fbcode() else "1"
+ )
+ == "1"
  )
 
  # should we stop a fusion to allow better tiling?