Add --topn option to select N largest shapes for each operator (#29)

Author: bertmaher

BackendBench/torchbench_suite.py

@@ -2,8 +2,8 @@
 Load aten inputs from serialized txt files.
 """
 
-import re
 import math
+import re
 from collections import defaultdict
 from pathlib import Path
 
@@ -33,13 +33,14 @@
 
 
 def _deserialize_tensor(size, dtype, stride=None, device="cuda"):
-    if stride is not None:
-        out = torch.empty_strided(size, stride, dtype=dtype, device=device)
-    else:
-        out = torch.empty(size, dtype=dtype, device=device)
+    kwargs = {}
     if dtype in _FLOATING_TYPES:
-        return out.copy_(make_tensor(size, dtype=dtype, device=device, low=0, high=1))
-    return out.copy_(make_tensor(size, dtype=dtype, device=device))
+        kwargs.update({"low": 0, "high": 1})
+    if stride is not None:
+        extent = 1 + sum((size - 1) * stride for size, stride in zip(size, stride))
+        data = make_tensor(extent, dtype=dtype, device=device, **kwargs)
+        return data.as_strided(size, stride)
+    return make_tensor(size, dtype=dtype, device=device, **kwargs)
 
 
 def _deserialize_args(inps):
@@ -63,20 +64,40 @@ def __init__(self, *args, **kwargs):
         self.kwargs = kwargs
 
 
+def _args_size(args):
+    size = 0
+    for arg in args:
+        if isinstance(arg, torch.Tensor):
+            size += arg.numel() * arg.element_size()
+        elif isinstance(arg, (tuple, list)):
+            size += _args_size(arg)
+    return size
+
+
 class TorchBenchOpTest:
-    def __init__(self, op, inputs):
+    def __init__(self, op, inputs, topn):
         self.op = eval(f"torch.ops.{op}")
         self.inputs = inputs
+        self.topn = topn
+
+    def tests(self):
+        inputs_and_sizes = []
+        for inp in self.inputs:
+            args, kwargs = _deserialize_args(inp)
+            size = _args_size(args) + _args_size(list(kwargs.values()))
+            inputs_and_sizes.append((size, inp))
+        ret = [x[1] for x in sorted(inputs_and_sizes, reverse=True)]
+        return ret if self.topn is None else ret[: self.topn]
 
     @property
     def correctness_tests(self):
-        for inp in self.inputs:
+        for inp in self.tests():
             args, kwargs = _deserialize_args(inp)
             yield TorchBenchTest(*args, **kwargs)
 
     @property
     def performance_tests(self):
-        for inp in self.inputs:
+        for inp in self.tests():
             args, kwargs = _deserialize_args(inp)
             yield TorchBenchTest(*args, **kwargs)
 
@@ -99,8 +120,9 @@ def _parse_inputs(filename, filter, op_inputs):
 
 
 class TorchBenchTestSuite:
-    def __init__(self, name, filename, filter=None):
+    def __init__(self, name, filename, filter=None, topn=None):
         self.name = name
+        self.topn = topn
         self.optests = defaultdict(list)
         if Path(filename).is_dir():
             for file_path in Path(filename).glob("**/*.txt"):
@@ -113,7 +135,21 @@ def __init__(self, name, filename, filter=None):
 
     def __iter__(self):
         for op, inputs in self.optests.items():
-            if any(s in op for s in ["embedding", "scatter", "gather", "index", "nll_loss"]):
+            if any(
+                s in op
+                for s in [
+                    "embedding",
+                    "scatter",
+                    "gather",
+                    "index",
+                    "nll_loss",
+                    "im2col_backward",
+                    "col2im_backward",
+                    "native_layer_norm_backward",
+                    "upsample_nearest2d_backward.vec",
+                    "upsample_bilinear2d_backward.vec",
+                ]
+            ):
                 # TODO: indexing ops need valid indices
                 continue
-            yield TorchBenchOpTest(op, inputs)
+            yield TorchBenchOpTest(op, inputs, self.topn)

scripts/main.py

@@ -7,10 +7,10 @@
 import BackendBench.eval as eval
 import click
 import torch
+from BackendBench.llm_client import ClaudeKernelGenerator
 from BackendBench.opinfo_suite import OpInfoTestSuite
-from BackendBench.torchbench_suite import TorchBenchTestSuite
 from BackendBench.suite import SmokeTestSuite
-from BackendBench.llm_client import ClaudeKernelGenerator
+from BackendBench.torchbench_suite import TorchBenchTestSuite
 
 logger = logging.getLogger(__name__)
 
@@ -53,6 +53,13 @@ def setup_logging(log_level):
     type=str,
     help="Comma-separated list of ops to run",
 )
+@click.option(
+    "--topn-inputs",
+    "--topn",
+    default=None,
+    type=int,
+    help="Select the top N largest inputs for each op (default: all inputs)",
+)
 @click.option(
     "--llm-max-attempts",
     default=5,
@@ -82,6 +89,7 @@ def cli(
     suite,
     backend,
     ops,
+    topn_inputs,
     llm_max_attempts,
     kernel_agent_workers,
     kernel_agent_max_rounds,
@@ -122,6 +130,7 @@ def cli(
             "torchbench",
             torchbench_data_path,
             filter=ops,
+            topn=topn_inputs,
         ),
     }[suite]()
 

test/test_torchbench_suite.py

@@ -0,0 +1,38 @@
+import torch
+from BackendBench.torchbench_suite import TorchBenchOpTest
+
+
+class TestOpTest:
+    def test_op_test(self):
+        op_test = TorchBenchOpTest(
+            "aten.relu.default", ["((T([32, 128, 512], f16, None, 'cpu'),), {})"], None
+        )
+        for test in op_test.correctness_tests:
+            args, kwargs = test.args, test.kwargs
+            arg, *extras = args
+            assert arg.shape == torch.Size([32, 128, 512])
+            assert arg.dtype == torch.float16
+            assert kwargs == {}
+            assert extras == []
+
+            torch.testing.assert_close(torch.relu(arg), op_test.op(arg))
+
+    def test_topn(self):
+        op_test = TorchBenchOpTest(
+            "aten.relu.default",
+            [
+                "((T([32, 128, 512], f16, None, 'cpu'),), {})",
+                "((T([32, 256, 512], f16, None, 'cpu'),), {})",
+            ],
+            1,
+        )
+        assert len(op_test.tests()) == 1
+        for test in op_test.correctness_tests:
+            args, kwargs = test.args, test.kwargs
+            arg, *extras = args
+            assert arg.shape == torch.Size([32, 256, 512])
+            assert arg.dtype == torch.float16
+            assert kwargs == {}
+            assert extras == []
+
+            torch.testing.assert_close(torch.relu(arg), op_test.op(arg))