Add an option to monitor YJIT stats (#330)

* Add an option to monitor YJIT stats

* Update README.md

---------

Co-authored-by: Maxime Chevalier-Boisvert <[email protected]>

README.md

@@ -210,6 +210,15 @@ PERF=record ruby --yjit-perf=map -Iharness-perf benchmarks/railsbench/benchmark.
 
 This is the only harness that uses `run_benchmark`'s argument, `num_itrs_hint`.
 
+### Printing YJIT stats
+
+The `--yjit-stats` option of `./run_benchmarks.rb` allows you to print the diff of YJIT stats counters
+after each iteration with the default harness.
+
+```
+./run_benchmarks.rb --yjit-stats=code_region_size,yjit_alloc_size
+```
+
 ## Measuring memory usage
 
 `--rss` option of `run_benchmarks.rb` allows you to measure RSS after benchmark iterations.

harness/harness-common.rb

@@ -91,10 +91,7 @@ def return_results(warmup_iterations, bench_iterations)
   # Collect our own peak mem usage as soon as reasonable after finishing the last iteration.
   rss = get_rss
   yjit_bench_results["rss"] = rss
-  puts "RSS: %.1fMiB" % (rss / 1024.0 / 1024.0)
-
   if maxrss = get_maxrss
-    puts "MAXRSS: %.1fMiB" % (maxrss / 1024.0 / 1024.0)
     yjit_bench_results["maxrss"] = maxrss
   end
 
@@ -103,15 +100,29 @@ def return_results(warmup_iterations, bench_iterations)
     yjit_bench_results["yjit_stats"] = yjit_stats
 
     formatted_stats = proc { |key| "%10s" % yjit_stats[key].to_s.reverse.scan(/\d{1,3}/).join(",").reverse }
-    puts "inline_code_size:   #{formatted_stats[:inline_code_size]}"
-    puts "outlined_code_size: #{formatted_stats[:outlined_code_size]}"
-    puts "code_region_size:   #{formatted_stats[:code_region_size]}"
-    puts "yjit_alloc_size:    #{formatted_stats[:yjit_alloc_size]}"
-    if yjit_stats.key?(:compile_time_ns)
-      puts "yjit_compile_time:  %.2fms" % (yjit_stats[:compile_time_ns] / 1_000_000.0).round(2)
+    yjit_stats_keys = [
+      *ENV.fetch("YJIT_BENCH_STATS", "").split(",").map(&:to_sym),
+      :inline_code_size,
+      :outlined_code_size,
+      :code_region_size,
+      :yjit_alloc_size,
+      :compile_time_ns,
+    ].uniq
+    yjit_stats_pads = yjit_stats_keys.map(&:size).max + 1
+    yjit_stats_keys.each do |key|
+      if key == :compile_time_ns
+        puts "#{"yjit_compile_time:".ljust(yjit_stats_pads)} %8.2fms" % (yjit_stats[:compile_time_ns] / 1_000_000.0).round(2)
+      else
+        puts "#{"#{key}:".ljust(yjit_stats_pads)} #{formatted_stats[key]}"
+      end
     end
   end
 
+  puts "RSS: %.1fMiB" % (rss / 1024.0 / 1024.0)
+  if maxrss
+    puts "MAXRSS: %.1fMiB" % (maxrss / 1024.0 / 1024.0)
+  end
+
   write_json_file(yjit_bench_results)
 end
 

harness/harness.rb

@@ -36,15 +36,33 @@ def run_benchmark(_num_itrs_hint, &block)
   times = []
   total_time = 0
   num_itrs = 0
+  header = "itr:   time"
 
+  # 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
+    yjit_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
 
     # NOTE: we may want to avoid this as it could trigger GC?
     time_ms = (1000 * time).to_i
-    puts "itr \##{num_itrs}: #{time_ms}ms"
+    itr_str = "%4s %6s" % ["##{num_itrs}:", "#{time_ms}ms"]
+
+    yjit_stats&.each do |key, old_value|
+      new_value = RubyVM::YJIT.runtime_stats(key)
+      diff = (new_value - old_value).to_s.reverse.scan(/\d{1,3}/).join(",").reverse
+      itr_str << " %#{key.size}s" % diff
+      yjit_stats[key] = new_value
+    end
 
+    puts itr_str
     # NOTE: we may want to preallocate an array and avoid append
     # We internally save the time in seconds to avoid loss of precision
     times << time

run_benchmarks.rb

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