[PATCH 3/3] KVM: arm64: Move stage-2 teardown into kvm_vm_s2_ops

[Fuad Tabba <fuad.tabba@linux.dev>]

Replace kvm_stage2_destroy()'s is_protected_kvm_enabled() and
KVM_PGT_FN() dispatch with a vm_stage2_destroy hook in kvm_vm_s2_ops,
like the other stage-2 operations. Both pKVM flavors share
pkvm_stage2_destroy(). Like kvm_vm_stage2_destroy(), it takes the range
from kvm_phys_size().

No functional change intended.

Signed-off-by: Fuad Tabba <fuad.tabba@linux.dev>
---
 arch/arm64/include/asm/kvm_host.h |  2 ++
 arch/arm64/kvm/mmu.c              | 24 +++++++++++++-----------
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/arch/arm64/include/asm/kvm_host.h b/arch/arm64/include/asm/kvm_host.h
index 97332ffcb1300..b36779bafae9b 100644
--- a/arch/arm64/include/asm/kvm_host.h
+++ b/arch/arm64/include/asm/kvm_host.h
@@ -233,6 +233,7 @@ struct kvm_vcpu_ops {
 
 struct kvm_gfn_range;
 struct kvm_s2_fault_desc;
+struct kvm_pgtable;
 
 struct kvm_vm_s2_ops {
 	bool (*vm_age_gfn)(struct kvm *kvm, struct kvm_gfn_range *range);
@@ -243,6 +244,7 @@ struct kvm_vm_s2_ops {
 	void (*vm_stage2_unmap_range)(struct kvm_s2_mmu *mmu,
 				      phys_addr_t start, u64 size,
 				      bool may_block);
+	void (*vm_stage2_destroy)(struct kvm_pgtable *pgt);
 	int (*vm_mem_abort)(const struct kvm_s2_fault_desc *s2fd);
 };
 
diff --git a/arch/arm64/kvm/mmu.c b/arch/arm64/kvm/mmu.c
index 10b425a0141be..9df3c80a5acd1 100644
--- a/arch/arm64/kvm/mmu.c
+++ b/arch/arm64/kvm/mmu.c
@@ -1029,17 +1029,16 @@ static void stage2_destroy_range(struct kvm_pgtable *pgt, phys_addr_t addr,
 	} while (addr = next, addr != end);
 }
 
-static void kvm_stage2_destroy(struct kvm_pgtable *pgt)
+static void kvm_vm_stage2_destroy(struct kvm_pgtable *pgt)
 {
-	unsigned int ia_bits = VTCR_EL2_IPA(pgt->mmu->vtcr);
+	stage2_destroy_range(pgt, 0, kvm_phys_size(pgt->mmu));
+	kvm_pgtable_stage2_destroy_pgd(pgt);
+}
 
-	/* pKVM walks its mapping tree, not the IPA space, and reschedules per mapping */
-	if (is_protected_kvm_enabled())
-		pkvm_pgtable_stage2_destroy_range(pgt, 0, BIT(ia_bits));
-	else
-		stage2_destroy_range(pgt, 0, BIT(ia_bits));
-
-	KVM_PGT_FN(kvm_pgtable_stage2_destroy_pgd)(pgt);
+static void pkvm_stage2_destroy(struct kvm_pgtable *pgt)
+{
+	pkvm_pgtable_stage2_destroy_range(pgt, 0, kvm_phys_size(pgt->mmu));
+	pkvm_pgtable_stage2_destroy_pgd(pgt);
 }
 
 /**
@@ -1126,7 +1125,7 @@ int kvm_init_stage2_mmu(struct kvm *kvm, struct kvm_s2_mmu *mmu, unsigned long t
 	return 0;
 
 out_destroy_pgtable:
-	kvm_stage2_destroy(pgt);
+	kvm->arch.vm_s2_ops->vm_stage2_destroy(pgt);
 	mmu->pgt = NULL;
 out_free_pgtable:
 	kfree(pgt);
@@ -1246,7 +1245,7 @@ void kvm_free_stage2_pgd(struct kvm_s2_mmu *mmu)
 	write_unlock(&kvm->mmu_lock);
 
 	if (pgt) {
-		kvm_stage2_destroy(pgt);
+		kvm->arch.vm_s2_ops->vm_stage2_destroy(pgt);
 		kfree(pgt);
 	}
 
@@ -3152,6 +3151,7 @@ static const struct kvm_vm_s2_ops protected_pkvm_vm_s2_ops = {
 	.vm_age_gfn			= no_age_gfn,
 	.vm_test_age_gfn		= no_age_gfn,
 	.vm_stage2_unmap_range		= no_stage2_unmap_range,
+	.vm_stage2_destroy		= pkvm_stage2_destroy,
 	.vm_mem_abort			= protected_pkvm_mem_abort,
 };
 
@@ -3161,6 +3161,7 @@ static const struct kvm_vm_s2_ops pkvm_vm_s2_ops = {
 	.vm_age_gfn			= pkvm_age_gfn,
 	.vm_test_age_gfn		= pkvm_test_age_gfn,
 	.vm_stage2_unmap_range		= pkvm_stage2_unmap_range,
+	.vm_stage2_destroy		= pkvm_stage2_destroy,
 	.vm_mem_abort			= kvm_vm_mem_abort,
 };
 
@@ -3170,6 +3171,7 @@ static const struct kvm_vm_s2_ops kvm_default_vm_s2_ops = {
 	.vm_age_gfn			= kvm_vm_age_gfn,
 	.vm_test_age_gfn		= kvm_vm_test_age_gfn,
 	.vm_stage2_unmap_range		= kvm_vm_stage2_unmap_range,
+	.vm_stage2_destroy		= kvm_vm_stage2_destroy,
 	.vm_mem_abort			= kvm_vm_mem_abort,
 };
 
-- 
2.47.3