1 files changed, 11 insertions, 2 deletions
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index ada3969b68..cbc5594c78 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -1820,7 +1820,14 @@ MpInitLibInitialize (
   // Wakeup APs to do some AP initialize sync (Microcode & MTRR)
   //
   if (CpuMpData->CpuCount > 1) {
-    CpuMpData->InitFlag = ApInitReconfig;
+    if (OldCpuMpData != NULL) {
+      //
+      // Only needs to use this flag for DXE phase to update the wake up
+      // buffer. Wakeup buffer allocated in PEI phase is no longer valid
+      // in DXE.
+      //
+      CpuMpData->InitFlag = ApInitReconfig;
+    }
     WakeUpAP (CpuMpData, TRUE, 0, ApInitializeSync, CpuMpData, TRUE);
     //
     // Wait for all APs finished initialization
@@ -1828,7 +1835,9 @@ MpInitLibInitialize (
     while (CpuMpData->FinishedCount < (CpuMpData->CpuCount - 1)) {
       CpuPause ();
     }
-    CpuMpData->InitFlag = ApInitDone;
+    if (OldCpuMpData != NULL) {
+      CpuMpData->InitFlag = ApInitDone;
+    }
     for (Index = 0; Index < CpuMpData->CpuCount; Index++) {
       SetApState (&CpuMpData->CpuData[Index], CpuStateIdle);
     }