[VPlan] Fix in-loop reduction chains using VPlan def-use chains (NFCI)

[fhahn]

Summary

Update adjustRecipesForReductions to directly use the VPlan def-use
chains for in-loop reductions to collect the reduction operations that
need adjusting.

This allows the removal of

• ReductionChainMap
• recording of recipes for instruction in the reduction chain
• removes late uses of getVPValue
• removes to need for removeVPValueFor.

[Ayal]

Thanks for following-up on this nice cleanup!
Looking over the code raises various initial thoughts, which could be pursued in several steps/patches.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1611

1775–1777

1782–1783

9152

?
Can also filter out non-isOrdered phi's when VF=1, to then iterate only over phi's of interest.

9156

Bypass this for loop if MinVF.isScalar()?

9166–9167

Would one SetVector work?

This represents the Chain - should it hold Recipes as Links, each residing inside the loop and having a single user (except last with live-out), rather than VPValues?

9170

Worth a comment explaining what's going to happen next.

9183–9186

9189

Above initialization of Worklist tolerates recipes with multiple defined values, but here only single valued recipes are expected?

9192

Hoist Kind and its assert for being a select-compare kind next to setting RdxDesc above?

9195

IndexOfFirstOperand?
Initialize? Comment what this variable holds? (Admittedly missing from current code base.)

9205–9206

Dead assert.

Worth a comment why we continue here - to skip over the compare of the select?
What about Fminimum/Fmaximum reductions which are isMinMaxRecurrenceKind but have a call to an intrinsic as their WidenRecipe, rather than a compare-select?

9211

comment appears out of date.

9219

9225

9228

Is this setting of Chain needed?

9251

Can this setting of Chain continue to be placed at the end of the loop below?

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9187

9188

Could this for loop be replaced by a simpler for (auto Link : Chain) { loop?

9189

9190

9193

9197

Comment deserves update - FMulAdd is also supported.
Potentially renamed isAnyOfRecurrenceKind() in D155786.

9216

It's somewhat alarming that we're indifferent to which operand is VecOp, in non commutative cases such as compare-select. The semantics of the compare/select are presumably captured by the min/max kind, rather than the condition itself (or its negation). Worth a comment, and ensuring test coverage of having VecOp in both positions.

9219

Or rather BasicBlock *CurrentLinkBB = CurrentLinkI->getParent();

Can also define and reuse: VPBasicBlock *CurrentLinkVPBB = CurrentLink->getParent();

9232

Should this breaking of fmuladd into an fadd recurrence fed by an fmul take place separately, as a preparatory step?
Should compare-select pairs be replaced by Min/Max recipes?
In order to convert an (initially widened) reduction recurrence into an in-loop one, it seems that isolating the reduction operation into a single, dedicated operation would be useful. It could also help canonicalize where to find the reduced value operands, aka vecOps.

9241

9262

Why/Is this needed?

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9156

Bypass this for loop if MinVF.isScalar()?

Ah, isOrdered reductions are considered inLoop even for VF=1 (but only them).

[fhahn]

Address latest comments, thanks! I tried to include most adjustments directly in this patch to keep things together.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

1611

Adjusted, thanks!

1775–1777

Updated, thanks!

1782–1783

Updated, thanks!

9152

Adjusted and collect only VPReductionPHIRecipe

9166–9167

Updated to use SetVector and use VPRecipeBase as element type, thanks!

9170

added a comment, thanks!

9183–9186

Updated, Thanks!

9188

Updated to use iterator based loop which skips the first element (aka PreviousLink)

9189

Updated, thanks!

Above initialization of Worklist tolerates recipes with multiple defined values, but here only single valued recipes are expected?

Added an assert in the collection loop

9190

Updated, thanks!

9192

Hoisted out of loop, thanks!

9193

Updated , thanks!

9195

Renamed & added comment. It should be guaranteed to be initialized before use.

9197

Updated the message; any reduction should be allowed, besides select/cmp ones.

9205–9206

Will update to have a check for the cmp only. In-loop reductions won't be used for those intrinsics at the moment. Will make sure there's a test. Added tests with min-max intrinsics in cc3986643696, for which in-loop reductions aren't applied yet

9211

Updated to also include FMulAdd

9216

Yep, that looks right, added a comment, thanks!

Added test coverage with flipped operands in cc3986643696.

9219

Updated, thanks!

9225

Updated, thanks!

9228

No, removed ,thanks!

9232

Might be good as follow up?

9251

Yep, moved, thanks!

9262

Not needed in the latest version, removed!

[Ayal]

Added some more comments to help clean this up more, can be applied now or later if preferred to keep the diff closer to a VPlan refactoring alone.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9164

nit: Should this filtering be applied above when inserting into InLoopReductionPhis?
Should unordered reductions be considered non-isInLoop in VPlans of VF=1?

9206

note: may assert that UserRecipe resides inside the vector loop, as the outside user is a LiveOut, at-least for now.

9208

9211

assert CurrentLink is a select recipe (whose first operand index is 1)?

9213

Treat FMulAdd separately from the binary operators - the operand along the chain is in index 2, go ahead and generate the FMul of indices 0 and 1 to produce VecOp? (Keeping the assert message intact ;-)

9215

nit: assert better placed above, asap, independent of filling Worklist.

9217

assert that VecOp != PreviousLinkV and that the other candidate does equal PreviousLinkV?
(CurrentLink->getOperand(IndexOfFirstOperand+IndexOfFirstOperand+1-VecOpId))

9221

Wonder if this information may already be recorded in VPlan - perhaps BlockInMask already exists?

9228

nit: would always considering the last two indices be (correct and) easier than considering either first two indices or indices 1 and 2? Excluding FMulAdd which deserves a separate treatment.

9231

9232

Sure, if not two.

9233–9236

assert VecOpId == 0?

9244

Could we use PreviousLink serve as/instead of CompareRecipe?

[fhahn]

Address latest comments, thanks!

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9164

Moved the filtering. Left scalar handling as is for now to keep with original code.

9206

Updated, thanks!

9208

Updated, thanks!

9211

added assert, thanks!

9213

Moved, thanks!

9215

Moved up, thanks!

9217

Added assert, thanks!

9221

Not easily yet, but I think as follow-up this could be handled by checking the predecessors of the block to find the predecessor with branch-on-mask.

9228

For the initial version, I think it would be good retain the original logic as much as possible. This also makes it easier to retain the asserts.

9233–9236

Not needed after code movement.

9244

Unfortunately I don't think so, but we can let dead-recipe removal take care of this removal. I removed the code.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9170–9177

9186

?

9192

9194

9210

If VF=1 CurrentLink should be a Replicate instead of a Widen(Call) recipe?

Better place the above assert(RecurrenceDescriptor::isFMulAddIntrinsic(CurrentLinkI)) here?

9244

If PreviousLink is set to Cmp before early-continuing above?

OTOH, dead-recipe removal could also reclaim CurrentLink.

[fhahn]

Address latest comments, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9170–9177

Reordered, thanks!

9186

Removed, thanks!

9192

Fixed, thanks!

9194

added, thanks!

9210

Moved assert and refined condition as suggested, thanks!

9231

Adjusted, thanks!

9244

Yep, updated, thanks!

[Ayal]

This looks good to me, thanks!
Adding last minor nits.

Would be good to see how adjustRecipesForReductions() may become a VPlan2VPlan pass, and apply fmuladd decomposition / minmax composition as separate, preparatory steps.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9216

Note: in this case PreviousLink stays the recipe before the compare, feeding one of the operands of the select.

9237

9244

nit: may be worth noting somewhere that this transformation may leave dead recipes for a subsequent pass to clean up.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9237

Sounds better, thanks!

9244

Added a note in the committed version, thanks!

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

public:
 
   /// Collect values we want to ignore in the cost model.
   void collectValuesToIgnore();
 
   /// Collect all element types in the loop for which widening is needed.
   void collectElementTypesForWidening();
 
   /// Split reductions into those that happen in the loop, and those that happen
-  /// outside. In loop reductions are collected into InLoopReductionChains.
+  /// outside. In loop reductions are collected into InLoopReductions.
   void collectInLoopReductions();
 
   /// Returns true if we should use strict in-order reductions for the given
   /// RdxDesc. This is true if the -enable-strict-reductions flag is passed,
   /// the IsOrdered flag of RdxDesc is set and we do not allow reordering
   /// of FP operations.
   bool useOrderedReductions(const RecurrenceDescriptor &RdxDesc) const {
     return !Hints->allowReordering() && RdxDesc.isOrdered();
public:
 
   /// Returns true if the instructions in this block requires predication
   /// for any reason, e.g. because tail folding now requires a predicate
   /// or because the block in the original loop was predicated.
   bool blockNeedsPredicationForAnyReason(BasicBlock *BB) const {
     return foldTailByMasking() || Legal->blockNeedsPredication(BB);
   }
 
-  /// A SmallMapVector to store the InLoop reduction op chains, mapping phi
-  /// nodes to the chain of instructions representing the reductions. Uses a
-  /// MapVector to ensure deterministic iteration order.
-  using ReductionChainMap =
-      SmallMapVector<PHINode *, SmallVector<Instruction *, 4>, 4>;
-
-  /// Return the chain of instructions representing an inloop reduction.
-  const ReductionChainMap &getInLoopReductionChains() const {
-    return InLoopReductionChains;
-  }
-
   /// Returns true if the Phi is part of an inloop reduction.
   bool isInLoopReduction(PHINode *Phi) const {
-    return InLoopReductionChains.count(Phi);
+    return InLoopReductions.count(Phi);
   }
 
   /// Estimate cost of an intrinsic call instruction CI if it were vectorized
   /// with factor VF.  Return the cost of the instruction, including
   /// scalarization overhead if it's needed.
   InstructionCost getVectorIntrinsicCost(CallInst *CI, ElementCount VF) const;
 
   /// Estimate cost of a call instruction CI if it were vectorized with factor
private:
   /// Holds the instructions known to be scalar after vectorization.
   /// The data is collected per VF.
   DenseMap<ElementCount, SmallPtrSet<Instruction *, 4>> Scalars;
 
   /// Holds the instructions (address computations) that are forced to be
   /// scalarized.
   DenseMap<ElementCount, SmallPtrSet<Instruction *, 4>> ForcedScalars;
 
   /// PHINodes of the reductions that should be expanded in-loop along with
   /// their associated chains of reduction operations, in program order from top
   /// (PHI) to bottom
-  ReductionChainMap InLoopReductionChains;
+  SmallPtrSet<PHINode *, 4> InLoopReductions;
 
   /// A Map of inloop reduction operations and their immediate chain operand.
   /// FIXME: This can be removed once reductions can be costed correctly in
   /// vplan. This was added to allow quick lookup to the inloop operations,
   /// without having to loop through InLoopReductionChains.
   DenseMap<Instruction *, Instruction *> InLoopReductionImmediateChains;
 
   /// Returns the expected difference in cost from scalarizing the expression
   /// feeding a predicated instruction \p PredInst. The instructions to
   /// scalarize and their scalar costs are collected in \p ScalarCosts. A
   /// non-negative return value implies the expression will be scalarized.
   /// Currently, only single-use chains are considered for scalarization.
   InstructionCost computePredInstDiscount(Instruction *PredInst,
LoopVectorizationCostModel::getInterleaveGroupCost(Instruction *I,
   return Cost;
 }
 
 std::optional<InstructionCost>
 LoopVectorizationCostModel::getReductionPatternCost(
     Instruction *I, ElementCount VF, Type *Ty, TTI::TargetCostKind CostKind) {
   using namespace llvm::PatternMatch;
   // Early exit for no inloop reductions
-  if (InLoopReductionChains.empty() || VF.isScalar() || !isa<VectorType>(Ty))
+  if (InLoopReductions.empty() || VF.isScalar() || !isa<VectorType>(Ty))
     return std::nullopt;
   auto *VectorTy = cast<VectorType>(Ty);
 
   // We are looking for a pattern of, and finding the minimal acceptable cost:
   //  reduce(mul(ext(A), ext(B))) or
   //  reduce(mul(A, B)) or
   //  reduce(ext(A)) or
   //  reduce(A).
if (!PreferInLoopReductions && !useOrderedReductions(RdxDesc) &&
       continue;
 
     // Check that we can correctly put the reductions into the loop, by
     // finding the chain of operations that leads from the phi to the loop
     // exit value.
     SmallVector<Instruction *, 4> ReductionOperations =
         RdxDesc.getReductionOpChain(Phi, TheLoop);
     bool InLoop = !ReductionOperations.empty();
+
     if (InLoop) {
-      InLoopReductionChains[Phi] = ReductionOperations;
+      InLoopReductions.insert(Phi);
       // Add the elements to InLoopReductionImmediateChains for cost modelling.
       Instruction *LastChain = Phi;
       for (auto *I : ReductionOperations) {
         InLoopReductionImmediateChains[I] = LastChain;
         LastChain = I;
       }
     }
     LLVM_DEBUG(dbgs() << "LV: Using " << (InLoop ? "inloop" : "out of loop")
std::optional<VPlanPtr> LoopVectorizationPlanner::tryToBuildVPlanWithVPRecipes(
 
   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.
   // ---------------------------------------------------------------------------
 
-  for (const auto &Reduction : CM.getInLoopReductionChains()) {
-    PHINode *Phi = Reduction.first;
-    RecurKind Kind =
-        Legal->getReductionVars().find(Phi)->second.getRecurrenceKind();
-    const SmallVector<Instruction *, 4> &ReductionOperations = Reduction.second;
-
-    RecipeBuilder.recordRecipeOf(Phi);
-    for (const auto &R : ReductionOperations) {
-      RecipeBuilder.recordRecipeOf(R);
-      // For min/max reductions, where we have a pair of icmp/select, we also
-      // need to record the ICmp recipe, so it can be removed later.
-      assert(!RecurrenceDescriptor::isSelectCmpRecurrenceKind(Kind) &&
-             "Only min/max recurrences allowed for inloop reductions");
-      if (RecurrenceDescriptor::isMinMaxRecurrenceKind(Kind))
-        RecipeBuilder.recordRecipeOf(cast<Instruction>(R->getOperand(0)));
-    }
-  }
-
   // For each interleave group which is relevant for this (possibly trimmed)
   // Range, add it to the set of groups to be later applied to the VPlan and add
   // placeholders for its members' Recipes which we'll be replacing with a
   // single VPInterleaveRecipe.
   for (InterleaveGroup<Instruction> *IG : IAI.getInterleaveGroups()) {
     auto applyIG = [IG, this](ElementCount VF) -> bool {
       bool Result = (VF.isVector() && // Query is illegal for VF == 1
                      CM.getWideningDecision(IG->getInsertPos(), VF) ==

 // Adjust the recipes for reductions. For in-loop 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. For other reductions, a select is introduced between the phi
 // and live-out recipes when folding the tail.
 void LoopVectorizationPlanner::adjustRecipesForReductions(
     VPBasicBlock *LatchVPBB, VPlanPtr &Plan, VPRecipeBuilder &RecipeBuilder,
     ElementCount MinVF) {
-  for (const auto &Reduction : CM.getInLoopReductionChains()) {
-    PHINode *Phi = Reduction.first;
-    const RecurrenceDescriptor &RdxDesc =
-        Legal->getReductionVars().find(Phi)->second;
-    const SmallVector<Instruction *, 4> &ReductionOperations = Reduction.second;
+  SmallVector<VPRecipeBase *> Phis;
+  for (VPRecipeBase &R :
+       Plan->getVectorLoopRegion()->getEntryBasicBlock()->phis())
+    Phis.push_back(&R);
 
-    if (MinVF.isScalar() && !CM.useOrderedReductions(RdxDesc))
+  for (VPRecipeBase *R : Phis) {
+    auto *PhiR = dyn_cast<VPReductionPHIRecipe>(R);
+    if (!PhiR || !PhiR->isInLoop())
       continue;
+    const RecurrenceDescriptor &RdxDesc = PhiR->getRecurrenceDescriptor();
+
+    if (MinVF.isScalar() && !PhiR->isOrdered())
+      continue;
+
+    SmallVector<VPValue *, 8> Worklist;
+    SmallPtrSet<VPValue *, 8> Visited;
+    Worklist.push_back(PhiR);
+    Visited.insert(PhiR);
+
+    for (unsigned I = 0; I != Worklist.size(); ++I) {
+      VPValue *Cur = Worklist[I];
+      for (VPUser *U : Cur->users()) {
+        auto *UserRecipe = dyn_cast<VPRecipeBase>(U);
+        if (!UserRecipe)
+          continue;
+        for (VPValue *V : UserRecipe->definedValues())
+          if (Visited.insert(V).second)
+            Worklist.push_back(V);
+      }
+    }
 
     // ReductionOperations are orders top-down from the phi's use to the
     // LoopExitValue. We keep a track of the previous item (the Chain) to tell
     // which of the two operands will remain scalar and which will be reduced.
     // For minmax the chain will be the select instructions.
-    Instruction *Chain = Phi;
-    for (Instruction *R : ReductionOperations) {
-      VPRecipeBase *WidenRecipe = RecipeBuilder.getRecipe(R);
+    VPRecipeBase *Chain = PhiR;
+    for (unsigned I = 1; I != Worklist.size(); ++I) {
+      VPValue *ChainOp = Chain->getVPSingleValue();
+      VPRecipeBase *WidenRecipe = Worklist[I]->getDefiningRecipe();
+
       RecurKind Kind = RdxDesc.getRecurrenceKind();
+      Instruction *R = WidenRecipe->getUnderlyingInstr();
 
-      VPValue *ChainOp = Plan->getVPValue(Chain);
       unsigned FirstOpId;
       assert(!RecurrenceDescriptor::isSelectCmpRecurrenceKind(Kind) &&
              "Only min/max recurrences allowed for inloop reductions");
       // Recognize a call to the llvm.fmuladd intrinsic.
       bool IsFMulAdd = (Kind == RecurKind::FMulAdd);
       assert((!IsFMulAdd || RecurrenceDescriptor::isFMulAddIntrinsic(R)) &&
              "Expected instruction to be a call to the llvm.fmuladd intrinsic");
       if (RecurrenceDescriptor::isMinMaxRecurrenceKind(Kind)) {
+        if (!isa<VPWidenSelectRecipe>(WidenRecipe))
+          continue;
         assert(isa<VPWidenSelectRecipe>(WidenRecipe) &&
                "Expected to replace a VPWidenSelectSC");
         FirstOpId = 1;
       } else {
         assert((MinVF.isScalar() || isa<VPWidenRecipe>(WidenRecipe) ||
                 (IsFMulAdd && isa<VPWidenCallRecipe>(WidenRecipe))) &&
                "Expected to replace a VPWidenSC");
         FirstOpId = 0;
       }
-      unsigned VecOpId =
-          R->getOperand(FirstOpId) == Chain ? FirstOpId + 1 : FirstOpId;
-      VPValue *VecOp = Plan->getVPValue(R->getOperand(VecOpId));
+      unsigned VecOpId = WidenRecipe->getOperand(FirstOpId) == ChainOp
+                             ? FirstOpId + 1
+                             : FirstOpId;
+      VPValue *VecOp = WidenRecipe->getOperand(VecOpId);
 
+      BasicBlock *BB = WidenRecipe->getUnderlyingInstr()->getParent();
       VPValue *CondOp = nullptr;
-      if (CM.blockNeedsPredicationForAnyReason(R->getParent())) {
+      if (CM.blockNeedsPredicationForAnyReason(BB)) {
         VPBuilder::InsertPointGuard Guard(Builder);
         Builder.setInsertPoint(WidenRecipe->getParent(),
                                WidenRecipe->getIterator());
         CondOp = RecipeBuilder.createBlockInMask(R->getParent(), *Plan);
       }
 
+      Chain = WidenRecipe;
       if (IsFMulAdd) {
         // If the instruction is a call to the llvm.fmuladd intrinsic then we
         // need to create an fmul recipe to use as the vector operand for the
         // fadd reduction.
         VPInstruction *FMulRecipe = new VPInstruction(
-            Instruction::FMul, {VecOp, Plan->getVPValue(R->getOperand(1))});
+            Instruction::FMul, {VecOp, WidenRecipe->getOperand(1)});
         FMulRecipe->setFastMathFlags(R->getFastMathFlags());
         WidenRecipe->getParent()->insert(FMulRecipe,
                                          WidenRecipe->getIterator());
         VecOp = FMulRecipe;
       }
       VPReductionRecipe *RedRecipe =
-          new VPReductionRecipe(&RdxDesc, R, ChainOp, VecOp, CondOp, &TTI);
-      WidenRecipe->getVPSingleValue()->replaceAllUsesWith(RedRecipe);
-      Plan->removeVPValueFor(R);
-      Plan->addVPValue(R, RedRecipe);
+          new VPReductionRecipe(&RdxDesc, WidenRecipe->getUnderlyingInstr(),
+                                ChainOp, VecOp, CondOp, &TTI);
+      VPRecipeBase *CompareRecipe =
+          WidenRecipe->getOperand(0)->getDefiningRecipe();
+
       // Append the recipe to the end of the VPBasicBlock because we need to
       // ensure that it comes after all of it's inputs, including CondOp.
       WidenRecipe->getParent()->appendRecipe(RedRecipe);
       WidenRecipe->getVPSingleValue()->replaceAllUsesWith(RedRecipe);
       WidenRecipe->eraseFromParent();
+      Chain = RedRecipe;
 
       if (RecurrenceDescriptor::isMinMaxRecurrenceKind(Kind)) {
-        VPRecipeBase *CompareRecipe =
-            RecipeBuilder.getRecipe(cast<Instruction>(R->getOperand(0)));
         assert(isa<VPWidenRecipe>(CompareRecipe) &&
                "Expected to replace a VPWidenSC");
         assert(cast<VPWidenRecipe>(CompareRecipe)->getNumUsers() == 0 &&
                "Expected no remaining users");
         CompareRecipe->eraseFromParent();
       }
-      Chain = R;
+
+      if (Chain == PhiR)
+        break;
     }
   }
 
   // If tail is folded by masking, introduce selects between the phi
   // and the live-out instruction of each reduction, at the beginning of the
   // dedicated latch block.
   if (CM.foldTailByMasking()) {
     Builder.setInsertPoint(LatchVPBB, LatchVPBB->begin());

llvm/lib/Transforms/Vectorize/VPlan.h

if (!Value2VPValue.count(V)) {
       VPValue *VPV = new VPValue(V);
       VPLiveInsToFree.push_back(VPV);
       addVPValue(V, VPV);
     }
 
     return getVPValue(V);
   }
 
-  void removeVPValueFor(Value *V) {
-    assert(Value2VPValueEnabled &&
-           "IR value to VPValue mapping may be out of date!");
-    Value2VPValue.erase(V);
-  }
-
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
   /// Print this VPlan to \p O.
   void print(raw_ostream &O) const;
 
   /// Print this VPlan in DOT format to \p O.
   void printDOT(raw_ostream &O) const;
 
   /// Dump the plan to stderr (for debugging).