[VPlan] Introduce VPWidenSelectRecipe (NFC).

[fhahn]

Summary

Widening a selects depends on whether the condition is loop invariant or
not. Rather than checking during codegen-time, the information can be
recorded at the VPlan construction time.

This was suggested as part of D76992, to reduce the reliance on
accessing the original underlying IR values.

[gilr]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6950

Can't be a phi at this point (actually this also holds for tryToWiden as all phi's are taken care of by tryToCreateRecipe, but that belongs in a separate patch).

6950

If statement can fold into the return expression.

llvm/lib/Transforms/Vectorize/VPlan.cpp

724

Should also print InvariantCond as e.g. in VPWidenGEPRecipe::print()

[fhahn]

Address comments, thanks!

Simplify willWiden check, print whether condition is invariant or not.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6950

Done! I think it would make sense to move all those checks into a separate function now, that there are multiple users. I'll do that as a follow up.

llvm/lib/Transforms/Vectorize/VPlan.cpp

724

Done, but I am not sure how to best print it specifically for the condition. I've added a note to the debug output.

[gilr]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6940

SI is not used beyond this point until recipe construction. Better to either use it consistently instead of I or only check isa<> here and cast<> at recipe construction.

6944

This can fold into willWiden.

llvm/lib/Transforms/Vectorize/VPlan.cpp

724

Agree, standardizing how to mark invariant operands would be easier for recipes using VPValues (e.g. just adding "[Inv]" before/after invariant operands).

725

Since it's an infrequent optimization it might be cleaner to only print if relevant, i.e. " << (InvariantCond ? " (condition is loop invariant)" : "").

[fhahn]

Address comments, thanks!

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

6940

changed to use it throughout.

[fhahn]

Fix type in comment.

[gilr]

LGTM, thanks!

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4420

A separate, independent nit, while we're here: get or create only values that are needed:

auto *ScalarCond = (InvariantCond ? getOrCreateScalarValue(I.getOperand(0), {0, 0}) : nullptr);

4423

Value *Cond = (InvariantCond ? ScalarCond : getOrCreateVectorValue(I.getOperand(0), Part));

4426

Value *Sel = Builder.CreateSelect(Cond, Op0, Op1);

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4420

+1, but it requires a few small test changes. I'll fix in a separate patch.

llvm/lib/Transforms/Vectorize/LoopVectorizationPlanner.h

 
 #include "VPlan.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 
 namespace llvm {
 
+class PredicatedScalarEvolution;
+
 /// VPlan-based builder utility analogous to IRBuilder.
 class VPBuilder {
   VPBasicBlock *BB = nullptr;
   VPBasicBlock::iterator InsertPt = VPBasicBlock::iterator();
 
   VPInstruction *createInstruction(unsigned Opcode,
                                    ArrayRef<VPValue *> Operands) {
     VPInstruction *Instr = new VPInstruction(Opcode, Operands);
class LoopVectorizationPlanner {
   LoopVectorizationLegality *Legal;
 
   /// The profitability analysis.
   LoopVectorizationCostModel &CM;
 
   /// The interleaved access analysis.
   InterleavedAccessInfo &IAI;
 
+  PredicatedScalarEvolution &PSE;
+
   SmallVector<VPlanPtr, 4> VPlans;
 
   /// This class is used to enable the VPlan to invoke a method of ILV. This is
   /// needed until the method is refactored out of ILV and becomes reusable.
   struct VPCallbackILV : public VPCallback {
     InnerLoopVectorizer &ILV;
 
     VPCallbackILV(InnerLoopVectorizer &ILV) : ILV(ILV) {}
class LoopVectorizationPlanner {
   unsigned BestVF = 0;
   unsigned BestUF = 0;
 
 public:
   LoopVectorizationPlanner(Loop *L, LoopInfo *LI, const TargetLibraryInfo *TLI,
                            const TargetTransformInfo *TTI,
                            LoopVectorizationLegality *Legal,
                            LoopVectorizationCostModel &CM,
-                           InterleavedAccessInfo &IAI)
-      : OrigLoop(L), LI(LI), TLI(TLI), TTI(TTI), Legal(Legal), CM(CM),
-        IAI(IAI) {}
+                           InterleavedAccessInfo &IAI,
+                           PredicatedScalarEvolution &PSE)
+      : OrigLoop(L), LI(LI), TLI(TLI), TTI(TTI), Legal(Legal), CM(CM), IAI(IAI),
+        PSE(PSE) {}
 
   /// Plan how to best vectorize, return the best VF and its cost, or None if
   /// vectorization and interleaving should be avoided up front.
   Optional<VectorizationFactor> plan(unsigned UserVF);
 
   /// Use the VPlan-native path to plan how to best vectorize, return the best
   /// VF and its cost.
   VectorizationFactor planInVPlanNativePath(unsigned UserVF);

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

public:
 
   /// Widen a single instruction within the innermost loop.
   void widenInstruction(Instruction &I);
 
   /// Widen a single call instruction within the innermost loop.
   void widenCallInstruction(CallInst &I, VPUser &ArgOperands,
                             VPTransformState &State);
 
+  /// Widen a single select instruction within the innermost loop.
+  void widenSelectInstruction(SelectInst &I, bool InvariantCond);
+
   /// Fix the vectorized code, taking care of header phi's, live-outs, and more.
   void fixVectorizedLoop();
 
   // Return true if any runtime check is added.
   bool areSafetyChecksAdded() { return AddedSafetyChecks; }
 
   /// A type for vectorized values in the new loop. Each value from the
   /// original loop, when vectorized, is represented by UF vector values in the

 }
 
 void InnerLoopVectorizer::widenInstruction(Instruction &I) {
   switch (I.getOpcode()) {
   case Instruction::Call:
   case Instruction::Br:
   case Instruction::PHI:
   case Instruction::GetElementPtr:
+  case Instruction::Select:
     llvm_unreachable("This instruction is handled by a different recipe.");
   case Instruction::UDiv:
   case Instruction::SDiv:
   case Instruction::SRem:
   case Instruction::URem:
   case Instruction::Add:
   case Instruction::FAdd:
   case Instruction::Sub:
for (unsigned Part = 0; Part < UF; ++Part) {
 
       // Use this vector value for all users of the original instruction.
       VectorLoopValueMap.setVectorValue(&I, Part, V);
       addMetadata(V, &I);
     }
 
     break;
   }
-  case Instruction::Select: {
-    // Widen selects.
-    // If the selector is loop invariant we can create a select
-    // instruction with a scalar condition. Otherwise, use vector-select.
-    auto *SE = PSE.getSE();
-    bool InvariantCond =
-        SE->isLoopInvariant(PSE.getSCEV(I.getOperand(0)), OrigLoop);
-    setDebugLocFromInst(Builder, &I);
-
-    // The condition can be loop invariant  but still defined inside the
-    // loop. This means that we can't just use the original 'cond' value.
-    // We have to take the 'vectorized' value and pick the first lane.
-    // Instcombine will make this a no-op.
-
-    auto *ScalarCond = getOrCreateScalarValue(I.getOperand(0), {0, 0});
-
-    for (unsigned Part = 0; Part < UF; ++Part) {
-      Value *Cond = getOrCreateVectorValue(I.getOperand(0), Part);
-      Value *Op0 = getOrCreateVectorValue(I.getOperand(1), Part);
-      Value *Op1 = getOrCreateVectorValue(I.getOperand(2), Part);
-      Value *Sel =
-          Builder.CreateSelect(InvariantCond ? ScalarCond : Cond, Op0, Op1);
-      VectorLoopValueMap.setVectorValue(&I, Part, Sel);
-      addMetadata(Sel, &I);
-    }
-
-    break;
-  }
-
   case Instruction::ICmp:
   case Instruction::FCmp: {
     // Widen compares. Generate vector compares.
     bool FCmp = (I.getOpcode() == Instruction::FCmp);
     auto *Cmp = cast<CmpInst>(&I);
     setDebugLocFromInst(Builder, Cmp);
     for (unsigned Part = 0; Part < UF; ++Part) {
       Value *A = getOrCreateVectorValue(Cmp->getOperand(0), Part);
#endif
       if (isa<FPMathOperator>(V))
         V->copyFastMathFlags(CI);
 
       VectorLoopValueMap.setVectorValue(&I, Part, V);
       addMetadata(V, &I);
   }
 }
 
+void InnerLoopVectorizer::widenSelectInstruction(SelectInst &I,
+                                                 bool InvariantCond) {
+  setDebugLocFromInst(Builder, &I);
+
+  // The condition can be loop invariant  but still defined inside the
+  // loop. This means that we can't just use the original 'cond' value.
+  // We have to take the 'vectorized' value and pick the first lane.
+  // Instcombine will make this a no-op.
+
+  auto *ScalarCond = getOrCreateScalarValue(I.getOperand(0), {0, 0});
+
+  for (unsigned Part = 0; Part < UF; ++Part) {
+    Value *Cond = getOrCreateVectorValue(I.getOperand(0), Part);
+    Value *Op0 = getOrCreateVectorValue(I.getOperand(1), Part);
+    Value *Op1 = getOrCreateVectorValue(I.getOperand(2), Part);
+    Value *Sel =
+        Builder.CreateSelect(InvariantCond ? ScalarCond : Cond, Op0, Op1);
+    VectorLoopValueMap.setVectorValue(&I, Part, Sel);
+    addMetadata(Sel, &I);
+  }
+}
+
 void LoopVectorizationCostModel::collectLoopScalars(unsigned VF) {
   // We should not collect Scalars more than once per VF. Right now, this
   // function is called from collectUniformsAndScalars(), which already does
   // this check. Collecting Scalars for VF=1 does not make any sense.
   assert(VF >= 2 && Scalars.find(VF) == Scalars.end() &&
          "This function should not be visited twice for the same VF");
 
   SmallSetVector<Instruction *, 8> Worklist;
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,
+  // 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);
+  };
+  if (!LoopVectorizationPlanner::getDecisionAndClampRange(willWiden, Range))
+    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;
 
   auto IsVectorizableOpcode = [](unsigned Opcode) {

 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;
   }
VPlanPtr LoopVectorizationPlanner::buildVPlanWithVPRecipes(
 
   // 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, Builder);
+  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.
if (Instruction *I = IG->getMember(i))
       O << " +\n"
         << Indent << "\"  " << VPlanIngredient(I) << " " << i << "\\l\"";
 }
 
 void VPWidenCallRecipe::execute(VPTransformState &State) {
   State.ILV->widenCallInstruction(Ingredient, User, State);
 }
 
+void VPWidenSelectRecipe::execute(VPTransformState &State) {
+  State.ILV->widenSelectInstruction(Ingredient, InvariantCond);
+}
+
 void VPWidenRecipe::execute(VPTransformState &State) {
   State.ILV->widenInstruction(Ingredient);
 }
 
 void VPWidenGEPRecipe::execute(VPTransformState &State) {
   State.ILV->widenGEP(GEP, State.UF, State.VF, IsPtrLoopInvariant,
                       IsIndexLoopInvariant);
 }
static bool processLoopInVPlanNativePath(
   ScalarEpilogueLowering SEL = getScalarEpilogueLowering(
       F, L, Hints, PSI, BFI, TTI, TLI, AC, LI, PSE.getSE(), DT, *LVL);
 
   LoopVectorizationCostModel CM(SEL, L, PSE, LI, LVL, *TTI, TLI, DB, AC, ORE, F,
                                 &Hints, IAI);
   // Use the planner for outer loop vectorization.
   // TODO: CM is not used at this point inside the planner. Turn CM into an
   // optional argument if we don't need it in the future.
-  LoopVectorizationPlanner LVP(L, LI, TLI, TTI, LVL, CM, IAI);
+  LoopVectorizationPlanner LVP(L, LI, TLI, TTI, LVL, CM, IAI, PSE);
 
   // Get user vectorization factor.
   const unsigned UserVF = Hints.getWidth();
 
   // Plan how to best vectorize, return the best VF and its cost.
   const VectorizationFactor VF = LVP.planInVPlanNativePath(UserVF);
 
   // If we are stress testing VPlan builds, do not attempt to generate vector
#endif /* NDEBUG */
   }
 
   // Use the cost model.
   LoopVectorizationCostModel CM(SEL, L, PSE, LI, &LVL, *TTI, TLI, DB, AC, ORE,
                                 F, &Hints, IAI);
   CM.collectValuesToIgnore();
 
   // Use the planner for vectorization.
-  LoopVectorizationPlanner LVP(L, LI, TLI, TTI, &LVL, CM, IAI);
+  LoopVectorizationPlanner LVP(L, LI, TLI, TTI, &LVL, CM, IAI, PSE);
 
   // Get user vectorization factor.
   unsigned UserVF = Hints.getWidth();
 
   // Plan how to best vectorize, return the best VF and its cost.
   Optional<VectorizationFactor> MaybeVF = LVP.plan(UserVF);
 
   VectorizationFactor VF = VectorizationFactor::Disabled();

llvm/lib/Transforms/Vectorize/VPRecipeBuilder.h

class VPRecipeBuilder {
   const TargetLibraryInfo *TLI;
 
   /// The legality analysis.
   LoopVectorizationLegality *Legal;
 
   /// The profitablity analysis.
   LoopVectorizationCostModel &CM;
 
+  PredicatedScalarEvolution &PSE;
+
   VPBuilder &Builder;
 
   /// When we if-convert we need to create edge masks. We have to cache values
   /// so that we don't end up with exponential recursion/IR. Note that
   /// if-conversion currently takes place during VPlan-construction, so these
   /// caches are only used at that stage.
   using EdgeMaskCacheTy =
       DenseMap<std::pair<BasicBlock *, BasicBlock *>, VPValue *>;
public:
   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);
 
   /// 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, VPBuilder &Builder)
-      : OrigLoop(OrigLoop), TLI(TLI), Legal(Legal), CM(CM), Builder(Builder) {}
+                  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);
 
   /// 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

llvm/lib/Transforms/Vectorize/VPlan.h

using VPRecipeTy = enum {
     VPReplicateSC,
     VPWidenCallSC,
     VPWidenCanonicalIVSC,
     VPWidenGEPSC,
     VPWidenIntOrFpInductionSC,
     VPWidenMemoryInstructionSC,
     VPWidenPHISC,
     VPWidenSC,
+    VPWidenSelectSC
   };
 
   VPRecipeBase(const unsigned char SC) : SubclassID(SC) {}
   virtual ~VPRecipeBase() = default;
 
   /// \return an ID for the concrete type of this object.
   /// This is used to implement the classof checks. This should not be used
   /// for any other purpose, as the values may change as LLVM evolves.
public:
   /// Produce a widened version of the call instruction.
   void execute(VPTransformState &State) override;
 
   /// Print the recipe.
   void print(raw_ostream &O, const Twine &Indent,
              VPSlotTracker &SlotTracker) const override;
 };
 
+/// A recipe for widening select instructions.
+class VPWidenSelectRecipe : public VPRecipeBase {
+private:
+  /// Hold the select to be widened.
+  SelectInst &Ingredient;
+
+  /// Is the condition of the select loop invariant?
+  bool InvariantCond;
+
+  /// Hold VPValues for the arguments of the call.
+  VPUser User;
+
+public:
+  VPWidenSelectRecipe(SelectInst &I, bool InvariantCond)
+      : VPRecipeBase(VPWidenSelectSC), Ingredient(I),
+        InvariantCond(InvariantCond) {}
+
+  ~VPWidenSelectRecipe() override = default;
+
+  /// Method to support type inquiry through isa, cast, and dyn_cast.
+  static inline bool classof(const VPRecipeBase *V) {
+    return V->getVPRecipeID() == VPRecipeBase::VPWidenSelectSC;
+  }
+
+  /// Produce a widened version of the select instruction.
+  void execute(VPTransformState &State) override;
+
+  /// Print the recipe.
+  void print(raw_ostream &O, const Twine &Indent,
+             VPSlotTracker &SlotTracker) const override;
+};
+
 /// A recipe for handling GEP instructions.
 class VPWidenGEPRecipe : public VPRecipeBase {
   GetElementPtrInst *GEP;
   bool IsPtrLoopInvariant;
   SmallBitVector IsIndexLoopInvariant;
 
 public:
   VPWidenGEPRecipe(GetElementPtrInst *GEP, Loop *OrigLoop)

llvm/lib/Transforms/Vectorize/VPlan.cpp

 }
 
 void VPWidenCallRecipe::print(raw_ostream &O, const Twine &Indent,
                               VPSlotTracker &SlotTracker) const {
   O << " +\n"
     << Indent << "\"WIDEN-CALL " << VPlanIngredient(&Ingredient) << "\\l\"";
 }
 
+void VPWidenSelectRecipe::print(raw_ostream &O, const Twine &Indent,
+                                VPSlotTracker &SlotTracker) const {
+  O << " +\n"
+    << Indent << "\"WIDEN-SELECT" << VPlanIngredient(&Ingredient)
+    << (InvariantCond ? " (condition is loop invariant)" : "") << "\\l\"";
+}
+
 void VPWidenRecipe::print(raw_ostream &O, const Twine &Indent,
                           VPSlotTracker &SlotTracker) const {
   O << " +\n" << Indent << "\"WIDEN\\l\"";
   O << "\"  " << VPlanIngredient(&Ingredient) << "\\l\"";
 }
 
 void VPWidenIntOrFpInductionRecipe::print(raw_ostream &O, const Twine &Indent,
                                           VPSlotTracker &SlotTracker) const {