message: Adds a utility method to MemoryUsageTracker to compute potential memory savings by rematerializing a sequence of instructions instead of just one.
diff: @@ -347,6 +347,10 @@ class MemoryUsageTracker {
  // if the given instruction is rematerialized.
  int64 MemoryReducedIfRematerialized(Item* item) const;
  + // Returns the number of bytes that the current memory usage will be reduced
  + // by if the given sequence of instructions is rematerialized.
  + int64 MemoryReducedIfRematerialized(const absl::Span<Item*>& items) const;
  +
  // Adjusts memory usage to account for the rematerialization of
  // original_item for all remaining unplaced uses. The rematerialization
  // is remat_item. This method should be called after the HLO graph has
  @@ -653,7 +657,7 @@ int64 MemoryUsageTracker::MemoryReducedIfRematerialized(Item* item) const {
  for (BufferId buffer_id : item->buffers_defined) {
  // Avoid rematerializing instructions with indirect uses as it is difficult
  // to reason about liveness after rematerializing the instruction.
  - // TODO(b/37714814): Consider rematerialzing instructions with indirect
  + // TODO(b/37714814): Consider rematerializing instructions with indirect
  // uses.
  if (buffers_.at(buffer_id).has_indirect_uses) {
  return 0;
  @@ -678,6 +682,60 @@ int64 MemoryUsageTracker::MemoryReducedIfRematerialized(Item* item) const {
  return memory_reduced;
  }
  +int64 MemoryUsageTracker::MemoryReducedIfRematerialized(
  + const absl::Span<Item*>& items) const {
  + CHECK_NE(in_progress_item_, nullptr);
  + int64 memory_reduced = 0;
  + absl::flat_hash_set<Item*> remat_candidates;
  +
  + for (Item* item : items) {
  + if (!item->placed || item == in_progress_item_) {
  + LOG(WARNING) << "Unplaced item or in progress item being checked for "
  + "rematerialization.";
  + return 0;
  + }
  +
  + // Compute the amount of memory reduced (if any) by rematerializing
  + // 'item->instruction'. The LogicalBuffers defined by 'item->instruction'
  + // will no longer be live at this program point, so initially set
  + // memory_reduced to the size of its defined values.
  + for (BufferId buffer_id : item->buffers_defined) {
  + // Avoid rematerializing instructions with indirect uses as it is
  + // difficult to reason about liveness after rematerializing the
  + // instruction.
  + // Avoid rematerializing instructions with live out buffers.
  + // TODO(mpurohit): Check why live_out buffers are an issue here.
  + if (buffers_.at(buffer_id).has_indirect_uses ||
  + buffers_.at(buffer_id).live_out) {
  + return 0;
  + }
  +
  + if (IsCurrentlyLive(buffer_id) && !IsInUse(buffer_id)) {
  + memory_reduced += AllocatedSize(buffer_id);
  + }
  + }
  +
  + // Account for any logical buffers whose live range must be extended across
  + // this program point.
  + for (BufferId buffer_id : item->buffers_used) {
  + if (!IsCurrentlyLive(buffer_id)) {
  + // This logical buffer is used by 'item->instruction' but is not live at
  + // this program point. Rematerializing 'item->instruction' will extend
  + // the buffer's live range across this program point unless it is
  + // defined by an instruction that is also being rematerialized.
  + Item* defining_instruction =
  + buffers_.at(buffer_id).defining_instruction;
  + if (!remat_candidates.contains(defining_instruction)) {
  + memory_reduced -= AllocatedSize(buffer_id);
  + }
  + }
  + }
  + remat_candidates.insert(item);
  + }
  +
  + return memory_reduced;
  +}
  +
  Status MemoryUsageTracker::AddRematerializedInstruction(Item* original_item,
  Item* remat_item) {
  VLOG(3) << "AddRematerializedInstruction: original_instruction = "
  