Turn GS.base validity into a global invariant

ostd/src/arch/riscv/boot/boot.S

@@ -28,7 +28,11 @@ _start:
     # 2. set sp (BSP only)
     lga    sp, boot_stack_top
 
-    # 3. jump to rust riscv_boot
+    # 3. set gp (CPU-local address)
+.extern __cpu_local_start
+    lga    gp, __cpu_local_start
+
+    # 4. jump to rust riscv_boot
     lga    t0, riscv_boot
     jr     t0
 

ostd/src/arch/riscv/cpu/local.rs

@@ -2,14 +2,6 @@
 
 //! Architecture dependent CPU-local information utilities.
 
-pub(crate) unsafe fn set_base(addr: u64) {
-    core::arch::asm!(
-        "mv gp, {addr}",
-        addr = in(reg) addr,
-        options(preserves_flags, nostack)
-    );
-}
-
 pub(crate) fn get_base() -> u64 {
     let mut gp;
     unsafe {

ostd/src/arch/x86/boot/ap_boot.S

@@ -5,10 +5,6 @@
 .global ap_boot_from_real_mode
 .global ap_boot_from_long_mode
 
-.extern boot_gdtr
-.extern boot_page_table_start
-.extern ap_early_entry
-
 .section ".ap_boot", "awx"
 .align 4096
 
@@ -18,6 +14,7 @@ MMIO_XAPIC_APICID = 0xFEE00020
 
 .macro setup_64bit_gdt_and_page_table eax
     // Use the 64-bit GDT.
+.extern boot_gdtr
     lgdt [boot_gdtr]
 
     // Enable PAE and PGE.
@@ -62,7 +59,7 @@ ap_boot_from_long_mode:
     retf // 32-bit far return
 .align 8
 retf_stack_bottom:
-.long ap_long_mode
+.long ap_long_mode_in_low_address
 .long 0x8
 retf_stack_top:
 
@@ -161,10 +158,10 @@ ap_protect:
     or eax, 1 << 31
     mov cr0, eax
 
-    ljmp 0x8, offset ap_long_mode
+    ljmp 0x8, offset ap_long_mode_in_low_address
 
 .code64
-ap_long_mode:
+ap_long_mode_in_low_address:
     mov ax, 0
     mov ds, ax
     mov ss, ax
@@ -172,26 +169,40 @@ ap_long_mode:
     mov fs, ax
     mov gs, ax
 
+    // Update RIP to use the virtual address.
+    mov rax, offset ap_long_mode
+    jmp rax
+
+.data
+// This is a pointer to be filled by the BSP when boot information
+// of all APs are allocated and initialized.
+.global __ap_boot_info_array_pointer
+.align 8
+__ap_boot_info_array_pointer:
+    .skip 8
+
+.text
+.code64
+ap_long_mode:
     // The local APIC ID is in the RDI.
     mov rax, rdi
-    shl rax, 3
+    shl rax, 4                   // 16-byte `PerApRawInfo`
 
+    mov rbx, [rip + __ap_boot_info_array_pointer]
     // Setup the stack.
-    mov rbx, [__ap_boot_stack_array_pointer]
+    mov rsp, [rbx + rax - 16]    // raw_info[cpu_id - 1].stack_top
-    mov rsp, [rbx + rax]
+    // Setup the GS base (the CPU-local address).
-    xor rbp, rbp
+    mov rax, [rbx + rax - 8]     // raw_info[cpu_id - 1].cpu_local
+    mov rdx, rax
+    shr rdx, 32          // EDX:EAX = raw_info.cpu_local
+    mov ecx, 0xC0000101  // ECX = GS.base
+    wrmsr
 
     // Go to Rust code.
+.extern ap_early_entry
+    xor rbp, rbp
     mov rax, offset ap_early_entry
     call rax
 
 .extern halt # bsp_boot.S
     jmp halt
-
-.data
-// This is a pointer to be filled by the BSP when boot stacks
-// of all APs are allocated and initialized.
-.global __ap_boot_stack_array_pointer
-.align 8
-__ap_boot_stack_array_pointer:
-    .skip 8

ostd/src/arch/x86/boot/bsp_boot.S

@@ -309,7 +309,18 @@ long_mode:
     sub rcx, rdi
     rep stosb
 
-    // Call the corresponding Rust entrypoint according to the boot entrypoint
+    // Clear RBP to stop the backtrace.
+    xor rbp, rbp
+
+    // Initialize the GS base to the CPU-local start address.
+.extern __cpu_local_start
+    lea rax, [rip + __cpu_local_start]
+    mov rdx, rax
+    shr rdx, 32          // EDX:EAX = __cpu_local_start
+    mov ecx, 0xC0000101  // ECX = GS.base
+    wrmsr
+
+    // Call the corresponding Rust entrypoint according to the boot entrypoint.
     pop rax
     cmp rax, ENTRYTYPE_MULTIBOOT
     je entry_type_multiboot
@@ -329,8 +340,6 @@ long_mode:
 entry_type_linux:
     pop rdi // boot_params ptr
 
-    xor rbp, rbp
-
     lea  rax, [rip + __linux_boot]  // jump into Rust code
     call rax
     jmp halt
@@ -339,8 +348,6 @@ entry_type_multiboot:
     pop rsi // the address of multiboot info
     pop rdi // multiboot magic
 
-    xor rbp, rbp
-
     lea  rax, [rip + __multiboot_entry]  // jump into Rust code
     call rax
     jmp halt
@@ -349,8 +356,6 @@ entry_type_multiboot2:
     pop rsi // the address of multiboot info
     pop rdi // multiboot magic
 
-    xor rbp, rbp
-
     lea  rax, [rip + __multiboot2_entry]  // jump into Rust code
     call rax
     jmp halt

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

@@ -37,9 +37,12 @@ use crate::{
             Level, TriggerMode,
         },
     },
-    boot::memory_region::{MemoryRegion, MemoryRegionType},
+    boot::{
+        memory_region::{MemoryRegion, MemoryRegionType},
+        smp::PerApRawInfo,
+    },
     if_tdx_enabled,
-    mm::{paddr_to_vaddr, PAGE_SIZE},
+    mm::{Paddr, PAGE_SIZE},
 };
 
 /// Counts the number of processors.
@@ -105,15 +108,16 @@ pub(crate) fn count_processors() -> Option<u32> {
 /// # Safety
 ///
 /// The caller must ensure that
-/// 1. we're in the boot context of the BSP, and
+/// 1. we're in the boot context of the BSP,
-/// 2. all APs have not yet been booted.
+/// 2. all APs have not yet been booted, and
-pub(crate) unsafe fn bringup_all_aps(num_cpus: u32) {
+/// 3. the arguments are valid to boot APs.
+pub(crate) unsafe fn bringup_all_aps(info_ptr: *mut PerApRawInfo, pt_ptr: Paddr, num_cpus: u32) {
     // SAFETY: The code and data to boot AP is valid to write because
     // there are no readers and we are the only writer at this point.
     unsafe {
         copy_ap_boot_code();
-        fill_boot_stack_array_ptr();
+        fill_boot_info_ptr(info_ptr);
-        fill_boot_pt_ptr();
+        fill_boot_pt_ptr(pt_ptr);
     }
 
     // SAFETY: We've properly prepared all the resources to boot APs.
@@ -166,39 +170,30 @@ unsafe fn copy_ap_boot_code() {
 /// # Safety
 ///
 /// The caller must ensure the pointer to be filled is valid to write.
-unsafe fn fill_boot_stack_array_ptr() {
+unsafe fn fill_boot_info_ptr(info_ptr: *mut PerApRawInfo) {
-    let pages = &crate::boot::smp::AP_BOOT_INFO
-        .get()
-        .unwrap()
-        .boot_stack_array;
-    let vaddr = paddr_to_vaddr(pages.start_paddr());
-
     extern "C" {
-        static mut __ap_boot_stack_array_pointer: usize;
+        static mut __ap_boot_info_array_pointer: *mut PerApRawInfo;
     }
 
     // SAFETY: The safety is upheld by the caller.
     unsafe {
-        __ap_boot_stack_array_pointer = vaddr;
+        __ap_boot_info_array_pointer = info_ptr;
     }
 }
 
 /// # Safety
 ///
 /// The caller must ensure the pointer to be filled is valid to write.
-unsafe fn fill_boot_pt_ptr() {
+unsafe fn fill_boot_pt_ptr(pt_ptr: Paddr) {
     extern "C" {
         static mut __boot_page_table_pointer: u32;
     }
 
-    let boot_pt = crate::mm::page_table::boot_pt::with_borrow(|pt| pt.root_address())
+    let pt_ptr32 = pt_ptr.try_into().unwrap();
-        .unwrap()
-        .try_into()
-        .unwrap();
 
     // SAFETY: The safety is upheld by the caller.
     unsafe {
-        __boot_page_table_pointer = boot_pt;
+        __boot_page_table_pointer = pt_ptr32;
     }
 }
 

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

@@ -4,20 +4,6 @@
 
 use x86_64::registers::segmentation::{Segment64, GS};
 
-/// Sets the base address for the CPU local storage by writing to the GS base model-specific register.
-/// This operation is marked as `unsafe` because it directly interfaces with low-level CPU registers.
-///
-/// # Safety
-///
-///  - This function is safe to call provided that the GS register is dedicated entirely for CPU local storage
-///    and is not concurrently accessed for other purposes.
-///  - The caller must ensure that `addr` is a valid address and properly aligned, as required by the CPU.
-///  - This function should only be called in contexts where the CPU is in a state to accept such changes,
-///    such as during processor initialization.
-pub(crate) unsafe fn set_base(addr: u64) {
-    GS::write_base(x86_64::addr::VirtAddr::new(addr));
-}
-
 /// Gets the base address for the CPU local storage by reading the GS base model-specific register.
 pub(crate) fn get_base() -> u64 {
     GS::read_base().as_u64()

ostd/src/boot/smp.rs

@@ -2,14 +2,14 @@
 
 //! Symmetric multiprocessing (SMP) boot support.
 
-use alloc::collections::BTreeMap;
+use alloc::{boxed::Box, vec::Vec};
 use core::sync::atomic::{AtomicBool, Ordering};
 
 use spin::Once;
 
 use crate::{
     arch::boot::smp::bringup_all_aps,
-    cpu::{self, num_cpus},
+    cpu,
     mm::{
         frame::{meta::KernelMeta, Segment},
         paddr_to_vaddr, FrameAllocOptions, PAGE_SIZE,
@@ -17,15 +17,15 @@ use crate::{
     task::Task,
 };
 
-pub(crate) static AP_BOOT_INFO: Once<ApBootInfo> = Once::new();
+static AP_BOOT_INFO: Once<ApBootInfo> = Once::new();
 
 const AP_BOOT_STACK_SIZE: usize = PAGE_SIZE * 64;
 
-pub(crate) struct ApBootInfo {
+struct ApBootInfo {
-    /// It holds the boot stack top pointers used by all APs.
+    /// Raw boot information for each AP.
-    pub(crate) boot_stack_array: Segment<KernelMeta>,
+    per_ap_raw_info: Segment<KernelMeta>,
-    /// `per_ap_info` maps each AP's ID to its associated boot information.
+    /// Boot information for each AP.
-    per_ap_info: BTreeMap<u32, PerApInfo>,
+    per_ap_info: Box<[PerApInfo]>,
 }
 
 struct PerApInfo {
@@ -38,6 +38,18 @@ struct PerApInfo {
     boot_stack_pages: Segment<KernelMeta>,
 }
 
+/// Raw boot information for APs.
+///
+/// This is "raw" information that the assembly code (run by APs at startup,
+/// before ever entering the Rust entry point) will directly access. So the
+/// layout is important. **Update the assembly code if the layout is changed!**
+#[repr(C)]
+#[derive(Clone, Copy)]
+pub(crate) struct PerApRawInfo {
+    stack_top: *mut u8,
+    cpu_local: *mut u8,
+}
+
 static AP_LATE_ENTRY: Once<fn()> = Once::new();
 
 /// Boots all application processors.
@@ -50,57 +62,70 @@ static AP_LATE_ENTRY: Once<fn()> = Once::new();
 ///
 /// This function can only be called in the boot context of the BSP where APs have
 /// not yet been booted.
-pub fn boot_all_aps() {
+pub(crate) unsafe fn boot_all_aps() {
-    let num_cpus = num_cpus() as u32;
+    let num_cpus = crate::cpu::num_cpus();
+
     if num_cpus == 1 {
         return;
     }
-    log::info!("Booting {} processors.", num_cpus - 1);
+    log::info!("Booting {} processors", num_cpus - 1);
 
-    // We currently assumes that bootstrap processor (BSP) have always the
+    // We currently assume that
-    // processor ID 0. And the processor ID starts from 0 to `num_cpus - 1`.
+    // 1. the bootstrap processor (BSP) has the processor ID 0;
+    // 2. the processor ID starts from `0` to `num_cpus - 1`.
 
-    AP_BOOT_INFO.call_once(|| {
+    let mut per_ap_info = Vec::new();
-        let mut per_ap_info = BTreeMap::new();
+
-        // Use two pages to place stack pointers of all APs, thus support up to 1024 APs.
+    let per_ap_raw_info = FrameAllocOptions::new()
-        let boot_stack_array = FrameAllocOptions::new()
+        .zeroed(false)
+        .alloc_segment_with(
+            num_cpus
+                .saturating_sub(1)
+                .checked_mul(core::mem::size_of::<PerApRawInfo>())
+                .unwrap()
+                .div_ceil(PAGE_SIZE),
+            |_| KernelMeta,
+        )
+        .unwrap();
+    let raw_info_ptr = paddr_to_vaddr(per_ap_raw_info.start_paddr()) as *mut PerApRawInfo;
+
+    for ap in 1..num_cpus {
+        let boot_stack_pages = FrameAllocOptions::new()
             .zeroed(false)
-            .alloc_segment_with(2, |_| KernelMeta)
+            .alloc_segment_with(AP_BOOT_STACK_SIZE / PAGE_SIZE, |_| KernelMeta)
             .unwrap();
-        assert!(num_cpus < 1024);
 
-        for ap in 1..num_cpus {
+        let raw_info = PerApRawInfo {
-            let boot_stack_pages = FrameAllocOptions::new()
+            stack_top: paddr_to_vaddr(boot_stack_pages.end_paddr()) as *mut u8,
-                .zeroed(false)
+            cpu_local: paddr_to_vaddr(crate::cpu::local::get_ap(ap.try_into().unwrap())) as *mut u8,
-                .alloc_segment_with(AP_BOOT_STACK_SIZE / PAGE_SIZE, |_| KernelMeta)
+        };
-                .unwrap();
-            let boot_stack_ptr = paddr_to_vaddr(boot_stack_pages.end_paddr());
-            let stack_array_ptr = paddr_to_vaddr(boot_stack_array.start_paddr()) as *mut u64;
-            // SAFETY: The `stack_array_ptr` is valid and aligned.
-            unsafe {
-                stack_array_ptr
-                    .add(ap as usize)
-                    .write_volatile(boot_stack_ptr as u64);
-            }
-            per_ap_info.insert(
-                ap,
-                PerApInfo {
-                    is_started: AtomicBool::new(false),
-                    boot_stack_pages,
-                },
-            );
-        }
 
-        ApBootInfo {
+        // SAFETY: The index is in range because we allocated enough memory.
-            boot_stack_array,
+        let ptr = unsafe { raw_info_ptr.add(ap - 1) };
-            per_ap_info,
+        // SAFETY: The memory is valid for writing because it was just allocated.
-        }
+        unsafe { ptr.write(raw_info) };
+
+        per_ap_info.push(PerApInfo {
+            is_started: AtomicBool::new(false),
+            boot_stack_pages,
+        });
+    }
+
+    assert!(!AP_BOOT_INFO.is_completed());
+    AP_BOOT_INFO.call_once(move || ApBootInfo {
+        per_ap_raw_info,
+        per_ap_info: per_ap_info.into_boxed_slice(),
     });
 
     log::info!("Booting all application processors...");
 
-    // SAFETY: The safety is upheld by the caller.
+    let info_ptr = paddr_to_vaddr(AP_BOOT_INFO.get().unwrap().per_ap_raw_info.start_paddr())
-    unsafe { bringup_all_aps(num_cpus) };
+        as *mut PerApRawInfo;
+    let pt_ptr = crate::mm::page_table::boot_pt::with_borrow(|pt| pt.root_address()).unwrap();
+    // SAFETY: It's the right time to boot APs (guaranteed by the caller) and
+    // the arguments are valid to boot APs (generated above).
+    unsafe { bringup_all_aps(info_ptr, pt_ptr, num_cpus as u32) };
+
     wait_for_all_aps_started();
 
     log::info!("All application processors started. The BSP continues to run.");
@@ -121,7 +146,6 @@ fn ap_early_entry(local_apic_id: u32) -> ! {
     // SAFETY: we are on the AP and they are only called once with the correct
     // CPU ID.
     unsafe {
-        cpu::local::init_on_ap(local_apic_id);
         cpu::set_this_cpu_id(local_apic_id);
     }
 
@@ -145,10 +169,7 @@ fn ap_early_entry(local_apic_id: u32) -> ! {
 
     // Mark the AP as started.
     let ap_boot_info = AP_BOOT_INFO.get().unwrap();
-    ap_boot_info
+    ap_boot_info.per_ap_info[local_apic_id as usize - 1]
-        .per_ap_info
-        .get(&local_apic_id)
-        .unwrap()
         .is_started
         .store(true, Ordering::Release);
 
@@ -166,7 +187,7 @@ fn wait_for_all_aps_started() {
         let ap_boot_info = AP_BOOT_INFO.get().unwrap();
         ap_boot_info
             .per_ap_info
-            .values()
+            .iter()
             .all(|info| info.is_started.load(Ordering::Acquire))
     }
 

ostd/src/cpu/local/cell.rs

@@ -110,7 +110,7 @@ impl<T: 'static> CpuLocalCell<T> {
     ///   must be taken to ensure that the borrowing rules are correctly
     ///   enforced, since the interrupts may come asynchronously.
     pub fn as_mut_ptr(&'static self) -> *mut T {
-        super::has_init::assert_true();
+        super::is_used::debug_set_true();
 
         let offset = {
             let bsp_va = self as *const _ as usize;

ostd/src/cpu/local/cpu_local.rs

@@ -109,7 +109,7 @@ impl<T: 'static> CpuLocal<T> {
     ///
     /// The caller must ensure that the reference to `self` is static.
     pub(crate) unsafe fn as_ptr(&'static self) -> *const T {
-        super::has_init::assert_true();
+        super::is_used::debug_set_true();
 
         let offset = self.get_offset();
 
@@ -140,16 +140,23 @@ impl<T: 'static + Sync> CpuLocal<T> {
     ///
     /// Panics if the CPU ID is out of range.
     pub fn get_on_cpu(&'static self, cpu_id: CpuId) -> &'static T {
-        super::has_init::assert_true();
+        super::is_used::debug_set_true();
+
+        let cpu_id = cpu_id.as_usize();
+
         // If on the BSP, just use the statically linked storage.
-        if cpu_id.as_usize() == 0 {
+        if cpu_id == 0 {
             return &self.0;
         }
 
         // SAFETY: Here we use `Once::get_unchecked` to make getting the CPU-
         // local base faster. The storages must be initialized here (since this
         // is not the BSP) so it is safe to do so.
-        let base = unsafe { super::CPU_LOCAL_STORAGES.get_unchecked().get(cpu_id) };
+        let base = unsafe {
+            *super::CPU_LOCAL_STORAGES
+                .get_unchecked()
+                .get_unchecked(cpu_id - 1)
+        };
         let base = crate::mm::paddr_to_vaddr(base);
 
         let offset = self.get_offset();

ostd/src/cpu/local/mod.rs

@@ -42,10 +42,7 @@ pub use cpu_local::{CpuLocal, CpuLocalDerefGuard};
 use spin::Once;
 
 use super::CpuId;
-use crate::{
+use crate::mm::{frame::allocator, paddr_to_vaddr, Paddr, PAGE_SIZE};
-    arch,
-    mm::{frame::allocator, paddr_to_vaddr, Paddr, PAGE_SIZE},
-};
 
 // These symbols are provided by the linker script.
 extern "C" {
@@ -53,147 +50,128 @@ extern "C" {
     fn __cpu_local_end();
 }
 
-/// The BSP initializes the CPU-local areas for APs.
+/// The CPU-local areas for APs.
-static CPU_LOCAL_STORAGES: Once<CpuLocalStoragePointers> = Once::new();
+static CPU_LOCAL_STORAGES: Once<&'static [Paddr]> = Once::new();
 
-struct CpuLocalStoragePointers(&'static mut [Paddr]);
+/// Copies the CPU-local data on the bootstrap processor (BSP)
+/// for application processors (APs).
+///
+/// # Safety
+///
+/// This function must be called in the boot context of the BSP, at a time
+/// when the APs have not yet booted.
+///
+/// The CPU-local data on the BSP must not be used before calling this
+/// function to copy it for the APs. Otherwise, the copied data will
+/// contain non-constant (also non-`Copy`) data, resulting in undefined
+/// behavior when it's loaded on the APs.
+pub(crate) unsafe fn copy_bsp_for_ap() {
+    let num_aps = super::num_cpus() - 1; // BSP does not need allocated storage.
+    if num_aps == 0 {
+        return;
+    }
 
-impl CpuLocalStoragePointers {
+    // Allocate a region to store the pointers to the CPU-local storage segments.
-    /// Allocates frames for storing CPU-local data for each AP.
+    let res = {
-    ///
+        let size = core::mem::size_of::<Paddr>()
-    /// # Safety
+            .checked_mul(num_aps)
-    ///
+            .unwrap()
-    /// The caller must ensure that the `num_cpus` matches the number of all
+            .align_up(PAGE_SIZE);
-    /// CPUs that will access CPU local storage.
-    ///
-    /// The BSP's CPU-local storage should not be touched before calling
-    /// this method, since the CPU-local storage for APs is copied from the
-    /// BSP's CPU-local storage.
-    unsafe fn allocate(num_cpus: usize) -> Self {
-        let num_aps = num_cpus - 1; // BSP does not need allocated storage.
-        assert!(num_aps > 0, "No APs to allocate CPU-local storage");
-
-        // Allocate a region to store the pointers to the CPU-local storage segments.
-        let size = (core::mem::size_of::<Paddr>() * num_aps).align_up(PAGE_SIZE);
         let addr =
             allocator::early_alloc(Layout::from_size_align(size, PAGE_SIZE).unwrap()).unwrap();
         let ptr = paddr_to_vaddr(addr) as *mut Paddr;
-        // SAFETY: The memory is allocated and the pointer is valid.
+
+        // SAFETY: The memory is properly allocated. We exclusively own it. So it's valid to write.
         unsafe {
             core::ptr::write_bytes(ptr as *mut u8, 0, size);
         }
-        // SAFETY: The memory would not be deallocated. And it is valid.
+        // SAFETY: The memory is properly allocated and initialized. We exclusively own it. We
-        let res = Self(unsafe { core::slice::from_raw_parts_mut(ptr, num_aps) });
+        // never deallocate it so it lives for '`static'. So we can create a mutable slice on it.
+        unsafe { core::slice::from_raw_parts_mut(ptr, num_aps) }
+    };
 
-        // Allocate the CPU-local storage segments for APs.
+    let bsp_base_va = __cpu_local_start as usize;
-        let bsp_base_va = __cpu_local_start as usize;
+    let bsp_end_va = __cpu_local_end as usize;
-        let bsp_end_va = __cpu_local_end as usize;
-        for id in 0..num_aps {
-            let ap_pages = {
-                let nbytes = (bsp_end_va - bsp_base_va).align_up(PAGE_SIZE);
-                allocator::early_alloc(Layout::from_size_align(nbytes, PAGE_SIZE).unwrap()).unwrap()
-            };
-            let ap_pages_ptr = paddr_to_vaddr(ap_pages) as *mut u8;
 
-            // SAFETY: The BSP has not initialized the CPU-local area, so the objects in
+    // Allocate the CPU-local storage segments for APs.
-            // in the `.cpu_local` section can be bitwise bulk copied to the AP's local
+    for res_addr_mut in res.iter_mut() {
-            // storage. The destination memory is allocated so it is valid to write to.
+        let nbytes = (bsp_end_va - bsp_base_va).align_up(PAGE_SIZE);
-            unsafe {
+        let ap_pages =
-                core::ptr::copy_nonoverlapping(
+            allocator::early_alloc(Layout::from_size_align(nbytes, PAGE_SIZE).unwrap()).unwrap();
-                    bsp_base_va as *const u8,
+        let ap_pages_ptr = paddr_to_vaddr(ap_pages) as *mut u8;
-                    ap_pages_ptr,
-                    bsp_end_va - bsp_base_va,
-                );
-            }
 
-            res.0[id] = ap_pages;
+        // SAFETY:
+        // 1. The source is valid to read because it has not been used before,
+        //    so it contains only constants.
+        // 2. The destination is valid to write because it is just allocated.
+        // 3. The memory is aligned because the alignment of `u8` is 1.
+        // 4. The two memory regions do not overlap because allocated memory
+        //    regions never overlap with the kernel data.
+        unsafe {
+            core::ptr::copy_nonoverlapping(bsp_base_va as *const u8, ap_pages_ptr, nbytes);
         }
 
-        res
+        *res_addr_mut = ap_pages;
     }
 
-    fn get(&self, cpu_id: CpuId) -> Paddr {
+    is_used::debug_assert_false();
-        let offset = cpu_id
-            .as_usize()
-            .checked_sub(1)
-            .expect("The BSP does not have allocated CPU-local storage");
 
-        let paddr = self.0[offset];
+    assert!(!CPU_LOCAL_STORAGES.is_completed());
-        assert!(
+    CPU_LOCAL_STORAGES.call_once(|| res);
-            paddr != 0,
-            "The CPU-local storage for CPU {} is not allocated",
-            cpu_id.as_usize()
-        );
-        paddr
-    }
 }
 
-/// Initializes the CPU local data for the bootstrap processor (BSP).
+/// Gets the pointer to the CPU-local storage for the given AP.
 ///
-/// # Safety
+/// # Panics
 ///
-/// This function can only called on the BSP, for once.
+/// This method will panic if the `cpu_id` does not represent an AP or the AP's CPU-local storage
-///
+/// has not been allocated.
-/// It must be guaranteed that the BSP will not access local data before
+pub(crate) fn get_ap(cpu_id: CpuId) -> Paddr {
-/// this function being called, otherwise copying non-constant values
+    let offset = cpu_id
-/// will result in pretty bad undefined behavior.
+        .as_usize()
-pub unsafe fn init_on_bsp() {
+        .checked_sub(1)
-    let num_cpus = super::num_cpus();
+        .expect("The BSP does not have allocated CPU-local storage");
 
-    if num_cpus > 1 {
+    let paddr = CPU_LOCAL_STORAGES
-        // SAFETY: The number of CPUs is correct. And other conditions are
-        // the caller's safety conditions.
-        let cpu_local_storages = unsafe { CpuLocalStoragePointers::allocate(num_cpus) };
-
-        CPU_LOCAL_STORAGES.call_once(|| cpu_local_storages);
-    }
-
-    arch::cpu::local::set_base(__cpu_local_start as usize as u64);
-
-    has_init::set_true();
-}
-
-/// Initializes the CPU local data for the application processor (AP).
-///
-/// # Safety
-///
-/// This function can only called on the AP.
-pub unsafe fn init_on_ap(cpu_id: u32) {
-    let ap_pages = CPU_LOCAL_STORAGES
         .get()
-        .unwrap()
+        .expect("No CPU-local storage has been allocated")[offset];
-        .get(CpuId::try_from(cpu_id as usize).unwrap());
+    assert_ne!(
-
+        paddr,
-    let ap_pages_ptr = paddr_to_vaddr(ap_pages) as *mut u32;
+        0,
-
+        "The CPU-local storage for CPU {} is not allocated",
-    // SAFETY: the memory will be dedicated to the AP. And we are on the AP.
+        cpu_id.as_usize(),
-    unsafe {
+    );
-        arch::cpu::local::set_base(ap_pages_ptr as u64);
+    paddr
-    }
 }
 
-mod has_init {
+mod is_used {
-    //! This module is used to detect the programming error of using the CPU-local
+    //! This module tracks whether any CPU-local variables are used.
-    //! mechanism before it is initialized. Such bugs have been found before and we
+    //!
-    //! do not want to repeat this error again. This module is only incurs runtime
+    //! [`copy_bsp_for_ap`] copies the CPU local data from the BSP
-    //! overhead if debug assertions are enabled.
+    //! to the APs, so it requires as a safety condition that the
+    //! CPU-local data has not been accessed before the copy. This
+    //! module provides utilities to check if the safety condition
+    //! is met, but only if debug assertions are enabled.
+    //!
+    //! [`copy_bsp_for_ap`]: super::copy_bsp_for_ap
+
     cfg_if::cfg_if! {
         if #[cfg(debug_assertions)] {
             use core::sync::atomic::{AtomicBool, Ordering};
 
-            static IS_INITIALIZED: AtomicBool = AtomicBool::new(false);
+            static IS_USED: AtomicBool = AtomicBool::new(false);
 
-            pub fn assert_true() {
+            pub fn debug_set_true() {
-                debug_assert!(IS_INITIALIZED.load(Ordering::Relaxed));
+                IS_USED.store(true, Ordering::Relaxed);
             }
 
-            pub fn set_true() {
+            pub fn debug_assert_false() {
-                IS_INITIALIZED.store(true, Ordering::Relaxed);
+                debug_assert!(!IS_USED.load(Ordering::Relaxed));
             }
         } else {
-            pub fn assert_true() {}
+            pub fn debug_set_true() {}
 
-            pub fn set_true() {}
+            pub fn debug_assert_false() {}
         }
     }
 }

ostd/src/lib.rs

@@ -84,11 +84,13 @@ unsafe fn init() {
 
     logger::init();
 
-    // SAFETY: They are only called once on BSP and ACPI has been initialized.
+    // SAFETY:
-    // No CPU local objects have been accessed by this far.
+    // 1. They are only called once in the boot context of the BSP.
+    // 2. The number of CPUs are available because ACPI has been initialized.
+    // 3. No CPU-local objects have been accessed yet.
     unsafe {
         cpu::init_num_cpus();
-        cpu::local::init_on_bsp();
+        cpu::local::copy_bsp_for_ap();
         cpu::set_this_cpu_id(0);
     }