[LV] Move exit cond simplification to separate transform.

[fhahn]

Summary

This sets the stage for D133017 by moving out the code that performs
VPlan based simplifications to a separate transform that takes the
chosen VF & UF as arguments.

The main advantage is that this transform runs before any changes to
the CFG are being made. This allows using SCEV without worrying about
making queries while the IR is in an incomplete state.

Note that this patch switches the reasoning to use SCEV, but still only
simplifies loops with constant trip counts. Using SCEV here is needed to
access the backedge taken count, because the trip count IR value has not
been created yet.

[Ayal]

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

464

TripCountV indeed hasn't been created yet, but perhaps there's a simpler way to check this condition, e.g., by checking if PSE.getBackedgeCount() is a constant greater than VF*UF-1, instead of incrementing it by 1 of type IdxTy; or trying to call SE.getSmallConstantTripCount(L); or recording it in VPlan, possibly in CanonicalIV recipe.

474

Worth an unconditional branch VPInstruction instead of the effort to generate a redundant True? Can be fixed independent of this patch, which is working a bit harder to generate it.

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

66

Missing documentation.

llvm/test/Transforms/LoopVectorize/lcssa-crashes.ll

16

Is this a regression from previously folding the conditional branch?

[Ayal]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

7623

nit: can drop "Concrete".

Restrict the VF range of BestPlan to include only BestVF?

When printing BestPlan would be good to now indicate that it applies to BestUF only.

[fhahn]

Address latest comments, thanks!

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

7623

nit: can drop "Concrete".

Done!

Restrict the VF range of BestPlan to include only BestVF?

Done!

When printing BestPlan would be good to now indicate that it applies to BestUF only.

Should I add a field for that?

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

464

That's what I started out with, but the trip count (BTC + 1) allows additional simplifications in D133017 as BTC's often involve a -1, which means it may be negative limiting SCEV simplifications.

474

It should probably be sufficient to just set a single successor. We will have to also replace & remove the header phis.

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

66

added, thanks!

llvm/test/Transforms/LoopVectorize/lcssa-crashes.ll

16

The test case is spurious, the vector body will never be executed, as the exit condition compares to undef.

[Ayal]

Looks good to me, adding various nits.

In the summary, perhaps indicate that the optimization is moved earlier to take place before creating the skeleton, as soon as the UF is available.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

7623

When printing BestPlan would be good to now indicate that it applies to BestUF only.

Should I add a field for that?

Probably, to then print it instead of "UF >= 1". Can also add setUF(), and an assert when executing (the epilog loop) that the plan used works for the current UF.

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

454

So we retain "if TC <= VF * UF" ?

456

nit: can fold this comment and logic into "canSimplifyLoopBranch(Term)" (independent of this patch).

466

nit: ExitCount >> TripCount ?

nit: where is createTripCountSCEV defined? Should VPlan or its loop Region provide it, similar to providing the canonical IV, or even have the canonical IV recipe provide it? When optimizing the loop branch recipe, one can also optimize the canonical IV recipe...

469

nit: is the isZero() check needed?

474

OK, right, there should be no need for an unconditional branch VPInstruction, once CFG is straightened, as in the 2nd TODO.

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

67

nit: may restrict >> restricts ?

llvm/test/Transforms/LoopVectorize/lcssa-crashes.ll

16

nit: worth fixing the test somehow, to clarify that no folding is lost here?

[fhahn]

Rebase and address nits, thanks! I am planning on landing this shortly.

[fhahn]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

7623

Thanks, I landed the following improvements separately:

• move generating of VFs & UFs as part of getName to ensure it always prints the up-to-date ones: 96296922b61f
• tracking of UF in VPlan: 5df34e971dd5
• add asserts in executePlan that plan has selected VF + UF: b7b1e5c96fe0

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

454

Yes, updated!

466

nit: ExitCount >> TripCount ?

Fixed, thanks!

nit: where is createTripCountSCEV defined?

At the moment it is a regular function, as it is also used outside Plan. But once the scaffolding-generation is also handled by VPlan it could be moved there.

469

Yes for now I think; if it is zero then the +1 overflowed and this check is the same as in the previous version of the code.

[Ayal]

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

469

Right, this is what the previous version was doing, and zero stands for a wrapped-around maximal BTC. But CheckMinIters bails out in such (Max iteration) cases nonetheless, w/o reaching the latch? Probably deserves a comment and/or could be simplified into C->isZero().

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

 void LoopVectorizationPlanner::executePlan(ElementCount BestVF, unsigned BestUF,
                                            VPlan &BestVPlan,
                                            InnerLoopVectorizer &ILV,
                                            DominatorTree *DT,
                                            bool IsEpilogueVectorization) {
   LLVM_DEBUG(dbgs() << "Executing best plan with VF=" << BestVF << ", UF=" << BestUF
                     << '\n');
 
+  if (!IsEpilogueVectorization)
+    VPlanTransforms::optimizeForConcreteVFAndUF(BestVPlan, BestVF, BestUF, PSE);
+
   // Perform the actual loop transformation.
 
   // 1. Set up the skeleton for vectorization, including vector pre-header and
   // middle block. The vector loop is created during VPlan execution.
   VPTransformState State{BestVF, BestUF, LI, DT, ILV.Builder, &ILV, &BestVPlan};
   Value *CanonicalIVStartValue;
   std::tie(State.CFG.PrevBB, CanonicalIVStartValue) =
       ILV.createVectorizedLoopSkeleton();

llvm/lib/Transforms/Vectorize/VPlan.cpp

VPActiveLaneMaskPHIRecipe *VPlan::getActiveLaneMaskPhi() {
   VPBasicBlock *Header = getVectorLoopRegion()->getEntryBasicBlock();
   for (VPRecipeBase &R : Header->phis()) {
     if (isa<VPActiveLaneMaskPHIRecipe>(&R))
       return cast<VPActiveLaneMaskPHIRecipe>(&R);
   }
   return nullptr;
 }
 
-static bool canSimplifyBranchOnCond(VPInstruction *Term) {
-  VPInstruction *Not = dyn_cast<VPInstruction>(Term->getOperand(0));
-  if (!Not || Not->getOpcode() != VPInstruction::Not)
-    return false;
-
-  VPInstruction *ALM = dyn_cast<VPInstruction>(Not->getOperand(0));
-  return ALM && ALM->getOpcode() == VPInstruction::ActiveLaneMask;
-}
-
 void VPlan::prepareToExecute(Value *TripCountV, Value *VectorTripCountV,
                              Value *CanonicalIVStartValue,
                              VPTransformState &State,
                              bool IsEpilogueVectorization) {
 
-  VPBasicBlock *ExitingVPBB = getVectorLoopRegion()->getExitingBasicBlock();
-  auto *Term = dyn_cast<VPInstruction>(&ExitingVPBB->back());
-  // Try to simplify the branch condition if TC <= VF * UF when preparing to
-  // execute the plan for the main vector loop. We only do this if the
-  // terminator is:
-  //  1. BranchOnCount, or
-  //  2. BranchOnCond where the input is Not(ActiveLaneMask).
-  if (!IsEpilogueVectorization && Term && isa<ConstantInt>(TripCountV) &&
-      (Term->getOpcode() == VPInstruction::BranchOnCount ||
-       (Term->getOpcode() == VPInstruction::BranchOnCond &&
-        canSimplifyBranchOnCond(Term)))) {
-    ConstantInt *C = cast<ConstantInt>(TripCountV);
-    uint64_t TCVal = C->getZExtValue();
-    if (TCVal && TCVal <= State.VF.getKnownMinValue() * State.UF) {
-      auto *BOC =
-          new VPInstruction(VPInstruction::BranchOnCond,
-                            {getOrAddExternalDef(State.Builder.getTrue())});
-      Term->eraseFromParent();
-      ExitingVPBB->appendRecipe(BOC);
-      // TODO: Further simplifications are possible
-      //      1. Replace inductions with constants.
-      //      2. Replace vector loop region with VPBasicBlock.
-    }
-  }
-
   // Check if the trip count is needed, and if so build it.
   if (TripCount && TripCount->getNumUsers()) {
     for (unsigned Part = 0, UF = State.UF; Part < UF; ++Part)
       State.set(TripCount, TripCountV, Part);
   }
 
   // Check if the backedge taken count is needed, and if so build it.
   if (BackedgeTakenCount && BackedgeTakenCount->getNumUsers()) {

llvm/lib/Transforms/Vectorize/VPlanTransforms.h

 
 namespace llvm {
 
 class InductionDescriptor;
 class Instruction;
 class PHINode;
 class ScalarEvolution;
 class Loop;
+class PredicatedScalarEvolution;
 class TargetLibraryInfo;
 
 struct VPlanTransforms {
   /// Replaces the VPInstructions in \p Plan with corresponding
   /// widen recipes.
   static void
   VPInstructionsToVPRecipes(Loop *OrigLoop, VPlanPtr &Plan,
                             function_ref<const InductionDescriptor *(PHINode *)>
struct VPlanTransforms {
   /// provide them by building scalar steps off of the canonical scalar IV and
   /// update the original IV's users. This is an optional optimization to reduce
   /// the needs of vector extracts.
   static void optimizeInductions(VPlan &Plan, ScalarEvolution &SE);
 
   /// Remove redundant EpxandSCEVRecipes in \p Plan's entry block by replacing
   /// them with already existing recipes expanding the same SCEV expression.
   static void removeRedundantExpandSCEVRecipes(VPlan &Plan);
+
+  static void optimizeForConcreteVFAndUF(VPlan &Plan, ElementCount VF,
+                                         unsigned UF,
+                                         PredicatedScalarEvolution &PSE);
 };
 
 } // namespace llvm
 
 #endif // LLVM_TRANSFORMS_VECTORIZE_VPLANTRANSFORMS_H

llvm/lib/Transforms/Vectorize/VPlanTransforms.cpp

for (VPRecipeBase &R :
 
     auto I = SCEV2VPV.insert({ExpR->getSCEV(), ExpR});
     if (I.second)
       continue;
     ExpR->replaceAllUsesWith(I.first->second);
     ExpR->eraseFromParent();
   }
 }
+
+static bool canSimplifyBranchOnCond(VPInstruction *Term) {
+  VPInstruction *Not = dyn_cast<VPInstruction>(Term->getOperand(0));
+  if (!Not || Not->getOpcode() != VPInstruction::Not)
+    return false;
+
+  VPInstruction *ALM = dyn_cast<VPInstruction>(Not->getOperand(0));
+  return ALM && ALM->getOpcode() == VPInstruction::ActiveLaneMask;
+}
+
+void VPlanTransforms::optimizeForConcreteVFAndUF(
+    VPlan &Plan, ElementCount VF, unsigned UF, PredicatedScalarEvolution &PSE) {
+  VPBasicBlock *ExitingVPBB =
+      Plan.getVectorLoopRegion()->getExitingBasicBlock();
+  auto *Term = dyn_cast<VPInstruction>(&ExitingVPBB->back());
+  // Try to simplify the branch condition if TC <= VF * UF when preparing to
+  // execute the plan for the main vector loop. We only do this if the
+  // terminator is:
+  //  1. BranchOnCount, or
+  //  2. BranchOnCond where the input is Not(ActiveLaneMask).
+  if (!Term || (Term->getOpcode() != VPInstruction::BranchOnCount &&
+                (Term->getOpcode() != VPInstruction::BranchOnCond ||
+                 !canSimplifyBranchOnCond(Term))))
+    return;
+
+  Type *IdxTy =
+      Plan.getCanonicalIV()->getStartValue()->getLiveInIRValue()->getType();
+  const SCEV *ExitCount = createTripCountSCEV(IdxTy, PSE);
+  auto *C = dyn_cast<SCEVConstant>(ExitCount);
+  ScalarEvolution &SE = *PSE.getSE();
+  if (!C || ExitCount->isZero() ||
+      C->getAPInt().getZExtValue() > VF.getKnownMinValue() * UF)
+    return;
+
+  LLVMContext &Ctx = SE.getContext();
+  auto *BOC =
+      new VPInstruction(VPInstruction::BranchOnCond,
+                        {Plan.getOrAddExternalDef(ConstantInt::getTrue(Ctx))});
+  Term->eraseFromParent();
+  ExitingVPBB->appendRecipe(BOC);
+  // TODO: Further simplifications are possible
+  //      1. Replace inductions with constants.
+  //      2. Replace vector loop region with VPBasicBlock.
+}

llvm/test/Transforms/LoopVectorize/lcssa-crashes.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -loop-vectorize -force-vector-interleave=1 -force-vector-width=4 -S %s | FileCheck %s
 
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
 
 define void @test() {
 ; CHECK-LABEL: @test(
 ; CHECK-NEXT:    br label [[FOR_BODY_LR_PH_I_I_I:%.*]]
 ; CHECK:       for.body.lr.ph.i.i.i:
 ; CHECK-NEXT:    br i1 true, label [[SCALAR_PH:%.*]], label [[VECTOR_PH:%.*]]
 ; CHECK:       vector.ph:
 ; CHECK-NEXT:    br label [[VECTOR_BODY:%.*]]
 ; CHECK:       vector.body:
 ; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[VECTOR_PH]] ], [ [[INDEX_NEXT:%.*]], [[VECTOR_BODY]] ]
 ; CHECK-NEXT:    [[INDEX_NEXT]] = add nuw i64 [[INDEX]], 4
-; CHECK-NEXT:    br i1 true, label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP0:![0-9]+]]
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i64 [[INDEX_NEXT]], 0
+; CHECK-NEXT:    br i1 [[TMP1]], label [[MIDDLE_BLOCK:%.*]], label [[VECTOR_BODY]], !llvm.loop [[LOOP0:![0-9]+]]
 ; CHECK:       middle.block:
 ; CHECK-NEXT:    [[CMP_N:%.*]] = icmp eq i64 1, 0
 ; CHECK-NEXT:    br i1 [[CMP_N]], label [[FOR_END_I_I_I:%.*]], label [[SCALAR_PH]]
 ; CHECK:       scalar.ph:
 ; CHECK-NEXT:    [[BC_RESUME_VAL:%.*]] = phi i64 [ 0, [[MIDDLE_BLOCK]] ], [ 0, [[FOR_BODY_LR_PH_I_I_I]] ]
 ; CHECK-NEXT:    br label [[FOR_BODY_I_I_I:%.*]]
 ; CHECK:       for.body.i.i.i:
 ; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[INDVARS_IV_NEXT:%.*]], [[FOR_INC_I_I_I:%.*]] ], [ [[BC_RESUME_VAL]], [[SCALAR_PH]] ]