From 9ee533479651d29a254a184a56d01809353f0bcc Mon Sep 17 00:00:00 2001
From: Jiaxin Wu <jiaxin.wu@intel.com>
Date: Wed, 26 Jun 2024 13:28:56 +0800
Subject: [PATCH] UefiCpuPkg/PiSmmCpuDxeSmm: Define mIsStandaloneMm to indicate
 SMM or MM

Define the mIsStandaloneMm to indicate it's the MM_STANDALONE MM CPU
driver or DXE_SMM_DRIVER SMM CPU driver execution.

With mIsStandaloneMm, GetMpInformationFromMpServices() can be skipped
for the MM CPU since it can not call the
GetMpInformationFromMpServices() due to the NON-SMM MP Services usage
for the MP Information.

Signed-off-by: Jiaxin Wu <jiaxin.wu@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Cc: Gerd Hoffmann <kraxel@redhat.com>
Cc: Star Zeng <star.zeng@intel.com>
Cc: Dun Tan <dun.tan@intel.com>
Cc: Hongbin1 Zhang <hongbin1.zhang@intel.com>
Cc: Wei6 Xu <wei6.xu@intel.com>
Cc: Yuanhao Xie <yuanhao.xie@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.c    | 11 +++++---
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.h    |  2 ++
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c    |  6 +++++
 .../PiSmmCpuDxeSmm/PiSmmCpuStandaloneMm.c     | 25 +++++++++++++++++++
 4 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.c b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.c
index 78bfc8a3e7..48b6fe62c4 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.c
@@ -625,9 +625,14 @@ GetMpInformation (
   HobCount              = 0;
 
   FirstMpInfo2Hob = GetFirstGuidHob (&gMpInformation2HobGuid);
-  if (FirstMpInfo2Hob == NULL) {
-    DEBUG ((DEBUG_INFO, "%a: [INFO] gMpInformation2HobGuid HOB not found.\n", __func__));
-    return GetMpInformationFromMpServices (NumberOfCpus, MaxNumberOfCpus);
+
+  if (mIsStandaloneMm) {
+    ASSERT (FirstMpInfo2Hob != NULL);
+  } else {
+    if (FirstMpInfo2Hob == NULL) {
+      DEBUG ((DEBUG_INFO, "%a: [INFO] gMpInformation2HobGuid HOB not found.\n", __func__));
+      return GetMpInformationFromMpServices (NumberOfCpus, MaxNumberOfCpus);
+    }
   }
 
   GuidHob = FirstMpInfo2Hob;
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.h
index 9272b9df3d..828080946e 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuCommon.h
@@ -253,6 +253,8 @@ typedef struct {
   LIST_ENTRY                        *FirstFreeToken;
 } SMM_CPU_PRIVATE_DATA;
 
+extern const BOOLEAN  mIsStandaloneMm;
+
 extern SMM_CPU_PRIVATE_DATA  *gSmmCpuPrivate;
 extern CPU_HOT_PLUG_DATA     mCpuHotPlugData;
 extern UINTN                 mMaxNumberOfCpus;
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c
index 8995b54226..97b2bc6efb 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c
@@ -12,6 +12,12 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include "PiSmmCpuCommon.h"
 #include <Library/UefiBootServicesTableLib.h>
 
+//
+// TRUE to indicate it's the MM_STANDALONE MM CPU driver.
+// FALSE to indicate it's the DXE_SMM_DRIVER SMM CPU driver.
+//
+const BOOLEAN  mIsStandaloneMm = FALSE;
+
 /**
   To get system port address of the SMI Command Port in FADT table.
 
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuStandaloneMm.c b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuStandaloneMm.c
index c0378f5bd2..235e53cde5 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuStandaloneMm.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuStandaloneMm.c
@@ -9,6 +9,12 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 
 #include "PiSmmCpuCommon.h"
 
+//
+// TRUE to indicate it's the MM_STANDALONE MM CPU driver.
+// FALSE to indicate it's the DXE_SMM_DRIVER SMM CPU driver.
+//
+const BOOLEAN  mIsStandaloneMm = TRUE;
+
 /**
   To get system port address of the SMI Command Port.
 
@@ -94,3 +100,22 @@ GetAcpiS3EnableFlag (
     mAcpiS3Enable = MmAcpiS3EnableHob->AcpiS3Enable;
   }
 }
+
+/**
+  Extract NumberOfCpus, MaxNumberOfCpus and EFI_PROCESSOR_INFORMATION.
+
+  @param[out] NumberOfCpus           Pointer to NumberOfCpus.
+  @param[out] MaxNumberOfCpus        Pointer to MaxNumberOfCpus.
+
+  @retval ProcessorInfo              Pointer to EFI_PROCESSOR_INFORMATION buffer.
+**/
+EFI_PROCESSOR_INFORMATION *
+GetMpInformationFromMpServices (
+  OUT UINTN  *NumberOfCpus,
+  OUT UINTN  *MaxNumberOfCpus
+  )
+{
+  ASSERT (FALSE);
+
+  return NULL;
+}