[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);
-
-  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);
   void buildVPlansWithVPRecipes(ElementCount MinVF, ElementCount 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(ElementCount MinVF,
   DenseMap<Instruction *, Instruction *> &SinkAfter = Legal->getSinkAfter();
   // Dead instructions do not need sinking. Remove them from SinkAfter.
   for (Instruction *I : DeadInstructions)
     SinkAfter.erase(I);
 
   ReversePostOrderTraversal<VPBlockBase *> RPOT(
       OriginalPlan.getEntry()->getEntryBasicBlock());
 
+  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) {
     VPBasicBlock *OriginalVPBB = Base->getEntryBasicBlock();
     // Introduce each ingredient into VPlan.
     for (auto I = OriginalVPBB->rbegin(), E = OriginalVPBB->rend(); I != E;) {
       VPRecipeBase *Ingredient = &*I++;
       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();
+        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);
+  }
+
   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) {
 
   // 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;
 
   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;