diff --git a/tokenizers/atomsplit/Cargo.toml b/tokenizers/atomsplit/Cargo.toml index fb8947f3f..dce664880 100644 --- a/tokenizers/atomsplit/Cargo.toml +++ b/tokenizers/atomsplit/Cargo.toml @@ -43,6 +43,10 @@ harness = false name = "classify" harness = false +[[bench]] +name = "literal" +harness = false + [[bench]] name = "class_runs" harness = false diff --git a/tokenizers/atomsplit/benches/literal.rs b/tokenizers/atomsplit/benches/literal.rs new file mode 100644 index 000000000..c165b3f6c --- /dev/null +++ b/tokenizers/atomsplit/benches/literal.rs @@ -0,0 +1,67 @@ +//! `Literal::matches` (one search per match) against `Literal::matches_into` (one scan per +//! text), across the delimiter densities pre-tokenizers see: a space about every six bytes of +//! English, a `▁` per word after a SentencePiece replace, and a pattern the text does not +//! contain at all (where `memmem`'s skip-ahead is at its best and the scan must keep up). +//! Run: cargo bench --bench literal +use atomsplit::literal::Literal; +use std::hint::black_box; +use std::time::Instant; + +fn best_ns_per_byte(text_len: usize, mut pass: impl FnMut() -> usize) -> f64 { + let iters = (16_000_000 / text_len.max(1)).clamp(4, 400) as u32; + for _ in 0..3 { + black_box(pass()); + } + let mut best = f64::INFINITY; + for _ in 0..9 { + let t = Instant::now(); + for _ in 0..iters { + black_box(pass()); + } + best = best.min(t.elapsed().as_nanos() as f64 / (iters as usize * text_len) as f64); + } + best +} + +fn main() { + let manifest = env!("CARGO_MANIFEST_DIR"); + let english = + std::fs::read_to_string(format!("{manifest}/../data/big.txt")).unwrap_or_default(); + let english: String = english.chars().take(2_000_000).collect(); + if english.is_empty() { + println!("empty corpus — data/big.txt missing? (make test downloads it)"); + return; + } + let metaspaced = english.replace(' ', "\u{2581}"); + + println!( + "{:<22} {:>9} {:>12} {:>12} {:>9} {:>12}", + "", "matches", "iterator", "batch", "speedup", "count-only" + ); + for (label, text, pattern) in [ + ("space in English", english.as_bytes(), " "), + ("▁ in metaspaced", metaspaced.as_bytes(), "\u{2581}"), + ("▁ absent", english.as_bytes(), "\u{2581}"), + ] { + let literal = Literal::new(pattern.as_bytes()).unwrap(); + let mut offsets: Vec = Vec::with_capacity(text.len()); + let mut buffer = vec![0u32; text.len() + 4]; + + let count = literal.matches(text).count(); + let iterator = best_ns_per_byte(text.len(), || { + offsets.clear(); + offsets.extend(literal.matches(text)); + offsets.len() + }); + let batch = best_ns_per_byte(text.len(), || literal.matches_into(text, &mut buffer)); + let counting = best_ns_per_byte(text.len(), || literal.count_matches(text)); + + println!( + "{label:<22} {count:>9} {:>7.2} GB/s {:>7.2} GB/s {:>8.2}x {:>7.2} GB/s", + 1.0 / iterator, + 1.0 / batch, + iterator / batch, + 1.0 / counting + ); + } +} diff --git a/tokenizers/atomsplit/src/lib.rs b/tokenizers/atomsplit/src/lib.rs index 5f22f52d0..13beac00a 100644 --- a/tokenizers/atomsplit/src/lib.rs +++ b/tokenizers/atomsplit/src/lib.rs @@ -28,6 +28,12 @@ mod simd_avx_classify; #[cfg(target_arch = "aarch64")] mod simd_classify; mod simd_fsm; +#[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") +))] +mod simd_literal; #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] mod simd_wasm_classify; pub mod tables; diff --git a/tokenizers/atomsplit/src/literal.rs b/tokenizers/atomsplit/src/literal.rs index c826b2dda..9e7106c1a 100644 --- a/tokenizers/atomsplit/src/literal.rs +++ b/tokenizers/atomsplit/src/literal.rs @@ -4,7 +4,15 @@ //! FSMs cut where the class changes ([`crate::fsm`]). //! //! The atom classification is unnecessary for pre-tokenizers splitting on an exact character or literal string: -//! We can use a simpler byte search looking at 16 bytes at a time. +//! a plain byte search finds the same cuts. Two searches cover the two match densities: +//! +//! - [`Literal::matches`] iterates with [`memmem`], which is built for needles that are rare in +//! the haystack: it skips ahead fast and restarts after every match. +//! - [`Literal::matches_into`] scans the whole text once and writes every offset into a caller +//! buffer. Pre-tokenizer delimiters are the opposite of rare (running English text has a space +//! about every six bytes), and at that density the per-match restarts dominate the iterator. +//! The scan answers a whole block of text at once and turns the answers into offsets; +//! `simd_literal.rs` explains how, step by step. use memchr::memmem; use std::fmt; @@ -53,4 +61,135 @@ impl Literal { pub fn matches<'t>(&'t self, text: &'t [u8]) -> impl Iterator + 't { self.finder.find_iter(text) } + + /// The same offsets as [`Literal::matches`], written into `out[..count]` in one scan of + /// `text`; returns the count. Use this when matches are frequent: the iterator restarts its + /// search machinery after every match, the scan never stops. + /// + /// # Preconditions + /// - `out.len() >= text.len() / pattern.len() + 4` (asserted). The division is the most + /// matches the text can hold; the `+ 4` is slack the SIMD path writes past the last match. + /// - `text.len()` must fit in `u32` (asserted), so that every offset does too. + pub fn matches_into(&self, text: &[u8], out: &mut [u32]) -> usize { + assert!( + out.len() >= text.len() / self.pattern().len() + 4, + "matches_into needs out.len() >= text.len() / pattern.len() + 4" + ); + assert!( + u32::try_from(text.len()).is_ok(), + "matches_into writes u32 offsets" + ); + #[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + ))] + if self.scannable() { + match *self.pattern() { + [a] => return crate::simd_literal::matches_into([a], text, out), + [a, b] => return crate::simd_literal::matches_into([a, b], text, out), + [a, b, c] => return crate::simd_literal::matches_into([a, b, c], text, out), + _ => unreachable!("scannable patterns are one to three bytes"), + } + } + let mut count = 0; + for position in self.finder.find_iter(text) { + out[count] = position as u32; + count += 1; + } + count + } + + /// How many matches [`Literal::matches`] would report. Counting runs the compare step + /// without producing any offsets, so it is several times faster than a full scan; count + /// first to size an output exactly, then fill it with a [`Literal::for_each_match`] pass. + #[must_use] + pub fn count_matches(&self, text: &[u8]) -> usize { + #[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + ))] + if self.scannable() { + match *self.pattern() { + [a] => return crate::simd_literal::count_matches([a], text), + [a, b] => return crate::simd_literal::count_matches([a, b], text), + [a, b, c] => return crate::simd_literal::count_matches([a, b, c], text), + _ => unreachable!("scannable patterns are one to three bytes"), + } + } + self.finder.find_iter(text).count() + } + + /// Calls `on_match` with the byte offset of every match, left to right: the same offsets + /// as [`Literal::matches`] at batch-scan speed, with no buffer for the caller to provide. + /// The scan streams through a small stack window, so its footprint stays flat however + /// long the text is. + pub fn for_each_match(&self, text: &[u8], mut on_match: impl FnMut(usize)) { + if !self.scannable() { + for position in self.finder.find_iter(text) { + on_match(position); + } + return; + } + self.scan_windows(text, |base, matches| { + for &position in matches { + on_match(base + position as usize); + } + }); + } + + /// Whether the batch scan covers this pattern. The scan emits every position where the + /// pattern matches, with no ordering check between them; that is only the non-overlapping + /// match list when the pattern cannot overlap itself, which one byte comparison per + /// length decides: a two-byte pattern overlaps itself when both bytes are equal, a + /// three-byte one when last equals first. Longer or self-overlapping patterns search + /// through the [`memmem`] engine instead. + #[cfg(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + ))] + fn scannable(&self) -> bool { + match *self.pattern() { + [_] => true, + [a, b] => a != b, + [a, _, c] => a != c, + _ => false, + } + } + + /// Without a SIMD kernel there is no batch scan; everything searches through [`memmem`]. + #[cfg(not(any( + target_arch = "aarch64", + target_arch = "x86_64", + all(target_arch = "wasm32", target_feature = "simd128") + )))] + fn scannable(&self) -> bool { + false + } + + /// Streams the batch scan window by window: [`Literal::matches_into`] fills a stack + /// buffer for one stretch of text, `on_window` consumes it (offsets are relative to the + /// window's `base`), and the next window starts where a match could no longer fit in the + /// previous one, so nothing at a window edge is missed or reported twice. + fn scan_windows(&self, text: &[u8], mut on_window: impl FnMut(usize, &[u32])) { + // 4KB of stack; a window of (1024 - 4) * pattern.len() bytes fills it exactly. + let mut buffer = [0u32; 1024]; + let width = self.pattern().len(); + let window = (buffer.len() - 4) * width; + let mut base = 0; + loop { + let end = usize::min(base + window, text.len()); + let count = self.matches_into(&text[base..end], &mut buffer); + on_window(base, &buffer[..count]); + if end == text.len() { + return; + } + // A match crossing `end` was not reported; the next window re-covers its + // possible starts. Matches of a scannable pattern never overlap, so the reports + // stay disjoint. + base = end - (width - 1); + } + } } diff --git a/tokenizers/atomsplit/src/simd_literal.rs b/tokenizers/atomsplit/src/simd_literal.rs new file mode 100644 index 000000000..78e6bfa58 --- /dev/null +++ b/tokenizers/atomsplit/src/simd_literal.rs @@ -0,0 +1,783 @@ +//! Finds every place a short pattern (up to three bytes) occurs in a text, in one pass. +//! This is the scan behind [`Literal::matches_into`](crate::literal::Literal::matches_into); +//! [`crate::literal`] explains when it is picked over the plain iterator. +//! +//! # Why a dedicated scan +//! +//! A pre-tokenizer splits on a delimiter that shows up every few bytes: running English text +//! has a space about every six bytes, SentencePiece text has a `▁` at every word. A +//! next-match engine like `memmem` is built for the opposite case, a needle that is rare, so +//! at this density it spends most of its time stopping and restarting: every match returns to +//! the caller, and every call re-enters the search machinery. This scan never stops. It +//! answers "where are ALL the matches in these 16 bytes?" with a handful of instructions, +//! block after block. +//! +//! # How: compare, mask, decode +//! +//! Two words come up in everything below. A SIMD **register** is an extra-wide CPU value +//! holding 16 bytes side by side; one instruction operates on all 16 at once. Each of those +//! 16 byte slots is a **lane**. A SIMD compare answers lane by lane: a lane becomes all ones +//! where the answer is yes, all zeros where it is no. +//! +//! **Step 1, compare.** Compare the text against the pattern's first byte (copied across all +//! 16 lanes), the text shifted by one against the second byte, and so on, then AND the +//! answers together: a lane that passed every compare is the start of a whole match, and +//! nothing needs a second look. Searching for `▁` (three bytes, `E2 96 81`) in +//! `"a…b▁c"`, where the `…` (`E2 80 A6`) shares a first byte with `▁` and acts as the decoy: +//! +//! ```text +//! offset 0 1 2 3 4 5 6 7 8 +//! text a E2 80 A6 b E2 96 81 c +//! └── … ──────┘ └── ▁ ──────┘ +//! text[j] == E2 ? · ✓ · · · ✓ · · · +//! text[j+1] == 96 ? · · · · · ✓ · · · +//! text[j+2] == 81 ? · · · · · ✓ · · · +//! AND of the three · · · · · ✓ · · · +//! ``` +//! +//! Only offset 5 passes all three rows: a match starts at 5, and the decoy at 1 is out after +//! the second row. +//! +//! **Step 2, mask.** The 16 lane answers are squeezed into one ordinary integer so plain +//! arithmetic can take over: bit `j` set means "a match starts at offset `j`". For the +//! example above the mask is `0b100000`, bit 5. Each target has a one-or-two instruction way +//! to build this integer; on NEON the bits come out spaced four positions apart instead of +//! one, a spacing the shared code carries along as `SHIFT` (the per-target layer below). +//! +//! **Step 3, decode.** Count-trailing-zeros on the mask gives the lowest set bit, which is the +//! next match offset; clear that bit and repeat. The one twist is that the code decodes four +//! offsets *without checking whether any bits are left*. Say a block starting at text offset +//! `base` had matches at offsets 2 and 5 (shown with the one-bit spacing): +//! +//! ```text +//! mask 0b100100 trailing zeros = 2 write base + 2 clear bit → 0b100000 +//! mask 0b100000 trailing zeros = 5 write base + 5 clear bit → 0 +//! mask 0 trailing zeros = 64 write garbage (slot 3) +//! mask 0 trailing zeros = 64 write garbage (slot 4) +//! cursor += 2 (how many bits were set) +//! ``` +//! +//! The two garbage slots sit past the cursor, so the next block's writes (or nothing) land on +//! them and they never reach the caller. The honest alternative, checking "still something +//! left?" before each decode, puts a branch in the hottest loop; with a match every four to +//! eight bytes that branch guesses wrong about once per block, which costs more than all the +//! compares together. +//! +//! Sparse text takes none of this: one instruction per 64-byte block answers "nothing here" +//! (see `any`), so a pattern that is rare or absent stays close to `memmem` speed. +//! +//! The compare trick is Wojciech Muła's ("SIMD-friendly algorithms for substring searching"); +//! the branch-free decode is how simdjson reads its masks. +//! +//! # The per-target layer +//! +//! Everything from the mask on is plain integer arithmetic and is shared. A target provides +//! four primitives in its `arch` module: +//! +//! - `splat`: one byte copied into every lane, the shape a compare wants. +//! - `match_starts`: step 1, the shifted compares ANDed together. +//! - `any`: "did anything in these 64 bytes match?", the cheap exit for sparse text. +//! - `bits`: step 2, lane answers to integer mask. +//! +//! x86 (`movemask`) and wasm (`bitmask`) each have an instruction that grabs one bit from +//! every lane, so offset `j` lands on bit `j`. NEON has no such instruction; its standard +//! substitute (Danila Kutenin's) leaves the bits four positions apart, so offset `j` lands +//! on bit `4 * j`. `SHIFT` is that spacing (`bit position >> SHIFT` recovers the offset), +//! and it is the only difference the shared code ever sees. +//! +//! On x86 the 16-byte baseline is SSE2, which every x86_64 CPU has; when the CPU reports +//! wider vectors, [`matches_into`] and [`count_matches`] pick a faster front end per call +//! (the `wide` module: AVX2 and AVX-512). +//! +//! The wasm layer is compile-checked but not run in CI; the other targets run the full +//! `tests/literal.rs` parity suite. + +#[cfg(target_arch = "aarch64")] +mod arch { + use core::arch::aarch64::*; + + /// 16 bytes of text in one register. + pub(super) type Chunk = uint8x16_t; + /// Offsets sit four bit positions apart in a mask from [`bits`]: offset `j` is bit `4 * j`. + pub(super) const SHIFT: u32 = 2; + + /// One byte copied into every lane. + #[inline(always)] + pub(super) fn splat(byte: u8) -> Chunk { + unsafe { vdupq_n_u8(byte) } + } + + /// All-ones in every lane where a whole `K`-byte match starts. + /// + /// # Safety + /// `p .. p + 16 + K - 1` must be readable: the shifted compares load 16 bytes from each + /// of `p`, `p + 1`, .., `p + K - 1`. + #[inline(always)] + pub(super) unsafe fn match_starts(p: *const u8, pattern: &[Chunk; K]) -> Chunk { + // SAFETY: the caller's contract above; each load of 16 bytes from `p + k`, `k < K`, + // stays inside that readable range. + let mut starts = unsafe { vceqq_u8(vld1q_u8(p), pattern[0]) }; + for (k, &byte) in pattern.iter().enumerate().skip(1) { + starts = unsafe { vandq_u8(starts, vceqq_u8(vld1q_u8(p.add(k)), byte)) }; + } + starts + } + + /// `true` when anything matched in the four chunks. + #[inline(always)] + pub(super) fn any(s0: Chunk, s1: Chunk, s2: Chunk, s3: Chunk) -> bool { + unsafe { vmaxvq_u8(vorrq_u8(vorrq_u8(s0, s1), vorrq_u8(s2, s3))) != 0 } + } + + /// The lane answers as one integer: one bit per lane, four bit positions apart. + /// + /// NEON cannot grab one bit from every lane in a single instruction. Instead `vshrn` + /// ("shift right and narrow") halves the register to 64 bits, shrinking each lane's + /// answer from eight bits to four (still all ones or all zeros); the `0x1111..` mask + /// then keeps one bit of each four, so clearing one bit in the decode drops exactly + /// one match. + #[inline(always)] + pub(super) fn bits(starts: Chunk) -> u64 { + unsafe { + let halved = vshrn_n_u16::<4>(vreinterpretq_u16_u8(starts)); + vget_lane_u64::<0>(vreinterpret_u64_u8(halved)) & 0x1111_1111_1111_1111 + } + } +} + +#[cfg(target_arch = "x86_64")] +mod arch { + use core::arch::x86_64::*; + + /// 16 bytes of text in one register. + pub(super) type Chunk = __m128i; + /// One bit per lane in a mask from [`bits`]: offset `j` is bit `j`, nothing to shift. + pub(super) const SHIFT: u32 = 0; + + /// One byte copied into every lane. + #[inline(always)] + pub(super) fn splat(byte: u8) -> Chunk { + unsafe { _mm_set1_epi8(byte as i8) } + } + + /// All-ones in every lane where a whole `K`-byte match starts. + /// + /// # Safety + /// `p .. p + 16 + K - 1` must be readable: the shifted compares load 16 bytes from each + /// of `p`, `p + 1`, .., `p + K - 1`. + #[inline(always)] + pub(super) unsafe fn match_starts(p: *const u8, pattern: &[Chunk; K]) -> Chunk { + // SAFETY: the caller's contract above; each load of 16 bytes from `p + k`, `k < K`, + // stays inside that readable range. + let mut starts = unsafe { _mm_cmpeq_epi8(_mm_loadu_si128(p.cast()), pattern[0]) }; + for (k, &byte) in pattern.iter().enumerate().skip(1) { + starts = unsafe { + _mm_and_si128( + starts, + _mm_cmpeq_epi8(_mm_loadu_si128(p.add(k).cast()), byte), + ) + }; + } + starts + } + + /// `true` when anything matched in the four chunks. + #[inline(always)] + pub(super) fn any(s0: Chunk, s1: Chunk, s2: Chunk, s3: Chunk) -> bool { + unsafe { _mm_movemask_epi8(_mm_or_si128(_mm_or_si128(s0, s1), _mm_or_si128(s2, s3))) != 0 } + } + + /// The lane answers as one integer: `movemask` grabs each lane's top bit, so offset `j` + /// is bit `j`. + #[inline(always)] + pub(super) fn bits(starts: Chunk) -> u64 { + unsafe { _mm_movemask_epi8(starts) as u32 as u64 } + } +} + +#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] +mod arch { + use core::arch::wasm32::*; + + /// 16 bytes of text in one register. + pub(super) type Chunk = v128; + /// One bit per lane in a mask from [`bits`]: offset `j` is bit `j`, nothing to shift. + pub(super) const SHIFT: u32 = 0; + + /// One byte copied into every lane. + #[inline(always)] + pub(super) fn splat(byte: u8) -> Chunk { + u8x16_splat(byte) + } + + /// All-ones in every lane where a whole `K`-byte match starts. + /// + /// # Safety + /// `p .. p + 16 + K - 1` must be readable: the shifted compares load 16 bytes from each + /// of `p`, `p + 1`, .., `p + K - 1`. + #[inline(always)] + pub(super) unsafe fn match_starts(p: *const u8, pattern: &[Chunk; K]) -> Chunk { + // SAFETY: the caller's contract above; each load of 16 bytes from `p + k`, `k < K`, + // stays inside that readable range. + let mut starts = unsafe { u8x16_eq(v128_load(p.cast()), pattern[0]) }; + for (k, &byte) in pattern.iter().enumerate().skip(1) { + starts = unsafe { v128_and(starts, u8x16_eq(v128_load(p.add(k).cast()), byte)) }; + } + starts + } + + /// `true` when anything matched in the four chunks. + #[inline(always)] + pub(super) fn any(s0: Chunk, s1: Chunk, s2: Chunk, s3: Chunk) -> bool { + v128_any_true(v128_or(v128_or(s0, s1), v128_or(s2, s3))) + } + + /// The lane answers as one integer: `bitmask` grabs each lane's top bit, so offset `j` + /// is bit `j`. + #[inline(always)] + pub(super) fn bits(starts: Chunk) -> u64 { + u8x16_bitmask(starts) as u64 + } +} + +/// Decodes every set bit of `bits` into a `base + offset` value at `cursor`, and returns the +/// cursor advanced by the true match count. +/// +/// The first four decodes run without checking `bits` (the module docs say why): once the mask +/// is used up, `trailing_zeros` is 64 and the slot gets a garbage offset, scratch that the +/// count never covers and a later call overwrites. Blocks with more than four matches finish +/// in the loop. +/// +/// # Safety +/// There must be room for `max(4, count_ones(bits))` writes between `cursor` and `end`: four +/// slots for the unconditional decodes, one per match when a block has more. Debug builds +/// check this bound before writing anything. +#[inline(always)] +unsafe fn decode(mut bits: u64, base: u32, cursor: *mut u32, end: *const u32) -> *mut u32 { + let count = bits.count_ones() as usize; + debug_assert!( + end as usize - cursor as usize >= count.max(4) * size_of::(), + "decode would write past the match buffer" + ); + // SAFETY: the caller's contract above gives these four slots. + for slot in 0..4 { + unsafe { *cursor.add(slot) = base + (bits.trailing_zeros() >> arch::SHIFT) }; + bits &= bits.wrapping_sub(1); + } + if bits != 0 { + let mut slot = 4; + // SAFETY: one slot per remaining match, still within the caller's `count` slots. + while bits != 0 { + unsafe { *cursor.add(slot) = base + (bits.trailing_zeros() >> arch::SHIFT) }; + bits &= bits.wrapping_sub(1); + slot += 1; + } + } + // SAFETY: `count` slots were just written, so one past the last of them is in bounds. + unsafe { cursor.add(count) } +} + +/// How many matches of `pattern` are in `text`: the compare and mask steps only, counting set +/// bits instead of decoding them, which makes counting several times faster than the full +/// scan. Same pattern rules as [`matches_into`]. +pub(crate) fn count_matches(pattern: [u8; K], text: &[u8]) -> usize { + const { + assert!( + K >= 1 && K <= 3, + "the scan is built for patterns of 1 to 3 bytes" + ); + } + let len = text.len(); + if len < K { + return 0; + } + #[cfg(target_arch = "x86_64")] + if let Some(count) = wide::count_matches_wide::(pattern, text) { + return count; + } + let p = text.as_ptr(); + let mut pattern_chunks = [arch::splat(0); K]; + for (chunk, &byte) in pattern_chunks.iter_mut().zip(&pattern) { + *chunk = arch::splat(byte); + } + let mut count = 0usize; + let mut i = 0usize; + + // The same block layout as `scan`, and the same SAFETY story for the reads: each loop + // condition bounds its loads inside `text`. Nothing here writes. + while i + 64 + (K - 1) <= len { + let (s0, s1, s2, s3) = unsafe { + ( + arch::match_starts::(p.add(i), &pattern_chunks), + arch::match_starts::(p.add(i + 16), &pattern_chunks), + arch::match_starts::(p.add(i + 32), &pattern_chunks), + arch::match_starts::(p.add(i + 48), &pattern_chunks), + ) + }; + if arch::any(s0, s1, s2, s3) { + count += (arch::bits(s0).count_ones() + + arch::bits(s1).count_ones() + + arch::bits(s2).count_ones() + + arch::bits(s3).count_ones()) as usize; + } + i += 64; + } + + count + unsafe { count_tail::(pattern, &pattern_chunks, text, i) } +} + +/// The counting twin of [`scan_tail`]: 16-byte rounds from `i`, the overlapped final block, +/// or the byte-by-byte walk for texts shorter than one block. +/// +/// # Safety +/// Reads are bounded by the loop conditions, exactly as in [`scan_tail`]. Nothing writes. +unsafe fn count_tail( + pattern: [u8; K], + pattern_chunks: &[arch::Chunk; K], + text: &[u8], + mut i: usize, +) -> usize { + let len = text.len(); + let p = text.as_ptr(); + let mut count = 0usize; + + // SAFETY: reads end at `i + 16 + (K - 1) <= len`. + while i + 16 + (K - 1) <= len { + let starts = unsafe { arch::match_starts::(p.add(i), pattern_chunks) }; + count += arch::bits(starts).count_ones() as usize; + i += 16; + } + + if i + K <= len { + if len >= 16 + (K - 1) { + // The overlapped final block: bits below `i` were counted above and are masked off. + // SAFETY: the block reads `base .. base + 16 + K - 1`, which ends exactly at `len`. + let base = len - 16 - (K - 1); + let starts = unsafe { arch::match_starts::(p.add(base), pattern_chunks) }; + let bits = arch::bits(starts) & (!0u64 << ((i - base) << arch::SHIFT)); + count += bits.count_ones() as usize; + } else { + while i + K <= len { + count += (text[i..i + K] == pattern[..]) as usize; + i += 1; + } + } + } + count +} + +/// The wider x86 front ends. Every x86_64 CPU has the SSE2 baseline above; these two rungs +/// are picked per call when the CPU reports the feature: +/// +/// - **AVX2** compares 32 bytes per instruction, so one round covers 64 bytes with two +/// compares and one combined mask. +/// - **AVX-512** compares 64 bytes straight into a bit mask (`_mm512_cmpeq_epi8_mask`, no +/// separate mask step) and decodes with `vpcompressd`: the CPU packs the matching offsets +/// itself, 16 at a time, replacing the whole count-trailing-zeros loop. +/// +/// Each rung is its own function because `#[target_feature]` applies per function; the bodies +/// mirror [`scan`] and [`count_matches`] round for round and share [`decode`], [`scan_tail`] +/// and [`count_tail`]. The AVX2 rung is exercised by the `tests/literal.rs` suite on any +/// AVX2 machine or emulator; the AVX-512 rung runs the same suite only on hardware that has +/// it, so run the tests on such a machine before trusting changes to it. +#[cfg(target_arch = "x86_64")] +mod wide { + use super::{arch, count_tail, decode, scan_tail}; + use core::arch::x86_64::*; + + /// Runs the widest scan this CPU supports; `None` means only the baseline exists. + #[inline] + pub(super) fn matches_into_wide( + pattern: [u8; K], + text: &[u8], + out: &mut [u32], + ) -> Option { + if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") { + // SAFETY: the features were just detected; the buffer contract is the caller's. + return Some(unsafe { scan_avx512::(pattern, text, out) }); + } + if is_x86_feature_detected!("avx2") { + // SAFETY: the feature was just detected; the buffer contract is the caller's. + return Some(unsafe { scan_avx2::(pattern, text, out) }); + } + None + } + + /// The counting twin of [`matches_into_wide`]. + #[inline] + pub(super) fn count_matches_wide( + pattern: [u8; K], + text: &[u8], + ) -> Option { + if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") { + // SAFETY: the features were just detected. Counting never writes. + return Some(unsafe { count_avx512::(pattern, text) }); + } + if is_x86_feature_detected!("avx2") { + // SAFETY: the feature was just detected. Counting never writes. + return Some(unsafe { count_avx2::(pattern, text) }); + } + None + } + + /// All-ones per lane where a whole `K`-byte match starts, for 32 positions at `p`. + /// + /// # Safety + /// AVX2 must be available, and `p .. p + 32 + K - 1` must be readable. + #[inline] + #[target_feature(enable = "avx2")] + unsafe fn match_starts32(p: *const u8, pattern: &[__m256i; K]) -> __m256i { + // SAFETY: the caller's contract above; each load of 32 bytes from `p + k`, `k < K`, + // stays inside that readable range. + let mut starts = _mm256_cmpeq_epi8(unsafe { _mm256_loadu_si256(p.cast()) }, pattern[0]); + for (k, &byte) in pattern.iter().enumerate().skip(1) { + starts = _mm256_and_si256( + starts, + _mm256_cmpeq_epi8(unsafe { _mm256_loadu_si256(p.add(k).cast()) }, byte), + ); + } + starts + } + + /// One byte copied into every lane, and the baseline chunks [`scan_tail`] wants. + #[target_feature(enable = "avx2")] + fn splat32(pattern: [u8; K]) -> ([__m256i; K], [arch::Chunk; K]) { + let mut wide = [_mm256_set1_epi8(0); K]; + let mut baseline = [arch::splat(0); K]; + for k in 0..K { + wide[k] = _mm256_set1_epi8(pattern[k] as i8); + baseline[k] = arch::splat(pattern[k]); + } + (wide, baseline) + } + + /// [`scan`](super::scan) with the AVX2 front end: 64 bytes per round as two 32-byte + /// compares, one combined mask, decoded in the same 16-bit quarters as the baseline so + /// [`decode`]'s four-slot unroll stays matched to typical match density. + /// + /// # Safety + /// AVX2 must be available, and `out` must satisfy + /// [`matches_into`](super::matches_into)'s buffer bound, which it asserts. + #[target_feature(enable = "avx2")] + unsafe fn scan_avx2(pattern: [u8; K], text: &[u8], out: &mut [u32]) -> usize { + let len = text.len(); + if len < K { + return 0; + } + let p = text.as_ptr(); + let (pattern_chunks, baseline) = splat32::(pattern); + let start = out.as_mut_ptr(); + // SAFETY: one past the buffer's last slot is a valid address to form (never read). + let end = unsafe { start.add(out.len()) } as *const u32; + let mut cursor = start; + let mut i = 0usize; + + // SAFETY: reads end at `i + 32 + 32 + (K - 1) <= len`; writes are the buffer + // contract, handed to `decode` as `end`. + while i + 64 + (K - 1) <= len { + let (s0, s1) = unsafe { + ( + match_starts32::(p.add(i), &pattern_chunks), + match_starts32::(p.add(i + 32), &pattern_chunks), + ) + }; + let bits = (_mm256_movemask_epi8(s0) as u32 as u64) + | ((_mm256_movemask_epi8(s1) as u32 as u64) << 32); + if bits != 0 { + unsafe { + cursor = decode(bits & 0xFFFF, i as u32, cursor, end); + cursor = decode((bits >> 16) & 0xFFFF, i as u32 + 16, cursor, end); + cursor = decode((bits >> 32) & 0xFFFF, i as u32 + 32, cursor, end); + cursor = decode(bits >> 48, i as u32 + 48, cursor, end); + } + } + i += 64; + } + // SAFETY: `scan_tail` continues under the same contract. + cursor = unsafe { scan_tail::(pattern, &baseline, text, i, cursor, end) }; + // SAFETY: `cursor` and `start` both point into `out`'s buffer. + unsafe { cursor.offset_from(start) as usize } + } + + /// The counting twin of [`scan_avx2`]. + /// + /// # Safety + /// AVX2 must be available. + #[target_feature(enable = "avx2")] + unsafe fn count_avx2(pattern: [u8; K], text: &[u8]) -> usize { + let len = text.len(); + if len < K { + return 0; + } + let p = text.as_ptr(); + let (pattern_chunks, baseline) = splat32::(pattern); + let mut count = 0usize; + let mut i = 0usize; + + // SAFETY: reads end at `i + 64 + (K - 1) <= len`. Nothing writes. + while i + 64 + (K - 1) <= len { + let (s0, s1) = unsafe { + ( + match_starts32::(p.add(i), &pattern_chunks), + match_starts32::(p.add(i + 32), &pattern_chunks), + ) + }; + count += (_mm256_movemask_epi8(s0).count_ones() + _mm256_movemask_epi8(s1).count_ones()) + as usize; + i += 64; + } + // SAFETY: `count_tail` continues under the same read bounds. + count + unsafe { count_tail::(pattern, &baseline, text, i) } + } + + /// The 64-byte match mask, straight from the compares: AVX-512 byte compares produce a + /// bit per lane natively, so there is no separate mask step to pay for. + /// + /// # Safety + /// AVX-512 F and BW must be available, and `p .. p + 64 + K - 1` must be readable. + #[inline] + #[target_feature(enable = "avx512f,avx512bw")] + unsafe fn match_bits64(p: *const u8, pattern: &[__m512i; K]) -> u64 { + // SAFETY: the caller's contract above; each load of 64 bytes from `p + k`, `k < K`, + // stays inside that readable range. + let mut bits = _mm512_cmpeq_epi8_mask(unsafe { _mm512_loadu_si512(p.cast()) }, pattern[0]); + for (k, &byte) in pattern.iter().enumerate().skip(1) { + bits &= _mm512_cmpeq_epi8_mask(unsafe { _mm512_loadu_si512(p.add(k).cast()) }, byte); + } + bits + } + + /// One byte copied into every lane, and the baseline chunks the shared tails want. + #[target_feature(enable = "avx512f")] + fn splat64(pattern: [u8; K]) -> ([__m512i; K], [arch::Chunk; K]) { + let mut wide = [_mm512_set1_epi8(0); K]; + let mut baseline = [arch::splat(0); K]; + for k in 0..K { + wide[k] = _mm512_set1_epi8(pattern[k] as i8); + baseline[k] = arch::splat(pattern[k]); + } + (wide, baseline) + } + + /// [`scan`](super::scan) with the AVX-512 front end. The decode is `vpcompressd` + /// (`_mm512_mask_compressstoreu_epi32`): give the CPU sixteen candidate offsets and a + /// mask, and it stores exactly the matching ones, packed. The count-trailing-zeros loop + /// and its density concerns disappear; four compress-stores cover a 64-byte round. + /// + /// # Safety + /// AVX-512 F and BW must be available, and `out` must satisfy + /// [`matches_into`](super::matches_into)'s buffer bound, which it asserts. + #[target_feature(enable = "avx512f,avx512bw")] + unsafe fn scan_avx512(pattern: [u8; K], text: &[u8], out: &mut [u32]) -> usize { + let len = text.len(); + if len < K { + return 0; + } + let p = text.as_ptr(); + let (pattern_chunks, baseline) = splat64::(pattern); + let start = out.as_mut_ptr(); + // SAFETY: one past the buffer's last slot is a valid address to form (never read). + let end = unsafe { start.add(out.len()) } as *const u32; + let mut cursor = start; + let mut i = 0usize; + // Lane numbers 0..16, the base of every offset vector. + let lanes = _mm512_set_epi32(15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0); + + // SAFETY: reads end at `i + 64 + (K - 1) <= len`. Each compress-store writes exactly + // the matches of its quarter, within the buffer contract. + while i + 64 + (K - 1) <= len { + let bits = unsafe { match_bits64::(p.add(i), &pattern_chunks) }; + if bits != 0 { + let base = _mm512_add_epi32(lanes, _mm512_set1_epi32(i as i32)); + for quarter in 0..4 { + let mask = ((bits >> (16 * quarter)) & 0xFFFF) as __mmask16; + let offsets = _mm512_add_epi32(base, _mm512_set1_epi32(16 * quarter)); + debug_assert!( + end as usize - cursor as usize + >= mask.count_ones() as usize * size_of::(), + "compress-store would write past the match buffer" + ); + unsafe { + _mm512_mask_compressstoreu_epi32(cursor.cast(), mask, offsets); + cursor = cursor.add(mask.count_ones() as usize); + } + } + } + i += 64; + } + // SAFETY: `scan_tail` continues under the same contract. + cursor = unsafe { scan_tail::(pattern, &baseline, text, i, cursor, end) }; + // SAFETY: `cursor` and `start` both point into `out`'s buffer. + unsafe { cursor.offset_from(start) as usize } + } + + /// The counting twin of [`scan_avx512`]. + /// + /// # Safety + /// AVX-512 F and BW must be available. + #[target_feature(enable = "avx512f,avx512bw")] + unsafe fn count_avx512(pattern: [u8; K], text: &[u8]) -> usize { + let len = text.len(); + if len < K { + return 0; + } + let p = text.as_ptr(); + let (pattern_chunks, baseline) = splat64::(pattern); + let mut count = 0usize; + let mut i = 0usize; + + // SAFETY: reads end at `i + 64 + (K - 1) <= len`. Nothing writes. + while i + 64 + (K - 1) <= len { + count += unsafe { match_bits64::(p.add(i), &pattern_chunks) }.count_ones() as usize; + i += 64; + } + // SAFETY: `count_tail` continues under the same read bounds. + count + unsafe { count_tail::(pattern, &baseline, text, i) } + } +} + +/// One-pass scan: writes the start offset of every match of `pattern` in `text` into `out` and +/// returns how many. The pattern must not be able to overlap itself, so that every matching +/// position belongs to the non-overlapping match list; +/// [`Literal::matches_into`](crate::literal::Literal::matches_into) routes self-overlapping +/// patterns away. +/// +/// The buffer bound is asserted here, not trusted from the caller: everything unsafe in this +/// module leans on it, so the check lives next to what it protects. +pub(crate) fn matches_into( + pattern: [u8; K], + text: &[u8], + out: &mut [u32], +) -> usize { + const { + assert!( + K >= 1 && K <= 3, + "the scan is built for patterns of 1 to 3 bytes" + ); + } + assert!( + out.len() >= text.len() / K + 4, + "the match buffer must hold text.len() / pattern.len() + 4 offsets" + ); + #[cfg(target_arch = "x86_64")] + if let Some(count) = wide::matches_into_wide::(pattern, text, out) { + return count; + } + // SAFETY: the assert above is exactly `scan`'s buffer contract. + unsafe { scan::(pattern, text, out) } +} + +/// # Safety +/// `out.len() >= text.len() / K + 4` must hold; [`matches_into`] asserts it. Every match found +/// counts one slot and matches cannot overlap, so at most `text.len() / K` slots are counted +/// and the four extra keep `decode`'s unconditional writes inside the buffer. Reads from +/// `text` are justified region by region below. +unsafe fn scan(pattern: [u8; K], text: &[u8], out: &mut [u32]) -> usize { + let len = text.len(); + if len < K { + return 0; + } + let p = text.as_ptr(); + let mut pattern_chunks = [arch::splat(0); K]; + for (chunk, &byte) in pattern_chunks.iter_mut().zip(&pattern) { + *chunk = arch::splat(byte); + } + let start = out.as_mut_ptr(); + // SAFETY: one past the buffer's last slot is a valid address to form (never read); + // only `decode`'s debug bound check uses it. + let end = unsafe { start.add(out.len()) } as *const u32; + let mut cursor = start; + let mut i = 0usize; + + // 64 bytes per round; `any` lets a match-free round skip the masks and decodes. + // SAFETY: the loop condition keeps every read inside `text`: the furthest `match_starts` + // begins at `p + i + 48` and reads `16 + K - 1` bytes, ending at `i + 64 + (K - 1) <= len`. + // Writes are `scan`'s buffer contract, handed to `decode` as `end`. + while i + 64 + (K - 1) <= len { + let (s0, s1, s2, s3) = unsafe { + ( + arch::match_starts::(p.add(i), &pattern_chunks), + arch::match_starts::(p.add(i + 16), &pattern_chunks), + arch::match_starts::(p.add(i + 32), &pattern_chunks), + arch::match_starts::(p.add(i + 48), &pattern_chunks), + ) + }; + if arch::any(s0, s1, s2, s3) { + unsafe { + cursor = decode(arch::bits(s0), i as u32, cursor, end); + cursor = decode(arch::bits(s1), i as u32 + 16, cursor, end); + cursor = decode(arch::bits(s2), i as u32 + 32, cursor, end); + cursor = decode(arch::bits(s3), i as u32 + 48, cursor, end); + } + } + i += 64; + } + + cursor = unsafe { scan_tail::(pattern, &pattern_chunks, text, i, cursor, end) }; + // SAFETY: `cursor` and `start` both point into `out`'s buffer. + unsafe { cursor.offset_from(start) as usize } +} + +/// The remainder every scan front end shares, once its 64-byte rounds got to `i`: 16-byte +/// rounds, then one overlapped final block, or a byte-by-byte walk when the whole text is +/// shorter than a block. Returns the advanced cursor. +/// +/// # Safety +/// Same contract as [`scan`]: reads are bounded by the loop conditions, and `cursor .. end` +/// must have room for every remaining match plus `decode`'s four-slot slack. +unsafe fn scan_tail( + pattern: [u8; K], + pattern_chunks: &[arch::Chunk; K], + text: &[u8], + mut i: usize, + mut cursor: *mut u32, + end: *const u32, +) -> *mut u32 { + let len = text.len(); + let p = text.as_ptr(); + + // 16 bytes per round over what the 64-byte rounds left. + // SAFETY: reads end at `i + 16 + (K - 1) <= len`; writes as in [`scan`]. + while i + 16 + (K - 1) <= len { + cursor = unsafe { + decode( + arch::bits(arch::match_starts::(p.add(i), pattern_chunks)), + i as u32, + cursor, + end, + ) + }; + i += 16; + } + + if i + K <= len { + if len >= 16 + (K - 1) { + // One last 16-byte block, moved back to end flush with the last position a match + // can still start at. Its low bits re-cover positions the rounds above already + // decoded, so they are masked off. + // SAFETY: the block reads `base .. base + 16 + K - 1`, which ends exactly at `len`; + // writes as in [`scan`]. + let base = len - 16 - (K - 1); + unsafe { + let mut bits = arch::bits(arch::match_starts::(p.add(base), pattern_chunks)); + bits &= !0u64 << ((i - base) << arch::SHIFT); + cursor = decode(bits, base as u32, cursor, end); + } + } else { + // The whole text is shorter than one block: byte-by-byte costs nothing here. + while i + K <= len { + if text[i..i + K] == pattern[..] { + debug_assert!(end as usize - cursor as usize >= size_of::()); + // SAFETY: one slot per match; the buffer holds a slot for every possible + // match plus four ([`matches_into`]'s assert). + unsafe { + *cursor = i as u32; + cursor = cursor.add(1); + } + } + i += 1; + } + } + } + cursor +} diff --git a/tokenizers/atomsplit/tests/literal.rs b/tokenizers/atomsplit/tests/literal.rs index 14ceed9a0..894dde46a 100644 --- a/tokenizers/atomsplit/tests/literal.rs +++ b/tokenizers/atomsplit/tests/literal.rs @@ -48,3 +48,203 @@ fn an_empty_pattern_is_rejected() { "an empty pattern matches everywhere" ); } + +// ---- the batch scan, `matches_into` ---- + +/// Run `matches_into` with a buffer of exactly the documented size and return what it wrote. +fn batch(literal: &Literal, text: &[u8]) -> Vec { + let mut out = vec![0u32; text.len() / literal.pattern().len() + 4]; + let count = literal.matches_into(text, &mut out); + out.truncate(count); + out +} + +fn iterated(literal: &Literal, text: &[u8]) -> Vec { + literal.matches(text).map(|p| p as u32).collect() +} + +#[test] +fn matches_into_reports_the_iterator_positions() { + let dash = Literal::new(b"-").unwrap(); + assert_eq!(batch(&dash, b"a-b--c"), [1, 3, 4]); + assert_eq!(batch(&dash, b"-starts and ends-"), [0, 16]); + assert_eq!(batch(&dash, b"none here"), []); + assert_eq!(batch(&dash, b""), []); + + let space = Literal::new(b" ").unwrap(); + let text = b"the quick brown fox jumps over the lazy dog".repeat(40); + assert_eq!(batch(&space, &text), iterated(&space, &text)); + + let metaspace = Literal::new("\u{2581}".as_bytes()).unwrap(); + let text = "\u{2581}the\u{2581}quick\u{2581}brown\u{2581}fox".repeat(40); + assert_eq!( + batch(&metaspace, text.as_bytes()), + iterated(&metaspace, text.as_bytes()) + ); +} + +/// The scan works in blocks; a match starting on either side of an internal block edge, or on +/// the last position where the pattern still fits, must not be missed or doubled. +#[test] +fn matches_into_is_seamless_across_block_boundaries() { + for pattern in [&b"-"[..], "\u{2581}".as_bytes()] { + let literal = Literal::new(pattern).unwrap(); + for start in [0, 13, 14, 15, 16, 17, 61, 62, 63, 64, 65, 127, 128] { + for len in [start + pattern.len(), 90, 129, 200] { + if start + pattern.len() > len { + continue; + } + let mut text = vec![b'x'; len]; + text[start..start + pattern.len()].copy_from_slice(pattern); + assert_eq!( + batch(&literal, &text), + [start as u32], + "pattern {pattern:?} at {start} in {len} bytes" + ); + } + } + } +} + +/// Every position matches: the largest count a text can produce, written into a buffer of +/// exactly the documented size. +#[test] +fn matches_into_handles_the_densest_text() { + let a = Literal::new(b"a").unwrap(); + let text = vec![b'a'; 300]; + assert_eq!(batch(&a, &text), (0..300).collect::>()); + + let metaspace = Literal::new("\u{2581}".as_bytes()).unwrap(); + let text = "\u{2581}".repeat(100); + assert_eq!( + batch(&metaspace, text.as_bytes()), + (0..100).map(|i| i * 3).collect::>() + ); +} + +#[test] +fn a_self_overlapping_pattern_keeps_non_overlapping_matches() { + let aa = Literal::new(b"aa").unwrap(); + assert_eq!(batch(&aa, b"aaa"), [0]); + let run = [b'a'; 100]; + assert_eq!(batch(&aa, &run), iterated(&aa, &run)); + + let aba = Literal::new(b"aba").unwrap(); + assert_eq!(batch(&aba, b"ababa"), [0]); +} + +#[test] +fn a_pattern_longer_than_three_bytes_still_matches() { + let literal = Literal::new(b"abcd").unwrap(); + let text = b"abcd xabcdx abcabcd".repeat(20); + assert_eq!(batch(&literal, &text), iterated(&literal, &text)); +} + +#[test] +#[should_panic(expected = "matches_into")] +fn an_undersized_buffer_is_rejected() { + let a = Literal::new(b"a").unwrap(); + let mut out = vec![0u32; 7]; // "aaaa" needs 4 / 1 + 4 = 8 + a.matches_into(b"aaaa", &mut out); +} + +// ---- the streaming pair, `count_matches` + `for_each_match` ---- + +fn streamed(literal: &Literal, text: &[u8]) -> Vec { + let mut out = Vec::new(); + literal.for_each_match(text, |start| out.push(start as u32)); + out +} + +/// The streaming scan works through a fixed window; a text much longer than any plausible +/// window must report the same matches as the iterator, including around every window edge. +#[test] +fn for_each_match_streams_the_iterator_offsets_across_windows() { + for pattern in [&b" "[..], "\u{2581}".as_bytes(), b"ab"] { + let literal = Literal::new(pattern).unwrap(); + let mut text = b"word ".repeat(4_000); // ~20KB, matches every 5 bytes + text.extend_from_slice("\u{2581}ab".repeat(2_000).as_bytes()); + assert_eq!(streamed(&literal, &text), iterated(&literal, &text)); + assert_eq!( + literal.count_matches(&text), + literal.matches(&text).count(), + "count for {pattern:?}" + ); + } +} + +/// The densest possible long text: every byte starts a match, in every window. +#[test] +fn for_each_match_handles_dense_windows() { + let a = Literal::new(b"a").unwrap(); + let text = vec![b'a'; 10_000]; + assert_eq!(streamed(&a, &text), (0..10_000).collect::>()); + assert_eq!(a.count_matches(&text), 10_000); +} + +#[test] +fn for_each_match_takes_the_iterator_route_for_uncovered_patterns() { + // Self-overlapping and longer-than-three patterns are exactly the ones the batch scan + // refuses; the streaming pair must still report the iterator's non-overlapping matches. + for pattern in [&b"aa"[..], b"aba", b"abcd"] { + let literal = Literal::new(pattern).unwrap(); + let text = b"aabaabcdaaaba".repeat(2_000); + assert_eq!(streamed(&literal, &text), iterated(&literal, &text)); + assert_eq!(literal.count_matches(&text), literal.matches(&text).count()); + } +} + +#[test] +fn streaming_agrees_with_the_iterator_on_random_lengths() { + let mut state = 0x1234_5678_9ABC_DEF0u64; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state + }; + let patterns: &[&[u8]] = &[b"a", b"ab", b"aa", b"aba", b"abc", b"abca"]; + for _ in 0..500 { + // lengths spread around plausible window sizes, so edges get hit both exactly and off by one + let len = (next() % 8_000) as usize; + let text: Vec = (0..len).map(|_| b'a' + (next() % 3) as u8).collect(); + let literal = Literal::new(patterns[(next() % patterns.len() as u64) as usize]).unwrap(); + assert_eq!( + streamed(&literal, &text), + iterated(&literal, &text), + "pattern {:?} len {len}", + literal.pattern() + ); + assert_eq!(literal.count_matches(&text), literal.matches(&text).count()); + } +} + +#[test] +fn matches_into_agrees_with_the_iterator_on_random_inputs() { + // xorshift64: deterministic, no dev-dependency needed. + let mut state = 0x9E37_79B9_7F4A_7C15u64; + let mut next = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state + }; + let patterns: &[&[u8]] = &[ + b"a", b"b", b"ab", b"aa", b"ba", b"aab", b"aba", b"abc", b"abca", + ]; + for _ in 0..20_000 { + let len = (next() % 300) as usize; + let alphabet = 2 + (next() % 2) as u8; + let text: Vec = (0..len) + .map(|_| b'a' + (next() % alphabet as u64) as u8) + .collect(); + let literal = Literal::new(patterns[(next() % patterns.len() as u64) as usize]).unwrap(); + assert_eq!( + batch(&literal, &text), + iterated(&literal, &text), + "pattern {:?} text {:?}", + literal.pattern(), + text + ); + } +} diff --git a/tokenizers/tk-encode/src/normalizers/replace.rs b/tokenizers/tk-encode/src/normalizers/replace.rs index cc4545b47..e23d20399 100644 --- a/tokenizers/tk-encode/src/normalizers/replace.rs +++ b/tokenizers/tk-encode/src/normalizers/replace.rs @@ -145,10 +145,22 @@ impl pipeline::Normalizer for Replace { Ok(match &self.search { Search::Literal(literal) => { let width = literal.pattern().len(); - let matches = literal - .matches(input.as_bytes()) - .map(|start| (start, start + width)); - replace_matches(input, &self.content, matches) + let count = literal.count_matches(input.as_bytes()); + if count == 0 { + return Ok(Cow::Borrowed(input)); + } + // Counting first buys the exact output length (every match swaps `width` + // bytes for the content), so the build below never reallocates. + let exact = input.len() - count * width + count * self.content.len(); + let mut replaced = String::with_capacity(exact); + let mut last_end = 0; + literal.for_each_match(input.as_bytes(), |start| { + replaced.push_str(&input[last_end..start]); + replaced.push_str(&self.content); + last_end = start + width; + }); + replaced.push_str(&input[last_end..]); + Cow::Owned(replaced) } Search::Regex(regex) => replace_matches(input, &self.content, regex.find_iter(input)), Search::Nothing => Cow::Borrowed(input), diff --git a/tokenizers/tk-encode/src/pre_tokenizers/split.rs b/tokenizers/tk-encode/src/pre_tokenizers/split.rs index 779c167f1..bc5682cad 100644 --- a/tokenizers/tk-encode/src/pre_tokenizers/split.rs +++ b/tokenizers/tk-encode/src/pre_tokenizers/split.rs @@ -201,13 +201,33 @@ impl pipeline::PreTokenizer for Split { ); return Ok(()); } - // Not a natively-routed GPT regex: fall-back to Literal or Regex search + // A plain-string pattern streams: the batch scan hands each delimiter offset straight + // to the span fold, and nothing is built in between. The count pass sizes `out` for + // the worst case, one span per match plus the pieces around them. + if let Search::Literal(literal) = &self.search { + let width = literal.pattern().len(); + let count = literal.count_matches(text.as_bytes()); + out.reserve(2 * count + 1); + let mut fold = pipeline::SplitFold::new(out, self.behavior); + let mut prev = 0; + literal.for_each_match(text.as_bytes(), |start| { + if prev != start { + fold.segment((prev, start), self.invert); + } + fold.segment((start, start + width), !self.invert); + prev = start + width; + }); + if prev != text.len() { + fold.segment((prev, text.len()), self.invert); + } + fold.finish(); + return Ok(()); + } + // Not a natively-routed GPT regex either: fall back to the Regex search let matches = match (&self.search, self.invert) { - (Search::Literal(literal), false) => literal.find_matches(text)?, - (Search::Literal(literal), true) => Invert(literal).find_matches(text)?, (Search::Regex(regex), false) => regex.find_matches(text)?, (Search::Regex(regex), true) => Invert(regex).find_matches(text)?, - (Search::Unavailable, _) => { + _ => { return Err( "this `Split` pattern needs a system-regex backend; enable the `fancy-regex` feature" .into(), @@ -415,6 +435,84 @@ mod tests { ); } + /// The literal path across every behavior, on a text with a delimiter at both ends and a + /// consecutive pair — the shapes the span fold can get wrong. + #[test] + fn pipeline_literal_covers_every_behavior() { + let text = "-a--b-"; + #[allow(clippy::type_complexity)] + let cases: Vec<(SplitDelimiterBehavior, Vec<(&str, (u32, u32))>)> = vec![ + (Removed, vec![("a", (1, 2)), ("b", (4, 5))]), + ( + Isolated, + vec![ + ("-", (0, 1)), + ("a", (1, 2)), + ("-", (2, 3)), + ("-", (3, 4)), + ("b", (4, 5)), + ("-", (5, 6)), + ], + ), + ( + MergedWithPrevious, + vec![("-", (0, 1)), ("a-", (1, 3)), ("-", (3, 4)), ("b-", (4, 6))], + ), + ( + MergedWithNext, + vec![("-a", (0, 2)), ("-", (2, 3)), ("-b", (3, 5)), ("-", (5, 6))], + ), + ( + Contiguous, + vec![ + ("-", (0, 1)), + ("a", (1, 2)), + ("--", (2, 4)), + ("b", (4, 5)), + ("-", (5, 6)), + ], + ), + ]; + for (behavior, expected) in cases { + assert_eq!( + pipeline_split("-".into(), behavior, false, text), + expected, + "behavior: {behavior:?}", + ); + } + } + + /// Inverted literal search: the gaps between delimiters become the matches. + #[test] + fn pipeline_literal_inverts() { + // Removed drops the (inverted) matches, keeping each delimiter. + assert_eq!( + pipeline_split("-".into(), Removed, true, "-a--b-"), + vec![("-", (0, 1)), ("-", (2, 3)), ("-", (3, 4)), ("-", (5, 6))], + ); + // Isolated keeps everything either way. + assert_eq!( + pipeline_split("-".into(), Isolated, true, "a-b"), + vec![("a", (0, 1)), ("-", (1, 2)), ("b", (2, 3))], + ); + } + + #[test] + fn pipeline_literal_edges() { + assert_eq!( + pipeline_split("-".into(), Removed, false, ""), + Vec::<(&str, (u32, u32))>::new(), + ); + assert_eq!( + pipeline_split("-".into(), Removed, false, "---"), + Vec::<(&str, (u32, u32))>::new(), + ); + assert_eq!( + pipeline_split("-".into(), Removed, false, "abc"), + vec![("abc", (0, 3))], + ); + } + #[test] fn pipeline_gpt2_uses_fsm_and_matches_legacy() { // The gpt2 pattern is recognized -> the pipeline path routes to the native diff --git a/tokenizers/tk-encode/src/tokenizer/pattern.rs b/tokenizers/tk-encode/src/tokenizer/pattern.rs index 5a147f5f4..a411ef265 100644 --- a/tokenizers/tk-encode/src/tokenizer/pattern.rs +++ b/tokenizers/tk-encode/src/tokenizer/pattern.rs @@ -43,22 +43,27 @@ impl Pattern for &Regex { } /// Searching for a plain string, does not need a regex engine: [`Literal`] scans the bytes. +/// Two batch-scan passes: [`Literal::count_matches`] sizes `splits` exactly (each match adds +/// itself and at most one gap before it, plus the final gap), then +/// [`Literal::for_each_match`] streams the offsets in. impl Pattern for &Literal { fn find_matches(&self, inside: &str) -> Result> { if inside.is_empty() { return Ok(vec![((0, 0), false)]); } + let width = self.pattern().len(); + let count = self.count_matches(inside.as_bytes()); + let mut splits = Vec::with_capacity(2 * count + 1); let mut prev = 0; - let mut splits = Vec::with_capacity(inside.len()); - for start in self.matches(inside.as_bytes()) { - let end = start + self.pattern().len(); + self.for_each_match(inside.as_bytes(), |start| { + let end = start + width; if prev != start { splits.push(((prev, start), false)); } splits.push(((start, end), true)); prev = end; - } + }); if prev != inside.len() { splits.push(((prev, inside.len()), false)) } diff --git a/tokenizers/tk-encode/src/tokenizer/pipeline.rs b/tokenizers/tk-encode/src/tokenizer/pipeline.rs index 16b7430b7..d81d3b12a 100644 --- a/tokenizers/tk-encode/src/tokenizer/pipeline.rs +++ b/tokenizers/tk-encode/src/tokenizer/pipeline.rs @@ -907,75 +907,101 @@ pub fn split_matches( matches: Vec<((usize, usize), bool)>, behavior: SplitDelimiterBehavior, ) { - use SplitDelimiterBehavior::*; + let mut fold = SplitFold::new(out, behavior); + for (offsets, is_match) in matches { + fold.segment(offsets, is_match); + } + fold.finish(); +} - // (offsets, should_remove) — mirrors `NormalizedString::split`. - let splits: Vec<((usize, usize), bool)> = match behavior { - Isolated => matches.into_iter().map(|(o, _)| (o, false)).collect(), - Removed => matches, // should_remove == is_match - Contiguous => { - let mut previous_match = false; - matches - .into_iter() - .fold(vec![], |mut acc, (offsets, is_match)| { - if is_match == previous_match { - if let Some(((_, end), _)) = acc.last_mut() { - *end = offsets.1; - } else { - acc.push((offsets, false)); - } - } else { - acc.push((offsets, false)); - } - previous_match = is_match; - acc - }) +/// The streaming form of [`split_matches`]: feed it the covering `(offsets, is_match)` +/// segments left to right with [`SplitFold::segment`], then call [`SplitFold::finish`]. +/// Callers that produce segments one at a time (the literal path in `Split`) drive it +/// directly and never build the segment vector. +/// +/// One span under construction and the previous segment's flag are the only state. A span +/// stays pending as long as the next segment might still extend it (a delimiter merging with +/// its neighbour, a contiguous run); a segment that cannot extend it flushes it to `out` and +/// takes its place. +pub(crate) struct SplitFold<'a> { + out: &'a mut Vec, + behavior: SplitDelimiterBehavior, + pending: Option<(usize, usize)>, + previous_match: bool, +} + +impl<'a> SplitFold<'a> { + pub(crate) fn new(out: &'a mut Vec, behavior: SplitDelimiterBehavior) -> Self { + Self { + out, + behavior, + pending: None, + previous_match: false, } - MergedWithPrevious => { - let mut previous_match = false; - matches - .into_iter() - .fold(vec![], |mut acc, (offsets, is_match)| { - if is_match && !previous_match { - if let Some(((_, end), _)) = acc.last_mut() { - *end = offsets.1; - } else { - acc.push((offsets, false)); - } - } else { - acc.push((offsets, false)); - } - previous_match = is_match; - acc - }) + } + + pub(crate) fn segment(&mut self, offsets: (usize, usize), is_match: bool) { + use SplitDelimiterBehavior::*; + match self.behavior { + Isolated => self.emit(Some(offsets)), + Removed => { + if !is_match { + self.emit(Some(offsets)); + } + } + // A run of equal flags becomes one span. + Contiguous => { + if is_match == self.previous_match { + self.extend_pending_to(offsets); + } else { + let done = self.pending.replace(offsets); + self.emit(done); + } + } + // A delimiter glues onto the piece before it; in a run of delimiters only the + // first one glues, the rest stand alone. + MergedWithPrevious => { + if is_match && !self.previous_match { + self.extend_pending_to(offsets); + } else { + let done = self.pending.replace(offsets); + self.emit(done); + } + } + // A delimiter glues onto the piece after it, so it stays pending until that + // piece arrives; a delimiter followed by another delimiter stands alone. + MergedWithNext => { + if !is_match && self.previous_match { + self.extend_pending_to(offsets); + } else { + let done = self.pending.replace(offsets); + self.emit(done); + } + } } - MergedWithNext => { - let mut previous_match = false; - let mut splits = - matches - .into_iter() - .rev() - .fold(vec![], |mut acc, (offsets, is_match)| { - if is_match && !previous_match { - if let Some(((start, _), _)) = acc.last_mut() { - *start = offsets.0; - } else { - acc.push((offsets, false)); - } - } else { - acc.push((offsets, false)); - } - previous_match = is_match; - acc - }); - splits.reverse(); - splits + self.previous_match = is_match; + } + + /// Flushes the last pending span. + pub(crate) fn finish(mut self) { + let last = self.pending.take(); + self.emit(last); + } + + /// Stretches the pending span to cover `offsets` too, or starts one from it. + fn extend_pending_to(&mut self, offsets: (usize, usize)) { + match &mut self.pending { + Some((_, end)) => *end = offsets.1, + None => self.pending = Some(offsets), } - }; + } - for ((start, end), should_remove) in splits { - if !should_remove && start != end { - out.push(Span { + /// Empty pieces are dropped, matching the fold in `NormalizedString::split`. + fn emit(&mut self, span: Option<(usize, usize)>) { + if let Some((start, end)) = span + && start != end + { + self.out.push(Span { start: start as u32, end: end as u32, });