outer heuristic

PaulZhang12 · PaulZhang12 · commit 48733bbe36dc · 2025-08-11T13:55:57.000-07:00
ghstack-source-id: e1d6fa6 Pull Request resolved: #159093
diff --git a/torch/_inductor/runtime/triton_heuristics.py b/torch/_inductor/runtime/triton_heuristics.py
@@ -2454,7 +2454,7 @@ def pointwise(
 
 
 def _reduction_configs(
-    *, size_hints: dict[str, int], inductor_meta: dict[str, Any]
+    *, size_hints: dict[str, int], inductor_meta: dict[str, Any], is_dynamic=False
 ) -> list[Config]:
     reduction_hint = inductor_meta.get("reduction_hint", None)
 
@@ -2507,12 +2507,40 @@ def make_config(x, r, num_warps=None, num_stages=1, register_intensive=False):
                 register_intensive=register_intensive,
             )
 
+    def make_outer_config():
+        max_x_block = 256
+        load_factor = inductor_meta.get("num_load", 0)
+        x = size_hints["x"]
+
+        if x <= 8 * 4096:
+            x_block = 8
+        else:
+            x_block = min(max_x_block, next_power_of_2(x // 4096))
+            if x_block < 64:
+                x_block = 64
+        if is_dynamic:
+            # Dynamic shapes introduce a lot register pressure for indexing
+            outer_r_block = (
+                1
+                if load_factor >= 3
+                else min(next_power_of_2(max(rnumel, 128) // 128), 16)
+            )
+        else:
+            # Try to do reduction in 1 pass
+            outer_r_block = min(next_power_of_2(rnumel), 128)
+
+        if x_block * outer_r_block > 4096:
+            x_block = 2048 // outer_r_block
+
+        # Set register intensive to true by default as we try to maximize tiles with heuristic
+        return make_config(x_block, outer_r_block, register_intensive=True)
+
     contiguous_config = make_config(
         1,
         min(rnumel, MAX_R0_BLOCK),
         register_intensive=register_intensive,
     )
-    outer_config = make_config(64, 8, register_intensive=register_intensive)
+    outer_config = make_outer_config()
     tiny_config = make_config(
         2 * (256 // rnumel) if rnumel <= 256 else 1,
         min(rnumel, MAX_R0_BLOCK),
@@ -2637,7 +2665,11 @@ def reduction(
 
     assert triton_meta is not None
 
-    configs = _reduction_configs(size_hints=size_hints, inductor_meta=inductor_meta)
+    is_dynamic = any("ks" in k for k in triton_meta["signature"].keys())
+    configs = _reduction_configs(
+        size_hints=size_hints, inductor_meta=inductor_meta, is_dynamic=is_dynamic
+    )
+
     configs = _maybe_filter_configs_for_tma_restrictions(inductor_meta, configs)
     return cached_autotune(
         size_hints,
