return (addr >> ilog2(SZ_1M)) & 0x1f;
 }
 
+static u64 rockchip_pcie_ep_align_addr(struct pci_epc *epc, u64 pci_addr,
+                                      size_t *pci_size, size_t *addr_offset)
+{
+       struct rockchip_pcie_ep *ep = epc_get_drvdata(epc);
+       size_t size = *pci_size;
+       u64 offset, mask;
+       int num_bits;
+
+       num_bits = rockchip_pcie_ep_ob_atu_num_bits(&ep->rockchip,
+                                                   pci_addr, size);
+       mask = (1ULL << num_bits) - 1;
+
+       offset = pci_addr & mask;
+       if (size + offset > SZ_1M)
+               size = SZ_1M - offset;
+
+       *pci_size = ALIGN(offset + size, ROCKCHIP_PCIE_AT_SIZE_ALIGN);
+       *addr_offset = offset;
+
+       return pci_addr & ~mask;
+}
+
 static int rockchip_pcie_ep_map_addr(struct pci_epc *epc, u8 fn, u8 vfn,
                                     phys_addr_t addr, u64 pci_addr,
                                     size_t size)
        .write_header   = rockchip_pcie_ep_write_header,
        .set_bar        = rockchip_pcie_ep_set_bar,
        .clear_bar      = rockchip_pcie_ep_clear_bar,
+       .align_addr     = rockchip_pcie_ep_align_addr,
        .map_addr       = rockchip_pcie_ep_map_addr,
        .unmap_addr     = rockchip_pcie_ep_unmap_addr,
        .set_msi        = rockchip_pcie_ep_set_msi,
 
 #define   ROCKCHIP_PCIE_EP_MSIX_CAP_CP_MASK            GENMASK(15, 8)
 #define ROCKCHIP_PCIE_EP_DUMMY_IRQ_ADDR                                0x1
 #define ROCKCHIP_PCIE_EP_PCI_LEGACY_IRQ_ADDR           0x3
+
+#define ROCKCHIP_PCIE_AT_MIN_NUM_BITS  8
+#define ROCKCHIP_PCIE_AT_MAX_NUM_BITS  20
+#define ROCKCHIP_PCIE_AT_SIZE_ALIGN    (1UL << ROCKCHIP_PCIE_AT_MIN_NUM_BITS)
+
 #define ROCKCHIP_PCIE_EP_FUNC_BASE(fn) \
        (PCIE_EP_PF_CONFIG_REGS_BASE + (((fn) << 12) & GENMASK(19, 12)))
 #define ROCKCHIP_PCIE_EP_VIRT_FUNC_BASE(fn) \