Optimize the space for `TlbFlushOp`s

This commit is contained in:
Zhang Junyang 2025-08-14 22:08:56 +08:00 committed by Tate, Hongliang Tian
parent d0b98130cf
commit 1452aab69c
5 changed files with 144 additions and 64 deletions

View File

@ -739,7 +739,7 @@ impl Vmar_ {
rss_delta.add(vm_mapping.rss_type(), num_copied as isize);
}
cur_cursor.flusher().issue_tlb_flush(TlbFlushOp::All);
cur_cursor.flusher().issue_tlb_flush(TlbFlushOp::for_all());
cur_cursor.flusher().dispatch_tlb_flush();
cur_cursor.flusher().sync_tlb_flush();
}

View File

@ -231,7 +231,7 @@ impl VmMapping {
if VmPerms::from(prop.flags).contains(required_perms) {
// The page fault is already handled maybe by other threads.
// Just flush the TLB and return.
TlbFlushOp::Range(va).perform_on_current();
TlbFlushOp::for_range(va).perform_on_current();
return Ok(());
}
assert!(is_write);
@ -253,7 +253,7 @@ impl VmMapping {
if self.is_shared || only_reference {
cursor.protect_next(PAGE_SIZE, |p| p.flags |= new_flags);
cursor.flusher().issue_tlb_flush(TlbFlushOp::Range(va));
cursor.flusher().issue_tlb_flush(TlbFlushOp::for_range(va));
cursor.flusher().dispatch_tlb_flush();
} else {
let new_frame = duplicate_frame(&frame)?;
@ -558,7 +558,7 @@ impl VmMapping {
let op = |p: &mut PageProperty| p.flags = perms.into();
while cursor.virt_addr() < range.end {
if let Some(va) = cursor.protect_next(range.end - cursor.virt_addr(), op) {
cursor.flusher().issue_tlb_flush(TlbFlushOp::Range(va));
cursor.flusher().issue_tlb_flush(TlbFlushOp::for_range(va));
} else {
break;
}

View File

@ -688,7 +688,7 @@ mod vmspace {
let mut cursor_mut = vmspace
.cursor_mut(&preempt_guard, &range)
.expect("Failed to create mutable cursor");
cursor_mut.flusher().issue_tlb_flush(TlbFlushOp::All);
cursor_mut.flusher().issue_tlb_flush(TlbFlushOp::for_all());
cursor_mut.flusher().dispatch_tlb_flush();
}

View File

@ -4,6 +4,7 @@
use alloc::vec::Vec;
use core::{
mem::MaybeUninit,
ops::Range,
sync::atomic::{AtomicBool, Ordering},
};
@ -14,6 +15,7 @@ use super::{
};
use crate::{
arch::irq,
const_assert,
cpu::{AtomicCpuSet, CpuSet, PinCurrentCpu},
cpu_local,
sync::{LocalIrqDisabled, SpinLock},
@ -150,39 +152,99 @@ impl<'a, G: PinCurrentCpu> TlbFlusher<'a, G> {
}
/// The operation to flush TLB entries.
///
/// The variants of this structure are:
/// - Flushing all TLB entries except for the global entries;
/// - Flushing the TLB entry associated with an address;
/// - Flushing the TLB entries for a specific range of virtual addresses;
///
/// This is a `struct` instead of an `enum` because if trivially representing
/// the three variants with an `enum`, it would be 24 bytes. To minimize the
/// memory footprint, we encode all three variants into an 8-byte integer.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TlbFlushOp {
/// Flush all TLB entries except for the global entries.
All,
/// Flush the TLB entry for the specified virtual address.
Address(Vaddr),
/// Flush the TLB entries for the specified virtual address range.
Range(Range<Vaddr>),
}
pub struct TlbFlushOp(Vaddr);
// We require the address to be page-aligned, so the in-page offset part of the
// address can be used to store the length. A sanity check to ensure that we
// don't allow ranged flush operations with a too long length.
const_assert!(TlbFlushOp::FLUSH_RANGE_NPAGES_MASK | (PAGE_SIZE - 1) == PAGE_SIZE - 1);
impl TlbFlushOp {
const FLUSH_ALL_VAL: Vaddr = Vaddr::MAX;
const FLUSH_RANGE_NPAGES_MASK: Vaddr =
(1 << (usize::BITS - FLUSH_ALL_PAGES_THRESHOLD.leading_zeros())) - 1;
/// Performs the TLB flush operation on the current CPU.
pub fn perform_on_current(&self) {
use crate::arch::mm::{
tlb_flush_addr, tlb_flush_addr_range, tlb_flush_all_excluding_global,
};
match self {
TlbFlushOp::All => tlb_flush_all_excluding_global(),
TlbFlushOp::Address(addr) => tlb_flush_addr(*addr),
TlbFlushOp::Range(range) => tlb_flush_addr_range(range),
match self.0 {
Self::FLUSH_ALL_VAL => tlb_flush_all_excluding_global(),
addr => {
let start = addr & !Self::FLUSH_RANGE_NPAGES_MASK;
let num_pages = addr & Self::FLUSH_RANGE_NPAGES_MASK;
debug_assert!((addr & (PAGE_SIZE - 1)) < FLUSH_ALL_PAGES_THRESHOLD);
debug_assert!(num_pages != 0);
if num_pages == 1 {
tlb_flush_addr(start);
} else {
tlb_flush_addr_range(&(start..start + num_pages * PAGE_SIZE));
}
}
}
}
fn optimize_for_large_range(self) -> Self {
match self {
TlbFlushOp::Range(range) => {
if range.len() > FLUSH_ALL_RANGE_THRESHOLD {
TlbFlushOp::All
} else {
TlbFlushOp::Range(range)
}
}
_ => self,
/// Creates a new TLB flush operation that flushes all TLB entries except
/// for the global entries.
pub const fn for_all() -> Self {
TlbFlushOp(Self::FLUSH_ALL_VAL)
}
/// Creates a new TLB flush operation that flushes the TLB entry associated
/// with the provided virtual address.
pub const fn for_single(addr: Vaddr) -> Self {
TlbFlushOp(addr | 1)
}
/// Creates a new TLB flush operation that flushes the TLB entries for the
/// specified virtual address range.
///
/// If the range is too large, the resulting [`TlbFlushOp`] will flush all
/// TLB entries instead.
///
/// # Panics
///
/// Panics if the range is not page-aligned or if the range is empty.
pub const fn for_range(range: Range<Vaddr>) -> Self {
assert!(
range.start % PAGE_SIZE == 0,
"Range start must be page-aligned"
);
assert!(range.end % PAGE_SIZE == 0, "Range end must be page-aligned");
assert!(range.start < range.end, "Range must not be empty");
let num_pages = (range.end - range.start) / PAGE_SIZE;
if num_pages >= FLUSH_ALL_PAGES_THRESHOLD {
return TlbFlushOp::for_all();
}
TlbFlushOp(range.start | (num_pages as Vaddr))
}
/// Returns the number of pages to flush.
///
/// If it returns `u32::MAX`, it means to flush all the entries. Otherwise
/// the return value should be less than [`FLUSH_ALL_PAGES_THRESHOLD`] and
/// non-zero.
fn num_pages(&self) -> u32 {
if self.0 == Self::FLUSH_ALL_VAL {
u32::MAX
} else {
debug_assert!((self.0 & (PAGE_SIZE - 1)) < FLUSH_ALL_PAGES_THRESHOLD);
let num_pages = (self.0 & Self::FLUSH_RANGE_NPAGES_MASK) as u32;
debug_assert!(num_pages != 0);
num_pages
}
}
}
@ -214,32 +276,38 @@ fn do_remote_flush() {
new_op_queue.flush_all();
}
/// If a TLB flushing request exceeds this threshold, we flush all.
const FLUSH_ALL_RANGE_THRESHOLD: usize = 32 * PAGE_SIZE;
/// If the number of pending requests exceeds this threshold, we flush all the
/// If the number of pending pages to flush exceeds this threshold, we flush all the
/// TLB entries instead of flushing them one by one.
const FLUSH_ALL_OPS_THRESHOLD: usize = 32;
const FLUSH_ALL_PAGES_THRESHOLD: usize = 32;
struct OpsStack {
ops: [Option<TlbFlushOp>; FLUSH_ALL_OPS_THRESHOLD],
need_flush_all: bool,
size: usize,
/// From 0 to `num_ops`, the array entry must be initialized.
ops: [MaybeUninit<TlbFlushOp>; FLUSH_ALL_PAGES_THRESHOLD],
num_ops: u32,
/// If this is `u32::MAX`, we should flush all entries irrespective of the
/// contents of `ops`. And in this case `num_ops` must be zero.
///
/// Otherwise, it counts the number of pages to flush in `ops`.
num_pages_to_flush: u32,
page_keeper: Vec<Frame<dyn AnyFrameMeta>>,
}
impl OpsStack {
const fn new() -> Self {
Self {
ops: [const { None }; FLUSH_ALL_OPS_THRESHOLD],
need_flush_all: false,
size: 0,
ops: [const { MaybeUninit::uninit() }; FLUSH_ALL_PAGES_THRESHOLD],
num_ops: 0,
num_pages_to_flush: 0,
page_keeper: Vec::new(),
}
}
fn is_empty(&self) -> bool {
!self.need_flush_all && self.size == 0
self.num_ops == 0 && self.num_pages_to_flush == 0
}
fn need_flush_all(&self) -> bool {
self.num_pages_to_flush == u32::MAX
}
fn push(&mut self, op: TlbFlushOp, drop_after_flush: Option<Frame<dyn AnyFrameMeta>>) {
@ -247,55 +315,67 @@ impl OpsStack {
self.page_keeper.push(frame);
}
if self.need_flush_all {
if self.need_flush_all() {
return;
}
let op = op.optimize_for_large_range();
if op == TlbFlushOp::All || self.size >= FLUSH_ALL_OPS_THRESHOLD {
self.need_flush_all = true;
self.size = 0;
let op_num_pages = op.num_pages();
if op == TlbFlushOp::for_all()
|| self.num_pages_to_flush + op_num_pages >= FLUSH_ALL_PAGES_THRESHOLD as u32
{
self.num_pages_to_flush = u32::MAX;
self.num_ops = 0;
return;
}
self.ops[self.size] = Some(op);
self.size += 1;
self.ops[self.num_ops as usize].write(op);
self.num_ops += 1;
self.num_pages_to_flush += op_num_pages;
}
fn push_from(&mut self, other: &OpsStack) {
self.page_keeper.extend(other.page_keeper.iter().cloned());
if self.need_flush_all {
if self.need_flush_all() {
return;
}
if other.need_flush_all || self.size + other.size >= FLUSH_ALL_OPS_THRESHOLD {
self.need_flush_all = true;
self.size = 0;
if other.need_flush_all()
|| self.num_pages_to_flush + other.num_pages_to_flush
>= FLUSH_ALL_PAGES_THRESHOLD as u32
{
self.num_pages_to_flush = u32::MAX;
self.num_ops = 0;
return;
}
for i in 0..other.size {
self.ops[self.size] = other.ops[i].clone();
self.size += 1;
for other_op in other.ops_iter() {
self.ops[self.num_ops as usize].write(other_op.clone());
self.num_ops += 1;
}
self.num_pages_to_flush += other.num_pages_to_flush;
}
fn flush_all(&mut self) {
if self.need_flush_all {
if self.need_flush_all() {
crate::arch::mm::tlb_flush_all_excluding_global();
} else {
for i in 0..self.size {
if let Some(op) = &self.ops[i] {
op.perform_on_current();
}
}
self.ops_iter().for_each(|op| {
op.perform_on_current();
});
}
self.clear_without_flush();
}
fn clear_without_flush(&mut self) {
self.need_flush_all = false;
self.size = 0;
self.num_pages_to_flush = 0;
self.num_ops = 0;
self.page_keeper.clear();
}
fn ops_iter(&self) -> impl Iterator<Item = &TlbFlushOp> {
self.ops.iter().take(self.num_ops as usize).map(|op| {
// SAFETY: From 0 to `num_ops`, the array entry must be initialized.
unsafe { op.assume_init_ref() }
})
}
}

View File

@ -309,7 +309,7 @@ impl<'a> CursorMut<'a> {
debug_assert_eq!(va, start_va);
let (old_frame, _) = item;
self.flusher
.issue_tlb_flush_with(TlbFlushOp::Address(start_va), old_frame.into());
.issue_tlb_flush_with(TlbFlushOp::for_single(start_va), old_frame.into());
self.flusher.dispatch_tlb_flush();
}
PageTableFrag::StrayPageTable { .. } => {
@ -351,7 +351,7 @@ impl<'a> CursorMut<'a> {
let (frame, _) = item;
num_unmapped += 1;
self.flusher
.issue_tlb_flush_with(TlbFlushOp::Address(va), frame.into());
.issue_tlb_flush_with(TlbFlushOp::for_single(va), frame.into());
}
PageTableFrag::StrayPageTable {
pt,
@ -361,7 +361,7 @@ impl<'a> CursorMut<'a> {
} => {
num_unmapped += num_frames;
self.flusher
.issue_tlb_flush_with(TlbFlushOp::Range(va..va + len), pt);
.issue_tlb_flush_with(TlbFlushOp::for_range(va..va + len), pt);
}
}
}