YJIT: shink local types from 8 to 4 bytes (#8597)

* Shink local types from 8 to 4 bytes, context from 21 to 17 bytes

Use repr(packed)

* Add comment about Type being limited to 4 bits

1 files changed, 79 insertions, 24 deletions

diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index dca228dc44..e493fa5466 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -42,7 +42,7 @@ pub type IseqIdx = u16;
 #[derive(Copy, Clone, Hash, PartialEq, Eq, Debug)]
 #[repr(u8)]
 pub enum Type {
-    Unknown,
+    Unknown = 0,
     UnknownImm,
     UnknownHeap,
     Nil,
@@ -64,6 +64,10 @@ pub enum Type {
 
     BlockParamProxy, // A special sentinel value indicating the block parameter should be read from
                      // the current surrounding cfp
+
+    // The context currently relies on types taking at most 4 bits (max value 15)
+    // to encode, so if we add any more, we will need to refactor the context,
+    // or we could remove HeapSymbol, which is currently unused.
 }
 
 // Default initialization
@@ -438,7 +442,8 @@ impl RegTemps {
 /// Code generation context
 /// Contains information we can use to specialize/optimize code
 /// There are a lot of context objects so we try to keep the size small.
-#[derive(Clone, Default, Eq, Hash, PartialEq, Debug)]
+#[derive(Copy, Clone, Default, Eq, Hash, PartialEq, Debug)]
+#[repr(packed)]
 pub struct Context {
     // Number of values currently on the temporary stack
     stack_size: u8,
@@ -457,7 +462,8 @@ pub struct Context {
     chain_depth_return_landing: u8,
 
     // Local variable types we keep track of
-    local_types: [Type; MAX_LOCAL_TYPES],
+    // We store 8 local types, requiring 4 bits each, for a total of 32 bits
+    local_types: u32,
 
     // Type we track for self
     self_type: Type,
@@ -1711,7 +1717,7 @@ impl Context {
                     MapToLocal => {
                         let idx = mapping.get_local_idx();
                         assert!((idx as usize) < MAX_LOCAL_TYPES);
-                        return self.local_types[idx as usize];
+                        return self.get_local_type(idx.into());
                     }
                 }
             }
@@ -1719,8 +1725,14 @@ impl Context {
     }
 
     /// Get the currently tracked type for a local variable
-    pub fn get_local_type(&self, idx: usize) -> Type {
-        *self.local_types.get(idx).unwrap_or(&Type::Unknown)
+    pub fn get_local_type(&self, local_idx: usize) -> Type {
+        if local_idx >= MAX_LOCAL_TYPES {
+            return Type::Unknown
+        } else {
+            // Each type is stored in 4 bits
+            let type_bits = (self.local_types >> (4 * local_idx)) & 0b1111;
+            unsafe { transmute::<u8, Type>(type_bits as u8) }
+        }
     }
 
     /// Upgrade (or "learn") the type of an instruction operand
@@ -1756,7 +1768,9 @@ impl Context {
                     MapToLocal => {
                         let idx = mapping.get_local_idx() as usize;
                         assert!(idx < MAX_LOCAL_TYPES);
-                        self.local_types[idx].upgrade(opnd_type);
+                        let mut new_type = self.get_local_type(idx);
+                        new_type.upgrade(opnd_type);
+                        self.set_local_type(idx, new_type);
                     }
                 }
             }
@@ -1814,26 +1828,26 @@ impl Context {
 
     /// Set the type of a local variable
     pub fn set_local_type(&mut self, local_idx: usize, local_type: Type) {
-        let ctx = self;
-
         // If type propagation is disabled, store no types
         if get_option!(no_type_prop) {
             return;
         }
 
         if local_idx >= MAX_LOCAL_TYPES {
-            return;
+            return
         }
 
         // If any values on the stack map to this local we must detach them
-        for mapping in ctx.temp_mapping.iter_mut() {
-            *mapping = match mapping.get_kind() {
-                MapToStack => *mapping,
-                MapToSelf => *mapping,
+        for mapping_idx in 0..self.temp_mapping.len() {
+            let mapping = self.temp_mapping[mapping_idx];
+            self.temp_mapping[mapping_idx] = match mapping.get_kind() {
+                MapToStack => mapping,
+                MapToSelf => mapping,
                 MapToLocal => {
                     let idx = mapping.get_local_idx();
                     if idx as usize == local_idx {
-                        TempMapping::map_to_stack(ctx.local_types[idx as usize])
+                        let local_type = self.get_local_type(local_idx.into());
+                        TempMapping::map_to_stack(local_type)
                     } else {
                         TempMapping::map_to_local(idx)
                     }
@@ -1841,7 +1855,12 @@ impl Context {
             }
         }
 
-        ctx.local_types[local_idx] = local_type;
+        // Update the type bits
+        let type_bits = local_type as u32;
+        assert!(type_bits <= 0b1111);
+        let mask_bits = (0b1111 as u32) << (4 * local_idx);
+        let shifted_bits = type_bits << (4 * local_idx);
+        self.local_types = (self.local_types & !mask_bits) | shifted_bits;
     }
 
     /// Erase local variable type information
@@ -1849,15 +1868,17 @@ impl Context {
     pub fn clear_local_types(&mut self) {
         // When clearing local types we must detach any stack mappings to those
         // locals. Even if local values may have changed, stack values will not.
-        for mapping in self.temp_mapping.iter_mut() {
+
+        for mapping_idx in 0..self.temp_mapping.len() {
+            let mapping = self.temp_mapping[mapping_idx];
             if mapping.get_kind() == MapToLocal {
-                let idx = mapping.get_local_idx();
-                *mapping = TempMapping::map_to_stack(self.local_types[idx as usize]);
+                let local_idx = mapping.get_local_idx() as usize;
+                self.temp_mapping[mapping_idx] = TempMapping::map_to_stack(self.get_local_type(local_idx));
             }
         }
 
         // Clear the local types
-        self.local_types = [Type::default(); MAX_LOCAL_TYPES];
+        self.local_types = 0;
     }
 
     /// Compute a difference score for two context objects
@@ -1902,9 +1923,9 @@ impl Context {
         };
 
         // For each local type we track
-        for i in 0..src.local_types.len() {
-            let t_src = src.local_types[i];
-            let t_dst = dst.local_types[i];
+        for i in 0.. MAX_LOCAL_TYPES {
+            let t_src = src.get_local_type(i);
+            let t_dst = dst.get_local_type(i);
             diff += match t_src.diff(t_dst) {
                 TypeDiff::Compatible(diff) => diff,
                 TypeDiff::Incompatible => return TypeDiff::Incompatible,
@@ -3238,11 +3259,45 @@ mod tests {
     use crate::core::*;
 
     #[test]
+    fn type_size() {
+        // Check that we can store types in 4 bits,
+        // and all local types in 32 bits
+        assert_eq!(mem::size_of::<Type>(), 1);
+        assert!(Type::BlockParamProxy as usize <= 0b1111);
+        assert!(MAX_LOCAL_TYPES * 4 <= 32);
+    }
+
+    #[test]
     fn tempmapping_size() {
         assert_eq!(mem::size_of::<TempMapping>(), 1);
     }
 
     #[test]
+    fn local_types() {
+        let mut ctx = Context::default();
+
+        for i in 0..MAX_LOCAL_TYPES {
+            ctx.set_local_type(i, Type::Fixnum);
+            assert_eq!(ctx.get_local_type(i), Type::Fixnum);
+            ctx.set_local_type(i, Type::BlockParamProxy);
+            assert_eq!(ctx.get_local_type(i), Type::BlockParamProxy);
+        }
+
+        ctx.set_local_type(0, Type::Fixnum);
+        ctx.clear_local_types();
+        assert!(ctx.get_local_type(0) == Type::Unknown);
+
+        // Make sure we don't accidentally set bits incorrectly
+        let mut ctx = Context::default();
+        ctx.set_local_type(0, Type::Fixnum);
+        assert_eq!(ctx.get_local_type(0), Type::Fixnum);
+        ctx.set_local_type(2, Type::Fixnum);
+        ctx.set_local_type(1, Type::BlockParamProxy);
+        assert_eq!(ctx.get_local_type(0), Type::Fixnum);
+        assert_eq!(ctx.get_local_type(2), Type::Fixnum);
+    }
+
+    #[test]
     fn tempmapping() {
         let t = TempMapping::map_to_stack(Type::Unknown);
         assert_eq!(t.get_kind(), MapToStack);
@@ -3259,7 +3314,7 @@ mod tests {
 
     #[test]
     fn context_size() {
-        assert_eq!(mem::size_of::<Context>(), 21);
+        assert_eq!(mem::size_of::<Context>(), 17);
     }
 
     #[test]