[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

2950–2951

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.

2952

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

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

2952

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

2939

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"?

2952

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(
   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.
   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 >> C   <u x --> true for C > 0.
+  // x udiv C <u x --> true 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->ugt(0)) ||
+      (match(LBO, m_UDiv(m_Specific(RHS), m_APInt(C))) && C->ugt(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

 ; 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_nonzero1(i32 %x) {
 ; CHECK-LABEL: @lshr_nonzero1(
 ; 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_nonzero2(i32 %x) {
 ; CHECK-LABEL: @lshr_nonzero2(
 ; 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_nonzero3(i32 %x) {
 ; CHECK-LABEL: @lshr_nonzero3(
 ; 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_nonzero4(i32 %x) {
 ; CHECK-LABEL: @lshr_nonzero4(
 ; 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_nonzero1(i32 %x) {
 ; CHECK-LABEL: @udiv_nonzero1(
 ; 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_nonzero2(i32 %x) {
 ; CHECK-LABEL: @udiv_nonzero2(
 ; 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_nonzero3(i32 %x) {
 ; CHECK-LABEL: @udiv_nonzero3(
 ; 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_nonzero4(i32 %x) {
 ; CHECK-LABEL: @udiv_nonzero4(
 ; 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
 }