| /external/tensorflow/tensorflow/dtensor/mlir/ |
| D | dtensor_allreduce_combine_optimization.cc | 197 mlir::TF::DTensorAllReduceOp& all_reduce = all_reduce_group[i]; in MergeAllReduceGroup() local
234 mlir::TF::DTensorAllReduceOp& all_reduce = all_reduce_group[i]; in MergeAllReduceGroup() local
257 mlir::TF::DTensorAllReduceOp& all_reduce = all_reduce_group[i]; in MergeAllReduceGroup() local
358 for (mlir::TF::DTensorAllReduceOp all_reduce : all_reduces) { in CombineAllReduceOpsOfSameTypeAndGroupAssignment() local
472 mlir::TF::DTensorAllReduceOp all_reduce = all_reduce_group[i]; in CombineAllReduceOpsOfSameTypeAndGroupAssignment() local
476 mlir::TF::DTensorAllReduceOp all_reduce = all_reduce_group[i]; in CombineAllReduceOpsOfSameTypeAndGroupAssignment() local
518 for (mlir::TF::DTensorAllReduceOp all_reduce : all_reduces) { in CombineAllReduceOpsOfSameType() local
554 cluster.GetBody().walk([&](mlir::TF::DTensorAllReduceOp all_reduce) { in runOnOperation()
|
| D | dtensor_allreduce_sum_optimization.cc | 154 mlir::TF::DTensorAllReduceOp all_reduce = in OptimizeAllReduceAndSum() local
290 mlir::TF::DTensorAllReduceOp all_reduce = in CheckWhileLoopOptimizationCriteria() local
367 const int output_index, mlir::TF::DTensorAllReduceOp all_reduce, in ExtractAllReduceFromWhileOp()
427 mlir::TF::DTensorAllReduceOp all_reduce; in OptimizeWhileLoopLazyAllReduce() local
|
| D | dtensor_allreduce_scatter_optimization.cc | 79 function.walk([&](mlir::TF::DTensorAllReduceOp all_reduce) { in ApplyOptimization()
|
| D | collectives.cc | 322 auto all_reduce = builder.create<mlir::TF::DTensorAllReduceOp>( in EmitAllReduce() local
|
| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | nvptx_compiler_test.cc | 53 HloInstruction* all_reduce = module->entry_computation()->root_instruction(); in TEST_F() local
83 HloInstruction* all_reduce = module->entry_computation()->root_instruction(); in TEST_F() local
|
| D | all_reduce_blueconnect_test.cc | 75 auto all_reduce = AllOf(op::Shape("f32[4]"), op::AllReduce(reduce_scatter), in TEST_F() local
117 auto all_reduce = AllOf(op::Shape("f32[2]"), op::AllReduce(reduce_scatter1), in TEST_F() local
162 auto all_reduce = AllOf(op::Shape("(f32[4], f32[8])"), in TEST_F() local
|
| D | all_reduce_blueconnect.cc | 111 const HloAllReduceInstruction& all_reduce, size_t num_devices_per_host) { in TryDecomposeReplicaGroups()
186 StatusOr<bool> TryDecomposeAllReduce(HloAllReduceInstruction* all_reduce, in TryDecomposeAllReduce()
|
| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | all_reduce_simplifier.cc | 41 [this](const HloInstruction* all_reduce) -> int64_t { in Run()
94 auto all_reduce = all_reduce_and_group_size.first; in Run() local
|
| D | while_loop_all_reduce_code_motion.cc | 159 HloInstruction* all_reduce, HloComputation* while_body, in IsAllReduceMovable()
170 auto all_reduce_is_summation = [](HloInstruction* all_reduce) -> bool { in IsAllReduceMovable()
464 HloComputation* while_body) -> MovableAllReduceContext { in IsAllReduceMovable()
829 HloInstruction* all_reduce = all_reduce_accumulations_pair.first; in Run() local
|
| D | all_reduce_contiguous.cc | 31 Status ReplaceWithContiguousAllReduce(HloAllReduceInstruction* all_reduce) { in ReplaceWithContiguousAllReduce()
|
| D | ar_crs_combiner.cc | 95 auto all_reduce = Cast<HloAllReduceInstruction>(hlo); in HasCombinableReplicaGroup() local
544 auto all_reduce = pair.ar; in RewriteGraph() local
|
| D | all_reduce_combiner_test.cc | 219 auto all_reduce = b.AddInstruction(HloInstruction::CreateAllReduce( in TEST_F() local
|
| D | hlo_cost_analysis_test.cc | 988 HloInstruction* all_reduce = module->entry_computation()->root_instruction(); in TEST_F() local
|
| D | bfloat16_propagation_test.cc | 212 HloInstruction* all_reduce = in TEST_F() local
|
| D | hlo_verifier.cc | 2091 auto all_reduce = DynCast<HloAllReduceInstruction>(instruction); in VerifyLayoutConstrainedAllReduce() local
|
| /external/tensorflow/tensorflow/dtensor/mlir/utils/ |
| D | collective_lowering.cc | 91 mlir::TF::DTensorAllReduceOp all_reduce, in EmitAllReduceForXla()
205 mlir::TF::DTensorAllReduceOp all_reduce, mlir::Value* value) { in LowerAllReduceOpImpl()
310 mlir::TF::DTensorAllReduceOp all_reduce) { in LowerAllReduceOp()
616 mlir::TF::DTensorAllReduceOp all_reduce = in LowerAllGatherOp() local
738 module.walk([&](mlir::TF::DTensorAllReduceOp all_reduce) { in runOnOperation()
743 for (auto& all_reduce : all_reduces) in runOnOperation() local
760 module.walk([&](mlir::TF::DTensorReduceScatterOp all_reduce) { in runOnOperation()
765 for (auto& all_reduce : all_reduces) in runOnOperation() local
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf_collective.cc | 142 auto all_reduce = builder.create<AllReduceOp>( in ConvertAllReduce() local
199 LogicalResult matchAndRewrite(TF::XlaAllReduceOp all_reduce, in matchAndRewrite()
251 LogicalResult matchAndRewrite(TF::CollectiveReduceV2Op all_reduce, in matchAndRewrite()
|
| D | mhlo_to_lhlo_with_xla.cc | 1175 auto* all_reduce = xla::Cast<xla::HloAllReduceInstruction>(instr); in EmitAllReduceOp() local
1200 auto* all_reduce = xla::Cast<xla::HloAllReduceInstruction>(instr); in EmitAllReduceStartOp() local
|
| /external/tensorflow/tensorflow/python/kernel_tests/ |
| D | collective_ops_test.py | 47 def all_reduce(t, group_size, group_key, instance_key, *args, **kwargs): member in CollectiveOpsV1
65 def all_reduce(t, group_size, group_key, instance_key, *args, **kwargs): member in CollectiveOpsV2
571 def all_reduce(device): function
607 def all_reduce(device): function
644 def all_reduce(device, device_index): function
682 def all_reduce(device, device_index): function
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | collective_ops.py | 19 def all_reduce(t, function
|
| /external/tensorflow/tensorflow/python/distribute/ |
| D | cross_device_utils.py | 332 def all_reduce( member in CollectiveReplicaLauncher
|
| D | vars_test.py | 1012 def all_reduce(): function
|
| D | tpu_strategy_test.py | 776 def all_reduce(x): function
|
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | gather_scatter_handler.cc | 935 auto all_reduce = in PartitionScatterIndexPassthroughPartition() local
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/ |
| D | hlo_function_importer.cc | 988 auto all_reduce = Cast<HloAllReduceInstruction>(instruction); in ImportInstructionImpl() local
|