[VPlan] Clean up tryToCreate(Widen)Recipe. (NFC)

[fhahn]

Summary

This patch includes some clean-ups to tryToCreateRecipe, suggested in
D77973.

It includes:

• Renaming tryToCreateRecipe to tryToCreateWidenRecipe.
• Move VPBB insertion logic to caller of tryToCreateWidenRecipe.
• Hoists instruction checks to tryToCreateWidenRecipe, making it clearer which instructions are handled by which recipe, simplifying the checks by using early exits.
• Split up handling of induction PHIs and truncates using inductions.

[Ayal]

Thanks for following up on this! Looks good to me with a couple of nits.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6867

No need to cast I into TruncInst.

6892

Orthogonal thought: perhaps if the constructor of VPBlendRecipe would receive two iterator_ranges, one for the values and one for the masks, and fuse them itself, tryToBlend() (or createBlendRecipe()) could be simplified using map_ranges.

7085

It would be clearer to further refactor the handling of Phi's into

if (Phi->getParent() != OrigLoop->getHeader())
  return tryToBlend(Phi, Plan); // Rename to createBlendRecipe()?
if (Recipe = tryToOptimizeInductionPHI(Phi))
  return Recipe;
return new VPWidenPHIRecipe(Phi); // Rename to VPWidenHeaderPHIRecipe()?

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

108

Above comment needs updating.

121

nit: [is] can

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

7085

Thanks, I'll include the changes in the commit. Renaming could be done as follow-up, together with the suggestion re BlendRecipe above.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

 //===- LoopVectorize.cpp - A Loop Vectorizer ------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
VPValue *VPRecipeBuilder::createBlockInMask(BasicBlock *BB, VPlanPtr &Plan) {
   }
 
   return BlockMaskCache[BB] = BlockMask;
 }
 
 VPWidenMemoryInstructionRecipe *
 VPRecipeBuilder::tryToWidenMemory(Instruction *I, VFRange &Range,
                                   VPlanPtr &Plan) {
-  if (!isa<LoadInst>(I) && !isa<StoreInst>(I))
-    return nullptr;
+  assert((isa<LoadInst>(I) || isa<StoreInst>(I)) &&
+         "Must be called with either a load or store");
 
   auto willWiden = [&](unsigned VF) -> bool {
     if (VF == 1)
       return false;
     LoopVectorizationCostModel::InstWidening Decision =
         CM.getWideningDecision(I, VF);
     assert(Decision != LoopVectorizationCostModel::CM_Unknown &&
            "CM decision should be taken at this point.");
if (LoadInst *Load = dyn_cast<LoadInst>(I))
     return new VPWidenMemoryInstructionRecipe(*Load, Addr, Mask);
 
   StoreInst *Store = cast<StoreInst>(I);
   VPValue *StoredValue = Plan->getOrAddVPValue(Store->getValueOperand());
   return new VPWidenMemoryInstructionRecipe(*Store, Addr, StoredValue, Mask);
 }
 
 VPWidenIntOrFpInductionRecipe *
-VPRecipeBuilder::tryToOptimizeInduction(Instruction *I, VFRange &Range) {
-  if (PHINode *Phi = dyn_cast<PHINode>(I)) {
-    // Check if this is an integer or fp induction. If so, build the recipe that
-    // produces its scalar and vector values.
-    InductionDescriptor II = Legal->getInductionVars().lookup(Phi);
-    if (II.getKind() == InductionDescriptor::IK_IntInduction ||
-        II.getKind() == InductionDescriptor::IK_FpInduction)
-      return new VPWidenIntOrFpInductionRecipe(Phi);
+VPRecipeBuilder::tryToOptimizeInductionPHI(PHINode *Phi) {
 
-    return nullptr;
-  }
 
+VPWidenIntOrFpInductionRecipe *
+VPRecipeBuilder::tryToOptimizeInductionTruncate(TruncInst *I, VFRange &Range) {
   // Optimize the special case where the source is a constant integer
   // induction variable. Notice that we can only optimize the 'trunc' case
   // because (a) FP conversions lose precision, (b) sext/zext may wrap, and
   // (c) other casts depend on pointer size.
 
   // Determine whether \p K is a truncation based on an induction variable that
   // can be optimized.
   auto isOptimizableIVTruncate =
       [&](Instruction *K) -> std::function<bool(unsigned)> {
     return
         [=](unsigned VF) -> bool { return CM.isOptimizableIVTruncate(K, VF); };
   };
 
-  if (isa<TruncInst>(I) && LoopVectorizationPlanner::getDecisionAndClampRange(
-                               isOptimizableIVTruncate(I), Range))
+  if (LoopVectorizationPlanner::getDecisionAndClampRange(
     return new VPWidenIntOrFpInductionRecipe(cast<PHINode>(I->getOperand(0)),
                                              cast<TruncInst>(I));
   return nullptr;
 }
 
-VPBlendRecipe *VPRecipeBuilder::tryToBlend(Instruction *I, VPlanPtr &Plan) {
-  PHINode *Phi = dyn_cast<PHINode>(I);
-  if (!Phi || Phi->getParent() == OrigLoop->getHeader())
+VPBlendRecipe *VPRecipeBuilder::tryToBlend(PHINode *Phi, VPlanPtr &Plan) {
+  if (Phi->getParent() == OrigLoop->getHeader())
     return nullptr;
 
   // We know that all PHIs in non-header blocks are converted into selects, so
   // we don't have to worry about the insertion order and we can just use the
   // builder. At this point we generate the predication tree. There may be
   // duplications since this is a simple recursive scan, but future
   // optimizations will clean it up.
 
   SmallVector<VPValue *, 2> Operands;
   unsigned NumIncoming = Phi->getNumIncomingValues();
   for (unsigned In = 0; In < NumIncoming; In++) {
     VPValue *EdgeMask =
       createEdgeMask(Phi->getIncomingBlock(In), Phi->getParent(), Plan);
     assert((EdgeMask || NumIncoming == 1) &&
            "Multiple predecessors with one having a full mask");
     Operands.push_back(Plan->getOrAddVPValue(Phi->getIncomingValue(In)));
     if (EdgeMask)
       Operands.push_back(EdgeMask);
   }
   return new VPBlendRecipe(Phi, Operands);
 }
 
-VPWidenCallRecipe *
-VPRecipeBuilder::tryToWidenCall(Instruction *I, VFRange &Range, VPlan &Plan) {
+VPWidenCallRecipe *VPRecipeBuilder::tryToWidenCall(CallInst *CI, VFRange &Range,
+                                                   VPlan &Plan) {
 
   bool IsPredicated = LoopVectorizationPlanner::getDecisionAndClampRange(
-      [&](unsigned VF) { return CM.isScalarWithPredication(I, VF); }, Range);
+      [this, CI](unsigned VF) { return CM.isScalarWithPredication(CI, VF); },
+      Range);
 
-  CallInst *CI = dyn_cast<CallInst>(I);
-  if (IsPredicated || !CI)
+  if (IsPredicated)
     return nullptr;
 
   Intrinsic::ID ID = getVectorIntrinsicIDForCall(CI, TLI);
   if (ID && (ID == Intrinsic::assume || ID == Intrinsic::lifetime_end ||
              ID == Intrinsic::lifetime_start || ID == Intrinsic::sideeffect))
     return nullptr;
 
   auto willWiden = [&](unsigned VF) -> bool {
auto WillScalarize = [this, I](unsigned VF) -> bool {
     return CM.isScalarAfterVectorization(I, VF) ||
            CM.isProfitableToScalarize(I, VF) ||
            CM.isScalarWithPredication(I, VF);
   };
   return !LoopVectorizationPlanner::getDecisionAndClampRange(WillScalarize,
                                                              Range);
 }
 
-VPWidenSelectRecipe *VPRecipeBuilder::tryToWidenSelect(Instruction *I) {
-  auto *SI = dyn_cast<SelectInst>(I);
-  if (!SI)
-    return nullptr;
-  auto *SE = PSE.getSE();
-  bool InvariantCond =
-      SE->isLoopInvariant(PSE.getSCEV(SI->getOperand(0)), OrigLoop);
-  // Success: widen this instruction.
-  return new VPWidenSelectRecipe(*SI, InvariantCond);
-}
-
 VPWidenRecipe *VPRecipeBuilder::tryToWiden(Instruction *I, VPlan &Plan) {
   auto IsVectorizableOpcode = [](unsigned Opcode) {
     switch (Opcode) {
     case Instruction::Add:
     case Instruction::And:
     case Instruction::AShr:
     case Instruction::BitCast:
     case Instruction::FAdd:
VPRegionBlock *VPRecipeBuilder::createReplicateRegion(Instruction *Instr,
   // Note: first set Entry as region entry and then connect successors starting
   // from it in order, to propagate the "parent" of each VPBasicBlock.
   VPBlockUtils::insertTwoBlocksAfter(Pred, Exit, BlockInMask, Entry);
   VPBlockUtils::connectBlocks(Pred, Exit);
 
   return Region;
 }
 
-bool VPRecipeBuilder::tryToCreateRecipe(Instruction *Instr, VFRange &Range,
-                                        VPlanPtr &Plan, VPBasicBlock *VPBB) {
-  VPRecipeBase *Recipe = nullptr;
-
+VPRecipeBase *VPRecipeBuilder::tryToCreateWidenRecipe(Instruction *Instr,
+                                                      VFRange &Range,
+                                                      VPlanPtr &Plan) {
   // First, check for specific widening recipes that deal with calls, memory
   // operations, inductions and Phi nodes.
-  if ((Recipe = tryToWidenCall(Instr, Range, *Plan)) ||
-      (Recipe = tryToWidenMemory(Instr, Range, Plan)) ||
-      (Recipe = tryToOptimizeInduction(Instr, Range)) ||
-      (Recipe = tryToBlend(Instr, Plan)) ||
-      (isa<PHINode>(Instr) &&
-       (Recipe = new VPWidenPHIRecipe(cast<PHINode>(Instr))))) {
-    setRecipe(Instr, Recipe);
-    VPBB->appendRecipe(Recipe);
-    return true;
+  if (auto *CI = dyn_cast<CallInst>(Instr))
+    return tryToWidenCall(CI, Range, *Plan);
+
+  if (isa<LoadInst>(Instr) || isa<StoreInst>(Instr))
+    return tryToWidenMemory(Instr, Range, Plan);
+
+  VPRecipeBase *Recipe;
+  if (auto Phi = dyn_cast<PHINode>(Instr)) {
+    if ((Recipe = tryToOptimizeInductionPHI(Phi)) ||
+        (Recipe = tryToBlend(Phi, Plan)))
+      return Recipe;
+
+    return new VPWidenPHIRecipe(Phi);
   }
 
-  // Calls and memory instructions are widened by the specialized recipes above,
-  // or scalarized.
-  if (isa<CallInst>(Instr) || isa<LoadInst>(Instr) || isa<StoreInst>(Instr))
-    return false;
+  if (isa<TruncInst>(Instr) &&
+      (Recipe = tryToOptimizeInductionTruncate(cast<TruncInst>(Instr), Range)))
+    return Recipe;
 
   if (!shouldWiden(Instr, Range))
-    return false;
+    return nullptr;
 
-  if ((Recipe = tryToWidenSelect(Instr)) ||
-      (isa<GetElementPtrInst>(Instr) &&
-       (Recipe =
-            new VPWidenGEPRecipe(cast<GetElementPtrInst>(Instr), OrigLoop))) ||
-      (Recipe = tryToWiden(Instr, *Plan))) {
-    setRecipe(Instr, Recipe);
-    VPBB->appendRecipe(Recipe);
-    return true;
+  if (auto GEP = dyn_cast<GetElementPtrInst>(Instr))
+    return new VPWidenGEPRecipe(GEP, OrigLoop);
+
+  if (auto *SI = dyn_cast<SelectInst>(Instr)) {
+    bool InvariantCond =
+        PSE.getSE()->isLoopInvariant(PSE.getSCEV(SI->getOperand(0)), OrigLoop);
+    return new VPWidenSelectRecipe(*SI, InvariantCond);
   }
 
-  return false;
+  return tryToWiden(Instr, *Plan);
 }
 
 void LoopVectorizationPlanner::buildVPlansWithVPRecipes(unsigned MinVF,
                                                         unsigned MaxVF) {
   assert(OrigLoop->empty() && "Inner loop expected.");
 
   // Collect conditions feeding internal conditional branches; they need to be
   // represented in VPlan for it to model masking.
for (Instruction &I : BB->instructionsWithoutDebug()) {
       Instruction *Instr = &I;
 
       // First filter out irrelevant instructions, to ensure no recipes are
       // built for them.
       if (isa<BranchInst>(Instr) ||
           DeadInstructions.find(Instr) != DeadInstructions.end())
         continue;
 
-      if (RecipeBuilder.tryToCreateRecipe(Instr, Range, Plan, VPBB))
+      if (auto Recipe =
+              RecipeBuilder.tryToCreateWidenRecipe(Instr, Range, Plan)) {
+        RecipeBuilder.setRecipe(Instr, Recipe);
+        VPBB->appendRecipe(Recipe);
         continue;
+      }
 
       // Otherwise, if all widening options failed, Instruction is to be
       // replicated. This may create a successor for VPBB.
       VPBasicBlock *NextVPBB = RecipeBuilder.handleReplication(
           Instr, Range, VPBB, PredInst2Recipe, Plan);
       if (NextVPBB != VPBB) {
         VPBB = NextVPBB;
         VPBB->setName(BB->hasName() ? BB->getName() + "." + Twine(VPBBsForBB++)

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

 //===- VPRecipeBuilder.h - Helper class to build recipes --------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
class VPRecipeBuilder {
   EdgeMaskCacheTy EdgeMaskCache;
   BlockMaskCacheTy BlockMaskCache;
 
   // VPlan-VPlan transformations support: Hold a mapping from ingredients to
   // their recipe. To save on memory, only do so for selected ingredients,
   // marked by having a nullptr entry in this map.
   DenseMap<Instruction *, VPRecipeBase *> Ingredient2Recipe;
 
+  /// Check if \p I can be widened at the start of \p Range and possibly
+  /// decrease the range such that the returned value holds for the entire \p
+  /// Range. The function should not be called for memory instructions or calls.
+  bool shouldWiden(Instruction *I, VFRange &Range) const;
+
+public:
   /// Set the recipe created for given ingredient. This operation is a no-op for
   /// ingredients that were not marked using a nullptr entry in the map.
   void setRecipe(Instruction *I, VPRecipeBase *R) {
     if (!Ingredient2Recipe.count(I))
       return;
     assert(Ingredient2Recipe[I] == nullptr &&
            "Recipe already set for ingredient");
     Ingredient2Recipe[I] = R;
   }
 
-  /// Check if \p I can be widened at the start of \p Range and possibly
-  /// decrease the range such that the returned value holds for the entire \p
-  /// Range. The function should not be called for memory instructions or calls.
-  bool shouldWiden(Instruction *I, VFRange &Range) const;
-
-public:
   /// A helper function that computes the predicate of the block BB, assuming
   /// that the header block of the loop is set to True. It returns the *entry*
   /// mask for the block BB.
   VPValue *createBlockInMask(BasicBlock *BB, VPlanPtr &Plan);
 
   /// A helper function that computes the predicate of the edge between SRC
   /// and DST.
   VPValue *createEdgeMask(BasicBlock *Src, BasicBlock *Dst, VPlanPtr &Plan);
public:
   VPRecipeBase *getRecipe(Instruction *I) {
     assert(Ingredient2Recipe.count(I) &&
            "Recording this ingredients recipe was not requested");
     assert(Ingredient2Recipe[I] != nullptr &&
            "Ingredient doesn't have a recipe");
     return Ingredient2Recipe[I];
   }
 
-  /// Check if \I is a memory instruction to be widened for \p Range.Start and
-  /// potentially masked. Such instructions are handled by a recipe that takes
-  /// an additional VPInstruction for the mask.
+  /// Check if the load or store instruction \p I should widened for \p
+  /// Range.Start and potentially masked. Such instructions are handled by a
+  /// recipe that takes an additional VPInstruction for the mask.
   VPWidenMemoryInstructionRecipe *
   tryToWidenMemory(Instruction *I, VFRange &Range, VPlanPtr &Plan);
 
   /// Check if an induction recipe should be constructed for \I within the given
   /// VF \p Range. If so build and return it. If not, return null. \p Range.End
   /// may be decreased to ensure same decision from \p Range.Start to
   /// \p Range.End.
-  VPWidenIntOrFpInductionRecipe *tryToOptimizeInduction(Instruction *I,
+  VPWidenIntOrFpInductionRecipe *tryToOptimizeInductionPHI(PHINode *Phi);
+
+  /// Optimize the special case where the operand of \p I is a constant integer
+  /// induction variable.
+  VPWidenIntOrFpInductionRecipe *tryToOptimizeInductionTruncate(TruncInst *I,
                                                                 VFRange &Range);
 
   /// Handle non-loop phi nodes. Currently all such phi nodes are turned into
   /// a sequence of select instructions as the vectorizer currently performs
   /// full if-conversion.
-  VPBlendRecipe *tryToBlend(Instruction *I, VPlanPtr &Plan);
+  VPBlendRecipe *tryToBlend(PHINode *Phi, VPlanPtr &Plan);
 
-  /// Handle call instruction. If \p I is a call that can be widened for \p
-  /// Range.Start, return a new VPWidenCallRecipe. Range.End may be decreased to
-  /// ensure same decision from \p Range.Start to \p Range.End.
-  VPWidenCallRecipe *tryToWidenCall(Instruction *I, VFRange &Range,
-                                    VPlan &Plan);
-  /// Check if \p I is a SelectInst and return a VPWidenSelectRecipe if it is.
-  /// The function should only be called if the cost-model indicates that
-  /// widening should be performed.
-  VPWidenSelectRecipe *tryToWidenSelect(Instruction *I);
+  /// Handle call instructions. If \p CI is can be widened for \p Range.Start,
+  /// return a new VPWidenCallRecipe. Range.End may be decreased to ensure same
+  /// decision from \p Range.Start to \p Range.End.
+  VPWidenCallRecipe *tryToWidenCall(CallInst *CI, VFRange &Range, VPlan &Plan);
 
   /// Check if \p I has an opcode that can be widened and return a VPWidenRecipe
   /// if it can. The function should only be called if the cost-model indicates
   /// that widening should be performed.
   VPWidenRecipe *tryToWiden(Instruction *I, VPlan &Plan);
 
   /// Create a replicating region for instruction \p I that requires
   /// predication. \p PredRecipe is a VPReplicateRecipe holding \p I.
   VPRegionBlock *createReplicateRegion(Instruction *I, VPRecipeBase *PredRecipe,
                                        VPlanPtr &Plan);
 
 public:
   VPRecipeBuilder(Loop *OrigLoop, const TargetLibraryInfo *TLI,
                   LoopVectorizationLegality *Legal,
                   LoopVectorizationCostModel &CM,
                   PredicatedScalarEvolution &PSE, VPBuilder &Builder)
       : OrigLoop(OrigLoop), TLI(TLI), Legal(Legal), CM(CM), PSE(PSE),
         Builder(Builder) {}
 
   /// Check if a recipe can be create for \p I withing the given VF \p Range.
-  /// If a recipe can be created, it adds it to \p VPBB.
-  bool tryToCreateRecipe(Instruction *Instr, VFRange &Range, VPlanPtr &Plan,
-                         VPBasicBlock *VPBB);
+  /// If a recipe can be created, return it. Otherwise return nullptr.
+  VPRecipeBase *tryToCreateWidenRecipe(Instruction *Instr, VFRange &Range,
+                                       VPlanPtr &Plan);
 
   /// Build a VPReplicationRecipe for \p I and enclose it within a Region if it
   /// is predicated. \return \p VPBB augmented with this new recipe if \p I is
   /// not predicated, otherwise \return a new VPBasicBlock that succeeds the new
   /// Region. Update the packing decision of predicated instructions if they
   /// feed \p I. Range.End may be decreased to ensure same recipe behavior from
   /// \p Range.Start to \p Range.End.
   VPBasicBlock *handleReplication(
       Instruction *I, VFRange &Range, VPBasicBlock *VPBB,
       DenseMap<Instruction *, VPReplicateRecipe *> &PredInst2Recipe,
       VPlanPtr &Plan);
 };
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_VECTORIZE_VPRECIPEBUILDER_H