| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | all_gather_decomposer.cc | 55 ag->operand(0)->shape().dimensions(ag->all_gather_dimension()); in DecomposeAllGather()
56 const int64 ag_size = ag->shape().dimensions(ag->all_gather_dimension()); in DecomposeAllGather()
122 start_indices[ag->all_gather_dimension()] = in DecomposeAllGather()
|
| D | hlo_verifier.cc | 253 TF_RET_CHECK(ag->all_gather_dimension() >= 0); in HandleAllGather()
254 TF_RET_CHECK(ag->all_gather_dimension() < ag->shape().rank()); in HandleAllGather()
255 TF_RET_CHECK(ag->all_gather_dimension() < ag->operand(0)->shape().rank()); in HandleAllGather()
258 ag->shape().dimensions(ag->all_gather_dimension()), in HandleAllGather()
259 ag->operand(0)->shape().dimensions(ag->all_gather_dimension())); in HandleAllGather()
276 ag->operand(0)->shape(), ag->all_gather_dimension(), in HandleAllGather()
|
| D | shape_inference.h | 131 int64 all_gather_dimension,
|
| D | shape_inference.cc | 2035 const Shape& operand_shape, int64 all_gather_dimension, int64 shard_count) { in InferAllGatherShape() argument
2036 TF_RET_CHECK(all_gather_dimension >= 0); in InferAllGatherShape()
2037 TF_RET_CHECK(all_gather_dimension < operand_shape.rank()); in InferAllGatherShape()
2040 shape.set_dimensions(all_gather_dimension, in InferAllGatherShape()
2041 shard_count * shape.dimensions(all_gather_dimension)); in InferAllGatherShape()
|
| D | hlo_instructions.cc | 619 const Shape& shape, HloInstruction* operand, int64 all_gather_dimension, in HloAllGatherInstruction() argument
624 all_gather_dimension_(all_gather_dimension),
643 shape, new_operands[0], all_gather_dimension(), replica_groups(), in CloneWithNewOperandsImpl()
661 all_gather_dimension_ == casted_other.all_gather_dimension() && in IdenticalSlowPathIgnoringChannelIdValues()
|
| D | hlo_instructions.h | 391 const Shape& shape, HloInstruction* operand, int64 all_gather_dimension,
398 int64 all_gather_dimension() const { return all_gather_dimension_; } in all_gather_dimension() function
|
| D | hlo_instruction.cc | 422 int64 all_gather_dimension = proto.dimensions(0); in CreateFromProto() local
424 shape, operands(0), all_gather_dimension, in CreateFromProto()
1044 const Shape& shape, HloInstruction* operand, int64 all_gather_dimension, in CreateAllGather() argument
1048 shape, operand, all_gather_dimension, replica_groups, constrain_layout, in CreateAllGather()
|
| D | hlo_instruction.h | 634 const Shape& shape, HloInstruction* operand, int64 all_gather_dimension,
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | canonicalize_all_gather_for_cse.cc | 54 if (!ag || ag->all_gather_dimension() != 0) { in RunOnComputation()
58 const int64 ag_dim = ag->all_gather_dimension(); in RunOnComputation()
|
| D | spmd_partitioner.h | 143 int64 channel_id, int64 all_gather_dimension)>
|
| D | spmd_partitioner_util.cc | 1659 int64 channel_id, int64 all_gather_dimension) { in GetPerGroupCollectiveOpsCreator() argument
1663 all_gather_dimension); in GetPerGroupCollectiveOpsCreator()
|
| D | spmd_partitioner.cc | 3474 int64 channel_id, int64 all_gather_dimension) { in GetDefaultCollectiveOpsCreator()
3486 ag_shape, operand, all_gather_dimension, device_groups, in GetDefaultCollectiveOpsCreator()
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | nccl_all_gather_thunk.cc | 47 return (Cast<HloAllGatherInstruction>(hlo)->all_gather_dimension() == 0) && in CanImplement()
58 return op.all_gather_dimension() == 0 && operands_are_supported; in CanImplement()
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/tests/ |
| D | lhlo_ops.mlir | 40 …use_global_device_ids = false, all_gather_dimension = 0 : i64} : (memref<2xf32>, memref<8xf32>) ->…
51 …use_global_device_ids = false, all_gather_dimension = 0 : i64} : (memref<2xf32>, memref<8xf32>) ->…
62 …use_global_device_ids = false, all_gather_dimension = 0 : i64} : (memref<2xf32>, memref<8xf32>) ->…
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | xla_builder.h | 734 XlaOp operand, int64 all_gather_dimension, int64 shard_count,
1297 friend XlaOp AllGather(XlaOp operand, int64 all_gather_dimension,
2189 XlaOp operand, int64 all_gather_dimension, int64 shard_count,
|
| D | xla_builder.cc | 2714 XlaOp XlaBuilder::AllGather(XlaOp operand, int64 all_gather_dimension, in AllGather() argument
2726 *operand_shape, all_gather_dimension, shard_count)); in AllGather()
2733 instr.add_dimensions(all_gather_dimension); in AllGather()
4653 XlaOp AllGather(const XlaOp operand, int64 all_gather_dimension, in AllGather() argument
4659 return operand.builder()->AllGather(operand, all_gather_dimension, in AllGather()
|
| /external/tensorflow/tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/ |
| D | lhlo_ops.td | 563 (ins I64Attr:$all_gather_dimension));
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/tests/hlo_to_lhlo_with_xla/ |
| D | hlo_text_to_lhlo_no_opt.hlotxt | 607 // CHECK_SAME: all_gather_dimension = 1 : i64
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | mhlo_to_lhlo_with_xla.cc | 1071 builder_.getI64IntegerAttr(all_gather->all_gather_dimension())); in EmitAllGatherOp()
|