diff --git a/drivers/pci/pci.c b/drivers/pci/pci.c
index e3145f02027..bbc82f08d1c 100644
--- a/drivers/pci/pci.c
+++ b/drivers/pci/pci.c
@@ -741,8 +741,8 @@ static int pci_save_pcie_state(struct pci_dev *dev)
 	u16 *cap;
 	u16 flags;
 
-	pos = pci_find_capability(dev, PCI_CAP_ID_EXP);
-	if (pos <= 0)
+	pos = pci_pcie_cap(dev);
+	if (!pos)
 		return 0;
 
 	save_state = pci_find_saved_cap(dev, PCI_CAP_ID_EXP);
@@ -1535,7 +1535,7 @@ void pci_enable_ari(struct pci_dev *dev)
 	if (!bridge || !bridge->is_pcie)
 		return;
 
-	pos = pci_find_capability(bridge, PCI_CAP_ID_EXP);
+	pos = pci_pcie_cap(bridge);
 	if (!pos)
 		return;
 
@@ -2140,7 +2140,7 @@ static int pcie_flr(struct pci_dev *dev, int probe)
 	u32 cap;
 	u16 status;
 
-	pos = pci_find_capability(dev, PCI_CAP_ID_EXP);
+	pos = pci_pcie_cap(dev);
 	if (!pos)
 		return -ENOTTY;
 
@@ -2489,7 +2489,7 @@ int pcie_get_readrq(struct pci_dev *dev)
 	int ret, cap;
 	u16 ctl;
 
-	cap = pci_find_capability(dev, PCI_CAP_ID_EXP);
+	cap = pci_pcie_cap(dev);
 	if (!cap)
 		return -EINVAL;
 
@@ -2519,7 +2519,7 @@ int pcie_set_readrq(struct pci_dev *dev, int rq)
 
 	v = (ffs(rq) - 8) << 12;
 
-	cap = pci_find_capability(dev, PCI_CAP_ID_EXP);
+	cap = pci_pcie_cap(dev);
 	if (!cap)
 		goto out;
 
diff --git a/drivers/pci/probe.c b/drivers/pci/probe.c
index 54b9f150148..2fdffc02a30 100644
--- a/drivers/pci/probe.c
+++ b/drivers/pci/probe.c
@@ -703,7 +703,7 @@ static void set_pcie_hotplug_bridge(struct pci_dev *pdev)
 	u16 reg16;
 	u32 reg32;
 
-	pos = pci_find_capability(pdev, PCI_CAP_ID_EXP);
+	pos = pci_pcie_cap(pdev);
 	if (!pos)
 		return;
 	pci_read_config_word(pdev, pos + PCI_EXP_FLAGS, &reg16);
@@ -920,7 +920,7 @@ int pci_cfg_space_size(struct pci_dev *dev)
 	if (class == PCI_CLASS_BRIDGE_HOST)
 		return pci_cfg_space_size_ext(dev);
 
-	pos = pci_find_capability(dev, PCI_CAP_ID_EXP);
+	pos = pci_pcie_cap(dev);
 	if (!pos) {
 		pos = pci_find_capability(dev, PCI_CAP_ID_PCIX);
 		if (!pos)