return ret;
 }
 
-static int kvmppc_get_hpa(struct kvm_vcpu *vcpu, unsigned long gpa,
-                         int writing, unsigned long *hpa,
+static int kvmppc_get_hpa(struct kvm_vcpu *vcpu, unsigned long mmu_seq,
+                         unsigned long gpa, int writing, unsigned long *hpa,
                          struct kvm_memory_slot **memslot_p)
 {
        struct kvm *kvm = vcpu->kvm;
        hva = __gfn_to_hva_memslot(memslot, gfn);
 
        /* Try to find the host pte for that virtual address */
-       ptep = __find_linux_pte(vcpu->arch.pgdir, hva, NULL, &shift);
+       ptep = find_kvm_host_pte(kvm, mmu_seq, hva, &shift);
        if (!ptep)
                return H_TOO_HARD;
        pte = kvmppc_read_update_linux_pte(ptep, writing);
        mmu_seq = kvm->mmu_notifier_seq;
        smp_rmb();
 
-       ret = kvmppc_get_hpa(vcpu, dest, 1, &pa, &memslot);
-       if (ret != H_SUCCESS)
-               return ret;
-
-       /* Check if we've been invalidated */
        arch_spin_lock(&kvm->mmu_lock.rlock.raw_lock);
-       if (mmu_notifier_retry(kvm, mmu_seq)) {
-               ret = H_TOO_HARD;
+
+       ret = kvmppc_get_hpa(vcpu, mmu_seq, dest, 1, &pa, &memslot);
+       if (ret != H_SUCCESS)
                goto out_unlock;
-       }
 
        /* Zero the page */
        for (i = 0; i < SZ_4K; i += L1_CACHE_BYTES, pa += L1_CACHE_BYTES)
        mmu_seq = kvm->mmu_notifier_seq;
        smp_rmb();
 
-       ret = kvmppc_get_hpa(vcpu, dest, 1, &dest_pa, &dest_memslot);
-       if (ret != H_SUCCESS)
-               return ret;
-       ret = kvmppc_get_hpa(vcpu, src, 0, &src_pa, NULL);
+       arch_spin_lock(&kvm->mmu_lock.rlock.raw_lock);
+       ret = kvmppc_get_hpa(vcpu, mmu_seq, dest, 1, &dest_pa, &dest_memslot);
        if (ret != H_SUCCESS)
-               return ret;
+               goto out_unlock;
 
-       /* Check if we've been invalidated */
-       arch_spin_lock(&kvm->mmu_lock.rlock.raw_lock);
-       if (mmu_notifier_retry(kvm, mmu_seq)) {
-               ret = H_TOO_HARD;
+       ret = kvmppc_get_hpa(vcpu, mmu_seq, src, 0, &src_pa, NULL);
+       if (ret != H_SUCCESS)
                goto out_unlock;
-       }
 
        /* Copy the page */
        memcpy((void *)dest_pa, (void *)src_pa, SZ_4K);