[InstCombine] Combine ZExt (B - A) + ZExt(A) to ZExt(B)

[StephenFan]

Summary

Combine ZExt (B - A) + ZExt(A) to ZExt(B)
https://alive2.llvm.org/ce/z/ESUwPi

[dmgreen]

Just a drive-by comment, but I believe it needs nuw on the sub.

[spatel]

In D139930#3991653, @dmgreen wrote:

Just a drive-by comment, but I believe it needs nuw on the sub.

Yes - the last test is a miscompile:
https://alive2.llvm.org/ce/z/xd8Pdp

[StephenFan]

Check NUW for sub.

[spatel]

Please pre-commit the baseline tests, so we just show the functional changes from this patch in the tests.

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1452

Use the raw constructor instead of Builder to preserve the value name:

return new ZextInst(B, LHS->getType());

llvm/test/Transforms/InstCombine/add.ll

2350

Add "_commute" to the test name to make the difference clear.

2355

Change types to increase test coverage. You could use weird types like <2 x i5>.
We also want to verify at the limits of the fold, so zext by just 1 bit from an i2 (or even i1?).
It would also be good to add extra uses to at least one test because the extra uses don't change the transform. See examples in this file of call void @use(i8).

2370

Add another (negative) test where the sub operands are swapped?

[StephenFan]

1. Use new instead of builder.
2. Rename to commute
3. Add zext <2 x i5> to <2 x i8> test
4. Add zext i2 to i3 test
5. Add test that has uses.
6. Add negative test of commuted sub operands.

[spatel]

I think this is good now, but please pre-commit the baseline tests.

[StephenFan]

Rebase.

[spatel]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

Instruction *InstCombinerImpl::visitAdd(BinaryOperator &I) {
 
   // (A & 2^C1) + A => A & (2^C1 - 1) iff bit C1 in A is a sign bit
   if (match(&I, m_c_Add(m_And(m_Value(A), m_APInt(C1)), m_Deferred(A))) &&
       C1->isPowerOf2() && (ComputeNumSignBits(A) > C1->countLeadingZeros())) {
     Constant *NewMask = ConstantInt::get(RHS->getType(), *C1 - 1);
     return BinaryOperator::CreateAnd(A, NewMask);
   }
 
+  // ZExt (B - A) + ZExt(A) --> ZExt(B)
+  if ((match(RHS, m_ZExt(m_Value(A))) &&
+       match(LHS, m_ZExt(m_NUWSub(m_Value(B), m_Specific(A))))) ||
+      (match(LHS, m_ZExt(m_Value(A))) &&
+       match(RHS, m_ZExt(m_NUWSub(m_Value(B), m_Specific(A))))))
+    return new ZExtInst(B, LHS->getType());
+
   // A+B --> A|B iff A and B have no bits set in common.
   if (haveNoCommonBitsSet(LHS, RHS, DL, &AC, &I, &DT))
     return BinaryOperator::CreateOr(LHS, RHS);
 
   if (Instruction *Ext = narrowMathIfNoOverflow(I))
     return Ext;
 
   // (add (xor A, B) (and A, B)) --> (or A, B)

llvm/test/Transforms/InstCombine/add.ll

 ;
   %m = mul i8 %x, -3
   %a = add i8 %m, 42
   ret i8 %a
 }
 
 define i16 @add_sub_zext(i8 %x, i8 %y) {
 ; CHECK-LABEL: @add_sub_zext(
-; CHECK-NEXT:    [[TMP1:%.*]] = sub nuw i8 [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[TMP1]] to i16
-; CHECK-NEXT:    [[TMP3:%.*]] = zext i8 [[X]] to i16
-; CHECK-NEXT:    [[TMP4:%.*]] = add nuw nsw i16 [[TMP2]], [[TMP3]]
-; CHECK-NEXT:    ret i16 [[TMP4]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[Y:%.*]] to i16
+; CHECK-NEXT:    ret i16 [[TMP1]]
 ;
   %1 = sub nuw i8 %y, %x
   %2 = zext i8 %1 to i16
   %3 = zext i8 %x to i16
   %4 = add i16 %2, %3
   ret i16 %4
 }
 
 define i16 @add_commute_sub_zext(i8 %x, i8 %y) {
 ; CHECK-LABEL: @add_commute_sub_zext(
-; CHECK-NEXT:    [[TMP1:%.*]] = sub nuw i8 [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[TMP1]] to i16
-; CHECK-NEXT:    [[TMP3:%.*]] = zext i8 [[X]] to i16
-; CHECK-NEXT:    [[TMP4:%.*]] = add nuw nsw i16 [[TMP3]], [[TMP2]]
-; CHECK-NEXT:    ret i16 [[TMP4]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i8 [[Y:%.*]] to i16
+; CHECK-NEXT:    ret i16 [[TMP1]]
 ;
   %1 = sub nuw i8 %y, %x
   %2 = zext i8 %1 to i16
   %3 = zext i8 %x to i16
   %4 = add i16 %3, %2
   ret i16 %4
 }
 
 define <2 x i8> @add_sub_2xi5_zext(<2 x i5> %x, <2 x i5> %y) {
 ; CHECK-LABEL: @add_sub_2xi5_zext(
-; CHECK-NEXT:    [[TMP1:%.*]] = sub nuw <2 x i5> [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[TMP2:%.*]] = zext <2 x i5> [[TMP1]] to <2 x i8>
-; CHECK-NEXT:    [[TMP3:%.*]] = zext <2 x i5> [[X]] to <2 x i8>
-; CHECK-NEXT:    [[TMP4:%.*]] = add nuw nsw <2 x i8> [[TMP3]], [[TMP2]]
-; CHECK-NEXT:    ret <2 x i8> [[TMP4]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext <2 x i5> [[Y:%.*]] to <2 x i8>
+; CHECK-NEXT:    ret <2 x i8> [[TMP1]]
 ;
   %1 = sub nuw <2 x i5> %y, %x
   %2 = zext <2 x i5> %1 to <2 x i8>
   %3 = zext <2 x i5> %x to <2 x i8>
   %4 = add <2 x i8> %3, %2
   ret <2 x i8> %4
 }
 
 
 define i3 @add_commute_sub_i2_zext_i3(i2 %x, i2 %y) {
 ; CHECK-LABEL: @add_commute_sub_i2_zext_i3(
-; CHECK-NEXT:    [[TMP1:%.*]] = sub nuw i2 [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[TMP2:%.*]] = zext i2 [[TMP1]] to i3
-; CHECK-NEXT:    [[TMP3:%.*]] = zext i2 [[X]] to i3
-; CHECK-NEXT:    [[TMP4:%.*]] = add nuw i3 [[TMP3]], [[TMP2]]
-; CHECK-NEXT:    ret i3 [[TMP4]]
+; CHECK-NEXT:    [[TMP1:%.*]] = zext i2 [[Y:%.*]] to i3
+; CHECK-NEXT:    ret i3 [[TMP1]]
 ;
   %1 = sub nuw i2 %y, %x
   %2 = zext i2 %1 to i3
   %3 = zext i2 %x to i3
   %4 = add i3 %3, %2
   ret i3 %4
 }
 
 define i16 @add_sub_use_zext(i8 %x, i8 %y) {
 ; CHECK-LABEL: @add_sub_use_zext(
 ; CHECK-NEXT:    [[TMP1:%.*]] = sub nuw i8 [[Y:%.*]], [[X:%.*]]
 ; CHECK-NEXT:    call void @use(i8 [[TMP1]])
-; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[TMP1]] to i16
-; CHECK-NEXT:    [[TMP3:%.*]] = zext i8 [[X]] to i16
-; CHECK-NEXT:    [[TMP4:%.*]] = add nuw nsw i16 [[TMP2]], [[TMP3]]
-; CHECK-NEXT:    ret i16 [[TMP4]]
+; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[Y]] to i16
+; CHECK-NEXT:    ret i16 [[TMP2]]
 ;
   %1 = sub nuw i8 %y, %x
   call void @use(i8 %1)
   %2 = zext i8 %1 to i16
   %3 = zext i8 %x to i16
   %4 = add i16 %2, %3
   ret i16 %4
 }
;
   %2 = zext i8 %1 to i16
   %3 = zext i8 %x to i16
   %4 = add i16 %3, %2
   ret i16 %4
 }
 
 define i16 @add_sub_zext_constant(i8 %x) {
 ; CHECK-LABEL: @add_sub_zext_constant(
-; CHECK-NEXT:    [[TMP1:%.*]] = sub nuw i8 -2, [[X:%.*]]
-; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[TMP1]] to i16
-; CHECK-NEXT:    [[TMP3:%.*]] = zext i8 [[X]] to i16
-; CHECK-NEXT:    [[TMP4:%.*]] = add nuw nsw i16 [[TMP2]], [[TMP3]]
-; CHECK-NEXT:    ret i16 [[TMP4]]
+; CHECK-NEXT:    ret i16 254
 ;
   %1 = sub nuw i8 254, %x
   %2 = zext i8 %1 to i16
   %3 = zext i8 %x to i16
   %4 = add i16 %2, %3
   ret i16 %4
 }