psci: fix affinity level upgrade issue

The psci implementation does not track target affinity level requests
specified during cpu_suspend calls correctly as per the following
example.

1. cpu0.cluster0 calls cpu_suspend with the target affinity level as 0
2. Only the cpu0.cluster0 is powered down while cluster0 remains
   powered up
3. cpu1.cluster0 calls cpu_off to power itself down to highest
   possible affinity level
4. cluster0 will be powered off even though cpu0.cluster0 does not
   allow cluster shutdown

This patch introduces reference counts at affinity levels > 0 to track
the number of cpus which want an affinity instance at level X to
remain powered up. This instance can be turned off only if its
reference count is 0. Cpus still undergo the normal state transitions
(ON, OFF, ON_PENDING, SUSPEND) but the higher levels can only be
either ON or OFF depending upon their reference count.

The above issue is thus fixed as follows:

1. cluster0's reference count is incremented by two when cpu0 and cpu1
   are initially powered on.

2. cpu0.cluster0 calls cpu_suspend with the target affinity level as
   0. This does not affect the cluster0 reference count.

3. Only the cpu0.cluster0 is powered down while cluster0 remains
   powered up as it has a non-zero reference count.

4. cpu1.cluster0 call cpu_off to power itself down to highest possible
   affinity level. This decrements the cluster0 reference count.

5. cluster0 is still not powered off since its reference count will at
   least be 1 due to the restriction placed by cpu0.

Change-Id: I433dfe82b946f5f6985b1602c2de87800504f7a9
diff --git a/common/psci/psci_afflvl_suspend.c b/common/psci/psci_afflvl_suspend.c
index b2dc732..2abcafb 100644
--- a/common/psci/psci_afflvl_suspend.c
+++ b/common/psci/psci_afflvl_suspend.c
@@ -91,6 +91,9 @@
 	/* Sanity check to safeguard against data corruption */
 	assert(cpu_node->level == MPIDR_AFFLVL0);
 
+	/* State management: mark this cpu as suspended */
+	psci_set_state(cpu_node, PSCI_STATE_SUSPEND);
+
 	/*
 	 * Generic management: Store the re-entry information for the
 	 * non-secure world
@@ -146,7 +149,7 @@
 	 * program the power controller etc.
 	 */
 	if (psci_plat_pm_ops->affinst_suspend) {
-		plat_state = psci_get_aff_phys_state(cpu_node);
+		plat_state = psci_get_phys_state(cpu_node);
 		rc = psci_plat_pm_ops->affinst_suspend(mpidr,
 						       psci_entrypoint,
 						       ns_entrypoint,
@@ -170,11 +173,14 @@
 	/* Sanity check the cluster level */
 	assert(cluster_node->level == MPIDR_AFFLVL1);
 
+	/* State management: Decrement the cluster reference count */
+	psci_set_state(cluster_node, PSCI_STATE_SUSPEND);
+
 	/*
 	 * Keep the physical state of this cluster handy to decide
 	 * what action needs to be taken
 	 */
-	plat_state = psci_get_aff_phys_state(cluster_node);
+	plat_state = psci_get_phys_state(cluster_node);
 
 	/*
 	 * Arch. management: Flush all levels of caches to PoC if the
@@ -221,11 +227,14 @@
 	/* Cannot go beyond this */
 	assert(system_node->level == MPIDR_AFFLVL2);
 
+	/* State management: Decrement the system reference count */
+	psci_set_state(system_node, PSCI_STATE_SUSPEND);
+
 	/*
 	 * Keep the physical state of the system handy to decide what
 	 * action needs to be taken
 	 */
-	plat_state = psci_get_aff_phys_state(system_node);
+	plat_state = psci_get_phys_state(system_node);
 
 	/*
 	 * Plat. Management : Allow the platform to do its bookeeping
@@ -324,7 +333,6 @@
 			int end_afflvl)
 {
 	int rc = PSCI_E_SUCCESS;
-	unsigned int prev_state;
 	mpidr_aff_map_nodes mpidr_nodes;
 
 	mpidr &= MPIDR_AFFINITY_MASK;
@@ -352,20 +360,6 @@
 				  end_afflvl,
 				  mpidr_nodes);
 
-	/*
-	 * Keep the old cpu state handy. It will be used to restore the
-	 * system to its original state in case something goes wrong
-	 */
-	prev_state = psci_get_state(mpidr_nodes[MPIDR_AFFLVL0]->state);
-
-	/*
-	 * State management: Update the state of each affinity instance
-	 * between the start and end affinity levels
-	 */
-	psci_change_state(mpidr_nodes,
-			  start_afflvl,
-			  end_afflvl,
-			  PSCI_STATE_SUSPEND);
 
 	/* Save the affinity level till which this cpu can be powered down */
 	psci_set_suspend_afflvl(mpidr_nodes[MPIDR_AFFLVL0], end_afflvl);
@@ -380,21 +374,6 @@
 					power_state);
 
 	/*
-	 * If an error is returned by a handler then restore the cpu state
-	 * to its original value. If the cpu state is restored then that
-	 * should result in the state of the higher affinity levels to
-	 * get restored as well.
-	 * TODO: We are not undoing any architectural or platform specific
-	 * operations that might have completed before encountering the
-	 * error. The system might not be in a stable state.
-	 */
-	if (rc != PSCI_E_SUCCESS)
-		psci_change_state(mpidr_nodes,
-				  start_afflvl,
-				  end_afflvl,
-				  prev_state);
-
-	/*
 	 * Release the locks corresponding to each affinity level in the
 	 * reverse order to which they were acquired.
 	 */
@@ -418,7 +397,7 @@
 	assert(cpu_node->level == MPIDR_AFFLVL0);
 
 	/* Ensure we have been woken up from a suspended state */
-	state = psci_get_state(cpu_node->state);
+	state = psci_get_state(cpu_node);
 	assert(state == PSCI_STATE_SUSPEND);
 
 	/*
@@ -431,7 +410,7 @@
 	if (psci_plat_pm_ops->affinst_suspend_finish) {
 
 		/* Get the physical state of this cpu */
-		plat_state = psci_get_phys_state(state);
+		plat_state = get_phys_state(state);
 		rc = psci_plat_pm_ops->affinst_suspend_finish(mpidr,
 							      cpu_node->level,
 							      plat_state);
@@ -465,6 +444,9 @@
 	 */
 	psci_get_ns_entry_info(index);
 
+	/* State management: mark this cpu as on */
+	psci_set_state(cpu_node, PSCI_STATE_ON);
+
 	/* Clean caches before re-entering normal world */
 	dcsw_op_louis(DCCSW);
 
@@ -489,13 +471,16 @@
 	if (psci_plat_pm_ops->affinst_suspend_finish) {
 
 		/* Get the physical state of this cpu */
-		plat_state = psci_get_aff_phys_state(cluster_node);
+		plat_state = psci_get_phys_state(cluster_node);
 		rc = psci_plat_pm_ops->affinst_suspend_finish(mpidr,
 							      cluster_node->level,
 							      plat_state);
 		assert(rc == PSCI_E_SUCCESS);
 	}
 
+	/* State management: Increment the cluster reference count */
+	psci_set_state(cluster_node, PSCI_STATE_ON);
+
 	return rc;
 }
 
@@ -524,13 +509,16 @@
 	if (psci_plat_pm_ops->affinst_suspend_finish) {
 
 		/* Get the physical state of the system */
-		plat_state = psci_get_aff_phys_state(system_node);
+		plat_state = psci_get_phys_state(system_node);
 		rc = psci_plat_pm_ops->affinst_suspend_finish(mpidr,
 							      system_node->level,
 							      plat_state);
 		assert(rc == PSCI_E_SUCCESS);
 	}
 
+	/* State management: Increment the system reference count */
+	psci_set_state(system_node, PSCI_STATE_ON);
+
 	return rc;
 }