#define destroy_context(mm)            do { } while(0)
 void check_and_switch_context(struct mm_struct *mm, unsigned int cpu);
 
-#define init_new_context(tsk,mm)       ({ atomic64_set(&mm->context.id, 0); 0; })
+#define init_new_context(tsk,mm)       ({ atomic64_set(&(mm)->context.id, 0); 0; })
 
 /*
  * This is called when "tsk" is about to enter lazy TLB mode.
 
 {
        efi_memory_desc_t *md;
 
+       init_new_context(NULL, &efi_mm);
+
        for_each_efi_memory_desc(&memmap, md) {
                u64 paddr, npages, size;
                pgprot_t prot;
                else
                        prot = PAGE_KERNEL;
 
-               create_pgd_mapping(&efi_mm, paddr, md->virt_addr, size, prot);
+               create_pgd_mapping(&efi_mm, paddr, md->virt_addr, size,
+                                  __pgprot(pgprot_val(prot) | PTE_NG));
        }
        return true;
 }
 
 static void efi_set_pgd(struct mm_struct *mm)
 {
-       if (mm == &init_mm)
-               cpu_set_reserved_ttbr0();
-       else
-               cpu_switch_mm(mm->pgd, mm);
-
-       local_flush_tlb_all();
-       if (icache_is_aivivt())
-               __local_flush_icache_all();
+       switch_mm(NULL, mm, NULL);
 }
 
 void efi_virtmap_load(void)