Merge branch 'master' of https://github.com/tianocore/edk2
authordavid wei <david.wei@intel.com>
Fri, 1 Jul 2016 07:05:48 +0000 (15:05 +0800)
committerdavid wei <david.wei@intel.com>
Fri, 1 Jul 2016 07:05:48 +0000 (15:05 +0800)
67 files changed:
BaseTools/Scripts/MemoryProfileSymbolGen.py [new file with mode: 0644]
EdkCompatibilityPkg/Compatibility/FrameworkHiiOnUefiHiiThunk/HiiDatabase.c
EdkCompatibilityPkg/Compatibility/FrameworkHiiOnUefiHiiThunk/Utility.c
IntelFrameworkModulePkg/Universal/StatusCode/RuntimeDxe/StatusCodeRuntimeDxe.h
IntelFsp2WrapperPkg/FspmWrapperPeim/FspmWrapperPeim.c
IntelFsp2WrapperPkg/FspsWrapperPeim/FspsWrapperPeim.c
MdeModulePkg/Application/MemoryProfileInfo/MemoryProfileInfo.c
MdeModulePkg/Application/MemoryProfileInfo/MemoryProfileInfo.inf
MdeModulePkg/Application/VariableInfo/VariableInfo.c
MdeModulePkg/Bus/Pci/SdMmcPciHcDxe/SdMmcPciHcDxe.c
MdeModulePkg/Bus/Pci/SdMmcPciHcDxe/SdMmcPciHci.c
MdeModulePkg/Bus/Sd/EmmcDxe/EmmcBlockIo.c
MdeModulePkg/Bus/Sd/EmmcDxe/EmmcDxe.c
MdeModulePkg/Bus/Sd/SdDxe/SdBlockIo.c
MdeModulePkg/Bus/Sd/SdDxe/SdDxe.c
MdeModulePkg/Core/Dxe/DxeMain.h
MdeModulePkg/Core/Dxe/DxeMain.inf
MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
MdeModulePkg/Core/Dxe/Mem/Page.c
MdeModulePkg/Core/Dxe/Mem/Pool.c
MdeModulePkg/Core/PiSmmCore/Page.c
MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
MdeModulePkg/Core/PiSmmCore/PiSmmCore.h
MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
MdeModulePkg/Core/PiSmmCore/PiSmmIpl.c
MdeModulePkg/Core/PiSmmCore/Pool.c
MdeModulePkg/Core/PiSmmCore/SmramProfileRecord.c
MdeModulePkg/Include/Guid/MemoryProfile.h
MdeModulePkg/Include/Library/MemoryProfileLib.h [new file with mode: 0644]
MdeModulePkg/Library/DxeCoreMemoryAllocationLib/DxeCoreMemoryAllocationLib.inf
MdeModulePkg/Library/DxeCoreMemoryAllocationLib/DxeCoreMemoryAllocationLib.uni
MdeModulePkg/Library/DxeCoreMemoryAllocationLib/DxeCoreMemoryAllocationProfileLib.inf [new file with mode: 0644]
MdeModulePkg/Library/DxeCoreMemoryAllocationLib/DxeCoreMemoryAllocationProfileLib.uni [new file with mode: 0644]
MdeModulePkg/Library/DxeCoreMemoryAllocationLib/DxeCoreMemoryProfileLib.c [new file with mode: 0644]
MdeModulePkg/Library/DxeCoreMemoryAllocationLib/DxeCoreMemoryProfileLibNull.c [new file with mode: 0644]
MdeModulePkg/Library/DxeCoreMemoryAllocationLib/DxeCoreMemoryProfileServices.h [new file with mode: 0644]
MdeModulePkg/Library/DxeCoreMemoryAllocationLib/MemoryAllocationLib.c
MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/MemoryAllocationLib.c
MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/PiSmmCoreMemoryAllocationLib.inf
MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/PiSmmCoreMemoryAllocationLib.uni
MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/PiSmmCoreMemoryAllocationProfileLib.inf [new file with mode: 0644]
MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/PiSmmCoreMemoryAllocationProfileLib.uni [new file with mode: 0644]
MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/PiSmmCoreMemoryProfileLib.c [new file with mode: 0644]
MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/PiSmmCoreMemoryProfileLibNull.c [new file with mode: 0644]
MdeModulePkg/Library/PiSmmCoreMemoryAllocationLib/PiSmmCoreMemoryProfileServices.h [new file with mode: 0644]
MdeModulePkg/Library/SmmMemoryAllocationProfileLib/MemoryAllocationLib.c [new file with mode: 0644]
MdeModulePkg/Library/SmmMemoryAllocationProfileLib/SmmMemoryAllocationProfileLib.inf [new file with mode: 0644]
MdeModulePkg/Library/SmmMemoryAllocationProfileLib/SmmMemoryAllocationProfileLib.uni [new file with mode: 0644]
MdeModulePkg/Library/SmmMemoryAllocationProfileLib/SmmMemoryProfileLib.c [new file with mode: 0644]
MdeModulePkg/Library/UefiMemoryAllocationProfileLib/DxeMemoryProfileLib.c [new file with mode: 0644]
MdeModulePkg/Library/UefiMemoryAllocationProfileLib/MemoryAllocationLib.c [new file with mode: 0644]
MdeModulePkg/Library/UefiMemoryAllocationProfileLib/UefiMemoryAllocationProfileLib.inf [new file with mode: 0644]
MdeModulePkg/Library/UefiMemoryAllocationProfileLib/UefiMemoryAllocationProfileLib.uni [new file with mode: 0644]
MdeModulePkg/MdeModulePkg.dec
MdeModulePkg/MdeModulePkg.dsc
MdeModulePkg/MdeModulePkg.uni
NetworkPkg/IpSecDxe/Ikev2/Sa.c
QuarkPlatformPkg/Acpi/Dxe/AcpiPlatform/AcpiPlatform.c
QuarkPlatformPkg/Acpi/Dxe/AcpiPlatform/MadtPlatform.c
QuarkPlatformPkg/Library/PlatformHelperLib/PlatformHelperDxe.c
QuarkPlatformPkg/Platform/Pei/PlatformInit/MrcWrapper.c
SecurityPkg/Tcg/Tcg2Dxe/Tcg2Dxe.c
SecurityPkg/Tcg/TrEEDxe/TrEEDxe.c
ShellPkg/Library/UefiHandleParsingLib/UefiHandleParsingLib.c
ShellPkg/Library/UefiShellLevel1CommandsLib/If.c
ShellPkg/Library/UefiShellLib/UefiShellLib.c
UefiCpuPkg/PiSmmCpuDxeSmm/X64/SmmProfileArch.c

diff --git a/BaseTools/Scripts/MemoryProfileSymbolGen.py b/BaseTools/Scripts/MemoryProfileSymbolGen.py
new file mode 100644 (file)
index 0000000..a9790d8
--- /dev/null
@@ -0,0 +1,281 @@
+##\r
+# Generate symbal for memory profile info.\r
+#\r
+# This tool depends on DIA2Dump.exe (VS) or nm (gcc) to parse debug entry.\r
+#\r
+# Copyright (c) 2016, Intel Corporation. All rights reserved.<BR>\r
+# This program and the accompanying materials are licensed and made available under\r
+# the terms and conditions of the BSD License that accompanies this distribution.\r
+# 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
+import os\r
+import re\r
+import sys\r
+from optparse import OptionParser\r
+\r
+versionNumber = "1.0"\r
+__copyright__ = "Copyright (c) 2016, Intel Corporation. All rights reserved."\r
+\r
+class Symbols:\r
+    def __init__(self):\r
+        self.listLineAddress = []\r
+        self.pdbName = ""\r
+        # Cache for function\r
+        self.functionName = ""\r
+        # Cache for line\r
+        self.sourceName = ""\r
+\r
+\r
+    def getSymbol (self, rva):\r
+        index = 0\r
+        lineName  = 0\r
+        sourceName = "??"\r
+        while index + 1 < self.lineCount :\r
+            if self.listLineAddress[index][0] <= rva and self.listLineAddress[index + 1][0] > rva :\r
+                offset = rva - self.listLineAddress[index][0]\r
+                functionName = self.listLineAddress[index][1]\r
+                lineName = self.listLineAddress[index][2]\r
+                sourceName = self.listLineAddress[index][3]\r
+                if lineName == 0 :\r
+                  return " (" + self.listLineAddress[index][1] + "() - " + ")"\r
+                else :\r
+                  return " (" + self.listLineAddress[index][1] + "() - " + sourceName + ":" + str(lineName) + ")"\r
+            index += 1\r
+\r
+        return " (unknown)"\r
+\r
+    def parse_debug_file(self, driverName, pdbName):\r
+        if cmp (pdbName, "") == 0 :\r
+            return\r
+        self.pdbName = pdbName;\r
+\r
+        try:\r
+            nmCommand = "nm"\r
+            nmLineOption = "-l"\r
+            print "parsing (debug) - " + pdbName\r
+            os.system ('%s %s %s > nmDump.line.log' % (nmCommand, nmLineOption, pdbName))\r
+        except :\r
+            print 'ERROR: nm command not available.  Please verify PATH'\r
+            return\r
+\r
+        #\r
+        # parse line\r
+        #\r
+        linefile = open("nmDump.line.log")\r
+        reportLines = linefile.readlines()\r
+        linefile.close()\r
+\r
+        # 000113ca T AllocatePool      c:\home\edk-ii\MdePkg\Library\UefiMemoryAllocationLib\MemoryAllocationLib.c:399\r
+        patchLineFileMatchString = "([0-9a-fA-F]{8})\s+[T|D|t|d]\s+(\w+)\s*((?:[a-zA-Z]:)?[\w+\-./_a-zA-Z0-9\\\\]*):?([0-9]*)"\r
+\r
+        for reportLine in reportLines:\r
+            #print "check - " + reportLine\r
+            match = re.match(patchLineFileMatchString, reportLine)\r
+            if match is not None:\r
+                #print "match - " + reportLine[:-1]\r
+                #print "0 - " + match.group(0)\r
+                #print "1 - " + match.group(1)\r
+                #print "2 - " + match.group(2)\r
+                #print "3 - " + match.group(3)\r
+                #print "4 - " + match.group(4)\r
+\r
+                rva = int (match.group(1), 16)\r
+                functionName = match.group(2)\r
+                sourceName = match.group(3)\r
+                if cmp (match.group(4), "") != 0 :\r
+                    lineName = int (match.group(4))\r
+                else :\r
+                    lineName = 0\r
+                self.listLineAddress.append ([rva, functionName, lineName, sourceName])\r
+\r
+        self.lineCount = len (self.listLineAddress)\r
+\r
+        self.listLineAddress = sorted(self.listLineAddress, key=lambda symbolAddress:symbolAddress[0])\r
+\r
+        #for key in self.listLineAddress :\r
+            #print "rva - " + "%x"%(key[0]) + ", func - " + key[1] + ", line - " + str(key[2]) + ", source - " + key[3]\r
+\r
+    def parse_pdb_file(self, driverName, pdbName):\r
+        if cmp (pdbName, "") == 0 :\r
+            return\r
+        self.pdbName = pdbName;\r
+\r
+        try:\r
+            #DIA2DumpCommand = "\"C:\\Program Files (x86)\Microsoft Visual Studio 14.0\\DIA SDK\\Samples\\DIA2Dump\\x64\\Debug\\Dia2Dump.exe\""\r
+            DIA2DumpCommand = "Dia2Dump.exe"\r
+            #DIA2SymbolOption = "-p"\r
+            DIA2LinesOption = "-l"\r
+            print "parsing (pdb) - " + pdbName\r
+            #os.system ('%s %s %s > DIA2Dump.symbol.log' % (DIA2DumpCommand, DIA2SymbolOption, pdbName))\r
+            os.system ('%s %s %s > DIA2Dump.line.log' % (DIA2DumpCommand, DIA2LinesOption, pdbName))\r
+        except :\r
+            print 'ERROR: DIA2Dump command not available.  Please verify PATH'\r
+            return\r
+\r
+        #\r
+        # parse line\r
+        #\r
+        linefile = open("DIA2Dump.line.log")\r
+        reportLines = linefile.readlines()\r
+        linefile.close()\r
+\r
+        #   ** GetDebugPrintErrorLevel\r
+        #      line 32 at [0000C790][0001:0000B790], len = 0x3 c:\home\edk-ii\mdepkg\library\basedebugprinterrorlevellib\basedebugprinterrorlevellib.c (MD5: 687C0AE564079D35D56ED5D84A6164CC)\r
+        #      line 36 at [0000C793][0001:0000B793], len = 0x5\r
+        #      line 37 at [0000C798][0001:0000B798], len = 0x2\r
+\r
+        patchLineFileMatchString = "\s+line ([0-9]+) at \[([0-9a-fA-F]{8})\]\[[0-9a-fA-F]{4}\:[0-9a-fA-F]{8}\], len = 0x[0-9a-fA-F]+\s*([\w+\-\:./_a-zA-Z0-9\\\\]*)\s*"\r
+        patchLineFileMatchStringFunc = "\*\*\s+(\w+)\s*"\r
+\r
+        for reportLine in reportLines:\r
+            #print "check line - " + reportLine\r
+            match = re.match(patchLineFileMatchString, reportLine)\r
+            if match is not None:\r
+                #print "match - " + reportLine[:-1]\r
+                #print "0 - " + match.group(0)\r
+                #print "1 - " + match.group(1)\r
+                #print "2 - " + match.group(2)\r
+                if cmp (match.group(3), "") != 0 :\r
+                    self.sourceName = match.group(3)\r
+                sourceName = self.sourceName\r
+                functionName = self.functionName\r
+\r
+                rva = int (match.group(2), 16)\r
+                lineName = int (match.group(1))\r
+                self.listLineAddress.append ([rva, functionName, lineName, sourceName])\r
+            else :\r
+                match = re.match(patchLineFileMatchStringFunc, reportLine)\r
+                if match is not None:\r
+                    self.functionName = match.group(1)\r
+\r
+        self.lineCount = len (self.listLineAddress)\r
+        self.listLineAddress = sorted(self.listLineAddress, key=lambda symbolAddress:symbolAddress[0])\r
+\r
+        #for key in self.listLineAddress :\r
+            #print "rva - " + "%x"%(key[0]) + ", func - " + key[1] + ", line - " + str(key[2]) + ", source - " + key[3]\r
+\r
+class SymbolsFile:\r
+    def __init__(self):\r
+        self.symbolsTable = {}\r
+\r
+symbolsFile = ""\r
+\r
+driverName = ""\r
+rvaName = ""\r
+symbolName = ""\r
+\r
+def getSymbolName(driverName, rva):\r
+    global symbolsFile\r
+\r
+    #print "driverName - " + driverName\r
+\r
+    try :\r
+        symbolList = symbolsFile.symbolsTable[driverName]\r
+        if symbolList is not None:\r
+            return symbolList.getSymbol (rva)\r
+        else:\r
+            return " (???)"\r
+    except Exception:\r
+        return " (???)"\r
+\r
+def processLine(newline):\r
+    global driverName\r
+    global rvaName\r
+\r
+    driverPrefixLen = len("Driver - ")\r
+    # get driver name\r
+    if cmp(newline[0:driverPrefixLen],"Driver - ") == 0 :\r
+        driverlineList = newline.split(" ")\r
+        driverName = driverlineList[2]\r
+        #print "Checking : ", driverName\r
+\r
+        # EDKII application output\r
+        pdbMatchString = "Driver - \w* \(Usage - 0x[0-9a-fA-F]+\) \(Pdb - ([:\-.\w\\\\/]*)\)\s*"\r
+        pdbName = ""\r
+        match = re.match(pdbMatchString, newline)\r
+        if match is not None:\r
+            #print "match - " + newline\r
+            #print "0 - " + match.group(0)\r
+            #print "1 - " + match.group(1)\r
+            pdbName = match.group(1)\r
+            #print "PDB - " + pdbName\r
+\r
+        symbolsFile.symbolsTable[driverName] = Symbols()\r
+\r
+        if cmp (pdbName[-3:], "pdb") == 0 :\r
+            symbolsFile.symbolsTable[driverName].parse_pdb_file (driverName, pdbName)\r
+        else :\r
+            symbolsFile.symbolsTable[driverName].parse_debug_file (driverName, pdbName)\r
+\r
+    elif cmp(newline,"") == 0 :\r
+        driverName = ""\r
+\r
+    # check entry line\r
+    if newline.find ("<==") != -1 :\r
+        entry_list = newline.split(" ")\r
+        rvaName = entry_list[4]\r
+        #print "rva : ", rvaName\r
+        symbolName = getSymbolName (driverName, int(rvaName, 16))\r
+    else :\r
+        rvaName = ""\r
+        symbolName = ""\r
+\r
+    if cmp(rvaName,"") == 0 :\r
+        return newline\r
+    else :\r
+        return newline + symbolName\r
+\r
+def myOptionParser():\r
+    usage = "%prog [--version] [-h] [--help] [-i inputfile [-o outputfile]]"\r
+    Parser = OptionParser(usage=usage, description=__copyright__, version="%prog " + str(versionNumber))\r
+    Parser.add_option("-i", "--inputfile", dest="inputfilename", type="string", help="The input memory profile info file output from MemoryProfileInfo application in MdeModulePkg")\r
+    Parser.add_option("-o", "--outputfile", dest="outputfilename", type="string", help="The output memory profile info file with symbol, MemoryProfileInfoSymbol.txt will be used if it is not specified")\r
+\r
+    (Options, args) = Parser.parse_args()\r
+    if Options.inputfilename is None:\r
+        Parser.error("no input file specified")\r
+    if Options.outputfilename is None:\r
+        Options.outputfilename = "MemoryProfileInfoSymbol.txt"\r
+    return Options\r
+\r
+def main():\r
+    global symbolsFile\r
+    global Options\r
+    Options = myOptionParser()\r
+\r
+    symbolsFile = SymbolsFile()\r
+\r
+    try :\r
+        file = open(Options.inputfilename)\r
+    except Exception:\r
+        print "fail to open " + Options.inputfilename\r
+        return 1\r
+    try :\r
+        newfile = open(Options.outputfilename, "w")\r
+    except Exception:\r
+        print "fail to open " + Options.outputfilename\r
+        return 1\r
+\r
+    try:\r
+        while 1:\r
+            line = file.readline()\r
+            if not line:\r
+                break\r
+            newline = line[:-1]\r
+\r
+            newline = processLine(newline)\r
+\r
+            newfile.write(newline)\r
+            newfile.write("\n")\r
+    finally:\r
+        file.close()\r
+        newfile.close()\r
+\r
+if __name__ == '__main__':\r
+    sys.exit(main())\r
index 386ff32..2d456da 100644 (file)
@@ -494,7 +494,7 @@ HiiGetSecondaryLanguages (
   UnicodeStrToAsciiStr (PrimaryLanguage, PrimaryLang639);\r
 \r
   PrimaryLang4646 = ConvertLanguagesIso639ToRfc4646 (PrimaryLang639);\r
-  ASSERT_EFI_ERROR (PrimaryLang4646 != NULL);\r
+  ASSERT (PrimaryLang4646 != NULL);\r
 \r
   SecLangCodes4646 = HiiGetSupportedSecondaryLanguages (UefiHiiHandle, PrimaryLang4646);\r
 \r
index f6d97f6..d269b8e 100644 (file)
@@ -153,7 +153,7 @@ ExportPackageLists (
                            &Size,\r
                            PackageListHdr\r
                            );\r
-  ASSERT_EFI_ERROR (Status != EFI_BUFFER_TOO_SMALL);\r
+  ASSERT (Status != EFI_BUFFER_TOO_SMALL);\r
   \r
   if (Status == EFI_BUFFER_TOO_SMALL) {\r
     PackageListHdr = AllocateZeroPool (Size);\r
index 9d0ed19..ae20f5b 100644 (file)
@@ -53,15 +53,6 @@ typedef struct {
 } DATAHUB_STATUSCODE_RECORD;\r
 \r
 \r
-//\r
-// Runtime memory status code worker definition\r
-//\r
-typedef struct {\r
-  UINT32   RecordIndex;\r
-  UINT32   NumberOfRecords;\r
-  UINT32   MaxRecordsNumber;\r
-} RUNTIME_MEMORY_STATUSCODE_HEADER;\r
-\r
 extern RUNTIME_MEMORY_STATUSCODE_HEADER  *mRtMemoryStatusCodeTable;\r
 \r
 /**\r
index 6144ad7..c98513e 100644 (file)
@@ -68,6 +68,11 @@ PeiFspMemoryInit (
   // Copy default FSP-M UPD data from Flash\r
   //\r
   FspmHeaderPtr = (FSP_INFO_HEADER *)FspFindFspHeader (PcdGet32 (PcdFspmBaseAddress));\r
+  DEBUG ((DEBUG_INFO, "FspmHeaderPtr - 0x%x\n", FspmHeaderPtr));\r
+  if (FspmHeaderPtr == NULL) {\r
+    return EFI_DEVICE_ERROR;\r
+  }\r
+\r
   FspmUpdDataPtr = (FSPM_UPD_COMMON *)AllocateZeroPool ((UINTN)FspmHeaderPtr->CfgRegionSize);\r
   ASSERT (FspmUpdDataPtr != NULL);\r
   SourceData = (UINTN *)((UINTN)FspmHeaderPtr->ImageBase + (UINTN)FspmHeaderPtr->CfgRegionOffset);\r
index 7a65ad7..c923690 100644 (file)
@@ -241,6 +241,11 @@ PeiMemoryDiscoveredNotify (
   // Copy default FSP-S UPD data from Flash\r
   //\r
   FspsHeaderPtr = (FSP_INFO_HEADER *)FspFindFspHeader (PcdGet32 (PcdFspsBaseAddress));\r
+  DEBUG ((DEBUG_INFO, "FspsHeaderPtr - 0x%x\n", FspsHeaderPtr));\r
+  if (FspsHeaderPtr == NULL) {\r
+    return EFI_DEVICE_ERROR;\r
+  }\r
+\r
   FspsUpdDataPtr = (FSPS_UPD_COMMON *)AllocateZeroPool ((UINTN)FspsHeaderPtr->CfgRegionSize);\r
   ASSERT (FspsUpdDataPtr != NULL);\r
   SourceData = (UINTN *)((UINTN)FspsHeaderPtr->ImageBase + (UINTN)FspsHeaderPtr->CfgRegionOffset);\r
index ea2a00b..48d8f46 100644 (file)
@@ -1,17 +1,16 @@
 /** @file\r
-  \r
+\r
   Copyright (c) 2014 - 2016, 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
+  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
+  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 <Base.h>\r
 #include <Uefi.h>\r
 #include <PiDxe.h>\r
 #include <Library/BaseLib.h>\r
@@ -23,7 +22,6 @@
 #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 <Guid/MemoryProfile.h>\r
 #include <Guid/PiSmmCommunicationRegionTable.h>\r
 \r
-CHAR16 *mActionString[] = {\r
-  L"Unknown",\r
-  L"AllocatePages",\r
-  L"FreePages",\r
-  L"AllocatePool",\r
-  L"FreePool",\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
+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
-/** \r
+LIST_ENTRY  mImageSummaryQueue = INITIALIZE_LIST_HEAD_VARIABLE (mImageSummaryQueue);\r
+MEMORY_PROFILE_CONTEXT_SUMMARY_DATA mMemoryProfileContextSummary;\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] == '\\') {\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
@@ -161,29 +235,24 @@ 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->Header.Length > sizeof (MEMORY_PROFILE_DRIVER_INFO)) {\r
+    GetShortPdbFileName ((CHAR8 *) (DriverInfo + 1), mNameString);\r
+    return mNameString;\r
   }\r
 \r
   if (!CompareGuid (&DriverInfo->FileName, &gZeroGuid)) {\r
@@ -203,17 +272,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
@@ -224,7 +296,7 @@ GetDriverNameString (
   @return Pointer to string.\r
 \r
 **/\r
-CHAR16 *\r
+CHAR8 *\r
 ProfileMemoryTypeToStr (\r
   IN EFI_MEMORY_TYPE    MemoryType\r
   )\r
@@ -248,12 +320,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 = sizeof (mSmmActionString) / sizeof (mSmmActionString[0]);\r
+  } else {\r
+    ActionString = mActionString;\r
+    ActionStringCount = sizeof (mActionString) / sizeof (mActionString[0]);\r
+  }\r
+\r
+  if ((UINTN) (UINT32) Action < ActionStringCount) {\r
+    return ActionString[Action];\r
+  }\r
+  for (Index = 0; Index < sizeof (mExtActionString) / sizeof (mExtActionString[0]); 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
@@ -262,20 +385,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
@@ -287,6 +420,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
@@ -294,12 +429,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
@@ -307,28 +444,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
@@ -340,13 +480,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
@@ -359,14 +502,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
@@ -375,7 +518,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
@@ -404,7 +547,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
@@ -433,7 +576,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
@@ -470,7 +613,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
@@ -513,6 +656,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
@@ -524,12 +671,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
@@ -538,7 +688,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
@@ -552,6 +702,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
@@ -564,10 +1025,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
@@ -575,6 +1038,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
@@ -584,13 +1056,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
@@ -604,20 +1072,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
@@ -638,6 +1125,7 @@ GetSmramProfileData (
   EFI_SMM_COMMUNICATE_HEADER                    *CommHeader;\r
   SMRAM_PROFILE_PARAMETER_GET_PROFILE_INFO      *CommGetProfileInfo;\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
@@ -648,6 +1136,10 @@ GetSmramProfileData (
   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
@@ -658,7 +1150,8 @@ GetSmramProfileData (
   MinimalSizeNeeded = sizeof (EFI_GUID) +\r
                       sizeof (UINTN) +\r
                       MAX (sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_INFO),\r
-                           sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET));\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
@@ -689,6 +1182,48 @@ GetSmramProfileData (
   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
@@ -704,14 +1239,12 @@ GetSmramProfileData (
 \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
+  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 = (UINTN) CommGetProfileInfo->ProfileSize;\r
@@ -720,10 +1253,10 @@ GetSmramProfileData (
   // Get Data\r
   //\r
   ProfileBuffer = AllocateZeroPool (ProfileSize);\r
-  if (ProfileBuffer == 0) {\r
+  if (ProfileBuffer == NULL) {\r
     Status = EFI_OUT_OF_RESOURCES;\r
     Print (L"SmramProfile: AllocateZeroPool (0x%x) for profile buffer - %r\n", ProfileSize, Status);\r
-    return Status;\r
+    goto Done;\r
   }\r
 \r
   CommHeader = (EFI_SMM_COMMUNICATE_HEADER *) &CommBuffer[0];\r
@@ -752,9 +1285,9 @@ GetSmramProfileData (
     ASSERT_EFI_ERROR (Status);\r
 \r
     if (CommGetProfileData->Header.ReturnStatus != 0) {\r
-      FreePool (ProfileBuffer);\r
+      Status = EFI_SUCCESS;\r
       Print (L"GetProfileData - 0x%x\n", CommGetProfileData->Header.ReturnStatus);\r
-      return EFI_SUCCESS;\r
+      goto Done;\r
     }\r
     CopyMem ((UINT8 *) ProfileBuffer + Offset, (VOID *) (UINTN) CommGetProfileData->ProfileBuffer, (UINTN) CommGetProfileData->ProfileSize);\r
   }\r
@@ -762,21 +1295,52 @@ GetSmramProfileData (
 \r
   Print (L"SmramProfileSize - 0x%x\n", ProfileSize);\r
   Print (L"======= SmramProfile begin =======\n");\r
-  DumpMemoryProfile ((PHYSICAL_ADDRESS) (UINTN) ProfileBuffer, 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 (ProfileBuffer);\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
@@ -788,7 +1352,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
index a1bc675..c512a3f 100644 (file)
@@ -46,9 +46,8 @@
   UefiLib\r
   MemoryAllocationLib\r
   DxeServicesLib\r
-  PeCoffGetEntryPointLib\r
   PrintLib\r
-  \r
+\r
 [Guids]\r
   ## SOMETIMES_CONSUMES   ## GUID # Locate protocol\r
   ## SOMETIMES_CONSUMES   ## GUID # SmiHandlerRegister\r
index df91c14..a88b7f1 100644 (file)
@@ -104,6 +104,7 @@ PrintInfoFromSmm (
   }\r
 \r
   CommBuffer = NULL;\r
+  RealCommSize = 0;\r
   Status = EfiGetSystemConfigurationTable (\r
              &gEdkiiPiSmmCommunicationRegionTableGuid,\r
              (VOID **) &PiSmmCommunicationRegionTable\r
index ed6b557..0be081d 100644 (file)
@@ -238,6 +238,7 @@ SdMmcPciHcEnumerateDevice (
   LIST_ENTRY                          *Link;\r
   LIST_ENTRY                          *NextLink;\r
   SD_MMC_HC_TRB                       *Trb;\r
+  EFI_TPL                             OldTpl;\r
 \r
   Private = (SD_MMC_HC_PRIVATE_DATA*)Context;\r
 \r
@@ -251,6 +252,7 @@ SdMmcPciHcEnumerateDevice (
         //\r
         // Signal all async task events at the slot with EFI_NO_MEDIA status.\r
         //\r
+        OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
         for (Link = GetFirstNode (&Private->Queue);\r
              !IsNull (&Private->Queue, Link);\r
              Link = NextLink) {\r
@@ -263,6 +265,7 @@ SdMmcPciHcEnumerateDevice (
             SdMmcFreeTrb (Trb);\r
           }\r
         }\r
+        gBS->RestoreTPL (OldTpl);\r
         //\r
         // Notify the upper layer the connect state change through ReinstallProtocolInterface.\r
         //\r
@@ -665,7 +668,7 @@ SdMmcPciHcDriverBindingStart (
   //\r
   Status = gBS->CreateEvent (\r
                   EVT_TIMER | EVT_NOTIFY_SIGNAL,\r
-                  TPL_CALLBACK,\r
+                  TPL_NOTIFY,\r
                   ProcessAsyncTaskList,\r
                   Private,\r
                   &Private->TimerEvent\r
@@ -961,7 +964,7 @@ SdMmcPassThruPassThru (
   // Wait async I/O list is empty before execute sync I/O operation.\r
   //\r
   while (TRUE) {\r
-    OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+    OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
     if (IsListEmpty (&Private->Queue)) {\r
       gBS->RestoreTPL (OldTpl);\r
       break;\r
@@ -1273,7 +1276,7 @@ SdMmcPassThruResetDevice (
   //\r
   // Free all async I/O requests in the queue\r
   //\r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
 \r
   for (Link = GetFirstNode (&Private->Queue);\r
        !IsNull (&Private->Queue, Link);\r
index 8978182..b4ff2af 100644 (file)
@@ -1315,7 +1315,7 @@ SdMmcCreateTrb (
   }\r
 \r
   if (Event != NULL) {\r
-    OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+    OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
     InsertTailList (&Private->Queue, &Trb->TrbList);\r
     gBS->RestoreTPL (OldTpl);\r
   }\r
index 5fe710d..fc705e1 100644 (file)
@@ -339,7 +339,7 @@ EmmcSetExtCsd (
   }\r
 \r
   SetExtCsdReq->Signature = EMMC_REQUEST_SIGNATURE;\r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   InsertTailList (&Partition->Queue, &SetExtCsdReq->Link);\r
   gBS->RestoreTPL (OldTpl);\r
   SetExtCsdReq->Packet.SdMmcCmdBlk    = &SetExtCsdReq->SdMmcCmdBlk;\r
@@ -361,7 +361,7 @@ EmmcSetExtCsd (
   if ((Token != NULL) && (Token->Event != NULL)) {\r
     Status = gBS->CreateEvent (\r
                     EVT_NOTIFY_SIGNAL,\r
-                    TPL_CALLBACK,\r
+                    TPL_NOTIFY,\r
                     AsyncIoCallback,\r
                     SetExtCsdReq,\r
                     &SetExtCsdReq->Event\r
@@ -382,7 +382,7 @@ Error:
     // The request and event will be freed in asynchronous callback for success case.\r
     //\r
     if (EFI_ERROR (Status) && (SetExtCsdReq != NULL)) {\r
-      OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&SetExtCsdReq->Link);\r
       gBS->RestoreTPL (OldTpl);\r
       if (SetExtCsdReq->Event != NULL) {\r
@@ -395,7 +395,7 @@ Error:
     // For synchronous operation, free request whatever the execution result is.\r
     //\r
     if (SetExtCsdReq != NULL) {\r
-      OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&SetExtCsdReq->Link);\r
       gBS->RestoreTPL (OldTpl);\r
       FreePool (SetExtCsdReq);\r
@@ -445,7 +445,7 @@ EmmcSetBlkCount (
   }\r
 \r
   SetBlkCntReq->Signature = EMMC_REQUEST_SIGNATURE;\r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   InsertTailList (&Partition->Queue, &SetBlkCntReq->Link);\r
   gBS->RestoreTPL (OldTpl);\r
   SetBlkCntReq->Packet.SdMmcCmdBlk    = &SetBlkCntReq->SdMmcCmdBlk;\r
@@ -463,7 +463,7 @@ EmmcSetBlkCount (
   if ((Token != NULL) && (Token->Event != NULL)) {\r
     Status = gBS->CreateEvent (\r
                     EVT_NOTIFY_SIGNAL,\r
-                    TPL_CALLBACK,\r
+                    TPL_NOTIFY,\r
                     AsyncIoCallback,\r
                     SetBlkCntReq,\r
                     &SetBlkCntReq->Event\r
@@ -484,7 +484,7 @@ Error:
     // The request and event will be freed in asynchronous callback for success case.\r
     //\r
     if (EFI_ERROR (Status) && (SetBlkCntReq != NULL)) {\r
-      OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&SetBlkCntReq->Link);\r
       gBS->RestoreTPL (OldTpl);\r
       if (SetBlkCntReq->Event != NULL) {\r
@@ -497,7 +497,7 @@ Error:
     // For synchronous operation, free request whatever the execution result is.\r
     //\r
     if (SetBlkCntReq != NULL) {\r
-      OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&SetBlkCntReq->Link);\r
       gBS->RestoreTPL (OldTpl);\r
       FreePool (SetBlkCntReq);\r
@@ -562,7 +562,7 @@ EmmcProtocolInOut (
   }\r
 \r
   ProtocolReq->Signature = EMMC_REQUEST_SIGNATURE;\r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   InsertTailList (&Partition->Queue, &ProtocolReq->Link);\r
   gBS->RestoreTPL (OldTpl);\r
   ProtocolReq->Packet.SdMmcCmdBlk    = &ProtocolReq->SdMmcCmdBlk;\r
@@ -596,7 +596,7 @@ EmmcProtocolInOut (
   if ((Token != NULL) && (Token->Event != NULL)) {\r
     Status = gBS->CreateEvent (\r
                     EVT_NOTIFY_SIGNAL,\r
-                    TPL_CALLBACK,\r
+                    TPL_NOTIFY,\r
                     AsyncIoCallback,\r
                     ProtocolReq,\r
                     &ProtocolReq->Event\r
@@ -617,7 +617,7 @@ Error:
     // The request and event will be freed in asynchronous callback for success case.\r
     //\r
     if (EFI_ERROR (Status) && (ProtocolReq != NULL)) {\r
-      OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&ProtocolReq->Link);\r
       gBS->RestoreTPL (OldTpl);\r
       if (ProtocolReq->Event != NULL) {\r
@@ -630,7 +630,7 @@ Error:
     // For synchronous operation, free request whatever the execution result is.\r
     //\r
     if (ProtocolReq != NULL) {\r
-      OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&ProtocolReq->Link);\r
       gBS->RestoreTPL (OldTpl);\r
       FreePool (ProtocolReq);\r
@@ -688,7 +688,7 @@ EmmcRwMultiBlocks (
   }\r
 \r
   RwMultiBlkReq->Signature = EMMC_REQUEST_SIGNATURE;\r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   InsertTailList (&Partition->Queue, &RwMultiBlkReq->Link);\r
   gBS->RestoreTPL (OldTpl);\r
   RwMultiBlkReq->Packet.SdMmcCmdBlk    = &RwMultiBlkReq->SdMmcCmdBlk;\r
@@ -730,7 +730,7 @@ EmmcRwMultiBlocks (
   if ((Token != NULL) && (Token->Event != NULL)) {\r
     Status = gBS->CreateEvent (\r
                     EVT_NOTIFY_SIGNAL,\r
-                    TPL_CALLBACK,\r
+                    TPL_NOTIFY,\r
                     AsyncIoCallback,\r
                     RwMultiBlkReq,\r
                     &RwMultiBlkReq->Event\r
@@ -751,7 +751,7 @@ Error:
     // The request and event will be freed in asynchronous callback for success case.\r
     //\r
     if (EFI_ERROR (Status) && (RwMultiBlkReq != NULL)) {\r
-      OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&RwMultiBlkReq->Link);\r
       gBS->RestoreTPL (OldTpl);\r
       if (RwMultiBlkReq->Event != NULL) {\r
@@ -764,7 +764,7 @@ Error:
     // For synchronous operation, free request whatever the execution result is.\r
     //\r
     if (RwMultiBlkReq != NULL) {\r
-      OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&RwMultiBlkReq->Link);\r
       gBS->RestoreTPL (OldTpl);\r
       FreePool (RwMultiBlkReq);\r
@@ -901,7 +901,7 @@ EmmcReadWrite (
     if (EFI_ERROR (Status)) {\r
       return Status;\r
     }\r
-    DEBUG ((EFI_D_INFO, "Emmc%a(): Lba 0x%x BlkNo 0x%x Event %p with %r\n", IsRead ? "Read" : "Write", Lba, BlockNum, Token->Event, Status));\r
+    DEBUG ((EFI_D_INFO, "Emmc%a(): Part %d Lba 0x%x BlkNo 0x%x Event %p with %r\n", IsRead ? "Read " : "Write", Partition->PartitionType, Lba, BlockNum, Token ? Token->Event : NULL, Status));\r
 \r
     Lba   += BlockNum;\r
     Buffer = (UINT8*)Buffer + BufferSize;\r
@@ -1069,7 +1069,7 @@ EmmcResetEx (
 \r
   Partition = EMMC_PARTITION_DATA_FROM_BLKIO2 (This);\r
 \r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   for (Link = GetFirstNode (&Partition->Queue);\r
        !IsNull (&Partition->Queue, Link);\r
        Link = NextLink) {\r
@@ -1629,7 +1629,7 @@ EmmcEraseBlockStart (
   }\r
 \r
   EraseBlockStart->Signature = EMMC_REQUEST_SIGNATURE;\r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   InsertTailList (&Partition->Queue, &EraseBlockStart->Link);\r
   gBS->RestoreTPL (OldTpl);\r
   EraseBlockStart->Packet.SdMmcCmdBlk    = &EraseBlockStart->SdMmcCmdBlk;\r
@@ -1652,7 +1652,7 @@ EmmcEraseBlockStart (
   if ((Token != NULL) && (Token->Event != NULL)) {\r
     Status = gBS->CreateEvent (\r
                     EVT_NOTIFY_SIGNAL,\r
-                    TPL_CALLBACK,\r
+                    TPL_NOTIFY,\r
                     AsyncIoCallback,\r
                     EraseBlockStart,\r
                     &EraseBlockStart->Event\r
@@ -1673,7 +1673,9 @@ Error:
     // The request and event will be freed in asynchronous callback for success case.\r
     //\r
     if (EFI_ERROR (Status) && (EraseBlockStart != NULL)) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&EraseBlockStart->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       if (EraseBlockStart->Event != NULL) {\r
         gBS->CloseEvent (EraseBlockStart->Event);\r
       }\r
@@ -1684,7 +1686,9 @@ Error:
     // For synchronous operation, free request whatever the execution result is.\r
     //\r
     if (EraseBlockStart != NULL) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&EraseBlockStart->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       FreePool (EraseBlockStart);\r
     }\r
   }\r
@@ -1732,7 +1736,7 @@ EmmcEraseBlockEnd (
   }\r
 \r
   EraseBlockEnd->Signature = EMMC_REQUEST_SIGNATURE;\r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   InsertTailList (&Partition->Queue, &EraseBlockEnd->Link);\r
   gBS->RestoreTPL (OldTpl);\r
   EraseBlockEnd->Packet.SdMmcCmdBlk    = &EraseBlockEnd->SdMmcCmdBlk;\r
@@ -1755,7 +1759,7 @@ EmmcEraseBlockEnd (
   if ((Token != NULL) && (Token->Event != NULL)) {\r
     Status = gBS->CreateEvent (\r
                     EVT_NOTIFY_SIGNAL,\r
-                    TPL_CALLBACK,\r
+                    TPL_NOTIFY,\r
                     AsyncIoCallback,\r
                     EraseBlockEnd,\r
                     &EraseBlockEnd->Event\r
@@ -1776,7 +1780,9 @@ Error:
     // The request and event will be freed in asynchronous callback for success case.\r
     //\r
     if (EFI_ERROR (Status) && (EraseBlockEnd != NULL)) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&EraseBlockEnd->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       if (EraseBlockEnd->Event != NULL) {\r
         gBS->CloseEvent (EraseBlockEnd->Event);\r
       }\r
@@ -1787,7 +1793,9 @@ Error:
     // For synchronous operation, free request whatever the execution result is.\r
     //\r
     if (EraseBlockEnd != NULL) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&EraseBlockEnd->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       FreePool (EraseBlockEnd);\r
     }\r
   }\r
@@ -1833,7 +1841,7 @@ EmmcEraseBlock (
   }\r
 \r
   EraseBlock->Signature = EMMC_REQUEST_SIGNATURE;\r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   InsertTailList (&Partition->Queue, &EraseBlock->Link);\r
   gBS->RestoreTPL (OldTpl);\r
   EraseBlock->Packet.SdMmcCmdBlk    = &EraseBlock->SdMmcCmdBlk;\r
@@ -1850,7 +1858,7 @@ EmmcEraseBlock (
   if ((Token != NULL) && (Token->Event != NULL)) {\r
     Status = gBS->CreateEvent (\r
                     EVT_NOTIFY_SIGNAL,\r
-                    TPL_CALLBACK,\r
+                    TPL_NOTIFY,\r
                     AsyncIoCallback,\r
                     EraseBlock,\r
                     &EraseBlock->Event\r
@@ -1871,7 +1879,9 @@ Error:
     // The request and event will be freed in asynchronous callback for success case.\r
     //\r
     if (EFI_ERROR (Status) && (EraseBlock != NULL)) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&EraseBlock->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       if (EraseBlock->Event != NULL) {\r
         gBS->CloseEvent (EraseBlock->Event);\r
       }\r
@@ -1882,7 +1892,9 @@ Error:
     // For synchronous operation, free request whatever the execution result is.\r
     //\r
     if (EraseBlock != NULL) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&EraseBlock->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       FreePool (EraseBlock);\r
     }\r
   }\r
index d3602a5..5040882 100644 (file)
@@ -443,77 +443,60 @@ InstallProtocolOnPartition (
     //\r
     // Install BlkIo/BlkIo2/Ssp for the specified partition\r
     //\r
-    Status = gBS->InstallMultipleProtocolInterfaces (\r
-                    &Partition->Handle,\r
-                    &gEfiDevicePathProtocolGuid,\r
-                    Partition->DevicePath,\r
-                    &gEfiBlockIoProtocolGuid,\r
-                    &Partition->BlockIo,\r
-                    &gEfiBlockIo2ProtocolGuid,\r
-                    &Partition->BlockIo2,\r
-                    NULL\r
-                    );\r
-    if (EFI_ERROR (Status)) {\r
-      goto Error;\r
-    }\r
-\r
     if (Partition->PartitionType != EmmcPartitionRPMB) {\r
-      Status = gBS->InstallProtocolInterface (\r
+      Status = gBS->InstallMultipleProtocolInterfaces (\r
                       &Partition->Handle,\r
+                      &gEfiDevicePathProtocolGuid,\r
+                      Partition->DevicePath,\r
+                      &gEfiBlockIoProtocolGuid,\r
+                      &Partition->BlockIo,\r
+                      &gEfiBlockIo2ProtocolGuid,\r
+                      &Partition->BlockIo2,\r
                       &gEfiEraseBlockProtocolGuid,\r
-                      EFI_NATIVE_INTERFACE,\r
-                      &Partition->EraseBlock\r
+                      &Partition->EraseBlock,\r
+                      NULL\r
                       );\r
       if (EFI_ERROR (Status)) {\r
-        gBS->UninstallMultipleProtocolInterfaces (\r
-               &Partition->Handle,\r
-               &gEfiDevicePathProtocolGuid,\r
-               Partition->DevicePath,\r
-               &gEfiBlockIoProtocolGuid,\r
-               &Partition->BlockIo,\r
-               &gEfiBlockIo2ProtocolGuid,\r
-               &Partition->BlockIo2,\r
-               NULL\r
-               );\r
         goto Error;\r
       }\r
-    }\r
 \r
-    if (((Partition->PartitionType == EmmcPartitionUserData) ||\r
-        (Partition->PartitionType == EmmcPartitionBoot1) ||\r
-        (Partition->PartitionType == EmmcPartitionBoot2)) &&\r
-        ((Device->Csd.Ccc & BIT10) != 0)) {\r
-      Status = gBS->InstallProtocolInterface (\r
-                      &Partition->Handle,\r
-                      &gEfiStorageSecurityCommandProtocolGuid,\r
-                      EFI_NATIVE_INTERFACE,\r
-                      &Partition->StorageSecurity\r
-                      );\r
-      if (EFI_ERROR (Status)) {\r
-        gBS->UninstallMultipleProtocolInterfaces (\r
-               &Partition->Handle,\r
-               &gEfiDevicePathProtocolGuid,\r
-               Partition->DevicePath,\r
-               &gEfiBlockIoProtocolGuid,\r
-               &Partition->BlockIo,\r
-               &gEfiBlockIo2ProtocolGuid,\r
-               &Partition->BlockIo2,\r
-               &gEfiEraseBlockProtocolGuid,\r
-               &Partition->EraseBlock,\r
-               NULL\r
-               );\r
-        goto Error;\r
+      if (((Partition->PartitionType == EmmcPartitionUserData) ||\r
+          (Partition->PartitionType == EmmcPartitionBoot1) ||\r
+          (Partition->PartitionType == EmmcPartitionBoot2)) &&\r
+          ((Device->Csd.Ccc & BIT10) != 0)) {\r
+        Status = gBS->InstallProtocolInterface (\r
+                        &Partition->Handle,\r
+                        &gEfiStorageSecurityCommandProtocolGuid,\r
+                        EFI_NATIVE_INTERFACE,\r
+                        &Partition->StorageSecurity\r
+                        );\r
+        if (EFI_ERROR (Status)) {\r
+          gBS->UninstallMultipleProtocolInterfaces (\r
+                 &Partition->Handle,\r
+                 &gEfiDevicePathProtocolGuid,\r
+                 Partition->DevicePath,\r
+                 &gEfiBlockIoProtocolGuid,\r
+                 &Partition->BlockIo,\r
+                 &gEfiBlockIo2ProtocolGuid,\r
+                 &Partition->BlockIo2,\r
+                 &gEfiEraseBlockProtocolGuid,\r
+                 &Partition->EraseBlock,\r
+                 NULL\r
+                 );\r
+          goto Error;\r
+        }\r
       }\r
+\r
+      gBS->OpenProtocol (\r
+             Device->Private->Controller,\r
+             &gEfiSdMmcPassThruProtocolGuid,\r
+             (VOID **) &(Device->Private->PassThru),\r
+             Device->Private->DriverBindingHandle,\r
+             Partition->Handle,\r
+             EFI_OPEN_PROTOCOL_BY_CHILD_CONTROLLER\r
+             );    \r
     }\r
 \r
-    gBS->OpenProtocol (\r
-           Device->Private->Controller,\r
-           &gEfiSdMmcPassThruProtocolGuid,\r
-           (VOID **) &(Device->Private->PassThru),\r
-           Device->Private->DriverBindingHandle,\r
-           Partition->Handle,\r
-           EFI_OPEN_PROTOCOL_BY_CHILD_CONTROLLER\r
-           );\r
   } else {\r
     Status = EFI_INVALID_PARAMETER;\r
   }\r
@@ -993,7 +976,6 @@ EmmcDxeDriverBindingStop (
   EFI_BLOCK_IO_PROTOCOL                  *BlockIo;\r
   EFI_BLOCK_IO2_PROTOCOL                 *BlockIo2;\r
   EFI_STORAGE_SECURITY_COMMAND_PROTOCOL  *StorageSecurity;\r
-  EFI_ERASE_BLOCK_PROTOCOL               *EraseBlock;\r
   LIST_ENTRY                             *Link;\r
   LIST_ENTRY                             *NextLink;\r
   EMMC_REQUEST                           *Request;\r
@@ -1120,6 +1102,8 @@ EmmcDxeDriverBindingStop (
                     &Partition->BlockIo,\r
                     &gEfiBlockIo2ProtocolGuid,\r
                     &Partition->BlockIo2,\r
+                    &gEfiEraseBlockProtocolGuid,\r
+                    &Partition->EraseBlock,\r
                     NULL\r
                     );\r
     if (EFI_ERROR (Status)) {\r
@@ -1135,38 +1119,6 @@ EmmcDxeDriverBindingStop (
       continue;\r
     }\r
 \r
-    //\r
-    // If Erase Block Protocol is installed, then uninstall this protocol.\r
-    //\r
-    Status = gBS->OpenProtocol (\r
-                    ChildHandleBuffer[Index],\r
-                    &gEfiEraseBlockProtocolGuid,\r
-                    (VOID **) &EraseBlock,\r
-                    This->DriverBindingHandle,\r
-                    Controller,\r
-                    EFI_OPEN_PROTOCOL_GET_PROTOCOL\r
-                    );\r
-\r
-    if (!EFI_ERROR (Status)) {\r
-      Status = gBS->UninstallProtocolInterface (\r
-                      ChildHandleBuffer[Index],\r
-                      &gEfiEraseBlockProtocolGuid,\r
-                      &Partition->EraseBlock\r
-                      );\r
-      if (EFI_ERROR (Status)) {\r
-        gBS->OpenProtocol (\r
-          Controller,\r
-          &gEfiSdMmcPassThruProtocolGuid,\r
-          (VOID **) &Partition->Device->Private->PassThru,\r
-          This->DriverBindingHandle,\r
-          ChildHandleBuffer[Index],\r
-          EFI_OPEN_PROTOCOL_BY_CHILD_CONTROLLER\r
-          );\r
-        AllChildrenStopped = FALSE;\r
-        continue;\r
-      }\r
-    }\r
-\r
     //\r
     // If Storage Security Command Protocol is installed, then uninstall this protocol.\r
     //\r
index d8b9459..516c3e7 100644 (file)
@@ -340,7 +340,7 @@ SdRwSingleBlock (
   }\r
 \r
   RwSingleBlkReq->Signature = SD_REQUEST_SIGNATURE;\r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   InsertTailList (&Device->Queue, &RwSingleBlkReq->Link);\r
   gBS->RestoreTPL (OldTpl);\r
   RwSingleBlkReq->Packet.SdMmcCmdBlk    = &RwSingleBlkReq->SdMmcCmdBlk;\r
@@ -382,7 +382,7 @@ SdRwSingleBlock (
   if ((Token != NULL) && (Token->Event != NULL)) {\r
     Status = gBS->CreateEvent (\r
                     EVT_NOTIFY_SIGNAL,\r
-                    TPL_CALLBACK,\r
+                    TPL_NOTIFY,\r
                     AsyncIoCallback,\r
                     RwSingleBlkReq,\r
                     &RwSingleBlkReq->Event\r
@@ -403,7 +403,9 @@ Error:
     // The request and event will be freed in asynchronous callback for success case.\r
     //\r
     if (EFI_ERROR (Status) && (RwSingleBlkReq != NULL)) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&RwSingleBlkReq->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       if (RwSingleBlkReq->Event != NULL) {\r
         gBS->CloseEvent (RwSingleBlkReq->Event);\r
       }\r
@@ -414,7 +416,9 @@ Error:
     // For synchronous operation, free request whatever the execution result is.\r
     //\r
     if (RwSingleBlkReq != NULL) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&RwSingleBlkReq->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       FreePool (RwSingleBlkReq);\r
     }\r
   }\r
@@ -468,7 +472,7 @@ SdRwMultiBlocks (
   }\r
 \r
   RwMultiBlkReq->Signature = SD_REQUEST_SIGNATURE;\r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   InsertTailList (&Device->Queue, &RwMultiBlkReq->Link);\r
   gBS->RestoreTPL (OldTpl);\r
   RwMultiBlkReq->Packet.SdMmcCmdBlk    = &RwMultiBlkReq->SdMmcCmdBlk;\r
@@ -510,7 +514,7 @@ SdRwMultiBlocks (
   if ((Token != NULL) && (Token->Event != NULL)) {\r
     Status = gBS->CreateEvent (\r
                     EVT_NOTIFY_SIGNAL,\r
-                    TPL_CALLBACK,\r
+                    TPL_NOTIFY,\r
                     AsyncIoCallback,\r
                     RwMultiBlkReq,\r
                     &RwMultiBlkReq->Event\r
@@ -531,7 +535,9 @@ Error:
     // The request and event will be freed in asynchronous callback for success case.\r
     //\r
     if (EFI_ERROR (Status) && (RwMultiBlkReq != NULL)) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&RwMultiBlkReq->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       if (RwMultiBlkReq->Event != NULL) {\r
         gBS->CloseEvent (RwMultiBlkReq->Event);\r
       }\r
@@ -542,7 +548,9 @@ Error:
     // For synchronous operation, free request whatever the execution result is.\r
     //\r
     if (RwMultiBlkReq != NULL) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&RwMultiBlkReq->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       FreePool (RwMultiBlkReq);\r
     }\r
   }\r
@@ -830,7 +838,7 @@ SdResetEx (
 \r
   Device = SD_DEVICE_DATA_FROM_BLKIO2 (This);\r
 \r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   for (Link = GetFirstNode (&Device->Queue);\r
        !IsNull (&Device->Queue, Link);\r
        Link = NextLink) {\r
@@ -1007,7 +1015,7 @@ SdEraseBlockStart (
   }\r
 \r
   EraseBlockStart->Signature = SD_REQUEST_SIGNATURE;\r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   InsertTailList (&Device->Queue, &EraseBlockStart->Link);\r
   gBS->RestoreTPL (OldTpl);\r
   EraseBlockStart->Packet.SdMmcCmdBlk    = &EraseBlockStart->SdMmcCmdBlk;\r
@@ -1030,7 +1038,7 @@ SdEraseBlockStart (
   if ((Token != NULL) && (Token->Event != NULL)) {\r
     Status = gBS->CreateEvent (\r
                     EVT_NOTIFY_SIGNAL,\r
-                    TPL_CALLBACK,\r
+                    TPL_NOTIFY,\r
                     AsyncIoCallback,\r
                     EraseBlockStart,\r
                     &EraseBlockStart->Event\r
@@ -1051,7 +1059,9 @@ Error:
     // The request and event will be freed in asynchronous callback for success case.\r
     //\r
     if (EFI_ERROR (Status) && (EraseBlockStart != NULL)) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&EraseBlockStart->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       if (EraseBlockStart->Event != NULL) {\r
         gBS->CloseEvent (EraseBlockStart->Event);\r
       }\r
@@ -1062,7 +1072,9 @@ Error:
     // For synchronous operation, free request whatever the execution result is.\r
     //\r
     if (EraseBlockStart != NULL) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&EraseBlockStart->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       FreePool (EraseBlockStart);\r
     }\r
   }\r
@@ -1107,7 +1119,7 @@ SdEraseBlockEnd (
   }\r
 \r
   EraseBlockEnd->Signature = SD_REQUEST_SIGNATURE;\r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   InsertTailList (&Device->Queue, &EraseBlockEnd->Link);\r
   gBS->RestoreTPL (OldTpl);\r
   EraseBlockEnd->Packet.SdMmcCmdBlk    = &EraseBlockEnd->SdMmcCmdBlk;\r
@@ -1130,7 +1142,7 @@ SdEraseBlockEnd (
   if ((Token != NULL) && (Token->Event != NULL)) {\r
     Status = gBS->CreateEvent (\r
                     EVT_NOTIFY_SIGNAL,\r
-                    TPL_CALLBACK,\r
+                    TPL_NOTIFY,\r
                     AsyncIoCallback,\r
                     EraseBlockEnd,\r
                     &EraseBlockEnd->Event\r
@@ -1151,7 +1163,9 @@ Error:
     // The request and event will be freed in asynchronous callback for success case.\r
     //\r
     if (EFI_ERROR (Status) && (EraseBlockEnd != NULL)) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&EraseBlockEnd->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       if (EraseBlockEnd->Event != NULL) {\r
         gBS->CloseEvent (EraseBlockEnd->Event);\r
       }\r
@@ -1162,7 +1176,9 @@ Error:
     // For synchronous operation, free request whatever the execution result is.\r
     //\r
     if (EraseBlockEnd != NULL) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&EraseBlockEnd->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       FreePool (EraseBlockEnd);\r
     }\r
   }\r
@@ -1205,7 +1221,7 @@ SdEraseBlock (
   }\r
 \r
   EraseBlock->Signature = SD_REQUEST_SIGNATURE;\r
-  OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+  OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
   InsertTailList (&Device->Queue, &EraseBlock->Link);\r
   gBS->RestoreTPL (OldTpl);\r
   EraseBlock->Packet.SdMmcCmdBlk    = &EraseBlock->SdMmcCmdBlk;\r
@@ -1222,7 +1238,7 @@ SdEraseBlock (
   if ((Token != NULL) && (Token->Event != NULL)) {\r
     Status = gBS->CreateEvent (\r
                     EVT_NOTIFY_SIGNAL,\r
-                    TPL_CALLBACK,\r
+                    TPL_NOTIFY,\r
                     AsyncIoCallback,\r
                     EraseBlock,\r
                     &EraseBlock->Event\r
@@ -1243,7 +1259,9 @@ Error:
     // The request and event will be freed in asynchronous callback for success case.\r
     //\r
     if (EFI_ERROR (Status) && (EraseBlock != NULL)) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&EraseBlock->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       if (EraseBlock->Event != NULL) {\r
         gBS->CloseEvent (EraseBlock->Event);\r
       }\r
@@ -1254,7 +1272,9 @@ Error:
     // For synchronous operation, free request whatever the execution result is.\r
     //\r
     if (EraseBlock != NULL) {\r
+      OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
       RemoveEntryList (&EraseBlock->Link);\r
+      gBS->RestoreTPL (OldTpl);\r
       FreePool (EraseBlock);\r
     }\r
   }\r
index 7c70c60..0cf9067 100644 (file)
@@ -800,7 +800,7 @@ SdDxeDriverBindingStop (
     //\r
     // Free all on-going async tasks.\r
     //\r
-    OldTpl = gBS->RaiseTPL (TPL_CALLBACK);\r
+    OldTpl = gBS->RaiseTPL (TPL_NOTIFY);\r
     for (Link = GetFirstNode (&Device->Queue);\r
          !IsNull (&Device->Queue, Link);\r
          Link = NextLink) {\r
index e6b9114..743221f 100644 (file)
@@ -2780,11 +2780,13 @@ MemoryProfileInstallProtocol (
   @param DriverEntry    Image info.\r
   @param FileType       Image file type.\r
 \r
-  @retval TRUE          Register success.\r
-  @retval FALSE         Register fail.\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 RegisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,\r
   IN EFI_FV_FILETYPE            FileType\r
@@ -2795,11 +2797,13 @@ RegisterMemoryProfileImage (
 \r
   @param DriverEntry    Image info.\r
 \r
-  @retval TRUE          Unregister success.\r
-  @retval FALSE         Unregister fail.\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_NOT_FOUND         The image is not found.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 UnregisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry\r
   );\r
@@ -2810,20 +2814,31 @@ UnregisterMemoryProfileImage (
   @param CallerAddress  Address of caller who call Allocate or Free.\r
   @param Action         This Allocate or Free action.\r
   @param MemoryType     Memory type.\r
+                        EfiMaxMemoryType means the MemoryType is unknown.\r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
+  @param ActionString   String for memory profile action.\r
+                        Only needed for user defined allocate action.\r
 \r
-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
+EFIAPI\r
 CoreUpdateProfile (\r
   IN EFI_PHYSICAL_ADDRESS   CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType,\r
   IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool\r
-  IN VOID                   *Buffer\r
+  IN VOID                   *Buffer,\r
+  IN CHAR8                  *ActionString OPTIONAL\r
   );\r
 \r
 /**\r
index e3e4d03..450da57 100644 (file)
   gEfiMdeModulePkgTokenSpaceGuid.PcdMaxEfiSystemTablePointerAddress         ## CONSUMES\r
   gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfileMemoryType                 ## CONSUMES\r
   gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfilePropertyMask               ## CONSUMES\r
+  gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfileDriverPath                 ## CONSUMES\r
   gEfiMdeModulePkgTokenSpaceGuid.PcdPropertiesTableEnable                   ## CONSUMES\r
 \r
 # [Hob]\r
index 6626e10..14c4959 100644 (file)
@@ -17,6 +17,9 @@
 \r
 #define IS_UEFI_MEMORY_PROFILE_ENABLED ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT0) != 0)\r
 \r
+#define GET_OCCUPIED_SIZE(ActualSize, Alignment) \\r
+  ((ActualSize) + (((Alignment) - ((ActualSize) & ((Alignment) - 1))) & ((Alignment) - 1)))\r
+\r
 typedef struct {\r
   UINT32                        Signature;\r
   MEMORY_PROFILE_CONTEXT        Context;\r
@@ -27,12 +30,14 @@ typedef struct {
   UINT32                        Signature;\r
   MEMORY_PROFILE_DRIVER_INFO    DriverInfo;\r
   LIST_ENTRY                    *AllocInfoList;\r
+  CHAR8                         *PdbString;\r
   LIST_ENTRY                    Link;\r
 } MEMORY_PROFILE_DRIVER_INFO_DATA;\r
 \r
 typedef struct {\r
   UINT32                        Signature;\r
   MEMORY_PROFILE_ALLOC_INFO     AllocInfo;\r
+  CHAR8                         *ActionString;\r
   LIST_ENTRY                    Link;\r
 } MEMORY_PROFILE_ALLOC_INFO_DATA;\r
 \r
@@ -58,7 +63,11 @@ GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA mMemoryProfileContext
 };\r
 GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA *mMemoryProfileContextPtr = NULL;\r
 \r
-BOOLEAN mMemoryProfileRecordingStatus = FALSE;\r
+EFI_LOCK mMemoryProfileLock = EFI_INITIALIZE_LOCK_VARIABLE (TPL_NOTIFY);\r
+BOOLEAN mMemoryProfileGettingStatus = FALSE;\r
+BOOLEAN mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;\r
+EFI_DEVICE_PATH_PROTOCOL *mMemoryProfileDriverPath;\r
+UINTN                    mMemoryProfileDriverPathSize;\r
 \r
 /**\r
   Get memory profile data.\r
@@ -69,6 +78,7 @@ BOOLEAN mMemoryProfileRecordingStatus = FALSE;
   @param[out]     ProfileBuffer     Profile buffer.\r
                       \r
   @return EFI_SUCCESS               Get the memory profile data successfully.\r
+  @return EFI_UNSUPPORTED           Memory profile is unsupported.\r
   @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data. \r
                                     ProfileSize is updated with the size required.\r
 \r
@@ -90,7 +100,9 @@ ProfileProtocolGetData (
   @param[in] ImageSize          Image size.\r
   @param[in] FileType           File type of the image.\r
 \r
-  @return EFI_SUCCESS           Register success.\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
   @return EFI_OUT_OF_RESOURCE   No enough resource for this register.\r
 \r
 **/\r
@@ -112,7 +124,9 @@ ProfileProtocolRegisterImage (
   @param[in] ImageBase          Image base address.\r
   @param[in] ImageSize          Image size.\r
 \r
-  @return EFI_SUCCESS           Unregister success.\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
   @return EFI_NOT_FOUND         The image is not found.\r
 \r
 **/\r
@@ -125,12 +139,107 @@ ProfileProtocolUnregisterImage (
   IN UINT64                             ImageSize\r
   );\r
 \r
+/**\r
+  Get memory profile recording state.\r
+\r
+  @param[in]  This              The EDKII_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[out] RecordingState    Recording state.\r
+\r
+  @return EFI_SUCCESS           Memory profile recording state is returned.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+  @return EFI_INVALID_PARAMETER RecordingState is NULL.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ProfileProtocolGetRecordingState (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,\r
+  OUT BOOLEAN                           *RecordingState\r
+  );\r
+\r
+/**\r
+  Set memory profile recording state.\r
+\r
+  @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] RecordingState     Recording state.\r
+\r
+  @return EFI_SUCCESS           Set memory profile recording state successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ProfileProtocolSetRecordingState (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,\r
+  IN BOOLEAN                            RecordingState\r
+  );\r
+\r
+/**\r
+  Record memory profile of multilevel caller.\r
+\r
+  @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] CallerAddress      Address of caller.\r
+  @param[in] Action             Memory profile action.\r
+  @param[in] MemoryType         Memory type.\r
+                                EfiMaxMemoryType means the MemoryType is unknown.\r
+  @param[in] Buffer             Buffer address.\r
+  @param[in] Size               Buffer size.\r
+  @param[in] ActionString       String for memory profile action.\r
+                                Only needed for user defined allocate action.\r
+\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ProfileProtocolRecord (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,\r
+  IN PHYSICAL_ADDRESS                   CallerAddress,\r
+  IN MEMORY_PROFILE_ACTION              Action,\r
+  IN EFI_MEMORY_TYPE                    MemoryType,\r
+  IN VOID                               *Buffer,\r
+  IN UINTN                              Size,\r
+  IN CHAR8                              *ActionString OPTIONAL\r
+  );\r
+\r
 EDKII_MEMORY_PROFILE_PROTOCOL mProfileProtocol = {\r
   ProfileProtocolGetData,\r
   ProfileProtocolRegisterImage,\r
-  ProfileProtocolUnregisterImage\r
+  ProfileProtocolUnregisterImage,\r
+  ProfileProtocolGetRecordingState,\r
+  ProfileProtocolSetRecordingState,\r
+  ProfileProtocolRecord,\r
 };\r
 \r
+/**\r
+  Acquire lock on mMemoryProfileLock.\r
+**/\r
+VOID\r
+CoreAcquireMemoryProfileLock (\r
+  VOID\r
+  )\r
+{\r
+  CoreAcquireLock (&mMemoryProfileLock);\r
+}\r
+\r
+/**\r
+  Release lock on mMemoryProfileLock.\r
+**/\r
+VOID\r
+CoreReleaseMemoryProfileLock (\r
+  VOID\r
+  )\r
+{\r
+  CoreReleaseLock (&mMemoryProfileLock);\r
+}\r
+\r
 /**\r
   Return memory profile context.\r
 \r
@@ -307,13 +416,27 @@ BuildDriverInfo (
   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   VOID                              *EntryPointInImage;\r
+  CHAR8                             *PdbString;\r
+  UINTN                             PdbSize;\r
+  UINTN                             PdbOccupiedSize;\r
+\r
+  PdbSize = 0;\r
+  PdbOccupiedSize = 0;\r
+  PdbString = NULL;\r
+  if (ImageBase != 0) {\r
+    PdbString = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageBase);\r
+    if (PdbString != NULL) {\r
+      PdbSize = AsciiStrSize (PdbString);\r
+      PdbOccupiedSize = GET_OCCUPIED_SIZE (PdbSize, sizeof (UINT64));\r
+    }\r
+  }\r
 \r
   //\r
   // Use CoreInternalAllocatePool() that will not update profile for this AllocatePool action.\r
   //\r
   Status = CoreInternalAllocatePool (\r
              EfiBootServicesData,\r
-             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY),\r
+             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY) + PdbSize,\r
              (VOID **) &DriverInfoData\r
              );\r
   if (EFI_ERROR (Status)) {\r
@@ -325,7 +448,7 @@ BuildDriverInfo (
   DriverInfo = &DriverInfoData->DriverInfo;\r
   DriverInfoData->Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;\r
   DriverInfo->Header.Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;\r
-  DriverInfo->Header.Length = sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
+  DriverInfo->Header.Length = (UINT16) (sizeof (MEMORY_PROFILE_DRIVER_INFO) + PdbOccupiedSize);\r
   DriverInfo->Header.Revision = MEMORY_PROFILE_DRIVER_INFO_REVISION;\r
   if (FileName != NULL) {\r
     CopyMem (&DriverInfo->FileName, FileName, sizeof (EFI_GUID));\r
@@ -349,6 +472,14 @@ BuildDriverInfo (
   DriverInfo->CurrentUsage = 0;\r
   DriverInfo->PeakUsage = 0;\r
   DriverInfo->AllocRecordCount = 0;\r
+  if (PdbSize != 0) {\r
+    DriverInfo->PdbStringOffset = (UINT16) sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
+    DriverInfoData->PdbString = (CHAR8 *) (DriverInfoData->AllocInfoList + 1);\r
+    CopyMem (DriverInfoData->PdbString, PdbString, PdbSize);\r
+  } else {\r
+    DriverInfo->PdbStringOffset = 0;\r
+    DriverInfoData->PdbString = NULL;\r
+  }\r
 \r
   InsertTailList (ContextData->DriverInfoList, &DriverInfoData->Link);\r
   ContextData->Context.ImageCount ++;\r
@@ -357,6 +488,65 @@ BuildDriverInfo (
   return DriverInfoData;\r
 }\r
 \r
+/**\r
+  Return if record for this driver is needed..\r
+\r
+  @param DriverFilePath     Driver file path.\r
+\r
+  @retval TRUE              Record for this driver is needed.\r
+  @retval FALSE             Record for this driver is not needed.\r
+\r
+**/\r
+BOOLEAN\r
+NeedRecordThisDriver (\r
+  IN EFI_DEVICE_PATH_PROTOCOL       *DriverFilePath\r
+  )\r
+{\r
+  EFI_DEVICE_PATH_PROTOCOL          *TmpDevicePath;\r
+  EFI_DEVICE_PATH_PROTOCOL          *DevicePathInstance;\r
+  UINTN                             DevicePathSize;\r
+  UINTN                             FilePathSize;\r
+\r
+  if (!IsDevicePathValid (mMemoryProfileDriverPath, mMemoryProfileDriverPathSize)) {\r
+    //\r
+    // Invalid Device Path means record all.\r
+    //\r
+    return TRUE;\r
+  }\r
+  \r
+  //\r
+  // Record FilePath without END node.\r
+  //\r
+  FilePathSize = GetDevicePathSize (DriverFilePath) - sizeof(EFI_DEVICE_PATH_PROTOCOL);\r
+\r
+  DevicePathInstance = mMemoryProfileDriverPath;\r
+  do {\r
+    //\r
+    // Find END node (it might be END_ENTIRE or END_INSTANCE).\r
+    //\r
+    TmpDevicePath = DevicePathInstance;\r
+    while (!IsDevicePathEndType (TmpDevicePath)) {\r
+      TmpDevicePath = NextDevicePathNode (TmpDevicePath);\r
+    }\r
+\r
+    //\r
+    // Do not compare END node.\r
+    //\r
+    DevicePathSize = (UINTN)TmpDevicePath - (UINTN)DevicePathInstance;\r
+    if ((FilePathSize == DevicePathSize) &&\r
+        (CompareMem (DriverFilePath, DevicePathInstance, DevicePathSize) == 0)) {\r
+      return TRUE;\r
+    }\r
+\r
+    //\r
+    // Get next instance.\r
+    //\r
+    DevicePathInstance = (EFI_DEVICE_PATH_PROTOCOL *)((UINTN)DevicePathInstance + DevicePathSize + DevicePathNodeLength(TmpDevicePath));\r
+  } while (DevicePathSubType (TmpDevicePath) != END_ENTIRE_DEVICE_PATH_SUBTYPE);\r
+\r
+  return FALSE;\r
+}\r
+\r
 /**\r
   Register DXE Core to memory profile.\r
 \r
@@ -376,6 +566,8 @@ RegisterDxeCore (
   EFI_PEI_HOB_POINTERS              DxeCoreHob;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   PHYSICAL_ADDRESS                  ImageBase;\r
+  UINT8                             TempBuffer[sizeof(MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof(EFI_DEVICE_PATH_PROTOCOL)];\r
+  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;\r
 \r
   ASSERT (ContextData != NULL);\r
 \r
@@ -394,6 +586,14 @@ RegisterDxeCore (
   }\r
   ASSERT (DxeCoreHob.Raw != NULL);\r
 \r
+  FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) TempBuffer;\r
+  EfiInitializeFwVolDevicepathNode (FilePath, &DxeCoreHob.MemoryAllocationModule->ModuleName);\r
+  SetDevicePathEndNode (FilePath + 1);\r
+\r
+  if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *) FilePath)) {\r
+    return FALSE;\r
+  }\r
+\r
   ImageBase = DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryBaseAddress;\r
   DriverInfoData = BuildDriverInfo (\r
                      ContextData,\r
@@ -433,7 +633,14 @@ MemoryProfileInit (
     return;\r
   }\r
 \r
-  mMemoryProfileRecordingStatus = TRUE;\r
+  mMemoryProfileGettingStatus = FALSE;\r
+  if ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT7) != 0) {\r
+    mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;\r
+  } else {\r
+    mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_ENABLE;\r
+  }\r
+  mMemoryProfileDriverPathSize = PcdGetSize (PcdMemoryProfileDriverPath);\r
+  mMemoryProfileDriverPath = AllocateCopyPool (mMemoryProfileDriverPathSize, PcdGetPtr (PcdMemoryProfileDriverPath));\r
   mMemoryProfileContextPtr = &mMemoryProfileContext;\r
 \r
   RegisterDxeCore (HobStart, &mMemoryProfileContext);\r
@@ -504,11 +711,13 @@ GetFileNameFromFilePath (
   @param DriverEntry    Image info.\r
   @param FileType       Image file type.\r
 \r
-  @retval TRUE          Register success.\r
-  @retval FALSE         Register fail.\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 RegisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,\r
   IN EFI_FV_FILETYPE            FileType\r
@@ -518,12 +727,16 @@ RegisterMemoryProfileImage (
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  if (!NeedRecordThisDriver (DriverEntry->Info.FilePath)) {\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = BuildDriverInfo (\r
@@ -536,10 +749,10 @@ RegisterMemoryProfileImage (
                      FileType\r
                      );\r
   if (DriverInfoData == NULL) {\r
-    return FALSE;\r
+    return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
@@ -586,45 +799,9 @@ GetMemoryProfileDriverInfoByFileNameAndAddress (
   return NULL;\r
 }\r
 \r
-/**\r
-  Search dummy image from memory profile.\r
-\r
-  @param ContextData    Memory profile context.\r
-\r
-  @return Pointer to memory profile driver info.\r
-\r
-**/\r
-MEMORY_PROFILE_DRIVER_INFO_DATA *\r
-FindDummyImage (\r
-  IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData\r
-  )\r
-{\r
-  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
-  LIST_ENTRY                        *DriverLink;\r
-  LIST_ENTRY                        *DriverInfoList;\r
-\r
-  DriverInfoList = ContextData->DriverInfoList;\r
-\r
-  for (DriverLink = DriverInfoList->ForwardLink;\r
-       DriverLink != DriverInfoList;\r
-       DriverLink = DriverLink->ForwardLink) {\r
-    DriverInfoData = CR (\r
-                   DriverLink,\r
-                   MEMORY_PROFILE_DRIVER_INFO_DATA,\r
-                   Link,\r
-                   MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
-                   );\r
-    if (CompareGuid (&gZeroGuid, &DriverInfoData->DriverInfo.FileName)) {\r
-      return DriverInfoData;\r
-    }\r
-  }\r
-\r
-  return BuildDriverInfo (ContextData, &gZeroGuid, 0, 0, 0, 0, 0);\r
-}\r
-\r
 /**\r
   Search image from memory profile.\r
-  It will return image, if (Address >= ImageBuffer) AND (Address < ImageBuffer + ImageSize)\r
+  It will return image, if (Address >= ImageBuffer) AND (Address < ImageBuffer + ImageSize).\r
 \r
   @param ContextData    Memory profile context.\r
   @param Address        Image or Function address.\r
@@ -661,10 +838,7 @@ GetMemoryProfileDriverInfoFromAddress (
     }\r
   }\r
 \r
-  //\r
-  // Should never come here.\r
-  //\r
-  return FindDummyImage (ContextData);\r
+  return NULL;\r
 }\r
 \r
 /**\r
@@ -672,11 +846,13 @@ GetMemoryProfileDriverInfoFromAddress (
 \r
   @param DriverEntry    Image info.\r
 \r
-  @retval TRUE          Unregister success.\r
-  @retval FALSE         Unregister fail.\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_NOT_FOUND         The image is not found.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 UnregisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA      *DriverEntry\r
   )\r
@@ -689,12 +865,16 @@ UnregisterMemoryProfileImage (
   VOID                              *EntryPointInImage;\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  if (!NeedRecordThisDriver (DriverEntry->Info.FilePath)) {\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = NULL;\r
@@ -716,12 +896,13 @@ UnregisterMemoryProfileImage (
     DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, ImageAddress);\r
   }\r
   if (DriverInfoData == NULL) {\r
-    return FALSE;\r
+    return EFI_NOT_FOUND;\r
   }\r
 \r
   ContextData->Context.TotalImageSize -= DriverInfoData->DriverInfo.ImageSize;\r
 \r
-  DriverInfoData->DriverInfo.ImageBase = 0;\r
+  // Keep the ImageBase for RVA calculation in Application.\r
+  //DriverInfoData->DriverInfo.ImageBase = 0;\r
   DriverInfoData->DriverInfo.ImageSize = 0;\r
 \r
   if (DriverInfoData->DriverInfo.PeakUsage == 0) {\r
@@ -733,7 +914,7 @@ UnregisterMemoryProfileImage (
     CoreInternalFreePool (DriverInfoData, NULL);\r
   }\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
@@ -803,55 +984,80 @@ GetProfileMemoryIndex (
   @param MemoryType     Memory type.\r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
+  @param ActionString   String for memory profile action.\r
 \r
-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 CoreUpdateProfileAllocate (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType,\r
   IN UINTN                  Size,\r
-  IN VOID                   *Buffer\r
+  IN VOID                   *Buffer,\r
+  IN CHAR8                  *ActionString OPTIONAL\r
   )\r
 {\r
   EFI_STATUS                        Status;\r
-  MEMORY_PROFILE_CONTEXT           *Context;\r
-  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;\r
-  MEMORY_PROFILE_ALLOC_INFO        *AllocInfo;\r
+  MEMORY_PROFILE_CONTEXT            *Context;\r
+  MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;\r
+  MEMORY_PROFILE_ALLOC_INFO         *AllocInfo;\r
   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;\r
   UINTN                             ProfileMemoryIndex;\r
+  MEMORY_PROFILE_ACTION             BasicAction;\r
+  UINTN                             ActionStringSize;\r
+  UINTN                             ActionStringOccupiedSize;\r
 \r
-  AllocInfoData = NULL;\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);\r
-  ASSERT (DriverInfoData != NULL);\r
+  if (DriverInfoData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  ActionStringSize = 0;\r
+  ActionStringOccupiedSize = 0;\r
+  if (ActionString != NULL) {\r
+    ActionStringSize = AsciiStrSize (ActionString);\r
+    ActionStringOccupiedSize = GET_OCCUPIED_SIZE (ActionStringSize, sizeof (UINT64));\r
+  }\r
 \r
   //\r
   // Use CoreInternalAllocatePool() that will not update profile for this AllocatePool action.\r
   //\r
+  AllocInfoData = NULL;\r
   Status = CoreInternalAllocatePool (\r
              EfiBootServicesData,\r
-             sizeof (*AllocInfoData),\r
+             sizeof (*AllocInfoData) + ActionStringSize,\r
              (VOID **) &AllocInfoData\r
              );\r
   if (EFI_ERROR (Status)) {\r
-    return FALSE;\r
+    return EFI_OUT_OF_RESOURCES;\r
   }\r
   ASSERT (AllocInfoData != NULL);\r
+\r
+  //\r
+  // Only update SequenceCount if and only if it is basic action.\r
+  //\r
+  if (Action == BasicAction) {\r
+    ContextData->Context.SequenceCount ++;\r
+  }\r
+\r
   AllocInfo = &AllocInfoData->AllocInfo;\r
   AllocInfoData->Signature      = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;\r
   AllocInfo->Header.Signature   = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;\r
-  AllocInfo->Header.Length      = sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
+  AllocInfo->Header.Length      = (UINT16) (sizeof (MEMORY_PROFILE_ALLOC_INFO) + ActionStringOccupiedSize);\r
   AllocInfo->Header.Revision    = MEMORY_PROFILE_ALLOC_INFO_REVISION;\r
   AllocInfo->CallerAddress      = CallerAddress;\r
   AllocInfo->SequenceId         = ContextData->Context.SequenceCount;\r
@@ -859,50 +1065,64 @@ CoreUpdateProfileAllocate (
   AllocInfo->MemoryType         = MemoryType;\r
   AllocInfo->Buffer             = (PHYSICAL_ADDRESS) (UINTN) Buffer;\r
   AllocInfo->Size               = Size;\r
+  if (ActionString != NULL) {\r
+    AllocInfo->ActionStringOffset = (UINT16) sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
+    AllocInfoData->ActionString = (CHAR8 *) (AllocInfoData + 1);\r
+    CopyMem (AllocInfoData->ActionString, ActionString, ActionStringSize);\r
+  } else {\r
+    AllocInfo->ActionStringOffset = 0;\r
+    AllocInfoData->ActionString = NULL;\r
+  }\r
 \r
   InsertTailList (DriverInfoData->AllocInfoList, &AllocInfoData->Link);\r
 \r
-  ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);\r
-\r
+  Context = &ContextData->Context;\r
   DriverInfo = &DriverInfoData->DriverInfo;\r
-  DriverInfo->CurrentUsage += Size;\r
-  if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {\r
-    DriverInfo->PeakUsage = DriverInfo->CurrentUsage;\r
-  }\r
-  DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;\r
-  if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {\r
-    DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];\r
-  }\r
   DriverInfo->AllocRecordCount ++;\r
 \r
-  Context = &ContextData->Context;\r
-  Context->CurrentTotalUsage += Size;\r
-  if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {\r
-    Context->PeakTotalUsage = Context->CurrentTotalUsage;\r
-  }\r
-  Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;\r
-  if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {\r
-    Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];\r
+  //\r
+  // Update summary if and only if it is basic action.\r
+  //\r
+  if (Action == BasicAction) {\r
+    ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);\r
+\r
+    DriverInfo->CurrentUsage += Size;\r
+    if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {\r
+      DriverInfo->PeakUsage = DriverInfo->CurrentUsage;\r
+    }\r
+    DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;\r
+    if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {\r
+      DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];\r
+    }\r
+\r
+    Context->CurrentTotalUsage += Size;\r
+    if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {\r
+      Context->PeakTotalUsage = Context->CurrentTotalUsage;\r
+    }\r
+    Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;\r
+    if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {\r
+      Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];\r
+    }\r
   }\r
-  Context->SequenceCount ++;\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
-  Get memory profile alloc info from memory profile\r
+  Get memory profile alloc info from memory profile.\r
 \r
-  @param DriverInfoData     Driver info\r
-  @param Action             This Free action\r
-  @param Size               Buffer size\r
-  @param Buffer             Buffer address\r
+  @param DriverInfoData     Driver info.\r
+  @param BasicAction        This Free basic action.\r
+  @param Size               Buffer size.\r
+  @param Buffer             Buffer address.\r
 \r
   @return Pointer to memory profile alloc info.\r
+\r
 **/\r
 MEMORY_PROFILE_ALLOC_INFO_DATA *\r
 GetMemoryProfileAllocInfoFromAddress (\r
   IN MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData,\r
-  IN MEMORY_PROFILE_ACTION              Action,\r
+  IN MEMORY_PROFILE_ACTION              BasicAction,\r
   IN UINTN                              Size,\r
   IN VOID                               *Buffer\r
   )\r
@@ -924,10 +1144,10 @@ GetMemoryProfileAllocInfoFromAddress (
                       MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
                       );\r
     AllocInfo = &AllocInfoData->AllocInfo;\r
-    if (AllocInfo->Action != Action) {\r
+    if ((AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK) != BasicAction) {\r
       continue;\r
     }\r
-    switch (Action) {\r
+    switch (BasicAction) {\r
       case MemoryProfileActionAllocatePages:\r
         if ((AllocInfo->Buffer <= (PHYSICAL_ADDRESS) (UINTN) Buffer) &&\r
             ((AllocInfo->Buffer + AllocInfo->Size) >= ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size))) {\r
@@ -956,11 +1176,12 @@ GetMemoryProfileAllocInfoFromAddress (
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
 \r
-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 CoreUpdateProfileFree (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
@@ -978,114 +1199,140 @@ CoreUpdateProfileFree (
   MEMORY_PROFILE_DRIVER_INFO_DATA  *ThisDriverInfoData;\r
   MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;\r
   UINTN                            ProfileMemoryIndex;\r
+  MEMORY_PROFILE_ACTION            BasicAction;\r
+  BOOLEAN                          Found;\r
+\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);\r
-  ASSERT (DriverInfoData != NULL);\r
 \r
-  switch (Action) {\r
-    case MemoryProfileActionFreePages:\r
-      AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
-      break;\r
-    case MemoryProfileActionFreePool:\r
-      AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
-      break;\r
-    default:\r
-      ASSERT (FALSE);\r
-      AllocInfoData = NULL;\r
-      break;\r
-  }\r
-  if (AllocInfoData == NULL) {\r
-    //\r
-    // Legal case, because driver A might free memory allocated by driver B, by some protocol.\r
-    //\r
-    DriverInfoList = ContextData->DriverInfoList;\r
-\r
-    for (DriverLink = DriverInfoList->ForwardLink;\r
-         DriverLink != DriverInfoList;\r
-         DriverLink = DriverLink->ForwardLink) {\r
-      ThisDriverInfoData = CR (\r
-                             DriverLink,\r
-                             MEMORY_PROFILE_DRIVER_INFO_DATA,\r
-                             Link,\r
-                             MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
-                             );\r
-      switch (Action) {\r
+  //\r
+  // Do not return if DriverInfoData == NULL here,\r
+  // because driver A might free memory allocated by driver B.\r
+  //\r
+\r
+  //\r
+  // Need use do-while loop to find all possible records,\r
+  // because one address might be recorded multiple times.\r
+  //\r
+  Found = FALSE;\r
+  AllocInfoData = NULL;\r
+  do {\r
+    if (DriverInfoData != NULL) {\r
+      switch (BasicAction) {\r
         case MemoryProfileActionFreePages:\r
-          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
+          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
           break;\r
         case MemoryProfileActionFreePool:\r
-          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
+          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
           break;\r
         default:\r
           ASSERT (FALSE);\r
           AllocInfoData = NULL;\r
           break;\r
       }\r
-      if (AllocInfoData != NULL) {\r
-        DriverInfoData = ThisDriverInfoData;\r
-        break;\r
-      }\r
     }\r
-\r
     if (AllocInfoData == NULL) {\r
       //\r
-      // No matched allocate operation is found for this free operation.\r
-      // It is because the specified memory type allocate operation has been\r
-      // filtered by CoreNeedRecordProfile(), but free operations have no\r
-      // memory type information, they can not be filtered by CoreNeedRecordProfile().\r
-      // Then, they will be filtered here.\r
+      // Legal case, because driver A might free memory allocated by driver B, by some protocol.\r
       //\r
-      return FALSE;\r
-    }\r
-  }\r
+      DriverInfoList = ContextData->DriverInfoList;\r
+\r
+      for (DriverLink = DriverInfoList->ForwardLink;\r
+           DriverLink != DriverInfoList;\r
+           DriverLink = DriverLink->ForwardLink) {\r
+        ThisDriverInfoData = CR (\r
+                               DriverLink,\r
+                               MEMORY_PROFILE_DRIVER_INFO_DATA,\r
+                               Link,\r
+                               MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
+                               );\r
+        switch (BasicAction) {\r
+          case MemoryProfileActionFreePages:\r
+            AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
+            break;\r
+          case MemoryProfileActionFreePool:\r
+            AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
+            break;\r
+          default:\r
+            ASSERT (FALSE);\r
+            AllocInfoData = NULL;\r
+            break;\r
+        }\r
+        if (AllocInfoData != NULL) {\r
+          DriverInfoData = ThisDriverInfoData;\r
+          break;\r
+        }\r
+      }\r
 \r
-  Context = &ContextData->Context;\r
-  DriverInfo = &DriverInfoData->DriverInfo;\r
-  AllocInfo = &AllocInfoData->AllocInfo;\r
+      if (AllocInfoData == NULL) {\r
+        //\r
+        // If (!Found), no matched allocate info is found for this free action.\r
+        // It is because the specified memory type allocate actions have been filtered by\r
+        // CoreNeedRecordProfile(), but free actions may have no memory type information,\r
+        // they can not be filtered by CoreNeedRecordProfile(). Then, they will be\r
+        // filtered here.\r
+        //\r
+        // If (Found), it is normal exit path.\r
+        return (Found ? EFI_SUCCESS : EFI_NOT_FOUND);\r
+      }\r
+    }\r
 \r
-  ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);\r
+    Found = TRUE;\r
 \r
-  Context->CurrentTotalUsage -= AllocInfo->Size;\r
-  Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
+    Context = &ContextData->Context;\r
+    DriverInfo = &DriverInfoData->DriverInfo;\r
+    AllocInfo = &AllocInfoData->AllocInfo;\r
 \r
-  DriverInfo->CurrentUsage -= AllocInfo->Size;\r
-  DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
-  DriverInfo->AllocRecordCount --;\r
+    DriverInfo->AllocRecordCount --;\r
+    //\r
+    // Update summary if and only if it is basic action.\r
+    //\r
+    if (AllocInfo->Action == (AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK)) {\r
+      ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);\r
 \r
-  RemoveEntryList (&AllocInfoData->Link);\r
+      Context->CurrentTotalUsage -= AllocInfo->Size;\r
+      Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
 \r
-  if (Action == MemoryProfileActionFreePages) {\r
-    if (AllocInfo->Buffer != (PHYSICAL_ADDRESS) (UINTN) Buffer) {\r
-      CoreUpdateProfileAllocate (\r
-        AllocInfo->CallerAddress,\r
-        MemoryProfileActionAllocatePages,\r
-        AllocInfo->MemoryType,\r
-        (UINTN) ((PHYSICAL_ADDRESS) (UINTN) Buffer - AllocInfo->Buffer),\r
-        (VOID *) (UINTN) AllocInfo->Buffer\r
-        );\r
-    }\r
-    if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)) {\r
-      CoreUpdateProfileAllocate (\r
-        AllocInfo->CallerAddress,\r
-        MemoryProfileActionAllocatePages,\r
-        AllocInfo->MemoryType,\r
-        (UINTN) ((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)),\r
-        (VOID *) ((UINTN) Buffer + Size)\r
-        );\r
+      DriverInfo->CurrentUsage -= AllocInfo->Size;\r
+      DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
     }\r
-  }\r
 \r
-  //\r
-  // Use CoreInternalFreePool() that will not update profile for this FreePool action.\r
-  //\r
-  CoreInternalFreePool (AllocInfoData, NULL);\r
+    RemoveEntryList (&AllocInfoData->Link);\r
+\r
+    if (BasicAction == MemoryProfileActionFreePages) {\r
+      if (AllocInfo->Buffer != (PHYSICAL_ADDRESS) (UINTN) Buffer) {\r
+        CoreUpdateProfileAllocate (\r
+          AllocInfo->CallerAddress,\r
+          AllocInfo->Action,\r
+          AllocInfo->MemoryType,\r
+          (UINTN) ((PHYSICAL_ADDRESS) (UINTN) Buffer - AllocInfo->Buffer),\r
+          (VOID *) (UINTN) AllocInfo->Buffer,\r
+          AllocInfoData->ActionString\r
+          );\r
+      }\r
+      if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)) {\r
+        CoreUpdateProfileAllocate (\r
+          AllocInfo->CallerAddress,\r
+          AllocInfo->Action,\r
+          AllocInfo->MemoryType,\r
+          (UINTN) ((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)),\r
+          (VOID *) ((UINTN) Buffer + Size),\r
+          AllocInfoData->ActionString\r
+          );\r
+      }\r
+    }\r
 \r
-  return TRUE;\r
+    //\r
+    // Use CoreInternalFreePool() that will not update profile for this FreePool action.\r
+    //\r
+    CoreInternalFreePool (AllocInfoData, NULL);\r
+  } while (TRUE);\r
 }\r
 \r
 /**\r
@@ -1094,62 +1341,93 @@ CoreUpdateProfileFree (
   @param CallerAddress  Address of caller who call Allocate or Free.\r
   @param Action         This Allocate or Free action.\r
   @param MemoryType     Memory type.\r
+                        EfiMaxMemoryType means the MemoryType is unknown.\r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
-\r
-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @param ActionString   String for memory profile action.\r
+                        Only needed for user defined allocate action.\r
+\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
+EFIAPI\r
 CoreUpdateProfile (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType,\r
   IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool\r
-  IN VOID                   *Buffer\r
+  IN VOID                   *Buffer,\r
+  IN CHAR8                  *ActionString OPTIONAL\r
   )\r
 {\r
+  EFI_STATUS                    Status;\r
   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
+  MEMORY_PROFILE_ACTION         BasicAction;\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
-  if (!mMemoryProfileRecordingStatus) {\r
-    return FALSE;\r
+  if (mMemoryProfileGettingStatus) {\r
+    return EFI_ACCESS_DENIED;\r
+  }\r
+\r
+  if (!mMemoryProfileRecordingEnable) {\r
+    return EFI_ABORTED;\r
   }\r
 \r
   //\r
-  // Only record limited MemoryType.\r
+  // Get the basic action to know how to process the record\r
   //\r
-  if (!CoreNeedRecordProfile (MemoryType)) {\r
-    return FALSE;\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
+\r
+  //\r
+  // EfiMaxMemoryType means the MemoryType is unknown.\r
+  //\r
+  if (MemoryType != EfiMaxMemoryType) {\r
+    //\r
+    // Only record limited MemoryType.\r
+    //\r
+    if (!CoreNeedRecordProfile (MemoryType)) {\r
+      return EFI_UNSUPPORTED;\r
+    }\r
   }\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
-  switch (Action) {\r
+  CoreAcquireMemoryProfileLock ();\r
+  switch (BasicAction) {\r
     case MemoryProfileActionAllocatePages:\r
-      CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);\r
+      Status = CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r
       break;\r
     case MemoryProfileActionFreePages:\r
-      CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);\r
+      Status = CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);\r
       break;\r
     case MemoryProfileActionAllocatePool:\r
-      CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);\r
+      Status = CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r
       break;\r
     case MemoryProfileActionFreePool:\r
-      CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);\r
+      Status = CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);\r
       break;\r
     default:\r
       ASSERT (FALSE);\r
+      Status = EFI_UNSUPPORTED;\r
       break;\r
   }\r
-  return TRUE;\r
+  CoreReleaseMemoryProfileLock ();\r
+\r
+  return Status;\r
 }\r
 \r
 ////////////////////\r
@@ -1167,8 +1445,11 @@ MemoryProfileGetDataSize (
 {\r
   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
+  MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;\r
   LIST_ENTRY                        *DriverInfoList;\r
   LIST_ENTRY                        *DriverLink;\r
+  LIST_ENTRY                        *AllocInfoList;\r
+  LIST_ENTRY                        *AllocLink;\r
   UINTN                             TotalSize;\r
 \r
 \r
@@ -1178,7 +1459,6 @@ MemoryProfileGetDataSize (
   }\r
 \r
   TotalSize = sizeof (MEMORY_PROFILE_CONTEXT);\r
-  TotalSize += sizeof (MEMORY_PROFILE_DRIVER_INFO) * (UINTN) ContextData->Context.ImageCount;\r
 \r
   DriverInfoList = ContextData->DriverInfoList;\r
   for (DriverLink = DriverInfoList->ForwardLink;\r
@@ -1190,7 +1470,20 @@ MemoryProfileGetDataSize (
                        Link,\r
                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
                        );\r
-    TotalSize += sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfoData->DriverInfo.AllocRecordCount;\r
+    TotalSize += DriverInfoData->DriverInfo.Header.Length;\r
+\r
+    AllocInfoList = DriverInfoData->AllocInfoList;\r
+    for (AllocLink = AllocInfoList->ForwardLink;\r
+         AllocLink != AllocInfoList;\r
+         AllocLink = AllocLink->ForwardLink) {\r
+      AllocInfoData = CR (\r
+                        AllocLink,\r
+                        MEMORY_PROFILE_ALLOC_INFO_DATA,\r
+                        Link,\r
+                        MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
+                        );\r
+      TotalSize += AllocInfoData->AllocInfo.Header.Length;\r
+    }\r
   }\r
 \r
   return TotalSize;\r
@@ -1217,6 +1510,8 @@ MemoryProfileCopyData (
   LIST_ENTRY                        *DriverLink;\r
   LIST_ENTRY                        *AllocInfoList;\r
   LIST_ENTRY                        *AllocLink;\r
+  UINTN                             PdbSize;\r
+  UINTN                             ActionStringSize;\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
@@ -1238,7 +1533,11 @@ MemoryProfileCopyData (
                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
                        );\r
     CopyMem (DriverInfo, &DriverInfoData->DriverInfo, sizeof (MEMORY_PROFILE_DRIVER_INFO));\r
-    AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) (DriverInfo + 1);\r
+    if (DriverInfo->PdbStringOffset != 0) {\r
+      PdbSize = AsciiStrSize (DriverInfoData->PdbString);\r
+      CopyMem ((VOID *) ((UINTN) DriverInfo + DriverInfo->PdbStringOffset), DriverInfoData->PdbString, PdbSize);\r
+    }\r
+    AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) ((UINTN) DriverInfo + DriverInfo->Header.Length);\r
 \r
     AllocInfoList = DriverInfoData->AllocInfoList;\r
     for (AllocLink = AllocInfoList->ForwardLink;\r
@@ -1251,10 +1550,14 @@ MemoryProfileCopyData (
                         MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
                         );\r
       CopyMem (AllocInfo, &AllocInfoData->AllocInfo, sizeof (MEMORY_PROFILE_ALLOC_INFO));\r
-      AllocInfo += 1;\r
+      if (AllocInfo->ActionStringOffset != 0) {\r
+        ActionStringSize = AsciiStrSize (AllocInfoData->ActionString);\r
+        CopyMem ((VOID *) ((UINTN) AllocInfo + AllocInfo->ActionStringOffset), AllocInfoData->ActionString, ActionStringSize);\r
+      }\r
+      AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) ((UINTN) AllocInfo + AllocInfo->Header.Length);\r
     }\r
 \r
-    DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) ((UINTN) (DriverInfo + 1) + sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfo->AllocRecordCount);\r
+    DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *)  AllocInfo;\r
   }\r
 }\r
 \r
@@ -1267,6 +1570,7 @@ MemoryProfileCopyData (
   @param[out]     ProfileBuffer     Profile buffer.\r
                       \r
   @return EFI_SUCCESS               Get the memory profile data successfully.\r
+  @return EFI_UNSUPPORTED           Memory profile is unsupported.\r
   @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data. \r
                                     ProfileSize is updated with the size required.\r
 \r
@@ -1281,28 +1585,28 @@ ProfileProtocolGetData (
 {\r
   UINTN                                 Size;\r
   MEMORY_PROFILE_CONTEXT_DATA           *ContextData;\r
-  BOOLEAN                               MemoryProfileRecordingStatus;\r
+  BOOLEAN                               MemoryProfileGettingStatus;\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
     return EFI_UNSUPPORTED;\r
   }\r
 \r
-  MemoryProfileRecordingStatus = mMemoryProfileRecordingStatus;\r
-  mMemoryProfileRecordingStatus = FALSE;\r
+  MemoryProfileGettingStatus = mMemoryProfileGettingStatus;\r
+  mMemoryProfileGettingStatus = TRUE;\r
 \r
   Size = MemoryProfileGetDataSize ();\r
 \r
   if (*ProfileSize < Size) {\r
     *ProfileSize = Size;\r
-    mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;\r
+    mMemoryProfileGettingStatus = MemoryProfileGettingStatus;\r
     return EFI_BUFFER_TOO_SMALL;\r
   }\r
 \r
   *ProfileSize = Size;\r
   MemoryProfileCopyData (ProfileBuffer);\r
 \r
-  mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;\r
+  mMemoryProfileGettingStatus = MemoryProfileGettingStatus;\r
   return EFI_SUCCESS;\r
 }\r
 \r
@@ -1315,8 +1619,10 @@ ProfileProtocolGetData (
   @param[in] ImageSize          Image size.\r
   @param[in] FileType           File type of the image.\r
 \r
-  @return EFI_SUCCESS           Register success.\r
-  @return EFI_OUT_OF_RESOURCE   No enough resource for this register.\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.\r
 \r
 **/\r
 EFI_STATUS\r
@@ -1342,7 +1648,7 @@ ProfileProtocolRegisterImage (
   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
   DriverEntry.ImageContext.ImageType = InternalPeCoffGetSubsystem ((VOID *) (UINTN) ImageBase);\r
 \r
-  return RegisterMemoryProfileImage (&DriverEntry, FileType) ? EFI_SUCCESS: EFI_OUT_OF_RESOURCES;\r
+  return RegisterMemoryProfileImage (&DriverEntry, FileType);\r
 }\r
 \r
 /**\r
@@ -1353,7 +1659,9 @@ ProfileProtocolRegisterImage (
   @param[in] ImageBase          Image base address.\r
   @param[in] ImageSize          Image size.\r
 \r
-  @return EFI_SUCCESS           Unregister success.\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
   @return EFI_NOT_FOUND         The image is not found.\r
 \r
 **/\r
@@ -1378,7 +1686,105 @@ ProfileProtocolUnregisterImage (
   ASSERT_EFI_ERROR (Status);\r
   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
 \r
-  return UnregisterMemoryProfileImage (&DriverEntry) ? EFI_SUCCESS: EFI_NOT_FOUND;\r
+  return UnregisterMemoryProfileImage (&DriverEntry);\r
+}\r
+\r
+/**\r
+  Get memory profile recording state.\r
+\r
+  @param[in]  This              The EDKII_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[out] RecordingState    Recording state.\r
+\r
+  @return EFI_SUCCESS           Memory profile recording state is returned.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+  @return EFI_INVALID_PARAMETER RecordingState is NULL.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ProfileProtocolGetRecordingState (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,\r
+  OUT BOOLEAN                           *RecordingState\r
+  )\r
+{\r
+  MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  if (RecordingState == NULL) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+  *RecordingState = mMemoryProfileRecordingEnable;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Set memory profile recording state.\r
+\r
+  @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] RecordingState     Recording state.\r
+\r
+  @return EFI_SUCCESS           Set memory profile recording state successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ProfileProtocolSetRecordingState (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,\r
+  IN BOOLEAN                            RecordingState\r
+  )\r
+{\r
+  MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  mMemoryProfileRecordingEnable = RecordingState;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Record memory profile of multilevel caller.\r
+\r
+  @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] CallerAddress      Address of caller.\r
+  @param[in] Action             Memory profile action.\r
+  @param[in] MemoryType         Memory type.\r
+                                EfiMaxMemoryType means the MemoryType is unknown.\r
+  @param[in] Buffer             Buffer address.\r
+  @param[in] Size               Buffer size.\r
+  @param[in] ActionString       String for memory profile action.\r
+                                Only needed for user defined allocate action.\r
+\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+ProfileProtocolRecord (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,\r
+  IN PHYSICAL_ADDRESS                   CallerAddress,\r
+  IN MEMORY_PROFILE_ACTION              Action,\r
+  IN EFI_MEMORY_TYPE                    MemoryType,\r
+  IN VOID                               *Buffer,\r
+  IN UINTN                              Size,\r
+  IN CHAR8                              *ActionString OPTIONAL\r
+  )\r
+{\r
+  return CoreUpdateProfile (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r
 }\r
 \r
 ////////////////////\r
index 898b722..b02bafb 100644 (file)
@@ -1335,7 +1335,14 @@ CoreAllocatePages (
 \r
   Status = CoreInternalAllocatePages (Type, MemoryType, NumberOfPages, Memory);\r
   if (!EFI_ERROR (Status)) {\r
-    CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionAllocatePages, MemoryType, EFI_PAGES_TO_SIZE (NumberOfPages), (VOID *) (UINTN) *Memory);\r
+    CoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionAllocatePages,\r
+      MemoryType,\r
+      EFI_PAGES_TO_SIZE (NumberOfPages),\r
+      (VOID *) (UINTN) *Memory,\r
+      NULL\r
+      );\r
     InstallMemoryAttributesTableOnMemoryAllocation (MemoryType);\r
   }\r
   return Status;\r
@@ -1444,7 +1451,14 @@ CoreFreePages (
 \r
   Status = CoreInternalFreePages (Memory, NumberOfPages, &MemoryType);\r
   if (!EFI_ERROR (Status)) {\r
-    CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionFreePages, MemoryType, EFI_PAGES_TO_SIZE (NumberOfPages), (VOID *) (UINTN) Memory);\r
+    CoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionFreePages,\r
+      MemoryType,\r
+      EFI_PAGES_TO_SIZE (NumberOfPages),\r
+      (VOID *) (UINTN) Memory,\r
+      NULL\r
+      );\r
     InstallMemoryAttributesTableOnMemoryAllocation (MemoryType);\r
   }\r
   return Status;\r
index 2980e22..6ef5fba 100644 (file)
@@ -276,7 +276,14 @@ CoreAllocatePool (
 \r
   Status = CoreInternalAllocatePool (PoolType, Size, Buffer);\r
   if (!EFI_ERROR (Status)) {\r
-    CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionAllocatePool, PoolType, Size, *Buffer);\r
+    CoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionAllocatePool,\r
+      PoolType,\r
+      Size,\r
+      *Buffer,\r
+      NULL\r
+      );\r
     InstallMemoryAttributesTableOnMemoryAllocation (PoolType);\r
   }\r
   return Status;\r
@@ -505,7 +512,14 @@ CoreFreePool (
 \r
   Status = CoreInternalFreePool (Buffer, &PoolType);\r
   if (!EFI_ERROR (Status)) {\r
-    CoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionFreePool, PoolType, 0, Buffer);\r
+    CoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionFreePool,\r
+      PoolType,\r
+      0,\r
+      Buffer,\r
+      NULL\r
+      );\r
     InstallMemoryAttributesTableOnMemoryAllocation (PoolType);\r
   }\r
   return Status;\r
index 9cc2a4c..5c04e8c 100644 (file)
@@ -1,7 +1,7 @@
 /** @file\r
   SMM Memory page management functions.\r
 \r
-  Copyright (c) 2009 - 2014, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>\r
   This program and the accompanying materials are licensed and made available \r
   under the terms and conditions of the BSD License which accompanies this \r
   distribution.  The full text of the license may be found at        \r
@@ -226,7 +226,14 @@ SmmAllocatePages (
 \r
   Status = SmmInternalAllocatePages (Type, MemoryType, NumberOfPages, Memory);\r
   if (!EFI_ERROR (Status)) {\r
-    SmmCoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionAllocatePages, MemoryType, EFI_PAGES_TO_SIZE (NumberOfPages), (VOID *) (UINTN) *Memory);\r
+    SmmCoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionAllocatePages,\r
+      MemoryType,\r
+      EFI_PAGES_TO_SIZE (NumberOfPages),\r
+      (VOID *) (UINTN) *Memory,\r
+      NULL\r
+      );\r
   }\r
   return Status;\r
 }\r
@@ -344,7 +351,14 @@ SmmFreePages (
 \r
   Status = SmmInternalFreePages (Memory, NumberOfPages);\r
   if (!EFI_ERROR (Status)) {\r
-    SmmCoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionFreePages, 0, EFI_PAGES_TO_SIZE (NumberOfPages), (VOID *) (UINTN) Memory);\r
+    SmmCoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionFreePages,\r
+      EfiMaxMemoryType,\r
+      EFI_PAGES_TO_SIZE (NumberOfPages),\r
+      (VOID *) (UINTN) Memory,\r
+      NULL\r
+      );\r
   }\r
   return Status;\r
 }\r
index 7245f20..551560b 100644 (file)
@@ -1,7 +1,7 @@
 /** @file\r
   SMM Core Main Entry Point\r
 \r
-  Copyright (c) 2009 - 2015, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>\r
   This program and the accompanying materials are licensed and made available \r
   under the terms and conditions of the BSD License which accompanies this \r
   distribution.  The full text of the license may be found at        \r
@@ -632,6 +632,7 @@ SmmMain (
   }\r
 \r
   RegisterSmramProfileHandler ();\r
+  SmramProfileInstallProtocol ();\r
 \r
   SmmCoreInstallLoadedImage ();\r
 \r
index 0e9c92a..000864d 100644 (file)
@@ -2,7 +2,7 @@
   The internal header file includes the common header files, defines\r
   internal structure and functions used by SmmCore module.\r
 \r
-  Copyright (c) 2009 - 2015, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>\r
   This program and the accompanying materials are licensed and made available \r
   under the terms and conditions of the BSD License which accompanies this \r
   distribution.  The full text of the license may be found at        \r
@@ -42,6 +42,7 @@
 #include <Library/BaseLib.h>\r
 #include <Library/BaseMemoryLib.h>\r
 #include <Library/PeCoffLib.h>\r
+#include <Library/PeCoffGetEntryPointLib.h>\r
 #include <Library/CacheMaintenanceLib.h>\r
 #include <Library/DebugLib.h>\r
 #include <Library/ReportStatusCodeLib.h>\r
@@ -884,17 +885,28 @@ SmramProfileInit (
   VOID\r
   );\r
 \r
+/**\r
+  Install SMRAM profile protocol.\r
+\r
+**/\r
+VOID\r
+SmramProfileInstallProtocol (\r
+  VOID\r
+  );\r
+\r
 /**\r
   Register SMM image to SMRAM profile.\r
 \r
   @param DriverEntry    SMM image info.\r
   @param RegisterToDxe  Register image to DXE.\r
 \r
-  @retval TRUE          Register success.\r
-  @retval FALSE         Register fail.\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 RegisterSmramProfileImage (\r
   IN EFI_SMM_DRIVER_ENTRY   *DriverEntry,\r
   IN BOOLEAN                RegisterToDxe\r
@@ -906,11 +918,13 @@ RegisterSmramProfileImage (
   @param DriverEntry        SMM image info.\r
   @param UnregisterToDxe    Unregister image from DXE.\r
 \r
-  @retval TRUE              Unregister success.\r
-  @retval FALSE             Unregister fail.\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_NOT_FOUND         The image is not found.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 UnregisterSmramProfileImage (\r
   IN EFI_SMM_DRIVER_ENTRY   *DriverEntry,\r
   IN BOOLEAN                UnregisterToDxe\r
@@ -922,20 +936,31 @@ UnregisterSmramProfileImage (
   @param CallerAddress  Address of caller who call Allocate or Free.\r
   @param Action         This Allocate or Free action.\r
   @param MemoryType     Memory type.\r
+                        EfiMaxMemoryType means the MemoryType is unknown.\r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
-\r
-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @param ActionString   String for memory profile action.\r
+                        Only needed for user defined allocate action.\r
+\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
+EFIAPI\r
 SmmCoreUpdateProfile (\r
-  IN EFI_PHYSICAL_ADDRESS CallerAddress,\r
-  IN MEMORY_PROFILE_ACTION Action,\r
-  IN EFI_MEMORY_TYPE      MemoryType, // Valid for AllocatePages/AllocatePool\r
-  IN UINTN                Size,       // Valid for AllocatePages/FreePages/AllocatePool\r
-  IN VOID                 *Buffer\r
+  IN PHYSICAL_ADDRESS       CallerAddress,\r
+  IN MEMORY_PROFILE_ACTION  Action,\r
+  IN EFI_MEMORY_TYPE        MemoryType, // Valid for AllocatePages/AllocatePool\r
+  IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool\r
+  IN VOID                   *Buffer,\r
+  IN CHAR8                  *ActionString OPTIONAL\r
   );\r
 \r
 /**\r
index 9c06b2a..f7e32c4 100644 (file)
@@ -1,7 +1,7 @@
 ## @file\r
 # This module provide an SMM CIS compliant implementation of SMM Core.\r
 #\r
-# Copyright (c) 2009 - 2015, Intel Corporation. All rights reserved.<BR>\r
+# Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>\r
 #\r
 # This program and the accompanying materials\r
 # are licensed and made available under the terms and conditions of the BSD License\r
@@ -48,6 +48,7 @@
   BaseLib\r
   BaseMemoryLib\r
   PeCoffLib\r
+  PeCoffGetEntryPointLib\r
   CacheMaintenanceLib  \r
   DebugLib\r
   ReportStatusCodeLib\r
@@ -81,6 +82,7 @@
   gEfiMdeModulePkgTokenSpaceGuid.PcdLoadModuleAtFixAddressEnable        ## CONSUMES\r
   gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfileMemoryType             ## CONSUMES\r
   gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfilePropertyMask           ## CONSUMES\r
+  gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfileDriverPath             ## CONSUMES\r
 \r
 [Guids]\r
   gAprioriGuid                                  ## SOMETIMES_CONSUMES   ## File\r
@@ -92,6 +94,8 @@
   ## SOMETIMES_CONSUMES   ## GUID # Locate protocol\r
   ## SOMETIMES_PRODUCES   ## GUID # SmiHandlerRegister\r
   gEdkiiMemoryProfileGuid\r
+  ## SOMETIMES_PRODUCES   ## GUID # Install protocol\r
+  gEdkiiSmmMemoryProfileGuid\r
   gZeroGuid                                     ## SOMETIMES_CONSUMES   ## GUID\r
 \r
 [UserExtensions.TianoCore."ExtraFiles"]\r
index acfcc83..acb852b 100644 (file)
@@ -1556,7 +1556,7 @@ SmmIplEntry (
     }\r
 \r
     if (gSmmCorePrivate->SmramRanges[Index].CpuStart >= BASE_1MB) {\r
-      if ((gSmmCorePrivate->SmramRanges[Index].CpuStart + gSmmCorePrivate->SmramRanges[Index].PhysicalSize) <= BASE_4GB) {\r
+      if ((gSmmCorePrivate->SmramRanges[Index].CpuStart + gSmmCorePrivate->SmramRanges[Index].PhysicalSize - 1) <= MAX_ADDRESS) {\r
         if (gSmmCorePrivate->SmramRanges[Index].PhysicalSize >= MaxSize) {\r
           MaxSize = gSmmCorePrivate->SmramRanges[Index].PhysicalSize;\r
           mCurrentSmramRange = &gSmmCorePrivate->SmramRanges[Index];\r
index 761988e..53ef3e0 100644 (file)
@@ -1,7 +1,7 @@
 /** @file\r
   SMM Memory pool management functions.\r
 \r
-  Copyright (c) 2009 - 2015, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>\r
   This program and the accompanying materials are licensed and made available \r
   under the terms and conditions of the BSD License which accompanies this \r
   distribution.  The full text of the license may be found at        \r
@@ -67,7 +67,7 @@ SmmInitializeMemoryServices (
       }\r
 \r
       if (SmramRanges[Index].CpuStart >= BASE_1MB) {\r
-        if ((SmramRanges[Index].CpuStart + SmramRanges[Index].PhysicalSize) <= BASE_4GB) {\r
+        if ((SmramRanges[Index].CpuStart + SmramRanges[Index].PhysicalSize - 1) <= MAX_ADDRESS) {\r
           if (SmramRanges[Index].PhysicalSize >= MaxSize) {\r
             MaxSize = SmramRanges[Index].PhysicalSize;\r
             CurrentSmramRangesIndex = Index;\r
@@ -260,7 +260,14 @@ SmmAllocatePool (
 \r
   Status = SmmInternalAllocatePool (PoolType, Size, Buffer);\r
   if (!EFI_ERROR (Status)) {\r
-    SmmCoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionAllocatePool, PoolType, Size, *Buffer);\r
+    SmmCoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionAllocatePool,\r
+      PoolType,\r
+      Size,\r
+      *Buffer,\r
+      NULL\r
+      );\r
   }\r
   return Status;\r
 }\r
@@ -319,7 +326,14 @@ SmmFreePool (
 \r
   Status = SmmInternalFreePool (Buffer);\r
   if (!EFI_ERROR (Status)) {\r
-    SmmCoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionFreePool, 0, 0, Buffer);\r
+    SmmCoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionFreePool,\r
+      EfiMaxMemoryType,\r
+      0,\r
+      Buffer,\r
+      NULL\r
+      );\r
   }\r
   return Status;\r
 }\r
index 57b34a0..281e382 100644 (file)
 #include "PiSmmCore.h"\r
 \r
 #define IS_SMRAM_PROFILE_ENABLED ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT1) != 0)\r
+#define IS_UEFI_MEMORY_PROFILE_ENABLED ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT0) != 0)\r
+\r
+#define GET_OCCUPIED_SIZE(ActualSize, Alignment) \\r
+  ((ActualSize) + (((Alignment) - ((ActualSize) & ((Alignment) - 1))) & ((Alignment) - 1)))\r
 \r
 typedef struct {\r
   UINT32                        Signature;\r
@@ -26,12 +30,14 @@ typedef struct {
   UINT32                        Signature;\r
   MEMORY_PROFILE_DRIVER_INFO    DriverInfo;\r
   LIST_ENTRY                    *AllocInfoList;\r
+  CHAR8                         *PdbString;\r
   LIST_ENTRY                    Link;\r
 } MEMORY_PROFILE_DRIVER_INFO_DATA;\r
 \r
 typedef struct {\r
   UINT32                        Signature;\r
   MEMORY_PROFILE_ALLOC_INFO     AllocInfo;\r
+  CHAR8                         *ActionString;\r
   LIST_ENTRY                    Link;\r
 } MEMORY_PROFILE_ALLOC_INFO_DATA;\r
 \r
@@ -72,7 +78,10 @@ GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA mSmramProfileContext =
 GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA *mSmramProfileContextPtr = NULL;\r
 \r
 BOOLEAN mSmramReadyToLock;\r
-BOOLEAN mSmramProfileRecordingStatus = FALSE;\r
+BOOLEAN mSmramProfileGettingStatus = FALSE;\r
+BOOLEAN mSmramProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;\r
+EFI_DEVICE_PATH_PROTOCOL *mSmramProfileDriverPath;\r
+UINTN                    mSmramProfileDriverPathSize;\r
 \r
 /**\r
   Dump SMRAM infromation.\r
@@ -83,6 +92,155 @@ DumpSmramInfo (
   VOID\r
   );\r
 \r
+/**\r
+  Get memory profile data.\r
+\r
+  @param[in]      This              The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in, out] ProfileSize       On entry, points to the size in bytes of the ProfileBuffer.\r
+                                    On return, points to the size of the data returned in ProfileBuffer.\r
+  @param[out]     ProfileBuffer     Profile buffer.\r
+                      \r
+  @return EFI_SUCCESS               Get the memory profile data successfully.\r
+  @return EFI_UNSUPPORTED           Memory profile is unsupported.\r
+  @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data. \r
+                                    ProfileSize is updated with the size required.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolGetData (\r
+  IN     EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN OUT UINT64                             *ProfileSize,\r
+     OUT VOID                               *ProfileBuffer\r
+  );\r
+\r
+/**\r
+  Register image to memory profile.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] FilePath           File path of the image.\r
+  @param[in] ImageBase          Image base address.\r
+  @param[in] ImageSize          Image size.\r
+  @param[in] FileType           File type of the image.\r
+\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCE   No enough resource for this register.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolRegisterImage (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,\r
+  IN PHYSICAL_ADDRESS                   ImageBase,\r
+  IN UINT64                             ImageSize,\r
+  IN EFI_FV_FILETYPE                    FileType\r
+  );\r
+\r
+/**\r
+  Unregister image from memory profile.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] FilePath           File path of the image.\r
+  @param[in] ImageBase          Image base address.\r
+  @param[in] ImageSize          Image size.\r
+\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_NOT_FOUND         The image is not found.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolUnregisterImage (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,\r
+  IN PHYSICAL_ADDRESS                   ImageBase,\r
+  IN UINT64                             ImageSize\r
+  );\r
+\r
+/**\r
+  Get memory profile recording state.\r
+\r
+  @param[in]  This              The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[out] RecordingState    Recording state.\r
+\r
+  @return EFI_SUCCESS           Memory profile recording state is returned.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+  @return EFI_INVALID_PARAMETER RecordingState is NULL.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolGetRecordingState (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  OUT BOOLEAN                           *RecordingState\r
+  );\r
+\r
+/**\r
+  Set memory profile recording state.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] RecordingState     Recording state.\r
+\r
+  @return EFI_SUCCESS           Set memory profile recording state successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolSetRecordingState (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN BOOLEAN                            RecordingState\r
+  );\r
+\r
+/**\r
+  Record memory profile of multilevel caller.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] CallerAddress      Address of caller.\r
+  @param[in] Action             Memory profile action.\r
+  @param[in] MemoryType         Memory type.\r
+                                EfiMaxMemoryType means the MemoryType is unknown.\r
+  @param[in] Buffer             Buffer address.\r
+  @param[in] Size               Buffer size.\r
+  @param[in] ActionString       String for memory profile action.\r
+                                Only needed for user defined allocate action.\r
+\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolRecord (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN PHYSICAL_ADDRESS                   CallerAddress,\r
+  IN MEMORY_PROFILE_ACTION              Action,\r
+  IN EFI_MEMORY_TYPE                    MemoryType,\r
+  IN VOID                               *Buffer,\r
+  IN UINTN                              Size,\r
+  IN CHAR8                              *ActionString OPTIONAL\r
+  );\r
+\r
+EDKII_SMM_MEMORY_PROFILE_PROTOCOL mSmmProfileProtocol = {\r
+  SmramProfileProtocolGetData,\r
+  SmramProfileProtocolRegisterImage,\r
+  SmramProfileProtocolUnregisterImage,\r
+  SmramProfileProtocolGetRecordingState,\r
+  SmramProfileProtocolSetRecordingState,\r
+  SmramProfileProtocolRecord,\r
+};\r
+\r
 /**\r
   Return SMRAM profile context.\r
 \r
@@ -239,7 +397,6 @@ InternalPeCoffGetEntryPoint (
   @param ImageSize      Image size.\r
   @param EntryPoint     Entry point of the image.\r
   @param ImageSubsystem Image subsystem of the image.\r
-\r
   @param FileType       File type of the image.\r
 \r
   @return Pointer to memory profile driver info.\r
@@ -260,13 +417,27 @@ BuildDriverInfo (
   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   VOID                              *EntryPointInImage;\r
+  CHAR8                             *PdbString;\r
+  UINTN                             PdbSize;\r
+  UINTN                             PdbOccupiedSize;\r
+\r
+  PdbSize = 0;\r
+  PdbOccupiedSize = 0;\r
+  PdbString = NULL;\r
+  if (ImageBase != 0) {\r
+    PdbString = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageBase);\r
+    if (PdbString != NULL) {\r
+      PdbSize = AsciiStrSize (PdbString);\r
+      PdbOccupiedSize = GET_OCCUPIED_SIZE (PdbSize, sizeof (UINT64));\r
+    }\r
+  }\r
 \r
   //\r
   // Use SmmInternalAllocatePool() that will not update profile for this AllocatePool action.\r
   //\r
   Status = SmmInternalAllocatePool (\r
              EfiRuntimeServicesData,\r
-             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY),\r
+             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY) + PdbSize,\r
              (VOID **) &DriverInfoData\r
              );\r
   if (EFI_ERROR (Status)) {\r
@@ -278,7 +449,7 @@ BuildDriverInfo (
   DriverInfo = &DriverInfoData->DriverInfo;\r
   DriverInfoData->Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;\r
   DriverInfo->Header.Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;\r
-  DriverInfo->Header.Length = sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
+  DriverInfo->Header.Length = (UINT16) (sizeof (MEMORY_PROFILE_DRIVER_INFO) + PdbOccupiedSize);\r
   DriverInfo->Header.Revision = MEMORY_PROFILE_DRIVER_INFO_REVISION;\r
   if (FileName != NULL) {\r
     CopyMem (&DriverInfo->FileName, FileName, sizeof (EFI_GUID));\r
@@ -302,6 +473,14 @@ BuildDriverInfo (
   DriverInfo->CurrentUsage = 0;\r
   DriverInfo->PeakUsage = 0;\r
   DriverInfo->AllocRecordCount = 0;\r
+  if (PdbSize != 0) {\r
+    DriverInfo->PdbStringOffset = (UINT16) sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
+    DriverInfoData->PdbString = (CHAR8 *) (DriverInfoData->AllocInfoList + 1);\r
+    CopyMem (DriverInfoData->PdbString, PdbString, PdbSize);\r
+  } else {\r
+    DriverInfo->PdbStringOffset = 0;\r
+    DriverInfoData->PdbString = NULL;\r
+  }\r
 \r
   InsertTailList (ContextData->DriverInfoList, &DriverInfoData->Link);\r
   ContextData->Context.ImageCount ++;\r
@@ -332,7 +511,7 @@ RegisterImageToDxe (
   MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;\r
   UINT8                             TempBuffer[sizeof (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof (EFI_DEVICE_PATH_PROTOCOL)];\r
 \r
-  if (IS_SMRAM_PROFILE_ENABLED) {\r
+  if (IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
 \r
     FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)TempBuffer;\r
     Status = gBS->LocateProtocol (&gEdkiiMemoryProfileGuid, NULL, (VOID **) &ProfileProtocol);\r
@@ -371,7 +550,7 @@ UnregisterImageFromDxe (
   MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;\r
   UINT8                             TempBuffer[sizeof (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof (EFI_DEVICE_PATH_PROTOCOL)];\r
 \r
-  if (IS_SMRAM_PROFILE_ENABLED) {\r
+  if (IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
 \r
     FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)TempBuffer;\r
     Status = gBS->LocateProtocol (&gEdkiiMemoryProfileGuid, NULL, (VOID *) &ProfileProtocol);\r
@@ -389,6 +568,65 @@ UnregisterImageFromDxe (
   }\r
 }\r
 \r
+/**\r
+  Return if record for this driver is needed..\r
+\r
+  @param DriverFilePath     Driver file path.\r
+\r
+  @retval TRUE              Record for this driver is needed.\r
+  @retval FALSE             Record for this driver is not needed.\r
+\r
+**/\r
+BOOLEAN\r
+NeedRecordThisDriver (\r
+  IN EFI_DEVICE_PATH_PROTOCOL       *DriverFilePath\r
+  )\r
+{\r
+  EFI_DEVICE_PATH_PROTOCOL    *TmpDevicePath;\r
+  EFI_DEVICE_PATH_PROTOCOL    *DevicePathInstance;\r
+  UINTN                       DevicePathSize;\r
+  UINTN                       FilePathSize;\r
+\r
+  if (!IsDevicePathValid (mSmramProfileDriverPath, mSmramProfileDriverPathSize)) {\r
+    //\r
+    // Invalid Device Path means record all.\r
+    //\r
+    return TRUE;\r
+  }\r
+  \r
+  //\r
+  // Record FilePath without end node.\r
+  //\r
+  FilePathSize = GetDevicePathSize (DriverFilePath) - sizeof(EFI_DEVICE_PATH_PROTOCOL);\r
+\r
+  DevicePathInstance = mSmramProfileDriverPath;\r
+  do {\r
+    //\r
+    // Find End node (it might be END_ENTIRE or END_INSTANCE)\r
+    //\r
+    TmpDevicePath = DevicePathInstance;\r
+    while (!IsDevicePathEndType (TmpDevicePath)) {\r
+      TmpDevicePath = NextDevicePathNode (TmpDevicePath);\r
+    }\r
+\r
+    //\r
+    // Do not compare END node\r
+    //\r
+    DevicePathSize = (UINTN)TmpDevicePath - (UINTN)DevicePathInstance;\r
+    if ((FilePathSize == DevicePathSize) &&\r
+        (CompareMem (DriverFilePath, DevicePathInstance, DevicePathSize) == 0)) {\r
+      return TRUE;\r
+    }\r
+\r
+    //\r
+    // Get next instance\r
+    //\r
+    DevicePathInstance = (EFI_DEVICE_PATH_PROTOCOL *)((UINTN)DevicePathInstance + DevicePathSize + DevicePathNodeLength(TmpDevicePath));\r
+  } while (DevicePathSubType (TmpDevicePath) != END_ENTIRE_DEVICE_PATH_SUBTYPE);\r
+\r
+  return FALSE;\r
+}\r
+\r
 /**\r
   Register SMM Core to SMRAM profile.\r
 \r
@@ -405,15 +643,16 @@ RegisterSmmCore (
 {\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   PHYSICAL_ADDRESS                  ImageBase;\r
+  UINT8                             TempBuffer[sizeof(MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof(EFI_DEVICE_PATH_PROTOCOL)];\r
+  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;\r
 \r
-  ASSERT (ContextData != NULL);\r
+  FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) TempBuffer;\r
+  EfiInitializeFwVolDevicepathNode (FilePath, &gEfiCallerIdGuid);\r
+  SetDevicePathEndNode (FilePath + 1);\r
 \r
-  RegisterImageToDxe (\r
-    &gEfiCallerIdGuid,\r
-    gSmmCorePrivate->PiSmmCoreImageBase,\r
-    gSmmCorePrivate->PiSmmCoreImageSize,\r
-    EFI_FV_FILETYPE_SMM_CORE\r
-    );\r
+  if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *) FilePath)) {\r
+    return FALSE;\r
+  }\r
 \r
   ImageBase = gSmmCorePrivate->PiSmmCoreImageBase;\r
   DriverInfoData = BuildDriverInfo (\r
@@ -443,6 +682,13 @@ SmramProfileInit (
 {\r
   MEMORY_PROFILE_CONTEXT_DATA *SmramProfileContext;\r
 \r
+  RegisterImageToDxe (\r
+    &gEfiCallerIdGuid,\r
+    gSmmCorePrivate->PiSmmCoreImageBase,\r
+    gSmmCorePrivate->PiSmmCoreImageSize,\r
+    EFI_FV_FILETYPE_SMM_CORE\r
+    );\r
+\r
   if (!IS_SMRAM_PROFILE_ENABLED) {\r
     return;\r
   }\r
@@ -452,7 +698,14 @@ SmramProfileInit (
     return;\r
   }\r
 \r
-  mSmramProfileRecordingStatus = TRUE;\r
+  mSmramProfileGettingStatus = FALSE;\r
+  if ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT7) != 0) {\r
+    mSmramProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;\r
+  } else {\r
+    mSmramProfileRecordingEnable = MEMORY_PROFILE_RECORDING_ENABLE;\r
+  }\r
+  mSmramProfileDriverPathSize = PcdGetSize (PcdMemoryProfileDriverPath);\r
+  mSmramProfileDriverPath = AllocateCopyPool (mSmramProfileDriverPathSize, PcdGetPtr (PcdMemoryProfileDriverPath));\r
   mSmramProfileContextPtr = &mSmramProfileContext;\r
 \r
   RegisterSmmCore (&mSmramProfileContext);\r
@@ -460,17 +713,76 @@ SmramProfileInit (
   DEBUG ((EFI_D_INFO, "SmramProfileInit SmramProfileContext - 0x%x\n", &mSmramProfileContext));\r
 }\r
 \r
+/**\r
+  Install SMRAM profile protocol.\r
+\r
+**/\r
+VOID\r
+SmramProfileInstallProtocol (\r
+  VOID\r
+  )\r
+{\r
+  EFI_HANDLE    Handle;\r
+  EFI_STATUS    Status;\r
+\r
+  if (!IS_SMRAM_PROFILE_ENABLED) {\r
+    return;\r
+  }\r
+\r
+  Handle = NULL;\r
+  Status = SmmInstallProtocolInterface (\r
+             &Handle,\r
+             &gEdkiiSmmMemoryProfileGuid,\r
+             EFI_NATIVE_INTERFACE,\r
+             &mSmmProfileProtocol\r
+             );\r
+  ASSERT_EFI_ERROR (Status);\r
+}\r
+\r
+/**\r
+  Get the GUID file name from the file path.\r
+\r
+  @param FilePath  File path.\r
+\r
+  @return The GUID file name from the file path.\r
+\r
+**/\r
+EFI_GUID *\r
+GetFileNameFromFilePath (\r
+  IN EFI_DEVICE_PATH_PROTOCOL   *FilePath\r
+  )\r
+{\r
+  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH     *ThisFilePath;\r
+  EFI_GUID                              *FileName;\r
+\r
+  FileName = NULL;\r
+  if (FilePath != NULL) {\r
+    ThisFilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) FilePath;\r
+    while (!IsDevicePathEnd (ThisFilePath)) {\r
+      FileName = EfiGetNameGuidFromFwVolDevicePathNode (ThisFilePath);\r
+      if (FileName != NULL) {\r
+        break;\r
+      }\r
+      ThisFilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) NextDevicePathNode (ThisFilePath);\r
+    }\r
+  }\r
+\r
+  return FileName;\r
+}\r
+\r
 /**\r
   Register SMM image to SMRAM profile.\r
 \r
   @param DriverEntry    SMM image info.\r
   @param RegisterToDxe  Register image to DXE.\r
 \r
-  @retval TRUE          Register success.\r
-  @retval FALSE         Register fail.\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 RegisterSmramProfileImage (\r
   IN EFI_SMM_DRIVER_ENTRY   *DriverEntry,\r
   IN BOOLEAN                RegisterToDxe\r
@@ -478,10 +790,8 @@ RegisterSmramProfileImage (
 {\r
   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
-\r
-  if (!IS_SMRAM_PROFILE_ENABLED) {\r
-    return FALSE;\r
-  }\r
+  UINT8                             TempBuffer[sizeof(MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof(EFI_DEVICE_PATH_PROTOCOL)];\r
+  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;\r
 \r
   if (RegisterToDxe) {\r
     RegisterImageToDxe (\r
@@ -492,9 +802,21 @@ RegisterSmramProfileImage (
       );\r
   }\r
 \r
+  if (!IS_SMRAM_PROFILE_ENABLED) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) TempBuffer;\r
+  EfiInitializeFwVolDevicepathNode (FilePath, &DriverEntry->FileName);\r
+  SetDevicePathEndNode (FilePath + 1);\r
+\r
+  if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *) FilePath)) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = BuildDriverInfo (\r
@@ -507,10 +829,10 @@ RegisterSmramProfileImage (
                      EFI_FV_FILETYPE_SMM\r
                      );\r
   if (DriverInfoData == NULL) {\r
-    return FALSE;\r
+    return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
@@ -557,42 +879,6 @@ GetMemoryProfileDriverInfoByFileNameAndAddress (
   return NULL;\r
 }\r
 \r
-/**\r
-  Search dummy image from SMRAM profile.\r
-\r
-  @param ContextData    Memory profile context.\r
-\r
-  @return Pointer to memory profile driver info.\r
-\r
-**/\r
-MEMORY_PROFILE_DRIVER_INFO_DATA *\r
-FindDummyImage (\r
-  IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData\r
-  )\r
-{\r
-  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
-  LIST_ENTRY                        *DriverLink;\r
-  LIST_ENTRY                        *DriverInfoList;\r
-\r
-  DriverInfoList = ContextData->DriverInfoList;\r
-\r
-  for (DriverLink = DriverInfoList->ForwardLink;\r
-       DriverLink != DriverInfoList;\r
-       DriverLink = DriverLink->ForwardLink) {\r
-    DriverInfoData = CR (\r
-                   DriverLink,\r
-                   MEMORY_PROFILE_DRIVER_INFO_DATA,\r
-                   Link,\r
-                   MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
-                   );\r
-    if (CompareGuid (&gZeroGuid, &DriverInfoData->DriverInfo.FileName)) {\r
-      return DriverInfoData;\r
-    }\r
-  }\r
-\r
-  return BuildDriverInfo (ContextData, &gZeroGuid, 0, 0, 0, 0, 0);\r
-}\r
-\r
 /**\r
   Search image from memory profile.\r
   It will return image, if (Address >= ImageBuffer) AND (Address < ImageBuffer + ImageSize)\r
@@ -632,10 +918,7 @@ GetMemoryProfileDriverInfoFromAddress (
     }\r
   }\r
 \r
-  //\r
-  // Should never come here.\r
-  //\r
-  return FindDummyImage (ContextData);\r
+  return NULL;\r
 }\r
 \r
 /**\r
@@ -644,11 +927,13 @@ GetMemoryProfileDriverInfoFromAddress (
   @param DriverEntry        SMM image info.\r
   @param UnregisterFromDxe  Unregister image from DXE.\r
 \r
-  @retval TRUE              Unregister success.\r
-  @retval FALSE             Unregister fail.\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_NOT_FOUND         The image is not found.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 UnregisterSmramProfileImage (\r
   IN EFI_SMM_DRIVER_ENTRY  *DriverEntry,\r
   IN BOOLEAN               UnregisterFromDxe\r
@@ -660,10 +945,8 @@ UnregisterSmramProfileImage (
   EFI_GUID                          *FileName;\r
   PHYSICAL_ADDRESS                  ImageAddress;\r
   VOID                              *EntryPointInImage;\r
-\r
-  if (!IS_SMRAM_PROFILE_ENABLED) {\r
-    return FALSE;\r
-  }\r
+  UINT8                             TempBuffer[sizeof(MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof(EFI_DEVICE_PATH_PROTOCOL)];\r
+  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;\r
 \r
   if (UnregisterFromDxe) {\r
     UnregisterImageFromDxe (\r
@@ -673,9 +956,21 @@ UnregisterSmramProfileImage (
       );\r
   }\r
 \r
+  if (!IS_SMRAM_PROFILE_ENABLED) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) TempBuffer;\r
+  EfiInitializeFwVolDevicepathNode (FilePath, &DriverEntry->FileName);\r
+  SetDevicePathEndNode (FilePath + 1);\r
+\r
+  if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *) FilePath)) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = NULL;\r
@@ -697,12 +992,13 @@ UnregisterSmramProfileImage (
     DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, ImageAddress);\r
   }\r
   if (DriverInfoData == NULL) {\r
-    return FALSE;\r
+    return EFI_NOT_FOUND;\r
   }\r
 \r
   ContextData->Context.TotalImageSize -= DriverInfoData->DriverInfo.ImageSize;\r
 \r
-  DriverInfoData->DriverInfo.ImageBase = 0;\r
+  // Keep the ImageBase for RVA calculation in Application.\r
+  //DriverInfoData->DriverInfo.ImageBase = 0;\r
   DriverInfoData->DriverInfo.ImageSize = 0;\r
 \r
   if (DriverInfoData->DriverInfo.PeakUsage == 0) {\r
@@ -714,7 +1010,7 @@ UnregisterSmramProfileImage (
     SmmInternalFreePool (DriverInfoData);\r
   }\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
@@ -807,54 +1103,80 @@ SmramProfileUpdateFreePages (
   @param MemoryType     Memory type.\r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
+  @param ActionString   String for memory profile action.\r
 \r
-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 SmmCoreUpdateProfileAllocate (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType,\r
   IN UINTN                  Size,\r
-  IN VOID                   *Buffer\r
+  IN VOID                   *Buffer,\r
+  IN CHAR8                  *ActionString OPTIONAL\r
   )\r
 {\r
   EFI_STATUS                        Status;\r
-  MEMORY_PROFILE_CONTEXT           *Context;\r
-  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;\r
-  MEMORY_PROFILE_ALLOC_INFO        *AllocInfo;\r
+  MEMORY_PROFILE_CONTEXT            *Context;\r
+  MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;\r
+  MEMORY_PROFILE_ALLOC_INFO         *AllocInfo;\r
   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;\r
   EFI_MEMORY_TYPE                   ProfileMemoryIndex;\r
+  MEMORY_PROFILE_ACTION             BasicAction;\r
+  UINTN                             ActionStringSize;\r
+  UINTN                             ActionStringOccupiedSize;\r
 \r
-  AllocInfoData = NULL;\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);\r
-  ASSERT (DriverInfoData != NULL);\r
+  if (DriverInfoData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  ActionStringSize = 0;\r
+  ActionStringOccupiedSize = 0;\r
+  if (ActionString != NULL) {\r
+    ActionStringSize = AsciiStrSize (ActionString);\r
+    ActionStringOccupiedSize = GET_OCCUPIED_SIZE (ActionStringSize, sizeof (UINT64));\r
+  }\r
 \r
   //\r
   // Use SmmInternalAllocatePool() that will not update profile for this AllocatePool action.\r
   //\r
+  AllocInfoData = NULL;\r
   Status = SmmInternalAllocatePool (\r
              EfiRuntimeServicesData,\r
-             sizeof (*AllocInfoData),\r
+             sizeof (*AllocInfoData) + ActionStringSize,\r
              (VOID **) &AllocInfoData\r
              );\r
   if (EFI_ERROR (Status)) {\r
-    return FALSE;\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
+  ASSERT (AllocInfoData != NULL);\r
+\r
+  //\r
+  // Only update SequenceCount if and only if it is basic action.\r
+  //\r
+  if (Action == BasicAction) {\r
+    ContextData->Context.SequenceCount ++;\r
   }\r
+\r
   AllocInfo = &AllocInfoData->AllocInfo;\r
   AllocInfoData->Signature      = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;\r
   AllocInfo->Header.Signature   = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;\r
-  AllocInfo->Header.Length      = sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
+  AllocInfo->Header.Length      = (UINT16) (sizeof (MEMORY_PROFILE_ALLOC_INFO) + ActionStringOccupiedSize);\r
   AllocInfo->Header.Revision    = MEMORY_PROFILE_ALLOC_INFO_REVISION;\r
   AllocInfo->CallerAddress      = CallerAddress;\r
   AllocInfo->SequenceId         = ContextData->Context.SequenceCount;\r
@@ -862,42 +1184,56 @@ SmmCoreUpdateProfileAllocate (
   AllocInfo->MemoryType         = MemoryType;\r
   AllocInfo->Buffer             = (PHYSICAL_ADDRESS) (UINTN) Buffer;\r
   AllocInfo->Size               = Size;\r
+  if (ActionString != NULL) {\r
+    AllocInfo->ActionStringOffset = (UINT16) sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
+    AllocInfoData->ActionString = (CHAR8 *) (AllocInfoData + 1);\r
+    CopyMem (AllocInfoData->ActionString, ActionString, ActionStringSize);\r
+  } else {\r
+    AllocInfo->ActionStringOffset = 0;\r
+    AllocInfoData->ActionString = NULL;\r
+  }\r
 \r
   InsertTailList (DriverInfoData->AllocInfoList, &AllocInfoData->Link);\r
 \r
-  ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);\r
-\r
+  Context = &ContextData->Context;\r
   DriverInfo = &DriverInfoData->DriverInfo;\r
-  DriverInfo->CurrentUsage += Size;\r
-  if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {\r
-    DriverInfo->PeakUsage = DriverInfo->CurrentUsage;\r
-  }\r
-  DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;\r
-  if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {\r
-    DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];\r
-  }\r
   DriverInfo->AllocRecordCount ++;\r
 \r
-  Context = &ContextData->Context;\r
-  Context->CurrentTotalUsage += Size;\r
-  if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {\r
-    Context->PeakTotalUsage = Context->CurrentTotalUsage;\r
-  }\r
-  Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;\r
-  if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {\r
-    Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];\r
+  //\r
+  // Update summary if and only if it is basic action.\r
+  //\r
+  if (Action == BasicAction) {\r
+    ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);\r
+\r
+    DriverInfo->CurrentUsage += Size;\r
+    if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {\r
+      DriverInfo->PeakUsage = DriverInfo->CurrentUsage;\r
+    }\r
+    DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;\r
+    if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {\r
+      DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];\r
+    }\r
+\r
+    Context->CurrentTotalUsage += Size;\r
+    if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {\r
+      Context->PeakTotalUsage = Context->CurrentTotalUsage;\r
+    }\r
+    Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;\r
+    if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {\r
+      Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];\r
+    }\r
+\r
+    SmramProfileUpdateFreePages (ContextData);\r
   }\r
-  Context->SequenceCount ++;\r
 \r
-  SmramProfileUpdateFreePages (ContextData);\r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
   Get memory profile alloc info from memory profile\r
 \r
   @param DriverInfoData     Driver info\r
-  @param Action             This Free action\r
+  @param BasicAction        This Free basic action\r
   @param Size               Buffer size\r
   @param Buffer             Buffer address\r
 \r
@@ -906,7 +1242,7 @@ SmmCoreUpdateProfileAllocate (
 MEMORY_PROFILE_ALLOC_INFO_DATA *\r
 GetMemoryProfileAllocInfoFromAddress (\r
   IN MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData,\r
-  IN MEMORY_PROFILE_ACTION              Action,\r
+  IN MEMORY_PROFILE_ACTION              BasicAction,\r
   IN UINTN                              Size,\r
   IN VOID                               *Buffer\r
   )\r
@@ -928,10 +1264,10 @@ GetMemoryProfileAllocInfoFromAddress (
                       MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
                       );\r
     AllocInfo = &AllocInfoData->AllocInfo;\r
-    if (AllocInfo->Action != Action) {\r
+    if ((AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK) != BasicAction) {\r
       continue;\r
     }\r
-    switch (Action) {\r
+    switch (BasicAction) {\r
       case MemoryProfileActionAllocatePages:\r
         if ((AllocInfo->Buffer <= (PHYSICAL_ADDRESS) (UINTN) Buffer) &&\r
             ((AllocInfo->Buffer + AllocInfo->Size) >= ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size))) {\r
@@ -960,11 +1296,12 @@ GetMemoryProfileAllocInfoFromAddress (
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
 \r
-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 SmmCoreUpdateProfileFree (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
@@ -982,114 +1319,140 @@ SmmCoreUpdateProfileFree (
   MEMORY_PROFILE_DRIVER_INFO_DATA  *ThisDriverInfoData;\r
   MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;\r
   EFI_MEMORY_TYPE                  ProfileMemoryIndex;\r
+  MEMORY_PROFILE_ACTION            BasicAction;\r
+  BOOLEAN                          Found;\r
+\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);\r
-  ASSERT (DriverInfoData != NULL);\r
 \r
-  switch (Action) {\r
-    case MemoryProfileActionFreePages:\r
-      AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
-      break;\r
-    case MemoryProfileActionFreePool:\r
-      AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
-      break;\r
-    default:\r
-      ASSERT (FALSE);\r
-      AllocInfoData = NULL;\r
-      break;\r
-  }\r
-  if (AllocInfoData == NULL) {\r
-    //\r
-    // Legal case, because driver A might free memory allocated by driver B, by some protocol.\r
-    //\r
-    DriverInfoList = ContextData->DriverInfoList;\r
-\r
-    for (DriverLink = DriverInfoList->ForwardLink;\r
-         DriverLink != DriverInfoList;\r
-         DriverLink = DriverLink->ForwardLink) {\r
-      ThisDriverInfoData = CR (\r
-                             DriverLink,\r
-                             MEMORY_PROFILE_DRIVER_INFO_DATA,\r
-                             Link,\r
-                             MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
-                             );\r
-      switch (Action) {\r
+  //\r
+  // Do not return if DriverInfoData == NULL here,\r
+  // because driver A might free memory allocated by driver B.\r
+  //\r
+\r
+  //\r
+  // Need use do-while loop to find all possible record,\r
+  // because one address might be recorded multiple times.\r
+  //\r
+  Found = FALSE;\r
+  AllocInfoData = NULL;\r
+  do {\r
+    if (DriverInfoData != NULL) {\r
+      switch (BasicAction) {\r
         case MemoryProfileActionFreePages:\r
-          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
+          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
           break;\r
         case MemoryProfileActionFreePool:\r
-          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
+          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
           break;\r
         default:\r
           ASSERT (FALSE);\r
           AllocInfoData = NULL;\r
           break;\r
       }\r
-      if (AllocInfoData != NULL) {\r
-        DriverInfoData = ThisDriverInfoData;\r
-        break;\r
-      }\r
     }\r
-\r
     if (AllocInfoData == NULL) {\r
       //\r
-      // No matched allocate operation is found for this free operation.\r
-      // It is because the specified memory type allocate operation has been\r
-      // filtered by CoreNeedRecordProfile(), but free operations have no\r
-      // memory type information, they can not be filtered by CoreNeedRecordProfile().\r
-      // Then, they will be filtered here.\r
+      // Legal case, because driver A might free memory allocated by driver B, by some protocol.\r
       //\r
-      return FALSE;\r
-    }\r
-  }\r
+      DriverInfoList = ContextData->DriverInfoList;\r
+\r
+      for (DriverLink = DriverInfoList->ForwardLink;\r
+           DriverLink != DriverInfoList;\r
+           DriverLink = DriverLink->ForwardLink) {\r
+        ThisDriverInfoData = CR (\r
+                               DriverLink,\r
+                               MEMORY_PROFILE_DRIVER_INFO_DATA,\r
+                               Link,\r
+                               MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
+                               );\r
+        switch (BasicAction) {\r
+          case MemoryProfileActionFreePages:\r
+            AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
+            break;\r
+          case MemoryProfileActionFreePool:\r
+            AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
+            break;\r
+          default:\r
+            ASSERT (FALSE);\r
+            AllocInfoData = NULL;\r
+            break;\r
+        }\r
+        if (AllocInfoData != NULL) {\r
+          DriverInfoData = ThisDriverInfoData;\r
+          break;\r
+        }\r
+      }\r
 \r
-  Context = &ContextData->Context;\r
-  DriverInfo = &DriverInfoData->DriverInfo;\r
-  AllocInfo = &AllocInfoData->AllocInfo;\r
+      if (AllocInfoData == NULL) {\r
+        //\r
+        // If (!Found), no matched allocate info is found for this free action.\r
+        // It is because the specified memory type allocate actions have been filtered by\r
+        // CoreNeedRecordProfile(), but free actions have no memory type information,\r
+        // they can not be filtered by CoreNeedRecordProfile(). Then, they will be\r
+        // filtered here.\r
+        //\r
+        // If (Found), it is normal exit path.\r
+        return (Found ? EFI_SUCCESS : EFI_NOT_FOUND);\r
+      }\r
+    }\r
 \r
-  ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);\r
+    Found = TRUE;\r
 \r
-  Context->CurrentTotalUsage -= AllocInfo->Size;\r
-  Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
+    Context = &ContextData->Context;\r
+    DriverInfo = &DriverInfoData->DriverInfo;\r
+    AllocInfo = &AllocInfoData->AllocInfo;\r
 \r
-  DriverInfo->CurrentUsage -= AllocInfo->Size;\r
-  DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
-  DriverInfo->AllocRecordCount --;\r
+    DriverInfo->AllocRecordCount --;\r
+    //\r
+    // Update summary if and only if it is basic action.\r
+    //\r
+    if (AllocInfo->Action == (AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK)) {\r
+      ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);\r
 \r
-  RemoveEntryList (&AllocInfoData->Link);\r
+      Context->CurrentTotalUsage -= AllocInfo->Size;\r
+      Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
 \r
-  if (Action == MemoryProfileActionFreePages) {\r
-    if (AllocInfo->Buffer != (PHYSICAL_ADDRESS) (UINTN) Buffer) {\r
-      SmmCoreUpdateProfileAllocate (\r
-        AllocInfo->CallerAddress,\r
-        MemoryProfileActionAllocatePages,\r
-        AllocInfo->MemoryType,\r
-        (UINTN) ((PHYSICAL_ADDRESS) (UINTN) Buffer - AllocInfo->Buffer),\r
-        (VOID *) (UINTN) AllocInfo->Buffer\r
-        );\r
-    }\r
-    if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)) {\r
-      SmmCoreUpdateProfileAllocate (\r
-        AllocInfo->CallerAddress,\r
-        MemoryProfileActionAllocatePages,\r
-        AllocInfo->MemoryType,\r
-        (UINTN) ((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)),\r
-        (VOID *) ((UINTN) Buffer + Size)\r
-        );\r
+      DriverInfo->CurrentUsage -= AllocInfo->Size;\r
+      DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
     }\r
-  }\r
 \r
-  //\r
-  // Use SmmInternalFreePool() that will not update profile for this FreePool action.\r
-  //\r
-  SmmInternalFreePool (AllocInfoData);\r
+    RemoveEntryList (&AllocInfoData->Link);\r
+\r
+    if (BasicAction == MemoryProfileActionFreePages) {\r
+      if (AllocInfo->Buffer != (PHYSICAL_ADDRESS) (UINTN) Buffer) {\r
+        SmmCoreUpdateProfileAllocate (\r
+          AllocInfo->CallerAddress,\r
+          AllocInfo->Action,\r
+          AllocInfo->MemoryType,\r
+          (UINTN) ((PHYSICAL_ADDRESS) (UINTN) Buffer - AllocInfo->Buffer),\r
+          (VOID *) (UINTN) AllocInfo->Buffer,\r
+          AllocInfoData->ActionString\r
+          );\r
+      }\r
+      if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)) {\r
+        SmmCoreUpdateProfileAllocate (\r
+          AllocInfo->CallerAddress,\r
+          AllocInfo->Action,\r
+          AllocInfo->MemoryType,\r
+          (UINTN) ((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)),\r
+          (VOID *) ((UINTN) Buffer + Size),\r
+          AllocInfoData->ActionString\r
+          );\r
+      }\r
+    }\r
 \r
-  return TRUE;\r
+    //\r
+    // Use SmmInternalFreePool() that will not update profile for this FreePool action.\r
+    //\r
+    SmmInternalFreePool (AllocInfoData);\r
+  } while (TRUE);\r
 }\r
 \r
 /**\r
@@ -1098,68 +1461,91 @@ SmmCoreUpdateProfileFree (
   @param CallerAddress  Address of caller who call Allocate or Free.\r
   @param Action         This Allocate or Free action.\r
   @param MemoryType     Memory type.\r
+                        EfiMaxMemoryType means the MemoryType is unknown.\r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
-\r
-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @param ActionString   String for memory profile action.\r
+                        Only needed for user defined allocate action.\r
+\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
+EFIAPI\r
 SmmCoreUpdateProfile (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType, // Valid for AllocatePages/AllocatePool\r
   IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool\r
-  IN VOID                   *Buffer\r
+  IN VOID                   *Buffer,\r
+  IN CHAR8                  *ActionString OPTIONAL\r
   )\r
 {\r
+  EFI_STATUS                    Status;\r
   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
+  MEMORY_PROFILE_ACTION         BasicAction;\r
 \r
   if (!IS_SMRAM_PROFILE_ENABLED) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
-  if (!mSmramProfileRecordingStatus) {\r
-    return FALSE;\r
+  if (mSmramProfileGettingStatus) {\r
+    return EFI_ACCESS_DENIED;\r
+  }\r
+\r
+  if (!mSmramProfileRecordingEnable) {\r
+    return EFI_ABORTED;\r
   }\r
 \r
+  //\r
+  // Get the basic action to know how to process the record\r
+  //\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
+\r
   //\r
   // Free operations have no memory type information, so skip the check.\r
   //\r
-  if ((Action == MemoryProfileActionAllocatePages) || (Action == MemoryProfileActionAllocatePool)) {\r
+  if ((BasicAction == MemoryProfileActionAllocatePages) || (BasicAction == MemoryProfileActionAllocatePool)) {\r
     //\r
     // Only record limited MemoryType.\r
     //\r
     if (!SmmCoreNeedRecordProfile (MemoryType)) {\r
-      return FALSE;\r
+      return EFI_UNSUPPORTED;\r
     }\r
   }\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
-  switch (Action) {\r
+  switch (BasicAction) {\r
     case MemoryProfileActionAllocatePages:\r
-      SmmCoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);\r
+      Status = SmmCoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r
       break;\r
     case MemoryProfileActionFreePages:\r
-      SmmCoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);\r
+      Status = SmmCoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);\r
       break;\r
     case MemoryProfileActionAllocatePool:\r
-      SmmCoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);\r
+      Status = SmmCoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r
       break;\r
     case MemoryProfileActionFreePool:\r
-      SmmCoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);\r
+      Status = SmmCoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);\r
       break;\r
     default:\r
       ASSERT (FALSE);\r
+      Status = EFI_UNSUPPORTED;\r
       break;\r
   }\r
 \r
-  return TRUE;\r
+  return Status;\r
 }\r
 \r
 /**\r
@@ -1192,17 +1578,20 @@ SmramProfileGetDataSize (
   VOID\r
   )\r
 {\r
-  MEMORY_PROFILE_CONTEXT_DATA      *ContextData;\r
-  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;\r
-  LIST_ENTRY                      *DriverInfoList;\r
-  LIST_ENTRY                      *DriverLink;\r
-  UINTN                           TotalSize;\r
-  LIST_ENTRY                      *Node;\r
-  LIST_ENTRY                      *FreePageList;\r
-  LIST_ENTRY                      *FreePoolList;\r
-  FREE_POOL_HEADER                *Pool;\r
-  UINTN                           PoolListIndex;\r
-  UINTN                           Index;\r
+  MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
+  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
+  MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;\r
+  LIST_ENTRY                        *DriverInfoList;\r
+  LIST_ENTRY                        *DriverLink;\r
+  LIST_ENTRY                        *AllocInfoList;\r
+  LIST_ENTRY                        *AllocLink;\r
+  UINTN                             TotalSize;\r
+  LIST_ENTRY                        *Node;\r
+  LIST_ENTRY                        *FreePageList;\r
+  LIST_ENTRY                        *FreePoolList;\r
+  FREE_POOL_HEADER                  *Pool;\r
+  UINTN                             PoolListIndex;\r
+  UINTN                             Index;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
@@ -1210,7 +1599,6 @@ SmramProfileGetDataSize (
   }\r
 \r
   TotalSize = sizeof (MEMORY_PROFILE_CONTEXT);\r
-  TotalSize += sizeof (MEMORY_PROFILE_DRIVER_INFO) * (UINTN) ContextData->Context.ImageCount;\r
 \r
   DriverInfoList = ContextData->DriverInfoList;\r
   for (DriverLink = DriverInfoList->ForwardLink;\r
@@ -1222,7 +1610,20 @@ SmramProfileGetDataSize (
                    Link,\r
                    MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
                    );\r
-    TotalSize += sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfoData->DriverInfo.AllocRecordCount;\r
+    TotalSize += DriverInfoData->DriverInfo.Header.Length;\r
+\r
+    AllocInfoList = DriverInfoData->AllocInfoList;\r
+    for (AllocLink = AllocInfoList->ForwardLink;\r
+         AllocLink != AllocInfoList;\r
+         AllocLink = AllocLink->ForwardLink) {\r
+      AllocInfoData = CR (\r
+                        AllocLink,\r
+                        MEMORY_PROFILE_ALLOC_INFO_DATA,\r
+                        Link,\r
+                        MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
+                        );\r
+      TotalSize += AllocInfoData->AllocInfo.Header.Length;\r
+    }\r
   }\r
 \r
 \r
@@ -1291,6 +1692,8 @@ SmramProfileCopyData (
   MEMORY_PROFILE_DESCRIPTOR       *MemoryProfileDescriptor;\r
   UINT64                          Offset;\r
   UINT64                          RemainingSize;\r
+  UINTN                           PdbSize;\r
+  UINTN                           ActionStringSize;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
@@ -1322,17 +1725,21 @@ SmramProfileCopyData (
                        Link,\r
                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
                        );\r
-    if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_DRIVER_INFO))) {\r
-      if (RemainingSize >= sizeof (MEMORY_PROFILE_DRIVER_INFO)) {\r
+    if (*ProfileOffset < (Offset + DriverInfoData->DriverInfo.Header.Length)) {\r
+      if (RemainingSize >= DriverInfoData->DriverInfo.Header.Length) {\r
         DriverInfo = ProfileBuffer;\r
         CopyMem (DriverInfo, &DriverInfoData->DriverInfo, sizeof (MEMORY_PROFILE_DRIVER_INFO));\r
-        RemainingSize -= sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
-        ProfileBuffer = (UINT8 *) ProfileBuffer + sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
+        if (DriverInfo->PdbStringOffset != 0) {\r
+          PdbSize = AsciiStrSize (DriverInfoData->PdbString);\r
+          CopyMem ((VOID *) ((UINTN) DriverInfo + DriverInfo->PdbStringOffset), DriverInfoData->PdbString, PdbSize);\r
+        }\r
+        RemainingSize -= DriverInfo->Header.Length;\r
+        ProfileBuffer = (UINT8 *) ProfileBuffer + DriverInfo->Header.Length;\r
       } else {\r
         goto Done;\r
       }\r
     }\r
-    Offset += sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
+    Offset += DriverInfoData->DriverInfo.Header.Length;\r
 \r
     AllocInfoList = DriverInfoData->AllocInfoList;\r
     for (AllocLink = AllocInfoList->ForwardLink;\r
@@ -1344,17 +1751,21 @@ SmramProfileCopyData (
                         Link,\r
                         MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
                         );\r
-      if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_ALLOC_INFO))) {\r
-        if (RemainingSize >= sizeof (MEMORY_PROFILE_ALLOC_INFO)) {\r
+      if (*ProfileOffset < (Offset + AllocInfoData->AllocInfo.Header.Length)) {\r
+        if (RemainingSize >= AllocInfoData->AllocInfo.Header.Length) {\r
           AllocInfo = ProfileBuffer;\r
           CopyMem (AllocInfo, &AllocInfoData->AllocInfo, sizeof (MEMORY_PROFILE_ALLOC_INFO));\r
-          RemainingSize -= sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
-          ProfileBuffer = (UINT8 *) ProfileBuffer + sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
+          if (AllocInfo->ActionStringOffset) {\r
+            ActionStringSize = AsciiStrSize (AllocInfoData->ActionString);\r
+            CopyMem ((VOID *) ((UINTN) AllocInfo + AllocInfo->ActionStringOffset), AllocInfoData->ActionString, ActionStringSize);\r
+          }\r
+          RemainingSize -= AllocInfo->Header.Length;\r
+          ProfileBuffer = (UINT8 *) ProfileBuffer + AllocInfo->Header.Length;\r
         } else {\r
           goto Done;\r
         }\r
       }\r
-      Offset += sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
+      Offset += AllocInfoData->AllocInfo.Header.Length;\r
     }\r
   }\r
 \r
@@ -1484,6 +1895,241 @@ Done:
   *ProfileOffset = Offset;\r
 }\r
 \r
+/**\r
+  Get memory profile data.\r
+\r
+  @param[in]      This              The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in, out] ProfileSize       On entry, points to the size in bytes of the ProfileBuffer.\r
+                                    On return, points to the size of the data returned in ProfileBuffer.\r
+  @param[out]     ProfileBuffer     Profile buffer.\r
+                      \r
+  @return EFI_SUCCESS               Get the memory profile data successfully.\r
+  @return EFI_UNSUPPORTED           Memory profile is unsupported.\r
+  @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data. \r
+                                    ProfileSize is updated with the size required.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolGetData (\r
+  IN     EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN OUT UINT64                             *ProfileSize,\r
+     OUT VOID                               *ProfileBuffer\r
+  )\r
+{\r
+  UINT64                                Size;\r
+  UINT64                                Offset;\r
+  MEMORY_PROFILE_CONTEXT_DATA           *ContextData;\r
+  BOOLEAN                               SmramProfileGettingStatus;\r
+\r
+  ContextData = GetSmramProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
+\r
+  Size = SmramProfileGetDataSize ();\r
+\r
+  if (*ProfileSize < Size) {\r
+    *ProfileSize = Size;\r
+    mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
+    return EFI_BUFFER_TOO_SMALL;\r
+  }\r
+\r
+  Offset = 0;\r
+  SmramProfileCopyData (ProfileBuffer, &Size, &Offset);\r
+  *ProfileSize = Size;\r
+\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Register image to memory profile.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] FilePath           File path of the image.\r
+  @param[in] ImageBase          Image base address.\r
+  @param[in] ImageSize          Image size.\r
+  @param[in] FileType           File type of the image.\r
+\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolRegisterImage (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,\r
+  IN PHYSICAL_ADDRESS                   ImageBase,\r
+  IN UINT64                             ImageSize,\r
+  IN EFI_FV_FILETYPE                    FileType\r
+  )\r
+{\r
+  EFI_STATUS                        Status;\r
+  EFI_SMM_DRIVER_ENTRY              DriverEntry;\r
+  VOID                              *EntryPointInImage;\r
+  EFI_GUID                          *Name;\r
+  \r
+  ZeroMem (&DriverEntry, sizeof (DriverEntry));\r
+  Name = GetFileNameFromFilePath (FilePath);\r
+  if (Name != NULL) {\r
+    CopyMem (&DriverEntry.FileName, Name, sizeof (EFI_GUID));\r
+  }\r
+  DriverEntry.ImageBuffer = ImageBase;\r
+  DriverEntry.NumberOfPage = EFI_SIZE_TO_PAGES ((UINTN) ImageSize);\r
+  Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) DriverEntry.ImageBuffer, &EntryPointInImage);\r
+  ASSERT_EFI_ERROR (Status);\r
+  DriverEntry.ImageEntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
+\r
+  return RegisterSmramProfileImage (&DriverEntry, FALSE);\r
+}\r
+\r
+/**\r
+  Unregister image from memory profile.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] FilePath           File path of the image.\r
+  @param[in] ImageBase          Image base address.\r
+  @param[in] ImageSize          Image size.\r
+\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_NOT_FOUND         The image is not found.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolUnregisterImage (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,\r
+  IN PHYSICAL_ADDRESS                   ImageBase,\r
+  IN UINT64                             ImageSize\r
+  )\r
+{\r
+  EFI_STATUS                        Status;\r
+  EFI_SMM_DRIVER_ENTRY              DriverEntry;\r
+  VOID                              *EntryPointInImage;\r
+  EFI_GUID                          *Name;\r
+\r
+  ZeroMem (&DriverEntry, sizeof (DriverEntry));\r
+  Name = GetFileNameFromFilePath (FilePath);\r
+  if (Name != NULL) {\r
+    CopyMem (&DriverEntry.FileName, Name, sizeof (EFI_GUID));\r
+  }\r
+  DriverEntry.ImageBuffer = ImageBase;\r
+  DriverEntry.NumberOfPage = EFI_SIZE_TO_PAGES ((UINTN) ImageSize);\r
+  Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) DriverEntry.ImageBuffer, &EntryPointInImage);\r
+  ASSERT_EFI_ERROR (Status);\r
+  DriverEntry.ImageEntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
+\r
+  return UnregisterSmramProfileImage (&DriverEntry, FALSE);\r
+}\r
+\r
+/**\r
+  Get memory profile recording state.\r
+\r
+  @param[in]  This              The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[out] RecordingState    Recording state.\r
+\r
+  @return EFI_SUCCESS           Memory profile recording state is returned.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+  @return EFI_INVALID_PARAMETER RecordingState is NULL.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolGetRecordingState (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  OUT BOOLEAN                           *RecordingState\r
+  )\r
+{\r
+  MEMORY_PROFILE_CONTEXT_DATA           *ContextData;\r
+\r
+  ContextData = GetSmramProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  if (RecordingState == NULL) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+  *RecordingState = mSmramProfileRecordingEnable;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Set memory profile recording state.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] RecordingState     Recording state.\r
+\r
+  @return EFI_SUCCESS           Set memory profile recording state successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolSetRecordingState (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN BOOLEAN                            RecordingState\r
+  )\r
+{\r
+  MEMORY_PROFILE_CONTEXT_DATA           *ContextData;\r
+\r
+  ContextData = GetSmramProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  mSmramProfileRecordingEnable = RecordingState;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Record memory profile of multilevel caller.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] CallerAddress      Address of caller.\r
+  @param[in] Action             Memory profile action.\r
+  @param[in] MemoryType         Memory type.\r
+                                EfiMaxMemoryType means the MemoryType is unknown.\r
+  @param[in] Buffer             Buffer address.\r
+  @param[in] Size               Buffer size.\r
+  @param[in] ActionString       String for memory profile action.\r
+                                Only needed for user defined allocate action.\r
+\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolRecord (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN PHYSICAL_ADDRESS                   CallerAddress,\r
+  IN MEMORY_PROFILE_ACTION              Action,\r
+  IN EFI_MEMORY_TYPE                    MemoryType,\r
+  IN VOID                               *Buffer,\r
+  IN UINTN                              Size,\r
+  IN CHAR8                              *ActionString OPTIONAL\r
+  )\r
+{\r
+  return SmmCoreUpdateProfile (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r
+}\r
+\r
 /**\r
   SMRAM profile handler to get profile info.\r
 \r
@@ -1496,20 +2142,20 @@ SmramProfileHandlerGetInfo (
   )\r
 {\r
   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
-  BOOLEAN                       SmramProfileRecordingStatus;\r
+  BOOLEAN                       SmramProfileGettingStatus;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
     return ;\r
   }\r
 \r
-  SmramProfileRecordingStatus = mSmramProfileRecordingStatus;\r
-  mSmramProfileRecordingStatus = FALSE;\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
 \r
   SmramProfileParameterGetInfo->ProfileSize = SmramProfileGetDataSize();\r
   SmramProfileParameterGetInfo->Header.ReturnStatus = 0;\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1527,15 +2173,15 @@ SmramProfileHandlerGetData (
   UINT64                                    ProfileOffset;\r
   SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA  SmramProfileGetData;\r
   MEMORY_PROFILE_CONTEXT_DATA               *ContextData;\r
-  BOOLEAN                                   SmramProfileRecordingStatus;\r
+  BOOLEAN                                   SmramProfileGettingStatus;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
     return ;\r
   }\r
 \r
-  SmramProfileRecordingStatus = mSmramProfileRecordingStatus;\r
-  mSmramProfileRecordingStatus = FALSE;\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
 \r
 \r
   CopyMem (&SmramProfileGetData, SmramProfileParameterGetData, sizeof (SmramProfileGetData));\r
@@ -1564,7 +2210,7 @@ SmramProfileHandlerGetData (
   SmramProfileParameterGetData->Header.ReturnStatus = 0;\r
 \r
 Done:\r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1580,15 +2226,15 @@ SmramProfileHandlerGetDataByOffset (
 {\r
   SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET    SmramProfileGetDataByOffset;\r
   MEMORY_PROFILE_CONTEXT_DATA                           *ContextData;\r
-  BOOLEAN                                               SmramProfileRecordingStatus;\r
+  BOOLEAN                                               SmramProfileGettingStatus;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
     return ;\r
   }\r
 \r
-  SmramProfileRecordingStatus = mSmramProfileRecordingStatus;\r
-  mSmramProfileRecordingStatus = FALSE;\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
 \r
 \r
   CopyMem (&SmramProfileGetDataByOffset, SmramProfileParameterGetDataByOffset, sizeof (SmramProfileGetDataByOffset));\r
@@ -1607,7 +2253,7 @@ SmramProfileHandlerGetDataByOffset (
   SmramProfileParameterGetDataByOffset->Header.ReturnStatus = 0;\r
 \r
 Done:\r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1624,7 +2270,6 @@ SmramProfileHandlerRegisterImage (
   EFI_STATUS                        Status;\r
   EFI_SMM_DRIVER_ENTRY              DriverEntry;\r
   VOID                              *EntryPointInImage;\r
-  BOOLEAN                           Ret;\r
 \r
   ZeroMem (&DriverEntry, sizeof (DriverEntry));\r
   CopyMem (&DriverEntry.FileName, &SmramProfileParameterRegisterImage->FileName, sizeof(EFI_GUID));\r
@@ -1634,8 +2279,8 @@ SmramProfileHandlerRegisterImage (
   ASSERT_EFI_ERROR (Status);\r
   DriverEntry.ImageEntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
 \r
-  Ret = RegisterSmramProfileImage (&DriverEntry, FALSE);\r
-  if (Ret) {\r
+  Status = RegisterSmramProfileImage (&DriverEntry, FALSE);\r
+  if (!EFI_ERROR (Status)) {\r
     SmramProfileParameterRegisterImage->Header.ReturnStatus = 0;\r
   }\r
 }\r
@@ -1654,7 +2299,6 @@ SmramProfileHandlerUnregisterImage (
   EFI_STATUS                        Status;\r
   EFI_SMM_DRIVER_ENTRY              DriverEntry;\r
   VOID                              *EntryPointInImage;\r
-  BOOLEAN                           Ret;\r
 \r
   ZeroMem (&DriverEntry, sizeof (DriverEntry));\r
   CopyMem (&DriverEntry.FileName, &SmramProfileParameterUnregisterImage->FileName, sizeof (EFI_GUID));\r
@@ -1664,8 +2308,8 @@ SmramProfileHandlerUnregisterImage (
   ASSERT_EFI_ERROR (Status);\r
   DriverEntry.ImageEntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
 \r
-  Ret = UnregisterSmramProfileImage (&DriverEntry, FALSE);\r
-  if (Ret) {\r
+  Status = UnregisterSmramProfileImage (&DriverEntry, FALSE);\r
+  if (!EFI_ERROR (Status)) {\r
     SmramProfileParameterUnregisterImage->Header.ReturnStatus = 0;\r
   }\r
 }\r
@@ -1697,6 +2341,7 @@ SmramProfileHandler (
 {\r
   SMRAM_PROFILE_PARAMETER_HEADER           *SmramProfileParameterHeader;\r
   UINTN                                    TempCommBufferSize;\r
+  SMRAM_PROFILE_PARAMETER_RECORDING_STATE  *ParameterRecordingState;\r
 \r
   DEBUG ((EFI_D_ERROR, "SmramProfileHandler Enter\n"));\r
 \r
@@ -1775,6 +2420,27 @@ SmramProfileHandler (
     }\r
     SmramProfileHandlerUnregisterImage ((SMRAM_PROFILE_PARAMETER_UNREGISTER_IMAGE *) (UINTN) CommBuffer);\r
     break;\r
+  case SMRAM_PROFILE_COMMAND_GET_RECORDING_STATE:\r
+    DEBUG ((EFI_D_ERROR, "SmramProfileHandlerGetRecordingState\n"));\r
+    if (TempCommBufferSize != sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE)) {\r
+      DEBUG ((EFI_D_ERROR, "SmramProfileHandler: SMM communication buffer size invalid!\n"));\r
+      return EFI_SUCCESS;\r
+    }\r
+    ParameterRecordingState = (SMRAM_PROFILE_PARAMETER_RECORDING_STATE *) (UINTN) CommBuffer;\r
+    ParameterRecordingState->RecordingState = mSmramProfileRecordingEnable;\r
+    ParameterRecordingState->Header.ReturnStatus = 0;\r
+    break;\r
+  case SMRAM_PROFILE_COMMAND_SET_RECORDING_STATE:\r
+    DEBUG ((EFI_D_ERROR, "SmramProfileHandlerSetRecordingState\n"));\r
+    if (TempCommBufferSize != sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE)) {\r
+      DEBUG ((EFI_D_ERROR, "SmramProfileHandler: SMM communication buffer size invalid!\n"));\r
+      return EFI_SUCCESS;\r
+    }\r
+    ParameterRecordingState = (SMRAM_PROFILE_PARAMETER_RECORDING_STATE *) (UINTN) CommBuffer;\r
+    mSmramProfileRecordingEnable = ParameterRecordingState->RecordingState;\r
+    ParameterRecordingState->Header.ReturnStatus = 0;\r
+    break;\r
+\r
   default:\r
     break;\r
   }\r
@@ -1821,15 +2487,15 @@ DumpSmramRange (
 {\r
   UINTN                         Index;\r
   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
-  BOOLEAN                       SmramProfileRecordingStatus;\r
+  BOOLEAN                       SmramProfileGettingStatus;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
     return ;\r
   }\r
 \r
-  SmramProfileRecordingStatus = mSmramProfileRecordingStatus;\r
-  mSmramProfileRecordingStatus = FALSE;\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
 \r
   DEBUG ((EFI_D_INFO, "FullSmramRange address - 0x%08x\n", mFullSmramRanges));\r
 \r
@@ -1846,7 +2512,7 @@ DumpSmramRange (
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile end =======\n"));\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1863,15 +2529,15 @@ DumpFreePagesList (
   FREE_PAGE_LIST                *Pages;\r
   UINTN                         Index;\r
   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
-  BOOLEAN                       SmramProfileRecordingStatus;\r
+  BOOLEAN                       SmramProfileGettingStatus;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
     return ;\r
   }\r
 \r
-  SmramProfileRecordingStatus = mSmramProfileRecordingStatus;\r
-  mSmramProfileRecordingStatus = FALSE;\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile begin =======\n"));\r
 \r
@@ -1888,7 +2554,7 @@ DumpFreePagesList (
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile end =======\n"));\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1900,21 +2566,21 @@ DumpFreePoolList (
   VOID\r
   )\r
 {\r
-  LIST_ENTRY                     *FreePoolList;\r
-  LIST_ENTRY                     *Node;\r
-  FREE_POOL_HEADER               *Pool;\r
-  UINTN                          Index;\r
-  UINTN                          PoolListIndex;\r
+  LIST_ENTRY                    *FreePoolList;\r
+  LIST_ENTRY                    *Node;\r
+  FREE_POOL_HEADER              *Pool;\r
+  UINTN                         Index;\r
+  UINTN                         PoolListIndex;\r
   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
-  BOOLEAN                       SmramProfileRecordingStatus;\r
+  BOOLEAN                       SmramProfileGettingStatus;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
     return ;\r
   }\r
 \r
-  SmramProfileRecordingStatus = mSmramProfileRecordingStatus;\r
-  mSmramProfileRecordingStatus = FALSE;\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile begin =======\n"));\r
 \r
@@ -1934,25 +2600,56 @@ DumpFreePoolList (
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile end =======\n"));\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
-GLOBAL_REMOVE_IF_UNREFERENCED CHAR16 *mActionString[] = {\r
-  L"Unknown",\r
-  L"AllocatePages",\r
-  L"FreePages",\r
-  L"AllocatePool",\r
-  L"FreePool",\r
+GLOBAL_REMOVE_IF_UNREFERENCED CHAR8 *mSmmActionString[] = {\r
+  "SmmUnknown",\r
+  "gSmst->SmmAllocatePages",\r
+  "gSmst->SmmFreePages",\r
+  "gSmst->SmmAllocatePool",\r
+  "gSmst->SmmFreePool",\r
 };\r
 \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
+GLOBAL_REMOVE_IF_UNREFERENCED CHAR8 mUserDefinedActionString[] = {"UserDefined-0x80000000"};\r
+\r
 typedef struct {\r
   EFI_MEMORY_TYPE   MemoryType;\r
-  CHAR16            *MemoryTypeStr;\r
+  CHAR            *MemoryTypeStr;\r
 } PROFILE_MEMORY_TYPE_STRING;\r
 \r
 GLOBAL_REMOVE_IF_UNREFERENCED PROFILE_MEMORY_TYPE_STRING mMemoryTypeString[] = {\r
-  {EfiRuntimeServicesCode, L"EfiRuntimeServicesCode"},\r
-  {EfiRuntimeServicesData, L"EfiRuntimeServicesData"}\r
+  {EfiRuntimeServicesCode, "EfiRuntimeServicesCode"},\r
+  {EfiRuntimeServicesData, "EfiRuntimeServicesData"}\r
 };\r
 \r
 /**\r
@@ -1963,7 +2660,7 @@ GLOBAL_REMOVE_IF_UNREFERENCED PROFILE_MEMORY_TYPE_STRING mMemoryTypeString[] = {
   @return Pointer to string.\r
 \r
 **/\r
-CHAR16 *\r
+CHAR8 *\r
 ProfileMemoryTypeToStr (\r
   IN EFI_MEMORY_TYPE    MemoryType\r
   )\r
@@ -1975,7 +2672,39 @@ ProfileMemoryTypeToStr (
     }\r
   }\r
 \r
-  return L"UnexpectedMemoryType";\r
+  return "UnexpectedMemoryType";\r
+}\r
+\r
+/**\r
+  Action to string.\r
+\r
+  @param[in] Action                     Profile action.\r
+\r
+  @return Pointer to string.\r
+\r
+**/\r
+CHAR8 *\r
+ProfileActionToStr (\r
+  IN MEMORY_PROFILE_ACTION  Action\r
+  )\r
+{\r
+  UINTN     Index;\r
+  UINTN     ActionStringCount;\r
+  CHAR8     **ActionString;\r
+\r
+  ActionString = mSmmActionString;\r
+  ActionStringCount = sizeof (mSmmActionString) / sizeof (mSmmActionString[0]);\r
+\r
+  if ((UINTN) (UINT32) Action < ActionStringCount) {\r
+    return ActionString[Action];\r
+  }\r
+  for (Index = 0; Index < sizeof (mExtActionString) / sizeof (mExtActionString[0]); Index++) {\r
+    if (mExtActionString[Index].Action == Action) {\r
+      return mExtActionString[Index].String;\r
+    }\r
+  }\r
+\r
+  return ActionString[0];\r
 }\r
 \r
 /**\r
@@ -1999,7 +2728,7 @@ DumpSmramProfile (
   LIST_ENTRY                        *AllocInfoList;\r
   UINTN                             AllocIndex;\r
   LIST_ENTRY                        *AllocLink;\r
-  BOOLEAN                           SmramProfileRecordingStatus;\r
+  BOOLEAN                           SmramProfileGettingStatus;\r
   UINTN                      &nbs