diff --git a/arch/x86/kvm/mmu.c b/arch/x86/kvm/mmu.c
index 337d173..e85eed6 100644
--- a/arch/x86/kvm/mmu.c
+++ b/arch/x86/kvm/mmu.c
@@ -2427,6 +2427,7 @@ static int set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
 {
 	u64 spte;
 	int ret = 0;
+	bool remap = is_rmap_spte(*sptep);
 
 	if (set_mmio_spte(vcpu->kvm, sptep, gfn, pfn, pte_access))
 		return 0;
@@ -2488,12 +2489,73 @@ static int set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
 		}
 	}
 
-	if (pte_access & ACC_WRITE_MASK)
-		mark_page_dirty(vcpu->kvm, gfn);
-
 set_pte:
 	if (mmu_spte_update(sptep, spte))
 		kvm_flush_remote_tlbs(vcpu->kvm);
+
+	if (!remap) {
+		if (rmap_add(vcpu, sptep, gfn) > RMAP_RECYCLE_THRESHOLD)
+			rmap_recycle(vcpu, sptep, gfn);
+
+		if (level > PT_PAGE_TABLE_LEVEL)
+			++vcpu->kvm->stat.lpages;
+	}
+
+	/*
+	 * The orders we require are:
+	 * 1) set spte to writable __before__ set the dirty bitmap.
+	 *    It makes sure that dirty-logging is not missed when do
+	 *    live migration at the final step where kvm should stop
+	 *    the guest and push the remaining dirty pages got from
+	 *    dirty-bitmap to the destination. The similar cases are
+	 *    in fast_pf_fix_direct_spte() and kvm_write_guest_page().
+	 *
+	 * 2) add the spte into rmap __before__ set the dirty bitmap.
+	 *
+	 * They can ensure we can find the writable spte on the rmap
+	 * when we do lockless write-protection since
+	 * kvm_vm_ioctl_get_dirty_log() write-protects the pages based
+	 * on its dirty-bitmap, otherwise these cases will happen:
+	 *
+	 *      CPU 0                         CPU 1
+	 *                              kvm ioctl doing get-dirty-pages
+	 * mark_page_dirty(gfn) which
+	 * set the gfn on the dirty maps
+	 *                              mask = xchg(dirty_bitmap, 0)
+	 *
+	 *                              try to write-protect gfns which
+	 *                              are set on "mask" then walk then
+	 *                              rmap, see no spte on that rmap
+	 * add the spte into rmap
+	 *
+	 * !!!!!! Then the page can be freely wrote but not recorded in
+	 * the dirty bitmap.
+	 *
+	 * And:
+	 *
+	 *      VCPU 0                        CPU 1
+	 *                                kvm ioctl doing get-dirty-pages
+	 * mark_page_dirty(gfn) which
+	 * set the gfn on the dirty maps
+	 *
+	 * add spte into rmap
+	 *                                mask = xchg(dirty_bitmap, 0)
+	 *
+	 *                                try to write-protect gfns which
+	 *                                are set on "mask" then walk then
+	 *                                rmap, see spte is on the ramp
+	 *                                but it is readonly or nonpresent
+	 * Mark spte writable
+	 *
+	 * !!!!!! Then the page can be freely wrote but not recorded in the
+	 * dirty bitmap.
+	 *
+	 * See the comments in kvm_vm_ioctl_get_dirty_log().
+	 */
+	smp_wmb();
+
+	if (pte_access & ACC_WRITE_MASK)
+		mark_page_dirty(vcpu->kvm, gfn);
 done:
 	return ret;
 }
@@ -2503,9 +2565,6 @@ static void mmu_set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
 			 int level, gfn_t gfn, pfn_t pfn, bool speculative,
 			 bool host_writable)
 {
-	int was_rmapped = 0;
-	int rmap_count;
-
 	pgprintk("%s: spte %llx write_fault %d gfn %llx\n", __func__,
 		 *sptep, write_fault, gfn);
 
@@ -2527,8 +2586,7 @@ static void mmu_set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
 				 spte_to_pfn(*sptep), pfn);
 			drop_spte(vcpu->kvm, sptep);
 			kvm_flush_remote_tlbs(vcpu->kvm);
-		} else
-			was_rmapped = 1;
+		}
 	}
 
 	if (set_spte(vcpu, sptep, pte_access, level, gfn, pfn, speculative,
@@ -2546,16 +2604,6 @@ static void mmu_set_spte(struct kvm_vcpu *vcpu, u64 *sptep,
 		 is_large_pte(*sptep)? "2MB" : "4kB",
 		 *sptep & PT_PRESENT_MASK ?"RW":"R", gfn,
 		 *sptep, sptep);
-	if (!was_rmapped && is_large_pte(*sptep))
-		++vcpu->kvm->stat.lpages;
-
-	if (is_shadow_present_pte(*sptep)) {
-		if (!was_rmapped) {
-			rmap_count = rmap_add(vcpu, sptep, gfn);
-			if (rmap_count > RMAP_RECYCLE_THRESHOLD)
-				rmap_recycle(vcpu, sptep, gfn);
-		}
-	}
 
 	kvm_release_pfn_clean(pfn);
 }
diff --git a/arch/x86/kvm/x86.c b/arch/x86/kvm/x86.c
index 573c6b3..4ac3a27 100644
--- a/arch/x86/kvm/x86.c
+++ b/arch/x86/kvm/x86.c
@@ -3566,6 +3566,16 @@ int kvm_vm_ioctl_get_dirty_log(struct kvm *kvm, struct kvm_dirty_log *log)
 		is_dirty = true;
 
 		mask = xchg(&dirty_bitmap[i], 0);
+		/*
+		 * smp_rmb();
+		 *
+		 * The read-barrier is implied in xchg() acting as a
+		 * full barrier and it ensures getting dirty bitmap
+		 * is before walking the rmap and spte write-protection.
+		 *
+		 * See the comments in set_spte().
+		 */
+
 		dirty_bitmap_buffer[i] = mask;
 
 		offset = i * BITS_PER_LONG;
