From 9cbff116813116db4b6652bb7474dec7257e28ef Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Fri, 5 Jan 2024 17:20:12 -0800
Subject: [PATCH] Add --bench and --warmup flags

---
 README.md         | 10 ++++++++++
 run_benchmarks.rb | 10 ++++++++++
 2 files changed, 20 insertions(+)

diff --git a/README.md b/README.md
index c97ac5f4..979c24d8 100644
--- a/README.md
+++ b/README.md
@@ -174,6 +174,16 @@ 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)
 
+You can also use `--warmup` and `--bench` to set these environment variables:
+
+```sh
+# same as: WARMUP_ITRS=0 MIN_BENCH_ITRS=1 MIN_BENCH_TIME=0 ./run_benchmarks.rb railsbench
+./run_benchmarks.rb railsbench --bench=1
+
+# 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
+```
+
 There is 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 7a8d739b..0b5eac17 100755
--- a/run_benchmarks.rb
+++ b/run_benchmarks.rb
@@ -347,6 +347,16 @@ def run_benchmarks(ruby:, ruby_description:, categories:, name_filters:, out_pat
     args.harness = v
   end
 
+  opts.on("--warmup=N", "the number of warmup iterations (default: 15)") do |n|
+    ENV["WARMUP_ITRS"] = n
+  end
+
+  opts.on("--bench=N", "the number of benchmark iterations (default: 10)") do |n|
+    ENV["WARMUP_ITRS"] ||= "0"
+    ENV["MIN_BENCH_ITRS"] = n
+    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