4 files changed, 30 insertions, 11 deletions

diff --git a/drivers/core/device-remove.c b/drivers/core/device-remove.c
index a7f77b4a21..cc0043b990 100644
--- a/drivers/core/device-remove.c
+++ b/drivers/core/device-remove.c
@@ -46,9 +46,10 @@ static int device_chld_unbind(struct udevice *dev)
 /**
  * device_chld_remove() - Stop all device's children
  * @dev:	The device whose children are to be removed
+ * @pre_os_remove: Flag, if this functions is called in the pre-OS stage
  * @return 0 on success, -ve on error
  */
-static int device_chld_remove(struct udevice *dev)
+static int device_chld_remove(struct udevice *dev, uint flags)
 {
 	struct udevice *pos, *n;
 	int ret;
@@ -56,7 +57,7 @@ static int device_chld_remove(struct udevice *dev)
 	assert(dev);
 
 	list_for_each_entry_safe(pos, n, &dev->child_head, sibling_node) {
-		ret = device_remove(pos);
+		ret = device_remove(pos, flags);
 		if (ret)
 			return ret;
 	}
@@ -151,7 +152,7 @@ void device_free(struct udevice *dev)
 	devres_release_probe(dev);
 }
 
-int device_remove(struct udevice *dev)
+int device_remove(struct udevice *dev, uint flags)
 {
 	const struct driver *drv;
 	int ret;
@@ -169,11 +170,17 @@ int device_remove(struct udevice *dev)
 	if (ret)
 		return ret;
 
-	ret = device_chld_remove(dev);
+	ret = device_chld_remove(dev, flags);
 	if (ret)
 		goto err;
 
-	if (drv->remove) {
+	/*
+	 * Remove the device if called with the "normal" remove flag set,
+	 * or if the remove flag matches any of the drivers remove flags
+	 */
+	if (drv->remove &&
+	    ((flags & DM_REMOVE_NORMAL) ||
+	     (flags & (drv->flags & DM_FLAG_ACTIVE_DMA)))) {
 		ret = drv->remove(dev);
 		if (ret)
 			goto err_remove;
@@ -187,10 +194,13 @@ int device_remove(struct udevice *dev)
 		}
 	}
 
-	device_free(dev);
+	if ((flags & DM_REMOVE_NORMAL) ||
+	    (flags & (drv->flags & DM_FLAG_ACTIVE_DMA))) {
+		device_free(dev);
 
-	dev->seq = -1;
-	dev->flags &= ~DM_FLAG_ACTIVATED;
+		dev->seq = -1;
+		dev->flags &= ~DM_FLAG_ACTIVATED;
+	}
 
 	return ret;
 
diff --git a/drivers/core/device.c b/drivers/core/device.c
index 70fcfc23e0..e1b0ebffc5 100644
--- a/drivers/core/device.c
+++ b/drivers/core/device.c
@@ -378,7 +378,7 @@ int device_probe(struct udevice *dev)
 
 	return 0;
 fail_uclass:
-	if (device_remove(dev)) {
+	if (device_remove(dev, DM_REMOVE_NORMAL)) {
 		dm_warn("%s: Device '%s' failed to remove on error path\n",
 			__func__, dev->name);
 	}
diff --git a/drivers/core/root.c b/drivers/core/root.c
index 33cfde6a5c..42679d047c 100644
--- a/drivers/core/root.c
+++ b/drivers/core/root.c
@@ -178,12 +178,21 @@ int dm_init(void)
 
 int dm_uninit(void)
 {
-	device_remove(dm_root());
+	device_remove(dm_root(), DM_REMOVE_NORMAL);
 	device_unbind(dm_root());
 
 	return 0;
 }
 
+#if CONFIG_IS_ENABLED(DM_DEVICE_REMOVE)
+int dm_remove_devices_flags(uint flags)
+{
+	device_remove(dm_root(), flags);
+
+	return 0;
+}
+#endif
+
 int dm_scan_platdata(bool pre_reloc_only)
 {
 	int ret;
diff --git a/drivers/core/uclass.c b/drivers/core/uclass.c
index 7de370644d..d94d43a98d 100644
--- a/drivers/core/uclass.c
+++ b/drivers/core/uclass.c
@@ -116,7 +116,7 @@ int uclass_destroy(struct uclass *uc)
 	while (!list_empty(&uc->dev_head)) {
 		dev = list_first_entry(&uc->dev_head, struct udevice,
 				       uclass_node);
-		ret = device_remove(dev);
+		ret = device_remove(dev, DM_REMOVE_NORMAL);
 		if (ret)
 			return ret;
 		ret = device_unbind(dev);