feat(psci): add support for PSCI_SET_SUSPEND_MODE

This patch adds a PSCI_SET_SUSPEND_MODE handler that validates the
request per section 5.20.2 of the PSCI spec (DEN0022D.b), and updates
the suspend mode to the requested mode.

This is conditionally compiled into the build depending on the value of
the `PSCI_OS_INIT_MODE` build option.

Change-Id: Iebf65f5f7846aef6b8643ad6082db99b4dcc4bef
Signed-off-by: Wing Li <wingers@google.com>
diff --git a/lib/psci/psci_common.c b/lib/psci/psci_common.c
index f233be1..5ce562f 100644
--- a/lib/psci/psci_common.c
+++ b/lib/psci/psci_common.c
@@ -76,6 +76,14 @@
 	(PLAT_MAX_PWR_LVL >= PSCI_CPU_PWR_LVL),
 	assert_platform_max_pwrlvl_check);
 
+#if PSCI_OS_INIT_MODE
+/*******************************************************************************
+ * The power state coordination mode used in CPU_SUSPEND.
+ * Defaults to platform-coordinated mode.
+ ******************************************************************************/
+suspend_mode_t psci_suspend_mode = PLAT_COORD;
+#endif
+
 /*
  * The plat_local_state used by the platform is one of these types: RUN,
  * RETENTION and OFF. The platform can define further sub-states for each type
@@ -154,7 +162,7 @@
 }
 
 /*******************************************************************************
- * This function verifies that the all the other cores in the system have been
+ * This function verifies that all the other cores in the system have been
  * turned OFF and the current CPU is the last running CPU in the system.
  * Returns true, if the current CPU is the last ON CPU or false otherwise.
  ******************************************************************************/
@@ -179,6 +187,23 @@
 }
 
 /*******************************************************************************
+ * This function verifies that all cores in the system have been turned ON.
+ * Returns true, if all CPUs are ON or false otherwise.
+ ******************************************************************************/
+static bool psci_are_all_cpus_on(void)
+{
+	unsigned int cpu_idx;
+
+	for (cpu_idx = 0; cpu_idx < psci_plat_core_count; cpu_idx++) {
+		if (psci_get_aff_info_state_by_idx(cpu_idx) == AFF_STATE_OFF) {
+			return false;
+		}
+	}
+
+	return true;
+}
+
+/*******************************************************************************
  * Routine to return the maximum power level to traverse to after a cpu has
  * been physically powered up. It is expected to be called immediately after
  * reset from assembler code.
@@ -1050,3 +1075,29 @@
 
 	return true;
 }
+
+/*******************************************************************************
+ * This function verifies that all cores in the system have been turned ON.
+ * Returns true, if all CPUs are ON or false otherwise.
+ *
+ * This API has following differences with psci_are_all_cpus_on
+ *  1. PSCI states are locked
+ ******************************************************************************/
+bool psci_are_all_cpus_on_safe(void)
+{
+	unsigned int this_core = plat_my_core_pos();
+	unsigned int parent_nodes[PLAT_MAX_PWR_LVL] = {0};
+
+	psci_get_parent_pwr_domain_nodes(this_core, PLAT_MAX_PWR_LVL, parent_nodes);
+
+	psci_acquire_pwr_domain_locks(PLAT_MAX_PWR_LVL, parent_nodes);
+
+	if (!psci_are_all_cpus_on()) {
+		psci_release_pwr_domain_locks(PLAT_MAX_PWR_LVL, parent_nodes);
+		return false;
+	}
+
+	psci_release_pwr_domain_locks(PLAT_MAX_PWR_LVL, parent_nodes);
+
+	return true;
+}
diff --git a/lib/psci/psci_main.c b/lib/psci/psci_main.c
index a631f3f..4b5fc81 100644
--- a/lib/psci/psci_main.c
+++ b/lib/psci/psci_main.c
@@ -370,6 +370,39 @@
 	return PSCI_E_SUCCESS;
 }
 
+#if PSCI_OS_INIT_MODE
+int psci_set_suspend_mode(unsigned int mode)
+{
+	if (psci_suspend_mode == mode) {
+		return PSCI_E_SUCCESS;
+	}
+
+	if (mode == PLAT_COORD) {
+		/* Check if the current CPU is the last ON CPU in the system */
+		if (!psci_is_last_on_cpu_safe()) {
+			return PSCI_E_DENIED;
+		}
+	}
+
+	if (mode == OS_INIT) {
+		/*
+		 * Check if all CPUs in the system are ON or if the current
+		 * CPU is the last ON CPU in the system.
+		 */
+		if (!(psci_are_all_cpus_on_safe() ||
+		      psci_is_last_on_cpu_safe())) {
+			return PSCI_E_DENIED;
+		}
+	}
+
+	psci_suspend_mode = mode;
+	psci_flush_dcache_range((uintptr_t)&psci_suspend_mode,
+				sizeof(psci_suspend_mode));
+
+	return PSCI_E_SUCCESS;
+}
+#endif
+
 /*******************************************************************************
  * PSCI top level handler for servicing SMCs.
  ******************************************************************************/
@@ -452,6 +485,12 @@
 		case PSCI_FEATURES:
 			ret = (u_register_t)psci_features(r1);
 			break;
+
+#if PSCI_OS_INIT_MODE
+		case PSCI_SET_SUSPEND_MODE:
+			ret = (u_register_t)psci_set_suspend_mode(r1);
+			break;
+#endif
 
 #if ENABLE_PSCI_STAT
 		case PSCI_STAT_RESIDENCY_AARCH32:
@@ -515,6 +554,12 @@
 			ret = (u_register_t)psci_system_suspend(x1, x2);
 			break;
 
+#if PSCI_OS_INIT_MODE
+		case PSCI_SET_SUSPEND_MODE:
+			ret = (u_register_t)psci_set_suspend_mode(x1);
+			break;
+#endif
+
 #if ENABLE_PSCI_STAT
 		case PSCI_STAT_RESIDENCY_AARCH64:
 			ret = psci_stat_residency(x1, (unsigned int) x2);
diff --git a/lib/psci/psci_private.h b/lib/psci/psci_private.h
index 6ca9ef6..73b161e 100644
--- a/lib/psci/psci_private.h
+++ b/lib/psci/psci_private.h
@@ -163,6 +163,16 @@
 	spinlock_t cpu_lock;
 } cpu_pd_node_t;
 
+#if PSCI_OS_INIT_MODE
+/*******************************************************************************
+ * The supported power state coordination modes that can be used in CPU_SUSPEND.
+ ******************************************************************************/
+typedef enum suspend_mode {
+	PLAT_COORD = 0,
+	OS_INIT = 1
+} suspend_mode_t;
+#endif
+
 /*******************************************************************************
  * The following are helpers and declarations of locks.
  ******************************************************************************/
@@ -260,6 +270,9 @@
 extern cpu_pd_node_t psci_cpu_pd_nodes[PLATFORM_CORE_COUNT];
 extern unsigned int psci_caps;
 extern unsigned int psci_plat_core_count;
+#if PSCI_OS_INIT_MODE
+extern suspend_mode_t psci_suspend_mode;
+#endif
 
 /*******************************************************************************
  * SPD's power management hooks registered with PSCI
diff --git a/lib/psci/psci_setup.c b/lib/psci/psci_setup.c
index 3cb4f7e..16d6e45 100644
--- a/lib/psci/psci_setup.c
+++ b/lib/psci/psci_setup.c
@@ -254,6 +254,9 @@
 			psci_caps |=  define_psci_cap(PSCI_CPU_SUSPEND_AARCH64);
 		if (psci_plat_pm_ops->get_sys_suspend_power_state != NULL)
 			psci_caps |=  define_psci_cap(PSCI_SYSTEM_SUSPEND_AARCH64);
+#if PSCI_OS_INIT_MODE
+		psci_caps |= define_psci_cap(PSCI_SET_SUSPEND_MODE);
+#endif
 	}
 	if (psci_plat_pm_ops->system_off != NULL)
 		psci_caps |=  define_psci_cap(PSCI_SYSTEM_OFF);