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

2 files changed, 27 insertions, 24 deletions

diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 0cfd7d9124..397ebb4566 100644
--- a/yjit/src/core.rs
+++ b/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,43 +898,44 @@ impl Context {
         ctx
     }
 
-    // Lookup the context in a cache of recently encoded/decoded contexts
-    fn cache_get(ctx: &Context) -> Option<u32>
+    // Store an entry in a cache of recently encoded/decoded contexts
+    fn cache_set(ctx: &Context, idx: u32)
     {
         unsafe {
             if CTX_CACHE == None {
-                return None;
+                let empty_tbl = [(Context::default(), 0); CTX_CACHE_SIZE];
+                CTX_CACHE = Some(Box::new(empty_tbl));
             }
 
-            let cache = CTX_CACHE.as_mut().unwrap();
-
             let mut hasher = DefaultHasher::new();
             ctx.hash(&mut hasher);
             let ctx_hash = hasher.finish() as usize;
-            let cache_entry = &cache[ctx_hash % CTX_CACHE_SIZE];
-
-            if cache_entry.0 == *ctx {
-                return Some(cache_entry.1);
-            }
 
-            return None;
+            let cache = CTX_CACHE.as_mut().unwrap();
+            cache[ctx_hash % CTX_CACHE_SIZE] = (*ctx, idx);
         }
     }
 
-    // Store an entry in a cache of recently encoded/decoded contexts
-    fn cache_set(ctx: &Context, idx: u32)
+    // Lookup the context in a cache of recently encoded/decoded contexts
+    fn cache_get(ctx: &Context) -> Option<u32>
     {
         unsafe {
             if CTX_CACHE == None {
-                CTX_CACHE = Some( [(Context::default(), 0); CTX_CACHE_SIZE] );
+                return None;
             }
 
+            let cache = CTX_CACHE.as_mut().unwrap();
+
             let mut hasher = DefaultHasher::new();
             ctx.hash(&mut hasher);
             let ctx_hash = hasher.finish() as usize;
+            let cache_entry = &cache[ctx_hash % CTX_CACHE_SIZE];
 
-            let cache = CTX_CACHE.as_mut().unwrap();
-            cache[ctx_hash % CTX_CACHE_SIZE] = (*ctx, idx);
+            if cache_entry.0 == *ctx {
+                return Some(cache_entry.1);
+            }
+
+            return None;
         }
     }
 
@@ -935,12 +943,6 @@ impl Context {
     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
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 6a7de68576..726e201b95 100644
--- a/yjit/src/stats.rs
+++ b/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);