_intel_psr_post_plane_update(state, crtc_state);
 }
 
-/**
- * psr_wait_for_idle - wait for PSR1 to idle
- * @intel_dp: Intel DP
- * @out_value: PSR status in case of failure
- *
- * Returns: 0 on success or -ETIMEOUT if PSR status does not idle.
- *
- */
-static int psr_wait_for_idle(struct intel_dp *intel_dp, u32 *out_value)
+static int _psr2_ready_for_pipe_update_locked(struct intel_dp *intel_dp)
+{
+       struct drm_i915_private *dev_priv = dp_to_i915(intel_dp);
+
+       /*
+        * Any state lower than EDP_PSR2_STATUS_STATE_DEEP_SLEEP is enough.
+        * As all higher states has bit 4 of PSR2 state set we can just wait for
+        * EDP_PSR2_STATUS_STATE_DEEP_SLEEP to be cleared.
+        */
+       return intel_de_wait_for_clear(dev_priv,
+                                      EDP_PSR2_STATUS(intel_dp->psr.transcoder),
+                                      EDP_PSR2_STATUS_STATE_DEEP_SLEEP, 50);
+}
+
+static int _psr1_ready_for_pipe_update_locked(struct intel_dp *intel_dp)
 {
        struct drm_i915_private *dev_priv = dp_to_i915(intel_dp);
 
         * exit training time + 1.5 ms of aux channel handshake. 50 ms is
         * defensive enough to cover everything.
         */
-       return __intel_wait_for_register(&dev_priv->uncore,
-                                        EDP_PSR_STATUS(intel_dp->psr.transcoder),
-                                        EDP_PSR_STATUS_STATE_MASK,
-                                        EDP_PSR_STATUS_STATE_IDLE, 2, 50,
-                                        out_value);
+       return intel_de_wait_for_clear(dev_priv,
+                                      EDP_PSR_STATUS(intel_dp->psr.transcoder),
+                                      EDP_PSR_STATUS_STATE_MASK, 50);
 }
 
 /**
- * intel_psr_wait_for_idle - wait for PSR1 to idle
+ * intel_psr_wait_for_idle - wait for PSR be ready for a pipe update
  * @new_crtc_state: new CRTC state
  *
  * This function is expected to be called from pipe_update_start() where it is
        for_each_intel_encoder_mask_with_psr(&dev_priv->drm, encoder,
                                             new_crtc_state->uapi.encoder_mask) {
                struct intel_dp *intel_dp = enc_to_intel_dp(encoder);
-               u32 psr_status;
+               int ret;
 
                mutex_lock(&intel_dp->psr.lock);
-               if (!intel_dp->psr.enabled || intel_dp->psr.psr2_enabled) {
+
+               if (!intel_dp->psr.enabled) {
                        mutex_unlock(&intel_dp->psr.lock);
                        continue;
                }
 
-               /* when the PSR1 is enabled */
-               if (psr_wait_for_idle(intel_dp, &psr_status))
-                       drm_err(&dev_priv->drm,
-                               "PSR idle timed out 0x%x, atomic update may fail\n",
-                               psr_status);
+               if (intel_dp->psr.psr2_enabled)
+                       ret = _psr2_ready_for_pipe_update_locked(intel_dp);
+               else
+                       ret = _psr1_ready_for_pipe_update_locked(intel_dp);
+
+               if (ret)
+                       drm_err(&dev_priv->drm, "PSR wait timed out, atomic update may fail\n");
+
                mutex_unlock(&intel_dp->psr.lock);
        }
 }
 
 #define  PSR_EVENT_LPSP_MODE_EXIT              (1 << 1)
 #define  PSR_EVENT_PSR_DISABLE                 (1 << 0)
 
-#define _PSR2_STATUS_A                 0x60940
-#define _PSR2_STATUS_EDP               0x6f940
-#define EDP_PSR2_STATUS(tran)          _MMIO_TRANS2(tran, _PSR2_STATUS_A)
-#define EDP_PSR2_STATUS_STATE_MASK     (0xf << 28)
-#define EDP_PSR2_STATUS_STATE_SHIFT    28
+#define _PSR2_STATUS_A                         0x60940
+#define _PSR2_STATUS_EDP                       0x6f940
+#define EDP_PSR2_STATUS(tran)                  _MMIO_TRANS2(tran, _PSR2_STATUS_A)
+#define EDP_PSR2_STATUS_STATE_MASK             REG_GENMASK(31, 28)
+#define EDP_PSR2_STATUS_STATE_DEEP_SLEEP       REG_FIELD_PREP(EDP_PSR2_STATUS_STATE_MASK, 0x8)
 
 #define _PSR2_SU_STATUS_A              0x60914
 #define _PSR2_SU_STATUS_EDP            0x6f914