Skip to content

Commit 1dc60b7

Browse files
factnnclaude
factnn
and
claude
committed
Add cudnn_convolution operator implementation, tests and benchmark
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 93e4d30 commit 1dc60b7

4 files changed

Lines changed: 256 additions & 0 deletions

File tree

‎src/flag_gems/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,7 @@ def torch_ge(v):
151151
("copysign", copysign),
152152
("copysign.out", copysign_out),
153153
("count_nonzero", count_nonzero),
154+
("cudnn_convolution", cudnn_convolution),
154155
("cummax", cummax),
155156
("cummin", cummin),
156157
("cumsum", cumsum),

‎src/flag_gems/ops/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@
8282
from flag_gems.ops.cos import cos, cos_
8383
from flag_gems.ops.cosh import cosh, cosh_, cosh_out
8484
from flag_gems.ops.count_nonzero import count_nonzero
85+
from flag_gems.ops.cudnn_convolution import cudnn_convolution
8586
from flag_gems.ops.cummax import cummax
8687
from flag_gems.ops.cummin import cummin
8788
from flag_gems.ops.cumsum import cumsum, cumsum_out, normed_cumsum
@@ -425,6 +426,7 @@
425426
"cosh_",
426427
"cosh_out",
427428
"count_nonzero",
429+
"cudnn_convolution",
428430
"cummax",
429431
"cummin",
430432
"cumsum",

‎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: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
import pytest
2+
import torch
3+
4+
import flag_gems
5+
6+
from .accuracy_utils import gems_assert_close, to_reference
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_accuracy_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+
ref_inp = to_reference(inp)
29+
torch.backends.cudnn.allow_tf32 = False
30+
weight = torch.randn(kernel, dtype=dtype, device=flag_gems.device)
31+
ref_weight = to_reference(weight)
32+
33+
ref_out = torch.cudnn_convolution(
34+
ref_inp,
35+
ref_weight,
36+
padding=[padding, padding],
37+
stride=[stride, stride],
38+
dilation=[dilation, dilation],
39+
groups=groups,
40+
benchmark=False,
41+
deterministic=False,
42+
allow_tf32=False,
43+
)
44+
45+
with flag_gems.use_gems():
46+
res_out = torch.cudnn_convolution(
47+
inp,
48+
weight,
49+
padding=[padding, padding],
50+
stride=[stride, stride],
51+
dilation=[dilation, dilation],
52+
groups=groups,
53+
benchmark=False,
54+
deterministic=False,
55+
allow_tf32=False,
56+
)
57+
58+
gems_assert_close(res_out, ref_out, dtype)
59+
60+
61+
SHAPE_CUDNN_CONV1D = [
62+
((32, 2, 4), (17, 2, 2)),
63+
((32, 15, 6), (17, 15, 2)),
64+
((64, 64, 64), (128, 64, 7)),
65+
]
66+
67+
68+
@pytest.mark.cudnn_convolution
69+
@pytest.mark.parametrize("shape, kernel", SHAPE_CUDNN_CONV1D)
70+
@pytest.mark.parametrize("stride", [1, 2])
71+
@pytest.mark.parametrize("padding", [0, 1])
72+
@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
73+
def test_accuracy_cudnn_convolution_1d(
74+
shape, kernel, stride, padding, dtype, monkeypatch
75+
):
76+
if flag_gems.vendor_name == "mthreads" and dtype == torch.float16:
77+
monkeypatch.setenv("MUSA_ENABLE_SQMMA", "1")
78+
79+
inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
80+
ref_inp = to_reference(inp)
81+
weight = torch.randn(kernel, dtype=dtype, device=flag_gems.device)
82+
ref_weight = to_reference(weight)
83+
84+
ref_out = torch.cudnn_convolution(
85+
ref_inp,
86+
ref_weight,
87+
padding=[padding],
88+
stride=[stride],
89+
dilation=[1],
90+
groups=1,
91+
benchmark=False,
92+
deterministic=False,
93+
allow_tf32=False,
94+
)
95+
96+
with flag_gems.use_gems():
97+
res_out = torch.cudnn_convolution(
98+
inp,
99+
weight,
100+
padding=[padding],
101+
stride=[stride],
102+
dilation=[1],
103+
groups=1,
104+
benchmark=False,
105+
deterministic=False,
106+
allow_tf32=False,
107+
)
108+
109+
gems_assert_close(res_out, ref_out, dtype)
110+
111+
112+
SHAPE_CUDNN_CONV3D = [
113+
((1, 2, 5, 5, 5), (1, 2, 3, 3, 3), 1),
114+
((2, 3, 9, 9, 9), (1, 3, 3, 3, 3), 1),
115+
]
116+
117+
118+
@pytest.mark.cudnn_convolution
119+
@pytest.mark.parametrize("shape, kernel, groups", SHAPE_CUDNN_CONV3D)
120+
@pytest.mark.parametrize("stride", [1, 2])
121+
@pytest.mark.parametrize("padding", [0, 1])
122+
@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
123+
@pytest.mark.parametrize("dilation", [1, 2])
124+
def test_accuracy_cudnn_convolution_3d(
125+
shape, kernel, stride, padding, groups, dtype, dilation, monkeypatch
126+
):
127+
if flag_gems.vendor_name == "mthreads" and dtype == torch.float16:
128+
monkeypatch.setenv("MUSA_ENABLE_SQMMA", "1")
129+
130+
inp = torch.randn(shape, dtype=dtype, device=flag_gems.device)
131+
ref_inp = to_reference(inp)
132+
torch.backends.cudnn.allow_tf32 = False
133+
weight = torch.randn(kernel, dtype=dtype, device=flag_gems.device)
134+
ref_weight = to_reference(weight)
135+
136+
ref_out = torch.cudnn_convolution(
137+
ref_inp,
138+
ref_weight,
139+
padding=[padding, padding, padding],
140+
stride=[stride, stride, stride],
141+
dilation=[dilation, dilation, dilation],
142+
groups=groups,
143+
benchmark=False,
144+
deterministic=False,
145+
allow_tf32=False,
146+
)
147+
148+
with flag_gems.use_gems():
149+
res_out = torch.cudnn_convolution(
150+
inp,
151+
weight,
152+
padding=[padding, padding, padding],
153+
stride=[stride, stride, stride],
154+
dilation=[dilation, dilation, dilation],
155+
groups=groups,
156+
benchmark=False,
157+
deterministic=False,
158+
allow_tf32=False,
159+
)
160+
161+
gems_assert_close(res_out, ref_out, dtype)

0 commit comments

Comments
 (0)