1 #include <ATen/ATen.h> 2 3 #include <benchmark/benchmark.h> 4 tensor_add(benchmark::State & state)5static void tensor_add(benchmark::State& state) { 6 const size_t batchSize = static_cast<size_t>(state.range(0)); 7 const size_t channels = static_cast<size_t>(state.range(1)); 8 9 at::Tensor a = at::rand({batchSize, channels}); 10 at::Tensor b = at::rand({batchSize, channels}); 11 at::Tensor c; 12 for (auto _ : state) { 13 c = a + b; 14 } 15 } 16 GenerateSizes(benchmark::internal::Benchmark * b)17static void GenerateSizes(benchmark::internal::Benchmark* b) { 18 b->ArgNames({"N", "C"}); 19 20 for (size_t n = 8; n < 1024;) { 21 for (size_t c = 8; c < 1024;) { 22 b->Args({n, c}); 23 c *= 2; 24 } 25 n *= 2; 26 } 27 } 28 29 BENCHMARK(tensor_add)->Apply(GenerateSizes); 30 BENCHMARK_MAIN(); 31