Make window prefetching in memory space assignment more deterministic.

PiperOrigin-RevId: 872073012

Author: mehrdadkhani

xla/service/memory_space_assignment/BUILD

@@ -618,6 +618,8 @@ cc_library(
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/container:linked_hash_map",
+        "@com_google_absl//absl/container:linked_hash_set",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/hash",
         "@com_google_absl//absl/log",

xla/service/memory_space_assignment/algorithm.cc

@@ -39,6 +39,7 @@ limitations under the License.
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/container/inlined_vector.h"
+#include "absl/container/linked_hash_map.h"
 #include "absl/functional/any_invocable.h"
 #include "absl/hash/hash.h"
 #include "absl/log/check.h"
@@ -7203,8 +7204,10 @@ absl::Status MsaAlgorithm::WindowPrefetch() {
   // cloned computation and use the cloned computation to determine the operand
   // span size.
 
-  // Map of the original instruction to a clone of the instruction.
-  absl::flat_hash_map<HloInstruction*, HloInstruction*> cloned_insts;
+  // Map of the original instruction to a clone of the instruction. Use a
+  // linked_hash_map to ensure deterministic traversal for memory space
+  // propagation and cleanup.
+  absl::linked_hash_map<HloInstruction*, HloInstruction*> cloned_insts;
   const std::vector<HloInstruction*>& instruction_sequence =
       hlo_live_range_.flattened_instruction_sequence().instructions();
   for (HloInstruction* instruction : instruction_sequence) {

xla/service/memory_space_assignment/algorithm.h

@@ -29,6 +29,8 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/container/linked_hash_set.h"
+
 // TODO(b/210891274): Use btree_map after build issue in Windows is resolved.
 #if defined(__GNUC__) || defined(__clang__)
 #include "absl/container/btree_map.h"
@@ -1362,8 +1364,9 @@ class MsaAlgorithm : public GlobalDecreasingSizeBestFitHeap<HloValue> {
   absl::flat_hash_map<HloPosition, std::vector<int64_t>>
       default_memory_coloring_requirements_;
 
-  // Set of HloUses that are in the default memory.
-  absl::flat_hash_set<HloUse> uses_in_default_memory_;
+  // Set of HloUses that are in the default memory. Using linked_hash_set for
+  // deterministic window prefetching results.
+  absl::linked_hash_set<HloUse> uses_in_default_memory_;
 };
 
 }  // namespace memory_space_assignment