[PredicateInfo] Infer operand bound from mul nuw square predicates

[ken-matsui]

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

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Ken Matsui (ken-matsui)

Changes

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

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Utils/PredicateInfo.cpp (+27)
• (added) llvm/test/Transforms/SCCP/assume-mul-nuw-square.ll (+20)

diff --git a/llvm/lib/Transforms/Utils/PredicateInfo.cpp b/llvm/lib/Transforms/Utils/PredicateInfo.cpp
index 27fed7340411b..c67d2c1e5cf5a 100644
--- a/llvm/lib/Transforms/Utils/PredicateInfo.cpp
+++ b/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) {
diff --git a/llvm/test/Transforms/SCCP/assume-mul-nuw-square.ll b/llvm/test/Transforms/SCCP/assume-mul-nuw-square.ll
new file mode 100644
index 0000000000000..caa72e2bb02c9
--- /dev/null
+++ b/llvm/test/Transforms/SCCP/assume-mul-nuw-square.ll
@@ -0,0 +1,20 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=sccp -S < %s | FileCheck %s
+
+declare void @llvm.assume(i1)
+
+define i1 @src(i16 %num, i16 %s) {
+; CHECK-LABEL: @src(
+; CHECK-NEXT:  start:
+; 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:    ret i1 true
+;
+start:
+  %mul = mul nuw i16 %s, %s
+  %cond = icmp ule i16 %mul, %num
+  call void @llvm.assume(i1 %cond)
+  %cmp = icmp ult i16 %s, 256
+  ret i1 %cmp
+}

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[nikic]

I'd rather handle this in ValueTracking or LazyValueInfo, especially for this case where the reasoning is in terms of the bit width only. We haven't infected PredicateInfo with this kind of special case handling yet, and I don't think this pattern is good motivation to start doing it.