[VPlan] Iterate over phi recipes to detect reductions to fix.

[fhahn]

Summary

After refactoring the phi recipes, we can now iterate over all header
phis in a VPlan to detect reductions when it comes to fixing them up
when tail folding.

This reduces the coupling with the cost model & legal by using the
information directly available in VPlan. It also removes a call to
getOrAddVPValue, which references the original IR value which may
become outdated after VPlan transformations.

[fhahn]

ping :)

[Ayal]

Reducing the coupling with the cost model by using information directly in VPlan is great. Have a few comments.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4307

number of operands must be 2, right?

this may be fragile, if some (cast?) vpinstruction is introduced between the VPReductionRecipe and the header phi it feeds? Since in-loop-reduction phi's require special code-gen handling, perhaps their recipes should be marked as such?

4722

ditto

9170–9171

Perhaps RecipeBuilder should allow checking if it has any (in-loop) reduction phi's, and iterate over them?

[fhahn]

Remove unnecessary check on number of operands, add assertion to catch if we miss an in-loop reduction in VPlan.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4307

number of operands must be 2, right?

It's not needed here, I removed it.

this may be fragile, if some (cast?) vpinstruction is introduced between the VPReductionRecipe and the header phi it feeds? Since in-loop-reduction phi's require special code-gen handling, perhaps their recipes should be marked as such?

Agreed, but not sure if we need to fix this issue before landing the patch?

I added an assertion in the code below to catch the case where we fail to identify an in-loop reduction in VPlan.

4722

here PhiR could also be a first order recurrence, so for now I think we need to at least make sure it's a reduction.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4307

wouldn't it be simpler, clearer and more robust if adjustRecipesForInLoopReductions() would turn on an "IsInLoop" bit inside VPWidenPhiRecipe, when converting its associated VPWiden[Select]Recipes into VPReductionRecipes?

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4307

It would be simpler, but I'm not sure if it is worth adding this extra coupling by adding a new InLoopReduction field to the general VPWidenPHIRecipe that needs to be maintained at this point? We could also just traverse the whole cycle if needed I think as in practice it should be only a very small number of steps. WDYT?

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4307

Agreed - the overhead of checking the cycle/chain to figure out if the reduction is in-loop or not, should be tolerable. It's about having recipes clearly model the code they represent up-front, during VPlan construction and/or planning, than whether to cache the results of this "isInLoop" query eagerly versus (re)running it lazily on demand during VPlan execution. A phi recipe representing an in-loop reduction is arguably conceptually distinct from a "regular" VPWidenPHIRecipe, with respect to the scalar-versus-vector data each operates on, and the epilogue each feeds. If the "cached result" needs to be updated, so may the cost of the recipe and of dependent recipes. Sounds reasonable?

4722

Agreed; "ditto" referred to asking here too directly if PhiR->isInLoopReduction().

9177

ditto ;-)

9180

This is to avoid redundant creation of the header block in-mask when the loop contains no non-in-loop reductions?
Worth doing in a separate patch?

9188

#ifndef NDEBUG?

perhaps an overkill - would CM.getInLoopReductionChains() eventually be discarded?

[fhahn]

Thanks for the latest comments!

This is just a rebase, will address comments separately!

[fhahn]

Add isInLoopReduction helper & use it instead isa<VPReductionRecipe>

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4307

I went ahead and added a VPWidenPHIRecipe::isInLoopReduction helper.

I'm not 100% sure this is in line with your suggestions, please let me know if I misunderstood and we should add a flag as well to start with :)

9188

Yes I think CM.getInLoopReductionChains is going away at some point. But until then it might be good to try to ensure they do not diverge. But I'd be more than happy to drop it as well, it might just be a bit too paranoid

[fhahn]

Clang-format.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4307

This API looks good to me, thanks!
Adding a flag to implement it, as argued above, would also be good.
As stated elsewhere, there are actually 3 types of reductions, including two types of in-loop reductions: ordered and unordered.

9177

nit: suffice to check if (!PhiR || PhiR->isInLoopReduction()) - the latter implies PhiR has a recurrence descriptor.

9180

There's a subtle change here: BlockInMask may currently be created redundantly here if there are reductions but no in-loop reductions; this modification eliminates this redundancy.
Probably insignificant, given foldTailByMasking?

9188

I'd vote for cleaner, paranoia free code in this case.

One could verify the integrity of in-loop reduction chains (or VPlan in general?), say before code-gen; i.e., that in-loop header phi recipes are consistent with reduction recipes.

9257

Wouldn't it be better, if in addition to converting widen recipes of in-loop-reduction chains into reduction recipes here, we also convert their header phi upfront:

auto *PhiR = cast<VPWidenPHIRecipe>(RecipeBuilder.getRecipe(Phi));
PhiR->setIsInLoopReduction();

?

(Making sure to record the recipe of Phi)

llvm/lib/Transforms/Vectorize/VPlan.cpp

1076

(being 'const' it can't even cache the result for repeated queries...)

[fhahn]

rebased after the other phi handling refactoring patches landed. The patch now just updates the tail folding handling for reduction values to iterate over the phis in the VPlan and use VPReductionPHIRecipe directly. This reduces some late callbacks to Legal and the cost model, as well as a call to getOrAddVPValue(), which could become out-of-date after VPlan transformations.

ping :)

[Ayal]

This does look pretty simpler now :-). Adding some more thoughts, could be addressed separately.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9166–9167

Perhaps wrap the following insert-select-for-non-in-loop-reduction-under-fold-tail along with the above, under adjustRecipesForReductions()? buildVPlanWithVPRecipes() has grown out of propotion.

9187–9190

Checking if (!Cond) seems redundant; if we want to avoid creating BlockInMask for the header block here unless there is at-least one reduction (although it will probably be created anyhow, especially under foldTail), may as well call createBlockInMask() repeatedly - it provides internal caching. I.e., should be renamed more accurately getOrCreateBlockInMask().

9190

VPValue *Red = PhiR->getBackedgeValue ();?

Would be good to note that (and remember how ;-) ILV::fixReduction() later hooks up to this select, as it seems pretty useless here.

[fhahn]

Address Ayal's comments, thanks!

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9187–9190

Checking if (!Cond) seems redundant; if we want to avoid creating BlockInMask for the header block here unless there is at-least one reduction (although it will probably be created anyhow, especially under foldTail), may as well call createBlockInMask() repeatedly - it provides internal caching. I.e., should be renamed more accurately getOrCreateBlockInMask().

Right, I just updated the code to unconditionally create the condition, relying on the caching.

Perhaps wrap the following insert-select-for-non-in-loop-reduction-under-fold-tail along with the above, under adjustRecipesForReductions()? buildVPlanWithVPRecipes() has grown out of propotion.

Done!

9190

I updated the code as suggested, but is the loop exit value always the backedge value? I added a note that the select here is only to model the def-use chain in VPlan (AFAIU)

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9190

is the loop exit value always the backedge value?

It appears so, taken from IVDescriptors.cpp/RecurrenceDescriptor::AddReductionVar():

// The instruction used by an outside user must be the last instruction
// before we feed back to the reduction phi. Otherwise, we loose VF-1
// operations on the value.
if (!is_contained(Phi->operands(), Cur))
  return false;
  
ExitInstruction = Cur;
continue;

Otherwise the generated the select between exit-instruction and phi may be erroneous.
It would be clearer if getLoopExitInstr() would return Phi's operand across the backedge, instead of recording ExitInstruction.

9197

Thanks! Rename adjustRecipesForInLoopReductions() to adjustRecipesForReductions()?

Would be good to iterate over all header (reduction) phi's once, checking for inLoop or not. Note that the select handles only non-inLoop reductions, under foldTail. Deserves a separate patch?

9314

Ah, ILV::fixReduction() actually looks for this select and hooks up to it, under foldTail:

Instruction *LoopExitInst = RdxDesc.getLoopExitInstr();
...
VPValue *LoopExitInstDef = State.Plan->getVPValue(LoopExitInst);
...
if (Cost->foldTailByMasking() && !PhiR->isInLoop()) {
  for (unsigned Part = 0; Part < UF; ++Part) {
    Value *VecLoopExitInst = State.get(LoopExitInstDef, Part);
    Value *Sel = nullptr;
    for (User *U : VecLoopExitInst->users()) {
      if (isa<SelectInst>(U)) {
        assert(!Sel && "Reduction exit feeding two selects");
        Sel = U;
      } else
        assert(isa<PHINode>(U) && "Reduction exit must feed Phi's or select");
    }
    assert(Sel && "Reduction exit feeds no select");
    State.reset(LoopExitInstDef, Sel, Part);
  
    // If the target can create a predicated operator for the reduction at no
    // extra cost in the loop (for example a predicated vadd), it can be
    // cheaper for the select to remain in the loop than be sunk out of it,
    // and so use the select value for the phi instead of the old
    // LoopExitValue.
    if (PreferPredicatedReductionSelect ||
        TTI->preferPredicatedReductionSelect(
            RdxDesc.getOpcode(), PhiTy,
            TargetTransformInfo::ReductionFlags())) {
      auto *VecRdxPhi =
          cast<PHINode>(State.get(PhiR->getVPSingleValue(), Part));
      VecRdxPhi->setIncomingValueForBlock(
          LI->getLoopFor(LoopVectorBody)->getLoopLatch(), Sel);
    }
  }
}

Simply take PhiR->getBackedgeValue() instead there as well?

PreferPredicatedReductionSelect should also hack the phi recipe and be taken care of by VPlan execute? (Here or in a separate patch.)

[fhahn]

Address comments: fix name, use getBackedgeValue at an additional place.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9190

Ok great!

9197

Thanks! Rename adjustRecipesForInLoopReductions() to adjustRecipesForReductions()?

Ah thanks, I thought I renamed it, but was wrong. should be updated now.

Would be good to iterate over all header (reduction) phi's once, checking for inLoop or not. Note that the select handles only non-inLoop reductions, under foldTail. Deserves a separate patch?

Can update as separate patch to keep the diff cleaner.

9314

Simply take PhiR->getBackedgeValue() instead there as well?

Should be updated there, thanks!

PreferPredicatedReductionSelect should also hack the phi recipe and be taken care of by VPlan execute? (Here or in a separate patch.)

I can do that in a separate one.

[Ayal]

Thanks! This looks good to me, with a minor comment comment.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

9313

Update comment? As noted below, fixReduction does rely on the select generated by this recipe, rather than introduce another one.

nit: extra space at the beginning of each line?

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

public:
   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);
-  }
-
   /// 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
   /// VF. Return the cost of the instruction, including scalarization overhead
   /// if it's needed. The flag NeedToScalarize shows if the call needs to be
void InnerLoopVectorizer::fixReduction(VPWidenPHIRecipe *PhiR,
   assert(Legal->isReductionVariable(OrigPhi) &&
          "Unable to find the reduction variable");
   RecurrenceDescriptor RdxDesc = *PhiR->getRecurrenceDescriptor();
 
   RecurKind RK = RdxDesc.getRecurrenceKind();
   TrackingVH<Value> ReductionStartValue = RdxDesc.getRecurrenceStartValue();
   Instruction *LoopExitInst = RdxDesc.getLoopExitInstr();
   setDebugLocFromInst(Builder, ReductionStartValue);
-  bool IsInLoopReductionPhi = Cost->isInLoopReduction(OrigPhi);
+  bool IsInLoopReductionPhi = PhiR->isInLoopReduction();
 
   VPValue *LoopExitInstDef = State.Plan->getVPValue(LoopExitInst);
   // This is the vector-clone of the value that leaves the loop.
   Type *VecTy = State.get(LoopExitInstDef, 0)->getType();
 
   // Wrap flags are in general invalid after vectorization, clear them.
   clearReductionWrapFlags(RdxDesc, State);
 
   // Fix the vector-loop phi.
void InnerLoopVectorizer::widenPHIInstruction(Instruction *PN,
   VPValue *StartVPV = PhiR->getStartValue();
   Value *StartV = StartVPV ? StartVPV->getLiveInIRValue() : nullptr;
   // In order to support recurrences we need to be able to vectorize Phi nodes.
   // Phi nodes have cycles, so we need to vectorize them in two stages. This is
   // stage #1: We create a new vector PHI node with no incoming edges. We'll use
   // this value when we vectorize all of the instructions that use the PHI.
   if (RdxDesc || Legal->isFirstOrderRecurrence(P)) {
     Value *Iden = nullptr;
-    bool ScalarPHI =
-        (State.VF.isScalar()) || Cost->isInLoopReduction(cast<PHINode>(PN));
+    bool IsInLoopReduction = PhiR->isInLoopReduction();
+    bool ScalarPHI = State.VF.isScalar() || IsInLoopReduction;
     Type *VecTy =
         ScalarPHI ? PN->getType() : VectorType::get(PN->getType(), State.VF);
 
     if (RdxDesc) {
       assert(Legal->isReductionVariable(P) && StartV &&
              "RdxDesc should only be set for reduction variables; in that case "
              "a StartV is also required");
       RecurKind RK = RdxDesc->getRecurrenceKind();
if (RdxDesc) {
           IRBuilderBase::InsertPointGuard IPBuilder(Builder);
           Builder.SetInsertPoint(LoopVectorPreHeader->getTerminator());
           Constant *Zero = Builder.getInt32(0);
           StartV = Builder.CreateInsertElement(Iden, StartV, Zero);
         }
       }
     }
 
-    bool IsOrdered = State.VF.isVector() &&
-                     Cost->isInLoopReduction(cast<PHINode>(PN)) &&
+    bool IsOrdered = State.VF.isVector() && IsInLoopReduction &&
                      useOrderedReductions(*RdxDesc);
 
     for (unsigned Part = 0; Part < State.UF; ++Part) {
       // This is phase one of vectorizing PHIs.
       if (Part > 0 && IsOrdered)
         return;
       Value *EntryPart = PHINode::Create(
           VecTy, 2, "vec.phi", &*LoopVectorBody->getFirstInsertionPt());
for (unsigned i = 0; i < IG->getFactor(); ++i)
           OriginalV->replaceAllUsesWith(VPIG->getVPValue(J));
           J++;
         }
         RecipeBuilder.getRecipe(Member)->eraseFromParent();
       }
   }
 
   // Adjust the recipes for any inloop reductions.
   if (Range.Start.isVector())
     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);
-    auto *Cond = RecipeBuilder.createBlockInMask(OrigLoop->getHeader(), Plan);
-    for (auto &Reduction : Legal->getReductionVars()) {
-      if (CM.isInLoopReduction(Reduction.first))
+  if (CM.foldTailByMasking()) {
+    VPValue *Cond = nullptr;
+    unsigned NumRegularReductions = 0;
+    for (VPRecipeBase &R : Plan->getEntry()->getEntryBasicBlock()->phis()) {
+      VPWidenPHIRecipe *PhiR = dyn_cast<VPWidenPHIRecipe>(&R);
+      if (!PhiR || !PhiR->getRecurrenceDescriptor() ||
+          PhiR->isInLoopReduction())
         continue;
-      VPValue *Phi = Plan->getOrAddVPValue(Reduction.first);
-      VPValue *Red = Plan->getOrAddVPValue(Reduction.second.getLoopExitInstr());
+
+      if (!Cond) {
+        Builder.setInsertPoint(VPBB);
+        Cond = RecipeBuilder.createBlockInMask(OrigLoop->getHeader(), Plan);
+      }
+      VPValue *Phi = PhiR;
+      VPValue *Red = Plan->getOrAddVPValue(
+          PhiR->getRecurrenceDescriptor()->getLoopExitInstr());
       Builder.createNaryOp(Instruction::Select, {Cond, Red, Phi});
+#ifndef NDEBUG
+      NumRegularReductions++;
+#endif
     }
+    (void)NumRegularReductions;
+    assert(!Range.Start.isVector() ||
+           NumRegularReductions == Legal->getReductionVars().size() -
+                                       CM.getInLoopReductionChains().size() &&
+               "failed to identify an in-loop reduction in VPlan");
   }
 
   std::string PlanName;
   raw_string_ostream RSO(PlanName);
   ElementCount VF = Range.Start;
   Plan->addVF(VF);
   RSO << "Initial VPlan for VF={" << VF;
   for (VF *= 2; ElementCount::isKnownLT(VF, Range.End); VF *= 2) {
     Plan->addVF(VF);

 
 // 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 LoopVectorizationPlanner::adjustRecipesForInLoopReductions(
     VPlanPtr &Plan, VPRecipeBuilder &RecipeBuilder) {
   for (auto &Reduction : CM.getInLoopReductionChains()) {
     PHINode *Phi = Reduction.first;
     RecurrenceDescriptor &RdxDesc = Legal->getReductionVars()[Phi];
     const SmallVector<Instruction *, 4> &ReductionOperations = Reduction.second;
 
     // 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) {
       }
       Chain = R;
     }
   }
 }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 void VPInterleaveRecipe::print(raw_ostream &O, const Twine &Indent,
                                VPSlotTracker &SlotTracker) const {
   O << Indent << "INTERLEAVE-GROUP with factor " << IG->getFactor() << " at ";
   IG->getInsertPos()->printAsOperand(O, false);
   O << ", ";
   getAddr()->printAsOperand(O, SlotTracker);
   VPValue *Mask = getMask();
   if (Mask) {
     O << ", ";
     Mask->printAsOperand(O, SlotTracker);
   }

llvm/lib/Transforms/Vectorize/VPlan.h

#endif
 
   /// Returns the incoming value from the loop backedge, if it is a reduction.
   VPValue *getBackedgeValue() {
     assert(RdxDesc && "second incoming value is only guaranteed to be backedge "
                       "value for reductions");
     return getOperand(1);
   }
 
+  const VPValue *getBackedgeValue() const {
+    assert(RdxDesc && "second incoming value is only guaranteed to be backedge "
+                      "value for reductions");
+    return getOperand(1);
+  }
+
   /// Adds a pair (\p IncomingV, \p IncomingBlock) to the phi.
   void addIncoming(VPValue *IncomingV, VPBasicBlock *IncomingBlock) {
     addOperand(IncomingV);
     IncomingBlocks.push_back(IncomingBlock);
   }
 
   /// Returns the \p I th incoming VPValue.
   VPValue *getIncomingValue(unsigned I) { return getOperand(I); }
 
   /// Returns the \p I th incoming VPBasicBlock.
   VPBasicBlock *getIncomingBlock(unsigned I) { return IncomingBlocks[I]; }
 
   RecurrenceDescriptor *getRecurrenceDescriptor() { return RdxDesc; }
+
+  /// Returns true if this phi is part of an in-loop reduction.
+  bool isInLoopReduction() const;
 };
 
 /// A recipe for vectorizing a phi-node as a sequence of mask-based select
 /// instructions.
 class VPBlendRecipe : public VPRecipeBase, public VPValue {
   PHINode *Phi;
 
 public:

llvm/lib/Transforms/Vectorize/VPlan.cpp

if (getNumOperands() != OriginalPhi->getNumOperands()) {
     return;
   }
 
   printAsOperand(O, SlotTracker);
   O << " = phi ";
   printOperands(O, SlotTracker);
 }
 
+bool VPWidenPHIRecipe::isInLoopReduction() const {
+  if (!RdxDesc)
+    return false;
+
+  SmallVector<const VPRecipeBase *, 4> WorkList;
+  // Walk backwards the backedge value until we either reach a VPReductionRecipe
+  // or a phi.
+  WorkList.push_back(cast<VPRecipeBase>(getBackedgeValue()->getDef()));
+  while (!WorkList.empty()) {
+    const VPRecipeBase *Current = WorkList.pop_back_val();
+    if (Current->isPhi())
+      continue;
+    if (isa<VPReductionRecipe>(Current))
+      return true;
+    for (VPValue *Op : Current->operands()) {
+      if (!Op->getDef())
+        continue;
+      if (auto *R = dyn_cast<VPRecipeBase>(Op->getDef()))
+        WorkList.push_back(R);
+    }
+  }
+  return false;
+}
+
 void VPBlendRecipe::print(raw_ostream &O, const Twine &Indent,
                           VPSlotTracker &SlotTracker) const {
   O << Indent << "BLEND ";
   Phi->printAsOperand(O, false);
   O << " =";
   if (getNumIncomingValues() == 1) {
     // Not a User of any mask: not really blending, this is a
     // single-predecessor phi.

llvm/lib/Transforms/Vectorize/VPlanValue.h

bool hasMoreThanOneUniqueUser() {
     while (Current != user_end() && *user_begin() == *Current)
       Current++;
     return Current != user_end();
   }
 
   void replaceAllUsesWith(VPValue *New);
 
   VPDef *getDef() { return Def; }
+  const VPDef *getDef() const { return Def; }
 
   /// Returns the underlying IR value, if this VPValue is defined outside the
   /// scope of VPlan. Returns nullptr if the VPValue is defined by a VPDef
   /// inside a VPlan.
   Value *getLiveInIRValue() {
     assert(!getDef() &&
            "VPValue is not a live-in; it is defined by a VPDef inside a VPlan");
     return getUnderlyingValue();