diff --git a/.cspell.dict/rust-more.txt b/.cspell.dict/rust-more.txt index b53639c3b41..459f60d04ca 100644 --- a/.cspell.dict/rust-more.txt +++ b/.cspell.dict/rust-more.txt @@ -20,6 +20,7 @@ cstring datelike deserializer deserializers +fcmp fdiv flamescope flate2 @@ -28,6 +29,8 @@ getres hasher hexf hexversion +iabs +iconst idents illumos ilog @@ -79,6 +82,7 @@ thiserror timelike timsort trai +uextend ulonglong unic unistd diff --git a/.cspell.json b/.cspell.json index 4ac238be18b..6ddd0fc1165 100644 --- a/.cspell.json +++ b/.cspell.json @@ -82,6 +82,7 @@ "deduped", "deoptimized", "deoptimize", + "deoptimizes", "emscripten", "excs", "fdigits", diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 7c56fc62076..e548239cb97 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -147,6 +147,26 @@ jobs: run: cargo build --locked --no-default-features --features ssl-openssl-vendor if: runner.os == 'Linux' + # Automatic compilation must not change what a program does, so the + # snippet has to pass identically with it on and off. + - name: Test aot build + run: | + cargo build --locked --features aot + target/debug/rustpython -X aot=1 extra_tests/snippets/aot.py + target/debug/rustpython -X aot=0 extra_tests/snippets/aot.py + target/debug/rustpython -X aot=1 extra_tests/snippets/jit.py + if: runner.os == 'Linux' + + # The snippets cover arithmetic. These cover the rest of what the + # automatic call path moves through: frames, tracebacks, tracing, and + # the threads that have to be able to leave a compiled loop. + - name: Test aot against the CPython suite + run: | + target/debug/rustpython -X aot=1 -m test \ + test_sys test_traceback test_sys_settrace test_monitoring \ + test_bdb test_trace test_exceptions test_generators test_threading + if: runner.os == 'Linux' + # - name: Install tk-dev for tkinter build # run: sudo apt-get update && sudo apt-get install -y tk-dev # if: runner.os == 'Linux' diff --git a/Cargo.lock b/Cargo.lock index f327b3ba419..ce13b6d7af0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3457,7 +3457,6 @@ dependencies = [ "cranelift", "cranelift-jit", "cranelift-module", - "libffi", "num-traits", "rustpython-compiler-core", "rustpython-derive", diff --git a/Cargo.toml b/Cargo.toml index f85f68b1d22..25cef6ad4b4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ stdlib = ["rustpython-stdlib", "rustpython-pylib", "encodings"] flame-it = ["rustpython-vm/flame-it", "rustpython-stdlib/flame-it", "flame", "flamescope"] freeze-stdlib = ["stdlib", "rustpython-vm/freeze-stdlib", "rustpython-pylib?/freeze-stdlib"] jit = ["rustpython-vm/jit"] +aot = ["rustpython-vm/aot"] threading = ["rustpython-vm/threading", "rustpython-stdlib/threading"] sqlite = ["rustpython-stdlib/sqlite"] ssl = ["host_env"] diff --git a/README.md b/README.md index 1c8c972edce..7af45d8991b 100644 --- a/README.md +++ b/README.md @@ -145,21 +145,70 @@ cargo build --release --target wasm32-wasip1 --features="freeze-stdlib" ### JIT (Just in time) compiler -RustPython has a **very** experimental JIT compiler that compile python functions into native code. +RustPython has a **very** experimental JIT compiler that compiles python functions into native code. +It comes in two forms: an automatic one that compiles a function once it has been called +often enough to be worth it, and an explicit `__jit__()` that compiles the one function +it is called on. #### Building -By default the JIT compiler isn't enabled, it's enabled with the `jit` cargo feature. +Neither is built by default. ```bash -cargo run --features jit +cargo run --features aot # automatic, and the explicit one with it +cargo run --features jit # explicit `__jit__()` alone ``` This requires autoconf, automake, libtool, and clang to be installed. -#### Using +#### Using the automatic compiler -To compile a function, call `__jit__()` on it. +A build that has it still has to be switched on, with `-X aot=1`, `RUSTPYTHON_AOT=1` +or `PYTHON_JIT=1`; `-X aot=0` and `=0` switch it back off. A function is then counted +as it is called and offered to the compiler once it has been called enough times to +repay one. The attempt happens once, so a function it turns down costs that one +attempt and is interpreted from then on; a function called only a handful of times +is never offered at all. + +`sys._jit.is_available()` reports whether the compiler was built in, `is_enabled()` +whether it is switched on, and `_stats()` returns `(compiled, rejected, deoptimized)` +for the functions it has looked at so far — a RustPython extension. + +#### What it compiles + +Scalar functions, and nothing else. The two forms differ in where the types come +from: `__jit__()` reads them off the annotations, and turns down a function without +them; the automatic path takes them from the arguments of the call that made the +function warm, so an unannotated function compiles as readily as an annotated one. +Nothing about that reads a `__annotations__` or runs an `__annotate__`. + +A guess about types is a guess: a later call whose arguments do not fit the compiled +signature is run by the interpreter instead. + +Taken: `int`, `float` and `bool` arguments, locals and return values; arithmetic, +comparison and boolean operators; `if`, `while`, and the assignments between them. + +Turned down: arguments of any other type, `*args`/`**kwargs`, closures, generators +and coroutines, `try`/`except`, attributes and methods, containers, `for`, calls to +anything but the function itself, and expressions that merge with an operand still +on the stack, such as a conditional expression. A call a function makes to itself +is compiled only by `__jit__()`; the automatic path turns those down too, because +the global it goes through can be rebound between one call and the next. + +Where a machine word runs out — an overflow, a division by zero, a shift past the +width, a power with no real answer — the compiled code hands the frame back at the +instruction it could not do, with the values it had, and the interpreter carries on +from there. The native code is dropped at that point, and the function is +interpreted afterwards. + +Compiled code runs with no python frame. That is why `sys._jit.is_active()` is +always `False`, why such a call reports no line and no return to `sys.settrace` or +`sys.monitoring`, and why a call is interpreted, and left uncompiled, while either +of those is installed. + +#### Using `__jit__()` + +To compile a single function, call `__jit__()` on it. This needs only the `jit` feature. ```python def foo(): diff --git a/crates/capi/src/pyframe.rs b/crates/capi/src/pyframe.rs index 611cf79b0b6..294f675a35c 100644 --- a/crates/capi/src/pyframe.rs +++ b/crates/capi/src/pyframe.rs @@ -14,8 +14,8 @@ pub unsafe extern "C" fn PyFrame_GetCode(frame: *mut PyFrameObject) -> *mut PyCo #[unsafe(no_mangle)] pub unsafe extern "C" fn PyFrame_GetLineNumber(frame: *mut PyFrameObject) -> c_int { - with_vm(|_vm| { - let lineno = unsafe { &*frame }.f_lineno(); + with_vm(|vm| { + let lineno = unsafe { &*frame }.f_lineno(vm); Ok(lineno.try_into().unwrap_or(c_int::MAX)) }) } diff --git a/crates/common/src/int.rs b/crates/common/src/int.rs index 06dc46c4098..ee7f170593c 100644 --- a/crates/common/src/int.rs +++ b/crates/common/src/int.rs @@ -5,6 +5,13 @@ use num_traits::{One, ToPrimitive, Zero}; #[must_use] pub fn true_div(numerator: &BigInt, denominator: &BigInt) -> f64 { + // A rational carries no signed zero, so `0 / -1` would round to `0.0`. A + // quotient of two differently signed operands is negative down to and + // including its zero, and only an exactly zero numerator loses that here: + // a quotient too small to represent still rounds to `-0.0` on its own. + if numerator.is_zero() && denominator.sign() == Sign::Minus { + return -0.0; + } let rational = Rational::from_integers_ref(numerator.into(), denominator.into()); match rational.rounding_into(RoundingMode::Nearest) { // returned value is $t::MAX but still less than the original diff --git a/crates/jit/Cargo.toml b/crates/jit/Cargo.toml index 5dcf0f4c31b..2471f587f23 100644 --- a/crates/jit/Cargo.toml +++ b/crates/jit/Cargo.toml @@ -15,7 +15,6 @@ rustpython-compiler-core = { workspace = true } num-traits = { workspace = true } thiserror = { workspace = true } -libffi = { workspace = true } cranelift = { workspace = true } cranelift-jit = { workspace = true } diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index a3c4ca800c4..dabf81c4cf8 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -1,5 +1,8 @@ // spell-checker: disable -use super::{JitCompileError, JitSig, JitType}; +use super::{ + DEOPT_HEADER_SLOTS, DEOPT_STATUS_NESTED, DeoptSite, JitCompileError, JitSig, JitType, + MAX_DEOPT_SLOTS, SLOT_SIZE, Safety, StackEntry, +}; use alloc::collections::BTreeSet; use cranelift::codegen::ir::FuncRef; use cranelift::prelude::*; @@ -10,19 +13,13 @@ use rustpython_compiler_core::bytecode::{ }; use std::collections::HashMap; -#[repr(u16)] -enum CustomTrapCode { - /// Raised when shifting by a negative number - NegativeShiftCount = 1, -} - #[derive(Clone)] struct Local { var: Variable, ty: JitType, } -#[derive(Debug)] +#[derive(Debug, Clone)] enum JitValue { Int(Value), Float(Value), @@ -59,20 +56,147 @@ impl JitValue { Self::None | Self::Null | Self::Tuple(_) | Self::FuncRef(_) => None, } } + + /// The cranelift value, without consuming the wrapper. + fn value(&self) -> Option { + match *self { + Self::Int(val) | Self::Float(val) | Self::Bool(val) => Some(val), + Self::None | Self::Null | Self::Tuple(_) | Self::FuncRef(_) => None, + } + } } -#[derive(Clone)] -struct DDValue { - hi: Value, - lo: Value, +/// What the module around a function being compiled hands it: the symbols it +/// may call, and the word its loops poll. +pub(crate) struct Externals { + /// `jit_powf`, imported into this function so `compile_fpow` can call it. + pub(crate) powf: FuncRef, + /// The byte a backward jump polls, or `None` to compile no poll at all. + pub(crate) safepoint: Option<&'static core::sync::atomic::AtomicU8>, } pub(crate) struct FunctionCompiler<'a, 'b> { builder: &'a mut FunctionBuilder<'b>, + /// The buffer a guard spills its record into, parameter 0 of the function. + deopt_ptr: Value, + /// The block every guard leaves through, created with the first one. + deopt_exit: Option, + /// What the module around this function supplies it with. + externals: Externals, + /// Bytecode offset the instruction being lowered would be re-entered at. + resume_offset: u32, stack: Vec, variables: Box<[Option]>, + /// One flag per varname slot, 1 once a local has been stored there. A local + /// declared inside a branch is not assigned on every path that reaches a + /// guard. They are all declared up front so that they can be zeroed in the + /// entry block, which is out of reach once compilation has started. + bound_flags: Box<[Variable]>, label_to_block: HashMap, + /// Whether any jump back to an earlier offset was lowered. Without one, + /// execution reaches a guard only by way of offsets below it, so every + /// store that can have run is already in `variables` and every site lists + /// all of them. + has_backward_jump: bool, + safety: Safety, pub(crate) sig: JitSig, + pub(crate) deopt_sites: Vec, +} + +fn jump_target_forward(offset: u32, caches: u32, arg: OpArg) -> Result { + let after = offset + .checked_add(1) + .and_then(|i| i.checked_add(caches)) + .ok_or(JitCompileError::BadBytecode)?; + let target = after + .checked_add(u32::from(arg)) + .ok_or(JitCompileError::BadBytecode)?; + Ok(Label::from_u32(target)) +} + +fn jump_target_backward(offset: u32, caches: u32, arg: OpArg) -> Result { + let after = offset + .checked_add(1) + .and_then(|i| i.checked_add(caches)) + .ok_or(JitCompileError::BadBytecode)?; + let target = after + .checked_sub(u32::from(arg)) + .ok_or(JitCompileError::BadBytecode)?; + Ok(Label::from_u32(target)) +} + +/// The offset a jump at `offset` lands on, or `None` when the instruction is +/// not a jump. Jump arguments are deltas, so this is the only way to learn +/// which offsets are branch targets - `CodeObject::label_targets` collects the +/// delta itself rather than the offset it points at. +pub(crate) fn instruction_target( + offset: u32, + instruction: Instruction, + arg: OpArg, +) -> Result, JitCompileError> { + let caches = instruction.cache_entries() as u32; + let target = match instruction { + Instruction::JumpForward { .. } => Some(jump_target_forward(offset, caches, arg)?), + Instruction::JumpBackward { .. } | Instruction::JumpBackwardNoInterrupt { .. } => { + Some(jump_target_backward(offset, caches, arg)?) + } + Instruction::PopJumpIfFalse { .. } + | Instruction::PopJumpIfTrue { .. } + | Instruction::PopJumpIfNone { .. } + | Instruction::PopJumpIfNotNone { .. } + | Instruction::ForIter { .. } + | Instruction::Send { .. } => Some(jump_target_forward(offset, caches, arg)?), + _ => None, + }; + Ok(target) +} + +/// Whether [`FunctionCompiler::add_instruction`] has a lowering for this opcode. +/// +/// This mirrors the match in that method so a caller can rule a code object out +/// before any compilation state is set up. It only has to be right in one +/// direction: claiming support for something the match rejects merely wastes a +/// compile attempt, and denying something it handles only costs an +/// optimization. Neither can produce wrong code. +pub(crate) const fn instruction_is_supported(instruction: Instruction) -> bool { + matches!( + instruction, + Instruction::BinaryOp { .. } + | Instruction::BuildTuple { .. } + | Instruction::Cache + | Instruction::Call { .. } + | Instruction::CallIntrinsic1 { .. } + | Instruction::CompareOp { .. } + | Instruction::CopyFreeVars { .. } + | Instruction::ExtendedArg + | Instruction::JumpBackward { .. } + | Instruction::JumpBackwardNoInterrupt { .. } + | Instruction::JumpForward { .. } + | Instruction::LoadConst { .. } + | Instruction::LoadFast { .. } + | Instruction::LoadFastBorrow { .. } + | Instruction::LoadFastBorrowLoadFastBorrow { .. } + | Instruction::LoadFastLoadFast { .. } + | Instruction::LoadGlobal { .. } + | Instruction::LoadSmallInt { .. } + | Instruction::MakeCell { .. } + | Instruction::Nop + | Instruction::NotTaken + | Instruction::PopJumpIfFalse { .. } + | Instruction::PopJumpIfTrue { .. } + | Instruction::PopTop + | Instruction::PushNull + | Instruction::Resume { .. } + | Instruction::ReturnValue + | Instruction::StoreFast { .. } + | Instruction::StoreFastLoadFast { .. } + | Instruction::StoreFastStoreFast { .. } + | Instruction::Swap { .. } + | Instruction::ToBool + | Instruction::UnaryNegative + | Instruction::UnaryNot + | Instruction::UnpackSequence { .. } + ) } impl<'a, 'b> FunctionCompiler<'a, 'b> { @@ -82,19 +206,40 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { arg_types: &[JitType], ret_type: Option, entry_block: Block, + safety: Safety, + externals: Externals, ) -> Self { + let params = builder.func.dfg.block_params(entry_block).to_vec(); + let (deopt_ptr, arg_params) = params.split_first().expect("the deopt buffer parameter"); + // The builder still sits in the entry block, which is the only place a + // flag can be given the value every path into the function starts from. + let bound_flags: Box<[Variable]> = (0..num_variables) + .map(|_| { + let flag = builder.declare_var(types::I8); + let unbound = builder.ins().iconst(types::I8, 0); + builder.def_var(flag, unbound); + flag + }) + .collect(); let mut compiler = Self { builder, + deopt_ptr: *deopt_ptr, + deopt_exit: None, + externals, + resume_offset: 0, stack: Vec::new(), variables: vec![None; num_variables].into_boxed_slice(), + bound_flags, label_to_block: HashMap::new(), + has_backward_jump: false, + safety, sig: JitSig { args: arg_types.to_vec(), ret: ret_type, }, + deopt_sites: Vec::new(), }; - let params = compiler.builder.func.dfg.block_params(entry_block).to_vec(); - for (i, (ty, val)) in arg_types.iter().zip(params).enumerate() { + for (i, (ty, val)) in arg_types.iter().zip(arg_params.iter().copied()).enumerate() { compiler .store_variable( (i as u32).into(), @@ -115,6 +260,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let builder = &mut self.builder; let ty = val.to_jit_type().ok_or(JitCompileError::NotSupported)?; let cranelift_ty = ty.to_cranelift().ok_or(JitCompileError::NotSupported)?; + let bound = self.bound_flags[idx]; let local = self.variables[idx].get_or_insert_with(|| { let var = builder.declare_var(cranelift_ty); Local { @@ -126,10 +272,173 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { Err(JitCompileError::NotSupported) } else { self.builder.def_var(local.var, val.into_value().unwrap()); + let is_bound = self.builder.ins().iconst(types::I8, 1); + self.builder.def_var(bound, is_bound); Ok(()) } } + /// Emit the two-way branch a guard is: when `cond` is non-zero the compiled + /// code gives up, spilling everything the interpreter needs to re-execute + /// this instruction from the start; otherwise it falls through and carries + /// on. + /// + /// `popped` is what this instruction has already taken off the stack, in + /// the order it must go back on. The interpreter re-executes the whole + /// instruction, so its operands have to be there. + fn deopt_branch(&mut self, cond: Value, popped: &[JitValue]) -> Result<(), JitCompileError> { + let live: Vec> = self + .variables + .iter() + .map(|local| local.as_ref().map(|l| l.ty.clone())) + .collect(); + // The callable and the null a call pushes beside it are the same on + // every path that reaches the guard, so they are described rather than + // written. Anything else without a slot encoding - a tuple, a bare + // `None` - is not statically known either, so a guard cannot be placed + // above one. + let mut entries = Vec::with_capacity(self.stack.len() + popped.len()); + let mut spilled = Vec::new(); + for val in self.stack.iter().chain(popped) { + match val { + JitValue::FuncRef(_) => entries.push(StackEntry::Callee), + JitValue::Null => entries.push(StackEntry::Null), + _ => { + let (ty, value) = val + .to_jit_type() + .zip(val.value()) + .ok_or(JitCompileError::NotSupported)?; + entries.push(StackEntry::Value(ty.clone())); + spilled.push((ty, value)); + } + } + } + + let slots = DEOPT_HEADER_SLOTS + live.iter().flatten().count() + spilled.len(); + if slots > MAX_DEOPT_SLOTS || live.len() > u64::BITS as usize { + return Err(JitCompileError::NotSupported); + } + + let site = self.deopt_sites.len(); + self.deopt_sites.push(DeoptSite { + offset: self.resume_offset, + locals: live.clone().into_boxed_slice(), + stack: entries.into_boxed_slice(), + // Provisional: only the whole function tells whether a store this + // site has not seen yet can run before its guard fires. + resumable: true, + safepoint: false, + }); + + self.deopt_if(cond, |this| { + let mut offset = DEOPT_HEADER_SLOTS; + let mut mask = this.builder.ins().iconst(types::I64, 0); + for (i, ty) in live.iter().enumerate() { + let Some(ty) = ty else { continue }; + let var = this.variables[i].as_ref().expect("live local").var; + let value = this.builder.use_var(var); + this.store_slot(ty, value, offset); + let bound = this.builder.use_var(this.bound_flags[i]); + let bound = this.builder.ins().uextend(types::I64, bound); + let bit = this.builder.ins().ishl_imm(bound, i as i64); + mask = this.builder.ins().bor(mask, bit); + offset += 1; + } + for (ty, value) in spilled { + this.store_slot(&ty, value, offset); + offset += 1; + } + this.store_raw(mask, 1); + let status = this.builder.ins().iconst(types::I64, site as i64 + 1); + this.store_raw(status, 0); + }); + Ok(()) + } + + /// Emit the poll a backward jump makes before it is taken: when the + /// interpreter has asked the running thread to leave the bytecode loop, + /// give up here, so the frame carries on interpreted from this same jump + /// and answers whatever the request was. + /// + /// A loop without one cannot be left. Compiled code runs no handler for a + /// pending signal, parks for no stop-the-world, and does not stop when the + /// interpreter shuts down, so a thread inside such a loop hangs the process + /// rather than being delayed by it. + fn compile_safepoint(&mut self) -> Result<(), JitCompileError> { + let Some(word) = self.externals.safepoint else { + return Ok(()); + }; + let ptr_type = self.builder.func.dfg.value_type(self.deopt_ptr); + let address = core::ptr::from_ref(word) as usize as i64; + let address = self.builder.ins().iconst(ptr_type, address); + // Deliberately not a `readonly` load: other threads and signal handlers + // write this byte while the loop runs, so what one iteration reads says + // nothing about what the next one will. + let pending = self + .builder + .ins() + .load(types::I8, MemFlags::trusted(), address, 0); + let tripped = self.builder.ins().icmp_imm(IntCC::NotEqual, pending, 0); + let site = self.deopt_sites.len(); + self.deopt_branch(tripped, &[])?; + // The site the branch above recorded. It spills the same record a + // guard's does; what differs is only that the code it left is still + // right for the values it was given. + self.deopt_sites[site].safepoint = true; + Ok(()) + } + + /// Emit "when `cond` is non-zero, leave through the shared deopt exit; + /// otherwise carry on". `spill` fills the block that is taken, which is + /// where a guard writes its record; a caller whose record is already + /// written leaves it empty. The builder is left in the fall-through block, + /// so lowering continues where it left off. + fn deopt_if(&mut self, cond: Value, spill: impl FnOnce(&mut Self)) { + let taken = self.builder.create_block(); + let carry_on = self.builder.create_block(); + self.builder.ins().brif(cond, taken, &[], carry_on, &[]); + + self.builder.switch_to_block(taken); + spill(self); + let exit = self.deopt_exit(); + self.builder.ins().jump(exit, &[]); + + self.builder.switch_to_block(carry_on); + } + + /// Write one value into the deopt buffer, in the 64-bit encoding + /// `AbiValue::from_slot` reads back. + fn store_slot(&mut self, ty: &JitType, value: Value, slot: usize) { + let value = match ty { + JitType::Int | JitType::Float => value, + // A flag occupies a whole slot, so the bits above it have to be zero. + JitType::Bool => self.builder.ins().uextend(types::I64, value), + // Neither a local nor a spilled stack entry can carry one: a local + // of no cranelift type is rejected when it is stored, and a `None` + // on the stack has no value to pair with its type. + JitType::None => return, + }; + self.store_raw(value, slot); + } + + fn store_raw(&mut self, value: Value, slot: usize) { + let offset = i32::try_from(slot * SLOT_SIZE).expect("slot count is capped"); + self.builder + .ins() + .store(MemFlags::trusted(), value, self.deopt_ptr, offset); + } + + /// The one block every guard leaves through. Its return is emitted at the + /// end of compilation, because until then the function's return type may + /// still be unknown. + fn deopt_exit(&mut self) -> Block { + #[expect(clippy::mut_mut, reason = "This seems like a false positive")] + let builder = &mut self.builder; + *self + .deopt_exit + .get_or_insert_with(|| builder.create_block()) + } + fn boolean_val(&mut self, val: JitValue) -> Result { match val { JitValue::Float(val) => { @@ -150,6 +459,23 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { } } + /// Require that the abstract stack is empty on an edge into a merge. + /// + /// The abstract stack is not reconciled where control flow merges: a merged + /// block keeps whichever predecessor's operands were lowered last, and its + /// entries are per-path SSA values, so even two predecessors of equal depth + /// describe different values. Only an empty stack is known to agree, which + /// is what every statement-level merge has - an `if`, an `if`/`else`, a + /// `while`. A conditional expression or a short-circuit operator merges + /// mid-expression, and is refused here rather than compiled wrongly. + fn require_empty_stack_at_merge(&self) -> Result<(), JitCompileError> { + if self.stack.is_empty() { + Ok(()) + } else { + Err(JitCompileError::NotSupported) + } + } + fn get_or_create_block(&mut self, label: Label) -> Block { #[expect(clippy::mut_mut, reason = "This seems like a false positive")] let builder = &mut self.builder; @@ -159,56 +485,6 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { .or_insert_with(|| builder.create_block()) } - fn jump_target_forward(offset: u32, caches: u32, arg: OpArg) -> Result { - let after = offset - .checked_add(1) - .and_then(|i| i.checked_add(caches)) - .ok_or(JitCompileError::BadBytecode)?; - let target = after - .checked_add(u32::from(arg)) - .ok_or(JitCompileError::BadBytecode)?; - Ok(Label::from_u32(target)) - } - - fn jump_target_backward( - offset: u32, - caches: u32, - arg: OpArg, - ) -> Result { - let after = offset - .checked_add(1) - .and_then(|i| i.checked_add(caches)) - .ok_or(JitCompileError::BadBytecode)?; - let target = after - .checked_sub(u32::from(arg)) - .ok_or(JitCompileError::BadBytecode)?; - Ok(Label::from_u32(target)) - } - - fn instruction_target( - offset: u32, - instruction: Instruction, - arg: OpArg, - ) -> Result, JitCompileError> { - let caches = instruction.cache_entries() as u32; - let target = match instruction { - Instruction::JumpForward { .. } => { - Some(Self::jump_target_forward(offset, caches, arg)?) - } - Instruction::JumpBackward { .. } | Instruction::JumpBackwardNoInterrupt { .. } => { - Some(Self::jump_target_backward(offset, caches, arg)?) - } - Instruction::PopJumpIfFalse { .. } - | Instruction::PopJumpIfTrue { .. } - | Instruction::PopJumpIfNone { .. } - | Instruction::PopJumpIfNotNone { .. } - | Instruction::ForIter { .. } - | Instruction::Send { .. } => Some(Self::jump_target_forward(offset, caches, arg)?), - _ => None, - }; - Ok(target) - } - pub(crate) fn compile( &mut self, func_ref: FuncRef, @@ -227,7 +503,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let mut target_arg_state = OpArgState::default(); for (offset, &raw_instr) in clean_instructions.iter().enumerate() { let (instruction, arg) = target_arg_state.get(raw_instr); - if let Some(target) = Self::instruction_target(offset as u32, instruction, arg)? { + if let Some(target) = instruction_target(offset as u32, instruction, arg)? { label_targets.insert(target); } } @@ -235,6 +511,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { // Track whether we have "returned" in the current block let mut in_unreachable_code = false; + let mut extended_start: Option = None; for (offset, &raw_instr) in clean_instructions.iter().enumerate() { let label = Label::from_u32(offset as u32); @@ -243,6 +520,13 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { // If this is a label that some earlier jump can target, // treat it as the start of a new reachable block: if label_targets.contains(&label) { + // Falling into a merge counts as an edge into it. An + // unreachable region has no edge, so its leftover operands are + // dropped rather than checked. + if !in_unreachable_code { + self.require_empty_stack_at_merge()?; + } + self.stack.clear(); // Create or get the block for this label: let target_block = self.get_or_create_block(label); @@ -276,9 +560,20 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { continue; } + // The oparg of an instruction preceded by EXTENDED_ARG is + // accumulated across the group, so a resume has to start where the + // group starts. + self.resume_offset = extended_start.unwrap_or(offset as u32); + // Actually compile this instruction: self.add_instruction(func_ref, bytecode, offset as u32, instruction, arg)?; + if matches!(instruction, Instruction::ExtendedArg) { + extended_start.get_or_insert(offset as u32); + } else { + extended_start = None; + } + // If that was an unconditional branch or return, mark future instructions unreachable match instruction { Instruction::ReturnValue @@ -301,6 +596,43 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { self.builder.ins().trap(TrapCode::user(0).unwrap()); } } + + // A site lists the locals the compiler had seen where its guard was + // lowered. A store to a varname first seen after it can still have run + // by the time that guard fires, but only by coming back down to it, so + // without a backward jump every site lists everything that can be + // bound. Where there is one, a site whose set of listed slots is short + // of the function's full set cannot describe the frame, and asks for a + // restart instead. + if self.has_backward_jump { + let bound: Vec = self.variables.iter().map(Option::is_some).collect(); + for site in &mut self.deopt_sites { + site.resumable = site + .locals + .iter() + .map(Option::is_some) + .eq(bound.iter().copied()); + } + } + + // The deopt exit returns whatever the function's signature ended up + // saying it returns; the value is never looked at, because the status + // says the call did not return one. + if let Some(exit) = self.deopt_exit { + self.builder.switch_to_block(exit); + match self.sig.ret.as_ref().and_then(JitType::to_cranelift) { + Some(ty) => { + let filler = match ty { + types::F64 => self.builder.ins().f64const(0.0), + ty => self.builder.ins().iconst(ty, 0), + }; + self.builder.ins().return_(&[filler]); + } + None => { + self.builder.ins().return_(&[]); + } + } + } Ok(()) } @@ -380,29 +712,56 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { JitValue::Int(b), ) => { let (out, carry) = self.builder.ins().sadd_overflow(a, b); - self.builder.ins().trapnz(carry, TrapCode::INTEGER_OVERFLOW); + // Overflow is not an error: it is where the interpreter + // stops using a machine word and starts using a bignum, + // so the operands go back to it intact. + self.deopt_branch(carry, &[JitValue::Int(a), JitValue::Int(b)])?; JitValue::Int(out) } ( BinaryOperator::Subtract | BinaryOperator::InplaceSubtract, JitValue::Int(a), JitValue::Int(b), - ) => JitValue::Int(self.compile_sub(a, b)), + ) => { + let out = self.compile_sub(a, b, &[JitValue::Int(a), JitValue::Int(b)])?; + JitValue::Int(out) + } ( BinaryOperator::FloorDivide | BinaryOperator::InplaceFloorDivide, JitValue::Int(a), JitValue::Int(b), - ) => JitValue::Int(self.builder.ins().sdiv(a, b)), + ) => JitValue::Int(self.compile_floor_div(a, b)?.0), ( BinaryOperator::TrueDivide | BinaryOperator::InplaceTrueDivide, JitValue::Int(a), JitValue::Int(b), ) => { - // Check if b == 0, If so trap with a division by zero error - self.builder - .ins() - .trapz(b, TrapCode::INTEGER_DIVISION_BY_ZERO); - // Else convert to float and divide + let operands = [JitValue::Int(a), JitValue::Int(b)]; + let by_zero = self.builder.ins().icmp_imm(IntCC::Equal, b, 0); + self.deopt_branch(by_zero, &operands)?; + + // `int.__truediv__` is correctly rounded. Converting both + // operands to double and dividing rounds twice, so it can be + // a ulp out as soon as either conversion is inexact - which + // is exactly when the operand does not fit in a double's + // significand. A magnitude of `1 << 53` is the largest that + // still converts exactly, so only what is past it deopts. + // `iabs` leaves `i64::MIN` negative, and the unsigned + // comparison reads that as `1 << 63`, which is past the + // bound - so the one value it cannot negate still deopts. + let too_wide = |compiler: &mut Self, v: Value| { + let magnitude = compiler.builder.ins().iabs(v); + compiler.builder.ins().icmp_imm( + IntCC::UnsignedGreaterThan, + magnitude, + 1 << 53, + ) + }; + let a_wide = too_wide(self, a); + let b_wide = too_wide(self, b); + let wide = self.builder.ins().bor(a_wide, b_wide); + self.deopt_branch(wide, &operands)?; + let a_float = self.builder.ins().fcvt_from_sint(types::F64, a); let b_float = self.builder.ins().fcvt_from_sint(types::F64, b); JitValue::Float(self.builder.ins().fdiv(a_float, b_float)) @@ -411,37 +770,57 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { BinaryOperator::Multiply | BinaryOperator::InplaceMultiply, JitValue::Int(a), JitValue::Int(b), - ) => JitValue::Int(self.builder.ins().imul(a, b)), + ) => { + let (out, carry) = self.builder.ins().smul_overflow(a, b); + self.deopt_branch(carry, &[JitValue::Int(a), JitValue::Int(b)])?; + JitValue::Int(out) + } ( BinaryOperator::Remainder | BinaryOperator::InplaceRemainder, JitValue::Int(a), JitValue::Int(b), - ) => JitValue::Int(self.builder.ins().srem(a, b)), + ) => JitValue::Int(self.compile_floor_div(a, b)?.1), ( BinaryOperator::Power | BinaryOperator::InplacePower, JitValue::Int(a), JitValue::Int(b), - ) => JitValue::Int(self.compile_ipow(a, b)), + ) => { + let operands = [JitValue::Int(a), JitValue::Int(b)]; + JitValue::Int(self.compile_ipow(a, b, &operands)?) + } ( - BinaryOperator::Lshift | BinaryOperator::Rshift, + BinaryOperator::Lshift + | BinaryOperator::InplaceLshift + | BinaryOperator::Rshift + | BinaryOperator::InplaceRshift, JitValue::Int(a), JitValue::Int(b), ) => { - // Shifts throw an exception if we have a negative shift count - // Remove all bits except the sign bit, and trap if its 1 (i.e. negative). - let sign = self.builder.ins().ushr_imm(b, 63); - self.builder.ins().trapnz( - sign, - TrapCode::user(CustomTrapCode::NegativeShiftCount as u8).unwrap(), - ); - - let out = - if matches!(op, BinaryOperator::Lshift | BinaryOperator::InplaceLshift) - { - self.builder.ins().ishl(a, b) - } else { - self.builder.ins().sshr(a, b) - }; + let operands = [JitValue::Int(a), JitValue::Int(b)]; + // A count outside `0..64` is not a machine shift at all: negative + // raises ValueError, and 64 or more is a well-defined answer the + // instruction does not give. An unsigned comparison catches both, + // because a negative count reads as huge. + let out_of_range = + self.builder + .ins() + .icmp_imm(IntCC::UnsignedGreaterThanOrEqual, b, 64); + self.deopt_branch(out_of_range, &operands)?; + + let left = + matches!(op, BinaryOperator::Lshift | BinaryOperator::InplaceLshift); + let out = if left { + let out = self.builder.ins().ishl(a, b); + // Shifting back has to give the operand again, or the bits + // that fell off the top are digits the interpreter would + // have kept. + let back = self.builder.ins().sshr(out, b); + let lost = self.builder.ins().icmp(IntCC::NotEqual, back, a); + self.deopt_branch(lost, &operands)?; + out + } else { + self.builder.ins().sshr(a, b) + }; JitValue::Int(out) } ( @@ -480,26 +859,51 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { BinaryOperator::TrueDivide | BinaryOperator::InplaceTrueDivide, JitValue::Float(a), JitValue::Float(b), - ) => JitValue::Float(self.builder.ins().fdiv(a, b)), + ) => { + let operands = [JitValue::Float(a), JitValue::Float(b)]; + let zero = self.builder.ins().f64const(0.0); + let by_zero = self.builder.ins().fcmp(FloatCC::Equal, b, zero); + self.deopt_branch(by_zero, &operands)?; + JitValue::Float(self.builder.ins().fdiv(a, b)) + } ( BinaryOperator::Power | BinaryOperator::InplacePower, JitValue::Float(a), JitValue::Float(b), - ) => JitValue::Float(self.compile_fpow(a, b)), + ) => { + let operands = [JitValue::Float(a), JitValue::Float(b)]; + JitValue::Float(self.compile_fpow(a, b, &operands)?) + } // Floats and Integers (_, JitValue::Int(a), JitValue::Float(b)) | (_, JitValue::Float(a), JitValue::Int(b)) => { - let operand_one = match a_type.unwrap() { + let a_ty = a_type.unwrap(); + let b_ty = b_type.unwrap(); + + let operand_one = match &a_ty { JitType::Int => self.builder.ins().fcvt_from_sint(types::F64, a), _ => a, }; - let operand_two = match b_type.unwrap() { + let operand_two = match &b_ty { JitType::Int => self.builder.ins().fcvt_from_sint(types::F64, b), _ => b, }; + // The original operands, for a guard to hand back on + // deopt - `operand_one`/`operand_two` above are the + // converted doubles and do not describe the stack + // the interpreter had. Only `TrueDivide` and `Power` + // ever guard, so this is built lazily rather than on + // every arm. + let operands = || { + [ + JitValue::from_type_and_value(a_ty.clone(), a), + JitValue::from_type_and_value(b_ty.clone(), b), + ] + }; + match op { BinaryOperator::Add | BinaryOperator::InplaceAdd => { JitValue::Float(self.builder.ins().fadd(operand_one, operand_two)) @@ -511,10 +915,18 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { JitValue::Float(self.builder.ins().fmul(operand_one, operand_two)) } BinaryOperator::TrueDivide | BinaryOperator::InplaceTrueDivide => { + let zero = self.builder.ins().f64const(0.0); + let by_zero = + self.builder.ins().fcmp(FloatCC::Equal, operand_two, zero); + self.deopt_branch(by_zero, &operands())?; JitValue::Float(self.builder.ins().fdiv(operand_one, operand_two)) } BinaryOperator::Power | BinaryOperator::InplacePower => { - JitValue::Float(self.compile_fpow(operand_one, operand_two)) + JitValue::Float(self.compile_fpow( + operand_one, + operand_two, + &operands(), + )?) } _ => return Err(JitCompileError::NotSupported), } @@ -533,11 +945,14 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { Instruction::Call { argc } => { let nargs = argc.get(arg); - let mut args = Vec::new(); + let mut args = Vec::with_capacity(nargs as usize + 1); for _ in 0..nargs { let arg = self.stack.pop().ok_or(JitCompileError::BadBytecode)?; args.push(arg.into_value().unwrap()); } + // Popping walks the arguments backwards. + args.reverse(); + args.insert(0, self.deopt_ptr); // Pop self_or_null (should be Null for JIT-compiled recursive calls) let self_or_null = self.stack.pop().ok_or(JitCompileError::BadBytecode)?; @@ -563,6 +978,30 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { (Some(ty), Some(val)) => JitValue::from_type_and_value(ty, val), _ => return Err(JitCompileError::NotSupported), }; + + // A nested frame that gave up has already written its + // record, and returned a filler in place of a result. + // Anything this frame computes from here on is built on + // that filler, so stop. + let status = self.builder.ins().load( + types::I64, + MemFlags::trusted(), + self.deopt_ptr, + 0, + ); + let nested = self.builder.ins().icmp_imm(IntCC::NotEqual, status, 0); + // The record standing in the buffer describes that + // frame, not this one. Overwrite the status so the + // resume path restarts the call rather than continuing + // this frame from an offset that belongs to another. + self.deopt_if(nested, |this| { + let sentinel = this + .builder + .ins() + .iconst(types::I64, DEOPT_STATUS_NESTED as i64); + this.store_raw(sentinel, 0); + }); + self.stack.push(val); Ok(()) @@ -651,8 +1090,16 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { Instruction::JumpBackward { .. } | Instruction::JumpBackwardNoInterrupt { .. } | Instruction::JumpForward { .. } => { - let target = Self::instruction_target(offset, instruction, arg)? + let target = instruction_target(offset, instruction, arg)? .ok_or(JitCompileError::BadBytecode)?; + let backward = target.as_u32() <= offset; + self.has_backward_jump |= backward; + self.require_empty_stack_at_merge()?; + // Only a jump to an earlier offset can loop, and only a loop + // can keep the thread from reaching the end of the function. + if backward { + self.compile_safepoint()?; + } let target_block = self.get_or_create_block(target); self.builder.ins().jump(target_block, &[]); Ok(()) @@ -703,22 +1150,30 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let oparg = namei.get(arg); let name = &bytecode.names[(oparg >> 1) as usize]; - if name.as_ref() != bytecode.obj_name.as_ref() { - Err(JitCompileError::NotSupported) - } else { - self.stack.push(JitValue::FuncRef(func_ref)); - if (oparg & 1) != 0 { - self.stack.push(JitValue::Null); - } - Ok(()) + // The only global with a lowering is this function itself, + // matched by name. Strict turns even that down: the interpreter + // reads the globals dict on every call, so rebinding the name - + // a decorator applied later, a test patching the module - makes + // the two disagree about what gets called. + let is_self_call = self.safety == Safety::Permissive + && name.as_ref() == bytecode.obj_name.as_ref(); + if !is_self_call { + return Err(JitCompileError::NotSupported); } + + self.stack.push(JitValue::FuncRef(func_ref)); + if (oparg & 1) != 0 { + self.stack.push(JitValue::Null); + } + Ok(()) } Instruction::Nop | Instruction::NotTaken => Ok(()), Instruction::PopJumpIfFalse { .. } => { let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?; let val = self.boolean_val(cond)?; - let then_label = Self::instruction_target(offset, instruction, arg)? + let then_label = instruction_target(offset, instruction, arg)? .ok_or(JitCompileError::BadBytecode)?; + self.require_empty_stack_at_merge()?; let then_block = self.get_or_create_block(then_label); let else_block = self.builder.create_block(); @@ -732,8 +1187,9 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { Instruction::PopJumpIfTrue { .. } => { let cond = self.stack.pop().ok_or(JitCompileError::BadBytecode)?; let val = self.boolean_val(cond)?; - let then_label = Self::instruction_target(offset, instruction, arg)? + let then_label = instruction_target(offset, instruction, arg)? .ok_or(JitCompileError::BadBytecode)?; + self.require_empty_stack_at_merge()?; let then_block = self.get_or_create_block(then_label); let else_block = self.builder.create_block(); @@ -807,9 +1263,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { Instruction::UnaryNegative => { match self.stack.pop().ok_or(JitCompileError::BadBytecode)? { JitValue::Int(val) => { - // Compile minus as 0 - val. + // Compile minus as 0 - val. The zero is not on the + // interpreter's stack, so only `val` is recorded. let zero = self.builder.ins().iconst(types::I64, 0); - let out = self.compile_sub(zero, val); + let out = self.compile_sub(zero, val, &[JitValue::Int(val)])?; self.stack.push(JitValue::Int(out)); Ok(()) } @@ -835,586 +1292,148 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { } } - fn compile_sub(&mut self, a: Value, b: Value) -> Value { + fn compile_sub( + &mut self, + a: Value, + b: Value, + popped: &[JitValue], + ) -> Result { let (out, carry) = self.builder.ins().ssub_overflow(a, b); - self.builder.ins().trapnz(carry, TrapCode::INTEGER_OVERFLOW); - out + self.deopt_branch(carry, popped)?; + Ok(out) } - /// Creates a double–double (DDValue) from a regular f64 constant. - /// The high part is set to x and the low part is set to 0.0. - fn dd_from_f64(&mut self, x: f64) -> DDValue { - DDValue { - hi: self.builder.ins().f64const(x), - lo: self.builder.ins().f64const(0.0), - } - } - - /// Creates a DDValue from a Value (assumed to represent an f64). - /// This function initializes the high part with x and the low part to 0.0. - fn dd_from_value(&mut self, x: Value) -> DDValue { - DDValue { - hi: x, - lo: self.builder.ins().f64const(0.0), - } - } - - /// Creates a DDValue from two f64 parts. - /// The 'hi' parameter sets the high part and 'lo' sets the low part. - fn dd_from_parts(&mut self, hi: f64, lo: f64) -> DDValue { - DDValue { - hi: self.builder.ins().f64const(hi), - lo: self.builder.ins().f64const(lo), - } - } - - /// Converts a DDValue back to a single f64 value by adding the high and low parts. - fn dd_to_f64(&mut self, dd: DDValue) -> Value { - self.builder.ins().fadd(dd.hi, dd.lo) - } - - /// Computes the negation of a DDValue. - /// It subtracts both the high and low parts from zero. - fn dd_neg(&mut self, dd: DDValue) -> DDValue { - let zero = self.builder.ins().f64const(0.0); - DDValue { - hi: self.builder.ins().fsub(zero, dd.hi), - lo: self.builder.ins().fsub(zero, dd.lo), - } - } - - /// Adds two DDValue numbers using error-free transformations to maintain extra precision. - /// It carefully adds the high parts, computes the rounding error, adds the low parts along with the error, - /// and then normalizes the result. - fn dd_add(&mut self, a: DDValue, b: DDValue) -> DDValue { - // Compute the sum of the high parts. - let s = self.builder.ins().fadd(a.hi, b.hi); - // Compute t = s - a.hi to capture part of the rounding error. - let t = self.builder.ins().fsub(s, a.hi); - // Compute the error e from the high part additions. - let s_minus_t = self.builder.ins().fsub(s, t); - let part1 = self.builder.ins().fsub(a.hi, s_minus_t); - let part2 = self.builder.ins().fsub(b.hi, t); - let e = self.builder.ins().fadd(part1, part2); - // Sum the low parts along with the error. - let lo = self.builder.ins().fadd(a.lo, b.lo); - let lo_sum = self.builder.ins().fadd(lo, e); - // Renormalize: add the low sum to s and compute a new low component. - let hi_new = self.builder.ins().fadd(s, lo_sum); - let hi_new_minus_s = self.builder.ins().fsub(hi_new, s); - let lo_new = self.builder.ins().fsub(lo_sum, hi_new_minus_s); - DDValue { - hi: hi_new, - lo: lo_new, - } - } - - /// Subtracts DDValue b from DDValue a by negating b and then using the addition function. - fn dd_sub(&mut self, a: DDValue, b: DDValue) -> DDValue { - let neg_b = self.dd_neg(b); - self.dd_add(a, neg_b) - } - - /// Multiplies two DDValue numbers using double–double arithmetic. - /// It calculates the high product, uses a fused multiply–add (FMA) to capture rounding error, - /// computes the cross products, and then normalizes the result. - fn dd_mul(&mut self, a: DDValue, b: DDValue) -> DDValue { - // p = a.hi * b.hi (primary product) - let p = self.builder.ins().fmul(a.hi, b.hi); - // err = fma(a.hi, b.hi, -p) recovers the rounding error. - let zero = self.builder.ins().f64const(0.0); - let neg_p = self.builder.ins().fsub(zero, p); - let err = self.builder.ins().fma(a.hi, b.hi, neg_p); - // Compute cross terms: a.hi*b.lo + a.lo*b.hi. - let a_hi_b_lo = self.builder.ins().fmul(a.hi, b.lo); - let a_lo_b_hi = self.builder.ins().fmul(a.lo, b.hi); - let cross = self.builder.ins().fadd(a_hi_b_lo, a_lo_b_hi); - // Sum p and the cross terms. - let s = self.builder.ins().fadd(p, cross); - // Isolate rounding error from the addition. - let t = self.builder.ins().fsub(s, p); - let s_minus_t = self.builder.ins().fsub(s, t); - let part1 = self.builder.ins().fsub(p, s_minus_t); - let part2 = self.builder.ins().fsub(cross, t); - let e = self.builder.ins().fadd(part1, part2); - // Include the error from the low parts multiplication. - let a_lo_b_lo = self.builder.ins().fmul(a.lo, b.lo); - let err_plus_e = self.builder.ins().fadd(err, e); - let lo_sum = self.builder.ins().fadd(err_plus_e, a_lo_b_lo); - // Renormalize the sum. - let hi_new = self.builder.ins().fadd(s, lo_sum); - let hi_new_minus_s = self.builder.ins().fsub(hi_new, s); - let lo_new = self.builder.ins().fsub(lo_sum, hi_new_minus_s); - DDValue { - hi: hi_new, - lo: lo_new, - } - } - - /// Multiplies a DDValue by a regular f64 (Value) using similar techniques as dd_mul. - /// It multiplies both the high and low parts by b, computes the rounding error, - /// and then renormalizes the result. - fn dd_mul_f64(&mut self, a: DDValue, b: Value) -> DDValue { - // p = a.hi * b (primary product) - let p = self.builder.ins().fmul(a.hi, b); - // Compute the rounding error using fma. - let zero = self.builder.ins().f64const(0.0); - let neg_p = self.builder.ins().fsub(zero, p); - let err = self.builder.ins().fma(a.hi, b, neg_p); - // Multiply the low part. - let cross = self.builder.ins().fmul(a.lo, b); - // Sum the primary product and the low multiplication. - let s = self.builder.ins().fadd(p, cross); - // Capture rounding error from addition. - let t = self.builder.ins().fsub(s, p); - let s_minus_t = self.builder.ins().fsub(s, t); - let part1 = self.builder.ins().fsub(p, s_minus_t); - let part2 = self.builder.ins().fsub(cross, t); - let e = self.builder.ins().fadd(part1, part2); - // Combine the error components. - let lo_sum = self.builder.ins().fadd(err, e); - // Renormalize to form the final double–double number. - let hi_new = self.builder.ins().fadd(s, lo_sum); - let hi_new_minus_s = self.builder.ins().fsub(hi_new, s); - let lo_new = self.builder.ins().fsub(lo_sum, hi_new_minus_s); - DDValue { - hi: hi_new, - lo: lo_new, - } - } - - /// Scales a DDValue by multiplying both its high and low parts by the given factor. - fn dd_scale(&mut self, dd: DDValue, factor: Value) -> DDValue { - DDValue { - hi: self.builder.ins().fmul(dd.hi, factor), - lo: self.builder.ins().fmul(dd.lo, factor), - } - } - - /// Approximates ln(1+f) using its Taylor series expansion in double–double arithmetic. - /// It computes the series ∑ (-1)^(i-1) * f^i / i from i = 1 to 1000 for high precision. - fn dd_ln_1p_series(&mut self, f: Value) -> DDValue { - // Convert f to a DDValue and initialize the sum and term. - let f_dd = self.dd_from_value(f); - let mut sum = f_dd.clone(); - let mut term = f_dd; - // Alternating sign starts at -1 for the second term. - let mut sign = -1.0_f64; - let range = 1000; - - // Loop over terms from i = 2 to 1000. - for i in 2..=range { - // Compute f^i by multiplying the previous term by f. - term = self.dd_mul_f64(term, f); - // Divide the term by i. - let inv_i = 1.0 / (i as f64); - let c_inv_i = self.builder.ins().f64const(inv_i); - let term_div = self.dd_mul_f64(term.clone(), c_inv_i); - // Multiply by the alternating sign. - let dd_sign = self.dd_from_f64(sign); - let to_add = self.dd_mul(dd_sign, term_div); - // Add the term to the cumulative sum. - sum = self.dd_add(sum, to_add); - // Flip the sign for the next term. - sign = -sign; - } - sum - } - - /// Computes the natural logarithm ln(x) in double–double arithmetic. - /// It first checks for domain errors (x ≤ 0 or NaN), then extracts the exponent - /// and mantissa from the bit-level representation of x. It computes ln(mantissa) using - /// the ln(1+f) series and adds k*ln2 to obtain ln(x). - fn dd_ln(&mut self, x: Value) -> DDValue { - // (A) Prepare a DDValue representing NaN. - let dd_nan = self.dd_from_f64(f64::NAN); - - // Build a zero constant for comparisons. - let zero_f64 = self.builder.ins().f64const(0.0); - - // Check if x is less than or equal to 0 or is NaN. - let cmp_le = self - .builder - .ins() - .fcmp(FloatCC::LessThanOrEqual, x, zero_f64); - let cmp_nan = self.builder.ins().fcmp(FloatCC::Unordered, x, x); - let need_nan = self.builder.ins().bor(cmp_le, cmp_nan); - - // (B) Reinterpret the bits of x as an integer. - let bits = self.builder.ins().bitcast(types::I64, MemFlags::new(), x); - - // (C) Extract the exponent (top 11 bits) from the bit representation. - let shift_52 = self.builder.ins().ushr_imm(bits, 52); - let exponent_mask = self.builder.ins().iconst(types::I64, 0x7FF); - let exponent = self.builder.ins().band(shift_52, exponent_mask); - - // k = exponent - 1023 (unbias the exponent). - let bias = self.builder.ins().iconst(types::I64, 1023); - let k_i64 = self.builder.ins().isub(exponent, bias); - - // (D) Extract the fraction (mantissa) from the lower 52 bits. - let fraction_mask = self.builder.ins().iconst(types::I64, 0x000F_FFFF_FFFF_FFFF); - let fraction_part = self.builder.ins().band(bits, fraction_mask); - - // (E) For normal numbers (exponent ≠ 0), add the implicit leading 1. - let implicit_one = self.builder.ins().iconst(types::I64, 1 << 52); - let zero_exp = self.builder.ins().icmp_imm(IntCC::Equal, exponent, 0); - let frac_one_bor = self.builder.ins().bor(fraction_part, implicit_one); - let fraction_with_leading_one = self.builder.ins().select( - zero_exp, - fraction_part, // For subnormals, do not add the implicit 1. - frac_one_bor, - ); - - // (F) Force the exponent bits to 1023, yielding a mantissa m in [1, 2). - let new_exp = self.builder.ins().iconst(types::I64, 0x3FF0_0000_0000_0000); - let fraction_bits = self.builder.ins().bor(fraction_with_leading_one, new_exp); - let m = self - .builder - .ins() - .bitcast(types::F64, MemFlags::new(), fraction_bits); - - // (G) Compute ln(m) using the series ln(1+f) with f = m - 1. - let one_f64 = self.builder.ins().f64const(1.0); - let f_val = self.builder.ins().fsub(m, one_f64); - let dd_ln_m = self.dd_ln_1p_series(f_val); - - // (H) Compute k*ln2 in double–double arithmetic. - let ln2_dd = self.dd_from_parts( - f64::from_bits(0x3fe62e42fefa39ef), - f64::from_bits(0x3c7abc9e3b39803f), - ); - let k_f64 = self.builder.ins().fcvt_from_sint(types::F64, k_i64); - let dd_ln2_k = self.dd_mul_f64(ln2_dd, k_f64); - - // Add ln(m) and k*ln2 to get the final ln(x). - let normal_result = self.dd_add(dd_ln_m, dd_ln2_k); - - // (I) If x was nonpositive or NaN, return NaN; otherwise, return the computed result. - let final_hi = self - .builder - .ins() - .select(need_nan, dd_nan.hi, normal_result.hi); - let final_lo = self - .builder - .ins() - .select(need_nan, dd_nan.lo, normal_result.lo); - - DDValue { - hi: final_hi, - lo: final_lo, - } - } - - /// Computes the exponential function exp(x) in double–double arithmetic. - /// It uses range reduction to write x = k*ln2 + r, computes exp(r) via a Taylor series, - /// scales the result by 2^k, and handles overflow by checking if k exceeds the maximum. - fn dd_exp(&mut self, dd: DDValue) -> DDValue { - // (A) Range reduction: Convert dd to a single f64 value. - let x = self.dd_to_f64(dd.clone()); - let ln2_f64 = self - .builder - .ins() - .f64const(f64::from_bits(0x3fe62e42fefa39ef)); - let div = self.builder.ins().fdiv(x, ln2_f64); - let half = self.builder.ins().f64const(0.5); - let div_plus_half = self.builder.ins().fadd(div, half); - // Rounding: floor(div + 0.5) gives the nearest integer k. - let k = self.builder.ins().fcvt_to_sint(types::I64, div_plus_half); - - // --- OVERFLOW CHECK --- - // Check if k is greater than the maximum exponent for finite doubles (1023). - let max_k = self.builder.ins().iconst(types::I64, 1023); - let is_overflow = self.builder.ins().icmp(IntCC::SignedGreaterThan, k, max_k); - - // Define infinity and zero for the overflow case. - let inf = self.builder.ins().f64const(f64::INFINITY); - let zero = self.builder.ins().f64const(0.0); - - // (B) Compute exp(x) normally when not overflowing. - // Compute k*ln2 in double–double arithmetic and subtract it from x. - let ln2_dd = self.dd_from_parts( - f64::from_bits(0x3fe62e42fefa39ef), - f64::from_bits(0x3c7abc9e3b39803f), - ); - let k_f64 = self.builder.ins().fcvt_from_sint(types::F64, k); - let k_ln2 = self.dd_mul_f64(ln2_dd, k_f64); - let r = self.dd_sub(dd, k_ln2); - - // Compute exp(r) using a Taylor series. - let mut sum = self.dd_from_f64(1.0); // Initialize sum to 1. - let mut term = self.dd_from_f64(1.0); // Initialize the first term to 1. - let n_terms = 1000; - for i in 1..=n_terms { - term = self.dd_mul(term, r.clone()); - let inv = 1.0 / (i as f64); - let inv_const = self.builder.ins().f64const(inv); - term = self.dd_mul_f64(term, inv_const); - sum = self.dd_add(sum, term.clone()); - } - - // Reconstruct the final result by scaling with 2^k. - let bias = self.builder.ins().iconst(types::I64, 1023); - let k_plus_bias = self.builder.ins().iadd(k, bias); - let shift_count = self.builder.ins().iconst(types::I64, 52); - let shifted = self.builder.ins().ishl(k_plus_bias, shift_count); - let two_to_k = self - .builder - .ins() - .bitcast(types::F64, MemFlags::new(), shifted); - let result = self.dd_scale(sum, two_to_k); - - // (C) If overflow was detected, return infinity; otherwise, return the computed value. - let final_hi = self.builder.ins().select(is_overflow, inf, result.hi); - let final_lo = self.builder.ins().select(is_overflow, zero, result.lo); - DDValue { - hi: final_hi, - lo: final_lo, - } + /// Floor division and its remainder, which are defined together: the + /// quotient rounds toward negative infinity and the remainder takes the + /// divisor's sign. `sdiv` and `srem` round toward zero and take the + /// dividend's sign, so both need a correction on the same condition. + /// + /// Two cases have no 64-bit answer at all and deoptimize: a zero divisor, + /// which raises, and `i64::MIN // -1`, whose quotient is one past the top + /// of the range. The guard is shared between the two results, so + /// `i64::MIN % -1` deopts alongside it even though `0` is a perfectly + /// good remainder - the pair is computed together. + fn compile_floor_div(&mut self, a: Value, b: Value) -> Result<(Value, Value), JitCompileError> { + let operands = [JitValue::Int(a), JitValue::Int(b)]; + let by_zero = self.builder.ins().icmp_imm(IntCC::Equal, b, 0); + self.deopt_branch(by_zero, &operands)?; + + let min = self.builder.ins().icmp_imm(IntCC::Equal, a, i64::MIN); + let neg_one = self.builder.ins().icmp_imm(IntCC::Equal, b, -1); + let overflows = self.builder.ins().band(min, neg_one); + self.deopt_branch(overflows, &operands)?; + + let quotient = self.builder.ins().sdiv(a, b); + // The remainder as `a - quotient * b` rather than a second `srem`: + // cranelift keeps a trapping division live even when nothing reads + // its result, so asking for both would cost two hardware divisions + // instead of one. The multiply cannot overflow: `quotient` truncates + // toward zero, so `|quotient * b| <= |a|` and `quotient * b` shares + // `a`'s sign, making the subtraction one between same-signed values + // with the smaller magnitude second. + let product = self.builder.ins().imul(quotient, b); + let remainder = self.builder.ins().isub(a, product); + + // The two disagree with Python exactly when the division was + // inexact and the operands had opposite signs. + let inexact = self.builder.ins().icmp_imm(IntCC::NotEqual, remainder, 0); + let mixed = self.builder.ins().bxor(a, b); + let mixed = self.builder.ins().icmp_imm(IntCC::SignedLessThan, mixed, 0); + let correct = self.builder.ins().band(inexact, mixed); + + let one = self.builder.ins().iconst(types::I64, 1); + let zero = self.builder.ins().iconst(types::I64, 0); + let quotient_adjust = self.builder.ins().select(correct, one, zero); + // `isub` cannot overflow: it only subtracts when `correct` holds, + // which requires `remainder != 0`. A quotient of `i64::MIN` is still + // reachable here - after the guards above, only `b == 1` produces + // it - but that division is exact, so `remainder` is `0` and + // `correct` is false there. + let quotient = self.builder.ins().isub(quotient, quotient_adjust); + + let remainder_adjust = self.builder.ins().select(correct, b, zero); + let remainder = self.builder.ins().iadd(remainder, remainder_adjust); + + Ok((quotient, remainder)) } - /// Computes the power function a^b (f_pow) for f64 values using double–double arithmetic for high precision. - /// It handles different cases for the base 'a': - /// - For a > 0: Computes exp(b * ln(a)). - /// - For a == 0: Handles special cases for 0^b, including returning 0, 1, or a domain error. - /// - For a < 0: Allows only an integer exponent b and adjusts the sign if b is odd. - fn compile_fpow(&mut self, a: Value, b: Value) -> Value { - let f64_ty = types::F64; - let i64_ty = types::I64; + /// Computes a raised to the power b by calling the same `f64::powf` the + /// interpreter's `float_pow` calls, once the two guards ahead of it in + /// `float_pow` rule out its other two outcomes: `ZeroDivisionError` and a + /// complex result. This is `float_pow` translated guard for guard. + fn compile_fpow( + &mut self, + a: Value, + b: Value, + operands: &[JitValue], + ) -> Result { let zero_f = self.builder.ins().f64const(0.0); - let one_f = self.builder.ins().f64const(1.0); - let nan_f = self.builder.ins().f64const(f64::NAN); - let inf_f = self.builder.ins().f64const(f64::INFINITY); - let neg_inf_f = self.builder.ins().f64const(f64::NEG_INFINITY); - - // Merge block for final result. - let merge_block = self.builder.create_block(); - self.builder.append_block_param(merge_block, f64_ty); - - // --- Edge Case 1: b == 0.0 → return 1.0 - let cmp_b_zero = self.builder.ins().fcmp(FloatCC::Equal, b, zero_f); - let b_zero_block = self.builder.create_block(); - let continue_block = self.builder.create_block(); - self.builder - .ins() - .brif(cmp_b_zero, b_zero_block, &[], continue_block, &[]); - self.builder.switch_to_block(b_zero_block); - self.builder.ins().jump(merge_block, &[one_f.into()]); - self.builder.switch_to_block(continue_block); - // --- Edge Case 2: b is NaN → return NaN - let cmp_b_nan = self.builder.ins().fcmp(FloatCC::Unordered, b, b); - let b_nan_block = self.builder.create_block(); - let continue_block2 = self.builder.create_block(); - self.builder - .ins() - .brif(cmp_b_nan, b_nan_block, &[], continue_block2, &[]); - self.builder.switch_to_block(b_nan_block); - self.builder.ins().jump(merge_block, &[nan_f.into()]); - self.builder.switch_to_block(continue_block2); - - // --- Edge Case 3: a == 0.0 → return 0.0 - let cmp_a_zero = self.builder.ins().fcmp(FloatCC::Equal, a, zero_f); - let a_zero_block = self.builder.create_block(); - let continue_block3 = self.builder.create_block(); - self.builder - .ins() - .brif(cmp_a_zero, a_zero_block, &[], continue_block3, &[]); - self.builder.switch_to_block(a_zero_block); - self.builder.ins().jump(merge_block, &[zero_f.into()]); - self.builder.switch_to_block(continue_block3); - - // --- Edge Case 4: a is NaN → return NaN - let cmp_a_nan = self.builder.ins().fcmp(FloatCC::Unordered, a, a); - let a_nan_block = self.builder.create_block(); - let continue_block4 = self.builder.create_block(); - self.builder - .ins() - .brif(cmp_a_nan, a_nan_block, &[], continue_block4, &[]); - self.builder.switch_to_block(a_nan_block); - self.builder.ins().jump(merge_block, &[nan_f.into()]); - self.builder.switch_to_block(continue_block4); - - // --- Edge Case 5: b == +infinity → return +infinity - let cmp_b_inf = self.builder.ins().fcmp(FloatCC::Equal, b, inf_f); - let b_inf_block = self.builder.create_block(); - let continue_block5 = self.builder.create_block(); - self.builder - .ins() - .brif(cmp_b_inf, b_inf_block, &[], continue_block5, &[]); - self.builder.switch_to_block(b_inf_block); - self.builder.ins().jump(merge_block, &[inf_f.into()]); - self.builder.switch_to_block(continue_block5); - - // --- Edge Case 6: b == -infinity → return 0.0 - let cmp_b_neg_inf = self.builder.ins().fcmp(FloatCC::Equal, b, neg_inf_f); - let b_neg_inf_block = self.builder.create_block(); - let continue_block6 = self.builder.create_block(); - self.builder - .ins() - .brif(cmp_b_neg_inf, b_neg_inf_block, &[], continue_block6, &[]); - self.builder.switch_to_block(b_neg_inf_block); - self.builder.ins().jump(merge_block, &[zero_f.into()]); - self.builder.switch_to_block(continue_block6); - - // --- Edge Case 7: a == +infinity → return +infinity - let cmp_a_inf = self.builder.ins().fcmp(FloatCC::Equal, a, inf_f); - let a_inf_block = self.builder.create_block(); - let continue_block7 = self.builder.create_block(); - self.builder - .ins() - .brif(cmp_a_inf, a_inf_block, &[], continue_block7, &[]); - self.builder.switch_to_block(a_inf_block); - self.builder.ins().jump(merge_block, &[inf_f.into()]); - self.builder.switch_to_block(continue_block7); - - // --- Edge Case 8: a == -infinity → check exponent parity - let cmp_a_neg_inf = self.builder.ins().fcmp(FloatCC::Equal, a, neg_inf_f); - let a_neg_inf_block = self.builder.create_block(); - let continue_block8 = self.builder.create_block(); - self.builder - .ins() - .brif(cmp_a_neg_inf, a_neg_inf_block, &[], continue_block8, &[]); - - self.builder.switch_to_block(a_neg_inf_block); - // a is -infinity here. First, ensure that b is an integer. + // v1.is_zero() && v2 < 0.0 -> ZeroDivisionError. Both comparisons are + // ordered, so neither a `-0.0` nor a nan operand takes this exit, and + // both `0.0` and `-0.0` count as a zero base. + let base_zero = self.builder.ins().fcmp(FloatCC::Equal, a, zero_f); + let exp_negative = self.builder.ins().fcmp(FloatCC::LessThan, b, zero_f); + let divides_by_zero = self.builder.ins().band(base_zero, exp_negative); + self.deopt_branch(divides_by_zero, operands)?; + + // v1 < 0.0 && v2.is_finite() && v2 != v2.floor() -> complex result. + // A value never sits below its own floor, so ordered `>` against it is + // that pair of conditions in one: an infinity equals its floor, and a + // nan is unordered against everything, itself included. + let base_negative = self.builder.ins().fcmp(FloatCC::LessThan, a, zero_f); let b_floor = self.builder.ins().floor(b); - let cmp_int = self.builder.ins().fcmp(FloatCC::Equal, b_floor, b); - let domain_error_blk = self.builder.create_block(); - let continue_neg_inf = self.builder.create_block(); - self.builder - .ins() - .brif(cmp_int, continue_neg_inf, &[], domain_error_blk, &[]); - - self.builder.switch_to_block(domain_error_blk); - self.builder.ins().jump(merge_block, &[nan_f.into()]); - - self.builder.switch_to_block(continue_neg_inf); - // b is an integer here; convert b_floor to an i64. - let b_i64 = self.builder.ins().fcvt_to_sint(i64_ty, b_floor); - let one_i = self.builder.ins().iconst(i64_ty, 1); - let remainder = self.builder.ins().band(b_i64, one_i); - let zero_i = self.builder.ins().iconst(i64_ty, 0); - let is_odd = self.builder.ins().icmp(IntCC::NotEqual, remainder, zero_i); - - // Create separate blocks for odd and even cases. - let odd_block = self.builder.create_block(); - let even_block = self.builder.create_block(); - self.builder.append_block_param(odd_block, f64_ty); - self.builder.append_block_param(even_block, f64_ty); - self.builder.ins().brif( - is_odd, - odd_block, - &[neg_inf_f.into()], - even_block, - &[inf_f.into()], - ); - - self.builder.switch_to_block(odd_block); - let phi_neg_inf = self.builder.block_params(odd_block)[0]; - self.builder.ins().jump(merge_block, &[phi_neg_inf.into()]); - - self.builder.switch_to_block(even_block); - let phi_inf = self.builder.block_params(even_block)[0]; - self.builder.ins().jump(merge_block, &[phi_inf.into()]); - - self.builder.switch_to_block(continue_block8); - - // --- Normal branch: neither a nor b hit the special cases. - // Here we branch based on the sign of a. - let cmp_lt = self.builder.ins().fcmp(FloatCC::LessThan, a, zero_f); - let a_neg_block = self.builder.create_block(); - let a_pos_block = self.builder.create_block(); - self.builder - .ins() - .brif(cmp_lt, a_neg_block, &[], a_pos_block, &[]); - - // ----- Case: a > 0: Compute a^b = exp(b * ln(a)) using double–double arithmetic. - self.builder.switch_to_block(a_pos_block); - let ln_a_dd = self.dd_ln(a); - let b_dd = self.dd_from_value(b); - let product_dd = self.dd_mul(ln_a_dd, b_dd); - let exp_dd = self.dd_exp(product_dd); - let pos_res = self.dd_to_f64(exp_dd); - self.builder.ins().jump(merge_block, &[pos_res.into()]); - - // ----- Case: a < 0: Only allow an integral exponent. - self.builder.switch_to_block(a_neg_block); - let b_floor = self.builder.ins().floor(b); - let cmp_int = self.builder.ins().fcmp(FloatCC::Equal, b_floor, b); - let neg_int_block = self.builder.create_block(); - let domain_error_blk = self.builder.create_block(); - self.builder - .ins() - .brif(cmp_int, neg_int_block, &[], domain_error_blk, &[]); - - // Domain error: non-integer exponent for negative base - self.builder.switch_to_block(domain_error_blk); - self.builder.ins().jump(merge_block, &[nan_f.into()]); + let fractional = self.builder.ins().fcmp(FloatCC::GreaterThan, b, b_floor); + let complex = self.builder.ins().band(base_negative, fractional); + self.deopt_branch(complex, operands)?; + + // ans = v1.powf(v2), through the exact function the interpreter + // calls - see `jit_powf`'s doc comment for why it is looked up by an + // explicit symbol rather than left for the JIT to resolve. + let call = self.builder.ins().call(self.externals.powf, &[a, b]); + let ans = match *self.builder.inst_results(call) { + [ans] => ans, + _ => return Err(JitCompileError::NotSupported), + }; - // For negative base with an integer exponent: - self.builder.switch_to_block(neg_int_block); + // ans.is_infinite() && !(v1.is_infinite() || v2.is_infinite()) -> + // OverflowError. An already-infinite operand (`inf ** 2.0`, + // `2.0 ** inf`) must keep answering with its infinity, hence the + // exemption. + let inf_f = self.builder.ins().f64const(f64::INFINITY); + let abs_ans = self.builder.ins().fabs(ans); + let ans_infinite = self.builder.ins().fcmp(FloatCC::Equal, abs_ans, inf_f); let abs_a = self.builder.ins().fabs(a); - let ln_abs_dd = self.dd_ln(abs_a); - let b_dd = self.dd_from_value(b); - let product_dd = self.dd_mul(ln_abs_dd, b_dd); - let exp_dd = self.dd_exp(product_dd); - let mag_val = self.dd_to_f64(exp_dd); - - let b_i64 = self.builder.ins().fcvt_to_sint(i64_ty, b_floor); - let one_i = self.builder.ins().iconst(i64_ty, 1); - let remainder = self.builder.ins().band(b_i64, one_i); - let zero_i = self.builder.ins().iconst(i64_ty, 0); - let is_odd = self.builder.ins().icmp(IntCC::NotEqual, remainder, zero_i); - - let odd_block = self.builder.create_block(); - let even_block = self.builder.create_block(); - // Append block parameters for both branches: - self.builder.append_block_param(odd_block, f64_ty); - self.builder.append_block_param(even_block, f64_ty); - // Pass mag_val to both branches: - self.builder.ins().brif( - is_odd, - odd_block, - &[mag_val.into()], - even_block, - &[mag_val.into()], - ); - - self.builder.switch_to_block(odd_block); - let phi_mag_val = self.builder.block_params(odd_block)[0]; - let neg_val = self.builder.ins().fneg(phi_mag_val); - self.builder.ins().jump(merge_block, &[neg_val.into()]); - - self.builder.switch_to_block(even_block); - let phi_mag_val_even = self.builder.block_params(even_block)[0]; - self.builder - .ins() - .jump(merge_block, &[phi_mag_val_even.into()]); - - // ----- Merge: Return the final result. - self.builder.switch_to_block(merge_block); - self.builder.block_params(merge_block)[0] + let a_infinite = self.builder.ins().fcmp(FloatCC::Equal, abs_a, inf_f); + let abs_b = self.builder.ins().fabs(b); + let b_infinite = self.builder.ins().fcmp(FloatCC::Equal, abs_b, inf_f); + let either_infinite = self.builder.ins().bor(a_infinite, b_infinite); + let overflowed = self.builder.ins().band_not(ans_infinite, either_infinite); + self.deopt_branch(overflowed, operands)?; + + Ok(ans) } - fn compile_ipow(&mut self, a: Value, b: Value) -> Value { + fn compile_ipow( + &mut self, + a: Value, + b: Value, + operands: &[JitValue], + ) -> Result { + // A negative exponent makes this a float; the loop below only + // computes non-negative integer powers. + let negative = self.builder.ins().icmp_imm(IntCC::SignedLessThan, b, 0); + self.deopt_branch(negative, operands)?; + let zero = self.builder.ins().iconst(types::I64, 0); let one_i64 = self.builder.ins().iconst(types::I64, 1); // Create required blocks - let check_negative = self.builder.create_block(); - let handle_negative = self.builder.create_block(); let loop_block = self.builder.create_block(); let continue_block = self.builder.create_block(); let exit_block = self.builder.create_block(); // Set up block parameters - self.builder.append_block_param(check_negative, types::I64); // exponent - self.builder.append_block_param(check_negative, types::I64); // base - - self.builder.append_block_param(handle_negative, types::I64); // abs(exponent) - self.builder.append_block_param(handle_negative, types::I64); // base - self.builder.append_block_param(loop_block, types::I64); // exponent self.builder.append_block_param(loop_block, types::I64); // result self.builder.append_block_param(loop_block, types::I64); // base @@ -1426,32 +1445,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { self.builder.append_block_param(continue_block, types::I64); // result self.builder.append_block_param(continue_block, types::I64); // base - // Initial jump to check if exponent is negative + // The exponent is known non-negative, so jump straight into the loop. self.builder .ins() - .jump(check_negative, &[b.into(), a.into()]); - - // Check if exponent is negative - self.builder.switch_to_block(check_negative); - let params = self.builder.block_params(check_negative); - let exp_check = params[0]; - let base_check = params[1]; - - let is_negative = self - .builder - .ins() - .icmp(IntCC::SignedLessThan, exp_check, zero); - self.builder.ins().brif( - is_negative, - handle_negative, - &[exp_check.into(), base_check.into()], - loop_block, - &[exp_check.into(), one_i64.into(), base_check.into()], - ); - - // Handle negative exponent (return 0 for integer exponentiation) - self.builder.switch_to_block(handle_negative); - self.builder.ins().jump(exit_block, &[zero.into()]); // Return 0 for negative exponents + .jump(loop_block, &[b.into(), one_i64.into(), a.into()]); // Loop block logic (square-and-multiply algorithm) self.builder.switch_to_block(loop_block); @@ -1480,12 +1477,27 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { // If exponent is odd, multiply result by base let is_odd = self.builder.ins().band_imm(exp_phi, 1); let is_odd = self.builder.ins().icmp_imm(IntCC::Equal, is_odd, 1); - let mul_result = self.builder.ins().imul(result_phi, base_phi); + + // Unlike the squaring below, this carry does not need masking to + // stay exact: `continue_block` only runs with `exp != 0`, so a clear + // low bit still forces `exp >= 2`, and once `|base| >= 2` an + // overflowing `result * base` means `result * base^exp` overflows + // too - a later guard always catches it - while with `|base| <= 1` + // the product cannot overflow at all. + let (mul_result, mul_carry) = self.builder.ins().smul_overflow(result_phi, base_phi); + self.deopt_branch(mul_carry, operands)?; let new_result = self.builder.ins().select(is_odd, mul_result, result_phi); - // Square the base and divide exponent by 2 - let squared_base = self.builder.ins().imul(base_phi, base_phi); + // The squared base is read only if there is another iteration to + // read it, and this mask is the one that is measurably load-bearing: + // dropping it deopts `2 ** 33` and `2 ** 62` on a squaring whose + // overflowed result the exit branch above never reads. + let (squared_base, square_carry) = self.builder.ins().smul_overflow(base_phi, base_phi); let new_exp = self.builder.ins().sshr_imm(exp_phi, 1); + let more = self.builder.ins().icmp_imm(IntCC::NotEqual, new_exp, 0); + let square_overflows = self.builder.ins().band(more, square_carry); + self.deopt_branch(square_overflows, operands)?; + self.builder.ins().jump( loop_block, &[new_exp.into(), new_result.into(), squared_base.into()], @@ -1496,12 +1508,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let res = self.builder.block_params(exit_block)[0]; // Seal all blocks - self.builder.seal_block(check_negative); - self.builder.seal_block(handle_negative); self.builder.seal_block(loop_block); self.builder.seal_block(continue_block); self.builder.seal_block(exit_block); - res + Ok(res) } } diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 0c700e93cf8..f5b35947608 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -2,13 +2,45 @@ mod instructions; extern crate alloc; +use alloc::collections::BTreeSet; use alloc::fmt; -use core::mem::ManuallyDrop; +use alloc::sync::Arc; +use core::mem::{self, ManuallyDrop}; +use core::sync::atomic::{AtomicU8, AtomicU64, Ordering}; use cranelift::prelude::*; use cranelift_jit::{JITBuilder, JITModule}; use cranelift_module::{FuncId, Linkage, Module, ModuleError}; -use instructions::FunctionCompiler; +use instructions::{Externals, FunctionCompiler}; use rustpython_compiler_core::bytecode; +use std::sync::{Mutex, PoisonError}; + +/// Arguments cross into compiled code through a flat buffer of 64-bit slots: +/// an int sign-extended, a bool as 0 or 1, a float as its bits. The buffer is +/// a fixed-size array so that a call allocates nothing, which caps how many +/// parameters a function can have and still be compiled. +const MAX_ARGS: usize = 16; +const SLOT_SIZE: usize = size_of::(); + +/// A guard that fires leaves its record in a second flat buffer: +/// +/// ```text +/// deopt[0] status: 0 when the call returned, `DEOPT_STATUS_NESTED` when it +/// left with no record of its own, otherwise the site index plus one +/// deopt[1] bound mask: bit i set when varname slot i holds a bound local +/// deopt[2..] the listed locals, then the value stack bottom to top +/// ``` +/// +/// Like the argument buffer it is a fixed-size array so that a call allocates +/// nothing, which caps how much state a guard can spill. +const MAX_DEOPT_SLOTS: usize = 64; +/// Slots taken by the status and the bound mask, before the record starts. +const DEOPT_HEADER_SLOTS: usize = 2; +/// Status for a frame that left because a nested frame gave up. Its record +/// belongs to that frame, so this one has nothing to resume from. +const DEOPT_STATUS_NESTED: u64 = u64::MAX; + +/// The entry point of a compiled function: `(args, ret, deopt)`. +type JitEntry = unsafe extern "C" fn(*const u64, *mut u64, *mut u64); #[derive(Debug, thiserror::Error)] #[non_exhaustive] @@ -27,6 +59,43 @@ impl From for JitCompileError { } } +/// Whether a self-call - a call matched to the function being compiled by +/// global name - may become a direct recursive call, or the whole function +/// is left to the interpreter instead. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Safety { + /// Reject the function outright rather than compile around the + /// self-call. The interpreter re-reads the global on every call, so a + /// rebound name - a decorator applied later, a test patching the module + /// - would make a compiled direct call disagree with it. + Strict, + /// Compile a self-call into a direct call. The callee is resolved once, + /// when the code is compiled, and a frame resumed from a guard keeps that + /// resolution - so rebinding the name is not observed by a call the + /// compiled code had already staged, even though the resumed frame runs + /// after the compiled code has left. Rebinding between two whole calls is + /// observed as usual, because the code is discarded when a guard fires. + /// + /// Making a site non-resumable whenever its stack holds a callee would + /// restore the older behaviour of re-reading the global, by restarting + /// instead of resuming. That is the worse trade: a restart re-runs the + /// function from the top, including self-calls that already returned, and + /// those run arbitrary Python and can have side effects. A stale callee + /// costs one extra invocation of the function that was there at compile + /// time; a restart costs every side effect the completed calls had. + Permissive, +} + +/// What compiled `**` on two floats calls, so the compiled answer comes from +/// the same function call the interpreter's `float_pow` makes rather than a +/// hand-rolled approximation of it. Registered on the builder by name below, +/// rather than left for the JIT to resolve `pow` through the platform's +/// libm - an explicit symbol is guaranteed to resolve, and guaranteed to be +/// this exact function. +extern "C" fn jit_powf(a: f64, b: f64) -> f64 { + a.powf(b) +} + #[derive(Debug, thiserror::Error, Eq, PartialEq)] #[non_exhaustive] pub enum JitArgumentError { @@ -39,27 +108,70 @@ pub enum JitArgumentError { struct Jit { builder_context: FunctionBuilderContext, ctx: codegen::Context, - module: JITModule, + /// `jit_powf`, declared once so every compiled function imports the same + /// symbol rather than redeclaring it. + powf_func: FuncId, + module: ManuallyDrop, } impl Jit { fn new() -> Self { - let builder = JITBuilder::new(cranelift_module::default_libcall_names()) + let mut builder = JITBuilder::new(cranelift_module::default_libcall_names()) .expect("Failed to build JITBuilder"); - let module = JITModule::new(builder); + builder.symbol("jit_powf", jit_powf as *const u8); + let mut module = JITModule::new(builder); + let mut powf_sig = module.make_signature(); + powf_sig.params.push(AbiParam::new(types::F64)); + powf_sig.params.push(AbiParam::new(types::F64)); + powf_sig.returns.push(AbiParam::new(types::F64)); + let powf_func = module + .declare_function("jit_powf", Linkage::Import, &powf_sig) + .expect("failed to declare jit_powf"); Self { builder_context: FunctionBuilderContext::new(), ctx: module.make_context(), - module, + powf_func, + module: ManuallyDrop::new(module), } } + /// Build one function into the module. The context is reset even when + /// compilation fails, so a rejected function leaves nothing behind for the + /// next one to trip over. fn build_function( &mut self, bytecode: &bytecode::CodeObject, args: &[JitType], ret: Option, - ) -> Result<(FuncId, JitSig), JitCompileError> { + unique: u64, + safety: Safety, + safepoint: Option<&'static AtomicU8>, + ) -> Result<(FuncId, JitSig, Vec), JitCompileError> { + let result = self.build_function_inner(bytecode, args, ret, unique, safety, safepoint); + self.module.clear_context(&mut self.ctx); + if result.is_err() { + // Only `FunctionBuilder::finalize` resets the builder context, and + // a rejected function never reaches it. Leaving it dirty makes the + // next `FunctionBuilder::new` panic. + self.builder_context = FunctionBuilderContext::new(); + } + result + } + + fn build_function_inner( + &mut self, + bytecode: &bytecode::CodeObject, + args: &[JitType], + ret: Option, + unique: u64, + safety: Safety, + safepoint: Option<&'static AtomicU8>, + ) -> Result<(FuncId, JitSig, Vec), JitCompileError> { + let ptr_type = self.module.target_config().pointer_type(); + // The deopt buffer comes first so that a guard can reach it without + // depending on how many parameters the function has. + self.ctx.func.signature.params.push(AbiParam::new(ptr_type)); + for arg in args { let arg = arg.to_cranelift().ok_or(JitCompileError::NotSupported)?; self.ctx.func.signature.params.push(AbiParam::new(arg)); @@ -70,119 +182,501 @@ impl Jit { } let id = self.module.declare_function( - &format!("jit_{}", bytecode.obj_name.as_ref()), + &format!("jit_{}_{unique}", bytecode.obj_name.as_ref()), Linkage::Export, &self.ctx.func.signature, )?; let func_ref = self.module.declare_func_in_func(id, &mut self.ctx.func); + let powf_func = self + .module + .declare_func_in_func(self.powf_func, &mut self.ctx.func); let mut builder = FunctionBuilder::new(&mut self.ctx.func, &mut self.builder_context); let entry_block = builder.create_block(); builder.append_block_params_for_function_params(entry_block); builder.switch_to_block(entry_block); - let sig = { + let (sig, deopt_sites) = { let mut compiler = FunctionCompiler::new( &mut builder, bytecode.varnames.len(), args, ret, entry_block, + safety, + Externals { + powf: powf_func, + safepoint, + }, ); compiler.compile(func_ref, bytecode)?; - compiler.sig + (compiler.sig, compiler.deopt_sites) }; builder.seal_all_blocks(); builder.finalize(); + // Compiling the body can widen the signature: a return type that was + // not annotated is only learned from the return statements. + let body_signature = self.ctx.func.signature.clone(); self.module.define_function(id, &mut self.ctx)?; - self.module.clear_context(&mut self.ctx); - Ok((id, sig)) + let entry = self.build_entry(id, body_signature, &sig, unique)?; + + Ok((entry, sig, deopt_sites)) + } + + /// Build the entry point callers go through: it takes a flat buffer of + /// 64-bit slots, unpacks them into the parameters the compiled body + /// actually takes, and writes the result back into the caller's slot. + /// + /// This is what makes the call a plain indirect call. Handing the same + /// arguments to a foreign-function library instead means describing the + /// signature and boxing every argument on every call. + fn build_entry( + &mut self, + target: FuncId, + target_signature: Signature, + sig: &JitSig, + unique: u64, + ) -> Result { + let ptr_type = self.module.target_config().pointer_type(); + // (args, ret, deopt) + self.ctx.func.signature.params.push(AbiParam::new(ptr_type)); + self.ctx.func.signature.params.push(AbiParam::new(ptr_type)); + self.ctx.func.signature.params.push(AbiParam::new(ptr_type)); + + let id = self.module.declare_function( + &format!("jit_entry_{unique}"), + Linkage::Export, + &self.ctx.func.signature, + )?; + + let callee = self.module.declare_func_in_func(target, &mut self.ctx.func); + // The import carries the signature the target was declared with, which + // is the one from before the body widened it. + let callee_signature = self.ctx.func.import_signature(target_signature); + self.ctx.func.dfg.ext_funcs[callee].signature = callee_signature; + + let mut builder = FunctionBuilder::new(&mut self.ctx.func, &mut self.builder_context); + let block = builder.create_block(); + builder.append_block_params_for_function_params(block); + builder.switch_to_block(block); + let args_ptr = builder.block_params(block)[0]; + let ret_ptr = builder.block_params(block)[1]; + let deopt_ptr = builder.block_params(block)[2]; + // Written before anything can read it, so the buffer never has to be zeroed. + let zero = builder.ins().iconst(types::I64, 0); + builder.ins().store(MemFlags::trusted(), zero, deopt_ptr, 0); + + let mut call_args = vec![deopt_ptr]; + for (i, ty) in sig.args.iter().enumerate() { + let offset = i32::try_from(i * SLOT_SIZE).map_err(|_| JitCompileError::NotSupported)?; + let mut load = |ty| { + builder + .ins() + .load(ty, MemFlags::trusted(), args_ptr, offset) + }; + call_args.push(match ty { + JitType::Int => load(types::I64), + JitType::Float => load(types::F64), + // A slot holds 0 or 1, so the low byte carries the whole value + // whichever end of the slot it sits at. + JitType::Bool => { + let slot = load(types::I64); + builder.ins().ireduce(types::I8, slot) + } + JitType::None => return Err(JitCompileError::NotSupported), + }); + } + + let call = builder.ins().call(callee, &call_args); + let returned = match sig.ret.as_ref().filter(|ty| ty.to_cranelift().is_some()) { + Some(ty) => match *builder.inst_results(call) { + [result] => Some((ty, result)), + _ => return Err(JitCompileError::NotSupported), + }, + None => None, + }; + if let Some((ty, result)) = returned { + let result = if *ty == JitType::Bool { + builder.ins().uextend(types::I64, result) + } else { + result + }; + builder.ins().store(MemFlags::trusted(), result, ret_ptr, 0); + } + builder.ins().return_(&[]); + + builder.seal_all_blocks(); + builder.finalize(); + + self.module.define_function(id, &mut self.ctx)?; + + Ok(id) } } +/// Owns the code memory of every function compiled through it. A +/// [`CompiledCode`] keeps its engine alive, so machine code is never freed +/// while something can still call it. +pub struct JitEngine { + jit: Mutex, + next_id: AtomicU64, + safepoint: Option<&'static AtomicU8>, +} + +impl JitEngine { + /// `safepoint` is the byte every compiled backward jump polls: non-zero + /// means the interpreter wants the running thread out of compiled code, be + /// it for a pending signal, a stop-the-world or its own shutdown. `None` + /// compiles no poll, for a caller with no interpreter behind the code; a + /// loop compiled that way runs to its end whatever happens meanwhile. + #[must_use] + pub fn new(safepoint: Option<&'static AtomicU8>) -> Arc { + Arc::new(Self { + jit: Mutex::new(Jit::new()), + next_id: AtomicU64::new(0), + safepoint, + }) + } + + pub fn compile( + self: &Arc, + bytecode: &bytecode::CodeObject, + args: &[JitType], + ret: Option, + safety: Safety, + ) -> Result { + if args.len() > MAX_ARGS { + return Err(JitCompileError::NotSupported); + } + // Symbol names must be unique within the module, and `obj_name` is not: + // any two `def f` in different scopes collide. + let unique = self.next_id.fetch_add(1, Ordering::Relaxed); + let mut jit = self.jit.lock().unwrap_or_else(PoisonError::into_inner); + let (id, sig, deopt_sites) = + jit.build_function(bytecode, args, ret, unique, safety, self.safepoint)?; + jit.module.finalize_definitions()?; + let code = jit.module.get_finalized_function(id); + drop(jit); + // SAFETY: `build_entry` defined this function with exactly this + // signature, and the engine keeps its code alive. + let entry = unsafe { mem::transmute::<*const u8, JitEntry>(code) }; + Ok(CompiledCode { + sig, + deopt_sites, + entry, + _engine: self.clone(), + }) + } +} + +impl Drop for JitEngine { + fn drop(&mut self) { + let jit = self.jit.get_mut().unwrap_or_else(PoisonError::into_inner); + // SAFETY: every CompiledCode holds an Arc to this engine, so no + // compiled function is reachable any more once we get here. + unsafe { ManuallyDrop::take(&mut jit.module).free_memory() } + } +} + +// The module is only ever touched under the mutex, and the code pointers it +// hands out stay valid for as long as the engine lives. +unsafe impl Send for JitEngine {} +unsafe impl Sync for JitEngine {} + +/// Whether the backend could plausibly compile this code object. +/// +/// A cheap pre-filter for callers that compile speculatively: one pass over the +/// bytecode, ruling out the shapes there is no lowering for at all. Passing is +/// not a promise that compilation will succeed - the argument types decide much +/// of that, and they are not visible here. +pub fn supports_code(code: &bytecode::CodeObject) -> bool { + // A frame is never built, so there is nowhere to put varargs, a generator's + // suspended state, or an exception handler's stack. + if code.flags.intersects( + bytecode::CodeFlags::VARARGS + | bytecode::CodeFlags::VARKEYWORDS + | bytecode::CodeFlags::GENERATOR + | bytecode::CodeFlags::COROUTINE + | bytecode::CodeFlags::ASYNC_GENERATOR, + ) { + return false; + } + if !code.exceptiontable.is_empty() { + return false; + } + // Cells and frees live past `varnames`, which is all the compiler allocates + // locals for, and are read through opcodes it has no lowering for. + if !code.cellvars.is_empty() || !code.freevars.is_empty() { + return false; + } + + // The compiler refuses a merge it cannot reconcile, which is every merge + // reached with a non-empty stack. Simulating the depth here keeps the + // automatic path from entering the backend only to be turned down: a + // conditional expression merges mid-expression, where a statement-level + // `if` or `while` merges with nothing on the stack. + // + // This has to walk what the compiler walks - the de-specialized stream, and + // branch targets resolved by the same function it resolves them with. A + // jump argument is a delta, and `CodeObject::label_targets` collects the + // delta rather than the offset it points at, so it cannot answer this. + let Ok(clean) = bytecode::CodeUnits::try_from(code.instructions.original_bytes().as_slice()) + else { + return false; + }; + + let mut targets = BTreeSet::new(); + let mut target_state = bytecode::OpArgState::default(); + for (offset, &word) in clean.iter().enumerate() { + let (instruction, arg) = target_state.get(word); + match instructions::instruction_target(offset as u32, instruction, arg) { + Ok(Some(target)) => { + targets.insert(target); + } + Ok(None) => {} + Err(_) => return false, + } + } + + let mut state = bytecode::OpArgState::default(); + let mut depth: i32 = 0; + for (offset, &word) in clean.iter().enumerate() { + let (instruction, arg) = state.get(word); + if !instructions::instruction_is_supported(instruction) { + return false; + } + // Falling into a branch target is an edge into it, carrying whatever + // is on the stack on the way in. + if depth != 0 && targets.contains(&bytecode::Label::from_u32(offset as u32)) { + return false; + } + + depth += instruction.stack_effect(u32::from(arg)); + // Losing track of the depth means the rest of the walk proves nothing. + if depth < 0 { + return false; + } + + // A jump's own edge carries the stack it leaves behind, which is after + // its effect - a conditional jump has popped the condition it tested by + // the time control leaves, and that is the depth the compiler checks. + let leaves_here = matches!( + instructions::instruction_target(offset as u32, instruction, arg), + Ok(Some(_)) + ); + if depth != 0 && leaves_here { + return false; + } + } + true +} + pub fn compile( bytecode: &bytecode::CodeObject, args: &[JitType], ret: Option, ) -> Result { - let mut jit = Jit::new(); - - let (id, sig) = jit.build_function(bytecode, args, ret)?; - - jit.module.finalize_definitions()?; - - let code = jit.module.get_finalized_function(id); - Ok(CompiledCode { - sig, - code, - module: ManuallyDrop::new(jit.module), - }) + JitEngine::new(None).compile(bytecode, args, ret, Safety::Permissive) } pub struct CompiledCode { sig: JitSig, - code: *const u8, - module: ManuallyDrop, + /// Indexed by the status a guard writes, less one. + deopt_sites: Vec, + entry: JitEntry, + /// Keeps the code memory alive; never read. + _engine: Arc, } impl CompiledCode { + #[must_use] pub fn args_builder(&self) -> ArgsBuilder<'_> { ArgsBuilder::new(self) } - pub fn invoke(&self, args: &[AbiValue]) -> Result, JitArgumentError> { + pub fn invoke(&self, args: &[AbiValue]) -> Result { if self.sig.args.len() != args.len() { return Err(JitArgumentError::WrongNumberOfArguments); } - let cif_args = self - .sig - .args + let mut slots = [0; MAX_ARGS]; + for ((slot, ty), value) in slots.iter_mut().zip(&self.sig.args).zip(args) { + type_check(ty, value)?; + *slot = value.to_slot(); + } + // SAFETY: the arity was checked above, and every slot was written from + // a value `type_check` matched against that parameter's type. + Ok(unsafe { self.invoke_raw(&slots) }) + } + + /// # Safety + /// `slots` must hold a value of the right type for each parameter. + unsafe fn invoke_raw(&self, slots: &[u64; MAX_ARGS]) -> Outcome { + let mut ret = 0; + // Only slot 0 is written before it is read, by the entry point itself; + // zeroing the rest would cost more per call than the code being called. + let mut deopt = core::mem::MaybeUninit::<[u64; MAX_DEOPT_SLOTS]>::uninit(); + let deopt_ptr = deopt.as_mut_ptr().cast::(); + // SAFETY: the entry point reads one slot per parameter, writes the return + // slot only when the signature says it returns something, and writes the + // deopt status before returning. + unsafe { (self.entry)(slots.as_ptr(), &raw mut ret, deopt_ptr) } + // SAFETY: the entry point stores the status first thing. + let status = unsafe { deopt_ptr.read() }; + if status != 0 { + // A status that is not the sentinel is the index of the site + // describing the record its guard just wrote into this buffer. + let site = + (status != DEOPT_STATUS_NESTED).then(|| &self.deopt_sites[status as usize - 1]); + let state = site.filter(|site| site.resumable).map(|site| { + // SAFETY: this is the site whose guard wrote the record, and + // the buffer it wrote is still in scope. + unsafe { Self::read_deopt(site, deopt_ptr) } + }); + return match state { + _ if site.is_some_and(|site| site.safepoint) => Outcome::Interrupted(state), + Some(state) => Outcome::Deopt(state), + None => Outcome::Restart, + }; + } + Outcome::Returned(match self.sig.ret.as_ref() { + Some(JitType::None) | None => None, + Some(ty) => Some(AbiValue::from_slot(ty, ret)), + }) + } + + /// # Safety + /// `deopt` must point at the buffer the guard for `site` wrote. + unsafe fn read_deopt(site: &DeoptSite, deopt: *const u64) -> DeoptState { + // SAFETY: the guard wrote the mask and then one slot per listed local + // and stack entry, in this order, and the site is the one it wrote for. + let read = |slot: usize| unsafe { deopt.add(slot).read() }; + let mask = read(1); + let mut slot = DEOPT_HEADER_SLOTS; + let mut locals = Vec::with_capacity(site.locals.len()); + for (i, ty) in site.locals.iter().enumerate() { + // A listed local takes a slot whether or not it is bound, so an + // unbound one still has to be stepped over. + let value = ty.as_ref().map(|ty| { + let value = (mask & (1 << i) != 0).then(|| AbiValue::from_slot(ty, read(slot))); + slot += 1; + value + }); + locals.push(value.flatten()); + } + let stack = site + .stack .iter() - .zip(args.iter()) - .map(|(ty, val)| type_check(ty, val).map(|_| val)) - .map(|v| v.map(AbiValue::to_libffi_arg)) - .collect::, _>>()?; - Ok(unsafe { self.invoke_raw(&cif_args) }) - } - - unsafe fn invoke_raw(&self, cif_args: &[libffi::middle::Arg<'_>]) -> Option { - unsafe { - let cif = self.sig.to_cif(); - let value = cif.call::( - libffi::middle::CodePtr::from_ptr(self.code as *const _), - cif_args, - ); - match self.sig.ret.as_ref() { - Some(JitType::None) | None => None, - Some(ty) => Some(value.to_typed(ty)), - } + .map(|entry| match entry { + StackEntry::Value(ty) => { + let value = StackValue::Value(AbiValue::from_slot(ty, read(slot))); + slot += 1; + value + } + StackEntry::Callee => StackValue::Callee, + StackEntry::Null => StackValue::Null, + }) + .collect(); + DeoptState { + offset: site.offset, + locals, + stack, } } } +/// What a call to compiled code did. +#[derive(Debug, PartialEq)] +pub enum Outcome { + Returned(Option), + /// A guard fired. The record is decoded here, off the hot path, so that + /// nothing borrows the buffer once it goes out of scope. + Deopt(DeoptState), + /// A guard fired, but left nothing this call can carry on from, so it has + /// to be run again from the start. Either the record in the buffer belongs + /// to a nested frame, or it belongs to a site that cannot describe every + /// local that can be bound where it fires. A poll a nested frame made + /// arrives here too: its status is overwritten by the frame it returned + /// to, which leaves nothing to tell the two apart. + Restart, + /// A backward jump polled and found the thread had been asked to leave the + /// bytecode loop. The code is still right for the values it was given, so + /// it stays installed; only this call finishes interpreted, from the + /// record where there is one and from the start where there is not. + Interrupted(Option), +} + +/// Everything the interpreter needs to pick up where the guard stopped. +#[derive(Debug, PartialEq)] +pub struct DeoptState { + /// Bytecode offset to resume at. + pub offset: u32, + /// One entry per varname slot; `None` where the local is not live or is + /// unbound. + pub locals: Vec>, + /// Bottom to top. + pub stack: Vec, +} + +/// One value-stack entry handed back to the interpreter. +#[derive(Debug, Clone, PartialEq)] +pub enum StackValue { + Value(AbiValue), + /// The function the compiled code belongs to. + Callee, + Null, +} + +/// What one guard spills, and how to read it back. +/// +/// The guard itself only writes values; everything about their shape is fixed +/// at compile time and kept here, so the record in the buffer needs no tags. +pub(crate) struct DeoptSite { + /// Bytecode offset of the instruction the guard belongs to. + pub(crate) offset: u32, + /// One entry per varname slot, in order; `None` where no local lives in + /// that slot. A listed local occupies a slot in the record even when the + /// bound mask says it is unassigned. + pub(crate) locals: Box<[Option]>, + /// Bottom to top, after the listed locals. + pub(crate) stack: Box<[StackEntry]>, + /// Whether a frame can be rebuilt from this site's record. A site lists + /// the locals the compiler had seen where the guard was lowered, which a + /// backward jump can leave short of what is bound where it fires. A local + /// the record omits is missing from the resumed frame, where anything + /// reading its fastlocals other than a `LoadFast` can see it go: + /// `f_locals`, a tracer stepping the frame, a debugger stopped in it. + pub(crate) resumable: bool, + /// Whether this is the poll a backward jump makes rather than a guard. A + /// guard fires because the compiled code is wrong for the values it was + /// given; a poll fires because the thread was asked to stop, which says + /// nothing about the code. Only the first is a reason to throw it away. + pub(crate) safepoint: bool, +} + +/// What one value-stack slot holds at a deopt site. Only `Value` occupies a +/// slot in the buffer; the other two are the same on every path that reaches +/// the site, so the interpreter can rebuild them. +pub(crate) enum StackEntry { + Value(JitType), + /// The compiled function itself, pushed by the self-reference lookup. + Callee, + /// The null a call pushes beside its callable. + Null, +} + struct JitSig { args: Vec, ret: Option, } -impl JitSig { - fn to_cif(&self) -> libffi::middle::Cif { - let ret = match self.ret { - Some(ref ty) => ty.to_libffi(), - None => libffi::middle::Type::void(), - }; - libffi::middle::Cif::new(self.args.iter().map(JitType::to_libffi), ret) - } -} - #[derive(Debug, Clone, PartialEq, Eq)] #[non_exhaustive] pub enum JitType { @@ -201,15 +695,6 @@ impl JitType { Self::None => None, } } - - fn to_libffi(&self) -> libffi::middle::Type { - match self { - Self::Int => libffi::middle::Type::i64(), - Self::Float => libffi::middle::Type::f64(), - Self::Bool => libffi::middle::Type::u8(), - Self::None => libffi::middle::Type::void(), - } - } } #[derive(Debug, Clone, PartialEq)] @@ -221,11 +706,21 @@ pub enum AbiValue { } impl AbiValue { - fn to_libffi_arg(&self) -> libffi::middle::Arg<'_> { - match self { - Self::Int(i) => libffi::middle::Arg::new(i), - Self::Float(f) => libffi::middle::Arg::new(f), - Self::Bool(b) => libffi::middle::Arg::new(b), + /// Pack into the 64-bit slot the entry point reads. + fn to_slot(&self) -> u64 { + match *self { + Self::Int(i) => i as u64, + Self::Float(f) => f.to_bits(), + Self::Bool(b) => b.into(), + } + } + + fn from_slot(ty: &JitType, slot: u64) -> Self { + match ty { + JitType::Int => Self::Int(slot as i64), + JitType::Float => Self::Float(f64::from_bits(slot)), + JitType::Bool => Self::Bool(slot != 0), + JitType::None => unreachable!("None has no slot"), } } } @@ -290,39 +785,6 @@ fn type_check(ty: &JitType, val: &AbiValue) -> Result<(), JitArgumentError> { } } -#[derive(Copy, Clone)] -union UnTypedAbiValue { - float: f64, - int: i64, - boolean: u8, - _void: (), -} - -impl UnTypedAbiValue { - unsafe fn to_typed(self, ty: &JitType) -> AbiValue { - unsafe { - match ty { - JitType::Int => AbiValue::Int(self.int), - JitType::Float => AbiValue::Float(self.float), - JitType::Bool => AbiValue::Bool(self.boolean != 0), - JitType::None => unreachable!("None has no ABI value"), - } - } - } -} - -// we don't actually ever touch CompiledCode til we drop it, it should be safe. -// TODO: confirm with wasmtime ppl that it's not unsound? -unsafe impl Send for CompiledCode {} -unsafe impl Sync for CompiledCode {} - -impl Drop for CompiledCode { - fn drop(&mut self) { - // SAFETY: The only pointer that this memory will also be dropped now - unsafe { ManuallyDrop::take(&mut self.module).free_memory() } - } -} - impl fmt::Debug for CompiledCode { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str("[compiled code]") @@ -330,7 +792,10 @@ impl fmt::Debug for CompiledCode { } pub struct ArgsBuilder<'a> { - values: Vec>, + slots: [u64; MAX_ARGS], + /// One bit per filled slot, so a caller can tell an argument it has + /// already placed from one still to come. + filled: u32, code: &'a CompiledCode, } @@ -338,44 +803,44 @@ impl<'a> ArgsBuilder<'a> { #[must_use] fn new(code: &'a CompiledCode) -> Self { Self { - values: vec![None; code.sig.args.len()], + slots: [0; MAX_ARGS], + filled: 0, code, } } pub fn set(&mut self, idx: usize, value: AbiValue) -> Result<(), JitArgumentError> { - type_check(&self.code.sig.args[idx], &value).map(|_| { - self.values[idx] = Some(value); + type_check(&self.code.sig.args[idx], &value).map(|()| { + self.slots[idx] = value.to_slot(); + self.filled |= 1 << idx; }) } #[must_use] pub fn is_set(&self, idx: usize) -> bool { - self.values[idx].is_some() + self.filled & (1 << idx) != 0 } #[must_use] pub fn into_args(self) -> Option> { - // Ensure all values are set - if self.values.iter().any(|v| v.is_none()) { - return None; - } - Some(Args { - values: self.values.into_iter().map(|v| v.unwrap()).collect(), + let wanted = (1 << self.code.sig.args.len()) - 1; + (self.filled == wanted).then_some(Args { + slots: self.slots, code: self.code, }) } } pub struct Args<'a> { - values: Vec, + slots: [u64; MAX_ARGS], code: &'a CompiledCode, } impl Args<'_> { #[must_use] - pub fn invoke(&self) -> Option { - let cif_args: Vec<_> = self.values.iter().map(AbiValue::to_libffi_arg).collect(); - unsafe { self.code.invoke_raw(&cif_args) } + pub fn invoke(&self) -> Outcome { + // SAFETY: `into_args` only hands out `Args` once every parameter has a + // slot, and `set` type-checked each one against the signature. + unsafe { self.code.invoke_raw(&self.slots) } } } diff --git a/crates/jit/tests/common.rs b/crates/jit/tests/common.rs index fe0e4bd33d4..deac8b78370 100644 --- a/crates/jit/tests/common.rs +++ b/crates/jit/tests/common.rs @@ -1,8 +1,9 @@ +use alloc::sync::Arc; use core::ops::ControlFlow; use rustpython_compiler_core::bytecode::{ CodeObject, ConstantData, Constants, Instruction, OpArg, OpArgState, }; -use rustpython_jit::{CompiledCode, JitType}; +use rustpython_jit::{CompiledCode, JitCompileError, JitEngine, JitType, Safety}; use rustpython_wtf8::{Wtf8, Wtf8Buf}; use std::collections::HashMap; @@ -14,6 +15,27 @@ pub(crate) struct Function { impl Function { pub(crate) fn compile(self) -> CompiledCode { + let (arg_types, ret_type) = self.signature(); + rustpython_jit::compile(&self.code, &arg_types, ret_type).expect("Compile failure") + } + + /// Compile onto a caller-owned engine, surfacing the error instead of panicking. + #[allow(dead_code)] + pub(crate) fn compile_on( + &self, + engine: &Arc, + safety: Safety, + ) -> Result { + let (arg_types, ret_type) = self.signature(); + engine.compile(&self.code, &arg_types, ret_type, safety) + } + + #[allow(dead_code)] + pub(crate) fn code(&self) -> &CodeObject { + &self.code + } + + fn signature(&self) -> (Vec, Option) { let mut arg_types = Vec::new(); for arg in self.code.arg_names().args { let arg_type = match self.annotations.get(AsRef::::as_ref(arg.as_str())) { @@ -38,7 +60,7 @@ impl Function { _ => None, }; - rustpython_jit::compile(&self.code, &arg_types, ret_type).expect("Compile failure") + (arg_types, ret_type) } } @@ -309,7 +331,7 @@ impl StackMachine { } } -macro_rules! jit_function { +macro_rules! py_function_def { ($func_name:ident => $($t:tt)*) => { { let code = rustpython_derive::py_compile!( @@ -319,9 +341,15 @@ macro_rules! jit_function { let code = code.decode(rustpython_compiler_core::bytecode::BasicBag); let mut machine = $crate::common::StackMachine::new(); machine.run(code); - machine.get_function(stringify!($func_name)).compile() + machine.get_function(stringify!($func_name)) } }; +} + +macro_rules! jit_function { + ($func_name:ident => $($t:tt)*) => { + py_function_def!($func_name => $($t)*).compile() + }; ($func_name:ident($($arg_name:ident:$arg_type:ty),*) -> $ret_type:ty => $($t:tt)*) => { { let jit_code = jit_function!($func_name => $($t)*); @@ -329,9 +357,22 @@ macro_rules! jit_function { move |$($arg_name:$arg_type),*| -> Result<$ret_type, rustpython_jit::JitArgumentError> { jit_code .invoke(&[$($arg_name.into()),*]) - .map(|ret| match ret { - Some(ret) => ret.try_into().expect("jit function returned unexpected type"), - None => panic!("jit function unexpectedly returned None") + .map(|outcome| match outcome { + rustpython_jit::Outcome::Returned(Some(ret)) => { + ret.try_into().expect("jit function returned unexpected type") + } + rustpython_jit::Outcome::Returned(None) => { + panic!("jit function unexpectedly returned None") + } + rustpython_jit::Outcome::Deopt(state) => { + panic!("jit function unexpectedly deoptimized: {state:?}") + } + rustpython_jit::Outcome::Restart => { + panic!("jit function unexpectedly asked to be restarted") + } + rustpython_jit::Outcome::Interrupted(state) => { + panic!("jit function unexpectedly interrupted: {state:?}") + } }) } } @@ -343,9 +384,20 @@ macro_rules! jit_function { move |$($arg_name:$arg_type),*| -> Result<(), rustpython_jit::JitArgumentError> { jit_code .invoke(&[$($arg_name.into()),*]) - .map(|ret| match ret { - Some(ret) => panic!("jit function unexpectedly returned a value {:?}", ret), - None => () + .map(|outcome| match outcome { + rustpython_jit::Outcome::Returned(None) => (), + rustpython_jit::Outcome::Returned(Some(ret)) => { + panic!("jit function unexpectedly returned a value {ret:?}") + } + rustpython_jit::Outcome::Deopt(state) => { + panic!("jit function unexpectedly deoptimized: {state:?}") + } + rustpython_jit::Outcome::Restart => { + panic!("jit function unexpectedly asked to be restarted") + } + rustpython_jit::Outcome::Interrupted(state) => { + panic!("jit function unexpectedly interrupted: {state:?}") + } }) } } diff --git a/crates/jit/tests/deopt_tests.rs b/crates/jit/tests/deopt_tests.rs new file mode 100644 index 00000000000..2cdcacb1dd2 --- /dev/null +++ b/crates/jit/tests/deopt_tests.rs @@ -0,0 +1,607 @@ +#[cfg(test)] +mod tests { + use rustpython_jit::{AbiValue, JitEngine, Outcome, Safety, StackValue}; + + fn int(value: i64) -> StackValue { + StackValue::Value(AbiValue::Int(value)) + } + + fn float(value: f64) -> StackValue { + StackValue::Value(AbiValue::Float(value)) + } + + /// Overflow is not an error; it is where the interpreter stops using a + /// machine word. The compiled code has to hand the operands back rather + /// than wrap or trap. + #[test] + fn addition_deopts_on_overflow() { + let code = jit_function! { add => r#" +def add(a: int, b: int) -> int: + return a + b +"# }; + assert_eq!( + code.invoke(&[1i64.into(), 2i64.into()]), + Ok(Outcome::Returned(Some(3i64.into()))) + ); + match code.invoke(&[i64::MAX.into(), 1i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(i64::MAX), int(1)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// Subtraction is the same guard as addition, mirrored: the operands go + /// back once the machine word can no longer hold the answer. + #[test] + fn subtraction_deopts_on_overflow() { + let code = jit_function! { sub => r#" +def sub(a: int, b: int) -> int: + return a - b +"# }; + assert_eq!( + code.invoke(&[5i64.into(), 3i64.into()]), + Ok(Outcome::Returned(Some(2i64.into()))) + ); + match code.invoke(&[i64::MIN.into(), 1i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(i64::MIN), int(1)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// Multiplication wraps just as readily as addition does, and needs the + /// same guard. + #[test] + fn multiplication_deopts_on_overflow() { + let code = jit_function! { mul => r#" +def mul(a: int, b: int) -> int: + return a * b +"# }; + assert_eq!( + code.invoke(&[3i64.into(), 4i64.into()]), + Ok(Outcome::Returned(Some(12i64.into()))) + ); + match code.invoke(&[i64::MAX.into(), 2i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(i64::MAX), int(2)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// Negation is lowered as `0 - a`, so it overflows exactly where + /// subtraction does. The zero is never on the interpreter's stack, so the + /// guard records only `a`. + #[test] + fn negation_deopts_on_overflow() { + let code = jit_function! { neg => r#" +def neg(a: int) -> int: + return -a +"# }; + assert_eq!( + code.invoke(&[5i64.into()]), + Ok(Outcome::Returned(Some((-5i64).into()))) + ); + match code.invoke(&[i64::MIN.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(i64::MIN)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// A shift count outside `0..64` is not a machine shift at all: negative + /// raises `ValueError`, and 64 or more is a well-defined answer the + /// instruction cannot give. A left shift that pushes bits off the top is + /// where the interpreter widens. + #[test] + fn left_shift_deopts_out_of_range_or_lossy() { + let code = jit_function! { shl => r#" +def shl(a: int, b: int) -> int: + return a << b +"# }; + assert_eq!( + code.invoke(&[1i64.into(), 2i64.into()]), + Ok(Outcome::Returned(Some(4i64.into()))) + ); + match code.invoke(&[1i64.into(), 64i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(1), int(64)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + match code.invoke(&[1i64.into(), (-1i64).into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(1), int(-1)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + match code.invoke(&[1i64.into(), 63i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(1), int(63)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// A negative count and a count of 64 or more are out of range for the + /// machine instruction the same way they are for a left shift. + #[test] + fn right_shift_deopts_out_of_range() { + let code = jit_function! { shr => r#" +def shr(a: int, b: int) -> int: + return a >> b +"# }; + assert_eq!( + code.invoke(&[8i64.into(), 1i64.into()]), + Ok(Outcome::Returned(Some(4i64.into()))) + ); + match code.invoke(&[8i64.into(), 64i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(8), int(64)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + match code.invoke(&[8i64.into(), (-1i64).into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(8), int(-1)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// A zero divisor has no machine answer: `sdiv`/`srem` trap where Python + /// raises ZeroDivisionError. `//` and `%` share the same guard. + #[test] + fn floor_divide_and_remainder_deopt_on_zero_divisor() { + let div = jit_function! { div => r#" +def div(a: int, b: int) -> int: + return a // b +"# }; + assert_eq!( + div.invoke(&[7i64.into(), 2i64.into()]), + Ok(Outcome::Returned(Some(3i64.into()))) + ); + match div.invoke(&[7i64.into(), 0i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(7), int(0)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + + let rem = jit_function! { rem => r#" +def rem(a: int, b: int) -> int: + return a % b +"# }; + match rem.invoke(&[7i64.into(), 0i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(7), int(0)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// `i64::MIN // -1`'s quotient is one past the top of the range - there + /// is no 64-bit answer, so it deoptimizes ahead of the `sdiv` that would + /// otherwise trap. The guard is shared with the remainder, so + /// `i64::MIN % -1` deopts too even though `0` is a perfectly good + /// answer for it. + #[test] + fn floor_divide_deopts_on_i64_min_over_negative_one() { + let code = jit_function! { div => r#" +def div(a: int, b: int) -> int: + return a // b +"# }; + match code.invoke(&[i64::MIN.into(), (-1i64).into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(i64::MIN), int(-1)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + + let rem = jit_function! { rem => r#" +def rem(a: int, b: int) -> int: + return a % b +"# }; + match rem.invoke(&[i64::MIN.into(), (-1i64).into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(i64::MIN), int(-1)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// A zero divisor raises ZeroDivisionError; there is no float to hand + /// back for it. + #[test] + fn true_divide_deopts_on_zero_divisor() { + let code = jit_function! { div => r#" +def div(a: int, b: int) -> float: + return a / b +"# }; + assert_eq!( + code.invoke(&[7i64.into(), 2i64.into()]), + Ok(Outcome::Returned(Some(3.5f64.into()))) + ); + match code.invoke(&[7i64.into(), 0i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(7), int(0)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// `int.__truediv__` is correctly rounded, but converting both operands + /// to `f64` and dividing rounds twice. That can be a ulp out as soon as + /// either conversion is inexact, which is exactly when an operand does + /// not fit a double's 53-bit significand. + #[test] + fn true_divide_deopts_on_wide_operand() { + let code = jit_function! { div => r#" +def div(a: int, b: int) -> float: + return a / b +"# }; + // `1 << 53` itself converts to a double exactly, so it stays compiled. + assert_eq!( + code.invoke(&[(1i64 << 53).into(), 1i64.into()]), + Ok(Outcome::Returned(Some(((1i64 << 53) as f64).into()))) + ); + match code.invoke(&[((1i64 << 53) + 1).into(), 1i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int((1i64 << 53) + 1), int(1)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + // The guard checks both operands - a wide divisor deopts as readily + // as a wide dividend. + match code.invoke(&[1i64.into(), ((1i64 << 53) + 1).into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(1), int((1i64 << 53) + 1)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + // `iabs` cannot negate `i64::MIN`; the unsigned comparison still + // places it past the bound. + match code.invoke(&[i64::MIN.into(), 1i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(i64::MIN), int(1)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// `iabs` of `i64::MIN` is `i64::MIN` again, which an unsigned comparison + /// reads as far above `1 << 53` - so it takes the wide-operand path too. + #[test] + fn true_divide_deopts_on_i64_min_operand() { + let code = jit_function! { div => r#" +def div(a: int, b: int) -> float: + return a / b +"# }; + match code.invoke(&[i64::MIN.into(), 1i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(i64::MIN), int(1)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// A negative exponent makes `**` a float, and the loop only computes + /// integers. + #[test] + fn power_deopts_on_negative_exponent() { + let code = jit_function! { pow => r#" +def pow(a: int, b: int) -> int: + return a ** b +"# }; + assert_eq!( + code.invoke(&[2i64.into(), 2i64.into()]), + Ok(Outcome::Returned(Some(4i64.into()))) + ); + match code.invoke(&[2i64.into(), (-2i64).into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(2), int(-2)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// `2 ** 64` does not fit an i64: the loop's final squaring carries into + /// the 65th bit even though every earlier iteration stayed in range. + #[test] + fn power_deopts_when_the_answer_does_not_fit() { + let code = jit_function! { pow => r#" +def pow(a: int, b: int) -> int: + return a ** b +"# }; + match code.invoke(&[2i64.into(), 64i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(2), int(64)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// A zero divisor raises ZeroDivisionError; `fdiv` would otherwise return + /// an infinity. `fcmp Equal` against `0.0` catches `-0.0` too, which + /// raises the same way. + #[test] + fn true_divide_deopts_on_float_zero_divisor() { + let code = jit_function! { div => r#" +def div(a: float, b: float) -> float: + return a / b +"# }; + assert_eq!( + code.invoke(&[4.0f64.into(), 2.0f64.into()]), + Ok(Outcome::Returned(Some(2.0f64.into()))) + ); + for divisor in [0.0f64, -0.0f64] { + match code.invoke(&[1.0f64.into(), divisor.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![float(1.0), float(divisor)]); + } + other => panic!("expected a deopt for {divisor}, got {other:?}"), + } + } + } + + /// The mixed int/float arm has its own `fdiv`, reached whenever exactly + /// one operand is already a float, and needs the same guard regardless of + /// which side that is - including a `-0.0` divisor, the same as the + /// float/float arm above. + #[test] + fn true_divide_deopts_on_mixed_zero_divisor() { + let int_over_float = jit_function! { div => r#" +def div(a: int, b: float) -> float: + return a / b +"# }; + for divisor in [0.0f64, -0.0f64] { + match int_over_float.invoke(&[1i64.into(), divisor.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(1), float(divisor)]); + } + other => panic!("expected a deopt for {divisor}, got {other:?}"), + } + } + + let float_over_int = jit_function! { div => r#" +def div(a: float, b: int) -> float: + return a / b +"# }; + match float_over_int.invoke(&[1.0f64.into(), 0i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![float(1.0), int(0)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// `0.0 ** negative` raises rather than returning an infinity. The + /// exponent is read by its value, so `-0.0` is not one of them - see + /// `basic_power` in float_tests.rs for the answer it gets instead. + #[test] + fn float_power_deopts_on_zero_base_negative_exponent() { + let code = jit_function! { pow => r#" +def pow(a: float, b: float) -> float: + return a ** b +"# }; + // A `-0.0` base is a zero base like any other. + for base in [0.0f64, -0.0f64] { + match code.invoke(&[base.into(), (-1.0f64).into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![float(base), float(-1.0)]); + } + other => panic!("expected a deopt for {base} ** -1.0, got {other:?}"), + } + } + } + + /// A negative base raised to a fractional power is complex. The base is + /// read by its value too, so `-0.0` is not one of those either. + #[test] + fn float_power_deopts_on_negative_base_fractional_exponent() { + let code = jit_function! { pow => r#" +def pow(a: float, b: float) -> float: + return a ** b +"# }; + match code.invoke(&[(-8.0f64).into(), 0.5f64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![float(-8.0), float(0.5)]); + } + other => panic!("expected a deopt for -8.0 ** 0.5, got {other:?}"), + } + } + + /// A finite base and exponent whose true power overflows a double raises + /// OverflowError rather than saturating to an infinity - unlike + /// `inf ** 2.0`, which correctly keeps returning `inf` and must not + /// deopt here. `1e100 ** 1e50` used to kill the process outright: an old + /// double–double implementation rounded `b * ln|a|` to an i64 with a + /// `fcvt_to_sint` that trapped once the product left i64 range, and + /// cranelift traps have no handler. Calling `f64::powf` directly has no + /// such trap, so it deoptimizes here like every other overflow shape. + #[test] + fn float_power_deopts_on_finite_base_overflow() { + let code = jit_function! { pow => r#" +def pow(a: float, b: float) -> float: + return a ** b +"# }; + assert_eq!( + code.invoke(&[f64::INFINITY.into(), 2.0f64.into()]), + Ok(Outcome::Returned(Some(f64::INFINITY.into()))) + ); + for (a, b) in [ + (1e308f64, 2.0f64), + (2.0f64, 1e300f64), + (-2.0f64, 1e300f64), + (1e100f64, 1e50f64), + (2.0f64, 1024.0f64), + (1e-308f64, -2.0f64), + (1e-308f64, -320.0f64), + (1e-100f64, -320.0f64), + (1e100f64, 4.0f64), + ] { + match code.invoke(&[a.into(), b.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![float(a), float(b)], "{a} ** {b}"); + } + other => panic!("expected a deopt for {a} ** {b}, got {other:?}"), + } + } + } + + /// A guard reports every live local, and reports a local that has not been + /// assigned on this path as unbound rather than inventing a value for it. + #[test] + fn a_guard_reports_the_live_state() { + let code = jit_function! { f => r#" +def f(a: int, b: int, c: bool) -> int: + if c: + d = 5 + return a + b +"# }; + match code.invoke(&[i64::MAX.into(), 1i64.into(), false.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.locals[0], Some(AbiValue::Int(i64::MAX))); + assert_eq!(state.locals[1], Some(AbiValue::Int(1))); + assert_eq!(state.locals[2], Some(AbiValue::Bool(false))); + // `d` was declared by the store inside the branch, but that + // branch did not run. + assert_eq!(state.locals[3], None); + assert_eq!(state.stack, vec![int(i64::MAX), int(1)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + // The same function, on the path that does assign `d`. + match code.invoke(&[i64::MAX.into(), 1i64.into(), true.into()]) { + Ok(Outcome::Deopt(state)) => assert_eq!(state.locals[3], Some(AbiValue::Int(5))), + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// A self-call leaves its callable and a null underneath the arguments + /// being evaluated. Neither has a slot in the buffer, but both are the same + /// on every path that reaches the guard, so the site describes them and the + /// function stays compilable. + #[test] + fn a_guard_under_a_self_call_describes_the_callable() { + let engine = JitEngine::new(None); + let f = py_function_def! { countdown => r#" +def countdown(a: int, b: int) -> int: + if a < 0: + return b + return countdown(a + b, b) +"# }; + let code = f + .compile_on(&engine, Safety::Permissive) + .expect("a guard below a self-call must not stop the function compiling"); + match code.invoke(&[i64::MAX.into(), 1i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!( + state.stack, + vec![StackValue::Callee, StackValue::Null, int(i64::MAX), int(1),] + ); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// A nested frame that gives up returns a filler in place of a result. The + /// caller has to stop rather than compute on it. + /// + /// Which of the two answers comes back depends on whose guard fired. + /// `blow(1)` overflows on its own addition, so its record describes the + /// frame that is asking for it. `blow(2)` reaches that overflow one frame + /// down, and the record standing in the buffer belongs to that frame - the + /// two frames run the same code, so every type in it lines up and a resume + /// would silently continue the outer frame from the inner frame's offset. + /// It has to come back as a restart instead. + #[test] + fn a_caller_stops_when_a_nested_frame_gives_up() { + let engine = JitEngine::new(None); + let f = py_function_def! { blow => r#" +def blow(n: int) -> int: + if n == 0: + return 4611686018427387904 + return blow(n - 1) + blow(n - 1) +"# }; + let code = f + .compile_on(&engine, Safety::Permissive) + .expect("should compile"); + assert_eq!( + code.invoke(&[0i64.into()]), + Ok(Outcome::Returned(Some(4611686018427387904i64.into()))) + ); + match code.invoke(&[1i64.into()]) { + Ok(Outcome::Deopt(state)) => assert_eq!( + state.stack, + vec![int(4611686018427387904), int(4611686018427387904)] + ), + other => panic!("expected a deopt, got {other:?}"), + } + assert_eq!(code.invoke(&[2i64.into()]), Ok(Outcome::Restart)); + } + + /// A guard lists the locals the compiler had seen where it was lowered, + /// which a backward jump can leave short: `extra` is stored further down + /// the loop body than the guard on the sum, yet it is bound by the time + /// that guard fires on a later iteration. Such a site cannot describe the + /// frame, so it asks for a restart rather than resuming without it. + /// + /// `extra` is never read, and could not be: a read the compiler cannot + /// prove bound is a `LoadFastCheck`, which has no lowering, so a function + /// that would observe the drop that way does not compile. What sees it is + /// anything reading the frame's fastlocals other than a `LoadFast` - + /// `f_locals`, a tracer, a debugger - which is why the snippet covering + /// this end to end goes through a traceback. + #[test] + fn a_site_that_cannot_describe_every_local_restarts() { + let code = jit_function! { late => r#" +def late(n: int, step: int) -> int: + total = 0 + while n > 0: + total = total + n * step + if n == 5: + extra = 1 + n = n - 1 + return total +"# }; + assert_eq!( + code.invoke(&[5i64.into(), 1i64.into()]), + Ok(Outcome::Returned(Some(15i64.into()))) + ); + assert_eq!( + code.invoke(&[5i64.into(), (1i64 << 60).into()]), + Ok(Outcome::Restart) + ); + } + + /// The same loop without the late store keeps its resume: every local it + /// can bind is established before the first guard is lowered, so no + /// backward jump can carry in one the sites do not list. + #[test] + fn a_loop_whose_locals_are_all_established_still_resumes() { + let code = jit_function! { mixed => r#" +def mixed(n: int, step: int) -> int: + total = 0 + while n > 0: + total = total + n * step + n = n - 1 + return total +"# }; + match code.invoke(&[5i64.into(), (1i64 << 60).into()]) { + Ok(Outcome::Deopt(state)) => { + // Two iterations in: `total` holds 5 << 60 and the multiply's + // 4 << 60 is on the stack under it, waiting for the addition + // that overflowed. + assert_eq!(state.locals[0], Some(AbiValue::Int(4))); + assert_eq!(state.locals[1], Some(AbiValue::Int(1 << 60))); + assert_eq!(state.locals[2], Some(AbiValue::Int(5 << 60))); + assert_eq!(state.stack, vec![int(5 << 60), int(4 << 60)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } +} diff --git a/crates/jit/tests/engine_tests.rs b/crates/jit/tests/engine_tests.rs new file mode 100644 index 00000000000..b4e9bad2626 --- /dev/null +++ b/crates/jit/tests/engine_tests.rs @@ -0,0 +1,126 @@ +#[cfg(test)] +mod tests { + use rustpython_jit::{JitCompileError, JitEngine, Outcome, Safety}; + + /// Two Python functions can share `obj_name`, so the module-level symbol has to + /// be made unique before they can live in one engine. + #[test] + fn same_name_functions_coexist() { + let engine = JitEngine::new(None); + let first = py_function_def!(foo => r#" + def foo(a: int, b: int) -> int: + return a + "#); + let second = py_function_def!(foo => r#" + def foo(a: int, b: int) -> int: + return b + "#); + + let first = first + .compile_on(&engine, Safety::Permissive) + .expect("first compile"); + let second = second + .compile_on(&engine, Safety::Permissive) + .expect("second compile of the same name"); + + assert_eq!( + first.invoke(&[3i64.into(), 4i64.into()]), + Ok(Outcome::Returned(Some(3i64.into()))) + ); + assert_eq!( + second.invoke(&[3i64.into(), 4i64.into()]), + Ok(Outcome::Returned(Some(4i64.into()))) + ); + } + + /// A rejected function must not leave half-built state in the shared context. + #[test] + fn failed_compile_does_not_poison_engine() { + let engine = JitEngine::new(None); + let unsupported = py_function_def!(unsupported => r#" + def unsupported(a: int) -> int: + return [a] + "#); + assert!(unsupported.compile_on(&engine, Safety::Permissive).is_err()); + + let good = py_function_def!(good => r#" + def good(a: int, b: int) -> int: + return a + b + "#); + let good = good + .compile_on(&engine, Safety::Permissive) + .expect("engine still usable after a rejected function"); + assert_eq!( + good.invoke(&[3i64.into(), 4i64.into()]), + Ok(Outcome::Returned(Some(7i64.into()))) + ); + } + + /// Compiled code outliving the caller's handle on the engine must still run. + #[test] + fn compiled_code_keeps_engine_alive() { + let code = { + let engine = JitEngine::new(None); + let f = py_function_def!(f => r#" + def f(a: int) -> int: + return a + "#); + f.compile_on(&engine, Safety::Permissive).expect("compile") + }; + assert_eq!( + code.invoke(&[7i64.into()]), + Ok(Outcome::Returned(Some(7i64.into()))) + ); + } + + /// Every parameter reaches the compiled body from its own slot, whatever + /// the types around it are. + #[test] + fn mixed_signature_round_trip() { + let pick_int = jit_function! { pick_int(a: i64, b: f64, c: bool, d: f64, e: i64) -> i64 => r#" + def pick_int(a: int, b: float, c: bool, d: float, e: int) -> int: + if c: + return a + return e + "# }; + assert_eq!(pick_int(1, 2.5, true, 4.5, 5), Ok(1)); + assert_eq!(pick_int(1, 2.5, false, 4.5, 5), Ok(5)); + + let pick_float = jit_function! { pick_float(a: i64, b: f64, c: bool, d: f64, e: i64) -> f64 => r#" + def pick_float(a: int, b: float, c: bool, d: float, e: int) -> float: + if c: + return b + return d + "# }; + assert_eq!(pick_float(1, 2.5, true, 4.5, 5), Ok(2.5)); + assert_eq!(pick_float(1, 2.5, false, 4.5, 5), Ok(4.5)); + } + + /// Arguments travel in a fixed-size buffer, so a function wider than the + /// buffer is turned down rather than compiled into an overrun. + #[test] + fn parameters_beyond_the_buffer_are_rejected() { + let engine = JitEngine::new(None); + let widest = py_function_def!(widest => r#" + def widest(a0: int, a1: int, a2: int, a3: int, a4: int, a5: int, a6: int, a7: int, a8: int, a9: int, a10: int, a11: int, a12: int, a13: int, a14: int, a15: int) -> int: + return a15 + "#); + let widest = widest + .compile_on(&engine, Safety::Permissive) + .expect("a function that fills the buffer still compiles"); + let args: Vec<_> = (0..16).map(|i| i64::from(i).into()).collect(); + assert_eq!( + widest.invoke(&args), + Ok(Outcome::Returned(Some(15i64.into()))) + ); + + let too_wide = py_function_def!(too_wide => r#" + def too_wide(a0: int, a1: int, a2: int, a3: int, a4: int, a5: int, a6: int, a7: int, a8: int, a9: int, a10: int, a11: int, a12: int, a13: int, a14: int, a15: int, a16: int) -> int: + return a16 + "#); + assert!(matches!( + too_wide.compile_on(&engine, Safety::Permissive), + Err(JitCompileError::NotSupported) + )); + } +} diff --git a/crates/jit/tests/float_tests.rs b/crates/jit/tests/float_tests.rs index f667b1e764a..9b688ad0809 100644 --- a/crates/jit/tests/float_tests.rs +++ b/crates/jit/tests/float_tests.rs @@ -118,109 +118,178 @@ mod tests { def pow(a:float, b: float): return a**b "##}; + // `**` calls `f64::powf` after its guards, the same function the + // interpreter's `float_pow` calls, so every case below is exact - + // there is no rounding step of this crate's own that a relative + // comparison would need to absorb. // Test base cases - assert_approx_eq!(pow(0.0, 0.0), Ok(1.0)); - assert_approx_eq!(pow(0.0, 1.0), Ok(0.0)); - assert_approx_eq!(pow(1.0, 0.0), Ok(1.0)); - assert_approx_eq!(pow(1.0, 1.0), Ok(1.0)); - assert_approx_eq!(pow(1.0, -1.0), Ok(1.0)); - assert_approx_eq!(pow(-1.0, 0.0), Ok(1.0)); - assert_approx_eq!(pow(-1.0, 1.0), Ok(-1.0)); - assert_approx_eq!(pow(-1.0, -1.0), Ok(-1.0)); - - // NaN and Infinity cases - assert_approx_eq!(pow(f64::NAN, 0.0), Ok(1.0)); - //assert_approx_eq!(pow(f64::NAN, 1.0), Ok(f64::NAN)); // Return the correct answer but fails compare - //assert_approx_eq!(pow(0.0, f64::NAN), Ok(f64::NAN)); // Return the correct answer but fails compare - assert_approx_eq!(pow(f64::INFINITY, 0.0), Ok(1.0)); - assert_approx_eq!(pow(f64::INFINITY, 1.0), Ok(f64::INFINITY)); - assert_approx_eq!(pow(f64::INFINITY, f64::INFINITY), Ok(f64::INFINITY)); + assert_bits_eq!(pow(0.0, 0.0), Ok(1.0f64)); + assert_bits_eq!(pow(0.0, 1.0), Ok(0.0f64)); + assert_bits_eq!(pow(1.0, 0.0), Ok(1.0f64)); + assert_bits_eq!(pow(1.0, 1.0), Ok(1.0f64)); + assert_bits_eq!(pow(1.0, -1.0), Ok(1.0f64)); + assert_bits_eq!(pow(-1.0, 0.0), Ok(1.0f64)); + assert_bits_eq!(pow(-1.0, 1.0), Ok(-1.0f64)); + assert_bits_eq!(pow(-1.0, -1.0), Ok(-1.0f64)); + + // NaN cases + assert_bits_eq!(pow(f64::NAN, 0.0), Ok(1.0f64)); + assert_bits_eq!(pow(f64::NAN, 2.0), Ok(f64::NAN)); + assert_bits_eq!(pow(0.0, f64::NAN), Ok(f64::NAN)); + assert_bits_eq!(pow(1.0, f64::NAN), Ok(1.0f64)); + + // An infinite exponent does not deoptimize by itself - only an + // overflowing finite-operand result does, see + // `float_power_deopts_on_finite_base_overflow` in deopt_tests.rs. + // `powf` answers these directly, matching the interpreter exactly + // because both call the same function. + assert_bits_eq!(pow(f64::INFINITY, f64::INFINITY), Ok(f64::INFINITY)); + assert_bits_eq!(pow(-1.0, f64::INFINITY), Ok(1.0f64)); + assert_bits_eq!(pow(-1.0, f64::NEG_INFINITY), Ok(1.0f64)); + assert_bits_eq!(pow(0.5, f64::INFINITY), Ok(0.0f64)); + assert_bits_eq!(pow(0.5, f64::NEG_INFINITY), Ok(f64::INFINITY)); + + // Infinity base cases: + assert_bits_eq!(pow(f64::INFINITY, 0.0), Ok(1.0f64)); + assert_bits_eq!(pow(f64::INFINITY, 1.0), Ok(f64::INFINITY)); + // An infinite base with a negative exponent correctly returns a + // signed zero rather than an infinity. + assert_bits_eq!(pow(f64::INFINITY, -2.0), Ok(0.0f64)); // Negative infinity cases: // For any exponent of 0.0, the result is 1.0. - assert_approx_eq!(pow(f64::NEG_INFINITY, 0.0), Ok(1.0)); + assert_bits_eq!(pow(f64::NEG_INFINITY, 0.0), Ok(1.0f64)); // For negative infinity base, when b is an odd integer, result is -infinity; // when b is even, result is +infinity. - assert_approx_eq!(pow(f64::NEG_INFINITY, 1.0), Ok(f64::NEG_INFINITY)); - assert_approx_eq!(pow(f64::NEG_INFINITY, 2.0), Ok(f64::INFINITY)); - assert_approx_eq!(pow(f64::NEG_INFINITY, 3.0), Ok(f64::NEG_INFINITY)); - // Exponent -infinity gives 0.0. - assert_approx_eq!(pow(f64::NEG_INFINITY, f64::NEG_INFINITY), Ok(0.0)); + assert_bits_eq!(pow(f64::NEG_INFINITY, 1.0), Ok(f64::NEG_INFINITY)); + assert_bits_eq!(pow(f64::NEG_INFINITY, 2.0), Ok(f64::INFINITY)); + assert_bits_eq!(pow(f64::NEG_INFINITY, 3.0), Ok(f64::NEG_INFINITY)); + // A negative odd exponent keeps the sign but flips the magnitude to + // a zero, the same as the positive-infinity-base case above. + assert_bits_eq!(pow(f64::NEG_INFINITY, -3.0), Ok(-0.0f64)); + // An infinite exponent is not special-cased for this base either. + assert_bits_eq!(pow(f64::NEG_INFINITY, f64::NEG_INFINITY), Ok(0.0f64)); + + // A negative zero base keeps its sign rather than being flattened to + // `+0.0`: `(-0.0) ** 3.0` is `-0.0`. + assert_bits_eq!(pow(-0.0, 3.0), Ok(-0.0f64)); + + // `-0.0` is a zero, not a negative number, on either side of the + // operator: it is neither the negative exponent that makes a zero + // base raise, nor the negative base that makes a fractional exponent + // complex. Both of those give up to the interpreter, so a compiled + // answer here is also the proof that neither guard fired. + assert_bits_eq!(pow(0.0, -0.0), Ok(1.0f64)); + assert_bits_eq!(pow(-0.0, -0.0), Ok(1.0f64)); + assert_bits_eq!(pow(-0.0, 0.5), Ok(0.0f64)); + // A nan exponent is not a fractional one, however negative the base. + assert_bits_eq!(pow(-2.0, f64::NAN), Ok(f64::NAN)); // Test positive float base, positive float exponent - assert_approx_eq!(pow(2.0, 2.0), Ok(4.0)); - assert_approx_eq!(pow(3.0, 3.0), Ok(27.0)); - assert_approx_eq!(pow(4.0, 4.0), Ok(256.0)); - assert_approx_eq!(pow(2.0, 3.0), Ok(8.0)); - assert_approx_eq!(pow(2.0, 4.0), Ok(16.0)); + assert_bits_eq!(pow(2.0, 2.0), Ok(4.0f64)); + assert_bits_eq!(pow(3.0, 3.0), Ok(27.0f64)); + assert_bits_eq!(pow(4.0, 4.0), Ok(256.0f64)); + assert_bits_eq!(pow(2.0, 3.0), Ok(8.0f64)); + assert_bits_eq!(pow(2.0, 4.0), Ok(16.0f64)); // Test negative float base, positive float exponent (integral exponents only) - assert_approx_eq!(pow(-2.0, 2.0), Ok(4.0)); - assert_approx_eq!(pow(-3.0, 3.0), Ok(-27.0)); - assert_approx_eq!(pow(-4.0, 4.0), Ok(256.0)); - assert_approx_eq!(pow(-2.0, 3.0), Ok(-8.0)); - assert_approx_eq!(pow(-2.0, 4.0), Ok(16.0)); + assert_bits_eq!(pow(-2.0, 2.0), Ok(4.0f64)); + assert_bits_eq!(pow(-3.0, 3.0), Ok(-27.0f64)); + assert_bits_eq!(pow(-4.0, 4.0), Ok(256.0f64)); + assert_bits_eq!(pow(-2.0, 3.0), Ok(-8.0f64)); + assert_bits_eq!(pow(-2.0, 4.0), Ok(16.0f64)); + // A negative base with an integral exponent is real, so the complex + // guard must not fire on it. + assert_bits_eq!(pow(-8.0, 2.0), Ok(64.0f64)); // Test positive float base, positive float exponent - assert_approx_eq!(pow(2.5, 2.0), Ok(6.25)); - assert_approx_eq!(pow(3.5, 3.0), Ok(42.875)); - assert_approx_eq!(pow(4.5, 4.0), Ok(410.0625)); - assert_approx_eq!(pow(2.5, 3.0), Ok(15.625)); - assert_approx_eq!(pow(2.5, 4.0), Ok(39.0625)); + assert_bits_eq!(pow(2.5, 2.0), Ok(6.25f64)); + assert_bits_eq!(pow(3.5, 3.0), Ok(42.875f64)); + assert_bits_eq!(pow(4.5, 4.0), Ok(410.0625f64)); + assert_bits_eq!(pow(2.5, 3.0), Ok(15.625f64)); + assert_bits_eq!(pow(2.5, 4.0), Ok(39.0625f64)); // Test negative float base, positive float exponent (integral exponents only) - assert_approx_eq!(pow(-2.5, 2.0), Ok(6.25)); - assert_approx_eq!(pow(-3.5, 3.0), Ok(-42.875)); - assert_approx_eq!(pow(-4.5, 4.0), Ok(410.0625)); - assert_approx_eq!(pow(-2.5, 3.0), Ok(-15.625)); - assert_approx_eq!(pow(-2.5, 4.0), Ok(39.0625)); + assert_bits_eq!(pow(-2.5, 2.0), Ok(6.25f64)); + assert_bits_eq!(pow(-3.5, 3.0), Ok(-42.875f64)); + assert_bits_eq!(pow(-4.5, 4.0), Ok(410.0625f64)); + assert_bits_eq!(pow(-2.5, 3.0), Ok(-15.625f64)); + assert_bits_eq!(pow(-2.5, 4.0), Ok(39.0625f64)); // Test positive float base, positive float exponent with non-integral exponents - assert_approx_eq!(pow(2.0, 2.5), Ok(5.656854249492381)); - assert_approx_eq!(pow(3.0, 3.5), Ok(46.76537180435969)); - assert_approx_eq!(pow(4.0, 4.5), Ok(512.0)); - assert_approx_eq!(pow(2.0, 3.5), Ok(11.313708498984761)); - assert_approx_eq!(pow(2.0, 4.5), Ok(22.627416997969522)); + assert_bits_eq!(pow(2.0, 2.5), Ok(5.656854249492381f64)); + assert_bits_eq!(pow(3.0, 3.5), Ok(46.76537180435969f64)); + assert_bits_eq!(pow(4.0, 4.5), Ok(512.0f64)); + assert_bits_eq!(pow(2.0, 3.5), Ok(11.313708498984761f64)); + assert_bits_eq!(pow(2.0, 4.5), Ok(22.627416997969522f64)); // Test positive float base, negative float exponent - assert_approx_eq!(pow(2.0, -2.5), Ok(0.1767766952966369)); - assert_approx_eq!(pow(3.0, -3.5), Ok(0.021383343303319473)); - assert_approx_eq!(pow(4.0, -4.5), Ok(0.001953125)); - assert_approx_eq!(pow(2.0, -3.5), Ok(0.08838834764831845)); - assert_approx_eq!(pow(2.0, -4.5), Ok(0.04419417382415922)); + assert_bits_eq!(pow(2.0, -2.5), Ok(0.1767766952966369f64)); + assert_bits_eq!(pow(3.0, -3.5), Ok(0.021383343303319473f64)); + assert_bits_eq!(pow(4.0, -4.5), Ok(0.001953125f64)); + assert_bits_eq!(pow(2.0, -3.5), Ok(0.08838834764831845f64)); + assert_bits_eq!(pow(2.0, -4.5), Ok(0.04419417382415922f64)); // Test negative float base, negative float exponent (integral exponents only) - assert_approx_eq!(pow(-2.0, -2.0), Ok(0.25)); - assert_approx_eq!(pow(-3.0, -3.0), Ok(-0.037037037037037035)); - assert_approx_eq!(pow(-4.0, -4.0), Ok(0.00390625)); - assert_approx_eq!(pow(-2.0, -3.0), Ok(-0.125)); - assert_approx_eq!(pow(-2.0, -4.0), Ok(0.0625)); - - // Currently negative float base with non-integral exponent is not supported: - // assert_approx_eq!(pow(-2.0, 2.5), Ok(5.656854249492381)); - // assert_approx_eq!(pow(-3.0, 3.5), Ok(-46.76537180435969)); - // assert_approx_eq!(pow(-4.0, 4.5), Ok(512.0)); - // assert_approx_eq!(pow(-2.0, -2.5), Ok(0.1767766952966369)); - // assert_approx_eq!(pow(-3.0, -3.5), Ok(0.021383343303319473)); - // assert_approx_eq!(pow(-4.0, -4.5), Ok(0.001953125)); - - // Extra cases **NOTE** these are not all working: - // * If they are commented in then they work - // * If they are commented out with a number that is the current return value it throws vs the expected value - // * If they are commented out with a "fail to run" that means I couldn't get them to work, could add a case for really big or small values - // 1e308^2.0 - assert_approx_eq!(pow(1e308, 2.0), Ok(f64::INFINITY)); - // 1e308^(1e-2) - assert_approx_eq!(pow(1e308, 1e-2), Ok(1202.2644346174131)); - // 1e-308^2.0 - //assert_approx_eq!(pow(1e-308, 2.0), Ok(0.0)); // --8.403311421507407 - // 1e-308^-2.0 - assert_approx_eq!(pow(1e-308, -2.0), Ok(f64::INFINITY)); - // 1e100^(1e50) - //assert_approx_eq!(pow(1e100, 1e50), Ok(1.0000000000000002e+150)); // fail to run (Crashes as "illegal hardware instruction") - // 1e50^(1e-100) - assert_approx_eq!(pow(1e50, 1e-100), Ok(1.0)); - // 1e308^(-1e2) - //assert_approx_eq!(pow(1e308, -1e2), Ok(0.0)); // 2.961801792837933e25 - // 1e-308^(1e2) - //assert_approx_eq!(pow(1e-308, 1e2), Ok(f64::INFINITY)); // 1.6692559244043896e46 - // 1e308^(-1e308) - // assert_approx_eq!(pow(1e308, -1e308), Ok(0.0)); // fail to run (Crashes as "illegal hardware instruction") - // 1e-308^(1e308) - // assert_approx_eq!(pow(1e-308, 1e308), Ok(0.0)); // fail to run (Crashes as "illegal hardware instruction") + assert_bits_eq!(pow(-2.0, -2.0), Ok(0.25f64)); + assert_bits_eq!(pow(-3.0, -3.0), Ok(-0.037037037037037035f64)); + assert_bits_eq!(pow(-4.0, -4.0), Ok(0.00390625f64)); + assert_bits_eq!(pow(-2.0, -3.0), Ok(-0.125f64)); + assert_bits_eq!(pow(-2.0, -4.0), Ok(0.0625f64)); + + // A negative base raised to a non-integral exponent is complex, + // which this crate cannot produce, so it deoptimizes instead - see + // `float_power_deopts_on_negative_base_fractional_exponent` in + // deopt_tests.rs. + + // Extreme magnitudes, finite on both sides: + assert_bits_eq!(pow(1e308, 1e-2), Ok(1202.2644346174131f64)); + assert_bits_eq!(pow(1e50, 1e-100), Ok(1.0f64)); + // 1e308 ** 2.0 overflows a finite base to an infinity, which raises + // OverflowError rather than saturating - see + // `float_power_deopts_on_finite_base_overflow` in deopt_tests.rs. + // Underflowing all the way to zero, in both directions, does not + // deoptimize - only an overflow to infinity does. + assert_bits_eq!(pow(1e-308, 2.0), Ok(0.0f64)); + assert_bits_eq!(pow(1e308, -1e2), Ok(0.0f64)); + assert_bits_eq!(pow(1e-308, 1e2), Ok(0.0f64)); + assert_bits_eq!(pow(1e308, -1e308), Ok(0.0f64)); + assert_bits_eq!(pow(1e-308, 1e308), Ok(0.0f64)); + } + + /// The lowering used to answer float `**` with a hand-rolled + /// double–double `ln`/`exp`, which lost whole significant digits on a + /// base far from 1: `1023.0 ** 1.0` came back as `1022.9277018310074`, + /// and the error stayed small enough at well-scaled inputs that + /// `assert_approx_eq!`'s relative tolerance in `basic_power` above never + /// caught it. Calling `f64::powf` directly cannot drift from the + /// interpreter this way: both sides run the same function on the same + /// bits. Four of these 24 pairs overflow a finite base and exponent to + /// an infinity and deoptimize instead of returning; see + /// `float_power_deopts_on_finite_base_overflow` in deopt_tests.rs. + #[test] + fn float_power_matches_far_from_one() { + let pow = jit_function! { pow(a:f64, b:f64) -> f64 => r##" + def pow(a:float, b: float): + return a**b + "##}; + assert_bits_eq!(pow(1023.0, 1.0), Ok(1023.0f64)); + assert_bits_eq!(pow(1023.0, 2.0), Ok(1046529.0f64)); + assert_bits_eq!(pow(1023.0, -2.0), Ok(9.555396935966418e-7f64)); + assert_bits_eq!(pow(1023.0, 4.0), Ok(1095222947841.0f64)); + assert_bits_eq!(pow(1023.0, -320.0), Ok(0.0f64)); + assert_bits_eq!(pow(1023.0, 0.5), Ok(31.984371183438952f64)); + assert_bits_eq!(pow(1e-308, 1.0), Ok(1e-308f64)); + assert_bits_eq!(pow(1e-308, 2.0), Ok(0.0f64)); + // (1e-308, -2.0) overflows - see deopt_tests.rs. + assert_bits_eq!(pow(1e-308, 4.0), Ok(0.0f64)); + // (1e-308, -320.0) overflows - see deopt_tests.rs. + assert_bits_eq!(pow(1e-308, 0.5), Ok(1e-154f64)); + assert_bits_eq!(pow(1e-100, 1.0), Ok(1e-100f64)); + assert_bits_eq!(pow(1e-100, 2.0), Ok(1e-200f64)); + assert_bits_eq!(pow(1e-100, -2.0), Ok(1e200f64)); + assert_bits_eq!(pow(1e-100, 4.0), Ok(0.0f64)); + // (1e-100, -320.0) overflows - see deopt_tests.rs. + assert_bits_eq!(pow(1e-100, 0.5), Ok(1e-50f64)); + assert_bits_eq!(pow(1e100, 1.0), Ok(1e100f64)); + assert_bits_eq!(pow(1e100, 2.0), Ok(1e200f64)); + assert_bits_eq!(pow(1e100, -2.0), Ok(1e-200f64)); + // (1e100, 4.0) overflows - see deopt_tests.rs. + assert_bits_eq!(pow(1e100, -320.0), Ok(0.0f64)); + assert_bits_eq!(pow(1e100, 0.5), Ok(1e50f64)); } #[test] @@ -231,11 +300,10 @@ mod tests { "## }; assert_approx_eq!(div(5.2, 2.0), Ok(2.6)); + assert_approx_eq!(div(4.0, 2.0), Ok(2.0)); assert_approx_eq!(div(3.4, -1.7), Ok(-2.0)); - assert_eq!(div(1.0, 0.0), Ok(f64::INFINITY)); - assert_eq!(div(1.0, -0.0), Ok(f64::NEG_INFINITY)); - assert_eq!(div(-1.0, 0.0), Ok(f64::NEG_INFINITY)); - assert_eq!(div(-1.0, -0.0), Ok(f64::INFINITY)); + // Division by zero raises rather than returning an infinity, so it + // deoptimizes instead - see deopt_tests.rs. assert_bits_eq!(div(-5.2, f64::NAN), Ok(f64::NAN)); assert_eq!(div(f64::INFINITY, 2.0), Ok(f64::INFINITY)); assert_bits_eq!(div(-2.0, f64::NEG_INFINITY), Ok(0.0f64)); @@ -253,10 +321,8 @@ mod tests { assert_approx_eq!(div(5.2, 2), Ok(2.6)); assert_approx_eq!(div(3.4, -1), Ok(-3.4)); - assert_eq!(div(1.0, 0), Ok(f64::INFINITY)); - assert_eq!(div(1.0, -0), Ok(f64::INFINITY)); - assert_eq!(div(-1.0, 0), Ok(f64::NEG_INFINITY)); - assert_eq!(div(-1.0, -0), Ok(f64::NEG_INFINITY)); + // Division by zero raises rather than returning an infinity, so it + // deoptimizes instead - see deopt_tests.rs. assert_eq!(div(f64::INFINITY, 2), Ok(f64::INFINITY)); assert_eq!(div(f64::NEG_INFINITY, 3), Ok(f64::NEG_INFINITY)); } diff --git a/crates/jit/tests/int_tests.rs b/crates/jit/tests/int_tests.rs index 23cf98aafe1..2f2a4ac5009 100644 --- a/crates/jit/tests/int_tests.rs +++ b/crates/jit/tests/int_tests.rs @@ -2,6 +2,36 @@ mod tests { use core::f64; + /// Floor division rounds toward negative infinity and the remainder takes + /// the divisor's sign, which is not what the machine instructions do. + #[test] + fn floor_div_and_remainder_follow_the_divisor() { + let div = jit_function! { div(a: i64, b: i64) -> i64 => r#" + def div(a: int, b: int) -> int: + return a // b + "# }; + assert_eq!(div(7, 2), Ok(3)); + assert_eq!(div(-7, 2), Ok(-4)); + assert_eq!(div(7, -2), Ok(-4)); + assert_eq!(div(-7, -2), Ok(3)); + assert_eq!(div(-6, 2), Ok(-3)); + // The correction's overflow argument rests on `i64::MIN`'s quotient + // only arising where the division is exact; pin both ends of that. + assert_eq!(div(i64::MIN, 1), Ok(i64::MIN)); + assert_eq!(div(i64::MIN, 3), Ok(-3074457345618258603)); + + let rem = jit_function! { rem(a: i64, b: i64) -> i64 => r#" + def rem(a: int, b: int) -> int: + return a % b + "# }; + assert_eq!(rem(7, 2), Ok(1)); + assert_eq!(rem(-7, 2), Ok(1)); + assert_eq!(rem(7, -2), Ok(-1)); + assert_eq!(rem(-7, -2), Ok(-1)); + assert_eq!(rem(-6, 2), Ok(0)); + assert_eq!(rem(i64::MIN, 3), Ok(1)); + } + #[test] fn basic_add() { let add = jit_function! { add(a:i64, b:i64) -> i64 => r##" @@ -65,10 +95,9 @@ mod tests { assert_eq!(div(1, 100000), Ok(0.00001)); assert_eq!(div(2, 3), Ok(0.6666666666666666)); assert_eq!(div(1, 3), Ok(0.3333333333333333)); - assert_eq!(div(i64::MAX, 2), Ok(4611686018427387904.0)); - assert_eq!(div(i64::MIN, 2), Ok(-4611686018427387904.0)); - assert_eq!(div(i64::MIN, -1), Ok(9223372036854775808.0)); // Overflow case - assert_eq!(div(i64::MIN, i64::MAX), Ok(-1.0)); + assert_eq!(div(1i64 << 53, 1), Ok((1i64 << 53) as f64)); + // An operand past `1 << 53` does not fit a double's significand and + // deoptimizes instead; see deopt_tests.rs. } #[test] @@ -119,7 +148,8 @@ mod tests { assert_eq!(modulo(12, 10), Ok(2)); assert_eq!(modulo(7, 10), Ok(7)); assert_eq!(modulo(-3, 1), Ok(0)); - assert_eq!(modulo(-5, 10), Ok(-5)); + // The remainder takes the divisor's sign, not the dividend's. + assert_eq!(modulo(-5, 10), Ok(5)); } #[test] @@ -132,6 +162,14 @@ mod tests { assert_eq!(power(10, 2), Ok(100)); assert_eq!(power(5, 1), Ok(5)); assert_eq!(power(1, 0), Ok(1)); + // Square-and-multiply squares the base every iteration and discards + // it once there is no further iteration to consume it, so the final + // squaring overflowing to 2^64 must not deoptimize an answer - 2^33 + // and 2^62 - that fits an i64 comfortably. That discard is guarded + // on "another iteration exists", not on the bit that would read the + // squared value. + assert_eq!(power(2, 33), Ok(8589934592)); + assert_eq!(power(2, 62), Ok(4611686018427387904)); } #[test] diff --git a/crates/jit/tests/lib.rs b/crates/jit/tests/lib.rs index aa5f0f22d64..a86a978eed8 100644 --- a/crates/jit/tests/lib.rs +++ b/crates/jit/tests/lib.rs @@ -1,7 +1,14 @@ +extern crate alloc; + #[macro_use] mod common; mod bool_tests; +mod deopt_tests; +mod engine_tests; mod float_tests; mod int_tests; mod misc_tests; mod none_tests; +mod safepoint_tests; +mod safety_tests; +mod support_tests; diff --git a/crates/jit/tests/misc_tests.rs b/crates/jit/tests/misc_tests.rs index 5404df0a769..2002f733d51 100644 --- a/crates/jit/tests/misc_tests.rs +++ b/crates/jit/tests/misc_tests.rs @@ -1,6 +1,6 @@ #[cfg(test)] mod tests { - use rustpython_jit::{AbiValue, JitArgumentError}; + use rustpython_jit::{AbiValue, JitArgumentError, JitCompileError, JitEngine, Outcome, Safety}; #[test] fn no_return_value() { @@ -33,7 +33,7 @@ mod tests { ); assert_eq!( func.invoke(&[AbiValue::Int(1), AbiValue::Float(2.0)]), - Ok(Some(AbiValue::Int(1))) + Ok(Outcome::Returned(Some(AbiValue::Int(1)))) ); } @@ -64,7 +64,10 @@ mod tests { let args = args_builder.into_args(); assert!(args.is_some()); - assert_eq!(args.unwrap().invoke(), Some(AbiValue::Int(1))); + assert_eq!( + args.unwrap().invoke(), + Outcome::Returned(Some(AbiValue::Int(1))) + ); } #[test] @@ -125,4 +128,24 @@ mod tests { assert_eq!(fib(10), Ok(89)); } + + /// A local read after being assigned on only one branch is the shape the + /// deopt spill has to survive: it reads every live local at every guard, + /// regardless of which branches actually ran. The bytecode compiler + /// substitutes `LOAD_FAST_CHECK` for `LOAD_FAST` on any local that is not + /// provably bound on every incoming path, and that opcode has no lowering + /// here, so this rejection is about a missing instruction, not about + /// whether the backend can read a partially-defined local. + #[test] + fn conditionally_defined_local_is_not_compiled() { + let engine = JitEngine::new(None); + let f = py_function_def!(f => r#" + def f(c: bool) -> int: + if c: + x = 1 + return x + "#); + let result = f.compile_on(&engine, Safety::Permissive); + assert!(matches!(result, Err(JitCompileError::NotSupported))); + } } diff --git a/crates/jit/tests/safepoint_tests.rs b/crates/jit/tests/safepoint_tests.rs new file mode 100644 index 00000000000..85fc77cfbf6 --- /dev/null +++ b/crates/jit/tests/safepoint_tests.rs @@ -0,0 +1,127 @@ +#[cfg(test)] +mod tests { + use core::sync::atomic::{AtomicU8, Ordering}; + use rustpython_jit::{AbiValue, JitEngine, Outcome, Safety}; + + /// A loop polls the word it was compiled against every time round, and + /// leaves as soon as it reads anything but zero. The record it leaves says + /// where the loop had got to, so the interpreter picks the same iteration + /// back up rather than starting the function again. + /// + /// Each test owns its word: the engines are separate, and a shared one + /// would leak the trip into whichever test happened to run beside it. + #[test] + fn a_loop_leaves_when_the_word_it_polls_is_set() { + static WORD: AtomicU8 = AtomicU8::new(0); + let engine = JitEngine::new(Some(&WORD)); + let f = py_function_def! { spin => r#" +def spin(n: int) -> int: + i = 0 + while i < n: + i = i + 1 + return i +"# }; + let code = f + .compile_on(&engine, Safety::Strict) + .expect("should compile"); + assert_eq!( + code.invoke(&[3i64.into()]), + Ok(Outcome::Returned(Some(3i64.into()))) + ); + + WORD.store(1, Ordering::Release); + match code.invoke(&[3i64.into()]) { + Ok(Outcome::Interrupted(Some(state))) => { + // One pass through the body, then the jump back polls. + assert_eq!(state.locals[0], Some(AbiValue::Int(3))); + assert_eq!(state.locals[1], Some(AbiValue::Int(1))); + assert!(state.stack.is_empty(), "{:?}", state.stack); + } + other => panic!("expected an interruption, got {other:?}"), + } + } + + /// The word is read on every iteration, not once on the way in: a loop + /// already running has to notice a word set while it runs. Nothing else + /// here proves the load survives into the loop body rather than being + /// hoisted out of it. + #[test] + fn a_running_loop_notices_the_word_being_set() { + static WORD: AtomicU8 = AtomicU8::new(0); + let engine = JitEngine::new(Some(&WORD)); + let f = py_function_def! { spin => r#" +def spin(n: int) -> int: + i = 0 + while i < n: + i = i + 1 + return i +"# }; + let code = f + .compile_on(&engine, Safety::Strict) + .expect("should compile"); + + // Long enough that the setter lands somewhere in the middle of it, and + // small enough that the test still ends if it does not. + let iterations = 1i64 << 32; + let setter = std::thread::spawn(|| { + std::thread::sleep(core::time::Duration::from_millis(20)); + WORD.store(1, Ordering::Release); + }); + let outcome = code.invoke(&[iterations.into()]); + setter.join().expect("the setter must not panic"); + + match outcome { + Ok(Outcome::Interrupted(Some(state))) => { + let Some(AbiValue::Int(reached)) = state.locals[1] else { + panic!("the counter must come back as an int: {state:?}"); + }; + assert!( + (0..iterations).contains(&reached), + "left mid-loop, not at either end: {reached}" + ); + } + other => panic!("expected an interruption, got {other:?}"), + } + WORD.store(0, Ordering::Release); + } + + /// A function with no loop can only run for as long as its own body, so it + /// polls nowhere and runs to its end however the word reads. + #[test] + fn a_straight_line_function_does_not_poll() { + static WORD: AtomicU8 = AtomicU8::new(1); + let engine = JitEngine::new(Some(&WORD)); + let f = py_function_def! { add => r#" +def add(a: int, b: int) -> int: + return a + b +"# }; + let code = f + .compile_on(&engine, Safety::Strict) + .expect("should compile"); + assert_eq!( + code.invoke(&[2i64.into(), 3i64.into()]), + Ok(Outcome::Returned(Some(5i64.into()))) + ); + } + + /// Compiled without a word there is nothing to poll, and the loop runs to + /// completion. This is what an engine with no interpreter behind it gets. + #[test] + fn a_loop_compiled_against_no_word_runs_to_the_end() { + let engine = JitEngine::new(None); + let f = py_function_def! { spin => r#" +def spin(n: int) -> int: + i = 0 + while i < n: + i = i + 1 + return i +"# }; + let code = f + .compile_on(&engine, Safety::Strict) + .expect("should compile"); + assert_eq!( + code.invoke(&[100i64.into()]), + Ok(Outcome::Returned(Some(100i64.into()))) + ); + } +} diff --git a/crates/jit/tests/safety_tests.rs b/crates/jit/tests/safety_tests.rs new file mode 100644 index 00000000000..cd22d6c53b3 --- /dev/null +++ b/crates/jit/tests/safety_tests.rs @@ -0,0 +1,233 @@ +#[cfg(test)] +mod tests { + use rustpython_jit::{JitEngine, Outcome, Safety}; + + /// Assert Strict rejects the function while Permissive still accepts it, + /// so the test cannot pass because of an unrelated compile failure. + macro_rules! assert_strict_rejects { + ($name:ident => $src:expr) => {{ + let engine = JitEngine::new(None); + let f = py_function_def!($name => $src); + assert!( + f.compile_on(&engine, Safety::Strict).is_err(), + concat!(stringify!($name), " should not compile under Strict") + ); + f.compile_on(&engine, Safety::Permissive).expect(concat!( + stringify!($name), + " is only meant to be rejected for being unsafe, but Permissive cannot compile it either" + )); + }}; + } + + macro_rules! assert_accepted { + ($safety:expr, $name:ident => $src:expr) => {{ + let engine = JitEngine::new(None); + let f = py_function_def!($name => $src); + f.compile_on(&engine, $safety) + .expect(concat!(stringify!($name), " should compile")) + }}; + } + + /// The operation used to be refused under Strict because its machine + /// code could trap or wrap. Both are guarded by a deopt now, so Strict + /// compiles the function and, given the input that used to be unsafe, + /// hands the operands back rather than trapping, wrapping, or otherwise + /// answering wrongly. An optional fourth and fifth argument also pin an + /// ordinary input's answer, so a guard that regressed to an + /// unconditional deopt could not leave every test in this file green. + macro_rules! assert_strict_deopts { + ($name:ident => $src:expr, $bad:expr) => {{ + let code = assert_accepted!(Safety::Strict, $name => $src); + match code.invoke(&$bad) { + Ok(Outcome::Deopt(_)) => {} + other => panic!( + "{} expected a deopt under Strict, got {other:?}", + stringify!($name) + ), + } + }}; + ($name:ident => $src:expr, $bad:expr, $good:expr, $expected:expr) => {{ + let code = assert_accepted!(Safety::Strict, $name => $src); + assert_eq!( + code.invoke(&$good), + Ok(Outcome::Returned(Some($expected.into()))), + "{} expected the ordinary input to still answer", + stringify!($name) + ); + match code.invoke(&$bad) { + Ok(Outcome::Deopt(_)) => {} + other => panic!( + "{} expected a deopt under Strict, got {other:?}", + stringify!($name) + ), + } + }}; + } + + #[test] + fn strict_compiles_int_add() { + assert_strict_deopts!(add => r#" +def add(a: int, b: int) -> int: + return a + b +"#, [i64::MAX.into(), 1i64.into()], [3i64.into(), 4i64.into()], 7i64); + } + + #[test] + fn strict_compiles_int_multiply() { + assert_strict_deopts!(mul => r#" +def mul(a: int, b: int) -> int: + return a * b +"#, [i64::MAX.into(), 2i64.into()]); + } + + #[test] + fn strict_compiles_int_floor_divide() { + assert_strict_deopts!(fdiv => r#" +def fdiv(a: int, b: int) -> int: + return a // b +"#, [7i64.into(), 0i64.into()]); + } + + #[test] + fn strict_compiles_int_true_divide() { + assert_strict_deopts!(true_divide => r#" +def true_divide(a: int, b: int) -> float: + return a / b +"#, [7i64.into(), 0i64.into()]); + } + + #[test] + fn strict_compiles_int_remainder() { + assert_strict_deopts!(rem => r#" +def rem(a: int, b: int) -> int: + return a % b +"#, [7i64.into(), 0i64.into()]); + } + + #[test] + fn strict_compiles_int_power() { + assert_strict_deopts!(pow => r#" +def pow(a: int, b: int) -> int: + return a ** b +"#, [2i64.into(), 64i64.into()]); + } + + #[test] + fn strict_compiles_int_shift() { + assert_strict_deopts!(shift => r#" +def shift(a: int, b: int) -> int: + return a << b +"#, [1i64.into(), 64i64.into()]); + } + + #[test] + fn strict_compiles_int_negate() { + assert_strict_deopts!(neg => r#" +def neg(a: int) -> int: + return -a +"#, [i64::MIN.into()]); + } + + /// `1.0 / 0.0` raises ZeroDivisionError; a bare `fdiv` would return inf. + #[test] + fn strict_compiles_float_divide() { + assert_strict_deopts!(fdiv => r#" +def fdiv(a: float, b: float) -> float: + return a / b +"#, [1.0f64.into(), 0.0f64.into()]); + } + + /// `(-8.0) ** 0.5` is complex in Python, which a compiled `**` cannot + /// produce. + #[test] + fn strict_compiles_float_power() { + assert_strict_deopts!(float_power => r#" +def float_power(a: float, b: float) -> float: + return a ** b +"#, [(-8.0f64).into(), 0.5f64.into()]); + } + + #[test] + fn strict_compiles_mixed_divide() { + assert_strict_deopts!(mixed => r#" +def mixed(a: int, b: float) -> float: + return a / b +"#, [1i64.into(), 0.0f64.into()]); + } + + /// Bitwise operations on two machine integers cannot leave the range. + #[test] + fn strict_allows_int_bitwise() { + let code = assert_accepted!(Safety::Strict, band => r#" +def band(a: int, b: int) -> int: + return a & b +"#); + assert_eq!( + code.invoke(&[6i64.into(), 3i64.into()]), + Ok(Outcome::Returned(Some(2i64.into()))) + ); + } + + #[test] + fn strict_allows_int_comparison() { + let code = assert_accepted!(Safety::Strict, lt => r#" +def lt(a: int, b: int) -> bool: + return a < b +"#); + assert_eq!( + code.invoke(&[1i64.into(), 2i64.into()]), + Ok(Outcome::Returned(Some(true.into()))) + ); + } + + #[test] + fn strict_allows_float_add_and_multiply() { + let code = assert_accepted!(Safety::Strict, poly => r#" +def poly(a: float, b: float) -> float: + return a * b + a - b +"#); + assert_eq!( + code.invoke(&[2.0f64.into(), 3.0f64.into()]), + Ok(Outcome::Returned(Some(5.0f64.into()))) + ); + } + + /// Mixing an int into float addition converts exactly the way the + /// interpreter does, so it stays available under Strict. + #[test] + fn strict_allows_mixed_add() { + let code = assert_accepted!(Safety::Strict, mixed => r#" +def mixed(a: int, b: float) -> float: + return a + b +"#); + assert_eq!( + code.invoke(&[2i64.into(), 0.5f64.into()]), + Ok(Outcome::Returned(Some(2.5f64.into()))) + ); + } + + /// The self-reference resolves by name, but the interpreter re-reads the + /// global on every call, so a rebound name would make them disagree. + /// This is the one place Strict and Permissive still differ. + #[test] + fn strict_rejects_self_recursion() { + assert_strict_rejects!(countdown => r#" +def countdown(a: float) -> float: + if a > 0.0: + return countdown(a - 1.0) + return a +"#); + } + + #[test] + fn permissive_still_compiles_int_arithmetic() { + let code = assert_accepted!(Safety::Permissive, add => r#" +def add(a: int, b: int) -> int: + return a + b +"#); + assert_eq!( + code.invoke(&[3i64.into(), 4i64.into()]), + Ok(Outcome::Returned(Some(7i64.into()))) + ); + } +} diff --git a/crates/jit/tests/support_tests.rs b/crates/jit/tests/support_tests.rs new file mode 100644 index 00000000000..3c2f2f0048f --- /dev/null +++ b/crates/jit/tests/support_tests.rs @@ -0,0 +1,118 @@ +#[cfg(test)] +mod tests { + use rustpython_jit::supports_code; + + macro_rules! assert_supported { + ($name:ident => $src:expr) => {{ + let f = py_function_def!($name => $src); + assert!( + supports_code(f.code()), + concat!(stringify!($name), " should pass the pre-filter") + ); + }}; + } + + macro_rules! assert_unsupported { + ($name:ident => $src:expr) => {{ + let f = py_function_def!($name => $src); + assert!( + !supports_code(f.code()), + concat!(stringify!($name), " should be rejected by the pre-filter") + ); + }}; + } + + #[test] + fn plain_arithmetic_is_supported() { + assert_supported!(add => r#" +def add(a: int, b: int) -> int: + return a + b +"#); + } + + #[test] + fn branches_and_loops_are_supported() { + assert_supported!(count => r#" +def count(n: int) -> int: + total = 0 + while n > 0: + if n > 5: + total = total + 1 + n = n - 1 + return total +"#); + } + + #[test] + fn mid_expression_merges_are_rejected() { + // The compiler cannot reconcile the value stack where control flow + // merges, so the pre-filter has to turn a merge reached mid-expression + // down. Without this the depth simulation could stop firing and the + // only symptom would be compile attempts that always fail. + // + // Every opcode below is one `instruction_is_supported` accepts, which + // is what leaves the merge as the only thing that can reject it. A + // short-circuit operator looks like the smaller shape for this and is + // not: `and` and `or` compile to `COPY`, which the opcode filter + // rejects on its own, so such a case passes whether or not this clause + // is here. Assigning the conditional expression rather than returning + // it matters too - codegen tail-duplicates one in return position, so + // `return (a if b else b) + 1` has no merge to reject. + assert_unsupported!(merge_mid_expression => r#" +def merge_mid_expression(a: int, b: int) -> int: + c = (a if b else b) + 1 + return c +"#); + } + + #[test] + fn varargs_are_rejected() { + assert_unsupported!(va => r#" +def va(*args) -> int: + return 1 +"#); + } + + #[test] + fn varkeywords_are_rejected() { + assert_unsupported!(vk => r#" +def vk(**kwargs) -> int: + return 1 +"#); + } + + #[test] + fn generators_are_rejected() { + assert_unsupported!(gen => r#" +def gen(n: int): + yield n +"#); + } + + #[test] + fn containers_are_rejected() { + assert_unsupported!(indexing => r#" +def indexing(a: int) -> int: + return [a][0] +"#); + } + + #[test] + fn attribute_access_is_rejected() { + assert_unsupported!(attr => r#" +def attr(a: int) -> int: + return a.bit_length() +"#); + } + + #[test] + fn exception_handling_is_rejected() { + assert_unsupported!(guarded => r#" +def guarded(a: int) -> int: + try: + return a + except ValueError: + return 0 +"#); + } +} diff --git a/crates/vm/Cargo.toml b/crates/vm/Cargo.toml index cc7c8dec2f3..d74ffe1f24c 100644 --- a/crates/vm/Cargo.toml +++ b/crates/vm/Cargo.toml @@ -19,6 +19,7 @@ vm-tracing-logging = [] flame-it = ["flame", "flamer"] freeze-stdlib = ["encodings"] jit = ["rustpython-jit"] +aot = ["jit"] threading = ["rustpython-common/threading"] gc = [] compiler = ["parser", "codegen", "rustpython-compiler"] diff --git a/crates/vm/src/builtins/code.rs b/crates/vm/src/builtins/code.rs index 9739cfe8e17..1b939be4277 100644 --- a/crates/vm/src/builtins/code.rs +++ b/crates/vm/src/builtins/code.rs @@ -489,6 +489,10 @@ pub struct PyCode { /// this code cannot leave the slot unbalanced, so `with_frame` skips the /// exc_info save/restore. Computed once by scanning the instruction stream. pub has_exc_handling: bool, + /// Cached verdict of the AOT pre-filter: 0 not yet run, 1 eligible, + /// 2 rejected. Threads that race here compute the same answer. + #[cfg(feature = "jit")] + pub aot_precheck: core::sync::atomic::AtomicU8, } impl Deref for PyCode { @@ -620,6 +624,8 @@ impl PyCode { monitoring_data: PyMutex::new(None), quickened: core::sync::atomic::AtomicBool::new(false), has_exc_handling, + #[cfg(feature = "jit")] + aot_precheck: core::sync::atomic::AtomicU8::new(0), } } diff --git a/crates/vm/src/builtins/float.rs b/crates/vm/src/builtins/float.rs index 2c07f9b81a3..2f7a6c96e6a 100644 --- a/crates/vm/src/builtins/float.rs +++ b/crates/vm/src/builtins/float.rs @@ -156,9 +156,13 @@ fn inner_divmod(v1: f64, v2: f64, vm: &VirtualMachine) -> PyResult<(f64, f64)> { } pub(crate) fn float_pow(v1: f64, v2: f64, vm: &VirtualMachine) -> PyResult { - if v1.is_zero() && v2.is_sign_negative() { + // Both tests are on the value, not the sign bit: `-0.0` is neither a + // negative exponent nor a negative base, so `0.0 ** -0.0` is 1.0 and + // `(-0.0) ** 0.5` is 0.0. An exponent that is not a number falls through + // to `powf`, which answers it the way `pow` is defined to. + if v1.is_zero() && v2 < 0.0 { Err(vm.new_zero_division_error("zero to a negative power")) - } else if v1.is_sign_negative() && (v2.floor() - v2).abs() > f64::EPSILON { + } else if v1 < 0.0 && v2.is_finite() && v2 != v2.floor() { let v1 = Complex64::new(v1, 0.); let v2 = Complex64::new(v2, 0.); Ok(super::complex::complex_pow(v1, v2, vm)?.to_pyobject(vm)) diff --git a/crates/vm/src/builtins/frame.rs b/crates/vm/src/builtins/frame.rs index 3704e070c34..e62f2c75fe5 100644 --- a/crates/vm/src/builtins/frame.rs +++ b/crates/vm/src/builtins/frame.rs @@ -462,6 +462,77 @@ impl FrameObject { } core::ptr::null() } + + /// Find the live source InterpreterFrame for a materialized FrameObject + /// on the chain of thread `tid`, which must not be this one — a thread + /// publishes its top frame for other threads, but keeps its own in TLS. + /// Returns null if that thread is not running the source frame. + /// + /// # Safety + /// Caller must hold the world stopped, so the owning thread is parked and + /// its chain is not being popped while it is walked. + #[cfg(feature = "threading")] + unsafe fn find_live_source_iframe_on( + &self, + tid: u64, + vm: &VirtualMachine, + ) -> *const crate::frame::InterpreterFrame { + let self_py_ptr = unsafe { Py::::from_payload_ptr(self) } as usize; + let registry = vm.state.thread_frames.lock(); + let Some(slot) = registry.get(&tid) else { + return core::ptr::null(); + }; + let mut cur = slot.top_iframe.load(Relaxed) as *const crate::frame::InterpreterFrame; + while !cur.is_null() { + if unsafe { (*cur).materialized.load(Relaxed) } == self_py_ptr { + return cur; + } + cur = unsafe { &*cur }.previous(); + } + core::ptr::null() + } + + /// The other thread that is still running the frame this object was + /// materialized from, or `None`. A source frame on this thread does not + /// count: `find_live_source_iframe` already covers this thread in full, + /// so anything it misses is running elsewhere or has returned. + #[cfg(feature = "threading")] + fn source_thread(&self) -> Option { + let tid = self.iframe().attached_tid(); + (tid != 0 && tid != crate::stdlib::_thread::get_ident()).then_some(tid) + } + + /// Where the frame this object stands for is executing. A materialized + /// frame's own copy only catches up when the source returns, so while the + /// source runs the position has to be read from it: a frame observed from + /// inside a call it made still reports that call's instruction. + fn live_lasti(&self, #[allow(unused)] vm: &VirtualMachine) -> u32 { + let live = self.find_live_source_iframe(); + if !live.is_null() { + return unsafe { (*live).lasti.load(Relaxed) }; + } + #[cfg(feature = "threading")] + if let Some(tid) = self.source_thread() { + return self.lasti_from_thread(tid, vm); + } + self.lasti() + } + + /// `live_lasti` for a source frame still running on thread `tid`. + #[cfg(feature = "threading")] + #[cold] + fn lasti_from_thread(&self, tid: u64, vm: &VirtualMachine) -> u32 { + vm.state.stop_the_world.stop_the_world(&vm.state); + scopeguard::defer! { vm.state.stop_the_world.start_the_world(&vm.state); } + // SAFETY: the world is stopped, so the owning thread is parked. + let live = unsafe { self.find_live_source_iframe_on(tid, vm) }; + if live.is_null() { + // The source frame returned between the read of `attached_tid` + // and the stop, so this object's own copy is up to date. + return self.lasti(); + } + unsafe { (*live).lasti.load(Relaxed) } + } } #[pyclass(flags(DISALLOW_INSTANTIATION), with(Py))] @@ -482,49 +553,26 @@ impl FrameObject { } #[pygetset] - fn f_lasti(&self) -> u32 { - // Return byte offset (each instruction is 2 bytes) for compatibility. - // For materialized frames, read live lasti from the source iframe on - // the TLS chain so f_lasti reflects the current execution position. - let live = self.find_live_source_iframe(); - let val = if !live.is_null() { - unsafe { (*live).lasti.load(Relaxed) } - } else { - self.lasti() - }; - val * 2 + fn f_lasti(&self, vm: &VirtualMachine) -> u32 { + // Byte offset — each instruction is 2 bytes. + self.live_lasti(vm) * 2 } #[pygetset] - pub fn f_lineno(&self) -> usize { - // If lasti is 0, execution hasn't started yet - use first line number - if self.lasti() == 0 { - return self - .iframe() - .code() - .first_line_number - .map_or(1, |n| n.get()); - } - // For executing frames (on the TLS chain), use prev_line which is - // updated at each bytecode instruction *before* the instruction - // runs. This gives the correct line even when observed mid-CALL - // (where lasti has already advanced past the CALL instruction). - let live = self.find_live_source_iframe(); - if !live.is_null() { - // Read live prev_line. Use read_volatile to bypass LLVM noalias - // on the &mut InterpreterFrame borrow held by the running frame. - let prev = unsafe { - let field_ptr = core::ptr::addr_of!((*live).prev_line); - core::ptr::read_volatile(field_ptr as *const u32) - }; - if prev > 0 { - return prev as usize; - } - } - // For returned frames, use lasti-based location lookup. This is - // correct for exception tracebacks where prev_line may have been - // updated by cleanup instructions after the exception. - self.current_location().line.get() + pub fn f_lineno(&self, vm: &VirtualMachine) -> usize { + let code = self.iframe().code(); + let first_line = || code.first_line_number.map_or(1, |n| n.get()); + // A running frame advances lasti past the instruction it is about to + // execute, so the line being executed is the one before it. + let lasti = self.live_lasti(vm); + let Some(idx) = lasti.checked_sub(1) else { + // Execution has not started. + return first_line(); + }; + // A live lasti can move between the read and the lookup. + code.locations + .get(idx as usize) + .map_or_else(first_line, |(loc, _)| loc.line.get()) } #[pygetset(setter)] @@ -759,6 +807,42 @@ impl FrameObject { } } +#[cfg(feature = "threading")] +impl Py { + /// `f_back` for a frame materialized from one that is still running on + /// thread `tid`. Such a copy carries no `previous` of its own, and the + /// chain it was taken from lives on that thread's stack. + #[cold] + fn back_from_thread(&self, tid: u64, vm: &VirtualMachine) -> Option> { + vm.state.stop_the_world.stop_the_world(&vm.state); + scopeguard::defer! { vm.state.stop_the_world.start_the_world(&vm.state); } + // SAFETY: the world is stopped, so the owning thread is parked. + let live = unsafe { self.find_live_source_iframe_on(tid, vm) }; + if live.is_null() { + // The source frame returned between the read of `attached_tid` + // and the stop, so its caller is recorded by now. + let retained = self.iframe().cold().retained_back.lock().clone(); + if let Some(frame) = &retained { + frame.mark_escaped(); + } + return retained; + } + let prev = unsafe { (*live).previous() }; + if prev.is_null() { + return None; + } + let prev_ref = unsafe { &*prev }; + if let Some(fo) = prev_ref.frame_obj() { + fo.mark_escaped(); + return Some(fo.to_owned()); + } + // SAFETY: the world is stopped, so the owning thread is parked. + let fo = unsafe { prev_ref.materialize_detached_chain(vm) }; + fo.mark_escaped(); + Some(fo) + } +} + #[pyclass] impl Py { #[pymethod] @@ -889,6 +973,10 @@ impl Py { frame.mark_escaped(); return Some(frame); } + #[cfg(feature = "threading")] + if let Some(tid) = self.source_thread() { + return self.back_from_thread(tid, vm); + } return None; } } diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index b8975cf1102..50a680ecb5f 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -1,4 +1,6 @@ #[cfg(feature = "jit")] +pub(crate) mod aot; +#[cfg(feature = "jit")] mod jit; use super::{ @@ -22,10 +24,12 @@ use crate::{ Representable, }, }; +#[cfg(feature = "jit")] +use core::sync::atomic::AtomicU8; use core::sync::atomic::{AtomicU32, Ordering::Relaxed}; use itertools::Itertools; #[cfg(feature = "jit")] -use rustpython_jit::CompiledCode; +use rustpython_jit::{CompiledCode, DeoptState, Outcome, Safety, StackValue}; fn format_missing_args( qualname: impl core::fmt::Display, @@ -79,6 +83,15 @@ pub struct PyFunction { func_version: AtomicU32, #[cfg(feature = "jit")] jitted_code: PyMutex>, + /// One of the `aot` state constants. Read on every call, so it is an + /// atomic rather than something behind `jitted_code`'s lock. + #[cfg(feature = "jit")] + jit_state: AtomicU8, + /// Calls made so far, counted only until the compiler has looked at this + /// function. Beside `jit_state` for the same reason: it is written on + /// every call while the function is still cold. + #[cfg(feature = "jit")] + jit_warmup: AtomicU32, } static FUNC_VERSION_COUNTER: AtomicU32 = AtomicU32::new(1); @@ -218,6 +231,10 @@ impl PyFunction { func_version: AtomicU32::new(next_func_version()), #[cfg(feature = "jit")] jitted_code: PyMutex::new(None), + #[cfg(feature = "jit")] + jit_state: AtomicU8::new(aot::UNTRIED), + #[cfg(feature = "jit")] + jit_warmup: AtomicU32::new(0), }; Ok(func) } @@ -550,38 +567,173 @@ impl Py { self.code.flags.contains(bytecode::CodeFlags::OPTIMIZED) } - /// Whether this function currently has native JIT code. Adaptive Python - /// call specializations must yield to that entry point. + /// Whether this call has to go through [`Self::invoke_with_locals`], which + /// is where native code is compiled and entered. Adaptive call + /// specializations skip that entry point, so they must yield to it - both + /// for a function that already has code and for one that has not had its + /// single automatic compile attempt yet. #[inline] - pub(crate) fn is_jitted(&self) -> bool { + pub(crate) fn requires_jit_entry(&self, vm: &VirtualMachine) -> bool { #[cfg(feature = "jit")] { - self.jitted_code.lock().is_some() + match self.jit_state.load(Relaxed) { + aot::COMPILED_AUTO | aot::COMPILED_MANUAL => true, + aot::UNTRIED => vm.state.config.settings.aot, + _ => false, + } } #[cfg(not(feature = "jit"))] { + let _ = vm; false } } + /// Drop native code and stop trying. + /// + /// A caller that only wants to stop guessing has to check where the code + /// came from first: an explicit `__jit__()` is a standing request, and + /// silently undoing it would be a surprise. A guard leaves no such + /// choice - the code is wrong for these values however it was asked + /// for, and keeping it means hitting the same guard on every retry. + #[cfg(feature = "jit")] + fn deoptimize(&self, vm: &VirtualMachine) { + self.jit_state.store(aot::REJECTED, Relaxed); + *self.jitted_code.lock() = None; + vm.state.aot_stats.deoptimized.fetch_add(1, Relaxed); + } + + /// Whether a deopt record fits the frame this function builds. + /// + /// Nothing between the guard and here measures the record against the code + /// object: the compiler decides the shape and the compiled code writes the + /// slots, and the VM takes both on trust. An oversized stack runs off the + /// end of the frame, where the push aborts the process instead of raising, + /// and an offset past the last instruction resumes into nothing. Check the + /// record against the frame it claims to describe and let a mismatch fall + /// back to running the call from the start, which costs the call over again + /// and gives up nothing else. + #[cfg(feature = "jit")] + fn deopt_state_fits(&self, state: &DeoptState) -> bool { + let code = &*self.code; + state.locals.len() <= code.localspluskinds.len() + && state.stack.len() <= code.max_stackdepth as usize + && (state.offset as usize) < code.instructions.len() + } + + /// Put a deopt record into a fresh frame: the locals as the guard saw + /// them, the operands the interrupted instruction had already pushed, and + /// the offset of that instruction, so the interpreter re-executes it whole. + #[cfg(feature = "jit")] + fn fill_locals_from_deopt( + &self, + iframe: &mut crate::frame::InterpreterFrame, + state: DeoptState, + vm: &VirtualMachine, + ) { + use crate::convert::ToPyObject; + + let fastlocals = iframe.localsplus.fastlocals_mut(); + for (local, value) in fastlocals.iter_mut().zip(state.locals) { + // `None` leaves the slot empty, which is what unbound means. + *local = value.map(|value| value.to_pyobject(vm)); + } + for entry in state.stack { + // The callable and the null beside it take no slot in the record: + // they are the same on every path that reaches the guard, so the + // site describes them and they are rebuilt here. + let value = match entry { + StackValue::Value(value) => Some(value.to_pyobject(vm)), + StackValue::Callee => Some(self.as_object().to_owned()), + StackValue::Null => None, + }; + iframe.localsplus.push_stack_opt(value); + } + iframe.set_lasti(state.offset); + } + pub fn invoke_with_locals( &self, func_args: FuncArgs, locals: Option, vm: &VirtualMachine, ) -> PyResult { + // Where a guard stopped, for the frame built below to carry on from. #[cfg(feature = "jit")] - if let Some(jitted_code) = self.jitted_code.lock().as_ref() { - use crate::convert::ToPyObject; - match jit::get_jit_args(self, &func_args, jitted_code, vm) { - Ok(args) => { - return Ok(args.invoke().to_pyobject(vm)); + let mut resume = None; + #[cfg(feature = "jit")] + 'compiled: { + // Compiled code runs no frame, so it reports no call, no line and + // no return. A tracer or a monitoring tool that asked for those + // events would not see the call at all, so while either is + // installed the call goes to the interpreter - and is not compiled + // on the way past, since the compilation would only be skipped. + if vm.use_tracing.get() || vm.state.monitoring_events.load() != 0 { + break 'compiled; + } + + let mut state = self.jit_state.load(Relaxed); + if state == aot::UNTRIED && vm.state.config.settings.aot { + state = aot::observe_call(self, &func_args, vm); + } + + if matches!(state, aot::COMPILED_AUTO | aot::COMPILED_MANUAL) { + // Run the call while holding the lock, but decide what to do + // about a failure after releasing it: giving the code back takes it. + let outcome = self.jitted_code.lock().as_ref().map(|jitted_code| { + jit::get_jit_args(self, &func_args, jitted_code, vm).map(|args| args.invoke()) + }); + match outcome { + Some(Ok(Outcome::Returned(ret))) => { + use crate::convert::ToPyObject; + return Ok(ret.to_pyobject(vm)); + } + Some(Ok(Outcome::Restart)) => { + // Nothing to carry on from: either a nested frame gave + // up, or the guard belongs to a site that cannot + // describe the frame. Run the call again from the + // start, which the opcodes with a lowering make + // unobservable - they touch nothing outside the frame. + self.deoptimize(vm); + } + Some(Ok(Outcome::Interrupted(state))) => { + // A backward jump found the thread had been asked out + // of the bytecode loop. That is no verdict on the code, + // which stays installed for the next call; only this + // call finishes interpreted, which is where the signal, + // the stop or the shutdown is answered. Without a + // record it starts over, on the same footing as a + // guard's restart. + resume = state.filter(|state| self.deopt_state_fits(state)); + } + Some(Ok(Outcome::Deopt(state))) => { + // The resume lands on the instruction the guard belongs + // to, so the interpreter re-executes it - including a + // recursive call, which would hit the same guard again + // if the code were still installed. + self.deoptimize(vm); + // A record that does not fit is treated as no record, + // which leaves the call to run from the start below. + if self.deopt_state_fits(&state) { + resume = Some(state); + } + } + Some(Err(err)) => { + info!( + "jit: function `{}` is falling back to being interpreted because of \ + the error: {}", + self.code.obj_name, err + ); + if state == aot::COMPILED_AUTO { + // Nobody asked for this one, and a function whose + // arguments do not fit pays twice over: the failed + // conversion on every call, plus the call + // specialization it displaced. + self.deoptimize(vm); + } + } + None => {} } - Err(err) => info!( - "jit: function `{}` is falling back to being interpreted because of the \ - error: {}", - self.code.obj_name, err - ), } } @@ -615,6 +767,17 @@ impl Py { use_datastack, vm, ); + #[cfg(feature = "jit")] + match resume { + Some(state) => { + // SAFETY: the frame was just built here, and nothing else + // holds it until it is run below. + let iframe = unsafe { frame.iframe_mut() }; + self.fill_locals_from_deopt(iframe, state, vm); + } + None => self.fill_locals_from_args(&frame, func_args, vm)?, + } + #[cfg(not(feature = "jit"))] self.fill_locals_from_args(&frame, func_args, vm)?; if is_gen || is_coro || is_async_gen { return Ok(self.make_generator_or_coro(frame, vm)); @@ -652,9 +815,17 @@ impl Py { self.closure.as_ref().map_or(&[], |c| c.as_slice()), vm, ); - let result = self - .fill_locals_from_args_iframe(iframe, func_args, vm) - .and_then(|()| vm.run_frame_fast(iframe)); + #[cfg(feature = "jit")] + let filled = match resume { + Some(state) => { + self.fill_locals_from_deopt(iframe, state, vm); + Ok(()) + } + None => self.fill_locals_from_args_iframe(iframe, func_args, vm), + }; + #[cfg(not(feature = "jit"))] + let filled = self.fill_locals_from_args_iframe(iframe, func_args, vm); + let result = filled.and_then(|()| vm.run_frame_fast(iframe)); // Release data stack memory — must happen on both success and error. unsafe { if let Some((base, size)) = iframe.release_datastack_frame() { @@ -934,6 +1105,7 @@ impl PyFunction { #[cfg(feature = "jit")] { *jit_guard = None; + self.jit_state.store(aot::UNTRIED, Relaxed); } self.func_version.store(0, Relaxed); Ok(()) @@ -1187,19 +1359,33 @@ impl PyFunction { Ok(()) } + /// Compile this function to native code, raising `JitError` if it cannot + /// be done. + /// + /// Unlike the automatic AOT path this compiles permissively, so integer + /// arithmetic that can trap or wrap is accepted. `force` compiles again over + /// code a function already has, and retries one the AOT path rejected. #[cfg(feature = "jit")] #[pymethod] - fn __jit__(zelf: PyRef, vm: &VirtualMachine) -> PyResult<()> { - let mut jit_guard = zelf.jitted_code.lock(); - if jit_guard.is_some() { + fn __jit__(zelf: PyRef, args: JitArgs, vm: &VirtualMachine) -> PyResult<()> { + let already_compiled = matches!( + zelf.jit_state.load(Relaxed), + aot::COMPILED_AUTO | aot::COMPILED_MANUAL + ); + if already_compiled && !args.force.unwrap_or(false) { return Ok(()); } + let arg_types = jit::get_jit_arg_types(&zelf, vm)?; let ret_type = jit::jit_ret_type(&zelf, vm)?; let code: &Py = &zelf.code; - let compiled = rustpython_jit::compile(&code.code, &arg_types, ret_type) + let compiled = vm + .state + .jit_engine + .compile(&code.code, &arg_types, ret_type, Safety::Permissive) .map_err(|err| jit::new_jit_error(err.to_string(), vm))?; - *jit_guard = Some(compiled); + *zelf.jitted_code.lock() = Some(compiled); + zelf.jit_state.store(aot::COMPILED_MANUAL, Relaxed); Ok(()) } } @@ -1239,6 +1425,13 @@ impl Representable for PyFunction { } } +#[cfg(feature = "jit")] +#[derive(FromArgs)] +struct JitArgs { + #[pyarg(any, optional)] + force: OptionalArg, +} + #[derive(FromArgs)] pub struct PyFunctionNewArgs { #[pyarg(positional)] @@ -1637,7 +1830,7 @@ pub(crate) fn vectorcall_function( let code: &Py = &zelf.code; let has_kwargs = kwnames.is_some_and(|kw| !kw.is_empty()); - if zelf.is_jitted() { + if zelf.requires_jit_entry(vm) { let func_args = if has_kwargs { FuncArgs::from_vectorcall_owned(args, nargs, kwnames) } else { diff --git a/crates/vm/src/builtins/function/aot.rs b/crates/vm/src/builtins/function/aot.rs new file mode 100644 index 00000000000..cdc77101126 --- /dev/null +++ b/crates/vm/src/builtins/function/aot.rs @@ -0,0 +1,107 @@ +//! Compiling functions to native code without being asked. +//! +//! `__jit__()` compiles one function on request, from the types its +//! annotations declare, and reports why it could not. The AOT path instead +//! waits for a function to be called often enough to be worth compiling and +//! takes its types from the call in front of it, which changes what the rules +//! have to be: a function compiled behind the caller's back must not answer +//! differently from the interpreter, and must not cost anything once it turns +//! out it cannot be compiled. + +use super::PyFunction; +use crate::{Py, VirtualMachine, builtins::PyCode, function::FuncArgs}; +use core::sync::atomic::{AtomicU64, Ordering::Relaxed}; +use rustpython_jit::{CompiledCode, Safety}; + +/// What the AOT path has done so far in this interpreter. Reported by +/// `sys._jit._stats()`, which is what tells us how much of a real workload +/// the current set of supported operations actually reaches. `rejected` +/// counts the functions the compiler was asked about, which is the warm ones +/// - a function called a few times and dropped is in neither number. +#[derive(Debug, Default)] +pub struct AotStats { + pub compiled: AtomicU64, + pub rejected: AtomicU64, + pub deoptimized: AtomicU64, +} + +/// Nothing has been attempted yet. +pub(super) const UNTRIED: u8 = 0; +/// Compiled by the AOT path. Gives up on the first argument mismatch. +pub(super) const COMPILED_AUTO: u8 = 1; +/// Compiled by an explicit `__jit__()`. Keeps its code across mismatches. +pub(super) const COMPILED_MANUAL: u8 = 2; +/// Will not be compiled again unless `__jit__()` asks. +pub(super) const REJECTED: u8 = 3; + +/// Calls a function makes before the compiler is asked to look at it. +/// +/// Compiling on the first call spends the eligibility scan on every function a +/// program calls once, which is most of them, and gets nothing back. Waiting +/// also decides what to specialize on: the types come from the call that +/// crosses this line, and a function called this often is being called with +/// the types it is meant for. +const WARMUP_CALLS: u32 = 64; + +const PRECHECK_ELIGIBLE: u8 = 1; +const PRECHECK_REJECTED: u8 = 2; + +/// Whether the backend could compile this code object, remembered on the code +/// object itself so the bytecode scan runs once however many function objects +/// are built from it - a decorator or a closure factory can build thousands. +fn code_is_eligible(code: &Py) -> bool { + match code.aot_precheck.load(Relaxed) { + PRECHECK_ELIGIBLE => true, + PRECHECK_REJECTED => false, + _ => { + let eligible = rustpython_jit::supports_code(&code.code); + let verdict = if eligible { + PRECHECK_ELIGIBLE + } else { + PRECHECK_REJECTED + }; + code.aot_precheck.store(verdict, Relaxed); + eligible + } + } +} + +fn try_compile( + func: &Py, + func_args: &FuncArgs, + vm: &VirtualMachine, +) -> Option { + let code: &Py = &func.code; + if !code_is_eligible(code) { + return None; + } + + // The return type is left to the compiler, which widens the signature to + // whatever the returns it lowered produce. + let arg_types = super::jit::observed_arg_types(func, func_args, vm).ok()?; + vm.state + .jit_engine + .compile(&code.code, &arg_types, None, Safety::Strict) + .ok() +} + +/// Count this call, and once the function is warm give it its one automatic +/// compile attempt. Returns the function's new state. +pub(super) fn observe_call(func: &Py, func_args: &FuncArgs, vm: &VirtualMachine) -> u8 { + if func.jit_warmup.fetch_add(1, Relaxed) < WARMUP_CALLS { + return UNTRIED; + } + + // Claim the function before compiling, so that two threads crossing the + // line together make one attempt rather than two. + func.jit_state.store(REJECTED, Relaxed); + + let Some(compiled) = try_compile(func, func_args, vm) else { + vm.state.aot_stats.rejected.fetch_add(1, Relaxed); + return REJECTED; + }; + *func.jitted_code.lock() = Some(compiled); + func.jit_state.store(COMPILED_AUTO, Relaxed); + vm.state.aot_stats.compiled.fetch_add(1, Relaxed); + COMPILED_AUTO +} diff --git a/crates/vm/src/builtins/function/jit.rs b/crates/vm/src/builtins/function/jit.rs index 96c1465d4f1..8489275ecf4 100644 --- a/crates/vm/src/builtins/function/jit.rs +++ b/crates/vm/src/builtins/function/jit.rs @@ -8,7 +8,7 @@ use crate::{ function::FuncArgs, }; use num_traits::ToPrimitive; -use rustpython_jit::{AbiValue, Args, CompiledCode, JitArgumentError, JitType}; +use rustpython_jit::{AbiValue, Args, ArgsBuilder, CompiledCode, JitArgumentError, JitType}; #[derive(Debug, thiserror::Error)] pub(super) enum ArgsError { @@ -31,9 +31,9 @@ pub(super) enum ArgsError { impl ToPyObject for AbiValue { fn to_pyobject(self, vm: &VirtualMachine) -> PyObjectRef { match self { - AbiValue::Int(i) => i.to_pyobject(vm), - AbiValue::Float(f) => f.to_pyobject(vm), - AbiValue::Bool(b) => b.to_pyobject(vm), + Self::Int(i) => i.to_pyobject(vm), + Self::Float(f) => f.to_pyobject(vm), + Self::Bool(b) => b.to_pyobject(vm), _ => unimplemented!(), } } @@ -152,18 +152,79 @@ fn get_jit_value(vm: &VirtualMachine, obj: &PyObject) -> Result( +/// Where a walk over a call's arguments puts what it finds. The two walks +/// resolve the same parameters from the same sources and must agree on which +/// calls they accept, so they share the walk and differ only here: one fills a +/// compiled function's slots, the other collects the types those slots hold. +trait ArgSink { + /// Take the object this parameter gets its value from. A parameter the + /// compiled code cannot carry is rejected here. + fn put(&mut self, index: usize, value: &PyObject, vm: &VirtualMachine) + -> Result<(), ArgsError>; + /// Whether the walk has already filled this parameter, which is what + /// decides between a default and what the call passed. + fn is_filled(&self, index: usize) -> bool; +} + +impl ArgSink for ArgsBuilder<'_> { + fn put( + &mut self, + index: usize, + value: &PyObject, + vm: &VirtualMachine, + ) -> Result<(), ArgsError> { + self.set(index, get_jit_value(vm, value)?)?; + Ok(()) + } + + fn is_filled(&self, index: usize) -> bool { + self.is_set(index) + } +} + +/// The types a call's arguments would give a compiled function's parameters, +/// in parameter order. +struct ObservedTypes(Vec>); + +impl ArgSink for ObservedTypes { + fn put( + &mut self, + index: usize, + value: &PyObject, + vm: &VirtualMachine, + ) -> Result<(), ArgsError> { + // Through `get_jit_value` rather than off the class directly, so that + // what is observed as an `int` is exactly what would later convert to + // one - an integer too wide for the machine is not this parameter's + // type, it is a call the compiled code could not have served. + let ty = match get_jit_value(vm, value)? { + AbiValue::Int(_) => JitType::Int, + AbiValue::Float(_) => JitType::Float, + AbiValue::Bool(_) => JitType::Bool, + _ => return Err(ArgsError::NonJitType), + }; + *self.0.get_mut(index).ok_or(ArgsError::WrongNumberOfArgs)? = Some(ty); + Ok(()) + } + + fn is_filled(&self, index: usize) -> bool { + self.0.get(index).is_some_and(Option::is_some) + } +} + +/// Resolve a call's arguments onto the parameters they fill, positional first, +/// then keyword, then the defaults for whatever is left - the order +/// `fill_locals_from_args` resolves them in, and the order a compiled +/// function's slots are laid out in. +/// +/// Unlike `fill_locals_from_args` this raises nothing: a call it turns down +/// goes to the interpreter, which raises whatever the call really deserves. +fn walk_arguments( func: &PyFunction, func_args: &FuncArgs, - jitted_code: &'a CompiledCode, + sink: &mut S, vm: &VirtualMachine, -) -> Result, ArgsError> { - let mut jit_args = jitted_code.args_builder(); +) -> Result<(), ArgsError> { let nargs = func_args.args.len(); let code: &Py = &func.code; @@ -177,7 +238,7 @@ pub(crate) fn get_jit_args<'a>( // Add positional arguments for i in 0..nargs { - jit_args.set(i, get_jit_value(vm, &func_args.args[i])?)?; + sink.put(i, &func_args.args[i], vm)?; } // Handle keyword arguments @@ -188,29 +249,32 @@ pub(crate) fn get_jit_args<'a>( // can never match one. let name = name.as_str().map_err(|_| ArgsError::NotAKeywordArg)?; if let Some(arg_idx) = arg_pos(arg_names.args, name) { - if jit_args.is_set(arg_idx) { + if sink.is_filled(arg_idx) { return Err(ArgsError::ArgPassedMultipleTimes); } - jit_args.set(arg_idx, get_jit_value(vm, value)?)?; + sink.put(arg_idx, value, vm)?; } else if let Some(kwarg_idx) = arg_pos(arg_names.kwonlyargs, name) { let arg_idx = kwarg_idx + arg_count as usize; - if jit_args.is_set(arg_idx) { + if sink.is_filled(arg_idx) { return Err(ArgsError::ArgPassedMultipleTimes); } - jit_args.set(arg_idx, get_jit_value(vm, value)?)?; + sink.put(arg_idx, value, vm)?; } else { return Err(ArgsError::NotAKeywordArg); } } - let (defaults, kwdefaults) = func.defaults_and_kwdefaults.lock().clone(); + // Held rather than cloned: filling a slot from a default reads the object + // but never runs Python code, so nothing can reach the lock again. + let defaults_and_kwdefaults = func.defaults_and_kwdefaults.lock(); + let (defaults, kwdefaults) = &*defaults_and_kwdefaults; // fill in positional defaults if let Some(defaults) = defaults { for (i, default) in defaults.iter().enumerate() { let arg_idx = i + arg_count as usize - defaults.len(); - if !jit_args.is_set(arg_idx) { - jit_args.set(arg_idx, get_jit_value(vm, default)?)?; + if !sink.is_filled(arg_idx) { + sink.put(arg_idx, default, vm)?; } } } @@ -219,15 +283,62 @@ pub(crate) fn get_jit_args<'a>( if let Some(kw_only_defaults) = kwdefaults { for (i, name) in arg_names.kwonlyargs.iter().enumerate() { let arg_idx = i + arg_count as usize; - if !jit_args.is_set(arg_idx) { + if !sink.is_filled(arg_idx) { let default = kw_only_defaults .get_item(&**name, vm) - .map_err(|_| ArgsError::NotAllArgsPassed) - .and_then(|obj| get_jit_value(vm, &obj))?; - jit_args.set(arg_idx, default)?; + .map_err(|_| ArgsError::NotAllArgsPassed)?; + sink.put(arg_idx, &default, vm)?; } } } + Ok(()) +} + +/// The parameter types this call would give the function, for the compiler to +/// specialize on. +/// +/// The automatic path takes its types from here rather than from annotations. +/// Almost nothing outside a type-checked codebase is annotated `int`, `float` +/// or `bool` on every parameter - seven functions in the whole standard +/// library are - while every call carries the types it is actually being made +/// with. A guess that turns out wrong costs a failed conversion and a fall +/// back to the interpreter, never a wrong answer. +#[cfg(feature = "jit")] +pub(super) fn observed_arg_types( + func: &PyFunction, + func_args: &FuncArgs, + vm: &VirtualMachine, +) -> Result, ArgsError> { + let code: &Py = &func.code; + if code + .flags + .intersects(CodeFlags::VARARGS | CodeFlags::VARKEYWORDS) + { + return Err(ArgsError::NonJitType); + } + let slots = code.arg_count as usize + code.arg_names().kwonlyargs.len(); + let mut observed = ObservedTypes(vec![None; slots]); + walk_arguments(func, func_args, &mut observed, vm)?; + observed + .0 + .into_iter() + .collect::>>() + .ok_or(ArgsError::NotAllArgsPassed) +} + +/// Like `fill_locals_from_args` but to populate arguments for calling a jit function. +/// This also doesn't do full error handling but instead return None if anything is wrong. In +/// that case it falls back to the executing the bytecode version which will call +/// `fill_locals_from_args` which will raise the actual exception if needed. +#[cfg(feature = "jit")] +pub(crate) fn get_jit_args<'a>( + func: &PyFunction, + func_args: &FuncArgs, + jitted_code: &'a CompiledCode, + vm: &VirtualMachine, +) -> Result, ArgsError> { + let mut jit_args = jitted_code.args_builder(); + walk_arguments(func, func_args, &mut jit_args, vm)?; jit_args.into_args().ok_or(ArgsError::NotAllArgsPassed) } diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index eef8fd6da4c..017a2b1944a 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -520,7 +520,13 @@ impl LocalsPlus { /// Push a PyObjectRef onto the evaluation stack. /// Panics on overflow. pub(crate) fn push_stack(&mut self, value: PyObjectRef) { - self.stack_try_push(Some(PyStackRef::new_owned(value))) + self.push_stack_opt(Some(value)); + } + + /// Push a value, or the null a call leaves beside its callable. + /// Panics on overflow. + pub(crate) fn push_stack_opt(&mut self, value: Option) { + self.stack_try_push(value.map(PyStackRef::new_owned)) .expect("stack overflow in push_stack"); } @@ -1091,6 +1097,12 @@ impl InterpreterFrame { self.lasti.load(Relaxed) } + /// Set the instruction index the frame runs from next. + #[inline(always)] + pub fn set_lasti(&self, lasti: u32) { + self.lasti.store(lasti, Relaxed); + } + /// Get the previous InterpreterFrame in the chain, or null. #[inline(always)] pub fn previous(&self) -> *const Self { @@ -1176,21 +1188,6 @@ impl InterpreterFrame { top } - /// Create a lightweight FrameObject with empty localsplus, suitable for - /// f_back chain building (retained_back). Unlike `materialize`, this does - /// NOT store into `temporary_refs` or set the `materialized` pointer, so - /// the returned FrameObject is only kept alive by the caller's `PyRef`. - /// This prevents non-GC-tracked `temporary_refs` on a stack-allocated - /// iframe from defeating cycle collection. - #[cold] - #[inline(never)] - pub(crate) fn materialize_chain(&self, vm: &VirtualMachine) -> FrameObjectRef { - if let Some(fo) = self.frame_obj() { - return fo.to_owned(); - } - self.materialize_slow_chain(vm) - } - #[cold] fn materialize_slow(&self, vm: &VirtualMachine) -> &Py { // Create a full FrameObject with its own InterpreterFrame copy. @@ -1285,10 +1282,10 @@ impl InterpreterFrame { unsafe { &*(fo_ptr as *const Py) } } - /// Like `materialize_slow` but with empty localsplus to avoid extra - /// refcounts on local variables. Only suitable for f_back chain building. - /// Returns an owned `PyRef` without storing into `temporary_refs` or - /// setting the `materialized` pointer, so GC can still detect cycles. + /// Like `materialize_slow` but detached: empty localsplus that nothing + /// ever fills, and no store into `temporary_refs` or the `materialized` + /// pointer, so the copy is kept alive only by the returned `PyRef`. + #[cfg(feature = "threading")] #[cold] fn materialize_slow_chain(&self, vm: &VirtualMachine) -> FrameObjectRef { let code: PyRef = self.code().to_owned(); @@ -3081,29 +3078,33 @@ impl ExecutingFrame<'_> { // (frames entered before sys.settrace() have trace=None). // Skip RESUME – it should not generate user-visible line events. if vm.use_tracing.get() - && self.trace_is_set(vm) && !matches!( self.code.instructions.read_op(idx), Instruction::Resume { .. } | Instruction::InstrumentedResume ) && let Some((loc, _)) = self.code.locations.get(idx) - && loc.line.get() as u32 != self.prev_line.get() { - self.prev_line.set(loc.line.get() as u32); - vm.trace_event(crate::protocol::TraceEvent::Line, None)?; - // Trace callback may have changed lasti via set_f_lineno. - // Re-read and restart the loop from the new position. - if self.lasti() != (idx as u32 + 1) { - // set_f_lineno defers stack unwinding because we hold - // the state mutex. Perform it now. - let pops = self.pending_stack_pops(); - if pops > 0 { - let from_stack = self.pending_unwind_from_stack(); - self.unwind_stack_for_lineno(pops as usize, from_stack, vm); - self.set_pending_stack_pops(0); - } - arg_state.reset(); - continue; + // The line is recorded even where this frame carries no trace + // function, so that a frame which starts being traced part way + // through does not report the line it is already on as new. + let line = loc.line.get() as u32; + let changed = line != self.prev_line.replace(line); + if changed && self.trace_is_set(vm) { + vm.trace_event(crate::protocol::TraceEvent::Line, None)?; + // Trace callback may have changed lasti via set_f_lineno. + // Re-read and restart the loop from the new position. + if self.lasti() != (idx as u32 + 1) { + // set_f_lineno defers stack unwinding because we hold + // the state mutex. Perform it now. + let pops = self.pending_stack_pops(); + if pops > 0 { + let from_stack = self.pending_unwind_from_stack(); + self.unwind_stack_for_lineno(pops as usize, from_stack, vm); + self.set_pending_stack_pops(0); + } + arg_state.reset(); + continue; + } } } let op = self.code.instructions.read_op(idx); @@ -3111,24 +3112,6 @@ impl ExecutingFrame<'_> { let mut do_extend_arg = false; let caches = op.cache_entries(); - // Always update prev_line so f_lineno returns the correct line - // even when the frame is observed mid-call (e.g. sys._getframe, - // warnings.warn). The lookup is a simple array index, so the - // cost is negligible. - // Update prev_line for f_lineno. Skip RESUME, ExtendedArg, - // and InstrumentedLine (it manages prev_line in its own handler; - // updating here first would defeat LINE de-duplication). - // Other instrumented opcodes update prev_line via - // execute_instrumented. - if !matches!( - op.into(), - Opcode::Resume | Opcode::ExtendedArg | Opcode::InstrumentedLine - ) && !op.is_instrumented() - && let Some((loc, _)) = self.code.locations.get(idx) - { - self.prev_line.set(loc.line.get() as u32); - } - if vm.use_tracing.get() { // Fire 'opcode' trace event for sys.settrace when f_trace_opcodes // is set. Skip RESUME and ExtendedArg @@ -5948,7 +5931,7 @@ impl ExecutingFrame<'_> { && func.func_version() == cached_version && cached_version != 0 { - if func.is_jitted() { + if func.requires_jit_entry(vm) { return self.execute_call_vectorcall(nargs, vm); } let effective_nargs = nargs + u32::from(self_or_null_is_some); @@ -6013,7 +5996,7 @@ impl ExecutingFrame<'_> { && func.func_version() == cached_version && cached_version != 0 { - if func.is_jitted() { + if func.requires_jit_entry(vm) { return self.execute_call_vectorcall(nargs, vm); } if !func.has_exact_argcount(nargs + 1) { @@ -6237,7 +6220,7 @@ impl ExecutingFrame<'_> { && func.func_version() == cached_version && cached_version != 0 { - if func.is_jitted() { + if func.requires_jit_entry(vm) { return self.execute_call_vectorcall(nargs, vm); } if self.specialization_call_recursion_guard(vm) { @@ -6276,7 +6259,7 @@ impl ExecutingFrame<'_> { && func.func_version() == cached_version && cached_version != 0 { - if func.is_jitted() { + if func.requires_jit_entry(vm) { return self.execute_call_vectorcall(nargs, vm); } if self.specialization_call_recursion_guard(vm) { @@ -6637,7 +6620,7 @@ impl ExecutingFrame<'_> { && func.func_version() == cached_version && cached_version != 0 { - if func.is_jitted() { + if func.requires_jit_entry(vm) { return self.execute_call_kw_vectorcall(nargs, vm); } if self.specialization_call_recursion_guard(vm) { @@ -6698,7 +6681,7 @@ impl ExecutingFrame<'_> { && func.func_version() == cached_version && cached_version != 0 { - if func.is_jitted() { + if func.requires_jit_entry(vm) { return self.execute_call_kw_vectorcall(nargs, vm); } let nargs_usize = nargs as usize; @@ -7297,20 +7280,6 @@ impl ExecutingFrame<'_> { instruction.is_instrumented(), "execute_instrumented called with non-instrumented opcode {instruction:?}" ); - // Update prev_line for f_lineno. The main bytecode loop skips - // instrumented opcodes to avoid interfering with LINE event - // de-duplication in InstrumentedLine. Update here instead, except - // for RESUME (prev_line must stay 0 for the first LINE event) and - // InstrumentedLine (manages prev_line in its own handler). - if !matches!( - instruction, - Instruction::InstrumentedResume | Instruction::InstrumentedLine - ) { - let idx = self.lasti() as usize - 1; - if let Some((loc, _)) = self.code.locations.get(idx) { - self.prev_line.set(loc.line.get() as u32); - } - } self.monitoring_mask = vm.state.monitoring_events.load(); match instruction { Instruction::InstrumentedResume => { @@ -7625,12 +7594,6 @@ impl ExecutingFrame<'_> { monitoring::fire_instruction(vm, self.code, offset)?; } - // Update prev_line for f_lineno since the bytecode loop's - // update skips all instrumented opcodes. - if let Some((loc, _)) = self.code.locations.get(idx) { - self.prev_line.set(loc.line.get() as u32); - } - // Re-dispatch to the real original opcode let original_op = Instruction::try_from(real_op_byte) .expect("invalid opcode in side-table chain"); @@ -9957,7 +9920,7 @@ impl ExecutingFrame<'_> { let callable = self.nth_value(nargs + 1); if let Some(func) = callable.downcast_ref_if_exact::(vm) { - if self.specialization_eval_frame_active(vm) || func.is_jitted() { + if self.specialization_eval_frame_active(vm) || func.requires_jit_entry(vm) { unsafe { self.code.instructions.write_adaptive_counter( cache_base, @@ -10020,7 +9983,7 @@ impl ExecutingFrame<'_> { .function_obj() .downcast_ref_if_exact::(vm) { - if self.specialization_eval_frame_active(vm) || func.is_jitted() { + if self.specialization_eval_frame_active(vm) || func.requires_jit_entry(vm) { unsafe { self.code.instructions.write_adaptive_counter( cache_base, @@ -10315,7 +10278,7 @@ impl ExecutingFrame<'_> { let callable = self.nth_value(nargs + 2); if let Some(func) = callable.downcast_ref_if_exact::(vm) { - if self.specialization_eval_frame_active(vm) || func.is_jitted() { + if self.specialization_eval_frame_active(vm) || func.requires_jit_entry(vm) { unsafe { self.code.instructions.write_adaptive_counter( cache_base, @@ -10366,7 +10329,7 @@ impl ExecutingFrame<'_> { .function_obj() .downcast_ref_if_exact::(vm) { - if self.specialization_eval_frame_active(vm) || func.is_jitted() { + if self.specialization_eval_frame_active(vm) || func.requires_jit_entry(vm) { unsafe { self.code.instructions.write_adaptive_counter( cache_base, diff --git a/crates/vm/src/signal.rs b/crates/vm/src/signal.rs index d1ddd8c3cdd..a46316a0f57 100644 --- a/crates/vm/src/signal.rs +++ b/crates/vm/src/signal.rs @@ -28,6 +28,27 @@ const QSBR_BIT: u8 = 1 << 1; /// allocation that tripped the threshold. #[cfg(feature = "threading")] const GC_BIT: u8 = 1 << 2; +/// A stop-the-world is in progress, so some thread has been asked to park. +#[cfg(feature = "threading")] +const STOP_BIT: u8 = 1 << 3; +/// An interpreter is shutting down, so every thread but the one driving that +/// shutdown must stop running bytecode. +#[cfg(feature = "threading")] +const FINALIZING_BIT: u8 = 1 << 4; + +/// Stop-the-world spans currently open. A single bit could not be cleared +/// safely: two interpreters can stop their own worlds at once, and whichever +/// finished first would clear the bit out from under the other. +#[cfg(feature = "threading")] +static STOP_REQUESTS: core::sync::atomic::AtomicUsize = core::sync::atomic::AtomicUsize::new(0); + +/// Shutdowns currently in progress, counted for the same reason as the +/// stop-the-world spans: a subinterpreter is finalized while the interpreter +/// that owns it is finalizing, and the word is shared by every interpreter in +/// the process. An interpreter that leaves the bit behind slows every one that +/// outlives it down to the per-instruction slow path. +#[cfg(feature = "threading")] +static FINALIZE_REQUESTS: core::sync::atomic::AtomicUsize = core::sync::atomic::AtomicUsize::new(0); #[expect( clippy::declare_interior_mutable_const, @@ -147,6 +168,61 @@ pub(crate) fn qsbr_bit_set() -> bool { EVAL_BREAKER.load(Ordering::Relaxed) & QSBR_BIT != 0 } +/// Open a stop-the-world span, so the word says a thread may have to park. +#[cfg(feature = "threading")] +pub(crate) fn begin_stop_request() { + if STOP_REQUESTS.fetch_add(1, Ordering::Release) == 0 { + EVAL_BREAKER.fetch_or(STOP_BIT, Ordering::Release); + } +} + +/// Close one, clearing the bit once the last one is closed. +#[cfg(feature = "threading")] +pub(crate) fn end_stop_request() { + if STOP_REQUESTS.fetch_sub(1, Ordering::Release) == 1 { + EVAL_BREAKER.fetch_and(!STOP_BIT, Ordering::Release); + } +} + +/// Close every span at once. Only a fork child, where the threads whose spans +/// they were do not exist any more, has spans nobody will ever close. +#[cfg(feature = "threading")] +pub(crate) fn reset_stop_requests() { + STOP_REQUESTS.store(0, Ordering::Release); + EVAL_BREAKER.fetch_and(!STOP_BIT, Ordering::Release); +} + +/// Open a shutdown span, so the word says a thread may have to stop. +#[cfg(feature = "threading")] +pub(crate) fn begin_finalize_request() { + if FINALIZE_REQUESTS.fetch_add(1, Ordering::Release) == 0 { + EVAL_BREAKER.fetch_or(FINALIZING_BIT, Ordering::Release); + } +} + +/// Close one, clearing the bit once the last one is closed. The threads the +/// span was opened for belong to an interpreter that no longer exists by the +/// time this runs. +#[cfg(feature = "threading")] +pub(crate) fn end_finalize_request() { + if FINALIZE_REQUESTS.fetch_sub(1, Ordering::Release) == 1 { + EVAL_BREAKER.fetch_and(!FINALIZING_BIT, Ordering::Release); + } +} + +/// The word compiled code polls at every backward jump. +/// +/// Every reason a thread has to leave the bytecode loop sets a bit here before +/// it becomes true of any one thread, so a zero word means no thread has to +/// leave. The converse does not hold - the bits say nothing about *which* +/// thread - which is why a poll that finds the word non-zero hands the frame +/// back to the interpreter to make the per-thread checks +/// `VirtualMachine::eval_breaker_tripped` makes. +#[cfg(feature = "jit")] +pub(crate) fn eval_breaker_word() -> &'static AtomicU8 { + &EVAL_BREAKER +} + /// Schedule an automatic collection to run at the next bytecode safepoint. #[cfg(feature = "threading")] pub(crate) fn schedule_gc() { diff --git a/crates/vm/src/stdlib/sys.rs b/crates/vm/src/stdlib/sys.rs index 85488656a2a..10f8b163ac8 100644 --- a/crates/vm/src/stdlib/sys.rs +++ b/crates/vm/src/stdlib/sys.rs @@ -9,25 +9,53 @@ pub(crate) use sys::{DOC, MAXSIZE, RUST_MULTIARCH, UnraisableHookArgsData, modul #[pymodule(name = "_jit")] mod sys_jit { + use crate::VirtualMachine; + /// Return True if the current Python executable supports JIT compilation, /// and False otherwise. #[pyfunction] const fn is_available() -> bool { - false // RustPython has no JIT + // The automatic compiler, which is what `is_enabled` reports and what + // `PYTHON_JIT` switches. Explicit `__jit__()` needs only the `jit` + // feature and is not what this pair describes. + cfg!(feature = "aot") } /// Return True if JIT compilation is enabled for the current Python process, /// and False otherwise. #[pyfunction] - const fn is_enabled() -> bool { - false // RustPython has no JIT + fn is_enabled(vm: &VirtualMachine) -> bool { + vm.state.config.settings.aot } /// Return True if the topmost Python frame is currently executing JIT code, /// and False otherwise. #[pyfunction] const fn is_active() -> bool { - false // RustPython has no JIT + // Compiled code runs without a Python frame and cannot call back into + // the interpreter, so no Python code can be running while it is. + false + } + + /// Return `(compiled, rejected, deoptimized)` for the functions the + /// automatic compiler has looked at. RustPython dialect. + #[pyfunction] + fn _stats(vm: &VirtualMachine) -> (u64, u64, u64) { + #[cfg(feature = "jit")] + { + use core::sync::atomic::Ordering::Relaxed; + let stats = &vm.state.aot_stats; + ( + stats.compiled.load(Relaxed), + stats.rejected.load(Relaxed), + stats.deoptimized.load(Relaxed), + ) + } + #[cfg(not(feature = "jit"))] + { + let _ = vm; + (0, 0, 0) + } } } diff --git a/crates/vm/src/vm/interpreter.rs b/crates/vm/src/vm/interpreter.rs index db806e0673e..5f38a72290d 100644 --- a/crates/vm/src/vm/interpreter.rs +++ b/crates/vm/src/vm/interpreter.rs @@ -162,6 +162,10 @@ where // Create PyGlobalState (≈ PyInterpreterState) let global_state = PyRc::new(PyGlobalState { gc: crate::gc_state::GcInterpreterState::new(&ctx), + #[cfg(feature = "jit")] + jit_engine: rustpython_jit::JitEngine::new(Some(crate::signal::eval_breaker_word())), + #[cfg(feature = "jit")] + aot_stats: Default::default(), interpreter_id, runtime_root_id, whence, @@ -693,6 +697,16 @@ impl Interpreter { // Now suppress unraisable exceptions from daemon threads and __del__ // methods during the rest of shutdown. vm.state.finalizing.store(true, Ordering::Release); + // Compiled code cannot read that flag, and a daemon thread inside a + // compiled loop would otherwise never leave it - shutdown would + // wait on a thread that had already been told to stop. The span + // ends with this function: the word is shared by every interpreter + // in the process, and one left set costs every interpreter that + // outlives this one the per-instruction slow path. + #[cfg(feature = "threading")] + crate::signal::begin_finalize_request(); + #[cfg(feature = "threading")] + scopeguard::defer! { crate::signal::end_finalize_request(); } // GC pass - collect cycles before module cleanup vm.state.gc.collect_force(2); diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index 8a47514c3e0..2c752cf3af1 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -279,6 +279,10 @@ impl StopTheWorldState { // registration (which also takes this lock), matching the // HEAD_LOCK-guarded stop-the-world bookkeeping. self.requested.store(true, Ordering::Release); + // Announce the span in the eval-breaker word before any thread's own + // stop bit is set, so no thread can be asked to park while the word + // still reads as nothing to do. + crate::signal::begin_stop_request(); let count = registry .keys() .filter(|&&thread_id| thread_id != requester) @@ -538,6 +542,9 @@ impl StopTheWorldState { drop(registry); self.thread_countdown.store(0, Ordering::Release); self.requester.store(0, Ordering::Relaxed); + // Every stop bit this span set has been cleared above, so the word can + // stop saying there is one. + crate::signal::end_stop_request(); #[cfg(debug_assertions)] self.debug_assert_all_non_requester_detached(state); // Release the exclusion last, ending the stop→start span so the next @@ -549,6 +556,7 @@ impl StopTheWorldState { /// Reset after fork in the child (only one thread alive). pub fn reset_after_fork(&self) { self.requested.store(false, Ordering::Relaxed); + crate::signal::reset_stop_requests(); self.world_stopped.store(false, Ordering::Relaxed); crate::common::lock::set_world_stopped(false); self.requester.store(0, Ordering::Relaxed); @@ -828,6 +836,12 @@ pub struct PyGlobalState { pub id_refcount: AtomicI64, /// When true, dropping the last ID ref destroys the interpreter. pub require_idref: AtomicBool, + /// Owns the machine code of every function compiled in this interpreter. + #[cfg(feature = "jit")] + pub jit_engine: alloc::sync::Arc, + /// What the AOT path has compiled, rejected, and given back. + #[cfg(feature = "jit")] + pub aot_stats: builtins::function::aot::AotStats, } impl PyGlobalState { @@ -2419,8 +2433,13 @@ impl VirtualMachine { // the frame object, so it is now readable from anywhere. fo.iframe().detach(); if !old_chain.is_null() { + // The frame object has to be attached to the caller, not + // a standalone copy of it: the caller runs this same + // block when it returns and finds only an attached one, + // so a copy would end the chain here and `f_back` would + // stop one link up from every escaped frame. let prev_iframe = unsafe { &*old_chain }; - let back_fo = prev_iframe.materialize_chain(self); + let back_fo = prev_iframe.materialize(self).to_owned(); *fo.iframe().cold().retained_back.lock() = Some(back_fo); } fo.iframe().owner.store( diff --git a/crates/vm/src/vm/setting.rs b/crates/vm/src/vm/setting.rs index 3c42ca0b6fc..c6b495dee73 100644 --- a/crates/vm/src/vm/setting.rs +++ b/crates/vm/src/vm/setting.rs @@ -64,6 +64,11 @@ pub struct Settings { // int import_time; /// -X no_debug_ranges: disable column info in bytecode pub code_debug_ranges: bool, + + /// -X aot, RUSTPYTHON_AOT: compile eligible functions to native code on + /// their first call instead of waiting for an explicit `__jit__()`. + /// Always false unless built with the `aot` feature. + pub aot: bool, // int show_ref_count; // int dump_refs; // wchar_t *dump_refs_file; @@ -210,6 +215,7 @@ impl Default for Settings { hash_seed: None, faulthandler: false, code_debug_ranges: true, + aot: cfg!(feature = "aot"), buffered_stdio: true, check_hash_pycs_mode: CheckHashPycsMode::Default, allow_external_library: cfg!(feature = "importlib"), diff --git a/crates/vm/src/warn.rs b/crates/vm/src/warn.rs index d5729858dcd..db2a861aa32 100644 --- a/crates/vm/src/warn.rs +++ b/crates/vm/src/warn.rs @@ -581,7 +581,7 @@ fn setup_context( ( f.iframe().globals().to_owned(), f.iframe().code().source_path(), - f.f_lineno(), + f.f_lineno(vm), ) } else if let Some(frame) = vm.current_frame() { // We have a frame but it wasn't found during stack walking diff --git a/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py new file mode 100644 index 00000000000..2d9d43b612c --- /dev/null +++ b/extra_tests/snippets/aot.py @@ -0,0 +1,268 @@ +# Automatic compilation must be invisible: every assertion below has to hold +# whether or not functions were compiled behind our back. +import sys + +# `_stats` is a RustPython addition, so this stays False under CPython even on a +# build whose own JIT is enabled. +AOT = sys._jit.is_enabled() and hasattr(sys._jit, "_stats") + +# The automatic path waits for a function to be called often enough to be worth +# compiling, and takes the types it specializes on from the call that crosses +# that line. So a function has to be warmed with the arguments the assertion +# below it means to test, or the assertion runs against the interpreter and +# proves nothing about compiled code. Comfortably more than the threshold: if +# it ever rises past this, the stat floors at the bottom say so. +WARMUP = 200 + + +def warm(f, *args): + for _ in range(WARMUP): + f(*args) + return f + + +def scale(a: float, b: float) -> float: + return a * b + a - b + + +def wide(a: int, b: int) -> int: + return a + b + + +def untyped(a, b): + return a + b + + +def variadic(*args) -> int: + total = 0 + for a in args: + total = total + a + return total + + +def closure_factory(n: int): + def inner(x: int) -> int: + return x + n + + return inner + + +def generator(n: int): + yield n + + +def guarded(a: int) -> int: + try: + return a + except ValueError: + return 0 + + +warm(scale, 2.0, 3.0) +assert scale(2.0, 3.0) == 5.0 + +# An integer where the warm-up saw a float does not fit the compiled signature. +# The call has to fall back rather than reinterpret the bits. +assert scale(2, 3.0) == 5.0 +assert scale(2.0, 3.0) == 5.0 + +# Integer arithmetic widens once the machine word can no longer hold the +# answer. `wide` compiles under the automatic path too now, and deoptimizes +# back to the interpreter exactly where that widening has to happen. +warm(wide, 3, 4) +assert wide(3, 4) == 7 +assert wide(2**62, 2**62) == 2**63 + + +# The guard above deoptimizes `wide`, and a deopt discards the compiled +# code and leaves the function permanently interpreted - so a third +# widening case on `wide` itself would never touch compiled code again. +# A fresh function gets its own compile attempt for it. +def wide2(a: int, b: int) -> int: + return a + b + + +warm(wide2, 3, 4) +assert wide2(-(2**63), -(2**63)) == -(2**64) + +# Shapes with no compiled form at all still behave, warm or cold. +assert untyped("a", "b") == "ab" +assert untyped(1, 2) == 3 +assert warm(variadic, 1, 2, 3)(1, 2, 3) == 6 +assert warm(closure_factory(5), 1)(1) == 6 +assert list(generator(7)) == [7] +assert warm(guarded, 3)(3) == 3 + +# A zero operand is ordinary for `scale`, which only multiplies and adds. +assert scale(1.0, 0.0) == 1.0 + + +def divide(a: float, b: float) -> float: + return a / b + + +warm(divide, 1.0, 2.0) +assert divide(1.0, 2.0) == 0.5 +try: + divide(1.0, 0.0) +except ZeroDivisionError: + pass +else: + raise AssertionError("expected ZeroDivisionError") + + +# A compiled self-call resolves the global by name, but the interpreter reads +# the globals dict on every call. Rebinding the name has to be observable. +def countdown(a: float) -> float: + if a > 0.0: + return countdown(a - 1.0) + return a + + +warm(countdown, 3.0) +assert countdown(3.0) == 0.0 +original_countdown = countdown + + +def countdown(a: float) -> float: + return -1.0 + + +warm(countdown, 3.0) +assert original_countdown(3.0) == -1.0 + + +def fib_iter(n: int) -> int: + a = 0 + b = 1 + i = 0 + while i < n: + a, b = b, a + b + i = i + 1 + return a + + +warm(fib_iter, 10) +assert fib_iter(10) == 55 +# Overflowing partway through is where the compiled loop hands back to the +# interpreter, which finishes it as a bignum. +assert fib_iter(95) == 31940434634990099905 + + +# Nothing the automatic path does runs Python code, annotations included. +# Under PEP 649 reading them means calling `__annotate__`, and a program that +# never asked for its own annotations must not have them evaluated behind its +# back - here that would raise, since the name is not defined yet. +def forward(a: NotDefinedYet) -> int: + return 1 + + +warm(forward, 1) +assert forward(1) == 1 +try: + forward.__annotations__ +except NameError: + pass +else: + raise AssertionError("expected the forward reference to still be unresolved") + + +if AOT: + # The same thing where evaluating the annotation would be unmissable. + class Boom(BaseException): + pass + + def explode(): + raise Boom + + def annotated(a: explode()) -> int: + return 1 + + warm(annotated, 1) + assert annotated(1) == 1 + try: + annotated.__annotations__ + except Boom: + pass + else: + raise AssertionError("expected the annotation to still be unevaluated") + + +# A frame that outlives its call has to keep resolving `f_back` past its +# immediate caller. Each frame on the way back gets a frame object only +# because the one below it returned and asked for one, so a link that stops +# after the first hop hides the entire stack behind it. +def innermost(): + return sys._getframe() + + +def middle(): + return innermost() + + +def outermost(): + return middle() + + +walked = [] +frame = outermost() +while frame is not None: + walked.append(frame.f_code.co_name) + frame = frame.f_back +assert walked[:4] == ["innermost", "middle", "outermost", ""], walked + + +# A compiled loop has to be leaveable. Native code that polls for nothing can +# never park for a stop-the-world, so a thread inside one holds up every +# collection the rest of the process asks for - not for a while, but for good. +# There is nothing to assert: the collection below either returns or it does +# not. +import gc +import threading +import time + + +def spin(n: int) -> int: + i = 0 + while i < n: + i = i + 1 + return i + + +# Warm on a count that returns at once, so the thread below enters a loop that +# is already compiled - which is the whole point of the check. +warm(spin, 0) + +spinning = threading.Event() + + +def keep_spinning(): + spinning.set() + # Further than this script will ever get. The thread is a daemon and is + # meant to be left exactly where it is. + spin(1 << 62) + + +threading.Thread(target=keep_spinning, daemon=True).start() +spinning.wait() +# Long enough that the thread is inside the loop rather than on its way in. +time.sleep(0.1) +gc.collect() + + +if AOT: + compiled, rejected, deoptimized = sys._jit._stats() + # `scale`, `wide`, `wide2`, `divide`, `fib_iter`, `spin`, `forward`, + # `annotated`, `warm` itself and the rebound `countdown` are what the + # automatic path takes above - the original self-recursive `countdown`, + # kept as `original_countdown`, is refused. These are floors, not exact + # counts, but they must not regress: a floor already met before a change + # cannot tell whether the gate came back. + assert compiled >= 7, (compiled, rejected, deoptimized) + # ... the int argument in `scale(2, 3.0)` handed it back, and + # `fib_iter(95)` overflows partway through its loop. + assert deoptimized >= 5, (compiled, rejected, deoptimized) + # Everything else was turned down rather than mis-compiled. `rejected` + # stays a loose floor: it moves with the feature set of the binary. + assert rejected >= 1, (compiled, rejected, deoptimized) + print("aot: compiled", compiled, "rejected", rejected, "deopt", deoptimized) diff --git a/extra_tests/snippets/jit.py b/extra_tests/snippets/jit.py index 887cbb50e7e..d6415dc3133 100644 --- a/extra_tests/snippets/jit.py +++ b/extra_tests/snippets/jit.py @@ -1,3 +1,6 @@ +import sys + + def foo(): a = 5 return 10 + a @@ -27,3 +30,322 @@ def tests(): bar.__jit__() baz.__jit__() tests() + + # A sum that stops fitting in a machine word is not an error: the compiled + # code hands its operands back and the interpreter answers with a bignum. + def add(a: int, b: int) -> int: + return a + b + + add.__jit__() + assert add(1, 2) == 3 + assert add(2**62, 2**62) == 2**63 + + # A negative exponent and an answer too wide for a machine word both + # send the operands back to the interpreter rather than answering 0 or + # wrapping. + def ipow(a: int, b: int) -> int: + return a**b + + ipow.__jit__() + assert ipow(2, 10) == 1024 + assert ipow(2, -2) == 0.25 + assert ipow(2, 33) == 8589934592 + + # A huge-magnitude exponent used to trap the process outright instead of + # deoptimizing; a finite base whose power overflows a double, a + # zero base with a negative exponent, and a negative zero base all + # deoptimize too, and the interpreter raises or answers exactly the way + # it would without the JIT. + def fpow(a: float, b: float) -> float: + return a**b + + fpow.__jit__() + assert fpow(2.0, 3.0) == 8.0 + assert fpow(-8.0, 2.0) == 64.0 + for base, exponent in [(2.0, 1e300), (-2.0, 1e300), (1e308, 2.0)]: + try: + fpow(base, exponent) + raise AssertionError("expected OverflowError") + except OverflowError: + pass + try: + fpow(0.0, -1.0) + raise AssertionError("expected ZeroDivisionError") + except ZeroDivisionError: + pass + assert str(fpow(-0.0, 3.0)) == "-0.0" + + # Division by zero raises rather than returning an infinity. + def fdiv(a: float, b: float) -> float: + return a / b + + fdiv.__jit__() + assert fdiv(4.0, 2.0) == 2.0 + try: + fdiv(1.0, 0.0) + raise AssertionError("expected ZeroDivisionError") + except ZeroDivisionError: + pass + + # A guard that fires with several locals live and a partial expression on + # the stack. The interpreter picks the addition up again with `total`, `n` + # and `step` as the guard saw them and `n * step` already computed, so a + # dropped local or a mis-ordered stack shows up as a wrong sum or an + # UnboundLocalError rather than as a crash. + def mixed(n: int, step: int) -> int: + total = 0 + while n > 0: + total = total + n * step + n = n - 1 + return total + + mixed.__jit__() + assert mixed(5, 1) == 15 + # `total` reaches 5 * 2**60 on the first iteration, which fits, and the + # second adds 2**62 to it, which does not - so the guard fires two + # iterations in, with the multiply's result already on the stack. + assert mixed(5, 2**60) == 17293822569102704640 + + # A frame that leaves because a *nested* frame gave up has no record of its + # own: the record in the buffer belongs to the frame that wrote it. + # `blow(1)` overflows on its own addition, and `blow(2)` reaches that + # overflow one frame down - continuing the outer frame from the inner + # frame's offset would add 2**62 to itself and answer 2**63. + def blow(n: int) -> int: + if n == 0: + return 4611686018427387904 + return blow(n - 1) + blow(n - 1) + + blow.__jit__() + assert blow(0) == 2**62 + assert blow(2) == 2**64 + + # Depth rather than fidelity: the multiply overflows about forty frames + # down, so forty callers each leave on a nested status instead of their own + # guard. It cannot tell a right resume from a wrong one - `grow` is + # tail-recursive, so what the outermost frame has left to do after the + # guard is exactly what the deepest frame has left to do, and both answer + # 3**50. `blow` above is what pins whose record gets honoured. + def grow(n: int, acc: int) -> int: + if n < 1: + return acc + return grow(n - 1, acc * 3) + + grow.__jit__() + assert grow(50, 1) == 717897987691852588770249 + + # A guard lists the locals the compiler had seen when it was lowered. + # `extra` is stored further down the loop body, so the guard on the sum + # cannot describe it - yet the backward jump means it is already bound by + # the time that guard fires on the second iteration. A site that cannot + # describe every bound local restarts the call instead of resuming + # without it. + # + # This reads the dropped local out of a traceback rather than out of the + # function, and it has to. The obvious version - `return total + extra` - + # cannot fail, because a read the compiler cannot prove bound compiles to + # LOAD_FAST_CHECK, which has no lowering, so no function that would + # observe the drop that way compiles in the first place. What does still + # see it is anything reading the frame's fastlocals other than a + # LOAD_FAST: `f_locals` here, a tracer stepping the resumed frame, a + # debugger stopped in it. + def late(n: int, step: int) -> int: + total = 0 + while n > 0: + total = total + n * step + if n == 5: + extra = 1 + n = n - 1 + return total // n + + late.__jit__() + try: + late(5, 2**60) + raise AssertionError("expected ZeroDivisionError") + except ZeroDivisionError as exc: + frame = exc.__traceback__ + while frame.tb_next is not None: + frame = frame.tb_next + locals_at_raise = frame.tb_frame.f_locals + assert locals_at_raise["total"] == 17293822569102704640, locals_at_raise + assert locals_at_raise["extra"] == 1, locals_at_raise + + # The point of all of the above: compiling a function must not change + # what it answers. Each `check` execs its own pair of functions so that + # `__jit__()` on the compiled one cannot affect the interpreted one. + def check(source, *args): + """Assert that compiling a function does not change what it answers.""" + interpreted_scope = {} + exec(source, interpreted_scope) + compiled_scope = {} + exec(source, compiled_scope) + compiled_scope["f"].__jit__() + + def call(f): + try: + return ("value", f(*args)) + except Exception as e: + return ("raised", type(e)) + + expected = call(interpreted_scope["f"]) + actual = call(compiled_scope["f"]) + assert expected == actual, (source.strip(), args, expected, actual) + if expected[0] == "value": + assert type(expected[1]) is type(actual[1]), ( + source.strip(), + args, + type(expected[1]), + type(actual[1]), + ) + + FLOOR = "def f(a: int, b: int) -> int:\n return a // b\n" + check(FLOOR, -7, 2) + check(FLOOR, 7, -2) + check(FLOOR, -7, -2) + check(FLOOR, 1, 0) + check(FLOOR, -(2**63), -1) + + MOD = "def f(a: int, b: int) -> int:\n return a % b\n" + check(MOD, -7, 2) + check(MOD, 7, -2) + check(MOD, 1, 0) + + MUL = "def f(a: int, b: int) -> int:\n return a * b\n" + check(MUL, 2**62, 4) + check(MUL, 3, 4) + + IPOW = "def f(a: int, b: int) -> int:\n return a ** b\n" + check(IPOW, 2, -2) + check(IPOW, 2, 64) + check(IPOW, 2, 10) + + DIV = "def f(a: int, b: int) -> float:\n return a / b\n" + check(DIV, 1, 0) + check(DIV, 2**60 + 1, 3) + check(DIV, 7, 2) + + FDIV = "def f(a: float, b: float) -> float:\n return a / b\n" + check(FDIV, 1.0, 0.0) + check(FDIV, 4.0, 2.0) + + FPOW = "def f(a: float, b: float) -> float:\n return a ** b\n" + check(FPOW, -8.0, 0.5) + check(FPOW, 0.0, -1.0) + check(FPOW, -8.0, 2.0) + + SHIFT = "def f(a: int, b: int) -> int:\n return a << b\n" + check(SHIFT, 1, 62) + check(SHIFT, 1, 63) + check(SHIFT, 1, 64) + check(SHIFT, 1, -1) + + RSHIFT = "def f(a: int, b: int) -> int:\n return a >> b\n" + check(RSHIFT, -8, 1) + check(RSHIFT, 1, 64) + check(RSHIFT, 1, -1) + + ADD = "def f(a: int, b: int) -> int:\n return a + b\n" + check(ADD, 2**62, 2**62) + check(ADD, 7, 3) + + # `Subtract`'s own arm calls `compile_sub(a, b, ...)` in call-site order; `NEG` + # below reaches the same helper through a separate arm, `compile_sub(zero, a, + # ...)`, with different operand order and arity. Covering both closes the gap + # a fix in one arm and not the other would leave open. + SUB = "def f(a: int, b: int) -> int:\n return a - b\n" + check(SUB, -(2**63), 1) + check(SUB, 7, 3) + + NEG = "def f(a: int) -> int:\n return -a\n" + check(NEG, 7) + check(NEG, -(2**63)) + + # `fib_iter(95)` overflows a signed 64-bit accumulator partway through, + # so the compiled run deoptimizes mid-loop and the interpreter finishes + # it. Before this, the same call killed the process with SIGILL. + def fib_iter(n: int) -> int: + a = 0 + b = 1 + i = 0 + while i < n: + a, b = b, a + b + i = i + 1 + return a + + assert fib_iter(95) == 31940434634990099905 + fib_iter.__jit__() + assert fib_iter(10) == 55 + assert fib_iter(95) == 31940434634990099905 + + # A self-call compiles only under `Safety::Permissive`, which is what + # `__jit__()` asks for. Automatic compilation uses `Strict`, which turns + # a self-call down because the interpreter re-reads the global on every + # call and rebinding the name has to stay observable (`instructions.rs`, + # the `LoadGlobal` arm). + def fib(n: int) -> int: + if n < 2: + return n + return fib(n - 1) + fib(n - 2) + + fib.__jit__() + assert fib(25) == 75025 + + # A recursive call with more than one argument: the compiler collects + # the arguments by popping, which walks them backwards, so this is the + # shape that catches them arriving in the wrong order. `countdown(3, 1)` + # is asymmetric in its two parameters, so passing them the wrong way + # round gives a different answer rather than the same one. + def countdown(a: int, b: int) -> int: + if a < 1: + return b + return countdown(a - 1, b * 2) + + countdown.__jit__() + assert countdown(3, 1) == 8 + assert countdown(0, 5) == 5 + + # Compiled code runs no frame, so it reports no call, no line and no + # return. A tracer installed while a function is compiled has to send + # the call back to the interpreter, or it observes nothing at all: this + # asserted an empty event list before the tracing test moved above the + # compiled entry. `sys.monitoring` is checked too because it is a + # separate switch that the same entry has to respect. + def traced(a: int, b: int) -> int: + c = a + b + return c * 2 + + traced.__jit__() + assert traced(3, 4) == 14 + + events = [] + + def tracer(frame, event, arg): + if frame.f_code.co_name == "traced": + events.append(event) + return tracer + + sys.settrace(tracer) + try: + assert traced(3, 4) == 14 + finally: + sys.settrace(None) + assert events[:1] == ["call"], events + assert events[-1:] == ["return"], events + + monitored = [] + mon = sys.monitoring + mon.use_tool_id(mon.PROFILER_ID, "jit snippet") + try: + mon.register_callback( + mon.PROFILER_ID, + mon.events.PY_START, + lambda code, offset: monitored.append(code.co_name), + ) + mon.set_events(mon.PROFILER_ID, mon.events.PY_START) + try: + assert traced(3, 4) == 14 + finally: + mon.set_events(mon.PROFILER_ID, 0) + finally: + mon.free_tool_id(mon.PROFILER_ID) + assert "traced" in monitored, monitored diff --git a/src/settings.rs b/src/settings.rs index 2074abd9710..81ce457ffbb 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -264,6 +264,25 @@ pub fn parse_opts() -> Result<(Settings, RunMode), lexopt::Error> { settings.check_hash_pycs_mode = args.check_hash_based_pycs; + // PYTHON_JIT is the spelling CPython gives this switch; RUSTPYTHON_AOT is + // the specific one and wins where both are set. A later -X aot overrides + // either. + for name in ["PYTHON_JIT", "RUSTPYTHON_AOT"] { + let Some(val) = get_env(name) else { continue }; + settings.aot = match val.to_str() { + Some("1") => true, + Some("0") => false, + _ => { + error!( + "Fatal Python error: config_init_aot: \ + {name}=N: N is missing or invalid\n\ + Python runtime state: preinitialized" + ); + std::process::exit(1); + } + }; + } + if let Some(val) = get_env("PYTHONUTF8") && let Some(val_str) = val.to_str() && !val_str.is_empty() @@ -306,6 +325,20 @@ pub fn parse_opts() -> Result<(Settings, RunMode), lexopt::Error> { } }; } + "aot" => { + settings.aot = match value { + None | Some("1") => true, + Some("0") => false, + _ => { + error!( + "Fatal Python error: config_init_aot: \ + -X aot=n: n is missing or invalid\n\ + Python runtime state: preinitialized" + ); + std::process::exit(1); + } + }; + } "no_sig_int" => settings.install_signal_handlers = false, "no_debug_ranges" => settings.code_debug_ranges = false, "int_max_str_digits" => { @@ -342,6 +375,10 @@ pub fn parse_opts() -> Result<(Settings, RunMode), lexopt::Error> { }); settings.xoptions.extend(xopts); + // No switch can turn on a compiler that was not built in, and + // `sys._jit.is_enabled()` reports this field verbatim. + settings.aot &= cfg!(feature = "aot"); + // Resolve utf8_mode if not explicitly set by PYTHONUTF8 or -X utf8. // Default to UTF-8 mode since RustPython's locale encoding detection // is incomplete. Users can set PYTHONUTF8=0 or -X utf8=0 to disable.