spl: mtd: Remove MTD device after loading images

Releasing the flash into proper state, after the loading completes,
is important for the next stage bootloader/kernel to be able to use
the MTD device. This would enable to reset the device for fresh
use by next boot stage.

Signed-off-by: Apurva Nandan <a-nandan@ti.com>

common/spl/spl_mtd.c

@@ -46,13 +46,17 @@ static int spl_mtd_load_image(struct spl_image_info *spl_image,
 	switch (bootdev->boot_device) {
 	case BOOT_DEVICE_SPINAND:
 		mtd = get_mtd_device_nm("spi-nand0");
-		if (IS_ERR_OR_NULL(mtd))
+		if (IS_ERR_OR_NULL(mtd)) {
 			printf("MTD device %s not found, ret %ld\n", "spi-nand",
 			       PTR_ERR(mtd));
+			err = PTR_ERR(mtd);
+			goto remove_mtd_device;
+		}
 		break;
 	default:
 		puts(SPL_TPL_PROMPT "Unsupported MTD Boot Device!\n");
-		return -EINVAL;
+		err = -EINVAL;
+		goto remove_mtd_device;
 	}
 
 	header = spl_get_load_buffer(0, sizeof(*header));
@@ -60,7 +64,7 @@ static int spl_mtd_load_image(struct spl_image_info *spl_image,
 	err = mtd_read(mtd, spl_mtd_get_uboot_offs(), sizeof(*header),
 		       &ret_len, (void *)header);
 	if (err)
-		return err;
+		goto remove_mtd_device;
 
 	if (IS_ENABLED(CONFIG_SPL_LOAD_FIT) &&
 	    image_get_magic(header) == FDT_MAGIC) {
@@ -70,25 +74,27 @@ static int spl_mtd_load_image(struct spl_image_info *spl_image,
 		load.filename = NULL;
 		load.bl_len = 1;
 		load.read = spl_mtd_fit_read;
-		return spl_load_simple_fit(spl_image, &load,
-					   spl_mtd_get_uboot_offs(), header);
+		err = spl_load_simple_fit(spl_image, &load,
+					  spl_mtd_get_uboot_offs(), header);
 	} else if (IS_ENABLED(CONFIG_SPL_LOAD_IMX_CONTAINER)) {
 		load.dev = mtd;
 		load.priv = NULL;
 		load.filename = NULL;
 		load.bl_len = 1;
 		load.read = spl_mtd_fit_read;
-		return spl_load_imx_container(spl_image, &load,
-					      spl_mtd_get_uboot_offs());
+		err = spl_load_imx_container(spl_image, &load,
+					     spl_mtd_get_uboot_offs());
 	} else {
 		err = spl_parse_image_header(spl_image, bootdev, header);
 		if (err)
-			return err;
-		return mtd_read(mtd, spl_mtd_get_uboot_offs(), spl_image->size,
-				&ret_len, (void *)(ulong)spl_image->load_addr);
+			goto remove_mtd_device;
+		err = mtd_read(mtd, spl_mtd_get_uboot_offs(), spl_image->size,
+			       &ret_len, (void *)(ulong)spl_image->load_addr);
 	}
 
-	return -EINVAL;
+remove_mtd_device:
+	mtd_remove(mtd);
+	return err;
 }
 
 SPL_LOAD_IMAGE_METHOD("SPINAND", 0, BOOT_DEVICE_SPINAND, spl_mtd_load_image);

drivers/mtd/mtd-uclass.c

@@ -11,6 +11,18 @@
 #include <errno.h>
 #include <mtd.h>
 
+/**
+ * mtd_remove - Remove the device @dev
+ *
+ * @dev: U-Boot device to probe
+ *
+ * @return 0 on success, an error otherwise.
+ */
+int mtd_remove(struct mtd_info *mtd)
+{
+	return device_remove(mtd->dev, DM_REMOVE_NORMAL);
+}
+
 /*
  * Implement a MTD uclass which should include most flash drivers.
  * The uclass private is pointed to mtd_info.

drivers/mtd/nand/spi/core.c

@@ -1300,26 +1300,25 @@ err_spinand_cleanup:
 	return ret;
 }
 
-#ifndef __UBOOT__
 static int spinand_remove(struct udevice *slave)
 {
-	struct spinand_device *spinand;
+	struct spinand_device *spinand = dev_get_priv(slave);
 	struct mtd_info *mtd;
-	int ret;
 
-	spinand = spi_mem_get_drvdata(slave);
 	mtd = spinand_to_mtd(spinand);
 	free(mtd->name);
 
+#ifndef __UBOOT__
 	ret = mtd_device_unregister(mtd);
 	if (ret)
 		return ret;
-
+#endif
 	spinand_cleanup(spinand);
 
 	return 0;
 }
 
+#ifndef __UBOOT__
 static const struct spi_device_id spinand_ids[] = {
 	{ .name = "spi-nand" },
 	{ /* sentinel */ },
@@ -1361,4 +1360,6 @@ U_BOOT_DRIVER(spinand) = {
 	.of_match = spinand_ids,
 	.priv_auto	= sizeof(struct spinand_device),
 	.probe = spinand_probe,
+	.remove = spinand_remove,
+	.flags = DM_FLAG_OS_PREPARE,
 };

include/mtd.h

@@ -11,6 +11,7 @@
 #include <linux/mtd/mtd.h>
 
 int mtd_probe_devices(void);
+int mtd_remove(struct mtd_info *mtd);
 
 void board_mtdparts_default(const char **mtdids, const char **mtdparts);