diff --git a/src/arch/src/aarch64/linux/regs.rs b/src/arch/src/aarch64/linux/regs.rs
index b932c9cc4..e5343c266 100644
--- a/src/arch/src/aarch64/linux/regs.rs
+++ b/src/arch/src/aarch64/linux/regs.rs
@@ -153,7 +153,7 @@ mod tests {
         let kvm = Kvm::new().unwrap();
         let vm = kvm.create_vm().unwrap();
         let vcpu = vm.create_vcpu(0).unwrap();
-        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000, 0);
+        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000, 0, None);
         let mem = GuestMemoryMmap::from_ranges(&regions).expect("Cannot initialize memory");
 
         match setup_regs(&vcpu, 0, 0x0, &mem).unwrap_err() {
diff --git a/src/arch/src/aarch64/mod.rs b/src/arch/src/aarch64/mod.rs
index 1a750e4f1..9ef4d2bdb 100644
--- a/src/arch/src/aarch64/mod.rs
+++ b/src/arch/src/aarch64/mod.rs
@@ -123,7 +123,7 @@ mod tests {
 
     #[test]
     fn test_regions_lt_1024gb() {
-        let (_mem_info, regions) = arch_memory_regions(1usize << 29, 0);
+        let (_mem_info, regions) = arch_memory_regions(1usize << 29, 0, None);
         assert_eq!(1, regions.len());
         assert_eq!(GuestAddress(super::layout::DRAM_MEM_START), regions[0].0);
         assert_eq!(1usize << 29, regions[0].1);
@@ -131,7 +131,7 @@ mod tests {
 
     #[test]
     fn test_regions_gt_1024gb() {
-        let (_mem_info, regions) = arch_memory_regions(1usize << 41, 0);
+        let (_mem_info, regions) = arch_memory_regions(1usize << 41, 0, None);
         assert_eq!(1, regions.len());
         assert_eq!(GuestAddress(super::layout::DRAM_MEM_START), regions[0].0);
         assert_eq!(super::layout::DRAM_MEM_MAX_SIZE, regions[0].1 as u64);
@@ -139,15 +139,15 @@ mod tests {
 
     #[test]
     fn test_get_fdt_addr() {
-        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE - 0x1000, 0);
+        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE - 0x1000, 0, None);
         let mem = GuestMemoryMmap::from_ranges(&regions).expect("Cannot initialize memory");
         assert_eq!(get_fdt_addr(&mem), layout::DRAM_MEM_START);
 
-        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE, 0);
+        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE, 0, None);
         let mem = GuestMemoryMmap::from_ranges(&regions).expect("Cannot initialize memory");
         assert_eq!(get_fdt_addr(&mem), layout::DRAM_MEM_START);
 
-        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000, 0);
+        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000, 0, None);
         let mem = GuestMemoryMmap::from_ranges(&regions).expect("Cannot initialize memory");
         assert_eq!(get_fdt_addr(&mem), 0x1000 + layout::DRAM_MEM_START);
     }
diff --git a/src/arch/src/riscv64/linux/regs.rs b/src/arch/src/riscv64/linux/regs.rs
index c99bad91e..2d91570c6 100644
--- a/src/arch/src/riscv64/linux/regs.rs
+++ b/src/arch/src/riscv64/linux/regs.rs
@@ -90,7 +90,7 @@ mod tests {
         let kvm = Kvm::new().unwrap();
         let vm = kvm.create_vm().unwrap();
         let vcpu = vm.create_vcpu(0).unwrap();
-        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000, 0);
+        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000, 0, None);
         let mem = GuestMemoryMmap::from_ranges(&regions).expect("Cannot initialize memory");
 
         match setup_regs(&vcpu, 0, 0x0, &mem).unwrap_err() {
diff --git a/src/arch/src/riscv64/mod.rs b/src/arch/src/riscv64/mod.rs
index ba47501c2..c51802cff 100644
--- a/src/arch/src/riscv64/mod.rs
+++ b/src/arch/src/riscv64/mod.rs
@@ -112,7 +112,7 @@ mod tests {
 
     #[test]
     fn test_regions_lt_1024gb() {
-        let (_mem_info, regions) = arch_memory_regions(1usize << 29, 0);
+        let (_mem_info, regions) = arch_memory_regions(1usize << 29, 0, None);
         assert_eq!(1, regions.len());
         assert_eq!(GuestAddress(super::layout::DRAM_MEM_START), regions[0].0);
         assert_eq!(1usize << 29, regions[0].1);
@@ -120,7 +120,7 @@ mod tests {
 
     #[test]
     fn test_regions_gt_1024gb() {
-        let (_mem_info, regions) = arch_memory_regions(1usize << 41, 0);
+        let (_mem_info, regions) = arch_memory_regions(1usize << 41, 0, None);
         assert_eq!(1, regions.len());
         assert_eq!(GuestAddress(super::layout::DRAM_MEM_START), regions[0].0);
         assert_eq!(super::layout::DRAM_MEM_MAX_SIZE, regions[0].1 as u64);
@@ -128,15 +128,15 @@ mod tests {
 
     #[test]
     fn test_get_fdt_addr() {
-        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE - 0x1000, 0);
+        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE - 0x1000, 0, None);
         let mem = GuestMemoryMmap::from_ranges(&regions).expect("Cannot initialize memory");
         assert_eq!(get_fdt_addr(&mem), layout::DRAM_MEM_START);
 
-        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE, 0);
+        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE, 0, None);
         let mem = GuestMemoryMmap::from_ranges(&regions).expect("Cannot initialize memory");
         assert_eq!(get_fdt_addr(&mem), layout::DRAM_MEM_START);
 
-        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000, 0);
+        let (_mem_info, regions) = arch_memory_regions(layout::FDT_MAX_SIZE + 0x1000, 0, None);
         let mem = GuestMemoryMmap::from_ranges(&regions).expect("Cannot initialize memory");
         assert_eq!(get_fdt_addr(&mem), 0x1000 + layout::DRAM_MEM_START);
     }