diff --git a/drivers/gpu/drm/i915/i915_drv.h b/drivers/gpu/drm/i915/i915_drv.h
index 5373b171bb96..a8d300280a2c 100644
--- a/drivers/gpu/drm/i915/i915_drv.h
+++ b/drivers/gpu/drm/i915/i915_drv.h
@@ -617,6 +617,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);
 };
 
 enum intel_pch {
diff --git a/drivers/gpu/drm/i915/intel_psr.c b/drivers/gpu/drm/i915/intel_psr.c
index fc7c36efd401..bcaac9e69f8c 100644
--- a/drivers/gpu/drm/i915/intel_psr.c
+++ b/drivers/gpu/drm/i915/intel_psr.c
@@ -838,53 +838,67 @@ 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)
+static void hsw_psr_exit(struct intel_dp *intel_dp)
 {
-	struct intel_dp *intel_dp = dev_priv->psr.enabled;
-	struct drm_crtc *crtc = dp_to_dig_port(intel_dp)->base.base.crtc;
+	struct intel_digital_port *dig_port = dp_to_dig_port(intel_dp);
+	struct drm_device *dev = dig_port->base.base.dev;
+	struct drm_i915_private *dev_priv = to_i915(dev);
+	u32 val;
+
+	if (dev_priv->psr.psr2_enabled) {
+		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);
+	}
+}
+
+static void vlv_psr_exit(struct intel_dp *intel_dp)
+{
+	struct intel_digital_port *dig_port = dp_to_dig_port(intel_dp);
+	struct drm_device *dev = dig_port->base.base.dev;
+	struct drm_i915_private *dev_priv = to_i915(dev);
+	struct drm_crtc *crtc = dig_port->base.base.crtc;
 	enum pipe pipe = to_intel_crtc(crtc)->pipe;
 	u32 val;
 
-	if (!dev_priv->psr.active)
-		return;
+	val = I915_READ(VLV_PSRCTL(pipe));
 
-	if (HAS_DDI(dev_priv)) {
-		if (dev_priv->psr.psr2_enabled) {
-			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);
 
-		/*
-		 * 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);
+}
 
-		/*
-		 * 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);
-	}
+static void intel_psr_exit(struct drm_i915_private *dev_priv)
+{
+	struct intel_dp *intel_dp = dev_priv->psr.enabled;
+
+	if (!dev_priv->psr.active)
+		return;
 
+	dev_priv->psr.exit(intel_dp);
 	dev_priv->psr.active = false;
 }
 
@@ -1094,6 +1108,7 @@ 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.has_hw_tracking = true;
 		dev_priv->psr.enable_source = hsw_psr_enable_source;
@@ -1101,5 +1116,6 @@ void intel_psr_init(struct drm_i915_private *dev_priv)
 		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;
 	}
 }
