UefiCpuPkg/MpInitLib: move SEV specific routines in AmdSev.c

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

Move all the SEV specific function in AmdSev.c.

No functional change intended.

Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Cc: Michael Roth <michael.roth@amd.com>
Cc: James Bottomley <jejb@linux.ibm.com>
Cc: Min Xu <min.m.xu@intel.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Tom Lendacky <thomas.lendacky@amd.com>
Cc: Jordan Justen <jordan.l.justen@intel.com>
Cc: Ard Biesheuvel <ardb+tianocore@kernel.org>
Cc: Erdem Aktas <erdemaktas@google.com>
Cc: Gerd Hoffmann <kraxel@redhat.com>
Reviewed-by: Ray Ni <ray.ni@intel.com>
Acked-by: Gerd Hoffmann <kraxel@redhat.com>
Suggested-by: Jiewen Yao <Jiewen.yao@intel.com>
Signed-off-by: Brijesh Singh <brijesh.singh@amd.com>

7 files changed, 421 insertions, 306 deletions

diff --git a/UefiCpuPkg/Library/MpInitLib/AmdSev.c b/UefiCpuPkg/Library/MpInitLib/AmdSev.c
new file mode 100644
index 0000000000..0e3c6e2310
--- /dev/null
+++ b/UefiCpuPkg/Library/MpInitLib/AmdSev.c
@@ -0,0 +1,245 @@
+/** @file

+  CPU MP Initialize helper function for AMD SEV.

+

+  Copyright (c) 2021, AMD Inc. All rights reserved.<BR>

+

+  SPDX-License-Identifier: BSD-2-Clause-Patent

+

+**/

+

+#include "MpLib.h"

+#include <Library/VmgExitLib.h>

+

+/**

+  Get Protected mode code segment with 16-bit default addressing

+  from current GDT table.

+

+  @return  Protected mode 16-bit code segment value.

+**/

+STATIC

+UINT16

+GetProtectedMode16CS (

+  VOID

+  )

+{

+  IA32_DESCRIPTOR          GdtrDesc;

+  IA32_SEGMENT_DESCRIPTOR  *GdtEntry;

+  UINTN                    GdtEntryCount;

+  UINT16                   Index;

+

+  Index = (UINT16)-1;

+  AsmReadGdtr (&GdtrDesc);

+  GdtEntryCount = (GdtrDesc.Limit + 1) / sizeof (IA32_SEGMENT_DESCRIPTOR);

+  GdtEntry      = (IA32_SEGMENT_DESCRIPTOR *)GdtrDesc.Base;

+  for (Index = 0; Index < GdtEntryCount; Index++) {

+    if ((GdtEntry->Bits.L == 0) &&

+        (GdtEntry->Bits.DB == 0) &&

+        (GdtEntry->Bits.Type > 8))

+    {

+      break;

+    }

+

+    GdtEntry++;

+  }

+

+  ASSERT (Index != GdtEntryCount);

+  return Index * 8;

+}

+

+/**

+  Get Protected mode code segment with 32-bit default addressing

+  from current GDT table.

+

+  @return  Protected mode 32-bit code segment value.

+**/

+STATIC

+UINT16

+GetProtectedMode32CS (

+  VOID

+  )

+{

+  IA32_DESCRIPTOR          GdtrDesc;

+  IA32_SEGMENT_DESCRIPTOR  *GdtEntry;

+  UINTN                    GdtEntryCount;

+  UINT16                   Index;

+

+  Index = (UINT16)-1;

+  AsmReadGdtr (&GdtrDesc);

+  GdtEntryCount = (GdtrDesc.Limit + 1) / sizeof (IA32_SEGMENT_DESCRIPTOR);

+  GdtEntry      = (IA32_SEGMENT_DESCRIPTOR *)GdtrDesc.Base;

+  for (Index = 0; Index < GdtEntryCount; Index++) {

+    if ((GdtEntry->Bits.L == 0) &&

+        (GdtEntry->Bits.DB == 1) &&

+        (GdtEntry->Bits.Type > 8))

+    {

+      break;

+    }

+

+    GdtEntry++;

+  }

+

+  ASSERT (Index != GdtEntryCount);

+  return Index * 8;

+}

+

+/**

+  Reset an AP when in SEV-ES mode.

+

+  If successful, this function never returns.

+

+  @param[in] Ghcb                 Pointer to the GHCB

+  @param[in] CpuMpData            Pointer to CPU MP Data

+

+**/

+VOID

+MpInitLibSevEsAPReset (

+  IN GHCB         *Ghcb,

+  IN CPU_MP_DATA  *CpuMpData

+  )

+{

+  EFI_STATUS  Status;

+  UINTN       ProcessorNumber;

+  UINT16      Code16, Code32;

+  AP_RESET    *APResetFn;

+  UINTN       BufferStart;

+  UINTN       StackStart;

+

+  Status = GetProcessorNumber (CpuMpData, &ProcessorNumber);

+  ASSERT_EFI_ERROR (Status);

+

+  Code16 = GetProtectedMode16CS ();

+  Code32 = GetProtectedMode32CS ();

+

+  if (CpuMpData->WakeupBufferHigh != 0) {

+    APResetFn = (AP_RESET *)(CpuMpData->WakeupBufferHigh + CpuMpData->AddressMap.SwitchToRealNoNxOffset);

+  } else {

+    APResetFn = (AP_RESET *)(CpuMpData->MpCpuExchangeInfo->BufferStart + CpuMpData->AddressMap.SwitchToRealOffset);

+  }

+

+  BufferStart = CpuMpData->MpCpuExchangeInfo->BufferStart;

+  StackStart  = CpuMpData->SevEsAPResetStackStart -

+                (AP_RESET_STACK_SIZE * ProcessorNumber);

+

+  //

+  // This call never returns.

+  //

+  APResetFn (BufferStart, Code16, Code32, StackStart);

+}

+

+/**

+  Allocate the SEV-ES AP jump table buffer.

+

+  @param[in, out]  CpuMpData  The pointer to CPU MP Data structure.

+**/

+VOID

+AllocateSevEsAPMemory (

+  IN OUT CPU_MP_DATA  *CpuMpData

+  )

+{

+  if (CpuMpData->SevEsAPBuffer == (UINTN)-1) {

+    CpuMpData->SevEsAPBuffer =

+      CpuMpData->SevEsIsEnabled ? GetSevEsAPMemory () : 0;

+  }

+}

+

+/**

+  Program the SEV-ES AP jump table buffer.

+

+  @param[in]  SipiVector  The SIPI vector used for the AP Reset

+**/

+VOID

+SetSevEsJumpTable (

+  IN UINTN  SipiVector

+  )

+{

+  SEV_ES_AP_JMP_FAR  *JmpFar;

+  UINT32             Offset, InsnByte;

+  UINT8              LoNib, HiNib;

+

+  JmpFar = (SEV_ES_AP_JMP_FAR *)(UINTN)FixedPcdGet32 (PcdSevEsWorkAreaBase);

+  ASSERT (JmpFar != NULL);

+

+  //

+  // Obtain the address of the Segment/Rip location in the workarea.

+  // This will be set to a value derived from the SIPI vector and will

+  // be the memory address used for the far jump below.

+  //

+  Offset  = FixedPcdGet32 (PcdSevEsWorkAreaBase);

+  Offset += sizeof (JmpFar->InsnBuffer);

+  LoNib   = (UINT8)Offset;

+  HiNib   = (UINT8)(Offset >> 8);

+

+  //

+  // Program the workarea (which is the initial AP boot address) with

+  // far jump to the SIPI vector (where XX and YY represent the

+  // address of where the SIPI vector is stored.

+  //

+  //   JMP FAR [CS:XXYY] => 2E FF 2E YY XX

+  //

+  InsnByte                       = 0;

+  JmpFar->InsnBuffer[InsnByte++] = 0x2E;  // CS override prefix

+  JmpFar->InsnBuffer[InsnByte++] = 0xFF;  // JMP (FAR)

+  JmpFar->InsnBuffer[InsnByte++] = 0x2E;  // ModRM (JMP memory location)

+  JmpFar->InsnBuffer[InsnByte++] = LoNib; // YY offset ...

+  JmpFar->InsnBuffer[InsnByte++] = HiNib; // XX offset ...

+

+  //

+  // Program the Segment/Rip based on the SIPI vector (always at least

+  // 16-byte aligned, so Rip is set to 0).

+  //

+  JmpFar->Rip     = 0;

+  JmpFar->Segment = (UINT16)(SipiVector >> 4);

+}

+

+/**

+  The function puts the AP in halt loop.

+

+  @param[in]  CpuMpData  The pointer to CPU MP Data structure.

+**/

+VOID

+SevEsPlaceApHlt (

+  CPU_MP_DATA  *CpuMpData

+  )

+{

+  MSR_SEV_ES_GHCB_REGISTER  Msr;

+  GHCB                      *Ghcb;

+  UINT64                    Status;

+  BOOLEAN                   DoDecrement;

+  BOOLEAN                   InterruptState;

+

+  DoDecrement = (BOOLEAN)(CpuMpData->InitFlag == ApInitConfig);

+

+  while (TRUE) {

+    Msr.GhcbPhysicalAddress = AsmReadMsr64 (MSR_SEV_ES_GHCB);

+    Ghcb                    = Msr.Ghcb;

+

+    VmgInit (Ghcb, &InterruptState);

+

+    if (DoDecrement) {

+      DoDecrement = FALSE;

+

+      //

+      // Perform the delayed decrement just before issuing the first

+      // VMGEXIT with AP_RESET_HOLD.

+      //

+      InterlockedDecrement ((UINT32 *)&CpuMpData->MpCpuExchangeInfo->NumApsExecuting);

+    }

+

+    Status = VmgExit (Ghcb, SVM_EXIT_AP_RESET_HOLD, 0, 0);

+    if ((Status == 0) && (Ghcb->SaveArea.SwExitInfo2 != 0)) {

+      VmgDone (Ghcb, InterruptState);

+      break;

+    }

+

+    VmgDone (Ghcb, InterruptState);

+  }

+

+  //

+  // Awakened in a new phase? Use the new CpuMpData

+  //

+  if (CpuMpData->NewCpuMpData != NULL) {

+    CpuMpData = CpuMpData->NewCpuMpData;

+  }

+

+  MpInitLibSevEsAPReset (Ghcb, CpuMpData);

+}

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf b/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf
index d34419c2a5..6e510aa891 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf
@@ -28,6 +28,7 @@
   X64/MpFuncs.nasm

 

 [Sources.common]

+  AmdSev.c

   MpEqu.inc

   DxeMpLib.c

   MpLib.c

diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index 2d8e9512ab..34555c0693 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -600,123 +600,6 @@ InitializeApData (
 }

 

 /**

-  Get Protected mode code segment with 16-bit default addressing

-  from current GDT table.

-

-  @return  Protected mode 16-bit code segment value.

-**/

-STATIC

-UINT16

-GetProtectedMode16CS (

-  VOID

-  )

-{

-  IA32_DESCRIPTOR          GdtrDesc;

-  IA32_SEGMENT_DESCRIPTOR  *GdtEntry;

-  UINTN                    GdtEntryCount;

-  UINT16                   Index;

-

-  Index = (UINT16)-1;

-  AsmReadGdtr (&GdtrDesc);

-  GdtEntryCount = (GdtrDesc.Limit + 1) / sizeof (IA32_SEGMENT_DESCRIPTOR);

-  GdtEntry      = (IA32_SEGMENT_DESCRIPTOR *)GdtrDesc.Base;

-  for (Index = 0; Index < GdtEntryCount; Index++) {

-    if ((GdtEntry->Bits.L == 0) &&

-        (GdtEntry->Bits.DB == 0) &&

-        (GdtEntry->Bits.Type > 8))

-    {

-      break;

-    }

-

-    GdtEntry++;

-  }

-

-  ASSERT (Index != GdtEntryCount);

-  return Index * 8;

-}

-

-/**

-  Get Protected mode code segment with 32-bit default addressing

-  from current GDT table.

-

-  @return  Protected mode 32-bit code segment value.

-**/

-STATIC

-UINT16

-GetProtectedMode32CS (

-  VOID

-  )

-{

-  IA32_DESCRIPTOR          GdtrDesc;

-  IA32_SEGMENT_DESCRIPTOR  *GdtEntry;

-  UINTN                    GdtEntryCount;

-  UINT16                   Index;

-

-  Index = (UINT16)-1;

-  AsmReadGdtr (&GdtrDesc);

-  GdtEntryCount = (GdtrDesc.Limit + 1) / sizeof (IA32_SEGMENT_DESCRIPTOR);

-  GdtEntry      = (IA32_SEGMENT_DESCRIPTOR *)GdtrDesc.Base;

-  for (Index = 0; Index < GdtEntryCount; Index++) {

-    if ((GdtEntry->Bits.L == 0) &&

-        (GdtEntry->Bits.DB == 1) &&

-        (GdtEntry->Bits.Type > 8))

-    {

-      break;

-    }

-

-    GdtEntry++;

-  }

-

-  ASSERT (Index != GdtEntryCount);

-  return Index * 8;

-}

-

-/**

-  Reset an AP when in SEV-ES mode.

-

-  If successful, this function never returns.

-

-  @param[in] Ghcb                 Pointer to the GHCB

-  @param[in] CpuMpData            Pointer to CPU MP Data

-

-**/

-STATIC

-VOID

-MpInitLibSevEsAPReset (

-  IN GHCB         *Ghcb,

-  IN CPU_MP_DATA  *CpuMpData

-  )

-{

-  EFI_STATUS  Status;

-  UINTN       ProcessorNumber;

-  UINT16      Code16, Code32;

-  AP_RESET    *APResetFn;

-  UINTN       BufferStart;

-  UINTN       StackStart;

-

-  Status = GetProcessorNumber (CpuMpData, &ProcessorNumber);

-  ASSERT_EFI_ERROR (Status);

-

-  Code16 = GetProtectedMode16CS ();

-  Code32 = GetProtectedMode32CS ();

-

-  if (CpuMpData->WakeupBufferHigh != 0) {

-    APResetFn = (AP_RESET *)(CpuMpData->WakeupBufferHigh + CpuMpData->AddressMap.SwitchToRealNoNxOffset);

-  } else {

-    APResetFn = (AP_RESET *)(CpuMpData->MpCpuExchangeInfo->BufferStart + CpuMpData->AddressMap.SwitchToRealOffset);

-  }

-

-  BufferStart = CpuMpData->MpCpuExchangeInfo->BufferStart;

-  StackStart  = CpuMpData->SevEsAPResetStackStart -

-                (AP_RESET_STACK_SIZE * ProcessorNumber);

-

-  //

-  // This call never returns.

-  //

-  APResetFn (BufferStart, Code16, Code32, StackStart);

-}

-

-/**

   This function will be called from AP reset code if BSP uses WakeUpAP.

 

   @param[in] ExchangeInfo     Pointer to the MP exchange info buffer

@@ -895,47 +778,7 @@ ApWakeupFunction (
       while (TRUE) {

         DisableInterrupts ();

         if (CpuMpData->SevEsIsEnabled) {

-          MSR_SEV_ES_GHCB_REGISTER  Msr;

-          GHCB                      *Ghcb;

-          UINT64                    Status;

-          BOOLEAN                   DoDecrement;

-          BOOLEAN                   InterruptState;

-

-          DoDecrement = (BOOLEAN)(CpuMpData->InitFlag == ApInitConfig);

-

-          while (TRUE) {

-            Msr.GhcbPhysicalAddress = AsmReadMsr64 (MSR_SEV_ES_GHCB);

-            Ghcb                    = Msr.Ghcb;

-

-            VmgInit (Ghcb, &InterruptState);

-

-            if (DoDecrement) {

-              DoDecrement = FALSE;

-

-              //

-              // Perform the delayed decrement just before issuing the first

-              // VMGEXIT with AP_RESET_HOLD.

-              //

-              InterlockedDecrement ((UINT32 *)&CpuMpData->MpCpuExchangeInfo->NumApsExecuting);

-            }

-

-            Status = VmgExit (Ghcb, SVM_EXIT_AP_RESET_HOLD, 0, 0);

-            if ((Status == 0) && (Ghcb->SaveArea.SwExitInfo2 != 0)) {

-              VmgDone (Ghcb, InterruptState);

-              break;

-            }

-

-            VmgDone (Ghcb, InterruptState);

-          }

-

-          //

-          // Awakened in a new phase? Use the new CpuMpData

-          //

-          if (CpuMpData->NewCpuMpData != NULL) {

-            CpuMpData = CpuMpData->NewCpuMpData;

-          }

-

-          MpInitLibSevEsAPReset (Ghcb, CpuMpData);

+          SevEsPlaceApHlt (CpuMpData);

         } else {

           CpuSleep ();

         }

@@ -1269,71 +1112,6 @@ FreeResetVector (
 }

 

 /**

-  Allocate the SEV-ES AP jump table buffer.

-

-  @param[in, out]  CpuMpData  The pointer to CPU MP Data structure.

-**/

-VOID

-AllocateSevEsAPMemory (

-  IN OUT CPU_MP_DATA  *CpuMpData

-  )

-{

-  if (CpuMpData->SevEsAPBuffer == (UINTN)-1) {

-    CpuMpData->SevEsAPBuffer =

-      CpuMpData->SevEsIsEnabled ? GetSevEsAPMemory () : 0;

-  }

-}

-

-/**

-  Program the SEV-ES AP jump table buffer.

-

-  @param[in]  SipiVector  The SIPI vector used for the AP Reset

-**/

-VOID

-SetSevEsJumpTable (

-  IN UINTN  SipiVector

-  )

-{

-  SEV_ES_AP_JMP_FAR  *JmpFar;

-  UINT32             Offset, InsnByte;

-  UINT8              LoNib, HiNib;

-

-  JmpFar = (SEV_ES_AP_JMP_FAR *)(UINTN)FixedPcdGet32 (PcdSevEsWorkAreaBase);

-  ASSERT (JmpFar != NULL);

-

-  //

-  // Obtain the address of the Segment/Rip location in the workarea.

-  // This will be set to a value derived from the SIPI vector and will

-  // be the memory address used for the far jump below.

-  //

-  Offset  = FixedPcdGet32 (PcdSevEsWorkAreaBase);

-  Offset += sizeof (JmpFar->InsnBuffer);

-  LoNib   = (UINT8)Offset;

-  HiNib   = (UINT8)(Offset >> 8);

-

-  //

-  // Program the workarea (which is the initial AP boot address) with

-  // far jump to the SIPI vector (where XX and YY represent the

-  // address of where the SIPI vector is stored.

-  //

-  //   JMP FAR [CS:XXYY] => 2E FF 2E YY XX

-  //

-  InsnByte                       = 0;

-  JmpFar->InsnBuffer[InsnByte++] = 0x2E;  // CS override prefix

-  JmpFar->InsnBuffer[InsnByte++] = 0xFF;  // JMP (FAR)

-  JmpFar->InsnBuffer[InsnByte++] = 0x2E;  // ModRM (JMP memory location)

-  JmpFar->InsnBuffer[InsnByte++] = LoNib; // YY offset ...

-  JmpFar->InsnBuffer[InsnByte++] = HiNib; // XX offset ...

-

-  //

-  // Program the Segment/Rip based on the SIPI vector (always at least

-  // 16-byte aligned, so Rip is set to 0).

-  //

-  JmpFar->Rip     = 0;

-  JmpFar->Segment = (UINT16)(SipiVector >> 4);

-}

-

-/**

   This function will be called by BSP to wakeup AP.

 

   @param[in] CpuMpData          Pointer to CPU MP Data

diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index a647772088..16b4d76d01 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -34,6 +34,9 @@
 #include <Library/PcdLib.h>

 #include <Library/MicrocodeLib.h>

 

+#include <Register/Amd/Fam17Msr.h>

+#include <Register/Amd/Ghcb.h>

+

 #include <Guid/MicrocodePatchHob.h>

 

 #define WAKEUP_AP_SIGNAL  SIGNATURE_32 ('S', 'T', 'A', 'P')

@@ -321,7 +324,7 @@ typedef struct {
                            from long mode to real mode.

 **/

 typedef

-VOID

+  VOID

 (EFIAPI AP_RESET)(

   IN UINTN    BufferStart,

   IN UINT16   Code16,

@@ -346,7 +349,7 @@ extern EFI_GUID  mCpuInitMpLibHobGuid;
   @param[in] PmCodeSegment   Protected mode code segment value.

 **/

 typedef

-VOID

+  VOID

 (EFIAPI *ASM_RELOCATE_AP_LOOP)(

   IN BOOLEAN                 MwaitSupport,

   IN UINTN                   ApTargetCState,

@@ -740,4 +743,34 @@ PlatformShadowMicrocode (
   IN OUT CPU_MP_DATA  *CpuMpData

   );

 

+/**

+  Allocate the SEV-ES AP jump table buffer.

+

+  @param[in, out]  CpuMpData  The pointer to CPU MP Data structure.

+**/

+VOID

+AllocateSevEsAPMemory (

+  IN OUT CPU_MP_DATA  *CpuMpData

+  );

+

+/**

+  Program the SEV-ES AP jump table buffer.

+

+  @param[in]  SipiVector  The SIPI vector used for the AP Reset

+**/

+VOID

+SetSevEsJumpTable (

+  IN UINTN  SipiVector

+  );

+

+/**

+  The function puts the AP in halt loop.

+

+  @param[in]  CpuMpData  The pointer to CPU MP Data structure.

+**/

+VOID

+SevEsPlaceApHlt (

+  CPU_MP_DATA  *CpuMpData

+  );

+

 #endif

diff --git a/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf b/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf
index 36fcb96b58..2cbd9b8b8a 100644
--- a/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf
+++ b/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf
@@ -28,6 +28,7 @@
   X64/MpFuncs.nasm

 

 [Sources.common]

+  AmdSev.c

   MpEqu.inc

   PeiMpLib.c

   MpLib.c

diff --git a/UefiCpuPkg/Library/MpInitLib/X64/AmdSev.nasm b/UefiCpuPkg/Library/MpInitLib/X64/AmdSev.nasm
new file mode 100644
index 0000000000..0ccafe25ec
--- /dev/null
+++ b/UefiCpuPkg/Library/MpInitLib/X64/AmdSev.nasm
@@ -0,0 +1,119 @@
+;------------------------------------------------------------------------------ ;

+; Copyright (c) 2021, AMD Inc. All rights reserved.<BR>

+; SPDX-License-Identifier: BSD-2-Clause-Patent

+;

+; Module Name:

+;

+;   AmdSev.nasm

+;

+; Abstract:

+;

+;   This provides helper used by the MpFunc.nasm. If AMD SEV-ES is active

+;   then helpers perform the additional setups (such as GHCB).

+;

+;-------------------------------------------------------------------------------

+

+%define SIZE_4KB    0x1000

+

+;

+; The function checks whether SEV-ES is enabled, if enabled

+; then setup the GHCB page.

+;

+SevEsSetupGhcb:

+    lea        edi, [esi + MP_CPU_EXCHANGE_INFO_FIELD (SevEsIsEnabled)]

+    cmp        byte [edi], 1        ; SevEsIsEnabled

+    jne        SevEsSetupGhcbExit

+

+    ;

+    ; program GHCB

+    ;   Each page after the GHCB is a per-CPU page, so the calculation programs

+    ;   a GHCB to be every 8KB.

+    ;

+    mov        eax, SIZE_4KB

+    shl        eax, 1                            ; EAX = SIZE_4K * 2

+    mov        ecx, ebx

+    mul        ecx                               ; EAX = SIZE_4K * 2 * CpuNumber

+    mov        edi, esi

+    add        edi, MP_CPU_EXCHANGE_INFO_FIELD (GhcbBase)

+    add        rax, qword [edi]

+    mov        rdx, rax

+    shr        rdx, 32

+    mov        rcx, 0xc0010130

+    wrmsr

+

+SevEsSetupGhcbExit:

+    OneTimeCallRet    SevEsSetupGhcb

+

+;

+; The function checks whether SEV-ES is enabled, if enabled, use

+; the GHCB

+;

+SevEsGetApicId:

+    lea        edi, [esi + MP_CPU_EXCHANGE_INFO_FIELD (SevEsIsEnabled)]

+    cmp        byte [edi], 1        ; SevEsIsEnabled

+    jne        SevEsGetApicIdExit

+

+    ;

+    ; Since we don't have a stack yet, we can't take a #VC

+    ; exception. Use the GHCB protocol to perform the CPUID

+    ; calls.

+    ;

+    mov        rcx, 0xc0010130

+    rdmsr

+    shl        rdx, 32

+    or         rax, rdx

+    mov        rdi, rax             ; RDI now holds the original GHCB GPA

+

+    mov        rdx, 0               ; CPUID function 0

+    mov        rax, 0               ; RAX register requested

+    or         rax, 4

+    wrmsr

+    rep vmmcall

+    rdmsr

+    cmp        edx, 0bh

+    jb         NoX2ApicSevEs        ; CPUID level below CPUID_EXTENDED_TOPOLOGY

+

+    mov        rdx, 0bh             ; CPUID function 0x0b

+    mov        rax, 040000000h      ; RBX register requested

+    or         rax, 4

+    wrmsr

+    rep vmmcall

+    rdmsr

+    test       edx, 0ffffh

+    jz         NoX2ApicSevEs        ; CPUID.0BH:EBX[15:0] is zero

+

+    mov        rdx, 0bh             ; CPUID function 0x0b

+    mov        rax, 0c0000000h      ; RDX register requested

+    or         rax, 4

+    wrmsr

+    rep vmmcall

+    rdmsr

+

+    ; Processor is x2APIC capable; 32-bit x2APIC ID is now in EDX

+    jmp        RestoreGhcb

+

+NoX2ApicSevEs:

+    ; Processor is not x2APIC capable, so get 8-bit APIC ID

+    mov        rdx, 1               ; CPUID function 1

+    mov        rax, 040000000h      ; RBX register requested

+    or         rax, 4

+    wrmsr

+    rep vmmcall

+    rdmsr

+    shr        edx, 24

+

+RestoreGhcb:

+    mov        rbx, rdx             ; Save x2APIC/APIC ID

+

+    mov        rdx, rdi             ; RDI holds the saved GHCB GPA

+    shr        rdx, 32

+    mov        eax, edi

+    wrmsr

+

+    mov        rdx, rbx

+

+    ; x2APIC ID or APIC ID is in EDX

+    jmp        GetProcessorNumber

+

+SevEsGetApicIdExit:

+    OneTimeCallRet    SevEsGetApicId

diff --git a/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm b/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
index 50df802d1f..f7f2937faf 100644
--- a/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
+++ b/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
@@ -15,6 +15,15 @@
 %include "MpEqu.inc"

 extern ASM_PFX(InitializeFloatingPointUnits)

 

+%macro  OneTimeCall 1

+    jmp     %1

+%1 %+ OneTimerCallReturn:

+%endmacro

+

+%macro  OneTimeCallRet 1

+    jmp     %1 %+ OneTimerCallReturn

+%endmacro

+

 DEFAULT REL

 

 SECTION .text

@@ -144,6 +153,12 @@ SkipEnable5LevelPaging:
     jmp far    [edi]

 

 BITS 64

+

+;

+; Required for the AMD SEV helper functions

+;

+%include "AmdSev.nasm"

+

 LongModeStart:

     mov        esi, ebx

     lea        edi, [esi + MP_CPU_EXCHANGE_INFO_FIELD (InitFlag)]

@@ -175,94 +190,17 @@ LongModeStart:
     add        rax, qword [edi]

     mov        rsp, rax

 

-    lea        edi, [esi + MP_CPU_EXCHANGE_INFO_FIELD (SevEsIsEnabled)]

-    cmp        byte [edi], 1        ; SevEsIsEnabled

-    jne        CProcedureInvoke

-

     ;

-    ; program GHCB

-    ;   Each page after the GHCB is a per-CPU page, so the calculation programs

-    ;   a GHCB to be every 8KB.

+    ;  Setup the GHCB when AMD SEV-ES active.

     ;

-    mov        eax, SIZE_4KB

-    shl        eax, 1                            ; EAX = SIZE_4K * 2

-    mov        ecx, ebx

-    mul        ecx                               ; EAX = SIZE_4K * 2 * CpuNumber

-    mov        edi, esi

-    add        edi, MP_CPU_EXCHANGE_INFO_FIELD (GhcbBase)

-    add        rax, qword [edi]

-    mov        rdx, rax

-    shr        rdx, 32

-    mov        rcx, 0xc0010130

-    wrmsr

+    OneTimeCall SevEsSetupGhcb

     jmp        CProcedureInvoke

 

 GetApicId:

-    lea        edi, [esi + MP_CPU_EXCHANGE_INFO_FIELD (SevEsIsEnabled)]

-    cmp        byte [edi], 1        ; SevEsIsEnabled

-    jne        DoCpuid

-

     ;

-    ; Since we don't have a stack yet, we can't take a #VC

-    ; exception. Use the GHCB protocol to perform the CPUID

-    ; calls.

+    ; Use the GHCB protocol to get the ApicId when SEV-ES is active.

     ;

-    mov        rcx, 0xc0010130

-    rdmsr

-    shl        rdx, 32

-    or         rax, rdx

-    mov        rdi, rax             ; RDI now holds the original GHCB GPA

-

-    mov        rdx, 0               ; CPUID function 0

-    mov        rax, 0               ; RAX register requested

-    or         rax, 4

-    wrmsr

-    rep vmmcall

-    rdmsr

-    cmp        edx, 0bh

-    jb         NoX2ApicSevEs        ; CPUID level below CPUID_EXTENDED_TOPOLOGY

-

-    mov        rdx, 0bh             ; CPUID function 0x0b

-    mov        rax, 040000000h      ; RBX register requested

-    or         rax, 4

-    wrmsr

-    rep vmmcall

-    rdmsr

-    test       edx, 0ffffh

-    jz         NoX2ApicSevEs        ; CPUID.0BH:EBX[15:0] is zero

-

-    mov        rdx, 0bh             ; CPUID function 0x0b

-    mov        rax, 0c0000000h      ; RDX register requested

-    or         rax, 4

-    wrmsr

-    rep vmmcall

-    rdmsr

-

-    ; Processor is x2APIC capable; 32-bit x2APIC ID is now in EDX

-    jmp        RestoreGhcb

-

-NoX2ApicSevEs:

-    ; Processor is not x2APIC capable, so get 8-bit APIC ID

-    mov        rdx, 1               ; CPUID function 1

-    mov        rax, 040000000h      ; RBX register requested

-    or         rax, 4

-    wrmsr

-    rep vmmcall

-    rdmsr

-    shr        edx, 24

-

-RestoreGhcb:

-    mov        rbx, rdx             ; Save x2APIC/APIC ID

-

-    mov        rdx, rdi             ; RDI holds the saved GHCB GPA

-    shr        rdx, 32

-    mov        eax, edi

-    wrmsr

-

-    mov        rdx, rbx

-

-    ; x2APIC ID or APIC ID is in EDX

-    jmp        GetProcessorNumber

+    OneTimeCall SevEsGetApicId

 

 DoCpuid:

     mov        eax, 0