Port run_benchmarks.rb --yjit-stats=STATS for ZJIT

harness/harness-common.rb

@@ -138,8 +138,8 @@ def return_results(warmup_iterations, bench_iterations)
     yjit_bench_results["maxrss"] = maxrss
   end
 
-  # If YJIT or ZJIT is enabled, show some of their stats at the end.
-  if yjit_stats
+  # If YJIT or ZJIT is enabled, show some of its stats unless it does by itself.
+  if yjit_stats && !RubyVM::YJIT.stats_enabled?
     yjit_bench_results["yjit_stats"] = yjit_stats
     stats_keys = [
       *ENV.fetch("YJIT_BENCH_STATS", "").split(",").map(&:to_sym),
@@ -150,9 +150,10 @@ def return_results(warmup_iterations, bench_iterations)
       :compile_time_ns,
     ].uniq
     puts "YJIT stats:"
-  elsif zjit_stats
+  elsif zjit_stats && !RubyVM::ZJIT.stats_enabled?
     yjit_bench_results["zjit_stats"] = zjit_stats
     stats_keys = [
+      *ENV.fetch("ZJIT_BENCH_STATS", "").split(",").map(&:to_sym),
       :compile_time_ns,
       :profile_time_ns,
       :gc_time_ns,

harness/harness.rb

@@ -40,16 +40,18 @@ def run_benchmark(_num_itrs_hint, &block)
 
   RubyVM::YJIT.reset_stats! if defined?(RubyVM::YJIT) && RubyVM::YJIT.enabled?
 
-  # If $YJIT_BENCH_STATS is given, print the diff of these stats at each iteration.
-  if ENV["YJIT_BENCH_STATS"]
-    yjit_stats = ENV["YJIT_BENCH_STATS"].split(",").map { |key| [key.to_sym, nil] }.to_h
+  # If $YJIT_BENCH_STATS or $ZJIT_BENCH_STATS is given, print the diff of these stats at each iteration.
+  if yjit_stats = ENV["YJIT_BENCH_STATS"]
+    yjit_stats = yjit_stats.split(",").map(&:to_sym).map { |key| [key, RubyVM::YJIT.runtime_stats(key)] }.to_h
     yjit_stats.each_key { |key| header << " #{key}" }
   end
+  if zjit_stats = ENV["ZJIT_BENCH_STATS"]
+    zjit_stats = zjit_stats.split(",").map(&:to_sym).map { |key| [key, RubyVM::ZJIT.stats(key)] }.to_h
+    zjit_stats.each_key { |key| header << " #{key}" }
+  end
 
   puts header
   begin
-    yjit_stats&.each_key { |key| yjit_stats[key] = RubyVM::YJIT.runtime_stats(key) }
-
     time = realtime(&block)
     num_itrs += 1
 
@@ -59,10 +61,14 @@ def run_benchmark(_num_itrs_hint, &block)
 
     yjit_stats&.each do |key, old_value|
       new_value = RubyVM::YJIT.runtime_stats(key)
-      diff = (new_value - old_value)
-      itr_str << " %#{key.size}s" % format_number(diff)
+      itr_str << " %#{key.size}s" % format_number(new_value - old_value)
       yjit_stats[key] = new_value
     end
+    zjit_stats&.each do |key, old_value|
+      new_value = RubyVM::ZJIT.stats(key)
+      itr_str << " %#{key.size}s" % format_number(new_value - old_value)
+      zjit_stats[key] = new_value
+    end
 
     puts itr_str
     # NOTE: we may want to preallocate an array and avoid append

run_benchmarks.rb

@@ -396,6 +396,10 @@ def run_benchmarks(ruby:, ruby_description:, categories:, name_filters:, out_pat
     ENV["YJIT_BENCH_STATS"] = str
   end
 
+  opts.on("--zjit-stats=STATS", "print ZJIT stats at each iteration for the default harness") do |str|
+    ENV["ZJIT_BENCH_STATS"] = str
+  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