/external/tensorflow/tensorflow/compiler/xla/service/ |
hlo_reachability.cc | 88 if (input->opcode() == HloOpcode::kAllReduce && input->all_reduce_id()) { 89 auto it = channel_group.find(*input->all_reduce_id()); 120 auto all_reduce_id = hlo->all_reduce_id(); local 121 if (all_reduce_id) { 122 auto it = channel_group.find(all_reduce_id.value());
|
hlo_module_group_util.cc | 65 metadata_.GetAllReduceGroup(*predecessor->all_reduce_id())) { 86 metadata_.GetAllReduceGroup(*instruction->all_reduce_id()); 142 metadata_.GetAllReduceGroup(*successor->all_reduce_id())) { 164 metadata_.GetAllReduceGroup(*instruction->all_reduce_id()); 264 instruction_group = metadata_.GetAllReduceGroup(*hlo->all_reduce_id());
|
ar_crs_combiner.cc | 249 int64 ar_id = *(instruction->all_reduce_id()); 293 auto all_reduce_id = it.first; local 304 all_reduce_map_.erase(all_reduce_id); 326 auto all_reduce_id = all_reduce->all_reduce_id(); local 371 next->set_all_reduce_id(all_reduce_id);
|
hlo_module_group_metadata.cc | 84 for (HloInstruction* instr : GetAllReduceGroup(*hlo->all_reduce_id())) { 235 int64 all_reduce_id) const { 236 auto it = all_reduce_map_.find(all_reduce_id); 318 TF_RET_CHECK(channel_id_map_.find(*hlo->all_reduce_id()) == 320 << "all_reduce_id " << *hlo->all_reduce_id() 322 all_reduce_map_[*hlo->all_reduce_id()].push_back(hlo); 323 max_channel_id_ = std::max(max_channel_id_, *hlo->all_reduce_id());
|
hlo_module_group_metadata.h | 140 // Returns the all-reduce instructions with the same all_reduce_id. 142 int64 all_reduce_id) const; 207 // Returns the maximum channel id or all_reduce_id used in the module group.
|
hlo.proto | 174 int64 all_reduce_id = 45;
|
hlo_computation.cc | 347 return inst->all_reduce_id(); 406 auto all_reduce_id = instruction->all_reduce_id(); local 407 if (all_reduce_id) { 408 channel_dependency_group[all_reduce_id.value()].push_back( [all...] |
hlo_instruction.cc | 379 absl::optional<int64> all_reduce_id; local 380 if (proto.all_reduce_id() > 0) { 381 all_reduce_id = proto.all_reduce_id(); 389 /*all_reduce_id=*/all_reduce_id); [all...] |
hlo_instruction.h | 484 // `all_reduce_id`: for Allreduce nodes from different modules, if they have 485 // the same all_reduce_id, they will be 'Allreduce'd. If empty, Allreduce will 491 absl::string_view barrier, const absl::optional<int64>& all_reduce_id); [all...] |
hlo_instructions.cc | 522 const absl::optional<int64>& all_reduce_id) 526 all_reduce_id_(all_reduce_id) { 531 const absl::optional<int64>& all_reduce_id) { 532 all_reduce_id_ = all_reduce_id; 551 return !all_reduce_id(); 562 result.push_back(StrCat("all_reduce_id=", *all_reduce_id_)); 575 all_reduce_id() == casted_other.all_reduce_id(); 584 all_reduce_id()); [all...] |
layout_assignment.cc | 512 int64 all_reduce_id = instruction->all_reduce_id().value(); 514 ->IsChannelConstrained(all_reduce_id)) { 522 ->LayoutShapeForChannel(buffer_shape, all_reduce_id); [all...] |
hlo_instructions.h | 318 absl::string_view barrier, const absl::optional<int64>& all_reduce_id); 325 absl::optional<int64> all_reduce_id() const { return all_reduce_id_; } function in class:xla::HloAllReduceInstruction 326 void set_all_reduce_id(const absl::optional<int64>& all_reduce_id); 352 // all_reduce_id, they will be 'Allreduce'd. If empty, Allreduce will not be [all...] |
hlo_parser_test.cc | [all...] |
hlo_verifier.cc | [all...] |
hlo_parser.cc | 822 optional<int64> all_reduce_id; local 828 attrs["all_reduce_id"] = {/*required=*/false, AttrTy::kInt64, 829 &all_reduce_id}; 839 all_reduce_id)); [all...] |