[PATCH] D100751: [VPlan] Properly handle sinking of replicate regions.

[Florian Hahn via Phabricator via llvm-commits]

fhahn created this revision.
fhahn added reviewers: rengolin, Ayal, dmgreen, gilr.
Herald added subscribers: tschuett, psnobl, rogfer01, bollu, hiraditya.
fhahn requested review of this revision.
Herald added a subscriber: vkmr.
Herald added a project: LLVM.

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.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D100751

Files:
  llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
  llvm/lib/Transforms/Vectorize/VPlan.cpp
  llvm/lib/Transforms/Vectorize/VPlan.h


Index: llvm/lib/Transforms/Vectorize/VPlan.h
===================================================================
--- llvm/lib/Transforms/Vectorize/VPlan.h
+++ llvm/lib/Transforms/Vectorize/VPlan.h
@@ -1509,6 +1509,11 @@
 
   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.
Index: llvm/lib/Transforms/Vectorize/VPlan.cpp
===================================================================
--- llvm/lib/Transforms/Vectorize/VPlan.cpp
+++ llvm/lib/Transforms/Vectorize/VPlan.cpp
@@ -403,6 +403,31 @@
   }
 }
 
+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 {
Index: llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
===================================================================
--- llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
+++ llvm/lib/Transforms/Vectorize/LoopVectorize.cpp
@@ -8976,6 +8976,35 @@
         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);
   }
 


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D100751.338477.patch
Type: text/x-patch
Size: 3605 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20210419/98e8fea2/attachment.bin>