perf: simd literal matching

tokenizers/atomsplit/Cargo.toml

@@ -26,6 +26,7 @@ path = "src/lib.rs"
 
 [dependencies]
 memchr = "2.8.2" # SIMD search: used for single-byte  search in CharDelimiterSplit (1.4–23× vs scalar), or multi-byte string pattern search in `literal`.
+wide = "1.6.0"
 # NOTE: classify tables live in src/atom_tables.rs (committed, generated). Regenerate after any atom
 # scheme change with `cargo run -p bitmap_gen`. No build script / build-dep — atomsplit builds clean.
 # onig is C (Oniguruma) and fancy-regex pulls it in transitively for benches only — neither

tokenizers/atomsplit/src/literal.rs

@@ -6,51 +6,146 @@
 //! 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.
 
-use memchr::memmem;
-use std::fmt;
+use std::{fmt, u32};
+use wide::{self, u8x32};
 
-/// The pattern handed to [`Literal::new`] was empty.
+/// The pattern handed to [`Literal::new`] cannot be searched for.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub struct EmptyPattern;
+pub enum InvalidPattern {
+    /// An empty pattern would match everywhere.
+    Empty,
+    /// Longer than [`Literal::MAX_PATTERN_LEN`] bytes.
+    TooLong,
+}
 
-impl fmt::Display for EmptyPattern {
+impl fmt::Display for InvalidPattern {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.write_str("an empty pattern matches everywhere")
+        match self {
+            Self::Empty => f.write_str("an empty pattern matches everywhere"),
+            Self::TooLong => write!(
+                f,
+                "patterns longer than {} bytes are not supported",
+                Literal::MAX_PATTERN_LEN
+            ),
+        }
     }
 }
 
-impl std::error::Error for EmptyPattern {}
+impl std::error::Error for InvalidPattern {}
 
-/// A literal string to split on.
-///
-/// The finder is boxed because it is large — a few hundred bytes of prefilter state on x86_64 — and
-/// callers store it inside enums whose other variants are tiny.
+/// A literal string to split on, at most [`Literal::MAX_PATTERN_LEN`] bytes.
 #[derive(Debug, Clone)]
 pub struct Literal {
-    finder: Box<memmem::Finder<'static>>,
+    pattern: [u8; Self::MAX_PATTERN_LEN],
+    pattern_len: u8,
 }
 
 impl Literal {
+    /// Longest supported pattern: the matcher compares one 32-byte block at a time, so a longer
+    /// pattern could never match inside a block.
+    pub const MAX_PATTERN_LEN: usize = 32;
+
     /// # Errors
-    /// If `pattern` is empty, which would match everywhere.
-    pub fn new(pattern: &[u8]) -> Result<Self, EmptyPattern> {
+    /// If `pattern` is empty (it would match everywhere) or longer than
+    /// [`Literal::MAX_PATTERN_LEN`] bytes.
+    pub fn new(pattern: &[u8]) -> Result<Self, InvalidPattern> {
         if pattern.is_empty() {
-            return Err(EmptyPattern);
+            return Err(InvalidPattern::Empty);
         }
+        if pattern.len() > Self::MAX_PATTERN_LEN {
+            return Err(InvalidPattern::TooLong);
+        }
+        let mut buf = [0u8; Self::MAX_PATTERN_LEN];
+        buf[..pattern.len()].copy_from_slice(pattern);
         Ok(Self {
-            finder: Box::new(memmem::Finder::new(pattern).into_owned()),
+            pattern: buf,
+            pattern_len: pattern.len() as u8,
         })
     }
 
     /// The string being searched for.
     #[must_use]
     pub fn pattern(&self) -> &[u8] {
-        self.finder.needle()
+        &self.pattern[..usize::from(self.pattern_len)]
     }
 
     /// Byte offset of every match, left to right. Matches never overlap, so `"aa"` is found once in
     /// `"aaa"` — the same matches a regex engine would report.
     pub fn matches<'t>(&'t self, text: &'t [u8]) -> impl Iterator<Item = usize> + 't {
-        self.finder.find_iter(text)
+        LiteralMatcher::new(text, self.pattern())
+    }
+}
+
+struct LiteralMatcher<'a> {
+    text: &'a [u8],
+    needle: &'a [u8],
+    needle_splats: Vec<u8x32>,
+    equality_mask: u32,
+    block_start: usize,
+    next_block: usize,
+    min_next_start: usize,
+}
+
+impl<'a> LiteralMatcher<'a> {
+    pub fn new(text: &'a [u8], needle: &'a [u8]) -> Self {
+        Self {
+            text,
+            needle,
+            needle_splats: needle.iter().map(|&b| u8x32::splat(b)).collect(),
+            equality_mask: 0,
+            block_start: 0,
+            next_block: 0,
+            min_next_start: 0,
+        }
+    }
+
+    pub fn load_more(&mut self) -> Option<usize> {
+        if self.next_block >= self.text.len() {
+            // We have read the whole text
+            return None;
+        }
+        self.block_start = self.next_block;
+        let remaining = &self.text[self.block_start..];
+        let len_to_load = remaining.len().min(32);
+        let mut padded_text = [0u8; 32];
+        padded_text[..len_to_load].copy_from_slice(&remaining[..len_to_load]);
+        let simd_bytes = u8x32::new(padded_text);
+        let mut mask = u32::MAX;
+        for (idx, &needle_splat) in self.needle_splats.iter().enumerate() {
+            mask &= simd_bytes.simd_eq(needle_splat).to_bitmask() >> idx;
+        }
+
+        mask &= u32::MAX >> (self.needle.len() - 1);
+        if len_to_load < 32 {
+            let valid = len_to_load.saturating_sub(self.needle.len() - 1);
+            mask &= (1u32 << valid) - 1
+        }
+        self.equality_mask = mask;
+        self.next_block += 32 - (self.needle.len() - 1);
+        Some(len_to_load)
+    }
+
+    pub fn next_match(&mut self) -> Option<usize> {
+        loop {
+            // Load more until a match or text has been entirely read
+            while self.equality_mask == 0 {
+                self.load_more()?;
+            }
+            let offset = self.equality_mask.trailing_zeros() as usize;
+            self.equality_mask &= self.equality_mask - 1;
+            let start = self.block_start + offset;
+            if start >= self.min_next_start {
+                self.min_next_start = start + self.needle.len();
+                return Some(start);
+            }
+        }
+    }
+}
+
+impl<'a> Iterator for LiteralMatcher<'a> {
+    type Item = usize;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.next_match()
     }
 }

tokenizers/atomsplit/tests/literal.rs

@@ -1,5 +1,5 @@
 //! Tests for the literal search. Kept out of `src/` so the core stays production-only.
-use atomsplit::literal::{EmptyPattern, Literal};
+use atomsplit::literal::{InvalidPattern, Literal};
 
 #[test]
 fn finds_every_match_and_nothing_else() {
@@ -29,22 +29,55 @@ fn matches_do_not_overlap() {
     assert_eq!(literal.matches(b"aaaa").collect::<Vec<_>>(), [0, 2]);
 }
 
-/// A `Literal` is stored inline in the normalizer and decoder enums, where every other variant is a
-/// handful of bytes. `memmem::Finder` itself is a few hundred, so it has to stay behind a pointer.
 #[test]
-fn a_literal_is_pointer_sized() {
-    assert_eq!(
-        size_of::<Literal>(),
-        size_of::<*const u8>(),
-        "a Literal must not carry its finder inline"
-    );
+fn finds_matches_beyond_the_first_simd_load() {
+    // The matcher reads the text 32 bytes at a time; a match in the second load checks that
+    // reported offsets count from the start of the text, not the start of the load.
+    let mut text = vec![b'x'; 40];
+    text.extend_from_slice(b"-yy");
+    let literal = Literal::new(b"-").unwrap();
+    assert_eq!(literal.matches(&text).collect::<Vec<_>>(), [40]);
+}
+
+#[test]
+fn finds_a_match_straddling_two_simd_loads() {
+    // "ab" placed so that 'a' is the last byte of the first 32-byte load and 'b' the first byte
+    // of the second: the match is invisible to either load alone.
+    let mut text = vec![b'x'; 31];
+    text.extend_from_slice(b"ab");
+    text.extend_from_slice(&[b'x'; 10]);
+    let literal = Literal::new(b"ab").unwrap();
+    assert_eq!(literal.matches(&text).collect::<Vec<_>>(), [31]);
+}
+
+#[test]
+fn finds_matches_at_every_load_alignment() {
+    // 120 bytes of "xab" put a match at every offset of the form 3k + 1, so some match starts at
+    // every position relative to a 32-byte load, including each straddle of a load boundary.
+    let text = b"xab".repeat(40);
+    let literal = Literal::new(b"ab").unwrap();
+    let expected: Vec<usize> = (0..40).map(|k| 3 * k + 1).collect();
+    assert_eq!(literal.matches(&text).collect::<Vec<_>>(), expected);
 }
 
 #[test]
 fn an_empty_pattern_is_rejected() {
-    assert_eq!(Literal::new(b"").unwrap_err(), EmptyPattern);
+    assert_eq!(Literal::new(b"").unwrap_err(), InvalidPattern::Empty);
     assert_eq!(
-        EmptyPattern.to_string(),
+        InvalidPattern::Empty.to_string(),
         "an empty pattern matches everywhere"
     );
 }
+
+#[test]
+fn patterns_longer_than_max_len_are_rejected() {
+    assert!(Literal::new(&[b'a'; Literal::MAX_PATTERN_LEN]).is_ok());
+    assert_eq!(
+        Literal::new(&[b'a'; Literal::MAX_PATTERN_LEN + 1]).unwrap_err(),
+        InvalidPattern::TooLong
+    );
+    assert_eq!(
+        InvalidPattern::TooLong.to_string(),
+        "patterns longer than 32 bytes are not supported"
+    );
+}

tokenizers/tk-encode/src/tokenizer/pattern.rs

@@ -1,6 +1,6 @@
 use crate::utils::SysRegex;
 use crate::{Offsets, Result};
-use atomsplit::literal::Literal;
+use atomsplit::literal::{InvalidPattern, Literal};
 use regex::Regex;
 
 /// Pattern used to split a NormalizedString
@@ -73,7 +73,8 @@ impl Pattern for &str {
         match Literal::new(self.as_bytes()) {
             Ok(literal) => (&literal).find_matches(inside),
             // An empty pattern would match everywhere, so it matches nowhere instead.
-            Err(_) => Ok(vec![((0, inside.len()), false)]),
+            Err(InvalidPattern::Empty) => Ok(vec![((0, inside.len()), false)]),
+            Err(e @ InvalidPattern::TooLong) => Err(e.into()),
         }
     }
 }