[llvm] [InstCombine] Eliminate icmp+zext pairs over phis more aggressively (PR #121767)

[Nikita Popov via llvm-commits]

https://github.com/nikic updated https://github.com/llvm/llvm-project/pull/121767

>From 92b826f61e0dad8003a188ac1a15b5f95228ddca Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Mon, 6 Jan 2025 15:03:33 +0100
Subject: [PATCH] [InstCombine] Eliminate icmp+zext pairs over phis more
 aggressively

---
 .../InstCombine/InstructionCombining.cpp      | 29 +++++++++++++++----
 llvm/test/Transforms/InstCombine/phi.ll       | 11 ++++---
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index f63de1f0d410e2..553435c937a70a 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -1822,12 +1822,29 @@ Instruction *InstCombinerImpl::foldOpIntoPhi(Instruction &I, PHINode *PN,
       continue;
     }
 
-    // If the only use of phi is comparing it with a constant then we can
-    // put this comparison in the incoming BB directly after a ucmp/scmp call
-    // because we know that it will simplify to a single icmp.
-    const APInt *Ignored;
-    if (isa<CmpIntrinsic>(InVal) && InVal->hasOneUser() &&
-        match(&I, m_ICmp(m_Specific(PN), m_APInt(Ignored)))) {
+    // Handle some cases that can't be fully simplified, but where we know that
+    // the two instructions will fold into one.
+    auto WillFold = [&]() {
+      if (!InVal->hasOneUser())
+        return false;
+
+      // icmp of ucmp/scmp with constant will fold to icmp.
+      const APInt *Ignored;
+      if (isa<CmpIntrinsic>(InVal) &&
+          match(&I, m_ICmp(m_Specific(PN), m_APInt(Ignored))))
+        return true;
+
+      // icmp eq zext(bool), 0 will fold to !bool.
+      if (isa<ZExtInst>(InVal) &&
+          cast<ZExtInst>(InVal)->getSrcTy()->isIntOrIntVectorTy(1) &&
+          match(&I,
+                m_SpecificICmp(ICmpInst::ICMP_EQ, m_Specific(PN), m_Zero())))
+        return true;
+
+      return false;
+    };
+
+    if (WillFold()) {
       OpsToMoveUseToIncomingBB.push_back(i);
       NewPhiValues.push_back(nullptr);
       continue;
diff --git a/llvm/test/Transforms/InstCombine/phi.ll b/llvm/test/Transforms/InstCombine/phi.ll
index 33849291d832fc..4756b4f30e5606 100644
--- a/llvm/test/Transforms/InstCombine/phi.ll
+++ b/llvm/test/Transforms/InstCombine/phi.ll
@@ -2828,13 +2828,13 @@ define i1 @test_zext_icmp_eq_0(i1 %a, i1 %b, i32 %c) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br i1 [[A:%.*]], label [[IF:%.*]], label [[ELSE:%.*]]
 ; CHECK:       if:
-; CHECK-NEXT:    [[B_EXT:%.*]] = zext i1 [[B:%.*]] to i32
+; CHECK-NEXT:    [[TMP0:%.*]] = xor i1 [[B:%.*]], true
 ; CHECK-NEXT:    br label [[JOIN:%.*]]
 ; CHECK:       else:
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i32 [[C:%.*]], 0
 ; CHECK-NEXT:    br label [[JOIN]]
 ; CHECK:       join:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ [[B_EXT]], [[IF]] ], [ [[C:%.*]], [[ELSE]] ]
-; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[PHI]], 0
+; CHECK-NEXT:    [[CMP:%.*]] = phi i1 [ [[TMP0]], [[IF]] ], [ [[TMP1]], [[ELSE]] ]
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
 entry:
@@ -2916,10 +2916,9 @@ define i1 @test_zext_icmp_eq_0_loop(i1 %c, i1 %b) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
-; CHECK-NEXT:    [[PHI:%.*]] = phi i32 [ 1, [[ENTRY:%.*]] ], [ [[EXT:%.*]], [[LOOP]] ]
-; CHECK-NEXT:    [[X:%.*]] = icmp eq i32 [[PHI]], 0
+; CHECK-NEXT:    [[X:%.*]] = phi i1 [ false, [[ENTRY:%.*]] ], [ [[TMP0:%.*]], [[LOOP]] ]
 ; CHECK-NEXT:    [[Y:%.*]] = and i1 [[X]], [[B:%.*]]
-; CHECK-NEXT:    [[EXT]] = zext i1 [[Y]] to i32
+; CHECK-NEXT:    [[TMP0]] = xor i1 [[Y]], true
 ; CHECK-NEXT:    br i1 [[C:%.*]], label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:
 ; CHECK-NEXT:    ret i1 [[X]]