[llvm] 0310f7f - [InstCombine] Fold `(add X, (sext/zext (icmp eq X, C)))`

[Noah Goldstein via llvm-commits]

Author: Noah Goldstein
Date: 2024-06-01T17:49:15-05:00
New Revision: 0310f7f2d0c56a5697710251cec9803cbf7b4d56

URL: https://github.com/llvm/llvm-project/commit/0310f7f2d0c56a5697710251cec9803cbf7b4d56
DIFF: https://github.com/llvm/llvm-project/commit/0310f7f2d0c56a5697710251cec9803cbf7b4d56.diff

LOG: [InstCombine] Fold `(add X, (sext/zext (icmp eq X, C)))`

We can convert this to a select based on the `(icmp eq X, C)`, then
constant fold the addition the true arm begin `(add C, (sext/zext 1))`
and the false arm being `(add X, 0)` e.g

    - `(select (icmp eq X, C), (add C, (sext/zext 1)), (add X, 0))`.

This is essentially a specialization of the only case that sees to
actually show up from #89020

Closes #93840

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
    llvm/test/Transforms/InstCombine/apint-shift.ll
    llvm/test/Transforms/InstCombine/fold-ext-eq-c-with-op.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
index bff09f5676680..8205b49dfbe2f 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp
@@ -1694,6 +1694,24 @@ Instruction *InstCombinerImpl::visitAdd(BinaryOperator &I) {
     return BinaryOperator::CreateOr(LHS, Zext);
   }
 
+  {
+    Value *Cond, *Ext;
+    Constant *C;
+    // (add X, (sext/zext (icmp eq X, C)))
+    //    -> (select (icmp eq X, C), (add C, (sext/zext 1)), X)
+    auto CondMatcher = m_CombineAnd(
+        m_Value(Cond), m_ICmp(Pred, m_Deferred(A), m_ImmConstant(C)));
+
+    if (match(&I,
+              m_c_Add(m_Value(A),
+                      m_CombineAnd(m_Value(Ext), m_ZExtOrSExt(CondMatcher)))) &&
+        Pred == ICmpInst::ICMP_EQ && Ext->hasOneUse()) {
+      Value *Add = isa<ZExtInst>(Ext) ? InstCombiner::AddOne(C)
+                                      : InstCombiner::SubOne(C);
+      return replaceInstUsesWith(I, Builder.CreateSelect(Cond, Add, A));
+    }
+  }
+
   if (Instruction *Ashr = foldAddToAshr(I))
     return Ashr;
 

diff  --git a/llvm/test/Transforms/InstCombine/apint-shift.ll b/llvm/test/Transforms/InstCombine/apint-shift.ll
index 05c3db70ce1ca..ecf9c4e9c4e69 100644
--- a/llvm/test/Transforms/InstCombine/apint-shift.ll
+++ b/llvm/test/Transforms/InstCombine/apint-shift.ll
@@ -240,8 +240,8 @@ define i23 @test11(i23 %x) {
 
 define i47 @test12(i47 %X) {
 ; CHECK-LABEL: @test12(
-; CHECK-NEXT:    [[SH2:%.*]] = and i47 [[X:%.*]], -256
-; CHECK-NEXT:    ret i47 [[SH2]]
+; CHECK-NEXT:    [[SH1:%.*]] = and i47 [[X:%.*]], -256
+; CHECK-NEXT:    ret i47 [[SH1]]
 ;
   %sh1 = ashr i47 %X, 8
   %sh2 = shl i47 %sh1, 8
@@ -250,8 +250,8 @@ define i47 @test12(i47 %X) {
 
 define <2 x i47> @test12_splat_vec(<2 x i47> %X) {
 ; CHECK-LABEL: @test12_splat_vec(
-; CHECK-NEXT:    [[SH2:%.*]] = and <2 x i47> [[X:%.*]], <i47 -256, i47 -256>
-; CHECK-NEXT:    ret <2 x i47> [[SH2]]
+; CHECK-NEXT:    [[SH1:%.*]] = and <2 x i47> [[X:%.*]], <i47 -256, i47 -256>
+; CHECK-NEXT:    ret <2 x i47> [[SH1]]
 ;
   %sh1 = ashr <2 x i47> %X, <i47 8, i47 8>
   %sh2 = shl <2 x i47> %sh1, <i47 8, i47 8>
@@ -564,14 +564,7 @@ define i40 @test26(i40 %A) {
 ; https://bugs.chromium.org/p/oss-fuzz/issues/detail?id=9880
 define i177 @ossfuzz_9880(i177 %X) {
 ; CHECK-LABEL: @ossfuzz_9880(
-; CHECK-NEXT:    [[A:%.*]] = alloca i177, align 8
-; CHECK-NEXT:    [[L1:%.*]] = load i177, ptr [[A]], align 4
-; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i177 [[L1]], -1
-; CHECK-NEXT:    [[B5_NEG:%.*]] = sext i1 [[TMP1]] to i177
-; CHECK-NEXT:    [[B14:%.*]] = add i177 [[L1]], [[B5_NEG]]
-; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i177 [[B14]], -1
-; CHECK-NEXT:    [[B1:%.*]] = zext i1 [[TMP2]] to i177
-; CHECK-NEXT:    ret i177 [[B1]]
+; CHECK-NEXT:    ret i177 0
 ;
   %A = alloca i177
   %L1 = load i177, ptr %A

diff  --git a/llvm/test/Transforms/InstCombine/fold-ext-eq-c-with-op.ll b/llvm/test/Transforms/InstCombine/fold-ext-eq-c-with-op.ll
index 848b0d189497f..4d02d492d2aa7 100644
--- a/llvm/test/Transforms/InstCombine/fold-ext-eq-c-with-op.ll
+++ b/llvm/test/Transforms/InstCombine/fold-ext-eq-c-with-op.ll
@@ -4,9 +4,7 @@
 declare void @use.i8(i8)
 define i8 @fold_add_zext_eq_0(i8 %x) {
 ; CHECK-LABEL: @fold_add_zext_eq_0(
-; CHECK-NEXT:    [[X_EQ:%.*]] = icmp eq i8 [[X:%.*]], 0
-; CHECK-NEXT:    [[X_EQ_EXT:%.*]] = zext i1 [[X_EQ]] to i8
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[X_EQ_EXT]], [[X]]
+; CHECK-NEXT:    [[R:%.*]] = call i8 @llvm.umax.i8(i8 [[X:%.*]], i8 1)
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %x_eq = icmp eq i8 %x, 0
@@ -19,8 +17,7 @@ define <2 x i8> @fold_add_sext_eq_4_6(<2 x i6> %xx) {
 ; CHECK-LABEL: @fold_add_sext_eq_4_6(
 ; CHECK-NEXT:    [[X:%.*]] = zext <2 x i6> [[XX:%.*]] to <2 x i8>
 ; CHECK-NEXT:    [[X_EQ:%.*]] = icmp eq <2 x i8> [[X]], <i8 4, i8 -127>
-; CHECK-NEXT:    [[X_EQ_EXT:%.*]] = sext <2 x i1> [[X_EQ]] to <2 x i8>
-; CHECK-NEXT:    [[R:%.*]] = add nsw <2 x i8> [[X_EQ_EXT]], [[X]]
+; CHECK-NEXT:    [[R:%.*]] = select <2 x i1> [[X_EQ]], <2 x i8> <i8 3, i8 -128>, <2 x i8> [[X]]
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %x = zext <2 x i6> %xx to <2 x i8>