[PATCH 4/4] KVM: x86: Load guest/host PKRU outside of the fastpath run loop

Sean Christopherson posted 4 patches 3 months, 1 week ago
There is a newer version of this series
[PATCH 4/4] KVM: x86: Load guest/host PKRU outside of the fastpath run loop
Posted by Sean Christopherson 3 months, 1 week ago
Move KVM's swapping of PKRU outside of the fastpath loop, as there is no
KVM code anywhere in the fastpath that accesses guest/userspace memory,
i.e. that can consume protection keys.

As documented by commit 1be0e61c1f25 ("KVM, pkeys: save/restore PKRU when
guest/host switches"), KVM just needs to ensure the host's PKRU is loaded
when KVM (or the kernel at-large) may access userspace memory.  And at the
time of commit 1be0e61c1f25, KVM didn't have a fastpath, and PKU was
strictly contained to VMX, i.e. there was no reason to swap PKRU outside
of vmx_vcpu_run().

Over time, the "need" to swap PKRU close to VM-Enter was likely falsely
solidified by the association with XFEATUREs in commit 37486135d3a7
("KVM: x86: Fix pkru save/restore when guest CR4.PKE=0, move it to x86.c"),
and XFEATURE swapping was in turn moved close to VM-Enter/VM-Exit as a
KVM hack-a-fix ution for an #MC handler bug by commit 1811d979c716
("x86/kvm: move kvm_load/put_guest_xcr0 into atomic context").

Deferring the PKRU loads shaves ~40 cycles off the fastpath for Intel,
and ~60 cycles for AMD.  E.g. using INVD in KVM-Unit-Test's vmexit.c,
with extra hacks to enable CR4.PKE and PKRU=(-1u & ~0x3), latency numbers
for AMD Turin go from ~1560 => ~1500, and for Intel Emerald Rapids, go
from ~810 => ~770.

Signed-off-by: Sean Christopherson <seanjc@google.com>
---
 arch/x86/kvm/svm/svm.c |  2 --
 arch/x86/kvm/vmx/vmx.c |  4 ----
 arch/x86/kvm/x86.c     | 14 ++++++++++----
 arch/x86/kvm/x86.h     |  2 --
 4 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/arch/x86/kvm/svm/svm.c b/arch/x86/kvm/svm/svm.c
index e8b158f73c79..e1fb853c263c 100644
--- a/arch/x86/kvm/svm/svm.c
+++ b/arch/x86/kvm/svm/svm.c
@@ -4260,7 +4260,6 @@ static __no_kcsan fastpath_t svm_vcpu_run(struct kvm_vcpu *vcpu, u64 run_flags)
 		svm_set_dr6(vcpu, DR6_ACTIVE_LOW);
 
 	clgi();
-	kvm_load_guest_xsave_state(vcpu);
 
 	/*
 	 * Hardware only context switches DEBUGCTL if LBR virtualization is
@@ -4303,7 +4302,6 @@ static __no_kcsan fastpath_t svm_vcpu_run(struct kvm_vcpu *vcpu, u64 run_flags)
 	    vcpu->arch.host_debugctl != svm->vmcb->save.dbgctl)
 		update_debugctlmsr(vcpu->arch.host_debugctl);
 
-	kvm_load_host_xsave_state(vcpu);
 	stgi();
 
 	/* Any pending NMI will happen here */
diff --git a/arch/x86/kvm/vmx/vmx.c b/arch/x86/kvm/vmx/vmx.c
index 123dae8cf46b..55d637cea84a 100644
--- a/arch/x86/kvm/vmx/vmx.c
+++ b/arch/x86/kvm/vmx/vmx.c
@@ -7465,8 +7465,6 @@ fastpath_t vmx_vcpu_run(struct kvm_vcpu *vcpu, u64 run_flags)
 	if (vcpu->guest_debug & KVM_GUESTDBG_SINGLESTEP)
 		vmx_set_interrupt_shadow(vcpu, 0);
 
-	kvm_load_guest_xsave_state(vcpu);
-
 	pt_guest_enter(vmx);
 
 	atomic_switch_perf_msrs(vmx);
@@ -7510,8 +7508,6 @@ fastpath_t vmx_vcpu_run(struct kvm_vcpu *vcpu, u64 run_flags)
 
 	pt_guest_exit(vmx);
 
-	kvm_load_host_xsave_state(vcpu);
-
 	if (is_guest_mode(vcpu)) {
 		/*
 		 * Track VMLAUNCH/VMRESUME that have made past guest state
diff --git a/arch/x86/kvm/x86.c b/arch/x86/kvm/x86.c
index b5c2879e3330..6924006f0796 100644
--- a/arch/x86/kvm/x86.c
+++ b/arch/x86/kvm/x86.c
@@ -1233,7 +1233,7 @@ static void kvm_load_host_xfeatures(struct kvm_vcpu *vcpu)
 	}
 }
 
-void kvm_load_guest_xsave_state(struct kvm_vcpu *vcpu)
+static void kvm_load_guest_pkru(struct kvm_vcpu *vcpu)
 {
 	if (vcpu->arch.guest_state_protected)
 		return;
@@ -1244,9 +1244,8 @@ void kvm_load_guest_xsave_state(struct kvm_vcpu *vcpu)
 	     kvm_is_cr4_bit_set(vcpu, X86_CR4_PKE)))
 		wrpkru(vcpu->arch.pkru);
 }
-EXPORT_SYMBOL_FOR_KVM_INTERNAL(kvm_load_guest_xsave_state);
 
-void kvm_load_host_xsave_state(struct kvm_vcpu *vcpu)
+static void kvm_load_host_pkru(struct kvm_vcpu *vcpu)
 {
 	if (vcpu->arch.guest_state_protected)
 		return;
@@ -1259,7 +1258,6 @@ void kvm_load_host_xsave_state(struct kvm_vcpu *vcpu)
 			wrpkru(vcpu->arch.host_pkru);
 	}
 }
-EXPORT_SYMBOL_FOR_KVM_INTERNAL(kvm_load_host_xsave_state);
 
 #ifdef CONFIG_X86_64
 static inline u64 kvm_guest_supported_xfd(struct kvm_vcpu *vcpu)
@@ -11331,6 +11329,12 @@ static int vcpu_enter_guest(struct kvm_vcpu *vcpu)
 
 	guest_timing_enter_irqoff();
 
+	/*
+	 * Swap PKRU with hardware breakpoints disabled to minimize the number
+	 * of flows where non-KVM code can run with guest state loaded.
+	 */
+	kvm_load_guest_pkru(vcpu);
+
 	for (;;) {
 		/*
 		 * Assert that vCPU vs. VM APICv state is consistent.  An APICv
@@ -11359,6 +11363,8 @@ static int vcpu_enter_guest(struct kvm_vcpu *vcpu)
 		++vcpu->stat.exits;
 	}
 
+	kvm_load_host_pkru(vcpu);
+
 	/*
 	 * Do this here before restoring debug registers on the host.  And
 	 * since we do this before handling the vmexit, a DR access vmexit
diff --git a/arch/x86/kvm/x86.h b/arch/x86/kvm/x86.h
index f3dc77f006f9..24c754b0db2e 100644
--- a/arch/x86/kvm/x86.h
+++ b/arch/x86/kvm/x86.h
@@ -622,8 +622,6 @@ static inline void kvm_machine_check(void)
 #endif
 }
 
-void kvm_load_guest_xsave_state(struct kvm_vcpu *vcpu);
-void kvm_load_host_xsave_state(struct kvm_vcpu *vcpu);
 int kvm_spec_ctrl_test_value(u64 value);
 int kvm_handle_memory_failure(struct kvm_vcpu *vcpu, int r,
 			      struct x86_exception *e);
-- 
2.51.1.930.gacf6e81ea2-goog
Re: [PATCH 4/4] KVM: x86: Load guest/host PKRU outside of the fastpath run loop
Posted by Jon Kohler 3 months, 1 week ago

> On Oct 30, 2025, at 6:42 PM, Sean Christopherson <seanjc@google.com> wrote:
> 
> !-------------------------------------------------------------------|
>  CAUTION: External Email
> 
> |-------------------------------------------------------------------!
> 
> Move KVM's swapping of PKRU outside of the fastpath loop, as there is no
> KVM code anywhere in the fastpath that accesses guest/userspace memory,
> i.e. that can consume protection keys.
> 
> As documented by commit 1be0e61c1f25 ("KVM, pkeys: save/restore PKRU when
> guest/host switches"), KVM just needs to ensure the host's PKRU is loaded
> when KVM (or the kernel at-large) may access userspace memory.  And at the
> time of commit 1be0e61c1f25, KVM didn't have a fastpath, and PKU was
> strictly contained to VMX, i.e. there was no reason to swap PKRU outside
> of vmx_vcpu_run().
> 
> Over time, the "need" to swap PKRU close to VM-Enter was likely falsely
> solidified by the association with XFEATUREs in commit 37486135d3a7
> ("KVM: x86: Fix pkru save/restore when guest CR4.PKE=0, move it to x86.c"),
> and XFEATURE swapping was in turn moved close to VM-Enter/VM-Exit as a
> KVM hack-a-fix ution for an #MC handler bug by commit 1811d979c716
> ("x86/kvm: move kvm_load/put_guest_xcr0 into atomic context").
> 
> Deferring the PKRU loads shaves ~40 cycles off the fastpath for Intel,
> and ~60 cycles for AMD.  E.g. using INVD in KVM-Unit-Test's vmexit.c,
> with extra hacks to enable CR4.PKE and PKRU=(-1u & ~0x3), latency numbers
> for AMD Turin go from ~1560 => ~1500, and for Intel Emerald Rapids, go
> from ~810 => ~770.
> 
> Signed-off-by: Sean Christopherson <seanjc@google.com>
> ---
> arch/x86/kvm/svm/svm.c |  2 --
> arch/x86/kvm/vmx/vmx.c |  4 ----
> arch/x86/kvm/x86.c     | 14 ++++++++++----
> arch/x86/kvm/x86.h     |  2 --
> 4 files changed, 10 insertions(+), 12 deletions(-)
> 
> diff --git a/arch/x86/kvm/svm/svm.c b/arch/x86/kvm/svm/svm.c
> index e8b158f73c79..e1fb853c263c 100644
> --- a/arch/x86/kvm/svm/svm.c
> +++ b/arch/x86/kvm/svm/svm.c
> @@ -4260,7 +4260,6 @@ static __no_kcsan fastpath_t svm_vcpu_run(struct kvm_vcpu *vcpu, u64 run_flags)
> svm_set_dr6(vcpu, DR6_ACTIVE_LOW);
> 
> clgi();
> - kvm_load_guest_xsave_state(vcpu);
> 
> /*
> * Hardware only context switches DEBUGCTL if LBR virtualization is
> @@ -4303,7 +4302,6 @@ static __no_kcsan fastpath_t svm_vcpu_run(struct kvm_vcpu *vcpu, u64 run_flags)
>    vcpu->arch.host_debugctl != svm->vmcb->save.dbgctl)
> update_debugctlmsr(vcpu->arch.host_debugctl);
> 
> - kvm_load_host_xsave_state(vcpu);
> stgi();
> 
> /* Any pending NMI will happen here */
> diff --git a/arch/x86/kvm/vmx/vmx.c b/arch/x86/kvm/vmx/vmx.c
> index 123dae8cf46b..55d637cea84a 100644
> --- a/arch/x86/kvm/vmx/vmx.c
> +++ b/arch/x86/kvm/vmx/vmx.c
> @@ -7465,8 +7465,6 @@ fastpath_t vmx_vcpu_run(struct kvm_vcpu *vcpu, u64 run_flags)
> if (vcpu->guest_debug & KVM_GUESTDBG_SINGLESTEP)
> vmx_set_interrupt_shadow(vcpu, 0);
> 
> - kvm_load_guest_xsave_state(vcpu);
> -
> pt_guest_enter(vmx);
> 
> atomic_switch_perf_msrs(vmx);
> @@ -7510,8 +7508,6 @@ fastpath_t vmx_vcpu_run(struct kvm_vcpu *vcpu, u64 run_flags)
> 
> pt_guest_exit(vmx);
> 
> - kvm_load_host_xsave_state(vcpu);
> -
> if (is_guest_mode(vcpu)) {
> /*
> * Track VMLAUNCH/VMRESUME that have made past guest state
> diff --git a/arch/x86/kvm/x86.c b/arch/x86/kvm/x86.c
> index b5c2879e3330..6924006f0796 100644
> --- a/arch/x86/kvm/x86.c
> +++ b/arch/x86/kvm/x86.c
> @@ -1233,7 +1233,7 @@ static void kvm_load_host_xfeatures(struct kvm_vcpu *vcpu)
> }
> }
> 
> -void kvm_load_guest_xsave_state(struct kvm_vcpu *vcpu)
> +static void kvm_load_guest_pkru(struct kvm_vcpu *vcpu)
> {
> if (vcpu->arch.guest_state_protected)
> return;
> @@ -1244,9 +1244,8 @@ void kvm_load_guest_xsave_state(struct kvm_vcpu *vcpu)
>     kvm_is_cr4_bit_set(vcpu, X86_CR4_PKE)))
> wrpkru(vcpu->arch.pkru);
> }
> -EXPORT_SYMBOL_FOR_KVM_INTERNAL(kvm_load_guest_xsave_state);
> 
> -void kvm_load_host_xsave_state(struct kvm_vcpu *vcpu)
> +static void kvm_load_host_pkru(struct kvm_vcpu *vcpu)
> {
> if (vcpu->arch.guest_state_protected)
> return;
> @@ -1259,7 +1258,6 @@ void kvm_load_host_xsave_state(struct kvm_vcpu *vcpu)
> wrpkru(vcpu->arch.host_pkru);
> }
> }
> -EXPORT_SYMBOL_FOR_KVM_INTERNAL(kvm_load_host_xsave_state);
> 
> #ifdef CONFIG_X86_64
> static inline u64 kvm_guest_supported_xfd(struct kvm_vcpu *vcpu)
> @@ -11331,6 +11329,12 @@ static int vcpu_enter_guest(struct kvm_vcpu *vcpu)
> 
> guest_timing_enter_irqoff();
> 
> + /*
> + * Swap PKRU with hardware breakpoints disabled to minimize the number
> + * of flows where non-KVM code can run with guest state loaded.
> + */
> + kvm_load_guest_pkru(vcpu);
> +

I was mocking this up after PUCK, and went down a similar-ish path, but was
thinking it might be interesting to have an x86 op called something to the effect of
“prepare_switch_to_guest_irqoff” and “prepare_switch_to_host_irqoff”, which
might make for a place to nestle any other sort of “needs to be done in atomic
context but doesn’t need to be done in the fast path” sort of stuff (if any).

One other one that caught my eye was the cr3 stuff that was moved out a while
ago, but then moved back with 1a7158101.

I haven’t gone through absolutely everything else in that tight loop code (and didn’t
get a chance to do the same for SVM code), but figured I’d put the idea out there
to see what you think.

To be clear, I’m totally OK with the series as-is, just thinking about perhaps future
ways to incrementally optimize here?

> for (;;) {
> /*
> * Assert that vCPU vs. VM APICv state is consistent.  An APICv
> @@ -11359,6 +11363,8 @@ static int vcpu_enter_guest(struct kvm_vcpu *vcpu)
> ++vcpu->stat.exits;
> }
> 
> + kvm_load_host_pkru(vcpu);
> +
> /*
> * Do this here before restoring debug registers on the host.  And
> * since we do this before handling the vmexit, a DR access vmexit
> diff --git a/arch/x86/kvm/x86.h b/arch/x86/kvm/x86.h
> index f3dc77f006f9..24c754b0db2e 100644
> --- a/arch/x86/kvm/x86.h
> +++ b/arch/x86/kvm/x86.h
> @@ -622,8 +622,6 @@ static inline void kvm_machine_check(void)
> #endif
> }
> 
> -void kvm_load_guest_xsave_state(struct kvm_vcpu *vcpu);
> -void kvm_load_host_xsave_state(struct kvm_vcpu *vcpu);
> int kvm_spec_ctrl_test_value(u64 value);
> int kvm_handle_memory_failure(struct kvm_vcpu *vcpu, int r,
>      struct x86_exception *e);
> -- 
> 2.51.1.930.gacf6e81ea2-goog
> 

Re: [PATCH 4/4] KVM: x86: Load guest/host PKRU outside of the fastpath run loop
Posted by Sean Christopherson 3 months, 1 week ago
On Fri, Oct 31, 2025, Jon Kohler wrote:
> > On Oct 30, 2025, at 6:42 PM, Sean Christopherson <seanjc@google.com> wrote:
> > + /*
> > + * Swap PKRU with hardware breakpoints disabled to minimize the number
> > + * of flows where non-KVM code can run with guest state loaded.
> > + */
> > + kvm_load_guest_pkru(vcpu);
> > +
> 
> I was mocking this up after PUCK, and went down a similar-ish path, but was
> thinking it might be interesting to have an x86 op called something to the effect of
> “prepare_switch_to_guest_irqoff” and “prepare_switch_to_host_irqoff”, which
> might make for a place to nestle any other sort of “needs to be done in atomic
> context but doesn’t need to be done in the fast path” sort of stuff (if any).

Hmm, I would say I'm flat out opposed to generic hooks of that nature.  For
anything that _needs_ to be modified with IRQs disabled, the ordering will matter
greatly.  E.g. we already have kvm_x86_ops.sync_pir_to_irr(), and that _must_ run
before kvm_vcpu_exit_request() if it triggers a late request.

And I also want to push for as much stuff as possible to be handled in common x86,
i.e. I want to actively encourage landing things like PKU and CET support in
common x86 instead of implementing support in one vendor and then having to churn
a pile of code to later move it to

> One other one that caught my eye was the cr3 stuff that was moved out a while
> ago, but then moved back with 1a7158101.
> 
> I haven’t gone through absolutely everything else in that tight loop code
> (and didn’t get a chance to do the same for SVM code), but figured I’d put
> the idea out there to see what you think.
> 
> To be clear, I’m totally OK with the series as-is, just thinking about
> perhaps future ways to incrementally optimize here?

To some extent, we're going to hit diminishing returns.  E.g. one of the reasons
I did a straight revert in commit 1a71581012dd is that were talking about a handful
of cycles difference.  E.g. as measured from the guest, eliding the CR3+CR4 checks
shaves 3-5 cycles.  From the host side it _looks_ like more (~20 cycles), but it's
hard to even measure accurately because just doing RDTSC affects the results.

For SVM, I don't see any obvious candidates.  E.g. pre_sev_run() has some code that
only needs to be done on the first iteration, but checking a flag or doing a static
CALL+RET is going to be just as costly as what's already there.

In short, the only flows that will benefit are relatively slow flows and/or flows
that aren't easily predicted by the CPU.  E.g. __get_current_cr3_fast() and
cr4_read_shadow() require CALL+RET and might not be super predictable?  But even
they are on the cusp of "who cares".

And that needs to be balanced against the probability of introducing bugs.  E.g.
this code _could_ be done only on the first iteration:

	if (vmx->ple_window_dirty) {
		vmx->ple_window_dirty = false;
		vmcs_write32(PLE_WINDOW, vmx->ple_window);
	}

but (a) checking vmx->ple_window_dirty is going to be super predictable after the
first iteration, (b) handling PLE exits in the fastpath would break things, and
(c) _if_ we want to optimize that code, it can/should be simply moved to
vmx_prepare_switch_to_guest() (but outside of the guest_state_loaded check).

All that said, I'm not totally opposed to shaving cycles.  Now that @run_flags
is a thing, it's actually trivially easy to optimize the CR3/CR4 checks (famous
last words):

diff --git a/arch/x86/include/asm/kvm_host.h b/arch/x86/include/asm/kvm_host.h
index 48598d017d6f..5cc1f0168b8a 100644
--- a/arch/x86/include/asm/kvm_host.h
+++ b/arch/x86/include/asm/kvm_host.h
@@ -1709,6 +1709,7 @@ enum kvm_x86_run_flags {
        KVM_RUN_FORCE_IMMEDIATE_EXIT    = BIT(0),
        KVM_RUN_LOAD_GUEST_DR6          = BIT(1),
        KVM_RUN_LOAD_DEBUGCTL           = BIT(2),
+       KVM_RUN_IS_FIRST_ITERATION      = BIT(3),
 };
 
 struct kvm_x86_ops {
diff --git a/arch/x86/kvm/vmx/vmx.c b/arch/x86/kvm/vmx/vmx.c
index 55d637cea84a..3deb20b8d0c5 100644
--- a/arch/x86/kvm/vmx/vmx.c
+++ b/arch/x86/kvm/vmx/vmx.c
@@ -7439,22 +7439,28 @@ fastpath_t vmx_vcpu_run(struct kvm_vcpu *vcpu, u64 run_flags)
                vmx_reload_guest_debugctl(vcpu);
 
        /*
-        * Refresh vmcs.HOST_CR3 if necessary.  This must be done immediately
-        * prior to VM-Enter, as the kernel may load a new ASID (PCID) any time
-        * it switches back to the current->mm, which can occur in KVM context
-        * when switching to a temporary mm to patch kernel code, e.g. if KVM
-        * toggles a static key while handling a VM-Exit.
+        * Refresh vmcs.HOST_CR3 if necessary.  This must be done after IRQs
+        * are disabled, i.e. not when preparing to switch to the guest, as the
+        * the kernel may load a new ASID (PCID) any time it switches back to
+        * the current->mm, which can occur in KVM context when switching to a
+        * temporary mm to patch kernel code, e.g. if KVM toggles a static key
+        * while handling a VM-Exit.
+        *
+        * Refresh host CR3 and CR4 only on the first iteration of the inner
+        * loop, as modifying CR3 or CR4 from NMI context is not allowed.
         */
-       cr3 = __get_current_cr3_fast();
-       if (unlikely(cr3 != vmx->loaded_vmcs->host_state.cr3)) {
-               vmcs_writel(HOST_CR3, cr3);
-               vmx->loaded_vmcs->host_state.cr3 = cr3;
-       }
+       if (run_flags & KVM_RUN_IS_FIRST_ITERATION) {
+               cr3 = __get_current_cr3_fast();
+               if (unlikely(cr3 != vmx->loaded_vmcs->host_state.cr3)) {
+                       vmcs_writel(HOST_CR3, cr3);
+                       vmx->loaded_vmcs->host_state.cr3 = cr3;
+               }
 
-       cr4 = cr4_read_shadow();
-       if (unlikely(cr4 != vmx->loaded_vmcs->host_state.cr4)) {
-               vmcs_writel(HOST_CR4, cr4);
-               vmx->loaded_vmcs->host_state.cr4 = cr4;
+               cr4 = cr4_read_shadow();
+               if (unlikely(cr4 != vmx->loaded_vmcs->host_state.cr4)) {
+                       vmcs_writel(HOST_CR4, cr4);
+                       vmx->loaded_vmcs->host_state.cr4 = cr4;
+               }
        }
 
        /* When single-stepping over STI and MOV SS, we must clear the
diff --git a/arch/x86/kvm/x86.c b/arch/x86/kvm/x86.c
index 6924006f0796..bff08f58c29a 100644
--- a/arch/x86/kvm/x86.c
+++ b/arch/x86/kvm/x86.c
@@ -11286,7 +11286,7 @@ static int vcpu_enter_guest(struct kvm_vcpu *vcpu)
                goto cancel_injection;
        }
 
-       run_flags = 0;
+       run_flags = KVM_RUN_IS_FIRST_ITERATION;
        if (req_immediate_exit) {
                run_flags |= KVM_RUN_FORCE_IMMEDIATE_EXIT;
                kvm_make_request(KVM_REQ_EVENT, vcpu);
Re: [PATCH 4/4] KVM: x86: Load guest/host PKRU outside of the fastpath run loop
Posted by Jon Kohler 3 months, 1 week ago

> On Oct 31, 2025, at 4:52 PM, Sean Christopherson <seanjc@google.com> wrote:
> 
> Hmm, I would say I'm flat out opposed to generic hooks of that nature.  For
> anything that _needs_ to be modified with IRQs disabled, the ordering will matter
> greatly.  E.g. we already have kvm_x86_ops.sync_pir_to_irr(), and that _must_ run
> before kvm_vcpu_exit_request() if it triggers a late request.
> 
> And I also want to push for as much stuff as possible to be handled in common x86,
> i.e. I want to actively encourage landing things like PKU and CET support in
> common x86 instead of implementing support in one vendor and then having to churn
> a pile of code to later move it to

Fair, agreed having things common-ized helps everyone

> All that said, I'm not totally opposed to shaving cycles.  Now that @run_flags
> is a thing, it's actually trivially easy to optimize the CR3/CR4 checks (famous
> last words):

A cycle saved is a cycle earned, perhaps? :)

> diff --git a/arch/x86/include/asm/kvm_host.h b/arch/x86/include/asm/kvm_host.h
> index 48598d017d6f..5cc1f0168b8a 100644
> --- a/arch/x86/include/asm/kvm_host.h
> +++ b/arch/x86/include/asm/kvm_host.h
> @@ -1709,6 +1709,7 @@ enum kvm_x86_run_flags {
>        KVM_RUN_FORCE_IMMEDIATE_EXIT    = BIT(0),
>        KVM_RUN_LOAD_GUEST_DR6          = BIT(1),
>        KVM_RUN_LOAD_DEBUGCTL           = BIT(2),
> +       KVM_RUN_IS_FIRST_ITERATION      = BIT(3),
> };

I like this approach, as it makes the code easier to grok what we want and when
> 
> struct kvm_x86_ops {
> diff --git a/arch/x86/kvm/vmx/vmx.c b/arch/x86/kvm/vmx/vmx.c
> index 55d637cea84a..3deb20b8d0c5 100644
> --- a/arch/x86/kvm/vmx/vmx.c
> +++ b/arch/x86/kvm/vmx/vmx.c
> @@ -7439,22 +7439,28 @@ fastpath_t vmx_vcpu_run(struct kvm_vcpu *vcpu, u64 run_flags)
>                vmx_reload_guest_debugctl(vcpu);
> 
>        /*
> -        * Refresh vmcs.HOST_CR3 if necessary.  This must be done immediately
> -        * prior to VM-Enter, as the kernel may load a new ASID (PCID) any time
> -        * it switches back to the current->mm, which can occur in KVM context
> -        * when switching to a temporary mm to patch kernel code, e.g. if KVM
> -        * toggles a static key while handling a VM-Exit.
> +        * Refresh vmcs.HOST_CR3 if necessary.  This must be done after IRQs
> +        * are disabled, i.e. not when preparing to switch to the guest, as the
> +        * the kernel may load a new ASID (PCID) any time it switches back to
> +        * the current->mm, which can occur in KVM context when switching to a
> +        * temporary mm to patch kernel code, e.g. if KVM toggles a static key
> +        * while handling a VM-Exit.
> +        *
> +        * Refresh host CR3 and CR4 only on the first iteration of the inner
> +        * loop, as modifying CR3 or CR4 from NMI context is not allowed.
>         */
> -       cr3 = __get_current_cr3_fast();
> -       if (unlikely(cr3 != vmx->loaded_vmcs->host_state.cr3)) {
> -               vmcs_writel(HOST_CR3, cr3);
> -               vmx->loaded_vmcs->host_state.cr3 = cr3;
> -       }
> +       if (run_flags & KVM_RUN_IS_FIRST_ITERATION) {
> +               cr3 = __get_current_cr3_fast();
> +               if (unlikely(cr3 != vmx->loaded_vmcs->host_state.cr3)) {
> +                       vmcs_writel(HOST_CR3, cr3);
> +                       vmx->loaded_vmcs->host_state.cr3 = cr3;
> +               }
> 
> -       cr4 = cr4_read_shadow();
> -       if (unlikely(cr4 != vmx->loaded_vmcs->host_state.cr4)) {
> -               vmcs_writel(HOST_CR4, cr4);
> -               vmx->loaded_vmcs->host_state.cr4 = cr4;
> +               cr4 = cr4_read_shadow();
> +               if (unlikely(cr4 != vmx->loaded_vmcs->host_state.cr4)) {
> +                       vmcs_writel(HOST_CR4, cr4);
> +                       vmx->loaded_vmcs->host_state.cr4 = cr4;
> +               }
>        }
> 
>        /* When single-stepping over STI and MOV SS, we must clear the
> diff --git a/arch/x86/kvm/x86.c b/arch/x86/kvm/x86.c
> index 6924006f0796..bff08f58c29a 100644
> --- a/arch/x86/kvm/x86.c
> +++ b/arch/x86/kvm/x86.c
> @@ -11286,7 +11286,7 @@ static int vcpu_enter_guest(struct kvm_vcpu *vcpu)
>                goto cancel_injection;
>        }
> 
> -       run_flags = 0;
> +       run_flags = KVM_RUN_IS_FIRST_ITERATION;
>        if (req_immediate_exit) {
>                run_flags |= KVM_RUN_FORCE_IMMEDIATE_EXIT;
>                kvm_make_request(KVM_REQ_EVENT, vcpu);
>