Suite of ops/shapes scraped from TorchBench (#26)

Author: bertmaher

BackendBench/backends.py

@@ -19,6 +19,21 @@ def __contains__(self, key):
         return True
 
 
+def _flag_gems_softmax(*args, **kwargs):
+    # half_to_float is not supported in flag_gems
+    import flag_gems
+
+    return flag_gems.ops.softmax(*args[:-1], **kwargs)
+
+
+def _flag_gems_layernorm(*args, **kwargs):
+    import flag_gems
+
+    x, m, v = flag_gems.ops.layer_norm(*args[:-1], **kwargs)
+    mv_shape = [*x.shape[:-1], 1]
+    return x, m.view(*mv_shape), v.view(*mv_shape)
+
+
 class FlagGemsBackend(Backend):
     def __init__(self) -> None:
         super().__init__("flaggems")
@@ -121,7 +136,7 @@ def __init__(self) -> None:
             torch.ops.aten.isnan.default: flag_gems.ops.isnan,
             torch.ops.aten.minimum.default: flag_gems.ops.minimum,
             torch.ops.aten.maximum.default: flag_gems.ops.maximum,
-            torch.ops.aten.native_layer_norm.default: flag_gems.ops.layer_norm,
+            torch.ops.aten.native_layer_norm.default: _flag_gems_layernorm,
             torch.ops.aten.native_layer_norm_backward.default: flag_gems.ops.layer_norm_backward,
             torch.ops.aten.le.Tensor: flag_gems.ops.le,
             torch.ops.aten.le.Scalar: flag_gems.ops.le_scalar,
@@ -177,7 +192,7 @@ def __init__(self) -> None:
             torch.ops.aten.silu_backward.default: flag_gems.ops.silu_backward,
             torch.ops.aten.sin.default: flag_gems.ops.sin,
             torch.ops.aten.sin_.default: flag_gems.ops.sin_,
-            torch.ops.aten._softmax.default: flag_gems.ops.softmax,
+            torch.ops.aten._softmax.default: _flag_gems_softmax,
             torch.ops.aten._softmax_backward_data.default: flag_gems.ops.softmax_backward,
             torch.ops.aten.sort.default: flag_gems.ops.sort,
             torch.ops.aten.sub.Tensor: flag_gems.ops.sub,

BackendBench/eval.py

@@ -2,29 +2,44 @@
 
 import torch
 
-from triton.testing import do_bench
+import triton.testing
 
 
 logger = logging.getLogger(__name__)
 
+EXC_MSG = """
+Exception raised for {op}:
+    args: {args}
+    kwargs: {kwargs}
+    exc: {exc}
+"""
+
+
+def format_tensor(t):
+    return f"{t.dtype}{list(t.shape)}"
+
+
+def format_args(args):
+    return [format_tensor(arg) if isinstance(arg, torch.Tensor) else arg for arg in args]
+
+
+def format_kwargs(kwargs):
+    return {k: format_tensor(v) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
+
+
+def format_exception(e, op, args, kwargs):
+    return EXC_MSG.format(op=op, args=format_args(args), kwargs=format_kwargs(kwargs), exc=e)
+
 
 def allclose(a, b):
     if isinstance(a, torch.Tensor):
-        torch.testing.assert_close(a, b, equal_nan=True)
+        torch.testing.assert_close(a, b, equal_nan=True, atol=1e-2, rtol=1e-2)
         return True
     if isinstance(a, (list, tuple)):
         return all(allclose(x, y) for x, y in zip(a, b))
     return a == b
 
 
-EXC_MSG = """
-Exception raised for {op}:
-    args: {args}
-    kwargs: {kwargs}
-    exc: {exc}
-"""
-
-
 def eval_correctness_test(op, impl, test):
     """Evaluate impl of op against test."""
     args, kwargs = test.args, test.kwargs
@@ -33,13 +48,16 @@ def eval_correctness_test(op, impl, test):
         res = impl(*args, **kwargs)
         return allclose(ref, res)
     except Exception as e:
-        logger.debug(EXC_MSG.format(op=op, args=args, kwargs=kwargs, exc=e))
+        logger.warning(format_exception(e, op, args, kwargs))
         return False
 
 
 def eval_correctness(op, impl, tests):
     correct, total = 0, 0
     for test in tests:
+        logging.debug(
+            f"Testing {op.__name__} with args {format_args(test.args)} and kwargs {format_kwargs(test.kwargs)}"
+        )
         if eval_correctness_test(op, impl, test):
             correct += 1
         total += 1
@@ -60,13 +78,20 @@ def cpu_bench(fn, num_runs=100):
 
 
 def eval_performance(op, impl, tests):
-    if torch.cuda.is_available():
-        base_times = [do_bench(lambda: op(*test.args, **test.kwargs)) for test in tests]
-        test_times = [do_bench(lambda: impl(*test.args, **test.kwargs)) for test in tests]
-    else:
-        base_times = [cpu_bench(lambda: op(*test.args, **test.kwargs)) for test in tests]
-        test_times = [cpu_bench(lambda: impl(*test.args, **test.kwargs)) for test in tests]
-
+    bench_fn = triton.testing.do_bench if torch.cuda.is_available() else cpu_bench
+    base_times = []
+    test_times = []
+    for test in tests:
+        logging.debug(
+            f"Benchmarking {op.__name__} with args {format_args(test.args)} and kwargs {format_kwargs(test.kwargs)}"
+        )
+        base_times.append(bench_fn(lambda: op(*test.args, **test.kwargs)))
+        try:
+            allclose(op(*test.args, **test.kwargs), impl(*test.args, **test.kwargs))
+        except Exception:
+            test_times.append(base_times[-1])
+            continue
+        test_times.append(bench_fn(lambda: impl(*test.args, **test.kwargs)))
     speedups = torch.tensor(test_times) / torch.tensor(base_times)
     return speedups.log().mean().exp()
 

BackendBench/torchbench_suite.py

@@ -0,0 +1,119 @@
+"""
+Load aten inputs from serialized txt files.
+"""
+
+import re
+import math
+from collections import defaultdict
+from pathlib import Path
+
+import torch
+from torch.testing import make_tensor
+
+
+dtype_abbrs = {
+    torch.bfloat16: "bf16",
+    torch.float64: "f64",
+    torch.float32: "f32",
+    torch.float16: "f16",
+    torch.complex32: "c32",
+    torch.complex64: "c64",
+    torch.complex128: "c128",
+    torch.int8: "i8",
+    torch.int16: "i16",
+    torch.int32: "i32",
+    torch.int64: "i64",
+    torch.bool: "b8",
+    torch.uint8: "u8",
+}
+
+dtype_abbrs_parsing = {value: key for key, value in dtype_abbrs.items()}
+
+_FLOATING_TYPES = [torch.float16, torch.bfloat16, torch.float32, torch.float64]
+
+
+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)
+    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))
+
+
+def _deserialize_args(inps):
+    inps = inps.strip().strip("'")
+    global_vals = {
+        "T": _deserialize_tensor,
+        "th": torch,
+        "inf": math.inf,
+        "torch": torch,
+        **dtype_abbrs_parsing,
+    }
+    # f strings introduce quotations we dont want
+    for key in dtype_abbrs_parsing:
+        inps = inps.replace(f"'{key}'", key)
+    return eval(inps.strip().strip("'").strip('"'), global_vals)
+
+
+class TorchBenchTest:
+    def __init__(self, *args, **kwargs):
+        self.args = args
+        self.kwargs = kwargs
+
+
+class TorchBenchOpTest:
+    def __init__(self, op, inputs):
+        self.op = eval(f"torch.ops.{op}")
+        self.inputs = inputs
+
+    @property
+    def correctness_tests(self):
+        for inp in self.inputs:
+            args, kwargs = _deserialize_args(inp)
+            yield TorchBenchTest(*args, **kwargs)
+
+    @property
+    def performance_tests(self):
+        for inp in self.inputs:
+            args, kwargs = _deserialize_args(inp)
+            yield TorchBenchTest(*args, **kwargs)
+
+
+def _parse_inputs(filename, filter, op_inputs):
+    op = None
+
+    with open(filename, "r") as f:
+        for line in f:
+            if m := re.match("Operator: (.*)", line):
+                op = m.group(1)
+                if op == "aten.sum.SymInt":
+                    op = "aten.sum.dim_IntList"
+            if m := re.match("cnt: \\d+, (.*)", line):
+                assert op is not None
+                args = m.group(1)
+                if filter is None or any(f in op for f in filter):
+                    op_inputs[op].append(args)
+    return op_inputs
+
+
+class TorchBenchTestSuite:
+    def __init__(self, name, filename, filter=None):
+        self.name = name
+        self.optests = defaultdict(list)
+        if Path(filename).is_dir():
+            for file_path in Path(filename).glob("**/*.txt"):
+                _parse_inputs(str(file_path), filter, self.optests)
+        else:
+            _parse_inputs(filename, filter, self.optests)
+        # Deduplicate the strings in self.optests
+        for op in self.optests:
+            self.optests[op] = list(set(self.optests[op]))
+
+    def __iter__(self):
+        for op, inputs in self.optests.items():
+            if any(s in op for s in ["embedding", "scatter", "gather", "index", "nll_loss"]):
+                # TODO: indexing ops need valid indices
+                continue
+            yield TorchBenchOpTest(op, inputs)

scripts/main.py

@@ -8,17 +8,37 @@
 import click
 import torch
 from BackendBench.opinfo_suite import OpInfoTestSuite
+from BackendBench.torchbench_suite import TorchBenchTestSuite
 from BackendBench.suite import SmokeTestSuite
 from BackendBench.llm_client import ClaudeKernelGenerator
 
 logger = logging.getLogger(__name__)
 
 
+def setup_logging(log_level):
+    """Configure logging with the specified level."""
+    numeric_level = getattr(logging, log_level.upper(), None)
+    if not isinstance(numeric_level, int):
+        raise ValueError(f"Invalid log level: {log_level}")
+
+    logging.basicConfig(
+        level=numeric_level,
+        format="[%(asctime)s][%(levelname)s][%(filename)s] %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+
 @click.command()
+@click.option(
+    "--log-level",
+    default=os.getenv("LOG_LEVEL", "INFO"),
+    type=click.Choice(["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], case_sensitive=False),
+    help="Set the logging level",
+)
 @click.option(
     "--suite",
     default="smoke",
-    type=click.Choice(["smoke", "opinfo"]),
+    type=click.Choice(["smoke", "opinfo", "torchbench"]),
     help="Which suite to run",
 )
 @click.option(
@@ -39,7 +59,14 @@
     type=int,
     help="Maximum attempts for LLM kernel generation with feedback",
 )
-def cli(suite, backend, ops, llm_max_attempts):
+@click.option(
+    "--torchbench-data-path",
+    default="third_party/tritonbench/tritonbench/data/input_configs",
+    type=str,
+    help="Path to TorchBench operator data",
+)
+def cli(log_level, suite, backend, ops, llm_max_attempts, torchbench_data_path):
+    setup_logging(log_level)
     if ops:
         ops = ops.split(",")
 
@@ -62,6 +89,11 @@ def cli(suite, backend, ops, llm_max_attempts):
             torch.bfloat16,
             filter=ops,
         ),
+        "torchbench": lambda: TorchBenchTestSuite(
+            "torchbench",
+            torchbench_data_path,
+            filter=ops,
+        ),
     }[suite]()
 
     overall_correctness = []