Rework block invalidation

* Optimize jit addr lookup

Author: Grarak

Cargo.lock

@@ -15,12 +15,13 @@ assets = "./static"
 [dependencies]
 bilge = { git = "https://github.com/Grarak/bilge.git" }
 chrono = "0.4.38"
-paste = "1.0.15"
-static_assertions = "1.1.0"
 dsvita_macros = { path = "macros" }
 gl = { git = "https://github.com/bjz/gl-rs" }
 glyph_brush = "0.7.9"
+lazy_static = "1.5.0"
 rust-ini = "0.21.1"
+static_assertions = "1.1.0"
+paste = "1.0.15"
 
 [target.armv7-unknown-linux-gnueabihf.dependencies]
 clap = { version = "4.5.17", features = ["cargo"] }

Cargo.toml

@@ -25,6 +25,6 @@ impl Main {
     }
 
     pub fn write_slice<T: Convert>(&mut self, addr_offset: u32, slice: &[T]) {
-        utils::write_to_mem_slice(self.main.deref_mut(), addr_offset & (regions::MAIN_MEMORY_SIZE - 1), slice);
+        utils::write_to_mem_slice(self.main.deref_mut(), (addr_offset & (regions::MAIN_MEMORY_SIZE - 1)) as usize, slice);
     }
 }

src/core/memory/main.rs

@@ -1,5 +1,5 @@
 use crate::core::cp15::TcmState;
-use crate::core::emu::{get_cp15, Emu};
+use crate::core::emu::{get_cp15, get_regs, Emu};
 use crate::core::memory::bios::{BiosArm7, BiosArm9};
 use crate::core::memory::io_arm7::IoArm7;
 use crate::core::memory::io_arm9::IoArm9;
@@ -15,7 +15,7 @@ use crate::core::memory::wram::Wram;
 use crate::core::spu::SoundSampler;
 use crate::core::CpuType;
 use crate::core::CpuType::ARM9;
-use crate::jit::jit_memory::JitMemory;
+use crate::jit::jit_memory::{JitMemory, JitRegion};
 use crate::logging::debug_println;
 use crate::utils::Convert;
 use std::intrinsics::{likely, unlikely};
@@ -36,7 +36,6 @@ pub struct Memory {
     pub jit: JitMemory,
     pub bios_arm9: BiosArm9,
     pub bios_arm7: BiosArm7,
-    pub current_jit_block_addr: u32, // Check if the jit block we are currently executing gets written to
     pub breakout_imm: bool,
     pub mmu_arm9: MmuArm9,
     pub mmu_arm7: MmuArm7,
@@ -66,7 +65,6 @@ impl Memory {
             jit: JitMemory::new(),
             bios_arm9: BiosArm9::new(),
             bios_arm7: BiosArm7::new(),
-            current_jit_block_addr: 0,
             breakout_imm: false,
             mmu_arm9: MmuArm9::new(),
             mmu_arm7: MmuArm7::new(),
@@ -199,8 +197,6 @@ impl Memory {
             }
         }
 
-        let mut invalidate_jit = || self.breakout_imm = self.jit.invalidate_block::<CPU>(aligned_addr, size_of::<T>(), self.current_jit_block_addr);
-
         match addr_base {
             regions::INSTRUCTION_TCM_OFFSET | regions::INSTRUCTION_TCM_MIRROR_OFFSET => match CPU {
                 ARM9 => {
@@ -209,7 +205,7 @@ impl Memory {
                         if aligned_addr < cp15.itcm_size && cp15.itcm_state != TcmState::Disabled {
                             self.tcm.write_itcm(aligned_addr, value);
                             debug_println!("{:?} itcm write at {:x} with value {:x}", CPU, aligned_addr, value.into(),);
-                            invalidate_jit();
+                            self.breakout_imm = self.jit.invalidate_block::<{ JitRegion::Itcm }>(aligned_addr, size_of::<T>(), get_regs!(emu, CPU).pc);
                         }
                     }
                 }
@@ -220,12 +216,12 @@ impl Memory {
             },
             regions::MAIN_MEMORY_OFFSET => {
                 self.main.write(addr_offset, value);
-                invalidate_jit();
+                self.breakout_imm = self.jit.invalidate_block::<{ JitRegion::Main }>(aligned_addr, size_of::<T>(), get_regs!(emu, CPU).pc);
             }
             regions::SHARED_WRAM_OFFSET => {
                 self.wram.write::<CPU, _>(addr_offset, value);
                 if CPU == ARM7 {
-                    invalidate_jit();
+                    self.breakout_imm = self.jit.invalidate_block::<{ JitRegion::Wram }>(aligned_addr, size_of::<T>(), get_regs!(emu, CPU).pc);
                 }
             }
             regions::IO_PORTS_OFFSET => match CPU {
@@ -247,7 +243,7 @@ impl Memory {
             regions::VRAM_OFFSET => {
                 self.vram.write::<CPU, _>(addr_offset, value);
                 if CPU == ARM7 {
-                    invalidate_jit();
+                    self.breakout_imm = self.jit.invalidate_block::<{ JitRegion::VramArm7 }>(aligned_addr, size_of::<T>(), get_regs!(emu, CPU).pc);
                 }
             }
             regions::OAM_OFFSET => self.oam.write(addr_offset, value),

src/core/memory/mem.rs

@@ -135,25 +135,25 @@ impl Rtc {
     fn update_date_time(&mut self) {
         let local_now = chrono::Local::now();
 
-        let year = local_now.year() as u32 % 100;
-        let month = local_now.month() as u8;
-        let day = local_now.day() as u8;
-        let (hour, is_pm) = {
-            let hour = local_now.hour();
-            ((if self.cnt & 0x2 == 0 { hour % 12 } else { hour }) as u8, hour >= 12)
-        };
-        let min = local_now.minute() as u8;
-        let sec = local_now.second() as u8;
-
-        // let year = 2000 as u32 % 100;
-        // let month = 1 as u8;
-        // let day = 1 as u8;
+        // let year = local_now.year() as u32 % 100;
+        // let month = local_now.month() as u8;
+        // let day = local_now.day() as u8;
         // let (hour, is_pm) = {
-        //     let hour = 11;
+        //     let hour = local_now.hour();
         //     ((if self.cnt & 0x2 == 0 { hour % 12 } else { hour }) as u8, hour >= 12)
         // };
-        // let min = 0 as u8;
-        // let sec = 0 as u8;
+        // let min = local_now.minute() as u8;
+        // let sec = local_now.second() as u8;
+
+        let year = 2000 as u32 % 100;
+        let month = 1 as u8;
+        let day = 1 as u8;
+        let (hour, is_pm) = {
+            let hour = 11;
+            ((if self.cnt & 0x2 == 0 { hour % 12 } else { hour }) as u8, hour >= 12)
+        };
+        let min = 0 as u8;
+        let sec = 0 as u8;
 
         self.date_time[0] = (((year / 10) << 4) | (year % 10)) as u8;
         self.date_time[1] = ((month / 10) << 4) | (month % 10);

src/core/rtc.rs

@@ -376,7 +376,7 @@ impl BasicBlock {
         }
     }
 
-    pub fn emit_opcodes(&self, asm: &mut BlockAsm, branch_placeholders: &mut Vec<usize>, opcodes_offset: usize) -> Vec<u32> {
+    pub fn emit_opcodes(&self, asm: &mut BlockAsm, branch_placeholders: &mut Vec<usize>, opcodes_offset: usize, used_host_regs: RegReserve) -> Vec<u32> {
         let mut opcodes = Vec::new();
         let mut inst_opcodes = Vec::new();
         for entry in self.insts_link.iter() {
@@ -397,7 +397,7 @@ impl BasicBlock {
             }
 
             inst_opcodes.clear();
-            inst.emit_opcode(&mut inst_opcodes, opcodes.len(), branch_placeholders, opcodes_offset);
+            inst.emit_opcode(&mut inst_opcodes, opcodes.len(), branch_placeholders, opcodes_offset, used_host_regs);
             opcodes.extend(&inst_opcodes);
         }
         opcodes

src/jit/assembler/basic_block.rs

@@ -6,7 +6,7 @@ use crate::jit::assembler::block_inst_list::BlockInstList;
 use crate::jit::assembler::block_reg_set::BlockRegSet;
 use crate::jit::assembler::{block_reg_allocator, BlockAsmBuf, BlockInst, BlockLabel, BlockOperand, BlockOperandShift, BlockReg, ANY_REG_LIMIT};
 use crate::jit::inst_info::InstInfo;
-use crate::jit::reg::{reg_reserve, Reg, RegReserve};
+use crate::jit::reg::{Reg, RegReserve};
 use crate::jit::{Cond, MemoryAmount, ShiftType};
 use crate::utils::{NoHashMap, NoHashSet};
 
@@ -55,7 +55,7 @@ pub struct BlockAsm<'a> {
 }
 
 impl<'a> BlockAsm<'a> {
-    pub fn new(thread_regs: &ThreadRegs, buf: &'a mut BlockAsmBuf) -> Self {
+    pub fn new(thread_regs: &ThreadRegs, host_sp_addr: usize, buf: &'a mut BlockAsmBuf) -> Self {
         buf.insts.clear();
         buf.guest_branches_mapping.clear();
         let thread_regs_addr_reg = BlockReg::Any(Reg::SPSR as u16 + 1);
@@ -82,9 +82,18 @@ impl<'a> BlockAsm<'a> {
             block_start: 0,
             inst_insert_index: None,
         };
-        instance.transfer_push(BlockReg::Fixed(Reg::SP), reg_reserve!(Reg::LR));
+
+        instance.buf.insts.push(BlockInst::Prologue);
+
+        let host_sp_addr_reg = thread_regs_addr_reg;
+        instance.mov(host_sp_addr_reg, host_sp_addr as u32);
+        instance.transfer_write(BlockReg::Fixed(Reg::SP), host_sp_addr_reg, 0, false, MemoryAmount::Word);
+
         instance.sub(BlockReg::Fixed(Reg::SP), BlockReg::Fixed(Reg::SP), ANY_REG_LIMIT as u32 * 4); // Reserve for spilled registers
+
         instance.mov(thread_regs_addr_reg, thread_regs.get_reg_start_addr() as *const _ as u32);
+        instance.restore_reg(Reg::CPSR);
+
         instance.block_start = instance.buf.insts.len();
         instance
     }
@@ -340,9 +349,9 @@ impl<'a> BlockAsm<'a> {
         });
     }
 
-    pub fn breakout(&mut self) {
-        self.add(BlockReg::Fixed(Reg::SP), BlockReg::Fixed(Reg::SP), ANY_REG_LIMIT as u32 * 4); // Remove reserve for spilled registers
-        self.transfer_pop(BlockReg::Fixed(Reg::SP), reg_reserve!(Reg::PC));
+    pub fn epilogue(&mut self) {
+        self.add(BlockReg::Fixed(Reg::SP), BlockReg::Fixed(Reg::SP), ANY_REG_LIMIT as u32 * 4);
+        self.buf.insts.push(BlockInst::Epilogue);
     }
 
     pub fn call(&mut self, func: impl Into<BlockOperand>) {
@@ -586,23 +595,13 @@ impl<'a> BlockAsm<'a> {
                     }
                     last_label = None;
                 }
-                // Breakout command
-                BlockInst::TransferMultiple {
-                    op: BlockTransferOp::Read,
-                    operand: BlockReg::Fixed(Reg::SP),
-                    regs,
-                    write_back: true,
-                    pre: false,
-                    add_to_base: true,
-                } => {
-                    if *regs == reg_reserve!(Reg::PC) {
-                        basic_blocks.push(BasicBlock::new(self, basic_block_start, current_node));
-                        basic_block_start = BlockInstList::deref(current_node).next;
-                        if let Some(last_label) = last_label {
-                            basic_block_label_mapping.insert(last_label.0, basic_blocks.len() - 1);
-                        }
-                        last_label = None;
+                BlockInst::Epilogue => {
+                    basic_blocks.push(BasicBlock::new(self, basic_block_start, current_node));
+                    basic_block_start = BlockInstList::deref(current_node).next;
+                    if let Some(last_label) = last_label {
+                        basic_block_label_mapping.insert(last_label.0, basic_blocks.len() - 1);
                     }
+                    last_label = None;
                 }
                 _ => {}
             }
@@ -631,15 +630,8 @@ impl<'a> BlockAsm<'a> {
                     }
                 }
                 // Don't add exit when last command in basic block is a breakout
-                BlockInst::TransferMultiple {
-                    op: BlockTransferOp::Read,
-                    operand: BlockReg::Fixed(Reg::SP),
-                    regs,
-                    write_back: true,
-                    pre: false,
-                    add_to_base: true,
-                } => {
-                    if *regs != reg_reserve!(Reg::PC) && i + 1 < basic_blocks_len {
+                BlockInst::Epilogue => {
+                    if i + 1 < basic_blocks_len {
                         basic_block.exit_blocks.insert(i + 1);
                     }
                 }
@@ -764,14 +756,19 @@ impl<'a> BlockAsm<'a> {
             self.buf.reg_allocator.global_mapping.insert(reg, Reg::None);
         }
 
+        let mut used_host_regs = RegReserve::new();
+
         for (i, basic_block) in basic_blocks.iter_mut().enumerate() {
             if i != 0 && basic_block.enter_blocks.is_empty() {
                 continue;
             }
             self.buf.reg_allocator.init_inputs(*basic_block.get_required_inputs());
             basic_block.allocate_regs(&mut self);
+            used_host_regs += self.buf.reg_allocator.dirty_regs;
         }
 
+        used_host_regs &= block_reg_allocator::ALLOCATION_REGS;
+
         let mut opcodes = Vec::new();
         let mut branch_placeholders = Vec::new();
         let mut opcodes_offset = Vec::with_capacity(basic_blocks.len());
@@ -780,7 +777,7 @@ impl<'a> BlockAsm<'a> {
                 continue;
             }
             opcodes_offset.push(opcodes.len());
-            opcodes.extend(basic_block.emit_opcodes(&mut self, &mut branch_placeholders, opcodes.len()));
+            opcodes.extend(basic_block.emit_opcodes(&mut self, &mut branch_placeholders, opcodes.len(), used_host_regs));
         }
 
         for branch_placeholder in branch_placeholders {

src/jit/assembler/block_asm.rs

@@ -3,7 +3,7 @@ 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;
 use crate::jit::assembler::block_reg_set::{block_reg_set, BlockRegSet};
-use crate::jit::assembler::{BlockLabel, BlockOperand, BlockOperandShift, BlockReg, BlockShift};
+use crate::jit::assembler::{block_reg_allocator, BlockLabel, BlockOperand, BlockOperandShift, BlockReg, BlockShift};
 use crate::jit::inst_info::{InstInfo, Operand, Shift, ShiftValue};
 use crate::jit::reg::{Reg, RegReserve};
 use crate::jit::{Cond, MemoryAmount, ShiftType};
@@ -194,6 +194,16 @@ impl BlockInst {
                 }
                 (inputs, outputs)
             }
+
+            BlockInst::Prologue => (
+                block_reg_set!(Some(BlockReg::Fixed(Reg::SP)), Some(BlockReg::Fixed(Reg::LR))),
+                block_reg_set!(Some(BlockReg::Fixed(Reg::SP))),
+            ),
+            BlockInst::Epilogue => (
+                block_reg_set!(Some(BlockReg::Fixed(Reg::SP))),
+                block_reg_set!(Some(BlockReg::Fixed(Reg::SP)), Some(BlockReg::Fixed(Reg::PC))),
+            ),
+
             BlockInst::Label { .. } | BlockInst::Branch { .. } | BlockInst::GuestPc(_) | BlockInst::Bkpt(_) => (block_reg_set!(), block_reg_set!()),
         }
     }
@@ -261,7 +271,7 @@ impl BlockInst {
                     Self::replace_reg(&mut regs_mapping[reg as usize], old, new);
                 }
             }
-            BlockInst::Label { .. } | BlockInst::Branch { .. } | BlockInst::GuestPc(_) | BlockInst::Bkpt(_) => {}
+            BlockInst::Label { .. } | BlockInst::Branch { .. } | BlockInst::GuestPc(_) | BlockInst::Bkpt(_) | BlockInst::Prologue | BlockInst::Epilogue => {}
         }
     }
 
@@ -315,7 +325,7 @@ impl BlockInst {
                     Self::replace_reg(&mut regs_mapping[reg as usize], old, new);
                 }
             }
-            BlockInst::Label { .. } | BlockInst::Branch { .. } | BlockInst::GuestPc(_) | BlockInst::Bkpt(_) => {}
+            BlockInst::Label { .. } | BlockInst::Branch { .. } | BlockInst::GuestPc(_) | BlockInst::Bkpt(_) | BlockInst::Prologue | BlockInst::Epilogue => {}
         }
     }
 
@@ -355,7 +365,7 @@ impl BlockInst {
                     Self::replace_reg(reg_mapping, old, new);
                 }
             }
-            BlockInst::Label { .. } | BlockInst::Branch { .. } | BlockInst::GuestPc(_) | BlockInst::Bkpt(_) => {}
+            BlockInst::Label { .. } | BlockInst::Branch { .. } | BlockInst::GuestPc(_) | BlockInst::Bkpt(_) | BlockInst::Prologue | BlockInst::Epilogue => {}
         }
     }
 
@@ -381,7 +391,7 @@ impl BlockInst {
         opcodes.push(LdrStrImm::str_offset_al(guest_reg, thread_regs_addr_reg, Reg::CPSR as u16 * 4));
     }
 
-    pub fn emit_opcode(&mut self, opcodes: &mut Vec<u32>, opcode_index: usize, branch_placeholders: &mut Vec<usize>, opcodes_offset: usize) {
+    pub fn emit_opcode(&mut self, opcodes: &mut Vec<u32>, opcode_index: usize, branch_placeholders: &mut Vec<usize>, opcodes_offset: usize, used_host_regs: RegReserve) {
         let alu_reg = |op: BlockAluOp, op0: BlockReg, op1: BlockReg, op2: BlockReg, shift: BlockShift, set_cond: bool| match shift.value {
             BlockOperand::Reg(shift_reg) => AluReg::generic(op as u8, op0.as_fixed(), op1.as_fixed(), op2.as_fixed(), shift.shift_type, shift_reg.as_fixed(), set_cond, Cond::AL),
             BlockOperand::Imm(shift_imm) => {
@@ -590,6 +600,11 @@ impl BlockInst {
                 opcodes.push(inst_info.assemble());
             }
 
+            // BlockInst::Prologue => opcodes.push(LdmStm::generic(Reg::SP, used_host_regs + Reg::LR, false, true, false, true, Cond::AL)),
+            // BlockInst::Epilogue => opcodes.push(LdmStm::generic(Reg::SP, used_host_regs + Reg::PC, true, true, true, false, Cond::AL)),
+            BlockInst::Prologue => opcodes.push(LdmStm::generic(Reg::SP, block_reg_allocator::ALLOCATION_REGS + Reg::LR, false, true, false, true, Cond::AL)),
+            BlockInst::Epilogue => opcodes.push(LdmStm::generic(Reg::SP, block_reg_allocator::ALLOCATION_REGS + Reg::PC, true, true, true, false, Cond::AL)),
+
             BlockInst::Label { .. } | BlockInst::GuestPc(_) => {}
         }
     }
@@ -713,6 +728,9 @@ pub enum BlockInst {
         inst: GuestInstInfo,
         regs_mapping: [BlockReg; Reg::None as usize],
     },
+
+    Prologue,
+    Epilogue,
 }
 
 impl Debug for BlockInst {
@@ -783,6 +801,8 @@ impl Debug for BlockInst {
             BlockInst::Bkpt(id) => write!(f, "Bkpt {id}"),
             BlockInst::GuestPc(pc) => write!(f, "GuestPc {pc:x}"),
             BlockInst::GenericGuestInst { inst, .. } => write!(f, "{inst:?}"),
+            BlockInst::Prologue => write!(f, "Prologue"),
+            BlockInst::Epilogue => write!(f, "Epilogue"),
         }
     }
 }