Skip to content

Commit c5d7a0d

Browse files
Schopenhauer-loves-Hegelfactnnclaudetengqm
committed
【KernelGen】Add cudnn_convolution operator (#1729)
* Add cudnn_convolution operator implementation, tests and benchmark Includes operators.yaml entry in alphabetical order. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Apply suggestions from code review Co-authored-by: Qiming Teng <tengqm@outlook.com> Signed-off-by: Zang Peiyu <166481866+factnn@users.noreply.github.com> * fix: apply black formatting to test_cudnn_convolution.py Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: remove to_reference for cudnn_convolution tests (CUDA-only op) torch.cudnn_convolution only supports CUDA backend, cannot run on CPU. Use the same GPU tensor for both reference and result. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: rename benchmark file to remove _perf suffix Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: move ref tensor to CPU before gems_assert_close torch.cudnn_convolution only runs on CUDA, so ref stays on GPU. gems_assert_close requires ref to be on CPU when TO_CPU is enabled. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: move both res and ref to CPU before comparison torch.cudnn_convolution is CUDA-only, so both tensors stay on GPU. gems_assert_close requires both on same device; move both to CPU. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: move cudnn_convolution to correct alphabetical position in operators.yaml Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Signed-off-by: Zang Peiyu <166481866+factnn@users.noreply.github.com> Co-authored-by: factnn <1050552884@qq.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Zang Peiyu <166481866+factnn@users.noreply.github.com> Co-authored-by: Qiming Teng <tengqm@outlook.com>
1 parent 0c9a160 commit c5d7a0d

6 files changed

Lines changed: 320 additions & 0 deletions

File tree

‎benchmark/test_cudnn_convolution.py‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
from typing import Generator
2+
3+
import pytest
4+
import torch
5+
6+
from . import base, consts, utils
7+
8+
9+
def cudnn_convolution_input_fn(shape, dtype, device):
10+
(
11+
batch,
12+
input_c,
13+
input_h,
14+
input_w,
15+
out_c,
16+
kernel_h,
17+
kernel_w,
18+
stride,
19+
padding,
20+
groups,
21+
) = shape
22+
input_shape = (batch, input_c, input_h, input_w)
23+
weight_shape = (out_c, input_c // groups, kernel_h, kernel_w)
24+
inp = utils.generate_tensor_input(input_shape, dtype, device)
25+
weight = utils.generate_tensor_input(weight_shape, dtype, device)
26+
27+
yield (
28+
inp,
29+
weight,
30+
[padding, padding],
31+
[stride, stride],
32+
[1, 1],
33+
groups,
34+
False,
35+
False,
36+
False,
37+
)
38+
39+
40+
class CudnnConv2dBenchmark(base.GenericBenchmark):
41+
def get_input_iter(self, dtype) -> Generator:
42+
shapes = [
43+
(32, 64, 128, 128, 32, 3, 3, 1, 2, 1),
44+
(32, 64, 210, 210, 16, 5, 5, 2, 1, 1),
45+
(16, 32, 12, 12, 24, 3, 3, 2, 1, 1),
46+
(16, 32, 24, 24, 24, 3, 3, 2, 2, 2),
47+
(16, 32, 24, 24, 24, 3, 3, 1, 2, 2),
48+
]
49+
50+
for shape in shapes:
51+
yield from self.input_fn(shape, dtype, self.device)
52+
53+
54+
@pytest.mark.cudnn_convolution
55+
def test_cudnn_convolution():
56+
bench = CudnnConv2dBenchmark(
57+
input_fn=cudnn_convolution_input_fn,
58+
op_name="cudnn_convolution",
59+
torch_op=torch.ops.aten.cudnn_convolution.default,
60+
dtypes=consts.FLOAT_DTYPES,
61+
)
62+
bench.run()

‎conf/operators.yaml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1349,6 +1349,18 @@ ops:
13491349
stages:
13501350
- stable: '2.0'
13511351
- removed: '3.0'
1352+
- id: cudnn_convolution
1353+
description: |
1354+
A wrapper for cuDNN convolution backend.
1355+
for:
1356+
- cudnn_convolution
1357+
labels:
1358+
- aten
1359+
- KernelGen
1360+
kind:
1361+
- NeuralNetwork
1362+
stages:
1363+
- beta: '5.1'
13521364
- id: cummax
13531365
description: |
13541366
Returns a named tuple `(values, indices)` where `values` is the cumulative maximum of elements

‎src/flag_gems/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,7 @@ def torch_ge(v):
184184
("copysign", copysign),
185185
("copysign.out", copysign_out),
186186
("count_nonzero", count_nonzero),
187+
("cudnn_convolution", cudnn_convolution),
187188
("cummax", cummax),
188189
("cummin", cummin),
189190
("cumprod", cumprod),

‎src/flag_gems/ops/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@
9090
from flag_gems.ops.cos import cos, cos_
9191
from flag_gems.ops.cosh import cosh, cosh_, cosh_out
9292
from flag_gems.ops.count_nonzero import count_nonzero
93+
from flag_gems.ops.cudnn_convolution import cudnn_convolution
9394
from flag_gems.ops.cummax import cummax
9495
from flag_gems.ops.cummin import cummin
9596
from flag_gems.ops.cumprod import cumprod, cumprod_
@@ -494,6 +495,7 @@
494495
"cosh_",
495496
"cosh_out",
496497
"count_nonzero",
498+
"cudnn_convolution",
497499
"cummax",
498500
"cummin",
499501
"cumprod",

‎src/flag_gems/ops/cudnn_convolution.py‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
import logging
2+
3+
from flag_gems.ops.conv1d import conv1d
4+
from flag_gems.ops.conv2d import conv2d
5+
from flag_gems.ops.conv3d import conv3d
6+
7+
logger = logging.getLogger(__name__)
8+
9+
10+
def cudnn_convolution(
11+
input,
12+
weight,
13+
padding,
14+
stride,
15+
dilation,
16+
groups,
17+
benchmark,
18+
deterministic,
19+
allow_tf32,
20+
):
21+
"""
22+
CUDNN convolution operation.
23+
24+
This is a lower-level convolution operation that does not include bias.
25+
It supports 1D, 2D, and 3D convolutions based on the input dimensionality.
26+
27+
Args:
28+
input: Input tensor of shape (N, C_in, *spatial_dims)
29+
weight: Weight tensor of shape (C_out, C_in/groups, *kernel_dims)
30+
padding: Padding for each spatial dimension
31+
stride: Stride for each spatial dimension
32+
dilation: Dilation for each spatial dimension
33+
groups: Number of groups for grouped convolution
34+
benchmark: cuDNN benchmark flag (ignored in Triton implementation)
35+
deterministic: cuDNN deterministic flag (ignored in Triton implementation)
36+
allow_tf32: Allow TF32 computation flag (ignored in Triton implementation)
37+
38+
Returns:
39+
Output tensor after convolution
40+
"""
41+
logger.debug("GEMS CUDNN_CONVOLUTION")
42+
43+
ndim = input.ndim - 2
44+
45+
# Extract values from lists if they are lists (cudnn_convolution receives lists)
46+
def extract_param(param, expected_len):
47+
if isinstance(param, (list, tuple)):
48+
if len(param) == expected_len:
49+
return param if expected_len > 1 else param[0]
50+
elif len(param) == 1:
51+
return param[0]
52+
return param
53+
54+
if ndim == 1:
55+
# For 1D convolution, extract single values from lists
56+
stride_val = extract_param(stride, 1)
57+
padding_val = extract_param(padding, 1)
58+
dilation_val = extract_param(dilation, 1)
59+
return conv1d(
60+
input,
61+
weight,
62+
bias=None,
63+
stride=stride_val,
64+
padding=padding_val,
65+
dilation=dilation_val,
66+
groups=groups,
67+
)
68+
elif ndim == 2:
69+
return conv2d(
70+
input,
71+
weight,
72+
bias=None,
73+
stride=stride,
74+
padding=padding,
75+
dilation=dilation,
76+
groups=groups,
77+
)
78+
elif ndim == 3:
79+
return conv3d(
80+
input,
81+
weight,
82+
bias=None,
83+
stride=stride,
84+
padding=padding,
85+
dilation=dilation,
86+
groups=groups,
87+
)
88+
else:
89+
raise ValueError(
90+
f"cudnn_convolution only supports 1D, 2D, and 3D convolutions, "
91+
f"got input with {ndim} spatial dimensions"
92+
)

‎tests/test_cudnn_convolution.py‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
import pytest
2+
import torch
3+
4+
import flag_gems
5+
6+
from .accuracy_utils import gems_assert_close
7+
8+
SHAPE_CUDNN_CONV2D = [
9+
((1, 2, 5, 5), (1, 2, 3, 3), 1),
10+
((2, 3, 9, 9), (1, 3, 3, 3), 1),
11+
((32, 8, 8, 8), (32, 8, 2, 2), 1),
12+
]
13+
14+
15+
@pytest.mark.cudnn_convolution
16+
@pytest.mark.parametrize("shape, kernel, groups", SHAPE_CUDNN_CONV2D)
17+
@pytest.mark.parametrize("stride", [1, 2])
18+
@pytest.mark.parametrize("padding", [0, 1])
19+
@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
20+
@pytest.mark.parametrize("dilation", [1, 2])
21+
def test_cudnn_convolution_2d(
22+
shape, kernel, stride, padding, groups, dtype, dilation, monkeypatch
23+
):
24+
if flag_gems.vendor_name == "mthreads" and dtype == torch.float16:
25+
monkeypatch.setenv("MUSA_ENABLE_SQMMA", "1")
26+
27+
inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
28+
weight = torch.randn(kernel, dtype=dtype, device=flag_gems.device)
29+
30+
ref_out = torch.cudnn_convolution(
31+
inp,
32+
weight,
33+
padding=[padding, padding],
34+
stride=[stride, stride],
35+
dilation=[dilation, dilation],
36+
groups=groups,
37+
benchmark=False,
38+
deterministic=False,
39+
allow_tf32=False,
40+
)
41+
42+
with flag_gems.use_gems():
43+
res_out = torch.cudnn_convolution(
44+
inp,
45+
weight,
46+
padding=[padding, padding],
47+
stride=[stride, stride],
48+
dilation=[dilation, dilation],
49+
groups=groups,
50+
benchmark=False,
51+
deterministic=False,
52+
allow_tf32=False,
53+
)
54+
55+
gems_assert_close(res_out.cpu(), ref_out.cpu(), dtype)
56+
57+
58+
SHAPE_CUDNN_CONV1D = [
59+
((32, 2, 4), (17, 2, 2)),
60+
((32, 15, 6), (17, 15, 2)),
61+
((64, 64, 64), (128, 64, 7)),
62+
]
63+
64+
65+
@pytest.mark.cudnn_convolution
66+
@pytest.mark.parametrize("shape, kernel", SHAPE_CUDNN_CONV1D)
67+
@pytest.mark.parametrize("stride", [1, 2])
68+
@pytest.mark.parametrize("padding", [0, 1])
69+
@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
70+
def test_cudnn_convolution_1d(shape, kernel, stride, padding, dtype, monkeypatch):
71+
if flag_gems.vendor_name == "mthreads" and dtype == torch.float16:
72+
monkeypatch.setenv("MUSA_ENABLE_SQMMA", "1")
73+
74+
inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
75+
weight = torch.randn(kernel, dtype=dtype, device=flag_gems.device)
76+
77+
ref_out = torch.cudnn_convolution(
78+
inp,
79+
weight,
80+
padding=[padding],
81+
stride=[stride],
82+
dilation=[1],
83+
groups=1,
84+
benchmark=False,
85+
deterministic=False,
86+
allow_tf32=False,
87+
)
88+
89+
with flag_gems.use_gems():
90+
res_out = torch.cudnn_convolution(
91+
inp,
92+
weight,
93+
padding=[padding],
94+
stride=[stride],
95+
dilation=[1],
96+
groups=1,
97+
benchmark=False,
98+
deterministic=False,
99+
allow_tf32=False,
100+
)
101+
102+
gems_assert_close(res_out.cpu(), ref_out.cpu(), dtype)
103+
104+
105+
SHAPE_CUDNN_CONV3D = [
106+
((1, 2, 5, 5, 5), (1, 2, 3, 3, 3), 1),
107+
((2, 3, 9, 9, 9), (1, 3, 3, 3, 3), 1),
108+
]
109+
110+
111+
@pytest.mark.cudnn_convolution
112+
@pytest.mark.parametrize("shape, kernel, groups", SHAPE_CUDNN_CONV3D)
113+
@pytest.mark.parametrize("stride", [1, 2])
114+
@pytest.mark.parametrize("padding", [0, 1])
115+
@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
116+
@pytest.mark.parametrize("dilation", [1, 2])
117+
def test_cudnn_convolution_3d(
118+
shape, kernel, stride, padding, groups, dtype, dilation, monkeypatch
119+
):
120+
if flag_gems.vendor_name == "mthreads" and dtype == torch.float16:
121+
monkeypatch.setenv("MUSA_ENABLE_SQMMA", "1")
122+
123+
inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
124+
weight = torch.randn(kernel, dtype=dtype, device=flag_gems.device)
125+
126+
ref_out = torch.cudnn_convolution(
127+
inp,
128+
weight,
129+
padding=[padding, padding, padding],
130+
stride=[stride, stride, stride],
131+
dilation=[dilation, dilation, dilation],
132+
groups=groups,
133+
benchmark=False,
134+
deterministic=False,
135+
allow_tf32=False,
136+
)
137+
138+
with flag_gems.use_gems():
139+
res_out = torch.cudnn_convolution(
140+
inp,
141+
weight,
142+
padding=[padding, padding, padding],
143+
stride=[stride, stride, stride],
144+
dilation=[dilation, dilation, dilation],
145+
groups=groups,
146+
benchmark=False,
147+
deterministic=False,
148+
allow_tf32=False,
149+
)
150+
151+
gems_assert_close(res_out.cpu(), ref_out.cpu(), dtype)

0 commit comments

Comments
 (0)