Allow omitting the "Bias" column in data table (#2412)

Summary: Pull Request resolved: #2412 Reviewed By: xuzhao9, sijiac Differential Revision: D60964532 Pulled By: bertmaher fbshipit-source-id: a6e73173478026179c1be51c73676b7e463aa028
pytorch · Aug 9, 2024 · cc5eaf7 · cc5eaf7
1 parent 59a3634
commit cc5eaf7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/torchbenchmark/operators/gemm/data_io.py b/torchbenchmark/operators/gemm/data_io.py
@@ -28,7 +28,7 @@ def read_shapes_from_csv(csv_path: str) -> List[List[int]]:
         reader = csv.DictReader(f)
         for row in reader:
             shape = [
-                int(row[f]) if row[f] else None
+                int(row.get(f)) if row.get(f) else None
                 for f in ("M", "N", "K", "Bias")
             ]
             shapes.append(shape)

diff --git a/torchbenchmark/operators/gemm/operator.py b/torchbenchmark/operators/gemm/operator.py
@@ -140,7 +140,7 @@ def triton_tma_persistent_cached_matmul(self, a, b, bias) -> Callable:
     def triton_ops_matmul(self, a, b, bias) -> Callable:
         if bias is None:
             return lambda: kernels.matmul(a, b)
-        return lambda: kernels.matmul(a, b, bias)
+        return lambda: kernels.matmul(a, b) + bias
 
     @register_benchmark(baseline=True)
     def aten_matmul(self, a, b, bias) -> Callable: