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

[Sean Christopherson <seanjc@google.com>]

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