diff --git a/src/riscv/lib/src/jit.rs b/src/riscv/lib/src/jit.rs index 30d9b40ac7c6a13e2c6e2d0f58dcdd8e844690d8..e6c157cbe14f8c43e29eb09cccfc17869eaf72c5 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 f9ba819e74c4ac7e5082f9353e1b8285af8d7f3f..975efdcd06cb9983fc758a855d8afe6f1d170609 100644 --- a/src/riscv/lib/src/machine_state/block_cache/block.rs +++ b/src/riscv/lib/src/machine_state/block_cache/block.rs @@ -4,15 +4,16 @@ // 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::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; @@ -248,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 Dispatch = 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. @@ -263,12 +264,12 @@ pub type Dispatch = 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: Dispatch, + 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 @@ -297,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::, Dispatch>(jitfn) } - } - None => Self::run_block_not_compiled, - }; - - self.dispatch = 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 @@ -339,26 +329,30 @@ impl InlineJit { } } -impl NewState for InlineJit { +impl, MC: MemoryConfig, M: JitStateAccess> NewState + for Jitted +{ fn new(manager: &mut M) -> Self where M: ManagerAlloc, { Self { fallback: Interpreted::new(manager), - dispatch: Self::run_block_interpreted, + dispatch: DispatchTarget::default(), } } } -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 M: ManagerWrite, { - self.dispatch = Self::run_block_interpreted; + self.dispatch.reset(); self.fallback.start_block() } @@ -366,7 +360,7 @@ impl Block for InlineJit { where M: ManagerWrite, { - self.dispatch = Self::run_block_interpreted; + self.dispatch.reset(); self.fallback.invalidate() } @@ -374,7 +368,7 @@ impl Block for InlineJit { where M: ManagerReadWrite, { - self.dispatch = Self::run_block_interpreted; + self.dispatch.reset(); self.fallback.reset() } @@ -382,7 +376,7 @@ impl Block for InlineJit { where M: ManagerReadWrite, { - self.dispatch = Self::run_block_interpreted; + self.dispatch.reset(); self.fallback.push_instr(instr) } @@ -396,7 +390,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 +418,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 } @@ -439,11 +435,13 @@ 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(), - 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 0000000000000000000000000000000000000000..fdb8e3bf0eb89d5e5539fd630fba2f5fea4ad7d1 --- /dev/null +++ b/src/riscv/lib/src/machine_state/block_cache/block/dispatch.rs @@ -0,0 +1,91 @@ +// 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::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, MC: MemoryConfig, M: JitStateAccess> { + fun: std::cell::Cell>, +} + +impl, MC: MemoryConfig, M: JitStateAccess> DispatchTarget { + /// Reset the dispatch target to the interpreted dispatch mechanism. + pub fn reset(&self) { + self.set(Jitted::run_block_interpreted); + } + + /// Set the dispatch target to use the given `block_run` function. + pub fn set(&self, fun: DispatchFn) { + self.fun.set(fun); + } + + /// Get the dispatch target's current `block_run` function. + pub fn get(&self) -> DispatchFn { + self.fun.get() + } +} + +impl, MC: MemoryConfig, M: JitStateAccess> Default + for DispatchTarget +{ + fn default() -> Self { + Self { + 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 a23cebe3141708f861b3b763f03f40486de9c5c9..bc50e7c6c6cc9485bbff9ee983e9db0bdaa82a46 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 fa674d55475bae35037d24617d38473fb86476b1..cd79126f4ce761c2066a8b110cba9159ea6de07d 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 acb0a075adbc92f6da47f3f07271324e29f9bf10..00d87f39f7168efce74d590b0c8d6bf57089c3eb 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"))]