Searched defs:SpmdBuilder (Results 1 – 3 of 3) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| H A D | convolution_handler.cc | 45 HloInstruction*, HloInstruction*, SpmdBuilder*, in PartitionConvolutionWithBatchGroupCount() argument
137 HloInstruction*, HloInstruction*, SpmdBuilder*, in PartitionConvolutionWithFeatureGroupCount() argument
231 HloInstruction*, HloInstruction*, SpmdBuilder*, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() argument
520 HloInstruction*, HloInstruction*, SpmdBuilder*, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS() argument
746 HloInstruction*, HloInstruction*, SpmdBuilder*, in PartitionConvolutionTiledOutput() argument
822 HloInstruction*, HloInstruction*, SpmdBuilder*, in PartitionConvolutionBaseCase() argument
970 HloInstruction*, HloInstruction*, SpmdBuilder*, in PartitionConvolution() argument
|
| H A D | dot_handler.cc | 453 HloInstruction*, HloInstruction*, SpmdBuilder*, in GetWindowedEinsumConfiguration() argument
762 HloInstruction*, HloInstruction*, SpmdBuilder*, in PartitionBaseCase() argument
2012 HloInstruction*, HloInstruction*, SpmdBuilder*, in PartitionDotGroupOnBatch() argument
2356 HloInstruction*, HloInstruction*, SpmdBuilder*, in PartitionDotGroupOnNonContracting() argument
2585 HloInstruction*, HloInstruction*, SpmdBuilder*, in PartitionDotGroupOnContracting() argument
2996 bool require_matching_devices_to_group, SpmdBuilder* b, in PrioritizeContractingDimensionsPartitioning() argument
3224 int64_t rhs_batch_partitions, SpmdBuilder* b, const Window& conv_window, in LhsIsBestMatchForNonContractingPartitioning() argument
3348 HloInstruction*, HloInstruction*, SpmdBuilder*, in PartitionDot() argument
3736 HloInstruction*, HloInstruction*, SpmdBuilder*, in PartitionDot() argument
3772 HloInstruction*, HloInstruction*, SpmdBuilder*, in HandleDotHelper() argument
|
| H A D | spmd_partitioner.h | 79 SpmdBuilder(const std::string& name, HloInstruction* hlo) in SpmdBuilder() function
|