[TRITON-BENCHMARKS] Add shape filter support (#4147)

Author: vlad-penkin

benchmarks/setup.py

@@ -122,11 +122,17 @@ def get_git_commit_hash(length=8):
         "matplotlib",
     ],
     package_dir={"triton_kernels_benchmark": "triton_kernels_benchmark"},
-    package_data={"triton_kernels_benchmark": ["xetla_kernel.cpython-*.so"]},
+    package_data={"triton_kernels_benchmark": [
+        "xetla_kernel.cpython-*.so",
+        "cutlass_kernel.cpython-*.so",
+    ]},
     cmdclass={
         "build_ext": build_ext,
     },
-    ext_modules=[CMakeExtension("triton_kernels_benchmark.xetla_kernel")],
+    ext_modules=[
+        CMakeExtension("triton_kernels_benchmark.xetla_kernel"),
+        CMakeExtension("triton_kernels_benchmark.cutlass_kernel"),
+    ],
     entry_points={
         "console_scripts": [
             "triton-benchmarks = triton_kernels_benchmark.benchmark_utils:main",

benchmarks/tests/test_entry_point.py

@@ -3,11 +3,10 @@
 import pytest
 
 from triton_kernels_benchmark.benchmark_testing import MarkArgs
-from triton_kernels_benchmark.becnhmark_config_templates import CONFIGS
 from triton_kernels_benchmark.benchmark_utils import BenchmarkCategory, BenchmarkConfigs
 
 ALL_CATEGORIES = {cat.value for cat in BenchmarkCategory}
-ALL_CONFIGS = {config.key: config for config in CONFIGS}
+ALL_CONFIGS = set(BenchmarkConfigs._get_all_configs().keys())  # pylint: disable=W0212
 
 
 @pytest.mark.parametrize(
@@ -22,7 +21,7 @@
         "providers_count",
     ),
     (
-        [True, ALL_CONFIGS, True, ALL_CATEGORIES, [], None, lambda x: x > 1, lambda x: x > 1],
+        [True, ALL_CONFIGS, True, ALL_CATEGORIES, [], None, lambda x: x > 1, lambda x: x >= 1],
         [True, {"softmax", "gemm"}, True, ALL_CATEGORIES, [], None, lambda x: x > 1, lambda x: x > 1],
         [True, {"softmax", "gemm"}, True, {"core", "gemm", "softmax"}, [], None, lambda x: x > 1, lambda x: x > 1],
         [False, {"softmax"}, False, {"optional"}, ["triton"], AssertionError, None, None],
@@ -49,6 +48,7 @@ def benchmark_configs():
             select_all=select_all,
             categories_filter=categories_filter,
             providers_filter=providers_filter,
+            shape_pattern=None,
             json_output=False,
             detailed_output=False,
             tag="",

benchmarks/tests/test_mocks.py

@@ -1,7 +1,6 @@
-from typing import Optional
+from typing import List, Optional
 
 import re
-
 import io
 
 import pytest
@@ -28,34 +27,52 @@
 }
 
 
+@pytest.fixture(autouse=True)
+def capture_hw_details_env(monkeypatch):
+    monkeypatch.setenv("GPU_DEVICE", "Intel(R) Data Center GPU Max 1100")
+    yield
+
+
 @pytest.mark.parametrize("command", ["run"])
-@pytest.mark.parametrize("benchmark", ["softmax"])
+@pytest.mark.parametrize(
+    "benchmark",
+    [
+        ["softmax"],
+        ["softmax", "--shape-pattern", "[256]"],
+        ["softmax", "--shape-pattern", "[*]"],
+    ],
+)
 @pytest.mark.parametrize("provider", ["triton", None])
-@pytest.mark.parametrize("n_runs", [None, 1, 2])
+@pytest.mark.parametrize("n_runs", [1, 2, 3])
 @pytest.mark.parametrize("show_details", [False, True])
 @pytest.mark.parametrize("json_output", [False, True])
+@pytest.mark.parametrize("reports", [False, True])
 def test_benchmark_run_monkeypatched(
     command: str,
-    benchmark: str,
+    benchmark: List[str],
     provider: Optional[str],
-    n_runs: Optional[int],
+    n_runs: int,
     show_details: bool,
     json_output: bool,
+    reports: bool,
     capsys,
+    tmp_path,
 ):
-    args = [command, benchmark]
+    args = [command] + benchmark
     if provider:
         args.extend(["--provider", provider])
-    if n_runs:
+    if n_runs > 1:
         args.extend(["--n_runs", str(n_runs)])
     if show_details:
         args.extend(["--show-details"])
     if json_output:
         args.extend(["--json"])
+    if reports:
+        args.extend(["--reports", str(tmp_path)])
 
     configs = BenchmarkConfigs.from_args(args)
     for config in configs.configs:
-        config.res_df = pd.read_csv(io.StringIO(PERFORMANCE_CSVS[config.key]))
+        config.res_df_list = [pd.read_csv(io.StringIO(PERFORMANCE_CSVS[config.key]))] * n_runs
     configs.run()
 
     captured_output = capsys.readouterr().out

benchmarks/tests/test_shape_pattern_parser.py

@@ -0,0 +1,140 @@
+import pytest
+
+from triton_kernels_benchmark.benchmark_utils import BenchmarkCategory, BenchmarkConfigs
+from triton_kernels_benchmark.benchmark_shapes_parser import ShapePatternParser
+
+CONFIG_SHAPES = [
+    shape for config in BenchmarkConfigs._get_configs(  # pylint: disable=W0212
+        configs_filter=BenchmarkConfigs._get_all_configs().keys(),  # pylint: disable=W0212
+        categories_filter=[cat.value for cat in BenchmarkCategory],
+        providers_filter=[],
+        shape_pattern=None,
+    ) for shape in config.shapes
+]
+
+
+@pytest.mark.parametrize(
+    "input_string, expected",
+    [
+        ("[1-2-3]", [1, 2, 3]),
+        ("[16-32-1024-64]", [16, 32, 1024, 64]),
+        ("[True-False-123]", ["True", "False", 123]),
+    ],
+)
+def test_parse_valid_shapes(input_string, expected):
+    assert ShapePatternParser.parse(input_string) == expected
+
+
+@pytest.mark.parametrize("known_shape", CONFIG_SHAPES)
+def test_parse_all_known_shapes(known_shape):
+    ShapePatternParser.parse(known_shape)
+
+
+@pytest.mark.parametrize(
+    "pattern, tokens",
+    [
+        ("[16-*-1024-*]", [16, "*", 1024, "*"]),
+    ],
+)
+def test_parse_valid_pattern_shapes(pattern, tokens):
+    parsed_tokens = ShapePatternParser.parse(pattern, pattern_shape=True)
+    assert tokens == parsed_tokens
+
+
+def test_parse_bracketed_star_without_pattern_shape_fails():
+    with pytest.raises(ValueError):
+        ShapePatternParser.parse("[*]")
+
+
+@pytest.mark.parametrize(
+    "invalid_string",
+    [
+        "1-2-3",
+        "[]",
+        "[ ]",
+        "[ - ]",
+        "[-]",
+        "[--]",
+        "[- -]",
+        "[1-2-@-4]",
+    ],
+)
+def test_parse_bracketed_invalid(invalid_string):
+    with pytest.raises(ValueError):
+        ShapePatternParser.parse(invalid_string, pattern_shape=True)
+    with pytest.raises(ValueError):
+        ShapePatternParser(invalid_string)
+
+
+@pytest.mark.parametrize(
+    "pattern",
+    [
+        "[1-2-3]",
+        "[a-b-c]",
+        "[*-b-3]",
+    ],
+)
+def test_init_valid(pattern):
+    parser = ShapePatternParser(pattern)
+    assert parser.pattern == pattern
+    assert len(parser.pattern_tokens) == parser.pattern_dims
+
+
+@pytest.fixture
+def sample_shapes():
+    return [
+        "[16-32-1024-64-False-bwd]",
+        "[16-64-1024-32-True-fwd]",
+        "[32-32-512-64-False-bwd]",
+        "[16-32-1024-64-False-fwd]",
+    ]
+
+
+def test_filter_exact_match(sample_shapes):  # pylint: disable=W0621
+    parser = ShapePatternParser("[16-32-1024-64-False-bwd]")
+    assert parser.filter_by_pattern(sample_shapes) == ["[16-32-1024-64-False-bwd]"]
+
+
+def test_filter_wildcard_match(sample_shapes):  # pylint: disable=W0621
+    parser = ShapePatternParser("[16-*-1024-64-False-*]")
+    expected = [
+        "[16-32-1024-64-False-bwd]",
+        "[16-32-1024-64-False-fwd]",
+    ]
+    assert parser.filter_by_pattern(sample_shapes) == expected
+
+
+@pytest.mark.parametrize(
+    "pattern, shapes",
+    [
+        ("[*]", ["bad-shape", "[foo]", "[123]"]),
+    ],
+)
+def test_filter_invalid_shape_strings(pattern, shapes):
+    parser = ShapePatternParser(pattern)
+    with pytest.raises(ValueError):
+        parser.filter_by_pattern(shapes)
+
+
+@pytest.mark.parametrize(
+    "pattern, shapes",
+    [
+        ("[1-2-3]", ["[1-2]"]),
+    ],
+)
+def test_filter_dims_mismatch(pattern, shapes):
+    parser = ShapePatternParser(pattern)
+    with pytest.raises(ValueError) as exc_info:
+        parser.filter_by_pattern(shapes)
+    assert "mismatch" in str(exc_info.value)
+
+
+@pytest.mark.parametrize(
+    "pattern, shapes",
+    [
+        ("[*]", ["[foo]", "[bar]", "[123]"]),
+    ],
+)
+def test_filter_single_wildcard_matches_all(pattern, shapes):
+    parser = ShapePatternParser(pattern)
+    assert parser.filter_by_pattern(shapes) == shapes

benchmarks/triton_kernels_benchmark/__init__.py

@@ -11,6 +11,8 @@
     BENCHMARKING_METHOD,
 )
 
+from .benchmark_shapes_parser import ShapePatternParser
+
 if BENCHMARKING_METHOD == "UPSTREAM_PYTORCH_PROFILER":
     os.environ["INJECT_PYTORCH"] = "True"
 
@@ -23,4 +25,5 @@
     "BenchmarkCategory",
     "BenchmarkConfig",
     "BENCHMARKING_METHOD",
+    "ShapePatternParser",
 ]

benchmarks/triton_kernels_benchmark/benchmark_shapes_parser.py

@@ -0,0 +1,63 @@
+from typing import List, Union
+from dataclasses import dataclass, field
+
+import re
+
+
+@dataclass
+class ShapePatternParser:
+    pattern: str
+    pattern_tokens: List[str] = field(init=False)
+    pattern_dims: int = field(init=False)
+
+    def __post_init__(self):
+        self.pattern_tokens = self.parse(self.pattern, pattern_shape=True)
+        self.pattern_dims = len(self.pattern_tokens)
+
+    def __str__(self):
+        return self.pattern
+
+    @staticmethod
+    def parse(shape_string: str, pattern_shape: bool = False) -> List[Union[int, str]]:
+        pattern_match = re.fullmatch(r"\[(.*)\]", shape_string)
+        if not pattern_match:
+            raise ValueError(
+                f"Invalid format: {shape_string!r}, only patterns similar to [16-*-1024-*-bwd] are supported", )
+        inner_string = pattern_match.group(1)
+        if not inner_string:
+            raise ValueError(f"Empty shape - {inner_string}")
+        tokens = inner_string.split("-")
+        if not any(tokens):
+            raise ValueError(f"Empty shape - {inner_string}")
+        result: List[Union[int, str]] = []
+        for token in tokens:
+            try:
+                result.append(int(token))
+            except ValueError:
+                if token.isalnum() or token == "*" and pattern_shape:
+                    result.append(token)
+                else:
+                    raise ValueError(  # pylint: disable=W0707
+                        f"Unsupported shape or shape pattern {shape_string}"
+                        "Each shape element could be either int, alphanumeric string or '*' in shape pattern")
+        return result
+
+    def matches_pattern(self, shape_string: str) -> bool:
+        tokens = self.parse(shape_string)
+        shape_dims = len(tokens)
+        if shape_dims != self.pattern_dims:
+            raise ValueError(f"Input shape dims {shape_dims} and pattern shape dims {self.pattern_dims} mismatch")
+        for pattern_token, token in zip(self.pattern_tokens, tokens):
+            if pattern_token == "*":
+                continue
+            if pattern_token != token:
+                break
+        else:
+            return True
+        return False
+
+    def __call__(self, shape_string: str) -> bool:
+        return self.matches_pattern(shape_string)
+
+    def filter_by_pattern(self, shape_strings: List[str]) -> List[str]:
+        return [shape_string for shape_string in shape_strings if self.matches_pattern(shape_string)]