diff --git a/src/arch/riscv32/mem.rs b/src/arch/riscv32/mem.rs
new file mode 100644
index 0000000..84040bf
--- /dev/null
+++ b/src/arch/riscv32/mem.rs
@@ -0,0 +1,6 @@
+use super::asm;
+
+pub fn update_kernel_sp(sp: usize) {
+    unsafe { core::arch::asm!("mv a0, {}", in(reg) sp) };
+    unsafe { asm::set_kernel_sp() };
+}
diff --git a/src/arch/riscv32/mod.rs b/src/arch/riscv32/mod.rs
index ed4df01..65da303 100644
--- a/src/arch/riscv32/mod.rs
+++ b/src/arch/riscv32/mod.rs
@@ -1,4 +1,5 @@
 pub mod asm;
+pub mod mem;
 pub mod scheduler;
 pub mod start;
 pub mod task;
diff --git a/src/boot.rs b/src/boot.rs
index 05805eb..9c17b40 100644
--- a/src/boot.rs
+++ b/src/boot.rs
@@ -7,7 +7,7 @@ use crate::{
     ktime::set_ktime_seconds,
     log,
     logs::LogLevel,
-    mem::{memory_init, update_kernel_sp},
+    mem::{mem_update_kernel_sp, memory_init},
     platform::platform_init,
 };
 
@@ -34,7 +34,7 @@ pub fn kernel_early_boot(core: usize, dtb_addr: usize) -> ! {
     );
     memory_init();
     log!(LogLevel::Debug, "Switch to new kernel stack...");
-    update_kernel_sp();
+    mem_update_kernel_sp();
     // Allow empty loop because it will never enter, just to make the fn never return without
     // warning
     #[allow(clippy::empty_loop)]
diff --git a/src/mem/mod.rs b/src/mem/mod.rs
index e9d1875..a0df021 100644
--- a/src/mem/mod.rs
+++ b/src/mem/mod.rs
@@ -19,12 +19,13 @@ Tests files:
 
 mod kernel;
 
-use core::{arch::asm, mem};
+use core::mem;
 
 use kernel::{__kernel_end, __kernel_start, KernelStack};
 
 use crate::{
-    arch, config::KERNEL_STACK_SIZE, log, logs::LogLevel, platform::mem::platform_init_mem,
+    arch::mem::update_kernel_sp, config::KERNEL_STACK_SIZE, log, logs::LogLevel,
+    platform::mem::platform_init_mem,
 };
 
 pub struct Memory {
@@ -137,12 +138,6 @@ pub fn memory_init() {
     }
 }
 
-#[unsafe(no_mangle)]
-pub fn update_kernel_sp() {
-    unsafe { asm!("mv a0, {}", in(reg) MEMORY.kernel_stack.top) };
-    unsafe { arch::asm::set_kernel_sp() };
-}
-
 pub fn mem_kernel_stack_info<'a>() -> &'a KernelStack {
     #[allow(static_mut_refs)]
     unsafe {
@@ -150,6 +145,11 @@ pub fn mem_kernel_stack_info<'a>() -> &'a KernelStack {
     }
 }
 
+pub fn mem_update_kernel_sp() {
+    let sp: usize = unsafe { MEMORY.kernel_stack.top };
+    update_kernel_sp(sp);
+}
+
 /// Return the hi and lo address of the RAM
 /// first index is hi, second is lo
 pub fn mem_reg_info() -> [usize; 2] {