UefiCpuPkg/DxeMpLib: Make sure APs in safe loop code

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>
Reviewed-by: Laszlo Ersek <lersek@redhat.com>
Tested-by: Laszlo Ersek <lersek@redhat.com>

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_SAFE_STACK_SIZE
+    mReservedTopOfApStack - ProcessorNumber * AP_SAFE_STACK_SIZE,
+    (UINTN) &mNumberToFinish
     );
   //
   // It should never reach here
@@ -282,7 +284,11 @@ MpInitChangeApLoopCallback (
   CpuMpData->SaveRestoreFlag = TRUE;
   CpuMpData->PmCodeSegment = GetProtectedModeCS ();
   CpuMpData->ApLoopMode = PcdGet8 (PcdCpuApLoopMode);
+  mNumberToFinish = CpuMpData->CpuCount - 1;
   WakeUpAP (CpuMpData, TRUE, 0, RelocateApLoop, NULL);
+  while (mNumberToFinish > 0) {
+    CpuPause ();
+  }
   DEBUG ((DEBUG_INFO, "%a() done!\n", __FUNCTION__));
 }
 

UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm

@@ -215,7 +215,7 @@ CProcedureInvoke:
 RendezvousFunnelProcEnd:
 
 ;-------------------------------------------------------------------------------------
-;  AsmRelocateApLoop (MwaitSupport, ApTargetCState, PmCodeSegment, TopOfApStack);
+;  AsmRelocateApLoop (MwaitSupport, ApTargetCState, PmCodeSegment, TopOfApStack, CountTofinish);
 ;-------------------------------------------------------------------------------------
 global ASM_PFX(AsmRelocateApLoop)
 ASM_PFX(AsmRelocateApLoop):
@@ -227,6 +227,8 @@ AsmRelocateApLoopStart:
     mov        ebp, esp
     mov        ebx, [eax + 8]      ; ApTargetCState
     mov        ecx, [eax + 4]      ; MwaitSupport
+    mov        eax, [eax + 20]     ; CountTofinish
+    lock dec   dword [eax]         ; (*CountTofinish)--
     cmp        cl,  1              ; Check mwait-monitor support
     jnz        HltLoop
 MwaitLoop:

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
   );
 
 /**

UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm

@@ -222,11 +222,13 @@ CProcedureInvoke:
 RendezvousFunnelProcEnd:
 
 ;-------------------------------------------------------------------------------------
-;  AsmRelocateApLoop (MwaitSupport, ApTargetCState, PmCodeSegment, TopOfApStack);
+;  AsmRelocateApLoop (MwaitSupport, ApTargetCState, PmCodeSegment, TopOfApStack, CountTofinish);
 ;-------------------------------------------------------------------------------------
 global ASM_PFX(AsmRelocateApLoop)
 ASM_PFX(AsmRelocateApLoop):
 AsmRelocateApLoopStart:
+    mov        rax, [rsp + 40]   ; CountTofinish
+    lock dec   dword [rax]       ; (*CountTofinish)--
     mov        rsp, r9
     push       rcx
     push       rdx