[InstCombine] Invert `add A, sext(B) --> sub A, zext(B)` canonicalization (to `sub A, zext B -> add A, sext B`)

[lebedev.ri]

Summary

D68408 proposes to greatly improve our negation sinking abilities.
But in current canonicalization, we produce sub A, zext(B),
which we will consider non-canonical and try to sink that negation,
undoing the existing canonicalization.
So unless we explicitly stop producing previous canonicalization,
we will have two conflicting folds, and will end up endlessly looping.

This inverts canonicalization, and adds back the obvious fold
that we'd miss:

• sub [nsw] Op0, sext/zext (bool Y) -> add [nsw] Op0, zext/sext (bool Y) https://rise4fun.com/Alive/xx4
• sext(bool) + C -> bool ? C - 1 : C https://rise4fun.com/Alive/fBl

It is obvious that @ossfuzz_9880() / @lshr_out_of_range()/@ashr_out_of_range() (oss-fuzz 4871)
are no longer folded as much, though it isn't obvious whether those failures are important?

[spatel]

The select transform is 1 that I thought about but never got around to implementing. We should have at least 1 minimal test for that pattern, but I don't see that in the diffs?

Did you confirm that codegen is equal or better for these cases (apart from the fuzzer tests - I agree that those are not important)? I think we have DAGCombiner reversals for this transform, but some targets that seem like they would benefit have not enabled the TLI hook.

[lebedev.ri]

In D71064#1770937, @spatel wrote:

The select transform is 1 that I thought about but never got around to implementing. We should have at least 1 minimal test for that pattern, but I don't see that in the diffs?

See llvm-project/llvm/test/Transforms/InstCombine/add.ll.

llvm/test/Transforms/InstCombine/add.ll

4–44 ↗

(On Diff #232330)

@spatel i believe this is the test coverage for select transform

[lebedev.ri]

In D71064#1770937, @spatel wrote:

Did you confirm that codegen is equal or better for these cases?
I think we have DAGCombiner reversals for this transform, but some targets that seem like they would benefit have not enabled the TLI hook.

I didn't yet. So, as far as i can tell, these are all the interesting cases:
(i'm looking at @t0_new_canon vs @t1_old_canon since that is the only change in this patch)
https://godbolt.org/z/vbb25Q - no regression for x86
https://godbolt.org/z/Htg7m5 - aarch64 also looks ok?
https://godbolt.org/z/xmP6JV - arm good?
https://godbolt.org/z/vNrNJ8 - thumb good?

So i'd say everything is already covered by backend undo folds?
Let me know if i'm missing the point here?

(apart from the fuzzer tests - I agree that those are not important)

Good to know.

[spatel]

In D71064#1770969, @lebedev.ri wrote:

In D71064#1770937, @spatel wrote:

Did you confirm that codegen is equal or better for these cases?
I think we have DAGCombiner reversals for this transform, but some targets that seem like they would benefit have not enabled the TLI hook.

I didn't yet. So, as far as i can tell, these are all the interesting cases:
(i'm looking at @t0_new_canon vs @t1_old_canon since that is the only change in this patch)
https://godbolt.org/z/vbb25Q - no regression for x86
https://godbolt.org/z/Htg7m5 - aarch64 also looks ok?
https://godbolt.org/z/xmP6JV - arm good?
https://godbolt.org/z/vNrNJ8 - thumb good?

So i'd say everything is already covered by backend undo folds?
Let me know if i'm missing the point here?

Scalar looks same all-around. Vector shows some potential diffs:
https://godbolt.org/z/y3E-mb

If I'm seeing it correctly, we always do better on the typical case where the bool vector is produced by a compare, but we might do worse if we don't have that cmp and don't have AssertSext knowledge.

[lebedev.ri]

In D71064#1771030, @spatel wrote:

In D71064#1770969, @lebedev.ri wrote:

In D71064#1770937, @spatel wrote:

Did you confirm that codegen is equal or better for these cases?
I think we have DAGCombiner reversals for this transform, but some targets that seem like they would benefit have not enabled the TLI hook.

I didn't yet. So, as far as i can tell, these are all the interesting cases:
(i'm looking at @t0_new_canon vs @t1_old_canon since that is the only change in this patch)
https://godbolt.org/z/vbb25Q - no regression for x86
https://godbolt.org/z/Htg7m5 - aarch64 also looks ok?
https://godbolt.org/z/xmP6JV - arm good?
https://godbolt.org/z/vNrNJ8 - thumb good?

So i'd say everything is already covered by backend undo folds?
Let me know if i'm missing the point here?

Scalar looks same all-around. Vector shows some potential diffs:
https://godbolt.org/z/y3E-mb

If I'm seeing it correctly, we always do better on the typical case where the bool vector is produced by a compare, but we might do worse if we don't have that cmp and don't have AssertSext knowledge.

So the comment is that the undo fold needs to be adjusted first, to fire for non-cmp i1 vectors on aarch64 and powerpc64le?

[lebedev.ri]

In D71064#1771063, @lebedev.ri wrote:

In D71064#1771030, @spatel wrote:

Scalar looks same all-around. Vector shows some potential diffs:
https://godbolt.org/z/y3E-mb

If I'm seeing it correctly, we always do better on the typical case where the bool vector is produced by a compare, but we might do worse if we don't have that cmp and don't have AssertSext knowledge.

So the comment is that the undo fold needs to be adjusted first, to fire for non-cmp i1 vectors on aarch64 and powerpc64le?

[spatel]

In D71064#1771063, @lebedev.ri wrote:

In D71064#1771030, @spatel wrote:

Scalar looks same all-around. Vector shows some potential diffs:
https://godbolt.org/z/y3E-mb

If I'm seeing it correctly, we always do better on the typical case where the bool vector is produced by a compare, but we might do worse if we don't have that cmp and don't have AssertSext knowledge.

So the comment is that the undo fold needs to be adjusted first, to fire for non-cmp i1 vectors on aarch64 and powerpc64le?

I don't think that's necessary as a preliminary step because we're improving what should be the common patterns (the examples that include the cmp).
LGTM.

[lebedev.ri]

In D71064#1771082, @spatel wrote:

In D71064#1771063, @lebedev.ri wrote:

In D71064#1771030, @spatel wrote:

Scalar looks same all-around. Vector shows some potential diffs:
https://godbolt.org/z/y3E-mb

If I'm seeing it correctly, we always do better on the typical case where the bool vector is produced by a compare, but we might do worse if we don't have that cmp and don't have AssertSext knowledge.

So the comment is that the undo fold needs to be adjusted first, to fire for non-cmp i1 vectors on aarch64 and powerpc64le?

I don't think that's necessary as a preliminary step because we're improving what should be the common patterns (the examples that include the cmp).
LGTM.

Ah, okay, good point :)
Thank you for the review.

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

Instruction *InstCombiner::foldAddWithConstant(BinaryOperator &Add) {
   if (match(Op0, m_OneUse(m_Sub(m_Value(X), m_Value(Y)))) &&
       match(Op1, m_AllOnes()))
     return BinaryOperator::CreateAdd(Builder.CreateNot(Y), X);
 
   // zext(bool) + C -> bool ? C + 1 : C
   if (match(Op0, m_ZExt(m_Value(X))) &&
       X->getType()->getScalarSizeInBits() == 1)
     return SelectInst::Create(X, AddOne(Op1C), Op1);
+  // sext(bool) + C -> bool ? C - 1 : C
+  if (match(Op0, m_SExt(m_Value(X))) &&
+      X->getType()->getScalarSizeInBits() == 1)
+    return SelectInst::Create(X, SubOne(Op1C), Op1);
 
   // ~X + C --> (C-1) - X
   if (match(Op0, m_Not(m_Value(X))))
     return BinaryOperator::CreateSub(SubOne(Op1C), X);
 
   const APInt *C;
   if (!match(Op1, m_APInt(C)))
     return nullptr;
if (match(LHS, m_Neg(m_Value(A)))) {
     // -A + -B --> -(A + B)
     if (match(RHS, m_Neg(m_Value(B))))
       return BinaryOperator::CreateNeg(Builder.CreateAdd(A, B));
 
     // -A + B --> B - A
     return BinaryOperator::CreateSub(RHS, A);
   }
 
-  // Canonicalize sext to zext for better value tracking potential.
-  // add A, sext(B) --> sub A, zext(B)
-  if (match(&I, m_c_Add(m_Value(A), m_OneUse(m_SExt(m_Value(B))))) &&
-      B->getType()->isIntOrIntVectorTy(1))
-    return BinaryOperator::CreateSub(A, Builder.CreateZExt(B, Ty));
-
   // A + -B  -->  A - B
   if (match(RHS, m_Neg(m_Value(B))))
     return BinaryOperator::CreateSub(LHS, B);
 
   if (Value *V = checkForNegativeOperand(I, Builder))
     return replaceInstUsesWith(I, V);
 
   // (A + 1) + ~B --> A - B
if (Op1->hasOneUse()) {
     // 'nuw' is dropped in favor of the canonical form.
     if (match(Op1, m_SExt(m_Value(Y))) &&
         Y->getType()->getScalarSizeInBits() == 1) {
       Value *Zext = Builder.CreateZExt(Y, I.getType());
       BinaryOperator *Add = BinaryOperator::CreateAdd(Op0, Zext);
       Add->setHasNoSignedWrap(I.hasNoSignedWrap());
       return Add;
     }
+    // sub [nsw] X, zext(bool Y) -> add [nsw] X, sext(bool Y)
+    // 'nuw' is dropped in favor of the canonical form.
+    if (match(Op1, m_ZExt(m_Value(Y))) && Y->getType()->isIntOrIntVectorTy(1)) {
+      Value *Sext = Builder.CreateSExt(Y, I.getType());
+      BinaryOperator *Add = BinaryOperator::CreateAdd(Op0, Sext);
+      Add->setHasNoSignedWrap(I.hasNoSignedWrap());
+      return Add;
+    }
 
     // X - A*-B -> X + A*B
     // X - -A*B -> X + A*B
     Value *A, *B;
     if (match(Op1, m_c_Mul(m_Value(A), m_Neg(m_Value(B)))))
       return BinaryOperator::CreateAdd(Op0, Builder.CreateMul(A, B));
 
     // X - A*C -> X + A*-C

llvm/test/Transforms/InstCombine/apint-shift.ll

 }
 
 ; OSS-Fuzz #9880
 ; https://bugs.chromium.org/p/oss-fuzz/issues/detail?id=9880
 define i177 @ossfuzz_9880(i177 %X) {
 ; CHECK-LABEL: @ossfuzz_9880(
 ; CHECK-NEXT:    [[A:%.*]] = alloca i177, align 8
 ; CHECK-NEXT:    [[L1:%.*]] = load i177, i177* [[A]], align 8
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i177 [[L1]], 0
-; CHECK-NEXT:    [[B1:%.*]] = zext i1 [[TMP1]] to i177
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i177 [[L1]], -1
+; CHECK-NEXT:    [[TMP2:%.*]] = sext i1 [[TMP1]] to i177
+; CHECK-NEXT:    [[B14:%.*]] = add i177 [[L1]], [[TMP2]]
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i177 [[B14]], -1
+; CHECK-NEXT:    [[B1:%.*]] = zext i1 [[TMP3]] to i177
 ; CHECK-NEXT:    ret i177 [[B1]]
 ;
   %A = alloca i177
   %L1 = load i177, i177* %A
   %B = or i177 0, -1
   %B5 = udiv i177 %L1, %B
   %B4 = add i177 %B5, %B
   %B2 = add i177 %B, %B4
   %B6 = mul i177 %B5, %B2
   %B20 = shl i177 %L1, %B6
   %B14 = sub i177 %B20, %B5
   %B1 = udiv i177 %B14, %B6
   ret i177 %B1
 }

llvm/test/Transforms/InstCombine/logical-select.ll

 }
 
 ; Allow the transform even if the mask values have multiple uses because
 ; there's still a net reduction of instructions from removing the and/and/or.
 
 define <4 x i32> @vec_sel_xor_multi_use(<4 x i32> %a, <4 x i32> %b, <4 x i1> %c) {
 ; CHECK-LABEL: @vec_sel_xor_multi_use(
 ; CHECK-NEXT:    [[TMP1:%.*]] = xor <4 x i1> [[C:%.*]], <i1 true, i1 false, i1 false, i1 false>
+; CHECK-NEXT:    [[MASK_FLIP1:%.*]] = sext <4 x i1> [[TMP1]] to <4 x i32>
 ; CHECK-NEXT:    [[TMP2:%.*]] = xor <4 x i1> [[C]], <i1 false, i1 true, i1 true, i1 true>
 ; CHECK-NEXT:    [[TMP3:%.*]] = select <4 x i1> [[TMP2]], <4 x i32> [[A:%.*]], <4 x i32> [[B:%.*]]
-; CHECK-NEXT:    [[TMP4:%.*]] = zext <4 x i1> [[TMP1]] to <4 x i32>
-; CHECK-NEXT:    [[ADD:%.*]] = sub <4 x i32> [[TMP3]], [[TMP4]]
+; CHECK-NEXT:    [[ADD:%.*]] = add <4 x i32> [[TMP3]], [[MASK_FLIP1]]
 ; CHECK-NEXT:    ret <4 x i32> [[ADD]]
 ;
   %mask = sext <4 x i1> %c to <4 x i32>
   %mask_flip1 = xor <4 x i32> %mask, <i32 -1, i32 0, i32 0, i32 0>
   %not_mask_flip1 = xor <4 x i32> %mask, <i32 0, i32 -1, i32 -1, i32 -1>
   %and1 = and <4 x i32> %not_mask_flip1, %a
   %and2 = and <4 x i32> %mask_flip1, %b
   %or = or <4 x i32> %and1, %and2

llvm/test/Transforms/InstCombine/select.ll

;
   %s = select i1 %cond, i32 %y, i32 %z
   %r = and i32 %x, %s
   ret i32 %r
 }
 
 define i32 @test42(i32 %x, i32 %y) {
 ; CHECK-LABEL: @test42(
 ; CHECK-NEXT:    [[COND:%.*]] = icmp eq i32 [[X:%.*]], 0
-; CHECK-NEXT:    [[TMP1:%.*]] = zext i1 [[COND]] to i32
-; CHECK-NEXT:    [[C:%.*]] = sub i32 [[Y:%.*]], [[TMP1]]
+; CHECK-NEXT:    [[B:%.*]] = sext i1 [[COND]] to i32
+; CHECK-NEXT:    [[C:%.*]] = add i32 [[B]], [[Y:%.*]]
 ; CHECK-NEXT:    ret i32 [[C]]
 ;
   %b = add i32 %y, -1
   %cond = icmp eq i32 %x, 0
   %c = select i1 %cond, i32 %b, i32 %y
   ret i32 %c
 }
 
 define <2 x i32> @test42vec(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @test42vec(
 ; CHECK-NEXT:    [[COND:%.*]] = icmp eq <2 x i32> [[X:%.*]], zeroinitializer
-; CHECK-NEXT:    [[TMP1:%.*]] = zext <2 x i1> [[COND]] to <2 x i32>
-; CHECK-NEXT:    [[C:%.*]] = sub <2 x i32> [[Y:%.*]], [[TMP1]]
+; CHECK-NEXT:    [[B:%.*]] = sext <2 x i1> [[COND]] to <2 x i32>
+; CHECK-NEXT:    [[C:%.*]] = add <2 x i32> [[B]], [[Y:%.*]]
 ; CHECK-NEXT:    ret <2 x i32> [[C]]
 ;
   %b = add <2 x i32> %y, <i32 -1, i32 -1>
   %cond = icmp eq <2 x i32> %x, zeroinitializer
   %c = select <2 x i1> %cond, <2 x i32> %b, <2 x i32> %y
   ret <2 x i32> %c
 }
 

llvm/test/Transforms/InstCombine/shift.ll

;
   %3 = ashr i32 %2, 1
   ret i32 %3
 }
 
 ; OSS Fuzz #4871
 ; https://bugs.chromium.org/p/oss-fuzz/issues/detail?id=4871
 define i177 @lshr_out_of_range(i177 %Y, i177** %A2) {
 ; CHECK-LABEL: @lshr_out_of_range(
-; CHECK-NEXT:    store i177** [[A2:%.*]], i177*** undef, align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp ne i177 [[Y:%.*]], -1
+; CHECK-NEXT:    [[B4:%.*]] = sext i1 [[TMP1]] to i177
+; CHECK-NEXT:    [[C8:%.*]] = icmp ult i177 [[B4]], [[Y]]
+; CHECK-NEXT:    [[TMP2:%.*]] = sext i1 [[C8]] to i64
+; CHECK-NEXT:    [[G18:%.*]] = getelementptr i177*, i177** [[A2:%.*]], i64 [[TMP2]]
+; CHECK-NEXT:    store i177** [[G18]], i177*** undef, align 8
 ; CHECK-NEXT:    ret i177 0
 ;
   %B5 = udiv i177 %Y, -1
   %B4 = add i177 %B5, -1
   %B2 = add i177 %B4, -1
   %B6 = mul i177 %B5, %B2
   %B3 = add i177 %B2, %B2
   %B10 = sub i177 %B5, %B3
   %B12 = lshr i177 %Y, %B6
   %C8 = icmp ugt i177 %B12, %B4
   %G18 = getelementptr i177*, i177** %A2, i1 %C8
   store i177** %G18, i177*** undef
   %B1 = udiv i177 %B10, %B6
   ret i177 %B1
 }
 
 ; OSS Fuzz #5032
 ; https://bugs.chromium.org/p/oss-fuzz/issues/detail?id=5032
 define void @ashr_out_of_range(i177* %A) {
 ; CHECK-LABEL: @ashr_out_of_range(
+; CHECK-NEXT:    [[L:%.*]] = load i177, i177* [[A:%.*]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i177 [[L]], -1
+; CHECK-NEXT:    [[B2:%.*]] = select i1 [[TMP1]], i64 -1, i64 -2
+; CHECK-NEXT:    [[G11:%.*]] = getelementptr i177, i177* [[A]], i64 [[B2]]
+; CHECK-NEXT:    [[L7:%.*]] = load i177, i177* [[G11]], align 4
+; CHECK-NEXT:    [[B36:%.*]] = select i1 [[TMP1]], i177 0, i177 [[L7]]
+; CHECK-NEXT:    [[C17:%.*]] = icmp sgt i177 [[B36]], [[L7]]
+; CHECK-NEXT:    [[TMP2:%.*]] = sext i1 [[C17]] to i64
+; CHECK-NEXT:    [[G62:%.*]] = getelementptr i177, i177* [[G11]], i64 [[TMP2]]
+; CHECK-NEXT:    [[TMP3:%.*]] = icmp eq i177 [[L7]], -1
+; CHECK-NEXT:    [[B28:%.*]] = select i1 [[TMP3]], i177 0, i177 [[L7]]
+; CHECK-NEXT:    store i177 [[B28]], i177* [[G62]], align 4
 ; CHECK-NEXT:    ret void
 ;
   %L = load i177, i177* %A
   %B5 = udiv i177 %L, -1
   %B4 = add i177 %B5, -1
   %B2 = add i177 %B4, -1
   %G11 = getelementptr i177, i177* %A, i177 %B2
   %L7 = load i177, i177* %G11

llvm/test/Transforms/InstCombine/zext-bool-add-sub.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -instcombine -S | FileCheck %s
 
 ; rdar://11748024
 
 define i32 @a(i1 zeroext %x, i1 zeroext %y) {
 ; CHECK-LABEL: @a(
+; CHECK-NEXT:    [[CONV3_NEG:%.*]] = sext i1 [[Y:%.*]] to i32
 ; CHECK-NEXT:    [[SUB:%.*]] = select i1 [[X:%.*]], i32 2, i32 1
-; CHECK-NEXT:    [[TMP1:%.*]] = zext i1 [[Y:%.*]] to i32
-; CHECK-NEXT:    [[ADD:%.*]] = sub nsw i32 [[SUB]], [[TMP1]]
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[SUB]], [[CONV3_NEG]]
 ; CHECK-NEXT:    ret i32 [[ADD]]
 ;
   %conv = zext i1 %x to i32
   %conv3 = zext i1 %y to i32
   %conv3.neg = sub i32 0, %conv3
   %sub = add i32 %conv, 1
   %add = add i32 %sub, %conv3.neg
   ret i32 %add

 ;
   %sext = sext i1 %y to i8
   %sub = sub nuw i8 %x, %sext
   ret i8 %sub
 }
 
 define i32 @sextbool_add(i1 %c, i32 %x) {
 ; CHECK-LABEL: @sextbool_add(
-; CHECK-NEXT:    [[TMP1:%.*]] = zext i1 [[C:%.*]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub i32 [[X:%.*]], [[TMP1]]
+; CHECK-NEXT:    [[B:%.*]] = sext i1 [[C:%.*]] to i32
+; CHECK-NEXT:    [[S:%.*]] = add i32 [[B]], [[X:%.*]]
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %b = sext i1 %c to i32
   %s = add i32 %b, %x
   ret i32 %s
 }
 
 define i32 @sextbool_add_commute(i1 %c, i32 %px) {
 ; CHECK-LABEL: @sextbool_add_commute(
 ; CHECK-NEXT:    [[X:%.*]] = urem i32 [[PX:%.*]], 42
-; CHECK-NEXT:    [[TMP1:%.*]] = zext i1 [[C:%.*]] to i32
-; CHECK-NEXT:    [[S:%.*]] = sub nsw i32 [[X]], [[TMP1]]
+; CHECK-NEXT:    [[B:%.*]] = sext i1 [[C:%.*]] to i32
+; CHECK-NEXT:    [[S:%.*]] = add nsw i32 [[X]], [[B]]
 ; CHECK-NEXT:    ret i32 [[S]]
 ;
   %x = urem i32 %px, 42 ; thwart complexity-based canonicalization
   %b = sext i1 %c to i32
   %s = add i32 %x, %b
   ret i32 %s
 }
 
;
   %b = sext i1 %c to i32
   call void @use32(i32 %b)
   %s = add i32 %b, %x
   ret i32 %s
 }
 
 define <4 x i32> @sextbool_add_vector(<4 x i1> %c, <4 x i32> %x) {
 ; CHECK-LABEL: @sextbool_add_vector(
-; CHECK-NEXT:    [[TMP1:%.*]] = zext <4 x i1> [[C:%.*]] to <4 x i32>
-; CHECK-NEXT:    [[S:%.*]] = sub <4 x i32> [[X:%.*]], [[TMP1]]
+; CHECK-NEXT:    [[B:%.*]] = sext <4 x i1> [[C:%.*]] to <4 x i32>
+; CHECK-NEXT:    [[S:%.*]] = add <4 x i32> [[B]], [[X:%.*]]
 ; CHECK-NEXT:    ret <4 x i32> [[S]]
 ;
   %b = sext <4 x i1> %c to <4 x i32>
   %s = add <4 x i32> %x, %b
   ret <4 x i32> %s
 }
 
 define i32 @zextbool_sub(i1 %c, i32 %x) {
;
   %b = zext i1 %c to i32
   call void @use32(i32 %b)
   %s = sub i32 %x, %b
   ret i32 %s
 }
 
 define <4 x i32> @zextbool_sub_vector(<4 x i1> %c, <4 x i32> %x) {
 ; CHECK-LABEL: @zextbool_sub_vector(
-; CHECK-NEXT:    [[B:%.*]] = zext <4 x i1> [[C:%.*]] to <4 x i32>
-; CHECK-NEXT:    [[S:%.*]] = sub <4 x i32> [[X:%.*]], [[B]]
+; CHECK-NEXT:    [[TMP1:%.*]] = sext <4 x i1> [[C:%.*]] to <4 x i32>
+; CHECK-NEXT:    [[S:%.*]] = add <4 x i32> [[TMP1]], [[X:%.*]]
 ; CHECK-NEXT:    ret <4 x i32> [[S]]
 ;
   %b = zext <4 x i1> %c to <4 x i32>
   %s = sub <4 x i32> %x, %b
   ret <4 x i32> %s
 }