Add ToString() methods for SlicedBufferInterval and BufferInterval, a…

…nd use them in VLOGGING.

PiperOrigin-RevId: 524997855

tensorflow/compiler/xla/service/BUILD

@@ -1595,6 +1595,7 @@ cc_library(
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/strings",
     ],
 )
 

tensorflow/compiler/xla/service/heap_simulator.cc

@@ -27,6 +27,8 @@ limitations under the License.
 #include "absl/container/btree_map.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
 #include "tensorflow/compiler/xla/comparison_util.h"
 #include "tensorflow/compiler/xla/hlo/ir/hlo_schedule.h"
 #include "tensorflow/compiler/xla/hlo/utils/hlo_live_range.h"
@@ -803,6 +805,35 @@ std::vector<Chunk> BufferIntervalTree::ChunksOverlappingInTime(
   return result;
 }
 
+template <typename BufferType>
+std::string
+GlobalDecreasingSizeBestFitHeap<BufferType>::BufferInterval::ToString() const {
+  return absl::StrCat("{ ",                                           //
+                      "buffer: ", buffer->ToString(), ", ",           //
+                      "size: ", size, ", ",                           //
+                      "start: ", start, ", ",                         //
+                      "end: ", end, ", ",                             //
+                      "num_colocations: ", colocations.size(), ", ",  //
+                      "need_allocation: ", need_allocation,           //
+                      " }");
+}
+
+template <typename BufferType>
+std::string GlobalDecreasingSizeBestFitHeap<
+    BufferType>::SlicedBufferInterval::ToString() const {
+  return absl::StrCat(
+      "{ full_buffer_interval: ", full_buffer_interval.ToString(),
+      ", sorted_slices: [ ",
+      absl::StrJoin(sorted_slices, ", ",
+                    [](std::string* out,
+                       const SlicedBufferInterval::IntervalSlice& slice) {
+                      absl::StrAppend(out, "{ size: ", slice.size,
+                                      ", allocation_start_time: ",
+                                      slice.allocation_start_time, " }");
+                    }),
+      " ] }");
+}
+
 template <typename BufferType>
 HeapSimulator::Result<BufferType>
 GlobalDecreasingSizeBestFitHeap<BufferType>::Finish() {
@@ -863,10 +894,9 @@ GlobalDecreasingSizeBestFitHeap<BufferType>::FindChunkCandidates(
   CHECK(sliced_buffer_interval.sorted_slices.empty())
       << "Chunk slicing is not yet supported.";
 
-  VLOG(1) << "Finding chunks for buffer: "
-          << buffer_interval.buffer->ToString();
-  VLOG(1) << "Size " << buffer_interval.size << ", start "
-          << buffer_interval.start << ", end " << buffer_interval.end;
+  VLOG(1) << "Finding chunks for sliced buffer interval: "
+          << sliced_buffer_interval.ToString();
+
   // Get all colocated buffers and gather all interferenced chunks.
   //
   // Imagine that we've already allocated three chunks : a, b and c.  And now

tensorflow/compiler/xla/service/heap_simulator.h

@@ -20,6 +20,7 @@ limitations under the License.
 #include <cstdint>
 #include <memory>
 #include <set>
+#include <string>
 #include <utility>
 #include <vector>
 
@@ -376,6 +377,9 @@ class GlobalDecreasingSizeBestFitHeap : public HeapAlgorithm<BufferType> {
 
   // BufferInterval stores a buffer's size and time interval.
   struct BufferInterval {
+    // Convenience method for use with debugging and logging.
+    std::string ToString() const;
+
     const BufferType* buffer;
     int64_t size;
     // Alloc time of the buffer.
@@ -438,6 +442,9 @@ class GlobalDecreasingSizeBestFitHeap : public HeapAlgorithm<BufferType> {
         : full_buffer_interval(buffer_interval) {}
     SlicedBufferInterval() = delete;
 
+    // Convenience method for use with debugging and logging.
+    std::string ToString() const;
+
     const BufferInterval& full_buffer_interval;
 
     // Describes allocations slices, after slice 0.