4 files changed, 32 insertions, 6 deletions

diff --git a/fs/proc/task_mmu.c b/fs/proc/task_mmu.c
index 60b914860f81..78fd3621f565 100644
--- a/fs/proc/task_mmu.c
+++ b/fs/proc/task_mmu.c
@@ -343,6 +343,8 @@ static int smaps_pte_range(pmd_t *pmd, unsigned long addr, unsigned long end,
         struct page *page;
         int mapcount;
 
+        split_huge_page_pmd(walk->mm, pmd);
+
         pte = pte_offset_map_lock(vma->vm_mm, pmd, addr, &ptl);
         for (; addr != end; pte++, addr += PAGE_SIZE) {
                 ptent = *pte;
@@ -467,6 +469,8 @@ static int clear_refs_pte_range(pmd_t *pmd, unsigned long addr,
         spinlock_t *ptl;
         struct page *page;
 
+        split_huge_page_pmd(walk->mm, pmd);
+
         pte = pte_offset_map_lock(vma->vm_mm, pmd, addr, &ptl);
         for (; addr != end; pte++, addr += PAGE_SIZE) {
                 ptent = *pte;
@@ -623,6 +627,8 @@ static int pagemap_pte_range(pmd_t *pmd, unsigned long addr, unsigned long end,
         pte_t *pte;
         int err = 0;
 
+        split_huge_page_pmd(walk->mm, pmd);
+
         /* find the first VMA at or above 'addr' */
         vma = find_vma(walk->mm, addr);
         for (; addr != end; addr += PAGE_SIZE) {
diff --git a/include/linux/mm.h b/include/linux/mm.h
index 901435e3a9a9..294104e0891d 100644
--- a/include/linux/mm.h
+++ b/include/linux/mm.h
@@ -914,6 +914,9 @@ unsigned long unmap_vmas(struct mmu_gather **tlb,
  * @pgd_entry: if set, called for each non-empty PGD (top-level) entry
  * @pud_entry: if set, called for each non-empty PUD (2nd-level) entry
  * @pmd_entry: if set, called for each non-empty PMD (3rd-level) entry
+ *             this handler is required to be able to handle
+ *             pmd_trans_huge() pmds.  They may simply choose to
+ *             split_huge_page() instead of handling it explicitly.
  * @pte_entry: if set, called for each non-empty PTE (4th-level) entry
  * @pte_hole: if set, called for each hole at all levels
  * @hugetlb_entry: if set, called for each hugetlb entry
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 9e0f05efd114..e1ee6ad9c971 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -4763,7 +4763,8 @@ static int mem_cgroup_count_precharge_pte_range(pmd_t *pmd,
         pte_t *pte;
         spinlock_t *ptl;
 
-        VM_BUG_ON(pmd_trans_huge(*pmd));
+        split_huge_page_pmd(walk->mm, pmd);
+
         pte = pte_offset_map_lock(vma->vm_mm, pmd, addr, &ptl);
         for (; addr != end; pte++, addr += PAGE_SIZE)
                 if (is_target_pte_for_mc(vma, addr, *pte, NULL))
@@ -4925,8 +4926,8 @@ static int mem_cgroup_move_charge_pte_range(pmd_t *pmd,
         pte_t *pte;
         spinlock_t *ptl;
 
+        split_huge_page_pmd(walk->mm, pmd);
 retry:
-        VM_BUG_ON(pmd_trans_huge(*pmd));
         pte = pte_offset_map_lock(vma->vm_mm, pmd, addr, &ptl);
         for (; addr != end; addr += PAGE_SIZE) {
                 pte_t ptent = *(pte++);
diff --git a/mm/pagewalk.c b/mm/pagewalk.c
index 7cfa6ae02303..c3450d533611 100644
--- a/mm/pagewalk.c
+++ b/mm/pagewalk.c
@@ -33,19 +33,35 @@ static int walk_pmd_range(pud_t *pud, unsigned long addr, unsigned long end,
 
         pmd = pmd_offset(pud, addr);
         do {
+again:
                 next = pmd_addr_end(addr, end);
-                split_huge_page_pmd(walk->mm, pmd);
-                if (pmd_none_or_clear_bad(pmd)) {
+                if (pmd_none(*pmd)) {
                         if (walk->pte_hole)
                                 err = walk->pte_hole(addr, next, walk);
                         if (err)
                                 break;
                         continue;
                 }
+                /*
+                 * This implies that each ->pmd_entry() handler
+                 * needs to know about pmd_trans_huge() pmds
+                 */
                 if (walk->pmd_entry)
                         err = walk->pmd_entry(pmd, addr, next, walk);
-                if (!err && walk->pte_entry)
-                        err = walk_pte_range(pmd, addr, next, walk);
+                if (err)
+                        break;
+
+                /*
+                 * Check this here so we only break down trans_huge
+                 * pages when we _need_ to
+                 */
+                if (!walk->pte_entry)
+                        continue;
+
+                split_huge_page_pmd(walk->mm, pmd);
+                if (pmd_none_or_clear_bad(pmd))
+                        goto again;
+                err = walk_pte_range(pmd, addr, next, walk);
                 if (err)
                         break;
         } while (pmd++, addr = next, addr != end);