Skip to content

Commit

Permalink
Refactor - removes usage of REGISTER_OTHER_SCRATCH (#597)
Browse files Browse the repository at this point in the history
* Adds support for RSP in Value::RegisterIndirect arguments.

* Pass value to store as RSP based RegisterIndirect argument.

* Reorders value to store block above.

* Replace usage of REGISTER_OTHER_SCRATCH with REGISTER_SCRATCH and stack_slot_of_value_to_store.

* Same for the case in which address translation is disabled.
  • Loading branch information
Lichtso authored Sep 27, 2024
1 parent 9d1a9a0 commit 7364447
Showing 1 changed file with 29 additions and 20 deletions.
49 changes: 29 additions & 20 deletions src/jit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,7 @@ const REGISTER_PTR_TO_VM: u8 = ARGUMENT_REGISTERS[0];
/// RBX: Program counter limit
const REGISTER_INSTRUCTION_METER: u8 = CALLEE_SAVED_REGISTERS[1];
/// R10: Other scratch register
const REGISTER_OTHER_SCRATCH: u8 = CALLER_SAVED_REGISTERS[7];
// const REGISTER_OTHER_SCRATCH: u8 = CALLER_SAVED_REGISTERS[7];
/// R11: Scratch register
const REGISTER_SCRATCH: u8 = CALLER_SAVED_REGISTERS[8];

Expand Down Expand Up @@ -974,7 +974,10 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
Value::RegisterIndirect(reg, offset, user_provided) => {
debug_assert!(!user_provided);
if is_stack_argument {
debug_assert_ne!(reg, RSP);
self.emit_ins(X86Instruction::push(reg, Some(X86IndirectAccess::Offset(offset))));
} else if reg == RSP {
self.emit_ins(X86Instruction::load(OperandSize::S64, RSP, dst, X86IndirectAccess::OffsetIndexShift(offset, RSP, 0)));
} else {
self.emit_ins(X86Instruction::load(OperandSize::S64, reg, dst, X86IndirectAccess::Offset(offset)));
}
Expand All @@ -984,6 +987,8 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
if is_stack_argument {
self.emit_ins(X86Instruction::push(reg, None));
self.emit_ins(X86Instruction::alu(OperandSize::S64, 0x81, 0, RSP, offset as i64, Some(X86IndirectAccess::OffsetIndexShift(0, RSP, 0))));
} else if reg == RSP {
self.emit_ins(X86Instruction::lea(OperandSize::S64, RSP, dst, Some(X86IndirectAccess::OffsetIndexShift(offset, RSP, 0))));
} else {
self.emit_ins(X86Instruction::lea(OperandSize::S64, reg, dst, Some(X86IndirectAccess::Offset(offset))));
}
Expand Down Expand Up @@ -1078,6 +1083,22 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
fn emit_address_translation(&mut self, dst: Option<u8>, vm_addr: Value, len: u64, value: Option<Value>) {
debug_assert_ne!(dst.is_some(), value.is_some());

let stack_slot_of_value_to_store = X86IndirectAccess::OffsetIndexShift(-112, RSP, 0);
match value {
Some(Value::Register(reg)) => {
self.emit_ins(X86Instruction::store(OperandSize::S64, reg, RSP, stack_slot_of_value_to_store));
}
Some(Value::Constant64(constant, user_provided)) => {
if user_provided && self.should_sanitize_constant(constant) {
self.emit_sanitized_load_immediate(OperandSize::S64, REGISTER_SCRATCH, constant);
} else {
self.emit_ins(X86Instruction::load_immediate(OperandSize::S64, REGISTER_SCRATCH, constant));
}
self.emit_ins(X86Instruction::store(OperandSize::S64, REGISTER_SCRATCH, RSP, stack_slot_of_value_to_store));
}
_ => {}
}

match vm_addr {
Value::RegisterPlusConstant64(reg, constant, user_provided) => {
if user_provided && self.should_sanitize_constant(constant) {
Expand All @@ -1100,20 +1121,6 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
},
}

match value {
Some(Value::Register(reg)) => {
self.emit_ins(X86Instruction::mov(OperandSize::S64, reg, REGISTER_OTHER_SCRATCH));
}
Some(Value::Constant64(constant, user_provided)) => {
if user_provided && self.should_sanitize_constant(constant) {
self.emit_sanitized_load_immediate(OperandSize::S64, REGISTER_OTHER_SCRATCH, constant);
} else {
self.emit_ins(X86Instruction::load_immediate(OperandSize::S64, REGISTER_OTHER_SCRATCH, constant));
}
}
_ => {}
}

if self.config.enable_address_translation {
let access_type = if value.is_none() { AccessType::Load } else { AccessType::Store };
let anchor = ANCHOR_TRANSLATE_MEMORY_ADDRESS + len.trailing_zeros() as usize + 4 * (access_type as usize);
Expand All @@ -1131,13 +1138,15 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
_ => unreachable!(),
}
} else {
self.emit_ins(X86Instruction::xchg(OperandSize::S64, RSP, REGISTER_MAP[0], Some(stack_slot_of_value_to_store))); // Save REGISTER_MAP[0] and retrieve value to store
match len {
1 => self.emit_ins(X86Instruction::store(OperandSize::S8, REGISTER_OTHER_SCRATCH, REGISTER_SCRATCH, X86IndirectAccess::Offset(0))),
2 => self.emit_ins(X86Instruction::store(OperandSize::S16, REGISTER_OTHER_SCRATCH, REGISTER_SCRATCH, X86IndirectAccess::Offset(0))),
4 => self.emit_ins(X86Instruction::store(OperandSize::S32, REGISTER_OTHER_SCRATCH, REGISTER_SCRATCH, X86IndirectAccess::Offset(0))),
8 => self.emit_ins(X86Instruction::store(OperandSize::S64, REGISTER_OTHER_SCRATCH, REGISTER_SCRATCH, X86IndirectAccess::Offset(0))),
1 => self.emit_ins(X86Instruction::store(OperandSize::S8, REGISTER_MAP[0], REGISTER_SCRATCH, X86IndirectAccess::Offset(0))),
2 => self.emit_ins(X86Instruction::store(OperandSize::S16, REGISTER_MAP[0], REGISTER_SCRATCH, X86IndirectAccess::Offset(0))),
4 => self.emit_ins(X86Instruction::store(OperandSize::S32, REGISTER_MAP[0], REGISTER_SCRATCH, X86IndirectAccess::Offset(0))),
8 => self.emit_ins(X86Instruction::store(OperandSize::S64, REGISTER_MAP[0], REGISTER_SCRATCH, X86IndirectAccess::Offset(0))),
_ => unreachable!(),
}
self.emit_ins(X86Instruction::xchg(OperandSize::S64, RSP, REGISTER_MAP[0], Some(stack_slot_of_value_to_store))); // Restore REGISTER_MAP[0]
}
}

Expand Down Expand Up @@ -1549,7 +1558,7 @@ impl<'a, C: ContextObject> JitCompiler<'a, C> {
};
self.emit_rust_call(Value::Constant64(store, false), &[
Argument { index: 3, value: Value::Register(REGISTER_SCRATCH) }, // Specify first as the src register could be overwritten by other arguments
Argument { index: 2, value: Value::Register(REGISTER_OTHER_SCRATCH) },
Argument { index: 2, value: Value::RegisterIndirect(RSP, -8, false) },
Argument { index: 4, value: Value::Constant64(0, false) }, // self.pc is set later
Argument { index: 1, value: Value::RegisterPlusConstant32(REGISTER_PTR_TO_VM, self.slot_in_vm(RuntimeEnvironmentSlot::MemoryMapping), false) },
Argument { index: 0, value: Value::RegisterPlusConstant32(REGISTER_PTR_TO_VM, self.slot_in_vm(RuntimeEnvironmentSlot::ProgramResult), false) },
Expand Down

0 comments on commit 7364447

Please sign in to comment.