diff --git a/kernel/src/context/aarch64.rs b/kernel/src/context/aarch64.rs
index 13b91241b..20bc01a44 100644
--- a/kernel/src/context/aarch64.rs
+++ b/kernel/src/context/aarch64.rs
@@ -7,7 +7,7 @@ pub struct ContextArgs {}
 /// Extended [Base] for AArch64.
 #[repr(C)]
 pub(super) struct Context {
-    base: Base, // Must be first field.
+    pub base: Base, // Must be first field.
 }
 
 impl Context {
@@ -19,11 +19,15 @@ impl Context {
         todo!();
     }
 
-    pub unsafe fn load_fixed_ptr<const O: usize, T>() -> *const T {
+    pub unsafe fn load_static_ptr<const O: usize, T>() -> *const T {
         todo!()
     }
 
-    pub unsafe fn load_usize<const O: usize>() -> usize {
+    pub unsafe fn load_ptr<const O: usize, T>() -> *const T {
+        todo!()
+    }
+
+    pub unsafe fn load_volatile_usize<const O: usize>() -> usize {
         todo!()
     }
 }
diff --git a/kernel/src/context/arc.rs b/kernel/src/context/arc.rs
index 146a05ce1..ec0d65785 100644
--- a/kernel/src/context/arc.rs
+++ b/kernel/src/context/arc.rs
@@ -15,9 +15,19 @@ use core::ops::Deref;
 pub struct BorrowedArc<T>(*const T); // A pointer make this type automatically !Send and !Sync.
 
 impl<T> BorrowedArc<T> {
+    /// # Safety
+    /// `v` must be owned by [Arc](alloc::sync::Arc) if not null.
+    pub(super) const unsafe fn new(v: *const T) -> Option<Self> {
+        if v.is_null() {
+            None
+        } else {
+            Some(Self(v))
+        }
+    }
+
     /// # Safety
     /// `v` must be owned by [Arc](alloc::sync::Arc).
-    pub(super) unsafe fn new(v: *const T) -> Self {
+    pub(super) const unsafe fn from_non_null(v: *const T) -> Self {
         Self(v)
     }
 
diff --git a/kernel/src/context/local.rs b/kernel/src/context/local.rs
index d54f3719c..dcfed1dfe 100644
--- a/kernel/src/context/local.rs
+++ b/kernel/src/context/local.rs
@@ -14,8 +14,6 @@ use core::ops::Deref;
 pub struct CpuLocal<T>(Vec<T>);
 
 impl<T> CpuLocal<T> {
-    /// # Context safety
-    /// This function does not require a CPU context on **stage 1** heap as long as `f` does not.
     pub fn new(mut f: impl FnMut(usize) -> T) -> Self {
         let len = config().max_cpu.get();
         let mut vec = Vec::with_capacity(len);
diff --git a/kernel/src/context/mod.rs b/kernel/src/context/mod.rs
index ca3a0a0e5..b3154abce 100644
--- a/kernel/src/context/mod.rs
+++ b/kernel/src/context/mod.rs
@@ -3,10 +3,12 @@ pub use self::arch::*;
 pub use self::local::*;
 
 use crate::proc::{ProcMgr, Thread};
+use crate::uma::Uma;
 use alloc::rc::Rc;
 use alloc::sync::Arc;
 use core::marker::PhantomData;
 use core::mem::offset_of;
+use core::ptr::null;
 use core::sync::atomic::Ordering;
 
 mod arc;
@@ -15,33 +17,32 @@ mod arc;
 mod arch;
 mod local;
 
-/// See `pcpu_init` on the PS4 for a reference.
-///
-/// # Context safety
-/// This function does not require a CPU context.
+/// See `pcpu_init` on the Orbis for a reference.
 ///
 /// # Safety
 /// - This function can be called only once per CPU.
 /// - `cpu` must be unique and valid.
-/// - `pmgr` must be the same for all context.
+/// - `setup` must return the same objects for all context.
 ///
-/// # Panics
-/// If `f` return. The reason we don't use `!` for a return type of `F` because it requires nightly
-/// Rust.
-pub unsafe fn run_with_context<R, F: FnOnce() -> R>(
+/// # Reference offsets
+/// | Version | Offset |
+/// |---------|--------|
+/// |PS4 11.00|0x08DA70|
+pub unsafe fn run_with_context(
     cpu: usize,
     td: Arc<Thread>,
-    pmgr: Arc<ProcMgr>,
     args: ContextArgs,
-    f: F,
+    setup: fn() -> ContextSetup,
+    main: fn() -> !,
 ) -> ! {
-    // We use a different mechanism here. The PS4 put all of pcpu at a global level but we put it on
-    // each CPU stack instead.
+    // We use a different mechanism here. The Orbis put all of pcpu at a global level but we put it
+    // on each CPU stack instead.
     let mut cx = Context::new(
         Base {
             cpu,
             thread: Arc::into_raw(td),
-            pmgr: Arc::into_raw(pmgr),
+            uma: null(),
+            pmgr: null(),
         },
         args,
     );
@@ -51,9 +52,13 @@ pub unsafe fn run_with_context<R, F: FnOnce() -> R>(
     // Prevent any code before and after this line to cross this line.
     core::sync::atomic::fence(Ordering::AcqRel);
 
-    f();
+    // Setup.
+    let r = setup();
+
+    cx.base.uma = Arc::into_raw(r.uma);
+    cx.base.pmgr = Arc::into_raw(r.pmgr);
 
-    panic!("return from a function passed to run_with_context() is not supported");
+    main();
 }
 
 /// # Interrupt safety
@@ -61,19 +66,33 @@ pub unsafe fn run_with_context<R, F: FnOnce() -> R>(
 pub fn current_thread() -> BorrowedArc<Thread> {
     // It does not matter if we are on a different CPU after we load the Context::thread because it
     // is going to be the same one since it represent the current thread.
-    unsafe { BorrowedArc::new(Context::load_fixed_ptr::<{ current_thread_offset() }, _>()) }
+    unsafe {
+        BorrowedArc::from_non_null(Context::load_static_ptr::<{ current_thread_offset() }, _>())
+    }
 }
 
 pub const fn current_thread_offset() -> usize {
     offset_of!(Base, thread)
 }
 
+/// Returns [`None`] if called from context setup function.
+///
 /// # Interrupt safety
-/// This function is interrupt safe.
-pub fn current_procmgr() -> BorrowedArc<ProcMgr> {
+/// This function can be called from interrupt handler.
+pub fn current_uma() -> Option<BorrowedArc<Uma>> {
+    // It does not matter if we are on a different CPU after we load the Context::uma because it is
+    // always the same for all CPU.
+    unsafe { BorrowedArc::new(Context::load_ptr::<{ offset_of!(Base, uma) }, _>()) }
+}
+
+/// Returns [`None`] if called from context setup function.
+///
+/// # Interrupt safety
+/// This function can be called from interrupt handle.
+pub fn current_procmgr() -> Option<BorrowedArc<ProcMgr>> {
     // It does not matter if we are on a different CPU after we load the Context::pmgr because it is
     // always the same for all CPU.
-    unsafe { BorrowedArc::new(Context::load_fixed_ptr::<{ offset_of!(Base, pmgr) }, _>()) }
+    unsafe { BorrowedArc::new(Context::load_ptr::<{ offset_of!(Base, pmgr) }, _>()) }
 }
 
 /// Pin the calling thread to one CPU.
@@ -96,6 +115,12 @@ pub fn pin_cpu() -> PinnedContext {
     }
 }
 
+/// Output of the context setup function.
+pub struct ContextSetup {
+    pub uma: Arc<Uma>,
+    pub pmgr: Arc<ProcMgr>,
+}
+
 /// Implementation of `pcpu` structure.
 ///
 /// Access to this structure must be done by **atomic reading or writing its field directly**. It is
@@ -118,6 +143,7 @@ pub fn pin_cpu() -> PinnedContext {
 struct Base {
     cpu: usize,            // pc_cpuid
     thread: *const Thread, // pc_curthread
+    uma: *const Uma,
     pmgr: *const ProcMgr,
 }
 
@@ -142,7 +168,7 @@ impl PinnedContext {
     /// Anything that derive from the returned value will invalid when this [`PinnedContext`]
     /// dropped.
     pub unsafe fn cpu(&self) -> usize {
-        Context::load_usize::<{ offset_of!(Base, cpu) }>()
+        Context::load_volatile_usize::<{ offset_of!(Base, cpu) }>()
     }
 }
 
diff --git a/kernel/src/context/x86_64.rs b/kernel/src/context/x86_64.rs
index 0ee8ee594..1dcf24b6b 100644
--- a/kernel/src/context/x86_64.rs
+++ b/kernel/src/context/x86_64.rs
@@ -19,9 +19,9 @@ pub struct ContextArgs {
 /// Extended [Base] for x86-64.
 #[repr(C)]
 pub(super) struct Context {
-    base: Base,        // Must be first field.
-    trap_rsp: *mut u8, // pc_rsp0
-    user_rsp: usize,   // pc_scratch_rsp
+    pub base: Base,        // Must be first field.
+    pub trap_rsp: *mut u8, // pc_rsp0
+    pub user_rsp: usize,   // pc_scratch_rsp
 }
 
 impl Context {
@@ -50,7 +50,7 @@ impl Context {
         wrmsr(0xc0000102, 0);
     }
 
-    pub unsafe fn load_fixed_ptr<const O: usize, T>() -> *const T {
+    pub unsafe fn load_static_ptr<const O: usize, T>() -> *const T {
         let mut v;
 
         asm!(
@@ -63,7 +63,20 @@ impl Context {
         v
     }
 
-    pub unsafe fn load_usize<const O: usize>() -> usize {
+    pub unsafe fn load_ptr<const O: usize, T>() -> *const T {
+        let mut v;
+
+        asm!(
+            "mov {out}, gs:[{off}]",
+            off = const O,
+            out = out(reg) v,
+            options(pure, readonly, preserves_flags, nostack)
+        );
+
+        v
+    }
+
+    pub unsafe fn load_volatile_usize<const O: usize>() -> usize {
         let mut v;
 
         asm!(
diff --git a/kernel/src/event/mod.rs b/kernel/src/event/mod.rs
index 9eda4072c..78fb8bbc1 100644
--- a/kernel/src/event/mod.rs
+++ b/kernel/src/event/mod.rs
@@ -19,9 +19,6 @@ impl<S> EventSet<S> {
 }
 
 impl<S: Default> Default for EventSet<S> {
-    /// # Context safety
-    /// This function does not require a CPU context as long as [`Default`] implementation on `S`
-    /// does not.
     fn default() -> Self {
         Self(Mutex::default())
     }
@@ -35,8 +32,6 @@ pub struct Event<T: EventType> {
 }
 
 impl<T: EventType> Default for Event<T> {
-    /// # Context safety
-    /// This function does not require a CPU context.
     fn default() -> Self {
         Self {
             subscribers: BTreeMap::new(),
diff --git a/kernel/src/lock/mutex/mod.rs b/kernel/src/lock/mutex/mod.rs
index 50c79a319..5e6fa50fc 100644
--- a/kernel/src/lock/mutex/mod.rs
+++ b/kernel/src/lock/mutex/mod.rs
@@ -15,9 +15,6 @@ pub struct Mutex<T> {
 
 impl<T> Mutex<T> {
     /// See `mtx_init` on the PS4 for a reference.
-    ///
-    /// # Context safety
-    /// This function does not require a CPU context.
     pub const fn new(data: T) -> Self {
         Self {
             data: UnsafeCell::new(data),
@@ -83,9 +80,6 @@ impl<T> Mutex<T> {
 }
 
 impl<T: Default> Default for Mutex<T> {
-    /// # Context safety
-    /// This function does not require a CPU context as long as [`Default`] implementation on `T`
-    /// does not.
     fn default() -> Self {
         Self::new(T::default())
     }
diff --git a/kernel/src/main.rs b/kernel/src/main.rs
index 756cbb0d4..9e8d491d5 100644
--- a/kernel/src/main.rs
+++ b/kernel/src/main.rs
@@ -1,13 +1,12 @@
 #![no_std]
 #![cfg_attr(not(test), no_main)]
 
-use self::context::current_procmgr;
+use self::context::{current_procmgr, ContextSetup};
 use self::imgact::Ps4Abi;
-use self::malloc::{KernelHeap, Stage2};
+use self::malloc::KernelHeap;
 use self::proc::{Fork, Proc, ProcAbi, ProcMgr, Thread};
 use self::sched::sleep;
 use self::uma::Uma;
-use alloc::boxed::Box;
 use alloc::sync::Arc;
 use core::mem::zeroed;
 use core::panic::PanicInfo;
@@ -45,7 +44,7 @@ extern crate alloc;
 /// 2. Interrupt is disabled.
 /// 3. Only main CPU can execute this function.
 ///
-/// See PS4 kernel entry point for a reference.
+/// See Orbis kernel entry point for a reference.
 #[allow(dead_code)]
 #[cfg_attr(target_os = "none", no_mangle)]
 unsafe extern "C" fn _start(env: &'static BootEnv, conf: &'static Config) -> ! {
@@ -65,21 +64,24 @@ unsafe extern "C" fn _start(env: &'static BootEnv, conf: &'static Config) -> ! {
     let proc0 = Arc::new(proc0);
     let thread0 = Thread::new_bare(proc0);
 
-    // Initialize foundations.
-    let uma = Uma::new();
-    let pmgr = ProcMgr::new();
-
     // Activate CPU context.
     let thread0 = Arc::new(thread0);
 
-    self::context::run_with_context(0, thread0, pmgr, cx, move || main(uma));
+    self::context::run_with_context(0, thread0, cx, setup, main);
+}
+
+fn setup() -> ContextSetup {
+    let uma = Uma::new();
+    let pmgr = ProcMgr::new();
+
+    ContextSetup { uma, pmgr }
 }
 
-fn main(mut uma: Uma) -> ! {
+fn main() -> ! {
     // Activate stage 2 heap.
     info!("Activating stage 2 heap.");
 
-    unsafe { KERNEL_HEAP.activate_stage2(Box::new(Stage2::new(&mut uma))) };
+    unsafe { KERNEL_HEAP.activate_stage2() };
 
     // Run sysinit vector. The PS4 use linker to put all sysinit functions in a list then loop the
     // list to execute all of it. We manually execute those functions instead for readability. This
@@ -91,7 +93,7 @@ fn main(mut uma: Uma) -> ! {
 
 /// See `create_init` function on the PS4 for a reference.
 fn create_init() {
-    let pmgr = current_procmgr();
+    let pmgr = current_procmgr().unwrap();
     let abi = Arc::new(Ps4Abi);
     let flags = Fork::new().with_copy_fd(true).with_create_process(true);
 
@@ -103,7 +105,7 @@ fn create_init() {
 /// See `scheduler` function on the PS4 for a reference.
 fn swapper() -> ! {
     // TODO: Subscribe to "system_suspend_phase2_pre_sync" and "system_resume_phase2" event.
-    let procs = current_procmgr();
+    let procs = current_procmgr().unwrap();
 
     loop {
         // TODO: Implement a call to vm_page_count_min().
diff --git a/kernel/src/malloc/mod.rs b/kernel/src/malloc/mod.rs
index 0c8db7909..478408be4 100644
--- a/kernel/src/malloc/mod.rs
+++ b/kernel/src/malloc/mod.rs
@@ -1,4 +1,4 @@
-pub use self::stage2::Stage2;
+use self::stage2::Stage2;
 use crate::lock::Mutex;
 use alloc::boxed::Box;
 use core::alloc::{GlobalAlloc, Layout};
@@ -39,16 +39,20 @@ impl KernelHeap {
 
     /// # Safety
     /// This must be called by main CPU and can be called only once.
-    pub unsafe fn activate_stage2(&self, stage2: Box<Stage2>) {
-        // What we are going here is highly unsafe. Do not edit this code unless you know what you
-        // are doing!
+    pub unsafe fn activate_stage2(&self) {
+        // Setup stage 2 using stage 1 heap.
+        let stage2 = Box::new(Stage2::new());
+
+        // What we are doing here is highly unsafe. Do not edit the code after this unless you know
+        // what you are doing!
         let stage = self.stage.get();
         let stage1 = match stage.read() {
             Stage::One(v) => Mutex::new(v.into_inner()),
             Stage::Two(_, _) => unreachable_unchecked(),
         };
 
-        // Switch to stage 2 WITHOUT dropping the value contained in Stage::One.
+        // Switch to stage 2 WITHOUT dropping the value contained in Stage::One. What we did here is
+        // moving the value from Stage::One to Stage::Two.
         stage.write(Stage::Two(stage2, stage1));
     }
 }
diff --git a/kernel/src/malloc/stage2.rs b/kernel/src/malloc/stage2.rs
index 634110efd..84261b4f2 100644
--- a/kernel/src/malloc/stage2.rs
+++ b/kernel/src/malloc/stage2.rs
@@ -1,6 +1,6 @@
 use crate::config::PAGE_SIZE;
-use crate::context::{current_thread, CpuLocal};
-use crate::uma::{Uma, UmaFlags, UmaZone};
+use crate::context::{current_thread, current_uma, CpuLocal};
+use crate::uma::{UmaFlags, UmaZone};
 use alloc::string::ToString;
 use alloc::sync::Arc;
 use alloc::vec::Vec;
@@ -29,10 +29,11 @@ impl Stage2 {
     /// | Version | Offset |
     /// |---------|--------|
     /// |PS4 11.00|0x1A4B80|
-    pub fn new(uma: &mut Uma) -> Self {
+    pub fn new() -> Self {
         // The possible of maximum alignment that Layout allowed is a bit before the most
         // significant bit of isize (e.g. 0x4000000000000000 on 64 bit system). So we can use
         // "size_of::<usize>() * 8 - 1" to get the size of array for all possible alignment.
+        let uma = current_uma().unwrap();
         let zones = core::array::from_fn(|align| {
             let mut zones = Vec::with_capacity(Self::KMEM_ZSIZE + 1);
             let mut last = 0;
diff --git a/kernel/src/proc/mod.rs b/kernel/src/proc/mod.rs
index f343aca07..741dddb1c 100644
--- a/kernel/src/proc/mod.rs
+++ b/kernel/src/proc/mod.rs
@@ -25,8 +25,6 @@ pub struct ProcMgr {
 }
 
 impl ProcMgr {
-    /// # Context safety
-    /// This function does not require a CPU context.
     pub fn new() -> Arc<Self> {
         let events = Arc::default();
 
diff --git a/kernel/src/uma/keg.rs b/kernel/src/uma/keg.rs
index b16ca695c..7ad344cc4 100644
--- a/kernel/src/uma/keg.rs
+++ b/kernel/src/uma/keg.rs
@@ -4,9 +4,6 @@ pub struct UmaKeg {}
 impl UmaKeg {
     /// See `keg_ctor` on the Orbis for a reference.
     ///
-    /// # Context safety
-    /// This function does not require a CPU context on **stage 1** heap.
-    ///
     /// # Reference offsets
     /// | Version | Offset |
     /// |---------|--------|
diff --git a/kernel/src/uma/mod.rs b/kernel/src/uma/mod.rs
index c962a9606..69ef5d758 100644
--- a/kernel/src/uma/mod.rs
+++ b/kernel/src/uma/mod.rs
@@ -1,6 +1,7 @@
 pub use self::zone::*;
 
 use alloc::string::String;
+use alloc::sync::Arc;
 use bitfield_struct::bitfield;
 use core::num::NonZero;
 
@@ -15,28 +16,22 @@ impl Uma {
     /// See `uma_startup` on the Orbis for a reference. Beware that our implementation cannot access
     /// the CPU context due to this function can be called before context activation.
     ///
-    /// # Context safety
-    /// This function does not require a CPU context on **stage 1** heap.
-    ///
     /// # Reference offsets
     /// | Version | Offset |
     /// |---------|--------|
     /// |PS4 11.00|0x13CA70|
-    pub fn new() -> Self {
-        Self {}
+    pub fn new() -> Arc<Self> {
+        Arc::new(Self {})
     }
 
     /// See `uma_zcreate` on the Orbis for a reference.
     ///
-    /// # Context safety
-    /// This function does not require a CPU context on **stage 1** heap.
-    ///
     /// # Reference offsets
     /// | Version | Offset |
     /// |---------|--------|
     /// |PS4 11.00|0x13DC80|
     pub fn create_zone(
-        &mut self,
+        &self,
         name: impl Into<String>,
         size: NonZero<usize>,
         align: Option<usize>,
diff --git a/kernel/src/uma/zone.rs b/kernel/src/uma/zone.rs
index 9750ba405..5d6119349 100644
--- a/kernel/src/uma/zone.rs
+++ b/kernel/src/uma/zone.rs
@@ -25,9 +25,6 @@ impl UmaZone {
 
     /// See `zone_ctor` on Orbis for a reference.
     ///
-    /// # Context safety
-    /// This function does not require a CPU context on **stage 1** heap.
-    ///
     /// # Reference offsets
     /// | Version | Offset |
     /// |---------|--------|