PSCI: Support hierarchical PM domains

index ffca4fc..4df7d45 100644
@@ -98,6 +98,39 @@ static int init_state_node(struct cpuidle_state *idle_state,
 }
 
 /*
+ * Get the state node at @idx. State node may be defined as domain's idle state
+ * if the CPU has its own domain or defined as CPU's idle state if it doesn't
+ * have a domain provider.
+ */
+static struct device_node *get_state_node(struct device_node *cpu_node,
+				unsigned int idx)
+{
+	struct device_node *dn;
+	bool cpu_has_domain = false;
+	struct of_phandle_args args;
+	const char *property;
+	int err;
+
+	err = of_parse_phandle_with_args(cpu_node, "power-domains",
+					"#power-domain-cells", 0, &args);
+	if (!err) {
+		dn = args.np;
+		err = of_count_phandle_with_args(dn, "domain-idle-states",
+								NULL);
+		cpu_has_domain = (err > 0);
+	}
+
+	if (cpu_has_domain) {
+		property = "domain-idle-states";
+	} else {
+		property = "cpu-idle-states";
+		dn = cpu_node;
+	}
+
+	return of_parse_phandle(dn, property, idx);
+}
+
+/*
  * Check that the idle state is uniform across all CPUs in the CPUidle driver
  * cpumask
  */

@@ -118,8 +151,7 @@ static bool idle_state_valid(struct device_node *state_node, unsigned int idx,
 	for (cpu = cpumask_next(cpumask_first(cpumask), cpumask);
 	     cpu < nr_cpu_ids; cpu = cpumask_next(cpu, cpumask)) {
 		cpu_node = of_cpu_device_node_get(cpu);
-		curr_state_node = of_parse_phandle(cpu_node, "cpu-idle-states",
-						   idx);
+		curr_state_node = get_state_node(cpu_node, idx);
 		if (state_node != curr_state_node)
 			valid = false;
 

@@ -176,7 +208,7 @@ int dt_init_idle_driver(struct cpuidle_driver *drv,
 	cpu_node = of_cpu_device_node_get(cpumask_first(cpumask));
 
 	for (i = 0; ; i++) {
-		state_node = of_parse_phandle(cpu_node, "cpu-idle-states", i);
+		state_node = get_state_node(cpu_node, i);
 		if (!state_node)
 			break;
 

index 6c60a50..ec922b8 100644
@@ -53,6 +53,7 @@
  * require cooperation with a Trusted OS driver.
  */
 static int resident_cpu = -1;
+static bool psci_has_osi;
 
 bool psci_tos_resident_on(int cpu)
 {

@@ -558,9 +559,8 @@ static int __init psci_0_2_init(struct device_node *np)
 	int err;
 
 	err = get_set_conduit_method(np);
-
 	if (err)
-		goto out_put_node;
+		return err;
 	/*
 	 * Starting with v0.2, the PSCI specification introduced a call
 	 * (PSCI_VERSION) that allows probing the firmware version, so

@@ -568,11 +568,7 @@ static int __init psci_0_2_init(struct device_node *np)
 	 * can be carried out according to the specific version reported
 	 * by firmware
 	 */
-	err = psci_probe();
-
-out_put_node:
-	of_node_put(np);
-	return err;
+	return psci_probe();
 }
 
 /*

@@ -584,9 +580,8 @@ static int __init psci_0_1_init(struct device_node *np)
 	int err;
 
 	err = get_set_conduit_method(np);
-
 	if (err)
-		goto out_put_node;
+		return err;
 
 	pr_info("Using PSCI v0.1 Function IDs from DT\n");
 

@@ -610,15 +605,31 @@ static int __init psci_0_1_init(struct device_node *np)
 		psci_ops.migrate = psci_migrate;
 	}
 
-out_put_node:
-	of_node_put(np);
 	return err;
 }
 
+static int __init psci_1_0_init(struct device_node *np)
+{
+	int ret;
+
+	ret = psci_0_2_init(np);
+	if (ret)
+		return ret;
+
+	/* Check if PSCI OSI mode is available */
+	ret = psci_features(psci_function_id[PSCI_FN_CPU_SUSPEND]);
+	if (ret & PSCI_1_0_OS_INITIATED) {
+		if (!psci_features(PSCI_1_0_FN_SET_SUSPEND_MODE))
+			psci_has_osi = true;
+	}
+
+	return 0;
+}
+
 static const struct of_device_id psci_of_match[] __initconst = {
 	{ .compatible = "arm,psci",	.data = psci_0_1_init},
 	{ .compatible = "arm,psci-0.2",	.data = psci_0_2_init},
-	{ .compatible = "arm,psci-1.0",	.data = psci_0_2_init},
+	{ .compatible = "arm,psci-1.0",	.data = psci_1_0_init},
 	{},
 };
 

@@ -627,6 +638,7 @@ int __init psci_dt_init(void)
 	struct device_node *np;
 	const struct of_device_id *matched_np;
 	psci_initcall_t init_fn;
+	int ret;
 
 	np = of_find_matching_node_and_match(NULL, psci_of_match, &matched_np);
 

@@ -634,7 +646,11 @@ int __init psci_dt_init(void)
 		return -ENODEV;
 
 	init_fn = (psci_initcall_t)matched_np->data;
-	return init_fn(np);
+	ret = init_fn(np);
+
+	of_node_put(np);
+
+	return ret;
 }
 
 #ifdef CONFIG_ACPI

index 3d7a0fc..7dd778e 100644
@@ -48,6 +48,7 @@
 
 #define PSCI_1_0_FN_PSCI_FEATURES		PSCI_0_2_FN(10)
 #define PSCI_1_0_FN_SYSTEM_SUSPEND		PSCI_0_2_FN(14)
+#define PSCI_1_0_FN_SET_SUSPEND_MODE		PSCI_0_2_FN(15)
 
 #define PSCI_1_0_FN64_SYSTEM_SUSPEND		PSCI_0_2_FN64(14)
 

@@ -93,6 +94,10 @@
 #define PSCI_1_0_FEATURES_CPU_SUSPEND_PF_MASK	\
 			(0x1 << PSCI_1_0_FEATURES_CPU_SUSPEND_PF_SHIFT)
 
+#define PSCI_1_0_OS_INITIATED			BIT(0)
+#define PSCI_1_0_SUSPEND_MODE_PC		0
+#define PSCI_1_0_SUSPEND_MODE_OSI		1
+
 /* PSCI return values (inclusive of all PSCI versions) */
 #define PSCI_RET_SUCCESS			0
 #define PSCI_RET_NOT_SUPPORTED			-1

index ec922b8..18ae62d 100644
@@ -16,6 +16,7 @@
 #include <linux/acpi.h>
 #include <linux/arm-smccc.h>
 #include <linux/cpuidle.h>
+#include <linux/cpu_domains.h>
 #include <linux/errno.h>
 #include <linux/linkage.h>
 #include <linux/of.h>

@@ -54,6 +55,18 @@
  */
 static int resident_cpu = -1;
 static bool psci_has_osi;
+static bool psci_has_osi_pd;
+static DEFINE_PER_CPU(u32, cluster_state_id);
+
+static inline u32 psci_get_composite_state_id(u32 cpu_state)
+{
+	return cpu_state | this_cpu_read(cluster_state_id);
+}
+
+static inline void psci_reset_composite_state_id(void)
+{
+	this_cpu_write(cluster_state_id, 0);
+}
 
 bool psci_tos_resident_on(int cpu)
 {

@@ -180,6 +193,8 @@ static int psci_cpu_on(unsigned long cpuid, unsigned long entry_point)
 
 	fn = psci_function_id[PSCI_FN_CPU_ON];
 	err = invoke_psci_fn(fn, cpuid, entry_point, 0);
+	/* Reset CPU cluster states */
+	psci_reset_composite_state_id();
 	return psci_to_linux_errno(err);
 }
 

@@ -251,6 +266,27 @@ static int __init psci_features(u32 psci_func_id)
 
 #ifdef CONFIG_CPU_IDLE
 static DEFINE_PER_CPU_READ_MOSTLY(u32 *, psci_power_state);
+static bool psci_suspend_mode_is_osi;
+
+static int psci_set_suspend_mode_osi(bool enable)
+{
+	int ret;
+	int mode;
+
+	if (enable && !psci_has_osi)
+		return -ENODEV;
+
+	if (enable == psci_suspend_mode_is_osi)
+		return 0;
+
+	mode = enable ? PSCI_1_0_SUSPEND_MODE_OSI : PSCI_1_0_SUSPEND_MODE_PC;
+	ret = invoke_psci_fn(PSCI_1_0_FN_SET_SUSPEND_MODE,
+			     mode, 0, 0);
+	if (!ret)
+		psci_suspend_mode_is_osi = enable;
+
+	return psci_to_linux_errno(ret);
+}
 
 static int psci_dt_cpu_init_idle(struct device_node *cpu_node, int cpu)
 {

@@ -353,6 +389,39 @@ static int __maybe_unused psci_acpi_cpu_init_idle(unsigned int cpu)
 }
 #endif
 
+static int psci_pd_populate_state_data(struct device_node *np, u32 *param)
+{
+	return of_property_read_u32(np, "arm,psci-suspend-param", param);
+}
+
+static int psci_pd_power_off(u32 idx, u32 param, const struct cpumask *mask)
+{
+	__this_cpu_add(cluster_state_id, param);
+	return 0;
+}
+
+const struct cpu_pd_ops psci_pd_ops = {
+	.populate_state_data = psci_pd_populate_state_data,
+	.power_off = psci_pd_power_off,
+};
+
+static int psci_cpu_osi_pd_init(int cpu)
+{
+	int ret;
+
+	if (!psci_has_osi_pd)
+		return 0;
+
+	ret = of_setup_cpu_pd_single(cpu, &psci_pd_ops);
+	if (!ret) {
+		ret = psci_set_suspend_mode_osi(true);
+		if (ret)
+			pr_warn("CPU%d: Error setting PSCI OSI mode\n", cpu);
+	}
+
+	return ret;
+}
+
 int psci_cpu_init_idle(unsigned int cpu)
 {
 	struct device_node *cpu_node;

@@ -368,6 +437,10 @@ int psci_cpu_init_idle(unsigned int cpu)
 	if (!acpi_disabled)
 		return psci_acpi_cpu_init_idle(cpu);
 
+	ret = psci_cpu_osi_pd_init(cpu);
+	if (ret)
+		return ret;
+
 	cpu_node = of_get_cpu_node(cpu, NULL);
 	if (!cpu_node)
 		return -ENODEV;

@@ -382,15 +455,17 @@ int psci_cpu_init_idle(unsigned int cpu)
 static int psci_suspend_finisher(unsigned long index)
 {
 	u32 *state = __this_cpu_read(psci_power_state);
+	u32 ext_state = psci_get_composite_state_id(state[index - 1]);
 
-	return psci_ops.cpu_suspend(state[index - 1],
-				    virt_to_phys(cpu_resume));
+	return psci_ops.cpu_suspend(ext_state, virt_to_phys(cpu_resume));
 }
 
 int psci_cpu_suspend_enter(unsigned long index)
 {
 	int ret;
 	u32 *state = __this_cpu_read(psci_power_state);
+	u32 ext_state = psci_get_composite_state_id(state[index - 1]);
+
 	/*
 	 * idle state index 0 corresponds to wfi, should never be called
 	 * from the cpu_suspend operations

@@ -399,10 +474,16 @@ int psci_cpu_suspend_enter(unsigned long index)
 		return -EINVAL;
 
 	if (!psci_power_state_loses_context(state[index - 1]))
-		ret = psci_ops.cpu_suspend(state[index - 1], 0);
+		ret = psci_ops.cpu_suspend(ext_state, 0);
 	else
 		ret = cpu_suspend(index, psci_suspend_finisher);
 
+	/*
+	 * Clear the CPU's cluster states, we start afresh after coming
+	 * out of idle.
+	 */
+	psci_reset_composite_state_id();
+
 	return ret;
 }
 

@@ -610,6 +691,7 @@ static int __init psci_0_1_init(struct device_node *np)
 
 static int __init psci_1_0_init(struct device_node *np)
 {
+	struct device_node *dn;
 	int ret;
 
 	ret = psci_0_2_init(np);

@@ -621,6 +703,11 @@ static int __init psci_1_0_init(struct device_node *np)
 	if (ret & PSCI_1_0_OS_INITIATED) {
 		if (!psci_features(PSCI_1_0_FN_SET_SUSPEND_MODE))
 			psci_has_osi = true;
+		/* Check if we power domains defined in the PSCI node */
+		dn = of_find_node_with_property(np, "#power-domain-cells");
+		if (dn)
+			psci_has_osi_pd = true;
+		of_node_put(dn);
 	}
 
 	return 0;

index 18ae62d..190d3a7 100644
@@ -292,15 +292,24 @@ static int psci_dt_cpu_init_idle(struct device_node *cpu_node, int cpu)
 {
 	int i, ret, count = 0;
 	u32 *psci_states;
-	struct device_node *state_node;
-
-	/* Count idle states */
-	while ((state_node = of_parse_phandle(cpu_node, "cpu-idle-states",
-					      count))) {
-		count++;
-		of_node_put(state_node);
+	struct device_node *state_node, *dn;
+	struct of_phandle_args args;
+	bool cpu_has_domain = false;
+	const char *property;
+
+	/* Is there a domain provider for this CPU? */
+	ret = of_parse_phandle_with_args(cpu_node, "power-domains",
+					"#power-domain-cells", 0, &args);
+	if (!ret) {
+		dn = args.np;
+		ret = of_count_phandle_with_args(dn, "domain-idle-states",
+								NULL);
+		cpu_has_domain = (ret > 0);
 	}
 
+	count = (cpu_has_domain) ? ret :
+		of_count_phandle_with_args(cpu_node, "cpu-idle-states", NULL);
+
 	if (!count)
 		return -ENODEV;
 

@@ -308,10 +317,17 @@ static int psci_dt_cpu_init_idle(struct device_node *cpu_node, int cpu)
 	if (!psci_states)
 		return -ENOMEM;
 
+	if (cpu_has_domain) {
+		property = "domain-idle-states";
+	} else {
+		property = "cpu-idle-states";
+		dn = cpu_node;
+	}
+
 	for (i = 0; i < count; i++) {
 		u32 state;
 
-		state_node = of_parse_phandle(cpu_node, "cpu-idle-states", i);
+		state_node = of_parse_phandle(dn, property, i);
 
 		ret = of_property_read_u32(state_node,
 					   "arm,psci-suspend-param",