Home
last modified time | relevance | path

Searched refs:original_hlo (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc47 const Window& conv_window, HloInstruction* original_hlo, in PartitionConvolutionWithBatchGroupCount() argument
49 TF_RET_CHECK(original_hlo->opcode() == HloOpcode::kConvolution); in PartitionConvolutionWithBatchGroupCount()
50 if (original_hlo->batch_group_count() == 1 || in PartitionConvolutionWithBatchGroupCount()
51 original_hlo->batch_group_count() < num_partitions) { in PartitionConvolutionWithBatchGroupCount()
55 const auto& dnums = original_hlo->convolution_dimension_numbers(); in PartitionConvolutionWithBatchGroupCount()
62 original_hlo->batch_group_count() != input_batch_size) { in PartitionConvolutionWithBatchGroupCount()
139 const Window& conv_window, HloInstruction* original_hlo, in PartitionConvolutionWithFeatureGroupCount() argument
141 TF_RET_CHECK(original_hlo->opcode() == HloOpcode::kConvolution); in PartitionConvolutionWithFeatureGroupCount()
142 if (original_hlo->feature_group_count() == 1 || in PartitionConvolutionWithFeatureGroupCount()
143 original_hlo->feature_group_count() < num_partitions) { in PartitionConvolutionWithFeatureGroupCount()
[all …]
Ddot_handler.cc473 const Window& conv_window, HloModule* module, HloInstruction* original_hlo, in PartitionBaseCase() argument
659 SpmdBuilder body_b("windowed_dot_general_body", original_hlo); in PartitionBaseCase()
755 if (original_hlo->opcode() == HloOpcode::kDot) { in PartitionBaseCase()
756 new_ddnums = original_hlo->dot_dimension_numbers(); in PartitionBaseCase()
867 if (original_hlo->opcode() == HloOpcode::kDot) { in PartitionBaseCase()
873 if (original_hlo->opcode() == HloOpcode::kDot) { in PartitionBaseCase()
902 if (original_hlo->opcode() == HloOpcode::kDot) { in PartitionBaseCase()
925 if (original_hlo->opcode() == HloOpcode::kDot) { in PartitionBaseCase()
952 ShapeUtil::MakeShape(original_hlo->shape().element_type(), new_dims); in PartitionBaseCase()
955 if (original_hlo->opcode() == HloOpcode::kDot) { in PartitionBaseCase()
[all …]
Dconvolution_handler.h35 const Window& conv_window, HloInstruction* original_hlo,