xref: /aosp_15_r20/external/pytorch/benchmarks/operator_benchmark/pt/configs.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1import operator_benchmark as op_bench
2
3
4"""
5Configs shared by multiple benchmarks
6"""
7
8
9def remove_cuda(config_list):
10    cuda_config = {"device": "cuda"}
11    return [config for config in config_list if cuda_config not in config]
12
13
14# Configs for conv-1d ops
15conv_1d_configs_short = op_bench.config_list(
16    attr_names=["IC", "OC", "kernel", "stride", "N", "L"],
17    attrs=[
18        [128, 256, 3, 1, 1, 64],
19        [256, 256, 3, 2, 4, 64],
20    ],
21    cross_product_configs={
22        "device": ["cpu", "cuda"],
23    },
24    tags=["short"],
25)
26
27conv_1d_configs_long = op_bench.cross_product_configs(
28    IC=[128, 512],
29    OC=[128, 512],
30    kernel=[3],
31    stride=[1, 2],
32    N=[8],
33    L=[128],
34    device=["cpu", "cuda"],
35    tags=["long"],
36)
37
38convtranspose_1d_configs_short = op_bench.config_list(
39    attr_names=["IC", "OC", "kernel", "stride", "N", "L"],
40    attrs=[
41        [2016, 1026, 1024, 256, 1, 224],
42    ],
43    cross_product_configs={
44        "device": ["cpu", "cuda"],
45    },
46    tags=["short"],
47)
48
49# Configs for Conv2d and ConvTranspose1d
50conv_2d_configs_short = op_bench.config_list(
51    attr_names=[
52        "IC",
53        "OC",
54        "kernel",
55        "stride",
56        "N",
57        "H",
58        "W",
59        "G",
60        "pad",
61    ],
62    attrs=[
63        [256, 256, 3, 1, 1, 16, 16, 1, 0],
64    ],
65    cross_product_configs={
66        "device": ["cpu", "cuda"],
67    },
68    tags=["short"],
69)
70
71conv_2d_configs_long = op_bench.cross_product_configs(
72    IC=[128, 256],
73    OC=[128, 256],
74    kernel=[3],
75    stride=[1, 2],
76    N=[4],
77    H=[32],
78    W=[32],
79    G=[1],
80    pad=[0],
81    device=["cpu", "cuda"],
82    tags=["long"],
83)
84
85# Configs for Conv2dPointwise
86conv_2d_pw_configs_short = op_bench.config_list(
87    attr_names=[
88        "IC",
89        "OC",
90        "stride",
91        "N",
92        "H",
93        "W",
94        "G",
95        "pad",
96    ],
97    attrs=[
98        [256, 256, 1, 1, 16, 16, 1, 0],
99    ],
100    cross_product_configs={
101        "device": ["cpu", "cuda"],
102    },
103    tags=["short"],
104)
105
106conv_2d_pw_configs_long = op_bench.cross_product_configs(
107    IC=[128, 256],
108    OC=[128, 256],
109    stride=[1, 2],
110    N=[4],
111    H=[32],
112    W=[32],
113    G=[1],
114    pad=[0],
115    device=["cpu", "cuda"],
116    tags=["long"],
117)
118
119# Configs for Conv3d and ConvTranspose3d
120conv_3d_configs_short = op_bench.config_list(
121    attr_names=["IC", "OC", "kernel", "stride", "N", "D", "H", "W"],
122    attrs=[
123        [64, 64, 3, 1, 8, 4, 16, 16],
124    ],
125    cross_product_configs={
126        "device": ["cpu", "cuda"],
127    },
128    tags=["short"],
129)
130
131linear_configs_short = op_bench.config_list(
132    attr_names=["N", "IN", "OUT"],
133    attrs=[
134        [1, 1, 1],
135        [4, 256, 128],
136        [16, 512, 256],
137    ],
138    cross_product_configs={
139        "device": ["cpu", "cuda"],
140    },
141    tags=["short"],
142)
143
144
145linear_configs_long = op_bench.cross_product_configs(
146    N=[32, 64], IN=[128, 512], OUT=[64, 128], device=["cpu", "cuda"], tags=["long"]
147)
148
149embeddingbag_short_configs = op_bench.cross_product_configs(
150    embeddingbags=[10, 120, 1000, 2300],
151    dim=[64],
152    mode=["sum"],
153    input_size=[8, 16, 64],
154    offset=[0],
155    sparse=[True, False],
156    include_last_offset=[True, False],
157    device=["cpu"],
158    tags=["short"],
159)
160
161embedding_short_configs = op_bench.cross_product_configs(
162    num_embeddings=[10, 120, 1000, 2300],
163    embedding_dim=[64],
164    input_size=[8, 16, 64],
165    device=["cpu"],
166    tags=["short"],
167)
168