From 5c877434f70e44955b658c5c93b91218eecf5344 Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Mon, 8 Jan 2024 13:33:47 -0800
Subject: [PATCH] Add --bench and --warmup flags (#261)

* Add --bench and --warmup flags

* Drop the --bench=1 shorthand

* Update the help description of --bench

* Update the help description of --warmup

Co-authored-by: Benoit Daloze <eregontp@gmail.com>

* Add --once option

---------

Co-authored-by: Benoit Daloze <eregontp@gmail.com>
---
 README.md         | 12 +++++++++++-
 run_benchmarks.rb | 15 +++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index c97ac5f4..e3ecf1b0 100644
--- a/README.md
+++ b/README.md
@@ -174,7 +174,17 @@ can be controlled by the following environment variables:
 * `MIN_BENCH_ITRS`: The minimum number of benchmark iterations (default: 10)
 * `MIN_BENCH_TIME`: The minimum seconds for benchmark (default: 10)
 
-There is a handy script for running benchmarks just once using
+You can also use `--warmup`, `--bench`, or `--once` to set these environment variables:
+
+```sh
+# same as: WARMUP_ITRS=2 MIN_BENCH_ITRS=3 MIN_BENCH_TIME=0 ./run_benchmarks.rb railsbench
+./run_benchmarks.rb railsbench --warmup=2 --bench=3
+
+# same as: WARMUP_ITRS=0 MIN_BENCH_ITRS=1 MIN_BENCH_TIME=0 ./run_benchmarks.rb railsbench
+./run_benchmarks.rb railsbench --once
+```
+
+There is also a handy script for running benchmarks just once using
 `WARMUP_ITRS=0 MIN_BENCH_ITRS=1 MIN_BENCH_TIME=0`, for example
 with the `--yjit-stats` command-line option:
 
diff --git a/run_benchmarks.rb b/run_benchmarks.rb
index ffd199ac..f6868ef4 100755
--- a/run_benchmarks.rb
+++ b/run_benchmarks.rb
@@ -355,6 +355,21 @@ def run_benchmarks(ruby:, ruby_description:, categories:, name_filters:, out_pat
     args.harness = v
   end
 
+  opts.on("--warmup=N", "the number of warmup iterations for the default harness (default: 15)") do |n|
+    ENV["WARMUP_ITRS"] = n
+  end
+
+  opts.on("--bench=N", "the number of benchmark iterations for the default harness (default: 10). Also defaults MIN_BENCH_TIME to 0.") do |n|
+    ENV["MIN_BENCH_ITRS"] = n
+    ENV["MIN_BENCH_TIME"] ||= "0"
+  end
+
+  opts.on("--once", "benchmarks only 1 iteration with no warmup for the default harness") do
+    ENV["WARMUP_ITRS"] = "0"
+    ENV["MIN_BENCH_ITRS"] = "1"
+    ENV["MIN_BENCH_TIME"] = "0"
+  end
+
   opts.on("--yjit_opts=OPT_STRING", "string of command-line options to run YJIT with (ignored if you use -e)") do |str|
     args.yjit_opts=str
   end