perf(context): pass actor sp through registers, not TLS

switch_to_actor takes the target sp in rdi and returns the actor's
next saved sp in rax (handed over by switch_to_scheduler's shim).
Deletes the ACTOR_SP thread-local and halves the helper calls per
one-way switch (2 -> 1); the scheduler loop also drops its
set_actor_sp/get_actor_sp TLS round-trips. SCHEDULER_SP stays: a
yielding actor at arbitrary call depth has no argument channel back.

asm before/after in outputs/history.md session 1-2. Tests: 354 pass.
This commit is contained in:
claude-asm-audit
2026-08-21 12:22:45 +00:00
parent 741c10337b
commit d7082eb266
3 changed files with 54 additions and 58 deletions
+42 -32
View File
@@ -4,16 +4,20 @@
//! actor running on its own mmap'd stack. The compiler cannot do this; the //! actor running on its own mmap'd stack. The compiler cannot do this; the
//! whole point of `#[unsafe(naked)]` is that we control every instruction. //! whole point of `#[unsafe(naked)]` is that we control every instruction.
//! //!
//! `SCHEDULER_SP` and `ACTOR_SP` are thread-locals holding each side's saved //! The actor's stack pointer travels in registers: `switch_to_actor` takes
//! stack pointer. `init_actor_stack` builds the initial stack so that the //! the target sp as its argument and returns the sp the actor saved when it
//! first `switch_to_actor` lands inside the entry function with `rsp % 16 == 8` //! next yielded (handed back in `rax` by `switch_to_scheduler`'s shim). Only
//! (the x86-64 ABI requirement at function entry). //! the *scheduler* sp lives in a thread-local — the yielding actor sits at
//! arbitrary call depth with no argument channel back to the scheduler, so
//! TLS is the one place it can find the way home. `init_actor_stack` builds
//! the initial stack so that the first `switch_to_actor` lands inside the
//! entry function with `rsp % 16 == 8` (the x86-64 ABI requirement at
//! function entry).
use std::cell::Cell; use std::cell::Cell;
thread_local! { thread_local! {
static SCHEDULER_SP: Cell<usize> = const { Cell::new(0) }; static SCHEDULER_SP: Cell<usize> = const { Cell::new(0) };
static ACTOR_SP: Cell<usize> = const { Cell::new(0) };
} }
fn get_scheduler_sp() -> usize { fn get_scheduler_sp() -> usize {
@@ -22,12 +26,6 @@ fn get_scheduler_sp() -> usize {
fn set_scheduler_sp(v: usize) { fn set_scheduler_sp(v: usize) {
SCHEDULER_SP.with(|c| c.set(v)) SCHEDULER_SP.with(|c| c.set(v))
} }
pub fn get_actor_sp() -> usize {
ACTOR_SP.with(|c| c.get())
}
pub fn set_actor_sp(v: usize) {
ACTOR_SP.with(|c| c.set(v))
}
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
// Initial stack layout // Initial stack layout
@@ -78,12 +76,24 @@ pub fn init_actor_stack(top: *mut u8, entry: extern "C-unwind" fn()) -> usize {
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
// Context switch shims // Context switch shims
// //
// Each shim: // switch_to_actor_asm (rdi = target actor sp, returns rax = the sp the actor
// 1. Pushes the six callee-saved integer registers. // saved when it next yielded):
// 2. Snaps rsp into rdi and calls the Rust helper that stores it. // 1. Pushes the six callee-saved integer registers (scheduler side).
// 3. Calls the Rust helper that returns the *other* side's saved rsp. // 2. Stashes the target sp in rbx — free scratch: the register's live
// 4. Moves that into rsp. // value is on the stack we just pushed to, and the pops below load the
// 5. Pops the six registers and rets. // *other* side's values anyway — then snaps rsp into rdi and calls the
// Rust helper that stores it in SCHEDULER_SP.
// 3. Installs the target sp and pops the actor's registers; `ret` lands
// where the actor yielded (or in `entry` on first resume).
//
// switch_to_scheduler_asm (no args; its "return value" materialises on the
// OTHER stack, as switch_to_actor's rax):
// 1. Pushes the six callee-saved integer registers (actor side).
// 2. Stashes its own rsp in rbx (same free-scratch argument), asks the
// Rust helper for SCHEDULER_SP.
// 3. Installs the scheduler sp, moves the saved actor sp into rax, pops
// the scheduler's registers and rets — completing the scheduler's
// `switch_to_actor(sp)` call with the actor's new sp as its result.
// //
// XMM registers are NOT saved here. We rely on every yield happening through // XMM registers are NOT saved here. We rely on every yield happening through
// a Rust call site, which means the compiler has spilled any live XMM state // a Rust call site, which means the compiler has spilled any live XMM state
@@ -94,31 +104,32 @@ pub fn init_actor_stack(top: *mut u8, entry: extern "C-unwind" fn()) -> usize {
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
#[unsafe(naked)] #[unsafe(naked)]
unsafe extern "C" fn switch_to_actor_asm() { unsafe extern "C" fn switch_to_actor_asm(actor_sp: usize) -> usize {
core::arch::naked_asm!( core::arch::naked_asm!(
"push rbx", "push rbp", "push r12", "push r13", "push r14", "push r15", "push rbx", "push rbp", "push r12", "push r13", "push r14", "push r15",
"mov rbx, rdi",
"mov rdi, rsp", "mov rdi, rsp",
"call {set_sched_sp}", "call {set_sched_sp}",
"call {get_actor_sp}", "mov rsp, rbx",
"mov rsp, rax",
"pop r15", "pop r14", "pop r13", "pop r12", "pop rbp", "pop rbx", "pop r15", "pop r14", "pop r13", "pop r12", "pop rbp", "pop rbx",
"ret", "ret",
set_sched_sp = sym set_scheduler_sp, set_sched_sp = sym set_scheduler_sp,
get_actor_sp = sym get_actor_sp,
); );
} }
/// Resume the actor whose sp is in `ACTOR_SP`. Returns when the actor yields. /// Resume the actor whose saved stack pointer is `actor_sp`. Returns when the
/// actor yields, with the stack pointer the actor saved as it did — store it
/// back into the slot for the next resume.
/// ///
/// # Safety /// # Safety
/// ///
/// The caller must be running on a scheduler thread with a valid actor stack /// `actor_sp` must be a valid saved actor stack pointer — either from
/// pointer installed in `ACTOR_SP` — either by `init_actor_stack` (first /// `init_actor_stack` (first resume) or the value a prior `switch_to_actor`
/// resume) or by a prior `switch_to_scheduler` (subsequent resumes). Resuming /// returned for this actor (subsequent resumes). Resuming with a stale or
/// with an unset or stale `ACTOR_SP` transfers control to an arbitrary address. /// forged sp transfers control to an arbitrary address. Must not be called
/// Must not be called from within an actor (only the scheduler side may resume). /// from within an actor (only the scheduler side may resume).
pub unsafe fn switch_to_actor() { pub unsafe fn switch_to_actor(actor_sp: usize) -> usize {
unsafe { switch_to_actor_asm() }; unsafe { switch_to_actor_asm(actor_sp) }
} }
/// Yield from the running actor back to its scheduler thread. Returns when the /// Yield from the running actor back to its scheduler thread. Returns when the
@@ -135,13 +146,12 @@ pub unsafe fn switch_to_actor() {
pub unsafe extern "C" fn switch_to_scheduler() { pub unsafe extern "C" fn switch_to_scheduler() {
core::arch::naked_asm!( core::arch::naked_asm!(
"push rbx", "push rbp", "push r12", "push r13", "push r14", "push r15", "push rbx", "push rbp", "push r12", "push r13", "push r14", "push r15",
"mov rdi, rsp", "mov rbx, rsp",
"call {set_actor_sp}",
"call {get_sched_sp}", "call {get_sched_sp}",
"mov rsp, rax", "mov rsp, rax",
"mov rax, rbx",
"pop r15", "pop r14", "pop r13", "pop r12", "pop rbp", "pop rbx", "pop r15", "pop r14", "pop r13", "pop r12", "pop rbp", "pop rbx",
"ret", "ret",
set_actor_sp = sym set_actor_sp,
get_sched_sp = sym get_scheduler_sp, get_sched_sp = sym get_scheduler_sp,
); );
} }
+2 -4
View File
@@ -116,7 +116,7 @@ use crate::actor::{
take_last_outcome, Actor, Outcome, take_last_outcome, Actor, Outcome,
}; };
use crate::channel::Sender; use crate::channel::Sender;
use crate::context::{get_actor_sp, set_actor_sp, switch_to_actor}; use crate::context::switch_to_actor;
use crate::io::IoThread; use crate::io::IoThread;
use crate::monitor::{Down, DownReason, MonitorId}; use crate::monitor::{Down, DownReason, MonitorId};
use crate::pid::Pid; use crate::pid::Pid;
@@ -2228,7 +2228,6 @@ fn schedule_loop(inner: &Arc<RuntimeInner>, slot_idx: usize) {
.current_pid_index .current_pid_index
.store(pid.index(), Ordering::Relaxed); .store(pid.index(), Ordering::Relaxed);
set_actor_sp(sp);
set_current_pid(pid); set_current_pid(pid);
crate::preempt::set_current_stop(stop_flag); crate::preempt::set_current_stop(stop_flag);
crate::preempt::set_current_slot(slot as *const Slot); crate::preempt::set_current_slot(slot as *const Slot);
@@ -2254,7 +2253,7 @@ fn schedule_loop(inner: &Arc<RuntimeInner>, slot_idx: usize) {
crate::causal::on_resume(slot); crate::causal::on_resume(slot);
crate::te!(crate::trace::Event::Resume(pid)); crate::te!(crate::trace::Event::Resume(pid));
unsafe { switch_to_actor() }; let saved_sp = unsafe { switch_to_actor(sp) };
PREEMPTION_ENABLED.with(|c| c.set(false)); PREEMPTION_ENABLED.with(|c| c.set(false));
// RFC 016 Chunk 2: charge the cycles this resume consumed to the actor // RFC 016 Chunk 2: charge the cycles this resume consumed to the actor
@@ -2268,7 +2267,6 @@ fn schedule_loop(inner: &Arc<RuntimeInner>, slot_idx: usize) {
crate::preempt::clear_current_slot(); crate::preempt::clear_current_slot();
let intent = YIELD_INTENT.with(|c| c.get()); let intent = YIELD_INTENT.with(|c| c.get());
let saved_sp = get_actor_sp();
slot.sp.store(saved_sp, Ordering::Relaxed); slot.sp.store(saved_sp, Ordering::Relaxed);
// RFC 019 §2: sampled high-water — one branch + at most one store // RFC 019 §2: sampled high-water — one branch + at most one store
// into the line the store above just dirtied. Relaxed and advisory; // into the line the store above just dirtied. Relaxed and advisory;
+10 -22
View File
@@ -1,9 +1,7 @@
//! Low-level context-switch tests. These poke `init_actor_stack` and the //! Low-level context-switch tests. These poke `init_actor_stack` and the
//! naked asm shims directly — no scheduler involved. //! naked asm shims directly — no scheduler involved.
use smarm::context::{ use smarm::context::{init_actor_stack, switch_to_actor, switch_to_scheduler};
get_actor_sp, init_actor_stack, set_actor_sp, switch_to_actor, switch_to_scheduler,
};
use smarm::stack::Stack; use smarm::stack::Stack;
use std::cell::Cell; use std::cell::Cell;
@@ -31,8 +29,7 @@ fn actor_runs_and_returns_to_scheduler() {
reset_log(); reset_log();
let stack = Stack::new(64 * 1024, 4096).unwrap(); let stack = Stack::new(64 * 1024, 4096).unwrap();
let sp = init_actor_stack(stack.top(), actor_simple); let sp = init_actor_stack(stack.top(), actor_simple);
set_actor_sp(sp); let _ = unsafe { switch_to_actor(sp) };
unsafe { switch_to_actor() };
assert_eq!(get_log(), 0x1); assert_eq!(get_log(), 0x1);
} }
@@ -48,12 +45,11 @@ fn actor_yields_and_resumes() {
reset_log(); reset_log();
let stack = Stack::new(64 * 1024, 4096).unwrap(); let stack = Stack::new(64 * 1024, 4096).unwrap();
let sp = init_actor_stack(stack.top(), actor_two_steps); let sp = init_actor_stack(stack.top(), actor_two_steps);
set_actor_sp(sp);
unsafe { switch_to_actor() }; let sp = unsafe { switch_to_actor(sp) };
assert_eq!(get_log(), 0x1, "after first resume"); assert_eq!(get_log(), 0x1, "after first resume");
unsafe { switch_to_actor() }; let _ = unsafe { switch_to_actor(sp) };
assert_eq!(get_log(), 0x1 | 0x2, "after second resume"); assert_eq!(get_log(), 0x1 | 0x2, "after second resume");
} }
@@ -96,10 +92,9 @@ extern "C-unwind" fn actor_reg_check() {
fn callee_saved_registers_survive_yield() { fn callee_saved_registers_survive_yield() {
let stack = Stack::new(64 * 1024, 4096).unwrap(); let stack = Stack::new(64 * 1024, 4096).unwrap();
let sp = init_actor_stack(stack.top(), actor_reg_check); let sp = init_actor_stack(stack.top(), actor_reg_check);
set_actor_sp(sp);
unsafe { unsafe {
switch_to_actor(); let sp = switch_to_actor(sp);
switch_to_actor(); let _ = switch_to_actor(sp);
} }
assert_eq!( assert_eq!(
REG_BEFORE.get().copied().unwrap(), REG_BEFORE.get().copied().unwrap(),
@@ -138,18 +133,11 @@ fn two_actors_dont_corrupt_each_other() {
let sp_a = init_actor_stack(stack_a.top(), actor_a); let sp_a = init_actor_stack(stack_a.top(), actor_a);
let sp_b = init_actor_stack(stack_b.top(), actor_b); let sp_b = init_actor_stack(stack_b.top(), actor_b);
set_actor_sp(sp_a); let sp_a = unsafe { switch_to_actor(sp_a) };
unsafe { switch_to_actor() }; let sp_b = unsafe { switch_to_actor(sp_b) };
let sp_a = get_actor_sp();
set_actor_sp(sp_b); let _ = unsafe { switch_to_actor(sp_a) };
unsafe { switch_to_actor() }; let _ = unsafe { switch_to_actor(sp_b) };
let sp_b = get_actor_sp();
set_actor_sp(sp_a);
unsafe { switch_to_actor() };
set_actor_sp(sp_b);
unsafe { switch_to_actor() };
assert_eq!(A_VAL.with(|c| c.get()), 0xA00D); assert_eq!(A_VAL.with(|c| c.get()), 0xA00D);
assert_eq!(B_VAL.with(|c| c.get()), 0xB00D); assert_eq!(B_VAL.with(|c| c.get()), 0xB00D);