unify chunk benchmark (#28892)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/28892
as title
Test Plan:
```
buck run mode/opt //caffe2/benchmarks/operator_benchmark/pt:chunk_test
# ----------------------------------------
# PyTorch/Caffe2 Operator Micro-benchmarks
# ----------------------------------------
# Tag : short
# Benchmarking PyTorch: chunks
# Mode: Eager
# Name: chunks_M256_N512_chunks2_cpu
# Input: M: 256, N: 512, chunks: 2, device: cpu
Forward Execution Time (us) : 4.098
Reviewed By: hl475
Differential Revision: D18227499
fbshipit-source-id: 72268b7fe94a7d92d6e47f58f33902a33367c68b
diff --git a/benchmarks/operator_benchmark/pt/chunk_test.py b/benchmarks/operator_benchmark/pt/chunk_test.py
index 0c3fde1..f30c8b9 100644
--- a/benchmarks/operator_benchmark/pt/chunk_test.py
+++ b/benchmarks/operator_benchmark/pt/chunk_test.py
@@ -11,17 +11,30 @@
# Configs for PT Chunk operator
-chunks_short_configs = op_bench.cross_product_configs(
- M=[256, 512],
- N=[512],
- chunks=[2],
- tags=['short']
+chunk_short_configs = op_bench.config_list(
+ attr_names=["M", "N", "chunks"],
+ attrs=[
+ [256, 512, 2],
+ [512, 512, 2],
+ ],
+ cross_product_configs={
+ 'device': ['cpu'],
+ },
+ tags=["short"],
+)
+
+chunks_long_configs = op_bench.cross_product_configs(
+ M=[128, 1024],
+ N=[128, 1024],
+ chunks=[2, 4],
+ device=['cpu'],
+ tags=['long']
)
class ChunkBenchmark(op_bench.TorchBenchmarkBase):
- def init(self, M, N, chunks):
- self.input_one = torch.rand(M, N)
+ def init(self, M, N, chunks, device):
+ self.input_one = torch.rand(M, N, device=device)
self.chunks = chunks
self.set_module_name('chunks')
@@ -29,7 +42,8 @@
return torch.chunk(self.input_one, self.chunks)
-op_bench.generate_pt_test(chunks_short_configs, ChunkBenchmark)
+op_bench.generate_pt_test(chunk_short_configs + chunks_long_configs,
+ ChunkBenchmark)
if __name__ == "__main__":
