[PATCH 3/3] UefiCpuPkg/DxeMpLib: Make sure APs in safe loop code

[Jeff Fan <jeff.fan@intel.com>]

Add one semaphore to make sure BSP to wait till all APs run in AP safe loop
code.

Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Feng Tian <feng.tian@intel.com>
Cc: Michael D Kinney <michael.d.kinney@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jeff Fan <jeff.fan@intel.com>
---
 UefiCpuPkg/Library/MpInitLib/DxeMpLib.c        | 8 +++++++-
 UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm | 2 ++
 UefiCpuPkg/Library/MpInitLib/MpLib.h           | 3 ++-
 UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm  | 2 ++
 4 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
index d0f9f7e..64c0c52 100644
--- a/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/DxeMpLib.c
@@ -27,6 +27,7 @@ EFI_EVENT        mLegacyBootEvent = NULL;
 volatile BOOLEAN mStopCheckAllApsStatus = TRUE;
 VOID             *mReservedApLoopFunc = NULL;
 UINTN            mReservedTopOfApStack;
+volatile UINT32  mNumberToFinish = 0;
 
 /**
   Get the pointer to CPU MP Data structure.
@@ -253,7 +254,8 @@ RelocateApLoop (
     MwaitSupport,
     CpuMpData->ApTargetCState,
     CpuMpData->PmCodeSegment,
-    mReservedTopOfApStack - ProcessorNumber * AP_SAFT_STACK_SIZE
+    mReservedTopOfApStack - ProcessorNumber * AP_SAFT_STACK_SIZE,
+    (UINTN) &mNumberToFinish
     );
   //
   // It should never reach here
@@ -282,8 +284,12 @@ MpInitChangeApLoopCallback (
   CpuMpData->SaveRestoreFlag = TRUE;
   CpuMpData->PmCodeSegment = GetProtectedModeCS ();
   CpuMpData->ApLoopMode = PcdGet8 (PcdCpuApLoopMode);
+  mNumberToFinish = CpuMpData->CpuCount - 1;
   WakeUpAP (CpuMpData, TRUE, 0, RelocateApLoop, NULL);
   DEBUG ((DEBUG_INFO, "%a() done!\n", __FUNCTION__));
+  while (mNumberToFinish > 0) {
+    CpuPause ();
+  }
 }
 
 /**
diff --git a/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm b/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
index 4e55760..6235748 100644
--- a/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
+++ b/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
@@ -222,6 +222,8 @@ AsmRelocateApLoopStart:
     mov        ecx, [ebp + 8]      ; MwaitSupport
     mov        ebx, [ebp + 12]     ; ApTargetCState
     mov        esp, [ebp + 20]     ; TopOfApStack
+    mov        eax, [ebp + 24]     ; CountTofinish
+    lock dec   dword [eax]         ; CountTofinish--
     cmp        cl,  1              ; Check mwait-monitor support
     jnz        HltLoop
 MwaitLoop:
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index e6dea18..49305ad 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -251,7 +251,8 @@ VOID
   IN BOOLEAN                 MwaitSupport,
   IN UINTN                   ApTargetCState,
   IN UINTN                   PmCodeSegment,
-  IN UINTN                   TopOfApStack
+  IN UINTN                   TopOfApStack,
+  IN UINTN                   NumberToFinish
   );
 
 /**
diff --git a/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm b/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
index 7c8fa45..deaee9e 100644
--- a/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
+++ b/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
@@ -227,6 +227,8 @@ AsmRelocateApLoopStart:
     push       rbp
     mov        rbp, rsp
     mov        rsp, r9
+    mov        rax, [rbp + 48]     ; CountTofinish
+    lock dec   dword [rax]         ; CountTofinish--
     push       rcx
     push       rdx
 
-- 
2.9.3.windows.2