diff --git a/benches/tailcall_baseline.py b/benches/tailcall_baseline.py new file mode 100644 index 00000000000..cd60b166708 --- /dev/null +++ b/benches/tailcall_baseline.py @@ -0,0 +1,268 @@ +"""Repeatable baseline benchmark for the Python-to-Python call trampoline. + +Run this with a release RustPython binary, not CPython: + + ./target/release/rustpython benches/tailcall_baseline.py + +Each timed loop runs inside one RustPython process. Cases are rotated between +rounds so that every case is sampled at different points in the run. +""" + +import dis +import sys +import time + + +DEFAULT_SAMPLES = 14 +DEFAULT_ITERATIONS = 1_000_000 +DEFAULT_DEEP_ITERATIONS = 10_000 +DEFAULT_DEEP_DEPTH = 100 + + +def add_one(value): + return value + 1 + + +class Adder: + def add_one(self, value): + return value + 1 + + +bound_add_one = Adder().add_one + + +def shallow_inner(value): + return value + 1 + + +def shallow_outer(value): + return shallow_inner(value) + + +def recursive_add(depth, value): + if depth: + return recursive_add(depth - 1, value) + return value + 1 + + +def bench_inline(iterations, _depth): + value = 0 + start = time.perf_counter_ns() + for _index in range(iterations): + value = value + 1 + elapsed = time.perf_counter_ns() - start + assert value == iterations + return elapsed + + +def bench_exact_function(iterations, _depth): + function = add_one + value = 0 + start = time.perf_counter_ns() + for _index in range(iterations): + value = function(value) + elapsed = time.perf_counter_ns() - start + assert value == iterations + return elapsed + + +def bench_exact_bound_method(iterations, _depth): + method = bound_add_one + value = 0 + start = time.perf_counter_ns() + for _index in range(iterations): + value = method(value) + elapsed = time.perf_counter_ns() - start + assert value == iterations + return elapsed + + +def bench_shallow_nested(iterations, _depth): + function = shallow_outer + value = 0 + start = time.perf_counter_ns() + for _index in range(iterations): + value = function(value) + elapsed = time.perf_counter_ns() - start + assert value == iterations + return elapsed + + +def bench_deep_recursive(iterations, depth): + function = recursive_add + value = 0 + start = time.perf_counter_ns() + for _index in range(iterations): + value = function(depth, value) + elapsed = time.perf_counter_ns() - start + assert value == iterations + return elapsed + + +def shallow_activation_generator(iterations): + """Time calls which each activate a fresh trampoline in their callee.""" + function = shallow_outer + value = 0 + start = time.perf_counter_ns() + for _index in range(iterations): + value = function(value) + elapsed = time.perf_counter_ns() - start + assert value == iterations + yield elapsed + + +def bench_shallow_activations(iterations, _depth): + # Generator frames do not issue TailCall themselves. Each shallow_outer() + # invocation therefore starts and finishes a new trampoline when it calls + # shallow_inner(), including a fresh frame_stack allocation. + return next(shallow_activation_generator(iterations)) + + +def deep_activation_generator(iterations, depth): + """Time deep calls which each allocate and spill a fresh frame stack.""" + function = recursive_add + value = 0 + start = time.perf_counter_ns() + for _index in range(iterations): + value = function(depth, value) + elapsed = time.perf_counter_ns() - start + assert value == iterations + yield elapsed + + +def bench_deep_activations(iterations, depth): + return next(deep_activation_generator(iterations, depth)) + + +def parse_positive_int(name, default): + prefix = "--" + name + "=" + for argument in sys.argv[1:]: + if argument.startswith(prefix): + value = int(argument[len(prefix) :]) + if value <= 0: + raise ValueError(prefix + " must be positive") + return value + return default + + +def median(values): + ordered = sorted(values) + midpoint = len(ordered) // 2 + if len(ordered) % 2: + return ordered[midpoint] + return (ordered[midpoint - 1] + ordered[midpoint]) / 2 + + +def require_instruction(function, opname): + instructions = dis.get_instructions(function, adaptive=True) + if not any(instruction.opname == opname for instruction in instructions): + raise RuntimeError(function.__name__ + " did not specialize to " + opname) + + +def verify_specializations(): + expected = [ + (bench_exact_function, "CALL_PY_EXACT_ARGS"), + (bench_exact_bound_method, "CALL_BOUND_METHOD_EXACT_ARGS"), + (bench_shallow_nested, "CALL_PY_EXACT_ARGS"), + (shallow_outer, "CALL_PY_EXACT_ARGS"), + (bench_deep_recursive, "CALL_PY_EXACT_ARGS"), + (recursive_add, "CALL_PY_EXACT_ARGS"), + (shallow_activation_generator, "CALL_PY_EXACT_ARGS"), + (deep_activation_generator, "CALL_PY_EXACT_ARGS"), + ] + for function, opname in expected: + require_instruction(function, opname) + return ";".join(function.__name__ + ":" + opname for function, opname in expected) + + +def main(): + if sys.implementation.name != "rustpython": + raise RuntimeError("run this benchmark with a release RustPython binary") + + samples = parse_positive_int("samples", DEFAULT_SAMPLES) + iterations = parse_positive_int("iterations", DEFAULT_ITERATIONS) + deep_iterations = parse_positive_int( + "deep-iterations", DEFAULT_DEEP_ITERATIONS + ) + deep_depth = parse_positive_int("deep-depth", DEFAULT_DEEP_DEPTH) + + cases = [ + ("inline", bench_inline, iterations, 0, 0), + ("exact_function", bench_exact_function, iterations, 0, 1), + ("exact_bound_method", bench_exact_bound_method, iterations, 0, 1), + ("shallow_nested_steady", bench_shallow_nested, iterations, 0, 2), + ( + "deep_recursive_steady", + bench_deep_recursive, + deep_iterations, + deep_depth, + deep_depth + 1, + ), + ( + "shallow_nested_activation", + bench_shallow_activations, + iterations, + 0, + 2, + ), + ( + "deep_recursive_activation", + bench_deep_activations, + deep_iterations, + deep_depth, + deep_depth + 1, + ), + ] + results = {name: [] for name, _function, _iterations, _depth, _calls in cases} + + # Warm every bytecode path before collecting the interleaved samples. + for _name, function, _iterations, depth, _calls in cases: + function(100, depth) + specializations = verify_specializations() + + print("benchmark=tailcall_baseline_v2") + print("implementation=" + sys.implementation.name) + print("version=" + sys.version.replace("\n", " ")) + print("samples=" + str(samples)) + print("iterations=" + str(iterations)) + print("deep_iterations=" + str(deep_iterations)) + print("deep_depth=" + str(deep_depth)) + print("specializations=" + specializations) + print("round,case,iterations,total_ns,ns_per_iteration") + + for round_index in range(samples): + offset = round_index % len(cases) + interleaved = cases[offset:] + cases[:offset] + for name, function, case_iterations, depth, _calls in interleaved: + total_ns = function(case_iterations, depth) + ns_per_iteration = total_ns / case_iterations + results[name].append(ns_per_iteration) + print( + str(round_index + 1) + + "," + + name + + "," + + str(case_iterations) + + "," + + str(total_ns) + + "," + + ("%.3f" % ns_per_iteration) + ) + + inline_median = median(results["inline"]) + print("case,median_ns_per_iteration,delta_vs_inline_ns,python_calls_per_iteration") + for name, _function, _iterations, _depth, calls in cases: + case_median = median(results[name]) + print( + name + + "," + + ("%.3f" % case_median) + + "," + + ("%.3f" % (case_median - inline_median)) + + "," + + str(calls) + ) + + +if __name__ == "__main__": + main() diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index ebb158c8f71..2fe1db5bdcd 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -10680,11 +10680,10 @@ impl ExecutingFrame<'_> { } } - // Pop the callable and transfer ownership to the trampoline via - // the VM side channel, avoiding a per-frame mutex lock on - // temporary_refs. + // Pop the callable and transfer ownership to the trampoline. This one + // reference keeps every field borrowed by the callee frame alive. let callable = self.pop_value(); - unsafe { &mut *vm.pending_tailcall_refs.get() }.push(callable); + vm.set_pending_tailcall_owner(callable); vm.set_pending_tailcall(callee_iframe); } @@ -10734,13 +10733,13 @@ impl ExecutingFrame<'_> { *dst = Some(arg); } self.pop_value_opt(); // null (self_or_null) - let callable = self.pop_value(); // callable (bound method) + self.pop_value(); // callable (bound method) fastlocals[0] = Some(bound_self); - // Transfer ownership to the trampoline via the VM side channel. - let refs = unsafe { &mut *vm.pending_tailcall_refs.get() }; - refs.push(bound_function); - refs.push(callable); + // The function owns every field borrowed by the callee frame. + // bound_self is owned by fastlocals; the bound-method object itself is + // no longer needed and was dropped above, matching the recursive path. + vm.set_pending_tailcall_owner(bound_function); vm.set_pending_tailcall(callee_iframe); } diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index e0a086c10db..51840fbdbe7 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -110,11 +110,11 @@ pub struct VirtualMachine { /// pointer here before returning `ExecutionResult::TailCall`. /// Access only via `set_pending_tailcall` / `take_pending_tailcall`. pending_tailcall_frame: Cell>, - /// Owned references that keep callee raw pointers valid during TailCall. - /// Set by `tailcall_prepare_frame`, drained by the trampoline into - /// its local `owned_refs` Vec. Uses UnsafeCell because the VM is - /// per-thread and this field is only accessed on the owning thread. - pub(crate) pending_tailcall_refs: core::cell::UnsafeCell>, + /// Owned reference that keeps callee raw pointers valid during TailCall. + /// Set by the exact-call handlers and moved into the trampoline's + /// `SuspendedFrame`. Uses UnsafeCell because the VM is per-thread and this + /// field is only accessed on the owning thread. + pending_tailcall_owner: core::cell::UnsafeCell>, } /// Non-owning frame pointer for the non-unix threading frames stack. @@ -828,11 +828,11 @@ pub(crate) struct IframeEntryState { struct SuspendedFrame { iframe: *mut crate::frame::InterpreterFrame, entry_state: IframeEntryState, - /// Owned references that keep callee's raw pointers (code, globals, - /// builtins borrowed from PyFunction) valid. Drained from - /// `vm.pending_tailcall_refs` when the callee's TailCall is consumed. + /// Function that owns the callee's raw pointers (code, globals, builtins, + /// closure, and func_obj). Moved from `vm.pending_tailcall_owner` when the + /// callee's TailCall is consumed. /// Dropped when this SuspendedFrame is popped (after callee returns/errors). - owned_refs: Vec, + callee_owner: PyObjectRef, /// True for the initial frame passed into the trampoline by the caller. /// The caller owns the datastack allocation for the entry frame, so the /// trampoline must NOT release it — only callee-allocated frames are @@ -955,7 +955,7 @@ impl VirtualMachine { callable_cache: CallableCache::default(), audit_hooks: RefCell::new(vec![]), pending_tailcall_frame: Cell::new(None), - pending_tailcall_refs: core::cell::UnsafeCell::new(Vec::with_capacity(2)), + pending_tailcall_owner: core::cell::UnsafeCell::new(None), }; if vm.state.hash_secret.hash_str("") @@ -1406,6 +1406,22 @@ impl VirtualMachine { .set(Some(PendingFrame(core::ptr::NonNull::from(iframe)))); } + /// Store the function that owns the fields borrowed by the pending callee. + #[inline(always)] + pub(crate) fn set_pending_tailcall_owner(&self, owner: PyObjectRef) { + let slot = unsafe { &mut *self.pending_tailcall_owner.get() }; + debug_assert!(slot.is_none(), "pending TailCall owner was not consumed"); + *slot = Some(owner); + } + + /// Take the pending callee owner, resetting the side channel. + #[inline(always)] + fn take_pending_tailcall_owner(&self) -> PyObjectRef { + unsafe { &mut *self.pending_tailcall_owner.get() } + .take() + .expect("TailCall without pending owner") + } + /// Take the pending tailcall frame pointer, resetting the side channel. #[inline(always)] fn take_pending_tailcall(&self) -> *mut crate::frame::InterpreterFrame { @@ -1466,14 +1482,11 @@ impl VirtualMachine { } let initial_ptr = self.take_pending_tailcall(); - // Drain the refs that keep the initial callee's raw pointers alive. - let initial_refs = unsafe { &mut *self.pending_tailcall_refs.get() } - .drain(..) - .collect(); + let initial_owner = self.take_pending_tailcall_owner(); frame_stack.push(SuspendedFrame { iframe: iframe as *mut crate::frame::InterpreterFrame, entry_state, - owned_refs: initial_refs, + callee_owner: initial_owner, is_entry: true, }); let mut action = Action::EnterCallee(initial_ptr); @@ -1498,13 +1511,11 @@ impl VirtualMachine { let result = crate::frame::run_iframe(callee, self); match result { Ok(ExecutionResult::TailCall) => { - let refs = unsafe { &mut *self.pending_tailcall_refs.get() } - .drain(..) - .collect(); + let callee_owner = self.take_pending_tailcall_owner(); frame_stack.push(SuspendedFrame { iframe: callee_ptr, entry_state: callee_entry, - owned_refs: refs, + callee_owner, is_entry: false, }); action = Action::EnterCallee(self.take_pending_tailcall()); @@ -1539,7 +1550,7 @@ impl VirtualMachine { let SuspendedFrame { iframe: caller_iframe_ptr, entry_state: caller_entry, - owned_refs: _caller_refs, + callee_owner, is_entry: caller_is_entry, } = caller; let caller_iframe = unsafe { &mut *caller_iframe_ptr }; @@ -1548,20 +1559,18 @@ impl VirtualMachine { let result = crate::frame::run_iframe(caller_iframe, self); match result { Ok(ExecutionResult::TailCall) => { - let refs = unsafe { &mut *self.pending_tailcall_refs.get() } - .drain(..) - .collect(); - drop(_caller_refs); + let next_callee_owner = self.take_pending_tailcall_owner(); + drop(callee_owner); frame_stack.push(SuspendedFrame { iframe: caller_iframe_ptr, entry_state: caller_entry, - owned_refs: refs, + callee_owner: next_callee_owner, is_entry: caller_is_entry, }); action = Action::EnterCallee(self.take_pending_tailcall()); } Ok(ExecutionResult::Return(value)) => { - drop(_caller_refs); + drop(callee_owner); self.exit_iframe(caller_entry); if !caller_is_entry { unsafe { @@ -1574,7 +1583,7 @@ impl VirtualMachine { } Ok(ExecutionResult::Yield(_)) => panic!("Yield in non-generator frame"), Err(exc) => { - drop(_caller_refs); + drop(callee_owner); self.exit_iframe(caller_entry); if !caller_is_entry { unsafe { @@ -1595,7 +1604,7 @@ impl VirtualMachine { let SuspendedFrame { iframe: caller_iframe_ptr, entry_state: caller_entry, - owned_refs: _caller_refs, + callee_owner, is_entry: caller_is_entry, } = caller; let caller_iframe = unsafe { &mut *caller_iframe_ptr }; @@ -1609,20 +1618,18 @@ impl VirtualMachine { let result = crate::frame::run_iframe(caller_iframe, self); match result { Ok(ExecutionResult::TailCall) => { - let refs = unsafe { &mut *self.pending_tailcall_refs.get() } - .drain(..) - .collect(); - drop(_caller_refs); + let next_callee_owner = self.take_pending_tailcall_owner(); + drop(callee_owner); frame_stack.push(SuspendedFrame { iframe: caller_iframe_ptr, entry_state: caller_entry, - owned_refs: refs, + callee_owner: next_callee_owner, is_entry: caller_is_entry, }); action = Action::EnterCallee(self.take_pending_tailcall()); } Ok(ExecutionResult::Return(value)) => { - drop(_caller_refs); + drop(callee_owner); self.exit_iframe(caller_entry); if !caller_is_entry { unsafe { @@ -1639,7 +1646,7 @@ impl VirtualMachine { panic!("Yield in non-generator frame") } Err(new_exc) => { - drop(_caller_refs); + drop(callee_owner); self.exit_iframe(caller_entry); if !caller_is_entry { unsafe { @@ -1655,7 +1662,7 @@ impl VirtualMachine { } } Ok(Some(ExecutionResult::Return(value))) => { - drop(_caller_refs); + drop(callee_owner); self.exit_iframe(caller_entry); if !caller_is_entry { unsafe { @@ -1670,7 +1677,7 @@ impl VirtualMachine { panic!("Unexpected execution result in trampoline unwind") } Err(new_exc) => { - drop(_caller_refs); + drop(callee_owner); self.exit_iframe(caller_entry); if !caller_is_entry { unsafe { diff --git a/crates/vm/src/vm/thread.rs b/crates/vm/src/vm/thread.rs index 1bab539a0a6..7845e4edd2b 100644 --- a/crates/vm/src/vm/thread.rs +++ b/crates/vm/src/vm/thread.rs @@ -1023,7 +1023,7 @@ impl VirtualMachine { callable_cache: self.callable_cache.clone(), audit_hooks: RefCell::new(vec![]), pending_tailcall_frame: Cell::new(None), - pending_tailcall_refs: core::cell::UnsafeCell::new(Vec::with_capacity(2)), + pending_tailcall_owner: core::cell::UnsafeCell::new(None), }; ThreadedVirtualMachine { vm } }