2 files changed, 5 insertions, 4 deletions

diff --git a/drivers/mtd/mtdcore.c b/drivers/mtd/mtdcore.c
index fac54a3fa3f1..00ebf7af7467 100644
--- a/drivers/mtd/mtdcore.c
+++ b/drivers/mtd/mtdcore.c
@@ -65,8 +65,8 @@ static void mtd_release(struct device *dev)
 static int mtd_cls_suspend(struct device *dev, pm_message_t state)
 {
         struct mtd_info *mtd = dev_to_mtd(dev);
-        
-        if (mtd->suspend)
+
+        if (mtd && mtd->suspend)
                 return mtd->suspend(mtd);
         else
                 return 0;
@@ -76,7 +76,7 @@ static int mtd_cls_resume(struct device *dev)
 {
         struct mtd_info *mtd = dev_to_mtd(dev);
         
-        if (mtd->resume)
+        if (mtd && mtd->resume)
                 mtd->resume(mtd);
         return 0;
 }
@@ -298,6 +298,7 @@ int add_mtd_device(struct mtd_info *mtd)
                         mtd->dev.class = &mtd_class;
                         mtd->dev.devt = MTD_DEVT(i);
                         dev_set_name(&mtd->dev, "mtd%d", i);
+                        dev_set_drvdata(&mtd->dev, mtd);
                         if (device_register(&mtd->dev) != 0) {
                                 mtd_table[i] = NULL;
                                 break;
diff --git a/include/linux/mtd/mtd.h b/include/linux/mtd/mtd.h
index 5675b63a0631..0f32a9b6ff55 100644
--- a/include/linux/mtd/mtd.h
+++ b/include/linux/mtd/mtd.h
@@ -251,7 +251,7 @@ struct mtd_info {
 
 static inline struct mtd_info *dev_to_mtd(struct device *dev)
 {
-        return dev ? container_of(dev, struct mtd_info, dev) : NULL;
+        return dev ? dev_get_drvdata(dev) : NULL;
 }
 
 static inline uint32_t mtd_div_by_eb(uint64_t sz, struct mtd_info *mtd)