[SCEV] Try harder to preserve NSW information for sext(sub) expressions

[aemerson]

Summary

This helps to preserve NSW flags of sext(sub) expressions by trying to push the sign extension onto the operands of a sext instruction.

[aemerson]

Ping.

[sanjoy]

Sorry for the delay, this one somehow slipped through the cracks.

lib/Analysis/ScalarEvolution.cpp

5966

Can you please use MatchBinaryOp here? That'll catch more cases than a literal sub instruction.

[aemerson]

Updated to use MatchBinaryOp.

[sanjoy]

lib/Analysis/ScalarEvolution.cpp

5975

Why not use BO->IsNSW?

test/Analysis/ScalarEvolution/flags-from-poison.ll

612

Please add a test case that uses Intrinsic::ssub_with_overflow.

[aemerson]

test/Analysis/ScalarEvolution/flags-from-poison.ll

612

I see the code in MatchBinaryOperator doesn't seem to check for guard edges of the ssub_with_overflow intrinsic's overflow return value like it does for sadd_with_overflow. Is there a reason why it doesn't do that, if not I'll go and add support for it.

Without that modification it looks like intrinsic test case won't work.

[sanjoy]

test/Analysis/ScalarEvolution/flags-from-poison.ll

612

Is there a reason why it doesn't do that, if not I'll go and add support for it.

I don't think there was any reason beyond that that code would not be exercised (but this patch changes that); please go ahead with the implementation.

[aemerson]

Added support for handling the nsw/nuw ssub.with.overflow intrinsic.

[sanjoy]

lgtm!

lib/Analysis/ScalarEvolution.cpp

if (auto *F = CI->getCalledFunction())
                           /* IsNUW = */ false);
         else
           return BinaryOp(Instruction::Add, CI->getArgOperand(0),
                           CI->getArgOperand(1), /* IsNSW = */ false,
                           /* IsNUW*/ true);
       }
 
       case Intrinsic::ssub_with_overflow:
-      case Intrinsic::usub_with_overflow:
+      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),
                         CI->getArgOperand(1));
       default:
         break;
       }
   }
const SCEV *ScalarEvolution::createSCEV(Value *V) {
   switch (U->getOpcode()) {
   case Instruction::Trunc:
     return getTruncateExpr(getSCEV(U->getOperand(0)), U->getType());
 
   case Instruction::ZExt:
     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:
     // BitCasts are no-op casts so we just eliminate the cast.
     if (isSCEVable(U->getType()) && isSCEVable(U->getOperand(0)->getType()))
       return getSCEV(U->getOperand(0));
     break;
 

test/Analysis/ScalarEvolution/flags-from-poison.ll

 loop:
   %i = phi i32 [ %nexti, %loop ], [ 0, %entry ]
 
 ; CHECK: %index32 =
 ; CHECK: --> {(-1 * %sub),+,1}<nsw>
   %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
   %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
 }
 
+; 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
 ; is not triggered, hence the strange basic block ordering.
 define void @test-sub-with-add(float* %input, i32 %offset, i32 %numIterations) {
 ; CHECK-LABEL: @test-sub-with-add
 entry:
   br label %loop

test/Transforms/LoopStrengthReduce/sext-ind-var.ll

 
 ; Test this pattern:
 ;
 ;   for (int i = 0; i < numIterations; ++i)
 ;     sum += ptr[i - offset];
 ;
 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
 ; CHECK-NOT: sext
 
 entry:
   br label %loop