[PATCH 0/4] KVM: x86/mmu: Add auto-prefetch to the TDP MMU

[PATCH 0/4] KVM: x86/mmu: Add auto-prefetch to the TDP MMU

[Sean Christopherson]

Add auto-prefetch to the TDP MMU, where "auto-prefetch" means the existing
prefetch in the shadow MMU where KVM grabs the surrounding ~8 SPTEs on a page
fault, when possible.

Rather than implement yet another custom solution, reuse the core functionality
from KVM_PRE_FAULT_MEMORY.  Doing so has several benefits, and a few drawbacks.
The big benefits are code reuse, reduced maintainence burden improved test
coverage (auto-prefetching exposed a variety of pre-existing issues in the
pre-fault code), and support for any backing store (the shadow MMU's approach
only works for fast-GUP, i.e. only for refcounted struct page memory).

The big drawback is that using the pre-fault code is "heavier" than the custom
prefetching code in the shadow MMU.  In a nutshell, this approach can't achieve
the same performance uplift at least not without additional tweaking (which may
or may not be acceptable; I genuinely haven't thought at all about how we could
close the gap, mainly because it's not clear that we *need* to close the gap).

There are a few subjective decisions in here.  The main one I want to call out
is that I chose to treat async #PF prefetching as manual/explicit prefetching,
versus automatic prefetching.  Manual prefetching will do more invasive
operations like punch through dirty logging and wait for I/O, whereas automatic
prefetching avoids spurious dirty logging and tries to avoid waiting on "major"
faults.  This series preserves the existing behavior, but I spent enough time
waffling on whether or not that was the right choice that I wanted to call it
out here.

Sean Christopherson (4):
  KVM: x86/mmu: Rework gfn_to_memslot_dirty_bitmap() into auto-prefetch
    helper
  KVM: x86/mmu: Differentiate between manual and automatic prefetching
  KVM: x86/mmu: Add automatic prefaulting/prefetching for TDP MMU faults
  KVM: x86/mmu: Use the generic auto-prefaulting to do prefetch for
    direct MMUs

 arch/x86/kvm/mmu/mmu.c          | 156 +++++++++++++++++++-------------
 arch/x86/kvm/mmu/mmu_internal.h |   1 +
 arch/x86/kvm/mmu/paging_tmpl.h  |   2 +-
 3 files changed, 97 insertions(+), 62 deletions(-)


base-commit: 7e60162b16f9da0d2fa83da5c1723b6853a0ac1d
-- 
2.56.0.313.g2479322e68-goog

[PATCH 1/4] KVM: x86/mmu: Rework gfn_to_memslot_dirty_bitmap() into auto-prefetch helper

[Sean Christopherson]

Rework gfn_to_memslot_dirty_bitmap() into a more generic helper that takes
a memslot to check if said a memslot is suitable for auto-prefetch.  "auto"
in this case means prefetch that's initiated by KVM without a corresponding
action from either the host or userspace.  Taking a memslot will allow
using the helper for other prefetch flows where KVM already has the memslot,
i.e. will allow reusing the dirty logging logic in particular, without
introducing potential TOCTOU bugs.

Return error numbers that align with KVM's uAPI for prefaulting so that KVM
is consistent in the unlikely case prefaulting ever ends up being subject
to the same restrictions (prefaulting isn't "auto" in that it's an explicit
request/command from userspace).

No functional change intended.

Signed-off-by: Sean Christopherson <seanjc@google.com>
---
 arch/x86/kvm/mmu/mmu.c | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/arch/x86/kvm/mmu/mmu.c b/arch/x86/kvm/mmu/mmu.c
index 5a2cc6a17a9b..393feeaddfb1 100644
--- a/arch/x86/kvm/mmu/mmu.c
+++ b/arch/x86/kvm/mmu/mmu.c
@@ -872,19 +872,19 @@ static void unaccount_nx_huge_page(struct kvm *kvm, struct kvm_mmu_page *sp,
 	untrack_possible_nx_huge_page(kvm, sp, mmu_type);
 }
 
-static struct kvm_memory_slot *gfn_to_memslot_dirty_bitmap(struct kvm_vcpu *vcpu,
-							   gfn_t gfn,
-							   bool no_dirty_log)
+static int kvm_check_memslot_usable_for_auto_prefetch(struct kvm_memory_slot *slot,
+						      unsigned int access)
 {
-	struct kvm_memory_slot *slot;
+	if (!slot)
+		return -ENOENT;
 
-	slot = kvm_vcpu_gfn_to_memslot(vcpu, gfn);
-	if (!slot || slot->flags & KVM_MEMSLOT_INVALID)
-		return NULL;
-	if (no_dirty_log && kvm_slot_dirty_track_enabled(slot))
-		return NULL;
+	if (slot->flags & KVM_MEMSLOT_INVALID)
+		return -EAGAIN;
 
-	return slot;
+	if ((access & ACC_WRITE_MASK) && kvm_slot_dirty_track_enabled(slot))
+		return -EACCES;
+
+	return 0;
 }
 
 /*
@@ -3231,8 +3231,8 @@ static bool kvm_mmu_prefetch_sptes(struct kvm_vcpu *vcpu, gfn_t gfn, u64 *sptep,
 	if (WARN_ON_ONCE(nr_pages > PTE_PREFETCH_NUM))
 		return false;
 
-	slot = gfn_to_memslot_dirty_bitmap(vcpu, gfn, access & ACC_WRITE_MASK);
-	if (!slot)
+	slot = kvm_vcpu_gfn_to_memslot(vcpu, gfn);
+	if (kvm_check_memslot_usable_for_auto_prefetch(slot, access))
 		return false;
 
 	nr_pages = kvm_prefetch_pages(slot, gfn, pages, nr_pages);
-- 
2.56.0.313.g2479322e68-goog

[PATCH 2/4] KVM: x86/mmu: Differentiate between manual and automatic prefetching

[Sean Christopherson]

Differentiate between "manual" and "automatic" prefetching, where automatic
means the prefetch is initiated by KVM without any corresponding userspace
request/command (KVM_PRE_FAULT_MEMORY) or guest access (async #PF).
Tracking auto-prefetch "faults" will allow implementing prefetching in the
TDP MMU using the standard page fault path, without breaking existing
prefetching rules in the shadow MMU, e.g. don't auto-prefetch if the slot
is being dirty logged.

Arguably, async #PFs could be considered automatic, not manual, but they
_are_ initiated in response to an explicit guest access, and unlike true
automatic prefetching, KVM generally wants to put more effort into handling
async #PFs, e.g. wants to punch through dirty logging or wait for I/O,
precisely because there was a prior guest access and so giving up likely
means a vCPU task will end up paying the cost.

No functional change intended.

Signed-off-by: Sean Christopherson <seanjc@google.com>
---
 arch/x86/kvm/mmu/mmu.c          | 26 ++++++++++++++++++--------
 arch/x86/kvm/mmu/mmu_internal.h |  1 +
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/arch/x86/kvm/mmu/mmu.c b/arch/x86/kvm/mmu/mmu.c
index 393feeaddfb1..3bd0d921778d 100644
--- a/arch/x86/kvm/mmu/mmu.c
+++ b/arch/x86/kvm/mmu/mmu.c
@@ -5031,9 +5031,15 @@ static int kvm_tdp_page_fault(struct kvm_vcpu *vcpu, struct kvm_page_fault *faul
 	return direct_page_fault(vcpu, fault);
 }
 
+enum kvm_prefetch_type {
+	KVM_PREFETCH_NONE,
+	KVM_PREFETCH_MANUAL,
+	KVM_PREFETCH_AUTO,
+};
+
 static int kvm_mmu_do_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa,
-				 u64 err, bool prefetch, int *emulation_type,
-				 u8 *level)
+				 u64 err, enum kvm_prefetch_type prefetch,
+				 int *emulation_type, u8 *level)
 {
 	struct kvm_page_fault fault = {
 		.addr = cr2_or_gpa,
@@ -5043,7 +5049,8 @@ static int kvm_mmu_do_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa,
 		.present = err & PFERR_PRESENT_MASK,
 		.rsvd = err & PFERR_RSVD_MASK,
 		.user = err & PFERR_USER_MASK,
-		.prefetch = prefetch,
+		.prefetch = prefetch != KVM_PREFETCH_NONE,
+		.auto_prefetch = prefetch == KVM_PREFETCH_AUTO,
 		.is_tdp = likely(vcpu->arch.mmu->page_fault == kvm_tdp_page_fault),
 		.nx_huge_page_workaround_enabled =
 			is_nx_huge_page_enabled(vcpu->kvm),
@@ -5097,7 +5104,8 @@ static int kvm_mmu_do_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa,
 
 
 static int kvm_tdp_page_prefault(struct kvm_vcpu *vcpu, gpa_t gpa,
-				 u64 error_code, u8 *level)
+				 u64 error_code,
+				 enum kvm_prefetch_type prefetch, u8 *level)
 {
 	int r;
 
@@ -5120,7 +5128,7 @@ static int kvm_tdp_page_prefault(struct kvm_vcpu *vcpu, gpa_t gpa,
 			return r;
 
 		cond_resched();
-		r = kvm_mmu_do_page_fault(vcpu, gpa, error_code, true, NULL, level);
+		r = kvm_mmu_do_page_fault(vcpu, gpa, error_code, prefetch, NULL, level);
 	} while (r == RET_PF_RETRY);
 
 	if (r < 0)
@@ -5175,7 +5183,8 @@ long kvm_arch_vcpu_pre_fault_memory(struct kvm_vcpu *vcpu,
 	 * Shadow paging uses GVA for kvm page fault, so restrict to
 	 * two-dimensional paging.
 	 */
-	r = kvm_tdp_page_prefault(vcpu, range->gpa | direct_bits, error_code, &level);
+	r = kvm_tdp_page_prefault(vcpu, range->gpa | direct_bits, error_code,
+				  KVM_PREFETCH_MANUAL, &level);
 	if (r < 0)
 		return r;
 
@@ -5207,7 +5216,7 @@ void kvm_arch_async_page_ready(struct kvm_vcpu *vcpu, struct kvm_async_pf *work)
 		return;
 
 	r = kvm_mmu_do_page_fault(vcpu, work->cr2_or_gpa, work->arch.error_code,
-				  true, NULL, NULL);
+				  KVM_PREFETCH_MANUAL, NULL, NULL);
 
 	/*
 	 * Account fixed page faults, otherwise they'll never be counted, but
@@ -6681,7 +6690,8 @@ int noinline kvm_mmu_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa, u64 err
 	if (r == RET_PF_INVALID) {
 		vcpu->stat.pf_taken++;
 
-		r = kvm_mmu_do_page_fault(vcpu, cr2_or_gpa, error_code, false,
+		r = kvm_mmu_do_page_fault(vcpu, cr2_or_gpa, error_code,
+					  KVM_PREFETCH_NONE,
 					  &emulation_type, NULL);
 		if (KVM_BUG_ON(r == RET_PF_INVALID, vcpu->kvm))
 			return -EIO;
diff --git a/arch/x86/kvm/mmu/mmu_internal.h b/arch/x86/kvm/mmu/mmu_internal.h
index c29002c60126..b02cdacd3e42 100644
--- a/arch/x86/kvm/mmu/mmu_internal.h
+++ b/arch/x86/kvm/mmu/mmu_internal.h
@@ -228,6 +228,7 @@ struct kvm_page_fault {
 	const gpa_t addr;
 	const u64 error_code;
 	const bool prefetch;
+	const bool auto_prefetch;
 
 	/* Derived from error_code.  */
 	const bool exec;
-- 
2.56.0.313.g2479322e68-goog

[PATCH 3/4] KVM: x86/mmu: Add automatic prefaulting/prefetching for TDP MMU faults

[Sean Christopherson]

Leverage the prefault code to add "automatic" prefetching of surrounding
SPTEs when handling TDP MMU page faults.  As pointed out by Marco, the TDP
MMU doesn't support prefetching SPTEs, which puts the TDP MMU at a severe
performance disadvantage relative to the shadow MMU when faulting in guest
memory at 4KiB granularity, and in general is a gap in the TDP MMU.

Use the prefault flow to prefetch SPTEs instead of the more minimal, custom
approach used by the shadow MMU.  Reusing the prefault flow sacrifices a
small amount of performance (the standard page fault path is "heavier"),
but is vastly simpler, supports backing memory that can't be GUP'd, and
shares its maintenance burden with KVM_PRE_FAULT_MEMORY.  E.g. multiple
flaws in x86's prefault code were practically impossible to hit with
userspace-driven prefaulting, but were hit almost immediately with
KVM-driven automatic prefetching.

Provide a writable module param to let userspace control the number of
pages to prefault, as KVM needs a control of some kind, and letting the
user choose how aggressively to prefetch is just as easy as hardcoding the
number of PTEs to prefetch.  I.e. supporting a variety of use cases is all
but free.  Default to PTE_PREFETCH_NUM as that's been KVM's behavior for
the shadow MMU for decades, i.e. even if it was arbitrary, it's gotten a
*lot* of mileage and so presumably is a reasonable default.

Instead of unconditionally bailing on hugepage faults, skip prefetching if
the original installed page size covers at least the number of pages the
user wants to prefetch.  In practice, it's unlikely hugepage prefetching
will see much, if any, use, but it's practically free to support, and will
allow for more experimentation and tuning.

Note!  Check for an incompatible memslot *before* invoking the specific
page fault handler, and thus before fast page faults are handled, otherwise
the fast page fault handler can "fix" a fault that never happened and log a
spurious, non-existent access.

Do NOT prefetch in response to protection faults, i.e. if there was already
a shadow-present SPTE, because prefetching in that scenario is unlikely to
be beneficial.  Note, this would also prevent spurious dirty logging via
fast page faults, but (obviously) only for writes, i.e. wouldn't block fast
handling for !PRESENT page faults when A/D bits are disabled.  I.e. the
aforementioned early check is still necessary to avoid spurious dirty log
entries (or KVM would need to explicitly treat auto-prefetch "faults" as
"can't be fast").

Lastly, don't go down the slow I/O path for auto-prefetch faults, as the
goal is for auto-prefetching to be light and fast, i.e. to pull in SPTEs
without significantly increasing the latency of handling page faults.

For now, use the auto-prefaulting code to the TDP MMU, i.e. defer
converting direct shadow MMUs to a separate change.

Reported-by: Marco Marangoni <mamarang@amazon.com>
Closes: https://lore.kernel.org/all/20260825140159.70997-1-mamarang@amazon.com
Signed-off-by: Sean Christopherson <seanjc@google.com>
---
 arch/x86/kvm/mmu/mmu.c | 64 +++++++++++++++++++++++++++++++++++++++---
 1 file changed, 60 insertions(+), 4 deletions(-)

diff --git a/arch/x86/kvm/mmu/mmu.c b/arch/x86/kvm/mmu/mmu.c
index 3bd0d921778d..c56a2eaa4646 100644
--- a/arch/x86/kvm/mmu/mmu.c
+++ b/arch/x86/kvm/mmu/mmu.c
@@ -122,8 +122,15 @@ static int max_huge_page_level __read_mostly;
 static int tdp_root_level __read_mostly;
 static int max_tdp_level __read_mostly;
 
+/*
+ * The exact number of PTEs that can be prefetched for the shadow MMU, and the
+ * default number of pages to prefault/prefetch pages for the TDP MMU.
+ */
 #define PTE_PREFETCH_NUM		8
 
+static unsigned int __read_mostly auto_prefault_nr_pages = PTE_PREFETCH_NUM;
+module_param(auto_prefault_nr_pages, uint, 0644);
+
 #include <trace/events/kvm.h>
 
 /* make pte_list_desc fit well in cache lines */
@@ -4709,10 +4716,12 @@ static int __kvm_mmu_faultin_pfn(struct kvm_vcpu *vcpu,
 	/*
 	 * If resolving the page failed because I/O is needed to fault-in the
 	 * page, then either set up an asynchronous #PF to do the I/O, or if
-	 * doing an async #PF isn't possible, retry with I/O allowed.  All
-	 * other failures are terminal, i.e. retrying won't help.
+	 * doing an async #PF isn't possible, retry with I/O allowed.  Never
+	 * go down the slow I/O path for automatic prefetching, as the goal is
+	 * for it to be light and fast.  All other failures are terminal, i.e.
+	 * retrying won't help.
 	 */
-	if (fault->pfn != KVM_PFN_ERR_NEEDS_IO)
+	if (fault->pfn != KVM_PFN_ERR_NEEDS_IO || fault->auto_prefetch)
 		return RET_PF_CONTINUE;
 
 	if (!fault->prefetch && kvm_can_do_async_pf(vcpu)) {
@@ -5072,6 +5081,18 @@ static int kvm_mmu_do_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa,
 		 */
 		fault.gfn = gpa_to_gfn(fault.addr) & ~kvm_gfn_direct_bits(vcpu->kvm);
 		fault.slot = kvm_vcpu_gfn_to_memslot(vcpu, fault.gfn);
+
+		if (fault.auto_prefetch) {
+			r = kvm_check_memslot_usable_for_auto_prefetch(fault.slot, ACC_ALL);
+			if (r)
+				return r;
+		}
+	} else {
+		/*
+		 * Indirect MMUs need to use a dedicated prefetch mechanism, as
+		 * the associated guest PTE factors into the prefetching.
+		 */
+		WARN_ON_ONCE(fault.auto_prefetch);
 	}
 
 	/*
@@ -6654,11 +6675,44 @@ static int kvm_mmu_write_protect_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa,
 	return RET_PF_EMULATE;
 }
 
+static void kvm_mmu_auto_prefault(struct kvm_vcpu *vcpu, gpa_t start,
+				  u64 error_code, u8 level)
+{
+	gfn_t nr_pages = READ_ONCE(auto_prefault_nr_pages);
+	int nr_pages_msb;
+	gfn_t i;
+
+	if (error_code & (PFERR_RSVD_MASK | PFERR_PRESENT_MASK))
+		return;
+
+	if (!tdp_mmu_enabled || vcpu->arch.mmu->page_fault != kvm_tdp_page_fault)
+		return;
+
+	nr_pages = min(nr_pages, KVM_PAGES_PER_HPAGE(PG_LEVEL_1G));
+	if (KVM_PAGES_PER_HPAGE(level) >= nr_pages)
+		return;
+
+	nr_pages_msb = find_last_bit((unsigned long *)&nr_pages, sizeof(nr_pages));
+
+	start = ALIGN_DOWN(start, gfn_to_gpa(BIT_ULL(nr_pages_msb)));
+
+	for (i = 0; i < nr_pages; i += KVM_PAGES_PER_HPAGE(level)) {
+		gpa_t gpa = start + gfn_to_gpa(i);
+
+		if (gpa < start || gpa_to_gfn(gpa) > kvm_mmu_max_gfn())
+			return;
+
+		if (kvm_tdp_page_prefault(vcpu, gpa, error_code, KVM_PREFETCH_AUTO, &level))
+			return;
+	}
+}
+
 int noinline kvm_mmu_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa, u64 error_code,
 		       void *insn, int insn_len)
 {
 	int r, emulation_type = EMULTYPE_PF;
 	bool direct = vcpu->arch.mmu->root_role.direct;
+	u8 level;
 
 	if (WARN_ON_ONCE(!VALID_PAGE(vcpu->arch.mmu->root.hpa)))
 		return RET_PF_RETRY;
@@ -6692,7 +6746,7 @@ int noinline kvm_mmu_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa, u64 err
 
 		r = kvm_mmu_do_page_fault(vcpu, cr2_or_gpa, error_code,
 					  KVM_PREFETCH_NONE,
-					  &emulation_type, NULL);
+					  &emulation_type, &level);
 		if (KVM_BUG_ON(r == RET_PF_INVALID, vcpu->kvm))
 			return -EIO;
 	}
@@ -6703,6 +6757,8 @@ int noinline kvm_mmu_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa, u64 err
 	if (r == RET_PF_WRITE_PROTECTED)
 		r = kvm_mmu_write_protect_fault(vcpu, cr2_or_gpa, error_code,
 						&emulation_type);
+	else if (r == RET_PF_FIXED)
+		kvm_mmu_auto_prefault(vcpu, cr2_or_gpa, error_code, level);
 
 	if (r == RET_PF_FIXED)
 		vcpu->stat.pf_fixed++;
-- 
2.56.0.313.g2479322e68-goog

[PATCH 4/4] KVM: x86/mmu: Use the generic auto-prefaulting to do prefetch for direct MMUs

[Sean Christopherson]

Use generic auto-prefaulting for direct shadow MMUs and drop the small bit
of custom prefetching that is unique to direct MMUs.  Direct shadow MMUs
are basically on life support, and the generic approach provides comparable
performance, plus it supported guest_memfd and PFN-backed memory.

Note!  Deliberately keep the core direct PTE prefetching logic, as it's
still needed to prefetch PTEs for direct shadow pages in indirect MMUs,
i.e. when the guest is using hugepages but KVM is not.

Opportunistically drop the double-underscores from direct_pte_prefetch()
now that only the version used by indirect MMUs exists.

Signed-off-by: Sean Christopherson <seanjc@google.com>
---
 arch/x86/kvm/mmu/mmu.c         | 48 ++++++----------------------------
 arch/x86/kvm/mmu/paging_tmpl.h |  2 +-
 2 files changed, 9 insertions(+), 41 deletions(-)

diff --git a/arch/x86/kvm/mmu/mmu.c b/arch/x86/kvm/mmu/mmu.c
index c56a2eaa4646..7f8beb1b307b 100644
--- a/arch/x86/kvm/mmu/mmu.c
+++ b/arch/x86/kvm/mmu/mmu.c
@@ -123,8 +123,8 @@ static int tdp_root_level __read_mostly;
 static int max_tdp_level __read_mostly;
 
 /*
- * The exact number of PTEs that can be prefetched for the shadow MMU, and the
- * default number of pages to prefault/prefetch pages for the TDP MMU.
+ * The exact number of PTEs that can be prefetched for indirect MMUs, and the
+ * default number of pages to prefault/prefetch pages for direct MMUs.
  */
 #define PTE_PREFETCH_NUM		8
 
@@ -3275,8 +3275,8 @@ static bool direct_pte_prefetch_many(struct kvm_vcpu *vcpu,
 	return kvm_mmu_prefetch_sptes(vcpu, gfn, start, end - start, access);
 }
 
-static void __direct_pte_prefetch(struct kvm_vcpu *vcpu,
-				  struct kvm_mmu_page *sp, u64 *sptep)
+static void direct_pte_prefetch(struct kvm_vcpu *vcpu, struct kvm_mmu_page *sp,
+				u64 *sptep)
 {
 	u64 *spte, *start = NULL;
 	int i;
@@ -3301,33 +3301,6 @@ static void __direct_pte_prefetch(struct kvm_vcpu *vcpu,
 		direct_pte_prefetch_many(vcpu, sp, start, spte);
 }
 
-static void direct_pte_prefetch(struct kvm_vcpu *vcpu, u64 *sptep)
-{
-	struct kvm_mmu_page *sp;
-
-	sp = sptep_to_sp(sptep);
-
-	/*
-	 * Without accessed bits, there's no way to distinguish between
-	 * actually accessed translations and prefetched, so disable pte
-	 * prefetch if accessed bits aren't available.
-	 */
-	if (sp_ad_disabled(sp))
-		return;
-
-	if (sp->role.level > PG_LEVEL_4K)
-		return;
-
-	/*
-	 * If addresses are being invalidated, skip prefetching to avoid
-	 * accidentally prefetching those addresses.
-	 */
-	if (unlikely(vcpu->kvm->mmu_invalidate_in_progress))
-		return;
-
-	__direct_pte_prefetch(vcpu, sp, sptep);
-}
-
 /*
  * Lookup the mapping level for @gfn in the current mm.
  *
@@ -3562,8 +3535,8 @@ static int direct_map(struct kvm_vcpu *vcpu, struct kvm_page_fault *fault)
 {
 	struct kvm_shadow_walk_iterator it;
 	struct kvm_mmu_page *sp;
-	int ret, access;
 	gfn_t base_gfn = fault->gfn;
+	int access;
 
 	kvm_mmu_hugepage_adjust(vcpu, fault);
 
@@ -3594,13 +3567,8 @@ static int direct_map(struct kvm_vcpu *vcpu, struct kvm_page_fault *fault)
 	if (KVM_BUG_ON(it.level != fault->goal_level, vcpu->kvm))
 		return -EIO;
 
-	ret = mmu_set_spte(vcpu, fault->slot, it.sptep, access,
-			   base_gfn, fault->pfn, fault);
-	if (ret == RET_PF_SPURIOUS)
-		return ret;
-
-	direct_pte_prefetch(vcpu, it.sptep);
-	return ret;
+	return mmu_set_spte(vcpu, fault->slot, it.sptep, access, base_gfn,
+			    fault->pfn, fault);
 }
 
 static void kvm_send_hwpoison_signal(struct kvm_memory_slot *slot, gfn_t gfn)
@@ -6685,7 +6653,7 @@ static void kvm_mmu_auto_prefault(struct kvm_vcpu *vcpu, gpa_t start,
 	if (error_code & (PFERR_RSVD_MASK | PFERR_PRESENT_MASK))
 		return;
 
-	if (!tdp_mmu_enabled || vcpu->arch.mmu->page_fault != kvm_tdp_page_fault)
+	if (vcpu->arch.mmu->page_fault != kvm_tdp_page_fault)
 		return;
 
 	nr_pages = min(nr_pages, KVM_PAGES_PER_HPAGE(PG_LEVEL_1G));
diff --git a/arch/x86/kvm/mmu/paging_tmpl.h b/arch/x86/kvm/mmu/paging_tmpl.h
index 8e350095508c..baaab791f107 100644
--- a/arch/x86/kvm/mmu/paging_tmpl.h
+++ b/arch/x86/kvm/mmu/paging_tmpl.h
@@ -631,7 +631,7 @@ static void FNAME(pte_prefetch)(struct kvm_vcpu *vcpu, struct guest_walker *gw,
 		return;
 
 	if (sp->role.direct)
-		return __direct_pte_prefetch(vcpu, sp, sptep);
+		return direct_pte_prefetch(vcpu, sp, sptep);
 
 	i = spte_index(sptep) & ~(PTE_PREFETCH_NUM - 1);
 	spte = sp->spt + i;
-- 
2.56.0.313.g2479322e68-goog

Re: [PATCH 3/4] KVM: x86/mmu: Add automatic prefaulting/prefetching for TDP MMU faults

[James Houghton]

On October 28, Sean Christopherson wrote:
> Leverage the prefault code to add "automatic" prefetching of surrounding
> SPTEs when handling TDP MMU page faults.  As pointed out by Marco, the TDP
> MMU doesn't support prefetching SPTEs, which puts the TDP MMU at a severe
> performance disadvantage relative to the shadow MMU when faulting in guest
> memory at 4KiB granularity, and in general is a gap in the TDP MMU.
> 
> Use the prefault flow to prefetch SPTEs instead of the more minimal, custom
> approach used by the shadow MMU.  Reusing the prefault flow sacrifices a
> small amount of performance (the standard page fault path is "heavier"),
> but is vastly simpler, supports backing memory that can't be GUP'd, and
> shares its maintenance burden with KVM_PRE_FAULT_MEMORY.  E.g. multiple
> flaws in x86's prefault code were practically impossible to hit with
> userspace-driven prefaulting, but were hit almost immediately with
> KVM-driven automatic prefetching.
> 
> Provide a writable module param to let userspace control the number of
> pages to prefault, as KVM needs a control of some kind, and letting the
> user choose how aggressively to prefetch is just as easy as hardcoding the
> number of PTEs to prefetch.  I.e. supporting a variety of use cases is all
> but free.  Default to PTE_PREFETCH_NUM as that's been KVM's behavior for
> the shadow MMU for decades, i.e. even if it was arbitrary, it's gotten a
> *lot* of mileage and so presumably is a reasonable default.
> 
> Instead of unconditionally bailing on hugepage faults, skip prefetching if
> the original installed page size covers at least the number of pages the
> user wants to prefetch.  In practice, it's unlikely hugepage prefetching
> will see much, if any, use, but it's practically free to support, and will
> allow for more experimentation and tuning.
> 
> Note!  Check for an incompatible memslot *before* invoking the specific
> page fault handler, and thus before fast page faults are handled, otherwise
> the fast page fault handler can "fix" a fault that never happened and log a
> spurious, non-existent access.
> 
> Do NOT prefetch in response to protection faults, i.e. if there was already
> a shadow-present SPTE, because prefetching in that scenario is unlikely to
> be beneficial.  Note, this would also prevent spurious dirty logging via
> fast page faults, but (obviously) only for writes, i.e. wouldn't block fast
> handling for !PRESENT page faults when A/D bits are disabled.  I.e. the
> aforementioned early check is still necessary to avoid spurious dirty log
> entries (or KVM would need to explicitly treat auto-prefetch "faults" as
> "can't be fast").
> 
> Lastly, don't go down the slow I/O path for auto-prefetch faults, as the
> goal is for auto-prefetching to be light and fast, i.e. to pull in SPTEs
> without significantly increasing the latency of handling page faults.
> 
> For now, use the auto-prefaulting code to the TDP MMU, i.e. defer
> converting direct shadow MMUs to a separate change.
> 
> Reported-by: Marco Marangoni <mamarang@amazon.com>
> Closes: https://lore.kernel.org/all/20260825140159.70997-1-mamarang@amazon.com
> Signed-off-by: Sean Christopherson <seanjc@google.com>
> ---
>  arch/x86/kvm/mmu/mmu.c | 64 +++++++++++++++++++++++++++++++++++++++---
>  1 file changed, 60 insertions(+), 4 deletions(-)
> 
> diff --git a/arch/x86/kvm/mmu/mmu.c b/arch/x86/kvm/mmu/mmu.c
> index 3bd0d921778d..c56a2eaa4646 100644
> --- a/arch/x86/kvm/mmu/mmu.c
> +++ b/arch/x86/kvm/mmu/mmu.c
> @@ -122,8 +122,15 @@ static int max_huge_page_level __read_mostly;
>  static int tdp_root_level __read_mostly;
>  static int max_tdp_level __read_mostly;
>  
> +/*
> + * The exact number of PTEs that can be prefetched for the shadow MMU, and the
> + * default number of pages to prefault/prefetch pages for the TDP MMU.
> + */
>  #define PTE_PREFETCH_NUM		8
>  
> +static unsigned int __read_mostly auto_prefault_nr_pages = PTE_PREFETCH_NUM;
> +module_param(auto_prefault_nr_pages, uint, 0644);
> +
>  #include <trace/events/kvm.h>
>  
>  /* make pte_list_desc fit well in cache lines */
> @@ -4709,10 +4716,12 @@ static int __kvm_mmu_faultin_pfn(struct kvm_vcpu *vcpu,
>  	/*
>  	 * If resolving the page failed because I/O is needed to fault-in the
>  	 * page, then either set up an asynchronous #PF to do the I/O, or if
> -	 * doing an async #PF isn't possible, retry with I/O allowed.  All
> -	 * other failures are terminal, i.e. retrying won't help.
> +	 * doing an async #PF isn't possible, retry with I/O allowed.  Never
> +	 * go down the slow I/O path for automatic prefetching, as the goal is
> +	 * for it to be light and fast.  All other failures are terminal, i.e.
> +	 * retrying won't help.
>  	 */
> -	if (fault->pfn != KVM_PFN_ERR_NEEDS_IO)
> +	if (fault->pfn != KVM_PFN_ERR_NEEDS_IO || fault->auto_prefetch)
>  		return RET_PF_CONTINUE;

Hi Sean,

I wasn't sure where to leave this feedback, but Marco's original patch[1]
used fast-only GUP instead of relying on normal GUP + FOLL_NOWAIT. I think
fast-only GUP is better here. (Correct me if I haven't understood this
series correctly.)

With normal GUP, we'll get allocate more pages than absolutely necessary,
and pages that are swapped out with zswap are swapped back in even for
FOLL_NOWAIT. And I'm a little bit concerned that normal GUP (even with
FOLL_NOWAIT) might stall page faults noticeably longer than before.

If you leave prefetching as normal GUP, maybe leave it off by default?

What do you think?

[1] https://lore.kernel.org/all/20260825140159.70997-1-mamarang@amazon.com/

>  
>  	if (!fault->prefetch && kvm_can_do_async_pf(vcpu)) {
> @@ -5072,6 +5081,18 @@ static int kvm_mmu_do_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa,
>  		 */
>  		fault.gfn = gpa_to_gfn(fault.addr) & ~kvm_gfn_direct_bits(vcpu->kvm);
>  		fault.slot = kvm_vcpu_gfn_to_memslot(vcpu, fault.gfn);
> +
> +		if (fault.auto_prefetch) {
> +			r = kvm_check_memslot_usable_for_auto_prefetch(fault.slot, ACC_ALL);
> +			if (r)
> +				return r;
> +		}
> +	} else {
> +		/*
> +		 * Indirect MMUs need to use a dedicated prefetch mechanism, as
> +		 * the associated guest PTE factors into the prefetching.
> +		 */
> +		WARN_ON_ONCE(fault.auto_prefetch);
>  	}
>  
>  	/*
> @@ -6654,11 +6675,44 @@ static int kvm_mmu_write_protect_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa,
>  	return RET_PF_EMULATE;
>  }
>  
> +static void kvm_mmu_auto_prefault(struct kvm_vcpu *vcpu, gpa_t start,
> +				  u64 error_code, u8 level)
> +{
> +	gfn_t nr_pages = READ_ONCE(auto_prefault_nr_pages);
> +	int nr_pages_msb;
> +	gfn_t i;
> +
> +	if (error_code & (PFERR_RSVD_MASK | PFERR_PRESENT_MASK))
> +		return;
> +
> +	if (!tdp_mmu_enabled || vcpu->arch.mmu->page_fault != kvm_tdp_page_fault)
> +		return;
> +
> +	nr_pages = min(nr_pages, KVM_PAGES_PER_HPAGE(PG_LEVEL_1G));
> +	if (KVM_PAGES_PER_HPAGE(level) >= nr_pages)
> +		return;
> +
> +	nr_pages_msb = find_last_bit((unsigned long *)&nr_pages, sizeof(nr_pages));
> +
> +	start = ALIGN_DOWN(start, gfn_to_gpa(BIT_ULL(nr_pages_msb)));
> +
> +	for (i = 0; i < nr_pages; i += KVM_PAGES_PER_HPAGE(level)) {
> +		gpa_t gpa = start + gfn_to_gpa(i);
> +
> +		if (gpa < start || gpa_to_gfn(gpa) > kvm_mmu_max_gfn())
> +			return;
> +
> +		if (kvm_tdp_page_prefault(vcpu, gpa, error_code, KVM_PREFETCH_AUTO, &level))
> +			return;
> +	}
> +}
> +
>  int noinline kvm_mmu_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa, u64 error_code,
>  		       void *insn, int insn_len)
>  {
>  	int r, emulation_type = EMULTYPE_PF;
>  	bool direct = vcpu->arch.mmu->root_role.direct;
> +	u8 level;
>  
>  	if (WARN_ON_ONCE(!VALID_PAGE(vcpu->arch.mmu->root.hpa)))
>  		return RET_PF_RETRY;
> @@ -6692,7 +6746,7 @@ int noinline kvm_mmu_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa, u64 err
>  
>  		r = kvm_mmu_do_page_fault(vcpu, cr2_or_gpa, error_code,
>  					  KVM_PREFETCH_NONE,
> -					  &emulation_type, NULL);
> +					  &emulation_type, &level);
>  		if (KVM_BUG_ON(r == RET_PF_INVALID, vcpu->kvm))
>  			return -EIO;
>  	}
> @@ -6703,6 +6757,8 @@ int noinline kvm_mmu_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa, u64 err
>  	if (r == RET_PF_WRITE_PROTECTED)
>  		r = kvm_mmu_write_protect_fault(vcpu, cr2_or_gpa, error_code,
>  						&emulation_type);
> +	else if (r == RET_PF_FIXED)
> +		kvm_mmu_auto_prefault(vcpu, cr2_or_gpa, error_code, level);
>  
>  	if (r == RET_PF_FIXED)
>  		vcpu->stat.pf_fixed++;
> -- 
> 2.56.0.313.g2479322e68-goog

Re: [PATCH 3/4] KVM: x86/mmu: Add automatic prefaulting/prefetching for TDP MMU faults

[Marangoni, Marco]

Hi Sean, thanks for the patch!
I run some benchmarks on it, and it definitely helps! This is the time to fault 128 MiB linearly after a snapshot restore, with an UFFD page handler that faults chunks on demand; medians of 60 restores. Baseline is kvm-x86/next. Each entry is nr_pages: change vs baseline on the patched kernel; negative means faster.

- UFFD 4 KiB handler (baseline 2228 ms) 0: +1%, 1: −5%, 2: −3%, 4: +1%, 8: −3%, 16: +1%, 32: 0%, 64: 0%, 128: −2%, 256: +1%, 512: +1%
- UFFD 512 KiB handler (baseline 543 ms) 0: 0%, 1: 0%, 2: −36%, 4: −51%, 8: −59%, 16: −64%, 32: −71%, 64: −67%, 128: −68%, 256: −67%, 512: −66%
- UFFD 2 MiB handler (baseline 513 ms) 0: 0%, 1: +1%, 2: −34%, 4: −50%, 8: −59%, 16: −64%, 32: −68%, 64: −68%, 128: −68%, 256: −68%, 512: −68%

8 is the series' default; 0 and 1 both disable prefetch and reproduce baseline. The 4 KiB row is within ±5%, i.e. noise.
It's measured on nested virtualization, so the improvement is probably less dramatic on metal. For context, my initial patch was around ~75% faster, so I'd say that performance-wise, your implementation is solid!

It seems like a prefetch of 32 is the sweet spot, but I agree that with James that with the normal GUP it's probably better to keep it off or very low by default.
From my perspective, using fast GUP has the problem that you cannot prefault "around" your page, but just ahead of it, so I don't mind going with normal GUP if it makes things easier for you.

I'd say that most setup would know the size at which UFFD copies the memory, so having to explicitly configure is okay.

For completeness, I also measured how much the patch improves at various UFFD chunks, with nr_pages at the same size:
  - 4 KiB: 2285 ms → 2321 ms (np1), +2% (np1 is "off", noise)
  - 8 KiB: 1435 ms → 1441 ms (np2), 0%
  - 16 KiB: 1036 ms → 860 ms (np4), −17%
  - 32 KiB: 743 ms → 469 ms (np8), −37%
  - 64 KiB: 645 ms → 319 ms (np16), −50%
  - 128 KiB: 587 ms → 242 ms (np32), −59%
  - 256 KiB: 561 ms → 200 ms (np64), −64%
  - 512 KiB: 543 ms → 165 ms (np128), −70%
  - 1 MiB: 534 ms → 171 ms (np256), −68%
  - 2 MiB: 523 ms → 151 ms (np512), −71%

These measurement highlight a possible limitation (why didn't 8KiB improve?): it seems like the prefetching doesn't kick in the first time a page is getting faulted in the chunk, is that expected? I asked an LLM to investigate, and it seems to be because with async PF, the first fault doesn't reach the kvm_mmu_auto_prefault code path.

---

And another question: this is just for x86. Do you think the same approach could be extended for ARM, or is that a different beast? Do you need some help to make that a reality as well?

Thanks again,
Marco

Re: [PATCH 3/4] KVM: x86/mmu: Add automatic prefaulting/prefetching for TDP MMU faults

[Sean Christopherson]

On Thu, Oct 29, 2026, Marco Marangoni wrote:
> Hi Sean, thanks for the patch!
> I run some benchmarks on it, and it definitely helps! This is the time to fault 128 MiB linearly after a snapshot restore, with an UFFD page handler that faults chunks on demand; medians of 60 restores. Baseline is kvm-x86/next. Each entry is nr_pages: change vs baseline on the patched kernel; negative means faster.
> 
> - UFFD 4 KiB handler (baseline 2228 ms) 0: +1%, 1: −5%, 2: −3%, 4: +1%, 8: −3%, 16: +1%, 32: 0%, 64: 0%, 128: −2%, 256: +1%, 512: +1%
> - UFFD 512 KiB handler (baseline 543 ms) 0: 0%, 1: 0%, 2: −36%, 4: −51%, 8: −59%, 16: −64%, 32: −71%, 64: −67%, 128: −68%, 256: −67%, 512: −66%
> - UFFD 2 MiB handler (baseline 513 ms) 0: 0%, 1: +1%, 2: −34%, 4: −50%, 8: −59%, 16: −64%, 32: −68%, 64: −68%, 128: −68%, 256: −68%, 512: −68%
> 
> 8 is the series' default; 0 and 1 both disable prefetch and reproduce
> baseline. The 4 KiB row is within ±5%, i.e. noise.  It's measured on nested
> virtualization, so the improvement is probably less dramatic on metal. For
> context, my initial patch was around ~75% faster, so I'd say that
> performance-wise, your implementation is solid!
> 
> It seems like a prefetch of 32 is the sweet spot, but I agree that with James
> that with the normal GUP it's probably better to keep it off or very low by
> default.

Yeah, agreed.  Triggering zswap or any other heavy operations could regress a
lot of configurations/workloads.

Hrm, but I don't want to lose support for guest_memfd without an hva.  Oh.  That's
a pre-existing bug, kvm_prefetch_pages() should reject kvm_memslot_is_gmem_only()
slots.  But that's a side topic.

Argh, FOLL_FAST_ONLY is internal-only.  But do we actually need it to be super
strictly fast-only?  Wouldn't FOLL_NOFAULT suffice?  IRQs are enabled so taking
mmap_sem is fine; it's not be ideal, but KVM will still try fast-GUP first, and
in practice performance is going to suffer no matter what if mmap_sem is
contended, e.g. arm64 *unconditionally* takes mmap_sem for read when handling
faults, and while it shows up in micro-benchmarks, e.g. mmu_stress_test, AFAIK
no one has complained otherwise.

So isn't the fix just this?

diff --git a/arch/x86/kvm/mmu/mmu.c b/arch/x86/kvm/mmu/mmu.c
index b673c8130edd..1fcf4fedc0bb 100644
--- a/arch/x86/kvm/mmu/mmu.c
+++ b/arch/x86/kvm/mmu/mmu.c
@@ -4684,6 +4684,9 @@ static int __kvm_mmu_faultin_pfn(struct kvm_vcpu *vcpu,
 {
 	unsigned int foll = fault->write ? FOLL_WRITE : 0;
 
+	if (fault->auto_prefetch)
+		foll |= FOLL_NOFAULT;
+
 	if (fault->is_private || kvm_memslot_is_gmem_only(fault->slot))
 		return kvm_mmu_faultin_pfn_gmem(vcpu, fault);

> From my perspective, using fast GUP has the problem that you cannot prefault
> "around" your page, but just ahead of it, so I don't mind going with normal
> GUP if it makes things easier for you.

I think I'm missing something, why would fast GUP prevent prefaulting "behind"?
E.g. is that effectively a limitation of the underlying backing store you're
using, a fundamental limitation somewhere in mm/, or a limitation in KVM itself?

> I'd say that most setup would know the size at which UFFD copies the memory,
> so having to explicitly configure is okay.
> 
> For completeness, I also measured how much the patch improves at various UFFD chunks, with nr_pages at the same size:
>   - 4 KiB: 2285 ms → 2321 ms (np1), +2% (np1 is "off", noise)
>   - 8 KiB: 1435 ms → 1441 ms (np2), 0%
>   - 16 KiB: 1036 ms → 860 ms (np4), −17%
>   - 32 KiB: 743 ms → 469 ms (np8), −37%
>   - 64 KiB: 645 ms → 319 ms (np16), −50%
>   - 128 KiB: 587 ms → 242 ms (np32), −59%
>   - 256 KiB: 561 ms → 200 ms (np64), −64%
>   - 512 KiB: 543 ms → 165 ms (np128), −70%
>   - 1 MiB: 534 ms → 171 ms (np256), −68%
>   - 2 MiB: 523 ms → 151 ms (np512), −71%
> 
> These measurement highlight a possible limitation (why didn't 8KiB improve?):
> it seems like the prefetching doesn't kick in the first time a page is
> getting faulted in the chunk, is that expected? I asked an LLM to
> investigate, and it seems to be because with async PF, the first fault
> doesn't reach the kvm_mmu_auto_prefault code path.

Hmm, right.  I deliberately hooked kvm_mmu_page_fault() instead of
kvm_mmu_do_page_fault() so that prefetching wouldn't trigger auto-prefetch, but
the existing shadow MMU prefetching does kick in for async #PFs, and it does make
sense to prefetch in that case if we buy my argument (from this changelog) that
async #PF should punch through dirty logging.

Auto-prefetching for KVM_PRE_FAULT_MEMORY feels "wrong", but we do prefetch in
that case for the shadow MMU, though I'm like 99.9% certain no one is running
TDP with the shadow MMU these days, so I don't think we should consider that a
valid datapoint.

Untested, but I think it'd be this:

diff --git a/arch/x86/kvm/mmu/mmu.c b/arch/x86/kvm/mmu/mmu.c
index b673c8130edd..c0ff50c52ad5 100644
--- a/arch/x86/kvm/mmu/mmu.c
+++ b/arch/x86/kvm/mmu/mmu.c
@@ -5024,6 +5024,9 @@ enum kvm_prefetch_type {
 	KVM_PREFETCH_AUTO,
 };
 
+static void kvm_mmu_auto_prefault(struct kvm_vcpu *vcpu, gpa_t start,
+				  u64 error_code, u8 level);
+
 static int kvm_mmu_do_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa,
 				 u64 err, enum kvm_prefetch_type prefetch,
 				 int *emulation_type, u8 *level)
@@ -5098,6 +5101,9 @@ static int kvm_mmu_do_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa,
 	if (level)
 		*level = fault.goal_level;
 
+	if (r == RET_PF_FIXED && prefetch != KVM_PREFETCH_AUTO)
+		kvm_mmu_auto_prefault(vcpu, cr2_or_gpa, err, fault.goal_level);
+
 	return r;
 }

or to disable it for KVM_PRE_FAULT_MEMORY:

diff --git a/arch/x86/kvm/mmu/mmu.c b/arch/x86/kvm/mmu/mmu.c
index b673c8130edd..ee31a35f6a55 100644
--- a/arch/x86/kvm/mmu/mmu.c
+++ b/arch/x86/kvm/mmu/mmu.c
@@ -5021,9 +5021,13 @@ static int kvm_tdp_page_fault(struct kvm_vcpu *vcpu, struct kvm_page_fault *faul
 enum kvm_prefetch_type {
 	KVM_PREFETCH_NONE,
 	KVM_PREFETCH_MANUAL,
+	KVM_PREFETCH_ASYNC,
 	KVM_PREFETCH_AUTO,
 };
 
+static void kvm_mmu_auto_prefault(struct kvm_vcpu *vcpu, gpa_t start,
+				  u64 error_code, u8 level);
+
 static int kvm_mmu_do_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa,
 				 u64 err, enum kvm_prefetch_type prefetch,
 				 int *emulation_type, u8 *level)
@@ -5098,6 +5102,10 @@ static int kvm_mmu_do_page_fault(struct kvm_vcpu *vcpu, gpa_t cr2_or_gpa,
 	if (level)
 		*level = fault.goal_level;
 
+	if (r == RET_PF_FIXED && prefetch != KVM_PREFETCH_AUTO &&
+	    prefetch != KVM_PREFETCH_MANUAL)
+		kvm_mmu_auto_prefault(vcpu, cr2_or_gpa, err, fault.goal_level);
+
 	return r;
 }
 
@@ -5215,7 +5223,7 @@ void kvm_arch_async_page_ready(struct kvm_vcpu *vcpu, struct kvm_async_pf *work)
 		return;
 
 	r = kvm_mmu_do_page_fault(vcpu, work->cr2_or_gpa, work->arch.error_code,
-				  KVM_PREFETCH_MANUAL, NULL, NULL);
+				  KVM_PREFETCH_ASYNC, NULL, NULL);
 
 	/*
 	 * Account fixed page faults, otherwise they'll never be counted, but