[PATCH v3 2/3] LoongArch: KVM: Implement vmid updating logic

[Bibo Mao]

VMID calculation method is the same with ASID on LoongArch, it is
percpu vmid calculation method. For every physical CPU, VMID of
different VM is different, and it is the same for different vCPUs
of the same VM.

When vCPU is scheduled on the physical CPU, it checked vmid of this
VM and the global cached vmid, and judge whether it is valid or not.

Signed-off-by: Bibo Mao <maobibo@xxxxxxxxxxx>
---
 arch/loongarch/include/asm/kvm_host.h |  7 ++++
 arch/loongarch/kvm/main.c             | 55 ++++++++++++++++++++++++---
 arch/loongarch/kvm/mmu.c              | 30 +++++++++++++--
 arch/loongarch/kvm/tlb.c              | 14 +++++++
 4 files changed, 97 insertions(+), 9 deletions(-)

diff --git a/arch/loongarch/include/asm/kvm_host.h b/arch/loongarch/include/asm/kvm_host.h
index 5c2f76d1e2f8..8273122604d8 100644
--- a/arch/loongarch/include/asm/kvm_host.h
+++ b/arch/loongarch/include/asm/kvm_host.h
@@ -79,6 +79,7 @@ struct kvm_arch_memory_slot {
 #define HOST_MAX_PMNUM			16
 struct kvm_context {
 	unsigned long vpid_cache;
+	unsigned long vmid_cache;
 	struct kvm_vcpu *last_vcpu;
 	/* Host PMU CSR */
 	u64 perf_ctrl[HOST_MAX_PMNUM];
@@ -132,6 +133,8 @@ struct kvm_arch {
 	unsigned long kvm_features;
 
 	s64 time_offset;
+	cpumask_t tlb_flush_pending;
+	unsigned long vmid[NR_CPUS];
 	struct kvm_context __percpu *vmcs;
 	struct loongarch_ipi *ipi;
 	struct loongarch_dmsintc *dmsintc;
@@ -319,6 +322,8 @@ bool kvm_arch_pmi_in_guest(struct kvm_vcpu *vcpu);
 int kvm_arch_vcpu_dump_regs(struct kvm_vcpu *vcpu);
 
 /* MMU handling */
+void kvm_flush_tlb_all_stage1(void);
+void kvm_flush_tlb_all_stage2(void);
 void kvm_flush_tlb_all(void);
 void kvm_flush_tlb_gpa(struct kvm_vcpu *vcpu, unsigned long gpa);
 int kvm_handle_mm_fault(struct kvm_vcpu *vcpu, unsigned long badv, bool write, int ecode);
@@ -353,6 +358,8 @@ static inline void kvm_arch_vcpu_block_finish(struct kvm_vcpu *vcpu) {}
 static inline void kvm_arch_free_memslot(struct kvm *kvm, struct kvm_memory_slot *slot) {}
 void kvm_check_vpid(struct kvm_vcpu *vcpu);
 enum hrtimer_restart kvm_swtimer_wakeup(struct hrtimer *timer);
+#define __KVM_HAVE_ARCH_FLUSH_REMOTE_TLBS
+int kvm_arch_flush_remote_tlbs(struct kvm *kvm);
 void kvm_arch_flush_remote_tlbs_memslot(struct kvm *kvm, const struct kvm_memory_slot *memslot);
 void kvm_init_vmcs(struct kvm *kvm);
 void kvm_exc_entry(void);
diff --git a/arch/loongarch/kvm/main.c b/arch/loongarch/kvm/main.c
index a98f0284134a..0099d0f0b803 100644
--- a/arch/loongarch/kvm/main.c
+++ b/arch/loongarch/kvm/main.c
@@ -218,7 +218,10 @@ static void kvm_update_vpid(struct kvm_vcpu *vcpu, int cpu)
 		++vpid; /* vpid 0 reserved for root */
 
 		/* start new vpid cycle */
-		kvm_flush_tlb_all();
+		if (!cpu_has_guestid)
+			kvm_flush_tlb_all();
+		else
+			kvm_flush_tlb_all_stage1();
 	}
 
 	context->vpid_cache = vpid;
@@ -278,15 +281,54 @@ static void __kvm_check_vpid(struct kvm_vcpu *vcpu)
 	}
 }
 
-static void __kvm_check_vmid(struct kvm_vcpu *vcpu)
+static void kvm_update_vmid(struct kvm_vcpu *vcpu, int cpu)
 {
 	unsigned long vmid;
+	struct kvm_context *context;
+
+	context = per_cpu_ptr(vcpu->kvm->arch.vmcs, cpu);
+	vmid = context->vmid_cache + 1;
+	if (!(vmid & vpid_mask)) {
+		/* finish round of vmid loop */
+		if (unlikely(!vmid))
+			vmid = vpid_mask + 1;
+
+		++vmid; /* vmid 0 reserved for root */
+
+		/* start new vmid cycle */
+		kvm_flush_tlb_all_stage2();
+	}
+
+	context->vmid_cache = vmid;
+	vcpu->kvm->arch.vmid[cpu] = vmid;
+}
 
-	vmid = vcpu->arch.vpid & vpid_mask;
-	if (vcpu->arch.tgid != vmid) {
-		vcpu->arch.tgid = vcpu->arch.vpid & vpid_mask;
-		kvm_clear_request(KVM_REQ_TLB_FLUSH_GPA, vcpu);
+static void __kvm_check_vmid(struct kvm_vcpu *vcpu)
+{
+	int cpu;
+	unsigned long ver, old, vmid;
+
+	/* On some machines like 3A5000, vmid needs the same with vpid */
+	if (!cpu_has_guestid) {
+		vmid = vcpu->arch.vpid & vpid_mask;
+		if (vcpu->arch.tgid != vmid) {
+			vcpu->arch.tgid = vcpu->arch.vpid & vpid_mask;
+			kvm_clear_request(KVM_REQ_TLB_FLUSH_GPA, vcpu);
+		}
+		return;
 	}
+
+	cpu = smp_processor_id();
+	if (cpumask_test_and_clear_cpu(cpu, &vcpu->kvm->arch.tlb_flush_pending))
+		vcpu->kvm->arch.vmid[cpu] = 0;
+
+	/* Check if the vmid is out of an older version */
+	ver = vcpu->kvm->arch.vmid[cpu] & ~vpid_mask;
+	old = this_cpu_ptr(vcpu->kvm->arch.vmcs)->vmid_cache & ~vpid_mask;
+	if (ver != old)
+		kvm_update_vmid(vcpu, cpu);
+
+	vcpu->arch.tgid = vcpu->kvm->arch.vmid[cpu] & vpid_mask;
 }
 
 void kvm_check_vpid(struct kvm_vcpu *vcpu)
@@ -392,6 +434,7 @@ static int kvm_loongarch_env_init(void)
 	for_each_possible_cpu(cpu) {
 		context = per_cpu_ptr(vmcs, cpu);
 		context->vpid_cache = vpid_mask + 1;
+		context->vmid_cache = vpid_mask + 1;
 		context->last_vcpu = NULL;
 	}
 
diff --git a/arch/loongarch/kvm/mmu.c b/arch/loongarch/kvm/mmu.c
index e104897aa532..e07ba28cc565 100644
--- a/arch/loongarch/kvm/mmu.c
+++ b/arch/loongarch/kvm/mmu.c
@@ -893,7 +893,12 @@ static int kvm_map_page(struct kvm_vcpu *vcpu, unsigned long gpa, bool write)
 		 * there is invalid tlb with small page
 		 * need flush these invalid tlbs for current vcpu
 		 */
-		kvm_make_request(KVM_REQ_TLB_FLUSH, vcpu);
+		if (!cpu_has_ptw) {
+			if (!cpu_has_guestid)
+				kvm_make_request(KVM_REQ_TLB_FLUSH, vcpu);
+			else
+				cpumask_set_cpu(vcpu->cpu, &vcpu->kvm->arch.tlb_flush_pending);
+		}
 		++kvm->stat.hugepages;
 	}  else if (kvm_pte_huge(*ptep) && write)
 		ptep = kvm_split_huge(vcpu, ptep, gfn);
@@ -929,8 +934,11 @@ int kvm_handle_mm_fault(struct kvm_vcpu *vcpu, unsigned long gpa, bool write, in
 		 *
 		 * With SW PTW, invalid TLB is added in TLB refill exception.
 		 */
-		vcpu->arch.flush_gpa = gpa;
-		kvm_make_request(KVM_REQ_TLB_FLUSH_GPA, vcpu);
+		if (!cpu_has_guestid) {
+			vcpu->arch.flush_gpa = gpa;
+			kvm_make_request(KVM_REQ_TLB_FLUSH_GPA, vcpu);
+		} else
+			cpumask_set_cpu(vcpu->cpu, &vcpu->kvm->arch.tlb_flush_pending);
 	}
 
 	return 0;
@@ -940,6 +948,22 @@ void kvm_arch_sync_dirty_log(struct kvm *kvm, struct kvm_memory_slot *memslot)
 {
 }
 
+int kvm_arch_flush_remote_tlbs(struct kvm *kvm)
+{
+	/*
+	 * Queue a TLB invalidation for each CPU to perform on next
+	 * vcpu loading
+	 */
+	if (cpu_has_guestid) {
+		cpumask_setall(&kvm->arch.tlb_flush_pending);
+		/* Be sure that other CPUS can watch the changes */
+		smp_wmb();
+	}
+
+	/* Return 1 continue to send ipi to running vCPUs */
+	return 1;
+}
+
 void kvm_arch_flush_remote_tlbs_memslot(struct kvm *kvm,
 					const struct kvm_memory_slot *memslot)
 {
diff --git a/arch/loongarch/kvm/tlb.c b/arch/loongarch/kvm/tlb.c
index afcacb62988a..994216c98ea2 100644
--- a/arch/loongarch/kvm/tlb.c
+++ b/arch/loongarch/kvm/tlb.c
@@ -21,6 +21,20 @@ void kvm_flush_tlb_all(void)
 	local_irq_restore(flags);
 }
 
+/* Invalidate all stage1 TLB entries including GVA-->GPA mappings */
+void kvm_flush_tlb_all_stage1(void)
+{
+	lockdep_assert_irqs_disabled();
+	invtlb_all(INVGTLB_ALLGID_GVA_TO_GPA, 0, 0);
+}
+
+/* Invalidate all stage2 TLB entries including GPA-->HPA  mappings */
+void kvm_flush_tlb_all_stage2(void)
+{
+	lockdep_assert_irqs_disabled();
+	invtlb_all(INVTLB_ALLGID_GPA_TO_HPA, 0, 0);
+}
+
 void kvm_flush_tlb_gpa(struct kvm_vcpu *vcpu, unsigned long gpa)
 {
 	unsigned int vmid;
-- 
2.39.3