Skip to content

Commit 93e4d30

Browse files
tengqmtengqm
authored
Split accuracy unit test for convolution operators (#2611)
1 parent 81e4613 commit 93e4d30

5 files changed

Lines changed: 398 additions & 671 deletions

File tree

‎tests/test_conv1d.py‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
import pytest
2+
import torch
3+
4+
import flag_gems
5+
6+
from . import accuracy_utils as utils
7+
8+
SHAPE_CONV1D = [
9+
((32, 2, 4), (17, 2, 2)),
10+
((32, 15, 6), (17, 15, 2)),
11+
((64, 64, 64), (128, 64, 7)),
12+
# ((32, 16, 1024), (1024, 16, 8)),
13+
# ((32, 12, 9), (17, 12, 3)),
14+
# ((32, 6, 6), (64, 6, 2)),
15+
]
16+
17+
SHAPE_CONV1D_DILATION = [
18+
((32, 2, 16), (17, 2, 3)),
19+
((32, 15, 32), (17, 15, 3)),
20+
((64, 64, 64), (128, 64, 3)),
21+
]
22+
23+
24+
@pytest.mark.conv1d
25+
@pytest.mark.parametrize("shape, kernel", SHAPE_CONV1D)
26+
@pytest.mark.parametrize("stride", [2])
27+
@pytest.mark.parametrize("padding", [1])
28+
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
29+
def test_conv1d(monkeypatch, shape, kernel, stride, padding, dtype):
30+
if flag_gems.vendor_name == "mthreads" and dtype == torch.float16:
31+
monkeypatch.env("MUSA_ENABLE_SQMMA", "1")
32+
33+
inp = torch.randn(shape, dtype=dtype, device=flag_gems.device, requires_grad=True)
34+
ref_inp = utils.to_reference(inp, True)
35+
weight = torch.randn(kernel, dtype=dtype, device=flag_gems.device)
36+
ref_weight = utils.to_reference(weight, True)
37+
ref_out = torch.nn.functional.conv1d(
38+
ref_inp, ref_weight, bias=None, stride=stride, padding=padding, dilation=1
39+
)
40+
41+
res_out = flag_gems.conv1d(
42+
inp, weight, bias=None, stride=stride, padding=padding, dilation=1
43+
)
44+
utils.gems_assert_close(res_out, ref_out, dtype)
45+
46+
47+
@pytest.mark.conv1d_padding
48+
@pytest.mark.skipif(flag_gems.vendor_name == "kunlunxin", reason="RESULT TODOFIX")
49+
@pytest.mark.parametrize("shape, kernel", SHAPE_CONV1D)
50+
@pytest.mark.parametrize("stride", [1])
51+
@pytest.mark.parametrize("padding", ["valid", "same"])
52+
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
53+
def test_accuracy_conv1d_padding(monkeypatch, shape, kernel, stride, padding, dtype):
54+
if flag_gems.vendor_name == "mthreads" and dtype == torch.float16:
55+
monkeypatch.env("MUSA_ENABLE_SQMMA", "1")
56+
57+
inp = torch.randn(shape, dtype=dtype, device=flag_gems.device, requires_grad=True)
58+
ref_inp = utils.to_reference(inp, True)
59+
weight = torch.randn(kernel, dtype=dtype, device=flag_gems.device)
60+
ref_weight = utils.to_reference(weight, True)
61+
ref_out = torch.nn.functional.conv1d(
62+
ref_inp, ref_weight, bias=None, stride=stride, padding=padding, dilation=1
63+
)
64+
65+
res_out = flag_gems.conv1d(
66+
inp, weight, bias=None, stride=stride, padding=padding, dilation=1
67+
)
68+
utils.gems_assert_close(res_out, ref_out, dtype)
69+
70+
71+
@pytest.mark.conv1d
72+
@pytest.mark.parametrize("shape, kernel", SHAPE_CONV1D_DILATION)
73+
@pytest.mark.parametrize("stride", [1])
74+
@pytest.mark.parametrize("padding", [0, 2])
75+
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
76+
@pytest.mark.parametrize("dilation", [1, 2, (1,), (2,)])
77+
def test_accuracy_conv1d_dilation(shape, kernel, stride, padding, dtype, dilation):
78+
"""Test conv1d with various dilation values, including tuple form.
79+
80+
This specifically tests the fix where conv1d must properly convert dilation
81+
to a 2D tuple before delegating to conv2d. Previously, passing dilation as
82+
a single-element tuple (e.g., (1,)) would cause a ValueError in conv2d.
83+
"""
84+
inp = torch.randn(shape, dtype=dtype, device=flag_gems.device, requires_grad=True)
85+
ref_inp = utils.to_reference(inp, True)
86+
weight = torch.randn(kernel, dtype=dtype, device=flag_gems.device)
87+
ref_weight = utils.to_reference(weight, True)
88+
89+
ref_out = torch.nn.functional.conv1d(
90+
ref_inp,
91+
ref_weight,
92+
bias=None,
93+
stride=stride,
94+
padding=padding,
95+
dilation=dilation,
96+
)
97+
98+
res_out = flag_gems.conv1d(
99+
inp, weight, bias=None, stride=stride, padding=padding, dilation=dilation
100+
)
101+
102+
utils.gems_assert_close(res_out, ref_out, dtype)

‎tests/test_conv2d.py‎

Lines changed: 185 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,185 @@
1+
import pytest
2+
import torch
3+
4+
import flag_gems
5+
6+
from . import accuracy_utils as utils
7+
8+
SHAPE_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+
# ((2, 2, 3, 3), (1, 2, 2, 2), 1),
13+
# ((18, 16, 4, 4), (16, 16, 2, 2), 1),
14+
# ((9, 16, 4, 4), (128, 4, 2, 2), 4),
15+
# ((32, 16, 8, 8), (32, 4, 4, 4), 4),
16+
# ((18, 16, 4, 4), (16, 8, 2, 2), 2),
17+
# ((9, 16, 4, 4), (128, 8, 2, 2), 2),
18+
# ((32, 8, 8, 8), (32, 8, 3, 3), 1),
19+
# ((18, 16, 5, 5), (16, 16, 3, 3), 1),
20+
# ((9, 16, 7, 7), (128, 4, 3, 3), 4),
21+
# ((32, 16, 9, 9), (32, 4, 5, 5), 4),
22+
# ((18, 16, 11, 11), (16, 8, 3, 3), 2),
23+
# ((9, 16, 6, 6), (128, 8, 3, 3), 2),
24+
]
25+
26+
27+
@pytest.mark.conv2d
28+
@pytest.mark.parametrize("shape, kernel,groups", SHAPE_CONV2D)
29+
@pytest.mark.parametrize("stride", [1, 2])
30+
@pytest.mark.parametrize("padding", [0, 1])
31+
@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
32+
@pytest.mark.parametrize("dilation", [1, 2])
33+
@pytest.mark.parametrize("bias", [True, False])
34+
def test_conv2d(
35+
monkeypatch, shape, kernel, stride, padding, groups, dtype, dilation, bias
36+
):
37+
if flag_gems.vendor_name == "mthreads" and dtype == torch.float16:
38+
monkeypatch.env("MUSA_ENABLE_SQMMA", "1")
39+
40+
if flag_gems.vendor_name == "hygon":
41+
monkeypatch.env("TRITON_HIP_USE_NEW_STREAM_PIPELINE", "0")
42+
43+
inp = torch.randn(shape, dtype=dtype, device=flag_gems.device, requires_grad=True)
44+
ref_inp = utils.to_reference(inp, True)
45+
torch.backends.cudnn.allow_tf32 = False
46+
weight = torch.randn(
47+
kernel, dtype=dtype, device=flag_gems.device, requires_grad=True
48+
)
49+
if bias is True:
50+
bias = torch.randn(
51+
[weight.shape[0]], dtype=dtype, device=flag_gems.device, requires_grad=True
52+
)
53+
bias_ref = utils.to_reference(bias, True)
54+
else:
55+
bias = None
56+
bias_ref = None
57+
58+
ref_weight = utils.to_reference(weight, True)
59+
ref_out = torch.nn.functional.conv2d(
60+
ref_inp,
61+
ref_weight,
62+
bias=bias_ref,
63+
groups=groups,
64+
stride=stride,
65+
padding=padding,
66+
dilation=dilation,
67+
).to(dtype)
68+
69+
res_out = flag_gems.conv2d(
70+
inp,
71+
weight,
72+
bias=bias,
73+
groups=groups,
74+
stride=stride,
75+
padding=padding,
76+
dilation=dilation,
77+
)
78+
79+
utils.gems_assert_close(res_out, ref_out, dtype)
80+
81+
out_grad = torch.randn_like(ref_out).to(flag_gems.device)
82+
83+
ref_grad = utils.to_reference(out_grad, True)
84+
if bias is not None:
85+
ref_in_grad, ref_weight_grad, ref_bias_grad = torch.autograd.grad(
86+
ref_out, (ref_inp, ref_weight, bias_ref), ref_grad
87+
)
88+
res_in_grad, res_weight_grad, res_bias_grad = torch.autograd.grad(
89+
res_out, (inp, weight, bias), out_grad
90+
)
91+
else:
92+
ref_in_grad, ref_weight_grad = torch.autograd.grad(
93+
ref_out, (ref_inp, ref_weight), ref_grad
94+
)
95+
res_in_grad, res_weight_grad = torch.autograd.grad(
96+
res_out, (inp, weight), out_grad
97+
)
98+
99+
utils.gems_assert_close(res_in_grad, ref_in_grad, dtype, reduce_dim=weight.shape[2])
100+
101+
utils.gems_assert_close(
102+
res_weight_grad, ref_weight_grad, dtype, reduce_dim=weight.shape[0]
103+
)
104+
if bias is not None:
105+
utils.gems_assert_close(res_bias_grad, ref_bias_grad, dtype)
106+
107+
108+
@pytest.mark.conv2d_padding
109+
@pytest.mark.skipif(flag_gems.vendor_name == "hygon", reason="RESULT TODOFIX")
110+
@pytest.mark.skipif(flag_gems.vendor_name == "kunlunxin", reason="RESULT TODOFIX")
111+
@pytest.mark.parametrize("shape, kernel,groups", SHAPE_CONV2D)
112+
@pytest.mark.parametrize("stride", [1])
113+
@pytest.mark.parametrize("padding", ["valid", "same"])
114+
@pytest.mark.parametrize("dtype", [torch.float16, torch.float32])
115+
@pytest.mark.parametrize("dilation", [1, 2])
116+
@pytest.mark.parametrize("bias", [True, False])
117+
def test_conv2d_padding(
118+
monkeypatch, shape, kernel, stride, padding, groups, dtype, dilation, bias
119+
):
120+
if flag_gems.vendor_name == "mthreads" and dtype == torch.float16:
121+
monkeypatch.env("MUSA_ENABLE_SQMMA", "1")
122+
123+
inp = torch.randn(shape, dtype=dtype, device=flag_gems.device, requires_grad=True)
124+
ref_inp = utils.to_reference(inp, True)
125+
torch.backends.cudnn.allow_tf32 = False
126+
weight = torch.randn(
127+
kernel, dtype=dtype, device=flag_gems.device, requires_grad=True
128+
)
129+
if bias is True:
130+
bias = torch.randn(
131+
[weight.shape[0]], dtype=dtype, device=flag_gems.device, requires_grad=True
132+
)
133+
bias_ref = utils.to_reference(bias, True)
134+
else:
135+
bias = None
136+
bias_ref = None
137+
138+
ref_weight = utils.to_reference(weight, True)
139+
ref_out = torch.nn.functional.conv2d(
140+
ref_inp,
141+
ref_weight,
142+
bias=bias_ref,
143+
groups=groups,
144+
stride=stride,
145+
padding=padding,
146+
dilation=dilation,
147+
).to(dtype)
148+
149+
res_out = flag_gems.conv2d(
150+
inp,
151+
weight,
152+
bias=bias,
153+
groups=groups,
154+
stride=stride,
155+
padding=padding,
156+
dilation=dilation,
157+
)
158+
159+
utils.gems_assert_close(res_out, ref_out, dtype)
160+
161+
out_grad = torch.randn_like(ref_out).to(flag_gems.device)
162+
163+
ref_grad = utils.to_reference(out_grad, True)
164+
if bias is not None:
165+
ref_in_grad, ref_weight_grad, ref_bias_grad = torch.autograd.grad(
166+
ref_out, (ref_inp, ref_weight, bias_ref), ref_grad
167+
)
168+
res_in_grad, res_weight_grad, res_bias_grad = torch.autograd.grad(
169+
res_out, (inp, weight, bias), out_grad
170+
)
171+
else:
172+
ref_in_grad, ref_weight_grad = torch.autograd.grad(
173+
ref_out, (ref_inp, ref_weight), ref_grad
174+
)
175+
res_in_grad, res_weight_grad = torch.autograd.grad(
176+
res_out, (inp, weight), out_grad
177+
)
178+
179+
utils.gems_assert_close(res_in_grad, ref_in_grad, dtype, reduce_dim=weight.shape[2])
180+
181+
utils.gems_assert_close(
182+
res_weight_grad, ref_weight_grad, dtype, reduce_dim=weight.shape[0]
183+
)
184+
if bias is not None:
185+
utils.gems_assert_close(res_bias_grad, ref_bias_grad, dtype)

0 commit comments

Comments
 (0)