Split benchmark for reduction operators

benchmark/performance_utils.py

@@ -10,6 +10,7 @@
 import yaml
 
 import flag_gems
+from flag_gems.utils import shape_utils
 
 from .attri_util import (
     BOOL_DTYPES,
@@ -552,6 +553,33 @@ def set_more_shapes(self):
         return [shape for shape in shapes if len(shape) == 2]
 
 
+class UnaryReductionBenchmark(Benchmark):
+    def set_more_metrics(self):
+        return ["gbps"]
+
+    def get_gbps(self, args, latency):
+        inp = args[0]
+        io_amount = sum([shape_utils.size_in_bytes(item) for item in [inp, inp]])
+        return io_amount * 1e-9 / (latency * 1e-3)
+
+    def set_more_shapes(self):
+        more_shapes_1d = [
+            (1025 * 1024,),
+            (1024 * 1024 * 1024,),
+        ]
+        more_shapes_2d = [(1024, 2**i) for i in range(0, 21, 4)]
+        more_shapes_3d = [(64, 2**i, 64) for i in range(0, 15, 4)]
+        return more_shapes_1d + more_shapes_2d + more_shapes_3d
+
+    def get_input_iter(self, cur_dtype) -> Generator:
+        for shape in self.shapes:
+            inp = generate_tensor_input(shape, cur_dtype, self.device)
+            if inp.ndim > 1:
+                yield inp, 1
+            else:
+                yield inp,
+
+
 def generate_tensor_input(shape, dtype, device):
     if dtype in FLOAT_DTYPES:
         return torch.randn(shape, dtype=dtype, device=device)

benchmark/test_all.py

@@ -0,0 +1,13 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+@pytest.mark.all
+def test_all():
+    bench = utils.UnaryReductionBenchmark(
+        op_name="all", torch_op=torch.all, dtypes=attr_utils.FLOAT_DTYPES
+    )
+    bench.run()

benchmark/test_amax.py

@@ -0,0 +1,13 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+@pytest.mark.test_amax
+def test_amax():
+    bench = utils.UnaryReductionBenchmark(
+        op_name="amax", torch_op=torch.amax, dtypes=attr_utils.FLOAT_DTYPES
+    )
+    bench.run()

benchmark/test_aminmax.py

@@ -0,0 +1,33 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+def aminmax_input_fn(shape, cur_dtype, device):
+    inp = utils.generate_tensor_input(shape, cur_dtype, device)
+    # Test dim=None (whole tensor reduction)
+    yield inp,
+    # Test dim=-1 (last dimension)
+    yield inp, {"dim": -1}
+    # Test dim=0 (first dimension)
+    if len(shape) > 1:
+        yield inp, {"dim": 0}
+
+
+class AminmaxBenchmark(utils.UnaryReductionBenchmark):
+    def get_input_iter(self, cur_dtype):
+        for shape in self.shapes:
+            yield from aminmax_input_fn(shape, cur_dtype, self.device)
+
+
+@pytest.mark.aminmax
+def test_aminmax():
+    bench = AminmaxBenchmark(
+        op_name="aminmax",
+        torch_op=torch.aminmax,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+
+    bench.run()

benchmark/test_any.py

@@ -0,0 +1,13 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+@pytest.mark.any
+def test_any():
+    bench = utils.UnaryReductionBenchmark(
+        op_name="any", torch_op=torch.any, dtypes=attr_utils.FLOAT_DTYPES
+    )
+    bench.run()

benchmark/test_argmax.py

@@ -0,0 +1,13 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+@pytest.mark.argmax
+def test_argmax():
+    bench = utils.UnaryReductionBenchmark(
+        op_name="argmax", torch_op=torch.argmax, dtypes=attr_utils.FLOAT_DTYPES
+    )
+    bench.run()

benchmark/test_argmin.py

@@ -0,0 +1,13 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+@pytest.mark.argmin
+def test_argmin():
+    bench = utils.UnaryReductionBenchmark(
+        op_name="argmin", torch_op=torch.argmin, dtypes=attr_utils.FLOAT_DTYPES
+    )
+    bench.run()

benchmark/test_avg_pool2d.py

@@ -0,0 +1,98 @@
+from typing import Generator
+
+import pytest
+import torch
+
+import flag_gems
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+class AvgPool2dBenchmark(utils.GenericBenchmark):
+    def get_input_iter(self, cur_dtype) -> Generator:
+        shapes_4d = [
+            (4, 3, 224, 224),  # Typical input image size
+            (16, 64, 56, 56),  # Early ResNet layer output
+            (32, 128, 28, 28),  # Mid ResNet layer output
+            (64, 256, 14, 14),  # Later ResNet layer output
+            (128, 512, 7, 7),  # Final ResNet layer output
+        ]
+
+        for shape in shapes_4d:
+            yield from self.input_fn(shape, cur_dtype, self.device)
+
+
+def avg_pool2d_input_fn(shape, dtype, device):
+    inp = utils.generate_tensor_input(shape, dtype, device)
+
+    # Common case
+    yield inp, {
+        "kernel_size": 3,
+        "stride": 2,
+        "padding": 1,
+        "ceil_mode": False,
+        "count_include_pad": True,
+        "divisor_override": None,
+    }
+
+    if utils.Config.bench_level == utils.BenchLevel.COMPREHENSIVE:
+        # With count_include_pad=False
+        yield inp, {
+            "kernel_size": 3,
+            "stride": 2,
+            "padding": 1,
+            "ceil_mode": False,
+            "count_include_pad": False,
+            "divisor_override": None,
+        }
+
+        # With ceil_mode
+        yield inp, {
+            "kernel_size": 3,
+            "stride": 2,
+            "padding": 1,
+            "ceil_mode": True,
+            "count_include_pad": True,
+            "divisor_override": None,
+        }
+
+        # With divisor_override
+        if shape[-2] >= 2 and shape[-1] >= 2:
+            yield inp, {
+                "kernel_size": 2,
+                "stride": 1,
+                "padding": 0,
+                "ceil_mode": False,
+                "count_include_pad": True,
+                "divisor_override": 3,
+            }
+
+
+@pytest.mark.avg_pool2d
+def test_avg_pool2d():
+    bench = AvgPool2dBenchmark(
+        input_fn=avg_pool2d_input_fn,
+        op_name="avg_pool2d",
+        torch_op=torch.ops.aten.avg_pool2d,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+    bench.run()
+
+
+@pytest.mark.skip(reason="Test case fails due to missing parameter self.")
+@pytest.mark.avg_pool2d_backward
+def test_avg_pool2d_backward():
+    if flag_gems.vendor_name == "mthreads":
+        dtypes = [torch.float32]
+    else:
+        dtypes = (attr_utils.FLOAT_DTYPES,)
+
+    bench = AvgPool2dBenchmark(
+        input_fn=avg_pool2d_input_fn,
+        op_name="avg_pool2d_backward",
+        torch_op=torch.ops.aten.avg_pool2d_backward,
+        dtypes=dtypes,
+        is_backward=True,
+    )
+    bench.run()

benchmark/test_bincount.py

@@ -0,0 +1,59 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+def bincount_input_fn(shape, dtype, device):
+    if shape[0] > 1_000_000:
+        return
+
+    n = shape[0]
+    for num_classes in [10, 256, 4096]:
+        inp = torch.randint(0, num_classes, (n,), dtype=torch.int64, device=device)
+
+        yield inp, {}
+
+        yield inp, {"minlength": max(512, num_classes * 2)}
+
+
+@pytest.mark.bincount
+def test_bincount():
+    bench = utils.GenericBenchmark(
+        input_fn=bincount_input_fn,
+        op_name="bincount",
+        torch_op=torch.bincount,
+        dtypes=[torch.float32],
+    )
+    bench.set_gems(flag_gems.bincount)
+    bench.run()
+
+
+def bincount_weighted_input_fn(shape, dtype, device):
+    if shape[0] > 1_000_000:
+        return
+
+    n = shape[0]
+    for num_classes in [10, 256, 4096]:
+        inp = torch.randint(0, num_classes, (n,), dtype=torch.int64, device=device)
+        weights = torch.randn((n,), dtype=dtype, device=device)
+
+        yield inp, {"weights": weights}
+
+        yield inp, {"weights": weights, "minlength": max(512, num_classes * 2)}
+
+
+@pytest.mark.bincount
+@pytest.mark.parametrize("dtype", attr_utils.FLOAT_DTYPES)
+def test_bincount_weighted(dtype):
+    bench = utils.GenericBenchmark(
+        input_fn=bincount_weighted_input_fn,
+        op_name=f"bincount_weighted_{str(dtype).split('.')[-1]}",
+        torch_op=torch.bincount,
+        dtypes=[dtype],
+    )
+    bench.set_gems(flag_gems.bincount)
+    bench.run()

benchmark/test_count_nonzero.py

@@ -0,0 +1,24 @@
+import random
+
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+@pytest.mark.count_nonzero
+def test_count_nonzero():
+    def count_nonzero_input_fn(shape, dtype, device):
+        inp = torch.randn(shape, dtype=dtype, device=device)
+        dim = random.choice([None, 0, 1])
+
+        yield inp, dim
+
+    bench = utils.GenericBenchmark2DOnly(
+        input_fn=count_nonzero_input_fn,
+        op_name="count_nonzero",
+        torch_op=torch.count_nonzero,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_cross_entropy_loss.py

@@ -0,0 +1,34 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+def cross_entropy_loss_input_fn(shape, cur_dtype, device):
+    inp = utils.generate_tensor_input(shape, cur_dtype, device)
+    target = torch.randint(0, shape[-1], (shape[0],), device=device)
+    yield inp, target
+
+    if utils.Config.bench_level == utils.BenchLevel.COMPREHENSIVE:
+        weight = torch.randn(shape[-1], dtype=cur_dtype, device=device)
+        yield inp, target, {"weight": weight, "ignore_index": 1, "reduction": "none"}
+        yield inp, target, {
+            "weight": weight,
+            "reduction": "sum",
+            "label_smoothing": 0.1,
+        }
+
+
+@pytest.mark.cross_entropy_loss
+def test_cross_entropy_loss():
+    bench = utils.GenericBenchmark2DOnly(
+        input_fn=cross_entropy_loss_input_fn,
+        op_name="cross_entropy_loss",
+        torch_op=torch.nn.functional.cross_entropy,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+    bench.set_gems(flag_gems.cross_entropy_loss)
+    bench.run()

benchmark/test_cummax.py

@@ -0,0 +1,22 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+def input_fn(shape, cur_dtype, device):
+    inp = utils.generate_tensor_input(shape, cur_dtype, device)
+    yield inp, 1
+
+
+@pytest.mark.cummax
+def test_cummax():
+    bench = utils.GenericBenchmark2DOnly(
+        input_fn=input_fn,
+        op_name="cummax",
+        torch_op=torch.cummax,
+        dtypes=attr_utils.FLOAT_DTYPES + attr_utils.INT_DTYPES,
+    )
+
+    bench.run()