Detect number of CPUs on init

Cargo.toml

@@ -8,6 +8,7 @@ members = [
 
 [workspace.package]
 version = "0.1.6"
+edition = "2021"
 authors = ["Yuekai Jia <equation618@gmail.com>", "hky1999 <976929993@qq.com>", "Su Mingxian <aarkegz@gmail.com>"]
 license = "GPL-3.0-or-later OR Apache-2.0 OR MulanPSL-2.0"
 homepage = "https://github.com/arceos-org/arceos"

percpu/Cargo.toml

@@ -1,9 +1,9 @@
 [package]
 name = "percpu"
-edition = "2021"
 description = "Define and access per-CPU data structures"
 documentation = "https://docs.rs/percpu"
 version.workspace = true
+edition.workspace = true
 authors.workspace = true
 license.workspace = true
 homepage.workspace = true

percpu/src/imp.rs

@@ -10,12 +10,20 @@ const fn align_up_64(val: usize) -> usize {
 #[cfg(not(target_os = "none"))]
 static PERCPU_AREA_BASE: spin::once::Once<usize> = spin::once::Once::new();
 
+unsafe extern "C" {
+    fn _percpu_start();
+    fn _percpu_end();
+    fn _percpu_load_start();
+    fn _percpu_load_end();
+}
+
+/// Returns the total size of per-CPU data area for all CPUs.
+fn percpu_area_total_size() -> usize {
+    _percpu_end as usize - _percpu_start as usize
+}
+
 /// Returns the per-CPU data area size for one CPU.
 pub fn percpu_area_size() -> usize {
-    extern "C" {
-        fn _percpu_load_start();
-        fn _percpu_load_end();
-    }
     // It seems that `_percpu_load_start as usize - _percpu_load_end as usize` will result in more instructions.
     use percpu_macros::percpu_symbol_offset;
     percpu_symbol_offset!(_percpu_load_end) - percpu_symbol_offset!(_percpu_load_start)
@@ -27,9 +35,6 @@ pub fn percpu_area_size() -> usize {
 pub fn percpu_area_base(cpu_id: usize) -> usize {
     cfg_if::cfg_if! {
         if #[cfg(target_os = "none")] {
-            extern "C" {
-                fn _percpu_start();
-            }
             let base = _percpu_start as usize;
         } else {
             let base = *PERCPU_AREA_BASE.get().unwrap();
@@ -39,21 +44,31 @@ pub fn percpu_area_base(cpu_id: usize) -> usize {
 }
 
 /// Initialize the per-CPU data area for `max_cpu_num` CPUs.
-pub fn init(max_cpu_num: usize) {
+///
+/// If `max_cpu_num` is `None`, it use the following formula to calculate the
+/// number of CPUs:
+///
+/// ```text
+/// max_cpu_num = (percpu_area_total_size / align_up(percpu_area_size, 64)
+/// ```
+///
+/// Returns the number of areas initialized.
+pub fn init(max_cpu_num: Option<usize>) -> usize {
     // avoid re-initialization.
     if IS_INIT
         .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
         .is_err()
     {
-        return;
+        return 0;
     }
 
     let size = percpu_area_size();
+    let total_size = percpu_area_total_size();
+    let max_cpu_num = max_cpu_num.unwrap_or(total_size / align_up_64(size));
 
     #[cfg(target_os = "linux")]
     {
         // we not load the percpu section in ELF, allocate them here.
-        let total_size = align_up_64(size) * max_cpu_num;
         let layout = std::alloc::Layout::from_size_align(total_size, 0x1000).unwrap();
         PERCPU_AREA_BASE.call_once(|| unsafe { std::alloc::alloc(layout) as usize });
     }
@@ -62,17 +77,13 @@ pub fn init(max_cpu_num: usize) {
     for i in 1..max_cpu_num {
         let secondary_base = percpu_area_base(i);
         #[cfg(target_os = "none")]
-        {
-            extern "C" {
-                fn _percpu_end();
-            }
-            assert!(secondary_base + size <= _percpu_end as usize);
-        }
+        assert!(secondary_base + size <= _percpu_end as usize);
         // copy per-cpu data of the primary CPU to other CPUs.
         unsafe {
             core::ptr::copy_nonoverlapping(base as *const u8, secondary_base as *mut u8, size);
         }
     }
+    max_cpu_num
 }
 
 /// Reads the architecture-specific per-CPU data register.

percpu/src/naive.rs

@@ -38,5 +38,9 @@ pub fn init_percpu_reg(_cpu_id: usize) {}
 
 /// Initialize the per-CPU data area for `max_cpu_num` CPUs.
 ///
-/// No effect for "sp-naive" use.
-pub fn init(_max_cpu_num: usize) {}
+/// Returns the number of areas initialized.
+///
+/// No effect for "sp-naive" use, just returns `1`.
+pub fn init(_max_cpu_num: Option<usize>) -> usize {
+    1
+}

percpu/tests/test_percpu.rs

@@ -40,7 +40,7 @@ fn test_percpu() {
 
     #[cfg(not(feature = "sp-naive"))]
     let base = {
-        init(4);
+        assert_eq!(init(None), 4);
         unsafe { write_percpu_reg(percpu_area_base(0)) };
 
         let base = read_percpu_reg();

percpu_macros/Cargo.toml

@@ -1,9 +1,9 @@
 [package]
 name = "percpu_macros"
-edition = "2021"
 description = "Macros to define and access a per-CPU data structure"
 documentation = "https://docs.rs/percpu_macros"
 version.workspace = true
+edition.workspace = true
 authors.workspace = true
 license.workspace = true
 homepage.workspace = true