xref: /aosp_15_r20/external/pytorch/benchmarks/operator_benchmark/pt/add_test.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1import operator_benchmark as op_bench
2
3import torch
4
5
6"""Microbenchmarks for add_ operator. Supports both Caffe2/PyTorch."""
7
8# Configs for PT add operator
9add_long_configs = op_bench.cross_product_configs(
10    M=[8, 128], N=[32, 64], K=[256, 512], device=["cpu", "cuda"], tags=["long"]
11)
12
13
14add_short_configs = op_bench.config_list(
15    attr_names=["M", "N", "K"],
16    attrs=[
17        [1, 1, 1],
18        [64, 64, 64],
19        [64, 64, 128],
20    ],
21    cross_product_configs={
22        "device": ["cpu", "cuda"],
23    },
24    tags=["short"],
25)
26
27
28class AddBenchmark(op_bench.TorchBenchmarkBase):
29    def init(self, M, N, K, device):
30        self.inputs = {
31            "input_one": torch.rand(
32                M, N, K, device=device, requires_grad=self.auto_set()
33            ),
34            "input_two": torch.rand(
35                M, N, K, device=device, requires_grad=self.auto_set()
36            ),
37        }
38        self.set_module_name("add")
39
40    def forward(self, input_one, input_two):
41        return torch.add(input_one, input_two)
42
43
44# The generated test names based on add_short_configs will be in the following pattern:
45# add_M8_N16_K32_devicecpu
46# add_M8_N16_K32_devicecpu_bwdall
47# add_M8_N16_K32_devicecpu_bwd1
48# add_M8_N16_K32_devicecpu_bwd2
49# ...
50# Those names can be used to filter tests.
51
52op_bench.generate_pt_test(add_long_configs + add_short_configs, AddBenchmark)
53op_bench.generate_pt_gradient_test(add_long_configs + add_short_configs, AddBenchmark)
54
55
56"""Mircobenchmark for addmm operator."""
57
58
59class AddmmBenchmark(op_bench.TorchBenchmarkBase):
60    def init(self, M, N, K, device):
61        self.inputs = {
62            "input_one": torch.rand(M, K, device=device, requires_grad=self.auto_set()),
63            "mat1": torch.rand(M, N, device=device, requires_grad=self.auto_set()),
64            "mat2": torch.rand(N, K, device=device, requires_grad=self.auto_set()),
65        }
66        self.set_module_name("addmm")
67
68    def forward(self, input_one, mat1, mat2):
69        return torch.addmm(input_one, mat1, mat2)
70
71
72op_bench.generate_pt_test(add_long_configs + add_short_configs, AddmmBenchmark)
73op_bench.generate_pt_gradient_test(add_long_configs + add_short_configs, AddmmBenchmark)
74
75
76"""Mircobenchmark for addr operator."""
77
78
79class AddrBenchmark(op_bench.TorchBenchmarkBase):
80    def init(self, M, N, device, dtype):
81        self.inputs = {
82            "input_one": torch.rand(
83                (M, N), device=device, requires_grad=self.auto_set(), dtype=dtype
84            ),
85            "vec1": torch.rand(
86                (M,), device=device, requires_grad=self.auto_set(), dtype=dtype
87            ),
88            "vec2": torch.rand(
89                (N,), device=device, requires_grad=self.auto_set(), dtype=dtype
90            ),
91        }
92        self.set_module_name("addr")
93
94    def forward(self, input_one, vec1, vec2):
95        return torch.addr(input_one, vec1, vec2)
96
97
98addr_configs = op_bench.cross_product_configs(
99    M=[8, 256],
100    N=[256, 16],
101    device=["cpu", "cuda"],
102    dtype=[torch.double, torch.half],
103    tags=["addr"],
104)
105
106op_bench.generate_pt_test(addr_configs, AddrBenchmark)
107op_bench.generate_pt_gradient_test(addr_configs, AddrBenchmark)
108
109
110"""Mircobenchmark for addbmm operator."""
111
112
113class AddbmmBenchmark(op_bench.TorchBenchmarkBase):
114    def init(self, B, M, N, K, device):
115        self.inputs = {
116            "input_one": torch.rand(
117                (M, N), device=device, requires_grad=self.auto_set()
118            ),
119            "batch1": torch.rand(
120                (B, M, K), device=device, requires_grad=self.auto_set()
121            ),
122            "batch2": torch.rand(
123                (
124                    B,
125                    K,
126                    N,
127                ),
128                device=device,
129                requires_grad=self.auto_set(),
130            ),
131        }
132        self.set_module_name("addbmm")
133
134    def forward(self, input_one, batch1, batch2):
135        return torch.addbmm(input_one, batch1, batch2)
136
137
138addbmm_configs = op_bench.cross_product_configs(
139    B=[2, 100],
140    M=[8, 256],
141    N=[256, 16],
142    K=[15, 16],
143    device=["cpu", "cuda"],
144    tags=["addbmm"],
145)
146
147op_bench.generate_pt_test(addbmm_configs, AddbmmBenchmark)
148op_bench.generate_pt_gradient_test(addbmm_configs, AddbmmBenchmark)
149
150if __name__ == "__main__":
151    op_bench.benchmark_runner.main()
152