diff --git a/cputlb.c b/cputlb.c
index b533f3f372..71cf64fd40 100644
--- a/cputlb.c
+++ b/cputlb.c
@@ -313,7 +313,7 @@ tb_page_addr_t get_page_addr_code(CPUArchState *env1, target_ulong addr)
         cpu_ldub_code(env1, addr);
     }
     pd = env1->iotlb[mmu_idx][page_index] & ~TARGET_PAGE_MASK;
-    mr = iotlb_to_region(pd);
+    mr = iotlb_to_region(&address_space_memory, pd);
     if (memory_region_is_unassigned(mr)) {
         CPUState *cpu = ENV_GET_CPU(env1);
         CPUClass *cc = CPU_GET_CLASS(cpu);
diff --git a/exec.c b/exec.c
index 05ad67d3fd..f4c987a862 100644
--- a/exec.c
+++ b/exec.c
@@ -1746,10 +1746,9 @@ static uint16_t dummy_section(PhysPageMap *map, MemoryRegion *mr)
     return phys_section_add(map, &section);
 }
 
-MemoryRegion *iotlb_to_region(hwaddr index)
+MemoryRegion *iotlb_to_region(AddressSpace *as, hwaddr index)
 {
-    return address_space_memory.dispatch->map.sections[
-           index & ~TARGET_PAGE_MASK].mr;
+    return as->dispatch->map.sections[index & ~TARGET_PAGE_MASK].mr;
 }
 
 static void io_mem_init(void)
diff --git a/include/exec/exec-all.h b/include/exec/exec-all.h
index 9d6c36b7a9..75fd1da16f 100644
--- a/include/exec/exec-all.h
+++ b/include/exec/exec-all.h
@@ -325,7 +325,7 @@ extern uintptr_t tci_tb_ptr;
 
 void phys_mem_set_alloc(void *(*alloc)(size_t));
 
-struct MemoryRegion *iotlb_to_region(hwaddr index);
+struct MemoryRegion *iotlb_to_region(AddressSpace *as, hwaddr index);
 bool io_mem_read(struct MemoryRegion *mr, hwaddr addr,
                  uint64_t *pvalue, unsigned size);
 bool io_mem_write(struct MemoryRegion *mr, hwaddr addr,
diff --git a/include/exec/softmmu_template.h b/include/exec/softmmu_template.h
index 8712dcd091..a946b35ca3 100644
--- a/include/exec/softmmu_template.h
+++ b/include/exec/softmmu_template.h
@@ -22,6 +22,7 @@
  * License along with this library; if not, see <http://www.gnu.org/licenses/>.
  */
 #include "qemu/timer.h"
+#include "exec/address-spaces.h"
 #include "exec/memory.h"
 
 #define DATA_SIZE (1 << SHIFT)
@@ -121,7 +122,7 @@ static inline DATA_TYPE glue(io_read, SUFFIX)(CPUArchState *env,
                                               uintptr_t retaddr)
 {
     uint64_t val;
-    MemoryRegion *mr = iotlb_to_region(physaddr);
+    MemoryRegion *mr = iotlb_to_region(&address_space_memory, physaddr);
 
     physaddr = (physaddr & TARGET_PAGE_MASK) + addr;
     env->mem_io_pc = retaddr;
@@ -327,7 +328,7 @@ static inline void glue(io_write, SUFFIX)(CPUArchState *env,
                                           target_ulong addr,
                                           uintptr_t retaddr)
 {
-    MemoryRegion *mr = iotlb_to_region(physaddr);
+    MemoryRegion *mr = iotlb_to_region(&address_space_memory, physaddr);
 
     physaddr = (physaddr & TARGET_PAGE_MASK) + addr;
     if (mr != &io_mem_rom && mr != &io_mem_notdirty && !can_do_io(env)) {