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

9 files changed, 168 insertions, 294 deletions

diff --git a/common/psci/psci_afflvl_off.c b/common/psci/psci_afflvl_off.c
index 0e78aa8..1d8f291 100644
--- a/common/psci/psci_afflvl_off.c
+++ b/common/psci/psci_afflvl_off.c
@@ -51,6 +51,9 @@ static int psci_afflvl0_off(unsigned long mpidr, aff_map_node *cpu_node)
 
 	assert(cpu_node->level == MPIDR_AFFLVL0);
 
+	/* State management: mark this cpu as turned off */
+	psci_set_state(cpu_node, PSCI_STATE_OFF);
+
 	/*
 	 * Generic management: Get the index for clearing any
 	 * lingering re-entry information
@@ -85,7 +88,7 @@ static int psci_afflvl0_off(unsigned long mpidr, aff_map_node *cpu_node)
 	if (psci_plat_pm_ops->affinst_off) {
 
 		/* Get the current physical state of this cpu */
-		plat_state = psci_get_aff_phys_state(cpu_node);
+		plat_state = psci_get_phys_state(cpu_node);
 		rc = psci_plat_pm_ops->affinst_off(mpidr,
 						   cpu_node->level,
 						   plat_state);
@@ -102,11 +105,14 @@ static int psci_afflvl1_off(unsigned long mpidr, aff_map_node *cluster_node)
 	/* 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_OFF);
+
 	/*
 	 * 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
@@ -136,11 +142,14 @@ static int psci_afflvl2_off(unsigned long mpidr, aff_map_node *system_node)
 	/* Cannot go beyond this level */
 	assert(system_node->level == MPIDR_AFFLVL2);
 
+	/* State management: Decrement the system reference count */
+	psci_set_state(system_node, PSCI_STATE_OFF);
+
 	/*
 	 * 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);
 
 	/* No arch. and generic bookeeping to do here currently */
 
@@ -219,7 +228,6 @@ int psci_afflvl_off(unsigned long mpidr,
 		    int end_afflvl)
 {
 	int rc = PSCI_E_SUCCESS;
-	unsigned int prev_state;
 	mpidr_aff_map_nodes mpidr_nodes;
 
 	mpidr &= MPIDR_AFFINITY_MASK;;
@@ -247,21 +255,6 @@ int psci_afflvl_off(unsigned long mpidr,
 				  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_OFF);
-
 	/* Perform generic, architecture and platform specific handling */
 	rc = psci_call_off_handlers(mpidr_nodes,
 				    start_afflvl,
@@ -269,21 +262,6 @@ int psci_afflvl_off(unsigned long mpidr,
 				    mpidr);
 
 	/*
-	 * 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.
 	 */
diff --git a/common/psci/psci_afflvl_on.c b/common/psci/psci_afflvl_on.c
index c1a2e99..83d47d5 100644
--- a/common/psci/psci_afflvl_on.c
+++ b/common/psci/psci_afflvl_on.c
@@ -46,12 +46,12 @@ typedef int (*afflvl_on_handler)(unsigned long,
  * This function checks whether a cpu which has been requested to be turned on
  * is OFF to begin with.
  ******************************************************************************/
-static int cpu_on_validate_state(unsigned int state)
+static int cpu_on_validate_state(aff_map_node *node)
 {
 	unsigned int psci_state;
 
 	/* Get the raw psci state */
-	psci_state = psci_get_state(state);
+	psci_state = psci_get_state(node);
 
 	if (psci_state == PSCI_STATE_ON || psci_state == PSCI_STATE_SUSPEND)
 		return PSCI_E_ALREADY_ON;
@@ -84,7 +84,7 @@ static int psci_afflvl0_on(unsigned long target_cpu,
 	 * Generic management: Ensure that the cpu is off to be
 	 * turned on
 	 */
-	rc = cpu_on_validate_state(cpu_node->state);
+	rc = cpu_on_validate_state(cpu_node);
 	if (rc != PSCI_E_SUCCESS)
 		return rc;
 
@@ -101,6 +101,9 @@ static int psci_afflvl0_on(unsigned long target_cpu,
 	/* Set the secure world (EL3) re-entry point after BL1 */
 	psci_entrypoint = (unsigned long) psci_aff_on_finish_entry;
 
+	/* State management: Set this cpu's state as ON PENDING */
+	psci_set_state(cpu_node, PSCI_STATE_ON_PENDING);
+
 	/*
 	 * Plat. management: Give the platform the current state
 	 * of the target cpu to allow it to perform the necessary
@@ -109,7 +112,7 @@ static int psci_afflvl0_on(unsigned long target_cpu,
 	if (psci_plat_pm_ops->affinst_on) {
 
 		/* Get the current physical state of this cpu */
-		plat_state = psci_get_aff_phys_state(cpu_node);
+		plat_state = psci_get_phys_state(cpu_node);
 		rc = psci_plat_pm_ops->affinst_on(target_cpu,
 						  psci_entrypoint,
 						  ns_entrypoint,
@@ -141,13 +144,15 @@ static int psci_afflvl1_on(unsigned long target_cpu,
 	 * management required
 	 */
 
+	/* State management: Is not required while turning a cluster on */
+
 	/*
 	 * Plat. management: Give the platform the current state
 	 * of the target cpu to allow it to perform the necessary
 	 * steps to power on.
 	 */
 	if (psci_plat_pm_ops->affinst_on) {
-		plat_state = psci_get_aff_phys_state(cluster_node);
+		plat_state = psci_get_phys_state(cluster_node);
 		psci_entrypoint = (unsigned long) psci_aff_on_finish_entry;
 		rc = psci_plat_pm_ops->affinst_on(target_cpu,
 						  psci_entrypoint,
@@ -181,13 +186,15 @@ static int psci_afflvl2_on(unsigned long target_cpu,
 	 * required
 	 */
 
+	/* State management: Is not required while turning a system on */
+
 	/*
 	 * Plat. management: Give the platform the current state
 	 * of the target cpu to allow it to perform the necessary
 	 * steps to power on.
 	 */
 	if (psci_plat_pm_ops->affinst_on) {
-		plat_state = psci_get_aff_phys_state(system_node);
+		plat_state = psci_get_phys_state(system_node);
 		psci_entrypoint = (unsigned long) psci_aff_on_finish_entry;
 		rc = psci_plat_pm_ops->affinst_on(target_cpu,
 						  psci_entrypoint,
@@ -299,19 +306,7 @@ int psci_afflvl_on(unsigned long target_cpu,
 				   target_cpu,
 				   entrypoint,
 				   context_id);
-	if (rc != PSCI_E_SUCCESS)
-		goto exit;
-
-	/*
-	 * State management: Update the state of each affinity instance
-	 * between the start and end affinity levels
-	 */
-	psci_change_state(target_cpu_nodes,
-			  start_afflvl,
-			  end_afflvl,
-			  PSCI_STATE_ON_PENDING);
 
-exit:
 	/*
 	 * This loop releases the lock corresponding to each affinity level
 	 * in the reverse order to which they were acquired.
@@ -336,7 +331,7 @@ static unsigned int psci_afflvl0_on_finish(unsigned long mpidr,
 	assert(cpu_node->level == MPIDR_AFFLVL0);
 
 	/* Ensure we have been explicitly woken up by another cpu */
-	state = psci_get_state(cpu_node->state);
+	state = psci_get_state(cpu_node);
 	assert(state == PSCI_STATE_ON_PENDING);
 
 	/*
@@ -348,7 +343,7 @@ static unsigned int psci_afflvl0_on_finish(unsigned long mpidr,
 	if (psci_plat_pm_ops->affinst_on_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_on_finish(mpidr,
 							 cpu_node->level,
 							 plat_state);
@@ -377,6 +372,9 @@ static unsigned int psci_afflvl0_on_finish(unsigned long mpidr,
 	index = cpu_node->data;
 	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);
 
@@ -401,13 +399,16 @@ static unsigned int psci_afflvl1_on_finish(unsigned long mpidr,
 	if (psci_plat_pm_ops->affinst_on_finish) {
 
 		/* Get the physical state of this cluster */
-		plat_state = psci_get_aff_phys_state(cluster_node);
+		plat_state = psci_get_phys_state(cluster_node);
 		rc = psci_plat_pm_ops->affinst_on_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;
 }
 
@@ -436,13 +437,16 @@ static unsigned int psci_afflvl2_on_finish(unsigned long mpidr,
 	if (psci_plat_pm_ops->affinst_on_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_on_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;
 }
 
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 @@ static int psci_afflvl0_suspend(unsigned long mpidr,
 	/* 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 @@ static int psci_afflvl0_suspend(unsigned long mpidr,
 	 * 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 @@ static int psci_afflvl1_suspend(unsigned long mpidr,
 	/* 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 @@ static int psci_afflvl2_suspend(unsigned long mpidr,
 	/* 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 psci_afflvl_suspend(unsigned long mpidr,
 			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 @@ int psci_afflvl_suspend(unsigned long mpidr,
 				  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 @@ int psci_afflvl_suspend(unsigned long mpidr,
 					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 @@ static unsigned int psci_afflvl0_suspend_finish(unsigned long mpidr,
 	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 @@ static unsigned int psci_afflvl0_suspend_finish(unsigned long mpidr,
 	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 @@ static unsigned int psci_afflvl0_suspend_finish(unsigned long mpidr,
 	 */
 	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 @@ static unsigned int psci_afflvl1_suspend_finish(unsigned long mpidr,
 	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 @@ static unsigned int psci_afflvl2_suspend_finish(unsigned long mpidr,
 	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;
 }
 
diff --git a/common/psci/psci_common.c b/common/psci/psci_common.c
index e9028cc..e6be2f8 100644
--- a/common/psci/psci_common.c
+++ b/common/psci/psci_common.c
@@ -93,7 +93,7 @@ int get_power_on_target_afflvl(unsigned long mpidr)
 	 * Call the handler in the suspend code if this cpu had been suspended.
 	 * Any other state is invalid.
 	 */
-	state = psci_get_state(node->state);
+	state = psci_get_state(node);
 	if (state == PSCI_STATE_ON_PENDING)
 		return get_max_afflvl();
 
@@ -214,164 +214,6 @@ int psci_validate_mpidr(unsigned long mpidr, int level)
 }
 
 /*******************************************************************************
- * Simple routine to determine the first affinity level instance that is present
- * between the start and end affinity levels. This helps to skip handling of
- * absent affinity levels while performing psci operations.
- * The start level can be > or <= to the end level depending upon whether this
- * routine is expected to search top down or bottom up.
- ******************************************************************************/
-int psci_get_first_present_afflvl(unsigned long mpidr,
-				  int start_afflvl,
-				  int end_afflvl,
-				  aff_map_node **node)
-{
-	int level;
-
-	/* Check whether we have to search up or down */
-	if (start_afflvl <= end_afflvl) {
-		for (level = start_afflvl; level <= end_afflvl; level++) {
-			*node = psci_get_aff_map_node(mpidr, level);
-			if (*node && ((*node)->state & PSCI_AFF_PRESENT))
-				break;
-		}
-	} else {
-		for (level = start_afflvl; level >= end_afflvl; level--) {
-			*node = psci_get_aff_map_node(mpidr, level);
-			if (*node && ((*node)->state & PSCI_AFF_PRESENT))
-				break;
-		}
-	}
-
-	return level;
-}
-
-/*******************************************************************************
- * Iteratively change the affinity state between the current and target affinity
- * levels. The target state matters only if we are starting from affinity level
- * 0 i.e. a cpu otherwise the state depends upon the state of the lower affinity
- * levels.
- ******************************************************************************/
-int psci_change_state(mpidr_aff_map_nodes mpidr_nodes,
-		      int start_afflvl,
-		      int end_afflvl,
-		      unsigned int tgt_state)
-{
-	int rc = PSCI_E_SUCCESS, level;
-	unsigned int state;
-	aff_map_node *node;
-
-	/*
-	 * Get a temp pointer to the node. It is not possible that affinity
-	 * level 0 is missing. Simply ignore higher missing levels.
-	 */
-	for (level = start_afflvl; level <= end_afflvl; level++) {
-
-		node = mpidr_nodes[level];
-		if (level == MPIDR_AFFLVL0) {
-			assert(node);
-			psci_set_state(node->state, tgt_state);
-		} else {
-			if (node == NULL)
-				continue;
-			state = psci_calculate_affinity_state(node);
-			psci_set_state(node->state, state);
-		}
-	}
-
-	/* If all went well then the cpu should be in the target state */
-	if (start_afflvl == MPIDR_AFFLVL0) {
-		node = mpidr_nodes[MPIDR_AFFLVL0];
-		state = psci_get_state(node->state);
-		assert(tgt_state == state);
-	}
-
-	return rc;
-}
-
-/*******************************************************************************
- * This routine does the heavy lifting for psci_change_state(). It examines the
- * state of each affinity instance at the next lower affinity level and decides
- * its final state accordingly. If a lower affinity instance is ON then the
- * higher affinity instance is ON. If all the lower affinity instances are OFF
- * then the higher affinity instance is OFF. If atleast one lower affinity
- * instance is SUSPENDED then the higher affinity instance is SUSPENDED. If only
- * a single lower affinity instance is ON_PENDING then the higher affinity
- * instance in ON_PENDING as well.
- ******************************************************************************/
-unsigned int psci_calculate_affinity_state(aff_map_node *aff_node)
-{
-	int ctr;
-	unsigned int aff_count, hi_aff_state;
-	unsigned long tempidr;
-	aff_map_node *lo_aff_node;
-
-	/* Cannot calculate lowest affinity state. It is simply assigned */
-	assert(aff_node->level > MPIDR_AFFLVL0);
-
-	/*
-	 * Find the number of affinity instances at level X-1 e.g. number of
-	 * cpus in a cluster. The level X state depends upon the state of each
-	 * instance at level X-1
-	 */
-	hi_aff_state = PSCI_STATE_OFF;
-	aff_count = plat_get_aff_count(aff_node->level - 1, aff_node->mpidr);
-	for (ctr = 0; ctr < aff_count; ctr++) {
-
-		/*
-		 * Create a mpidr for each lower affinity level (X-1). Use their
-		 * states to influence the higher affinity state (X).
-		 */
-		tempidr = mpidr_set_aff_inst(aff_node->mpidr,
-					     ctr,
-					     aff_node->level - 1);
-		lo_aff_node = psci_get_aff_map_node(tempidr,
-						    aff_node->level - 1);
-		assert(lo_aff_node);
-
-		/* Continue only if the cpu exists within the cluster */
-		if (!(lo_aff_node->state & PSCI_AFF_PRESENT))
-			continue;
-
-		switch (psci_get_state(lo_aff_node->state)) {
-
-		/*
-		 * If any lower affinity is on within the cluster, then
-		 * the higher affinity is on.
-		 */
-		case PSCI_STATE_ON:
-			return PSCI_STATE_ON;
-
-		/*
-		 * At least one X-1 needs to be suspended for X to be suspended
-		 * but it is effectively on for the affinity_info call.
-		 * SUSPEND > ON_PENDING > OFF.
-		 */
-		case PSCI_STATE_SUSPEND:
-			hi_aff_state = PSCI_STATE_SUSPEND;
-			continue;
-
-		/*
-		 * Atleast one X-1 needs to be on_pending & the rest off for X
-		 * to be on_pending. ON_PENDING > OFF.
-		 */
-		case PSCI_STATE_ON_PENDING:
-			if (hi_aff_state != PSCI_STATE_SUSPEND)
-				hi_aff_state = PSCI_STATE_ON_PENDING;
-			continue;
-
-		/* Higher affinity is off if all lower affinities are off. */
-		case PSCI_STATE_OFF:
-			continue;
-
-		default:
-			assert(0);
-		}
-	}
-
-	return hi_aff_state;
-}
-
-/*******************************************************************************
  * This function retrieves all the stashed information needed to correctly
  * resume a cpu's execution in the non-secure state after it has been physically
  * powered on i.e. turned ON or resumed from SUSPEND
@@ -518,23 +360,83 @@ int psci_set_ns_entry_info(unsigned int index,
 }
 
 /*******************************************************************************
+ * This function takes a pointer to an affinity node in the topology tree and
+ * returns its state. State of a non-leaf node needs to be calculated.
+ ******************************************************************************/
+unsigned short psci_get_state(aff_map_node *node)
+{
+	assert(node->level >= MPIDR_AFFLVL0 && node->level <= MPIDR_MAX_AFFLVL);
+
+	/* A cpu node just contains the state which can be directly returned */
+	if (node->level == MPIDR_AFFLVL0)
+		return (node->state >> PSCI_STATE_SHIFT) & PSCI_STATE_MASK;
+
+	/*
+	 * For an affinity level higher than a cpu, the state has to be
+	 * calculated. It depends upon the value of the reference count
+	 * which is managed by each node at the next lower affinity level
+	 * e.g. for a cluster, each cpu increments/decrements the reference
+	 * count. If the reference count is 0 then the affinity level is
+	 * OFF else ON.
+	 */
+	if (node->ref_count)
+		return PSCI_STATE_ON;
+	else
+		return PSCI_STATE_OFF;
+}
+
+/*******************************************************************************
+ * This function takes a pointer to an affinity node in the topology tree and
+ * a target state. State of a non-leaf node needs to be converted to a reference
+ * count. State of a leaf node can be set directly.
+ ******************************************************************************/
+void psci_set_state(aff_map_node *node, unsigned short state)
+{
+	assert(node->level >= MPIDR_AFFLVL0 && node->level <= MPIDR_MAX_AFFLVL);
+
+	/*
+	 * For an affinity level higher than a cpu, the state is used
+	 * to decide whether the reference count is incremented or
+	 * decremented. Entry into the ON_PENDING state does not have
+	 * effect.
+	 */
+	if (node->level > MPIDR_AFFLVL0) {
+		switch (state) {
+		case PSCI_STATE_ON:
+			node->ref_count++;
+			break;
+		case PSCI_STATE_OFF:
+		case PSCI_STATE_SUSPEND:
+			node->ref_count--;
+			break;
+		case PSCI_STATE_ON_PENDING:
+			/*
+			 * An affinity level higher than a cpu will not undergo
+			 * a state change when it is about to be turned on
+			 */
+			return;
+		default:
+			assert(0);
+		}
+	} else {
+		node->state &= ~(PSCI_STATE_MASK << PSCI_STATE_SHIFT);
+		node->state |= (state & PSCI_STATE_MASK) << PSCI_STATE_SHIFT;
+	}
+}
+
+/*******************************************************************************
  * An affinity level could be on, on_pending, suspended or off. These are the
  * logical states it can be in. Physically either it is off or on. When it is in
  * the state on_pending then it is about to be turned on. It is not possible to
  * tell whether that's actually happenned or not. So we err on the side of
  * caution & treat the affinity level as being turned off.
  ******************************************************************************/
-inline unsigned int psci_get_phys_state(unsigned int aff_state)
+unsigned short psci_get_phys_state(aff_map_node *node)
 {
-	return (aff_state != PSCI_STATE_ON ? PSCI_STATE_OFF : PSCI_STATE_ON);
-}
-
-unsigned int psci_get_aff_phys_state(aff_map_node *aff_node)
-{
-	unsigned int aff_state;
+	unsigned int state;
 
-	aff_state = psci_get_state(aff_node->state);
-	return psci_get_phys_state(aff_state);
+	state = psci_get_state(node);
+	return get_phys_state(state);
 }
 
 /*******************************************************************************
@@ -630,15 +532,6 @@ void psci_afflvl_power_on_finish(unsigned long mpidr,
 	assert (rc == PSCI_E_SUCCESS);
 
 	/*
-	 * 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_ON);
-
-	/*
 	 * This loop releases the lock corresponding to each affinity level
 	 * in the reverse order to which they were acquired.
 	 */
diff --git a/common/psci/psci_main.c b/common/psci/psci_main.c
index 2edf77b..a70a21a 100644
--- a/common/psci/psci_main.c
+++ b/common/psci/psci_main.c
@@ -142,13 +142,18 @@ int psci_affinity_info(unsigned long target_affinity,
 	unsigned int aff_state;
 	aff_map_node *node;
 
-	if (lowest_affinity_level > get_max_afflvl()) {
-		goto exit;
-	}
+	if (lowest_affinity_level > get_max_afflvl())
+		return rc;
 
 	node = psci_get_aff_map_node(target_affinity, lowest_affinity_level);
 	if (node && (node->state & PSCI_AFF_PRESENT)) {
-		aff_state = psci_get_state(node->state);
+
+		/*
+		 * TODO: For affinity levels higher than 0 i.e. cpu, the
+		 * state will always be either ON or OFF. Need to investigate
+		 * how critical is it to support ON_PENDING here.
+		 */
+		aff_state = psci_get_state(node);
 
 		/* A suspended cpu is available & on for the OS */
 		if (aff_state == PSCI_STATE_SUSPEND) {
@@ -157,7 +162,7 @@ int psci_affinity_info(unsigned long target_affinity,
 
 		rc = aff_state;
 	}
-exit:
+
 	return rc;
 }
 
diff --git a/common/psci/psci_private.h b/common/psci/psci_private.h
index 1cc7104..9b5c552 100644
--- a/common/psci/psci_private.h
+++ b/common/psci/psci_private.h
@@ -57,8 +57,9 @@ typedef struct {
  ******************************************************************************/
 typedef struct {
 	unsigned long mpidr;
+	unsigned short ref_count;
 	unsigned char state;
-	char level;
+	unsigned char level;
 	unsigned int data;
 	bakery_lock lock;
 } aff_map_node;
@@ -100,12 +101,11 @@ extern afflvl_power_on_finisher psci_afflvl_sus_finish_handlers[];
  ******************************************************************************/
 /* Private exported functions from psci_common.c */
 extern int get_max_afflvl(void);
-extern unsigned int psci_get_phys_state(unsigned int);
-extern unsigned int psci_get_aff_phys_state(aff_map_node *);
-extern unsigned int psci_calculate_affinity_state(aff_map_node *);
+extern unsigned short psci_get_state(aff_map_node *node);
+extern unsigned short psci_get_phys_state(aff_map_node *node);
+extern void psci_set_state(aff_map_node *node, unsigned short state);
 extern void psci_get_ns_entry_info(unsigned int index);
 extern unsigned long mpidr_set_aff_inst(unsigned long,unsigned char, int);
-extern int psci_change_state(mpidr_aff_map_nodes, int, int, unsigned int);
 extern int psci_validate_mpidr(unsigned long, int);
 extern int get_power_on_target_afflvl(unsigned long mpidr);
 extern void psci_afflvl_power_on_finish(unsigned long,
@@ -115,9 +115,6 @@ extern void psci_afflvl_power_on_finish(unsigned long,
 extern int psci_set_ns_entry_info(unsigned int index,
 				  unsigned long entrypoint,
 				  unsigned long context_id);
-extern int psci_get_first_present_afflvl(unsigned long,
-					 int, int,
-					 aff_map_node **);
 extern int psci_check_afflvl_range(int start_afflvl, int end_afflvl);
 extern void psci_acquire_afflvl_locks(unsigned long mpidr,
 				      int start_afflvl,
diff --git a/common/psci/psci_setup.c b/common/psci/psci_setup.c
index e01789f..c0d29f2 100644
--- a/common/psci/psci_setup.c
+++ b/common/psci/psci_setup.c
@@ -157,11 +157,16 @@ static void psci_init_aff_map_node(unsigned long mpidr,
 	 */
 	state = plat_get_aff_state(level, mpidr);
 	psci_aff_map[idx].state = state;
-	if (state & PSCI_AFF_PRESENT) {
-		psci_set_state(psci_aff_map[idx].state, PSCI_STATE_OFF);
-	}
 
 	if (level == MPIDR_AFFLVL0) {
+
+		/*
+		 * Mark the cpu as OFF. Higher affinity level reference counts
+		 * have already been memset to 0
+		 */
+		if (state & PSCI_AFF_PRESENT)
+			psci_set_state(&psci_aff_map[idx], PSCI_STATE_OFF);
+
 		/* Ensure that we have not overflowed the psci_ns_einfo array */
 		assert(psci_ns_einfo_idx < PSCI_NUM_AFFS);
 
@@ -299,15 +304,14 @@ void psci_setup(unsigned long mpidr)
 	 * this is the primary cpu.
 	 */
 	mpidr &= MPIDR_AFFINITY_MASK;
-	for (afflvl = max_afflvl; afflvl >= MPIDR_AFFLVL0; afflvl--) {
+	for (afflvl = MPIDR_AFFLVL0; afflvl <= max_afflvl; afflvl++) {
 
 		node = psci_get_aff_map_node(mpidr, afflvl);
 		assert(node);
 
 		/* Mark each present node as ON. */
-		if (node->state & PSCI_AFF_PRESENT) {
-			psci_set_state(node->state, PSCI_STATE_ON);
-		}
+		if (node->state & PSCI_AFF_PRESENT)
+			psci_set_state(node, PSCI_STATE_ON);
 	}
 
 	rc = platform_setup_pm(&psci_plat_pm_ops);
diff --git a/docs/change-log.md b/docs/change-log.md
index 50e1111..1ad631e 100644
--- a/docs/change-log.md
+++ b/docs/change-log.md
@@ -104,6 +104,11 @@ Detailed changes since last release
     automatically detected by the make file when they are added to the plat
     directory.
 
+*   An issue in the PSCI implementation has been fixed which could result in the
+    power down of an affinity instance at level X even though at least one
+    affinity instance at level X - 1 does not allow this.
+
+
 ARM Trusted Firmware - version 0.2
 ==================================
 
diff --git a/include/psci.h b/include/psci.h
index 6784612..5da78ee 100644
--- a/include/psci.h
+++ b/include/psci.h
@@ -100,10 +100,7 @@
  * could in one of the 4 further defined states.
  ******************************************************************************/
 #define PSCI_STATE_SHIFT	1
-#define PSCI_STATE_MASK		0x7
-#define psci_get_state(x)	(x >> PSCI_STATE_SHIFT) & PSCI_STATE_MASK
-#define psci_set_state(x,y)	x &= ~(PSCI_STATE_MASK << PSCI_STATE_SHIFT); \
-				x |= (y & PSCI_STATE_MASK) << PSCI_STATE_SHIFT;
+#define PSCI_STATE_MASK		0xff
 
 #define PSCI_AFF_ABSENT		0x0
 #define PSCI_AFF_PRESENT	0x1
@@ -112,6 +109,9 @@
 #define PSCI_STATE_ON_PENDING	0x2
 #define PSCI_STATE_SUSPEND	0x3
 
+#define get_phys_state(x)	(x != PSCI_STATE_ON ? \
+				 PSCI_STATE_OFF : PSCI_STATE_ON)
+
 /* Number of affinity instances whose state this psci imp. can track */
 #define PSCI_NUM_AFFS		32ull