#include <linux/io.h>
 #include <asm/opal.h>
 
-/*
- * For static allocation of some of the structures.
- */
-#define IMC_MAX_PMUS                   32
-
 /*
  * Compatibility macros for IMC devices
  */
 extern int init_imc_pmu(struct device_node *parent,
                                struct imc_pmu *pmu_ptr, int pmu_id);
 extern void thread_imc_disable(void);
+extern int get_max_nest_dev(void);
 #endif /* __ASM_POWERPC_IMC_PMU_H */
 
  */
 static DEFINE_MUTEX(nest_init_lock);
 static DEFINE_PER_CPU(struct imc_pmu_ref *, local_nest_imc_refc);
-static struct imc_pmu *per_nest_pmu_arr[IMC_MAX_PMUS];
+static struct imc_pmu **per_nest_pmu_arr;
 static cpumask_t nest_imc_cpumask;
 struct imc_pmu_ref *nest_imc_refc;
 static int nest_pmus;
 static void nest_change_cpu_context(int old_cpu, int new_cpu)
 {
        struct imc_pmu **pn = per_nest_pmu_arr;
-       int i;
 
        if (old_cpu < 0 || new_cpu < 0)
                return;
 
-       for (i = 0; *pn && i < IMC_MAX_PMUS; i++, pn++)
+       while (*pn) {
                perf_pmu_migrate_context(&(*pn)->pmu, old_cpu, new_cpu);
+               pn++;
+       }
 }
 
 static int ppc_nest_imc_cpu_offline(unsigned int cpu)
                kfree(pmu_ptr->attr_groups[IMC_EVENT_ATTR]->attrs);
        kfree(pmu_ptr->attr_groups[IMC_EVENT_ATTR]);
        kfree(pmu_ptr);
+       kfree(per_nest_pmu_arr);
        return;
 }
 
                        return -ENOMEM;
 
                /* Needed for hotplug/migration */
+               if (!per_nest_pmu_arr) {
+                       per_nest_pmu_arr = kcalloc(get_max_nest_dev() + 1,
+                                               sizeof(struct imc_pmu *),
+                                               GFP_KERNEL);
+                       if (!per_nest_pmu_arr)
+                               return -ENOMEM;
+               }
                per_nest_pmu_arr[pmu_index] = pmu_ptr;
                break;
        case IMC_DOMAIN_CORE:
 
        put_online_cpus();
 }
 
+int get_max_nest_dev(void)
+{
+       struct device_node *node;
+       u32 pmu_units = 0, type;
+
+       for_each_compatible_node(node, NULL, IMC_DTB_UNIT_COMPAT) {
+               if (of_property_read_u32(node, "type", &type))
+                       continue;
+
+               if (type == IMC_TYPE_CHIP)
+                       pmu_units++;
+       }
+
+       return pmu_units;
+}
+
 static int opal_imc_counters_probe(struct platform_device *pdev)
 {
        struct device_node *imc_dev = pdev->dev.of_node;