[PATCH v2] UefiCpuPkg/MpInitLib: Simplify logic in SwitchBsp

["Zhiguang Liu" <zhiguang.liu@intel.com>]

When switch bsp, old bsp and new bsp put CR0/CR4 into stack, and put IDT
and GDT register into a structure. After they exchange their stack, they
restore these registers. This logic is now implemented by assembly code.
This patch aims to reuse (Save/Restore)VolatileRegisters function to
replace such assembly code for better code readability.

Cc: Eric Dong <eric.dong@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Rahul Kumar <rahul1.kumar@intel.com>
Signed-off-by: Zhiguang Liu <zhiguang.liu@intel.com>
---
 .../Library/MpInitLib/Ia32/MpFuncs.nasm       | 20 +--------
 UefiCpuPkg/Library/MpInitLib/MpLib.c          | 35 ++++++++++++++-
 UefiCpuPkg/Library/MpInitLib/MpLib.h          | 43 +++++++++----------
 UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm | 21 ---------
 4 files changed, 56 insertions(+), 63 deletions(-)

diff --git a/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm b/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
index 28301bb8f0..1d67f510e9 100644
--- a/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
+++ b/UefiCpuPkg/Library/MpInitLib/Ia32/MpFuncs.nasm
@@ -284,15 +284,8 @@ ASM_PFX(AsmExchangeRole):
     ; edi contains OthersInfo pointer
     mov        edi, [ebp + 28h]
 
-    ;Store EFLAGS, GDTR and IDTR register to stack
+    ;Store EFLAGS to stack
     pushfd
-    mov        eax, cr4
-    push       eax       ; push cr4 firstly
-    mov        eax, cr0
-    push       eax
-
-    sgdt       [esi + CPU_EXCHANGE_ROLE_INFO.Gdtr]
-    sidt       [esi + CPU_EXCHANGE_ROLE_INFO.Idtr]
 
     ; Store the its StackPointer
     mov        [esi + CPU_EXCHANGE_ROLE_INFO.StackPointer],esp
@@ -308,13 +301,6 @@ WaitForOtherStored:
     jmp        WaitForOtherStored
 
 OtherStored:
-    ; Since another CPU already stored its state, load them
-    ; load GDTR value
-    lgdt       [edi + CPU_EXCHANGE_ROLE_INFO.Gdtr]
-
-    ; load IDTR value
-    lidt       [edi + CPU_EXCHANGE_ROLE_INFO.Idtr]
-
     ; load its future StackPointer
     mov        esp, [edi + CPU_EXCHANGE_ROLE_INFO.StackPointer]
 
@@ -331,10 +317,6 @@ WaitForOtherLoaded:
 
 OtherLoaded:
     ; since the other CPU already get the data it want, leave this procedure
-    pop        eax
-    mov        cr0, eax
-    pop        eax
-    mov        cr4, eax
     popfd
 
     popad
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.c b/UefiCpuPkg/Library/MpInitLib/MpLib.c
index 8d1f24370a..041a32e659 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.c
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.c
@@ -1,7 +1,7 @@
 /** @file
   CPU MP Initialize Library common functions.
 
-  Copyright (c) 2016 - 2021, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2016 - 2022, Intel Corporation. All rights reserved.<BR>
   Copyright (c) 2020, AMD Inc. All rights reserved.<BR>
 
   SPDX-License-Identifier: BSD-2-Clause-Patent
@@ -15,6 +15,29 @@
 
 EFI_GUID  mCpuInitMpLibHobGuid = CPU_INIT_MP_LIB_HOB_GUID;
 
+/**
+  Save the volatile registers required to be restored following INIT IPI.
+
+  @param[out]  VolatileRegisters    Returns buffer saved the volatile resisters
+**/
+VOID
+SaveVolatileRegisters (
+  OUT CPU_VOLATILE_REGISTERS  *VolatileRegisters
+  );
+
+/**
+  Restore the volatile registers following INIT IPI.
+
+  @param[in]  VolatileRegisters   Pointer to volatile resisters
+  @param[in]  IsRestoreDr         TRUE:  Restore DRx if supported
+                                  FALSE: Do not restore DRx
+**/
+VOID
+RestoreVolatileRegisters (
+  IN CPU_VOLATILE_REGISTERS  *VolatileRegisters,
+  IN BOOLEAN                 IsRestoreDr
+  );
+
 /**
   The function will check if BSP Execute Disable is enabled.
 
@@ -83,7 +106,12 @@ FutureBSPProc (
   CPU_MP_DATA  *DataInHob;
 
   DataInHob = (CPU_MP_DATA *)Buffer;
+  //
+  // Save and restore volatile registers when switch BSP
+  //
+  SaveVolatileRegisters (&DataInHob->APInfo.VolatileRegisters);
   AsmExchangeRole (&DataInHob->APInfo, &DataInHob->BSPInfo);
+  RestoreVolatileRegisters (&DataInHob->APInfo.VolatileRegisters, FALSE);
 }
 
 /**
@@ -2233,7 +2261,12 @@ SwitchBSPWorker (
   //
   WakeUpAP (CpuMpData, FALSE, ProcessorNumber, FutureBSPProc, CpuMpData, TRUE);
 
+  //
+  // Save and restore volatile registers when switch BSP
+  //
+  SaveVolatileRegisters (&CpuMpData->BSPInfo.VolatileRegisters);
   AsmExchangeRole (&CpuMpData->BSPInfo, &CpuMpData->APInfo);
+  RestoreVolatileRegisters (&CpuMpData->BSPInfo.VolatileRegisters, FALSE);
 
   //
   // Set the BSP bit of MSR_IA32_APIC_BASE on new BSP
diff --git a/UefiCpuPkg/Library/MpInitLib/MpLib.h b/UefiCpuPkg/Library/MpInitLib/MpLib.h
index 974fb76019..47b722cb2f 100644
--- a/UefiCpuPkg/Library/MpInitLib/MpLib.h
+++ b/UefiCpuPkg/Library/MpInitLib/MpLib.h
@@ -68,14 +68,31 @@ typedef struct {
   UINTN    Size;
 } MICROCODE_PATCH_INFO;
 
+//
+// CPU volatile registers around INIT-SIPI-SIPI
+//
+typedef struct {
+  UINTN              Cr0;
+  UINTN              Cr3;
+  UINTN              Cr4;
+  UINTN              Dr0;
+  UINTN              Dr1;
+  UINTN              Dr2;
+  UINTN              Dr3;
+  UINTN              Dr6;
+  UINTN              Dr7;
+  IA32_DESCRIPTOR    Gdtr;
+  IA32_DESCRIPTOR    Idtr;
+  UINT16             Tr;
+} CPU_VOLATILE_REGISTERS;
+
 //
 // CPU exchange information for switch BSP
 //
 typedef struct {
-  UINT8              State;        // offset 0
-  UINTN              StackPointer; // offset 4 / 8
-  IA32_DESCRIPTOR    Gdtr;         // offset 8 / 16
-  IA32_DESCRIPTOR    Idtr;         // offset 14 / 26
+  UINT8                     State;             // offset 0
+  UINTN                     StackPointer;      // offset 4 / 8
+  CPU_VOLATILE_REGISTERS    VolatileRegisters; // offset 8 / 16
 } CPU_EXCHANGE_ROLE_INFO;
 
 //
@@ -112,24 +129,6 @@ typedef enum {
   CpuStateDisabled
 } CPU_STATE;
 
-//
-// CPU volatile registers around INIT-SIPI-SIPI
-//
-typedef struct {
-  UINTN              Cr0;
-  UINTN              Cr3;
-  UINTN              Cr4;
-  UINTN              Dr0;
-  UINTN              Dr1;
-  UINTN              Dr2;
-  UINTN              Dr3;
-  UINTN              Dr6;
-  UINTN              Dr7;
-  IA32_DESCRIPTOR    Gdtr;
-  IA32_DESCRIPTOR    Idtr;
-  UINT16             Tr;
-} CPU_VOLATILE_REGISTERS;
-
 //
 // AP related data
 //
diff --git a/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm b/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
index cd95b03da8..b7f8d48504 100644
--- a/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
+++ b/UefiCpuPkg/Library/MpInitLib/X64/MpFuncs.nasm
@@ -482,22 +482,13 @@ ASM_PFX(AsmExchangeRole):
     push       r14
     push       r15
 
-    mov        rax, cr0
-    push       rax
-
-    mov        rax, cr4
-    push       rax
-
     ; rsi contains MyInfo pointer
     mov        rsi, rcx
 
     ; rdi contains OthersInfo pointer
     mov        rdi, rdx
 
-    ;Store EFLAGS, GDTR and IDTR regiter to stack
     pushfq
-    sgdt       [rsi + CPU_EXCHANGE_ROLE_INFO.Gdtr]
-    sidt       [rsi + CPU_EXCHANGE_ROLE_INFO.Idtr]
 
     ; Store the its StackPointer
     mov        [rsi + CPU_EXCHANGE_ROLE_INFO.StackPointer], rsp
@@ -513,12 +504,6 @@ WaitForOtherStored:
     jmp        WaitForOtherStored
 
 OtherStored:
-    ; Since another CPU already stored its state, load them
-    ; load GDTR value
-    lgdt       [rdi + CPU_EXCHANGE_ROLE_INFO.Gdtr]
-
-    ; load IDTR value
-    lidt       [rdi + CPU_EXCHANGE_ROLE_INFO.Idtr]
 
     ; load its future StackPointer
     mov        rsp, [rdi + CPU_EXCHANGE_ROLE_INFO.StackPointer]
@@ -538,12 +523,6 @@ OtherLoaded:
     ; since the other CPU already get the data it want, leave this procedure
     popfq
 
-    pop        rax
-    mov        cr4, rax
-
-    pop        rax
-    mov        cr0, rax
-
     pop        r15
     pop        r14
     pop        r13
-- 
2.31.1.windows.1