[X86] Fold paired post-mask users of a constant-data PSHUFB into transformed LUT shuffles
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
## Summary
LLVM currently leaves the following AVX2 pattern as one `vpshufb`, two
`vpand`, and one `vpsrlw`:
```text
x = pshufb(constant_lut, control)
lo = bitcast(x) & 0x3333
hi = (bitcast(x) >> 2) & 0x3333
```
For the constant table below, both post-processing branches can be folded into
the lookup tables:
```text
lo = pshufb(lut & 0x33 byte-wise, control)
hi = pshufb((lut >> 2) & 0x33 byte-wise, control)
```
The resulting AVX2 sequence uses two `vpshufb` instructions and removes the
two ANDs and word shift. Local LLVM-MCA results are favorable on Haswell and
Skylake, although the transform has a small constant/code-size cost and should
be target-profitability-aware.
## Real-world motivation
This occurs in Faiss v1.15.0, commit
`20f14b31a6d54e243a3d1de6ae193fc4c3ec18ed`, in
`faiss/utils/simd_impl/partitioning_simdlib256.h` inside `compute_accu2<1>`.
The optimized instance is in:
```text
_ZN5faiss16simd_histogram_8ILNS_9SIMDLevelE1EEEvPKtitiPi
```
The source table and consumers are visible around
[`partitioning_simdlib256.h:595-625`](https://github.com/facebookresearch/faiss/blob/20f14b31a6d54e243a3d1de6ae193fc4c3ec18ed/faiss/utils/simd_impl/partitioning_simdlib256.h#L595-L625):
```cpp
alignas(32) static const uint8_t shifts[32] = {
1, 16, 0, 0, 4, 64, 0, 0, 0, 0, 1, 16, 0, 0, 4, 64,
1, 16, 0, 0, 4, 64, 0, 0, 0, 0, 1, 16, 0, 0, 4, 64};
a4lo += a2 & simd16uint16(0x3333);
a4hi += (a2 >> 2) & simd16uint16(0x3333);
```
The reproducer isolates the `N=1` switch-tail case where `a2` is one lookup
result. It does not propose the same local rewrite for the `N=2/3` cases where
`a2` is a sum of multiple lookup results.
## Minimal source IR (`src.ll`)
```llvm
source_filename = "faiss-pshufb-lut-postmask-issue-src"
target triple = "x86_64-unknown-linux-gnu"
define void @faiss_pshufb_lut_postmask(
ptr %lo_out, ptr %hi_out, <32 x i8> %control) {
entry:
%lookup = call <32 x i8> @llvm.x86.avx2.pshuf.b(
<32 x i8> ,
<32 x i8> %control)
%words = bitcast <32 x i8> %lookup to <16 x i16>
%lo = and <16 x i16> %words, splat (i16 13107)
%shift = lshr <16 x i16> %words, splat (i16 2)
%hi = and <16 x i16> %shift, splat (i16 13107)
store <16 x i16> %lo, ptr %lo_out, align 32
store <16 x i16> %hi, ptr %hi_out, align 32
ret void
}
declare <32 x i8> @llvm.x86.avx2.pshuf.b(<32 x i8>, <32 x i8>)
```
## Proposed target IR (`tgt.ll`)
```llvm
source_filename = "faiss-pshufb-lut-postmask-issue-tgt"
target triple = "x86_64-unknown-linux-gnu"
define void @faiss_pshufb_lut_postmask(
ptr %lo_out, ptr %hi_out, <32 x i8> %control) {
entry:
%lo_bytes = call <32 x i8> @llvm.x86.avx2.pshuf.b(
<32 x i8> ,
<32 x i8> %control)
%hi_bytes = call <32 x i8> @llvm.x86.avx2.pshuf.b(
<32 x i8> ,
<32 x i8> %control)
store <32 x i8> %lo_bytes, ptr %lo_out, align 32
store <32 x i8> %hi_bytes, ptr %hi_out, align 32
ret void
}
declare <32 x i8> @llvm.x86.avx2.pshuf.b(<32 x i8>, <32 x i8>)
```
## Correctness rationale
For each 16-byte PSHUFB lane:
```text
T = [1,16,0,0,4,64,0,0, 0,0,1,16,0,0,4,64]
TA = [1,16,0,0,0, 0,0,0, 0,0,1,16,0,0,0, 0]
TB = [0, 0,0,0,1,16,0,0, 0,0,0, 0,0,0,1,16]
```
`TA[i] = T[i] & 0x33` and `TB[i] = (T[i] >> 2) & 0x33`, with the latter
shift performed per byte.
For a little-endian result word `w = L | (H << 8)`:
```text
w & 0x3333
= (L & 0x33) | ((H & 0x33) << 8)
(w >> 2) & 0x3333
= ((L >> 2) & 0x33) | (((H >> 2) & 0x33) << 8)
```
Bits crossing from `H` into the low byte under the word shift land in bit
positions 6-7, which `0x33` clears. Both target shuffles use the same control,
so PSHUFB's control-high-bit zeroing behavior is preserved.
## Cost-model evidence
LLVM 24 `llc` and `llvm-mca --iterations=100` on the complete minimal function:
| CPU | Form | Instructions / 100 | uOps / 100 | Cycles / 100 | Block RThroughput |
| --- | --- | ---: | ---: | ---: | ---: |
| Haswell | current | 1000 | 1700 | 509 | 4.3 |
| Haswell | proposed | 800 | 1500 | 409 | 3.8 |
| Skylake | current | 1000 | 1700 | 311 | 2.8 |
| Skylake | proposed | 800 | 1500 | 309 | 2.5 |
The function body decreases from 10 to 8 instructions and from 17 to 15
modeled uOps per iteration. The proposed form increases loaded constant payload
from 20 to 32 bytes and total text-plus-constant size from 132 to 136 bytes.
It also increases shuffle-resource pressure: Haswell's modeled port-5-class
pressure changes from 1.80 to 2.01 per iteration. This suggests a target-aware
profitability check rather than an unconditional transform.
The Faiss occurrence is an `N=1` tail, so this evidence establishes local
code-generation value but does not claim a measurable end-to-end Faiss speedup.
Compiler-explorer : https://compiler-explorer.com/z/dWKT3M1bK
Contributor guide
Assessment
This issue has not been assessed yet.