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

[Craig Topper via llvm-commits]

================
@@ -68,6 +72,179 @@ static cl::opt<unsigned>
                           cl::desc("The maximum length of a constant string to "
                                    "inline a memchr call."));
 
+/// Try to fold a select-based split cttz pattern into a single full-width cttz.
+///
+///   %lo = trunc iN %val to i(N/2)
+///   %cmp = icmp eq i(N/2) %lo, 0
+///   %shr = lshr iN %val, N/2
+///   %hi = trunc iN %shr to i(N/2)
+///   %cttz_hi = call i(N/2) @llvm.cttz.i(N/2)(i(N/2) %hi, ...)
+///   %hi_plus = add/or_disjoint i(N/2) %cttz_hi, N/2
+///   %cttz_lo = call i(N/2) @llvm.cttz.i(N/2)(i(N/2) %lo, ...)
+///   %result = select i1 %cmp, i(N/2) %hi_plus, i(N/2) %cttz_lo
+/// -->
+///   %cttz_wide = call iN @llvm.cttz.iN(iN %val, i1 false)
+///   %result = trunc iN %cttz_wide to i(N/2)
+/// Alive proof (for i64/i32):  https://alive2.llvm.org/ce/z/-s14-s
+static bool foldSelectSplitCTTZ(Instruction &I) {
+  Value *Cond, *TrueVal, *FalseVal;
+  if (!match(&I, m_Select(m_Value(Cond), m_Value(TrueVal), m_Value(FalseVal))))
+    return false;
+
+  Type *HalfTy = I.getType();
+  if (!HalfTy->isIntegerTy())
+    return false;
+  unsigned HalfWidth = HalfTy->getIntegerBitWidth();
+  unsigned FullWidth = HalfWidth * 2;
+
+  // select (icmp eq (trunc SrcVal to i(N/2)), 0), HiResult, LoResult
+  // Or select (icmp ne ...), LoResult, HiResult
+  Value *LoTrunc;
+  Value *HiResult, *LoResult;
+  if (match(Cond,
+            m_SpecificICmp(CmpInst::ICMP_EQ, m_Value(LoTrunc), m_ZeroInt()))) {
+    HiResult = TrueVal;
+    LoResult = FalseVal;
+  } else if (match(Cond, m_SpecificICmp(CmpInst::ICMP_NE, m_Value(LoTrunc),
+                                        m_ZeroInt()))) {
+    HiResult = FalseVal;
+    LoResult = TrueVal;
+  } else {
+    return false;
+  }
+
+  // LoTrunc: trunc iN SrcVal to i(N/2)
+  Value *SrcVal;
+  if (!match(LoTrunc, m_Trunc(m_Value(SrcVal))))
+    return false;
+  if (!SrcVal->getType()->isIntegerTy(FullWidth))
+    return false;
+
+  // LoResult: cttz(trunc(SrcVal), _),  must use same truncated value
+  Value *LoCttzArg;
+  if (!match(LoResult, m_OneUse(m_Intrinsic<Intrinsic::cttz>(m_Value(LoCttzArg),
+                                                             m_Value()))))
+    return false;
+  if (LoCttzArg != LoTrunc)
+    return false;
+
+  // HiResult: add/or_disjoint(cttz(trunc(lshr(SrcVal, N/2)), _), N/2)
+  Value *CttzHiCall;
+  if (!match(HiResult, m_OneUse(m_AddLike(m_Value(CttzHiCall),
+                                          m_SpecificInt(HalfWidth)))))
+    return false;
+
+  Value *HiCttzArg;
+  if (!match(CttzHiCall, m_OneUse(m_Intrinsic<Intrinsic::cttz>(
+                             m_Value(HiCttzArg), m_Value()))))
+    return false;
+
+  if (!match(HiCttzArg,
+             m_Trunc(m_LShr(m_Specific(SrcVal), m_SpecificInt(HalfWidth)))))
+    return false;
+
+  // Match successful.
+  IRBuilder<> Builder(&I);
+  Value *CttzWide = Builder.CreateIntrinsic(
+      Intrinsic::cttz, {SrcVal->getType()}, {SrcVal, Builder.getFalse()});
+  Value *Trunc = Builder.CreateTrunc(CttzWide, HalfTy);
----------------
topperc wrote:

> Would it make sense to add a guard like if (HalfWidth < 2) (or < 4) to exclude these types?

I think so.

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