[IndVarSimplify] Teach calculatePostIncRange to take guards into account

[sanjoy]

lgtm with nits inline

lib/Transforms/Scalar/IndVarSimplify.cpp

890

Add a one-liner here on what this is for.

1610

Can you use make_range here?

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

152

Can you also add a test case where you've not broken the block? That is, the loop has only one basic block?

lib/Transforms/Scalar/IndVarSimplify.cpp

class WidenIV {
   Type *WideType;
 
   // Context
   LoopInfo        *LI;
   Loop            *L;
   ScalarEvolution *SE;
   DominatorTree   *DT;
 
+  bool HasGuards;
+
   // Result
   PHINode *WidePhi;
   Instruction *WideInc;
   const SCEV *WideIncExpr;
   SmallVectorImpl<WeakVH> &DeadInsts;
 
   SmallPtrSet<Instruction *,16> Widened;
   SmallVector<NarrowIVDefUse, 8> NarrowIVUsers;
void updatePostIncRangeInfo(Value *Def, Instruction *UseI, ConstantRange R) {
     } else {
       It->second = R.intersectWith(It->second);
     }
   }
 
 public:
   WidenIV(const WideIVInfo &WI, LoopInfo *LInfo,
           ScalarEvolution *SEv, DominatorTree *DTree,
-          SmallVectorImpl<WeakVH> &DI) :
+          SmallVectorImpl<WeakVH> &DI, bool HasGuards) :
     OrigPhi(WI.NarrowIV),
     WideType(WI.WidestNativeType),
     LI(LInfo),
     L(LI->getLoopFor(OrigPhi->getParent())),
     SE(SEv),
     DT(DTree),
+    HasGuards(HasGuards),
     WidePhi(nullptr),
     WideInc(nullptr),
     WideIncExpr(nullptr),
     DeadInsts(DI) {
     assert(L->getHeader() == OrigPhi->getParent() && "Phi must be an IV");
     ExtendKindMap[OrigPhi] = WI.IsSigned ? SignExtended : ZeroExtended;
   }
 
void WidenIV::calculatePostIncRange(Instruction *NarrowDef,
 
   Value *NarrowDefLHS;
   const APInt *NarrowDefRHS;
   if (!match(NarrowDef, m_NSWAdd(m_Value(NarrowDefLHS),
                                  m_APInt(NarrowDefRHS))) ||
       !NarrowDefRHS->isNonNegative())
     return;
 
-  auto UpdateRangeFromCondition = [&] (Value *Condition) {
+  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);
   };
 
+  auto UpdateRangeFromGuards = [&](Instruction *Ctx) {
+    if (!HasGuards)
+      return;
+
+    for (BasicBlock::reverse_iterator I = Ctx->getIterator().getReverse(),
+                 E = Ctx->getParent()->rend(); I != E; I++) {
+      Value *Cond = nullptr;
+      if (!match(&*I, m_Intrinsic<Intrinsic::experimental_guard>(
+                 m_Value(Cond))))
+        continue;
+      UpdateRangeFromCondition(Cond);
+    }
+  };
+
+  UpdateRangeFromGuards(NarrowUser);
+
   BasicBlock *NarrowUserBB = NarrowUser->getParent();
   for (auto *DTB = (*DT)[NarrowUserBB]->getIDom();
        L->contains(DTB->getBlock());
        DTB = DTB->getIDom()) {
     auto TI = DTB->getBlock()->getTerminator();
+    UpdateRangeFromGuards(TI);
 
     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())))

 ///
 /// Sign/Zero extend elimination is interleaved with IV simplification.
 ///
 void IndVarSimplify::simplifyAndExtend(Loop *L,
                                        SCEVExpander &Rewriter,
                                        LoopInfo *LI) {
   SmallVector<WideIVInfo, 8> WideIVs;
 
+  auto *GuardDecl = L->getBlocks()[0]->getModule()->getFunction(
+          Intrinsic::getName(Intrinsic::experimental_guard));
+  bool HasGuards = GuardDecl && !GuardDecl->use_empty();
+
   SmallVector<PHINode*, 8> LoopPhis;
   for (BasicBlock::iterator I = L->getHeader()->begin(); isa<PHINode>(I); ++I) {
     LoopPhis.push_back(cast<PHINode>(I));
   }
   // Each round of simplification iterates through the SimplifyIVUsers worklist
   // for all current phis, then determines whether any IVs can be
   // widened. Widening adds new phis to LoopPhis, inducing another round of
   // simplification on the wide IVs.
do {
       Changed |= simplifyUsersOfIV(CurrIV, SE, DT, LI, DeadInsts, &Visitor);
 
       if (Visitor.WI.WidestNativeType) {
         WideIVs.push_back(Visitor.WI);
       }
     } while(!LoopPhis.empty());
 
     for (; !WideIVs.empty(); WideIVs.pop_back()) {
-      WidenIV Widener(WideIVs.back(), LI, SE, DT, DeadInsts);
+      WidenIV Widener(WideIVs.back(), LI, SE, DT, DeadInsts, HasGuards);
       if (PHINode *WidePhi = Widener.createWideIV(Rewriter)) {
         Changed = true;
         LoopPhis.push_back(WidePhi);
       }
     }
   }
 }
 

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

for.inc:
   br i1 %cmp, label %for.body, label %for.end
 
 for.end:
   br label %exit
 
 exit:
   ret void
 }
+
+declare void @llvm.experimental.guard(i1, ...)
+
+define void @test_guard_in_the_same_bb(i32* %base, i32 %limit, i32 %start) {
+; CHECK-LABEL: @test_guard_in_the_same_bb(
+; CHECK-NOT: trunc
+; CHECK-NOT: icmp slt i32
+
+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
+  %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:
+  call void(i1, ...) @llvm.experimental.guard(i1 %within_limits) [ "deopt"() ]
+  %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_guard_in_idom(i32* %base, i32 %limit, i32 %start) {
+; CHECK-LABEL: @test_guard_in_idom(
+; CHECK-NOT: trunc
+; CHECK-NOT: icmp slt i32
+
+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
+  call void(i1, ...) @llvm.experimental.guard(i1 %within_limits) [ "deopt"() ]
+  %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
+}