From 12fddfb626eab5b6693bbb1ff6e042d50627faad Mon Sep 17 00:00:00 2001
From: Russell Power <power@google.com>
Date: Tue, 12 Sep 2023 18:37:57 -0700
Subject: [PATCH] Use buffer ids to provide an explicitly stable sort.

PiperOrigin-RevId: 564893903
---
 third_party/xla/xla/service/buffer_assignment.cc | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc
index 4e0d16e3dae..ae14bdf5f73 100644
--- a/third_party/xla/xla/service/buffer_assignment.cc
+++ b/third_party/xla/xla/service/buffer_assignment.cc
@@ -1465,7 +1465,7 @@ Status BufferAssigner::AssignBuffersForComputations(
     }
   }
 
-  absl::c_stable_sort(
+  absl::c_sort(
       sorted_buffers, [&post_order_position, &alias_analysis, assignment](
                           const HloBuffer* a, const HloBuffer* b) {
         // Primary sort is by decreasing buffer size.
@@ -1487,7 +1487,16 @@ Status BufferAssigner::AssignBuffersForComputations(
         };
         const HloValue* a_min = *absl::c_min_element(a->values(), compare);
         const HloValue* b_min = *absl::c_min_element(b->values(), compare);
-        return compare(a_min, b_min);
+        if (post_order_position.at(a_min->instruction()) <
+            post_order_position.at(b_min->instruction())) {
+          return true;
+        } else if (post_order_position.at(a_min->instruction()) >
+                   post_order_position.at(b_min->instruction())) {
+          return false;
+        }
+
+        // Use buffer ids to break ties and ensure a stable ordering.
+        return a->id() < b->id();
       });
 
   std::vector<BufferAllocation::Index> allocation_indices;
-- 
GitLab