Skip to content

Commit ed0d0a7

Browse files
jataylojithunnair-amd
authored andcommitted
[SWDEV-539076] Initial naive foreach autotune support (#2377)
Adds initial autotuning for foreach support required for https://ontrack-internal.amd.com/browse/SWDEV-539076 4x improvement for some kernels Before: triton_for_fused_18.kd 🔍 | 4.986 ms | 4.986 ms | 2.493 ms | 2 |   triton_for_fused_6.kd 🔍 | 0.098 ms | 0.098 ms | 0.049 ms | 2 |   triton_for_fused_7.kd 🔍 | 0.036 ms | 0.036 ms | 0.018 ms | 2 |   After: triton_for_fused_18.kd 🔍 | 1.273 ms | 1.273 ms | 0.636 ms | 2 |   triton_for_fused_6.kd 🔍 | 0.044 ms | 0.044 ms | 0.022 ms | 2 |   triton_for_fused_7.kd 🔍 | 0.024 ms | 0.024 ms | 0.012 ms | 2 |   (cherry picked from commit f07b7f7)
1 parent 22d02e8 commit ed0d0a7

File tree

2 files changed

+13
-4
lines changed

2 files changed

+13
-4
lines changed

torch/_inductor/codegen/triton_combo_kernel.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -614,7 +614,7 @@ def jit_line(
614614
if heuristics == "foreach":
615615
heuristics_line = f"""
616616
@triton_heuristics.foreach(
617-
num_warps={self.num_warps},
617+
filename=__file__,
618618
triton_meta={triton_meta!r},
619619
inductor_meta={inductor_meta!r},
620620
)

torch/_inductor/runtime/triton_heuristics.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2779,20 +2779,29 @@ def user_autotune(
27792779
)
27802780

27812781

2782-
def foreach(triton_meta, num_warps, filename=None, inductor_meta=None):
2782+
def foreach(triton_meta, filename=None, inductor_meta=None):
27832783
"""
27842784
Compile a triton foreach kernel
27852785
"""
2786+
configs = []
2787+
if disable_pointwise_autotuning(inductor_meta) and not (
2788+
inductor_meta.get("max_autotune") or
2789+
inductor_meta.get("max_autotune_pointwise")
2790+
):
2791+
configs.append(triton.Config({}, num_stages=1, num_warps=8))
2792+
else:
2793+
for warps in [1, 2, 4, 8]:
2794+
configs.append(triton.Config({}, num_stages=1, num_warps=warps))
2795+
27862796
return cached_autotune(
27872797
None,
2788-
[triton.Config({}, num_stages=1, num_warps=num_warps)],
2798+
configs,
27892799
triton_meta=triton_meta,
27902800
inductor_meta=inductor_meta,
27912801
heuristic_type=HeuristicType.TEMPLATE,
27922802
filename=filename,
27932803
)
27942804

2795-
27962805
@dataclasses.dataclass
27972806
class GridExpr:
27982807
"""Generate code for grid size expressions in launcher"""

0 commit comments

Comments
 (0)