Home
last modified time | relevance | path

Searched refs:all_reduce_id (Results 1 – 15 of 15) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_reachability.cc88 if (input->opcode() == HloOpcode::kAllReduce && input->all_reduce_id()) { in Build()
89 auto it = channel_group.find(*input->all_reduce_id()); in Build()
120 auto all_reduce_id = hlo->all_reduce_id(); in Build() local
121 if (all_reduce_id) { in Build()
122 auto it = channel_group.find(all_reduce_id.value()); in Build()
Dar_crs_combiner.cc249 int64 ar_id = *(instruction->all_reduce_id()); in GroupAllReducesById()
293 auto all_reduce_id = it.first; in KeepProvablyEqualInstructionGroups() local
304 all_reduce_map_.erase(all_reduce_id); in KeepProvablyEqualInstructionGroups()
326 auto all_reduce_id = all_reduce->all_reduce_id(); in RewriteGraph() local
371 next->set_all_reduce_id(all_reduce_id); in RewriteGraph()
Dhlo_module_group_util.cc65 metadata_.GetAllReduceGroup(*predecessor->all_reduce_id())) { in GlobalPredecessors()
86 metadata_.GetAllReduceGroup(*instruction->all_reduce_id()); in GlobalPredecessors()
142 metadata_.GetAllReduceGroup(*successor->all_reduce_id())) { in GlobalSuccessors()
164 metadata_.GetAllReduceGroup(*instruction->all_reduce_id()); in GlobalSuccessors()
264 instruction_group = metadata_.GetAllReduceGroup(*hlo->all_reduce_id()); in VisitTopologicalOrder()
Dhlo_module_group_metadata.cc84 for (HloInstruction* instr : GetAllReduceGroup(*hlo->all_reduce_id())) { in Build()
235 int64 all_reduce_id) const { in GetAllReduceGroup()
236 auto it = all_reduce_map_.find(all_reduce_id); in GetAllReduceGroup()
318 TF_RET_CHECK(channel_id_map_.find(*hlo->all_reduce_id()) == in RecordInstructions()
320 << "all_reduce_id " << *hlo->all_reduce_id() in RecordInstructions()
322 all_reduce_map_[*hlo->all_reduce_id()].push_back(hlo); in RecordInstructions()
323 max_channel_id_ = std::max(max_channel_id_, *hlo->all_reduce_id()); in RecordInstructions()
Dhlo_computation.cc347 return inst->all_reduce_id(); in ComputeInstructionPostOrder()
406 auto all_reduce_id = instruction->all_reduce_id(); in ComputeChannelDependencies() local
407 if (all_reduce_id) { in ComputeChannelDependencies()
408 channel_dependency_group[all_reduce_id.value()].push_back( in ComputeChannelDependencies()
Dhlo_instruction.cc379 absl::optional<int64> all_reduce_id; in CreateFromProto() local
380 if (proto.all_reduce_id() > 0) { in CreateFromProto()
381 all_reduce_id = proto.all_reduce_id(); in CreateFromProto()
389 /*all_reduce_id=*/all_reduce_id); in CreateFromProto()
836 const absl::optional<int64>& all_reduce_id) { in CreateAllReduce() argument
839 all_reduce_id); in CreateAllReduce()
1249 return all_reduce_id().has_value(); in HasSideEffectNoRecurse()
2183 return opcode() == HloOpcode::kAllReduce && all_reduce_id(); in IsCrossModuleAllReduce()
2187 return opcode() == HloOpcode::kAllReduce && !all_reduce_id(); in IsCrossReplicaAllReduce()
3429 absl::optional<int64> HloInstruction::all_reduce_id() const { in all_reduce_id() function in xla::HloInstruction
[all …]
Dhlo_module_group_metadata.h142 int64 all_reduce_id) const;
Dhlo_instructions.cc522 const absl::optional<int64>& all_reduce_id) in HloAllReduceInstruction() argument
526 all_reduce_id_(all_reduce_id) { in HloAllReduceInstruction()
531 const absl::optional<int64>& all_reduce_id) { in set_all_reduce_id() argument
532 all_reduce_id_ = all_reduce_id; in set_all_reduce_id()
551 return !all_reduce_id(); in IsNoop()
575 all_reduce_id() == casted_other.all_reduce_id(); in IdenticalSlowPath()
584 all_reduce_id()); in CloneWithNewOperandsImpl()
Dhlo_instruction.h491 absl::string_view barrier, const absl::optional<int64>& all_reduce_id);
1532 absl::optional<int64> all_reduce_id() const;
1533 void set_all_reduce_id(const absl::optional<int64>& all_reduce_id);
Dlayout_assignment.cc512 int64 all_reduce_id = instruction->all_reduce_id().value(); in AddMandatoryConstraints() local
514 ->IsChannelConstrained(all_reduce_id)) { in AddMandatoryConstraints()
522 ->LayoutShapeForChannel(buffer_shape, all_reduce_id); in AddMandatoryConstraints()
1863 ->ConstrainChannel(instruction->all_reduce_id().value(), in ConstrainChannelLayouts()
Dhlo_instructions.h318 absl::string_view barrier, const absl::optional<int64>& all_reduce_id);
325 absl::optional<int64> all_reduce_id() const { return all_reduce_id_; } in all_reduce_id() function
326 void set_all_reduce_id(const absl::optional<int64>& all_reduce_id);
Dhlo_verifier.cc1416 if (crs->all_reduce_id().has_value()) { in HandleAllReduce()
1417 TF_RET_CHECK(crs->all_reduce_id().value() > 0) in HandleAllReduce()
Dhlo.proto174 int64 all_reduce_id = 45; field
Dhlo_parser_test.cc1387 crs.1 = f32[8]{0} all-reduce(input), replica_groups={{0}}, all_reduce_id=1, to_apply=add in CreateTestCases()
1388 ROOT crs.0 = f32[8]{0} all-reduce(input), replica_groups={{0}}, all_reduce_id=1, to_apply=add in CreateTestCases()
Dhlo_parser.cc822 optional<int64> all_reduce_id; in ParseInstructionRhs() local
829 &all_reduce_id}; in ParseInstructionRhs()
839 all_reduce_id)); in ParseInstructionRhs()