5 files changed, 25 insertions, 14 deletions

diff --git a/mm/debug.c b/mm/debug.c
index d69cb5a7ba9a..3eb3ac2fcee7 100644
--- a/mm/debug.c
+++ b/mm/debug.c
@@ -173,7 +173,7 @@ void dump_mm(const struct mm_struct *mm)
                 "get_unmapped_area %p\n"
 #endif
                 "mmap_base %lu mmap_legacy_base %lu highest_vm_end %lu\n"
-                "pgd %p mm_users %d mm_count %d nr_ptes %lu map_count %d\n"
+                "pgd %p mm_users %d mm_count %d nr_ptes %lu nr_pmds %lu map_count %d\n"
                 "hiwater_rss %lx hiwater_vm %lx total_vm %lx locked_vm %lx\n"
                 "pinned_vm %lx shared_vm %lx exec_vm %lx stack_vm %lx\n"
                 "start_code %lx end_code %lx start_data %lx end_data %lx\n"
@@ -206,6 +206,7 @@ void dump_mm(const struct mm_struct *mm)
                 mm->pgd, atomic_read(&mm->mm_users),
                 atomic_read(&mm->mm_count),
                 atomic_long_read((atomic_long_t *)&mm->nr_ptes),
+                mm_nr_pmds((struct mm_struct *)mm),
                 mm->map_count,
                 mm->hiwater_rss, mm->hiwater_vm, mm->total_vm, mm->locked_vm,
                 mm->pinned_vm, mm->shared_vm, mm->exec_vm, mm->stack_vm,
diff --git a/mm/hugetlb.c b/mm/hugetlb.c
index fd28d6ba5e5d..0a9ac6c26832 100644
--- a/mm/hugetlb.c
+++ b/mm/hugetlb.c
@@ -3598,6 +3598,7 @@ pte_t *huge_pmd_share(struct mm_struct *mm, unsigned long addr, pud_t *pud)
                 if (saddr) {
                         spte = huge_pte_offset(svma->vm_mm, saddr);
                         if (spte) {
+                                mm_inc_nr_pmds(mm);
                                 get_page(virt_to_page(spte));
                                 break;
                         }
@@ -3609,11 +3610,13 @@ pte_t *huge_pmd_share(struct mm_struct *mm, unsigned long addr, pud_t *pud)
 
         ptl = huge_pte_lockptr(hstate_vma(vma), mm, spte);
         spin_lock(ptl);
-        if (pud_none(*pud))
+        if (pud_none(*pud)) {
                 pud_populate(mm, pud,
                                 (pmd_t *)((unsigned long)spte & PAGE_MASK));
-        else
+        } else {
                 put_page(virt_to_page(spte));
+                mm_inc_nr_pmds(mm);
+        }
         spin_unlock(ptl);
 out:
         pte = (pte_t *)pmd_alloc(mm, pud, addr);
@@ -3644,6 +3647,7 @@ int huge_pmd_unshare(struct mm_struct *mm, unsigned long *addr, pte_t *ptep)
 
         pud_clear(pud);
         put_page(virt_to_page(ptep));
+        mm_dec_nr_pmds(mm);
         *addr = ALIGN(*addr, HPAGE_SIZE * PTRS_PER_PTE) - HPAGE_SIZE;
         return 1;
 }
diff --git a/mm/memory.c b/mm/memory.c
index d63849b5188f..bbe6a73a899d 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -428,6 +428,7 @@ static inline void free_pmd_range(struct mmu_gather *tlb, pud_t *pud,
         pmd = pmd_offset(pud, start);
         pud_clear(pud);
         pmd_free_tlb(tlb, pmd, start);
+        mm_dec_nr_pmds(tlb->mm);
 }
 
 static inline void free_pud_range(struct mmu_gather *tlb, pgd_t *pgd,
@@ -3322,15 +3323,17 @@ int __pmd_alloc(struct mm_struct *mm, pud_t *pud, unsigned long address)
 
         spin_lock(&mm->page_table_lock);
 #ifndef __ARCH_HAS_4LEVEL_HACK
-        if (pud_present(*pud))          /* Another has populated it */
-                pmd_free(mm, new);
-        else
+        if (!pud_present(*pud)) {
+                mm_inc_nr_pmds(mm);
                 pud_populate(mm, pud, new);
-#else
-        if (pgd_present(*pud))          /* Another has populated it */
+        } else  /* Another has populated it */
                 pmd_free(mm, new);
-        else
+#else
+        if (!pgd_present(*pud)) {
+                mm_inc_nr_pmds(mm);
                 pgd_populate(mm, pud, new);
+        } else /* Another has populated it */
+                pmd_free(mm, new);
 #endif /* __ARCH_HAS_4LEVEL_HACK */
         spin_unlock(&mm->page_table_lock);
         return 0;
diff --git a/mm/mmap.c b/mm/mmap.c
index 14d84666e8ba..6a7d36d133fb 100644
--- a/mm/mmap.c
+++ b/mm/mmap.c
@@ -2853,7 +2853,9 @@ void exit_mmap(struct mm_struct *mm)
         vm_unacct_memory(nr_accounted);
 
         WARN_ON(atomic_long_read(&mm->nr_ptes) >
-                        (FIRST_USER_ADDRESS+PMD_SIZE-1)>>PMD_SHIFT);
+                        round_up(FIRST_USER_ADDRESS, PMD_SIZE) >> PMD_SHIFT);
+        WARN_ON(mm_nr_pmds(mm) >
+                        round_up(FIRST_USER_ADDRESS, PUD_SIZE) >> PUD_SHIFT);
 }
 
 /* Insert vm structure into process list sorted by address
diff --git a/mm/oom_kill.c b/mm/oom_kill.c
index b8df76ee2be3..642f38cb175a 100644
--- a/mm/oom_kill.c
+++ b/mm/oom_kill.c
@@ -169,8 +169,8 @@ unsigned long oom_badness(struct task_struct *p, struct mem_cgroup *memcg,
          * The baseline for the badness score is the proportion of RAM that each
          * task's rss, pagetable and swap space use.
          */
-        points = get_mm_rss(p->mm) + atomic_long_read(&p->mm->nr_ptes) +
-                 get_mm_counter(p->mm, MM_SWAPENTS);
+        points = get_mm_rss(p->mm) + get_mm_counter(p->mm, MM_SWAPENTS) +
+                atomic_long_read(&p->mm->nr_ptes) + mm_nr_pmds(p->mm);
         task_unlock(p);
 
         /*
@@ -351,7 +351,7 @@ static void dump_tasks(struct mem_cgroup *memcg, const nodemask_t *nodemask)
         struct task_struct *p;
         struct task_struct *task;
 
-        pr_info("[ pid ]   uid  tgid total_vm      rss nr_ptes swapents oom_score_adj name\n");
+        pr_info("[ pid ]   uid  tgid total_vm      rss nr_ptes nr_pmds swapents oom_score_adj name\n");
         rcu_read_lock();
         for_each_process(p) {
                 if (oom_unkillable_task(p, memcg, nodemask))
@@ -367,10 +367,11 @@ static void dump_tasks(struct mem_cgroup *memcg, const nodemask_t *nodemask)
                         continue;
                 }
 
-                pr_info("[%5d] %5d %5d %8lu %8lu %7ld %8lu         %5hd %s\n",
+                pr_info("[%5d] %5d %5d %8lu %8lu %7ld %7ld %8lu         %5hd %s\n",
                         task->pid, from_kuid(&init_user_ns, task_uid(task)),
                         task->tgid, task->mm->total_vm, get_mm_rss(task->mm),
                         atomic_long_read(&task->mm->nr_ptes),
+                        mm_nr_pmds(task->mm),
                         get_mm_counter(task->mm, MM_SWAPENTS),
                         task->signal->oom_score_adj, task->comm);
                 task_unlock(task);