[llvm] [InstCombine] Fold out-of-range bits for squaring signed integers (PR #153484)

[Nikita Popov via llvm-commits]

================
@@ -423,6 +423,48 @@ static void computeKnownBitsMul(const Value *Op0, const Value *Op1, bool NSW,
     Known.makeNonNegative();
   else if (isKnownNegative && !Known.isNonNegative())
     Known.makeNegative();
+
+  // Additional logic: If both operands are the same sign- or zero-extended
+  // value from a small integer, and the multiplication is (sext x) * (sext x)
+  // or (zext x) * (zext x), then the result cannot set bits above the maximum
+  // possible square. This allows InstCombine and other passes to fold (x * x) &
+  // (1 << N) to 0 when N is out of range.
+  const Value *A = nullptr;
+  // Only handle the case where both operands are the same extension of the same
+  // value.
+  if ((match(Op0, m_SExt(m_Value(A))) && match(Op1, m_SExt(m_Specific(A)))) ||
+      (match(Op0, m_ZExt(m_Value(A))) && match(Op1, m_ZExt(m_Specific(A))))) {
----------------
nikic wrote:

I don't think the zext handling here is useful. This will be handled by the generic code.

For the sext case, we know that the result is non-negative (due to self-multiply) and that we have a certain number of sign bits (due to multiply of sext), so together we know that the sign bits are actually zero bits.

I think the principled thing to do here would be, for self-multiplies, to call ComputeNumSignBits() and then set all those bits to zero.

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