Col-major for mat2

Summary: as the title Reviewed By: xuzhao9 Differential Revision: D58097785 fbshipit-source-id: 056e78bf04388416302ea0e6835d353fb3ff469e
pytorch · Jun 7, 2024 · 9cd6e86 · 9cd6e86
1 parent 1a986c5
commit 9cd6e86
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/torchbenchmark/operators/addmm/data_io.py b/torchbenchmark/operators/addmm/data_io.py
@@ -12,5 +12,6 @@ def parse_args(args: List[str]) -> argparse.Namespace:
  parser.add_argument("--k", type=int)
  parser.add_argument("--n", type=int)
  parser.add_argument("--input", type=str)
+ parser.add_argument("--col-major", type=bool, default=False)
  args = parser.parse_args(args)
  return args
diff --git a/torchbenchmark/operators/addmm/operator.py b/torchbenchmark/operators/addmm/operator.py
@@ -77,6 +77,7 @@ def __init__(self, mode: str, device: str, extra_args: Optional[List[str]] = Non
  self.shapes = [(addmm_args.m, addmm_args.k, addmm_args.n)]
  else:
  self.shapes = BUILDIN_SHAPES
+ self.col_major = addmm_args.col_major
 
  @register_benchmark()
  def triton_addmm(self, a, mat1, mat2) -> Callable:
@@ -140,6 +141,8 @@ def get_input_iter(self) -> Generator:
  mat2 = torch.randn(
  (k, n), device=self.device, dtype=self.dtype
  ).requires_grad_(False)
+ if self.col_major:
+ mat2 = mat2.T.contiguous().T
  yield a, mat1, mat2
 
  def _get_accuracy(self, fn: Callable, baseline_fn: Callable) -> bool: