From f4812c6864dce4c269f78ce3931fbee99515cbc8 Mon Sep 17 00:00:00 2001
From: Fuad Tabba <tabba@google.com>
Date: Wed, 11 Oct 2023 13:55:44 +0100
Subject: [PATCH] ANDROID: KVM: arm64: Store hyp address in the host fp state
 array

Store the hyp address in kvm_arm_hyp_host_fp_state[], to avoid having to
calculate it with kern_hyp_va() on every access.

Bug: 303684934
Signed-off-by: Fuad Tabba <tabba@google.com>
Change-Id: I52902c297f9b957a8d035be942e3cbeb32fed0a2
---
 arch/arm64/kvm/arm.c           | 21 +++++++++++++++++++++
 arch/arm64/kvm/hyp/nvhe/pkvm.c |  6 +++++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/arch/arm64/kvm/arm.c b/arch/arm64/kvm/arm.c
index de68a5c2cfa2..c9428ce23d80 100644
--- a/arch/arm64/kvm/arm.c
+++ b/arch/arm64/kvm/arm.c
@@ -2127,6 +2127,20 @@ static int init_pkvm_host_fp_state(void)
 	return 0;
 }
 
+/*
+ * Finalizes the initialization of hyp mode, once everything else is initialized
+ * and the initialziation process cannot fail.
+ */
+static void finalize_init_hyp_mode(void)
+{
+	int cpu;
+
+	for_each_possible_cpu(cpu) {
+		kvm_nvhe_sym(kvm_arm_hyp_host_fp_state)[cpu] =
+			kern_hyp_va(kvm_nvhe_sym(kvm_arm_hyp_host_fp_state)[cpu]);
+	}
+}
+
 /**
  * Inits Hyp-mode on all online CPUs
  */
@@ -2456,6 +2470,13 @@ int kvm_arch_init(void *opaque)
 		kvm_info("Hyp mode initialized successfully\n");
 	}
 
+	/*
+	 * This should be called after initialization is done and failure isn't
+	 * possible anymore.
+	 */
+	if (!in_hyp_mode)
+		finalize_init_hyp_mode();
+
 	return 0;
 
 out_hyp:
diff --git a/arch/arm64/kvm/hyp/nvhe/pkvm.c b/arch/arm64/kvm/hyp/nvhe/pkvm.c
index 8a74afc5da5f..b3acb60a2d6c 100644
--- a/arch/arm64/kvm/hyp/nvhe/pkvm.c
+++ b/arch/arm64/kvm/hyp/nvhe/pkvm.c
@@ -45,7 +45,11 @@ unsigned long __ro_after_init kvm_arm_hyp_host_fp_state[NR_CPUS];
 
 static void *__get_host_fpsimd_bytes(void)
 {
-	return kern_hyp_va((void *) kvm_arm_hyp_host_fp_state[hyp_smp_processor_id()]);
+	/*
+	 * The addresses in this array have been converted to hyp addresses
+	 * in finalize_init_hyp_mode().
+	 */
+	return (void *)kvm_arm_hyp_host_fp_state[hyp_smp_processor_id()];
 }
 
 struct user_fpsimd_state *get_host_fpsimd_state(struct kvm_vcpu *vcpu)