/* Architecture __weak replacement functions */
 void __init mem_encrypt_init(void);
 
+void add_encrypt_protection_map(void);
+
 /*
  * The __sme_pa() and __sme_pa_nodebug() macros are meant for use when
  * writing to or comparing values from the cr3 register.  Having the
 
 
 #endif /* __ASSEMBLY__ */
 
-/*         xwr */
-#define __P000 PAGE_NONE
-#define __P001 PAGE_READONLY
-#define __P010 PAGE_COPY
-#define __P011 PAGE_COPY
-#define __P100 PAGE_READONLY_EXEC
-#define __P101 PAGE_READONLY_EXEC
-#define __P110 PAGE_COPY_EXEC
-#define __P111 PAGE_COPY_EXEC
-
-#define __S000 PAGE_NONE
-#define __S001 PAGE_READONLY
-#define __S010 PAGE_SHARED
-#define __S011 PAGE_SHARED
-#define __S100 PAGE_READONLY_EXEC
-#define __S101 PAGE_READONLY_EXEC
-#define __S110 PAGE_SHARED_EXEC
-#define __S111 PAGE_SHARED_EXEC
-
 /*
  * early identity mapping  pte attrib macros.
  */
 
 #include <asm/tlbflush.h>
 #include <asm/fixmap.h>
 #include <asm/setup.h>
+#include <asm/mem_encrypt.h>
 #include <asm/bootparam.h>
 #include <asm/set_memory.h>
 #include <asm/cacheflush.h>
 
 void __init sme_early_init(void)
 {
-       unsigned int i;
-
        if (!sme_me_mask)
                return;
 
        __supported_pte_mask = __sme_set(__supported_pte_mask);
 
        /* Update the protection map with memory encryption mask */
-       for (i = 0; i < ARRAY_SIZE(protection_map); i++)
-               protection_map[i] = pgprot_encrypted(protection_map[i]);
+       add_encrypt_protection_map();
 
        x86_platform.guest.enc_status_change_prepare = amd_enc_status_change_prepare;
        x86_platform.guest.enc_status_change_finish  = amd_enc_status_change_finish;
 
 #include <linux/export.h>
 #include <linux/mm.h>
 #include <asm/pgtable.h>
+#include <asm/mem_encrypt.h>
+
+static pgprot_t protection_map[16] __ro_after_init = {
+       [VM_NONE]                                       = PAGE_NONE,
+       [VM_READ]                                       = PAGE_READONLY,
+       [VM_WRITE]                                      = PAGE_COPY,
+       [VM_WRITE | VM_READ]                            = PAGE_COPY,
+       [VM_EXEC]                                       = PAGE_READONLY_EXEC,
+       [VM_EXEC | VM_READ]                             = PAGE_READONLY_EXEC,
+       [VM_EXEC | VM_WRITE]                            = PAGE_COPY_EXEC,
+       [VM_EXEC | VM_WRITE | VM_READ]                  = PAGE_COPY_EXEC,
+       [VM_SHARED]                                     = PAGE_NONE,
+       [VM_SHARED | VM_READ]                           = PAGE_READONLY,
+       [VM_SHARED | VM_WRITE]                          = PAGE_SHARED,
+       [VM_SHARED | VM_WRITE | VM_READ]                = PAGE_SHARED,
+       [VM_SHARED | VM_EXEC]                           = PAGE_READONLY_EXEC,
+       [VM_SHARED | VM_EXEC | VM_READ]                 = PAGE_READONLY_EXEC,
+       [VM_SHARED | VM_EXEC | VM_WRITE]                = PAGE_SHARED_EXEC,
+       [VM_SHARED | VM_EXEC | VM_WRITE | VM_READ]      = PAGE_SHARED_EXEC
+};
+
+void add_encrypt_protection_map(void)
+{
+       unsigned int i;
+
+       for (i = 0; i < ARRAY_SIZE(protection_map); i++)
+               protection_map[i] = pgprot_encrypted(protection_map[i]);
+}
 
 pgprot_t vm_get_page_prot(unsigned long vm_flags)
 {