Enable testing preconfigured internal shapes

Differential Revision: D85905941

Pull Request resolved: #610

Author: njriasan

tritonbench/operators/addmm/data_io.py

@@ -11,5 +11,6 @@ def parse_args(args: List[str]) -> argparse.Namespace:
     parser.add_argument("--col-major", type=bool, default=False)
     parser.add_argument("--large-k-shapes", type=bool, default=False)
     parser.add_argument("--bias-1D-y", type=bool, default=False)
+    parser.add_argument("--config", type=str, default=None)
     args = parser.parse_args(args)
     return args

tritonbench/operators/addmm/operator.py

@@ -5,6 +5,7 @@
 import torch
 import torch._inductor.config as inductor_config
 import triton
+from tritonbench.utils.env_utils import is_fbcode
 
 try:
     from hammer.ops.triton.triton_hstu_linear import triton_addmm
@@ -26,6 +27,11 @@
 
 from .data_io import parse_args
 
+if is_fbcode():
+    from tritonbench.utils.fb.addmm_prod import get_prod_shapes
+else:
+    get_prod_shapes = lambda x: None
+
 
 # Shape encoding information: (M, K, N, BIAS_1D_Y)
 BUILDIN_SHAPES = [
@@ -88,7 +94,10 @@ def __init__(
     ):
         super().__init__(tb_args, extra_args)
         addmm_args = parse_args(self.extra_args)
-        if addmm_args.m and addmm_args.n and addmm_args.k and addmm_args.bias_1D_y:
+        prod_shapes = get_prod_shapes(addmm_args.config)
+        if prod_shapes:
+            self.shapes = prod_shapes
+        elif addmm_args.m and addmm_args.n and addmm_args.k and addmm_args.bias_1D_y:
             self.shapes = [
                 (addmm_args.m, addmm_args.k, addmm_args.n, addmm_args.bias_1D_y)
             ]