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

[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))))) {
----------------
Aethezz wrote:

Ok thanks, i have added this into last commit. One question: currently my code uses match while other parts of this function use Op0 == Op1. Should we only handle the explicit self-multiply case (x * x), or also consider cases where both operands are sign-extensions of the same value?

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