[InstCombine] Extend sub(sext(add(X, Y)), sext(add(X, Z))) -> sub(Y, Z) fold to adds that mutually wrap

[dmgreen]

Summary

This combine, added in D157598, folds the sub of two extended adds if the adds are no wrap, as we know the extends are simple. This isn't just valid for no-wrap though - it is also valid if the two adds wrap at the same time. i.e given the following code:

;  %a = srem i32 %X, 3
;  switch i32 %a, label %.loopexit477 [
;    i32 1, label %preheader
;    i32 2, label %.loopexit477.loopexit667
;  ]
;
;preheader:
;  %b = add i32 %X, 1
;  %c = sext i32 %b to i64
;  %d = add i32 %X, 2
;  %e = sext i32 %d to i64
;  %f = sub nsw i64 %e, %c

The sub is provably '1' because due to knowing from the dominating condition that X % 3 == 1, and that for i32 the two adds will either both not wrap (for most values of X), or both wrap if X==2147483647. https://alive2.llvm.org/ce/z/JvnY-o

So this patch adds a isKnownMutuallyWrapping function to detect when two adds will wrap at the same time, and from the mod case detects if the two adds will wrap using (C3 - (Max % C4) + C1) / C4 == (C3 - (Max % C4) + C2) / C4, where C4 is the mod denom and C3 is the remainder.
For Signed: https://alive2.llvm.org/ce/z/N93Njy
And Unsigned: https://alive2.llvm.org/ce/z/A9gzbK (note this is more generally true for the conditions here, but this patch only tries to prove it for the mutually wrapping case).

llvm/include/llvm/Analysis/ValueTracking.h

 /// i.e.  x == sext(trunc(x to MaxSignificantBits) to bitwidth(x)).
 /// Similar to the APInt::getSignificantBits function.
 unsigned ComputeMaxSignificantBits(const Value *Op, const DataLayout &DL,
                                    unsigned Depth = 0,
                                    AssumptionCache *AC = nullptr,
                                    const Instruction *CxtI = nullptr,
                                    const DominatorTree *DT = nullptr);
 
+/// Return true if V1 and V2 are known to mutually wrap - that is they wrap at
+/// the same time. For example i8's (a+1) and (a+2) would mutually signed wrap
+/// if the a%3==1, as they either both do not wrap or both wrap at the same
+/// time. NSWWrap specifies whether it checks for signed (true) or unsigned
+/// (false) wrapping.
+bool isKnownMutuallyWrapping(const Value *V1, const Value *V2, bool NSWWrap,
+                             const Instruction *ContextI = nullptr);
+
 /// Map a call instruction to an intrinsic ID.  Libcalls which have equivalent
 /// intrinsics are treated as-if they were intrinsics.
 Intrinsic::ID getIntrinsicForCallSite(const CallBase &CB,
                                       const TargetLibraryInfo *TLI);
 
 /// Returns a pair of values, which if passed to llvm.is.fpclass, returns the
 /// same result as an fcmp with the given operands.
 ///

llvm/lib/Analysis/ValueTracking.cpp

Intrinsic::ID llvm::getIntrinsicForCallSite(const CallBase &CB,
   case LibFunc_sqrtf:
   case LibFunc_sqrtl:
     return Intrinsic::sqrt;
   }
 
   return Intrinsic::not_intrinsic;
 }
 
+static std::pair<Value *, bool>
+getDomPredecessorCondition(const Instruction *ContextI);
+
+bool llvm::isKnownMutuallyWrapping(const Value *V1, const Value *V2,
+                                   bool NSWWrap, const Instruction *ContextI) {
+  if (!match(V1, m_Add(m_Value(), m_Value())) ||
+      !match(V2, m_Add(m_Value(), m_Value())))
+    return false;
+
+  // Check the easy case where both Adds are NSW/NUW.
+  if (NSWWrap && cast<OverflowingBinaryOperator>(V1)->hasNoSignedWrap() &&
+      cast<OverflowingBinaryOperator>(V2)->hasNoSignedWrap())
+    return true;
+  if (!NSWWrap && cast<OverflowingBinaryOperator>(V1)->hasNoUnsignedWrap() &&
+      cast<OverflowingBinaryOperator>(V2)->hasNoUnsignedWrap())
+    return true;
+
+  // Check for X+C1 and X+C2
+  Value *X;
+  const APInt *C1, *C2, *Rem, *Denom;
+  if (ContextI && match(V1, m_Add(m_Value(X), m_APInt(C1))) &&
+      match(V2, m_Add(m_Specific(X), m_APInt(C2))) && C1->isNonNegative() &&
+      C2->isNonNegative()) {
+    auto PredCond = getDomPredecessorCondition(ContextI);
+    if (PredCond.first) {
+      ICmpInst::Predicate Pred;
+      // X+C1 and X+C2 will mutually wrap at the same time if we know the range
+      // of X is a given by a modulo and that the two adds will always fall on
+      // the same side of wrap:
+      // Limit = signed/unsigned max + 1
+      // (Rem - (Limit % Denom) + C1) / Denom = (Rem - (Limit % Denom) + C2) / Denom
+      if (match(PredCond.first,
+                m_ICmp(Pred, m_SRem(m_Specific(X), m_APInt(Denom)),
+                       m_APInt(Rem))) &&
+          Pred == ICmpInst::ICMP_EQ && !Denom->isNegative()) {
+        unsigned BW = C1->getBitWidth();
+        APInt Limit =
+            NSWWrap ? APInt::getSignedMaxValue(BW) : APInt::getAllOnes(BW);
+        Limit = (Limit.zext(BW + 1) + 1).urem(Denom->zext(BW + 1)).trunc(BW);
+        APInt Start = *Rem - Limit;
+        return (Start + *C1).udiv(*Denom) == (Start + *C2).udiv(*Denom);
+      }
+    }
+  }
+
+  return false;
+}
+
 /// Deprecated, use computeKnownFPClass instead.
 ///
 /// If \p SignBitOnly is true, test for a known 0 sign bit rather than a
 /// standard ordered compare. e.g. make -0.0 olt 0.0 be true because of the sign
 /// bit despite comparing equal.
 static bool cannotBeOrderedLessThanZeroImpl(const Value *V,
                                             const DataLayout &DL,
                                             const TargetLibraryInfo *TLI,

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

if (I.hasNoUnsignedWrap() || I.hasNoSignedWrap()) {
       Value *Sub =
           Builder.CreateSub(X, Y, "sub", /*HasNUW=*/false, /*HasNSW=*/true);
       Value *Call =
           Builder.CreateBinaryIntrinsic(Intrinsic::abs, Sub, Builder.getTrue());
       return replaceInstUsesWith(I, Call);
     }
   }
 
-  // sext(X + Y) - sext(X + Z) -> sext(Y) - sext(Z)
-  if ((match(Op0, m_SExt(m_NSWAdd(m_Value(X), m_Value(Y)))) &&
+  // sext(X + Y) - sext(X + Z) -> sext(Y) - sext(Z) if X+Y and X+Z are known to
+  // both wrap at the same time.
+  if (((match(Op0, m_SExt(m_Add(m_Value(X), m_Value(Y)))) &&
         match(Op1, m_SExt(m_c_Add(m_Specific(X), m_Value(Z))))) ||
-      (match(Op0, m_SExt(m_NSWAdd(m_Value(Y), m_Value(X)))) &&
-       match(Op1, m_SExt(m_c_Add(m_Specific(X), m_Value(Z)))))) {
-    Value *Add1 = cast<Instruction>(Op1)->getOperand(0);
-    if ((Op0->hasOneUse() || Op1->hasOneUse() ||
-         (isa<Constant>(Y) && isa<Constant>(Z))) &&
-        cast<OverflowingBinaryOperator>(Add1)->hasNoSignedWrap()) {
-      Value *S1 = Builder.CreateSExt(Y, Ty);
-      Value *S2 = Builder.CreateSExt(Z, Ty);
-      return BinaryOperator::CreateSub(S1, S2);
-    }
-  }
-  if ((match(Op0, m_ZExt(m_NUWAdd(m_Value(X), m_Value(Y)))) &&
+       (match(Op0, m_SExt(m_Add(m_Value(Y), m_Value(X)))) &&
+        match(Op1, m_SExt(m_c_Add(m_Specific(X), m_Value(Z)))))) &&
+      (Op0->hasOneUse() || Op1->hasOneUse() ||
+      isKnownMutuallyWrapping(cast<Instruction>(Op0)->getOperand(0),
+                              cast<Instruction>(Op1)->getOperand(0), true,
+                              &I)) {
+  if (((match(Op0, m_ZExt(m_Add(m_Value(X), m_Value(Y)))) &&
         match(Op1, m_ZExt(m_c_Add(m_Specific(X), m_Value(Z))))) ||
-      (match(Op0, m_ZExt(m_NUWAdd(m_Value(Y), m_Value(X)))) &&
-       match(Op1, m_ZExt(m_c_Add(m_Specific(X), m_Value(Z)))))) {
-    Value *Add1 = cast<Instruction>(Op1)->getOperand(0);
-    if ((Op0->hasOneUse() || Op1->hasOneUse() ||
-         (isa<Constant>(Y) && isa<Constant>(Z))) &&
-        cast<OverflowingBinaryOperator>(Add1)->hasNoUnsignedWrap()) {
-      Value *S1 = Builder.CreateZExt(Y, Ty);
-      Value *S2 = Builder.CreateZExt(Z, Ty);
-      return BinaryOperator::CreateSub(S1, S2);
-    }
-  }
+       (match(Op0, m_ZExt(m_Add(m_Value(Y), m_Value(X)))) &&
+        match(Op1, m_ZExt(m_c_Add(m_Specific(X), m_Value(Z)))))) &&
+      (Op0->hasOneUse() || Op1->hasOneUse() ||
+      isKnownMutuallyWrapping(cast<Instruction>(Op0)->getOperand(0),
+                              cast<Instruction>(Op1)->getOperand(0), false,
+                              &I)) {
 
   if (Instruction *Res = foldBinOpOfSelectAndCastOfSelectCondition(I))
     return Res;
 
   return TryToNarrowDeduceFlags();
 }
 
 /// This eliminates floating-point negation in either 'fneg(X)' or

llvm/test/Transforms/InstCombine/sub-ext-add.ll

 
 ; Valid
 define i32 @implied_3mod1_sext_1_2(i8 %a) {
 ; CHECK-LABEL: @implied_3mod1_sext_1_2(
 ; CHECK-NEXT:    [[B:%.*]] = srem i8 [[A:%.*]], 3
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i8 [[B]], 1
 ; CHECK-NEXT:    br i1 [[C]], label [[BB:%.*]], label [[END:%.*]]
 ; CHECK:       bb:
-; CHECK-NEXT:    [[A1:%.*]] = add i8 [[A]], 1
-; CHECK-NEXT:    [[Z1:%.*]] = sext i8 [[A1]] to i32
-; CHECK-NEXT:    [[A2:%.*]] = add i8 [[A]], 2
-; CHECK-NEXT:    [[Z2:%.*]] = sext i8 [[A2]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[Z2]], [[Z1]]
-; CHECK-NEXT:    ret i32 [[S]]
+; CHECK-NEXT:    ret i32 1
 ; CHECK:       end:
 ; CHECK-NEXT:    ret i32 0
 ;
   %b = srem i8 %a, 3
   %c = icmp eq i8 %b, 1
   br i1 %c, label %bb, label %end
 
 bb:

 
 ; Valid
 define i32 @implied_3mod1_zext_1_2(i8 %a) {
 ; CHECK-LABEL: @implied_3mod1_zext_1_2(
 ; CHECK-NEXT:    [[B:%.*]] = srem i8 [[A:%.*]], 3
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i8 [[B]], 1
 ; CHECK-NEXT:    br i1 [[C]], label [[BB:%.*]], label [[END:%.*]]
 ; CHECK:       bb:
-; CHECK-NEXT:    [[A1:%.*]] = add i8 [[A]], 1
-; CHECK-NEXT:    [[Z1:%.*]] = zext i8 [[A1]] to i32
-; CHECK-NEXT:    [[A2:%.*]] = add i8 [[A]], 2
-; CHECK-NEXT:    [[Z2:%.*]] = zext i8 [[A2]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[Z2]], [[Z1]]
-; CHECK-NEXT:    ret i32 [[S]]
+; CHECK-NEXT:    ret i32 1
 ; CHECK:       end:
 ; CHECK-NEXT:    ret i32 0
 ;
   %b = srem i8 %a, 3
   %c = icmp eq i8 %b, 1
   br i1 %c, label %bb, label %end
 
 bb:

 
 ; Valid
 define i32 @implied_3mod0_zext_1_2(i8 %a) {
 ; CHECK-LABEL: @implied_3mod0_zext_1_2(
 ; CHECK-NEXT:    [[B:%.*]] = srem i8 [[A:%.*]], 3
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i8 [[B]], 0
 ; CHECK-NEXT:    br i1 [[C]], label [[BB:%.*]], label [[END:%.*]]
 ; CHECK:       bb:
-; CHECK-NEXT:    [[A1:%.*]] = add i8 [[A]], 1
-; CHECK-NEXT:    [[Z1:%.*]] = zext i8 [[A1]] to i32
-; CHECK-NEXT:    [[A2:%.*]] = add i8 [[A]], 2
-; CHECK-NEXT:    [[Z2:%.*]] = zext i8 [[A2]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[Z2]], [[Z1]]
-; CHECK-NEXT:    ret i32 [[S]]
+; CHECK-NEXT:    ret i32 1
 ; CHECK:       end:
 ; CHECK-NEXT:    ret i32 0
 ;
   %b = srem i8 %a, 3
   %c = icmp eq i8 %b, 0
   br i1 %c, label %bb, label %end
 
 bb:

 
 ; Doesn't handle negative C4
 define i32 @implied_m3mod1_sext_1_2(i8 %a) {
 ; CHECK-LABEL: @implied_m3mod1_sext_1_2(
 ; CHECK-NEXT:    [[B:%.*]] = srem i8 [[A:%.*]], 3
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i8 [[B]], 1
 ; CHECK-NEXT:    br i1 [[C]], label [[BB:%.*]], label [[END:%.*]]
 ; CHECK:       bb:
-; CHECK-NEXT:    [[A1:%.*]] = add i8 [[A]], 1
-; CHECK-NEXT:    [[Z1:%.*]] = sext i8 [[A1]] to i32
-; CHECK-NEXT:    [[A2:%.*]] = add i8 [[A]], 2
-; CHECK-NEXT:    [[Z2:%.*]] = sext i8 [[A2]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[Z2]], [[Z1]]
-; CHECK-NEXT:    ret i32 [[S]]
+; CHECK-NEXT:    ret i32 1
 ; CHECK:       end:
 ; CHECK-NEXT:    ret i32 0
 ;
   %b = srem i8 %a, -3
   %c = icmp eq i8 %b, 1
   br i1 %c, label %bb, label %end
 
 bb: