[PredicateInfo] Infer operand bound from mul nuw square predicates

A mul nuw X, X used in an assume/branch condition cannot overflow (or
the condition would be poison, which is UB for assumes and control
flow), which implies:

  X < 2^ceil(bitwidth(X)/2) (e.g., i16: X < 256).

Fixes #122412

Author: ken-matsui

llvm/lib/Transforms/Utils/PredicateInfo.cpp

@@ -345,6 +345,14 @@ void collectCmpOps(CmpInst *Comparison, SmallVectorImpl<Value *> &CmpOperands) {
 
   CmpOperands.push_back(Op0);
   CmpOperands.push_back(Op1);
+
+  auto AddSquareOp = [&CmpOperands](Value *V) {
+    Value *SquareOp;
+    if (match(V, m_NUWMul(m_Value(SquareOp), m_Deferred(SquareOp))))
+      CmpOperands.push_back(SquareOp);
+  };
+  AddSquareOp(Op0);
+  AddSquareOp(Op1);
 }
 
 // Add Op, PB to the list of value infos for Op, and mark Op to be renamed.
@@ -735,6 +743,25 @@ std::optional<PredicateConstraint> PredicateBase::getConstraint() const {
       return std::nullopt;
     }
 
+    // A `mul nuw RenamedOp, RenamedOp` used by the condition cannot overflow
+    // (or the condition would be poison, which is UB for assumes and control
+    // flow), which implies `RenamedOp < 2^ceil(bitwidth(RenamedOp)/2)`
+    // (e.g., i16: `RenamedOp < 256`).
+    auto IsNuwSquareOfRenamedOp = [this](Value *V) {
+      return match(V, m_NUWMul(m_Specific(RenamedOp), m_Specific(RenamedOp)));
+    };
+    if (RenamedOp->getType()->isIntegerTy() &&
+        (IsNuwSquareOfRenamedOp(Cmp->getOperand(0)) ||
+         IsNuwSquareOfRenamedOp(Cmp->getOperand(1)))) {
+      unsigned BitWidth = RenamedOp->getType()->getScalarSizeInBits();
+      unsigned LimitBits = (BitWidth + 1) / 2;
+      if (LimitBits < BitWidth) {
+        APInt Upper = APInt::getOneBitSet(BitWidth, LimitBits);
+        return {
+            {CmpInst::ICMP_ULT, ConstantInt::get(RenamedOp->getType(), Upper)}};
+      }
+    }
+
     CmpInst::Predicate Pred;
     Value *OtherOp;
     if (Cmp->getOperand(0) == RenamedOp) {

llvm/test/Transforms/SCCP/assume-mul-nuw-square.ll

@@ -9,8 +9,7 @@ define i1 @src(i16 %num, i16 %s) {
 ; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i16 [[S:%.*]], [[S]]
 ; CHECK-NEXT:    [[COND:%.*]] = icmp ule i16 [[MUL]], [[NUM:%.*]]
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[COND]])
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult i16 [[S]], 256
-; CHECK-NEXT:    ret i1 [[CMP]]
+; CHECK-NEXT:    ret i1 true
 ;
 start:
   %mul = mul nuw i16 %s, %s