fix TB matmul integration

Author: yf225

benchmarks/run.py

@@ -85,8 +85,8 @@
     "gemm": (
         "tritonbench.operators.gemm.operator",
         [
-            ("examples.matmul", "matmul"),
-            ("examples.matmul_split_k", "matmul_split_k"),
+            ("examples.matmul", "matmul_tritonbench"),
+            ("examples.matmul_split_k", "matmul_split_k_tritonbench"),
         ],
     ),
 }

examples/matmul.py

@@ -50,6 +50,14 @@ def check(m: int, k: int, n: int) -> None:
     run_example(kernel_with_bias, expected_with_bias, (x, y))
 
 
+def matmul_tritonbench(a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
+    """Wrapper for tritonbench compatibility that handles bias."""
+    if bias is not None:
+        return matmul(a, b, lambda acc, tile: acc + bias[tile[1]])
+    else:
+        return matmul(a, b)
+
+
 def main() -> None:
     check(1024, 1024, 1024)
 

examples/matmul_split_k.py

@@ -58,6 +58,14 @@ def check(m: int, k: int, n: int) -> None:
     run_example(kernel_with_bias, expected_with_bias, (x, y), atol=1)
 
 
+def matmul_split_k_tritonbench(a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
+    """Wrapper for tritonbench compatibility that handles bias."""
+    if bias is not None:
+        return matmul_split_k(a, b, lambda acc, tile: acc + bias[tile[1]])
+    else:
+        return matmul_split_k(a, b)
+
+
 def main() -> None:
     check(64, 32768, 64)