--- a/drivers/mtd/devices/block2mtd.c
+++ b/drivers/mtd/devices/block2mtd.c
@@ -32,6 +32,8 @@
 #include <linux/slab.h>
 #include <linux/major.h>
 
+static const char * const block2mtd_probe_types[] = { "cmdlinepart", NULL };
+
 /* Info for the block device */
 struct block2mtd_dev {
 	struct list_head list;
@@ -227,6 +229,7 @@ static struct block2mtd_dev *add_device(
 #endif
 	const fmode_t mode = FMODE_READ | FMODE_WRITE | FMODE_EXCL;
 	struct block_device *bdev = ERR_PTR(-ENODEV);
+	struct mtd_part_parser_data ppdata = { 0 };
 	struct block2mtd_dev *dev;
 	struct mtd_partition *part;
 	char *name;
@@ -307,11 +310,7 @@ static struct block2mtd_dev *add_device(
 	dev->mtd.priv = dev;
 	dev->mtd.owner = THIS_MODULE;
 
-	part = kzalloc(sizeof(struct mtd_partition), GFP_KERNEL);
-	part->name = name;
-	part->offset = 0;
-	part->size = dev->mtd.size;
-	if (mtd_device_register(&dev->mtd, part, 1)) {
+	if (mtd_device_parse_register(&dev->mtd, block2mtd_probe_types, &ppdata, NULL, 0)) {
 		/* Device didn't get added, so free the entry */
 		goto err_destroy_mutex;
 	}