Searched refs:left_halo_size_function (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner_util.cc | 455 OffsetCalculation left_halo_size_function = in TileToPartialReplicateHaloExchange() local
467 result, left_halo_size_function, right_halo_size_function, dim, in TileToPartialReplicateHaloExchange()
554 OffsetCalculation left_halo_size_function = in PadFromPartialReplicateShape() local
567 result, left_halo_size_function, right_halo_size_function, dim, in PadFromPartialReplicateShape()
799 HloInstruction* hlo, const OffsetCalculation& left_halo_size_function, in ExchangeHalo() argument
809 int64 max_left_halo_size = left_halo_size_function.MaxInRange(1, shard_count); in ExchangeHalo()
924 const OffsetCalculation& left_halo_size_function, in ExchangeHaloAndGetValidData() argument
933 ExchangeHalo(hlo, left_halo_size_function, right_halo_size_function, dim, in ExchangeHaloAndGetValidData()
940 int64 max_left_halo_size = left_halo_size_function.MaxInRange(1, shard_count); in ExchangeHaloAndGetValidData()
954 left_halo_size_function; in ExchangeHaloAndGetValidData()
[all …]
|
| D | spmd_partitioner_util.h | 218 HloInstruction* hlo, const OffsetCalculation& left_halo_size_function,
254 const OffsetCalculation& left_halo_size_function,
|
| D | fft_handler.cc | 64 OffsetCalculation left_halo_size_function = in PadEachPartitionWithHaloExchange() local
80 ExchangeHalo(hlo, left_halo_size_function, right_halo_size_function, in PadEachPartitionWithHaloExchange()
|