【KernelGen】Add cudnn_convolution operator

benchmark/test_cudnn_convolution.py

@@ -0,0 +1,62 @@
+from typing import Generator
+
+import pytest
+import torch
+
+from . import base, consts, utils
+
+
+def cudnn_convolution_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)
+    inp = utils.generate_tensor_input(input_shape, dtype, device)
+    weight = utils.generate_tensor_input(weight_shape, dtype, device)
+
+    yield (
+        inp,
+        weight,
+        [padding, padding],
+        [stride, stride],
+        [1, 1],
+        groups,
+        False,
+        False,
+        False,
+    )
+
+
+class CudnnConv2dBenchmark(base.GenericBenchmark):
+    def get_input_iter(self, dtype) -> Generator:
+        shapes = [
+            (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),
+        ]
+
+        for shape in shapes:
+            yield from self.input_fn(shape, dtype, self.device)
+
+
+@pytest.mark.cudnn_convolution
+def test_cudnn_convolution():
+    bench = CudnnConv2dBenchmark(
+        input_fn=cudnn_convolution_input_fn,
+        op_name="cudnn_convolution",
+        torch_op=torch.ops.aten.cudnn_convolution.default,
+        dtypes=consts.FLOAT_DTYPES,
+    )
+    bench.run()

conf/operators.yaml

@@ -1337,6 +1337,18 @@ ops:
     stages:
       - stable: '2.0'
       - removed: '3.0'
+  - id: cudnn_convolution
+    description: |
+      A wrapper for cuDNN convolution backend.
+    for:
+      - cudnn_convolution
+    labels:
+      - aten
+      - KernelGen
+    kind:
+      - NeuralNetwork
+    stages:
+      - beta: '5.1'
   - id: cummax
     description: |
       Returns a named tuple `(values, indices)` where `values` is the cumulative maximum of elements

src/flag_gems/__init__.py

@@ -169,6 +169,7 @@ def torch_ge(v):
     ("copysign", copysign),
     ("copysign.out", copysign_out),
     ("count_nonzero", count_nonzero),
+    ("cudnn_convolution", cudnn_convolution),
     ("cummax", cummax),
     ("cummin", cummin),
     ("cumprod", cumprod),

src/flag_gems/ops/__init__.py

@@ -89,6 +89,7 @@
 from flag_gems.ops.cos import cos, cos_
 from flag_gems.ops.cosh import cosh, cosh_, cosh_out
 from flag_gems.ops.count_nonzero import count_nonzero
+from flag_gems.ops.cudnn_convolution import cudnn_convolution
 from flag_gems.ops.cummax import cummax
 from flag_gems.ops.cummin import cummin
 from flag_gems.ops.cumprod import cumprod, cumprod_
@@ -483,6 +484,7 @@
     "cosh_",
     "cosh_out",
     "count_nonzero",
+    "cudnn_convolution",
     "cummax",
     "cummin",
     "cumprod",

src/flag_gems/ops/cudnn_convolution.py

@@ -0,0 +1,92 @@
+import logging
+
+from flag_gems.ops.conv1d import conv1d
+from flag_gems.ops.conv2d import conv2d
+from flag_gems.ops.conv3d import conv3d
+
+logger = logging.getLogger(__name__)
+
+
+def cudnn_convolution(
+    input,
+    weight,
+    padding,
+    stride,
+    dilation,
+    groups,
+    benchmark,
+    deterministic,
+    allow_tf32,
+):
+    """
+    CUDNN convolution operation.
+
+    This is a lower-level convolution operation that does not include bias.
+    It supports 1D, 2D, and 3D convolutions based on the input dimensionality.
+
+    Args:
+        input: Input tensor of shape (N, C_in, *spatial_dims)
+        weight: Weight tensor of shape (C_out, C_in/groups, *kernel_dims)
+        padding: Padding for each spatial dimension
+        stride: Stride for each spatial dimension
+        dilation: Dilation for each spatial dimension
+        groups: Number of groups for grouped convolution
+        benchmark: cuDNN benchmark flag (ignored in Triton implementation)
+        deterministic: cuDNN deterministic flag (ignored in Triton implementation)
+        allow_tf32: Allow TF32 computation flag (ignored in Triton implementation)
+
+    Returns:
+        Output tensor after convolution
+    """
+    logger.debug("GEMS CUDNN_CONVOLUTION")
+
+    ndim = input.ndim - 2
+
+    # Extract values from lists if they are lists (cudnn_convolution receives lists)
+    def extract_param(param, expected_len):
+        if isinstance(param, (list, tuple)):
+            if len(param) == expected_len:
+                return param if expected_len > 1 else param[0]
+            elif len(param) == 1:
+                return param[0]
+        return param
+
+    if ndim == 1:
+        # For 1D convolution, extract single values from lists
+        stride_val = extract_param(stride, 1)
+        padding_val = extract_param(padding, 1)
+        dilation_val = extract_param(dilation, 1)
+        return conv1d(
+            input,
+            weight,
+            bias=None,
+            stride=stride_val,
+            padding=padding_val,
+            dilation=dilation_val,
+            groups=groups,
+        )
+    elif ndim == 2:
+        return conv2d(
+            input,
+            weight,
+            bias=None,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+        )
+    elif ndim == 3:
+        return conv3d(
+            input,
+            weight,
+            bias=None,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+        )
+    else:
+        raise ValueError(
+            f"cudnn_convolution only supports 1D, 2D, and 3D convolutions, "
+            f"got input with {ndim} spatial dimensions"
+        )

tests/test_cudnn_convolution.py

@@ -0,0 +1,151 @@
+import pytest
+import torch
+
+import flag_gems
+
+from .accuracy_utils import gems_assert_close
+
+SHAPE_CUDNN_CONV2D = [
+    ((1, 2, 5, 5), (1, 2, 3, 3), 1),
+    ((2, 3, 9, 9), (1, 3, 3, 3), 1),
+    ((32, 8, 8, 8), (32, 8, 2, 2), 1),
+]
+
+
+@pytest.mark.cudnn_convolution
+@pytest.mark.parametrize("shape, kernel, groups", SHAPE_CUDNN_CONV2D)
+@pytest.mark.parametrize("stride", [1, 2])
+@pytest.mark.parametrize("padding", [0, 1])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
+@pytest.mark.parametrize("dilation", [1, 2])
+def test_cudnn_convolution_2d(
+    shape, kernel, stride, padding, groups, dtype, dilation, monkeypatch
+):
+    if flag_gems.vendor_name == "mthreads" and dtype == torch.float16:
+        monkeypatch.setenv("MUSA_ENABLE_SQMMA", "1")
+
+    inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    weight = torch.randn(kernel, dtype=dtype, device=flag_gems.device)
+
+    ref_out = torch.cudnn_convolution(
+        inp,
+        weight,
+        padding=[padding, padding],
+        stride=[stride, stride],
+        dilation=[dilation, dilation],
+        groups=groups,
+        benchmark=False,
+        deterministic=False,
+        allow_tf32=False,
+    )
+
+    with flag_gems.use_gems():
+        res_out = torch.cudnn_convolution(
+            inp,
+            weight,
+            padding=[padding, padding],
+            stride=[stride, stride],
+            dilation=[dilation, dilation],
+            groups=groups,
+            benchmark=False,
+            deterministic=False,
+            allow_tf32=False,
+        )
+
+    gems_assert_close(res_out.cpu(), ref_out.cpu(), dtype)
+
+
+SHAPE_CUDNN_CONV1D = [
+    ((32, 2, 4), (17, 2, 2)),
+    ((32, 15, 6), (17, 15, 2)),
+    ((64, 64, 64), (128, 64, 7)),
+]
+
+
+@pytest.mark.cudnn_convolution
+@pytest.mark.parametrize("shape, kernel", SHAPE_CUDNN_CONV1D)
+@pytest.mark.parametrize("stride", [1, 2])
+@pytest.mark.parametrize("padding", [0, 1])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
+def test_cudnn_convolution_1d(shape, kernel, stride, padding, dtype, monkeypatch):
+    if flag_gems.vendor_name == "mthreads" and dtype == torch.float16:
+        monkeypatch.setenv("MUSA_ENABLE_SQMMA", "1")
+
+    inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    weight = torch.randn(kernel, dtype=dtype, device=flag_gems.device)
+
+    ref_out = torch.cudnn_convolution(
+        inp,
+        weight,
+        padding=[padding],
+        stride=[stride],
+        dilation=[1],
+        groups=1,
+        benchmark=False,
+        deterministic=False,
+        allow_tf32=False,
+    )
+
+    with flag_gems.use_gems():
+        res_out = torch.cudnn_convolution(
+            inp,
+            weight,
+            padding=[padding],
+            stride=[stride],
+            dilation=[1],
+            groups=1,
+            benchmark=False,
+            deterministic=False,
+            allow_tf32=False,
+        )
+
+    gems_assert_close(res_out.cpu(), ref_out.cpu(), dtype)
+
+
+SHAPE_CUDNN_CONV3D = [
+    ((1, 2, 5, 5, 5), (1, 2, 3, 3, 3), 1),
+    ((2, 3, 9, 9, 9), (1, 3, 3, 3, 3), 1),
+]
+
+
+@pytest.mark.cudnn_convolution
+@pytest.mark.parametrize("shape, kernel, groups", SHAPE_CUDNN_CONV3D)
+@pytest.mark.parametrize("stride", [1, 2])
+@pytest.mark.parametrize("padding", [0, 1])
+@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
+@pytest.mark.parametrize("dilation", [1, 2])
+def test_cudnn_convolution_3d(
+    shape, kernel, stride, padding, groups, dtype, dilation, monkeypatch
+):
+    if flag_gems.vendor_name == "mthreads" and dtype == torch.float16:
+        monkeypatch.setenv("MUSA_ENABLE_SQMMA", "1")
+
+    inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
+    weight = torch.randn(kernel, dtype=dtype, device=flag_gems.device)
+
+    ref_out = torch.cudnn_convolution(
+        inp,
+        weight,
+        padding=[padding, padding, padding],
+        stride=[stride, stride, stride],
+        dilation=[dilation, dilation, dilation],
+        groups=groups,
+        benchmark=False,
+        deterministic=False,
+        allow_tf32=False,
+    )
+
+    with flag_gems.use_gems():
+        res_out = torch.cudnn_convolution(
+            inp,
+            weight,
+            padding=[padding, padding, padding],
+            stride=[stride, stride, stride],
+            dilation=[dilation, dilation, dilation],
+            groups=groups,
+            benchmark=False,
+            deterministic=False,
+            allow_tf32=False,
+        )
+
+    gems_assert_close(res_out.cpu(), ref_out.cpu(), dtype)