Split benchmarks for convolution operators (#2670)

Author: tengqm

benchmark/test_conv1d.py

@@ -0,0 +1,62 @@
+import pytest
+import torch
+
+import flag_gems
+from benchmark.performance_utils import GenericBenchmark
+
+
+class Conv1DBenchmark(GenericBenchmark):
+    def set_more_shapes(self):
+        return [
+            (32, 64, 512, 64, 3, 1, 0, 1),
+            (64, 48, 1024, 128, 5, 2, 2, 1),
+            (16, 24, 2048, 96, 7, 1, 3, 2),
+            (8, 8, 8192, 16, 11, 4, 5, 1),
+            (4, 4, 16384, 4, 15, 2, 7, 1),
+            (32, 64, 512, 64, 3, 1, "valid", 1),
+            (64, 48, 1024, 128, 5, 2, "valid", 1),
+            (16, 24, 2048, 96, 7, 1, "same", 2),
+            (8, 8, 8192, 16, 11, 1, "same", 1),
+        ]
+
+
+def conv1d_input_fn(shape, dtype, device):
+    (
+        batch,
+        input_c,
+        input_l,
+        out_c,
+        kernel,
+        stride,
+        padding,
+        groups,
+    ) = shape
+    input_shape = (batch, input_c, input_l)
+    weight_shape = (out_c, input_c // groups, kernel)
+    input = torch.randn(size=input_shape, device=device, dtype=dtype)
+    weight = torch.randn(size=weight_shape, device=device, dtype=dtype)
+
+    yield {
+        "input": input,
+        "weight": weight,
+        "bias": None,
+        "groups": groups,
+        "stride": stride,
+        "padding": padding,
+    },
+
+
+@pytest.mark.conv1d
+def test_conv1d():
+    torch.backends.cudnn.allow_tf32 = False
+    bench = Conv1DBenchmark(
+        input_fn=conv1d_input_fn,
+        op_name="conv1d",
+        torch_op=torch.nn.functional.conv1d,
+        dtypes=[
+            torch.float16,
+            torch.float32,
+        ],  # Exclude bfloat16 due to cuDNN limitations
+    )
+    bench.set_gems(flag_gems.conv1d)
+    bench.run()

benchmark/test_conv2d.py

@@ -0,0 +1,66 @@
+import pytest
+import torch
+
+import flag_gems
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.performance_utils import GenericBenchmark
+
+
+class Conv2DBenchmark(GenericBenchmark):
+    def set_more_shapes(self):
+        return [
+            (32, 64, 128, 128, 32, 3, 3, 1, 2, 1),
+            (32, 64, 210, 210, 16, 5, 5, 2, 1, 1),
+            (16, 32, 12, 12, 24, 3, 3, 2, 1, 1),
+            (16, 32, 24, 24, 24, 3, 3, 2, 2, 2),
+            (16, 32, 24, 24, 24, 3, 3, 1, 2, 2),
+            (16, 32, 12, 12, 24, 3, 3, 2, "valid", 1),
+            (32, 64, 128, 128, 32, 3, 3, 1, "valid", 1),
+            (16, 32, 24, 24, 24, 3, 3, 1, "same", 2),
+            (32, 64, 210, 210, 16, 5, 5, 1, "same", 1),
+        ]
+
+
+def _input_fn(shape, dtype, device):
+    (
+        batch,
+        input_c,
+        input_h,
+        input_w,
+        out_c,
+        kernel_h,
+        kernel_w,
+        stride,
+        padding,
+        groups,
+    ) = shape
+    input_shape = (batch, input_c, input_h, input_w)
+    weight_shape = (out_c, input_c // groups, kernel_h, kernel_w)
+    input = torch.randn(size=input_shape, device=device, dtype=dtype)
+    weight = torch.randn(size=weight_shape, device=device, dtype=dtype)
+
+    yield {
+        "input": input,
+        "weight": weight,
+        "bias": None,
+        "groups": groups,
+        "stride": stride,
+        "padding": padding,
+    },
+
+
+@pytest.mark.conv2d
+def test_conv2d(monkeypatch):
+    if flag_gems.vendor_name == "hygon":
+        monkeypatch.setenv("TRITON_HIP_USE_NEW_STREAM_PIPELINE", "0")
+
+    torch.backends.cudnn.allow_tf32 = False
+    bench = Conv2DBenchmark(
+        input_fn=_input_fn,
+        op_name="conv2d",
+        torch_op=torch.nn.functional.conv2d,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.set_gems(flag_gems.conv2d)
+
+    bench.run()

benchmark/test_conv3d.py

@@ -0,0 +1,53 @@
+import pytest
+import torch
+
+import flag_gems
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.performance_utils import GenericBenchmark
+
+
+class Conv3DBenchmark(GenericBenchmark):
+    def set_more_shapes(self):
+        return None
+
+
+@pytest.mark.conv3d
+def test_conv3d():
+    def conv3d_input_fn(shape, dtype, device):
+        (
+            batch,
+            input_c,
+            input_d,
+            input_h,
+            input_w,
+            out_c,
+            kernel_d,
+            kernel_h,
+            kernel_w,
+            stride,
+            padding,
+            groups,
+        ) = shape
+        input_shape = (batch, input_c, input_d, input_h, input_w)
+        weight_shape = (out_c, input_c // groups, kernel_d, kernel_h, kernel_w)
+        input = torch.randn(size=input_shape, device=device, dtype=dtype)
+        weight = torch.randn(size=weight_shape, device=device, dtype=dtype)
+
+        yield {
+            "input": input,
+            "weight": weight,
+            "bias": None,
+            "groups": groups,
+            "stride": stride,
+            "padding": padding,
+        },
+
+    torch.backends.cudnn.allow_tf32 = False
+    bench = Conv3DBenchmark(
+        op_name="conv3d",
+        input_fn=conv3d_input_fn,
+        torch_op=torch.nn.functional.conv3d,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.set_gems(flag_gems.conv3d)
+    bench.run()

benchmark/test_conv_depthwise2d.py

@@ -0,0 +1,58 @@
+import pytest
+import torch
+
+import flag_gems
+from benchmark.performance_utils import GenericBenchmark
+
+
+class ConvDepthwise2DBenchmark(GenericBenchmark):
+    def set_more_shapes(self):
+        # Additional shapes for COMPREHENSIVE mode
+        return [
+            (1, 64, 224, 224, 3, 3, 1, 1, 1),
+            (1, 128, 112, 112, 5, 5, 2, 2, 1),
+        ]
+
+
+def _input_fn(shape, dtype, device):
+    (
+        batch,
+        channels,
+        input_h,
+        input_w,
+        kernel_h,
+        kernel_w,
+        stride,
+        padding,
+        dilation,
+    ) = shape
+    input_shape = (batch, channels, input_h, input_w)
+    weight_shape = (channels, 1, kernel_h, kernel_w)
+    input_tensor = torch.randn(size=input_shape, device=device, dtype=dtype)
+    weight = torch.randn(size=weight_shape, device=device, dtype=dtype)
+
+    # Pass as positional args since the first arg is named 'self' in aten op
+    yield (
+        input_tensor,
+        weight,
+        [kernel_h, kernel_w],
+        None,  # bias
+        [stride, stride],
+        [padding, padding],
+        [dilation, dilation],
+    )
+
+
+@pytest.mark.conv_depthwise2d
+def test_conv_depthwise2d():
+    torch.backends.cudnn.allow_tf32 = False
+
+    bench = ConvDepthwise2DBenchmark(
+        op_name="_conv_depthwise2d",
+        input_fn=_input_fn,
+        torch_op=torch.ops.aten._conv_depthwise2d,
+        gems_op=flag_gems._conv_depthwise2d,
+        dtypes=[torch.float16, torch.float32],
+    )
+
+    bench.run()