From 18c0b8a8461867d6d04477dc7dce2e25775ef5e9 Mon Sep 17 00:00:00 2001 From: Emma Turner Date: Wed, 30 Apr 2025 14:57:51 +0100 Subject: [PATCH 1/2] RISC-V: add dispatch target wrapping the block_run pointer --- .../src/machine_state/block_cache/block.rs | 32 +++++++------ .../block_cache/block/dispatch.rs | 45 +++++++++++++++++++ 2 files changed, 63 insertions(+), 14 deletions(-) create mode 100644 src/riscv/lib/src/machine_state/block_cache/block/dispatch.rs diff --git a/src/riscv/lib/src/machine_state/block_cache/block.rs b/src/riscv/lib/src/machine_state/block_cache/block.rs index f9ba819e74c4..943f58a069df 100644 --- a/src/riscv/lib/src/machine_state/block_cache/block.rs +++ b/src/riscv/lib/src/machine_state/block_cache/block.rs @@ -4,8 +4,10 @@ // SPDX-License-Identifier: MIT //! Switching of execution strategy for blocks. -//! -//! Currently just for interperation only, but will expand to cover JIT. + +mod dispatch; + +use dispatch::DispatchTarget; use super::CACHE_INSTR; use super::ICallPlaced; @@ -248,7 +250,7 @@ impl Clone for Interpreted { /// /// Internally, this may be interpreted, just-in-time compiled, or do /// additional work over just execution. -pub type Dispatch = unsafe extern "C" fn( +pub type DispatchFn = unsafe extern "C" fn( &mut InlineJit, &mut MachineCoreState, Address, @@ -265,7 +267,7 @@ pub type Dispatch = unsafe extern "C" fn( /// Blocks are compiled upon calling [`Block::run_block`], in a *stop the world* fashion. pub struct InlineJit { fallback: Interpreted, - dispatch: Dispatch, + dispatch: DispatchTarget, } impl InlineJit { @@ -303,12 +305,12 @@ impl InlineJit { // last parameters. These are both pointers, and ignored by the JitFn. // // It's therefore safe to cast these to thin-pointers to any type. - unsafe { std::mem::transmute::, Dispatch>(jitfn) } + unsafe { std::mem::transmute::, DispatchFn>(jitfn) } } None => Self::run_block_not_compiled, }; - self.dispatch = fun; + self.dispatch.set(fun); // Safety: the block builder passed to this function is always the same for the // lifetime of the block @@ -346,7 +348,7 @@ impl NewState for InlineJit { { Self { fallback: Interpreted::new(manager), - dispatch: Self::run_block_interpreted, + dispatch: DispatchTarget::default(), } } } @@ -358,7 +360,7 @@ impl Block for InlineJit { where M: ManagerWrite, { - self.dispatch = Self::run_block_interpreted; + self.dispatch.reset(); self.fallback.start_block() } @@ -366,7 +368,7 @@ impl Block for InlineJit { where M: ManagerWrite, { - self.dispatch = Self::run_block_interpreted; + self.dispatch.reset(); self.fallback.invalidate() } @@ -374,7 +376,7 @@ impl Block for InlineJit { where M: ManagerReadWrite, { - self.dispatch = Self::run_block_interpreted; + self.dispatch.reset(); self.fallback.reset() } @@ -382,7 +384,7 @@ impl Block for InlineJit { where M: ManagerReadWrite, { - self.dispatch = Self::run_block_interpreted; + self.dispatch.reset(); self.fallback.push_instr(instr) } @@ -396,7 +398,7 @@ impl Block for InlineJit { fn bind(allocated: AllocatedOf) -> Self { Self { fallback: Interpreted::bind(allocated), - dispatch: Self::run_block_interpreted, + dispatch: DispatchTarget::default(), } } @@ -424,9 +426,11 @@ impl Block for InlineJit { { let mut result = Ok(()); + let fun = self.dispatch.get(); + // Safety: the block builder is always the same instance, guarantee-ing that any // jit-compiled function is still alive. - unsafe { (self.dispatch)(self, core, instr_pc, steps, &mut result, block_builder) }; + unsafe { (fun)(self, core, instr_pc, steps, &mut result, block_builder) }; result } @@ -443,7 +447,7 @@ impl Clone for InlineJit Self { Self { fallback: self.fallback.clone(), - dispatch: Self::run_block_interpreted, + dispatch: DispatchTarget::default(), } } } diff --git a/src/riscv/lib/src/machine_state/block_cache/block/dispatch.rs b/src/riscv/lib/src/machine_state/block_cache/block/dispatch.rs new file mode 100644 index 000000000000..96461531c682 --- /dev/null +++ b/src/riscv/lib/src/machine_state/block_cache/block/dispatch.rs @@ -0,0 +1,45 @@ +// SPDX-FileCopyrightText: 2025 TriliTech + +//! Dispatching of blocks under JIT is done via hot-swappable +//! function pointers. +//! +//! This module exposes wrappers for the style of dispatch and compilation that is done. +//! +//! Currently, this is only 'inline' jit, but will soon be expanded to 'outline' jit also; +//! where 'outline' means any JIT compilation occurs in a separate thread. + +use super::InlineJit; +use crate::jit::state_access::JitStateAccess; +use crate::machine_state::memory::MemoryConfig; + +/// Dispatch target that wraps a [`DispatchFn`]. +/// +/// This is the target used for compilation - see [`DispatchCompiler::compile`]. +pub struct DispatchTarget { + fun: std::cell::Cell>, +} + +impl DispatchTarget { + /// Reset the dispatch target to the interpreted dispatch mechanism. + pub fn reset(&self) { + self.set(InlineJit::run_block_interpreted); + } + + /// Set the dispatch target to use the given `block_run` function. + pub fn set(&self, fun: super::DispatchFn) { + self.fun.set(fun); + } + + /// Get the dispatch target's current `block_run` function. + pub fn get(&self) -> super::DispatchFn { + self.fun.get() + } +} + +impl Default for DispatchTarget { + fn default() -> Self { + Self { + fun: std::cell::Cell::new(InlineJit::run_block_interpreted), + } + } +} -- GitLab From 606393b73977c970fe0d8f9039e8ee2aa8bab79b Mon Sep 17 00:00:00 2001 From: Emma Turner Date: Wed, 30 Apr 2025 15:08:40 +0100 Subject: [PATCH 2/2] RISC-V: move block compilation behind trait --- src/riscv/lib/src/jit.rs | 4 +- .../src/machine_state/block_cache/block.rs | 42 ++++++------ .../block_cache/block/dispatch.rs | 64 ++++++++++++++++--- src/riscv/lib/src/stepper/test/interpreter.rs | 5 +- src/riscv/lib/tests/test_regression.rs | 4 +- src/riscv/sandbox/src/commands/run.rs | 2 +- 6 files changed, 81 insertions(+), 40 deletions(-) diff --git a/src/riscv/lib/src/jit.rs b/src/riscv/lib/src/jit.rs index 30d9b40ac7c6..e6c157cbe14f 100644 --- a/src/riscv/lib/src/jit.rs +++ b/src/riscv/lib/src/jit.rs @@ -37,7 +37,7 @@ use crate::traps::EnvironException; /// Alias for the function signature produced by the JIT compilation. /// -/// This must have the same Abi as [`Dispatch`], which is used by +/// This must have the same Abi as [`DispatchFn`], which is used by /// the block dispatch mechanism in the block cache. /// /// The JitFn does not inspect the first and last parameters here, however. @@ -47,7 +47,7 @@ use crate::traps::EnvironException; /// with pointers to `c_void` - which in the C abi map to the same parameter type as the /// thin-references to the actual variables passed. /// -/// [`Dispatch`]: crate::machine_state::block_cache::block::Dispatch +/// [`DispatchFn`]: crate::machine_state::block_cache::block::DispatchFn pub type JitFn = unsafe extern "C" fn( // ignored *const c_void, diff --git a/src/riscv/lib/src/machine_state/block_cache/block.rs b/src/riscv/lib/src/machine_state/block_cache/block.rs index 943f58a069df..975efdcd06cb 100644 --- a/src/riscv/lib/src/machine_state/block_cache/block.rs +++ b/src/riscv/lib/src/machine_state/block_cache/block.rs @@ -7,14 +7,13 @@ mod dispatch; +use dispatch::DispatchCompiler; use dispatch::DispatchTarget; use super::CACHE_INSTR; use super::ICallPlaced; use super::run_instr; use crate::default::ConstDefault; -use crate::jit::JIT; -use crate::jit::JitFn; use crate::jit::state_access::JitStateAccess; use crate::machine_state::MachineCoreState; use crate::machine_state::ProgramCounterUpdate; @@ -250,13 +249,13 @@ impl Clone for Interpreted { /// /// Internally, this may be interpreted, just-in-time compiled, or do /// additional work over just execution. -pub type DispatchFn = unsafe extern "C" fn( - &mut InlineJit, +pub type DispatchFn = unsafe extern "C" fn( + &mut Jitted, &mut MachineCoreState, Address, &mut usize, &mut Result<(), EnvironException>, - &mut as Block>::BlockBuilder, + &mut as Block>::BlockBuilder, ); /// Blocks that are compiled to native code for execution, when possible. @@ -265,12 +264,12 @@ pub type DispatchFn = unsafe extern "C" fn( /// unsupported instructions, a fallback to [`Interpreted`] mode occurs. /// /// Blocks are compiled upon calling [`Block::run_block`], in a *stop the world* fashion. -pub struct InlineJit { +pub struct Jitted, MC: MemoryConfig, M: JitStateAccess> { fallback: Interpreted, - dispatch: DispatchTarget, + dispatch: DispatchTarget, } -impl InlineJit { +impl, MC: MemoryConfig, M: JitStateAccess> Jitted { /// The default initial dispatcher for inline jit. /// /// This will run the block in interpreted mode by default, but will attempt to JIT-compile @@ -299,18 +298,7 @@ impl InlineJit { .map(|i| i.read_stored()) .collect::>(); - let fun = match block_builder.0.compile(&instr) { - Some(jitfn) => { - // Safety: the two function signatures are identical, apart from the first and - // last parameters. These are both pointers, and ignored by the JitFn. - // - // It's therefore safe to cast these to thin-pointers to any type. - unsafe { std::mem::transmute::, DispatchFn>(jitfn) } - } - None => Self::run_block_not_compiled, - }; - - self.dispatch.set(fun); + let fun = block_builder.0.compile(&mut self.dispatch, instr); // Safety: the block builder passed to this function is always the same for the // lifetime of the block @@ -341,7 +329,9 @@ impl InlineJit { } } -impl NewState for InlineJit { +impl, MC: MemoryConfig, M: JitStateAccess> NewState + for Jitted +{ fn new(manager: &mut M) -> Self where M: ManagerAlloc, @@ -353,8 +343,10 @@ impl NewState for InlineJit { } } -impl Block for InlineJit { - type BlockBuilder = (JIT, InterpretedBlockBuilder); +impl, MC: MemoryConfig, M: JitStateAccess> Block + for Jitted +{ + type BlockBuilder = (D, InterpretedBlockBuilder); fn start_block(&mut self) where @@ -443,7 +435,9 @@ impl Block for InlineJit { } } -impl Clone for InlineJit { +impl, MC: MemoryConfig, M: JitStateAccess + ManagerClone> Clone + for Jitted +{ fn clone(&self) -> Self { Self { fallback: self.fallback.clone(), diff --git a/src/riscv/lib/src/machine_state/block_cache/block/dispatch.rs b/src/riscv/lib/src/machine_state/block_cache/block/dispatch.rs index 96461531c682..fdb8e3bf0eb8 100644 --- a/src/riscv/lib/src/machine_state/block_cache/block/dispatch.rs +++ b/src/riscv/lib/src/machine_state/block_cache/block/dispatch.rs @@ -8,38 +8,84 @@ //! Currently, this is only 'inline' jit, but will soon be expanded to 'outline' jit also; //! where 'outline' means any JIT compilation occurs in a separate thread. -use super::InlineJit; +use super::DispatchFn; +use super::Jitted; +use crate::jit::JIT; +use crate::jit::JitFn; use crate::jit::state_access::JitStateAccess; +use crate::machine_state::instruction::Instruction; use crate::machine_state::memory::MemoryConfig; /// Dispatch target that wraps a [`DispatchFn`]. /// /// This is the target used for compilation - see [`DispatchCompiler::compile`]. -pub struct DispatchTarget { - fun: std::cell::Cell>, +pub struct DispatchTarget, MC: MemoryConfig, M: JitStateAccess> { + fun: std::cell::Cell>, } -impl DispatchTarget { +impl, MC: MemoryConfig, M: JitStateAccess> DispatchTarget { /// Reset the dispatch target to the interpreted dispatch mechanism. pub fn reset(&self) { - self.set(InlineJit::run_block_interpreted); + self.set(Jitted::run_block_interpreted); } /// Set the dispatch target to use the given `block_run` function. - pub fn set(&self, fun: super::DispatchFn) { + pub fn set(&self, fun: DispatchFn) { self.fun.set(fun); } /// Get the dispatch target's current `block_run` function. - pub fn get(&self) -> super::DispatchFn { + pub fn get(&self) -> DispatchFn { self.fun.get() } } -impl Default for DispatchTarget { +impl, MC: MemoryConfig, M: JitStateAccess> Default + for DispatchTarget +{ fn default() -> Self { Self { - fun: std::cell::Cell::new(InlineJit::run_block_interpreted), + fun: std::cell::Cell::new(Jitted::run_block_interpreted), } } } + +/// A compiler that can JIT-compile blocks of instructions, and hot-swap the execution of +/// said block in the given dispatch target. +pub trait DispatchCompiler: Default + Sized { + /// Compile a block, hot-swapping the `run_block` function contained in `target` in + /// the process. This could be to an interpreted execution method, and/or jit-compiled + /// function. + /// + /// NB - the hot-swapping of JIT-compiled blocks may occur at any time, and is not + /// guaranteed to be contained within the call-time of this function. (This is true for + /// outline jit, especially). + fn compile( + &mut self, + target: &mut DispatchTarget, + instr: Vec, + ) -> DispatchFn; +} + +impl DispatchCompiler for JIT { + fn compile( + &mut self, + target: &mut DispatchTarget, + instr: Vec, + ) -> DispatchFn { + let fun = match self.compile(&instr) { + Some(jitfn) => { + // Safety: the two function signatures are identical, apart from the first and + // last parameters. These are both pointers, and ignored by the JitFn. + // + // It's therefore safe to cast these to thin-pointers to any type. + unsafe { std::mem::transmute::, DispatchFn>(jitfn) } + } + None => Jitted::run_block_not_compiled, + }; + + target.set(fun); + + fun + } +} diff --git a/src/riscv/lib/src/stepper/test/interpreter.rs b/src/riscv/lib/src/stepper/test/interpreter.rs index a23cebe31417..bc50e7c6c6cc 100644 --- a/src/riscv/lib/src/stepper/test/interpreter.rs +++ b/src/riscv/lib/src/stepper/test/interpreter.rs @@ -10,9 +10,10 @@ use std::ops::Bound; use goldenfile::Mint; use paste::paste; +use crate::jit::JIT; use crate::machine_state::block_cache::block::Block; -use crate::machine_state::block_cache::block::InlineJit; use crate::machine_state::block_cache::block::Interpreted; +use crate::machine_state::block_cache::block::Jitted; use crate::machine_state::memory::M1M; use crate::machine_state::mode::Mode; use crate::machine_state::registers::XRegister; @@ -94,7 +95,7 @@ fn interpret_test_with_check(path: &str, exit_mode: Mode, check_xregs: &[(XRegis /// For the JIT, we run it twice - the first run to build up the blocks, and the /// second to run with these blocks already compiled (so that we actually use them). fn inline_jit_test_with_check(path: &str, exit_mode: Mode, check_xregs: &[(XRegister, u64)]) { - type BlockImpl = InlineJit; + type BlockImpl = Jitted, M1M, Owned>; let block_builder = Default::default(); let block_builder = diff --git a/src/riscv/lib/tests/test_regression.rs b/src/riscv/lib/tests/test_regression.rs index fa674d55475b..cd79126f4ce7 100644 --- a/src/riscv/lib/tests/test_regression.rs +++ b/src/riscv/lib/tests/test_regression.rs @@ -11,9 +11,9 @@ use std::path::PathBuf; use octez_riscv::jit::JIT; use octez_riscv::machine_state::DefaultCacheLayouts; use octez_riscv::machine_state::block_cache::block::Block; -use octez_riscv::machine_state::block_cache::block::InlineJit; use octez_riscv::machine_state::block_cache::block::Interpreted; use octez_riscv::machine_state::block_cache::block::InterpretedBlockBuilder; +use octez_riscv::machine_state::block_cache::block::Jitted; use octez_riscv::machine_state::memory::M64M; use octez_riscv::pvm::PvmHooks; use octez_riscv::state_backend::owned_backend::Owned; @@ -77,7 +77,7 @@ fn test_regression( // This needs to run *after* the previous interpreted test. Otherwise, we run into trouble when // checking and updating the golden files. - test_regression_for_block::>( + test_regression_for_block::>( (JIT::::new().unwrap(), InterpretedBlockBuilder), golden_dir, kernel_path, diff --git a/src/riscv/sandbox/src/commands/run.rs b/src/riscv/sandbox/src/commands/run.rs index acb0a075adbc..00d87f39f716 100644 --- a/src/riscv/sandbox/src/commands/run.rs +++ b/src/riscv/sandbox/src/commands/run.rs @@ -31,7 +31,7 @@ type BlockImplInner = block::Interpreted; /// Inner execution strategy for blocks. #[cfg(feature = "inline-jit")] -type BlockImplInner = block::InlineJit; +type BlockImplInner = block::Jitted, M1G, Owned>; /// Executor of blocks #[cfg(not(feature = "metrics"))] -- GitLab