[llvm] [InstCombine] Fold square into comparison with constant (reopenend) (PR #197665)
Macsen Casaus via llvm-commits
llvm-commits at lists.llvm.org
Thu May 14 05:17:38 PDT 2026
https://github.com/macsencasaus created https://github.com/llvm/llvm-project/pull/197665
Closes #196233
https://alive2.llvm.org/ce/z/PQzYhb
>From b228b05c2c4040989f32991c086449bf75075f48 Mon Sep 17 00:00:00 2001
From: Macsen Casaus <macsencasaus at gmail.com>
Date: Fri, 8 May 2026 07:04:22 +0000
Subject: [PATCH 1/5] pre-commit test
---
llvm/test/Transforms/InstCombine/icmp-mul.ll | 48 ++++++++++++++++++++
1 file changed, 48 insertions(+)
diff --git a/llvm/test/Transforms/InstCombine/icmp-mul.ll b/llvm/test/Transforms/InstCombine/icmp-mul.ll
index 1e4876d5cd569..d484eb4cf9923 100644
--- a/llvm/test/Transforms/InstCombine/icmp-mul.ll
+++ b/llvm/test/Transforms/InstCombine/icmp-mul.ll
@@ -91,6 +91,54 @@ define i1 @squared_nsw_sgt0(i5 %x) {
ret i1 %r
}
+define i1 @squared_nuw_ult_sqr(i8 %x) {
+; CHECK-LABEL: @squared_nuw_ult_sqr(
+; CHECK-NEXT: [[M:%.*]] = mul nuw i8 [[X:%.*]], [[X]]
+; CHECK-NEXT: [[R:%.*]] = icmp ult i8 [[M]], 9
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul nuw i8 %x, %x
+ %r = icmp ult i8 %m, 9
+ ret i1 %r
+}
+
+define i1 @squared_nuw_eq_sqr(i8 %x) {
+; CHECK-LABEL: @squared_nuw_eq_sqr(
+; CHECK-NEXT: [[M:%.*]] = mul nuw i8 [[X:%.*]], [[X]]
+; CHECK-NEXT: [[R:%.*]] = icmp eq i8 [[M]], 9
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul nuw i8 %x, %x
+ %r = icmp eq i8 %m, 9
+ ret i1 %r
+}
+
+; negative test - signed compare
+
+define i1 @squared_nuw_slt_sqr(i8 %x) {
+; CHECK-LABEL: @squared_nuw_slt_sqr(
+; CHECK-NEXT: [[M:%.*]] = mul nuw i8 [[X:%.*]], [[X]]
+; CHECK-NEXT: [[R:%.*]] = icmp slt i8 [[M]], 9
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul nuw i8 %x, %x
+ %r = icmp slt i8 %m, 9
+ ret i1 %r
+}
+
+; negative test - no nuw
+
+define i1 @squared_ult_sqr(i8 %x) {
+; CHECK-LABEL: @squared_ult_sqr(
+; CHECK-NEXT: [[M:%.*]] = mul i8 [[X:%.*]], [[X]]
+; CHECK-NEXT: [[R:%.*]] = icmp ult i8 [[M]], 9
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul i8 %x, %x
+ %r = icmp ult i8 %m, 9
+ ret i1 %r
+}
+
; Tests for slt/ult
define i1 @slt_positive_multip_rem_zero(i8 %x) {
>From f26a6f25c41f40007734c73cc4f24df1612d1ff4 Mon Sep 17 00:00:00 2001
From: Macsen Casaus <macsencasaus at gmail.com>
Date: Fri, 8 May 2026 07:12:38 +0000
Subject: [PATCH 2/5] [InstCombine] Fold square into comparison with constant
---
llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp | 9 +++++++++
llvm/test/Transforms/InstCombine/icmp-mul.ll | 6 ++----
2 files changed, 11 insertions(+), 4 deletions(-)
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index 0b1f683e10408..4426ea346949f 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -2195,6 +2195,15 @@ Instruction *InstCombinerImpl::foldICmpMulConstant(ICmpInst &Cmp,
(Mul->hasNoUnsignedWrap() || Mul->hasNoSignedWrap()))
return new ICmpInst(Pred, X, ConstantInt::getNullValue(MulTy));
+ // If unsigned compare or equality comparison of self multiply and square,
+ // compare square roots
+ if (!Cmp.isSigned() && X == Mul->getOperand(1) && Mul->hasNoUnsignedWrap() &&
+ !C.isNegative()) {
+ APInt R = C.sqrt();
+ if (C == R * R)
+ return new ICmpInst(Pred, X, ConstantInt::get(MulTy, R));
+ }
+
const APInt *MulC;
if (!match(Mul->getOperand(1), m_APInt(MulC)))
return nullptr;
diff --git a/llvm/test/Transforms/InstCombine/icmp-mul.ll b/llvm/test/Transforms/InstCombine/icmp-mul.ll
index d484eb4cf9923..40259cbe6977f 100644
--- a/llvm/test/Transforms/InstCombine/icmp-mul.ll
+++ b/llvm/test/Transforms/InstCombine/icmp-mul.ll
@@ -93,8 +93,7 @@ define i1 @squared_nsw_sgt0(i5 %x) {
define i1 @squared_nuw_ult_sqr(i8 %x) {
; CHECK-LABEL: @squared_nuw_ult_sqr(
-; CHECK-NEXT: [[M:%.*]] = mul nuw i8 [[X:%.*]], [[X]]
-; CHECK-NEXT: [[R:%.*]] = icmp ult i8 [[M]], 9
+; CHECK-NEXT: [[R:%.*]] = icmp ult i8 [[X:%.*]], 3
; CHECK-NEXT: ret i1 [[R]]
;
%m = mul nuw i8 %x, %x
@@ -104,8 +103,7 @@ define i1 @squared_nuw_ult_sqr(i8 %x) {
define i1 @squared_nuw_eq_sqr(i8 %x) {
; CHECK-LABEL: @squared_nuw_eq_sqr(
-; CHECK-NEXT: [[M:%.*]] = mul nuw i8 [[X:%.*]], [[X]]
-; CHECK-NEXT: [[R:%.*]] = icmp eq i8 [[M]], 9
+; CHECK-NEXT: [[R:%.*]] = icmp eq i8 [[X:%.*]], 3
; CHECK-NEXT: ret i1 [[R]]
;
%m = mul nuw i8 %x, %x
>From c2c34c2df412c1cc99f34abd6ad9d74b530df822 Mon Sep 17 00:00:00 2001
From: Macsen Casaus <macsencasaus at gmail.com>
Date: Fri, 8 May 2026 16:17:55 +0000
Subject: [PATCH 3/5] Remove negativity check
---
.../lib/Transforms/InstCombine/InstCombineCompares.cpp | 3 +--
llvm/test/Transforms/InstCombine/icmp-mul.ll | 10 ++++++++++
2 files changed, 11 insertions(+), 2 deletions(-)
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index 4426ea346949f..2a7d15ed4d274 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -2197,8 +2197,7 @@ Instruction *InstCombinerImpl::foldICmpMulConstant(ICmpInst &Cmp,
// If unsigned compare or equality comparison of self multiply and square,
// compare square roots
- if (!Cmp.isSigned() && X == Mul->getOperand(1) && Mul->hasNoUnsignedWrap() &&
- !C.isNegative()) {
+ if (!Cmp.isSigned() && X == Mul->getOperand(1) && Mul->hasNoUnsignedWrap()) {
APInt R = C.sqrt();
if (C == R * R)
return new ICmpInst(Pred, X, ConstantInt::get(MulTy, R));
diff --git a/llvm/test/Transforms/InstCombine/icmp-mul.ll b/llvm/test/Transforms/InstCombine/icmp-mul.ll
index 40259cbe6977f..ba2fe44243f71 100644
--- a/llvm/test/Transforms/InstCombine/icmp-mul.ll
+++ b/llvm/test/Transforms/InstCombine/icmp-mul.ll
@@ -111,6 +111,16 @@ define i1 @squared_nuw_eq_sqr(i8 %x) {
ret i1 %r
}
+define i1 @squared_nuw_eq_negative_sqr(i8 %x) {
+; CHECK-LABEL: @squared_nuw_eq_negative_sqr(
+; CHECK-NEXT: [[R:%.*]] = icmp eq i8 [[X:%.*]], 15
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul nuw i8 %x, %x
+ %r = icmp eq i8 %m, -31
+ ret i1 %r
+}
+
; negative test - signed compare
define i1 @squared_nuw_slt_sqr(i8 %x) {
>From ae0af6fa02bb3bea863dfd9f8d23fe5db67f0f22 Mon Sep 17 00:00:00 2001
From: Macsen Casaus <macsencasaus at gmail.com>
Date: Sat, 9 May 2026 02:52:59 +0000
Subject: [PATCH 4/5] Add equality always false case and sqrt floor/ceil case
---
.../InstCombine/InstCombineCompares.cpp | 55 ++++++++++++---
llvm/test/Transforms/InstCombine/icmp-mul.ll | 69 +++++++++++++++++++
2 files changed, 113 insertions(+), 11 deletions(-)
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index 2a7d15ed4d274..60b7ebfad86b4 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -2188,19 +2188,52 @@ Instruction *InstCombinerImpl::foldICmpMulConstant(ICmpInst &Cmp,
Type *MulTy = Mul->getType();
Value *X = Mul->getOperand(0);
- // If there's no overflow:
- // X * X == 0 --> X == 0
- // X * X != 0 --> X != 0
- if (Cmp.isEquality() && C.isZero() && X == Mul->getOperand(1) &&
- (Mul->hasNoUnsignedWrap() || Mul->hasNoSignedWrap()))
- return new ICmpInst(Pred, X, ConstantInt::getNullValue(MulTy));
-
- // If unsigned compare or equality comparison of self multiply and square,
- // compare square roots
- if (!Cmp.isSigned() && X == Mul->getOperand(1) && Mul->hasNoUnsignedWrap()) {
+ // If comparing a square with a constant, try simplifying to comparing square
+ // roots
+ if (X == Mul->getOperand(1)) {
APInt R = C.sqrt();
- if (C == R * R)
+ bool IsSqr = C == R * R;
+
+ // X * X eq/ne C
+ if (Cmp.isEquality() &&
+ (Mul->hasNoUnsignedWrap() || (Mul->hasNoSignedWrap() && C.isZero()))) {
+
+ // If constant is not a square, eq/ne is false/true respectively
+ if (!IsSqr)
+ return replaceInstUsesWith(
+ Cmp, ConstantInt::getBool(Cmp.getType(),
+ CmpInst::isFalseWhenEqual(Pred)));
+
return new ICmpInst(Pred, X, ConstantInt::get(MulTy, R));
+ }
+
+ // If the multiply does not wrap
+ // X * X pred C --> X pred R
+ if (Cmp.isUnsigned() && Mul->hasNoUnsignedWrap()) {
+
+ if (IsSqr)
+ return new ICmpInst(Pred, X, ConstantInt::get(MulTy, R));
+
+ // If C is not a square, we use floor/ceil of sqrt(C)
+ //
+ // If LT or LE, we need R to be an overestimate of sqrt(C),
+ // then use the strict predicate (LT->LT, LE->LT).
+ //
+ // If GT or GE, we need R to be an underestimate of sqrt(C),
+ // then use the strict predicate (GT->GT, GE->GT).
+ bool NeedOverestimate = ICmpInst::isLT(Pred) || ICmpInst::isLE(Pred);
+
+ // Is R currently an underestimate or overestimate of sqrt(C)?
+ bool Overestimate = C.ult(R * R);
+
+ if (Overestimate && !NeedOverestimate)
+ --R;
+ if (!Overestimate && NeedOverestimate)
+ ++R;
+
+ return new ICmpInst(Cmp.getStrictPredicate(), X,
+ ConstantInt::get(MulTy, R));
+ }
}
const APInt *MulC;
diff --git a/llvm/test/Transforms/InstCombine/icmp-mul.ll b/llvm/test/Transforms/InstCombine/icmp-mul.ll
index ba2fe44243f71..e4fc0a6f5264a 100644
--- a/llvm/test/Transforms/InstCombine/icmp-mul.ll
+++ b/llvm/test/Transforms/InstCombine/icmp-mul.ll
@@ -121,6 +121,15 @@ define i1 @squared_nuw_eq_negative_sqr(i8 %x) {
ret i1 %r
}
+define i1 @squared_eq_non_sqr(i8 %x) {
+; CHECK-LABEL: @squared_eq_non_sqr(
+; CHECK-NEXT: ret i1 false
+;
+ %m = mul nuw i8 %x, %x
+ %r = icmp eq i8 %m, 10
+ ret i1 %r
+}
+
; negative test - signed compare
define i1 @squared_nuw_slt_sqr(i8 %x) {
@@ -147,6 +156,66 @@ define i1 @squared_ult_sqr(i8 %x) {
ret i1 %r
}
+define i1 @squared_ult_est_sqr(i8 %x) {
+; CHECK-LABEL: @squared_ult_est_sqr(
+; CHECK-NEXT: [[R:%.*]] = icmp ult i8 [[X:%.*]], 4
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul nuw i8 %x, %x
+ %r = icmp ult i8 %m, 10
+ ret i1 %r
+}
+
+define i1 @squared_ule_est_sqr(i8 %x) {
+; CHECK-LABEL: @squared_ule_est_sqr(
+; CHECK-NEXT: [[R:%.*]] = icmp ult i8 [[X:%.*]], 4
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul nuw i8 %x, %x
+ %r = icmp ule i8 %m, 10
+ ret i1 %r
+}
+
+define i1 @squared_ule_est_sqr2(i8 %x) {
+; CHECK-LABEL: @squared_ule_est_sqr2(
+; CHECK-NEXT: [[R:%.*]] = icmp ult i8 [[X:%.*]], 4
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul nuw i8 %x, %x
+ %r = icmp ule i8 %m, 15
+ ret i1 %r
+}
+
+define i1 @squared_ugt_est_sqr(i8 %x) {
+; CHECK-LABEL: @squared_ugt_est_sqr(
+; CHECK-NEXT: [[R:%.*]] = icmp ugt i8 [[X:%.*]], 3
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul nuw i8 %x, %x
+ %r = icmp ugt i8 %m, 10
+ ret i1 %r
+}
+
+define i1 @squared_uge_est_sqr(i8 %x) {
+; CHECK-LABEL: @squared_uge_est_sqr(
+; CHECK-NEXT: [[R:%.*]] = icmp ugt i8 [[X:%.*]], 3
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul nuw i8 %x, %x
+ %r = icmp uge i8 %m, 10
+ ret i1 %r
+}
+
+define i1 @squared_uge_est_sqr2(i8 %x) {
+; CHECK-LABEL: @squared_uge_est_sqr2(
+; CHECK-NEXT: [[R:%.*]] = icmp ugt i8 [[X:%.*]], 3
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul nuw i8 %x, %x
+ %r = icmp uge i8 %m, 15
+ ret i1 %r
+}
+
; Tests for slt/ult
define i1 @slt_positive_multip_rem_zero(i8 %x) {
>From 48dd8e396d8ad94f6fabff02b1279e5c6f2f93cc Mon Sep 17 00:00:00 2001
From: Macsen Casaus <macsencasaus at gmail.com>
Date: Tue, 12 May 2026 21:41:37 +0000
Subject: [PATCH 5/5] Ignore R * R overflow case & add vector test case
---
.../InstCombine/InstCombineCompares.cpp | 71 ++++++++++---------
llvm/test/Transforms/InstCombine/icmp-mul.ll | 45 ++++++++++++
2 files changed, 83 insertions(+), 33 deletions(-)
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
index 60b7ebfad86b4..8affa706c18e2 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp
@@ -2190,49 +2190,54 @@ Instruction *InstCombinerImpl::foldICmpMulConstant(ICmpInst &Cmp,
// If comparing a square with a constant, try simplifying to comparing square
// roots
- if (X == Mul->getOperand(1)) {
+ if (X == Mul->getOperand(1) && !Cmp.isSigned()) {
APInt R = C.sqrt();
- bool IsSqr = C == R * R;
+ bool Overflow;
+ APInt R2 = R.umul_ov(R, Overflow);
- // X * X eq/ne C
- if (Cmp.isEquality() &&
- (Mul->hasNoUnsignedWrap() || (Mul->hasNoSignedWrap() && C.isZero()))) {
+ if (!Overflow) {
+ bool IsSqr = C == R2;
- // If constant is not a square, eq/ne is false/true respectively
- if (!IsSqr)
- return replaceInstUsesWith(
- Cmp, ConstantInt::getBool(Cmp.getType(),
- CmpInst::isFalseWhenEqual(Pred)));
+ // X * X eq/ne C
+ if (Cmp.isEquality() && (Mul->hasNoUnsignedWrap() ||
+ (Mul->hasNoSignedWrap() && C.isZero()))) {
- return new ICmpInst(Pred, X, ConstantInt::get(MulTy, R));
- }
+ // If constant is not a square, eq/ne is false/true respectively
+ if (!IsSqr)
+ return replaceInstUsesWith(
+ Cmp, ConstantInt::getBool(Cmp.getType(),
+ CmpInst::isFalseWhenEqual(Pred)));
- // If the multiply does not wrap
- // X * X pred C --> X pred R
- if (Cmp.isUnsigned() && Mul->hasNoUnsignedWrap()) {
-
- if (IsSqr)
return new ICmpInst(Pred, X, ConstantInt::get(MulTy, R));
+ }
- // If C is not a square, we use floor/ceil of sqrt(C)
- //
- // If LT or LE, we need R to be an overestimate of sqrt(C),
- // then use the strict predicate (LT->LT, LE->LT).
- //
- // If GT or GE, we need R to be an underestimate of sqrt(C),
- // then use the strict predicate (GT->GT, GE->GT).
- bool NeedOverestimate = ICmpInst::isLT(Pred) || ICmpInst::isLE(Pred);
+ // If the multiply does not wrap
+ // X * X pred C --> X pred R
+ if (Mul->hasNoUnsignedWrap()) {
+
+ if (IsSqr)
+ return new ICmpInst(Pred, X, ConstantInt::get(MulTy, R));
- // Is R currently an underestimate or overestimate of sqrt(C)?
- bool Overestimate = C.ult(R * R);
+ // If C is not a square, we use floor/ceil of sqrt(C)
+ //
+ // If LT or LE, we need R to be an overestimate of sqrt(C),
+ // then use the strict predicate (LT->LT, LE->LT).
+ //
+ // If GT or GE, we need R to be an underestimate of sqrt(C),
+ // then use the strict predicate (GT->GT, GE->GT).
+ bool NeedOverestimate = ICmpInst::isLT(Pred) || ICmpInst::isLE(Pred);
- if (Overestimate && !NeedOverestimate)
- --R;
- if (!Overestimate && NeedOverestimate)
- ++R;
+ // Is R currently an underestimate or overestimate of sqrt(C)?
+ bool Overestimate = C.ult(R2);
- return new ICmpInst(Cmp.getStrictPredicate(), X,
- ConstantInt::get(MulTy, R));
+ if (Overestimate && !NeedOverestimate)
+ --R;
+ if (!Overestimate && NeedOverestimate)
+ ++R;
+
+ return new ICmpInst(Cmp.getStrictPredicate(), X,
+ ConstantInt::get(MulTy, R));
+ }
}
}
diff --git a/llvm/test/Transforms/InstCombine/icmp-mul.ll b/llvm/test/Transforms/InstCombine/icmp-mul.ll
index e4fc0a6f5264a..14b6fda3e5ac5 100644
--- a/llvm/test/Transforms/InstCombine/icmp-mul.ll
+++ b/llvm/test/Transforms/InstCombine/icmp-mul.ll
@@ -130,6 +130,16 @@ define i1 @squared_eq_non_sqr(i8 %x) {
ret i1 %r
}
+define <2 x i1> @squared_nuw_sqr_v(<2 x i8> %x) {
+; CHECK-LABEL: @squared_nuw_sqr_v(
+; CHECK-NEXT: [[R:%.*]] = icmp ult <2 x i8> [[X:%.*]], splat (i8 3)
+; CHECK-NEXT: ret <2 x i1> [[R]]
+;
+ %m = mul nuw <2 x i8> %x, %x
+ %r = icmp ult <2 x i8> %m, <i8 9, i8 9>
+ ret <2 x i1> %r
+}
+
; negative test - signed compare
define i1 @squared_nuw_slt_sqr(i8 %x) {
@@ -143,6 +153,41 @@ define i1 @squared_nuw_slt_sqr(i8 %x) {
ret i1 %r
}
+; negative test - close to unsigned integer type max value
+
+define i1 @squared_nuw_ult_sqr_u32_max(i32 %x) {
+; CHECK-LABEL: @squared_nuw_ult_sqr_u32_max(
+; CHECK-NEXT: [[M:%.*]] = mul nuw i32 [[X:%.*]], [[X]]
+; CHECK-NEXT: [[R:%.*]] = icmp ne i32 [[M]], -1
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul nuw i32 %x, %x
+ %r = icmp ult i32 %m, -1
+ ret i1 %r
+}
+
+define i1 @squared_nuw_ult_sqr_u64_max(i64 %x) {
+; CHECK-LABEL: @squared_nuw_ult_sqr_u64_max(
+; CHECK-NEXT: [[M:%.*]] = mul nuw i64 [[X:%.*]], [[X]]
+; CHECK-NEXT: [[R:%.*]] = icmp ult i64 [[M]], -2
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul nuw i64 %x, %x
+ %r = icmp ult i64 %m, -2
+ ret i1 %r
+}
+
+define i1 @squared_nuw_ult_sqr_u128_max(i128 %x) {
+; CHECK-LABEL: @squared_nuw_ult_sqr_u128_max(
+; CHECK-NEXT: [[M:%.*]] = mul nuw i128 [[X:%.*]], [[X]]
+; CHECK-NEXT: [[R:%.*]] = icmp ult i128 [[M]], -3
+; CHECK-NEXT: ret i1 [[R]]
+;
+ %m = mul nuw i128 %x, %x
+ %r = icmp ult i128 %m, -3
+ ret i1 %r
+}
+
; negative test - no nuw
define i1 @squared_ult_sqr(i8 %x) {
More information about the llvm-commits
mailing list