[XLA:CollectivePipeliner] Fix two issues:

1) Accept transpose as a formatting op in ForwardSink.
2) Do not stop when a large collective was sunk in the previous iteration. Instead, delay sinking large collectives while sinking small collectives level by level and run an additional sinking iteration dedicated to large collectives at the end.

PiperOrigin-RevId: 833425046

Author: seherellis

xla/service/collective_pipeliner.cc

@@ -348,7 +348,7 @@ CheckStoreIntoSliceIsCompatible(
     if (direction ==
         collective_pipeliner_utils::PipeliningDirection::kForwardSink) {
       // TODO(maggioni): Support these ops in forward sink.
-      if (HloPredicateIsOp<HloOpcode::kConcatenate, HloOpcode::kGetTupleElement,
+      if (HloPredicateIsOp<HloOpcode::kGetTupleElement,
                            HloOpcode::kReduceScatter>(i)) {
         return false;
       }
@@ -819,13 +819,17 @@ class WhileLoopAnalysis {
       HloInstruction* while_instr, int64_t max_pipelining_per_loop,
       bool pipeline_use_tree, bool process_different_sized_options,
       TuplePointsToAnalysis* tuple_points_to_analysis,
-      std::optional<ConstantValue> known_start = std::nullopt)
+      std::optional<ConstantValue> known_start = std::nullopt,
+      bool delay_sinking_large_collectives = false,
+      int64_t collective_size_threshold = INT64_MAX)
       : while_(while_instr),
         loop_start_(known_start),
         max_pipelining_per_loop_(max_pipelining_per_loop),
         tuple_points_to_analysis_(tuple_points_to_analysis),
         pipeline_use_tree_(pipeline_use_tree),
-        process_different_sized_options_(process_different_sized_options) {}
+        process_different_sized_options_(process_different_sized_options),
+        delay_sinking_large_collectives_(delay_sinking_large_collectives),
+        collective_size_threshold_(collective_size_threshold) {}
   std::optional<ConstantValue> GetLoopIterationCount() const;
   std::optional<ConstantValue> GetLoopStart() const;
   std::optional<ConstantValue> GetLoopIncrement() const;
@@ -926,6 +930,8 @@ class WhileLoopAnalysis {
 
   bool pipeline_use_tree_;
   bool process_different_sized_options_;
+  bool delay_sinking_large_collectives_;
+  int64_t collective_size_threshold_;
 };
 
 int64_t WhileLoopAnalysis::GetDUSIndex(const HloInstruction* dus) const {
@@ -1368,6 +1374,17 @@ void WhileLoopAnalysis::CollectCollectivesToMove(
     if (!should_process(instr)) {
       continue;
     }
+    if (delay_sinking_large_collectives_ &&
+        direction ==
+            collective_pipeliner_utils::PipeliningDirection::kForwardSink &&
+        ShapeUtil::ElementsIn(instr->shape()) >= collective_size_threshold_) {
+      VLOG(1) << "Delay sinking " << instr->name() << " because its size "
+              << ShapeUtil::ElementsIn(instr->shape())
+              << " is greater than the threshold "
+              << collective_size_threshold_;
+      continue;
+    }
+
     if (direction ==
             collective_pipeliner_utils::PipeliningDirection::kForward ||
         direction ==
@@ -2462,6 +2479,17 @@ absl::Status TransformFormattingOp(
     pipelined_map[formatting_op] = expanded_transpose;
     return absl::OkStatus();
   }
+  if (formatting_op->opcode() == HloOpcode::kConcatenate) {
+    HloConcatenateInstruction* concat =
+        Cast<HloConcatenateInstruction>(formatting_op);
+    HloInstruction* expanded_concat =
+        loop_computation->AddInstruction(HloInstruction::CreateConcatenate(
+            ComputeFullOutputShape(to_move, formatting_op->shape()),
+            collect_operands(formatting_op),
+            concat->concatenate_dimension() + 1));
+    pipelined_map[formatting_op] = expanded_concat;
+    return absl::OkStatus();
+  }
   return absl::InvalidArgumentError(
       absl::StrCat("Unsupported instruction ", formatting_op->ToString()));
 }
@@ -3233,24 +3261,6 @@ static absl::Status TransformLoopBackward(
   return absl::OkStatus();
 }
 
-bool IsForwardSinkIterationFeasible(HloInstruction* while_inst,
-                                    int64_t collective_size_threshold) {
-  for (HloInstruction* inst :
-       while_inst->while_body()->root_instruction()->operands()) {
-    if (inst->opcode() == HloOpcode::kDynamicUpdateSlice &&
-        inst->operand(1)->IsCustomCall(
-            CollectivePipeliner::kSunkByPreviousStep)) {
-      HloInstruction* cc = inst->mutable_operand(1);
-      if (ShapeUtil::ElementsIn(cc->shape()) >= collective_size_threshold) {
-        VLOG(1) << "Encountered a large collective which was sunk by the "
-                   "previous step, should stop the iteration.";
-        return false;
-      }
-    }
-  }
-  return true;
-}
-
 absl::StatusOr<bool> CollectivePipeliner::RunPipeliner(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
@@ -3280,7 +3290,9 @@ absl::StatusOr<bool> CollectivePipeliner::RunPipeliner(
       auto loop_analysis = std::make_unique<WhileLoopAnalysis>(
           instruction, config_.max_pipelining_per_loop,
           config_.pipeline_use_tree, config_.process_different_sized_ops,
-          tuple_points_to_analysis.get());
+          tuple_points_to_analysis.get(), /*known_start=*/std::nullopt,
+          config_.delay_sinking_large_collectives,
+          config_.collective_size_threshold_to_delay_sinking);
       loop_analysis->ComputeLoopStatistics();
       if (loop_analysis->GetLoopIterationCount() &&
           loop_analysis->GetLoopIterationCount()->GetUnsignedValue() > 1) {
@@ -3297,12 +3309,6 @@ absl::StatusOr<bool> CollectivePipeliner::RunPipeliner(
   for (auto& [instruction, loop_analysis] : loop_analyses) {
     VLOG(1) << "While iterations: "
             << loop_analysis->GetLoopIterationCount()->ToString();
-    if (config_.pipelining_direction ==
-            collective_pipeliner_utils::PipeliningDirection::kForwardSink &&
-        !IsForwardSinkIterationFeasible(
-            instruction, config_.collective_size_threshold_to_stop_sinking)) {
-      continue;
-    }
     loop_analysis->CollectCollectivesToMove(
         config_.level_to_operate_on, config_.pipelining_direction,
         config_.should_process, config_.acceptable_formatting,
@@ -3393,20 +3399,24 @@ absl::StatusOr<bool> CollectivePipeliner::RunImpl(
     return RunPipeliner(module, execution_threads);
   }
 
-  // If the pipelining direction is kForwardSink, run the pipeliner until it
-  // does not change the module anymore. The maximum number of iterations should
-  // be equal to the maximum number of pipelineable collectives in a chain of
-  // users plus one. In each iteration, we pipeline the last pipelineable
-  // collectives, which do not have any other pipelineable collectives in their
-  // user subtree.
+  // If the pipelining direction is kForwardSink, first run the pipeliner on
+  // small collectives iteratively until it does not change the module anymore.
+  // In each iteration, we pipeline the last pipelineable collectives, which do
+  // not have any other pipelineable collectives in their user subtrees. Then
+  // run the pipeliner one last time on the large collectives.
   bool changed = true;
   int64_t iter = 0;
   while (changed) {
     TF_ASSIGN_OR_RETURN(changed, RunPipeliner(module, execution_threads));
-    VLOG(1) << "Finished running pipeliner's iteration: " << iter;
+    VLOG(1) << "Finished running pipeliner's iteration for small collectives: "
+            << iter;
     iter++;
   }
-  return iter > 1;
+  config_.delay_sinking_large_collectives = false;
+  TF_ASSIGN_OR_RETURN(changed, RunPipeliner(module, execution_threads));
+  VLOG(1) << "Finished running pipeliner's iteration for large collectives: "
+          << iter;
+  return iter > 1 || changed;
 }
 
 }  // namespace xla

xla/service/collective_pipeliner.h

@@ -108,7 +108,8 @@ class CollectivePipeliner : public HloModulePass {
     bool should_add_loop_invariant_op_in_chain = false;
     // Postprocessing hook which runs for every successfully pipelined op.
     HloPostprocessor postprocess_pipelined_ops;
-    int64_t collective_size_threshold_to_stop_sinking = INT64_MAX;
+    int64_t collective_size_threshold_to_delay_sinking = INT64_MAX;
+    bool delay_sinking_large_collectives = true;
   };
   static const char* const kInsertedByPreviousStep;
   static const char* const kSunkByPreviousStep;
@@ -155,7 +156,7 @@ class CollectivePipeliner : public HloModulePass {
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
 
  private:
-  const Config config_;
+  Config config_;
 };
 
 }  // namespace xla

xla/service/collective_pipeliner_test.cc

@@ -108,7 +108,7 @@ absl::StatusOr<bool> RunOptimizer(
     CollectivePipeliner::HloPostprocessor postprocess_backward_peeled_trailing =
         {},
     bool should_add_loop_invariant_op_in_chain = false,
-    int64_t collective_size_threshold_to_stop_sinking = INT64_MAX) {
+    int64_t collective_size_threshold_to_delay_sinking = INT64_MAX) {
   CollectivePipeliner::Config config = {
       /*level_to_operate_on=*/level_to_operate_on,
       /*max_pipelining_per_loop=*/INT64_MAX,
@@ -125,7 +125,7 @@ absl::StatusOr<bool> RunOptimizer(
       postprocess_backward_rotated, postprocess_backward_peeled_trailing,
       should_add_loop_invariant_op_in_chain,
       /*postprocess_pipelined_ops=*/{},
-      collective_size_threshold_to_stop_sinking};
+      collective_size_threshold_to_delay_sinking};
   HloPassPipeline pass("optimizer");
   pass.AddPass<HloVerifier>(/*layout_sensitive=*/false,
                             /*allow_mixed_precision=*/false);
@@ -3751,7 +3751,7 @@ ENTRY entry {
 }
 
 TEST_F(CollectivePipelinerTest,
-       ForwardSinkDependentPipelineableCollectivesDoNotPipeline) {
+       ForwardSinkDoNotStopPipeliningAfterLargeCollectives) {
   constexpr absl::string_view hlo_string = R"(
 HloModule module
 
@@ -3762,71 +3762,66 @@ add {
 }
 
 add.1 {
-  lhs.1 = bf16[] parameter(0)
-  rhs.1 = bf16[] parameter(1)
-  ROOT add.1 = bf16[] add(lhs.1, rhs.1)
-}
-
-while_body.clone {
-  sink_param.1 = (s32[], bf16[3,8,128]{2,1,0}, bf16[3,8,128]{2,1,0}, bf16[3,1,8,128]{3,2,1,0}) parameter(0)
-  get-tuple-element.0 = s32[] get-tuple-element(sink_param.1), index=0
-  constant.5 = s32[] constant(1)
-  add.2 = s32[] add(get-tuple-element.0, constant.5)
-  get-tuple-element.1 = bf16[3,8,128]{2,1,0} get-tuple-element(sink_param.1), index=1
-  get-tuple-element.2 = bf16[3,8,128]{2,1,0} get-tuple-element(sink_param.1), index=2
-  get-tuple-element.3 = bf16[3,1,8,128]{3,2,1,0} get-tuple-element(sink_param.1), index=3
-  constant.6 = s32[] constant(3)
-  subtract.0 = s32[] subtract(constant.6, get-tuple-element.0)
-  constant.7 = s32[] constant(-1)
-  add.3 = s32[] add(subtract.0, constant.7)
-  constant.8 = s32[] constant(0)
-  compare.0 = pred[] compare(add.3, constant.8), direction=LT
-  constant.9 = s32[] constant(2)
-  add.4 = s32[] add(subtract.0, constant.9)
-  select.0 = s32[] select(compare.0, add.4, add.3)
-  dynamic-slice.0 = bf16[1,8,128]{2,1,0} dynamic-slice(get-tuple-element.2, select.0, constant.8, constant.8), dynamic_slice_sizes={1,8,128}
-  mul.1 = bf16[1,8,128]{2,1,0} multiply(dynamic-slice.0, dynamic-slice.0)
-  ar.0 = bf16[1,8,128]{2,1,0} all-reduce(mul.1), channel_id=1, replica_groups={}, to_apply=add
-  b.0 = bf16[1,8,128,32]{3,2,1,0} broadcast(ar.0), dimensions={0,1,2}
-  constant.10 = bf16[] constant(0)
-  reduce.1 = bf16[1,8,128]{2,1,0} reduce(b.0, constant.10), dimensions={3}, to_apply=add.1
-  reshape.1 = bf16[1,1,8,128]{3,2,1,0} reshape(reduce.1)
-  custom-call.2 = bf16[1,1,8,128]{3,2,1,0} custom-call(reshape.1), custom_call_target="SunkByPreviousStep"
-  constant.12 = s32[] constant(0)
-  dynamic-update-slice.1 = bf16[3,1,8,128]{3,2,1,0} dynamic-update-slice(get-tuple-element.3, custom-call.2, select.0, constant.12, constant.12, constant.12)
-  ROOT tuple.3 = (s32[], bf16[3,8,128]{2,1,0}, bf16[3,8,128]{2,1,0}, bf16[3,1,8,128]{3,2,1,0}) tuple(add.2, get-tuple-element.1, get-tuple-element.2, dynamic-update-slice.1)
+  lhs = bf16[] parameter(0)
+  rhs = bf16[] parameter(1)
+  ROOT add = bf16[] add(lhs, rhs)
 }
 
-while_cond.clone {
-  sink_param = (s32[], bf16[3,8,128]{2,1,0}, bf16[3,8,128]{2,1,0}, bf16[3,1,8,128]{3,2,1,0}) parameter(0)
-  gte.1 = s32[] get-tuple-element(sink_param), index=0
-  constant.13 = s32[] constant(3)
-  ROOT cmp.1 = pred[] compare(gte.1, constant.13), direction=LT
+while_cond {
+  param = (s32[], bf16[3,8,128], bf16[3,8,128], bf16[3,128,128], bf16[3,128,128]) parameter(0)
+  gte = s32[] get-tuple-element(param), index=0
+  constant.1 = s32[] constant(3)
+  ROOT cmp = pred[] compare(gte, constant.1), direction=LT
+}
+
+while_body {
+  param = (s32[], bf16[3,8,128], bf16[3,8,128], bf16[3,128,128], bf16[3,128,128]) parameter(0)
+  get-tuple-element.394 = s32[] get-tuple-element(param), index=0
+  get-tuple-element.395 = bf16[3,8,128] get-tuple-element(param), index=1
+  get-tuple-element.35 = bf16[3,8,128] get-tuple-element(param), index=2
+  constant.2557 = s32[] constant(1)
+  add.230 = s32[] add(get-tuple-element.394, constant.2557)
+  constant.2559 = s32[] constant(3)
+  subtract.139 = s32[] subtract(constant.2559, get-tuple-element.394)
+  constant.2560 = s32[] constant(-1)
+  add.231 = s32[] add(subtract.139, constant.2560)
+  constant.2561 = s32[] constant(0)
+  compare.747 = pred[] compare(add.231, constant.2561), direction=LT
+  constant.2562 = s32[] constant(2)
+  add.232 = s32[] add(subtract.139, constant.2562)
+  select.1348 = s32[] select(compare.747, add.232, add.231)
+  dynamic-slice.99 = bf16[1,8,128] dynamic-slice(get-tuple-element.35, select.1348, constant.2561, constant.2561), dynamic_slice_sizes={1,8,128}
+  mul = bf16[1,8,128] multiply(dynamic-slice.99, dynamic-slice.99)
+  ar.1 = bf16[1,8,128] all-reduce(mul), replica_groups={}, to_apply=add, channel_id=1
+  b.1 = bf16[1,8,128,32] broadcast(ar.1), dimensions={0,1,2}
+  constant = bf16[] constant(0)
+  reduce = bf16[1,8,128] reduce(b.1, constant), dimensions={3}, to_apply=add.1
+  ar.2 = bf16[1,8,128] all-reduce(reduce), replica_groups={}, to_apply=add, channel_id=2
+  c1 = bf16[] constant(2.0)
+  bc = bf16[1,8,128] broadcast(c1)
+  mul1 = bf16[1,8,128] multiply(ar.2, bc)
+  mul3 = bf16[1,8,128] multiply(mul1, ar.2)
+  dynamic-update-slice.35 = bf16[3,8,128] dynamic-update-slice(get-tuple-element.395, mul3, select.1348, constant.2561, constant.2561)
+  get-tuple-element.396 = bf16[3,128,128] get-tuple-element(param), index=3
+  get-tuple-element.36 = bf16[3,128,128] get-tuple-element(param), index=4
+  dynamic-slice.100 = bf16[1,128,128] dynamic-slice(get-tuple-element.36, select.1348, constant.2561, constant.2561), dynamic_slice_sizes={1,128,128}
+  large-ar = bf16[1,128,128] all-reduce(dynamic-slice.100), replica_groups={}, to_apply=add, channel_id=3
+  dynamic-update-slice.36 = bf16[3,128,128] dynamic-update-slice(get-tuple-element.396, large-ar, select.1348, constant.2561, constant.2561)
+  ROOT tuple = (s32[], bf16[3,8,128], bf16[3,8,128], bf16[3,128,128], bf16[3,128,128]) tuple(add.230, dynamic-update-slice.35, get-tuple-element.35, dynamic-update-slice.36, get-tuple-element.36)
 }
 
 ENTRY entry {
   c0 = s32[] constant(0)
-  p0 = bf16[3,8,128]{2,1,0} parameter(0)
-  constant.2 = bf16[] constant(0)
-  broadcast = bf16[3,1,8,128]{3,2,1,0} broadcast(constant.2), dimensions={}
-  tuple.2 = (s32[], bf16[3,8,128]{2,1,0}, bf16[3,8,128]{2,1,0}, bf16[3,1,8,128]{3,2,1,0}) tuple(c0, p0, p0, broadcast)
-  while.1 = (s32[], bf16[3,8,128]{2,1,0}, bf16[3,8,128]{2,1,0}, bf16[3,1,8,128]{3,2,1,0}) while(tuple.2), condition=while_cond.clone, body=while_body.clone
-  get-tuple-element.5 = s32[] get-tuple-element(while.1), index=0
-  get-tuple-element.4 = bf16[3,1,8,128]{3,2,1,0} get-tuple-element(while.1), index=3
-  ar.4 = bf16[3,1,8,128]{3,2,1,0} all-reduce(get-tuple-element.4), channel_id=3, replica_groups={}, to_apply=add
-  c1.3 = bf16[] constant(2)
-  broadcast.1 = bf16[3,1,8,128]{3,2,1,0} broadcast(c1.3), dimensions={}
-  mul1.2 = bf16[3,1,8,128]{3,2,1,0} multiply(ar.4, broadcast.1)
-  mul3.2 = bf16[3,1,8,128]{3,2,1,0} multiply(mul1.2, ar.4)
-  reshape.2 = bf16[3,8,128]{2,1,0} reshape(mul3.2)
-  get-tuple-element.6 = bf16[3,8,128]{2,1,0} get-tuple-element(while.1), index=2
-  tuple.4 = (s32[], bf16[3,8,128]{2,1,0}, bf16[3,8,128]{2,1,0}) tuple(get-tuple-element.5, reshape.2, get-tuple-element.6)
-  ROOT gte1 = bf16[3,8,128]{2,1,0} get-tuple-element(tuple.4), index=1
+  p0 = bf16[3,8,128] parameter(0)
+  p1 = bf16[3,128,128] parameter(1)
+  tuple = (s32[], bf16[3,8,128], bf16[3,8,128], bf16[3,128,128], bf16[3,128,128]) tuple(c0, p0, p0, p1, p1)
+  while = (s32[], bf16[3,8,128], bf16[3,8,128], bf16[3,128,128], bf16[3,128,128]) while(tuple), condition=while_cond, body=while_body
+  ROOT gte1 = bf16[3,8,128] get-tuple-element(while), index=1
 }
 )";
   config_.set_use_spmd_partitioning(true);
   auto module = ParseAndReturnUnverifiedModule(hlo_string, config_).value();
-  EXPECT_FALSE(
+  EXPECT_TRUE(
       RunOptimizer(
           module.get(), /*last_run=*/false,
           /*level_to_operate_on=*/0,
@@ -3842,8 +3837,15 @@ ENTRY entry {
           /*postprocess_backward_rotated=*/{},
           /*postprocess_backward_peeled_trailing=*/{},
           /*should_add_loop_invariant_op_in_chain=*/false,
-          /*collective_size_threshold_to_stop_sinking=*/1024)
+          /*collective_size_threshold_to_delay_sinking=*/2048)
           .value());
+  XLA_VLOG_LINES(1, module->ToString());
+  const HloInstruction* while_instr =
+      FindInstruction(module.get(), HloOpcode::kWhile);
+  EXPECT_TRUE(absl::c_none_of(while_instr->while_body()->instructions(),
+                              [](const HloInstruction* instr) {
+                                return instr->opcode() == HloOpcode::kAllReduce;
+                              }));
 }
 
 TEST_F(CollectivePipelinerTest, ForwardSinkFirstDimNotMatchingLoopCount) {