feat: add percpu variables

Cargo.lock

@@ -15,6 +15,7 @@ dependencies = [
 name = "arch"
 version = "0.1.0"
 dependencies = [
+ "percpu",
  "x86_64",
 ]
 
@@ -120,9 +121,9 @@ dependencies = [
 
 [[package]]
 name = "libc"
-version = "0.2.159"
+version = "0.2.164"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "561d97a539a36e26a9a5fad1ea11a3039a67714694aaa379433e580854bc3dc5"
+checksum = "433bfe06b8c75da9b2e3fbea6e5329ff87748f0b144ef75306e674c3f6f7c13f"
 
 [[package]]
 name = "libloading"
@@ -172,11 +173,28 @@ dependencies = [
  "minimal-lexical",
 ]
 
+[[package]]
+name = "percpu"
+version = "0.1.0"
+dependencies = [
+ "percpu-macros",
+ "x86_64",
+]
+
+[[package]]
+name = "percpu-macros"
+version = "0.1.0"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "prettyplease"
-version = "0.2.22"
+version = "0.2.25"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "479cf940fbbb3426c32c5d5176f62ad57549a0bb84773423ba8be9d089f5faba"
+checksum = "64d1ec885c64d0457d564db4ec299b2dae3f9c02808b8ad9c3a089c591b18033"
 dependencies = [
  "proc-macro2",
  "syn",
@@ -184,9 +202,9 @@ dependencies = [
 
 [[package]]
 name = "proc-macro2"
-version = "1.0.87"
+version = "1.0.92"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b3e4daa0dcf6feba26f985457cdf104d4b4256fc5a09547140f3631bb076b19a"
+checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0"
 dependencies = [
  "unicode-ident",
 ]
@@ -202,9 +220,9 @@ dependencies = [
 
 [[package]]
 name = "regex"
-version = "1.11.0"
+version = "1.11.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "38200e5ee88914975b69f657f0801b6f6dccafd44fd9326302a4aaeecfacb1d8"
+checksum = "b544ef1b4eac5dc2db33ea63606ae9ffcfac26c1416a2806ae0bf5f56b201191"
 dependencies = [
  "aho-corasick",
  "memchr",
@@ -214,9 +232,9 @@ dependencies = [
 
 [[package]]
 name = "regex-automata"
-version = "0.4.8"
+version = "0.4.9"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "368758f23274712b504848e9d5a6f010445cc8b87a7cdb4d7cbee666c1288da3"
+checksum = "809e8dc61f6de73b46c85f4c96486310fe304c434cfa43669d7b40f711150908"
 dependencies = [
  "aho-corasick",
  "memchr",
@@ -258,9 +276,9 @@ dependencies = [
 
 [[package]]
 name = "syn"
-version = "2.0.79"
+version = "2.0.89"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "89132cd0bf050864e1d38dc3bbc07a0eb8e7530af26344d3d2bbbef83499f590"
+checksum = "44d46482f1c1c87acd84dea20c1bf5ebff4c757009ed6bf19cfd36fb10e92c4e"
 dependencies = [
  "proc-macro2",
  "quote",
@@ -269,9 +287,9 @@ dependencies = [
 
 [[package]]
 name = "unicode-ident"
-version = "1.0.13"
+version = "1.0.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e91b56cd4cadaeb79bbf1a5645f6b4f8dc5bde8834ad5894a8db35fda9efa1fe"
+checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83"
 
 [[package]]
 name = "windows-targets"

Cargo.toml

@@ -14,8 +14,9 @@ lazy_static = { version = "1.5.0", features = ["spin_no_std"] }
 spin = "0.9.8"
 
 [features]
-default = []
+default = ["smp"]
 debug_syscall = []
+smp = []
 
 [build-dependencies]
 bindgen = "0.70.1"

arch/Cargo.lock

@@ -6,9 +6,60 @@ version = 3
 name = "arch"
 version = "0.1.0"
 dependencies = [
+ "percpu",
  "x86_64",
 ]
 
+[[package]]
+name = "percpu"
+version = "0.1.0"
+dependencies = [
+ "percpu-macros",
+]
+
+[[package]]
+name = "percpu-macros"
+version = "0.1.0"
+dependencies = [
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.92"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.37"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "syn"
+version = "2.0.89"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "44d46482f1c1c87acd84dea20c1bf5ebff4c757009ed6bf19cfd36fb10e92c4e"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83"
+
 [[package]]
 name = "x86_64"
 version = "0.1.0"

arch/Cargo.toml

@@ -5,3 +5,4 @@ edition = "2021"
 
 [dependencies]
 x86_64 = { path="./x86_64" }
+percpu = { path="./percpu" }

arch/percpu/Cargo.lock

@@ -0,0 +1,53 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 3
+
+[[package]]
+name = "percpu"
+version = "0.1.0"
+dependencies = [
+ "percpu-macros",
+]
+
+[[package]]
+name = "percpu-macros"
+version = "0.1.0"
+dependencies = [
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.92"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.37"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "syn"
+version = "2.0.89"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "44d46482f1c1c87acd84dea20c1bf5ebff4c757009ed6bf19cfd36fb10e92c4e"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83"

arch/percpu/Cargo.toml

@@ -0,0 +1,8 @@
+[package]
+name = "percpu"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+percpu-macros = { path = "macros" }
+x86_64 = { path = "../x86_64" }

arch/percpu/macros/Cargo.lock

@@ -0,0 +1,47 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 3
+
+[[package]]
+name = "percpu-macros"
+version = "0.1.0"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.92"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.37"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "syn"
+version = "2.0.89"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "44d46482f1c1c87acd84dea20c1bf5ebff4c757009ed6bf19cfd36fb10e92c4e"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83"

arch/percpu/macros/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "percpu-macros"
+version = "0.1.0"
+edition = "2021"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+syn = { version = "2.0", features = ["full"] }
+quote = "1.0"
+proc-macro2 = "1.0"

arch/percpu/macros/src/arch.rs

@@ -0,0 +1,22 @@
+use proc_macro2::TokenStream;
+use quote::quote;
+use syn::{Ident, Type};
+
+/// Get the base address for percpu variables of the current thread.
+pub fn get_percpu_pointer(percpu: &Ident, ty: &Type) -> TokenStream {
+    quote! {
+        #[cfg(target_arch = "x86_64")]
+        {
+            let base: *mut #ty;
+            ::core::arch::asm!(
+                "mov %gs:0, {address}",
+                "add ${percpu_pointer}, {address}",
+                percpu_pointer = sym #percpu,
+                address = out(reg) base,
+                options(att_syntax)
+            );
+            base
+        }
+    }
+    .into()
+}

arch/percpu/macros/src/lib.rs

@@ -0,0 +1,86 @@
+extern crate proc_macro;
+
+use proc_macro::TokenStream;
+use quote::{format_ident, quote};
+use syn::{parse_macro_input, ItemStatic};
+
+mod arch;
+
+#[proc_macro_attribute]
+pub fn define_percpu(attrs: TokenStream, item: TokenStream) -> TokenStream {
+    if !attrs.is_empty() {
+        panic!("`define_percpu` attribute does not take any arguments");
+    }
+
+    let item = parse_macro_input!(item as ItemStatic);
+    let vis = &item.vis;
+    let ident = &item.ident;
+    let ty = &item.ty;
+    let expr = &item.expr;
+
+    if !["bool", "u8", "u16", "u32", "u64", "usize"].contains(&quote!(#ty).to_string().as_str()) {
+        panic!("`define_percpu` only supports bool, u8, u16, u32, u64 and usize");
+    }
+
+    let inner_ident = format_ident!("_percpu_inner_{}", ident);
+    let access_ident = format_ident!("_access_{}", ident);
+
+    let integer_methods = match quote!(#ty).to_string().as_str() {
+        "bool" => quote! {},
+        _ => quote! {
+            pub fn add(&self, value: #ty) {
+                *unsafe { self.as_mut() } += value;
+            }
+
+            pub fn sub(&self, value: #ty) {
+                *unsafe { self.as_mut() } -= value;
+            }
+        },
+    };
+
+    let as_ptr = arch::get_percpu_pointer(&inner_ident, &ty);
+
+    quote! {
+        #[link_section = ".percpu"]
+        #[allow(non_upper_case_globals)]
+        static mut #inner_ident: #ty = #expr;
+        #[allow(non_camel_case_types)]
+        #vis struct #access_ident;
+        #vis static #ident: #access_ident = #access_ident;
+
+        impl #access_ident {
+            pub unsafe fn as_ptr(&self) -> *mut #ty {
+                #as_ptr
+            }
+
+            pub fn get(&self) -> #ty {
+                unsafe { self.as_ptr().read() }
+            }
+
+            pub fn set(&self, value: #ty) {
+                unsafe { self.as_ptr().write(value) }
+            }
+
+            /// # Safety
+            /// This function is unsafe because it allows for immutable aliasing of the percpu
+            /// variable.
+            /// Make sure that preempt is disabled when calling this function.
+            pub unsafe fn as_ref(&self) -> & #ty {
+                // SAFETY: This is safe because `as_ptr()` is guaranteed to be valid.
+                self.as_ptr().as_ref().unwrap()
+            }
+
+            /// # Safety
+            /// This function is unsafe because it allows for mutable aliasing of the percpu
+            /// variable.
+            /// Make sure that preempt is disabled when calling this function.
+            pub unsafe fn as_mut(&self) -> &mut #ty {
+                // SAFETY: This is safe because `as_ptr()` is guaranteed to be valid.
+                self.as_ptr().as_mut().unwrap()
+            }
+
+            #integer_methods
+        }
+    }
+    .into()
+}

arch/percpu/src/arch.rs

@@ -0,0 +1,25 @@
+pub unsafe fn save_percpu_pointer(percpu_area_base: *mut ()) {
+    #[cfg(target_arch = "x86_64")]
+    x86_64::task::wrmsr(0xC0000101, percpu_area_base as u64);
+
+    #[cfg(not(target_arch = "x86_64"))]
+    compile_error!("unsupported architecture");
+}
+
+pub unsafe fn set_percpu_area_thiscpu(percpu_area_base: *mut ()) {
+    use core::arch::asm;
+
+    save_percpu_pointer(percpu_area_base);
+
+    #[cfg(target_arch = "x86_64")]
+    {
+        asm!(
+            "movq {}, %gs:0",
+            in(reg) percpu_area_base,
+            options(att_syntax)
+        );
+    }
+
+    #[cfg(not(target_arch = "x86_64"))]
+    compile_error!("unsupported architecture");
+}

arch/percpu/src/lib.rs

@@ -0,0 +1,6 @@
+#![no_std]
+
+mod arch;
+
+pub use arch::set_percpu_area_thiscpu;
+pub use percpu_macros::define_percpu;

arch/src/lib.rs

@@ -45,6 +45,9 @@ pub mod task {
     pub fn context_switch_light(current_task_sp: *mut usize, next_task_sp: *mut usize) {
         x86_64::task::context_switch_light(current_task_sp, next_task_sp);
     }
+
+    #[cfg(target_arch = "x86_64")]
+    pub use x86_64::task::{rdmsr, wrmsr};
 }
 
 pub mod interrupt {
@@ -91,4 +94,5 @@ pub mod io {
     }
 }
 
+pub use percpu::{define_percpu, set_percpu_area_thiscpu};
 pub use x86_64;

arch/x86_64/src/task.rs

@@ -24,6 +24,40 @@ pub fn freeze() -> ! {
     }
 }
 
+#[inline(always)]
+pub fn rdmsr(msr: u32) -> u64 {
+    let edx: u32;
+    let eax: u32;
+
+    unsafe {
+        asm!(
+            "rdmsr",
+            in("ecx") msr,
+            out("eax") eax,
+            out("edx") edx,
+            options(att_syntax),
+        );
+    }
+
+    (edx as u64) << 32 | eax as u64
+}
+
+#[inline(always)]
+pub fn wrmsr(msr: u32, value: u64) {
+    let eax = value as u32;
+    let edx = (value >> 32) as u32;
+
+    unsafe {
+        asm!(
+            "wrmsr",
+            in("ecx") msr,
+            in("eax") eax,
+            in("edx") edx,
+            options(att_syntax),
+        );
+    }
+}
+
 global_asm!(
     r"
     .macro movcfi reg, offset

src/asm/interrupt.s

@@ -42,6 +42,11 @@ ISR_stub:
 	.cfi_def_cfa_offset 0x18
 	.cfi_offset %rsp, 0x10
 
+	cmpq $0x08, 24(%rsp)
+	je 1f
+	swapgs
+
+1:
 	sub $0x78, %rsp
 	.cfi_def_cfa_offset 0x90
 
@@ -101,6 +106,11 @@ ISR_stub_restore:
 	add $0x88, %rsp
 	.cfi_def_cfa_offset 0x08
 
+	cmpq $0x08, 8(%rsp)
+	je 1f
+	swapgs
+
+1:
 	iretq
 	.cfi_endproc
 

src/kernel.ld

@@ -77,6 +77,8 @@ SECTIONS
         QUAD(ABSOLUTE(_fix_start));
         FIX_END = .;
         QUAD(ABSOLUTE(_fix_end));
+        PERCPU_PAGES = .;
+        QUAD(_PERCPU_PAGES);
 
         . = ALIGN(0x1000);
         RODATA_END = .;
@@ -102,10 +104,11 @@ SECTIONS
     DATA_PAGES = (DATA_END - DATA_START) / 0x1000;
 
     _PERCPU_DATA_START = .;
-    .percpu 0 (NOLOAD) : AT(LOADADDR(.data) + SIZEOF(.data))
+    .percpu 0 : AT(LOADADDR(.data) + SIZEOF(.data))
     {
         PERCPU_START = .;
-        . += 0x10; /* Reserved for x86 percpu pointer */
+        QUAD(0); /* Reserved for x86 percpu pointer */
+        QUAD(0);
 
         *(.percpu .percpu*)
 
@@ -114,7 +117,7 @@ SECTIONS
     } > KPERCPU
     _PERCPU_LENGTH = PERCPU_END - PERCPU_START;
 
-    PERCPU_PAGES = _PERCPU_LENGTH / 0x1000;
+    _PERCPU_PAGES = _PERCPU_LENGTH / 0x1000;
 
     .bss :
     {
@@ -126,14 +129,14 @@ SECTIONS
         BSS_END = .;
     } > KBSS
 
-    KIMAGE_PAGES = TEXT_PAGES + RODATA_PAGES + PERCPU_PAGES + DATA_PAGES;
+    KIMAGE_PAGES = TEXT_PAGES + RODATA_PAGES + _PERCPU_PAGES + DATA_PAGES;
     BSS_PAGES = (BSS_END - BSS_START) / 0x1000;
     KERNEL_MAGIC = 0x01145140;
 
     KIMAGE_32K_COUNT = (KIMAGE_PAGES * 0x1000 + 32 * 1024 - 1) / (32 * 1024);
 
     .eh_frame :
-        AT(LOADADDR(.data) + SIZEOF(.data))
+        AT(LOADADDR(.percpu) + SIZEOF(.percpu))
     {
         KEEP(*(.eh_frame*))
         . = ALIGN(0x1000);

src/kernel.rs

@@ -9,6 +9,9 @@ pub mod timer;
 pub mod user;
 pub mod vfs;
 
+#[cfg(feature = "smp")]
+pub mod smp;
+
 mod chardev;
 mod terminal;
 

src/kernel/smp.rs

@@ -0,0 +1,3 @@
+mod percpu;
+
+pub use percpu::{alloc_percpu_area, set_percpu_area};

src/kernel/smp/percpu.rs

@@ -0,0 +1,32 @@
+use crate::kernel::mem::{paging::Page, phys::PhysPtr as _};
+
+/// # Safety
+/// The memory allocated by this function will never be freed and can only be used
+/// for per-cpu area.
+pub unsafe fn alloc_percpu_area() -> *mut () {
+    extern "C" {
+        static PERCPU_PAGES: usize;
+        fn _PERCPU_DATA_START();
+    }
+    assert_eq!(
+        unsafe { PERCPU_PAGES },
+        1,
+        "We support only 1 page per-cpu variables for now"
+    );
+
+    let page = Page::alloc_one();
+    unsafe {
+        page.as_cached()
+            .as_ptr::<u8>()
+            .copy_from_nonoverlapping(_PERCPU_DATA_START as *const _, page.len())
+    };
+
+    let addr = page.as_cached().as_ptr();
+    core::mem::forget(page);
+
+    addr
+}
+
+pub unsafe fn set_percpu_area(area: *mut ()) {
+    arch::set_percpu_area_thiscpu(area)
+}

src/kernel/task/thread.rs

@@ -1,5 +1,4 @@
 use core::{
-    arch::asm,
     cell::RefCell,
     cmp,
     sync::atomic::{self, AtomicU32},
@@ -193,7 +192,8 @@ struct ThreadInner {
     name: Arc<[u8]>,
 
     /// Thread TLS descriptor 32-bit
-    tls_desc32: u64,
+    tls_desc32: Option<u64>,
+    tls_base: Option<u64>,
 
     /// User pointer
     /// Store child thread's tid when child thread returns to user space.
@@ -812,7 +812,8 @@ impl Thread {
             state: Spin::new(ThreadState::Preparing),
             inner: Spin::new(ThreadInner {
                 name,
-                tls_desc32: 0,
+                tls_desc32: None,
+                tls_base: None,
                 set_child_tid: 0,
             }),
         });
@@ -842,6 +843,7 @@ impl Thread {
             inner: Spin::new(ThreadInner {
                 name: other_inner.name.clone(),
                 tls_desc32: other_inner.tls_desc32,
+                tls_base: other_inner.tls_base,
                 set_child_tid: other_inner.set_child_tid,
             }),
         });
@@ -896,23 +898,19 @@ impl Thread {
     }
 
     pub fn load_thread_area32(&self) {
+        const IA32_KERNEL_GS_BASE: u32 = 0xc0000102;
+
         let inner = self.inner.lock();
-        if inner.tls_desc32 == 0 {
-            return;
+
+        if let Some(desc32) = inner.tls_desc32 {
+            // SAFETY: `tls32` should be per cpu.
+            let tls32_addr = CachedPP::new(0x0 + 7 * 8);
+            tls32_addr.as_mut::<u64>().clone_from(&desc32);
         }
 
-        // SAFETY: `tls32` should be per cpu.
-        let tls32_addr = CachedPP::new(0x0 + 7 * 8);
-        tls32_addr.as_mut::<u64>().clone_from(&inner.tls_desc32);
-
-        unsafe {
-            asm!(
-                "mov %gs, %ax",
-                "mov %ax, %gs",
-                out("ax") _,
-                options(att_syntax)
-            )
-        };
+        if let Some(base) = inner.tls_base {
+            arch::x86_64::task::wrmsr(IA32_KERNEL_GS_BASE, base);
+        }
     }
 
     pub fn set_thread_area(&self, desc: &mut UserDescriptor) -> KResult<()> {
@@ -920,9 +918,17 @@ impl Thread {
 
         // Clear the TLS area if it is not present.
         if desc.flags.is_read_exec_only() && !desc.flags.is_present() {
-            if desc.limit != 0 && desc.base != 0 {
-                CheckedUserPointer::new(desc.base as _, desc.limit as _)?.zero()?;
+            if desc.limit == 0 || desc.base == 0 {
+                return Ok(());
             }
+
+            let len = if desc.flags.is_limit_in_pages() {
+                (desc.limit as usize) << 12
+            } else {
+                desc.limit as usize
+            };
+
+            CheckedUserPointer::new(desc.base as _, len)?.zero()?;
             return Ok(());
         }
 
@@ -931,17 +937,19 @@ impl Thread {
         }
         desc.entry = 7;
 
-        inner.tls_desc32 = desc.limit as u64 & 0xffff;
-        inner.tls_desc32 |= (desc.base as u64 & 0xffffff) << 16;
-        inner.tls_desc32 |= 0x4_0_f2_000000_0000;
-        inner.tls_desc32 |= (desc.limit as u64 & 0xf_0000) << (48 - 16);
+        let mut desc32 = desc.limit as u64 & 0xffff;
+        desc32 |= (desc.base as u64 & 0xffffff) << 16;
+        desc32 |= 0x4_0_f2_000000_0000;
+        desc32 |= (desc.limit as u64 & 0xf_0000) << (48 - 16);
 
         if desc.flags.is_limit_in_pages() {
-            inner.tls_desc32 |= 1 << 55;
+            desc32 |= 1 << 55;
         }
 
-        inner.tls_desc32 |= (desc.base as u64 & 0xff_000000) << (56 - 24);
+        desc32 |= (desc.base as u64 & 0xff_000000) << (56 - 24);
 
+        inner.tls_desc32 = Some(desc32);
+        inner.tls_base = Some(desc.base as u64);
         Ok(())
     }
 

src/lib.rs

@@ -23,6 +23,7 @@ mod rcu;
 mod sync;
 
 use alloc::ffi::CString;
+use arch::task::rdmsr;
 use core::{
     alloc::{GlobalAlloc, Layout},
     arch::{asm, global_asm},
@@ -110,22 +111,6 @@ extern "C" {
     fn boot_semaphore();
 }
 
-fn rdmsr(msr: u32) -> u64 {
-    let edx: u32;
-    let eax: u32;
-
-    unsafe {
-        asm!(
-            "rdmsr",
-            in("ecx") msr,
-            out("eax") eax,
-            out("edx") edx,
-        )
-    };
-
-    (edx as u64) << 32 | eax as u64
-}
-
 fn bootstrap_cpus() {
     let apic_base = rdmsr(0x1b);
     assert_eq!(apic_base & 0x800, 0x800, "LAPIC not enabled");
@@ -163,6 +148,11 @@ pub extern "C" fn rust_kinit(early_kstack_pfn: usize) -> ! {
     // We don't call global constructors.
     // Rust doesn't need that, and we're not going to use global variables in C++.
 
+    unsafe {
+        let area = kernel::smp::alloc_percpu_area();
+        kernel::smp::set_percpu_area(area);
+    }
+
     kernel::interrupt::init().unwrap();
 
     // TODO: Move this to rust.
@@ -259,8 +249,7 @@ extern "C" fn init_process(early_kstack_pfn: usize) {
 
     unsafe {
         asm!(
-            "mov %ax, %fs",
-            "mov %ax, %gs",
+            "swapgs",
             "mov ${ds}, %rax",
             "mov %ax, %ds",
             "mov %ax, %es",

src/sync.rs

@@ -5,26 +5,26 @@ pub mod spin;
 pub mod strategy;
 
 pub mod preempt {
-    use core::sync::atomic::{compiler_fence, AtomicUsize, Ordering};
+    use core::sync::atomic::{compiler_fence, Ordering};
 
-    /// TODO: This should be per cpu.
-    static PREEMPT_COUNT: AtomicUsize = AtomicUsize::new(0);
+    #[arch::define_percpu]
+    static PREEMPT_COUNT: usize = 0;
 
     #[inline(always)]
     pub fn disable() {
-        PREEMPT_COUNT.fetch_add(1, Ordering::Relaxed);
+        PREEMPT_COUNT.add(1);
         compiler_fence(Ordering::SeqCst);
     }
 
     #[inline(always)]
     pub fn enable() {
         compiler_fence(Ordering::SeqCst);
-        PREEMPT_COUNT.fetch_sub(1, Ordering::Relaxed);
+        PREEMPT_COUNT.sub(1);
     }
 
     #[inline(always)]
     pub fn count() -> usize {
-        PREEMPT_COUNT.load(Ordering::Relaxed)
+        PREEMPT_COUNT.get()
     }
 }