diff --git a/Silicon/Hisilicon/Hi1610/Drivers/PcieInit1610/PcieInitDxe.inf b/Silicon/Hisilicon/Hi1610/Drivers/PcieInit1610/PcieInitDxe.inf
index 94d0fc8c028b..d1efdf39131a 100644
--- a/Silicon/Hisilicon/Hi1610/Drivers/PcieInit1610/PcieInitDxe.inf
+++ b/Silicon/Hisilicon/Hi1610/Drivers/PcieInit1610/PcieInitDxe.inf
@@ -48,10 +48,11 @@ [LibraryClasses]
   OemMiscLib
 
 [Protocols]
-  #gEfiPcieRootBridgeProtocolGuid
   gEfiFirmwareVolume2ProtocolGuid
+  gEfiPciHostBridgeResourceAllocationProtocolGuid
   gEfiPciIoProtocolGuid
   gEfiPciPlatformProtocolGuid
+  gEfiPciRootBridgeIoProtocolGuid
 
 [Pcd]
   gHisiTokenSpaceGuid.PcdPcieRootBridgeMask
diff --git a/Silicon/Hisilicon/Drivers/PciHostBridgeDxe/PciHostBridge.h b/Silicon/Hisilicon/Drivers/PciHostBridgeDxe/PciHostBridge.h
index c04361fceee6..435385491a17 100644
--- a/Silicon/Hisilicon/Drivers/PciHostBridgeDxe/PciHostBridge.h
+++ b/Silicon/Hisilicon/Drivers/PciHostBridgeDxe/PciHostBridge.h
@@ -401,10 +401,6 @@ PreprocessController (
 #define EFI_RESOURCE_NONEXISTENT   0xFFFFFFFFFFFFFFFFULL
 #define EFI_RESOURCE_LESS          0xFFFFFFFFFFFFFFFEULL
 
-#define	INVALID_CAPABILITY_00       0x00
-#define	INVALID_CAPABILITY_FF       0xFF
-#define	PCI_CAPABILITY_POINTER_MASK 0xFC
-
 //
 // Driver Instance Data Prototypes
 //
@@ -521,8 +517,4 @@ RootBridgeConstructor (
   IN UINT32                             Seg
   );
 
-VOID
-EnlargeAtuConfig0 (
-  IN EFI_PCI_HOST_BRIDGE_RESOURCE_ALLOCATION_PROTOCOL *This
-  );
 #endif
diff --git a/Silicon/Hisilicon/Hi1610/Drivers/PcieInit1610/PcieInitLib.h b/Silicon/Hisilicon/Hi1610/Drivers/PcieInit1610/PcieInitLib.h
index ab0c7ab8bfa7..ead926b4c4f3 100644
--- a/Silicon/Hisilicon/Hi1610/Drivers/PcieInit1610/PcieInitLib.h
+++ b/Silicon/Hisilicon/Hi1610/Drivers/PcieInit1610/PcieInitLib.h
@@ -248,6 +248,11 @@ EFI_STATUS PcieSetDBICS2Enable(UINT32 HostBridgeNum, UINT32 Port, UINT32 Enable)
 
 VOID InitAtu (PCI_ROOT_BRIDGE_RESOURCE_APPETURE *Private);
 
+VOID
+EnlargeAtuConfig0 (
+  IN EFI_HANDLE HostBridge
+  );
+
 EFI_STATUS
 PciPlatformDriverEntry (
   IN EFI_HANDLE                         ImageHandle,
diff --git a/Silicon/Hisilicon/Drivers/PciHostBridgeDxe/PciHostBridge.c b/Silicon/Hisilicon/Drivers/PciHostBridgeDxe/PciHostBridge.c
index 9fa3f8409813..e3d3988a64c1 100644
--- a/Silicon/Hisilicon/Drivers/PciHostBridgeDxe/PciHostBridge.c
+++ b/Silicon/Hisilicon/Drivers/PciHostBridgeDxe/PciHostBridge.c
@@ -839,7 +839,6 @@ NotifyPhase(
 
   case EfiPciHostBridgeEndEnumeration:
     PCIE_DEBUG("Case EfiPciHostBridgeEndEnumeration\n");
-    EnlargeAtuConfig0 (This);
     break;
 
   case EfiPciHostBridgeBeginBusAllocation:
diff --git a/Silicon/Hisilicon/Drivers/PciHostBridgeDxe/PciRootBridgeIo.c b/Silicon/Hisilicon/Drivers/PciHostBridgeDxe/PciRootBridgeIo.c
index 273a322ee48f..3c265ea43378 100644
--- a/Silicon/Hisilicon/Drivers/PciHostBridgeDxe/PciRootBridgeIo.c
+++ b/Silicon/Hisilicon/Drivers/PciHostBridgeDxe/PciRootBridgeIo.c
@@ -2218,81 +2218,3 @@ RootBridgeIoConfiguration (
   return EFI_SUCCESS;
 }
 
-BOOLEAN
-PcieCheckAriFwdEn (
-  UINTN  PciBaseAddr
-  )
-{
-  UINT8   PciPrimaryStatus;
-  UINT8   CapabilityOffset;
-  UINT8   CapId;
-  UINT8   TempData;
-
-  PciPrimaryStatus = MmioRead16 (PciBaseAddr + PCI_PRIMARY_STATUS_OFFSET);
-
-  if (PciPrimaryStatus & EFI_PCI_STATUS_CAPABILITY) {
-    CapabilityOffset = MmioRead8 (PciBaseAddr + PCI_CAPBILITY_POINTER_OFFSET);
-    CapabilityOffset &= PCI_CAPABILITY_POINTER_MASK;
-
-    while ((CapabilityOffset != INVALID_CAPABILITY_00) && (CapabilityOffset != INVALID_CAPABILITY_FF)) {
-      CapId = MmioRead8 (PciBaseAddr + CapabilityOffset);
-      if (CapId == EFI_PCI_CAPABILITY_ID_PCIEXP) {
-        break;
-      }
-      CapabilityOffset = MmioRead8 (PciBaseAddr + CapabilityOffset + 1);
-      CapabilityOffset &= PCI_CAPABILITY_POINTER_MASK;
-    }
-  } else {
-    PCIE_DEBUG ("[%a:%d] - No PCIE Capability.\n", __FUNCTION__, __LINE__);
-    return FALSE;
-  }
-
-  if ((CapabilityOffset == INVALID_CAPABILITY_FF) || (CapabilityOffset == INVALID_CAPABILITY_00)) {
-    PCIE_DEBUG ("[%a:%d] - No PCIE Capability.\n", __FUNCTION__, __LINE__);
-    return FALSE;
-  }
-
-  TempData = MmioRead16 (PciBaseAddr + CapabilityOffset +
-                          EFI_PCIE_CAPABILITY_DEVICE_CONTROL_2_OFFSET);
-  TempData &= EFI_PCIE_CAPABILITY_DEVICE_CAPABILITIES_2_ARI_FORWARDING;
-
-  if (TempData == EFI_PCIE_CAPABILITY_DEVICE_CAPABILITIES_2_ARI_FORWARDING) {
-    return TRUE;
-  } else {
-    return FALSE;
-  }
-}
-
-VOID
-EnlargeAtuConfig0 (
-  IN EFI_PCI_HOST_BRIDGE_RESOURCE_ALLOCATION_PROTOCOL *This
-  )
-{
-  UINTN                           RbPciBase;
-  UINT64                          MemLimit;
-  LIST_ENTRY                      *List;
-  PCI_HOST_BRIDGE_INSTANCE        *HostBridgeInstance;
-  PCI_ROOT_BRIDGE_INSTANCE        *RootBridgeInstance;
-
-  PCIE_DEBUG ("In Enlarge RP iatu Config 0.\n");
-
-  HostBridgeInstance = INSTANCE_FROM_RESOURCE_ALLOCATION_THIS (This);
-  List = HostBridgeInstance->Head.ForwardLink;
-
-  while (List != &HostBridgeInstance->Head) {
-    PCIE_DEBUG ("HostBridge has data.\n");
-    RootBridgeInstance = DRIVER_INSTANCE_FROM_LIST_ENTRY (List);
-
-    RbPciBase = RootBridgeInstance->RbPciBar;
-
-    // Those ARI FWD Enable Root Bridge, need enlarge iatu window.
-    if (PcieCheckAriFwdEn (RbPciBase)) {
-      MemLimit = GetPcieCfgAddress (RootBridgeInstance->Ecam,
-                                    RootBridgeInstance->BusBase + 2, 0, 0, 0)
-	               - 1;
-      MmioWrite32 (RbPciBase + IATU_OFFSET + IATU_VIEW_POINT, 1);
-      MmioWrite32 (RbPciBase + IATU_OFFSET + IATU_REGION_BASE_LIMIT, (UINT32) MemLimit);
-    }
-    List = List->ForwardLink;
-  }
-}
diff --git a/Silicon/Hisilicon/Hi1610/Drivers/PcieInit1610/PcieInitAtu.c b/Silicon/Hisilicon/Hi1610/Drivers/PcieInit1610/PcieInitAtu.c
index 9eb3451f7402..f2365b5f3e49 100644
--- a/Silicon/Hisilicon/Hi1610/Drivers/PcieInit1610/PcieInitAtu.c
+++ b/Silicon/Hisilicon/Hi1610/Drivers/PcieInit1610/PcieInitAtu.c
@@ -13,13 +13,20 @@
 *
 **/
 
+#include <IndustryStandard/Acpi.h>
 #include <Library/OemMiscLib.h>
 #include <Library/PcdLib.h>
 #include <Library/PciExpressLib.h>
 #include <Library/PlatformPciLib.h>
 #include <Library/UefiBootServicesTableLib.h>
+#include <Protocol/PciRootBridgeIo.h>
+#include <Protocol/PciHostBridgeResourceAllocation.h>
 #include "PcieInit.h"
 
+#define INVALID_CAPABILITY_00       0x00
+#define INVALID_CAPABILITY_FF       0xFF
+#define PCI_CAPABILITY_POINTER_MASK 0xFC
+
 STATIC
 UINT64
 GetPcieCfgAddress (
@@ -34,6 +41,53 @@ GetPcieCfgAddress (
 }
 
 STATIC
+PCI_ROOT_BRIDGE_RESOURCE_APPETURE *
+GetAppetureByRootBridgeIo (
+    IN  EFI_PCI_ROOT_BRIDGE_IO_PROTOCOL  *RootBridge
+    )
+{
+  EFI_STATUS Status;
+  EFI_ACPI_ADDRESS_SPACE_DESCRIPTOR *Configuration = NULL;
+  UINTN Hb;
+  UINTN Rb;
+
+  Status = RootBridge->Configuration (
+      RootBridge,
+      (VOID **)&Configuration
+      );
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "[%a:%d] RootBridgeIo->Configuration failed %r\n",
+          __FUNCTION__, __LINE__, Status));
+    return NULL;
+  };
+
+  while (Configuration->Desc == ACPI_ADDRESS_SPACE_DESCRIPTOR) {
+    if (Configuration->ResType == ACPI_ADDRESS_SPACE_TYPE_BUS) {
+      break;
+    }
+    Configuration++;
+  }
+
+  if (Configuration->Desc != ACPI_ADDRESS_SPACE_DESCRIPTOR) {
+    DEBUG ((DEBUG_ERROR, "[%a:%d] Can't find bus descriptor\n", __FUNCTION__, __LINE__));
+    return NULL;
+  }
+
+  for (Hb = 0; Hb < PCIE_MAX_HOSTBRIDGE; Hb++) {
+    for (Rb = 0; Rb < PCIE_MAX_ROOTBRIDGE; Rb++) {
+      if (RootBridge->SegmentNumber == mResAppeture[Hb][Rb].Segment &&
+          Configuration->AddrRangeMin >= mResAppeture[Hb][Rb].BusBase &&
+          Configuration->AddrRangeMax <= mResAppeture[Hb][Rb].BusLimit) {
+        return &mResAppeture[Hb][Rb];
+      }
+    }
+  }
+
+  DEBUG ((DEBUG_ERROR, "[%a:%d] Can't find PCI appeture\n", __FUNCTION__, __LINE__));
+  return NULL;
+}
+
+STATIC
 VOID
 SetAtuConfig0RW (
     PCI_ROOT_BRIDGE_RESOURCE_APPETURE *Private,
@@ -61,7 +115,9 @@ SetAtuConfig0RW (
   }
 }
 
-void SetAtuConfig1RW (
+STATIC
+VOID
+SetAtuConfig1RW (
     PCI_ROOT_BRIDGE_RESOURCE_APPETURE *Private,
     UINT32 Index
     )
@@ -87,7 +143,9 @@ void SetAtuConfig1RW (
   }
 }
 
-void SetAtuIoRW(UINT64 RbPciBase,UINT64 IoBase,UINT64 CpuIoRegionLimit, UINT64 CpuIoRegionBase, UINT32 Index)
+STATIC
+VOID
+SetAtuIoRW (UINT64 RbPciBase,UINT64 IoBase,UINT64 CpuIoRegionLimit, UINT64 CpuIoRegionBase, UINT32 Index)
 {
 
     MmioWrite32 (RbPciBase + IATU_OFFSET + IATU_VIEW_POINT, Index);
@@ -109,7 +167,9 @@ void SetAtuIoRW(UINT64 RbPciBase,UINT64 IoBase,UINT64 CpuIoRegionLimit, UINT64 C
     }
 }
 
-void SetAtuMemRW(UINT64 RbPciBase,UINT64 MemBase,UINT64 CpuMemRegionLimit, UINT64 CpuMemRegionBase, UINT32 Index)
+STATIC
+VOID
+SetAtuMemRW(UINT64 RbPciBase,UINT64 MemBase,UINT64 CpuMemRegionLimit, UINT64 CpuMemRegionBase, UINT32 Index)
 {
 
     MmioWrite32 (RbPciBase + IATU_OFFSET + IATU_VIEW_POINT, Index);
@@ -131,7 +191,8 @@ void SetAtuMemRW(UINT64 RbPciBase,UINT64 MemBase,UINT64 CpuMemRegionLimit, UINT6
     }
 }
 
-VOID InitAtu (PCI_ROOT_BRIDGE_RESOURCE_APPETURE *Private)
+VOID
+InitAtu (PCI_ROOT_BRIDGE_RESOURCE_APPETURE *Private)
 {
   SetAtuMemRW (Private->RbPciBar, Private->PciRegionBase, Private->PciRegionLimit, Private->CpuMemRegionBase, 0);
   SetAtuConfig0RW (Private, 1);
@@ -139,3 +200,110 @@ VOID InitAtu (PCI_ROOT_BRIDGE_RESOURCE_APPETURE *Private)
   SetAtuIoRW (Private->RbPciBar, Private->IoBase, Private->IoLimit, Private->CpuIoRegionBase, 3);
 }
 
+STATIC
+BOOLEAN
+PcieCheckAriFwdEn (
+  UINTN  PciBaseAddr
+  )
+{
+  UINT8   PciPrimaryStatus;
+  UINT8   CapabilityOffset;
+  UINT8   CapId;
+  UINT8   TempData;
+
+  PciPrimaryStatus = MmioRead16 (PciBaseAddr + PCI_PRIMARY_STATUS_OFFSET);
+
+  if (PciPrimaryStatus & EFI_PCI_STATUS_CAPABILITY) {
+    CapabilityOffset = MmioRead8 (PciBaseAddr + PCI_CAPBILITY_POINTER_OFFSET);
+    CapabilityOffset &= PCI_CAPABILITY_POINTER_MASK;
+
+    while ((CapabilityOffset != INVALID_CAPABILITY_00) && (CapabilityOffset != INVALID_CAPABILITY_FF)) {
+      CapId = MmioRead8 (PciBaseAddr + CapabilityOffset);
+      if (CapId == EFI_PCI_CAPABILITY_ID_PCIEXP) {
+        break;
+      }
+      CapabilityOffset = MmioRead8 (PciBaseAddr + CapabilityOffset + 1);
+      CapabilityOffset &= PCI_CAPABILITY_POINTER_MASK;
+    }
+  } else {
+    return FALSE;
+  }
+
+  if ((CapabilityOffset == INVALID_CAPABILITY_FF) || (CapabilityOffset == INVALID_CAPABILITY_00)) {
+    return FALSE;
+  }
+
+  TempData = MmioRead16 (PciBaseAddr + CapabilityOffset +
+                          EFI_PCIE_CAPABILITY_DEVICE_CONTROL_2_OFFSET);
+  TempData &= EFI_PCIE_CAPABILITY_DEVICE_CAPABILITIES_2_ARI_FORWARDING;
+
+  if (TempData == EFI_PCIE_CAPABILITY_DEVICE_CAPABILITIES_2_ARI_FORWARDING) {
+    return TRUE;
+  } else {
+    return FALSE;
+  }
+}
+
+VOID
+EnlargeAtuConfig0 (
+  IN EFI_HANDLE HostBridge
+  )
+{
+  EFI_PCI_HOST_BRIDGE_RESOURCE_ALLOCATION_PROTOCOL    *ResAlloc = NULL;
+  EFI_STATUS                                          Status;
+  EFI_HANDLE RootBridgeHandle = NULL;
+  EFI_PCI_ROOT_BRIDGE_IO_PROTOCOL *RootBridgeIo = NULL;
+  PCI_ROOT_BRIDGE_RESOURCE_APPETURE *Appeture;
+  UINTN                           RbPciBase;
+  UINT64                          MemLimit;
+
+  DEBUG ((DEBUG_INFO, "In Enlarge RP iATU Config 0.\n"));
+
+  Status = gBS->HandleProtocol (
+      HostBridge,
+      &gEfiPciHostBridgeResourceAllocationProtocolGuid,
+      (VOID **)&ResAlloc
+      );
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "[%a:%d] - HandleProtocol failed %r\n", __FUNCTION__,
+          __LINE__, Status));
+    return;
+  }
+
+  while (TRUE) {
+    Status = ResAlloc->GetNextRootBridge (
+        ResAlloc,
+        &RootBridgeHandle
+        );
+    if (EFI_ERROR (Status)) {
+      break;
+    }
+    Status = gBS->HandleProtocol (
+        RootBridgeHandle,
+        &gEfiPciRootBridgeIoProtocolGuid,
+        (VOID **)&RootBridgeIo
+        );
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR, "[%a:%d] - HandleProtocol failed %r\n", __FUNCTION__, __LINE__, Status));
+      // This should never happen so that it is a fatal error and we don't try
+      // to continue
+      break;
+    }
+
+    Appeture = GetAppetureByRootBridgeIo (RootBridgeIo);
+    if (Appeture == NULL) {
+      DEBUG ((DEBUG_ERROR, "[%a:%d] Get appeture failed\n", __FUNCTION__,
+            __LINE__));
+      continue;
+    }
+
+    RbPciBase = Appeture->RbPciBar;
+    // Those ARI FWD Enable Root Bridge, need enlarge iATU window.
+    if (PcieCheckAriFwdEn (RbPciBase)) {
+      MemLimit = GetPcieCfgAddress (Appeture->Ecam, Appeture->BusBase + 2, 0, 0, 0) - 1;
+      MmioWrite32 (RbPciBase + IATU_OFFSET + IATU_VIEW_POINT, 1);
+      MmioWrite32 (RbPciBase + IATU_OFFSET + IATU_REGION_BASE_LIMIT, (UINT32) MemLimit);
+    }
+  }
+}
+
