add large-k shapes to addmm

Summary: Add shapes from the diff linked in T186231930

Reviewed By: xuzhao9

Differential Revision: D60489313

fbshipit-source-id: c49483dc164affc0c383966d164fd4c9a5239807

torchbenchmark/operators/addmm/data_io.py

@@ -13,5 +13,6 @@ def parse_args(args: List[str]) -> argparse.Namespace:
     parser.add_argument("--n", type=int)
     parser.add_argument("--input", type=str)
     parser.add_argument("--col-major", type=bool, default=False)
+    parser.add_argument("--large-k-shapes", type=bool, default=False)
     args = parser.parse_args(args)
     return args

torchbenchmark/operators/addmm/operator.py

@@ -1,5 +1,6 @@
-import os
 import argparse
+import itertools
+import os
 from typing import Any, Callable, Generator, List, Optional, Tuple
 
 import numpy
@@ -64,6 +65,8 @@
     (20068, 1536, 512),
 ]
 
+# M=13, K=2^6..2^25, N=2
+LARGE_K_SHAPES = list(itertools.product([13], [2**i for i in range(6, 26)], [2]))
 
 class Operator(BenchmarkOperator):
     DEFAULT_METRICS = ["tflops", "best_config"]
@@ -74,6 +77,8 @@ def __init__(self, tb_args: argparse.Namespace, extra_args: Optional[List[str]]
         addmm_args = parse_args(self.extra_args)
         if addmm_args.m and addmm_args.n and addmm_args.k:
             self.shapes = [(addmm_args.m, addmm_args.k, addmm_args.n)]
+        elif addmm_args.large_k_shapes:
+            self.shapes = LARGE_K_SHAPES
         else:
             self.shapes = BUILDIN_SHAPES
         self.col_major = addmm_args.col_major