#define CPU_FTR_STCX_CHECKS_ADDRESS    LONG_ASM_CONST(0x0004000000000000)
 #define CPU_FTR_POPCNTB                        LONG_ASM_CONST(0x0008000000000000)
 #define CPU_FTR_POPCNTD                        LONG_ASM_CONST(0x0010000000000000)
-/* Free                                        LONG_ASM_CONST(0x0020000000000000) */
+#define CPU_FTR_PKEY                   LONG_ASM_CONST(0x0020000000000000)
 #define CPU_FTR_VMX_COPY               LONG_ASM_CONST(0x0040000000000000)
 #define CPU_FTR_TM                     LONG_ASM_CONST(0x0080000000000000)
 #define CPU_FTR_CFAR                   LONG_ASM_CONST(0x0100000000000000)
            CPU_FTR_DSCR | CPU_FTR_SAO  | CPU_FTR_ASYM_SMT | \
            CPU_FTR_STCX_CHECKS_ADDRESS | CPU_FTR_POPCNTB | CPU_FTR_POPCNTD | \
            CPU_FTR_CFAR | CPU_FTR_HVMODE | \
-           CPU_FTR_VMX_COPY | CPU_FTR_HAS_PPR | CPU_FTR_DABRX)
+           CPU_FTR_VMX_COPY | CPU_FTR_HAS_PPR | CPU_FTR_DABRX | CPU_FTR_PKEY)
 #define CPU_FTRS_POWER8 (CPU_FTR_USE_TB | CPU_FTR_LWSYNC | \
            CPU_FTR_PPCAS_ARCH_V2 | CPU_FTR_CTRL | CPU_FTR_ARCH_206 |\
            CPU_FTR_MMCRA | CPU_FTR_SMT | \
            CPU_FTR_STCX_CHECKS_ADDRESS | CPU_FTR_POPCNTB | CPU_FTR_POPCNTD | \
            CPU_FTR_CFAR | CPU_FTR_HVMODE | CPU_FTR_VMX_COPY | \
            CPU_FTR_DBELL | CPU_FTR_HAS_PPR | CPU_FTR_DAWR | \
-           CPU_FTR_ARCH_207S | CPU_FTR_TM_COMP)
+           CPU_FTR_ARCH_207S | CPU_FTR_TM_COMP | CPU_FTR_PKEY)
 #define CPU_FTRS_POWER8E (CPU_FTRS_POWER8 | CPU_FTR_PMAO_BUG)
 #define CPU_FTRS_POWER8_DD1 (CPU_FTRS_POWER8 & ~CPU_FTR_DBELL)
 #define CPU_FTRS_POWER9 (CPU_FTR_USE_TB | CPU_FTR_LWSYNC | \
            CPU_FTR_STCX_CHECKS_ADDRESS | CPU_FTR_POPCNTB | CPU_FTR_POPCNTD | \
            CPU_FTR_CFAR | CPU_FTR_HVMODE | CPU_FTR_VMX_COPY | \
            CPU_FTR_DBELL | CPU_FTR_HAS_PPR | CPU_FTR_DAWR | \
-           CPU_FTR_ARCH_207S | CPU_FTR_TM_COMP | CPU_FTR_ARCH_300)
+           CPU_FTR_ARCH_207S | CPU_FTR_TM_COMP | CPU_FTR_ARCH_300 | \
+           CPU_FTR_PKEY)
 #define CPU_FTRS_POWER9_DD1 ((CPU_FTRS_POWER9 | CPU_FTR_POWER9_DD1) & \
                             (~CPU_FTR_SAO))
 #define CPU_FTRS_POWER9_DD2_0 CPU_FTRS_POWER9
 
 #define _ASM_POWERPC_KEYS_H
 
 #include <linux/jump_label.h>
+#include <asm/firmware.h>
 
 DECLARE_STATIC_KEY_TRUE(pkey_disabled);
 extern int pkeys_total; /* total pkeys as per device tree */
 }
 
 extern void pkey_mm_init(struct mm_struct *mm);
+extern bool arch_supports_pkeys(int cap);
+extern unsigned int arch_usable_pkeys(void);
 extern void thread_pkey_regs_save(struct thread_struct *thread);
 extern void thread_pkey_regs_restore(struct thread_struct *new_thread,
                                     struct thread_struct *old_thread);
 
  */
 
 #include <asm/mman.h>
+#include <asm/setup.h>
 #include <linux/pkeys.h>
+#include <linux/of_device.h>
 
 DEFINE_STATIC_KEY_TRUE(pkey_disabled);
 bool pkey_execute_disable_supported;
 int  pkeys_total;              /* Total pkeys as per device tree */
+bool pkeys_devtree_defined;    /* pkey property exported by device tree */
 u32  initial_allocation_mask;  /* Bits set for reserved keys */
 u64  pkey_amr_uamor_mask;      /* Bits in AMR/UMOR not to be touched */
 u64  pkey_iamr_mask;           /* Bits in AMR not to be touched */
 #define PKEY_REG_BITS (sizeof(u64)*8)
 #define pkeyshift(pkey) (PKEY_REG_BITS - ((pkey+1) * AMR_BITS_PER_PKEY))
 
+static void scan_pkey_feature(void)
+{
+       u32 vals[2];
+       struct device_node *cpu;
+
+       cpu = of_find_node_by_type(NULL, "cpu");
+       if (!cpu)
+               return;
+
+       if (of_property_read_u32_array(cpu,
+                       "ibm,processor-storage-keys", vals, 2))
+               return;
+
+       /*
+        * Since any pkey can be used for data or execute, we will just treat
+        * all keys as equal and track them as one entity.
+        */
+       pkeys_total = be32_to_cpu(vals[0]);
+       pkeys_devtree_defined = true;
+}
+
+static inline bool pkey_mmu_enabled(void)
+{
+       if (firmware_has_feature(FW_FEATURE_LPAR))
+               return pkeys_total;
+       else
+               return cpu_has_feature(CPU_FTR_PKEY);
+}
+
 int pkey_initialize(void)
 {
        int os_reserved, i;
                     __builtin_popcountl(ARCH_VM_PKEY_FLAGS >> VM_PKEY_SHIFT)
                                != (sizeof(u64) * BITS_PER_BYTE));
 
+       /* scan the device tree for pkey feature */
+       scan_pkey_feature();
+
        /*
-        * Disable the pkey system till everything is in place. A subsequent
-        * patch will enable it.
+        * Let's assume 32 pkeys on P8 bare metal, if its not defined by device
+        * tree. We make this exception since skiboot forgot to expose this
+        * property on power8.
         */
-       static_branch_enable(&pkey_disabled);
-
-       /* Lets assume 32 keys */
-       pkeys_total = 32;
+       if (!pkeys_devtree_defined && !firmware_has_feature(FW_FEATURE_LPAR) &&
+                       cpu_has_feature(CPU_FTRS_POWER8))
+               pkeys_total = 32;
 
        /*
         * Adjust the upper limit, based on the number of bits supported by
        pkeys_total = min_t(int, pkeys_total,
                        (ARCH_VM_PKEY_FLAGS >> VM_PKEY_SHIFT));
 
+       if (!pkey_mmu_enabled() || radix_enabled() || !pkeys_total)
+               static_branch_enable(&pkey_disabled);
+       else
+               static_branch_disable(&pkey_disabled);
+
+       if (static_branch_likely(&pkey_disabled))
+               return 0;
+
        /*
-        * Disable execute_disable support for now. A subsequent patch will
-        * enable it.
+        * The device tree cannot be relied to indicate support for
+        * execute_disable support. Instead we use a PVR check.
         */
-       pkey_execute_disable_supported = false;
+       if (pvr_version_is(PVR_POWER7) || pvr_version_is(PVR_POWER7p))
+               pkey_execute_disable_supported = false;
+       else
+               pkey_execute_disable_supported = true;
 
 #ifdef CONFIG_PPC_4K_PAGES
        /*