u64  pkey_amr_mask;            /* Bits in AMR not to be touched */
 u64  pkey_iamr_mask;           /* Bits in AMR not to be touched */
 u64  pkey_uamor_mask;          /* Bits in UMOR not to be touched */
+int  execute_only_key = 2;
 
 #define AMR_BITS_PER_PKEY 2
 #define AMR_RD_BIT 0x1UL
 #else
        os_reserved = 0;
 #endif
-       initial_allocation_mask  = (0x1 << 0) | (0x1 << 1);
+       initial_allocation_mask  = (0x1 << 0) | (0x1 << 1) |
+                                       (0x1 << execute_only_key);
 
        /* register mask is in BE format */
        pkey_amr_mask = ~0x0ul;
 
        pkey_iamr_mask = ~0x0ul;
        pkey_iamr_mask &= ~(0x3ul << pkeyshift(0));
+       pkey_iamr_mask &= ~(0x3ul << pkeyshift(execute_only_key));
 
        pkey_uamor_mask = ~0x0ul;
        pkey_uamor_mask &= ~(0x3ul << pkeyshift(0));
+       pkey_uamor_mask &= ~(0x3ul << pkeyshift(execute_only_key));
 
        /* mark the rest of the keys as reserved and hence unavailable */
        for (i = (pkeys_total - os_reserved); i < pkeys_total; i++) {
                pkey_uamor_mask &= ~(0x3ul << pkeyshift(i));
        }
 
+       if (unlikely((pkeys_total - os_reserved) <= execute_only_key)) {
+               /*
+                * Insufficient number of keys to support
+                * execute only key. Mark it unavailable.
+                * Any AMR, UAMOR, IAMR bit set for
+                * this key is irrelevant since this key
+                * can never be allocated.
+                */
+               execute_only_key = -1;
+       }
+
        return 0;
 }
 
        if (static_branch_likely(&pkey_disabled))
                return;
        mm_pkey_allocation_map(mm) = initial_allocation_mask;
-       /* -1 means unallocated or invalid */
-       mm->context.execute_only_pkey = -1;
+       mm->context.execute_only_pkey = execute_only_key;
 }
 
 static inline u64 read_amr(void)
 
 int __execute_only_pkey(struct mm_struct *mm)
 {
-       bool need_to_set_mm_pkey = false;
-       int execute_only_pkey = mm->context.execute_only_pkey;
-       int ret;
-
-       /* Do we need to assign a pkey for mm's execute-only maps? */
-       if (execute_only_pkey == -1) {
-               /* Go allocate one to use, which might fail */
-               execute_only_pkey = mm_pkey_alloc(mm);
-               if (execute_only_pkey < 0)
-                       return -1;
-               need_to_set_mm_pkey = true;
-       }
-
-       /*
-        * We do not want to go through the relatively costly dance to set AMR
-        * if we do not need to. Check it first and assume that if the
-        * execute-only pkey is readwrite-disabled than we do not have to set it
-        * ourselves.
-        */
-       if (!need_to_set_mm_pkey && !pkey_allows_readwrite(execute_only_pkey))
-               return execute_only_pkey;
-
-       /*
-        * Set up AMR so that it denies access for everything other than
-        * execution.
-        */
-       ret = __arch_set_user_pkey_access(current, execute_only_pkey,
-                                         PKEY_DISABLE_ACCESS |
-                                         PKEY_DISABLE_WRITE);
-       /*
-        * If the AMR-set operation failed somehow, just return 0 and
-        * effectively disable execute-only support.
-        */
-       if (ret) {
-               mm_pkey_free(mm, execute_only_pkey);
-               return -1;
-       }
-
-       /* We got one, store it and use it from here on out */
-       if (need_to_set_mm_pkey)
-               mm->context.execute_only_pkey = execute_only_pkey;
-       return execute_only_pkey;
+       return mm->context.execute_only_pkey;
 }
 
 static inline bool vma_is_pkey_exec_only(struct vm_area_struct *vma)