diff --git a/drivers/gpu/drm/msm/msm_drv.c b/drivers/gpu/drm/msm/msm_drv.c
index 5e758d95751a..aa145d6f6469 100644
--- a/drivers/gpu/drm/msm/msm_drv.c
+++ b/drivers/gpu/drm/msm/msm_drv.c
@@ -1070,18 +1070,15 @@ static int msm_pm_suspend(struct device *dev)
 {
 	struct drm_device *ddev = dev_get_drvdata(dev);
 	struct msm_drm_private *priv = ddev->dev_private;
-	struct msm_kms *kms = priv->kms;
 
-	/* TODO: Use atomic helper suspend/resume */
-	if (kms && kms->funcs && kms->funcs->pm_suspend)
-		return kms->funcs->pm_suspend(dev);
-
-	drm_kms_helper_poll_disable(ddev);
+	if (WARN_ON(priv->pm_state))
+		drm_atomic_state_put(priv->pm_state);
 
 	priv->pm_state = drm_atomic_helper_suspend(ddev);
 	if (IS_ERR(priv->pm_state)) {
-		drm_kms_helper_poll_enable(ddev);
-		return PTR_ERR(priv->pm_state);
+		int ret = PTR_ERR(priv->pm_state);
+		DRM_ERROR("Failed to suspend dpu, %d\n", ret);
+		return ret;
 	}
 
 	return 0;
@@ -1091,16 +1088,16 @@ static int msm_pm_resume(struct device *dev)
 {
 	struct drm_device *ddev = dev_get_drvdata(dev);
 	struct msm_drm_private *priv = ddev->dev_private;
-	struct msm_kms *kms = priv->kms;
+	int ret;
 
-	/* TODO: Use atomic helper suspend/resume */
-	if (kms && kms->funcs && kms->funcs->pm_resume)
-		return kms->funcs->pm_resume(dev);
+	if (WARN_ON(!priv->pm_state))
+		return -ENOENT;
 
-	drm_atomic_helper_resume(ddev, priv->pm_state);
-	drm_kms_helper_poll_enable(ddev);
+	ret = drm_atomic_helper_resume(ddev, priv->pm_state);
+	if (!ret)
+		priv->pm_state = NULL;
 
-	return 0;
+	return ret;
 }
 #endif
 
diff --git a/drivers/gpu/drm/msm/msm_kms.h b/drivers/gpu/drm/msm/msm_kms.h
index fd88cebb6adb..2b81b43a4bab 100644
--- a/drivers/gpu/drm/msm/msm_kms.h
+++ b/drivers/gpu/drm/msm/msm_kms.h
@@ -67,9 +67,6 @@ struct msm_kms_funcs {
 	void (*set_encoder_mode)(struct msm_kms *kms,
 				 struct drm_encoder *encoder,
 				 bool cmd_mode);
-	/* pm suspend/resume hooks */
-	int (*pm_suspend)(struct device *dev);
-	int (*pm_resume)(struct device *dev);
 	/* cleanup: */
 	void (*destroy)(struct msm_kms *kms);
 #ifdef CONFIG_DEBUG_FS