With NV, "stage2 mmu" becomes ambiguous, it could refer to the canonical
one or one of the nested ones. Modify kvm_uninit_stage2_mmu() so that it
can uninitialize both the canonical s2 mmu and nested s2 mmus.

Signed-off-by: Wei-Lin Chang <weilin.chang@arm.com>
---
 arch/arm64/include/asm/kvm_mmu.h | 2 +-
 arch/arm64/kvm/mmu.c             | 8 +++++---
 arch/arm64/kvm/nested.c          | 4 ++--
 3 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/arch/arm64/include/asm/kvm_mmu.h b/arch/arm64/include/asm/kvm_mmu.h
index 960b6aed4ffa..4b49aabecbf3 100644
--- a/arch/arm64/include/asm/kvm_mmu.h
+++ b/arch/arm64/include/asm/kvm_mmu.h
@@ -176,7 +176,7 @@ void kvm_stage2_wp_range(struct kvm_s2_mmu *mmu, phys_addr_t addr, phys_addr_t e
 
 void stage2_unmap_vm(struct kvm *kvm);
 int kvm_init_stage2_mmu(struct kvm *kvm, struct kvm_s2_mmu *mmu, unsigned long type);
-void kvm_uninit_stage2_mmu(struct kvm *kvm);
+void kvm_uninit_stage2_mmu(struct kvm *kvm, struct kvm_s2_mmu *mmu);
 void kvm_free_stage2(struct kvm_s2_mmu *mmu);
 int kvm_phys_addr_ioremap(struct kvm *kvm, phys_addr_t guest_ipa,
 			  phys_addr_t pa, unsigned long size, bool writable);
diff --git a/arch/arm64/kvm/mmu.c b/arch/arm64/kvm/mmu.c
index b19b9a9b3c27..a2affd70eca6 100644
--- a/arch/arm64/kvm/mmu.c
+++ b/arch/arm64/kvm/mmu.c
@@ -1018,10 +1018,12 @@ int kvm_init_stage2_mmu(struct kvm *kvm, struct kvm_s2_mmu *mmu, unsigned long t
 	return err;
 }
 
-void kvm_uninit_stage2_mmu(struct kvm *kvm)
+void kvm_uninit_stage2_mmu(struct kvm *kvm, struct kvm_s2_mmu *mmu)
 {
-	kvm_free_stage2(&kvm->arch.mmu);
-	kvm_mmu_free_memory_cache(&kvm->arch.mmu.split_page_cache);
+	kvm_free_stage2(mmu);
+
+	if (!kvm_is_nested_s2_mmu(kvm, mmu))
+		kvm_mmu_free_memory_cache(&mmu->split_page_cache);
 }
 
 static void stage2_unmap_memslot(struct kvm *kvm,
diff --git a/arch/arm64/kvm/nested.c b/arch/arm64/kvm/nested.c
index 772d922cf0ee..5eba94d6cc67 100644
--- a/arch/arm64/kvm/nested.c
+++ b/arch/arm64/kvm/nested.c
@@ -1194,12 +1194,12 @@ void kvm_arch_flush_shadow_all(struct kvm *kvm)
 		struct kvm_s2_mmu *mmu = &kvm->arch.nested_mmus[i];
 
 		if (!WARN_ON(atomic_read(&mmu->refcnt)))
-			kvm_free_stage2(mmu);
+			kvm_uninit_stage2_mmu(kvm, mmu);
 	}
 	kvfree(kvm->arch.nested_mmus);
 	kvm->arch.nested_mmus = NULL;
 	kvm->arch.nested_mmus_size = 0;
-	kvm_uninit_stage2_mmu(kvm);
+	kvm_uninit_stage2_mmu(kvm, &kvm->arch.mmu);
 }
 
 /*
-- 
2.43.0