Make MSA more deterministic.

tensorflower-gardener · copybara-github · commit 83f1074619a9 · 2024-05-01T15:16:50.000-07:00
PiperOrigin-RevId: 629852083
diff --git 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(

Original file line number	Diff line number	Diff line change
`@@ -3514,10 +3514,17 @@ void MsaAlgorithm::FinalizeAllocations(`
`3514`	`3514`	`}`
`3515`	`3515`	`}`
`3516`	`3516`	`}`
	`3517`	`+ // The colocation_map is a hash table using a pointer as a key. Process its`
	`3518`	`+ // values in some sorted order to get deterministic results.`
	`3519`	`+ std::vector<std::pair<const AliasedOffset, std::vector<Allocation>>>`
	`3520`	`+ sorted_colocations(colocation_map.begin(), colocation_map.end());`
	`3521`	`+ absl::c_sort(sorted_colocations, [](const auto& a, const auto& b) {`
	`3522`	`+ return a.first->offset < b.first->offset;`
	`3523`	`+ });`
`3517`	`3524`	`// The allocations that have the same AliasedOffset need to be colocated.`
`3518`	`3525`	`// Export these to repack_allocation_blocks_ so that we can repack them to`
`3519`	`3526`	`// reduce fragmentation.`
`3520`		`- for (auto& colocation : colocation_map) {`
	`3527`	`+ for (auto& colocation : sorted_colocations) {`
`3521`	`3528`	`std::vector<AllocationBlock*> colocations;`
`3522`	`3529`	`for (Allocation* colocated_allocation : colocation.second) {`
`3523`	`3530`	`repack_allocation_blocks_.push_back(MakeRepackAllocationBlock(`