unify unary ops benchmark (#28913)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/28913
as title
Test Plan:
```
buck run mode/opt //caffe2/benchmarks/operator_benchmark/pt:unary_test
# ----------------------------------------
# PyTorch/Caffe2 Operator Micro-benchmarks
# ----------------------------------------
# Tag : short
# Benchmarking PyTorch: abs
# Mode: Eager
# Name: abs_M512_N512_cpu
# Input: M: 512, N: 512, device: cpu
Forward Execution Time (us) : 90.233
...
Reviewed By: hl475
Differential Revision: D18231641
fbshipit-source-id: 3093db47d0356b927768f15dc63af6ad8aadd430
diff --git a/benchmarks/operator_benchmark/pt/unary_test.py b/benchmarks/operator_benchmark/pt/unary_test.py
index 5ae5dee..3544244 100644
--- a/benchmarks/operator_benchmark/pt/unary_test.py
+++ b/benchmarks/operator_benchmark/pt/unary_test.py
@@ -13,25 +13,26 @@
# Configs for pointwise unary ops
unary_ops_configs_short = op_bench.config_list(
+ attr_names=['M', 'N'],
attrs=[
[512, 512],
],
- attr_names=['M', 'N'],
+ cross_product_configs={
+ 'device': ['cpu'],
+ },
tags=['short']
)
-unary_ops_configs_long = op_bench.config_list(
- attrs=[
- [256, 256],
- [1024, 1024],
- ],
- attr_names=['M', 'N'],
+unary_ops_configs_long = op_bench.cross_product_configs(
+ M=[256, 1024],
+ N=[256, 1024],
+ device=['cpu'],
tags=['long']
)
class UnaryOpBenchmark(op_bench.TorchBenchmarkBase):
- def init(self, M, N, op_func):
- self.input_one = torch.rand(M, N)
+ def init(self, M, N, device, op_func):
+ self.input_one = torch.rand(M, N, device=device)
self.op_func = op_func
def forward(self):
