[alexreinking]

Vectorized lookups into small tables currently compile poorly on CPUs: a vectorized mux(idx, {c0, ..., cn-1}) becomes a chain of n compares and selects, and a load from a small Buffer becomes a gather (one scalar load per lane). Most CPUs have a register table-lookup instruction that does either in one instruction per vector: ARM tbl, x86 pshufb/vpermb, WebAssembly i8x16.swizzle. This matters for quantized-weight formats that decode 4-bit codes through a 16-entry int8 codebook (e.g. GGML's mxfp4 and iq4_nl), which GGML's NEON kernels decode with a single vqtbl1q_s8 per 16 codes.

This PR adds a pass, optimize_small_table_lookups (in OptimizeShuffles.{h,cpp}), that turns such lookups into dynamic_shuffle(table, index, 0, n - 1), and lowers dynamic_shuffle on ARM, x86 and WebAssembly.

What is recognized.

• mux(index, {c0, ..., cn-1}) where every value is a constant. The table is a constant vector. If constant_integer_bounds can't prove the index lies in [0, n - 1], it is clamped (as unsigned) to n - 1, which matches mux's semantics: any out-of-range index, including a negative one, selects the last value.
• A Load from a Buffer embedded in the pipeline (image defined, no Parameter) whose index is proven to lie in a window of at most n elements inside the buffer's allocation. The table is a dense vector load of that window, and the index is rebased to it.
  • The contents of embedded buffers are not treated as compile-time constants. The JIT passes embedded buffers by reference on every call, so their contents can change between realizations, and the test checks this.
  • Only their shape is relied on: it is fixed at compile time and asserted on entry to the pipeline. So the window load never needs padding and never reads out of bounds. This sidesteps the input-buffer padding question that Support Call::dynamic_shuffle for LUT in SVE2 #9087 has to deal with.
• If the index range can't be proven, the IR is left alone.

Where. The pass runs at the start of CodeGen_LLVM::compile_func, the same place #9087 runs optimize_shuffles for SVE2 and Hexagon runs its own. The query is a new virtual, int CodeGen_LLVM::max_dynamic_shuffle_table_size(const Type &elem) const: the largest table the target can profitably index, or 0 to disable. The default is 0, so every target without an override keeps today's code. That includes Hexagon, which has its own dynamic_shuffle lowering, and SVE, which #9087 targets. So recognition is generic, and lowering is per target.

Lowering.

Target	Condition	Instruction	Max table
AArch64	NEON, fixed-width vectors	tbl1–tbl4 (8- or 16-lane results)	64 bytes
ARMv7	NEON	vtbl1–vtbl4	32 bytes
x86	SSE4.1 (SSSE3)	pshufb	16 bytes
x86	AVX2 / AVX-512	vpshufb ymm/zmm, table replicated per 128-bit lane	16 bytes
x86	AVX-512 VBMI (avx512_cannonlake)	vpermb	64 bytes
WebAssembly	wasm_simd128	i8x16.swizzle	16 bytes

• Only 8-bit element types are handled for now; wider tables are left alone.
• Wider vectors are split into slices, and the table is zero-padded to whole registers.
• A shared helper, CodeGen_LLVM::codegen_dynamic_shuffle, does the slicing and padding.
• CodeGen_LLVM also gets a generic, slow but correct lowering of dynamic_shuffle: one variable extractelement per lane.

Relationship to #9087. That draft PR generalizes optimize_shuffles (the bounded indirect-load → dynamic_shuffle pass) to SVE2 and lowers dynamic_shuffle to SVE2 tbl. This PR is complementary and is designed so the two can converge:

• It emits the same dynamic_shuffle(table, u8 index, 0, n - 1) form, and the backends also cast the index to u8 in codegen as Support Call::dynamic_shuffle for LUT in SVE2 #9087 does.
• It does not touch optimize_shuffles.
• CodeGen_ARM::max_dynamic_shuffle_table_size returns 0 under SVE (with a TODO), so Support Call::dynamic_shuffle for LUT in SVE2 #9087's SVE2 lowering can claim constant tables there too.

Effect. In a Halide reimplementation of GGML's mxfp4 kernels (apps/ggml on a branch), the codes are now decoded through mux(nibble, kvalues_mxfp4) instead of arithmetic. That now compiles to one tbl per 16 codes, as in GGML's NEON code (vqtbl1q_s8).

mxfp4 · q8_0 kernel (Apple M-series, 1 thread)	Before	After	GGML
vec_dot inner loop, instructions per 2 blocks (64 codes)	85	49	~51
vec_dot speed relative to GGML, K = 4096	0.37×	0.88×	1×
gemv speed relative to GGML, 4096×4096	n/a	0.93×	1×

Nothing else in the app's kernels changed: rebuilding its ~340 AOT objects with this patch, but without the app change, gave byte-identical objects.

Tests

New tests:

• correctness_small_table_lookup (JIT) covers:

  • mux with an in-range index, a let-bound index, an unbounded u8 index, and a possibly negative i32 index;
  • embedded buffers with a clamped index, a non-zero min and a sub-range window, and two dimensions;
  • an index not provable at compile time (bounded by a Param), which must be left alone;
  • mutating embedded buffers between realizations.

  It runs these for table sizes 2–65, vector widths 8/16/64, int8/uint8, plus uint16/float, which are not transformed. It also checks the generic dynamic_shuffle fallback.

  Targets: the host, plus host-without-NEON on ARM, plus SSE4.1-only and baseline on x86 hosts, plus no-SIMD on wasm.

• simd_op_check_{arm,x86,wasm}:

  • ARM: tbl with 1–4 table registers on arm-64, and vtbl.8 with 1/2/4 registers on arm-32;
  • x86: pshufb, vpshufb ymm/zmm, and vpermb;
  • wasm: i8x16.swizzle, replacing a TODO.

Commands run (macOS arm64 host, LLVM 22.1, cmake --preset macOS):

Command	Result
ctest -j 8 -R "^(correctness_(mux|gather|aligned_split_mux_phase|many_inlined_selects|multi_way_select|unroll_huge_mux|non_vector_aligned_embeded_buffer|vector_cast|vector_math|simd_op_check_.*|register_shuffle|hexagon_scatter|lazy_convolution|convolution|decompose_vector_shuffle|vector_shuffle|fallback_vscale_sve|sve_codegen_predicated|sme_streaming|small_table_lookup)|performance_clamped_vector_load|generator_aot_shuffler)$"	23 passed; 5 skipped (hexagon_scatter, register_shuffle, and simd_op_check_{hvx,powerpc,riscv}, whose targets this LLVM lacks)
HL_JIT_TARGET=wasm-32-wasmrt-wasm_simd128 ./test/correctness/correctness_small_table_lookup	Success
With Homebrew LLVM 23.1 (which has Hexagon/PowerPC/RISC-V): ./test/correctness/correctness_simd_op_check_hvx, ..._powerpc, ..._riscv	Success (all three)
pre-commit hooks on commit (clang-format, keep-sorted, codespell, CMake checks)	Passed

Not run:

• The full test suite.
• x86 execution. On this arm64 host, simd_op_check_x86 checks the generated assembly only, and the correctness test's x86 target variants run only on x86 hosts.
• ARMv7 execution.

Fixes: no tracking issue.

Breaking changes

None. No public API changes. Generated code changes only on ARM (NEON), x86 (SSE4.1+) and WebAssembly (SIMD128), for vectorized 8-bit lookups into small constant tables and embedded buffers; results are unchanged.

Checklist

• Tests added or updated (not required for docs, CI config, or typo fixes)
• Documentation updated (if public API changed) — no public API changes; the new Internal pass and hook are documented in their headers.
• Python bindings updated (if public API changed) — not applicable, no public API changes.
• Benchmarks are included here if the change is intended to affect performance. See "Effect" above.
• Commits include AI attribution where applicable (see Code of Conduct). Written with GitHub Copilot; the commit carries a Co-authored-by: Copilot trailer.

[alexreinking]

@copilot Please fix the clang-tidy failure.

[Copilot]

@copilot Please fix the clang-tidy failure.

Added args.reserve(regs + 1) to resolve the clang-tidy warning. Commit: c7785fa.

[codecov]

Codecov Report

❌ Patch coverage is 85.52036% with 32 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.21%. Comparing base (0416644) to head (8441b83).

Files with missing lines	Patch %	Lines
src/OptimizeShuffles.cpp	88.88%	1 Missing and 12 partials ⚠️
src/CodeGen_LLVM.cpp	79.41%	5 Missing and 2 partials ⚠️
src/Lower.cpp	33.33%	2 Missing and 2 partials ⚠️
src/CodeGen_C.cpp	82.35%	1 Missing and 2 partials ⚠️
src/CodeGen_WebAssembly.cpp	66.66%	0 Missing and 2 partials ⚠️
src/CodeGen_X86.cpp	91.66%	0 Missing and 2 partials ⚠️
src/CodeGen_ARM.cpp	94.11%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@                       Coverage Diff                        @@
##           alexreinking/arm-field-shift    #9581      +/-   ##
================================================================
+ Coverage                         73.13%   73.21%   +0.08%     
================================================================
  Files                               263      263              
  Lines                             82586    82810     +224     
  Branches                          20195    20272      +77     
================================================================
+ Hits                              60397    60633     +236     
+ Misses                            15894    15867      -27     
- Partials                           6295     6310      +15     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[abadams]

I think instead of the lambda arg and simple_call_intrin this should just use the vector slicing already done by call_intrin

[abadams]

I think that this block of code should be abstracted (or just documented) as can_prove_index_span_safe_to_load(op, ConstantInterval), and should handle pipeline-internal fixed-size allocations in addition to embedded buffers.

[abadams]

It would be great if this could be a lowering pass. If it really does depend on internal codegen state, I think it would clearer if this was a CodeGen_LLVM * pointer that you call the virtual method on, and if that's the case it should live in CodeGen_Internal.cpp/h