4 files changed, 5 insertions, 59 deletions
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
index 06ffc6d..83f5bf9 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/MpService.c
@@ -770,7 +770,7 @@ Gen4GPageTable (
   //
   // Allocate the page table
   //
-  PageTable = AllocatePageTableMemory (ExtraPages + 5 + PagesNeeded);
+  PageTable = AllocatePages (ExtraPages + 5 + PagesNeeded);
   ASSERT (PageTable != NULL);
 
   PageTable = (VOID *)((UINTN)PageTable + EFI_PAGES_TO_SIZE (ExtraPages));
@@ -872,7 +872,7 @@ SetCacheability (
     //
     // Allocate a page from SMRAM
     //
-    NewPageTableAddress = AllocatePageTableMemory (1);
+    NewPageTableAddress = AllocatePages (1);
     ASSERT (NewPageTableAddress != NULL);
 
     NewPageTable = (UINT64 *)NewPageTableAddress;
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c
index fb46551..670a5cf 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.c
@@ -970,9 +970,9 @@ PiCpuSmmEntry (
   //
   BufferPages = EFI_SIZE_TO_PAGES (SIZE_32KB + TileSize * (mMaxNumberOfCpus - 1));
   if ((FamilyId == 4) || (FamilyId == 5)) {
-    Buffer = AllocateAlignedPages (BufferPages, SIZE_32KB);
+    Buffer = AllocateAlignedCodePages (BufferPages, SIZE_32KB);
   } else {
-    Buffer = AllocateAlignedPages (BufferPages, SIZE_4KB);
+    Buffer = AllocateAlignedCodePages (BufferPages, SIZE_4KB);
   }
   ASSERT (Buffer != NULL);
   DEBUG ((EFI_D_INFO, "SMRAM SaveState Buffer (0x%08x, 0x%08x)\n", Buffer, EFI_PAGES_TO_SIZE(BufferPages)));
@@ -1416,35 +1416,6 @@ ConfigSmmCodeAccessCheck (
 }
 
 /**
-  This API provides a way to allocate memory for page table.
-
-  This API can be called more once to allocate memory for page tables.
-
-  Allocates the number of 4KB pages of type EfiRuntimeServicesData and returns a pointer to the
-  allocated buffer.  The buffer returned is aligned on a 4KB boundary.  If Pages is 0, then NULL
-  is returned.  If there is not enough memory remaining to satisfy the request, then NULL is
-  returned.
-
-  @param  Pages                 The number of 4 KB pages to allocate.
-
-  @return A pointer to the allocated buffer or NULL if allocation fails.
-
-**/
-VOID *
-AllocatePageTableMemory (
-  IN UINTN           Pages
-  )
-{
-  VOID  *Buffer;
-
-  Buffer = SmmCpuFeaturesAllocatePageTableMemory (Pages);
-  if (Buffer != NULL) {
-    return Buffer;
-  }
-  return AllocatePages (Pages);
-}
-
-/**
   Perform the remaining tasks.
 
 **/
@@ -1469,8 +1440,6 @@ PerformRemainingTasks (
     //
     ConfigSmmCodeAccessCheck ();
 
-    SmmCpuFeaturesCompleteSmmReadyToLock ();
-
     //
     // Clean SMM ready to lock flag
     //
@@ -1496,8 +1465,6 @@ PerformPreTasks (
     //
     ConfigSmmCodeAccessCheck ();
 
-    SmmCpuFeaturesCompleteSmmReadyToLock ();
-
     mRestoreSmmConfigurationInS3 = FALSE;
   }
 }
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
index f2a9165..66d85d8 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/PiSmmCpuDxeSmm.h
@@ -718,25 +718,4 @@ VOID
 DumpModuleInfoByIp (
   IN  UINTN              CallerIpAddress
   );
-
-/**
-  This API provides a way to allocate memory for page table.
-
-  This API can be called more once to allocate memory for page tables.
-
-  Allocates the number of 4KB pages of type EfiRuntimeServicesData and returns a pointer to the
-  allocated buffer.  The buffer returned is aligned on a 4KB boundary.  If Pages is 0, then NULL
-  is returned.  If there is not enough memory remaining to satisfy the request, then NULL is
-  returned.
-
-  @param  Pages                 The number of 4 KB pages to allocate.
-
-  @return A pointer to the allocated buffer or NULL if allocation fails.
-
-**/
-VOID *
-AllocatePageTableMemory (
-  IN UINTN           Pages
-  );
-
 #endif
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c
index ff4e28e..8ddde9a 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmProfile.c
@@ -552,7 +552,7 @@ InitPaging (
           //
           ASSERT (Address == (*Pte & PHYSICAL_ADDRESS_MASK));
 
-          Pt = AllocatePageTableMemory (1);
+          Pt = AllocatePages (1);
           ASSERT (Pt != NULL);
 
           // Split it