From 7507dadc02810e8b95341fd7bb42b9244f5955f8 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Sun, 16 Nov 2025 15:16:18 -0800 Subject: [PATCH] Properly tiebreak ForceDelay PiperOrigin-RevId: 833065628 --- third_party/xla/xla/service/latency_hiding_scheduler.cc | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/third_party/xla/xla/service/latency_hiding_scheduler.cc b/third_party/xla/xla/service/latency_hiding_scheduler.cc index 4841d9b90f3..3ca8f126f97 100644 --- a/third_party/xla/xla/service/latency_hiding_scheduler.cc +++ b/third_party/xla/xla/service/latency_hiding_scheduler.cc @@ -1278,10 +1278,13 @@ class ReadySetLt { HloGraphNode* bn = b.node; // Schedule according to ForceEarly. CMP_PROPERTY(GetForceEarly(), "kForceEarly"); - // Schedule according to highest ForceDelay first. + // Schedule according to ForceDelay, if exactly one of the two instructions + // has ForceDelay set. + CMP_EXPLICIT(!an->GetForceDelay(), !bn->GetForceDelay(), "kForceDelay"); + // Schedule according to highest ForceDelay first, if both instructions + // have ForceDelay set. CMP_EXPLICIT(-an->GetForceDelayPriority(), -bn->GetForceDelayPriority(), "kForceDelayPriority"); - CMP_EXPLICIT(!an->GetForceDelay(), !bn->GetForceDelay(), "kForceDelay"); // Use the preference value (comes from a heuristic) to choose between // the two candidates. If two preferences are the same regular LHS logic // will run as usual, we take advantage of this fact when initializing