UefiCpuPkg/CpuMpPei: Register callback on End Of Pei PPI

Add CpuMpEndOfPeiCallback () to restore wakeup buffer data on S3 path and flag
flag wakeup buffer to be un-used type on normal boot path. Set one EndOfPei
flag save/restore wakeup buffer when wakeup APs every time.

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jeff Fan <jeff.fan@intel.com>
Reviewed-by: Feng Tian <feng.tian@intel.com>
Reviewed-by: Jiewen Yao <jiewen.yao@intel.com>

git-svn-id: https://svn.code.sf.net/p/edk2/code/trunk/edk2@18014 6f19259b-4bc3-4df7-8a09-765794883524

4 files changed, 177 insertions, 0 deletions

diff --git a/UefiCpuPkg/CpuMpPei/CpuMpPei.c b/UefiCpuPkg/CpuMpPei/CpuMpPei.c
index 40e62e0405..97b7e99c9a 100644
--- a/UefiCpuPkg/CpuMpPei/CpuMpPei.c
+++ b/UefiCpuPkg/CpuMpPei/CpuMpPei.c
@@ -38,6 +38,12 @@ GLOBAL_REMOVE_IF_UNREFERENCED IA32_DESCRIPTOR mGdt = {
   (UINTN) mGdtEntries
   };
 
+GLOBAL_REMOVE_IF_UNREFERENCED EFI_PEI_NOTIFY_DESCRIPTOR mNotifyList = {
+  (EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
+  &gEfiEndOfPeiSignalPpiGuid,
+  CpuMpEndOfPeiCallback
+};
+
 /**
   Sort the APIC ID of all processors.
 
@@ -317,6 +323,20 @@ BackupAndPrepareWakeupBuffer(
     PeiCpuMpData->AddressMap.RendezvousFunnelSize
     );
 }
+
+/**
+  Restore wakeup buffer data.
+
+  @param PeiCpuMpData        Pointer to PEI CPU MP Data
+**/
+VOID
+RestoreWakeupBuffer(
+  IN PEI_CPU_MP_DATA         *PeiCpuMpData
+  )
+{
+  CopyMem ((VOID *) PeiCpuMpData->WakeupBuffer, (VOID *) PeiCpuMpData->BackupBuffer, PeiCpuMpData->BackupBufferSize);
+}
+
 /**
   This function will get CPU count in the system.
 
@@ -381,6 +401,7 @@ PrepareAPStartupVector (
   AsmGetAddressMap (&AddressMap);
   WakeupBufferSize = AddressMap.RendezvousFunnelSize + sizeof (MP_CPU_EXCHANGE_INFO);
   WakeupBuffer     = GetWakeupBuffer ((WakeupBufferSize + SIZE_4KB - 1) & ~(SIZE_4KB - 1));
+  ASSERT (WakeupBuffer != (UINTN) -1);
   DEBUG ((EFI_D_INFO, "CpuMpPei: WakeupBuffer = 0x%x\n", WakeupBuffer));
 
   //
@@ -409,6 +430,7 @@ PrepareAPStartupVector (
   PeiCpuMpData->CpuData                  = (PEI_CPU_DATA *) (PeiCpuMpData->MpCpuExchangeInfo + 1);
   PeiCpuMpData->CpuData[0].ApicId        = GetInitialApicId ();
   PeiCpuMpData->CpuData[0].Health.Uint32 = 0;
+  PeiCpuMpData->EndOfPeiFlag             = FALSE;
   CopyMem (&PeiCpuMpData->AddressMap, &AddressMap, sizeof (MP_ASSEMBLY_ADDRESS_MAP));
 
   //
@@ -418,6 +440,70 @@ PrepareAPStartupVector (
 
   return PeiCpuMpData;
 }
+
+/**
+  Notify function on End Of Pei PPI.
+
+  On S3 boot, this function will restore wakeup buffer data.
+  On normal boot, this function will flag wakeup buffer to be un-used type.
+
+  @param  PeiServices        The pointer to the PEI Services Table.
+  @param  NotifyDescriptor   Address of the notification descriptor data structure.
+  @param  Ppi                Address of the PPI that was installed.
+
+  @retval EFI_SUCCESS        When everything is OK.
+
+**/
+EFI_STATUS
+EFIAPI
+CpuMpEndOfPeiCallback (
+  IN      EFI_PEI_SERVICES        **PeiServices,
+  IN EFI_PEI_NOTIFY_DESCRIPTOR    *NotifyDescriptor,
+  IN VOID                         *Ppi
+  )
+{
+  EFI_STATUS                Status;
+  EFI_BOOT_MODE             BootMode;
+  PEI_CPU_MP_DATA           *PeiCpuMpData;
+  EFI_PEI_HOB_POINTERS      Hob;
+  EFI_HOB_MEMORY_ALLOCATION *MemoryHob;
+
+  DEBUG ((EFI_D_INFO, "CpuMpPei: CpuMpEndOfPeiCallback () invokded\n"));
+
+  Status = PeiServicesGetBootMode (&BootMode);
+  ASSERT_EFI_ERROR (Status);
+
+  PeiCpuMpData = GetMpHobData ();
+  ASSERT (PeiCpuMpData != NULL);
+
+  if (BootMode != BOOT_ON_S3_RESUME) {
+    //
+    // Get the HOB list for processing
+    //
+    Hob.Raw = GetHobList ();
+    //
+    // Collect memory ranges
+    //
+    while (!END_OF_HOB_LIST (Hob)) {
+      if (Hob.Header->HobType == EFI_HOB_TYPE_MEMORY_ALLOCATION) {
+        MemoryHob = Hob.MemoryAllocation;
+        if(MemoryHob->AllocDescriptor.MemoryBaseAddress == PeiCpuMpData->WakeupBuffer) {
+          //
+          // Flag this HOB type to un-used
+          //
+          GET_HOB_TYPE (Hob) = EFI_HOB_TYPE_UNUSED;
+          break;
+        }
+      }
+      Hob.Raw = GET_NEXT_HOB (Hob);
+    }
+  } else {
+    RestoreWakeupBuffer (PeiCpuMpData);
+    PeiCpuMpData->EndOfPeiFlag = TRUE;
+  }
+  return EFI_SUCCESS;
+}
+
 /**
   The Entry point of the MP CPU PEIM.
 
@@ -466,6 +552,11 @@ CpuMpPeimInit (
   //
   CollectBistDataFromPpi (PeiServices, PeiCpuMpData);
   //
+  // register an event for EndOfPei
+  //
+  Status  = PeiServicesNotifyPpi (&mNotifyList);
+  ASSERT_EFI_ERROR (Status);
+  //
   // Install CPU MP PPI
   //
   Status = PeiServicesInstallPpi(&mPeiCpuMpPpiDesc);
diff --git a/UefiCpuPkg/CpuMpPei/CpuMpPei.h b/UefiCpuPkg/CpuMpPei/CpuMpPei.h
index c241349a97..8d01ac6433 100644
--- a/UefiCpuPkg/CpuMpPei/CpuMpPei.h
+++ b/UefiCpuPkg/CpuMpPei/CpuMpPei.h
@@ -20,6 +20,7 @@
 #include <Ppi/MpServices.h>
 #include <Ppi/SecPlatformInformation.h>
 #include <Ppi/SecPlatformInformation2.h>
+#include <Ppi/EndOfPeiPhase.h>
 
 #include <Register/LocalApic.h>
 
@@ -133,6 +134,7 @@ struct _PEI_CPU_MP_DATA {
   UINTN                          ApFunction;
   UINTN                          ApFunctionArgument;
   volatile UINT32                FinishedCount;
+  BOOLEAN                        EndOfPeiFlag;
   BOOLEAN                        InitFlag;
   CPU_EXCHANGE_ROLE_INFO         BSPInfo;
   CPU_EXCHANGE_ROLE_INFO         APInfo;
@@ -177,6 +179,47 @@ AsmCliHltLoop (
   );
 
 /**
+  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
+  );
+
+/**
+  Restore wakeup buffer data.
+
+  @param PeiCpuMpData        Pointer to PEI CPU MP Data
+**/
+VOID
+RestoreWakeupBuffer(
+  IN PEI_CPU_MP_DATA         *PeiCpuMpData
+  );
+
+/**
+  Notify function on End Of Pei PPI.
+
+  On S3 boot, this function will restore wakeup buffer data.
+  On normal boot, this function will flag wakeup buffer to be un-used type.
+
+  @param  PeiServices        The pointer to the PEI Services Table.
+  @param  NotifyDescriptor   Address of the notification descriptor data structure.
+  @param  Ppi                Address of the PPI that was installed.
+
+  @retval EFI_SUCCESS        When everything is OK.
+
+**/
+EFI_STATUS
+EFIAPI
+CpuMpEndOfPeiCallback (
+  IN      EFI_PEI_SERVICES        **PeiServices,
+  IN EFI_PEI_NOTIFY_DESCRIPTOR    *NotifyDescriptor,
+  IN VOID                         *Ppi
+  );
+
+/**
   This function will be called by BSP to wakeup AP.
 
   @param PeiCpuMpData       Pointer to PEI CPU MP Data
diff --git a/UefiCpuPkg/CpuMpPei/CpuMpPei.inf b/UefiCpuPkg/CpuMpPei/CpuMpPei.inf
index 7160c0e05c..92c64f9e4e 100644
--- a/UefiCpuPkg/CpuMpPei/CpuMpPei.inf
+++ b/UefiCpuPkg/CpuMpPei/CpuMpPei.inf
@@ -69,6 +69,7 @@
 
 [Ppis]
   gEfiPeiMpServicesPpiGuid                      ## PRODUCES
+  gEfiEndOfPeiSignalPpiGuid                     ## NOTIFY
   gEfiSecPlatformInformationPpiGuid             ## SOMETIMES_CONSUMES
   ## SOMETIMES_CONSUMES
   ## SOMETIMES_PRODUCES
diff --git a/UefiCpuPkg/CpuMpPei/PeiMpServices.c b/UefiCpuPkg/CpuMpPei/PeiMpServices.c
index 5c20c1874e..4689d2fcc0 100644
--- a/UefiCpuPkg/CpuMpPei/PeiMpServices.c
+++ b/UefiCpuPkg/CpuMpPei/PeiMpServices.c
@@ -480,6 +480,13 @@ PeiStartupAllAPs (
     return EFI_NOT_READY;
   }
 
+  if (PeiCpuMpData->EndOfPeiFlag) {
+    //
+    // Backup original data and copy AP reset vector in it
+    //
+    BackupAndPrepareWakeupBuffer(PeiCpuMpData);
+  }
+
   WaitCountNumber = TimeoutInMicroSeconds / CPU_CHECK_AP_INTERVAL + 1;
   WaitCountIndex = 0;
   FinishedCount = &PeiCpuMpData->FinishedCount;
@@ -531,6 +538,13 @@ PeiStartupAllAPs (
     }
   }
 
+  if (PeiCpuMpData->EndOfPeiFlag) {
+    //
+    // Restore original data
+    //
+    RestoreWakeupBuffer(PeiCpuMpData);
+  }
+
   return Status;
 }
 
@@ -626,6 +640,13 @@ PeiStartupThisAP (
     return EFI_INVALID_PARAMETER;
   }
 
+  if (PeiCpuMpData->EndOfPeiFlag) {
+    //
+    // Backup original data and copy AP reset vector in it
+    //
+    BackupAndPrepareWakeupBuffer(PeiCpuMpData);
+  }
+
   WaitCountNumber = TimeoutInMicroseconds / CPU_CHECK_AP_INTERVAL + 1;
   WaitCountIndex = 0;
   FinishedCount = &PeiCpuMpData->FinishedCount;
@@ -651,6 +672,13 @@ PeiStartupThisAP (
     }
   }
 
+  if (PeiCpuMpData->EndOfPeiFlag) {
+    //
+    // Backup original data and copy AP reset vector in it
+    //
+    RestoreWakeupBuffer(PeiCpuMpData);
+  }
+
   return Status;
 }
 
@@ -749,6 +777,13 @@ PeiSwitchBSP (
   PeiCpuMpData->BSPInfo.State = CPU_SWITCH_STATE_IDLE;
   PeiCpuMpData->APInfo.State  = CPU_SWITCH_STATE_IDLE;
 
+  if (PeiCpuMpData->EndOfPeiFlag) {
+    //
+    // Backup original data and copy AP reset vector in it
+    //
+    BackupAndPrepareWakeupBuffer(PeiCpuMpData);
+  }
+
   //
   // Need to wakeUp AP (future BSP).
   //
@@ -756,6 +791,13 @@ PeiSwitchBSP (
 
   AsmExchangeRole (&PeiCpuMpData->BSPInfo, &PeiCpuMpData->APInfo);
 
+  if (PeiCpuMpData->EndOfPeiFlag) {
+    //
+    // Backup original data and copy AP reset vector in it
+    //
+    RestoreWakeupBuffer(PeiCpuMpData);
+  }
+
   //
   // Set the BSP bit of MSR_IA32_APIC_BASE on new BSP
   //