[ValueTracking] `isKnownToBeAPowerOfTwo` ignores `llvm.assume` of a disjunction of equalities with power-of-two constants
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
Discovered this when investigating https://github.com/rust-lang/rust/issues/162513 with LLM's help.
`assume(X == 8 || X == 16 || X == 32 || X == 64)` states more than `assume(ctpop(X) == 1)`: a value set of four powers of two rather than any power of two. LLVM uses the latter to turn `udiv C, X` into `lshr C, cttz(X)` (since #115767), but the former contributes nothing and the division stays.
Any fact that holds under every disjunct of an assumed `or` holds unconditionally, so this is the dual of the existing handling of `assume(A and B)`.
## Reproduction (Rust)
```rust
#[derive(Clone, Copy)]
#[repr(u8)]
pub enum Width { W8 = 8, W16 = 16, W32 = 32, W64 = 64 }
#[no_mangle]
pub fn from_value_set(w: Width) -> u32 {
let w = w as u8;
// SAFETY: Every variant is one of these
unsafe { core::hint::assert_unchecked(w == 8 || w == 16 || w == 32 || w == 64) };
512 / u32::from(w)
}
#[no_mangle]
pub fn from_pow2(w: Width) -> u32 {
let w = w as u8;
// SAFETY: Every variant is a power of two
unsafe { core::hint::assert_unchecked(w.is_power_of_two()) };
512 / u32::from(w)
}
```
`rustc -O --emit asm`:
```asm
from_value_set: ; identical to the version with no assume at all
movzbl %dil, %ecx
movw $512, %ax
xorl %edx, %edx
divw %cx
movzwl %ax, %eax
retq
from_pow2:
rep bsfl %edi, %ecx
movl $512, %eax
shrl %cl, %eax
retq
```
Optimized IR of `from_pow2` for reference, showing the fold that `from_value_set` should get too:
```llvm
%0 = tail call range(i8 1, 8) i8 @llvm.ctpop.i8(i8 %s)
%_3 = icmp eq i8 %0, 1
tail call void @llvm.assume(i1 %_3)
%1 = tail call range(i8 0, 9) i8 @llvm.cttz.i8(i8 %s, i1 true)
%2 = zext nneg i8 %1 to i32
%_01 = lshr exact i32 512, %2
```
## Reproduction (LLVM IR)
This is what the Rust source reaches LLVM as, after SROA and SimplifyCFG fold the front-end's short-circuit diamond into an `or` chain:
```llvm
declare void @llvm.assume(i1)
; expected: `lshr i32 512, cttz(%x)`; actual: `udiv` survives to `divw`
define i32 @from_value_set(i8 range(i8 8, 65) %x) {
%c1 = icmp eq i8 %x, 8
%c2 = icmp eq i8 %x, 16
%c3 = icmp eq i8 %x, 32
%c4 = icmp eq i8 %x, 64
%o1 = or i1 %c1, %c2
%o2 = or i1 %c3, %c4
%o = or i1 %o1, %o2
call void @llvm.assume(i1 %o)
%z = zext nneg i8 %x to i32
%d = udiv i32 512, %z
ret i32 %d
}
; folds today, for contrast
define i32 @from_pow2(i8 range(i8 8, 65) %x) {
%p = call i8 @llvm.ctpop.i8(i8 %x)
%c = icmp eq i8 %p, 1
call void @llvm.assume(i1 %c)
%z = zext nneg i8 %x to i32
%d = udiv i32 512, %z
ret i32 %d
}
declare i8 @llvm.ctpop.i8(i8)
```
## Analysis (LLM-generated)
`computeKnownBitsFromCond` and `isKnownToBeAPowerOfTwo` walk `assume(A and B)` into the facts of `A` and `B`, but there is no rule for `assume(A or B)`. For a disjunction, the intersection of what each disjunct implies still holds: here every `icmp eq X, C` has `C` a power of two, so `X` is a power of two. Two shapes would cover the practical cases:
* `or` of `icmp eq X, C_i` with all `C_i` powers of two, for `isKnownToBeAPowerOfTwo`;
* the same, producing the intersection of `KnownBits` of all `C_i` (`0b0??000`-style knowledge), and a union `ConstantRange`, for `computeKnownBitsFromCond`.
The disjunction is the natural encoding for a front end that knows an exact value set (C-like or Rust enums with power-of-two values: element widths, register group sizes, alignments). When the front end keeps a `switch` over those constants instead, LLVM already gets it right, so this only matters where the value arrives as a plain integer with a fact attached.
## Related
* #115767 / #121386: strength reduction once a power of two is known, which is what makes the
`assume(ctpop == 1)` form work.
* #92074: `isKnownToBeAPowerOfTwo` recursion depth through phis.
* [rust-lang/rust#162513](https://github.com/rust-lang/rust/issues/162513): the front-end side, where the `match` producing the value is folded into a cast so that only a range reaches LLVM.
Contributor guide
Research direction
Search the LLVM sources for computeKnownBitsFromCond and isKnownToBeAPowerOfTwo, then reproduce the supplied LLVM IR case. Add coverage for an llvm.assume containing an or-chain of equality comparisons and verify that the power-of-two fact is propagated so the division can be strength-reduced.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100