#define KVM_ARCH_WANT_MMU_NOTIFIER
 int kvm_unmap_hva(struct kvm *kvm, unsigned long hva);
 int kvm_age_hva(struct kvm *kvm, unsigned long hva);
+int kvm_test_age_hva(struct kvm *kvm, unsigned long hva);
 void kvm_set_spte_hva(struct kvm *kvm, unsigned long hva, pte_t pte);
 int cpuid_maxphyaddr(struct kvm_vcpu *vcpu);
 int kvm_cpu_has_interrupt(struct kvm_vcpu *vcpu);
 
        return young;
 }
 
+static int kvm_test_age_rmapp(struct kvm *kvm, unsigned long *rmapp,
+                             unsigned long data)
+{
+       u64 *spte;
+       int young = 0;
+
+       /*
+        * If there's no access bit in the secondary pte set by the
+        * hardware it's up to gup-fast/gup to set the access bit in
+        * the primary pte or in the page structure.
+        */
+       if (!shadow_accessed_mask)
+               goto out;
+
+       spte = rmap_next(kvm, rmapp, NULL);
+       while (spte) {
+               u64 _spte = *spte;
+               BUG_ON(!(_spte & PT_PRESENT_MASK));
+               young = _spte & PT_ACCESSED_MASK;
+               if (young) {
+                       young = 1;
+                       break;
+               }
+               spte = rmap_next(kvm, rmapp, spte);
+       }
+out:
+       return young;
+}
+
 #define RMAP_RECYCLE_THRESHOLD 1000
 
 static void rmap_recycle(struct kvm_vcpu *vcpu, u64 *spte, gfn_t gfn)
        return kvm_handle_hva(kvm, hva, 0, kvm_age_rmapp);
 }
 
+int kvm_test_age_hva(struct kvm *kvm, unsigned long hva)
+{
+       return kvm_handle_hva(kvm, hva, 0, kvm_test_age_rmapp);
+}
+
 #ifdef MMU_DEBUG
 static int is_empty_shadow_page(u64 *spt)
 {
 
 #include <linux/mm.h>
 #include <linux/vmstat.h>
 #include <linux/highmem.h>
+#include <linux/swap.h>
 
 #include <asm/pgtable.h>
 
                VM_BUG_ON(!pfn_valid(pte_pfn(pte)));
                page = pte_page(pte);
                get_page(page);
+               SetPageReferenced(page);
                pages[*nr] = page;
                (*nr)++;
 
        VM_BUG_ON(page != compound_head(page));
        VM_BUG_ON(page_count(page) == 0);
        atomic_add(nr, &page->_count);
+       SetPageReferenced(page);
 }
 
 static inline void get_huge_page_tail(struct page *page)
 
                                 struct mm_struct *mm,
                                 unsigned long address);
 
+       /*
+        * test_young is called to check the young/accessed bitflag in
+        * the secondary pte. This is used to know if the page is
+        * frequently used without actually clearing the flag or tearing
+        * down the secondary mapping on the page.
+        */
+       int (*test_young)(struct mmu_notifier *mn,
+                         struct mm_struct *mm,
+                         unsigned long address);
+
        /*
         * change_pte is called in cases that pte mapping to page is changed:
         * for example, when ksm remaps pte to point to a new shared page.
 extern void __mmu_notifier_release(struct mm_struct *mm);
 extern int __mmu_notifier_clear_flush_young(struct mm_struct *mm,
                                          unsigned long address);
+extern int __mmu_notifier_test_young(struct mm_struct *mm,
+                                    unsigned long address);
 extern void __mmu_notifier_change_pte(struct mm_struct *mm,
                                      unsigned long address, pte_t pte);
 extern void __mmu_notifier_invalidate_page(struct mm_struct *mm,
        return 0;
 }
 
+static inline int mmu_notifier_test_young(struct mm_struct *mm,
+                                         unsigned long address)
+{
+       if (mm_has_notifiers(mm))
+               return __mmu_notifier_test_young(mm, address);
+       return 0;
+}
+
 static inline void mmu_notifier_change_pte(struct mm_struct *mm,
                                           unsigned long address, pte_t pte)
 {
        return 0;
 }
 
+static inline int mmu_notifier_test_young(struct mm_struct *mm,
+                                         unsigned long address)
+{
+       return 0;
+}
+
 static inline void mmu_notifier_change_pte(struct mm_struct *mm,
                                           unsigned long address, pte_t pte)
 {
 
                VM_BUG_ON(PageLRU(page));
 
                /* If there is no mapped pte young don't collapse the page */
-               if (pte_young(pteval))
+               if (pte_young(pteval) || PageReferenced(page) ||
+                   mmu_notifier_test_young(vma->vm_mm, address))
                        referenced = 1;
        }
        if (unlikely(!referenced))
                /* cannot use mapcount: can't collapse if there's a gup pin */
                if (page_count(page) != 1)
                        goto out_unmap;
-               if (pte_young(pteval))
+               if (pte_young(pteval) || PageReferenced(page) ||
+                   mmu_notifier_test_young(vma->vm_mm, address))
                        referenced = 1;
        }
        if (referenced)
 
        return young;
 }
 
+int __mmu_notifier_test_young(struct mm_struct *mm,
+                             unsigned long address)
+{
+       struct mmu_notifier *mn;
+       struct hlist_node *n;
+       int young = 0;
+
+       rcu_read_lock();
+       hlist_for_each_entry_rcu(mn, n, &mm->mmu_notifier_mm->list, hlist) {
+               if (mn->ops->test_young) {
+                       young = mn->ops->test_young(mn, mm, address);
+                       if (young)
+                               break;
+               }
+       }
+       rcu_read_unlock();
+
+       return young;
+}
+
 void __mmu_notifier_change_pte(struct mm_struct *mm, unsigned long address,
                               pte_t pte)
 {
 
        return young;
 }
 
+static int kvm_mmu_notifier_test_young(struct mmu_notifier *mn,
+                                      struct mm_struct *mm,
+                                      unsigned long address)
+{
+       struct kvm *kvm = mmu_notifier_to_kvm(mn);
+       int young, idx;
+
+       idx = srcu_read_lock(&kvm->srcu);
+       spin_lock(&kvm->mmu_lock);
+       young = kvm_test_age_hva(kvm, address);
+       spin_unlock(&kvm->mmu_lock);
+       srcu_read_unlock(&kvm->srcu, idx);
+
+       return young;
+}
+
 static void kvm_mmu_notifier_release(struct mmu_notifier *mn,
                                     struct mm_struct *mm)
 {
        .invalidate_range_start = kvm_mmu_notifier_invalidate_range_start,
        .invalidate_range_end   = kvm_mmu_notifier_invalidate_range_end,
        .clear_flush_young      = kvm_mmu_notifier_clear_flush_young,
+       .test_young             = kvm_mmu_notifier_test_young,
        .change_pte             = kvm_mmu_notifier_change_pte,
        .release                = kvm_mmu_notifier_release,
 };