prism_q/sim/
dynamic.rs

1//! Per-shot execution of a [`DynamicProgram`].
2
3use rand::SeedableRng;
4
5use super::branching::{self, BranchLimits};
6use super::dispatch::{BackendPlan, ExecutionPlan};
7use super::{
8    BackendKind, CountsResult, RunMetadata, Seeded, ShotsResult, Unseeded, backend_metadata,
9    expand_for_backend, fuse_for_backend, mix_seed, resolve, run_counts_with, run_shots_with,
10    validate_explicit_backend,
11};
12use crate::backend::statevector::StatevectorBackend;
13use crate::backend::{Backend, max_statevector_qubits};
14use crate::circuit::Instruction;
15use crate::circuit::dynamic::{Action, BasicBlock, ClassicalValue, DynamicProgram, Terminator};
16use crate::error::{PrismError, Result};
17
18/// Blocks a shot may run before [`SimulateProgram::shots`] stops it, unless
19/// [`SimulateProgram::max_steps`] sets another bound.
20const DEFAULT_MAX_STEPS: u64 = 1_000_000;
21
22/// Builder for running a [`DynamicProgram`], the counterpart of
23/// [`Simulate`](super::Simulate) for programs with runtime control flow.
24pub struct SimulateProgram<'p, SeedState> {
25    program: &'p DynamicProgram,
26    kind: BackendKind,
27    seed: SeedState,
28    max_steps: u64,
29}
30
31/// Start a run of `program` on [`BackendKind::Auto`].
32///
33/// A program with no runtime control flow runs as its one circuit does, through
34/// [`simulate`](super::simulate), so it keeps every sampling shortcut and the
35/// same seeded shots. Any other program has each block fused for the backend
36/// once. On the host statevector the shots then walk the graph together,
37/// splitting binomially at each measurement as circuits with mid-circuit
38/// measurements do, so the program evolves once per distinct outcome history.
39/// On any other backend every shot walks the graph on a fresh state, seeded as
40/// the per-shot circuit routes seed shot `i`.
41///
42/// Auto routing reads every block at once, so a program holding only Clifford
43/// gates runs on the stabilizer tableau, and a runtime rotation counts as a
44/// non-Clifford gate. Stabilizer rank, Pauli propagation, and the distributed
45/// statevector have no per-shot state to walk a graph on and decline.
46pub fn simulate_program(program: &DynamicProgram) -> SimulateProgram<'_, Unseeded> {
47    SimulateProgram {
48        program,
49        kind: BackendKind::Auto,
50        seed: Unseeded,
51        max_steps: DEFAULT_MAX_STEPS,
52    }
53}
54
55impl<'p, SeedState> SimulateProgram<'p, SeedState> {
56    pub fn backend(mut self, kind: BackendKind) -> Self {
57        self.kind = kind;
58        self
59    }
60
61    /// Blocks one shot may run before it stops with
62    /// [`PrismError::StepLimit`], one million by default. Every block entered
63    /// counts, the empty test at the head of a loop included.
64    pub fn max_steps(mut self, max_steps: u64) -> Self {
65        self.max_steps = max_steps;
66        self
67    }
68}
69
70impl<'p> SimulateProgram<'p, Unseeded> {
71    pub fn seed(self, seed: u64) -> SimulateProgram<'p, Seeded> {
72        SimulateProgram {
73            program: self.program,
74            kind: self.kind,
75            seed: Seeded { seed },
76            max_steps: self.max_steps,
77        }
78    }
79}
80
81impl SimulateProgram<'_, Seeded> {
82    /// Run `num_shots` shots and collect each one's classical bits.
83    ///
84    /// # Errors
85    ///
86    /// [`PrismError::StepLimit`] when a shot runs past the step bound,
87    /// [`PrismError::IncompatibleBackend`] for a backend that cannot walk the
88    /// graph or hold the program's gates, and anything a block's execution
89    /// raises.
90    pub fn shots(self, num_shots: usize) -> Result<ShotsResult> {
91        let seed = self.seed.seed;
92        if let Some(circuit) = self.program.static_circuit() {
93            return run_shots_with(self.kind, circuit, num_shots, seed);
94        }
95        let prepared = Prepared::new(self.program, &self.kind, seed, self.max_steps)?;
96        prepared.run(&self.kind, num_shots, seed)
97    }
98
99    /// Run `num_shots` shots and count each distinct outcome.
100    ///
101    /// # Errors
102    ///
103    /// Same conditions as [`SimulateProgram::shots`].
104    pub fn sample_counts(self, num_shots: usize) -> Result<CountsResult> {
105        let num_classical_bits = self.program.num_classical_bits();
106        if let Some(circuit) = self.program.static_circuit() {
107            let (counts, metadata) =
108                run_counts_with(self.kind, circuit, num_shots, self.seed.seed)?;
109            return Ok(CountsResult {
110                counts,
111                num_classical_bits,
112                metadata,
113            });
114        }
115        let shots = self.shots(num_shots)?;
116        Ok(CountsResult {
117            counts: shots.counts(),
118            num_classical_bits,
119            metadata: shots.metadata,
120        })
121    }
122}
123
124struct PreparedBlock<'p> {
125    instructions: Vec<Instruction>,
126    block: &'p BasicBlock,
127}
128
129/// A program fused for one backend plan, ready to run any number of shots.
130pub(super) struct Prepared<'p> {
131    program: &'p DynamicProgram,
132    plan: BackendPlan,
133    blocks: Vec<PreparedBlock<'p>>,
134    initial: Vec<ClassicalValue>,
135    max_steps: u64,
136}
137
138impl<'p> Prepared<'p> {
139    pub(super) fn new(
140        program: &'p DynamicProgram,
141        kind: &BackendKind,
142        seed: u64,
143        max_steps: u64,
144    ) -> Result<Self> {
145        let summary = program.summary_circuit();
146        if !kind.is_auto() {
147            validate_explicit_backend(kind, &summary)?;
148        }
149        let plan = match resolve(kind, &summary, false) {
150            #[cfg(feature = "distributed")]
151            ExecutionPlan::Backend(BackendPlan::Distributed(_)) => {
152                return Err(decline(
153                    kind,
154                    "the distributed statevector runs its ranks in lockstep on one circuit",
155                ));
156            }
157            ExecutionPlan::Backend(plan) => plan,
158            _ => {
159                return Err(decline(
160                    kind,
161                    "this engine keeps no per-shot state to walk a control-flow graph on",
162                ));
163            }
164        };
165        let probe = plan.build(seed);
166        let blocks = program
167            .blocks()
168            .iter()
169            .map(|block| {
170                let expanded = expand_for_backend(&*probe, &block.circuit);
171                let fused = fuse_for_backend(&*probe, &expanded);
172                PreparedBlock {
173                    instructions: fused.into_owned().instructions,
174                    block,
175                }
176            })
177            .collect();
178        Ok(Self {
179            program,
180            plan,
181            blocks,
182            initial: program
183                .variables()
184                .iter()
185                .map(|variable| variable.initial)
186                .collect(),
187            max_steps,
188        })
189    }
190
191    fn run(&self, kind: &BackendKind, num_shots: usize, seed: u64) -> Result<ShotsResult> {
192        let num_qubits = self.program.num_qubits();
193        if self.plan.is_host_statevector() && num_qubits < max_statevector_qubits() {
194            #[cfg(feature = "parallel")]
195            let parallel =
196                super::shots_split_across_workers(kind, &[(self.plan.resolved(), num_qubits)])
197                    && num_qubits >= branching::MIN_QUBITS_FOR_PARALLEL_SPLITS;
198            #[cfg(not(feature = "parallel"))]
199            let parallel = false;
200            let limits = BranchLimits::for_width(num_qubits, parallel);
201            return branching::branch_program_shots(self, num_shots, seed, &limits);
202        }
203        self.run_per_shot(kind, num_shots, seed)
204    }
205
206    fn run_per_shot(&self, kind: &BackendKind, num_shots: usize, seed: u64) -> Result<ShotsResult> {
207        let reuse = self.plan.is_host_statevector();
208        let shot = |state: &mut ShotState, i: usize| -> Result<(Vec<bool>, RunMetadata)> {
209            let shot_seed = mix_seed(seed, i);
210            if reuse {
211                let backend = state
212                    .statevector
213                    .get_or_insert_with(|| StatevectorBackend::new(shot_seed));
214                backend.rng = rand_chacha::ChaCha8Rng::seed_from_u64(shot_seed);
215                let bits = self.run_shot(backend, &mut state.vars)?;
216                return Ok((bits, backend_metadata(backend)));
217            }
218            let mut backend = self.plan.build(shot_seed);
219            let bits = self.run_shot(&mut *backend, &mut state.vars)?;
220            Ok((bits, backend_metadata(&*backend)))
221        };
222
223        #[cfg(feature = "parallel")]
224        if num_shots > 1
225            && super::shots_split_across_workers(
226                kind,
227                &[(self.plan.resolved(), self.program.num_qubits())],
228            )
229        {
230            use rayon::prelude::*;
231            let runs: Vec<Result<(Vec<bool>, RunMetadata)>> = (0..num_shots)
232                .into_par_iter()
233                .map_init(ShotState::default, |state, i| shot(state, i))
234                .collect();
235            return self.fold(runs);
236        }
237        #[cfg(not(feature = "parallel"))]
238        let _ = kind;
239        let mut state = ShotState::default();
240        self.fold((0..num_shots).map(|i| shot(&mut state, i)))
241    }
242
243    fn fold(
244        &self,
245        runs: impl IntoIterator<Item = Result<(Vec<bool>, RunMetadata)>>,
246    ) -> Result<ShotsResult> {
247        let runs = runs.into_iter();
248        let mut shots = Vec::with_capacity(runs.size_hint().0);
249        let mut metadata = RunMetadata::exact(self.plan.resolved());
250        for (i, run) in runs.enumerate() {
251            let (bits, shot_metadata) = run?;
252            if i == 0 {
253                metadata = shot_metadata;
254            } else {
255                metadata.weaken_with(&shot_metadata);
256            }
257            shots.push(bits);
258        }
259        Ok(
260            ShotsResult::from_shots(shots, self.program.num_classical_bits())
261                .with_metadata(metadata),
262        )
263    }
264
265    /// Walk the graph once from a fresh state, returning the classical bits.
266    fn run_shot(
267        &self,
268        backend: &mut dyn Backend,
269        vars: &mut Vec<ClassicalValue>,
270    ) -> Result<Vec<bool>> {
271        backend.init(self.program.num_qubits(), self.program.num_classical_bits())?;
272        vars.clear();
273        vars.extend_from_slice(&self.initial);
274        let mut steps = 0;
275        self.count_step(0, &mut steps)?;
276        self.apply_block(0, backend)?;
277        self.run_from(0, steps, backend, vars)?;
278        Ok(backend.classical_results().to_vec())
279    }
280
281    /// Carry a shot whose state has run block `at`'s circuit to the end of the
282    /// graph, `steps` blocks entered so far.
283    pub(super) fn run_from(
284        &self,
285        mut at: usize,
286        mut steps: u64,
287        backend: &mut dyn Backend,
288        vars: &mut [ClassicalValue],
289    ) -> Result<()> {
290        while let Some(next) = self.finish_block(at, backend, vars)? {
291            self.count_step(next, &mut steps)?;
292            self.apply_block(next, backend)?;
293            at = next;
294        }
295        Ok(())
296    }
297
298    fn apply_block(&self, at: usize, backend: &mut dyn Backend) -> Result<()> {
299        let instructions = &self.blocks[at].instructions;
300        if !instructions.is_empty() {
301            backend.apply_instructions(instructions)?;
302        }
303        Ok(())
304    }
305
306    /// Count entry into block `at` against the step bound.
307    pub(super) fn count_step(&self, at: usize, steps: &mut u64) -> Result<()> {
308        if *steps == self.max_steps {
309            return Err(self.step_limit(at));
310        }
311        *steps += 1;
312        Ok(())
313    }
314
315    /// Run block `at`'s actions and terminator, returning the next block or
316    /// `None` at the end of the graph.
317    pub(super) fn finish_block(
318        &self,
319        at: usize,
320        backend: &mut dyn Backend,
321        vars: &mut [ClassicalValue],
322    ) -> Result<Option<usize>> {
323        let block = self.blocks[at].block;
324        let variables = self.program.variables();
325        for action in &block.actions {
326            match action {
327                Action::Assign { var, value } => {
328                    let evaluated = value.eval(backend.classical_results(), vars)?;
329                    vars[var.index()] = variables[var.index()].ty.store(evaluated);
330                }
331                Action::Rotation {
332                    kind,
333                    targets,
334                    angle,
335                } => {
336                    let angle = angle.eval(backend.classical_results(), vars)?.to_float();
337                    if !angle.is_finite() {
338                        return Err(PrismError::InvalidParameter {
339                            message: format!(
340                                "a runtime {kind:?} angle evaluated to {angle}, which no \
341                                 gate can carry"
342                            ),
343                        });
344                    }
345                    backend.apply(&Instruction::Gate {
346                        gate: kind.gate(angle),
347                        targets: targets.clone(),
348                    })?;
349                }
350            }
351        }
352        Ok(match &block.terminator {
353            Terminator::Jump(target) => Some(target.index()),
354            Terminator::Branch {
355                condition,
356                then,
357                otherwise,
358            } => {
359                if condition.eval(backend.classical_results(), vars)?.truthy() {
360                    Some(then.index())
361                } else {
362                    Some(otherwise.index())
363                }
364            }
365            Terminator::End => None,
366        })
367    }
368
369    /// Block `at`'s fused circuit, and whether the shot ends with it: no
370    /// actions follow and its terminator ends the graph.
371    pub(super) fn block_body(&self, at: usize) -> (&[Instruction], bool) {
372        let prepared = &self.blocks[at];
373        let ends =
374            prepared.block.actions.is_empty() && prepared.block.terminator == Terminator::End;
375        (&prepared.instructions, ends)
376    }
377
378    pub(super) fn initial_vars(&self) -> &[ClassicalValue] {
379        &self.initial
380    }
381
382    pub(super) fn num_qubits(&self) -> usize {
383        self.program.num_qubits()
384    }
385
386    pub(super) fn num_classical_bits(&self) -> usize {
387        self.program.num_classical_bits()
388    }
389
390    fn step_limit(&self, at: usize) -> PrismError {
391        PrismError::StepLimit {
392            region: match &self.blocks[at].block.loop_name {
393                Some(name) => format!("loop `{name}`"),
394                None => format!("block {at}"),
395            },
396            max_steps: self.max_steps,
397        }
398    }
399}
400
401/// What one worker keeps across the shots it runs.
402#[derive(Default)]
403struct ShotState {
404    statevector: Option<StatevectorBackend>,
405    vars: Vec<ClassicalValue>,
406}
407
408fn decline(kind: &BackendKind, reason: &str) -> PrismError {
409    PrismError::IncompatibleBackend {
410        backend: format!("{kind:?}"),
411        reason: format!("a dynamic program runs once per shot, and {reason}"),
412    }
413}
414
415#[cfg(test)]
416mod tests {
417    use super::*;
418    use crate::circuit::dynamic::{BinaryOp, ClassicalExpr, ClassicalType, DynamicProgramBuilder};
419    use crate::gates::Gate;
420    use crate::sim::ResolvedBackend;
421
422    fn rus() -> DynamicProgram {
423        let mut b = DynamicProgramBuilder::new(2, 1);
424        b.add_gate(Gate::H, &[0]).add_gate(Gate::Cx, &[0, 1]);
425        b.add_measure(0, 0);
426        b.begin_while("retry", ClassicalExpr::Bit(0));
427        b.add_reset(0).add_gate(Gate::H, &[0]).add_measure(0, 0);
428        b.end().unwrap();
429        b.build().unwrap()
430    }
431
432    #[test]
433    fn repeat_until_success_always_ends_on_zero() {
434        let shots = simulate_program(&rus()).seed(7).shots(200).unwrap();
435        assert!(shots.shots.iter().all(|shot| !shot[0]));
436        assert_eq!(shots.metadata.backend, ResolvedBackend::Stabilizer);
437    }
438
439    #[test]
440    fn a_loop_past_the_step_bound_names_itself() {
441        let mut b = DynamicProgramBuilder::new(1, 0);
442        b.begin_while("forever", ClassicalExpr::from(true));
443        b.add_gate(Gate::X, &[0]);
444        b.end().unwrap();
445        let program = b.build().unwrap();
446        let err = simulate_program(&program)
447            .max_steps(50)
448            .seed(1)
449            .shots(2)
450            .unwrap_err();
451        match err {
452            PrismError::StepLimit { region, max_steps } => {
453                assert_eq!(region, "loop `forever`");
454                assert_eq!(max_steps, 50);
455            }
456            other => panic!("expected a step limit, got {other:?}"),
457        }
458    }
459
460    #[test]
461    fn a_reused_statevector_matches_a_fresh_one_per_shot() {
462        let mut b = DynamicProgramBuilder::new(2, 2);
463        let n = b.declare("n", ClassicalType::Uint { width: 4 }, 0i64.into());
464        b.add_gate(Gate::Ry(0.7), &[0]).add_gate(Gate::Cx, &[0, 1]);
465        b.add_measure(0, 0);
466        b.begin_while(
467            "count",
468            ClassicalExpr::binary(BinaryOp::Lt, n.into(), 3i64.into()),
469        );
470        b.add_gate(Gate::Ry(0.4), &[1]).add_measure(1, 1);
471        b.assign(
472            n,
473            ClassicalExpr::binary(BinaryOp::Add, n.into(), 1i64.into()),
474        );
475        b.end().unwrap();
476        let program = b.build().unwrap();
477        let prepared = Prepared::new(&program, &BackendKind::Statevector, 3, 100).unwrap();
478        let reused = prepared
479            .run_per_shot(&BackendKind::Statevector, 64, 3)
480            .unwrap();
481        let fresh: Vec<Vec<bool>> = (0..64)
482            .map(|i| {
483                let mut backend = StatevectorBackend::new(mix_seed(3, i));
484                prepared.run_shot(&mut backend, &mut Vec::new()).unwrap()
485            })
486            .collect();
487        assert_eq!(reused.shots, fresh);
488    }
489}
prism_q/sim/dynamic.rs

prism_q/sim/
dynamic.rs
