}
 
 #ifdef CONFIG_PGSTE
-static pmd_t *pmd_alloc_map(struct mm_struct *mm, unsigned long addr)
+static int pmd_lookup(struct mm_struct *mm, unsigned long addr, pmd_t **pmdp)
 {
+       struct vm_area_struct *vma;
        pgd_t *pgd;
        p4d_t *p4d;
        pud_t *pud;
-       pmd_t *pmd;
+
+       /* We need a valid VMA, otherwise this is clearly a fault. */
+       vma = vma_lookup(mm, addr);
+       if (!vma)
+               return -EFAULT;
 
        pgd = pgd_offset(mm, addr);
-       p4d = p4d_alloc(mm, pgd, addr);
-       if (!p4d)
-               return NULL;
-       pud = pud_alloc(mm, p4d, addr);
-       if (!pud)
-               return NULL;
-       pmd = pmd_alloc(mm, pud, addr);
-       return pmd;
+       if (!pgd_present(*pgd))
+               return -ENOENT;
+
+       p4d = p4d_offset(pgd, addr);
+       if (!p4d_present(*p4d))
+               return -ENOENT;
+
+       pud = pud_offset(p4d, addr);
+       if (!pud_present(*pud))
+               return -ENOENT;
+
+       /* Large PUDs are not supported yet. */
+       if (pud_large(*pud))
+               return -EFAULT;
+
+       *pmdp = pmd_offset(pud, addr);
+       return 0;
 }
 #endif
 
        pmd_t *pmdp;
        pte_t *ptep;
 
-       pmdp = pmd_alloc_map(mm, addr);
-       if (unlikely(!pmdp))
+       if (pmd_lookup(mm, addr, &pmdp))
                return -EFAULT;
 
        ptl = pmd_lock(mm, pmdp);
        pte_t *ptep;
        int cc = 0;
 
-       pmdp = pmd_alloc_map(mm, addr);
-       if (unlikely(!pmdp))
+       if (pmd_lookup(mm, addr, &pmdp))
                return -EFAULT;
 
        ptl = pmd_lock(mm, pmdp);
        pmd_t *pmdp;
        pte_t *ptep;
 
-       pmdp = pmd_alloc_map(mm, addr);
-       if (unlikely(!pmdp))
+       /*
+        * If we don't have a PTE table and if there is no huge page mapped,
+        * the storage key is 0.
+        */
+       *key = 0;
+
+       switch (pmd_lookup(mm, addr, &pmdp)) {
+       case -ENOENT:
+               return 0;
+       case 0:
+               break;
+       default:
                return -EFAULT;
+       }
 
        ptl = pmd_lock(mm, pmdp);
        if (!pmd_present(*pmdp)) {
-               /* Not yet mapped memory has a zero key */
                spin_unlock(ptl);
-               *key = 0;
                return 0;
        }