[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

481

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.

491

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–17

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

481

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.

491

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–17

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

471

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

473

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

483

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...

486

nit: is the isZero() check needed?

491

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–17

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

471

Yes, updated!

483

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.

486

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

486

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

assert(BestVPlan.hasVF(BestVF) &&
          "Trying to execute plan with unsupported VF");
   assert(BestVPlan.hasUF(BestUF) &&
          "Trying to execute plan with unsupported UF");
 
   LLVM_DEBUG(dbgs() << "Executing best plan with VF=" << BestVF << ", UF=" << BestUF
                     << '\n');
 
   // Workaround!  Compute the trip count of the original loop and cache it
   // before we start modifying the CFG.  This code has a systemic problem
   // wherein it tries to run analysis over partially constructed IR; this is
   // wrong, and not simply for SCEV.  The trip count of the original loop
   // simply happens to be prone to hitting this in practice.  In theory, we
   // can hit the same issue for any SCEV, or ValueTracking query done during
   // mutation.  See PR49900.
   ILV.getOrCreateTripCount(OrigLoop->getLoopPreheader());
 
+  if (!IsEpilogueVectorization)
+    VPlanTransforms::optimizeForVFAndUF(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.h

public:
   VPValue &getVectorTripCount() { return VectorTripCount; }
 
   /// Mark the plan to indicate that using Value2VPValue is not safe any
   /// longer, because it may be stale.
   void disableValue2VPValue() { Value2VPValueEnabled = false; }
 
   void addVF(ElementCount VF) { VFs.insert(VF); }
 
+  void setVF(ElementCount VF) {
+    assert(hasVF(VF) && "Cannot set VF not already in plan");
+    VFs.clear();
+    VFs.insert(VF);
+  }
+
   bool hasVF(ElementCount VF) { return VFs.count(VF); }
 
   bool hasScalarVFOnly() const { return VFs.size() == 1 && VFs[0].isScalar(); }
 
   bool hasUF(unsigned UF) const { return UFs.empty() || UFs.contains(UF); }
 
   void setUF(unsigned UF) {
     assert(hasUF(UF) && "Cannot set the UF not already in plan");

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);
+
+  /// Optimize \p Plan based on \p BestVF and \p BestUF. This may restrict the
+  /// resulting plan to \p BestVF and \p BestUF.
+  static void optimizeForVFAndUF(VPlan &Plan, ElementCount BestVF,
+                                 unsigned BestUF,
+                                 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::optimizeForVFAndUF(VPlan &Plan, ElementCount BestVF,
+                                         unsigned BestUF,
+                                         PredicatedScalarEvolution &PSE) {
+  assert(Plan.hasVF(BestVF) && "BestVF is not available in Plan");
+  assert(Plan.hasUF(BestUF) && "BestUF is not available in Plan");
+  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 *TripCount = createTripCountSCEV(IdxTy, PSE);
+  auto *C = dyn_cast<SCEVConstant>(TripCount);
+  ScalarEvolution &SE = *PSE.getSE();
+  if (!C || TripCount->isZero() ||
+      C->getAPInt().getZExtValue() > BestVF.getKnownMinValue() * BestUF)
+    return;
+
+  LLVMContext &Ctx = SE.getContext();
+  auto *BOC =
+      new VPInstruction(VPInstruction::BranchOnCond,
+                        {Plan.getOrAddExternalDef(ConstantInt::getTrue(Ctx))});
+  Term->eraseFromParent();
+  ExitingVPBB->appendRecipe(BOC);
+  Plan.setVF(BestVF);
+  Plan.setUF(BestUF);
+  // 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 -passes=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]] ]