[PATCH v2] KVM: arm64: Resume stage-2 unmap where the walk stopped

[Haseeb Ashraf <haseebashraf091@gmail.com>]

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

kvm_vm_stage2_unmap_range() unmaps a range one block-sized chunk at a
time. The walker steps over an invalid entry in one go, but the next
walk starts at the following chunk, so an invalid entry larger than a
chunk is visited once for every chunk it covers. Unmapping the whole
IPA space of a shadow stage-2, as emulating TLBI VMALLS12E1IS and TLBI
ALLE1IS does, takes 2^18 walks with 4K pages and a 48-bit IPA space,
however little it maps. Each walk also ends with a range TLBI when
FEAT_TLBIRANGE and FWB are available.

On an AmpereOne A192-32X (192 CPUs, Xen on KVM, 4K pages, 48-bit IPA),
a full-range unmap of one shadow stage-2 took 0.82-0.94 s (mean 0.91 s
over 192 unmaps) on 6.18.29. With this change and the mainline commits
it builds on backported, it took 2.1-27.8 ms (mean 3.8 ms over 128
unmaps), with 512 to 1024 walks instead of 262144. A TLBI VMALLS12E1IS
trap that unmapped two shadow stage-2s took up to 1.85 s before.

Add kvm_pgtable_stage2_unmap_range(), which returns the address its
walk reached, and have the unmap loop resume from there.
kvm_pgtable_stage2_unmap() still returns 0 or an error, for its
callers at EL2. The walk handles every entry it steps over, including
invalid entries that carry state, so nothing has to be looked up
beforehand.

The unmap of a shadow stage-2 also drops the guest stage-2 records of
the whole span the walk covered: nothing is mapped in the part it
stepped over, so any record there is stale. Each walk that remains
ends with the same deferred TLB invalidation as before. The spans that
were skipped have no valid entry, so there is nothing to invalidate.

Suggested-by: Oliver Upton <oupton@kernel.org>
Signed-off-by: Haseeb Ashraf <haseeb.ashraf@siemens.com>
---
 arch/arm64/include/asm/kvm_pgtable.h | 18 ++++++++
 arch/arm64/kvm/hyp/pgtable.c         | 22 ++++++++--
 arch/arm64/kvm/mmu.c                 | 63 +++++++++++++++++++++++-----
 3 files changed, 90 insertions(+), 13 deletions(-)

diff --git a/arch/arm64/include/asm/kvm_pgtable.h b/arch/arm64/include/asm/kvm_pgtable.h
index 1c9c57c..4ed096f 100644
--- a/arch/arm64/include/asm/kvm_pgtable.h
+++ b/arch/arm64/include/asm/kvm_pgtable.h
@@ -725,6 +725,24 @@ int kvm_pgtable_stage2_annotate(struct kvm_pgtable *pgt, u64 addr, u64 size,
  */
 int kvm_pgtable_stage2_unmap(struct kvm_pgtable *pgt, u64 addr, u64 size);
 
+/**
+ * kvm_pgtable_stage2_unmap_range() - Remove a mapping from a guest stage-2
+ *                                    page-table, and report where the walk
+ *                                    stopped.
+ * @pgt:	Page-table structure initialised by kvm_pgtable_stage2_init*().
+ * @addr:	Intermediate physical address from which to remove the mapping.
+ * @size:	Size of the mapping.
+ *
+ * Same as kvm_pgtable_stage2_unmap(), but the walk can end past @addr + @size
+ * when it steps over an invalid entry that extends beyond the range. Nothing
+ * is mapped up to the address that is returned, so a caller that unmaps a
+ * larger range in chunks can resume from it.
+ *
+ * Return: the address the walk reached, which can be past @addr + @size, or a
+ * negative error code on failure.
+ */
+s64 kvm_pgtable_stage2_unmap_range(struct kvm_pgtable *pgt, u64 addr, u64 size);
+
 /**
  * kvm_pgtable_stage2_wrprotect() - Write-protect guest stage-2 address range
  *                                  without TLB invalidation.
diff --git a/arch/arm64/kvm/hyp/pgtable.c b/arch/arm64/kvm/hyp/pgtable.c
index 9c83837..b02f548 100644
--- a/arch/arm64/kvm/hyp/pgtable.c
+++ b/arch/arm64/kvm/hyp/pgtable.c
@@ -1197,7 +1197,7 @@ static int stage2_unmap_walker(const struct kvm_pgtable_visit_ctx *ctx,
 	return 0;
 }
 
-int kvm_pgtable_stage2_unmap(struct kvm_pgtable *pgt, u64 addr, u64 size)
+s64 kvm_pgtable_stage2_unmap_range(struct kvm_pgtable *pgt, u64 addr, u64 size)
 {
 	int ret;
 	struct kvm_pgtable_walker walker = {
@@ -1205,13 +1205,29 @@ int kvm_pgtable_stage2_unmap(struct kvm_pgtable *pgt, u64 addr, u64 size)
 		.arg	= pgt,
 		.flags	= KVM_PGTABLE_WALK_LEAF | KVM_PGTABLE_WALK_TABLE_POST,
 	};
+	struct kvm_pgtable_walk_data data = {
+		.walker	= &walker,
+		.start	= ALIGN_DOWN(addr, PAGE_SIZE),
+		.addr	= ALIGN_DOWN(addr, PAGE_SIZE),
+		.end	= PAGE_ALIGN(data.addr + size),
+	};
 
-	ret = kvm_pgtable_walk(pgt, addr, size, &walker);
+	ret = _kvm_pgtable_walk(pgt, &data);
 	if (stage2_unmap_defer_tlb_flush(pgt))
 		/* Perform the deferred TLB invalidations */
 		kvm_tlb_flush_vmid_range(pgt->mmu, addr, size);
 
-	return ret;
+	if (ret)
+		return ret;
+
+	return data.addr;
+}
+
+int kvm_pgtable_stage2_unmap(struct kvm_pgtable *pgt, u64 addr, u64 size)
+{
+	s64 ret = kvm_pgtable_stage2_unmap_range(pgt, addr, size);
+
+	return ret < 0 ? ret : 0;
 }
 
 struct stage2_attr_data {
diff --git a/arch/arm64/kvm/mmu.c b/arch/arm64/kvm/mmu.c
index 1db10d8..55a3cb5 100644
--- a/arch/arm64/kvm/mmu.c
+++ b/arch/arm64/kvm/mmu.c
@@ -319,22 +319,64 @@ static void invalidate_icache_guest_page(void *va, size_t size)
 	__invalidate_icache_guest_page(va, size);
 }
 
-static int kvm_pgtable_stage2_unmap_tracked(struct kvm_pgtable *pgt,
+static s64 kvm_pgtable_stage2_unmap_tracked(struct kvm_pgtable *pgt,
 					    u64 addr, u64 size);
 
+/*
+ * Unmap [start, end) one chunk at a time, resuming from the address that each
+ * walk reached. An invalid entry that is larger than a chunk is then stepped
+ * over once, rather than once for every chunk that it covers.
+ */
+static int stage2_apply_unmap_range(struct kvm_s2_mmu *mmu, phys_addr_t start,
+				    phys_addr_t end,
+				    s64 (*fn)(struct kvm_pgtable *, u64, u64),
+				    bool resched)
+{
+	struct kvm *kvm = kvm_s2_mmu_to_kvm(mmu);
+	bool lock_dropped = false;
+	phys_addr_t addr = start;
+	u64 next;
+	s64 ret;
+
+	do {
+		struct kvm_pgtable *pgt = mmu->pgt;
+
+		/*
+		 * We may be raced on PGT teardown when we release the
+		 * kvm->mmu_lock. That's fine as the PGT is legitimately no
+		 * longer present.
+		 */
+		if (!pgt)
+			return lock_dropped ? 0 : -EINVAL;
+
+		next = stage2_range_addr_end(addr, end);
+		ret = fn(pgt, addr, next - addr);
+		if (ret < 0)
+			return ret;
+
+		next = min_t(u64, ret, end);
+		if (resched && next != end) {
+			cond_resched_rwlock_write(&kvm->mmu_lock);
+			lock_dropped = true;
+		}
+	} while (addr = next, addr != end);
+
+	return 0;
+}
+
 static void kvm_vm_stage2_unmap_range(struct kvm_s2_mmu *mmu,
 				      phys_addr_t start,
 				      u64 size, bool may_block)
 {
 	struct kvm *kvm = kvm_s2_mmu_to_kvm(mmu);
-	int (*fn)(struct kvm_pgtable *, u64, u64);
+	s64 (*fn)(struct kvm_pgtable *pgt, u64 addr, u64 size);
 
 	if (kvm_is_nested_s2_mmu(kvm, mmu))
 		fn = kvm_pgtable_stage2_unmap_tracked;
 	else
-		fn = kvm_pgtable_stage2_unmap;
+		fn = kvm_pgtable_stage2_unmap_range;
 
-	WARN_ON(stage2_apply_range(mmu, start, start + size, fn, may_block));
+	WARN_ON(stage2_apply_unmap_range(mmu, start, start + size, fn, may_block));
 }
 
 static void pkvm_stage2_unmap_range(struct kvm_s2_mmu *mmu,
@@ -375,16 +417,17 @@ static void no_stage2_unmap_range(struct kvm_s2_mmu *mmu,
  * does.
  */
 
-static int kvm_pgtable_stage2_unmap_tracked(struct kvm_pgtable *pgt, u64 addr, u64 size)
+static s64 kvm_pgtable_stage2_unmap_tracked(struct kvm_pgtable *pgt, u64 addr, u64 size)
 {
-	int ret;
+	s64 ret;
 
-	ret = kvm_pgtable_stage2_unmap(pgt, addr, size);
-	if (ret)
+	ret = kvm_pgtable_stage2_unmap_range(pgt, addr, size);
+	if (ret < 0)
 		return ret;
 
-	kvm_remove_guest_s2_mappings(pgt->mmu, addr, size);
-	return 0;
+	/* The walk may have stepped over invalid entries: any record there is stale. */
+	kvm_remove_guest_s2_mappings(pgt->mmu, addr, ret - addr);
+	return ret;
 }
 
 /**
-- 
2.43.0