[SCEV] Preserve NSW information for sext(subtract).

Pushes the sext onto the operands of a Sub if NSW is present.
Also adds support for propagating the nowrap flags of the
llvm.ssub.with.overflow intrinsic during analysis.

Differential Revision: https://reviews.llvm.org/D35256

llvm-svn: 310117

3 files changed, 65 insertions, 6 deletions

diff --git a/llvm/lib/Analysis/ScalarEvolution.cpp b/llvm/lib/Analysis/ScalarEvolution.cpp
index d3baf532045..0d1c6d16f41 100644
--- a/llvm/lib/Analysis/ScalarEvolution.cpp
+++ b/llvm/lib/Analysis/ScalarEvolution.cpp
@@ -4168,10 +4168,21 @@ static Optional<BinaryOp> MatchBinaryOp(Value *V, DominatorTree &DT) {
       }
 
       case Intrinsic::ssub_with_overflow:
-      case Intrinsic::usub_with_overflow:
-        return BinaryOp(Instruction::Sub, CI->getArgOperand(0),
-                        CI->getArgOperand(1));
+      case Intrinsic::usub_with_overflow: {
+        if (!isOverflowIntrinsicNoWrap(cast<IntrinsicInst>(CI), DT))
+          return BinaryOp(Instruction::Sub, CI->getArgOperand(0),
+                          CI->getArgOperand(1));
 
+        // The same reasoning as sadd/uadd above.
+        if (F->getIntrinsicID() == Intrinsic::ssub_with_overflow)
+          return BinaryOp(Instruction::Sub, CI->getArgOperand(0),
+                          CI->getArgOperand(1), /* IsNSW = */ true,
+                          /* IsNUW = */ false);
+        else
+          return BinaryOp(Instruction::Sub, CI->getArgOperand(0),
+                          CI->getArgOperand(1), /* IsNSW = */ false,
+                          /* IsNUW = */ true);
+      }
       case Intrinsic::smul_with_overflow:
       case Intrinsic::umul_with_overflow:
         return BinaryOp(Instruction::Mul, CI->getArgOperand(0),
@@ -5952,6 +5963,21 @@ const SCEV *ScalarEvolution::createSCEV(Value *V) {
     return getZeroExtendExpr(getSCEV(U->getOperand(0)), U->getType());
 
   case Instruction::SExt:
+    if (auto BO = MatchBinaryOp(U->getOperand(0), DT)) {
+      // The NSW flag of a subtract does not always survive the conversion to
+      // A + (-1)*B.  By pushing sign extension onto its operands we are much
+      // more likely to preserve NSW and allow later AddRec optimisations.
+      //
+      // NOTE: This is effectively duplicating this logic from getSignExtend:
+      //   sext((A + B + ...)<nsw>) --> (sext(A) + sext(B) + ...)<nsw>
+      // but by that point the NSW information has potentially been lost.
+      if (BO->Opcode == Instruction::Sub && BO->IsNSW) {
+        Type *Ty = U->getType();
+        auto *V1 = getSignExtendExpr(getSCEV(BO->LHS), Ty);
+        auto *V2 = getSignExtendExpr(getSCEV(BO->RHS), Ty);
+        return getMinusSCEV(V1, V2, SCEV::FlagNSW);
+      }
+    }
     return getSignExtendExpr(getSCEV(U->getOperand(0)), U->getType());
 
   case Instruction::BitCast:
diff --git a/llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll b/llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll
index 15c679a5f10..8733a857cba 100644
--- a/llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll
+++ b/llvm/test/Analysis/ScalarEvolution/flags-from-poison.ll
@@ -609,7 +609,7 @@ loop:
   %index32 = sub nsw i32 %i, %sub
 
 ; CHECK: %index64 =
-; CHECK: --> {(sext i32 (-1 * %sub) to i64),+,1}<nsw>
+; CHECK: --> {(-1 * (sext i32 %sub to i64))<nsw>,+,1}<nsw
   %index64 = sext i32 %index32 to i64
 
   %ptr = getelementptr inbounds float, float* %input, i64 %index64
@@ -621,6 +621,40 @@ exit:
   ret void
 }
 
+; Example checking that a sext is pushed onto a sub's operands if the sub is an
+; overflow intrinsic.
+define void @test-sext-sub(float* %input, i32 %sub, i32 %numIterations) {
+; CHECK-LABEL: @test-sext-sub
+entry:
+  br label %loop
+loop:
+  %i = phi i32 [ %nexti, %cont ], [ 0, %entry ]
+
+; CHECK: %val = extractvalue { i32, i1 } %ssub, 0
+; CHECK: --> {(-1 * %sub),+,1}<nw>
+  %ssub = tail call { i32, i1 } @llvm.ssub.with.overflow.i32(i32 %i, i32 %sub)
+  %val = extractvalue { i32, i1 } %ssub, 0
+  %ovfl = extractvalue { i32, i1 } %ssub, 1
+  br i1 %ovfl, label %trap, label %cont
+
+trap:
+  tail call void @llvm.trap()
+  unreachable
+
+cont:
+; CHECK: %index64 =
+; CHECK: --> {(-1 * (sext i32 %sub to i64))<nsw>,+,1}<nsw
+  %index64 = sext i32 %val to i64
+
+  %ptr = getelementptr inbounds float, float* %input, i64 %index64
+  %nexti = add nsw i32 %i, 1
+  %f = load float, float* %ptr, align 4
+  %exitcond = icmp eq i32 %nexti, %numIterations
+  br i1 %exitcond, label %exit, label %loop
+exit:
+  ret void
+}
+
 ; Two adds with a sub in the middle and the sub should have nsw. There is
 ; a special case for sequential adds/subs and this test covers that. We have to
 ; put the final add first in the program since otherwise the special case
diff --git a/llvm/test/Transforms/LoopStrengthReduce/sext-ind-var.ll b/llvm/test/Transforms/LoopStrengthReduce/sext-ind-var.ll
index 3cf8f536fa7..21e72b29831 100644
--- a/llvm/test/Transforms/LoopStrengthReduce/sext-ind-var.ll
+++ b/llvm/test/Transforms/LoopStrengthReduce/sext-ind-var.ll
@@ -47,8 +47,7 @@ exit:
 ;
 define float @testsub(float* %input, i32 %offset, i32 %numIterations) {
 ; CHECK-LABEL: @testsub
-; CHECK: sub i32 0, %offset
-; CHECK: sext i32
+; CHECK: sext i32 %offset to i64
 ; CHECK: loop:
 ; CHECK-DAG: phi float*
 ; CHECK-DAG: phi i32