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

[Yunbo Ni via llvm-commits]

================
@@ -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()) {
----------------
cardigan1008 wrote:

That signed-negativity check is too strong for equality and unsigned comparisons. For example, in i8, the constant `-31` is the unsigned value 225, and 225 is exactly `15*15`. Because `mul nuw i8 %x, %x` is poison for `%x > 15`, every defined input is in the monotone unsigned square domain `[0, 15]`. Therefore `icmp eq (x*x), 225` is equivalent to `icmp eq x, 15`, and `icmp ult (x*x), 225` is equivalent to `icmp ult x, 15` on the source-defined domain. Similarly, in i5, `-16` is unsigned 16, which is `4*4`, and `mul nuw` restricts defined `%x` to `[0,5]`.

Alive2 proof: https://alive2.llvm.org/ce/z/vydwL2

> Found with [Archer](https://github.com/cuhk-s3/Archer). Please let me know if anything is wrong.

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