From 83f1074619a99ffa2a945c63d95c48d6fafa29d0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 1 May 2024 15:12:05 -0700
Subject: [PATCH] Make MSA more deterministic.

PiperOrigin-RevId: 629852083
---
 xla/service/memory_space_assignment/algorithm.cc | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/xla/service/memory_space_assignment/algorithm.cc b/xla/service/memory_space_assignment/algorithm.cc
index 7d64b8fcca23d..6ce1596af649b 100644
--- a/xla/service/memory_space_assignment/algorithm.cc
+++ b/xla/service/memory_space_assignment/algorithm.cc
@@ -3514,10 +3514,17 @@ void MsaAlgorithm::FinalizeAllocations(
       }
     }
   }
+  // The colocation_map is a hash table using a pointer as a key. Process its
+  // values in some sorted order to get deterministic results.
+  std::vector<std::pair<const AliasedOffset*, std::vector<Allocation*>>>
+      sorted_colocations(colocation_map.begin(), colocation_map.end());
+  absl::c_sort(sorted_colocations, [](const auto& a, const auto& b) {
+    return a.first->offset < b.first->offset;
+  });
   // The allocations that have the same AliasedOffset need to be colocated.
   // Export these to repack_allocation_blocks_ so that we can repack them to
   // reduce fragmentation.
-  for (auto& colocation : colocation_map) {
+  for (auto& colocation : sorted_colocations) {
     std::vector<AllocationBlock*> colocations;
     for (Allocation* colocated_allocation : colocation.second) {
       repack_allocation_blocks_.push_back(MakeRepackAllocationBlock(