1 files changed, 11 insertions, 11 deletions

diff --git a/x86_64/src/lib.rs b/x86_64/src/lib.rs
index 9c960d7..c866e56 100644
--- a/x86_64/src/lib.rs
+++ b/x86_64/src/lib.rs
@@ -313,13 +313,13 @@ impl arch::LinuxArch for X8664arch {
         ) -> std::result::Result<Vec<(Box<dyn PciDevice>, Option<Minijail>)>, E>,
         E: StdError + 'static,
     {
-        let mut resources =
-            Self::get_resource_allocator(components.memory_size, components.wayland_dmabuf);
         let has_bios = match components.vm_image {
             VmImage::Bios(_) => true,
             _ => false,
         };
         let mem = Self::setup_memory(components.memory_size, has_bios)?;
+        let mut resources = Self::get_resource_allocator(&mem, components.wayland_dmabuf);
+
         let kvm = Kvm::new().map_err(Error::CreateKvm)?;
         let mut vm = Self::create_vm(&kvm, split_irqchip, mem.clone())?;
 
@@ -578,17 +578,17 @@ impl X8664arch {
         Ok(None)
     }
 
-    /// This returns the first page frame number for use by the balloon driver.
+    /// This returns the start address of device memory
     ///
     /// # Arguments
     ///
-    /// * `mem_size` - the size in bytes of physical ram for the guest
-    fn get_base_dev_pfn(mem_size: u64) -> u64 {
+    /// * mem: The memory to be used by the guest
+    fn get_dev_memory_base(mem: &GuestMemory) -> u64 {
         // Put device memory at a 2MB boundary after physical memory or 4gb, whichever is greater.
-        const MB: u64 = 1024 * 1024;
-        const GB: u64 = 1024 * MB;
-        let mem_size_round_2mb = (mem_size + 2 * MB - 1) / (2 * MB) * (2 * MB);
-        std::cmp::max(mem_size_round_2mb, 4 * GB) / sys_util::pagesize() as u64
+        const MB: u64 = 1 << 20;
+        const GB: u64 = 1 << 30;
+        let ram_end_round_2mb = (mem.end_addr().offset() + 2 * MB - 1) / (2 * MB) * (2 * MB);
+        std::cmp::max(ram_end_round_2mb, 4 * GB)
     }
 
     /// This returns a minimal kernel command for this architecture
@@ -604,9 +604,9 @@ impl X8664arch {
     }
 
     /// Returns a system resource allocator.
-    fn get_resource_allocator(mem_size: u64, gpu_allocation: bool) -> SystemAllocator {
+    fn get_resource_allocator(mem: &GuestMemory, gpu_allocation: bool) -> SystemAllocator {
         const MMIO_BASE: u64 = 0xe0000000;
-        let device_addr_start = Self::get_base_dev_pfn(mem_size) * sys_util::pagesize() as u64;
+        let device_addr_start = Self::get_dev_memory_base(mem);
         SystemAllocator::builder()
             .add_io_addresses(0xc000, 0x10000)
             .add_mmio_addresses(MMIO_BASE, 0x100000)