#include <linux/numa.h>
 #include <linux/perf_event.h>
 #include <linux/ptrace.h>
+#include <linux/vmalloc.h>
 
 #include <trace/events/kmem.h>
 
 #include <asm/tlb.h>
 #include <asm/tlbflush.h>
 
+#include "pgalloc-track.h"
 #include "internal.h"
 
 #if defined(LAST_CPUPID_NOT_IN_PAGE_FLAGS) && !defined(CONFIG_COMPILE_TEST)
 
 static int apply_to_pte_range(struct mm_struct *mm, pmd_t *pmd,
                                     unsigned long addr, unsigned long end,
-                                    pte_fn_t fn, void *data, bool create)
+                                    pte_fn_t fn, void *data, bool create,
+                                    pgtbl_mod_mask *mask)
 {
        pte_t *pte;
        int err = 0;
 
        if (create) {
                pte = (mm == &init_mm) ?
-                       pte_alloc_kernel(pmd, addr) :
+                       pte_alloc_kernel_track(pmd, addr, mask) :
                        pte_alloc_map_lock(mm, pmd, addr, &ptl);
                if (!pte)
                        return -ENOMEM;
                                break;
                }
        } while (addr += PAGE_SIZE, addr != end);
+       *mask |= PGTBL_PTE_MODIFIED;
 
        arch_leave_lazy_mmu_mode();
 
 
 static int apply_to_pmd_range(struct mm_struct *mm, pud_t *pud,
                                     unsigned long addr, unsigned long end,
-                                    pte_fn_t fn, void *data, bool create)
+                                    pte_fn_t fn, void *data, bool create,
+                                    pgtbl_mod_mask *mask)
 {
        pmd_t *pmd;
        unsigned long next;
        BUG_ON(pud_huge(*pud));
 
        if (create) {
-               pmd = pmd_alloc(mm, pud, addr);
+               pmd = pmd_alloc_track(mm, pud, addr, mask);
                if (!pmd)
                        return -ENOMEM;
        } else {
                next = pmd_addr_end(addr, end);
                if (create || !pmd_none_or_clear_bad(pmd)) {
                        err = apply_to_pte_range(mm, pmd, addr, next, fn, data,
-                                                create);
+                                                create, mask);
                        if (err)
                                break;
                }
 
 static int apply_to_pud_range(struct mm_struct *mm, p4d_t *p4d,
                                     unsigned long addr, unsigned long end,
-                                    pte_fn_t fn, void *data, bool create)
+                                    pte_fn_t fn, void *data, bool create,
+                                    pgtbl_mod_mask *mask)
 {
        pud_t *pud;
        unsigned long next;
        int err = 0;
 
        if (create) {
-               pud = pud_alloc(mm, p4d, addr);
+               pud = pud_alloc_track(mm, p4d, addr, mask);
                if (!pud)
                        return -ENOMEM;
        } else {
                next = pud_addr_end(addr, end);
                if (create || !pud_none_or_clear_bad(pud)) {
                        err = apply_to_pmd_range(mm, pud, addr, next, fn, data,
-                                                create);
+                                                create, mask);
                        if (err)
                                break;
                }
 
 static int apply_to_p4d_range(struct mm_struct *mm, pgd_t *pgd,
                                     unsigned long addr, unsigned long end,
-                                    pte_fn_t fn, void *data, bool create)
+                                    pte_fn_t fn, void *data, bool create,
+                                    pgtbl_mod_mask *mask)
 {
        p4d_t *p4d;
        unsigned long next;
        int err = 0;
 
        if (create) {
-               p4d = p4d_alloc(mm, pgd, addr);
+               p4d = p4d_alloc_track(mm, pgd, addr, mask);
                if (!p4d)
                        return -ENOMEM;
        } else {
                next = p4d_addr_end(addr, end);
                if (create || !p4d_none_or_clear_bad(p4d)) {
                        err = apply_to_pud_range(mm, p4d, addr, next, fn, data,
-                                                create);
+                                                create, mask);
                        if (err)
                                break;
                }
                                 void *data, bool create)
 {
        pgd_t *pgd;
-       unsigned long next;
+       unsigned long start = addr, next;
        unsigned long end = addr + size;
+       pgtbl_mod_mask mask = 0;
        int err = 0;
 
        if (WARN_ON(addr >= end))
                next = pgd_addr_end(addr, end);
                if (!create && pgd_none_or_clear_bad(pgd))
                        continue;
-               err = apply_to_p4d_range(mm, pgd, addr, next, fn, data, create);
+               err = apply_to_p4d_range(mm, pgd, addr, next, fn, data, create, &mask);
                if (err)
                        break;
        } while (pgd++, addr = next, addr != end);
 
+       if (mask & ARCH_PAGE_TABLE_SYNC_MASK)
+               arch_sync_kernel_mappings(start, start + size);
+
        return err;
 }