Searched refs:all_reduce_map_ (Results 1 – 4 of 4) sorted by relevance
259 CHECK(all_reduce_map_.find(ar_id) == all_reduce_map_.end()); in GroupAllReducesById()261 auto prev_pair = all_reduce_map_[prev_ar_id].back(); in GroupAllReducesById()266 all_reduce_map_.erase(prev_ar_id); in GroupAllReducesById()268 all_reduce_map_[ar_id].push_back(pair); in GroupAllReducesById()276 if (all_reduce_map_.find(ar_id) != all_reduce_map_.end()) { in GroupAllReducesById()277 int64 prev_distance = all_reduce_map_[ar_id].back().distance; in GroupAllReducesById()283 all_reduce_map_[ar_id].push_back(pair); in GroupAllReducesById()292 for (auto it : all_reduce_map_) { in KeepProvablyEqualInstructionGroups()304 all_reduce_map_.erase(all_reduce_id); in KeepProvablyEqualInstructionGroups()318 if (all_reduce_map_.empty()) { in RewriteGraph()[all …]
236 auto it = all_reduce_map_.find(all_reduce_id); in GetAllReduceGroup()237 CHECK(it != all_reduce_map_.end()); in GetAllReduceGroup()322 all_reduce_map_[*hlo->all_reduce_id()].push_back(hlo); in RecordInstructions()331 TF_RET_CHECK(all_reduce_map_.find(hlo->channel_id()) == in RecordInstructions()332 all_reduce_map_.end()) in RecordInstructions()378 VLOG(2) << "Created " << all_reduce_map_.size() << " all-reduce groups"; in RecordInstructions()
142 absl::flat_hash_map<int64, std::vector<ArCrsPair>> all_reduce_map_; variable
277 absl::flat_hash_map<int64, std::vector<HloInstruction*>> all_reduce_map_; variable