From 3fbf9df39ad9c9338a2f9b94b4dd87c1c446ec95 Mon Sep 17 00:00:00 2001
From: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Date: Thu, 11 Jul 2024 15:01:05 -0400
Subject: [PATCH] YJIT: increase context cache size to 1024 redux (#11140)

* YJIT: increase context cache size to 1024 redux

* Move context hashing code outside of unsafe block

* Avoid allocating large table on the stack, which would cause a stack overflow

Co-authored by Alan Wu @XrXr
---
 yjit/src/core.rs | 28 +++++++++++++++-------------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 7ddb65fd90..699fcae840 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -842,7 +842,7 @@ enum CtxOp {
 }
 
 // Number of entries in the context cache
-const CTX_CACHE_SIZE: usize = 512;
+const CTX_CACHE_SIZE: usize = 1024;
 
 // Cache of the last contexts encoded
 // Empirically this saves a few percent of memory
@@ -906,18 +906,20 @@ impl Context {
     // Store an entry in a cache of recently encoded/decoded contexts
     fn cache_set(ctx: &Context, idx: u32)
     {
+        // Compute the hash for this context
+        let mut hasher = DefaultHasher::new();
+        ctx.hash(&mut hasher);
+        let ctx_hash = hasher.finish() as usize;
+
         unsafe {
             // Lazily initialize the context cache
             if CTX_CACHE == None {
-                let empty_tbl = [(Context::default(), 0); CTX_CACHE_SIZE];
-                CTX_CACHE = Some(Box::new(empty_tbl));
+                // Here we use the vec syntax to avoid allocating the large table on the stack,
+                // as this can cause a stack overflow
+                let tbl = vec![(Context::default(), 0); CTX_CACHE_SIZE].into_boxed_slice().try_into().unwrap();
+                CTX_CACHE = Some(tbl);
             }
 
-            // Compute the hash for this context
-            let mut hasher = DefaultHasher::new();
-            ctx.hash(&mut hasher);
-            let ctx_hash = hasher.finish() as usize;
-
             // Write a cache entry for this context
             let cache = CTX_CACHE.as_mut().unwrap();
             cache[ctx_hash % CTX_CACHE_SIZE] = (*ctx, idx);
@@ -927,6 +929,11 @@ impl Context {
     // Lookup the context in a cache of recently encoded/decoded contexts
     fn cache_get(ctx: &Context) -> Option<u32>
     {
+        // Compute the hash for this context
+        let mut hasher = DefaultHasher::new();
+        ctx.hash(&mut hasher);
+        let ctx_hash = hasher.finish() as usize;
+
         unsafe {
             if CTX_CACHE == None {
                 return None;
@@ -934,11 +941,6 @@ impl Context {
 
             let cache = CTX_CACHE.as_mut().unwrap();
 
-            // Compute the hash for this context
-            let mut hasher = DefaultHasher::new();
-            ctx.hash(&mut hasher);
-            let ctx_hash = hasher.finish() as usize;
-
             // Check that the context for this cache entry mmatches
             let cache_entry = &cache[ctx_hash % CTX_CACHE_SIZE];
             if cache_entry.0 == *ctx {