From d9e1f9d7591b0d3e4df3c0e33ffd8984fb1632a5 Mon Sep 17 00:00:00 2001
From: Sanjay Patel <spatel@rotateright.com>
Date: Thu, 8 Sep 2022 10:44:49 -0400
Subject: [PATCH] [InstCombine] Fold icmp of truncated left shift

(trunc (1 << Y) to iN) == 0 --> Y u>= N
(trunc (1 << Y) to iN) != 0 --> Y u<  N

These can be generalized in several ways as noted by the TODO
items, but this handles the pattern in the motivating bug report.

Fixes #51889

Differential Revision: https://reviews.llvm.org/D115480
---
 .../Transforms/InstCombine/InstCombineCompares.cpp | 12 ++++++++++
 llvm/test/Transforms/InstCombine/icmp-trunc.ll     | 26 +++++++++++-----------
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index febfc13..35b9e37 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -1551,6 +1551,18 @@ Instruction *InstCombinerImpl::foldICmpTruncConstant(ICmpInst &Cmp,
   Type *SrcTy = X->getType();
   unsigned DstBits = Trunc->getType()->getScalarSizeInBits(),
            SrcBits = SrcTy->getScalarSizeInBits();
+
+  // (trunc (1 << Y) to iN) == 0 --> Y u>= N
+  // (trunc (1 << Y) to iN) != 0 --> Y u<  N
+  // TODO: Handle any shifted constant by subtracting trailing zeros.
+  // TODO: Handle non-equality predicates.
+  // TODO: Handle compare to power-of-2 (non-zero) constant.
+  Value *Y;
+  if (Cmp.isEquality() && C.isZero() && match(X, m_Shl(m_One(), m_Value(Y)))) {
+    auto NewPred = Pred == Cmp.ICMP_EQ ? Cmp.ICMP_UGE : Cmp.ICMP_ULT;
+    return new ICmpInst(NewPred, Y, ConstantInt::get(SrcTy, DstBits));
+  }
+
   if (Cmp.isEquality() && Trunc->hasOneUse()) {
     // Canonicalize to a mask and wider compare if the wide type is suitable:
     // (trunc X to i8) == C --> (X & 0xff) == (zext C)
diff --git a/llvm/test/Transforms/InstCombine/icmp-trunc.ll b/llvm/test/Transforms/InstCombine/icmp-trunc.ll
index a6715be..5d4ccf7 100644
--- a/llvm/test/Transforms/InstCombine/icmp-trunc.ll
+++ b/llvm/test/Transforms/InstCombine/icmp-trunc.ll
@@ -356,15 +356,9 @@ define i1 @trunc_ne_i64_i10(i64 %x) {
 }
 
 define i1 @shl1_trunc_eq0(i32 %a) {
-; DL64-LABEL: @shl1_trunc_eq0(
-; DL64-NEXT:    [[R:%.*]] = icmp ugt i32 [[A:%.*]], 15
-; DL64-NEXT:    ret i1 [[R]]
-;
-; DL8-LABEL: @shl1_trunc_eq0(
-; DL8-NEXT:    [[SHL:%.*]] = shl i32 1, [[A:%.*]]
-; DL8-NEXT:    [[T:%.*]] = trunc i32 [[SHL]] to i16
-; DL8-NEXT:    [[R:%.*]] = icmp eq i16 [[T]], 0
-; DL8-NEXT:    ret i1 [[R]]
+; CHECK-LABEL: @shl1_trunc_eq0(
+; CHECK-NEXT:    [[R:%.*]] = icmp ugt i32 [[A:%.*]], 15
+; CHECK-NEXT:    ret i1 [[R]]
 ;
   %shl = shl i32 1, %a
   %t = trunc i32 %shl to i16
@@ -374,9 +368,7 @@ define i1 @shl1_trunc_eq0(i32 %a) {
 
 define <2 x i1> @shl1_trunc_ne0(<2 x i8> %a) {
 ; CHECK-LABEL: @shl1_trunc_ne0(
-; CHECK-NEXT:    [[SHL:%.*]] = shl <2 x i8> <i8 1, i8 poison>, [[A:%.*]]
-; CHECK-NEXT:    [[T:%.*]] = trunc <2 x i8> [[SHL]] to <2 x i5>
-; CHECK-NEXT:    [[R:%.*]] = icmp ne <2 x i5> [[T]], zeroinitializer
+; CHECK-NEXT:    [[R:%.*]] = icmp ult <2 x i8> [[A:%.*]], <i8 5, i8 5>
 ; CHECK-NEXT:    ret <2 x i1> [[R]]
 ;
   %shl = shl <2 x i8> <i8 1, i8 poison>, %a
@@ -404,7 +396,7 @@ define i1 @shl1_trunc_ne0_use2(i37 %a) {
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i37 1, [[A:%.*]]
 ; CHECK-NEXT:    [[T:%.*]] = trunc i37 [[SHL]] to i8
 ; CHECK-NEXT:    call void @use(i8 [[T]])
-; CHECK-NEXT:    [[R:%.*]] = icmp ne i8 [[T]], 0
+; CHECK-NEXT:    [[R:%.*]] = icmp ult i37 [[A]], 8
 ; CHECK-NEXT:    ret i1 [[R]]
 ;
   %shl = shl i37 1, %a
@@ -414,6 +406,8 @@ define i1 @shl1_trunc_ne0_use2(i37 %a) {
   ret i1 %r
 }
 
+; TODO: A > 4
+
 define i1 @shl2_trunc_eq0(i9 %a) {
 ; CHECK-LABEL: @shl2_trunc_eq0(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i9 2, [[A:%.*]]
@@ -427,6 +421,8 @@ define i1 @shl2_trunc_eq0(i9 %a) {
   ret i1 %r
 }
 
+; TODO: A < 5
+
 define i1 @shl1_trunc_sgt0(i9 %a) {
 ; CHECK-LABEL: @shl1_trunc_sgt0(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i9 1, [[A:%.*]]
@@ -440,6 +436,8 @@ define i1 @shl1_trunc_sgt0(i9 %a) {
   ret i1 %r
 }
 
+; TODO: A == 0
+
 define i1 @shl1_trunc_eq1(i64 %a) {
 ; CHECK-LABEL: @shl1_trunc_eq1(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i64 1, [[A:%.*]]
@@ -455,6 +453,8 @@ define i1 @shl1_trunc_eq1(i64 %a) {
   ret i1 %r
 }
 
+; TODO: A != 5
+
 define i1 @shl1_trunc_ne32(i8 %a) {
 ; CHECK-LABEL: @shl1_trunc_ne32(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl i8 1, [[A:%.*]]
-- 
2.7.4