[RFC PATCH 23/26] kvm: arm64: Intercept host's CPU_SUSPEND PSCI SMCs

[David Brazdil <dbrazdil@google.com>]

Add a handler of CPU_SUSPEND host PSCI SMCs. When invoked, it determines
whether the requested power state loses context, ie. whether it is
indistinguishable from a WHI or whether it is a deeper sleep state that
behaves like a CPU_OFF+CPU_ON.

If it's the former, it forwards the call to EL3 and returns to the host
after waking up.

If it's the latter, it saves r0,pc of the host into and makes the same
call to EL3 with the hyp CPU entry point. When the core wakes up, EL2
state is initialized before dropping back to EL1.

Signed-off-by: David Brazdil <dbrazdil@google.com>
---
 arch/arm64/kvm/arm.c           |  2 ++
 arch/arm64/kvm/hyp/nvhe/psci.c | 49 +++++++++++++++++++++++++++++++++-
 drivers/firmware/psci/psci.c   |  9 -------
 include/uapi/linux/psci.h      |  7 +++++
 4 files changed, 57 insertions(+), 10 deletions(-)

diff --git a/arch/arm64/kvm/arm.c b/arch/arm64/kvm/arm.c
index 166975999ead..6fbda652200b 100644
--- a/arch/arm64/kvm/arm.c
+++ b/arch/arm64/kvm/arm.c
@@ -1521,9 +1521,11 @@ static void init_psci(void)
 {
 	extern u32 kvm_nvhe_sym(kvm_host_psci_version);
 	extern u32 kvm_nvhe_sym(kvm_host_psci_function_id)[PSCI_FN_MAX];
+	extern u32 kvm_nvhe_sym(kvm_host_psci_cpu_suspend_feature);
 	int cpu;
 
 	kvm_nvhe_sym(kvm_host_psci_version) = psci_driver_version;
+	kvm_nvhe_sym(kvm_host_psci_cpu_suspend_feature) = psci_cpu_suspend_feature;
 	memcpy(kvm_nvhe_sym(kvm_host_psci_function_id),
 		psci_function_id, sizeof(psci_function_id));
 
diff --git a/arch/arm64/kvm/hyp/nvhe/psci.c b/arch/arm64/kvm/hyp/nvhe/psci.c
index 42ee5effa827..4899c8319bb4 100644
--- a/arch/arm64/kvm/hyp/nvhe/psci.c
+++ b/arch/arm64/kvm/hyp/nvhe/psci.c
@@ -21,6 +21,7 @@
 /* Config options set by the host. */
 u32 kvm_host_psci_version = PSCI_VERSION(0, 0);
 u32 kvm_host_psci_function_id[PSCI_FN_MAX];
+u32 kvm_host_psci_cpu_suspend_feature;
 s64 hyp_physvirt_offset;
 
 #define __hyp_pa(x) ((phys_addr_t)(x) + hyp_physvirt_offset)
@@ -83,6 +84,20 @@ static __noreturn unsigned long psci_forward_noreturn(struct kvm_cpu_context *ho
 	hyp_panic(); /* unreachable */
 }
 
+static bool psci_has_ext_power_state(void)
+{
+	return kvm_host_psci_cpu_suspend_feature & PSCI_1_0_FEATURES_CPU_SUSPEND_PF_MASK;
+}
+
+static bool psci_power_state_loses_context(u32 state)
+{
+	const u32 mask = psci_has_ext_power_state() ?
+					PSCI_1_0_EXT_POWER_STATE_TYPE_MASK :
+					PSCI_0_2_POWER_STATE_TYPE_MASK;
+
+	return state & mask;
+}
+
 static unsigned int find_cpu_id(u64 mpidr)
 {
 	int i;
@@ -106,6 +121,34 @@ static phys_addr_t cpu_entry_pa(void)
 	return kern_va - kimage_voffset;
 }
 
+static int psci_cpu_suspend(u64 func_id, struct kvm_cpu_context *host_ctxt)
+{
+	u64 power_state = host_ctxt->regs.regs[1];
+	unsigned long pc = host_ctxt->regs.regs[2];
+	unsigned long r0 = host_ctxt->regs.regs[3];
+	hyp_spinlock_t *cpu_lock;
+	struct vcpu_reset_state	*cpu_reset;
+	struct kvm_nvhe_init_params *cpu_params;
+
+	if (!psci_power_state_loses_context(power_state)) {
+		/* This power state has the same semantics as WFI. */
+		return psci_call(PSCI_0_2_FN64_CPU_SUSPEND, 0, 0, 0);
+	}
+
+	cpu_lock = this_cpu_ptr(&psci_cpu_lock);
+	cpu_reset = this_cpu_ptr(&psci_cpu_reset);
+	cpu_params = this_cpu_ptr(&kvm_init_params);
+
+	/* Resuming from this state has the same semantics as CPU_ON. */
+	hyp_spin_lock(cpu_lock);
+	*cpu_reset = (struct vcpu_reset_state){
+		.pc = pc,
+		.r0 = r0,
+	};
+	hyp_spin_unlock(cpu_lock);
+	return psci_call(func_id, power_state, cpu_entry_pa(), __hyp_pa(cpu_params));
+}
+
 static int psci_cpu_off(u64 func_id, struct kvm_cpu_context *host_ctxt)
 {
 	hyp_spinlock_t *cpu_lock = this_cpu_ptr(&psci_cpu_lock);
@@ -193,7 +236,9 @@ static int psci_cpu_on(u64 func_id, struct kvm_cpu_context *host_ctxt)
 
 static unsigned long psci_0_1_handler(u64 func_id, struct kvm_cpu_context *host_ctxt)
 {
-	if (func_id == kvm_host_psci_function_id[PSCI_FN_CPU_OFF])
+	if (func_id == kvm_host_psci_function_id[PSCI_FN_CPU_SUSPEND])
+		return psci_cpu_suspend(func_id, host_ctxt);
+	else if (func_id == kvm_host_psci_function_id[PSCI_FN_CPU_OFF])
 		return psci_cpu_off(func_id, host_ctxt);
 	else if (func_id == kvm_host_psci_function_id[PSCI_FN_CPU_ON])
 		return psci_cpu_on(func_id, host_ctxt);
@@ -216,6 +261,8 @@ static unsigned long psci_0_2_handler(u64 func_id, struct kvm_cpu_context *host_
 	case PSCI_0_2_FN_SYSTEM_RESET:
 		psci_forward_noreturn(host_ctxt);
 		unreachable();
+	case PSCI_0_2_FN64_CPU_SUSPEND:
+		return psci_cpu_suspend(func_id, host_ctxt);
 	case PSCI_0_2_FN_CPU_OFF:
 		return psci_cpu_off(func_id, host_ctxt);
 	case PSCI_0_2_FN64_CPU_ON:
diff --git a/drivers/firmware/psci/psci.c b/drivers/firmware/psci/psci.c
index b6ad237b1518..387e24409da7 100644
--- a/drivers/firmware/psci/psci.c
+++ b/drivers/firmware/psci/psci.c
@@ -62,15 +62,6 @@ static psci_fn *invoke_psci_fn;
 
 u32 psci_function_id[PSCI_FN_MAX];
 
-#define PSCI_0_2_POWER_STATE_MASK		\
-				(PSCI_0_2_POWER_STATE_ID_MASK | \
-				PSCI_0_2_POWER_STATE_TYPE_MASK | \
-				PSCI_0_2_POWER_STATE_AFFL_MASK)
-
-#define PSCI_1_0_EXT_POWER_STATE_MASK		\
-				(PSCI_1_0_EXT_POWER_STATE_ID_MASK | \
-				PSCI_1_0_EXT_POWER_STATE_TYPE_MASK)
-
 u32 psci_cpu_suspend_feature;
 static bool psci_system_reset2_supported;
 
diff --git a/include/uapi/linux/psci.h b/include/uapi/linux/psci.h
index 0d52b8dbe8c2..df3d85ce86f7 100644
--- a/include/uapi/linux/psci.h
+++ b/include/uapi/linux/psci.h
@@ -65,6 +65,10 @@
 #define PSCI_0_2_POWER_STATE_AFFL_SHIFT		24
 #define PSCI_0_2_POWER_STATE_AFFL_MASK		\
 				(0x3 << PSCI_0_2_POWER_STATE_AFFL_SHIFT)
+#define PSCI_0_2_POWER_STATE_MASK					\
+				(PSCI_0_2_POWER_STATE_ID_MASK |		\
+				 PSCI_0_2_POWER_STATE_TYPE_MASK |	\
+				 PSCI_0_2_POWER_STATE_AFFL_MASK)
 
 /* PSCI extended power state encoding for CPU_SUSPEND function */
 #define PSCI_1_0_EXT_POWER_STATE_ID_MASK	0xfffffff
@@ -72,6 +76,9 @@
 #define PSCI_1_0_EXT_POWER_STATE_TYPE_SHIFT	30
 #define PSCI_1_0_EXT_POWER_STATE_TYPE_MASK	\
 				(0x1 << PSCI_1_0_EXT_POWER_STATE_TYPE_SHIFT)
+#define PSCI_1_0_EXT_POWER_STATE_MASK					\
+				(PSCI_1_0_EXT_POWER_STATE_ID_MASK |	\
+				 PSCI_1_0_EXT_POWER_STATE_TYPE_MASK)
 
 /* PSCI v0.2 affinity level state returned by AFFINITY_INFO */
 #define PSCI_0_2_AFFINITY_LEVEL_ON		0
-- 
2.29.1.341.ge80a0c044ae-goog