diff --git a/arch/arm64/include/asm/virtcca_cvm_smc.h b/arch/arm64/include/asm/virtcca_cvm_smc.h
index 14b4785900945ae25dff7cc79de33bc87747917d..11dbd1cf9a294a819143a110b11dc106c8e1e62d 100644
--- a/arch/arm64/include/asm/virtcca_cvm_smc.h
+++ b/arch/arm64/include/asm/virtcca_cvm_smc.h
@@ -241,10 +241,10 @@ static inline unsigned long tsi_mig_integrity_checksum_init(unsigned long dst_rd
 	return res.a0;
 }
 
-static inline unsigned long tsi_mig_integrity_checksum_loop(unsigned long rd)
+static inline unsigned long tsi_mig_integrity_checksum_loop(unsigned long guest_rd, unsigned long thread_id)
 {
 	struct arm_smccc_res res;
-	arm_smccc_1_1_smc(SMC_TSI_MIG_INTEGRITY_CHECKSUM_LOOP, rd, &res);
+	arm_smccc_1_1_smc(SMC_TSI_MIG_INTEGRITY_CHECKSUM_LOOP, guest_rd, thread_id, &res);
 
 	return res.a0;
 }
diff --git a/arch/arm64/include/uapi/asm/virtcca_cvm_tsi.h b/arch/arm64/include/uapi/asm/virtcca_cvm_tsi.h
index bffde57be1c35bf1633299ea527f0021179f5c98..637551c4acc58fae2fabdd515e9adc93877e60eb 100644
--- a/arch/arm64/include/uapi/asm/virtcca_cvm_tsi.h
+++ b/arch/arm64/include/uapi/asm/virtcca_cvm_tsi.h
@@ -104,6 +104,11 @@ typedef struct virtcca_migvm_info {
 
 #define TMM_GET_MIGRATION_INFO _IOWR(TSI_MAGIC, 3, struct virtcca_migvm_info)
 
-#define TMM_GET_MIGVM_MEM_CHECKSUM _IOW(TSI_MAGIC, 4, unsigned long)
+typedef struct virtcca_migvm_checksum_info {
+	unsigned long guest_rd;
+	unsigned long thread_id;
+} virtcca_migvm_checksum_info_t;
+
+#define TMM_GET_MIGVM_MEM_CHECKSUM _IOW(TSI_MAGIC, 4, struct virtcca_migvm_checksum_info)
 
 #endif  /* __ASM_VIRTCCA_CVM_TSI_H_ */
diff --git a/arch/arm64/kernel/virtcca_cvm_tsi.c b/arch/arm64/kernel/virtcca_cvm_tsi.c
index 9f7bef5dad7a56931ad55eacca9bd7914ea53b0a..271597fca441b1b122b538dce5f3232417455324 100644
--- a/arch/arm64/kernel/virtcca_cvm_tsi.c
+++ b/arch/arm64/kernel/virtcca_cvm_tsi.c
@@ -23,7 +23,7 @@ static int tmm_get_tsi_version(struct virtcca_cvm_tsi_version __user *arg);
 static int tmm_get_attestation_token(struct virtcca_cvm_attestation_cmd __user *arg);
 static int tmm_get_device_cert(struct virtcca_device_cert __user *arg);
 static int tmm_get_set_migration_info(struct virtcca_migvm_info __user *arg);
-static int tmm_migvm_mem_checksum_loop(unsigned long rd);
+static int tmm_migvm_mem_checksum_loop(struct virtcca_migvm_checksum_info checksum_info);
 
 static const struct file_operations tmm_tsi_fops = {
 	.owner		  = THIS_MODULE,
@@ -70,7 +70,7 @@ static void __exit tmm_tsi_exit(void)
 static long tmm_tsi_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
 {
 	int ret;
-	unsigned long rd;
+	struct virtcca_migvm_checksum_info checksum_info;
 	void __user *argp = (void __user *)arg;
 
 	switch (cmd) {
@@ -87,11 +87,11 @@ static long tmm_tsi_ioctl(struct file *file, unsigned int cmd, unsigned long arg
 		ret = tmm_get_set_migration_info((struct virtcca_migvm_info *)arg);
 		break;
 	case TMM_GET_MIGVM_MEM_CHECKSUM:
-		if (copy_from_user(&rd, argp, sizeof(unsigned long))) {
+		if (copy_from_user(&checksum_info, argp, sizeof(struct virtcca_migvm_checksum_info))) {
 			pr_err("tmm_tsi: mem checksum copy data from user failed\n");
 			return -ENOTTY;
 		}
-		ret = tmm_migvm_mem_checksum_loop(rd);
+		ret = tmm_migvm_mem_checksum_loop(checksum_info);
 		break;
 	default:
 		pr_err("tmm_tsi: unknown ioctl command (0x%x)!\n", cmd);
@@ -101,11 +101,11 @@ static long tmm_tsi_ioctl(struct file *file, unsigned int cmd, unsigned long arg
 	return ret;
 }
 
-static int tmm_migvm_mem_checksum_loop(unsigned long rd)
+static int tmm_migvm_mem_checksum_loop(struct virtcca_migvm_checksum_info checksum_info)
 {
 	unsigned long ret;
 
-	ret = tsi_mig_integrity_checksum_loop(rd);
+	ret = tsi_mig_integrity_checksum_loop(checksum_info.guest_rd, checksum_info.thread_id);
 
 	return ret;
 }
diff --git a/arch/arm64/kvm/virtcca_mig.c b/arch/arm64/kvm/virtcca_mig.c
index e2b28b958d834aabf91aa781b20212d9be14a1e2..aa9e299e87ae553bd231def7be99395dc21f8214 100644
--- a/arch/arm64/kvm/virtcca_mig.c
+++ b/arch/arm64/kvm/virtcca_mig.c
@@ -1150,6 +1150,7 @@ static int virtcca_mig_import_state_tec(struct kvm *kvm, struct virtcca_mig_stre
 	uint64_t ret;
 	uint64_t npages;
 	uint16_t vcpu_idx;
+	unsigned long timeout = jiffies + msecs_to_jiffies(TMI_IMPORT_TIMEOUT_MS);
 
 	if (copy_from_user(&npages, (void __user *)data, sizeof(uint64_t)))
 		return -EFAULT;
@@ -1160,11 +1161,18 @@ static int virtcca_mig_import_state_tec(struct kvm *kvm, struct virtcca_mig_stre
 	vcpu = kvm_get_vcpu(kvm, vcpu_idx);
 	tec = &vcpu->arch.tec;
 
-	ret = tmi_import_tec(tec->tec, stream->mbmd.hpa_and_size, stream->page_list.info.val, stream_info.val);
-	if (ret != TMI_SUCCESS) {
-		pr_err("%s: failed, err=%llx\n", __func__, ret);
-		return -EIO;
-	}
+	do {
+		ret = tmi_import_tec(tec->tec, stream->mbmd.hpa_and_size, stream->page_list.info.val, stream_info.val);
+		if (ret != TMI_SUCCESS && ret != TMI_IMPORT_INCOMPLETE) {
+			pr_err("%s: failed, err=%llx\n", __func__, ret);
+			return -EIO;
+		}
+		if (time_after(jiffies, timeout)) {
+			pr_err("tmi_import_commit timeout (%d ms)", TMI_IMPORT_TIMEOUT_MS);
+			return ETIMEDOUT;
+		}
+		msleep(1);
+	} while (ret == TMI_IMPORT_INCOMPLETE);
 
 	return 0;
 }
diff --git a/virt/kvm/kvm_main.c b/virt/kvm/kvm_main.c
index 0aa225f7e11f4e90460ab1556d67a71db13589e7..93899085632c5581791eeceb8bb7f5aabb76d718 100644
--- a/virt/kvm/kvm_main.c
+++ b/virt/kvm/kvm_main.c
@@ -2281,11 +2281,6 @@ static int kvm_get_dirty_log_protect(struct kvm *kvm, struct kvm_dirty_log *log)
 	n = kvm_dirty_bitmap_bytes(memslot);
 	flush = false;
 
-#ifdef CONFIG_HISI_VIRTCCA_HOST
-	if (!kvm_is_realm(kvm))
-		kvm->manual_dirty_log_protect = false;
-#endif
-
 	if (kvm->manual_dirty_log_protect) {
 		/*
 		 * Unlike kvm_get_dirty_log, we always return false in *flush,