Re: [PATCH v2 1/4] KVM: arm64: nv: Keep reference on stage-2 MMU when scheduled out

[Sean Christopherson <seanjc@google.com>]

On Mon, Oct 07, 2024, Oliver Upton wrote:
> If a vCPU is scheduling out and not in WFI emulation, it is highly
> likely it will get scheduled again soon and reuse the MMU it had before.
> Dropping the MMU at vcpu_put() can have some unfortunate consequences,
> as the MMU could get reclaimed and used in a different context, forcing
> another 'cold start' on an otherwise active MMU.
> 
> Avoid that altogether by keeping a reference on the MMU if the vCPU is
> scheduling out, ensuring that another vCPU cannot reclaim it while the
> current vCPU is away. Since there are more MMUs than vCPUs, this does
> not affect the guarantee that an unused MMU is available at any time.
> 
> Furthermore, this makes the vcpu->arch.hw_mmu ~stable in preemptible
> code, at least for where it matters in the stage-2 abort path. Yes, the
> MMU can change across WFI emulation, but there isn't even a use case
> where this would matter.
> 
> Signed-off-by: Oliver Upton <oliver.upton@linux.dev>
> ---
>  arch/arm64/kvm/nested.c | 21 ++++++++++++++++++---
>  1 file changed, 18 insertions(+), 3 deletions(-)
> 
> diff --git a/arch/arm64/kvm/nested.c b/arch/arm64/kvm/nested.c
> index f9e30dd34c7a..df670c14e1c6 100644
> --- a/arch/arm64/kvm/nested.c
> +++ b/arch/arm64/kvm/nested.c
> @@ -663,6 +663,13 @@ void kvm_init_nested_s2_mmu(struct kvm_s2_mmu *mmu)
>  
>  void kvm_vcpu_load_hw_mmu(struct kvm_vcpu *vcpu)
>  {
> +	/*
> +	 * The vCPU kept its reference on the MMU after the last put, keep
> +	 * rolling with it.
> +	 */
> +	if (vcpu->arch.hw_mmu)
> +		return;
> +
>  	if (is_hyp_ctxt(vcpu)) {
>  		vcpu->arch.hw_mmu = &vcpu->kvm->arch.mmu;
>  	} else {
> @@ -674,10 +681,18 @@ void kvm_vcpu_load_hw_mmu(struct kvm_vcpu *vcpu)
>  
>  void kvm_vcpu_put_hw_mmu(struct kvm_vcpu *vcpu)
>  {
> -	if (kvm_is_nested_s2_mmu(vcpu->kvm, vcpu->arch.hw_mmu)) {
> +	/*
> +	 * Keep a reference on the associated stage-2 MMU if the vCPU is
> +	 * scheduling out and not in WFI emulation, suggesting it is likely to
> +	 * reuse the MMU sometime soon.
> +	 */
> +	if (vcpu->scheduled_out && !vcpu_get_flag(vcpu, IN_WFI))
> +		return;

Assuming KVM arm64 supports halt-polling, I think it makes more sense to check
kvm_vcpu_is_blocking() instead of IN_WFI.  That way, KVM will keep the MMU resident
if the vCPU happens to be preempted and halt-polling is successful.

And somewhat of a side topic, calling kvm_vgic_put() from kvm_arch_vcpu_blocking()
instead of kvm_vcpu_wfi() would provide the same optimization for the GIC, as KVM
would only need to put/reload the vGIC if the vCPU is actually scheduled out.

Unfortunately, using kvm_vcpu_is_blocking() in vgic_v4_put() won't do the right
thing, as KVM deliberately calls prepare_to_rcuwait() _after_ kvm_arch_vcpu_blocking(),
because x86's AVIC implementation relies on kvm_vcpu_is_blocking() to be false
when reloading after unblocking.

Aha!  With vcpu->scheduled_out, I think we can swap the order and make it more
obvious what's going on in the AVIC code at the same time.  With that, I think
KVM arm64 could drop IN_WFI and use kvm_vcpu_is_blocking() throughout.  E.g.

diff --git a/arch/x86/kvm/svm/avic.c b/arch/x86/kvm/svm/avic.c
index 4b74ea91f4e6..bba93d1d65ae 100644
--- a/arch/x86/kvm/svm/avic.c
+++ b/arch/x86/kvm/svm/avic.c
@@ -1038,13 +1038,12 @@ void avic_vcpu_load(struct kvm_vcpu *vcpu, int cpu)
                return;
 
        /*
-        * No need to update anything if the vCPU is blocking, i.e. if the vCPU
-        * is being scheduled in after being preempted.  The CPU entries in the
-        * Physical APIC table and IRTE are consumed iff IsRun{ning} is '1'.
-        * If the vCPU was migrated, its new CPU value will be stuffed when the
-        * vCPU unblocks.
+        * No need to update anything if the vCPU is preempted while blocking.
+        * The CPU entries in the Physical APIC table and IRTE are consumed iff
+        * IsRun{ning} is '1'.  If the vCPU was migrated, its new CPU value
+        * will be stuffed when the vCPU unblocks.
         */
-       if (kvm_vcpu_is_blocking(vcpu))
+       if (kvm_vcpu_is_blocking(vcpu) && vcpu->scheduled_out)
                return;
 
        /*