[IndVarSimplify] Use control-dependent range information to prove non-negativity

[apilipenko]

Summary

This change is motivated by the case when IndVarSimplify doesn't widen a comparison of IV increment because it can't prove IV increment being non-negative. We end up with a redundant trunc of the widened increment on this example.

for.body:
  %i = phi i32 [ %start, %for.body.lr.ph ], [ %i.inc, %for.inc ]
  %within_limits = icmp ult i32 %i, 64
  br i1 %within_limits, label %continue, label %for.end

continue:
  %i.i64 = zext i32 %i to i64
  %arrayidx = getelementptr inbounds i32, i32* %base, i64 %i.i64
  %val = load i32, i32* %arrayidx, align 4
  br label %for.inc

for.inc:
  %i.inc = add nsw nuw i32 %i, 1
  %cmp = icmp slt i32 %i.inc, %limit
  br i1 %cmp, label %for.body, label %for.end

There is a range check inside of the loop which guarantees the IV to be non-negative. NSW on the increment guarantees that the increment is also non-negative. ScalarEvolution can't prove that the increments non-negative because it doesn't reason about control dependencies inside of the loop. Teach IndVarSimplify to use the range check to prove non-negativity of loop increments.

It's difficult to gather control-dependent information on demand, because when it's needed some of the dominating conditions might be already widened. That's why collect control-dependent ranges before any transformation occurs for all increments of the IV and all uses of these increments.

[sanjoy]

lib/Transforms/Scalar/IndVarSimplify.cpp

927

Current seems unused?

930

You can use PostIncRangeInfos.insert({Key, R}); here.

I'd also drop the braces.

1474

End comment with period.

1521

s/in the form of/of the form/

1594

I'd s/ConstantRange/auto here.

1607

auto * here.

I'd actually just remove TI altogether -- and instead do:

auto *BI = dyn_cast<BranchInst>(DTB->getBlock()->getTerminator());

1618

Add a TODO here about also exploiting false edges in the future (feel free to implement it right now too -- it won't be that much more code).

1623

I'd perhaps use a SmallPtrSet<16> here.

test/Transforms/IndVarSimplify/post-inc-range.ll

8

Can you please add a test case that checks that we do the right thing for a comparison with, say, (%i * 3) + 1 where %i is the IV and there is a range check on %i * 3? That is, check that we're actually using the worklist algorithm above.

[apilipenko]

lib/Transforms/Scalar/IndVarSimplify.cpp

1607

I'll need TI for guards here (see D25739).

[apilipenko]

Address Sanjoy's comments, put under a flag.

[sanjoy]

lgtm

lib/Transforms/Scalar/IndVarSimplify.cpp

class WidenIV {
 
   enum ExtendKind { ZeroExtended, SignExtended, Unknown };
   // A map tracking the kind of extension used to widen each narrow IV 
   // and narrow IV user.
   // Key: pointer to a narrow IV or IV user.
   // Value: the kind of extension used to widen this Instruction.
   DenseMap<AssertingVH<Instruction>, ExtendKind> ExtendKindMap;
 
+  typedef std::pair<AssertingVH<Value>, AssertingVH<Instruction>> DefUserPair;
+  // A map with control-dependent ranges for post increment IV uses. The key is
+  // a pair of IV def and a use of this def denoting the context. The value is
+  // a ConstantRange representing possible values of the def at the given
+  // context.
+  DenseMap<DefUserPair, ConstantRange> PostIncRangeInfos;
+
+  Optional<ConstantRange> getPostIncRangeInfo(Value *Def,
+                                              Instruction *UseI) {
+    DefUserPair Key(Def, UseI);
+    auto It = PostIncRangeInfos.find(Key);
+    return It == PostIncRangeInfos.end()
+               ? Optional<ConstantRange>(None)
+               : Optional<ConstantRange>(It->second);
+  }
+
+  // Calculates PostIncRangeInfos map for the given IV
+  void calculatePostIncRanges(PHINode *OrigPhi);
+  void calculatePostIncRange(Instruction *NarrowDef, Instruction *NarrowUser);
+  void updatePostIncRangeInfo(Value *Def, Instruction *UseI, ConstantRange R) {
+    DefUserPair Key(Def, UseI);
+    ConstantRange Current(R.getBitWidth(), /*isFullSet=*/true);
+    auto It = PostIncRangeInfos.find(Key);
+    if (It == PostIncRangeInfos.end()) {
+      PostIncRangeInfos.insert(std::make_pair(Key, R));
+    } else {
+      It->second = R.intersectWith(It->second);
+    }
+  }
+
 public:
   WidenIV(const WideIVInfo &WI, LoopInfo *LInfo,
           ScalarEvolution *SEv, DominatorTree *DTree,
           SmallVectorImpl<WeakVH> &DI) :
     OrigPhi(WI.NarrowIV),
     WideType(WI.WidestNativeType),
     LI(LInfo),
     L(LI->getLoopFor(OrigPhi->getParent())),
Instruction *WidenIV::widenIVUse(NarrowIVDefUse DU, SCEVExpander &Rewriter) {
   // Returning WideUse pushes it on the worklist.
   return WideUse;
 }
 
 /// Add eligible users of NarrowDef to NarrowIVUsers.
 ///
 void WidenIV::pushNarrowIVUsers(Instruction *NarrowDef, Instruction *WideDef) {
   const SCEV *NarrowSCEV = SE->getSCEV(NarrowDef);
-  bool NeverNegative =
+  bool NonNegativeDef =
       SE->isKnownPredicate(ICmpInst::ICMP_SGE, NarrowSCEV,
                            SE->getConstant(NarrowSCEV->getType(), 0));
   for (User *U : NarrowDef->users()) {
     Instruction *NarrowUser = cast<Instruction>(U);
 
     // Handle data flow merges and bizarre phi cycles.
     if (!Widened.insert(NarrowUser).second)
       continue;
 
-    NarrowIVUsers.emplace_back(NarrowDef, NarrowUser, WideDef, NeverNegative);
+    bool NonNegativeUse = false;
+    if (!NonNegativeDef) {
+      // We might have a control-dependent range information for this context
+      if (auto RangeInfo = getPostIncRangeInfo(NarrowDef, NarrowUser))
+        NonNegativeUse = RangeInfo->getSignedMin().isNonNegative();
+    }
+
+    NarrowIVUsers.emplace_back(NarrowDef, NarrowUser, WideDef,
+                               NonNegativeDef || NonNegativeUse);
   }
 }
 
 /// Process a single induction variable. First use the SCEVExpander to create a
 /// wide induction variable that evaluates to the same recurrence as the
 /// original narrow IV. Then use a worklist to forward traverse the narrow IV's
 /// def-use chain. After widenIVUse has processed all interesting IV users, the
 /// narrow IV will be isolated for removal by DeleteDeadPHIs.
PHINode *WidenIV::createWideIV(SCEVExpander &Rewriter) {
   // An AddRec must have loop-invariant operands. Since this AddRec is
   // materialized by a loop header phi, the expression cannot have any post-loop
   // operands, so they must dominate the loop header.
   assert(
       SE->properlyDominates(AddRec->getStart(), L->getHeader()) &&
       SE->properlyDominates(AddRec->getStepRecurrence(*SE), L->getHeader()) &&
       "Loop header phi recurrence inputs do not dominate the loop");
 
+  // Iterate over IV uses (including transitive ones) looking for IV increments
+  // in the form of 'add nsw %iv, <const>'. For each increment and each use of
+  // the increment calculate control-dependent range information basing on
+  // dominating conditions inside of the loop (e.g. a range check inside of the
+  // loop). Calculated ranges are stored in PostIncRangeInfos map.
+  //
+  // Control-dependent range information is later used to prove that a narrow
+  // definition is not negative (see pushNarrowIVUsers). It's difficult to do
+  // this on demand because when pushNarrowIVUsers needs this information some
+  // of the dominating conditions might be already widened.
+  calculatePostIncRanges(OrigPhi);
+
   // The rewriter provides a value for the desired IV expression. This may
   // either find an existing phi or materialize a new one. Either way, we
   // expect a well-formed cyclic phi-with-increments. i.e. any operand not part
   // of the phi-SCC dominates the loop entry.
   Instruction *InsertPt = &L->getHeader()->front();
   WidePhi = cast<PHINode>(Rewriter.expandCodeFor(AddRec, WideType, InsertPt));
 
   // Remembering the WideIV increment generated by SCEVExpander allows
while (!NarrowIVUsers.empty()) {
 
     // widenIVUse may have removed the def-use edge.
     if (DU.NarrowDef->use_empty())
       DeadInsts.emplace_back(DU.NarrowDef);
   }
   return WidePhi;
 }
 
+void WidenIV::calculatePostIncRange(Instruction *NarrowDef,
+                                    Instruction *NarrowUser) {
+  using namespace llvm::PatternMatch;
+
+  Value *NarrowDefLHS;
+  const APInt *NarrowDefRHS;
+  if (!match(NarrowDef, m_NSWAdd(m_Value(NarrowDefLHS),
+                                 m_APInt(NarrowDefRHS))) ||
+      !NarrowDefRHS->isNonNegative())
+    return;
+
+  auto UpdateRangeFromCondition = [&] (Value *Condition) {
+    CmpInst::Predicate P;
+    Value *CmpRHS;
+    if (!match(Condition, m_ICmp(P, m_Specific(NarrowDefLHS),
+                                 m_Value(CmpRHS))))
+      return;
+
+    ConstantRange CmpRHSRange = SE->getSignedRange(SE->getSCEV(CmpRHS));
+    ConstantRange CmpConstrainedLHSRange =
+            ConstantRange::makeAllowedICmpRegion(P, CmpRHSRange);
+    ConstantRange NarrowDefRange =
+            CmpConstrainedLHSRange.addWithNoSignedWrap(*NarrowDefRHS);
+
+    updatePostIncRangeInfo(NarrowDef, NarrowUser, NarrowDefRange);
+  };
+
+  BasicBlock *NarrowUserBB = NarrowUser->getParent();
+  for (auto *DTB = (*DT)[NarrowUserBB]->getIDom();
+       L->contains(DTB->getBlock());
+       DTB = DTB->getIDom()) {
+    auto TI = DTB->getBlock()->getTerminator();
+
+    auto *BI = dyn_cast<BranchInst>(TI);
+    if (!BI || !BI->isConditional())
+      continue;
+
+    BasicBlockEdge BBE(BI->getParent(), BI->getSuccessor(0));
+    if (!(BBE.isSingleEdge() &&
+          DT->dominates(BBE, NarrowUser->getParent())))
+      continue;
+
+    UpdateRangeFromCondition(BI->getCondition());
+  }
+}
+
+void WidenIV::calculatePostIncRanges(PHINode *OrigPhi) {
+  DenseSet<Instruction *> Visited;
+  SmallVector<Instruction *, 6> Worklist;
+  Worklist.push_back(OrigPhi);
+  Visited.insert(OrigPhi);
+
+  while (!Worklist.empty()) {
+    Instruction *NarrowDef = Worklist.pop_back_val();
+
+    for (Use &U : NarrowDef->uses()) {
+      auto *NarrowUser = cast<Instruction>(U.getUser());
+
+      // Don't go looking outside the current loop.
+      auto *NarrowUserLoop = (*LI)[NarrowUser->getParent()];
+      if (!NarrowUserLoop || !L->contains(NarrowUserLoop))
+        continue;
+
+      if (!Visited.insert(NarrowUser).second)
+        continue;
+
+      Worklist.push_back(NarrowUser);
+
+      calculatePostIncRange(NarrowDef, NarrowUser);
+    }
+  }
+}
+
 //===----------------------------------------------------------------------===//
 //  Live IV Reduction - Minimize IVs live across the loop.
 //===----------------------------------------------------------------------===//
 
 
 //===----------------------------------------------------------------------===//
 //  Simplification of IV users based on SCEV evaluation.
 //===----------------------------------------------------------------------===//

test/Transforms/IndVarSimplify/post-inc-range.ll

+; RUN: opt < %s -indvars -S | FileCheck %s
+
+target datalayout = "p:64:64:64-n32:64"
+
+; When the IV in this loop is widened we want to widen this use as well:
+; icmp slt i32 %i.inc, %limit
+; In order to do this indvars need to prove that the narrow IV def (%i.inc)
+; is not-negative from the range check inside of the loop.
+define void @test(i32* %base, i32 %limit, i32 %start) {
+; CHECK-LABEL: @test(
+; CHECK-NOT: trunc
+
+for.body.lr.ph:
+  br label %for.body
+
+for.body:
+  %i = phi i32 [ %start, %for.body.lr.ph ], [ %i.inc, %for.inc ]
+  %within_limits = icmp ult i32 %i, 64
+  br i1 %within_limits, label %continue, label %for.end
+
+continue:
+  %i.i64 = zext i32 %i to i64
+  %arrayidx = getelementptr inbounds i32, i32* %base, i64 %i.i64
+  %val = load i32, i32* %arrayidx, align 4
+  br label %for.inc
+
+for.inc:
+  %i.inc = add nsw nuw i32 %i, 1
+  %cmp = icmp slt i32 %i.inc, %limit
+  br i1 %cmp, label %for.body, label %for.end
+
+for.end:
+  br label %exit
+
+exit:
+  ret void
+}
+
+define void @test_range_metadata(i32* %array_length_ptr, i32* %base,
+                                 i32 %limit, i32 %start) {
+; CHECK-LABEL: @test_range_metadata(
+; CHECK-NOT: trunc
+
+for.body.lr.ph:
+  br label %for.body
+
+for.body:
+  %i = phi i32 [ %start, %for.body.lr.ph ], [ %i.inc, %for.inc ]
+  %array_length = load i32, i32* %array_length_ptr, !range !{i32 0, i32 64 }
+  %within_limits = icmp ult i32 %i, %array_length
+  br i1 %within_limits, label %continue, label %for.end
+
+continue:
+  %i.i64 = zext i32 %i to i64
+  %arrayidx = getelementptr inbounds i32, i32* %base, i64 %i.i64
+  %val = load i32, i32* %arrayidx, align 4
+  br label %for.inc
+
+for.inc:
+  %i.inc = add nsw nuw i32 %i, 1
+  %cmp = icmp slt i32 %i.inc, %limit
+  br i1 %cmp, label %for.body, label %for.end
+
+for.end:
+  br label %exit
+
+exit:
+  ret void
+}
+
+; Negative version of the test above, we don't know anything about
+; array_length_ptr range.
+define void @test_neg(i32* %array_length_ptr, i32* %base,
+                      i32 %limit, i32 %start) {
+; CHECK-LABEL: @test_neg(
+; CHECK: trunc i64
+
+for.body.lr.ph:
+  br label %for.body
+
+for.body:
+  %i = phi i32 [ %start, %for.body.lr.ph ], [ %i.inc, %for.inc ]
+  %array_length = load i32, i32* %array_length_ptr
+  %within_limits = icmp ult i32 %i, %array_length
+  br i1 %within_limits, label %continue, label %for.end
+
+continue:
+  %i.i64 = zext i32 %i to i64
+  %arrayidx = getelementptr inbounds i32, i32* %base, i64 %i.i64
+  %val = load i32, i32* %arrayidx, align 4
+  br label %for.inc
+
+for.inc:
+  %i.inc = add nsw nuw i32 %i, 1
+  %cmp = icmp slt i32 %i.inc, %limit
+  br i1 %cmp, label %for.body, label %for.end
+
+for.end:
+  br label %exit
+
+exit:
+  ret void
+}