[IndVarSimplify] Do not expand unsafe expressions in IndVarSimplify

[mkazantsev]

Summary

When using SCEVExpander, passes are obliged to make sure that what they request
to expand is safe in terms of expander. We are going to enforce this adding an assert
that it is safe to expand what we request.

This patch makes sure that we only expand what is safe to expand in IndVarSimplify,
in particular that we don't expand division by potential zero.

[sanjoy]

lib/Transforms/Scalar/IndVarSimplify.cpp

1589

Hm, I'm not sold on this. If AR is expandable (which we know since we got it from a PHI in the IR, right?) sext(AR) and zext(AR) should also be expandable.

test/Transforms/IndVarSimplify/2017-10-24_unsafe_division.ll

39

I'm not sure if I understand this test case - if %iv.trunc is 0 then doesn't the function already divide by 0 (and thus have UB, making any transformation correct by definition)?

lib/Transforms/Scalar/IndVarSimplify.cpp

PHINode *WidenIV::createWideIV(SCEVExpander &Rewriter) {
   if (UsePostIncrementRanges)
     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();
+
+  // Make sure it is safe to expand at this point. For example, that we don't
+  // have division by value which is potentially zero.
+  if (!isSafeToExpand(AddRec, *SE))
+    return nullptr;
+
   WidePhi = cast<PHINode>(Rewriter.expandCodeFor(AddRec, WideType, InsertPt));
 
   // Remembering the WideIV increment generated by SCEVExpander allows
   // widenIVUse to reuse it when widening the narrow IV's increment. We don't
   // employ a general reuse mechanism because the call above is the only call to
   // SCEVExpander. Henceforth, we produce 1-to-1 narrow to wide uses.
   if (BasicBlock *LatchBlock = L->getLoopLatch()) {
     WideInc =
if (IndVar->getType()->isPointerTy() && !IVCount->getType()->isPointerTy()) {
     IRBuilder<> Builder(BI);
     assert(SE->isLoopInvariant(IVLimit, L) &&
            "Computed iteration count is not loop invariant!");
     // Ensure that we generate the same type as IndVar, or a smaller integer
     // type. In the presence of null pointer values, we have an integer type
     // SCEV expression (IVInit) for a pointer type IV value (IndVar).
     Type *LimitTy = IVCount->getType()->isPointerTy() ?
       IndVar->getType() : IVCount->getType();
+
+    // We may end up with unsafe expanding, do not deal with it.
+    if (!isSafeToExpand(IVLimit, *SE))
+      return nullptr;
+
     return Rewriter.expandCodeFor(IVLimit, LimitTy, BI);
   }
 }
 
 /// This method rewrites the exit condition of the loop to be a canonical !=
 /// comparison against the incremented loop induction variable.  This pass is
 /// able to rewrite the exit tests of any loop where the SCEV analysis can
 /// determine a loop-invariant trip count of the loop, which is actually a much
IVCount = SE->getAddExpr(BackedgeTakenCount,
                              SE->getOne(BackedgeTakenCount->getType()));
     // The BackedgeTaken expression contains the number of times that the
     // backedge branches to the loop header.  This is one less than the
     // number of times the loop executes, so use the incremented indvar.
     CmpIndVar = IndVar->getIncomingValueForBlock(L->getExitingBlock());
   }
 
   Value *ExitCnt = genLoopLimit(IndVar, IVCount, L, Rewriter, SE);
+  // We may end up with nullptr if we failed to expand something.
+  if (!ExitCnt)
+    return nullptr;
+
   assert(ExitCnt->getType()->isPointerTy() ==
              IndVar->getType()->isPointerTy() &&
          "genLoopLimit missed a cast");
 
   // Insert a new icmp_ne or icmp_eq instruction before the branch.
   BranchInst *BI = cast<BranchInst>(L->getExitingBlock()->getTerminator());
   ICmpInst::Predicate P;
   if (L->contains(BI->getSuccessor(0)))
if (!DisableLFTR && canExpandBackedgeTakenCount(L, SE, Rewriter) &&
       needsLFTR(L, DT)) {
     PHINode *IndVar = FindLoopCounter(L, BackedgeTakenCount, SE, DT);
     if (IndVar) {
       // Check preconditions for proper SCEVExpander operation. SCEV does not
       // express SCEVExpander's dependencies, such as LoopSimplify. Instead any
       // pass that uses the SCEVExpander must do it. This does not work well for
       // loop passes because SCEVExpander makes assumptions about all loops,
       // while LoopPassManager only forces the current loop to be simplified.
-      //
-      // FIXME: SCEV expansion has no way to bail out, so the caller must
-      // explicitly check any assumptions made by SCEV. Brittle.
       const SCEVAddRecExpr *AR = dyn_cast<SCEVAddRecExpr>(BackedgeTakenCount);
       if (!AR || AR->getLoop()->getLoopPreheader())
         (void)linearFunctionTestReplace(L, BackedgeTakenCount, IndVar,
                                         Rewriter);
     }
   }
   // Clear the rewriter cache, because values that are in the rewriter's cache
   // can be deleted in the loop below, causing the AssertingVH in the cache to

test/Transforms/IndVarSimplify/2017-10-24_unsafe_division.ll

+; RUN: opt -S -indvars < %s | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128-ni:1"
+target triple = "x86_64-unknown-linux-gnu"
+
+; In this example, we may end up expanding %unsafe. The value of %iv.trunc may
+; potentially be zero even despite %iv is always positive, and in %unsafe we
+; divide by it. In terms of SCEV expander, expanding such a value is incorrect.
+
+define void @test_01() {
+
+; CHECK-LABEL: test_01
+; CHECK:       entry:
+; CHECK-NEXT:    br label %outer_loop
+; CHECK:       outer_loop:
+; CHECK-NEXT:    %iv = phi i64 [ 5, %entry ], [ %iv.next, %outer_backedge ]
+; CHECK-NEXT:    %iv.trunc = trunc i64 %iv to i32
+; CHECK-NEXT:    %unsafe = udiv i32 224, %iv.trunc
+; CHECK-NEXT:    br label %inner_loop
+; CHECK:       inner_loop:
+; CHECK-NEXT:    %unsafe_phi = phi i32 [ %unsafe, %outer_loop ], [ %unsafe_phi.next, %inner_loop ]
+; CHECK-NEXT:    %iv.inner = phi i32 [ 0, %outer_loop ], [ %iv.inner.next, %inner_loop ]
+; CHECK-NEXT:    %unsafe_phi.next = add nuw nsw i32 %unsafe_phi, 1
+; CHECK-NEXT:    %expand.me = zext i32 %unsafe_phi.next to i64
+; CHECK-NEXT:    %iv.inner.next = add nuw nsw i32 %iv.inner, 1
+; CHECK-NEXT:    %cond.inner = icmp ugt i32 %iv.inner, 1643
+; CHECK-NEXT:    br i1 %cond.inner, label %outer_backedge, label %inner_loop
+; CHECK:       outer_backedge:
+; CHECK-NEXT:    %iv.next = add nuw nsw i64 %iv, 1
+; CHECK-NEXT:    br label %outer_loop
+
+entry:
+  br label %outer_loop
+
+outer_loop:
+  %iv = phi i64 [ 5, %entry ], [ %iv.next, %outer_backedge ]
+  %iv.next = add nuw nsw i64 %iv, 1
+  %iv.trunc = trunc i64 %iv to i32
+  %unsafe = udiv i32 224, %iv.trunc
+  br label %inner_loop
+
+inner_loop:
+  %unsafe_phi = phi i32 [ %unsafe, %outer_loop ], [ %unsafe_phi.next, %inner_loop ]
+  %iv.inner = phi i32 [ 0, %outer_loop ], [ %iv.inner.next, %inner_loop ]
+  %unsafe_phi.next = add nuw nsw i32 %unsafe_phi, 1
+  %expand.me = zext i32 %unsafe_phi.next to i64
+  %iv.inner.next = add nuw nsw i32 %iv.inner, 1
+  %cond.inner = icmp ugt i32 %iv.inner, 1643
+  br i1 %cond.inner, label %outer_backedge, label %inner_loop
+
+outer_backedge:
+  br label %outer_loop
+}

test/Transforms/IndVarSimplify/udiv.ll

 
 for.end.loopexit:                                 ; preds = %for.body
   br label %for.end
 
 for.end:                                          ; preds = %for.end.loopexit, %entry
   ret void
 }
 
+
+; Same as foo_02, but we divide by non-constant value which can be zero. But we
+; do not need to generate any new instructions for that. SCEV expander will
+; prohibit expanding this.
+
+; CHECK-LABEL: @foo_03(
+; CHECK: for.body.preheader:
+; CHECK: udiv
+
+define void @foo_03(double* %p, i64 %n) nounwind {
+entry:
+  %denom = load i64, i64* undef, align 4, !range !1
+  %div0 = udiv i64 %n, %denom                          ; <i64> [#uses=1]
+  %div1 = add i64 %div0, 1
+  %cmp2 = icmp ult i64 0, %div1                   ; <i1> [#uses=1]
+  br i1 %cmp2, label %for.body.preheader, label %for.end
+
+for.body.preheader:                               ; preds = %entry
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.preheader, %for.body
+  %i.03 = phi i64 [ %inc, %for.body ], [ 0, %for.body.preheader ] ; <i64> [#uses=2]
+  %arrayidx = getelementptr inbounds double, double* %p, i64 %i.03 ; <double*> [#uses=1]
+  store double 0.000000e+00, double* %arrayidx
+  %inc = add i64 %i.03, 1                         ; <i64> [#uses=2]
+  %divx = udiv i64 %n, %denom                           ; <i64> [#uses=1]
+  %div = add i64 %divx, 1
+  %cmp = icmp ult i64 %inc, %div                  ; <i1> [#uses=1]
+  br i1 %cmp, label %for.body, label %for.end.loopexit
+
+for.end.loopexit:                                 ; preds = %for.body
+  br label %for.end
+
+for.end:                                          ; preds = %for.end.loopexit, %entry
+  ret void
+}
+
 !0 = !{i64 1, i64 10}
+!1 = !{i64 0, i64 10}