]> git.proxmox.com Git - mirror_edk2.git/blobdiff - UefiCpuPkg/Library/MpInitLib/MpLib.c
UefiCpuPkg: Move AsmRelocateApLoopStart from Mpfuncs.nasm to AmdSev.nasm
[mirror_edk2.git] / UefiCpuPkg / Library / MpInitLib / MpLib.c
index 44a011ba75dea8362ebecc33205257a0f4098367..e5dc852ed95fd23a6fbd6f437cbc1e062e42e969 100644 (file)
@@ -1,7 +1,7 @@
 /** @file\r
   CPU MP Initialize Library common functions.\r
 \r
-  Copyright (c) 2016 - 2021, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2016 - 2022, Intel Corporation. All rights reserved.<BR>\r
   Copyright (c) 2020, AMD Inc. All rights reserved.<BR>\r
 \r
   SPDX-License-Identifier: BSD-2-Clause-Patent\r
@@ -9,12 +9,35 @@
 **/\r
 \r
 #include "MpLib.h"\r
-#include <Library/VmgExitLib.h>\r
+#include <Library/CcExitLib.h>\r
 #include <Register/Amd/Fam17Msr.h>\r
 #include <Register/Amd/Ghcb.h>\r
 \r
 EFI_GUID  mCpuInitMpLibHobGuid = CPU_INIT_MP_LIB_HOB_GUID;\r
 \r
+/**\r
+  Save the volatile registers required to be restored following INIT IPI.\r
+\r
+  @param[out]  VolatileRegisters    Returns buffer saved the volatile resisters\r
+**/\r
+VOID\r
+SaveVolatileRegisters (\r
+  OUT CPU_VOLATILE_REGISTERS  *VolatileRegisters\r
+  );\r
+\r
+/**\r
+  Restore the volatile registers following INIT IPI.\r
+\r
+  @param[in]  VolatileRegisters   Pointer to volatile resisters\r
+  @param[in]  IsRestoreDr         TRUE:  Restore DRx if supported\r
+                                  FALSE: Do not restore DRx\r
+**/\r
+VOID\r
+RestoreVolatileRegisters (\r
+  IN CPU_VOLATILE_REGISTERS  *VolatileRegisters,\r
+  IN BOOLEAN                 IsRestoreDr\r
+  );\r
+\r
 /**\r
   The function will check if BSP Execute Disable is enabled.\r
 \r
@@ -83,7 +106,12 @@ FutureBSPProc (
   CPU_MP_DATA  *DataInHob;\r
 \r
   DataInHob = (CPU_MP_DATA *)Buffer;\r
+  //\r
+  // Save and restore volatile registers when switch BSP\r
+  //\r
+  SaveVolatileRegisters (&DataInHob->APInfo.VolatileRegisters);\r
   AsmExchangeRole (&DataInHob->APInfo, &DataInHob->BSPInfo);\r
+  RestoreVolatileRegisters (&DataInHob->APInfo.VolatileRegisters, FALSE);\r
 }\r
 \r
 /**\r
@@ -295,10 +323,12 @@ GetApLoopMode (
       ApLoopMode = ApInHltLoop;\r
     }\r
 \r
-    if (ConfidentialComputingGuestHas (CCAttrAmdSevEs)) {\r
+    if (ConfidentialComputingGuestHas (CCAttrAmdSevEs) &&\r
+        !ConfidentialComputingGuestHas (CCAttrAmdSevSnp))\r
+    {\r
       //\r
-      // For SEV-ES, force AP in Hlt-loop mode in order to use the GHCB\r
-      // protocol for starting APs\r
+      // For SEV-ES (SEV-SNP is also considered SEV-ES), force AP in Hlt-loop\r
+      // mode in order to use the GHCB protocol for starting APs\r
       //\r
       ApLoopMode = ApInHltLoop;\r
     }\r
@@ -569,6 +599,7 @@ InitializeApData (
 {\r
   CPU_INFO_IN_HOB                *CpuInfoInHob;\r
   MSR_IA32_PLATFORM_ID_REGISTER  PlatformIdMsr;\r
+  AP_STACK_DATA                  *ApStackData;\r
 \r
   CpuInfoInHob                                = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;\r
   CpuInfoInHob[ProcessorNumber].InitialApicId = GetInitialApicId ();\r
@@ -576,6 +607,12 @@ InitializeApData (
   CpuInfoInHob[ProcessorNumber].Health        = BistData;\r
   CpuInfoInHob[ProcessorNumber].ApTopOfStack  = ApTopOfStack;\r
 \r
+  //\r
+  // AP_STACK_DATA is stored at the top of AP Stack\r
+  //\r
+  ApStackData         = (AP_STACK_DATA *)((UINTN)ApTopOfStack - sizeof (AP_STACK_DATA));\r
+  ApStackData->MpData = CpuMpData;\r
+\r
   CpuMpData->CpuData[ProcessorNumber].Waiting    = FALSE;\r
   CpuMpData->CpuData[ProcessorNumber].CpuHealthy = (BistData == 0) ? TRUE : FALSE;\r
 \r
@@ -621,6 +658,7 @@ ApWakeupFunction (
   CPU_INFO_IN_HOB   *CpuInfoInHob;\r
   UINT64            ApTopOfStack;\r
   UINTN             CurrentApicMode;\r
+  AP_STACK_DATA     *ApStackData;\r
 \r
   //\r
   // AP finished assembly code and begin to execute C code\r
@@ -646,7 +684,9 @@ ApWakeupFunction (
       // This is first time AP wakeup, get BIST information from AP stack\r
       //\r
       ApTopOfStack = CpuMpData->Buffer + (ProcessorNumber + 1) * CpuMpData->CpuApStackSize;\r
-      BistData     = *(UINT32 *)((UINTN)ApTopOfStack - sizeof (UINTN));\r
+      ApStackData  = (AP_STACK_DATA *)((UINTN)ApTopOfStack - sizeof (AP_STACK_DATA));\r
+      BistData     = (UINT32)ApStackData->Bist;\r
+\r
       //\r
       // CpuMpData->CpuData[0].VolatileRegisters is initialized based on BSP environment,\r
       //   to initialize AP in InitConfig path.\r
@@ -763,7 +803,7 @@ ApWakeupFunction (
       // to allow the APs to issue an AP_RESET_HOLD before the BSP possibly\r
       // performs another INIT-SIPI-SIPI sequence.\r
       //\r
-      if (!CpuMpData->SevEsIsEnabled) {\r
+      if (!CpuMpData->UseSevEsAPMethod) {\r
         InterlockedDecrement ((UINT32 *)&CpuMpData->MpCpuExchangeInfo->NumApsExecuting);\r
       }\r
     }\r
@@ -777,7 +817,7 @@ ApWakeupFunction (
       //\r
       while (TRUE) {\r
         DisableInterrupts ();\r
-        if (CpuMpData->SevEsIsEnabled) {\r
+        if (CpuMpData->UseSevEsAPMethod) {\r
           SevEsPlaceApHlt (CpuMpData);\r
         } else {\r
           CpuSleep ();\r
@@ -846,6 +886,30 @@ WaitApWakeup (
   }\r
 }\r
 \r
+/**\r
+  Calculate the size of the reset vector.\r
+\r
+  @param[in]  AddressMap   The pointer to Address Map structure.\r
+  @param[out] SizeBelow1Mb Return the size of below 1MB memory for AP reset area.\r
+  @param[out] SizeAbove1Mb Return the size of abvoe 1MB memory for AP reset area.\r
+**/\r
+STATIC\r
+VOID\r
+GetApResetVectorSize (\r
+  IN  MP_ASSEMBLY_ADDRESS_MAP  *AddressMap,\r
+  OUT UINTN                    *SizeBelow1Mb OPTIONAL,\r
+  OUT UINTN                    *SizeAbove1Mb OPTIONAL\r
+  )\r
+{\r
+  if (SizeBelow1Mb != NULL) {\r
+    *SizeBelow1Mb = AddressMap->ModeTransitionOffset + sizeof (MP_CPU_EXCHANGE_INFO);\r
+  }\r
+\r
+  if (SizeAbove1Mb != NULL) {\r
+    *SizeAbove1Mb = AddressMap->RendezvousFunnelSize - AddressMap->ModeTransitionOffset;\r
+  }\r
+}\r
+\r
 /**\r
   This function will fill the exchange info structure.\r
 \r
@@ -900,6 +964,13 @@ FillExchangeInfoData (
   ExchangeInfo->SevSnpIsEnabled = CpuMpData->SevSnpIsEnabled;\r
   ExchangeInfo->GhcbBase        = (UINTN)CpuMpData->GhcbBase;\r
 \r
+  //\r
+  // Populate SEV-ES specific exchange data.\r
+  //\r
+  if (ExchangeInfo->SevSnpIsEnabled) {\r
+    FillExchangeInfoDataSevEs (ExchangeInfo);\r
+  }\r
+\r
   //\r
   // Get the BSP's data of GDT and IDT\r
   //\r
@@ -922,26 +993,7 @@ FillExchangeInfoData (
     Size     -= sizeof (IA32_SEGMENT_DESCRIPTOR);\r
   }\r
 \r
-  //\r
-  // Copy all 32-bit code and 64-bit code into memory with type of\r
-  // EfiBootServicesCode to avoid page fault if NX memory protection is enabled.\r
-  //\r
-  if (CpuMpData->WakeupBufferHigh != 0) {\r
-    Size = CpuMpData->AddressMap.RendezvousFunnelSize +\r
-           CpuMpData->AddressMap.SwitchToRealSize -\r
-           CpuMpData->AddressMap.ModeTransitionOffset;\r
-    CopyMem (\r
-      (VOID *)CpuMpData->WakeupBufferHigh,\r
-      CpuMpData->AddressMap.RendezvousFunnelAddress +\r
-      CpuMpData->AddressMap.ModeTransitionOffset,\r
-      Size\r
-      );\r
-\r
-    ExchangeInfo->ModeTransitionMemory = (UINT32)CpuMpData->WakeupBufferHigh;\r
-  } else {\r
-    ExchangeInfo->ModeTransitionMemory = (UINT32)\r
-                                         (ExchangeInfo->BufferStart + CpuMpData->AddressMap.ModeTransitionOffset);\r
-  }\r
+  ExchangeInfo->ModeTransitionMemory = (UINT32)CpuMpData->WakeupBufferHigh;\r
 \r
   ExchangeInfo->ModeHighMemory = ExchangeInfo->ModeTransitionMemory +\r
                                  (UINT32)ExchangeInfo->ModeOffset -\r
@@ -982,8 +1034,7 @@ BackupAndPrepareWakeupBuffer (
   CopyMem (\r
     (VOID *)CpuMpData->WakeupBuffer,\r
     (VOID *)CpuMpData->AddressMap.RendezvousFunnelAddress,\r
-    CpuMpData->AddressMap.RendezvousFunnelSize +\r
-    CpuMpData->AddressMap.SwitchToRealSize\r
+    CpuMpData->BackupBufferSize - sizeof (MP_CPU_EXCHANGE_INFO)\r
     );\r
 }\r
 \r
@@ -1004,59 +1055,39 @@ RestoreWakeupBuffer (
     );\r
 }\r
 \r
-/**\r
-  Calculate the size of the reset vector.\r
-\r
-  @param[in]  AddressMap  The pointer to Address Map structure.\r
-\r
-  @return                 Total amount of memory required for the AP reset area\r
-**/\r
-STATIC\r
-UINTN\r
-GetApResetVectorSize (\r
-  IN MP_ASSEMBLY_ADDRESS_MAP  *AddressMap\r
-  )\r
-{\r
-  UINTN  Size;\r
-\r
-  Size = AddressMap->RendezvousFunnelSize +\r
-         AddressMap->SwitchToRealSize +\r
-         sizeof (MP_CPU_EXCHANGE_INFO);\r
-\r
-  return Size;\r
-}\r
-\r
 /**\r
   Allocate reset vector buffer.\r
 \r
   @param[in, out]  CpuMpData  The pointer to CPU MP Data structure.\r
 **/\r
 VOID\r
-AllocateResetVector (\r
+AllocateResetVectorBelow1Mb (\r
   IN OUT CPU_MP_DATA  *CpuMpData\r
   )\r
 {\r
-  UINTN  ApResetVectorSize;\r
   UINTN  ApResetStackSize;\r
 \r
   if (CpuMpData->WakeupBuffer == (UINTN)-1) {\r
-    ApResetVectorSize = GetApResetVectorSize (&CpuMpData->AddressMap);\r
-\r
-    CpuMpData->WakeupBuffer      = GetWakeupBuffer (ApResetVectorSize);\r
+    CpuMpData->WakeupBuffer      = GetWakeupBuffer (CpuMpData->BackupBufferSize);\r
     CpuMpData->MpCpuExchangeInfo = (MP_CPU_EXCHANGE_INFO *)(UINTN)\r
-                                   (CpuMpData->WakeupBuffer +\r
-                                    CpuMpData->AddressMap.RendezvousFunnelSize +\r
-                                    CpuMpData->AddressMap.SwitchToRealSize);\r
-    CpuMpData->WakeupBufferHigh = GetModeTransitionBuffer (\r
-                                    CpuMpData->AddressMap.RendezvousFunnelSize +\r
-                                    CpuMpData->AddressMap.SwitchToRealSize -\r
-                                    CpuMpData->AddressMap.ModeTransitionOffset\r
-                                    );\r
+                                   (CpuMpData->WakeupBuffer + CpuMpData->BackupBufferSize - sizeof (MP_CPU_EXCHANGE_INFO));\r
+    DEBUG ((\r
+      DEBUG_INFO,\r
+      "AP Vector: 16-bit = %p/%x, ExchangeInfo = %p/%x\n",\r
+      CpuMpData->WakeupBuffer,\r
+      CpuMpData->BackupBufferSize - sizeof (MP_CPU_EXCHANGE_INFO),\r
+      CpuMpData->MpCpuExchangeInfo,\r
+      sizeof (MP_CPU_EXCHANGE_INFO)\r
+      ));\r
     //\r
     // The AP reset stack is only used by SEV-ES guests. Do not allocate it\r
-    // if SEV-ES is not enabled.\r
+    // if SEV-ES is not enabled. An SEV-SNP guest is also considered\r
+    // an SEV-ES guest, but uses a different method of AP startup, eliminating\r
+    // the need for the allocation.\r
     //\r
-    if (ConfidentialComputingGuestHas (CCAttrAmdSevEs)) {\r
+    if (ConfidentialComputingGuestHas (CCAttrAmdSevEs) &&\r
+        !ConfidentialComputingGuestHas (CCAttrAmdSevSnp))\r
+    {\r
       //\r
       // Stack location is based on ProcessorNumber, so use the total number\r
       // of processors for calculating the total stack area.\r
@@ -1107,7 +1138,7 @@ FreeResetVector (
   // perform the restore as this will overwrite memory which has data\r
   // needed by SEV-ES.\r
   //\r
-  if (!CpuMpData->SevEsIsEnabled) {\r
+  if (!CpuMpData->UseSevEsAPMethod) {\r
     RestoreWakeupBuffer (CpuMpData);\r
   }\r
 }\r
@@ -1146,7 +1177,7 @@ WakeUpAP (
       (CpuMpData->InitFlag   != ApInitDone))\r
   {\r
     ResetVectorRequired = TRUE;\r
-    AllocateResetVector (CpuMpData);\r
+    AllocateResetVectorBelow1Mb (CpuMpData);\r
     AllocateSevEsAPMemory (CpuMpData);\r
     FillExchangeInfoData (CpuMpData);\r
     SaveLocalApicTimerSetting (CpuMpData);\r
@@ -1186,7 +1217,7 @@ WakeUpAP (
 \r
     if (ResetVectorRequired) {\r
       //\r
-      // For SEV-ES, the initial AP boot address will be defined by\r
+      // For SEV-ES and SEV-SNP, the initial AP boot address will be defined by\r
       // PcdSevEsWorkAreaBase. The Segment/Rip must be the jump address\r
       // from the original INIT-SIPI-SIPI.\r
       //\r
@@ -1196,8 +1227,14 @@ WakeUpAP (
 \r
       //\r
       // Wakeup all APs\r
+      //   Must use the INIT-SIPI-SIPI method for initial configuration in\r
+      //   order to obtain the APIC ID.\r
       //\r
-      SendInitSipiSipiAllExcludingSelf ((UINT32)ExchangeInfo->BufferStart);\r
+      if (CpuMpData->SevSnpIsEnabled && (CpuMpData->InitFlag != ApInitConfig)) {\r
+        SevSnpCreateAP (CpuMpData, -1);\r
+      } else {\r
+        SendInitSipiSipiAllExcludingSelf ((UINT32)ExchangeInfo->BufferStart);\r
+      }\r
     }\r
 \r
     if (CpuMpData->InitFlag == ApInitConfig) {\r
@@ -1288,7 +1325,7 @@ WakeUpAP (
       CpuInfoInHob = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;\r
 \r
       //\r
-      // For SEV-ES, the initial AP boot address will be defined by\r
+      // For SEV-ES and SEV-SNP, the initial AP boot address will be defined by\r
       // PcdSevEsWorkAreaBase. The Segment/Rip must be the jump address\r
       // from the original INIT-SIPI-SIPI.\r
       //\r
@@ -1296,10 +1333,14 @@ WakeUpAP (
         SetSevEsJumpTable (ExchangeInfo->BufferStart);\r
       }\r
 \r
-      SendInitSipiSipi (\r
-        CpuInfoInHob[ProcessorNumber].ApicId,\r
-        (UINT32)ExchangeInfo->BufferStart\r
-        );\r
+      if (CpuMpData->SevSnpIsEnabled && (CpuMpData->InitFlag != ApInitConfig)) {\r
+        SevSnpCreateAP (CpuMpData, (INTN)ProcessorNumber);\r
+      } else {\r
+        SendInitSipiSipi (\r
+          CpuInfoInHob[ProcessorNumber].ApicId,\r
+          (UINT32)ExchangeInfo->BufferStart\r
+          );\r
+      }\r
     }\r
 \r
     //\r
@@ -1776,7 +1817,8 @@ MpInitLibInitialize (
   UINT8                    ApLoopMode;\r
   UINT8                    *MonitorBuffer;\r
   UINTN                    Index;\r
-  UINTN                    ApResetVectorSize;\r
+  UINTN                    ApResetVectorSizeBelow1Mb;\r
+  UINTN                    ApResetVectorSizeAbove1Mb;\r
   UINTN                    BackupBufferAddr;\r
   UINTN                    ApIdtBase;\r
 \r
@@ -1790,18 +1832,26 @@ MpInitLibInitialize (
   ASSERT (MaxLogicalProcessorNumber != 0);\r
 \r
   AsmGetAddressMap (&AddressMap);\r
-  ApResetVectorSize = GetApResetVectorSize (&AddressMap);\r
-  ApStackSize       = PcdGet32 (PcdCpuApStackSize);\r
-  ApLoopMode        = GetApLoopMode (&MonitorFilterSize);\r
+  GetApResetVectorSize (&AddressMap, &ApResetVectorSizeBelow1Mb, &ApResetVectorSizeAbove1Mb);\r
+  ApStackSize = PcdGet32 (PcdCpuApStackSize);\r
+  //\r
+  // ApStackSize must be power of 2\r
+  //\r
+  ASSERT ((ApStackSize & (ApStackSize - 1)) == 0);\r
+  ApLoopMode = GetApLoopMode (&MonitorFilterSize);\r
 \r
   //\r
   // Save BSP's Control registers for APs.\r
   //\r
   SaveVolatileRegisters (&VolatileRegisters);\r
 \r
-  BufferSize  = ApStackSize * MaxLogicalProcessorNumber;\r
+  BufferSize = ApStackSize * MaxLogicalProcessorNumber;\r
+  //\r
+  // Allocate extra ApStackSize to let AP stack align on ApStackSize bounday\r
+  //\r
+  BufferSize += ApStackSize;\r
   BufferSize += MonitorFilterSize * MaxLogicalProcessorNumber;\r
-  BufferSize += ApResetVectorSize;\r
+  BufferSize += ApResetVectorSizeBelow1Mb;\r
   BufferSize  = ALIGN_VALUE (BufferSize, 8);\r
   BufferSize += VolatileRegisters.Idtr.Limit + 1;\r
   BufferSize += sizeof (CPU_MP_DATA);\r
@@ -1809,13 +1859,13 @@ MpInitLibInitialize (
   MpBuffer    = AllocatePages (EFI_SIZE_TO_PAGES (BufferSize));\r
   ASSERT (MpBuffer != NULL);\r
   ZeroMem (MpBuffer, BufferSize);\r
-  Buffer = (UINTN)MpBuffer;\r
+  Buffer = ALIGN_VALUE ((UINTN)MpBuffer, ApStackSize);\r
 \r
   //\r
-  //  The layout of the Buffer is as below:\r
+  //  The layout of the Buffer is as below (lower address on top):\r
   //\r
-  //    +--------------------+ <-- Buffer\r
-  //        AP Stacks (N)\r
+  //    +--------------------+ <-- Buffer (Pointer of CpuMpData is stored in the top of each AP's stack.)\r
+  //        AP Stacks (N)                 (StackTop = (RSP + ApStackSize) & ~ApStackSize))\r
   //    +--------------------+ <-- MonitorBuffer\r
   //    AP Monitor Filters (N)\r
   //    +--------------------+ <-- BackupBufferAddr (CpuMpData->BackupBuffer)\r
@@ -1823,7 +1873,7 @@ MpInitLibInitialize (
   //    +--------------------+\r
   //           Padding\r
   //    +--------------------+ <-- ApIdtBase (8-byte boundary)\r
-  //           AP IDT          All APs share one separate IDT. So AP can get address of CPU_MP_DATA from IDT Base.\r
+  //           AP IDT          All APs share one separate IDT.\r
   //    +--------------------+ <-- CpuMpData\r
   //         CPU_MP_DATA\r
   //    +--------------------+ <-- CpuMpData->CpuData\r
@@ -1834,12 +1884,12 @@ MpInitLibInitialize (
   //\r
   MonitorBuffer               = (UINT8 *)(Buffer + ApStackSize * MaxLogicalProcessorNumber);\r
   BackupBufferAddr            = (UINTN)MonitorBuffer + MonitorFilterSize * MaxLogicalProcessorNumber;\r
-  ApIdtBase                   = ALIGN_VALUE (BackupBufferAddr + ApResetVectorSize, 8);\r
+  ApIdtBase                   = ALIGN_VALUE (BackupBufferAddr + ApResetVectorSizeBelow1Mb, 8);\r
   CpuMpData                   = (CPU_MP_DATA *)(ApIdtBase + VolatileRegisters.Idtr.Limit + 1);\r
   CpuMpData->Buffer           = Buffer;\r
   CpuMpData->CpuApStackSize   = ApStackSize;\r
   CpuMpData->BackupBuffer     = BackupBufferAddr;\r
-  CpuMpData->BackupBufferSize = ApResetVectorSize;\r
+  CpuMpData->BackupBufferSize = ApResetVectorSizeBelow1Mb;\r
   CpuMpData->WakeupBuffer     = (UINTN)-1;\r
   CpuMpData->CpuCount         = 1;\r
   CpuMpData->BspNumber        = 0;\r
@@ -1848,17 +1898,23 @@ MpInitLibInitialize (
   CpuMpData->CpuData          = (CPU_AP_DATA *)(CpuMpData + 1);\r
   CpuMpData->CpuInfoInHob     = (UINT64)(UINTN)(CpuMpData->CpuData + MaxLogicalProcessorNumber);\r
   InitializeSpinLock (&CpuMpData->MpLock);\r
-  CpuMpData->SevEsIsEnabled  = ConfidentialComputingGuestHas (CCAttrAmdSevEs);\r
-  CpuMpData->SevSnpIsEnabled = ConfidentialComputingGuestHas (CCAttrAmdSevSnp);\r
-  CpuMpData->SevEsAPBuffer   = (UINTN)-1;\r
-  CpuMpData->GhcbBase        = PcdGet64 (PcdGhcbBase);\r
+  CpuMpData->SevEsIsEnabled   = ConfidentialComputingGuestHas (CCAttrAmdSevEs);\r
+  CpuMpData->SevSnpIsEnabled  = ConfidentialComputingGuestHas (CCAttrAmdSevSnp);\r
+  CpuMpData->SevEsAPBuffer    = (UINTN)-1;\r
+  CpuMpData->GhcbBase         = PcdGet64 (PcdGhcbBase);\r
+  CpuMpData->UseSevEsAPMethod = CpuMpData->SevEsIsEnabled && !CpuMpData->SevSnpIsEnabled;\r
+\r
+  if (CpuMpData->SevSnpIsEnabled) {\r
+    ASSERT ((PcdGet64 (PcdGhcbHypervisorFeatures) & GHCB_HV_FEATURES_SNP_AP_CREATE) == GHCB_HV_FEATURES_SNP_AP_CREATE);\r
+  }\r
 \r
   //\r
   // Make sure no memory usage outside of the allocated buffer.\r
+  // (ApStackSize - (Buffer - (UINTN)MpBuffer)) is the redundant caused by alignment\r
   //\r
   ASSERT (\r
     (CpuMpData->CpuInfoInHob + sizeof (CPU_INFO_IN_HOB) * MaxLogicalProcessorNumber) ==\r
-    Buffer + BufferSize\r
+    (UINTN)MpBuffer + BufferSize - (ApStackSize - Buffer + (UINTN)MpBuffer)\r
     );\r
 \r
   //\r
@@ -1896,6 +1952,19 @@ MpInitLibInitialize (
       (UINT32 *)(MonitorBuffer + MonitorFilterSize * Index);\r
   }\r
 \r
+  //\r
+  // Copy all 32-bit code and 64-bit code into memory with type of\r
+  // EfiBootServicesCode to avoid page fault if NX memory protection is enabled.\r
+  //\r
+  CpuMpData->WakeupBufferHigh = AllocateCodeBuffer (ApResetVectorSizeAbove1Mb);\r
+  CopyMem (\r
+    (VOID *)CpuMpData->WakeupBufferHigh,\r
+    CpuMpData->AddressMap.RendezvousFunnelAddress +\r
+    CpuMpData->AddressMap.ModeTransitionOffset,\r
+    ApResetVectorSizeAbove1Mb\r
+    );\r
+  DEBUG ((DEBUG_INFO, "AP Vector: non-16-bit = %p/%x\n", CpuMpData->WakeupBufferHigh, ApResetVectorSizeAbove1Mb));\r
+\r
   //\r
   // Enable the local APIC for Virtual Wire Mode.\r
   //\r
@@ -2211,7 +2280,12 @@ SwitchBSPWorker (
   //\r
   WakeUpAP (CpuMpData, FALSE, ProcessorNumber, FutureBSPProc, CpuMpData, TRUE);\r
 \r
+  //\r
+  // Save and restore volatile registers when switch BSP\r
+  //\r
+  SaveVolatileRegisters (&CpuMpData->BSPInfo.VolatileRegisters);\r
   AsmExchangeRole (&CpuMpData->BSPInfo, &CpuMpData->APInfo);\r
+  RestoreVolatileRegisters (&CpuMpData->BSPInfo.VolatileRegisters, FALSE);\r
 \r
   //\r
   // Set the BSP bit of MSR_IA32_APIC_BASE on new BSP\r