[PATCH v9 02/32] UefiCpuPkg/MpInitLib: move SEV specific routines in AmdSev.c

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

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>
Suggested-by: Jiewen Yao <Jiewen.yao@intel.com>
Signed-off-by: Brijesh Singh <brijesh.singh@amd.com>
---
 UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf |   1 +
 UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf |   1 +
 UefiCpuPkg/Library/MpInitLib/MpLib.h          |  33 +++
 UefiCpuPkg/Library/MpInitLib/AmdSev.c         | 239 ++++++++++++++++++
 UefiCpuPkg/Library/MpInitLib/MpLib.c          | 218 +---------------
 UefiCpuPkg/Library/MpInitLib/X64/AmdSev.nasm  | 119 +++++++++
 UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm | 100 ++------
 7 files changed, 413 insertions(+), 298 deletions(-)
 create mode 100644 UefiCpuPkg/Library/MpInitLib/AmdSev.c
 create mode 100644 UefiCpuPkg/Library/MpInitLib/X64/AmdSev.nasm

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf b/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf
index d34419c2a524..6e510aa89120 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpInitLib.inf
@@ -28,6 +28,7 @@ [Sources.X64]
   X64/MpFuncs.nasm
 
 [Sources.common]
+  AmdSev.c
   MpEqu.inc
   DxeMpLib.c
   MpLib.c
diff --git a/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf b/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf
index 36fcb96b5852..2cbd9b8b8acc 100644
--- a/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf
+++ b/UefiCpuPkg/Library/MpInitLib/PeiMpInitLib.inf
@@ -28,6 +28,7 @@ [Sources.X64]
   X64/MpFuncs.nasm
 
 [Sources.common]
+  AmdSev.c
   MpEqu.inc
   PeiMpLib.c
   MpLib.c
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index e88a5355c983..3d4446df8ce6 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')
@@ -741,5 +744,35 @@ 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/AmdSev.c b/UefiCpuPkg/Library/MpInitLib/AmdSev.c
new file mode 100644
index 000000000000..7dbf117c2b71
--- /dev/null
+++ b/UefiCpuPkg/Library/MpInitLib/AmdSev.c
@@ -0,0 +1,239 @@
+/** @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/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index b9a06747edbf..890945bc5994 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -596,117 +596,6 @@ InitializeApData (
   SetApState (&CpuMpData->CpuData[ProcessorNumber], CpuStateIdle);
 }
 
-/**
-  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.
 
@@ -884,47 +773,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 ();
         }
@@ -1252,71 +1101,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.
 
diff --git a/UefiCpuPkg/Library/MpInitLib/X64/AmdSev.nasm b/UefiCpuPkg/Library/MpInitLib/X64/AmdSev.nasm
new file mode 100644
index 000000000000..0ccafe25eca4
--- /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 50df802d1fca..f7f2937fafad 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
-- 
2.25.1