[PATCH 2/5] UefiCpuPkg: Contiguous memory allocation and code clean-up.

["Yuanhao Xie" <yuanhao.xie@intel.com>]

Contiguous memory allocation and
This patch includes the code refactoring to eliminate the duplication,
non-descriptive variable, etc.

The memory is calculated taking into account the size difference of
RelocateApLoopFunc under different cases.

Allocate the memory for stacks and AP loop at contiguous address.

Cc: Guo Dong <guo.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Sean Rhodes <sean@starlabs.systems>
Cc: James Lu <james.lu@intel.com>
Cc: Gua Guo <gua.guo@intel.com>
Signed-off-by: Yuanhao Xie <yuanhao.xie@intel.com>
---
 UefiCpuPkg/Library/MpInitLib/DxeMpLib.c | 157 ++++++++++++------------
 UefiCpuPkg/Library/MpInitLib/MpLib.h    |   9 ++
 2 files changed, 85 insertions(+), 81 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index fd94652e5b..7b51c5c7f7 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
@@ -20,14 +20,14 @@
 
 #define  AP_SAFE_STACK_SIZE  128
 
-CPU_MP_DATA       *mCpuMpData                  = NULL;
-EFI_EVENT         mCheckAllApsEvent            = NULL;
-EFI_EVENT         mMpInitExitBootServicesEvent = NULL;
-EFI_EVENT         mLegacyBootEvent             = NULL;
-volatile BOOLEAN  mStopCheckAllApsStatus       = TRUE;
-VOID              *mReservedApLoopFunc         = NULL;
-UINTN             mReservedTopOfApStack;
-volatile UINT32   mNumberToFinish = 0;
+CPU_MP_DATA             *mCpuMpData                  = NULL;
+EFI_EVENT               mCheckAllApsEvent            = NULL;
+EFI_EVENT               mMpInitExitBootServicesEvent = NULL;
+EFI_EVENT               mLegacyBootEvent             = NULL;
+volatile BOOLEAN        mStopCheckAllApsStatus       = TRUE;
+UINTN                   mReservedTopOfApStack;
+volatile UINT32         mNumberToFinish = 0;
+RELOCATE_AP_LOOP_ENTRY  mReservedApLoop;
 
 //
 // Begin wakeup buffer allocation below 0x88000
@@ -380,8 +380,6 @@ RelocateApLoop (
 {
   CPU_MP_DATA                 *CpuMpData;
   BOOLEAN                     MwaitSupport;
-  ASM_RELOCATE_AP_LOOP        AsmRelocateApLoopFunc;
-  ASM_RELOCATE_AP_LOOP_AMD64  AsmRelocateApLoopFuncAmd64;
   UINTN                       ProcessorNumber;
   UINTN                       StackStart;
 
@@ -390,30 +388,28 @@ RelocateApLoop (
   MwaitSupport = IsMwaitSupport ();
   if (StandardSignatureIsAuthenticAMD () && (sizeof (UINTN) == sizeof (UINT64))) {
     StackStart = CpuMpData->UseSevEsAPMethod ? CpuMpData->SevEsAPResetStackStart : mReservedTopOfApStack;
-    AsmRelocateApLoopFuncAmd64 = (ASM_RELOCATE_AP_LOOP)(UINTN)mReservedApLoopFunc;
-    AsmRelocateApLoopFuncAmd64 (
-      MwaitSupport,
-      CpuMpData->ApTargetCState,
-      CpuMpData->PmCodeSegment,
-      StackStart - ProcessorNumber * AP_SAFE_STACK_SIZE,
-      (UINTN)&mNumberToFinish,
-      CpuMpData->Pm16CodeSegment,
-      CpuMpData->SevEsAPBuffer,
-      CpuMpData->WakeupBuffer
-      );
+    mReservedApLoop.Amd64Entry (
+                      MwaitSupport,
+                      CpuMpData->ApTargetCState,
+                      CpuMpData->PmCodeSegment,
+                      StackStart - ProcessorNumber * AP_SAFE_STACK_SIZE,
+                      (UINTN)&mNumberToFinish,
+                      CpuMpData->Pm16CodeSegment,
+                      CpuMpData->SevEsAPBuffer,
+                      CpuMpData->WakeupBuffer
+                      );
   } else {
-    StackStart            = mReservedTopOfApStack;
-    AsmRelocateApLoopFunc = (ASM_RELOCATE_AP_LOOP)(UINTN)mReservedApLoopFunc;
-    AsmRelocateApLoopFunc (
-      MwaitSupport,
-      CpuMpData->ApTargetCState,
-      CpuMpData->PmCodeSegment,
-      StackStart - ProcessorNumber * AP_SAFE_STACK_SIZE,
-      (UINTN)&mNumberToFinish,
-      CpuMpData->Pm16CodeSegment,
-      CpuMpData->SevEsAPBuffer,
-      CpuMpData->WakeupBuffer
-      );
+    StackStart = mReservedTopOfApStack;
+    mReservedApLoop.GenericEntry (
+                      MwaitSupport,
+                      CpuMpData->ApTargetCState,
+                      CpuMpData->PmCodeSegment,
+                      StackStart - ProcessorNumber * AP_SAFE_STACK_SIZE,
+                      (UINTN)&mNumberToFinish,
+                      CpuMpData->Pm16CodeSegment,
+                      CpuMpData->SevEsAPBuffer,
+                      CpuMpData->WakeupBuffer
+                      );
   }
 
   //
@@ -477,12 +473,15 @@ InitMpGlobalData (
   )
 {
   EFI_STATUS                       Status;
-  EFI_PHYSICAL_ADDRESS             Address;
-  UINTN                            ApSafeBufferSize;
+  MP_ASSEMBLY_ADDRESS_MAP          *AddressMap;
+  UINTN                            AllocSize;
   UINTN                            Index;
   EFI_GCD_MEMORY_SPACE_DESCRIPTOR  MemDesc;
   UINTN                            StackBase;
   CPU_INFO_IN_HOB                  *CpuInfoInHob;
+  EFI_PHYSICAL_ADDRESS             Address;
+  UINT8                            *ApLoopFuncData;
+  UINTN                            ApLoopFuncSize;
 
   SaveCpuMpData (CpuMpData);
 
@@ -537,6 +536,21 @@ InitMpGlobalData (
     }
   }
 
+  AddressMap = &CpuMpData->AddressMap;
+  if (StandardSignatureIsAuthenticAMD () && (sizeof (UINTN) == sizeof (UINT64))) {
+    //
+    // 64-bit AMD Processor
+    //
+    ApLoopFuncData = AddressMap->RelocateApLoopFuncAddressAmd64;
+    ApLoopFuncSize = AddressMap->RelocateApLoopFuncSizeAmd64;
+  } else {
+    //
+    // Intel Processor (32-bit or 64-bit), or 32-bit AMD Processor
+    //
+    ApLoopFuncData = AddressMap->RelocateApLoopFuncAddress;
+    ApLoopFuncSize = AddressMap->RelocateApLoopFuncSize;
+  }
+
   //
   // Avoid APs access invalid buffer data which allocated by BootServices,
   // so we will allocate reserved data for AP loop code. We also need to
@@ -545,26 +559,32 @@ InitMpGlobalData (
   // Allocating it in advance since memory services are not available in
   // Exit Boot Services callback function.
   //
-  ApSafeBufferSize = EFI_PAGES_TO_SIZE (
-                       EFI_SIZE_TO_PAGES (
-                         CpuMpData->AddressMap.RelocateApLoopFuncSize
-                         )
-                       );
+  // +------------+
+  // | Ap Loop    |
+  // +------------+
+  // | Stack * N  |
+  // +------------+ (low address)
+  //
   Address = BASE_4GB - 1;
-  Status  = gBS->AllocatePages (
-                   AllocateMaxAddress,
-                   EfiReservedMemoryType,
-                   EFI_SIZE_TO_PAGES (ApSafeBufferSize),
-                   &Address
-                   );
+  STATIC_ASSERT ((AP_SAFE_STACK_SIZE & (CPU_STACK_ALIGNMENT - 1)) == 0, "AP_SAFE_STACK_SIZE is not aligned with CPU_STACK_ALIGNMENT");
+  AllocSize = EFI_PAGES_TO_SIZE (
+                EFI_SIZE_TO_PAGES (
+                  CpuMpData->CpuCount * AP_SAFE_STACK_SIZE + ApLoopFuncSize
+                  )
+                );
+  Status = gBS->AllocatePages (
+                  AllocateMaxAddress,
+                  EfiReservedMemoryType,
+                  EFI_SIZE_TO_PAGES (AllocSize),
+                  &Address
+                  );
   ASSERT_EFI_ERROR (Status);
-
-  mReservedApLoopFunc = (VOID *)(UINTN)Address;
-  ASSERT (mReservedApLoopFunc != NULL);
-
+  // If a memory range has the EFI_MEMORY_XP attribute, OS loader
+  // may set the IA32_EFER.NXE (No-eXecution Enable) bit in IA32_EFER MSR,
+  // then set the XD (eXecution Disable) bit in the CPU PAE page table.
   //
-  // Make sure that the buffer memory is executable if NX protection is enabled
-  // for EfiReservedMemoryType.
+  // Here is to make sure that the memory is executable if NX protection is
+  // enabled for EfiReservedMemoryType.
   //
   // TODO: Check EFI_MEMORY_XP bit set or not once it's available in DXE GCD
   //       service.
@@ -573,40 +593,15 @@ InitMpGlobalData (
   if (!EFI_ERROR (Status)) {
     gDS->SetMemorySpaceAttributes (
            Address,
-           ApSafeBufferSize,
+           ApLoopFuncSize,
            MemDesc.Attributes & (~EFI_MEMORY_XP)
            );
   }
 
-  ApSafeBufferSize = EFI_PAGES_TO_SIZE (
-                       EFI_SIZE_TO_PAGES (
-                         CpuMpData->CpuCount * AP_SAFE_STACK_SIZE
-                         )
-                       );
-  Address = BASE_4GB - 1;
-  Status  = gBS->AllocatePages (
-                   AllocateMaxAddress,
-                   EfiReservedMemoryType,
-                   EFI_SIZE_TO_PAGES (ApSafeBufferSize),
-                   &Address
-                   );
-  ASSERT_EFI_ERROR (Status);
-
-  mReservedTopOfApStack = (UINTN)Address + ApSafeBufferSize;
+  mReservedTopOfApStack = ((UINTN)Address + CpuMpData->CpuCount * AP_SAFE_STACK_SIZE);
   ASSERT ((mReservedTopOfApStack & (UINTN)(CPU_STACK_ALIGNMENT - 1)) == 0);
-  if (StandardSignatureIsAuthenticAMD () && (sizeof (UINTN) == sizeof (UINT64))) {
-    CopyMem (
-      mReservedApLoopFunc,
-      CpuMpData->AddressMap.RelocateApLoopFuncAddress,
-      CpuMpData->AddressMap.RelocateApLoopFuncSize
-      );
-  } else {
-    CopyMem (
-      mReservedApLoopFunc,
-      CpuMpData->AddressMap.RelocateApLoopFuncAddressAmd64,
-      CpuMpData->AddressMap.RelocateApLoopFuncSizeAmd64
-      );
-  }
+  mReservedApLoop.Data = (VOID *)mReservedTopOfApStack;
+  CopyMem (mReservedApLoop.Data, ApLoopFuncData, ApLoopFuncSize);
 
   Status = gBS->CreateEvent (
                   EVT_TIMER | EVT_NOTIFY_SIGNAL,
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index 4f6146e30d..f1b9063e78 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -401,6 +401,15 @@ typedef
   IN UINTN                   WakeupBuffer
   );
 
+//
+// Union holds the relocate APs loop entries for different cases
+//
+typedef union {
+  VOID                          *Data;
+  ASM_RELOCATE_AP_LOOP_AMD64    Amd64Entry;   // 64-bit AMD Processor
+  ASM_RELOCATE_AP_LOOP          GenericEntry; // Intel Processor (32-bit or 64-bit), or 32-bit AMD Processor
+} RELOCATE_AP_LOOP_ENTRY;
+
 /**
   Assembly code to get starting address and size of the rendezvous entry for APs.
   Information for fixing a jump instruction in the code is also returned.
-- 
2.36.1.windows.1