Speed up HeapSimulator's UniqueOperandSourceBuffers.

Instead of constructing a temporary set, adding all its members to
our big set, then flattening into a vector, add the members to our
vector directly, then sort it and remove duplicates at the end.

PiperOrigin-RevId: 161565289

1 files changed, 15 insertions, 7 deletions

diff --git a/tensorflow/compiler/xla/service/heap_simulator.cc b/tensorflow/compiler/xla/service/heap_simulator.cc
index c662cec9c7..840be603bf 100644
--- a/tensorflow/compiler/xla/service/heap_simulator.cc
+++ b/tensorflow/compiler/xla/service/heap_simulator.cc
@@ -35,18 +35,26 @@ namespace {
 std::vector<const LogicalBuffer*> UniqueOperandSourceBuffers(
     const HloInstruction* instruction,
     const TuplePointsToAnalysis& points_to_analysis) {
-  FlatSet<const LogicalBuffer*> buffers;
+  std::vector<const LogicalBuffer*> buffers;
   for (const HloInstruction* operand : instruction->operands()) {
-    FlatSet<const LogicalBuffer*> sources =
-        points_to_analysis.GetPointsToSet(operand).CreateFlattenedSet();
-    buffers.insert(sources.begin(), sources.end());
+    points_to_analysis.GetPointsToSet(operand).ForEachElement(
+        [&](const ShapeIndex& /*index*/,
+            const std::vector<const LogicalBuffer*>& points_to) {
+          buffers.insert(buffers.end(), points_to.begin(), points_to.end());
+        });
   }
-  std::vector<const LogicalBuffer*> sorted(buffers.begin(), buffers.end());
-  std::sort(sorted.begin(), sorted.end(),
+
+  // Sort and then remove duplicates from buffers.
+  std::sort(buffers.begin(), buffers.end(),
             [](const LogicalBuffer* a, const LogicalBuffer* b) {
               return a->id() < b->id();
             });
-  return sorted;
+  buffers.erase(std::unique(buffers.begin(), buffers.end(),
+                            [](const LogicalBuffer* a, const LogicalBuffer* b) {
+                              return a->id() == b->id();
+                            }),
+                buffers.end());
+  return buffers;
 }
 
 }  // namespace