Reduce memory size of jit mappings

Author: Grarak

src/jit/assembler/arm/alu_assembler.rs

@@ -590,6 +590,33 @@ impl Bfi {
     }
 }
 
+#[bitsize(32)]
+#[derive(FromBits)]
+pub struct Ubfx {
+    rn: u4,
+    id: u3,
+    lsb: u5,
+    rd: u4,
+    widthm1: u5,
+    id2: u7,
+    cond: u4,
+}
+
+impl Ubfx {
+    #[inline]
+    pub fn create(rd: Reg, rn: Reg, lsb: u8, width: u8, cond: Cond) -> u32 {
+        u32::from(crate::jit::assembler::arm::alu_assembler::Bfi::new(
+            u4::new(rn as u8),
+            u3::new(0b101),
+            u5::new(lsb),
+            u4::new(rd as u8),
+            u5::new(width - 1),
+            u7::new(0b0111111),
+            u4::new(cond as u8),
+        ))
+    }
+}
+
 // TODO Add const asserts once const features has been added back to rust
 // https://github.com/rust-lang/rust/issues/110395
 //const_assert_eq!(lookup_opcode(Self::add(0, 0, 0, 0)).0 as u8, And as u8);

src/jit/assembler/block_asm.rs

@@ -350,7 +350,15 @@ impl<'a> BlockAsm<'a> {
             operands: [op0.into(), op1.into()],
             lsb,
             width,
-        })
+        });
+    }
+
+    pub fn ubfx(&mut self, op0: impl Into<BlockReg>, op1: impl Into<BlockReg>, lsb: u8, width: u8) {
+        self.insert_inst(BlockInstKind::Ubfx {
+            operands: [op0.into(), op1.into()],
+            lsb,
+            width,
+        });
     }
 
     pub fn muls_guest_thumb_pc_aligned(&mut self, op0: impl Into<BlockReg>, op1: impl Into<BlockReg>, op2: impl Into<BlockOperandShift>) {

src/jit/assembler/block_inst.rs

@@ -1,4 +1,4 @@
-use crate::jit::assembler::arm::alu_assembler::{AluImm, AluReg, AluShiftImm, Bfc, Bfi, MulReg};
+use crate::jit::assembler::arm::alu_assembler::{AluImm, AluReg, AluShiftImm, Bfc, Bfi, MulReg, Ubfx};
 use crate::jit::assembler::arm::branch_assembler::Bx;
 use crate::jit::assembler::arm::transfer_assembler::{LdmStm, LdrStrImm, LdrStrImmSBHD, LdrStrReg, LdrStrRegSBHD, Mrs, Msr};
 use crate::jit::assembler::arm::Bkpt;
@@ -168,6 +168,11 @@ pub enum BlockInstKind {
         lsb: u8,
         width: u8,
     },
+    Ubfx {
+        operands: [BlockReg; 2],
+        lsb: u8,
+        width: u8,
+    },
     Mul {
         operands: [BlockOperandShift; 3],
         set_cond: BlockAluSetCond,
@@ -287,7 +292,7 @@ impl BlockInstKind {
                 BlockSystemRegOp::Msr => (block_reg_set!(operand.try_as_reg()), block_reg_set!()),
             },
             BlockInstKind::Bfc { operand, .. } => (block_reg_set!(Some(*operand)), block_reg_set!(Some(*operand))),
-            BlockInstKind::Bfi { operands, .. } => (block_reg_set!(Some(operands[0]), Some(operands[1])), block_reg_set!(Some(operands[0]))),
+            BlockInstKind::Bfi { operands, .. } | BlockInstKind::Ubfx { operands, .. } => (block_reg_set!(Some(operands[0]), Some(operands[1])), block_reg_set!(Some(operands[0]))),
 
             BlockInstKind::SaveContext { .. } => (block_reg_set!(), block_reg_set!()),
             BlockInstKind::SaveReg {
@@ -433,6 +438,7 @@ impl BlockInstKind {
                 Self::replace_reg(&mut operands[0], old, new);
                 Self::replace_reg(&mut operands[1], old, new);
             }
+            BlockInstKind::Ubfx { operands, .. } => Self::replace_reg(&mut operands[1], old, new),
             BlockInstKind::SaveContext { .. } => {
                 unreachable!()
             }
@@ -485,7 +491,7 @@ impl BlockInstKind {
                 }
             }
             BlockInstKind::Bfc { operand, .. } => Self::replace_reg(operand, old, new),
-            BlockInstKind::Bfi { operands, .. } => Self::replace_reg(&mut operands[0], old, new),
+            BlockInstKind::Bfi { operands, .. } | BlockInstKind::Ubfx { operands, .. } => Self::replace_reg(&mut operands[0], old, new),
             BlockInstKind::SaveContext { .. } => {}
             BlockInstKind::SaveReg { guest_reg, reg_mapped, .. } => {
                 if *guest_reg == Reg::CPSR {
@@ -659,6 +665,7 @@ impl BlockInstKind {
             },
             BlockInstKind::Bfc { operand, lsb, width } => opcodes.push(Bfc::create(operand.as_fixed(), *lsb, *width, Cond::AL)),
             BlockInstKind::Bfi { operands, lsb, width } => opcodes.push(Bfi::create(operands[0].as_fixed(), operands[1].as_fixed(), *lsb, *width, Cond::AL)),
+            BlockInstKind::Ubfx { operands, lsb, width } => opcodes.push(Ubfx::create(operands[0].as_fixed(), operands[1].as_fixed(), *lsb, *width, Cond::AL)),
             BlockInstKind::Mul { operands, set_cond, .. } => match operands[2].operand {
                 BlockOperand::Reg(reg) => opcodes.push(MulReg::mul(
                     operands[0].as_reg().as_fixed(),
@@ -844,6 +851,7 @@ impl Debug for BlockInstKind {
             BlockInstKind::SystemReg { op, operand } => write!(f, "{op:?} {operand:?}"),
             BlockInstKind::Bfc { operand, lsb, width } => write!(f, "Bfc {operand:?}, {lsb}, {width}"),
             BlockInstKind::Bfi { operands, lsb, width } => write!(f, "Bfi {:?}, {:?}, {lsb}, {width}", operands[0], operands[1]),
+            BlockInstKind::Ubfx { operands, lsb, width } => write!(f, "Ubfx {:?}, {:?}, {lsb}, {width}", operands[0], operands[1]),
             BlockInstKind::Mul { operands, set_cond, thumb_pc_aligned } => write!(f, "Mul{set_cond:?} {operands:?}, align pc: {thumb_pc_aligned}"),
             BlockInstKind::Label { label, guest_pc } => {
                 let guest_pc = match guest_pc {

src/jit/jit_asm_common_funs.rs

@@ -51,7 +51,8 @@ impl<const CPU: CpuType> JitAsmCommonFuns<CPU> {
         let map_entry_base_ptr_reg = block_asm.new_reg();
 
         block_asm.mov(map_ptr_reg, map_ptr as u32);
-        block_asm.mov(map_index_reg, (aligned_target_reg.into(), ShiftType::Lsr, BlockOperand::from(jit_memory_map::BLOCK_SHIFT as u32 + 1)));
+        const ADDR_SHIFT: u8 = jit_memory_map::BLOCK_SHIFT as u8 + 1;
+        block_asm.ubfx(map_index_reg, aligned_target_reg, ADDR_SHIFT, 28 - ADDR_SHIFT);
         block_asm.load_u32(map_entry_base_ptr_reg, map_ptr_reg, (map_index_reg.into(), ShiftType::Lsl, BlockOperand::from(2)));
         let block_size_mask_reg = map_index_reg;
         block_asm.mov(block_size_mask_reg, (jit_memory_map::BLOCK_SIZE as u32 - 1) << 2);

src/jit/jit_memory.rs

@@ -5,7 +5,7 @@ use crate::jit::jit_memory_map::JitMemoryMap;
 use crate::logging::debug_println;
 use crate::mmap::Mmap;
 use crate::utils;
-use crate::utils::{HeapMem, HeapMemU32};
+use crate::utils::{HeapMem, HeapMemU8};
 use lazy_static::lazy_static;
 use paste::paste;
 use std::intrinsics::unlikely;
@@ -89,10 +89,10 @@ create_jit_blocks!(
 
 #[derive(Default)]
 pub struct JitLiveRanges {
-    pub itcm: HeapMemU32<{ (regions::INSTRUCTION_TCM_SIZE / JIT_LIVE_RANGE_PAGE_SIZE / 32) as usize }>,
-    pub main: HeapMemU32<{ (regions::MAIN_MEMORY_SIZE / JIT_LIVE_RANGE_PAGE_SIZE / 32) as usize }>,
-    pub wram: HeapMemU32<{ ((regions::SHARED_WRAM_SIZE + regions::ARM7_WRAM_SIZE) / JIT_LIVE_RANGE_PAGE_SIZE / 32) as usize }>,
-    pub vram_arm7: HeapMemU32<{ (vram::ARM7_SIZE / JIT_LIVE_RANGE_PAGE_SIZE / 32) as usize }>,
+    pub itcm: HeapMemU8<{ (regions::INSTRUCTION_TCM_SIZE / JIT_LIVE_RANGE_PAGE_SIZE / 8) as usize }>,
+    pub main: HeapMemU8<{ (regions::MAIN_MEMORY_SIZE / JIT_LIVE_RANGE_PAGE_SIZE / 8) as usize }>,
+    pub wram: HeapMemU8<{ ((regions::SHARED_WRAM_SIZE + regions::ARM7_WRAM_SIZE) / JIT_LIVE_RANGE_PAGE_SIZE / 8) as usize }>,
+    pub vram_arm7: HeapMemU8<{ (vram::ARM7_SIZE / JIT_LIVE_RANGE_PAGE_SIZE / 8) as usize }>,
 }
 
 #[cfg(target_os = "linux")]
@@ -191,10 +191,10 @@ impl JitMemory {
                 $entries[entries_index] = JitEntry(jit_entry_addr);
                 assert_eq!(ptr::addr_of!($entries[entries_index]), self.jit_memory_map.get_jit_entry::<CPU>(guest_pc));
 
-                // >> 5 for u32 (each bit represents a page)
-                let live_ranges_index = ((guest_pc >> JIT_LIVE_RANGE_PAGE_SIZE_SHIFT) >> 5) as usize;
+                // >> 3 for u8 (each bit represents a page)
+                let live_ranges_index = ((guest_pc >> JIT_LIVE_RANGE_PAGE_SIZE_SHIFT) >> 3) as usize;
                 let live_ranges_index = live_ranges_index % $live_ranges.len();
-                let live_ranges_bit = (guest_pc >> JIT_LIVE_RANGE_PAGE_SIZE_SHIFT) & 31;
+                let live_ranges_bit = (guest_pc >> JIT_LIVE_RANGE_PAGE_SIZE_SHIFT) & 0x7;
                 $live_ranges[live_ranges_index] |= 1 << live_ranges_bit;
                 assert_eq!(ptr::addr_of!($live_ranges[live_ranges_index]), self.jit_memory_map.get_live_range::<CPU>(guest_pc));
 
@@ -212,7 +212,7 @@ impl JitMemory {
             }};
         }
 
-        let jit_addr = match CPU {
+        match CPU {
             ARM9 => match guest_pc & 0xFF000000 {
                 regions::INSTRUCTION_TCM_OFFSET | regions::INSTRUCTION_TCM_MIRROR_OFFSET => insert!(self.jit_entries.itcm, self.jit_live_ranges.itcm),
                 regions::MAIN_MEMORY_OFFSET => insert!(self.jit_entries.main_arm9, self.jit_live_ranges.main),
@@ -224,9 +224,7 @@ impl JitMemory {
                 regions::VRAM_OFFSET => insert!(self.jit_entries.vram_arm7, self.jit_live_ranges.vram_arm7),
                 _ => todo!("{:x}", guest_pc),
             },
-        };
-
-        jit_addr
+        }
     }
 
     pub fn get_jit_start_addr<const CPU: CpuType>(&self, guest_pc: u32) -> *const extern "C" fn(bool) {
@@ -237,7 +235,7 @@ impl JitMemory {
         macro_rules! invalidate {
             ($guest_addr:expr, $live_range:ident, $cpu:expr, [$(($cpu_entry:expr, $entries:ident)),+]) => {{
                 let live_range = unsafe { self.jit_memory_map.get_live_range::<{ $cpu }>($guest_addr).as_mut_unchecked() };
-                let live_ranges_bit = ($guest_addr >> JIT_LIVE_RANGE_PAGE_SIZE_SHIFT) & 31;
+                let live_ranges_bit = ($guest_addr >> JIT_LIVE_RANGE_PAGE_SIZE_SHIFT) & 0x7;
                 if unlikely(*live_range & (1 << live_ranges_bit) != 0) {
                     *live_range &= !(1 << live_ranges_bit);
 

src/jit/jit_memory_map.rs

@@ -4,17 +4,23 @@ use crate::jit::jit_memory::{JitEntries, JitEntry, JitLiveRanges, BIOS_UNINTERRU
 use crate::utils::HeapMemU32;
 use CpuType::{ARM7, ARM9};
 
+// ARM9 Bios starts at 0xFFFF0000, but just treat everything above OAM region as bios
+// Also omit 0xF msb to save more memory
+const MEMORY_RANGE_ARM9: u32 = 0x10000000;
+const MEMORY_RANGE_ARM7: u32 = regions::OAM_OFFSET;
+
 pub const BLOCK_SHIFT: usize = 13;
 pub const BLOCK_SIZE: usize = 1 << BLOCK_SHIFT;
-const SIZE: usize = (1 << 31) / BLOCK_SIZE;
-const LIVE_RANGES_SIZE: usize = 1 << (32 - JIT_LIVE_RANGE_PAGE_SIZE_SHIFT - 5);
+const SIZE_ARM9: usize = (MEMORY_RANGE_ARM9 >> 1) as usize / BLOCK_SIZE;
+const SIZE_ARM7: usize = (MEMORY_RANGE_ARM7 >> 1) as usize / BLOCK_SIZE;
+const LIVE_RANGES_SIZE: usize = (regions::OAM_OFFSET >> (JIT_LIVE_RANGE_PAGE_SIZE_SHIFT + 3)) as usize;
 
 const BIOS_UNINTERRUPT_ENTRIES_ARM9: [JitEntry; BLOCK_SIZE] = [BIOS_UNINTERRUPT_ENTRY_ARM9; BLOCK_SIZE];
 const BIOS_UNINTERRUPT_ENTRIES_ARM7: [JitEntry; BLOCK_SIZE] = [BIOS_UNINTERRUPT_ENTRY_ARM7; BLOCK_SIZE];
 
 pub struct JitMemoryMap {
-    map_arm9: HeapMemU32<SIZE>,
-    map_arm7: HeapMemU32<SIZE>,
+    map_arm9: HeapMemU32<SIZE_ARM9>,
+    map_arm7: HeapMemU32<SIZE_ARM7>,
     live_ranges_map_arm9: HeapMemU32<LIVE_RANGES_SIZE>,
     live_ranges_map_arm7: HeapMemU32<LIVE_RANGES_SIZE>,
 }
@@ -28,36 +34,39 @@ impl JitMemoryMap {
             live_ranges_map_arm7: HeapMemU32::new(),
         };
 
-        for i in 0..SIZE {
+        macro_rules! get_ptr {
+            ($addr:expr, $entries:expr) => {{
+                (unsafe { $entries.as_ptr().add(($addr >> 1) % $entries.len()) } as u32)
+            }};
+        }
+
+        for i in 0..SIZE_ARM9 {
             let addr = (i << BLOCK_SHIFT) << 1;
             let arm9_ptr = &mut instance.map_arm9[i];
-            let arm7_ptr = &mut instance.map_arm7[i];
 
-            macro_rules! get_ptr {
-                ($entries:expr) => {{
-                    (unsafe { $entries.as_ptr().add((addr >> 1) % $entries.len()) } as u32)
-                }};
+            match (addr as u32) & 0x0F000000 {
+                regions::INSTRUCTION_TCM_OFFSET | regions::INSTRUCTION_TCM_MIRROR_OFFSET => *arm9_ptr = get_ptr!(addr, entries.itcm),
+                regions::MAIN_MEMORY_OFFSET => *arm9_ptr = get_ptr!(addr, entries.main_arm9),
+                0x0F000000 => *arm9_ptr = BIOS_UNINTERRUPT_ENTRIES_ARM9.as_ptr() as u32,
+                _ => {}
             }
+        }
 
-            match (addr as u32) & 0xFF000000 {
-                0 => {
-                    *arm9_ptr = get_ptr!(entries.itcm);
-                    *arm7_ptr = BIOS_UNINTERRUPT_ENTRIES_ARM7.as_ptr() as u32;
-                }
-                regions::INSTRUCTION_TCM_MIRROR_OFFSET => *arm9_ptr = get_ptr!(entries.itcm),
-                regions::MAIN_MEMORY_OFFSET => {
-                    *arm9_ptr = get_ptr!(entries.main_arm9);
-                    *arm7_ptr = get_ptr!(entries.main_arm7);
-                }
-                regions::SHARED_WRAM_OFFSET => *arm7_ptr = get_ptr!(entries.wram),
-                regions::VRAM_OFFSET => *arm7_ptr = get_ptr!(entries.vram_arm7),
-                0xFF000000 => *arm9_ptr = BIOS_UNINTERRUPT_ENTRIES_ARM9.as_ptr() as u32,
+        for i in 0..SIZE_ARM7 {
+            let addr = (i << BLOCK_SHIFT) << 1;
+            let arm7_ptr = &mut instance.map_arm7[i];
+
+            match (addr as u32) & 0x0F000000 {
+                0 => *arm7_ptr = BIOS_UNINTERRUPT_ENTRIES_ARM7.as_ptr() as u32,
+                regions::MAIN_MEMORY_OFFSET => *arm7_ptr = get_ptr!(addr, entries.main_arm7),
+                regions::SHARED_WRAM_OFFSET => *arm7_ptr = get_ptr!(addr, entries.wram),
+                regions::VRAM_OFFSET => *arm7_ptr = get_ptr!(addr, entries.vram_arm7),
                 _ => {}
             }
         }
 
         for i in 0..LIVE_RANGES_SIZE {
-            let addr = i << (JIT_LIVE_RANGE_PAGE_SIZE_SHIFT + 5);
+            let addr = i << (JIT_LIVE_RANGE_PAGE_SIZE_SHIFT + 3);
             let arm9_ptr = &mut instance.live_ranges_map_arm9[i];
             let arm7_ptr = &mut instance.live_ranges_map_arm7[i];
 
@@ -68,9 +77,7 @@ impl JitMemoryMap {
             }
 
             match (addr as u32) & 0xFF000000 {
-                0 => {
-                    *arm9_ptr = get_ptr!(live_ranges.itcm);
-                }
+                0 => *arm9_ptr = get_ptr!(live_ranges.itcm),
                 regions::INSTRUCTION_TCM_MIRROR_OFFSET => *arm9_ptr = get_ptr!(live_ranges.itcm),
                 regions::MAIN_MEMORY_OFFSET => {
                     *arm9_ptr = get_ptr!(live_ranges.main);
@@ -86,7 +93,7 @@ impl JitMemoryMap {
     }
 
     pub fn get_jit_entry<const CPU: CpuType>(&self, addr: u32) -> *mut JitEntry {
-        let addr = addr >> 1;
+        let addr = (addr & 0x0FFFFFFF) >> 1;
         macro_rules! get_jit_entry {
             ($map:expr) => {{
                 unsafe { ($map[(addr >> BLOCK_SHIFT) as usize] as *mut JitEntry).add((addr as usize) & (BLOCK_SIZE - 1)) }
@@ -98,10 +105,10 @@ impl JitMemoryMap {
         }
     }
 
-    pub fn get_live_range<const CPU: CpuType>(&self, addr: u32) -> *mut u32 {
+    pub fn get_live_range<const CPU: CpuType>(&self, addr: u32) -> *mut u8 {
         match CPU {
-            ARM9 => self.live_ranges_map_arm9[(addr >> (JIT_LIVE_RANGE_PAGE_SIZE_SHIFT + 5)) as usize] as _,
-            ARM7 => self.live_ranges_map_arm7[(addr >> (JIT_LIVE_RANGE_PAGE_SIZE_SHIFT + 5)) as usize] as _,
+            ARM9 => self.live_ranges_map_arm9[(addr >> (JIT_LIVE_RANGE_PAGE_SIZE_SHIFT + 3)) as usize] as _,
+            ARM7 => self.live_ranges_map_arm7[(addr >> (JIT_LIVE_RANGE_PAGE_SIZE_SHIFT + 3)) as usize] as _,
         }
     }