1import operator_benchmark as op_bench 2 3import torch 4 5 6"""Microbenchmarks for sum reduction operator.""" 7 8# Configs for PT add operator 9sum_configs = op_bench.cross_product_configs( 10 R=[64, 256], # Length of reduced dimension 11 V=[32, 512], # Length of other dimension 12 dim=[0, 1], 13 contiguous=[True, False], 14 device=["cpu", "cuda"], 15 tags=["short"], 16) + op_bench.cross_product_configs( 17 R=[1024, 8192], 18 V=[512, 1024], 19 dim=[0, 1], 20 contiguous=[True, False], 21 device=["cpu", "cuda"], 22 tags=["long"], 23) 24 25 26class SumBenchmark(op_bench.TorchBenchmarkBase): 27 def init(self, R, V, dim, contiguous, device): 28 shape = (R, V) if dim == 0 else (V, R) 29 tensor = torch.rand(shape, device=device) 30 31 if not contiguous: 32 storage = torch.empty([s * 2 for s in shape], device=device) 33 storage[::2, ::2] = tensor 34 self.input_tensor = storage[::2, ::2] 35 else: 36 self.input_tensor = tensor 37 38 self.inputs = {"input_tensor": self.input_tensor, "dim": dim} 39 self.set_module_name("sum") 40 41 def forward(self, input_tensor, dim: int): 42 return input_tensor.sum(dim=dim) 43 44 45op_bench.generate_pt_test(sum_configs, SumBenchmark) 46 47if __name__ == "__main__": 48 op_bench.benchmark_runner.main() 49