X-Git-Url: https://git.proxmox.com/?p=mirror_edk2.git;a=blobdiff_plain;f=MdeModulePkg%2FCore%2FPiSmmCore%2FDispatcher.c;h=178681ec90aecc58fd7d46c093afc3943dfefc74;hp=a4686306f6e2de9bde78f69cb87806cb6048c437;hb=84edd20bd0756ef5719835498d4283435d6b5e77;hpb=dad83a8c1207e585b2cb417d0680fe2e9452d262

diff --git a/MdeModulePkg/Core/PiSmmCore/Dispatcher.c b/MdeModulePkg/Core/PiSmmCore/Dispatcher.c
index a4686306f6..178681ec90 100644
--- a/MdeModulePkg/Core/PiSmmCore/Dispatcher.c
+++ b/MdeModulePkg/Core/PiSmmCore/Dispatcher.c
@@ -874,10 +874,12 @@ SmmDispatcher (
       //
       // For each SMM driver, pass NULL as ImageHandle
       //
+      RegisterSmramProfileImage (DriverEntry, TRUE);
       PERF_START (DriverEntry->ImageHandle, "StartImage:", NULL, 0);
       Status = ((EFI_IMAGE_ENTRY_POINT)(UINTN)DriverEntry->ImageEntryPoint)(DriverEntry->ImageHandle, gST);
       PERF_END (DriverEntry->ImageHandle, "StartImage:", NULL, 0);
       if (EFI_ERROR(Status)){
+        UnregisterSmramProfileImage (DriverEntry, TRUE);
         SmmFreePages(DriverEntry->ImageBuffer, DriverEntry->NumberOfPage);
       }