llvm / llvm/llvm-project

[X86] Fold paired post-mask users of a constant-data PSHUFB into transformed LUT shuffles

Open
#216,906 1 comment 1 reaction 1 assignee Claimed by @ParkHanbum View on GitHub
backend:X86 missed-optimization
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.