Take advantage of `assume`s on `x * x` (with `nuw`)

[scottmcm]

Rust stabilized its isqrt method today, so I was looking at its codegen.

The output for

let root = u16::isqrt(num);

currently has a manual

assume(root <= 255)

I was thinking that it'd be nice to give LLVM a bit more specific information, like

assume(root * root <= num)

since that's (a partial) definition of what isqrt does, and strictly more specific.

Sadly, that assume today doesn't seem to be worth bothering adding, since opt seems like it doesn't take advantage of it even in a "simple" situation that doesn't need to know anything about num: https://llvm.godbolt.org/z/zoq4T5fKr

But Alive2 confirms that it would be allowed to https://alive2.llvm.org/ce/z/ko8HYR

define i1 @src(i16 noundef %num, i16 noundef %s) noundef zeroext {
start:
  %_4 = mul nuw i16 noundef %s, noundef %s
  %cond = icmp ule i16 %_4, noundef %num
  assume i1 %cond
  %_0 = icmp ult i16 noundef %s, 256
  ret i1 %_0
}
=>
define i1 @tgt(i16 noundef %num, i16 noundef %s) noundef zeroext {
start:
  ret i1 1
}
Transformation seems to be correct!

---

Original Rust code I used to get the LLVM IR: https://rust.godbolt.org/z/vM8qj4xjf

pub unsafe fn isqrt_facts(num: u16, s: u16) -> bool {
    std::hint::assert_unchecked(u16::unchecked_mul(s, s) <= num);
    s <= 255
}

cc rust-lang/rust#116226

[dtcxzyw]

We can perform this optimization in CVP/SCCP.

[llvmbot]

Hi!

This issue may be a good introductory issue for people new to working on LLVM. If you would like to work on this issue, your first steps are:

1. Check that no other contributor has already been assigned to this issue. If you believe that no one is actually working on it despite an assignment, ping the person. After one week without a response, the assignee may be changed.
2. In the comments of this issue, request for it to be assigned to you, or just create a pull request after following the steps below. Mention this issue in the description of the pull request.
3. Fix the issue locally.
4. Run the test suite locally. Remember that the subdirectories under test/ create fine-grained testing targets, so you can e.g. use make check-clang-ast to only run Clang's AST tests.
5. Create a Git commit.
6. Run git clang-format HEAD~1 to format your changes.
7. Open a pull request to the upstream repository on GitHub. Detailed instructions can be found in GitHub's documentation. Mention this issue in the description of the pull request.

If you have any further questions about this issue, don't hesitate to ask via a comment in the thread below.

[llvmbot]

@llvm/issue-subscribers-good-first-issue

Author: None (scottmcm)

Rust stabilized its [`isqrt`](https://doc.rust-lang.org/nightly/std/primitive.u16.html#method.isqrt) method [today](https://blog.rust-lang.org/2025/01/09/Rust-1.84.0.html#stabilized-apis), so I was looking at its codegen.

The output for

let root = u16::isqrt(num);

currently has a manual

assume(root <= 255)

I was thinking that it'd be nice to give LLVM a bit more specific information, like

assume(root * root <= num)

since that's (a partial) definition of what isqrt does, and strictly more specific.

Sadly, that assume today doesn't seem to be worth bothering adding, since opt seems like it doesn't take advantage of it even in a "simple" situation that doesn't need to know anything about num: <https://llvm.godbolt.org/z/zoq4T5fKr>

But Alive2 confirms that it would be allowed to <https://alive2.llvm.org/ce/z/ko8HYR>

define i1 @<!-- -->src(i16 noundef %num, i16 noundef %s) noundef zeroext {
start:
  %_4 = mul nuw i16 noundef %s, noundef %s
  %cond = icmp ule i16 %_4, noundef %num
  assume i1 %cond
  %_0 = icmp ult i16 noundef %s, 256
  ret i1 %_0
}
=&gt;
define i1 @<!-- -->tgt(i16 noundef %num, i16 noundef %s) noundef zeroext {
start:
  ret i1 1
}
Transformation seems to be correct!

---

Original Rust code I used to get the LLVM IR: <https://rust.godbolt.org/z/vM8qj4xjf>

pub unsafe fn isqrt_facts(num: u16, s: u16) -> bool {
    std::hint::assert_unchecked(u16::unchecked_mul(s, s) <= num);
    s <= 255
}

cc rust-lang/rust#116226

[ScaredLaptop]

@dtcxzyw If you could assign this to me, going to take a stab at it.

[el-ev]

@ScaredLaptop Are you still working on this?