diff --git a/MdeModulePkg/Bus/Ufs/UfsBlockIoPei/DmaMem.c b/MdeModulePkg/Bus/Ufs/UfsBlockIoPei/DmaMem.c
index 62aeaf86c8..5aa046873c 100644
--- a/MdeModulePkg/Bus/Ufs/UfsBlockIoPei/DmaMem.c
+++ b/MdeModulePkg/Bus/Ufs/UfsBlockIoPei/DmaMem.c
@@ -2,7 +2,7 @@
   The DMA memory help function.
 
   Copyright (c) 2017, Intel Corporation. All rights reserved.<BR>
-
+  Copyright (c) 1985 - 2022, American Megatrends International LLC. <BR>
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -240,10 +240,16 @@ IoMmuInit (
   VOID
   )
 {
-  PeiServicesLocatePpi (
-    &gEdkiiIoMmuPpiGuid,
-    0,
-    NULL,
-    (VOID **)&mIoMmu
-    );
+  EFI_STATUS  Status;
+
+  Status =   PeiServicesLocatePpi (
+               &gEdkiiIoMmuPpiGuid,
+               0,
+               NULL,
+               (VOID **)&mIoMmu
+               );
+
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_INFO, "Locate mIoMmu Ppi is failed!!!\n"));
+  }
 }
diff --git a/MdeModulePkg/Bus/Ufs/UfsBlockIoPei/UfsHcMem.c b/MdeModulePkg/Bus/Ufs/UfsBlockIoPei/UfsHcMem.c
index b7dd9b6acd..562a77676c 100644
--- a/MdeModulePkg/Bus/Ufs/UfsBlockIoPei/UfsHcMem.c
+++ b/MdeModulePkg/Bus/Ufs/UfsBlockIoPei/UfsHcMem.c
@@ -1,6 +1,7 @@
 /** @file
 
 Copyright (c) 2014 - 2018, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 1985 - 2022, American Megatrends International LLC. <BR>
 
 SPDX-License-Identifier: BSD-2-Clause-Patent
 
@@ -18,7 +19,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 **/
 UFS_PEIM_MEM_BLOCK *
 UfsPeimAllocMemBlock (
-  IN  UINTN  Pages
+  IN UINTN  Pages
   )
 {
   UFS_PEIM_MEM_BLOCK    *Block;
@@ -107,8 +108,8 @@ UfsPeimFreeMemBlock (
 **/
 VOID *
 UfsPeimAllocMemFromBlock (
-  IN  UFS_PEIM_MEM_BLOCK  *Block,
-  IN  UINTN               Units
+  IN UFS_PEIM_MEM_BLOCK  *Block,
+  IN UINTN               Units
   )
 {
   UINTN  Byte;
@@ -221,7 +222,7 @@ UfsPeimIsMemBlockEmpty (
 **/
 EFI_STATUS
 UfsPeimInitMemPool (
-  IN  UFS_PEIM_HC_PRIVATE_DATA  *Private
+  IN UFS_PEIM_HC_PRIVATE_DATA  *Private
   )
 {
   UFS_PEIM_MEM_POOL  *Pool;
@@ -292,8 +293,8 @@ UfsPeimFreeMemPool (
 **/
 VOID *
 UfsPeimAllocateMem (
-  IN  UFS_PEIM_MEM_POOL  *Pool,
-  IN  UINTN              Size
+  IN UFS_PEIM_MEM_POOL  *Pool,
+  IN UINTN              Size
   )
 {
   UFS_PEIM_MEM_BLOCK  *Head;
@@ -414,6 +415,10 @@ UfsPeimFreeMem (
   //
   ASSERT (Block != NULL);
 
+  if (Block == NULL) {
+    return;
+  }
+
   //
   // Release the current memory block if it is empty and not the head
   //