FlagGems/tests/test_cudnn_convolution.py at cbf2212f611f166634c846d30b920c0fe1e26e76 · Schopenhauer-loves-Hegel/FlagGems · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
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)