Split benchmark for concat and stack (#2671)

Author: tengqm

benchmark/test_cat.py

@@ -0,0 +1,45 @@
+from typing import Generator
+
+import pytest
+import torch
+
+from benchmark.attri_util import FLOAT_DTYPES, INT_DTYPES, BenchLevel
+from benchmark.performance_utils import Benchmark, Config, generate_tensor_input
+
+
+def _input_fn(shape, dtype, device):
+    inp1 = generate_tensor_input(shape, dtype, device)
+    inp2 = generate_tensor_input(shape, dtype, device)
+    inp3 = generate_tensor_input(shape, dtype, device)
+    yield [inp1, inp2, inp3], {"dim": 0},
+
+    if Config.bench_level == BenchLevel.COMPREHENSIVE:
+        yield [inp1, inp2, inp3], {"dim": -1},
+
+
+class CatBenchmark(Benchmark):
+    def __init__(self, *args, **kwargs):
+        self.input_fn = kwargs.pop("input_fn", _input_fn)
+        super().__init__(*args, **kwargs)
+
+    def get_input_iter(self, dtype) -> Generator:
+        for shape in self.shapes:
+            yield from self.input_fn(shape, dtype, self.device)
+
+    def set_more_shapes(self):
+        more_shapes_2d = [[1024, 2**i] for i in range(1, 11, 4)]
+        more_shapes_3d = [[64, 64, 2**i] for i in range(0, 8, 4)]
+
+        return more_shapes_2d + more_shapes_3d
+
+
+@pytest.mark.skip("Benchmark test fails: issue #2673")
+@pytest.mark.cat
+def test_cat():
+    bench = CatBenchmark(
+        op_name="cat",
+        input_fn=_input_fn,
+        torch_op=torch.cat,
+        dtypes=FLOAT_DTYPES + INT_DTYPES,
+    )
+    bench.run()

benchmark/test_hstack.py

@@ -0,0 +1,40 @@
+from typing import Generator
+
+import pytest
+import torch
+
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.performance_utils import Benchmark, generate_tensor_input
+
+
+def _input_fn(shape, dtype, device):
+    inp1 = generate_tensor_input(shape, dtype, device)
+    inp2 = generate_tensor_input(shape, dtype, device)
+    inp3 = generate_tensor_input(shape, dtype, device)
+
+    yield [inp1, inp2, inp3],
+
+
+class HStackBenchmark(Benchmark):
+    def __init__(self, *args, **kwargs):
+        self.input_fn = kwargs.pop("input_fn", _input_fn)
+        super().__init__(*args, **kwargs)
+
+    def get_input_iter(self, dtype) -> Generator:
+        for shape in self.shapes:
+            yield from self.input_fn(shape, dtype, self.device)
+
+    def set_more_shapes(self):
+        more_shapes_2d = [[1024, 2**i] for i in range(1, 11, 4)]
+        more_shapes_3d = [[64, 64, 2**i] for i in range(0, 8, 4)]
+
+        return more_shapes_2d + more_shapes_3d
+
+
+@pytest.mark.skip("Benchmark test fails: issue #2673")
+@pytest.mark.vstack
+def test_vstack():
+    bench = HStackBenchmark(
+        op_name="hstack", input_fn=_input_fn, torch_op=torch.hstack, dtypes=FLOAT_DTYPES
+    )
+    bench.run()

benchmark/test_repeat.py

@@ -0,0 +1,40 @@
+import pytest
+import torch
+
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.performance_utils import GenericBenchmark, generate_tensor_input
+
+
+class RepeatBenchmark(GenericBenchmark):
+    """
+    RepeatBenchmark designed to evaluate tensor repeat operations along specified dimensions.
+    This includes operations like tile, repeat, and repeat_interval.
+    Due to potential memory limitations, benchmark sizes need to be carefully controlled.
+
+    Notably, when the input size is set to (1024, 1024, 1024) and the repeat dimensions
+    are set to [1, 1, 2], the system encountered an "illegal memory access" error.
+    To avoid such issues, we constrain the benchmark input sizes for these operations
+    to prevent excessive memory usage.
+    """
+
+    def set_more_shapes(self):
+        return [(16, 256, 256), (512, 512, 512), (64, 64, 64, 64)]
+
+
+def _input_fn(shape, cur_dtype, device):
+    inp1 = generate_tensor_input(shape, cur_dtype, device)
+    inp2 = [1] * len(shape)
+    inp2[0] = 2
+
+    yield inp1, inp2,
+
+
+@pytest.mark.repeat
+def test_repeat():
+    bench = RepeatBenchmark(
+        op_name="repeat",
+        input_fn=_input_fn,
+        torch_op=torch.Tensor.repeat,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_repeat_interleave.py

@@ -0,0 +1,94 @@
+import pytest
+import torch
+
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.performance_utils import GenericBenchmark, generate_tensor_input
+
+
+class RepeatInterleaveBenchmark(GenericBenchmark):
+    """
+    Due to potential memory limitations, benchmark sizes need to be carefully controlled.
+
+    Notably, when the input size is set to (1024, 1024, 1024) and the repeat dimensions
+    are set to [1, 1, 2], the system encountered an "illegal memory access" error.
+    To avoid such issues, we constrain the benchmark input sizes for these operations
+    to prevent excessive memory usage.
+    """
+
+    def set_more_shapes(self):
+        return [
+            (16, 256, 256),
+            (512, 512, 512),
+            (64, 64, 64, 64),
+        ]
+
+
+# repeat_interleave.self_int(Tensor self, SymInt repeats, int? dim=None, *, SymInt? output_size=None) -> Tensor
+def repeat_interleave_self_int_input_fn(shape, dtype, device):
+    inp = generate_tensor_input(shape, dtype, device)
+    repeats = 3
+    yield inp, repeats,
+
+
+@pytest.mark.repeat_interleave
+def test_repeat_interleave_self_int():
+    bench = RepeatInterleaveBenchmark(
+        input_fn=repeat_interleave_self_int_input_fn,
+        op_name="repeat_interleave.self_int",
+        torch_op=torch.repeat_interleave,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.run()
+
+
+# repeat_interleave.self_Tensor(Tensor self, Tensor repeats, int? dim=None, *, SymInt? output_size=None) -> Tensor
+def repeat_interleave_self_tensor_input_fn(shape, dtype, device):
+    inp = generate_tensor_input(shape, dtype, device)
+    repeats = torch.randint(
+        low=0,
+        high=0x1F,  # control the repeats number here
+        size=[
+            shape[0],
+        ],
+        device=device,
+    )
+    dim = 0
+    yield inp, repeats, dim
+
+
+@pytest.mark.skip(reason="This test case runs out of memory: issue #2674")
+@pytest.mark.repeat_interleave
+def test_repeat_interleave_self_tensor():
+    bench = RepeatInterleaveBenchmark(
+        op_name="repeat_interleave.self_tensor",
+        input_fn=repeat_interleave_self_tensor_input_fn,
+        torch_op=torch.repeat_interleave,
+        dtypes=[torch.int32],
+    )
+    bench.run()
+
+
+# repeat_interleave.Tensor(Tensor repeats, *, SymInt? output_size=None) -> Tensor
+def repeat_interleave_tensor_input_fn(shape, dtype, device):
+    repeats = torch.randint(
+        low=0,
+        high=0x1F,  # control the repeats number here
+        size=[
+            shape[0],
+        ],
+        device=device,
+    )
+    yield repeats,
+
+
+@pytest.mark.skip(reason="This test case runs out of memory: issue #2674")
+@pytest.mark.repeat_interleave
+def test_repeat_interleave_tensor():
+    bench = RepeatInterleaveBenchmark(
+        op_name="repeat_interleave.tensor",
+        input_fn=repeat_interleave_tensor_input_fn,
+        torch_op=torch.repeat_interleave,
+        dtypes=[torch.int32],
+    )
+
+    bench.run()

benchmark/test_stack.py

@@ -0,0 +1,41 @@
+from typing import Generator
+
+import pytest
+import torch
+
+from benchmark.attri_util import FLOAT_DTYPES, BenchLevel
+from benchmark.performance_utils import Benchmark, Config, generate_tensor_input
+
+
+class StackBenchmark(Benchmark):
+    def __init__(self, *args, input_fn, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.input_fn = input_fn
+
+    def get_input_iter(self, cur_dtype) -> Generator:
+        for shape in self.shapes:
+            yield from self.input_fn(shape, cur_dtype, self.device)
+
+    def set_more_shapes(self):
+        more_shapes_2d = [(1024, 2**i) for i in range(1, 11, 4)]
+        more_shapes_3d = [(64, 64, 2**i) for i in range(0, 8, 4)]
+        return more_shapes_2d + more_shapes_3d
+
+
+def _input_fn(shape, dtype, device):
+    inp1 = generate_tensor_input(shape, dtype, device)
+    inp2 = generate_tensor_input(shape, dtype, device)
+    inp3 = generate_tensor_input(shape, dtype, device)
+    yield [inp1, inp2, inp3], {"dim": 0},
+
+    if Config.bench_level == BenchLevel.COMPREHENSIVE:
+        yield [inp1, inp2, inp3], {"dim": -1},
+
+
+@pytest.mark.skip(reason="CUDA error - illegal memory access: issue #2675")
+@pytest.mark.stack
+def test_stack():
+    bench = StackBenchmark(
+        op_name="stack", input_fn=_input_fn, torch_op=torch.stack, dtypes=FLOAT_DTYPES
+    )
+    bench.run()