[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(ElementCount MinVF, ElementCount 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<Instruction *> &DeadInstructions,
-      const DenseMap<Instruction *, Instruction *> &SinkAfter);
-
   /// Returns a new VPlan based on \p OriginalPlan, but with all VPInstruction
   /// in \p OriginalPlan lowered to widened recipes.
-  VPlanPtr
-  convertToVPRecipes(VPlan &OriginalPlan, VFRange &Range,
-                     const DenseMap<Instruction *, Instruction *> &SinkAfter);
+  VPlanPtr convertToVPRecipes(VPlan &OriginalPlan, VFRange &Range);
 
   /// 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(ElementCount MinVF, ElementCount MaxVF,
                                 VPlan &OriginalPlan);
 
   /// Adjust the recipes for any inloop reductions. The chain of instructions

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

void LoopVectorizationPlanner::buildVPlansWithVPRecipes(ElementCount MinVF,
   auto &ConditionalAssumes = Legal->getConditionalAssumes();
   DeadInstructions.insert(ConditionalAssumes.begin(), ConditionalAssumes.end());
 
   DenseMap<Instruction *, Instruction *> &SinkAfter = Legal->getSinkAfter();
   // Dead instructions do not need sinking. Remove them from SinkAfter.
   for (Instruction *I : DeadInstructions)
     SinkAfter.erase(I);
 
+  DenseMap<Instruction *, VPInstruction *> Inst2VPInst;
+  for (auto &Entry : SinkAfter) {
+    Inst2VPInst[Entry.first] = nullptr;
+    Inst2VPInst[Entry.second] = nullptr;
+  }
+
   // Remove dead instructions from original VPlan.
   auto *DummyVal = new VPValue();
   OriginalPlan.addExternalDef(DummyVal);
   for (VPBlockBase *Base :
        depth_first<VPBlockBase *>(TopRegion->getEntryBasicBlock())) {
     VPBasicBlock *OriginalVPBB = Base->getEntryBasicBlock();
     for (VPRecipeBase &Ingredient :
          make_early_inc_range(reverse(*OriginalVPBB))) {
       VPInstruction *VPInst = dyn_cast<VPInstruction>(&Ingredient);
       if (!VPInst)
         continue;
       Instruction *Instr = VPInst->getUnderlyingInstr();
 
       if (DeadInstructions.contains(Instr) || isa<DbgInfoIntrinsic>(Instr)) {
         Ingredient.getVPValue()->replaceAllUsesWith(DummyVal);
         Ingredient.eraseFromParent();
       }
-    }
+
+      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);
   }
 
   auto MaxVFPlusOne = MaxVF.getWithIncrement(1);
   for (ElementCount VF = MinVF; ElementCount::isKnownLT(VF, MaxVFPlusOne);) {
     VFRange SubRange = {VF, MaxVFPlusOne};
-    VPlans.push_back(convertToVPRecipes(OriginalPlan, SubRange, SinkAfter));
+    VPlans.push_back(convertToVPRecipes(OriginalPlan, SubRange));
     VF = SubRange.End;
   }
 }
 
-VPlanPtr LoopVectorizationPlanner::convertToVPRecipes(
-    VPlan &OriginalPlan, VFRange &Range,
-    const DenseMap<Instruction *, Instruction *> &SinkAfter) {
+VPlanPtr LoopVectorizationPlanner::convertToVPRecipes(VPlan &OriginalPlan,
+                                                      VFRange &Range) {
 
   SmallPtrSet<const InterleaveGroup<Instruction> *, 1> InterleaveGroups;
 
   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;
     RecurKind Kind = Legal->getReductionVars()[Phi].getRecurrenceKind();
     const SmallVector<Instruction *, 4> &ReductionOperations = Reduction.second;
 
     RecipeBuilder.recordRecipeOf(Phi);
     for (auto &R : ReductionOperations) {
       RecipeBuilder.recordRecipeOf(R);
for (VPBlockBase *Base : RPOT) {
     }
   }
 
   // ---------------------------------------------------------------------------
   // 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;
-      }
-    }
-    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;