#define PRTS_MASK      0x1f            /* process table size field */
 #define PRTB_MASK      0x0ffffffffffff000UL
 
+/* Number of supported LPID bits */
+extern unsigned int mmu_lpid_bits;
+
 /* Number of supported PID bits */
 extern unsigned int mmu_pid_bits;
 
 #define PRTB_SIZE_SHIFT        (mmu_pid_bits + 4)
 #define PRTB_ENTRIES   (1ul << mmu_pid_bits)
 
-/*
- * Power9 currently only support 64K partition table size.
- */
-#define PATB_SIZE_SHIFT        16
+#define PATB_SIZE_SHIFT        (mmu_lpid_bits + 4)
+#define PATB_ENTRIES   (1ul << mmu_lpid_bits)
 
 typedef unsigned long mm_context_id_t;
 struct spinlock;
 
        unsigned long patb_size = 1UL << PATB_SIZE_SHIFT;
        unsigned long ptcr;
 
-       BUILD_BUG_ON_MSG((PATB_SIZE_SHIFT > 36), "Partition table size too large.");
        /* Initialize the Partition Table with no entries */
        partition_tb = memblock_alloc(patb_size, patb_size);
        if (!partition_tb)
                panic("%s: Failed to allocate %lu bytes align=0x%lx\n",
                      __func__, patb_size, patb_size);
 
-       /*
-        * update partition table control register,
-        * 64 K size.
-        */
        ptcr = __pa(partition_tb) | (PATB_SIZE_SHIFT - 12);
        set_ptcr_when_no_uv(ptcr);
        powernv_set_nmmu_ptcr(ptcr);
 
 
 #include <trace/events/thp.h>
 
-unsigned int mmu_pid_bits;
 unsigned int mmu_base_pid;
 unsigned long radix_mem_block_size __ro_after_init;
 
                                                -1, PAGE_KERNEL));
        }
 
-       /* Find out how many PID bits are supported */
        if (!cpu_has_feature(CPU_FTR_HVMODE) &&
                        cpu_has_feature(CPU_FTR_P9_RADIX_PREFETCH_BUG)) {
                /*
                 * Older versions of KVM on these machines perfer if the
                 * guest only uses the low 19 PID bits.
                 */
-               if (!mmu_pid_bits)
-                       mmu_pid_bits = 19;
-       } else {
-               if (!mmu_pid_bits)
-                       mmu_pid_bits = 20;
+               mmu_pid_bits = 19;
        }
        mmu_base_pid = 1;
 
        if (type == NULL || strcmp(type, "cpu") != 0)
                return 0;
 
-       /* Find MMU PID size */
-       prop = of_get_flat_dt_prop(node, "ibm,mmu-pid-bits", &size);
-       if (prop && size == 4)
-               mmu_pid_bits = be32_to_cpup(prop);
-
        /* Grab page size encodings */
        prop = of_get_flat_dt_prop(node, "ibm,processor-radix-AP-encodings", &size);
        if (!prop)
 
 #endif /* CONFIG_SPARSEMEM_VMEMMAP */
 
 #ifdef CONFIG_PPC_BOOK3S_64
+unsigned int mmu_lpid_bits;
+unsigned int mmu_pid_bits;
+
 static bool disable_radix = !IS_ENABLED(CONFIG_PPC_RADIX_MMU_DEFAULT);
 
 static int __init parse_disable_radix(char *p)
        }
 }
 
+static int __init dt_scan_mmu_pid_width(unsigned long node,
+                                          const char *uname, int depth,
+                                          void *data)
+{
+       int size = 0;
+       const __be32 *prop;
+       const char *type = of_get_flat_dt_prop(node, "device_type", NULL);
+
+       /* We are scanning "cpu" nodes only */
+       if (type == NULL || strcmp(type, "cpu") != 0)
+               return 0;
+
+       /* Find MMU LPID, PID register size */
+       prop = of_get_flat_dt_prop(node, "ibm,mmu-lpid-bits", &size);
+       if (prop && size == 4)
+               mmu_lpid_bits = be32_to_cpup(prop);
+
+       prop = of_get_flat_dt_prop(node, "ibm,mmu-pid-bits", &size);
+       if (prop && size == 4)
+               mmu_pid_bits = be32_to_cpup(prop);
+
+       if (!mmu_pid_bits && !mmu_lpid_bits)
+               return 0;
+
+       return 1;
+}
+
 void __init mmu_early_init_devtree(void)
 {
+       bool hvmode = !!(mfmsr() & MSR_HV);
+
        /* Disable radix mode based on kernel command line. */
        if (disable_radix)
                cur_cpu_spec->mmu_features &= ~MMU_FTR_TYPE_RADIX;
 
+       of_scan_flat_dt(dt_scan_mmu_pid_width, NULL);
+       if (hvmode && !mmu_lpid_bits) {
+               if (early_cpu_has_feature(CPU_FTR_ARCH_207S))
+                       mmu_lpid_bits = 12; /* POWER8-10 */
+               else
+                       mmu_lpid_bits = 10; /* POWER7 */
+       }
+       if (!mmu_pid_bits) {
+               if (early_cpu_has_feature(CPU_FTR_ARCH_300))
+                       mmu_pid_bits = 20; /* POWER9-10 */
+       }
+
        /*
         * Check /chosen/ibm,architecture-vec-5 if running as a guest.
         * When running bare-metal, we can use radix if we like
         * even though the ibm,architecture-vec-5 property created by
         * skiboot doesn't have the necessary bits set.
         */
-       if (!(mfmsr() & MSR_HV))
+       if (!hvmode)
                early_check_vec5();
 
        if (early_radix_enabled()) {