[InstCombine] Fold 'check for [no] signed truncation' pattern

[lebedev.ri]

Summary

PR38149

As discussed in https://reviews.llvm.org/D49179#1158957 and later,
the IR for 'check for [no] signed truncation' pattern can be improved:
https://rise4fun.com/Alive/gBf
^ that pattern will be produced by Implicit Integer Truncation sanitizer,
https://reviews.llvm.org/D48958 https://bugs.llvm.org/show_bug.cgi?id=21530
in signed case, therefore it is probably a good idea to improve it.

The DAGCombine will reverse this transform, see
https://reviews.llvm.org/D49266

This transform is surprisingly frustrating.
This does not deal with non-splat shift amounts, or with undef shift amounts.
I've outlined what i think the solution should be:

// Potential handling of non-splats: for each element:
//  * if both are undef, replace with constant 0.
//    Because (1<<0) is OK and is 1, and ((1<<0)>>1) is also OK and is 0.
//  * if both are not undef, and are different, bailout.
//  * else, only one is undef, then pick the non-undef one.

This is a re-commit, as the original patch, committed in rL337190
was reverted in rL337344 as it broke chromium build:
https://bugs.llvm.org/show_bug.cgi?id=38204 and
https://crbug.com/864832
Proofs that the fixed folds are ok: https://rise4fun.com/Alive/VYM

[spatel]

LGTM

[lebedev.ri]

In D49320#1163592, @spatel wrote:

LGTM

Thank you for the review.
I'm not sure about the case with undef's.
Right now, honestly, i only care about the simple scalar case, so i'll just leave that unimplemented for now..

[lebedev.ri]

Was reverted in rL337344 by @inglorion:

Revert "[InstCombine] Fold 'check for [no] signed truncation' pattern"

This reverts r337190 (and a few follow-up commits), which caused the
Chromium build to fail. See
https://bugs.llvm.org/show_bug.cgi?id=38204 and
https://crbug.com/864832

Not sure what is the problem yet.

[lebedev.ri]

Update the diff in the differential with the straight un-revert for now.
Will look into the problem next.

[lebedev.ri]

Rebase on top of more tests, that show miscompile.

[lebedev.ri]

Operate on APInt's directly, fixes PR38204.

Once this handled non-splats, this will have to not operate on these APInt's,
but operate on ConstantInt vectors directly anyway.

lib/Transforms/InstCombine/InstCombineCompares.cpp

case ICmpInst::Predicate::ICMP_SLE:
     break;
   default:
     llvm_unreachable("All possible folds are handled.");
   }
 
   return Builder.CreateICmp(DstPred, X, M);
 }
 
+/// Some comparisons can be simplified.
+/// In this case, we are looking for comparisons that look like
+/// a check for a lossy signed truncation.
+/// Folds:   (MaskedBits is a constant.)
+///   ((%x << MaskedBits) a>> MaskedBits) SrcPred %x
+/// Into:
+///   (add %x, (1 << (KeptBits-1))) DstPred (1 << KeptBits)
+/// Where  KeptBits = bitwidth(%x) - MaskedBits
+static Value *
+foldICmpWithTruncSignExtendedVal(ICmpInst &I,
+                                 InstCombiner::BuilderTy &Builder) {
+  ICmpInst::Predicate SrcPred;
+  Value *X;
+  const APInt *C0, *C1; // FIXME: non-splats, potentially with undef.
+  // We are ok with 'shl' having multiple uses, but 'ashr' must be one-use.
+  if (!match(&I, m_c_ICmp(SrcPred,
+                          m_OneUse(m_AShr(m_Shl(m_Value(X), m_APInt(C0)),
+                                          m_APInt(C1))),
+                          m_Deferred(X))))
+    return nullptr;
+
+  // Potential handling of non-splats: for each element:
+  //  * if both are undef, replace with constant 0.
+  //    Because (1<<0) is OK and is 1, and ((1<<0)>>1) is also OK and is 0.
+  //  * if both are not undef, and are different, bailout.
+  //  * else, only one is undef, then pick the non-undef one.
+
+  // The shift amount must be equal.
+  if (*C0 != *C1)
+    return nullptr;
+  const APInt &MaskedBits = *C0;
+  assert(MaskedBits != 0 && "shift by zero should be folded away already.");
+
+  ICmpInst::Predicate DstPred;
+  switch (SrcPred) {
+  case ICmpInst::Predicate::ICMP_EQ:
+    // ((%x << MaskedBits) a>> MaskedBits) == %x
+    //   =>
+    // (add %x, (1 << (KeptBits-1))) u< (1 << KeptBits)
+    DstPred = ICmpInst::Predicate::ICMP_ULT;
+    break;
+  case ICmpInst::Predicate::ICMP_NE:
+    // ((%x << MaskedBits) a>> MaskedBits) != %x
+    //   =>
+    // (add %x, (1 << (KeptBits-1))) u>= (1 << KeptBits)
+    DstPred = ICmpInst::Predicate::ICMP_UGE;
+    break;
+  // FIXME: are more folds possible?
+  default:
+    return nullptr;
+  }
+
+  auto *XType = X->getType();
+  const unsigned XBitWidth = XType->getScalarSizeInBits();
+  const APInt BitWidth = APInt(XBitWidth, XBitWidth);
+  assert(BitWidth.ugt(MaskedBits) && "shifts should leave some bits untouched");
+
+  // KeptBits = bitwidth(%x) - MaskedBits
+  const APInt KeptBits = BitWidth - MaskedBits;
+  assert(KeptBits.ugt(0) && KeptBits.ult(BitWidth) && "unreachable");
+  // ICmpCst = (1 << KeptBits)
+  const APInt ICmpCst = APInt(XBitWidth, 1).shl(KeptBits);
+  assert(ICmpCst.isPowerOf2());
+  // AddCst = (1 << (KeptBits-1))
+  const APInt AddCst = ICmpCst.lshr(1);
+  assert(AddCst.ult(ICmpCst) && AddCst.isPowerOf2());
+
+  // T0 = add %x, AddCst
+  Value *T0 = Builder.CreateAdd(X, ConstantInt::get(XType, AddCst));
+  // T1 = T0 DstPred ICmpCst
+  Value *T1 = Builder.CreateICmp(DstPred, T0, ConstantInt::get(XType, ICmpCst));
+
+  return T1;
+}
+
 /// Try to fold icmp (binop), X or icmp X, (binop).
 /// TODO: A large part of this logic is duplicated in InstSimplify's
 /// simplifyICmpWithBinOp(). We should be able to share that and avoid the code
 /// duplication.
 Instruction *InstCombiner::foldICmpBinOp(ICmpInst &I) {
   Value *Op0 = I.getOperand(0), *Op1 = I.getOperand(1);
 
   // Special logic for binary operators.
if (match(BO0, BitwiseAnd) && Pred == ICmpInst::ICMP_ULT) {
       auto *Zero = Constant::getNullValue(BO0->getType());
       return new ICmpInst(ICmpInst::ICMP_NE, Op1, Zero);
     }
   }
 
   if (Value *V = foldICmpWithLowBitMaskedVal(I, Builder))
     return replaceInstUsesWith(I, V);
 
+  if (Value *V = foldICmpWithTruncSignExtendedVal(I, Builder))
+    return replaceInstUsesWith(I, V);
+
   return nullptr;
 }
 
 /// Fold icmp Pred min|max(X, Y), X.
 static Instruction *foldICmpWithMinMax(ICmpInst &Cmp) {
   ICmpInst::Predicate Pred = Cmp.getPredicate();
   Value *Op0 = Cmp.getOperand(0);
   Value *X = Cmp.getOperand(1);

test/Transforms/InstCombine/canonicalize-lack-of-signed-truncation-check.ll

 ; Where  KeptBits = bitwidth(%x) - MaskedBits
 
 ; ============================================================================ ;
 ; Basic positive tests
 ; ============================================================================ ;
 
 define i1 @p0(i8 %x) {
 ; CHECK-LABEL: @p0(
-; CHECK-NEXT:    [[TMP0:%.*]] = shl i8 [[X:%.*]], 5
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr exact i8 [[TMP0]], 5
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i8 [[TMP1]], [[X]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[X:%.*]], 4
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i8 [[TMP1]], 8
 ; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %tmp0 = shl i8 %x, 5
   %tmp1 = ashr exact i8 %tmp0, 5
   %tmp2 = icmp eq i8 %tmp1, %x
   ret i1 %tmp2
 }
 
 ; Big unusual bit width, https://bugs.llvm.org/show_bug.cgi?id=38204
 define i1 @pb(i65 %x) {
 ; CHECK-LABEL: @pb(
-; CHECK-NEXT:    [[TMP0:%.*]] = shl i65 [[X:%.*]], 1
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr exact i65 [[TMP0]], 1
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i65 [[TMP1]], [[X]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i65 [[X:%.*]], 9223372036854775808
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp sgt i65 [[TMP1]], -1
 ; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %tmp0 = shl i65 %x, 1
   %tmp1 = ashr exact i65 %tmp0, 1
   %tmp2 = icmp eq i65 %x, %tmp1
   ret i1 %tmp2
 }
 
 ; ============================================================================ ;
 ; Vector tests
 ; ============================================================================ ;
 
 define <2 x i1> @p1_vec_splat(<2 x i8> %x) {
 ; CHECK-LABEL: @p1_vec_splat(
-; CHECK-NEXT:    [[TMP0:%.*]] = shl <2 x i8> [[X:%.*]], <i8 5, i8 5>
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr exact <2 x i8> [[TMP0]], <i8 5, i8 5>
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq <2 x i8> [[TMP1]], [[X]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add <2 x i8> [[X:%.*]], <i8 4, i8 4>
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult <2 x i8> [[TMP1]], <i8 8, i8 8>
 ; CHECK-NEXT:    ret <2 x i1> [[TMP2]]
 ;
   %tmp0 = shl <2 x i8> %x, <i8 5, i8 5>
   %tmp1 = ashr exact <2 x i8> %tmp0, <i8 5, i8 5>
   %tmp2 = icmp eq <2 x i8> %tmp1, %x
   ret <2 x i1> %tmp2
 }
 

 ; Commutativity tests.
 ; ============================================================================ ;
 
 declare i8 @gen8()
 
 define i1 @c0() {
 ; CHECK-LABEL: @c0(
 ; CHECK-NEXT:    [[X:%.*]] = call i8 @gen8()
-; CHECK-NEXT:    [[TMP0:%.*]] = shl i8 [[X]], 5
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr exact i8 [[TMP0]], 5
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i8 [[X]], [[TMP1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[X]], 4
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i8 [[TMP1]], 8
 ; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %x = call i8 @gen8()
   %tmp0 = shl i8 %x, 5
   %tmp1 = ashr exact i8 %tmp0, 5
   %tmp2 = icmp eq i8 %x, %tmp1 ; swapped order
   ret i1 %tmp2
 }
 
 ; ============================================================================ ;
 ; One-use tests.
 ; ============================================================================ ;
 
 declare void @use8(i8)
 
 define i1 @n_oneuse0(i8 %x) {
 ; CHECK-LABEL: @n_oneuse0(
 ; CHECK-NEXT:    [[TMP0:%.*]] = shl i8 [[X:%.*]], 5
 ; CHECK-NEXT:    call void @use8(i8 [[TMP0]])
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr exact i8 [[TMP0]], 5
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i8 [[TMP1]], [[X]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[X]], 4
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i8 [[TMP1]], 8
 ; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %tmp0 = shl i8 %x, 5
   call void @use8(i8 %tmp0)
   %tmp1 = ashr exact i8 %tmp0, 5
   %tmp2 = icmp eq i8 %tmp1, %x
   ret i1 %tmp2
 }

test/Transforms/InstCombine/canonicalize-signed-truncation-check.ll

 ; Where  KeptBits = bitwidth(%x) - MaskedBits
 
 ; ============================================================================ ;
 ; Basic positive tests
 ; ============================================================================ ;
 
 define i1 @p0(i8 %x) {
 ; CHECK-LABEL: @p0(
-; CHECK-NEXT:    [[TMP0:%.*]] = shl i8 [[X:%.*]], 5
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr exact i8 [[TMP0]], 5
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i8 [[TMP1]], [[X]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[X:%.*]], 4
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ugt i8 [[TMP1]], 7
 ; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %tmp0 = shl i8 %x, 5
   %tmp1 = ashr exact i8 %tmp0, 5
   %tmp2 = icmp ne i8 %tmp1, %x
   ret i1 %tmp2
 }
 
 ; Big unusual bit width, https://bugs.llvm.org/show_bug.cgi?id=38204
 define i1 @pb(i65 %x) {
 ; CHECK-LABEL: @pb(
-; CHECK-NEXT:    [[TMP0:%.*]] = shl i65 [[X:%.*]], 1
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr exact i65 [[TMP0]], 1
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i65 [[TMP1]], [[X]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i65 [[X:%.*]], 9223372036854775808
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp slt i65 [[TMP1]], 0
 ; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %tmp0 = shl i65 %x, 1
   %tmp1 = ashr exact i65 %tmp0, 1
   %tmp2 = icmp ne i65 %x, %tmp1
   ret i1 %tmp2
 }
 
 ; ============================================================================ ;
 ; Vector tests
 ; ============================================================================ ;
 
 define <2 x i1> @p1_vec_splat(<2 x i8> %x) {
 ; CHECK-LABEL: @p1_vec_splat(
-; CHECK-NEXT:    [[TMP0:%.*]] = shl <2 x i8> [[X:%.*]], <i8 5, i8 5>
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr exact <2 x i8> [[TMP0]], <i8 5, i8 5>
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne <2 x i8> [[TMP1]], [[X]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add <2 x i8> [[X:%.*]], <i8 4, i8 4>
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ugt <2 x i8> [[TMP1]], <i8 7, i8 7>
 ; CHECK-NEXT:    ret <2 x i1> [[TMP2]]
 ;
   %tmp0 = shl <2 x i8> %x, <i8 5, i8 5>
   %tmp1 = ashr exact <2 x i8> %tmp0, <i8 5, i8 5>
   %tmp2 = icmp ne <2 x i8> %tmp1, %x
   ret <2 x i1> %tmp2
 }
 

 ; Commutativity tests.
 ; ============================================================================ ;
 
 declare i8 @gen8()
 
 define i1 @c0() {
 ; CHECK-LABEL: @c0(
 ; CHECK-NEXT:    [[X:%.*]] = call i8 @gen8()
-; CHECK-NEXT:    [[TMP0:%.*]] = shl i8 [[X]], 5
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr exact i8 [[TMP0]], 5
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i8 [[X]], [[TMP1]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[X]], 4
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ugt i8 [[TMP1]], 7
 ; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %x = call i8 @gen8()
   %tmp0 = shl i8 %x, 5
   %tmp1 = ashr exact i8 %tmp0, 5
   %tmp2 = icmp ne i8 %x, %tmp1 ; swapped order
   ret i1 %tmp2
 }
 
 ; ============================================================================ ;
 ; One-use tests.
 ; ============================================================================ ;
 
 declare void @use8(i8)
 
 define i1 @n_oneuse0(i8 %x) {
 ; CHECK-LABEL: @n_oneuse0(
 ; CHECK-NEXT:    [[TMP0:%.*]] = shl i8 [[X:%.*]], 5
 ; CHECK-NEXT:    call void @use8(i8 [[TMP0]])
-; CHECK-NEXT:    [[TMP1:%.*]] = ashr exact i8 [[TMP0]], 5
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp ne i8 [[TMP1]], [[X]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[X]], 4
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ugt i8 [[TMP1]], 7
 ; CHECK-NEXT:    ret i1 [[TMP2]]
 ;
   %tmp0 = shl i8 %x, 5
   call void @use8(i8 %tmp0)
   %tmp1 = ashr exact i8 %tmp0, 5
   %tmp2 = icmp ne i8 %tmp1, %x
   ret i1 %tmp2
 }