[llvm] [InstCombine] Fold square into comparison with constant (PR #196480)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Macsen Casaus (macsencasaus)

<details>
<summary>Changes</summary>

Closes #<!-- -->196233 
https://alive2.llvm.org/ce/z/5kht6g

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


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineCompares.cpp (+9) 
- (modified) llvm/test/Transforms/InstCombine/icmp-mul.ll (+46) 


``````````diff
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 1e4876d5cd569..40259cbe6977f 100644
--- a/llvm/test/Transforms/InstCombine/icmp-mul.ll
+++ b/llvm/test/Transforms/InstCombine/icmp-mul.ll
@@ -91,6 +91,52 @@ 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:    [[R:%.*]] = icmp ult i8 [[X:%.*]], 3
+; 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:    [[R:%.*]] = icmp eq i8 [[X:%.*]], 3
+; 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) {

``````````

</details>


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