From 89fb13e29aef535e2d7e655271b9b979cc31c742 Mon Sep 17 00:00:00 2001 From: smarm-dev Date: Fri, 29 May 2026 06:40:05 +0000 Subject: [PATCH] feat(arch): port context switch + cycle counter to aarch64 Extract the x86-64-specific context-switch shims, initial-stack layout, and cycle counter out of context.rs into a target_arch-gated src/arch/ module, and add an aarch64 (AAPCS64) backend. - src/arch/mod.rs: shared SCHEDULER_SP/ACTOR_SP thread-locals + the four-item backend contract (init_actor_stack, switch_to_{actor, scheduler}, read_cycle_counter); compile_error! on other arches. - src/arch/x86_64.rs: existing implementation verbatim, plus the rdtsc cycle counter relocated here from preempt.rs. - src/arch/aarch64.rs: x19-x30 stp/ldp context switch, lr-based return, CNTVCT_EL0 cycle counter (isb-serialised). - context.rs: now a thin facade re-exporting crate::arch; public surface unchanged for scheduler/preempt/tests. - preempt.rs: read_cycle_counter delegates to crate::arch; per-arch DEFAULT_TIMESLICE_CYCLES (TSC and CNTVCT have different units). - tests/context.rs: cfg-gate the x86 callee-saved-register probe and add the aarch64 x23-x26 sibling. - .cargo/config.toml: aarch64-unknown-linux-gnu cross linker. No functional change on x86-64: arch/x86_64.rs is byte-for-byte the old context.rs body, and tests/preempt.rs (incl. the XMM-not-saved guard from 14dc7a7) is untouched. --- .cargo/config.toml | 2 + src/arch/aarch64.rs | 165 ++++++++++++++++++++++++++++++++++++++++++++ src/arch/mod.rs | 54 +++++++++++++++ src/arch/x86_64.rs | 118 +++++++++++++++++++++++++++++++ src/context.rs | 122 +++----------------------------- src/lib.rs | 1 + src/preempt.rs | 20 ++++-- tests/context.rs | 34 +++++++++ 8 files changed, 399 insertions(+), 117 deletions(-) create mode 100644 .cargo/config.toml create mode 100644 src/arch/aarch64.rs create mode 100644 src/arch/mod.rs create mode 100644 src/arch/x86_64.rs diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..3c32d25 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,2 @@ +[target.aarch64-unknown-linux-gnu] +linker = "aarch64-linux-gnu-gcc" diff --git a/src/arch/aarch64.rs b/src/arch/aarch64.rs new file mode 100644 index 0000000..245c41b --- /dev/null +++ b/src/arch/aarch64.rs @@ -0,0 +1,165 @@ +//! Cooperative context switching and cycle counter, aarch64 (AAPCS64). +//! +//! The aarch64 mirror of the x86-64 backend. The protocol is identical — save +//! callee-saved state, swap stack pointers via the shared `*_SP` thread-locals, +//! restore, return — but the mechanics differ from x86 in two ways worth +//! stating up front, because they drive the stack layout: +//! +//! 1. **Return is via the link register, not the stack.** x86 `ret` pops the +//! return address off the stack; aarch64 `ret` jumps to whatever is in +//! `x30` (lr). So the entry point is parked in the saved-`x30` slot, and +//! the shim restores `x30` and `ret`s to it. There is no "ret target word" +//! sitting on the stack the way there is on x86. +//! +//! 2. **sp must be 16-byte aligned at all times**, not just at call entry. +//! We save 12 registers (x19–x28, x29, x30) = 96 bytes, already a multiple +//! of 16, so the frame is aligned by construction. +//! +//! FP/SIMD registers (v8–v15, whose low 64 bits are callee-saved under AAPCS64) +//! are NOT saved, for the same reason the x86 backend skips XMM: every yield +//! goes through a Rust call boundary, so the compiler has already spilled any +//! live vector state. If we ever yield from a non-call-boundary, this breaks. + +use super::{get_actor_sp, get_scheduler_sp, set_actor_sp, set_scheduler_sp}; + +// --------------------------------------------------------------------------- +// Initial stack layout +// +// Twelve 8-byte slots, all zero except the x30 slot which holds `entry`. The +// first `switch_to_actor` loads x19–x28/x29/x30 back and `ret`s to x30, +// landing at the top of `entry` with sp 16-aligned (the AAPCS64 requirement). +// +// Layout (high → low), relative to aligned_top = top & ~15: +// +// aligned_top - 8 : x30 = entry ← `ret` jumps here. +// aligned_top - 16 : x29 = 0 (fp) +// aligned_top - 24 : x28 = 0 +// aligned_top - 32 : x27 = 0 +// aligned_top - 40 : x26 = 0 +// aligned_top - 48 : x25 = 0 +// aligned_top - 56 : x24 = 0 +// aligned_top - 64 : x23 = 0 +// aligned_top - 72 : x22 = 0 +// aligned_top - 80 : x21 = 0 +// aligned_top - 88 : x20 = 0 +// aligned_top - 96 : x19 = 0 ← initial sp (16-aligned) +// +// The restore order in the shim (ldp pairs, low→high) must match this exactly. +// --------------------------------------------------------------------------- + +pub fn init_actor_stack(top: *mut u8, entry: extern "C-unwind" fn()) -> usize { + unsafe { + let aligned_top = top as usize & !15; + let mut sp = aligned_top; + sp -= 8; (sp as *mut usize).write(entry as usize); // x30 (lr) → entry + sp -= 8; (sp as *mut usize).write(0); // x29 (fp) + sp -= 8; (sp as *mut usize).write(0); // x28 + sp -= 8; (sp as *mut usize).write(0); // x27 + sp -= 8; (sp as *mut usize).write(0); // x26 + sp -= 8; (sp as *mut usize).write(0); // x25 + sp -= 8; (sp as *mut usize).write(0); // x24 + sp -= 8; (sp as *mut usize).write(0); // x23 + sp -= 8; (sp as *mut usize).write(0); // x22 + sp -= 8; (sp as *mut usize).write(0); // x21 + sp -= 8; (sp as *mut usize).write(0); // x20 + sp -= 8; (sp as *mut usize).write(0); // x19 ← initial sp + sp + } +} + +// --------------------------------------------------------------------------- +// Context switch shims +// +// Each shim: +// 1. Pushes x19–x28, x29, x30 as six 16-byte stp pairs (sp pre-decrement). +// 2. Moves sp into x0 and calls the Rust helper that stores it. +// 3. Calls the Rust helper that returns the *other* side's saved sp. +// 4. Moves that into sp. +// 5. Restores the twelve registers and rets (to the restored x30). +// +// The helper calls clobber x30 themselves, but that's fine: the outgoing x30 +// was already saved to the stack in step 1, and step 5 restores the incoming +// side's x30 before `ret`. The push/pop register order is paired so that the +// `ldp` sequence reverses the `stp` sequence exactly. +// --------------------------------------------------------------------------- + +#[unsafe(naked)] +unsafe extern "C" fn switch_to_actor_asm() { + core::arch::naked_asm!( + "stp x19, x20, [sp, #-96]!", + "stp x21, x22, [sp, #16]", + "stp x23, x24, [sp, #32]", + "stp x25, x26, [sp, #48]", + "stp x27, x28, [sp, #64]", + "stp x29, x30, [sp, #80]", + "mov x0, sp", + "bl {set_sched_sp}", + "bl {get_actor_sp}", + "mov sp, x0", + "ldp x21, x22, [sp, #16]", + "ldp x23, x24, [sp, #32]", + "ldp x25, x26, [sp, #48]", + "ldp x27, x28, [sp, #64]", + "ldp x29, x30, [sp, #80]", + "ldp x19, x20, [sp], #96", + "ret", + 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. +pub unsafe fn switch_to_actor() { + unsafe { switch_to_actor_asm() }; +} + +#[unsafe(naked)] +pub unsafe extern "C" fn switch_to_scheduler() { + core::arch::naked_asm!( + "stp x19, x20, [sp, #-96]!", + "stp x21, x22, [sp, #16]", + "stp x23, x24, [sp, #32]", + "stp x25, x26, [sp, #48]", + "stp x27, x28, [sp, #64]", + "stp x29, x30, [sp, #80]", + "mov x0, sp", + "bl {set_actor_sp}", + "bl {get_sched_sp}", + "mov sp, x0", + "ldp x21, x22, [sp, #16]", + "ldp x23, x24, [sp, #32]", + "ldp x25, x26, [sp, #48]", + "ldp x27, x28, [sp, #64]", + "ldp x29, x30, [sp, #80]", + "ldp x19, x20, [sp], #96", + "ret", + set_actor_sp = sym set_actor_sp, + get_sched_sp = sym get_scheduler_sp, + ); +} + +// --------------------------------------------------------------------------- +// Cycle counter +// +// CNTVCT_EL0 is the virtual count register, readable from EL0 on Linux. `isb` +// serialises so we don't sample the counter before prior instructions retire +// (the aarch64 analogue of x86's `lfence` before rdtsc). +// +// Unit: CNTVCT ticks, whose frequency is CNTFRQ_EL0 — typically 1–50 MHz, +// NOT the CPU clock. This is a different and much coarser unit than the x86 +// TSC, which is why `TIMESLICE_CYCLES` must be tuned separately for aarch64. +// --------------------------------------------------------------------------- + +#[inline(always)] +pub fn read_cycle_counter() -> u64 { + let cnt: u64; + unsafe { + core::arch::asm!( + "isb", + "mrs {cnt}, cntvct_el0", + cnt = out(reg) cnt, + options(nostack, nomem), + ); + } + cnt +} diff --git a/src/arch/mod.rs b/src/arch/mod.rs new file mode 100644 index 0000000..a611631 --- /dev/null +++ b/src/arch/mod.rs @@ -0,0 +1,54 @@ +//! Architecture abstraction layer. +//! +//! Everything that depends on the target ISA lives behind this module: the +//! cooperative context-switch shims, the initial-stack layout, and the +//! cycle counter used by the preemption timeslice. The rest of the runtime +//! talks only to the API re-exported here and never names a register. +//! +//! Each backend (`x86_64`, `aarch64`) exposes the same four items: +//! +//! - `init_actor_stack(top, entry) -> usize` +//! Build a fresh actor stack so the first `switch_to_actor` lands +//! inside `entry` with the ABI-correct stack alignment. +//! - `switch_to_actor()` +//! Resume the actor whose sp is in `ACTOR_SP`; returns when it yields. +//! - `switch_to_scheduler()` +//! Yield from the current actor back to its scheduler thread. +//! - `read_cycle_counter() -> u64` +//! Monotonic per-core cycle counter for the timeslice clock. The unit +//! is ISA-defined (x86 TSC ticks vs. aarch64 virtual-counter ticks), +//! so `TIMESLICE_CYCLES` must be tuned per-arch — see `preempt`. +//! +//! The `SCHEDULER_SP` / `ACTOR_SP` thread-locals are shared across backends +//! and live here, since the saved-stack-pointer protocol is identical +//! regardless of which registers the shim happens to push. + +use std::cell::Cell; + +thread_local! { + static SCHEDULER_SP: Cell = const { Cell::new(0) }; + static ACTOR_SP: Cell = const { Cell::new(0) }; +} + +// Used by the naked shims (via `sym`) and by the scheduler/tests through the +// re-exports below. `pub` because the asm references them as symbols. +pub(crate) fn get_scheduler_sp() -> usize { SCHEDULER_SP.with(|c| c.get()) } +pub(crate) fn set_scheduler_sp(v: usize) { 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)) } + +#[cfg(target_arch = "x86_64")] +mod x86_64; +#[cfg(target_arch = "x86_64")] +pub use x86_64::{init_actor_stack, read_cycle_counter, switch_to_actor, switch_to_scheduler}; + +#[cfg(target_arch = "aarch64")] +mod aarch64; +#[cfg(target_arch = "aarch64")] +pub use aarch64::{init_actor_stack, read_cycle_counter, switch_to_actor, switch_to_scheduler}; + +#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))] +compile_error!( + "smarm's context-switch and cycle-counter layer supports only x86_64 and \ + aarch64. Add a backend in src/arch/ for this target." +); diff --git a/src/arch/x86_64.rs b/src/arch/x86_64.rs new file mode 100644 index 0000000..bcd54e9 --- /dev/null +++ b/src/arch/x86_64.rs @@ -0,0 +1,118 @@ +//! Cooperative context switching and cycle counter, x86-64 (SysV AMD64). +//! +//! Two naked-asm functions move execution between a scheduler thread and an +//! 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. +//! +//! `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 super::{get_actor_sp, get_scheduler_sp, set_actor_sp, set_scheduler_sp}; + +// --------------------------------------------------------------------------- +// Initial stack layout +// +// We start from aligned_top = top & ~15, then bias by 8 and push seven 8-byte +// slots (downward). The first `switch_to_actor` pops r15..rbx and `ret`s — +// landing in `entry` with rsp % 16 == 8 (the x86-64 ABI state at the point +// just after a `call`, which is what `entry` is compiled to expect). +// +// Layout (high → low), relative to aligned_top = top & ~15. Note the entry +// slot is at aligned_top - 16, NOT aligned_top - 8: the function does +// `(top & ~15) - 8` and *then* a `-= 8` before the first write, so the first +// stored word lands at aligned_top - 16. Verified by single-stepping the +// `ret` under llmdbg: entry sits at an address with %16 == 0, so the post-ret +// rsp is %16 == 8. +// +// aligned_top - 8 : (unused padding; keeps entry's slot %16 == 0) +// aligned_top - 16 : entry ptr ← `ret` target. Post-ret: rsp % 16 == 8. +// aligned_top - 24 : rbx = 0 +// aligned_top - 32 : rbp = 0 +// aligned_top - 40 : r12 = 0 +// aligned_top - 48 : r13 = 0 +// aligned_top - 56 : r14 = 0 +// aligned_top - 64 : r15 = 0 ← initial rsp +// --------------------------------------------------------------------------- + +pub fn init_actor_stack(top: *mut u8, entry: extern "C-unwind" fn()) -> usize { + unsafe { + let mut sp = (top as usize & !15) - 8; + sp -= 8; (sp as *mut usize).write(entry as usize); // ret target + sp -= 8; (sp as *mut usize).write(0); // rbx + sp -= 8; (sp as *mut usize).write(0); // rbp + sp -= 8; (sp as *mut usize).write(0); // r12 + sp -= 8; (sp as *mut usize).write(0); // r13 + sp -= 8; (sp as *mut usize).write(0); // r14 + sp -= 8; (sp as *mut usize).write(0); // r15 + sp + } +} + +// --------------------------------------------------------------------------- +// Context switch shims +// +// Each shim: +// 1. Pushes the six callee-saved integer registers. +// 2. Snaps rsp into rdi and calls the Rust helper that stores it. +// 3. Calls the Rust helper that returns the *other* side's saved rsp. +// 4. Moves that into rsp. +// 5. Pops the six registers and rets. +// +// 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 +// to the stack before we get here. (This is the same argument the compiler +// uses internally — callee-saved regs are what survive a `call`, and the +// SysV AMD64 ABI says XMM0–15 are all caller-saved.) If we ever yield from +// a place that isn't a Rust call boundary, this assumption breaks. +// --------------------------------------------------------------------------- + +#[unsafe(naked)] +unsafe extern "C" fn switch_to_actor_asm() { + core::arch::naked_asm!( + "push rbx", "push rbp", "push r12", "push r13", "push r14", "push r15", + "mov rdi, rsp", + "call {set_sched_sp}", + "call {get_actor_sp}", + "mov rsp, rax", + "pop r15", "pop r14", "pop r13", "pop r12", "pop rbp", "pop rbx", + "ret", + 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. +pub unsafe fn switch_to_actor() { + unsafe { switch_to_actor_asm() }; +} + +#[unsafe(naked)] +pub unsafe extern "C" fn switch_to_scheduler() { + core::arch::naked_asm!( + "push rbx", "push rbp", "push r12", "push r13", "push r14", "push r15", + "mov rdi, rsp", + "call {set_actor_sp}", + "call {get_sched_sp}", + "mov rsp, rax", + "pop r15", "pop r14", "pop r13", "pop r12", "pop rbp", "pop rbx", + "ret", + set_actor_sp = sym set_actor_sp, + get_sched_sp = sym get_scheduler_sp, + ); +} + +// --------------------------------------------------------------------------- +// Cycle counter +// +// `lfence` serialises the instruction stream so we don't measure time before +// prior instructions retire. Unit: TSC ticks (≈ CPU base clock). +// --------------------------------------------------------------------------- + +#[inline(always)] +pub fn read_cycle_counter() -> u64 { + unsafe { + core::arch::asm!("lfence", options(nostack, nomem, preserves_flags)); + core::arch::x86_64::_rdtsc() + } +} diff --git a/src/context.rs b/src/context.rs index cf0ea86..589f25f 100644 --- a/src/context.rs +++ b/src/context.rs @@ -1,114 +1,14 @@ -//! Cooperative context switching, x86-64. +//! Cooperative context switching — public façade over the arch backend. //! -//! Two naked-asm functions move execution between a scheduler thread and an -//! 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. +//! The real implementation lives in `crate::arch`, selected per `target_arch`. +//! This module re-exports the stable surface (`init_actor_stack`, the two +//! `switch_to_*` shims, and the `*_actor_sp` accessors) so the scheduler, +//! `preempt`, and the `tests/context.rs` integration tests keep importing +//! `smarm::context::{...}` unchanged regardless of which ISA is built. //! -//! `SCHEDULER_SP` and `ACTOR_SP` are thread-locals holding each side's saved -//! stack pointer. `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). +//! See `crate::arch` for the contract every backend implements, and +//! `crate::arch::{x86_64, aarch64}` for the per-ISA register choreography. -use std::cell::Cell; - -thread_local! { - static SCHEDULER_SP: Cell = const { Cell::new(0) }; - static ACTOR_SP: Cell = const { Cell::new(0) }; -} - -fn get_scheduler_sp() -> usize { SCHEDULER_SP.with(|c| c.get()) } -fn set_scheduler_sp(v: usize) { 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 -// -// We start from aligned_top = top & ~15, then bias by 8 and push seven 8-byte -// slots (downward). The first `switch_to_actor` pops r15..rbx and `ret`s — -// landing in `entry` with rsp % 16 == 8 (the x86-64 ABI state at the point -// just after a `call`, which is what `entry` is compiled to expect). -// -// Layout (high → low), relative to aligned_top = top & ~15. Note the entry -// slot is at aligned_top - 16, NOT aligned_top - 8: the function does -// `(top & ~15) - 8` and *then* a `-= 8` before the first write, so the first -// stored word lands at aligned_top - 16. Verified by single-stepping the -// `ret` under llmdbg: entry sits at an address with %16 == 0, so the post-ret -// rsp is %16 == 8. -// -// aligned_top - 8 : (unused padding; keeps entry's slot %16 == 0) -// aligned_top - 16 : entry ptr ← `ret` target. Post-ret: rsp % 16 == 8. -// aligned_top - 24 : rbx = 0 -// aligned_top - 32 : rbp = 0 -// aligned_top - 40 : r12 = 0 -// aligned_top - 48 : r13 = 0 -// aligned_top - 56 : r14 = 0 -// aligned_top - 64 : r15 = 0 ← initial rsp -// --------------------------------------------------------------------------- - -pub fn init_actor_stack(top: *mut u8, entry: extern "C-unwind" fn()) -> usize { - unsafe { - let mut sp = (top as usize & !15) - 8; - sp -= 8; (sp as *mut usize).write(entry as usize); // ret target - sp -= 8; (sp as *mut usize).write(0); // rbx - sp -= 8; (sp as *mut usize).write(0); // rbp - sp -= 8; (sp as *mut usize).write(0); // r12 - sp -= 8; (sp as *mut usize).write(0); // r13 - sp -= 8; (sp as *mut usize).write(0); // r14 - sp -= 8; (sp as *mut usize).write(0); // r15 - sp - } -} - -// --------------------------------------------------------------------------- -// Context switch shims -// -// Each shim: -// 1. Pushes the six callee-saved integer registers. -// 2. Snaps rsp into rdi and calls the Rust helper that stores it. -// 3. Calls the Rust helper that returns the *other* side's saved rsp. -// 4. Moves that into rsp. -// 5. Pops the six registers and rets. -// -// 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 -// to the stack before we get here. (This is the same argument the compiler -// uses internally — callee-saved regs are what survive a `call`, and the -// SysV AMD64 ABI says XMM0–15 are all caller-saved.) If we ever yield from -// a place that isn't a Rust call boundary, this assumption breaks. -// --------------------------------------------------------------------------- - -#[unsafe(naked)] -unsafe extern "C" fn switch_to_actor_asm() { - core::arch::naked_asm!( - "push rbx", "push rbp", "push r12", "push r13", "push r14", "push r15", - "mov rdi, rsp", - "call {set_sched_sp}", - "call {get_actor_sp}", - "mov rsp, rax", - "pop r15", "pop r14", "pop r13", "pop r12", "pop rbp", "pop rbx", - "ret", - 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. -pub unsafe fn switch_to_actor() { - unsafe { switch_to_actor_asm() }; -} - -#[unsafe(naked)] -pub unsafe extern "C" fn switch_to_scheduler() { - core::arch::naked_asm!( - "push rbx", "push rbp", "push r12", "push r13", "push r14", "push r15", - "mov rdi, rsp", - "call {set_actor_sp}", - "call {get_sched_sp}", - "mov rsp, rax", - "pop r15", "pop r14", "pop r13", "pop r12", "pop rbp", "pop rbx", - "ret", - set_actor_sp = sym set_actor_sp, - get_sched_sp = sym get_scheduler_sp, - ); -} +pub use crate::arch::{ + get_actor_sp, init_actor_stack, set_actor_sp, switch_to_actor, switch_to_scheduler, +}; diff --git a/src/lib.rs b/src/lib.rs index 3faa229..b4dabf1 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,6 +11,7 @@ //! //! See `LOOM.md` for the design intent and the deferred-for-later list. +pub mod arch; pub mod stack; pub mod context; pub mod preempt; diff --git a/src/preempt.rs b/src/preempt.rs index e8b88c1..b8c35f9 100644 --- a/src/preempt.rs +++ b/src/preempt.rs @@ -30,7 +30,17 @@ use std::cell::Cell; use std::sync::atomic::{AtomicBool, Ordering}; pub const DEFAULT_ALLOC_INTERVAL: u32 = 128; + +// The timeslice is measured in `rdtsc()` ticks, and the tick *unit* differs by +// ISA (see `crate::arch::read_cycle_counter`). Both constants target ~100µs. +// x86-64 : TSC ≈ CPU base clock; 300_000 ticks ≈ 100µs on a 3 GHz core. +// aarch64: CNTVCT runs at CNTFRQ_EL0, commonly ~24 MHz; 2_400 ticks ≈ 100µs. +// A real implementation would read the frequency at startup and compute this; +// for now the constant is calibrated per-arch. Tune on-device if needed. +#[cfg(target_arch = "x86_64")] pub const DEFAULT_TIMESLICE_CYCLES: u64 = 300_000; // ≈ 100µs on a 3 GHz CPU +#[cfg(target_arch = "aarch64")] +pub const DEFAULT_TIMESLICE_CYCLES: u64 = 2_400; // ≈ 100µs at a 24 MHz CNTVCT thread_local! { /// While `false`, the allocator hook is a no-op. @@ -116,14 +126,12 @@ pub fn reset_timeslice() { TIMESLICE_START.with(|c| c.set(rdtsc())); } +/// Per-core cycle counter for the timeslice clock. Delegates to the active +/// arch backend (`rdtsc` on x86-64, `CNTVCT_EL0` on aarch64). The tick unit +/// is ISA-defined, so `DEFAULT_TIMESLICE_CYCLES` is calibrated per-arch below. #[inline(always)] pub fn rdtsc() -> u64 { - unsafe { - // SAFETY: x86-64 only. `lfence` serialises the instruction stream so - // we don't measure time before prior instructions retire. - core::arch::asm!("lfence", options(nostack, nomem, preserves_flags)); - core::arch::x86_64::_rdtsc() - } + crate::arch::read_cycle_counter() } pub struct PreemptingAllocator; diff --git a/tests/context.rs b/tests/context.rs index 150bcb1..397d1d3 100644 --- a/tests/context.rs +++ b/tests/context.rs @@ -58,6 +58,7 @@ use std::sync::OnceLock; static REG_BEFORE: OnceLock<[u64; 4]> = OnceLock::new(); static REG_AFTER: OnceLock<[u64; 4]> = OnceLock::new(); +#[cfg(target_arch = "x86_64")] extern "C-unwind" fn actor_reg_check() { unsafe { let s0: u64 = 0xAAAA_BBBB_0000_0001; @@ -83,6 +84,39 @@ extern "C-unwind" fn actor_reg_check() { } } +// aarch64 sibling: same intent, AAPCS64 callee-saved integer registers. +// LLVM reserves x19 (and uses x29/x30 as fp/lr), so those can't be named as +// asm operands. We use x23–x26 instead — callee-saved and bindable — moving +// known values into them, yielding, then reading them back, exactly mirroring +// the x86 test's move-into-r12..r15 pattern. The shim saves x23–x26 in its +// `stp x23,x24` / `stp x25,x26` pairs, so a clean round-trip proves they +// survive the switch. +#[cfg(target_arch = "aarch64")] +extern "C-unwind" fn actor_reg_check() { + unsafe { + let s0: u64 = 0xAAAA_BBBB_0000_0001; + let s1: u64 = 0xCCCC_DDDD_0000_0002; + let s2: u64 = 0xEEEE_FFFF_0000_0003; + let s3: u64 = 0x1111_2222_0000_0004; + + core::arch::asm!( + "mov x23, {s0}", "mov x24, {s1}", "mov x25, {s2}", "mov x26, {s3}", + s0 = in(reg) s0, s1 = in(reg) s1, s2 = in(reg) s2, s3 = in(reg) s3, + out("x23") _, out("x24") _, out("x25") _, out("x26") _, + ); + REG_BEFORE.set([s0, s1, s2, s3]).ok(); + switch_to_scheduler(); + + let a0: u64; let a1: u64; let a2: u64; let a3: u64; + core::arch::asm!( + "mov {a0}, x23", "mov {a1}, x24", "mov {a2}, x25", "mov {a3}, x26", + a0 = out(reg) a0, a1 = out(reg) a1, a2 = out(reg) a2, a3 = out(reg) a3, + ); + REG_AFTER.set([a0, a1, a2, a3]).ok(); + switch_to_scheduler(); + } +} + #[test] fn callee_saved_registers_survive_yield() { let stack = Stack::new(64 * 1024).unwrap();