Home
last modified time | relevance | path

Searched refs:replica_groups (Results 1 – 25 of 29) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/
Dar_crs_combiner_test.cc454 auto replica_groups_before = crs_before->replica_groups(); in TEST_F()
464 auto replica_groups_after = crs_after->replica_groups(); in TEST_F()
504 auto replica_groups_before = crs_before->replica_groups(); in TEST_F()
513 auto replica_groups_after = crs_after->replica_groups(); in TEST_F()
572 auto replica_groups_before = crs_before->replica_groups(); in TEST_F()
582 auto replica_groups_after = crs_after->replica_groups(); in TEST_F()
640 auto replica_groups_before = crs_before->replica_groups(); in TEST_F()
651 auto replica_groups_after = crs_after->replica_groups(); in TEST_F()
682 auto replica_groups_before = crs_before->replica_groups(); in TEST_F()
692 auto replica_groups_after = crs_after->replica_groups(); in TEST_F()
[all …]
Dcollective_ops_utils.cc47 int64 device_ordinal, absl::Span<const ReplicaGroup> replica_groups, in GetParticipatingReplicas() argument
53 if (replica_groups.empty()) { in GetParticipatingReplicas()
65 for (const ReplicaGroup& g : replica_groups) { in GetParticipatingReplicas()
Dbfloat16_normalization_test.cc286 std::vector<ReplicaGroup> replica_groups(1); in TEST_F() local
287 replica_groups[0].add_replica_ids(0); in TEST_F()
288 replica_groups[0].add_replica_ids(1); in TEST_F()
291 replica_groups, absl::nullopt)); in TEST_F()
315 std::vector<ReplicaGroup> replica_groups(1); in TEST_F() local
316 replica_groups[0].add_replica_ids(0); in TEST_F()
317 replica_groups[0].add_replica_ids(1); in TEST_F()
320 replica_groups, absl::nullopt)); in TEST_F()
Dall_reduce_simplifier.cc60 if (all_reduce->replica_groups().empty()) { in Run()
64 for (const auto& group : all_reduce->replica_groups()) { in Run()
Dhlo_replication_analysis.cc66 return hlo->replica_groups().empty() || hlo->replica_groups().size() == 1; in DetermineHloInstructionIsReplicated()
72 return hlo->replica_groups().empty() || hlo->replica_groups().size() == 1; in DetermineHloInstructionIsReplicated()
Dhlo_verifier_test.cc853 string ReplicaGroupsStr(std::vector<std::vector<int64>> replica_groups) { in ReplicaGroupsStr() argument
855 for (const auto& g : replica_groups) { in ReplicaGroupsStr()
863 std::vector<std::vector<int64>> replica_groups) { in MakeAllReduceComputation() argument
876 kTemplate, {{"REPLICA_GROUPS", ReplicaGroupsStr(replica_groups)}})); in MakeAllReduceComputation()
911 std::vector<std::vector<int64>> replica_groups) { in MakeAllToAllComputation() argument
925 kTemplate, {{"REPLICA_GROUPS", ReplicaGroupsStr(replica_groups)}})); in MakeAllToAllComputation()
Dhlo_instructions.cc515 const std::vector<ReplicaGroup>& replica_groups, in HloCollectiveInstruction() argument
518 replica_groups_(replica_groups) { in HloCollectiveInstruction()
536 StrCat("replica_groups=", ReplicaGroupsToString(replica_groups()))); in ExtraAttributesToStringImpl()
547 absl::c_equal(replica_groups(), casted_other.replica_groups(), in IdenticalSlowPath()
556 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout, in HloAllReduceInstruction() argument
559 replica_groups, channel_id), in HloAllReduceInstruction()
565 for (auto replica_group : replica_groups()) { in IsNoop()
604 shape, new_operands, to_apply(), replica_groups(), constrain_layout(), in CloneWithNewOperandsImpl()
610 const std::vector<ReplicaGroup>& replica_groups, in HloAllToAllInstruction() argument
614 replica_groups, channel_id), in HloAllToAllInstruction()
[all …]
Dhlo_parser.cc283 bool ParseReplicaGroupsOnly(std::vector<ReplicaGroup>* replica_groups);
413 std::vector<ReplicaGroup> replica_groups; in CreateReplicaGroups() local
414 absl::c_transform(groups, std::back_inserter(replica_groups), in CreateReplicaGroups()
420 return replica_groups; in CreateReplicaGroups()
865 std::vector<ReplicaGroup> replica_groups; in ParseInstructionRhs() local
867 replica_groups = CreateReplicaGroups(*tmp_groups); in ParseInstructionRhs()
870 shape, operands, *to_apply, replica_groups, in ParseInstructionRhs()
887 std::vector<ReplicaGroup> replica_groups; in ParseInstructionRhs() local
889 replica_groups = CreateReplicaGroups(*tmp_groups); in ParseInstructionRhs()
896 shape, operands, replica_groups, channel_id, split_dimension)); in ParseInstructionRhs()
[all …]
Dhlo_parser_test.cc1410 ROOT crs = f32[8]{0} all-reduce(input), replica_groups={}, to_apply=add
1428 ROOT all-reduce = f32[128,32]{0,1} all-reduce(input), replica_groups={{0,1},{2,3}}, to_apply=add
1446 ROOT crs = f32[8]{0} all-reduce(input), replica_groups={}, constrain_layout=true, to_apply=add
1464 crs.1 = f32[8]{0} all-reduce(input), channel_id=1, replica_groups={{0}}, to_apply=add
1465 ROOT crs.0 = f32[8]{0} all-reduce(input), channel_id=1, replica_groups={{0}}, to_apply=add
1477 ROOT a2a = (f32[128,32]{0,1}) all-to-all(input), replica_groups={}
1490 ROOT a2a = (f32[128,32]{0,1}, f32[128,32]{0,1}) all-to-all(p0, p1), replica_groups={{1,2},{3,0}}
Dhlo_instruction.cc401 std::vector<ReplicaGroup>(proto.replica_groups().begin(), in CreateFromProto()
402 proto.replica_groups().end()), in CreateFromProto()
424 std::vector<ReplicaGroup>(proto.replica_groups().begin(), in CreateFromProto()
425 proto.replica_groups().end()), in CreateFromProto()
913 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout, in CreateAllReduce() argument
916 shape, operands, reduce_computation, replica_groups, constrain_layout, in CreateAllReduce()
922 const std::vector<ReplicaGroup>& replica_groups, in CreateAllToAll() argument
926 shape, operands, replica_groups, channel_id, split_dimension); in CreateAllToAll()
3394 string ReplicaGroupsToString(const std::vector<ReplicaGroup>& replica_groups) { in ReplicaGroupsToString() argument
3396 replica_group_str.reserve(replica_groups.size()); in ReplicaGroupsToString()
[all …]
Dhlo_instruction.h610 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout,
643 const std::vector<ReplicaGroup>& replica_groups,
1705 const std::vector<ReplicaGroup>& replica_groups() const;
2032 string ReplicaGroupsToString(const std::vector<ReplicaGroup>& replica_groups);
Dhlo_verifier.cc192 for (const ReplicaGroup& g : hlo->replica_groups()) { in CheckReplicaGroups()
231 if (hlo->replica_groups().empty()) { in HandleAllToAll()
241 const int64 split_count = hlo->replica_groups().empty() in HandleAllToAll()
243 : hlo->replica_groups()[0].replica_ids_size(); in HandleAllToAll()
244 for (const ReplicaGroup& g : hlo->replica_groups()) { in HandleAllToAll()
Dcollective_ops_utils.h42 int64 device_ordinal, absl::Span<const ReplicaGroup> replica_groups,
Dhlo_instructions.h311 const std::vector<ReplicaGroup>& replica_groups() const { in replica_groups() function
319 const std::vector<ReplicaGroup>& replica_groups,
339 const std::vector<ReplicaGroup>& replica_groups, bool constrain_layout,
384 const std::vector<ReplicaGroup>& replica_groups,
Dar_crs_combiner.cc76 instruction->replica_groups().size() == num_replicas_ && in MatchesArCrsPattern()
/external/tensorflow/tensorflow/compiler/xla/python/
Dxla_client.py1008 def AllReduce(self, operand, computation, replica_groups=None): argument
1022 replica_groups_protos = _get_replica_groups_protos(replica_groups)
1030 replica_groups=None): argument
1046 replica_groups_protos = _get_replica_groups_protos(replica_groups)
1047 if not replica_groups:
1050 split_count = len(replica_groups[0])
1051 if not all(split_count == len(g) for g in replica_groups):
1056 def CrossReplicaSum(self, operand, replica_groups=None): argument
1069 replica_groups_protos = _get_replica_groups_protos(replica_groups)
1901 def _get_replica_groups_protos(replica_groups): argument
[all …]
/external/tensorflow/tensorflow/compiler/xla/tests/
Dcollective_ops_test.cc44 const Shape& shape, std::vector<std::vector<int64>> replica_groups, in MakeCrsModule() argument
65 for (const auto& g : replica_groups) { in MakeCrsModule()
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder.h512 absl::Span<const ReplicaGroup> replica_groups = {});
516 absl::Span<const ReplicaGroup> replica_groups = {},
522 const std::vector<ReplicaGroup>& replica_groups);
923 absl::Span<const ReplicaGroup> replica_groups);
925 absl::Span<const ReplicaGroup> replica_groups,
930 const std::vector<ReplicaGroup>& replica_groups);
1654 absl::Span<const ReplicaGroup> replica_groups = {});
1676 absl::Span<const ReplicaGroup> replica_groups = {},
1683 const std::vector<ReplicaGroup>& replica_groups = {});
Dxla_builder.cc2084 XlaOp operand, absl::Span<const ReplicaGroup> replica_groups) { in CrossReplicaSum() argument
2108 return AllReduce(operand, computation, replica_groups, in CrossReplicaSum()
2114 absl::Span<const ReplicaGroup> replica_groups, in AllReduce() argument
2166 for (const ReplicaGroup& group : replica_groups) { in AllReduce()
2191 const std::vector<ReplicaGroup>& replica_groups) { in AllToAll() argument
2227 for (const ReplicaGroup& group : replica_groups) { in AllToAll()
3302 absl::Span<const ReplicaGroup> replica_groups) { in CrossReplicaSum() argument
3303 return operand.builder()->CrossReplicaSum(operand, replica_groups); in CrossReplicaSum()
3307 absl::Span<const ReplicaGroup> replica_groups, in AllReduce() argument
3310 return operand.builder()->AllReduce(operand, computation, replica_groups, in AllReduce()
[all …]
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dmlir_hlo_to_hlo.cc487 auto replica_groups = Convert_replica_groups(op.replica_groups()); in ExportXlaOp() local
490 xla::AllReduce(value_map[op.operand()], computation, replica_groups, in ExportXlaOp()
496 replica_groups, channel_id); in ExportXlaOp()
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/translate/
Dexport.mlir25 replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : tensor<2x4xi64>,
39 // CHECK-SAME: replica_groups={{[{][{]}}0,2,4,6},{1,3,5,7{{[}][}]}}
346 …%1 = "xla_hlo.cross-replica-sum"(%arg0) {replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : te…
356 // CHECK-SAME: replica_groups={{[{][{]}}0,2,4,6},{1,3,5,7{{[}][}]}}
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dnccl_all_reduce_thunk.cc497 GetParticipatingReplicas(device_ordinal, instr->replica_groups(), in ExecuteOnStream()
/external/tensorflow/tensorflow/compiler/mlir/xla/ir/
Dhlo_ops.td527 I64ElementsAttr:$replica_groups,
544 I64ElementsAttr:$replica_groups
861 I64ElementsAttr:$replica_groups
Dhlo_ops_base.td685 - replica_groups: each ReplicaGroup contains a list of replica id
689 `replica_groups` = {{1,2,3}, {4,5,0}} means that an AllToAll will be applied
693 also 4, 5, 0. If `replica_groups` is empty, all replicas belong to one
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/
Dops.mlir22 replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
35 replica_groups = dense<[[0, 1, 2, 3, 4]]> : tensor<1x5xi64>
48 replica_groups = dense<[[0, 1, 2, 3, 4]]> : tensor<1x5xi64>

12