FlagGems/benchmark/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
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()