[PATCH v3 1/2] UefiCpuPkg/MpInitLib: Restore IDT context for APs.

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

REF: https://bugzilla.tianocore.org/show_bug.cgi?id=2683

This patch fixes an assertion because AP can't find the CpuMpData.
When AP is waken up through Init-Sipi-Sipi, AP's IDT should
be restored to pre-allocated buffer so AP can get the CpuMpData
through the IDT base address.
Current code already has logic to handle this when CpuMpData->
InitFlag is ApInitConfig but misses the logic
when CpuMpData->InitFlag is ApInitReconfig.
This patch fixes this gap.

Cc: Ray Ni <ray.ni@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Chandana Kumar <chandana.c.kumar@intel.com>
Signed-off-by: Eric Dong <eric.dong@intel.com>
---
V3:
  Remove invalid save volatile registers process. Refine restore
  volatile registers process.

V2: 
  Enhance code to remove CpuMpData->ApLoopMode == ApInHltLoop check. 

 UefiCpuPkg/Library/MpInitLib/MpLib.c | 34 +++++++++++++++++++---------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index 64a4c3546e..7fd757b428 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -686,18 +686,31 @@ ApWakeupFunction (
         WAKEUP_AP_SIGNAL,
         0
         );
-      if (CpuMpData->ApLoopMode == ApInHltLoop) {
-        //
-        // Restore AP's volatile registers saved
-        //
-        RestoreVolatileRegisters (&CpuMpData->CpuData[ProcessorNumber].VolatileRegisters, TRUE);
-      } else {
+
+      if (CpuMpData->InitFlag == ApInitReconfig) {
         //
-        // The CPU driver might not flush TLB for APs on spot after updating
-        // page attributes. AP in mwait loop mode needs to take care of it when
-        // woken up.
+        // ApInitReconfig happens when:
+        // 1. AP is re-enabled after it's disabled, in either PEI or DXE phase.
+        // 2. AP is initialized in DXE phase. 
+        // In either case, use the volatile registers value derived from BSP.
+        // NOTE: IDTR.BASE stored in CpuMpData->CpuData[0].VolatileRegisters points to a
+        //   different IDT shared by all APs.
         //
-        CpuFlushTlb ();
+        RestoreVolatileRegisters (&CpuMpData->CpuData[0].VolatileRegisters, FALSE);
+      }  else {
+        if (CpuMpData->ApLoopMode == ApInHltLoop) {
+          //
+          // Restore AP's volatile registers saved before AP is halted
+          //
+          RestoreVolatileRegisters (&CpuMpData->CpuData[ProcessorNumber].VolatileRegisters, TRUE);
+        } else {
+          //
+          // The CPU driver might not flush TLB for APs on spot after updating
+          // page attributes. AP in mwait loop mode needs to take care of it when
+          // woken up.
+          //
+          CpuFlushTlb ();
+        }
       }
 
       if (GetApState (&CpuMpData->CpuData[ProcessorNumber]) == CpuStateReady) {
@@ -1780,7 +1793,6 @@ MpInitLibInitialize (
       InitializeSpinLock(&CpuMpData->CpuData[Index].ApLock);
       CpuMpData->CpuData[Index].CpuHealthy = (CpuInfoInHob[Index].Health == 0)? TRUE:FALSE;
       CpuMpData->CpuData[Index].ApFunction = 0;
-      CopyMem (&CpuMpData->CpuData[Index].VolatileRegisters, &VolatileRegisters, sizeof (CPU_VOLATILE_REGISTERS));
     }
   }
 
-- 
2.23.0.windows.1