simd: BLAKE3 shuffle surface on U32x16, all six backends - #267
Conversation
BLAKE3's pure-Rust AVX2 backend (rust_avx2.rs) is hash_many at DEGREE = 8, so it is U32x8-shaped rather than U32x16-shaped. Porting it off raw core::arch needs exactly six operations beyond what avx2_int_type! generates: rotate_left, the 32- and 64-bit unpacks, and the two 128-bit half concatenations its transpose network uses. Every body is a plain scalar index loop, deliberately. The codegen oracle measured this exact shape: a fixed two-source permutation written as an index loop compiles to a real packed shuffle (vpermd + vpblendd), and a transpose composed from such primitives emits vpunpcklqdq / vpermq / vinserti128 -- the same instruction family the hand-written intrinsics use. No unsafe, no core::arch, no intrinsic override earned. Semantics are x86's, INCLUDING the per-128-bit-lane split, on every backend. That is not an x86 leak: BLAKE3's transpose is defined in terms of it, so a backend that "helpfully" used whole-vector interleave would compute a different permutation and produce wrong hashes with no compile error. Added on the avx2 arm (where U32x8 is avx2_int_type!-generated) and mirrored on the scalar arm; avx512 re-exports the avx2 type and gets it for free. Only two backends define U32x8 -- neon/wasm/nightly do not, and do not need to, since rust_avx2.rs is x86-gated and blake3 has separate backends for those. Only the two permute2x128 immediates BLAKE3 uses (0x20, 0x31) are exposed, as named methods rather than a generic const-IMM permute: a full 2x128 permute is a much larger surface whose remaining immediates would have no caller and no parity test. Two tests, and the first is the load-bearing one: - u32x8_shuffles_match_x86_intrinsics asserts each method against the REAL intrinsic it claims to reproduce, and the four rotates against upstream's exact `srli | slli` form rather than a reformulation of our own. Gated on target_feature = "avx2". - u32x8_shuffle_surface_is_lane_exact pins the same contract lane-by-lane with no intrinsics, so scalar/avx512/any future arm cannot diverge on a machine where the oracle is unavailable. Control-tested: rewriting interleave_lo_u32 as the "obvious" whole-vector form -- the precise mistake the per-128-lane semantics exist to prevent -- fails BOTH tests. They discriminate; they are not vacuous assertions. Recorded for the port that follows: upstream's rot16/rot8 use the shift-or form deliberately, citing LLVM bug 44379 and a measured preference on recent x86, whereas a scalar rotate_left may fold to vpshufb for byte-granular amounts. That is a codegen difference, not a semantic one -- values are identical either way -- but it is why the port needs a throughput comparison and not only a parity test.
|
Warning Review limit reached
Next review available in: 10 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (7)
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 |
Bugbot couldn't run - usage limit reachedBugbot is counted against Cursor usage for this user or team, and this run hit a usage or spend limit. A user or team admin can review and increase usage limits in the Cursor dashboard. (requestId: serverGenReqId_557c19b9-a31d-4c0f-9a2e-c2a986a821b5) |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5287b6ab9a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let b: [u32; 8] = [20, 21, 22, 23, 24, 25, 26, 27]; | ||
| let (va, vb) = (U32x8::from_array(a), U32x8::from_array(b)); | ||
|
|
||
| assert_eq!(va.interleave_lo_u32(vb).to_array(), [10, 20, 11, 21, 14, 24, 15, 25]); |
There was a problem hiding this comment.
Mirror the U32x8 methods in the nightly backend
When nightly-simd is enabled, crate::simd::U32x8 is re-exported from src/simd_nightly/u_word_types.rs, whose U32x8 implementation has none of the newly added shuffle methods or rotate_left. Consequently this unconditional test—and any downstream BLAKE3 port using the new surface—fails to compile with method-not-found errors under that feature, even though nightly-simd is intended to replace the other backends. Add the same surface to the portable-SIMD implementation rather than only the AVX2 and scalar types.
Useful? React with 👍 / 👎.
Supersedes this branch's first commit, which put the surface on U32x8. That was wrong and the operator called it: two __m256i fit in one U32x16. BLAKE3's AVX2 backend is DEGREE 8 over __m256i, and I treated that as requiring a matching 8-wide type. It does not. Every operation in hash_many is either lane-wise (add / xor / rotate) or confined WITHIN a 128- or 256-bit lane -- the unpacks are per-128, permute2x128 is per-256 -- so two 8-lane groups packed into one U32x16 never interact. The algorithm runs on both at once as DEGREE 16 with no cross-talk. This is also what the 2026-07-28 ruling already said: a half-width type standing in for the lane the substrate actually uses is an absolute no-go. Practical consequences of being on the right type: - rotate_left already exists on U32x16 on every backend, so only the six shuffles were missing rather than seven ops. - U32x8 is absent from neon, wasm and nightly; U32x16 is on all six. The surface is now uniform, which also answers codex's P1 on #267 (the nightly arm re-exports its own U32x8 without the added methods, so the previous commit's unconditional test could not compile under nightly-simd). - Every backend's U32x16 has to_array/from_array, so ONE identical body serves all six -- no per-backend storage access. Semantics are x86's applied to each 256-bit half independently: interleave_{lo,hi}_u{32,64} reproduce _mm256_unpack{lo,hi}_epi{32,64} within each 128-bit quad, concat_{lo,hi}_halves reproduce _mm256_permute2x128_si256(_, _, 0x20 / 0x31) within each 256-bit half. The per-lane structure is load-bearing: BLAKE3's transpose is defined in terms of it, so a whole-vector interleave would compute a different permutation and produce wrong hashes with no compile error. Bodies are plain index loops, per the oracle's measurement that fixed two-source permutations written that way compile to real packed shuffles. No unsafe, no core::arch, no intrinsic override earned. Two tests. u32x16_blake3_shuffles_match_x86_intrinsics_per_half asserts each half of the U32x16 result against the 256-bit intrinsic applied to that half's inputs, and the four rotates against upstream's exact `srli | slli` form. u32x16_blake3_shuffles_are_lane_exact pins the same permutation with no intrinsics so scalar/avx512/neon/wasm/nightly are held to it where the x86 oracle cannot run. Control-tested: rewriting interleave_lo_u32 as a whole-vector interleave -- precisely the mistake the quad semantics prevent -- fails both tests. 31 simd tests pass; clippy clean; fmt clean.
|
@codex — your P1 is fixed in You were correct that the nightly arm re-exports its own BLAKE3's AVX2 backend is That resolves the P1 structurally rather than by addition:
Semantics are x86's applied to each 256-bit half independently, and the tests assert exactly that: each half of the 31 Generated by Claude Code |
|
To use Codex here, create an environment for this repo. |
Step 1 of replacing BLAKE3's inline AVX2 with the ndarray polyfill.
The mistake, and why it was one
BLAKE3's AVX2 backend is
hash_manyatDEGREE = 8over__m256i, and I treated that as requiring a matching 8-wide type. It doesn't. Every operation inhash_manyis either lane-wise (add / xor / rotate) or confined within a 128- or 256-bit lane — the unpacks are per-128,permute2x128is per-256. So two 8-lane groups packed into oneU32x16never interact, and the algorithm runs on both at once asDEGREE = 16with no cross-talk.Which is what the 2026-07-28 ruling already said: a half-width type standing in for the lane the substrate actually uses is an absolute no-go. I had it recorded and still reached for
U32x8.Being on the right type changes three practical things:
rotate_leftalready exists onU32x16on every backend, so only six shuffles were missing, not seven ops.U32x8is absent from neon, wasm and nightly;U32x16is on all six. The surface is uniform — which also resolves Codex's P1 on the previous commit: the nightly arm re-exports its ownU32x8without the added methods, so the unconditional test could not compile undernightly-simd.U32x16hasto_array/from_array, so one identical body serves all six with no per-backend storage access.What's added
Six methods on
U32x16, in all six backends (neon's gated#[cfg(target_arch = "aarch64")]to match its type):interleave_lo_u32/interleave_hi_u32_mm256_unpack{lo,hi}_epi32within each 128-bit quadinterleave_lo_u64/interleave_hi_u64_mm256_unpack{lo,hi}_epi64within each 128-bit quadconcat_lo_halves/concat_hi_halves_mm256_permute2x128_si256(_, _, 0x20 / 0x31)Bodies are plain index loops, per the oracle's measurement that fixed two-source permutations written that way compile to real packed shuffles (
vpermd/vpblendd; composed,vpunpcklqdq/vpermq/vinserti128). Nounsafe, nocore::arch, no intrinsic override earned.Only the two
permute2x128immediates BLAKE3 uses are exposed, as named methods rather than a genericconst IMMpermute — the rest would have no caller and no parity test.The per-lane structure is load-bearing
A whole-vector interleave would look reasonable and compute a different permutation, producing wrong hashes with no compile error. So both tests exist to catch exactly that:
u32x16_blake3_shuffles_match_x86_intrinsics_per_half— asserts each half of theU32x16result against the 256-bit intrinsic applied to that half's inputs, and the four rotates against upstream's exactsrli | slliform rather than a reformulation of mine.u32x16_blake3_shuffles_are_lane_exact— pins the same permutation with no intrinsics, so scalar/avx512/neon/wasm/nightly are held to it where the x86 oracle can't run.Control-tested. Rewriting
interleave_lo_u32as a whole-vector interleave fails both:31
simd::tests pass; clippy and fmt clean.Known blocker for the follow-up, not for this PR
The
rust_avx2.rsrewrite itself hits a cargo cycle: rootndarraydepends onblake3(11 files undersrc/hpc/), soblake3 → ndarray::simdcloses the loop. Cargo names it and silently falls back to registry blake3 with[[patch.unused]]— the port would compile and have zero effect. Breaking that edge is a separate architectural decision. This PR stands on its own regardless of how it's resolved.Also noted: upstream's
rot16/rot8use the shift-or form deliberately (LLVM bug 44379), while a scalarrotate_leftmay fold tovpshufb. Same values, different codegen — which is why the port needs a throughput comparison and not only a parity test. Still unrun.