Lines Matching refs:thunk
25 const Thunk& thunk, const HloInstruction& operand,
26 const std::unordered_map<const HloInstruction*, Thunk*>& hlo_to_thunk) {
28 // If `operand` is mapped to a thunk, adds `operand` to `thunk`'s dependency
32 stream_assignment_->StreamNumberForHlo(*thunk.hlo_instruction())) {
33 depends_on_[&thunk].push_back(FindOrDie(hlo_to_thunk, &operand));
36 // If `operand` doesn't need a thunk (e.g. bitcast), continue with its
39 AddDependenciesOnTransitiveOperands(thunk, *operand_of_operand,
51 std::unordered_map<const HloInstruction*, Thunk*> hlo_to_thunk;
52 for (const auto& thunk : *thunks_) {
53 InsertOrDie(&hlo_to_thunk, thunk->hlo_instruction(), thunk.get());
62 for (const Thunk* thunk : thunk_total_order_) {
63 const auto* dst = thunk->hlo_instruction();
66 AddDependenciesOnTransitiveOperands(*thunk, *src, hlo_to_thunk);
81 std::unordered_map<const Thunk*, int> thunk_to_total_order;
97 // Suppose thunk T1 and T3 are scheduled on stream S1, and T2 and T4 are on
103 // last_dependency[S1][S2] indicates the last thunk (with the maximum order
105 // S1 thunk depends on a S2 thunk ordered <=last_dependency[S1][S2], that is a
108 for (const Thunk* dst : thunk_total_order_) {
115 std::list<const Thunk*>& sources = FindOrDie(depends_on_, dst);
117 const Thunk* src = *iter;
135 const std::list<const Thunk*>& ThunkSchedule::DependsOn(
136 const Thunk* thunk) const {
137 if (depends_on_.count(thunk)) {
138 return FindOrDie(depends_on_, thunk);
146 for (Thunk* thunk : thunk_total_order_) {
148 thunk->hlo_instruction()->ToString(), "\n");
152 const Thunk* dependent = entry.first;
153 for (const Thunk* dependency : entry.second) {
