[llvm] [ValueTracking] Infer operand bound from mul nuw square predicates (PR #173127)

Ken Matsui via llvm-commits llvm-commits at lists.llvm.org
Sun Jul 12 18:59:34 PDT 2026


https://github.com/ken-matsui updated https://github.com/llvm/llvm-project/pull/173127

>From 19737e89823f2347cbc161c0fedc80e6dff40961 Mon Sep 17 00:00:00 2001
From: Ken Matsui <github at kmts.me>
Date: Wed, 17 Dec 2025 19:21:16 -0500
Subject: [PATCH 1/2] Add baseline tests for upcoming patch

---
 .../mul-nuw-square.ll                         | 118 ++++++++++++++++++
 1 file changed, 118 insertions(+)
 create mode 100644 llvm/test/Transforms/CorrelatedValuePropagation/mul-nuw-square.ll

diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/mul-nuw-square.ll b/llvm/test/Transforms/CorrelatedValuePropagation/mul-nuw-square.ll
new file mode 100644
index 0000000000000..91abcbf012de8
--- /dev/null
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/mul-nuw-square.ll
@@ -0,0 +1,118 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -passes=correlated-propagation -S < %s | FileCheck %s
+
+declare void @llvm.assume(i1)
+
+; %mul ule 120 implies %s u< 11 (11 * 11 == 121 > 120).
+define i1 @assume_mul_nuw_square_i8(i8 %s) {
+; CHECK-LABEL: define i1 @assume_mul_nuw_square_i8(
+; CHECK-SAME: i8 [[S:%.*]]) {
+; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i8 [[S]], [[S]]
+; CHECK-NEXT:    [[COND:%.*]] = icmp ule i8 [[MUL]], 120
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[S]], 11
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %mul = mul nuw i8 %s, %s
+  %cond = icmp ule i8 %mul, 120
+  call void @llvm.assume(i1 %cond)
+  %cmp = icmp ult i8 %s, 11
+  ret i1 %cmp
+}
+
+; %s may be 10 (10 * 10 == 100 ule 120), so %s u< 10 must not fold.
+define i1 @assume_mul_nuw_square_i8_tight(i8 %s) {
+; CHECK-LABEL: define i1 @assume_mul_nuw_square_i8_tight(
+; CHECK-SAME: i8 [[S:%.*]]) {
+; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i8 [[S]], [[S]]
+; CHECK-NEXT:    [[COND:%.*]] = icmp ule i8 [[MUL]], 120
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[S]], 10
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %mul = mul nuw i8 %s, %s
+  %cond = icmp ule i8 %mul, 120
+  call void @llvm.assume(i1 %cond)
+  %cmp = icmp ult i8 %s, 10
+  ret i1 %cmp
+}
+
+; Without a constant bound, the square not overflowing implies %s u< 16 on
+; both edges.
+define i1 @branch_mul_nuw_square(i8 %s, i8 %num) {
+; CHECK-LABEL: define i1 @branch_mul_nuw_square(
+; CHECK-SAME: i8 [[S:%.*]], i8 [[NUM:%.*]]) {
+; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i8 [[S]], [[S]]
+; CHECK-NEXT:    [[COND:%.*]] = icmp ule i8 [[MUL]], [[NUM]]
+; CHECK-NEXT:    br i1 [[COND]], label %[[TRUE:.*]], label %[[FALSE:.*]]
+; CHECK:       [[TRUE]]:
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[S]], 16
+; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK:       [[FALSE]]:
+; CHECK-NEXT:    [[CMP2:%.*]] = icmp ult i8 [[S]], 16
+; CHECK-NEXT:    ret i1 [[CMP2]]
+;
+  %mul = mul nuw i8 %s, %s
+  %cond = icmp ule i8 %mul, %num
+  br i1 %cond, label %true, label %false
+
+true:
+  %cmp = icmp ult i8 %s, 16
+  ret i1 %cmp
+
+false:
+  %cmp2 = icmp ult i8 %s, 16
+  ret i1 %cmp2
+}
+
+; The square may be on the icmp RHS. Without a constant bound, the square not
+; overflowing i5 implies %s u< 6 (6 * 6 == 36 > 31).
+define i1 @assume_mul_nuw_square_i5_rhs(i5 %s, i5 %num) {
+; CHECK-LABEL: define i1 @assume_mul_nuw_square_i5_rhs(
+; CHECK-SAME: i5 [[S:%.*]], i5 [[NUM:%.*]]) {
+; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i5 [[S]], [[S]]
+; CHECK-NEXT:    [[COND:%.*]] = icmp uge i5 [[NUM]], [[MUL]]
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i5 [[S]], 6
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %mul = mul nuw i5 %s, %s
+  %cond = icmp uge i5 %num, %mul
+  call void @llvm.assume(i1 %cond)
+  %cmp = icmp ult i5 %s, 6
+  ret i1 %cmp
+}
+
+; negative test: missing nuw on the multiply.
+define i1 @assume_mul_square_no_nuw(i8 %s) {
+; CHECK-LABEL: define i1 @assume_mul_square_no_nuw(
+; CHECK-SAME: i8 [[S:%.*]]) {
+; CHECK-NEXT:    [[MUL:%.*]] = mul i8 [[S]], [[S]]
+; CHECK-NEXT:    [[COND:%.*]] = icmp ule i8 [[MUL]], 120
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[S]], 16
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %mul = mul i8 %s, %s
+  %cond = icmp ule i8 %mul, 120
+  call void @llvm.assume(i1 %cond)
+  %cmp = icmp ult i8 %s, 16
+  ret i1 %cmp
+}
+
+; negative test: multiply is not a square.
+define i1 @assume_mul_nuw_not_square(i8 %s, i8 %t) {
+; CHECK-LABEL: define i1 @assume_mul_nuw_not_square(
+; CHECK-SAME: i8 [[S:%.*]], i8 [[T:%.*]]) {
+; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i8 [[S]], [[T]]
+; CHECK-NEXT:    [[COND:%.*]] = icmp ule i8 [[MUL]], 120
+; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[S]], 16
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %mul = mul nuw i8 %s, %t
+  %cond = icmp ule i8 %mul, 120
+  call void @llvm.assume(i1 %cond)
+  %cmp = icmp ult i8 %s, 16
+  ret i1 %cmp
+}

>From dd8064d34d2aaf7257b1841b84bb804b89005e2b Mon Sep 17 00:00:00 2001
From: Ken Matsui <github at kmts.me>
Date: Sun, 21 Dec 2025 00:07:45 -0500
Subject: [PATCH 2/2] [LVI] Infer ranges from mul nuw square conditions

A non-poison comparison involving `mul nuw X, X` implies that the
multiplication does not overflow.  This bounds X by:

  X <= floor(sqrt(2^bitwidth(X) - 1)) (e.g., i16: X <= 255)

An unsigned constant comparison can tighten the bound, e.g.,
`X * X <= 120` implies `X <= 10`.
---
 llvm/include/llvm/IR/ConstantRange.h          |  3 ++
 llvm/lib/Analysis/LazyValueInfo.cpp           | 30 +++++++++++++++++++
 llvm/lib/Analysis/ValueTracking.cpp           |  8 +++++
 llvm/lib/IR/ConstantRange.cpp                 | 10 +++++++
 .../mul-nuw-square.ll                         | 14 ++++-----
 llvm/unittests/IR/ConstantRangeTest.cpp       |  5 ++++
 6 files changed, 61 insertions(+), 9 deletions(-)

diff --git a/llvm/include/llvm/IR/ConstantRange.h b/llvm/include/llvm/IR/ConstantRange.h
index 0d36fddac505f..e553aab2b93ee 100644
--- a/llvm/include/llvm/IR/ConstantRange.h
+++ b/llvm/include/llvm/IR/ConstantRange.h
@@ -575,6 +575,9 @@ class [[nodiscard]] ConstantRange {
   /// Calculate ctpop range.
   LLVM_ABI ConstantRange ctpop() const;
 
+  /// Calculate sqrtFloor range.  See APInt::sqrtFloor().
+  LLVM_ABI ConstantRange sqrtFloor() const;
+
   /// Represents whether an operation on the given constant range is known to
   /// always or never overflow.
   enum class OverflowResult {
diff --git a/llvm/lib/Analysis/LazyValueInfo.cpp b/llvm/lib/Analysis/LazyValueInfo.cpp
index 053144c42341f..c8f76ccd4b0e9 100644
--- a/llvm/lib/Analysis/LazyValueInfo.cpp
+++ b/llvm/lib/Analysis/LazyValueInfo.cpp
@@ -1355,6 +1355,33 @@ static ValueLatticeElement getValueFromICmpCtpop(ICmpInst::Predicate Pred,
       ConstantRange::getNonEmpty(std::move(ValMin), ValMax + 1));
 }
 
+/// Get the unsigned range for \p V from a `mul nuw V, V` comparison.
+static std::optional<ConstantRange>
+getRangeForNUWMulSquare(const Value *V, CmpInst::Predicate Pred,
+                        const Value *LHS, const Value *RHS) {
+  if (!V->getType()->isIntegerTy())
+    return std::nullopt;
+
+  if (!match(LHS, m_NUWMul(m_Specific(V), m_Specific(V)))) {
+    if (!match(RHS, m_NUWMul(m_Specific(V), m_Specific(V))))
+      return std::nullopt;
+
+    Pred = CmpInst::getSwappedPredicate(Pred);
+    RHS = LHS;
+  }
+
+  ConstantRange MulCR =
+      ConstantRange::getFull(V->getType()->getScalarSizeInBits());
+  const APInt *C;
+  if (match(RHS, m_APInt(C)))
+    MulCR = ConstantRange::makeExactICmpRegion(Pred, *C);
+
+  ConstantRange Res = MulCR.sqrtFloor();
+  if (Res.isFullSet())
+    return std::nullopt;
+  return Res;
+}
+
 std::optional<ValueLatticeElement> LazyValueInfoImpl::getValueFromICmpCondition(
     Value *Val, ICmpInst *ICI, bool isTrueDest, bool UseBlockValue) {
   Value *LHS = ICI->getOperand(0);
@@ -1378,6 +1405,9 @@ std::optional<ValueLatticeElement> LazyValueInfoImpl::getValueFromICmpCondition(
     return ValueLatticeElement::getOverdefined();
 
   unsigned BitWidth = Ty->getScalarSizeInBits();
+  if (auto Range = getRangeForNUWMulSquare(Val, EdgePred, LHS, RHS))
+    return ValueLatticeElement::getRange(*Range);
+
   APInt Offset(BitWidth, 0);
   if (matchICmpOperand(Offset, LHS, Val, EdgePred))
     return getValueFromSimpleICmpCondition(EdgePred, RHS, Offset, ICI,
diff --git a/llvm/lib/Analysis/ValueTracking.cpp b/llvm/lib/Analysis/ValueTracking.cpp
index efc14f9a639da..734f1199cb099 100644
--- a/llvm/lib/Analysis/ValueTracking.cpp
+++ b/llvm/lib/Analysis/ValueTracking.cpp
@@ -10736,6 +10736,14 @@ void llvm::findValuesAffectedByCondition(
         }
       }
 
+      auto AddNuwSquareOperand = [&AddAffected](Value *Op) {
+        Value *SquareOp = nullptr;
+        if (match(Op, m_NUWMul(m_Value(SquareOp), m_Deferred(SquareOp))))
+          AddAffected(SquareOp);
+      };
+      AddNuwSquareOperand(A);
+      AddNuwSquareOperand(B);
+
       if (HasRHSC && match(A, m_Ctpop(m_Value(X))))
         AddAffected(X);
     } else if (match(V, m_FCmp(Pred, m_Value(A), m_Value(B)))) {
diff --git a/llvm/lib/IR/ConstantRange.cpp b/llvm/lib/IR/ConstantRange.cpp
index 74b329b59bd73..ffe51cdcbe2ba 100644
--- a/llvm/lib/IR/ConstantRange.cpp
+++ b/llvm/lib/IR/ConstantRange.cpp
@@ -2173,6 +2173,16 @@ ConstantRange ConstantRange::ctpop() const {
   return CR1.unionWith(CR2);
 }
 
+ConstantRange ConstantRange::sqrtFloor() const {
+  if (isEmptySet())
+    return getEmpty();
+
+  // sqrtFloor is monotonic, so the output range is composed by the result of
+  // sqrtFloor of the two extremes.
+  return getNonEmpty(getUnsignedMin().sqrtFloor(),
+                     getUnsignedMax().sqrtFloor() + 1);
+}
+
 ConstantRange::OverflowResult ConstantRange::unsignedAddMayOverflow(
     const ConstantRange &Other) const {
   if (isEmptySet() || Other.isEmptySet())
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/mul-nuw-square.ll b/llvm/test/Transforms/CorrelatedValuePropagation/mul-nuw-square.ll
index 91abcbf012de8..415ba6bcf555d 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/mul-nuw-square.ll
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/mul-nuw-square.ll
@@ -10,8 +10,7 @@ define i1 @assume_mul_nuw_square_i8(i8 %s) {
 ; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i8 [[S]], [[S]]
 ; CHECK-NEXT:    [[COND:%.*]] = icmp ule i8 [[MUL]], 120
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[S]], 11
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 true
 ;
   %mul = mul nuw i8 %s, %s
   %cond = icmp ule i8 %mul, 120
@@ -27,7 +26,7 @@ define i1 @assume_mul_nuw_square_i8_tight(i8 %s) {
 ; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i8 [[S]], [[S]]
 ; CHECK-NEXT:    [[COND:%.*]] = icmp ule i8 [[MUL]], 120
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[S]], 10
+; CHECK-NEXT:    [[CMP:%.*]] = icmp samesign ult i8 [[S]], 10
 ; CHECK-NEXT:    ret i1 [[CMP]]
 ;
   %mul = mul nuw i8 %s, %s
@@ -46,11 +45,9 @@ define i1 @branch_mul_nuw_square(i8 %s, i8 %num) {
 ; CHECK-NEXT:    [[COND:%.*]] = icmp ule i8 [[MUL]], [[NUM]]
 ; CHECK-NEXT:    br i1 [[COND]], label %[[TRUE:.*]], label %[[FALSE:.*]]
 ; CHECK:       [[TRUE]]:
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i8 [[S]], 16
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 true
 ; CHECK:       [[FALSE]]:
-; CHECK-NEXT:    [[CMP2:%.*]] = icmp ult i8 [[S]], 16
-; CHECK-NEXT:    ret i1 [[CMP2]]
+; CHECK-NEXT:    ret i1 true
 ;
   %mul = mul nuw i8 %s, %s
   %cond = icmp ule i8 %mul, %num
@@ -73,8 +70,7 @@ define i1 @assume_mul_nuw_square_i5_rhs(i5 %s, i5 %num) {
 ; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i5 [[S]], [[S]]
 ; CHECK-NEXT:    [[COND:%.*]] = icmp uge i5 [[NUM]], [[MUL]]
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i5 [[S]], 6
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 true
 ;
   %mul = mul nuw i5 %s, %s
   %cond = icmp uge i5 %num, %mul
diff --git a/llvm/unittests/IR/ConstantRangeTest.cpp b/llvm/unittests/IR/ConstantRangeTest.cpp
index 872f220a34842..099a43bba34c4 100644
--- a/llvm/unittests/IR/ConstantRangeTest.cpp
+++ b/llvm/unittests/IR/ConstantRangeTest.cpp
@@ -2822,6 +2822,11 @@ TEST_F(ConstantRangeTest, Ctpop) {
       [](const APInt &N) { return APInt(N.getBitWidth(), N.popcount()); });
 }
 
+TEST_F(ConstantRangeTest, SqrtFloor) {
+  TestUnaryOpExhaustive([](const ConstantRange &CR) { return CR.sqrtFloor(); },
+                        [](const APInt &N) { return N.sqrtFloor(); });
+}
+
 TEST_F(ConstantRangeTest, castOps) {
   ConstantRange A(APInt(16, 66), APInt(16, 128));
   ConstantRange FpToI8 = A.castOp(Instruction::FPToSI, 8);



More information about the llvm-commits mailing list