YJIT: shrink Context from 29 to 21 bytes by reducing space used by TempMapping (#8321)

* YJIT: merge tempmapping and temp types into a single-byte encoding

YJIT: refactor to shrink Context by 8 bytes

* Add tests, fix bug in TempMapping::map_to_local()

* Update yjit/src/core.rs

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>

* Update yjit/src/core.rs

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>

* Fewer transmutes where `as` would suffice. Also repr(u8)

* Update yjit/src/core.rs

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>

* Update yjit/src/core.rs

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>

* Update yjit/src/core.rs

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>

---------

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>
Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

2 files changed, 161 insertions, 107 deletions

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 354b8e5fd2..c506b4360c 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -342,12 +342,14 @@ fn verify_ctx(jit: &JITState, ctx: &Context) {
     // Verify stack operand types
     let top_idx = cmp::min(ctx.get_stack_size(), MAX_TEMP_TYPES as u8);
     for i in 0..top_idx {
-        let (learned_mapping, learned_type) = ctx.get_opnd_mapping(StackOpnd(i));
+        let learned_mapping = ctx.get_opnd_mapping(StackOpnd(i));
+        let learned_type = ctx.get_opnd_type(StackOpnd(i));
+
         let stack_val = jit.peek_at_stack(ctx, i as isize);
         let val_type = Type::from(stack_val);
 
-        match learned_mapping {
-            TempMapping::MapToSelf => {
+        match learned_mapping.get_kind() {
+            TempMappingKind::MapToSelf => {
                 if self_val != stack_val {
                     panic!(
                         "verify_ctx: stack value was mapped to self, but values did not match!\n  stack: {}\n  self: {}",
@@ -356,8 +358,8 @@ fn verify_ctx(jit: &JITState, ctx: &Context) {
                     );
                 }
             }
-            TempMapping::MapToLocal(local_idx) => {
-                let local_idx: u8 = local_idx.into();
+            TempMappingKind::MapToLocal => {
+                let local_idx: u8 = learned_mapping.get_local_idx().into();
                 let local_val = jit.peek_at_local(local_idx.into());
                 if local_val != stack_val {
                     panic!(
@@ -368,7 +370,7 @@ fn verify_ctx(jit: &JITState, ctx: &Context) {
                     );
                 }
             }
-            TempMapping::MapToStack => {}
+            TempMappingKind::MapToStack => {}
         }
 
         // If the actual type differs from the learned type
@@ -1009,9 +1011,9 @@ fn gen_dup(
     _ocb: &mut OutlinedCb,
 ) -> Option<CodegenStatus> {
     let dup_val = asm.stack_opnd(0);
-    let (mapping, tmp_type) = asm.ctx.get_opnd_mapping(dup_val.into());
+    let mapping = asm.ctx.get_opnd_mapping(dup_val.into());
 
-    let loc0 = asm.stack_push_mapping((mapping, tmp_type));
+    let loc0 = asm.stack_push_mapping(mapping);
     asm.mov(loc0, dup_val);
 
     Some(KeepCompiling)
@@ -2327,7 +2329,7 @@ fn gen_setinstancevariable(
         return None;
     }
 
-    let (_, stack_type) = asm.ctx.get_opnd_mapping(StackOpnd(0));
+    let stack_type = asm.ctx.get_opnd_type(StackOpnd(0));
 
     // Check if the comptime class uses a custom allocator
     let custom_allocator = unsafe { rb_get_alloc_func(comptime_val_klass) };
@@ -8968,8 +8970,8 @@ mod tests {
 
         let status = gen_swap(&mut jit, &mut asm, &mut ocb);
 
-        let (_, tmp_type_top) = asm.ctx.get_opnd_mapping(StackOpnd(0));
-        let (_, tmp_type_next) = asm.ctx.get_opnd_mapping(StackOpnd(1));
+        let tmp_type_top = asm.ctx.get_opnd_type(StackOpnd(0));
+        let tmp_type_next = asm.ctx.get_opnd_type(StackOpnd(1));
 
         assert_eq!(status, Some(KeepCompiling));
         assert_eq!(tmp_type_top, Type::Fixnum);
@@ -8981,7 +8983,7 @@ mod tests {
         let (mut jit, _context, mut asm, mut cb, mut ocb) = setup_codegen();
         let status = gen_putnil(&mut jit, &mut asm, &mut ocb);
 
-        let (_, tmp_type_top) = asm.ctx.get_opnd_mapping(StackOpnd(0));
+        let tmp_type_top = asm.ctx.get_opnd_type(StackOpnd(0));
 
         assert_eq!(status, Some(KeepCompiling));
         assert_eq!(tmp_type_top, Type::Nil);
@@ -9000,7 +9002,7 @@ mod tests {
 
         let status = gen_putobject(&mut jit, &mut asm, &mut ocb);
 
-        let (_, tmp_type_top) = asm.ctx.get_opnd_mapping(StackOpnd(0));
+        let tmp_type_top = asm.ctx.get_opnd_type(StackOpnd(0));
 
         assert_eq!(status, Some(KeepCompiling));
         assert_eq!(tmp_type_top, Type::True);
@@ -9020,7 +9022,7 @@ mod tests {
 
         let status = gen_putobject(&mut jit, &mut asm, &mut ocb);
 
-        let (_, tmp_type_top) = asm.ctx.get_opnd_mapping(StackOpnd(0));
+        let tmp_type_top = asm.ctx.get_opnd_type(StackOpnd(0));
 
         assert_eq!(status, Some(KeepCompiling));
         assert_eq!(tmp_type_top, Type::Fixnum);
@@ -9034,7 +9036,7 @@ mod tests {
         jit.opcode = YARVINSN_putobject_INT2FIX_0_.as_usize();
         let status = gen_putobject_int2fix(&mut jit, &mut asm, &mut ocb);
 
-        let (_, tmp_type_top) = asm.ctx.get_opnd_mapping(StackOpnd(0));
+        let tmp_type_top = asm.ctx.get_opnd_type(StackOpnd(0));
 
         // Right now we're not testing the generated machine code to make sure a literal 1 or 0 was pushed. I've checked locally.
         assert_eq!(status, Some(KeepCompiling));
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index c0577f08e8..805b7dd685 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -18,13 +18,14 @@ use std::cell::*;
 use std::collections::HashSet;
 use std::fmt;
 use std::mem;
+use std::mem::transmute;
 use std::ops::Range;
 use std::rc::Rc;
 use mem::MaybeUninit;
 use std::ptr;
 use ptr::NonNull;
 use YARVOpnd::*;
-use TempMapping::*;
+use TempMappingKind::*;
 use crate::invariants::*;
 
 // Maximum number of temp value types we keep track of
@@ -39,6 +40,7 @@ pub type IseqIdx = u16;
 
 // Represent the type of a value (local/stack/self) in YJIT
 #[derive(Copy, Clone, Hash, PartialEq, Eq, Debug)]
+#[repr(u8)]
 pub enum Type {
     Unknown,
     UnknownImm,
@@ -286,63 +288,92 @@ pub enum TypeDiff {
     Incompatible,
 }
 
-// Potential mapping of a value on the temporary stack to
-// self, a local variable or constant so that we can track its type
 #[derive(Copy, Clone, Eq, Hash, PartialEq, Debug)]
-pub enum TempMapping {
-    MapToStack, // Normal stack value
-    MapToSelf,  // Temp maps to the self operand
-    MapToLocal(LocalIndex), // Temp maps to a local variable with index
-                //ConstMapping,         // Small constant (0, 1, 2, Qnil, Qfalse, Qtrue)
+#[repr(u8)]
+pub enum TempMappingKind
+{
+    MapToStack = 0,
+    MapToSelf = 1,
+    MapToLocal = 2,
 }
 
-// Index used by MapToLocal. Using this instead of u8 makes TempMapping 1 byte.
+// Potential mapping of a value on the temporary stack to
+// self, a local variable or constant so that we can track its type
+//
+// The highest two bits represent TempMappingKind, and the rest of
+// the bits are used differently across different kinds.
+// * MapToStack: The lowest 5 bits are used for mapping Type.
+// * MapToSelf: The remaining bits are not used; the type is stored in self_type.
+// * MapToLocal: The lowest 3 bits store the index of a local variable.
 #[derive(Copy, Clone, Eq, Hash, PartialEq, Debug)]
-pub enum LocalIndex {
-    Local0,
-    Local1,
-    Local2,
-    Local3,
-    Local4,
-    Local5,
-    Local6,
-    Local7,
-}
+pub struct TempMapping(u8);
 
-impl From<LocalIndex> for u8 {
-    fn from(idx: LocalIndex) -> Self {
-        match idx {
-            LocalIndex::Local0 => 0,
-            LocalIndex::Local1 => 1,
-            LocalIndex::Local2 => 2,
-            LocalIndex::Local3 => 3,
-            LocalIndex::Local4 => 4,
-            LocalIndex::Local5 => 5,
-            LocalIndex::Local6 => 6,
-            LocalIndex::Local7 => 7,
-        }
+impl TempMapping {
+    pub fn map_to_stack(t: Type) -> TempMapping
+    {
+        let kind_bits = TempMappingKind::MapToStack as u8;
+        let type_bits = t as u8;
+        assert!(type_bits <= 0b11111);
+        let bits = (kind_bits << 6) | (type_bits & 0b11111);
+        TempMapping(bits)
     }
-}
 
-impl From<u8> for LocalIndex {
-    fn from(idx: u8) -> Self {
-        match idx {
-            0 => LocalIndex::Local0,
-            1 => LocalIndex::Local1,
-            2 => LocalIndex::Local2,
-            3 => LocalIndex::Local3,
-            4 => LocalIndex::Local4,
-            5 => LocalIndex::Local5,
-            6 => LocalIndex::Local6,
-            7 => LocalIndex::Local7,
-            _ => unreachable!("{idx} was larger than {MAX_LOCAL_TYPES}"),
+    pub fn map_to_self() -> TempMapping
+    {
+        let kind_bits = TempMappingKind::MapToSelf as u8;
+        let bits = kind_bits << 6;
+        TempMapping(bits)
+    }
+
+    pub fn map_to_local(local_idx: u8) -> TempMapping
+    {
+        let kind_bits = TempMappingKind::MapToLocal as u8;
+        assert!(local_idx <= 0b111);
+        let bits = (kind_bits << 6) | (local_idx & 0b111);
+        TempMapping(bits)
+    }
+
+    pub fn without_type(&self) -> TempMapping
+    {
+        if self.get_kind() != TempMappingKind::MapToStack {
+            return *self;
         }
+
+        TempMapping::map_to_stack(Type::Unknown)
+    }
+
+    pub fn get_kind(&self) -> TempMappingKind
+    {
+        // Take the two highest bits
+        let TempMapping(bits) = self;
+        let kind_bits = bits >> 6;
+        assert!(kind_bits <= 2);
+        unsafe { transmute::<u8, TempMappingKind>(kind_bits) }
+    }
+
+    pub fn get_type(&self) -> Type
+    {
+        assert!(self.get_kind() == TempMappingKind::MapToStack);
+
+        // Take the 5 lowest bits
+        let TempMapping(bits) = self;
+        let type_bits = bits & 0b11111;
+        unsafe { transmute::<u8, Type>(type_bits) }
+    }
+
+    pub fn get_local_idx(&self) -> u8
+    {
+        assert!(self.get_kind() == TempMappingKind::MapToLocal);
+
+        // Take the 3 lowest bits
+        let TempMapping(bits) = self;
+        bits & 0b111
     }
 }
 
 impl Default for TempMapping {
     fn default() -> Self {
-        MapToStack
+        TempMapping::map_to_stack(Type::Unknown)
     }
 }
 
@@ -425,9 +456,6 @@ pub struct Context {
     // Local variable types we keep track of
     local_types: [Type; MAX_LOCAL_TYPES],
 
-    // Temporary variable types we keep track of
-    temp_types: [Type; MAX_TEMP_TYPES],
-
     // Type we track for self
     self_type: Type,
 
@@ -1647,10 +1675,11 @@ impl Context {
 
                 let mapping = self.temp_mapping[stack_idx];
 
-                match mapping {
+                match mapping.get_kind() {
                     MapToSelf => self.self_type,
-                    MapToStack => self.temp_types[(self.stack_size - 1 - idx) as usize],
-                    MapToLocal(idx) => {
+                    MapToStack => mapping.get_type(),
+                    MapToLocal => {
+                        let idx = mapping.get_local_idx();
                         assert!((idx as usize) < MAX_LOCAL_TYPES);
                         return self.local_types[idx as usize];
                     }
@@ -1687,11 +1716,15 @@ impl Context {
 
                 let mapping = self.temp_mapping[stack_idx];
 
-                match mapping {
+                match mapping.get_kind() {
                     MapToSelf => self.self_type.upgrade(opnd_type),
-                    MapToStack => self.temp_types[stack_idx].upgrade(opnd_type),
-                    MapToLocal(idx) => {
-                        let idx = idx as usize;
+                    MapToStack => {
+                        let mut temp_type = mapping.get_type();
+                        temp_type.upgrade(opnd_type);
+                        self.temp_mapping[stack_idx] = TempMapping::map_to_stack(temp_type);
+                    }
+                    MapToLocal => {
+                        let idx = mapping.get_local_idx() as usize;
                         assert!(idx < MAX_LOCAL_TYPES);
                         self.local_types[idx].upgrade(opnd_type);
                     }
@@ -1705,29 +1738,29 @@ impl Context {
     This is can be used with stack_push_mapping or set_opnd_mapping to copy
     a stack value's type while maintaining the mapping.
     */
-    pub fn get_opnd_mapping(&self, opnd: YARVOpnd) -> (TempMapping, Type) {
+    pub fn get_opnd_mapping(&self, opnd: YARVOpnd) -> TempMapping {
         let opnd_type = self.get_opnd_type(opnd);
 
         match opnd {
-            SelfOpnd => (MapToSelf, opnd_type),
+            SelfOpnd => TempMapping::map_to_self(),
             StackOpnd(idx) => {
                 assert!(idx < self.stack_size);
                 let stack_idx = (self.stack_size - 1 - idx) as usize;
 
                 if stack_idx < MAX_TEMP_TYPES {
-                    (self.temp_mapping[stack_idx], opnd_type)
+                    self.temp_mapping[stack_idx]
                 } else {
                     // We can't know the source of this stack operand, so we assume it is
                     // a stack-only temporary. type will be UNKNOWN
                     assert!(opnd_type == Type::Unknown);
-                    (MapToStack, opnd_type)
+                    TempMapping::map_to_stack(opnd_type)
                 }
             }
         }
     }
 
     /// Overwrite both the type and mapping of a stack operand.
-    pub fn set_opnd_mapping(&mut self, opnd: YARVOpnd, (mapping, opnd_type): (TempMapping, Type)) {
+    pub fn set_opnd_mapping(&mut self, opnd: YARVOpnd, mapping: TempMapping) {
         match opnd {
             SelfOpnd => unreachable!("self always maps to self"),
             StackOpnd(idx) => {
@@ -1745,9 +1778,6 @@ impl Context {
                 }
 
                 self.temp_mapping[stack_idx] = mapping;
-
-                // Only used when mapping == MAP_STACK
-                self.temp_types[stack_idx] = opnd_type;
             }
         }
     }
@@ -1766,16 +1796,16 @@ impl Context {
         }
 
         // If any values on the stack map to this local we must detach them
-        for (i, mapping) in ctx.temp_mapping.iter_mut().enumerate() {
-            *mapping = match *mapping {
-                MapToStack => MapToStack,
-                MapToSelf => MapToSelf,
-                MapToLocal(idx) => {
+        for mapping in ctx.temp_mapping.iter_mut() {
+            *mapping = match mapping.get_kind() {
+                MapToStack => *mapping,
+                MapToSelf => *mapping,
+                MapToLocal => {
+                    let idx = mapping.get_local_idx();
                     if idx as usize == local_idx {
-                        ctx.temp_types[i] = ctx.local_types[idx as usize];
-                        MapToStack
+                        TempMapping::map_to_stack(ctx.local_types[idx as usize])
                     } else {
-                        MapToLocal(idx)
+                        TempMapping::map_to_local(idx)
                     }
                 }
             }
@@ -1789,14 +1819,10 @@ 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 (i, mapping) in self.temp_mapping.iter_mut().enumerate() {
-            *mapping = match *mapping {
-                MapToStack => MapToStack,
-                MapToSelf => MapToSelf,
-                MapToLocal(idx) => {
-                    self.temp_types[i] = self.local_types[idx as usize];
-                    MapToStack
-                }
+        for mapping in self.temp_mapping.iter_mut() {
+            if mapping.get_kind() == MapToLocal {
+                let idx = mapping.get_local_idx();
+                *mapping = TempMapping::map_to_stack(self.local_types[idx as usize]);
             }
         }
 
@@ -1853,12 +1879,12 @@ impl Context {
 
         // For each value on the temp stack
         for i in 0..src.stack_size {
-            let (src_mapping, src_type) = src.get_opnd_mapping(StackOpnd(i));
-            let (dst_mapping, dst_type) = dst.get_opnd_mapping(StackOpnd(i));
+            let src_mapping = src.get_opnd_mapping(StackOpnd(i));
+            let dst_mapping = dst.get_opnd_mapping(StackOpnd(i));
 
             // If the two mappings aren't the same
             if src_mapping != dst_mapping {
-                if dst_mapping == MapToStack {
+                if dst_mapping.get_kind() == MapToStack {
                     // We can safely drop information about the source of the temp
                     // stack operand.
                     diff += 1;
@@ -1867,6 +1893,9 @@ impl Context {
                 }
             }
 
+            let src_type = src.get_opnd_type(StackOpnd(i));
+            let dst_type = dst.get_opnd_type(StackOpnd(i));
+
             diff += match src_type.diff(dst_type) {
                 TypeDiff::Compatible(diff) => diff,
                 TypeDiff::Incompatible => return TypeDiff::Incompatible,
@@ -1896,10 +1925,10 @@ impl Context {
 impl Assembler {
     /// Push one new value on the temp stack with an explicit mapping
     /// Return a pointer to the new stack top
-    pub fn stack_push_mapping(&mut self, (mapping, temp_type): (TempMapping, Type)) -> Opnd {
+    pub fn stack_push_mapping(&mut self, mapping: TempMapping) -> Opnd {
         // If type propagation is disabled, store no types
         if get_option!(no_type_prop) {
-            return self.stack_push_mapping((mapping, Type::Unknown));
+            return self.stack_push_mapping(mapping.without_type());
         }
 
         let stack_size: usize = self.ctx.stack_size.into();
@@ -1907,9 +1936,9 @@ impl Assembler {
         // Keep track of the type and mapping of the value
         if stack_size < MAX_TEMP_TYPES {
             self.ctx.temp_mapping[stack_size] = mapping;
-            self.ctx.temp_types[stack_size] = temp_type;
 
-            if let MapToLocal(idx) = mapping {
+            if mapping.get_kind() == MapToLocal {
+                let idx = mapping.get_local_idx();
                 assert!((idx as usize) < MAX_LOCAL_TYPES);
             }
         }
@@ -1928,12 +1957,12 @@ impl Assembler {
     /// Push one new value on the temp stack
     /// Return a pointer to the new stack top
     pub fn stack_push(&mut self, val_type: Type) -> Opnd {
-        return self.stack_push_mapping((MapToStack, val_type));
+        return self.stack_push_mapping(TempMapping::map_to_stack(val_type));
     }
 
     /// Push the self value on the stack
     pub fn stack_push_self(&mut self) -> Opnd {
-        return self.stack_push_mapping((MapToSelf, Type::Unknown));
+        return self.stack_push_mapping(TempMapping::map_to_self());
     }
 
     /// Push a local variable on the stack
@@ -1942,7 +1971,7 @@ impl Assembler {
             return self.stack_push(Type::Unknown);
         }
 
-        return self.stack_push_mapping((MapToLocal((local_idx as u8).into()), Type::Unknown));
+        return self.stack_push_mapping(TempMapping::map_to_local((local_idx as u8).into()));
     }
 
     // Pop N values off the stack
@@ -1957,8 +1986,7 @@ impl Assembler {
             let idx: usize = (self.ctx.stack_size as usize) - i - 1;
 
             if idx < MAX_TEMP_TYPES {
-                self.ctx.temp_types[idx] = Type::Unknown;
-                self.ctx.temp_mapping[idx] = MapToStack;
+                self.ctx.temp_mapping[idx] = TempMapping::map_to_stack(Type::Unknown);
             }
         }
 
@@ -1976,7 +2004,6 @@ impl Assembler {
 
         for i in method_name_index..(self.ctx.stack_size - 1) as usize {
             if i + 1 < MAX_TEMP_TYPES {
-                self.ctx.temp_types[i] = self.ctx.temp_types[i + 1];
                 self.ctx.temp_mapping[i] = self.ctx.temp_mapping[i + 1];
             }
         }
@@ -3154,6 +3181,31 @@ mod tests {
     use crate::core::*;
 
     #[test]
+    fn tempmapping_size() {
+        assert_eq!(mem::size_of::<TempMapping>(), 1);
+    }
+
+    #[test]
+    fn tempmapping() {
+        let t = TempMapping::map_to_stack(Type::Unknown);
+        assert_eq!(t.get_kind(), MapToStack);
+        assert_eq!(t.get_type(), Type::Unknown);
+
+        let t = TempMapping::map_to_stack(Type::TString);
+        assert_eq!(t.get_kind(), MapToStack);
+        assert_eq!(t.get_type(), Type::TString);
+
+        let t = TempMapping::map_to_local(7);
+        assert_eq!(t.get_kind(), MapToLocal);
+        assert_eq!(t.get_local_idx(), 7);
+    }
+
+    #[test]
+    fn context_size() {
+        assert_eq!(mem::size_of::<Context>(), 21);
+    }
+
+    #[test]
     fn types() {
         // Valid src => dst
         assert_eq!(Type::Unknown.diff(Type::Unknown), TypeDiff::Compatible(0));