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
|
import operator_benchmark as op_bench
import torch
import random
from typing import List
"""Microbenchmarks for Stack operator"""
# Configs for PT stack operator
stack_configs_static_runtime = op_bench.config_list(
attr_names=['sizes', 'N'],
attrs=[
[(20, 40), 5],
[(1, 40), 5],
],
cross_product_configs={
'device': ['cpu', 'cuda'],
'dim': list(range(3))
},
tags=['static_runtime'],
)
stack_configs_short = op_bench.config_list(
attr_names=['sizes', 'N'],
attrs=[
[(1, 1, 1), 2], # noqa: E241
[(512, 512, 2), 2], # noqa: E241
[(128, 1024, 2), 2], # noqa: E241
],
cross_product_configs={
'device': ['cpu', 'cuda'],
'dim': list(range(4))
},
tags=['short'],
)
stack_configs_long = op_bench.config_list(
attr_names=['sizes', 'N'],
attrs=[
[(2**10, 2**10, 2), 2], # noqa: E241
[(2**10+1, 2**10-1, 2), 2], # noqa: E226,E241
[(2**10, 2**10, 2), 2], # noqa: E241
],
cross_product_configs={
'device': ['cpu', 'cuda'],
'dim': list(range(4))
},
tags=['long'],
)
# There is a different codepath on CUDA for >4 dimensions
stack_configs_multidim = op_bench.config_list(
attr_names=['sizes', 'N'],
attrs=[
[(2**6, 2**5, 2**2, 2**4, 2**5), 2], # noqa: E241
[(2**4, 2**5, 2**2, 2**4, 2**5), 8], # noqa: E241
[(2**3+1, 2**5-1, 2**2+1, 2**4-1, 2**5+1), 17], # noqa: E226,E241
],
cross_product_configs={
'device': ['cpu', 'cuda'],
'dim': list(range(6))
},
tags=['multidim'],
)
class StackBenchmark(op_bench.TorchBenchmarkBase):
def init(self, sizes, N, dim, device):
random.seed(42)
inputs = []
gen_sizes = []
if type(sizes) == list and N == -1:
gen_sizes = sizes
else:
for i in range(N):
gen_sizes.append([old_size() if callable(old_size) else old_size for old_size in sizes])
for s in gen_sizes:
inputs.append(torch.rand(s, device=device))
result = torch.rand(gen_sizes[0], device=device)
self.inputs = {
"result": result,
"inputs": inputs,
"dim": dim
}
self.set_module_name('stack')
def forward(self, result: torch.Tensor, inputs: List[torch.Tensor], dim: int):
return torch.stack(inputs, dim=dim, out=result)
op_bench.generate_pt_test(stack_configs_static_runtime +
stack_configs_short +
stack_configs_long +
stack_configs_multidim,
StackBenchmark)
if __name__ == "__main__":
op_bench.benchmark_runner.main()
|
