diff --git a/drivers/gpu/drm/i915/display/intel_de.c b/drivers/gpu/drm/i915/display/intel_de.c
index fce92535bd6a..d2a418da2d54 100644
--- a/drivers/gpu/drm/i915/display/intel_de.c
+++ b/drivers/gpu/drm/i915/display/intel_de.c
@@ -3,10 +3,85 @@
  * Copyright © 2026 Intel Corporation
  */
 
+#include <linux/delay.h>
+
 #include <drm/drm_print.h>
 
 #include "intel_de.h"
 
+static int __intel_de_wait_for_register(struct intel_display *display,
+					i915_reg_t reg, u32 mask, u32 value,
+					unsigned int timeout_us,
+					u32 (*read)(struct intel_display *display, i915_reg_t reg),
+					u32 *out_val, bool is_atomic)
+{
+	const ktime_t end = ktime_add_us(ktime_get_raw(), timeout_us);
+	int wait_max = 1000;
+	int wait = 10;
+	u32 reg_value;
+	int ret;
+
+	might_sleep_if(!is_atomic);
+
+	if (timeout_us <= 10) {
+		is_atomic = true;
+		wait = 1;
+	}
+
+	for (;;) {
+		bool expired = ktime_after(ktime_get_raw(), end);
+
+		/* guarantee the condition is evaluated after timeout expired */
+		barrier();
+
+		reg_value = read(display, reg);
+		if ((reg_value & mask) == value) {
+			ret = 0;
+			break;
+		}
+
+		if (expired) {
+			ret = -ETIMEDOUT;
+			break;
+		}
+
+		if (is_atomic)
+			udelay(wait);
+		else
+			usleep_range(wait, wait << 1);
+
+		if (wait < wait_max)
+			wait <<= 1;
+	}
+
+	if (out_val)
+		*out_val = reg_value;
+
+	return ret;
+}
+
+static int intel_de_wait_for_register(struct intel_display *display,
+				      i915_reg_t reg, u32 mask, u32 value,
+				      unsigned int fast_timeout_us,
+				      unsigned int slow_timeout_us,
+				      u32 (*read)(struct intel_display *display, i915_reg_t reg),
+				      u32 *out_value, bool is_atomic)
+{
+	int ret = -EINVAL;
+
+	if (fast_timeout_us)
+		ret = __intel_de_wait_for_register(display, reg, mask, value,
+						   fast_timeout_us, read,
+						   out_value, is_atomic);
+
+	if (ret && slow_timeout_us)
+		ret = __intel_de_wait_for_register(display, reg, mask, value,
+						   slow_timeout_us, read,
+						   out_value, is_atomic);
+
+	return ret;
+}
+
 int intel_de_wait_us(struct intel_display *display, i915_reg_t reg,
 		     u32 mask, u32 value, unsigned int timeout_us,
 		     u32 *out_value)
@@ -15,8 +90,10 @@ int intel_de_wait_us(struct intel_display *display, i915_reg_t reg,
 
 	intel_dmc_wl_get(display, reg);
 
-	ret = __intel_wait_for_register(__to_uncore(display), reg, mask,
-					value, timeout_us, 0, out_value);
+	ret = intel_de_wait_for_register(display, reg, mask, value,
+					 timeout_us, 0,
+					 intel_de_read,
+					 out_value, false);
 
 	intel_dmc_wl_put(display, reg);
 
@@ -31,8 +108,10 @@ int intel_de_wait_ms(struct intel_display *display, i915_reg_t reg,
 
 	intel_dmc_wl_get(display, reg);
 
-	ret = __intel_wait_for_register(__to_uncore(display), reg, mask,
-					value, 2, timeout_ms, out_value);
+	ret = intel_de_wait_for_register(display, reg, mask, value,
+					 2, timeout_ms * 1000,
+					 intel_de_read,
+					 out_value, false);
 
 	intel_dmc_wl_put(display, reg);
 
@@ -43,16 +122,20 @@ int intel_de_wait_fw_ms(struct intel_display *display, i915_reg_t reg,
 			u32 mask, u32 value, unsigned int timeout_ms,
 			u32 *out_value)
 {
-	return __intel_wait_for_register_fw(__to_uncore(display), reg, mask,
-					    value, 2, timeout_ms, out_value);
+	return intel_de_wait_for_register(display, reg, mask, value,
+					  2, timeout_ms * 1000,
+					  intel_de_read_fw,
+					  out_value, false);
 }
 
 int intel_de_wait_fw_us_atomic(struct intel_display *display, i915_reg_t reg,
 			       u32 mask, u32 value, unsigned int timeout_us,
 			       u32 *out_value)
 {
-	return __intel_wait_for_register_fw(__to_uncore(display), reg, mask,
-					    value, timeout_us, 0, out_value);
+	return intel_de_wait_for_register(display, reg, mask, value,
+					  timeout_us, 0,
+					  intel_de_read_fw,
+					  out_value, true);
 }
 
 int intel_de_wait_for_set_us(struct intel_display *display, i915_reg_t reg,
diff --git a/drivers/gpu/drm/xe/compat-i915-headers/intel_uncore.h b/drivers/gpu/drm/xe/compat-i915-headers/intel_uncore.h
index a8cfd65119e0..08d7ab933672 100644
--- a/drivers/gpu/drm/xe/compat-i915-headers/intel_uncore.h
+++ b/drivers/gpu/drm/xe/compat-i915-headers/intel_uncore.h
@@ -98,37 +98,6 @@ static inline u32 intel_uncore_rmw(struct intel_uncore *uncore,
 	return xe_mmio_rmw32(__compat_uncore_to_mmio(uncore), reg, clear, set);
 }
 
-static inline int
-__intel_wait_for_register(struct intel_uncore *uncore, i915_reg_t i915_reg,
-			  u32 mask, u32 value, unsigned int fast_timeout_us,
-			  unsigned int slow_timeout_ms, u32 *out_value)
-{
-	struct xe_reg reg = XE_REG(i915_mmio_reg_offset(i915_reg));
-	bool atomic;
-
-	/*
-	 * Replicate the behavior from i915 here, in which sleep is not
-	 * performed if slow_timeout_ms == 0. This is necessary because
-	 * of some paths in display code where waits are done in atomic
-	 * context.
-	 */
-	atomic = !slow_timeout_ms && fast_timeout_us > 0;
-
-	return xe_mmio_wait32(__compat_uncore_to_mmio(uncore), reg, mask, value,
-			      fast_timeout_us + 1000 * slow_timeout_ms,
-			      out_value, atomic);
-}
-
-static inline int
-__intel_wait_for_register_fw(struct intel_uncore *uncore, i915_reg_t i915_reg,
-			     u32 mask, u32 value, unsigned int fast_timeout_us,
-			     unsigned int slow_timeout_ms, u32 *out_value)
-{
-	return __intel_wait_for_register(uncore, i915_reg, mask, value,
-					 fast_timeout_us, slow_timeout_ms,
-					 out_value);
-}
-
 static inline u32 intel_uncore_read_fw(struct intel_uncore *uncore,
 				       i915_reg_t i915_reg)
 {