gigatoken 0.1.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (75) hide show
  1. checksums.yaml +7 -0
  2. data/Cargo.lock +3016 -0
  3. data/Cargo.toml +135 -0
  4. data/LICENSE +21 -0
  5. data/README.md +141 -0
  6. data/exe/gigatoken +9 -0
  7. data/ext/gigatoken/Cargo.toml +24 -0
  8. data/ext/gigatoken/extconf.rb +19 -0
  9. data/ext/gigatoken/src/error.rs +20 -0
  10. data/ext/gigatoken/src/gvl.rs +122 -0
  11. data/ext/gigatoken/src/lib.rs +50 -0
  12. data/ext/gigatoken/src/sentencepiece.rs +205 -0
  13. data/ext/gigatoken/src/sources.rs +207 -0
  14. data/ext/gigatoken/src/tokenizer.rs +571 -0
  15. data/lib/gigatoken/cli/bench.rb +64 -0
  16. data/lib/gigatoken/cli/support.rb +132 -0
  17. data/lib/gigatoken/cli/validate.rb +55 -0
  18. data/lib/gigatoken/cli.rb +18 -0
  19. data/lib/gigatoken/hub.rb +242 -0
  20. data/lib/gigatoken/packed_result.rb +48 -0
  21. data/lib/gigatoken/tokenizer.rb +117 -0
  22. data/lib/gigatoken/version.rb +5 -0
  23. data/lib/gigatoken.rb +29 -0
  24. data/rust-toolchain.toml +8 -0
  25. data/src/batch.rs +1808 -0
  26. data/src/bindings/bridge.rs +396 -0
  27. data/src/bindings/hub.rs +42 -0
  28. data/src/bindings/matcher.rs +114 -0
  29. data/src/bindings/mod.rs +14 -0
  30. data/src/bindings/padding.rs +177 -0
  31. data/src/bindings/pretokenize.rs +53 -0
  32. data/src/bindings/sources.rs +273 -0
  33. data/src/bindings/train.rs +125 -0
  34. data/src/bpe/mod.rs +1217 -0
  35. data/src/bpe/pretoken_cache.rs +495 -0
  36. data/src/bpe/sentencepiece.rs +1485 -0
  37. data/src/bpe/tiktoken.rs +2555 -0
  38. data/src/bpe_train.rs +351 -0
  39. data/src/input/decompress.rs +11 -0
  40. data/src/input/file_source.rs +514 -0
  41. data/src/input/jsonl.rs +94 -0
  42. data/src/input/mod.rs +333 -0
  43. data/src/input/parquet.rs +303 -0
  44. data/src/lib.rs +578 -0
  45. data/src/load_tokenizer/hf.rs +1036 -0
  46. data/src/load_tokenizer/hub.rs +344 -0
  47. data/src/load_tokenizer/mod.rs +3 -0
  48. data/src/load_tokenizer/tiktoken.rs +87 -0
  49. data/src/main.rs +95 -0
  50. data/src/pretokenize/fast/cl100k.rs +426 -0
  51. data/src/pretokenize/fast/cl100k_family.rs +891 -0
  52. data/src/pretokenize/fast/deepseek_v3.rs +605 -0
  53. data/src/pretokenize/fast/kimi.rs +281 -0
  54. data/src/pretokenize/fast/mask.rs +1486 -0
  55. data/src/pretokenize/fast/mod.rs +446 -0
  56. data/src/pretokenize/fast/nemotron.rs +138 -0
  57. data/src/pretokenize/fast/o200k.rs +347 -0
  58. data/src/pretokenize/fast/o200k_family.rs +1734 -0
  59. data/src/pretokenize/fast/olmo3.rs +505 -0
  60. data/src/pretokenize/fast/qwen2.rs +429 -0
  61. data/src/pretokenize/fast/qwen3_5.rs +541 -0
  62. data/src/pretokenize/fast/r50k.rs +1250 -0
  63. data/src/pretokenize/mod.rs +1079 -0
  64. data/src/pretokenize/options.rs +188 -0
  65. data/src/pretokenize/pretoken.rs +20 -0
  66. data/src/pretokenize/pretokenize_traits.rs +49 -0
  67. data/src/pretokenize/reference/avx512.rs +522 -0
  68. data/src/pretokenize/reference/combinator.rs +572 -0
  69. data/src/pretokenize/reference/mod.rs +28 -0
  70. data/src/pretokenize/reference/simd.rs +852 -0
  71. data/src/pretokenize/reference/state_machine.rs +365 -0
  72. data/src/pretokenize/unicode.rs +546 -0
  73. data/src/test_hub.rs +28 -0
  74. data/src/token.rs +42 -0
  75. metadata +161 -0
@@ -0,0 +1,1250 @@
1
+ //! Fast pretokenizer for the GPT-2 (r50k_base) regex:
2
+ //! `'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+`
3
+ //!
4
+ //! On aarch64 (NEON) and x86_64 with AVX-512 or AVX2 (runtime-detected)
5
+ //! the iterator runs a simdjson-style mask scanner: 64-byte
6
+ //! batches are classified with SIMD into per-byte u64 class masks, the
7
+ //! token-boundary bits are derived with shifted-mask algebra in scalar
8
+ //! registers (log step 17; the original vector-register algebra of step
9
+ //! 15 measured the same and was retired for the simpler form), and the
10
+ //! walker pops one bit per token — no per-token dispatch branches, which
11
+ //! sidesteps the ~8 cy/token branch-miss floor of the scalar scanner
12
+ //! (log step 13). Apostrophes get a contraction bit-fixup; batches with
13
+ //! any non-ASCII (~21% of OWT) take `extended_masks`, which classifies
14
+ //! every unicode char with the packed table so it joins the same
15
+ //! algebra. Chars straddling a batch edge are resolved with lookahead
16
+ //! and a prev-char walk-back, so bad zones (scalar re-derivation) remain
17
+ //! only for edge-straddling whitespace, contractions at the batch edge,
18
+ //! and invalid UTF-8 — ~0.4% of batches. Measured 2,460-2,600 MB/s on
19
+ //! 1 GB OWT (pretokenize_profile, min-of-N interleaved; 2,132 at step
20
+ //! 15, 983 for the scalar scanner).
21
+ //!
22
+ //! The scalar path (`advance_pos`, SWAR letter runs + arithmetic
23
+ //! predicates) remains the reference implementation, the no-SIMD
24
+ //! fallback, and the executor for bad zones and buffer tails.
25
+ //!
26
+ //! `advance_pos` is a pure free function (`(bytes, pos) -> end`) rather than
27
+ //! a `&mut self` method: keeping the cursor in a register instead of writing
28
+ //! `self.pos` at every scan step shortens the per-token dependency chain and
29
+ //! is worth ~30% throughput. Non-ASCII characters are classified with one
30
+ //! packed-table load (`unicode::class_of`) on a hand-rolled UTF-8 decode.
31
+ //!
32
+ //! A windowed multi-cursor variant (finding guaranteed boundaries and running
33
+ //! 2-4 independent `advance_pos` chains interleaved, queueing token ends) was
34
+ //! benchmarked at 0.80-0.95x of this streaming version: the queue traffic and
35
+ //! interleaved branch history cost more than the extra ILP recovers.
36
+
37
+ use super::mask::{self, MaskScheme, MaskState};
38
+ use super::{
39
+ decode_cp, is_ascii_ws, is_digit, is_letter, scan_digits_from, scan_letters_from,
40
+ scan_other_from,
41
+ };
42
+ use crate::pretokenize::unicode::{self, CharClass};
43
+ use crate::pretokenize::Pretoken;
44
+
45
+ // -----------------------------------------------------------------------
46
+ // FastR50kPretokenizer
47
+ // -----------------------------------------------------------------------
48
+
49
+ /// Boundary and bad-zone bitmasks for `bytes[scan..scan+64]` (requires
50
+ /// `scan + 64 <= bytes.len()`). Bit `k` of `usable` = a trustworthy token
51
+ /// start at `scan + k`; `bad` marks bytes whose boundaries must be
52
+ /// re-derived by `advance_pos`, and no token may be emitted across an
53
+ /// unresolved bad zone.
54
+ ///
55
+ /// NEON classifies the ASCII classes (letter, digit, space, whitespace:
56
+ /// 4 movemasks; apostrophe and non-ASCII behind horizontal any-tests)
57
+ /// and the boundary bits come from u64 shifted-mask algebra in scalar
58
+ /// registers, as in `cl100k_family::batch_masks`. Batches with any
59
+ /// non-ASCII byte (~21% on OWT, mostly curly quotes) take
60
+ /// [`extended_masks`]. Inlining that path here measured 0.98x, and
61
+ /// `#[inline(never)]` on this whole function 0.94x — the split keeps the
62
+ /// walker's register allocation clean (step 15's lesson) while the hot
63
+ /// ASCII algebra stays inline.
64
+ #[cfg(target_arch = "aarch64")]
65
+ #[inline(always)]
66
+ fn batch_masks(bytes: &[u8], scan: usize) -> (u64, u64) {
67
+ use std::arch::aarch64::*;
68
+ let len = bytes.len();
69
+ if scan + 70 > len {
70
+ // Not enough lookahead for the batch-edge char classification
71
+ // (up to a 4-byte char starting at scan + 66); scalar batch.
72
+ return (0, u64::MAX);
73
+ }
74
+ unsafe {
75
+ let p = bytes.as_ptr().add(scan);
76
+ let zero = vdupq_n_u8(0);
77
+ let mut lv = [zero; 4];
78
+ let mut dv = [zero; 4];
79
+ let mut sv = [zero; 4];
80
+ let mut wsv = [zero; 4];
81
+ let mut hiv = [zero; 4];
82
+ let mut apv = [zero; 4];
83
+ for i in 0..4 {
84
+ let v = vld1q_u8(p.add(16 * i));
85
+ let lowered = vorrq_u8(v, vdupq_n_u8(0x20));
86
+ lv[i] = vcleq_u8(vsubq_u8(lowered, vdupq_n_u8(b'a')), vdupq_n_u8(25));
87
+ dv[i] = vcleq_u8(vsubq_u8(v, vdupq_n_u8(b'0')), vdupq_n_u8(9));
88
+ sv[i] = vceqq_u8(v, vdupq_n_u8(b' '));
89
+ wsv[i] = vorrq_u8(
90
+ sv[i],
91
+ vcleq_u8(vsubq_u8(v, vdupq_n_u8(9)), vdupq_n_u8(4)),
92
+ );
93
+ hiv[i] = vcltzq_s8(vreinterpretq_s8_u8(v));
94
+ apv[i] = vceqq_u8(v, vdupq_n_u8(b'\''));
95
+ }
96
+
97
+ let lb = mask::movemask64(lv[0], lv[1], lv[2], lv[3]);
98
+ let db = mask::movemask64(dv[0], dv[1], dv[2], dv[3]);
99
+ let s64 = mask::movemask64(sv[0], sv[1], sv[2], sv[3]);
100
+ let wsa = mask::movemask64(wsv[0], wsv[1], wsv[2], wsv[3]);
101
+ // Apostrophes only matter for the contraction fixup below.
102
+ let ap_any = vorrq_u8(vorrq_u8(apv[0], apv[1]), vorrq_u8(apv[2], apv[3]));
103
+ let ap64 = if vmaxvq_u8(ap_any) != 0 {
104
+ mask::movemask64(apv[0], apv[1], apv[2], apv[3])
105
+ } else {
106
+ 0
107
+ };
108
+
109
+ // Any non-ASCII byte routes to the extended classifier, which
110
+ // reuses the ASCII masks computed above.
111
+ let hi_any = vorrq_u8(vorrq_u8(hiv[0], hiv[1]), vorrq_u8(hiv[2], hiv[3]));
112
+ if vmaxvq_u8(hi_any) != 0 {
113
+ let hi64 = mask::movemask64(hiv[0], hiv[1], hiv[2], hiv[3]);
114
+ return extended_masks(bytes, scan, lb, db, s64, wsa, hi64, ap64);
115
+ }
116
+
117
+ ascii_batch_algebra(bytes, scan, lb, db, s64, wsa, ap64)
118
+ }
119
+ }
120
+
121
+ /// x86 counterpart of the NEON `batch_masks`, monomorphized on the SIMD
122
+ /// tier: the classification is [`mask::ascii_masks_avx512`] (one 64-byte
123
+ /// load and one k-register compare per class) or
124
+ /// [`mask::ascii_masks_avx2`] (two 32-byte loads, compare + vpmovmskb per
125
+ /// class); the boundary algebra and the extended (non-ASCII) path are the
126
+ /// same shared scalar code. `#[inline(always)]` (with no `target_feature`
127
+ /// of its own) so the body fuses into whichever feature region calls it —
128
+ /// the tier-monomorphized fill wrappers
129
+ /// (`MaskState::fill_spans_two_phase_{avx512,avx2}_crc`), where LLVM's
130
+ /// cost model declined to inline the previous `#[target_feature]` form
131
+ /// and left a call per 64-byte batch, or the runtime-dispatched
132
+ /// `MaskScheme::batch_masks` that `next_span` uses.
133
+ ///
134
+ /// # Safety
135
+ ///
136
+ /// The selected tier must have been runtime-detected
137
+ /// ([`mask::avx512_scanner_available`] /
138
+ /// [`mask::avx2_scanner_available`]).
139
+ #[cfg(target_arch = "x86_64")]
140
+ #[inline(always)]
141
+ unsafe fn batch_masks_x86<const AVX512: bool>(bytes: &[u8], scan: usize) -> (u64, u64) {
142
+ let len = bytes.len();
143
+ if scan + 70 > len {
144
+ // Not enough lookahead for the batch-edge char classification
145
+ // (up to a 4-byte char starting at scan + 66); scalar batch.
146
+ return (0, u64::MAX);
147
+ }
148
+ let am = if AVX512 {
149
+ // SAFETY: the caller detected the AVX-512 tier (fn contract).
150
+ unsafe { mask::ascii_masks_avx512(bytes, scan) }
151
+ } else {
152
+ // SAFETY: the caller detected the AVX2 tier (fn contract).
153
+ unsafe { mask::ascii_masks_avx2(bytes, scan) }
154
+ };
155
+ let wsa = am.s | am.wt | am.n;
156
+ if am.hi != 0 {
157
+ // SAFETY: both detected tiers include the BMI1/BMI2/LZCNT/POPCNT
158
+ // bit features `extended_masks` re-declares (fn contract).
159
+ return unsafe { extended_masks(bytes, scan, am.l, am.d, am.s, wsa, am.hi, am.ap) };
160
+ }
161
+ ascii_batch_algebra(bytes, scan, am.l, am.d, am.s, wsa, am.ap)
162
+ }
163
+
164
+ /// Pure-ASCII boundary algebra shared by the NEON and AVX-512 batch
165
+ /// classifiers (the batch has no non-ASCII byte; `wsa` = all ASCII
166
+ /// whitespace). Everything here is platform-independent u64 bit math.
167
+ #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
168
+ #[inline(always)]
169
+ fn ascii_batch_algebra(
170
+ bytes: &[u8],
171
+ scan: usize,
172
+ lb: u64,
173
+ db: u64,
174
+ s64: u64,
175
+ wsa: u64,
176
+ ap64: u64,
177
+ ) -> (u64, u64) {
178
+ let ob = !(lb | db | wsa); // hi == 0 on this path
179
+
180
+ // Bit-0 carries from the char before the batch. This batch is
181
+ // pure ASCII, so a multi-byte prev char always ends exactly at
182
+ // the boundary and the walk-back gives true carries — no bad
183
+ // zone.
184
+ let (pl, pd, ps, pws, po) = if scan == 0 {
185
+ (0, 0, 0, 0, 0)
186
+ } else {
187
+ carries_at(bytes, scan)
188
+ };
189
+
190
+ let cont_same =
191
+ (lb & ((lb << 1) | pl)) | (db & ((db << 1) | pd)) | (ob & ((ob << 1) | po));
192
+ let after_sp = (s64 << 1) | ps;
193
+ let nb = !wsa & !cont_same & !after_sp;
194
+
195
+ // Ws-run split (`\s+(?!\S)`); bit 63 needs the real lookahead
196
+ // char. The ASCII case is branchless — "is byte 63 ws" is a
197
+ // ~20% coin flip on natural text, so testing it costs a
198
+ // mispredict every few batches. Only a non-ASCII lookahead
199
+ // byte (rare) branches, for the table-backed ws check.
200
+ let mut split_ok = wsa & (!wsa >> 1); // bit 63: shifted-in 0
201
+ let nb64 = bytes[scan + 64]; // in bounds: scan + 70 <= len
202
+ if nb64 < 0x80 {
203
+ split_ok |= (u64::from(!is_ascii_ws(nb64)) << 63) & wsa;
204
+ } else if wsa >> 63 != 0
205
+ // SAFETY: this classifier's scan + 70 <= len batch guard puts the
206
+ // decode at scan + 64 in bounds (needs scan + 68 <= len).
207
+ && unsafe { mask::nn_at_full(bytes, scan + 64) }
208
+ {
209
+ split_ok |= 1 << 63;
210
+ }
211
+ let pwsb = (wsa << 1) | pws;
212
+ let wsboundary = wsa & (!pwsb | split_ok);
213
+ let mut boundary = nb | wsboundary;
214
+
215
+ let mut bad = 0u64;
216
+
217
+ // Contraction fixup (see extended_masks for the rules).
218
+ if ap64 != 0 {
219
+ let mut cand = ap64 & boundary;
220
+ while cand != 0 {
221
+ let i = cand.trailing_zeros() as usize;
222
+ cand &= cand - 1;
223
+ if i >= 61 {
224
+ bad |= u64::MAX << i;
225
+ break;
226
+ }
227
+ let k = match bytes[scan + i + 1] {
228
+ b's' | b'd' | b'm' | b't' => 2,
229
+ b'l' if bytes[scan + i + 2] == b'l' => 3,
230
+ b'v' if bytes[scan + i + 2] == b'e' => 3,
231
+ b'r' if bytes[scan + i + 2] == b'e' => 3,
232
+ _ => 0,
233
+ };
234
+ if k != 0 {
235
+ boundary &= !(1u64 << (i + 1));
236
+ boundary |= 1u64 << (i + k);
237
+ }
238
+ }
239
+ }
240
+ (boundary & !bad, bad)
241
+ }
242
+
243
+ /// Slow(er) path for batches containing non-ASCII: every unicode char in
244
+ /// (or straddling into/out of) the batch is classified with the packed
245
+ /// table via [`mask::classify_uni_chars`] — the same lookup the scalar
246
+ /// path would do — and joins the per-byte effective class masks, so
247
+ /// byte-adjacency == char-adjacency and the u64 boundary algebra applies
248
+ /// unchanged. Takes the ASCII class masks the caller already computed
249
+ /// (`ws64` = all ASCII whitespace).
250
+ ///
251
+ /// Bad zones remain only for whitespace chars straddling a batch edge
252
+ /// (their `\s+(?!\S)` bookkeeping crosses the boundary), stray
253
+ /// continuation bytes (invalid UTF-8), and contractions at the batch
254
+ /// edge. An earlier version pattern-matched common leads with ~95 vector
255
+ /// ops (`mask::unicode_leads`) before falling back to per-char bad
256
+ /// zones; the direct table loop (typical hi batch: 1-3 unicode chars,
257
+ /// table hot in cache) plus edge-char resolution was worth ~13% end to
258
+ /// end. `#[inline(never)]`: inlining this into the walker wrecks the
259
+ /// clean path's register allocation (step 15).
260
+ ///
261
+ /// On x86_64 the `target_feature` re-declaration keeps the bit-scan
262
+ /// loops on tzcnt/lzcnt/blsr in a baseline (non-native) build — this
263
+ /// function is out-of-line, so without it the ~21% of OWT batches
264
+ /// landing here would compile against baseline x86-64 even though every
265
+ /// caller is a SIMD batch classifier. Only the bit features are enabled
266
+ /// (not the callers' vector features): both the AVX-512 and AVX2 tiers
267
+ /// call this, so it must never emit instructions beyond the AVX2 tier's
268
+ /// set. Measured neutral on the OWT mask-compute diagnostic
269
+ /// (ASCII-dominated); kept for codegen parity on non-ASCII-heavy corpora
270
+ /// where this path dominates.
271
+ #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
272
+ #[cfg_attr(
273
+ target_arch = "x86_64",
274
+ target_feature(enable = "bmi1,bmi2,lzcnt,popcnt")
275
+ )]
276
+ #[inline(never)]
277
+ #[allow(clippy::too_many_arguments)]
278
+ fn extended_masks(
279
+ bytes: &[u8],
280
+ scan: usize,
281
+ l64: u64,
282
+ d64: u64,
283
+ s64: u64,
284
+ ws64: u64,
285
+ hi64: u64,
286
+ ap64: u64,
287
+ ) -> (u64, u64) {
288
+ let wsa = ws64;
289
+
290
+ // Class-table LazyLock resolved once; the per-char classify below is
291
+ // then a bare slice index.
292
+ let ct = unicode::ClassTable::get();
293
+ let class = move |cp| ct.class_of(cp);
294
+
295
+ // Bit-0 carries via the prev-char walk-back; a char straddling into
296
+ // this batch claims its continuation bytes with its class. Without
297
+ // this, every batch following a unicode char became a bad zone at
298
+ // bit 0, and a bad zone costs ~800 cycles in walker re-entries and
299
+ // cold scalar gaps.
300
+ let mut claim = mask::UniClasses::default();
301
+ let (pl, pd, ps, pws, po) = if scan == 0 {
302
+ (0, 0, 0, 0, 0)
303
+ } else if bytes[scan - 1] < 0x80 {
304
+ carries_at(bytes, scan)
305
+ } else {
306
+ // SAFETY: scan > 0 on this branch, and the classifier's
307
+ // scan + 70 <= len batch guard covers pos + 3 <= len.
308
+ let (cls, _lead, end) = unsafe { mask::char_through(bytes, scan, class) };
309
+ let chm = if end > scan { (1u64 << (end - scan)) - 1 } else { 0 };
310
+ claim.cont = chm;
311
+ match cls {
312
+ CharClass::Letter => {
313
+ claim.l = chm;
314
+ (1, 0, 0, 0, 0)
315
+ }
316
+ CharClass::Number => {
317
+ claim.n = chm;
318
+ (0, 1, 0, 0, 0)
319
+ }
320
+ CharClass::Other => {
321
+ claim.o = chm;
322
+ (0, 0, 0, 0, 1)
323
+ }
324
+ CharClass::Whitespace => {
325
+ // A ws char straddling in defers to the scalar path (its
326
+ // run-split bookkeeping needs the pre-batch extent) but
327
+ // still marks its true class for neighbors' algebra.
328
+ claim.ws = chm;
329
+ claim.resid = chm;
330
+ (0, 0, u64::from(bytes[scan - 1] == b' '), 1, 0)
331
+ }
332
+ }
333
+ };
334
+
335
+ // SAFETY: this classifier's scan + 70 <= len batch guard is exactly
336
+ // `classify_uni_chars`' contract.
337
+ let uni = unsafe {
338
+ mask::classify_uni_chars::<true, false>(bytes, scan, hi64 & !claim.cont, class)
339
+ };
340
+
341
+ // Effective per-byte classes: every byte of a classified char carries
342
+ // the char's class, so the same algebra as the pure-ASCII path
343
+ // applies.
344
+ let lb = l64 | claim.l | uni.l;
345
+ let db = d64 | claim.n | uni.n;
346
+ let wsb = wsa | claim.ws | uni.ws;
347
+ let ob = !(l64 | d64 | wsa | hi64) | claim.o | uni.o;
348
+ let contm = claim.cont | uni.cont;
349
+ let resid = claim.resid | uni.resid;
350
+
351
+ let cont_same =
352
+ (lb & ((lb << 1) | pl)) | (db & ((db << 1) | pd)) | (ob & ((ob << 1) | po));
353
+ let after_sp = (s64 << 1) | ps;
354
+ let nb = !wsb & !cont_same & !after_sp & !contm;
355
+
356
+ // Ws-run split: char-length-aware "followed by non-ws" test. All ws
357
+ // chars whose lookahead crosses the batch edge look at byte 64: an
358
+ // ASCII ws at 63, a 2-byte ws led at 62, a 3-byte ws led at 61
359
+ // (later leads straddle out and are already bad zones). The ASCII
360
+ // case is branchless as in the fast path; multi-byte edge leads are
361
+ // rare enough to branch.
362
+ let nn = !wsb;
363
+ let mut split_ok = (wsa & (nn >> 1)) | (uni.w2 & (nn >> 2)) | (uni.w3 & (nn >> 3));
364
+ let ws_leads = wsa | uni.w2 | uni.w3;
365
+ let edge_mb = (uni.w2 & (1 << 62)) | (uni.w3 & (1 << 61));
366
+ let nb64 = bytes[scan + 64]; // in bounds: scan + 70 <= len
367
+ if nb64 < 0x80 && edge_mb == 0 {
368
+ split_ok = (split_ok & !(1 << 63)) | ((u64::from(!is_ascii_ws(nb64)) << 63) & wsa);
369
+ } else {
370
+ let edge = edge_mb | ((1 << 63) & wsa);
371
+ if edge != 0 {
372
+ // SAFETY: this classifier's scan + 70 <= len batch guard puts
373
+ // the decode at scan + 64 in bounds (needs scan + 68 <= len).
374
+ if unsafe { mask::nn_at_full(bytes, scan + 64) } {
375
+ split_ok |= edge;
376
+ } else {
377
+ split_ok &= !edge;
378
+ }
379
+ }
380
+ }
381
+ let pwsb = (wsb << 1) | pws;
382
+ let wsboundary = ws_leads & (!pwsb | split_ok);
383
+ let mut boundary = nb | wsboundary;
384
+
385
+ let mut bad = resid | resid << 1 | resid >> 1;
386
+
387
+ // Contraction fixup: an apostrophe at a token start absorbs an
388
+ // s/d/m/t/ll/ve/re suffix. One that could reach past bit 63 defers
389
+ // to the scalar path (the next batch cannot see the moved boundary).
390
+ let mut cand = ap64 & boundary & !bad;
391
+ while cand != 0 {
392
+ let i = cand.trailing_zeros() as usize;
393
+ cand &= cand - 1;
394
+ if i >= 61 {
395
+ bad |= u64::MAX << i;
396
+ break;
397
+ }
398
+ let k = match bytes[scan + i + 1] {
399
+ b's' | b'd' | b'm' | b't' => 2,
400
+ b'l' if bytes[scan + i + 2] == b'l' => 3,
401
+ b'v' if bytes[scan + i + 2] == b'e' => 3,
402
+ b'r' if bytes[scan + i + 2] == b'e' => 3,
403
+ _ => 0,
404
+ };
405
+ if k != 0 {
406
+ boundary &= !(1u64 << (i + 1));
407
+ boundary |= 1u64 << (i + k);
408
+ }
409
+ }
410
+
411
+ (boundary & !bad, bad)
412
+ }
413
+
414
+ /// `(pl, pd, ps, pws, po)` boundary carries for the char ending at
415
+ /// `scan - 1` (`scan > 0`), multi-byte aware via [`mask::char_through`].
416
+ /// `ps` (the ` ?` absorb) is ASCII 0x20 only. The ASCII case (almost
417
+ /// every call) is branchless — a class if-chain here is a per-batch
418
+ /// mispredict on natural text. Callers are the batch classifiers, whose
419
+ /// `scan + 70 <= len` guard covers the walk-back decode.
420
+ #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
421
+ #[inline(always)]
422
+ fn carries_at(bytes: &[u8], scan: usize) -> (u64, u64, u64, u64, u64) {
423
+ let b = bytes[scan - 1];
424
+ if b < 0x80 {
425
+ let (l, d, w) = (is_letter(b), is_digit(b), is_ascii_ws(b));
426
+ let bit = |c: bool| u64::from(c);
427
+ return (bit(l), bit(d), bit(b == b' '), bit(w), bit(!l && !d && !w));
428
+ }
429
+ // SAFETY: scan > 0 (this fn's caller contract), and the calling batch
430
+ // classifier's scan + 70 <= len guard covers pos + 3 <= len.
431
+ match unsafe { mask::char_through(bytes, scan, unicode::class_of) }.0 {
432
+ CharClass::Letter => (1, 0, 0, 0, 0),
433
+ CharClass::Number => (0, 1, 0, 0, 0),
434
+ CharClass::Whitespace => (0, 0, 0, 1, 0),
435
+ CharClass::Other => (0, 0, 0, 0, 1),
436
+ }
437
+ }
438
+ pub(crate) struct R50kScheme;
439
+
440
+ impl MaskScheme for R50kScheme {
441
+ #[inline(always)]
442
+ fn advance(bytes: &[u8], pos: usize) -> usize {
443
+ advance_pos(bytes, pos)
444
+ }
445
+
446
+ #[cfg(target_arch = "aarch64")]
447
+ #[inline(always)]
448
+ fn batch_masks(bytes: &[u8], scan: usize) -> (u64, u64) {
449
+ batch_masks(bytes, scan)
450
+ }
451
+
452
+ #[cfg(target_arch = "x86_64")]
453
+ #[inline(always)]
454
+ unsafe fn batch_masks_x86<const AVX512: bool>(bytes: &[u8], scan: usize) -> (u64, u64) {
455
+ // SAFETY: the caller detected the tier (trait contract).
456
+ unsafe { batch_masks_x86::<AVX512>(bytes, scan) }
457
+ }
458
+ }
459
+
460
+ /// With SIMD support (aarch64 NEON, or x86_64 AVX-512/AVX2 detected at
461
+ /// runtime), iteration runs on the mask scanner above via the shared
462
+ /// [`MaskState`] batch walker; elsewhere every token takes `advance_pos`.
463
+ pub struct FastR50kPretokenizer<'a> {
464
+ bytes: &'a [u8],
465
+ state: MaskState,
466
+ }
467
+
468
+ impl<'a> FastR50kPretokenizer<'a> {
469
+ #[inline]
470
+ pub fn new(bytes: &'a [u8]) -> Self {
471
+ Self::with_pos(bytes, 0)
472
+ }
473
+
474
+ /// Resume iteration at a byte offset previously returned by [`Self::pos`].
475
+ /// Used by the Python bindings, which re-borrow the underlying buffer on
476
+ /// every `__next__` call.
477
+ #[inline]
478
+ pub fn with_pos(bytes: &'a [u8], pos: usize) -> Self {
479
+ Self { bytes, state: MaskState::new(pos) }
480
+ }
481
+
482
+ /// Current position as a byte offset into the input.
483
+ #[inline]
484
+ pub fn pos(&self) -> usize {
485
+ self.state.pos
486
+ }
487
+ }
488
+
489
+ impl<'a> Iterator for FastR50kPretokenizer<'a> {
490
+ type Item = Pretoken<'a>;
491
+
492
+ #[inline]
493
+ fn next(&mut self) -> Option<Pretoken<'a>> {
494
+ let (start, end) = self.state.next_span::<R50kScheme>(self.bytes)?;
495
+ Some(Pretoken(&self.bytes[start..end]))
496
+ }
497
+ }
498
+
499
+ super::impl_mask_pretoken_spans!(FastR50kPretokenizer, R50kScheme);
500
+
501
+ /// Advance past one token starting at `start`; returns the token's end.
502
+ /// `start` must be < `bytes.len()` and a valid token start.
503
+ /// Uses direct comparison chains instead of LUT + jump table to avoid
504
+ /// GOT indirection and improve branch prediction on common patterns.
505
+ ///
506
+ /// Byte loads here look redundant but are effectively free: their addresses
507
+ /// depend only on `start`, so they issue in parallel under speculation. A
508
+ /// variant that did one u64 load and extracted bytes/scanned letters
509
+ /// in-register measured 0.84x — the shifts serialize after the load, while
510
+ /// independent L1 loads don't.
511
+ #[inline(always)]
512
+ fn advance_pos(bytes: &[u8], start: usize) -> usize {
513
+ let len = bytes.len();
514
+ let b0 = unsafe { *bytes.get_unchecked(start) };
515
+
516
+ // Bare ASCII letter start (~5% of OWT tokens; most words carry a space)
517
+ if is_letter(b0) {
518
+ return scan_letters_from(bytes, start + 1);
519
+ }
520
+
521
+ // Hot path: space before content (~78% of tokens, ~75% space+letters)
522
+ if b0 == b' ' {
523
+ if start + 1 < len {
524
+ let b1 = unsafe { *bytes.get_unchecked(start + 1) };
525
+ if is_letter(b1) {
526
+ return scan_letters_from(bytes, start + 2);
527
+ }
528
+ if is_digit(b1) {
529
+ return scan_digits_from(bytes, start + 2);
530
+ }
531
+ if b1 >= 0x80 {
532
+ let (cp, l) = unsafe { decode_cp(bytes, start + 1) };
533
+ let p = start + 1 + l;
534
+ return match unicode::class_of(cp) {
535
+ CharClass::Letter => scan_letters_from(bytes, p),
536
+ CharClass::Number => scan_digits_from(bytes, p),
537
+ CharClass::Whitespace => advance_ws(bytes, p, start),
538
+ CharClass::Other => scan_other_from(bytes, p),
539
+ };
540
+ }
541
+ if is_ascii_ws(b1) {
542
+ return advance_ws(bytes, start + 1, start);
543
+ }
544
+ return scan_other_from(bytes, start + 2);
545
+ }
546
+ return start + 1;
547
+ }
548
+
549
+ // Non-ASCII
550
+ if b0 >= 0x80 {
551
+ let (cp, l) = unsafe { decode_cp(bytes, start) };
552
+ let p = start + l;
553
+ return match unicode::class_of(cp) {
554
+ CharClass::Letter => scan_letters_from(bytes, p),
555
+ CharClass::Number => scan_digits_from(bytes, p),
556
+ CharClass::Whitespace => advance_ws(bytes, p, start),
557
+ CharClass::Other => scan_other_from(bytes, p),
558
+ };
559
+ }
560
+
561
+ // Digit
562
+ if is_digit(b0) {
563
+ return scan_digits_from(bytes, start + 1);
564
+ }
565
+
566
+ // Apostrophe / contraction
567
+ if b0 == b'\'' {
568
+ match bytes.get(start + 1) {
569
+ Some(b's' | b'd' | b'm' | b't') => return start + 2,
570
+ Some(b'l') if bytes.get(start + 2) == Some(&b'l') => return start + 3,
571
+ Some(b'v') if bytes.get(start + 2) == Some(&b'e') => return start + 3,
572
+ Some(b'r') if bytes.get(start + 2) == Some(&b'e') => return start + 3,
573
+ _ => return scan_other_from(bytes, start + 1),
574
+ }
575
+ }
576
+
577
+ // Whitespace (tab, newline, etc.)
578
+ if b0.wrapping_sub(9) < 5 {
579
+ return advance_ws(bytes, start + 1, start);
580
+ }
581
+
582
+ // Other (punctuation, symbols)
583
+ scan_other_from(bytes, start + 1)
584
+ }
585
+
586
+ /// Advance through whitespace. `scan_pos` is where to continue scanning,
587
+ /// `token_start` is where the token began (for the split-off-last-char logic).
588
+ #[inline(always)]
589
+ fn advance_ws(bytes: &[u8], scan_pos: usize, token_start: usize) -> usize {
590
+ let len = bytes.len();
591
+ let mut p = scan_pos;
592
+ while p < len {
593
+ let b = unsafe { *bytes.get_unchecked(p) };
594
+ if is_ascii_ws(b) {
595
+ p += 1;
596
+ } else if b >= 0x80 {
597
+ let (cp, l) = unsafe { decode_cp(bytes, p) };
598
+ if unicode::class_of(cp) == CharClass::Whitespace {
599
+ p += l;
600
+ } else {
601
+ break;
602
+ }
603
+ } else {
604
+ break;
605
+ }
606
+ }
607
+ if p < len {
608
+ let ws_bytes = p - token_start;
609
+ if ws_bytes >= 2 {
610
+ let mut last = p - 1;
611
+ while last > token_start && unsafe { *bytes.get_unchecked(last) } & 0xC0 == 0x80 {
612
+ last -= 1;
613
+ }
614
+ if last > token_start {
615
+ return last;
616
+ }
617
+ }
618
+ }
619
+ p
620
+ }
621
+
622
+ #[cfg(test)]
623
+ mod tests {
624
+ use super::*;
625
+
626
+ #[test]
627
+ fn fast_matches_state_machine_owt() {
628
+ let data_dir = std::env::home_dir().unwrap().join("data");
629
+ let all_bytes = std::fs::read(data_dir.join("owt_train.txt"))
630
+ .expect("Could not read ~/data/owt_train.txt");
631
+ let max = 5_000_000.min(all_bytes.len());
632
+ let mut end = max;
633
+ while end > 0 && std::str::from_utf8(&all_bytes[..end]).is_err() {
634
+ end -= 1;
635
+ }
636
+ let input = &all_bytes[..end];
637
+
638
+ let mut sm = crate::pretokenize::PretokenizerIter::new(input);
639
+ let mut fast = FastR50kPretokenizer::new(input);
640
+ let mut idx = 0usize;
641
+
642
+ loop {
643
+ match (sm.next(), fast.next()) {
644
+ (Some(a), Some(b)) => {
645
+ assert_eq!(
646
+ a.0, b.0,
647
+ "Mismatch at token {idx}: sm={:?} fast={:?}",
648
+ String::from_utf8_lossy(a.0),
649
+ String::from_utf8_lossy(b.0),
650
+ );
651
+ }
652
+ (None, None) => break,
653
+ (Some(a), None) => panic!("SM extra at {idx}: {:?}", String::from_utf8_lossy(a.0)),
654
+ (None, Some(b)) => {
655
+ panic!("Fast extra at {idx}: {:?}", String::from_utf8_lossy(b.0))
656
+ }
657
+ }
658
+ idx += 1;
659
+ }
660
+ eprintln!("All {idx} tokens match.");
661
+ }
662
+
663
+ fn load_owt(max_bytes: usize) -> Vec<u8> {
664
+ let data_dir = std::env::home_dir().unwrap().join("data");
665
+ let all_bytes = std::fs::read(data_dir.join("owt_train.txt"))
666
+ .expect("Could not read ~/data/owt_train.txt");
667
+ let mut end = max_bytes.min(all_bytes.len());
668
+ while end > 0 && std::str::from_utf8(&all_bytes[..end]).is_err() {
669
+ end -= 1;
670
+ }
671
+ all_bytes[..end].to_vec()
672
+ }
673
+
674
+ /// Reference iterator: plain `advance_pos` loop (the pre-mask-scanner
675
+ /// implementation, itself validated against the state machine by
676
+ /// `fast_matches_state_machine_owt`). The differential tests check the
677
+ /// shipped mask-scanner iterator against it token for token.
678
+ struct ScalarIter<'a> {
679
+ bytes: &'a [u8],
680
+ pos: usize,
681
+ }
682
+ impl<'a> ScalarIter<'a> {
683
+ fn new(bytes: &'a [u8]) -> Self {
684
+ Self { bytes, pos: 0 }
685
+ }
686
+ }
687
+ impl<'a> Iterator for ScalarIter<'a> {
688
+ type Item = Pretoken<'a>;
689
+ #[inline]
690
+ fn next(&mut self) -> Option<Pretoken<'a>> {
691
+ let start = self.pos;
692
+ if start >= self.bytes.len() {
693
+ return None;
694
+ }
695
+ let end = advance_pos(self.bytes, start);
696
+ self.pos = end;
697
+ Some(Pretoken(&self.bytes[start..end]))
698
+ }
699
+ }
700
+
701
+ /// Test-local wrapper over the shared batch walker, kept because
702
+ /// `cargo test` builds skip fat LTO: shipped symbols are not inlined
703
+ /// into test loops and measure ~1.4x slow, so the interleaved perf
704
+ /// harness needs a locally instantiated iterator. Must produce exactly
705
+ /// the shipped tokenization.
706
+ #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
707
+ struct MaskIter<'a> {
708
+ bytes: &'a [u8],
709
+ state: MaskState,
710
+ }
711
+
712
+ #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
713
+ impl<'a> MaskIter<'a> {
714
+ fn new(bytes: &'a [u8]) -> Self {
715
+ Self { bytes, state: MaskState::new(0) }
716
+ }
717
+ }
718
+
719
+ #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
720
+ impl<'a> Iterator for MaskIter<'a> {
721
+ type Item = Pretoken<'a>;
722
+ #[inline]
723
+ fn next(&mut self) -> Option<Pretoken<'a>> {
724
+ let (start, end) = self.state.next_span::<R50kScheme>(self.bytes)?;
725
+ Some(Pretoken(&self.bytes[start..end]))
726
+ }
727
+ }
728
+
729
+ /// The batch classifier must actually engage on any machine with the
730
+ /// assumed feature sets (aarch64 NEON; x86_64 AVX-512 or AVX2):
731
+ /// on plain ASCII text it must report real token starts and no bad
732
+ /// zones. Guards against a broken runtime detection or classifier
733
+ /// silently passing the differential tests via the scalar fallback.
734
+ #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
735
+ #[test]
736
+ fn batch_classifier_engages_on_ascii() {
737
+ if !mask::simd_scanner_available() {
738
+ eprintln!("no SIMD scanner on this CPU; skipping");
739
+ return;
740
+ }
741
+ // Longer than scan + 70: the classifier needs byte-64+ lookahead.
742
+ let text = b"The quick brown fox jumps over the lazy dog while 42 geese watch on quietly";
743
+ let (usable, bad) = R50kScheme::batch_masks(text, 0);
744
+ assert_eq!(bad, 0, "plain ASCII must produce no bad zones");
745
+ let mut starts = vec![];
746
+ let mut p = 0;
747
+ while p < text.len() {
748
+ starts.push(p);
749
+ p = advance_pos(text, p);
750
+ }
751
+ for i in 0..64usize {
752
+ if usable >> i & 1 == 1 {
753
+ assert!(starts.contains(&i), "usable bit {i} is not a token start");
754
+ }
755
+ }
756
+ assert!(usable.count_ones() >= 10, "classifier found too few boundaries");
757
+ }
758
+
759
+ /// Token-for-token differential check of `MaskIter` vs the shipped
760
+ /// iterator on crafted edge cases.
761
+ #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
762
+ #[test]
763
+ fn mask_iter_matches_shipped_edge_cases() {
764
+ let cases: Vec<Vec<u8>> = vec![
765
+ b"".to_vec(),
766
+ b" ".to_vec(),
767
+ b"a".to_vec(),
768
+ b"hello world".to_vec(),
769
+ b" double spaces ".to_vec(),
770
+ b"a\n\nb".to_vec(),
771
+ b"a \n b".to_vec(),
772
+ b"tabs\tand\nnewlines\r\n end".to_vec(),
773
+ b"don't can't we'll they've you're I'm he's 'tis 'twas".to_vec(),
774
+ b"DON'T CAN'T 'S 'LL".to_vec(),
775
+ b"x'y z' 'a '' ' ".to_vec(),
776
+ b"3.14 100,000 2nd a1b2".to_vec(),
777
+ b"!!! ?! #hashtag @user (paren) [brack]".to_vec(),
778
+ "café résumé naïve".as_bytes().to_vec(),
779
+ "日本語のテキスト and English".as_bytes().to_vec(),
780
+ "space\u{00A0}nbsp \u{00A0} runs".as_bytes().to_vec(),
781
+ "emoji 🎉🎊 mix".as_bytes().to_vec(),
782
+ "µ§±².5 ×÷".as_bytes().to_vec(),
783
+ b"ws at end ".to_vec(),
784
+ b" ws at start".to_vec(),
785
+ // Exactly chunk-sized and chunk-straddling patterns.
786
+ b"abcdefghijklmnop".to_vec(),
787
+ b"abcdefghijklmno ".to_vec(),
788
+ b"abcdefghijklmn 'll xyz".to_vec(),
789
+ b"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".to_vec(),
790
+ b"a b c d e f g h i j k l m n o p q r s t u v w x".to_vec(),
791
+ [b"word ".repeat(10), b"\xE2\x80\x82ws".to_vec()].concat(),
792
+ ];
793
+ for (ci, case) in cases.iter().enumerate() {
794
+ let shipped: Vec<&[u8]> = ScalarIter::new(case).map(|t| t.0).collect();
795
+ let masked: Vec<&[u8]> = FastR50kPretokenizer::new(case).map(|t| t.0).collect();
796
+ assert_eq!(
797
+ shipped,
798
+ masked,
799
+ "case {ci} diverged: {:?}",
800
+ String::from_utf8_lossy(case)
801
+ );
802
+ }
803
+ }
804
+
805
+ /// Differential fuzz: random mixes of letters, digits, ws, punctuation,
806
+ /// apostrophes, and multi-byte UTF-8 at every length 0..~200.
807
+ #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
808
+ #[test]
809
+ fn mask_iter_matches_shipped_fuzz() {
810
+ let pieces: &[&str] = &[
811
+ "a", "B", "z", "9", "0", " ", " ", "\n", "\t", "\r\n", "'", "'s", "'ll", "'re",
812
+ "!", ".", ",", "(", "é", "ß", "日", "🎉", "\u{00A0}", "\u{2003}", "word", "12",
813
+ "’", "’s", "“", "”", "–", "—", "…", "\u{2009}", "\u{200B}", "\u{2028}",
814
+ "\u{202F}", "×", "÷", "«", "µ", "café", "éé", "naïve", "Α", "а", "ſ", "'ſ",
815
+ "\u{661}\u{662}", "\u{FF11}", "क", "\u{940}", "\u{1D54F}", "€", "™", "\u{301}",
816
+ ];
817
+ let mut state = 0x243F6A8885A308D3u64;
818
+ let mut rng = move || {
819
+ state ^= state << 13;
820
+ state ^= state >> 7;
821
+ state ^= state << 17;
822
+ state
823
+ };
824
+ for round in 0..4000 {
825
+ let target = (round % 200) + 1;
826
+ let mut buf = Vec::new();
827
+ while buf.len() < target {
828
+ buf.extend_from_slice(pieces[(rng() % pieces.len() as u64) as usize].as_bytes());
829
+ }
830
+ let shipped: Vec<&[u8]> = ScalarIter::new(&buf).map(|t| t.0).collect();
831
+ let masked: Vec<&[u8]> = FastR50kPretokenizer::new(&buf).map(|t| t.0).collect();
832
+ assert_eq!(
833
+ shipped,
834
+ masked,
835
+ "round {round} diverged on {:?}",
836
+ String::from_utf8_lossy(&buf)
837
+ );
838
+ }
839
+ }
840
+
841
+ /// Differential check on the FULL OWT file (~11.9 GB), token for token.
842
+ #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
843
+ #[test]
844
+ #[ignore]
845
+ fn mask_iter_matches_shipped_owt_full() {
846
+ let input = load_owt(usize::MAX);
847
+ eprintln!("loaded {} bytes", input.len());
848
+ let mut shipped = ScalarIter::new(&input);
849
+ let mut masked = FastR50kPretokenizer::new(&input);
850
+ let mut idx = 0usize;
851
+ loop {
852
+ match (shipped.next(), masked.next()) {
853
+ (Some(a), Some(b)) => {
854
+ if a.0 != b.0 {
855
+ panic!(
856
+ "token {idx} diverged: scalar={:?} masked={:?}",
857
+ String::from_utf8_lossy(a.0),
858
+ String::from_utf8_lossy(b.0)
859
+ );
860
+ }
861
+ }
862
+ (None, None) => break,
863
+ (a, b) => panic!("length mismatch at {idx}: {:?} vs {:?}", a.is_some(), b.is_some()),
864
+ }
865
+ idx += 1;
866
+ }
867
+ eprintln!("all {idx} tokens match on full OWT");
868
+ }
869
+
870
+ /// Differential check on real OWT (100 MB), token for token.
871
+ #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
872
+ #[test]
873
+ #[ignore]
874
+ fn mask_iter_matches_shipped_owt() {
875
+ let input = load_owt(100_000_000);
876
+ let mut shipped = ScalarIter::new(&input);
877
+ let mut masked = FastR50kPretokenizer::new(&input);
878
+ let mut idx = 0usize;
879
+ loop {
880
+ match (shipped.next(), masked.next()) {
881
+ (Some(a), Some(b)) => assert_eq!(
882
+ a.0,
883
+ b.0,
884
+ "token {idx} diverged: shipped={:?} masked={:?}",
885
+ String::from_utf8_lossy(a.0),
886
+ String::from_utf8_lossy(b.0)
887
+ ),
888
+ (None, None) => break,
889
+ (a, b) => panic!("length mismatch at {idx}: {:?} vs {:?}", a.is_some(), b.is_some()),
890
+ }
891
+ idx += 1;
892
+ }
893
+ eprintln!("all {idx} tokens match");
894
+ }
895
+
896
+ /// Iterator-level interleaved A/B: local copy of the shipped scalar
897
+ /// iterator vs the mask scanner. Both variants are LOCAL copies: test
898
+ /// builds skip fat LTO, so shipped symbols are not inlined here and
899
+ /// measure ~1.4x slow (see project memory / step 14 of the log).
900
+ #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
901
+ #[test]
902
+ #[ignore]
903
+ fn ab_r50k_mask_iter_interleaved() {
904
+ fn drive_scalar(input: &[u8]) -> (usize, u64) {
905
+ let mut n = 0usize;
906
+ let mut acc = 0u64;
907
+ for t in ScalarIter::new(input) {
908
+ std::hint::black_box(t.0);
909
+ acc = acc.wrapping_add(t.0.len() as u64);
910
+ n += 1;
911
+ }
912
+ (n, acc)
913
+ }
914
+ fn drive_mask(input: &[u8]) -> (usize, u64) {
915
+ let mut n = 0usize;
916
+ let mut acc = 0u64;
917
+ for t in MaskIter::new(input) {
918
+ std::hint::black_box(t.0);
919
+ acc = acc.wrapping_add(t.0.len() as u64);
920
+ n += 1;
921
+ }
922
+ (n, acc)
923
+ }
924
+
925
+ let input = load_owt(100_000_000);
926
+ let mb = input.len() as f64 / 1e6;
927
+ let a = drive_scalar(&input);
928
+ let b = drive_mask(&input);
929
+ assert_eq!(a, b, "variants disagree");
930
+
931
+ let (mut best_a, mut best_b) = (f64::INFINITY, f64::INFINITY);
932
+ for round in 0..7 {
933
+ let t = std::time::Instant::now();
934
+ std::hint::black_box(drive_scalar(&input));
935
+ let da = t.elapsed().as_secs_f64();
936
+ let t = std::time::Instant::now();
937
+ std::hint::black_box(drive_mask(&input));
938
+ let db = t.elapsed().as_secs_f64();
939
+ best_a = best_a.min(da);
940
+ best_b = best_b.min(db);
941
+ eprintln!(
942
+ "round {round}: scalar {:.0} MB/s | mask {:.0} MB/s",
943
+ mb / da,
944
+ mb / db
945
+ );
946
+ }
947
+ eprintln!(
948
+ "best: scalar {:.0} MB/s | mask {:.0} MB/s | mask/scalar {:.3}x",
949
+ mb / best_a,
950
+ mb / best_b,
951
+ best_a / best_b
952
+ );
953
+ }
954
+
955
+ /// Diagnostics for the mask scanner: fallback rate on real OWT, and
956
+ /// mask-vs-scalar throughput on a sanitized buffer (apostrophes and
957
+ /// non-ASCII replaced) that never triggers the fallback.
958
+ #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
959
+ #[test]
960
+ #[ignore]
961
+ fn mask_iter_diagnostics() {
962
+ let input = load_owt(100_000_000);
963
+
964
+ // Bad-zone census: dirty batches, and bytes covered by the first
965
+ // gap (prefix end to resume point) as a scalar-work proxy.
966
+ let (mut batches, mut dirty_batches, mut gap_bytes) = (0usize, 0usize, 0u64);
967
+ let mut scan = 0usize;
968
+ while scan + 64 <= input.len() {
969
+ let (usable, bad) = R50kScheme::batch_masks(&input, scan);
970
+ batches += 1;
971
+ if bad != 0 {
972
+ dirty_batches += 1;
973
+ let first_bad = bad.trailing_zeros();
974
+ let rest = usable & (u64::MAX << first_bad);
975
+ let resume = if rest != 0 { rest.trailing_zeros() } else { 64 };
976
+ gap_bytes += u64::from(resume - first_bad);
977
+ }
978
+ scan += 64;
979
+ }
980
+ eprintln!(
981
+ "batches: {batches}, dirty: {dirty_batches} ({:.2}%), first-gap bytes {:.2}%",
982
+ 100.0 * dirty_batches as f64 / batches as f64,
983
+ 100.0 * gap_bytes as f64 / input.len() as f64
984
+ );
985
+
986
+ // Sanitized buffer: no apostrophes, no high bytes.
987
+ let clean: Vec<u8> = input
988
+ .iter()
989
+ .map(|&b| if b >= 0x80 || b == b'\'' { b'x' } else { b })
990
+ .collect();
991
+ let mb = clean.len() as f64 / 1e6;
992
+
993
+ fn drive<'a, I: Iterator<Item = Pretoken<'a>>>(it: I) -> (usize, u64) {
994
+ let mut n = 0usize;
995
+ let mut acc = 0u64;
996
+ for t in it {
997
+ std::hint::black_box(t.0);
998
+ acc = acc.wrapping_add(t.0.len() as u64);
999
+ n += 1;
1000
+ }
1001
+ (n, acc)
1002
+ }
1003
+ assert_eq!(
1004
+ drive(ScalarIter::new(&clean)),
1005
+ drive(MaskIter::new(&clean))
1006
+ );
1007
+ let (mut best_a, mut best_b) = (f64::INFINITY, f64::INFINITY);
1008
+ for round in 0..5 {
1009
+ let t = std::time::Instant::now();
1010
+ std::hint::black_box(drive(ScalarIter::new(&clean)));
1011
+ let da = t.elapsed().as_secs_f64();
1012
+ let t = std::time::Instant::now();
1013
+ std::hint::black_box(drive(MaskIter::new(&clean)));
1014
+ let db = t.elapsed().as_secs_f64();
1015
+ best_a = best_a.min(da);
1016
+ best_b = best_b.min(db);
1017
+ eprintln!(
1018
+ "clean round {round}: scalar {:.0} MB/s | mask {:.0} MB/s",
1019
+ mb / da,
1020
+ mb / db
1021
+ );
1022
+ }
1023
+ eprintln!(
1024
+ "clean best: scalar {:.0} | mask {:.0} MB/s | {:.3}x",
1025
+ mb / best_a,
1026
+ mb / best_b,
1027
+ best_a / best_b
1028
+ );
1029
+ }
1030
+
1031
+ #[test]
1032
+ #[ignore]
1033
+ fn r50k_token_stats_owt() {
1034
+ let input = load_owt(10_000_000);
1035
+ let mut counts = [0usize; 8]; // letter, sp+letter, sp+digit, sp+other, digit, ws, apos, other/nonascii
1036
+ let mut letter_tokens = 0usize;
1037
+ let mut in_window = 0usize;
1038
+ let mut total = 0usize;
1039
+ let mut pos = 0usize;
1040
+ while pos < input.len() {
1041
+ let end = advance_pos(&input, pos);
1042
+ let b0 = input[pos];
1043
+ let idx = if is_letter(b0) {
1044
+ 0
1045
+ } else if b0 == b' ' {
1046
+ match input.get(pos + 1) {
1047
+ Some(&b) if is_letter(b) => 1,
1048
+ Some(&b) if is_digit(b) => 2,
1049
+ _ => 3,
1050
+ }
1051
+ } else if is_digit(b0) {
1052
+ 4
1053
+ } else if is_ascii_ws(b0) {
1054
+ 5
1055
+ } else if b0 == b'\'' {
1056
+ 6
1057
+ } else {
1058
+ 7
1059
+ };
1060
+ counts[idx] += 1;
1061
+ if idx == 0 || idx == 1 {
1062
+ letter_tokens += 1;
1063
+ if end - pos <= 8 {
1064
+ in_window += 1;
1065
+ }
1066
+ }
1067
+ total += 1;
1068
+ pos = end;
1069
+ }
1070
+ let pct = |n: usize| 100.0 * n as f64 / total as f64;
1071
+ eprintln!("total tokens: {total}");
1072
+ for (name, n) in [
1073
+ "letter", "sp+letter", "sp+digit", "sp+other", "digit", "ws", "apos", "other/hi",
1074
+ ]
1075
+ .iter()
1076
+ .zip(counts)
1077
+ {
1078
+ eprintln!("{name:>10}: {n:>9} ({:.1}%)", pct(n));
1079
+ }
1080
+ eprintln!(
1081
+ "letter tokens resolved in 8-byte window: {:.1}%",
1082
+ 100.0 * in_window as f64 / letter_tokens as f64
1083
+ );
1084
+ eprintln!(
1085
+ "avg token len: {:.2} bytes",
1086
+ input.len() as f64 / total as f64
1087
+ );
1088
+ }
1089
+
1090
+ /// Best-of-5 throughput of `advance_pos` over `bytes`, in MB/s.
1091
+ fn measure(bytes: &[u8], label: &str) -> f64 {
1092
+ let mb = bytes.len() as f64 / 1e6;
1093
+ let (n, _) = drive(bytes, advance_pos);
1094
+ let mut best = f64::INFINITY;
1095
+ for _ in 0..5 {
1096
+ let t = std::time::Instant::now();
1097
+ std::hint::black_box(drive(bytes, advance_pos));
1098
+ best = best.min(t.elapsed().as_secs_f64());
1099
+ }
1100
+ let mbs = mb / best;
1101
+ // cycles/token estimate assumes ~4.5 GHz P-core
1102
+ let cpt = 4.5e9 * best / n as f64;
1103
+ eprintln!(
1104
+ "{label:>22}: {mbs:>5.0} MB/s | {:.2} B/token | ~{cpt:.1} cy/token",
1105
+ bytes.len() as f64 / n as f64
1106
+ );
1107
+ mbs
1108
+ }
1109
+
1110
+ #[test]
1111
+ #[ignore]
1112
+ fn r50k_prediction_floor() {
1113
+ // Pure single-path floor: one token shape, perfectly predictable.
1114
+ let hello: Vec<u8> = b" hello".repeat(16_000_000 / 6);
1115
+ measure(&hello, "' hello' x N");
1116
+
1117
+ // Same multiset of real OWT tokens, predictable vs shuffled order.
1118
+ // Space-prefixed tokens are adjacency-safe: concatenating them in any
1119
+ // order re-tokenizes to exactly the same tokens.
1120
+ let input = load_owt(100_000_000);
1121
+ let bag: Vec<&[u8]> = FastR50kPretokenizer::new(&input)
1122
+ .map(|t| t.0)
1123
+ .filter(|t| t[0] == b' ' && t.len() > 1)
1124
+ .collect();
1125
+ eprintln!("bag: {} space-prefixed tokens", bag.len());
1126
+
1127
+ let mut sorted = bag.clone();
1128
+ sorted.sort_unstable();
1129
+ let predictable: Vec<u8> = sorted.concat();
1130
+ drop(sorted);
1131
+
1132
+ let mut shuffled = bag;
1133
+ let mut state = 0x9E3779B97F4A7C15u64;
1134
+ let mut next = || {
1135
+ state ^= state >> 12;
1136
+ state ^= state << 25;
1137
+ state ^= state >> 27;
1138
+ state.wrapping_mul(0x2545F4914F6CDD1D)
1139
+ };
1140
+ for i in (1..shuffled.len()).rev() {
1141
+ let j = (next() % (i as u64 + 1)) as usize;
1142
+ shuffled.swap(i, j);
1143
+ }
1144
+ let unpredictable: Vec<u8> = shuffled.concat();
1145
+ let n_bag = shuffled.len();
1146
+ drop(shuffled);
1147
+
1148
+ assert_eq!(drive(&predictable, advance_pos).0, n_bag);
1149
+ assert_eq!(drive(&unpredictable, advance_pos).0, n_bag);
1150
+
1151
+ measure(&predictable, "sorted (predictable)");
1152
+ measure(&unpredictable, "shuffled (same bag)");
1153
+ measure(&input, "natural OWT");
1154
+ }
1155
+
1156
+ fn drive(bytes: &[u8], f: impl Fn(&[u8], usize) -> usize) -> (usize, u64) {
1157
+ let mut pos = 0usize;
1158
+ let mut n = 0usize;
1159
+ let mut acc = 0u64;
1160
+ while pos < bytes.len() {
1161
+ let end = f(bytes, pos);
1162
+ acc = acc.wrapping_add(end as u64);
1163
+ n += 1;
1164
+ pos = end;
1165
+ }
1166
+ (n, acc)
1167
+ }
1168
+
1169
+ #[test]
1170
+ #[ignore]
1171
+ fn aa_r50k_advance_interleaved() {
1172
+ // A/A control: same implementation through two separately
1173
+ // monomorphized drivers, to gauge code-layout noise in this harness.
1174
+ let input = load_owt(100_000_000);
1175
+ let mb = input.len() as f64 / 1e6;
1176
+ let copy_a = advance_pos;
1177
+ let copy_b = |b: &[u8], s: usize| advance_pos(b, s);
1178
+ std::hint::black_box(drive(&input, copy_a));
1179
+ std::hint::black_box(drive(&input, copy_b));
1180
+ let (mut best_a, mut best_b) = (f64::INFINITY, f64::INFINITY);
1181
+ for round in 0..7 {
1182
+ let t = std::time::Instant::now();
1183
+ std::hint::black_box(drive(&input, copy_a));
1184
+ let da = t.elapsed().as_secs_f64();
1185
+ let t = std::time::Instant::now();
1186
+ std::hint::black_box(drive(&input, copy_b));
1187
+ let db = t.elapsed().as_secs_f64();
1188
+ best_a = best_a.min(da);
1189
+ best_b = best_b.min(db);
1190
+ eprintln!("round {round}: A {:.0} MB/s | A' {:.0} MB/s", mb / da, mb / db);
1191
+ }
1192
+ eprintln!(
1193
+ "best: A {:.0} MB/s | A' {:.0} MB/s | ratio {:.3}x",
1194
+ mb / best_a,
1195
+ mb / best_b,
1196
+ best_a / best_b
1197
+ );
1198
+ }
1199
+
1200
+ /// Interleaved same-binary A/B harness (min-of-7): swap an experimental
1201
+ /// `advance_pos` candidate in below and run with `--ignored --nocapture`.
1202
+ /// Run `aa_r50k_advance_interleaved` first as the layout-noise control.
1203
+ /// (Isolated A/B of the packed class table vs the old chained ICU
1204
+ /// predicates measured 0.90x for the ICU version, i.e. the table alone
1205
+ /// is worth ~1.11x.)
1206
+ #[test]
1207
+ #[ignore]
1208
+ fn ab_r50k_advance_interleaved() {
1209
+ let input = load_owt(100_000_000);
1210
+ let mb = input.len() as f64 / 1e6;
1211
+ eprintln!("input: {mb:.1} MB");
1212
+
1213
+ // Experimental candidate under test; placeholder = shipped impl.
1214
+ let experimental = |b: &[u8], s: usize| advance_pos(b, s);
1215
+
1216
+ // Warmup + equivalence check
1217
+ let (n_a, acc_a) = drive(&input, advance_pos);
1218
+ let (n_b, acc_b) = drive(&input, experimental);
1219
+ assert_eq!((n_a, acc_a), (n_b, acc_b), "variants disagree");
1220
+
1221
+ let mut best_a = f64::INFINITY;
1222
+ let mut best_b = f64::INFINITY;
1223
+ for round in 0..7 {
1224
+ let t = std::time::Instant::now();
1225
+ let r = drive(&input, advance_pos);
1226
+ let da = t.elapsed().as_secs_f64();
1227
+ std::hint::black_box(r);
1228
+
1229
+ let t = std::time::Instant::now();
1230
+ let r = drive(&input, experimental);
1231
+ let db = t.elapsed().as_secs_f64();
1232
+ std::hint::black_box(r);
1233
+
1234
+ best_a = best_a.min(da);
1235
+ best_b = best_b.min(db);
1236
+ eprintln!(
1237
+ "round {round}: shipped {:.0} MB/s | experimental {:.0} MB/s",
1238
+ mb / da,
1239
+ mb / db
1240
+ );
1241
+ }
1242
+ eprintln!(
1243
+ "best: shipped {:.0} MB/s | experimental {:.0} MB/s | ratio {:.3}x",
1244
+ mb / best_a,
1245
+ mb / best_b,
1246
+ best_a / best_b
1247
+ );
1248
+ }
1249
+
1250
+ }