unsigned int shift,
                        const struct kvm_memory_slot *memslot,
                        unsigned int lpid);
-extern bool kvmppc_hv_handle_set_rc(struct kvm *kvm, pgd_t *pgtable,
+extern bool kvmppc_hv_handle_set_rc(struct kvm *kvm, bool nested,
                                    bool writing, unsigned long gpa,
                                    unsigned int lpid);
 extern int kvmppc_book3s_instantiate_page(struct kvm_vcpu *vcpu,
 
        return pte;
 }
 
+extern pte_t *find_kvm_nested_guest_pte(struct kvm *kvm, unsigned long lpid,
+                                       unsigned long ea, unsigned *hshift);
+
 #endif /* CONFIG_KVM_BOOK3S_HV_POSSIBLE */
 
 #endif /* __ASM_KVM_BOOK3S_64_H__ */
 
        return ret;
 }
 
-bool kvmppc_hv_handle_set_rc(struct kvm *kvm, pgd_t *pgtable, bool writing,
+bool kvmppc_hv_handle_set_rc(struct kvm *kvm, bool nested, bool writing,
                             unsigned long gpa, unsigned int lpid)
 {
        unsigned long pgflags;
        pgflags = _PAGE_ACCESSED;
        if (writing)
                pgflags |= _PAGE_DIRTY;
-       /*
-        * We are walking the secondary (partition-scoped) page table here.
-        * We can do this without disabling irq because the Linux MM
-        * subsystem doesn't do THP splits and collapses on this tree.
-        */
-       ptep = __find_linux_pte(pgtable, gpa, NULL, &shift);
+
+       if (nested)
+               ptep = find_kvm_nested_guest_pte(kvm, lpid, gpa, &shift);
+       else
+               ptep = find_kvm_secondary_pte(kvm, gpa, &shift);
+
        if (ptep && pte_present(*ptep) && (!writing || pte_write(*ptep))) {
                kvmppc_radix_update_pte(kvm, ptep, 0, pgflags, gpa, shift);
                return true;
        /* Failed to set the reference/change bits */
        if (dsisr & DSISR_SET_RC) {
                spin_lock(&kvm->mmu_lock);
-               if (kvmppc_hv_handle_set_rc(kvm, kvm->arch.pgtable,
-                                           writing, gpa, kvm->arch.lpid))
+               if (kvmppc_hv_handle_set_rc(kvm, false, writing,
+                                           gpa, kvm->arch.lpid))
                        dsisr &= ~DSISR_SET_RC;
                spin_unlock(&kvm->mmu_lock);
 
 
        return kvm->arch.nested_guests[lpid];
 }
 
-static pte_t *find_kvm_nested_guest_pte(struct kvm *kvm, unsigned long lpid,
-                                       unsigned long ea, unsigned *hshift)
+pte_t *find_kvm_nested_guest_pte(struct kvm *kvm, unsigned long lpid,
+                                unsigned long ea, unsigned *hshift)
 {
        struct kvm_nested_guest *gp;
        pte_t *pte;
        return pte;
 }
 
-
 static inline bool kvmhv_n_rmap_is_equal(u64 rmap_1, u64 rmap_2)
 {
        return !((rmap_1 ^ rmap_2) & (RMAP_NESTED_LPID_MASK |
 
        spin_lock(&kvm->mmu_lock);
        /* Set the rc bit in the pte of our (L0) pgtable for the L1 guest */
-       ret = kvmppc_hv_handle_set_rc(kvm, kvm->arch.pgtable, writing,
-                                    gpte.raddr, kvm->arch.lpid);
+       ret = kvmppc_hv_handle_set_rc(kvm, false, writing,
+                                     gpte.raddr, kvm->arch.lpid);
        if (!ret) {
                ret = -EINVAL;
                goto out_unlock;
        }
 
        /* Set the rc bit in the pte of the shadow_pgtable for the nest guest */
-       ret = kvmppc_hv_handle_set_rc(kvm, gp->shadow_pgtable, writing, n_gpa,
-                                     gp->shadow_lpid);
+       ret = kvmppc_hv_handle_set_rc(kvm, true, writing,
+                                     n_gpa, gp->shadow_lpid);
        if (!ret)
                ret = -EINVAL;
        else