cuda-kernel-study/06_reduction/benchmarks.py at main · junstar92/cuda-kernel-study · 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
import argparse

import torch
import triton
from reduction import (
    reduction_atomic,
    reduction_atomic_shm,
    reduction_cg,
    reduction_cg_unroll_2,
    reduction_cub,
    reduction_shm,
    reduction_shuffle,
    reduction_unroll,
)

DEVICE = torch.device("cuda", index=0)


@triton.testing.perf_report(
    triton.testing.Benchmark(
        x_names=["N"],
        x_vals=[2**i for i in range(16, 26, 1)],
        x_log=True,
        line_arg="provider",
        line_vals=[
            "torch",
            "atomic_global",
            "atomic_shm",
            "shm",
            "shuffle",
            "unroll",
            "cg",
            "cg_unroll_2",
            "cub",
        ],
        line_names=[
            "Torch",
            "Atomic (global)",
            "Atomic (shared)",
            "Shared",
            "Shuffle",
            "Unroll",
            "CG",
            "CG (unroll 2)",
            "CUB",
        ],
        styles=[
            ("red", "-"),
            ("blue", "--"),
            ("green", "--"),
            ("purple", "--"),
            ("orange", "--"),
            ("brown", "--"),
            ("pink", "--"),
            ("gray", "--"),
            ("olive", "--"),
            ("cyan", "--"),
        ],
        ylabel="Time (ms)",
        plot_name="Reduction Performance (ms)",
        args={},
    )
)
def benchmark(N, provider, dtype):
    if dtype.is_floating_point:
        x = torch.randn(N, dtype=dtype)
    else:
        x = torch.randint(-100, 100, (N,), dtype=dtype)
    out = torch.tensor(0, dtype=dtype)

    quantiles = [0.5, 0.1, 0.9]
    if provider == "torch":
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: x.sum(), quantiles=quantiles)
    if provider == "atomic_global":
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: reduction_atomic(out, x), quantiles=quantiles)
    if provider == "atomic_shm":
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: reduction_atomic_shm(out, x), quantiles=quantiles)
    if provider == "shm":
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: reduction_shm(out, x), quantiles=quantiles)
    if provider == "shuffle":
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: reduction_shuffle(out, x), quantiles=quantiles)
    if provider == "unroll":
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: reduction_unroll(out, x), quantiles=quantiles)
    if provider == "cg":
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: reduction_cg(out, x), quantiles=quantiles)
    if provider == "cg_unroll_2":
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: reduction_cg_unroll_2(out, x), quantiles=quantiles)
    if provider == "cub":
        ms, min_ms, max_ms = triton.testing.do_bench(lambda: reduction_cub(out, x), quantiles=quantiles)

    return ms, min_ms, max_ms


class DtypeAction(argparse.Action):
    def __init__(self, option_strings, dest, nargs=None, **kwargs):
        if nargs is not None:
            raise ValueError("nargs not allowed")
        super().__init__(option_strings, dest, **kwargs)

    def __call__(self, parser, namespace, values, option_string=None):
        assert isinstance(values, str), "values must be string"
        if values == "float32":
            dtype = torch.float32
        elif values == "float16":
            dtype = torch.float16
        elif values == "bfloat16":
            dtype = torch.bfloat16
        elif values == "int32":
            dtype = torch.int32
        elif values == "int64":
            dtype = torch.int64
        else:
            assert False, f"unsupported type: {values}"

        setattr(namespace, self.dest, dtype)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--dtype",
        default=torch.float32,
        choices=["float32", "double", "int32"],
        action=DtypeAction,
    )
    parser.add_argument("--save", action="store_true", help="save results")
    args = parser.parse_args()

    torch.manual_seed(0)
    torch.set_default_device(DEVICE)

    benchmark.run(print_data=True, save_path="results" if args.save else None, dtype=args.dtype)