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 @@
 
 	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 @@
 	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 @@
 	/* 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 @@
 	/* 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 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 @@
 				  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 @@
 				    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 @@
  * 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 @@
 	 * 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 @@
 	/* 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 @@
 	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 @@
 	 * 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 @@
 	 * 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,20 +306,8 @@
 				   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 @@
 	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 @@
 	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 @@
 	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 @@
 	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 @@
 	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 @@
 	/* 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;
 }
 
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 @@
 	 * 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 @@
 }
 
 /*******************************************************************************
- * 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 @@
 }
 
 /*******************************************************************************
+ * 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 state;
 
-unsigned int psci_get_aff_phys_state(aff_map_node *aff_node)
-{
-	unsigned int aff_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 @@
 	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 @@
 	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 @@
 
 		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 {
 	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 @@
  ******************************************************************************/
 /* 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 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 @@
 	 */
 	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 @@
 	 * 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);