[llvm] [KnownBits] Improve mul when one operand is a power of two or zero (PR #227687)

[via llvm-commits]

================
@@ -1154,6 +1154,23 @@ KnownBits KnownBits::mul(const KnownBits &LHS, const KnownBits &RHS,
   Res.Zero |= (~BottomKnown).getLoBits(ResultBitsKnown);
   Res.One = BottomKnown.getLoBits(ResultBitsKnown);
 
+  // If all bits of one operand are known zero except bit K, that operand is
+  // either zero or 1 << K, so the product is either zero or the other operand
+  // shifted left by K. Any bit that is zero in the shifted operand is zero in
+  // the result, even if the known bits are not contiguous. If bit K is known
+  // one, the product is exactly the shifted operand.
+  auto AddBitsForPow2OrZero = [&](const KnownBits &Pow2OrZero,
+                                  const KnownBits &Other) {
+    if (Pow2OrZero.Zero.popcount() != BitWidth - 1)
+      return;
+    unsigned K = Pow2OrZero.Zero.countr_one();
+    Res.Zero |= Other.Zero.shl(K) | APInt::getLowBitsSet(BitWidth, K);
+    if (!Pow2OrZero.One.isZero())
+      Res.One |= Other.One.shl(K);
----------------
aswinkaliesrm wrote:

Good question. For a general constant the product is a sum of shifted copies of the other operand, so the carries make an optimal result harder. Adding the shifted copies with KnownBits::add would be correct but not optimal. I'd prefer to keep this PR to the zero/power-of-two case, and I can look into the general constant case separately.

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