Move CpuSet out of ostd::arch and implement this_cpu

Co-authored-by: Chuandong Li <lichuand@pku.edu.cn>

ostd/src/arch/x86/boot/mod.rs

@@ -22,6 +22,8 @@ mod linux_boot;
 mod multiboot;
 mod multiboot2;
 
+pub mod smp;
+
 use core::arch::global_asm;
 
 global_asm!(include_str!("boot.S"));

ostd/src/arch/x86/boot/smp.rs

@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: MPL-2.0
+
+//! Multiprocessor Boot Support
+//!
+//! The MP initialization protocol defines two classes of processors:
+//! the bootstrap processor (BSP) and the application processors (APs).
+//! Following a power-up or RESET of an MP system, system hardware dynamically
+//! selects one of the processors on the system bus as the BSP. The remaining
+//! processors are designated as APs.
+//!
+//! The BSP executes the BIOS's boot-strap code to configure the APIC environment,
+//! sets up system-wide data structures. Up to now, BSP has completed most of the
+//! initialization of the OS, but APs has not been awakened.
+
+use acpi::platform::{PlatformInfo, ProcessorInfo};
+
+use crate::arch::x86::kernel::acpi::ACPI_TABLES;
+
+/// Get processor information
+///
+/// This function needs to be called after the OS initializes the ACPI table.
+pub(crate) fn get_processor_info() -> Option<ProcessorInfo> {
+    if !ACPI_TABLES.is_completed() {
+        return None;
+    }
+    Some(
+        PlatformInfo::new(&*ACPI_TABLES.get().unwrap().lock())
+            .unwrap()
+            .processor_info
+            .unwrap(),
+    )
+}

ostd/src/arch/x86/cpu/mod.rs

@@ -4,17 +4,12 @@
 
 pub mod local;
 
-use alloc::vec::Vec;
 use core::{
     arch::x86_64::{_fxrstor, _fxsave},
     fmt::Debug,
 };
 
 use bitflags::bitflags;
-use bitvec::{
-    prelude::{BitVec, Lsb0},
-    slice::IterOnes,
-};
 use log::debug;
 #[cfg(feature = "intel_tdx")]
 use tdx_guest::tdcall;
@@ -29,86 +24,6 @@ use crate::{
     user::{ReturnReason, UserContextApi, UserContextApiInternal},
 };
 
-/// Returns the number of CPUs.
-pub fn num_cpus() -> u32 {
-    // FIXME: we only start one cpu now.
-    1
-}
-
-/// Returns the ID of this CPU.
-pub fn this_cpu() -> u32 {
-    // FIXME: we only start one cpu now.
-    0
-}
-
-/// A set of CPUs.
-#[derive(Default)]
-pub struct CpuSet {
-    bitset: BitVec,
-}
-
-impl CpuSet {
-    /// Creates a new `CpuSet` with all CPUs included.
-    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 included.
-    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 with identifier `cpu_id` to the `CpuSet`.
-    pub fn add(&mut self, cpu_id: u32) {
-        self.bitset.set(cpu_id as usize, true);
-    }
-
-    /// Adds multiple CPUs from `cpu_ids` to the `CpuSet`.
-    pub fn add_from_vec(&mut self, cpu_ids: Vec<u32>) {
-        for cpu_id in cpu_ids {
-            self.add(cpu_id)
-        }
-    }
-
-    /// Adds all available CPUs to the `CpuSet`.
-    pub fn add_all(&mut self) {
-        self.bitset.fill(true);
-    }
-
-    /// Removes a CPU with identifier `cpu_id` from the `CpuSet`.
-    pub fn remove(&mut self, cpu_id: u32) {
-        self.bitset.set(cpu_id as usize, false);
-    }
-
-    /// Removes multiple CPUs from `cpu_ids` from the `CpuSet`.
-    pub fn remove_from_vec(&mut self, cpu_ids: Vec<u32>) {
-        for cpu_id in cpu_ids {
-            self.remove(cpu_id);
-        }
-    }
-
-    /// Clears the `CpuSet`, removing all CPUs.
-    pub fn clear(&mut self) {
-        self.bitset.fill(false);
-    }
-
-    /// Checks if the `CpuSet` contains a specific CPU.
-    pub fn contains(&self, cpu_id: u32) -> bool {
-        self.bitset.get(cpu_id as usize).as_deref() == Some(&true)
-    }
-
-    /// Returns an iterator over the set CPUs.
-    pub fn iter(&self) -> IterOnes<'_, usize, Lsb0> {
-        self.bitset.iter_ones()
-    }
-}
-
 /// Cpu context, including both general-purpose registers and floating-point registers.
 #[derive(Clone, Default, Copy, Debug)]
 #[repr(C)]

ostd/src/arch/x86/mod.rs

@@ -32,10 +32,8 @@ use {
     ::tdx_guest::{init_tdx, tdcall::InitError, tdx_is_enabled},
 };
 
-pub(crate) fn before_all_init() {
-    enable_common_cpu_features();
-    serial::init();
 #[cfg(feature = "intel_tdx")]
+pub(crate) fn check_tdx_init() {
     match init_tdx() {
         Ok(td_info) => {
             early_println!(
@@ -138,7 +136,7 @@ fn has_avx512() -> bool {
     cpuid_result.ebx & (1 << 16) != 0
 }
 
-fn enable_common_cpu_features() {
+pub(crate) fn enable_cpu_features() {
     use x86_64::registers::{control::Cr4Flags, model_specific::EferFlags, xcontrol::XCr0Flags};
     let mut cr4 = x86_64::registers::control::Cr4::read();
     cr4 |= Cr4Flags::FSGSBASE

ostd/src/cpu/mod.rs

@@ -9,3 +9,110 @@ cfg_if::cfg_if! {
         pub use crate::arch::x86::cpu::*;
     }
 }
+
+use alloc::vec::Vec;
+
+use bitvec::{
+    prelude::{BitVec, Lsb0},
+    slice::IterOnes,
+};
+use spin::Once;
+
+use crate::{arch::boot::smp::get_processor_info, cpu};
+
+/// The number of CPUs.
+pub static NUM_CPUS: Once<u32> = Once::new();
+
+/// Initializes the number of CPUs.
+pub fn init() {
+    let processor_info = get_processor_info();
+    let num_processors = match processor_info {
+        Some(info) => info.application_processors.len() + 1,
+        None => 1,
+    };
+    NUM_CPUS.call_once(|| num_processors as u32);
+}
+
+/// Returns the number of CPUs.
+pub fn num_cpus() -> u32 {
+    *NUM_CPUS.get().unwrap()
+}
+
+/// Returns the ID of this CPU.
+///
+/// The CPU ID is strategically placed at the beginning of the CPU local storage area.
+pub fn this_cpu() -> u32 {
+    // SAFETY: the cpu ID is stored at the beginning of the cpu local area, provided
+    // by the linker script.
+    unsafe { (cpu::local::get_base() as usize as *mut u32).read() }
+}
+
+/// 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 a list of CPUs to the set.
+    pub fn add_from_vec(&mut self, cpu_ids: Vec<u32>) {
+        for cpu_id in cpu_ids {
+            self.add(cpu_id)
+        }
+    }
+
+    /// 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 a list of CPUs from the set.
+    pub fn remove_from_vec(&mut self, cpu_ids: Vec<u32>) {
+        for cpu_id in cpu_ids {
+            self.remove(cpu_id);
+        }
+    }
+
+    /// 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)
+    }
+
+    /// Iterates over the CPUs in the set.
+    pub fn iter(&self) -> IterOnes<'_, usize, Lsb0> {
+        self.bitset.iter_ones()
+    }
+}

ostd/src/lib.rs

@@ -57,7 +57,11 @@ pub use self::{cpu::cpu_local::CpuLocal, error::Error, prelude::Result};
 /// make inter-initialization-dependencies more clear and reduce usages of
 /// boot stage only global variables.
 pub fn init() {
-    arch::before_all_init();
+    arch::enable_cpu_features();
+    arch::serial::init();
+
+    #[cfg(feature = "intel_tdx")]
+    arch::check_tdx_init();
 
     // SAFETY: This function is called only once and only on the BSP.
     unsafe { cpu::cpu_local::early_init_bsp_local_base() };
@@ -77,6 +81,9 @@ pub fn init() {
 
     trap::init();
     arch::after_all_init();
+
+    cpu::init();
+
     bus::init();
 
     mm::kspace::activate_kernel_page_table();