I’ve been working on a relatively basic RISC-V interpreter and tried to optimize its performance. Once I learned about SROA existence, I decided to refactor register file from a flat array to a data structure that is SROA-friendly: a few hot registers as separate fields, the rest are in a cold array:
pub struct ContractRegisters {
ra: u64,
sp: u64,
s0: u64,
a0: u64,
a1: u64,
a2: u64,
a3: u64,
cold_registers: [u64; 22],
}
impl ContractRegisters {
#[inline(always)]
pub fn stack_registers(&mut self) -> StackRegisters<'_> {
StackRegisters {
ra: self.ra,
sp: self.sp,
s0: self.s0,
a0: self.a0,
a1: self.a1,
a2: self.a2,
a3: self.a3,
original: self,
}
}
}
pub struct StackRegisters<'a> {
ra: u64,
sp: u64,
s0: u64,
a0: u64,
a1: u64,
a2: u64,
a3: u64,
original: &'a mut ContractRegisters,
}
impl Drop for StackRegisters<'_> {
#[inline(always)]
fn drop(&mut self) {
self.original.ra = self.ra;
self.original.sp = self.sp;
self.original.s0 = self.s0;
self.original.a0 = self.a0;
self.original.a1 = self.a1;
self.original.a2 = self.a2;
self.original.a3 = self.a3;
}
}
impl StackRegisters<'_> {
#[inline(always)]
fn read(&self, reg: ContractRegister) -> u64 {
match reg {
ContractRegister::Zero => {
// Always zero
0
}
ContractRegister::Ra => self.ra,
ContractRegister::Sp => self.sp,
ContractRegister::S0 => self.s0,
ContractRegister::A0 => self.a0,
ContractRegister::A1 => self.a1,
ContractRegister::A2 => self.a2,
ContractRegister::A3 => self.a3,
reg => {
cold_path();
// SAFETY: register offset is always within bounds
*unsafe {
self.original
.cold_registers
.get_unchecked(usize::from(reg as u8))
}
}
}
}
#[inline(always)]
fn write(&mut self, reg: ContractRegister, value: u64) {
match reg {
ContractRegister::Zero => {
// Writes are ignored
}
ContractRegister::Ra => {
self.ra = value;
}
ContractRegister::Sp => {
self.sp = value;
}
ContractRegister::S0 => {
self.s0 = value;
}
ContractRegister::A0 => {
self.a0 = value;
}
ContractRegister::A1 => {
self.a1 = value;
}
ContractRegister::A2 => {
self.a2 = value;
}
ContractRegister::A3 => {
self.a3 = value;
}
reg => {
cold_path();
// SAFETY: register offset is always within bounds
*unsafe {
self.original
.cold_registers
.get_unchecked_mut(usize::from(reg as u8))
} = value;
}
}
}
}
#[derive(Clone, Copy)]
#[repr(u8)]
pub enum ContractRegister {
Zero = 255,
Ra = 127,
Sp = 128,
T0 = 0,
T1 = 1,
T2 = 2,
S0 = 129,
S1 = 3,
A0 = 130,
A1 = 131,
A2 = 132,
A3 = 133,
A4 = 20,
A5 = 21,
A6 = 4,
A7 = 5,
S2 = 6,
S3 = 7,
S4 = 8,
S5 = 9,
S6 = 10,
S7 = 11,
S8 = 12,
S9 = 13,
S10 = 14,
S11 = 15,
T3 = 16,
T4 = 17,
T5 = 18,
T6 = 19,
}
I then create a stack variable before executing instructions for SROA promotion and copy data back after the loop:
fn execute_experiment(
state: &mut BasicInterpreterState,
) -> Result<(), ExecutionError> {
// Use local variable for the register file so that SROA (Scalar Replacement of Aggregates)
// pass in LLVM can promote hot registers to native registers
let mut regs = state.regs.stack_registers();
loop {
let instruction = match state.instruction_fetcher.fetch_instruction(&state.memory)? {
FetchInstructionResult::Instruction(instruction) => instruction,
FetchInstructionResult::ControlFlow(ControlFlow::Continue(())) => {
continue;
}
FetchInstructionResult::ControlFlow(ControlFlow::Break(())) => {
break;
}
};
match instruction.execute(
&mut regs,
&mut state.ext_state,
&mut state.memory,
&mut state.instruction_fetcher,
&mut state.system_instruction_handler,
)? {
ControlFlow::Continue(()) => {
continue;
}
ControlFlow::Break(()) => {
break;
}
}
}
Ok(())
}
Execution method looks something like this:
impl ContractInstruction {
#[inline(always)]
pub fn execute(
self,
regs: &mut StackRegisters<'_>,
_ext_state: &mut (),
memory: &mut TestMemory,
program_counter: &mut EagerTestInstructionFetcher,
system_instruction_handler: &mut NoopRv64SystemInstructionHandler,
) -> Result<ControlFlow<()>, ExecutionError> {
match self {
Self::Add { rd, rs1, rs2 } => {
let value = regs.read(rs1).wrapping_add(regs.read(rs2));
regs.write(rd, value);
Ok(ControlFlow::Continue(()))
}
Self::Sub { rd, rs1, rs2 } => {
let value = regs.read(rs1).wrapping_sub(regs.read(rs2));
regs.write(rd, value);
Ok(ControlFlow::Continue(()))
}
// ...
}
}
}
According to LLVM IR promotion to SROA worked, but due to a large switch on the instruction enum, each RISC-V register resulted in phi node in most of the branches.
The result is ~2x lower performance than memory read/write despite expectation of the opposite (load-to-store conflicts surfaced during profiling and I wanted to have hot RISC-V registers in native registers).
Remarks uncovered this:
// 121 spills 1.940628e+01 total spills cost 218 reloads 3.220546e+01 total reloads cost 446 virtual registers copies 2.721900e+02 total copies cost generated in function
fn execute_experiment(
state: &mut BasicInterpreterState,
) -> Result<(), ExecutionError> {
//
// 117 spills 1.540628e+01 total spills cost 213 reloads 2.820546e+01 total reloads cost 446 virtual registers copies 2.721900e+02 total copies cost generated in loop
loop {
//
}
Ok(())
}
And this is for interpreter of the base RISC-V ISA with no extensions. With extensions it is a lot worse.
Here is a reduced example on compiler explorer: Compiler Explorer
This is a common pattern for things like interpreters and I’m far from the first one to hit this (Building the fastest Lua interpreter.. automatically! |).
And with a simple interpreter in mostly safe Rust I can’t exactly do calls with custom calling conventions, especially in cross-platform way.
Is this is something that can be optimized better in LLVM? Anything I can do in the meantime?
It seems like this should be possible, but my understanding of all this is very superficial.