Take advantage of assumes 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.