#include <asm/cpuidle.h>
 
+#include "cpuidle-psci.h"
 #include "dt_idle_states.h"
 
-static DEFINE_PER_CPU_READ_MOSTLY(u32 *, psci_power_state);
+struct psci_cpuidle_data {
+       u32 *psci_states;
+       struct device *dev;
+};
+
+static DEFINE_PER_CPU_READ_MOSTLY(struct psci_cpuidle_data, psci_cpuidle_data);
 
 static int psci_enter_idle_state(struct cpuidle_device *dev,
                                struct cpuidle_driver *drv, int idx)
 {
-       u32 *state = __this_cpu_read(psci_power_state);
+       u32 *state = __this_cpu_read(psci_cpuidle_data.psci_states);
 
        return CPU_PM_CPU_IDLE_ENTER_PARAM(psci_cpu_suspend_enter,
                                           idx, state[idx]);
        int i, ret = 0;
        u32 *psci_states;
        struct device_node *state_node;
+       struct psci_cpuidle_data *data = per_cpu_ptr(&psci_cpuidle_data, cpu);
 
        state_count++; /* Add WFI state too */
        psci_states = kcalloc(state_count, sizeof(*psci_states), GFP_KERNEL);
                goto free_mem;
        }
 
-       /* Idle states parsed correctly, initialize per-cpu pointer */
-       per_cpu(psci_power_state, cpu) = psci_states;
+       /* Currently limit the hierarchical topology to be used in OSI mode. */
+       if (psci_has_osi_support()) {
+               data->dev = psci_dt_attach_cpu(cpu);
+               if (IS_ERR(data->dev)) {
+                       ret = PTR_ERR(data->dev);
+                       goto free_mem;
+               }
+       }
+
+       /* Idle states parsed correctly, store them in the per-cpu struct. */
+       data->psci_states = psci_states;
        return 0;
 
 free_mem: