[InstSimplify] Fold X {lshr,udiv} C <u X --> true for nonzero X, non-identity C

[erikdesjardins]

Summary

This eliminates the bounds check in Rust code like

pub fn mid(data: &[i32]) -> i32 {
  if data.is_empty() { return 0; }
  return data[data.len()/2];
}

(from https://blog.sigplan.org/2021/11/18/undefined-behavior-deserves-a-better-reputation/)

Alive proofs:
lshr https://alive2.llvm.org/ce/z/nyTu8D
udiv https://alive2.llvm.org/ce/z/CNUZH7

[erikdesjardins]

tweak summary formatting

[spatel]

llvm/lib/Analysis/InstructionSimplify.cpp

2977–2978

Spell out the other preds/patterns for completeness? Might want to improve the comments on the block above here while we're at it to show the transform instead of just matched pattern.

2979

Is that TODO feasible? The lshr half could use isKnownNonZero, but how would we deal with ">1" for udiv?

llvm/test/Transforms/InstSimplify/compare.ll

580–581

Put the predicate in the test name rather than incrementing the '1' suffix. That makes it clearer that we're cycling through the preds. Add a couple of negative tests for signed preds?

[erikdesjardins]

make comments more explicit about matched patterns
use more general C != 0 and C != 1 instead of C >u 0 and C >u 1 respectively

[erikdesjardins]

llvm/lib/Analysis/InstructionSimplify.cpp

2979

I think something like

KnownBits Known = computeKnownBits(Divisor, ...);
if (KnownBits::ne(Known, KnownBits::makeConstant(1)))

would technically work.

I don't know if this would help much in practice though; I can remove "divisor" from the comment if you prefer.

[spatel]

LGTM

llvm/lib/Analysis/InstructionSimplify.cpp

2963–2966

Nit: here and below, we're using "u" or "udiv" to make the unsigned requirements explicit, but we're using the ambiguous ">>" for the shift op.
Better to make that ">>u" or spell out "lshr"?

2979

It's fine to leave it. I was just wondering how we'd do that. :)
There's also a question of whether it's worth the compile-time, but these patterns are probably rare enough that it is ok (assuming it would fire on real code enough times to make a difference).

[erikdesjardins]

use >>u to make shift type explicit

[erikdesjardins]

Thanks!

I don't have write access, please commit this (and tests in parent revision https://reviews.llvm.org/D114280) on my behalf: Erik Desjardins <[e-mail]>

llvm/lib/Analysis/InstructionSimplify.cpp

static Value *simplifyICmpWithBinOpOnLHS(
   // icmp pred (urem X, Y), X
   if (match(LBO, m_URem(m_Specific(RHS), m_Value()))) {
     if (Pred == ICmpInst::ICMP_ULE)
       return getTrue(ITy);
     if (Pred == ICmpInst::ICMP_UGT)
       return getFalse(ITy);
   }
 
-  // x >> y <=u x
-  // x udiv y <=u x.
+  // x >>u y <=u x --> true.
+  // x >>u y >u  x --> false.
+  // x udiv y <=u x --> true.
+  // x udiv y >u  x --> false.
   if (match(LBO, m_LShr(m_Specific(RHS), m_Value())) ||
       match(LBO, m_UDiv(m_Specific(RHS), m_Value()))) {
     // icmp pred (X op Y), X
     if (Pred == ICmpInst::ICMP_UGT)
       return getFalse(ITy);
     if (Pred == ICmpInst::ICMP_ULE)
       return getTrue(ITy);
   }
 
+  // If x is nonzero:
+  // x >>u C <u  x --> true  for C != 0.
+  // x >>u C !=  x --> true  for C != 0.
+  // x >>u C >=u x --> false for C != 0.
+  // x >>u C ==  x --> false for C != 0.
+  // x udiv C <u  x --> true  for C != 1.
+  // x udiv C !=  x --> true  for C != 1.
+  // x udiv C >=u x --> false for C != 1.
+  // x udiv C ==  x --> false for C != 1.
+  // TODO: allow non-constant shift amount/divisor
+  const APInt *C;
+  if ((match(LBO, m_LShr(m_Specific(RHS), m_APInt(C))) && *C != 0) ||
+      (match(LBO, m_UDiv(m_Specific(RHS), m_APInt(C))) && *C != 1)) {
+    if (isKnownNonZero(RHS, Q.DL, 0, Q.AC, Q.CxtI, Q.DT)) {
+      switch (Pred) {
+      default:
+        break;
+      case ICmpInst::ICMP_EQ:
+      case ICmpInst::ICMP_UGE:
+        return getFalse(ITy);
+      case ICmpInst::ICMP_NE:
+      case ICmpInst::ICMP_ULT:
+        return getTrue(ITy);
+      case ICmpInst::ICMP_UGT:
+      case ICmpInst::ICMP_ULE:
+        // UGT/ULE are handled by the more general case just above
+        llvm_unreachable("Unexpected UGT/ULE, should have been handled");
+      }
+    }
+  }
+
   // (x*C1)/C2 <= x for C1 <= C2.
   // This holds even if the multiplication overflows: Assume that x != 0 and
   // arithmetic is modulo M. For overflow to occur we must have C1 >= M/x and
   // thus C2 >= M/x. It follows that (x*C1)/C2 <= (M-1)/C2 <= ((M-1)*x)/M < x.
   //
   // Additionally, either the multiplication and division might be represented
   // as shifts:
   // (x*C1)>>C2 <= x for C1 < 2**C2.

llvm/test/Transforms/InstSimplify/compare.ll

 define i1 @lshr7(i32 %X, i32 %Y) {
 ; CHECK-LABEL: @lshr7(
 ; CHECK-NEXT:    ret i1 true
 ;
   %A = lshr i32 %X, %Y
   %C = icmp uge i32 %X, %A
   ret i1 %C
 }
 
 define i1 @lshr_nonzero_eq(i32 %x) {
 ; CHECK-LABEL: @lshr_nonzero_eq(
 ; CHECK-NEXT:    [[X_NE_0:%.*]] = icmp ne i32 [[X:%.*]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[X_NE_0]])
-; CHECK-NEXT:    [[LHS:%.*]] = lshr i32 [[X]], 1
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[LHS]], [[X]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %x_ne_0 = icmp ne i32 %x, 0
   call void @llvm.assume(i1 %x_ne_0)
   %lhs = lshr i32 %x, 1
   %cmp = icmp eq i32 %lhs, %x
   ret i1 %cmp
 }
 
 define i1 @lshr_nonzero_uge(i32 %x) {
 ; CHECK-LABEL: @lshr_nonzero_uge(
 ; CHECK-NEXT:    [[X_NE_0:%.*]] = icmp ne i32 [[X:%.*]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[X_NE_0]])
-; CHECK-NEXT:    [[LHS:%.*]] = lshr i32 [[X]], 1
-; CHECK-NEXT:    [[CMP:%.*]] = icmp uge i32 [[LHS]], [[X]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %x_ne_0 = icmp ne i32 %x, 0
   call void @llvm.assume(i1 %x_ne_0)
   %lhs = lshr i32 %x, 1
   %cmp = icmp uge i32 %lhs, %x
   ret i1 %cmp
 }
 
 define i1 @lshr_nonzero_ne(i32 %x) {
 ; CHECK-LABEL: @lshr_nonzero_ne(
 ; CHECK-NEXT:    [[X_NE_0:%.*]] = icmp ne i32 [[X:%.*]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[X_NE_0]])
-; CHECK-NEXT:    [[LHS:%.*]] = lshr i32 [[X]], 1
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[LHS]], [[X]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 true
 ;
   %x_ne_0 = icmp ne i32 %x, 0
   call void @llvm.assume(i1 %x_ne_0)
   %lhs = lshr i32 %x, 1
   %cmp = icmp ne i32 %lhs, %x
   ret i1 %cmp
 }
 
 define i1 @lshr_nonzero_ult(i32 %x) {
 ; CHECK-LABEL: @lshr_nonzero_ult(
 ; CHECK-NEXT:    [[X_NE_0:%.*]] = icmp ne i32 [[X:%.*]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[X_NE_0]])
-; CHECK-NEXT:    [[LHS:%.*]] = lshr i32 [[X]], 1
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[LHS]], [[X]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 true
 ;
   %x_ne_0 = icmp ne i32 %x, 0
   call void @llvm.assume(i1 %x_ne_0)
   %lhs = lshr i32 %x, 1
   %cmp = icmp ult i32 %lhs, %x
   ret i1 %cmp
 }
 
;
   %C = icmp uge i32 %X, %A
   ret i1 %C
 }
 
 define i1 @udiv_nonzero_eq(i32 %x) {
 ; CHECK-LABEL: @udiv_nonzero_eq(
 ; CHECK-NEXT:    [[X_NE_0:%.*]] = icmp ne i32 [[X:%.*]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[X_NE_0]])
-; CHECK-NEXT:    [[LHS:%.*]] = udiv i32 [[X]], 3
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[LHS]], [[X]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %x_ne_0 = icmp ne i32 %x, 0
   call void @llvm.assume(i1 %x_ne_0)
   %lhs = udiv i32 %x, 3
   %cmp = icmp eq i32 %lhs, %x
   ret i1 %cmp
 }
 
 define i1 @udiv_nonzero_uge(i32 %x) {
 ; CHECK-LABEL: @udiv_nonzero_uge(
 ; CHECK-NEXT:    [[X_NE_0:%.*]] = icmp ne i32 [[X:%.*]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[X_NE_0]])
-; CHECK-NEXT:    [[LHS:%.*]] = udiv i32 [[X]], 3
-; CHECK-NEXT:    [[CMP:%.*]] = icmp uge i32 [[LHS]], [[X]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 false
 ;
   %x_ne_0 = icmp ne i32 %x, 0
   call void @llvm.assume(i1 %x_ne_0)
   %lhs = udiv i32 %x, 3
   %cmp = icmp uge i32 %lhs, %x
   ret i1 %cmp
 }
 
 define i1 @udiv_nonzero_ne(i32 %x) {
 ; CHECK-LABEL: @udiv_nonzero_ne(
 ; CHECK-NEXT:    [[X_NE_0:%.*]] = icmp ne i32 [[X:%.*]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[X_NE_0]])
-; CHECK-NEXT:    [[LHS:%.*]] = udiv i32 [[X]], 3
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[LHS]], [[X]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 true
 ;
   %x_ne_0 = icmp ne i32 %x, 0
   call void @llvm.assume(i1 %x_ne_0)
   %lhs = udiv i32 %x, 3
   %cmp = icmp ne i32 %lhs, %x
   ret i1 %cmp
 }
 
 define i1 @udiv_nonzero_ult(i32 %x) {
 ; CHECK-LABEL: @udiv_nonzero_ult(
 ; CHECK-NEXT:    [[X_NE_0:%.*]] = icmp ne i32 [[X:%.*]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[X_NE_0]])
-; CHECK-NEXT:    [[LHS:%.*]] = udiv i32 [[X]], 3
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i32 [[LHS]], [[X]]
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 true
 ;
   %x_ne_0 = icmp ne i32 %x, 0
   call void @llvm.assume(i1 %x_ne_0)
   %lhs = udiv i32 %x, 3
   %cmp = icmp ult i32 %lhs, %x
   ret i1 %cmp
 }