diff --git a/drivers/gpu/drm/armada/armada_drv.c b/drivers/gpu/drm/armada/armada_drv.c
index 4f2b28354915..069f64533ac3 100644
--- a/drivers/gpu/drm/armada/armada_drv.c
+++ b/drivers/gpu/drm/armada/armada_drv.c
@@ -128,6 +128,7 @@ static int armada_drm_load(struct drm_device *dev, unsigned long flags)
 		return -ENOMEM;
 	}
 
+	platform_set_drvdata(dev->platformdev, dev);
 	dev->dev_private = priv;
 
 	/* Get the implementation specific driver data. */
@@ -376,7 +377,7 @@ static int armada_drm_probe(struct platform_device *pdev)
 
 static int armada_drm_remove(struct platform_device *pdev)
 {
-	drm_platform_exit(&armada_drm_driver, pdev);
+	drm_put_dev(platform_get_drvdata(pdev));
 	return 0;
 }