| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/ |
| H A 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()
|
| H A 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
|
| H A D | dtensor_allreduce_scatter_optimization.cc | 79 function.walk([&](mlir::TF::DTensorAllReduceOp all_reduce) { in ApplyOptimization()
|
| H A D | collectives.cc | 322 auto all_reduce = builder.create<mlir::TF::DTensorAllReduceOp>( in EmitAllReduce() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/utils/ |
| H A 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
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/kernel_tests/ |
| H A 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
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| H A 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
|
| H A 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
|
| H A 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()
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
| H A 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
|
| H A 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
|
| H A D | ar_crs_combiner.cc | 95 auto all_reduce = Cast<HloAllReduceInstruction>(hlo); in HasCombinableReplicaGroup() local
544 auto all_reduce = pair.ar; in RewriteGraph() local
|
| H A D | all_reduce_contiguous.cc | 31 Status ReplaceWithContiguousAllReduce(HloAllReduceInstruction* all_reduce) { in ReplaceWithContiguousAllReduce()
|
| H A D | all_reduce_combiner_test.cc | 219 auto all_reduce = b.AddInstruction(HloInstruction::CreateAllReduce( in TEST_F() local
|
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| H A 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()
|
| H A 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
|
| /aosp_15_r20/external/pytorch/torch/cuda/ |
| H A D | nccl.py | 71 def all_reduce(inputs, outputs=None, op=SUM, streams=None, comms=None): function
|
| /aosp_15_r20/external/pytorch/test/distributed/tensor/parallel/ |
| H A D | test_tp_examples.py | 47 reduce_scatter, all_gather, all_reduce = ( variable
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/ops/ |
| H A D | collective_ops.py | 19 def all_reduce(t, function
|
| /aosp_15_r20/external/tensorflow/tensorflow/python/distribute/ |
| H A D | cross_device_utils.py | 332 def all_reduce( member in CollectiveReplicaLauncher
|
| H A D | vars_test.py | 1012 def all_reduce(): function
|
| /aosp_15_r20/external/pytorch/torch/distributed/checkpoint/ |
| H A D | utils.py | 194 def all_reduce( member in _DistWrapper
|
| /aosp_15_r20/external/pytorch/torch/distributed/nn/ |
| H A D | functional.py | 205 def all_reduce(tensor, op=ReduceOp.SUM, group=group.WORLD): function
|
| /aosp_15_r20/external/pytorch/torch/distributed/_tools/ |
| H A D | fsdp2_mem_tracker.py | 517 def all_reduce( function
|
| /aosp_15_r20/external/pytorch/torch/csrc/distributed/c10d/ |
| H A D | Functional.cpp | 119 at::Tensor all_reduce( in all_reduce() function
|