diff --git a/drivers/mtd/mtdcore.c b/drivers/mtd/mtdcore.c
index 06ec9f836ae5..f11f55f02413 100644
--- a/drivers/mtd/mtdcore.c
+++ b/drivers/mtd/mtdcore.c
@@ -214,12 +214,23 @@ struct mtd_info *get_mtd_device(struct mtd_info *mtd, int num)
 			ret = NULL;
 	}
 
-	if (ret && !try_module_get(ret->owner))
+	if (!ret)
+		goto out_unlock;
+
+	if (!try_module_get(ret->owner)) {
 		ret = NULL;
+		goto out_unlock;
+	}
 
-	if (ret)
-		ret->usecount++;
+	if (ret->get_device && ret->get_device(ret)) {
+		module_put(ret->owner);
+		ret = NULL;
+		goto out_unlock;
+	}
 
+	ret->usecount++;
+
+out_unlock:
 	mutex_unlock(&mtd_table_mutex);
 	return ret;
 }
@@ -235,8 +246,8 @@ struct mtd_info *get_mtd_device(struct mtd_info *mtd, int num)
 
 struct mtd_info *get_mtd_device_nm(const char *name)
 {
-	int i;
-	struct mtd_info *mtd = ERR_PTR(-ENODEV);
+	int i, err = -ENODEV;
+	struct mtd_info *mtd = NULL;
 
 	mutex_lock(&mtd_table_mutex);
 
@@ -247,17 +258,27 @@ struct mtd_info *get_mtd_device_nm(const char *name)
 		}
 	}
 
-	if (i == MAX_MTD_DEVICES)
+	if (!mtd)
 		goto out_unlock;
 
 	if (!try_module_get(mtd->owner))
 		goto out_unlock;
 
-	mtd->usecount++;
+	if (mtd->get_device) {
+		err = mtd->get_device(mtd);
+		if (err)
+			goto out_put;
+	}
 
-out_unlock:
+	mtd->usecount++;
 	mutex_unlock(&mtd_table_mutex);
 	return mtd;
+
+out_put:
+	module_put(mtd->owner);
+out_unlock:
+	mutex_unlock(&mtd_table_mutex);
+	return ERR_PTR(err);
 }
 
 void put_mtd_device(struct mtd_info *mtd)
@@ -266,6 +287,8 @@ void put_mtd_device(struct mtd_info *mtd)
 
 	mutex_lock(&mtd_table_mutex);
 	c = --mtd->usecount;
+	if (mtd->put_device)
+		mtd->put_device(mtd);
 	mutex_unlock(&mtd_table_mutex);
 	BUG_ON(c < 0);
 
diff --git a/include/linux/mtd/mtd.h b/include/linux/mtd/mtd.h
index 89e937dfef55..d644e57703ad 100644
--- a/include/linux/mtd/mtd.h
+++ b/include/linux/mtd/mtd.h
@@ -207,6 +207,13 @@ struct mtd_info {
 
 	struct module *owner;
 	int usecount;
+
+	/* If the driver is something smart, like UBI, it may need to maintain
+	 * its own reference counting. The below functions are only for driver.
+	 * The driver may register its callbacks. These callbacks are not
+	 * supposed to be called by MTD users */
+	int (*get_device) (struct mtd_info *mtd);
+	void (*put_device) (struct mtd_info *mtd);
 };