fast atomsplit::literal

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

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<usize> = 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
+        );
+    }
+}

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;

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<Item = usize> + '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);
+        }
+    }
 }