{
        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;
 {
        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
 
 
        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