Home
last modified time | relevance | path

Searched refs:SPMDCollectiveOpsCreator (Results 1 – 5 of 5) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner.h113 struct SPMDCollectiveOpsCreator { struct
148 SPMDCollectiveOpsCreator GetDefaultCollectiveOpsCreator(int64 num_partitions,
188 SPMDCollectiveOpsCreator collective_ops_creator) in SpmdPartitioner()
213 const SPMDCollectiveOpsCreator& collectives_creator);
220 const SPMDCollectiveOpsCreator& collectives_creator,
228 const SPMDCollectiveOpsCreator& collective_ops_creator,
235 const SPMDCollectiveOpsCreator& collectives_creator, bool per_dim_ag);
239 const SPMDCollectiveOpsCreator& collectives_creator,
250 SPMDCollectiveOpsCreator collective_ops_creator_;
288 SPMDCollectiveOpsCreator collective_ops_creator;
[all …]
Dspmd_partitioner_util.h221 const SPMDCollectiveOpsCreator& collective_ops_creator,
231 const SPMDCollectiveOpsCreator& collective_ops_creator,
260 const SPMDCollectiveOpsCreator& collective_ops_creator,
370 const SPMDCollectiveOpsCreator& collective_ops_creator,
393 const SPMDCollectiveOpsCreator& collective_ops_creator,
Dfft_handler.cc54 const SPMDCollectiveOpsCreator& collective_ops_creator, in PadEachPartitionWithHaloExchange()
161 const SPMDCollectiveOpsCreator& collective_ops_creator, in ShuffleDataWithAllToAll()
229 const SPMDCollectiveOpsCreator& collective_ops_creator, in GetFinalFftUsingCollectivePermute()
Dspmd_partitioner_util.cc412 const SPMDCollectiveOpsCreator& collective_ops_creator, in TileToPartialReplicateHaloExchange()
506 const SPMDCollectiveOpsCreator& collective_ops_creator, in PadFromPartialReplicateShape()
802 const SPMDCollectiveOpsCreator& collective_ops_creator, in ExchangeHalo()
904 const SPMDCollectiveOpsCreator& collective_ops_creator, in ExchangeHalo()
930 const SPMDCollectiveOpsCreator& collective_ops_creator, in ExchangeHaloAndGetValidData()
1590 SPMDCollectiveOpsCreator GetPerGroupCollectiveOpsCreator( in GetPerGroupCollectiveOpsCreator()
1591 const SPMDCollectiveOpsCreator& creator, in GetPerGroupCollectiveOpsCreator()
1593 SPMDCollectiveOpsCreator result; in GetPerGroupCollectiveOpsCreator()
Dspmd_partitioner.cc1353 const SPMDCollectiveOpsCreator& collective_ops_creator, in SpmdPartitioningVisitor()
3411 SPMDCollectiveOpsCreator GetDefaultCollectiveOpsCreator(int64 num_partitions, in GetDefaultCollectiveOpsCreator()
3502 const SPMDCollectiveOpsCreator& collectives_creator) { in AllGatherShards()
3511 const SPMDCollectiveOpsCreator& collectives_creator, bool per_dim_ag) { in AllGatherShardsInternal()
3601 const SPMDCollectiveOpsCreator& collectives_creator, in AllReduceAlongShardingDims()
3611 const SPMDCollectiveOpsCreator& collectives_creator, in AllReduceAlongShardingDimsInternal()
3643 const SPMDCollectiveOpsCreator& collective_ops_creator, in CreateVisitor()