Split norm benchmarks

benchmark/test_batch_norm.py

@@ -0,0 +1,51 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+class NormBenchmark(utils.GenericBenchmark):
+    # TODO: add new metric
+
+    def set_more_shapes(self):
+        return [
+            # 3D shapes represented as [batch_size, channels, hidden_size]
+            (16, 16, 64),
+            (16, 16, 1024),
+            (16, 16, 4098),
+            # 4D shapes represented as [batch_size, channels, H, W]
+            (1, 8, 4, 4),
+            (16, 8, 128, 128),
+        ]
+
+
+def input_fn(shape, dtype, device):
+    C = shape[1]
+    inp = torch.randn(shape, dtype=dtype, device=device)
+    weight = torch.randn((C,), dtype=dtype, device=device)
+    bias = torch.randn((C,), dtype=dtype, device=device)
+    running_mean = None
+    running_var = None
+    training = True
+    momentum = 0.1
+    eps = 1e-5
+    cudnn_enabled = True
+    yield inp, weight, bias, running_mean, running_var, training, momentum, eps, cudnn_enabled
+
+    if utils.Config.bench_level == utils.BenchLevel.COMPREHENSIVE:
+        running_mean = torch.randn((C,), dtype=dtype, device=device)
+        running_var = torch.randn((C,), dtype=dtype, device=device)
+        yield inp, weight, bias, running_mean, running_var, training, momentum, eps, cudnn_enabled
+
+
+@pytest.mark.batch_norm
+def test_batch_norm():
+    bench = NormBenchmark(
+        op_name="batch_norm",
+        input_fn=input_fn,
+        torch_op=torch.batch_norm,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+
+    bench.run()

benchmark/test_batch_norm_backward.py

@@ -0,0 +1,96 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import attr_util as attr_utils
+from . import performance_utils as utils
+
+
+# TODO(Qiming): Consolidate this to a base package
+class NormBenchmark(utils.GenericBenchmark):
+    # TODO: add new metric
+
+    def set_more_shapes(self):
+        return [
+            # 3D shapes represented as [batch_size, channels, hidden_size]
+            (16, 16, 64),
+            (16, 16, 1024),
+            (16, 16, 4098),
+            # 4D shapes represented as [batch_size, channels, H, W]
+            (1, 8, 4, 4),
+            (16, 8, 128, 128),
+        ]
+
+
+def batchnorm_input_fn(shape, dtype, device):
+    C = shape[1]
+    inp = torch.randn(shape, dtype=dtype, device=device)
+    weight = torch.randn((C,), dtype=dtype, device=device)
+    bias = torch.randn((C,), dtype=dtype, device=device)
+    running_mean = None
+    running_var = None
+    training = True
+    momentum = 0.1
+    eps = 1e-5
+    cudnn_enabled = True
+    yield inp, weight, bias, running_mean, running_var, training, momentum, eps, cudnn_enabled
+
+    if utils.Config.bench_level == utils.BenchLevel.COMPREHENSIVE:
+        running_mean = torch.randn((C,), dtype=dtype, device=device)
+        running_var = torch.randn((C,), dtype=dtype, device=device)
+        yield inp, weight, bias, running_mean, running_var, training, momentum, eps, cudnn_enabled
+
+
+@pytest.mark.batch_norm_backward
+def test_batch_norm_backward():
+    def batch_norm_backward_input_fn(shape, dtype, device):
+        for forward_args in batchnorm_input_fn(shape, dtype, device):
+            (
+                inp,
+                weight,
+                bias,
+                running_mean,
+                running_var,
+                training,
+                _,
+                eps,
+                _,
+            ) = forward_args
+
+            grad_output = torch.randn_like(inp)
+            channels = weight.shape[0] if weight is not None else inp.shape[1]
+
+            if running_mean is None:
+                running_mean = torch.zeros(channels, dtype=dtype, device=device)
+            if running_var is None:
+                running_var = torch.ones(channels, dtype=dtype, device=device)
+
+            save_mean = torch.randn(channels, dtype=torch.float32, device=device)
+            save_invstd = torch.randn(channels, dtype=torch.float32, device=device)
+            output_mask = [True, weight is not None, bias is not None]
+
+            yield (
+                grad_output,
+                inp,
+                weight,
+                running_mean,
+                running_var,
+                save_mean,
+                save_invstd,
+                training,
+                eps,
+                output_mask,
+            )
+
+    bench = NormBenchmark(
+        input_fn=batch_norm_backward_input_fn,
+        op_name="native_batch_norm_backward",
+        torch_op=torch.ops.aten.native_batch_norm_backward,
+        dtypes=[torch.float32]
+        if flag_gems.vendor_name == "mthreads"
+        else attr_utils.FLOAT_DTYPES,
+    )
+    bench.set_gems(flag_gems.batch_norm_backward)
+
+    bench.run()

benchmark/test_group_norm.py

@@ -0,0 +1,55 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+# TODO(Qiming): Extract this to a base class
+class NormBenchmark(utils.GenericBenchmark):
+    # TODO: add new metric
+
+    def set_more_shapes(self):
+        return [
+            # 3D shapes represented as [batch_size, channels, hidden_size]
+            (16, 16, 64),
+            (16, 16, 1024),
+            (16, 16, 4098),
+            # 4D shapes represented as [batch_size, channels, H, W]
+            (1, 8, 4, 4),
+            (16, 8, 128, 128),
+        ]
+
+
+def group_norm_input_fn(shape, dtype, device):
+    inp = torch.randn(shape, dtype=dtype, device=device)
+    channel = shape[1]
+    weight = torch.randn(
+        [
+            channel,
+        ],
+        dtype=dtype,
+        device=device,
+    )
+    bias = torch.randn(
+        [
+            channel,
+        ],
+        dtype=dtype,
+        device=device,
+    )
+    yield inp, channel // 2, weight, bias
+
+    if utils.Config.bench_level == utils.BenchLevel.COMPREHENSIVE:
+        yield inp, channel, weight, bias
+
+
+@pytest.mark.group_norm
+def test_group_norm():
+    bench = NormBenchmark(
+        input_fn=group_norm_input_fn,
+        op_name="group_norm",
+        torch_op=torch.nn.functional.group_norm,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+    bench.run()

benchmark/test_instance_norm.py

@@ -0,0 +1,61 @@
+import pytest
+import torch
+
+import flag_gems
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+class NormBenchmark(utils.GenericBenchmark):
+    # TODO: add new metric
+
+    def set_more_shapes(self):
+        return [
+            # 3D shapes represented as [batch_size, channels, hidden_size]
+            (16, 16, 64),
+            (16, 16, 1024),
+            (16, 16, 4098),
+            # 4D shapes represented as [batch_size, channels, H, W]
+            (1, 8, 4, 4),
+            (16, 8, 128, 128),
+        ]
+
+
+def input_fn(shape, dtype, device):
+    C = shape[1]
+    inp = torch.randn(shape, dtype=dtype, device=device)
+    weight = torch.randn((C,), dtype=dtype, device=device)
+    bias = torch.randn((C,), dtype=dtype, device=device)
+    running_mean = None
+    running_var = None
+    use_input_stats = True
+    momentum = 0.1
+    eps = 1e-5
+    cudnn_enabled = True
+    yield inp, weight, bias, running_mean, running_var, use_input_stats, momentum, eps, cudnn_enabled
+    if utils.Config.bench_level == utils.BenchLevel.COMPREHENSIVE:
+        running_mean = torch.randn((C,), dtype=dtype, device=device)
+        running_var = torch.randn((C,), dtype=dtype, device=device)
+        yield inp, weight, bias, running_mean, running_var, use_input_stats, momentum, eps, cudnn_enabled
+
+
+@pytest.mark.instance_norm
+def test_instance_norm(monkeypatch):
+    if flag_gems.vendor_name == "kunlunxin" and utils.SkipVersion("torch", "<2.5"):
+        pytest.skip(
+            "BF16 is not supported in XPytorch 2.0. Please upgrade your PyTorch version >= 2.5"
+        )
+
+    if flag_gems.vendor_name == "mthreads":
+        # Compatible with older versions of LLVM
+        monkeypatch.env("DISABLE_LLVM_OPT", "1")
+
+    bench = NormBenchmark(
+        op_name="instance_norm",
+        input_fn=input_fn,
+        torch_op=torch.instance_norm,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+    bench.set_gems(flag_gems.instance_norm)
+    bench.run()

benchmark/test_layer_norm.py

@@ -0,0 +1,40 @@
+import pytest
+import torch
+
+from . import attri_util as attr_utils
+from . import performance_utils as utils
+
+
+# TODO(Qiming): Extract this to a base class
+class NormBenchmark(utils.GenericBenchmark):
+    # TODO: add new metric
+
+    def set_more_shapes(self):
+        return [
+            # 3D shapes represented as [batch_size, channels, hidden_size]
+            (16, 16, 64),
+            (16, 16, 1024),
+            (16, 16, 4098),
+            # 4D shapes represented as [batch_size, channels, H, W]
+            (1, 8, 4, 4),
+            (16, 8, 128, 128),
+        ]
+
+
+def input_fn(shape, dtype, device):
+    inp = torch.randn(shape, dtype=dtype, device=device)
+    layer_shape = shape[1:]
+    weight = torch.randn(layer_shape, dtype=dtype, device=device)
+    bias = torch.randn(layer_shape, dtype=dtype, device=device)
+    yield inp, layer_shape, weight, bias
+
+
+@pytest.mark.layer_norm
+def test_layer_norm():
+    bench = NormBenchmark(
+        op_name="layer_norm",
+        input_fn=input_fn,
+        torch_op=torch.layer_norm,
+        dtypes=attr_utils.FLOAT_DTYPES,
+    )
+    bench.run()