[Patch V2 2/2] UefiCpuPkg: Keep 4GB limitation of memory allocation

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

Keep 4GB limitation of memory allocation if APs need transferring to
32bit mode before handoff to the OS.

Bugzilla: https://bugzilla.tianocore.org/show_bug.cgi?id=4234

Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Signed-off-by: Yuanhao Xie <yuanhao.xie@intel.com>
---
 UefiCpuPkg/Library/MpInitLib/DxeMpLib.c | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index acbbf155c0..e4c42e34ce 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
@@ -480,6 +480,7 @@ InitMpGlobalData (
   EFI_GCD_MEMORY_SPACE_DESCRIPTOR  MemDesc;
   UINTN                            StackBase;
   CPU_INFO_IN_HOB                  *CpuInfoInHob;
+  EFI_PHYSICAL_ADDRESS             Address;
 
   SaveCpuMpData (CpuMpData);
 
@@ -561,13 +562,25 @@ InitMpGlobalData (
   ASSERT ((AP_SAFE_STACK_SIZE & (CPU_STACK_ALIGNMENT - 1)) == 0);
 
   mReservedApLoopFunc = (VOID *)(mReservedTopOfApStack + CpuMpData->CpuCount * AP_SAFE_STACK_SIZE);
-  if (StandardSignatureIsAuthenticAMD ()) {
+  if (StandardSignatureIsAuthenticAMD () && (sizeof (UINTN) == sizeof (UINT64))) {
+    Address = BASE_4GB - 1;
+    Status  = gBS->AllocatePages (
+                     AllocateMaxAddress,
+                     EfiReservedMemoryType,
+                     EFI_SIZE_TO_PAGES (ApSafeBufferSize),
+                     &Address
+                     );
+    ASSERT_EFI_ERROR (Status);
+    mReservedApLoopFunc = (VOID *)((UINTN)Address + CpuMpData->CpuCount * AP_SAFE_STACK_SIZE);
     CopyMem (
       mReservedApLoopFunc,
       CpuMpData->AddressMap.RelocateApLoopFuncAddressAmd,
       CpuMpData->AddressMap.RelocateApLoopFuncSizeAmd
       );
   } else {
+    Address = (UINTN)AllocateReservedPages (EFI_SIZE_TO_PAGES (ApSafeBufferSize));
+    ASSERT (Address != 0);
+    mReservedApLoopFunc = (VOID *)((UINTN)Address + CpuMpData->CpuCount * AP_SAFE_STACK_SIZE);
     CopyMem (
       mReservedApLoopFunc,
       CpuMpData->AddressMap.RelocateApLoopFuncAddress,
@@ -575,12 +588,14 @@ InitMpGlobalData (
       );
 
     mApPageTable = CreatePageTable (
-                     mReservedTopOfApStack,
+                     (UINTN)Address,
                      ApSafeBufferSize
                      );
   }
 
-  mReservedTopOfApStack += CpuMpData->CpuCount * AP_SAFE_STACK_SIZE;
+  mReservedTopOfApStack = (UINTN)Address + CpuMpData->CpuCount * AP_SAFE_STACK_SIZE;
+  ASSERT ((mReservedTopOfApStack & (UINTN)(CPU_STACK_ALIGNMENT - 1)) == 0);
+  ASSERT ((AP_SAFE_STACK_SIZE & (CPU_STACK_ALIGNMENT - 1)) == 0);
 
   Status = gBS->CreateEvent (
                   EVT_TIMER | EVT_NOTIFY_SIGNAL,
-- 
2.36.1.windows.1