[InstCombine] Missed fold: `lshr C1, (sub (cttz X), C2)` to `lshr (C1 << C2), (cttz X)`
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
Another artifact discovered when investigating https://github.com/rust-lang/rust/issues/162513
When a `switch` over power-of-two constants is turned into a shift map, the result is `C1 >> (cttz(X) - C2)`, which keeps the subtraction even though `(C1 << C2) >> cttz(X)` is equivalent whenever `C1 << C2` does not overflow and the shift is in range. The subtraction costs an instruction on every path that reaches it, in code whose whole point was replacing a division.
Reproduction:
```rust
#[no_mangle]
pub fn expected(x: u8) -> u32 {
let x = (x & 0xF8) | 0x80;
512 >> x.trailing_zeros()
}
// `x` is nonzero with its low three bits clear, so `cttz(x)` is in `3..=7`
#[no_mangle]
pub fn from_sub(x: u8) -> u32 {
let x = (x & 0xF8) | 0x80;
64 >> (x.trailing_zeros() - 3)
}
```
https://rust.godbolt.org/z/afPbf4GY1
Generated assembly:
```asm
expected:
and dil, 120
or dil, -128
movzx eax, dil
rep bsf ecx, eax
mov eax, 512
shr eax, cl
ret
from_sub:
and dil, 120
or dil, -128
movzx eax, dil
rep bsf ecx, eax
add cl, 29 ; the `- 3`, modulo the shift width
mov eax, 64
shr eax, cl
ret
```
Equivalent IR shape:
```llvm
define i32 @from_sub(i8 %x) {
%m = and i8 %x, -8
%v = or i8 %m, -128
%tz = call i8 @llvm.cttz.i8(i8 %v, i1 true) ; range(i8 3, 8)
%z = zext nneg i8 %tz to i32
%s = sub nuw nsw i32 %z, 3
%r = lshr i32 64, %s
ret i32 %r
}
declare i8 @llvm.cttz.i8(i8, i1)
```
Expected: `%r = lshr i32 512, %z` (`64 << 3 = 512`, no overflow; `%z <= 7 < 32`).
## Where it comes from (LLM-generated)
This is what a `switch` mapping `{8, 16, 32, 64}` to `512 / value` becomes after SimplifyCFG builds the shift map and InstCombine folds the division: `C1 >> (cttz(X) - log2(min))`. The general fold `lshr C1, (sub X, C2)` to `lshr (C1 << C2), X` is valid when `C1 << C2` does not overflow and `X` is known to be at least `C2` (so the original shift amount was non-negative) and below the bit width; the `sub nuw` plus range information provides both. The same applies to `shl` with `lshr`/`ashr` of the constant when it is exact.
A Rust-level version of the same `switch` is in https://github.com/rust-lang/rust/issues/162513, where the front end folds the `switch` into a cast first; with `-Zmir-enable-passes=-MatchBranchSimplification` it reaches LLVM as a `switch` and produces exactly the `from_sub` code above.
Contributor guide
Assessment
This issue has not been assessed yet.