message: [XLA] Fix async copy statistics to include non-entry computations.
diff: @@ -1847,9 +1847,9 @@ MemorySpaceAssignment::CalculateAsyncCopyStats() const {
  int64 current_copies = 0;
  TF_ASSIGN_OR_RETURN(std::unique_ptr<HloDataflowAnalysis> dataflow_analysis,
  HloDataflowAnalysis::Run(*module_));
  - for (HloInstruction* instruction : module_->schedule()
  - .sequence(module_->entry_computation())
  - .instructions()) {
  + for (const HloComputation* computation :
  + module_->MakeNonfusionComputations()) {
  + for (HloInstruction* instruction : computation->instructions()) {
  if (instruction->opcode() == HloOpcode::kCopyStart) {
  current_copies++;
  } else if (instruction->opcode() == HloOpcode::kCopyDone) {
  @@ -1868,6 +1868,7 @@ MemorySpaceAssignment::CalculateAsyncCopyStats() const {
  stats.max_outstanding_async_copies =
  std::max(stats.max_outstanding_async_copies, current_copies);
  }
  + }
  return stats;
  }
  