[PATCH 2/3] UefiCpuPkg/PiSmmCpu: Refine CR2 save/restore logic

["Ni, Ray" <ray.ni@intel.com>]

Because IsStaticPageTableEnabled() is added for both IA32 and x64
build, the CR2 save/restore logic can be refined:
1. Remove arch specific SaveCr2() /  RestoreCr2() implementation;
2. Conditionally save and restore CR2 in SmiRendezvous().

Signed-off-by: Ray Ni <ray.ni@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Vanguput Narendra K <narendra.k.vanguput@intel.com>
---
 UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c   | 25 -------------------
 UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c      |  8 ++++--
 UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h | 22 ----------------
 UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c    | 29 ----------------------
 4 files changed, 6 insertions(+), 78 deletions(-)

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
index 2a9af4b77d..cae23d6d1d 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/Ia32/PageTbl.c
@@ -327,28 +327,3 @@ SetPageTableAttributes (
   return ;
 }
 
-/**
-  This function returns with no action for 32 bit.
-
-  @param[out]  *Cr2  Pointer to variable to hold CR2 register value.
-**/
-VOID
-SaveCr2 (
-  OUT UINTN  *Cr2
-  )
-{
-  return ;
-}
-
-/**
-  This function returns with no action for 32 bit.
-
-  @param[in]  Cr2  Value to write into CR2 register.
-**/
-VOID
-RestoreCr2 (
-  IN UINTN  Cr2
-  )
-{
-  return ;
-}
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
index ef16997547..5d0124b368 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
@@ -1575,7 +1575,9 @@ SmiRendezvous (
   // when using on-demand paging for above 4G memory.
   //
   Cr2 = 0;
-  SaveCr2 (&Cr2);
+  if (!IsStaticPageTableEnabled ()) {
+    Cr2 = AsmReadCr2 ();
+  }
 
   //
   // Call the user register Startup function first.
@@ -1725,7 +1727,9 @@ Exit:
   //
   // Restore Cr2
   //
-  RestoreCr2 (Cr2);
+  if (!IsStaticPageTableEnabled ()) {
+    AsmWriteCr2 (Cr2);
+  }
 }
 
 /**
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
index 14b7676c16..5a97733def 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
@@ -1282,28 +1282,6 @@ IsStaticPageTableEnabled (
   )
 ;
 
-/**
-  This function reads CR2 register when on-demand paging is enabled
-  for 64 bit and no action for 32 bit.
-
-  @param[out]  *Cr2  Pointer to variable to hold CR2 register value.
-**/
-VOID
-SaveCr2 (
-  OUT UINTN  *Cr2
-  );
-
-/**
-  This function writes into CR2 register when on-demand paging is enabled
-  for 64 bit and no action for 32 bit.
-
-  @param[in]  Cr2  Value to write into CR2 register.
-**/
-VOID
-RestoreCr2 (
-  IN UINTN  Cr2
-  );
-
 /**
   Schedule a procedure to run on the specified CPU.
 
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
index 18e3f9e08d..8259b01a95 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/X64/PageTbl.c
@@ -1209,32 +1209,3 @@ SetPageTableAttributes (
   return ;
 }
 
-/**
-  This function reads CR2 register when on-demand paging is enabled.
-
-  @param[out]  *Cr2  Pointer to variable to hold CR2 register value.
-**/
-VOID
-SaveCr2 (
-  OUT UINTN  *Cr2
-  )
-{
-  if (!mCpuSmmStaticPageTable) {
-    *Cr2 = AsmReadCr2 ();
-  }
-}
-
-/**
-  This function restores CR2 register when on-demand paging is enabled.
-
-  @param[in]  Cr2  Value to write into CR2 register.
-**/
-VOID
-RestoreCr2 (
-  IN UINTN  Cr2
-  )
-{
-  if (!mCpuSmmStaticPageTable) {
-    AsmWriteCr2 (Cr2);
-  }
-}
-- 
2.21.0.windows.1