[llvm] [AggressiveInstCombine] Fold split-width i32 cttz/ctlz patterns into wide i64 intrinsics (PR #192296)

[Craig Topper via llvm-commits]

================
@@ -68,6 +76,439 @@ static cl::opt<unsigned>
                           cl::desc("The maximum length of a constant string to "
                                    "inline a memchr call."));
 
+/// Try to fold a phi that merges two 32-bit cttz operations into a single
+/// 64-bit cttz. This pattern arises when a 64-bit cttz is implemented as
+/// two 32-bit cttz operations with a branch on whether the low half is zero:
+///
+/// GuardBB:
+///   %lo = trunc i64 %val to i32
+///   %cmp = icmp eq i32 %lo, 0
+///   br i1 %cmp, label %HiBB, label %LoBB
+///
+/// HiBB:
+///   %shr = lshr i64 %val, 32
+///   %hi = trunc i64 %shr to i32
+///   %cttz_hi = call i32 @llvm.cttz.i32(i32 %hi, i1 false)
+///   %val_is_zero = icmp eq i64 %val, 0
+///   %cttz_hi_plus32 = or i32 %cttz_hi, 32
+///   %hi_result = select i1 %val_is_zero, i32 ZeroResult, i32 %cttz_hi_plus32
+///   br label %PhiBB
+///
+/// LoBB:
+///   %cttz_lo = call i32 @llvm.cttz.i32(i32 %lo, i1 true)
+///   br label %PhiBB
+///
+/// PhiBB:
+///   %result = phi i32 [%hi_result, %HiBB], [%cttz_lo, %LoBB]
+/// -->
+///   %cttz64 = call i64 @llvm.cttz.i64(i64 %val, i1 true)
+///   %trunc = trunc i64 %cttz64 to i32
+///   %val_is_zero = icmp eq i64 %val, 0
+///   %result = select i1 %val_is_zero, i32 ZeroResult, i32 %trunc
+///
+/// Alive proof:  https://alive2.llvm.org/ce/z/FHgJpq
+static bool foldSplitWidthCTTZ(Instruction &I, const DominatorTree &DT) {
+  if (I.getOpcode() != Instruction::PHI || I.getNumOperands() != 2)
+    return false;
+
+  if (!I.getType()->isIntegerTy(32))
+    return false;
+
+  PHINode &Phi = cast<PHINode>(I);
+
+  // Try both orderings of phi operands.
+  for (unsigned HiIdx = 0; HiIdx < 2; ++HiIdx) {
+    unsigned LoIdx = 1 - HiIdx;
+    Value *HiVal = Phi.getIncomingValue(HiIdx);
+    Value *LoVal = Phi.getIncomingValue(LoIdx);
+    BasicBlock *HiBB = Phi.getIncomingBlock(HiIdx);
+    BasicBlock *LoBB = Phi.getIncomingBlock(LoIdx);
+
+    // Match LoVal: cttz.i32(trunc(SrcVal to i32), _)
+    Value *LoTrunc;
+    if (!match(LoVal,
+               m_Intrinsic<Intrinsic::cttz>(m_Value(LoTrunc), m_Value())))
+      continue;
+
+    // LoTrunc must be a trunc of an i64 value.
+    Value *SrcVal = nullptr;
+    if (!match(LoTrunc, m_Trunc(m_Value(SrcVal))))
+      continue;
+    if (!SrcVal->getType()->isIntegerTy(64))
+      continue;
+
+    // Match HiVal: select(icmp eq SrcVal, 0, C,
+    //                     or(cttz.i32(trunc(lshr(SrcVal, 32)), _), 32))
+    // where C is the result when the full value is zero.
+    Value *SelectCond, *SelectTrue, *SelectFalse;
+    if (!match(HiVal, m_Select(m_Value(SelectCond), m_Value(SelectTrue),
+                               m_Value(SelectFalse))))
+      continue;
+
+    // SelectCond: icmp eq SrcVal, 0
+    if (!match(SelectCond, m_SpecificICmp(CmpInst::ICMP_EQ, m_Specific(SrcVal),
+                                          m_ZeroInt())))
+      continue;
+
+    // SelectTrue: the zero-input result
+    const APInt *ZeroResult;
+    if (!match(SelectTrue, m_APInt(ZeroResult)))
+      continue;
+
+    // SelectFalse: or(cttz.i32(trunc(lshr(SrcVal, 32)), _), 32)
+    // Also accept add instead of or.
+    Value *CttzHiCall;
+    if (!match(SelectFalse, m_c_Or(m_Value(CttzHiCall), m_SpecificInt(32))) &&
+        !match(SelectFalse, m_c_Add(m_Value(CttzHiCall), m_SpecificInt(32))))
+      continue;
+
+    // CttzHiCall: cttz.i32(trunc(lshr(SrcVal, 32)), _)
+    Value *HiTrunc;
+    if (!match(CttzHiCall,
+               m_Intrinsic<Intrinsic::cttz>(m_Value(HiTrunc), m_Value())))
+      continue;
+
+    if (!match(HiTrunc, m_Trunc(m_LShr(m_Specific(SrcVal), m_SpecificInt(32)))))
+      continue;
+
+    // Both HiBB and LoBB share a single predecessor
+    // that branches on whether the low 32 bits are zero.
+    BasicBlock *GuardBB = LoBB->getSinglePredecessor();
+    if (!GuardBB || GuardBB != HiBB->getSinglePredecessor())
+      continue;
+
+    Instruction *GuardTerm = GuardBB->getTerminator();
+    // Guard: br (icmp eq (trunc SrcVal to i32), 0), HiBB, LoBB
+    if (!match(GuardTerm,
+               m_Br(m_SpecificICmp(CmpInst::ICMP_EQ,
+                                   m_Trunc(m_Specific(SrcVal)), m_ZeroInt()),
+                    m_SpecificBB(HiBB), m_SpecificBB(LoBB))))
+      continue;
+
+    // Verify SrcVal dominates the guard.
----------------
topperc wrote:

Why wouldn't it dominate? We're in SSA form.

https://github.com/llvm/llvm-project/pull/192296