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
|
import operator_benchmark as op_bench
"""
Configs shared by multiple benchmarks
"""
def remove_cuda(config_list):
cuda_config = {'device': 'cuda'}
return [config for config in config_list if cuda_config not in config]
# Configs for conv-1d ops
conv_1d_configs_short = op_bench.config_list(
attr_names=[
'IC', 'OC', 'kernel', 'stride', 'N', 'L'
],
attrs=[
[128, 256, 3, 1, 1, 64],
[256, 256, 3, 2, 4, 64],
],
cross_product_configs={
'device': ['cpu', 'cuda'],
},
tags=['short']
)
conv_1d_configs_long = op_bench.cross_product_configs(
IC=[128, 512],
OC=[128, 512],
kernel=[3],
stride=[1, 2],
N=[8],
L=[128],
device=['cpu', 'cuda'],
tags=["long"]
)
# Configs for Conv2d and ConvTranspose1d
conv_2d_configs_short = op_bench.config_list(
attr_names=[
'IC', 'OC', 'kernel', 'stride', 'N', 'H', 'W', 'G', 'pad',
],
attrs=[
[256, 256, 3, 1, 1, 16, 16, 1, 0],
],
cross_product_configs={
'device': ['cpu', 'cuda'],
},
tags=['short']
)
conv_2d_configs_long = op_bench.cross_product_configs(
IC=[128, 256],
OC=[128, 256],
kernel=[3],
stride=[1, 2],
N=[4],
H=[32],
W=[32],
G=[1],
pad=[0],
device=['cpu', 'cuda'],
tags=["long"]
)
# Configs for Conv3d and ConvTranspose3d
conv_3d_configs_short = op_bench.config_list(
attr_names=[
'IC', 'OC', 'kernel', 'stride', 'N', 'D', 'H', 'W'
],
attrs=[
[64, 64, 3, 1, 8, 4, 16, 16],
],
cross_product_configs={
'device': ['cpu', 'cuda'],
},
tags=['short']
)
linear_configs_short = op_bench.config_list(
attr_names=["N", "IN", "OUT"],
attrs=[
[1, 1, 1],
[4, 256, 128],
[16, 512, 256],
],
cross_product_configs={
'device': ['cpu', 'cuda'],
},
tags=["short"]
)
linear_configs_long = op_bench.cross_product_configs(
N=[32, 64],
IN=[128, 512],
OUT=[64, 128],
device=['cpu', 'cuda'],
tags=["long"]
)
embeddingbag_short_configs = op_bench.cross_product_configs(
embeddingbags=[10, 120, 1000, 2300],
dim=[64],
mode=['sum'],
input_size=[8, 16, 64],
offset=[0],
sparse=[True, False],
include_last_offset=[True, False],
device=['cpu'],
tags=['short']
)
embedding_short_configs = op_bench.cross_product_configs(
num_embeddings=[10, 120, 1000, 2300],
embedding_dim=[64],
input_size=[8, 16, 64],
device=['cpu'],
tags=['short']
)
|
