Fix bug in peak buffer accounting in buffer assignment.

Buffer assignment keeps track of the set of logical buffers which are live at
the point of peak memory usage for each allocation. Previously colocated
buffers were not properly accounted for. This CL addresses this problem.

PiperOrigin-RevId: 195001567

1 files changed, 12 insertions, 11 deletions

diff --git a/tensorflow/compiler/xla/service/buffer_assignment.h b/tensorflow/compiler/xla/service/buffer_assignment.h
index 3086d0e2ca..15fd905e8d 100644
--- a/tensorflow/compiler/xla/service/buffer_assignment.h
+++ b/tensorflow/compiler/xla/service/buffer_assignment.h
@@ -206,17 +206,15 @@ class BufferAllocation {
     return heap_traces_;
   }
 
-  // Compute and return the LogicalBuffers which are live at the point of peak
-  // memory usage for the given allocation. The point of peak memory usage is
-  // the point at which the total size of all live logical buffers is
-  // maximal. If peak memory is reached at multiple points, the set of logical
-  // buffers live at the earliest maximal point is returned. The vector is
-  // stabily asserted by LogicalBuffer::Index.
-  //
-  // The return value is a pair of total size of the logical buffers at peak,
-  // and the buffers themselves.
-  std::pair<int64, std::vector<const LogicalBuffer*>>
-  ComputePeakMemoryLogicalBuffers() const;
+  // Returns the LogicalBuffers which are live at the point of peak memory usage
+  // for this allocation. The point of peak memory usage is the point at which
+  // the total size of all live logical buffers is maximal. If peak memory is
+  // reached at multiple points, the set of logical buffers live at the earliest
+  // maximal point is returned. The vector is stabily sorted by
+  // LogicalBuffer::Index.
+  const std::vector<const LogicalBuffer*>& PeakMemoryLogicalBuffers() const {
+    return peak_buffers_;
+  }
 
   // Get the number of bytes lost to fragmentation. This is equal to the
   // difference between the size of the allocation and the size of the maximal
@@ -291,6 +289,9 @@ class BufferAllocation {
 
   int64 fragmentation_bytes_ = 0;
   std::vector<HeapSimulatorTrace> heap_traces_;
+
+  // Set of buffers live at the point of peak memory usage for this allocation.
+  std::vector<const LogicalBuffer*> peak_buffers_;
 };
 
 // Add stream operators for nicer output of CHECK/RET_CHECK failures.