KVM: MMU: fix large page breakage

index 95c12bc..9cd75ee 100644
@@ -1047,25 +1047,24 @@ static void mmu_set_spte(struct kvm_vcpu *vcpu, u64 *shadow_pte,
 		 write_fault, user_fault, gfn);
 
 	if (is_rmap_pte(*shadow_pte)) {
-		if (page != spte_to_page(*shadow_pte))
-			rmap_remove(vcpu->kvm, shadow_pte);
-		else
-			was_rmapped = 1;
-	}
-
-	/*
-	 * If we overwrite a PTE page pointer with a 2MB PMD, unlink
-	 * the parent of the now unreachable PTE.
-	 */
-	if (largepage) {
-		if (was_rmapped && !is_large_pte(*shadow_pte)) {
+		/*
+		 * If we overwrite a PTE page pointer with a 2MB PMD, unlink
+		 * the parent of the now unreachable PTE.
+		 */
+		if (largepage && !is_large_pte(*shadow_pte)) {
 			struct kvm_mmu_page *child;
 			u64 pte = *shadow_pte;
 
 			child = page_header(pte & PT64_BASE_ADDR_MASK);
 			mmu_page_remove_parent_pte(child, shadow_pte);
+		} else if (page != spte_to_page(*shadow_pte))
+			rmap_remove(vcpu->kvm, shadow_pte);
+		else {
+			if (largepage)
+				was_rmapped = is_large_pte(*shadow_pte);
+			else
+				was_rmapped = 1;
 		}
-		was_rmapped = is_large_pte(*shadow_pte);
 	}
 
 	/*