[Patch 2/2] UefiCpuPkg/MpInitLib: Enhance waiting for AP initialization logic.

[Eric Dong <eric.dong@intel.com>]

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>
---
 UefiCpuPkg/Library/MpInitLib/Ia32/MpEqu.inc    |  1 +
 UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm |  6 ++++++
 UefiCpuPkg/Library/MpInitLib/MpLib.c           | 20 ++++++++++++++------
 UefiCpuPkg/Library/MpInitLib/MpLib.h           |  1 +
 UefiCpuPkg/Library/MpInitLib/X64/MpEqu.inc     |  3 ++-
 UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm  |  6 ++++++
 6 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/Ia32/MpEqu.inc b/UefiCpuPkg/Library/MpInitLib/Ia32/MpEqu.inc
index 976af1f..bdfe0d3 100644
--- a/UefiCpuPkg/Library/MpInitLib/Ia32/MpEqu.inc
+++ b/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
 
diff --git a/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm b/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
index 1b9c6a6..2b6c27d 100644
--- a/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
+++ b/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
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index db923c9..48f930b 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/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
       //
-      TimedWaitForApFinish (
-        CpuMpData,
-        PcdGet32 (PcdCpuMaxLogicalProcessorNumber) - 1,
-        PcdGet32 (PcdCpuApInitTimeOutInMicroSeconds)
-        );
+      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
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index e41d2db..d13d5c0 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/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;
diff --git a/UefiCpuPkg/Library/MpInitLib/X64/MpEqu.inc b/UefiCpuPkg/Library/MpInitLib/X64/MpEqu.inc
index 114f4e0..d255ca5 100644
--- a/UefiCpuPkg/Library/MpInitLib/X64/MpEqu.inc
+++ b/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
 
diff --git a/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm b/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
index 4ada649..21d2786 100644
--- a/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
+++ b/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
-- 
2.7.0.windows.1