]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Application/MemoryProfileInfo/MemoryProfileInfo.c
IntelSiliconPkg: Replace BSD License with BSD+Patent License
[mirror_edk2.git] / MdeModulePkg / Application / MemoryProfileInfo / MemoryProfileInfo.c
index f69a3e77486a82368f127e42b5280fb2c3608865..072ae72dcd7614295e602e83c2bf0821145c8ceb 100644 (file)
@@ -1,13 +1,13 @@
 /** @file\r
-  \r
-  Copyright (c) 2014 - 2015, Intel Corporation. All rights reserved.<BR>\r
-  This program and the accompanying materials                          \r
-  are licensed and made available under the terms and conditions of the BSD License         \r
-  which accompanies this distribution.  The full text of the license may be found at        \r
-  http://opensource.org/licenses/bsd-license.php                                            \r
 \r
-  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,                     \r
-  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.             \r
+  Copyright (c) 2014 - 2017, Intel Corporation. All rights reserved.<BR>\r
+  This program and the accompanying materials\r
+  are licensed and made available under the terms and conditions of the BSD License\r
+  which accompanies this distribution.  The full text of the license may be found at\r
+  http://opensource.org/licenses/bsd-license.php\r
+\r
+  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,\r
+  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r
 \r
 **/\r
 \r
 #include <Library/UefiRuntimeServicesTableLib.h>\r
 #include <Library/DebugLib.h>\r
 #include <Library/DxeServicesLib.h>\r
-#include <Library/PeCoffGetEntryPointLib.h>\r
 #include <Library/PrintLib.h>\r
 \r
 #include <Protocol/SmmCommunication.h>\r
 #include <Protocol/SmmAccess2.h>\r
 \r
-#include <Guid/ZeroGuid.h>\r
 #include <Guid/MemoryProfile.h>\r
+#include <Guid/PiSmmCommunicationRegionTable.h>\r
+\r
+CHAR8 *mActionString[] = {\r
+  "Unknown",\r
+  "gBS->AllocatePages",\r
+  "gBS->FreePages",\r
+  "gBS->AllocatePool",\r
+  "gBS->FreePool",\r
+};\r
+\r
+CHAR8 *mSmmActionString[] = {\r
+  "SmmUnknown",\r
+  "gSmst->SmmAllocatePages",\r
+  "gSmst->SmmFreePages",\r
+  "gSmst->SmmAllocatePool",\r
+  "gSmst->SmmFreePool",\r
+};\r
 \r
-CHAR16 *mActionString[] = {\r
-  L"Unknown",\r
-  L"AllocatePages",\r
-  L"FreePages",\r
-  L"AllocatePool",\r
-  L"FreePool",\r
+typedef struct {\r
+  MEMORY_PROFILE_ACTION  Action;\r
+  CHAR8                 *String;\r
+} ACTION_STRING;\r
+\r
+ACTION_STRING mExtActionString[] = {\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_PAGES,                    "Lib:AllocatePages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_PAGES,            "Lib:AllocateRuntimePages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_PAGES,           "Lib:AllocateReservedPages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_FREE_PAGES,                        "Lib:FreePages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ALIGNED_PAGES,            "Lib:AllocateAlignedPages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ALIGNED_RUNTIME_PAGES,    "Lib:AllocateAlignedRuntimePages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ALIGNED_RESERVED_PAGES,   "Lib:AllocateAlignedReservedPages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_FREE_ALIGNED_PAGES,                "Lib:FreeAlignedPages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_POOL,                     "Lib:AllocatePool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_POOL,             "Lib:AllocateRuntimePool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_POOL,            "Lib:AllocateReservedPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_FREE_POOL,                         "Lib:FreePool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ZERO_POOL,                "Lib:AllocateZeroPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_ZERO_POOL,        "Lib:AllocateRuntimeZeroPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_ZERO_POOL,       "Lib:AllocateReservedZeroPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_COPY_POOL,                "Lib:AllocateCopyPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_COPY_POOL,        "Lib:AllocateRuntimeCopyPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_COPY_POOL,       "Lib:AllocateReservedCopyPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_REALLOCATE_POOL,                   "Lib:ReallocatePool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_REALLOCATE_RUNTIME_POOL,           "Lib:ReallocateRuntimePool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_REALLOCATE_RESERVED_POOL,          "Lib:ReallocateReservedPool"},\r
 };\r
 \r
-CHAR16 *mMemoryTypeString[] = {\r
-  L"EfiReservedMemoryType",\r
-  L"EfiLoaderCode",\r
-  L"EfiLoaderData",\r
-  L"EfiBootServicesCode",\r
-  L"EfiBootServicesData",\r
-  L"EfiRuntimeServicesCode",\r
-  L"EfiRuntimeServicesData",\r
-  L"EfiConventionalMemory",\r
-  L"EfiUnusableMemory",\r
-  L"EfiACPIReclaimMemory",\r
-  L"EfiACPIMemoryNVS",\r
-  L"EfiMemoryMappedIO",\r
-  L"EfiMemoryMappedIOPortSpace",\r
-  L"EfiPalCode",\r
-  L"EfiPersistentMemory",\r
-  L"EfiOSReserved",\r
-  L"EfiOemReserved",\r
+CHAR8 mUserDefinedActionString[] = {"UserDefined-0x80000000"};\r
+\r
+CHAR8 *mMemoryTypeString[] = {\r
+  "EfiReservedMemoryType",\r
+  "EfiLoaderCode",\r
+  "EfiLoaderData",\r
+  "EfiBootServicesCode",\r
+  "EfiBootServicesData",\r
+  "EfiRuntimeServicesCode",\r
+  "EfiRuntimeServicesData",\r
+  "EfiConventionalMemory",\r
+  "EfiUnusableMemory",\r
+  "EfiACPIReclaimMemory",\r
+  "EfiACPIMemoryNVS",\r
+  "EfiMemoryMappedIO",\r
+  "EfiMemoryMappedIOPortSpace",\r
+  "EfiPalCode",\r
+  "EfiPersistentMemory",\r
+  "EfiOSReserved",\r
+  "EfiOemReserved",\r
 };\r
 \r
-CHAR16 *mSubsystemString[] = {\r
-  L"Unknown",\r
-  L"NATIVE",\r
-  L"WINDOWS_GUI",\r
-  L"WINDOWS_CUI",\r
-  L"Unknown",\r
-  L"Unknown",\r
-  L"Unknown",\r
-  L"POSIX_CUI",\r
-  L"Unknown",\r
-  L"WINDOWS_CE_GUI",\r
-  L"EFI_APPLICATION",\r
-  L"EFI_BOOT_SERVICE_DRIVER",\r
-  L"EFI_RUNTIME_DRIVER",\r
-  L"EFI_ROM",\r
-  L"XBOX",\r
-  L"Unknown",\r
+CHAR8 *mSubsystemString[] = {\r
+  "Unknown",\r
+  "NATIVE",\r
+  "WINDOWS_GUI",\r
+  "WINDOWS_CUI",\r
+  "Unknown",\r
+  "Unknown",\r
+  "Unknown",\r
+  "POSIX_CUI",\r
+  "Unknown",\r
+  "WINDOWS_CE_GUI",\r
+  "EFI_APPLICATION",\r
+  "EFI_BOOT_SERVICE_DRIVER",\r
+  "EFI_RUNTIME_DRIVER",\r
+  "EFI_ROM",\r
+  "XBOX",\r
+  "Unknown",\r
 };\r
 \r
-CHAR16 *mFileTypeString[] = {\r
-  L"Unknown",\r
-  L"RAW",\r
-  L"FREEFORM",\r
-  L"SECURITY_CORE",\r
-  L"PEI_CORE",\r
-  L"DXE_CORE",\r
-  L"PEIM",\r
-  L"DRIVER",\r
-  L"COMBINED_PEIM_DRIVER",\r
-  L"APPLICATION",\r
-  L"SMM",\r
-  L"FIRMWARE_VOLUME_IMAGE",\r
-  L"COMBINED_SMM_DXE",\r
-  L"SMM_CORE",\r
+CHAR8 *mFileTypeString[] = {\r
+  "Unknown",\r
+  "RAW",\r
+  "FREEFORM",\r
+  "SECURITY_CORE",\r
+  "PEI_CORE",\r
+  "DXE_CORE",\r
+  "PEIM",\r
+  "DRIVER",\r
+  "COMBINED_PEIM_DRIVER",\r
+  "APPLICATION",\r
+  "SMM",\r
+  "FIRMWARE_VOLUME_IMAGE",\r
+  "COMBINED_SMM_DXE",\r
+  "SMM_CORE",\r
 };\r
 \r
-#define PROFILE_NAME_STRING_LENGTH  36\r
-CHAR16 mNameString[PROFILE_NAME_STRING_LENGTH + 1];\r
+#define PROFILE_NAME_STRING_LENGTH  64\r
+CHAR8 mNameString[PROFILE_NAME_STRING_LENGTH + 1];\r
+\r
+//\r
+// Profile summary information\r
+//\r
+#define MEMORY_PROFILE_ALLOC_SUMMARY_INFO_SIGNATURE SIGNATURE_32 ('M','P','A','S')\r
+#define MEMORY_PROFILE_ALLOC_SUMMARY_INFO_REVISION 0x0001\r
+\r
+typedef struct {\r
+  MEMORY_PROFILE_COMMON_HEADER  Header;\r
+  PHYSICAL_ADDRESS              CallerAddress;\r
+  MEMORY_PROFILE_ACTION         Action;\r
+  CHAR8                         *ActionString;\r
+  UINT32                        AllocateCount;\r
+  UINT64                        TotalSize;\r
+} MEMORY_PROFILE_ALLOC_SUMMARY_INFO;\r
+\r
+typedef struct {\r
+  UINT32                            Signature;\r
+  MEMORY_PROFILE_ALLOC_SUMMARY_INFO AllocSummaryInfo;\r
+  LIST_ENTRY                        Link;\r
+} MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA;\r
+\r
+typedef struct {\r
+  UINT32                        Signature;\r
+  MEMORY_PROFILE_DRIVER_INFO    *DriverInfo;\r
+  LIST_ENTRY                    *AllocSummaryInfoList;\r
+  LIST_ENTRY                    Link;\r
+} MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA;\r
+\r
+typedef struct {\r
+  UINT32                        Signature;\r
+  MEMORY_PROFILE_CONTEXT        *Context;\r
+  LIST_ENTRY                    *DriverSummaryInfoList;\r
+} MEMORY_PROFILE_CONTEXT_SUMMARY_DATA;\r
+\r
+LIST_ENTRY  mImageSummaryQueue = INITIALIZE_LIST_HEAD_VARIABLE (mImageSummaryQueue);\r
+MEMORY_PROFILE_CONTEXT_SUMMARY_DATA mMemoryProfileContextSummary;\r
 \r
-/** \r
+/**\r
   Get the file name portion of the Pdb File Name.\r
-  \r
+\r
   The portion of the Pdb File Name between the last backslash and\r
-  either a following period or the end of the string is converted\r
-  to Unicode and copied into UnicodeBuffer.  The name is truncated,\r
-  if necessary, to ensure that UnicodeBuffer is not overrun.\r
-  \r
+  either a following period or the end of the string is copied into\r
+  AsciiBuffer.  The name is truncated, if necessary, to ensure that\r
+  AsciiBuffer is not overrun.\r
+\r
   @param[in]  PdbFileName     Pdb file name.\r
-  @param[out] UnicodeBuffer   The resultant Unicode File Name.\r
-  \r
+  @param[out] AsciiBuffer     The resultant Ascii File Name.\r
+\r
 **/\r
 VOID\r
 GetShortPdbFileName (\r
   IN  CHAR8     *PdbFileName,\r
-  OUT CHAR16    *UnicodeBuffer\r
+  OUT CHAR8     *AsciiBuffer\r
   )\r
 {\r
-  UINTN IndexA;     // Current work location within an ASCII string.\r
-  UINTN IndexU;     // Current work location within a Unicode string.\r
+  UINTN IndexPdb;     // Current work location within a Pdb string.\r
+  UINTN IndexBuffer;  // Current work location within a Buffer string.\r
   UINTN StartIndex;\r
   UINTN EndIndex;\r
 \r
-  ZeroMem (UnicodeBuffer, (PROFILE_NAME_STRING_LENGTH + 1) * sizeof (CHAR16));\r
+  ZeroMem (AsciiBuffer, PROFILE_NAME_STRING_LENGTH + 1);\r
 \r
   if (PdbFileName == NULL) {\r
-    StrnCpyS (UnicodeBuffer, PROFILE_NAME_STRING_LENGTH + 1, L" ", 1);\r
+    AsciiStrnCpyS (AsciiBuffer, PROFILE_NAME_STRING_LENGTH + 1, " ", 1);\r
   } else {\r
     StartIndex = 0;\r
     for (EndIndex = 0; PdbFileName[EndIndex] != 0; EndIndex++);\r
-    for (IndexA = 0; PdbFileName[IndexA] != 0; IndexA++) {\r
-      if (PdbFileName[IndexA] == '\\') {\r
-        StartIndex = IndexA + 1;\r
+    for (IndexPdb = 0; PdbFileName[IndexPdb] != 0; IndexPdb++) {\r
+      if ((PdbFileName[IndexPdb] == '\\') || (PdbFileName[IndexPdb] == '/')) {\r
+        StartIndex = IndexPdb + 1;\r
       }\r
 \r
-      if (PdbFileName[IndexA] == '.') {\r
-        EndIndex = IndexA;\r
+      if (PdbFileName[IndexPdb] == '.') {\r
+        EndIndex = IndexPdb;\r
       }\r
     }\r
 \r
-    IndexU = 0;\r
-    for (IndexA = StartIndex; IndexA < EndIndex; IndexA++) {\r
-      UnicodeBuffer[IndexU] = (CHAR16) PdbFileName[IndexA];\r
-      IndexU++;\r
-      if (IndexU >= PROFILE_NAME_STRING_LENGTH) {\r
-        UnicodeBuffer[PROFILE_NAME_STRING_LENGTH] = 0;\r
+    IndexBuffer = 0;\r
+    for (IndexPdb = StartIndex; IndexPdb < EndIndex; IndexPdb++) {\r
+      AsciiBuffer[IndexBuffer] = PdbFileName[IndexPdb];\r
+      IndexBuffer++;\r
+      if (IndexBuffer >= PROFILE_NAME_STRING_LENGTH) {\r
+        AsciiBuffer[PROFILE_NAME_STRING_LENGTH] = 0;\r
         break;\r
       }\r
     }\r
   }\r
 }\r
 \r
-/** \r
+/**\r
   Get a human readable name for an image.\r
   The following methods will be tried orderly:\r
     1. Image PDB\r
@@ -159,32 +234,27 @@ GetShortPdbFileName (
 \r
   @param[in] DriverInfo Pointer to memory profile driver info.\r
 \r
-  @post The resulting Unicode name string is stored in the mNameString global array.\r
+  @return The resulting Ascii name string is stored in the mNameString global array.\r
 \r
 **/\r
-VOID\r
+CHAR8 *\r
 GetDriverNameString (\r
  IN MEMORY_PROFILE_DRIVER_INFO  *DriverInfo\r
  )\r
 {\r
   EFI_STATUS                  Status;\r
-  CHAR8                       *PdbFileName;\r
   CHAR16                      *NameString;\r
   UINTN                       StringSize;\r
 \r
   //\r
   // Method 1: Get the name string from image PDB\r
   //\r
-  if ((DriverInfo->ImageBase != 0) && (DriverInfo->FileType != EFI_FV_FILETYPE_SMM) && (DriverInfo->FileType != EFI_FV_FILETYPE_SMM_CORE)) {\r
-    PdbFileName = PeCoffLoaderGetPdbPointer ((VOID *) (UINTN) DriverInfo->ImageBase);\r
-\r
-    if (PdbFileName != NULL) {\r
-      GetShortPdbFileName (PdbFileName, mNameString);\r
-      return;\r
-    }\r
+  if (DriverInfo->PdbStringOffset != 0) {\r
+    GetShortPdbFileName ((CHAR8 *) ((UINTN) DriverInfo + DriverInfo->PdbStringOffset), mNameString);\r
+    return mNameString;\r
   }\r
 \r
-  if (!CompareGuid (&DriverInfo->FileName, &gZeroGuid)) {\r
+  if (!IsZeroGuid (&DriverInfo->FileName)) {\r
     //\r
     // Try to get the image's FFS UI section by image GUID\r
     //\r
@@ -201,17 +271,20 @@ GetDriverNameString (
       //\r
       // Method 2: Get the name string from FFS UI section\r
       //\r
-      StrnCpyS (mNameString, PROFILE_NAME_STRING_LENGTH + 1, NameString, PROFILE_NAME_STRING_LENGTH);\r
-      mNameString[PROFILE_NAME_STRING_LENGTH] = 0;\r
+      if (StrLen (NameString) > PROFILE_NAME_STRING_LENGTH) {\r
+        NameString[PROFILE_NAME_STRING_LENGTH] = 0;\r
+      }\r
+      UnicodeStrToAsciiStrS (NameString, mNameString, sizeof (mNameString));\r
       FreePool (NameString);\r
-      return;\r
+      return mNameString;\r
     }\r
   }\r
 \r
   //\r
   // Method 3: Get the name string from image GUID\r
   //\r
-  UnicodeSPrint (mNameString, sizeof (mNameString), L"%g", &DriverInfo->FileName);\r
+  AsciiSPrint (mNameString, sizeof (mNameString), "%g", &DriverInfo->FileName);\r
+  return mNameString;\r
 }\r
 \r
 /**\r
@@ -222,7 +295,7 @@ GetDriverNameString (
   @return Pointer to string.\r
 \r
 **/\r
-CHAR16 *\r
+CHAR8 *\r
 ProfileMemoryTypeToStr (\r
   IN EFI_MEMORY_TYPE    MemoryType\r
   )\r
@@ -246,12 +319,63 @@ ProfileMemoryTypeToStr (
   return mMemoryTypeString[Index];\r
 }\r
 \r
+/**\r
+  Action to string.\r
+\r
+  @param[in] Action                     Profile action.\r
+  @param[in] UserDefinedActionString    Pointer to user defined action string.\r
+  @param[in] IsForSmm                   TRUE  - SMRAM profile.\r
+                                        FALSE - UEFI memory profile.\r
+\r
+  @return Pointer to string.\r
+\r
+**/\r
+CHAR8 *\r
+ProfileActionToStr (\r
+  IN MEMORY_PROFILE_ACTION  Action,\r
+  IN CHAR8                  *UserDefinedActionString,\r
+  IN BOOLEAN                IsForSmm\r
+  )\r
+{\r
+  UINTN     Index;\r
+  UINTN     ActionStringCount;\r
+  CHAR8     **ActionString;\r
+\r
+  if (IsForSmm) {\r
+    ActionString = mSmmActionString;\r
+    ActionStringCount = ARRAY_SIZE (mSmmActionString);\r
+  } else {\r
+    ActionString = mActionString;\r
+    ActionStringCount = ARRAY_SIZE (mActionString);\r
+  }\r
+\r
+  if ((UINTN) (UINT32) Action < ActionStringCount) {\r
+    return ActionString[Action];\r
+  }\r
+  for (Index = 0; Index < ARRAY_SIZE (mExtActionString); Index++) {\r
+    if (mExtActionString[Index].Action == Action) {\r
+      return mExtActionString[Index].String;\r
+    }\r
+  }\r
+  if ((Action & MEMORY_PROFILE_ACTION_USER_DEFINED_MASK) != 0) {\r
+    if (UserDefinedActionString != NULL) {\r
+      return UserDefinedActionString;\r
+    }\r
+    AsciiSPrint (mUserDefinedActionString, sizeof (mUserDefinedActionString), "UserDefined-0x%08x", Action);\r
+    return mUserDefinedActionString;\r
+  }\r
+\r
+  return ActionString[0];\r
+}\r
+\r
 /**\r
   Dump memory profile allocate information.\r
 \r
   @param[in] DriverInfo         Pointer to memory profile driver info.\r
   @param[in] AllocIndex         Memory profile alloc info index.\r
   @param[in] AllocInfo          Pointer to memory profile alloc info.\r
+  @param[in] IsForSmm           TRUE  - SMRAM profile.\r
+                                FALSE - UEFI memory profile.\r
 \r
   @return Pointer to next memory profile alloc info.\r
 \r
@@ -260,20 +384,30 @@ MEMORY_PROFILE_ALLOC_INFO *
 DumpMemoryProfileAllocInfo (\r
   IN MEMORY_PROFILE_DRIVER_INFO *DriverInfo,\r
   IN UINTN                      AllocIndex,\r
-  IN MEMORY_PROFILE_ALLOC_INFO  *AllocInfo\r
+  IN MEMORY_PROFILE_ALLOC_INFO  *AllocInfo,\r
+  IN BOOLEAN                    IsForSmm\r
   )\r
 {\r
+  CHAR8     *ActionString;\r
+\r
   if (AllocInfo->Header.Signature != MEMORY_PROFILE_ALLOC_INFO_SIGNATURE) {\r
     return NULL;\r
   }\r
+\r
+  if (AllocInfo->ActionStringOffset != 0) {\r
+    ActionString = (CHAR8 *) ((UINTN) AllocInfo + AllocInfo->ActionStringOffset);\r
+  } else {\r
+    ActionString = NULL;\r
+  }\r
+\r
   Print (L"    MEMORY_PROFILE_ALLOC_INFO (0x%x)\n", AllocIndex);\r
   Print (L"      Signature     - 0x%08x\n", AllocInfo->Header.Signature);\r
   Print (L"      Length        - 0x%04x\n", AllocInfo->Header.Length);\r
-  Print (L"      Revision      - 0x%04x\n", AllocInfo->Header.Revision);  \r
+  Print (L"      Revision      - 0x%04x\n", AllocInfo->Header.Revision);\r
   Print (L"      CallerAddress - 0x%016lx (Offset: 0x%08x)\n", AllocInfo->CallerAddress, (UINTN) (AllocInfo->CallerAddress - DriverInfo->ImageBase));\r
   Print (L"      SequenceId    - 0x%08x\n", AllocInfo->SequenceId);\r
-  Print (L"      Action        - 0x%08x (%s)\n", AllocInfo->Action, mActionString[(AllocInfo->Action < sizeof(mActionString)/sizeof(mActionString[0])) ? AllocInfo->Action : 0]);\r
-  Print (L"      MemoryType    - 0x%08x (%s)\n", AllocInfo->MemoryType, ProfileMemoryTypeToStr (AllocInfo->MemoryType));\r
+  Print (L"      Action        - 0x%08x (%a)\n", AllocInfo->Action, ProfileActionToStr (AllocInfo->Action, ActionString, IsForSmm));\r
+  Print (L"      MemoryType    - 0x%08x (%a)\n", AllocInfo->MemoryType, ProfileMemoryTypeToStr (AllocInfo->MemoryType));\r
   Print (L"      Buffer        - 0x%016lx\n", AllocInfo->Buffer);\r
   Print (L"      Size          - 0x%016lx\n", AllocInfo->Size);\r
 \r
@@ -285,6 +419,8 @@ DumpMemoryProfileAllocInfo (
 \r
   @param[in] DriverIndex        Memory profile driver info index.\r
   @param[in] DriverInfo         Pointer to memory profile driver info.\r
+  @param[in] IsForSmm           TRUE  - SMRAM profile.\r
+                                FALSE - UEFI memory profile.\r
 \r
   @return Pointer to next memory profile driver info.\r
 \r
@@ -292,12 +428,14 @@ DumpMemoryProfileAllocInfo (
 MEMORY_PROFILE_DRIVER_INFO *\r
 DumpMemoryProfileDriverInfo (\r
   IN UINTN                      DriverIndex,\r
-  IN MEMORY_PROFILE_DRIVER_INFO *DriverInfo\r
+  IN MEMORY_PROFILE_DRIVER_INFO *DriverInfo,\r
+  IN BOOLEAN                    IsForSmm\r
   )\r
 {\r
   UINTN                         TypeIndex;\r
   MEMORY_PROFILE_ALLOC_INFO     *AllocInfo;\r
   UINTN                         AllocIndex;\r
+  CHAR8                         *NameString;\r
 \r
   if (DriverInfo->Header.Signature != MEMORY_PROFILE_DRIVER_INFO_SIGNATURE) {\r
     return NULL;\r
@@ -305,28 +443,31 @@ DumpMemoryProfileDriverInfo (
   Print (L"  MEMORY_PROFILE_DRIVER_INFO (0x%x)\n", DriverIndex);\r
   Print (L"    Signature               - 0x%08x\n", DriverInfo->Header.Signature);\r
   Print (L"    Length                  - 0x%04x\n", DriverInfo->Header.Length);\r
-  Print (L"    Revision                - 0x%04x\n", DriverInfo->Header.Revision);  \r
-  GetDriverNameString (DriverInfo);\r
-  Print (L"    FileName                - %s\n", &mNameString);\r
+  Print (L"    Revision                - 0x%04x\n", DriverInfo->Header.Revision);\r
+  NameString = GetDriverNameString (DriverInfo);\r
+  Print (L"    FileName                - %a\n", NameString);\r
+  if (DriverInfo->PdbStringOffset != 0) {\r
+    Print (L"    Pdb                     - %a\n", (CHAR8 *) ((UINTN) DriverInfo + DriverInfo->PdbStringOffset));\r
+  }\r
   Print (L"    ImageBase               - 0x%016lx\n", DriverInfo->ImageBase);\r
   Print (L"    ImageSize               - 0x%016lx\n", DriverInfo->ImageSize);\r
   Print (L"    EntryPoint              - 0x%016lx\n", DriverInfo->EntryPoint);\r
-  Print (L"    ImageSubsystem          - 0x%04x (%s)\n", DriverInfo->ImageSubsystem, mSubsystemString[(DriverInfo->ImageSubsystem < sizeof(mSubsystemString)/sizeof(mSubsystemString[0])) ? DriverInfo->ImageSubsystem : 0]);\r
-  Print (L"    FileType                - 0x%02x (%s)\n", DriverInfo->FileType, mFileTypeString[(DriverInfo->FileType < sizeof(mFileTypeString)/sizeof(mFileTypeString[0])) ? DriverInfo->FileType : 0]);\r
+  Print (L"    ImageSubsystem          - 0x%04x (%a)\n", DriverInfo->ImageSubsystem, mSubsystemString[(DriverInfo->ImageSubsystem < sizeof(mSubsystemString)/sizeof(mSubsystemString[0])) ? DriverInfo->ImageSubsystem : 0]);\r
+  Print (L"    FileType                - 0x%02x (%a)\n", DriverInfo->FileType, mFileTypeString[(DriverInfo->FileType < sizeof(mFileTypeString)/sizeof(mFileTypeString[0])) ? DriverInfo->FileType : 0]);\r
   Print (L"    CurrentUsage            - 0x%016lx\n", DriverInfo->CurrentUsage);\r
   Print (L"    PeakUsage               - 0x%016lx\n", DriverInfo->PeakUsage);\r
   for (TypeIndex = 0; TypeIndex < sizeof (DriverInfo->CurrentUsageByType) / sizeof (DriverInfo->CurrentUsageByType[0]); TypeIndex++) {\r
     if ((DriverInfo->CurrentUsageByType[TypeIndex] != 0) ||\r
         (DriverInfo->PeakUsageByType[TypeIndex] != 0)) {\r
-      Print (L"    CurrentUsage[0x%02x]      - 0x%016lx (%s)\n", TypeIndex, DriverInfo->CurrentUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]);\r
-      Print (L"    PeakUsage[0x%02x]         - 0x%016lx (%s)\n", TypeIndex, DriverInfo->PeakUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]);\r
+      Print (L"    CurrentUsage[0x%02x]      - 0x%016lx (%a)\n", TypeIndex, DriverInfo->CurrentUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]);\r
+      Print (L"    PeakUsage[0x%02x]         - 0x%016lx (%a)\n", TypeIndex, DriverInfo->PeakUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]);\r
     }\r
   }\r
   Print (L"    AllocRecordCount        - 0x%08x\n", DriverInfo->AllocRecordCount);\r
 \r
   AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) ((UINTN) DriverInfo + DriverInfo->Header.Length);\r
   for (AllocIndex = 0; AllocIndex < DriverInfo->AllocRecordCount; AllocIndex++) {\r
-    AllocInfo = DumpMemoryProfileAllocInfo (DriverInfo, AllocIndex, AllocInfo);\r
+    AllocInfo = DumpMemoryProfileAllocInfo (DriverInfo, AllocIndex, AllocInfo, IsForSmm);\r
     if (AllocInfo == NULL) {\r
       return NULL;\r
     }\r
@@ -338,13 +479,16 @@ DumpMemoryProfileDriverInfo (
   Dump memory profile context information.\r
 \r
   @param[in] Context            Pointer to memory profile context.\r
+  @param[in] IsForSmm           TRUE  - SMRAM profile.\r
+                                FALSE - UEFI memory profile.\r
 \r
   @return Pointer to the end of memory profile context buffer.\r
 \r
 **/\r
 VOID *\r
 DumpMemoryProfileContext (\r
-  IN MEMORY_PROFILE_CONTEXT     *Context\r
+  IN MEMORY_PROFILE_CONTEXT     *Context,\r
+  IN BOOLEAN                    IsForSmm\r
   )\r
 {\r
   UINTN                         TypeIndex;\r
@@ -357,14 +501,14 @@ DumpMemoryProfileContext (
   Print (L"MEMORY_PROFILE_CONTEXT\n");\r
   Print (L"  Signature                     - 0x%08x\n", Context->Header.Signature);\r
   Print (L"  Length                        - 0x%04x\n", Context->Header.Length);\r
-  Print (L"  Revision                      - 0x%04x\n", Context->Header.Revision);  \r
+  Print (L"  Revision                      - 0x%04x\n", Context->Header.Revision);\r
   Print (L"  CurrentTotalUsage             - 0x%016lx\n", Context->CurrentTotalUsage);\r
   Print (L"  PeakTotalUsage                - 0x%016lx\n", Context->PeakTotalUsage);\r
   for (TypeIndex = 0; TypeIndex < sizeof (Context->CurrentTotalUsageByType) / sizeof (Context->CurrentTotalUsageByType[0]); TypeIndex++) {\r
     if ((Context->CurrentTotalUsageByType[TypeIndex] != 0) ||\r
         (Context->PeakTotalUsageByType[TypeIndex] != 0)) {\r
-      Print (L"  CurrentTotalUsage[0x%02x]       - 0x%016lx (%s)\n", TypeIndex, Context->CurrentTotalUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]);\r
-      Print (L"  PeakTotalUsage[0x%02x]          - 0x%016lx (%s)\n", TypeIndex, Context->PeakTotalUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]);\r
+      Print (L"  CurrentTotalUsage[0x%02x]       - 0x%016lx (%a)\n", TypeIndex, Context->CurrentTotalUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]);\r
+      Print (L"  PeakTotalUsage[0x%02x]          - 0x%016lx (%a)\n", TypeIndex, Context->PeakTotalUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]);\r
     }\r
   }\r
   Print (L"  TotalImageSize                - 0x%016lx\n", Context->TotalImageSize);\r
@@ -373,7 +517,7 @@ DumpMemoryProfileContext (
 \r
   DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) ((UINTN) Context + Context->Header.Length);\r
   for (DriverIndex = 0; DriverIndex < Context->ImageCount; DriverIndex++) {\r
-    DriverInfo = DumpMemoryProfileDriverInfo (DriverIndex, DriverInfo);\r
+    DriverInfo = DumpMemoryProfileDriverInfo (DriverIndex, DriverInfo, IsForSmm);\r
     if (DriverInfo == NULL) {\r
       return NULL;\r
     }\r
@@ -402,7 +546,7 @@ DumpMemoryProfileDescriptor (
   Print (L"  MEMORY_PROFILE_DESCRIPTOR (0x%x)\n", DescriptorIndex);\r
   Print (L"    Signature               - 0x%08x\n", Descriptor->Header.Signature);\r
   Print (L"    Length                  - 0x%04x\n", Descriptor->Header.Length);\r
-  Print (L"    Revision                - 0x%04x\n", Descriptor->Header.Revision);  \r
+  Print (L"    Revision                - 0x%04x\n", Descriptor->Header.Revision);\r
   Print (L"    Address                 - 0x%016lx\n", Descriptor->Address);\r
   Print (L"    Size                    - 0x%016lx\n", Descriptor->Size);\r
 \r
@@ -431,7 +575,7 @@ DumpMemoryProfileFreeMemory (
   Print (L"MEMORY_PROFILE_FREE_MEMORY\n");\r
   Print (L"  Signature                     - 0x%08x\n", FreeMemory->Header.Signature);\r
   Print (L"  Length                        - 0x%04x\n", FreeMemory->Header.Length);\r
-  Print (L"  Revision                      - 0x%04x\n", FreeMemory->Header.Revision);  \r
+  Print (L"  Revision                      - 0x%04x\n", FreeMemory->Header.Revision);\r
   Print (L"  TotalFreeMemoryPages          - 0x%016lx\n", FreeMemory->TotalFreeMemoryPages);\r
   Print (L"  FreeMemoryEntryCount          - 0x%08x\n", FreeMemory->FreeMemoryEntryCount);\r
 \r
@@ -468,7 +612,7 @@ DumpMemoryProfileMemoryRange (
   Print (L"MEMORY_PROFILE_MEMORY_RANGE\n");\r
   Print (L"  Signature                     - 0x%08x\n", MemoryRange->Header.Signature);\r
   Print (L"  Length                        - 0x%04x\n", MemoryRange->Header.Length);\r
-  Print (L"  Revision                      - 0x%04x\n", MemoryRange->Header.Revision);  \r
+  Print (L"  Revision                      - 0x%04x\n", MemoryRange->Header.Revision);\r
   Print (L"  MemoryRangeCount              - 0x%08x\n", MemoryRange->MemoryRangeCount);\r
 \r
   Descriptor = (MEMORY_PROFILE_DESCRIPTOR *) ((UINTN) MemoryRange + MemoryRange->Header.Length);\r
@@ -511,6 +655,10 @@ ScanMemoryProfileBySignature (
       //\r
       return (VOID *) CommonHeader;\r
     }\r
+    if (CommonHeader->Length == 0) {\r
+      ASSERT (FALSE);\r
+      return NULL;\r
+    }\r
     CommonHeader = (MEMORY_PROFILE_COMMON_HEADER *) ((UINTN) CommonHeader + CommonHeader->Length);\r
   }\r
 \r
@@ -522,12 +670,15 @@ ScanMemoryProfileBySignature (
 \r
   @param[in] ProfileBuffer      Memory profile base address.\r
   @param[in] ProfileSize        Memory profile size.\r
+  @param[in] IsForSmm           TRUE  - SMRAM profile.\r
+                                FALSE - UEFI memory profile.\r
 \r
 **/\r
 VOID\r
 DumpMemoryProfile (\r
   IN PHYSICAL_ADDRESS           ProfileBuffer,\r
-  IN UINT64                     ProfileSize\r
+  IN UINT64                     ProfileSize,\r
+  IN BOOLEAN                    IsForSmm\r
   )\r
 {\r
   MEMORY_PROFILE_CONTEXT        *Context;\r
@@ -536,7 +687,7 @@ DumpMemoryProfile (
 \r
   Context = (MEMORY_PROFILE_CONTEXT *) ScanMemoryProfileBySignature (ProfileBuffer, ProfileSize, MEMORY_PROFILE_CONTEXT_SIGNATURE);\r
   if (Context != NULL) {\r
-    DumpMemoryProfileContext (Context);\r
+    DumpMemoryProfileContext (Context, IsForSmm);\r
   }\r
 \r
   FreeMemory = (MEMORY_PROFILE_FREE_MEMORY *) ScanMemoryProfileBySignature (ProfileBuffer, ProfileSize, MEMORY_PROFILE_FREE_MEMORY_SIGNATURE);\r
@@ -550,6 +701,317 @@ DumpMemoryProfile (
   }\r
 }\r
 \r
+/**\r
+  Get Allocate summary information structure by caller address.\r
+\r
+  @param[in] CallerAddress          Caller address.\r
+  @param[in] DriverSummaryInfoData  Driver summary information data structure.\r
+\r
+  @return Allocate summary information structure by caller address.\r
+\r
+**/\r
+MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA *\r
+GetAllocSummaryInfoByCallerAddress (\r
+  IN PHYSICAL_ADDRESS                           CallerAddress,\r
+  IN MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA    *DriverSummaryInfoData\r
+  )\r
+{\r
+  LIST_ENTRY                                    *AllocSummaryInfoList;\r
+  LIST_ENTRY                                    *AllocSummaryLink;\r
+  MEMORY_PROFILE_ALLOC_SUMMARY_INFO             *AllocSummaryInfo;\r
+  MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA        *AllocSummaryInfoData;\r
+\r
+  AllocSummaryInfoList = DriverSummaryInfoData->AllocSummaryInfoList;\r
+\r
+  for (AllocSummaryLink = AllocSummaryInfoList->ForwardLink;\r
+       AllocSummaryLink != AllocSummaryInfoList;\r
+       AllocSummaryLink = AllocSummaryLink->ForwardLink) {\r
+    AllocSummaryInfoData = CR (\r
+                             AllocSummaryLink,\r
+                             MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA,\r
+                             Link,\r
+                             MEMORY_PROFILE_ALLOC_SUMMARY_INFO_SIGNATURE\r
+                             );\r
+    AllocSummaryInfo = &AllocSummaryInfoData->AllocSummaryInfo;\r
+    if (AllocSummaryInfo->CallerAddress == CallerAddress) {\r
+      return AllocSummaryInfoData;\r
+    }\r
+  }\r
+  return NULL;\r
+}\r
+\r
+/**\r
+  Create Allocate summary information structure and\r
+  link to Driver summary information data structure.\r
+\r
+  @param[in, out] DriverSummaryInfoData Driver summary information data structure.\r
+  @param[in]      AllocInfo             Pointer to memory profile alloc info.\r
+\r
+  @return Pointer to next memory profile alloc info.\r
+\r
+**/\r
+MEMORY_PROFILE_ALLOC_INFO *\r
+CreateAllocSummaryInfo (\r
+  IN OUT MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA    *DriverSummaryInfoData,\r
+  IN MEMORY_PROFILE_ALLOC_INFO                      *AllocInfo\r
+  )\r
+{\r
+  MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA            *AllocSummaryInfoData;\r
+  MEMORY_PROFILE_ALLOC_SUMMARY_INFO                 *AllocSummaryInfo;\r
+\r
+  if (AllocInfo->Header.Signature != MEMORY_PROFILE_ALLOC_INFO_SIGNATURE) {\r
+    return NULL;\r
+  }\r
+\r
+  AllocSummaryInfoData = GetAllocSummaryInfoByCallerAddress (AllocInfo->CallerAddress, DriverSummaryInfoData);\r
+  if (AllocSummaryInfoData == NULL) {\r
+    AllocSummaryInfoData = AllocatePool (sizeof (*AllocSummaryInfoData));\r
+    if (AllocSummaryInfoData == NULL) {\r
+      return NULL;\r
+    }\r
+\r
+    AllocSummaryInfoData->Signature = MEMORY_PROFILE_ALLOC_SUMMARY_INFO_SIGNATURE;\r
+    AllocSummaryInfo = &AllocSummaryInfoData->AllocSummaryInfo;\r
+    AllocSummaryInfo->Header.Signature = MEMORY_PROFILE_ALLOC_SUMMARY_INFO_SIGNATURE;\r
+    AllocSummaryInfo->Header.Length    = sizeof (*AllocSummaryInfo);\r
+    AllocSummaryInfo->Header.Revision  = MEMORY_PROFILE_ALLOC_SUMMARY_INFO_REVISION;\r
+    AllocSummaryInfo->CallerAddress = AllocInfo->CallerAddress;\r
+    AllocSummaryInfo->Action        = AllocInfo->Action;\r
+    if (AllocInfo->ActionStringOffset != 0) {\r
+      AllocSummaryInfo->ActionString = (CHAR8 *) ((UINTN) AllocInfo + AllocInfo->ActionStringOffset);\r
+    } else {\r
+      AllocSummaryInfo->ActionString  = NULL;\r
+    }\r
+    AllocSummaryInfo->AllocateCount = 0;\r
+    AllocSummaryInfo->TotalSize     = 0;\r
+    InsertTailList (DriverSummaryInfoData->AllocSummaryInfoList, &AllocSummaryInfoData->Link);\r
+  }\r
+  AllocSummaryInfo = &AllocSummaryInfoData->AllocSummaryInfo;\r
+  AllocSummaryInfo->AllocateCount ++;\r
+  AllocSummaryInfo->TotalSize += AllocInfo->Size;\r
+\r
+  return (MEMORY_PROFILE_ALLOC_INFO *) ((UINTN) AllocInfo + AllocInfo->Header.Length);\r
+}\r
+\r
+/**\r
+  Create Driver summary information structure and\r
+  link to Context summary information data structure.\r
+\r
+  @param[in, out] ContextSummaryData    Context summary information data structure.\r
+  @param[in]      DriverInfo            Pointer to memory profile driver info.\r
+\r
+  @return Pointer to next memory profile driver info.\r
+\r
+**/\r
+MEMORY_PROFILE_DRIVER_INFO *\r
+CreateDriverSummaryInfo (\r
+  IN OUT MEMORY_PROFILE_CONTEXT_SUMMARY_DATA    *ContextSummaryData,\r
+  IN MEMORY_PROFILE_DRIVER_INFO                 *DriverInfo\r
+  )\r
+{\r
+  MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA       *DriverSummaryInfoData;\r
+  MEMORY_PROFILE_ALLOC_INFO                     *AllocInfo;\r
+  UINTN                                         AllocIndex;\r
+\r
+  if (DriverInfo->Header.Signature != MEMORY_PROFILE_DRIVER_INFO_SIGNATURE) {\r
+    return NULL;\r
+  }\r
+\r
+  DriverSummaryInfoData = AllocatePool (sizeof (*DriverSummaryInfoData) + sizeof (LIST_ENTRY));\r
+  if (DriverSummaryInfoData == NULL) {\r
+    return NULL;\r
+  }\r
+  DriverSummaryInfoData->Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;\r
+  DriverSummaryInfoData->DriverInfo = DriverInfo;\r
+  DriverSummaryInfoData->AllocSummaryInfoList = (LIST_ENTRY *) (DriverSummaryInfoData + 1);\r
+  InitializeListHead (DriverSummaryInfoData->AllocSummaryInfoList);\r
+  InsertTailList (ContextSummaryData->DriverSummaryInfoList, &DriverSummaryInfoData->Link);\r
+\r
+  AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) ((UINTN) DriverInfo + DriverInfo->Header.Length);\r
+  for (AllocIndex = 0; AllocIndex < DriverInfo->AllocRecordCount; AllocIndex++) {\r
+    AllocInfo = CreateAllocSummaryInfo (DriverSummaryInfoData, AllocInfo);\r
+    if (AllocInfo == NULL) {\r
+      return NULL;\r
+    }\r
+  }\r
+  return (MEMORY_PROFILE_DRIVER_INFO *) AllocInfo;\r
+}\r
+\r
+/**\r
+  Create Context summary information structure.\r
+\r
+  @param[in] ProfileBuffer      Memory profile base address.\r
+  @param[in] ProfileSize        Memory profile size.\r
+\r
+  @return Context summary information structure.\r
+\r
+**/\r
+MEMORY_PROFILE_CONTEXT_SUMMARY_DATA *\r
+CreateContextSummaryData (\r
+  IN PHYSICAL_ADDRESS           ProfileBuffer,\r
+  IN UINT64                     ProfileSize\r
+  )\r
+{\r
+  MEMORY_PROFILE_CONTEXT        *Context;\r
+  MEMORY_PROFILE_DRIVER_INFO    *DriverInfo;\r
+  UINTN                         DriverIndex;\r
+\r
+  Context = (MEMORY_PROFILE_CONTEXT *) ScanMemoryProfileBySignature (ProfileBuffer, ProfileSize, MEMORY_PROFILE_CONTEXT_SIGNATURE);\r
+  if (Context == NULL) {\r
+    return NULL;\r
+  }\r
+\r
+  mMemoryProfileContextSummary.Signature = MEMORY_PROFILE_CONTEXT_SIGNATURE;\r
+  mMemoryProfileContextSummary.Context = Context;\r
+  mMemoryProfileContextSummary.DriverSummaryInfoList = &mImageSummaryQueue;\r
+\r
+  DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) ((UINTN) Context + Context->Header.Length);\r
+  for (DriverIndex = 0; DriverIndex < Context->ImageCount; DriverIndex++) {\r
+    DriverInfo = CreateDriverSummaryInfo (&mMemoryProfileContextSummary, DriverInfo);\r
+    if (DriverInfo == NULL) {\r
+      return NULL;\r
+    }\r
+  }\r
+\r
+  return &mMemoryProfileContextSummary;\r
+}\r
+\r
+/**\r
+  Dump Context summary information.\r
+\r
+  @param[in] ContextSummaryData Context summary information data.\r
+  @param[in] IsForSmm           TRUE  - SMRAM profile.\r
+                                FALSE - UEFI memory profile.\r
+\r
+**/\r
+VOID\r
+DumpContextSummaryData (\r
+  IN MEMORY_PROFILE_CONTEXT_SUMMARY_DATA    *ContextSummaryData,\r
+  IN BOOLEAN                                IsForSmm\r
+  )\r
+{\r
+  MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA   *DriverSummaryInfoData;\r
+  MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA    *AllocSummaryInfoData;\r
+  LIST_ENTRY                                *DriverSummaryInfoList;\r
+  LIST_ENTRY                                *DriverSummaryLink;\r
+  LIST_ENTRY                                *AllocSummaryInfoList;\r
+  LIST_ENTRY                                *AllocSummaryLink;\r
+  MEMORY_PROFILE_DRIVER_INFO                *DriverInfo;\r
+  MEMORY_PROFILE_ALLOC_SUMMARY_INFO         *AllocSummaryInfo;\r
+  CHAR8                                     *NameString;\r
+\r
+  if (ContextSummaryData == NULL) {\r
+    return ;\r
+  }\r
+\r
+  Print (L"\nSummary Data:\n");\r
+\r
+  DriverSummaryInfoList = ContextSummaryData->DriverSummaryInfoList;\r
+  for (DriverSummaryLink = DriverSummaryInfoList->ForwardLink;\r
+       DriverSummaryLink != DriverSummaryInfoList;\r
+       DriverSummaryLink = DriverSummaryLink->ForwardLink) {\r
+    DriverSummaryInfoData = CR (\r
+                              DriverSummaryLink,\r
+                              MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA,\r
+                              Link,\r
+                              MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
+                              );\r
+    DriverInfo = DriverSummaryInfoData->DriverInfo;\r
+\r
+    NameString = GetDriverNameString (DriverInfo);\r
+    Print (L"\nDriver - %a (Usage - 0x%08x)", NameString, DriverInfo->CurrentUsage);\r
+    if (DriverInfo->CurrentUsage == 0) {\r
+      Print (L"\n");\r
+      continue;\r
+    }\r
+\r
+    if (DriverInfo->PdbStringOffset != 0) {\r
+      Print (L" (Pdb - %a)\n", (CHAR8 *) ((UINTN) DriverInfo + DriverInfo->PdbStringOffset));\r
+    } else {\r
+      Print (L"\n");\r
+    }\r
+    Print (L"Caller List:\n");\r
+    Print(L"  Count            Size                   RVA              Action\n");\r
+    Print(L"==========  ==================     ================== (================================)\n");\r
+    AllocSummaryInfoList = DriverSummaryInfoData->AllocSummaryInfoList;\r
+    for (AllocSummaryLink = AllocSummaryInfoList->ForwardLink;\r
+         AllocSummaryLink != AllocSummaryInfoList;\r
+         AllocSummaryLink = AllocSummaryLink->ForwardLink) {\r
+      AllocSummaryInfoData = CR (\r
+                               AllocSummaryLink,\r
+                               MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA,\r
+                               Link,\r
+                               MEMORY_PROFILE_ALLOC_SUMMARY_INFO_SIGNATURE\r
+                               );\r
+      AllocSummaryInfo = &AllocSummaryInfoData->AllocSummaryInfo;\r
+\r
+      Print(L"0x%08x  0x%016lx <== 0x%016lx",\r
+        AllocSummaryInfo->AllocateCount,\r
+        AllocSummaryInfo->TotalSize,\r
+        AllocSummaryInfo->CallerAddress - DriverInfo->ImageBase\r
+        );\r
+      Print (L" (%a)\n", ProfileActionToStr (AllocSummaryInfo->Action, AllocSummaryInfo->ActionString, IsForSmm));\r
+    }\r
+  }\r
+  return ;\r
+}\r
+\r
+/**\r
+  Destroy Context summary information.\r
+\r
+  @param[in, out] ContextSummaryData    Context summary information data.\r
+\r
+**/\r
+VOID\r
+DestroyContextSummaryData (\r
+  IN OUT MEMORY_PROFILE_CONTEXT_SUMMARY_DATA    *ContextSummaryData\r
+  )\r
+{\r
+  MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA       *DriverSummaryInfoData;\r
+  MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA        *AllocSummaryInfoData;\r
+  LIST_ENTRY                                    *DriverSummaryInfoList;\r
+  LIST_ENTRY                                    *DriverSummaryLink;\r
+  LIST_ENTRY                                    *AllocSummaryInfoList;\r
+  LIST_ENTRY                                    *AllocSummaryLink;\r
+\r
+  if (ContextSummaryData == NULL) {\r
+    return ;\r
+  }\r
+\r
+  DriverSummaryInfoList = ContextSummaryData->DriverSummaryInfoList;\r
+  for (DriverSummaryLink = DriverSummaryInfoList->ForwardLink;\r
+       DriverSummaryLink != DriverSummaryInfoList;\r
+       ) {\r
+    DriverSummaryInfoData = CR (\r
+                              DriverSummaryLink,\r
+                              MEMORY_PROFILE_DRIVER_SUMMARY_INFO_DATA,\r
+                              Link,\r
+                              MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
+                              );\r
+    DriverSummaryLink = DriverSummaryLink->ForwardLink;\r
+\r
+    AllocSummaryInfoList = DriverSummaryInfoData->AllocSummaryInfoList;\r
+    for (AllocSummaryLink = AllocSummaryInfoList->ForwardLink;\r
+         AllocSummaryLink != AllocSummaryInfoList;\r
+         ) {\r
+      AllocSummaryInfoData = CR (\r
+                               AllocSummaryLink,\r
+                               MEMORY_PROFILE_ALLOC_SUMMARY_INFO_DATA,\r
+                               Link,\r
+                               MEMORY_PROFILE_ALLOC_SUMMARY_INFO_SIGNATURE\r
+                               );\r
+      AllocSummaryLink = AllocSummaryLink->ForwardLink;\r
+\r
+      RemoveEntryList (&AllocSummaryInfoData->Link);\r
+      FreePool (AllocSummaryInfoData);\r
+    }\r
+\r
+    RemoveEntryList (&DriverSummaryInfoData->Link);\r
+    FreePool (DriverSummaryInfoData);\r
+  }\r
+  return ;\r
+}\r
+\r
 /**\r
   Get and dump UEFI memory profile data.\r
 \r
@@ -562,10 +1024,12 @@ GetUefiMemoryProfileData (
   VOID\r
   )\r
 {\r
-  EFI_STATUS                    Status;\r
-  EDKII_MEMORY_PROFILE_PROTOCOL *ProfileProtocol;\r
-  VOID                          *Data;\r
-  UINT64                        Size;\r
+  EFI_STATUS                          Status;\r
+  EDKII_MEMORY_PROFILE_PROTOCOL       *ProfileProtocol;\r
+  VOID                                *Data;\r
+  UINT64                              Size;\r
+  MEMORY_PROFILE_CONTEXT_SUMMARY_DATA *MemoryProfileContextSummaryData;\r
+  BOOLEAN                             RecordingState;\r
 \r
   Status = gBS->LocateProtocol (&gEdkiiMemoryProfileGuid, NULL, (VOID **) &ProfileProtocol);\r
   if (EFI_ERROR (Status)) {\r
@@ -573,6 +1037,15 @@ GetUefiMemoryProfileData (
     return Status;\r
   }\r
 \r
+  //\r
+  // Set recording state if needed.\r
+  //\r
+  RecordingState = MEMORY_PROFILE_RECORDING_DISABLE;\r
+  Status = ProfileProtocol->GetRecordingState (ProfileProtocol, &RecordingState);\r
+  if (RecordingState == MEMORY_PROFILE_RECORDING_ENABLE) {\r
+    ProfileProtocol->SetRecordingState (ProfileProtocol, MEMORY_PROFILE_RECORDING_DISABLE);\r
+  }\r
+\r
   Size = 0;\r
   Data = NULL;\r
   Status = ProfileProtocol->GetData (\r
@@ -582,13 +1055,9 @@ GetUefiMemoryProfileData (
                               );\r
   if (Status != EFI_BUFFER_TOO_SMALL) {\r
     Print (L"UefiMemoryProfile: GetData - %r\n", Status);\r
-    return Status;\r
+    goto Done;\r
   }\r
 \r
-  //\r
-  // Add one sizeof (MEMORY_PROFILE_ALLOC_INFO) to Size for this AllocatePool action.\r
-  //\r
-  Size = Size + sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
   Data = AllocateZeroPool ((UINTN) Size);\r
   if (Data == NULL) {\r
     Status = EFI_OUT_OF_RESOURCES;\r
@@ -602,20 +1071,39 @@ GetUefiMemoryProfileData (
                               Data\r
                               );\r
   if (EFI_ERROR (Status)) {\r
-    FreePool (Data);\r
     Print (L"UefiMemoryProfile: GetData - %r\n", Status);\r
-    return Status;\r
+    goto Done;\r
   }\r
 \r
 \r
   Print (L"UefiMemoryProfileSize - 0x%x\n", Size);\r
   Print (L"======= UefiMemoryProfile begin =======\n");\r
-  DumpMemoryProfile ((PHYSICAL_ADDRESS) (UINTN) Data, Size);\r
+  DumpMemoryProfile ((PHYSICAL_ADDRESS) (UINTN) Data, Size, FALSE);\r
+\r
+  //\r
+  // Dump summary information\r
+  //\r
+  MemoryProfileContextSummaryData = CreateContextSummaryData ((PHYSICAL_ADDRESS) (UINTN) Data, Size);\r
+  if (MemoryProfileContextSummaryData != NULL) {\r
+    DumpContextSummaryData (MemoryProfileContextSummaryData, FALSE);\r
+    DestroyContextSummaryData (MemoryProfileContextSummaryData);\r
+  }\r
+\r
   Print (L"======= UefiMemoryProfile end =======\n\n\n");\r
 \r
-  FreePool (Data);\r
+Done:\r
+  if (Data != NULL) {\r
+    FreePool (Data);\r
+  }\r
 \r
-  return EFI_SUCCESS;\r
+  //\r
+  // Restore recording state if needed.\r
+  //\r
+  if (RecordingState == MEMORY_PROFILE_RECORDING_ENABLE) {\r
+    ProfileProtocol->SetRecordingState (ProfileProtocol, MEMORY_PROFILE_RECORDING_ENABLE);\r
+  }\r
+\r
+  return Status;\r
 }\r
 \r
 /**\r
@@ -635,10 +1123,22 @@ GetSmramProfileData (
   UINT8                                         *CommBuffer;\r
   EFI_SMM_COMMUNICATE_HEADER                    *CommHeader;\r
   SMRAM_PROFILE_PARAMETER_GET_PROFILE_INFO      *CommGetProfileInfo;\r
-  SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA      *CommGetProfileData;\r
-  UINT64                                        ProfileSize;\r
-  PHYSICAL_ADDRESS                              ProfileBuffer;\r
+  SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET *CommGetProfileData;\r
+  SMRAM_PROFILE_PARAMETER_RECORDING_STATE       *CommRecordingState;\r
+  UINTN                                         ProfileSize;\r
+  VOID                                          *ProfileBuffer;\r
   EFI_SMM_COMMUNICATION_PROTOCOL                *SmmCommunication;\r
+  UINTN                                         MinimalSizeNeeded;\r
+  EDKII_PI_SMM_COMMUNICATION_REGION_TABLE       *PiSmmCommunicationRegionTable;\r
+  UINT32                                        Index;\r
+  EFI_MEMORY_DESCRIPTOR                         *Entry;\r
+  VOID                                          *Buffer;\r
+  UINTN                                         Size;\r
+  UINTN                                         Offset;\r
+  MEMORY_PROFILE_CONTEXT_SUMMARY_DATA           *MemoryProfileContextSummaryData;\r
+  BOOLEAN                                       RecordingState;\r
+\r
+  ProfileBuffer = NULL;\r
 \r
   Status = gBS->LocateProtocol (&gEfiSmmCommunicationProtocolGuid, NULL, (VOID **) &SmmCommunication);\r
   if (EFI_ERROR (Status)) {\r
@@ -646,14 +1146,83 @@ GetSmramProfileData (
     return Status;\r
   }\r
 \r
-  CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA);\r
-  CommBuffer = AllocateZeroPool (CommSize);\r
-  if (CommBuffer == NULL) {\r
-    Status = EFI_OUT_OF_RESOURCES;\r
-    Print (L"SmramProfile: AllocateZeroPool (0x%x) for comm buffer - %r\n", CommSize, Status);\r
+  MinimalSizeNeeded = sizeof (EFI_GUID) +\r
+                      sizeof (UINTN) +\r
+                      MAX (sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_INFO),\r
+                           MAX (sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET),\r
+                                sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE)));\r
+  MinimalSizeNeeded += MAX (sizeof (MEMORY_PROFILE_CONTEXT),\r
+                            MAX (sizeof (MEMORY_PROFILE_DRIVER_INFO),\r
+                                 MAX (sizeof (MEMORY_PROFILE_ALLOC_INFO),\r
+                                      MAX (sizeof (MEMORY_PROFILE_DESCRIPTOR),\r
+                                           MAX (sizeof (MEMORY_PROFILE_FREE_MEMORY),\r
+                                                sizeof (MEMORY_PROFILE_MEMORY_RANGE))))));\r
+\r
+  Status = EfiGetSystemConfigurationTable (\r
+             &gEdkiiPiSmmCommunicationRegionTableGuid,\r
+             (VOID **) &PiSmmCommunicationRegionTable\r
+             );\r
+  if (EFI_ERROR (Status)) {\r
+    DEBUG ((EFI_D_ERROR, "SmramProfile: Get PiSmmCommunicationRegionTable - %r\n", Status));\r
+    return Status;\r
+  }\r
+  ASSERT (PiSmmCommunicationRegionTable != NULL);\r
+  Entry = (EFI_MEMORY_DESCRIPTOR *) (PiSmmCommunicationRegionTable + 1);\r
+  Size = 0;\r
+  for (Index = 0; Index < PiSmmCommunicationRegionTable->NumberOfEntries; Index++) {\r
+    if (Entry->Type == EfiConventionalMemory) {\r
+      Size = EFI_PAGES_TO_SIZE ((UINTN) Entry->NumberOfPages);\r
+      if (Size >= MinimalSizeNeeded) {\r
+        break;\r
+      }\r
+    }\r
+    Entry = (EFI_MEMORY_DESCRIPTOR *) ((UINT8 *) Entry + PiSmmCommunicationRegionTable->DescriptorSize);\r
+  }\r
+  ASSERT (Index < PiSmmCommunicationRegionTable->NumberOfEntries);\r
+  CommBuffer = (UINT8 *) (UINTN) Entry->PhysicalStart;\r
+\r
+  //\r
+  // Set recording state if needed.\r
+  //\r
+  RecordingState = MEMORY_PROFILE_RECORDING_DISABLE;\r
+\r
+  CommHeader = (EFI_SMM_COMMUNICATE_HEADER *) &CommBuffer[0];\r
+  CopyMem (&CommHeader->HeaderGuid, &gEdkiiMemoryProfileGuid, sizeof (gEdkiiMemoryProfileGuid));\r
+  CommHeader->MessageLength = sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE);\r
+\r
+  CommRecordingState = (SMRAM_PROFILE_PARAMETER_RECORDING_STATE *) &CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];\r
+  CommRecordingState->Header.Command      = SMRAM_PROFILE_COMMAND_GET_RECORDING_STATE;\r
+  CommRecordingState->Header.DataLength   = sizeof (*CommRecordingState);\r
+  CommRecordingState->Header.ReturnStatus = (UINT64)-1;\r
+  CommRecordingState->RecordingState      = MEMORY_PROFILE_RECORDING_DISABLE;\r
+\r
+  CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + CommHeader->MessageLength;\r
+  Status = SmmCommunication->Communicate (SmmCommunication, CommBuffer, &CommSize);\r
+  if (EFI_ERROR (Status)) {\r
+    DEBUG ((EFI_D_ERROR, "SmramProfile: SmmCommunication - %r\n", Status));\r
     return Status;\r
   }\r
 \r
+  if (CommRecordingState->Header.ReturnStatus != 0) {\r
+    Print (L"SmramProfile: GetRecordingState - 0x%0x\n", CommRecordingState->Header.ReturnStatus);\r
+    return EFI_SUCCESS;\r
+  }\r
+  RecordingState = CommRecordingState->RecordingState;\r
+  if (RecordingState == MEMORY_PROFILE_RECORDING_ENABLE) {\r
+    CommHeader = (EFI_SMM_COMMUNICATE_HEADER *) &CommBuffer[0];\r
+    CopyMem (&CommHeader->HeaderGuid, &gEdkiiMemoryProfileGuid, sizeof (gEdkiiMemoryProfileGuid));\r
+    CommHeader->MessageLength = sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE);\r
+\r
+    CommRecordingState = (SMRAM_PROFILE_PARAMETER_RECORDING_STATE *) &CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];\r
+    CommRecordingState->Header.Command      = SMRAM_PROFILE_COMMAND_SET_RECORDING_STATE;\r
+    CommRecordingState->Header.DataLength   = sizeof (*CommRecordingState);\r
+    CommRecordingState->Header.ReturnStatus = (UINT64)-1;\r
+    CommRecordingState->RecordingState      = MEMORY_PROFILE_RECORDING_DISABLE;\r
+\r
+    CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + CommHeader->MessageLength;\r
+    SmmCommunication->Communicate (SmmCommunication, CommBuffer, &CommSize);\r
+  }\r
+\r
   //\r
   // Get Size\r
   //\r
@@ -669,71 +1238,108 @@ GetSmramProfileData (
 \r
   CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + CommHeader->MessageLength;\r
   Status = SmmCommunication->Communicate (SmmCommunication, CommBuffer, &CommSize);\r
-  if (EFI_ERROR (Status)) {\r
-    FreePool (CommBuffer);\r
-    DEBUG ((EFI_D_ERROR, "SmramProfile: SmmCommunication - %r\n", Status));\r
-    return Status;\r
-  }\r
+  ASSERT_EFI_ERROR (Status);\r
 \r
   if (CommGetProfileInfo->Header.ReturnStatus != 0) {\r
+    Status = EFI_SUCCESS;\r
     Print (L"SmramProfile: GetProfileInfo - 0x%0x\n", CommGetProfileInfo->Header.ReturnStatus);\r
-    return EFI_SUCCESS;\r
+    goto Done;\r
   }\r
 \r
-  ProfileSize = CommGetProfileInfo->ProfileSize;\r
+  ProfileSize = (UINTN) CommGetProfileInfo->ProfileSize;\r
 \r
   //\r
   // Get Data\r
   //\r
-  ProfileBuffer = (PHYSICAL_ADDRESS) (UINTN) AllocateZeroPool ((UINTN) ProfileSize);\r
-  if (ProfileBuffer == 0) {\r
-    FreePool (CommBuffer);\r
+  ProfileBuffer = AllocateZeroPool (ProfileSize);\r
+  if (ProfileBuffer == NULL) {\r
     Status = EFI_OUT_OF_RESOURCES;\r
-    Print (L"SmramProfile: AllocateZeroPool (0x%x) for profile buffer - %r\n", (UINTN) ProfileSize, Status);\r
-    return Status;\r
+    Print (L"SmramProfile: AllocateZeroPool (0x%x) for profile buffer - %r\n", ProfileSize, Status);\r
+    goto Done;\r
   }\r
 \r
   CommHeader = (EFI_SMM_COMMUNICATE_HEADER *) &CommBuffer[0];\r
   CopyMem (&CommHeader->HeaderGuid, &gEdkiiMemoryProfileGuid, sizeof(gEdkiiMemoryProfileGuid));\r
-  CommHeader->MessageLength = sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA);\r
+  CommHeader->MessageLength = sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET);\r
 \r
-  CommGetProfileData = (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA *) &CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];\r
-  CommGetProfileData->Header.Command      = SMRAM_PROFILE_COMMAND_GET_PROFILE_DATA;\r
+  CommGetProfileData = (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET *) &CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];\r
+  CommGetProfileData->Header.Command      = SMRAM_PROFILE_COMMAND_GET_PROFILE_DATA_BY_OFFSET;\r
   CommGetProfileData->Header.DataLength   = sizeof (*CommGetProfileData);\r
   CommGetProfileData->Header.ReturnStatus = (UINT64)-1;\r
-  CommGetProfileData->ProfileSize         = ProfileSize;\r
-  CommGetProfileData->ProfileBuffer       = ProfileBuffer;\r
 \r
   CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + CommHeader->MessageLength;\r
-  Status = SmmCommunication->Communicate (SmmCommunication, CommBuffer, &CommSize);\r
-  ASSERT_EFI_ERROR (Status);\r
+  Buffer = (UINT8 *) CommHeader + CommSize;\r
+  Size -= CommSize;\r
+\r
+  CommGetProfileData->ProfileBuffer       = (PHYSICAL_ADDRESS) (UINTN) Buffer;\r
+  CommGetProfileData->ProfileOffset       = 0;\r
+  while (CommGetProfileData->ProfileOffset < ProfileSize) {\r
+    Offset = (UINTN) CommGetProfileData->ProfileOffset;\r
+    if (Size <= (ProfileSize - CommGetProfileData->ProfileOffset)) {\r
+      CommGetProfileData->ProfileSize = (UINT64) Size;\r
+    } else {\r
+      CommGetProfileData->ProfileSize = (UINT64) (ProfileSize - CommGetProfileData->ProfileOffset);\r
+    }\r
+    Status = SmmCommunication->Communicate (SmmCommunication, CommBuffer, &CommSize);\r
+    ASSERT_EFI_ERROR (Status);\r
 \r
-  if (CommGetProfileData->Header.ReturnStatus != 0) {\r
-    FreePool ((VOID *) (UINTN) CommGetProfileData->ProfileBuffer);\r
-    FreePool (CommBuffer);\r
-    Print (L"GetProfileData - 0x%x\n", CommGetProfileData->Header.ReturnStatus);\r
-    return EFI_SUCCESS;\r
+    if (CommGetProfileData->Header.ReturnStatus != 0) {\r
+      Status = EFI_SUCCESS;\r
+      Print (L"GetProfileData - 0x%x\n", CommGetProfileData->Header.ReturnStatus);\r
+      goto Done;\r
+    }\r
+    CopyMem ((UINT8 *) ProfileBuffer + Offset, (VOID *) (UINTN) CommGetProfileData->ProfileBuffer, (UINTN) CommGetProfileData->ProfileSize);\r
   }\r
 \r
 \r
-  Print (L"SmramProfileSize - 0x%x\n", CommGetProfileData->ProfileSize);\r
+  Print (L"SmramProfileSize - 0x%x\n", ProfileSize);\r
   Print (L"======= SmramProfile begin =======\n");\r
-  DumpMemoryProfile (CommGetProfileData->ProfileBuffer, CommGetProfileData->ProfileSize);\r
+  DumpMemoryProfile ((PHYSICAL_ADDRESS) (UINTN) ProfileBuffer, ProfileSize, TRUE);\r
+\r
+  //\r
+  // Dump summary information\r
+  //\r
+  MemoryProfileContextSummaryData = CreateContextSummaryData ((PHYSICAL_ADDRESS) (UINTN) ProfileBuffer, ProfileSize);\r
+  if (MemoryProfileContextSummaryData != NULL) {\r
+    DumpContextSummaryData (MemoryProfileContextSummaryData, TRUE);\r
+    DestroyContextSummaryData (MemoryProfileContextSummaryData);\r
+  }\r
+\r
   Print (L"======= SmramProfile end =======\n\n\n");\r
 \r
-  FreePool ((VOID *) (UINTN) CommGetProfileData->ProfileBuffer);\r
-  FreePool (CommBuffer);\r
+Done:\r
+  if (ProfileBuffer != NULL) {\r
+    FreePool (ProfileBuffer);\r
+  }\r
 \r
-  return EFI_SUCCESS;\r
+  //\r
+  // Restore recording state if needed.\r
+  //\r
+  if (RecordingState == MEMORY_PROFILE_RECORDING_ENABLE) {\r
+    CommHeader = (EFI_SMM_COMMUNICATE_HEADER *) &CommBuffer[0];\r
+    CopyMem (&CommHeader->HeaderGuid, &gEdkiiMemoryProfileGuid, sizeof (gEdkiiMemoryProfileGuid));\r
+    CommHeader->MessageLength = sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE);\r
+\r
+    CommRecordingState = (SMRAM_PROFILE_PARAMETER_RECORDING_STATE *) &CommBuffer[OFFSET_OF (EFI_SMM_COMMUNICATE_HEADER, Data)];\r
+    CommRecordingState->Header.Command      = SMRAM_PROFILE_COMMAND_SET_RECORDING_STATE;\r
+    CommRecordingState->Header.DataLength   = sizeof (*CommRecordingState);\r
+    CommRecordingState->Header.ReturnStatus = (UINT64)-1;\r
+    CommRecordingState->RecordingState      = MEMORY_PROFILE_RECORDING_ENABLE;\r
+\r
+    CommSize = sizeof (EFI_GUID) + sizeof (UINTN) + CommHeader->MessageLength;\r
+    SmmCommunication->Communicate (SmmCommunication, CommBuffer, &CommSize);\r
+  }\r
+\r
+  return Status;\r
 }\r
 \r
 /**\r
   The user Entry Point for Application. The user code starts with this function\r
   as the real entry point for the image goes into a library that calls this function.\r
 \r
-  @param[in] ImageHandle    The firmware allocated handle for the EFI image.  \r
+  @param[in] ImageHandle    The firmware allocated handle for the EFI image.\r
   @param[in] SystemTable    A pointer to the EFI System Table.\r
-  \r
+\r
   @retval EFI_SUCCESS       The entry point is executed successfully.\r
   @retval other             Some error occurs when executing this entry point.\r
 \r
@@ -745,7 +1351,7 @@ UefiMain (
   IN EFI_SYSTEM_TABLE   *SystemTable\r
   )\r
 {\r
-  EFI_STATUS                    Status;\r
+  EFI_STATUS                     Status;\r
 \r
   Status = GetUefiMemoryProfileData ();\r
   if (EFI_ERROR (Status)) {\r