[SCEV] Add nuw/nsw to mul ops in StrengthenNoWrapFlags where safe.

[jlebar]

Summary

Previously we would add them for adds, but not multiplies.

[jlebar]

Fix all remaining tests.

[sanjoy]

lgtm!

[jlebar]

Thank you for the reviews!

Here goes nothing...

llvm/lib/Analysis/ScalarEvolution.cpp

StrengthenNoWrapFlags(ScalarEvolution *SE, SCEVTypes Type,
   };
 
   if (SignOrUnsignWrap == SCEV::FlagNSW && all_of(Ops, IsKnownNonNegative))
     Flags =
         ScalarEvolution::setFlags(Flags, (SCEV::NoWrapFlags)SignOrUnsignMask);
 
   SignOrUnsignWrap = ScalarEvolution::maskFlags(Flags, SignOrUnsignMask);
 
-  if (SignOrUnsignWrap != SignOrUnsignMask && Type == scAddExpr &&
-      Ops.size() == 2 && isa<SCEVConstant>(Ops[0])) {
+  if (SignOrUnsignWrap != SignOrUnsignMask &&
+      (Type == scAddExpr || Type == scMulExpr) && Ops.size() == 2 &&
+      isa<SCEVConstant>(Ops[0])) {
 
-    // (A + C) --> (A + C)<nsw> if the addition does not sign overflow
-    // (A + C) --> (A + C)<nuw> if the addition does not unsign overflow
+    auto Opcode = [&] {
+      switch (Type) {
+      case scAddExpr:
+        return Instruction::Add;
+      case scMulExpr:
+        return Instruction::Mul;
+      default:
+        llvm_unreachable("Unexpected SCEV op.");
+      }
+    }();
 
     const APInt &C = cast<SCEVConstant>(Ops[0])->getAPInt();
+
+    // (A <opcode> C) --> (A <opcode> C)<nsw> if the op doesn't sign overflow.
     if (!(SignOrUnsignWrap & SCEV::FlagNSW)) {
       auto NSWRegion = ConstantRange::makeGuaranteedNoWrapRegion(
-          Instruction::Add, C, OBO::NoSignedWrap);
+          Opcode, C, OBO::NoSignedWrap);
       if (NSWRegion.contains(SE->getSignedRange(Ops[1])))
         Flags = ScalarEvolution::setFlags(Flags, SCEV::FlagNSW);
     }
+
+    // (A <opcode> C) --> (A <opcode> C)<nuw> if the op doesn't unsign overflow.
     if (!(SignOrUnsignWrap & SCEV::FlagNUW)) {
       auto NUWRegion = ConstantRange::makeGuaranteedNoWrapRegion(
           Instruction::Add, C, OBO::NoUnsignedWrap);
       if (NUWRegion.contains(SE->getUnsignedRange(Ops[1])))
         Flags = ScalarEvolution::setFlags(Flags, SCEV::FlagNUW);
     }
   }
 

llvm/test/Analysis/Delinearization/a.ll

 ; RUN: opt < %s -analyze -delinearize | FileCheck %s
 ;
 ; void foo(long n, long m, long o, int A[n][m][o]) {
 ;   for (long i = 0; i < n; i++)
 ;     for (long j = 0; j < m; j++)
 ;       for (long k = 0; k < o; k++)
 ;         A[2*i+3][3*j-4][5*k+7] = 1;
 ; }
 
 ; AddRec: {{{(28 + (4 * (-4 + (3 * %m)) * %o) + %A),+,(8 * %m * %o)}<%for.i>,+,(12 * %o)}<%for.j>,+,20}<%for.k>
 ; CHECK: Base offset: %A
 ; CHECK: ArrayDecl[UnknownSize][%m][%o] with elements of 4 bytes.
-; CHECK: ArrayRef[{3,+,2}<%for.i>][{-4,+,3}<nw><%for.j>][{7,+,5}<nw><%for.k>]
+; CHECK: ArrayRef[{3,+,2}<nw><%for.i>][{-4,+,3}<nw><%for.j>][{7,+,5}<nw><%for.k>]
 
 define void @foo(i64 %n, i64 %m, i64 %o, i32* nocapture %A) #0 {
 entry:
   %cmp32 = icmp sgt i64 %n, 0
   br i1 %cmp32, label %for.cond1.preheader.lr.ph, label %for.end17
 
 for.cond1.preheader.lr.ph:                        ; preds = %entry
   %cmp230 = icmp sgt i64 %m, 0

llvm/test/Analysis/ScalarEvolution/sext-mul.ll

 ; RUN: opt < %s -analyze -scalar-evolution | FileCheck %s
 
 ; CHECK: %tmp9 = shl i64 %tmp8, 33
-; CHECK-NEXT: --> {{.*}} Exits: (-8589934592 + (8589934592 * (zext i32 %arg2 to i64)))
+; CHECK-NEXT: --> {{.*}} Exits: (-8589934592 + (8589934592 * (zext i32 %arg2 to i64))<nuw>)
 ; CHECK: %tmp10 = ashr exact i64 %tmp9, 32
 ; CHECK-NEXT: --> {{.*}} Exits: (sext i32 (-2 + (2 * %arg2)) to i64)
 ; CHECK: %tmp11 = getelementptr inbounds i32, i32* %arg, i64 %tmp10
 ; CHECK-NEXT: --> {{.*}} Exits: ((4 * (sext i32 (-2 + (2 * %arg2)) to i64)) + %arg)
 ; CHECK:  %tmp14 = or i64 %tmp10, 1
 ; CHECK-NEXT: --> {{.*}} Exits: (1 + (sext i32 (-2 + (2 * %arg2)) to i64))<nsw>
 ; CHECK: %tmp15 = getelementptr inbounds i32, i32* %arg, i64 %tmp14
 ; CHECK-NEXT: --> {{.*}} Exits: (4 + (4 * (sext i32 (-2 + (2 * %arg2)) to i64)) + %arg)
bb7:                                              ; preds = %bb7, %bb3
   %tmp17 = mul nsw i32 %tmp16, %arg1
   store i32 %tmp17, i32* %tmp15, align 4
   %tmp18 = add nuw nsw i64 %tmp8, 1
   %tmp19 = icmp eq i64 %tmp18, %tmp4
   br i1 %tmp19, label %bb5, label %bb7
 }
 
 ; CHECK: %t10 = ashr exact i128 %t9, 1
-; CHECK-NEXT: --> {{.*}} Exits: (sext i127 (-633825300114114700748351602688 + (633825300114114700748351602688 * (zext i32 %arg5 to i127))) to i128)
+; CHECK-NEXT: --> {{.*}} Exits: (sext i127 (-633825300114114700748351602688 + (633825300114114700748351602688 * (zext i32 %arg5 to i127))<nuw>) to i128)
 ; CHECK: %t14 = or i128 %t10, 1
-; CHECK-NEXT: --> {{.*}} Exits: (1 + (sext i127 (-633825300114114700748351602688 + (633825300114114700748351602688 * (zext i32 %arg5 to i127))) to i128))<nsw>
+; CHECK-NEXT: --> {{.*}} Exits: (1 + (sext i127 (-633825300114114700748351602688 + (633825300114114700748351602688 * (zext i32 %arg5 to i127))<nuw>) to i128))<nsw>
 ; CHECK: Loop %bb7: backedge-taken count is (-1 + (zext i32 %arg5 to i128))<nsw>
 ; CHECK-NEXT: Loop %bb7: max backedge-taken count is -1
 ; CHECK-NEXT: Loop %bb7: Predicated backedge-taken count is (-1 + (zext i32 %arg5 to i128))<nsw>
 
 define void @goo(i32* nocapture %arg3, i32 %arg4, i32 %arg5) {
 bb:
   %t = icmp sgt i32 %arg5, 0
   br i1 %t, label %bb3, label %bb6

llvm/test/Transforms/IndVarSimplify/2011-09-10-widen-nsw.ll

for.body11:                                       ; preds = %entry
   br label %for.body153
 
 for.body153:                                      ; preds = %for.body153, %for.body11
   br i1 undef, label %for.body170, label %for.body153
 
 ; CHECK: add nuw nsw i64 %indvars.iv, 1
 ; CHECK: sub nsw i64 %indvars.iv, 2
 ; CHECK: sub nsw i64 4, %indvars.iv
-; CHECK: mul nsw i64 %indvars.iv, 8
+; CHECK: mul nuw nsw i64 %indvars.iv, 8
 for.body170:                                      ; preds = %for.body170, %for.body153
   %i2.19 = phi i32 [ %add249, %for.body170 ], [ 0, %for.body153 ]
 
   %add = add nsw i32 %i2.19, 1
   %add.idxprom = sext i32 %add to i64
 
   %sub = sub nsw i32 %i2.19, 2
   %sub.idxprom = sext i32 %sub to i64