[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