diff --git a/drivers/gpu/drm/i915/i915_drv.h b/drivers/gpu/drm/i915/i915_drv.h
index 82a106b1bdbc..44b551f1576d 100644
--- a/drivers/gpu/drm/i915/i915_drv.h
+++ b/drivers/gpu/drm/i915/i915_drv.h
@@ -780,6 +780,7 @@ struct i915_psr {
 	void (*enable_sink)(struct intel_dp *);
 	void (*activate)(struct intel_dp *);
 	void (*setup_vsc)(struct intel_dp *, const struct intel_crtc_state *);
+	void (*exit)(struct intel_dp *intel_dp, bool wait_exit);
 };
 
 enum intel_pch {
diff --git a/drivers/gpu/drm/i915/intel_psr.c b/drivers/gpu/drm/i915/intel_psr.c
index 2ef374f936b9..e8c32c3afb0e 100644
--- a/drivers/gpu/drm/i915/intel_psr.c
+++ b/drivers/gpu/drm/i915/intel_psr.c
@@ -551,25 +551,12 @@ static void vlv_psr_disable(struct intel_dp *intel_dp,
 	struct intel_crtc *crtc = to_intel_crtc(old_crtc_state->base.crtc);
 	uint32_t val;
 
-	if (dev_priv->psr.active) {
-		/* Put VLV PSR back to PSR_state 0 (disabled). */
-		if (intel_wait_for_register(dev_priv,
-					    VLV_PSRSTAT(crtc->pipe),
-					    VLV_EDP_PSR_IN_TRANS,
-					    0,
-					    1))
-			WARN(1, "PSR transition took longer than expected\n");
+	dev_priv->psr.exit(intel_dp, true);
 
-		val = I915_READ(VLV_PSRCTL(crtc->pipe));
-		val &= ~VLV_EDP_PSR_ACTIVE_ENTRY;
-		val &= ~VLV_EDP_PSR_ENABLE;
-		val &= ~VLV_EDP_PSR_MODE_MASK;
-		I915_WRITE(VLV_PSRCTL(crtc->pipe), val);
-
-		dev_priv->psr.active = false;
-	} else {
-		WARN_ON(vlv_is_psr_active_on_pipe(dev, crtc->pipe));
-	}
+	val = I915_READ(VLV_PSRCTL(crtc->pipe));
+	val &= ~VLV_EDP_PSR_ENABLE;
+	val &= ~VLV_EDP_PSR_MODE_MASK;
+	I915_WRITE(VLV_PSRCTL(crtc->pipe), val);
 }
 
 static void hsw_psr_disable(struct intel_dp *intel_dp,
@@ -579,44 +566,100 @@ static void hsw_psr_disable(struct intel_dp *intel_dp,
 	struct drm_device *dev = intel_dig_port->base.base.dev;
 	struct drm_i915_private *dev_priv = to_i915(dev);
 
-	if (dev_priv->psr.active) {
-		i915_reg_t psr_status;
-		u32 psr_status_mask;
-
-		if (dev_priv->psr.aux_frame_sync)
-			drm_dp_dpcd_writeb(&intel_dp->aux,
-					DP_SINK_DEVICE_AUX_FRAME_SYNC_CONF,
-					0);
+	dev_priv->psr.exit(intel_dp, true);
+}
 
-		if (dev_priv->psr.psr2_support) {
-			psr_status = EDP_PSR2_STATUS;
-			psr_status_mask = EDP_PSR2_STATUS_STATE_MASK;
+static void hsw_psr_exit(struct intel_dp *intel_dp, bool wait_exit)
+{
+	struct intel_digital_port *intel_dig_port = dp_to_dig_port(intel_dp);
+	struct drm_device *dev = intel_dig_port->base.base.dev;
+	struct drm_i915_private *dev_priv = to_i915(dev);
+	i915_reg_t psr_status;
+	u32 psr_status_mask;
 
-			I915_WRITE(EDP_PSR2_CTL,
-				   I915_READ(EDP_PSR2_CTL) &
-				   ~(EDP_PSR2_ENABLE | EDP_SU_TRACK_ENABLE));
+	if (!dev_priv->psr.active) {
+		if (dev_priv->psr.psr2_support)
+			WARN_ON(I915_READ(EDP_PSR2_CTL) & EDP_PSR2_ENABLE);
+		else
+			WARN_ON(I915_READ(EDP_PSR_CTL) & EDP_PSR_ENABLE);
+		return;
+	}
 
-		} else {
-			psr_status = EDP_PSR_STATUS;
-			psr_status_mask = EDP_PSR_STATUS_STATE_MASK;
+	if (dev_priv->psr.aux_frame_sync)
+		drm_dp_dpcd_writeb(&intel_dp->aux,
+				   DP_SINK_DEVICE_AUX_FRAME_SYNC_CONF, 0);
 
-			I915_WRITE(EDP_PSR_CTL,
-				   I915_READ(EDP_PSR_CTL) & ~EDP_PSR_ENABLE);
-		}
+	if (dev_priv->psr.psr2_support) {
+		psr_status = EDP_PSR2_STATUS;
+		psr_status_mask = EDP_PSR2_STATUS_STATE_MASK;
 
-		/* Wait till PSR is idle */
-		if (intel_wait_for_register(dev_priv,
-					    psr_status, psr_status_mask, 0,
-					    2000))
-			DRM_ERROR("Timed out waiting for PSR Idle State\n");
+		I915_WRITE(EDP_PSR2_CTL,
+			   I915_READ(EDP_PSR2_CTL) &
+			   ~(EDP_PSR2_ENABLE | EDP_SU_TRACK_ENABLE));
 
-		dev_priv->psr.active = false;
 	} else {
-		if (dev_priv->psr.psr2_support)
-			WARN_ON(I915_READ(EDP_PSR2_CTL) & EDP_PSR2_ENABLE);
-		else
-			WARN_ON(I915_READ(EDP_PSR_CTL) & EDP_PSR_ENABLE);
+		psr_status = EDP_PSR_STATUS;
+		psr_status_mask = EDP_PSR_STATUS_STATE_MASK;
+
+		I915_WRITE(EDP_PSR_CTL,
+			   I915_READ(EDP_PSR_CTL) & ~EDP_PSR_ENABLE);
+	}
+
+	/* Wait till PSR is idle */
+	if (wait_exit && intel_wait_for_register(dev_priv, psr_status,
+						 psr_status_mask, 0, 2000))
+		DRM_ERROR("Timed out waiting for PSR Idle State\n");
+
+	dev_priv->psr.active = false;
+}
+
+static void vlv_psr_exit(struct intel_dp *intel_dp, bool wait_exit)
+{
+	struct intel_digital_port *intel_dig_port = dp_to_dig_port(intel_dp);
+	struct drm_device *dev = intel_dig_port->base.base.dev;
+	struct drm_i915_private *dev_priv = to_i915(dev);
+	struct intel_crtc *crtc = to_intel_crtc(intel_dig_port->base.base.crtc);
+	uint32_t val;
+
+	if (!dev_priv->psr.active) {
+		WARN_ON(vlv_is_psr_active_on_pipe(dev, crtc->pipe));
+		return;
 	}
+
+	if (intel_wait_for_register(dev_priv, VLV_PSRSTAT(crtc->pipe),
+				    VLV_EDP_PSR_IN_TRANS, 0, 1))
+		WARN(1, "PSR transition took longer than expected\n");
+
+	/*
+	 * Here we do the transition directly from
+	 * PSR_state 3 (active - no Remote Frame Buffer (RFB) update) to
+	 * PSR_state 5 (exit).
+	 * PSR State 4 (active with single frame update) can be skipped.
+	 * On PSR_state 5 (exit) Hardware is responsible to transition
+	 * back to PSR_state 1 (inactive).
+	 * Now we are at Same state after vlv_psr_enable_source.
+	 */
+	val = I915_READ(VLV_PSRCTL(crtc->pipe));
+	val &= ~VLV_EDP_PSR_ACTIVE_ENTRY;
+	I915_WRITE(VLV_PSRCTL(crtc->pipe), val);
+
+	/*
+	 * Send AUX wake up - Spec says after transitioning to PSR
+	 * active we have to send AUX wake up by writing 01h in DPCD
+	 * 600h of sink device.
+	 * XXX: This might slow down the transition, but without this
+	 * HW doesn't complete the transition to PSR_state 1 and we
+	 * never get the screen updated.
+	 */
+	drm_dp_dpcd_writeb(&intel_dp->aux, DP_SET_POWER, DP_SET_POWER_D0);
+
+	if (wait_exit && intel_wait_for_register(dev_priv,
+						 VLV_PSRSTAT(crtc->pipe),
+						 VLV_EDP_PSR_INACTIVE,
+						 VLV_EDP_PSR_INACTIVE, 2000))
+		DRM_ERROR("Timed out waiting for PSR inactive state\n");
+
+	dev_priv->psr.active = false;
 }
 
 /**
@@ -718,60 +761,6 @@ static void intel_psr_work(struct work_struct *work)
 	mutex_unlock(&dev_priv->psr.lock);
 }
 
-static void intel_psr_exit(struct drm_i915_private *dev_priv)
-{
-	struct intel_dp *intel_dp = dev_priv->psr.enabled;
-	struct drm_crtc *crtc = dp_to_dig_port(intel_dp)->base.base.crtc;
-	enum pipe pipe = to_intel_crtc(crtc)->pipe;
-	u32 val;
-
-	if (!dev_priv->psr.active)
-		return;
-
-	if (HAS_DDI(dev_priv)) {
-		if (dev_priv->psr.aux_frame_sync)
-			drm_dp_dpcd_writeb(&intel_dp->aux,
-					DP_SINK_DEVICE_AUX_FRAME_SYNC_CONF,
-					0);
-		if (dev_priv->psr.psr2_support) {
-			val = I915_READ(EDP_PSR2_CTL);
-			WARN_ON(!(val & EDP_PSR2_ENABLE));
-			I915_WRITE(EDP_PSR2_CTL, val & ~EDP_PSR2_ENABLE);
-		} else {
-			val = I915_READ(EDP_PSR_CTL);
-			WARN_ON(!(val & EDP_PSR_ENABLE));
-			I915_WRITE(EDP_PSR_CTL, val & ~EDP_PSR_ENABLE);
-		}
-	} else {
-		val = I915_READ(VLV_PSRCTL(pipe));
-
-		/*
-		 * Here we do the transition drirectly from
-		 * PSR_state 3 (active - no Remote Frame Buffer (RFB) update) to
-		 * PSR_state 5 (exit).
-		 * PSR State 4 (active with single frame update) can be skipped.
-		 * On PSR_state 5 (exit) Hardware is responsible to transition
-		 * back to PSR_state 1 (inactive).
-		 * Now we are at Same state after vlv_psr_enable_source.
-		 */
-		val &= ~VLV_EDP_PSR_ACTIVE_ENTRY;
-		I915_WRITE(VLV_PSRCTL(pipe), val);
-
-		/*
-		 * Send AUX wake up - Spec says after transitioning to PSR
-		 * active we have to send AUX wake up by writing 01h in DPCD
-		 * 600h of sink device.
-		 * XXX: This might slow down the transition, but without this
-		 * HW doesn't complete the transition to PSR_state 1 and we
-		 * never get the screen updated.
-		 */
-		drm_dp_dpcd_writeb(&intel_dp->aux, DP_SET_POWER,
-				   DP_SET_POWER_D0);
-	}
-
-	dev_priv->psr.active = false;
-}
-
 /**
  * intel_psr_single_frame_update - Single Frame Update
  * @dev_priv: i915 device
@@ -854,7 +843,7 @@ void intel_psr_invalidate(struct drm_i915_private *dev_priv,
 	dev_priv->psr.busy_frontbuffer_bits |= frontbuffer_bits;
 
 	if (frontbuffer_bits)
-		intel_psr_exit(dev_priv);
+		dev_priv->psr.exit(dev_priv->psr.enabled, false);
 
 	mutex_unlock(&dev_priv->psr.lock);
 }
@@ -895,7 +884,7 @@ void intel_psr_flush(struct drm_i915_private *dev_priv,
 
 	/* By definition flush = invalidate + flush */
 	if (frontbuffer_bits)
-		intel_psr_exit(dev_priv);
+		dev_priv->psr.exit(dev_priv->psr.enabled, false);
 
 	if (!dev_priv->psr.active && !dev_priv->psr.busy_frontbuffer_bits)
 		if (!work_busy(&dev_priv->psr.work.work))
@@ -956,11 +945,13 @@ void intel_psr_init(struct drm_i915_private *dev_priv)
 		dev_priv->psr.enable_sink = vlv_psr_enable_sink;
 		dev_priv->psr.activate = vlv_psr_activate;
 		dev_priv->psr.setup_vsc = vlv_psr_setup_vsc;
+		dev_priv->psr.exit = vlv_psr_exit;
 	} else {
 		dev_priv->psr.enable_source = hsw_psr_enable_source;
 		dev_priv->psr.disable_source = hsw_psr_disable;
 		dev_priv->psr.enable_sink = hsw_psr_enable_sink;
 		dev_priv->psr.activate = hsw_psr_activate;
 		dev_priv->psr.setup_vsc = hsw_psr_setup_vsc;
+		dev_priv->psr.exit = hsw_psr_exit;
 	}
 }
