Use sparse sorted Vec for IseqProfile instead of dense pre-allocated Vecs

Replace the two dense Vec fields (opnd_types and num_profiles) that were
pre-allocated to iseq_size for every ISEQ with a single sorted
Vec<(u32, ProfileEntry)> that only allocates entries for instructions
that actually get profiled. Lookups use binary search.

On Lobsters benchmark: zjit_alloc_bytes drops from 38.9 MB to 20.7 MB
(46.8% reduction) with no performance regression.

Shrink Distribution counts from usize to u16

The counts in Distribution only need to track frequency during profiling
(typically up to num_profiles=30). u16 (max 65535) is more than enough
and saves 6 bytes per count slot.

Distribution<ProfiledType, 4> shrinks from 104 to 74 bytes.
On Lobsters: zjit_alloc_bytes drops from 20.7 MB to 18.5 MB.

wip

Revert "wip"

This reverts commit a3fe018442fb267c9350dcdb027eddf30bcc08bb.

ZJIT: Restore accidentally removed comments

ZJIT: Remove unused iseq_size parameter from IseqProfile::new

ZJIT: Move insn_idx into ProfileEntry struct

ZJIT: Add per-rev result caching to zjit_diff and build with stats

ZJIT: Split Distribution into Mono/Poly variants to reduce profile memory

ZJIT: Shrink NumProfiles from u32 to u16

Revert "ZJIT: Split Distribution into Mono/Poly variants to reduce profile memory"

This reverts commit 01a87703413f218d25973afd894b272edb51c12f.

.

Author: tekknolagi

tool/zjit_diff.rb

@@ -2,6 +2,7 @@
 # frozen_string_literal: true
 
 require 'fileutils'
+require 'json'
 require 'optparse'
 require 'tmpdir'
 require 'logger'
@@ -11,6 +12,7 @@
 GitRef = Struct.new(:ref, :commit_hash)
 
 RUBIES_DIR = File.join(Dir.home, '.zjit-diff')
+RESULTS_CACHE_DIR = File.join(Dir.home, '.zjit-diff', 'results')
 BEFORE_NAME = 'ruby-zjit-before'
 AFTER_NAME = 'ruby-zjit-after'
 
@@ -51,17 +53,52 @@ def bench!
 
   private
 
-  def run_benchmarks(ruby_bench_path)
+  def results_cache_key(rev_hash)
+    key_parts = [rev_hash, @options[:bench_args], @options[:name_filters]].inspect
+    Digest::MD5.hexdigest(key_parts)
+  end
+
+  def run_rev(ruby_bench_path, name, rev_hash)
+    cache_file = File.join(RESULTS_CACHE_DIR, "#{results_cache_key(rev_hash)}.json")
+
+    if File.exist?(cache_file) && !@options[:force_rerun]
+      LOG.info("Using cached results for #{name} from #{cache_file}")
+      return JSON.parse(File.read(cache_file))
+    end
+
+    out_name = File.join('data', "zjit_diff_#{name}")
     Dir.chdir(ruby_bench_path) do
       @runner.cmd({ 'RUBIES_DIR' => RUBIES_DIR },
                   './run_benchmarks.rb',
                   '--chruby',
-                  "before::#{@before_hash} --zjit-stats;after::#{@after_hash} --zjit-stats",
+                  "#{name}::#{rev_hash} --zjit-stats",
                   '--out-name',
-                  DATA_FILENAME,
+                  out_name,
                   *@options[:bench_args],
                   *@options[:name_filters])
+    end
+
+    result = JSON.parse(File.read(File.join(ruby_bench_path, "#{out_name}.json")))
+    FileUtils.mkdir_p(RESULTS_CACHE_DIR)
+    File.write(cache_file, JSON.generate(result))
+    LOG.info("Cached results for #{name} to #{cache_file}")
+    result
+  end
+
+  def run_benchmarks(ruby_bench_path)
+    before_data = run_rev(ruby_bench_path, 'before', @before_hash)
+    after_data = run_rev(ruby_bench_path, 'after', @after_hash)
+
+    merged = {
+      'metadata' => before_data['metadata'].merge(after_data['metadata']),
+      'raw_data' => before_data['raw_data'].merge(after_data['raw_data']),
+    }
 
+    merged_path = File.join(ruby_bench_path, "#{DATA_FILENAME}.json")
+    FileUtils.mkdir_p(File.dirname(merged_path))
+    File.write(merged_path, JSON.generate(merged))
+
+    Dir.chdir(ruby_bench_path) do
       @runner.cmd('./misc/zjit_diff.rb', "#{DATA_FILENAME}.json", out: $stdout)
     end
   end
@@ -98,7 +135,7 @@ def initialize(name:, ref:, runner:, force_rebuild: false)
 
   def build!
     Dir.chdir(@path) do
-      configure_cmd_args = ['--enable-zjit=dev', '--disable-install-doc']
+      configure_cmd_args = ['--enable-zjit=stats', '--disable-install-doc']
       if macos?
         brew_prefixes = BREW_REQUIRED_PACKAGES.map do |pkg|
           `brew --prefix #{pkg}`.strip
@@ -157,7 +194,7 @@ def clean!
   end
 
   if Dir.exist?(RUBIES_DIR)
-    LOG.info("Removing ruby installations from #{RUBIES_DIR}")
+    LOG.info("Removing ruby installations and cached results from #{RUBIES_DIR}")
     FileUtils.rm_rf(RUBIES_DIR)
   end
 
@@ -231,6 +268,10 @@ def parse_ref(ref)
       options[:force_rebuild] = true
     end
 
+    opts.on('--force-rerun', 'Force re-running benchmarks even if cached results exist') do
+      options[:force_rerun] = true
+    end
+
     opts.on('--quiet', 'Silence output of commands except for benchmark result') do
       options[:quiet] = true
     end

zjit/src/distribution.rs

@@ -1,5 +1,7 @@
 //! Type frequency distribution tracker.
 
+use crate::options::NumProfiles;
+
 /// This implementation was inspired by the type feedback module from Google's S6, which was
 /// written in C++ for use with Python. This is a new implementation in Rust created for use with
 /// Ruby instead of Python.
@@ -8,9 +10,9 @@ pub struct Distribution<T: Copy + PartialEq + Default, const N: usize> {
     /// buckets and counts have the same length
     /// `buckets[0]` is always the most common item
     buckets: [T; N],
-    counts: [usize; N],
+    counts: [NumProfiles; N],
     /// if there is no more room, increment the fallback
-    other: usize,
+    other: NumProfiles,
     // TODO(max): Add count disparity, which can help determine when to reset the distribution
 }
 
@@ -23,13 +25,13 @@ impl<T: Copy + PartialEq + Default, const N: usize> Distribution<T, N> {
         for (bucket, count) in self.buckets.iter_mut().zip(self.counts.iter_mut()) {
             if *bucket == item || *count == 0 {
                 *bucket = item;
-                *count += 1;
+                *count = count.saturating_add(1);
                 // Keep the most frequent item at the front
                 self.bubble_up();
                 return;
             }
         }
-        self.other += 1;
+        self.other = self.other.saturating_add(1);
     }
 
     /// Keep the highest counted bucket at index 0
@@ -87,7 +89,7 @@ impl<T: Copy + PartialEq + Default + std::fmt::Debug, const N: usize> Distributi
                 assert!(first_count >= count, "First count should be the largest");
             }
         }
-        let num_seen = dist.counts.iter().sum::<usize>() + dist.other;
+        let num_seen = dist.counts.iter().map(|&c| c as usize).sum::<usize>() + dist.other as usize;
         let kind = if dist.other == 0 {
             // Seen <= N types total
             if dist.counts[0] == 0 {

zjit/src/options.rs

@@ -8,7 +8,7 @@ use std::collections::HashSet;
 
 /// Default --zjit-num-profiles
 const DEFAULT_NUM_PROFILES: NumProfiles = 5;
-pub type NumProfiles = u32;
+pub type NumProfiles = u16;
 
 /// Default --zjit-call-threshold. This should be large enough to avoid compiling
 /// warmup code, but small enough to perform well on micro-benchmarks.

zjit/src/payload.rs

@@ -14,9 +14,9 @@ pub struct IseqPayload {
 }
 
 impl IseqPayload {
-    fn new(iseq_size: u32) -> Self {
+    fn new() -> Self {
         Self {
-            profile: IseqProfile::new(iseq_size),
+            profile: IseqProfile::new(),
             versions: vec![],
         }
     }
@@ -87,8 +87,7 @@ pub fn get_or_create_iseq_payload_ptr(iseq: IseqPtr) -> *mut IseqPayload {
             // We drop the payload with Box::from_raw when the GC frees the ISEQ and calls us.
             // NOTE(alan): Sometimes we read from an ISEQ without ever writing to it.
             // We allocate in those cases anyways.
-            let iseq_size = get_iseq_encoded_size(iseq);
-            let new_payload = IseqPayload::new(iseq_size);
+            let new_payload = IseqPayload::new();
             let new_payload = Box::into_raw(Box::new(new_payload));
             rb_iseq_set_zjit_payload(iseq, new_payload as VoidPtr);
 

zjit/src/profile.rs

@@ -104,8 +104,9 @@ fn profile_insn(bare_opcode: ruby_vminsn_type, ec: EcPtr) {
     }
 
     // Once we profile the instruction num_profiles times, we stop profiling it.
-    profile.num_profiles[profiler.insn_idx] = profile.num_profiles[profiler.insn_idx].saturating_add(1);
-    if profile.num_profiles[profiler.insn_idx] == get_option!(num_profiles) {
+    let entry = profile.entry_mut(profiler.insn_idx);
+    entry.num_profiles = entry.num_profiles.saturating_add(1);
+    if entry.num_profiles == get_option!(num_profiles) {
         unsafe { rb_zjit_iseq_insn_set(profiler.iseq, profiler.insn_idx as u32, bare_opcode); }
     }
 }
@@ -118,12 +119,12 @@ pub type TypeDistributionSummary = DistributionSummary<ProfiledType, DISTRIBUTIO
 
 /// Profile the Type of top-`n` stack operands
 fn profile_operands(profiler: &mut Profiler, profile: &mut IseqProfile, n: usize) {
-    let types = &mut profile.opnd_types[profiler.insn_idx];
-    if types.is_empty() {
-        types.resize(n, TypeDistribution::new());
+    let entry = profile.entry_mut(profiler.insn_idx);
+    if entry.opnd_types.is_empty() {
+        entry.opnd_types.resize(n, TypeDistribution::new());
     }
 
-    for (i, profile_type) in types.iter_mut().enumerate() {
+    for (i, profile_type) in entry.opnd_types.iter_mut().enumerate() {
         let obj = profiler.peek_at_stack((n - i - 1) as isize);
         // TODO(max): Handle GC-hidden classes like Array, Hash, etc and make them look normal or
         // drop them or something
@@ -134,33 +135,33 @@ fn profile_operands(profiler: &mut Profiler, profile: &mut IseqProfile, n: usize
 }
 
 fn profile_self(profiler: &mut Profiler, profile: &mut IseqProfile) {
-    let types = &mut profile.opnd_types[profiler.insn_idx];
-    if types.is_empty() {
-        types.resize(1, TypeDistribution::new());
+    let entry = profile.entry_mut(profiler.insn_idx);
+    if entry.opnd_types.is_empty() {
+        entry.opnd_types.resize(1, TypeDistribution::new());
     }
     let obj = profiler.peek_at_self();
     // TODO(max): Handle GC-hidden classes like Array, Hash, etc and make them look normal or
     // drop them or something
     let ty = ProfiledType::new(obj);
     VALUE::from(profiler.iseq).write_barrier(ty.class());
-    types[0].observe(ty);
+    entry.opnd_types[0].observe(ty);
 }
 
 fn profile_block_handler(profiler: &mut Profiler, profile: &mut IseqProfile) {
-    let types = &mut profile.opnd_types[profiler.insn_idx];
-    if types.is_empty() {
-        types.resize(1, TypeDistribution::new());
+    let entry = profile.entry_mut(profiler.insn_idx);
+    if entry.opnd_types.is_empty() {
+        entry.opnd_types.resize(1, TypeDistribution::new());
     }
     let obj = profiler.peek_at_block_handler();
     let ty = ProfiledType::object(obj);
     VALUE::from(profiler.iseq).write_barrier(ty.class());
-    types[0].observe(ty);
+    entry.opnd_types[0].observe(ty);
 }
 
 fn profile_getblockparamproxy(profiler: &mut Profiler, profile: &mut IseqProfile) {
-    let types = &mut profile.opnd_types[profiler.insn_idx];
-    if types.is_empty() {
-        types.resize(1, TypeDistribution::new());
+    let entry = profile.entry_mut(profiler.insn_idx);
+    if entry.opnd_types.is_empty() {
+        entry.opnd_types.resize(1, TypeDistribution::new());
     }
 
     let level = profiler.insn_opnd(1).as_u32();
@@ -170,7 +171,7 @@ fn profile_getblockparamproxy(profiler: &mut Profiler, profile: &mut IseqProfile
 
     let ty = ProfiledType::object(untagged);
     VALUE::from(profiler.iseq).write_barrier(ty.class());
-    types[0].observe(ty);
+    entry.opnd_types[0].observe(ty);
 }
 
 fn profile_invokesuper(profiler: &mut Profiler, profile: &mut IseqProfile) {
@@ -374,30 +375,61 @@ impl ProfiledType {
     }
 }
 
+/// Per-instruction profile entry, stored sparsely in a sorted Vec.
 #[derive(Debug)]
-pub struct IseqProfile {
-    /// Type information of YARV instruction operands, indexed by the instruction index
-    opnd_types: Vec<Vec<TypeDistribution>>,
+struct ProfileEntry {
+    /// YARV instruction index
+    insn_idx: u32,
+    /// Type information of YARV instruction operands
+    opnd_types: Vec<TypeDistribution>,
+    /// Number of profiled executions for this YARV instruction
+    num_profiles: NumProfiles,
+}
 
-    /// Number of profiled executions for each YARV instruction, indexed by the instruction index
-    num_profiles: Vec<NumProfiles>,
+#[derive(Debug)]
+pub struct IseqProfile {
+    /// Sparse storage of per-instruction profile data, sorted by instruction index.
+    /// Only instructions that have actually been profiled have entries here.
+    entries: Vec<ProfileEntry>,
 
     /// Method entries for `super` calls (stored as VALUE to be GC-safe)
     super_cme: HashMap<usize, TypeDistribution>
 }
 
 impl IseqProfile {
-    pub fn new(iseq_size: u32) -> Self {
+    pub fn new() -> Self {
         Self {
-            opnd_types: vec![vec![]; iseq_size as usize],
-            num_profiles: vec![0; iseq_size as usize],
+            entries: Vec::new(),
             super_cme: HashMap::new(),
         }
     }
 
+    /// Get or create a mutable profile entry for the given instruction index.
+    fn entry_mut(&mut self, insn_idx: usize) -> &mut ProfileEntry {
+        let idx = insn_idx as u32;
+        match self.entries.binary_search_by_key(&idx, |e| e.insn_idx) {
+            Ok(i) => &mut self.entries[i],
+            Err(i) => {
+                self.entries.insert(i, ProfileEntry {
+                    insn_idx: idx,
+                    opnd_types: Vec::new(),
+                    num_profiles: 0,
+                });
+                &mut self.entries[i]
+            }
+        }
+    }
+
+    /// Get a profile entry for the given instruction index (read-only).
+    fn entry(&self, insn_idx: usize) -> Option<&ProfileEntry> {
+        let idx = insn_idx as u32;
+        self.entries.binary_search_by_key(&idx, |e| e.insn_idx)
+            .ok().map(|i| &self.entries[i])
+    }
+
     /// Get profiled operand types for a given instruction index
     pub fn get_operand_types(&self, insn_idx: usize) -> Option<&[TypeDistribution]> {
-        self.opnd_types.get(insn_idx).map(|v| &**v)
+        self.entry(insn_idx).map(|e| e.opnd_types.as_slice()).filter(|s| !s.is_empty())
     }
 
     pub fn get_super_method_entry(&self, insn_idx: usize) -> Option<*const rb_callable_method_entry_t> {
@@ -413,8 +445,8 @@ impl IseqProfile {
 
     /// Run a given callback with every object in IseqProfile
     pub fn each_object(&self, callback: impl Fn(VALUE)) {
-        for operands in &self.opnd_types {
-            for distribution in operands {
+        for entry in &self.entries {
+            for distribution in &entry.opnd_types {
                 for profiled_type in distribution.each_item() {
                     // If the type is a GC object, call the callback
                     callback(profiled_type.class);
@@ -431,8 +463,8 @@ impl IseqProfile {
 
     /// Run a given callback with a mutable reference to every object in IseqProfile.
     pub fn each_object_mut(&mut self, callback: impl Fn(&mut VALUE)) {
-        for operands in &mut self.opnd_types {
-            for distribution in operands {
+        for entry in &mut self.entries {
+            for distribution in &mut entry.opnd_types {
                 for ref mut profiled_type in distribution.each_item_mut() {
                     // If the type is a GC object, call the callback
                     callback(&mut profiled_type.class);