perf(encode): add NEON kernel for the Row tag scan#306
Merged
Conversation
Completes the cross-architecture coverage of the vectorized Row tag scan: aarch64 was left on the scalar reference. Add a NEON kernel that compares the row tags against the broadcast tag with vceqq_u8, then folds the 16 per-lane high bits into a 16-bit movemask via the progressive shift-right-accumulate sequence (NEON has no direct movemask). Resolved once per RowMatchGenerator and cached, same as the x86 kernels. A unit test asserts the NEON mask is byte-identical to the scalar reference across row widths 16/32/64 and a spread of tags, so match selection (and the compressed output) is unchanged on aarch64. Correctness validated on aarch64 (96 lib tests + the equivalence test). Throughput is not bench-claimed here (the project's only timing host is x86); the kernel applies the identical algorithmic transform as the x86 SSE2/AVX2 kernels, which measured -6.61% on L5 compress. Part of #247.
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Pull request overview
Adds an AArch64 NEON implementation of the Row match-finder “tag scan” mask build, bringing Row’s SIMD coverage in line with the x86 SSE2/AVX2 work from #305 while keeping scalar as the fallback.
Changes:
- Extend
RowTagKernelwith an AArch64Neonvariant and add std/no-std feature detection for selecting it. - Implement
row_tag_match_mask_neon()usingvceqq_u8+ a NEON movemask emulation to produce the same bitmask as the scalar/x86 paths. - Add an AArch64-only unit test to assert NEON mask output matches the scalar reference across supported row widths and tags.
The movemask packing reinterprets the comparison vector through u16/u32/u64 lanes, which groups bytes by native order — correct only on little-endian. On big-endian aarch64 the lane regrouping would pack the wrong bits and produce an incorrect match mask. Gate the Neon tag-kernel variant, its detection, dispatch arm, implementation, and equivalence test on target_endian = "little"; big-endian aarch64 falls back to the scalar reference (which is byte-order agnostic).
Closed
29 tasks
polaz
added a commit
that referenced
this pull request
Jun 1, 2026
* refactor(encode): reshape row tag-scan into a match bitmask The Row match finder's candidate walk re-read each slot's tag byte and branched on it (one compare + branch per walked slot). Reshape it: compare all row_entries tags against the target tag up front into a u64 bitmask (bit j = tags[j] == tag), then the head-relative walk tests a single mask bit per slot. row_entries clamps to 16/32/64 so the mask fits a u64. Output is byte-identical (same slots, same head-relative order, same max_walk cap and target-len early exit) — this only isolates the tag comparison into a form a SIMD kernel can compute with one vector compare plus movemask. The vector kernels land next. Part of #247. * perf(encode): vectorize the Row match-finder tag scan The Row finder's tag-match mask was built one byte-compare at a time. On x86 compute it with a vector compare + movemask instead: SSE2 (_mm_cmpeq_epi8) over 16-byte chunks, AVX2 (_mm256_cmpeq_epi8) over 32-byte chunks, covering all row widths (16/32/64, all multiples of 16). The kernel is resolved once per RowMatchGenerator (CPU features are fixed per process) and cached, so the per-position row_candidate hot path does no repeated feature query. Non-x86 / pre-SSE2 builds keep the scalar reference (no NEON kernel yet — aarch64 stays on scalar, unchanged). A unit test asserts the SSE2/AVX2 masks are byte-identical to the scalar reference across widths 16/32/64 and a spread of tags, so match selection (and thus the compressed output) is unchanged. row_candidate was ~38% of L5 (greedy/Row) encode self-time on i9, almost all in this scalar tag scan. Part of #247. * test(encode): gate the row tag-mask test on the std feature The SIMD-vs-scalar equivalence test skips kernels the host lacks via std::arch::is_x86_feature_detected!, which is std-only. Add feature = "std" to the test module's cfg so a no-std test build does not try to reference the std-only probe, matching how RowTagKernel::detect gates the same call. * perf(encode): add NEON kernel for the Row tag scan (#306) * perf(encode): add NEON kernel for the Row tag scan Completes the cross-architecture coverage of the vectorized Row tag scan: aarch64 was left on the scalar reference. Add a NEON kernel that compares the row tags against the broadcast tag with vceqq_u8, then folds the 16 per-lane high bits into a 16-bit movemask via the progressive shift-right-accumulate sequence (NEON has no direct movemask). Resolved once per RowMatchGenerator and cached, same as the x86 kernels. A unit test asserts the NEON mask is byte-identical to the scalar reference across row widths 16/32/64 and a spread of tags, so match selection (and the compressed output) is unchanged on aarch64. Correctness validated on aarch64 (96 lib tests + the equivalence test). Throughput is not bench-claimed here (the project's only timing host is x86); the kernel applies the identical algorithmic transform as the x86 SSE2/AVX2 kernels, which measured -6.61% on L5 compress. Part of #247. * fix(encode): gate the NEON row tag-mask on little-endian aarch64 The movemask packing reinterprets the comparison vector through u16/u32/u64 lanes, which groups bytes by native order — correct only on little-endian. On big-endian aarch64 the lane regrouping would pack the wrong bits and produce an incorrect match mask. Gate the Neon tag-kernel variant, its detection, dispatch arm, implementation, and equivalence test on target_endian = "little"; big-endian aarch64 falls back to the scalar reference (which is byte-order agnostic). * fix(encode): no_std RowTagKernel::detect must not dead-code or unreachable The no_std kernel detection used #[cfg]-gated unconditional return blocks. On x86_64 (SSE2 baseline) and aarch64 (NEON baseline) the SSE2/NEON return is always compiled, making the trailing scalar fallback unreachable_code and leaving the AVX2 variant never-constructed dead_code — both hard errors under -D warnings, which broke the no-std CI lint. Resolve the features through if cfg!(target_feature = ...) instead, mirroring the cpu_kernel detect path: cfg! const-folds to the same codegen but keeps the scalar arm reachable and every variant constructed. Also guard RowTagKernel::match_mask with a debug_assert that the row width is 16/32/64 so a future row-width change can't silently produce a wrong mask. Verified: no-std lib check clean (no unreachable/dead_code) on x86_64-unknown-linux-gnu and aarch64; 88/88 row+roundtrip+equivalence tests, clippy + fmt clean. * fix(encode): guard no_std RowTagKernel::detect fallback with allow On baseline-feature targets (x86_64 SSE2, aarch64 NEON) the no_std `if cfg!(target_feature = ...)` can const-fold to an unconditional return, making the trailing scalar fallback unreachable under -D warnings. Add #[allow(unreachable_code)] on the fallback (mirroring cpu_kernel::detect_cpu_kernel) so the trim build stays clean regardless of how cfg! folds. * perf(encode): keep on-the-fly tag check for the scalar row kernel Building the full row bitmask scans every row_entries tag up front, which is free for the SIMD kernels (one vector compare) but does more work than the old per-walk `row_tags[idx] == tag` check on the scalar fallback when search_depth < row_entries (e.g. row_entries=32, search_depth=16). Gate the mask build on a non-scalar kernel and keep the on-the-fly compare for RowTagKernel::Scalar, so scalar targets (non-SIMD ISAs, big-endian aarch64) don't regress. SIMD targets are unchanged (mask + bit test). Output is identical: the mask bit equals the on-the-fly compare.
This was referenced Jun 1, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Stacked on #305 (x86 Row tag-scan vectorization). Completes the cross-architecture coverage: aarch64 was left on the scalar reference. Add a NEON kernel that compares the row tags against the broadcast tag with
vceqq_u8, then folds the 16 per-lane high bits into a 16-bit movemask via the progressive shift-right-accumulate sequence (NEON has no directmovemask). Resolved once perRowMatchGeneratorand cached, same shape as the x86 kernels. The NEON kernel is gated ontarget_endian = "little"— the movemask packing reinterprets the lanes through u16/u32/u64, which is only correct on little-endian; big-endian aarch64 falls back to the scalar reference.Verification
neon_tag_mask_matches_scalarasserts the NEON mask is byte-identical to the scalar reference across row widths 16/32/64 and a spread of tags, so match selection (and the compressed output) is unchanged on aarch64.Throughput is not bench-claimed here: the project's only sanctioned timing host is x86, so this carries no aarch64 perf number. The kernel applies the identical algorithmic transform as the x86 SSE2/AVX2 kernels in #305 (replace up to
row_entriesscalar byte-compares + branches with one vector compare + movemask), which measured −6.61% on L5 compress.Part of #247.