perf(word-cache): scan the tag window with shrn, not an emulated movemask - #2377
Open
ArthurZucker wants to merge 3 commits into
Open
perf(word-cache): scan the tag window with shrn, not an emulated movemask#2377ArthurZucker wants to merge 3 commits into
ArthurZucker wants to merge 3 commits into
Conversation
The module and `LookupKey` docs described a 128-bit key: a tag bit at 127, a 7-bit length, 120 bits of word, "127 bits of hash", and a correctness guarantee for words up to 15 bytes. `LookupKey` is a `u64`. The guarantee is real but the bound is `INLINE_KEY_BYTES` (7), not 15, and the hash arm carries 64 bits, not 127. The `WordCacheSlot` diagrams were wrong the same way: they placed the payload at offset 16 and stopped at 29. It is at 8, `ids_len` is at 20, and 11 of the 32 bytes are padding -- which is exactly what widening the key to a `u128` would consume, so that option is noted where the bound is. Comments only; no generated code changes.
`probe_emit_keyed` built the full `InsertPlacement` before testing the home
slot, but the fast path only needs the home index and the key -- `tag` is
read by `lookup_placed` alone. `InsertPlacement` is 24 bytes and AAPCS
passes anything over 16 indirectly, so the whole thing was materialized on
the stack on every cache hit for the benefit of the miss path:
stp x8, x0, [sp, #160] <- every hit
strb w9, [sp, #176] <- every hit
ldr x9, [x8] ; cmp x9, x0 ; b.ne <miss>
Splitting `placement_from` into `home_index` and `tag_of` lets the hit path
compute only what it reads and the cold branch assemble the rest. `stp`
27 -> 22 and `strb` 9 -> 6 across `tokenize_spans`.
On its own this measures as nothing (-0.1% to -0.7%, inside noise), and the
reason is worth writing down: the placement still has to be materialized
for the *call* to `lookup_placed`, so the stores move rather than
disappear. It pays only once the next commit shrinks that function enough
for LLVM to inline it. Incremental gain on top of that commit, in points:
+1.1 llama-3/code, +0.9 llama-3/english, +2.3 gpt2/code, +0.8 gpt2/english.
The two belong together.
`find_matches_and_first_empty` took two `wide::i8x16::to_bitmask()` calls.
On aarch64 that is not a movemask -- NEON has no `pmovmskb`, so `wide`
emulates one per call with a 16-byte constant load, an `and`, a `tbl`, a
cross-lane `addv` and a SIMD->GPR `fmov`:
cmlt v1, v0, #0 ; and v1, v1, v2 ; tbl v1, {v1}, v3 ; addv h1, v1 ; fmov w9, s1
Both results are on the critical path before the candidate loop can start,
and `addv` is a cross-lane reduction feeding a cross-domain move -- the two
shapes to avoid on Apple cores.
`shrn` narrows the 16 comparison lanes to 16 nibbles in one instruction, so
the mask becomes a `u64` of nibbles and a lane index is
`trailing_zeros() / 4`. `LANE_BITS`/`pop_lowest_lane` keep the one-bit and
four-bit representations interchangeable, and non-NEON targets keep the
`wide` path.
`tbl`, `addv` and the constant `and` all drop to zero and `ldr` goes
115 -> 100. `lookup_placed` also gets small enough that LLVM now inlines
it, so its `bl` and its indirect 24-byte struct argument go too:
573 -> 546 instructions across `tokenize_spans` + `lookup_placed`.
Min of 8 full-corpus passes, 5 interleaved rounds, token ids byte-identical
against the reference throughout:
alone with the previous commit
llama-3/code +5.1% +6.2%
llama-3/english +4.1% +5.0%
gpt2/code +5.2% +7.5%
gpt2/english +4.2% +5.0%
gpt2/chinese +1.4% +1.3%
llama-3/chinese +0.2% +0.4%
CJK is flat by construction: a 3-byte-per-char pretoken is past
`INLINE_KEY_BYTES`, takes the hashed arm, and barely reaches the window.
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
1 participant
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.
NEON has no
pmovmskb, sowide::to_bitmaskemulates one withtblplus across-lane
addv— and the window scan did it twice per lookup.shrnreplacesboth. +4–7.5% on code/english, flat on CJK. Token ids byte-identical.
Measured
Min of 8 full-corpus passes, 5 interleaved rounds (separately built binaries, rotating
order),
ids_hashidentical across every combination:CJK is flat by construction: a 3-byte-per-char pretoken is past
INLINE_KEY_BYTES,takes the hashed arm, and barely reaches the tag window.
The placement commit is worth nothing on its own and only pays in combination — it
elides stores that were unavoidable while
lookup_placedwas still an out-of-linecall.
shrnshrinks that function enough for LLVM to inline it, and then they go.Codegen
tokenize_spans+lookup_placed, 573 → 546 instructions:tbl.16baddv.8hand.16b(selectbit)shrn.8bldrstp/strbNotes
widepath;LANE_BITS/pop_lowest_lanekeep theone-bit and four-bit masks interchangeable. Verified by running the
word_cachetests under
--target x86_64-apple-darwinas well.LookupKeyis au64, but the module and type docs described a 128-bit key with a tag bit, a15-byte exactness bound and "127 bits of hash", and the slot diagrams put the
payload at offset 16 when it is at 8.