From 26c4abde58a063223b2ed5aaa7cc3fb9182f4f2d Mon Sep 17 00:00:00 2001 From: Zhang Junyang Date: Mon, 23 Sep 2024 09:57:17 +0800 Subject: [PATCH] Refactor `CpuSet` and introduce `AtomicCpuSet` --- Cargo.lock | 1 + ostd/Cargo.toml | 1 + ostd/src/cpu/mod.rs | 72 +----------- ostd/src/cpu/set.rs | 270 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 274 insertions(+), 70 deletions(-) create mode 100644 ostd/src/cpu/set.rs diff --git a/Cargo.lock b/Cargo.lock index 3d5c78f1..dd987e29 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1123,6 +1123,7 @@ dependencies = [ "owo-colors 3.5.0", "riscv", "sbi-rt", + "smallvec", "spin 0.9.8", "static_assertions", "tdx-guest", diff --git a/ostd/Cargo.toml b/ostd/Cargo.toml index ba1c3010..0dca04dd 100644 --- a/ostd/Cargo.toml +++ b/ostd/Cargo.toml @@ -38,6 +38,7 @@ ostd-test = { path = "libs/ostd-test", version = "0.9.2" } owo-colors = { version = "3", optional = true } ostd-pod = { git = "https://github.com/asterinas/ostd-pod", rev = "c4644be", version = "0.1.1" } spin = "0.9.4" +smallvec = "1.13.2" static_assertions = "1.1.0" unwinding = { version = "0.2.3", default-features = false, features = ["fde-gnu-eh-frame-hdr", "hide-trace", "panic", "personality", "unwinder"] } volatile = { version = "0.4.5", features = ["unstable"] } diff --git a/ostd/src/cpu/mod.rs b/ostd/src/cpu/mod.rs index 6525d4cf..1c50d15a 100644 --- a/ostd/src/cpu/mod.rs +++ b/ostd/src/cpu/mod.rs @@ -3,6 +3,7 @@ //! CPU-related definitions. pub mod local; +pub mod set; cfg_if::cfg_if! { if #[cfg(target_arch = "x86_64")] { @@ -12,7 +13,7 @@ cfg_if::cfg_if! { } } -use bitvec::prelude::BitVec; +pub use set::{AtomicCpuSet, CpuSet}; use spin::Once; use crate::{ @@ -88,72 +89,3 @@ cpu_local_cell! { /// The number of the current CPU. static CURRENT_CPU: u32 = u32::MAX; } - -/// A subset of all CPUs in the system. -/// -/// This structure can be used to mask out a subset of CPUs in the system. -#[derive(Clone, Debug, Default)] -pub struct CpuSet { - bitset: BitVec, -} - -impl CpuSet { - /// Creates a new `CpuSet` with all CPUs in the system. - pub fn new_full() -> Self { - let num_cpus = num_cpus(); - let mut bitset = BitVec::with_capacity(num_cpus as usize); - bitset.resize(num_cpus as usize, true); - Self { bitset } - } - - /// Creates a new `CpuSet` with no CPUs in the system. - pub fn new_empty() -> Self { - let num_cpus = num_cpus(); - let mut bitset = BitVec::with_capacity(num_cpus as usize); - bitset.resize(num_cpus as usize, false); - Self { bitset } - } - - /// Adds a CPU to the set. - pub fn add(&mut self, cpu_id: u32) { - self.bitset.set(cpu_id as usize, true); - } - - /// Adds all CPUs to the set. - pub fn add_all(&mut self) { - self.bitset.fill(true); - } - - /// Removes a CPU from the set. - pub fn remove(&mut self, cpu_id: u32) { - self.bitset.set(cpu_id as usize, false); - } - - /// Removes all CPUs from the set. - pub fn clear(&mut self) { - self.bitset.fill(false); - } - - /// Returns true if the set contains the specified CPU. - pub fn contains(&self, cpu_id: u32) -> bool { - self.bitset.get(cpu_id as usize).as_deref() == Some(&true) - } - - /// Returns the number of CPUs in the set. - pub fn count(&self) -> usize { - self.bitset.count_ones() - } - - /// Iterates over the CPUs in the set. - pub fn iter(&self) -> impl Iterator + '_ { - self.bitset.iter_ones().map(|idx| idx as u32) - } -} - -impl From for CpuSet { - fn from(cpu_id: u32) -> Self { - let mut set = Self::new_empty(); - set.add(cpu_id); - set - } -} diff --git a/ostd/src/cpu/set.rs b/ostd/src/cpu/set.rs new file mode 100644 index 00000000..01db7bb2 --- /dev/null +++ b/ostd/src/cpu/set.rs @@ -0,0 +1,270 @@ +// SPDX-License-Identifier: MPL-2.0 + +//! This module contains the implementation of the CPU set and atomic CPU set. + +use core::sync::atomic::{AtomicU64, Ordering}; + +use smallvec::SmallVec; +use static_assertions::const_assert_eq; + +use super::num_cpus; + +/// A subset of all CPUs in the system. +#[derive(Clone, Debug, Default)] +pub struct CpuSet { + // A bitset representing the CPUs in the system. + bits: SmallVec<[InnerPart; NR_PARTS_NO_ALLOC]>, +} + +type InnerPart = u64; + +const BITS_PER_PART: usize = core::mem::size_of::() * 8; +const NR_PARTS_NO_ALLOC: usize = 2; + +const fn part_idx(cpu_id: u32) -> usize { + cpu_id as usize / BITS_PER_PART +} + +const fn bit_idx(cpu_id: u32) -> usize { + cpu_id as usize % BITS_PER_PART +} + +const fn parts_for_cpus(num_cpus: usize) -> usize { + num_cpus.div_ceil(BITS_PER_PART) +} + +impl CpuSet { + /// Creates a new `CpuSet` with all CPUs in the system. + pub fn new_full() -> Self { + let mut ret = Self::with_capacity_val(num_cpus() as usize, !0); + ret.clear_nonexistent_cpu_bits(); + ret + } + + /// Creates a new `CpuSet` with no CPUs in the system. + pub fn new_empty() -> Self { + Self::with_capacity_val(num_cpus() as usize, 0) + } + + /// Adds a CPU to the set. + pub fn add(&mut self, cpu_id: u32) { + let part_idx = part_idx(cpu_id); + let bit_idx = bit_idx(cpu_id); + if part_idx >= self.bits.len() { + self.bits.resize(part_idx + 1, 0); + } + self.bits[part_idx] |= 1 << bit_idx; + } + + /// Removes a CPU from the set. + pub fn remove(&mut self, cpu_id: u32) { + let part_idx = part_idx(cpu_id); + let bit_idx = bit_idx(cpu_id); + if part_idx < self.bits.len() { + self.bits[part_idx] &= !(1 << bit_idx); + } + } + + /// Returns true if the set contains the specified CPU. + pub fn contains(&self, cpu_id: u32) -> bool { + let part_idx = part_idx(cpu_id); + let bit_idx = bit_idx(cpu_id); + part_idx < self.bits.len() && (self.bits[part_idx] & (1 << bit_idx)) != 0 + } + + /// Returns the number of CPUs in the set. + pub fn count(&self) -> usize { + self.bits + .iter() + .map(|part| part.count_ones() as usize) + .sum() + } + + /// Adds all CPUs to the set. + pub fn add_all(&mut self) { + self.bits.fill(!0); + self.clear_nonexistent_cpu_bits(); + } + + /// Removes all CPUs from the set. + pub fn clear(&mut self) { + self.bits.fill(0); + } + + /// Iterates over the CPUs in the set. + pub fn iter(&self) -> impl Iterator + '_ { + self.bits.iter().enumerate().flat_map(|(part_idx, &part)| { + (0..BITS_PER_PART).filter_map(move |bit_idx| { + if (part & (1 << bit_idx)) != 0 { + Some((part_idx * BITS_PER_PART + bit_idx) as u32) + } else { + None + } + }) + }) + } + + /// Only for internal use. The set cannot contain non-existent CPUs. + fn with_capacity_val(num_cpus: usize, val: InnerPart) -> Self { + let num_parts = parts_for_cpus(num_cpus); + let mut bits = SmallVec::with_capacity(num_parts); + bits.resize(num_parts, val); + Self { bits } + } + + fn clear_nonexistent_cpu_bits(&mut self) { + let num_cpus = num_cpus() as usize; + if num_cpus % BITS_PER_PART != 0 { + let num_parts = parts_for_cpus(num_cpus); + self.bits[num_parts - 1] &= (1 << (num_cpus % BITS_PER_PART)) - 1; + } + } +} + +impl From for CpuSet { + fn from(cpu_id: u32) -> Self { + let mut set = Self::new_empty(); + set.add(cpu_id); + set + } +} + +/// A subset of all CPUs in the system with atomic operations. +/// +/// It provides atomic operations for each CPU in the system. When the +/// operation contains multiple CPUs, the ordering is not guaranteed. +pub struct AtomicCpuSet { + bits: SmallVec<[AtomicInnerPart; NR_PARTS_NO_ALLOC]>, +} + +type AtomicInnerPart = AtomicU64; +const_assert_eq!(core::mem::size_of::() * 8, BITS_PER_PART); + +impl AtomicCpuSet { + /// Creates a new `AtomicCpuSet` with an initial value. + pub fn new(value: CpuSet) -> Self { + let bits = value.bits.into_iter().map(AtomicU64::new).collect(); + Self { bits } + } + + /// Loads the value of the set. + /// + /// This operation can only be done in the [`Ordering::Relaxed`] memory + /// order. It cannot be used to synchronize anything between CPUs. + pub fn load(&self) -> CpuSet { + let bits = self + .bits + .iter() + .map(|part| part.load(Ordering::Relaxed)) + .collect(); + CpuSet { bits } + } + + /// Stores a new value to the set. + /// + /// This operation can only be done in the [`Ordering::Relaxed`] memory + /// order. It cannot be used to synchronize anything between CPUs. + pub fn store(&self, value: CpuSet) { + for (part, new_part) in self.bits.iter().zip(value.bits) { + part.store(new_part, Ordering::Relaxed); + } + } + + /// Atomically adds a CPU with the given ordering. + pub fn add(&self, cpu_id: u32, ordering: Ordering) { + let part_idx = part_idx(cpu_id); + let bit_idx = bit_idx(cpu_id); + if part_idx < self.bits.len() { + self.bits[part_idx].fetch_or(1 << bit_idx, ordering); + } + } + + /// Atomically removes a CPU with the given ordering. + pub fn remove(&self, cpu_id: u32, ordering: Ordering) { + let part_idx = part_idx(cpu_id); + let bit_idx = bit_idx(cpu_id); + if part_idx < self.bits.len() { + self.bits[part_idx].fetch_and(!(1 << bit_idx), ordering); + } + } + + /// Atomically checks if the set contains the specified CPU. + pub fn contains(&self, cpu_id: u32, ordering: Ordering) -> bool { + let part_idx = part_idx(cpu_id); + let bit_idx = bit_idx(cpu_id); + part_idx < self.bits.len() && (self.bits[part_idx].load(ordering) & (1 << bit_idx)) != 0 + } +} + +#[cfg(ktest)] +mod test { + use super::*; + use crate::prelude::*; + + #[ktest] + fn test_full_cpu_set_iter_is_all() { + let set = CpuSet::new_full(); + let num_cpus = num_cpus(); + let all_cpus = (0..num_cpus).collect::>(); + let set_cpus = set.iter().collect::>(); + + assert!(set_cpus.len() == num_cpus as usize); + assert_eq!(set_cpus, all_cpus); + } + + #[ktest] + fn test_full_cpu_set_contains_all() { + let set = CpuSet::new_full(); + let num_cpus = num_cpus(); + for cpu_id in 0..num_cpus { + assert!(set.contains(cpu_id)); + } + } + + #[ktest] + fn test_empty_cpu_set_iter_is_empty() { + let set = CpuSet::new_empty(); + let set_cpus = set.iter().collect::>(); + assert!(set_cpus.is_empty()); + } + + #[ktest] + fn test_empty_cpu_set_contains_none() { + let set = CpuSet::new_empty(); + let num_cpus = num_cpus(); + for cpu_id in 0..num_cpus { + assert!(!set.contains(cpu_id)); + } + } + + #[ktest] + fn test_atomic_cpu_set_multiple_sizes() { + for test_num_cpus in [1, 3, 12, 64, 96, 99, 128, 256, 288, 1024] { + let set = CpuSet::with_capacity_val(test_num_cpus, 0); + let atomic_set = AtomicCpuSet::new(set); + + for cpu_id in 0..test_num_cpus as u32 { + assert!(!atomic_set.contains(cpu_id, Ordering::Relaxed)); + if cpu_id % 3 == 0 { + atomic_set.add(cpu_id, Ordering::Relaxed); + } + } + + let loaded = atomic_set.load(); + for cpu_id in loaded.iter() { + if cpu_id % 3 == 0 { + assert!(loaded.contains(cpu_id)); + } else { + assert!(!loaded.contains(cpu_id)); + } + } + + atomic_set.store(CpuSet::with_capacity_val(test_num_cpus, 0)); + + for cpu_id in 0..test_num_cpus as u32 { + assert!(!atomic_set.contains(cpu_id, Ordering::Relaxed)); + atomic_set.add(cpu_id, Ordering::Relaxed); + } + } + } +}