diff --git a/guest/rialto/src/main.rs b/guest/rialto/src/main.rs
index 456af7f..8095a1f 100644
--- a/guest/rialto/src/main.rs
+++ b/guest/rialto/src/main.rs
@@ -32,8 +32,6 @@
 use core::num::NonZeroUsize;
 use core::slice;
 use diced_open_dice::{bcc_handover_parse, DiceArtifacts};
-use hypervisor_backends::get_mem_sharer;
-use libfdt::FdtError;
 use log::{debug, error, info};
 use service_vm_comm::{ServiceVmRequest, VmType};
 use service_vm_fake_chain::service_vm;
@@ -50,7 +48,10 @@
     generate_image_header,
     layout::{self, crosvm},
     main,
-    memory::{MemoryTracker, PageTable, MEMORY, PAGE_SIZE, SIZE_128KB},
+    memory::{
+        init_shared_pool, map_rodata, map_rodata_outside_main_memory, resize_available_memory,
+        switch_to_dynamic_page_tables, PageTable, PAGE_SIZE, SIZE_128KB,
+    },
     power::reboot,
     virtio::{
         pci::{self, PciTransportIterator, VirtIOSocket},
@@ -91,60 +92,40 @@
     info!("Welcome to Rialto!");
     let page_table = new_page_table()?;
 
-    MEMORY.lock().replace(MemoryTracker::new(
-        page_table,
-        crosvm::MEM_START..layout::MAX_VIRT_ADDR,
-        crosvm::MMIO_RANGE,
-    ));
+    switch_to_dynamic_page_tables(page_table);
 
-    let fdt_range = MEMORY
-        .lock()
-        .as_mut()
-        .unwrap()
-        .alloc(fdt_addr, NonZeroUsize::new(crosvm::FDT_MAX_SIZE).unwrap())?;
+    let fdt_size = NonZeroUsize::new(crosvm::FDT_MAX_SIZE).unwrap();
+    map_rodata(fdt_addr, fdt_size)?;
     // SAFETY: The tracker validated the range to be in main memory, mapped, and not overlap.
-    let fdt = unsafe { slice::from_raw_parts(fdt_range.start as *mut u8, fdt_range.len()) };
+    let fdt = unsafe { slice::from_raw_parts(fdt_addr as *mut u8, fdt_size.into()) };
     // We do not need to validate the DT since it is already validated in pvmfw.
     let fdt = libfdt::Fdt::from_slice(fdt)?;
 
     let memory_range = fdt.first_memory_range()?;
-    MEMORY.lock().as_mut().unwrap().shrink(&memory_range).inspect_err(|_| {
+    resize_available_memory(&memory_range).inspect_err(|_| {
         error!("Failed to use memory range value from DT: {memory_range:#x?}");
     })?;
 
-    if let Some(mem_sharer) = get_mem_sharer() {
-        let granule = mem_sharer.granule()?;
-        MEMORY.lock().as_mut().unwrap().init_dynamic_shared_pool(granule).inspect_err(|_| {
-            error!("Failed to initialize dynamically shared pool.");
-        })?;
-    } else if let Ok(Some(swiotlb_info)) = SwiotlbInfo::new_from_fdt(fdt) {
-        let range = swiotlb_info.fixed_range().ok_or_else(|| {
-            error!("Pre-shared pool range not specified in swiotlb node");
-            Error::from(FdtError::BadValue)
-        })?;
-        MEMORY.lock().as_mut().unwrap().init_static_shared_pool(range).inspect_err(|_| {
-            error!("Failed to initialize pre-shared pool.");
-        })?;
-    } else {
-        info!("No MEM_SHARE capability detected or swiotlb found: allocating buffers from heap.");
-        MEMORY.lock().as_mut().unwrap().init_heap_shared_pool().inspect_err(|_| {
-            error!("Failed to initialize heap-based pseudo-shared pool.");
-        })?;
-    }
+    let swiotlb_range = SwiotlbInfo::new_from_fdt(fdt)
+        .inspect_err(|_| {
+            error!("Rialto failed when access swiotlb");
+        })?
+        .and_then(|info| info.fixed_range());
+    init_shared_pool(swiotlb_range).inspect_err(|_| {
+        error!("Failed to initialize shared pool.");
+    })?;
 
     let bcc_handover: Box<dyn DiceArtifacts> = match vm_type(fdt)? {
         VmType::ProtectedVm => {
             let dice_range = read_dice_range_from(fdt)?;
             info!("DICE range: {dice_range:#x?}");
-            // SAFETY: This region was written by pvmfw in its writable_data region. The region
-            // has no overlap with the main memory region and is safe to be mapped as read-only
-            // data.
-            let res = unsafe {
-                MEMORY.lock().as_mut().unwrap().alloc_range_outside_main_memory(&dice_range)
-            };
-            res.inspect_err(|_| {
-                error!("Failed to use DICE range from DT: {dice_range:#x?}");
-            })?;
+            let dice_size = dice_range.len().try_into().unwrap();
+            // SAFETY: The DICE memory region has been generated by pvmfw and doesn't overlap.
+            unsafe { map_rodata_outside_main_memory(dice_range.start, dice_size) }.inspect_err(
+                |_| {
+                    error!("Failed to use DICE range from DT: {dice_range:#x?}");
+                },
+            )?;
             let dice_start = dice_range.start as *const u8;
             // SAFETY: There's no memory overlap and the region is mapped as read-only data.
             let bcc_handover = unsafe { slice::from_raw_parts(dice_start, dice_range.len()) };
@@ -157,8 +138,7 @@
 
     let pci_info = PciInfo::from_fdt(fdt)?;
     debug!("PCI: {pci_info:#x?}");
-    let mut pci_root = pci::initialize(pci_info, MEMORY.lock().as_mut().unwrap())
-        .map_err(Error::PciInitializationFailed)?;
+    let mut pci_root = pci::initialize(pci_info).map_err(Error::PciInitializationFailed)?;
     debug!("PCI root: {pci_root:#x?}");
     let socket_device = find_socket_device::<HalImpl>(&mut pci_root)?;
     debug!("Found socket device: guest cid = {:?}", socket_device.guest_cid());