UefiCpuPkg/MpInitLib: Enhance waiting for AP initialization logic.

Current logic always waiting for a specific value to collect all APs
count. This logic may caused some platforms cost too much time to
wait for time out.
This patch add new logic to collect APs count. It adds new variable
NumApsExecuting to detect whether all APs have finished initialization.
Each AP let NumApsExecuting++ when begin to initialize itself and let
NumApsExecuting-- when it finish the initialization. BSP base on whether
NumApsExecuting == 0  to finished the collect AP process.

Cc: Ruiyu Ni <ruiyu.ni@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Eric Dong <eric.dong@intel.com>
Reviewed-by: Ruiyu Ni <ruiyu.ni@intel.com>
Reviewed-by: Jeff Fan <vanjeff_919@hotmail.com>

UefiCpuPkg/Library/MpInitLib/Ia32/MpEqu.inc

@@ -40,4 +40,5 @@ EnableExecuteDisableLocation  equ        LockLocation + 30h
 Cr3Location                   equ        LockLocation + 34h
 InitFlagLocation              equ        LockLocation + 38h
 CpuInfoLocation               equ        LockLocation + 3Ch
+NumApsExecutingLocation       equ        LockLocation + 40h
 

UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm

@@ -86,6 +86,12 @@ Flat32Start:                                   ; protected mode entry point
 
     mov        esi, ebx
 
+    ; Increment the number of APs executing here as early as possible
+    ; This is decremented in C code when AP is finished executing
+    mov        edi, esi
+    add        edi, NumApsExecutingLocation
+    lock inc   dword [edi]
+
     mov         edi, esi
     add         edi, EnableExecuteDisableLocation
     cmp         byte [edi], 0

UefiCpuPkg/Library/MpInitLib/MpLib.c

@@ -662,6 +662,7 @@ ApWakeupFunction (
     // AP finished executing C code
     //
     InterlockedIncrement ((UINT32 *) &CpuMpData->FinishedCount);
+    InterlockedDecrement ((UINT32 *) &CpuMpData->MpCpuExchangeInfo->NumApsExecuting);
 
     //
     // Place AP is specified loop mode
@@ -765,6 +766,7 @@ FillExchangeInfoData (
 
   ExchangeInfo->CFunction       = (UINTN) ApWakeupFunction;
   ExchangeInfo->ApIndex         = 0;
+  ExchangeInfo->NumApsExecuting = 0;
   ExchangeInfo->InitFlag        = (UINTN) CpuMpData->InitFlag;
   ExchangeInfo->CpuInfo         = (CPU_INFO_IN_HOB *) (UINTN) CpuMpData->CpuInfoInHob;
   ExchangeInfo->CpuMpData       = CpuMpData;
@@ -934,13 +936,19 @@ WakeUpAP (
     }
     if (CpuMpData->InitFlag == ApInitConfig) {
       //
-      // Wait for all potential APs waken up in one specified period
+      // Wait for one potential AP waken up in one specified period
       //
+      if (CpuMpData->CpuCount == 0) {
         TimedWaitForApFinish (
           CpuMpData,
           PcdGet32 (PcdCpuMaxLogicalProcessorNumber) - 1,
           PcdGet32 (PcdCpuApInitTimeOutInMicroSeconds)
           );
+      }
+
+      while (CpuMpData->MpCpuExchangeInfo->NumApsExecuting != 0) {
+        CpuPause();
+      }
     } else {
       //
       // Wait all APs waken up if this is not the 1st broadcast of SIPI

UefiCpuPkg/Library/MpInitLib/MpLib.h

@@ -176,6 +176,7 @@ typedef struct {
   UINTN                 Cr3;
   UINTN                 InitFlag;
   CPU_INFO_IN_HOB       *CpuInfo;
+  UINTN                 NumApsExecuting;
   CPU_MP_DATA           *CpuMpData;
   UINTN                 InitializeFloatingPointUnitsAddress;
 } MP_CPU_EXCHANGE_INFO;

UefiCpuPkg/Library/MpInitLib/X64/MpEqu.inc

@@ -40,5 +40,6 @@ EnableExecuteDisableLocation  equ        LockLocation + 5Ch
 Cr3Location                   equ        LockLocation + 64h
 InitFlagLocation              equ        LockLocation + 6Ch
 CpuInfoLocation               equ        LockLocation + 74h
-InitializeFloatingPointUnitsAddress equ  LockLocation + 84h
+NumApsExecutingLocation       equ        LockLocation + 7Ch
+InitializeFloatingPointUnitsAddress equ  LockLocation + 8Ch
 

UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm

@@ -124,6 +124,12 @@ LongModeStart:
     cmp        qword [edi], 1       ; ApInitConfig
     jnz        GetApicId
 
+    ; Increment the number of APs executing here as early as possible
+    ; This is decremented in C code when AP is finished executing
+    mov        edi, esi
+    add        edi, NumApsExecutingLocation
+    lock inc   dword [edi]
+
     ; AP init
     mov        edi, esi
     add        edi, LockLocation