[PATCH v1 1/2] UefiCpuPkg: Extend SMM CPU Service with rendezvous support.

["Li, Zhihao" <zhihao.li@intel.com>]

From: Zhihao Li <zhihao.li@intel.com>

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

This patch extends the SMM CPU Service protocol with new interface
SmmWaitForAllProcessor(), which can be used by SMI handler to optionally
wait for other APs to complete SMM rendezvous in relaxed AP mode.

A new library SmmCpuRendezvousLib is provided to abstract the service
into library API to simple SMI handler code.

Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>

Signed-off-by: Zhihao Li <zhihao.li@intel.com>
---
 UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.c   | 109 ++++++++++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/CpuService.c                         |  65 ++++++++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c                          |  14 ++-
 UefiCpuPkg/PiSmmCpuDxeSmm/SyncTimer.c                          |   2 +-
 UefiCpuPkg/Include/Library/SmmCpuRendezvousLib.h               |  27 +++++
 UefiCpuPkg/Include/Protocol/SmmCpuService.h                    |  40 +++++++
 UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.inf |  32 ++++++
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h                     |  28 +++++
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf                   |   3 +-
 UefiCpuPkg/UefiCpuPkg.dec                                      |   5 +-
 10 files changed, 318 insertions(+), 7 deletions(-)

diff --git a/UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.c b/UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.c
new file mode 100644
index 0000000000..3c5cd51d0c
--- /dev/null
+++ b/UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.c
@@ -0,0 +1,109 @@
+/** @file
+SMM CPU Rendezvous library header file.
+
+Copyright (c) 2021, Intel Corporation. All rights reserved.<BR>
+SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+
+#include <Base.h>
+#include <Uefi.h>
+#include <Library/BaseLib.h>
+#include <Library/DebugLib.h>
+#include <Library/PcdLib.h>
+#include <Library/SmmServicesTableLib.h>
+#include <Protocol/SmmCpuService.h>
+#include <Library/SmmCpuRendezvousLib.h>
+
+STATIC EDKII_SMM_CPU_SERVICE_PROTOCOL  *mSmmCpuService = NULL;
+
+/**
+  This routine wait for all AP processors to arrive in SMM.
+
+  @param  BlockingMode  Blocking mode or non-blocking mode.
+
+  @retval TRUE   All processors checked in to SMM
+  @retval FALSE  Some processor not checked in to SMM
+
+**/
+EFI_STATUS
+EFIAPI
+SmmWaitForAllProcessor (
+  IN  BOOLEAN  BlockingMode
+  )
+{
+  EFI_STATUS Status;
+
+  if (mSmmCpuService == NULL) {
+    return TRUE;
+  }
+
+  Status = mSmmCpuService->WaitForAllProcessor (
+              mSmmCpuService,
+              BlockingMode
+              );
+  return EFI_ERROR(Status) ? FALSE : TRUE;
+}
+
+/**
+  Register status code callback function only when Report Status Code protocol
+  is installed.
+
+  @param Protocol       Points to the protocol's unique identifier.
+  @param Interface      Points to the interface instance.
+  @param Handle         The handle on which the interface was installed.
+
+  @retval EFI_SUCCESS   Notification runs successfully.
+
+**/
+EFI_STATUS
+EFIAPI
+SmmCpuServiceProtocolNotify (
+  IN CONST EFI_GUID        *Protocol,
+  IN VOID                  *Interface,
+  IN EFI_HANDLE            Handle
+  )
+{
+  EFI_STATUS                   Status;
+
+  Status = gSmst->SmmLocateProtocol (
+                    &gEdkiiSmmCpuServiceProtocolGuid,
+                    NULL,
+                    (VOID **) &mSmmCpuService
+                    );
+  ASSERT_EFI_ERROR (Status);
+
+  return EFI_SUCCESS;
+}
+
+/**
+  The constructor function
+
+  @param[in]  ImageHandle  The firmware allocated handle for the EFI image.
+  @param[in]  SystemTable  A pointer to the EFI System Table.
+
+  @retval EFI_SUCCESS      The constructor always returns EFI_SUCCESS.
+
+**/
+EFI_STATUS
+EFIAPI
+SmmCpuRendezvousLibConstructor (
+  IN EFI_HANDLE        ImageHandle,
+  IN EFI_SYSTEM_TABLE  *SystemTable
+  )
+{
+  EFI_STATUS  Status;
+  VOID        *Registration;
+
+  Status = gSmst->SmmLocateProtocol (&gEdkiiSmmCpuServiceProtocolGuid, NULL, (VOID **) &mSmmCpuService);
+  if (EFI_ERROR (Status)) {
+    Status = gSmst->SmmRegisterProtocolNotify (
+                &gEdkiiSmmCpuServiceProtocolGuid,
+                SmmCpuServiceProtocolNotify,
+                &Registration
+                );
+    ASSERT_EFI_ERROR (Status);
+  }
+  return EFI_SUCCESS;
+}
\ No newline at end of file
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/CpuService.c b/UefiCpuPkg/PiSmmCpuDxeSmm/CpuService.c
index 5d624f8e9e..34019c24ff 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/CpuService.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/CpuService.c
@@ -20,6 +20,19 @@ EFI_SMM_CPU_SERVICE_PROTOCOL  mSmmCpuService = {
   SmmRegisterExceptionHandler
 };
 
+//
+// EDKII SMM CPU Service Protocol instance
+//
+EDKII_SMM_CPU_SERVICE_PROTOCOL  mEdkiiSmmCpuService = {
+  SmmGetProcessorInfo,
+  SmmSwitchBsp,
+  SmmAddProcessor,
+  SmmRemoveProcessor,
+  SmmWhoAmI,
+  SmmRegisterExceptionHandler,
+  SmmWaitForAllProcessor
+};
+
 /**
   Gets processor information on the requested processor at the instant this call is made.
 
@@ -365,5 +378,57 @@ InitializeSmmCpuServices (
                     &mSmmCpuService
                     );
   ASSERT_EFI_ERROR (Status);
+
+  Status = gSmst->SmmInstallProtocolInterface (
+                    &Handle,
+                    &gEdkiiSmmCpuServiceProtocolGuid,
+                    EFI_NATIVE_INTERFACE,
+                    &mEdkiiSmmCpuService
+                    );
+  ASSERT_EFI_ERROR (Status);
+  return Status;
+}
+
+/**
+  Wait for all processors enterring SMM until all CPUs are already synchronized or not.
+
+  @param  This                   A pointer to the EDKII_SMM_CPU_SERVICE_PROTOCOL instance.
+  @param  BlockingMode           Blocking mode or non-blocking mode.
+
+  @retval EFI_SUCCESS    All avaiable APs arrived.
+  @retval EFI_TIMEOUT    Wait for all APs until timeout.
+**/
+EFI_STATUS
+EFIAPI
+SmmWaitForAllProcessor (
+  IN  EDKII_SMM_CPU_SERVICE_PROTOCOL  *This,
+  IN  BOOLEAN                          BlockingMode
+  )
+{
+  EFI_STATUS                Status;
+
+  //
+  // Return success immediately if all CPUs are already synchronized.
+  //
+  if (mSmmMpSyncData->AllApArrivedWithException) {
+    Status = EFI_SUCCESS;
+    goto ON_EXIT;
+  }
+
+  if (!BlockingMode) {
+    Status = EFI_TIMEOUT;
+    goto ON_EXIT;
+  }
+
+  //
+  // There are some APs outside SMM, Wait for all avaiable APs to arrive.
+  //
+  SmmWaitForApArrival (BlockingMode);
+  Status = mSmmMpSyncData->AllApArrivedWithException ? EFI_SUCCESS : EFI_TIMEOUT;
+
+ON_EXIT:
+  if (!mSmmMpSyncData->AllApArrivedWithException){
+    DEBUG ((EFI_D_INFO, "EdkiiSmmWaitForAllApArrival: Timeout to wait all APs arrival\n"));
+  }
   return Status;
 }
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
index 882dee4fe2..9c7b16728a 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
@@ -261,7 +261,7 @@ IsLmceSignaled (
 **/
 VOID
 SmmWaitForApArrival (
-  VOID
+  IN  BOOLEAN  BlockingMode
   )
 {
   UINT64   Timer;
@@ -270,7 +270,14 @@ SmmWaitForApArrival (
   BOOLEAN  LmceSignal;
 
   ASSERT (*mSmmMpSyncData->Counter <= mNumberOfCpus);
-
+  
+  //
+  // if block is False, do not wait and return immediately.
+  //
+  if (!BlockingMode){
+    return;
+  }
+  
   LmceEn     = FALSE;
   LmceSignal = FALSE;
   if (mMachineCheckSupported) {
@@ -511,7 +518,7 @@ BSPHandler (
     //
     // Wait for APs to arrive
     //
-    SmmWaitForApArrival ();
+    SmmWaitForApArrival(TRUE);
 
     //
     // Lock the counter down and retrieve the number of APs
@@ -1886,6 +1893,7 @@ InitializeMpSyncData (
     *mSmmMpSyncData->Counter       = 0;
     *mSmmMpSyncData->InsideSmm     = FALSE;
     *mSmmMpSyncData->AllCpusInSync = FALSE;
+    mSmmMpSyncData->AllApArrivedWithException = FALSE;
 
     for (CpuIndex = 0; CpuIndex < gSmmCpuPrivate->SmmCoreEntryContext.NumberOfCpus; CpuIndex++) {
       mSmmMpSyncData->CpuData[CpuIndex].Busy =
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SyncTimer.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SyncTimer.c
index 0c070c5736..844263f889 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SyncTimer.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SyncTimer.c
@@ -8,7 +8,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 
 #include "PiSmmCpuDxeSmm.h"
 
-UINT64  mTimeoutTicker = 0;
+UINT64   mTimeoutTicker = 0;
 //
 //  Number of counts in a roll-over cycle of the performance counter.
 //
diff --git a/UefiCpuPkg/Include/Library/SmmCpuRendezvousLib.h b/UefiCpuPkg/Include/Library/SmmCpuRendezvousLib.h
new file mode 100644
index 0000000000..f245c3a1c9
--- /dev/null
+++ b/UefiCpuPkg/Include/Library/SmmCpuRendezvousLib.h
@@ -0,0 +1,27 @@
+/** @file
+SMM CPU Rendezvous library header file.
+
+Copyright (c) 2021, Intel Corporation. All rights reserved.<BR>
+SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#ifndef _SMM_CPU_RENDEZVOUS_H_
+#define _SMM_CPU_RENDEZVOUS_H_
+
+/**
+  This routine wait for all AP processors to arrive in SMM.
+
+  @param  BlockingMode  Blocking mode or non-blocking mode.
+  
+  @retval TRUE   All processors checked in to SMM
+  @retval FALSE  Some processor not checked in to SMM
+
+**/
+EFI_STATUS
+EFIAPI
+SmmWaitForAllProcessor (
+  IN  BOOLEAN BlockingMode
+  );
+
+#endif
diff --git a/UefiCpuPkg/Include/Protocol/SmmCpuService.h b/UefiCpuPkg/Include/Protocol/SmmCpuService.h
index 952767afce..0ace5356a9 100644
--- a/UefiCpuPkg/Include/Protocol/SmmCpuService.h
+++ b/UefiCpuPkg/Include/Protocol/SmmCpuService.h
@@ -200,4 +200,44 @@ struct _EFI_SMM_CPU_SERVICE_PROTOCOL {
 
 extern EFI_GUID  gEfiSmmCpuServiceProtocolGuid;
 
+//
+//  EDKII_SMM_CPU_SERVICE_PROTOCOL extends the EFI_SMM_CPU_SERVICE_PROTOCOL
+//  with rendezvous service support.
+//
+#define EDKII_SMM_CPU_SERVICE_PROTOCOL_GUID \
+  { \
+    0xaa00d50b, 0x4911, 0x428f, {0xb9, 0x1a, 0xa5, 0x9d, 0xdb, 0x13, 0xe2, 0x4c} \
+  }
+
+typedef struct _EDKII_SMM_CPU_SERVICE_PROTOCOL EDKII_SMM_CPU_SERVICE_PROTOCOL;
+
+/**
+  Wait for all APs to arrive SMM mode in given timeout constraint.
+
+  @param  This                  A pointer to the SMM_CPU_SERVICE_PROTOCOL instance.
+  @param  BlockingMode          Block or non-block mode.
+
+  @retval EFI_SUCCESS           All APs have arrived SMM mode except SMI disabled APs.
+  @retval EFI_TIMEOUT           There are APs not in SMM mode in given timeout constraint.
+
+**/
+typedef
+EFI_STATUS
+(EFIAPI *EDKII_WAIT_FOR_ALL_PROCESSOR) (
+  IN  EDKII_SMM_CPU_SERVICE_PROTOCOL  *This,
+  IN  BOOLEAN                         BlockingMode
+  );
+
+struct _EDKII_SMM_CPU_SERVICE_PROTOCOL {
+  EFI_SMM_GET_PROCESSOR_INFO          GetProcessorInfo;
+  EFI_SMM_SWITCH_BSP                  SwitchBsp;
+  EFI_SMM_ADD_PROCESSOR               AddProcessor;
+  EFI_SMM_REMOVE_PROCESSOR            RemoveProcessor;
+  EFI_SMM_WHOAMI                      WhoAmI;
+  EFI_SMM_REGISTER_EXCEPTION_HANDLER  RegisterExceptionHandler;
+  EDKII_WAIT_FOR_ALL_PROCESSOR        WaitForAllProcessor;
+};
+
+extern EFI_GUID gEdkiiSmmCpuServiceProtocolGuid;
+
 #endif
diff --git a/UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.inf b/UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.inf
new file mode 100644
index 0000000000..aff77c5a18
--- /dev/null
+++ b/UefiCpuPkg/Library/SmmCpuRendezvousLib/SmmCpuRendezvousLib.inf
@@ -0,0 +1,32 @@
+## @file
+# Component description file for CPU SMM Rendezvous check library
+##
+
+[Defines]
+  INF_VERSION                    = 0x00010005
+  BASE_NAME                      = SmmCpuRendezvousLib
+  FILE_GUID                      = 1509Bb36-9Ba4-438B-B195-Ac5914Db14E2
+  MODULE_TYPE                    = DXE_SMM_DRIVER
+  VERSION_STRING                 = 1.0
+  LIBRARY_CLASS                  = SmmCpuRendezvousLib
+  CONSTRUCTOR                    = SmmCpuRendezvousLibConstructor
+
+[Sources]
+  SmmCpuRendezvousLib.c
+
+[Packages]
+  MdePkg/MdePkg.dec
+  UefiCpuPkg/UefiCpuPkg.dec
+
+[LibraryClasses]
+  BaseLib
+  DebugLib
+  PcdLib
+  SmmServicesTableLib
+
+[Pcd]
+  gUefiCpuPkgTokenSpaceGuid.PcdCpuSmmApSyncTimeout                 ## CONSUMES
+
+[Protocols]
+  gEdkiiSmmCpuServiceProtocolGuid
+
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
index 26d07c5b5e..c6f31ace77 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
@@ -428,6 +428,7 @@ typedef struct {
   volatile SMM_CPU_SYNC_MODE    EffectiveSyncMode;
   volatile BOOLEAN              SwitchBsp;
   volatile BOOLEAN              *CandidateBsp;
+  volatile BOOLEAN              AllApArrivedWithException;
   EFI_AP_PROCEDURE              StartupProcedure;
   VOID                          *StartupProcArgs;
 } SMM_DISPATCHER_MP_SYNC_DATA;
@@ -1488,4 +1489,31 @@ IsRestrictedMemoryAccess (
   VOID
   );
 
+/**
+  Choose blocking or non-blocking mode to Wait for all APs
+
+  @param  This                  A pointer to the SMM_CPU_SERVICE_PROTOCOL instance.
+  @param  BlockingMode          Blocking or non-blocking mode
+
+  @retval EFI_SUCCESS           All APs have arrived SMM mode except SMI disabled APs.
+  @retval EFI_TIMEOUT           There are APs not in SMM mode in given timeout constraint.
+
+**/
+EFI_STATUS
+EFIAPI
+SmmWaitForAllProcessor (
+  IN  EDKII_SMM_CPU_SERVICE_PROTOCOL  *This,
+  IN  BOOLEAN                          BlockingMode
+  );
+
+/**
+  Choose blocking or non-blocking mode to wait for all APs. True for Blocking and false for not. 
+  Insure when this function returns, no AP will execute normal mode code before entering SMM, except SMI disabled APs.
+
+**/
+VOID
+SmmWaitForApArrival (
+  IN  BOOLEAN  BlockingMode                     
+  );
+
 #endif
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf
index 0e88071c70..1af7280d18 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.inf
@@ -107,7 +107,8 @@
   gEfiSmmReadyToLockProtocolGuid           ## NOTIFY
   gEfiSmmCpuServiceProtocolGuid            ## PRODUCES
   gEdkiiSmmMemoryAttributeProtocolGuid     ## PRODUCES
-  gEfiMmMpProtocolGuid                    ## PRODUCES
+  gEfiMmMpProtocolGuid                     ## PRODUCES
+  gEdkiiSmmCpuServiceProtocolGuid          ## PRODUCES
 
 [Guids]
   gEfiAcpiVariableGuid                     ## SOMETIMES_CONSUMES ## HOB # it is used for S3 boot.
diff --git a/UefiCpuPkg/UefiCpuPkg.dec b/UefiCpuPkg/UefiCpuPkg.dec
index 7de66fde67..c170e7d75d 100644
--- a/UefiCpuPkg/UefiCpuPkg.dec
+++ b/UefiCpuPkg/UefiCpuPkg.dec
@@ -77,7 +77,8 @@
 
 [Protocols]
   ## Include/Protocol/SmmCpuService.h
-  gEfiSmmCpuServiceProtocolGuid  = { 0x1d202cab, 0xc8ab, 0x4d5c, { 0x94, 0xf7, 0x3c, 0xfc, 0xc0, 0xd3, 0xd3, 0x35 }}
+  gEfiSmmCpuServiceProtocolGuid   = { 0x1d202cab, 0xc8ab, 0x4d5c, { 0x94, 0xf7, 0x3c, 0xfc, 0xc0, 0xd3, 0xd3, 0x35 }}
+  gEdkiiSmmCpuServiceProtocolGuid = { 0xaa00d50b, 0x4911, 0x428f, { 0xb9, 0x1a, 0xa5, 0x9d, 0xdb, 0x13, 0xe2, 0x4c }}
 
   ## Include/Protocol/SmMonitorInit.h
   gEfiSmMonitorInitProtocolGuid  = { 0x228f344d, 0xb3de, 0x43bb, { 0xa4, 0xd7, 0xea, 0x20, 0xb, 0x1b, 0x14, 0x82 }}
@@ -304,7 +305,7 @@
   #   0x00  - Traditional CPU synchronization method.<BR>
   #   0x01  - Relaxed CPU synchronization method.<BR>
   # @Prompt SMM CPU Synchronization Method.
-  gUefiCpuPkgTokenSpaceGuid.PcdCpuSmmSyncMode|0x00|UINT8|0x60000014
+  gUefiCpuPkgTokenSpaceGuid.PcdCpuSmmSyncMode|0x01|UINT8|0x60000014
 
   ## Specifies the On-demand clock modulation duty cycle when ACPI feature is enabled.
   # @Prompt The encoded values for target duty cycle modulation.
-- 
2.26.2.windows.1