[codex] perf: extend small dictionary take across NEON value widths

[joseph-isaacs]

Superseded by #10317. The expanded NEON kernels and benchmarks have been consolidated there in commit 04f4c4e, alongside its broader x86 implementation. The original branch is retained for benchmark provenance.

---

Summary

Extend small-dictionary take beyond the recently merged AVX2 and NEON 16-entry byte tables. Keep the AVX2/VBMI extension to 32 byte values and broaden NEON to fixed-width values and larger dictionaries with u8 codes.

Physical value width	NEON dictionary limit	Minimum output rows	Examples
1 byte	256	64	u8, i8, narrow decimals
2 bytes	128	64	u16, i16, f16, narrow decimals
4 bytes	64	64	u32, i32, f32, narrow decimals
8 bytes	32	1,024	u64, i64, f64, narrow decimals

Every cardinality within each limit is supported, including non-powers of two. Wider code types and unsupported value widths/cardinalities retain the existing fallback. The 16- and 32-byte record experiments did not beat scalar copying, so those widths remain on the fallback. Variable-length strings and bit-packed booleans are outside this fixed-width path.

Implementation

• AVX2: two VPSHUFB lookups and a blend for 17–32 byte values.
• AVX-512 VBMI: VPERMB for 17–32 byte values when supported.
• NEON byte values: one- through four-register TBL, then multiple 64-byte banks combined with OR for dictionaries through 256 bytes. Independent maximum-code accumulators shorten the bounds-check dependency chain.
• NEON wider values: transpose the dictionary into a table for each byte position, look up 16 codes, and interleave the results using ST2/ST4. This avoids repeating a multi-bank lookup for every output byte.
• Preserve checked indices, scalar tails, allocation behavior, and logical types. Keep native NEON table helpers separate from dispatch.
• Add a standalone comparison against generic fearless_simd 1.0.0. Its dependency is outside the production workspace. It directly includes the production NEON loops and implements matching byte-table and transposed-table generic kernels.

NEON — Apple M3 Max

End-to-end dictionary execution, including allocation, one million uniform non-null u8 codes. Times are microseconds. Rust 1.98.0, native aarch64 macOS, no CPU pinning.

Baseline is 3c0a33e with the expanded benchmark source. It already has the earlier PR's 32-entry byte fast path; larger byte tables and wider values use the scalar default. Candidate is cac5eca.

Value type	Dictionary values	Baseline	Candidate	Speedup
u8	2	36.27	20.06	1.81×
u8	4	36.30	20.02	1.81×
u8	8	36.23	19.90	1.82×
u8	16	36.11	19.99	1.81×
u8	32	36.32	21.70	1.67×
u8	48	279.40	20.74	13.47×
u8	64	281.00	22.60	12.43×
u8	96	282.30	45.96	6.14×
u8	128	281.90	46.15	6.11×
u8	192	280.10	68.35	4.10×
u8	256	281.60	95.80	2.94×
u16	64	280.10	47.18	5.94×
u16	128	281.90	90.60	3.11×
u32	32	280.40	68.76	4.08×
u32	64	280.70	106.80	2.63×
u64	16	557.30	151.50	3.68×
u64	32	560.40	166.50	3.37×
f16	128	281.60	89.10	3.16×
f32	64	281.30	105.80	2.66×
f64	32	559.20	156.40	3.58×

Entries are medians of three run medians, with 31 samples of four iterations. Run order was baseline/candidate, candidate/baseline, baseline/candidate. All 1,408 cases ran in each executable: 11 primitive types × 16 cardinalities × four lengths × two distributions.

The host showed scheduling variation, including in unchanged fallback cases. For example, u64/32 candidate run medians were 166.5, 156.4, and 325.3 µs. Small differences at 64 rows should not be treated as stable wins. These measurements establish gains on this M3 Max; they are not measurements of other NEON implementations.

Native NEON versus generic fearless_simd

Kernel-only comparison on the same M3 Max with Rust 1.98.0, one million uniform codes. Times are microseconds. Output is preallocated; table preparation, dispatch, bounds checks, lookup, and tails are timed. The generator differs from the end-to-end benchmark, so compare columns within this table.

Value width	Dictionary values	Checked scalar	Native NEON	Generic fearless_simd
1	16	301.65	18.71	34.88
1	32	298.15	18.71	18.69
1	48	298.65	18.75	20.23
1	64	300.02	20.44	20.12
1	128	295.08	43.54	43.29
1	256	296.29	89.29	88.73
2	64	277.21	43.54	43.38
2	128	277.58	83.54	85.12
4	32	278.08	63.37	62.63
4	64	277.29	99.96	99.71
8	16	278.46	144.83	164.42
8	32	278.00	149.92	148.33
16	2	309.54	571.25	327.50
32	2	388.50	1037.88	910.04

The 16-/32-byte native entries are experimental kernels, not enabled production paths. Most enabled cases are effectively tied between native and generic SIMD. The byte/16 native loop uses four independent accumulators while the generic version uses a single 16-lane vector; this comparison applies to the supplied kernels, not to the libraries in general.

Each entry is the median of three run medians. Each run rotates the three variants over 31 samples, with clamp(2_000_000 / output_bytes, 1, 256) iterations per sample. All 504 datasets and setup correctness checks ran successfully.

AVX — Intel Xeon 6975P-C

Measured over SSH on vortex-ec2, pinned to CPU 2 with taskset. These are the previously collected measurements from a0aaec6, before the module cleanup; they are not a measurement of the latest PR head.

Dictionary values	Scalar fallback	Upstream d97fb1c	Measured x86 candidate
2	292.8	38.86	37.53
4	286.8	38.84	38.39
8	287.1	38.84	38.37
16	281.6	38.85	38.35
32	284.5	286.40	37.38

The 32-value case is approximately 7.7× faster. A separately measured AVX2-only variant took 37.62 µs, so this improvement is available without AVX-512.

The scalar diagnostic variant bypassed small-table dispatch; the AVX2-only diagnostic variant disabled VBMI selection. Each variant was built and saved as a separate executable. Baseline and diagnostic entries are medians of three runs in alternating order, with 200 samples of 20 iterations. The final x86 dispatch was measured once with those settings.

Reproduction and validation

For the end-to-end comparison, keep take_small_dict.rs identical in the baseline and candidate and save each emitted executable:

cargo bench -p vortex-array --bench take_small_dict --no-run
<saved-benchmark-binary> --bench --sample-count 31 --sample-size 4

For the generic comparison, on AArch64:

cargo run --release --manifest-path benchmarks/take-small-dict/Cargo.toml > results.jsonl

See the comparison README for timing boundaries, data generation, and dispatch limits.

• Native baseline/candidate benchmark builds succeeded; all 1,408 scalar-reference checks passed in all six end-to-end runs.
• The standalone benchmark built with Rust 1.98.0. All three runs passed comparisons across record widths, non-power-of-two counts, unaligned inputs, vector tails, and invalid-code rejection.
• Earlier x86 builds and 40 reference checks passed; the VBMI binary emitted vpermb. No new x86 performance measurements were collected for this extension.
• Unit tests, formatting, and Clippy were not run.

[codspeed]

Merging this PR will regress 3 benchmarks

⚠️ Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 3 improved benchmarks
❌ 3 regressed benchmarks
✅ 2095 untouched benchmarks
🆕 120 new benchmarks
⏩ 518 skipped benchmarks¹

Warning

Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
❌	WallTime	mul_u64_nonnull_neon	29.1 µs	39.6 µs	-26.38%
❌	WallTime	multiply_shapes_neon[(32768, PerRowPerRow)]	32.9 µs	38.8 µs	-15.34%
❌	WallTime	mul_i64_nonnull_neon	33 µs	38.7 µs	-14.91%
⚡	Simulation	density_sweep_dense_runs[0.9]	56.1 µs	18.1 µs	×3.1
⚡	Simulation	decode_primitives[u8, (2000, 32)]	44.1 µs	34.1 µs	+29.45%
⚡	Simulation	decode_primitives[u8, (1000, 32)]	36.7 µs	32.5 µs	+12.72%
🆕	WallTime	skewed_avx512[16, 1000000]	N/A	51 µs	N/A
🆕	WallTime	skewed_avx512[16, 1024]	N/A	1 µs	N/A
🆕	WallTime	skewed_avx512[16, 64]	N/A	985 ns	N/A
🆕	WallTime	skewed_avx512[16, 65536]	N/A	3.3 µs	N/A
🆕	WallTime	skewed_avx512[2, 1000000]	N/A	51.2 µs	N/A
🆕	WallTime	skewed_avx512[2, 1024]	N/A	1 µs	N/A
🆕	WallTime	skewed_avx512[2, 64]	N/A	1 µs	N/A
🆕	WallTime	skewed_avx512[2, 65536]	N/A	3.1 µs	N/A
🆕	WallTime	skewed_avx512[32, 1000000]	N/A	49.5 µs	N/A
🆕	WallTime	skewed_avx512[32, 1024]	N/A	1 µs	N/A
🆕	WallTime	skewed_avx512[32, 64]	N/A	994 ns	N/A
🆕	WallTime	skewed_avx512[32, 65536]	N/A	3 µs	N/A
🆕	WallTime	skewed_avx512[4, 1000000]	N/A	51.1 µs	N/A
🆕	WallTime	skewed_avx512[4, 1024]	N/A	1 µs	N/A
...	...	...	...	...	...

ℹ️ Only the first 20 benchmarks are displayed. Go to the app to view all benchmarks.

Tip

Investigate this regression by commenting @codspeedbot fix this regression on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing ji/small-dict-take-kernel (3c0a33e) with develop (5079a20)}

Footnotes

1. 518 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

[codspeed]

Too many benchmarks in a single upload

The performance report could not be generated because there were too many benchmarks in a single upload to CodSpeed. We recommend sharding your benchmarks into smaller uploads, with max 1000 benchmarks per upload. See the documentation for more information.