[VPlan] Properly handle sinking of replicate regions.

[fhahn]

Summary

This patch updates the code that sinks recipes required for first-order
recurrences to properly handle replicate-regions. At the moment, the
code would just move the replicate recipe out of its replicate-region,
producing an invalid VPlan.

When sinking a recipe in a replicate-region, we have to sink the whole
region. To do that, we first need to split the block at the target
recipe and move the region in between.

This patch also adds a splitAt helper to VPBasicBlock to split a
VPBasicBlock at a given iterator.

I think we also need to deal with the case when the target is a
replicate region also. I'll put up a patch once I constructed a test
for that scenario.

Fixes PR50009.

[fhahn]

Add missing tests.

[Ayal]

This looks good, thanks for fixing, adding a couple of nits.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8966

I think we also need to deal with the case when the target is a
replicate region also. I'll put up a patch once I constructed a test
for that scenario.

Agreed. Worth adding an assert in the meanwhile?

9005

better early-exit the simple Sink->moveAfter(Target) non-pred-rep case first?

llvm/lib/Transforms/Vectorize/VPlan.cpp

407

error message

llvm/lib/Transforms/Vectorize/VPlan.h

1513

mobing

[fhahn]

Address comments, thanks!

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8966

That's a great idea, I added an assert.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8966

Actually this case is already handled, not sure if I was thinking about something else. I also updated the logic to check directly whether sink is in a replicate region, same as for target.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8966

Worth mentioning where the case having both sink and target predicated is handled(?)

Better rename Region >> TargetRegion.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8966

Ah right, I missed the *also* when re-reading. So the unhandled case is when both sink and target are in separate regions! Let me look into adding an assert for that and update the naming.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

8966

I pushed 01c26d4e048c & 337d7652823f to rename the variable and add an extra assert.

[ebrevnov]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9024

Hi Florian,

Just want to give you a heads up that we observe crash downstream at this line. In particular, std::next(Target->getIterator()) returns null. So far I have internal reproducer only. Maybe you can figure out the problem faster than I come up with upstreamable reproducer. Thanks in advance.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9024

Thanks for the heads-up! I'll try to see if I can come up with a test case tomorrow.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9024

I managed to construct a test case that requires splitting at the end of the block, which crashes at the moment, because the assert dereferences splitAt unconditionally. Should be fixed in bdada7546e6b

Please let me know if there's still an issue downstream.

[ebrevnov]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9024

Thanks for quick turn around. Will check and let you know.

[ebrevnov]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9024

Thanks for quick turn around. Will check and let you know.

That does fix our issue. Thanks!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

VPlanPtr LoopVectorizationPlanner::buildVPlanWithVPRecipes(
   // ---------------------------------------------------------------------------
   // Transform initial VPlan: Apply previously taken decisions, in order, to
   // bring the VPlan to its final state.
   // ---------------------------------------------------------------------------
 
   // Apply Sink-After legal constraints.
   for (auto &Entry : SinkAfter) {
     VPRecipeBase *Sink = RecipeBuilder.getRecipe(Entry.first);
     VPRecipeBase *Target = RecipeBuilder.getRecipe(Entry.second);
     // 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 *Region =
             dyn_cast_or_null<VPRegionBlock>(Target->getParent()->getParent())) {
       if (Region->isReplicator()) {
         assert(Region->getNumSuccessors() == 1 && "Expected SESE region!");
         VPBasicBlock *NextBlock =
             cast<VPBasicBlock>(Region->getSuccessors().front());
         Sink->moveBefore(*NextBlock, NextBlock->getFirstNonPhi());
         continue;
       }
     }
+
+    // If the source is a predicate replicate recipe, we need to move the whole
+    // replicate region, which should only contain a single recipe in the main
+    // block.
+    auto *RepR = dyn_cast<VPReplicateRecipe>(Sink);
+    if (RepR && RepR->isPredicated()) {
+      auto *Region =
+          dyn_cast_or_null<VPRegionBlock>(RepR->getParent()->getParent());
+      assert(Region->isReplicator() && RepR->getParent()->size() == 1 &&
+             "parent must be a replicator with a single recipe");
+      auto *SplitBlock =
+          Target->getParent()->splitAt(std::next(Target->getIterator()));
+
+      auto *Pred = Region->getSinglePredecessor();
+      auto *Succ = Region->getSingleSuccessor();
+      VPBlockUtils::disconnectBlocks(Pred, Region);
+      VPBlockUtils::disconnectBlocks(Region, Succ);
+      VPBlockUtils::connectBlocks(Pred, Succ);
+
+      auto *SplitPred = SplitBlock->getSinglePredecessor();
+
+      VPBlockUtils::disconnectBlocks(SplitPred, SplitBlock);
+      VPBlockUtils::connectBlocks(SplitPred, Region);
+      VPBlockUtils::connectBlocks(Region, SplitBlock);
+      if (VPBB == SplitPred)
+        VPBB = SplitBlock;
+      continue;
+    }
+
     Sink->moveAfter(Target);
   }
 
   // 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;
     for (unsigned i = 0; i < IG->getFactor(); ++i)
       if (auto *SI = dyn_cast_or_null<StoreInst>(IG->getMember(i)))
         StoredValues.push_back(Plan->getOrAddVPValue(SI->getOperand(0)));
 
     auto *VPIG = new VPInterleaveRecipe(IG, Recipe->getAddr(), StoredValues,
                                         Recipe->getMask());
     VPIG->insertBefore(Recipe);
     unsigned J = 0;
     for (unsigned i = 0; i < IG->getFactor(); ++i)
       if (Instruction *Member = IG->getMember(i)) {
         if (!Member->getType()->isVoidTy()) {
           VPValue *OriginalV = Plan->getVPValue(Member);
           Plan->removeVPValueFor(Member);
           Plan->addVPValue(Member, VPIG->getVPValue(J));
           OriginalV->replaceAllUsesWith(VPIG->getVPValue(J));

llvm/lib/Transforms/Vectorize/VPlan.h

public:
   /// this VPBasicBlock, thereby "executing" the VPlan.
   void execute(struct VPTransformState *State) override;
 
   /// Return the position of the first non-phi node recipe in the block.
   iterator getFirstNonPhi();
 
   void dropAllReferences(VPValue *NewValue) override;
 
+  /// Split current block at \p SplitAt by inserting a new block between the
+  /// current block and its successors and mobing all recipes starting at
+  /// SplitAt to the new block. Returns the new block.
+  VPBasicBlock *splitAt(iterator SplitAt);
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print this VPBsicBlock to \p O, prefixing all lines with \p Indent. \p
   /// SlotTracker is used to print unnamed VPValue's using consequtive numbers.
   ///
   /// Note that the numbering is applied to the whole VPlan, so printing
   /// individual blocks is consistent with the whole VPlan printing.
   void print(raw_ostream &O, const Twine &Indent,
              VPSlotTracker &SlotTracker) const override;

llvm/lib/Transforms/Vectorize/VPlan.cpp

for (VPRecipeBase &R : Recipes) {
     for (auto *Def : R.definedValues())
       Def->replaceAllUsesWith(NewValue);
 
     for (unsigned I = 0, E = R.getNumOperands(); I != E; I++)
       R.setOperand(I, NewValue);
   }
 }
 
+VPBasicBlock *VPBasicBlock::splitAt(iterator SplitAt) {
+  assert(SplitAt->getParent() == this);
+
+  SmallVector<VPBlockBase *, 2> Succs(getSuccessors().begin(),
+                                      getSuccessors().end());
+  // First, disconnect the current block from its successors.
+  for (VPBlockBase *Succ : Succs)
+    VPBlockUtils::disconnectBlocks(this, Succ);
+
+  // Create new empty block after the block to split.
+  auto *SplitBlock = new VPBasicBlock(getName() + ".split");
+  VPBlockUtils::insertBlockAfter(SplitBlock, this);
+
+  // Add successors for block to split to new block.
+  for (VPBlockBase *Succ : Succs)
+    VPBlockUtils::connectBlocks(SplitBlock, Succ);
+
+  // Finally, move the recipes starting at SplitAt to new block.
+  for (VPRecipeBase &ToMove :
+       make_early_inc_range(make_range(SplitAt, this->end())))
+    ToMove.moveBefore(*SplitBlock, SplitBlock->end());
+
+  return SplitBlock;
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 void VPBasicBlock::print(raw_ostream &O, const Twine &Indent,
                          VPSlotTracker &SlotTracker) const {
   O << Indent << getName() << ":\n";
   if (const VPValue *Pred = getPredicate()) {
     O << Indent << "BlockPredicate:";
     Pred->printAsOperand(O, SlotTracker);
     if (const auto *PredInst = dyn_cast<VPInstruction>(Pred))