[FlagGems Operator Development Competition] Optimize conv_transpose2d forward

benchmark/test_conv_transpose2d.py

@@ -0,0 +1,76 @@
+from typing import Generator
+
+import pytest
+import torch
+
+import flag_gems
+from benchmark.attri_util import FLOAT_DTYPES
+from benchmark.performance_utils import GenericBenchmark
+
+
+class ConvTranspose2DBenchmark(GenericBenchmark):
+    def set_more_shapes(self):
+        return [
+            (16, 32, 24, 24, 24, 3, 3, 1, 1, 0, 2, 1),
+            (16, 32, 24, 24, 24, 3, 3, 2, 1, 1, 2, 1),
+            (32, 64, 64, 64, 32, 3, 3, 2, 1, 1, 1, 1),
+            (32, 64, 128, 128, 32, 5, 5, 2, 2, 1, 1, 1),
+            (4, 32, 128, 128, 32, 5, 5, 3, 2, 0, 1, 1),
+            (16, 32, 24, 24, 24, 3, 3, 3, 1, 0, 2, 1),
+            (16, 32, 24, 24, 24, 3, 3, 4, 1, 0, 2, 1),
+            (8, 32, 32, 32, 32, 3, 3, 4, 1, 0, 1, 1),
+            (4, 32, 32, 32, 32, 3, 3, 8, 1, 0, 1, 1),
+        ]
+
+    def get_input_iter(self, cur_dtype) -> Generator:
+        for shape in self.set_more_shapes():
+            yield from self.input_fn(shape, cur_dtype, self.device)
+
+
+def _input_fn(shape, dtype, device):
+    (
+        batch,
+        input_c,
+        input_h,
+        input_w,
+        out_c,
+        kernel_h,
+        kernel_w,
+        stride,
+        padding,
+        output_padding,
+        groups,
+        dilation,
+    ) = shape
+    input_shape = (batch, input_c, input_h, input_w)
+    weight_shape = (input_c, out_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,
+        "stride": stride,
+        "padding": padding,
+        "output_padding": output_padding,
+        "groups": groups,
+        "dilation": dilation,
+    },
+
+
+@pytest.mark.conv_transpose2d
+def test_conv_transpose2d():
+    def gems_conv_transpose2d(**kwargs):
+        with torch.no_grad():
+            return flag_gems.conv_transpose2d(**kwargs)
+
+    torch.backends.cudnn.allow_tf32 = False
+    bench = ConvTranspose2DBenchmark(
+        input_fn=_input_fn,
+        op_name="conv_transpose2d",
+        torch_op=torch.nn.functional.conv_transpose2d,
+        dtypes=FLOAT_DTYPES,
+    )
+    bench.set_gems(gems_conv_transpose2d)
+    bench.run()

src/flag_gems/__init__.py

@@ -140,6 +140,7 @@ def torch_ge(v):
     ("conv1d.padding", conv1d),
     ("conv2d", conv2d),
     ("conv2d.padding", conv2d),
+    ("conv_transpose2d.input", conv_transpose2d),
     ("conv3d", conv3d),
     ("conv3d.padding", conv3d),
     (

src/flag_gems/ops/__init__.py

@@ -80,6 +80,7 @@
 from flag_gems.ops.conv2d import conv2d
 from flag_gems.ops.conv3d import conv3d
 from flag_gems.ops.conv_depthwise2d import _conv_depthwise2d
+from flag_gems.ops.conv_transpose2d import conv_transpose2d
 from flag_gems.ops.copy import copy, copy_
 from flag_gems.ops.copysign import copysign, copysign_out
 from flag_gems.ops.cos import cos, cos_
@@ -427,6 +428,7 @@
     "conv1d",
     "conv2d",
     "conv3d",
+    "conv_transpose2d",
     "copy",
     "copy_",
     "copysign",