[PATCH v6 29/29] UefiCpuPkg/MpInitLib: Use SEV-SNP AP Creation NAE event to launch APs

["Brijesh Singh" <brijesh.singh@amd.com>]

From: Tom Lendacky <thomas.lendacky@amd.com>

BZ: https://bugzilla.tianocore.org/show_bug.cgi?id=3275

Use the SEV-SNP AP Creation NAE event to create and launch APs under
SEV-SNP. This capability will be advertised in the SEV Hypervisor
Feature Support PCD (PcdSevEsHypervisorFeatures).

Cc: Michael Roth <michael.roth@amd.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Cc: James Bottomley <jejb@linux.ibm.com>
Cc: Min Xu <min.m.xu@intel.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Tom Lendacky <thomas.lendacky@amd.com>
Cc: Jordan Justen <jordan.l.justen@intel.com>
Cc: Ard Biesheuvel <ardb+tianocore@kernel.org>
Cc: Erdem Aktas <erdemaktas@google.com>
Signed-off-by: Tom Lendacky <thomas.lendacky@amd.com>
Signed-off-by: Brijesh Singh <brijesh.singh@amd.com>
---
 UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf |   3 +
 UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf |   5 +-
 UefiCpuPkg/Library/MpInitLib/MpLib.h          |  17 ++
 UefiCpuPkg/Library/MpInitLib/DxeMpLib.c       |  11 +-
 .../MpInitLib/Ia32/SevSnpRmpAdjustInternal.c  |  31 ++
 UefiCpuPkg/Library/MpInitLib/MpLib.c          | 275 ++++++++++++++++--
 .../MpInitLib/X64/SevSnpRmpAdjustInternal.c   |  44 +++
 7 files changed, 361 insertions(+), 25 deletions(-)
 create mode 100644 UefiCpuPkg/Library/MpInitLib/Ia32/SevSnpRmpAdjustInternal.c
 create mode 100644 UefiCpuPkg/Library/MpInitLib/X64/SevSnpRmpAdjustInternal.c

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf b/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf
index 48d7dfa4450f..b9ce05e81b54 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf
@@ -23,9 +23,11 @@ [Defines]
 
 [Sources.IA32]
   Ia32/MpFuncs.nasm
+  Ia32/SevSnpRmpAdjustInternal.c
 
 [Sources.X64]
   X64/MpFuncs.nasm
+  X64/SevSnpRmpAdjustInternal.c
 
 [Sources.common]
   MpEqu.inc
@@ -72,6 +74,7 @@ [Pcd]
   gUefiCpuPkgTokenSpaceGuid.PcdCpuApLoopMode                           ## CONSUMES
   gUefiCpuPkgTokenSpaceGuid.PcdCpuApTargetCstate                       ## SOMETIMES_CONSUMES
   gUefiCpuPkgTokenSpaceGuid.PcdCpuApStatusCheckIntervalInMicroSeconds  ## CONSUMES
+  gUefiCpuPkgTokenSpaceGuid.PcdGhcbHypervisorFeatures                  ## CONSUMES
   gUefiCpuPkgTokenSpaceGuid.PcdSevEsIsEnabled                          ## CONSUMES
   gUefiCpuPkgTokenSpaceGuid.PcdSevEsWorkAreaBase                       ## SOMETIMES_CONSUMES
   gEfiMdeModulePkgTokenSpaceGuid.PcdCpuStackGuard                      ## CONSUMES
diff --git a/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf b/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf
index ab8279df596f..35057ac07cbb 100644
--- a/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf
+++ b/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf
@@ -23,9 +23,11 @@ [Defines]
 
 [Sources.IA32]
   Ia32/MpFuncs.nasm
+  Ia32/SevSnpRmpAdjustInternal.c
 
 [Sources.X64]
   X64/MpFuncs.nasm
+  X64/SevSnpRmpAdjustInternal.c
 
 [Sources.common]
   MpEqu.inc
@@ -62,10 +64,11 @@ [Pcd]
   gUefiCpuPkgTokenSpaceGuid.PcdCpuMicrocodePatchRegionSize         ## CONSUMES
   gUefiCpuPkgTokenSpaceGuid.PcdCpuApLoopMode                       ## CONSUMES
   gUefiCpuPkgTokenSpaceGuid.PcdCpuApTargetCstate                   ## SOMETIMES_CONSUMES
+  gUefiCpuPkgTokenSpaceGuid.PcdGhcbHypervisorFeatures              ## CONSUMES
   gUefiCpuPkgTokenSpaceGuid.PcdSevEsIsEnabled                      ## CONSUMES
   gUefiCpuPkgTokenSpaceGuid.PcdSevEsWorkAreaBase                   ## SOMETIMES_CONSUMES
-  gEfiMdeModulePkgTokenSpaceGuid.PcdGhcbBase                       ## CONSUMES
   gUefiCpuPkgTokenSpaceGuid.PcdSevSnpIsEnabled                     ## CONSUMES
+  gEfiMdeModulePkgTokenSpaceGuid.PcdGhcbBase                       ## CONSUMES
 
 [Ppis]
   gEdkiiPeiShadowMicrocodePpiGuid        ## SOMETIMES_CONSUMES
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index 19e91bf7d74e..a046de8e5c68 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -15,6 +15,7 @@
 
 #include <Register/Intel/Cpuid.h>
 #include <Register/Amd/Cpuid.h>
+#include <Register/Amd/Ghcb.h>
 #include <Register/Intel/Msr.h>
 #include <Register/Intel/LocalApic.h>
 #include <Register/Intel/Microcode.h>
@@ -146,6 +147,7 @@ typedef struct {
   UINT8                          PlatformId;
   UINT64                         MicrocodeEntryAddr;
   UINT32                         MicrocodeRevision;
+  SEV_ES_SAVE_AREA               *SevEsSaveArea;
 } CPU_AP_DATA;
 
 //
@@ -290,6 +292,7 @@ struct _CPU_MP_DATA {
 
   BOOLEAN                        SevEsIsEnabled;
   BOOLEAN                        SevSnpIsEnabled;
+  BOOLEAN                        UseSevEsAPMethod;
   UINTN                          SevEsAPBuffer;
   UINTN                          SevEsAPResetStackStart;
   CPU_MP_DATA                    *NewCpuMpData;
@@ -744,5 +747,19 @@ PlatformShadowMicrocode (
   IN OUT CPU_MP_DATA             *CpuMpData
   );
 
+/**
+  Issue RMPADJUST to adjust the VMSA attribute of an SEV-SNP page.
+
+  @param[in]  PageAddress
+  @param[in]  VmsaPage
+
+  @return  RMPADJUST return value
+**/
+UINT32
+SevSnpRmpAdjust (
+  IN  EFI_PHYSICAL_ADDRESS  PageAddress,
+  IN  BOOLEAN               VmsaPage
+  );
+
 #endif
 
diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index 93fc63bf93e3..15e797cd0990 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
@@ -93,7 +93,12 @@ GetWakeupBuffer (
   EFI_PHYSICAL_ADDRESS    StartAddress;
   EFI_MEMORY_TYPE         MemoryType;
 
-  if (PcdGetBool (PcdSevEsIsEnabled)) {
+  //
+  // An SEV-ES-only guest requires the memory to be reserved. SEV-SNP, which
+  // is also considered SEV-ES, uses a different AP startup method, though,
+  // which does not have the same requirement.
+  //
+  if (PcdGetBool (PcdSevEsIsEnabled) && !PcdGetBool (PcdSevSnpIsEnabled)) {
     MemoryType = EfiReservedMemoryType;
   } else {
     MemoryType = EfiBootServicesData;
@@ -373,7 +378,7 @@ RelocateApLoop (
   MpInitLibWhoAmI (&ProcessorNumber);
   CpuMpData    = GetCpuMpData ();
   MwaitSupport = IsMwaitSupport ();
-  if (CpuMpData->SevEsIsEnabled) {
+  if (CpuMpData->UseSevEsAPMethod) {
     StackStart = CpuMpData->SevEsAPResetStackStart;
   } else {
     StackStart = mReservedTopOfApStack;
@@ -422,7 +427,7 @@ MpInitChangeApLoopCallback (
     CpuPause ();
   }
 
-  if (CpuMpData->SevEsIsEnabled && (CpuMpData->WakeupBuffer != (UINTN) -1)) {
+  if (CpuMpData->UseSevEsAPMethod && (CpuMpData->WakeupBuffer != (UINTN) -1)) {
     //
     // There are APs present. Re-use reserved memory area below 1MB from
     // WakeupBuffer as the area to be used for transitioning to 16-bit mode
diff --git a/UefiCpuPkg/Library/MpInitLib/Ia32/SevSnpRmpAdjustInternal.c b/UefiCpuPkg/Library/MpInitLib/Ia32/SevSnpRmpAdjustInternal.c
new file mode 100644
index 000000000000..5c9ee08b2117
--- /dev/null
+++ b/UefiCpuPkg/Library/MpInitLib/Ia32/SevSnpRmpAdjustInternal.c
@@ -0,0 +1,31 @@
+/** @file
+
+  RMPADJUST helper function.
+
+  Copyright (c) 2021, AMD Incorporated. All rights reserved.<BR>
+
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include "MpLib.h"
+
+/**
+  Issue RMPADJUST to adjust the VMSA attribute of an SEV-SNP page.
+
+  @param[in]  PageAddress
+  @param[in]  VmsaPage
+
+  @return  RMPADJUST return value
+**/
+UINT32
+SevSnpRmpAdjust (
+  IN  EFI_PHYSICAL_ADDRESS  PageAddress,
+  IN  BOOLEAN               VmsaPage
+  )
+{
+  //
+  // RMPADJUST is not supported in 32-bit mode
+  //
+  return RETURN_UNSUPPORTED;
+}
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index c2b00a1f04ff..01f11ab1d9b3 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -295,10 +295,10 @@ GetApLoopMode (
       ApLoopMode = ApInHltLoop;
     }
 
-    if (PcdGetBool (PcdSevEsIsEnabled)) {
+    if (PcdGetBool (PcdSevEsIsEnabled) && !PcdGetBool (PcdSevSnpIsEnabled)) {
       //
-      // For SEV-ES, force AP in Hlt-loop mode in order to use the GHCB
-      // protocol for starting APs
+      // For SEV-ES (SEV-SNP is also considered SEV-ES), force AP in Hlt-loop
+      // mode in order to use the GHCB protocol for starting APs
       //
       ApLoopMode = ApInHltLoop;
     }
@@ -869,7 +869,7 @@ ApWakeupFunction (
       // to allow the APs to issue an AP_RESET_HOLD before the BSP possibly
       // performs another INIT-SIPI-SIPI sequence.
       //
-      if (!CpuMpData->SevEsIsEnabled) {
+      if (!CpuMpData->UseSevEsAPMethod) {
         InterlockedDecrement ((UINT32 *) &CpuMpData->MpCpuExchangeInfo->NumApsExecuting);
       }
     }
@@ -883,7 +883,7 @@ ApWakeupFunction (
       //
       while (TRUE) {
         DisableInterrupts ();
-        if (CpuMpData->SevEsIsEnabled) {
+        if (CpuMpData->UseSevEsAPMethod) {
           MSR_SEV_ES_GHCB_REGISTER  Msr;
           GHCB                      *Ghcb;
           UINT64                    Status;
@@ -1207,9 +1207,11 @@ AllocateResetVector (
                                     );
     //
     // The AP reset stack is only used by SEV-ES guests. Do not allocate it
-    // if SEV-ES is not enabled.
+    // if SEV-ES is not enabled. An SEV-SNP guest is also considered
+    // an SEV-ES guest, but uses a different method of AP startup, eliminating
+    // the need for the allocation.
     //
-    if (PcdGetBool (PcdSevEsIsEnabled)) {
+    if (PcdGetBool (PcdSevEsIsEnabled) && !PcdGetBool (PcdSevSnpIsEnabled)) {
       //
       // Stack location is based on ProcessorNumber, so use the total number
       // of processors for calculating the total stack area.
@@ -1259,7 +1261,7 @@ FreeResetVector (
   // perform the restore as this will overwrite memory which has data
   // needed by SEV-ES.
   //
-  if (!CpuMpData->SevEsIsEnabled) {
+  if (!CpuMpData->UseSevEsAPMethod) {
     RestoreWakeupBuffer (CpuMpData);
   }
 }
@@ -1276,7 +1278,7 @@ AllocateSevEsAPMemory (
 {
   if (CpuMpData->SevEsAPBuffer == (UINTN) -1) {
     CpuMpData->SevEsAPBuffer =
-      CpuMpData->SevEsIsEnabled ? GetSevEsAPMemory () : 0;
+      CpuMpData->UseSevEsAPMethod ? GetSevEsAPMemory () : 0;
   }
 }
 
@@ -1329,6 +1331,222 @@ SetSevEsJumpTable (
   JmpFar->Segment = (UINT16) (SipiVector >> 4);
 }
 
+/**
+  Create an SEV-SNP AP save area (VMSA) for use in running the vCPU.
+
+  @param[in]  CpuMpData        Pointer to CPU MP Data
+  @param[in]  CpuData          Pointer to CPU AP Data
+  @param[in]  ApicId           APIC ID of the vCPU
+**/
+STATIC
+VOID
+SevSnpCreateSaveArea (
+  IN CPU_MP_DATA               *CpuMpData,
+  IN CPU_AP_DATA               *CpuData,
+  UINT32                       ApicId
+  )
+{
+  SEV_ES_SAVE_AREA          *SaveArea;
+  IA32_CR0                  ApCr0;
+  IA32_CR0                  ResetCr0;
+  IA32_CR4                  ApCr4;
+  IA32_CR4                  ResetCr4;
+  UINTN                     StartIp;
+  UINT8                     SipiVector;
+  UINT32                    RmpAdjustStatus;
+  UINT64                    VmgExitStatus;
+  MSR_SEV_ES_GHCB_REGISTER  Msr;
+  GHCB                      *Ghcb;
+  BOOLEAN                   InterruptState;
+  UINT64                    ExitInfo1;
+  UINT64                    ExitInfo2;
+
+  //
+  // Allocate a single page for the SEV-ES Save Area and initialize it.
+  //
+  SaveArea = AllocateReservedPages (1);
+  if (!SaveArea) {
+    return;
+  }
+  ZeroMem (SaveArea, EFI_PAGE_SIZE);
+
+  //
+  // Propogate the CR0.NW and CR0.CD setting to the AP
+  //
+  ResetCr0.UintN = 0x00000010;
+  ApCr0.UintN = CpuData->VolatileRegisters.Cr0;
+  if (ApCr0.Bits.NW) {
+    ResetCr0.Bits.NW = 1;
+  }
+  if (ApCr0.Bits.CD) {
+    ResetCr0.Bits.CD = 1;
+  }
+
+  //
+  // Propagate the CR4.MCE setting to the AP
+  //
+  ResetCr4.UintN = 0;
+  ApCr4.UintN = CpuData->VolatileRegisters.Cr4;
+  if (ApCr4.Bits.MCE) {
+    ResetCr4.Bits.MCE = 1;
+  }
+
+  //
+  // Convert the start IP into a SIPI Vector
+  //
+  StartIp = CpuMpData->MpCpuExchangeInfo->BufferStart;
+  SipiVector = (UINT8) (StartIp >> 12);
+
+  //
+  // Set the CS:RIP value based on the start IP
+  //
+  SaveArea->Cs.Base = SipiVector << 12;
+  SaveArea->Cs.Selector = SipiVector << 8;
+  SaveArea->Cs.Limit = 0xFFFF;
+  SaveArea->Cs.Attributes.Bits.Present = 1;
+  SaveArea->Cs.Attributes.Bits.Sbit = 1;
+  SaveArea->Cs.Attributes.Bits.Type = SEV_ES_RESET_CODE_SEGMENT_TYPE;
+  SaveArea->Rip = StartIp & 0xFFF;
+
+  //
+  // Set the remaining values as defined in APM for INIT
+  //
+  SaveArea->Ds.Limit = 0xFFFF;
+  SaveArea->Ds.Attributes.Bits.Present = 1;
+  SaveArea->Ds.Attributes.Bits.Sbit = 1;
+  SaveArea->Ds.Attributes.Bits.Type = SEV_ES_RESET_DATA_SEGMENT_TYPE;
+  SaveArea->Es = SaveArea->Ds;
+  SaveArea->Fs = SaveArea->Ds;
+  SaveArea->Gs = SaveArea->Ds;
+  SaveArea->Ss = SaveArea->Ds;
+
+  SaveArea->Gdtr.Limit = 0xFFFF;
+  SaveArea->Ldtr.Limit = 0xFFFF;
+  SaveArea->Ldtr.Attributes.Bits.Present = 1;
+  SaveArea->Ldtr.Attributes.Bits.Type = SEV_ES_RESET_LDT_TYPE;
+  SaveArea->Idtr.Limit = 0xFFFF;
+  SaveArea->Tr.Limit = 0xFFFF;
+  SaveArea->Ldtr.Attributes.Bits.Present = 1;
+  SaveArea->Ldtr.Attributes.Bits.Type = SEV_ES_RESET_TSS_TYPE;
+
+  SaveArea->Efer   = 0x1000;
+  SaveArea->Cr4    = ResetCr4.UintN;
+  SaveArea->Cr0    = ResetCr0.UintN;
+  SaveArea->Dr7    = 0x0400;
+  SaveArea->Dr6    = 0xFFFF0FF0;
+  SaveArea->Rflags = 0x0002;
+  SaveArea->GPat   = 0x0007040600070406ULL;
+  SaveArea->XCr0   = 0x0001;
+  SaveArea->Mxcsr  = 0x1F80;
+  SaveArea->X87Ftw = 0x5555;
+  SaveArea->X87Fcw = 0x0040;
+
+  //
+  // Set the SEV-SNP specific fields for the save area:
+  //   VMPL - always VMPL0
+  //   SEV_FEATURES - equivalent to the SEV_STATUS MSR right shifted 2 bits
+  //
+  SaveArea->Vmpl        = 0;
+  SaveArea->SevFeatures = AsmReadMsr64 (MSR_SEV_STATUS) >> 2;
+
+  //
+  // To turn the page into a recognized VMSA page, issue RMPADJUST:
+  //   Target VMPL but numerically higher than current VMPL
+  //   Target PermissionMask is not used
+  //
+  RmpAdjustStatus = SevSnpRmpAdjust (
+                      (EFI_PHYSICAL_ADDRESS) (UINTN) SaveArea,
+                      TRUE
+                      );
+  ASSERT (RmpAdjustStatus == 0);
+
+  ExitInfo1 = (UINT64) ApicId << 32;
+  ExitInfo1 |= SVM_VMGEXIT_SNP_AP_CREATE;
+  ExitInfo2 = (UINT64) (UINTN) SaveArea;
+
+  Msr.GhcbPhysicalAddress = AsmReadMsr64 (MSR_SEV_ES_GHCB);
+  Ghcb = Msr.Ghcb;
+
+  VmgInit (Ghcb, &InterruptState);
+  Ghcb->SaveArea.Rax = SaveArea->SevFeatures;
+  VmgSetOffsetValid (Ghcb, GhcbRax);
+  VmgExitStatus = VmgExit (
+                    Ghcb,
+                    SVM_EXIT_SNP_AP_CREATION,
+                    ExitInfo1,
+                    ExitInfo2
+                    );
+  VmgDone (Ghcb, InterruptState);
+
+  ASSERT (VmgExitStatus == 0);
+  if (VmgExitStatus != 0) {
+    RmpAdjustStatus = SevSnpRmpAdjust (
+                        (EFI_PHYSICAL_ADDRESS) (UINTN) SaveArea,
+                        FALSE
+                        );
+    if (RmpAdjustStatus == 0) {
+      FreePages (SaveArea, 1);
+    } else {
+      DEBUG ((DEBUG_INFO, "SEV-SNP: RMPADJUST failed, leaking VMSA page\n"));
+    }
+
+    SaveArea = NULL;
+  }
+
+  if (CpuData->SevEsSaveArea) {
+    RmpAdjustStatus = SevSnpRmpAdjust (
+                        (EFI_PHYSICAL_ADDRESS) (UINTN) CpuData->SevEsSaveArea,
+                        FALSE
+                        );
+    if (RmpAdjustStatus == 0) {
+      FreePages (CpuData->SevEsSaveArea, 1);
+    } else {
+      DEBUG ((DEBUG_INFO, "SEV-SNP: RMPADJUST failed, leaking VMSA page\n"));
+    }
+  }
+
+  CpuData->SevEsSaveArea = SaveArea;
+}
+
+/**
+  Create SEV-SNP APs.
+
+  @param[in]  CpuMpData        Pointer to CPU MP Data
+  @param[in]  ProcessorNumber  The handle number of specified processor
+                               (-1 for all APs)
+**/
+STATIC
+VOID
+SevSnpCreateAP (
+  IN CPU_MP_DATA               *CpuMpData,
+  IN INTN                      ProcessorNumber
+  )
+{
+  CPU_INFO_IN_HOB  *CpuInfoInHob;
+  CPU_AP_DATA      *CpuData;
+  UINTN            Index;
+  UINT32           ApicId;
+
+  ASSERT (CpuMpData->MpCpuExchangeInfo->BufferStart < 0x100000);
+
+  CpuInfoInHob = (CPU_INFO_IN_HOB *) (UINTN) CpuMpData->CpuInfoInHob;
+
+  if (ProcessorNumber < 0) {
+    for (Index = 0; Index < CpuMpData->CpuCount; Index++) {
+      if (Index != CpuMpData->BspNumber) {
+        CpuData = &CpuMpData->CpuData[Index];
+        ApicId = CpuInfoInHob[Index].ApicId,
+        SevSnpCreateSaveArea (CpuMpData, CpuData, ApicId);
+      }
+    }
+  } else {
+    Index = (UINTN) ProcessorNumber;
+    CpuData = &CpuMpData->CpuData[Index];
+    ApicId = CpuInfoInHob[ProcessorNumber].ApicId,
+    SevSnpCreateSaveArea (CpuMpData, CpuData, ApicId);
+  }
+}
+
 /**
   This function will be called by BSP to wakeup AP.
 
@@ -1360,7 +1578,7 @@ WakeUpAP (
   ResetVectorRequired = FALSE;
 
   if (CpuMpData->WakeUpByInitSipiSipi ||
-      CpuMpData->InitFlag   != ApInitDone) {
+      CpuMpData->InitFlag != ApInitDone) {
     ResetVectorRequired = TRUE;
     AllocateResetVector (CpuMpData);
     AllocateSevEsAPMemory (CpuMpData);
@@ -1401,7 +1619,7 @@ WakeUpAP (
     }
     if (ResetVectorRequired) {
       //
-      // For SEV-ES, the initial AP boot address will be defined by
+      // For SEV-ES and SEV-SNP, the initial AP boot address will be defined by
       // PcdSevEsWorkAreaBase. The Segment/Rip must be the jump address
       // from the original INIT-SIPI-SIPI.
       //
@@ -1411,8 +1629,14 @@ WakeUpAP (
 
       //
       // Wakeup all APs
+      //   Must use the INIT-SIPI-SIPI method for initial configuration in
+      //   order to obtain the APIC ID.
       //
-      SendInitSipiSipiAllExcludingSelf ((UINT32) ExchangeInfo->BufferStart);
+      if (CpuMpData->SevSnpIsEnabled && CpuMpData->InitFlag != ApInitConfig) {
+        SevSnpCreateAP (CpuMpData, -1);
+      } else {
+        SendInitSipiSipiAllExcludingSelf ((UINT32) ExchangeInfo->BufferStart);
+      }
     }
     if (CpuMpData->InitFlag == ApInitConfig) {
       if (PcdGet32 (PcdCpuBootLogicalProcessorNumber) > 0) {
@@ -1502,7 +1726,7 @@ WakeUpAP (
       CpuInfoInHob = (CPU_INFO_IN_HOB *) (UINTN) CpuMpData->CpuInfoInHob;
 
       //
-      // For SEV-ES, the initial AP boot address will be defined by
+      // For SEV-ES and SEV-SNP, the initial AP boot address will be defined by
       // PcdSevEsWorkAreaBase. The Segment/Rip must be the jump address
       // from the original INIT-SIPI-SIPI.
       //
@@ -1510,10 +1734,14 @@ WakeUpAP (
         SetSevEsJumpTable (ExchangeInfo->BufferStart);
       }
 
-      SendInitSipiSipi (
-        CpuInfoInHob[ProcessorNumber].ApicId,
-        (UINT32) ExchangeInfo->BufferStart
-        );
+      if (CpuMpData->SevSnpIsEnabled && CpuMpData->InitFlag != ApInitConfig) {
+        SevSnpCreateAP (CpuMpData, (INTN) ProcessorNumber);
+      } else {
+        SendInitSipiSipi (
+          CpuInfoInHob[ProcessorNumber].ApicId,
+          (UINT32) ExchangeInfo->BufferStart
+          );
+      }
     }
     //
     // Wait specified AP waken up
@@ -2044,10 +2272,15 @@ MpInitLibInitialize (
   CpuMpData->CpuData          = (CPU_AP_DATA *) (CpuMpData + 1);
   CpuMpData->CpuInfoInHob     = (UINT64) (UINTN) (CpuMpData->CpuData + MaxLogicalProcessorNumber);
   InitializeSpinLock(&CpuMpData->MpLock);
-  CpuMpData->SevEsIsEnabled = PcdGetBool (PcdSevEsIsEnabled);
-  CpuMpData->SevSnpIsEnabled = PcdGetBool (PcdSevSnpIsEnabled);
-  CpuMpData->SevEsAPBuffer  = (UINTN) -1;
-  CpuMpData->GhcbBase       = PcdGet64 (PcdGhcbBase);
+  CpuMpData->SevEsIsEnabled   = PcdGetBool (PcdSevEsIsEnabled);
+  CpuMpData->SevSnpIsEnabled  = PcdGetBool (PcdSevSnpIsEnabled);
+  CpuMpData->SevEsAPBuffer    = (UINTN) -1;
+  CpuMpData->GhcbBase         = PcdGet64 (PcdGhcbBase);
+  CpuMpData->UseSevEsAPMethod = CpuMpData->SevEsIsEnabled && !CpuMpData->SevSnpIsEnabled;
+
+  if (CpuMpData->SevSnpIsEnabled) {
+    ASSERT ((PcdGet64 (PcdGhcbHypervisorFeatures) & GHCB_HV_FEATURES_SNP_AP_CREATE) == GHCB_HV_FEATURES_SNP_AP_CREATE);
+  }
 
   //
   // Make sure no memory usage outside of the allocated buffer.
diff --git a/UefiCpuPkg/Library/MpInitLib/X64/SevSnpRmpAdjustInternal.c b/UefiCpuPkg/Library/MpInitLib/X64/SevSnpRmpAdjustInternal.c
new file mode 100644
index 000000000000..3048742facb5
--- /dev/null
+++ b/UefiCpuPkg/Library/MpInitLib/X64/SevSnpRmpAdjustInternal.c
@@ -0,0 +1,44 @@
+/** @file
+
+  RMPADJUST helper function.
+
+  Copyright (c) 2021, AMD Incorporated. All rights reserved.<BR>
+
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include "MpLib.h"
+
+/**
+  Issue RMPADJUST to adjust the VMSA attribute of an SEV-SNP page.
+
+  @param[in]  PageAddress
+  @param[in]  VmsaPage
+
+  @return  RMPADJUST return value
+**/
+UINT32
+SevSnpRmpAdjust (
+  IN  EFI_PHYSICAL_ADDRESS  PageAddress,
+  IN  BOOLEAN               VmsaPage
+  )
+{
+  UINT64  Rdx;
+
+  //
+  // The RMPADJUST instruction is used to set or clear the VMSA bit for a
+  // page. The VMSA change is only made when running at VMPL0 and is ignored
+  // otherwise. If too low a target VMPL is specified, the instruction can
+  // succeed without changing the VMSA bit when not running at VMPL0. Using a
+  // target VMPL level of 1, RMPADJUST will return a FAIL_PERMISSION error if
+  // not running at VMPL0, thus ensuring that the VMSA bit is set appropriately
+  // when no error is returned.
+  //
+  Rdx = 1;
+  if (VmsaPage) {
+    Rdx |= RMPADJUST_VMSA_PAGE_BIT;
+  }
+
+  return AsmRmpAdjust ((UINT64) PageAddress, 0, Rdx);
+}
-- 
2.17.1