Benchmark fixes.

Author: maleadt

examples/batchmatmul.jl

@@ -74,7 +74,7 @@ function prepare(; benchmark::Bool=false,
     )
 end
 
-function run(data; tm::Int=64, tn::Int=64, tk::Int=64, nruns::Int=1, warmup::Int=0)
+function run(data; tm::Int=128, tn::Int=128, tk::Int=64, nruns::Int=1, warmup::Int=0)
     (; A, B, C, M, N, Batch) = data
     grid = (cld(M, tm), cld(N, tn), Batch)
 

examples/batchmatmul.py

@@ -29,6 +29,11 @@ def batchmatmul_cutile_kernel(A, B, C, tm: ct.Constant[int], tn: ct.Constant[int
         b = ct.load(B, index=(pid_batch, k, bidy), shape=(1, tk, tn), padding_mode=zero_pad)
         b = ct.reshape(b, (tk, tn))
 
+        # Convert to TF32 for tensor cores (Float32 inputs only)
+        if A.dtype == ct.float32:
+            a = ct.astype(a, ct.tfloat32)
+            b = ct.astype(b, ct.tfloat32)
+
         accumulator = ct.mma(a, b, acc=accumulator)
 
     result = ct.astype(accumulator, C.dtype)
@@ -61,7 +66,7 @@ def prepare(*, benchmark: bool = False, Batch: int = None, M: int = None, K: int
     }
 
 
-def run(data, *, tm: int = 64, tn: int = 64, tk: int = 64, nruns: int = 1, warmup: int = 0):
+def run(data, *, tm: int = 128, tn: int = 128, tk: int = 64, nruns: int = 1, warmup: int = 0):
     """Run batch matmul kernel with timing."""
     A, B, C = data["A"], data["B"], data["C"]
     Batch, M, N = data["Batch"], data["M"], data["N"]

examples/transpose.py

@@ -12,7 +12,7 @@ def transpose_cutile_kernel(input, output, tile_m: ct.Constant[int], tile_n: ct.
     pid_m = ct.bid(0)
     pid_n = ct.bid(1)
     tile = ct.load(input, index=(pid_m, pid_n), shape=(tile_m, tile_n))
-    tile_t = transpose(tile)
+    tile_t = ct.transpose(tile)
     ct.store(output, index=(pid_n, pid_m), tile=tile_t)