devres_release_probe(dev);
 }
 
+static bool flags_remove(uint flags, uint drv_flags)
+{
+       if ((flags & DM_REMOVE_NORMAL) ||
+           (flags & (drv_flags & (DM_FLAG_ACTIVE_DMA | DM_FLAG_OS_PREPARE))))
+               return true;
+
+       return false;
+}
+
 int device_remove(struct udevice *dev, uint flags)
 {
        const struct driver *drv;
         * 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)))) {
+       if (drv->remove && flags_remove(flags, drv->flags)) {
                ret = drv->remove(dev);
                if (ret)
                        goto err_remove;
                }
        }
 
-       if ((flags & DM_REMOVE_NORMAL) ||
-           (flags & (drv->flags & DM_FLAG_ACTIVE_DMA))) {
+       if (flags_remove(flags, drv->flags)) {
                device_free(dev);
 
                dev->seq = -1;
 
  */
 #define DM_FLAG_ACTIVE_DMA             (1 << 9)
 
+/*
+ * Call driver remove function to do some final configuration, before
+ * U-Boot exits and the OS is started
+ */
+#define DM_FLAG_OS_PREPARE             (1 << 10)
+
 /*
  * One or multiple of these flags are passed to device_remove() so that
  * a selective device removal as specified by the remove-stage and the
        /* Remove devices with active DMA */
        DM_REMOVE_ACTIVE_DMA = DM_FLAG_ACTIVE_DMA,
 
+       /* Remove devices which need some final OS preparation steps */
+       DM_REMOVE_OS_PREPARE = DM_FLAG_OS_PREPARE,
+
        /* Add more use cases here */
 
        /* Remove devices with any active flag */
-       DM_REMOVE_ACTIVE_ALL = DM_REMOVE_ACTIVE_DMA,
+       DM_REMOVE_ACTIVE_ALL = DM_REMOVE_ACTIVE_DMA | DM_REMOVE_OS_PREPARE,
 };
 
 /**