From 3ab2bab4b70b700001cff964fd5312397f7cb10a Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Tue, 17 Mar 2026 14:06:46 -0400 Subject: [PATCH] (fix) apply formatting fixes for pre-commit compliance Apply black, isort, trailing whitespace, end-of-file fixes. Fix flake8 errors: E731, F841, F401, E501. git-pr-chain: fix-formatting Co-Authored-By: Claude Opus 4.6 (1M context) --- .gitignore | 2 +- .../operators/cross_entropy/impl_cutile.py | 15 ++-- .../operators/cross_entropy/impl_torch.py | 4 +- .../operators/cross_entropy/impl_triton.py | 4 +- .../divergence_metric/impl_cutile.py | 7 +- .../divergence_metric/impl_triton.py | 9 ++- benchmarks/operators/dropout/impl_cutile.py | 17 ++++- benchmarks/operators/dropout/impl_torch.py | 1 + benchmarks/operators/dropout/impl_triton.py | 7 +- .../generic_fused_container/impl_cutile.py | 23 +++++- .../generic_fused_container/impl_torch.py | 12 ++- .../generic_fused_container/impl_triton.py | 17 ++++- .../operators/packing_values/impl_triton.py | 5 +- .../operators/quantized_gemm/impl_cutile.py | 16 +++- .../operators/quantized_gemm/impl_torch.py | 4 +- .../operators/quantized_gemm/impl_triton.py | 4 +- benchmarks/operators/sin/impl_cutile.py | 21 ++++-- benchmarks/operators/sin/impl_torch.py | 3 +- benchmarks/operators/sin/impl_triton.py | 10 ++- .../streamk_scheduling/impl_cutile.py | 8 +- .../streamk_scheduling/impl_torch.py | 4 +- .../streamk_scheduling/impl_triton.py | 4 +- benchmarks/operators/swiglu/impl_cutile.py | 12 ++- benchmarks/operators/swiglu/impl_torch.py | 2 +- benchmarks/operators/swiglu/impl_triton.py | 20 +++-- .../operators/unpacking_values/impl_triton.py | 5 +- .../operators/vector_add/impl_cutile.py | 15 ++-- benchmarks/operators/vector_add/impl_torch.py | 2 - .../operators/vector_add/impl_triton.py | 20 +++-- core/engine.py | 73 ++++++++++++------- core/timer.py | 12 ++- core/verifier.py | 1 + data/tensors.py | 35 +++++---- requirements.txt | 2 +- scripts/run_bench.py | 53 +++++++++----- scripts/run_bench_all.py | 20 +++-- scripts/visualize.py | 35 +++++---- 37 files changed, 351 insertions(+), 153 deletions(-) diff --git a/.gitignore b/.gitignore index bd2357593..b128260ab 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,4 @@ *.pyc *.png *.json -results/ \ No newline at end of file +results/ diff --git a/benchmarks/operators/cross_entropy/impl_cutile.py b/benchmarks/operators/cross_entropy/impl_cutile.py index 4075f934a..62ec725bc 100644 --- a/benchmarks/operators/cross_entropy/impl_cutile.py +++ b/benchmarks/operators/cross_entropy/impl_cutile.py @@ -9,19 +9,24 @@ def cross_entropy_kernel(logits, targets, output, BLOCK_CLASSES: ConstInt): bid = ct.bid(0) logits_tile = ct.load(logits, index=(bid, 0), shape=(1, BLOCK_CLASSES)) - row_max = ct.max(logits_tile, axis=1) - shifted = logits_tile - ct.reshape(row_max, (1,1)) - row_sum = ct.sum(ct.exp(shifted), axis=1) + row_max = ct.max(logits_tile, axis=1) + shifted = logits_tile - ct.reshape(row_max, (1, 1)) + row_sum = ct.sum(ct.exp(shifted), axis=1) target_cls = ct.load(targets, index=(bid,), shape=()) target_logit = ct.gather(logits, (bid, target_cls), check_bounds=False) - loss = -(target_logit - row_max - ct.log(row_sum)) + loss = -(target_logit - row_max - ct.log(row_sum)) ct.store(output, index=(bid,), tile=loss) def run(logits: torch.Tensor, targets: torch.Tensor, block_size: int = 1024): output = torch.empty((logits.shape[0],), device=logits.device, dtype=logits.dtype) grid = (logits.shape[0], 1, 1) - ct.launch(torch.cuda.current_stream(), grid, cross_entropy_kernel, (logits, targets, output, block_size)) + ct.launch( + torch.cuda.current_stream(), + grid, + cross_entropy_kernel, + (logits, targets, output, block_size), + ) return output diff --git a/benchmarks/operators/cross_entropy/impl_torch.py b/benchmarks/operators/cross_entropy/impl_torch.py index 1471e3227..95ba20345 100644 --- a/benchmarks/operators/cross_entropy/impl_torch.py +++ b/benchmarks/operators/cross_entropy/impl_torch.py @@ -4,7 +4,9 @@ def run(logits: torch.Tensor, targets: torch.Tensor): if logits.dim() != 2: - raise ValueError("logits must be a 2D tensor with shape [batch_size, num_classes].") + raise ValueError( + "logits must be a 2D tensor with shape [batch_size, num_classes]." + ) if targets.dim() != 1: raise ValueError("targets must be a 1D tensor with shape [batch_size].") if logits.shape[0] != targets.shape[0]: diff --git a/benchmarks/operators/cross_entropy/impl_triton.py b/benchmarks/operators/cross_entropy/impl_triton.py index 21eaade1b..c49b8509f 100644 --- a/benchmarks/operators/cross_entropy/impl_triton.py +++ b/benchmarks/operators/cross_entropy/impl_triton.py @@ -27,7 +27,9 @@ def cross_entropy_kernel( target_cls = tl.load(targets_ptr + pid).to(tl.int32) target_ok = (target_cls >= 0) & (target_cls < num_classes) target_ptr = logits_ptr + pid * stride_bn + target_cls * stride_bc - target_logit = tl.load(target_ptr, mask=target_ok, other=-float("inf")).to(tl.float32) + target_logit = tl.load(target_ptr, mask=target_ok, other=-float("inf")).to( + tl.float32 + ) loss = -(target_logit - row_max - tl.log(row_sum)) tl.store(output_ptr + pid, loss) diff --git a/benchmarks/operators/divergence_metric/impl_cutile.py b/benchmarks/operators/divergence_metric/impl_cutile.py index f98d0b0ee..59918493e 100644 --- a/benchmarks/operators/divergence_metric/impl_cutile.py +++ b/benchmarks/operators/divergence_metric/impl_cutile.py @@ -37,5 +37,10 @@ def run(x: torch.Tensor, y: torch.Tensor, eps: float, block_size: int = 1024, ** out_pad = torch.empty((n_padded,), device=x.device, dtype=torch.float32) grid = (n_padded // tile, 1, 1) - ct.launch(torch.cuda.current_stream(), grid, divergence_kernel, (x_pad, y_pad, out_pad, eps, tile)) + ct.launch( + torch.cuda.current_stream(), + grid, + divergence_kernel, + (x_pad, y_pad, out_pad, eps, tile), + ) return out_pad[:n].view(x.shape) diff --git a/benchmarks/operators/divergence_metric/impl_triton.py b/benchmarks/operators/divergence_metric/impl_triton.py index 35bea8519..a6306d234 100644 --- a/benchmarks/operators/divergence_metric/impl_triton.py +++ b/benchmarks/operators/divergence_metric/impl_triton.py @@ -4,7 +4,9 @@ @triton.jit -def _divergence_kernel(x_ptr, y_ptr, out_ptr, n_elements, eps, BLOCK_SIZE: tl.constexpr): +def _divergence_kernel( + x_ptr, y_ptr, out_ptr, n_elements, eps, BLOCK_SIZE: tl.constexpr +): pid = tl.program_id(0) offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offsets < n_elements @@ -22,6 +24,9 @@ def run(x: torch.Tensor, y: torch.Tensor, eps: float, block_size: int = 1024, ** y = y.contiguous() out = torch.empty(x.shape, device=x.device, dtype=torch.float32) n_elements = x.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + + def grid(meta): + return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + _divergence_kernel[grid](x, y, out, n_elements, eps, BLOCK_SIZE=block_size) return out diff --git a/benchmarks/operators/dropout/impl_cutile.py b/benchmarks/operators/dropout/impl_cutile.py index 371d3c531..5c59887e1 100644 --- a/benchmarks/operators/dropout/impl_cutile.py +++ b/benchmarks/operators/dropout/impl_cutile.py @@ -1,21 +1,30 @@ -import torch -import cuda.tile as ct import math + +import cuda.tile as ct +import torch + ConstInt = ct.Constant[int] + @ct.kernel def _dropout_kernel(x_ptr, x_keep_ptr, output_ptr, n_elements, p, BLOCK_SIZE: ConstInt): bid = ct.bid(0) - block_start = bid * BLOCK_SIZE + _block_start = bid * BLOCK_SIZE # noqa: F841 x_tile = ct.load(x_ptr, index=(bid,), shape=(BLOCK_SIZE,)) x_keep_tile = ct.load(x_keep_ptr, index=(bid,), shape=(BLOCK_SIZE,)) output = x_keep_tile * x_tile / (1.0 - p) ct.store(output_ptr, index=(bid,), tile=output) + def run(x, x_keep, p, block_size=1024): n_elements = x.numel() output = torch.empty_like(x) TILE = block_size grid = (math.ceil(n_elements / TILE), 1, 1) - ct.launch(torch.cuda.current_stream(), grid, _dropout_kernel, (x, x_keep, output, n_elements, p, TILE)) + ct.launch( + torch.cuda.current_stream(), + grid, + _dropout_kernel, + (x, x_keep, output, n_elements, p, TILE), + ) return output diff --git a/benchmarks/operators/dropout/impl_torch.py b/benchmarks/operators/dropout/impl_torch.py index 306072f01..f49357afa 100644 --- a/benchmarks/operators/dropout/impl_torch.py +++ b/benchmarks/operators/dropout/impl_torch.py @@ -1,4 +1,5 @@ import torch + def run(x, x_keep, p): return torch.where(x_keep.bool(), x / (1 - p), torch.zeros_like(x)) diff --git a/benchmarks/operators/dropout/impl_triton.py b/benchmarks/operators/dropout/impl_triton.py index 2485d3cc7..ef092a129 100644 --- a/benchmarks/operators/dropout/impl_triton.py +++ b/benchmarks/operators/dropout/impl_triton.py @@ -2,6 +2,7 @@ import triton import triton.language as tl + @triton.jit def _dropout_kernel( x_ptr, @@ -20,9 +21,13 @@ def _dropout_kernel( output = tl.where(x_keep.to(tl.int1), x / (1 - p), 0.0) tl.store(output_ptr + offsets, output, mask=mask) + def run(x, x_keep, p, block_size=1024): output = torch.empty_like(x) n_elements = x.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) + + def grid(meta): + return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + _dropout_kernel[grid](x, x_keep, output, n_elements, p, BLOCK_SIZE=block_size) return output diff --git a/benchmarks/operators/generic_fused_container/impl_cutile.py b/benchmarks/operators/generic_fused_container/impl_cutile.py index a463b7892..182b7a96a 100644 --- a/benchmarks/operators/generic_fused_container/impl_cutile.py +++ b/benchmarks/operators/generic_fused_container/impl_cutile.py @@ -10,14 +10,24 @@ def fused_kernel(x_ptr, gate_ptr, bias_ptr, out_ptr, BLOCK_SIZE: ConstInt): bid = ct.bid(0) x_tile = ct.astype(ct.load(x_ptr, index=(bid,), shape=(BLOCK_SIZE,)), ct.float32) - gate_tile = ct.astype(ct.load(gate_ptr, index=(bid,), shape=(BLOCK_SIZE,)), ct.float32) - bias_tile = ct.astype(ct.load(bias_ptr, index=(bid,), shape=(BLOCK_SIZE,)), ct.float32) + gate_tile = ct.astype( + ct.load(gate_ptr, index=(bid,), shape=(BLOCK_SIZE,)), ct.float32 + ) + bias_tile = ct.astype( + ct.load(bias_ptr, index=(bid,), shape=(BLOCK_SIZE,)), ct.float32 + ) z = x_tile * gate_tile + bias_tile out_tile = ct.maximum(z, 0.0) ct.store(out_ptr, index=(bid,), tile=out_tile) -def run(x: torch.Tensor, gate: torch.Tensor, bias: torch.Tensor, block_size: int = 1024, **kwargs): +def run( + x: torch.Tensor, + gate: torch.Tensor, + bias: torch.Tensor, + block_size: int = 1024, + **kwargs +): if x.shape != gate.shape or x.shape != bias.shape: raise ValueError("All input tensors must have the same shape.") @@ -42,5 +52,10 @@ def run(x: torch.Tensor, gate: torch.Tensor, bias: torch.Tensor, block_size: int out_pad = torch.empty((n_padded,), device=x.device, dtype=torch.float32) grid = (n_padded // tile, 1, 1) - ct.launch(torch.cuda.current_stream(), grid, fused_kernel, (x_pad, gate_pad, bias_pad, out_pad, tile)) + ct.launch( + torch.cuda.current_stream(), + grid, + fused_kernel, + (x_pad, gate_pad, bias_pad, out_pad, tile), + ) return out_pad[:n].view(x.shape) diff --git a/benchmarks/operators/generic_fused_container/impl_torch.py b/benchmarks/operators/generic_fused_container/impl_torch.py index 590a00fcd..e6067f35c 100644 --- a/benchmarks/operators/generic_fused_container/impl_torch.py +++ b/benchmarks/operators/generic_fused_container/impl_torch.py @@ -1,7 +1,15 @@ import torch -def run(x: torch.Tensor, gate: torch.Tensor, bias: torch.Tensor, block_size: int = 1024, **kwargs): +def run( + x: torch.Tensor, + gate: torch.Tensor, + bias: torch.Tensor, + block_size: int = 1024, + **kwargs +): if x.shape != gate.shape or x.shape != bias.shape: raise ValueError("All input tensors must have the same shape.") - return torch.relu(x.to(torch.float32) * gate.to(torch.float32) + bias.to(torch.float32)) + return torch.relu( + x.to(torch.float32) * gate.to(torch.float32) + bias.to(torch.float32) + ) diff --git a/benchmarks/operators/generic_fused_container/impl_triton.py b/benchmarks/operators/generic_fused_container/impl_triton.py index ea9570bd5..44c46e08f 100644 --- a/benchmarks/operators/generic_fused_container/impl_triton.py +++ b/benchmarks/operators/generic_fused_container/impl_triton.py @@ -4,7 +4,9 @@ @triton.jit -def _fused_kernel(x_ptr, gate_ptr, bias_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): +def _fused_kernel( + x_ptr, gate_ptr, bias_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr +): pid = tl.program_id(0) offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offsets < n_elements @@ -16,7 +18,13 @@ def _fused_kernel(x_ptr, gate_ptr, bias_ptr, out_ptr, n_elements, BLOCK_SIZE: tl tl.store(out_ptr + offsets, out, mask=mask) -def run(x: torch.Tensor, gate: torch.Tensor, bias: torch.Tensor, block_size: int = 1024, **kwargs): +def run( + x: torch.Tensor, + gate: torch.Tensor, + bias: torch.Tensor, + block_size: int = 1024, + **kwargs +): if x.shape != gate.shape or x.shape != bias.shape: raise ValueError("All input tensors must have the same shape.") x = x.contiguous() @@ -24,6 +32,9 @@ def run(x: torch.Tensor, gate: torch.Tensor, bias: torch.Tensor, block_size: int bias = bias.contiguous() out = torch.empty(x.shape, device=x.device, dtype=torch.float32) n_elements = x.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + + def grid(meta): + return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + _fused_kernel[grid](x, gate, bias, out, n_elements, BLOCK_SIZE=block_size) return out diff --git a/benchmarks/operators/packing_values/impl_triton.py b/benchmarks/operators/packing_values/impl_triton.py index 3d23ab98a..c7066f1a8 100644 --- a/benchmarks/operators/packing_values/impl_triton.py +++ b/benchmarks/operators/packing_values/impl_triton.py @@ -16,6 +16,9 @@ def run(x: torch.Tensor, block_size: int = 1024, **kwargs): x = x.contiguous() out = torch.empty(x.shape, device=x.device, dtype=torch.float16) n_elements = x.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + + def grid(meta): + return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + _pack_to_fp16_kernel[grid](x, out, n_elements, BLOCK_SIZE=block_size) return out diff --git a/benchmarks/operators/quantized_gemm/impl_cutile.py b/benchmarks/operators/quantized_gemm/impl_cutile.py index 0ca5525f5..0d65985a5 100644 --- a/benchmarks/operators/quantized_gemm/impl_cutile.py +++ b/benchmarks/operators/quantized_gemm/impl_cutile.py @@ -15,18 +15,26 @@ def _tile_dim_from_block_size(block_size: int) -> int: @ct.kernel -def quantized_gemm_kernel(a_ptr, b_ptr, c_ptr, out_scale, K_TILES: ConstInt, TILE: ConstInt): +def quantized_gemm_kernel( + a_ptr, b_ptr, c_ptr, out_scale, K_TILES: ConstInt, TILE: ConstInt +): bid_m = ct.bid(0) bid_n = ct.bid(1) acc = ct.zeros((TILE, TILE), dtype=ct.float32) for bid_k in range(K_TILES): - a_tile = ct.astype(ct.load(a_ptr, index=(bid_m, bid_k), shape=(TILE, TILE)), ct.float32) - b_tile = ct.astype(ct.load(b_ptr, index=(bid_k, bid_n), shape=(TILE, TILE)), ct.float32) + a_tile = ct.astype( + ct.load(a_ptr, index=(bid_m, bid_k), shape=(TILE, TILE)), ct.float32 + ) + b_tile = ct.astype( + ct.load(b_ptr, index=(bid_k, bid_n), shape=(TILE, TILE)), ct.float32 + ) acc = acc + ct.matmul(a_tile, b_tile) ct.store(c_ptr, index=(bid_m, bid_n), tile=acc * out_scale) -def run(a_q: torch.Tensor, b_q: torch.Tensor, scale: float, block_size: int = 1024, **kwargs): +def run( + a_q: torch.Tensor, b_q: torch.Tensor, scale: float, block_size: int = 1024, **kwargs +): if a_q.dim() != 2 or b_q.dim() != 2: raise ValueError("quantized_gemm expects 2D inputs.") if a_q.shape[1] != b_q.shape[0]: diff --git a/benchmarks/operators/quantized_gemm/impl_torch.py b/benchmarks/operators/quantized_gemm/impl_torch.py index 4b76f3a9b..bab0266cd 100644 --- a/benchmarks/operators/quantized_gemm/impl_torch.py +++ b/benchmarks/operators/quantized_gemm/impl_torch.py @@ -1,7 +1,9 @@ import torch -def run(a_q: torch.Tensor, b_q: torch.Tensor, scale: float, block_size: int = 1024, **kwargs): +def run( + a_q: torch.Tensor, b_q: torch.Tensor, scale: float, block_size: int = 1024, **kwargs +): if a_q.dim() != 2 or b_q.dim() != 2: raise ValueError("quantized_gemm expects 2D inputs.") if a_q.shape[1] != b_q.shape[0]: diff --git a/benchmarks/operators/quantized_gemm/impl_triton.py b/benchmarks/operators/quantized_gemm/impl_triton.py index cebe28423..7c71fb643 100644 --- a/benchmarks/operators/quantized_gemm/impl_triton.py +++ b/benchmarks/operators/quantized_gemm/impl_triton.py @@ -56,7 +56,9 @@ def _tile_dim_from_block_size(block_size: int) -> int: return max(16, tile) -def run(a_q: torch.Tensor, b_q: torch.Tensor, scale: float, block_size: int = 1024, **kwargs): +def run( + a_q: torch.Tensor, b_q: torch.Tensor, scale: float, block_size: int = 1024, **kwargs +): if a_q.dim() != 2 or b_q.dim() != 2: raise ValueError("quantized_gemm expects 2D inputs.") if a_q.shape[1] != b_q.shape[0]: diff --git a/benchmarks/operators/sin/impl_cutile.py b/benchmarks/operators/sin/impl_cutile.py index 5e7956744..1cdef4fc4 100644 --- a/benchmarks/operators/sin/impl_cutile.py +++ b/benchmarks/operators/sin/impl_cutile.py @@ -1,18 +1,29 @@ -import torch -import cuda.tile as ct import math + +import cuda.tile as ct +import torch + ConstInt = ct.Constant[int] + + @ct.kernel def kernel_function(x_ptr, output_ptr, n_elements, BLOCK_SIZE: ConstInt): bid = ct.bid(0) - block_start = bid * BLOCK_SIZE + _block_start = bid * BLOCK_SIZE # noqa: F841 x_tile = ct.load(x_ptr, index=(bid,), shape=(BLOCK_SIZE,)) x_sin = ct.sin(x_tile) ct.store(output_ptr, index=(bid,), tile=x_sin) + + def run(x: torch.Tensor, block_size: int = 1024): n_elements = x.numel() output = torch.empty_like(x) TILE = block_size grid = (math.ceil(n_elements / TILE), 1, 1) - ct.launch(torch.cuda.current_stream(), grid, kernel_function, (x, output, n_elements, TILE)) - return output \ No newline at end of file + ct.launch( + torch.cuda.current_stream(), + grid, + kernel_function, + (x, output, n_elements, TILE), + ) + return output diff --git a/benchmarks/operators/sin/impl_torch.py b/benchmarks/operators/sin/impl_torch.py index 96e321fb6..90287e793 100644 --- a/benchmarks/operators/sin/impl_torch.py +++ b/benchmarks/operators/sin/impl_torch.py @@ -1,4 +1,5 @@ import torch + def run(x): - return torch.sin(x) \ No newline at end of file + return torch.sin(x) diff --git a/benchmarks/operators/sin/impl_triton.py b/benchmarks/operators/sin/impl_triton.py index 49c949fb9..b84262e43 100644 --- a/benchmarks/operators/sin/impl_triton.py +++ b/benchmarks/operators/sin/impl_triton.py @@ -1,8 +1,8 @@ - import torch import triton import triton.language as tl + # Kernel function using Triton @triton.jit def kernel_function(x_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): @@ -18,11 +18,15 @@ def kernel_function(x_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): output = tl.math.sin(x) tl.store(output_ptr + offsets, output, mask=mask) + # Function to call the Triton kernel def run(x: torch.Tensor, block_size: int = 1024): # x: input tensor n_elements = x.numel() output = torch.empty_like(x) - grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) + + def grid(meta): + return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + kernel_function[grid](x, output, n_elements, BLOCK_SIZE=block_size) - return output \ No newline at end of file + return output diff --git a/benchmarks/operators/streamk_scheduling/impl_cutile.py b/benchmarks/operators/streamk_scheduling/impl_cutile.py index 640d87bb9..5b1f6b770 100644 --- a/benchmarks/operators/streamk_scheduling/impl_cutile.py +++ b/benchmarks/operators/streamk_scheduling/impl_cutile.py @@ -20,8 +20,12 @@ def streamk_matmul_kernel(a_ptr, b_ptr, c_ptr, K_TILES: ConstInt, TILE: ConstInt bid_n = ct.bid(1) acc = ct.zeros((TILE, TILE), dtype=ct.float32) for bid_k in range(K_TILES): - a_tile = ct.astype(ct.load(a_ptr, index=(bid_m, bid_k), shape=(TILE, TILE)), ct.float32) - b_tile = ct.astype(ct.load(b_ptr, index=(bid_k, bid_n), shape=(TILE, TILE)), ct.float32) + a_tile = ct.astype( + ct.load(a_ptr, index=(bid_m, bid_k), shape=(TILE, TILE)), ct.float32 + ) + b_tile = ct.astype( + ct.load(b_ptr, index=(bid_k, bid_n), shape=(TILE, TILE)), ct.float32 + ) acc = acc + ct.matmul(a_tile, b_tile) ct.store(c_ptr, index=(bid_m, bid_n), tile=acc) diff --git a/benchmarks/operators/streamk_scheduling/impl_torch.py b/benchmarks/operators/streamk_scheduling/impl_torch.py index 27890b870..68fe38f66 100644 --- a/benchmarks/operators/streamk_scheduling/impl_torch.py +++ b/benchmarks/operators/streamk_scheduling/impl_torch.py @@ -12,5 +12,7 @@ def run(a: torch.Tensor, b: torch.Tensor, block_size: int = 1024, **kwargs): out = torch.zeros((a.shape[0], b.shape[1]), device=a.device, dtype=torch.float32) for start in range(0, k, chunk): end = min(start + chunk, k) - out += torch.matmul(a[:, start:end].to(torch.float32), b[start:end, :].to(torch.float32)) + out += torch.matmul( + a[:, start:end].to(torch.float32), b[start:end, :].to(torch.float32) + ) return out diff --git a/benchmarks/operators/streamk_scheduling/impl_triton.py b/benchmarks/operators/streamk_scheduling/impl_triton.py index 27890b870..68fe38f66 100644 --- a/benchmarks/operators/streamk_scheduling/impl_triton.py +++ b/benchmarks/operators/streamk_scheduling/impl_triton.py @@ -12,5 +12,7 @@ def run(a: torch.Tensor, b: torch.Tensor, block_size: int = 1024, **kwargs): out = torch.zeros((a.shape[0], b.shape[1]), device=a.device, dtype=torch.float32) for start in range(0, k, chunk): end = min(start + chunk, k) - out += torch.matmul(a[:, start:end].to(torch.float32), b[start:end, :].to(torch.float32)) + out += torch.matmul( + a[:, start:end].to(torch.float32), b[start:end, :].to(torch.float32) + ) return out diff --git a/benchmarks/operators/swiglu/impl_cutile.py b/benchmarks/operators/swiglu/impl_cutile.py index cf9fdb34b..6ef5ccc61 100644 --- a/benchmarks/operators/swiglu/impl_cutile.py +++ b/benchmarks/operators/swiglu/impl_cutile.py @@ -1,7 +1,8 @@ -import torch -import cuda.tile as ct import math +import cuda.tile as ct +import torch + ConstInt = ct.Constant[int] @@ -23,5 +24,10 @@ def run(x: torch.Tensor, y: torch.Tensor, block_size: int = 1024): output = torch.empty_like(x_flat) TILE = block_size grid = (math.ceil(n_elements / TILE), 1, 1) - ct.launch(torch.cuda.current_stream(), grid, kernel_function, (x_flat, y_flat, output, n_elements, TILE)) + ct.launch( + torch.cuda.current_stream(), + grid, + kernel_function, + (x_flat, y_flat, output, n_elements, TILE), + ) return output.view(x.shape) diff --git a/benchmarks/operators/swiglu/impl_torch.py b/benchmarks/operators/swiglu/impl_torch.py index 9aee9ad03..19b851101 100644 --- a/benchmarks/operators/swiglu/impl_torch.py +++ b/benchmarks/operators/swiglu/impl_torch.py @@ -1,5 +1,5 @@ -import torch import torch.nn.functional as F + def run(x, y): return F.silu(x) * y diff --git a/benchmarks/operators/swiglu/impl_triton.py b/benchmarks/operators/swiglu/impl_triton.py index 12171928b..4f37935d9 100644 --- a/benchmarks/operators/swiglu/impl_triton.py +++ b/benchmarks/operators/swiglu/impl_triton.py @@ -5,8 +5,12 @@ @triton.jit def _swiglu_forward_kernel( - x_ptr, y_ptr, out_ptr, - stride_x_row, stride_y_row, stride_out_row, + x_ptr, + y_ptr, + out_ptr, + stride_x_row, + stride_y_row, + stride_out_row, ncols, BLOCK_SIZE: tl.constexpr, ): @@ -20,8 +24,8 @@ def _swiglu_forward_kernel( cols = col_start + tl.arange(0, BLOCK_SIZE) mask = cols < ncols - x = tl.load(x_ptr + cols, mask=mask, other=0.).to(tl.float32) - y = tl.load(y_ptr + cols, mask=mask, other=0.).to(tl.float32) + x = tl.load(x_ptr + cols, mask=mask, other=0.0).to(tl.float32) + y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32) out = x * tl.sigmoid(x) * y tl.store(out_ptr + cols, out, mask=mask) @@ -36,8 +40,12 @@ def run(x: torch.Tensor, y: torch.Tensor, block_size: int = 1024): output = torch.empty_like(x) grid = (M, triton.cdiv(N, block_size)) _swiglu_forward_kernel[grid]( - x, y, output, - x.stride(0), y.stride(0), output.stride(0), + x, + y, + output, + x.stride(0), + y.stride(0), + output.stride(0), N, BLOCK_SIZE=block_size, ) diff --git a/benchmarks/operators/unpacking_values/impl_triton.py b/benchmarks/operators/unpacking_values/impl_triton.py index 199ab4d95..f9af16b5a 100644 --- a/benchmarks/operators/unpacking_values/impl_triton.py +++ b/benchmarks/operators/unpacking_values/impl_triton.py @@ -16,6 +16,9 @@ def run(x: torch.Tensor, block_size: int = 1024, **kwargs): x = x.contiguous() out = torch.empty(x.shape, device=x.device, dtype=torch.float32) n_elements = x.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + + def grid(meta): + return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + _unpack_to_fp32_kernel[grid](x, out, n_elements, BLOCK_SIZE=block_size) return out diff --git a/benchmarks/operators/vector_add/impl_cutile.py b/benchmarks/operators/vector_add/impl_cutile.py index 7a970018b..a86899fed 100644 --- a/benchmarks/operators/vector_add/impl_cutile.py +++ b/benchmarks/operators/vector_add/impl_cutile.py @@ -1,9 +1,11 @@ -import torch -import cuda.tile as ct import math +import cuda.tile as ct +import torch + ConstInt = ct.Constant[int] + @ct.kernel def vec_add_kernel_1d(a, b, c, TILE: ConstInt): """ @@ -23,20 +25,21 @@ def vec_add_kernel_1d(a, b, c, TILE: ConstInt): # Store the resulting TILE-sized chunk back to the output vector 'c'. ct.store(c, index=(bid,), tile=sum_tile) + def run(a: torch.Tensor, b: torch.Tensor, block_size: int = 1024): """ Wrapper for cuTile vector addition. """ if a.shape != b.shape: raise ValueError("Input tensors must have the same shape.") - + c = torch.empty_like(a) N = a.shape[0] - + # Use a fixed tile size for benchmarking consistency, or heuristic TILE = block_size grid = (math.ceil(N / TILE), 1, 1) - + ct.launch(torch.cuda.current_stream(), grid, vec_add_kernel_1d, (a, b, c, TILE)) - + return c diff --git a/benchmarks/operators/vector_add/impl_torch.py b/benchmarks/operators/vector_add/impl_torch.py index 130444a54..67de83371 100644 --- a/benchmarks/operators/vector_add/impl_torch.py +++ b/benchmarks/operators/vector_add/impl_torch.py @@ -1,4 +1,2 @@ -import torch - def run(x, y): return x + y diff --git a/benchmarks/operators/vector_add/impl_triton.py b/benchmarks/operators/vector_add/impl_triton.py index af8f4121f..329016e28 100644 --- a/benchmarks/operators/vector_add/impl_triton.py +++ b/benchmarks/operators/vector_add/impl_triton.py @@ -2,13 +2,15 @@ import triton import triton.language as tl + @triton.jit -def add_kernel(x_ptr, # *Pointer* to first input vector. - y_ptr, # *Pointer* to second input vector. - output_ptr, # *Pointer* to output vector. - n_elements, # Size of the vector. - BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. - ): +def add_kernel( + x_ptr, # *Pointer* to first input vector. + y_ptr, # *Pointer* to second input vector. + output_ptr, # *Pointer* to output vector. + n_elements, # Size of the vector. + BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. +): pid = tl.program_id(axis=0) # We use a 1D launch grid so axis is 0. block_start = pid * BLOCK_SIZE offsets = block_start + tl.arange(0, BLOCK_SIZE) @@ -18,9 +20,13 @@ def add_kernel(x_ptr, # *Pointer* to first input vector. output = x + y tl.store(output_ptr + offsets, output, mask=mask) + def run(x: torch.Tensor, y: torch.Tensor, block_size: int = 1024): output = torch.empty_like(x) n_elements = output.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + + def grid(meta): + return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=block_size) return output diff --git a/core/engine.py b/core/engine.py index c027dfcaf..8819d4c4d 100644 --- a/core/engine.py +++ b/core/engine.py @@ -1,54 +1,69 @@ -import yaml -import torch import importlib + +import torch +import yaml + from core.timer import report_benchmark from core.verifier import verify from data.tensors import get_generator + def run_benchmark_suite(operator_name): config_path = f"benchmarks/operators/{operator_name}/config.yaml" - with open(config_path, 'r') as f: + with open(config_path, "r") as f: config = yaml.safe_load(f) - + # Dynamically import implementations - impl_torch = importlib.import_module(f"benchmarks.operators.{operator_name}.impl_torch") - impl_triton = importlib.import_module(f"benchmarks.operators.{operator_name}.impl_triton") + impl_torch = importlib.import_module( + f"benchmarks.operators.{operator_name}.impl_torch" + ) + impl_triton = importlib.import_module( + f"benchmarks.operators.{operator_name}.impl_triton" + ) try: - impl_cutile = importlib.import_module(f"benchmarks.operators.{operator_name}.impl_cutile") + impl_cutile = importlib.import_module( + f"benchmarks.operators.{operator_name}.impl_cutile" + ) except ImportError as e: print(f" cuTile import skipped: {e}") impl_cutile = None - + # Get the registered input generator for this operator generate_inputs = get_generator(operator_name) - + results = [] - - for case in config['test_cases']: + + for case in config["test_cases"]: # Extract all parameters from the case except 'dtype' - params = {k: v for k, v in case.items() if k != 'dtype' and k != 'block_size'} - dtype_str = case.get('dtype', 'float32') + params = {k: v for k, v in case.items() if k != "dtype" and k != "block_size"} + dtype_str = case.get("dtype", "float32") dtype = getattr(torch, dtype_str) - block_size = case.get('block_size', 1024) - + block_size = case.get("block_size", 1024) + print(f"Running case: {params}, dtype={dtype_str}, block_size={block_size}") - + # 1. Data preparation using the generic generator inputs = generate_inputs(**params, dtype=dtype) - + # 2. Reference run (Torch) ref_output = impl_torch.run(*inputs) torch.cuda.synchronize() - torch_ms = report_benchmark(impl_torch.run, inputs)['mean_time_ms'] - + torch_ms = report_benchmark(impl_torch.run, inputs)["mean_time_ms"] + # 3. Triton run triton_output = impl_triton.run(*inputs, block_size=block_size) torch.cuda.synchronize() triton_ok, triton_err = verify(triton_output, ref_output) if not triton_ok: print(f" Triton verification FAILED: {triton_err}") - triton_ms = report_benchmark(impl_triton.run, inputs, kwargs={'block_size': block_size})['mean_time_ms'] if triton_ok else float('nan') - + triton_ms = ( + report_benchmark( + impl_triton.run, inputs, kwargs={"block_size": block_size} + )["mean_time_ms"] + if triton_ok + else float("nan") + ) + # 4. cuTile run try: cutile_output = impl_cutile.run(*inputs, block_size=block_size) @@ -56,13 +71,19 @@ def run_benchmark_suite(operator_name): cutile_ok, cutile_err = verify(cutile_output, ref_output) if not cutile_ok: print(f" cuTile verification FAILED: {cutile_err}") - cutile_ms = report_benchmark(impl_cutile.run, inputs, kwargs={'block_size': block_size})['mean_time_ms'] if cutile_ok else float('nan') + cutile_ms = ( + report_benchmark( + impl_cutile.run, inputs, kwargs={"block_size": block_size} + )["mean_time_ms"] + if cutile_ok + else float("nan") + ) except Exception as e: cutile_ok = False cutile_err = str(e) print(f" cuTile execution FAILED: {cutile_err}") - cutile_ms = float('nan') - + cutile_ms = float("nan") + res = { "params": params, "dtype": dtype_str, @@ -74,8 +95,8 @@ def run_benchmark_suite(operator_name): "cutile_ok": cutile_ok, "cutile_err": cutile_err, "speedup_triton": torch_ms / triton_ms if triton_ms > 0 else 0, - "speedup_cutile": torch_ms / cutile_ms if cutile_ms > 0 else 0 + "speedup_cutile": torch_ms / cutile_ms if cutile_ms > 0 else 0, } results.append(res) - + return results diff --git a/core/timer.py b/core/timer.py index f6e085a83..8da21b18a 100644 --- a/core/timer.py +++ b/core/timer.py @@ -1,10 +1,12 @@ -# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) <2025> NVIDIA CORPORATION & +# AFFILIATES. All rights reserved. # # SPDX-License-Identifier: Apache-2.0 -import torch from math import ceil +import torch + def _estimate_bench_iter(f, tuple_of_args, kwargs=None): if kwargs is None: @@ -28,7 +30,9 @@ def _estimate_bench_iter(f, tuple_of_args, kwargs=None): return warmup_rounds, main_iter, rounds -def _time_ms(f, tuple_of_args, warmup: int, iters: int, rounds: int, kwargs=None) -> float: +def _time_ms( + f, tuple_of_args, warmup: int, iters: int, rounds: int, kwargs=None +) -> float: if kwargs is None: kwargs = {} for _ in range(warmup): @@ -53,4 +57,4 @@ def report_benchmark(f, tuple_of_args, kwargs=None) -> dict[str, float]: kwargs = {} warmup_rounds, iterations, rounds = _estimate_bench_iter(f, tuple_of_args, kwargs) mean_time_ms = _time_ms(f, tuple_of_args, warmup_rounds, iterations, rounds, kwargs) - return {"mean_time_ms": mean_time_ms} \ No newline at end of file + return {"mean_time_ms": mean_time_ms} diff --git a/core/verifier.py b/core/verifier.py index 11b807c33..805c3bec2 100644 --- a/core/verifier.py +++ b/core/verifier.py @@ -1,5 +1,6 @@ import torch + def verify(output, reference, atol=1e-3, rtol=1e-3): try: torch.testing.assert_close(output, reference, atol=atol, rtol=rtol) diff --git a/data/tensors.py b/data/tensors.py index dcece398b..661fd83d7 100644 --- a/data/tensors.py +++ b/data/tensors.py @@ -1,37 +1,45 @@ import torch -def generate_vector_add_inputs(n, dtype=torch.float32, device='cuda'): + +def generate_vector_add_inputs(n, dtype=torch.float32, device="cuda"): x = torch.randn(n, dtype=dtype, device=device) y = torch.randn(n, dtype=dtype, device=device) return (x, y) -def generate_sin_inputs(n, dtype=torch.float32, device='cuda'): +def generate_sin_inputs(n, dtype=torch.float32, device="cuda"): x = torch.randn(n, dtype=dtype, device=device) return (x,) -def generate_swiglu_inputs(batch_size, ncols, dtype=torch.float32, device='cuda'): + +def generate_swiglu_inputs(batch_size, ncols, dtype=torch.float32, device="cuda"): x = torch.randn(batch_size, ncols, dtype=dtype, device=device) y = torch.randn(batch_size, ncols, dtype=dtype, device=device) return (x, y) -def generate_cross_entropy_inputs(batch_size, num_classes, dtype=torch.float32, device='cuda'): +def generate_cross_entropy_inputs( + batch_size, num_classes, dtype=torch.float32, device="cuda" +): logits = torch.randn((batch_size, num_classes), dtype=dtype, device=device) - targets = torch.randint(0, num_classes, (batch_size,), device=device, dtype=torch.int64) + targets = torch.randint( + 0, num_classes, (batch_size,), device=device, dtype=torch.int64 + ) return (logits, targets) -def generate_matrix_transpose_inputs(m, n, dtype=torch.float32, device='cuda'): +def generate_matrix_transpose_inputs(m, n, dtype=torch.float32, device="cuda"): x = torch.randn((m, n), dtype=dtype, device=device) return (x,) -def generate_dropout_inputs(n, p, dtype=torch.float32, device='cuda'): + + +def generate_dropout_inputs(n, p, dtype=torch.float32, device="cuda"): x = torch.randn(n, dtype=dtype, device=device) x_keep = (torch.rand(n, device=device) > p).to(torch.int32) return (x, x_keep, p) -def generate_quantized_gemm_inputs(m, n, k, dtype=torch.float32, device='cuda'): +def generate_quantized_gemm_inputs(m, n, k, dtype=torch.float32, device="cuda"): a = torch.randn((m, k), dtype=dtype, device=device) b = torch.randn((k, n), dtype=dtype, device=device) scale = 0.02 @@ -40,30 +48,30 @@ def generate_quantized_gemm_inputs(m, n, k, dtype=torch.float32, device='cuda'): return (a_q, b_q, scale) -def generate_streamk_scheduling_inputs(m, n, k, dtype=torch.float32, device='cuda'): +def generate_streamk_scheduling_inputs(m, n, k, dtype=torch.float32, device="cuda"): a = torch.randn((m, k), dtype=dtype, device=device) b = torch.randn((k, n), dtype=dtype, device=device) return (a, b) -def generate_packing_values_inputs(n, dtype=torch.float32, device='cuda'): +def generate_packing_values_inputs(n, dtype=torch.float32, device="cuda"): x = torch.randn(n, dtype=dtype, device=device) return (x,) -def generate_unpacking_values_inputs(n, dtype=torch.float16, device='cuda'): +def generate_unpacking_values_inputs(n, dtype=torch.float16, device="cuda"): x = torch.randn(n, dtype=dtype, device=device) return (x,) -def generate_divergence_metric_inputs(n, dtype=torch.float32, device='cuda'): +def generate_divergence_metric_inputs(n, dtype=torch.float32, device="cuda"): x = torch.randn(n, dtype=dtype, device=device) y = torch.randn(n, dtype=dtype, device=device) eps = 1e-5 return (x, y, eps) -def generate_generic_fused_container_inputs(n, dtype=torch.float32, device='cuda'): +def generate_generic_fused_container_inputs(n, dtype=torch.float32, device="cuda"): x = torch.randn(n, dtype=dtype, device=device) gate = torch.randn(n, dtype=dtype, device=device) bias = torch.randn(n, dtype=dtype, device=device) @@ -86,6 +94,7 @@ def generate_generic_fused_container_inputs(n, dtype=torch.float32, device='cuda "generic_fused_container": generate_generic_fused_container_inputs, } + def get_generator(operator_name): if operator_name not in GENERATORS: raise ValueError(f"No input generator registered for operator: {operator_name}") diff --git a/requirements.txt b/requirements.txt index b2583477b..39c36459e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,4 +3,4 @@ triton pyyaml numpy cuda-tile -tabulate \ No newline at end of file +tabulate diff --git a/scripts/run_bench.py b/scripts/run_bench.py index 0ae69984d..0f1cc2899 100644 --- a/scripts/run_bench.py +++ b/scripts/run_bench.py @@ -6,41 +6,60 @@ from core.engine import run_benchmark_suite + def main(): parser = argparse.ArgumentParser(description="Run TileBench benchmarks") - parser.add_argument("--operator", type=str, default="vector_add", help="Operator to benchmark") - parser.add_argument("--output", type=str, default="results/logs/results.json", help="Output JSON file") + parser.add_argument( + "--operator", type=str, default="vector_add", help="Operator to benchmark" + ) + parser.add_argument( + "--output", + type=str, + default="results/logs/results.json", + help="Output JSON file", + ) args = parser.parse_args() - + print(f"Starting benchmark for operator: {args.operator}") results = run_benchmark_suite(args.operator) - + os.makedirs(os.path.dirname(args.output), exist_ok=True) - with open(args.output, 'w') as f: + with open(args.output, "w") as f: json.dump(results, f, indent=4) - + print(f"Results saved to {args.output}") - + # Print a summary table print("\nSummary:") rows = [] for r in results: - rows.append([ - str(r["params"]), - r["dtype"], - f"{r['torch_ms']:.4f}", - f"{r['triton_ms']:.4f}", - f"{r['cutile_ms']:.4f}", - f"{r['speedup_triton']:.2f}", - f"{r['speedup_cutile']:.2f}", - ]) + rows.append( + [ + str(r["params"]), + r["dtype"], + f"{r['torch_ms']:.4f}", + f"{r['triton_ms']:.4f}", + f"{r['cutile_ms']:.4f}", + f"{r['speedup_triton']:.2f}", + f"{r['speedup_cutile']:.2f}", + ] + ) print( tabulate( rows, - headers=["Params", "Dtype", "Torch(ms)", "Triton(ms)", "cuTile(ms)", "Speedup(T)", "Speedup(C)"], + headers=[ + "Params", + "Dtype", + "Torch(ms)", + "Triton(ms)", + "cuTile(ms)", + "Speedup(T)", + "Speedup(C)", + ], tablefmt="simple", ) ) + if __name__ == "__main__": main() diff --git a/scripts/run_bench_all.py b/scripts/run_bench_all.py index 61388e8da..68cd0f350 100644 --- a/scripts/run_bench_all.py +++ b/scripts/run_bench_all.py @@ -40,7 +40,9 @@ def create_run_dir(base_dir: Path, run_name: str = "") -> Path: return run_dir -def copy_profile_artifacts(patterns: List[str], run_dir: Path, repo_root: Path) -> List[str]: +def copy_profile_artifacts( + patterns: List[str], run_dir: Path, repo_root: Path +) -> List[str]: copied = [] profiles_dir = run_dir / "profiles" profiles_dir.mkdir(parents=True, exist_ok=True) @@ -85,7 +87,9 @@ def copy_profile_artifacts(patterns: List[str], run_dir: Path, repo_root: Path) def main() -> int: - parser = argparse.ArgumentParser(description="Run all TileBench operators and save results per run.") + parser = argparse.ArgumentParser( + description="Run all TileBench operators and save results per run." + ) parser.add_argument( "--results-root", type=str, @@ -143,7 +147,9 @@ def main() -> int: print(f"Run directory: {run_dir}") with run_log_path.open("w", encoding="utf-8") as run_log: - run_log.write(f"[{_now_utc_str()}] Starting run for {len(operators)} operators\n") + run_log.write( + f"[{_now_utc_str()}] Starting run for {len(operators)} operators\n" + ) for op in operators: print(f"=== Running {op} ===") run_log.write(f"\n[{_now_utc_str()}] START operator={op}\n") @@ -153,7 +159,9 @@ def main() -> int: with out_path.open("w", encoding="utf-8") as f: json.dump(results, f, indent=2) manifest["operators_succeeded"].append(op) - run_log.write(f"[{_now_utc_str()}] DONE operator={op} output={out_path}\n") + run_log.write( + f"[{_now_utc_str()}] DONE operator={op} output={out_path}\n" + ) except Exception: err = traceback.format_exc() manifest["operators_failed"].append({"operator": op, "error": err}) @@ -162,7 +170,9 @@ def main() -> int: run_log.write(f"[{_now_utc_str()}] FAIL operator={op}\n{err}\n") print(f" FAILED: {op}") - copied_profiles = copy_profile_artifacts(args.profile_artifact, run_dir, repo_root) + copied_profiles = copy_profile_artifacts( + args.profile_artifact, run_dir, repo_root + ) manifest["artifacts"]["profiles"] = copied_profiles manifest["finished_at_utc"] = _now_utc_str() diff --git a/scripts/visualize.py b/scripts/visualize.py index f88f526ee..4bb7ea008 100644 --- a/scripts/visualize.py +++ b/scripts/visualize.py @@ -1,37 +1,40 @@ import json + import matplotlib.pyplot as plt import numpy as np + def plot_results(json_path, output_png): - with open(json_path, 'r') as f: + with open(json_path, "r") as f: results = json.load(f) - + def format_params(params): return ", ".join([f"{k}={v}" for k, v in params.items()]) labels = [f"{format_params(r['params'])}\n{r['dtype']}" for r in results] - torch_times = [r['torch_ms'] for r in results] - triton_times = [r['triton_ms'] for r in results] - cutile_times = [r['cutile_ms'] for r in results] - + torch_times = [r["torch_ms"] for r in results] + triton_times = [r["triton_ms"] for r in results] + cutile_times = [r["cutile_ms"] for r in results] + x = np.arange(len(labels)) width = 0.25 - + fig, ax = plt.subplots(figsize=(14, 7)) - ax.bar(x - width, torch_times, width, label='PyTorch') - ax.bar(x, triton_times, width, label='Triton') - ax.bar(x + width, cutile_times, width, label='cuTile') - - ax.set_ylabel('Time (ms)') - ax.set_title('Performance Comparison: PyTorch vs Triton vs cuTile') + ax.bar(x - width, torch_times, width, label="PyTorch") + ax.bar(x, triton_times, width, label="Triton") + ax.bar(x + width, cutile_times, width, label="cuTile") + + ax.set_ylabel("Time (ms)") + ax.set_title("Performance Comparison: PyTorch vs Triton vs cuTile") ax.set_xticks(x) - ax.set_xticklabels(labels, rotation=45, ha='right') + ax.set_xticklabels(labels, rotation=45, ha="right") ax.legend() - ax.set_yscale('log') # Use log scale because of large differences - + ax.set_yscale("log") # Use log scale because of large differences + plt.tight_layout() plt.savefig(output_png) print(f"Plot saved to {output_png}") + if __name__ == "__main__": plot_results("results/logs/results.json", "results/comparison.png")