diff --git a/drivers/gpu/drm/i915/display/intel_crtc.c b/drivers/gpu/drm/i915/display/intel_crtc.c
index 037fc140b585c..2c8d564e73182 100644
--- a/drivers/gpu/drm/i915/display/intel_crtc.c
+++ b/drivers/gpu/drm/i915/display/intel_crtc.c
@@ -205,6 +205,10 @@ static void intel_crtc_destroy(struct drm_crtc *_crtc)
 	cpu_latency_qos_remove_request(&crtc->vblank_pm_qos);
 
 	drm_crtc_cleanup(&crtc->base);
+
+#if IS_ENABLED(CONFIG_DRM_I915_DEBUG_RUNTIME_PM)
+	drm_WARN_ON(crtc->base.dev, crtc->enabled_power_domains.wakerefs);
+#endif
 	kfree(crtc);
 }
 
diff --git a/drivers/gpu/drm/i915/display/intel_display_power.c b/drivers/gpu/drm/i915/display/intel_display_power.c
index 4c1de91e56ff9..db235b79c9629 100644
--- a/drivers/gpu/drm/i915/display/intel_display_power.c
+++ b/drivers/gpu/drm/i915/display/intel_display_power.c
@@ -830,20 +830,85 @@ void intel_display_power_put_unchecked(struct drm_i915_private *dev_priv,
 }
 #endif
 
+#if IS_ENABLED(CONFIG_DRM_I915_DEBUG_RUNTIME_PM)
+static void
+add_domain_to_set(struct drm_i915_private *i915,
+		  struct intel_display_power_domain_set *power_domain_set,
+		  enum intel_display_power_domain domain,
+		  intel_wakeref_t wf)
+{
+	drm_WARN_ON(&i915->drm, test_bit(domain, power_domain_set->mask.bits));
+
+	if (!power_domain_set->wakerefs)
+		power_domain_set->wakerefs = kcalloc(POWER_DOMAIN_NUM,
+						     sizeof(*power_domain_set->wakerefs),
+						     GFP_KERNEL);
+
+	if (power_domain_set->wakerefs)
+		power_domain_set->wakerefs[domain] = wf;
+
+	set_bit(domain, power_domain_set->mask.bits);
+}
+
+static intel_wakeref_t
+remove_domain_from_set(struct drm_i915_private *i915,
+		       struct intel_display_power_domain_set *power_domain_set,
+		       enum intel_display_power_domain domain)
+{
+	intel_wakeref_t wf;
+
+	drm_WARN_ON(&i915->drm, !test_bit(domain, power_domain_set->mask.bits));
+
+	clear_bit(domain, power_domain_set->mask.bits);
+
+	if (!power_domain_set->wakerefs)
+		return -1;
+
+	wf = fetch_and_zero(&power_domain_set->wakerefs[domain]);
+
+	if (bitmap_empty(power_domain_set->mask.bits, POWER_DOMAIN_NUM)) {
+		kfree(power_domain_set->wakerefs);
+		power_domain_set->wakerefs = NULL;
+	}
+
+	return wf;
+
+}
+#else
+static void
+add_domain_to_set(struct drm_i915_private *i915,
+		  struct intel_display_power_domain_set *power_domain_set,
+		  enum intel_display_power_domain domain,
+		  intel_wakeref_t wf)
+{
+	drm_WARN_ON(&i915->drm, test_bit(domain, power_domain_set->mask.bits));
+
+	set_bit(domain, power_domain_set->mask.bits);
+}
+
+static intel_wakeref_t
+remove_domain_from_set(struct drm_i915_private *i915,
+		       struct intel_display_power_domain_set *power_domain_set,
+		       enum intel_display_power_domain domain)
+{
+	drm_WARN_ON(&i915->drm, !test_bit(domain, power_domain_set->mask.bits));
+
+	clear_bit(domain, power_domain_set->mask.bits);
+
+	return -1;
+}
+#endif
+
 void
 intel_display_power_get_in_set(struct drm_i915_private *i915,
 			       struct intel_display_power_domain_set *power_domain_set,
 			       enum intel_display_power_domain domain)
 {
-	intel_wakeref_t __maybe_unused wf;
-
-	drm_WARN_ON(&i915->drm, test_bit(domain, power_domain_set->mask.bits));
+	intel_wakeref_t wf;
 
 	wf = intel_display_power_get(i915, domain);
-#if IS_ENABLED(CONFIG_DRM_I915_DEBUG_RUNTIME_PM)
-	power_domain_set->wakerefs[domain] = wf;
-#endif
-	set_bit(domain, power_domain_set->mask.bits);
+
+	add_domain_to_set(i915, power_domain_set, domain, wf);
 }
 
 bool
@@ -853,16 +918,11 @@ intel_display_power_get_in_set_if_enabled(struct drm_i915_private *i915,
 {
 	intel_wakeref_t wf;
 
-	drm_WARN_ON(&i915->drm, test_bit(domain, power_domain_set->mask.bits));
-
 	wf = intel_display_power_get_if_enabled(i915, domain);
 	if (!wf)
 		return false;
 
-#if IS_ENABLED(CONFIG_DRM_I915_DEBUG_RUNTIME_PM)
-	power_domain_set->wakerefs[domain] = wf;
-#endif
-	set_bit(domain, power_domain_set->mask.bits);
+	add_domain_to_set(i915, power_domain_set, domain, wf);
 
 	return true;
 }
@@ -874,17 +934,10 @@ intel_display_power_put_mask_in_set(struct drm_i915_private *i915,
 {
 	enum intel_display_power_domain domain;
 
-	drm_WARN_ON(&i915->drm,
-		    !bitmap_subset(mask->bits, power_domain_set->mask.bits, POWER_DOMAIN_NUM));
-
 	for_each_power_domain(domain, mask) {
-		intel_wakeref_t __maybe_unused wf = -1;
+		intel_wakeref_t wf = remove_domain_from_set(i915, power_domain_set, domain);
 
-#if IS_ENABLED(CONFIG_DRM_I915_DEBUG_RUNTIME_PM)
-		wf = fetch_and_zero(&power_domain_set->wakerefs[domain]);
-#endif
 		intel_display_power_put(i915, domain, wf);
-		clear_bit(domain, power_domain_set->mask.bits);
 	}
 }
 
diff --git a/drivers/gpu/drm/i915/display/intel_display_power.h b/drivers/gpu/drm/i915/display/intel_display_power.h
index 1e77e52c87fec..662123d260a7a 100644
--- a/drivers/gpu/drm/i915/display/intel_display_power.h
+++ b/drivers/gpu/drm/i915/display/intel_display_power.h
@@ -147,7 +147,7 @@ struct i915_power_domains {
 struct intel_display_power_domain_set {
 	struct intel_power_domain_mask mask;
 #ifdef CONFIG_DRM_I915_DEBUG_RUNTIME_PM
-	intel_wakeref_t wakerefs[POWER_DOMAIN_NUM];
+	intel_wakeref_t *wakerefs;
 #endif
 };
 
