tls: add support of thread-local storage

api/axfeat/Cargo.toml

@@ -27,6 +27,7 @@ alloc-tlsf = ["axalloc/tlsf"]
 alloc-slab = ["axalloc/slab"]
 alloc-buddy = ["axalloc/buddy"]
 paging = ["alloc", "axhal/paging", "axruntime/paging"]
+tls = ["alloc", "axhal/tls", "axruntime/tls", "axtask?/tls"]
 
 # Multi-threading and scheduler
 multitask = ["alloc", "axtask/multitask", "axsync/multitask", "axruntime/multitask"]

api/axfeat/src/lib.rs

@@ -13,6 +13,7 @@
 //!     - `alloc-slab`: Use the slab allocator.
 //!     - `alloc-buddy`: Use the buddy system allocator.
 //!     - `paging`: Enable page table manipulation.
+//!     - `tls`: Enable thread-local storage.
 //! - Task management
 //!     - `multitask`: Enable multi-threading support.
 //!     - `sched_fifo`: Use the FIFO cooperative scheduler.

crates/percpu/src/imp.rs

@@ -71,7 +71,7 @@ pub fn get_local_thread_pointer() -> usize {
                     unimplemented!()
                 };
             } else if #[cfg(any(target_arch = "riscv32", target_arch = "riscv64"))] {
-                core::arch::asm!("mv {}, tp", out(reg) tp)
+                core::arch::asm!("mv {}, gp", out(reg) tp)
             } else if #[cfg(target_arch = "aarch64")] {
                 core::arch::asm!("mrs {}, TPIDR_EL1", out(reg) tp)
             }
@@ -105,7 +105,7 @@ pub fn set_local_thread_pointer(cpu_id: usize) {
                 }
                 SELF_PTR.write_current_raw(tp);
             } else if #[cfg(any(target_arch = "riscv32", target_arch = "riscv64"))] {
-                core::arch::asm!("mv tp, {}", in(reg) tp)
+                core::arch::asm!("mv gp, {}", in(reg) tp)
             } else if #[cfg(target_arch = "aarch64")] {
                 core::arch::asm!("msr TPIDR_EL1, {}", in(reg) tp)
             }

crates/percpu/src/lib.rs

@@ -10,6 +10,12 @@
 //! pointer register to obtain the corresponding per-CPU data area, and then add
 //! an offset to access the corresponding field.
 //!
+//! # Notes
+//!
+//! Since RISC-V does not provide separate thread pointer registers for user and
+//! kernel mode, we temporarily use the `gp` register to point to the per-CPU data
+//! area, while the `tp` register is used for thread-local storage.
+//!
 //! # Examples
 //!
 //! ```no_run

crates/percpu_macros/src/arch.rs

@@ -57,7 +57,7 @@ pub fn gen_current_ptr(symbol: &Ident, ty: &Type) -> proc_macro2::TokenStream {
             #[cfg(target_arch = "aarch64")]
             ::core::arch::asm!("mrs {}, TPIDR_EL1", out(reg) base);
             #[cfg(any(target_arch = "riscv32", target_arch = "riscv64"))]
-            ::core::arch::asm!("mv {}, tp", out(reg) base);
+            ::core::arch::asm!("mv {}, gp", out(reg) base);
             (base + self.offset()) as *const #ty
         }
     })
@@ -77,7 +77,7 @@ pub fn gen_read_current_raw(symbol: &Ident, ty: &Type) -> proc_macro2::TokenStre
     let rv64_asm = quote! {
         ::core::arch::asm!(
             "lui {0}, %hi({VAR})",
-            "add {0}, {0}, tp",
+            "add {0}, {0}, gp",
             concat!(#rv64_op, " {0}, %lo({VAR})({0})"),
             out(reg) value,
             VAR = sym #symbol,
@@ -154,7 +154,7 @@ pub fn gen_write_current_raw(symbol: &Ident, val: &Ident, ty: &Type) -> proc_mac
     let rv64_code = quote! {
         ::core::arch::asm!(
             "lui {0}, %hi({VAR})",
-            "add {0}, {0}, tp",
+            "add {0}, {0}, gp",
             concat!(#rv64_op, " {1}, %lo({VAR})({0})"),
             out(reg) _,
             in(reg) #val as #ty_fixup,

modules/axhal/Cargo.toml

@@ -11,9 +11,12 @@ documentation = "https://rcore-os.github.io/arceos/axhal/index.html"
 
 [features]
 smp = []
+alloc = []
 fp_simd = []
 paging = ["axalloc", "page_table"]
 irq = []
+tls = ["alloc"]
+default = []
 
 [dependencies]
 log = "0.4"

modules/axhal/src/arch/aarch64/context.rs

@@ -75,9 +75,10 @@ impl TaskContext {
 
     /// Initializes the context for a new task, with the given entry point and
     /// kernel stack.
-    pub fn init(&mut self, entry: usize, kstack_top: VirtAddr) {
+    pub fn init(&mut self, entry: usize, kstack_top: VirtAddr, tls_area: VirtAddr) {
         self.sp = kstack_top.as_usize() as u64;
         self.lr = entry as u64;
+        self.tpidr_el0 = tls_area.as_usize() as u64;
     }
 
     /// Switches to another task.
@@ -171,6 +172,7 @@ unsafe extern "C" fn fpstate_switch(_current_fpstate: &mut FpState, _next_fpstat
         msr     fpcr, x9
         msr     fpsr, x10
 
+        isb
         ret",
         options(noreturn),
     )

modules/axhal/src/arch/aarch64/mod.rs

@@ -115,3 +115,23 @@ pub fn set_exception_vector_base(vbar_el1: usize) {
 pub fn flush_dcache_line(vaddr: VirtAddr) {
     unsafe { asm!("dc ivac, {0:x}; dsb sy; isb", in(reg) vaddr.as_usize()) };
 }
+
+/// Reads the thread pointer of the current CPU.
+///
+/// It is used to implement TLS (Thread Local Storage).
+#[inline]
+pub fn read_thread_pointer() -> usize {
+    TPIDR_EL0.get() as usize
+}
+
+/// Writes the thread pointer of the current CPU.
+///
+/// It is used to implement TLS (Thread Local Storage).
+///
+/// # Safety
+///
+/// This function is unsafe as it changes the CPU states.
+#[inline]
+pub unsafe fn write_thread_pointer(tpidr_el0: usize) {
+    TPIDR_EL0.set(tpidr_el0 as _)
+}

modules/axhal/src/arch/riscv/context.rs

@@ -84,6 +84,8 @@ pub struct TaskContext {
     pub s9: usize,
     pub s10: usize,
     pub s11: usize,
+
+    pub tp: usize,
     // TODO: FP states
 }
 
@@ -95,18 +97,24 @@ impl TaskContext {
 
     /// Initializes the context for a new task, with the given entry point and
     /// kernel stack.
-    pub fn init(&mut self, entry: usize, kstack_top: VirtAddr) {
+    pub fn init(&mut self, entry: usize, kstack_top: VirtAddr, tls_area: VirtAddr) {
         self.sp = kstack_top.as_usize();
         self.ra = entry;
+        self.tp = tls_area.as_usize();
     }
 
     /// Switches to another task.
     ///
     /// It first saves the current task's context from CPU to this place, and then
     /// restores the next task's context from `next_ctx` to CPU.
     pub fn switch_to(&mut self, next_ctx: &Self) {
+        #[cfg(feature = "tls")]
+        {
+            self.tp = super::read_thread_pointer();
+            unsafe { super::write_thread_pointer(next_ctx.tp) };
+        }
         unsafe {
-            // TODO: switch TLS
+            // TODO: switch FP states
             context_switch(self, next_ctx)
         }
     }

modules/axhal/src/arch/riscv/mod.rs

@@ -85,3 +85,25 @@ pub fn flush_tlb(vaddr: Option<VirtAddr>) {
 pub fn set_trap_vector_base(stvec: usize) {
     unsafe { stvec::write(stvec, stvec::TrapMode::Direct) }
 }
+
+/// Reads the thread pointer of the current CPU.
+///
+/// It is used to implement TLS (Thread Local Storage).
+#[inline]
+pub fn read_thread_pointer() -> usize {
+    let tp;
+    unsafe { core::arch::asm!("mv {}, tp", out(reg) tp) };
+    tp
+}
+
+/// Writes the thread pointer of the current CPU.
+///
+/// It is used to implement TLS (Thread Local Storage).
+///
+/// # Safety
+///
+/// This function is unsafe as it changes the CPU states.
+#[inline]
+pub unsafe fn write_thread_pointer(tp: usize) {
+    core::arch::asm!("mv tp, {}", in(reg) tp)
+}

modules/axhal/src/arch/x86_64/context.rs

@@ -136,6 +136,8 @@ pub struct TaskContext {
     pub kstack_top: VirtAddr,
     /// `RSP` after all callee-saved registers are pushed.
     pub rsp: u64,
+    /// Thread Local Storage (TLS).
+    pub fs_base: usize,
     /// Extended states, i.e., FP/SIMD states.
     #[cfg(feature = "fp_simd")]
     pub ext_state: ExtendedState,
@@ -147,14 +149,15 @@ impl TaskContext {
         Self {
             kstack_top: VirtAddr::from(0),
             rsp: 0,
+            fs_base: 0,
             #[cfg(feature = "fp_simd")]
             ext_state: ExtendedState::default(),
         }
     }
 
     /// Initializes the context for a new task, with the given entry point and
     /// kernel stack.
-    pub fn init(&mut self, entry: usize, kstack_top: VirtAddr) {
+    pub fn init(&mut self, entry: usize, kstack_top: VirtAddr, tls_area: VirtAddr) {
         unsafe {
             // x86_64 calling convention: the stack must be 16-byte aligned before
             // calling a function. That means when entering a new task (`ret` in `context_switch`
@@ -171,6 +174,7 @@ impl TaskContext {
             self.rsp = frame_ptr as u64;
         }
         self.kstack_top = kstack_top;
+        self.fs_base = tls_area.as_usize();
     }
 
     /// Switches to another task.
@@ -183,10 +187,12 @@ impl TaskContext {
             self.ext_state.save();
             next_ctx.ext_state.restore();
         }
-        unsafe {
-            // TODO: swtich tls
-            context_switch(&mut self.rsp, &next_ctx.rsp)
+        #[cfg(feature = "tls")]
+        {
+            self.fs_base = super::read_thread_pointer();
+            unsafe { super::write_thread_pointer(next_ctx.fs_base) };
         }
+        unsafe { context_switch(&mut self.rsp, &next_ctx.rsp) }
     }
 }
 

modules/axhal/src/arch/x86_64/mod.rs

@@ -8,7 +8,7 @@ mod trap;
 use core::arch::asm;
 
 use memory_addr::{PhysAddr, VirtAddr};
-use x86::{controlregs, tlb};
+use x86::{controlregs, msr, tlb};
 use x86_64::instructions::interrupts;
 
 pub use self::context::{ExtendedState, FxsaveArea, TaskContext, TrapFrame};
@@ -88,3 +88,23 @@ pub fn flush_tlb(vaddr: Option<VirtAddr>) {
         unsafe { tlb::flush_all() }
     }
 }
+
+/// Reads the thread pointer of the current CPU.
+///
+/// It is used to implement TLS (Thread Local Storage).
+#[inline]
+pub fn read_thread_pointer() -> usize {
+    unsafe { msr::rdmsr(msr::IA32_FS_BASE) as usize }
+}
+
+/// Writes the thread pointer of the current CPU.
+///
+/// It is used to implement TLS (Thread Local Storage).
+///
+/// # Safety
+///
+/// This function is unsafe as it changes the CPU states.
+#[inline]
+pub unsafe fn write_thread_pointer(fs_base: usize) {
+    unsafe { msr::wrmsr(msr::IA32_FS_BASE, fs_base as u64) }
+}

modules/axhal/src/lib.rs

@@ -42,6 +42,9 @@ pub mod mem;
 pub mod time;
 pub mod trap;
 
+#[cfg(feature = "tls")]
+pub mod tls;
+
 #[cfg(feature = "irq")]
 pub mod irq;