[LV] Apply sink-after transform on initial VPlan.

[fhahn]

rebased

[fhahn]

rebased and simplified.

[fhahn]

Putting this on hold for now, to first focus on sinking scalar operands as VPlan-to-VPlan transform: D100258

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

protected:
   void buildVPlans(unsigned MinVF, unsigned MaxVF);
 
 private:
   /// Build a VPlan according to the information gathered by Legal. \return a
   /// VPlan for vectorization factors \p Range.Start and up to \p Range.End
   /// exclusive, possibly decreasing \p Range.End.
   VPlanPtr buildVPlan(const VFRange &Range);
 
-  /// Build a VPlan using VPRecipes according to the information gather by
-  /// Legal. This method is only used for the legacy inner loop vectorizer.
-  VPlanPtr buildVPlanWithVPRecipes(
-      VFRange &Range, SmallPtrSetImpl<Value *> &NeedDef,
-      SmallPtrSetImpl<Instruction *> &DeadInstructions,
-      const DenseMap<Instruction *, Instruction *> &SinkAfter);
-
-  VPlanPtr
-  convertToVPRecipes(VPlan &OriginalPlan, VFRange &Range,
-                     SmallPtrSetImpl<Value *> &NeedDef,
-                     const DenseMap<Instruction *, Instruction *> &SinkAfter);
+  VPlanPtr convertToVPRecipes(VPlan &OriginalPlan, VFRange &Range,
+                              SmallPtrSetImpl<Value *> &NeedDef);
 
   /// Build VPlans for power-of-2 VF's between \p MinVF and \p MaxVF inclusive,
   /// according to the information gathered by Legal when it checked if it is
   /// legal to vectorize the loop. This method creates VPlans using VPRecipes.
-  void buildVPlansWithVPRecipes(unsigned MinVF, unsigned MaxVF);
   void buildVPlansWithVPRecipes(unsigned MinVF, unsigned MaxVF,
                                 VPlan &OriginalPlan);
 
   /// Adjust the recipes for any inloop reductions. The chain of instructions
   /// leading from the loop exit instr to the phi need to be converted to
   /// reductions, with one operand being vector and the other being the scalar
   /// reduction chain.
   void adjustRecipesForInLoopReductions(VPlanPtr &Plan,
                                         VPRecipeBuilder &RecipeBuilder);
 };
 
 } // namespace llvm
 
 #endif // LLVM_TRANSFORMS_VECTORIZE_LOOPVECTORIZATIONPLANNER_H

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

void LoopVectorizationPlanner::buildVPlansWithVPRecipes(unsigned MinVF,
   DenseMap<Instruction *, Instruction *> &SinkAfter = Legal->getSinkAfter();
   // Dead instructions do not need sinking. Remove them from SinkAfter.
   for (Instruction *I : DeadInstructions)
     SinkAfter.erase(I);
 
   auto *TopRegion = cast<VPRegionBlock>(OriginalPlan.getEntry());
   ReversePostOrderTraversal<VPBlockBase *> RPOT(TopRegion->getEntry());
 
+  DenseMap<Instruction *, VPInstruction *> Inst2VPInst;
+  for (auto &Entry : SinkAfter) {
+    Inst2VPInst[Entry.first] = nullptr;
+    Inst2VPInst[Entry.second] = nullptr;
+  }
+
   auto *DummyVal = new VPValue();
   OriginalPlan.addExternalDef(DummyVal);
   for (VPBlockBase *Base : RPOT) {
     // Do not widen instructions in pre-header and exit blocks.
     if (Base->getNumPredecessors() == 0 || Base->getNumSuccessors() == 0)
       continue;
 
     VPBasicBlock *OriginalVPBB = Base->getEntryBasicBlock();
     OriginalVPBB->setCondBit(nullptr);
     // Introduce each ingredient into VPlan.
     for (auto I = OriginalVPBB->rbegin(), E = OriginalVPBB->rend(); I != E;) {
       VPRecipeBase *Ingredient = &*I++;
       VPInstruction *VPInst = cast<VPInstruction>(Ingredient);
       Instruction *Instr = VPInst->getUnderlyingInstr();
 
       if (DeadInstructions.contains(Instr) || isa<DbgInfoIntrinsic>(Instr)) {
         Ingredient->getVPValue()->replaceAllUsesWith(DummyVal);
         Ingredient->eraseFromParent();
-      }
-    }
+        continue;
+      }
+
+      auto Iter = Inst2VPInst.find(Instr);
+      if (Iter != Inst2VPInst.end())
+        Iter->second = VPInst;
+
+  for (auto &Entry : SinkAfter) {
+    VPRecipeBase *Sink = Inst2VPInst.find(Entry.first)->second;
+    VPRecipeBase *Target = Inst2VPInst.find(Entry.second)->second;
+    Sink->moveAfter(Target);
   }
 
   for (unsigned VF = MinVF; VF < MaxVF + 1;) {
     VFRange SubRange = {VF, MaxVF + 1};
-    VPlans.push_back(
-        convertToVPRecipes(OriginalPlan, SubRange, NeedDef, SinkAfter));
+    VPlans.push_back(convertToVPRecipes(OriginalPlan, SubRange, NeedDef));
     VF = SubRange.End;
   }
 }
 
 VPlanPtr LoopVectorizationPlanner::convertToVPRecipes(
-    VPlan &OriginalPlan, VFRange &Range, SmallPtrSetImpl<Value *> &NeedDef,
-    const DenseMap<Instruction *, Instruction *> &SinkAfter) {
+    VPlan &OriginalPlan, VFRange &Range, SmallPtrSetImpl<Value *> &NeedDef) {
 
   // Hold a mapping from predicated instructions to their recipes, in order to
   // fix their AlsoPack behavior if a user is determined to replicate and use a
   // scalar instead of vector value.
   DenseMap<Instruction *, VPReplicateRecipe *> PredInst2Recipe;
 
   SmallPtrSet<const InterleaveGroup<Instruction> *, 1> InterleaveGroups;
   SmallPtrSet<Instruction *, 8> DeadInterleaveGroupMembers;
 
   VPRecipeBuilder RecipeBuilder(OrigLoop, TLI, Legal, CM, PSE, Builder);
 
   // ---------------------------------------------------------------------------
   // Pre-construction: record ingredients whose recipes we'll need to further
   // process after constructing the initial VPlan.
   // ---------------------------------------------------------------------------
 
-  // Mark instructions we'll need to sink later and their targets as
-  // ingredients whose recipe we'll need to record.
-  for (auto &Entry : SinkAfter) {
-    RecipeBuilder.recordRecipeOf(Entry.first);
-    RecipeBuilder.recordRecipeOf(Entry.second);
-  }
   for (auto &Reduction : CM.getInLoopReductionChains()) {
     PHINode *Phi = Reduction.first;
     RecurrenceDescriptor::RecurrenceKind Kind =
         Legal->getReductionVars()[Phi].getRecurrenceKind();
     const SmallVector<Instruction *, 4> &ReductionOperations = Reduction.second;
 
     RecipeBuilder.recordRecipeOf(Phi);
     for (auto &R : ReductionOperations) {
for (InterleaveGroup<Instruction> *IG : IAI.getInterleaveGroups()) {
     for (unsigned i = 0; i < IG->getFactor(); i++)
       if (Instruction *Member = IG->getMember(i)) {
         RecipeBuilder.recordRecipeOf(Member);
         if (Member != IG->getInsertPos())
           DeadInterleaveGroupMembers.insert(Member);
       }
   };
 
-  auto skipDeadInterleaveMembers =
-      [&DeadInterleaveGroupMembers](Instruction *I) {
-        BasicBlock *BB = I->getParent();
-        for (auto &I : make_range(I->getIterator(), BB->end()))
-          if (!DeadInterleaveGroupMembers.contains(&I))
-            return &I;
-        llvm_unreachable("Need to find a valid insert point");
-      };
-  // Mark instructions we'll need to sink later and their targets as
-  // ingredients whose recipe we'll need to record.
-  for (auto &Entry : SinkAfter) {
-    RecipeBuilder.recordRecipeOf(skipDeadInterleaveMembers(Entry.first));
-    RecipeBuilder.recordRecipeOf(skipDeadInterleaveMembers(Entry.second));
-  }
-
   // ---------------------------------------------------------------------------
   // Build initial VPlan: Scan the body of the loop in a topological order to
   // visit each basic block after having visited its predecessor basic blocks.
   // ---------------------------------------------------------------------------
 
   // Create a dummy pre-entry VPBasicBlock to start building the VPlan.
   auto Plan = std::make_unique<VPlan>();
   VPBasicBlock *VPBB = new VPBasicBlock("Pre-Entry");
VPlanPtr LoopVectorizationPlanner::convertToVPRecipes(
   VPBlockUtils::disconnectBlocks(PreEntry, Entry);
   delete PreEntry;
 
   // ---------------------------------------------------------------------------
   // 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(skipDeadInterleaveMembers(Entry.first));
-    VPRecipeBase *Target =
-        RecipeBuilder.getRecipe(skipDeadInterleaveMembers(Entry.second));
-    Sink->moveAfter(Target);
-  }
-
   // Adjust the recipes for any inloop reductions.
   if (Range.Start > 1)
     adjustRecipesForInLoopReductions(Plan, RecipeBuilder);
 
   // Finally, if tail is folded by masking, introduce selects between the phi
   // and the live-out instruction of each reduction, at the end of the latch.
   if (CM.foldTailByMasking() && !Legal->getReductionVars().empty()) {
     Builder.setInsertPoint(VPBB);