fixup!  KVM: arm64: Support TLB invalidation in guest context

- Save the old MMU context rather than the new one (!)

Signed-off-by: Will Deacon <will@kernel.org>
diff --git a/arch/arm64/kvm/hyp/nvhe/tlb.c b/arch/arm64/kvm/hyp/nvhe/tlb.c
index 9d66cfb..eb97294 100644
--- a/arch/arm64/kvm/hyp/nvhe/tlb.c
+++ b/arch/arm64/kvm/hyp/nvhe/tlb.c
@@ -33,10 +33,17 @@ static void enter_vmid_context(struct kvm_s2_mmu *mmu,
 	 * to do.
 	 */
 	if (vcpu) {
+		/* We're in guest context */
 		if (mmu == vcpu->arch.hw_mmu || WARN_ON(mmu != host_s2_mmu))
 			return;
-	} else if (mmu == host_s2_mmu) {
-		return;
+
+		cxt->mmu = vcpu->arch.hw_mmu;
+	} else {
+		/* We're in host context */
+		if (mmu == host_s2_mmu)
+			return;
+
+		cxt->mmu = host_s2_mmu;
 	}
 
 	/*
@@ -61,7 +68,6 @@ static void enter_vmid_context(struct kvm_s2_mmu *mmu,
 	else
 		dsb(ish);
 
-	cxt->mmu = mmu;
 	if (cpus_have_final_cap(ARM64_WORKAROUND_SPECULATIVE_AT)) {
 		u64 val;