From f6afd876631e2f7d1e0eeda70c2bf299c5d83deb Mon Sep 17 00:00:00 2001
From: Khor Swee Aun <swee.aun.khor@intel.com>
Date: Mon, 20 Jan 2025 22:38:47 -0800
Subject: [PATCH] UefiCpuPkg/PiSmmCpuDxeSmm: SmmWaitForApArrival first sync
 check

Restructured the SmmWaitForApArrival first CPU synchronization
conditional checks and added new IsCpuSyncAlwaysNeeded check to determine
whether the sync should be executed unconditionally when a SMI occurs.

The first CPU synchronization shall executed unconditionally if the new
IsCpuSyncAlwaysNeeded check return TRUE. Otherwise, first CPU
synchronization is not executed unconditionally, and the decision to
synchronize should be based on the system configuration and status.

Signed-off-by: Khor Swee Aun <swee.aun.khor@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
index 546c94bae5..b19906b610 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
@@ -234,10 +234,9 @@ SmmWaitForApArrival (
 {
   UINT64   Timer;
   UINTN    Index;
-  BOOLEAN  LmceEn;
-  BOOLEAN  LmceSignal;
   UINT32   DelayedCount;
   UINT32   BlockedCount;
+  BOOLEAN  SyncNeeded;
 
   PERF_FUNCTION_BEGIN ();
 
@@ -246,11 +245,9 @@ SmmWaitForApArrival (
 
   ASSERT (SmmCpuSyncGetArrivedCpuCount (mSmmMpSyncData->SyncContext) <= mNumberOfCpus);
 
-  LmceEn     = FALSE;
-  LmceSignal = FALSE;
-  if (mMachineCheckSupported) {
-    LmceEn     = IsLmceOsEnabled ();
-    LmceSignal = IsLmceSignaled ();
+  SyncNeeded = IsCpuSyncAlwaysNeeded ();
+  if (!SyncNeeded) {
+    SyncNeeded = !(mMachineCheckSupported && IsLmceOsEnabled () && IsLmceSignaled ());
   }
 
   //
@@ -268,7 +265,7 @@ SmmWaitForApArrival (
   // Sync with APs 1st timeout
   //
   for (Timer = StartSyncTimer ();
-       !IsSyncTimerTimeout (Timer, mTimeoutTicker) && !(LmceEn && LmceSignal);
+       !IsSyncTimerTimeout (Timer, mTimeoutTicker) && SyncNeeded;
        )
   {
     mSmmMpSyncData->AllApArrivedWithException = AllCpusInSmmExceptBlockedDisabled ();