Method AddAssignment

tensorflow/compiler/xla/service/buffer_assignment.cc:235–264

}

void BufferAllocation::AddAssignment(const HloValue& buffer, int64 offset,
                                     int64 size) {
  VLOG(4) << "Adding the following buffer to allocation #" << index()
          << absl::StrFormat(" (size=%d, offset=%d) %s", size, offset,
                             buffer.ToShortString());
  CHECK(!assigned_buffers_.contains(&buffer))
      << "LogicalBuffer " << buffer << " already assigned to allocation "
      << index_;
  CHECK_LE(offset, size_) << "LogicalBuffer " << buffer
                          << " offset out of range";
  CHECK_LE(offset + size, size_)
      << "LogicalBuffer " << buffer
      << " size out of range at offset: " << offset << " with size: " << size;
  CHECK_EQ(buffer.color(), color())
      << "Buffer color " << buffer.color() << " for buffer " << buffer
      << " does not match allocation color " << color() << ".";
  OffsetSize offset_size;
  offset_size.offset = offset;
  offset_size.size = size;
  assigned_buffers_.emplace(&buffer, offset_size);
  // For debugging purposes, store the assigned memory space in the
  // instruction's layout.
  for (HloPosition position : buffer.positions()) {
    Shape* shape = ShapeUtil::GetMutableSubshape(
        position.instruction->mutable_shape(), position.index);
    if (shape->has_layout()) {
      shape->mutable_layout()->set_memory_space(buffer.color().value());
    }
  }
}

BufferAllocationProto BufferAllocation::ToProto() const {
  BufferAllocationProto proto;