feat: refactor memory management and integrate memblock

- Updated physical memory management (PMM) to utilize memblock for memory layout and frame allocation.
- Enhanced the virtual memory manager (VMM) to allocate the frame map from memblock, ensuring it starts at the correct virtual address without relocation.
- Improved frame allocation and deallocation functions to handle physical addresses more consistently.
- Adjusted the slab allocator to correctly convert between physical and virtual addresses.
- Refined page fault handling to ensure proper frame allocation and address translation.
- Streamlined the initialization process for PMM and VMM, ensuring a more robust setup of memory structures.

Co-authored-by: Copilot <copilot@github.com>

6 files changed, 735 insertions, 425 deletions

diff --git a/na/include/lume/memblock.h b/na/include/lume/memblock.h
new file mode 100644
index 0000000..41baf34
--- /dev/null
+++ b/na/include/lume/memblock.h
@@ -0,0 +1,53 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+/* Copyright (C) 2026 Richard Qin */
+
+#pragma once
+
+/*
+ * memblock — Boot-time physical memory allocator
+ *
+ * A simple region-based allocator active from early boot until
+ * the buddy system (PMM) takes over.  Inspired by Linux's memblock.
+ *
+ * Lifetime:
+ *   1. memblock_init()       — discovers RAM via FDT, reserves kernel image
+ *   2. vmm_init()            — uses memblock_alloc for page table pages & frame_map
+ *   3. pmm_init()            — absorbs remaining free regions into buddy
+ *   4. memblock_retire()     — marks memblock inactive; all alloc via PMM
+ *
+ * This eliminates the need for post-init pointer relocation when
+ * rebinding frame_map to a VMEMMAP virtual address.
+ */
+
+#include <lume/types.h>
+
+inline constexpr int kMemblockMaxRegions = 128;
+
+/* Initialize memblock: scan FDT for RAM, reserve kernel image + firmware. */
+void memblock_init(uint64 fdt_paddr);
+
+/* Mark a physical range as reserved (will not be handed to buddy). */
+void memblock_reserve(uint64 base, uint64 size);
+
+/* Allocate `size` bytes aligned to `align` from top of free memory.
+ * Returns PA, or 0 on failure.  Used for frame_map. */
+uint64 memblock_alloc_top(uint64 size, uint64 align);
+
+/* Allocate `size` bytes aligned to `align` from bottom of free memory.
+ * Returns PA, or 0 on failure.  Used for page table pages. */
+uint64 memblock_alloc(uint64 size, uint64 align);
+
+/* Query memory layout (same info as FDT). */
+uint64 memblock_get_mem_base();
+uint64 memblock_get_mem_size();
+
+/* Is memblock still the active allocator? */
+bool memblock_is_active();
+
+/* Mark memblock as retired (PMM takes over). */
+void memblock_retire();
+
+/* Iterate all free (unreserved) regions in ascending order.
+ * Used by pmm_init() to feed free pages into the buddy system. */
+void memblock_for_each_free(void (*cb)(uint64 base, uint64 size, void* ctx),
+                            void* ctx);
diff --git a/na/mm/memblock.cc b/na/mm/memblock.cc
new file mode 100644
index 0000000..ac71472
--- /dev/null
+++ b/na/mm/memblock.cc
@@ -0,0 +1,217 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+/* Copyright (C) 2026 Richard Qin */
+
+/*
+ * memblock.cc — Boot-time physical memory allocator
+ *
+ * Simple region tracker: one memory region (from FDT) plus an array
+ * of reserved regions.  Allocation searches for gaps between reserved
+ * regions.  All operations are O(n_reserved), which is fine for the
+ * ~20 allocations that happen during early boot.
+ *
+ * Reference: inspired by Linux kernel's mm/memblock.c
+ */
+
+#include <lume/memblock.h>
+#include <lume/fdt.h>
+#include <lume/klog.h>
+#include <lume/kprintf.h>
+#include <lume/addr.h>
+
+extern "C" char _stext[], _kernel_end[];
+
+/* ------------------------------------------------------------------ */
+/*  Module state                                                       */
+/* ------------------------------------------------------------------ */
+
+static struct {
+    uint64 mem_base;
+    uint64 mem_size;
+
+    struct {
+        uint64 base;
+        uint64 size;
+    } reserved[kMemblockMaxRegions];
+    int    reserved_count;
+    bool   active;
+} g_mb;
+
+/* ------------------------------------------------------------------ */
+/*  Internal helpers                                                    */
+/* ------------------------------------------------------------------ */
+
+/* Insert a reserved region, keeping the array sorted by base. */
+static void insert_reserved(uint64 base, uint64 size)
+{
+    if (g_mb.reserved_count >= kMemblockMaxRegions) {
+        kernel_panic("memblock: too many reserved regions");
+    }
+
+    int i = g_mb.reserved_count;
+    while (i > 0 && g_mb.reserved[i - 1].base > base) {
+        g_mb.reserved[i] = g_mb.reserved[i - 1];
+        i--;
+    }
+    g_mb.reserved[i] = {base, size};
+    g_mb.reserved_count++;
+}
+
+/* Check if [base, base+size) overlaps any reserved region. */
+static bool overlaps_reserved(uint64 base, uint64 size)
+{
+    uint64 end = base + size;
+    for (int i = 0; i < g_mb.reserved_count; i++) {
+        uint64 rb = g_mb.reserved[i].base;
+        uint64 re = rb + g_mb.reserved[i].size;
+        if (base < re && end > rb) {
+            return true;
+        }
+    }
+    return false;
+}
+
+/* Find the end of the first reserved region that covers `addr`. */
+static uint64 reserved_end_covering(uint64 addr)
+{
+    for (int i = 0; i < g_mb.reserved_count; i++) {
+        uint64 rb = g_mb.reserved[i].base;
+        uint64 re = rb + g_mb.reserved[i].size;
+        if (addr >= rb && addr < re) {
+            return re;
+        }
+    }
+    return addr;
+}
+
+/* ------------------------------------------------------------------ */
+/*  Public API                                                         */
+/* ------------------------------------------------------------------ */
+
+void memblock_init(uint64 fdt_paddr)
+{
+    (void)fdt_paddr;
+
+    /* 1. Get memory layout from FDT */
+    uint64 mem_base = 0;
+    uint64 mem_size = 0;
+    fdt_early_get_mem_info(&mem_base, &mem_size);
+    if (mem_size == 0) {
+        kernel_panic("memblock_init: FDT reports zero memory");
+    }
+
+    g_mb.mem_base = mem_base;
+    g_mb.mem_size = mem_size;
+    g_mb.reserved_count = 0;
+    g_mb.active = true;
+
+    /* 2. Reserve firmware region [mem_base, kernel_start).
+     *    OpenSBI occupies the first 2MB of RAM on RISC-V virt. */
+    uint64 ks_pa = va_to_pa(virt_addr(
+        reinterpret_cast<uint64>(_stext))).raw;
+    if (ks_pa > mem_base) {
+        memblock_reserve(mem_base, ks_pa - mem_base);
+    }
+
+    /* 3. Reserve kernel image [_stext, _kernel_end) */
+    uint64 ke_pa = va_to_pa(virt_addr(
+        reinterpret_cast<uint64>(_kernel_end))).raw;
+    memblock_reserve(ks_pa, ke_pa - ks_pa);
+
+    kprintf("[memblock] RAM [0x%llx, 0x%llx) = %llu MB\n",
+            mem_base, mem_base + mem_size, mem_size / (1024 * 1024));
+    kprintf("[memblock] kernel reserved [0x%llx, 0x%llx)\n", ks_pa, ke_pa);
+}
+
+void memblock_reserve(uint64 base, uint64 size)
+{
+    uint64 aligned_base = page_align_down(base);
+    uint64 aligned_end  = page_align_up(base + size);
+    insert_reserved(aligned_base, aligned_end - aligned_base);
+}
+
+uint64 memblock_alloc_top(uint64 size, uint64 align)
+{
+    uint64 mem_end = g_mb.mem_base + g_mb.mem_size;
+
+    /* Search from top of RAM downward */
+    uint64 end   = page_align_down(mem_end);
+    uint64 start = (end - size) & ~(align - 1);
+
+    while (start >= g_mb.mem_base && start + size <= mem_end) {
+        if (!overlaps_reserved(start, size)) {
+            memblock_reserve(start, size);
+            return start;
+        }
+        if (start < align) {
+            break;
+        }
+        start -= align;
+    }
+
+    kernel_panic("memblock_alloc_top: out of memory");
+    return 0;
+}
+
+uint64 memblock_alloc(uint64 size, uint64 align)
+{
+    uint64 mem_end = g_mb.mem_base + g_mb.mem_size;
+    uint64 start   = page_align_up(g_mb.mem_base);
+
+    while (start + size <= mem_end) {
+        start = (start + align - 1) & ~(align - 1);
+        if (start + size > mem_end) {
+            break;
+        }
+
+        if (!overlaps_reserved(start, size)) {
+            memblock_reserve(start, size);
+            return start;
+        }
+
+        /* Skip past the overlapping reserved region */
+        uint64 skip = reserved_end_covering(start);
+        if (skip <= start) {
+            start += align;
+        } else {
+            start = skip;
+        }
+    }
+
+    kernel_panic("memblock_alloc: out of memory");
+    return 0;
+}
+
+uint64 memblock_get_mem_base() { return g_mb.mem_base; }
+uint64 memblock_get_mem_size() { return g_mb.mem_size; }
+
+bool memblock_is_active() { return g_mb.active; }
+void memblock_retire()    { g_mb.active = false; }
+
+void memblock_for_each_free(void (*cb)(uint64 base, uint64 size, void* ctx),
+                            void* ctx)
+{
+    uint64 pos     = g_mb.mem_base;
+    uint64 mem_end = g_mb.mem_base + g_mb.mem_size;
+
+    for (int i = 0; i < g_mb.reserved_count && pos < mem_end; i++) {
+        uint64 rb = g_mb.reserved[i].base;
+        if (rb < pos) {
+            /* Overlapping/adjacent — skip */
+            uint64 re = g_mb.reserved[i].base + g_mb.reserved[i].size;
+            if (re > pos) {
+                pos = re;
+            }
+            continue;
+        }
+        if (rb > pos) {
+            /* Free gap [pos, rb) */
+            cb(pos, rb - pos, ctx);
+        }
+        pos = g_mb.reserved[i].base + g_mb.reserved[i].size;
+    }
+
+    /* Trailing free region after last reserved */
+    if (pos < mem_end) {
+        cb(pos, mem_end - pos, ctx);
+    }
+}
diff --git a/na/mm/pmm.cc b/na/mm/pmm.cc
index 00f998c..1400448 100644
--- a/na/mm/pmm.cc
+++ b/na/mm/pmm.cc
@@ -19,43 +19,42 @@
  * Reference: docs/specs/pmm.md
  */
 
-#include <lume/pmm.h>
+#include <arch/config.h>
+#include <arch/cpu.h>
+#include <kernel/sync/spinlock.h>
 #include <lume/addr.h>
 #include <lume/config.h>
-#include <lume/fdt.h>
-#include <lume/new.h>
 #include <lume/klog.h>
+#include <lume/kprintf.h>
+#include <lume/memblock.h>
+#include <lume/new.h>
+#include <lume/pmm.h>
 #include <lume/types.h>
-#include <arch/cpu.h>
-#include <kernel/sync/spinlock.h>
-
-/* ------------------------------------------------------------------ */
-/*  Linker symbols                                                    */
-/* ------------------------------------------------------------------ */
-
-extern "C" char _kernel_end[];
+#include <lume/vmm.h>
 
 /* ------------------------------------------------------------------ */
 /*  Per-CPU page cache                                                */
 /* ------------------------------------------------------------------ */
 
 struct PerCpuCache {
-    list_node free_list;   // Single-frame free list
-    uint32 count;          // Number of cached frames
+  list_node free_list; // Single-frame free list
+  uint32 count;        // Number of cached frames
 };
 
 /* ------------------------------------------------------------------ */
 /*  Module-private state                                              */
 /* ------------------------------------------------------------------ */
 
-static Spinlock        g_pmm_lock;
-static list_node       g_free_areas[kMaxOrder];   // Buddy free lists [0..10]
-static PerCpuCache     g_pcp[kMaxCpus];
+static Spinlock g_pmm_lock;
+static list_node g_free_areas[kMaxOrder]; // Buddy free lists [0..10]
+static PerCpuCache g_pcp[kMaxCpus];
 
-static Frame*          g_frame_map = nullptr;      // Flat array of Frame descriptors
-static uint64          g_mem_base  = 0;            // Physical memory base (from FDT)
-static uint64          g_mem_size  = 0;            // Total physical memory size
-static uint64          g_num_frames = 0;           // Total managed frames
+static Frame *g_frame_map = nullptr; // Flat array of Frame descriptors
+static uint64 g_frame_map_pa = 0;    // Physical base of frame_map storage
+static uint64 g_frame_map_bytes = 0; // Byte size of frame_map storage
+static uint64 g_mem_base = 0;        // Physical memory base (from FDT)
+static uint64 g_mem_size = 0;        // Total physical memory size
+static uint64 g_num_frames = 0;      // Total managed frames
 
 uint32 g_pcp_high_watermark = 0;
 uint32 g_pcp_batch_size = 0;
@@ -68,33 +67,28 @@ uint32 g_pcp_batch_size = 0;
  * Physical Frame Number: index into the frame_map array.
  * PFN = (pa - g_mem_base) / kPageSize
  */
-static inline uint64 pa_to_frame_index(uint64 pa)
-{
-    return (pa - g_mem_base) / kPageSize;
+static inline uint64 pa_to_frame_index(uint64 pa) {
+  return (pa - g_mem_base) / kPageSize;
 }
 
-static inline uint64 frame_index_to_pa(uint64 idx)
-{
-    return g_mem_base + idx * kPageSize;
+static inline uint64 frame_index_to_pa(uint64 idx) {
+  return g_mem_base + idx * kPageSize;
 }
 
-uint64 frame_to_pa(const Frame* frame)
-{
-    uint64 idx = static_cast<uint64>(frame - g_frame_map);
-    if (frame < g_frame_map || idx >= g_num_frames)
-        kernel_panic("frame_to_pa: frame outside frame_map bounds");
-    return frame_index_to_pa(idx);
+PhysAddr frame_to_pa(const Frame *frame) {
+  uint64 idx = static_cast<uint64>(frame - g_frame_map);
+  if (frame < g_frame_map || idx >= g_num_frames)
+    kernel_panic("frame_to_pa: frame outside frame_map bounds");
+  return phys_addr(frame_index_to_pa(idx));
 }
 
-Frame* pa_to_frame(uint64 pa)
-{
-    uint64 idx = pa_to_frame_index(pa);
-    return &g_frame_map[idx];
+Frame *pa_to_frame(PhysAddr pa) {
+  uint64 idx = pa_to_frame_index(pa.raw);
+  return &g_frame_map[idx];
 }
 
-uint64 frame_to_va(const Frame* frame)
-{
-    return pa_to_va(frame_to_pa(frame));
+VirtAddr frame_to_va(const Frame *frame) {
+  return pa_to_va(frame_to_pa(frame));
 }
 
 /* ------------------------------------------------------------------ */
@@ -105,40 +99,37 @@ uint64 frame_to_va(const Frame* frame)
  * Find a frame's buddy at a given order.
  * Buddy PFN = PFN ^ (1 << order).
  */
-static inline Frame* buddy_of(Frame* frame, uint8 order)
-{
-    uint64 pfn = static_cast<uint64>(frame - g_frame_map);
-    uint64 buddy_pfn = pfn ^ (1ULL << order);
-    if (buddy_pfn >= g_num_frames)
-        return nullptr;
-    return &g_frame_map[buddy_pfn];
+static inline Frame *buddy_of(Frame *frame, uint8 order) {
+  uint64 pfn = static_cast<uint64>(frame - g_frame_map);
+  uint64 buddy_pfn = pfn ^ (1ULL << order);
+  if (buddy_pfn >= g_num_frames)
+    return nullptr;
+  return &g_frame_map[buddy_pfn];
 }
 
 /*
  * Insert a free block into the buddy free list at the given order.
  */
-static void buddy_list_add(Frame* frame, uint8 order)
-{
-    frame->state = FrameState::Free;
-    frame->order = order;
-    list_node* head = &g_free_areas[order];
-    list_node* node = &frame->free.free_link;
-    node->next = head->next;
-    node->prev = head;
-    head->next->prev = node;
-    head->next = node;
+static void buddy_list_add(Frame *frame, uint8 order) {
+  frame->state = FrameState::Free;
+  frame->order = order;
+  list_node *head = &g_free_areas[order];
+  list_node *node = &frame->free.free_link;
+  node->next = head->next;
+  node->prev = head;
+  head->next->prev = node;
+  head->next = node;
 }
 
 /*
  * Remove a block from whatever buddy free list it is on.
  */
-static void buddy_list_del(Frame* frame)
-{
-    list_node* node = &frame->free.free_link;
-    node->prev->next = node->next;
-    node->next->prev = node->prev;
-    node->next = node;
-    node->prev = node;
+static void buddy_list_del(Frame *frame) {
+  list_node *node = &frame->free.free_link;
+  node->prev->next = node->next;
+  node->next->prev = node->prev;
+  node->next = node;
+  node->prev = node;
 }
 
 /*
@@ -146,66 +137,64 @@ static void buddy_list_del(Frame* frame)
  * Splits higher-order blocks if no exact match is available.
  * Returns nullptr if no memory available.
  */
-static Frame* buddy_alloc(uint8 order)
-{
-    for (int cur = order; cur < kMaxOrder; cur++) {
-        list_node* head = &g_free_areas[cur];
-        if (head->next == head)
-            continue; // Empty at this order
-
-        // Pop the first block from this order's free list
-        list_node* node = head->next;
-        // Compute Frame* from the free_link member offset
-        Frame* frame = reinterpret_cast<Frame*>(
-            reinterpret_cast<uintptr>(node) -
-            __builtin_offsetof(Frame, free.free_link));
-        buddy_list_del(frame);
-
-        // Split down to the requested order
-        while (cur > order) {
-            cur--;
-            // The upper half becomes a free buddy
-            Frame* split = &g_frame_map[(frame - g_frame_map) + (1ULL << cur)];
-            buddy_list_add(split, static_cast<uint8>(cur));
-        }
-
-        frame->state = FrameState::Allocated;
-        frame->order = order;
-        frame->refcount.store(1, __ATOMIC_RELAXED);
-        return frame;
+static Frame *buddy_alloc(uint8 order) {
+  for (int cur = order; cur < kMaxOrder; cur++) {
+    list_node *head = &g_free_areas[cur];
+    if (head->next == head)
+      continue; // Empty at this order
+
+    // Pop the first block from this order's free list
+    list_node *node = head->next;
+    // Compute Frame* from the free_link member offset
+    Frame *frame =
+        reinterpret_cast<Frame *>(reinterpret_cast<uintptr>(node) -
+                                  __builtin_offsetof(Frame, free.free_link));
+    buddy_list_del(frame);
+
+    // Split down to the requested order
+    while (cur > order) {
+      cur--;
+      // The upper half becomes a free buddy
+      Frame *split = &g_frame_map[(frame - g_frame_map) + (1ULL << cur)];
+      buddy_list_add(split, static_cast<uint8>(cur));
     }
-    return nullptr;
+
+    frame->state = FrameState::Allocated;
+    frame->order = order;
+    frame->refcount.store(1, __ATOMIC_RELAXED);
+    return frame;
+  }
+  return nullptr;
 }
 
 /*
  * Return a block of 2^order frames to the buddy system.
  * Coalesces with its buddy recursively up to kMaxOrder-1.
  */
-static void buddy_free(Frame* frame, uint8 order)
-{
-    uint64 pfn = static_cast<uint64>(frame - g_frame_map);
-
-    while (order < kMaxOrder - 1) {
-        Frame* buddy = buddy_of(frame, order);
-        if (!buddy)
-            break;
-        // Buddy must be free AND at the same order to coalesce
-        if (buddy->state != FrameState::Free || buddy->order != order)
-            break;
-
-        // Remove buddy from its free list and merge
-        buddy_list_del(buddy);
-
-        // The merged block starts at the lower PFN
-        uint64 buddy_pfn = static_cast<uint64>(buddy - g_frame_map);
-        if (buddy_pfn < pfn) {
-            frame = buddy;
-            pfn = buddy_pfn;
-        }
-        order++;
+static void buddy_free(Frame *frame, uint8 order) {
+  uint64 pfn = static_cast<uint64>(frame - g_frame_map);
+
+  while (order < kMaxOrder - 1) {
+    Frame *buddy = buddy_of(frame, order);
+    if (!buddy)
+      break;
+    // Buddy must be free AND at the same order to coalesce
+    if (buddy->state != FrameState::Free || buddy->order != order)
+      break;
+
+    // Remove buddy from its free list and merge
+    buddy_list_del(buddy);
+
+    // The merged block starts at the lower PFN
+    uint64 buddy_pfn = static_cast<uint64>(buddy - g_frame_map);
+    if (buddy_pfn < pfn) {
+      frame = buddy;
+      pfn = buddy_pfn;
     }
+    order++;
+  }
 
-    buddy_list_add(frame, order);
+  buddy_list_add(frame, order);
 }
 
 /* ------------------------------------------------------------------ */
@@ -216,166 +205,157 @@ static void buddy_free(Frame* frame, uint8 order)
  * Refill the PCP from the global buddy system.
  * Acquires g_pmm_lock internally.
  */
-static void refill_pcp(PerCpuCache* pcp)
-{
-    LockGuard guard(g_pmm_lock);
-    for (uint32 i = 0; i < g_pcp_batch_size; i++) {
-        Frame* f = buddy_alloc(0);
-        if (!f)
-            break;
-        // Put onto PCP list (reuse the free_link for PCP chain)
-        f->state = FrameState::PcpCached;
-        list_node* node = &f->free.free_link;
-        node->next = pcp->free_list.next;
-        node->prev = &pcp->free_list;
-        pcp->free_list.next->prev = node;
-        pcp->free_list.next = node;
-        pcp->count++;
-    }
+static void refill_pcp(PerCpuCache *pcp) {
+  LockGuard guard(g_pmm_lock);
+  for (uint32 i = 0; i < g_pcp_batch_size; i++) {
+    Frame *f = buddy_alloc(0);
+    if (!f)
+      break;
+    // Put onto PCP list (reuse the free_link for PCP chain)
+    f->state = FrameState::PcpCached;
+    list_node *node = &f->free.free_link;
+    node->next = pcp->free_list.next;
+    node->prev = &pcp->free_list;
+    pcp->free_list.next->prev = node;
+    pcp->free_list.next = node;
+    pcp->count++;
+  }
 }
 
 /*
  * Drain half the PCP back to the global buddy system.
  * Acquires g_pmm_lock internally.
  */
-static void drain_pcp(PerCpuCache* pcp)
-{
-    LockGuard guard(g_pmm_lock);
-    int to_drain = pcp->count / 2;
-    for (int i = 0; i < to_drain; i++) {
-        list_node* head = &pcp->free_list;
-        if (head->next == head)
-            break;
-        list_node* node = head->next;
-        // Unlink from PCP
-        node->prev->next = node->next;
-        node->next->prev = node->prev;
-        pcp->count--;
-
-        Frame* f = reinterpret_cast<Frame*>(
-            reinterpret_cast<uintptr>(node) -
-            __builtin_offsetof(Frame, free.free_link));
-        buddy_free(f, 0);
-    }
+static void drain_pcp(PerCpuCache *pcp) {
+  LockGuard guard(g_pmm_lock);
+  int to_drain = pcp->count / 2;
+  for (int i = 0; i < to_drain; i++) {
+    list_node *head = &pcp->free_list;
+    if (head->next == head)
+      break;
+    list_node *node = head->next;
+    // Unlink from PCP
+    node->prev->next = node->next;
+    node->next->prev = node->prev;
+    pcp->count--;
+
+    Frame *f =
+        reinterpret_cast<Frame *>(reinterpret_cast<uintptr>(node) -
+                                  __builtin_offsetof(Frame, free.free_link));
+    buddy_free(f, 0);
+  }
 }
 
 /* ------------------------------------------------------------------ */
 /*  Public API: single-frame (fast path via PCP)                      */
 /* ------------------------------------------------------------------ */
 
-Frame* pmm_alloc_frame()
-{
-    /*
-     * Disable interrupts to prevent reentry on the same CPU
-     * (e.g. timer IRQ handler calling pmm_alloc_frame while we
-     * are mid-way through manipulating the PCP list).
-     */
-    bool was_on = arch::cpu::intr_enabled();
-    arch::cpu::intr_off();
-
-    uint64 cpu = arch::cpu::id();
-    PerCpuCache* pcp = &g_pcp[cpu];
-
-    if (pcp->count == 0)
-        refill_pcp(pcp);
-
-    Frame* frame = nullptr;
-    if (pcp->count > 0) {
-        list_node* node = pcp->free_list.next;
-        // Unlink
-        node->prev->next = node->next;
-        node->next->prev = node->prev;
-        pcp->count--;
-
-        frame = reinterpret_cast<Frame*>(
-            reinterpret_cast<uintptr>(node) -
-            __builtin_offsetof(Frame, free.free_link));
-
-        // Validate the frame is within the frame_map bounds
-        if (frame < g_frame_map || frame >= g_frame_map + g_num_frames) {
-            kernel_panic("pmm_alloc_frame: PCP returned out-of-bounds frame");
-        }
+Frame *pmm_alloc_frame() {
+  /*
+   * Disable interrupts to prevent reentry on the same CPU
+   * (e.g. timer IRQ handler calling pmm_alloc_frame while we
+   * are mid-way through manipulating the PCP list).
+   */
+  bool was_on = arch::cpu::intr_enabled();
+  arch::cpu::intr_off();
+
+  uint64 cpu = arch::cpu::id();
+  PerCpuCache *pcp = &g_pcp[cpu];
+
+  if (pcp->count == 0)
+    refill_pcp(pcp);
+
+  Frame *frame = nullptr;
+  if (pcp->count > 0) {
+    list_node *node = pcp->free_list.next;
+    // Unlink
+    node->prev->next = node->next;
+    node->next->prev = node->prev;
+    pcp->count--;
+
+    frame =
+        reinterpret_cast<Frame *>(reinterpret_cast<uintptr>(node) -
+                                  __builtin_offsetof(Frame, free.free_link));
 
-        frame->state = FrameState::Allocated;
-        frame->order = 0;
-        frame->refcount.store(1, __ATOMIC_RELAXED);
+    // Validate the frame is within the frame_map bounds
+    if (frame < g_frame_map || frame >= g_frame_map + g_num_frames) {
+      kernel_panic("pmm_alloc_frame: PCP returned out-of-bounds frame");
     }
 
-    if (was_on)
-        arch::cpu::intr_on();
-    return frame;
+    frame->state = FrameState::Allocated;
+    frame->order = 0;
+    frame->refcount.store(1, __ATOMIC_RELAXED);
+  }
+
+  if (was_on)
+    arch::cpu::intr_on();
+  return frame;
 }
 
-void pmm_free_frame(Frame* frame)
-{
-    if (!frame)
-        kernel_panic("pmm_free_frame: null frame");
+void pmm_free_frame(Frame *frame) {
+  if (!frame)
+    kernel_panic("pmm_free_frame: null frame");
 
-    bool was_on = arch::cpu::intr_enabled();
-    arch::cpu::intr_off();
+  bool was_on = arch::cpu::intr_enabled();
+  arch::cpu::intr_off();
 
-    uint64 cpu = arch::cpu::id();
-    PerCpuCache* pcp = &g_pcp[cpu];
+  uint64 cpu = arch::cpu::id();
+  PerCpuCache *pcp = &g_pcp[cpu];
 
-    // Push onto PCP
-    frame->state = FrameState::PcpCached;
-    frame->order = 0;
-    list_node* node = &frame->free.free_link;
-    node->next = pcp->free_list.next;
-    node->prev = &pcp->free_list;
-    pcp->free_list.next->prev = node;
-    pcp->free_list.next = node;
-    pcp->count++;
+  // Push onto PCP
+  frame->state = FrameState::PcpCached;
+  frame->order = 0;
+  list_node *node = &frame->free.free_link;
+  node->next = pcp->free_list.next;
+  node->prev = &pcp->free_list;
+  pcp->free_list.next->prev = node;
+  pcp->free_list.next = node;
+  pcp->count++;
 
-    // Drain if over high watermark
-    if (pcp->count >= g_pcp_high_watermark)
-        drain_pcp(pcp);
+  // Drain if over high watermark
+  if (pcp->count >= g_pcp_high_watermark)
+    drain_pcp(pcp);
 
-    if (was_on)
-        arch::cpu::intr_on();
+  if (was_on)
+    arch::cpu::intr_on();
 }
 
 /* ------------------------------------------------------------------ */
 /*  Public API: multi-frame (slow path, direct buddy)                 */
 /* ------------------------------------------------------------------ */
 
-Frame* pmm_alloc_frames(uint8 order)
-{
-    if (order >= kMaxOrder)
-        return nullptr;
-    LockGuard guard(g_pmm_lock);
-    return buddy_alloc(order);
+Frame *pmm_alloc_frames(uint8 order) {
+  if (order >= kMaxOrder)
+    return nullptr;
+  LockGuard guard(g_pmm_lock);
+  return buddy_alloc(order);
 }
 
-void pmm_free_frames(Frame* frame, uint8 order)
-{
-    if (!frame)
-        kernel_panic("pmm_free_frames: null frame");
-    if (order >= kMaxOrder)
-        kernel_panic("pmm_free_frames: order out of range");
-    LockGuard guard(g_pmm_lock);
-    buddy_free(frame, order);
+void pmm_free_frames(Frame *frame, uint8 order) {
+  if (!frame)
+    kernel_panic("pmm_free_frames: null frame");
+  if (order >= kMaxOrder)
+    kernel_panic("pmm_free_frames: order out of range");
+  LockGuard guard(g_pmm_lock);
+  buddy_free(frame, order);
 }
 
 /* ------------------------------------------------------------------ */
 /*  Public API: reference counting                                    */
 /* ------------------------------------------------------------------ */
 
-void frame_incref(Frame* frame)
-{
-    frame->refcount.fetch_add(1);
-}
-
-void frame_decref(Frame* frame)
-{
-    uint32 old = frame->refcount.fetch_sub(1);
-    if (old == 1) {
-        // Refcount dropped to 0 — return to pool based on originally allocated order
-        if (frame->order == 0)
-            pmm_free_frame(frame);
-        else
-            pmm_free_frames(frame, frame->order);
-    }
+void frame_incref(Frame *frame) { frame->refcount.fetch_add(1); }
+
+void frame_decref(Frame *frame) {
+  uint32 old = frame->refcount.fetch_sub(1);
+  if (old == 1) {
+    // Refcount dropped to 0 — return to pool based on originally allocated
+    // order
+    if (frame->order == 0)
+      pmm_free_frame(frame);
+    else
+      pmm_free_frames(frame, frame->order);
+  }
 }
 
 /* ------------------------------------------------------------------ */
@@ -386,110 +366,104 @@ void frame_decref(Frame* frame)
  * Helper: zero out a range of bytes.  We cannot use memset here
  * because lib/string may not be linked yet.
  */
-static void zero_range(void* start, uint64 bytes)
-{
-    uint8* p = static_cast<uint8*>(start);
-    for (uint64 i = 0; i < bytes; i++)
-        p[i] = 0;
+static void zero_range(void *start, uint64 bytes) {
+  uint8 *p = static_cast<uint8 *>(start);
+  for (uint64 i = 0; i < bytes; i++)
+    p[i] = 0;
 }
 
-void pmm_init()
-{
-    /* 1. Query FDT for physical memory layout */
-    uint64 mem_base = 0, mem_size = 0;
-    fdt_early_get_mem_info(&mem_base, &mem_size);
-
-    if (mem_size == 0)
-        kernel_panic("pmm_init: FDT reports zero memory");
-
-    g_mem_base = mem_base;
-    g_mem_size = mem_size;
-
-    uint64 mem_end = mem_base + mem_size;
-    g_num_frames = mem_size / kPageSize;
-
-    /*
-     * Dynamically calculate PCP parameters based on memory size.
-     * Rule of thumb: batch_size = max(16, total_pages / (cpu_count * 1024))
-     * This avoids depleting global memory on small RAM systems.
-     */
-    uint32 calc_batch = static_cast<uint32>(g_num_frames / (kMaxCpus * 1024));
-    g_pcp_batch_size = (calc_batch > 16) ? calc_batch : 16;
-    g_pcp_high_watermark = g_pcp_batch_size * 2;
-
-    /*
-     * 2. Place the frame_map array.
-     *
-     * The frame_map lives right after the kernel image (_kernel_end),
-     * page-aligned.  We compute its size and mark that region as used.
-     *
-     * _kernel_end is a virtual address; convert to physical for arithmetic,
-     * then back to virtual for the actual pointer.
-     */
-    /*
-     * _kernel_end is linked at a virtual address (0xFFFFFFC0...),
-     * but if the kernel is still running via identity mapping (before
-     * vmm_init removes it), PC-relative `la` resolves to the physical
-     * address instead.  Detect which case we're in.
-     */
-    uint64 kernel_end_raw = reinterpret_cast<uint64>(_kernel_end);
-    uint64 kernel_end_pa = ensure_pa(kernel_end_raw);
-    uint64 frame_map_pa = (kernel_end_pa + kPageSize - 1) & ~(kPageSize - 1);
-    uint64 frame_map_bytes = g_num_frames * sizeof(Frame);
-    uint64 frame_map_end_pa = (frame_map_pa + frame_map_bytes + kPageSize - 1)
-                              & ~(kPageSize - 1);
-
-    if (frame_map_end_pa > mem_end)
-        kernel_panic("pmm_init: not enough memory for frame_map");
-
-    g_frame_map = reinterpret_cast<Frame*>(pa_to_va(frame_map_pa));
-
-    /* 3. Zero out the entire frame_map */
-    zero_range(g_frame_map, frame_map_bytes);
-
-    /* 4. Initialize buddy free list heads */
-    g_pmm_lock.init("pmm");
-    for (int i = 0; i < kMaxOrder; i++)
-        g_free_areas[i].init();
-
-    /* 5. Initialize PCP list heads */
-    for (int i = 0; i < kMaxCpus; i++) {
-        g_pcp[i].free_list.init();
-        g_pcp[i].count = 0;
-    }
-
-    /*
-     * 6. Feed all free pages into the buddy system.
-     *
-     * Pages before frame_map_end (kernel + frame_map) are reserved.
-     * We iterate from the first free page to the end of memory,
-     * inserting at the highest possible order for efficiency.
-     */
-    uint64 free_start_pfn = pa_to_frame_index(frame_map_end_pa);
-    uint64 end_pfn = g_num_frames;
-
-    uint64 pfn = free_start_pfn;
-    while (pfn < end_pfn) {
-        // Find the largest order block that is naturally aligned at this PFN
-        uint8 order = 0;
-        for (int o = kMaxOrder - 1; o >= 0; o--) {
+void pmm_init() {
+  /* 1. Get memory layout from memblock (already scanned FDT) */
+  g_mem_base = memblock_get_mem_base();
+  g_mem_size = memblock_get_mem_size();
+  g_num_frames = vmm_get_num_frames();
+
+  if (g_mem_size == 0) {
+    kernel_panic("pmm_init: memblock reports zero memory");
+  }
+
+  /* 2. PCP parameters */
+  uint32 calc_batch = static_cast<uint32>(g_num_frames / (kMaxCpus * 1024));
+  g_pcp_batch_size = (calc_batch > 16) ? calc_batch : 16;
+  g_pcp_high_watermark = g_pcp_batch_size * 2;
+
+  /* 3. frame_map was allocated by vmm_init at vmemmap VA.
+   *    No relocation needed — frame_map starts life at the right VA. */
+  g_frame_map_pa = vmm_get_frame_map_pa();
+  g_frame_map_bytes = vmm_get_frame_map_size();
+  g_frame_map = reinterpret_cast<Frame *>(arch::kVmemmapBase);
+
+  /* 4. Zero out the entire frame_map */
+  zero_range(g_frame_map, g_frame_map_bytes);
+
+  /* 5. Initialize buddy free list heads */
+  g_pmm_lock.init("pmm");
+  for (int i = 0; i < kMaxOrder; i++) {
+    g_free_areas[i].init();
+  }
+
+  /* 6. Initialize PCP list heads */
+  for (int i = 0; i < kMaxCpus; i++) {
+    g_pcp[i].free_list.init();
+    g_pcp[i].count = 0;
+  }
+
+  /* 7. Feed free regions from memblock into the buddy system.
+   *    memblock knows about all reserved regions (firmware, kernel,
+   *    page table pages, frame_map).  Everything else is free. */
+  struct FeedCtx {
+    uint64 total_pages;
+  };
+  FeedCtx ctx{0};
+
+  memblock_for_each_free(
+      [](uint64 base, uint64 size, void *arg) {
+        auto *c = static_cast<FeedCtx *>(arg);
+        uint64 start_pfn = (base - memblock_get_mem_base()) / kPageSize;
+        uint64 end_pfn = start_pfn + size / kPageSize;
+
+        /* Access frame_map via the extern g_frame_map (already at vmemmap VA)
+         */
+        extern Frame *g_frame_map;
+        extern list_node g_free_areas[];
+
+        uint64 pfn = start_pfn;
+        while (pfn < end_pfn) {
+          uint8 order = 0;
+          for (int o = kMaxOrder - 1; o >= 0; o--) {
             uint64 block_size = 1ULL << o;
             if ((pfn & (block_size - 1)) == 0 && pfn + block_size <= end_pfn) {
-                order = static_cast<uint8>(o);
-                break;
+              order = static_cast<uint8>(o);
+              break;
             }
+          }
+
+          Frame *f = &g_frame_map[pfn];
+          f->state = FrameState::Free;
+          f->order = order;
+          f->refcount.store(0, __ATOMIC_RELAXED);
+          f->free.free_link.init();
+
+          /* Inline buddy_list_add since we're inside a lambda */
+          list_node *head = &g_free_areas[order];
+          list_node *node = &f->free.free_link;
+          node->next = head->next;
+          node->prev = head;
+          head->next->prev = node;
+          head->next = node;
+
+          pfn += (1ULL << order);
+          c->total_pages += (1ULL << order);
         }
+      },
+      &ctx);
 
-        Frame* f = &g_frame_map[pfn];
-        f->state = FrameState::Free;
-        f->order = order;
-        f->refcount.store(0, __ATOMIC_RELAXED);
-        f->free.free_link.init();
-        buddy_list_add(f, order);
+  kprintf("[pmm] buddy: %llu free pages (%llu MB)\n", ctx.total_pages,
+          (ctx.total_pages * kPageSize) / (1024 * 1024));
 
-        pfn += (1ULL << order);
-    }
+  /* 8. Retire memblock — all future allocations go through PMM */
+  memblock_retire();
 }
 
-uint64 pmm_get_mem_base() { return g_mem_base; }
+PhysAddr pmm_get_mem_base() { return phys_addr(g_mem_base); }
 uint64 pmm_get_mem_size() { return g_mem_size; }
diff --git a/na/mm/slab.cc b/na/mm/slab.cc
index ef15b6a..98029aa 100644
--- a/na/mm/slab.cc
+++ b/na/mm/slab.cc
@@ -36,7 +36,7 @@
  */
 static inline void* frame_to_page_va(Frame* f)
 {
-    return reinterpret_cast<void*>(pa_to_va(frame_to_pa(f)));
+    return reinterpret_cast<void*>(pa_to_va(frame_to_pa(f)).raw);
 }
 
 /* Round up x to the next multiple of align (align must be power of 2) */
@@ -193,9 +193,9 @@ void KmemCache::free(void* obj)
     uint64 obj_addr = reinterpret_cast<uint64>(obj);
 
     /* Convert VA to PA. Like pmm.cc, handle both PA and VA cases. */
-    uint64 obj_pa = ensure_pa(obj_addr);
+    uint64 obj_pa = va_to_pa(virt_addr(obj_addr)).raw;
     uint64 page_pa = obj_pa & ~(kPageSize - 1);
-    Frame* f = pa_to_frame(page_pa);
+    Frame* f = pa_to_frame(phys_addr(page_pa));
 
     if (f->state != FrameState::Slab || f->slab.cache != this)
         kernel_panic("slab free: object does not belong to this cache");
@@ -255,7 +255,7 @@ void* kmalloc(uint32 size)
         Frame* f = pmm_alloc_frames(order);
         if (!f)
             return nullptr;
-        return reinterpret_cast<void*>(pa_to_va(frame_to_pa(f)));
+        return reinterpret_cast<void*>(pa_to_va(frame_to_pa(f)).raw);
     }
 
     KmemCache* cache = find_cache(size);
@@ -270,9 +270,9 @@ void kfree(void* ptr)
         return;
 
     uint64 addr = reinterpret_cast<uint64>(ptr);
-    uint64 pa = ensure_pa(addr);
+    uint64 pa = va_to_pa(virt_addr(addr)).raw;
     uint64 page_pa = pa & ~(kPageSize - 1);
-    Frame* f = pa_to_frame(page_pa);
+    Frame* f = pa_to_frame(phys_addr(page_pa));
 
     if (f->state == FrameState::Slab) {
         KmemCache* cache = f->slab.cache;
diff --git a/na/mm/vma.cc b/na/mm/vma.cc
index e27853e..02f92b0 100644
--- a/na/mm/vma.cc
+++ b/na/mm/vma.cc
@@ -158,7 +158,7 @@ int vmm_handle_page_fault(VmSpace* space, uint64 fault_addr, uint64 cause) {
     Frame* f = pmm_alloc_frame();
     if (!f) return -ENOMEM;
 
-    uint64 pa = frame_to_pa(f);
+    uint64 pa = frame_to_pa(f).raw;
 
     /* Translate MI permission flags to hardware PTE bits (MD layer) */
     uint64 pte_perm = pmap::vm_perm_to_pte(vma->perm);
diff --git a/na/mm/vmm.cc b/na/mm/vmm.cc
index 50c7920..041413d 100644
--- a/na/mm/vmm.cc
+++ b/na/mm/vmm.cc
@@ -4,24 +4,30 @@
 /*
  * vmm.cc — Virtual Memory Manager (Phase 1: Kernel Page Table)
  *
- * Builds a fine-grained kernel page table (4KB pages, precise
- * permissions per section), activates it, and retires the early_pgdir
- * identity mapping from entry.S.
+ * Builds a fine-grained kernel page table, activates it, and retires
+ * the early_pgdir identity mapping from entry.S.
  *
- * This module is machine-independent (MI).  All architecture-specific
- * page table manipulation goes through the pmap:: namespace.
+ * With memblock integration, the initialization order is:
+ *   memblock_init → vmm_init (allocates frame_map from memblock,
+ *   maps vmemmap, activates page table) → pmm_init (buddy system
+ *   starts with frame_map already at vmemmap VA — no relocation).
  *
  * Reference: docs/specs/vmm.md
  */
 
-#include <lume/types.h>
+#include <arch/config.h>
+#include <arch/mmu.h>
+#include <arch/pmap.h>
 #include <lume/addr.h>
-#include <lume/pmm.h>
-#include <lume/klog.h>
+#include <lume/console.h>
 #include <lume/fdt.h>
+#include <lume/frame.h>
+#include <lume/klog.h>
+#include <lume/kprintf.h>
+#include <lume/memblock.h>
+#include <lume/pmm.h>
+#include <lume/types.h>
 #include <lume/vmm.h>
-#include <arch/pmap.h>
-#include <arch/mmu.h>
 
 /* Linker-exported section boundaries (virtual addresses) */
 extern "C" char _stext[], _etext[];
@@ -33,92 +39,152 @@ extern "C" char _kernel_end[];
 /* Global kernel page table root PA */
 static uint64 g_kernel_pgtbl = 0;
 
-/* Convert a linker symbol to PA. Handles both VA and PA cases. */
-static inline uint64 sym_to_pa(const void *sym)
-{
-    return ensure_pa(reinterpret_cast<uint64>(sym));
+/* VMM exports for PMM: frame_map location after vmemmap mapping. */
+static uint64 g_vmemmap_frame_map_pa = 0;
+static uint64 g_vmemmap_frame_map_size = 0;
+static uint64 g_vmemmap_num_frames = 0;
+
+uint64 vmm_get_frame_map_pa() { return g_vmemmap_frame_map_pa; }
+uint64 vmm_get_frame_map_size() { return g_vmemmap_frame_map_size; }
+uint64 vmm_get_num_frames() { return g_vmemmap_num_frames; }
+
+/* Convert a linker symbol to PA (kernel is executing in higher-half VA). */
+static inline uint64 sym_to_pa(const void *sym) {
+  return va_to_pa(virt_addr(reinterpret_cast<uint64>(sym))).raw;
 }
 
-/* Map a PA range into the kernel page table with given permissions.
- * All addresses must be page-aligned. Panics on failure. */
-static void map_range(uint64 root, uint64 pa_start, uint64 pa_end, uint64 perm)
-{
-    for (uint64 pa = pa_start; pa < pa_end; pa += kPageSize) {
-        uint64 va = pa_to_va(pa);
-        int ret = pmap::map(root, va, pa, perm | pmap::PTE_G);
-        if (ret != 0)
-            kernel_panic("vmm_init: pmap::map failed");
+/* Map a PA range into the kernel page table via the direct map. */
+static void map_range(uint64 root, uint64 pa_start, uint64 pa_end,
+                      uint64 perm) {
+  for (uint64 pa = pa_start; pa < pa_end; pa += kPageSize) {
+    uint64 va = pa_to_va(phys_addr(pa)).raw;
+    int ret = pmap::map(root, va, pa, perm | pmap::PTE_G);
+    if (ret != 0) {
+      kernel_panic("vmm_init: pmap::map failed");
     }
+  }
 }
 
-
-
-void vmm_init()
-{
-    /* 1. Allocate new root page table */
-    uint64 root = pmap::create();
-    if (!root)
-        kernel_panic("vmm_init: failed to create root page table");
-
-    /* 2. Compute PA boundaries of each kernel section */
-    uint64 text_start  = page_align_down(sym_to_pa(_stext));
-    uint64 text_end    = page_align_up(sym_to_pa(_etext));
-    uint64 ro_start    = page_align_down(sym_to_pa(_srodata));
-    uint64 ro_end      = page_align_up(sym_to_pa(_erodata));
-    uint64 data_start  = page_align_down(sym_to_pa(_sdata));
-    uint64 bss_end     = page_align_up(sym_to_pa(_ebss));
-
-    /* Memory extent from PMM */
-    uint64 mem_base = pmm_get_mem_base();
-    uint64 mem_end  = mem_base + pmm_get_mem_size();
-
-    /* 3. Map kernel .text (RX) */
-    map_range(root, text_start, text_end,
-              pmap::PTE_R | pmap::PTE_X);
-
-    /* 4. Map .rodata (R) */
-    map_range(root, ro_start, ro_end,
-              pmap::PTE_R);
-
-    /* 5. Map .data + .bss (RW) */
-    map_range(root, data_start, bss_end,
-              pmap::PTE_R | pmap::PTE_W);
-
-    /* 6. Map remaining physical memory after kernel (RW)
-     * This covers frame_map, slab objects, and all free pages. */
-    uint64 after_kernel = page_align_up(sym_to_pa(_kernel_end));
-    if (after_kernel < mem_end) {
-        map_range(root, after_kernel, mem_end,
-                  pmap::PTE_R | pmap::PTE_W);
+/* Map a PA range to an explicit VA base, using 2MB pages where possible. */
+static void map_range_at_huge(uint64 root, uint64 va_base, uint64 pa_start,
+                              uint64 pa_end, uint64 perm) {
+  constexpr uint64 k2MB = 2ULL * 1024 * 1024;
+  uint64 pa = pa_start;
+  uint64 va = va_base;
+
+  /* Map 2MB-aligned chunks with mega pages */
+  while (pa + k2MB <= pa_end) {
+    if ((va & (k2MB - 1)) == 0 && (pa & (k2MB - 1)) == 0) {
+      int ret = pmap::map_2mb(root, va, pa, perm | pmap::PTE_G);
+      if (ret != 0) {
+        kernel_panic("vmm_init: pmap::map_2mb failed");
+      }
+      pa += k2MB;
+      va += k2MB;
+    } else {
+      break;
     }
-    /* 7. MMIO regions will be mapped later via vmm_map_kernel_mmio by architecture/drivers */
-
-    /* 8. Activate the new page table */
-    pmap::activate(root);
+  }
 
-    g_kernel_pgtbl = root;
+  /* Map remaining 4KB tail */
+  while (pa < pa_end) {
+    int ret = pmap::map(root, va, pa, perm | pmap::PTE_G);
+    if (ret != 0) {
+      kernel_panic("vmm_init: vmemmap 4KB pmap::map failed");
+    }
+    pa += kPageSize;
+    va += kPageSize;
+  }
 }
 
-void vmm_init_ap()
-{
-    /* AP just switches to the kernel page table built by BSP */
-    if (!g_kernel_pgtbl)
-        kernel_panic("vmm_init_ap: kernel page table not ready");
-    pmap::activate(g_kernel_pgtbl);
+void vmm_init() {
+  /* 1. Allocate new root page table (from memblock) */
+  uint64 root = pmap::create();
+  if (!root) {
+    kernel_panic("vmm_init: failed to create root page table");
+  }
+
+  /* 2. Compute PA boundaries of each kernel section */
+  uint64 text_start = page_align_down(sym_to_pa(_stext));
+  uint64 text_end = page_align_up(sym_to_pa(_etext));
+  uint64 ro_start = page_align_down(sym_to_pa(_srodata));
+  uint64 ro_end = page_align_up(sym_to_pa(_erodata));
+  uint64 data_start = page_align_down(sym_to_pa(_sdata));
+  uint64 bss_end = page_align_up(sym_to_pa(_ebss));
+
+  /* Memory extent from memblock */
+  uint64 mem_base = memblock_get_mem_base();
+  uint64 mem_size = memblock_get_mem_size();
+  uint64 mem_end = mem_base + mem_size;
+
+  /* 3. Map kernel .text (RX) */
+  map_range(root, text_start, text_end, pmap::PTE_R | pmap::PTE_X);
+
+  /* 4. Map .rodata (R) */
+  map_range(root, ro_start, ro_end, pmap::PTE_R);
+
+  /* 5. Map .data + .bss (RW) */
+  map_range(root, data_start, bss_end, pmap::PTE_R | pmap::PTE_W);
+
+  /* 6. Map remaining physical memory after kernel (RW) */
+  uint64 after_kernel = page_align_up(sym_to_pa(_kernel_end));
+  if (after_kernel < mem_end) {
+    map_range(root, after_kernel, mem_end, pmap::PTE_R | pmap::PTE_W);
+  }
+
+  /* 7. Allocate frame_map from memblock (top of RAM) and map at VMEMMAP.
+   *
+   * Since memblock is active (PMM not yet initialized), pmap::map
+   * allocates page table pages from memblock too.  frame_map starts
+   * life at the vmemmap VA — no relocation ever needed.
+   */
+  uint64 num_frames = mem_size / kPageSize;
+  uint64 frame_map_bytes = num_frames * sizeof(Frame);
+  uint64 frame_map_span = page_align_up(frame_map_bytes);
+  uint64 frame_map_pa = memblock_alloc_top(frame_map_span, kPageSize);
+
+  if (!frame_map_pa) {
+    kernel_panic("vmm_init: memblock failed to allocate frame_map");
+  }
+
+  kprintf("[vmm] vmemmap: %llu frames, %llu KB at PA 0x%llx\n", num_frames,
+          frame_map_span / 1024, frame_map_pa);
+
+  /* Map frame_map at kVmemmapBase, using 2MB pages where possible */
+  map_range_at_huge(root, arch::kVmemmapBase, frame_map_pa,
+                    frame_map_pa + frame_map_span, pmap::PTE_R | pmap::PTE_W);
+
+  /* Export for pmm_init() */
+  g_vmemmap_frame_map_pa = frame_map_pa;
+  g_vmemmap_frame_map_size = frame_map_bytes;
+  g_vmemmap_num_frames = num_frames;
+
+  /* 8. Pre-map console MMIO before page table switch */
+  g_kernel_pgtbl = root;
+  console_init();
+
+  /* 9. Activate the new page table */
+  pmap::activate(root);
+
+  kprintf("[vmm] vmm_init complete (vmemmap at 0x%llx)\n", arch::kVmemmapBase);
 }
 
-void vmm_map_kernel_mmio(uint64 pa, uint64 size)
-{
-    if (!g_kernel_pgtbl)
-        kernel_panic("vmm_map_kernel_mmio: kernel page table not ready");
-    
-    uint64 aligned_pa = page_align_down(pa);
-    uint64 aligned_size = page_align_up(pa + size) - aligned_pa;
-    
-    map_range(g_kernel_pgtbl, aligned_pa, aligned_pa + aligned_size,
-              pmap::PTE_R | pmap::PTE_W | pmap::PTE_G);
-    arch::mmu::flush_tlb_all();
+void vmm_init_ap() {
+  if (!g_kernel_pgtbl) {
+    kernel_panic("vmm_init_ap: kernel page table not ready");
+  }
+  pmap::activate(g_kernel_pgtbl);
 }
 
+void vmm_map_kernel_mmio(uint64 pa, uint64 size) {
+  if (!g_kernel_pgtbl) {
+    kernel_panic("vmm_map_kernel_mmio: kernel page table not ready");
+  }
 
+  uint64 aligned_pa = page_align_down(pa);
+  uint64 aligned_size = page_align_up(pa + size) - aligned_pa;
 
+  map_range(g_kernel_pgtbl, aligned_pa, aligned_pa + aligned_size,
+            pmap::PTE_R | pmap::PTE_W | pmap::PTE_G);
+  arch::mmu::flush_tlb_all();
+}