diff --git a/drivers/pci/hotplug/pciehp.h b/drivers/pci/hotplug/pciehp.h
index a4b294fddc93941dd813d4feb1ae2e46abc00f9c..1fe9cd746904254742deae7fe31841a9d668ab84 100644
--- a/drivers/pci/hotplug/pciehp.h
+++ b/drivers/pci/hotplug/pciehp.h
@@ -213,4 +213,9 @@ static inline const char *slot_name(struct slot *slot)
 	return hotplug_slot_name(slot->hotplug_slot);
 }
 
+static inline struct pci_dev *ctrl_dev(struct controller *ctrl)
+{
+	return ctrl->pcie->port;
+}
+
 #endif				/* _PCIEHP_H */
diff --git a/drivers/pci/hotplug/pciehp_ctrl.c b/drivers/pci/hotplug/pciehp_ctrl.c
index 5c26328da66ccc1d1d2908940ad69b8eafc81c83..2d549c97ac420bdc17d42d59929e85ee2d94ba3d 100644
--- a/drivers/pci/hotplug/pciehp_ctrl.c
+++ b/drivers/pci/hotplug/pciehp_ctrl.c
@@ -146,6 +146,7 @@ void pciehp_queue_pushbutton_work(struct work_struct *work)
 	struct slot *p_slot = container_of(work, struct slot, work.work);
 	struct controller *ctrl = p_slot->ctrl;
 	int events = p_slot->work.data;
+	struct pci_dev *rpdev = ctrl_dev(ctrl)->rpdev;
 
 	mutex_lock(&p_slot->lock);
 	switch (p_slot->state) {
@@ -160,8 +161,11 @@ void pciehp_queue_pushbutton_work(struct work_struct *work)
 			atomic_or(events, &ctrl->pending_events);
 			if (!pciehp_poll_mode)
 				irq_wake_thread(ctrl->pcie->irq, ctrl);
-		} else
-			slot_being_removed_rescanned = 0;
+		} else {
+			if (rpdev)
+				clear_bit(0,
+					  &rpdev->slot_being_removed_rescanned);
+		}
 		break;
 	}
 	mutex_unlock(&p_slot->lock);
@@ -170,6 +174,7 @@ void pciehp_queue_pushbutton_work(struct work_struct *work)
 void pciehp_handle_button_press(struct slot *p_slot)
 {
 	struct controller *ctrl = p_slot->ctrl;
+	struct pci_dev *rpdev = ctrl_dev(ctrl)->rpdev;
 
 	mutex_lock(&p_slot->lock);
 	switch (p_slot->state) {
@@ -209,12 +214,14 @@ void pciehp_handle_button_press(struct slot *p_slot)
 		pciehp_set_attention_status(p_slot, 0);
 		ctrl_info(ctrl, "Slot(%s): Action canceled due to button press\n",
 			  slot_name(p_slot));
-		slot_being_removed_rescanned = 0;
+		if (rpdev)
+			clear_bit(0, &rpdev->slot_being_removed_rescanned);
 		break;
 	default:
 		ctrl_err(ctrl, "Slot(%s): Ignoring invalid state %#x\n",
 			 slot_name(p_slot), p_slot->state);
-		slot_being_removed_rescanned = 0;
+		if (rpdev)
+			clear_bit(0, &rpdev->slot_being_removed_rescanned);
 		break;
 	}
 	mutex_unlock(&p_slot->lock);
@@ -223,6 +230,7 @@ void pciehp_handle_button_press(struct slot *p_slot)
 void pciehp_handle_disable_request(struct slot *slot)
 {
 	struct controller *ctrl = slot->ctrl;
+	struct pci_dev *rpdev = ctrl_dev(ctrl)->rpdev;
 
 	mutex_lock(&slot->lock);
 	switch (slot->state) {
@@ -235,7 +243,8 @@ void pciehp_handle_disable_request(struct slot *slot)
 	mutex_unlock(&slot->lock);
 
 	ctrl->request_result = pciehp_disable_slot(slot, SAFE_REMOVAL);
-	slot_being_removed_rescanned = 0;
+	if (rpdev)
+		clear_bit(0, &rpdev->slot_being_removed_rescanned);
 }
 
 void pciehp_handle_presence_or_link_change(struct slot *slot, u32 events)
@@ -243,6 +252,7 @@ void pciehp_handle_presence_or_link_change(struct slot *slot, u32 events)
 	struct controller *ctrl = slot->ctrl;
 	bool present, link_active;
 	bool removal = SAFE_REMOVAL;
+	struct pci_dev *rpdev = ctrl_dev(ctrl)->rpdev;
 
 	/*
 	 * If the slot is on and presence or link has changed, turn it off.
@@ -285,7 +295,8 @@ void pciehp_handle_presence_or_link_change(struct slot *slot, u32 events)
 	link_active = pciehp_check_link_active(ctrl);
 	if (!present && !link_active) {
 		mutex_unlock(&slot->lock);
-		slot_being_removed_rescanned = 0;
+		if (rpdev)
+			clear_bit(0, &rpdev->slot_being_removed_rescanned);
 		return;
 	}
 
@@ -308,7 +319,8 @@ void pciehp_handle_presence_or_link_change(struct slot *slot, u32 events)
 		mutex_unlock(&slot->lock);
 		break;
 	}
-	slot_being_removed_rescanned = 0;
+	if (rpdev)
+		clear_bit(0, &rpdev->slot_being_removed_rescanned);
 }
 
 static int __pciehp_enable_slot(struct slot *p_slot)
@@ -435,8 +447,10 @@ int pciehp_sysfs_disable_slot(struct slot *p_slot)
 {
 	struct controller *ctrl = p_slot->ctrl;
 	struct pci_dev *pdev = ctrl->pcie->port;
+	struct pci_dev *rpdev = pdev->rpdev;
 
-	if (test_and_set_bit(0, &slot_being_removed_rescanned)) {
+	if (rpdev && test_and_set_bit(0,
+				&rpdev->slot_being_removed_rescanned)) {
 		ctrl_info(ctrl, "Slot(%s): Slot is being removed or rescanned, please try later!\n",
 			  slot_name(p_slot));
 		return -EINVAL;
@@ -452,7 +466,8 @@ int pciehp_sysfs_disable_slot(struct slot *p_slot)
 		pciehp_handle_disable_request(p_slot);
 		up_read(&ctrl->reset_lock);
 		pci_config_pm_runtime_put(pdev);
-		slot_being_removed_rescanned = 0;
+		if (rpdev)
+			clear_bit(0, &rpdev->slot_being_removed_rescanned);
 		return ctrl->request_result;
 	case POWEROFF_STATE:
 		ctrl_info(ctrl, "Slot(%s): Already in powering off state\n",
@@ -471,7 +486,8 @@ int pciehp_sysfs_disable_slot(struct slot *p_slot)
 	}
 	mutex_unlock(&p_slot->lock);
 
-	slot_being_removed_rescanned = 0;
+	if (rpdev)
+		clear_bit(0, &rpdev->slot_being_removed_rescanned);
 
 	return -ENODEV;
 }
diff --git a/drivers/pci/hotplug/pciehp_hpc.c b/drivers/pci/hotplug/pciehp_hpc.c
index 43bb78122795a9de7df5488886746b1a5f660fe8..7ec78081cc28ecd65dc9af8c1c82fce25118e041 100644
--- a/drivers/pci/hotplug/pciehp_hpc.c
+++ b/drivers/pci/hotplug/pciehp_hpc.c
@@ -27,11 +27,6 @@
 #include "../pci.h"
 #include "pciehp.h"
 
-static inline struct pci_dev *ctrl_dev(struct controller *ctrl)
-{
-	return ctrl->pcie->port;
-}
-
 static irqreturn_t pciehp_isr(int irq, void *dev_id);
 static irqreturn_t pciehp_ist(int irq, void *dev_id);
 static int pciehp_poll(void *data);
@@ -632,6 +627,7 @@ static irqreturn_t pciehp_ist(int irq, void *dev_id)
 {
 	struct controller *ctrl = (struct controller *)dev_id;
 	struct pci_dev *pdev = ctrl_dev(ctrl);
+	struct pci_dev *rpdev = pdev->rpdev;
 	struct slot *slot = ctrl->slot;
 	irqreturn_t ret;
 	u32 events;
@@ -659,7 +655,8 @@ static irqreturn_t pciehp_ist(int irq, void *dev_id)
 	if (events & PCI_EXP_SLTSTA_ABP) {
 		ctrl_info(ctrl, "Slot(%s): Attention button pressed\n",
 			  slot_name(slot));
-		if (!test_and_set_bit(0, &slot_being_removed_rescanned))
+		if (!rpdev || (rpdev && !test_and_set_bit(0,
+					&rpdev->slot_being_removed_rescanned)))
 			pciehp_handle_button_press(slot);
 		else {
 			if (slot->state == BLINKINGOFF_STATE ||
@@ -686,7 +683,8 @@ static irqreturn_t pciehp_ist(int irq, void *dev_id)
 	 */
 	down_read(&ctrl->reset_lock);
 	if (events & DISABLE_SLOT) {
-		if (!test_and_set_bit(0, &slot_being_removed_rescanned))
+		if (!rpdev || (rpdev && !test_and_set_bit(0,
+					&rpdev->slot_being_removed_rescanned)))
 			pciehp_handle_disable_request(slot);
 		else {
 			if (slot->state == BLINKINGOFF_STATE ||
@@ -712,7 +710,8 @@ static irqreturn_t pciehp_ist(int irq, void *dev_id)
 			}
 		}
 	} else if (events & (PCI_EXP_SLTSTA_PDC | PCI_EXP_SLTSTA_DLLSC)) {
-		if (!test_and_set_bit(0, &slot_being_removed_rescanned))
+		if (!rpdev || (rpdev && !test_and_set_bit(0,
+					&rpdev->slot_being_removed_rescanned)))
 			pciehp_handle_presence_or_link_change(slot, events);
 		else {
 			if (slot->state == BLINKINGOFF_STATE ||
diff --git a/drivers/pci/pci-sysfs.c b/drivers/pci/pci-sysfs.c
index 09e2066a93d405ca150ce749c4caab950f0f8a85..391a811ba3445fafc3209a557f3f3326353ba655 100644
--- a/drivers/pci/pci-sysfs.c
+++ b/drivers/pci/pci-sysfs.c
@@ -470,19 +470,33 @@ static ssize_t remove_store(struct device *dev, struct device_attribute *attr,
 			    const char *buf, size_t count)
 {
 	unsigned long val;
+	struct pci_dev *rpdev = to_pci_dev(dev)->rpdev;
 
 	if (kstrtoul(buf, 0, &val) < 0)
 		return -EINVAL;
 
-	if (test_and_set_bit(0, &slot_being_removed_rescanned)) {
+	if (rpdev && test_and_set_bit(0,
+				&rpdev->slot_being_removed_rescanned)) {
 		pr_info("Slot is being removed or rescanned, please try later!\n");
 		return -EINVAL;
 	}
 
+	/*
+	 * if 'dev' is root port itself, 'pci_stop_and_remove_bus_device()' may
+	 * free the 'rpdev', but we need to clear
+	 * 'rpdev->slot_being_removed_rescanned' in the end. So get 'rpdev' to
+	 * avoid possible 'use-after-free'.
+	 */
+	if (rpdev)
+		pci_dev_get(rpdev);
+
 	if (val && device_remove_file_self(dev, attr))
 		pci_stop_and_remove_bus_device_locked(to_pci_dev(dev));
 
-	slot_being_removed_rescanned = 0;
+	if (rpdev) {
+		clear_bit(0, &rpdev->slot_being_removed_rescanned);
+		pci_dev_put(rpdev);
+	}
 
 	return count;
 }
diff --git a/drivers/pci/probe.c b/drivers/pci/probe.c
index f89af91f667525ebd8c9eef8c38c89b2d7d8caac..1320b48a50ac52cc02a52e1c378c26a648bec851 100644
--- a/drivers/pci/probe.c
+++ b/drivers/pci/probe.c
@@ -2928,6 +2928,11 @@ void pci_device_add(struct pci_dev *dev, struct pci_bus *bus)
 	/* Set up MSI IRQ domain */
 	pci_set_msi_domain(dev);
 
+	if (pci_pcie_type(dev) == PCI_EXP_TYPE_ROOT_PORT)
+		dev->rpdev = dev;
+	else
+		dev->rpdev = pcie_find_root_port(dev);
+
 	/* Notifier could use PCI capabilities */
 	dev->match_driver = false;
 	ret = device_add(&dev->dev);
diff --git a/drivers/pci/remove.c b/drivers/pci/remove.c
index 6d15ad0488d21815dc8dec9b4ab5c6aebeac2442..e9c6b120cf451331dc294f50a3ac1315cd37c2c3 100644
--- a/drivers/pci/remove.c
+++ b/drivers/pci/remove.c
@@ -3,11 +3,6 @@
 #include <linux/module.h>
 #include "pci.h"
 
-/*
- * When a slot is being removed/rescanned, this flag is set.
- */
-unsigned long slot_being_removed_rescanned;
-
 static void pci_free_resources(struct pci_dev *dev)
 {
 	int i;
diff --git a/include/linux/pci.h b/include/linux/pci.h
index afc4506974b6132336c694971aceb31e728548e9..7a63613c83f5b77e54a5e7f11946be596883c0ea 100644
--- a/include/linux/pci.h
+++ b/include/linux/pci.h
@@ -450,6 +450,12 @@ struct pci_dev {
 	char		*driver_override; /* Driver name to force a match */
 
 	unsigned long	priv_flags;	/* Private flags for the PCI driver */
+	/*
+	 * This flag is only set on root ports. When a slot below a root port
+	 * is being removed or rescanned, this flag is set.
+	 */
+	unsigned long	slot_being_removed_rescanned;
+	struct pci_dev	*rpdev; /* root port pci_dev */
 
 	KABI_RESERVE(1)
 	KABI_RESERVE(2)
@@ -916,8 +922,6 @@ extern struct bus_type pci_bus_type;
  * code, or PCI core code. */
 extern struct list_head pci_root_buses;	/* List of all known PCI buses */
 
-extern unsigned long slot_being_removed_rescanned;
-
 /* Some device drivers need know if PCI is initiated */
 int no_pci_devices(void);