| /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...] |
