perf(scratch): reuse Scratch buffers across encode() calls

tokenizers/tk-encode/src/models/bpe/model.rs

@@ -1438,6 +1438,24 @@ mod tests {
             assert!(pipeline_ids(&pipeline, "").is_empty());
         }
 
+        // The pool hands the SAME scratch to successive encodes. State left behind by one
+        // call — an undrained merge queue, a stale word buffer — would corrupt every call
+        // after it. Drive several inputs through one scratch and check each still matches
+        // the reference model.
+        #[test]
+        fn reused_scratch_matches_fresh() {
+            let bpe = hello_builder().build().unwrap();
+            let reference = bpe.clone();
+            let model = PipelineBPE::from_bpe(bpe, false).unwrap();
+            let mut scratch = model.init_scratch();
+            for input in ["hello", "hell", "helo", "oleh", "hello", "", "hxe"] {
+                let mut out = Vec::new();
+                pipeline::Model::tokenize_pipeline(&model, input, &mut scratch, &mut out).unwrap();
+                let got: Vec<u32> = out.iter().map(|t| t.id).collect();
+                assert_eq!(got, reference_ids(&reference, input), "{input:?}");
+            }
+        }
+
         #[test]
         fn unknown_char_without_unk_is_dropped() {
             let bpe = hello_builder().build().unwrap();

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

@@ -1,5 +1,7 @@
 use std::cell::RefCell;
 use std::convert::TryInto;
+use std::mem;
+use std::sync::{Mutex, PoisonError};
 use std::{borrow::Cow, convert::TryFrom};
 
 use atomsplit::classify::classify;
@@ -391,6 +393,78 @@ pub struct PipelineTokenizer {
     pre_tokenizer: PipelinePreTokenizer,
     model: PipelineModel,
     post_processor: PipelinePostProcessor,
+    scratch_pool: ScratchPool,
+}
+
+/// A pool of [`PipelineModelScratch`].
+///
+/// When calling [`PipelineTokenizer::encode`], an instance of [`PipelineModelScratch`] is taken out of this pool
+/// and given to the tokenizer. When the encoding is done, the scratch buffer is returned to the pool and can be
+/// reused by later calls.
+///
+/// The reusability matters because the scratch buffer may hold cache structures which are more useful when reused,
+/// and less importantly it saves an extra allocation for an fresh buffer every time.
+struct ScratchPool(Mutex<Vec<PipelineModelScratch>>);
+
+impl ScratchPool {
+    fn new() -> Self {
+        Self(Mutex::new(Vec::new()))
+    }
+
+    /// Get a scratch buffer from the pool, wrapped in a [`ScratchGuard`].
+    /// When the [`ScratchGuard`] gets dropped, the scratch buffer is pushed back to the pool.
+    fn get<'a>(&'a self, model: &PipelineModel) -> ScratchGuard<'a> {
+        // The Mutex lock is held just long enough to pop the scratch out of the pool
+        let taken = self.0.lock().unwrap_or_else(PoisonError::into_inner).pop();
+        ScratchGuard {
+            // If there was no scratch buffer available in the pool, we build.a fresh one
+            scratch: taken.unwrap_or_else(|| model.init_scratch()),
+            pool: self,
+        }
+    }
+
+    #[cfg(test)]
+    fn len(&self) -> usize {
+        self.0.lock().unwrap_or_else(PoisonError::into_inner).len()
+    }
+}
+
+/// A wrapper around [`PipelineModelScratch`].
+/// Implements [`Deref`] and [`DerefMut`], so it behaves as [`PipelineModelScratch`].
+///
+/// When it gets dropped, it pushes [`Self::scratch`] back into the shared [`Self::pool`] so it can
+/// get reused by a later call to [`PipelineTokenizer::encode`].
+///
+/// TODO @McPatate : The Mutex can create contention, to be replaced by a better access pattern
+struct ScratchGuard<'a> {
+    scratch: PipelineModelScratch,
+    pool: &'a ScratchPool,
+}
+
+impl Drop for ScratchGuard<'_> {
+    fn drop(&mut self) {
+        // Steals the scratch buffer from self, replaces it with PipelineModelScratch::default()
+        let scratch = mem::take(&mut self.scratch);
+        // Push the scratch back in the pool
+        self.pool
+            .0
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .push(scratch);
+    }
+}
+
+impl std::ops::Deref for ScratchGuard<'_> {
+    type Target = PipelineModelScratch;
+    fn deref(&self) -> &PipelineModelScratch {
+        &self.scratch
+    }
+}
+
+impl std::ops::DerefMut for ScratchGuard<'_> {
+    fn deref_mut(&mut self) -> &mut PipelineModelScratch {
+        &mut self.scratch
+    }
 }
 
 impl TryFrom<&Tokenizer> for PipelineTokenizer {
@@ -493,6 +567,7 @@ impl TryFrom<&Tokenizer> for PipelineTokenizer {
                 .map(PipelinePostProcessor::try_from)
                 .transpose()?
                 .unwrap_or_default(),
+            scratch_pool: ScratchPool::new(),
         })
     }
 }
@@ -523,7 +598,7 @@ impl PipelineTokenizer {
     pub fn encode(&self, input: &str, add_special_tokens: bool) -> Result<Vec<PipelineToken>> {
         let mut output = Vec::new();
         let mut pre_tokens = Vec::new();
-        let mut scratch = self.model.init_scratch();
+        let mut scratch = self.scratch_pool.get(&self.model);
 
         self.encode_generic::<{ Self::STAGE_POSTPROCESS }>(
             input,
@@ -972,11 +1047,19 @@ impl Model for PipelineModel {
     }
 }
 
+/// A set of buffers and other state the model needs to encode efficiently,
+/// reused among calls to [`PipelineTokenizer::encode`].
+///
+/// Each model gets its own variant.
+#[derive(Default)]
 pub enum PipelineModelScratch {
     BPE(BpeScratch),
     WordLevel(()),
     WordPiece(WordPieceScratch),
     Unigram(UnigramScratch),
+    /// We need a default value to able to use [`mem::take`] in [`ScratchGuard::drop`]
+    #[default]
+    None,
 }
 
 impl ModelScratch for PipelineModelScratch {}
@@ -1314,4 +1397,101 @@ mod tests {
         let err = conversion_error(&tok);
         assert!(err.contains("not supported"), "{}", err);
     }
+
+    /// A BPE pipeline that merges "hello" into the single id 7.
+    fn hello_pipeline() -> PipelineTokenizer {
+        use crate::models::bpe::{BpeBuilder, Merges, Vocab};
+
+        let vocab: Vocab = [
+            ("h", 0u32),
+            ("e", 1),
+            ("l", 2),
+            ("o", 3),
+            ("he", 4),
+            ("hel", 5),
+            ("hell", 6),
+            ("hello", 7),
+        ]
+        .into_iter()
+        .map(|(s, i)| (s.to_string(), i))
+        .collect();
+        let merges: Merges = vec![
+            ("h".to_string(), "e".to_string()),
+            ("he".to_string(), "l".to_string()),
+            ("hel".to_string(), "l".to_string()),
+            ("hell".to_string(), "o".to_string()),
+        ];
+        let bpe = BpeBuilder::default()
+            .vocab_and_merges(vocab, merges)
+            .build()
+            .unwrap();
+        PipelineTokenizer::try_from(&Tokenizer::new(bpe)).unwrap()
+    }
+
+    // The pool exists so ONE `&self` tokenizer can be shared across rayon workers. Encode
+    // the same input from thousands of threads through a single instance; each must get a
+    // private scratch and produce the sequential result. Two threads sharing a scratch
+    // would corrupt some of them — and this only compiles if `PipelineTokenizer: Sync`,
+    // which the pool has to preserve.
+    #[test]
+    fn encode_shared_across_threads() {
+        use rayon::prelude::*;
+
+        let pipeline = hello_pipeline();
+
+        let want: Vec<u32> = pipeline
+            .encode("hello", false)
+            .unwrap()
+            .iter()
+            .map(|t| t.id)
+            .collect();
+        assert_eq!(want, vec![7]);
+
+        let all_match = (0..10_000u32).into_par_iter().all(|_| {
+            pipeline
+                .encode("hello", false)
+                .unwrap()
+                .iter()
+                .map(|t| t.id)
+                .collect::<Vec<_>>()
+                == want
+        });
+        assert!(all_match);
+    }
+
+    // Reusing scratches is the whole point of the pool, so it must not build one per call:
+    // one thread encoding in a loop has to keep coming back to the same scratch, and a
+    // burst of N threads must leave at most N behind for later calls to use.
+    #[test]
+    fn scratches_are_reused_rather_than_piling_up() {
+        use std::sync::Barrier;
+
+        let pipeline = hello_pipeline();
+        for _ in 0..1000 {
+            pipeline.encode("hello", false).unwrap();
+        }
+        assert_eq!(pipeline.scratch_pool.len(), 1);
+
+        let threads = 64;
+        let all_holding = Barrier::new(threads);
+        std::thread::scope(|scope| {
+            for _ in 0..threads {
+                scope.spawn(|| {
+                    let scratch = pipeline.scratch_pool.get(&pipeline.model);
+                    all_holding.wait();
+                    drop(scratch);
+                });
+            }
+        });
+
+        let after_burst = pipeline.scratch_pool.len();
+        assert!(
+            after_burst <= threads,
+            "{after_burst} scratches kept for {threads} threads"
+        );
+        for _ in 0..1000 {
+            pipeline.encode("hello", false).unwrap();
+        }
+        assert_eq!(pipeline.scratch_pool.len(), after_burst);
+    }
 }