Add some split-k shapes to the benchmark harness

Summary: Available via the arg `--splitk` although I'm not sure if that's the best way to do things...

Reviewed By: xuzhao9, chenyang78

Differential Revision: D56197655

fbshipit-source-id: d30303b7402faa057722197481e185580199aba0

Author: bertmaher

torchbenchmark/operators/gemm/data_io.py

@@ -12,6 +12,7 @@ def parse_args(args: List[str]) -> argparse.Namespace:
     parser.add_argument("--k", default=8, type=int)
     parser.add_argument("--n", default=8, type=int)
     parser.add_argument("--input", default=None, type=str)
+    parser.add_argument("--splitk", action="store_true", default=False)
     args = parser.parse_args(args)
     return args
 

torchbenchmark/operators/gemm/operator.py

@@ -67,6 +67,11 @@
     (4096, 4096, 4096, None),
 ]
 
+SPLIT_K_SHAPES = [
+    (m, k, m, None)
+    for m in [128 * i for i in range(1, 5)]
+    for k in [2048 * i for i in range(1, 9)]
+]
 
 class Operator(BenchmarkOperator):
     DEFAULT_METRICS = ["latency", "speedup", "accuracy"]
@@ -81,6 +86,8 @@ def __init__(self, mode: str, device: str, extra_args: List[str] = []):
             self.tbargs = parse_args(self.extra_args)
             if self.tbargs.input:
                 self.shapes = read_shapes_from_csv(self.tbargs.input)
+            elif self.tbargs.splitk:
+                self.shapes = SPLIT_K_SHAPES
             else:
                 self.shapes = [(self.tb_args.m, self.tbargs.k, self.tbargs.n)]
             self.DEFAULT_NUM_BATCH = len(self.shapes)