diff --git a/drivers/gpu/drm/i915/i915_drv.h b/drivers/gpu/drm/i915/i915_drv.h
index a10f31c9e4a9..5494582fdfea 100644
--- a/drivers/gpu/drm/i915/i915_drv.h
+++ b/drivers/gpu/drm/i915/i915_drv.h
@@ -397,6 +397,7 @@ enum intel_display_power_domain {
 	POWER_DOMAIN_AUX_C,
 	POWER_DOMAIN_AUX_D,
 	POWER_DOMAIN_GMBUS,
+	POWER_DOMAIN_VBLANK,
 	POWER_DOMAIN_MODESET,
 	POWER_DOMAIN_GT_IRQ,
 	POWER_DOMAIN_INIT,
@@ -1476,7 +1477,14 @@ struct i915_power_well {
 			bool has_vga:1;
 			bool has_fuses:1;
 		} hsw;
+		struct {
+			bool was_disabled;
+		} dc_off;
 	};
+
+	spinlock_t lock;
+	bool supports_atomic_ctx;
+
 	const struct i915_power_well_ops *ops;
 };
 
diff --git a/drivers/gpu/drm/i915/intel_drv.h b/drivers/gpu/drm/i915/intel_drv.h
index 30f791f89d64..164e62cb047b 100644
--- a/drivers/gpu/drm/i915/intel_drv.h
+++ b/drivers/gpu/drm/i915/intel_drv.h
@@ -1797,6 +1797,9 @@ bool intel_display_power_get_if_enabled(struct drm_i915_private *dev_priv,
 					enum intel_display_power_domain domain);
 void intel_display_power_put(struct drm_i915_private *dev_priv,
 			     enum intel_display_power_domain domain);
+void intel_display_power_vblank_get(struct drm_i915_private *dev_priv,
+				    bool *needs_restore);
+void intel_display_power_vblank_put(struct drm_i915_private *dev_priv);
 
 static inline void
 assert_rpm_device_not_suspended(struct drm_i915_private *dev_priv)
diff --git a/drivers/gpu/drm/i915/intel_runtime_pm.c b/drivers/gpu/drm/i915/intel_runtime_pm.c
index 992caec1fbc4..9cbb332cb418 100644
--- a/drivers/gpu/drm/i915/intel_runtime_pm.c
+++ b/drivers/gpu/drm/i915/intel_runtime_pm.c
@@ -56,6 +56,16 @@ static struct i915_power_well *
 lookup_power_well(struct drm_i915_private *dev_priv,
 		  enum i915_power_well_id power_well_id);
 
+/* optimize for the case when this function is called from atomic context,
+ * although this is unlikely */
+#define power_well_lock(power_well, flags)			\
+	if (likely(power_well->supports_atomic_ctx))		\
+		spin_lock_irqsave(&power_well->lock, flags)
+
+#define power_well_unlock(power_well, flags)			\
+	if (likely(power_well->supports_atomic_ctx))		\
+		spin_unlock_irqrestore(&power_well->lock, flags)
+
 const char *
 intel_display_power_domain_str(enum intel_display_power_domain domain)
 {
@@ -126,6 +136,8 @@ intel_display_power_domain_str(enum intel_display_power_domain domain)
 		return "AUX_D";
 	case POWER_DOMAIN_GMBUS:
 		return "GMBUS";
+	case POWER_DOMAIN_VBLANK:
+		return "VBLANK";
 	case POWER_DOMAIN_INIT:
 		return "INIT";
 	case POWER_DOMAIN_MODESET:
@@ -141,6 +153,9 @@ intel_display_power_domain_str(enum intel_display_power_domain domain)
 static void intel_power_well_enable(struct drm_i915_private *dev_priv,
 				    struct i915_power_well *power_well)
 {
+	if (power_well->supports_atomic_ctx)
+		assert_spin_locked(&power_well->lock);
+
 	DRM_DEBUG_KMS("enabling %s\n", power_well->name);
 	power_well->ops->enable(dev_priv, power_well);
 	power_well->hw_enabled = true;
@@ -149,19 +164,34 @@ static void intel_power_well_enable(struct drm_i915_private *dev_priv,
 static void intel_power_well_disable(struct drm_i915_private *dev_priv,
 				     struct i915_power_well *power_well)
 {
+	if (power_well->supports_atomic_ctx)
+		assert_spin_locked(&power_well->lock);
+
 	DRM_DEBUG_KMS("disabling %s\n", power_well->name);
 	power_well->hw_enabled = false;
 	power_well->ops->disable(dev_priv, power_well);
 }
 
-static void intel_power_well_get(struct drm_i915_private *dev_priv,
+
+static void __intel_power_well_get(struct drm_i915_private *dev_priv,
 				 struct i915_power_well *power_well)
 {
 	if (!power_well->count++)
 		intel_power_well_enable(dev_priv, power_well);
 }
 
-static void intel_power_well_put(struct drm_i915_private *dev_priv,
+
+static void intel_power_well_get(struct drm_i915_private *dev_priv,
+				 struct i915_power_well *power_well)
+{
+	unsigned long flags = 0;
+
+	power_well_lock(power_well, flags);
+	__intel_power_well_get(dev_priv, power_well);
+	power_well_unlock(power_well, flags);
+}
+
+static void __intel_power_well_put(struct drm_i915_private *dev_priv,
 				 struct i915_power_well *power_well)
 {
 	WARN(!power_well->count, "Use count on power well %s is already zero",
@@ -171,6 +201,16 @@ static void intel_power_well_put(struct drm_i915_private *dev_priv,
 		intel_power_well_disable(dev_priv, power_well);
 }
 
+static void intel_power_well_put(struct drm_i915_private *dev_priv,
+				 struct i915_power_well *power_well)
+{
+	unsigned long flags = 0;
+
+	power_well_lock(power_well, flags);
+	__intel_power_well_put(dev_priv, power_well);
+	power_well_unlock(power_well, flags);
+}
+
 /**
  * __intel_display_power_is_enabled - unlocked check for a power domain
  * @dev_priv: i915 device instance
@@ -726,6 +766,7 @@ static void gen9_dc_off_power_well_disable(struct drm_i915_private *dev_priv,
 		skl_enable_dc6(dev_priv);
 	else if (dev_priv->csr.allowed_dc_mask & DC_STATE_EN_UPTO_DC5)
 		gen9_enable_dc5(dev_priv);
+	power_well->dc_off.was_disabled = true;
 }
 
 static void i9xx_power_well_sync_hw_noop(struct drm_i915_private *dev_priv,
@@ -1443,6 +1484,63 @@ static void chv_pipe_power_well_disable(struct drm_i915_private *dev_priv,
 	chv_set_pipe_power_well(dev_priv, power_well, false);
 }
 
+#define CAN_PSR(dev_priv) (HAS_PSR(dev_priv) && dev_priv->psr.sink_support)
+void intel_display_power_vblank_get(struct drm_i915_private *dev_priv,
+				    bool *needs_restore)
+{
+	struct i915_power_domains *power_domains  = &dev_priv->power_domains;
+	struct i915_power_well *power_well;
+
+	*needs_restore = false;
+
+	if (!HAS_CSR(dev_priv))
+		return;
+
+	if (!CAN_PSR(dev_priv))
+		return;
+
+	intel_runtime_pm_get_noresume(dev_priv);
+
+	for_each_power_domain_well(dev_priv, power_well, BIT_ULL(POWER_DOMAIN_VBLANK)) {
+		unsigned long flags = 0;
+
+		power_well_lock(power_well, flags);
+		__intel_power_well_get(dev_priv, power_well);
+		*needs_restore = power_well->dc_off.was_disabled;
+		power_well->dc_off.was_disabled = false;
+		power_well_unlock(power_well, flags);
+	}
+
+	atomic_inc(&power_domains->domain_use_count[POWER_DOMAIN_VBLANK]);
+}
+
+void intel_display_power_vblank_put(struct drm_i915_private *dev_priv)
+{
+	struct i915_power_domains *power_domains = &dev_priv->power_domains;
+	struct i915_power_well *power_well;
+
+	if (!HAS_CSR(dev_priv))
+		return;
+
+	if (!CAN_PSR(dev_priv))
+		return;
+
+	WARN(atomic_dec_return(&power_domains->domain_use_count[POWER_DOMAIN_VBLANK]) < 0,
+	     "Use count on domain %s was already zero\n",
+	     intel_display_power_domain_str(POWER_DOMAIN_VBLANK));
+
+	for_each_power_domain_well_rev(dev_priv, power_well, BIT_ULL(POWER_DOMAIN_VBLANK)) {
+		unsigned long flags = 0;
+
+		power_well_lock(power_well, flags);
+		__intel_power_well_put(dev_priv, power_well);
+		power_well_unlock(power_well, flags);
+	}
+
+	intel_runtime_pm_put(dev_priv);
+}
+#undef CAN_PSR
+
 static void
 __intel_display_power_get_domain(struct drm_i915_private *dev_priv,
 				 enum intel_display_power_domain domain)
@@ -1482,6 +1580,38 @@ void intel_display_power_get(struct drm_i915_private *dev_priv,
 	mutex_unlock(&power_domains->lock);
 }
 
+static bool dc_off_get_if_enabled(struct drm_i915_private *dev_priv,
+				  enum intel_display_power_domain domain)
+{
+	struct i915_power_well *power_well;
+	bool is_enabled;
+	unsigned long flags = 0;
+
+	power_well = lookup_power_well(dev_priv, SKL_DISP_PW_DC_OFF);
+	if (!power_well || !(power_well->domains & domain))
+		return true;
+
+	power_well_lock(power_well, flags);
+	is_enabled = power_well->hw_enabled;
+	if (is_enabled)
+		__intel_power_well_get(dev_priv, power_well);
+	power_well_unlock(power_well, flags);
+
+	return is_enabled;
+}
+
+static void dc_off_put(struct drm_i915_private *dev_priv,
+		       enum intel_display_power_domain domain)
+{
+	struct i915_power_well *power_well;
+
+	power_well = lookup_power_well(dev_priv, SKL_DISP_PW_DC_OFF);
+	if (!power_well || !(power_well->domains & domain))
+		return;
+
+	intel_power_well_put(dev_priv, power_well);
+}
+
 /**
  * intel_display_power_get_if_enabled - grab a reference for an enabled display power domain
  * @dev_priv: i915 device instance
@@ -1498,20 +1628,25 @@ bool intel_display_power_get_if_enabled(struct drm_i915_private *dev_priv,
 					enum intel_display_power_domain domain)
 {
 	struct i915_power_domains *power_domains = &dev_priv->power_domains;
-	bool is_enabled;
+	bool is_enabled = false;
+
 
 	if (!intel_runtime_pm_get_if_in_use(dev_priv))
 		return false;
 
 	mutex_lock(&power_domains->lock);
 
+	if (!dc_off_get_if_enabled(dev_priv, domain))
+		goto out;
+
 	if (__intel_display_power_is_enabled(dev_priv, domain)) {
 		__intel_display_power_get_domain(dev_priv, domain);
 		is_enabled = true;
-	} else {
-		is_enabled = false;
 	}
 
+	dc_off_put(dev_priv, domain);
+
+out:
 	mutex_unlock(&power_domains->lock);
 
 	if (!is_enabled)
@@ -1709,6 +1844,7 @@ void intel_display_power_put(struct drm_i915_private *dev_priv,
 	BIT_ULL(POWER_DOMAIN_GT_IRQ) |			\
 	BIT_ULL(POWER_DOMAIN_MODESET) |			\
 	BIT_ULL(POWER_DOMAIN_AUX_A) |			\
+	BIT_ULL(POWER_DOMAIN_VBLANK) |			\
 	BIT_ULL(POWER_DOMAIN_INIT))
 
 #define BXT_DISPLAY_POWERWELL_2_POWER_DOMAINS (		\
@@ -1732,6 +1868,7 @@ void intel_display_power_put(struct drm_i915_private *dev_priv,
 	BIT_ULL(POWER_DOMAIN_GT_IRQ) |			\
 	BIT_ULL(POWER_DOMAIN_MODESET) |			\
 	BIT_ULL(POWER_DOMAIN_AUX_A) |			\
+	BIT_ULL(POWER_DOMAIN_VBLANK) |			\
 	BIT_ULL(POWER_DOMAIN_INIT))
 #define BXT_DPIO_CMN_A_POWER_DOMAINS (			\
 	BIT_ULL(POWER_DOMAIN_PORT_DDI_A_LANES) |		\
@@ -1791,6 +1928,7 @@ void intel_display_power_put(struct drm_i915_private *dev_priv,
 	BIT_ULL(POWER_DOMAIN_GT_IRQ) |			\
 	BIT_ULL(POWER_DOMAIN_MODESET) |			\
 	BIT_ULL(POWER_DOMAIN_AUX_A) |			\
+	BIT_ULL(POWER_DOMAIN_VBLANK) |			\
 	BIT_ULL(POWER_DOMAIN_INIT))
 
 #define CNL_DISPLAY_POWERWELL_2_POWER_DOMAINS (		\
@@ -1838,6 +1976,7 @@ void intel_display_power_put(struct drm_i915_private *dev_priv,
 	CNL_DISPLAY_POWERWELL_2_POWER_DOMAINS |		\
 	BIT_ULL(POWER_DOMAIN_MODESET) |			\
 	BIT_ULL(POWER_DOMAIN_AUX_A) |			\
+	BIT_ULL(POWER_DOMAIN_VBLANK) |			\
 	BIT_ULL(POWER_DOMAIN_INIT))
 
 static const struct i915_power_well_ops i9xx_always_on_power_well_ops = {
@@ -2071,9 +2210,12 @@ bool intel_display_power_well_is_enabled(struct drm_i915_private *dev_priv,
 {
 	struct i915_power_well *power_well;
 	bool ret;
+	unsigned long flags = 0;
 
 	power_well = lookup_power_well(dev_priv, power_well_id);
+	power_well_lock(power_well, flags);
 	ret = power_well->ops->is_enabled(dev_priv, power_well);
+	power_well_unlock(power_well, flags);
 
 	return ret;
 }
@@ -2108,6 +2250,7 @@ static struct i915_power_well skl_power_wells[] = {
 		.domains = SKL_DISPLAY_DC_OFF_POWER_DOMAINS,
 		.ops = &gen9_dc_off_power_well_ops,
 		.id = SKL_DISP_PW_DC_OFF,
+		.supports_atomic_ctx = true,
 	},
 	{
 		.name = "power well 2",
@@ -2168,6 +2311,7 @@ static struct i915_power_well bxt_power_wells[] = {
 		.domains = BXT_DISPLAY_DC_OFF_POWER_DOMAINS,
 		.ops = &gen9_dc_off_power_well_ops,
 		.id = SKL_DISP_PW_DC_OFF,
+		.supports_atomic_ctx = true,
 	},
 	{
 		.name = "power well 2",
@@ -2223,6 +2367,7 @@ static struct i915_power_well glk_power_wells[] = {
 		.domains = GLK_DISPLAY_DC_OFF_POWER_DOMAINS,
 		.ops = &gen9_dc_off_power_well_ops,
 		.id = SKL_DISP_PW_DC_OFF,
+		.supports_atomic_ctx = true,
 	},
 	{
 		.name = "power well 2",
@@ -2347,6 +2492,7 @@ static struct i915_power_well cnl_power_wells[] = {
 		.domains = CNL_DISPLAY_DC_OFF_POWER_DOMAINS,
 		.ops = &gen9_dc_off_power_well_ops,
 		.id = SKL_DISP_PW_DC_OFF,
+		.supports_atomic_ctx = true,
 	},
 	{
 		.name = "power well 2",
@@ -2475,6 +2621,7 @@ static void assert_power_well_ids_unique(struct drm_i915_private *dev_priv)
 int intel_power_domains_init(struct drm_i915_private *dev_priv)
 {
 	struct i915_power_domains *power_domains = &dev_priv->power_domains;
+	struct i915_power_well *power_well;
 
 	i915_modparams.disable_power_well =
 		sanitize_disable_power_well_option(dev_priv,
@@ -2512,6 +2659,10 @@ int intel_power_domains_init(struct drm_i915_private *dev_priv)
 		set_power_wells(power_domains, i9xx_always_on_power_well);
 	}
 
+	for_each_power_well(dev_priv, power_well)
+		if (power_well->supports_atomic_ctx)
+			spin_lock_init(&power_well->lock);
+
 	assert_power_well_ids_unique(dev_priv);
 
 	return 0;
@@ -2559,9 +2710,14 @@ static void intel_power_domains_sync_hw(struct drm_i915_private *dev_priv)
 
 	mutex_lock(&power_domains->lock);
 	for_each_power_well(dev_priv, power_well) {
+		unsigned long flags = 0;
+
+		power_well_lock(power_well, flags);
+
 		power_well->ops->sync_hw(dev_priv, power_well);
 		power_well->hw_enabled = power_well->ops->is_enabled(dev_priv,
 								     power_well);
+		power_well_unlock(power_well, flags);
 	}
 	mutex_unlock(&power_domains->lock);
 }
@@ -3034,16 +3190,18 @@ void intel_power_domains_suspend(struct drm_i915_private *dev_priv)
 		bxt_display_core_uninit(dev_priv);
 }
 
-static void intel_power_domains_dump_info(struct drm_i915_private *dev_priv)
+static void intel_power_domains_dump_info(struct drm_i915_private *dev_priv,
+					  const int *power_well_use)
 {
 	struct i915_power_domains *power_domains = &dev_priv->power_domains;
 	struct i915_power_well *power_well;
+	int i = 0;
 
 	for_each_power_well(dev_priv, power_well) {
 		enum intel_display_power_domain domain;
 
 		DRM_DEBUG_DRIVER("%-25s %d\n",
-				 power_well->name, power_well->count);
+				 power_well->name, power_well_use[i++]);
 
 		for_each_power_domain(domain, power_well->domains)
 			DRM_DEBUG_DRIVER("  %-23s %d\n",
@@ -3067,6 +3225,7 @@ void intel_power_domains_verify_state(struct drm_i915_private *dev_priv)
 	struct i915_power_domains *power_domains = &dev_priv->power_domains;
 	struct i915_power_well *power_well;
 	bool dump_domain_info;
+	int power_well_use[dev_priv->power_domains.power_well_count];
 
 	mutex_lock(&power_domains->lock);
 
@@ -3075,6 +3234,16 @@ void intel_power_domains_verify_state(struct drm_i915_private *dev_priv)
 		enum intel_display_power_domain domain;
 		int domains_count;
 		bool enabled;
+		int well_count, i = 0;
+		unsigned long flags = 0;
+
+
+		power_well_lock(power_well, flags);
+		well_count = power_well->count;
+		enabled = power_well->ops->is_enabled(dev_priv, power_well);
+		power_well_unlock(power_well, flags);
+
+		power_well_use[i++] = well_count;
 
 		/*
 		 * Power wells not belonging to any domain (like the MISC_IO
@@ -3084,20 +3253,19 @@ void intel_power_domains_verify_state(struct drm_i915_private *dev_priv)
 		if (!power_well->domains)
 			continue;
 
-		enabled = power_well->ops->is_enabled(dev_priv, power_well);
-		if ((power_well->count || power_well->always_on) != enabled)
+
+		if ((well_count || power_well->always_on) != enabled)
 			DRM_ERROR("power well %s state mismatch (refcount %d/enabled %d)",
-				  power_well->name, power_well->count, enabled);
+				  power_well->name, well_count, enabled);
 
 		domains_count = 0;
 		for_each_power_domain(domain, power_well->domains)
 			domains_count += atomic_read(&power_domains->domain_use_count[domain]);
 
-		if (power_well->count != domains_count) {
+		if (well_count != domains_count) {
 			DRM_ERROR("power well %s refcount/domain refcount mismatch "
 				  "(refcount %d/domains refcount %d)\n",
-				  power_well->name, power_well->count,
-				  domains_count);
+				  power_well->name, well_count, domains_count);
 			dump_domain_info = true;
 		}
 	}
@@ -3106,7 +3274,7 @@ void intel_power_domains_verify_state(struct drm_i915_private *dev_priv)
 		static bool dumped;
 
 		if (!dumped) {
-			intel_power_domains_dump_info(dev_priv);
+			intel_power_domains_dump_info(dev_priv, power_well_use);
 			dumped = true;
 		}
 	}
