[InstCombine] Fold out-of-range bits for squaring signed integers #153484

Aethezz · 2025-08-13T20:14:25Z

Fixes an issue where bits next to the sign bit were not constant-folded when squaring a sign- or zero-extended small integer. Added logic to detect when both operands of a multiplication are the same extended value, allowing InstCombine to mark bits above the maximum possible square as known zero. This enables correct folding of (x * x) & (1 << N) to 0 when N is out of range.

Proof: https://alive2.llvm.org/ce/z/YGou44

Fixes #152061

llvmbot · 2025-08-13T20:14:58Z

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-analysis

Author: None (Aethezz)

Changes

Fixes an issue where bits next to the sign bit were not constant-folded when squaring a sign- or zero-extended small integer. Added logic to detect when both operands of a multiplication are the same extended value, allowing InstCombine to mark bits above the maximum possible square as known zero. This enables correct folding of (x * x) & (1 << N) to 0 when N is out of range.

Fixes #152061

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

2 Files Affected:

(modified) llvm/lib/Analysis/ValueTracking.cpp (+43)
(modified) llvm/test/Analysis/ValueTracking/known-bits.ll (+33)

diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index af85ce4077ec8..7a973140f6075 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -423,6 +423,49 @@ 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.
+  using namespace PatternMatch;
+  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))))) {
+    Type *FromTy = A->getType();
+    Type *ToTy = Op0->getType();
+    if (FromTy->isIntegerTy() && ToTy->isIntegerTy() &&
+        FromTy->getScalarSizeInBits() < ToTy->getScalarSizeInBits()) {
+      unsigned FromBits = FromTy->getScalarSizeInBits();
+      unsigned ToBits = ToTy->getScalarSizeInBits();
+      // For both signed and unsigned, the maximum absolute value is max(|min|,
+      // |max|)
+      APInt minVal(FromBits, 0), maxVal(FromBits, 0);
+      bool isSigned = isa<SExtInst>(Op0);
+      if (isSigned) {
+        minVal = APInt::getSignedMinValue(FromBits);
+        maxVal = APInt::getSignedMaxValue(FromBits);
+      } else {
+        minVal = APInt::getMinValue(FromBits);
+        maxVal = APInt::getMaxValue(FromBits);
+      }
+      APInt absMin = minVal.abs();
+      APInt absMax = maxVal.abs();
+      APInt maxAbs = absMin.ugt(absMax) ? absMin : absMax;
+      APInt maxSquare = maxAbs.zext(ToBits);
+      maxSquare = maxSquare * maxSquare;
+      // All bits above the highest set bit in maxSquare are known zero.
+      unsigned MaxBit = maxSquare.isZero() ? 0 : maxSquare.logBase2();
+      if (MaxBit + 1 < ToBits) {
+        APInt KnownZeroMask =
+            APInt::getHighBitsSet(ToBits, ToBits - (MaxBit + 1));
+        Known.Zero |= KnownZeroMask;
+      }
+    }
+  }
 }
 
 void llvm::computeKnownBitsFromRangeMetadata(const MDNode &Ranges,
diff --git a/llvm/test/Analysis/ValueTracking/known-bits.ll b/llvm/test/Analysis/ValueTracking/known-bits.ll
index 5b71402a96f0d..d9f119bd0d146 100644
--- a/llvm/test/Analysis/ValueTracking/known-bits.ll
+++ b/llvm/test/Analysis/ValueTracking/known-bits.ll
@@ -49,3 +49,36 @@ define i1 @vec_reverse_known_bits_demanded_fail(<4 x i8> %xx) {
   %r = icmp slt i8 %ele, 0
   ret i1 %r
 }
+
+; Test known bits for (sext i8 x) * (sext i8 x)
+; RUN: opt -passes=instcombine < %s -S | FileCheck %s --check-prefix=SEXT_SQUARE
+
+define i1 @sext_square_bit31(i8 %x) {
+; SEXT_SQUARE-LABEL: @sext_square_bit31(
+; SEXT_SQUARE-NEXT:    ret i1 false
+  %sx = sext i8 %x to i32
+  %mul = mul nsw i32 %sx, %sx
+  %and = and i32 %mul, 2147483648 ; 1 << 31
+  %cmp = icmp ne i32 %and, 0
+  ret i1 %cmp
+}
+
+define i1 @sext_square_bit30(i8 %x) {
+; SEXT_SQUARE-LABEL: @sext_square_bit30(
+; SEXT_SQUARE-NEXT:    ret i1 false
+  %sx = sext i8 %x to i32
+  %mul = mul nsw i32 %sx, %sx
+  %and = and i32 %mul, 1073741824 ; 1 << 30
+  %cmp = icmp ne i32 %and, 0
+  ret i1 %cmp
+}
+
+define i1 @sext_square_bit14(i8 %x) {
+; SEXT_SQUARE-LABEL: @sext_square_bit14(
+; SEXT_SQUARE-NOT: ret i1 false
+  %sx = sext i8 %x to i32
+  %mul = mul nsw i32 %sx, %sx
+  %and = and i32 %mul, 16384 ; 1 << 14
+  %cmp = icmp ne i32 %and, 0
+  ret i1 %cmp
+}

nikic · 2025-08-17T14:49:45Z

llvm/lib/Analysis/ValueTracking.cpp

+  // 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))))) {


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.

Oh right, the zext is redundant. I’ve updated the code so that for self-multiplies using sext, we now call ComputeNumSignBits() to determine the number of sign bits and mark them as known zero.

Actually, after reviewing the previous commit, how should we call ComputeNumSignBits() and set the corresponding bits to zero? In this function, we only track known bits and don’t explicitly compute the product, so it’s unclear how to determine the exact number of sign bits.

I’ve made another commit that reverts to the previous approach using max/min value boundaries and removed the zext handling for now.

You can use the same logic as ComputeNumSignBits:

llvm-project/llvm/lib/Analysis/ValueTracking.cpp

Lines 4278 to 4280 in 92a91f7

unsigned OutValidBits =

(TyBits - SignBitsOp0 + 1) + (TyBits - SignBitsOp1 + 1);

return OutValidBits > TyBits ? 1 : TyBits - OutValidBits + 1;

Adjusted for the case where the sign bits are the same for both operands:

unsigned OutValidBits = 2 * (TyBits - SignBits + 1); unsigned OutSignBits = OutValidBits > TyBits ? 1 : TyBits - OutValidBits + 1;

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?

It's not necessary to handle sign extensions of the same value, as CSE will convert this into one sign extension used in both operands. So we should use Op0 == Op1.

Oh ok, i moved it into the selfmultiply handling instead which uses Op0 == Op1

github-actions · 2025-08-19T17:46:08Z

✅ With the latest revision this PR passed the C/C++ code formatter.

nikic · 2025-08-19T19:04:33Z

llvm/test/Analysis/ValueTracking/known-bits.ll

 }
+
+; Test known bits for (sext i8 x) * (sext i8 x)
+; RUN: opt -passes=instcombine < %s -S | FileCheck %s --check-prefix=SEXT_SQUARE


Do not add an extra run line to this test. If this does not fold through -passes=instsimplify, then this should be tested inside llvm/test/Transforms/InstCombine.

This doesn't fold with instsimplify so I moved it into llvm/test/Transforms/InstCombine

llvm/lib/Analysis/ValueTracking.cpp

nikic · 2025-08-19T19:25:06Z

~~Proof: https://alive2.llvm.org/ce/z/8qoQo-~~

Edit: Fixed proof: https://alive2.llvm.org/ce/z/YGou44

nikic · 2025-08-20T11:00:33Z

llvm/lib/Analysis/ValueTracking.cpp

This is the same for both cases. I think it would be cleaner if you just added an extra if (SelfMultiply) {} after the current KnownBits::mul() call.

After testing, I found that adding the extra if (SelfMultiply) {} causes the update_test_checks.py run to produce unfolded IR. I think it has to do with the previous if (SelfMultiply) where SelfMultiply &= isGuaranteedNotToBeUndef(Op0, Q.AC, Q.CxtI, Q.DT, Depth + 1); returns false. This causes our logic to not run since SelfMultiply is now false.

Edit: I printed out SelfMultiply and it indeed becomes false after that line, however the IR is folded with -O2 still

if (SelfMultiply) where SelfMultiply &= isGuaranteedNotToBeUndef(Op0, Q.AC, Q.CxtI, Q.DT, Depth + 1); returns false.

Can you add noundef to the parameter %x?

Thanks, adding noundef to %x resolved the issue.

llvm/test/Analysis/ValueTracking/known-bits.ll

llvm/test/Transforms/InstCombine/sext.ll

nikic · 2025-08-20T11:03:59Z

llvm/test/Transforms/InstCombine/sext.ll

It looks like this one is incorrect: https://alive2.llvm.org/ce/z/y8D-Yg

Though I don't really get why this case is folding. We should have TyBits = 32, SignBits = 9, OutValidBits = 16, OutSignBits = 17. So I would not expect bit 1<<14 to be set to zero.

Oh wait, I didn't notice that this one is actually using CHECK-NOT, not CHECK. So everything is correct.

…tion to drop ternary and moved test to llvm/test/Transforms/InstCombine since wasn't folded with instsimplifiy

Aethezz · 2025-09-03T13:23:55Z

@nikic

nikic

LGTM

Added additional logic to fold (x * x) masks for out-of-range bits

98e536c

Aethezz requested a review from nikic as a code owner August 13, 2025 20:14

llvmbot added the llvm:analysis Includes value tracking, cost tables and constant folding label Aug 13, 2025

Aethezz changed the title ~~[InstCombine][missed-optimizations] Fold out-of-range bits for squaring signed integers~~ [InstCombine] Fold out-of-range bits for squaring signed integers Aug 13, 2025

Aethezz added 2 commits August 13, 2025 16:51

Merge branch 'main' into Fix-llvm#152061

b3b264c

remove redundant namespace

f5f6d15

nikic reviewed Aug 17, 2025

View reviewed changes

nikic requested a review from dtcxzyw August 17, 2025 14:49

Aethezz added 3 commits August 18, 2025 17:30

use ComputeNumSignBits() instead and remove zext handling

246d414

remove previous comment

0c58e22

temporary revert back to previous changes and remove zext handling

0ff6997

Aethezz force-pushed the Fix-#152061 branch from 28fe831 to fd26986 Compare August 19, 2025 17:46

Added logic to compute max number of valid and sign bits and set to zero

72cd125

Aethezz force-pushed the Fix-#152061 branch from fd26986 to 72cd125 Compare August 19, 2025 18:01

nikic reviewed Aug 19, 2025

View reviewed changes

llvm/lib/Analysis/ValueTracking.cpp Outdated Show resolved Hide resolved

llvmbot added llvm:instcombine Covers the InstCombine, InstSimplify and AggressiveInstCombine passes llvm:transforms labels Aug 19, 2025

nikic mentioned this pull request Aug 20, 2025

Task submission dtcxzyw/llvm-opt-benchmark#1312

Open

zyw-bot mentioned this pull request Aug 20, 2025

pre-commit: PR153484 dtcxzyw/llvm-opt-benchmark#2695

Closed

nikic mentioned this pull request Aug 20, 2025

Fuzz PR153484 dtcxzyw/llvm-mutation-based-fuzz-service#83

Closed

nikic reviewed Aug 20, 2025

View reviewed changes

removed match and moved into selfmultiply condition, simplified condi…

514f267

…tion to drop ternary and moved test to llvm/test/Transforms/InstCombine since wasn't folded with instsimplifiy

Aethezz force-pushed the Fix-#152061 branch 2 times, most recently from c85b6a4 to 98c0200 Compare August 20, 2025 18:19

fixed formatting on testcase

54acb35

Aethezz force-pushed the Fix-#152061 branch from 98c0200 to 54acb35 Compare August 20, 2025 18:28

Aethezz requested a review from nikic August 26, 2025 20:27

made two if statements and added noundef to testcases

579b510

nikic approved these changes Sep 5, 2025

View reviewed changes

nikic merged commit ef1539c into llvm:main Sep 5, 2025
9 checks passed

	unsigned OutValidBits =
	(TyBits - SignBitsOp0 + 1) + (TyBits - SignBitsOp1 + 1);
	return OutValidBits > TyBits ? 1 : TyBits - OutValidBits + 1;

[InstCombine] Fold out-of-range bits for squaring signed integers #153484

[InstCombine] Fold out-of-range bits for squaring signed integers #153484

Uh oh!

Conversation

Aethezz commented Aug 13, 2025 • edited by nikic Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Aug 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Aug 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

nikic commented Aug 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Aethezz Aug 20, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Aethezz commented Sep 3, 2025

Uh oh!

nikic left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Aethezz commented Aug 13, 2025 •

edited by nikic

Loading

llvmbot commented Aug 13, 2025 •

edited

Loading

github-actions bot commented Aug 19, 2025 •

edited

Loading

nikic commented Aug 19, 2025 •

edited

Loading

Aethezz Aug 20, 2025 •

edited

Loading