llvm / llvm/llvm-project

[ValueTracking] `isKnownToBeAPowerOfTwo` ignores `llvm.assume` of a disjunction of equalities with power-of-two constants

Open
#222,278 2 comments 0 reactions 0 assignees View on GitHub
llvm:analysis missed-optimization
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.