Method ThunkSchedule

tensorflow/compiler/xla/service/gpu/thunk_schedule.cc:50–83

}

ThunkSchedule::ThunkSchedule(
    std::unique_ptr<ThunkSequence> thunks,
    std::unique_ptr<StreamAssignment> stream_assignment,
    const std::vector<HloInstruction*>& hlo_total_order)
    : thunks_(std::move(thunks)),
      stream_assignment_(std::move(stream_assignment)) {
  absl::flat_hash_map<const HloInstruction*, Thunk*> hlo_to_thunk;
  for (const auto& thunk : *thunks_) {
    InsertOrDie(&hlo_to_thunk, thunk->hlo_instruction(), thunk.get());
  }

  for (HloInstruction* hlo : hlo_total_order) {
    if (Thunk** thunk = tensorflow::gtl::FindOrNull(hlo_to_thunk, hlo)) {
      thunk_total_order_.push_back(*thunk);
    }
  }

  for (const Thunk* thunk : thunk_total_order_) {
    const auto* dst = thunk->hlo_instruction();
    CHECK(stream_assignment_->HasStreamAssigned(*dst));
    for (const auto* src : dst->operands()) {
      AddDependenciesOnTransitiveOperands(*thunk, *src, hlo_to_thunk);
    }
  }

  RemoveRedundantDependencyEdges();

  // Compute `depended_by_`, the inverse of `depends_on_`.
  for (const auto& dependency : depends_on_) {
    for (const auto* depended : dependency.second) {
      depended_by_.insert(depended);
    }
  }
}

void ThunkSchedule::RemoveRedundantDependencyEdges() {
  std::unordered_map<const Thunk*, int> thunk_to_total_order;