}
 #endif
 
+/*
+ * For a prot_numa update we only hold mmap_sem for read so there is a
+ * potential race with faulting where a pmd was temporarily none. This
+ * function checks for a transhuge pmd under the appropriate lock. It
+ * returns a pte if it was successfully locked or NULL if it raced with
+ * a transhuge insertion.
+ */
+static pte_t *lock_pte_protection(struct vm_area_struct *vma, pmd_t *pmd,
+                       unsigned long addr, int prot_numa, spinlock_t **ptl)
+{
+       pte_t *pte;
+       spinlock_t *pmdl;
+
+       /* !prot_numa is protected by mmap_sem held for write */
+       if (!prot_numa)
+               return pte_offset_map_lock(vma->vm_mm, pmd, addr, ptl);
+
+       pmdl = pmd_lock(vma->vm_mm, pmd);
+       if (unlikely(pmd_trans_huge(*pmd) || pmd_none(*pmd))) {
+               spin_unlock(pmdl);
+               return NULL;
+       }
+
+       pte = pte_offset_map_lock(vma->vm_mm, pmd, addr, ptl);
+       spin_unlock(pmdl);
+       return pte;
+}
+
 static unsigned long change_pte_range(struct vm_area_struct *vma, pmd_t *pmd,
                unsigned long addr, unsigned long end, pgprot_t newprot,
                int dirty_accountable, int prot_numa)
        spinlock_t *ptl;
        unsigned long pages = 0;
 
-       pte = pte_offset_map_lock(mm, pmd, addr, &ptl);
+       pte = lock_pte_protection(vma, pmd, addr, prot_numa, &ptl);
+       if (!pte)
+               return 0;
+
        arch_enter_lazy_mmu_mode();
        do {
                oldpte = *pte;
                                                pages += HPAGE_PMD_NR;
                                                nr_huge_updates++;
                                        }
+
+                                       /* huge pmd was handled */
                                        continue;
                                }
                        }
                        /* fall through, the trans huge pmd just split */
                }
-               VM_BUG_ON(pmd_trans_huge(*pmd));
                this_pages = change_pte_range(vma, pmd, addr, next, newprot,
                                 dirty_accountable, prot_numa);
                pages += this_pages;