[indvars] Canonicalize exit conditions to unsigned using range info

[reames]

Summary

This is a companion patch to D111836. They share a bunch of common code, and whichever one lands second will be rebased over the other.

This patch duplicates a bit of logic we apply to comparisons encountered during the IV users walk to conditions which feed exit conditions. Why? simplifyAndExtend has a very limited list of users it walks. In particular, in the examples is stops at the zext and never visits the icmp. (Because we can't fold the zext to an addrec yet in SCEV.) Being willing to visit when we haven't simplified regresses multiple tests (seemingly because of less optimal results when computing trip counts).

Note that this can be trivially extended to multiple exiting blocks. I'm leaving that to a future patch (solely to cut down on the number of versions of the same code in review at once.)

[nikic]

My main question here would be why we need this in IndVars. This doesn't really seem to be related to IV or loop based reasoning, but is a generic range based transform. Even InstCombine can handle this (based on known bits): https://godbolt.org/z/drGfM4K9o And as Roman mentioned on the other review, CVP could do this based on ranges as well.

Now, if this case simply fell out of existing code (i.e. if the "visiting more users" approach worked) I'd understand, but it's not super clear to me why we need to go out of the way to handle this here. Is this addressing some kind of phase ordering issue?

[reames]

In D111896#3067338, @nikic wrote:

My main question here would be why we need this in IndVars. This doesn't really seem to be related to IV or loop based reasoning, but is a generic range based transform. Even InstCombine can handle this (based on known bits): https://godbolt.org/z/drGfM4K9o And as Roman mentioned on the other review, CVP could do this based on ranges as well.

Now, if this case simply fell out of existing code (i.e. if the "visiting more users" approach worked) I'd understand, but it's not super clear to me why we need to go out of the way to handle this here. Is this addressing some kind of phase ordering issue?

Well, the motivation for doing this here is that the companion patch *is* inherently loop based. It feels silly to leave cases on the floor just because we might get them elsewhere when we already have a need for a loop based transform in exactly this place.

[mkazantsev]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1428

Loop exiting branch cannot be unconditional.

1445

isRelational && isSigned ?

1848

Do we really need this? We don't fold any blocks here, just replacing signed predicate with unsigned. How can this be a problem for caches?

[reames]

address review comments

[reames]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1848

We don't, though for somewhat subtle reasons. In theory, switching the predicate could let us compute a trip count, but since we haven't *changed* the trip count of the loop, at worst we've left SCEV in an imprecise state.

I removed the invalidation for now, we may want to add it back if that imprecision turns out to be a problem later.

[mkazantsev]

Looks fine now, thanks! Suggestion for follow-up in comments.

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1441

Can we end up with non-canonicalized comparison here? I.e. LHS is invariant and RHS is not. If so, we can swap them. OK if done in follow-up.

[reames]

I went ahead and pushed a follow on to this without review. 0836a105 extends this code to handle multiple exiting blocks. The change is trivial; as noted in the review description here, it was left out only to make keeping a couple of reviews in sync reasonable.

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1441

I think we can. To be honest, I'd rather handle that with a general canonicalization towards invariant on rhs, but that'll require a bit more thought. I won't get to this immediately, but I wrote it down to come back to in a few days.

[reames]

In D111896#3067338, @nikic wrote:

Now, if this case simply fell out of existing code (i.e. if the "visiting more users" approach worked) I'd understand, but it's not super clear to me why we need to go out of the way to handle this here. Is this addressing some kind of phase ordering issue?

I went back and took another look at the simplifyAndExtend option. I think I managed to get that working; we'll see if the compile time is reasonable. See https://reviews.llvm.org/D112170.

Assuming that works, we'll have some code to cleanup in Scalar/IndVarSimplify.cpp. Some of this transform will disappear, but so will a good amount of the existing code in optimizeLoopExits.

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

 #include "llvm/Transforms/Utils/LoopUtils.h"
 #include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"
 #include "llvm/Transforms/Utils/SimplifyIndVar.h"
 #include <cassert>
 #include <cstdint>
 #include <utility>
 
 using namespace llvm;
+using namespace PatternMatch;
 
 #define DEBUG_TYPE "indvars"
 
 STATISTIC(NumWidened     , "Number of indvars widened");
 STATISTIC(NumReplaced    , "Number of exit values replaced");
 STATISTIC(NumLFTR        , "Number of loop exit tests replaced");
 STATISTIC(NumElimExt     , "Number of IV sign/zero extends eliminated");
 STATISTIC(NumElimIV      , "Number of congruent IVs eliminated");
class IndVarSimplify {
 
   SmallVector<WeakTrackingVH, 16> DeadInsts;
   bool WidenIndVars;
 
   bool handleFloatingPointIV(Loop *L, PHINode *PH);
   bool rewriteNonIntegerIVs(Loop *L);
 
   bool simplifyAndExtend(Loop *L, SCEVExpander &Rewriter, LoopInfo *LI);
+  /// See if we can convert an exit condition from signed to unsigned.
+  /// (See inline comment about why this is duplicated from simplifyAndExtend)
+  bool canonicalizeExitCondition(Loop *L);
   /// Try to eliminate loop exits based on analyzeable exit counts
   bool optimizeLoopExits(Loop *L, SCEVExpander &Rewriter);
   /// Try to form loop invariant tests for loop exits by changing how many
   /// iterations of the loop run when that is unobservable.
   bool predicateLoopExits(Loop *L, SCEVExpander &Rewriter);
 
   bool rewriteFirstIterationLoopExitValues(Loop *L);
 

 
 static bool optimizeLoopExitWithUnknownExitCount(
     const Loop *L, BranchInst *BI, BasicBlock *ExitingBB,
     const SCEV *MaxIter, bool Inverted, bool SkipLastIter,
     ScalarEvolution *SE, SCEVExpander &Rewriter,
     SmallVectorImpl<WeakTrackingVH> &DeadInsts) {
   ICmpInst::Predicate Pred;
   Value *LHS, *RHS;
-  using namespace PatternMatch;
   BasicBlock *TrueSucc, *FalseSucc;
   if (!match(BI, m_Br(m_ICmp(Pred, m_Value(LHS), m_Value(RHS)),
                       m_BasicBlock(TrueSucc), m_BasicBlock(FalseSucc))))
     return false;
 
   assert((L->contains(TrueSucc) != L->contains(FalseSucc)) &&
          "Not a loop exit!");
 
if (SE->isKnownPredicateAt(LIP->Pred, LIP->LHS, LIP->RHS, BI))
     foldExit(L, ExitingBB, Inverted, DeadInsts);
   else
     replaceWithInvariantCond(L, ExitingBB, LIP->Pred, LIP->LHS, LIP->RHS,
                              Rewriter, DeadInsts);
 
   return true;
 }
 
+bool IndVarSimplify::canonicalizeExitCondition(Loop *L) {
+  // Note: This is duplicating a particular part on SimplifyIndVars reasoning.
+  // We need to duplicate it because given icmp zext(small-iv), C, IVUsers
+  // never reaches the icmp since the zext doesn't fold to an AddRec unless
+  // it already has flags.  The alternative to this would be to extending the
+  // set of "interesting" IV users to include the icmp, but doing that
+  // regresses results in practice by querying SCEVs before trip counts which
+  // rely on them which results in SCEV caching sub-optimal answers.  The
+  // concern about caching sub-optimal results is why we only query SCEVs of
+  // the loop invariant RHS here.
+
+  auto *ExitingBB = L->getExitingBlock();
+  if (!ExitingBB)
+    return false;
+  auto *BI = dyn_cast<BranchInst>(ExitingBB->getTerminator());
+  if (!BI)
+    return false;
+  assert(BI->isConditional() && "exit branch must be conditional");
+
+  auto *ICmp = dyn_cast<ICmpInst>(BI->getCondition());
+  if (!ICmp)
+    return false;
+
+  auto *LHS = ICmp->getOperand(0);
+  auto *RHS = ICmp->getOperand(1);
+  // Avoid computing SCEVs in the loop to avoid poisoning cache with
+  // sub-optimal results.
+  if (!L->isLoopInvariant(RHS))
+    return false;
+
+  // Match (icmp signed-cond zext, RHS)
+  Value *LHSOp = nullptr;
+  if (!match(LHS, m_ZExt(m_Value(LHSOp))) || !ICmp->isSigned())
+    return false;
+
+  const DataLayout &DL = ExitingBB->getModule()->getDataLayout();
+  const unsigned InnerBitWidth = DL.getTypeSizeInBits(LHSOp->getType());
+  const unsigned OuterBitWidth = DL.getTypeSizeInBits(RHS->getType());
+  auto FullCR = ConstantRange::getFull(InnerBitWidth);
+  FullCR = FullCR.zeroExtend(OuterBitWidth);
+  if (!FullCR.contains(SE->getUnsignedRange(SE->getSCEV(RHS))))
+    return false;
+
+  // We have now matched icmp signed-cond zext(X), zext(Y'), and can thus
+  // replace the signed condition with the unsigned version.
+  ICmp->setPredicate(ICmp->getUnsignedPredicate());
+  return true;
+}
+
 bool IndVarSimplify::optimizeLoopExits(Loop *L, SCEVExpander &Rewriter) {
   SmallVector<BasicBlock*, 16> ExitingBlocks;
   L->getExitingBlocks(ExitingBlocks);
 
   // Remove all exits which aren't both rewriteable and execute on every
   // iteration.
   llvm::erase_if(ExitingBlocks, [&](BasicBlock *ExitingBB) {
     // If our exitting block exits multiple loops, we can only rewrite the
if (int Rewrites = rewriteLoopExitValues(L, LI, TLI, SE, TTI, Rewriter, DT,
       NumReplaced += Rewrites;
       Changed = true;
     }
   }
 
   // Eliminate redundant IV cycles.
   NumElimIV += Rewriter.replaceCongruentIVs(L, DT, DeadInsts);
 
+  if (canonicalizeExitCondition(L))
+    // We've changed the predicate, but have not changed exit counts, or the
+    // values which can flow through any SCEV.  i.e, no invalidation needed.
+    Changed = true;
+
   // Try to eliminate loop exits based on analyzeable exit counts
   if (optimizeLoopExits(L, Rewriter))  {
     Changed = true;
     // Given we've changed exit counts, notify SCEV
     // Some nested loops may share same folded exit basic block,
     // thus we need to notify top most loop.
     SE->forgetTopmostLoop(L);
   }
 

llvm/test/Transforms/IndVarSimplify/finite-exit-comparisons.ll

 define void @slt_constant_rhs(i16 %n.raw, i8 %start) mustprogress {
 ; CHECK-LABEL: @slt_constant_rhs(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i8 [ [[IV_NEXT:%.*]], [[FOR_BODY]] ], [ [[START:%.*]], [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[IV_NEXT]] = add i8 [[IV]], 1
 ; CHECK-NEXT:    [[ZEXT:%.*]] = zext i8 [[IV_NEXT]] to i16
-; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i16 [[ZEXT]], 254
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i16 [[ZEXT]], 254
 ; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY]], label [[FOR_END:%.*]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %entry, %for.body

 ; CHECK-LABEL: @slt_non_constant_rhs_no_mustprogress(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[N:%.*]] = and i16 [[N_RAW:%.*]], 255
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i8 [ [[IV_NEXT:%.*]], [[FOR_BODY]] ], [ 0, [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[IV_NEXT]] = add i8 [[IV]], 1
 ; CHECK-NEXT:    [[ZEXT:%.*]] = zext i8 [[IV_NEXT]] to i16
-; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i16 [[ZEXT]], [[N]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i16 [[ZEXT]], [[N]]
 ; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY]], label [[FOR_END:%.*]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   %n = and i16 %n.raw, 255
   br label %for.body
 

 define void @sgt_constant_rhs(i16 %n.raw, i8 %start) mustprogress {
 ; CHECK-LABEL: @sgt_constant_rhs(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i8 [ [[IV_NEXT:%.*]], [[FOR_BODY]] ], [ [[START:%.*]], [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[IV_NEXT]] = add i8 [[IV]], 1
 ; CHECK-NEXT:    [[ZEXT:%.*]] = zext i8 [[IV_NEXT]] to i16
-; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i16 [[ZEXT]], 254
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i16 [[ZEXT]], 254
 ; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY]], label [[FOR_END:%.*]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %entry, %for.body

 define void @sle_constant_rhs(i16 %n.raw, i8 %start) mustprogress {
 ; CHECK-LABEL: @sle_constant_rhs(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i8 [ [[IV_NEXT:%.*]], [[FOR_BODY]] ], [ [[START:%.*]], [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[IV_NEXT]] = add i8 [[IV]], 1
 ; CHECK-NEXT:    [[ZEXT:%.*]] = zext i8 [[IV_NEXT]] to i16
-; CHECK-NEXT:    [[CMP:%.*]] = icmp sle i16 [[ZEXT]], 254
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ule i16 [[ZEXT]], 254
 ; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY]], label [[FOR_END:%.*]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %entry, %for.body

 define void @sge_constant_rhs(i16 %n.raw, i8 %start) mustprogress {
 ; CHECK-LABEL: @sge_constant_rhs(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.body:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i8 [ [[IV_NEXT:%.*]], [[FOR_BODY]] ], [ [[START:%.*]], [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[IV_NEXT]] = add i8 [[IV]], 1
 ; CHECK-NEXT:    [[ZEXT:%.*]] = zext i8 [[IV_NEXT]] to i16
-; CHECK-NEXT:    [[CMP:%.*]] = icmp sge i16 [[ZEXT]], 254
+; CHECK-NEXT:    [[CMP:%.*]] = icmp uge i16 [[ZEXT]], 254
 ; CHECK-NEXT:    br i1 [[CMP]], label [[FOR_BODY]], label [[FOR_END:%.*]]
 ; CHECK:       for.end:
 ; CHECK-NEXT:    ret void
 ;
 entry:
   br label %for.body
 
 for.body:                                         ; preds = %entry, %for.body