Searched refs:SpmdPartitioner (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner.h | 182 class SpmdPartitioner : public HloModulePass {
184 SpmdPartitioner(int64 num_partitions, int64 num_replicas,
186 SpmdPartitioner(int64 num_partitions, int64 num_replicas, in SpmdPartitioner() function
291 SpmdPartitioner* partitioner;
406 SpmdPartitionerOptions options, SpmdPartitioner* partitioner);
546 SpmdPartitioner* partitioner_;
|
| D | spmd_partitioner.cc | 1355 SpmdPartitioner* partitioner) in SpmdPartitioningVisitor()
3493 SpmdPartitioner::SpmdPartitioner(int64 num_partitions, int64 num_replicas, in SpmdPartitioner() function in xla::spmd::SpmdPartitioner
3495 : SpmdPartitioner( in SpmdPartitioner()
3499 HloInstruction* SpmdPartitioner::AllGatherShards( in AllGatherShards()
3508 HloInstruction* SpmdPartitioner::AllGatherShardsInternal( in AllGatherShardsInternal()
3598 HloInstruction* SpmdPartitioner::AllReduceAlongShardingDims( in AllReduceAlongShardingDims()
3608 HloInstruction* SpmdPartitioner::AllReduceAlongShardingDimsInternal( in AllReduceAlongShardingDimsInternal()
3632 StatusOr<bool> SpmdPartitioner::PartitionComputation( in PartitionComputation()
3641 std::unique_ptr<SpmdPartitioningVisitor> SpmdPartitioner::CreateVisitor( in CreateVisitor()
3651 StatusOr<bool> SpmdPartitioner::Run(HloModule* module) { in Run()
[all …]
|
| D | spmd_partitioner_test.cc | 60 pass.AddPass<SpmdPartitioner>(num_devices, /*num_replicas=*/1, options, in PartitionComputation()
|