return &slot->arch.rmap[level - PG_LEVEL_4K][idx];
 }
 
-static struct kvm_rmap_head *gfn_to_rmap(struct kvm *kvm, gfn_t gfn,
-                                        struct kvm_mmu_page *sp)
-{
-       struct kvm_memslots *slots;
-       struct kvm_memory_slot *slot;
-
-       slots = kvm_memslots_for_spte_role(kvm, sp->role);
-       slot = __gfn_to_memslot(slots, gfn);
-       return __gfn_to_rmap(gfn, sp->role.level, slot);
-}
-
 static bool rmap_can_add(struct kvm_vcpu *vcpu)
 {
        struct kvm_mmu_memory_cache *mc;
 
 static int rmap_add(struct kvm_vcpu *vcpu, u64 *spte, gfn_t gfn)
 {
+       struct kvm_memory_slot *slot;
        struct kvm_mmu_page *sp;
        struct kvm_rmap_head *rmap_head;
 
        sp = sptep_to_sp(spte);
        kvm_mmu_page_set_gfn(sp, spte - sp->spt, gfn);
-       rmap_head = gfn_to_rmap(vcpu->kvm, gfn, sp);
+       slot = kvm_vcpu_gfn_to_memslot(vcpu, gfn);
+       rmap_head = __gfn_to_rmap(gfn, sp->role.level, slot);
        return pte_list_add(vcpu, spte, rmap_head);
 }
 
+
 static void rmap_remove(struct kvm *kvm, u64 *spte)
 {
+       struct kvm_memslots *slots;
+       struct kvm_memory_slot *slot;
        struct kvm_mmu_page *sp;
        gfn_t gfn;
        struct kvm_rmap_head *rmap_head;
 
        sp = sptep_to_sp(spte);
        gfn = kvm_mmu_page_get_gfn(sp, spte - sp->spt);
-       rmap_head = gfn_to_rmap(kvm, gfn, sp);
+
+       /*
+        * Unlike rmap_add and rmap_recycle, rmap_remove does not run in the
+        * context of a vCPU so have to determine which memslots to use based
+        * on context information in sp->role.
+        */
+       slots = kvm_memslots_for_spte_role(kvm, sp->role);
+
+       slot = __gfn_to_memslot(slots, gfn);
+       rmap_head = __gfn_to_rmap(gfn, sp->role.level, slot);
+
        __pte_list_remove(spte, rmap_head);
 }
 
 
 static void rmap_recycle(struct kvm_vcpu *vcpu, u64 *spte, gfn_t gfn)
 {
+       struct kvm_memory_slot *slot;
        struct kvm_rmap_head *rmap_head;
        struct kvm_mmu_page *sp;
 
        sp = sptep_to_sp(spte);
-
-       rmap_head = gfn_to_rmap(vcpu->kvm, gfn, sp);
+       slot = kvm_vcpu_gfn_to_memslot(vcpu, gfn);
+       rmap_head = __gfn_to_rmap(gfn, sp->role.level, slot);
 
        kvm_unmap_rmapp(vcpu->kvm, rmap_head, NULL, gfn, sp->role.level, __pte(0));
        kvm_flush_remote_tlbs_with_address(vcpu->kvm, sp->gfn,