From 4b7bd4c591a81a290b31e9d1a94c4b8be787989e Mon Sep 17 00:00:00 2001
From: "Liu, Zhiguang" <Zhiguang.Liu@intel.com>
Date: Fri, 26 Aug 2022 15:39:01 +0800
Subject: UefiCpuPkg: Enhance logic in InitializeMpExceptionStackSwitchHandlers

Parallelly run the function to SeparateExceptionStacks for all CPUs and
allocate buffers together for better performance.

Cc: Eric Dong <eric.dong@intel.com>
Reviewed-by: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Signed-off-by: Zhiguang Liu <zhiguang.liu@intel.com>
---
 UefiCpuPkg/CpuDxe/CpuMp.c | 104 +++++++++++++++++++++++++++++-----------------
 1 file changed, 65 insertions(+), 39 deletions(-)

(limited to 'UefiCpuPkg/CpuDxe')

diff --git a/UefiCpuPkg/CpuDxe/CpuMp.c b/UefiCpuPkg/CpuDxe/CpuMp.c
index f3ca813d2a..e7575d9b80 100644
--- a/UefiCpuPkg/CpuDxe/CpuMp.c
+++ b/UefiCpuPkg/CpuDxe/CpuMp.c
@@ -600,8 +600,9 @@ CollectBistDataFromHob (
 // Structure for InitializeSeparateExceptionStacks
 //
 typedef struct {
-  VOID     *Buffer;
-  UINTN    *BufferSize;
+  VOID          *Buffer;
+  UINTN         BufferSize;
+  EFI_STATUS    Status;
 } EXCEPTION_STACK_SWITCH_CONTEXT;
 
 /**
@@ -620,9 +621,18 @@ InitializeExceptionStackSwitchHandlers (
   )
 {
   EXCEPTION_STACK_SWITCH_CONTEXT  *SwitchStackData;
+  UINTN                           Index;
 
+  MpInitLibWhoAmI (&Index);
   SwitchStackData = (EXCEPTION_STACK_SWITCH_CONTEXT *)Buffer;
-  InitializeSeparateExceptionStacks (SwitchStackData->Buffer, SwitchStackData->BufferSize);
+
+  //
+  // This may be called twice for each Cpu. Only run InitializeSeparateExceptionStacks
+  // if this is the first call or the first call failed because of size too small.
+  //
+  if ((SwitchStackData[Index].Status == EFI_NOT_STARTED) || (SwitchStackData[Index].Status == EFI_BUFFER_TOO_SMALL)) {
+    SwitchStackData[Index].Status = InitializeSeparateExceptionStacks (SwitchStackData[Index].Buffer, &SwitchStackData[Index].BufferSize);
+  }
 }
 
 /**
@@ -638,53 +648,69 @@ InitializeMpExceptionStackSwitchHandlers (
   )
 {
   UINTN                           Index;
-  UINTN                           Bsp;
-  EXCEPTION_STACK_SWITCH_CONTEXT  SwitchStackData;
+  EXCEPTION_STACK_SWITCH_CONTEXT  *SwitchStackData;
   UINTN                           BufferSize;
+  EFI_STATUS                      Status;
+  UINT8                           *Buffer;
 
-  SwitchStackData.BufferSize = &BufferSize;
-  MpInitLibWhoAmI (&Bsp);
-
+  SwitchStackData = AllocateZeroPool (mNumberOfProcessors * sizeof (EXCEPTION_STACK_SWITCH_CONTEXT));
+  ASSERT (SwitchStackData != NULL);
   for (Index = 0; Index < mNumberOfProcessors; ++Index) {
-    SwitchStackData.Buffer = NULL;
-    BufferSize             = 0;
+    //
+    // Because the procedure may runs multiple times, use the status EFI_NOT_STARTED
+    // to indicate the procedure haven't been run yet.
+    //
+    SwitchStackData[Index].Status = EFI_NOT_STARTED;
+  }
 
-    if (Index == Bsp) {
-      InitializeExceptionStackSwitchHandlers (&SwitchStackData);
+  Status = MpInitLibStartupAllCPUs (
+             InitializeExceptionStackSwitchHandlers,
+             0,
+             SwitchStackData
+             );
+  ASSERT_EFI_ERROR (Status);
+
+  BufferSize = 0;
+  for (Index = 0; Index < mNumberOfProcessors; ++Index) {
+    if (SwitchStackData[Index].Status == EFI_BUFFER_TOO_SMALL) {
+      ASSERT (SwitchStackData[Index].BufferSize != 0);
+      BufferSize += SwitchStackData[Index].BufferSize;
     } else {
-      //
-      // AP might need different buffer size from BSP.
-      //
-      MpInitLibStartupThisAP (InitializeExceptionStackSwitchHandlers, Index, NULL, 0, (VOID *)&SwitchStackData, NULL);
+      ASSERT (SwitchStackData[Index].Status == EFI_SUCCESS);
+      ASSERT (SwitchStackData[Index].BufferSize == 0);
     }
+  }
 
-    if (BufferSize == 0) {
-      continue;
+  if (BufferSize != 0) {
+    Buffer = AllocateRuntimeZeroPool (BufferSize);
+    ASSERT (Buffer != NULL);
+    BufferSize = 0;
+    for (Index = 0; Index < mNumberOfProcessors; ++Index) {
+      if (SwitchStackData[Index].Status == EFI_BUFFER_TOO_SMALL) {
+        SwitchStackData[Index].Buffer = (VOID *)(&Buffer[BufferSize]);
+        BufferSize                   += SwitchStackData[Index].BufferSize;
+        DEBUG ((
+          DEBUG_INFO,
+          "Buffer[cpu%lu] for InitializeExceptionStackSwitchHandlers: 0x%lX with size 0x%lX\n",
+          (UINT64)(UINTN)Index,
+          (UINT64)(UINTN)SwitchStackData[Index].Buffer,
+          (UINT64)(UINTN)SwitchStackData[Index].BufferSize
+          ));
+      }
     }
 
-    SwitchStackData.Buffer = AllocateRuntimeZeroPool (BufferSize);
-    ASSERT (SwitchStackData.Buffer != NULL);
-    DEBUG ((
-      DEBUG_INFO,
-      "Buffer[cpu%lu] for InitializeExceptionStackSwitchHandlers: 0x%lX with size 0x%x\n",
-      (UINT64)(UINTN)Index,
-      (UINT64)(UINTN)SwitchStackData.Buffer,
-      (UINT32)BufferSize
-      ));
-
-    if (Index == Bsp) {
-      InitializeExceptionStackSwitchHandlers (&SwitchStackData);
-    } else {
-      MpInitLibStartupThisAP (
-        InitializeExceptionStackSwitchHandlers,
-        Index,
-        NULL,
-        0,
-        (VOID *)&SwitchStackData,
-        NULL
-        );
+    Status = MpInitLibStartupAllCPUs (
+               InitializeExceptionStackSwitchHandlers,
+               0,
+               SwitchStackData
+               );
+    ASSERT_EFI_ERROR (Status);
+    for (Index = 0; Index < mNumberOfProcessors; ++Index) {
+      ASSERT (SwitchStackData[Index].Status == EFI_SUCCESS);
     }
   }
+
+  FreePool (SwitchStackData);
 }
 
 /**
-- 
cgit v1.2.3