[InstCombine] improve demanded bits for Sub operand 0

[spatel]

Summary

This is copying the code that was added for add with D130075. (That patch removed a fallthrough in the cases, but we can probably still share at least some code again as a follow-up cleanup, but I didn't want to risk it here.)

The reasoning is similar to the carry propagation for add: if we don't demand low bits of the subtraction and the subtrahend (aka RHS or operand 1) is known zero in those low bits, then there can't be any borrowing required from the higher bits of operand 0, so they don't matter.

Also, the no-wrap flags can be propagated (and I think that should be true for add too).

Here's an attempt to prove that in Alive2:
https://alive2.llvm.org/ce/z/xqh7Pa
(can add nsw or nuw from src and tgt, and it should still pass)

[bcl5980]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineSimplifyDemanded.cpp

case Instruction::Sub: {
     APInt DemandedFromOps = APInt::getLowBitsSet(BitWidth, BitWidth - NLZ);
     if (ShrinkDemandedConstant(I, 1, DemandedFromOps) ||
         SimplifyDemandedBits(I, 1, DemandedFromOps, RHSKnown, Depth + 1))
       return disableWrapFlagsBasedOnUnusedHighBits(I, NLZ);
     if (ShrinkDemandedConstant(I, 0, DemandedFromOps) ||
         SimplifyDemandedBits(I, 0, DemandedFromOps, LHSKnown, Depth + 1))
       return disableWrapFlagsBasedOnUnusedHighBits(I, NLZ);
 
+    // If low order bits are not demanded and are known to be zero in RHS,
+    // then we don't need to demand them from LHS, since they can't cause a
+    // borrow from any bits that are demanded in the result.
+    unsigned NTZ = (~DemandedMask & RHSKnown.Zero).countTrailingOnes();
+    APInt DemandedFromLHS = DemandedFromOps;
+    DemandedFromLHS.clearLowBits(NTZ);
+    if (ShrinkDemandedConstant(I, 0, DemandedFromLHS) ||
+        SimplifyDemandedBits(I, 0, DemandedFromLHS, LHSKnown, Depth + 1))
+      return I;
+
     // If we are known to be subtracting zeros from every bit below
     // the highest demanded bit, we just return the other side.
     if (DemandedFromOps.isSubsetOf(RHSKnown.Zero))
       return I->getOperand(0);
     // We can't do this with the LHS for subtraction, unless we are only
     // demanding the LSB.
     if (DemandedFromOps.isOne() && DemandedFromOps.isSubsetOf(LHSKnown.Zero))
       return I->getOperand(1);

llvm/test/Transforms/InstCombine/sub.ll

;
   %m = and i8 %x, -64 ; 0xC0
   %a = sub i8 %m, %y
   call void @use8(i8 %a)
   %s = sub i8 %a, %z
   %r = shl i8 %s, 2
   ret i8 %r
 }
 
+; sub becomes negate and combines with shl
+
 define i8 @shrink_sub_from_constant_lowbits(i8 %x) {
 ; CHECK-LABEL: @shrink_sub_from_constant_lowbits(
-; CHECK-NEXT:    [[X000:%.*]] = shl i8 [[X:%.*]], 3
-; CHECK-NEXT:    [[SUB:%.*]] = sub i8 7, [[X000]]
-; CHECK-NEXT:    [[R:%.*]] = and i8 [[SUB]], -8
-; CHECK-NEXT:    ret i8 [[R]]
+; CHECK-NEXT:    [[X000_NEG:%.*]] = mul i8 [[X:%.*]], -8
+; CHECK-NEXT:    ret i8 [[X000_NEG]]
 ;
   %x000 = shl i8 %x, 3   ; 3 low bits are known zero
   %sub = sub i8 7, %x000
   %r = and i8 %sub, -8   ; 3 low bits are not demanded
   ret i8 %r
 }
 
+; negative test - extra use prevents shrinking '7'
+
 define i8 @shrink_sub_from_constant_lowbits_uses(i8 %x) {
 ; CHECK-LABEL: @shrink_sub_from_constant_lowbits_uses(
 ; CHECK-NEXT:    [[X000:%.*]] = shl i8 [[X:%.*]], 3
 ; CHECK-NEXT:    [[SUB:%.*]] = sub i8 7, [[X000]]
 ; CHECK-NEXT:    call void @use8(i8 [[SUB]])
 ; CHECK-NEXT:    [[R:%.*]] = and i8 [[SUB]], -8
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x000 = shl i8 %x, 3   ; 3 low bits are known zero
   %sub = sub i8 7, %x000
   call void @use8(i8 %sub)
   %r = and i8 %sub, -8   ; 3 low bits are not demanded
   ret i8 %r
 }
 
+; safe to clear 3 low bits (2 higher bits remain set)
+
 define i8 @shrink_sub_from_constant_lowbits2(i8 %x) {
 ; CHECK-LABEL: @shrink_sub_from_constant_lowbits2(
 ; CHECK-NEXT:    [[X000:%.*]] = and i8 [[X:%.*]], -8
-; CHECK-NEXT:    [[SUB:%.*]] = sub nsw i8 30, [[X000]]
+; CHECK-NEXT:    [[SUB:%.*]] = sub nsw i8 24, [[X000]]
 ; CHECK-NEXT:    [[R:%.*]] = and i8 [[SUB]], -16
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x000 = and i8 %x, -8   ; 3 low bits are known zero
   %sub = sub nsw i8 30, %x000 ; 0b0001_1110
   %r = and i8 %sub, -16   ; 4 low bits are not demanded
   ret i8 %r
 }
 
+; safe to clear 3 low bits (2 higher bits remain set)
+
 define <2 x i8> @shrink_sub_from_constant_lowbits3(<2 x i8> %x) {
 ; CHECK-LABEL: @shrink_sub_from_constant_lowbits3(
 ; CHECK-NEXT:    [[X0000:%.*]] = shl <2 x i8> [[X:%.*]], <i8 4, i8 4>
-; CHECK-NEXT:    [[SUB:%.*]] = sub nuw <2 x i8> <i8 31, i8 31>, [[X0000]]
-; CHECK-NEXT:    [[R:%.*]] = lshr <2 x i8> [[SUB]], <i8 3, i8 3>
+; CHECK-NEXT:    [[SUB:%.*]] = sub nuw <2 x i8> <i8 24, i8 24>, [[X0000]]
+; CHECK-NEXT:    [[R:%.*]] = lshr exact <2 x i8> [[SUB]], <i8 3, i8 3>
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %x0000 = shl <2 x i8> %x, <i8 4, i8 4>     ; 4 low bits are known zero
   %sub = sub nuw <2 x i8> <i8 31, i8 31>, %x0000
   %r = lshr <2 x i8> %sub, <i8 3, i8 3>      ; 3 low bits are not demanded
   ret <2 x i8> %r
 }
 
+; eliminate the mask of y or the mask of the result
+
 define i8 @demand_sub_from_variable_lowbits(i8 %x, i8 %y) {
 ; CHECK-LABEL: @demand_sub_from_variable_lowbits(
 ; CHECK-NEXT:    [[X000:%.*]] = shl i8 [[X:%.*]], 3
-; CHECK-NEXT:    [[Y000:%.*]] = and i8 [[Y:%.*]], -8
-; CHECK-NEXT:    [[SUB:%.*]] = sub i8 [[Y000]], [[X000]]
-; CHECK-NEXT:    ret i8 [[SUB]]
+; CHECK-NEXT:    [[SUB:%.*]] = sub i8 [[Y:%.*]], [[X000]]
+; CHECK-NEXT:    [[R:%.*]] = and i8 [[SUB]], -8
+; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x000 = shl i8 %x, 3   ; 3 low bits are known zero
   %y000 = and i8 %y, -8
   %sub = sub i8 %y000, %x000
   %r = and i8 %sub, -8   ; 3 low bits are not demanded
   ret i8 %r
 }
 
+; setting the low 3 bits of y doesn't change anything
+
 define i8 @demand_sub_from_variable_lowbits2(i8 %x, i8 %y) {
 ; CHECK-LABEL: @demand_sub_from_variable_lowbits2(
 ; CHECK-NEXT:    [[X0000:%.*]] = shl i8 [[X:%.*]], 4
-; CHECK-NEXT:    [[Y111:%.*]] = or i8 [[Y:%.*]], 7
-; CHECK-NEXT:    [[SUB:%.*]] = sub nuw nsw i8 [[Y111]], [[X0000]]
+; CHECK-NEXT:    [[SUB:%.*]] = sub nuw nsw i8 [[Y:%.*]], [[X0000]]
 ; CHECK-NEXT:    [[R:%.*]] = lshr i8 [[SUB]], 4
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x0000 = shl i8 %x, 4   ; 4 low bits are known zero
   %y111 = or i8 %y, 7
   %sub = sub nsw nuw i8 %y111, %x0000
   %r = lshr i8 %sub, 4    ; 4 low bits are not demanded
   ret i8 %r
 }
 
+; negative test - the mask of y removes an extra bit, so that instruction is needed
+
 define i8 @demand_sub_from_variable_lowbits3(i8 %x, i8 %y) {
 ; CHECK-LABEL: @demand_sub_from_variable_lowbits3(
 ; CHECK-NEXT:    [[X0000:%.*]] = shl i8 [[X:%.*]], 4
 ; CHECK-NEXT:    [[Y00000:%.*]] = and i8 [[Y:%.*]], -32
 ; CHECK-NEXT:    [[SUB:%.*]] = sub i8 [[Y00000]], [[X0000]]
 ; CHECK-NEXT:    [[R:%.*]] = lshr exact i8 [[SUB]], 4
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x0000 = shl i8 %x, 4   ; 4 low bits are known zero
   %y00000 = and i8 %y, -32
   %sub = sub i8 %y00000, %x0000
   %r = lshr i8 %sub, 4    ; 4 low bits are not demanded
   ret i8 %r
 }