[VPlan] Move widening check for non-memory/non-calls to function (NFC).

[fhahn]

Summary

After introducing VPWidenSelectRecipe, the duplicated logic can be
shared.

[Ayal]

Yes, this follows D77972 naturally :-)

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6935

Can add !isa<PHINode>(I) to this assertion list after sinking tryToWidenSelect() as suggested below.

7093

How about doing here (tryToCreateRecipe should have a better name btw):

  if (!shouldWiden(I, Range))
    return false;

  if ((Recipe = tryToWidenSelect(Instr)) ||
      (isa<GetElementPtrInst>(Instr) &&
       (Recipe = new VPWidenGEPRecipe(GEP, OrigLoop)))) ||
     (Recipe = tryToWiden(Instr, *Plan))) {
    setRecipe(Instr, Recipe);
    VPBB->appendRecipe(Recipe);
    return true;
  }

  return false;
}

[fhahn]

Consolidate widening checks for GEP, Select and general widening as suggested in tryToCreateRecipe.

[Ayal]

Thanks for cleaning this up! Final couple of nits.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6935

Retain the !isa<Call, Load, Store> alongside !isa<PHI>?
To avoid using this method for such opcodes, as their widening decision is more complex.

6942

More natural to reverse the conditions and rename willWiden to willScalarize, returning !getDecisionAndClampRange(willScalarize, Range)?

6966

Branches should also find their way out of IsVectorizableOpcode and into an assert; no recipes are created for them.

[fhahn]

Change willWiden to willScalarize.

(This patch depends on D77973)

[fhahn]

Avoid calling shouldWiden for calls, loads and stores by exiting earlier. They should be widened by the specialized recipes or scalarized.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6935

I've brought it back, but with the new code structure, it also required adding an early exit for calls, loads & stores, but I think that makes sense. What do you think?

6966

Will submit as a follow-up commit once this one lands.

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6935

Ahh, right. Perhaps better to early-exit here:

// Calls and memory instructions that were not widened earlier should be scalarized.
if (isa<CallInst>(Instr) || isa<LoadInst>(Instr) || isa<StoreInst>(Instr))
  return false;
assert(!isa<PHINode>(I) && ...);

6966

Sure.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6935

Is there an advantage to have the early exit here? To me it seems to fit slightly better at the call site of shouldWiden, because there it is clearer what is happening (i.e. the code to handle calls, stores, loads is just above.)

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6935

Your original version is fine, as approved. Advantage of last suggestion was merely to avoid an obvious assert immediately following an early-exiting.

Thoughts of further cleanup:

• have tryToCreate[Widen]Recipe return a VPRecipeBase* instead of bool,
• have its caller take care of { setRecipe(Instr, Recipe); VPBB->appendRecipe(Recipe); }; either using another method of VPRecipeBuilder or exposing its setRecipe() to LVP.
• split tryToOptimizeInduction in two: first part that handles PHI and second that handles Trunc
• refactor out various isa's from tryTo's and fold the early-exits closer:

VPRecipeBase *VPRecipeBuilder::tryToCreateWidenRecipe(Instruction *Instr, VFRange &Range, VPlanPtr &Plan) {
  if (auto *CI = dyn_cast<CallInst>(Instr))
    return tryToWidenCall(CI, Range, *Plan);
  if (isa<LoadInst>(Instr) || isa<StoreInst>(Instr))
    return tryToWidenMemory(Instr, Range, Plan);
  if (auto *PHI = dyn_cast<PHINode>(Instr)) {
    if (auto *Recipe = tryToOptimizeInduction(PHI, Range))
      return Recipe;
    if (auto *Recipe = tryToBlend(PHI, Plan))
      return Recipe;
    return new VPWidenPHIRecipe(PHI);
  }
  if (auto *Trunc = dyn_cast<TruncInst>(Instr))
    if (auto *Recipe = tryToOptimizeIVTruncate(Trunc, Range))
      return Recipe;
  
  if (!shouldWiden(Instr, Range))
    return nullptr;

  if (auto *GEP = dyn_cast<GetElementPtrInst>(Instr))
    return new VPWidenGEPRecipe(GEP, OrigLoop);
  if (auto *Select = dyn_cast<SelectInst<(Instr))
    return tryToWidenSelect(Instr);
  return tryToWiden(Instr, *Plan);
}

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6935

SGTM, I just wanted to keep this relatively self-contained. It would be great if you could also take another look at D77972, which is required for this patch to land :)

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

VPRecipeBuilder::tryToWidenCall(Instruction *I, VFRange &Range, VPlan &Plan) {
   // Success: widen this call.
   auto VPValues = map_range(CI->arg_operands(), [&Plan](Value *Op) {
     return Plan.getOrAddVPValue(Op);
   });
 
   return new VPWidenCallRecipe(*CI, VPValues);
 }
 
-VPWidenSelectRecipe *VPRecipeBuilder::tryToWidenSelect(Instruction *I,
-                                                       VFRange &Range) {
-  auto *SI = dyn_cast<SelectInst>(I);
-  if (!SI)
-    return nullptr;
-
-  // SI should be widened, unless it is scalar after vectorization,
+bool VPRecipeBuilder::shouldWiden(Instruction *I, VFRange &Range) const {
+  assert(!isa<PHINode>(I) && !isa<LoadInst>(I) && !isa<StoreInst>(I) &&
+         "Instruction should have been handled earlier");
+  // Instruction should be widened, unless it is scalar after vectorization,
   // scalarization is profitable or it is predicated.
-  auto willWiden = [this, SI](unsigned VF) -> bool {
-    return !CM.isScalarAfterVectorization(SI, VF) &&
-           !CM.isProfitableToScalarize(SI, VF) &&
-           !CM.isScalarWithPredication(SI, VF);
+  auto WillScalarize = [this, I](unsigned VF) -> bool {
+    return CM.isScalarAfterVectorization(I, VF) ||
+           CM.isProfitableToScalarize(I, VF) ||
+           CM.isScalarWithPredication(I, VF);
   };
-  if (!LoopVectorizationPlanner::getDecisionAndClampRange(willWiden, Range))
-    return nullptr;
+  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, VFRange &Range) {
-  bool IsPredicated = LoopVectorizationPlanner::getDecisionAndClampRange(
-      [&](unsigned VF) { return CM.isScalarWithPredication(I, VF); }, Range);
-
-  if (IsPredicated)
-    return nullptr;
-
-  assert(!isa<PHINode>(I) && "PHIs should have been handled earlier");
+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::Br:
     case Instruction::FAdd:
     case Instruction::FCmp:
     case Instruction::FDiv:
     case Instruction::FMul:
     case Instruction::FNeg:
     case Instruction::FPExt:
     case Instruction::FPToSI:
     case Instruction::FPToUI:
case Instruction::ZExt:
       return true;
     }
     return false;
   };
 
   if (!IsVectorizableOpcode(I->getOpcode()))
     return nullptr;
 
-  auto willWiden = [&](unsigned VF) -> bool {
-    return !CM.isScalarAfterVectorization(I, VF) &&
-           !CM.isProfitableToScalarize(I, VF);
-  };
-
-  if (!LoopVectorizationPlanner::getDecisionAndClampRange(willWiden, Range))
-    return nullptr;
-
   // Success: widen this instruction.
   return new VPWidenRecipe(*I);
 }
 
 VPBasicBlock *VPRecipeBuilder::handleReplication(
     Instruction *I, VFRange &Range, VPBasicBlock *VPBB,
     DenseMap<Instruction *, VPReplicateRecipe *> &PredInst2Recipe,
     VPlanPtr &Plan) {

 bool VPRecipeBuilder::tryToCreateRecipe(Instruction *Instr, VFRange &Range,
                                         VPlanPtr &Plan, VPBasicBlock *VPBB) {
   VPRecipeBase *Recipe = nullptr;
 
   // 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 = tryToWidenSelect(Instr, Range)) ||
       (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;
   }
 
-  // Handle GEP widening.
-  if (GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(Instr)) {
-    auto Scalarize = [&](unsigned VF) {
-      return CM.isScalarWithPredication(Instr, VF) ||
-             CM.isScalarAfterVectorization(Instr, VF) ||
-             CM.isProfitableToScalarize(Instr, VF);
-    };
-    if (LoopVectorizationPlanner::getDecisionAndClampRange(Scalarize, Range))
-      return false;
-    VPWidenGEPRecipe *Recipe = new VPWidenGEPRecipe(GEP, OrigLoop);
-    setRecipe(Instr, Recipe);
-    VPBB->appendRecipe(Recipe);
-    return true;
-  }
+  // 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 (!shouldWiden(Instr, Range))
 
-  // Check if Instr is to be widened by a general VPWidenRecipe, after
-  // having first checked for specific widening recipes.
-  if ((Recipe = tryToWiden(Instr, Range))) {
+  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;
   }
 
   return false;
 }
 

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

class VPRecipeBuilder {
   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.
public:
   /// full if-conversion.
   VPBlendRecipe *tryToBlend(Instruction *I, 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);
-
-  VPWidenSelectRecipe *tryToWidenSelect(Instruction *I, VFRange &Range);
-
-  /// Check if \p I can be widened within the given VF \p Range. If \p I can be
-  /// widened for \p Range.Start, build a new VPWidenRecipe and return it.
-  /// Range.End may be decreased to ensure same decision from \p Range.Start to
-  /// \p Range.End.
-  VPWidenRecipe *tryToWiden(Instruction *I, VFRange &Range);
+  /// 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);
+
+  /// 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,