Searched refs:hlo_to_thunk (Results 1 – 2 of 2) sorted by relevance
31 const absl::flat_hash_map<const HloInstruction*, Thunk*>& hlo_to_thunk) { in AddDependenciesOnTransitiveOperands() argument32 if (hlo_to_thunk.contains(&operand)) { in AddDependenciesOnTransitiveOperands()38 depends_on_[&thunk].push_back(FindOrDie(hlo_to_thunk, &operand)); in AddDependenciesOnTransitiveOperands()45 hlo_to_thunk); in AddDependenciesOnTransitiveOperands()56 absl::flat_hash_map<const HloInstruction*, Thunk*> hlo_to_thunk; in ThunkSchedule() local58 InsertOrDie(&hlo_to_thunk, thunk->hlo_instruction(), thunk.get()); in ThunkSchedule()62 if (Thunk** thunk = tensorflow::gtl::FindOrNull(hlo_to_thunk, hlo)) { in ThunkSchedule()71 AddDependenciesOnTransitiveOperands(*thunk, *src, hlo_to_thunk); in ThunkSchedule()
82 const absl::flat_hash_map<const HloInstruction*, Thunk*>& hlo_to_thunk);