message: [BufferAssignment] Hoist the calculation for eligibility of buffer set merging.
  This avoids calling corresponding logic O(n^2) times, where n is number of colocated buffer sets, which can be expensive when there are many of colocated buffer sets.
diff: @@ -1434,11 +1434,8 @@ BufferAssigner::MergeColocatedBufferSets(
  computation == module->entry_computation();
  };
  - // Returns true if the two colocated buffer sets (specified by their indices
  - // into the colocated_buffer_sets) can be merged into a single set.
  - auto cannot_merge_buffer_sets = [&colocated_buffer_sets, &buffer_liveness,
  - &buffer_size,
  - &is_entry_parameter](int64 i, int64 j) {
  + std::vector<bool> set_can_be_merged(colocated_buffer_sets.size(), true);
  +
  // Do not merge if one of the sets includes live outs, entry parameters or
  // constants.
  //
  @@ -1453,16 +1450,26 @@ BufferAssigner::MergeColocatedBufferSets(
  // the caller of the executable so we can't write to entry parameters
  // either, and the argument for not merging constants also applies to entry
  // parameters.
  - for (int64 key : {i, j}) {
  - for (auto& buffer : colocated_buffer_sets[key]) {
  + for (int64 i = 0; i < colocated_buffer_sets.size(); ++i) {
  + for (auto& buffer : colocated_buffer_sets[i]) {
  if (buffer_liveness.MaybeLiveOut(*buffer) ||
  is_entry_parameter(*buffer) ||
  buffer->instruction()->opcode() == HloOpcode::kConstant) {
  - return true;
  + set_can_be_merged[i] = false;
  + break;
  }
  }
  }
  + // Returns true if the two colocated buffer sets (specified by their indices
  + // into the colocated_buffer_sets) can be merged into a single set.
  + auto cannot_merge_buffer_sets = [&colocated_buffer_sets, &buffer_liveness,
  + &buffer_size,
  + &set_can_be_merged](int64 i, int64 j) {
  + if (!set_can_be_merged[i] || !set_can_be_merged[j]) {
  + return true;
  + }
  +
  // Colocated sets satisfy the invariant that all buffers within a set have
  // the same size. That means we need to check whether the size is the same
  // between the two sets, but also that it's enough to look at just one
  