]> 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 ff09a0e9e7e843f2332a0d78fb7b24a138414ee5..e5dc852ed95fd23a6fbd6f437cbc1e062e42e969 100644 (file)
@@ -1,20 +1,42 @@
 /** @file\r
   CPU MP Initialize Library common functions.\r
 \r
-  Copyright (c) 2016 - 2018, Intel Corporation. All rights reserved.<BR>\r
-  This program and the accompanying materials\r
-  are licensed and made available under the terms and conditions of the BSD License\r
-  which accompanies this distribution.  The full text of the license may be found at\r
-  http://opensource.org/licenses/bsd-license.php\r
+  Copyright (c) 2016 - 2022, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2020, AMD Inc. All rights reserved.<BR>\r
 \r
-  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,\r
-  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r
+  SPDX-License-Identifier: BSD-2-Clause-Patent\r
 \r
 **/\r
 \r
 #include "MpLib.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
+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
@@ -38,7 +60,7 @@ IsBspExecuteDisableEnabled (
   BOOLEAN                     Enabled;\r
   IA32_CR0                    Cr0;\r
 \r
-  Enabled = FALSE;\r
+  Enabled   = FALSE;\r
   Cr0.UintN = AsmReadCr0 ();\r
   if (Cr0.Bits.PG != 0) {\r
     //\r
@@ -78,13 +100,18 @@ IsBspExecuteDisableEnabled (
 VOID\r
 EFIAPI\r
 FutureBSPProc (\r
-  IN  VOID            *Buffer\r
+  IN  VOID  *Buffer\r
   )\r
 {\r
-  CPU_MP_DATA         *DataInHob;\r
+  CPU_MP_DATA  *DataInHob;\r
 \r
-  DataInHob = (CPU_MP_DATA *) Buffer;\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
@@ -96,7 +123,7 @@ FutureBSPProc (
 **/\r
 CPU_STATE\r
 GetApState (\r
-  IN  CPU_AP_DATA     *CpuData\r
+  IN  CPU_AP_DATA  *CpuData\r
   )\r
 {\r
   return CpuData->State;\r
@@ -110,8 +137,8 @@ GetApState (
 **/\r
 VOID\r
 SetApState (\r
-  IN  CPU_AP_DATA     *CpuData,\r
-  IN  CPU_STATE       State\r
+  IN  CPU_AP_DATA  *CpuData,\r
+  IN  CPU_STATE    State\r
   )\r
 {\r
   AcquireSpinLock (&CpuData->ApLock);\r
@@ -126,7 +153,7 @@ SetApState (
 **/\r
 VOID\r
 SaveLocalApicTimerSetting (\r
-  IN CPU_MP_DATA   *CpuMpData\r
+  IN CPU_MP_DATA  *CpuMpData\r
   )\r
 {\r
   //\r
@@ -148,7 +175,7 @@ SaveLocalApicTimerSetting (
 **/\r
 VOID\r
 SyncLocalApicTimerSetting (\r
-  IN CPU_MP_DATA   *CpuMpData\r
+  IN CPU_MP_DATA  *CpuMpData\r
   )\r
 {\r
   //\r
@@ -173,10 +200,10 @@ SyncLocalApicTimerSetting (
 **/\r
 VOID\r
 SaveVolatileRegisters (\r
-  OUT CPU_VOLATILE_REGISTERS    *VolatileRegisters\r
+  OUT CPU_VOLATILE_REGISTERS  *VolatileRegisters\r
   )\r
 {\r
-  CPUID_VERSION_INFO_EDX        VersionInfoEdx;\r
+  CPUID_VERSION_INFO_EDX  VersionInfoEdx;\r
 \r
   VolatileRegisters->Cr0 = AsmReadCr0 ();\r
   VolatileRegisters->Cr3 = AsmReadCr3 ();\r
@@ -210,16 +237,16 @@ SaveVolatileRegisters (
 **/\r
 VOID\r
 RestoreVolatileRegisters (\r
-  IN CPU_VOLATILE_REGISTERS    *VolatileRegisters,\r
-  IN BOOLEAN                   IsRestoreDr\r
+  IN CPU_VOLATILE_REGISTERS  *VolatileRegisters,\r
+  IN BOOLEAN                 IsRestoreDr\r
   )\r
 {\r
-  CPUID_VERSION_INFO_EDX        VersionInfoEdx;\r
-  IA32_TSS_DESCRIPTOR           *Tss;\r
+  CPUID_VERSION_INFO_EDX  VersionInfoEdx;\r
+  IA32_TSS_DESCRIPTOR     *Tss;\r
 \r
-  AsmWriteCr0 (VolatileRegisters->Cr0);\r
   AsmWriteCr3 (VolatileRegisters->Cr3);\r
   AsmWriteCr4 (VolatileRegisters->Cr4);\r
+  AsmWriteCr0 (VolatileRegisters->Cr0);\r
 \r
   if (IsRestoreDr) {\r
     AsmCpuid (CPUID_VERSION_INFO, NULL, NULL, NULL, &VersionInfoEdx.Uint32);\r
@@ -239,8 +266,9 @@ RestoreVolatileRegisters (
 \r
   AsmWriteGdtr (&VolatileRegisters->Gdtr);\r
   AsmWriteIdtr (&VolatileRegisters->Idtr);\r
-  if (VolatileRegisters->Tr != 0 &&\r
-      VolatileRegisters->Tr < VolatileRegisters->Gdtr.Limit) {\r
+  if ((VolatileRegisters->Tr != 0) &&\r
+      (VolatileRegisters->Tr < VolatileRegisters->Gdtr.Limit))\r
+  {\r
     Tss = (IA32_TSS_DESCRIPTOR *)(VolatileRegisters->Gdtr.Base +\r
                                   VolatileRegisters->Tr);\r
     if (Tss->Bits.P == 1) {\r
@@ -261,7 +289,7 @@ IsMwaitSupport (
   VOID\r
   )\r
 {\r
-  CPUID_VERSION_INFO_ECX        VersionInfoEcx;\r
+  CPUID_VERSION_INFO_ECX  VersionInfoEcx;\r
 \r
   AsmCpuid (CPUID_VERSION_INFO, NULL, NULL, &VersionInfoEcx.Uint32, NULL);\r
   return (VersionInfoEcx.Bits.MONITOR == 1) ? TRUE : FALSE;\r
@@ -276,11 +304,11 @@ IsMwaitSupport (
 **/\r
 UINT8\r
 GetApLoopMode (\r
-  OUT UINT32     *MonitorFilterSize\r
+  OUT UINT32  *MonitorFilterSize\r
   )\r
 {\r
-  UINT8                         ApLoopMode;\r
-  CPUID_MONITOR_MWAIT_EBX       MonitorMwaitEbx;\r
+  UINT8                    ApLoopMode;\r
+  CPUID_MONITOR_MWAIT_EBX  MonitorMwaitEbx;\r
 \r
   ASSERT (MonitorFilterSize != NULL);\r
 \r
@@ -294,6 +322,16 @@ GetApLoopMode (
       //\r
       ApLoopMode = ApInHltLoop;\r
     }\r
+\r
+    if (ConfidentialComputingGuestHas (CCAttrAmdSevEs) &&\r
+        !ConfidentialComputingGuestHas (CCAttrAmdSevSnp))\r
+    {\r
+      //\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
   }\r
 \r
   if (ApLoopMode != ApInMwaitLoop) {\r
@@ -320,20 +358,20 @@ GetApLoopMode (
 **/\r
 VOID\r
 SortApicId (\r
-  IN CPU_MP_DATA   *CpuMpData\r
+  IN CPU_MP_DATA  *CpuMpData\r
   )\r
 {\r
-  UINTN             Index1;\r
-  UINTN             Index2;\r
-  UINTN             Index3;\r
-  UINT32            ApicId;\r
-  CPU_INFO_IN_HOB   CpuInfo;\r
-  UINT32            ApCount;\r
-  CPU_INFO_IN_HOB   *CpuInfoInHob;\r
-  volatile UINT32   *StartupApSignal;\r
-\r
-  ApCount = CpuMpData->CpuCount - 1;\r
-  CpuInfoInHob = (CPU_INFO_IN_HOB *) (UINTN) CpuMpData->CpuInfoInHob;\r
+  UINTN            Index1;\r
+  UINTN            Index2;\r
+  UINTN            Index3;\r
+  UINT32           ApicId;\r
+  CPU_INFO_IN_HOB  CpuInfo;\r
+  UINT32           ApCount;\r
+  CPU_INFO_IN_HOB  *CpuInfoInHob;\r
+  volatile UINT32  *StartupApSignal;\r
+\r
+  ApCount      = CpuMpData->CpuCount - 1;\r
+  CpuInfoInHob = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;\r
   if (ApCount != 0) {\r
     for (Index1 = 0; Index1 < ApCount; Index1++) {\r
       Index3 = Index1;\r
@@ -347,6 +385,7 @@ SortApicId (
           ApicId = CpuInfoInHob[Index2].ApicId;\r
         }\r
       }\r
+\r
       if (Index3 != Index1) {\r
         CopyMem (&CpuInfo, &CpuInfoInHob[Index3], sizeof (CPU_INFO_IN_HOB));\r
         CopyMem (\r
@@ -359,7 +398,7 @@ SortApicId (
         //\r
         // Also exchange the StartupApSignal.\r
         //\r
-        StartupApSignal = CpuMpData->CpuData[Index3].StartupApSignal;\r
+        StartupApSignal                            = CpuMpData->CpuData[Index3].StartupApSignal;\r
         CpuMpData->CpuData[Index3].StartupApSignal =\r
           CpuMpData->CpuData[Index1].StartupApSignal;\r
         CpuMpData->CpuData[Index1].StartupApSignal = StartupApSignal;\r
@@ -372,7 +411,7 @@ SortApicId (
     ApicId = GetInitialApicId ();\r
     for (Index1 = 0; Index1 < CpuMpData->CpuCount; Index1++) {\r
       if (CpuInfoInHob[Index1].ApicId == ApicId) {\r
-        CpuMpData->BspNumber = (UINT32) Index1;\r
+        CpuMpData->BspNumber = (UINT32)Index1;\r
         break;\r
       }\r
     }\r
@@ -405,12 +444,16 @@ ApInitializeSync (
   )\r
 {\r
   CPU_MP_DATA  *CpuMpData;\r
+  UINTN        ProcessorNumber;\r
+  EFI_STATUS   Status;\r
 \r
-  CpuMpData = (CPU_MP_DATA *) Buffer;\r
+  CpuMpData = (CPU_MP_DATA *)Buffer;\r
+  Status    = GetProcessorNumber (CpuMpData, &ProcessorNumber);\r
+  ASSERT_EFI_ERROR (Status);\r
   //\r
   // Load microcode on AP\r
   //\r
-  MicrocodeDetect (CpuMpData, FALSE);\r
+  MicrocodeDetect (CpuMpData, ProcessorNumber);\r
   //\r
   // Sync BSP's MTRR table to AP\r
   //\r
@@ -428,20 +471,20 @@ ApInitializeSync (
 **/\r
 EFI_STATUS\r
 GetProcessorNumber (\r
-  IN CPU_MP_DATA               *CpuMpData,\r
-  OUT UINTN                    *ProcessorNumber\r
+  IN CPU_MP_DATA  *CpuMpData,\r
+  OUT UINTN       *ProcessorNumber\r
   )\r
 {\r
-  UINTN                   TotalProcessorNumber;\r
-  UINTN                   Index;\r
-  CPU_INFO_IN_HOB         *CpuInfoInHob;\r
-  UINT32                  CurrentApicId;\r
+  UINTN            TotalProcessorNumber;\r
+  UINTN            Index;\r
+  CPU_INFO_IN_HOB  *CpuInfoInHob;\r
+  UINT32           CurrentApicId;\r
 \r
-  CpuInfoInHob = (CPU_INFO_IN_HOB *) (UINTN) CpuMpData->CpuInfoInHob;\r
+  CpuInfoInHob = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;\r
 \r
   TotalProcessorNumber = CpuMpData->CpuCount;\r
-  CurrentApicId = GetApicId ();\r
-  for (Index = 0; Index < TotalProcessorNumber; Index ++) {\r
+  CurrentApicId        = GetApicId ();\r
+  for (Index = 0; Index < TotalProcessorNumber; Index++) {\r
     if (CpuInfoInHob[Index].ApicId == CurrentApicId) {\r
       *ProcessorNumber = Index;\r
       return EFI_SUCCESS;\r
@@ -460,44 +503,60 @@ GetProcessorNumber (
 **/\r
 UINTN\r
 CollectProcessorCount (\r
-  IN CPU_MP_DATA         *CpuMpData\r
+  IN CPU_MP_DATA  *CpuMpData\r
   )\r
 {\r
-  UINTN                  Index;\r
+  UINTN            Index;\r
+  CPU_INFO_IN_HOB  *CpuInfoInHob;\r
+  BOOLEAN          X2Apic;\r
 \r
   //\r
   // Send 1st broadcast IPI to APs to wakeup APs\r
   //\r
-  CpuMpData->InitFlag     = ApInitConfig;\r
-  CpuMpData->X2ApicEnable = FALSE;\r
-  WakeUpAP (CpuMpData, TRUE, 0, NULL, NULL);\r
+  CpuMpData->InitFlag = ApInitConfig;\r
+  WakeUpAP (CpuMpData, TRUE, 0, NULL, NULL, TRUE);\r
   CpuMpData->InitFlag = ApInitDone;\r
-  ASSERT (CpuMpData->CpuCount <= PcdGet32 (PcdCpuMaxLogicalProcessorNumber));\r
   //\r
-  // Wait for all APs finished the initialization\r
+  // When InitFlag == ApInitConfig, WakeUpAP () guarantees all APs are checked in.\r
+  // FinishedCount is the number of check-in APs.\r
   //\r
-  while (CpuMpData->FinishedCount < (CpuMpData->CpuCount - 1)) {\r
-    CpuPause ();\r
-  }\r
+  CpuMpData->CpuCount = CpuMpData->FinishedCount + 1;\r
+  ASSERT (CpuMpData->CpuCount <= PcdGet32 (PcdCpuMaxLogicalProcessorNumber));\r
 \r
+  //\r
+  // Enable x2APIC mode if\r
+  //  1. Number of CPU is greater than 255; or\r
+  //  2. There are any logical processors reporting an Initial APIC ID of 255 or greater.\r
+  //\r
+  X2Apic = FALSE;\r
   if (CpuMpData->CpuCount > 255) {\r
     //\r
     // If there are more than 255 processor found, force to enable X2APIC\r
     //\r
-    CpuMpData->X2ApicEnable = TRUE;\r
+    X2Apic = TRUE;\r
+  } else {\r
+    CpuInfoInHob = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;\r
+    for (Index = 0; Index < CpuMpData->CpuCount; Index++) {\r
+      if (CpuInfoInHob[Index].InitialApicId >= 0xFF) {\r
+        X2Apic = TRUE;\r
+        break;\r
+      }\r
+    }\r
   }\r
-  if (CpuMpData->X2ApicEnable) {\r
+\r
+  if (X2Apic) {\r
     DEBUG ((DEBUG_INFO, "Force x2APIC mode!\n"));\r
     //\r
     // Wakeup all APs to enable x2APIC mode\r
     //\r
-    WakeUpAP (CpuMpData, TRUE, 0, ApFuncEnableX2Apic, NULL);\r
+    WakeUpAP (CpuMpData, TRUE, 0, ApFuncEnableX2Apic, NULL, TRUE);\r
     //\r
     // Wait for all known APs finished\r
     //\r
     while (CpuMpData->FinishedCount < (CpuMpData->CpuCount - 1)) {\r
       CpuPause ();\r
     }\r
+\r
     //\r
     // Enable x2APIC on BSP\r
     //\r
@@ -509,6 +568,7 @@ CollectProcessorCount (
       SetApState (&CpuMpData->CpuData[Index], CpuStateIdle);\r
     }\r
   }\r
+\r
   DEBUG ((DEBUG_INFO, "APIC MODE is %d\n", GetApicMode ()));\r
   //\r
   // Sort BSP/Aps by CPU APIC ID in ascending order\r
@@ -531,33 +591,48 @@ CollectProcessorCount (
 **/\r
 VOID\r
 InitializeApData (\r
-  IN OUT CPU_MP_DATA      *CpuMpData,\r
-  IN     UINTN            ProcessorNumber,\r
-  IN     UINT32           BistData,\r
-  IN     UINT64           ApTopOfStack\r
+  IN OUT CPU_MP_DATA  *CpuMpData,\r
+  IN     UINTN        ProcessorNumber,\r
+  IN     UINT32       BistData,\r
+  IN     UINT64       ApTopOfStack\r
   )\r
 {\r
-  CPU_INFO_IN_HOB          *CpuInfoInHob;\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                                = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;\r
   CpuInfoInHob[ProcessorNumber].InitialApicId = GetInitialApicId ();\r
   CpuInfoInHob[ProcessorNumber].ApicId        = GetApicId ();\r
   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
-  if (CpuInfoInHob[ProcessorNumber].InitialApicId >= 0xFF) {\r
-    //\r
-    // Set x2APIC mode if there are any logical processor reporting\r
-    // an Initial APIC ID of 255 or greater.\r
-    //\r
-    AcquireSpinLock(&CpuMpData->MpLock);\r
-    CpuMpData->X2ApicEnable = TRUE;\r
-    ReleaseSpinLock(&CpuMpData->MpLock);\r
+\r
+  //\r
+  // NOTE: PlatformId is not relevant on AMD platforms.\r
+  //\r
+  if (!StandardSignatureIsAuthenticAMD ()) {\r
+    PlatformIdMsr.Uint64                           = AsmReadMsr64 (MSR_IA32_PLATFORM_ID);\r
+    CpuMpData->CpuData[ProcessorNumber].PlatformId = (UINT8)PlatformIdMsr.Bits.PlatformId;\r
   }\r
 \r
-  InitializeSpinLock(&CpuMpData->CpuData[ProcessorNumber].ApLock);\r
+  AsmCpuid (\r
+    CPUID_VERSION_INFO,\r
+    &CpuMpData->CpuData[ProcessorNumber].ProcessorSignature,\r
+    NULL,\r
+    NULL,\r
+    NULL\r
+    );\r
+\r
+  InitializeSpinLock (&CpuMpData->CpuData[ProcessorNumber].ApLock);\r
   SetApState (&CpuMpData->CpuData[ProcessorNumber], CpuStateIdle);\r
 }\r
 \r
@@ -570,19 +645,20 @@ InitializeApData (
 VOID\r
 EFIAPI\r
 ApWakeupFunction (\r
-  IN MP_CPU_EXCHANGE_INFO      *ExchangeInfo,\r
-  IN UINTN                     ApIndex\r
+  IN MP_CPU_EXCHANGE_INFO  *ExchangeInfo,\r
+  IN UINTN                 ApIndex\r
   )\r
 {\r
-  CPU_MP_DATA                *CpuMpData;\r
-  UINTN                      ProcessorNumber;\r
-  EFI_AP_PROCEDURE           Procedure;\r
-  VOID                       *Parameter;\r
-  UINT32                     BistData;\r
-  volatile UINT32            *ApStartupSignalBuffer;\r
-  CPU_INFO_IN_HOB            *CpuInfoInHob;\r
-  UINT64                     ApTopOfStack;\r
-  UINTN                      CurrentApicMode;\r
+  CPU_MP_DATA       *CpuMpData;\r
+  UINTN             ProcessorNumber;\r
+  EFI_AP_PROCEDURE  Procedure;\r
+  VOID              *Parameter;\r
+  UINT32            BistData;\r
+  volatile UINT32   *ApStartupSignalBuffer;\r
+  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
@@ -603,20 +679,14 @@ ApWakeupFunction (
   CurrentApicMode = GetApicMode ();\r
   while (TRUE) {\r
     if (CpuMpData->InitFlag == ApInitConfig) {\r
-      //\r
-      // Add CPU number\r
-      //\r
-      InterlockedIncrement ((UINT32 *) &CpuMpData->CpuCount);\r
       ProcessorNumber = ApIndex;\r
       //\r
       // 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
-      //\r
-      // Do some AP initialize sync\r
-      //\r
-      ApInitializeSync (CpuMpData);\r
+      ApTopOfStack = CpuMpData->Buffer + (ProcessorNumber + 1) * CpuMpData->CpuApStackSize;\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
@@ -635,27 +705,40 @@ ApWakeupFunction (
       //\r
       ApStartupSignalBuffer = CpuMpData->CpuData[ProcessorNumber].StartupApSignal;\r
       InterlockedCompareExchange32 (\r
-        (UINT32 *) ApStartupSignalBuffer,\r
+        (UINT32 *)ApStartupSignalBuffer,\r
         WAKEUP_AP_SIGNAL,\r
         0\r
         );\r
-      if (CpuMpData->ApLoopMode == ApInHltLoop) {\r
+\r
+      if (CpuMpData->InitFlag == ApInitReconfig) {\r
         //\r
-        // Restore AP's volatile registers saved\r
+        // ApInitReconfig happens when:\r
+        // 1. AP is re-enabled after it's disabled, in either PEI or DXE phase.\r
+        // 2. AP is initialized in DXE phase.\r
+        // In either case, use the volatile registers value derived from BSP.\r
+        // NOTE: IDTR.BASE stored in CpuMpData->CpuData[0].VolatileRegisters points to a\r
+        //   different IDT shared by all APs.\r
         //\r
-        RestoreVolatileRegisters (&CpuMpData->CpuData[ProcessorNumber].VolatileRegisters, TRUE);\r
+        RestoreVolatileRegisters (&CpuMpData->CpuData[0].VolatileRegisters, FALSE);\r
       } else {\r
-        //\r
-        // The CPU driver might not flush TLB for APs on spot after updating\r
-        // page attributes. AP in mwait loop mode needs to take care of it when\r
-        // woken up.\r
-        //\r
-        CpuFlushTlb ();\r
+        if (CpuMpData->ApLoopMode == ApInHltLoop) {\r
+          //\r
+          // Restore AP's volatile registers saved before AP is halted\r
+          //\r
+          RestoreVolatileRegisters (&CpuMpData->CpuData[ProcessorNumber].VolatileRegisters, TRUE);\r
+        } else {\r
+          //\r
+          // The CPU driver might not flush TLB for APs on spot after updating\r
+          // page attributes. AP in mwait loop mode needs to take care of it when\r
+          // woken up.\r
+          //\r
+          CpuFlushTlb ();\r
+        }\r
       }\r
 \r
       if (GetApState (&CpuMpData->CpuData[ProcessorNumber]) == CpuStateReady) {\r
         Procedure = (EFI_AP_PROCEDURE)CpuMpData->CpuData[ProcessorNumber].ApFunction;\r
-        Parameter = (VOID *) CpuMpData->CpuData[ProcessorNumber].ApFunctionArgument;\r
+        Parameter = (VOID *)CpuMpData->CpuData[ProcessorNumber].ApFunctionArgument;\r
         if (Procedure != NULL) {\r
           SetApState (&CpuMpData->CpuData[ProcessorNumber], CpuStateBusy);\r
           //\r
@@ -666,19 +749,20 @@ ApWakeupFunction (
           // Invoke AP function here\r
           //\r
           Procedure (Parameter);\r
-          CpuInfoInHob = (CPU_INFO_IN_HOB *) (UINTN) CpuMpData->CpuInfoInHob;\r
+          CpuInfoInHob = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;\r
           if (CpuMpData->SwitchBspFlag) {\r
             //\r
             // Re-get the processor number due to BSP/AP maybe exchange in AP function\r
             //\r
             GetProcessorNumber (CpuMpData, &ProcessorNumber);\r
-            CpuMpData->CpuData[ProcessorNumber].ApFunction = 0;\r
+            CpuMpData->CpuData[ProcessorNumber].ApFunction         = 0;\r
             CpuMpData->CpuData[ProcessorNumber].ApFunctionArgument = 0;\r
-            ApStartupSignalBuffer = CpuMpData->CpuData[ProcessorNumber].StartupApSignal;\r
-            CpuInfoInHob[ProcessorNumber].ApTopOfStack = CpuInfoInHob[CpuMpData->NewBspNumber].ApTopOfStack;\r
+            ApStartupSignalBuffer                                  = CpuMpData->CpuData[ProcessorNumber].StartupApSignal;\r
+            CpuInfoInHob[ProcessorNumber].ApTopOfStack             = CpuInfoInHob[CpuMpData->NewBspNumber].ApTopOfStack;\r
           } else {\r
-            if (CpuInfoInHob[ProcessorNumber].ApicId != GetApicId () ||\r
-                CpuInfoInHob[ProcessorNumber].InitialApicId != GetInitialApicId ()) {\r
+            if ((CpuInfoInHob[ProcessorNumber].ApicId != GetApicId ()) ||\r
+                (CpuInfoInHob[ProcessorNumber].InitialApicId != GetInitialApicId ()))\r
+            {\r
               if (CurrentApicMode != GetApicMode ()) {\r
                 //\r
                 // If APIC mode change happened during AP function execution,\r
@@ -696,40 +780,60 @@ ApWakeupFunction (
             }\r
           }\r
         }\r
-        SetApState (&CpuMpData->CpuData[ProcessorNumber], CpuStateIdle);\r
+\r
+        SetApState (&CpuMpData->CpuData[ProcessorNumber], CpuStateFinished);\r
       }\r
     }\r
 \r
+    if (CpuMpData->ApLoopMode == ApInHltLoop) {\r
+      //\r
+      // Save AP volatile registers\r
+      //\r
+      SaveVolatileRegisters (&CpuMpData->CpuData[ProcessorNumber].VolatileRegisters);\r
+    }\r
+\r
     //\r
     // AP finished executing C code\r
     //\r
-    InterlockedIncrement ((UINT32 *) &CpuMpData->FinishedCount);\r
-    InterlockedDecrement ((UINT32 *) &CpuMpData->MpCpuExchangeInfo->NumApsExecuting);\r
+    InterlockedIncrement ((UINT32 *)&CpuMpData->FinishedCount);\r
+\r
+    if (CpuMpData->InitFlag == ApInitConfig) {\r
+      //\r
+      // Delay decrementing the APs executing count when SEV-ES is enabled\r
+      // 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->UseSevEsAPMethod) {\r
+        InterlockedDecrement ((UINT32 *)&CpuMpData->MpCpuExchangeInfo->NumApsExecuting);\r
+      }\r
+    }\r
 \r
     //\r
     // Place AP is specified loop mode\r
     //\r
     if (CpuMpData->ApLoopMode == ApInHltLoop) {\r
-      //\r
-      // Save AP volatile registers\r
-      //\r
-      SaveVolatileRegisters (&CpuMpData->CpuData[ProcessorNumber].VolatileRegisters);\r
       //\r
       // Place AP in HLT-loop\r
       //\r
       while (TRUE) {\r
         DisableInterrupts ();\r
-        CpuSleep ();\r
+        if (CpuMpData->UseSevEsAPMethod) {\r
+          SevEsPlaceApHlt (CpuMpData);\r
+        } else {\r
+          CpuSleep ();\r
+        }\r
+\r
         CpuPause ();\r
       }\r
     }\r
+\r
     while (TRUE) {\r
       DisableInterrupts ();\r
       if (CpuMpData->ApLoopMode == ApInMwaitLoop) {\r
         //\r
         // Place AP in MWAIT-loop\r
         //\r
-        AsmMonitor ((UINTN) ApStartupSignalBuffer, 0, 0);\r
+        AsmMonitor ((UINTN)ApStartupSignalBuffer, 0, 0);\r
         if (*ApStartupSignalBuffer != WAKEUP_AP_SIGNAL) {\r
           //\r
           // Check AP start-up signal again.\r
@@ -765,7 +869,7 @@ ApWakeupFunction (
 **/\r
 VOID\r
 WaitApWakeup (\r
-  IN volatile UINT32        *ApStartupSignalBuffer\r
+  IN volatile UINT32  *ApStartupSignalBuffer\r
   )\r
 {\r
   //\r
@@ -773,14 +877,39 @@ WaitApWakeup (
   // Otherwise, write StartupApSignal again till AP waken up.\r
   //\r
   while (InterlockedCompareExchange32 (\r
-          (UINT32 *) ApStartupSignalBuffer,\r
-          WAKEUP_AP_SIGNAL,\r
-          WAKEUP_AP_SIGNAL\r
-          ) != 0) {\r
+           (UINT32 *)ApStartupSignalBuffer,\r
+           WAKEUP_AP_SIGNAL,\r
+           WAKEUP_AP_SIGNAL\r
+           ) != 0)\r
+  {\r
     CpuPause ();\r
   }\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
@@ -789,80 +918,86 @@ WaitApWakeup (
 **/\r
 VOID\r
 FillExchangeInfoData (\r
-  IN CPU_MP_DATA               *CpuMpData\r
+  IN CPU_MP_DATA  *CpuMpData\r
   )\r
 {\r
-  volatile MP_CPU_EXCHANGE_INFO    *ExchangeInfo;\r
-  UINTN                            Size;\r
-  IA32_SEGMENT_DESCRIPTOR          *Selector;\r
+  volatile MP_CPU_EXCHANGE_INFO  *ExchangeInfo;\r
+  UINTN                          Size;\r
+  IA32_SEGMENT_DESCRIPTOR        *Selector;\r
+  IA32_CR4                       Cr4;\r
 \r
-  ExchangeInfo                  = CpuMpData->MpCpuExchangeInfo;\r
-  ExchangeInfo->Lock            = 0;\r
-  ExchangeInfo->StackStart      = CpuMpData->Buffer;\r
-  ExchangeInfo->StackSize       = CpuMpData->CpuApStackSize;\r
-  ExchangeInfo->BufferStart     = CpuMpData->WakeupBuffer;\r
-  ExchangeInfo->ModeOffset      = CpuMpData->AddressMap.ModeEntryOffset;\r
+  ExchangeInfo              = CpuMpData->MpCpuExchangeInfo;\r
+  ExchangeInfo->StackStart  = CpuMpData->Buffer;\r
+  ExchangeInfo->StackSize   = CpuMpData->CpuApStackSize;\r
+  ExchangeInfo->BufferStart = CpuMpData->WakeupBuffer;\r
+  ExchangeInfo->ModeOffset  = CpuMpData->AddressMap.ModeEntryOffset;\r
 \r
-  ExchangeInfo->CodeSegment     = AsmReadCs ();\r
-  ExchangeInfo->DataSegment     = AsmReadDs ();\r
+  ExchangeInfo->CodeSegment = AsmReadCs ();\r
+  ExchangeInfo->DataSegment = AsmReadDs ();\r
 \r
-  ExchangeInfo->Cr3             = AsmReadCr3 ();\r
+  ExchangeInfo->Cr3 = AsmReadCr3 ();\r
 \r
-  ExchangeInfo->CFunction       = (UINTN) ApWakeupFunction;\r
+  ExchangeInfo->CFunction       = (UINTN)ApWakeupFunction;\r
   ExchangeInfo->ApIndex         = 0;\r
   ExchangeInfo->NumApsExecuting = 0;\r
-  ExchangeInfo->InitFlag        = (UINTN) CpuMpData->InitFlag;\r
-  ExchangeInfo->CpuInfo         = (CPU_INFO_IN_HOB *) (UINTN) CpuMpData->CpuInfoInHob;\r
+  ExchangeInfo->InitFlag        = (UINTN)CpuMpData->InitFlag;\r
+  ExchangeInfo->CpuInfo         = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;\r
   ExchangeInfo->CpuMpData       = CpuMpData;\r
 \r
   ExchangeInfo->EnableExecuteDisable = IsBspExecuteDisableEnabled ();\r
 \r
   ExchangeInfo->InitializeFloatingPointUnitsAddress = (UINTN)InitializeFloatingPointUnits;\r
 \r
+  //\r
+  // We can check either CPUID(7).ECX[bit16] or check CR4.LA57[bit12]\r
+  //  to determin whether 5-Level Paging is enabled.\r
+  // CPUID(7).ECX[bit16] shows CPU's capability, CR4.LA57[bit12] shows\r
+  // current system setting.\r
+  // Using latter way is simpler because it also eliminates the needs to\r
+  //  check whether platform wants to enable it.\r
+  //\r
+  Cr4.UintN                        = AsmReadCr4 ();\r
+  ExchangeInfo->Enable5LevelPaging = (BOOLEAN)(Cr4.Bits.LA57 == 1);\r
+  DEBUG ((DEBUG_INFO, "%a: 5-Level Paging = %d\n", gEfiCallerBaseName, ExchangeInfo->Enable5LevelPaging));\r
+\r
+  ExchangeInfo->SevEsIsEnabled  = CpuMpData->SevEsIsEnabled;\r
+  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
-  AsmReadGdtr ((IA32_DESCRIPTOR *) &ExchangeInfo->GdtrProfile);\r
-  AsmReadIdtr ((IA32_DESCRIPTOR *) &ExchangeInfo->IdtrProfile);\r
+  AsmReadGdtr ((IA32_DESCRIPTOR *)&ExchangeInfo->GdtrProfile);\r
+  AsmReadIdtr ((IA32_DESCRIPTOR *)&ExchangeInfo->IdtrProfile);\r
 \r
   //\r
   // Find a 32-bit code segment\r
   //\r
   Selector = (IA32_SEGMENT_DESCRIPTOR *)ExchangeInfo->GdtrProfile.Base;\r
-  Size = ExchangeInfo->GdtrProfile.Limit + 1;\r
+  Size     = ExchangeInfo->GdtrProfile.Limit + 1;\r
   while (Size > 0) {\r
-    if (Selector->Bits.L == 0 && Selector->Bits.Type >= 8) {\r
+    if ((Selector->Bits.L == 0) && (Selector->Bits.Type >= 8)) {\r
       ExchangeInfo->ModeTransitionSegment =\r
         (UINT16)((UINTN)Selector - ExchangeInfo->GdtrProfile.Base);\r
       break;\r
     }\r
+\r
     Selector += 1;\r
-    Size -= sizeof (IA32_SEGMENT_DESCRIPTOR);\r
+    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.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
-                         (UINT32)CpuMpData->AddressMap.ModeTransitionOffset;\r
+                                 (UINT32)ExchangeInfo->ModeOffset -\r
+                                 (UINT32)CpuMpData->AddressMap.ModeTransitionOffset;\r
   ExchangeInfo->ModeHighSegment = (UINT16)ExchangeInfo->CodeSegment;\r
 }\r
 \r
@@ -876,9 +1011,9 @@ FillExchangeInfoData (
 **/\r
 VOID\r
 TimedWaitForApFinish (\r
-  IN CPU_MP_DATA               *CpuMpData,\r
-  IN UINT32                    FinishedApLimit,\r
-  IN UINT32                    TimeLimit\r
+  IN CPU_MP_DATA  *CpuMpData,\r
+  IN UINT32       FinishedApLimit,\r
+  IN UINT32       TimeLimit\r
   );\r
 \r
 /**\r
@@ -887,19 +1022,19 @@ TimedWaitForApFinish (
   @param[in]  CpuMpData  The pointer to CPU MP Data structure.\r
 **/\r
 VOID\r
-BackupAndPrepareWakeupBuffer(\r
-  IN CPU_MP_DATA              *CpuMpData\r
+BackupAndPrepareWakeupBuffer (\r
+  IN CPU_MP_DATA  *CpuMpData\r
   )\r
 {\r
   CopyMem (\r
-    (VOID *) CpuMpData->BackupBuffer,\r
-    (VOID *) CpuMpData->WakeupBuffer,\r
+    (VOID *)CpuMpData->BackupBuffer,\r
+    (VOID *)CpuMpData->WakeupBuffer,\r
     CpuMpData->BackupBufferSize\r
     );\r
   CopyMem (\r
-    (VOID *) CpuMpData->WakeupBuffer,\r
-    (VOID *) CpuMpData->AddressMap.RendezvousFunnelAddress,\r
-    CpuMpData->AddressMap.RendezvousFunnelSize\r
+    (VOID *)CpuMpData->WakeupBuffer,\r
+    (VOID *)CpuMpData->AddressMap.RendezvousFunnelAddress,\r
+    CpuMpData->BackupBufferSize - sizeof (MP_CPU_EXCHANGE_INFO)\r
     );\r
 }\r
 \r
@@ -909,13 +1044,13 @@ BackupAndPrepareWakeupBuffer(
   @param[in]  CpuMpData  The pointer to CPU MP Data structure.\r
 **/\r
 VOID\r
-RestoreWakeupBuffer(\r
-  IN CPU_MP_DATA              *CpuMpData\r
+RestoreWakeupBuffer (\r
+  IN CPU_MP_DATA  *CpuMpData\r
   )\r
 {\r
   CopyMem (\r
-    (VOID *) CpuMpData->WakeupBuffer,\r
-    (VOID *) CpuMpData->BackupBuffer,\r
+    (VOID *)CpuMpData->WakeupBuffer,\r
+    (VOID *)CpuMpData->BackupBuffer,\r
     CpuMpData->BackupBufferSize\r
     );\r
 }\r
@@ -926,24 +1061,64 @@ RestoreWakeupBuffer(
   @param[in, out]  CpuMpData  The pointer to CPU MP Data structure.\r
 **/\r
 VOID\r
-AllocateResetVector (\r
-  IN OUT CPU_MP_DATA          *CpuMpData\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
+    CpuMpData->WakeupBuffer      = GetWakeupBuffer (CpuMpData->BackupBufferSize);\r
+    CpuMpData->MpCpuExchangeInfo = (MP_CPU_EXCHANGE_INFO *)(UINTN)\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. 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
+        !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
+      //\r
+      ApResetStackSize = (AP_RESET_STACK_SIZE *\r
+                          PcdGet32 (PcdCpuMaxLogicalProcessorNumber));\r
+\r
+      //\r
+      // Invoke GetWakeupBuffer a second time to allocate the stack area\r
+      // below 1MB. The returned buffer will be page aligned and sized and\r
+      // below the previously allocated buffer.\r
+      //\r
+      CpuMpData->SevEsAPResetStackStart = GetWakeupBuffer (ApResetStackSize);\r
 \r
-  if (CpuMpData->WakeupBuffer == (UINTN) -1) {\r
-    ApResetVectorSize = CpuMpData->AddressMap.RendezvousFunnelSize +\r
-                          sizeof (MP_CPU_EXCHANGE_INFO);\r
+      //\r
+      // Check to be sure that the "allocate below" behavior hasn't changed.\r
+      // This will also catch a failed allocation, as "-1" is returned on\r
+      // failure.\r
+      //\r
+      if (CpuMpData->SevEsAPResetStackStart >= CpuMpData->WakeupBuffer) {\r
+        DEBUG ((\r
+          DEBUG_ERROR,\r
+          "SEV-ES AP reset stack is not below wakeup buffer\n"\r
+          ));\r
 \r
-    CpuMpData->WakeupBuffer      = GetWakeupBuffer (ApResetVectorSize);\r
-    CpuMpData->MpCpuExchangeInfo = (MP_CPU_EXCHANGE_INFO *) (UINTN)\r
-                    (CpuMpData->WakeupBuffer + CpuMpData->AddressMap.RendezvousFunnelSize);\r
-    CpuMpData->WakeupBufferHigh  = GetModeTransitionBuffer (\r
-                                    CpuMpData->AddressMap.RendezvousFunnelSize -\r
-                                    CpuMpData->AddressMap.ModeTransitionOffset\r
-                                    );\r
+        ASSERT (FALSE);\r
+        CpuDeadLoop ();\r
+      }\r
+    }\r
   }\r
+\r
   BackupAndPrepareWakeupBuffer (CpuMpData);\r
 }\r
 \r
@@ -954,10 +1129,18 @@ AllocateResetVector (
 **/\r
 VOID\r
 FreeResetVector (\r
-  IN CPU_MP_DATA              *CpuMpData\r
+  IN CPU_MP_DATA  *CpuMpData\r
   )\r
 {\r
-  RestoreWakeupBuffer (CpuMpData);\r
+  //\r
+  // If SEV-ES is enabled, the reset area is needed for AP parking and\r
+  // and AP startup in the OS, so the reset area is reserved. Do not\r
+  // perform the restore as this will overwrite memory which has data\r
+  // needed by SEV-ES.\r
+  //\r
+  if (!CpuMpData->UseSevEsAPMethod) {\r
+    RestoreWakeupBuffer (CpuMpData);\r
+  }\r
 }\r
 \r
 /**\r
@@ -969,29 +1152,33 @@ FreeResetVector (
   @param[in] ProcessorNumber    The handle number of specified processor\r
   @param[in] Procedure          The function to be invoked by AP\r
   @param[in] ProcedureArgument  The argument to be passed into AP function\r
+  @param[in] WakeUpDisabledAps  Whether need to wake up disabled APs in broadcast mode.\r
 **/\r
 VOID\r
 WakeUpAP (\r
-  IN CPU_MP_DATA               *CpuMpData,\r
-  IN BOOLEAN                   Broadcast,\r
-  IN UINTN                     ProcessorNumber,\r
-  IN EFI_AP_PROCEDURE          Procedure,              OPTIONAL\r
-  IN VOID                      *ProcedureArgument      OPTIONAL\r
+  IN CPU_MP_DATA       *CpuMpData,\r
+  IN BOOLEAN           Broadcast,\r
+  IN UINTN             ProcessorNumber,\r
+  IN EFI_AP_PROCEDURE  Procedure               OPTIONAL,\r
+  IN VOID              *ProcedureArgument      OPTIONAL,\r
+  IN BOOLEAN           WakeUpDisabledAps\r
   )\r
 {\r
-  volatile MP_CPU_EXCHANGE_INFO    *ExchangeInfo;\r
-  UINTN                            Index;\r
-  CPU_AP_DATA                      *CpuData;\r
-  BOOLEAN                          ResetVectorRequired;\r
-  CPU_INFO_IN_HOB                  *CpuInfoInHob;\r
+  volatile MP_CPU_EXCHANGE_INFO  *ExchangeInfo;\r
+  UINTN                          Index;\r
+  CPU_AP_DATA                    *CpuData;\r
+  BOOLEAN                        ResetVectorRequired;\r
+  CPU_INFO_IN_HOB                *CpuInfoInHob;\r
 \r
   CpuMpData->FinishedCount = 0;\r
-  ResetVectorRequired = FALSE;\r
+  ResetVectorRequired      = FALSE;\r
 \r
   if (CpuMpData->WakeUpByInitSipiSipi ||\r
-      CpuMpData->InitFlag   != ApInitDone) {\r
+      (CpuMpData->InitFlag   != ApInitDone))\r
+  {\r
     ResetVectorRequired = TRUE;\r
-    AllocateResetVector (CpuMpData);\r
+    AllocateResetVectorBelow1Mb (CpuMpData);\r
+    AllocateSevEsAPMemory (CpuMpData);\r
     FillExchangeInfoData (CpuMpData);\r
     SaveLocalApicTimerSetting (CpuMpData);\r
   }\r
@@ -1010,39 +1197,108 @@ WakeUpAP (
     for (Index = 0; Index < CpuMpData->CpuCount; Index++) {\r
       if (Index != CpuMpData->BspNumber) {\r
         CpuData = &CpuMpData->CpuData[Index];\r
-        CpuData->ApFunction         = (UINTN) Procedure;\r
-        CpuData->ApFunctionArgument = (UINTN) ProcedureArgument;\r
+        //\r
+        // All AP(include disabled AP) will be woke up by INIT-SIPI-SIPI, but\r
+        // the AP procedure will be skipped for disabled AP because AP state\r
+        // is not CpuStateReady.\r
+        //\r
+        if ((GetApState (CpuData) == CpuStateDisabled) && !WakeUpDisabledAps) {\r
+          continue;\r
+        }\r
+\r
+        CpuData->ApFunction         = (UINTN)Procedure;\r
+        CpuData->ApFunctionArgument = (UINTN)ProcedureArgument;\r
         SetApState (CpuData, CpuStateReady);\r
         if (CpuMpData->InitFlag != ApInitConfig) {\r
-          *(UINT32 *) CpuData->StartupApSignal = WAKEUP_AP_SIGNAL;\r
+          *(UINT32 *)CpuData->StartupApSignal = WAKEUP_AP_SIGNAL;\r
         }\r
       }\r
     }\r
+\r
     if (ResetVectorRequired) {\r
+      //\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
+      if (CpuMpData->SevEsIsEnabled) {\r
+        SetSevEsJumpTable (ExchangeInfo->BufferStart);\r
+      }\r
+\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
-      //\r
-      // Here support two methods to collect AP count through adjust\r
-      // PcdCpuApInitTimeOutInMicroSeconds values.\r
-      //\r
-      // one way is set a value to just let the first AP to start the\r
-      // initialization, then through the later while loop to wait all Aps\r
-      // finsh the initialization.\r
-      // The other way is set a value to let all APs finished the initialzation.\r
-      // In this case, the later while loop is useless.\r
-      //\r
-      TimedWaitForApFinish (\r
-        CpuMpData,\r
-        PcdGet32 (PcdCpuMaxLogicalProcessorNumber) - 1,\r
-        PcdGet32 (PcdCpuApInitTimeOutInMicroSeconds)\r
-        );\r
+      if (PcdGet32 (PcdCpuBootLogicalProcessorNumber) > 0) {\r
+        //\r
+        // The AP enumeration algorithm below is suitable only when the\r
+        // platform can tell us the *exact* boot CPU count in advance.\r
+        //\r
+        // The wait below finishes only when the detected AP count reaches\r
+        // (PcdCpuBootLogicalProcessorNumber - 1), regardless of how long that\r
+        // takes. If at least one AP fails to check in (meaning a platform\r
+        // hardware bug), the detection hangs forever, by design. If the actual\r
+        // boot CPU count in the system is higher than\r
+        // PcdCpuBootLogicalProcessorNumber (meaning a platform\r
+        // misconfiguration), then some APs may complete initialization after\r
+        // the wait finishes, and cause undefined behavior.\r
+        //\r
+        TimedWaitForApFinish (\r
+          CpuMpData,\r
+          PcdGet32 (PcdCpuBootLogicalProcessorNumber) - 1,\r
+          MAX_UINT32 // approx. 71 minutes\r
+          );\r
+      } else {\r
+        //\r
+        // The AP enumeration algorithm below is suitable for two use cases.\r
+        //\r
+        // (1) The check-in time for an individual AP is bounded, and APs run\r
+        //     through their initialization routines strongly concurrently. In\r
+        //     particular, the number of concurrently running APs\r
+        //     ("NumApsExecuting") is never expected to fall to zero\r
+        //     *temporarily* -- it is expected to fall to zero only when all\r
+        //     APs have checked-in.\r
+        //\r
+        //     In this case, the platform is supposed to set\r
+        //     PcdCpuApInitTimeOutInMicroSeconds to a low-ish value (just long\r
+        //     enough for one AP to start initialization). The timeout will be\r
+        //     reached soon, and remaining APs are collected by watching\r
+        //     NumApsExecuting fall to zero. If NumApsExecuting falls to zero\r
+        //     mid-process, while some APs have not completed initialization,\r
+        //     the behavior is undefined.\r
+        //\r
+        // (2) The check-in time for an individual AP is unbounded, and/or APs\r
+        //     may complete their initializations widely spread out. In\r
+        //     particular, some APs may finish initialization before some APs\r
+        //     even start.\r
+        //\r
+        //     In this case, the platform is supposed to set\r
+        //     PcdCpuApInitTimeOutInMicroSeconds to a high-ish value. The AP\r
+        //     enumeration will always take that long (except when the boot CPU\r
+        //     count happens to be maximal, that is,\r
+        //     PcdCpuMaxLogicalProcessorNumber). All APs are expected to\r
+        //     check-in before the timeout, and NumApsExecuting is assumed zero\r
+        //     at timeout. APs that miss the time-out may cause undefined\r
+        //     behavior.\r
+        //\r
+        TimedWaitForApFinish (\r
+          CpuMpData,\r
+          PcdGet32 (PcdCpuMaxLogicalProcessorNumber) - 1,\r
+          PcdGet32 (PcdCpuApInitTimeOutInMicroSeconds)\r
+          );\r
 \r
-      while (CpuMpData->MpCpuExchangeInfo->NumApsExecuting != 0) {\r
-        CpuPause();\r
+        while (CpuMpData->MpCpuExchangeInfo->NumApsExecuting != 0) {\r
+          CpuPause ();\r
+        }\r
       }\r
     } else {\r
       //\r
@@ -1056,22 +1312,37 @@ WakeUpAP (
       }\r
     }\r
   } else {\r
-    CpuData = &CpuMpData->CpuData[ProcessorNumber];\r
-    CpuData->ApFunction         = (UINTN) Procedure;\r
-    CpuData->ApFunctionArgument = (UINTN) ProcedureArgument;\r
+    CpuData                     = &CpuMpData->CpuData[ProcessorNumber];\r
+    CpuData->ApFunction         = (UINTN)Procedure;\r
+    CpuData->ApFunctionArgument = (UINTN)ProcedureArgument;\r
     SetApState (CpuData, CpuStateReady);\r
     //\r
     // Wakeup specified AP\r
     //\r
     ASSERT (CpuMpData->InitFlag != ApInitConfig);\r
-    *(UINT32 *) CpuData->StartupApSignal = WAKEUP_AP_SIGNAL;\r
+    *(UINT32 *)CpuData->StartupApSignal = WAKEUP_AP_SIGNAL;\r
     if (ResetVectorRequired) {\r
-      CpuInfoInHob = (CPU_INFO_IN_HOB *) (UINTN) CpuMpData->CpuInfoInHob;\r
-      SendInitSipiSipi (\r
-        CpuInfoInHob[ProcessorNumber].ApicId,\r
-        (UINT32) ExchangeInfo->BufferStart\r
-        );\r
+      CpuInfoInHob = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;\r
+\r
+      //\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
+      if (CpuMpData->SevEsIsEnabled) {\r
+        SetSevEsJumpTable (ExchangeInfo->BufferStart);\r
+      }\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
     // Wait specified AP waken up\r
     //\r
@@ -1111,8 +1382,8 @@ CalculateTimeout (
   OUT UINT64  *CurrentTime\r
   )\r
 {\r
-  UINT64 TimeoutInSeconds;\r
-  UINT64 TimestampCounterFreq;\r
+  UINT64  TimeoutInSeconds;\r
+  UINT64  TimestampCounterFreq;\r
 \r
   //\r
   // Read the current value of the performance counter\r
@@ -1197,25 +1468,30 @@ CheckTimeout (
   if (Timeout == 0) {\r
     return FALSE;\r
   }\r
+\r
   GetPerformanceCounterProperties (&Start, &End);\r
   Cycle = End - Start;\r
   if (Cycle < 0) {\r
     Cycle = -Cycle;\r
   }\r
+\r
   Cycle++;\r
-  CurrentTime = GetPerformanceCounter();\r
-  Delta = (INT64) (CurrentTime - *PreviousTime);\r
+  CurrentTime = GetPerformanceCounter ();\r
+  Delta       = (INT64)(CurrentTime - *PreviousTime);\r
   if (Start > End) {\r
     Delta = -Delta;\r
   }\r
+\r
   if (Delta < 0) {\r
     Delta += Cycle;\r
   }\r
-  *TotalTime += Delta;\r
+\r
+  *TotalTime   += Delta;\r
   *PreviousTime = CurrentTime;\r
   if (*TotalTime > Timeout) {\r
     return TRUE;\r
   }\r
+\r
   return FALSE;\r
 }\r
 \r
@@ -1229,9 +1505,9 @@ CheckTimeout (
 **/\r
 VOID\r
 TimedWaitForApFinish (\r
-  IN CPU_MP_DATA               *CpuMpData,\r
-  IN UINT32                    FinishedApLimit,\r
-  IN UINT32                    TimeLimit\r
+  IN CPU_MP_DATA  *CpuMpData,\r
+  IN UINT32       FinishedApLimit,\r
+  IN UINT32       TimeLimit\r
   )\r
 {\r
   //\r
@@ -1242,7 +1518,7 @@ TimedWaitForApFinish (
     return;\r
   }\r
 \r
-  CpuMpData->TotalTime = 0;\r
+  CpuMpData->TotalTime    = 0;\r
   CpuMpData->ExpectedTime = CalculateTimeout (\r
                               TimeLimit,\r
                               &CpuMpData->CurrentTime\r
@@ -1252,7 +1528,8 @@ TimedWaitForApFinish (
             &CpuMpData->CurrentTime,\r
             &CpuMpData->TotalTime,\r
             CpuMpData->ExpectedTime\r
-            )) {\r
+            ))\r
+  {\r
     CpuPause ();\r
   }\r
 \r
@@ -1281,18 +1558,19 @@ TimedWaitForApFinish (
 **/\r
 VOID\r
 ResetProcessorToIdleState (\r
-  IN UINTN                     ProcessorNumber\r
+  IN UINTN  ProcessorNumber\r
   )\r
 {\r
-  CPU_MP_DATA           *CpuMpData;\r
+  CPU_MP_DATA  *CpuMpData;\r
 \r
   CpuMpData = GetCpuMpData ();\r
 \r
   CpuMpData->InitFlag = ApInitReconfig;\r
-  WakeUpAP (CpuMpData, FALSE, ProcessorNumber, NULL, NULL);\r
+  WakeUpAP (CpuMpData, FALSE, ProcessorNumber, NULL, NULL, TRUE);\r
   while (CpuMpData->FinishedCount < 1) {\r
     CpuPause ();\r
   }\r
+\r
   CpuMpData->InitFlag = ApInitDone;\r
 \r
   SetApState (&CpuMpData->CpuData[ProcessorNumber], CpuStateIdle);\r
@@ -1311,11 +1589,11 @@ ResetProcessorToIdleState (
 **/\r
 EFI_STATUS\r
 GetNextWaitingProcessorNumber (\r
-  OUT UINTN                    *NextProcessorNumber\r
+  OUT UINTN  *NextProcessorNumber\r
   )\r
 {\r
-  UINTN           ProcessorNumber;\r
-  CPU_MP_DATA     *CpuMpData;\r
+  UINTN        ProcessorNumber;\r
+  CPU_MP_DATA  *CpuMpData;\r
 \r
   CpuMpData = GetCpuMpData ();\r
 \r
@@ -1342,11 +1620,11 @@ GetNextWaitingProcessorNumber (
 **/\r
 EFI_STATUS\r
 CheckThisAP (\r
-  IN UINTN        ProcessorNumber\r
+  IN UINTN  ProcessorNumber\r
   )\r
 {\r
-  CPU_MP_DATA     *CpuMpData;\r
-  CPU_AP_DATA     *CpuData;\r
+  CPU_MP_DATA  *CpuMpData;\r
+  CPU_AP_DATA  *CpuData;\r
 \r
   CpuMpData = GetCpuMpData ();\r
   CpuData   = &CpuMpData->CpuData[ProcessorNumber];\r
@@ -1359,10 +1637,12 @@ CheckThisAP (
   //\r
   // If the AP finishes for StartupThisAP(), return EFI_SUCCESS.\r
   //\r
-  if (GetApState(CpuData) == CpuStateIdle) {\r
+  if (GetApState (CpuData) == CpuStateFinished) {\r
     if (CpuData->Finished != NULL) {\r
       *(CpuData->Finished) = TRUE;\r
     }\r
+\r
+    SetApState (CpuData, CpuStateIdle);\r
     return EFI_SUCCESS;\r
   } else {\r
     //\r
@@ -1372,6 +1652,7 @@ CheckThisAP (
       if (CpuData->Finished != NULL) {\r
         *(CpuData->Finished) = FALSE;\r
       }\r
+\r
       //\r
       // Reset failed AP to idle state\r
       //\r
@@ -1380,6 +1661,7 @@ CheckThisAP (
       return EFI_TIMEOUT;\r
     }\r
   }\r
+\r
   return EFI_NOT_READY;\r
 }\r
 \r
@@ -1398,12 +1680,12 @@ CheckAllAPs (
   VOID\r
   )\r
 {\r
-  UINTN           ProcessorNumber;\r
-  UINTN           NextProcessorNumber;\r
-  UINTN           ListIndex;\r
-  EFI_STATUS      Status;\r
-  CPU_MP_DATA     *CpuMpData;\r
-  CPU_AP_DATA     *CpuData;\r
+  UINTN        ProcessorNumber;\r
+  UINTN        NextProcessorNumber;\r
+  UINTN        ListIndex;\r
+  EFI_STATUS   Status;\r
+  CPU_MP_DATA  *CpuMpData;\r
+  CPU_AP_DATA  *CpuData;\r
 \r
   CpuMpData = GetCpuMpData ();\r
 \r
@@ -1423,9 +1705,10 @@ CheckAllAPs (
     // Only BSP and corresponding AP access this unit of CPU Data. This means the AP will not modify the\r
     // value of state after setting the it to CpuStateIdle, so BSP can safely make use of its value.\r
     //\r
-    if (GetApState(CpuData) == CpuStateIdle) {\r
-      CpuMpData->RunningCount ++;\r
+    if (GetApState (CpuData) == CpuStateFinished) {\r
+      CpuMpData->RunningCount--;\r
       CpuMpData->CpuData[ProcessorNumber].Waiting = FALSE;\r
+      SetApState (CpuData, CpuStateIdle);\r
 \r
       //\r
       // If in Single Thread mode, then search for the next waiting AP for execution.\r
@@ -1437,11 +1720,12 @@ CheckAllAPs (
           WakeUpAP (\r
             CpuMpData,\r
             FALSE,\r
-            (UINT32) NextProcessorNumber,\r
+            (UINT32)NextProcessorNumber,\r
             CpuMpData->Procedure,\r
-            CpuMpData->ProcArguments\r
+            CpuMpData->ProcArguments,\r
+            TRUE\r
             );\r
-         }\r
+        }\r
       }\r
     }\r
   }\r
@@ -1449,7 +1733,7 @@ CheckAllAPs (
   //\r
   // If all APs finish, return EFI_SUCCESS.\r
   //\r
-  if (CpuMpData->RunningCount == CpuMpData->StartCount) {\r
+  if (CpuMpData->RunningCount == 0) {\r
     return EFI_SUCCESS;\r
   }\r
 \r
@@ -1457,18 +1741,21 @@ CheckAllAPs (
   // If timeout expires, report timeout.\r
   //\r
   if (CheckTimeout (\r
-       &CpuMpData->CurrentTime,\r
-       &CpuMpData->TotalTime,\r
-       CpuMpData->ExpectedTime)\r
-       ) {\r
+        &CpuMpData->CurrentTime,\r
+        &CpuMpData->TotalTime,\r
+        CpuMpData->ExpectedTime\r
+        )\r
+      )\r
+  {\r
     //\r
     // If FailedCpuList is not NULL, record all failed APs in it.\r
     //\r
     if (CpuMpData->FailedCpuList != NULL) {\r
       *CpuMpData->FailedCpuList =\r
-         AllocatePool ((CpuMpData->StartCount - CpuMpData->FinishedCount + 1) * sizeof (UINTN));\r
+        AllocatePool ((CpuMpData->RunningCount + 1) * sizeof (UINTN));\r
       ASSERT (*CpuMpData->FailedCpuList != NULL);\r
     }\r
+\r
     ListIndex = 0;\r
 \r
     for (ProcessorNumber = 0; ProcessorNumber < CpuMpData->CpuCount; ProcessorNumber++) {\r
@@ -1486,11 +1773,14 @@ CheckAllAPs (
         }\r
       }\r
     }\r
+\r
     if (CpuMpData->FailedCpuList != NULL) {\r
       (*CpuMpData->FailedCpuList)[ListIndex] = END_OF_CPU_LIST;\r
     }\r
+\r
     return EFI_TIMEOUT;\r
   }\r
+\r
   return EFI_NOT_READY;\r
 }\r
 \r
@@ -1527,32 +1817,41 @@ 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
-  VOID                     *MicrocodePatchInRam;\r
 \r
   OldCpuMpData = GetCpuMpDataFromGuidedHob ();\r
   if (OldCpuMpData == NULL) {\r
-    MaxLogicalProcessorNumber = PcdGet32(PcdCpuMaxLogicalProcessorNumber);\r
+    MaxLogicalProcessorNumber = PcdGet32 (PcdCpuMaxLogicalProcessorNumber);\r
   } else {\r
     MaxLogicalProcessorNumber = OldCpuMpData->CpuCount;\r
   }\r
+\r
   ASSERT (MaxLogicalProcessorNumber != 0);\r
 \r
   AsmGetAddressMap (&AddressMap);\r
-  ApResetVectorSize = AddressMap.RendezvousFunnelSize + sizeof (MP_CPU_EXCHANGE_INFO);\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
+  // 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
@@ -1560,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
@@ -1574,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
@@ -1583,60 +1882,40 @@ MpInitLibInitialize (
   //      CPU_INFO_IN_HOB (N)\r
   //    +--------------------+\r
   //\r
-  MonitorBuffer    = (UINT8 *) (Buffer + ApStackSize * MaxLogicalProcessorNumber);\r
-  BackupBufferAddr = (UINTN) MonitorBuffer + MonitorFilterSize * MaxLogicalProcessorNumber;\r
-  ApIdtBase        = ALIGN_VALUE (BackupBufferAddr + ApResetVectorSize, 8);\r
-  CpuMpData        = (CPU_MP_DATA *) (ApIdtBase + VolatileRegisters.Idtr.Limit + 1);\r
+  MonitorBuffer               = (UINT8 *)(Buffer + ApStackSize * MaxLogicalProcessorNumber);\r
+  BackupBufferAddr            = (UINTN)MonitorBuffer + MonitorFilterSize * MaxLogicalProcessorNumber;\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->WakeupBuffer     = (UINTN) -1;\r
+  CpuMpData->BackupBufferSize = ApResetVectorSizeBelow1Mb;\r
+  CpuMpData->WakeupBuffer     = (UINTN)-1;\r
   CpuMpData->CpuCount         = 1;\r
   CpuMpData->BspNumber        = 0;\r
   CpuMpData->WaitEvent        = NULL;\r
   CpuMpData->SwitchBspFlag    = FALSE;\r
-  CpuMpData->CpuData          = (CPU_AP_DATA *) (CpuMpData + 1);\r
-  CpuMpData->CpuInfoInHob     = (UINT64) (UINTN) (CpuMpData->CpuData + MaxLogicalProcessorNumber);\r
-  CpuMpData->MicrocodePatchRegionSize = PcdGet64 (PcdCpuMicrocodePatchRegionSize);\r
-  //\r
-  // If platform has more than one CPU, relocate microcode to memory to reduce\r
-  // loading microcode time.\r
-  //\r
-  MicrocodePatchInRam = NULL;\r
-  if (MaxLogicalProcessorNumber > 1) {\r
-    MicrocodePatchInRam = AllocatePages (\r
-                            EFI_SIZE_TO_PAGES (\r
-                              (UINTN)CpuMpData->MicrocodePatchRegionSize\r
-                              )\r
-                            );\r
-  }\r
-  if (MicrocodePatchInRam == NULL) {\r
-    //\r
-    // there is only one processor, or no microcode patch is available, or\r
-    // memory allocation failed\r
-    //\r
-    CpuMpData->MicrocodePatchAddress = PcdGet64 (PcdCpuMicrocodePatchAddress);\r
-  } else {\r
-    //\r
-    // there are multiple processors, and a microcode patch is available, and\r
-    // memory allocation succeeded\r
-    //\r
-    CopyMem (\r
-      MicrocodePatchInRam,\r
-      (VOID *)(UINTN)PcdGet64 (PcdCpuMicrocodePatchAddress),\r
-      (UINTN)CpuMpData->MicrocodePatchRegionSize\r
-      );\r
-    CpuMpData->MicrocodePatchAddress = (UINTN)MicrocodePatchInRam;\r
-  }\r
+  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->UseSevEsAPMethod = CpuMpData->SevEsIsEnabled && !CpuMpData->SevSnpIsEnabled;\r
 \r
-  InitializeSpinLock(&CpuMpData->MpLock);\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 ((CpuMpData->CpuInfoInHob + sizeof (CPU_INFO_IN_HOB) * MaxLogicalProcessorNumber) ==\r
-          Buffer + BufferSize);\r
+  ASSERT (\r
+    (CpuMpData->CpuInfoInHob + sizeof (CPU_INFO_IN_HOB) * MaxLogicalProcessorNumber) ==\r
+    (UINTN)MpBuffer + BufferSize - (ApStackSize - Buffer + (UINTN)MpBuffer)\r
+    );\r
 \r
   //\r
   // Duplicate BSP's IDT to APs.\r
@@ -1644,6 +1923,10 @@ MpInitLibInitialize (
   //\r
   CopyMem ((VOID *)ApIdtBase, (VOID *)VolatileRegisters.Idtr.Base, VolatileRegisters.Idtr.Limit + 1);\r
   VolatileRegisters.Idtr.Base = ApIdtBase;\r
+  //\r
+  // Don't pass BSP's TR to APs to avoid AP init failure.\r
+  //\r
+  VolatileRegisters.Tr = 0;\r
   CopyMem (&CpuMpData->CpuData[0].VolatileRegisters, &VolatileRegisters, sizeof (VolatileRegisters));\r
   //\r
   // Set BSP basic information\r
@@ -1668,14 +1951,20 @@ MpInitLibInitialize (
     CpuMpData->CpuData[Index].StartupApSignal =\r
       (UINT32 *)(MonitorBuffer + MonitorFilterSize * Index);\r
   }\r
+\r
   //\r
-  // Load Microcode on BSP\r
-  //\r
-  MicrocodeDetect (CpuMpData, TRUE);\r
-  //\r
-  // Store BSP's MTRR setting\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
-  MtrrGetAllMtrrs (&CpuMpData->MtrrTable);\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
@@ -1693,38 +1982,100 @@ MpInitLibInitialize (
     // APs have been wakeup before, just get the CPU Information\r
     // from HOB\r
     //\r
-    CpuMpData->CpuCount  = OldCpuMpData->CpuCount;\r
-    CpuMpData->BspNumber = OldCpuMpData->BspNumber;\r
-    CpuMpData->InitFlag  = ApInitReconfig;\r
-    CpuMpData->CpuInfoInHob = OldCpuMpData->CpuInfoInHob;\r
-    CpuInfoInHob = (CPU_INFO_IN_HOB *) (UINTN) CpuMpData->CpuInfoInHob;\r
+    OldCpuMpData->NewCpuMpData = CpuMpData;\r
+    CpuMpData->CpuCount        = OldCpuMpData->CpuCount;\r
+    CpuMpData->BspNumber       = OldCpuMpData->BspNumber;\r
+    CpuMpData->CpuInfoInHob    = OldCpuMpData->CpuInfoInHob;\r
+    CpuInfoInHob               = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;\r
     for (Index = 0; Index < CpuMpData->CpuCount; Index++) {\r
-      InitializeSpinLock(&CpuMpData->CpuData[Index].ApLock);\r
-      if (CpuInfoInHob[Index].InitialApicId >= 255 || Index > 254) {\r
-        CpuMpData->X2ApicEnable = TRUE;\r
-      }\r
-      CpuMpData->CpuData[Index].CpuHealthy = (CpuInfoInHob[Index].Health == 0)? TRUE:FALSE;\r
+      InitializeSpinLock (&CpuMpData->CpuData[Index].ApLock);\r
+      CpuMpData->CpuData[Index].CpuHealthy = (CpuInfoInHob[Index].Health == 0) ? TRUE : FALSE;\r
       CpuMpData->CpuData[Index].ApFunction = 0;\r
-      CopyMem (&CpuMpData->CpuData[Index].VolatileRegisters, &VolatileRegisters, sizeof (CPU_VOLATILE_REGISTERS));\r
     }\r
-    if (MaxLogicalProcessorNumber > 1) {\r
+  }\r
+\r
+  if (!GetMicrocodePatchInfoFromHob (\r
+         &CpuMpData->MicrocodePatchAddress,\r
+         &CpuMpData->MicrocodePatchRegionSize\r
+         ))\r
+  {\r
+    //\r
+    // The microcode patch information cache HOB does not exist, which means\r
+    // the microcode patches data has not been loaded into memory yet\r
+    //\r
+    ShadowMicrocodeUpdatePatch (CpuMpData);\r
+  }\r
+\r
+  //\r
+  // Detect and apply Microcode on BSP\r
+  //\r
+  MicrocodeDetect (CpuMpData, CpuMpData->BspNumber);\r
+  //\r
+  // Store BSP's MTRR setting\r
+  //\r
+  MtrrGetAllMtrrs (&CpuMpData->MtrrTable);\r
+\r
+  //\r
+  // Wakeup APs to do some AP initialize sync (Microcode & MTRR)\r
+  //\r
+  if (CpuMpData->CpuCount > 1) {\r
+    if (OldCpuMpData != NULL) {\r
       //\r
-      // Wakeup APs to do some AP initialize sync\r
+      // Only needs to use this flag for DXE phase to update the wake up\r
+      // buffer. Wakeup buffer allocated in PEI phase is no longer valid\r
+      // in DXE.\r
       //\r
-      WakeUpAP (CpuMpData, TRUE, 0, ApInitializeSync, CpuMpData);\r
+      CpuMpData->InitFlag = ApInitReconfig;\r
+    }\r
+\r
+    WakeUpAP (CpuMpData, TRUE, 0, ApInitializeSync, CpuMpData, TRUE);\r
+    //\r
+    // Wait for all APs finished initialization\r
+    //\r
+    while (CpuMpData->FinishedCount < (CpuMpData->CpuCount - 1)) {\r
+      CpuPause ();\r
+    }\r
+\r
+    if (OldCpuMpData != NULL) {\r
+      CpuMpData->InitFlag = ApInitDone;\r
+    }\r
+\r
+    for (Index = 0; Index < CpuMpData->CpuCount; Index++) {\r
+      SetApState (&CpuMpData->CpuData[Index], CpuStateIdle);\r
+    }\r
+  }\r
+\r
+  //\r
+  // Dump the microcode revision for each core.\r
+  //\r
+  DEBUG_CODE_BEGIN ();\r
+  UINT32  ThreadId;\r
+  UINT32  ExpectedMicrocodeRevision;\r
+\r
+  CpuInfoInHob = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;\r
+  for (Index = 0; Index < CpuMpData->CpuCount; Index++) {\r
+    GetProcessorLocationByApicId (CpuInfoInHob[Index].InitialApicId, NULL, NULL, &ThreadId);\r
+    if (ThreadId == 0) {\r
       //\r
-      // Wait for all APs finished initialization\r
+      // MicrocodeDetect() loads microcode in first thread of each core, so,\r
+      // CpuMpData->CpuData[Index].MicrocodeEntryAddr is initialized only for first thread of each core.\r
       //\r
-      while (CpuMpData->FinishedCount < (CpuMpData->CpuCount - 1)) {\r
-        CpuPause ();\r
-      }\r
-      CpuMpData->InitFlag = ApInitDone;\r
-      for (Index = 0; Index < CpuMpData->CpuCount; Index++) {\r
-        SetApState (&CpuMpData->CpuData[Index], CpuStateIdle);\r
+      ExpectedMicrocodeRevision = 0;\r
+      if (CpuMpData->CpuData[Index].MicrocodeEntryAddr != 0) {\r
+        ExpectedMicrocodeRevision = ((CPU_MICROCODE_HEADER *)(UINTN)CpuMpData->CpuData[Index].MicrocodeEntryAddr)->UpdateRevision;\r
       }\r
+\r
+      DEBUG ((\r
+        DEBUG_INFO,\r
+        "CPU[%04d]: Microcode revision = %08x, expected = %08x\n",\r
+        Index,\r
+        CpuMpData->CpuData[Index].MicrocodeRevision,\r
+        ExpectedMicrocodeRevision\r
+        ));\r
     }\r
   }\r
 \r
+  DEBUG_CODE_END ();\r
   //\r
   // Initialize global data for MP support\r
   //\r
@@ -1758,12 +2109,19 @@ MpInitLibGetProcessorInfo (
   OUT EFI_HEALTH_FLAGS           *HealthData  OPTIONAL\r
   )\r
 {\r
-  CPU_MP_DATA            *CpuMpData;\r
-  UINTN                  CallerNumber;\r
-  CPU_INFO_IN_HOB        *CpuInfoInHob;\r
+  CPU_MP_DATA      *CpuMpData;\r
+  UINTN            CallerNumber;\r
+  CPU_INFO_IN_HOB  *CpuInfoInHob;\r
+  UINTN            OriginalProcessorNumber;\r
 \r
-  CpuMpData = GetCpuMpData ();\r
-  CpuInfoInHob = (CPU_INFO_IN_HOB *) (UINTN) CpuMpData->CpuInfoInHob;\r
+  CpuMpData    = GetCpuMpData ();\r
+  CpuInfoInHob = (CPU_INFO_IN_HOB *)(UINTN)CpuMpData->CpuInfoInHob;\r
+\r
+  //\r
+  // Lower 24 bits contains the actual processor number.\r
+  //\r
+  OriginalProcessorNumber = ProcessorNumber;\r
+  ProcessorNumber        &= BIT24 - 1;\r
 \r
   //\r
   // Check whether caller processor is BSP\r
@@ -1781,14 +2139,16 @@ MpInitLibGetProcessorInfo (
     return EFI_NOT_FOUND;\r
   }\r
 \r
-  ProcessorInfoBuffer->ProcessorId = (UINT64) CpuInfoInHob[ProcessorNumber].ApicId;\r
+  ProcessorInfoBuffer->ProcessorId = (UINT64)CpuInfoInHob[ProcessorNumber].ApicId;\r
   ProcessorInfoBuffer->StatusFlag  = 0;\r
   if (ProcessorNumber == CpuMpData->BspNumber) {\r
     ProcessorInfoBuffer->StatusFlag |= PROCESSOR_AS_BSP_BIT;\r
   }\r
+\r
   if (CpuMpData->CpuData[ProcessorNumber].CpuHealthy) {\r
     ProcessorInfoBuffer->StatusFlag |= PROCESSOR_HEALTH_STATUS_BIT;\r
   }\r
+\r
   if (GetApState (&CpuMpData->CpuData[ProcessorNumber]) == CpuStateDisabled) {\r
     ProcessorInfoBuffer->StatusFlag &= ~PROCESSOR_ENABLED_BIT;\r
   } else {\r
@@ -1805,6 +2165,18 @@ MpInitLibGetProcessorInfo (
     &ProcessorInfoBuffer->Location.Thread\r
     );\r
 \r
+  if ((OriginalProcessorNumber & CPU_V2_EXTENDED_TOPOLOGY) != 0) {\r
+    GetProcessorLocation2ByApicId (\r
+      CpuInfoInHob[ProcessorNumber].ApicId,\r
+      &ProcessorInfoBuffer->ExtendedInformation.Location2.Package,\r
+      &ProcessorInfoBuffer->ExtendedInformation.Location2.Die,\r
+      &ProcessorInfoBuffer->ExtendedInformation.Location2.Tile,\r
+      &ProcessorInfoBuffer->ExtendedInformation.Location2.Module,\r
+      &ProcessorInfoBuffer->ExtendedInformation.Location2.Core,\r
+      &ProcessorInfoBuffer->ExtendedInformation.Location2.Thread\r
+      );\r
+  }\r
+\r
   if (HealthData != NULL) {\r
     HealthData->Uint32 = CpuInfoInHob[ProcessorNumber].Health;\r
   }\r
@@ -1825,8 +2197,8 @@ MpInitLibGetProcessorInfo (
 **/\r
 EFI_STATUS\r
 SwitchBSPWorker (\r
-  IN UINTN                     ProcessorNumber,\r
-  IN BOOLEAN                   EnableOldBSP\r
+  IN UINTN    ProcessorNumber,\r
+  IN BOOLEAN  EnableOldBSP\r
   )\r
 {\r
   CPU_MP_DATA                  *CpuMpData;\r
@@ -1899,21 +2271,26 @@ SwitchBSPWorker (
   //\r
   // Clear the BSP bit of MSR_IA32_APIC_BASE\r
   //\r
-  ApicBaseMsr.Uint64 = AsmReadMsr64 (MSR_IA32_APIC_BASE);\r
+  ApicBaseMsr.Uint64   = AsmReadMsr64 (MSR_IA32_APIC_BASE);\r
   ApicBaseMsr.Bits.BSP = 0;\r
   AsmWriteMsr64 (MSR_IA32_APIC_BASE, ApicBaseMsr.Uint64);\r
 \r
   //\r
   // Need to wakeUp AP (future BSP).\r
   //\r
-  WakeUpAP (CpuMpData, FALSE, ProcessorNumber, FutureBSPProc, CpuMpData);\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
   //\r
-  ApicBaseMsr.Uint64 = AsmReadMsr64 (MSR_IA32_APIC_BASE);\r
+  ApicBaseMsr.Uint64   = AsmReadMsr64 (MSR_IA32_APIC_BASE);\r
   ApicBaseMsr.Bits.BSP = 1;\r
   AsmWriteMsr64 (MSR_IA32_APIC_BASE, ApicBaseMsr.Uint64);\r
   ProgramVirtualWireMode ();\r
@@ -1921,7 +2298,7 @@ SwitchBSPWorker (
   //\r
   // Wait for old BSP finished AP task\r
   //\r
-  while (GetApState (&CpuMpData->CpuData[CallerNumber]) != CpuStateIdle) {\r
+  while (GetApState (&CpuMpData->CpuData[CallerNumber]) != CpuStateFinished) {\r
     CpuPause ();\r
   }\r
 \r
@@ -1934,10 +2311,11 @@ SwitchBSPWorker (
   } else {\r
     SetApState (&CpuMpData->CpuData[CallerNumber], CpuStateIdle);\r
   }\r
+\r
   //\r
   // Save new BSP number\r
   //\r
-  CpuMpData->BspNumber = (UINT32) ProcessorNumber;\r
+  CpuMpData->BspNumber = (UINT32)ProcessorNumber;\r
 \r
   //\r
   // Restore interrupt state.\r
@@ -1967,13 +2345,13 @@ SwitchBSPWorker (
 **/\r
 EFI_STATUS\r
 EnableDisableApWorker (\r
-  IN  UINTN                     ProcessorNumber,\r
-  IN  BOOLEAN                   EnableAP,\r
-  IN  UINT32                    *HealthFlag OPTIONAL\r
+  IN  UINTN    ProcessorNumber,\r
+  IN  BOOLEAN  EnableAP,\r
+  IN  UINT32   *HealthFlag OPTIONAL\r
   )\r
 {\r
-  CPU_MP_DATA               *CpuMpData;\r
-  UINTN                     CallerNumber;\r
+  CPU_MP_DATA  *CpuMpData;\r
+  UINTN        CallerNumber;\r
 \r
   CpuMpData = GetCpuMpData ();\r
 \r
@@ -2001,7 +2379,7 @@ EnableDisableApWorker (
 \r
   if (HealthFlag != NULL) {\r
     CpuMpData->CpuData[ProcessorNumber].CpuHealthy =\r
-          (BOOLEAN) ((*HealthFlag & PROCESSOR_HEALTH_STATUS_BIT) != 0);\r
+      (BOOLEAN)((*HealthFlag & PROCESSOR_HEALTH_STATUS_BIT) != 0);\r
   }\r
 \r
   return EFI_SUCCESS;\r
@@ -2026,10 +2404,10 @@ EnableDisableApWorker (
 EFI_STATUS\r
 EFIAPI\r
 MpInitLibWhoAmI (\r
-  OUT UINTN                    *ProcessorNumber\r
+  OUT UINTN  *ProcessorNumber\r
   )\r
 {\r
-  CPU_MP_DATA           *CpuMpData;\r
+  CPU_MP_DATA  *CpuMpData;\r
 \r
   if (ProcessorNumber == NULL) {\r
     return EFI_INVALID_PARAMETER;\r
@@ -2063,15 +2441,15 @@ MpInitLibWhoAmI (
 EFI_STATUS\r
 EFIAPI\r
 MpInitLibGetNumberOfProcessors (\r
-  OUT UINTN                     *NumberOfProcessors,       OPTIONAL\r
-  OUT UINTN                     *NumberOfEnabledProcessors OPTIONAL\r
+  OUT UINTN  *NumberOfProcessors        OPTIONAL,\r
+  OUT UINTN  *NumberOfEnabledProcessors OPTIONAL\r
   )\r
 {\r
-  CPU_MP_DATA             *CpuMpData;\r
-  UINTN                   CallerNumber;\r
-  UINTN                   ProcessorNumber;\r
-  UINTN                   EnabledProcessorNumber;\r
-  UINTN                   Index;\r
+  CPU_MP_DATA  *CpuMpData;\r
+  UINTN        CallerNumber;\r
+  UINTN        ProcessorNumber;\r
+  UINTN        EnabledProcessorNumber;\r
+  UINTN        Index;\r
 \r
   CpuMpData = GetCpuMpData ();\r
 \r
@@ -2091,13 +2469,14 @@ MpInitLibGetNumberOfProcessors (
   EnabledProcessorNumber = 0;\r
   for (Index = 0; Index < ProcessorNumber; Index++) {\r
     if (GetApState (&CpuMpData->CpuData[Index]) != CpuStateDisabled) {\r
-      EnabledProcessorNumber ++;\r
+      EnabledProcessorNumber++;\r
     }\r
   }\r
 \r
   if (NumberOfProcessors != NULL) {\r
     *NumberOfProcessors = ProcessorNumber;\r
   }\r
+\r
   if (NumberOfEnabledProcessors != NULL) {\r
     *NumberOfEnabledProcessors = EnabledProcessorNumber;\r
   }\r
@@ -2105,7 +2484,6 @@ MpInitLibGetNumberOfProcessors (
   return EFI_SUCCESS;\r
 }\r
 \r
-\r
 /**\r
   Worker function to execute a caller provided function on all enabled APs.\r
 \r
@@ -2117,6 +2495,7 @@ MpInitLibGetNumberOfProcessors (
                                       number.  If FALSE, then all the enabled APs\r
                                       execute the function specified by Procedure\r
                                       simultaneously.\r
+  @param[in]  ExcludeBsp              Whether let BSP also trig this task.\r
   @param[in]  WaitEvent               The event created by the caller with CreateEvent()\r
                                       service.\r
   @param[in]  TimeoutInMicroseconds   Indicates the time limit in microseconds for\r
@@ -2138,23 +2517,24 @@ MpInitLibGetNumberOfProcessors (
 \r
 **/\r
 EFI_STATUS\r
-StartupAllAPsWorker (\r
-  IN  EFI_AP_PROCEDURE          Procedure,\r
-  IN  BOOLEAN                   SingleThread,\r
-  IN  EFI_EVENT                 WaitEvent               OPTIONAL,\r
-  IN  UINTN                     TimeoutInMicroseconds,\r
-  IN  VOID                      *ProcedureArgument      OPTIONAL,\r
-  OUT UINTN                     **FailedCpuList         OPTIONAL\r
+StartupAllCPUsWorker (\r
+  IN  EFI_AP_PROCEDURE  Procedure,\r
+  IN  BOOLEAN           SingleThread,\r
+  IN  BOOLEAN           ExcludeBsp,\r
+  IN  EFI_EVENT         WaitEvent               OPTIONAL,\r
+  IN  UINTN             TimeoutInMicroseconds,\r
+  IN  VOID              *ProcedureArgument      OPTIONAL,\r
+  OUT UINTN             **FailedCpuList         OPTIONAL\r
   )\r
 {\r
-  EFI_STATUS              Status;\r
-  CPU_MP_DATA             *CpuMpData;\r
-  UINTN                   ProcessorCount;\r
-  UINTN                   ProcessorNumber;\r
-  UINTN                   CallerNumber;\r
-  CPU_AP_DATA             *CpuData;\r
-  BOOLEAN                 HasEnabledAp;\r
-  CPU_STATE               ApState;\r
+  EFI_STATUS   Status;\r
+  CPU_MP_DATA  *CpuMpData;\r
+  UINTN        ProcessorCount;\r
+  UINTN        ProcessorNumber;\r
+  UINTN        CallerNumber;\r
+  CPU_AP_DATA  *CpuData;\r
+  BOOLEAN      HasEnabledAp;\r
+  CPU_STATE    ApState;\r
 \r
   CpuMpData = GetCpuMpData ();\r
 \r
@@ -2162,7 +2542,7 @@ StartupAllAPsWorker (
     *FailedCpuList = NULL;\r
   }\r
 \r
-  if (CpuMpData->CpuCount == 1) {\r
+  if ((CpuMpData->CpuCount == 1) && ExcludeBsp) {\r
     return EFI_NOT_STARTED;\r
   }\r
 \r
@@ -2205,16 +2585,16 @@ StartupAllAPsWorker (
     }\r
   }\r
 \r
-  if (!HasEnabledAp) {\r
+  if (!HasEnabledAp && ExcludeBsp) {\r
     //\r
-    // If no enabled AP exists, return EFI_NOT_STARTED.\r
+    // If no enabled AP exists and not include Bsp to do the procedure, return EFI_NOT_STARTED.\r
     //\r
     return EFI_NOT_STARTED;\r
   }\r
 \r
-  CpuMpData->StartCount = 0;\r
+  CpuMpData->RunningCount = 0;\r
   for (ProcessorNumber = 0; ProcessorNumber < ProcessorCount; ProcessorNumber++) {\r
-    CpuData = &CpuMpData->CpuData[ProcessorNumber];\r
+    CpuData          = &CpuMpData->CpuData[ProcessorNumber];\r
     CpuData->Waiting = FALSE;\r
     if (ProcessorNumber != CpuMpData->BspNumber) {\r
       if (CpuData->State == CpuStateIdle) {\r
@@ -2222,7 +2602,7 @@ StartupAllAPsWorker (
         // Mark this processor as responsible for current calling.\r
         //\r
         CpuData->Waiting = TRUE;\r
-        CpuMpData->StartCount++;\r
+        CpuMpData->RunningCount++;\r
       }\r
     }\r
   }\r
@@ -2231,29 +2611,36 @@ StartupAllAPsWorker (
   CpuMpData->ProcArguments = ProcedureArgument;\r
   CpuMpData->SingleThread  = SingleThread;\r
   CpuMpData->FinishedCount = 0;\r
-  CpuMpData->RunningCount  = 0;\r
   CpuMpData->FailedCpuList = FailedCpuList;\r
   CpuMpData->ExpectedTime  = CalculateTimeout (\r
                                TimeoutInMicroseconds,\r
                                &CpuMpData->CurrentTime\r
                                );\r
-  CpuMpData->TotalTime     = 0;\r
-  CpuMpData->WaitEvent     = WaitEvent;\r
+  CpuMpData->TotalTime = 0;\r
+  CpuMpData->WaitEvent = WaitEvent;\r
 \r
   if (!SingleThread) {\r
-    WakeUpAP (CpuMpData, TRUE, 0, Procedure, ProcedureArgument);\r
+    WakeUpAP (CpuMpData, TRUE, 0, Procedure, ProcedureArgument, FALSE);\r
   } else {\r
     for (ProcessorNumber = 0; ProcessorNumber < ProcessorCount; ProcessorNumber++) {\r
       if (ProcessorNumber == CallerNumber) {\r
         continue;\r
       }\r
+\r
       if (CpuMpData->CpuData[ProcessorNumber].Waiting) {\r
-        WakeUpAP (CpuMpData, FALSE, ProcessorNumber, Procedure, ProcedureArgument);\r
+        WakeUpAP (CpuMpData, FALSE, ProcessorNumber, Procedure, ProcedureArgument, TRUE);\r
         break;\r
       }\r
     }\r
   }\r
 \r
+  if (!ExcludeBsp) {\r
+    //\r
+    // Start BSP.\r
+    //\r
+    Procedure (ProcedureArgument);\r
+  }\r
+\r
   Status = EFI_SUCCESS;\r
   if (WaitEvent == NULL) {\r
     do {\r
@@ -2289,18 +2676,18 @@ StartupAllAPsWorker (
 **/\r
 EFI_STATUS\r
 StartupThisAPWorker (\r
-  IN  EFI_AP_PROCEDURE          Procedure,\r
-  IN  UINTN                     ProcessorNumber,\r
-  IN  EFI_EVENT                 WaitEvent               OPTIONAL,\r
-  IN  UINTN                     TimeoutInMicroseconds,\r
-  IN  VOID                      *ProcedureArgument      OPTIONAL,\r
-  OUT BOOLEAN                   *Finished               OPTIONAL\r
+  IN  EFI_AP_PROCEDURE  Procedure,\r
+  IN  UINTN             ProcessorNumber,\r
+  IN  EFI_EVENT         WaitEvent               OPTIONAL,\r
+  IN  UINTN             TimeoutInMicroseconds,\r
+  IN  VOID              *ProcedureArgument      OPTIONAL,\r
+  OUT BOOLEAN           *Finished               OPTIONAL\r
   )\r
 {\r
-  EFI_STATUS              Status;\r
-  CPU_MP_DATA             *CpuMpData;\r
-  CPU_AP_DATA             *CpuData;\r
-  UINTN                   CallerNumber;\r
+  EFI_STATUS   Status;\r
+  CPU_MP_DATA  *CpuMpData;\r
+  CPU_AP_DATA  *CpuData;\r
+  UINTN        CallerNumber;\r
 \r
   CpuMpData = GetCpuMpData ();\r
 \r
@@ -2354,13 +2741,13 @@ StartupThisAPWorker (
   // BSP saves data for CheckAPsStatus(), and returns EFI_SUCCESS.\r
   // CheckAPsStatus() will check completion and timeout periodically.\r
   //\r
-  CpuData = &CpuMpData->CpuData[ProcessorNumber];\r
+  CpuData               = &CpuMpData->CpuData[ProcessorNumber];\r
   CpuData->WaitEvent    = WaitEvent;\r
   CpuData->Finished     = Finished;\r
   CpuData->ExpectedTime = CalculateTimeout (TimeoutInMicroseconds, &CpuData->CurrentTime);\r
   CpuData->TotalTime    = 0;\r
 \r
-  WakeUpAP (CpuMpData, FALSE, ProcessorNumber, Procedure, ProcedureArgument);\r
+  WakeUpAP (CpuMpData, FALSE, ProcessorNumber, Procedure, ProcedureArgument, TRUE);\r
 \r
   //\r
   // If WaitEvent is NULL, execute in blocking mode.\r
@@ -2386,16 +2773,128 @@ GetCpuMpDataFromGuidedHob (
   VOID\r
   )\r
 {\r
-  EFI_HOB_GUID_TYPE       *GuidHob;\r
-  VOID                    *DataInHob;\r
-  CPU_MP_DATA             *CpuMpData;\r
+  EFI_HOB_GUID_TYPE  *GuidHob;\r
+  VOID               *DataInHob;\r
+  CPU_MP_DATA        *CpuMpData;\r
 \r
   CpuMpData = NULL;\r
-  GuidHob = GetFirstGuidHob (&mCpuInitMpLibHobGuid);\r
+  GuidHob   = GetFirstGuidHob (&mCpuInitMpLibHobGuid);\r
   if (GuidHob != NULL) {\r
     DataInHob = GET_GUID_HOB_DATA (GuidHob);\r
-    CpuMpData = (CPU_MP_DATA *) (*(UINTN *) DataInHob);\r
+    CpuMpData = (CPU_MP_DATA *)(*(UINTN *)DataInHob);\r
   }\r
+\r
   return CpuMpData;\r
 }\r
 \r
+/**\r
+  This service executes a caller provided function on all enabled CPUs.\r
+\r
+  @param[in]  Procedure               A pointer to the function to be run on\r
+                                      enabled APs of the system. See type\r
+                                      EFI_AP_PROCEDURE.\r
+  @param[in]  TimeoutInMicroseconds   Indicates the time limit in microseconds for\r
+                                      APs to return from Procedure, either for\r
+                                      blocking or non-blocking mode. Zero means\r
+                                      infinity. TimeoutInMicroseconds is ignored\r
+                                      for BSP.\r
+  @param[in]  ProcedureArgument       The parameter passed into Procedure for\r
+                                      all APs.\r
+\r
+  @retval EFI_SUCCESS             In blocking mode, all CPUs have finished before\r
+                                  the timeout expired.\r
+  @retval EFI_SUCCESS             In non-blocking mode, function has been dispatched\r
+                                  to all enabled CPUs.\r
+  @retval EFI_DEVICE_ERROR        Caller processor is AP.\r
+  @retval EFI_NOT_READY           Any enabled APs are busy.\r
+  @retval EFI_NOT_READY           MP Initialize Library is not initialized.\r
+  @retval EFI_TIMEOUT             In blocking mode, the timeout expired before\r
+                                  all enabled APs have finished.\r
+  @retval EFI_INVALID_PARAMETER   Procedure is NULL.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+MpInitLibStartupAllCPUs (\r
+  IN  EFI_AP_PROCEDURE  Procedure,\r
+  IN  UINTN             TimeoutInMicroseconds,\r
+  IN  VOID              *ProcedureArgument      OPTIONAL\r
+  )\r
+{\r
+  return StartupAllCPUsWorker (\r
+           Procedure,\r
+           FALSE,\r
+           FALSE,\r
+           NULL,\r
+           TimeoutInMicroseconds,\r
+           ProcedureArgument,\r
+           NULL\r
+           );\r
+}\r
+\r
+/**\r
+  The function check if the specified Attr is set.\r
+\r
+  @param[in]  CurrentAttr   The current attribute.\r
+  @param[in]  Attr          The attribute to check.\r
+\r
+  @retval  TRUE      The specified Attr is set.\r
+  @retval  FALSE     The specified Attr is not set.\r
+\r
+**/\r
+STATIC\r
+BOOLEAN\r
+AmdMemEncryptionAttrCheck (\r
+  IN  UINT64                             CurrentAttr,\r
+  IN  CONFIDENTIAL_COMPUTING_GUEST_ATTR  Attr\r
+  )\r
+{\r
+  switch (Attr) {\r
+    case CCAttrAmdSev:\r
+      //\r
+      // SEV is automatically enabled if SEV-ES or SEV-SNP is active.\r
+      //\r
+      return CurrentAttr >= CCAttrAmdSev;\r
+    case CCAttrAmdSevEs:\r
+      //\r
+      // SEV-ES is automatically enabled if SEV-SNP is active.\r
+      //\r
+      return CurrentAttr >= CCAttrAmdSevEs;\r
+    case CCAttrAmdSevSnp:\r
+      return CurrentAttr == CCAttrAmdSevSnp;\r
+    default:\r
+      return FALSE;\r
+  }\r
+}\r
+\r
+/**\r
+  Check if the specified confidential computing attribute is active.\r
+\r
+  @param[in]  Attr          The attribute to check.\r
+\r
+  @retval TRUE   The specified Attr is active.\r
+  @retval FALSE  The specified Attr is not active.\r
+\r
+**/\r
+BOOLEAN\r
+EFIAPI\r
+ConfidentialComputingGuestHas (\r
+  IN  CONFIDENTIAL_COMPUTING_GUEST_ATTR  Attr\r
+  )\r
+{\r
+  UINT64  CurrentAttr;\r
+\r
+  //\r
+  // Get the current CC attribute.\r
+  //\r
+  CurrentAttr = PcdGet64 (PcdConfidentialComputingGuestAttr);\r
+\r
+  //\r
+  // If attr is for the AMD group then call AMD specific checks.\r
+  //\r
+  if (((RShiftU64 (CurrentAttr, 8)) & 0xff) == 1) {\r
+    return AmdMemEncryptionAttrCheck (CurrentAttr, Attr);\r
+  }\r
+\r
+  return (CurrentAttr == Attr);\r
+}\r