Merge branch 'master' of git://git.denx.de/u-boot-spi

This is the PR for SPI-NAND changes along with few spi changes.

[trini: Re-sync changes for ls1012afrwy_qspi*_defconfig]
Signed-off-by: Tom Rini <trini@konsulko.com>

cmd/Kconfig

@@ -866,6 +866,12 @@ config CMD_MMC_SWRITE
 	  Enable support for the "mmc swrite" command to write Android sparse
 	  images to eMMC.
 
+config CMD_MTD
+	bool "mtd"
+	select MTD_PARTITIONS
+	help
+	  MTD commands support.
+
 config CMD_NAND
 	bool "nand"
 	default y if NAND_SUNXI
@@ -1714,18 +1720,22 @@ config CMD_MTDPARTS
 	bool "MTD partition support"
 	select MTD_DEVICE if (CMD_NAND || NAND)
 	help
-	  MTD partition support
+	  MTD partitioning tool support.
+	  It is strongly encouraged to avoid using this command
+	  anymore along with 'sf', 'nand', 'onenand'. One can still
+	  declare the partitions in the mtdparts environment variable
+	  but better use the MTD stack and the 'mtd' command instead.
 
 config MTDIDS_DEFAULT
 	string "Default MTD IDs"
-	depends on CMD_MTDPARTS || CMD_NAND || CMD_FLASH
+	depends on CMD_MTD || CMD_MTDPARTS || CMD_NAND || CMD_FLASH
 	help
 	  Defines a default MTD IDs list for use with MTD partitions in the
 	  Linux MTD command line partitions format.
 
 config MTDPARTS_DEFAULT
 	string "Default MTD partition scheme"
-	depends on CMD_MTDPARTS || CMD_NAND || CMD_FLASH
+	depends on CMD_MTD || CMD_MTDPARTS || CMD_NAND || CMD_FLASH
 	help
 	  Defines a default MTD partitioning scheme in the Linux MTD command
 	  line partitions format
@@ -1855,6 +1865,8 @@ config CMD_UBI
 	  capabilities. Please, consult the MTD web site for more details
 	  (www.linux-mtd.infradead.org). Activate this option if you want
 	  to use U-Boot UBI commands.
+	  It is also strongly encouraged to also enable CONFIG_MTD to get full
+	  partition support.
 
 config CMD_UBIFS
 	tristate "Enable UBIFS - Unsorted block images filesystem commands"

cmd/Makefile

@@ -93,6 +93,7 @@ obj-$(CONFIG_CMD_MISC) += misc.o
 obj-$(CONFIG_CMD_MMC) += mmc.o
 obj-$(CONFIG_CMD_MMC_SPI) += mmc_spi.o
 obj-$(CONFIG_MP) += mp.o
+obj-$(CONFIG_CMD_MTD) += mtd.o
 obj-$(CONFIG_CMD_MTDPARTS) += mtdparts.o
 obj-$(CONFIG_CMD_NAND) += nand.o
 obj-$(CONFIG_CMD_NET) += net.o

cmd/mtd.c

@@ -0,0 +1,473 @@
+// SPDX-License-Identifier:  GPL-2.0+
+/*
+ * mtd.c
+ *
+ * Generic command to handle basic operations on any memory device.
+ *
+ * Copyright: Bootlin, 2018
+ * Author: Miquèl Raynal <miquel.raynal@bootlin.com>
+ */
+
+#include <command.h>
+#include <common.h>
+#include <console.h>
+#include <malloc.h>
+#include <mapmem.h>
+#include <mtd.h>
+
+static uint mtd_len_to_pages(struct mtd_info *mtd, u64 len)
+{
+	do_div(len, mtd->writesize);
+
+	return len;
+}
+
+static bool mtd_is_aligned_with_min_io_size(struct mtd_info *mtd, u64 size)
+{
+	return !do_div(size, mtd->writesize);
+}
+
+static bool mtd_is_aligned_with_block_size(struct mtd_info *mtd, u64 size)
+{
+	return !do_div(size, mtd->erasesize);
+}
+
+static void mtd_dump_buf(const u8 *buf, uint len, uint offset)
+{
+	int i, j;
+
+	for (i = 0; i < len; ) {
+		printf("0x%08x:\t", offset + i);
+		for (j = 0; j < 8; j++)
+			printf("%02x ", buf[i + j]);
+		printf(" ");
+		i += 8;
+		for (j = 0; j < 8; j++)
+			printf("%02x ", buf[i + j]);
+		printf("\n");
+		i += 8;
+	}
+}
+
+static void mtd_dump_device_buf(struct mtd_info *mtd, u64 start_off,
+				const u8 *buf, u64 len, bool woob)
+{
+	bool has_pages = mtd->type == MTD_NANDFLASH ||
+		mtd->type == MTD_MLCNANDFLASH;
+	int npages = mtd_len_to_pages(mtd, len);
+	uint page;
+
+	if (has_pages) {
+		for (page = 0; page < npages; page++) {
+			u64 data_off = page * mtd->writesize;
+
+			printf("\nDump %d data bytes from 0x%08llx:\n",
+			       mtd->writesize, start_off + data_off);
+			mtd_dump_buf(&buf[data_off],
+				     mtd->writesize, start_off + data_off);
+
+			if (woob) {
+				u64 oob_off = page * mtd->oobsize;
+
+				printf("Dump %d OOB bytes from page at 0x%08llx:\n",
+				       mtd->oobsize, start_off + data_off);
+				mtd_dump_buf(&buf[len + oob_off],
+					     mtd->oobsize, 0);
+			}
+		}
+	} else {
+		printf("\nDump %lld data bytes from 0x%llx:\n",
+		       len, start_off);
+		mtd_dump_buf(buf, len, start_off);
+	}
+}
+
+static void mtd_show_parts(struct mtd_info *mtd, int level)
+{
+	struct mtd_info *part;
+	int i;
+
+	list_for_each_entry(part, &mtd->partitions, node) {
+		for (i = 0; i < level; i++)
+			printf("\t");
+		printf("  - 0x%012llx-0x%012llx : \"%s\"\n",
+		       part->offset, part->offset + part->size, part->name);
+
+		mtd_show_parts(part, level + 1);
+	}
+}
+
+static void mtd_show_device(struct mtd_info *mtd)
+{
+	/* Device */
+	printf("* %s\n", mtd->name);
+#if defined(CONFIG_DM)
+	if (mtd->dev) {
+		printf("  - device: %s\n", mtd->dev->name);
+		printf("  - parent: %s\n", mtd->dev->parent->name);
+		printf("  - driver: %s\n", mtd->dev->driver->name);
+	}
+#endif
+
+	/* MTD device information */
+	printf("  - type: ");
+	switch (mtd->type) {
+	case MTD_RAM:
+		printf("RAM\n");
+		break;
+	case MTD_ROM:
+		printf("ROM\n");
+		break;
+	case MTD_NORFLASH:
+		printf("NOR flash\n");
+		break;
+	case MTD_NANDFLASH:
+		printf("NAND flash\n");
+		break;
+	case MTD_DATAFLASH:
+		printf("Data flash\n");
+		break;
+	case MTD_UBIVOLUME:
+		printf("UBI volume\n");
+		break;
+	case MTD_MLCNANDFLASH:
+		printf("MLC NAND flash\n");
+		break;
+	case MTD_ABSENT:
+	default:
+		printf("Unknown\n");
+		break;
+	}
+
+	printf("  - block size: 0x%x bytes\n", mtd->erasesize);
+	printf("  - min I/O: 0x%x bytes\n", mtd->writesize);
+
+	if (mtd->oobsize) {
+		printf("  - OOB size: %u bytes\n", mtd->oobsize);
+		printf("  - OOB available: %u bytes\n", mtd->oobavail);
+	}
+
+	if (mtd->ecc_strength) {
+		printf("  - ECC strength: %u bits\n", mtd->ecc_strength);
+		printf("  - ECC step size: %u bytes\n", mtd->ecc_step_size);
+		printf("  - bitflip threshold: %u bits\n",
+		       mtd->bitflip_threshold);
+	}
+
+	printf("  - 0x%012llx-0x%012llx : \"%s\"\n",
+	       mtd->offset, mtd->offset + mtd->size, mtd->name);
+
+	/* MTD partitions, if any */
+	mtd_show_parts(mtd, 1);
+}
+
+/* Logic taken from fs/ubifs/recovery.c:is_empty() */
+static bool mtd_oob_write_is_empty(struct mtd_oob_ops *op)
+{
+	int i;
+
+	for (i = 0; i < op->len; i++)
+		if (op->datbuf[i] != 0xff)
+			return false;
+
+	for (i = 0; i < op->ooblen; i++)
+		if (op->oobbuf[i] != 0xff)
+			return false;
+
+	return true;
+}
+
+static int do_mtd_list(void)
+{
+	struct mtd_info *mtd;
+	int dev_nb = 0;
+
+	/* Ensure all devices (and their partitions) are probed */
+	mtd_probe_devices();
+
+	printf("List of MTD devices:\n");
+	mtd_for_each_device(mtd) {
+		if (!mtd_is_partition(mtd))
+			mtd_show_device(mtd);
+
+		dev_nb++;
+	}
+
+	if (!dev_nb) {
+		printf("No MTD device found\n");
+		return CMD_RET_FAILURE;
+	}
+
+	return CMD_RET_SUCCESS;
+}
+
+static int mtd_special_write_oob(struct mtd_info *mtd, u64 off,
+				 struct mtd_oob_ops *io_op,
+				 bool write_empty_pages, bool woob)
+{
+	int ret = 0;
+
+	/*
+	 * By default, do not write an empty page.
+	 * Skip it by simulating a successful write.
+	 */
+	if (!write_empty_pages && mtd_oob_write_is_empty(io_op)) {
+		io_op->retlen = mtd->writesize;
+		io_op->oobretlen = woob ? mtd->oobsize : 0;
+	} else {
+		ret = mtd_write_oob(mtd, off, io_op);
+	}
+
+	return ret;
+}
+
+static int do_mtd(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
+{
+	struct mtd_info *mtd;
+	const char *cmd;
+	char *mtd_name;
+
+	/* All MTD commands need at least two arguments */
+	if (argc < 2)
+		return CMD_RET_USAGE;
+
+	/* Parse the command name and its optional suffixes */
+	cmd = argv[1];
+
+	/* List the MTD devices if that is what the user wants */
+	if (strcmp(cmd, "list") == 0)
+		return do_mtd_list();
+
+	/*
+	 * The remaining commands require also at least a device ID.
+	 * Check the selected device is valid. Ensure it is probed.
+	 */
+	if (argc < 3)
+		return CMD_RET_USAGE;
+
+	mtd_name = argv[2];
+	mtd_probe_devices();
+	mtd = get_mtd_device_nm(mtd_name);
+	if (IS_ERR_OR_NULL(mtd)) {
+		printf("MTD device %s not found, ret %ld\n",
+		       mtd_name, PTR_ERR(mtd));
+		return CMD_RET_FAILURE;
+	}
+	put_mtd_device(mtd);
+
+	argc -= 3;
+	argv += 3;
+
+	/* Do the parsing */
+	if (!strncmp(cmd, "read", 4) || !strncmp(cmd, "dump", 4) ||
+	    !strncmp(cmd, "write", 5)) {
+		bool has_pages = mtd->type == MTD_NANDFLASH ||
+				 mtd->type == MTD_MLCNANDFLASH;
+		bool dump, read, raw, woob, write_empty_pages;
+		struct mtd_oob_ops io_op = {};
+		uint user_addr = 0, npages;
+		u64 start_off, off, len, remaining, default_len;
+		u32 oob_len;
+		u8 *buf;
+		int ret;
+
+		dump = !strncmp(cmd, "dump", 4);
+		read = dump || !strncmp(cmd, "read", 4);
+		raw = strstr(cmd, ".raw");
+		woob = strstr(cmd, ".oob");
+		write_empty_pages = !has_pages || strstr(cmd, ".dontskipff");
+
+		if (!dump) {
+			if (!argc)
+				return CMD_RET_USAGE;
+
+			user_addr = simple_strtoul(argv[0], NULL, 16);
+			argc--;
+			argv++;
+		}
+
+		start_off = argc > 0 ? simple_strtoul(argv[0], NULL, 16) : 0;
+		if (!mtd_is_aligned_with_min_io_size(mtd, start_off)) {
+			printf("Offset not aligned with a page (0x%x)\n",
+			       mtd->writesize);
+			return CMD_RET_FAILURE;
+		}
+
+		default_len = dump ? mtd->writesize : mtd->size;
+		len = argc > 1 ? simple_strtoul(argv[1], NULL, 16) :
+				 default_len;
+		if (!mtd_is_aligned_with_min_io_size(mtd, len)) {
+			len = round_up(len, mtd->writesize);
+			printf("Size not on a page boundary (0x%x), rounding to 0x%llx\n",
+			       mtd->writesize, len);
+		}
+
+		remaining = len;
+		npages = mtd_len_to_pages(mtd, len);
+		oob_len = woob ? npages * mtd->oobsize : 0;
+
+		if (dump)
+			buf = kmalloc(len + oob_len, GFP_KERNEL);
+		else
+			buf = map_sysmem(user_addr, 0);
+
+		if (!buf) {
+			printf("Could not map/allocate the user buffer\n");
+			return CMD_RET_FAILURE;
+		}
+
+		if (has_pages)
+			printf("%s %lld byte(s) (%d page(s)) at offset 0x%08llx%s%s%s\n",
+			       read ? "Reading" : "Writing", len, npages, start_off,
+			       raw ? " [raw]" : "", woob ? " [oob]" : "",
+			       !read && write_empty_pages ? " [dontskipff]" : "");
+		else
+			printf("%s %lld byte(s) at offset 0x%08llx\n",
+			       read ? "Reading" : "Writing", len, start_off);
+
+		io_op.mode = raw ? MTD_OPS_RAW : MTD_OPS_AUTO_OOB;
+		io_op.len = has_pages ? mtd->writesize : len;
+		io_op.ooblen = woob ? mtd->oobsize : 0;
+		io_op.datbuf = buf;
+		io_op.oobbuf = woob ? &buf[len] : NULL;
+
+		/* Search for the first good block after the given offset */
+		off = start_off;
+		while (mtd_block_isbad(mtd, off))
+			off += mtd->erasesize;
+
+		/* Loop over the pages to do the actual read/write */
+		while (remaining) {
+			/* Skip the block if it is bad */
+			if (mtd_is_aligned_with_block_size(mtd, off) &&
+			    mtd_block_isbad(mtd, off)) {
+				off += mtd->erasesize;
+				continue;
+			}
+
+			if (read)
+				ret = mtd_read_oob(mtd, off, &io_op);
+			else
+				ret = mtd_special_write_oob(mtd, off, &io_op,
+							    write_empty_pages,
+							    woob);
+
+			if (ret) {
+				printf("Failure while %s at offset 0x%llx\n",
+				       read ? "reading" : "writing", off);
+				return CMD_RET_FAILURE;
+			}
+
+			off += io_op.retlen;
+			remaining -= io_op.retlen;
+			io_op.datbuf += io_op.retlen;
+			io_op.oobbuf += io_op.oobretlen;
+		}
+
+		if (!ret && dump)
+			mtd_dump_device_buf(mtd, start_off, buf, len, woob);
+
+		if (dump)
+			kfree(buf);
+		else
+			unmap_sysmem(buf);
+
+		if (ret) {
+			printf("%s on %s failed with error %d\n",
+			       read ? "Read" : "Write", mtd->name, ret);
+			return CMD_RET_FAILURE;
+		}
+
+	} else if (!strcmp(cmd, "erase")) {
+		bool scrub = strstr(cmd, ".dontskipbad");
+		struct erase_info erase_op = {};
+		u64 off, len;
+		int ret;
+
+		off = argc > 0 ? simple_strtoul(argv[0], NULL, 16) : 0;
+		len = argc > 1 ? simple_strtoul(argv[1], NULL, 16) : mtd->size;
+
+		if (!mtd_is_aligned_with_block_size(mtd, off)) {
+			printf("Offset not aligned with a block (0x%x)\n",
+			       mtd->erasesize);
+			return CMD_RET_FAILURE;
+		}
+
+		if (!mtd_is_aligned_with_block_size(mtd, len)) {
+			printf("Size not a multiple of a block (0x%x)\n",
+			       mtd->erasesize);
+			return CMD_RET_FAILURE;
+		}
+
+		printf("Erasing 0x%08llx ... 0x%08llx (%d eraseblock(s))\n",
+		       off, off + len - 1, mtd_div_by_eb(len, mtd));
+
+		erase_op.mtd = mtd;
+		erase_op.addr = off;
+		erase_op.len = len;
+		erase_op.scrub = scrub;
+
+		while (erase_op.len) {
+			ret = mtd_erase(mtd, &erase_op);
+
+			/* Abort if its not a bad block error */
+			if (ret != -EIO)
+				break;
+
+			printf("Skipping bad block at 0x%08llx\n",
+			       erase_op.fail_addr);
+
+			/* Skip bad block and continue behind it */
+			erase_op.len -= erase_op.fail_addr - erase_op.addr;
+			erase_op.len -= mtd->erasesize;
+			erase_op.addr = erase_op.fail_addr + mtd->erasesize;
+		}
+
+		if (ret && ret != -EIO)
+			return CMD_RET_FAILURE;
+	} else if (!strcmp(cmd, "bad")) {
+		loff_t off;
+
+		if (!mtd_can_have_bb(mtd)) {
+			printf("Only NAND-based devices can have bad blocks\n");
+			return CMD_RET_SUCCESS;
+		}
+
+		printf("MTD device %s bad blocks list:\n", mtd->name);
+		for (off = 0; off < mtd->size; off += mtd->erasesize)
+			if (mtd_block_isbad(mtd, off))
+				printf("\t0x%08llx\n", off);
+	} else {
+		return CMD_RET_USAGE;
+	}
+
+	return CMD_RET_SUCCESS;
+}
+
+static char mtd_help_text[] =
+#ifdef CONFIG_SYS_LONGHELP
+	"- generic operations on memory technology devices\n\n"
+	"mtd list\n"
+	"mtd read[.raw][.oob]                  <name> <addr> [<off> [<size>]]\n"
+	"mtd dump[.raw][.oob]                  <name>        [<off> [<size>]]\n"
+	"mtd write[.raw][.oob][.dontskipff]    <name> <addr> [<off> [<size>]]\n"
+	"mtd erase[.dontskipbad]               <name>        [<off> [<size>]]\n"
+	"\n"
+	"Specific functions:\n"
+	"mtd bad                               <name>\n"
+	"\n"
+	"With:\n"
+	"\t<name>: NAND partition/chip name\n"
+	"\t<addr>: user address from/to which data will be retrieved/stored\n"
+	"\t<off>: offset in <name> in bytes (default: start of the part)\n"
+	"\t\t* must be block-aligned for erase\n"
+	"\t\t* must be page-aligned otherwise\n"
+	"\t<size>: length of the operation in bytes (default: the entire device)\n"
+	"\t\t* must be a multiple of a block for erase\n"
+	"\t\t* must be a multiple of a page otherwise (special case: default is a page with dump)\n"
+	"\n"
+	"The .dontskipff option forces writing empty pages, don't use it if unsure.\n"
+#endif
+	"";
+
+U_BOOT_CMD(mtd, 10, 1, do_mtd, "MTD utils", mtd_help_text);

cmd/mtdparts.c

@@ -37,14 +37,14 @@
  * mtdids=<idmap>[,<idmap>,...]
  *
  * <idmap>    := <dev-id>=<mtd-id>
- * <dev-id>   := 'nand'|'nor'|'onenand'<dev-num>
+ * <dev-id>   := 'nand'|'nor'|'onenand'|'spi-nand'<dev-num>
  * <dev-num>  := mtd device number, 0...
  * <mtd-id>   := unique device tag used by linux kernel to find mtd device (mtd->name)
  *
  *
  * 'mtdparts' - partition list
  *
- * mtdparts=mtdparts=<mtd-def>[;<mtd-def>...]
+ * mtdparts=[mtdparts=]<mtd-def>[;<mtd-def>...]
  *
  * <mtd-def>  := <mtd-id>:<part-def>[,<part-def>...]
  * <mtd-id>   := unique device tag used by linux kernel to find mtd device (mtd->name)
@@ -62,11 +62,11 @@
  *
  * 1 NOR Flash, with 1 single writable partition:
  * mtdids=nor0=edb7312-nor
- * mtdparts=mtdparts=edb7312-nor:-
+ * mtdparts=[mtdparts=]edb7312-nor:-
  *
  * 1 NOR Flash with 2 partitions, 1 NAND with one
  * mtdids=nor0=edb7312-nor,nand0=edb7312-nand
- * mtdparts=mtdparts=edb7312-nor:256k(ARMboot)ro,-(root);edb7312-nand:-(home)
+ * mtdparts=[mtdparts=]edb7312-nor:256k(ARMboot)ro,-(root);edb7312-nand:-(home)
  *
  */
 
@@ -177,13 +177,16 @@ static u64 memsize_parse (const char *const ptr, const char **retptr)
 		case 'G':
 		case 'g':
 			ret <<= 10;
+			/* Fallthrough */
 		case 'M':
 		case 'm':
 			ret <<= 10;
+			/* Fallthrough */
 		case 'K':
 		case 'k':
 			ret <<= 10;
 			(*retptr)++;
+			/* Fallthrough */
 		default:
 			break;
 	}
@@ -336,7 +339,7 @@ static int part_validate_eraseblock(struct mtdids *id, struct part_info *part)
 
 	if (!mtd->numeraseregions) {
 		/*
-		 * Only one eraseregion (NAND, OneNAND or uniform NOR),
+		 * Only one eraseregion (NAND, SPI-NAND, OneNAND or uniform NOR),
 		 * checking for alignment is easy here
 		 */
 		offset = part->offset;
@@ -1027,7 +1030,7 @@ static struct mtdids* id_find_by_mtd_id(const char *mtd_id, unsigned int mtd_id_
 }
 
 /**
- * Parse device id string <dev-id> := 'nand'|'nor'|'onenand'<dev-num>,
+ * Parse device id string <dev-id> := 'nand'|'nor'|'onenand'|'spi-nand'<dev-num>,
  * return device type and number.
  *
  * @param id string describing device id
@@ -1051,6 +1054,9 @@ int mtd_id_parse(const char *id, const char **ret_id, u8 *dev_type,
 	} else if (strncmp(p, "onenand", 7) == 0) {
 		*dev_type = MTD_DEV_TYPE_ONENAND;
 		p += 7;
+	} else if (strncmp(p, "spi-nand", 8) == 0) {
+		*dev_type = MTD_DEV_TYPE_SPINAND;
+		p += 8;
 	} else {
 		printf("incorrect device type in %s\n", id);
 		return 1;
@@ -1093,9 +1099,6 @@ static int generate_mtdparts(char *buf, u32 buflen)
 		return 0;
 	}
 
-	strcpy(p, "mtdparts=");
-	p += 9;
-
 	list_for_each(dentry, &devices) {
 		dev = list_entry(dentry, struct mtd_device, link);
 
@@ -1566,11 +1569,9 @@ static int parse_mtdparts(const char *const mtdparts)
 	if (!p)
 		p = mtdparts;
 
-	if (strncmp(p, "mtdparts=", 9) != 0) {
-		printf("mtdparts variable doesn't start with 'mtdparts='\n");
-		return err;
-	}
-	p += 9;
+	/* Skip the useless prefix, if any */
+	if (strncmp(p, "mtdparts=", 9) == 0)
+		p += 9;
 
 	while (*p != '\0') {
 		err = 1;
@@ -1633,7 +1634,7 @@ static int parse_mtdids(const char *const ids)
 	while(p && (*p != '\0')) {
 
 		ret = 1;
-		/* parse 'nor'|'nand'|'onenand'<dev-num> */
+		/* parse 'nor'|'nand'|'onenand'|'spi-nand'<dev-num> */
 		if (mtd_id_parse(p, &p, &type, &num) != 0)
 			break;
 
@@ -2109,7 +2110,7 @@ static char mtdparts_help_text[] =
 	"'mtdids' - linux kernel mtd device id <-> u-boot device id mapping\n\n"
 	"mtdids=<idmap>[,<idmap>,...]\n\n"
 	"<idmap>    := <dev-id>=<mtd-id>\n"
-	"<dev-id>   := 'nand'|'nor'|'onenand'<dev-num>\n"
+	"<dev-id>   := 'nand'|'nor'|'onenand'|'spi-nand'<dev-num>\n"
 	"<dev-num>  := mtd device number, 0...\n"
 	"<mtd-id>   := unique device tag used by linux kernel to find mtd device (mtd->name)\n\n"
 	"'mtdparts' - partition list\n\n"

cmd/ubi.c

@@ -15,6 +15,7 @@
 #include <command.h>
 #include <exports.h>
 #include <memalign.h>
+#include <mtd.h>
 #include <nand.h>
 #include <onenand_uboot.h>
 #include <linux/mtd/mtd.h>
@@ -27,24 +28,8 @@
 #undef ubi_msg
 #define ubi_msg(fmt, ...) printf("UBI: " fmt "\n", ##__VA_ARGS__)
 
-#define DEV_TYPE_NONE		0
-#define DEV_TYPE_NAND		1
-#define DEV_TYPE_ONENAND	2
-#define DEV_TYPE_NOR		3
-
 /* Private own data */
 static struct ubi_device *ubi;
-static char buffer[80];
-static int ubi_initialized;
-
-struct selected_dev {
-	char part_name[80];
-	int selected;
-	int nr;
-	struct mtd_info *mtd_info;
-};
-
-static struct selected_dev ubi_dev;
 
 #ifdef CONFIG_CMD_UBIFS
 #include <ubifs_uboot.h>
@@ -408,43 +393,24 @@ int ubi_volume_read(char *volume, char *buf, size_t size)
 	return err;
 }
 
-static int ubi_dev_scan(struct mtd_info *info, char *ubidev,
-		const char *vid_header_offset)
+static int ubi_dev_scan(struct mtd_info *info, const char *vid_header_offset)
 {
-	struct mtd_device *dev;
-	struct part_info *part;
-	struct mtd_partition mtd_part;
 	char ubi_mtd_param_buffer[80];
-	u8 pnum;
 	int err;
 
-	if (find_dev_and_part(ubidev, &dev, &pnum, &part) != 0)
-		return 1;
+	if (!vid_header_offset)
+		sprintf(ubi_mtd_param_buffer, "%s", info->name);
+	else
+		sprintf(ubi_mtd_param_buffer, "%s,%s", info->name,
+			vid_header_offset);
 
-	sprintf(buffer, "mtd=%d", pnum);
-	memset(&mtd_part, 0, sizeof(mtd_part));
-	mtd_part.name = buffer;
-	mtd_part.size = part->size;
-	mtd_part.offset = part->offset;
-	add_mtd_partitions(info, &mtd_part, 1);
-
-	strcpy(ubi_mtd_param_buffer, buffer);
-	if (vid_header_offset)
-		sprintf(ubi_mtd_param_buffer, "mtd=%d,%s", pnum,
-				vid_header_offset);
 	err = ubi_mtd_param_parse(ubi_mtd_param_buffer, NULL);
-	if (err) {
-		del_mtd_partitions(info);
+	if (err)
 		return -err;
-	}
 
 	err = ubi_init();
-	if (err) {
-		del_mtd_partitions(info);
+	if (err)
 		return -err;
-	}
-
-	ubi_initialized = 1;
 
 	return 0;
 }
@@ -469,50 +435,33 @@ int ubi_detach(void)
 	/*
 	 * Call ubi_exit() before re-initializing the UBI subsystem
 	 */
-	if (ubi_initialized) {
+	if (ubi)
 		ubi_exit();
-		del_mtd_partitions(ubi_dev.mtd_info);
-		ubi_initialized = 0;
-	}
 
-	ubi_dev.selected = 0;
+	ubi = NULL;
+
 	return 0;
 }
 
 int ubi_part(char *part_name, const char *vid_header_offset)
 {
+	struct mtd_info *mtd;
 	int err = 0;
-	char mtd_dev[16];
-	struct mtd_device *dev;
-	struct part_info *part;
-	u8 pnum;
 
 	ubi_detach();
-	/*
-	 * Search the mtd device number where this partition
-	 * is located
-	 */
-	if (find_dev_and_part(part_name, &dev, &pnum, &part)) {
+
+	mtd_probe_devices();
+	mtd = get_mtd_device_nm(part_name);
+	if (IS_ERR(mtd)) {
 		printf("Partition %s not found!\n", part_name);
 		return 1;
 	}
-	sprintf(mtd_dev, "%s%d", MTD_DEV_TYPE(dev->id->type), dev->id->num);
-	ubi_dev.mtd_info = get_mtd_device_nm(mtd_dev);
-	if (IS_ERR(ubi_dev.mtd_info)) {
-		printf("Partition %s not found on device %s!\n", part_name,
-		       mtd_dev);
-		return 1;
-	}
+	put_mtd_device(mtd);
 
-	ubi_dev.selected = 1;
-
-	strcpy(ubi_dev.part_name, part_name);
-	err = ubi_dev_scan(ubi_dev.mtd_info, ubi_dev.part_name,
-			vid_header_offset);
+	err = ubi_dev_scan(mtd, vid_header_offset);
 	if (err) {
 		printf("UBI init error %d\n", err);
 		printf("Please check, if the correct MTD partition is used (size big enough?)\n");
-		ubi_dev.selected = 0;
 		return err;
 	}
 
@@ -543,13 +492,13 @@ static int do_ubi(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
 
 		/* Print current partition */
 		if (argc == 2) {
-			if (!ubi_dev.selected) {
-				printf("Error, no UBI device/partition selected!\n");
+			if (!ubi) {
+				printf("Error, no UBI device selected!\n");
 				return 1;
 			}
 
-			printf("Device %d: %s, partition %s\n",
-			       ubi_dev.nr, ubi_dev.mtd_info->name, ubi_dev.part_name);
+			printf("Device %d: %s, MTD partition %s\n",
+			       ubi->ubi_num, ubi->ubi_name, ubi->mtd->name);
 			return 0;
 		}
 
@@ -562,8 +511,8 @@ static int do_ubi(cmd_tbl_t *cmdtp, int flag, int argc, char * const argv[])
 		return ubi_part(argv[2], vid_header_offset);
 	}
 
-	if ((strcmp(argv[1], "part") != 0) && (!ubi_dev.selected)) {
-		printf("Error, no UBI device/partition selected!\n");
+	if ((strcmp(argv[1], "part") != 0) && !ubi) {
+		printf("Error, no UBI device selected!\n");
 		return 1;
 	}