Segment-wise SIMD shuffle under a variable permutation

Problem. Given a 64-byte vector v (one cache line, one zmm), a 64-bit boundary mask m (bit i set ⇔ a segment starts at byte i; bit 0 always set; k = popcount(m) segments), and a variable permutation π of {0..k-1}: produce the segments of v concatenated in the order π[0], π[1], …, π[k-1]. Both m and π vary per call — nothing can be precompiled into shuffle constants.

Hardware. Intel Xeon, family 6 model 207 = Emerald Rapids (Golden Cove cores, Sapphire Rapids pipeline), with AVX-512 VBMI/VBMI2/BITALG/VPOPCNTDQ. That's the fun regime: vpermb (arbitrary 64-byte permute), vpcompressb, vpexpandb all operate on a full cache line in one register.

1. Algorithm

Everything reduces to building the byte-level gather index idx[j] = source position of output byte j, then one vpermb idx, v finishes the job. The interesting part is building idx in registers, without per-segment loops.

Derivation (all arithmetic on byte lanes, mod 256):

starts   = vpcompressb(m,  iota)      // position of each set bit = old
                                      // segment starts, packed to lanes 0..k-1
ends+1   = vpcompressb(m2, iota+1)    // m2 = (m >> 1) | bit63: last byte of
                                      // each segment; +1 folded into the iota
lens     = (ends+1) - starts
plens    = vpermb(π, lens)            // segment lengths in output order
pstarts  = vpermb(π, starts)          // old starts in output order
incl     = prefix_sum(plens)          // inclusive, per byte lane
newstart = incl - plens               // output-side segment starts
delta    = pstarts - newstart         // per-segment gather offset:
                                      //   idx[j] = j + delta[outseg(j)]

To broadcast delta[t] across each output segment without knowing segment extents per lane, difference-encode it and let a second prefix sum do the propagation:

ddelta   = delta - (delta shifted up one lane)     // ddelta[0] = delta[0]
nb       = bitmask with bits at newstart[t]        // output boundary mask
seeded   = vpexpandb(nb, ddelta)                   // ddelta[t] lands exactly
                                                   // at position newstart[t]
deltapos = prefix_sum(seeded)                      // telescopes to
                                                   // delta[outseg(j)] at all j
idx      = iota + deltapos
out      = vpermb(idx, v)

Two cute facts make this sound:

mod-256 wraparound is free. delta and ddelta wrap in byte lanes, but the telescoping sum restores delta[outseg(j)] mod 256, and vpermb only reads the low 6 bits of each index — so mod 256 ⊇ mod 64 is exactly enough. No saturation or widening needed anywhere.
vpcompressb of iota is a "positions of set bits" instruction. Both segment starts and ends come from two compresses; lengths are one subtract.

The one step with no clean SIMD form is nb (a scatter of bits to positions given by byte values). Two implementations, both kept:

scalar: for t: nb |= 1 << incl[t] — O(k), fine for small k;
vector: zero-extend incl in 8-byte chunks (vpmovzxbq), variable-shift 1 by each (vpsllvq, port 0), OR-reduce. ~O(k/8) with the serial chain replaced by a reduction tree. Padding lanes ≥ k must be zero-masked — their mod-256 garbage can alias back into 0..63 and set stray bits (found by differential test, naturally).

Correctness for all variants: differential testing against the scalar reference, 64 × 20 000 random (v, m, π) cases covering every k = 1..64. All variants agree.

2. Final kernel (V2c, solver-guided layout)

#include <immintrin.h>
#include <stdint.h>

#define N 64

static __m512i bcast7;        // per-128-bit lane {7×8, 15×8}
static __m512i shift1_idx;    // vpermb indices for 1-byte cross-lane shift
static __mmask64 shift1_msk;

static void init_tables(void) {
  uint8_t b[N];
  for (int i = 0; i < N; i++) b[i] = (uint8_t)((i % 16 < 8) ? 7 : 15);
  bcast7 = _mm512_loadu_si512(b);
  for (int i = 0; i < N; i++) b[i] = (uint8_t)((i - 1) & 63);
  shift1_idx = _mm512_loadu_si512(b);
  shift1_msk = ~1ull;
}

// Inclusive byte prefix sum over 64 lanes, mod 256.
// In-qword levels run as qword *bit* shifts on port 0; only the cross-qword
// carry touches the shuffle port (1× vpshufb + 4× valignq).
static inline __m512i prefix_sum_bytes(__m512i x) {
  const __m512i z = _mm512_setzero_si512();
  x = _mm512_add_epi8(x, _mm512_slli_epi64(x, 8));    // p0
  x = _mm512_add_epi8(x, _mm512_slli_epi64(x, 16));   // p0
  x = _mm512_add_epi8(x, _mm512_slli_epi64(x, 32));   // p0
  __m512i t = _mm512_shuffle_epi8(x, bcast7);         // qword totals, lat 1
  __m512i e = _mm512_alignr_epi64(t, z, 7);           // e[L] = tot[L-1]
  e = _mm512_add_epi8(e, _mm512_alignr_epi64(e, z, 7));
  e = _mm512_add_epi8(e, _mm512_alignr_epi64(e, z, 6));
  e = _mm512_add_epi8(e, _mm512_alignr_epi64(e, z, 4));
  return _mm512_add_epi8(x, e);                       // + exclusive lane carry
}

void seg_shuffle(const uint8_t *v, uint64_t m, const uint8_t *pi /* 64B,
                 lanes >= k ignored */, int k, uint8_t *out) {
  const __m512i iota = _mm512_set_epi8(
      63,62,61,60,59,58,57,56,55,54,53,52,51,50,49,48,
      47,46,45,44,43,42,41,40,39,38,37,36,35,34,33,32,
      31,30,29,28,27,26,25,24,23,22,21,20,19,18,17,16,
      15,14,13,12,11,10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0);
  const __m512i iota1 = _mm512_add_epi8(iota, _mm512_set1_epi8(1));

  __m512i starts = _mm512_maskz_compress_epi8(m, iota);
  uint64_t m2 = (m >> 1) | (1ull << 63);
  __m512i ends1 = _mm512_maskz_compress_epi8(m2, iota1);
  __m512i lens = _mm512_sub_epi8(ends1, starts);

  __m512i piv = _mm512_loadu_si512(pi);
  __m512i plens = _mm512_permutexvar_epi8(piv, lens);
  __m512i pstarts = _mm512_permutexvar_epi8(piv, starts);

  __m512i incl = prefix_sum_bytes(plens);
  __m512i newstart = _mm512_sub_epi8(incl, plens);
  __m512i delta = _mm512_sub_epi8(pstarts, newstart);
  __m512i dsh = _mm512_maskz_permutexvar_epi8(shift1_msk, shift1_idx, delta);
  __m512i ddelta = _mm512_sub_epi8(delta, dsh);

  // output boundary mask nb (bits at newstart[t]) via vpsllvq chunks
  uint64_t valid = (k > 1) ? (((uint64_t)1 << (k - 1)) - 1) : 0;
  __m512i one = _mm512_set1_epi64(1), acc = _mm512_setzero_si512();
  uint8_t inclb[N] __attribute__((aligned(64)));
  _mm512_store_si512(inclb, incl);
  for (int c = 0; c * 8 < k - 1; c++) {
    __mmask8 mk = (__mmask8)(valid >> (c * 8));
    __m512i q = _mm512_cvtepu8_epi64(
        _mm_loadl_epi64((const __m128i *)(inclb + c * 8)));
    acc = _mm512_or_si512(acc, _mm512_maskz_sllv_epi64(mk, one, q));
  }
  uint64_t nb = 1 | (uint64_t)_mm512_reduce_or_epi64(acc);

  __m512i seeded = _mm512_maskz_expand_epi8(nb, ddelta);
  __m512i deltapos = prefix_sum_bytes(seeded);
  __m512i idx = _mm512_add_epi8(iota, deltapos);

  _mm512_storeu_si512(out,
      _mm512_permutexvar_epi8(idx, _mm512_loadu_si512(v)));
}

The companion perm.c additionally contains: the scalar reference (shuffle_scalar, tzcnt + memcpy), a scalar-index-build + vpermb hybrid (V1), the pre-optimization SIMD kernels (V2 with the all-vpermb prefix sum, V2b with vectorized nb), the apply-only lower bound, the differential tester, and the benchmark harness.

3. Solver-driven layout optimization

Since m and π are variable, there's nothing to superoptimize into constants — the leverage is in scheduling: which instruction forms to use so the dependency DAG fits the Golden Cove port structure. Two-step approach:

Cost model from llvm-mca. llvm-mca -mcpu=emeraldrapids probes on instruction snippets gave the per-op latency/port table (matching the SPR scheduler model):

instruction	uops	lat	port
`vpermb`	1	3	p5
`vpermb {k}{z}`	1	5	p5
`vpcompressb {z}`	2	6	2×p5
`vpexpandb {z}`	2	8	2×p5
`vpaddb/vpsubb/vporq`	1	1	p0 or p5
`valignq`	1	3	p5
`vpslldq` / `vpshufb` (zmm)	1	1	p5
`vpsllq imm` / `vpsllvq`	1	1	p0
`vpmovzxbq` (reg / mem)	1 / 2	3 / 11	p5 (+load)
`vextracti64x4`	1	3	p5
`kmovq k, r`	1	1	p5 (!)

Three facts drove the redesign: zero-masked vpermb costs latency 5, not 3, in this model; every 512-bit shuffle is p5-only (the kernel is fundamentally p5-bound); and qword bit-shifts (vpsllq) run on p0.

z3 resource-constrained scheduling (sched.py, run under uv run --with z3-solver). Each design's DAG is encoded with integer start times, t[op] ≥ t[dep] + lat[dep] edges, port capacity as distinct issue cycles per port (ALU-class ops get a Boolean choosing p0 vs p5), and Optimize.minimize(makespan). Steady-state throughput is bounded separately by port counts (p5 dominates). Comparing the original prefix sum (A: 6 × zero-masked vpermb + add) against the restructured one (B: 3 in-qword levels as vpsllq+vpaddb on p0, cross-qword carry as vpshufb + valignqs):

design	k	p5 uops	z3 min makespan (latency)	port bound (cyc/call)
A (V2b)	8	28	124	~28
A (V2b)	64	35	132	~35
B (V2c)	8	24	90	~26
B (V2c)	64	31	98	~36

So the solver's message: design B buys ~34 cycles of dependency latency and 4 p5 uops; steady-state throughput moves less because the kernel stays p5-bound. Both predictions were borne out (below). One caveat worth stating plainly: on a big out-of-order core the solver's schedule itself is irrelevant (the hardware reschedules); its value is ranking structural alternatives — which instruction forms, which decomposition — under an honest latency/port model.

4. Measurements

Emerald Rapids, gcc 13.3 -O3 -march=native, rdtscp, median of 25 runs over batches of 4096 random (v, m, π) per k. Cycles per call.

Throughput (independent calls):

k	scalar memcpy	scalar-idx + vpermb	V2 (simd idx)	V2b (+vec nb)	V2c (solver layout)	apply-only
1	16	23	25	25	23	9.5
2	49	78	29	35	32	8.5
4	123	204	28	32	30	8.5
8	249	317	31	34	32	8.5
16	463	551	38	33	31	8.5
32	766	922	59	40	34	8.5
48	820	734	76	48	42	8.5
64	635	414	99	49	47	8.5

Latency mode (serially dependent calls: output feeds the next input; m, π fixed, so only the data path serializes — construction still pipelines):

k	V2	V2b	V2c
8	33	33	30
32	57	40	37
64	97	54	45

Cross-checks: llvm-mca on the compiled V2c body predicts block reciprocal throughput 29.0 cycles — measured ~30–32 at small k. The apply-only floor (~8.5 cycles = load + vpermb + store) shows what you get if (m, π) is reused and the index vector amortized: build once, then ~9 cycles per line.

Headline: near-flat 23–47 cycles across the entire k range, vs scalar 16→920. At k = 32 that's a 22× speedup; even against the better scalar variant at k = 64 it's ~9×.

5. What worked

vpcompressb(m, iota) as a one-instruction "enumerate set-bit positions" — the cleanest way to get segment starts and ends into lanes.
Difference-encoded delta + vpexpandb + prefix sum to broadcast per-segment offsets across variable-length runs with no loop. Letting byte lanes wrap mod 256 and relying on vpermb's mod-64 index read made the whole thing overflow-proof for free.
Solver-guided restructuring: moving the three in-qword prefix-sum levels to port 0 (vpsllq bit-shifts instead of cross-lane byte shuffles). This is the single change behind the V2b→V2c gains, and it came directly out of the latency table (maskz vpermb = 5) plus the z3 makespan comparison — not something obvious from reading the intrinsics.
Vectorized nb (vpmovzxbq/vpsllvq/OR-reduce) replacing the scalar 1 << incl[t] loop: turned the O(k) serial tail from ~1.3 cycles/segment into a short reduction — the k = 64 case went 99→49 before the layout work.
Differential testing as the safety net — it caught both real bugs: the missing shift-1 step in the Hillis–Steele lane carry (a classic), and stray nb bits from unmasked padding lanes wrapping mod 256.

6. What didn't work / dead ends

Scalar index build + vpermb (V1) loses to plain memcpy at mid k. Writing 64 index bytes through nested variable-length loops is more branch misprediction than memcpy'ing the segments directly. It only wins at k = 64 (all loops degenerate to length 1). Building the index is the entire problem; applying it is trivial.
No scatter, no clean nb. Everything else in the kernel is branch-free SIMD; the output boundary mask is a bit-scatter (bits at positions given by byte values) and AVX-512 simply has no byte/bit scatter. Attempted routes that all collapsed back into needing nb or an inverse permutation: computing per-position segment ids via vpopcntb-style rank; building the scatter permutation source-side and inverting it (inverting a permutation is a scatter); vpshufbitqmb (BITALG) only selects bits within qwords. The vpsllvq reduction is the honest workaround.
vpmovzxbq from memory doesn't dodge the shuffle port — the memory form still carries a p5 uop (lat 11 total), so feeding nb chunks from the spilled incl buffer saves nothing over the register form. mca killed that idea before it was written.
vpermt2b/vpermi2b with a zero source as an unmasked substitute for zero-masked vpermb: 3 uops (2×p5 + p0), lat 5 — strictly worse. Same story for masked valignq variants.
z3 as a prover, first attempt: a bounded correctness proof of the ddelta construction (symbolic mask + symbolic π, BitVec(8) semantics) verified k ≤ 4 at N = 8 quickly but slowed badly at higher k — and was scrapped in favor of differential testing, with the solver redeployed for scheduling, where it actually earns its keep here.
Latency benchmark subtlety: chaining only v between calls doesn't serialize the index-construction pipeline (it depends on m, π), so the measured serial-mode numbers understate the 124→90 model delta. The model's prediction shows up as the growing V2→V2c gap at high k, where the nb chain lengthens the critical path.

7. Reproduce

gcc -O3 -march=native perm.c -o perm && ./perm     # diff tests + benchmarks
uv run --with z3-solver sched.py                   # scheduling comparison
llvm-mca-18 -mcpu=emeraldrapids v2c_body.s         # port-pressure cross-check

Portability notes: requires AVX-512 VBMI + VBMI2 (Ice Lake+; Zen 4/5 also qualify, with different port economics — Zen 4 double-pumps 512-bit ops, so the p0/p5 split argument changes). An AVX2 fallback would work on 32-byte lines: vpermd handles dword-granular cross-lane moves, byte-granular needs 2×vpshufb + blend, and compress/expand must be emulated via pshufb-table lookups per 8-bit mask chunk — sketched but not implemented.

chadbrewbaker/perm.md

Select an option

No results found