message: [XLA] [NFC] Add more information to buffer assignment: print parameter and output shape
diff: @@ -298,6 +298,38 @@ static bool CompareHloValuesById(const HloValue* a, const HloValue* b) {
  return a->id() < b->id();
  }
  +// Returns parameter instruction corresponding to the allocation or nullptr.
  +static const HloInstruction* GetEntryParameterInstruction(
  + const BufferAllocation& alloc) {
  + for (const auto& p : alloc.assigned_buffers()) {
  + const HloValue* value = p.first;
  + const HloInstruction* instr = value->instruction();
  + if (instr->opcode() == HloOpcode::kParameter &&
  + instr->parent() == instr->parent()->parent()->entry_computation()) {
  + return instr;
  + }
  + }
  + return nullptr;
  +}
  +
  +// Returns root module output instruction corresponding to the allocation or
  +// nullptr.
  +static const HloInstruction* GetOutputInstruction(
  + const BufferAllocation& alloc) {
  + for (const auto& p : alloc.assigned_buffers()) {
  + const HloValue* value = p.first;
  + for (const HloPosition& position : value->positions()) {
  + const HloInstruction* instr = position.instruction;
  + if (position.index.empty() &&
  + instr->parent()->root_instruction() == instr &&
  + instr->parent()->IsEntryComputation()) {
  + return instr;
  + }
  + }
  + }
  + return nullptr;
  +}
  +
  string BufferAllocation::ToString() const {
  string output;
  StrAppendFormat(&output, "allocation %d: %p, size %d", index_, this, size());
  @@ -305,8 +337,15 @@ string BufferAllocation::ToString() const {
  StrAppend(&output, ", color ", color().value());
  }
  if (is_entry_computation_parameter()) {
  - StrAppend(&output, ", parameter ", parameter_number(), " at ShapeIndex ",
  - param_shape_index().ToString());
  + const HloInstruction* param = GetEntryParameterInstruction(*this);
  + CHECK(param);
  + StrAppend(&output, ", parameter ", parameter_number(), ", shape |",
  + param->shape().ToString(/*print_layout=*/false),
  + "| at ShapeIndex ", param_shape_index().ToString());
  + }
  + if (const HloInstruction* instr = GetOutputInstruction(*this)) {
  + StrAppend(&output, ", output shape is |",
  + instr->shape().ToString(/*print_layout=*/false), "|");
  }
  if (is_constant()) {
  StrAppend(&output, ", constant");
  