#define RMAP_RECYCLE_THRESHOLD 1000
 
-static void rmap_add(struct kvm_vcpu *vcpu, u64 *spte, gfn_t gfn)
+static void rmap_add(struct kvm_vcpu *vcpu, struct kvm_memory_slot *slot,
+                    u64 *spte, gfn_t gfn)
 {
-       struct kvm_memory_slot *slot;
        struct kvm_mmu_page *sp;
        struct kvm_rmap_head *rmap_head;
        int rmap_count;
 
        sp = sptep_to_sp(spte);
        kvm_mmu_page_set_gfn(sp, spte - sp->spt, gfn);
-       slot = kvm_vcpu_gfn_to_memslot(vcpu, gfn);
        rmap_head = gfn_to_rmap(gfn, sp->role.level, slot);
        rmap_count = pte_list_add(vcpu, spte, rmap_head);
 
        return 0;
 }
 
-static int mmu_set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
-                       unsigned int pte_access, gfn_t gfn,
+static int mmu_set_spte(struct kvm_vcpu *vcpu, struct kvm_memory_slot *slot,
+                       u64 *sptep, unsigned int pte_access, gfn_t gfn,
                        kvm_pfn_t pfn, struct kvm_page_fault *fault)
 {
        struct kvm_mmu_page *sp = sptep_to_sp(sptep);
        if (!was_rmapped) {
                WARN_ON_ONCE(ret == RET_PF_SPURIOUS);
                kvm_update_page_stats(vcpu->kvm, level, 1);
-               rmap_add(vcpu, sptep, gfn);
+               rmap_add(vcpu, slot, sptep, gfn);
        }
 
        return ret;
 }
 
-static kvm_pfn_t pte_prefetch_gfn_to_pfn(struct kvm_vcpu *vcpu, gfn_t gfn,
-                                    bool no_dirty_log)
-{
-       struct kvm_memory_slot *slot;
-
-       slot = gfn_to_memslot_dirty_bitmap(vcpu, gfn, no_dirty_log);
-       if (!slot)
-               return KVM_PFN_ERR_FAULT;
-
-       return gfn_to_pfn_memslot_atomic(slot, gfn);
-}
-
 static int direct_pte_prefetch_many(struct kvm_vcpu *vcpu,
                                    struct kvm_mmu_page *sp,
                                    u64 *start, u64 *end)
                return -1;
 
        for (i = 0; i < ret; i++, gfn++, start++) {
-               mmu_set_spte(vcpu, start, access, gfn,
+               mmu_set_spte(vcpu, slot, start, access, gfn,
                             page_to_pfn(pages[i]), NULL);
                put_page(pages[i]);
        }
        if (WARN_ON_ONCE(it.level != fault->goal_level))
                return -EFAULT;
 
-       ret = mmu_set_spte(vcpu, it.sptep, ACC_ALL,
+       ret = mmu_set_spte(vcpu, fault->slot, it.sptep, ACC_ALL,
                           base_gfn, fault->pfn, fault);
        if (ret == RET_PF_SPURIOUS)
                return ret;
 
 FNAME(prefetch_gpte)(struct kvm_vcpu *vcpu, struct kvm_mmu_page *sp,
                     u64 *spte, pt_element_t gpte, bool no_dirty_log)
 {
+       struct kvm_memory_slot *slot;
        unsigned pte_access;
        gfn_t gfn;
        kvm_pfn_t pfn;
        gfn = gpte_to_gfn(gpte);
        pte_access = sp->role.access & FNAME(gpte_access)(gpte);
        FNAME(protect_clean_gpte)(vcpu->arch.mmu, &pte_access, gpte);
-       pfn = pte_prefetch_gfn_to_pfn(vcpu, gfn,
+
+       slot = gfn_to_memslot_dirty_bitmap(vcpu, gfn,
                        no_dirty_log && (pte_access & ACC_WRITE_MASK));
+       if (!slot)
+               return false;
+
+       pfn = gfn_to_pfn_memslot_atomic(slot, gfn);
        if (is_error_pfn(pfn))
                return false;
 
-       mmu_set_spte(vcpu, spte, pte_access, gfn, pfn, NULL);
+       mmu_set_spte(vcpu, slot, spte, pte_access, gfn, pfn, NULL);
        kvm_release_pfn_clean(pfn);
        return true;
 }
        if (WARN_ON_ONCE(it.level != fault->goal_level))
                return -EFAULT;
 
-       ret = mmu_set_spte(vcpu, it.sptep, gw->pte_access,
+       ret = mmu_set_spte(vcpu, fault->slot, it.sptep, gw->pte_access,
                           base_gfn, fault->pfn, fault);
        if (ret == RET_PF_SPURIOUS)
                return ret;