mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
dd3a77bc96
Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/105928 Approved by: https://github.com/albanD
50 lines
1.3 KiB
Python
50 lines
1.3 KiB
Python
import benchmark_caffe2 as op_bench_c2
|
|
from benchmark_caffe2 import Caffe2BenchmarkBase # noqa: F401
|
|
from caffe2.python import core
|
|
|
|
import operator_benchmark as op_bench
|
|
|
|
|
|
"""Microbenchmarks for element-wise Add operator. Supports both Caffe2/PyTorch."""
|
|
|
|
# Configs for C2 add operator
|
|
add_long_configs = op_bench.cross_product_configs(
|
|
M=[8, 64, 128],
|
|
N=range(2, 10, 3),
|
|
K=[2**x for x in range(0, 3)],
|
|
dtype=["int", "float"],
|
|
tags=["long"],
|
|
)
|
|
|
|
|
|
add_short_configs = op_bench.config_list(
|
|
attrs=[
|
|
[8, 16, 32, "int"],
|
|
[16, 16, 64, "float"],
|
|
[64, 64, 128, "int"],
|
|
],
|
|
attr_names=["M", "N", "K", "dtype"],
|
|
tags=["short"],
|
|
)
|
|
|
|
|
|
class AddBenchmark(op_bench_c2.Caffe2BenchmarkBase):
|
|
def init(self, M, N, K, dtype):
|
|
self.input_one = self.tensor([M, N, K], dtype)
|
|
self.input_two = self.tensor([M, N, K], dtype)
|
|
self.output = self.tensor([M, N, K], dtype)
|
|
self.set_module_name("add")
|
|
|
|
def forward(self):
|
|
op = core.CreateOperator(
|
|
"Add", [self.input_one, self.input_two], self.output, **self.args
|
|
)
|
|
return op
|
|
|
|
|
|
op_bench_c2.generate_c2_test(add_long_configs + add_short_configs, AddBenchmark)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
op_bench.benchmark_runner.main()
|