1 files changed, 6 insertions, 23 deletions
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
index d67f036aea..766eb1246a 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
@@ -50,22 +50,6 @@ SetPageTableBase (
 }
 
 /**
-  Return page table base.
-
-  @return page table base.
-**/
-UINTN
-GetPageTableBase (
-  VOID
-  )
-{
-  if (mInternalCr3 != 0) {
-    return mInternalCr3;
-  }
-  return (AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64);
-}
-
-/**
   Return length according to page attributes.
 
   @param[in]  PageAttributes   The page attribute of the page entry.
@@ -131,21 +115,20 @@ GetPageTableEntry (
   UINT64                *L3PageTable;
   UINT64                *L4PageTable;
   UINT64                *L5PageTable;
-  IA32_CR4              Cr4;
+  UINTN                 PageTableBase;
   BOOLEAN               Enable5LevelPaging;
 
+  GetPageTable (&PageTableBase, &Enable5LevelPaging);
+
   Index5 = ((UINTN)RShiftU64 (Address, 48)) & PAGING_PAE_INDEX_MASK;
   Index4 = ((UINTN)RShiftU64 (Address, 39)) & PAGING_PAE_INDEX_MASK;
   Index3 = ((UINTN)Address >> 30) & PAGING_PAE_INDEX_MASK;
   Index2 = ((UINTN)Address >> 21) & PAGING_PAE_INDEX_MASK;
   Index1 = ((UINTN)Address >> 12) & PAGING_PAE_INDEX_MASK;
 
-  Cr4.UintN = AsmReadCr4 ();
-  Enable5LevelPaging = (BOOLEAN) (Cr4.Bits.LA57 == 1);
-
   if (sizeof(UINTN) == sizeof(UINT64)) {
     if (Enable5LevelPaging) {
-      L5PageTable = (UINT64 *)GetPageTableBase ();
+      L5PageTable = (UINT64 *)PageTableBase;
       if (L5PageTable[Index5] == 0) {
         *PageAttribute = PageNone;
         return NULL;
@@ -153,7 +136,7 @@ GetPageTableEntry (
 
       L4PageTable = (UINT64 *)(UINTN)(L5PageTable[Index5] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);
     } else {
-      L4PageTable = (UINT64 *)GetPageTableBase ();
+      L4PageTable = (UINT64 *)PageTableBase;
     }
     if (L4PageTable[Index4] == 0) {
       *PageAttribute = PageNone;
@@ -162,7 +145,7 @@ GetPageTableEntry (
 
     L3PageTable = (UINT64 *)(UINTN)(L4PageTable[Index4] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);
   } else {
-    L3PageTable = (UINT64 *)GetPageTableBase ();
+    L3PageTable = (UINT64 *)PageTableBase;
   }
   if (L3PageTable[Index3] == 0) {
     *PageAttribute = PageNone;