[InstCombine] Transform sub(sext(add(X, Y)), sext(add(X, Z))) to sub(Y, Z) given nsw adds

[dmgreen]

Summary

This adds combines for sext(X +nsw Y) - sext(X +nsw Z) -> sext(Y) - sext(Z) and zext(X +nuw Y) - zext(X +nuw Z) -> zext(Y) - zext(Z). This is especially useful when Y and Z are constants.

See https://alive2.llvm.org/ce/z/XEazaS and https://alive2.llvm.org/ce/z/kHGPtU.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

2529

This assumes the common operand is on the LHS. This will be the case for constants, but may not hold for non-constants. We should handle the commuted case as well.

2532

It looks like it's legal to transfer the nsw flag to the sub here (https://alive2.llvm.org/ce/z/cYJh4x). But not in the nuw case (https://alive2.llvm.org/ce/z/jHsEcP).

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

3–4

Drop data layout.

5–6

Could you please change the variable names to something more meaningful?

565

Please add multi-use tests. I think this fold needs one-use checks.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

2532

Okay, I see in the tests that the nsw will get inferred anyway, so I guess there's no need to make it explicit.

[dmgreen]

Thanks for the review. Updated for one use checks and commutivity. This might start to get stranger in a followup but this part should hopefully be sensible enough on it's own.

[nikic]

LGTM

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

49

This test case looks the same as subaddnuw_c1?

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)))) &&
+       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(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);
+    }
+  }
+
   if (Instruction *Res = foldBinOpOfSelectAndCastOfSelectCondition(I))
     return Res;
 
   return TryToNarrowDeduceFlags();
 }
 
 /// This eliminates floating-point negation in either 'fneg(X)' or
 /// 'fsub(-0.0, X)' form by combining into a constant operand.

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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -passes=instcombine -S | FileCheck %s
 
 define i32 @subaddnuw(i8 %a, i8 %C1, i8 %C2) {
 ; CHECK-LABEL: @subaddnuw(
-; CHECK-NEXT:    [[A1:%.*]] = add nuw i8 [[A:%.*]], [[C1:%.*]]
-; CHECK-NEXT:    [[Z1:%.*]] = zext i8 [[A1]] to i32
-; CHECK-NEXT:    [[A2:%.*]] = add nuw i8 [[A]], [[C2:%.*]]
-; CHECK-NEXT:    [[Z2:%.*]] = zext i8 [[A2]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[Z2]], [[Z1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[C2:%.*]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[C1:%.*]] to i32
+; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %a1 = add nuw i8 %a, %C1
   %z1 = zext i8 %a1 to i32
   %a2 = add nuw i8 %a, %C2
   %z2 = zext i8 %a2 to i32
   %s = sub i32 %z2, %z1
   ret i32 %s
 }
 
 define i32 @subaddnuw_c1(i8 %a, i8 %C1, i8 %C2) {
 ; CHECK-LABEL: @subaddnuw_c1(
-; CHECK-NEXT:    [[A1:%.*]] = add nuw i8 [[C1:%.*]], [[A:%.*]]
-; CHECK-NEXT:    [[Z1:%.*]] = zext i8 [[A1]] to i32
-; CHECK-NEXT:    [[A2:%.*]] = add nuw i8 [[A]], [[C2:%.*]]
-; CHECK-NEXT:    [[Z2:%.*]] = zext i8 [[A2]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[Z2]], [[Z1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[C2:%.*]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[C1:%.*]] to i32
+; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %a1 = add nuw i8 %C1, %a
   %z1 = zext i8 %a1 to i32
   %a2 = add nuw i8 %a, %C2
   %z2 = zext i8 %a2 to i32
   %s = sub i32 %z2, %z1
   ret i32 %s
 }
 
 define i32 @subaddnuw_c2(i8 %a, i8 %C1, i8 %C2) {
 ; CHECK-LABEL: @subaddnuw_c2(
-; CHECK-NEXT:    [[A1:%.*]] = add nuw i8 [[C1:%.*]], [[A:%.*]]
-; CHECK-NEXT:    [[Z1:%.*]] = zext i8 [[A1]] to i32
-; CHECK-NEXT:    [[A2:%.*]] = add nuw i8 [[C2:%.*]], [[A]]
-; CHECK-NEXT:    [[Z2:%.*]] = zext i8 [[A2]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[Z2]], [[Z1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[C2:%.*]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[C1:%.*]] to i32
+; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %a1 = add nuw i8 %C1, %a
   %z1 = zext i8 %a1 to i32
   %a2 = add nuw i8 %C2, %a
   %z2 = zext i8 %a2 to i32
   %s = sub i32 %z2, %z1
   ret i32 %s
 }
 
 define i32 @subaddnuw_c3(i8 %a, i8 %C1, i8 %C2) {
 ; CHECK-LABEL: @subaddnuw_c3(
-; CHECK-NEXT:    [[A1:%.*]] = add nuw i8 [[C1:%.*]], [[A:%.*]]
-; CHECK-NEXT:    [[Z1:%.*]] = zext i8 [[A1]] to i32
-; CHECK-NEXT:    [[A2:%.*]] = add nuw i8 [[A]], [[C2:%.*]]
-; CHECK-NEXT:    [[Z2:%.*]] = zext i8 [[A2]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[Z2]], [[Z1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[C2:%.*]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[C1:%.*]] to i32
+; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %a1 = add nuw i8 %C1, %a
   %z1 = zext i8 %a1 to i32
   %a2 = add nuw i8 %a, %C2
   %z2 = zext i8 %a2 to i32
   %s = sub i32 %z2, %z1
   ret i32 %s
 }
 
 define i32 @subaddnuw_const(i8 %a) {
 ; CHECK-LABEL: @subaddnuw_const(
-; CHECK-NEXT:    [[A1:%.*]] = add nuw i8 [[A:%.*]], 1
-; CHECK-NEXT:    [[Z1:%.*]] = zext i8 [[A1]] to i32
-; CHECK-NEXT:    [[A2:%.*]] = add nuw 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
 ;
   %a1 = add nuw i8 %a, 1
   %z1 = zext i8 %a1 to i32
   %a2 = add nuw i8 %a, 2
   %z2 = zext i8 %a2 to i32
   %s = sub i32 %z2, %z1
   ret i32 %s
 }
 
 define i32 @subaddnsw(i8 %a, i8 %C1, i8 %C2) {
 ; CHECK-LABEL: @subaddnsw(
-; CHECK-NEXT:    [[A1:%.*]] = add nsw i8 [[A:%.*]], [[C1:%.*]]
-; CHECK-NEXT:    [[Z1:%.*]] = sext i8 [[A1]] to i32
-; CHECK-NEXT:    [[A2:%.*]] = add nsw i8 [[A]], [[C2:%.*]]
-; CHECK-NEXT:    [[Z2:%.*]] = sext i8 [[A2]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[Z2]], [[Z1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = sext i8 [[C2:%.*]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = sext i8 [[C1:%.*]] to i32
+; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %a1 = add nsw i8 %a, %C1
   %z1 = sext i8 %a1 to i32
   %a2 = add nsw i8 %a, %C2
   %z2 = sext i8 %a2 to i32
   %s = sub i32 %z2, %z1
   ret i32 %s
 }
 
 define i32 @subaddnsw_c(i8 %a) {
 ; CHECK-LABEL: @subaddnsw_c(
-; CHECK-NEXT:    [[A1:%.*]] = add nsw i8 [[A:%.*]], 1
-; CHECK-NEXT:    [[Z1:%.*]] = sext i8 [[A1]] to i32
-; CHECK-NEXT:    [[A2:%.*]] = add nsw 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
 ;
   %a1 = add nsw i8 %a, 1
   %z1 = sext i8 %a1 to i32
   %a2 = add nsw i8 %a, 2
   %z2 = sext i8 %a2 to i32
   %s = sub i32 %z2, %z1
   ret i32 %s
 }
;
   %z2 = zext i8 %a2 to i32
   %s = sub i32 %z2, %z1
   ret i32 %s
 }
 
 define i32 @subaddnuw_useadd(i8 %a, i8 %C1, i8 %C2) {
 ; CHECK-LABEL: @subaddnuw_useadd(
 ; CHECK-NEXT:    [[A1:%.*]] = add nuw i8 [[A:%.*]], [[C1:%.*]]
-; CHECK-NEXT:    [[Z1:%.*]] = zext i8 [[A1]] to i32
-; CHECK-NEXT:    [[A2:%.*]] = add nuw i8 [[A]], [[C2:%.*]]
-; CHECK-NEXT:    [[Z2:%.*]] = zext i8 [[A2]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[Z2]], [[Z1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[C2:%.*]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[C1]] to i32
+; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    call void @use8(i8 [[A1]])
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %a1 = add nuw i8 %a, %C1
   %z1 = zext i8 %a1 to i32
   %a2 = add nuw i8 %a, %C2
   %z2 = zext i8 %a2 to i32
   %s = sub i32 %z2, %z1
   call void @use8(i8 %a1)
   ret i32 %s
 }
 
 define i32 @subaddnuw_usezext(i8 %a, i8 %C1, i8 %C2) {
 ; CHECK-LABEL: @subaddnuw_usezext(
 ; CHECK-NEXT:    [[A1:%.*]] = add nuw i8 [[A:%.*]], [[C1:%.*]]
 ; CHECK-NEXT:    [[Z1:%.*]] = zext i8 [[A1]] to i32
-; CHECK-NEXT:    [[A2:%.*]] = add nuw i8 [[A]], [[C2:%.*]]
-; CHECK-NEXT:    [[Z2:%.*]] = zext i8 [[A2]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[Z2]], [[Z1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[C2:%.*]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[C1]] to i32
+; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    call void @use32(i32 [[Z1]])
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %a1 = add nuw i8 %a, %C1
   %z1 = zext i8 %a1 to i32
   %a2 = add nuw i8 %a, %C2
   %z2 = zext i8 %a2 to i32
   %s = sub i32 %z2, %z1
   call void @use32(i32 %z1)
   ret i32 %s
 }
 
 define i32 @subaddnuw_twouseadd(i8 %a, i8 %C1, i8 %C2) {
 ; CHECK-LABEL: @subaddnuw_twouseadd(
 ; CHECK-NEXT:    [[A1:%.*]] = add nuw i8 [[A:%.*]], [[C1:%.*]]
-; CHECK-NEXT:    [[Z1:%.*]] = zext i8 [[A1]] to i32
 ; CHECK-NEXT:    [[A2:%.*]] = add nuw i8 [[A]], [[C2:%.*]]
-; CHECK-NEXT:    [[Z2:%.*]] = zext i8 [[A2]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[Z2]], [[Z1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[C2]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[C1]] to i32
+; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[TMP1]], [[TMP2]]
 ; CHECK-NEXT:    call void @use8(i8 [[A1]])
 ; CHECK-NEXT:    call void @use8(i8 [[A2]])
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %a1 = add nuw i8 %a, %C1
   %z1 = zext i8 %a1 to i32
   %a2 = add nuw i8 %a, %C2
   %z2 = zext i8 %a2 to i32

 }
 
 define i32 @subaddnuw_twousezext_c(i8 %a) {
 ; CHECK-LABEL: @subaddnuw_twousezext_c(
 ; CHECK-NEXT:    [[A1:%.*]] = add nuw i8 [[A:%.*]], 1
 ; CHECK-NEXT:    [[Z1:%.*]] = zext i8 [[A1]] to i32
 ; CHECK-NEXT:    [[A2:%.*]] = add nuw i8 [[A]], 2
 ; CHECK-NEXT:    [[Z2:%.*]] = zext i8 [[A2]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[Z2]], [[Z1]]
 ; CHECK-NEXT:    call void @use32(i32 [[Z1]])
 ; CHECK-NEXT:    call void @use32(i32 [[Z2]])
-; CHECK-NEXT:    ret i32 [[S]]
+; CHECK-NEXT:    ret i32 1
 ;
   %a1 = add nuw i8 %a, 1
   %z1 = zext i8 %a1 to i32
   %a2 = add nuw i8 %a, 2
   %z2 = zext i8 %a2 to i32
   %s = sub i32 %z2, %z1
   call void @use32(i32 %z1)
   call void @use32(i32 %z2)
bb:
   %z1 = zext i8 %a1 to i32
   %a2 = add i8 %a, 2
   %z2 = zext i8 %a2 to i32
   %s = sub i32 %z2, %z1
   ret i32 %s
 
 end:
   ret i32 0
 }