From e2e0112584d72d578f12018377c2e121a752f7d3 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Sat, 11 Jan 2025 12:09:10 +0800 Subject: [PATCH] Update the benchmark input dimensions to vary based on batch size instead of height and width --- conv2d.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/conv2d.py b/conv2d.py index 7c6ad0c..6bb2174 100644 --- a/conv2d.py +++ b/conv2d.py @@ -215,8 +215,8 @@ def grid(meta): @triton.testing.perf_report( triton.testing.Benchmark( - x_names=["h", "w"], - x_vals=[8 * i for i in range(2, 33)], + x_names=["n"], + x_vals=[2**i for i in range(11)], line_arg="provider", line_vals=["ninetoothed", "torch", "triton"], line_names=["NineToothed", "PyTorch", "Triton"], @@ -226,9 +226,9 @@ def grid(meta): args={}, ) ) - def benchmark(h, w, provider): - n, c, _, _ = 64, 3, h, w - k, _, r, s = 64, c, 3, 3 + def benchmark(n, provider): + _, c, h, w = n, 512, 14, 14 + k, _, r, s = 512, c, 3, 3 dtype = torch.float16 input = torch.randn((n, c, h, w), dtype=dtype, device="cuda") filter = torch.randn((k, c, r, s), dtype=dtype, device="cuda")