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

@@ -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:

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)