From b2eb0c98ccda9639bb21367dc1a6ceb21da3d4c9 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 24 Aug 2026 11:01:15 +0900 Subject: [PATCH 01/57] jit: compile into a shared module owned by JitEngine `compile` built a fresh `JITModule` per function and handed its ownership to `CompiledCode`, which freed the module on drop. Introduce `JitEngine`, which holds one module behind a mutex and hands out `CompiledCode` values that keep the engine alive through an `Arc`; the module memory is freed in `JitEngine::drop`. The free `compile` function stays as a wrapper over a single-use engine. Sharing the module surfaced two states that a per-function module could not reach: - Symbol names came from `obj_name`, which repeats across functions and collided as `Duplicate definition`. Names now carry a per-engine counter. - A rejected function left the codegen context populated and the `FunctionBuilderContext` un-finalized, so the next `FunctionBuilder::new` panicked. `build_function` now clears the codegen context on every path and replaces the builder context when compilation fails. Assisted-by: Claude --- crates/jit/src/lib.rs | 110 ++++++++++++++++++++++++------- crates/jit/tests/common.rs | 30 +++++++-- crates/jit/tests/engine_tests.rs | 70 ++++++++++++++++++++ crates/jit/tests/lib.rs | 3 + 4 files changed, 184 insertions(+), 29 deletions(-) create mode 100644 crates/jit/tests/engine_tests.rs diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 0c700e93cf8..1f4deda2c51 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -3,12 +3,15 @@ mod instructions; extern crate alloc; use alloc::fmt; +use alloc::sync::Arc; use core::mem::ManuallyDrop; +use core::sync::atomic::{AtomicU64, Ordering}; use cranelift::prelude::*; use cranelift_jit::{JITBuilder, JITModule}; use cranelift_module::{FuncId, Linkage, Module, ModuleError}; use instructions::FunctionCompiler; use rustpython_compiler_core::bytecode; +use std::sync::{Mutex, PoisonError}; #[derive(Debug, thiserror::Error)] #[non_exhaustive] @@ -39,7 +42,7 @@ pub enum JitArgumentError { struct Jit { builder_context: FunctionBuilderContext, ctx: codegen::Context, - module: JITModule, + module: ManuallyDrop, } impl Jit { @@ -50,15 +53,37 @@ impl Jit { Self { builder_context: FunctionBuilderContext::new(), ctx: module.make_context(), - module, + 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, + unique: u64, + ) -> Result<(FuncId, JitSig), JitCompileError> { + let result = self.build_function_inner(bytecode, args, ret, unique); + 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, ) -> Result<(FuncId, JitSig), JitCompileError> { for arg in args { let arg = arg.to_cranelift().ok_or(JitCompileError::NotSupported)?; @@ -70,7 +95,7 @@ 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, )?; @@ -101,38 +126,80 @@ impl Jit { self.module.define_function(id, &mut self.ctx)?; - self.module.clear_context(&mut self.ctx); - Ok((id, sig)) } } +/// 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, +} + +impl JitEngine { + #[must_use] + pub fn new() -> Arc { + Arc::new(Self { + jit: Mutex::new(Jit::new()), + next_id: AtomicU64::new(0), + }) + } + + pub fn compile( + self: &Arc, + bytecode: &bytecode::CodeObject, + args: &[JitType], + ret: Option, + ) -> Result { + // 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) = jit.build_function(bytecode, args, ret, unique)?; + jit.module.finalize_definitions()?; + let code = jit.module.get_finalized_function(id); + drop(jit); + Ok(CompiledCode { + sig, + code, + _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 {} + 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().compile(bytecode, args, ret) } pub struct CompiledCode { sig: JitSig, code: *const u8, - module: ManuallyDrop, + /// Keeps the code memory alive; never read. + _engine: Arc, } impl CompiledCode { + #[must_use] pub fn args_builder(&self) -> ArgsBuilder<'_> { ArgsBuilder::new(self) } @@ -316,13 +383,6 @@ impl UnTypedAbiValue { 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]") diff --git a/crates/jit/tests/common.rs b/crates/jit/tests/common.rs index fe0e4bd33d4..bd4e09bb19d 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}; use rustpython_wtf8::{Wtf8, Wtf8Buf}; use std::collections::HashMap; @@ -14,6 +15,21 @@ 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, + ) -> Result { + let (arg_types, ret_type) = self.signature(); + engine.compile(&self.code, &arg_types, ret_type) + } + + 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 +54,7 @@ impl Function { _ => None, }; - rustpython_jit::compile(&self.code, &arg_types, ret_type).expect("Compile failure") + (arg_types, ret_type) } } @@ -309,7 +325,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 +335,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)*); diff --git a/crates/jit/tests/engine_tests.rs b/crates/jit/tests/engine_tests.rs new file mode 100644 index 00000000000..29b1ad61c91 --- /dev/null +++ b/crates/jit/tests/engine_tests.rs @@ -0,0 +1,70 @@ +#[cfg(test)] +mod tests { + use rustpython_jit::JitEngine; + + /// 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(); + 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).expect("first compile"); + let second = second + .compile_on(&engine) + .expect("second compile of the same name"); + + assert_eq!( + first.invoke(&[3i64.into(), 4i64.into()]), + Ok(Some(3i64.into())) + ); + assert_eq!( + second.invoke(&[3i64.into(), 4i64.into()]), + Ok(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(); + let unsupported = py_function_def!(unsupported => r#" + def unsupported(a: int) -> int: + return [a] + "#); + assert!(unsupported.compile_on(&engine).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) + .expect("engine still usable after a rejected function"); + assert_eq!( + good.invoke(&[3i64.into(), 4i64.into()]), + Ok(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(); + let f = py_function_def!(f => r#" + def f(a: int) -> int: + return a + "#); + f.compile_on(&engine).expect("compile") + }; + assert_eq!(code.invoke(&[7i64.into()]), Ok(Some(7i64.into()))); + } +} diff --git a/crates/jit/tests/lib.rs b/crates/jit/tests/lib.rs index aa5f0f22d64..275924a8ce7 100644 --- a/crates/jit/tests/lib.rs +++ b/crates/jit/tests/lib.rs @@ -1,6 +1,9 @@ +extern crate alloc; + #[macro_use] mod common; mod bool_tests; +mod engine_tests; mod float_tests; mod int_tests; mod misc_tests; From 5d91aa77d31d516a8e8a16bc595fbecdc170b3a5 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 24 Aug 2026 11:04:40 +0900 Subject: [PATCH 02/57] jit: add a Strict safety level that rejects diverging operations Compiled arithmetic does not always answer the way the interpreter does: - Integer `+`, `-` and unary `-` trap on overflow, `//`, `%` and `/` trap on a zero divisor, and `<<`/`>>` trap on a negative count. Nothing installs a trap handler, so each of these aborts the process where the interpreter would widen to a big integer or raise. - Integer `*` is a bare `imul` and wraps silently. - Float `/` is a bare `fdiv`, returning inf instead of raising ZeroDivisionError, and float `**` neither raises for `0.0 ** -1.0` nor produces the complex result Python gives for a negative base. `Safety::Strict` rejects those; `Safety::Permissive` keeps compiling them and stays the behaviour of the free `compile` function and of `__jit__`. Integer bitwise operations, comparisons, float `+ - *`, and int-to-float mixed `+ - *` are faithful and remain available under Strict. Rejection tests assert that Permissive compiles the same source, so they cannot pass on an unrelated compile failure. Assisted-by: Claude --- crates/jit/src/instructions.rs | 64 ++++++++++- crates/jit/src/lib.rs | 22 +++- crates/jit/tests/common.rs | 5 +- crates/jit/tests/engine_tests.rs | 14 ++- crates/jit/tests/lib.rs | 1 + crates/jit/tests/safety_tests.rs | 185 +++++++++++++++++++++++++++++++ 6 files changed, 279 insertions(+), 12 deletions(-) create mode 100644 crates/jit/tests/safety_tests.rs diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index a3c4ca800c4..05e0ff060f4 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -1,5 +1,5 @@ // spell-checker: disable -use super::{JitCompileError, JitSig, JitType}; +use super::{JitCompileError, JitSig, JitType, Safety}; use alloc::collections::BTreeSet; use cranelift::codegen::ir::FuncRef; use cranelift::prelude::*; @@ -72,9 +72,59 @@ pub(crate) struct FunctionCompiler<'a, 'b> { stack: Vec, variables: Box<[Option]>, label_to_block: HashMap, + safety: Safety, pub(crate) sig: JitSig, } +/// Whether the machine code emitted for `op` answers the way the interpreter +/// does for every pair of values of these types. +/// +/// Integer arithmetic either traps on overflow and division by zero - and a +/// trap has no handler, so it takes the process down instead of raising - or +/// wraps where Python would widen to an arbitrary-precision integer. +/// Float `/` is a bare `fdiv` with none of the checks that raise +/// ZeroDivisionError, and float `**` neither raises for `0.0 ** -1.0` nor +/// produces the complex result Python gives a negative base. +fn binary_op_is_faithful(op: BinaryOperator, a: Option<&JitType>, b: Option<&JitType>) -> bool { + let traps_or_wraps_on_ints = matches!( + op, + BinaryOperator::Add + | BinaryOperator::InplaceAdd + | BinaryOperator::Subtract + | BinaryOperator::InplaceSubtract + | BinaryOperator::Multiply + | BinaryOperator::InplaceMultiply + | BinaryOperator::TrueDivide + | BinaryOperator::InplaceTrueDivide + | BinaryOperator::FloorDivide + | BinaryOperator::InplaceFloorDivide + | BinaryOperator::Remainder + | BinaryOperator::InplaceRemainder + | BinaryOperator::Power + | BinaryOperator::InplacePower + | BinaryOperator::Lshift + | BinaryOperator::InplaceLshift + | BinaryOperator::Rshift + | BinaryOperator::InplaceRshift + ); + let diverges_on_floats = matches!( + op, + BinaryOperator::TrueDivide + | BinaryOperator::InplaceTrueDivide + | BinaryOperator::Power + | BinaryOperator::InplacePower + ); + + match (a, b) { + (Some(JitType::Int), Some(JitType::Int)) => !traps_or_wraps_on_ints, + // `(Int, Int)` is taken by the arm above, so it cannot land here. + (Some(JitType::Float | JitType::Int), Some(JitType::Float)) + | (Some(JitType::Float), Some(JitType::Int)) => !diverges_on_floats, + // Any other combination has no lowering at all and is rejected anyway. + _ => true, + } +} + impl<'a, 'b> FunctionCompiler<'a, 'b> { pub(crate) fn new( builder: &'a mut FunctionBuilder<'b>, @@ -82,12 +132,14 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { arg_types: &[JitType], ret_type: Option, entry_block: Block, + safety: Safety, ) -> Self { let mut compiler = Self { builder, stack: Vec::new(), variables: vec![None; num_variables].into_boxed_slice(), label_to_block: HashMap::new(), + safety, sig: JitSig { args: arg_types.to_vec(), ret: ret_type, @@ -373,6 +425,12 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let a_type = a.to_jit_type(); let b_type = b.to_jit_type(); + if self.safety == Safety::Strict + && !binary_op_is_faithful(op, a_type.as_ref(), b_type.as_ref()) + { + return Err(JitCompileError::NotSupported); + } + let val = match (op, a, b) { ( BinaryOperator::Add | BinaryOperator::InplaceAdd, @@ -806,6 +864,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { } Instruction::UnaryNegative => { match self.stack.pop().ok_or(JitCompileError::BadBytecode)? { + // Lowered as `0 - val`, which traps on i64::MIN. + JitValue::Int(_) if self.safety == Safety::Strict => { + Err(JitCompileError::NotSupported) + } JitValue::Int(val) => { // Compile minus as 0 - val. let zero = self.builder.ins().iconst(types::I64, 0); diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 1f4deda2c51..74827b0ea68 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -30,6 +30,18 @@ impl From for JitCompileError { } } +/// How far the compiled code is allowed to diverge from interpreted semantics. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Safety { + /// Reject every operation whose machine code can trap, wrap, or otherwise + /// answer differently from the interpreter. Traps have no handler and kill + /// the process, and a wrapped integer is a silently wrong result, so code + /// that was compiled without being asked for must not reach either. + Strict, + /// Compile everything the backend supports. + Permissive, +} + #[derive(Debug, thiserror::Error, Eq, PartialEq)] #[non_exhaustive] pub enum JitArgumentError { @@ -66,8 +78,9 @@ impl Jit { args: &[JitType], ret: Option, unique: u64, + safety: Safety, ) -> Result<(FuncId, JitSig), JitCompileError> { - let result = self.build_function_inner(bytecode, args, ret, unique); + let result = self.build_function_inner(bytecode, args, ret, unique, safety); self.module.clear_context(&mut self.ctx); if result.is_err() { // Only `FunctionBuilder::finalize` resets the builder context, and @@ -84,6 +97,7 @@ impl Jit { args: &[JitType], ret: Option, unique: u64, + safety: Safety, ) -> Result<(FuncId, JitSig), JitCompileError> { for arg in args { let arg = arg.to_cranelift().ok_or(JitCompileError::NotSupported)?; @@ -114,6 +128,7 @@ impl Jit { args, ret, entry_block, + safety, ); compiler.compile(func_ref, bytecode)?; @@ -152,12 +167,13 @@ impl JitEngine { bytecode: &bytecode::CodeObject, args: &[JitType], ret: Option, + safety: Safety, ) -> Result { // 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) = jit.build_function(bytecode, args, ret, unique)?; + let (id, sig) = jit.build_function(bytecode, args, ret, unique, safety)?; jit.module.finalize_definitions()?; let code = jit.module.get_finalized_function(id); drop(jit); @@ -188,7 +204,7 @@ pub fn compile( args: &[JitType], ret: Option, ) -> Result { - JitEngine::new().compile(bytecode, args, ret) + JitEngine::new().compile(bytecode, args, ret, Safety::Permissive) } pub struct CompiledCode { diff --git a/crates/jit/tests/common.rs b/crates/jit/tests/common.rs index bd4e09bb19d..e0e6d4ee5b2 100644 --- a/crates/jit/tests/common.rs +++ b/crates/jit/tests/common.rs @@ -3,7 +3,7 @@ use core::ops::ControlFlow; use rustpython_compiler_core::bytecode::{ CodeObject, ConstantData, Constants, Instruction, OpArg, OpArgState, }; -use rustpython_jit::{CompiledCode, JitCompileError, JitEngine, JitType}; +use rustpython_jit::{CompiledCode, JitCompileError, JitEngine, JitType, Safety}; use rustpython_wtf8::{Wtf8, Wtf8Buf}; use std::collections::HashMap; @@ -24,9 +24,10 @@ impl Function { 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) + engine.compile(&self.code, &arg_types, ret_type, safety) } fn signature(&self) -> (Vec, Option) { diff --git a/crates/jit/tests/engine_tests.rs b/crates/jit/tests/engine_tests.rs index 29b1ad61c91..bd076861f27 100644 --- a/crates/jit/tests/engine_tests.rs +++ b/crates/jit/tests/engine_tests.rs @@ -1,6 +1,6 @@ #[cfg(test)] mod tests { - use rustpython_jit::JitEngine; + use rustpython_jit::{JitEngine, 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. @@ -16,9 +16,11 @@ mod tests { return b "#); - let first = first.compile_on(&engine).expect("first compile"); + let first = first + .compile_on(&engine, Safety::Permissive) + .expect("first compile"); let second = second - .compile_on(&engine) + .compile_on(&engine, Safety::Permissive) .expect("second compile of the same name"); assert_eq!( @@ -39,14 +41,14 @@ mod tests { def unsupported(a: int) -> int: return [a] "#); - assert!(unsupported.compile_on(&engine).is_err()); + 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) + .compile_on(&engine, Safety::Permissive) .expect("engine still usable after a rejected function"); assert_eq!( good.invoke(&[3i64.into(), 4i64.into()]), @@ -63,7 +65,7 @@ mod tests { def f(a: int) -> int: return a "#); - f.compile_on(&engine).expect("compile") + f.compile_on(&engine, Safety::Permissive).expect("compile") }; assert_eq!(code.invoke(&[7i64.into()]), Ok(Some(7i64.into()))); } diff --git a/crates/jit/tests/lib.rs b/crates/jit/tests/lib.rs index 275924a8ce7..c16b5645ba1 100644 --- a/crates/jit/tests/lib.rs +++ b/crates/jit/tests/lib.rs @@ -8,3 +8,4 @@ mod float_tests; mod int_tests; mod misc_tests; mod none_tests; +mod safety_tests; diff --git a/crates/jit/tests/safety_tests.rs b/crates/jit/tests/safety_tests.rs new file mode 100644 index 00000000000..da320191065 --- /dev/null +++ b/crates/jit/tests/safety_tests.rs @@ -0,0 +1,185 @@ +#[cfg(test)] +mod tests { + use rustpython_jit::{JitEngine, Safety}; + + /// Every operation whose machine code can trap or wrap. There is no trap + /// handler, so a trap kills the process instead of raising. + /// Assert Strict rejects the function *and* that 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(); + 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(); + let f = py_function_def!($name => $src); + f.compile_on(&engine, $safety) + .expect(concat!(stringify!($name), " should compile")) + }}; + } + + #[test] + fn strict_rejects_int_add() { + assert_strict_rejects!(add => r#" +def add(a: int, b: int) -> int: + return a + b +"#); + } + + #[test] + fn strict_rejects_int_multiply() { + assert_strict_rejects!(mul => r#" +def mul(a: int, b: int) -> int: + return a * b +"#); + } + + #[test] + fn strict_rejects_int_floor_divide() { + assert_strict_rejects!(fdiv => r#" +def fdiv(a: int, b: int) -> int: + return a // b +"#); + } + + #[test] + fn strict_rejects_int_true_divide() { + assert_strict_rejects!(true_divide => r#" +def true_divide(a: int, b: int) -> float: + return a / b +"#); + } + + #[test] + fn strict_rejects_int_remainder() { + assert_strict_rejects!(rem => r#" +def rem(a: int, b: int) -> int: + return a % b +"#); + } + + #[test] + fn strict_rejects_int_power() { + assert_strict_rejects!(pow => r#" +def pow(a: int, b: int) -> int: + return a ** b +"#); + } + + #[test] + fn strict_rejects_int_shift() { + assert_strict_rejects!(shift => r#" +def shift(a: int, b: int) -> int: + return a << b +"#); + } + + #[test] + fn strict_rejects_int_negate() { + assert_strict_rejects!(neg => r#" +def neg(a: int) -> int: + return -a +"#); + } + + /// `1.0 / 0.0` raises ZeroDivisionError; `fdiv` returns inf. + #[test] + fn strict_rejects_float_divide() { + assert_strict_rejects!(fdiv => r#" +def fdiv(a: float, b: float) -> float: + return a / b +"#); + } + + /// `(-1.0) ** 0.5` is complex in Python and `0.0 ** -1.0` raises. + #[test] + fn strict_rejects_float_power() { + assert_strict_rejects!(float_power => r#" +def float_power(a: float, b: float) -> float: + return a ** b +"#); + } + + #[test] + fn strict_rejects_mixed_divide() { + assert_strict_rejects!(mixed => r#" +def mixed(a: int, b: float) -> float: + return a / b +"#); + } + + /// 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(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(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(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(Some(2.5f64.into())) + ); + } + + #[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(Some(7i64.into())) + ); + } +} From 31e83df1cf89e9b962bbe82c8882dc94a7f199bb Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 24 Aug 2026 11:06:53 +0900 Subject: [PATCH 03/57] jit: add supports_code as a pre-filter for speculative callers A caller that compiles without being asked needs to rule out hopeless code objects before paying for annotation lookup and codegen setup. `supports_code` makes one pass over the bytecode and rejects the shapes the compiler has no lowering for: varargs, generators and coroutines, a non-empty exception table, cells and frees, and any unsupported opcode. The opcode predicate mirrors the match in `add_instruction` and only has to be right in one direction, which the doc comment records: a wrong "yes" wastes a compile attempt and a wrong "no" costs an optimization, and neither produces wrong code. Assisted-by: Claude --- crates/jit/src/instructions.rs | 48 ++++++++++++++++ crates/jit/src/lib.rs | 34 +++++++++++ crates/jit/tests/common.rs | 5 ++ crates/jit/tests/lib.rs | 1 + crates/jit/tests/support_tests.rs | 96 +++++++++++++++++++++++++++++++ 5 files changed, 184 insertions(+) create mode 100644 crates/jit/tests/support_tests.rs diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 05e0ff060f4..39550e9c0e3 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -76,6 +76,54 @@ pub(crate) struct FunctionCompiler<'a, 'b> { pub(crate) sig: JitSig, } +/// 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 { .. } + ) +} + /// Whether the machine code emitted for `op` answers the way the interpreter /// does for every pair of values of these types. /// diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 74827b0ea68..9b237d6eb5a 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -199,6 +199,40 @@ impl Drop for JitEngine { 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; + } + + let mut state = bytecode::OpArgState::default(); + code.instructions.iter().all(|&word| { + let (instruction, _) = state.get(word); + instructions::instruction_is_supported(instruction) + }) +} + pub fn compile( bytecode: &bytecode::CodeObject, args: &[JitType], diff --git a/crates/jit/tests/common.rs b/crates/jit/tests/common.rs index e0e6d4ee5b2..1a35f2078f7 100644 --- a/crates/jit/tests/common.rs +++ b/crates/jit/tests/common.rs @@ -30,6 +30,11 @@ impl Function { 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 { diff --git a/crates/jit/tests/lib.rs b/crates/jit/tests/lib.rs index c16b5645ba1..8438ef51e61 100644 --- a/crates/jit/tests/lib.rs +++ b/crates/jit/tests/lib.rs @@ -9,3 +9,4 @@ mod int_tests; mod misc_tests; mod none_tests; mod safety_tests; +mod support_tests; diff --git a/crates/jit/tests/support_tests.rs b/crates/jit/tests/support_tests.rs new file mode 100644 index 00000000000..7b8b6fa5213 --- /dev/null +++ b/crates/jit/tests/support_tests.rs @@ -0,0 +1,96 @@ +#[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 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 +"#); + } +} From 0cd8c40944f63a7b13cf5895be2e63aafbe74101 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 24 Aug 2026 11:42:37 +0900 Subject: [PATCH 04/57] vm: add an aot feature that compiles functions on their first call `__jit__()` had to be called by hand. Under the new `aot` feature every function gets one automatic compile attempt the first time it is called, and `__jit__()` stays available. Automatic compilation answers to different rules than a requested one, so it is a separate path: - It compiles with `Safety::Strict`, which turns down anything that can trap or wrap. `__jit__()` still compiles permissively and still raises `JitError` when it cannot; `__jit__(force=True)` compiles again over code a function already has, and retries one the automatic path turned down. - Failure is silent. Reading `__annotations__` runs `__annotate__`, which is Python code that can raise, so the attempt is made only after `supports_code` has ruled out the shapes with no compiled form, and the function claims itself before evaluating anything that can call back into it. - A function whose arguments do not fit its compiled signature is handed back to the interpreter on the first mismatch instead of retrying the conversion on every call. Only automatically compiled code is handed back; `__jit__()` is a standing request. Call state moves from the `jitted_code` mutex to a `jit_state` atomic, so the per-call check is a relaxed load rather than a lock. The frame's specialization sites ask `requires_jit_entry`, which also yields for a function that has not had its attempt yet - otherwise a specialized call would skip the entry point where compilation happens. The bytecode pre-filter verdict is cached on the code object, so it runs once no matter how many functions are built from it. The engine now lives on `PyGlobalState`, so one module holds the code for the whole interpreter instead of one per function. `-X aot=0|1` and `RUSTPYTHON_AOT=0|1` toggle it; the feature sets the default. `sys._jit.is_available()` and `is_enabled()` now report the truth instead of a `false` stub, and `sys._jit._stats()` returns `(compiled, rejected, deoptimized)`. Assisted-by: Claude --- Cargo.toml | 1 + crates/vm/Cargo.toml | 1 + crates/vm/src/builtins/code.rs | 6 ++ crates/vm/src/builtins/function.rs | 115 ++++++++++++++++++++----- crates/vm/src/builtins/function/aot.rs | 90 +++++++++++++++++++ crates/vm/src/frame.rs | 20 ++--- crates/vm/src/stdlib/sys.rs | 33 ++++++- crates/vm/src/vm/interpreter.rs | 4 + crates/vm/src/vm/mod.rs | 6 ++ crates/vm/src/vm/setting.rs | 6 ++ extra_tests/snippets/aot.py | 92 ++++++++++++++++++++ src/settings.rs | 28 ++++++ 12 files changed, 367 insertions(+), 35 deletions(-) create mode 100644 crates/vm/src/builtins/function/aot.rs create mode 100644 extra_tests/snippets/aot.py 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/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/function.rs b/crates/vm/src/builtins/function.rs index b8975cf1102..ab7b740dcae 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, Safety}; fn format_missing_args( qualname: impl core::fmt::Display, @@ -79,6 +83,10 @@ 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, } static FUNC_VERSION_COUNTER: AtomicU32 = AtomicU32::new(1); @@ -218,6 +226,8 @@ 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), }; Ok(func) } @@ -550,20 +560,39 @@ 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. + /// + /// Only for code the AOT path compiled: an explicit `__jit__()` is a + /// standing request, and silently undoing it would be a surprise. + #[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); + } + pub fn invoke_with_locals( &self, func_args: FuncArgs, @@ -571,17 +600,39 @@ impl Py { vm: &VirtualMachine, ) -> PyResult { #[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 state = self.jit_state.load(Relaxed); + if state == aot::UNTRIED && vm.state.config.settings.aot { + state = aot::compile_on_first_call(self, 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(ret)) => { + use crate::convert::ToPyObject; + return Ok(ret.to_pyobject(vm)); + } + 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 - ), } } @@ -934,6 +985,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 +1239,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 +1305,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 +1710,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..97b8f336661 --- /dev/null +++ b/crates/vm/src/builtins/function/aot.rs @@ -0,0 +1,90 @@ +//! Compiling functions to native code without being asked. +//! +//! `__jit__()` compiles one function on request and reports why it could not. +//! The AOT path instead tries every function on its first call, 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}; +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. +#[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; + +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, vm: &VirtualMachine) -> Option { + let code: &Py = &func.code; + if !code_is_eligible(code) { + return None; + } + + // Reading `__annotations__` runs `__annotate__`, which is Python code and + // can raise for a forward reference. Nobody asked for these annotations, so + // whatever it raises is discarded along with the compile attempt. The + // pre-filter above keeps this off the vast majority of functions. + let arg_types = super::jit::get_jit_arg_types(func, vm).ok()?; + let ret_type = super::jit::jit_ret_type(func, vm).ok()?; + + vm.state + .jit_engine + .compile(&code.code, &arg_types, ret_type, Safety::Strict) + .ok() +} + +/// Give `func` its one automatic compile attempt and return its new state. +pub(super) fn compile_on_first_call(func: &Py, vm: &VirtualMachine) -> u8 { + // Claim the function before running anything that can re-enter it: + // evaluating `__annotate__` calls Python, which can reach this same + // function, and a reentrant attempt has to interpret rather than recurse. + func.jit_state.store(REJECTED, Relaxed); + + let Some(compiled) = try_compile(func, 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/frame.rs b/crates/vm/src/frame.rs index eef8fd6da4c..b73bdb49ea8 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -5948,7 +5948,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 +6013,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 +6237,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 +6276,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 +6637,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 +6698,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; @@ -9957,7 +9957,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 +10020,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 +10315,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 +10366,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/stdlib/sys.rs b/crates/vm/src/stdlib/sys.rs index 85488656a2a..c31004a40fb 100644 --- a/crates/vm/src/stdlib/sys.rs +++ b/crates/vm/src/stdlib/sys.rs @@ -9,25 +9,50 @@ 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 + cfg!(feature = "jit") } /// 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..7e16a0ee636 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(), + #[cfg(feature = "jit")] + aot_stats: Default::default(), interpreter_id, runtime_root_id, whence, diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index 8a47514c3e0..50b305bf3c5 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -828,6 +828,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: std::sync::Arc, + /// What the AOT path has compiled, rejected, and given back. + #[cfg(feature = "jit")] + pub aot_stats: builtins::function::aot::AotStats, } impl PyGlobalState { 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/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py new file mode 100644 index 00000000000..08ceab77c1a --- /dev/null +++ b/extra_tests/snippets/aot.py @@ -0,0 +1,92 @@ +# Automatic compilation must be invisible: every assertion below has to hold +# whether or not functions were compiled behind our back. +import sys + + +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 + + +assert scale(2.0, 3.0) == 5.0 + +# An integer where a float was declared 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 must widen. Compiled, `+` traps on overflow and `*` wraps, +# so the automatic path has to leave these alone. +assert wide(3, 4) == 7 +assert wide(2**62, 2**62) == 2**63 +assert wide(-(2**63), -(2**63)) == -(2**64) + +# Shapes with no compiled form at all still behave. +assert untyped("a", "b") == "ab" +assert untyped(1, 2) == 3 +assert variadic(1, 2, 3) == 6 +assert closure_factory(5)(1) == 6 +assert list(generator(7)) == [7] +assert guarded(3) == 3 + +# Division by zero raises rather than returning inf or killing the process. +try: + scale(1.0, 0.0) +except ZeroDivisionError: + raise AssertionError("scale does not divide") + + +def divide(a: float, b: float) -> float: + return a / b + + +assert divide(1.0, 2.0) == 0.5 +try: + divide(1.0, 0.0) +except ZeroDivisionError: + pass +else: + raise AssertionError("expected ZeroDivisionError") + + +if sys._jit.is_enabled(): + compiled, rejected, deoptimized = sys._jit._stats() + # `scale` is the one function above the automatic path can take. + assert compiled >= 1, (compiled, rejected, deoptimized) + # ... and the int argument in `scale(2, 3.0)` handed it back. + assert deoptimized >= 1, (compiled, rejected, deoptimized) + # Everything else was turned down rather than mis-compiled. + assert rejected >= 1, (compiled, rejected, deoptimized) + print("aot: compiled", compiled, "rejected", rejected, "deopt", deoptimized) diff --git a/src/settings.rs b/src/settings.rs index 2074abd9710..b4df5f8fec2 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -306,6 +306,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" => { @@ -355,6 +369,20 @@ pub fn parse_opts() -> Result<(Settings, RunMode), lexopt::Error> { if env_bool("PYTHONNODEBUGRANGES") { settings.code_debug_ranges = false; } + if let Some(val) = get_env("RUSTPYTHON_AOT") { + settings.aot = match val.to_str() { + Some("1") => true, + Some("0") => false, + _ => { + error!( + "Fatal Python error: config_init_aot: \ + RUSTPYTHON_AOT=N: N is missing or invalid\n\ + Python runtime state: preinitialized" + ); + std::process::exit(1); + } + }; + } if let Some(val) = get_env("PYTHON_THREAD_INHERIT_CONTEXT") { settings.thread_inherit_context = match val.to_str() { Some("1") => true, From dcedf971cacc3e4ebfea1729218ba65b8c5ca2ed Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 24 Aug 2026 11:51:42 +0900 Subject: [PATCH 05/57] jit: reject the compiled self-call under Strict `LoadGlobal` resolves the one global it accepts - the function itself - by comparing names. The interpreter reads the globals dict on every call, so once the name is rebound the two disagree about what runs: a decorator applied after the definition, or a test patching the module, and the compiled code keeps calling its old self. Strict now turns down `LoadGlobal`, which costs it self-recursion. `__jit__()` compiles permissively and is unchanged. The snippet asserts a rebound name is observable through a recursive function, so loosening this without a real guard fails the test. Assisted-by: Claude --- crates/jit/src/instructions.rs | 9 ++++++++- crates/jit/tests/safety_tests.rs | 12 ++++++++++++ extra_tests/snippets/aot.py | 19 +++++++++++++++++++ 3 files changed, 39 insertions(+), 1 deletion(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 39550e9c0e3..f27305dedb6 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -809,7 +809,14 @@ 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() { + // The only global that resolves here is this function itself, + // and it resolves by name. 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. + if self.safety == Safety::Strict { + Err(JitCompileError::NotSupported) + } else if name.as_ref() != bytecode.obj_name.as_ref() { Err(JitCompileError::NotSupported) } else { self.stack.push(JitValue::FuncRef(func_ref)); diff --git a/crates/jit/tests/safety_tests.rs b/crates/jit/tests/safety_tests.rs index da320191065..308da6742fb 100644 --- a/crates/jit/tests/safety_tests.rs +++ b/crates/jit/tests/safety_tests.rs @@ -171,6 +171,18 @@ def mixed(a: int, b: float) -> float: ); } + /// The self-reference resolves by name, but the interpreter re-reads the + /// global on every call, so a rebound name would make them disagree. + #[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#" diff --git a/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py index 08ceab77c1a..53aa210186f 100644 --- a/extra_tests/snippets/aot.py +++ b/extra_tests/snippets/aot.py @@ -81,6 +81,25 @@ def divide(a: float, b: float) -> float: 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 + + +assert countdown(3.0) == 0.0 +original_countdown = countdown + + +def countdown(a: float) -> float: + return -1.0 + + +assert original_countdown(3.0) == -1.0 + + if sys._jit.is_enabled(): compiled, rejected, deoptimized = sys._jit._stats() # `scale` is the one function above the automatic path can take. From 6d5811f2deec454afe4768c87b379d621613b425 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 24 Aug 2026 11:51:44 +0900 Subject: [PATCH 06/57] ci: build with aot and run the snippet in both modes Automatic compilation is only correct if it changes nothing, so the snippet has to pass identically with `-X aot=1` and `-X aot=0`. Also runs the existing jit snippet against the same build. Assisted-by: Claude --- .github/workflows/ci.yaml | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 7c56fc62076..38a09553527 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -147,6 +147,16 @@ 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 --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' + # - name: Install tk-dev for tkinter build # run: sudo apt-get update && sudo apt-get install -y tk-dev # if: runner.os == 'Linux' From 595717acebf6e8f4b27e53a22b3f0e8dbba365a2 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 24 Aug 2026 11:54:22 +0900 Subject: [PATCH 07/57] vm: let an interrupt out of a speculative annotation read Automatic compilation reads `__annotations__`, which under PEP 649 runs `__annotate__` - Python code that can raise. Discarding every error along with the compile attempt also discarded a KeyboardInterrupt or SystemExit that happened to land there, losing a signal the program was owed. Errors that are `Exception` subclasses stay discarded, since a forward reference raising NameError is no business of a compile attempt nobody asked for. Anything else propagates. The snippet covers both: a forward reference leaves the function interpreted and its annotations still unresolved, and a BaseException from an annotation expression reaches the caller. Assisted-by: Claude --- crates/vm/src/builtins/function.rs | 2 +- crates/vm/src/builtins/function/aot.rs | 47 ++++++++++++++++++-------- extra_tests/snippets/aot.py | 35 +++++++++++++++++++ 3 files changed, 68 insertions(+), 16 deletions(-) diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index ab7b740dcae..a7ecedb1def 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -603,7 +603,7 @@ impl Py { { let mut state = self.jit_state.load(Relaxed); if state == aot::UNTRIED && vm.state.config.settings.aot { - state = aot::compile_on_first_call(self, vm); + state = aot::compile_on_first_call(self, vm)?; } if matches!(state, aot::COMPILED_AUTO | aot::COMPILED_MANUAL) { diff --git a/crates/vm/src/builtins/function/aot.rs b/crates/vm/src/builtins/function/aot.rs index 97b8f336661..9db83adf94b 100644 --- a/crates/vm/src/builtins/function/aot.rs +++ b/crates/vm/src/builtins/function/aot.rs @@ -7,7 +7,7 @@ //! anything once it turns out it cannot be compiled. use super::PyFunction; -use crate::{Py, VirtualMachine, builtins::PyCode}; +use crate::{AsObject, Py, PyResult, VirtualMachine, builtins::PyCode}; use core::sync::atomic::{AtomicU64, Ordering::Relaxed}; use rustpython_jit::{CompiledCode, Safety}; @@ -53,38 +53,55 @@ fn code_is_eligible(code: &Py) -> bool { } } -fn try_compile(func: &Py, vm: &VirtualMachine) -> Option { +/// Discard the reason a speculative compile failed, unless it is something +/// the program has to see. +/// +/// Reading annotations runs `__annotate__`, which is Python code: a forward +/// reference raises NameError and is none of our business, but a +/// KeyboardInterrupt or SystemExit that happens to land there belongs to the +/// program, not to a compile attempt nobody asked for. +fn ignore_speculative_error(result: PyResult, vm: &VirtualMachine) -> PyResult> { + match result { + Ok(value) => Ok(Some(value)), + Err(err) if err.fast_isinstance(vm.ctx.exceptions.exception_type) => Ok(None), + Err(err) => Err(err), + } +} + +fn try_compile(func: &Py, vm: &VirtualMachine) -> PyResult> { let code: &Py = &func.code; if !code_is_eligible(code) { - return None; + return Ok(None); } - // Reading `__annotations__` runs `__annotate__`, which is Python code and - // can raise for a forward reference. Nobody asked for these annotations, so - // whatever it raises is discarded along with the compile attempt. The - // pre-filter above keeps this off the vast majority of functions. - let arg_types = super::jit::get_jit_arg_types(func, vm).ok()?; - let ret_type = super::jit::jit_ret_type(func, vm).ok()?; + let Some(arg_types) = ignore_speculative_error(super::jit::get_jit_arg_types(func, vm), vm)? + else { + return Ok(None); + }; + let Some(ret_type) = ignore_speculative_error(super::jit::jit_ret_type(func, vm), vm)? else { + return Ok(None); + }; - vm.state + Ok(vm + .state .jit_engine .compile(&code.code, &arg_types, ret_type, Safety::Strict) - .ok() + .ok()) } /// Give `func` its one automatic compile attempt and return its new state. -pub(super) fn compile_on_first_call(func: &Py, vm: &VirtualMachine) -> u8 { +pub(super) fn compile_on_first_call(func: &Py, vm: &VirtualMachine) -> PyResult { // Claim the function before running anything that can re-enter it: // evaluating `__annotate__` calls Python, which can reach this same // function, and a reentrant attempt has to interpret rather than recurse. func.jit_state.store(REJECTED, Relaxed); - let Some(compiled) = try_compile(func, vm) else { + let Some(compiled) = try_compile(func, vm)? else { vm.state.aot_stats.rejected.fetch_add(1, Relaxed); - return REJECTED; + return Ok(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 + Ok(COMPILED_AUTO) } diff --git a/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py index 53aa210186f..92740ccd95b 100644 --- a/extra_tests/snippets/aot.py +++ b/extra_tests/snippets/aot.py @@ -100,6 +100,41 @@ def countdown(a: float) -> float: assert original_countdown(3.0) == -1.0 +# Automatic compilation reads annotations, which under PEP 649 means running +# `__annotate__`. A name that is not defined yet raises there, and that is +# none of the program's business: it never asked for its annotations. +def forward(a: NotDefinedYet) -> int: + return 1 + + +assert forward(1) == 1 +try: + forward.__annotations__ +except NameError: + pass +else: + raise AssertionError("expected the forward reference to still be unresolved") + + +if sys._jit.is_enabled(): + # ... but an interrupt that lands in `__annotate__` belongs to the program. + class Boom(BaseException): + pass + + def explode(): + raise Boom + + def annotated(a: explode()) -> int: + return 1 + + try: + annotated(1) + except Boom: + pass + else: + raise AssertionError("expected a BaseException to reach the caller") + + if sys._jit.is_enabled(): compiled, rejected, deoptimized = sys._jit._stats() # `scale` is the one function above the automatic path can take. From a425fa106c65f9f94938e0e82432648cf74ce106 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 24 Aug 2026 12:08:33 +0900 Subject: [PATCH 08/57] extra_tests: gate the aot snippet on a RustPython-only symbol Every snippet also runs under CPython, which has its own `sys._jit`. On a CPython built with its JIT enabled, `is_enabled()` is true and the checks below it would reach for `_stats`, which only RustPython has. Gate on both. Assisted-by: Claude --- extra_tests/snippets/aot.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py index 92740ccd95b..2823974ca4f 100644 --- a/extra_tests/snippets/aot.py +++ b/extra_tests/snippets/aot.py @@ -2,6 +2,10 @@ # 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") + def scale(a: float, b: float) -> float: return a * b + a - b @@ -116,7 +120,7 @@ def forward(a: NotDefinedYet) -> int: raise AssertionError("expected the forward reference to still be unresolved") -if sys._jit.is_enabled(): +if AOT: # ... but an interrupt that lands in `__annotate__` belongs to the program. class Boom(BaseException): pass @@ -135,7 +139,7 @@ def annotated(a: explode()) -> int: raise AssertionError("expected a BaseException to reach the caller") -if sys._jit.is_enabled(): +if AOT: compiled, rejected, deoptimized = sys._jit._stats() # `scale` is the one function above the automatic path can take. assert compiled >= 1, (compiled, rejected, deoptimized) From 8c8ecdddb6d4f786f76161635450bee2a061382e Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 24 Aug 2026 12:42:30 +0900 Subject: [PATCH 09/57] jit: build the libffi cif once per compiled function `invoke_raw` called `JitSig::to_cif` on every invocation, rebuilding the libffi description of the signature - and the allocation behind it - for each call. Build it when the function is compiled and keep it in `CompiledCode`. Assisted-by: Claude --- crates/jit/src/lib.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 9b237d6eb5a..c5f2d03bbbb 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -177,8 +177,12 @@ impl JitEngine { jit.module.finalize_definitions()?; let code = jit.module.get_finalized_function(id); drop(jit); + // Built once: describing the signature to libffi allocates, and doing + // it per call costs more than the compiled body saves. + let cif = sig.to_cif(); Ok(CompiledCode { sig, + cif, code, _engine: self.clone(), }) @@ -243,6 +247,7 @@ pub fn compile( pub struct CompiledCode { sig: JitSig, + cif: libffi::middle::Cif, code: *const u8, /// Keeps the code memory alive; never read. _engine: Arc, @@ -272,8 +277,7 @@ impl CompiledCode { unsafe fn invoke_raw(&self, cif_args: &[libffi::middle::Arg<'_>]) -> Option { unsafe { - let cif = self.sig.to_cif(); - let value = cif.call::( + let value = self.cif.call::( libffi::middle::CodePtr::from_ptr(self.code as *const _), cif_args, ); From b836043d46edb8f4e15a513d9a2266dd91a30ea7 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 24 Aug 2026 12:43:09 +0900 Subject: [PATCH 10/57] jit: fold the two LoadGlobal rejections into one condition Turning down the self-call under Strict left two branches with identical bodies. State the one condition that admits a self-call instead. Assisted-by: Claude --- crates/jit/src/instructions.rs | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index f27305dedb6..fb1b38423f7 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -809,22 +809,22 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let oparg = namei.get(arg); let name = &bytecode.names[(oparg >> 1) as usize]; - // The only global that resolves here is this function itself, - // and it resolves by name. 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. - if self.safety == Safety::Strict { - Err(JitCompileError::NotSupported) - } else 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 { .. } => { From 69532d1d8a86705b646b218d382ac8143b78dd03 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 24 Aug 2026 22:21:16 +0900 Subject: [PATCH 11/57] jit: call compiled code through a generated entry point Every compiled function now gets a second entry point compiled alongside it, taking a flat buffer of 64-bit slots and unpacking them into the parameters the body takes. Calling it is an indirect call to a plain function pointer, so the three per-call allocations and the libffi dependency are gone, along with the union used to read the result back. Arguments travel in a fixed-size array, so a function with more than 16 parameters is now rejected instead of compiled. Assisted-by: Claude --- .cspell.dict/rust-more.txt | 1 + Cargo.lock | 1 - crates/jit/Cargo.toml | 1 - crates/jit/src/lib.rs | 251 ++++++++++++++++++++----------- crates/jit/tests/engine_tests.rs | 50 +++++- 5 files changed, 209 insertions(+), 95 deletions(-) diff --git a/.cspell.dict/rust-more.txt b/.cspell.dict/rust-more.txt index b53639c3b41..cbd2510a923 100644 --- a/.cspell.dict/rust-more.txt +++ b/.cspell.dict/rust-more.txt @@ -79,6 +79,7 @@ thiserror timelike timsort trai +uextend ulonglong unic unistd 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/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/lib.rs b/crates/jit/src/lib.rs index c5f2d03bbbb..535fbe89ffb 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -4,7 +4,7 @@ extern crate alloc; use alloc::fmt; use alloc::sync::Arc; -use core::mem::ManuallyDrop; +use core::mem::{self, ManuallyDrop}; use core::sync::atomic::{AtomicU64, Ordering}; use cranelift::prelude::*; use cranelift_jit::{JITBuilder, JITModule}; @@ -13,6 +13,16 @@ use instructions::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::(); + +/// The entry point of a compiled function: `(args, ret)`. +type JitEntry = unsafe extern "C" fn(*const u64, *mut u64); + #[derive(Debug, thiserror::Error)] #[non_exhaustive] pub enum JitCompileError { @@ -139,9 +149,100 @@ impl Jit { 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); + + let entry = self.build_entry(id, body_signature, &sig, unique)?; + + Ok((entry, sig)) + } + + /// 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) + 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 mut call_args = Vec::with_capacity(sig.args.len()); + 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, sig)) + Ok(id) } } @@ -169,6 +270,9 @@ impl JitEngine { 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); @@ -177,13 +281,12 @@ impl JitEngine { jit.module.finalize_definitions()?; let code = jit.module.get_finalized_function(id); drop(jit); - // Built once: describing the signature to libffi allocates, and doing - // it per call costs more than the compiled body saves. - let cif = sig.to_cif(); + // 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, - cif, - code, + entry, _engine: self.clone(), }) } @@ -247,8 +350,7 @@ pub fn compile( pub struct CompiledCode { sig: JitSig, - cif: libffi::middle::Cif, - code: *const u8, + entry: JitEntry, /// Keeps the code memory alive; never read. _engine: Arc, } @@ -264,27 +366,24 @@ impl CompiledCode { return Err(JitArgumentError::WrongNumberOfArguments); } - let cif_args = self - .sig - .args - .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) }) + 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(); + } + Ok(unsafe { self.invoke_raw(&slots) }) } - unsafe fn invoke_raw(&self, cif_args: &[libffi::middle::Arg<'_>]) -> Option { - unsafe { - let value = self.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)), - } + /// # Safety + /// `slots` must hold a value of the right type for each parameter. + unsafe fn invoke_raw(&self, slots: &[u64; MAX_ARGS]) -> Option { + let mut ret = 0; + // SAFETY: the entry point reads one slot per parameter and writes the + // return slot only when the signature says it returns something. + unsafe { (self.entry)(slots.as_ptr(), &raw mut ret) } + match self.sig.ret.as_ref() { + Some(JitType::None) | None => None, + Some(ty) => Some(AbiValue::from_slot(ty, ret)), } } } @@ -294,16 +393,6 @@ struct JitSig { 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 { @@ -322,15 +411,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)] @@ -342,11 +422,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"), } } } @@ -411,32 +501,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 fmt::Debug for CompiledCode { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str("[compiled code]") @@ -444,7 +508,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, } @@ -452,44 +519,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) } + // 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/engine_tests.rs b/crates/jit/tests/engine_tests.rs index bd076861f27..55c43b2d1c4 100644 --- a/crates/jit/tests/engine_tests.rs +++ b/crates/jit/tests/engine_tests.rs @@ -1,6 +1,6 @@ #[cfg(test)] mod tests { - use rustpython_jit::{JitEngine, Safety}; + use rustpython_jit::{JitCompileError, JitEngine, 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. @@ -69,4 +69,52 @@ mod tests { }; assert_eq!(code.invoke(&[7i64.into()]), Ok(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(); + 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(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) + )); + } } From 5009642d7c68ada277f25d21e50e1cde85b483c6 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 24 Aug 2026 22:21:44 +0900 Subject: [PATCH 12/57] vm: trim the jit call path and its build warnings Hold `defaults_and_kwdefaults` across the argument fill instead of cloning the pair on every call; filling a slot from a default reads the object but never runs Python code. Import `Arc` from `alloc` and name the `AbiValue` variants through `Self`, both of which clippy only reports when the jit feature is on. Assisted-by: Claude --- crates/vm/src/builtins/function/jit.rs | 11 +++++++---- crates/vm/src/vm/mod.rs | 2 +- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/crates/vm/src/builtins/function/jit.rs b/crates/vm/src/builtins/function/jit.rs index 96c1465d4f1..fe3bd91799e 100644 --- a/crates/vm/src/builtins/function/jit.rs +++ b/crates/vm/src/builtins/function/jit.rs @@ -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!(), } } @@ -203,7 +203,10 @@ pub(crate) fn get_jit_args<'a>( } } - 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 { diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index 50b305bf3c5..c7bf04a47be 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -830,7 +830,7 @@ pub struct PyGlobalState { pub require_idref: AtomicBool, /// Owns the machine code of every function compiled in this interpreter. #[cfg(feature = "jit")] - pub jit_engine: std::sync::Arc, + 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, From 1c88aa9b2fefbf1040a56bd856dc7599e8a90582 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Wed, 26 Aug 2026 01:44:12 +0900 Subject: [PATCH 13/57] vm: compute f_lineno from the instruction pointer The interpreter loop updated `prev_line` from the locations table on every instruction so that `f_lineno` could read it. `lasti` is already advanced past the instruction being executed before that instruction runs, so `locations[lasti - 1]` is its line; `f_lineno` now reads the live `lasti` and looks the line up on demand. That leaves `prev_line` carrying only the line a LINE event last fired at, so the updates the instrumented paths made for `f_lineno` are gone as well. Assisted-by: Claude --- crates/vm/src/builtins/frame.rs | 46 ++++++++++++++------------------- crates/vm/src/frame.rs | 38 --------------------------- 2 files changed, 19 insertions(+), 65 deletions(-) diff --git a/crates/vm/src/builtins/frame.rs b/crates/vm/src/builtins/frame.rs index 3704e070c34..07a9785bf7c 100644 --- a/crates/vm/src/builtins/frame.rs +++ b/crates/vm/src/builtins/frame.rs @@ -497,34 +497,26 @@ impl FrameObject { #[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 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. The live + // position has to be read rather than this object's copy: a frame + // observed from inside a call it made still reports that call's line. 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() + let lasti = if live.is_null() { + self.lasti() + } else { + unsafe { (*live).lasti.load(Relaxed) } + }; + 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)] diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index b73bdb49ea8..60514cf2973 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -3111,24 +3111,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 @@ -7297,20 +7279,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 +7593,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"); From 11f98794856d486752a30a31fe33894c158dd963 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sun, 30 Aug 2026 20:13:51 +0900 Subject: [PATCH 14/57] jit: pin down what a conditionally-defined local compiles to Assisted-by: Claude --- crates/jit/tests/misc_tests.rs | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/crates/jit/tests/misc_tests.rs b/crates/jit/tests/misc_tests.rs index 5404df0a769..28fcaf4f083 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, Safety}; #[test] fn no_return_value() { @@ -125,4 +125,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(); + 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))); + } } From 8b40262023f983e1c0a0884f7a770b4c745ea232 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sun, 30 Aug 2026 20:24:25 +0900 Subject: [PATCH 15/57] jit: pass a deopt buffer through the call boundary The entry point takes a third pointer, zeroes its status slot, and hands it to the body as the body's first parameter. Recursive calls forward it, and pass their arguments the right way round. `invoke` now returns an `Outcome` of a normal return or a `DeoptState`; nothing writes a non-zero status yet, so the deopt arm is unreachable. The interpreter answers a deopt by dropping the compiled code and running the call again from the start, whether the code was compiled on its own or on request. Assisted-by: Claude --- .cspell.dict/rust-more.txt | 1 + crates/jit/src/instructions.rs | 13 ++++- crates/jit/src/lib.rs | 88 ++++++++++++++++++++++++++---- crates/jit/tests/common.rs | 24 ++++++-- crates/jit/tests/engine_tests.rs | 18 ++++-- crates/jit/tests/misc_tests.rs | 9 ++- crates/jit/tests/safety_tests.rs | 12 ++-- crates/vm/src/builtins/function.rs | 19 +++++-- 8 files changed, 144 insertions(+), 40 deletions(-) diff --git a/.cspell.dict/rust-more.txt b/.cspell.dict/rust-more.txt index cbd2510a923..e5e1550af78 100644 --- a/.cspell.dict/rust-more.txt +++ b/.cspell.dict/rust-more.txt @@ -28,6 +28,7 @@ getres hasher hexf hexversion +iconst idents illumos ilog diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index fb1b38423f7..466dc555f95 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -69,6 +69,8 @@ struct DDValue { 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, stack: Vec, variables: Box<[Option]>, label_to_block: HashMap, @@ -182,8 +184,11 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { entry_block: Block, safety: Safety, ) -> 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"); let mut compiler = Self { builder, + deopt_ptr: *deopt_ptr, stack: Vec::new(), variables: vec![None; num_variables].into_boxed_slice(), label_to_block: HashMap::new(), @@ -193,8 +198,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { ret: ret_type, }, }; - 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(), @@ -639,11 +643,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)?; diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 535fbe89ffb..7fd27671c65 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -20,8 +20,23 @@ use std::sync::{Mutex, PoisonError}; const MAX_ARGS: usize = 16; const SLOT_SIZE: usize = size_of::(); -/// The entry point of a compiled function: `(args, ret)`. -type JitEntry = unsafe extern "C" fn(*const u64, *mut u64); +/// A guard that fires leaves its record in a second flat buffer: +/// +/// ```text +/// deopt[0] status: 0 when the call returned, otherwise the site index plus one +/// deopt[1] bound mask: bit i set when the site's i-th listed local is bound +/// 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. +#[expect(dead_code, reason = "read once a guard writes a record")] +const DEOPT_HEADER_SLOTS: usize = 2; + +/// 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] @@ -109,6 +124,11 @@ impl Jit { unique: u64, safety: Safety, ) -> Result<(FuncId, JitSig), 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)); @@ -175,7 +195,8 @@ impl Jit { unique: u64, ) -> Result { let ptr_type = self.module.target_config().pointer_type(); - // (args, ret) + // (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)); @@ -197,8 +218,12 @@ impl Jit { 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::with_capacity(sig.args.len()); + 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| { @@ -361,7 +386,7 @@ impl CompiledCode { 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); } @@ -376,18 +401,57 @@ impl CompiledCode { /// # Safety /// `slots` must hold a value of the right type for each parameter. - unsafe fn invoke_raw(&self, slots: &[u64; MAX_ARGS]) -> Option { + unsafe fn invoke_raw(&self, slots: &[u64; MAX_ARGS]) -> Outcome { let mut ret = 0; - // SAFETY: the entry point reads one slot per parameter and writes the - // return slot only when the signature says it returns something. - unsafe { (self.entry)(slots.as_ptr(), &raw mut ret) } - match self.sig.ret.as_ref() { + // 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 { + // SAFETY: the site describes exactly which slots the guard wrote. + return Outcome::Deopt(unsafe { self.read_deopt(status, deopt_ptr) }); + } + Outcome::Returned(match self.sig.ret.as_ref() { Some(JitType::None) | None => None, Some(ty) => Some(AbiValue::from_slot(ty, ret)), - } + }) + } + + /// # Safety + /// `status` must be a status this code's guards can produce, and `deopt` must + /// point at the buffer they wrote. + unsafe fn read_deopt(&self, _status: u64, _deopt: *const u64) -> DeoptState { + unreachable!("nothing writes a non-zero status yet") } } +/// 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), +} + +/// 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, +} + struct JitSig { args: Vec, ret: Option, @@ -554,7 +618,7 @@ pub struct Args<'a> { impl Args<'_> { #[must_use] - pub fn invoke(&self) -> Option { + 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 1a35f2078f7..e2dd53f1786 100644 --- a/crates/jit/tests/common.rs +++ b/crates/jit/tests/common.rs @@ -357,9 +357,16 @@ 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:?}") + } }) } } @@ -371,9 +378,14 @@ 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:?}") + } }) } } diff --git a/crates/jit/tests/engine_tests.rs b/crates/jit/tests/engine_tests.rs index 55c43b2d1c4..009fecfecf0 100644 --- a/crates/jit/tests/engine_tests.rs +++ b/crates/jit/tests/engine_tests.rs @@ -1,6 +1,6 @@ #[cfg(test)] mod tests { - use rustpython_jit::{JitCompileError, JitEngine, Safety}; + 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. @@ -25,11 +25,11 @@ mod tests { assert_eq!( first.invoke(&[3i64.into(), 4i64.into()]), - Ok(Some(3i64.into())) + Ok(Outcome::Returned(Some(3i64.into()))) ); assert_eq!( second.invoke(&[3i64.into(), 4i64.into()]), - Ok(Some(4i64.into())) + Ok(Outcome::Returned(Some(4i64.into()))) ); } @@ -52,7 +52,7 @@ mod tests { .expect("engine still usable after a rejected function"); assert_eq!( good.invoke(&[3i64.into(), 4i64.into()]), - Ok(Some(7i64.into())) + Ok(Outcome::Returned(Some(7i64.into()))) ); } @@ -67,7 +67,10 @@ mod tests { "#); f.compile_on(&engine, Safety::Permissive).expect("compile") }; - assert_eq!(code.invoke(&[7i64.into()]), Ok(Some(7i64.into()))); + assert_eq!( + code.invoke(&[7i64.into()]), + Ok(Outcome::Returned(Some(7i64.into()))) + ); } /// Every parameter reaches the compiled body from its own slot, whatever @@ -106,7 +109,10 @@ mod tests { .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(Some(15i64.into()))); + 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: diff --git a/crates/jit/tests/misc_tests.rs b/crates/jit/tests/misc_tests.rs index 28fcaf4f083..d8eda16e90c 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, JitCompileError, JitEngine, Safety}; + 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] diff --git a/crates/jit/tests/safety_tests.rs b/crates/jit/tests/safety_tests.rs index 308da6742fb..cf02ff6dd85 100644 --- a/crates/jit/tests/safety_tests.rs +++ b/crates/jit/tests/safety_tests.rs @@ -1,6 +1,6 @@ #[cfg(test)] mod tests { - use rustpython_jit::{JitEngine, Safety}; + use rustpython_jit::{JitEngine, Outcome, Safety}; /// Every operation whose machine code can trap or wrap. There is no trap /// handler, so a trap kills the process instead of raising. @@ -129,7 +129,7 @@ def band(a: int, b: int) -> int: "#); assert_eq!( code.invoke(&[6i64.into(), 3i64.into()]), - Ok(Some(2i64.into())) + Ok(Outcome::Returned(Some(2i64.into()))) ); } @@ -141,7 +141,7 @@ def lt(a: int, b: int) -> bool: "#); assert_eq!( code.invoke(&[1i64.into(), 2i64.into()]), - Ok(Some(true.into())) + Ok(Outcome::Returned(Some(true.into()))) ); } @@ -153,7 +153,7 @@ def poly(a: float, b: float) -> float: "#); assert_eq!( code.invoke(&[2.0f64.into(), 3.0f64.into()]), - Ok(Some(5.0f64.into())) + Ok(Outcome::Returned(Some(5.0f64.into()))) ); } @@ -167,7 +167,7 @@ def mixed(a: int, b: float) -> float: "#); assert_eq!( code.invoke(&[2i64.into(), 0.5f64.into()]), - Ok(Some(2.5f64.into())) + Ok(Outcome::Returned(Some(2.5f64.into()))) ); } @@ -191,7 +191,7 @@ def add(a: int, b: int) -> int: "#); assert_eq!( code.invoke(&[3i64.into(), 4i64.into()]), - Ok(Some(7i64.into())) + Ok(Outcome::Returned(Some(7i64.into()))) ); } } diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index a7ecedb1def..a3c7ab11864 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -29,7 +29,7 @@ use core::sync::atomic::AtomicU8; use core::sync::atomic::{AtomicU32, Ordering::Relaxed}; use itertools::Itertools; #[cfg(feature = "jit")] -use rustpython_jit::{CompiledCode, Safety}; +use rustpython_jit::{CompiledCode, Outcome, Safety}; fn format_missing_args( qualname: impl core::fmt::Display, @@ -584,8 +584,11 @@ impl Py { /// Drop native code and stop trying. /// - /// Only for code the AOT path compiled: an explicit `__jit__()` is a - /// standing request, and silently undoing it would be a surprise. + /// 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); @@ -613,10 +616,18 @@ impl Py { jit::get_jit_args(self, &func_args, jitted_code, vm).map(|args| args.invoke()) }); match outcome { - Some(Ok(ret)) => { + Some(Ok(Outcome::Returned(ret))) => { use crate::convert::ToPyObject; return Ok(ret.to_pyobject(vm)); } + Some(Ok(Outcome::Deopt(_))) => { + // Until the frame can be rebuilt from the record, a guard + // means running the call again from the start. The opcodes + // with a lowering have no effect outside the frame, so + // redoing the work is not observable - but the code has to + // go first, or the second attempt hits the same guard. + self.deoptimize(vm); + } Some(Err(err)) => { info!( "jit: function `{}` is falling back to being interpreted because of \ From eb40ea9c0deec76b46d0181787f3cc6102ce2573 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sun, 30 Aug 2026 21:14:36 +0900 Subject: [PATCH 16/57] jit: spill locals and the value stack at a deopt site A site records the resume offset, the type of every live local, and one entry per value-stack slot; the guard writes their values, a bound mask, and the site index into the deopt buffer, then leaves through a shared exit block. A stack entry that is the same on every path reaching the site - the callable a self-call pushes, and the null beside it - is described by the site instead of occupying a slot. Integer addition is the first operation to use it: it hands its operands back where the sum stops fitting in 64 bits, instead of trapping. The jit snippet covers that case, so the fallback to interpreting the call is exercised by the test suite. A self-call shares the caller's buffer, so after one the caller checks the status and leaves through the same exit when a nested frame gave up, rather than computing on the filler that frame returned. It leaves the nested frame's record standing. Assisted-by: Claude --- crates/jit/src/instructions.rs | 216 +++++++++++++++++++++++++++++++- crates/jit/src/lib.rs | 95 ++++++++++++-- crates/jit/tests/deopt_tests.rs | 119 ++++++++++++++++++ crates/jit/tests/lib.rs | 1 + extra_tests/snippets/jit.py | 9 ++ 5 files changed, 426 insertions(+), 14 deletions(-) create mode 100644 crates/jit/tests/deopt_tests.rs diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 466dc555f95..2c77bfa76d1 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, Safety}; +use super::{ + DEOPT_HEADER_SLOTS, DeoptSite, JitCompileError, JitSig, JitType, MAX_DEOPT_SLOTS, SLOT_SIZE, + Safety, StackEntry, +}; use alloc::collections::BTreeSet; use cranelift::codegen::ir::FuncRef; use cranelift::prelude::*; @@ -22,7 +25,7 @@ struct Local { ty: JitType, } -#[derive(Debug)] +#[derive(Debug, Clone)] enum JitValue { Int(Value), Float(Value), @@ -59,6 +62,14 @@ 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)] @@ -71,11 +82,21 @@ 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, + /// 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, safety: Safety, pub(crate) sig: JitSig, + pub(crate) deopt_sites: Vec, } /// Whether [`FunctionCompiler::add_instruction`] has a lowering for this opcode. @@ -186,17 +207,31 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { ) -> 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, + resume_offset: 0, stack: Vec::new(), variables: vec![None; num_variables].into_boxed_slice(), + bound_flags, label_to_block: HashMap::new(), safety, sig: JitSig { args: arg_types.to_vec(), ret: ret_type, }, + deopt_sites: Vec::new(), }; for (i, (ty, val)) in arg_types.iter().zip(arg_params.iter().copied()).enumerate() { compiler @@ -219,6 +254,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 { @@ -230,10 +266,136 @@ 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(), + }); + + 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 "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) => { @@ -339,6 +501,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); @@ -380,9 +543,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 @@ -405,6 +579,25 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { self.builder.ins().trap(TrapCode::user(0).unwrap()); } } + + // 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(()) } @@ -490,7 +683,10 @@ 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) } ( @@ -676,6 +872,20 @@ 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 and leave its record standing. + 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); + self.deopt_if(nested, |_| {}); + self.stack.push(val); Ok(()) diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 7fd27671c65..3fb057bb6a3 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -32,7 +32,6 @@ const SLOT_SIZE: usize = size_of::(); /// 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. -#[expect(dead_code, reason = "read once a guard writes a record")] const DEOPT_HEADER_SLOTS: usize = 2; /// The entry point of a compiled function: `(args, ret, deopt)`. @@ -104,7 +103,7 @@ impl Jit { ret: Option, unique: u64, safety: Safety, - ) -> Result<(FuncId, JitSig), JitCompileError> { + ) -> Result<(FuncId, JitSig, Vec), JitCompileError> { let result = self.build_function_inner(bytecode, args, ret, unique, safety); self.module.clear_context(&mut self.ctx); if result.is_err() { @@ -123,7 +122,7 @@ impl Jit { ret: Option, unique: u64, safety: Safety, - ) -> Result<(FuncId, JitSig), JitCompileError> { + ) -> 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. @@ -151,7 +150,7 @@ impl Jit { 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(), @@ -163,7 +162,7 @@ impl Jit { compiler.compile(func_ref, bytecode)?; - compiler.sig + (compiler.sig, compiler.deopt_sites) }; builder.seal_all_blocks(); @@ -177,7 +176,7 @@ impl Jit { let entry = self.build_entry(id, body_signature, &sig, unique)?; - Ok((entry, sig)) + Ok((entry, sig, deopt_sites)) } /// Build the entry point callers go through: it takes a flat buffer of @@ -302,7 +301,7 @@ impl JitEngine { // 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) = jit.build_function(bytecode, args, ret, unique, safety)?; + let (id, sig, deopt_sites) = jit.build_function(bytecode, args, ret, unique, safety)?; jit.module.finalize_definitions()?; let code = jit.module.get_finalized_function(id); drop(jit); @@ -311,6 +310,7 @@ impl JitEngine { let entry = unsafe { mem::transmute::<*const u8, JitEntry>(code) }; Ok(CompiledCode { sig, + deopt_sites, entry, _engine: self.clone(), }) @@ -375,6 +375,8 @@ pub fn compile( pub struct CompiledCode { sig: JitSig, + /// Indexed by the status a guard writes, less one. + deopt_sites: Vec, entry: JitEntry, /// Keeps the code memory alive; never read. _engine: Arc, @@ -414,7 +416,9 @@ impl CompiledCode { // SAFETY: the entry point stores the status first thing. let status = unsafe { deopt_ptr.read() }; if status != 0 { - // SAFETY: the site describes exactly which slots the guard wrote. + // SAFETY: a non-zero status is written only by a guard of this + // code, and it is the index of the site that describes the record + // that guard just wrote into this buffer. return Outcome::Deopt(unsafe { self.read_deopt(status, deopt_ptr) }); } Outcome::Returned(match self.sig.ret.as_ref() { @@ -426,8 +430,42 @@ impl CompiledCode { /// # Safety /// `status` must be a status this code's guards can produce, and `deopt` must /// point at the buffer they wrote. - unsafe fn read_deopt(&self, _status: u64, _deopt: *const u64) -> DeoptState { - unreachable!("nothing writes a non-zero status yet") + unsafe fn read_deopt(&self, status: u64, deopt: *const u64) -> DeoptState { + let site = &self.deopt_sites[status as usize - 1]; + // 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() + .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, + } } } @@ -449,7 +487,42 @@ pub struct DeoptState { /// unbound. pub locals: Vec>, /// Bottom to top. - pub stack: Vec, + 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]>, +} + +/// 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 { diff --git a/crates/jit/tests/deopt_tests.rs b/crates/jit/tests/deopt_tests.rs new file mode 100644 index 00000000000..186d740c888 --- /dev/null +++ b/crates/jit/tests/deopt_tests.rs @@ -0,0 +1,119 @@ +#[cfg(test)] +mod tests { + use rustpython_jit::{AbiValue, JitEngine, Outcome, Safety, StackValue}; + + fn int(value: i64) -> StackValue { + StackValue::Value(AbiValue::Int(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:?}"), + } + } + + /// 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(); + 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: the record the callee + /// wrote is the one the interpreter needs. + #[test] + fn a_caller_stops_when_a_nested_frame_gives_up() { + let engine = JitEngine::new(); + 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()))) + ); + // `blow(1)` overflows. `blow(2)` reaches that overflow one frame down, + // so the guard that fires is the inner frame's and its record is what + // comes back. + for n in [1i64, 2] { + match code.invoke(&[n.into()]) { + Ok(Outcome::Deopt(state)) => assert_eq!( + state.stack, + vec![int(4611686018427387904), int(4611686018427387904)], + "n = {n}" + ), + other => panic!("expected a deopt for n = {n}, got {other:?}"), + } + } + } +} diff --git a/crates/jit/tests/lib.rs b/crates/jit/tests/lib.rs index 8438ef51e61..58697e5dc3e 100644 --- a/crates/jit/tests/lib.rs +++ b/crates/jit/tests/lib.rs @@ -3,6 +3,7 @@ extern crate alloc; #[macro_use] mod common; mod bool_tests; +mod deopt_tests; mod engine_tests; mod float_tests; mod int_tests; diff --git a/extra_tests/snippets/jit.py b/extra_tests/snippets/jit.py index 887cbb50e7e..1e98efbb376 100644 --- a/extra_tests/snippets/jit.py +++ b/extra_tests/snippets/jit.py @@ -27,3 +27,12 @@ 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 From 5f3d697a869675f73f1e393438b2b24b1001bc31 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sun, 30 Aug 2026 21:49:56 +0900 Subject: [PATCH 17/57] jit: deoptimize instead of trapping or wrapping on int overflow Subtraction, multiplication, negation and the shifts hand the operands back to the interpreter where the answer stops fitting in 64 bits, where a shift count is out of range, or where a left shift loses bits. Assisted-by: Claude --- crates/jit/src/instructions.rs | 82 ++++++++++++--------- crates/jit/tests/deopt_tests.rs | 121 +++++++++++++++++++++++++++++++ crates/jit/tests/safety_tests.rs | 1 + 3 files changed, 171 insertions(+), 33 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 2c77bfa76d1..11954df041a 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -13,12 +13,6 @@ 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, @@ -693,7 +687,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 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), @@ -717,7 +714,11 @@ 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), @@ -729,25 +730,38 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { JitValue::Int(b), ) => JitValue::Int(self.compile_ipow(a, b)), ( - 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) } ( @@ -1136,14 +1150,11 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { } Instruction::UnaryNegative => { match self.stack.pop().ok_or(JitCompileError::BadBytecode)? { - // Lowered as `0 - val`, which traps on i64::MIN. - JitValue::Int(_) if self.safety == Safety::Strict => { - Err(JitCompileError::NotSupported) - } 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(()) } @@ -1169,10 +1180,15 @@ 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. diff --git a/crates/jit/tests/deopt_tests.rs b/crates/jit/tests/deopt_tests.rs index 186d740c888..1b35c1b408a 100644 --- a/crates/jit/tests/deopt_tests.rs +++ b/crates/jit/tests/deopt_tests.rs @@ -27,6 +27,127 @@ def add(a: int, b: int) -> int: } } + /// 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 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] diff --git a/crates/jit/tests/safety_tests.rs b/crates/jit/tests/safety_tests.rs index cf02ff6dd85..934e9519acd 100644 --- a/crates/jit/tests/safety_tests.rs +++ b/crates/jit/tests/safety_tests.rs @@ -87,6 +87,7 @@ def shift(a: int, b: int) -> int: } #[test] + #[ignore = "Task 7 makes Strict accept this"] fn strict_rejects_int_negate() { assert_strict_rejects!(neg => r#" def neg(a: int) -> int: From 4982e3121a7bf7a6ac55f2c09e313c156e0f1dea Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sun, 30 Aug 2026 21:51:19 +0900 Subject: [PATCH 18/57] jit: correct the bound-mask comment in the deopt buffer layout The bit index is the varname slot, including slots whose entry is unlisted, not the position among listed locals. Assisted-by: Claude --- crates/jit/src/lib.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 3fb057bb6a3..dd3fd68447a 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -24,7 +24,7 @@ const SLOT_SIZE: usize = size_of::(); /// /// ```text /// deopt[0] status: 0 when the call returned, otherwise the site index plus one -/// deopt[1] bound mask: bit i set when the site's i-th listed local is bound +/// 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 /// ``` /// From f0dd84d58be5b1b08a2425654abc01ac44bb9e65 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sun, 30 Aug 2026 22:03:06 +0900 Subject: [PATCH 19/57] jit: floor int division and take the divisor's sign for the remainder sdiv rounds toward zero and srem takes the dividend's sign; both are corrected on the condition they disagree. A zero divisor, i64::MIN // -1, and a true division whose operands do not fit a double's significand deoptimize. Assisted-by: Claude --- .cspell.dict/rust-more.txt | 1 + .cspell.json | 1 + crates/jit/src/instructions.rs | 70 ++++++++++++++++++--- crates/jit/tests/deopt_tests.rs | 106 ++++++++++++++++++++++++++++++++ crates/jit/tests/int_tests.rs | 34 ++++++++-- 5 files changed, 200 insertions(+), 12 deletions(-) diff --git a/.cspell.dict/rust-more.txt b/.cspell.dict/rust-more.txt index e5e1550af78..c989a26482d 100644 --- a/.cspell.dict/rust-more.txt +++ b/.cspell.dict/rust-more.txt @@ -28,6 +28,7 @@ getres hasher hexf hexversion +iabs iconst idents illumos 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/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 11954df041a..e709dc157b8 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -695,17 +695,34 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 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. + let inexact = |compiler: &mut Self, v: Value| { + let magnitude = compiler.builder.ins().iabs(v); + compiler.builder.ins().icmp_imm( + IntCC::UnsignedGreaterThanOrEqual, + magnitude, + 1 << 53, + ) + }; + let a_wide = inexact(self, a); + let b_wide = inexact(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)) @@ -723,7 +740,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 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), @@ -1191,6 +1208,45 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { Ok(out) } + /// 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. + 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); + let remainder = self.builder.ins().srem(a, b); + + // 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); + 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)) + } + /// 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 { diff --git a/crates/jit/tests/deopt_tests.rs b/crates/jit/tests/deopt_tests.rs index 1b35c1b408a..4c9645d85ea 100644 --- a/crates/jit/tests/deopt_tests.rs +++ b/crates/jit/tests/deopt_tests.rs @@ -148,6 +148,112 @@ def shr(a: int, b: int) -> int: } } + /// 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. + #[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:?}"), + } + } + + /// 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 +"# }; + assert_eq!( + code.invoke(&[((1i64 << 53) - 1).into(), 1i64.into()]), + Ok(Outcome::Returned(Some((((1i64 << 53) - 1) as f64).into()))) + ); + match code.invoke(&[(1i64 << 53).into(), 1i64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(1i64 << 53), 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 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] diff --git a/crates/jit/tests/int_tests.rs b/crates/jit/tests/int_tests.rs index 23cf98aafe1..7085c602352 100644 --- a/crates/jit/tests/int_tests.rs +++ b/crates/jit/tests/int_tests.rs @@ -2,6 +2,31 @@ 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)); + + 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)); + } + #[test] fn basic_add() { let add = jit_function! { add(a:i64, b:i64) -> i64 => r##" @@ -65,10 +90,8 @@ 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)); + // Operands at or past i64::MAX / i64::MIN do not fit a double's + // significand and deoptimize instead; see deopt_tests.rs. } #[test] @@ -119,7 +142,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] From e812d68eb1dcaeaff72e143b66d82c31997ae6bb Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sun, 30 Aug 2026 22:30:13 +0900 Subject: [PATCH 20/57] jit: derive the floor-division remainder from the quotient srem duplicated the sdiv the quotient already computed: cranelift keeps a trapping division live even when its result goes unused, so a % b cost two hardware divisions where one suffices. Replaced it with a - quotient * b, which cannot overflow because the quotient truncates toward zero. Also: tested the wide-operand guard's divisor half (only the dividend side had a test), pinned i64::MIN // 1, i64::MIN // 3 and their remainders, tested that the shared guard deopts i64::MIN % -1 too, corrected the basic_div comment's threshold, renamed the true-division closure to avoid colliding with compile_floor_div's local of the same name, and corrected the comment justifying the quotient correction's overflow safety. Assisted-by: Claude --- crates/jit/src/instructions.rs | 25 ++++++++++++++++++++----- crates/jit/tests/deopt_tests.rs | 23 ++++++++++++++++++++++- crates/jit/tests/int_tests.rs | 9 +++++++-- 3 files changed, 49 insertions(+), 8 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index e709dc157b8..9c4eb305244 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -710,7 +710,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { // a ulp out as soon as either conversion is inexact - which // is exactly when the operand does not fit in a double's // significand. - let inexact = |compiler: &mut Self, v: Value| { + let too_wide = |compiler: &mut Self, v: Value| { let magnitude = compiler.builder.ins().iabs(v); compiler.builder.ins().icmp_imm( IntCC::UnsignedGreaterThanOrEqual, @@ -718,8 +718,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 1 << 53, ) }; - let a_wide = inexact(self, a); - let b_wide = inexact(self, b); + 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)?; @@ -1215,7 +1215,9 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { /// /// 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. + /// 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); @@ -1227,7 +1229,15 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { self.deopt_branch(overflows, &operands)?; let quotient = self.builder.ins().sdiv(a, b); - let remainder = self.builder.ins().srem(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. @@ -1239,6 +1249,11 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 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); diff --git a/crates/jit/tests/deopt_tests.rs b/crates/jit/tests/deopt_tests.rs index 4c9645d85ea..f2605263edc 100644 --- a/crates/jit/tests/deopt_tests.rs +++ b/crates/jit/tests/deopt_tests.rs @@ -181,7 +181,9 @@ def rem(a: int, b: int) -> int: /// `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. + /// 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#" @@ -194,6 +196,17 @@ def div(a: int, b: int) -> int: } 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 @@ -236,6 +249,14 @@ def div(a: int, b: int) -> float: } 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).into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(1), int(1i64 << 53)]); + } + other => panic!("expected a deopt, got {other:?}"), + } } /// `iabs` of `i64::MIN` is `i64::MIN` again, which an unsigned comparison diff --git a/crates/jit/tests/int_tests.rs b/crates/jit/tests/int_tests.rs index 7085c602352..0cf31961e03 100644 --- a/crates/jit/tests/int_tests.rs +++ b/crates/jit/tests/int_tests.rs @@ -15,6 +15,10 @@ mod tests { 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: @@ -25,6 +29,7 @@ mod tests { 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] @@ -90,8 +95,8 @@ 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)); - // Operands at or past i64::MAX / i64::MIN do not fit a double's - // significand and deoptimize instead; see deopt_tests.rs. + // An operand at or past `1 << 53` does not fit a double's + // significand and deoptimizes instead; see deopt_tests.rs. } #[test] From 2b731905efdaa78d2039935b50c6716086a341e9 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sun, 30 Aug 2026 23:01:24 +0900 Subject: [PATCH 21/57] jit: deoptimize where a power or a float division leaves the reals Integer exponentiation stops answering 0 for a negative exponent and guards every multiplication in its loop. Float division and float exponentiation deoptimize where the interpreter raises or returns a complex number. Assisted-by: Claude --- .cspell.dict/rust-more.txt | 1 + crates/jit/src/instructions.rs | 128 ++++++++++++++++++++------------ crates/jit/tests/deopt_tests.rs | 121 ++++++++++++++++++++++++++++++ crates/jit/tests/float_tests.rs | 16 ++-- crates/jit/tests/int_tests.rs | 6 ++ 5 files changed, 217 insertions(+), 55 deletions(-) diff --git a/.cspell.dict/rust-more.txt b/.cspell.dict/rust-more.txt index c989a26482d..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 diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 9c4eb305244..af536b9eb90 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -745,7 +745,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 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::InplaceLshift @@ -817,22 +820,42 @@ 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(); + + // The original operands, for a guard to hand back on deopt - + // `operand_one`/`operand_two` below are the converted doubles + // and do not describe the stack the interpreter had. + let operands = [ + JitValue::from_type_and_value(a_ty.clone(), a), + JitValue::from_type_and_value(b_ty.clone(), b), + ]; + + 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, }; @@ -848,10 +871,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), } @@ -1599,7 +1630,12 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { /// - 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 { + fn compile_fpow( + &mut self, + a: Value, + b: Value, + operands: &[JitValue], + ) -> Result { let f64_ty = types::F64; let i64_ty = types::I64; let zero_f = self.builder.ins().f64const(0.0); @@ -1608,6 +1644,18 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let inf_f = self.builder.ins().f64const(f64::INFINITY); let neg_inf_f = self.builder.ins().f64const(f64::NEG_INFINITY); + // 0.0 ** negative raises rather than returning an infinity. + 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)?; + // A negative base raised to a fractional power is complex. + let base_negative = self.builder.ins().fcmp(FloatCC::LessThan, a, zero_f); + let truncated = self.builder.ins().trunc(b); + let fractional = self.builder.ins().fcmp(FloatCC::NotEqual, truncated, b); + let complex = self.builder.ins().band(base_negative, fractional); + self.deopt_branch(complex, operands)?; + // Merge block for final result. let merge_block = self.builder.create_block(); self.builder.append_block_param(merge_block, f64_ty); @@ -1815,27 +1863,29 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { // ----- Merge: Return the final result. self.builder.switch_to_block(merge_block); - self.builder.block_params(merge_block)[0] + Ok(self.builder.block_params(merge_block)[0]) } - 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, and the result of a widening + // loop stops fitting in 64 bits quickly. + 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 @@ -1847,32 +1897,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); @@ -1901,12 +1929,20 @@ 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); + + // The result product is kept only when this bit of the exponent is set. + let (mul_result, mul_carry) = self.builder.ins().smul_overflow(result_phi, base_phi); + let mul_overflows = self.builder.ins().band(is_odd, mul_carry); + self.deopt_branch(mul_overflows, 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. + 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()], @@ -1917,12 +1953,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/tests/deopt_tests.rs b/crates/jit/tests/deopt_tests.rs index f2605263edc..95ad2155dac 100644 --- a/crates/jit/tests/deopt_tests.rs +++ b/crates/jit/tests/deopt_tests.rs @@ -6,6 +6,10 @@ mod tests { 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. @@ -275,6 +279,123 @@ def div(a: int, b: int) -> float: } } + /// 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. + #[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 +"# }; + match int_over_float.invoke(&[1i64.into(), 0.0f64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![int(1), float(0.0)]); + } + other => panic!("expected a deopt, 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. + #[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 +"# }; + match code.invoke(&[0.0f64.into(), (-1.0f64).into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![float(0.0), float(-1.0)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// A negative base raised to a fractional power is complex. + #[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, 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] diff --git a/crates/jit/tests/float_tests.rs b/crates/jit/tests/float_tests.rs index f667b1e764a..30512bf8847 100644 --- a/crates/jit/tests/float_tests.rs +++ b/crates/jit/tests/float_tests.rs @@ -158,6 +158,9 @@ mod tests { 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)); + // A negative base with an integral exponent is real, so the complex + // guard must not fire on it. + assert_approx_eq!(pow(-8.0, 2.0), Ok(64.0)); // 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)); @@ -231,11 +234,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 +255,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 0cf31961e03..218f538525e 100644 --- a/crates/jit/tests/int_tests.rs +++ b/crates/jit/tests/int_tests.rs @@ -161,6 +161,12 @@ 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 computes a product every iteration and + // discards it whenever the corresponding exponent bit is unset, so + // the final squaring overflowing to 2^64 must not deoptimize an + // answer - 2^33 and 2^62 - that fits an i64 comfortably. + assert_eq!(power(2, 33), Ok(8589934592)); + assert_eq!(power(2, 62), Ok(4611686018427387904)); } #[test] From cf22801a4b72e63c932680266d67bd1bce47d34b Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Sun, 30 Aug 2026 23:58:47 +0900 Subject: [PATCH 22/57] jit: close the float power trap and its remaining wrong answers compile_fpow was neither trap-free nor interpreter-faithful. A large-magnitude exponent made dd_exp round b * ln|a| to an i64 outside range, and fcvt_to_sint traps there with no handler - 2.0 ** 1e300 and similar killed the process. A finite base whose power overflowed a double returned an infinity instead of raising OverflowError. An infinite or NaN exponent answered on b alone, ignoring the base, so several pairs returned a value the interpreter would not. A negative zero base lost its sign in the zero-base fast path. Three guards close these, all before compile_fpow's normal blocks or after its merge, deoptimizing rather than computing a wrong answer: bounding |b| under 1024 (bounding the product dd_exp rounds, and written as the negation of an ordered LessThan so an unordered NaN exponent deopts too), a negative-zero base check by bit pattern, and a post-merge check that a finite base never produces an infinite result. With the fractional-exponent guard already excluding non-integral exponents for a negative base, the two now-unreachable domain-error branches in compile_fpow are removed. In compile_ipow, the result multiply's overflow no longer needs masking by the exponent's low bit: continue_block only runs with exp != 0, so a clear low bit still forces exp >= 2, and an overflowing result * base with |base| >= 2 means result * base^exp overflows too, so a later guard always catches it; with |base| <= 1 the product cannot overflow at all. The squaring's mask stays, since dropping it does deopt answers that fit an i64. The mixed int/float arm's operands array, used only by TrueDivide and Power, is now built lazily rather than for every arithmetic op. The mixed-arm zero-divisor test gained -0.0 coverage to match the float/float test. Assisted-by: Claude --- crates/jit/src/instructions.rs | 123 +++++++++++++++++++------------- crates/jit/tests/deopt_tests.rs | 106 +++++++++++++++++++++++++-- crates/jit/tests/float_tests.rs | 12 ++-- crates/jit/tests/int_tests.rs | 10 +-- extra_tests/snippets/jit.py | 47 ++++++++++++ 5 files changed, 234 insertions(+), 64 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index af536b9eb90..193dacff0b2 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -842,24 +842,29 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let a_ty = a_type.unwrap(); let b_ty = b_type.unwrap(); - // The original operands, for a guard to hand back on deopt - - // `operand_one`/`operand_two` below are the converted doubles - // and do not describe the stack the interpreter had. - let operands = [ - JitValue::from_type_and_value(a_ty.clone(), a), - JitValue::from_type_and_value(b_ty.clone(), b), - ]; - - let operand_one = match a_ty { + let operand_one = match &a_ty { JitType::Int => self.builder.ins().fcvt_from_sint(types::F64, a), _ => a, }; - let operand_two = match b_ty { + 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)) @@ -874,14 +879,14 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 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)?; + 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, - &operands, + &operands(), )?) } _ => return Err(JitCompileError::NotSupported), @@ -1628,8 +1633,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { /// 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. + /// - For a == 0: Handles special cases for 0^b, returning 0 or 1 - a + /// negative exponent has already deopted. + /// - For a < 0: Adjusts the sign if b is odd - a fractional exponent has + /// already deopted. fn compile_fpow( &mut self, a: Value, @@ -1644,6 +1651,24 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let inf_f = self.builder.ins().f64const(f64::INFINITY); let neg_inf_f = self.builder.ins().f64const(f64::NEG_INFINITY); + // Below, `dd_exp` rounds `b * ln|a|` to an i64 exponent with + // `fcvt_to_sint`, which traps once that product leaves i64 range + // (roughly 6.4e18) - and cranelift traps have no handler. A finite + // base's `ln|a|` is at most ~709.8, so bounding `|b|` under 1024 + // keeps the product under ~7.3e5, nowhere near the trap; the two + // `fcvt_to_sint` calls on `b` further down share the same bound. + // `UnorderedOrGreaterThanOrEqual` is the exact negation of an + // ordered `LessThan`, so an unordered (NaN) exponent deopts too + // rather than comparing false on both sides. + let exponent_bound = self.builder.ins().f64const(1024.0); + let abs_b = self.builder.ins().fabs(b); + let exponent_out_of_range = self.builder.ins().fcmp( + FloatCC::UnorderedOrGreaterThanOrEqual, + abs_b, + exponent_bound, + ); + self.deopt_branch(exponent_out_of_range, operands)?; + // 0.0 ** negative raises rather than returning an infinity. let base_zero = self.builder.ins().fcmp(FloatCC::Equal, a, zero_f); let exp_negative = self.builder.ins().fcmp(FloatCC::LessThan, b, zero_f); @@ -1655,6 +1680,12 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let fractional = self.builder.ins().fcmp(FloatCC::NotEqual, truncated, b); let complex = self.builder.ins().band(base_negative, fractional); self.deopt_branch(complex, operands)?; + // A negative zero base loses its sign below - Edge Case 3 always + // returns +0.0 - but `(-0.0) ** 3.0` is `-0.0`. + let base_bits = self.builder.ins().bitcast(types::I64, MemFlags::new(), a); + let base_bits_nonzero = self.builder.ins().icmp_imm(IntCC::NotEqual, base_bits, 0); + let negative_zero_base = self.builder.ins().band(base_zero, base_bits_nonzero); + self.deopt_branch(negative_zero_base, operands)?; // Merge block for final result. let merge_block = self.builder.create_block(); @@ -1746,19 +1777,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { .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. + // a is -infinity here, and the fractional-exponent guard above has + // already deopted any negative base (-infinity included) paired + // with a non-integral b, so b is always an integer at this point. 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); @@ -1807,22 +1829,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 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. + // ----- Case: a < 0: Only an integral exponent reaches here - the + // fractional-exponent guard above has already deopted the rest. 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()]); - - // For negative base with an integer exponent: - self.builder.switch_to_block(neg_int_block); 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); @@ -1861,9 +1871,19 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { .ins() .jump(merge_block, &[phi_mag_val_even.into()]); - // ----- Merge: Return the final result. + // ----- Merge: Return the final result, unless a finite base + // produced an infinity - the interpreter raises OverflowError there + // instead of saturating. An already-infinite base (`inf ** 2.0`) + // must keep returning its infinity, hence the base-finite half. self.builder.switch_to_block(merge_block); - Ok(self.builder.block_params(merge_block)[0]) + let result = self.builder.block_params(merge_block)[0]; + let abs_result = self.builder.ins().fabs(result); + let result_infinite = self.builder.ins().fcmp(FloatCC::Equal, abs_result, inf_f); + let abs_a = self.builder.ins().fabs(a); + let base_finite = self.builder.ins().fcmp(FloatCC::LessThan, abs_a, inf_f); + let overflowed = self.builder.ins().band(result_infinite, base_finite); + self.deopt_branch(overflowed, operands)?; + Ok(result) } fn compile_ipow( @@ -1872,8 +1892,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { b: Value, operands: &[JitValue], ) -> Result { - // A negative exponent makes this a float, and the result of a widening - // loop stops fitting in 64 bits quickly. + // 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)?; @@ -1930,13 +1950,20 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let is_odd = self.builder.ins().band_imm(exp_phi, 1); let is_odd = self.builder.ins().icmp_imm(IntCC::Equal, is_odd, 1); - // The result product is kept only when this bit of the exponent is set. + // 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); - let mul_overflows = self.builder.ins().band(is_odd, mul_carry); - self.deopt_branch(mul_overflows, operands)?; + self.deopt_branch(mul_carry, operands)?; let new_result = self.builder.ins().select(is_odd, mul_result, result_phi); - // The squared base is read only if there is another iteration to read it. + // 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); diff --git a/crates/jit/tests/deopt_tests.rs b/crates/jit/tests/deopt_tests.rs index 95ad2155dac..e1ec7513b9d 100644 --- a/crates/jit/tests/deopt_tests.rs +++ b/crates/jit/tests/deopt_tests.rs @@ -340,18 +340,21 @@ def div(a: float, b: float) -> float: /// 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. + /// 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 "# }; - match int_over_float.invoke(&[1i64.into(), 0.0f64.into()]) { - Ok(Outcome::Deopt(state)) => { - assert_eq!(state.stack, vec![int(1), float(0.0)]); + 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:?}"), } - other => panic!("expected a deopt, got {other:?}"), } let float_over_int = jit_function! { div => r#" @@ -396,6 +399,99 @@ def pow(a: float, b: float) -> float: } } + /// `dd_exp` rounds `b * ln|a|` to an i64 with `fcvt_to_sint`, which traps + /// once that leaves i64 range - and cranelift traps have no handler, so + /// `2.0 ** 1e300` and `(-2.0) ** 1e300` used to kill the process + /// outright. `2.0 ** 1024.0` sits exactly on the bound's threshold, and + /// an infinite exponent is caught by the same bound regardless of how it + /// got there. + #[test] + fn float_power_deopts_on_an_out_of_range_exponent() { + let code = jit_function! { pow => r#" +def pow(a: float, b: float) -> float: + return a ** b +"# }; + for (a, b) in [ + (2.0f64, 1e300f64), + (-2.0f64, 1e300f64), + (2.0f64, 1024.0f64), + (-1.0f64, f64::INFINITY), + (0.5f64, f64::INFINITY), + (0.5f64, f64::NEG_INFINITY), + (-1.0f64, f64::NEG_INFINITY), + ] { + 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:?}"), + } + } + } + + /// The out-of-range bound is written as the negation of an ordered + /// `LessThan`, so an unordered (NaN) exponent deopts the same way a huge + /// one does. `NaN != NaN`, so this gets its own test rather than joining + /// the table above. + #[test] + fn float_power_deopts_on_a_nan_exponent() { + let code = jit_function! { pow => r#" +def pow(a: float, b: float) -> float: + return a ** b +"# }; + match code.invoke(&[1.0f64.into(), f64::NAN.into()]) { + Ok(Outcome::Deopt(state)) => match &state.stack[..] { + [ + StackValue::Value(AbiValue::Float(a)), + StackValue::Value(AbiValue::Float(b)), + ] => { + assert_eq!(*a, 1.0); + assert!(b.is_nan()); + } + other => panic!("unexpected stack shape: {other:?}"), + }, + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// A finite base 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. + #[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()))) + ); + match code.invoke(&[1e308f64.into(), 2.0f64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![float(1e308), float(2.0)]); + } + other => panic!("expected a deopt, got {other:?}"), + } + } + + /// A negative zero base loses its sign in the zero-base edge case, which + /// always returns `+0.0`; `(-0.0) ** 3.0` is `-0.0`. + #[test] + fn float_power_deopts_on_negative_zero_base() { + let code = jit_function! { pow => r#" +def pow(a: float, b: float) -> float: + return a ** b +"# }; + match code.invoke(&[(-0.0f64).into(), 3.0f64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![float(-0.0), float(3.0)]); + } + other => panic!("expected a deopt, 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] diff --git a/crates/jit/tests/float_tests.rs b/crates/jit/tests/float_tests.rs index 30512bf8847..4b469003de5 100644 --- a/crates/jit/tests/float_tests.rs +++ b/crates/jit/tests/float_tests.rs @@ -134,7 +134,8 @@ mod tests { //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)); + // An infinite exponent deoptimizes regardless of the base - see + // deopt_tests.rs. // 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)); @@ -143,8 +144,6 @@ mod tests { 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)); // Test positive float base, positive float exponent assert_approx_eq!(pow(2.0, 2.0), Ok(4.0)); @@ -204,14 +203,13 @@ mod tests { // * 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^2.0 overflows a finite base to an infinity, which raises + // OverflowError rather than saturating - see deopt_tests.rs. // 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)); + // 1e-308^-2.0 overflows the same way as 1e308^2.0 above. // 1e100^(1e50) //assert_approx_eq!(pow(1e100, 1e50), Ok(1.0000000000000002e+150)); // fail to run (Crashes as "illegal hardware instruction") // 1e50^(1e-100) diff --git a/crates/jit/tests/int_tests.rs b/crates/jit/tests/int_tests.rs index 218f538525e..4150037a8ac 100644 --- a/crates/jit/tests/int_tests.rs +++ b/crates/jit/tests/int_tests.rs @@ -161,10 +161,12 @@ 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 computes a product every iteration and - // discards it whenever the corresponding exponent bit is unset, so - // the final squaring overflowing to 2^64 must not deoptimize an - // answer - 2^33 and 2^62 - that fits an i64 comfortably. + // 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)); } diff --git a/extra_tests/snippets/jit.py b/extra_tests/snippets/jit.py index 1e98efbb376..2fd34a2bbcd 100644 --- a/extra_tests/snippets/jit.py +++ b/extra_tests/snippets/jit.py @@ -36,3 +36,50 @@ def add(a: int, b: int) -> int: 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 From eb8026183e4f21e75a6982f4cb52d9b0e42cc55e Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 00:01:39 +0900 Subject: [PATCH 23/57] jit: revive the 1e100 ** 1e50 crash as a deopt test An old comment in float_tests.rs documented this case as crashing with an illegal hardware instruction and commented it out rather than fixing it; the out-of-range-exponent guard added in the previous commit closes it. Moved it into float_power_deopts_on_an_out_of_range_exponent instead of leaving it disabled. The expected value the old comment recorded (1.0000000000000002e+150) was wrong regardless - the true result overflows a double, which is why the interpreter raises. Assisted-by: Claude --- crates/jit/tests/deopt_tests.rs | 11 +++++++---- crates/jit/tests/float_tests.rs | 3 +-- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/crates/jit/tests/deopt_tests.rs b/crates/jit/tests/deopt_tests.rs index e1ec7513b9d..952b2967d5c 100644 --- a/crates/jit/tests/deopt_tests.rs +++ b/crates/jit/tests/deopt_tests.rs @@ -401,10 +401,12 @@ def pow(a: float, b: float) -> float: /// `dd_exp` rounds `b * ln|a|` to an i64 with `fcvt_to_sint`, which traps /// once that leaves i64 range - and cranelift traps have no handler, so - /// `2.0 ** 1e300` and `(-2.0) ** 1e300` used to kill the process - /// outright. `2.0 ** 1024.0` sits exactly on the bound's threshold, and - /// an infinite exponent is caught by the same bound regardless of how it - /// got there. + /// `2.0 ** 1e300`, `(-2.0) ** 1e300`, and `1e100 ** 1e50` used to kill + /// the process outright (the last of those was a known crash: an old + /// comment in `float_tests.rs` documented it and commented the case out + /// rather than fixing it). `2.0 ** 1024.0` sits exactly on the bound's + /// threshold, and an infinite exponent is caught by the same bound + /// regardless of how it got there. #[test] fn float_power_deopts_on_an_out_of_range_exponent() { let code = jit_function! { pow => r#" @@ -414,6 +416,7 @@ def pow(a: float, b: float) -> float: for (a, b) in [ (2.0f64, 1e300f64), (-2.0f64, 1e300f64), + (1e100f64, 1e50f64), (2.0f64, 1024.0f64), (-1.0f64, f64::INFINITY), (0.5f64, f64::INFINITY), diff --git a/crates/jit/tests/float_tests.rs b/crates/jit/tests/float_tests.rs index 4b469003de5..d5d763ebc30 100644 --- a/crates/jit/tests/float_tests.rs +++ b/crates/jit/tests/float_tests.rs @@ -210,8 +210,7 @@ mod tests { // 1e-308^2.0 //assert_approx_eq!(pow(1e-308, 2.0), Ok(0.0)); // --8.403311421507407 // 1e-308^-2.0 overflows the same way as 1e308^2.0 above. - // 1e100^(1e50) - //assert_approx_eq!(pow(1e100, 1e50), Ok(1.0000000000000002e+150)); // fail to run (Crashes as "illegal hardware instruction") + // 1e100^(1e50) has an out-of-range exponent - see deopt_tests.rs. // 1e50^(1e-100) assert_approx_eq!(pow(1e50, 1e-100), Ok(1.0)); // 1e308^(-1e2) From 6fee87451d6c3539fa37fb03d3d389672805d764 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 00:13:40 +0900 Subject: [PATCH 24/57] jit: remove compile_fpow's edge cases now unreachable under the exponent bound The |b| < 1024 guard already deopts any NaN or infinite exponent, which made Edge Cases 2, 5, and 6 (b is NaN, b == +infinity, b == -infinity) dead code - the same situation as the domain-error blocks removed earlier for a negative base. Removed all three blocks along with their brif splits, leaving the fall-through chain from Edge Case 1 through Edge Case 8 contiguous. Pinned nan ** 2.0 in float_tests.rs::basic_power to confirm Edge Case 4 (a is NaN) still answers, since only the exponent is bounded, not the base. Added a removal note next to (-infinity) ** (-infinity), which was missing one from the previous commit. Assisted-by: Claude --- crates/jit/src/instructions.rs | 36 +++++---------------------------- crates/jit/tests/float_tests.rs | 6 +++++- 2 files changed, 10 insertions(+), 32 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 193dacff0b2..17e6bf91554 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -1702,16 +1702,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 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 2 (b is NaN) is unreachable: the out-of-range-exponent + // guard above already deopts any NaN exponent. // --- Edge Case 3: a == 0.0 → return 0.0 let cmp_a_zero = self.builder.ins().fcmp(FloatCC::Equal, a, zero_f); @@ -1735,27 +1727,9 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 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 Cases 5 and 6 (b == +/-infinity) are unreachable: the + // out-of-range-exponent guard above already deopts any infinite + // exponent (`fabs(+/-inf) >= 1024.0`). // --- Edge Case 7: a == +infinity → return +infinity let cmp_a_inf = self.builder.ins().fcmp(FloatCC::Equal, a, inf_f); diff --git a/crates/jit/tests/float_tests.rs b/crates/jit/tests/float_tests.rs index d5d763ebc30..b6f0384e40d 100644 --- a/crates/jit/tests/float_tests.rs +++ b/crates/jit/tests/float_tests.rs @@ -130,7 +130,9 @@ mod tests { // 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 + // A NaN base with a bounded, non-zero exponent still reaches Edge + // Case 4 - only the exponent is bounded, not the base. + assert_bits_eq!(pow(f64::NAN, 2.0), Ok(f64::NAN)); //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)); @@ -144,6 +146,8 @@ mod tests { 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)); + // An infinite exponent deoptimizes regardless of the base, the same + // as above - (-infinity) ** (-infinity) moved to deopt_tests.rs. // Test positive float base, positive float exponent assert_approx_eq!(pow(2.0, 2.0), Ok(4.0)); From ade2a6aff9ff40104ffa4b6ce07a6b3326c4d1c9 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 00:30:33 +0900 Subject: [PATCH 25/57] jit: let Strict compile arithmetic now that it deoptimizes Every arithmetic rejection existed because the machine code could trap or wrap. Both are guarded, so Strict and Permissive now differ only in whether a self-call may be compiled into a direct call. Assisted-by: Claude --- crates/jit/src/instructions.rs | 55 ------------------ crates/jit/src/lib.rs | 13 +++-- crates/jit/tests/safety_tests.rs | 97 +++++++++++++++++++------------- extra_tests/snippets/aot.py | 5 +- 4 files changed, 67 insertions(+), 103 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 17e6bf91554..868d5463890 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -141,55 +141,6 @@ pub(crate) const fn instruction_is_supported(instruction: Instruction) -> bool { ) } -/// Whether the machine code emitted for `op` answers the way the interpreter -/// does for every pair of values of these types. -/// -/// Integer arithmetic either traps on overflow and division by zero - and a -/// trap has no handler, so it takes the process down instead of raising - or -/// wraps where Python would widen to an arbitrary-precision integer. -/// Float `/` is a bare `fdiv` with none of the checks that raise -/// ZeroDivisionError, and float `**` neither raises for `0.0 ** -1.0` nor -/// produces the complex result Python gives a negative base. -fn binary_op_is_faithful(op: BinaryOperator, a: Option<&JitType>, b: Option<&JitType>) -> bool { - let traps_or_wraps_on_ints = matches!( - op, - BinaryOperator::Add - | BinaryOperator::InplaceAdd - | BinaryOperator::Subtract - | BinaryOperator::InplaceSubtract - | BinaryOperator::Multiply - | BinaryOperator::InplaceMultiply - | BinaryOperator::TrueDivide - | BinaryOperator::InplaceTrueDivide - | BinaryOperator::FloorDivide - | BinaryOperator::InplaceFloorDivide - | BinaryOperator::Remainder - | BinaryOperator::InplaceRemainder - | BinaryOperator::Power - | BinaryOperator::InplacePower - | BinaryOperator::Lshift - | BinaryOperator::InplaceLshift - | BinaryOperator::Rshift - | BinaryOperator::InplaceRshift - ); - let diverges_on_floats = matches!( - op, - BinaryOperator::TrueDivide - | BinaryOperator::InplaceTrueDivide - | BinaryOperator::Power - | BinaryOperator::InplacePower - ); - - match (a, b) { - (Some(JitType::Int), Some(JitType::Int)) => !traps_or_wraps_on_ints, - // `(Int, Int)` is taken by the arm above, so it cannot land here. - (Some(JitType::Float | JitType::Int), Some(JitType::Float)) - | (Some(JitType::Float), Some(JitType::Int)) => !diverges_on_floats, - // Any other combination has no lowering at all and is rejected anyway. - _ => true, - } -} - impl<'a, 'b> FunctionCompiler<'a, 'b> { pub(crate) fn new( builder: &'a mut FunctionBuilder<'b>, @@ -664,12 +615,6 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let a_type = a.to_jit_type(); let b_type = b.to_jit_type(); - if self.safety == Safety::Strict - && !binary_op_is_faithful(op, a_type.as_ref(), b_type.as_ref()) - { - return Err(JitCompileError::NotSupported); - } - let val = match (op, a, b) { ( BinaryOperator::Add | BinaryOperator::InplaceAdd, diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index dd3fd68447a..0a6a7088f43 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -54,15 +54,16 @@ impl From for JitCompileError { } } -/// How far the compiled code is allowed to diverge from interpreted semantics. +/// Whether a call matched to the function being compiled by global name may +/// become a direct recursive call rather than being left to the interpreter. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Safety { - /// Reject every operation whose machine code can trap, wrap, or otherwise - /// answer differently from the interpreter. Traps have no handler and kill - /// the process, and a wrapped integer is a silently wrong result, so code - /// that was compiled without being asked for must not reach either. + /// Reject a 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 everything the backend supports. + /// Compile a self-call into a direct call, trusting the name to still + /// resolve to this function for as long as the compiled code runs. Permissive, } diff --git a/crates/jit/tests/safety_tests.rs b/crates/jit/tests/safety_tests.rs index 934e9519acd..4f78cc1a2c7 100644 --- a/crates/jit/tests/safety_tests.rs +++ b/crates/jit/tests/safety_tests.rs @@ -2,10 +2,8 @@ mod tests { use rustpython_jit::{JitEngine, Outcome, Safety}; - /// Every operation whose machine code can trap or wrap. There is no trap - /// handler, so a trap kills the process instead of raising. - /// Assert Strict rejects the function *and* that Permissive still accepts - /// it, so the test cannot pass because of an unrelated compile failure. + /// 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(); @@ -30,95 +28,113 @@ mod tests { }}; } + /// 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. + 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) + ), + } + }}; + } + #[test] - fn strict_rejects_int_add() { - assert_strict_rejects!(add => r#" + 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()]); } #[test] - fn strict_rejects_int_multiply() { - assert_strict_rejects!(mul => r#" + 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_rejects_int_floor_divide() { - assert_strict_rejects!(fdiv => r#" + 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_rejects_int_true_divide() { - assert_strict_rejects!(true_divide => r#" + 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_rejects_int_remainder() { - assert_strict_rejects!(rem => r#" + 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_rejects_int_power() { - assert_strict_rejects!(pow => r#" + 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_rejects_int_shift() { - assert_strict_rejects!(shift => r#" + 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] - #[ignore = "Task 7 makes Strict accept this"] - fn strict_rejects_int_negate() { - assert_strict_rejects!(neg => r#" + 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; `fdiv` returns inf. + /// `1.0 / 0.0` raises ZeroDivisionError; a bare `fdiv` would return inf. #[test] - fn strict_rejects_float_divide() { - assert_strict_rejects!(fdiv => r#" + 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()]); } - /// `(-1.0) ** 0.5` is complex in Python and `0.0 ** -1.0` raises. + /// `(-8.0) ** 0.5` is complex in Python, which a compiled `**` cannot + /// produce. #[test] - fn strict_rejects_float_power() { - assert_strict_rejects!(float_power => r#" + 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_rejects_mixed_divide() { - assert_strict_rejects!(mixed => r#" + 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. @@ -174,6 +190,7 @@ def mixed(a: int, b: float) -> float: /// 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#" diff --git a/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py index 2823974ca4f..7f7d06571ba 100644 --- a/extra_tests/snippets/aot.py +++ b/extra_tests/snippets/aot.py @@ -51,8 +51,9 @@ def guarded(a: int) -> int: assert scale(2, 3.0) == 5.0 assert scale(2.0, 3.0) == 5.0 -# Integer arithmetic must widen. Compiled, `+` traps on overflow and `*` wraps, -# so the automatic path has to leave these alone. +# 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. assert wide(3, 4) == 7 assert wide(2**62, 2**62) == 2**63 assert wide(-(2**63), -(2**63)) == -(2**64) From d512846e8ee6d77feb275bcc07d092e3626e9503 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 02:08:04 +0900 Subject: [PATCH 26/57] jit: replace compile_fpow's double-double pow with a call to f64::powf crates/jit/tests/float_tests.rs's basic_power test used a relative-epsilon comparison that absorbed a bug in the double-double ln/exp implementation: compiled float ** lost whole significant digits on a base far from 1 (1023.0 ** 1.0 answered 1022.9277018310074), had no underflow check in dd_exp so an underflowing result came back as a wrong-signed number of the wrong magnitude, and an infinite base with a negative exponent answered +/-inf instead of +/-0.0. compile_fpow now keeps only the two guards that mirror float_pow's special cases - zero base with a negative-sign exponent, and a negative-sign base with a fractional exponent, both read by sign bit rather than by value so a -0.0 operand is caught the same as a negative one - then calls the same f64::powf the interpreter calls through an explicit symbol (jit_powf, registered on the JITBuilder and imported per compiled function) rather than leaving the JIT to resolve pow through the platform's libm. A guard after the call catches a finite-operand result that overflowed to infinity, which raises OverflowError rather than saturating. This removes DDValue and every dd_* helper (dd_from_f64, dd_from_value, dd_from_parts, dd_to_f64, dd_neg, dd_add, dd_sub, dd_mul, dd_mul_f64, dd_scale, dd_ln_1p_series, dd_ln, dd_exp), and the |b| < 1024 exponent bound and negative-zero-base guard added in an earlier round, both now subsumed by powf's own behavior and the guards above. float_tests.rs's basic_power switches from assert_approx_eq! to assert_bits_eq! throughout, since a call to f64::powf is exact by construction; four cases that used to be commented out as wrong or crashing now return the interpreter's exact answer. A new float_power_matches_far_from_one test sweeps 24 (base, exponent) pairs at magnitudes far from 1, 20 of which return and 4 of which overflow to infinity and deoptimize. deopt_tests.rs's float power tests are updated for the new guard shapes: the zero-base and negative-base tests now also cover a -0.0 operand, float_power_deopts_on_finite_base_overflow gains the overflow cases that used to deopt on the removed exponent bound (including the historical 1e100 ** 1e50 crash case), and the tests for the removed exponent bound, NaN exponent, and negative-zero base are gone since those inputs now return rather than deopt. Also: extra_tests/snippets/aot.py's comment on the automatic-compile count is corrected (scale, wide, divide, and the rebound countdown, not just scale); the Safety doc comment in lib.rs now says Strict rejects the whole function containing a self-call rather than just the call; and safety_tests.rs's assert_strict_deopts! macro gains an optional good-input case, used by strict_compiles_int_add, so a regression to an unconditional deopt cannot leave every test in the file green. Assisted-by: Claude --- crates/jit/src/instructions.rs | 604 +++---------------------------- crates/jit/src/lib.rs | 42 ++- crates/jit/tests/deopt_tests.rs | 122 ++----- crates/jit/tests/float_tests.rs | 236 +++++++----- crates/jit/tests/safety_tests.rs | 22 +- extra_tests/snippets/aot.py | 3 +- 6 files changed, 295 insertions(+), 734 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 868d5463890..52d90567cf7 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -66,18 +66,14 @@ impl JitValue { } } -#[derive(Clone)] -struct DDValue { - hi: Value, - lo: Value, -} - 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, + /// `jit_powf`, imported into this function so `compile_fpow` can call it. + powf_func: FuncRef, /// Bytecode offset the instruction being lowered would be re-entered at. resume_offset: u32, stack: Vec, @@ -149,6 +145,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { ret_type: Option, entry_block: Block, safety: Safety, + powf_func: FuncRef, ) -> 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"); @@ -166,6 +163,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { builder, deopt_ptr: *deopt_ptr, deopt_exit: None, + powf_func, resume_offset: 0, stack: Vec::new(), variables: vec![None; num_variables].into_boxed_slice(), @@ -1243,566 +1241,78 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { Ok((quotient, remainder)) } - /// 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, - } + /// The sign bit of an f64, matching `f64::is_sign_negative` - true for + /// `-0.0` and a negative NaN, not just an ordinary negative number. + /// `fcmp` compares values and cannot see this bit; reinterpreting the + /// bits as a signed integer and testing that sign is what + /// `f64::is_sign_negative` itself does. + fn is_sign_negative(&mut self, v: Value) -> Value { + let bits = self.builder.ins().bitcast(types::I64, MemFlags::new(), v); + self.builder.ins().icmp_imm(IntCC::SignedLessThan, bits, 0) } - /// 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, - } - } - - /// 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, returning 0 or 1 - a - /// negative exponent has already deopted. - /// - For a < 0: Adjusts the sign if b is odd - a fractional exponent has - /// already deopted. + /// 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, + /// including the parts that look wrong - the first guard below reads the + /// SIGN BIT of the exponent, not its value, because that is what + /// `is_sign_negative` does, so `0.0 ** -0.0` deopts exactly as + /// `0.0 ** -1.0` does. fn compile_fpow( &mut self, a: Value, b: Value, operands: &[JitValue], ) -> Result { - let f64_ty = types::F64; - let i64_ty = types::I64; 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); - - // Below, `dd_exp` rounds `b * ln|a|` to an i64 exponent with - // `fcvt_to_sint`, which traps once that product leaves i64 range - // (roughly 6.4e18) - and cranelift traps have no handler. A finite - // base's `ln|a|` is at most ~709.8, so bounding `|b|` under 1024 - // keeps the product under ~7.3e5, nowhere near the trap; the two - // `fcvt_to_sint` calls on `b` further down share the same bound. - // `UnorderedOrGreaterThanOrEqual` is the exact negation of an - // ordered `LessThan`, so an unordered (NaN) exponent deopts too - // rather than comparing false on both sides. - let exponent_bound = self.builder.ins().f64const(1024.0); - let abs_b = self.builder.ins().fabs(b); - let exponent_out_of_range = self.builder.ins().fcmp( - FloatCC::UnorderedOrGreaterThanOrEqual, - abs_b, - exponent_bound, - ); - self.deopt_branch(exponent_out_of_range, operands)?; - // 0.0 ** negative raises rather than returning an infinity. + // v1.is_zero() && v2.is_sign_negative() -> ZeroDivisionError. 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); + let exp_sign_negative = self.is_sign_negative(b); + let divides_by_zero = self.builder.ins().band(base_zero, exp_sign_negative); self.deopt_branch(divides_by_zero, operands)?; - // A negative base raised to a fractional power is complex. - let base_negative = self.builder.ins().fcmp(FloatCC::LessThan, a, zero_f); - let truncated = self.builder.ins().trunc(b); - let fractional = self.builder.ins().fcmp(FloatCC::NotEqual, truncated, b); - let complex = self.builder.ins().band(base_negative, fractional); - self.deopt_branch(complex, operands)?; - // A negative zero base loses its sign below - Edge Case 3 always - // returns +0.0 - but `(-0.0) ** 3.0` is `-0.0`. - let base_bits = self.builder.ins().bitcast(types::I64, MemFlags::new(), a); - let base_bits_nonzero = self.builder.ins().icmp_imm(IntCC::NotEqual, base_bits, 0); - let negative_zero_base = self.builder.ins().band(base_zero, base_bits_nonzero); - self.deopt_branch(negative_zero_base, operands)?; - - // 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) is unreachable: the out-of-range-exponent - // guard above already deopts any NaN exponent. - - // --- 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 Cases 5 and 6 (b == +/-infinity) are unreachable: the - // out-of-range-exponent guard above already deopts any infinite - // exponent (`fabs(+/-inf) >= 1024.0`). - - // --- 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, and the fractional-exponent guard above has - // already deopted any negative base (-infinity included) paired - // with a non-integral b, so b is always an integer at this point. + // v1.is_sign_negative() && (v2.floor() - v2).abs() > f64::EPSILON -> + // complex result - a `-0.0` base is caught by its sign bit here too, + // the same as the exponent above. + let base_sign_negative = self.is_sign_negative(a); let b_floor = self.builder.ins().floor(b); - // 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 + let diff = self.builder.ins().fsub(b_floor, b); + let abs_diff = self.builder.ins().fabs(diff); + let epsilon = self.builder.ins().f64const(f64::EPSILON); + let fractional = 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 an integral exponent reaches here - the - // fractional-exponent guard above has already deopted the rest. - self.builder.switch_to_block(a_neg_block); - let b_floor = self.builder.ins().floor(b); - 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()], - ); + .fcmp(FloatCC::GreaterThan, abs_diff, epsilon); + let complex = self.builder.ins().band(base_sign_negative, fractional); + self.deopt_branch(complex, operands)?; - 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()]); + // 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.powf_func, &[a, b]); + let ans = match *self.builder.inst_results(call) { + [ans] => ans, + _ => return Err(JitCompileError::NotSupported), + }; - 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, unless a finite base - // produced an infinity - the interpreter raises OverflowError there - // instead of saturating. An already-infinite base (`inf ** 2.0`) - // must keep returning its infinity, hence the base-finite half. - self.builder.switch_to_block(merge_block); - let result = self.builder.block_params(merge_block)[0]; - let abs_result = self.builder.ins().fabs(result); - let result_infinite = self.builder.ins().fcmp(FloatCC::Equal, abs_result, inf_f); + // 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 base_finite = self.builder.ins().fcmp(FloatCC::LessThan, abs_a, inf_f); - let overflowed = self.builder.ins().band(result_infinite, base_finite); + 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(result) + + Ok(ans) } fn compile_ipow( diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 0a6a7088f43..4ba75db9f50 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -54,19 +54,31 @@ impl From for JitCompileError { } } -/// Whether a call matched to the function being compiled by global name may -/// become a direct recursive call rather than being left to the interpreter. +/// 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 a 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. + /// 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, trusting the name to still /// resolve to this function for as long as the compiled code runs. 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 { @@ -79,17 +91,29 @@ pub enum JitArgumentError { struct Jit { builder_context: FunctionBuilderContext, ctx: codegen::Context, + /// `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(), + powf_func, module: ManuallyDrop::new(module), } } @@ -145,6 +169,9 @@ impl Jit { )?; 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(); @@ -159,6 +186,7 @@ impl Jit { ret, entry_block, safety, + powf_func, ); compiler.compile(func_ref, bytecode)?; diff --git a/crates/jit/tests/deopt_tests.rs b/crates/jit/tests/deopt_tests.rs index 952b2967d5c..ecc24fddd36 100644 --- a/crates/jit/tests/deopt_tests.rs +++ b/crates/jit/tests/deopt_tests.rs @@ -369,59 +369,72 @@ def div(a: float, b: int) -> float: } } - /// `0.0 ** negative` raises rather than returning an infinity. + /// `0.0 ** negative` raises rather than returning an infinity. This + /// reads the sign bit of the exponent, not its value, so `-0.0` deopts + /// it exactly as `-1.0` does. #[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 "# }; - match code.invoke(&[0.0f64.into(), (-1.0f64).into()]) { - Ok(Outcome::Deopt(state)) => { - assert_eq!(state.stack, vec![float(0.0), float(-1.0)]); + for exponent in [-1.0f64, -0.0f64] { + match code.invoke(&[0.0f64.into(), exponent.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![float(0.0), float(exponent)]); + } + other => panic!("expected a deopt for 0.0 ** {exponent}, got {other:?}"), } - other => panic!("expected a deopt, got {other:?}"), } } - /// A negative base raised to a fractional power is complex. + /// A negative base raised to a fractional power is complex. The base is + /// read by its sign bit too, so a `-0.0` base is caught the same way a + /// `-8.0` one is. #[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)]); + for base in [-8.0f64, -0.0f64] { + match code.invoke(&[base.into(), 0.5f64.into()]) { + Ok(Outcome::Deopt(state)) => { + assert_eq!(state.stack, vec![float(base), float(0.5)]); + } + other => panic!("expected a deopt for {base} ** 0.5, got {other:?}"), } - other => panic!("expected a deopt, got {other:?}"), } } - /// `dd_exp` rounds `b * ln|a|` to an i64 with `fcvt_to_sint`, which traps - /// once that leaves i64 range - and cranelift traps have no handler, so - /// `2.0 ** 1e300`, `(-2.0) ** 1e300`, and `1e100 ** 1e50` used to kill - /// the process outright (the last of those was a known crash: an old - /// comment in `float_tests.rs` documented it and commented the case out - /// rather than fixing it). `2.0 ** 1024.0` sits exactly on the bound's - /// threshold, and an infinite exponent is caught by the same bound - /// regardless of how it got there. + /// 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_an_out_of_range_exponent() { + 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), - (-1.0f64, f64::INFINITY), - (0.5f64, f64::INFINITY), - (0.5f64, f64::NEG_INFINITY), - (-1.0f64, f64::NEG_INFINITY), + (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)) => { @@ -432,69 +445,6 @@ def pow(a: float, b: float) -> float: } } - /// The out-of-range bound is written as the negation of an ordered - /// `LessThan`, so an unordered (NaN) exponent deopts the same way a huge - /// one does. `NaN != NaN`, so this gets its own test rather than joining - /// the table above. - #[test] - fn float_power_deopts_on_a_nan_exponent() { - let code = jit_function! { pow => r#" -def pow(a: float, b: float) -> float: - return a ** b -"# }; - match code.invoke(&[1.0f64.into(), f64::NAN.into()]) { - Ok(Outcome::Deopt(state)) => match &state.stack[..] { - [ - StackValue::Value(AbiValue::Float(a)), - StackValue::Value(AbiValue::Float(b)), - ] => { - assert_eq!(*a, 1.0); - assert!(b.is_nan()); - } - other => panic!("unexpected stack shape: {other:?}"), - }, - other => panic!("expected a deopt, got {other:?}"), - } - } - - /// A finite base 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. - #[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()))) - ); - match code.invoke(&[1e308f64.into(), 2.0f64.into()]) { - Ok(Outcome::Deopt(state)) => { - assert_eq!(state.stack, vec![float(1e308), float(2.0)]); - } - other => panic!("expected a deopt, got {other:?}"), - } - } - - /// A negative zero base loses its sign in the zero-base edge case, which - /// always returns `+0.0`; `(-0.0) ** 3.0` is `-0.0`. - #[test] - fn float_power_deopts_on_negative_zero_base() { - let code = jit_function! { pow => r#" -def pow(a: float, b: float) -> float: - return a ** b -"# }; - match code.invoke(&[(-0.0f64).into(), 3.0f64.into()]) { - Ok(Outcome::Deopt(state)) => { - assert_eq!(state.stack, vec![float(-0.0), float(3.0)]); - } - other => panic!("expected a deopt, 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] diff --git a/crates/jit/tests/float_tests.rs b/crates/jit/tests/float_tests.rs index b6f0384e40d..895865e51d3 100644 --- a/crates/jit/tests/float_tests.rs +++ b/crates/jit/tests/float_tests.rs @@ -118,113 +118,167 @@ 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)); - // A NaN base with a bounded, non-zero exponent still reaches Edge - // Case 4 - only the exponent is bounded, not the base. + 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_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)); - // An infinite exponent deoptimizes regardless of the base - see - // deopt_tests.rs. + 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)); - // An infinite exponent deoptimizes regardless of the base, the same - // as above - (-infinity) ** (-infinity) moved to deopt_tests.rs. + 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)); // 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_approx_eq!(pow(-8.0, 2.0), Ok(64.0)); + 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 overflows a finite base to an infinity, which raises - // OverflowError rather than saturating - see deopt_tests.rs. - // 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 overflows the same way as 1e308^2.0 above. - // 1e100^(1e50) has an out-of-range exponent - see deopt_tests.rs. - // 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] diff --git a/crates/jit/tests/safety_tests.rs b/crates/jit/tests/safety_tests.rs index 4f78cc1a2c7..747577fac8a 100644 --- a/crates/jit/tests/safety_tests.rs +++ b/crates/jit/tests/safety_tests.rs @@ -32,7 +32,9 @@ mod tests { /// 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. + /// 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); @@ -44,6 +46,22 @@ mod tests { ), } }}; + ($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] @@ -51,7 +69,7 @@ mod tests { assert_strict_deopts!(add => r#" def add(a: int, b: int) -> int: return a + b -"#, [i64::MAX.into(), 1i64.into()]); +"#, [i64::MAX.into(), 1i64.into()], [3i64.into(), 4i64.into()], 7i64); } #[test] diff --git a/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py index 7f7d06571ba..e996ed95a15 100644 --- a/extra_tests/snippets/aot.py +++ b/extra_tests/snippets/aot.py @@ -142,7 +142,8 @@ def annotated(a: explode()) -> int: if AOT: compiled, rejected, deoptimized = sys._jit._stats() - # `scale` is the one function above the automatic path can take. + # `scale`, `wide`, `divide`, and the rebound `countdown` at line 101 are + # what the automatic path takes above. assert compiled >= 1, (compiled, rejected, deoptimized) # ... and the int argument in `scale(2, 3.0)` handed it back. assert deoptimized >= 1, (compiled, rejected, deoptimized) From d6f289dd4b330d1520dd52427e60cec033f7bec8 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 02:17:19 +0900 Subject: [PATCH 27/57] jit: give aot.py's third widening case its own compiled function, raise the stat floors A deopt discards a PyFunction's compiled code and leaves it permanently interpreted, so the third wide() assertion never touched compiled code after the second one deopted (measured: compiled +0, deopt +0). Moved it to a new wide2() so it gets its own compile attempt. The compiled/deoptimized floors at the bottom were both already met before Strict was allowed to compile arithmetic, so they could not have caught that gate regressing. Raised them to what this file's current functions actually produce (compiled >= 5, deoptimized >= 4); left rejected as a loose floor since it moves with the binary's feature set. Assisted-by: Claude --- extra_tests/snippets/aot.py | 28 ++++++++++++++++++++++------ 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py index e996ed95a15..988612624a1 100644 --- a/extra_tests/snippets/aot.py +++ b/extra_tests/snippets/aot.py @@ -56,7 +56,17 @@ def guarded(a: int) -> int: # back to the interpreter exactly where that widening has to happen. assert wide(3, 4) == 7 assert wide(2**62, 2**62) == 2**63 -assert wide(-(2**63), -(2**63)) == -(2**64) + + +# 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 + + +assert wide2(-(2**63), -(2**63)) == -(2**64) # Shapes with no compiled form at all still behave. assert untyped("a", "b") == "ab" @@ -142,11 +152,17 @@ def annotated(a: explode()) -> int: if AOT: compiled, rejected, deoptimized = sys._jit._stats() - # `scale`, `wide`, `divide`, and the rebound `countdown` at line 101 are - # what the automatic path takes above. - assert compiled >= 1, (compiled, rejected, deoptimized) + # `scale`, `wide`, `wide2`, `divide`, and the rebound `countdown` at line + # 101 are what the automatic path takes above. These are floors, not + # exact counts, but they must not regress: the point of letting Strict + # compile arithmetic was to take shapes it used to refuse outright, and a + # floor already met before that change could not tell if the gate came + # back. `compiled` and `deoptimized` are pinned to what this file + # currently produces (`compiled 5 ... deopt 4`). + assert compiled >= 5, (compiled, rejected, deoptimized) # ... and the int argument in `scale(2, 3.0)` handed it back. - assert deoptimized >= 1, (compiled, rejected, deoptimized) - # Everything else was turned down rather than mis-compiled. + assert deoptimized >= 4, (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) From 3f2b1c889ae24d75e06026c5d0a18ce898093a73 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 02:50:13 +0900 Subject: [PATCH 28/57] vm: resume the interpreter from a deopt record A guard that fires discards the compiled code and hands its record to a fresh frame: the locals as the guard saw them, the operands already on the stack, and the offset of the instruction to re-execute. The callable and the null a call leaves beside it take no slot in the record and are rebuilt from the site's description, which needs LocalsPlus to be able to push a null. Two shapes have no record this frame can use, and both now reach the VM as a third Outcome variant that runs the call again from the start. A self-recursive function shares one deopt buffer down the whole recursion, so a caller leaving because a nested frame gave up used to read that frame's site index; decoded against the same function's site table every type lined up, and resuming would have continued the outermost frame from the deepest frame's offset. The check after a self-call now overwrites the status with a sentinel. 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. Sites are compared against the function's final set at the end of compilation, and only where a backward jump was lowered - without one, execution reaches a guard only by way of offsets below it. Such a local is never read back after a resume, because a read the compiler cannot prove bound is a LOAD_FAST_CHECK and no function containing one compiles at all, but it stays visible through f_locals on a traceback frame. a_caller_stops_when_a_nested_frame_gives_up asserted the inner frame's record for both blow(1) and blow(2); the second is now the restart. Assisted-by: Claude --- crates/jit/src/instructions.rs | 46 +++++++++++++++-- crates/jit/src/lib.rs | 33 +++++++++--- crates/jit/tests/common.rs | 6 +++ crates/jit/tests/deopt_tests.rs | 82 +++++++++++++++++++++++++----- crates/vm/src/builtins/function.rs | 81 +++++++++++++++++++++++++---- crates/vm/src/frame.rs | 8 ++- extra_tests/snippets/jit.py | 77 ++++++++++++++++++++++++++++ 7 files changed, 297 insertions(+), 36 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 52d90567cf7..ba31cc98a1c 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -1,7 +1,7 @@ // spell-checker: disable use super::{ - DEOPT_HEADER_SLOTS, DeoptSite, JitCompileError, JitSig, JitType, MAX_DEOPT_SLOTS, SLOT_SIZE, - Safety, StackEntry, + 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; @@ -84,6 +84,11 @@ pub(crate) struct FunctionCompiler<'a, 'b> { /// 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, @@ -169,6 +174,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 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(), @@ -261,6 +267,9 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 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, }); self.deopt_if(cond, |this| { @@ -523,6 +532,24 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { } } + // 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. @@ -886,7 +913,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { // 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 and leave its record standing. + // that filler, so stop. let status = self.builder.ins().load( types::I64, MemFlags::trusted(), @@ -894,7 +921,17 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 0, ); let nested = self.builder.ins().icmp_imm(IntCC::NotEqual, status, 0); - self.deopt_if(nested, |_| {}); + // 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); @@ -986,6 +1023,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { | Instruction::JumpForward { .. } => { let target = Self::instruction_target(offset, instruction, arg)? .ok_or(JitCompileError::BadBytecode)?; + self.has_backward_jump |= target.as_u32() <= offset; let target_block = self.get_or_create_block(target); self.builder.ins().jump(target_block, &[]); Ok(()) diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 4ba75db9f50..00d727a05eb 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -33,6 +33,9 @@ const SLOT_SIZE: usize = size_of::(); 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); @@ -445,10 +448,17 @@ impl CompiledCode { // SAFETY: the entry point stores the status first thing. let status = unsafe { deopt_ptr.read() }; if status != 0 { - // SAFETY: a non-zero status is written only by a guard of this - // code, and it is the index of the site that describes the record - // that guard just wrote into this buffer. - return Outcome::Deopt(unsafe { self.read_deopt(status, deopt_ptr) }); + // 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]) + .filter(|site| site.resumable); + return match site { + // SAFETY: this is the site whose guard wrote the record, and + // the buffer it wrote is still in scope. + Some(site) => Outcome::Deopt(unsafe { Self::read_deopt(site, deopt_ptr) }), + None => Outcome::Restart, + }; } Outcome::Returned(match self.sig.ret.as_ref() { Some(JitType::None) | None => None, @@ -457,10 +467,8 @@ impl CompiledCode { } /// # Safety - /// `status` must be a status this code's guards can produce, and `deopt` must - /// point at the buffer they wrote. - unsafe fn read_deopt(&self, status: u64, deopt: *const u64) -> DeoptState { - let site = &self.deopt_sites[status as usize - 1]; + /// `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() }; @@ -505,6 +513,11 @@ pub enum Outcome { /// 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. + Restart, } /// Everything the interpreter needs to pick up where the guard stopped. @@ -541,6 +554,10 @@ pub(crate) struct DeoptSite { 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. + pub(crate) resumable: bool, } /// What one value-stack slot holds at a deopt site. Only `Value` occupies a diff --git a/crates/jit/tests/common.rs b/crates/jit/tests/common.rs index e2dd53f1786..a53daea81e1 100644 --- a/crates/jit/tests/common.rs +++ b/crates/jit/tests/common.rs @@ -367,6 +367,9 @@ macro_rules! jit_function { rustpython_jit::Outcome::Deopt(state) => { panic!("jit function unexpectedly deoptimized: {state:?}") } + rustpython_jit::Outcome::Restart => { + panic!("jit function unexpectedly asked to be restarted") + } }) } } @@ -386,6 +389,9 @@ macro_rules! jit_function { rustpython_jit::Outcome::Deopt(state) => { panic!("jit function unexpectedly deoptimized: {state:?}") } + rustpython_jit::Outcome::Restart => { + panic!("jit function unexpectedly asked to be restarted") + } }) } } diff --git a/crates/jit/tests/deopt_tests.rs b/crates/jit/tests/deopt_tests.rs index ecc24fddd36..dda4db9d792 100644 --- a/crates/jit/tests/deopt_tests.rs +++ b/crates/jit/tests/deopt_tests.rs @@ -502,8 +502,15 @@ def countdown(a: int, b: int) -> int: } /// A nested frame that gives up returns a filler in place of a result. The - /// caller has to stop rather than compute on it: the record the callee - /// wrote is the one the interpreter needs. + /// 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(); @@ -520,18 +527,67 @@ def blow(n: int) -> int: code.invoke(&[0i64.into()]), Ok(Outcome::Returned(Some(4611686018427387904i64.into()))) ); - // `blow(1)` overflows. `blow(2)` reaches that overflow one frame down, - // so the guard that fires is the inner frame's and its record is what - // comes back. - for n in [1i64, 2] { - match code.invoke(&[n.into()]) { - Ok(Outcome::Deopt(state)) => assert_eq!( - state.stack, - vec![int(4611686018427387904), int(4611686018427387904)], - "n = {n}" - ), - other => panic!("expected a deopt for n = {n}, got {other:?}"), + 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. + #[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/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index a3c7ab11864..446aad904ac 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -29,7 +29,7 @@ use core::sync::atomic::AtomicU8; use core::sync::atomic::{AtomicU32, Ordering::Relaxed}; use itertools::Itertools; #[cfg(feature = "jit")] -use rustpython_jit::{CompiledCode, Outcome, Safety}; +use rustpython_jit::{CompiledCode, DeoptState, Outcome, Safety, StackValue}; fn format_missing_args( qualname: impl core::fmt::Display, @@ -596,12 +596,46 @@ impl Py { vm.state.aot_stats.deoptimized.fetch_add(1, Relaxed); } + /// 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.lasti.store(state.offset, Relaxed); + } + 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")] + let mut resume = None; #[cfg(feature = "jit")] { let mut state = self.jit_state.load(Relaxed); @@ -620,13 +654,21 @@ impl Py { use crate::convert::ToPyObject; return Ok(ret.to_pyobject(vm)); } - Some(Ok(Outcome::Deopt(_))) => { - // Until the frame can be rebuilt from the record, a guard - // means running the call again from the start. The opcodes - // with a lowering have no effect outside the frame, so - // redoing the work is not observable - but the code has to - // go first, or the second attempt hits the same guard. + 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::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); + resume = Some(state); } Some(Err(err)) => { info!( @@ -677,6 +719,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)); @@ -714,9 +767,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() { diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 60514cf2973..0ff33572256 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"); } diff --git a/extra_tests/snippets/jit.py b/extra_tests/snippets/jit.py index 2fd34a2bbcd..7ccb338fb21 100644 --- a/extra_tests/snippets/jit.py +++ b/extra_tests/snippets/jit.py @@ -83,3 +83,80 @@ def fdiv(a: float, b: float) -> float: 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 local a site cannot + # describe is never read back after a resume, because a read the compiler + # cannot prove bound is a LOAD_FAST_CHECK and no function containing one + # compiles at all; it stays visible through `f_locals` though, so a site + # that cannot describe every bound local restarts the call instead of + # resuming without 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 From 9f586ea1b38a6dfa64a6da764ee9014fc2c13f14 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 02:54:21 +0900 Subject: [PATCH 29/57] jit: name the whole observer class for a non-resumable site A local a site cannot describe is missing from the resumed frame, where anything reading its fastlocals other than a LOAD_FAST sees it go - f_locals, a tracer stepping the frame, a debugger stopped in it. The comments named only the traceback the snippet happens to read it out of. The snippet also now says why it goes through a traceback at all: the obvious `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 observing the drop that way compiles. Also drops the line reference from aot.py's comment on the stat floors, which pointed at the self-recursive countdown that is refused rather than the rebound one that compiles, and moves with every edit to the file. Assisted-by: Claude --- crates/jit/src/lib.rs | 5 ++++- crates/jit/tests/deopt_tests.rs | 7 +++++++ extra_tests/snippets/aot.py | 7 ++++--- extra_tests/snippets/jit.py | 18 ++++++++++++------ 4 files changed, 27 insertions(+), 10 deletions(-) diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 00d727a05eb..509e288cad7 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -556,7 +556,10 @@ pub(crate) struct DeoptSite { 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. + /// 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, } diff --git a/crates/jit/tests/deopt_tests.rs b/crates/jit/tests/deopt_tests.rs index dda4db9d792..1ef381be55e 100644 --- a/crates/jit/tests/deopt_tests.rs +++ b/crates/jit/tests/deopt_tests.rs @@ -542,6 +542,13 @@ def blow(n: int) -> int: /// 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#" diff --git a/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py index 988612624a1..9266513ebc6 100644 --- a/extra_tests/snippets/aot.py +++ b/extra_tests/snippets/aot.py @@ -152,9 +152,10 @@ def annotated(a: explode()) -> int: if AOT: compiled, rejected, deoptimized = sys._jit._stats() - # `scale`, `wide`, `wide2`, `divide`, and the rebound `countdown` at line - # 101 are what the automatic path takes above. These are floors, not - # exact counts, but they must not regress: the point of letting Strict + # `scale`, `wide`, `wide2`, `divide`, and the rebound `countdown` are what + # the automatic path takes above - the original self-recursive one, kept as + # `original_countdown`, is refused. These are floors, not exact + # counts, but they must not regress: the point of letting Strict # compile arithmetic was to take shapes it used to refuse outright, and a # floor already met before that change could not tell if the gate came # back. `compiled` and `deoptimized` are pinned to what this file diff --git a/extra_tests/snippets/jit.py b/extra_tests/snippets/jit.py index 7ccb338fb21..6e334f0d29f 100644 --- a/extra_tests/snippets/jit.py +++ b/extra_tests/snippets/jit.py @@ -134,12 +134,18 @@ def grow(n: int, acc: int) -> int: # 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 local a site cannot - # describe is never read back after a resume, because a read the compiler - # cannot prove bound is a LOAD_FAST_CHECK and no function containing one - # compiles at all; it stays visible through `f_locals` though, so a site - # that cannot describe every bound local restarts the call instead of - # resuming without it. + # 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: From 388ed1cc4f2bb8f41a95325616716a134bf1c2a6 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 07:35:59 +0900 Subject: [PATCH 30/57] jit: document slot 0's third status and Permissive's resolution point The buffer-ABI block described slot 0 as either zero or a site index plus one; a frame that leaves with no record of its own writes a third value. Safety::Permissive scoped its trust to "for as long as the compiled code runs", which a frame resumed from a guard outlives. It now states the contract: the callee is resolved once, when the code is compiled, and a resumed frame keeps that resolution, so rebinding the name is not observed by a call the compiled code had already staged. Rebinding between two whole calls is still observed, because a guard discards the code. Also records why making a callee-bearing site non-resumable is the worse fix: restarting re-runs self-calls that already returned, and those run arbitrary Python. Assisted-by: Claude --- crates/jit/src/lib.rs | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 509e288cad7..2f5cce80e0a 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -23,7 +23,8 @@ 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, otherwise the site index plus one +/// 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 /// ``` @@ -67,8 +68,20 @@ pub enum Safety { /// 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, trusting the name to still - /// resolve to this function for as long as the compiled code runs. + /// 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, } From 78862a2233fecfc789efa718464933e0e1b22ef8 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 07:39:32 +0900 Subject: [PATCH 31/57] extra_tests: assert compiled arithmetic answers what the interpreter does Every guarded operation is checked against the same source run interpreted, including the cases that used to answer wrongly or take the process down. fib_iter(95) is the regression: it overflows partway through and now finishes in the interpreter. fib_iter(80) stays inside 64 bits, so aot=1 runs the compiled loop to completion instead of deoptimizing. 20000 calls, min of 3 runs each, on a shared machine under sustained unrelated load (multiple rustc processes from other sessions, load average 80-100 across 109 users) that did not clear: compiled (aot=1): 0.0042-0.0053s, stable across two independent runs interpreted (aot=0): 0.395-1.92s, varying by ~5x with the contention The compiled figure is a real measurement. The interpreted figure is contention-dominated and is a ceiling on the true interpreted time, not a measurement of it - the two lowest readings taken (0.395s and 1.065s) both come from the same noisy baseline, not from two different speeds. So the speedup is at least two orders of magnitude on this shape; the exact multiple is not measurable on this machine right now. Assisted-by: Claude --- extra_tests/snippets/aot.py | 39 ++++++++---- extra_tests/snippets/jit.py | 122 ++++++++++++++++++++++++++++++++++++ 2 files changed, 150 insertions(+), 11 deletions(-) diff --git a/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py index 9266513ebc6..4e3a43895f9 100644 --- a/extra_tests/snippets/aot.py +++ b/extra_tests/snippets/aot.py @@ -115,6 +115,22 @@ def countdown(a: float) -> float: 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 + + +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 + + # Automatic compilation reads annotations, which under PEP 649 means running # `__annotate__`. A name that is not defined yet raises there, and that is # none of the program's business: it never asked for its annotations. @@ -152,17 +168,18 @@ def annotated(a: explode()) -> int: if AOT: compiled, rejected, deoptimized = sys._jit._stats() - # `scale`, `wide`, `wide2`, `divide`, and the rebound `countdown` are what - # the automatic path takes above - the original self-recursive one, kept as - # `original_countdown`, is refused. These are floors, not exact - # counts, but they must not regress: the point of letting Strict - # compile arithmetic was to take shapes it used to refuse outright, and a - # floor already met before that change could not tell if the gate came - # back. `compiled` and `deoptimized` are pinned to what this file - # currently produces (`compiled 5 ... deopt 4`). - assert compiled >= 5, (compiled, rejected, deoptimized) - # ... and the int argument in `scale(2, 3.0)` handed it back. - assert deoptimized >= 4, (compiled, rejected, deoptimized) + # `scale`, `wide`, `wide2`, `divide`, `fib_iter`, and the rebound + # `countdown` are what the automatic path takes above - the original + # self-recursive one, kept as `original_countdown`, is refused. These + # are floors, not exact counts, but they must not regress: the point of + # letting Strict compile arithmetic was to take shapes it used to + # refuse outright, and a floor already met before that change could not + # tell if the gate came back. `compiled` and `deoptimized` are pinned to + # what this file currently produces (`compiled 6 ... deopt 5`). + assert compiled >= 6, (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) diff --git a/extra_tests/snippets/jit.py b/extra_tests/snippets/jit.py index 6e334f0d29f..3c398fa1704 100644 --- a/extra_tests/snippets/jit.py +++ b/extra_tests/snippets/jit.py @@ -166,3 +166,125 @@ def late(n: int, step: int) -> int: 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) + + 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 From fa4c09eacb6160b0852e0350d76a92cad3da20a1 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 08:14:10 +0900 Subject: [PATCH 32/57] extra_tests: cover binary add and subtract in the equivalence table The check() table had no row for binary a + b or a - b. Add's overflow case was only reached incidentally through fib_iter, whose shape could change without anyone noticing the coverage went with it. Subtract had no Python-level coverage at all: Subtract's own arm calls compile_sub(a, b, ...) in call-site order, and UnaryNegative reaches the same helper through a separate arm with different operand order and arity, so NEG's existing coverage does not stand in for it. Assisted-by: Claude --- extra_tests/snippets/jit.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/extra_tests/snippets/jit.py b/extra_tests/snippets/jit.py index 3c398fa1704..1d91dcd7bf2 100644 --- a/extra_tests/snippets/jit.py +++ b/extra_tests/snippets/jit.py @@ -241,6 +241,18 @@ def call(f): 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)) From 2c7def5ac9562636c2288de20e6f97425160041f Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 08:44:54 +0900 Subject: [PATCH 33/57] jit: refuse a control-flow merge reached with a non-empty stack The abstract value stack is not reconciled where control flow merges. A merged block kept whichever predecessor's operands were lowered last, and its entries are per-path SSA values, so a merge reached with operands live produced a wrong answer rather than a rejection: def g(a, b): return (a if b else b) + (b if a else a) g(1, 2) returned 5 compiled and 3 interpreted. Every edge into a merge now has to arrive with the stack empty, which is what a statement-level `if`, `if`/`else` or `while` has, and what a conditional expression or a short-circuit operator does not. `supports_code` simulates the same depth so the automatic path stops before the backend rather than inside it. Assisted-by: Claude --- crates/jit/src/instructions.rs | 27 +++++++++++++++++++++++++++ crates/jit/src/lib.rs | 29 +++++++++++++++++++++++++---- 2 files changed, 52 insertions(+), 4 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index ba31cc98a1c..0c2034f7b72 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -368,6 +368,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; @@ -462,6 +479,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); @@ -1024,6 +1048,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let target = Self::instruction_target(offset, instruction, arg)? .ok_or(JitCompileError::BadBytecode)?; self.has_backward_jump |= target.as_u32() <= offset; + self.require_empty_stack_at_merge()?; let target_block = self.get_or_create_block(target); self.builder.ins().jump(target_block, &[]); Ok(()) @@ -1097,6 +1122,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let val = self.boolean_val(cond)?; let then_label = Self::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(); @@ -1112,6 +1138,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { let val = self.boolean_val(cond)?; let then_label = Self::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(); diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 2f5cce80e0a..4a4f218d07b 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -403,11 +403,32 @@ pub fn supports_code(code: &bytecode::CodeObject) -> b 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 or a short-circuit operator merges mid-expression, + // where a statement-level `if` or `while` merges with nothing on the stack. + let targets = code.label_targets(); let mut state = bytecode::OpArgState::default(); - code.instructions.iter().all(|&word| { - let (instruction, _) = state.get(word); - instructions::instruction_is_supported(instruction) - }) + let mut depth: i32 = 0; + for (offset, &word) in code.instructions.iter().enumerate() { + let (instruction, arg) = state.get(word); + if !instructions::instruction_is_supported(instruction) { + return false; + } + if depth != 0 + && (targets.contains(&bytecode::Label::from_u32(offset as u32)) + || instruction.label_arg().is_some()) + { + 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; + } + } + true } pub fn compile( From 9311270d55211d56d461bca47a2b588389cb78db Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 08:47:59 +0900 Subject: [PATCH 34/57] vm: check a deopt record against the frame before resuming from it The record's stack was pushed onto a fresh frame slot by slot with no bound check. A stack longer than the code object's `max_stackdepth` runs off the end of the frame, where `push_stack_opt` panics rather than raising, so an ill-formed record aborted the process. The record's local count, stack depth and resume offset are now measured against the code object. A record that does not fit is discarded and the call runs from the start instead. Assisted-by: Claude --- crates/vm/src/builtins/function.rs | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index 446aad904ac..dc732ccba9a 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -596,6 +596,24 @@ impl Py { 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. @@ -668,7 +686,11 @@ impl Py { // recursive call, which would hit the same guard again // if the code were still installed. self.deoptimize(vm); - resume = Some(state); + // 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!( From d4bd40b6253447e80a3c5a5c99b1ee6f3362b323 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 08:52:56 +0900 Subject: [PATCH 35/57] vm: interpret a call while a tracer or a monitoring tool is installed Compiled code runs no frame, so a compiled function reported no call, no line and no return: `sys.settrace`, `sys.setprofile` and `sys.monitoring` all observed nothing for it. The compiled entry now tests `use_tracing` and the monitoring event mask before it is taken, and the call falls through to the interpreter while either is set. The function is left compiled, so it is used again once the tracer is removed. Assisted-by: Claude --- crates/vm/src/builtins/function.rs | 11 ++++++- extra_tests/snippets/jit.py | 49 ++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index dc732ccba9a..8882843ed80 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -655,7 +655,16 @@ impl Py { #[cfg(feature = "jit")] 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::compile_on_first_call(self, vm)?; diff --git a/extra_tests/snippets/jit.py b/extra_tests/snippets/jit.py index 1d91dcd7bf2..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 @@ -300,3 +303,49 @@ def countdown(a: int, b: int) -> int: 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 From 10ef3a27ff9d8adbf6a406985e78d036a282528c Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 08:53:48 +0900 Subject: [PATCH 36/57] jit: document the safety of the checked invoke path `CompiledCode::invoke` checked the arity and every slot's type and then called `invoke_raw` with no note of it. `Args::invoke` already carries the equivalent comment. Assisted-by: Claude --- crates/jit/src/lib.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 4a4f218d07b..1f562e0fddb 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -464,6 +464,8 @@ impl CompiledCode { 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) }) } From 9628fe2bbae6b39d657d57275abbafbc8e620e4f Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 08:53:55 +0900 Subject: [PATCH 37/57] ci: build the aot step against the locked dependency versions Every other build step in the job passes `--locked`. Without it this step may resolve newer dependency versions than the lockfile pins. Assisted-by: Claude --- .github/workflows/ci.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 38a09553527..73bf48fce66 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -151,7 +151,7 @@ jobs: # snippet has to pass identically with it on and off. - name: Test aot build run: | - cargo build --features aot + 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 From ec93de2969067c04eb4a6d3be355d60edadef146 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 10:02:05 +0900 Subject: [PATCH 38/57] jit: make the pre-filter's merge rejection actually fire `supports_code` detected a branch with `code.label_targets()` and `Instruction::label_arg`. Neither answers the question: a jump argument is a delta, `label_targets` collects that delta rather than the offset it points at, and `label_arg` reported `None` for the conditional jumps here. Both halves of the merge clause were therefore dead, and the pre-filter passed every mid-expression merge on to the backend, which rejected it. The walk now resolves targets with `instruction_target`, the function the compiler resolves them with, over the same de-specialized stream the compiler consumes. `instruction_target` and the two `jump_target_*` helpers move out of `FunctionCompiler` to be reachable from here. A jump's outgoing edge is checked after the instruction's stack effect rather than before, since a conditional jump has popped the condition it tested by the time control leaves it - which is the depth the compiler checks, and what a `while` loop's `POP_JUMP_IF_FALSE` needs to pass. Assisted-by: Claude --- crates/jit/src/instructions.rs | 106 ++++++++++++++++----------------- crates/jit/src/lib.rs | 50 +++++++++++++--- 2 files changed, 94 insertions(+), 62 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 0c2034f7b72..5680246e637 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -101,6 +101,54 @@ pub(crate) struct FunctionCompiler<'a, 'b> { /// 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. +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) +} + pub(crate) const fn instruction_is_supported(instruction: Instruction) -> bool { matches!( instruction, @@ -394,56 +442,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, @@ -462,7 +460,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); } } @@ -1045,7 +1043,7 @@ 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)?; self.has_backward_jump |= target.as_u32() <= offset; self.require_empty_stack_at_merge()?; @@ -1120,7 +1118,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 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); @@ -1136,7 +1134,7 @@ 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); diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 1f562e0fddb..5a380b08374 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -2,6 +2,7 @@ mod instructions; extern crate alloc; +use alloc::collections::BTreeSet; use alloc::fmt; use alloc::sync::Arc; use core::mem::{self, ManuallyDrop}; @@ -406,27 +407,60 @@ pub fn supports_code(code: &bytecode::CodeObject) -> b // 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 or a short-circuit operator merges mid-expression, - // where a statement-level `if` or `while` merges with nothing on the stack. - let targets = code.label_targets(); + // 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 code.instructions.iter().enumerate() { + for (offset, &word) in clean.iter().enumerate() { let (instruction, arg) = state.get(word); if !instructions::instruction_is_supported(instruction) { return false; } - if depth != 0 - && (targets.contains(&bytecode::Label::from_u32(offset as u32)) - || instruction.label_arg().is_some()) - { + // 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 } From 3f6f7286f8677f3781be05e7cb7cade159c99041 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 10:02:11 +0900 Subject: [PATCH 39/57] jit: pin the pre-filter's rejection of a mid-expression merge `branches_and_loops_are_supported` covers the shapes the merge clause has to keep accepting; nothing covered the shapes it has to reject, so the clause could stop firing with every test still passing. Assisted-by: Claude --- crates/jit/tests/support_tests.rs | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/crates/jit/tests/support_tests.rs b/crates/jit/tests/support_tests.rs index 7b8b6fa5213..3c2f2f0048f 100644 --- a/crates/jit/tests/support_tests.rs +++ b/crates/jit/tests/support_tests.rs @@ -43,6 +43,28 @@ def count(n: int) -> int: "#); } + #[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#" From 46802e3ef863af2c00e4068b3115c954ab83f2a6 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 20:20:08 +0900 Subject: [PATCH 40/57] vm: report the automatic compiler from sys._jit `is_available` answers for the `aot` feature rather than `jit`, which is what `is_enabled` reports and what the switches below reach; explicit `__jit__()` needs only `jit` and is a separate path. PYTHON_JIT joins RUSTPYTHON_AOT as a spelling of the same switch. Both are read before the -X options so an explicit `-X aot` wins, and the flag is masked with the feature so nothing can turn on a compiler that was not built in. Assisted-by: Claude --- crates/vm/src/stdlib/sys.rs | 5 ++++- src/settings.rs | 37 +++++++++++++++++++++++-------------- 2 files changed, 27 insertions(+), 15 deletions(-) diff --git a/crates/vm/src/stdlib/sys.rs b/crates/vm/src/stdlib/sys.rs index c31004a40fb..10f8b163ac8 100644 --- a/crates/vm/src/stdlib/sys.rs +++ b/crates/vm/src/stdlib/sys.rs @@ -15,7 +15,10 @@ mod sys_jit { /// and False otherwise. #[pyfunction] const fn is_available() -> bool { - cfg!(feature = "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, diff --git a/src/settings.rs b/src/settings.rs index b4df5f8fec2..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() @@ -356,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. @@ -369,20 +392,6 @@ pub fn parse_opts() -> Result<(Settings, RunMode), lexopt::Error> { if env_bool("PYTHONNODEBUGRANGES") { settings.code_debug_ranges = false; } - if let Some(val) = get_env("RUSTPYTHON_AOT") { - settings.aot = match val.to_str() { - Some("1") => true, - Some("0") => false, - _ => { - error!( - "Fatal Python error: config_init_aot: \ - RUSTPYTHON_AOT=N: N is missing or invalid\n\ - Python runtime state: preinitialized" - ); - std::process::exit(1); - } - }; - } if let Some(val) = get_env("PYTHON_THREAD_INHERIT_CONTEXT") { settings.thread_inherit_context = match val.to_str() { Some("1") => true, From 15c17398347f2d1e213d7d9d29edaf3e1565a29f Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 20:20:25 +0900 Subject: [PATCH 41/57] jit: move a doc comment onto the function it describes The comment for `instruction_is_supported` was left on `jump_target_forward` when both were lifted out of the impl block. Also collapses a run of spaces inside a `concat!`-assembled panic message. Assisted-by: Claude --- crates/jit/src/instructions.rs | 14 +++++++------- crates/jit/tests/safety_tests.rs | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 5680246e637..e3df9263636 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -94,13 +94,6 @@ pub(crate) struct FunctionCompiler<'a, 'b> { pub(crate) deopt_sites: Vec, } -/// 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. fn jump_target_forward(offset: u32, caches: u32, arg: OpArg) -> Result { let after = offset .checked_add(1) @@ -149,6 +142,13 @@ pub(crate) fn instruction_target( 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, diff --git a/crates/jit/tests/safety_tests.rs b/crates/jit/tests/safety_tests.rs index 747577fac8a..56eb09c7a1a 100644 --- a/crates/jit/tests/safety_tests.rs +++ b/crates/jit/tests/safety_tests.rs @@ -14,7 +14,7 @@ mod tests { ); 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" + " is only meant to be rejected for being unsafe, but Permissive cannot compile it either" )); }}; } From 713eab5e60c67d4a78a7c2b4e54dbc85e09296a5 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 20:20:34 +0900 Subject: [PATCH 42/57] jit: keep a true-division operand of 1 << 53 compiled The guard deoptimizes where converting an operand to a double is inexact, but `1 << 53` is the largest magnitude that still converts exactly, so the bound is strict rather than inclusive. `iabs` leaves `i64::MIN` negative and the unsigned comparison reads that as `1 << 63`, which stays past the bound. Assisted-by: Claude --- crates/jit/src/instructions.rs | 8 ++++++-- crates/jit/tests/deopt_tests.rs | 21 +++++++++++++++------ crates/jit/tests/int_tests.rs | 5 +++-- 3 files changed, 24 insertions(+), 10 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index e3df9263636..3aa150b68b7 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -701,11 +701,15 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { // 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. + // 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::UnsignedGreaterThanOrEqual, + IntCC::UnsignedGreaterThan, magnitude, 1 << 53, ) diff --git a/crates/jit/tests/deopt_tests.rs b/crates/jit/tests/deopt_tests.rs index 1ef381be55e..5c348cc390e 100644 --- a/crates/jit/tests/deopt_tests.rs +++ b/crates/jit/tests/deopt_tests.rs @@ -243,21 +243,30 @@ def div(a: int, b: int) -> float: 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) - 1).into(), 1i64.into()]), - Ok(Outcome::Returned(Some((((1i64 << 53) - 1) as f64).into()))) + code.invoke(&[(1i64 << 53).into(), 1i64.into()]), + Ok(Outcome::Returned(Some(((1i64 << 53) as f64).into()))) ); - match code.invoke(&[(1i64 << 53).into(), 1i64.into()]) { + match code.invoke(&[((1i64 << 53) + 1).into(), 1i64.into()]) { Ok(Outcome::Deopt(state)) => { - assert_eq!(state.stack, vec![int(1i64 << 53), int(1)]); + 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).into()]) { + match code.invoke(&[1i64.into(), ((1i64 << 53) + 1).into()]) { Ok(Outcome::Deopt(state)) => { - assert_eq!(state.stack, vec![int(1), int(1i64 << 53)]); + 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:?}"), } diff --git a/crates/jit/tests/int_tests.rs b/crates/jit/tests/int_tests.rs index 4150037a8ac..2f2a4ac5009 100644 --- a/crates/jit/tests/int_tests.rs +++ b/crates/jit/tests/int_tests.rs @@ -95,8 +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)); - // An operand at or past `1 << 53` does not fit a double's - // significand and deoptimizes instead; see deopt_tests.rs. + 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] From 9de0d3cae77a2ff8f1eab18ecb9c070bb89e0926 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 20:20:48 +0900 Subject: [PATCH 43/57] extra_tests: assert scale's value instead of a division it does not do `scale` multiplies and adds, so the ZeroDivisionError the block guarded against could never be raised and the check could not fail. The float division by zero it described is already covered by `divide` below. Assisted-by: Claude --- extra_tests/snippets/aot.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py index 4e3a43895f9..6371efa51e0 100644 --- a/extra_tests/snippets/aot.py +++ b/extra_tests/snippets/aot.py @@ -76,11 +76,8 @@ def wide2(a: int, b: int) -> int: assert list(generator(7)) == [7] assert guarded(3) == 3 -# Division by zero raises rather than returning inf or killing the process. -try: - scale(1.0, 0.0) -except ZeroDivisionError: - raise AssertionError("scale does not divide") +# 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: From 83d4ae1221289b6bafb8b4e47eb00b86de54f673 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 20:20:49 +0900 Subject: [PATCH 44/57] ci: build the benchmarks with automatic compilation on Unlike the default build, so the benchmarks measure what the eligibility check costs on the first call of every function. Assisted-by: Claude --- .github/workflows/codspeed.yaml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/codspeed.yaml b/.github/workflows/codspeed.yaml index a9bc341408e..2b83ba8e6ad 100644 --- a/.github/workflows/codspeed.yaml +++ b/.github/workflows/codspeed.yaml @@ -62,8 +62,10 @@ jobs: - name: Install cargo-codspeed run: cargo install cargo-codspeed --locked --version ^5 + # Benchmarks run with automatic compilation on, unlike the default + # build, so its cost on the first call of every function is measured. - name: Build the benchmark targets - run: cargo codspeed build --locked --measurement-mode simulation -p rustpython -p rustpython-sre_engine + run: cargo codspeed build --locked --measurement-mode simulation --features rustpython/aot -p rustpython -p rustpython-sre_engine - name: Save cache if: ${{ github.ref == 'refs/heads/main' }} # only save on main From 5a89f78273b75f0572dee6ffeec8d14d395d71d1 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Mon, 31 Aug 2026 20:41:53 +0900 Subject: [PATCH 45/57] vm: record the traced line even where a frame carries no trace function Computing f_lineno from the instruction pointer removed the per-instruction prev_line update, but prev_line is still the key the line event de-duplicates on. It then only advanced when an event fired, so a frame that started being traced part way through - a caller with no trace function of its own, stepped into from a callee's return - compared the line it was already on against a stale value and reported it as new. The line is now recorded for every instruction executed while a tracer is installed, whether or not this frame fires events. Untraced execution, which is what dropping the update was for, is unaffected. Assisted-by: Claude --- crates/vm/src/frame.rs | 38 +++++++++++++++++++++----------------- 1 file changed, 21 insertions(+), 17 deletions(-) diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 0ff33572256..f403080640f 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -3087,29 +3087,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); From 2fcd657024af56d8522932422ef0e64a99dee576 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Wed, 2 Sep 2026 14:22:35 +0900 Subject: [PATCH 46/57] vm: attach the frame object to the caller when an escaped frame returns `exit_iframe` recorded the caller of a returning materialized frame through `materialize_chain`, which hands back a standalone copy when the caller has no frame object of its own. Nothing links that copy to the caller, so when the caller returns in turn it finds no frame object to extend, and the chain ends one link up: `f_back` reached the caller and then None. Materialize the caller instead, so the caller runs the same block on its own return and links further up, and drop `materialize_chain`, which had no other caller. `materialize_slow_chain` keeps making detached copies for the cross-thread reader that wants them, and is now `threading`-only. Assisted-by: Claude --- crates/vm/src/frame.rs | 23 ++++------------------- crates/vm/src/vm/mod.rs | 7 ++++++- extra_tests/snippets/aot.py | 24 ++++++++++++++++++++++++ 3 files changed, 34 insertions(+), 20 deletions(-) diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index f403080640f..a4a13a9ea1b 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -1182,21 +1182,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. @@ -1291,10 +1276,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(); diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index c7bf04a47be..0bd774aadc2 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -2425,8 +2425,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/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py index 6371efa51e0..34b15728d15 100644 --- a/extra_tests/snippets/aot.py +++ b/extra_tests/snippets/aot.py @@ -163,6 +163,30 @@ def annotated(a: explode()) -> int: raise AssertionError("expected a BaseException to reach the caller") +# 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 + + if AOT: compiled, rejected, deoptimized = sys._jit._stats() # `scale`, `wide`, `wide2`, `divide`, `fib_iter`, and the rebound From 32cf3a0a5f83c18cfbaf8b0f47503abe5fa07253 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Wed, 2 Sep 2026 14:22:43 +0900 Subject: [PATCH 47/57] vm: resolve f_back from the thread still running the frame MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A frame object materialized from a running frame carries no `previous` of its own, and the walk that stands in for it only covers this thread's chain. A frame belonging to another thread therefore reported no caller at all — `sys._current_exceptions()` handed back a traceback whose frame was the whole stack. `attached_tid` already names the thread that is still running the source frame. Take it as the gate: stop the world, find the source on that thread's published chain, and answer from its caller, materializing the rest of that chain when the caller has no frame object yet. Assisted-by: Claude --- crates/vm/src/builtins/frame.rs | 79 +++++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/crates/vm/src/builtins/frame.rs b/crates/vm/src/builtins/frame.rs index 07a9785bf7c..90a2cc6acb7 100644 --- a/crates/vm/src/builtins/frame.rs +++ b/crates/vm/src/builtins/frame.rs @@ -462,6 +462,45 @@ 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) + } } #[pyclass(flags(DISALLOW_INSTANTIATION), with(Py))] @@ -751,6 +790,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] @@ -881,6 +956,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; } } From a40ba854b7ac66c3f0eddbc887183e294806dc57 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Wed, 2 Sep 2026 14:22:52 +0900 Subject: [PATCH 48/57] vm: read a running frame's position from the thread running it `f_lasti` and `f_lineno` read the live instruction pointer through the walk of this thread's chain, and fell back to the frame object's own copy when it missed. That copy only catches up when the source frame returns, so a frame belonging to another thread reported wherever it stood when it was materialized rather than where it is now. Give both getters the same treatment `f_back` has: consult the thread named by `attached_tid` under stop-the-world. The shared read is now `live_lasti`, which is what each getter was open-coding. Assisted-by: Claude --- crates/capi/src/pyframe.rs | 4 +-- crates/vm/src/builtins/frame.rs | 59 +++++++++++++++++++++------------ crates/vm/src/warn.rs | 2 +- 3 files changed, 41 insertions(+), 24 deletions(-) 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/vm/src/builtins/frame.rs b/crates/vm/src/builtins/frame.rs index 90a2cc6acb7..e62f2c75fe5 100644 --- a/crates/vm/src/builtins/frame.rs +++ b/crates/vm/src/builtins/frame.rs @@ -501,6 +501,38 @@ impl FrameObject { 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))] @@ -521,33 +553,18 @@ 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 { + 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. The live - // position has to be read rather than this object's copy: a frame - // observed from inside a call it made still reports that call's line. - let live = self.find_live_source_iframe(); - let lasti = if live.is_null() { - self.lasti() - } else { - unsafe { (*live).lasti.load(Relaxed) } - }; + // 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(); 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 From 39fb1ebd9bc8bba192b7b72bb135b0f0de31d78f Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Wed, 2 Sep 2026 14:22:53 +0900 Subject: [PATCH 49/57] ci: run part of the CPython suite with automatic compilation on The aot job ran two snippets, which cover arithmetic and nothing else. The frame, traceback and tracing machinery that the automatic call path moves through went untested with the feature on. Assisted-by: Claude --- .github/workflows/ci.yaml | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 73bf48fce66..afaff27ed8a 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -157,6 +157,15 @@ jobs: 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 and tracing. + - 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 + 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' From 30ef72b1d9b4ecdd4b7585a3adca8cdaad1b1b4b Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Wed, 2 Sep 2026 22:18:33 +0900 Subject: [PATCH 50/57] vm: set a resumed frame's offset through a method on the frame `fill_locals_from_deopt` stored `lasti` directly, which needs the atomic trait in scope where that field is a plain cell. `InterpreterFrame` now carries the setter its getter was already paired with, and a build with `jit` but without `threading` compiles again. Assisted-by: Claude --- crates/vm/src/builtins/function.rs | 2 +- crates/vm/src/frame.rs | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index 8882843ed80..84a07fd5088 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -642,7 +642,7 @@ impl Py { }; iframe.localsplus.push_stack_opt(value); } - iframe.lasti.store(state.offset, Relaxed); + iframe.set_lasti(state.offset); } pub fn invoke_with_locals( diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index a4a13a9ea1b..017a2b1944a 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -1097,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 { From 8a8fa85ee598fc39b70af119e968f71d46dd62ac Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Wed, 2 Sep 2026 22:20:06 +0900 Subject: [PATCH 51/57] jit: poll the eval breaker at every backward jump A compiled loop ran to its end whatever the interpreter wanted of the thread. It answered no pending signal, parked for no stop-the-world, and did not stop when the interpreter began shutting down, so a `gc.collect()` from another thread or the shutdown that joins a daemon thread waited on a thread that could not reach a safepoint. The process hung rather than being delayed. Every reason a thread has to leave the bytecode loop now sets a bit in the eval-breaker word before it becomes true of any single thread: a stop-the-world span while it is open, and shutdown once it starts. The word is process-global, so a bit says a thread may have to stop, not which one. A backward jump loads that byte, and leaves through the deopt exit when it is not zero. Leaving that way is no verdict on the code, so the code stays installed and only the interrupted call finishes interpreted: from the record where the site can describe the frame, and from the start where it cannot. `Outcome` gains `Interrupted` to tell the two apart from a guard, which still drops the code. The engine is given the word to poll when it is built, and compiles no poll at all without one. The fall-through costs a load, a compare and a branch per iteration. Assisted-by: Claude --- .github/workflows/ci.yaml | 5 +- crates/jit/src/instructions.rs | 61 +++++++++++-- crates/jit/src/lib.rs | 54 +++++++++--- crates/jit/tests/common.rs | 6 ++ crates/jit/tests/deopt_tests.rs | 4 +- crates/jit/tests/engine_tests.rs | 8 +- crates/jit/tests/lib.rs | 1 + crates/jit/tests/misc_tests.rs | 2 +- crates/jit/tests/safepoint_tests.rs | 127 ++++++++++++++++++++++++++++ crates/jit/tests/safety_tests.rs | 4 +- crates/vm/src/builtins/function.rs | 10 +++ crates/vm/src/signal.rs | 56 ++++++++++++ crates/vm/src/vm/interpreter.rs | 7 +- crates/vm/src/vm/mod.rs | 8 ++ extra_tests/snippets/aot.py | 40 ++++++++- 15 files changed, 359 insertions(+), 34 deletions(-) create mode 100644 crates/jit/tests/safepoint_tests.rs diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index afaff27ed8a..e548239cb97 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -158,12 +158,13 @@ jobs: if: runner.os == 'Linux' # The snippets cover arithmetic. These cover the rest of what the - # automatic call path moves through: frames, tracebacks and tracing. + # 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_bdb test_trace test_exceptions test_generators test_threading if: runner.os == 'Linux' # - name: Install tk-dev for tkinter build diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 3aa150b68b7..422c8461674 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -66,14 +66,23 @@ impl JitValue { } } +/// 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, - /// `jit_powf`, imported into this function so `compile_fpow` can call it. - powf_func: FuncRef, + /// 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, @@ -198,7 +207,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { ret_type: Option, entry_block: Block, safety: Safety, - powf_func: FuncRef, + 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"); @@ -216,7 +225,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { builder, deopt_ptr: *deopt_ptr, deopt_exit: None, - powf_func, + externals, resume_offset: 0, stack: Vec::new(), variables: vec![None; num_variables].into_boxed_slice(), @@ -318,6 +327,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { // 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| { @@ -345,6 +355,39 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { 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 @@ -1049,8 +1092,14 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { | Instruction::JumpForward { .. } => { let target = instruction_target(offset, instruction, arg)? .ok_or(JitCompileError::BadBytecode)?; - self.has_backward_jump |= target.as_u32() <= offset; + 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(()) @@ -1358,7 +1407,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { // 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.powf_func, &[a, b]); + 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), diff --git a/crates/jit/src/lib.rs b/crates/jit/src/lib.rs index 5a380b08374..f5b35947608 100644 --- a/crates/jit/src/lib.rs +++ b/crates/jit/src/lib.rs @@ -6,11 +6,11 @@ use alloc::collections::BTreeSet; use alloc::fmt; use alloc::sync::Arc; use core::mem::{self, ManuallyDrop}; -use core::sync::atomic::{AtomicU64, Ordering}; +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}; @@ -145,8 +145,9 @@ impl Jit { ret: Option, unique: u64, safety: Safety, + safepoint: Option<&'static AtomicU8>, ) -> Result<(FuncId, JitSig, Vec), JitCompileError> { - let result = self.build_function_inner(bytecode, args, ret, unique, safety); + 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 @@ -164,6 +165,7 @@ impl Jit { 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 @@ -203,7 +205,10 @@ impl Jit { ret, entry_block, safety, - powf_func, + Externals { + powf: powf_func, + safepoint, + }, ); compiler.compile(func_ref, bytecode)?; @@ -322,14 +327,21 @@ impl Jit { 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() -> Arc { + pub fn new(safepoint: Option<&'static AtomicU8>) -> Arc { Arc::new(Self { jit: Mutex::new(Jit::new()), next_id: AtomicU64::new(0), + safepoint, }) } @@ -347,7 +359,8 @@ impl JitEngine { // 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)?; + 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); @@ -470,7 +483,7 @@ pub fn compile( args: &[JitType], ret: Option, ) -> Result { - JitEngine::new().compile(bytecode, args, ret, Safety::Permissive) + JitEngine::new(None).compile(bytecode, args, ret, Safety::Permissive) } pub struct CompiledCode { @@ -520,13 +533,16 @@ impl CompiledCode { 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]) - .filter(|site| site.resumable); - return match site { + 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. - Some(site) => Outcome::Deopt(unsafe { Self::read_deopt(site, deopt_ptr) }), + 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, }; } @@ -586,8 +602,15 @@ pub enum Outcome { /// 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. + /// 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. @@ -631,6 +654,11 @@ pub(crate) struct DeoptSite { /// 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 diff --git a/crates/jit/tests/common.rs b/crates/jit/tests/common.rs index a53daea81e1..deac8b78370 100644 --- a/crates/jit/tests/common.rs +++ b/crates/jit/tests/common.rs @@ -370,6 +370,9 @@ macro_rules! jit_function { rustpython_jit::Outcome::Restart => { panic!("jit function unexpectedly asked to be restarted") } + rustpython_jit::Outcome::Interrupted(state) => { + panic!("jit function unexpectedly interrupted: {state:?}") + } }) } } @@ -392,6 +395,9 @@ macro_rules! jit_function { 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 index 5c348cc390e..79c3bbffc9b 100644 --- a/crates/jit/tests/deopt_tests.rs +++ b/crates/jit/tests/deopt_tests.rs @@ -489,7 +489,7 @@ def f(a: int, b: int, c: bool) -> int: /// function stays compilable. #[test] fn a_guard_under_a_self_call_describes_the_callable() { - let engine = JitEngine::new(); + let engine = JitEngine::new(None); let f = py_function_def! { countdown => r#" def countdown(a: int, b: int) -> int: if a < 0: @@ -522,7 +522,7 @@ def countdown(a: int, b: int) -> int: /// It has to come back as a restart instead. #[test] fn a_caller_stops_when_a_nested_frame_gives_up() { - let engine = JitEngine::new(); + let engine = JitEngine::new(None); let f = py_function_def! { blow => r#" def blow(n: int) -> int: if n == 0: diff --git a/crates/jit/tests/engine_tests.rs b/crates/jit/tests/engine_tests.rs index 009fecfecf0..b4e9bad2626 100644 --- a/crates/jit/tests/engine_tests.rs +++ b/crates/jit/tests/engine_tests.rs @@ -6,7 +6,7 @@ mod tests { /// be made unique before they can live in one engine. #[test] fn same_name_functions_coexist() { - let engine = JitEngine::new(); + let engine = JitEngine::new(None); let first = py_function_def!(foo => r#" def foo(a: int, b: int) -> int: return a @@ -36,7 +36,7 @@ mod tests { /// 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(); + let engine = JitEngine::new(None); let unsupported = py_function_def!(unsupported => r#" def unsupported(a: int) -> int: return [a] @@ -60,7 +60,7 @@ mod tests { #[test] fn compiled_code_keeps_engine_alive() { let code = { - let engine = JitEngine::new(); + let engine = JitEngine::new(None); let f = py_function_def!(f => r#" def f(a: int) -> int: return a @@ -100,7 +100,7 @@ mod tests { /// buffer is turned down rather than compiled into an overrun. #[test] fn parameters_beyond_the_buffer_are_rejected() { - let engine = JitEngine::new(); + 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 diff --git a/crates/jit/tests/lib.rs b/crates/jit/tests/lib.rs index 58697e5dc3e..a86a978eed8 100644 --- a/crates/jit/tests/lib.rs +++ b/crates/jit/tests/lib.rs @@ -9,5 +9,6 @@ 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 d8eda16e90c..2002f733d51 100644 --- a/crates/jit/tests/misc_tests.rs +++ b/crates/jit/tests/misc_tests.rs @@ -138,7 +138,7 @@ mod tests { /// whether the backend can read a partially-defined local. #[test] fn conditionally_defined_local_is_not_compiled() { - let engine = JitEngine::new(); + let engine = JitEngine::new(None); let f = py_function_def!(f => r#" def f(c: bool) -> int: if c: 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 index 56eb09c7a1a..cd22d6c53b3 100644 --- a/crates/jit/tests/safety_tests.rs +++ b/crates/jit/tests/safety_tests.rs @@ -6,7 +6,7 @@ mod tests { /// so the test cannot pass because of an unrelated compile failure. macro_rules! assert_strict_rejects { ($name:ident => $src:expr) => {{ - let engine = JitEngine::new(); + let engine = JitEngine::new(None); let f = py_function_def!($name => $src); assert!( f.compile_on(&engine, Safety::Strict).is_err(), @@ -21,7 +21,7 @@ mod tests { macro_rules! assert_accepted { ($safety:expr, $name:ident => $src:expr) => {{ - let engine = JitEngine::new(); + let engine = JitEngine::new(None); let f = py_function_def!($name => $src); f.compile_on(&engine, $safety) .expect(concat!(stringify!($name), " should compile")) diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index 84a07fd5088..423c842d8a5 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -689,6 +689,16 @@ impl Py { // 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 diff --git a/crates/vm/src/signal.rs b/crates/vm/src/signal.rs index d1ddd8c3cdd..33678bfbaf5 100644 --- a/crates/vm/src/signal.rs +++ b/crates/vm/src/signal.rs @@ -28,6 +28,19 @@ 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; +/// The interpreter is shutting down, so every thread but the main one must +/// stop running bytecode. Set once and never cleared. +#[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); #[expect( clippy::declare_interior_mutable_const, @@ -147,6 +160,49 @@ 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); +} + +/// Record that the interpreter has started shutting down. +#[cfg(feature = "threading")] +pub(crate) fn set_finalizing_bit() { + EVAL_BREAKER.fetch_or(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/vm/interpreter.rs b/crates/vm/src/vm/interpreter.rs index 7e16a0ee636..acdcead72b2 100644 --- a/crates/vm/src/vm/interpreter.rs +++ b/crates/vm/src/vm/interpreter.rs @@ -163,7 +163,7 @@ where let global_state = PyRc::new(PyGlobalState { gc: crate::gc_state::GcInterpreterState::new(&ctx), #[cfg(feature = "jit")] - jit_engine: rustpython_jit::JitEngine::new(), + jit_engine: rustpython_jit::JitEngine::new(Some(crate::signal::eval_breaker_word())), #[cfg(feature = "jit")] aot_stats: Default::default(), interpreter_id, @@ -697,6 +697,11 @@ 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. + #[cfg(feature = "threading")] + crate::signal::set_finalizing_bit(); // 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 0bd774aadc2..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); diff --git a/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py index 34b15728d15..915d64e4e4b 100644 --- a/extra_tests/snippets/aot.py +++ b/extra_tests/snippets/aot.py @@ -187,17 +187,51 @@ def outermost(): 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 + + +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`, and the rebound + # `scale`, `wide`, `wide2`, `divide`, `fib_iter`, `spin`, and the rebound # `countdown` are what the automatic path takes above - the original # self-recursive one, kept as `original_countdown`, is refused. These # are floors, not exact counts, but they must not regress: the point of # letting Strict compile arithmetic was to take shapes it used to # refuse outright, and a floor already met before that change could not # tell if the gate came back. `compiled` and `deoptimized` are pinned to - # what this file currently produces (`compiled 6 ... deopt 5`). - assert compiled >= 6, (compiled, rejected, deoptimized) + # what this file currently produces (`compiled 7 ... deopt 5`). + 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) From 053f6303c4e2765f609c04bbe58a966e4a09951a Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Wed, 2 Sep 2026 22:20:45 +0900 Subject: [PATCH 52/57] vm: decide a float power by its operands' values, not their sign bits `float_pow` asked `is_sign_negative`, which is true of `-0.0`, so `0.0 ** -0.0` raised ZeroDivisionError where it is 1.0, and `(-0.0) ** 0.5` went to the complex branch where it is 0.0. Both tests now compare against zero, and whether the exponent is an integer is decided by whether it sits above its own floor - which answers an infinity and a nan without a case of its own, where subtracting and comparing against an epsilon answered a nan by accident and got exponents a hair above an integer wrong. `compile_fpow` is those guards translated one for one, so it moves with them. The sign-bit helper it needed has no other caller and goes with it. Assisted-by: Claude --- crates/jit/src/instructions.rs | 43 ++++++++++----------------------- crates/jit/tests/deopt_tests.rs | 28 ++++++++++----------- crates/jit/tests/float_tests.rs | 11 +++++++++ crates/vm/src/builtins/float.rs | 8 ++++-- 4 files changed, 43 insertions(+), 47 deletions(-) diff --git a/crates/jit/src/instructions.rs b/crates/jit/src/instructions.rs index 422c8461674..dabf81c4cf8 100644 --- a/crates/jit/src/instructions.rs +++ b/crates/jit/src/instructions.rs @@ -1357,24 +1357,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { Ok((quotient, remainder)) } - /// The sign bit of an f64, matching `f64::is_sign_negative` - true for - /// `-0.0` and a negative NaN, not just an ordinary negative number. - /// `fcmp` compares values and cannot see this bit; reinterpreting the - /// bits as a signed integer and testing that sign is what - /// `f64::is_sign_negative` itself does. - fn is_sign_negative(&mut self, v: Value) -> Value { - let bits = self.builder.ins().bitcast(types::I64, MemFlags::new(), v); - self.builder.ins().icmp_imm(IntCC::SignedLessThan, bits, 0) - } - /// 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, - /// including the parts that look wrong - the first guard below reads the - /// SIGN BIT of the exponent, not its value, because that is what - /// `is_sign_negative` does, so `0.0 ** -0.0` deopts exactly as - /// `0.0 ** -1.0` does. + /// complex result. This is `float_pow` translated guard for guard. fn compile_fpow( &mut self, a: Value, @@ -1383,25 +1369,22 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> { ) -> Result { let zero_f = self.builder.ins().f64const(0.0); - // v1.is_zero() && v2.is_sign_negative() -> ZeroDivisionError. + // 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_sign_negative = self.is_sign_negative(b); - let divides_by_zero = self.builder.ins().band(base_zero, exp_sign_negative); + 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.is_sign_negative() && (v2.floor() - v2).abs() > f64::EPSILON -> - // complex result - a `-0.0` base is caught by its sign bit here too, - // the same as the exponent above. - let base_sign_negative = self.is_sign_negative(a); + // 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 diff = self.builder.ins().fsub(b_floor, b); - let abs_diff = self.builder.ins().fabs(diff); - let epsilon = self.builder.ins().f64const(f64::EPSILON); - let fractional = self - .builder - .ins() - .fcmp(FloatCC::GreaterThan, abs_diff, epsilon); - let complex = self.builder.ins().band(base_sign_negative, fractional); + 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 diff --git a/crates/jit/tests/deopt_tests.rs b/crates/jit/tests/deopt_tests.rs index 79c3bbffc9b..2cdcacb1dd2 100644 --- a/crates/jit/tests/deopt_tests.rs +++ b/crates/jit/tests/deopt_tests.rs @@ -378,41 +378,39 @@ def div(a: float, b: int) -> float: } } - /// `0.0 ** negative` raises rather than returning an infinity. This - /// reads the sign bit of the exponent, not its value, so `-0.0` deopts - /// it exactly as `-1.0` does. + /// `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 "# }; - for exponent in [-1.0f64, -0.0f64] { - match code.invoke(&[0.0f64.into(), exponent.into()]) { + // 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(0.0), float(exponent)]); + assert_eq!(state.stack, vec![float(base), float(-1.0)]); } - other => panic!("expected a deopt for 0.0 ** {exponent}, got {other:?}"), + 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 sign bit too, so a `-0.0` base is caught the same way a - /// `-8.0` one 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 "# }; - for base in [-8.0f64, -0.0f64] { - match code.invoke(&[base.into(), 0.5f64.into()]) { - Ok(Outcome::Deopt(state)) => { - assert_eq!(state.stack, vec![float(base), float(0.5)]); - } - other => panic!("expected a deopt for {base} ** 0.5, got {other:?}"), + 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:?}"), } } diff --git a/crates/jit/tests/float_tests.rs b/crates/jit/tests/float_tests.rs index 895865e51d3..9b688ad0809 100644 --- a/crates/jit/tests/float_tests.rs +++ b/crates/jit/tests/float_tests.rs @@ -173,6 +173,17 @@ mod tests { // `+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_bits_eq!(pow(2.0, 2.0), Ok(4.0f64)); assert_bits_eq!(pow(3.0, 3.0), Ok(27.0f64)); 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)) From 6fa0b7f42febbd5a417addb8fe3a91c45163f68c Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Wed, 2 Sep 2026 22:20:59 +0900 Subject: [PATCH 53/57] common: keep the sign of a zero quotient A rational carries no signed zero, so `true_div` rounded `0 / -1` to `0.0` where it is `-0.0`. Only an exactly zero numerator lost the sign this way: a quotient too small to represent already rounded to the right signed zero on its own. The compiled path, which divides in floating point, had it right and disagreed with the interpreter. Assisted-by: Claude --- crates/common/src/int.rs | 7 +++++++ 1 file changed, 7 insertions(+) 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 From c31469b7baaa9e1a6468701c5e4d6c98fdac94cf Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Wed, 2 Sep 2026 22:21:07 +0900 Subject: [PATCH 54/57] docs: describe what the automatic compiler takes and turns down The README's JIT section covered only `__jit__()`. It now also covers the `aot` feature, the switches that turn it on, the shapes it compiles and the shapes it refuses, what happens where a machine word runs out, and why compiled code is invisible to a tracer. Assisted-by: Claude --- README.md | 51 ++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 46 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 1c8c972edce..dcdbc8540dd 100644 --- a/README.md +++ b/README.md @@ -145,21 +145,62 @@ 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 tries every function once, as it is first called, 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. Every function is then +offered to the compiler the first time it is called. The attempt happens once, +so a function it turns down costs that one attempt and is interpreted from then on. + +`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 + +Annotated scalar functions, and nothing else. The argument and return types come +from the annotations rather than from the values a call arrives with, so a function +without them is turned down however it is called. + +Taken: `int`, `float` and `bool` arguments, locals and return values; arithmetic, +comparison and boolean operators; `if`, `while`, and the assignments between them. + +Turned down: missing annotations, `*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(): From 09b35839e791cb530ce85dbd9099d136b54ef6e7 Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Thu, 3 Sep 2026 01:45:41 +0900 Subject: [PATCH 55/57] vm: end the finalizing span with the shutdown that opened it The bit is in a word every interpreter in the process shares, and nothing cleared it. After one interpreter finalized - a subinterpreter closed with `_interpreters.destroy`, say - every interpreter that outlived it read a non-zero word on every instruction and took the eval-breaker slow path for the rest of the process: 78 ms against 212 ms for the same integer loop. Count the spans the way the stop-the-world ones are counted, for the same reason: a subinterpreter is finalized while the interpreter that owns it is finalizing, so the bit belongs to the last span to close, not the first. Assisted-by: Claude --- crates/vm/src/signal.rs | 30 +++++++++++++++++++++++++----- crates/vm/src/vm/interpreter.rs | 9 +++++++-- 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/crates/vm/src/signal.rs b/crates/vm/src/signal.rs index 33678bfbaf5..a46316a0f57 100644 --- a/crates/vm/src/signal.rs +++ b/crates/vm/src/signal.rs @@ -31,8 +31,8 @@ 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; -/// The interpreter is shutting down, so every thread but the main one must -/// stop running bytecode. Set once and never cleared. +/// 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; @@ -42,6 +42,14 @@ const FINALIZING_BIT: u8 = 1 << 4; #[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, reason = "workaround for const array repeat limitation (rust issue #79270)" @@ -184,10 +192,22 @@ pub(crate) fn reset_stop_requests() { EVAL_BREAKER.fetch_and(!STOP_BIT, Ordering::Release); } -/// Record that the interpreter has started shutting down. +/// Open a shutdown span, so the word says a thread may have to stop. #[cfg(feature = "threading")] -pub(crate) fn set_finalizing_bit() { - EVAL_BREAKER.fetch_or(FINALIZING_BIT, Ordering::Release); +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. diff --git a/crates/vm/src/vm/interpreter.rs b/crates/vm/src/vm/interpreter.rs index acdcead72b2..5f38a72290d 100644 --- a/crates/vm/src/vm/interpreter.rs +++ b/crates/vm/src/vm/interpreter.rs @@ -699,9 +699,14 @@ impl Interpreter { 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. + // 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::set_finalizing_bit(); + 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); From cb70c1dbafe1e6a1b01c7a4a132dc78dc340eeaa Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Thu, 3 Sep 2026 01:45:58 +0900 Subject: [PATCH 56/57] ci: build the benchmarks without automatic compilation again The perf gate compares a branch against main, and main builds without the feature, so a branch that turns it on for the benchmark job compares two configurations rather than two revisions. Merging it would leave every later comparison measuring a build nobody ships. Reverts the benchmark half of 60ee4c9d7; the aot path is still run by the CPython-suite step and both snippets. Assisted-by: Claude --- .github/workflows/codspeed.yaml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.github/workflows/codspeed.yaml b/.github/workflows/codspeed.yaml index 2b83ba8e6ad..a9bc341408e 100644 --- a/.github/workflows/codspeed.yaml +++ b/.github/workflows/codspeed.yaml @@ -62,10 +62,8 @@ jobs: - name: Install cargo-codspeed run: cargo install cargo-codspeed --locked --version ^5 - # Benchmarks run with automatic compilation on, unlike the default - # build, so its cost on the first call of every function is measured. - name: Build the benchmark targets - run: cargo codspeed build --locked --measurement-mode simulation --features rustpython/aot -p rustpython -p rustpython-sre_engine + run: cargo codspeed build --locked --measurement-mode simulation -p rustpython -p rustpython-sre_engine - name: Save cache if: ${{ github.ref == 'refs/heads/main' }} # only save on main From 831b29a1e2acb8640f14c668e4f4c30e9fc0de2e Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Thu, 3 Sep 2026 11:09:54 +0900 Subject: [PATCH 57/57] vm: compile a function once it is warm, on the types its calls pass The automatic path read the argument types off the annotations, which is why it compiled almost nothing: seven functions in the whole of `Lib/` annotate every parameter and the return `int`, `float` or `bool`, and the compiler turns all seven down for other reasons. A call carries the types the function is actually being used with, and the compiled code already type-checks every argument it is handed, so a guess that turns out wrong costs a fall back to the interpreter rather than a wrong answer. Compiling on the first call also spent the eligibility scan on every function a program calls once, which is most of them. Count the calls instead and offer the function to the compiler when the count says a compile could be repaid; the types come from the call that crosses that line. Both paths resolve a call's arguments onto its parameters the same way and have to agree on which calls they accept, so they now share the walk. The sink is a trait rather than a returned vector to keep the per-call path allocating nothing. Nothing in the automatic path runs Python code any more, so a speculative compile can no longer pull an `__annotate__` forward. `__jit__()` still reads annotations. Measured with `sys._jit._stats()`: three unannotated numeric functions that could not be compiled at all before now are, and run 1.7x, 2.2x and 27x faster. Importing 25 stdlib modules asks the compiler about 87 functions where it used to ask about 1882, and the same import-heavy startup now costs the same with the feature on as with it off, against 6.8% more before. Assisted-by: Claude --- README.md | 28 +++-- crates/vm/src/builtins/function.rs | 9 +- crates/vm/src/builtins/function/aot.rs | 88 +++++++-------- crates/vm/src/builtins/function/jit.rs | 150 +++++++++++++++++++++---- extra_tests/snippets/aot.py | 65 +++++++---- 5 files changed, 245 insertions(+), 95 deletions(-) diff --git a/README.md b/README.md index dcdbc8540dd..7af45d8991b 100644 --- a/README.md +++ b/README.md @@ -146,8 +146,9 @@ cargo build --release --target wasm32-wasip1 --features="freeze-stdlib" ### JIT (Just in time) compiler RustPython has a **very** experimental JIT compiler that compiles python functions into native code. -It comes in two forms: an automatic one that tries every function once, as it is first called, and -an explicit `__jit__()` that compiles the one function it is called on. +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 @@ -163,9 +164,11 @@ This requires autoconf, automake, libtool, and clang to be installed. #### Using the automatic compiler 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. Every function is then -offered to the compiler the first time it is called. The attempt happens once, -so a function it turns down costs that one attempt and is interpreted from then on. +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)` @@ -173,15 +176,20 @@ for the functions it has looked at so far — a RustPython extension. #### What it compiles -Annotated scalar functions, and nothing else. The argument and return types come -from the annotations rather than from the values a call arrives with, so a function -without them is turned down however it is called. +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: missing annotations, `*args`/`**kwargs`, closures, generators and -coroutines, `try`/`except`, attributes and methods, containers, `for`, calls to +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 diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index 423c842d8a5..50a680ecb5f 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -87,6 +87,11 @@ pub struct PyFunction { /// 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); @@ -228,6 +233,8 @@ impl PyFunction { jitted_code: PyMutex::new(None), #[cfg(feature = "jit")] jit_state: AtomicU8::new(aot::UNTRIED), + #[cfg(feature = "jit")] + jit_warmup: AtomicU32::new(0), }; Ok(func) } @@ -667,7 +674,7 @@ impl Py { let mut state = self.jit_state.load(Relaxed); if state == aot::UNTRIED && vm.state.config.settings.aot { - state = aot::compile_on_first_call(self, vm)?; + state = aot::observe_call(self, &func_args, vm); } if matches!(state, aot::COMPILED_AUTO | aot::COMPILED_MANUAL) { diff --git a/crates/vm/src/builtins/function/aot.rs b/crates/vm/src/builtins/function/aot.rs index 9db83adf94b..cdc77101126 100644 --- a/crates/vm/src/builtins/function/aot.rs +++ b/crates/vm/src/builtins/function/aot.rs @@ -1,19 +1,23 @@ //! Compiling functions to native code without being asked. //! -//! `__jit__()` compiles one function on request and reports why it could not. -//! The AOT path instead tries every function on its first call, 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. +//! `__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::{AsObject, Py, PyResult, VirtualMachine, builtins::PyCode}; +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. +/// 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, @@ -30,6 +34,15 @@ 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; @@ -53,55 +66,42 @@ fn code_is_eligible(code: &Py) -> bool { } } -/// Discard the reason a speculative compile failed, unless it is something -/// the program has to see. -/// -/// Reading annotations runs `__annotate__`, which is Python code: a forward -/// reference raises NameError and is none of our business, but a -/// KeyboardInterrupt or SystemExit that happens to land there belongs to the -/// program, not to a compile attempt nobody asked for. -fn ignore_speculative_error(result: PyResult, vm: &VirtualMachine) -> PyResult> { - match result { - Ok(value) => Ok(Some(value)), - Err(err) if err.fast_isinstance(vm.ctx.exceptions.exception_type) => Ok(None), - Err(err) => Err(err), - } -} - -fn try_compile(func: &Py, vm: &VirtualMachine) -> PyResult> { +fn try_compile( + func: &Py, + func_args: &FuncArgs, + vm: &VirtualMachine, +) -> Option { let code: &Py = &func.code; if !code_is_eligible(code) { - return Ok(None); + return None; } - let Some(arg_types) = ignore_speculative_error(super::jit::get_jit_arg_types(func, vm), vm)? - else { - return Ok(None); - }; - let Some(ret_type) = ignore_speculative_error(super::jit::jit_ret_type(func, vm), vm)? else { - return Ok(None); - }; - - Ok(vm - .state + // 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, ret_type, Safety::Strict) - .ok()) + .compile(&code.code, &arg_types, None, Safety::Strict) + .ok() } -/// Give `func` its one automatic compile attempt and return its new state. -pub(super) fn compile_on_first_call(func: &Py, vm: &VirtualMachine) -> PyResult { - // Claim the function before running anything that can re-enter it: - // evaluating `__annotate__` calls Python, which can reach this same - // function, and a reentrant attempt has to interpret rather than recurse. +/// 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, vm)? else { + let Some(compiled) = try_compile(func, func_args, vm) else { vm.state.aot_stats.rejected.fetch_add(1, Relaxed); - return Ok(REJECTED); + return REJECTED; }; *func.jitted_code.lock() = Some(compiled); func.jit_state.store(COMPILED_AUTO, Relaxed); vm.state.aot_stats.compiled.fetch_add(1, Relaxed); - Ok(COMPILED_AUTO) + COMPILED_AUTO } diff --git a/crates/vm/src/builtins/function/jit.rs b/crates/vm/src/builtins/function/jit.rs index fe3bd91799e..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 { @@ -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,16 +249,16 @@ 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); } @@ -212,8 +273,8 @@ pub(crate) fn get_jit_args<'a>( 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)?; } } } @@ -222,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/extra_tests/snippets/aot.py b/extra_tests/snippets/aot.py index 915d64e4e4b..2d9d43b612c 100644 --- a/extra_tests/snippets/aot.py +++ b/extra_tests/snippets/aot.py @@ -6,6 +6,20 @@ # 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 @@ -44,9 +58,10 @@ def guarded(a: int) -> int: return 0 +warm(scale, 2.0, 3.0) assert scale(2.0, 3.0) == 5.0 -# An integer where a float was declared does not fit the compiled signature. +# 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 @@ -54,6 +69,7 @@ def guarded(a: int) -> int: # 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 @@ -66,15 +82,16 @@ 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. +# Shapes with no compiled form at all still behave, warm or cold. assert untyped("a", "b") == "ab" assert untyped(1, 2) == 3 -assert variadic(1, 2, 3) == 6 -assert closure_factory(5)(1) == 6 +assert warm(variadic, 1, 2, 3)(1, 2, 3) == 6 +assert warm(closure_factory(5), 1)(1) == 6 assert list(generator(7)) == [7] -assert guarded(3) == 3 +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 @@ -84,6 +101,7 @@ 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) @@ -101,6 +119,7 @@ def countdown(a: float) -> float: return a +warm(countdown, 3.0) assert countdown(3.0) == 0.0 original_countdown = countdown @@ -109,6 +128,7 @@ def countdown(a: float) -> float: return -1.0 +warm(countdown, 3.0) assert original_countdown(3.0) == -1.0 @@ -122,19 +142,22 @@ def fib_iter(n: int) -> int: 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 -# Automatic compilation reads annotations, which under PEP 649 means running -# `__annotate__`. A name that is not defined yet raises there, and that is -# none of the program's business: it never asked for its annotations. +# 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__ @@ -145,7 +168,7 @@ def forward(a: NotDefinedYet) -> int: if AOT: - # ... but an interrupt that lands in `__annotate__` belongs to the program. + # The same thing where evaluating the annotation would be unmissable. class Boom(BaseException): pass @@ -155,12 +178,14 @@ def explode(): def annotated(a: explode()) -> int: return 1 + warm(annotated, 1) + assert annotated(1) == 1 try: - annotated(1) + annotated.__annotations__ except Boom: pass else: - raise AssertionError("expected a BaseException to reach the caller") + raise AssertionError("expected the annotation to still be unevaluated") # A frame that outlives its call has to keep resolving `f_back` past its @@ -204,6 +229,10 @@ def spin(n: int) -> int: 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() @@ -223,14 +252,12 @@ def keep_spinning(): if AOT: compiled, rejected, deoptimized = sys._jit._stats() - # `scale`, `wide`, `wide2`, `divide`, `fib_iter`, `spin`, and the rebound - # `countdown` are what the automatic path takes above - the original - # self-recursive one, kept as `original_countdown`, is refused. These - # are floors, not exact counts, but they must not regress: the point of - # letting Strict compile arithmetic was to take shapes it used to - # refuse outright, and a floor already met before that change could not - # tell if the gate came back. `compiled` and `deoptimized` are pinned to - # what this file currently produces (`compiled 7 ... deopt 5`). + # `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.