Make stim sampling parallel by default

.gitignore

@@ -1,4 +1,5 @@
 /target
+/debug
 .DS_Store
 .vscode/
 *__pycache__*

crates/ppvm-python-native/Cargo.toml

@@ -13,6 +13,6 @@ bnum = "0.13.0"
 num = "0.4.3"
 paste = "1.0.15"
 ppvm-runtime = { version = "0.1.0", path = "../ppvm-runtime" }
-ppvm-stim = { version = "0.1.0", path = "../ppvm-stim" }
+ppvm-stim = { version = "0.1.0", path = "../ppvm-stim", features = ["rayon"] }
 ppvm-tableau = { version = "0.1.0", path = "../ppvm-tableau" }
 pyo3 = { version = "0.27.0", features = ["multiple-pymethods"] }

crates/ppvm-python-native/src/interface_tableau.rs

@@ -263,38 +263,40 @@ macro_rules! create_interface {
             }
 
             /// Multi-shot sampling: builds a fresh tableau per shot.
+            ///
+            /// Shots run in parallel on rayon's global thread pool (GIL
+            /// released), falling back to serial for small batches. Shot `i`
+            /// is seeded with `seed + i` when `seed` is given, so results are
+            /// reproducible and independent of the thread count; set the
+            /// `RAYON_NUM_THREADS` environment variable to control the pool
+            /// size.
             #[staticmethod]
             #[pyo3(signature = (prog, n_qubits, min_abs_coeff = 1e-10, num_shots = 1, seed = None))]
             pub fn sample(
+                py: Python<'_>,
                 prog: &crate::stim_program::PyStimProgram,
                 n_qubits: usize,
                 min_abs_coeff: f64,
                 num_shots: usize,
                 seed: Option<u64>,
             ) -> pyo3::PyResult<Vec<Vec<u8>>> {
-                let mut next_seed = seed;
-                let count = prog.measurement_count();
-                Ok((0..num_shots)
-                    .map(|_| {
-                        let s = next_seed;
-                        if let Some(ref mut v) = next_seed {
-                            *v = v.wrapping_add(1);
-                        }
-                        let mut tab = match s {
+                let raw = py
+                    .detach(|| {
+                        ppvm_stim::sample(&prog.0, num_shots, |i| match seed {
                             Some(s) => GeneralizedTableau::<$type, $indexType>::new_with_seed(
                                 n_qubits,
                                 min_abs_coeff,
-                                s,
+                                s.wrapping_add(i as u64),
                             ),
                             None => GeneralizedTableau::<$type, $indexType>::new(
                                 n_qubits,
                                 min_abs_coeff,
                             ),
-                        };
-                        let mut shot = Vec::with_capacity(count);
-                        ppvm_stim::execute_prepared(&prog.instructions, &mut tab, &mut shot);
-                        shot
+                        })
                     })
+                    .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("{e}")))?;
+                Ok(raw
+                    .into_iter()
                     .map(|shot| {
                         shot.into_iter()
                             .map(crate::interface_tableau::measurement_to_u8)

crates/ppvm-stim/Cargo.toml

@@ -3,6 +3,11 @@ name = "ppvm-stim"
 version = "0.1.0"
 edition = "2024"
 
+[features]
+default = []
+# Parallelise multi-shot sampling across the global rayon thread pool.
+rayon = ["dep:rayon"]
+
 [dependencies]
 ppvm-runtime = { version = "0.1.0", path = "../ppvm-runtime" }
 ppvm-tableau = { version = "0.1.0", path = "../ppvm-tableau" }
@@ -11,6 +16,7 @@ itertools = "0.14"
 num = "0.4.3"
 thiserror = "2"
 stim-parser = { version = "0.1.0", path = "../stim-parser" }
+rayon = { version = "1", optional = true }
 
 [dev-dependencies]
 # Used by benchmarks.

crates/ppvm-stim/src/executor.rs

@@ -52,11 +52,152 @@ where
     Ok(results)
 }
 
-/// Execute many shots, building a fresh tableau per shot via `make_tableau`.
+/// Execute many shots serially, building a fresh tableau for shot `i` via
+/// `make_tableau(i)`.
+///
+/// The shot index lets callers derive a deterministic per-shot seed (e.g.
+/// `seed + i`) so results are independent of evaluation order — the same
+/// factory then yields identical results from `sample_parallel` (when the `rayon` feature is enabled).
+pub fn sample_serial<T, I, C, F>(
+    program: &ExtendedProgram,
+    num_shots: usize,
+    make_tableau: F,
+) -> Result<Vec<Vec<Option<bool>>>, ExecError>
+where
+    T: Config,
+    <<T as Config>::Storage as BitView>::Store: PrimInt,
+    C: SparseVector<Complex<T::Coeff>, I> + std::fmt::Debug,
+    T::Coeff: One
+        + Zero
+        + Clone
+        + num::Num
+        + ToPrimitive
+        + std::fmt::Debug
+        + std::ops::Mul<f64>
+        + PartialOrd<f64>
+        + Send
+        + Sync,
+    Complex<T::Coeff>: std::ops::Mul<Output = Complex<T::Coeff>>
+        + From<Complex64>
+        + std::ops::MulAssign
+        + std::ops::AddAssign
+        + One
+        + ComplexFloat
+        + Copy,
+    I: TableauIndex + Debug + Send + Sync,
+    F: Fn(usize) -> GeneralizedTableau<T, I, C>,
+{
+    prepare(program)?;
+    let count = program.measurement_count();
+    Ok((0..num_shots)
+        .map(|i| {
+            let mut tab = make_tableau(i);
+            let mut results = Vec::with_capacity(count);
+            execute_prepared(&program.instructions, &mut tab, &mut results);
+            results
+        })
+        .collect())
+}
+
+/// Execute many shots, building a fresh tableau for shot `i` via
+/// `make_tableau(i)`.
+///
+/// When the `rayon` feature is enabled this dispatches to `sample_parallel`
+/// for batches large enough to amortise thread-scheduling overhead, and to
+/// [`sample_serial`] otherwise. Without the feature it is always serial. Use
+/// [`sample_serial`] / `sample_parallel` directly to force one or the other.
+#[cfg(feature = "rayon")]
+pub fn sample<T, I, C, F>(
+    program: &ExtendedProgram,
+    num_shots: usize,
+    make_tableau: F,
+) -> Result<Vec<Vec<Option<bool>>>, ExecError>
+where
+    T: Config,
+    <<T as Config>::Storage as BitView>::Store: PrimInt,
+    C: SparseVector<Complex<T::Coeff>, I> + std::fmt::Debug,
+    T::Coeff: One
+        + Zero
+        + Clone
+        + num::Num
+        + ToPrimitive
+        + std::fmt::Debug
+        + std::ops::Mul<f64>
+        + PartialOrd<f64>
+        + Send
+        + Sync,
+    Complex<T::Coeff>: std::ops::Mul<Output = Complex<T::Coeff>>
+        + From<Complex64>
+        + std::ops::MulAssign
+        + std::ops::AddAssign
+        + One
+        + ComplexFloat
+        + Copy,
+    I: TableauIndex + Debug + Send + Sync,
+    F: Fn(usize) -> GeneralizedTableau<T, I, C> + Sync,
+{
+    // Below ~4 shots per thread, rayon's scheduling overhead outweighs
+    // the gain, so stay serial; with a single thread there is no upside.
+    let n_threads = rayon::current_num_threads();
+    if n_threads <= 1 || num_shots < 4 * n_threads {
+        sample_serial(program, num_shots, make_tableau)
+    } else {
+        sample_parallel(program, num_shots, make_tableau)
+    }
+}
+
+/// Execute many shots, building a fresh tableau for shot `i` via
+/// `make_tableau(i)`.
+///
+/// Without the `rayon` feature execution is always serial. Enable the
+/// `rayon` feature to get a version that can dispatch to parallel execution
+/// for large batches. Use [`sample_serial`] to force serial execution
+/// regardless of the feature.
+#[cfg(not(feature = "rayon"))]
 pub fn sample<T, I, C, F>(
     program: &ExtendedProgram,
     num_shots: usize,
-    mut make_tableau: F,
+    make_tableau: F,
+) -> Result<Vec<Vec<Option<bool>>>, ExecError>
+where
+    T: Config,
+    <<T as Config>::Storage as BitView>::Store: PrimInt,
+    C: SparseVector<Complex<T::Coeff>, I> + std::fmt::Debug,
+    T::Coeff: One
+        + Zero
+        + Clone
+        + num::Num
+        + ToPrimitive
+        + std::fmt::Debug
+        + std::ops::Mul<f64>
+        + PartialOrd<f64>
+        + Send
+        + Sync,
+    Complex<T::Coeff>: std::ops::Mul<Output = Complex<T::Coeff>>
+        + From<Complex64>
+        + std::ops::MulAssign
+        + std::ops::AddAssign
+        + One
+        + ComplexFloat
+        + Copy,
+    I: TableauIndex + Debug + Send + Sync,
+    F: Fn(usize) -> GeneralizedTableau<T, I, C>,
+{
+    sample_serial(program, num_shots, make_tableau)
+}
+
+/// Execute many shots in parallel across the global rayon thread pool,
+/// building a fresh tableau for shot `i` via `make_tableau(i)`.
+///
+/// Per-shot seeds derived from `i` make the result independent of how rayon
+/// schedules the work, so a seeded factory yields the same shots (in the same
+/// order) as [`sample_serial`]. Thread count follows rayon's global pool
+/// (set `RAYON_NUM_THREADS` to override).
+#[cfg(feature = "rayon")]
+pub fn sample_parallel<T, I, C, F>(
+    program: &ExtendedProgram,
+    num_shots: usize,
+    make_tableau: F,
 ) -> Result<Vec<Vec<Option<bool>>>, ExecError>
 where
     T: Config,
@@ -80,13 +221,15 @@ where
         + ComplexFloat
         + Copy,
     I: TableauIndex + Debug + Send + Sync,
-    F: FnMut() -> GeneralizedTableau<T, I, C>,
+    F: Fn(usize) -> GeneralizedTableau<T, I, C> + Sync,
 {
+    use rayon::prelude::*;
     prepare(program)?;
     let count = program.measurement_count();
     Ok((0..num_shots)
-        .map(|_| {
-            let mut tab = make_tableau();
+        .into_par_iter()
+        .map(|i| {
+            let mut tab = make_tableau(i);
             let mut results = Vec::with_capacity(count);
             execute_prepared(&program.instructions, &mut tab, &mut results);
             results

crates/ppvm-stim/src/lib.rs

@@ -21,7 +21,7 @@
 //! use ppvm_tableau::prelude::*;
 //!
 //! let prog = parse_extended(circuit_src)?;
-//! let shots = sample(&prog, 10_000, || {
+//! let shots = sample(&prog, 10_000, |_| {
 //!     GeneralizedTableau::<_, usize, _>::new(n_qubits, 1e-10)
 //! })?;
 //! # Ok::<(), ppvm_stim::Error>(())
@@ -38,7 +38,9 @@ pub mod prepare;
 
 pub use stim_parser::prelude::*;
 
-pub use executor::{execute, execute_prepared, sample};
+#[cfg(feature = "rayon")]
+pub use executor::sample_parallel;
+pub use executor::{execute, execute_prepared, sample, sample_serial};
 pub use prepare::{ExecError, prepare};
 
 use std::path::{Path, PathBuf};

crates/ppvm-stim/tests/executor.rs

@@ -110,7 +110,7 @@ fn measurement_buffer_is_pre_sized() {
 fn sample_runs_n_shots_each_with_fresh_tableau() {
     use ppvm_stim::sample;
     let prog = parse_extended("X 0\nM 0").unwrap();
-    let shots = sample::<_, _, _, _>(&prog, 5, || {
+    let shots = sample::<_, _, _, _>(&prog, 5, |_| {
         GeneralizedTableau::<ByteFxHashF64<1>, usize>::new(1, 1e-10)
     })
     .unwrap();
@@ -124,7 +124,7 @@ fn sample_runs_n_shots_each_with_fresh_tableau() {
 fn sample_zero_shots_returns_empty() {
     use ppvm_stim::sample;
     let prog = parse_extended("X 0\nM 0").unwrap();
-    let shots = sample::<_, _, _, _>(&prog, 0, || {
+    let shots = sample::<_, _, _, _>(&prog, 0, |_| {
         GeneralizedTableau::<ByteFxHashF64<1>, usize>::new(1, 1e-10)
     })
     .unwrap();
@@ -137,10 +137,8 @@ fn sample_random_h_distribution_within_3_sigma() {
     use ppvm_stim::sample;
     let prog = parse_extended("H 0\nM 0").unwrap();
     let n = 4096;
-    let mut seed_counter: u64 = 0;
-    let shots = sample::<_, _, _, _>(&prog, n, || {
-        seed_counter += 1;
-        GeneralizedTableau::<ByteFxHashF64<1>, usize>::new_with_seed(1, 1e-10, seed_counter)
+    let shots = sample::<_, _, _, _>(&prog, n, |i| {
+        GeneralizedTableau::<ByteFxHashF64<1>, usize>::new_with_seed(1, 1e-10, i as u64 + 1)
     })
     .unwrap();
     let ones = shots.iter().filter(|s| s[0] == Some(true)).count();
@@ -383,10 +381,8 @@ fn measure_noise_distribution_within_3_sigma() {
     // So recorded == 0 with probability 0.3 over many shots.
     let prog = parse_extended("X 0\nMZ(0.3) 0").unwrap();
     let n = 4096usize;
-    let mut seed_counter: u64 = 0;
-    let shots = sample::<_, _, _, _>(&prog, n, || {
-        seed_counter += 1;
-        GeneralizedTableau::<ByteFxHashF64<1>, usize>::new_with_seed(1, 1e-10, seed_counter)
+    let shots = sample::<_, _, _, _>(&prog, n, |i| {
+        GeneralizedTableau::<ByteFxHashF64<1>, usize>::new_with_seed(1, 1e-10, i as u64 + 1)
     })
     .unwrap();
     let zeros = shots.iter().filter(|s| s[0] == Some(false)).count();
@@ -437,10 +433,8 @@ fn mpad_noise_distribution_within_3_sigma() {
     // MPAD(0.3) 0 — pad value is 0; recorded bit flips to 1 with prob 0.3.
     let prog = parse_extended("MPAD(0.3) 0").unwrap();
     let n = 4096usize;
-    let mut seed_counter: u64 = 0;
-    let shots = sample::<_, _, _, _>(&prog, n, || {
-        seed_counter += 1;
-        GeneralizedTableau::<ByteFxHashF64<1>, usize>::new_with_seed(1, 1e-10, seed_counter)
+    let shots = sample::<_, _, _, _>(&prog, n, |i| {
+        GeneralizedTableau::<ByteFxHashF64<1>, usize>::new_with_seed(1, 1e-10, i as u64 + 1)
     })
     .unwrap();
     let ones = shots.iter().filter(|s| s[0] == Some(true)).count();
@@ -451,3 +445,21 @@ fn mpad_noise_distribution_within_3_sigma() {
         "got {ones} ones, expected mean {mean} +/- 3*{std}"
     );
 }
+
+#[cfg(feature = "rayon")]
+#[test]
+fn sample_parallel_matches_sample_serial_for_seeded_factory() {
+    use ppvm_stim::{sample_parallel, sample_serial};
+    // A randomising circuit (H then measure on two qubits, with a noisy
+    // readout) so each shot's outcome depends on its RNG draws. With a
+    // per-shot seed derived from the index, the parallel result must match
+    // the serial result shot-for-shot regardless of thread scheduling.
+    let prog = parse_extended("H 0\nH 1\nM 0\nMZ(0.2) 1").unwrap();
+    let n = 1024usize;
+    let factory = |i: usize| {
+        GeneralizedTableau::<ByteFxHashF64<1>, usize>::new_with_seed(2, 1e-10, i as u64 + 1)
+    };
+    let serial = sample_serial::<_, _, _, _>(&prog, n, factory).unwrap();
+    let parallel = sample_parallel::<_, _, _, _>(&prog, n, factory).unwrap();
+    assert_eq!(serial, parallel);
+}

ppvm-python/src/ppvm/generalized_tableau.py

@@ -223,6 +223,14 @@ def sample(
 
         Each shot starts from a fresh tableau, so this is the right entry
         point for multi-shot sampling.
+
+        Shots run in parallel across CPU cores (the GIL is released during
+        sampling), with a serial fallback for small batches. When ``seed`` is
+        given, shot ``i`` uses ``(seed + i) % 2**64`` (wrapping ``u64``
+        arithmetic), so results are reproducible and independent of the number
+        of threads. Set the ``RAYON_NUM_THREADS`` environment variable before
+        the first call to control the pool size (it defaults to the number of
+        logical cores).
         """
         N_interface = (n_qubits + 63) // 64
         native_cls = getattr(ppvm_python_native, f"GeneralizedTableau{N_interface}")
@@ -237,7 +245,11 @@ def sample_stim(
     num_shots: int = 1,
     seed: int | None = None,
 ) -> list[list[MeasurementResult]]:
-    """Multi-shot sampling — module-level alias for ``GeneralizedTableau.sample``."""
+    """Multi-shot sampling — module-level alias for ``GeneralizedTableau.sample``.
+
+    Shots are sampled in parallel across CPU cores with the GIL released; see
+    :meth:`GeneralizedTableau.sample` for seeding and ``RAYON_NUM_THREADS``.
+    """
     return GeneralizedTableau.sample(
         prog, n_qubits, min_abs_coeff=min_abs_coeff, num_shots=num_shots, seed=seed
     )