try changes from #3036

Signed-off-by: Anatoly Myachev <[email protected]>
intel · Dec 18, 2024 · 2a4b818 · 2a4b818
1 parent 8cf03e9
commit 2a4b818
Show file tree

Hide file tree

Showing 11 changed files with 23 additions and 3 deletions.
diff --git a/benchmarks/triton_kernels_benchmark/__init__.py b/benchmarks/triton_kernels_benchmark/__init__.py
@@ -1,4 +1,4 @@
-from .benchmark_testing import do_bench, assert_close, perf_report, Benchmark, USE_IPEX_OPTION, BENCHMARKING_METHOD  # type: ignore # noqa: F401
+from .benchmark_testing import do_bench, make_do_bench_for_autotune, assert_close, perf_report, Benchmark, USE_IPEX_OPTION, BENCHMARKING_METHOD  # type: ignore # noqa: F401
 
 if USE_IPEX_OPTION or BENCHMARKING_METHOD == "UPSTREAM_PYTORCH_PROFILER":
     from triton.runtime import driver

diff --git a/benchmarks/triton_kernels_benchmark/benchmark_testing.py b/benchmarks/triton_kernels_benchmark/benchmark_testing.py
@@ -237,6 +237,14 @@ def extract_kernels(funcs):
     raise NotImplementedError(f"BENCHMARKING_METHOD: {BENCHMARKING_METHOD} isn't implemented")
 
 
+def make_do_bench_for_autotune():
+
+    def autotuner_do_bench(*args, **kwargs):
+        return do_bench(*args, n_warmup=10, n_repeat=10, **kwargs)
+
+    return autotuner_do_bench
+
+
 def assert_close(x, y, atol=None, rtol=None, err_msg=""):
     import numpy as np
     import torch

diff --git a/benchmarks/triton_kernels_benchmark/flash_attention_fwd_benchmark.py b/benchmarks/triton_kernels_benchmark/flash_attention_fwd_benchmark.py
@@ -164,7 +164,7 @@ def _attn_fwd(Q, K, V, sm_scale, M, Out,  #
     for w in [8, 16, 32] \
     ]
 
-tuner = triton.autotune(configs, key=['N_CTX', 'BLOCK_DMODEL'])
+tuner = triton.autotune(configs, key=['N_CTX', 'BLOCK_DMODEL'], do_bench=benchmark_suit.make_do_bench_for_autotune())
 tune_attn_fwd = tuner(_attn_fwd)
 
 

diff --git a/benchmarks/triton_kernels_benchmark/fused_softmax.py b/benchmarks/triton_kernels_benchmark/fused_softmax.py
@@ -50,6 +50,7 @@ def naive_softmax(x):
         triton.Config({"threads_per_warp": 16}, num_warps=4),
     ],
     key=["BLOCK_SIZE_X", "BLOCK_SIZE_Y"],
+    do_bench=benchmark_suit.make_do_bench_for_autotune(),
 )
 @triton.jit
 def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols, BLOCK_SIZE_X: tl.constexpr,

diff --git a/benchmarks/triton_kernels_benchmark/gemm_benchmark.py b/benchmarks/triton_kernels_benchmark/gemm_benchmark.py
@@ -43,6 +43,7 @@
             num_stages=s, num_warps=32) for s in [2, 3]
     ],
     key=['M', 'N', 'K'],
+    do_bench=benchmark_suit.make_do_bench_for_autotune(),
 )
 @triton.jit
 def matmul_kernel_with_block_pointers(
@@ -116,6 +117,7 @@ def matmul_kernel_with_block_pointers(
             num_stages=s, num_warps=4) for s in [2]
     ],
     key=['M', 'N', 'K'],
+    do_bench=benchmark_suit.make_do_bench_for_autotune(),
 )
 @triton.jit
 def matmul_kernel_with_block_pointers_batched(

diff --git a/benchmarks/triton_kernels_benchmark/gemm_postop_addmatrix_benchmark.py b/benchmarks/triton_kernels_benchmark/gemm_postop_addmatrix_benchmark.py
@@ -35,6 +35,7 @@
             num_stages=2, num_warps=32),
     ],
     key=['M', 'N', 'K'],
+    do_bench=benchmark_suit.make_do_bench_for_autotune(),
 )
 @triton.jit
 def matmul_kernel_with_block_pointers(
@@ -109,6 +110,7 @@ def matmul_kernel_with_block_pointers(
             num_stages=2, num_warps=4),
     ],
     key=['M', 'N', 'K'],
+    do_bench=benchmark_suit.make_do_bench_for_autotune(),
 )
 @triton.jit
 def matmul_kernel_with_block_pointers_batched(

diff --git a/benchmarks/triton_kernels_benchmark/gemm_postop_gelu_benchmark.py b/benchmarks/triton_kernels_benchmark/gemm_postop_gelu_benchmark.py
@@ -54,6 +54,7 @@ def gelu(x):
             num_stages=2, num_warps=32),
     ],
     key=['M', 'N', 'K'],
+    do_bench=benchmark_suit.make_do_bench_for_autotune(),
 )
 @triton.jit
 def matmul_kernel_with_block_pointers(
@@ -122,6 +123,7 @@ def matmul_kernel_with_block_pointers(
             num_stages=2, num_warps=4),
     ],
     key=['M', 'N', 'K'],
+    do_bench=benchmark_suit.make_do_bench_for_autotune(),
 )
 @triton.jit
 def matmul_kernel_with_block_pointers_batched(

diff --git a/benchmarks/triton_kernels_benchmark/gemm_preop_exp_benchmark.py b/benchmarks/triton_kernels_benchmark/gemm_preop_exp_benchmark.py
@@ -36,6 +36,7 @@
             num_stages=2, num_warps=32),
     ],
     key=['M', 'N', 'K'],
+    do_bench=benchmark_suit.make_do_bench_for_autotune(),
 )
 @triton.jit
 def matmul_kernel_with_block_pointers(
@@ -107,6 +108,7 @@ def matmul_kernel_with_block_pointers(
             num_stages=2, num_warps=4),
     ],
     key=['M', 'N', 'K'],
+    do_bench=benchmark_suit.make_do_bench_for_autotune(),
 )
 @triton.jit
 def matmul_kernel_with_block_pointers_batched(

diff --git a/benchmarks/triton_kernels_benchmark/gemm_splitk_benchmark.py b/benchmarks/triton_kernels_benchmark/gemm_splitk_benchmark.py
@@ -15,6 +15,7 @@
                       num_stages=4, num_warps=32),
     ],
     key=['M', 'N', 'K'],
+    do_bench=benchmark_suit.make_do_bench_for_autotune(),
 )
 @triton.jit
 def _kernel(A, B, C,  #

diff --git a/benchmarks/triton_kernels_benchmark/gemm_streamk_benchmark.py b/benchmarks/triton_kernels_benchmark/gemm_streamk_benchmark.py
@@ -107,6 +107,7 @@ def mac_loop(
             num_stages=2, num_warps=32),
     ],
     key=['M', 'N', 'K'],
+    do_bench=benchmark_suit.make_do_bench_for_autotune(),
 )
 @triton.jit
 def first_wave(
@@ -143,6 +144,7 @@ def first_wave(
             num_stages=2, num_warps=32),
     ],
     key=['M', 'N', 'K'],
+    do_bench=benchmark_suit.make_do_bench_for_autotune(),
 )
 @triton.jit
 def full_tiles(

diff --git a/python/triton/runtime/autotuner.py b/python/triton/runtime/autotuner.py
@@ -357,7 +357,7 @@ def kernel(x_ptr, x_size, **META):
     def decorator(fn):
         return Autotuner(fn, fn.arg_names, configs, key, reset_to_zero, restore_value, pre_hook=pre_hook,
                          post_hook=post_hook, prune_configs_by=prune_configs_by, warmup=warmup, rep=rep,
-                         use_cuda_graph=use_cuda_graph)
+                         use_cuda_graph=use_cuda_graph, do_bench=do_bench)
 
     return decorator