This patch aims on mProcessorInfo global variable, which can be
defined as local variable in SmmRelocateBases(). With this patch,
no need to allocate the memory for all CPUs to store the
Processor Info.

Cc: Ray Ni <ray.ni@intel.com>
Cc: Zeng Star <star.zeng@intel.com>
Cc: Gerd Hoffmann <kraxel@redhat.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Signed-off-by: Jiaxin Wu <jiaxin.wu@intel.com>
---
 .../Library/SmmRelocationLib/SmmRelocationLib.c    | 88 ++++++++--------------
 1 file changed, 32 insertions(+), 56 deletions(-)

diff --git a/UefiCpuPkg/Library/SmmRelocationLib/SmmRelocationLib.c b/UefiCpuPkg/Library/SmmRelocationLib/SmmRelocationLib.c
index 3694a07cbb..86df66a280 100644
--- a/UefiCpuPkg/Library/SmmRelocationLib/SmmRelocationLib.c
+++ b/UefiCpuPkg/Library/SmmRelocationLib/SmmRelocationLib.c
@@ -14,15 +14,10 @@
 #include "InternalSmmRelocationLib.h"
 
 UINTN  mMaxNumberOfCpus = 1;
 UINTN  mNumberOfCpus    = 1;
 
-//
-// Record all Processors Info
-//
-EFI_PROCESSOR_INFORMATION  *mProcessorInfo = NULL;
-
 //
 // IDT used during SMM Init
 //
 IA32_DESCRIPTOR  gcSmmInitIdtr;
 
@@ -34,10 +29,15 @@ UINT64  mSmBase;
 //
 // SmBase Rebased flag for current CPU
 //
 volatile BOOLEAN  mRebased;
 
+//
+// CpuIndex for current CPU
+//
+UINTN  mCpuIndex;
+
 /**
   This function will get the SmBase for CpuIndex.
 
   @param[in]   CpuIndex            The processor index.
   @param[in]   SmmRelocationStart  The start address of Smm relocated memory in SMRAM.
@@ -138,37 +138,26 @@ VOID
 EFIAPI
 SmmInitHandler (
   VOID
   )
 {
-  UINT32  ApicId;
-  UINTN   Index;
-
   //
   // Update SMM IDT entries' code segment and load IDT
   //
   AsmWriteIdtr (&gcSmmInitIdtr);
-  ApicId = GetApicId ();
 
-  for (Index = 0; Index < mNumberOfCpus; Index++) {
-    if (ApicId == (UINT32)mProcessorInfo[Index].ProcessorId) {
-      //
-      // Configure SmBase.
-      //
-      ConfigureSmBase (mSmBase);
-
-      //
-      // Hook return after RSM to set SMM re-based flag
-      // SMM re-based flag can't be set before RSM, because SMM save state context might be override
-      // by next AP flow before it take effect.
-      //
-      SemaphoreHook (Index, &mRebased);
-      return;
-    }
-  }
+  //
+  // Configure SmBase.
+  //
+  ConfigureSmBase (mSmBase);
 
-  ASSERT (FALSE);
+  //
+  // Hook return after RSM to set SMM re-based flag
+  // SMM re-based flag can't be set before RSM, because SMM save state context might be override
+  // by next AP flow before it take effect.
+  //
+  SemaphoreHook (mCpuIndex, &mRebased);
 }
 
 /**
   Relocate SmmBases for each processor.
   Execute on first boot and all S3 resumes
@@ -185,17 +174,19 @@ SmmRelocateBases (
   IN EDKII_PEI_MP_SERVICES2_PPI  *MpServices2,
   IN EFI_PHYSICAL_ADDRESS        SmmRelocationStart,
   IN UINTN                       TileSize
   )
 {
-  UINT8                 BakBuf[BACK_BUF_SIZE];
-  SMRAM_SAVE_STATE_MAP  BakBuf2;
-  SMRAM_SAVE_STATE_MAP  *CpuStatePtr;
-  UINT8                 *U8Ptr;
-  UINTN                 Index;
-  UINTN                 BspIndex;
-  UINT32                BspApicId;
+  EFI_STATUS                 Status;
+  UINT8                      BakBuf[BACK_BUF_SIZE];
+  SMRAM_SAVE_STATE_MAP       BakBuf2;
+  SMRAM_SAVE_STATE_MAP       *CpuStatePtr;
+  UINT8                      *U8Ptr;
+  UINTN                      Index;
+  UINTN                      BspIndex;
+  UINT32                     BspApicId;
+  EFI_PROCESSOR_INFORMATION  ProcessorInfo;
 
   //
   // Make sure the reserved size is large enough for procedure SmmInitTemplate.
   //
   ASSERT (sizeof (BakBuf) >= gcSmmInitSize);
@@ -230,14 +221,18 @@ SmmRelocateBases (
   // Relocate SM bases for all APs
   // This is APs' 1st SMI - rebase will be done here, and APs' default SMI handler will be overridden by gcSmmInitTemplate
   //
   BspIndex = (UINTN)-1;
   for (Index = 0; Index < mNumberOfCpus; Index++) {
-    if (BspApicId != (UINT32)mProcessorInfo[Index].ProcessorId) {
-      mRebased = FALSE;
-      mSmBase  = GetSmBase (Index, SmmRelocationStart, TileSize);
-      SendSmiIpi ((UINT32)mProcessorInfo[Index].ProcessorId);
+    Status = MpServices2->GetProcessorInfo (MpServices2, Index | CPU_V2_EXTENDED_TOPOLOGY, &ProcessorInfo);
+    ASSERT_EFI_ERROR (Status);
+
+    if (BspApicId != (UINT32)ProcessorInfo.ProcessorId) {
+      mRebased  = FALSE;
+      mSmBase   = GetSmBase (Index, SmmRelocationStart, TileSize);
+      mCpuIndex = Index;
+      SendSmiIpi ((UINT32)ProcessorInfo.ProcessorId);
       //
       // Wait for this AP to finish its 1st SMI
       //
       while (!mRebased) {
       }
@@ -443,11 +438,10 @@ SmmRelocationInit (
   UINTN                 TileSize;
   UINT64                SmmRelocationSize;
   EFI_PHYSICAL_ADDRESS  SmmRelocationStart;
   UINTN                 SmmStackSize;
   UINT8                 *SmmStacks;
-  UINTN                 Index;
 
   SmmRelocationStart = 0;
   SmmStacks          = NULL;
 
   DEBUG ((DEBUG_INFO, "SmmRelocationInit Start \n"));
@@ -473,28 +467,10 @@ SmmRelocationInit (
     mMaxNumberOfCpus = mNumberOfCpus;
   }
 
   ASSERT (mNumberOfCpus <= mMaxNumberOfCpus);
 
-  //
-  // Retrieve the Processor Info for all CPUs
-  //
-  mProcessorInfo = (EFI_PROCESSOR_INFORMATION *)AllocatePool (sizeof (EFI_PROCESSOR_INFORMATION) * mMaxNumberOfCpus);
-  if (mProcessorInfo == NULL) {
-    Status = EFI_OUT_OF_RESOURCES;
-    goto ON_EXIT;
-  }
-
-  for (Index = 0; Index < mMaxNumberOfCpus; Index++) {
-    if (Index < mNumberOfCpus) {
-      Status = MpServices2->GetProcessorInfo (MpServices2, Index | CPU_V2_EXTENDED_TOPOLOGY, &mProcessorInfo[Index]);
-      if (EFI_ERROR (Status)) {
-        goto ON_EXIT;
-      }
-    }
-  }
-
   //
   // Calculate SmmRelocationSize for all of the tiles.
   //
   // The CPU save state and code for the SMI entry point are tiled within an SMRAM
   // allocated buffer. The minimum size of this buffer for a uniprocessor system
--
2.16.2.windows.1