3 files changed, 135 insertions, 2 deletions

diff --git a/UefiCpuPkg/CpuMpPei/CpuMpPei.c b/UefiCpuPkg/CpuMpPei/CpuMpPei.c
index 0e34f26afc..45243d804e 100644
--- a/UefiCpuPkg/CpuMpPei/CpuMpPei.c
+++ b/UefiCpuPkg/CpuMpPei/CpuMpPei.c
@@ -111,6 +111,88 @@ GetWakeupBuffer (
 }
 
 /**
+  Get available system memory below 1MB by specified size.
+
+  @param PeiCpuMpData        Pointer to PEI CPU MP Data
+**/
+VOID
+BackupAndPrepareWakeupBuffer(
+  IN PEI_CPU_MP_DATA         *PeiCpuMpData
+  )
+{
+  CopyMem (
+    (VOID *) PeiCpuMpData->BackupBuffer,
+    (VOID *) PeiCpuMpData->WakeupBuffer,
+    PeiCpuMpData->BackupBufferSize
+    );
+  CopyMem (
+    (VOID *) PeiCpuMpData->WakeupBuffer,
+    (VOID *) PeiCpuMpData->AddressMap.RendezvousFunnelAddress,
+    PeiCpuMpData->AddressMap.RendezvousFunnelSize
+    );
+}
+/**
+  Prepare for AP wakeup buffer and copy AP reset code into it.
+
+  Get wakeup buffer below 1MB. Allocate memory for CPU MP Data and APs Stack.
+
+  @return   Pointer to PEI CPU MP Data
+**/
+PEI_CPU_MP_DATA *
+PrepareAPStartupVector (
+  VOID
+  )
+{
+  EFI_STATUS                    Status;
+  UINT32                        MaxCpuCount;
+  PEI_CPU_MP_DATA               *PeiCpuMpData;
+  EFI_PHYSICAL_ADDRESS          Buffer;
+  UINTN                         BufferSize;
+  UINTN                         WakeupBuffer;
+  UINTN                         WakeupBufferSize;
+  MP_ASSEMBLY_ADDRESS_MAP       AddressMap;
+
+  AsmGetAddressMap (&AddressMap);
+  WakeupBufferSize = AddressMap.RendezvousFunnelSize + sizeof (MP_CPU_EXCHANGE_INFO);
+  WakeupBuffer     = GetWakeupBuffer ((WakeupBufferSize + SIZE_4KB - 1) & ~(SIZE_4KB - 1));
+  DEBUG ((EFI_D_INFO, "CpuMpPei: WakeupBuffer = 0x%x\n", WakeupBuffer));
+
+  //
+  // Allocate Pages for APs stack, CPU MP Data and backup buffer for wakeup buffer
+  //
+  MaxCpuCount = PcdGet32(PcdCpuMaxLogicalProcessorNumber);
+  BufferSize  = PcdGet32 (PcdCpuApStackSize) * MaxCpuCount + sizeof (PEI_CPU_MP_DATA)
+                  + WakeupBufferSize + sizeof (PEI_CPU_DATA) * MaxCpuCount;
+  Status = PeiServicesAllocatePages (
+             EfiBootServicesData,
+             EFI_SIZE_TO_PAGES (BufferSize),
+             &Buffer
+             );
+  ASSERT_EFI_ERROR (Status);
+
+  PeiCpuMpData = (PEI_CPU_MP_DATA *) (UINTN) (Buffer + PcdGet32 (PcdCpuApStackSize) * MaxCpuCount);
+  PeiCpuMpData->Buffer            = (UINTN) Buffer;
+  PeiCpuMpData->CpuApStackSize    = PcdGet32 (PcdCpuApStackSize);
+  PeiCpuMpData->WakeupBuffer      = WakeupBuffer;
+  PeiCpuMpData->BackupBuffer      = (UINTN)PeiCpuMpData + sizeof (PEI_CPU_MP_DATA);
+  PeiCpuMpData->BackupBufferSize  = WakeupBufferSize;
+  PeiCpuMpData->MpCpuExchangeInfo = (MP_CPU_EXCHANGE_INFO *) (UINTN) (WakeupBuffer + AddressMap.RendezvousFunnelSize);
+
+  PeiCpuMpData->CpuCount                 = 1;
+  PeiCpuMpData->BspNumber                = 0;
+  PeiCpuMpData->CpuData                  = (PEI_CPU_DATA *) (PeiCpuMpData->MpCpuExchangeInfo + 1);
+  PeiCpuMpData->CpuData[0].ApicId        = GetInitialApicId ();
+  PeiCpuMpData->CpuData[0].Health.Uint32 = 0;
+  CopyMem (&PeiCpuMpData->AddressMap, &AddressMap, sizeof (MP_ASSEMBLY_ADDRESS_MAP));
+
+  //
+  // Backup original data and copy AP reset code in it
+  //
+  BackupAndPrepareWakeupBuffer(PeiCpuMpData);
+
+  return PeiCpuMpData;
+}
+/**
   The Entry point of the MP CPU PEIM.
 
   This function will wakeup APs and collect CPU AP count and install the
@@ -130,11 +212,16 @@ CpuMpPeimInit (
   )
 {
 
+  PEI_CPU_MP_DATA      *PeiCpuMpData;
 
   //
   // Load new GDT table on BSP
   //
   AsmInitializeGdt (&mGdt);
+  //
+  // Get wakeup buffer and copy AP reset code in it
+  //
+  PeiCpuMpData = PrepareAPStartupVector ();
 
   return EFI_SUCCESS;
 }
diff --git a/UefiCpuPkg/CpuMpPei/CpuMpPei.h b/UefiCpuPkg/CpuMpPei/CpuMpPei.h
index 7c96084ca2..1219e4ea5a 100644
--- a/UefiCpuPkg/CpuMpPei/CpuMpPei.h
+++ b/UefiCpuPkg/CpuMpPei/CpuMpPei.h
@@ -17,11 +17,25 @@
 
 #include <PiPei.h>
 
+#include <Ppi/SecPlatformInformation.h>
 
 #include <Library/BaseLib.h>
+#include <Library/BaseMemoryLib.h>
+#include <Library/DebugLib.h>
 #include <Library/HobLib.h>
+#include <Library/LocalApicLib.h>
+#include <Library/PcdLib.h>
 #include <Library/PeimEntryPoint.h>
+#include <Library/PeiServicesLib.h>
 #include <Library/UefiCpuLib.h>
+//
+// AP state
+//
+typedef enum {
+  CpuStateIdle,
+  CpuStateBusy,
+  CpuStateDisabled
+} CPU_STATE;
 
 //
 // AP reset code information
@@ -33,7 +47,9 @@ typedef struct {
   UINTN             RendezvousFunnelSize;
 } MP_ASSEMBLY_ADDRESS_MAP;
 
-#pragma pack(1)
+typedef struct _PEI_CPU_MP_DATA  PEI_CPU_MP_DATA;
+
+#pragma pack()
 
 typedef union {
   struct {
@@ -73,6 +89,29 @@ typedef struct {
 
 #pragma pack()
 
+typedef struct {
+  UINT32                         ApicId;
+  EFI_HEALTH_FLAGS               Health;
+  CPU_STATE                      State;
+  BOOLEAN                        CpuHealthy;
+} PEI_CPU_DATA;
+
+//
+// PEI CPU MP Data save in memory
+//
+struct _PEI_CPU_MP_DATA {
+  UINT32                         CpuCount;
+  UINT32                         BspNumber;
+  UINTN                          Buffer;
+  UINTN                          CpuApStackSize;
+  MP_ASSEMBLY_ADDRESS_MAP        AddressMap;
+  UINTN                          WakeupBuffer;
+  UINTN                          BackupBuffer;
+  UINTN                          BackupBufferSize;
+  PEI_CPU_DATA                   *CpuData;
+  volatile MP_CPU_EXCHANGE_INFO  *MpCpuExchangeInfo;
+};
+
 /**
   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.
diff --git a/UefiCpuPkg/CpuMpPei/CpuMpPei.inf b/UefiCpuPkg/CpuMpPei/CpuMpPei.inf
index 9fb9e945a8..a04556fb02 100644
--- a/UefiCpuPkg/CpuMpPei/CpuMpPei.inf
+++ b/UefiCpuPkg/CpuMpPei/CpuMpPei.inf
@@ -49,12 +49,19 @@
 
 [LibraryClasses]
   BaseLib
+  BaseMemoryLib
+  DebugLib
   HobLib
+  LocalApicLib
+  PcdLib
   PeimEntryPoint
+  PeiServicesLib
   UefiCpuLib
 
 
-
+[Pcd]
+  gUefiCpuPkgTokenSpaceGuid.PcdCpuMaxLogicalProcessorNumber
+  gUefiCpuPkgTokenSpaceGuid.PcdCpuApStackSize
 
 [Depex]
   gEfiPeiMemoryDiscoveredPpiGuid