[PATCH v9 3/3] drm/xe/lrc: Fix ABA race on engine migration in context timestamp read

xe_lrc_context_timestamp() reads the engine id once via
xe_lrc_engine_id(lrc) and uses it to fetch the live CTX_TIMESTAMP MMIO
register, then re-checks the LRC-stored value to detect whether the
context switched out while the MMIO read was in flight. That check
only confirms the context is (still/again) active - it does not
confirm it is active on the *same* engine the MMIO read targeted.

If the context is saved and restored onto a different engine between
the initial engine id read and the final activity check, the
CONTEXT_ACTIVE sentinel will be observed again (now for the new
engine), and the stale MMIO value read from the old, now-unrelated
engine is returned as if it were valid.

Pin the engine id used for the MMIO read and re-validate it against
the current engine id after the final activity check. If the engine
changed, retry the whole read (bounded by a small retry count) instead
of trusting a timestamp sampled from an unrelated context, falling
back to the last cached value if the context keeps migrating.

Assisted-by: GitHub-Copilot:claude-sonnet-5
Signed-off-by: Gajendra Uttamchand <[email protected]>
---
 drivers/gpu/drm/xe/xe_lrc.c | 47 ++++++++++++++++++++-----------------
 1 file changed, 26 insertions(+), 21 deletions(-)

diff --git a/drivers/gpu/drm/xe/xe_lrc.c b/drivers/gpu/drm/xe/xe_lrc.c
index 6f0247bdf8e4..3603642ed0e6 100644
--- a/drivers/gpu/drm/xe/xe_lrc.c
+++ b/drivers/gpu/drm/xe/xe_lrc.c
@@ -2742,39 +2742,44 @@ static u64 xe_lrc_update_multi_queue_timestamp(struct xe_lrc *lrc, u64 *old_ts)
 	return lrc->queue_timestamp;
 }

+/*
+ * Bound the number of times we retry the full read sequence when a
+ * context migration between engines is detected. A small bound (3)
+ * prevents long loops; if we exhaust retries we fall back to the
+ * last cached `lrc->ctx_timestamp`.
+ */
+#define CTX_TIMESTAMP_MAX_RETRIES 3
+
 static u64 xe_lrc_context_timestamp(struct xe_lrc *lrc)
 {
-	u64 reg_ts, new_ts = lrc->ctx_timestamp;
+	u64 reg_ts;
 	u64 stored;
+	u32 engine_id;
+	int retries = CTX_TIMESTAMP_MAX_RETRIES;

 	/* CTX_TIMESTAMP mmio read is invalid on VF, so return the LRC value */
 	if (IS_SRIOV_VF(lrc_to_xe(lrc)))
 		return xe_lrc_ctx_timestamp(lrc);

-	/*
-	 * Safely read CTX_TIMESTAMP: check the LRC-stored value before and
-	 * after the MMIO read to avoid a TOCTOU where a context switch makes the
-	 * MMIO value stale. If the LRC value is not `CONTEXT_ACTIVE` return it;
-	 * otherwise accept the MMIO value only if the context remained active.
-	 */
+	do {
+		/* Read LRC-stored timestamp before/after MMIO to avoid TOCTOU. */
+		stored = xe_lrc_ctx_timestamp(lrc);
+		if (stored != CONTEXT_ACTIVE)
+			return stored;

-	stored = xe_lrc_ctx_timestamp(lrc);
-	if (stored != CONTEXT_ACTIVE)
-		return stored;
+		engine_id = xe_lrc_engine_id(lrc);
+		if (get_ctx_timestamp(lrc, engine_id, &reg_ts))
+			continue;

-	/* Context is active: read the live timestamp from the engine's MMIO register. */
-	if (!get_ctx_timestamp(lrc, xe_lrc_engine_id(lrc), &reg_ts))
-		new_ts = reg_ts;
+		stored = xe_lrc_ctx_timestamp(lrc);
+		if (stored != CONTEXT_ACTIVE)
+			return stored;

-	/* Re-check the LRC-stored timestamp: if the context switched out while
-	 * reading MMIO the hardware saved the canonical timestamp into the LRC
-	 * during context-save, so return that value instead of the MMIO read.
-	 */
-	stored = xe_lrc_ctx_timestamp(lrc);
-	if (stored != CONTEXT_ACTIVE)
-		return stored;
+		if (xe_lrc_engine_id(lrc) == engine_id)
+			return reg_ts;
+	} while (--retries);

-	return new_ts;
+	return lrc->ctx_timestamp;
 }

 static u64 xe_lrc_update_context_timestamp(struct xe_lrc *lrc, u64 *old_ts)
--
2.43.0