}
 
        if (n_removal > 0) {
-               u64 size = efi.memmap.nr_map - n_removal;
+               struct efi_memory_map_data data = {
+                       .phys_map = efi.memmap.phys_map,
+                       .desc_version = efi.memmap.desc_version,
+                       .desc_size = efi.memmap.desc_size,
+                       .size = data.desc_size * (efi.memmap.nr_map - n_removal),
+                       .flags = 0,
+               };
 
                pr_warn("Removing %d invalid memory map entries.\n", n_removal);
-               efi_memmap_install(efi.memmap.phys_map, size);
+               efi_memmap_install(&data);
        }
 }
 
 
  */
 void __init efi_arch_mem_reserve(phys_addr_t addr, u64 size)
 {
-       phys_addr_t new_phys, new_size;
+       struct efi_memory_map_data data = { 0 };
        struct efi_mem_range mr;
        efi_memory_desc_t md;
        int num_entries;
        num_entries = efi_memmap_split_count(&md, &mr.range);
        num_entries += efi.memmap.nr_map;
 
-       new_size = efi.memmap.desc_size * num_entries;
-
-       new_phys = efi_memmap_alloc(num_entries);
-       if (!new_phys) {
+       if (efi_memmap_alloc(num_entries, &data) != 0) {
                pr_err("Could not allocate boot services memmap\n");
                return;
        }
 
-       new = early_memremap(new_phys, new_size);
+       new = early_memremap(data.phys_map, data.size);
        if (!new) {
                pr_err("Failed to map new boot services memmap\n");
                return;
        }
 
        efi_memmap_insert(&efi.memmap, new, &mr);
-       early_memunmap(new, new_size);
+       early_memunmap(new, data.size);
 
-       efi_memmap_install(new_phys, num_entries);
+       efi_memmap_install(&data);
        e820__range_update(addr, size, E820_TYPE_RAM, E820_TYPE_RESERVED);
        e820__update_table(e820_table);
 }
 
 void __init efi_free_boot_services(void)
 {
-       phys_addr_t new_phys, new_size;
+       struct efi_memory_map_data data = { 0 };
        efi_memory_desc_t *md;
        int num_entries = 0;
        void *new, *new_md;
        if (!num_entries)
                return;
 
-       new_size = efi.memmap.desc_size * num_entries;
-       new_phys = efi_memmap_alloc(num_entries);
-       if (!new_phys) {
+       if (efi_memmap_alloc(num_entries, &data) != 0) {
                pr_err("Failed to allocate new EFI memmap\n");
                return;
        }
 
-       new = memremap(new_phys, new_size, MEMREMAP_WB);
+       new = memremap(data.phys_map, data.size, MEMREMAP_WB);
        if (!new) {
                pr_err("Failed to map new EFI memmap\n");
                return;
 
        memunmap(new);
 
-       if (efi_memmap_install(new_phys, num_entries)) {
+       if (efi_memmap_install(&data) != 0) {
                pr_err("Could not install new EFI memmap\n");
                return;
        }
 
 
 void __init efi_fake_memmap(void)
 {
+       struct efi_memory_map_data data = { 0 };
        int new_nr_map = efi.memmap.nr_map;
        efi_memory_desc_t *md;
-       phys_addr_t new_memmap_phy;
        void *new_memmap;
        int i;
 
        }
 
        /* allocate memory for new EFI memmap */
-       new_memmap_phy = efi_memmap_alloc(new_nr_map);
-       if (!new_memmap_phy)
+       if (efi_memmap_alloc(new_nr_map, &data) != 0)
                return;
 
        /* create new EFI memmap */
-       new_memmap = early_memremap(new_memmap_phy,
-                                   efi.memmap.desc_size * new_nr_map);
+       new_memmap = early_memremap(data.phys_map, data.size);
        if (!new_memmap) {
-               memblock_free(new_memmap_phy, efi.memmap.desc_size * new_nr_map);
+               memblock_free(data.phys_map, data.size);
                return;
        }
 
                efi_memmap_insert(&efi.memmap, new_memmap, &efi_fake_mems[i]);
 
        /* swap into new EFI memmap */
-       early_memunmap(new_memmap, efi.memmap.desc_size * new_nr_map);
+       early_memunmap(new_memmap, data.size);
 
-       efi_memmap_install(new_memmap_phy, new_nr_map);
+       efi_memmap_install(&data);
 
        /* print new EFI memmap */
        efi_print_memmap();
 
 /**
  * efi_memmap_alloc - Allocate memory for the EFI memory map
  * @num_entries: Number of entries in the allocated map.
+ * @data: efi memmap installation parameters
  *
  * Depending on whether mm_init() has already been invoked or not,
  * either memblock or "normal" page allocation is used.
  * Returns the physical address of the allocated memory map on
  * success, zero on failure.
  */
-phys_addr_t __init efi_memmap_alloc(unsigned int num_entries)
+int __init efi_memmap_alloc(unsigned int num_entries,
+               struct efi_memory_map_data *data)
 {
-       unsigned long size = num_entries * efi.memmap.desc_size;
-
-       if (slab_is_available())
-               return __efi_memmap_alloc_late(size);
+       /* Expect allocation parameters are zero initialized */
+       WARN_ON(data->phys_map || data->size);
+
+       data->size = num_entries * efi.memmap.desc_size;
+       data->desc_version = efi.memmap.desc_version;
+       data->desc_size = efi.memmap.desc_size;
+       data->flags &= ~(EFI_MEMMAP_SLAB | EFI_MEMMAP_MEMBLOCK);
+       data->flags |= efi.memmap.flags & EFI_MEMMAP_LATE;
+
+       if (slab_is_available()) {
+               data->flags |= EFI_MEMMAP_SLAB;
+               data->phys_map = __efi_memmap_alloc_late(data->size);
+       } else {
+               data->flags |= EFI_MEMMAP_MEMBLOCK;
+               data->phys_map = __efi_memmap_alloc_early(data->size);
+       }
 
-       return __efi_memmap_alloc_early(size);
+       if (!data->phys_map)
+               return -ENOMEM;
+       return 0;
 }
 
 /**
  *
  * Returns zero on success, a negative error code on failure.
  */
-static int __init
-__efi_memmap_init(struct efi_memory_map_data *data)
+static int __init __efi_memmap_init(struct efi_memory_map_data *data)
 {
        struct efi_memory_map map;
        phys_addr_t phys_map;
 
 /**
  * efi_memmap_install - Install a new EFI memory map in efi.memmap
- * @addr: Physical address of the memory map
- * @nr_map: Number of entries in the memory map
+ * @ctx: map allocation parameters (address, size, flags)
  *
  * Unlike efi_memmap_init_*(), this function does not allow the caller
  * to switch from early to late mappings. It simply uses the existing
  *
  * Returns zero on success, a negative error code on failure.
  */
-int __init efi_memmap_install(phys_addr_t addr, unsigned int nr_map)
+int __init efi_memmap_install(struct efi_memory_map_data *data)
 {
-       struct efi_memory_map_data data;
-       unsigned long flags;
-
        efi_memmap_unmap();
 
-       data.phys_map = addr;
-       data.size = efi.memmap.desc_size * nr_map;
-       data.desc_version = efi.memmap.desc_version;
-       data.desc_size = efi.memmap.desc_size;
-       data.flags = efi.memmap.flags & EFI_MEMMAP_LATE;
-
-       return __efi_memmap_init(&data);
+       return __efi_memmap_init(data);
 }
 
 /**
 
 
 /*
  * Architecture independent structure for describing a memory map for the
- * benefit of efi_memmap_init_early(), saving us the need to pass four
- * parameters.
+ * benefit of efi_memmap_init_early(), and for passing context between
+ * efi_memmap_alloc() and efi_memmap_install().
  */
 struct efi_memory_map_data {
        phys_addr_t phys_map;
        unsigned long desc_version;
        unsigned long desc_size;
 #define EFI_MEMMAP_LATE (1UL << 0)
+#define EFI_MEMMAP_MEMBLOCK (1UL << 1)
+#define EFI_MEMMAP_SLAB (1UL << 2)
        unsigned long flags;
 };
 
 #endif
 extern void __iomem *efi_lookup_mapped_addr(u64 phys_addr);
 
-extern phys_addr_t __init efi_memmap_alloc(unsigned int num_entries);
+extern int __init efi_memmap_alloc(unsigned int num_entries,
+                                  struct efi_memory_map_data *data);
 extern int __init efi_memmap_init_early(struct efi_memory_map_data *data);
 extern int __init efi_memmap_init_late(phys_addr_t addr, unsigned long size);
 extern void __init efi_memmap_unmap(void);
-extern int __init efi_memmap_install(phys_addr_t addr, unsigned int nr_map);
+extern int __init efi_memmap_install(struct efi_memory_map_data *data);
 extern int __init efi_memmap_split_count(efi_memory_desc_t *md,
                                         struct range *range);
 extern void __init efi_memmap_insert(struct efi_memory_map *old_memmap,