2 files changed, 11 insertions, 4 deletions

diff --git a/arch/i386/mm/hugetlbpage.c b/arch/i386/mm/hugetlbpage.c
index 3b099f32b948..57c486f0e896 100644
--- a/arch/i386/mm/hugetlbpage.c
+++ b/arch/i386/mm/hugetlbpage.c
@@ -22,12 +22,21 @@ pte_t *huge_pte_alloc(struct mm_struct *mm, unsigned long addr)
 {
         pgd_t *pgd;
         pud_t *pud;
-        pmd_t *pmd = NULL;
+        pmd_t *pmd;
+        pte_t *pte = NULL;
 
         pgd = pgd_offset(mm, addr);
         pud = pud_alloc(mm, pgd, addr);
         pmd = pmd_alloc(mm, pud, addr);
-        return (pte_t *) pmd;
+
+        if (!pmd)
+                goto out;
+
+        pte = (pte_t *) pmd;
+        if (!pte_none(*pte) && !pte_huge(*pte))
+                hugetlb_clean_stale_pgtable(pte);
+out:
+        return pte;
 }
 
 pte_t *huge_pte_offset(struct mm_struct *mm, unsigned long addr)
diff --git a/mm/hugetlb.c b/mm/hugetlb.c
index 6bf720bc662c..901ac523a1c3 100644
--- a/mm/hugetlb.c
+++ b/mm/hugetlb.c
@@ -360,8 +360,6 @@ int hugetlb_prefault(struct address_space *mapping, struct vm_area_struct *vma)
                         ret = -ENOMEM;
                         goto out;
                 }
-                if (! pte_none(*pte))
-                        hugetlb_clean_stale_pgtable(pte);
 
                 idx = ((addr - vma->vm_start) >> HPAGE_SHIFT)
                         + (vma->vm_pgoff >> (HPAGE_SHIFT - PAGE_SHIFT));