[InstCombine] Use KnownBits for lshr/add -> (a + b < a)

[Pierre-vh]

Summary

Split from D138814; This patch allows any value for the add operands as long
as they have the right amount of leading zeroes.

Depends on D138814

[Pierre-vh]

Add m_OneUse to match

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

846

Should that be "Only the K trailing bits..." or "The 'width - K' leading bits must be zero."?

[Pierre-vh]

Rebase, address comment

[Pierre-vh]

Forgot some changes

[foad]

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

869–871

Should be '>' instead of '!=' surely? It would be fine if these values happen to have more leading zeros than required.

[Pierre-vh]

Address comments

[spatel]

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

510

Demanded bits should always zap cases like this.
A more interesting test would have just one side of the add with extra leading zeros:
https://alive2.llvm.org/ce/z/c4_YAE

So we'd be back to creating 2 extra instructions vs. the original sequence. That doesn't seem like a good trade-off in IR.

530

wrong hex value comment - 0xFFFF

[arsenm]

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

530

There is a rarely used syntax to use hex constants in the IR but I never remember what it is

[lebedev.ri]

(marking as reviewed)

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

510

Presumably we can simply avoid that unprofitable case,
since e.g. lshr_16_add_known_16_leading_zeroes() looks like an improvement?

[Pierre-vh]

Comments

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

510

It does removes it, not sure why it didn't there. Maybe I forgot to rebase?

[spatel]

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

510

Yes - that seems like it would be ok; just restrict the pre-condition with ShAmt to equality (it was like that in an earlier draft?).

510

My comment might not have been clear:

1. If both sides of the "add" have more leading zeros than the shift amount (the test as shown in this patch @lshr_16_add_known_17_leading_zeroes), we already reduce that "0" via SimplifyDemandedBits(). AFAIK, this patch does not change anything with an example like that.
2. If only one side of the "add" has more leading zeros than the shift amount, then we probably can't get rid of the "and" mask by creating a "trunc". That's what I am seeing in the example I posted ( https://alive2.llvm.org/ce/z/c4_YAE ), and that's a test that I'd like to see included with this patch.

[Pierre-vh]

• Only allow == ShAmt
• Add tests with 16/17 + 32/33 leading zeroes

[spatel]

LGTM

[Pierre-vh]

@lebedev.ri can you please re-review? Thank you

[Pierre-vh]

Ping :)

[Pierre-vh]

@lebedev.ri please re-review - I can't land unless you approve the diff :)

[nikic]

Would you mind pre-commit the new/update tests? The test diff is a bit hard to follow right now.

Do I understand right that we will replace two ands, one add and a lshr with two truncs, a uaddo and a zext? It's not super obvious that this is beneficial, but I guess the hope is that this will interact positively with other transforms (e.g. by pushing the trunc through operations and eliminating it entirely, or something)?

[Pierre-vh]

We don't use uaddo but rather (a + b) < b. The intention is indeed to:

• Positively interact with other transforms
• Get a better canonical form that's more easily recognizable by backends (end goal being to solve https://github.com/RadeonOpenCompute/ROCm/issues/488)

Initially we were using uaddo but it was decided to not have an intrinsic as a canonical form - I don't remember exactly why.

[Pierre-vh]

Rebase + ping

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);
 
+/// Get the upper bound on bit size for this Value \p Op as an unsigned integer.
+/// i.e.  x == zext(trunc(x to MaxSignificantBits) to bitwidth(x)).
+unsigned ComputeMaxUnsignedSignificantBits(const Value *Op,
+                                           const DataLayout &DL,
+                                           unsigned Depth = 0,
+                                           AssumptionCache *AC = nullptr,
+                                           const Instruction *CxtI = nullptr,
+                                           const DominatorTree *DT = 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);
 
 /// Return true if we can prove that the specified FP value is never equal to
 /// -0.0.
 bool CannotBeNegativeZero(const Value *V, const TargetLibraryInfo *TLI,

llvm/include/llvm/Transforms/InstCombine/InstCombiner.h

unsigned ComputeNumSignBits(const Value *Op, unsigned Depth = 0,
     return llvm::ComputeNumSignBits(Op, DL, Depth, &AC, CxtI, &DT);
   }
 
   unsigned ComputeMaxSignificantBits(const Value *Op, unsigned Depth = 0,
                                      const Instruction *CxtI = nullptr) const {
     return llvm::ComputeMaxSignificantBits(Op, DL, Depth, &AC, CxtI, &DT);
   }
 
+  unsigned
+  ComputeMaxUnsignedSignificantBits(const Value *Op, unsigned Depth = 0,
+                                    const Instruction *CxtI = nullptr) const {
+    return llvm::ComputeMaxUnsignedSignificantBits(Op, DL, Depth, &AC, CxtI,
+                                                   &DT);
+  }
+
   OverflowResult computeOverflowForUnsignedMul(const Value *LHS,
                                                const Value *RHS,
                                                const Instruction *CxtI) const {
     return llvm::computeOverflowForUnsignedMul(LHS, RHS, DL, &AC, CxtI, &DT);
   }
 
   OverflowResult computeOverflowForSignedMul(const Value *LHS, const Value *RHS,
                                              const Instruction *CxtI) const {

llvm/lib/Analysis/ValueTracking.cpp

 unsigned llvm::ComputeMaxSignificantBits(const Value *V, const DataLayout &DL,
                                          unsigned Depth, AssumptionCache *AC,
                                          const Instruction *CxtI,
                                          const DominatorTree *DT) {
   unsigned SignBits = ComputeNumSignBits(V, DL, Depth, AC, CxtI, DT);
   return V->getType()->getScalarSizeInBits() - SignBits + 1;
 }
 
+unsigned llvm::ComputeMaxUnsignedSignificantBits(
+    const Value *V, const DataLayout &DL, unsigned Depth, AssumptionCache *AC,
+    const Instruction *CxtI, const DominatorTree *DT) {
+  KnownBits KB = computeKnownBits(V, DL, Depth, AC, CxtI, DT);
+  return V->getType()->getScalarSizeInBits() - KB.countMinLeadingZeros();
+}
+
 static void computeKnownBitsAddSub(bool Add, const Value *Op0, const Value *Op1,
                                    bool NSW, const APInt &DemandedElts,
                                    KnownBits &KnownOut, KnownBits &Known2,
                                    unsigned Depth, const Query &Q) {
   computeKnownBits(Op1, DemandedElts, KnownOut, Depth + 1, Q);
 
   // If one operand is unknown and we have no nowrap information,
   // the result will be unknown independently of the second operand.

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

if (!isa<Constant>(TrueVal) && FBO->getOperand(0) == TrueVal &&
       return SelectInst::Create(Cond, NewShift, NewOp);
     }
   }
 
   return nullptr;
 }
 
 // Tries to perform
-//    (lshr (add (zext X), (zext Y)), K)
-//      -> (icmp ult (add X, Y), X)
+//    (lshr (add X, Y), K)
+//      -> (icmp ult (add (trunc X), (trunc Y)), (trunc X))
 //    where
-//      - The add's operands are zexts from a K-bits integer to a bigger type.
+//      - Only the K trailing bits of X and Y can be non-zero.
 //      - The add is only used by the shr, or by iK (or narrower) truncates.
-//      - The lshr type has more than 2 bits (other types are boolean math).
-//      - K > 1
+//      - K > 1 as we don't want to deal with boolean math here.
 //    note that
 //      - The resulting add cannot have nuw/nsw, else on overflow we get a
 //        poison value and the transform isn't legal anymore.
 Instruction *InstCombinerImpl::foldLShrOverflowBit(BinaryOperator &I) {
   assert(I.getOpcode() == Instruction::LShr);
 
   Value *Add = I.getOperand(0);
   Value *ShiftAmt = I.getOperand(1);
   Type *Ty = I.getType();
 
-  if (Ty->getScalarSizeInBits() < 3)
-    return nullptr;
-
   const APInt *ShAmtAPInt = nullptr;
   Value *X = nullptr, *Y = nullptr;
   if (!match(ShiftAmt, m_APInt(ShAmtAPInt)) ||
-      !match(Add,
-             m_Add(m_OneUse(m_ZExt(m_Value(X))), m_OneUse(m_ZExt(m_Value(Y))))))
+      !match(Add, m_Add(m_OneUse(m_Value(X)), m_OneUse(m_Value(Y)))))
     return nullptr;
 
   const unsigned ShAmt = ShAmtAPInt->getZExtValue();
   if (ShAmt == 1)
     return nullptr;
 
-  // X/Y are zexts from `ShAmt`-sized ints.
-  if (X->getType()->getScalarSizeInBits() != ShAmt ||
-      Y->getType()->getScalarSizeInBits() != ShAmt)
+  if (ComputeMaxUnsignedSignificantBits(X, 0, &I) > ShAmt ||
+      ComputeMaxUnsignedSignificantBits(Y, 0, &I) > ShAmt)
     return nullptr;
 
   // Make sure that `Add` is only used by `I` and `ShAmt`-truncates.
   if (!Add->hasOneUse()) {
     for (User *U : Add->users()) {
       if (U == &I)
         continue;
 
       TruncInst *Trunc = dyn_cast<TruncInst>(U);
       if (!Trunc || Trunc->getType()->getScalarSizeInBits() > ShAmt)
         return nullptr;
     }
   }
 
   // Insert at Add so that the newly created `NarrowAdd` will dominate it's
   // users (i.e. `Add`'s users).
   Instruction *AddInst = cast<Instruction>(Add);
   Builder.SetInsertPoint(AddInst);
 
+  Type *OpTy = Builder.getIntNTy(ShAmt);
+  X = Builder.CreateTrunc(X, OpTy);
+  Y = Builder.CreateTrunc(Y, OpTy);
+
   Value *NarrowAdd = Builder.CreateAdd(X, Y, "add.narrowed");
   Value *Overflow =
       Builder.CreateICmpULT(NarrowAdd, X, "add.narrowed.overflow");
 
   // Replace the uses of the original add with a zext of the
   // NarrowAdd's result. Note that all users at this stage are known to
   // be ShAmt-sized truncs, or the lshr itself.
   if (!Add->hasOneUse())

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

;
   %add = add i32 %zext.a, %zext.b
   %lshr = lshr i32 %add, 16
   %otheruse = or i32 %lshr, %zext.a
   ret i32 %otheruse
 }
 
 define i32 @lshr_16_add_known_16_leading_zeroes(i32 %a, i32 %b) {
 ; CHECK-LABEL: @lshr_16_add_known_16_leading_zeroes(
-; CHECK-NEXT:    [[A16:%.*]] = and i32 [[A:%.*]], 65535
-; CHECK-NEXT:    [[B16:%.*]] = and i32 [[B:%.*]], 65535
-; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw i32 [[A16]], [[B16]]
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 [[ADD]], 16
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[A:%.*]] to i16
+; CHECK-NEXT:    [[TMP2:%.*]] = trunc i32 [[B:%.*]] to i16
+; CHECK-NEXT:    [[TMP3:%.*]] = xor i16 [[TMP1]], -1
+; CHECK-NEXT:    [[ADD_NARROWED_OVERFLOW:%.*]] = icmp ult i16 [[TMP3]], [[TMP2]]
+; CHECK-NEXT:    [[LSHR:%.*]] = zext i1 [[ADD_NARROWED_OVERFLOW]] to i32
 ; CHECK-NEXT:    ret i32 [[LSHR]]
 ;
-  %a16 = and i32 %a, 65535 ; 0x65535
-  %b16 = and i32 %b, 65535 ; 0x65535
+  %a16 = and i32 %a, 65535 ; 0xFFFF
+  %b16 = and i32 %b, 65535 ; 0xFFFF
+  %add = add i32 %a16, %b16
+  %lshr = lshr i32 %add, 16
+  ret i32 %lshr
+}
+
+define i32 @lshr_16_add_known_17_leading_zeroes(i32 %a, i32 %b) {
+; CHECK-LABEL: @lshr_16_add_known_17_leading_zeroes(
+; CHECK-NEXT:    ret i32 0
+;
+  %a16 = and i32 %a, 4095 ; 0xFFF
+  %b16 = and i32 %b, 4095 ; 0xFFF
   %add = add i32 %a16, %b16
   %lshr = lshr i32 %add, 16
   ret i32 %lshr
 }
 
 define i32 @lshr_16_add_not_known_16_leading_zeroes(i32 %a, i32 %b) {
 ; CHECK-LABEL: @lshr_16_add_not_known_16_leading_zeroes(
 ; CHECK-NEXT:    [[A16:%.*]] = and i32 [[A:%.*]], 131071
 ; CHECK-NEXT:    [[B16:%.*]] = and i32 [[B:%.*]], 65535
 ; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw i32 [[A16]], [[B16]]
 ; CHECK-NEXT:    [[LSHR:%.*]] = lshr i32 [[ADD]], 16
 ; CHECK-NEXT:    ret i32 [[LSHR]]
 ;
   %a16 = and i32 %a, 131071 ; 0x1FFFF
   %b16 = and i32 %b, 65535 ; 0x65535
   %add = add i32 %a16, %b16
   %lshr = lshr i32 %add, 16
   ret i32 %lshr
 }
 
 define i64 @lshr_32_add_zext_basic(i32 %a, i32 %b) {
 ; CHECK-LABEL: @lshr_32_add_zext_basic(
 ; CHECK-NEXT:    [[TMP1:%.*]] = xor i32 [[A:%.*]], -1
;
   %zext.b = zext i16 %b to i64
   %add = add i64 %zext.a, %zext.b
   %lshr = lshr i64 %add, 16
   ret i64 %lshr
 }
 
 define i64 @lshr_32_add_known_32_leading_zeroes(i64 %a, i64 %b) {
 ; CHECK-LABEL: @lshr_32_add_known_32_leading_zeroes(
-; CHECK-NEXT:    [[A32:%.*]] = and i64 [[A:%.*]], 4294967295
-; CHECK-NEXT:    [[B32:%.*]] = and i64 [[B:%.*]], 4294967295
-; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw i64 [[A32]], [[B32]]
-; CHECK-NEXT:    [[LSHR:%.*]] = lshr i64 [[ADD]], 32
+; CHECK-NEXT:    [[TMP1:%.*]] = trunc i64 [[A:%.*]] to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = trunc i64 [[B:%.*]] to i32
+; CHECK-NEXT:    [[TMP3:%.*]] = xor i32 [[TMP1]], -1
+; CHECK-NEXT:    [[ADD_NARROWED_OVERFLOW:%.*]] = icmp ult i32 [[TMP3]], [[TMP2]]
+; CHECK-NEXT:    [[LSHR:%.*]] = zext i1 [[ADD_NARROWED_OVERFLOW]] to i64
 ; CHECK-NEXT:    ret i64 [[LSHR]]
 ;
   %a32 = and i64 %a, 4294967295 ; 0xFFFFFFFF
   %b32 = and i64 %b, 4294967295 ; 0xFFFFFFFF
   %add = add i64 %a32, %b32
   %lshr = lshr i64 %add, 32
   ret i64 %lshr
 }
 
+define i64 @lshr_32_add_known_33_leading_zeroes(i64 %a, i64 %b) {
+; CHECK-LABEL: @lshr_32_add_known_33_leading_zeroes(
+; CHECK-NEXT:    ret i64 0
+;
+  %a32 = and i64 %a, 268435455 ; 0xFFFFFFF
+  %b32 = and i64 %b, 268435455 ; 0xFFFFFFF
+  %add = add i64 %a32, %b32
+  %lshr = lshr i64 %add, 32
+  ret i64 %lshr
+}
+
 define i64 @lshr_32_add_not_known_32_leading_zeroes(i64 %a, i64 %b) {
 ;
 ; CHECK-LABEL: @lshr_32_add_not_known_32_leading_zeroes(
 ; CHECK-NEXT:    [[A32:%.*]] = and i64 [[A:%.*]], 8589934591
 ; CHECK-NEXT:    [[B32:%.*]] = and i64 [[B:%.*]], 4294967295
 ; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw i64 [[A32]], [[B32]]
 ; CHECK-NEXT:    [[LSHR:%.*]] = lshr i64 [[ADD]], 32
 ; CHECK-NEXT:    ret i64 [[LSHR]]