Searched refs:SpmdPartitioningVisitor (Results 1 – 6 of 6) sorted by relevance
1351 SpmdPartitioningVisitor::SpmdPartitioningVisitor( in SpmdPartitioningVisitor() function in xla::spmd::SpmdPartitioningVisitor1368 Status SpmdPartitioningVisitor::DefaultAction(HloInstruction* hlo) { in DefaultAction()1407 Status SpmdPartitioningVisitor::Preprocess(HloInstruction* hlo) { in Preprocess()1452 Status SpmdPartitioningVisitor::Postprocess(HloInstruction* hlo) { in Postprocess()1472 Status SpmdPartitioningVisitor::HandleElementwise(HloInstruction* hlo) { in HandleElementwise()1485 Status SpmdPartitioningVisitor::HandleConcatenate(HloInstruction* hlo) { in HandleConcatenate()1579 Status SpmdPartitioningVisitor::HandleSlice(HloInstruction* hlo) { in HandleSlice()1641 Status SpmdPartitioningVisitor::HandleSort(HloInstruction* hlo) { in HandleSort()1775 Status SpmdPartitioningVisitor::HandleCustomCall(HloInstruction* hlo) { in HandleCustomCall()1946 Status SpmdPartitioningVisitor::HandleTranspose(HloInstruction* hlo) { in HandleTranspose()[all …]
132 SpmdPartitioningVisitor* visitor);183 SpmdPartitioningVisitor* visitor) { in ParititonPassthroughOperand()218 SpmdPartitioningVisitor* visitor) { in ParititonTrivialIndexedOperandDimension()318 SpmdPartitioningVisitor* visitor) { in PartitionIndexParallelDimensions()465 SpmdPartitioningVisitor* visitor) { in PartitionGather()507 Status SpmdPartitioningVisitor::HandleScatter(HloInstruction* hlo) { in HandleScatter()671 Status SpmdPartitioningVisitor::HandleGather(HloInstruction* hlo) { in HandleGather()
180 class SpmdPartitioningVisitor; variable226 virtual std::unique_ptr<SpmdPartitioningVisitor> CreateVisitor(400 class SpmdPartitioningVisitor : public DfsHloVisitorWithDefault {402 SpmdPartitioningVisitor(
42 Status SpmdPartitioningVisitor::HandleDot(HloInstruction* hlo) { in HandleDot()481 std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>* in PartitionBaseCase()1645 std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>*1660 std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>* in PartitionDotGroupOnBatch()1994 std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>* in PartitionDotGroupOnNonContracting()2089 std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>* in PartitionDotGroupOnContracting()2441 std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>* in PartitionDot()2807 std::vector<SpmdPartitioningVisitor::WindowedDotGeneralLoop>* in PartitionDot()2835 Status SpmdPartitioningVisitor::HandleDotHelper( in HandleDotHelper()3400 Status SpmdPartitioningVisitor::DoCodeMotionForWindowedDotGeneralLoops( in DoCodeMotionForWindowedDotGeneralLoops()
351 Status SpmdPartitioningVisitor::HandleFft(HloInstruction* hlo) { in HandleFft()
987 Status SpmdPartitioningVisitor::HandleConvolution(HloInstruction* hlo) { in HandleConvolution()