YJIT: add context cache size stat, lazily allocate cache

* YJIT: add context cache size stat
* Allocate the context cache in a box so CRuby doesn't pay overhead
* Add an extra debug assertion

yjit.rb

@@ -390,6 +390,7 @@ module RubyVM::YJIT
       out.puts "yjit_alloc_size:       " + format_number(13, stats[:yjit_alloc_size]) if stats.key?(:yjit_alloc_size)
 
       bytes_per_context = stats[:context_data_bytes].fdiv(stats[:num_contexts_encoded])
+      out.puts "context_cache_bytes:   " + format_number(13, stats[:context_cache_bytes])
       out.puts "context_data_bytes:    " + format_number(13, stats[:context_data_bytes])
       out.puts "num_contexts_encoded:  " + format_number(13, stats[:num_contexts_encoded])
       out.puts "bytes_per_context:     " + ("%13.2f" % bytes_per_context)

yjit/src/core.rs

@@ -841,12 +841,17 @@ enum CtxOp {
     EndOfCode,
 }
 
+// Number of entries in the context cache
 const CTX_CACHE_SIZE: usize = 512;
 
 // Cache of the last contexts encoded
 // Empirically this saves a few percent of memory
 // We can experiment with varying the size of this cache
-static mut CTX_CACHE: Option<[(Context, u32); CTX_CACHE_SIZE]> = None;
+pub type CtxCacheTbl = [(Context, u32); CTX_CACHE_SIZE];
+static mut CTX_CACHE: Option<Box<CtxCacheTbl>> = None;
+
+// Size of the context cache in bytes
+pub const CTX_CACHE_BYTES: usize = std::mem::size_of::<CtxCacheTbl>();
 
 impl Context {
     pub fn encode(&self) -> u32 {
@@ -857,12 +862,14 @@ impl Context {
         }
 
         if let Some(idx) = Self::cache_get(self) {
+            debug_assert!(Self::decode(idx) == *self);
             return idx;
         }
 
         let context_data = CodegenGlobals::get_context_data();
 
-        // Offset 0 is reserved for the default context
+        // Make sure we don't use offset 0 because
+        // it's is reserved for the default context
         if context_data.num_bits() == 0 {
             context_data.push_u1(0);
         }
@@ -891,6 +898,24 @@ impl Context {
         ctx
     }
 
+    // Store an entry in a cache of recently encoded/decoded contexts
+    fn cache_set(ctx: &Context, idx: u32)
+    {
+        unsafe {
+            if CTX_CACHE == None {
+                let empty_tbl = [(Context::default(), 0); CTX_CACHE_SIZE];
+                CTX_CACHE = Some(Box::new(empty_tbl));
+            }
+
+            let mut hasher = DefaultHasher::new();
+            ctx.hash(&mut hasher);
+            let ctx_hash = hasher.finish() as usize;
+
+            let cache = CTX_CACHE.as_mut().unwrap();
+            cache[ctx_hash % CTX_CACHE_SIZE] = (*ctx, idx);
+        }
+    }
+
     // Lookup the context in a cache of recently encoded/decoded contexts
     fn cache_get(ctx: &Context) -> Option<u32>
     {
@@ -914,33 +939,10 @@ impl Context {
         }
     }
 
-    // Store an entry in a cache of recently encoded/decoded contexts
-    fn cache_set(ctx: &Context, idx: u32)
-    {
-        unsafe {
-            if CTX_CACHE == None {
-                CTX_CACHE = Some( [(Context::default(), 0); CTX_CACHE_SIZE] );
-            }
-
-            let mut hasher = DefaultHasher::new();
-            ctx.hash(&mut hasher);
-            let ctx_hash = hasher.finish() as usize;
-
-            let cache = CTX_CACHE.as_mut().unwrap();
-            cache[ctx_hash % CTX_CACHE_SIZE] = (*ctx, idx);
-        }
-    }
-
     // Encode into a compressed context representation in a bit vector
     fn encode_into(&self, bits: &mut BitVector) -> usize {
         let start_idx = bits.num_bits();
 
-        // NOTE: this value is often zero or falls within
-        // a small range, so could be compressed
-        //println!("stack_size={}", self.stack_size);
-        //println!("sp_offset={}", self.sp_offset);
-        //println!("chain_depth_and_flags={}", self.chain_depth_and_flags);
-
         // Most of the time, the stack size is small and sp offset has the same value
         if (self.stack_size as i64) == (self.sp_offset as i64) && self.stack_size < 4 {
             // One single bit to signify a compact stack_size/sp_offset encoding

yjit/src/stats.rs

@@ -746,6 +746,7 @@ fn rb_yjit_gen_stats_dict() -> VALUE {
         // How many bytes we are using to store context data
         let context_data = CodegenGlobals::get_context_data();
         hash_aset_usize!(hash, "context_data_bytes", context_data.num_bytes());
+        hash_aset_usize!(hash, "context_cache_bytes", crate::core::CTX_CACHE_BYTES);
 
         // VM instructions count
         hash_aset_usize!(hash, "vm_insns_count", rb_vm_insns_count as usize);