--- a/src/percpu.rs +++ b/src/percpu.rs @@ -5,9 +5,13 @@ use core::{ cell::{Cell, RefCell}, hint, - sync::atomic::{AtomicBool, AtomicPtr, Ordering}, + sync::atomic::{AtomicBool, AtomicPtr, AtomicU32, AtomicU64, Ordering}, }; +/// Maximum number of pages to flush individually using INVLPG before falling +/// back to a full TLB flush (CR3 reload). +const TLB_RANGE_THRESHOLD: u32 = 32; + use rmm::Arch; use syscall::PtraceFlags; @@ -41,6 +45,23 @@ /// Counts how many times the scheduler MCS lock acquisition was contended. pub mcs_contention_count: Cell, + /// TLB shootdown range: start virtual address (page-aligned). + /// Set to 0 for a full flush. Only valid when `wants_tlb_shootdown` is true. + pub tlb_flush_start: AtomicU64, + /// TLB shootdown range: number of pages to invalidate. + pub tlb_flush_count: AtomicU32, + + /// Priority inheritance donation. When another CPU is blocked waiting on a + /// lock this CPU holds, the blocked CPU may donate its priority here. + /// `u32::MAX` means no donation; otherwise it's a priority level (0-39). + pub pi_donated_prio: AtomicU32, + + /// Cached priority of the currently-running context on this CPU. + /// Set by the scheduler when selecting a new context. Read by the MCS + /// lock during priority donation — avoids acquiring the context RwLock + /// from the spin loop. Default 39 (lowest priority). + pub current_prio: Cell, + // TODO: Put mailbox queues here, e.g. for TLB shootdown? Just be sure to 128-byte align it // first to avoid cache invalidation. pub profiling: Option<&'static crate::profiling::RingBuffer>, @@ -64,6 +85,15 @@ ALL_PERCPU_BLOCKS[id.get() as usize].store(block, Ordering::Release) } +/// Get a reference to another CPU's PercpuBlock by logical CPU ID. +pub fn get_for_cpu(id: LogicalCpuId) -> Option<&'static PercpuBlock> { + unsafe { + ALL_PERCPU_BLOCKS[id.get() as usize] + .load(Ordering::Acquire) + .as_ref() + } +} + pub fn get_all_stats() -> Vec<(LogicalCpuId, CpuStatsData)> { let mut res = ALL_PERCPU_BLOCKS .iter() @@ -108,6 +138,9 @@ core::hint::spin_loop(); } } + // Full flush — clear range info + percpublock.tlb_flush_start.store(0, Ordering::Relaxed); + percpublock.tlb_flush_count.store(0, Ordering::Relaxed); crate::ipi::ipi_single(crate::ipi::IpiKind::Tlb, percpublock); } else { @@ -138,20 +171,114 @@ hint::spin_loop(); } } + // Full flush — clear range info + percpublock.tlb_flush_start.store(0, Ordering::Relaxed); + percpublock.tlb_flush_count.store(0, Ordering::Relaxed); } // Single broadcast IPI to all other CPUs using destination shorthand crate::ipi::ipi(crate::ipi::IpiKind::Tlb, crate::ipi::IpiTarget::Other); } } + +/// Range-based TLB shootdown IPI. Only invalidates the specified virtual address +/// range using INVLPG per page for ranges up to TLB_RANGE_THRESHOLD pages. +/// Falls back to full flush for larger ranges. +pub fn shootdown_tlb_ipi_range(target: Option, start: usize, count: usize) { + if cfg!(not(feature = "multi_core")) { + return; + } + + let start_aligned = start as u64 & !0xFFF; + let count_u32 = count as u32; + let use_range = count_u32 > 0 && count_u32 <= TLB_RANGE_THRESHOLD; + + let set_range = |percpublock: &PercpuBlock| { + if use_range { + percpublock.tlb_flush_start.store(start_aligned, Ordering::Release); + percpublock.tlb_flush_count.store(count_u32, Ordering::Release); + } else { + percpublock.tlb_flush_start.store(0, Ordering::Release); + percpublock.tlb_flush_count.store(0, Ordering::Release); + } + }; + + if let Some(target) = target { + let my_percpublock = PercpuBlock::current(); + assert_ne!(target, my_percpublock.cpu_id); + + let Some(percpublock) = (unsafe { + ALL_PERCPU_BLOCKS[target.get() as usize] + .load(Ordering::Acquire) + .as_ref() + }) else { + return; + }; + #[expect(clippy::bool_comparison)] + while percpublock.wants_tlb_shootdown.swap(true, Ordering::Release) == true { + while percpublock.wants_tlb_shootdown.load(Ordering::Relaxed) == true { + my_percpublock.maybe_handle_tlb_shootdown(); + hint::spin_loop(); + } + } + set_range(percpublock); + crate::ipi::ipi_single(crate::ipi::IpiKind::Tlb, percpublock); + } else { + let my_percpublock = PercpuBlock::current(); + for id in 0..crate::cpu_count() { + let target_id = LogicalCpuId::new(id); + if target_id == my_percpublock.cpu_id { + continue; + } + let Some(percpublock) = (unsafe { + ALL_PERCPU_BLOCKS[id as usize] + .load(Ordering::Acquire) + .as_ref() + }) else { + continue; + }; + #[expect(clippy::bool_comparison)] + while percpublock.wants_tlb_shootdown.swap(true, Ordering::Release) == true { + while percpublock.wants_tlb_shootdown.load(Ordering::Relaxed) == true { + my_percpublock.maybe_handle_tlb_shootdown(); + hint::spin_loop(); + } + } + set_range(percpublock); + } + crate::ipi::ipi(crate::ipi::IpiKind::Tlb, crate::ipi::IpiTarget::Other); + } +} impl PercpuBlock { + /// Return the effective scheduling priority, accounting for priority inheritance. + /// Lower number = higher priority (0-39 range). + pub fn effective_prio(&self, context_prio: usize) -> usize { + let donated = self.pi_donated_prio.load(Ordering::Relaxed); + if donated < context_prio as u32 { + donated as usize + } else { + context_prio + } + } + pub fn maybe_handle_tlb_shootdown(&self) { #[expect(clippy::bool_comparison)] if self.wants_tlb_shootdown.swap(false, Ordering::Relaxed) == false { return; } - // TODO: Finer-grained flush - crate::memory::RmmA::invalidate_all(); + let start = self.tlb_flush_start.load(Ordering::Acquire); + let count = self.tlb_flush_count.load(Ordering::Acquire); + + if start != 0 && count > 0 && count <= TLB_RANGE_THRESHOLD { + // Range-based flush using INVLPG per page — cheaper than full CR3 reload. + for i in 0..count { + let addr = start + (i as u64) * 4096; + crate::memory::RmmA::invalidate(rmm::VirtualAddress::new(addr as usize)); + } + } else { + // Full TLB flush (CR3 reload) for large ranges or global shootdowns. + crate::memory::RmmA::invalidate_all(); + } if let Some(addrsp) = &*self.current_addrsp.borrow() { addrsp.tlb_ack.fetch_add(1, Ordering::Release); @@ -223,6 +350,10 @@ last_queue: Cell::new(39), mcs_sched_node: McsNode::new(), mcs_contention_count: Cell::new(0), + tlb_flush_start: AtomicU64::new(0), + tlb_flush_count: AtomicU32::new(0), + pi_donated_prio: AtomicU32::new(u32::MAX), + current_prio: Cell::new(39), ptrace_flags: Cell::new(PtraceFlags::empty()), ptrace_session: RefCell::new(None), inside_syscall: Cell::new(false),