Skip to content

Commit 8c58805

Browse files
jataylonaromero77amd
authored andcommitted
Bug fix and optimisation for persistent reduction kernel tuning
Original PR had incorrect indentation. Updated PR such that autotune will always add tiny configs, otherwise use the hinted configs only.
1 parent 2d423a3 commit 8c58805

File tree

1 file changed

+14
-14
lines changed

1 file changed

+14
-14
lines changed

torch/_inductor/runtime/triton_heuristics.py

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2780,20 +2780,20 @@ def _persistent_reduction_configs(
27802780
elif reduction_hint == ReductionHint.OUTER:
27812781
configs = configs[-1:]
27822782

2783-
if reduction_hint == ReductionHint.OUTER_TINY:
2784-
tiny_configs = [
2785-
triton_config_reduction(
2786-
size_hints,
2787-
2 * (256 // rnumel) if rnumel <= 256 else 1,
2788-
rnumel,
2789-
)
2790-
]
2791-
if max_autotune_enabled:
2792-
for tconfig in tiny_configs:
2793-
if tconfig not in configs:
2794-
configs.append(tconfig)
2795-
else:
2796-
configs = tiny_configs
2783+
tiny_configs = [
2784+
triton_config_reduction(
2785+
size_hints,
2786+
2 * (256 // rnumel) if rnumel <= 256 else 1,
2787+
rnumel,
2788+
)
2789+
]
2790+
2791+
if max_autotune_enabled:
2792+
for conf in tiny_configs:
2793+
if conf not in configs:
2794+
configs.append(conf)
2795+
elif reduction_hint == ReductionHint.OUTER_TINY:
2796+
configs = tiny_configs
27972797

27982798
for c in configs:
27992799
# we don't need Rn_BLOCK for persistent reduction

0 commit comments

Comments
 (0)