Home
last modified time | relevance | path

Searched refs:constrain_layout (Results 1 – 25 of 28) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_dce.cc55 maybe_collective_op->constrain_layout())) { in IsRemovableWhile()
78 !maybe_collective_op->constrain_layout()) || in RunOnComputation()
Dasync_collective_creator.cc67 ar->replica_groups(), ar->constrain_layout(), ar->channel_id(), in Run()
97 ag->replica_groups(), ag->constrain_layout(), ag->channel_id(), in Run()
Dreduce_scatter_decomposer.cc61 rs->replica_groups(), rs->constrain_layout(), channel_id, in Run()
Dall_gather_decomposer.cc76 /*constrain_layout=*/ag->constrain_layout(), ag->channel_id(), in DecomposeAllGather()
Dhlo_instructions.cc740 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout, in HloCollectiveInstruction() argument
744 constrain_layout_(constrain_layout) { in HloCollectiveInstruction()
778 constrain_layout() == casted_other.constrain_layout() && in IdenticalSlowPathIgnoringChannelIdValues()
788 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout, in HloAllGatherInstruction() argument
791 constrain_layout, channel_id), in HloAllGatherInstruction()
812 constrain_layout(), channel_id(), use_global_device_ids()); in CloneWithNewOperandsImpl()
837 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout, in HloAllReduceInstructionBase() argument
840 constrain_layout, channel_id), in HloAllReduceInstructionBase()
873 constrain_layout() == casted_other.constrain_layout() && in IdenticalSlowPathIgnoringChannelIdValues()
893 constrain_layout(), channel_id(), use_global_device_ids()); in CloneWithNewOperandsImpl()
[all …]
Dhlo_instruction.cc484 proto.constrain_layout(), channel_id, in CreateFromProto()
491 proto.constrain_layout(), channel_id, in CreateFromProto()
516 replica_groups, proto.constrain_layout(), in CreateFromProto()
524 proto.constrain_layout(), channel_id, proto.use_global_device_ids(), in CreateFromProto()
529 replica_groups, proto.constrain_layout(), in CreateFromProto()
553 /*constrain_layout=*/proto.constrain_layout(), in CreateFromProto()
719 if (proto.constrain_layout()) { in CreateFromProto()
1260 bool constrain_layout, const std::optional<int64_t>& channel_id, in CreateAllGather() argument
1264 replica_groups, constrain_layout, channel_id, use_global_device_ids); in CreateAllGather()
1271 bool constrain_layout, const std::optional<int64_t>& channel_id, in CreateAllGatherStart() argument
[all …]
Dhlo_query.cc136 DynCast<HloCollectiveInstruction>(hlo)->constrain_layout()) { in ContainsLayoutConstrainedCollective()
Dall_to_all_decomposer.cc44 if (all_to_all->constrain_layout()) { in InstructionMatchesPattern()
Dhlo_instruction.h686 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout,
699 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout,
718 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout,
727 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout,
747 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout,
778 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout,
Dhlo_parser.cc1432 optional<bool> constrain_layout; in CreateInstruction() local
1440 &constrain_layout}; in CreateInstruction()
1454 constrain_layout ? *constrain_layout : false, channel_id, in CreateInstruction()
1459 constrain_layout ? *constrain_layout : false, channel_id, in CreateInstruction()
1469 optional<bool> constrain_layout; in CreateInstruction() local
1478 &constrain_layout}; in CreateInstruction()
1496 constrain_layout ? *constrain_layout : false, channel_id, in CreateInstruction()
1501 constrain_layout ? *constrain_layout : false, channel_id, in CreateInstruction()
1507 constrain_layout ? *constrain_layout : false, channel_id, in CreateInstruction()
1519 optional<bool> constrain_layout; in CreateInstruction() local
[all …]
Dwhile_loop_all_reduce_code_motion.cc640 old_all_reduce->constrain_layout(), in CreateSinkedAllReduces()
789 if (all_reduce_instruction->constrain_layout()) { in Run()
Dhlo_instructions.h551 bool constrain_layout() const { return constrain_layout_; } in constrain_layout() function
559 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout,
580 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout,
625 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout,
682 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout,
719 absl::Span<const ReplicaGroup> replica_groups, bool constrain_layout,
Dreduce_scatter_utils.cc273 if (!ar->shape().IsArray() || ar->constrain_layout() || in MatchReduceScatter()
Dwhile_loop_all_reduce_code_motion_test.cc124 EXPECT_FALSE(moved_all_reduce->constrain_layout()); in TEST_F()
Dhlo.proto262 // constrain_layout is true then the 'shape' field must contain a layout, and
265 bool constrain_layout = 56; field
Dhlo_verifier.cc2095 if (reference->constrain_layout() != all_reduce->constrain_layout()) { in VerifyLayoutConstrainedAllReduce()
Dhlo_parser_test.cc1736 ROOT crs = f32[8]{0} all-reduce(input), replica_groups={}, constrain_layout=true, to_apply=add in CreateTestCases()
1816 …ROOT ag = f32[128,128]{0,1} all-gather(input), replica_groups={}, constrain_layout=true, dimension… in CreateTestCases()
Dhlo_rematerialization.cc75 return !collective->constrain_layout(); in IsRematerializable()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_reduce_scatter_creator.cc87 ar->constrain_layout(), channel_id, ar->use_global_device_ids(), in Run()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dcanonicalize_all_gather_for_cse.cc84 ag->replica_groups(), ag->constrain_layout(), new_channel_id, in RunOnComputation()
Dschedule_aware_collective_ops_cse.cc71 if (coll->constrain_layout()) { in MayConsiderCollective()
/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/tests/Dialect/lhlo/
Dops.mlir12 {channel_id = #mhlo.channel_handle<handle = 1, type = 0>, constrain_layout = false,
27 {channel_id = #mhlo.channel_handle<handle = 1, type = 0>, constrain_layout = false,
51 constrain_layout = false, replica_groups = dense<0> : tensor<1x1xi64>,
61 {channel_id = #mhlo.channel_handle<handle = 1, type = 0>, constrain_layout = false,
72 {channel_id = #mhlo.channel_handle<handle = 1, type = 0>, constrain_layout = false,
83 {channel_id = #mhlo.channel_handle<handle = 1, type = 0>, constrain_layout = false,
789 constrain_layout = true,
/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/include/mlir-hlo/Dialect/lhlo_gpu/IR/
Dlhlo_gpu_ops.td177 DefaultValuedAttr<BoolAttr, "false">:$constrain_layout,
/external/tensorflow/tensorflow/compiler/mlir/tfrt/tests/lhlo_to_tfrt/
Dlhlo_gpu_to_tfrt_gpu.mlir184 constrain_layout = true,
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/hlo_to_lhlo_with_xla/
Dhlo_text_to_lhlo_no_opt.hlotxt544 // CHECK-SAME: constrain_layout = false

12