[LV] Support sinking recipe in replicate region after another region.

[fhahn]

Summary

This patch handles sinking a replicate region after another replicate
region. In that case, we can connect the sink region after the target
region. This properly handles the case for which an assertion has been
added in 337d7652823f.

Fixes https://bugs.chromium.org/p/oss-fuzz/issues/detail?id=34842.

[fhahn]

Rebase & ping :)

[Ayal]

Minor comments and suggestions. This should compete the handling of sinking a recipe after another, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9198

Set SinkRegion and TargetRegion once at the outset?

9200

These asserts are better placed inside the lambda?

9201

Have this simplest and most common (!SinkRegion && !TargetRegion) case early-exit first?
Followed by the next simple (!SinkRegion && TargetRegion) case; both move the Sink recipe via Sink->move*();
Last deal with SinkRegion by first disconnecting it and then placing it between an AfterBlock and its single successor, where AfterBlock is produced by splitting the Target block if !TargetRegion.

9203

// If the recipe >> // The recipe

"recpliate"

9205–9206

same assert...

9205–9206

This part of disconnecting SinkRegion could be done once for both TargetRegion and !TargetRegion cases.

9209

same asserts...

9222

// If the target >> // The target

llvm/test/Transforms/LoopVectorize/first-order-recurrence-sink-replicate-region.ll

388

%rem and %rem.div both need to sink after Previous = %recur.next (which is invariant...), and do so by first having %rem move after Previous and then having %rem.div move after %rem, where this last move is the one that moves a replicated region after another, right?

[fhahn]

Thanks for all the comments Ayal! I updated the code and removed some redundancies as suggested.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9200

I moved the SESE one. The other one would be a bit more work, because currently the lambda does not have access to Target's/Sink's parents. A general helper to check for replicate regions in the future may be helpful.

9201

I restructured the code to handle the simplest case first, then SinkRegion and then TargetRegion.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9200

Lambda can check if R->getParent() has size 1, right? This should be asserted for both sink and target.

9226

Handle the "SinkRegion && !TargetRegion" case here, splitting TargetBlock etc., instead of last below?

9228

Handle this relatively simple "!SinkRegion && TargetRegion" case 2nd, after the simplest !SinkRegion && !TargetRegion? (I.e., can check first if (!SinkRegion), etc.)

[fhahn]

Restructure checks as suggested and moved the assertion to lambda, thanks!

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9200

Yes, I moved the assertion,

9226

Done, thanks!

9228

I restructured the code, it should be a bit more compact now.

[fhahn]

ping :)

[Ayal]

Looks good to me, thanks for accommodating the changes!
Adding a minor error message suggestion.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9194

"parent of R must be a replicator with a single recipe" >>
"A recipe in an original replicator region must be the only recipe in its block"?
("original" - because this does not hold for merged replicator regions).

[fhahn]

Thanks!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9194

Sounds good, I'll update the message when committing!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

VPlanPtr LoopVectorizationPlanner::buildVPlanWithVPRecipes(
   // Apply Sink-After legal constraints.
   for (auto &Entry : SinkAfter) {
     VPRecipeBase *Sink = RecipeBuilder.getRecipe(Entry.first);
     VPRecipeBase *Target = RecipeBuilder.getRecipe(Entry.second);
 
     auto GetReplicateRegion = [](VPRecipeBase *R) -> VPRegionBlock * {
       auto *Region =
           dyn_cast_or_null<VPRegionBlock>(R->getParent()->getParent());
-      if (Region && Region->isReplicator())
+      if (Region && Region->isReplicator()) {
+        assert(Region->getNumSuccessors() == 1 &&
+               Region->getNumPredecessors() == 1 && "Expected SESE region!");
         return Region;
+      }
       return nullptr;
     };
-
-    // If the target is in a replication region, make sure to move Sink to the
-    // block after it, not into the replication region itself.
-    if (auto *TargetRegion = GetReplicateRegion(Target)) {
-      assert(TargetRegion->getNumSuccessors() == 1 && "Expected SESE region!");
-      assert(!GetReplicateRegion(Sink) &&
-             "cannot sink a region into another region yet");
-      VPBasicBlock *NextBlock =
-          cast<VPBasicBlock>(TargetRegion->getSuccessors().front());
-      Sink->moveBefore(*NextBlock, NextBlock->getFirstNonPhi());
-      continue;
-    }
-
+    auto *TargetRegion = GetReplicateRegion(Target);
     auto *SinkRegion = GetReplicateRegion(Sink);
     // Unless the sink source is in a replicate region, sink the recipe
     // directly.
-    if (!SinkRegion) {
+    if (!SinkRegion && !TargetRegion) {
       Sink->moveAfter(Target);
       continue;
     }
 
-    // If the sink source is in a replicate region, we need to move the whole
-    // replicate region, which should only contain a single recipe in the main
-    // block.
+    if (SinkRegion) {
+      // The recipe to sink is in a replicate region. Unhook the region from the
+      // CFG.
       assert(Sink->getParent()->size() == 1 &&
              "parent must be a replicator with a single recipe");
+
+      auto *SinkPred = SinkRegion->getSinglePredecessor();
+      auto *SinkSucc = SinkRegion->getSingleSuccessor();
+      VPBlockUtils::disconnectBlocks(SinkPred, SinkRegion);
+      VPBlockUtils::disconnectBlocks(SinkRegion, SinkSucc);
+      VPBlockUtils::connectBlocks(SinkPred, SinkSucc);
+
+      if (TargetRegion) {
+        // The target recipe is also in a replicate region, move the sink region
+        // after the target region.
+        auto *TargetSucc = TargetRegion->getSingleSuccessor();
+        VPBlockUtils::disconnectBlocks(TargetRegion, TargetSucc);
+        VPBlockUtils::connectBlocks(TargetRegion, SinkRegion);
+        VPBlockUtils::connectBlocks(SinkRegion, TargetSucc);
+        continue;
+      }
+    }
+    if (TargetRegion) {
+      // The target is in a replication region, make sure to move Sink to
+      // the block after it, not into the replication region itself.
+      VPBasicBlock *NextBlock =
+          cast<VPBasicBlock>(TargetRegion->getSuccessors().front());
+      Sink->moveBefore(*NextBlock, NextBlock->getFirstNonPhi());
+    } else {
+      // The sink source is in a replicate region, we need to move the whole
+      // replicate region, which should only contain a single recipe in the main
+      // block.
       auto *SplitBlock =
           Target->getParent()->splitAt(std::next(Target->getIterator()));
 
-    auto *Pred = SinkRegion->getSinglePredecessor();
-    auto *Succ = SinkRegion->getSingleSuccessor();
-    VPBlockUtils::disconnectBlocks(Pred, SinkRegion);
-    VPBlockUtils::disconnectBlocks(SinkRegion, Succ);
-    VPBlockUtils::connectBlocks(Pred, Succ);
-
       auto *SplitPred = SplitBlock->getSinglePredecessor();
 
       VPBlockUtils::disconnectBlocks(SplitPred, SplitBlock);
       VPBlockUtils::connectBlocks(SplitPred, SinkRegion);
       VPBlockUtils::connectBlocks(SinkRegion, SplitBlock);
       if (VPBB == SplitPred)
         VPBB = SplitBlock;
     }
+  }
 
   // Interleave memory: for each Interleave Group we marked earlier as relevant
   // for this VPlan, replace the Recipes widening its memory instructions with a
   // single VPInterleaveRecipe at its insertion point.
   for (auto IG : InterleaveGroups) {
     auto *Recipe = cast<VPWidenMemoryInstructionRecipe>(
         RecipeBuilder.getRecipe(IG->getInsertPos()));
     SmallVector<VPValue *, 4> StoredValues;

llvm/test/Transforms/LoopVectorize/first-order-recurrence-sink-replicate-region.ll

 ; REQUIRES: asserts
-; RUN: opt < %s -loop-vectorize -force-vector-width=2 -force-vector-interleave=1 -debug-only=loop-vectorize 2>&1 | FileCheck %s
+; RUN: opt < %s -loop-vectorize -force-vector-width=2 -force-vector-interleave=1 -disable-output -debug-only=loop-vectorize 2>&1 | FileCheck %s
 
 target datalayout = "e-m:e-i64:64-i128:128-n32:64-S128"
 
 ; Test cases for PR50009, which require sinking a replicate-region due to a
 ; first-order recurrence.
 
 define void @sink_replicate_region_1(i32 %x, i8* %ptr) optsize {
 ; CHECK-LABEL: sink_replicate_region_1
loop:
   %add = add i32 %add.1, %conv.lv.2
   %iv.next = add nsw i32 %iv, 1
   %ec = icmp eq i32 %iv.next, 20001
   br i1 %ec, label %exit, label %loop
 
 exit:
   ret void
 }
+
+; Test case that requires sinking a recipe in a replicate region after another replicate region.
+define void @sink_replicate_region_after_replicate_region(i32* %ptr, i32 %x, i8 %y) optsize {
+; CHECK-LABEL: sink_replicate_region_after_replicate_region
+; CHECK:      VPlan 'Initial VPlan for VF={2},UF>=1' {
+; CHECK-NEXT: loop:
+; CHECK-NEXT:   WIDEN-PHI %recur = phi 0, %recur.next
+; CHECK-NEXT:   WIDEN-INDUCTION %iv = phi 0, %iv.next
+; CHECK-NEXT:   EMIT vp<%3> = icmp ule ir<%iv> vp<%0>
+; CHECK-NEXT: Successor(s): loop.0
+; CHECK-EMPTY:
+; CHECK-NEXT: loop.0:
+; CHECK-NEXT: Successor(s): loop.1
+; CHECK-EMPTY:
+; CHECK-NEXT:  loop.1:
+; CHECK-NEXT:   WIDEN ir<%recur.next> = sext ir<%y>
+; CHECK-NEXT: Successor(s): pred.srem
+; CHECK-EMPTY:
+; CHECK-NEXT: <xVFxUF> pred.srem: {
+; CHECK-NEXT:   pred.srem.entry:
+; CHECK-NEXT:     BRANCH-ON-MASK vp<%3>
+; CHECK-NEXT:   Successor(s): pred.srem.if, pred.srem.continue
+; CHECK-NEXT:   CondBit: vp<%3> (loop)
+; CHECK-EMPTY:
+; CHECK-NEXT:   pred.srem.if:
+; CHECK-NEXT:     REPLICATE ir<%rem> = srem ir<%recur>, ir<%x>
+; CHECK-NEXT:   Successor(s): pred.srem.continue
+; CHECK-EMPTY:
+; CHECK-NEXT:   pred.srem.continue:
+; CHECK-NEXT:     PHI-PREDICATED-INSTRUCTION vp<%6> = ir<%rem>
+; CHECK-NEXT:   No successors
+; CHECK-NEXT: }
+
+; CHECK:      <xVFxUF> pred.sdiv: {
+; CHECK-NEXT:  pred.sdiv.entry:
+; CHECK-NEXT:    BRANCH-ON-MASK vp<%3>
+; CHECK-NEXT:  Successor(s): pred.sdiv.if, pred.sdiv.continue
+; CHECK-NEXT:  CondBit: vp<%3> (loop)
+; CHECK-EMPTY:
+; CHECK-NEXT:       pred.sdiv.if:
+; CHECK-NEXT:    REPLICATE ir<%rem.div> = sdiv ir<20>, vp<%6>
+; CHECK-NEXT:  Successor(s): pred.sdiv.continue
+; CHECK-EMPTY:
+; CHECK-NEXT:       pred.sdiv.continue:
+; CHECK-NEXT:    PHI-PREDICATED-INSTRUCTION vp<%8> = ir<%rem.div>
+; CHECK-NEXT:  No successors
+; CHECK-NEXT: }
+; CHECK-NEXT: Successor(s): loop.1.split
+; CHECK-EMPTY:
+; CHECK-NEXT:      loop.1.split:
+; CHECK-NEXT: Successor(s): pred.store
+; CHECK-EMPTY:
+; CHECK-NEXT:       <xVFxUF> pred.store: {
+; CHECK-NEXT:   pred.store.entry:
+; CHECK-NEXT:     BRANCH-ON-MASK vp<%3>
+; CHECK-NEXT:   Successor(s): pred.store.if, pred.store.continue
+; CHECK-NEXT:   CondBit: vp<%3> (loop)
+; CHECK-EMPTY:
+; CHECK-NEXT:   pred.store.if:
+; CHECK-NEXT:     REPLICATE ir<%gep> = getelementptr ir<%ptr>, ir<%iv>
+; CHECK-NEXT:     REPLICATE store vp<%8>, ir<%gep>
+; CHECK-NEXT:   Successor(s): pred.store.continue
+; CHECK-EMPTY:
+; CHECK-NEXT:   pred.store.continue:
+; CHECK-NEXT:   No successors
+; CHECK-NEXT: }
+
+; CHECK:      loop.2:
+; CHECK-NEXT: No successors
+; CHECK-NEXT: }
+;
+entry:
+  br label %loop
+
+loop:                                             ; preds = %loop, %entry
+  %recur = phi i32 [ 0, %entry ], [ %recur.next, %loop ]
+  %iv = phi i32 [ 0, %entry ], [ %iv.next, %loop ]
+  %rem = srem i32 %recur, %x
+  %rem.div = sdiv i32 20, %rem
+  %recur.next = sext i8 %y to i32
+  %gep = getelementptr i32, i32* %ptr, i32 %iv
+  store i32 %rem.div, i32* %gep
+  %iv.next = add nsw i32 %iv, 1
+  %C = icmp sgt i32 %iv.next, %recur.next
+  br i1 %C, label %exit, label %loop
+
+exit:                                             ; preds = %loop
+  ret void
+}