Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.h134 class OffsetCalculation; variable
145 OffsetCalculation operator-(
172 class OffsetCalculation {
174 OffsetCalculation() : opcode_(HloOpcode::kCopy), copy_from_() {} in OffsetCalculation() function
175 explicit OffsetCalculation( in OffsetCalculation() function
178 OffsetCalculation(const OffsetCalculation& copy_from) { *this = copy_from; } in OffsetCalculation() function
179 OffsetCalculation(HloOpcode opcode, in OffsetCalculation() function
183 lhs_(absl::make_unique<OffsetCalculation>(lhs)), in OffsetCalculation()
184 rhs_(absl::make_unique<OffsetCalculation>(rhs)) {} in OffsetCalculation()
185 OffsetCalculation(HloOpcode opcode, const OffsetCalculation& lhs, in OffsetCalculation() function
[all …]
Dspmd_partitioner_util.cc455 OffsetCalculation left_halo_size_function = in TileToPartialReplicateHaloExchange()
456 OffsetCalculation(MultiplyAddDivideOffsetCalculation( in TileToPartialReplicateHaloExchange()
461 OffsetCalculation right_halo_size_function = in TileToPartialReplicateHaloExchange()
462 OffsetCalculation(MultiplyAddDivideOffsetCalculation(0, 0, 1)); in TileToPartialReplicateHaloExchange()
483 OffsetCalculation start_offset_on_padded_concat_calculation = in TileToPartialReplicateHaloExchange()
484 OffsetCalculation(MultiplyAddDivideOffsetCalculation( in TileToPartialReplicateHaloExchange()
554 OffsetCalculation left_halo_size_function = in PadFromPartialReplicateShape()
555 OffsetCalculation(MultiplyAddDivideOffsetCalculation(0, 0, 1)); in PadFromPartialReplicateShape()
559 OffsetCalculation right_halo_size_function = in PadFromPartialReplicateShape()
560 OffsetCalculation(MultiplyAddDivideOffsetCalculation( in PadFromPartialReplicateShape()
[all …]
Dconvolution_handler.cc340 std::vector<OffsetCalculation> left_halo_size_functions( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
342 std::vector<OffsetCalculation> right_halo_size_functions( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
380 OffsetCalculation(MultiplyAddDivideOffsetCalculation( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
386 OffsetCalculation(MultiplyAddDivideOffsetCalculation( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
411 OffsetCalculation(HloOpcode::kMultiply, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
413 OffsetCalculation(MultiplyAddDivideOffsetCalculation( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
415 OffsetCalculation(MultiplyAddDivideOffsetCalculation( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
437 (OffsetCalculation( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
479 OffsetCalculation(MultiplyAddDivideOffsetCalculation( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
649 std::vector<OffsetCalculation> left_halo_size_functions( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
[all …]
Dfft_handler.cc64 OffsetCalculation left_halo_size_function = in PadEachPartitionWithHaloExchange()
65 OffsetCalculation(MultiplyAddDivideOffsetCalculation(0, 0, 1)); in PadEachPartitionWithHaloExchange()
72 OffsetCalculation right_halo_size_function = in PadEachPartitionWithHaloExchange()
73 OffsetCalculation(MultiplyAddDivideOffsetCalculation( in PadEachPartitionWithHaloExchange()
92 OffsetCalculation start_offset_on_padded_concat_calculation = in PadEachPartitionWithHaloExchange()
93 OffsetCalculation(MultiplyAddDivideOffsetCalculation( in PadEachPartitionWithHaloExchange()
Dspmd_partitioner.cc776 std::vector<OffsetCalculation> left_halo_size_functions(base_shape_.rank()); in ReshardAsWindowedInput()
777 std::vector<OffsetCalculation> right_halo_size_functions(base_shape_.rank()); in ReshardAsWindowedInput()
3216 std::vector<OffsetCalculation> data_left_halo_sizes(hlo->shape().rank()); in HandleSelectAndScatter()
3217 std::vector<OffsetCalculation> data_right_halo_sizes(hlo->shape().rank()); in HandleSelectAndScatter()
3218 std::vector<OffsetCalculation> source_left_halo_sizes(hlo->shape().rank()); in HandleSelectAndScatter()
3219 std::vector<OffsetCalculation> source_right_halo_sizes(hlo->shape().rank()); in HandleSelectAndScatter()
3264 OffsetCalculation(MultiplyAddDivideOffsetCalculation( in HandleSelectAndScatter()
3266 OffsetCalculation( in HandleSelectAndScatter()
3270 OffsetCalculation( in HandleSelectAndScatter()
3273 OffsetCalculation(MultiplyAddDivideOffsetCalculation( in HandleSelectAndScatter()