| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | all_to_all_decomposer.cc | 39 auto* all_to_all = DynCast<HloAllToAllInstruction>(instruction); in InstructionMatchesPattern() local
40 if (all_to_all == nullptr) { in InstructionMatchesPattern()
44 if (all_to_all->constrain_layout()) { in InstructionMatchesPattern()
47 if (all_to_all->shape().IsTuple()) { in InstructionMatchesPattern()
53 return all_to_all->shape().rank() < min_array_rank_; in InstructionMatchesPattern()
57 auto* all_to_all = Cast<HloAllToAllInstruction>(instruction); in ExpandInstruction() local
58 int64 split_dim = *all_to_all->split_dimension(); in ExpandInstruction()
60 all_to_all->replica_groups().empty() in ExpandInstruction()
62 : all_to_all->replica_groups()[0].replica_ids_size(); in ExpandInstruction()
64 all_to_all->shape().dimensions(split_dim) / all_to_all_group_size; in ExpandInstruction()
[all …]
|
| D | hlo_verifier.cc | 292 auto* all_to_all = Cast<HloAllToAllInstruction>(hlo); in HandleAllToAll() local
295 TF_RET_CHECK(all_to_all != nullptr); in HandleAllToAll()
296 if (all_to_all->split_dimension()) { in HandleAllToAll()
319 if (all_to_all->split_dimension()) { in HandleAllToAll()
323 hlo->operand(0)->shape(), *all_to_all->split_dimension(), in HandleAllToAll()
324 *all_to_all->split_dimension(), split_count)); in HandleAllToAll()
|
| /external/tensorflow/tensorflow/python/tpu/ops/ |
| D | tpu_ops.py | 45 def all_to_all(x, function
69 return gen_tpu_ops.all_to_all(
84 gen_tpu_ops.all_to_all(
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | fft_handler.cc | 167 auto all_to_all = collective_ops_creator.create_cross_partition_all_to_all( in ShuffleDataWithAllToAll() local
169 return all_to_all; in ShuffleDataWithAllToAll()
|
| D | spmd_partitioner_test.cc | 1144 auto all_to_all = in TEST_F() local
1146 auto reshard_lhs = AllOf(op::Reshape(op::Transpose(all_to_all)), in TEST_F()
1586 auto all_to_all = in TEST_F() local
1588 auto reshard = AllOf(op::Reshape(op::Transpose(all_to_all))); in TEST_F()
1621 auto all_to_all = in TEST_F() local
1623 auto reshard = op::Reshape(op::Transpose(all_to_all)); in TEST_F()
4939 auto all_to_all = AllOf(op::Shape("f32[4,4,2,4,4]"), op::AllToAll(reshape)); in TEST_F() local
4940 auto xpose = AllOf(op::Shape("f32[2,4,4,4,4]"), op::Transpose(all_to_all)); in TEST_F()
4963 auto all_to_all = op::AllToAll( in TEST_F() local
4966 AllOf(op::Shape("f32[2,4]"), op::Reshape(op::Transpose(all_to_all))); in TEST_F()
[all …]
|
| D | spmd_partitioner.cc | 1184 auto all_to_all = in ReshardWithAllToAll() local
1193 for (int64 i = 0; i < all_to_all->shape().rank(); ++i) { in ReshardWithAllToAll()
1203 ShapeInference::InferTransposeShape(all_to_all->shape(), permutation) in ReshardWithAllToAll()
1205 all_to_all, permutation)); in ReshardWithAllToAll()
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | tpu_strategy.py | 1324 result = tpu_ops.all_to_all(
1565 unordered_output = tpu_ops.all_to_all(
1578 replica_ids = tpu_ops.all_to_all(
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | mhlo_to_lhlo_with_xla.cc | 1051 auto* all_to_all = xla::Cast<xla::HloAllToAllInstruction>(instr); in EmitAllToAllOp() local
1054 if (all_to_all->split_dimension().has_value()) { in EmitAllToAllOp()
1056 builder_.getI64IntegerAttr(*all_to_all->split_dimension())); in EmitAllToAllOp()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | xla_builder.cc | 2894 XlaOp all_to_all, in AllToAllArray()
2897 return all_to_all; in AllToAllArray()
2908 all_to_all = Reshape(all_to_all, sizes); in AllToAllArray()
2918 all_to_all = Transpose(all_to_all, permutation); in AllToAllArray()
2919 return Reshape(all_to_all_shape, all_to_all); in AllToAllArray()
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/tests/ |
| D | lhlo_ops.mlir | 48 "lmhlo.all_to_all"(%input0, %output)
59 "lmhlo.all_to_all"(%input0, %output)
|
| D | ops.mlir | 18 %0 = "mhlo.all_to_all"(%data) {
31 %0 = "mhlo.all_to_all"(%data) {
44 %0 = "mhlo.all_to_all"(%data) {
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/ |
| D | lhlo_ops.td | 572 def LHLO_AllToAllOp : LHLO_CollectiveCommunicationOp<"all_to_all">,
|
| D | hlo_ops.td | 599 def HLO_AllToAllOp : HLO_Op<"all_to_all",
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/tests/hlo_to_lhlo_with_xla/ |
| D | hlo_text_to_lhlo_no_opt.hlotxt | 619 // CHECK: "lmhlo.all_to_all"
|
| /external/tensorflow/tensorflow/python/tpu/ |
| D | tensor_tracer.py | 1574 return tpu_ops.all_to_all(
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/tests/ |
| D | legalize-tf.mlir | 4885 // CHECK: mhlo.all_to_all
|