ax_percpu/
imp.rs

1use core::sync::atomic::{AtomicBool, Ordering};
2
3use ax_percpu_macros::percpu_symbol_vma;
4
5static IS_INIT: AtomicBool = AtomicBool::new(false);
6
7const fn align_up_64(val: usize) -> usize {
8    const SIZE_64BIT: usize = 0x40;
9    (val + SIZE_64BIT - 1) & !(SIZE_64BIT - 1)
10}
11
12#[cfg(feature = "host-test")]
13static PERCPU_AREA_BASE: spin::once::Once<usize> = spin::once::Once::new();
14
15unsafe extern "C" {
16    fn _percpu_start();
17    fn _percpu_end();
18    // WARNING: `_percpu_load_start`/`_percpu_load_end` (i.e. symbols in the
19    // `.percpu` section) must be used with `percpu_symbol_vma!` macro to get
20    // their VMA addresses. Casting them directly to `usize` may lead to
21    // unexpected results, including:
22    // - Rust assuming they are valid pointers and optimizing code based on that
23    //   assumption (they are non-zero), causing unexpected runtime errors;
24    // - Link-time errors because they are too far away from the program counter
25    //.  (when Rust uses PC-relative addressing).
26    // See https://github.com/arceos-org/percpu/issues/18 for more details.
27    fn _percpu_load_start();
28    fn _percpu_load_end();
29}
30
31/// Returns the number of per-CPU data areas reserved.
32pub fn percpu_area_num() -> usize {
33    (_percpu_end as *const () as usize - _percpu_start as *const () as usize)
34        / align_up_64(percpu_area_size())
35}
36
37/// Returns the per-CPU data area size for one CPU.
38pub fn percpu_area_size() -> usize {
39    percpu_symbol_vma!(_percpu_load_end) - percpu_symbol_vma!(_percpu_load_start)
40}
41
42/// Returns the base address of the per-CPU data area on the given CPU.
43///
44/// if `cpu_id` is 0, it returns the base address of all per-CPU data areas.
45pub fn percpu_area_base(cpu_id: usize) -> usize {
46    cfg_if::cfg_if! {
47        if #[cfg(feature = "host-test")] {
48            let base = *PERCPU_AREA_BASE.get().unwrap();
49        } else {
50            let base = _percpu_start as *const () as usize;
51        }
52    }
53    base + cpu_id * align_up_64(percpu_area_size())
54}
55
56/// Initialize all per-CPU data areas.
57///
58/// The number of areas is determined by the following formula:
59///
60/// ```text
61/// (percpu_section_size / align_up(percpu_area_size, 64)
62/// ```
63///
64/// Returns the number of areas initialized. If this function has been called
65/// before, it does nothing and returns 0.
66pub fn init() -> usize {
67    // avoid re-initialization.
68    if IS_INIT
69        .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
70        .is_err()
71    {
72        return 0;
73    }
74
75    #[cfg(not(feature = "non-zero-vma"))]
76    {
77        // `_percpu_load_start as *const () as usize` cannot be used here because
78        // rust will assume a `*const ()` is a valid pointer and will not be 0,
79        // causing unexpected `0 != 0` assertion failure.
80        assert_eq!(
81            percpu_symbol_vma!(_percpu_load_start),
82            0,
83            "The `.percpu` section must be loaded at VMA address 0 when feature \"non-zero-vma\" \
84             is disabled"
85        )
86    }
87
88    #[cfg(feature = "host-test")]
89    {
90        // we not load the ax-percpu section in ELF, allocate them here.
91        let total_size = _percpu_end as *const () as usize - _percpu_start as *const () as usize;
92        let layout = std::alloc::Layout::from_size_align(total_size, 0x1000).unwrap();
93        PERCPU_AREA_BASE.call_once(|| unsafe { std::alloc::alloc(layout) as usize });
94    }
95
96    let base = percpu_area_base(0);
97    let size = percpu_area_size();
98    let num = percpu_area_num();
99    for i in 1..num {
100        let secondary_base = percpu_area_base(i);
101        #[cfg(not(feature = "host-test"))]
102        assert!(secondary_base + size <= _percpu_end as *const () as usize);
103        // copy per-cpu data of the primary CPU to other CPUs.
104        unsafe {
105            core::ptr::copy_nonoverlapping(base as *const u8, secondary_base as *mut u8, size);
106        }
107    }
108    num
109}
110
111/// Reads the architecture-specific per-CPU data register.
112///
113/// This register is used to hold the per-CPU data base on each CPU.
114pub fn read_percpu_reg() -> usize {
115    let tp: usize;
116    unsafe {
117        cfg_if::cfg_if! {
118            if #[cfg(target_arch = "x86_64")] {
119                tp = if cfg!(feature = "host-test") {
120                    SELF_PTR.read_current_raw()
121                } else {
122                    x86::msr::rdmsr(x86::msr::IA32_GS_BASE) as usize
123                };
124            } else if #[cfg(any(target_arch = "riscv32", target_arch = "riscv64"))] {
125                core::arch::asm!("mv {}, gp", out(reg) tp)
126            } else if #[cfg(all(target_arch = "aarch64", not(feature = "arm-el2")))] {
127                core::arch::asm!("mrs {}, TPIDR_EL1", out(reg) tp)
128            } else if #[cfg(all(target_arch = "aarch64", feature = "arm-el2"))] {
129                core::arch::asm!("mrs {}, TPIDR_EL2", out(reg) tp)
130            } else if #[cfg(target_arch = "loongarch64")] {
131                // Register Convention
132                // https://docs.kernel.org/arch/loongarch/introduction.html#gprs
133                core::arch::asm!("move {}, $r21", out(reg) tp)
134            } else if #[cfg(target_arch = "arm")] {
135                core::arch::asm!("mrc p15, 0, {}, c13, c0, 3", out(reg) tp)
136            }
137        }
138    }
139    cfg_if::cfg_if! {
140        if #[cfg(feature = "non-zero-vma")] {
141            tp + percpu_symbol_vma!(_percpu_load_start)
142        } else {
143            tp
144        }
145    }
146}
147
148/// Writes the architecture-specific per-CPU data register.
149///
150/// This register is used to hold the per-CPU data base on each CPU.
151///
152/// # Safety
153///
154/// This function is unsafe because it writes the low-level register directly.
155pub unsafe fn write_percpu_reg(tp: usize) {
156    cfg_if::cfg_if! {
157        if #[cfg(feature = "non-zero-vma")] {
158            let tp = tp - percpu_symbol_vma!(_percpu_load_start);
159        }
160    };
161
162    unsafe {
163        cfg_if::cfg_if! {
164            if #[cfg(target_arch = "x86_64")] {
165                if cfg!(feature = "host-test") {
166                    const ARCH_SET_GS: u32 = 0x1001;
167                    const SYS_ARCH_PRCTL: u32 = 158;
168                    core::arch::asm!(
169                        "syscall",
170                        in("eax") SYS_ARCH_PRCTL,
171                        in("edi") ARCH_SET_GS,
172                        in("rsi") tp,
173                    );
174                } else {
175                    x86::msr::wrmsr(x86::msr::IA32_GS_BASE, tp as u64);
176                }
177                SELF_PTR.write_current_raw(tp);
178            } else if #[cfg(any(target_arch = "riscv32", target_arch = "riscv64"))] {
179                core::arch::asm!("mv gp, {}", in(reg) tp)
180            } else if #[cfg(all(target_arch = "aarch64", not(feature = "arm-el2")))] {
181                core::arch::asm!("msr TPIDR_EL1, {}", in(reg) tp)
182            } else if #[cfg(all(target_arch = "aarch64", feature = "arm-el2"))] {
183                core::arch::asm!("msr TPIDR_EL2, {}", in(reg) tp)
184            } else if #[cfg(target_arch = "loongarch64")] {
185                core::arch::asm!("move $r21, {}", in(reg) tp)
186            } else if #[cfg(target_arch = "arm")] {
187                core::arch::asm!("mcr p15, 0, {}, c13, c0, 3", in(reg) tp)
188            }
189        }
190    }
191}
192
193/// Initializes the per-CPU data register.
194///
195/// It is equivalent to `write_percpu_reg(percpu_area_base(cpu_id))`, which set
196/// the architecture-specific per-CPU data register to the base address of the
197/// corresponding per-CPU data area.
198///
199/// `cpu_id` indicates which per-CPU data area to use.
200pub fn init_percpu_reg(cpu_id: usize) {
201    let tp = percpu_area_base(cpu_id);
202    unsafe { write_percpu_reg(tp) }
203}
204
205/// To use `ax_percpu::__priv::NoPreemptGuard::new()` and `ax_percpu::percpu_area_base()` in macro expansion.
206#[allow(unused_imports)]
207use crate as ax_percpu;
208
209/// On x86, we use `gs:SELF_PTR` to store the address of the per-CPU data area base.
210#[cfg(target_arch = "x86_64")]
211#[unsafe(no_mangle)]
212#[ax_percpu_macros::def_percpu]
213static SELF_PTR: usize = 0;
ax_percpu/imp.rs

ax_percpu/
imp.rs
