kernel: add cpu_local proc macro

Its like the `thread_local` macro; makes the code much cleaner :^)

Signed-off-by: Anhad Singh <[email protected]>

Author: Andy-Python-Programmer

src/.cargo/kernel.ld

@@ -30,6 +30,14 @@ SECTIONS
         *(.rodata .rodata.*)
     } :rodata
 
+    .cpu_local : {
+        __cpu_local_start = .;
+        KEEP(*(.cpu_local_self_ptr));
+        KEEP(*(.cpu_local_tss));
+        KEEP(*(.cpu_local));
+        __cpu_local_end = .;
+    }
+
     /* Move to the next memory page for .data */
     . += CONSTANT(MAXPAGESIZE);
 

src/aero_kernel/src/arch/x86_64/apic.rs

@@ -18,8 +18,8 @@
 use core::ptr;
 use core::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
 
+use crate::arch::interrupts;
 use crate::arch::interrupts::InterruptStack;
-use crate::arch::{interrupts, tls};
 use crate::mem::paging::{PhysAddr, VirtAddr};
 use raw_cpuid::{CpuId, FeatureInfo};
 use spin::Once;
@@ -108,6 +108,9 @@ fn lapic_error_handler(_stack: &mut InterruptStack) {
     log::error!("ESR={:#0x}", self::get_local_apic().get_esr());
 }
 
+#[cpu_local]
+static mut LAPIC_TIMER_FREQUENCY: u32 = 0;
+
 pub struct LocalApic {
     address: VirtAddr,
     apic_type: ApicType,
@@ -194,7 +197,7 @@ impl LocalApic {
     pub fn timer_oneshot(&mut self, vec: u8, us: usize) {
         self.timer_stop();
 
-        let lapic_timer_frequency = tls::get_percpu().lapic_timer_frequency;
+        let lapic_timer_frequency = unsafe { *LAPIC_TIMER_FREQUENCY };
         let ticks = us * (lapic_timer_frequency / 1000000) as usize;
 
         unsafe {
@@ -225,7 +228,7 @@ impl LocalApic {
             let pit_ticks = initial_pit_tick - final_pit_tick;
             let timer_frequency = (SAMPLES / pit_ticks as u32) * time::PIT_DIVIDEND as u32;
 
-            tls::get_percpu().lapic_timer_frequency = timer_frequency;
+            *LAPIC_TIMER_FREQUENCY = timer_frequency;
         }
 
         self.timer_stop();

src/aero_kernel/src/arch/x86_64/cpu_local.rs

@@ -0,0 +1,98 @@
+// Copyright (C) 2021-2023 The Aero Project Developers.
+//
+// This file is part of The Aero Project.
+//
+// Aero is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// Aero is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with Aero. If not, see <https://www.gnu.org/licenses/>.
+
+use core::alloc::Layout;
+use core::ops::{Deref, DerefMut};
+
+use crate::mem::paging::VirtAddr;
+use crate::utils::LinkerSymbol;
+
+use super::io;
+
+extern "C" {
+    static __cpu_local_start: LinkerSymbol<u8>;
+    static __cpu_local_end: LinkerSymbol<u8>;
+}
+
+#[repr(C)]
+pub struct CpuLocal<T>(T);
+
+impl<T> CpuLocal<T> {
+    pub const fn new(val: T) -> Self {
+        Self(val)
+    }
+
+    pub fn addr(&self) -> VirtAddr {
+        let val: u64;
+
+        unsafe {
+            // gs:[0] -> SELF_PTR
+            asm!(
+                "mov {}, qword ptr gs:[0]",
+                lateout(reg) val,
+                options(nostack, preserves_flags, pure, readonly),
+            );
+        }
+
+        let self_addr = VirtAddr::new(self as *const _ as u64);
+        let section_addr = VirtAddr::new(unsafe { &__cpu_local_start as *const _ as u64 });
+
+        let offset = self_addr - section_addr;
+        VirtAddr::new(val) + offset
+    }
+}
+
+impl<T> Deref for CpuLocal<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        unsafe { &*self.addr().as_ptr() }
+    }
+}
+
+impl<T> DerefMut for CpuLocal<T> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        unsafe { &mut *self.addr().as_mut_ptr() }
+    }
+}
+
+/// The GS register holds a pointer to CPU-local data at a fixed offset of 0. While this approach
+/// requires an additional memory lookup for accessing the data, it enables better code optimization
+/// since the subsequent access is just a normal memory access. Considering the size of the
+/// CPU-local data, this optimization is beneficial.
+#[cpu_local(subsection = "self_ptr")]
+static SELF_PTR: u64 = 0;
+
+#[cpu_local]
+static mut CPUID: usize = 0;
+
+pub fn init(cpu_id: usize) {
+    unsafe {
+        let start = VirtAddr::new(&__cpu_local_start as *const _ as u64);
+        let end = VirtAddr::new(&__cpu_local_end as *const _ as u64);
+        let size = end - start;
+
+        let layout = Layout::from_size_align_unchecked(size as _, 64);
+        let data = alloc::alloc::alloc_zeroed(layout);
+
+        core::ptr::copy_nonoverlapping::<u8>(start.as_ptr(), data, size as usize);
+        *data.cast::<u64>() = data as u64;
+
+        io::wrmsr(io::IA32_GS_BASE, data as u64);
+        *CPUID = cpu_id;
+    }
+}

src/aero_kernel/src/arch/x86_64/gdt.rs

@@ -29,10 +29,6 @@ use core::mem;
 
 use alloc::alloc::alloc_zeroed;
 
-use crate::arch::tls::PerCpuData;
-
-use super::{io, tls};
-
 bitflags::bitflags! {
     /// Specifies which element to load into a segment from
     /// descriptor tables (i.e., is a index to LDT or GDT table
@@ -283,12 +279,16 @@ pub struct Tss {
     pub iomap_base: u16, // offset 0x66
 }
 
-// Processor Control Region
-#[repr(C)]
-pub struct Kpcr {
-    pub tss: Tss,
-    pub cpu_local: PerCpuData,
-}
+#[cpu_local(subsection = "tss")]
+pub static mut TSS: Tss = Tss {
+    reserved: 0,
+    rsp: [0; 3],
+    reserved2: 0,
+    ist: [0; 7],
+    reserved3: 0,
+    reserved4: 0,
+    iomap_base: 0,
+};
 
 /// Initialize the bootstrap GDT which is required to initialize TLS (Thread Local Storage)
 /// support so, after the kernel heap we will map the TLS section and initialize the *actual* GDT
@@ -315,14 +315,6 @@ pub fn init_boot() {
     }
 }
 
-pub fn get_task_state_segment() -> &'static mut Tss {
-    &mut get_kpcr().tss
-}
-
-pub fn get_kpcr() -> &'static mut Kpcr {
-    unsafe { &mut *(io::rdmsr(io::IA32_GS_BASE) as *mut Kpcr) }
-}
-
 static STK: [u8; 4096 * 16] = [0; 4096 * 16];
 
 /// Initialize the *actual* GDT stored in TLS.
@@ -345,14 +337,13 @@ pub fn init() {
     gdt.copy_from_slice(&GDT);
 
     unsafe {
-        let tss_ref = get_task_state_segment();
-        let tss_ptr = tss_ref as *mut Tss;
+        let tss_ptr = TSS.addr().as_mut_ptr::<Tss>();
 
         gdt[GdtEntryType::TSS as usize].set_offset(tss_ptr as u32);
         gdt[GdtEntryType::TSS as usize].set_limit(mem::size_of::<Tss>() as u32);
         gdt[GdtEntryType::TSS_HI as usize].set_raw((tss_ptr as u64) >> 32);
 
-        tss_ref.rsp[0] = STK.as_ptr().offset(4096 * 16) as u64;
+        TSS.rsp[0] = STK.as_ptr().offset(4096 * 16) as u64;
 
         let gdt_descriptor = GdtDescriptor::new(
             (mem::size_of::<[GdtEntry; GDT_ENTRY_COUNT]>() - 1) as u16,
@@ -371,9 +362,9 @@ pub fn init() {
         load_tss(SegmentSelector::new(GdtEntryType::TSS, Ring::Ring0));
     }
 
-    // Now we update the per-cpu storage to store a reference
-    // to the per-cpu GDT.
-    tls::get_percpu().gdt = gdt;
+    // // Now we update the per-cpu storage to store a reference
+    // // to the per-cpu GDT.
+    // tls::get_percpu().gdt = gdt;
 }
 
 #[inline(always)]

src/aero_kernel/src/arch/x86_64/interrupts/exceptions.rs

@@ -18,12 +18,14 @@
 use super::{io, InterruptErrorStack};
 
 use crate::arch::controlregs;
-use crate::arch::tls::get_percpu;
-use crate::mem::paging::PageFaultErrorCode;
+use crate::mem::paging::{PageFaultErrorCode, VirtAddr};
 
 use crate::unwind;
 use crate::userland::scheduler;
 
+#[cpu_local]
+pub static mut PF_RESUME: VirtAddr = VirtAddr::new(0);
+
 const LOG_PF_PTABLE: bool = true;
 
 macro interrupt_exception(fn $name:ident() => $message:expr) {
@@ -121,9 +123,9 @@ pub fn breakpoint(stack: &mut InterruptErrorStack) {
 }
 
 pub(super) fn page_fault(stack: &mut InterruptErrorStack) {
-    let pf_resume = get_percpu().pf_resume.as_ptr::<u8>();
-    if !pf_resume.is_null() {
-        stack.stack.iret.rip = pf_resume as u64;
+    let pf_resume = unsafe { *PF_RESUME };
+    if !pf_resume.is_zero() {
+        stack.stack.iret.rip = pf_resume.as_u64();
         return;
     }
 

src/aero_kernel/src/arch/x86_64/interrupts/mod.rs

@@ -15,7 +15,7 @@
 // You should have received a copy of the GNU General Public License
 // along with Aero. If not, see <https://www.gnu.org/licenses/>.
 
-mod exceptions;
+pub mod exceptions;
 mod idt;
 
 use core::sync::atomic::{AtomicUsize, Ordering};

src/aero_kernel/src/arch/x86_64/mod.rs

@@ -15,6 +15,8 @@
 // You should have received a copy of the GNU General Public License
 // along with Aero. If not, see <https://www.gnu.org/licenses/>.
 
+pub mod cpu_local;
+
 pub mod apic;
 pub mod controlregs;
 pub mod gdt;
@@ -181,7 +183,8 @@ extern "C" fn arch_aero_main() -> ! {
     acpi::init(rsdp);
     log::info!("loaded ACPI");
 
-    tls::init(0);
+    tls::init();
+    cpu_local::init(0);
     log::info!("loaded TLS");
 
     crate::unwind::set_panic_hook_ready(true);
@@ -209,7 +212,8 @@ extern "C" fn x86_64_aero_ap_main(boot_info: *const SmpInfo) -> ! {
     gdt::init_boot();
     log::info!("AP{}: loaded boot GDT", ap_id);
 
-    tls::init(ap_id);
+    tls::init();
+    cpu_local::init(ap_id);
     log::info!("AP{}: loaded TLS", ap_id);
 
     gdt::init();

src/aero_kernel/src/arch/x86_64/syscall_handler.asm

@@ -23,8 +23,8 @@ extern x86_64_do_syscall
 extern x86_64_check_sysenter
 global x86_64_syscall_handler
 
-%define TSS_TEMP_USTACK_OFF 0x1c
-%define TSS_RSP0_OFF        0x04
+%define TSS_TEMP_USTACK_OFF 0x24
+%define TSS_RSP0_OFF        0xc
 
 %define USERLAND_SS         0x23
 %define USERLAND_CS         0x2b

src/aero_kernel/src/arch/x86_64/task.rs

@@ -548,7 +548,7 @@ pub fn arch_task_spinup(from: &mut ArchTask, to: &ArchTask) {
     unsafe {
         // Load the new thread's kernel stack pointer everywhere it's needed.
         let kstackp = to.context_switch_rsp.as_u64();
-        super::gdt::get_task_state_segment().rsp[0] = kstackp;
+        super::gdt::TSS.rsp[0] = kstackp;
         io::wrmsr(io::IA32_SYSENTER_ESP, kstackp);
 
         // Preserve and restore the %fs, %gs bases.

src/aero_kernel/src/arch/x86_64/tls.rs

@@ -23,15 +23,8 @@
 //! * <https://wiki.osdev.org/Thread_Local_Storage>
 //! * <https://doc.rust-lang.org/std/thread/struct.LocalKey.html>
 
-use core::alloc::Layout;
-
-use alloc::alloc::alloc_zeroed;
 use alloc::vec::Vec;
 
-use super::gdt::*;
-use super::io;
-
-use crate::mem::paging::VirtAddr;
 use crate::utils::sync::Mutex;
 
 use raw_cpuid::FeatureInfo;
@@ -110,37 +103,11 @@ pub struct CpuInfo {
     pub features: Vec<&'static &'static str>,
 }
 
-#[repr(C)]
-pub struct PerCpuData {
-    pub cpuid: usize,
-    pub lapic_timer_frequency: u32,
-
-    pub(super) gdt: &'static mut [GdtEntry],
-    pub(super) pf_resume: VirtAddr,
-}
-
-/// SAFETY: The GS base should point to the kernel PCR.
 pub fn get_cpuid() -> usize {
-    get_percpu().cpuid
+    0
 }
 
-/// SAFETY: The GS base should point to the kernel PCR.
-pub fn get_percpu() -> &'static mut PerCpuData {
-    &mut get_kpcr().cpu_local
-}
-
-pub fn init(cpuid: usize) {
-    // NOTE: Inside kernel space, the GS base will always point to the CPU local data and when
-    // jumping to userland `swapgs` is called making the GS base point to the userland TLS data.
-    unsafe {
-        let kpcr_layout = Layout::new::<Kpcr>();
-
-        let kpcr_ptr = alloc_zeroed(kpcr_layout) as *mut Kpcr;
-        io::wrmsr(io::IA32_GS_BASE, kpcr_ptr as u64);
-    }
-
-    get_percpu().cpuid = cpuid;
-
+pub fn init() {
     let cpuid = raw_cpuid::CpuId::new();
 
     let features = cpuid