#define __pa(x)  ((unsigned long)(x))
 #define __va(x)  ((void *)((unsigned long)(x)))
 
-/*
- * The pgtable.h and mm/ident_map.c includes make use of the SME related
- * information which is not used in the compressed image support. Un-define
- * the SME support to avoid any compile and link errors.
- */
-#undef CONFIG_AMD_MEM_ENCRYPT
-
 /* No PAGE_TABLE_ISOLATION support needed either: */
 #undef CONFIG_PAGE_TABLE_ISOLATION
 
 /* Locates and clears a region for a new top level page table. */
 void initialize_identity_maps(void)
 {
-       unsigned long sev_me_mask = get_sev_encryption_mask();
+       /* If running as an SEV guest, the encryption mask is required. */
+       set_sev_encryption_mask();
 
        /* Init mapping_info with run-time function/buffer pointers. */
        mapping_info.alloc_pgt_page = alloc_pgt_page;
        mapping_info.context = &pgt_data;
-       mapping_info.page_flag = __PAGE_KERNEL_LARGE_EXEC | sev_me_mask;
-       mapping_info.kernpg_flag = _KERNPG_TABLE | sev_me_mask;
+       mapping_info.page_flag = __PAGE_KERNEL_LARGE_EXEC | sme_me_mask;
+       mapping_info.kernpg_flag = _KERNPG_TABLE;
 
        /*
         * It should be impossible for this not to already be true,
 
 ENDPROC(get_sev_encryption_bit)
 
        .code64
-ENTRY(get_sev_encryption_mask)
-       xor     %rax, %rax
-
+ENTRY(set_sev_encryption_mask)
 #ifdef CONFIG_AMD_MEM_ENCRYPT
        push    %rbp
        push    %rdx
        testl   %eax, %eax
        jz      .Lno_sev_mask
 
-       xor     %rdx, %rdx
-       bts     %rax, %rdx              /* Create the encryption mask */
-       mov     %rdx, %rax              /* ... and return it */
+       bts     %rax, sme_me_mask(%rip) /* Create the encryption mask */
 
 .Lno_sev_mask:
        movq    %rbp, %rsp              /* Restore original stack pointer */
        pop     %rbp
 #endif
 
+       xor     %rax, %rax
        ret
-ENDPROC(get_sev_encryption_mask)
+ENDPROC(set_sev_encryption_mask)
 
        .data
 enc_bit:
        .int    0xffffffff
+
+#ifdef CONFIG_AMD_MEM_ENCRYPT
+       .balign 8
+GLOBAL(sme_me_mask)
+       .quad   0
+#endif