#define PHY_PLL_LOCK_WAIT_USLEEP_MAX   200
 #define PHY_PLL_LOCK_WAIT_TIMEOUT      (2000 * PHY_PLL_LOCK_WAIT_USLEEP_MAX)
 
-/* PCIe Root Complex registers (memory-mapped) */
-#define PCIE_RC_IMX6_MSI_CAP                   0x50
-#define PCIE_RC_LCR                            0x7c
-#define PCIE_RC_LCR_MAX_LINK_SPEEDS_GEN1       0x1
-#define PCIE_RC_LCR_MAX_LINK_SPEEDS_GEN2       0x2
-#define PCIE_RC_LCR_MAX_LINK_SPEEDS_MASK       0xf
-
-#define PCIE_RC_LCSR                           0x80
-
 /* PCIe Port Logic registers (memory-mapped) */
 #define PL_OFFSET 0x700
 
 {
        struct dw_pcie *pci = imx6_pcie->pci;
        struct device *dev = pci->dev;
+       u8 offset = dw_pcie_find_capability(pci, PCI_CAP_ID_EXP);
        u32 tmp;
        int ret;
 
         * started in Gen2 mode, there is a possibility the devices on the
         * bus will not be detected at all.  This happens with PCIe switches.
         */
-       tmp = dw_pcie_readl_dbi(pci, PCIE_RC_LCR);
-       tmp &= ~PCIE_RC_LCR_MAX_LINK_SPEEDS_MASK;
-       tmp |= PCIE_RC_LCR_MAX_LINK_SPEEDS_GEN1;
-       dw_pcie_writel_dbi(pci, PCIE_RC_LCR, tmp);
+       tmp = dw_pcie_readl_dbi(pci, offset + PCI_EXP_LNKCAP);
+       tmp &= ~PCI_EXP_LNKCAP_SLS;
+       tmp |= PCI_EXP_LNKCAP_SLS_2_5GB;
+       dw_pcie_writel_dbi(pci, offset + PCI_EXP_LNKCAP, tmp);
 
        /* Start LTSSM. */
        imx6_pcie_ltssm_enable(dev);
 
        if (imx6_pcie->link_gen == 2) {
                /* Allow Gen2 mode after the link is up. */
-               tmp = dw_pcie_readl_dbi(pci, PCIE_RC_LCR);
-               tmp &= ~PCIE_RC_LCR_MAX_LINK_SPEEDS_MASK;
-               tmp |= PCIE_RC_LCR_MAX_LINK_SPEEDS_GEN2;
-               dw_pcie_writel_dbi(pci, PCIE_RC_LCR, tmp);
+               tmp = dw_pcie_readl_dbi(pci, offset + PCI_EXP_LNKCAP);
+               tmp &= ~PCI_EXP_LNKCAP_SLS;
+               tmp |= PCI_EXP_LNKCAP_SLS_5_0GB;
+               dw_pcie_writel_dbi(pci, offset + PCI_EXP_LNKCAP, tmp);
 
                /*
                 * Start Directed Speed Change so the best possible
                dev_info(dev, "Link: Gen2 disabled\n");
        }
 
-       tmp = dw_pcie_readl_dbi(pci, PCIE_RC_LCSR);
-       dev_info(dev, "Link up, Gen%i\n", (tmp >> 16) & 0xf);
+       tmp = dw_pcie_readw_dbi(pci, offset + PCI_EXP_LNKSTA);
+       dev_info(dev, "Link up, Gen%i\n", tmp & PCI_EXP_LNKSTA_CLS);
        return 0;
 
 err_reset_phy:
                return ret;
 
        if (pci_msi_enabled()) {
-               val = dw_pcie_readw_dbi(pci, PCIE_RC_IMX6_MSI_CAP +
-                                       PCI_MSI_FLAGS);
+               u8 offset = dw_pcie_find_capability(pci, PCI_CAP_ID_MSI);
+               val = dw_pcie_readw_dbi(pci, offset + PCI_MSI_FLAGS);
                val |= PCI_MSI_FLAGS_ENABLE;
-               dw_pcie_writew_dbi(pci, PCIE_RC_IMX6_MSI_CAP + PCI_MSI_FLAGS,
-                                  val);
+               dw_pcie_writew_dbi(pci, offset + PCI_MSI_FLAGS, val);
        }
 
        return 0;