(fix) apply formatting fixes for pre-commit compliance

.gitignore

@@ -1,4 +1,4 @@
 *.pyc
 *.png
 *.json
-results/
+results/

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

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]:

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)

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)

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

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

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))

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

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)

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)
+    )

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,14 +18,23 @@ 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()
     gate = gate.contiguous()
     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

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

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]:

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]:

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]:

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
+    ct.launch(
+        torch.cuda.current_stream(),
+        grid,
+        kernel_function,
+        (x, output, n_elements, TILE),
+    )
+    return output

benchmarks/operators/sin/impl_torch.py

@@ -1,4 +1,5 @@
 import torch
 
+
 def run(x):
-    return torch.sin(x)
+    return torch.sin(x)

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
+    return output