[RFC PATCH 07/24] KVM: SEV: Track ASID->vCPU on vCPU load

[Yosry Ahmed <yosry.ahmed@linux.dev>]

Check for changes in the ASID to vCPU mapping on vCPU load instead of
doing it on vCPU run. This should be sufficient and more efficient, and
is needed to allow generalizing the tracking and making it more
expensive.

Signed-off-by: Yosry Ahmed <yosry.ahmed@linux.dev>
---
 arch/x86/kvm/svm/sev.c | 13 ++++---------
 arch/x86/kvm/svm/svm.c | 13 +++++++++++++
 arch/x86/kvm/svm/svm.h |  1 +
 3 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/arch/x86/kvm/svm/sev.c b/arch/x86/kvm/svm/sev.c
index ddb4d5b211ed7..3ef0dfdbb34d2 100644
--- a/arch/x86/kvm/svm/sev.c
+++ b/arch/x86/kvm/svm/sev.c
@@ -224,7 +224,7 @@ static int sev_asid_new(struct kvm_sev_info *sev)
 	return ret;
 }
 
-static unsigned int sev_get_asid(struct kvm *kvm)
+unsigned int sev_get_asid(struct kvm *kvm)
 {
 	return to_kvm_sev_info(kvm)->asid;
 }
@@ -3453,7 +3453,6 @@ void sev_es_unmap_ghcb(struct vcpu_svm *svm)
 
 int pre_sev_run(struct vcpu_svm *svm, int cpu)
 {
-	struct svm_cpu_data *sd = per_cpu_ptr(&svm_data, cpu);
 	struct kvm *kvm = svm->vcpu.kvm;
 	unsigned int asid = sev_get_asid(kvm);
 
@@ -3469,16 +3468,12 @@ int pre_sev_run(struct vcpu_svm *svm, int cpu)
 	svm->asid = asid;
 
 	/*
-	 * Flush guest TLB:
-	 *
-	 * 1) when different vCPU for the same ASID is to be run on the same host CPU.
-	 * 2) or this VMCB was executed on different host CPU in previous VMRUNs.
+	 * Flush guest TLB if the VMCB was executed on a differet host CPU in
+	 * previous VMRUNs.
 	 */
-	if (sd->sev_vcpus[asid] == &svm->vcpu &&
-	    svm->vcpu.arch.last_vmentry_cpu == cpu)
+	if (svm->vcpu.arch.last_vmentry_cpu == cpu)
 		return 0;
 
-	sd->sev_vcpus[asid] = &svm->vcpu;
 	vmcb_set_flush_asid(svm->vmcb);
 	vmcb_mark_dirty(svm->vmcb, VMCB_ASID);
 	return 0;
diff --git a/arch/x86/kvm/svm/svm.c b/arch/x86/kvm/svm/svm.c
index 1156ca97fd798..e6e380411fbec 100644
--- a/arch/x86/kvm/svm/svm.c
+++ b/arch/x86/kvm/svm/svm.c
@@ -1554,6 +1554,7 @@ static void svm_prepare_host_switch(struct kvm_vcpu *vcpu)
 
 static void svm_vcpu_load(struct kvm_vcpu *vcpu, int cpu)
 {
+	unsigned int asid;
 	struct vcpu_svm *svm = to_svm(vcpu);
 	struct svm_cpu_data *sd = per_cpu_ptr(&svm_data, cpu);
 
@@ -1568,6 +1569,18 @@ static void svm_vcpu_load(struct kvm_vcpu *vcpu, int cpu)
 	}
 	if (kvm_vcpu_apicv_active(vcpu))
 		avic_vcpu_load(vcpu, cpu);
+
+	if (sev_guest(vcpu->kvm)) {
+		/*
+		 * Flush the TLB when a different vCPU using the same ASID is
+		 * run on the same CPU.
+		 */
+		asid = sev_get_asid(vcpu->kvm);
+		if (sd->sev_vcpus[asid] != vcpu) {
+			sd->sev_vcpus[asid] = vcpu;
+			kvm_make_request(KVM_REQ_TLB_FLUSH, vcpu);
+		}
+	}
 }
 
 static void svm_vcpu_put(struct kvm_vcpu *vcpu)
diff --git a/arch/x86/kvm/svm/svm.h b/arch/x86/kvm/svm/svm.h
index 4ea6c61c3b048..ca38a233fa24c 100644
--- a/arch/x86/kvm/svm/svm.h
+++ b/arch/x86/kvm/svm/svm.h
@@ -768,6 +768,7 @@ void sev_es_vcpu_reset(struct vcpu_svm *svm);
 void sev_vcpu_deliver_sipi_vector(struct kvm_vcpu *vcpu, u8 vector);
 void sev_es_prepare_switch_to_guest(struct vcpu_svm *svm, struct sev_es_save_area *hostsa);
 void sev_es_unmap_ghcb(struct vcpu_svm *svm);
+unsigned int sev_get_asid(struct kvm *kvm);
 
 #ifdef CONFIG_KVM_AMD_SEV
 int sev_mem_enc_ioctl(struct kvm *kvm, void __user *argp);
-- 
2.49.0.395.g12beb8f557-goog