hlo_dfs_reachability_test: use batch mode in Build benchmark

- Before:
----------------------------------------------------------------------------
Benchmark                                  Time             CPU   Iterations
----------------------------------------------------------------------------
BM_HloDfsReachabilityBuild/1             143 ns          143 ns     10314343
BM_HloDfsReachabilityBuild/64           1776 ns         1776 ns       783254
BM_HloDfsReachabilityBuild/128          3565 ns         3564 ns       391936
BM_HloDfsReachabilityBuild/256          7446 ns         7445 ns       189334
BM_HloDfsReachabilityBuild/512         15044 ns        15041 ns        92625
BM_HloDfsReachabilityBuild/4096       158882 ns       158865 ns         8625
BM_HloDfsReachabilityBuild/32768     1709590 ns      1709401 ns          825
BM_HloDfsReachabilityBuild/262144   37142558 ns     37136529 ns           38

- After:
----------------------------------------------------------------------------
Benchmark                                  Time             CPU   Iterations
----------------------------------------------------------------------------
BM_HloDfsReachabilityBuild/1             135 ns          135 ns     10330052
BM_HloDfsReachabilityBuild/64           27.8 ns         27.8 ns     50131200
BM_HloDfsReachabilityBuild/128          28.7 ns         28.7 ns     47522432
BM_HloDfsReachabilityBuild/256          29.2 ns         29.2 ns     47567104
BM_HloDfsReachabilityBuild/512          29.9 ns         29.9 ns     46650880
BM_HloDfsReachabilityBuild/4096         41.1 ns         41.1 ns     36945920
BM_HloDfsReachabilityBuild/32768        59.5 ns         59.5 ns     23035904
BM_HloDfsReachabilityBuild/262144        154 ns          154 ns      8912896

PiperOrigin-RevId: 619340769
This commit is contained in:
Emilio Cota 2024-03-26 16:06:21 -07:00 committed by TensorFlower Gardener
parent 0e6c3baaa5
commit 9fd1dc32fe

View File

@ -168,8 +168,9 @@ class HloDfsReachabilityBenchmark {
}; };
void BM_HloDfsReachabilityBuild(benchmark::State& state) { void BM_HloDfsReachabilityBuild(benchmark::State& state) {
HloDfsReachabilityBenchmark bm(state.range(0), state.name()); int num_nodes = state.range(0);
for (auto s : state) { HloDfsReachabilityBenchmark bm(num_nodes, state.name());
while (state.KeepRunningBatch(num_nodes)) {
benchmark::DoNotOptimize(bm.Build()); benchmark::DoNotOptimize(bm.Build());
} }
} }