[RFC PATCH 0/2] KVM: arm64: nv: Skip empty spans when unmapping a shadow S2

[RFC PATCH 0/2] KVM: arm64: nv: Skip empty spans when unmapping a shadow S2

[Haseeb Ashraf]

Emulating TLBI VMALLS12E1IS, emulating TLBI ALLE1IS and recycling a
shadow stage-2 all unmap the whole nested IPA space in chunks of the
largest block size, even when the table is empty. With a 48-bit PARange
and 4kB pages that is 262144 unmap calls per shadow MMU, each ending in a
range TLBI when FEAT_TLBIRANGE and FWB are available.

Measured cost: one full-range unmap of a shadow MMU takes 0.87-1.06 s on
AmpereOne (Xen on KVM, 4kB pages, function_graph). On Graviton4
(Xen-on-KVM, 48-bit IPA) a single VMALLS12E1IS trap took 1.46-9.50 s over
twelve samples (median 3.09 s), all of it in stage2_apply_range().

Patch 1 adds kvm_stage2_unmap_all(), which skips spans covered by a zero
descriptor at least as large as an unmap chunk, keeps the tracked unmap
for everything else, removes the guest stage-2 records of skipped spans,
and ends with the VMID-wide TLBI that VMALLS12E1IS requires. It converts
the VMALLS12E1IS handler. Patch 2 converts kvm_nested_s2_unmap() (ALLE1IS,
whole-VM unmap) and the pending unmap of a recycled shadow MMU.

The series is based on master, which already has the guest stage-2
tracking. It overlaps with the unmerged kvm-arm64/unmap-vmall work, which
defines a kvm_stage2_unmap_all() with the same name and signature; if that
lands first, patch 1 becomes "skip empty spans in kvm_stage2_unmap_all()".
As far as I can tell, that WIP helper does not touch the guest stage-2
tracking records, which this series does.

All twelve Graviton4 samples, in ftrace order:

  22) $ 2278347 us  |  } /* handle_vmalls12e1is */
  21) $ 3229689 us  |  } /* handle_vmalls12e1is */
  13) $ 7244813 us  |  } /* handle_vmalls12e1is */
   1) $ 2950198 us  |  } /* handle_vmalls12e1is */
   6) $ 5785126 us  |  } /* handle_vmalls12e1is */
  25) $ 2484152 us  |  } /* handle_vmalls12e1is */
   4) $ 2505404 us  |  } /* handle_vmalls12e1is */
  16) $ 1464441 us  |  } /* handle_vmalls12e1is */
  28) $ 1935604 us  |  } /* handle_vmalls12e1is */
  19) $ 9495486 us  |  } /* handle_vmalls12e1is */
  14) $ 4151658 us  |  } /* handle_vmalls12e1is */
  18) $ 9223579 us  |  } /* handle_vmalls12e1is */

Tooling: the patches were drafted with AI assistance (SiemensGPT).

Testing: build-tested on the master branch and
boot-tested with Linux 6.18.29.

[RFC PATCH 1/2] KVM: arm64: nv: Skip empty spans when unmapping a shadow S2

[Haseeb Ashraf]

From: Haseeb Ashraf <haseeb.ashraf@siemens.com>

Emulating TLBI VMALLS12E1IS unmaps the whole shadow stage-2 of the
guest hypervisor's current VMID by passing the complete nested IPA
range to kvm_stage2_unmap_range(). That walks the range in chunks of
the largest block size and calls kvm_pgtable_stage2_unmap() on each of
them, mapped or not: 262144 chunks per shadow MMU with 4kB pages and a
48-bit PARange, each ending in a range TLBI when FEAT_TLBIRANGE and
FWB are available.

Guest hypervisors such as Xen issue this TLBI while creating domains,
and an empty table is charged for the whole IPA space. A full-range
unmap of one shadow MMU takes 0.87-1.06 s on AmpereOne (Xen on KVM,
4kB pages, function_graph). On Graviton4 (Xen-on-KVM, 48-bit IPA),
twelve VMALLS12E1IS traps took 1.46-9.50 s each (median 3.09 s), all
of it in stage2_apply_range():

  22)   4.180 us    |              } /* kvm_pgtable_stage2_unmap */
  22) $ 2278329 us  |            } /* stage2_apply_range */
  22) $ 2278329 us  |          } /* __unmap_stage2_range */
  22) $ 2278330 us  |        } /* kvm_stage2_unmap_range */
  22) $ 2278331 us  |      } /* s2_mmu_unmap_range */
  22) $ 2278337 us  |    } /* kvm_s2_mmu_iterate_by_vmid */
  22) $ 2278347 us  |  } /* handle_vmalls12e1is */

Add kvm_stage2_unmap_all(). It uses kvm_pgtable_get_leaf() to find the
entry covering the current address and, if that is a zero descriptor
at least as large as an unmap chunk, skips to the end of the span
without entering the unmap walker. Everything else is unmapped chunk
by chunk as before, so populated chunks keep their unmap, cache
maintenance and table reclamation ordering. Non-zero invalid
descriptors are never skipped, as they can carry software state.

Populated chunks use the tracked unmap, and the guest stage-2 records
of a skipped span are removed under the same mmu_lock hold: nothing is
mapped there, so they are stale. Dropping them after the walk would be
wrong if mmu_lock was released, as a concurrent fault could have
mapped something behind it.

The helper ends with a VMID-wide TLBI: VMALLS12E1IS also invalidates
stage-1 translations, which the per-chunk invalidations used to do
even for empty tables. mmu_lock may be dropped between spans if the
caller allows blocking, as for the range unmap.

Only the VMALLS12E1IS handler is converted here.

Signed-off-by: Haseeb Ashraf <haseeb.ashraf@siemens.com>
---
 arch/arm64/include/asm/kvm_mmu.h |  1 +
 arch/arm64/kvm/mmu.c             | 80 ++++++++++++++++++++++++++++++++
 arch/arm64/kvm/sys_regs.c        | 20 ++++----
 3 files changed, 91 insertions(+), 10 deletions(-)

diff --git a/arch/arm64/include/asm/kvm_mmu.h b/arch/arm64/include/asm/kvm_mmu.h
index 32392abf2e21..45655a669bd7 100644
--- a/arch/arm64/include/asm/kvm_mmu.h
+++ b/arch/arm64/include/asm/kvm_mmu.h
@@ -184,6 +184,7 @@ void __init free_hyp_pgds(void);
 
 void kvm_stage2_unmap_range(struct kvm_s2_mmu *mmu, phys_addr_t start,
 			    u64 size, bool may_block);
+void kvm_stage2_unmap_all(struct kvm_s2_mmu *mmu, bool may_block);
 void kvm_stage2_flush_range(struct kvm_s2_mmu *mmu, phys_addr_t addr, phys_addr_t end);
 void kvm_stage2_wp_range(struct kvm_s2_mmu *mmu, phys_addr_t addr, phys_addr_t end);
 
diff --git a/arch/arm64/kvm/mmu.c b/arch/arm64/kvm/mmu.c
index 1db10d8d6db7..af45c4226b31 100644
--- a/arch/arm64/kvm/mmu.c
+++ b/arch/arm64/kvm/mmu.c
@@ -409,6 +409,86 @@ void kvm_stage2_unmap_range(struct kvm_s2_mmu *mmu, phys_addr_t start,
 	kvm->arch.vm_s2_ops->vm_stage2_unmap_range(mmu, start, size, may_block);
 }
 
+/**
+ * kvm_stage2_unmap_all -- Unmap a whole shadow stage-2
+ * @mmu:       The shadow stage-2 MMU to unmap
+ * @may_block: Whether or not we are permitted to block
+ *
+ * Remove every mapping from a shadow stage-2, as unmapping its whole IPA
+ * space with kvm_stage2_unmap_range() would, but without entering the
+ * unmap walker for the (usually vast) parts of the IPA space that are
+ * already empty. A span covered by a zero descriptor that is at least as
+ * large as an unmap chunk is skipped; everything else is unmapped chunk
+ * by chunk, exactly as before.
+ *
+ * Like kvm_stage2_unmap_range(), this must be called with mmu_lock held
+ * for write, and may drop it between spans if @may_block is set.
+ */
+void kvm_stage2_unmap_all(struct kvm_s2_mmu *mmu, bool may_block)
+{
+	struct kvm *kvm = kvm_s2_mmu_to_kvm(mmu);
+	phys_addr_t addr = 0, end = kvm_phys_size(mmu);
+	bool lock_dropped = false;
+
+	lockdep_assert_held_write(&kvm->mmu_lock);
+
+	if (WARN_ON_ONCE(!kvm_is_nested_s2_mmu(kvm, mmu)))
+		return;
+
+	do {
+		struct kvm_pgtable *pgt = mmu->pgt;
+		phys_addr_t next;
+		kvm_pte_t pte;
+		s8 level;
+
+		/*
+		 * We may be raced on PGT teardown when we drop mmu_lock.
+		 * That's fine as the PGT is legitimately no longer present.
+		 */
+		if (!pgt) {
+			WARN_ON_ONCE(!lock_dropped);
+			return;
+		}
+
+		/*
+		 * A zero descriptor at, or above, the unmap chunk level
+		 * covers an empty span. Do not skip non-zero invalid
+		 * descriptors, as the unmap walker may have to clear the
+		 * software state they carry. Finer-grained leaves keep the
+		 * chunk-sized unmap, so that a populated table is not
+		 * processed one page at a time.
+		 */
+		if (!kvm_pgtable_get_leaf(pgt, addr, &pte, &level, 0) &&
+		    !pte && level <= KVM_PGTABLE_MIN_BLOCK_LEVEL) {
+			next = __stage2_range_addr_end(addr, end,
+						       kvm_granule_size(level));
+			/* Nothing is mapped here, so any record is stale. */
+			kvm_remove_guest_s2_mappings(mmu, addr, next - addr);
+		} else {
+			next = stage2_range_addr_end(addr, end);
+			if (WARN_ON_ONCE(kvm_pgtable_stage2_unmap_tracked(pgt, addr,
+									  next - addr)))
+				break;
+		}
+
+		if (may_block && next != end) {
+			cond_resched_rwlock_write(&kvm->mmu_lock);
+			lock_dropped = true;
+		}
+		addr = next;
+	} while (addr != end);
+
+	/*
+	 * Skipping empty spans must not skip the invalidation that the
+	 * guest operation implies: a TLBI VMALLS12E1IS also invalidates
+	 * stage-1 translations, and a context that does not use the shadow
+	 * page tables can have TLB entries tagged with this MMU's VMID
+	 * even though they are empty. The per-chunk range invalidations
+	 * used to provide this as a side effect.
+	 */
+	kvm_call_hyp(__kvm_tlb_flush_vmid, mmu);
+}
+
 void kvm_stage2_flush_range(struct kvm_s2_mmu *mmu, phys_addr_t addr, phys_addr_t end)
 {
 	stage2_apply_range_resched(mmu, addr, end, KVM_PGT_FN(kvm_pgtable_stage2_flush));
diff --git a/arch/arm64/kvm/sys_regs.c b/arch/arm64/kvm/sys_regs.c
index 850eba167aad..58c221e7c710 100644
--- a/arch/arm64/kvm/sys_regs.c
+++ b/arch/arm64/kvm/sys_regs.c
@@ -4124,26 +4124,26 @@ static void s2_mmu_unmap_range(struct kvm_s2_mmu *mmu,
 	kvm_stage2_unmap_range(mmu, info->range.start, info->range.size, true);
 }
 
+static void s2_mmu_unmap_all(struct kvm_s2_mmu *mmu,
+			     const union tlbi_info *info)
+{
+	/* See s2_mmu_unmap_range() for the rules when dropping mmu_lock. */
+	kvm_stage2_unmap_all(mmu, true);
+}
+
 static bool handle_vmalls12e1is(struct kvm_vcpu *vcpu, struct sys_reg_params *p,
 				const struct sys_reg_desc *r)
 {
 	u32 sys_encoding = sys_insn(p->Op0, p->Op1, p->CRn, p->CRm, p->Op2);
-	u64 limit, vttbr;
+	u64 vttbr;
 
 	if (!kvm_supported_tlbi_s12_op(vcpu, sys_encoding))
 		return undef_access(vcpu, p, r);
 
 	vttbr = vcpu_read_sys_reg(vcpu, VTTBR_EL2);
-	limit = BIT_ULL(kvm_get_pa_bits(vcpu->kvm));
 
-	kvm_s2_mmu_iterate_by_vmid(vcpu->kvm, get_vmid(vttbr),
-				   &(union tlbi_info) {
-					   .range = {
-						   .start = 0,
-						   .size = limit,
-					   },
-				   },
-				   s2_mmu_unmap_range);
+	kvm_s2_mmu_iterate_by_vmid(vcpu->kvm, get_vmid(vttbr), NULL,
+				   s2_mmu_unmap_all);
 
 	return true;
 }
-- 
2.43.0

[RFC PATCH 2/2] KVM: arm64: nv: Use kvm_stage2_unmap_all() for shadow S2 teardown

[Haseeb Ashraf]

From: Haseeb Ashraf <haseeb.ashraf@siemens.com>

kvm_nested_s2_unmap() and the pending unmap that is completed at vcpu
entry both drop a complete shadow stage-2 by unmapping its whole IPA
range chunk by chunk, with the same cost for an empty table as the
TLBI VMALLS12E1IS handler: 262144 unmap calls per MMU with 4kB pages
and a 48-bit PARange.

kvm_nested_s2_unmap() is used when emulating TLBI ALLE1IS and when the
whole VM is unmapped. The pending unmap runs on the first entry after
a shadow MMU has been recycled for a different context, which a guest
hypervisor running more contexts than there are shadow MMUs can
trigger repeatedly.

Switch both to kvm_stage2_unmap_all(). The may_block rules are
unchanged: kvm_nested_s2_unmap() forwards its argument, and the
pending unmap, which only holds mmu_lock, still allows blocking. The
VNCR invalidation and the synchronous completion of the pending unmap
are left as they are.

Unmaps of a canonical IPA range, such as MMU notifier invalidations
and memslot removal, no longer go through kvm_nested_s2_unmap() since
the guest stage-2 tracking was introduced, and are not affected.

Signed-off-by: Haseeb Ashraf <haseeb.ashraf@siemens.com>
---
 arch/arm64/kvm/nested.c | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/arch/arm64/kvm/nested.c b/arch/arm64/kvm/nested.c
index 4d83281a7a28..875eb9c312f9 100644
--- a/arch/arm64/kvm/nested.c
+++ b/arch/arm64/kvm/nested.c
@@ -1441,7 +1441,7 @@ void kvm_nested_s2_unmap(struct kvm *kvm, bool may_block)
 		struct kvm_s2_mmu *mmu = kvm->arch.nested_mmus[i];
 
 		if (kvm_s2_mmu_valid(mmu))
-			kvm_stage2_unmap_range(mmu, 0, kvm_phys_size(mmu), may_block);
+			kvm_stage2_unmap_all(mmu, may_block);
 	}
 
 	kvm_invalidate_vncr_ipa_all(kvm);
@@ -2160,7 +2160,7 @@ void check_nested_vcpu_requests(struct kvm_vcpu *vcpu)
 
 		write_lock(&vcpu->kvm->mmu_lock);
 		if (mmu->pending_unmap) {
-			kvm_stage2_unmap_range(mmu, 0, kvm_phys_size(mmu), true);
+			kvm_stage2_unmap_all(mmu, true);
 			mmu->pending_unmap = false;
 		}
 		write_unlock(&vcpu->kvm->mmu_lock);
-- 
2.43.0