diff --git a/CHANGELOG.md b/CHANGELOG.md index db3e077f..8f2ac2b9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - The samplers now default to `strategy="active_statevector"`; pass `strategy="cat5"` for the previous stabilizer-rank default. `CompiledStateProbs` keeps `"cat5"`. (#203) +- `strategy="active_statevector"` no longer uses JAX: it samples on the CPU + with numba kernels, also when JAX has a GPU, and trajectory sharing is on + by default (`trajectory_sharing=None` means True for this strategy and + False for the others). Amplitudes are complex128. Postselected + detectors of a component prune its shots with sharing off as well, and also + when no direct detector is postselected. (#199) +- Without a `batch_size`, a statevector call is split into one batch per + numba thread whenever each batch gets at least 4096 shots, and smaller + calls run as one batch on the parallel kernels. An explicit `batch_size` + keeps samples independent of the thread count. (#199) +- Noise sampling runs in numba kernels, so the samples for a given seed + differ from earlier versions; `numba>=0.61` is now a dependency. (#199) - Tsim now uses zxparam instead of pyzx-param. Circuits are built directly as quantum/classical ZX diagrams (no more composing with the adjoint), and reduction and stabilizer-rank decomposition run in zxparam. - `Circuit.get_graph()` and `Circuit.get_sampling_graph()` return `zxparam.Graph`. - In `Circuit.get_graph()` (and the `"zx"` diagram), every detector and observable has a classical output boundary, with ports after the qubit outputs, so the graph's tensor is the circuit's channel together with the detector and observable bits. diff --git a/pyproject.toml b/pyproject.toml index 4fddd037..ffeeb112 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ dependencies = [ "ipywidgets>=8.0.0", "jax>=0.4.38", "lxml>=5.0.0", + "numba>=0.61", "numpy>=1.25.0", "psutil>=5.9.0", "stim>=1.0.0", diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index d1022f40..98436cae 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -865,7 +865,7 @@ def compile_sampler( *, strategy: SamplingStrategy = "active_statevector", seed: int | None = None, - trajectory_sharing: bool = False, + trajectory_sharing: bool | None = None, ) -> CompiledMeasurementSampler: """Compile circuit into a measurement sampler. @@ -873,9 +873,12 @@ def compile_sampler( strategy: "active_statevector" (the default) to simulate each non-direct component on a dense state vector over its active qubits, or a stabilizer rank decomposition strategy ("cat5", - "bss", "cutting"). + "bss", "cutting"). "active_statevector" samples on the CPU, up + to ``numba.get_num_threads()`` batches at once, each with its + own memory. trajectory_sharing: With ``strategy="active_statevector"``, simulate shots - with equal histories on one lane (experimental). + with equal histories on one lane. Defaults to True for + ``"active_statevector"`` and False otherwise. seed: Random seed for the sampler. IMPORTANT: Currently, the sampler will only produce deterministic samples for fixed batch size. If deterministic samples are needed, the batch size should be set @@ -896,7 +899,7 @@ def compile_detector_sampler( *, strategy: SamplingStrategy = "active_statevector", seed: int | None = None, - trajectory_sharing: bool = False, + trajectory_sharing: bool | None = None, ) -> CompiledDetectorSampler: """Compile circuit into a detector sampler. @@ -909,9 +912,12 @@ def compile_detector_sampler( strategy: "active_statevector" (the default) to simulate each non-direct component on a dense state vector over its active qubits, or a stabilizer rank decomposition strategy ("cat5", - "bss", "cutting"). + "bss", "cutting"). "active_statevector" samples on the CPU, up + to ``numba.get_num_threads()`` batches at once, each with its + own memory. trajectory_sharing: With ``strategy="active_statevector"``, simulate shots - with equal histories on one lane (experimental). + with equal histories on one lane. Defaults to True for + ``"active_statevector"`` and False otherwise. seed: Random seed for the sampler. IMPORTANT: Currently, the sampler will only produce deterministic samples for fixed batch size. If deterministic samples are needed, the batch size should be set diff --git a/src/tsim/compile/pipeline.py b/src/tsim/compile/pipeline.py index 1496742a..bb99ee37 100644 --- a/src/tsim/compile/pipeline.py +++ b/src/tsim/compile/pipeline.py @@ -107,10 +107,10 @@ def compile_program( return CompiledProgram( components=tuple(compiled_components), - direct_f_indices=jnp.array([e[1] for e in direct_entries], dtype=jnp.int32), - direct_flips=jnp.array([e[2] for e in direct_entries], dtype=jnp.bool_), - output_order=jnp.asarray(output_order), - output_reindex=None if is_identity else jnp.asarray(reindex), + direct_f_indices=np.array([e[1] for e in direct_entries], dtype=np.int32), + direct_flips=np.array([e[2] for e in direct_entries], dtype=np.bool_), + output_order=output_order, + output_reindex=None if is_identity else reindex, num_outputs=prepared.num_outputs, num_detectors=prepared.num_detectors, ) diff --git a/src/tsim/core/types.py b/src/tsim/core/types.py index 14ecd6bb..95d81d99 100644 --- a/src/tsim/core/types.py +++ b/src/tsim/core/types.py @@ -111,9 +111,9 @@ class CompiledProgram: """ components: tuple[CompiledComponent | StatevectorComponent, ...] - direct_f_indices: Array - direct_flips: Array - output_order: Array - output_reindex: Array | None + direct_f_indices: np.ndarray + direct_flips: np.ndarray + output_order: np.ndarray + output_reindex: np.ndarray | None num_outputs: int num_detectors: int diff --git a/src/tsim/noise/channels.py b/src/tsim/noise/channels.py index 58f32d0f..f8c40e33 100644 --- a/src/tsim/noise/channels.py +++ b/src/tsim/noise/channels.py @@ -4,9 +4,20 @@ from collections import defaultdict from dataclasses import dataclass +from itertools import combinations +from typing import NamedTuple +import numba import numpy as np +from tsim.utils.kernels import jit, kernel +from tsim.utils.rng import next_skip, next_uniform + +# NOTE: shots are sampled in blocks of this many rows, in parallel, each block +# from its own random stream; a block's rows stay in cache while every channel +# visits them. +_BLOCK_SHOTS = 4096 + @dataclass class Channel: @@ -453,12 +464,20 @@ def absorb_subset_channels(channels: list[Channel], max_bits: int = 4) -> list[C max_bits: Maximum number of bits allowed per channel Returns: - List with no channel being a strict subset of another + List in which no channel is a strict subset of another of at most + ``max_bits`` signatures """ - # Sort by number of bits (largest first) for efficient processing + # NOTE: largest first, so that every channel is absorbed into the first + # channel that contains it. channels = sorted(channels, key=lambda c: -len(c.unique_col_ids)) + # NOTE: a channel of at most max_bits signatures has few strict subsets, so + # they are looked up instead of compared with every other channel. + by_signatures: dict[frozenset[int], list[int]] = defaultdict(list) + for j, channel in enumerate(channels): + by_signatures[frozenset(channel.unique_col_ids)].append(j) + result: list[Channel] = [] absorbed: set[int] = set() @@ -466,22 +485,25 @@ def absorb_subset_channels(channels: list[Channel], max_bits: int = 4) -> list[C if i in absorbed: continue - set_i = set(channel_i.unique_col_ids) - - # Try to absorb smaller channels into this one + signatures = sorted(set(channel_i.unique_col_ids)) current_probs = channel_i.probs.copy() current_col_ids = channel_i.unique_col_ids - for j, channel_j in enumerate(channels): - if j <= i or j in absorbed: - continue - - set_j = set(channel_j.unique_col_ids) - - # Check if j is a strict subset of i - if set_j < set_i and len(set_i) <= max_bits: + if len(signatures) <= max_bits: + subsets = ( + frozenset(subset) + for size in range(len(signatures)) + for subset in combinations(signatures, size) + ) + subset_channels = sorted( + j + for subset in subsets + for j in by_signatures.get(subset, ()) + if j > i and j not in absorbed + ) + for j in subset_channels: # Expand channel_j to match channel_i's signatures and convolve - expanded_j = expand_channel(channel_j, current_col_ids) + expanded_j = expand_channel(channels[j], current_col_ids) current_probs = xor_convolve(current_probs, expanded_j.probs) absorbed.add(j) @@ -513,6 +535,35 @@ def simplify_channels( return channels +def unique_columns(matrix: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Return ``np.unique(matrix, axis=1, return_inverse=True)``, with a 1-D inverse. + + ``np.unique`` sorts every column, comparing whole columns, which is slow for + a large sparse matrix. Here a dict over the bit-packed columns finds the + distinct ones, and only those are sorted. Big-endian packing keeps the + lexicographic order of the columns, so the result is the same. + + Args: + matrix: Binary matrix with entries 0 or 1 (nonzero entries count as 1). + + Returns: + The distinct columns in lexicographic order, and the index of each + column of ``matrix`` among them. + + """ + packed = np.ascontiguousarray(np.packbits(np.asarray(matrix) != 0, axis=0).T) + ids: dict[bytes, int] = {} + column_ids = np.array( + [ids.setdefault(column.tobytes(), len(ids)) for column in packed], + dtype=np.int64, + ) + _, first = np.unique(column_ids, return_index=True) + _, order = np.unique(packed[first], axis=0, return_index=True) + rank = np.empty_like(order) + rank[order] = np.arange(len(order)) + return matrix[:, first[order]], rank[column_ids] + + class ChannelSampler: """Samples from multiple error channels and transforms to a reduced basis. @@ -528,9 +579,9 @@ class ChannelSampler: Channels are automatically simplified by: 1. Removing bits e_i that do not affect any f-variable (i.e. all-zero columns in error_transform) 2. Folding duplicate column IDs, i.e. channels whose column signatures contain duplicate IDs. - 2. Merging channels with identical column signatures, i.e. channels whose corresponding + 3. Merging channels with identical column signatures, i.e. channels whose corresponding columns in error_transform are identical. - 3. Absorbing channels whose signatures are subsets of others, i.e. channels whose corresponding + 4. Absorbing channels whose signatures are subsets of others, i.e. channels whose corresponding columns in error_transform are a strict subset of another channel's columns. Example: @@ -560,7 +611,7 @@ def __init__( seed: Random seed for sampling. If None, a random seed is generated. """ - unique_cols, inverse = np.unique(error_transform, axis=1, return_inverse=True) + unique_cols, inverse = unique_columns(error_transform) # Signature matrix: each row is a unique column signature signature_matrix = unique_cols.T # shape (num_signatures, num_f) @@ -584,88 +635,127 @@ def __init__( self._rng = np.random.default_rng( seed if seed is not None else np.random.default_rng().integers(0, 2**30) ) - self._sparse_data = self._precompute_sparse( - self.channels, self.signature_matrix - ) + self._tables = _fire_tables(self.channels, self.signature_matrix) - @staticmethod - def _precompute_sparse( - channels: list[Channel], signature_matrix: np.ndarray - ) -> list[tuple[float, np.ndarray, np.ndarray]]: - """Precompute per-channel data for geometric-skip sampling. + def block_seeds(self, num_samples: int) -> np.ndarray: + """Draw the seeds of the blocks of the next ``num_samples`` samples.""" + num_blocks = -(-num_samples // _BLOCK_SHOTS) + return self._rng.integers(0, 2**64, size=num_blocks, dtype=np.uint64) - For each channel with non-trivial fire probability, computes: - - p_fire: probability of any non-identity outcome - - cond_cdf: conditional CDF over non-identity outcomes - - xor_patterns: precomputed XOR output patterns per outcome + def sample( + self, num_samples: int = 1, seeds: np.ndarray | None = None + ) -> np.ndarray: + """Sample from all error channels and transform to new error basis. + + Uses geometric-skip sampling, optimized for low-noise regimes where + P(non-identity) << 1 per channel. Blocks of shots are sampled in + parallel, each from its own random stream, so the samples depend on the + seed and ``num_samples`` but not on the number of threads. Args: - channels: List of noise channels to precompute data for. - signature_matrix: Binary matrix of shape (num_e, num_f) mapping - error-variable columns to output f-variables. + num_samples: Number of samples to draw. + seeds: The blocks' seeds from ``block_seeds(num_samples)``; drawn + here if ``None``. Returns: - List of (p_fire, cond_cdf, xor_patterns) tuples, one per channel - with non-trivial fire probability. ``p_fire`` is a float, - ``cond_cdf`` is a float64 array of shape (n_outcomes - 1,), and - ``xor_patterns`` is a uint8 array of shape (n_outcomes - 1, num_f). + NumPy array of shape (num_samples, num_f) with uint8 values indicating + which f-variables are set for each sample. """ - data: list[tuple[float, np.ndarray, np.ndarray]] = [] - for ch in channels: - probs = ch.probs.astype(np.float64) - p_fire = 1.0 - float(probs[0]) - n_outcomes = len(probs) - - if p_fire <= 1e-15 or n_outcomes <= 1: - continue - - cond_cdf = np.cumsum(probs[1:] / p_fire, dtype=np.float64) - cond_cdf /= cond_cdf[-1] - - col_ids = np.asarray(ch.unique_col_ids) - num_bits = len(col_ids) - outcomes = np.arange(1, n_outcomes) - bits_mask = ((outcomes[:, None] >> np.arange(num_bits)) & 1).astype( - np.uint8 + num_f = self.signature_matrix.shape[1] + result = np.empty((num_samples, num_f), dtype=np.uint8) + if seeds is None: + seeds = self.block_seeds(num_samples) + num_blocks = -(-num_samples // _BLOCK_SHOTS) + if len(seeds) != num_blocks: + raise ValueError( + f"{num_samples} samples need {num_blocks} block seeds, " + f"got {len(seeds)}" ) - xor_patterns = (bits_mask @ signature_matrix[col_ids] % 2).astype(np.uint8) + work = num_samples * len(self._tables.p_fire) + _sample_blocks(work, result, seeds, *self._tables) + return result - data.append((p_fire, cond_cdf, xor_patterns)) - return data - def sample(self, num_samples: int = 1) -> np.ndarray: - """Sample from all error channels and transform to new error basis. +class _FireTables(NamedTuple): + """Geometric-skip sampling tables of the channels that can fire. - Uses geometric-skip sampling, optimized for low-noise regimes where - P(non-identity) << 1 per channel. + Channel ``c`` fires with probability ``p_fire[c]``. Its non-identity + outcomes are ``outcome_start[c] … outcome_start[c+1] - 1``, picked by the + conditional CDF ``cdf`` over those indices, and outcome ``k`` flips the + f-variables ``flip_cols[flip_start[k] : flip_start[k+1]]``. + """ - Args: - num_samples: Number of samples to draw. + p_fire: np.ndarray + outcome_start: np.ndarray + cdf: np.ndarray + flip_start: np.ndarray + flip_cols: np.ndarray - Returns: - NumPy array of shape (num_samples, num_f) with uint8 values indicating - which f-variables are set for each sample. - """ - num_outputs = self.signature_matrix.shape[1] - result = np.zeros((num_samples, num_outputs), dtype=np.uint8) +def _fire_tables(channels: list[Channel], signature_matrix: np.ndarray) -> _FireTables: + """Return the sampling tables of ``channels``. - for p_fire, cond_cdf, xor_pats in self._sparse_data: - expected = num_samples * p_fire - sigma = np.sqrt(expected * (1.0 - p_fire)) - # At 7 sigma, we undersample in about 1 out of 1e12 cases - n_draws = int(expected + 7.0 * sigma) + 100 + A channel fires with the total probability of its non-identity outcomes, + clipped to [0, 1]: ``1 - probs[0]`` can exceed 1 by a rounding error, and it + loses precision for small probabilities. - positions = np.cumsum(self._rng.geometric(p_fire, size=n_draws)) - 1 - positions = positions[positions < num_samples] + Args: + channels: The noise channels. + signature_matrix: Binary matrix of shape (num_signatures, num_f) mapping + each column ID to the f-variables it flips. - if len(positions) == 0: - continue + """ + p_fire: list[float] = [] + cdf: list[float] = [] + flip_cols: list[int] = [] + outcome_start, flip_start = [0], [0] + for channel in channels: + probs = channel.probs.astype(np.float64) + conditional = np.cumsum(np.clip(probs[1:], 0, None)) + fire = min(float(conditional[-1]), 1.0) if len(conditional) else 0.0 + if fire <= 1e-15: + continue + p_fire.append(fire) + cdf.extend(conditional / conditional[-1]) + outcome_start.append(len(cdf)) + col_ids = np.asarray(channel.unique_col_ids) + outcomes = np.arange(1, len(probs)) + bits = (outcomes[:, None] >> np.arange(len(col_ids))) & 1 + for flips in bits @ signature_matrix[col_ids] % 2: + flip_cols.extend(np.flatnonzero(flips).tolist()) + flip_start.append(len(flip_cols)) + return _FireTables( + p_fire=np.array(p_fire, dtype=np.float64), + outcome_start=np.array(outcome_start, dtype=np.int64), + cdf=np.array(cdf, dtype=np.float64), + flip_start=np.array(flip_start, dtype=np.int64), + flip_cols=np.array(flip_cols, dtype=np.int64), + ) - outcome_idx = np.searchsorted( - cond_cdf, self._rng.uniform(size=len(positions)) - ) - result[positions] ^= xor_pats[outcome_idx] - return result +@jit +def _sample_block(out, state, p_fire, outcome_start, cdf, flip_start, flip_cols): + """Clear the rows of ``out``, then XOR in every channel's outcomes.""" + out[:] = 0 + for c in range(len(p_fire)): + log_keep = np.log1p(-p_fire[c]) + first, last = outcome_start[c], outcome_start[c + 1] + row = -1.0 + while True: + state, skip = next_skip(state, log_keep) + row += 1.0 + skip + if row >= out.shape[0]: + break + state, u = next_uniform(state) + k = first + np.searchsorted(cdf[first:last], u, side="right") + for j in range(flip_start[k], flip_start[k + 1]): + out[int(row), flip_cols[j]] ^= 1 + + +@kernel +def _sample_blocks(out, seeds, p_fire, outcome_start, cdf, flip_start, flip_cols): + """Sample block ``b`` of ``_BLOCK_SHOTS`` rows from the stream ``seeds[b]``.""" + for b in numba.prange(len(seeds)): + rows = out[b * _BLOCK_SHOTS : (b + 1) * _BLOCK_SHOTS] + _sample_block(rows, seeds[b], p_fire, outcome_start, cdf, flip_start, flip_cols) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 4e1374b4..9cbfbae0 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -5,10 +5,11 @@ import warnings from collections import deque from math import ceil -from typing import TYPE_CHECKING, Literal, overload +from typing import TYPE_CHECKING, Any, Literal, Protocol, overload import jax import jax.numpy as jnp +import numba import numpy as np import psutil @@ -17,10 +18,16 @@ from tsim.core.graph import prepare_graph from tsim.core.types import CompiledComponent, CompiledProgram, SamplingStrategy from tsim.noise.channels import ChannelSampler -from tsim.statevector.execute import lane_bytes, sample_statevector_component from tsim.statevector.program import STABILIZER_RANK_HINT, StatevectorComponent -from tsim.statevector.shared import Sharing, sample_shared_component, sharing_layout +from tsim.statevector.shared import ( + SharingLayout, + lane_bytes, + sample_shared_component, + sharing_layout, +) from tsim.utils.cuda_helpers import copy_d2h +from tsim.utils.kernels import kernel, run_tasks +from tsim.utils.rng import GOLDEN if TYPE_CHECKING: from jax import Array as PRNGKey @@ -30,6 +37,18 @@ # NOTE: beyond 2^16 lanes, statevector batches cost memory without being faster. STATEVECTOR_MAX_BATCH_SIZE = 2**16 +# NOTE: without a given batch size, a call is split into a batch per thread when +# every batch gets at least this many shots, whose fixed costs would dominate +# otherwise. A call too small for that stays one batch, which the parallel +# kernels spread over the threads instead. +MIN_THREAD_BATCH = 2**12 + +# NOTE: postselection screens shots for direct discards in chunks of at most +# this many, several chunks per thread at a time, so that calls of a few batches +# still keep every thread busy and slow cores take fewer chunks. +SCREEN_SHOTS = 2**14 +SCREEN_CHUNKS_PER_THREAD = 2 + def _sample_component( component: CompiledComponent, @@ -98,7 +117,7 @@ def _sample_component_jit( def sample_component( - component: CompiledComponent | StatevectorComponent, + component: CompiledComponent, f_params: jax.Array, key: PRNGKey, ) -> tuple[jax.Array, PRNGKey, jax.Array]: @@ -114,89 +133,375 @@ def sample_component( (batch_size, num_outputs_for_component). """ - if isinstance(component, StatevectorComponent): - return sample_statevector_component(component, f_params, key) # Skip JIT for small components (overhead not worth it) if len(component.output_indices) <= 1: return _sample_component(component, f_params, key) return _sample_component_jit(component, f_params, key) +def _check_norm(max_norm_deviation: float | jax.Array, threshold: float) -> None: + """Raise on a vanishing normalization; warn when it drifts beyond ``threshold``.""" + if np.isclose(max_norm_deviation, 1): + raise ValueError( + "A vanishing marginal probability distribution was encountered (normalization 0). " + "This is likely the result of an underflow error. Please report this " + "as a bug at https://github.com/QuEraComputing/tsim/issues/new." + ) # pragma: no cover + # NOTE: written so that a NaN deviation also warns. + if not max_norm_deviation <= threshold: + warnings.warn( + "A marginal probability was not normalized correctly " + f"(normalization deviated from 1 by {max_norm_deviation:.1e}). " + "This is likely a floating point precision issue.", + stacklevel=3, + ) + + def sample_program( - program: CompiledProgram, - f_params: jax.Array, - key: PRNGKey, - sharing: tuple[Sharing | None, ...] | None = None, - sharing_stats: deque[dict] | None = None, + program: CompiledProgram, f_params: jax.Array, key: PRNGKey ) -> jax.Array: - """Sample all outputs from a compiled program. + """Sample all outputs of a program of stabilizer-rank components with JAX. Args: program: The compiled program to sample from. f_params: Error parameters, shape (batch_size, num_f_params). key: JAX random key. - sharing: One sharing state per component, to sample statevector - components with shared trajectories; ``None`` samples one lane per - shot. - sharing_stats: If given, the shared executor appends its lane counts. Returns: Samples array of shape (batch_size, num_outputs), reordered to match the original output indices. """ - results: list[jax.Array] = [] - if program.num_outputs == 0: - batch_size = f_params.shape[0] - return jnp.zeros((batch_size, 0), dtype=jnp.bool_) - + return jnp.zeros((f_params.shape[0], 0), dtype=jnp.bool_) + results = [] if len(program.direct_f_indices) > 0: direct_bits = ( f_params[:, program.direct_f_indices].astype(jnp.bool_) ^ program.direct_flips ) results.append(direct_bits) - - for index, component in enumerate(program.components): - shared = sharing[index] if sharing is not None else None - if isinstance(component, StatevectorComponent) and shared is not None: - samples, key, max_norm_deviation, stats = sample_shared_component( - component, shared, f_params, key - ) - if sharing_stats is not None: - sharing_stats.append(stats) - else: - samples, key, max_norm_deviation = sample_component( - component, f_params, key - ) - if np.isclose(max_norm_deviation, 1): - raise ValueError( - "A vanishing marginal probability distribution was encountered (normalization 0). " - "This is likely the result of an underflow error. Please report this " - "as a bug at https://github.com/QuEraComputing/tsim/issues/new." - ) # pragma: no cover - if max_norm_deviation > component.norm_warning_threshold: - warnings.warn( - "A marginal probability was not normalized correctly " - f"(normalization deviated from 1 by {max_norm_deviation:.1e}). " - "This is likely a floating point precision issue.", - stacklevel=2, - ) + for component in program.components: + assert isinstance(component, CompiledComponent) + samples, key, max_norm_deviation = sample_component(component, f_params, key) + _check_norm(max_norm_deviation, component.norm_warning_threshold) results.append(samples) - combined = jnp.concatenate(results, axis=1) if program.output_reindex is not None: combined = combined[:, program.output_reindex] return combined +def sample_statevector_program( + program: CompiledProgram, + f_params: np.ndarray, + seed: int, + layouts: tuple[SharingLayout, ...], + direct_flips: np.ndarray, + *, + share: bool, + sharing_stats: deque[dict] | None = None, +) -> np.ndarray: + """Sample all outputs of a program of statevector components, without JAX. + + Args: + program: The compiled program to sample from. + f_params: Error parameters, shape (batch_size, num_f_params). + seed: Seed of the batch's random numbers, below 2^64. + layouts: The sharing layout of each component. + direct_flips: The flip of each direct output, with its reference bit. + share: Whether shots with equal histories share a lane. + sharing_stats: If given, the executor appends its lane counts. + + Returns: + Samples array of shape (batch_size, num_outputs) in the original order. + + """ + rows = f_params.shape[0] + results = [] + for index, (component, layout) in enumerate( + zip(program.components, layouts, strict=True) + ): + assert isinstance(component, StatevectorComponent) + component_seed = (seed + index * int(GOLDEN)) % 2**64 + samples, max_norm_deviation, stats = sample_shared_component( + component, layout, f_params, component_seed, share + ) + if sharing_stats is not None: + sharing_stats.append(stats) + _check_norm(max_norm_deviation, component.norm_warning_threshold) + results.append(samples) + component_bits = ( + np.concatenate(results, axis=1) if results else np.zeros((rows, 0), np.bool_) + ) + out = np.empty((rows, program.num_outputs), dtype=np.bool_) + _gather_outputs( + out.size, + f_params, + program.direct_f_indices, + direct_flips, + component_bits, + program.output_order, + out, + ) + return out + + +@kernel +def _write_direct_detectors(f, f_indices, flips, columns, postselect, out): + """Clear every row of ``out`` and write its direct detector ``j`` at ``columns[j]``. + + Detector ``j`` is ``f[row, f_indices[j]] ^ flips[j]``. Returns which rows + have a postselected detector of 1. + """ + discarded = np.zeros(f.shape[0], dtype=np.bool_) + for row in numba.prange(f.shape[0]): + out[row] = False + for j in range(len(f_indices)): + bit = (f[row, f_indices[j]] != 0) ^ flips[j] + out[row, columns[j]] = bit + if postselect[j] and bit: + discarded[row] = True + return discarded + + +@kernel +def _gather_outputs(f, f_indices, flips, component_bits, order, out): + """Write every row's outputs into ``out`` in their original order. + + Output ``k`` of the concatenation goes to column ``order[k]``: first the + direct outputs ``f[row, f_indices[j]] ^ flips[j]``, then ``component_bits``. + """ + num_direct = len(f_indices) + for row in numba.prange(f.shape[0]): + for j in range(num_direct): + out[row, order[j]] = (f[row, f_indices[j]] != 0) ^ flips[j] + for j in range(component_bits.shape[1]): + out[row, order[num_direct + j]] = component_bits[row, j] + + def _nbytes(tree: object) -> int: + if isinstance(tree, StatevectorComponent): + return sum(array.nbytes for array in tree.tape) return sum( v.nbytes for v in jax.tree_util.tree_leaves(tree) if isinstance(v, jax.Array) ) +class _Backend(Protocol): + """Samples batches of a program's non-direct components. + + Attributes: + prunes: Whether shots that a masked component detector discards stop + being simulated. + max_batch_size: The largest batch worth sampling at once, if any. + + """ + + prunes: bool + max_batch_size: int | None + + def max_threads(self) -> int: + """Return how many batches may run at once.""" + ... + + def next_seed(self) -> Any: + """Return the seed of the next batch, in the order of the batches.""" + ... + + def memory_available(self) -> int: + """Return the free memory of the device that holds the batches.""" + ... + + def peak_bytes_per_shot(self, num_f: int) -> int: + """Return the peak memory of one shot of a batch.""" + ... + + def sample( + self, + f_params: np.ndarray, + seed: Any, + *, + mask: np.ndarray | None = None, + flip: np.ndarray | None = None, + pad_to: int | None = None, + ) -> Any: + """Return the outputs of one batch, XOR ``flip`` (one bit per output). + + ``mask`` marks the postselected detectors and ``pad_to`` the number of + rows the batch may be padded to. + """ + ... + + def concatenate(self, batches: list) -> np.ndarray: + """Return the batches' outputs as one host array.""" + ... + + +class _StabilizerRankBackend: + """Samples stabilizer-rank programs with JAX, one batch at a time.""" + + prunes: bool = False + max_batch_size: int | None = None + + def __init__(self, program: CompiledProgram, seed: int) -> None: + self._program = program + self._key = jax.random.key(seed) + + def max_threads(self) -> int: + return 1 + + def next_seed(self) -> PRNGKey: + self._key, key = jax.random.split(self._key) + return key + + def memory_available(self) -> int: + device = jax.devices()[0] + if device.platform != "gpu": + return psutil.virtual_memory().available + stats = device.memory_stats() + return stats.get("bytes_limit", 8 * 1024**3) - stats.get("bytes_in_use", 0) + + def peak_bytes_per_shot(self, num_f: int) -> int: + peak = 0 + for component in self._program.components: + assert isinstance(component, CompiledComponent) + for circuit in component.compiled_scalar_graphs: + G = circuit.num_graphs + max_a = circuit.node_phases.phases.shape[1] + max_b = circuit.phase_terms.coeffs.shape[1] + max_c = circuit.pi_products.psi_params.shape[1] + max_d = circuit.phase_pairs.alpha.shape[1] + max_e = circuit.float_terms.table.shape[1] + largest = max(max_a * 16, max_b * 4, max_c * 4, max_d * 16, max_e * 16) + peak = max(peak, G * largest * 3) + return peak + + def sample( + self, + f_params: np.ndarray, + seed: PRNGKey, + *, + mask: np.ndarray | None = None, + flip: np.ndarray | None = None, + pad_to: int | None = None, + ) -> jax.Array: + """Sample the batch, padded with its first row to ``pad_to`` rows. + + Padding keeps the batch shapes, and so JAX's compilations, few. + """ + count = len(f_params) + if pad_to is not None and count < pad_to: + padding = np.repeat(f_params[:1], pad_to - count, axis=0) + f_params = np.concatenate([f_params, padding]) + samples = sample_program(self._program, jnp.asarray(f_params), seed) + if len(f_params) > count: + samples = samples[:count] + return samples if flip is None else samples ^ flip + + def concatenate(self, batches: list) -> np.ndarray: + # NOTE: one device-to-host copy of the concatenation is much faster than + # one per batch for large outputs. + combined = batches[0] if len(batches) == 1 else jnp.concatenate(batches) + jax.block_until_ready(combined) + return copy_d2h(combined) + + +class _StatevectorBackend: + """Samples statevector programs with numba kernels on the host. + + Up to ``numba.get_num_threads()`` batches run at once. + """ + + prunes: bool = True + max_batch_size: int | None = STATEVECTOR_MAX_BATCH_SIZE + + def __init__( + self, program: CompiledProgram, seed: int, num_detectors: int, share: bool + ) -> None: + self._program = program + self._num_detectors = num_detectors + self._share = share + self._seeds = np.random.default_rng([seed, 1]) + self._layouts: dict[bytes, tuple[tuple[SharingLayout, ...], np.ndarray]] = {} + self.sharing_stats: deque[dict] = deque(maxlen=10_000) + + def max_threads(self) -> int: + return numba.get_num_threads() + + def next_seed(self) -> int: + return int(self._seeds.integers(0, 2**64, dtype=np.uint64)) + + def memory_available(self) -> int: + return psutil.virtual_memory().available + + def peak_bytes_per_shot(self, num_f: int) -> int: + # NOTE: plus the f-params and the output records, twice over. + records = 2 * num_f + 2 * self._program.num_outputs + return max( + ( + lane_bytes(component) + records + for component in self._program.components + if isinstance(component, StatevectorComponent) + ), + default=0, + ) + + def sample( + self, + f_params: np.ndarray, + seed: int, + *, + mask: np.ndarray | None = None, + flip: np.ndarray | None = None, + pad_to: int | None = None, + ) -> np.ndarray: + """Sample the batch; shots that a ``mask`` detector discards are pruned.""" + layouts, direct_flips = self.layouts(mask, flip) + return sample_statevector_program( + self._program, + f_params, + seed, + layouts, + direct_flips, + share=self._share, + sharing_stats=self.sharing_stats, + ) + + def layouts( + self, mask: np.ndarray | None, flip: np.ndarray | None + ) -> tuple[tuple[SharingLayout, ...], np.ndarray]: + """Return the components' sharing layouts and the direct outputs' flips.""" + program = self._program + postselected = np.zeros(program.num_outputs, dtype=np.bool_) + if mask is not None: + postselected[: self._num_detectors] = mask + if flip is None: + flip = np.zeros(program.num_outputs, dtype=np.bool_) + key = postselected.tobytes() + flip.tobytes() + if key not in self._layouts: + layouts = [] + for component in program.components: + assert isinstance(component, StatevectorComponent) + outputs = list(component.output_indices) + layouts.append( + sharing_layout(component, postselected[outputs], flip[outputs]) + ) + direct = program.output_order[: len(program.direct_f_indices)] + self._layouts[key] = (tuple(layouts), program.direct_flips ^ flip[direct]) + return self._layouts[key] + + def concatenate(self, batches: list) -> np.ndarray: + return np.concatenate(batches) + + +def _check_sizes(shots: int, batch_size: int | None) -> None: + if shots < 0: + raise ValueError(f"shots must be non-negative, got {shots}") + if batch_size is not None and batch_size < 1: + raise ValueError(f"batch_size must be at least 1, got {batch_size}") + + class _CompiledSamplerBase: """Base class for compiled samplers with common initialization logic.""" @@ -208,7 +513,7 @@ def __init__( mode: Literal["sequential", "joint"], strategy: SamplingStrategy = "active_statevector", seed: int | None = None, - trajectory_sharing: bool = False, + trajectory_sharing: bool | None = None, ): """Initialize the sampler by compiling the circuit. @@ -219,27 +524,28 @@ def __init__( strategy: "active_statevector" (the default) to simulate each non-direct component on a dense state vector over its active qubits, or a stabilizer rank decomposition strategy ("cat5", - "bss", "cutting"). + "bss", "cutting"). "active_statevector" samples on the CPU, up + to ``numba.get_num_threads()`` batches at once, each with its + own memory. seed: Random seed. If None, a random seed is generated. Note that deterministic results are only guaranteed for a fixed batch size and fixed reference sample settings. trajectory_sharing: With ``strategy="active_statevector"``, simulate shots - with equal histories on one lane (experimental). + with equal histories on one lane. Defaults to True for + ``"active_statevector"`` and False otherwise. Raises: ValueError: If ``trajectory_sharing`` is combined with a stabilizer-rank strategy. """ + if trajectory_sharing is None: + trajectory_sharing = strategy == "active_statevector" if trajectory_sharing and strategy != "active_statevector": raise ValueError("trajectory_sharing needs strategy='active_statevector'") if seed is None: seed = int(np.random.default_rng().integers(0, 2**30)) - - self._key = jax.random.key(seed) self._trajectory_sharing = trajectory_sharing - self._sharing: dict[bytes, tuple[Sharing | None, ...]] = {} - self.sharing_stats: deque[dict] = deque(maxlen=10_000) prepared = prepare_graph(circuit, sample_detectors=sample_detectors) self._program = compile_program(prepared, mode=mode, strategy=strategy) @@ -253,34 +559,31 @@ def __init__( self.circuit = circuit self._num_detectors = prepared.num_detectors - - prog = self._program - self._direct_f_indices = np.asarray(prog.direct_f_indices) - self._direct_flips = np.asarray(prog.direct_flips, dtype=np.bool_) - self._direct_reindex = ( - np.asarray(prog.output_reindex) if prog.output_reindex is not None else None + self._backend: _Backend = ( + _StatevectorBackend( + self._program, seed, self._num_detectors, trajectory_sharing + ) + if strategy == "active_statevector" + else _StabilizerRankBackend(self._program, seed) ) - # Zero-copy fast path: f-indices are 0..n-1, no flips, no reindex. - # Hit by typical surface-code detector circuits at low noise. - n_direct = len(self._direct_f_indices) - self._direct_zero_copy = ( - n_direct > 0 - and self._direct_reindex is None - and not self._direct_flips.any() - and np.array_equal(self._direct_f_indices, np.arange(n_direct)) + + program = self._program + # NOTE: _gather_outputs writes column output_order[k] without bounds checks. + assert np.array_equal( + np.sort(program.output_order), np.arange(program.num_outputs) ) - self._direct_global_indices = np.asarray( - prog.output_order[:n_direct], dtype=np.int32 + columns = program.output_order[: len(program.direct_f_indices)] + is_detector = columns < self._num_detectors + self._direct_detectors = ( + program.direct_f_indices[is_detector], + program.direct_flips[is_detector], + columns[is_detector], ) - self._direct_output_mask = np.zeros(prog.num_outputs, dtype=np.bool_) - if n_direct > 0: - self._direct_output_mask[self._direct_global_indices] = True - self._direct_detector_mask = self._direct_output_mask[ - : self._num_detectors - ].copy() + self._direct_detector_mask = np.zeros(self._num_detectors, dtype=np.bool_) + self._direct_detector_mask[columns[is_detector]] = True statevector = [ - c for c in self._program.components if isinstance(c, StatevectorComponent) + c for c in program.components if isinstance(c, StatevectorComponent) ] if statevector: peak = self._peak_bytes_per_sample() @@ -293,95 +596,22 @@ def __init__( f"{STABILIZER_RANK_HINT}" ) - def _sharing_arguments( - self, - postselection_mask: np.ndarray | None, - detector_reference: np.ndarray | None = None, - ) -> dict: - """Return the extra ``sample_program`` arguments: none without sharing. - - Args: - postselection_mask: Postselected detectors, whose lanes are pruned. - detector_reference: Reference bits XOR-ed into the detectors before - the postselection test, if any. - - """ - if not self._trajectory_sharing: - return {} - num_detectors = self._num_detectors - mask = np.zeros(num_detectors, dtype=np.bool_) - if postselection_mask is not None: - mask = np.asarray(postselection_mask, dtype=np.bool_) - flip = np.zeros(num_detectors, dtype=np.bool_) - if detector_reference is not None: - flip = np.asarray(detector_reference[:num_detectors], dtype=np.bool_) - key = mask.tobytes() + flip.tobytes() - if key not in self._sharing: - self._sharing[key] = tuple( - ( - Sharing( - sharing_layout( - component, - np.array( - [ - i < num_detectors and bool(mask[i]) - for i in component.output_indices - ] - ), - np.array( - [ - i < num_detectors and bool(flip[i]) - for i in component.output_indices - ] - ), - ) - ) - if isinstance(component, StatevectorComponent) - else None - ) - for component in self._program.components - ) - return {"sharing": self._sharing[key], "sharing_stats": self.sharing_stats} - - def _compute_direct_outputs(self, f_params_np: np.ndarray) -> np.ndarray: - """Scatter direct output bits into a full (batch, num_outputs) bool array. - - Non-direct columns are zero. The zero-copy fast path applies when - direct indices are 0..n-1, there are no flips, and no reindex — - i.e. the common surface-code case. - """ - batch = f_params_np.shape[0] - num_outputs = self._program.num_outputs - n_direct = len(self._direct_f_indices) - if n_direct == 0: - return np.zeros((batch, num_outputs), dtype=np.bool_) - if self._direct_zero_copy and n_direct == num_outputs: - return f_params_np[:, :n_direct].view(np.bool_).copy() - raw = ( - f_params_np[:, :n_direct].view(np.bool_) - if self._direct_zero_copy - else (f_params_np[:, self._direct_f_indices] ^ self._direct_flips).view( - np.bool_ - ) - ) - out = np.zeros((batch, num_outputs), dtype=np.bool_) - out[:, self._direct_global_indices] = raw - return out - def _compute_reference_sample(self) -> np.ndarray: """Return the noiseless reference sample (all f_params = 0). Does not advance the channel sampler RNG. """ + program = self._program + if not program.components: + reference = np.zeros(program.num_outputs, dtype=np.bool_) + reference[program.output_order[: len(program.direct_f_indices)]] = ( + program.direct_flips + ) + return reference num_f = self._channel_sampler.signature_matrix.shape[1] f_ref = np.zeros((1, num_f), dtype=np.uint8) - if not self._program.components: - return self._compute_direct_outputs(f_ref)[0] - self._key, subkey = jax.random.split(self._key) - return np.asarray( - sample_program(self._program, jnp.asarray(f_ref), subkey)[0], - dtype=np.bool_, - ) + samples = self._backend.sample(f_ref, self._backend.next_seed()) + return np.asarray(samples[0], dtype=np.bool_) def _resolve_batch_size( self, @@ -390,57 +620,61 @@ def _resolve_batch_size( *, compute_reference: bool, ) -> int: - """Choose a uniform JAX batch size for ``shots`` samples.""" + """Choose a uniform batch size for ``shots`` samples. + + Without a ``batch_size``, batches are as large as memory allows, but split + into a batch per thread (``_threads``) when each gets at least + ``MIN_THREAD_BATCH`` shots, so that every thread has work. + """ if batch_size is None: - max_batch_size = self._estimate_batch_size() - num_batches = max(1, ceil(shots / max_batch_size)) - batch_size = ceil(shots / num_batches) + batch_size = self._auto_batch_size(shots) if compute_reference and batch_size * ceil(shots / batch_size) == shots: batch_size += 1 return batch_size + def _auto_batch_size(self, shots: int) -> int: + """Return the size of equal batches of ``shots``, one or more per thread.""" + num_batches = max(1, ceil(shots / self._estimate_batch_size())) + return ceil(shots / max(num_batches, self._per_thread(shots))) + + def _per_thread(self, shots: int) -> int: + """Return the thread count if ``shots`` fill a batch per thread, else 1.""" + threads = self._threads(1) + return threads if shots >= threads * MIN_THREAD_BATCH else 1 + def _peak_bytes_per_sample(self) -> int: """Estimate peak device memory per sample from compiled program structure.""" num_f = self._channel_sampler.signature_matrix.shape[1] - peak = 0 - for component in self._program.components: - if isinstance(component, StatevectorComponent): - # NOTE: plus the f-params and the output records, on host and device. - lane = lane_bytes(component) + 2 * num_f + 2 * self._program.num_outputs - peak = max(peak, lane) - continue - for circuit in component.compiled_scalar_graphs: - G = circuit.num_graphs - max_a = circuit.node_phases.phases.shape[1] - max_b = circuit.phase_terms.coeffs.shape[1] - max_c = circuit.pi_products.psi_params.shape[1] - max_d = circuit.phase_pairs.alpha.shape[1] - max_e = circuit.float_terms.table.shape[1] - largest = max(max_a * 16, max_b * 4, max_c * 4, max_d * 16, max_e * 16) - peak = max(peak, G * largest * 3) - return max(peak, 1) + return max(self._backend.peak_bytes_per_shot(num_f), 1) def _memory_budget(self) -> int: """Return the memory one batch may use: half of the free device memory.""" - device = jax.devices()[0] - if device.platform == "gpu": - stats = device.memory_stats() - available = stats.get("bytes_limit", 8 * 1024**3) - stats.get( - "bytes_in_use", 0 - ) - else: - available = psutil.virtual_memory().available - return int(available * 0.5) # conservative + return int(self._backend.memory_available() * 0.5) # conservative + + def _threads(self, batch_size: int) -> int: + """Return how many batches of ``batch_size`` shots run at once. + + That is the backend's limit, as long as the batches fit in the memory + budget. + """ + # NOTE: with sharing a batch may still hold a lane per shot. + fits = self._memory_budget() // (self._peak_bytes_per_sample() * batch_size) + return max(1, min(self._backend.max_threads(), fits)) def _estimate_batch_size(self) -> int: """Estimate the largest batch size that fits in available device memory. - Programs with statevector components use at most - ``STATEVECTOR_MAX_BATCH_SIZE`` lanes per batch. + Without sharing, every batch that runs at once holds a lane per shot. With + sharing, lanes are allocated as trajectories split, so the estimate + leaves room for a lane per shot in one batch only. The backend may cap + the batch size further. """ - batch_size = max(1, self._memory_budget() // self._peak_bytes_per_sample()) - if any(isinstance(c, StatevectorComponent) for c in self._program.components): - batch_size = min(batch_size, STATEVECTOR_MAX_BATCH_SIZE) + peak = self._peak_bytes_per_sample() + if not self._trajectory_sharing: + peak *= self._threads(1) + batch_size = max(1, self._memory_budget() // peak) + if self._backend.max_batch_size is not None: + batch_size = min(batch_size, self._backend.max_batch_size) return batch_size @overload @@ -480,11 +714,7 @@ def _sample_batches( Samples array, or (samples, reference) tuple when compute_reference=True. """ - if shots < 0: - raise ValueError(f"shots must be non-negative, got {shots}") - if batch_size is not None and batch_size < 1: - raise ValueError(f"batch_size must be at least 1, got {batch_size}") - + _check_sizes(shots, batch_size) if shots == 0: empty = np.empty((0, self._program.num_outputs), dtype=np.bool_) if compute_reference: @@ -498,47 +728,32 @@ def _sample_batches( return samples, reference return samples - if batch_size is None: - max_batch_size = self._estimate_batch_size() - num_batches = max(1, ceil(shots / max_batch_size)) - batch_size = ceil(shots / num_batches) - else: - num_batches = ceil(shots / batch_size) + # NOTE: with a reference sample the batch size grows by one when the + # batches would hold exactly ``shots``, so that the first batch's + # reference sample fits without another batch shape for JIT. + batch_size = self._resolve_batch_size( + shots, batch_size, compute_reference=compute_reference + ) + num_batches = ceil((shots + compute_reference) / batch_size) + backend = self._backend + tasks = [ + (index, self._channel_sampler.block_seeds(batch_size), backend.next_seed()) + for index in range(num_batches) + ] - if compute_reference and batch_size * num_batches == shots: - # Bump batch_size so the first batch's reference sample fits - # within existing batches (keeps shapes uniform for JIT). - batch_size += 1 + def sample_batch(task: tuple) -> jax.Array | np.ndarray: + index, block_seeds, seed = task + f_params = self._channel_sampler.sample(batch_size, block_seeds) + if compute_reference and index == 0: + f_params[0] = 0 + return backend.sample(f_params, seed) - batches: list[jax.Array] = [] + batches = run_tasks(sample_batch, tasks, self._threads(batch_size)) reference: np.ndarray | None = None - - for _ in range(num_batches): - f_params_np = self._channel_sampler.sample(batch_size) - - if compute_reference and reference is None: - f_params_np[0] = 0 - - f_params = jnp.asarray(f_params_np) - self._key, subkey = jax.random.split(self._key) - samples = sample_program( - self._program, f_params, subkey, **self._sharing_arguments(None) - ) - - if compute_reference and reference is None: - reference = np.asarray(samples[0]) - samples = samples[1:] - - batches.append(samples) - - # Concatenate on device, then a single d2h. The prior - # np.concatenate(batches) triggered per-batch __array__ (one d2h each) - # plus a host-side memcpy into a fresh numpy buffer for the concat - # output. For big bool tensors (e.g. 500k shots × 528 detector bits) - # the host memcpy alone was ~1 s on top of the PCIe transfer. - combined = batches[0] if len(batches) == 1 else jnp.concatenate(batches, axis=0) - jax.block_until_ready(combined) - result = copy_d2h(combined)[:shots] + if compute_reference: + reference = np.asarray(batches[0][0]) + batches[0] = batches[0][1:] + result = backend.concatenate(batches)[:shots] if compute_reference: assert reference is not None @@ -551,138 +766,143 @@ def _sample_batches_with_postselection( batch_size: int | None, *, postselection_mask: np.ndarray, - compute_reference: bool = False, - xor_detector_ref: bool = False, - ) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]: - """Sample with postselection, skipping JAX for direct discarded shots. - - Shots discarded by a direct masked detector are filled with their - direct detector columns and ``False`` elsewhere; JAX is never called for - those shots. Survivors are buffered until a full batch of - ``batch_size`` is ready, then dispatched to ``sample_program`` in one - call. The final partial batch is padded to keep the JAX batch size - fixed (avoiding recompilation) and the padding rows are discarded. - """ - if shots < 0: - raise ValueError(f"shots must be non-negative, got {shots}") - if batch_size is not None and batch_size < 1: - raise ValueError(f"batch_size must be at least 1, got {batch_size}") + referenced: np.ndarray | None = None, + ) -> np.ndarray: + """Sample with postselection, skipping the components for direct discards. - num_outputs = self._program.num_outputs - if shots == 0: - empty = np.empty((0, num_outputs), dtype=np.bool_) - empty_discarded = np.empty(0, dtype=np.bool_) - if compute_reference: - return empty, np.zeros(num_outputs, dtype=np.bool_), empty_discarded - return empty, None, empty_discarded + A shot that a masked direct detector discards keeps its direct detector + columns and reads ``False`` elsewhere. If the backend prunes, a shot that + a masked component detector discards also reads ``False`` in the outputs + set after that. The other shots are sampled in batches of ``batch_size``. + Without a ``batch_size``, batches are as large as memory allows, and the + last survivors are split into a batch per thread. - postselect_direct = postselection_mask & self._direct_detector_mask + Args: + shots: Number of samples to draw. + batch_size: Survivors per batch. Auto-determined if None. + postselection_mask: The postselected detectors. + referenced: The outputs that a noiseless reference sample is XOR-ed + into, before the postselection test, if any. - if not self._program.components: - samples = self._sample_direct(shots) - if compute_reference: - reference = self._compute_reference_sample() - if xor_detector_ref: - samples[:, : self._num_detectors] ^= reference[ - : self._num_detectors - ] - return samples, reference, np.zeros(shots, dtype=np.bool_) - return samples, None, np.zeros(shots, dtype=np.bool_) + """ + _check_sizes(shots, batch_size) + num_outputs = self._program.num_outputs + if shots == 0: + return np.empty((0, num_outputs), dtype=np.bool_) + split_last = batch_size is None if batch_size is None: - batch_size = self._resolve_batch_size( - shots, batch_size, compute_reference=False + batch_size = self._estimate_batch_size() + flip = None + if referenced is not None: + flip = self._compute_reference_sample() & referenced + f_indices, flips, columns = self._direct_detectors + if flip is not None: + flips = flips ^ flip[columns] + postselect = postselection_mask[columns] + result = np.empty((shots, num_outputs), dtype=np.bool_) + backend = self._backend + chunk = min(batch_size, SCREEN_SHOTS) + + def screen(task: tuple[int, np.ndarray]) -> tuple[np.ndarray, np.ndarray]: + """Write the direct detectors of the chunk at ``start``; return its survivors.""" + start, block_seeds = task + stop = min(start + chunk, shots) + f_params = self._channel_sampler.sample(stop - start, block_seeds) + discarded = _write_direct_detectors( + (stop - start) * len(columns), + f_params, + f_indices, + flips, + columns, + postselect, + result[start:stop], ) + survivors = np.flatnonzero(~discarded) + return f_params[survivors], start + survivors - reference: np.ndarray | None = None - if compute_reference: - reference = self._compute_reference_sample() - - result = np.zeros((shots, num_outputs), dtype=np.bool_) - was_discarded = np.zeros(shots, dtype=np.bool_) - survivor_f_buf: list[np.ndarray] = [] - survivor_idx_buf: list[int] = [] - shot_idx = 0 - - sharing = self._sharing_arguments( - postselection_mask, reference if xor_detector_ref else None - ) - - def _dispatch(f_batch: np.ndarray, indices: list[int], n_valid: int) -> None: - self._key, subkey = jax.random.split(self._key) - jax_out = np.asarray( - sample_program(self._program, jnp.asarray(f_batch), subkey, **sharing) + def dispatch(task: tuple) -> None: + f_batch, shot_batch, seed, rows = task + samples = backend.sample( + f_batch, seed, mask=postselection_mask, flip=flip, pad_to=rows ) - result[indices[:n_valid]] = jax_out[:n_valid] - - def _flush(*, final: bool = False) -> None: - nonlocal survivor_f_buf, survivor_idx_buf - while len(survivor_f_buf) >= batch_size: - _dispatch( - np.stack(survivor_f_buf[:batch_size]), - survivor_idx_buf[:batch_size], - batch_size, + result[shot_batch] = np.asarray(samples) + + threads = self._threads(batch_size) + num_f = self._channel_sampler.signature_matrix.shape[1] + pending_f = [np.empty((0, num_f), dtype=np.uint8)] + pending_shot = [np.empty(0, dtype=np.int64)] + pending = 0 + starts = range(0, shots, chunk) + screened_chunks = 0 + # NOTE: chunks are screened until every thread has a full batch of + # survivors. A batch starts every batch_size survivors, and all seeds + # are drawn in order, so the thread count doesn't change the samples. + while screened_chunks < len(starts) or pending: + while screened_chunks < len(starts) and pending < threads * batch_size: + tasks = [ + ( + start, + self._channel_sampler.block_seeds(min(chunk, shots - start)), + ) + for start in starts[ + screened_chunks : screened_chunks + + SCREEN_CHUNKS_PER_THREAD * threads + ] + ] + screened_chunks += len(tasks) + for f_survivors, shot_survivors in run_tasks(screen, tasks, threads): + pending_f.append(f_survivors) + pending_shot.append(shot_survivors) + pending += len(shot_survivors) + ready = pending + size, rows = batch_size, batch_size + if screened_chunks < len(starts): + ready -= ready % (threads * batch_size) + elif split_last and ready: + # NOTE: JAX pads these batches to a power of two, so that few + # batch shapes are compiled. + per_thread = self._per_thread(ready) + size = ceil(ready / max(1, ceil(ready / batch_size), per_thread)) + rows = min(batch_size, 1 << (size - 1).bit_length()) + all_f, all_shot = np.concatenate(pending_f), np.concatenate(pending_shot) + batches = [ + ( + all_f[start : start + size], + all_shot[start : start + size], + backend.next_seed(), + rows, ) - survivor_f_buf = survivor_f_buf[batch_size:] - survivor_idx_buf = survivor_idx_buf[batch_size:] - - if final and survivor_f_buf: - n_valid = len(survivor_f_buf) - f_stack = np.stack(survivor_f_buf) - f_batch = np.empty((batch_size, f_stack.shape[1]), dtype=f_stack.dtype) - f_batch[:n_valid] = f_stack - f_batch[n_valid:] = f_stack[0] - _dispatch(f_batch, survivor_idx_buf, n_valid) - survivor_f_buf = [] - survivor_idx_buf = [] - - while shot_idx < shots: - chunk = min(batch_size, shots - shot_idx) - f_params_np = self._channel_sampler.sample(chunk) - direct_full = self._compute_direct_outputs(f_params_np) - det_cols = direct_full[:, : self._num_detectors] - if xor_detector_ref and reference is not None: - det_cols = det_cols ^ reference[: self._num_detectors] - - discarded = (det_cols & postselect_direct).any(axis=1) - - result[shot_idx : shot_idx + chunk, : self._num_detectors] = direct_full[ - :, : self._num_detectors + for start in range(0, ready, size) ] - was_discarded[shot_idx : shot_idx + chunk] = discarded - - survivor_local = np.flatnonzero(~discarded) - if survivor_local.size: - survivor_f_buf.extend(f_params_np[survivor_local]) - survivor_idx_buf.extend((shot_idx + survivor_local).tolist()) - - shot_idx += chunk - _flush() - - _flush(final=True) - - if xor_detector_ref and reference is not None: - det_ref = reference[: self._num_detectors] - survivors = ~was_discarded - result[survivors, : self._num_detectors] ^= det_ref - result[was_discarded, : self._num_detectors] ^= ( - det_ref & self._direct_detector_mask - ) - - if compute_reference: - assert reference is not None - return result, reference, was_discarded - return result, None, was_discarded + run_tasks(dispatch, batches, threads) + pending_f, pending_shot = [all_f[ready:]], [all_shot[ready:]] + pending -= ready + return result def _sample_direct(self, shots: int) -> np.ndarray: - """Fast path when all components are direct (pure numpy, no JAX).""" + """Sample a program without components, whose outputs are all direct.""" + program = self._program f_params = self._channel_sampler.sample(shots) - if self._direct_zero_copy: - return f_params[:, : len(self._direct_f_indices)].view(np.bool_) - result = f_params[:, self._direct_f_indices] ^ self._direct_flips - if self._direct_reindex is not None: - result = result[:, self._direct_reindex] - return result.view(np.bool_) + num_direct = len(program.direct_f_indices) + # NOTE: typical surface-code detector circuits need no copy. + if ( + program.output_reindex is None + and not program.direct_flips.any() + and np.array_equal(program.direct_f_indices, np.arange(num_direct)) + ): + return f_params[:, :num_direct].view(np.bool_) + out = np.empty((shots, program.num_outputs), dtype=np.bool_) + _gather_outputs( + out.size, + f_params, + program.direct_f_indices, + program.direct_flips, + np.zeros((shots, 0), dtype=np.bool_), + program.output_order, + out, + ) + return out def __repr__(self) -> str: """Return a string representation with compilation statistics.""" @@ -775,7 +995,7 @@ def __init__( *, strategy: SamplingStrategy = "active_statevector", seed: int | None = None, - trajectory_sharing: bool = False, + trajectory_sharing: bool | None = None, ): """Create a measurement sampler. @@ -784,7 +1004,8 @@ def __init__( strategy: "active_statevector" (the default), or a stabilizer rank decomposition strategy ("cat5", "bss", "cutting"). trajectory_sharing: With ``strategy="active_statevector"``, simulate shots - with equal histories on one lane (experimental). + with equal histories on one lane. Defaults to True for + ``"active_statevector"`` and False otherwise. seed: Random seed for the sampler. IMPORTANT: Currently, the sampler will only produce deterministic samples for fixed batch size. If deterministic samples are needed, the batch size should be set @@ -807,9 +1028,12 @@ def sample(self, shots: int, *, batch_size: int | None = None) -> np.ndarray: shots: The number of times to sample every measurement in the circuit. batch_size: The number of samples to process in each batch. Defaults to None, which automatically chooses a batch size based on available - memory. When using a GPU, setting this explicitly can help fully - utilize VRAM for maximum performance. NOTE: Changing the batch size - will affect reproducibility even with a fixed seed. + memory. With ``strategy="active_statevector"`` a call is split + into a batch per numba thread when each gets at least 4096 shots. + With a stabilizer-rank strategy on a GPU, setting this explicitly + can help fully utilize VRAM. + NOTE: Changing the batch size will affect reproducibility even with + a fixed seed. Returns: A numpy array containing the measurement samples. @@ -834,7 +1058,7 @@ def __init__( *, strategy: SamplingStrategy = "active_statevector", seed: int | None = None, - trajectory_sharing: bool = False, + trajectory_sharing: bool | None = None, ): """Create a detector sampler. @@ -843,7 +1067,8 @@ def __init__( strategy: "active_statevector" (the default), or a stabilizer rank decomposition strategy ("cat5", "bss", "cutting"). trajectory_sharing: With ``strategy="active_statevector"``, simulate shots - with equal histories on one lane (experimental). + with equal histories on one lane. Defaults to True for + ``"active_statevector"`` and False otherwise. seed: Random seed for the sampler. IMPORTANT: Currently, the sampler will only produce deterministic samples for fixed batch size and fixed reference sample settings. If deterministic samples are @@ -912,9 +1137,12 @@ def sample( shots: The number of times to sample every detector in the circuit. batch_size: The number of samples to process in each batch. Defaults to None, which automatically chooses a batch size based on available - memory. When using a GPU, setting this explicitly can help fully - utilize VRAM for maximum performance. NOTE: Changing the batch size - will affect reproducibility even with a fixed seed. + memory. With ``strategy="active_statevector"`` a call is split + into a batch per numba thread when each gets at least 4096 shots. + With a stabilizer-rank strategy on a GPU, setting this explicitly + can help fully utilize VRAM. + NOTE: Changing the batch size will affect reproducibility even with + a fixed seed. separate_observables: Defaults to False. When set to True, the return value is a (detection_events, observable_flips) tuple instead of a flat detection_events array. @@ -934,11 +1162,14 @@ def sample( only be used when observables are deterministic. Otherwise, it can unpredictably change the results. postselection_mask: Optional boolean array of length ``num_detectors``. - When set, shots where any masked direct detector fires skip the JAX - sampling loop. All ``shots`` rows are still returned: survivors contain - the full sample, while discarded rows retain direct detector columns - and fill component columns with ``False``. Re-apply the mask to the - returned detector columns to recover surviving shots. + When set, shots where any masked direct detector fires are not + sampled further. All ``shots`` rows are still returned: survivors + contain the full sample, while discarded rows retain direct detector + columns and fill component columns with ``False``; with + ``strategy="active_statevector"``, a shot discarded by a masked + component detector also reads ``False`` in the outputs set after + it. Re-apply the mask to the returned detector columns to recover + surviving shots. Returns: A numpy array or tuple of numpy arrays containing the samples. @@ -957,49 +1188,43 @@ def sample( compute_reference = ( use_detector_reference_sample or use_observable_reference_sample ) + num_detectors = self._num_detectors if postselection_mask is not None: mask = np.asarray(postselection_mask, dtype=np.bool_) - if mask.shape != (self._num_detectors,): + if mask.shape != (num_detectors,): raise ValueError( - f"postselection_mask must have shape ({self._num_detectors},), " + f"postselection_mask must have shape ({num_detectors},), " f"got {mask.shape}" ) - if postselection_mask is not mask: - postselection_mask = mask - if ( - not (postselection_mask & self._direct_detector_mask).any() - or not self._program.components - ): - postselection_mask = None + # NOTE: without pruning, only direct detectors can skip any work. + if not self._backend.prunes: + skips = mask & self._direct_detector_mask + else: + skips = mask + postselection_mask = ( + mask if skips.any() and self._program.components else None + ) if postselection_mask is not None: + referenced = None if compute_reference: - samples, reference, direct_discarded = ( - self._sample_batches_with_postselection( - shots, - batch_size, - postselection_mask=postselection_mask, - compute_reference=True, - xor_detector_ref=use_detector_reference_sample, - ) - ) - assert reference is not None - num_detectors = self._num_detectors - if use_observable_reference_sample: - obs_ref = reference[num_detectors:] - samples[~direct_discarded, num_detectors:] ^= obs_ref - else: - samples, _, _ = self._sample_batches_with_postselection( - shots, - batch_size, - postselection_mask=postselection_mask, + is_detector = np.arange(self._program.num_outputs) < num_detectors + referenced = np.where( + is_detector, + use_detector_reference_sample, + use_observable_reference_sample, ) + samples = self._sample_batches_with_postselection( + shots, + batch_size, + postselection_mask=postselection_mask, + referenced=referenced, + ) elif compute_reference: samples, reference = self._sample_batches( shots, batch_size, compute_reference=True ) - num_detectors = self._num_detectors if use_detector_reference_sample: samples[:, :num_detectors] ^= reference[:num_detectors] if use_observable_reference_sample: @@ -1007,7 +1232,6 @@ def sample( else: samples = self._sample_batches(shots, batch_size) - num_detectors = self._num_detectors det_samples = samples[:, :num_detectors] obs_samples = samples[:, num_detectors:] diff --git a/src/tsim/statevector/__init__.py b/src/tsim/statevector/__init__.py index 98638f20..fa884400 100644 --- a/src/tsim/statevector/__init__.py +++ b/src/tsim/statevector/__init__.py @@ -3,9 +3,9 @@ A non-direct component of the reduced sampling graph is read as rotations and measurements in the frame of the initial ``|0…0⟩`` state (``pbc``), reordered to keep the dense array small (``schedule``), planned on a fixed -number of array axes (``plan``), lowered to JAX arrays (``program``) and run -lane by lane (``execute``). As in Clifft [1], the dense array only spans the -active qubits. +number of array axes (``plan``), lowered to NumPy arrays (``program``) and run +by numba kernels on the CPU, with shots that share a history on one lane +(``shared``). As in Clifft [1], the dense array only spans the active qubits. References: [1] B. A. Chase, F. Labib, "Clifft: Fast Exact Simulation of Near-Clifford diff --git a/src/tsim/statevector/execute.py b/src/tsim/statevector/execute.py deleted file mode 100644 index e24e02b6..00000000 --- a/src/tsim/statevector/execute.py +++ /dev/null @@ -1,302 +0,0 @@ -"""JAX executor for statevector components. - -``run`` interprets a component's action tape with one ``lax.fori_loop``; each -step dispatches on the action's kind with ``lax.switch``. The loop carry is a -``LaneState``: one lane per shot, each with ``2^W`` amplitudes and its symbols -(f-bits, then branch bits as they are sampled), from which ``outputs`` reads -the output bits once the tape has run. All randomness goes through -``choose_outcome``, and ``count`` holds the number of shots a lane stands for, -so lanes can be shared between shots without changing the kernels. -""" - -from __future__ import annotations - -from functools import partial -from typing import NamedTuple - -import jax -import jax.numpy as jnp -from jax import Array - -from tsim.statevector.program import ( - MEASURE, - MEASURE_DIAG, - MEASURE_DORMANT, - ROTATE, - ROTATE_DIAG, - StatevectorComponent, -) - -_I_POWERS = (1, 1j, -1, -1j) -_TINY = 1e-30 -_TWO_TO_32 = 2.0**32 - - -class LaneState(NamedTuple): - """Per-lane state of the executor; the leading axis is the lane. - - Attributes: - amps: Amplitudes, shape ``(L, 2^W)``. - symbols: f-bits, then branch bits, packed into uint32 words, - shape ``(L, num_words)``. - count: Shots each lane stands for, shape ``(L,)``. - max_norm_deviation: Largest ``|‖ψ‖² - 1|`` seen at a measurement. - - """ - - amps: Array - symbols: Array - count: Array # pyright: ignore[reportIncompatibleMethodOverride] - max_norm_deviation: Array - - -def lane_bytes(component: StatevectorComponent) -> int: - """Return the peak device memory of one lane, temporary amplitudes included.""" - # NOTE: XLA's memory analysis of the executor shows four amplitude arrays. - amplitudes = 4 * 2**component.width * jnp.dtype(component.dtype).itemsize - return amplitudes + 4 * component.num_words + component.num_outputs - - -def choose_outcome(p0: Array, p1: Array, count: Array, key: Array) -> Array: - """Return each lane's outcome: True with probability ``p1 = 1 - p0``. - - The less likely outcome is drawn, so that both outcomes keep the relative - precision of their probability: in float32, ``p1`` alone rounds to 1 once - ``p0`` is below about 6·10⁻⁸. Two uint32 draws per lane form a uniform 64-bit - fraction, compared with ``p·2^32`` (an exact scaling) for the smaller ``p``, - so the resolution is 2^-64 for any float dtype, and an outcome of - probability 0 is never drawn. ``count`` is unused while every lane is one - shot. - """ - del count - high, low = jax.random.bits(key, (2, *p1.shape), dtype=jnp.uint32) - one_is_rare = p1 <= p0 - scaled = jnp.clip(jnp.minimum(p0, p1), 0, 0.5) * _TWO_TO_32 - whole = jnp.floor(scaled) - high_bound = whole.astype(jnp.uint32) - low_bound = jnp.floor((scaled - whole) * _TWO_TO_32).astype(jnp.uint32) - rare = (high < high_bound) | ((high == high_bound) & (low < low_bound)) - return rare == one_is_rare - - -def pack_bits(bits: Array, num_words: int) -> Array: - """Pack boolean columns into uint32 words: column s is bit s % 32 of word s // 32.""" - lanes, num_bits = bits.shape - padded = jnp.zeros((lanes, num_words * 32), dtype=jnp.uint32) - padded = padded.at[:, :num_bits].set(bits.astype(jnp.uint32)) - shifts = jnp.arange(32, dtype=jnp.uint32) - words = padded.reshape(lanes, num_words, 32) << shifts - return jnp.sum(words, axis=-1, dtype=jnp.uint32) - - -def initial_state(component: StatevectorComponent, f_bits: Array) -> LaneState: - """Return lanes in ``|0…0⟩`` with the given f-bits, shape ``(L, num_f)``.""" - lanes = f_bits.shape[0] - dtype = jnp.dtype(component.dtype) - amps = jnp.zeros((lanes, 2**component.width), dtype=dtype).at[:, 0].set(1) - return LaneState( - amps=amps, - symbols=pack_bits(f_bits, component.num_words), - count=jnp.ones(lanes, dtype=jnp.int32), - max_norm_deviation=jnp.zeros((), dtype=jnp.finfo(dtype).dtype), - ) - - -def _parity(values: Array) -> Array: - return (jax.lax.population_count(values) & 1).astype(jnp.bool_) - - -def _affine(symbols: Array, rows: Array, constant: Array) -> Array: - """Return ``constant ⊕ parity(symbols ∧ rows)`` over the last (word) axis.""" - ones = jnp.sum(jax.lax.population_count(symbols & rows), axis=-1) - return ((ones & 1) == 1) ^ constant - - -def outputs(component: StatevectorComponent, state: LaneState) -> Array: - """Return each lane's output bits, shape ``(L, num_outputs)``.""" - return _affine( - state.symbols[:, None, :], component.output_row, component.output_const - ) - - -def _index(component: StatevectorComponent) -> Array: - return jnp.arange(2**component.width, dtype=jnp.uint32) - - -def _sign(component: StatevectorComponent, state: LaneState, a: Array) -> Array: - """Return each lane's sign bit of rotation ``a``.""" - return _affine(state.symbols, component.sign_row[a], component.sign_const[a]) - - -def _apply_pauli(amps: Array, index: Array, x: Array, z: Array) -> Array: - """Return P·ψ lane by lane, for P = i^|x∧z| X^x Z^z.""" - flipped = index ^ x - quarter_turns = (jax.lax.population_count(x & z) & 3).astype(jnp.int32) - i_power = jnp.asarray(_I_POWERS, dtype=amps.dtype)[quarter_turns] - signs = jnp.where(_parity(flipped & z), -1, 1).astype(amps.dtype) - return i_power * signs * jnp.take(amps, flipped.astype(jnp.int32), axis=1) - - -def _rotate_diag( - component: StatevectorComponent, state: LaneState, a: Array, key: Array -) -> LaneState: - del key - odd = _parity(_index(component) & component.z_mask[a]) - flip = _sign(component, state, a)[:, None] ^ odd[None, :] - factor = component.cos[a] - 1j * component.sin[a] * jnp.where(flip, -1, 1) - return state._replace(amps=(state.amps * factor).astype(state.amps.dtype)) - - -def _rotate( - component: StatevectorComponent, state: LaneState, a: Array, key: Array -) -> LaneState: - del key - paulied = _apply_pauli( - state.amps, _index(component), component.x_mask[a], component.z_mask[a] - ) - negate = jnp.where(_sign(component, state, a), -1, 1)[:, None] - amps = component.cos[a] * state.amps - 1j * component.sin[a] * negate * paulied - return state._replace(amps=amps.astype(state.amps.dtype)) - - -def _squared(values: Array) -> Array: - return jnp.real(values) ** 2 + jnp.imag(values) ** 2 - - -def _set_branch( - component: StatevectorComponent, state: LaneState, a: Array, outcome: Array -) -> LaneState: - """Set the measurement's branch symbol to ``outcome``.""" - branch = component.branch[a] - word = branch // 32 - bit = jnp.uint32(1) << (branch % 32).astype(jnp.uint32) - flags = jnp.where(outcome, bit, jnp.uint32(0)) - symbols = state.symbols.at[:, word].set(state.symbols[:, word] | flags) - return state._replace(symbols=symbols) - - -def _sample_outcome( - state: LaneState, a: Array, key: Array, weight_0: Array, weight_1: Array -) -> tuple[Array, Array]: - """Return each lane's outcome and its weight. - - Each outcome's probability is its own weight over the total, so neither is - computed as one minus the other. - """ - total = jnp.maximum(weight_0 + weight_1, _TINY) - key = jax.random.fold_in(key, a) - outcome = choose_outcome(weight_0 / total, weight_1 / total, state.count, key) - return outcome, jnp.where(outcome, weight_1, weight_0) - - -def _collapse( - component: StatevectorComponent, - state: LaneState, - a: Array, - outcome: Array, - projected: Array, - norm: Array, -) -> LaneState: - """Keep the pivot-0 half of the normalized ``projected``; branch on ``outcome``.""" - keep = ((_index(component) >> component.pivot[a].astype(jnp.uint32)) & 1) == 0 - deviation = jnp.maximum(state.max_norm_deviation, jnp.max(jnp.abs(norm - 1))) - state = state._replace( - amps=jnp.where(keep[None, :], projected, 0).astype(state.amps.dtype), - max_norm_deviation=deviation.astype(state.max_norm_deviation.dtype), - ) - return _set_branch(component, state, a, outcome) - - -def _measure( - component: StatevectorComponent, state: LaneState, a: Array, key: Array -) -> LaneState: - amps = state.amps - paulied = _apply_pauli( - amps, _index(component), component.x_mask[a], component.z_mask[a] - ) - plus, minus = amps + paulied, amps - paulied - weight_0 = jnp.sum(_squared(plus), axis=1) / 4 - weight_1 = jnp.sum(_squared(minus), axis=1) / 4 - outcome, weight = _sample_outcome(state, a, key, weight_0, weight_1) - scale = jax.lax.rsqrt(jnp.maximum(2 * weight, _TINY))[:, None] - projected = jnp.where(outcome[:, None], minus, plus) * scale - return _collapse(component, state, a, outcome, projected, weight_0 + weight_1) - - -def _measure_diag( - component: StatevectorComponent, state: LaneState, a: Array, key: Array -) -> LaneState: - index = _index(component) - amps = state.amps - odd = _parity(index & component.z_mask[a]) - weights = _squared(amps) - weight_0 = jnp.sum(jnp.where(odd[None, :], 0, weights), axis=1) - weight_1 = jnp.sum(jnp.where(odd[None, :], weights, 0), axis=1) - outcome, weight = _sample_outcome(state, a, key, weight_0, weight_1) - bit = jnp.uint32(1) << component.pivot[a].astype(jnp.uint32) - partner = jnp.take(amps, (index ^ bit).astype(jnp.int32), axis=1) - gathered = jnp.where(odd[None, :] == outcome[:, None], amps, partner) - scale = jax.lax.rsqrt(jnp.maximum(weight, _TINY))[:, None] - projected = gathered * scale - return _collapse(component, state, a, outcome, projected, weight_0 + weight_1) - - -def _measure_dormant( - component: StatevectorComponent, state: LaneState, a: Array, key: Array -) -> LaneState: - half = jnp.full(state.count.shape, 0.5, dtype=state.max_norm_deviation.dtype) - outcome = choose_outcome(half, half, state.count, jax.random.fold_in(key, a)) - return _set_branch(component, state, a, outcome) - - -_KERNELS = { - ROTATE_DIAG: _rotate_diag, - ROTATE: _rotate, - MEASURE_DIAG: _measure_diag, - MEASURE: _measure, - MEASURE_DORMANT: _measure_dormant, -} - - -def run( - component: StatevectorComponent, - state: LaneState, - key: Array, - start: int | Array, - stop: int | Array, -) -> LaneState: - """Run actions ``start … stop-1`` of the tape on every lane. - - A component without actions returns ``state`` unchanged. - """ - if component.num_actions == 0: - return state - branches = [partial(_KERNELS[kind], component) for kind in range(len(_KERNELS))] - - def step(a: Array, carry: LaneState) -> LaneState: - return jax.lax.switch(component.kind[a], branches, carry, a, key) - - return jax.lax.fori_loop(start, stop, step, state) - - -@jax.jit -def sample_statevector_component( - component: StatevectorComponent, f_params: Array, key: Array -) -> tuple[Array, Array, Array]: - """Sample one component's outputs; the contract of ``sampler.sample_component``. - - Args: - component: The compiled component. - f_params: Error-basis bits of every shot, shape ``(B, num_f_global)``. - key: JAX random key. - - Returns: - ``(samples, next_key, max_norm_deviation)`` with ``samples`` of shape - ``(B, num_outputs)``. - - """ - f_bits = f_params[:, component.f_selection].astype(jnp.bool_) - key, subkey = jax.random.split(key) - state = initial_state(component, f_bits) - state = run(component, state, subkey, 0, component.num_actions) - return outputs(component, state), key, state.max_norm_deviation diff --git a/src/tsim/statevector/plan.py b/src/tsim/statevector/plan.py index ae95a9e9..ec8d4174 100644 --- a/src/tsim/statevector/plan.py +++ b/src/tsim/statevector/plan.py @@ -215,11 +215,10 @@ def _promote(self, j: int, dormant_x: list[int]) -> None: raise RuntimeError( f"the plan needs more active axes than the schedule's width {self.width}" ) - in_array = [q for q in dormant_x if q < self.width] - if in_array: - axis = in_array[0] - else: - axis = free[0] + # NOTE: the lowest free axis keeps the active axes low, so the extent of + # the array that later actions touch tracks the active width. + axis = free[0] + if axis not in dormant_x: self._apply("SWAP", axis, dormant_x[0]) x, _, _ = self._virtual(j) for q in np.flatnonzero(x): diff --git a/src/tsim/statevector/program.py b/src/tsim/statevector/program.py index ed51a150..6aea2ce9 100644 --- a/src/tsim/statevector/program.py +++ b/src/tsim/statevector/program.py @@ -1,21 +1,19 @@ """Compilation of one component for ``strategy="active_statevector"``. ``plan_component`` reads the Pauli ops off the component, schedules and -plans them; ``lower`` turns the plan into the arrays of a ``StatevectorComponent``: -an action tape that ``execute.run`` interprets, and the outputs as parities of -the symbols the tape leaves behind. +plans them; ``lower`` turns the plan into the NumPy arrays of a +``StatevectorComponent``: an action tape that ``shared.sample_shared_component`` +runs, and the outputs as parities of the symbols the tape leaves behind. """ from __future__ import annotations from collections.abc import Iterable, Sequence +from dataclasses import dataclass from math import ceil +from typing import ClassVar, NamedTuple -import equinox as eqx -import jax -import jax.numpy as jnp import numpy as np -from jax import Array from zxparam import Graph from tsim.core.params import ParamLayout @@ -31,22 +29,15 @@ MAX_WIDTH = 30 -# NOTE: outcome probabilities are normalized by the measured weights, so only -# drift beyond rounding (~1e-7 per rotation in complex64) warns. -_NORM_WARNING_THRESHOLDS = {"complex64": 1e-3, "complex128": 1e-9} - STABILIZER_RANK_HINT = ( "Use a stabilizer-rank strategy such as strategy='cat5' for this circuit." ) -class StatevectorComponent(eqx.Module): - """A component compiled for ``strategy="active_statevector"``: an action tape. +class Tape(NamedTuple): + """A component's actions and outputs as NumPy arrays, one row per action. Attributes: - output_indices: Global output index of each output slot. - width: Number of array axes W; lanes hold ``2^W`` amplitudes. - dtype: ``"complex64"`` or ``"complex128"``. f_selection: Global f-bit index of each component f-bit. kind: Kernel of each action (``ROTATE_DIAG`` … ``MEASURE_DORMANT``). x_mask: X mask of the action's Pauli on the array axes. @@ -54,6 +45,8 @@ class StatevectorComponent(eqx.Module): cos: cos θ of each rotation. sin: sin θ of each rotation. pivot: Axis an active measurement resets to ``|0⟩``. + extent: One plus the highest axis the action can touch: amplitudes at + ``2^extent`` and above are zero before and after it. branch: Symbol a measurement sets, -1 for a rotation. sign_const: Constant of each rotation's sign. sign_row: Symbol bits of each rotation's sign, shape ``(A, num_words)``. @@ -62,42 +55,58 @@ class StatevectorComponent(eqx.Module): """ - output_indices: tuple[int, ...] = eqx.field(static=True) - width: int = eqx.field(static=True) - dtype: str = eqx.field(static=True) - f_selection: Array - kind: Array - x_mask: Array - z_mask: Array - cos: Array - sin: Array - pivot: Array - branch: Array - sign_const: Array - sign_row: Array - output_const: Array - output_row: Array + f_selection: np.ndarray + kind: np.ndarray + x_mask: np.ndarray + z_mask: np.ndarray + cos: np.ndarray + sin: np.ndarray + pivot: np.ndarray + extent: np.ndarray + branch: np.ndarray + sign_const: np.ndarray + sign_row: np.ndarray + output_const: np.ndarray + output_row: np.ndarray + + +@dataclass(frozen=True, eq=False) +class StatevectorComponent: + """A component compiled for ``strategy="active_statevector"``. + + Its lanes hold ``2^width`` complex128 amplitudes. + + Attributes: + output_indices: Global output index of each output slot. + width: Number of array axes W. + tape: The actions and outputs. + + """ + + output_indices: tuple[int, ...] + width: int + tape: Tape + + # NOTE: outcome probabilities are normalized by the measured weights, so + # only drift beyond rounding warns. + norm_warning_threshold: ClassVar[float] = 1e-9 + """Largest normalization deviation that sampling tolerates without a warning.""" @property def num_actions(self) -> int: """Return the number of actions on the tape.""" - return int(self.kind.shape[0]) + return int(self.tape.kind.shape[0]) @property def num_words(self) -> int: """Return the number of uint32 words that hold a lane's symbols.""" - return int(self.sign_row.shape[1]) + return int(self.tape.sign_row.shape[1]) @property def num_outputs(self) -> int: """Return the number of output slots.""" return len(self.output_indices) - @property - def norm_warning_threshold(self) -> float: - """Return the largest normalization deviation that passes without a warning.""" - return _NORM_WARNING_THRESHOLDS[self.dtype] - def plan_component( component: Graph, output_indices: Sequence[int], layout: ParamLayout @@ -136,7 +145,7 @@ def _pack(terms: Iterable[int], num_words: int) -> np.ndarray: return row -def lower(plan: Plan, layout: ComponentLayout, dtype: str) -> StatevectorComponent: +def lower(plan: Plan, layout: ComponentLayout) -> StatevectorComponent: """Return the action tape of ``plan``.""" num_actions = len(plan.actions) num_words = max(1, ceil(plan.num_symbols / 32)) @@ -145,15 +154,18 @@ def lower(plan: Plan, layout: ComponentLayout, dtype: str) -> StatevectorCompone z_mask = np.zeros(num_actions, dtype=np.uint32) angle = np.zeros(num_actions, dtype=np.float64) pivot = np.zeros(num_actions, dtype=np.int32) + extent = np.zeros(num_actions, dtype=np.int32) branch = np.full(num_actions, -1, dtype=np.int32) sign_const = np.zeros(num_actions, dtype=bool) sign_row = np.zeros((num_actions, num_words), dtype=np.uint32) + active = 0 for a, action in enumerate(plan.actions): if isinstance(action, Rotate): kind[a] = ROTATE if action.x else ROTATE_DIAG x_mask[a], z_mask[a], angle[a] = action.x, action.z, action.angle sign_const[a] = action.sign.constant sign_row[a] = _pack(action.sign.terms, num_words) + active |= action.x | action.z elif isinstance(action, MeasureActive): kind[a] = MEASURE if action.x else MEASURE_DIAG x_mask[a], z_mask[a] = action.x, action.z @@ -161,27 +173,28 @@ def lower(plan: Plan, layout: ComponentLayout, dtype: str) -> StatevectorCompone else: kind[a] = MEASURE_DORMANT branch[a] = action.branch + extent[a] = active.bit_length() + if isinstance(action, MeasureActive): + active &= ~(1 << action.pivot) output_row = np.zeros((len(plan.outputs), num_words), dtype=np.uint32) for slot, output in enumerate(plan.outputs): output_row[slot] = _pack(output.terms, num_words) - real = np.finfo(dtype).dtype - return StatevectorComponent( - output_indices=layout.output_indices, - width=plan.width, - dtype=dtype, - f_selection=jnp.asarray(np.array(layout.f_selection, dtype=np.int32)), - kind=jnp.asarray(kind), - x_mask=jnp.asarray(x_mask), - z_mask=jnp.asarray(z_mask), - cos=jnp.asarray(np.cos(angle).astype(real)), - sin=jnp.asarray(np.sin(angle).astype(real)), - pivot=jnp.asarray(pivot), - branch=jnp.asarray(branch), - sign_const=jnp.asarray(sign_const), - sign_row=jnp.asarray(sign_row), - output_const=jnp.asarray([output.constant for output in plan.outputs], bool), - output_row=jnp.asarray(output_row), + tape = Tape( + f_selection=np.array(layout.f_selection, dtype=np.int32), + kind=kind, + x_mask=x_mask, + z_mask=z_mask, + cos=np.cos(angle), + sin=np.sin(angle), + pivot=pivot, + extent=extent, + branch=branch, + sign_const=sign_const, + sign_row=sign_row, + output_const=np.array([output.constant for output in plan.outputs], bool), + output_row=output_row, ) + return StatevectorComponent(layout.output_indices, plan.width, tape) def compile_statevector_component( @@ -189,8 +202,6 @@ def compile_statevector_component( ) -> StatevectorComponent: """Compile one non-direct component for ``strategy="active_statevector"``. - Amplitudes are complex128 when JAX has 64-bit types enabled, else complex64. - Args: component: A connected component of the prepared sampling graph. output_indices: Global output index of each of the component's outputs. @@ -198,5 +209,4 @@ def compile_statevector_component( """ planned, component_layout = plan_component(component, output_indices, layout) - dtype = jax.dtypes.canonicalize_dtype(np.complex128).name - return lower(planned, component_layout, dtype) + return lower(planned, component_layout) diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index f40b1751..50ad6e91 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -1,4 +1,4 @@ -"""Trajectory-sharing executor for statevector components (investigation prototype). +"""Trajectory-sharing executor for statevector components, on the CPU with numba. Shots that agree on every symbol that reaches the amplitudes share one lane. The only per-shot input to the amplitudes is the sign of a rotation (outputs are @@ -6,43 +6,38 @@ every rotation sign. With postselection the key also holds the f-bit part of every postselected output, so that a lane is pruned as a whole as soon as one of its postselected outputs is known to be 1: at the action that sets the output's -last branch bit (the first measurement for an output of f-bits only). +last branch bit, or before the first action for an output of f-bits only. A lane holds amplitudes, the symbols of one of its shots (its f-bits agree with every member on the key, and its branch bits are the lane's outcomes), a shot count and a range in the key-sorted order of the shots. At a measurement the lane -draws how many of its shots see the rarer outcome (``_rare_drawn``, which -resolves the chance of at least one to 2^-64 like ``choose_outcome``) and -splits: the child takes the last n shots of the range. Shots in a lane agree on -the key, so they are exchangeable, and the sort breaks ties at random, so each -range is in uniformly random order: this assignment gives every shot -independent outputs. Pruned lanes move to a dead table that keeps their range, -outcomes and the pruning action but no amplitudes. A dormant measurement splits -lanes the same way (p = 1/2) when its branch reaches a key sign; otherwise it is -a fair coin drawn per shot at the end. - -Capacities are static. The dead table holds one entry per shot, so it never -fills. When the lane table is full the loop stops before the action, the host -doubles it and resumes; every lane slot draws from its own key, so the redone -action draws the same numbers. After the tape each shot's outputs are evaluated -from its own f-bits, its lane's outcomes and its own dormant coins; a pruned -shot's outputs that need a branch bit set after its pruning action are False. +draws how many of its shots see the rarer outcome (a binomial whose chance of a +first hit is resolved to 2^-64) and splits: the child takes the last n shots of +the range. Shots in a lane agree on the key, so they are exchangeable, and the +sort breaks ties at random, so each range is in uniformly random order: this +assignment gives every shot independent outputs. Pruned lanes keep their range, +outcomes and pruning action but drop their amplitudes. A dormant measurement +splits lanes the same way (p = 1/2) when its branch reaches a key sign; otherwise +it is a fair coin drawn per shot at the end. + +One numba function samples a batch without the GIL: it groups the shots, runs +the tape on the live lanes and the ``2^extent`` amplitudes each action can +touch, and evaluates the outputs. When the extent grows, the newly touched +amplitudes are cleared, because a slot reused for another lane can hold stale +values there. After the tape each shot's outputs are evaluated from its own +f-bits, its lane's outcomes and its own dormant coins, XOR its reference bit; a +pruned shot's outputs that need a branch bit set after its pruning action are +False. """ from __future__ import annotations -from dataclasses import dataclass -from functools import partial from typing import NamedTuple -import equinox as eqx -import jax -import jax.numpy as jnp +import numba import numpy as np -from jax import Array -from tsim.statevector import execute -from tsim.statevector.execute import LaneState, pack_bits +import tsim.utils.kernels as kernels from tsim.statevector.program import ( MEASURE, MEASURE_DIAG, @@ -51,56 +46,55 @@ ROTATE_DIAG, StatevectorComponent, ) +from tsim.utils.kernels import jit, kernel from tsim.utils.linalg import find_basis +from tsim.utils.rng import GOLDEN, next_skip, splitmix64 -_MIN_CAPACITY = 16 +_TINY = 1e-30 +_I_POWERS = (1, 1j, -1, -1j) -class SharingLayout(eqx.Module): - """Which symbols the shots of a lane must share, from the sign and output rows. +class SharingLayout(NamedTuple): + """Which symbols the shots of a lane must share, and which outputs prune them. Attributes: key_rows: f-bit masks whose parities form the grouping key, a basis of the f-parts of every rotation sign and every postselected output, shape ``(K, num_words)``. lane_dormant: Dormant measurements that split lanes, shape ``(A,)``. + prune_first: The postselected outputs of f-bits only, tested before the + first action, shape ``(num_outputs,)``. pruned: Actions after which some postselected output is tested, shape ``(A,)``. prune_outputs: The postselected outputs tested after each action, shape ``(A, num_outputs)``. - prune_flip: Reference bit XOR-ed into each output before the - postselection test, shape ``(num_outputs,)``. + output_const: Constant of each output XOR its reference bit, shape + ``(num_outputs,)``. lane_branches: Branch bits a lane carries, shape ``(num_words,)``. shot_dormant: Branch bits of the dormant measurements drawn per shot, shape ``(num_words,)``. ready: The action that sets each output's last branch bit, -1 for an output of f-bits only, shape ``(num_outputs,)``. - num_key_words: Number of uint32 words of a packed key. + i_power: ``i^|x∧z|`` of each action's Pauli ``X^x Z^z``, shape ``(A,)``. + factor: Each action's diagonal rotation factors ``cos θ ∓ i sin θ``, + shape ``(A, 2)``. + coupling: Each action's rotation coupling ``-i sin θ i^|x∧z|``, shape + ``(A,)``. """ - key_rows: Array - lane_dormant: Array - pruned: Array - prune_outputs: Array - prune_flip: Array - lane_branches: Array - shot_dormant: Array - ready: Array - num_key_words: int = eqx.field(static=True) - - -@dataclass -class Sharing: - """A component's sharing layout and the largest lane table it has needed. - - Every new capacity compiles the lane loop again (about a second), so narrow - components start each batch at ``lane_capacity`` and settle on one compiled - loop. - """ - - layout: SharingLayout - lane_capacity: int = 0 + key_rows: np.ndarray + lane_dormant: np.ndarray + prune_first: np.ndarray + pruned: np.ndarray + prune_outputs: np.ndarray + output_const: np.ndarray + lane_branches: np.ndarray + shot_dormant: np.ndarray + ready: np.ndarray + i_power: np.ndarray + factor: np.ndarray + coupling: np.ndarray def _unpack(words: np.ndarray, num_words: int) -> np.ndarray: @@ -129,37 +123,29 @@ def sharing_layout( postselected: Whether each output slot is postselected, shape ``(num_outputs,)``. ``None`` prunes nothing. flipped: Reference bit of each output slot, XOR-ed into the output before - the postselection test (``use_detector_reference_sample``), shape - ``(num_outputs,)``. ``None`` flips nothing. + the postselection test and before a pruned shot's later outputs are + cleared, shape ``(num_outputs,)``. ``None`` flips nothing. """ + tape = component.tape num_words = component.num_words - num_f = len(component.f_selection) - kind = np.asarray(component.kind) - branch = np.asarray(component.branch) - num_actions = len(kind) - sign_bits = _unpack(np.asarray(component.sign_row), num_words) - output_bits = _unpack(np.asarray(component.output_row), num_words) + num_f = len(tape.f_selection) + kind, branch = tape.kind, tape.branch num_outputs = component.num_outputs + sign_bits = _unpack(tape.sign_row, num_words) + output_bits = _unpack(tape.output_row, num_words) + no_outputs = np.zeros(num_outputs, dtype=bool) postselected = np.asarray( - np.zeros(num_outputs) if postselected is None else postselected, dtype=bool - ) - flipped = np.asarray( - np.zeros(num_outputs) if flipped is None else flipped, dtype=bool + no_outputs if postselected is None else postselected, dtype=bool ) + flipped = np.asarray(no_outputs if flipped is None else flipped, dtype=bool) setter = np.full(32 * num_words, -1, dtype=np.int32) setter[branch[branch >= 0]] = np.flatnonzero(branch >= 0) ready = np.max(np.where(output_bits, setter[None, :], -1), axis=1, initial=-1) - measurements = np.flatnonzero(np.isin(kind, (MEASURE_DIAG, MEASURE))) - if len(measurements) == 0: - measurements = np.flatnonzero(kind == MEASURE_DORMANT) - first_measurement = measurements[0] if len(measurements) else -1 - prune_at = np.where(ready >= 0, ready, first_measurement) - prune_outputs = np.zeros((num_actions, num_outputs), dtype=bool) - for o in np.flatnonzero(postselected & (prune_at >= 0)): - prune_outputs[prune_at[o], o] = True - pruned = prune_outputs.any(axis=1) + prune_outputs = np.zeros((len(kind), num_outputs), dtype=bool) + for o in np.flatnonzero(postselected & (ready >= 0)): + prune_outputs[ready[o], o] = True rotation = np.isin(kind, (ROTATE_DIAG, ROTATE)) keyed_bits = np.vstack([sign_bits[rotation], output_bits[postselected]]) @@ -168,569 +154,636 @@ def sharing_layout( if len(key_bits): key_bits, _ = find_basis(key_bits.astype(np.uint8)) key_bits = np.asarray(key_bits, dtype=bool) - num_key_words = max(1, -(-len(key_bits) // 32)) + # NOTE: a pruning action sets a branch bit of a postselected output, so its + # branch reaches the key. reaches_key = np.asarray(np.any(keyed_bits, axis=0), dtype=bool) dormant = kind == MEASURE_DORMANT - lane_dormant = dormant & (pruned | ((branch >= 0) & reaches_key[branch])) + lane_dormant = dormant & reaches_key[branch] active = np.isin(kind, (MEASURE_DIAG, MEASURE)) lane_bits = np.zeros(32 * num_words, dtype=bool) - lane_bits[branch[(active | lane_dormant) & (branch >= 0)]] = True + lane_bits[branch[active | lane_dormant]] = True shot_bits = np.zeros(32 * num_words, dtype=bool) - shot_bits[branch[dormant & ~lane_dormant & (branch >= 0)]] = True + shot_bits[branch[dormant & ~lane_dormant]] = True + + quarter_turns = [ + int(x & z).bit_count() & 3 + for x, z in zip(tape.x_mask, tape.z_mask, strict=True) + ] + i_power = np.array([_I_POWERS[q] for q in quarter_turns], dtype=np.complex128) + cos, sin = tape.cos, tape.sin return SharingLayout( - key_rows=jnp.asarray(_pack(key_bits, num_words).reshape(-1, num_words)), - lane_dormant=jnp.asarray(lane_dormant), - pruned=jnp.asarray(pruned), - prune_outputs=jnp.asarray(prune_outputs), - prune_flip=jnp.asarray(flipped), - lane_branches=jnp.asarray(_pack(lane_bits, num_words)), - shot_dormant=jnp.asarray(_pack(shot_bits, num_words)), - ready=jnp.asarray(ready.astype(np.int32)), - num_key_words=num_key_words, + key_rows=_pack(key_bits, num_words).reshape(-1, num_words), + lane_dormant=lane_dormant, + prune_first=postselected & (ready < 0), + pruned=np.asarray(prune_outputs.any(axis=1)), + prune_outputs=prune_outputs, + output_const=tape.output_const ^ flipped, + lane_branches=_pack(lane_bits, num_words), + shot_dormant=_pack(shot_bits, num_words), + ready=ready.astype(np.int64), + i_power=i_power, + factor=np.stack([cos - 1j * sin, cos + 1j * sin], axis=1), + coupling=-1j * sin * i_power, ) -class Lanes(NamedTuple): - """Lane and dead tables of the shared executor. - - Live lanes occupy slots ``0 … used-1`` of the lane table; dead lanes occupy - slots ``0 … dead_used-1`` of the dead table, which has one slot per shot. +@jit +def _parity(value): + """Return the parity of the bits of a non-negative integer below 2^32.""" + value ^= value >> 16 + value ^= value >> 8 + value ^= value >> 4 + value ^= value >> 2 + value ^= value >> 1 + return value & 1 + + +@jit +def _chunks(work, parallel_work, threads): + """Return how many chunks a loop of ``work`` runs in: 1 below ``parallel_work``.""" + if work < parallel_work: + return 1 + return min(work, 4 * threads) + + +@jit +def _copy_lane(amps, target, source, span): + """Copy the first ``span`` amplitudes of lane ``source`` to lane ``target``.""" + for i in range(span): + amps[target, i] = amps[source, i] + + +@jit +def _clear(amps, lanes, lo, hi): + """Set amplitudes ``lo … hi-1`` of the first ``lanes`` lanes to 0.""" + for lane in range(lanes): + for i in range(lo, hi): + amps[lane, i] = 0 + + +@jit +def _grown(amps, rows, lanes, span): + """Return ``amps`` with ``rows`` rows, keeping the lanes' first ``span``.""" + out = np.zeros((rows, amps.shape[1]), dtype=amps.dtype) + for lane in range(lanes): + for i in range(span): + out[lane, i] = amps[lane, i] + return out + + +@jit +def _pack_range(f, selection, out, lo, hi): + """Pack the selected f-bits of shots ``lo … hi-1``: bit j is in word j // 32.""" + for s in range(lo, hi): + for j in range(len(selection)): + if f[s, selection[j]]: + out[s, j >> 5] |= np.uint32(1) << np.uint32(j & 31) + + +@jit +def _key_range(symbols, key_rows, out, lo, hi): + """Pack the parities of shots ``lo … hi-1`` with the key rows into words.""" + for s in range(lo, hi): + for k in range(key_rows.shape[0]): + acc = 0 + for w in range(key_rows.shape[1]): + acc ^= symbols[s, w] & key_rows[k, w] + if _parity(acc): + out[s, k >> 5] |= np.uint32(1) << np.uint32(k & 31) + + +@jit +def _signs(symbols, lanes, row, const): + """Return each lane's sign bit: ``const ^ parity(symbols[l] & row)``.""" + flip = np.zeros(lanes, dtype=np.int64) + for lane in range(lanes): + acc = 0 + for w in range(len(row)): + acc ^= symbols[lane, w] & row[w] + flip[lane] = _parity(acc) ^ const + return flip + + +@jit +def _rotate_range(amps, lo, hi, shift, flip, diagonal, x, z, cos, factor, coupling): + """Apply one rotation to the entries ``lo … hi-1`` of ``lanes × 2^shift``. + + A diagonal rotation multiplies amplitude i by ``factor[sign ^ parity(i & z)]``, + the others mix amplitudes i and ``i ^ x`` with ``cos`` and ``±coupling``. """ - - amps: Array - symbols: Array - count: Array # pyright: ignore[reportIncompatibleMethodOverride] - start: Array - used: Array - dead_symbols: Array - dead_count: Array - dead_start: Array - dead_action: Array - dead_used: Array - action: Array - full: Array - max_norm_deviation: Array - - -def _parity(values: Array) -> Array: - return (jnp.sum(jax.lax.population_count(values), axis=-1) & 1) == 1 - - -@jax.jit -def key_words(layout: SharingLayout, f_packed: Array) -> Array: - """Return each shot's packed key, shape ``(B, num_key_words)``.""" - num_shots = f_packed.shape[0] - if layout.key_rows.shape[0] == 0: - return jnp.zeros((num_shots, layout.num_key_words), dtype=jnp.uint32) - - def key_bit(k, words): - bit = _parity(f_packed & layout.key_rows[k]).astype(jnp.uint32) - word = k // 32 - shift = (k % 32).astype(jnp.uint32) - return words.at[:, word].set(words[:, word] | (bit << shift)) - - return jax.lax.fori_loop( - 0, - layout.key_rows.shape[0], - key_bit, - jnp.zeros((num_shots, layout.num_key_words), dtype=jnp.uint32), - ) - - -def group_on_host( - words: np.ndarray, rng: np.random.Generator -) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """Sort shots by key and return ``(order, start, count)`` of each group. + span = 1 << shift + low = x & -x + for t in range(lo, hi): + lane, i = t >> shift, t & (span - 1) + if diagonal: + amps[lane, i] *= factor[flip[lane] ^ _parity(i & z)] + continue + if i & low: + continue + j = i ^ x + c = -coupling if flip[lane] else coupling + p, q = amps[lane, i], amps[lane, j] + amps[lane, i] = cos * p + c * (-q if _parity(j & z) else q) + amps[lane, j] = cos * q + c * (-p if _parity(i & z) else p) + + +@jit +def _weights_range(amps, lo, hi, span, x, z, i_power, w0, w1): + """Write the outcome weights of measuring ``X^x Z^z`` for lanes ``lo … hi-1``.""" + for lane in range(lo, hi): + for i in range(span): + a = amps[lane, i] + if x == 0: + if _parity(i & z): + w1[lane] += a.real**2 + a.imag**2 + else: + w0[lane] += a.real**2 + a.imag**2 + else: + j = i ^ x + b = i_power * amps[lane, j] + if _parity(j & z): + b = -b + plus, minus = a + b, a - b + w0[lane] += (plus.real**2 + plus.imag**2) / 4 + w1[lane] += (minus.real**2 + minus.imag**2) / 4 + + +@jit +def _project_range(amps, lo, hi, span, x, z, pivot, i_power, outcome, weight): + """Project lanes ``lo … hi-1`` onto their outcome and free the pivot. + + The pivot-1 half of the amplitudes becomes 0: the planner moved the measured + Pauli onto the pivot, so the pivot ends in ``|0⟩``. + """ + bit = 1 << pivot + for lane in range(lo, hi): + if x == 0: + scale = 1 / np.sqrt(max(weight[lane], _TINY)) + else: + scale = 1 / np.sqrt(max(2 * weight[lane], _TINY)) + for i in range(span): + if i & bit: + continue + if x == 0: + j = i ^ bit + keep = _parity(i & z) == outcome[lane] + amps[lane, i] = (amps[lane, i] if keep else amps[lane, j]) * scale + else: + j = i ^ x + b = i_power * amps[lane, j] + if _parity(j & z): + b = -b + if outcome[lane]: + b = -b + amps[lane, i] = (amps[lane, i] + b) * scale + amps[lane, j] = 0 + + +@jit +def _draw_ones(count, w0, w1, state): + """Return how many of each lane's shots see outcome 1. + + The rarer outcome's count is a Binomial(count, q) drawn by geometric skips, + so the chance of a first hit, 1 - (1 - q)^count, is exact to about 2^-64. + """ + ones = np.zeros(len(count), dtype=np.int64) + for lane in range(len(count)): + total = w0[lane] + w1[lane] + q = min(w0[lane], w1[lane]) / total if total > _TINY else 0.0 + hits = 0 + if q > 0: + log_keep = np.log1p(-q) + position = -1.0 + while True: + state, skip = next_skip(state, log_keep) + position += 1.0 + skip + if position >= count[lane]: + break + hits += 1 + ones[lane] = hits if w1[lane] <= w0[lane] else count[lane] - hits + return ones + + +@jit +def _group(keys, state): + """Sort shots by key; return ``(order, start, count, state)`` of the groups. Ties are broken at random (a random permutation, then stable sorts), so the - shots of a group appear in uniformly random order. numpy's radix sort is an - order of magnitude faster than XLA's sort on CPU; a GPU build would sort on - the device. + shots of a group appear in uniformly random order. """ - num_shots = words.shape[0] - order = rng.permutation(num_shots) - for column in range(words.shape[1]): - order = order[np.argsort(words[order, column], kind="stable")] - ordered = words[order] - boundary = np.ones(num_shots, dtype=bool) - boundary[1:] = np.any(ordered[1:] != ordered[:-1], axis=1) - start = np.flatnonzero(boundary) - count = np.diff(np.append(start, num_shots)) - return order.astype(np.int32), start.astype(np.int32), count.astype(np.int32) - - -@partial(jax.jit, static_argnames=("capacity",)) -def initial_lanes( - component: StatevectorComponent, - f_packed: Array, - order: Array, - start: Array, - count: Array, - num_groups: Array, - capacity: int, -) -> Lanes: - """Return one lane per group in ``|0…0⟩``, with an empty dead table.""" - num_shots = f_packed.shape[0] - padding = max(0, capacity - start.shape[0]) - start = jnp.pad(start, (0, padding), constant_values=num_shots)[:capacity] - count = jnp.pad(count, (0, padding))[:capacity] - dtype = jnp.dtype(component.dtype) - representative = order[jnp.minimum(start, num_shots - 1)] - symbols = jnp.where( - (jnp.arange(capacity) < num_groups)[:, None], f_packed[representative], 0 - ) - return Lanes( - amps=jnp.zeros((capacity, 2**component.width), dtype=dtype).at[:, 0].set(1), - symbols=symbols.astype(jnp.uint32), - count=count, - start=start, - used=jnp.asarray(num_groups, dtype=jnp.int32), - dead_symbols=jnp.zeros((num_shots, component.num_words), jnp.uint32), - dead_count=jnp.zeros(num_shots, dtype=jnp.int32), - dead_start=jnp.zeros(num_shots, dtype=jnp.int32), - dead_action=jnp.zeros(num_shots, dtype=jnp.int32), - dead_used=jnp.zeros((), dtype=jnp.int32), - action=jnp.zeros((), dtype=jnp.int32), - full=jnp.zeros((), dtype=jnp.bool_), - max_norm_deviation=jnp.zeros((), dtype=jnp.finfo(dtype).dtype), - ) - - -def _lane_state(lanes: Lanes) -> LaneState: - return LaneState( - amps=lanes.amps, - symbols=lanes.symbols, - count=lanes.count, - max_norm_deviation=lanes.max_norm_deviation, - ) - - -def _rare_drawn(key: Array, count: Array, q: Array) -> Array: - """Return a Binomial(count, q) draw per lane slot, for q ≤ 1/2. - - Whether any shot sees the rare outcome is decided by ``choose_outcome`` - (resolution 2^-64) with P(≥ 1) = 1 - (1 - q)^count; float32 binomial - samplers cut that chance to a 2^-23 grid. Given at least one, the first such - shot J follows a truncated geometric law and the remaining count - J shots - a binomial. Slot s draws from ``fold_in(key, s)``, so a lane's draw does not - depend on the table size. + num_shots = keys.shape[0] + order = np.arange(num_shots) + for i in range(num_shots - 1, 0, -1): + state, z = splitmix64(state) + j = z % np.uint64(i + 1) + order[i], order[j] = order[j], order[i] + for column in range(keys.shape[1]): + order = order[np.argsort(keys[order, column], kind="mergesort")] + start = np.empty(num_shots, dtype=np.int64) + groups = 0 + for p in range(num_shots): + new = p == 0 + for column in range(keys.shape[1]): + new |= keys[order[p], column] != keys[order[p - 1], column] + if new: + start[groups] = p + groups += 1 + start = start[:groups] + count = np.empty(groups, dtype=np.int64) + for g in range(groups): + count[g] = (start[g + 1] if g + 1 < groups else num_shots) - start[g] + return order, start, count, state + + +@jit +def _prune(amps, span, size, dead, tests, output_row, output_const, lanes, graves, a): + """Move the lanes with a tested output of 1 to the dead lanes; return the counts. + + ``tests`` marks the outputs to test, ``lanes`` holds the live lanes' symbols, + starts and counts, and ``graves`` the same and the pruning action of the dead + lanes. The live lanes keep their order in rows ``0 … kept-1``. """ - slots = jnp.arange(count.shape[0]) - keys = jax.vmap(lambda s: jax.random.split(jax.random.fold_in(key, s), 3))(slots) - trials = count.astype(q.dtype) - log_keep = jnp.log1p(-q) - p_any = -jnp.expm1(trials * log_keep) - p_none = jnp.exp(trials * log_keep) - once = jnp.ones(1, dtype=jnp.int32) - hit = jax.vmap( - lambda k, p0, p1: execute.choose_outcome(p0[None], p1[None], once, k)[0] - )(keys[:, 0], p_none, p_any) - - def after_first_hit() -> Array: - uniform = jax.vmap(lambda k: jax.random.uniform(k, dtype=q.dtype))(keys[:, 1]) - first = jnp.ceil( - jnp.log1p(-uniform * p_any) / jnp.where(log_keep < 0, log_keep, -1) - ) - first = jnp.clip(first, 1, jnp.maximum(trials, 1)) - rest = jax.vmap(jax.random.binomial)(keys[:, 2], trials - first, q) - return jnp.where(hit, 1 + jnp.round(rest), 0).astype(jnp.int32) - - # NOTE: rare outcomes rarely happen, so most actions skip the binomials. - rare = jax.lax.cond( - jnp.any(hit), after_first_hit, lambda: jnp.zeros(count.shape, jnp.int32) - ) - return jnp.clip(rare, 0, count) + symbols, start, count = lanes + dead_symbols, dead_start, dead_count, dead_action = graves + kept = 0 + for lane in range(size): + dies = False + for o in range(len(tests)): + if not tests[o]: + continue + acc = 0 + for w in range(symbols.shape[1]): + acc ^= symbols[lane, w] & output_row[o, w] + if _parity(acc) ^ output_const[o]: + dies = True + break + if dies: + dead_symbols[dead] = symbols[lane] + dead_start[dead], dead_count[dead] = start[lane], count[lane] + dead_action[dead] = a + dead += 1 + continue + if kept != lane: + _copy_lane(amps, kept, lane, span) + symbols[kept] = symbols[lane] + start[kept], count[kept] = start[lane], count[lane] + kept += 1 + return kept, dead + + +@jit +def _poison(amps, size, span): + """Fill every amplitude the executor must not read with NaN (for tests).""" + amps[:size, span:] = np.nan + amps[size:] = np.nan + + +@jit +def _shot_range( + symbols, order, owner, entry_symbols, lane_branches, dormant, seed, out, lo, hi +): + """Write the symbols of the shots at positions ``lo … hi-1``. + + A shot's symbols are its f-bits, its lane's branches and its dormant coins. + """ + num_words = symbols.shape[1] + for position in range(lo, hi): + shot, entry = order[position], owner[position] + for w in range(num_words): + word = symbols[shot, w] | (entry_symbols[entry, w] & lane_branches[w]) + if dormant[w]: + n = np.uint64(position * num_words + w) + _, coin = splitmix64(seed + n * GOLDEN) + word |= np.uint32(coin & np.uint64(0xFFFFFFFF)) & dormant[w] + out[shot, w] = word + + +@jit +def _output_range( + shot_symbols, output_row, output_const, ready, pruned_at, out, lo, hi +): + """Write the outputs of shots ``lo … hi-1``; outputs set after pruning are 0.""" + for s in range(lo, hi): + for o in range(output_row.shape[0]): + acc = 0 + for w in range(output_row.shape[1]): + acc ^= shot_symbols[s, w] & output_row[o, w] + out[s, o] = bool(_parity(acc) ^ output_const[o]) and ( + ready[o] <= pruned_at[s] + ) -def _ones_drawn(key: Array, count: Array, weight_0: Array, weight_1: Array) -> Array: - """Return how many of each lane's shots see outcome 1, from the outcome weights.""" - total = jnp.maximum(weight_0 + weight_1, execute._TINY) - rare_is_one = weight_1 <= weight_0 - q = jnp.clip(jnp.minimum(weight_0, weight_1) / total, 0, 0.5) - rare = _rare_drawn(key, count, q) - return jnp.where(rare_is_one, rare, count - rare) +@kernel +def _execute(f, tape, layout, width, seed, share, parallel_work, threads, poison): + """Sample one batch of shots with shared trajectories. + ``tape`` is a ``Tape`` and ``layout`` a ``SharingLayout``. Lanes are rows + ``0 … size-1`` of arrays with a row per shot (every lane holds a shot), but + the amplitudes grow as lanes split. Each loop runs in chunks across threads + once its work reaches ``parallel_work``. -def _split(lanes: Lanes, ones: Array) -> tuple[Lanes, Array, Array]: - """Split every live lane whose shots see both outcomes. + Returns: + ``(samples, deviation, groups, peak, lanes, dead)``. - Returns the split lanes, each slot's outcome and the source slot of each - slot's amplitudes (its parent for a new child). """ - capacity = lanes.count.shape[0] - slots = jnp.arange(capacity) - live = slots < lanes.used - ones = jnp.where(live, ones, 0) - split = live & (ones > 0) & (ones < lanes.count) - child = lanes.used + jnp.cumsum(split) - 1 - target = jnp.where(split, child, capacity) - source = slots.at[target].set(slots, mode="drop") - is_child = jnp.zeros(capacity, dtype=jnp.bool_).at[target].set(True, mode="drop") - outcome = is_child | (live & (ones == lanes.count) & (lanes.count > 0)) - parent_count = lanes.count[source] - parent_ones = ones[source] - count = jnp.where( - is_child, parent_ones, jnp.where(split, lanes.count - ones, lanes.count) + # NOTE: numba's parallel loops cannot read the fields of a named tuple. + selection, output_row = tape.f_selection, tape.output_row + output_const, ready = layout.output_const, layout.ready + key_rows, lane_branches, shot_dormant = ( + layout.key_rows, + layout.lane_branches, + layout.shot_dormant, ) - start = jnp.where( - is_child, lanes.start[source] + parent_count - parent_ones, lanes.start + shots, num_words = f.shape[0], lane_branches.shape[0] + num_actions = tape.kind.shape[0] + # NOTE: SplitMix64 counts up from its state, so seeds that differ by a few + # GOLDEN steps would share outputs; their mixed output starts a fresh stream. + _, state = splitmix64(np.uint64(seed)) + + symbols = np.zeros((shots, num_words), dtype=np.uint32) + num_key_words = (key_rows.shape[0] + 31) >> 5 if share else 0 + keys = np.zeros((shots, num_key_words), dtype=np.uint32) + key_work = key_rows.size if share else 0 + chunks = _chunks(shots * (len(selection) + key_work), parallel_work, threads) + if chunks == 1: + _pack_range(f, selection, symbols, 0, shots) + if share: + _key_range(symbols, key_rows, keys, 0, shots) + else: + for c in numba.prange(chunks): + lo, hi = shots * c // chunks, shots * (c + 1) // chunks + _pack_range(f, selection, symbols, lo, hi) + if share: + _key_range(symbols, key_rows, keys, lo, hi) + if share: + order, start, count, state = _group(keys, state) + else: + order = np.arange(shots) + start = np.arange(shots) + count = np.ones(shots, dtype=np.int64) + groups = len(start) + + # NOTE: lane and dead-lane rows are written before they are read. + lane_symbols = np.empty((shots, num_words), dtype=np.uint32) + lane_start = np.empty(shots, dtype=np.int64) + lane_count = np.empty(shots, dtype=np.int64) + amps = np.zeros((min(max(16, 2 * groups), shots), 1 << width), dtype=np.complex128) + for g in range(groups): + lane_symbols[g] = symbols[order[start[g]]] + lane_start[g], lane_count[g] = start[g], count[g] + amps[g, 0] = 1 + lanes = (lane_symbols, lane_start, lane_count) + graves = ( + np.empty((shots, num_words), dtype=np.uint32), + np.empty(shots, dtype=np.int64), + np.empty(shots, dtype=np.int64), + np.empty(shots, dtype=np.int64), ) - split_lanes = lanes._replace( - symbols=lanes.symbols[source], - count=count.astype(jnp.int32), - start=start.astype(jnp.int32), - used=(lanes.used + jnp.sum(split)).astype(jnp.int32), - ) - return split_lanes, outcome, source - - -def _set_branch(component: StatevectorComponent, symbols, a, outcome): - branch = component.branch[a] - word = jnp.maximum(branch, 0) // 32 - bit = jnp.where( - branch >= 0, - jnp.uint32(1) << (jnp.maximum(branch, 0) % 32).astype(jnp.uint32), - jnp.uint32(0), - ) - flags = jnp.where(outcome, bit, jnp.uint32(0)) - return symbols.at[:, word].set(symbols[:, word] | flags) - - -def _prune( - component: StatevectorComponent, layout: SharingLayout, lanes: Lanes, a: Array -) -> Lanes: - """Move live lanes that fail the postselection after action ``a`` to the dead table.""" - capacity = lanes.count.shape[0] - num_shots = lanes.dead_count.shape[0] - live = jnp.arange(capacity) < lanes.used - values = execute.outputs(component, _lane_state(lanes)) ^ layout.prune_flip - dies = live & jnp.any(values & layout.prune_outputs[a], axis=1) - dead_slot = jnp.where(dies, lanes.dead_used + jnp.cumsum(dies) - 1, num_shots) - keeps = live & ~dies - kept_slot = jnp.where(keeps, jnp.cumsum(keeps) - 1, capacity) - - def compact(values): - return jnp.zeros_like(values).at[kept_slot].set(values, mode="drop") - - def bury(dead, values): - return dead.at[dead_slot].set(values, mode="drop") - - return lanes._replace( - amps=compact(lanes.amps), - symbols=compact(lanes.symbols), - count=compact(lanes.count), - start=compact(lanes.start), - used=jnp.sum(keeps).astype(jnp.int32), - dead_symbols=bury(lanes.dead_symbols, lanes.symbols), - dead_count=bury(lanes.dead_count, lanes.count), - dead_start=bury(lanes.dead_start, lanes.start), - dead_action=bury(lanes.dead_action, jnp.full(capacity, a, dtype=jnp.int32)), - dead_used=(lanes.dead_used + jnp.sum(dies)).astype(jnp.int32), - ) - - -def _finish(component, layout, lanes, split, a) -> Lanes: - """Prune ``split`` and advance, or flag ``lanes`` if the lane table is full.""" - pruned = jax.lax.cond( - layout.pruned[a], - lambda: _prune(component, layout, split, a), - lambda: split, - ) - full = split.used > lanes.count.shape[0] - advanced = pruned._replace(action=(a + 1).astype(jnp.int32)) - flagged = lanes._replace(full=jnp.asarray(True)) - return jax.tree.map(lambda old, new: jnp.where(full, old, new), flagged, advanced) - - -def _measure_weights( - component: StatevectorComponent, amps: Array, a: Array -) -> tuple[Array, Array]: - """Return each lane's outcome weights ``(w0, w1)`` of ``MEASURE`` action ``a``.""" - paulied = execute._apply_pauli( - amps, execute._index(component), component.x_mask[a], component.z_mask[a] - ) - weight_0 = jnp.sum(execute._squared(amps + paulied), axis=1) / 4 - weight_1 = jnp.sum(execute._squared(amps - paulied), axis=1) / 4 - return weight_0, weight_1 - - -def _measure_project( - component: StatevectorComponent, amps: Array, a: Array, outcome: Array, weight -) -> Array: - """Return the amplitudes after ``MEASURE`` ``a`` gave ``outcome`` of ``weight``.""" - index = execute._index(component) - paulied = execute._apply_pauli( - amps, index, component.x_mask[a], component.z_mask[a] + size, dead = _prune( + amps, + 1, + groups, + 0, + layout.prune_first, + output_row, + output_const, + lanes, + graves, + -1, ) - projected = jnp.where(outcome[:, None], amps - paulied, amps + paulied) - scale = jax.lax.rsqrt(jnp.maximum(2 * weight, execute._TINY))[:, None] - keep = ((index >> component.pivot[a].astype(jnp.uint32)) & 1) == 0 - return jnp.where(keep[None, :], projected * scale, 0) - - -def _measure_diag_weights( - component: StatevectorComponent, amps: Array, a: Array -) -> tuple[Array, Array]: - """Return each lane's outcome weights ``(w0, w1)`` of ``MEASURE_DIAG`` ``a``.""" - odd = execute._parity(execute._index(component) & component.z_mask[a]) - weights = execute._squared(amps) - weight_0 = jnp.sum(jnp.where(odd[None, :], 0, weights), axis=1) - weight_1 = jnp.sum(jnp.where(odd[None, :], weights, 0), axis=1) - return weight_0, weight_1 - - -def _measure_diag_project( - component: StatevectorComponent, amps: Array, a: Array, outcome: Array, weight -) -> Array: - """Return the amplitudes after ``MEASURE_DIAG`` ``a`` gave ``outcome``.""" - index = execute._index(component) - odd = execute._parity(index & component.z_mask[a]) - pivot = component.pivot[a].astype(jnp.uint32) - partner = jnp.take(amps, (index ^ (jnp.uint32(1) << pivot)).astype(jnp.int32), 1) - gathered = jnp.where(odd[None, :] == outcome[:, None], amps, partner) - scale = jax.lax.rsqrt(jnp.maximum(weight, execute._TINY))[:, None] - keep = ((index >> pivot) & 1) == 0 - return jnp.where(keep[None, :], gathered * scale, 0) - - -def _measure_kernel(weights, project): - def kernel(component, layout, lanes, a, key): - weight_0, weight_1 = weights(component, lanes.amps, a) - live = jnp.arange(lanes.count.shape[0]) < lanes.used - deviation = jnp.max(jnp.where(live, jnp.abs(weight_0 + weight_1 - 1), 0)) - ones = _ones_drawn(jax.random.fold_in(key, a), lanes.count, weight_0, weight_1) - split, outcome, source = _split(lanes, ones) - amps = lanes.amps[source] - weight = jnp.where(outcome, weight_1[source], weight_0[source]) - split = split._replace( - amps=project(component, amps, a, outcome, weight).astype(amps.dtype), - symbols=_set_branch(component, split.symbols, a, outcome), - max_norm_deviation=jnp.maximum( - lanes.max_norm_deviation, - deviation.astype(lanes.max_norm_deviation.dtype), - ), - ) - return _finish(component, layout, lanes, split, a) - - return kernel - - -def _rotate_kernel(kernel): - def apply(component, layout, lanes, a, key): - del layout - amps = kernel(component, _lane_state(lanes), a, key).amps - return lanes._replace(amps=amps, action=(a + 1).astype(jnp.int32)) - - return apply - + peak, region, deviation = groups, 0, 0.0 + + for a in range(num_actions): + if size == 0: + break + if poison: + _poison(amps, size, 1 << region) + shift = tape.extent[a] + span = 1 << shift + if shift > region: + _clear(amps, size, 1 << region, span) + region = shift + kind = tape.kind[a] + x, z = np.int64(tape.x_mask[a]), np.int64(tape.z_mask[a]) + if kind in (ROTATE_DIAG, ROTATE): + flip = _signs( + lane_symbols, size, tape.sign_row[a], np.int64(tape.sign_const[a]) + ) + diagonal = kind == ROTATE_DIAG + cos, factor, coupling = tape.cos[a], layout.factor[a], layout.coupling[a] + work = size * span + chunks = _chunks(work, parallel_work, threads) + if chunks == 1: + _rotate_range( + amps, 0, work, shift, flip, diagonal, x, z, cos, factor, coupling + ) + else: + for c in numba.prange(chunks): + lo, hi = work * c // chunks, work * (c + 1) // chunks + _rotate_range( + amps, lo, hi, shift, flip, diagonal, x, z, cos, factor, coupling + ) + continue + dormant = kind == MEASURE_DORMANT + if dormant and not layout.lane_dormant[a]: + continue + half = 0.5 if dormant else 0.0 + w0 = np.full(size, half) + w1 = np.full(size, half) + i_power, pivot = layout.i_power[a], tape.pivot[a] + if not dormant: + chunks = _chunks(size * span, parallel_work, threads) + if chunks == 1: + _weights_range(amps, 0, size, span, x, z, i_power, w0, w1) + else: + for c in numba.prange(chunks): + lo, hi = size * c // chunks, size * (c + 1) // chunks + _weights_range(amps, lo, hi, span, x, z, i_power, w0, w1) + for lane in range(size): + d = abs(w0[lane] + w1[lane] - 1) + # NOTE: written so that a NaN deviation sticks. + if d > deviation or d != d: + deviation = d + state, draw = splitmix64(state) + ones = _draw_ones(lane_count[:size], w0, w1, draw) + + new = size + for lane in range(size): + if 0 < ones[lane] < lane_count[lane]: + new += 1 + if new > amps.shape[0]: + rows = min(max(new, 2 * amps.shape[0]), shots) + amps = _grown(amps, rows, size, span) + outcome = np.zeros(new, dtype=np.bool_) + weight = np.zeros(new) + child = size + for lane in range(size): + n, total = ones[lane], lane_count[lane] + outcome[lane] = n == total + weight[lane] = w1[lane] if n == total else w0[lane] + if 0 < n < total: + _copy_lane(amps, child, lane, span) + lane_symbols[child] = lane_symbols[lane] + lane_count[child] = n + lane_start[child] = lane_start[lane] + total - n + lane_count[lane] = total - n + outcome[child] = True + weight[child] = w1[lane] + child += 1 + size = new + if not dormant: + chunks = _chunks(size * span, parallel_work, threads) + if chunks == 1: + _project_range( + amps, 0, size, span, x, z, pivot, i_power, outcome, weight + ) + else: + for c in numba.prange(chunks): + lo, hi = size * c // chunks, size * (c + 1) // chunks + _project_range( + amps, lo, hi, span, x, z, pivot, i_power, outcome, weight + ) + bit = tape.branch[a] + if bit >= 0: + for lane in range(size): + if outcome[lane]: + lane_symbols[lane, bit >> 5] |= np.uint32(1) << np.uint32(bit & 31) + peak = max(peak, size) + if layout.pruned[a]: + tests = layout.prune_outputs[a] + size, dead = _prune( + amps, + span, + size, + dead, + tests, + output_row, + output_const, + lanes, + graves, + a, + ) -def _dormant(component, layout, lanes, a, key): - def split_lanes(lanes): - half = jnp.full(lanes.count.shape, 0.5, dtype=lanes.max_norm_deviation.dtype) - ones = _ones_drawn(jax.random.fold_in(key, a), lanes.count, half, half) - split, outcome, source = _split(lanes, ones) - split = split._replace( - amps=lanes.amps[source], - symbols=_set_branch(component, split.symbols, a, outcome), + # NOTE: the live and dead lanes' ranges tile the sorted positions. + dead_symbols, dead_start, dead_count, dead_action = graves + entry_symbols = np.zeros((size + dead, num_words), dtype=np.uint32) + entry_pruned = np.empty(size + dead, dtype=np.int64) + owner = np.empty(shots, dtype=np.int64) + for e in range(size + dead): + if e < size: + entry_symbols[e] = lane_symbols[e] + entry_pruned[e] = num_actions + first, last = lane_start[e], lane_start[e] + lane_count[e] + else: + entry_symbols[e] = dead_symbols[e - size] + entry_pruned[e] = dead_action[e - size] + first = dead_start[e - size] + last = first + dead_count[e - size] + for p in range(first, last): + owner[p] = e + pruned_at = np.empty(shots, dtype=np.int64) + for p in range(shots): + pruned_at[order[p]] = entry_pruned[owner[p]] + state, coins = splitmix64(state) + shot_symbols = np.empty_like(symbols) + samples = np.zeros((shots, output_row.shape[0]), dtype=np.bool_) + chunks = _chunks(shots * (num_words + output_row.size), parallel_work, threads) + if chunks == 1: + _shot_range( + symbols, + order, + owner, + entry_symbols, + lane_branches, + shot_dormant, + coins, + shot_symbols, + 0, + shots, ) - return _finish(component, layout, lanes, split, a) - - def skip(lanes): - return lanes._replace(action=(a + 1).astype(jnp.int32)) - - return jax.lax.cond(layout.lane_dormant[a], split_lanes, skip, lanes) - - -# NOTE: the order must match the kind codes ROTATE_DIAG … MEASURE_DORMANT in -# program.py. -_KERNELS = ( - _rotate_kernel(execute._rotate_diag), - _rotate_kernel(execute._rotate), - _measure_kernel(_measure_diag_weights, _measure_diag_project), - _measure_kernel(_measure_weights, _measure_project), - _dormant, -) - - -@jax.jit -def run_lanes( - component: StatevectorComponent, - layout: SharingLayout, - lanes: Lanes, - key: Array, - stop: Array, -) -> Lanes: - """Run actions ``lanes.action … stop-1``, or until the lane table is full.""" - branches = [partial(kernel, component, layout) for kernel in _KERNELS] - - def step(lanes: Lanes) -> Lanes: - a = lanes.action - return jax.lax.switch(component.kind[a], branches, lanes, a, key) - - def going(lanes: Lanes) -> Array: - return (lanes.action < stop) & ~lanes.full - - return jax.lax.while_loop(going, step, lanes) - - -@partial(jax.jit, static_argnames=("capacity",)) -def resize(lanes: Lanes, capacity: int) -> Lanes: - """Return ``lanes`` with a new lane-table capacity (the live lanes must fit).""" - - def fit(values): - if capacity <= values.shape[0]: - return values[:capacity] - padding = jnp.zeros( - (capacity - values.shape[0], *values.shape[1:]), values.dtype + _output_range( + shot_symbols, output_row, output_const, ready, pruned_at, samples, 0, shots ) - return jnp.concatenate([values, padding]) - - return lanes._replace( - amps=fit(lanes.amps), - symbols=fit(lanes.symbols), - count=fit(lanes.count), - start=fit(lanes.start), - full=jnp.zeros((), dtype=jnp.bool_), - ) - + else: + for c in numba.prange(chunks): + lo, hi = shots * c // chunks, shots * (c + 1) // chunks + _shot_range( + symbols, + order, + owner, + entry_symbols, + lane_branches, + shot_dormant, + coins, + shot_symbols, + lo, + hi, + ) + for c in numba.prange(chunks): + lo, hi = shots * c // chunks, shots * (c + 1) // chunks + _output_range( + shot_symbols, + output_row, + output_const, + ready, + pruned_at, + samples, + lo, + hi, + ) + return samples, deviation, groups, peak, size, dead -@jax.jit -def reconstruct( - component: StatevectorComponent, - layout: SharingLayout, - order: Array, - lanes: Lanes, - f_packed: Array, - key: Array, -) -> Array: - """Return every shot's outputs, in the original shot order.""" - num_shots = f_packed.shape[0] - capacity = lanes.count.shape[0] - live = jnp.concatenate( - [jnp.arange(capacity) < lanes.used, jnp.arange(num_shots) < lanes.dead_used] - ) - count = jnp.concatenate([lanes.count, lanes.dead_count]) - start = jnp.where( - live & (count > 0), jnp.concatenate([lanes.start, lanes.dead_start]), num_shots - ) - symbols = jnp.concatenate([lanes.symbols, lanes.dead_symbols]) - pruned_at = jnp.concatenate( - [jnp.full(capacity, component.num_actions, jnp.int32), lanes.dead_action] - ) - # The ranges partition the positions, so each position's entry is the one - # starting at the latest range start at or before it. - position = jnp.arange(num_shots) - starting = ( - jnp.full(num_shots, -1, jnp.int32) - .at[start] - .set(jnp.arange(start.shape[0], dtype=jnp.int32), mode="drop") - ) - latest_start = jax.lax.cummax(jnp.where(starting >= 0, position, -1)) - entry = starting[latest_start] - coins = jax.random.bits(key, (num_shots, component.num_words), dtype=jnp.uint32) - shot_symbols = ( - f_packed[order] - | (symbols[entry] & layout.lane_branches) - | (coins & layout.shot_dormant) - ) - values = execute._affine( - shot_symbols[:, None, :], component.output_row, component.output_const - ) - written = layout.ready[None, :] <= pruned_at[entry][:, None] - values = values & written - return jnp.zeros_like(values).at[order].set(values) +# NOTE: tests set this to fill every amplitude the executor must not read with NaN. +_POISON = False -def _capacity(needed: int) -> int: - return max(_MIN_CAPACITY, 1 << int(np.ceil(np.log2(max(needed, 1))))) +def lane_bytes(component: StatevectorComponent) -> int: + """Return the peak memory of one lane of ``sample_shared_component``. -@partial(jax.jit, static_argnames=("num_words",)) -def _packed_f(f_params: Array, f_selection: Array, num_words: int) -> Array: - return pack_bits(f_params[:, f_selection].astype(jnp.bool_), num_words) + That is three rows of amplitudes (the old and the new array while lanes + grow) and the per-shot symbols, ranges and outputs. + """ + amplitudes = 3 * 2**component.width * np.dtype(np.complex128).itemsize + return amplitudes + 24 * component.num_words + 128 + component.num_outputs def sample_shared_component( component: StatevectorComponent, - sharing: Sharing, - f_params: Array, - key: Array, - *, - segment: int | None = None, -) -> tuple[Array, Array, Array, dict]: - """Sample one component's outputs with shared trajectories. + layout: SharingLayout, + f_params: np.ndarray, + seed: int, + share: bool = True, +) -> tuple[np.ndarray, float, dict]: + """Sample one component's outputs, with shared trajectories if ``share``. Args: component: The compiled component. - sharing: Its sharing layout and capacity hint (updated in place). + layout: Its sharing layout. f_params: Error-basis bits of every shot, shape ``(B, num_f_global)``. - key: JAX random key. - segment: Actions between host checks that shrink the lane table to the - live lanes. Defaults to 32 for ``W ≥ 8`` and to the whole tape below, - where shrinking saves less array work than its compilations cost. + seed: Seed of the batch's random numbers, below 2^64. + share: Whether shots with equal keys share a lane; otherwise every shot + has its own. Returns: - ``(samples, next_key, max_norm_deviation, stats)``, where ``samples`` has - shape ``(B, num_outputs)`` and ``stats`` counts groups and lanes. + ``(samples, max_norm_deviation, stats)``, where ``samples`` has shape + ``(B, num_outputs)`` and ``stats`` counts groups and lanes. """ - layout = sharing.layout - num_actions = component.num_actions - if segment is None: - segment = 32 if component.width >= 8 else max(num_actions, 1) - f_packed = _packed_f(f_params, component.f_selection, component.num_words) - num_shots = f_packed.shape[0] - key, order_key, lane_key, coin_key = jax.random.split(key, 4) - rng = np.random.default_rng(np.asarray(jax.random.bits(order_key, (4,)))) - order_np, start_np, count_np = group_on_host( - np.asarray(key_words(layout, f_packed)), rng - ) - groups = len(start_np) - sticky = sharing.lane_capacity if segment >= num_actions else 0 - capacity = max(_capacity(2 * groups), sticky) - start_padded = np.full(capacity, num_shots, dtype=np.int32) - count_padded = np.zeros(capacity, dtype=np.int32) - start_padded[:groups], count_padded[:groups] = start_np, count_np - order = jnp.asarray(order_np) - lanes = initial_lanes( - component, - f_packed, - order, - jnp.asarray(start_padded), - jnp.asarray(count_padded), - jnp.asarray(groups, dtype=jnp.int32), - capacity=capacity, + f = np.asarray(f_params) + samples, deviation, groups, peak, lanes, dead = _execute( + f.shape[0] << component.width, + f, + component.tape, + layout, + component.width, + np.uint64(seed), + share, + kernels.PARALLEL_WORK, + numba.get_num_threads(), + _POISON, ) - resumes = peak = dead_used = 0 - largest = capacity - action, used = 0, groups - while action < num_actions: - stop = jnp.asarray(min(action + segment, num_actions), dtype=jnp.int32) - lanes = run_lanes(component, layout, lanes, lane_key, stop) - full, action, used, dead_used = ( - int(value) - for value in jax.device_get( - (lanes.full, lanes.action, lanes.used, lanes.dead_used) - ) - ) - peak = max(peak, used) - capacity = lanes.count.shape[0] - if full: - resumes += 1 - lanes = resize(lanes, capacity=2 * capacity) - elif action < num_actions and _capacity(2 * used) < capacity: - lanes = resize(lanes, capacity=_capacity(2 * used)) - largest = max(largest, lanes.count.shape[0]) - sharing.lane_capacity = max(sharing.lane_capacity, largest) - samples = reconstruct(component, layout, order, lanes, f_packed, coin_key) - stats = { - "groups": groups, - "peak lanes": peak, - "lanes": used, - "dead": dead_used, - "lane_capacity": lanes.count.shape[0], - "resumes": resumes, - } - return samples, key, lanes.max_norm_deviation, stats + stats = {"groups": groups, "peak lanes": peak, "lanes": lanes, "dead": dead} + return samples, float(deviation), stats diff --git a/src/tsim/utils/kernels.py b/src/tsim/utils/kernels.py new file mode 100644 index 00000000..574a16be --- /dev/null +++ b/src/tsim/utils/kernels.py @@ -0,0 +1,137 @@ +"""Numba kernels that run in parallel only when the work pays for it.""" + +import functools +import os +import threading +import types +from collections.abc import Callable, Iterable +from concurrent.futures import ThreadPoolExecutor, wait +from typing import TypeVar + +import numba + +# NOTE: waking numba's threads costs about 0.1 ms, more than many kernel calls +# take serially. Numba's default threading layer aborts the process when two +# Python threads run parallel code at once, so parallel calls hold a lock; each +# one uses every core anyway. +PARALLEL_WORK = 1 << 16 +_LOCK = threading.Lock() +_TASK_THREAD = threading.local() + +_Task = TypeVar("_Task") +_Result = TypeVar("_Result") + + +def _renamed(function: Callable, suffix: str) -> Callable: + """Return a copy of ``function`` whose name ends in ``suffix``.""" + renamed = types.FunctionType( + function.__code__, + function.__globals__, + function.__name__ + suffix, + function.__defaults__, + function.__closure__, + ) + renamed.__qualname__ = function.__qualname__ + suffix + renamed.__doc__ = function.__doc__ + return renamed + + +def _prange_only() -> dict[str, bool]: + """Return numba's parallel options for explicit ``prange`` loops only. + + Numba would also run array expressions such as ``np.zeros`` or slice fills + in parallel loops, each a thread launch, even inside a serial stretch. + """ + options = ("comprehension", "reduction", "inplace_binop", "setitem", "numpy") + return {**dict.fromkeys(options, False), "stencil": False, "fusion": False} + + +def jit(function: Callable, *, parallel: bool = False) -> Callable: + """Compile ``function`` with numba, cached on disk where numba can write. + + The compiled function releases the GIL, so Python threads can run it at once. + Numba refuses to define a cached function when no cache directory is + writable (e.g. a read-only install without ``NUMBA_CACHE_DIR``), so such + installs compile in every process instead. + """ + # NOTE: numba names its cache files after the function and ignores nogil when + # it reads them, so the name carries the flag: an entry cached without nogil + # would keep holding the GIL. + function = _renamed(function, "_nogil") + options = _prange_only() if parallel else False + try: + return numba.njit(cache=True, parallel=options, nogil=True)(function) + except RuntimeError: + return numba.njit(parallel=options, nogil=True)(function) + + +def kernel(function: Callable) -> Callable: + """Compile ``function`` serially and in parallel; ``run(work, *args)`` picks one. + + ``work`` estimates the call's cost in simple operations; calls below + ``PARALLEL_WORK`` run serially, and so do calls on the worker threads of + ``run_tasks``. In the serial compilation every ``numba.prange`` is a plain + range, so both compute the same result. + """ + serial = jit(function) + # NOTE: numba names its cache files after the function, so the parallel + # compilation needs another name to keep its own cache entry. + parallel = jit(_renamed(function, "_prange"), parallel=True) + + def run(work: int, *args): + if work < PARALLEL_WORK or getattr(_TASK_THREAD, "active", False): + return serial(*args) + with _LOCK: + return parallel(*args) + + run.__doc__ = function.__doc__ + return run + + +def _mark_task_thread() -> None: + _TASK_THREAD.active = True + + +@functools.cache +def _pool() -> ThreadPoolExecutor: + # NOTE: one pool for every call, as large as numba's thread count can get. + return ThreadPoolExecutor( + numba.config.NUMBA_NUM_THREADS, # pyright: ignore[reportAttributeAccessIssue] + initializer=_mark_task_thread, + ) + + +def _after_fork_in_child() -> None: + """Start over in a forked child: the parent's threads do not exist there.""" + global _LOCK + _LOCK = threading.Lock() + _pool.cache_clear() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_after_fork_in_child) + + +def run_tasks( + function: Callable[[_Task], _Result], tasks: Iterable[_Task], threads: int +) -> list[_Result]: + """Return ``[function(task) for task in tasks]``, on up to ``threads`` threads. + + With more than one thread the tasks run on worker threads, and every kernel + in them runs serially: the tasks keep the cores busy, and independent tasks + scale far better than numba's threads inside each kernel. Otherwise, and in + a task, the tasks run in the calling thread. Every task finishes before the + first error, if any, is raised. + """ + tasks = list(tasks) + if min(threads, len(tasks)) <= 1 or getattr(_TASK_THREAD, "active", False): + return [function(task) for task in tasks] + slots = threading.Semaphore(threads) + + def run(task: _Task) -> _Result: + with slots: + return function(task) + + futures = [_pool().submit(run, task) for task in tasks] + wait(futures) + return [future.result() for future in futures] diff --git a/src/tsim/utils/rng.py b/src/tsim/utils/rng.py new file mode 100644 index 00000000..8da74ec2 --- /dev/null +++ b/src/tsim/utils/rng.py @@ -0,0 +1,47 @@ +"""SplitMix64 random numbers for numba kernels. + +A kernel threads a uint64 state through these functions. Seeding each block of +work with its own 64-bit seed makes its numbers independent of how the blocks +are scheduled across threads. +""" + +import numpy as np + +from tsim.utils.kernels import jit + +GOLDEN = np.uint64(0x9E3779B97F4A7C15) + + +@jit +def splitmix64(state: np.uint64) -> tuple[np.uint64, np.uint64]: + """Advance a SplitMix64 state; return it and the next 64-bit output. + + ``splitmix64(seed + n * GOLDEN)`` is the stream's ``n + 1``-th output, so a + kernel can also read the stream at any index. + """ + state += GOLDEN + z = (state ^ (state >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9) + z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB) + return state, z ^ (z >> np.uint64(31)) + + +@jit +def next_uniform(state: np.uint64) -> tuple[np.uint64, float]: + """Return the advanced state and a uniform float in [0, 1) with 53 bits.""" + state, z = splitmix64(state) + return state, (z >> np.uint64(11)) * 2.0**-53 + + +@jit +def next_skip(state: np.uint64, log_keep: float) -> tuple[np.uint64, float]: + """Return the advanced state and the failures before the next success. + + The trials succeed with probability p and ``log_keep = log1p(-p)``. The + uniform behind the draw has 64-bit resolution near 0, so the chance of no + failure is exact for p down to about 2^-64. The count is a float: it is inf + when the uniform rounds to 1, so callers compare it before converting it. + """ + if log_keep == -np.inf: + return state, 0.0 + state, z = splitmix64(state) + return state, np.floor(np.log1p(-(z * 2.0**-64)) / log_keep) diff --git a/test/integration/test_sampler.py b/test/integration/test_sampler.py index c4970b77..c7e3e64d 100644 --- a/test/integration/test_sampler.py +++ b/test/integration/test_sampler.py @@ -117,9 +117,9 @@ def test_sample_program_raises_on_component_norm_deviation(monkeypatch): ) program = CompiledProgram( components=components, - direct_f_indices=jnp.array([], dtype=jnp.int32), - direct_flips=jnp.array([], dtype=jnp.bool_), - output_order=jnp.array([0, 1]), + direct_f_indices=np.array([], dtype=np.int32), + direct_flips=np.array([], dtype=np.bool_), + output_order=np.array([0, 1]), output_reindex=None, num_outputs=2, num_detectors=0, diff --git a/test/integration/test_sampler_circuits.py b/test/integration/test_sampler_circuits.py index adcb51fb..485bce86 100644 --- a/test/integration/test_sampler_circuits.py +++ b/test/integration/test_sampler_circuits.py @@ -33,8 +33,9 @@ def test_detector_sampler_bell_state_with_measurement_error(): """) sampler = c.compile_detector_sampler(seed=1) - d = sampler.sample(10) - assert np.count_nonzero(d) == 4 + shots = 100_000 + d = sampler.sample(shots) + assert abs(np.count_nonzero(d) / shots - 0.3) < 5 * np.sqrt(0.3 * 0.7 / shots) def test_t_gate(): diff --git a/test/integration/test_statevector_sampler.py b/test/integration/test_statevector_sampler.py index 3e6efd5b..4c0363b8 100644 --- a/test/integration/test_statevector_sampler.py +++ b/test/integration/test_statevector_sampler.py @@ -4,7 +4,6 @@ from test.helpers.gen import gen_stim_circuit from test.helpers.statevector import CIRCUITS, assert_counts_match -import jax import jax.numpy as jnp import numpy as np import pytest @@ -15,8 +14,8 @@ from tsim.core.graph import prepare_graph from tsim.core.types import SamplingStrategy from tsim.external.vec_sim.vec_sampler import VecSampler -from tsim.statevector.execute import lane_bytes from tsim.statevector.program import StatevectorComponent +from tsim.statevector.shared import lane_bytes @pytest.mark.parametrize("name", sorted(CIRCUITS)) @@ -104,54 +103,34 @@ def test_probabilities_reject_statevector(): compile_program(prepared, mode="joint", strategy="active_statevector") -def test_complex128_sampling_matches_stabilizer_rank(): - circuit = Circuit(CIRCUITS["t_ghz"]) - expected = circuit.compile_detector_sampler(strategy="cat5", seed=1).sample( - 10_000, append_observables=True - ) - with jax.enable_x64(True): - sampler = circuit.compile_detector_sampler( - strategy="active_statevector", seed=2 - ) - assert all( - isinstance(c, StatevectorComponent) and c.dtype == "complex128" - for c in sampler._program.components - ) - got = sampler.sample(10_000, append_observables=True) - assert_counts_match(got, expected) - - -def _norm_warnings(monkeypatch, strategy, deviation, *, x64=False): +def _norm_warnings(monkeypatch, strategy, deviation): sample = sampler_module.sample_component + sample_shared = sampler_module.sample_shared_component def drifting(component, f_params, key): samples, key, _ = sample(component, f_params, key) return samples, key, jnp.asarray(deviation) + def drifting_shared(*args): + samples, _, stats = sample_shared(*args) + return samples, deviation, stats + monkeypatch.setattr(sampler_module, "sample_component", drifting) - with jax.enable_x64(x64): - sampler = Circuit(CIRCUITS["t_ghz"]).compile_detector_sampler( - strategy=strategy, seed=0 - ) - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - sampler.sample(10) + monkeypatch.setattr(sampler_module, "sample_shared_component", drifting_shared) + sampler = Circuit(CIRCUITS["t_ghz"]).compile_detector_sampler( + strategy=strategy, seed=0 + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + sampler.sample(10) return [w for w in caught if "not normalized" in str(w.message)] -@pytest.mark.parametrize( - "x64, deviation, warns", - [ - (False, 5e-4, False), - (False, 2e-3, True), - (True, 1e-10, False), - (True, 1e-8, True), - ], -) -def test_statevector_norm_warnings_depend_on_the_amplitude_dtype( - monkeypatch, x64, deviation, warns +@pytest.mark.parametrize("deviation, warns", [(1e-10, False), (1e-8, True)]) +def test_statevector_norm_warnings_have_a_complex128_threshold( + monkeypatch, deviation, warns ): - caught = _norm_warnings(monkeypatch, "active_statevector", deviation, x64=x64) + caught = _norm_warnings(monkeypatch, "active_statevector", deviation) assert bool(caught) == warns @@ -176,13 +155,13 @@ def test_the_automatic_batch_size_of_a_statevector_sampler_is_capped(monkeypatch assert sampler._estimate_batch_size() == 2**16 assert sampler._resolve_batch_size(10**7, None, compute_reference=False) <= 2**16 lanes = [] - sample_program = sampler_module.sample_program + sample_program = sampler_module.sample_statevector_program - def recording(program, f_params, key): + def recording(program, f_params, *args, **kwargs): lanes.append(f_params.shape[0]) - return sample_program(program, f_params, key) + return sample_program(program, f_params, *args, **kwargs) - monkeypatch.setattr(sampler_module, "sample_program", recording) + monkeypatch.setattr(sampler_module, "sample_statevector_program", recording) sampler.sample(2**16 + 1) sampler.sample( 2**16 + 1, postselection_mask=np.ones(circuit.num_detectors, dtype=bool) diff --git a/test/integration/test_trajectory_sharing.py b/test/integration/test_trajectory_sharing.py index aac80e76..99b45257 100644 --- a/test/integration/test_trajectory_sharing.py +++ b/test/integration/test_trajectory_sharing.py @@ -1,4 +1,4 @@ -"""trajectory_sharing=True end to end, against one lane per shot.""" +"""trajectory_sharing=True end to end, against the stabilizer-rank strategy.""" from test.helpers.gen import gen_stim_circuit from test.helpers.statevector import CIRCUITS, assert_counts_match @@ -6,11 +6,12 @@ import numpy as np import pytest +import tsim.sampler as sampler_module +import tsim.utils.kernels as kernels from tsim.circuit import Circuit -# NOTE: t_ghz plus a direct detector: postselection only takes the pruning path -# when it masks a direct detector, and t_ghz's random detectors give reference -# bits of 1 for some seeds. +# NOTE: t_ghz plus a direct detector (detector 2); t_ghz's random detectors give +# reference bits of 1 for some seeds. GHZ_WITH_DIRECT_DETECTOR = """ R 0 1 2 3 H 0 @@ -32,6 +33,42 @@ """ +# NOTE: flipped direct detectors and observables, unflipped ones, and two T +# components, so that the outputs are reordered. +DIRECT_FLIPS = """ +R 0 1 2 3 4 5 +H 0 +T 0 +H 0 +X_ERROR(0.3) 1 +X 2 +X_ERROR(0.2) 2 +M 2 +DETECTOR rec[-1] +M 0 1 +DETECTOR rec[-2] +DETECTOR rec[-1] +X 3 +X_ERROR(0.1) 3 +M 3 +OBSERVABLE_INCLUDE(0) rec[-1] +X_ERROR(0.4) 4 +M 4 +OBSERVABLE_INCLUDE(1) rec[-1] +H 5 +T 5 +H 5 +M 5 +OBSERVABLE_INCLUDE(2) rec[-1] +""" + + +def _sharing_stats(sampler): + backend = sampler._backend + assert isinstance(backend, sampler_module._StatevectorBackend) + return backend.sharing_stats + + def _accepted(sampler, shots, mask, *, reference=False): detectors, observables = sampler.sample( shots, @@ -57,8 +94,11 @@ def test_pruning_through_the_sampler_keeps_the_accepted_distribution(reference): ) if sampler._compute_reference_sample()[:2].any(): break + # NOTE: the baseline needs this sampler's reference bits, so it only stops sharing. rate, expected = _accepted( - circuit.compile_detector_sampler(strategy="active_statevector", seed=seed), + circuit.compile_detector_sampler( + strategy="active_statevector", seed=seed, trajectory_sharing=False + ), shots, mask, reference=reference, @@ -67,19 +107,19 @@ def test_pruning_through_the_sampler_keeps_the_accepted_distribution(reference): strategy="active_statevector", seed=seed, trajectory_sharing=True ) shared_rate, got = _accepted(sampler, shots, mask, reference=reference) - assert any(stats["dead"] > 0 for stats in sampler.sharing_stats) + assert any(stats["dead"] > 0 for stats in _sharing_stats(sampler)) sigma = np.sqrt(rate * (1 - rate) * 2 / shots) assert abs(shared_rate - rate) <= 5 * sigma + 3 / shots assert_counts_match(got, expected) @pytest.mark.parametrize("name", sorted(CIRCUITS)) -def test_shared_detector_sampling_matches_one_lane_per_shot(name): +def test_shared_detector_sampling_matches_stabilizer_rank(name): circuit = Circuit(CIRCUITS[name]) shots = 20_000 - expected = circuit.compile_detector_sampler( - strategy="active_statevector", seed=1 - ).sample(shots, batch_size=5_000, append_observables=True) + expected = circuit.compile_detector_sampler(strategy="cat5", seed=1).sample( + shots, batch_size=5_000, append_observables=True + ) sampler = circuit.compile_detector_sampler( strategy="active_statevector", seed=2, trajectory_sharing=True ) @@ -88,10 +128,10 @@ def test_shared_detector_sampling_matches_one_lane_per_shot(name): @pytest.mark.parametrize("name", sorted(CIRCUITS)) -def test_shared_measurement_sampling_matches_one_lane_per_shot(name): +def test_shared_measurement_sampling_matches_stabilizer_rank(name): circuit = Circuit(CIRCUITS[name]) shots = 20_000 - expected = circuit.compile_sampler(strategy="active_statevector", seed=1).sample( + expected = circuit.compile_sampler(strategy="cat5", seed=1).sample( shots, batch_size=5_000 ) sampler = circuit.compile_sampler( @@ -108,9 +148,7 @@ def test_pruned_postselection_keeps_the_accepted_distribution(name): shots = 20_000 mask = np.ones(circuit.num_detectors, dtype=bool) rate, expected = _accepted( - circuit.compile_detector_sampler(strategy="active_statevector", seed=1), - shots, - mask, + circuit.compile_detector_sampler(strategy="cat5", seed=1), shots, mask ) shared_rate, got = _accepted( circuit.compile_detector_sampler( @@ -130,7 +168,7 @@ def test_shared_sampling_of_random_noisy_circuits(seed): stim_circuit = gen_stim_circuit(3, 60, p_t=3, p_x_error=0.5, p_r_z=0.5, seed=seed) circuit = Circuit.from_stim_program(stim_circuit) shots = 4_096 - expected = circuit.compile_sampler(strategy="active_statevector", seed=seed).sample( + expected = circuit.compile_sampler(strategy="cat5", seed=seed).sample( shots, batch_size=shots ) sampler = circuit.compile_sampler( @@ -139,8 +177,224 @@ def test_shared_sampling_of_random_noisy_circuits(seed): assert_counts_match(sampler.sample(shots, batch_size=shots), expected) +# NOTE: two identical, disconnected components, each with two measurements. +TWO_COMPONENTS = """ +H 0 +T 0 +H 0 +M 0 +H 0 +T 0 +H 0 +M 0 +H 1 +T 1 +H 1 +M 1 +H 1 +T 1 +H 1 +M 1 +""" + + +@pytest.mark.parametrize("sharing", [False, True]) +def test_components_draw_independent_random_numbers(sharing): + batch, batches = 64, 1_500 + sampler = Circuit(TWO_COMPONENTS).compile_sampler( + strategy="active_statevector", seed=3, trajectory_sharing=sharing + ) + assert len(sampler._program.components) == 2 + m = sampler.sample(batch * batches, batch_size=batch).astype(int) + first, second = m[:, :2], m[:, 2:] + shots = len(m) + for i in range(2): + for j in range(2): + correlation = np.corrcoef(first[:, i], second[:, j])[0, 1] + assert abs(correlation) < 5 / np.sqrt(shots) + # NOTE: shots in one batch must not share draws across components either. + per_batch = m.reshape(batches, batch, 4).sum(axis=1) + correlation = np.corrcoef(per_batch[:, :2].sum(1), per_batch[:, 2:].sum(1))[0, 1] + assert abs(correlation) < 5 / np.sqrt(batches) + + +def test_trajectory_sharing_is_the_statevector_default(): + circuit = Circuit(CIRCUITS["t_ghz"]) + assert circuit.compile_detector_sampler()._trajectory_sharing + assert circuit.compile_sampler()._trajectory_sharing + assert not circuit.compile_detector_sampler(strategy="cat5")._trajectory_sharing + + def test_trajectory_sharing_needs_the_statevector_strategy(): with pytest.raises(ValueError, match="trajectory_sharing"): Circuit(CIRCUITS["t_ghz"]).compile_detector_sampler( strategy="cat5", trajectory_sharing=True ) + + +@pytest.mark.parametrize("sharing", [False, True]) +@pytest.mark.parametrize("screen_shots", [None, 64]) +@pytest.mark.parametrize("postselect", [False, True]) +@pytest.mark.parametrize("reference", [False, True]) +def test_samples_do_not_depend_on_the_thread_count( + postselect, reference, screen_shots, sharing, monkeypatch +): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + mask = np.ones(circuit.num_detectors, dtype=bool) if postselect else None + if screen_shots is not None: + monkeypatch.setattr(sampler_module, "SCREEN_SHOTS", screen_shots) + + def sample(threads): + sampler = circuit.compile_detector_sampler( + strategy="active_statevector", seed=3, trajectory_sharing=sharing + ) + sampler._threads = lambda batch_size: threads + return [ + sampler.sample( + 5_003, + batch_size=300, + append_observables=True, + postselection_mask=mask, + use_detector_reference_sample=reference, + ) + for _ in range(2) + ] + + # NOTE: the serial baseline runs numba's parallel kernels, which the worker + # threads replace by serial ones. + with monkeypatch.context() as patch: + patch.setattr(kernels, "PARALLEL_WORK", 0) + expected = sample(1) + for threads in (2, 5): + for got, want in zip(sample(threads), expected, strict=True): + np.testing.assert_array_equal(got, want) + + +def test_shared_measurement_samples_do_not_depend_on_the_thread_count(): + circuit = Circuit(CIRCUITS["midcircuit_feedforward"]) + + def sample(threads): + sampler = circuit.compile_sampler( + strategy="active_statevector", seed=4, trajectory_sharing=True + ) + sampler._threads = lambda batch_size: threads + return sampler.sample(2_001, batch_size=250) + + np.testing.assert_array_equal(sample(4), sample(1)) + + +def _recording_rows(sampler, monkeypatch) -> list[int]: + rows = [] + original = sampler_module.sample_statevector_program + + def spy(program, f_params, *args, **kwargs): + assert isinstance(f_params, np.ndarray) + rows.append(f_params.shape[0]) + return original(program, f_params, *args, **kwargs) + + monkeypatch.setattr(sampler_module, "sample_statevector_program", spy) + return rows + + +def test_shared_postselection_samples_the_survivors_without_padding(monkeypatch): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + sampler = circuit.compile_detector_sampler( + strategy="active_statevector", seed=5, trajectory_sharing=True + ) + sampler._threads = lambda batch_size: 1 + counts = _recording_rows(sampler, monkeypatch) + mask = np.ones(circuit.num_detectors, dtype=bool) + detectors = sampler.sample(5_003, batch_size=300, postselection_mask=mask) + + assert counts[:-1] == [300] * (len(counts) - 1) + assert 0 < counts[-1] < 300 + # NOTE: detector 2 is the circuit's only direct detector. + assert sum(counts) == int((~detectors[:, 2]).sum()) + + +@pytest.mark.parametrize("sharing", [False, True]) +@pytest.mark.parametrize( + ("shots", "expected"), + [(40_000, [10_000] * 4), (12_000, [12_000]), (6_000, [6_000])], +) +def test_default_batches_give_every_thread_one(sharing, shots, expected, monkeypatch): + sampler = Circuit(GHZ_WITH_DIRECT_DETECTOR).compile_detector_sampler( + strategy="active_statevector", seed=7, trajectory_sharing=sharing + ) + sampler._threads = lambda batch_size: 4 + rows = _recording_rows(sampler, monkeypatch) + assert sampler.sample(shots).shape[0] == shots + assert sorted(rows) == expected + + +@pytest.mark.parametrize("sharing", [False, True]) +def test_default_postselection_splits_the_survivors_over_the_threads( + sharing, monkeypatch +): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + sampler = circuit.compile_detector_sampler( + strategy="active_statevector", seed=8, trajectory_sharing=sharing + ) + sampler._threads = lambda batch_size: 4 + rows = _recording_rows(sampler, monkeypatch) + mask = np.ones(circuit.num_detectors, dtype=bool) + detectors = sampler.sample(40_000, postselection_mask=mask) + + survivors = int((~detectors[:, 2]).sum()) + size = -(-survivors // 4) + assert len(rows) == 4 + assert sum(rows) == survivors and max(rows) == size + + +@pytest.mark.parametrize("rows", [1, 300, 70_000]) +@pytest.mark.parametrize("name", [*sorted(CIRCUITS), "direct_flips"]) +def test_direct_outputs_land_in_their_columns(name, rows): + text = DIRECT_FLIPS if name == "direct_flips" else CIRCUITS[name] + sampler = Circuit(text).compile_detector_sampler( + strategy="active_statevector", seed=1 + ) + program = sampler._program + if name == "direct_flips": + assert program.direct_flips.any() + assert program.output_reindex is not None + f_params = sampler._channel_sampler.sample(rows) + samples = sampler._backend.sample(f_params, 2) + columns = program.output_order[: len(program.direct_f_indices)] + expected = (f_params[:, program.direct_f_indices] != 0) ^ program.direct_flips + np.testing.assert_array_equal(samples[:, columns], expected) + + +def test_masks_of_component_detectors_alone_prune_lanes(): + sampler = Circuit(GHZ_WITH_DIRECT_DETECTOR).compile_detector_sampler( + strategy="active_statevector", seed=1 + ) + sampler.sample(20_000, postselection_mask=np.array([True, True, False])) + assert any(stats["dead"] > 0 for stats in _sharing_stats(sampler)) + + +@pytest.mark.parametrize("sharing", [False, True]) +def test_reference_bits_leave_the_outputs_of_pruned_shots_false(sharing, monkeypatch): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + mask = np.ones(circuit.num_detectors, dtype=bool) + + def sample(reference): + sampler = circuit.compile_detector_sampler( + strategy="active_statevector", seed=4, trajectory_sharing=sharing + ) + monkeypatch.setattr(sampler, "_compute_reference_sample", lambda: reference) + return sampler.sample( + 20_000, + postselection_mask=mask, + separate_observables=True, + use_observable_reference_sample=True, + ) + + detectors, raw = sample(np.zeros(4, dtype=bool)) + _, flipped = sample(np.array([False, False, False, True])) + raw, flipped = raw[:, 0], flipped[:, 0] + kept = ~detectors.any(axis=1) + np.testing.assert_array_equal(flipped[kept], ~raw[kept]) + # NOTE: a shot pruned before the observable is set reads False either way. + cleared = ~raw & ~flipped + assert cleared.any() and not cleared[kept].any() + np.testing.assert_array_equal(cleared | (flipped == ~raw), np.ones_like(kept)) diff --git a/test/unit/noise/test_channels.py b/test/unit/noise/test_channels.py index 162f75a3..35708b5f 100644 --- a/test/unit/noise/test_channels.py +++ b/test/unit/noise/test_channels.py @@ -3,8 +3,11 @@ from numpy.testing import assert_allclose from tsim.noise.channels import ( + _BLOCK_SHOTS, Channel, ChannelSampler, + _fire_tables, + _sample_blocks, absorb_subset_channels, channel_num_bits, correlated_error_probs, @@ -18,6 +21,7 @@ pauli_channel_2_probs, reduce_null_bits, simplify_channels, + unique_columns, xor_convolve, ) @@ -213,9 +217,7 @@ def _sample_channels(channels, matrix, n_samples, seed=42): sampler.channels = channels sampler.signature_matrix = np.asarray(matrix, dtype=np.uint8) sampler._rng = np.random.default_rng(seed) - sampler._sparse_data = ChannelSampler._precompute_sparse( - channels, sampler.signature_matrix - ) + sampler._tables = _fire_tables(channels, sampler.signature_matrix) return sampler.sample(n_samples) @@ -751,6 +753,69 @@ def test_absorb_duplicate_subset_channel_uses_xor_expansion(self): assert_allclose(absorbed[0].probs, [1.0, 0.0, 0.0, 0.0]) +def _absorb_by_comparing_every_pair(channels, max_bits): + """The original quadratic ``absorb_subset_channels``, as a reference.""" + channels = sorted(channels, key=lambda c: -len(c.unique_col_ids)) + result, absorbed = [], set() + for i, channel_i in enumerate(channels): + if i in absorbed: + continue + set_i = set(channel_i.unique_col_ids) + probs = channel_i.probs.copy() + for j, channel_j in enumerate(channels): + if j <= i or j in absorbed: + continue + if set(channel_j.unique_col_ids) < set_i and len(set_i) <= max_bits: + expanded = expand_channel(channel_j, channel_i.unique_col_ids) + probs = xor_convolve(probs, expanded.probs) + absorbed.add(j) + result.append(Channel(probs=probs, unique_col_ids=channel_i.unique_col_ids)) + return result + + +class TestAbsorbMatchesPairwiseComparison: + """The indexed absorption gives the pairwise algorithm's channels bitwise.""" + + @pytest.mark.parametrize("max_bits", [2, 4]) + @pytest.mark.parametrize("seed", range(5)) + def test_random_channels(self, seed, max_bits): + rng = np.random.default_rng(seed) + channels = [] + for _ in range(60): + size = int(rng.integers(0, 6)) + col_ids = tuple(sorted(rng.choice(7, size=size, replace=False).tolist())) + probs = rng.random(2**size) + channels.append(Channel(probs=probs / probs.sum(), unique_col_ids=col_ids)) + got = absorb_subset_channels(channels, max_bits) + expected = _absorb_by_comparing_every_pair(channels, max_bits) + assert [c.unique_col_ids for c in got] == [c.unique_col_ids for c in expected] + for a, b in zip(got, expected, strict=True): + np.testing.assert_array_equal(a.probs, b.probs) + + +class TestUniqueColumns: + """unique_columns returns exactly what np.unique(axis=1) returns.""" + + @pytest.mark.parametrize("rows", [0, 1, 7, 8, 9, 64, 100]) + @pytest.mark.parametrize("cols", [0, 1, 5, 200]) + def test_matches_numpy(self, rows, cols): + rng = np.random.default_rng(rows * 1000 + cols) + distinct = (rng.random((rows, max(1, cols // 3))) < 0.2).astype(np.uint8) + matrix = distinct[:, rng.integers(0, distinct.shape[1], cols)] + got_cols, got_inverse = unique_columns(matrix) + expected_cols, expected_inverse = np.unique(matrix, axis=1, return_inverse=True) + np.testing.assert_array_equal(got_cols, expected_cols) + np.testing.assert_array_equal(got_inverse, np.ravel(expected_inverse)) + + def test_accepts_float_matrices(self): + matrix = np.array([[1.0, 0.0, 1.0], [0.0, 2.0, 0.0]]) + got_cols, got_inverse = unique_columns(matrix) + np.testing.assert_array_equal(got_cols, [[0.0, 1.0], [2.0, 0.0]]) + np.testing.assert_array_equal(got_inverse, [1, 0, 1]) + sampler = ChannelSampler([error_probs(1)], np.eye(1)) + np.testing.assert_array_equal(sampler.sample(3), np.ones((3, 1))) + + class TestSimplifyChannels: """Tests for the full simplify_channels function.""" @@ -1047,6 +1112,61 @@ def test_channel_simplification(self): expected = 3 * 0.1 * 0.9**2 + 0.1**3 assert_allclose(freq, expected, rtol=0.05) + def test_two_bit_channel_joint_distribution_across_blocks(self): + """Every outcome of a 2-bit channel has its probability, in every block.""" + probs = np.array([0.5, 0.2, 0.25, 0.05]) + transform = np.eye(2, dtype=np.uint8) + sampler = ChannelSampler([probs], transform, seed=3) + num_samples = 5 * _BLOCK_SHOTS + 123 + samples = sampler.sample(num_samples) + outcome = samples[:, 0] + 2 * samples[:, 1] + counts = np.bincount(outcome, minlength=4) + sigma = np.sqrt(num_samples * probs * (1 - probs)) + assert np.all(np.abs(counts - num_samples * probs) <= 5 * sigma) + + def test_blocks_draw_independent_streams(self): + """Every block of shots has its own random stream.""" + sampler = ChannelSampler( + [error_probs(0.5)], np.ones((1, 1), dtype=np.uint8), seed=1 + ) + blocks = sampler.sample(3 * _BLOCK_SHOTS).reshape(3, _BLOCK_SHOTS) + assert not np.array_equal(blocks[0], blocks[1]) + assert not np.array_equal(blocks[1], blocks[2]) + + def test_identity_rounded_below_zero_still_fires_every_row(self): + """``1 - probs[0]`` above 1 by a rounding error means a certain error.""" + probs = np.array([-1.25e-16, 0.33, 0.11, 0.56]) + sampler = ChannelSampler([probs], np.eye(2, dtype=np.uint8), seed=2) + num_samples = 100_000 + rates = sampler.sample(num_samples).mean(axis=0) + expected = np.array([0.33 + 0.56, 0.11 + 0.56]) + sigma = np.sqrt(expected * (1 - expected) / num_samples) + assert np.all(np.abs(rates - expected) <= 5 * sigma) + + def test_certain_error_flips_every_row(self): + """A channel that always fires reaches every row, last partial block too.""" + sampler = ChannelSampler( + [error_probs(1.0)], np.ones((1, 1), dtype=np.uint8), seed=0 + ) + assert np.all(sampler.sample(2 * _BLOCK_SHOTS + 7) == 1) + + def test_seed_reproduces_samples(self): + """Equal seeds give equal samples; consecutive calls differ.""" + probs = [error_probs(0.1), error_probs(0.2)] + transform = np.array([[1, 0], [1, 1]], dtype=np.uint8) + first = ChannelSampler(probs, transform, seed=5) + second = ChannelSampler(probs, transform, seed=5) + a, b = first.sample(10_000), second.sample(10_000) + assert np.array_equal(a, b) + assert not np.array_equal(first.sample(10_000), a) + + def test_zero_samples(self): + """Sampling no shots returns an empty array with one column per f.""" + sampler = ChannelSampler( + [error_probs(0.1)], np.ones((1, 1), dtype=np.uint8), seed=0 + ) + assert sampler.sample(0).shape == (0, 1) + def test_empty_transform(self): """Test with no f-variables (empty transform).""" probs = [error_probs(0.1)] @@ -1056,3 +1176,40 @@ def test_empty_transform(self): samples = sampler.sample(100) assert samples.shape == (100, 0) + + def test_drawn_block_seeds_reproduce_samples(self): + """sample(n) equals sample(n, block_seeds(n)), and seed draws split freely.""" + probs = [error_probs(0.1), error_probs(0.2)] + transform = np.array([[1, 0], [1, 1]], dtype=np.uint8) + num_samples = 3 * _BLOCK_SHOTS + 5 + expected = ChannelSampler(probs, transform, seed=8).sample(num_samples) + drawn = ChannelSampler(probs, transform, seed=8) + seeds = drawn.block_seeds(num_samples) + assert np.array_equal(drawn.sample(num_samples, seeds), expected) + + whole = ChannelSampler(probs, transform, seed=9).block_seeds(5 * _BLOCK_SHOTS) + split = ChannelSampler(probs, transform, seed=9) + parts = [ + split.block_seeds(2 * _BLOCK_SHOTS), + split.block_seeds(3 * _BLOCK_SHOTS), + ] + assert np.array_equal(np.concatenate(parts), whole) + + def test_wrong_number_of_block_seeds_raises(self): + sampler = ChannelSampler( + [error_probs(0.1)], np.ones((1, 1), dtype=np.uint8), seed=0 + ) + with pytest.raises(ValueError, match="block seeds"): + sampler.sample(_BLOCK_SHOTS + 1, sampler.block_seeds(_BLOCK_SHOTS)) + + def test_sampled_blocks_clear_their_rows(self): + """Sampling overwrites whatever the output array held before.""" + probs = [error_probs(0.1), error_probs(0.2)] + transform = np.array([[1, 0], [1, 1]], dtype=np.uint8) + num_samples = 2 * _BLOCK_SHOTS + 9 + sampler = ChannelSampler(probs, transform, seed=4) + seeds = sampler.block_seeds(num_samples) + expected = sampler.sample(num_samples, seeds) + out = np.ones_like(expected) + _sample_blocks(0, out, seeds, *sampler._tables) + assert np.array_equal(out, expected) diff --git a/test/unit/statevector/test_execute.py b/test/unit/statevector/test_execute.py deleted file mode 100644 index cb30ea1f..00000000 --- a/test/unit/statevector/test_execute.py +++ /dev/null @@ -1,346 +0,0 @@ -from test.helpers.statevector import ( - CIRCUITS, - apply_pauli, - assert_matches_distribution, - f_patterns, - output_indices, - pauli_matrix, - plan_distribution, - quantum_components, - reference_measure, -) - -import jax -import jax.numpy as jnp -import numpy as np -import pytest - -import tsim.statevector.execute as execute -from tsim.statevector.execute import ( - initial_state, - outputs, - pack_bits, - run, - sample_statevector_component, -) -from tsim.statevector.pbc import ComponentLayout -from tsim.statevector.plan import AffineRow, MeasureActive, MeasureDormant, Plan, Rotate -from tsim.statevector.program import ( - MEASURE, - MEASURE_DIAG, - MEASURE_DORMANT, - compile_statevector_component, - lower, - plan_component, -) - - -def _component(actions, *, width, num_f=1, outputs=(), dtype="complex64"): - planned = Plan( - width=width, num_f=num_f, actions=tuple(actions), outputs=tuple(outputs) - ) - layout = ComponentLayout(tuple(range(num_f)), {}, tuple(range(len(outputs)))) - return lower(planned, layout, dtype) - - -def _random_states(lanes, width, seed=0): - rng = np.random.default_rng(seed) - psi = rng.normal(size=(lanes, 2**width)) + 1j * rng.normal(size=(lanes, 2**width)) - return psi / np.linalg.norm(psi, axis=1, keepdims=True) - - -def _measuring(x, z, pivot, psi, f_bits=None): - """Return a component that measures ``x``, ``z``, and lanes in the states ``psi``. - - Its one output is the outcome XOR the f-bit. - """ - component = _component( - [MeasureActive(x, z, pivot, branch=1)], - width=3, - outputs=[AffineRow(False, frozenset({0, 1}))], - ) - if f_bits is None: - f_bits = jnp.zeros((len(psi), 1), dtype=bool) - amps = jnp.asarray(psi, jnp.complex64) - return component, initial_state(component, f_bits)._replace(amps=amps) - - -def _only_component(name): - (component,), layout = quantum_components(CIRCUITS[name], sample_detectors=True) - return component, output_indices(component), layout - - -def test_pack_bits(): - bits = np.zeros((2, 40), dtype=bool) - bits[0, [0, 31, 32, 39]] = True - bits[1, 5] = True - words = np.asarray(pack_bits(jnp.asarray(bits), 2)) - assert words.dtype == np.uint32 - assert words.tolist() == [[1 | (1 << 31), 1 | (1 << 7)], [1 << 5, 0]] - - -@pytest.mark.parametrize("dtype", ["complex64", "complex128"]) -@pytest.mark.parametrize("x, z", [(0, 0b101), (0b011, 0), (0b110, 0b011), (1, 1)]) -def test_rotations_match_dense_matrices(x, z, dtype): - with jax.enable_x64(dtype == "complex128"): - sign = AffineRow(False, frozenset({0})) - component = _component([Rotate(x, z, 0.37, sign)], width=3, dtype=dtype) - psi = _random_states(2, 3) - f_bits = jnp.array([[False], [True]]) - amps = jnp.asarray(psi, jnp.dtype(dtype)) - state = initial_state(component, f_bits)._replace(amps=amps) - out = run(component, state, jax.random.key(0), 0, 1) - assert out.amps.dtype == jnp.dtype(dtype) - pauli = pauli_matrix(x, z, 3) - atol = 1e-5 if dtype == "complex64" else 1e-12 - for lane, negate in ((0, 1), (1, -1)): - rotation = np.cos(0.37) * np.eye(8) - 1j * np.sin(0.37) * negate * pauli - np.testing.assert_allclose( - np.asarray(out.amps[lane]), rotation @ psi[lane], atol=atol - ) - - -_MEASUREMENTS = [ - (0b011, 0b110, 0), - (0b010, 0b011, 1), - (0b110, 0b010, 2), - (0, 0b101, 2), - (0, 0b011, 0), -] - - -@pytest.mark.parametrize("x, z, pivot", _MEASUREMENTS) -def test_measurements_match_the_reference_kernel(monkeypatch, x, z, pivot): - psi = _random_states(3, 3) - f_bits = jnp.array([[False], [False], [True]]) - component, state = _measuring(x, z, pivot, psi, f_bits) - outcomes = jnp.array([False, True, True]) - monkeypatch.setattr(execute, "choose_outcome", lambda p0, p1, count, key: outcomes) - out = run(component, state, jax.random.key(0), 0, 1) - for lane, outcome in enumerate([0, 1, 1]): - expected, _ = reference_measure(psi[lane], x, z, pivot, outcome) - np.testing.assert_allclose(np.asarray(out.amps[lane]), expected, atol=1e-5) - assert np.asarray(outputs(component, out)[:, 0]).tolist() == [False, True, False] - assert ((np.asarray(out.symbols[:, 0]) >> 1) & 1).tolist() == [0, 1, 1] - - -def test_outputs_are_parities_of_the_symbols(monkeypatch): - component = _component( - [MeasureDormant(branch=1)], - width=0, - outputs=[AffineRow(True, frozenset({1})), AffineRow(False, frozenset({0, 1}))], - ) - monkeypatch.setattr( - execute, "choose_outcome", lambda p0, p1, count, key: jnp.array([False, True]) - ) - state = initial_state(component, jnp.array([[True], [True]])) - out = run(component, state, jax.random.key(0), 0, 1) - assert np.asarray(outputs(component, out)).tolist() == [ - [True, True], - [False, False], - ] - - -def test_outputs_read_symbols_beyond_the_first_word(monkeypatch): - component = _component( - [MeasureDormant(branch=40)], - width=0, - num_f=40, - outputs=[ - AffineRow(False, frozenset({40})), - AffineRow(True, frozenset({3, 40})), - ], - ) - assert component.num_words == 2 - monkeypatch.setattr( - execute, "choose_outcome", lambda p0, p1, count, key: jnp.array([False, True]) - ) - f_bits = np.zeros((2, 40), dtype=bool) - f_bits[1, 3] = True - state = initial_state(component, jnp.asarray(f_bits)) - out = run(component, state, jax.random.key(0), 0, 1) - assert np.asarray(outputs(component, out)).tolist() == [[False, True], [True, True]] - - -def test_a_component_without_actions_samples_no_outputs(): - component = _component([], width=0, num_f=0) - assert component.num_actions == 0 - state = initial_state(component, jnp.zeros((3, 0), dtype=bool)) - out = run(component, state, jax.random.key(0), 0, component.num_actions) - for got, initial in zip(out, state, strict=True): - np.testing.assert_array_equal(np.asarray(got), np.asarray(initial)) - samples, _, deviation = sample_statevector_component( - component, jnp.zeros((3, 0), dtype=jnp.uint8), jax.random.key(0) - ) - assert samples.shape == (3, 0) - assert float(deviation) == 0.0 - - -@pytest.mark.parametrize( - "name", - [ - "t_ghz", - "midcircuit_feedforward", - "two_components", - "clifford_component", - "rotations", - ], -) -def test_sampling_matches_the_plan_distribution(name): - shots = 20_000 - components, layout = quantum_components(CIRCUITS[name], sample_detectors=True) - for component in components: - indices = output_indices(component) - planned, component_layout = plan_component(component, indices, layout) - f_selection = component_layout.f_selection - compiled = lower(planned, component_layout, "complex64") - num_columns = max(f_selection, default=-1) + 1 - for f in f_patterns(planned.num_f, limit=2): - f_params = np.zeros((shots, num_columns), dtype=np.uint8) - if f_selection: - f_params[:, list(f_selection)] = f - samples, _, deviation = sample_statevector_component( - compiled, jnp.asarray(f_params), jax.random.key(7) - ) - assert float(deviation) < 1e-4 - assert_matches_distribution( - np.asarray(samples), plan_distribution(planned, f) - ) - - -@pytest.mark.parametrize("x64", [False, True]) -def test_amplitudes_are_complex128_with_64_bit_types(x64): - dtype = "complex128" if x64 else "complex64" - with jax.enable_x64(x64): - compiled = compile_statevector_component(*_only_component("t_ghz")) - assert compiled.dtype == dtype - num_f = compiled.f_selection.shape[0] - state = initial_state(compiled, jnp.zeros((4, num_f), dtype=bool)) - out = run(compiled, state, jax.random.key(0), 0, compiled.num_actions) - assert out.amps.dtype == jnp.dtype(dtype) - - -def _count_outcomes(p1, dtype, draws, chunk=2**22): - @jax.jit - def ones_in_chunk(key): - p0s = jnp.full(chunk, 1 - p1, dtype=dtype) - p1s = jnp.full(chunk, p1, dtype=dtype) - count = jnp.ones(chunk, dtype=jnp.int32) - return jnp.sum(execute.choose_outcome(p0s, p1s, count, key)) - - return sum( - int(ones_in_chunk(jax.random.key(seed))) for seed in range(draws // chunk) - ) - - -@pytest.mark.parametrize("p1", [0.0, 1e-15]) -def test_choose_outcome_never_returns_true_below_its_resolution(p1): - assert _count_outcomes(p1, jnp.float32, 2**26) == 0 - - -@pytest.mark.parametrize("dtype", ["float32", "float64"]) -def test_choose_outcome_always_returns_true_for_p1_one(dtype): - with jax.enable_x64(dtype == "float64"): - assert _count_outcomes(1.0, jnp.dtype(dtype), 2**22) == 2**22 - - -@pytest.mark.parametrize("dtype", ["float32", "float64"]) -@pytest.mark.parametrize("p1", [0.5, 0.3, 0.7, 1e-3, 1 - 1e-3, 1e-6]) -def test_choose_outcome_frequency_matches_p1(p1, dtype): - draws = 2**24 - with jax.enable_x64(dtype == "float64"): - ones = _count_outcomes(p1, jnp.dtype(dtype), draws) - expected = draws * p1 - assert abs(ones - expected) <= 5 * np.sqrt(draws * p1 * (1 - p1)) + 1 - - -def _projected(vectors, x, z, eigenvalue): - projected = np.array([v + eigenvalue * apply_pauli(v, x, z) for v in vectors]) - return projected / np.linalg.norm(projected, axis=1, keepdims=True) - - -def _run_recording_probabilities(monkeypatch, component, state): - """Run one measurement; return the ``(p0, p1)`` it sampled from and the state.""" - recorded = [] - choose = execute.choose_outcome - - def recording(p0, p1, count, key): - recorded.append((np.asarray(p0, np.float64), np.asarray(p1, np.float64))) - return choose(p0, p1, count, key) - - monkeypatch.setattr(execute, "choose_outcome", recording) - with jax.disable_jit(): - out = run(component, state, jax.random.key(3), 0, 1) - (probabilities,) = recorded - return probabilities, out - - -@pytest.mark.parametrize("eigenvalue", [1, -1]) -@pytest.mark.parametrize("x, z, pivot", _MEASUREMENTS) -def test_measuring_an_eigenstate_never_samples_the_impossible_outcome( - monkeypatch, x, z, pivot, eigenvalue -): - lanes = 4096 - # NOTE: the noise is of the size rotations leave in complex64 amplitudes. - rounding = 1e-7 * _random_states(lanes, 3, seed=1) - psi = _projected(_random_states(lanes, 3), x, z, eigenvalue) + rounding - component, state = _measuring(x, z, pivot, psi) - probabilities, out = _run_recording_probabilities(monkeypatch, component, state) - impossible = probabilities[1 if eigenvalue == 1 else 0] - assert impossible.max() < 1e-12 - assert np.all(np.asarray(outputs(component, out)[:, 0]) == (eigenvalue == -1)) - - -@pytest.mark.parametrize("unlikely", [0, 1]) -@pytest.mark.parametrize("x, z, pivot", _MEASUREMENTS) -def test_a_rare_outcome_keeps_its_probability_in_complex64( - monkeypatch, x, z, pivot, unlikely -): - # NOTE: 1 - p1 would round to 0 or a multiple of 6e-8 below these probabilities. - lanes = 64 - vectors = _random_states(lanes, 3) - likely = _projected(vectors, x, z, (-1) ** (1 - unlikely)) - rare = _projected(vectors, x, z, (-1) ** unlikely) - p = np.logspace(-9, -7, lanes)[:, None] - component, state = _measuring( - x, z, pivot, np.sqrt(1 - p) * likely + np.sqrt(p) * rare - ) - probabilities, _ = _run_recording_probabilities(monkeypatch, component, state) - np.testing.assert_allclose(probabilities[unlikely], p[:, 0], rtol=1e-2) - - -@pytest.mark.parametrize("unlikely", [0, 1]) -@pytest.mark.parametrize("x, z, pivot", _MEASUREMENTS) -def test_an_unlikely_outcome_leaves_a_normalized_state( - monkeypatch, x, z, pivot, unlikely -): - lanes = 64 - vectors = _random_states(lanes, 3) - likely = _projected(vectors, x, z, (-1) ** (1 - unlikely)) - rare = _projected(vectors, x, z, (-1) ** unlikely) - p = np.logspace(-5, -3, lanes)[:, None] - component, state = _measuring( - x, z, pivot, np.sqrt(1 - p) * likely + np.sqrt(p) * rare - ) - outcomes = jnp.full(lanes, bool(unlikely)) - monkeypatch.setattr(execute, "choose_outcome", lambda p0, p1, count, key: outcomes) - out = run(component, state, jax.random.key(0), 0, 1) - norms = np.linalg.norm(np.asarray(out.amps, dtype=np.complex128), axis=1) ** 2 - assert np.abs(norms - 1).max() < 1e-5 - - -def test_running_a_tape_in_two_parts_equals_running_it_whole(): - compiled = compile_statevector_component(*_only_component("midcircuit_feedforward")) - lanes, split, stop = 512, 4, compiled.num_actions - kinds = np.asarray(compiled.kind) - measurements = np.isin(kinds, [MEASURE_DIAG, MEASURE, MEASURE_DORMANT]) - assert measurements[:split].any() and measurements[split:].any() - rng = np.random.default_rng(0) - num_f = compiled.f_selection.shape[0] - f_bits = jnp.asarray(rng.integers(0, 2, (lanes, num_f)).astype(bool)) - state = initial_state(compiled, f_bits) - key = jax.random.key(5) - whole = run(compiled, state, key, 0, stop) - parts = run(compiled, run(compiled, state, key, 0, split), key, split, stop) - for got, expected in zip(parts, whole, strict=True): - np.testing.assert_array_equal(np.asarray(got), np.asarray(expected)) diff --git a/test/unit/statevector/test_program.py b/test/unit/statevector/test_program.py index c68aaf3c..77b51cce 100644 --- a/test/unit/statevector/test_program.py +++ b/test/unit/statevector/test_program.py @@ -69,28 +69,29 @@ def test_lower_encodes_every_action(): ), ) layout = ComponentLayout(f_selection=(4,), f_position={}, output_indices=(2, 5)) - compiled = lower(planned, layout, "complex64") - assert compiled.kind.tolist() == [ + compiled = lower(planned, layout) + tape = compiled.tape + assert tape.kind.tolist() == [ ROTATE_DIAG, ROTATE, MEASURE_DIAG, MEASURE, MEASURE_DORMANT, ] - assert compiled.x_mask.tolist() == [0, 1, 0, 1, 0] - assert compiled.z_mask.tolist() == [3, 2, 2, 0, 0] - np.testing.assert_allclose(compiled.cos[:2], np.cos([0.25, -0.5]), rtol=1e-6) - np.testing.assert_allclose(compiled.sin[:2], np.sin([0.25, -0.5]), rtol=1e-6) - assert compiled.pivot.tolist()[2:4] == [1, 0] - assert compiled.branch.tolist() == [-1, -1, 1, 2, 3] - assert compiled.sign_const.tolist()[:2] == [True, False] - assert compiled.sign_row[:2, 0].tolist() == [1, 0] - assert compiled.output_const.tolist() == [False, True] - assert compiled.output_row[:, 0].tolist() == [0b1010, 0b0101] + assert tape.x_mask.tolist() == [0, 1, 0, 1, 0] + assert tape.z_mask.tolist() == [3, 2, 2, 0, 0] + np.testing.assert_allclose(tape.cos[:2], np.cos([0.25, -0.5])) + np.testing.assert_allclose(tape.sin[:2], np.sin([0.25, -0.5])) + assert tape.pivot.tolist()[2:4] == [1, 0] + assert tape.branch.tolist() == [-1, -1, 1, 2, 3] + assert tape.sign_const.tolist()[:2] == [True, False] + assert tape.sign_row[:2, 0].tolist() == [1, 0] + assert tape.output_const.tolist() == [False, True] + assert tape.output_row[:, 0].tolist() == [0b1010, 0b0101] assert compiled.num_words == 1 - assert compiled.f_selection.tolist() == [4] + assert tape.f_selection.tolist() == [4] assert compiled.output_indices == (2, 5) - assert (compiled.width, compiled.num_outputs, compiled.dtype) == (2, 2, "complex64") + assert (compiled.width, compiled.num_outputs) == (2, 2) assert compiled.num_actions == 5 @@ -101,9 +102,9 @@ def test_lower_packs_symbols_beyond_32_into_more_words(): actions=(), outputs=(AffineRow(False, frozenset({0, 33, 39})),), ) - compiled = lower(planned, ComponentLayout(tuple(range(40)), {}, (0,)), "complex64") + compiled = lower(planned, ComponentLayout(tuple(range(40)), {}, (0,))) assert compiled.num_words == 2 - assert compiled.output_row.tolist() == [[1, (1 << 1) | (1 << 7)]] + assert compiled.tape.output_row.tolist() == [[1, (1 << 1) | (1 << 7)]] def test_too_wide_components_are_rejected(monkeypatch): diff --git a/test/unit/statevector/test_shared.py b/test/unit/statevector/test_shared.py index b77b4bfb..4f1a81f2 100644 --- a/test/unit/statevector/test_shared.py +++ b/test/unit/statevector/test_shared.py @@ -1,3 +1,4 @@ +from dataclasses import replace from test.helpers.statevector import ( CIRCUITS, assert_matches_distribution, @@ -7,18 +8,14 @@ quantum_components, ) -import jax -import jax.numpy as jnp import numpy as np import pytest +import stim import tsim.statevector.shared as shared -from tsim.statevector.program import lower, plan_component -from tsim.statevector.shared import ( - Sharing, - sample_shared_component, - sharing_layout, -) +import tsim.utils.kernels as kernels +from tsim.statevector.program import MEASURE, MEASURE_DIAG, lower, plan_component +from tsim.statevector.shared import sample_shared_component, sharing_layout NAMES = [ "t_ghz", @@ -32,11 +29,15 @@ def _compiled(name): - components, layout = quantum_components(CIRCUITS[name], sample_detectors=True) + yield from _compiled_text(CIRCUITS[name]) + + +def _compiled_text(text): + components, layout = quantum_components(text, sample_detectors=True) for component in components: indices = output_indices(component) planned, component_layout = plan_component(component, indices, layout) - yield planned, component_layout, lower(planned, component_layout, "complex64") + yield planned, component_layout, lower(planned, component_layout) def _f_params(component_layout, patterns, shots_each): @@ -77,20 +78,23 @@ def _assert_counts(rows, discarded, expected): ) +@pytest.mark.parametrize("share", [True, False]) @pytest.mark.parametrize("name", NAMES) -def test_shared_sampling_matches_the_plan_distribution(name): +def test_sampling_matches_the_plan_distribution(name, share): shots = 20_000 for planned, component_layout, compiled in _compiled(name): - sharing = Sharing(sharing_layout(compiled)) + layout = sharing_layout(compiled) for f in f_patterns(planned.num_f, limit=2): f_params = _f_params(component_layout, [f], shots) - samples, _, deviation, _ = sample_shared_component( - compiled, sharing, jnp.asarray(f_params), jax.random.key(7) + samples, deviation, _ = sample_shared_component( + compiled, layout, np.asarray(f_params), 7, share ) assert float(deviation) < 1e-4 - assert_matches_distribution( - np.asarray(samples), plan_distribution(planned, f) - ) + distribution = plan_distribution(planned, f) + assert_matches_distribution(np.asarray(samples), distribution) + # NOTE: an outcome of probability 0 is never drawn, not even rarely. + seen = {tuple(row) for row in np.asarray(samples).astype(int).tolist()} + assert all(distribution.get(key, 0.0) > 1e-12 for key in seen) @pytest.mark.parametrize("name", NAMES) @@ -100,11 +104,11 @@ def test_mixed_noise_patterns_share_one_batch(name): patterns = f_patterns(planned.num_f, limit=3) f_params = _f_params(component_layout, patterns, shots_each) permutation = np.random.default_rng(1).permutation(len(f_params)) - samples, _, _, _ = sample_shared_component( + samples, _, _ = sample_shared_component( compiled, - Sharing(sharing_layout(compiled)), - jnp.asarray(f_params[permutation]), - jax.random.key(3), + sharing_layout(compiled), + np.asarray(f_params[permutation]), + 3, ) samples = np.asarray(samples)[np.argsort(permutation)] for i, f in enumerate(patterns): @@ -112,139 +116,208 @@ def test_mixed_noise_patterns_share_one_batch(name): assert_matches_distribution(rows, plan_distribution(planned, f)) +@pytest.mark.parametrize("share", [True, False]) @pytest.mark.parametrize("flip", [False, True]) @pytest.mark.parametrize("name", NAMES) -def test_pruning_discards_exactly_the_postselected_failures(name, flip): +def test_pruning_discards_exactly_the_postselected_failures(name, flip, share): shots_each = 8_000 for planned, component_layout, compiled in _compiled(name): postselected = np.ones(compiled.num_outputs, dtype=bool) postselected[-1] = False flipped = np.zeros(compiled.num_outputs, dtype=bool) flipped[0] = flip - sharing = Sharing(sharing_layout(compiled, postselected, flipped)) + layout = sharing_layout(compiled, postselected, flipped) patterns = f_patterns(planned.num_f, limit=3) f_params = _f_params(component_layout, patterns, shots_each) - samples, _, _, _ = sample_shared_component( - compiled, sharing, jnp.asarray(f_params), jax.random.key(5) + samples, _, _ = sample_shared_component( + compiled, layout, np.asarray(f_params), 5, share ) samples = np.asarray(samples) for i, f in enumerate(patterns): rows = samples[i * shots_each : (i + 1) * shots_each] - discarded = np.asarray(((rows ^ flipped) & postselected).any(axis=1)) + discarded = np.asarray((rows & postselected).any(axis=1)) expected = _postselected_distribution( plan_distribution(planned, f), postselected, flipped ) - _assert_counts(rows, discarded, expected) + _assert_counts(rows ^ flipped, discarded, expected) +@pytest.mark.parametrize("flip", [False, True]) @pytest.mark.parametrize("name", NAMES) -def test_pruned_shots_report_false_after_their_pruning_action(name): +def test_pruned_shots_report_false_after_their_pruning_action(name, flip): for planned, component_layout, compiled in _compiled(name): postselected = np.ones(compiled.num_outputs, dtype=bool) - sharing = Sharing(sharing_layout(compiled, postselected)) + flipped = np.full(compiled.num_outputs, flip) + layout = sharing_layout(compiled, postselected, flipped) f_params = _f_params(component_layout, f_patterns(planned.num_f, 3), 4_000) - samples, _, _, _ = sample_shared_component( - compiled, sharing, jnp.asarray(f_params), jax.random.key(9) + samples, _, _ = sample_shared_component( + compiled, layout, np.asarray(f_params), 9 ) rows = np.asarray(samples) - ready = np.asarray(sharing.layout.ready) - tested = np.asarray(sharing.layout.prune_outputs) + ready = layout.ready + tested = layout.prune_outputs prune_at = np.where(tested.any(axis=0), tested.argmax(axis=0), ready) pruning = np.min(np.where(rows, prune_at[None, :], np.inf), axis=1) late = rows & (ready[None, :] > pruning[:, None]) assert not np.any(late), "an output set after the pruning action is 1" -def test_full_tables_grow_and_resume(monkeypatch): - monkeypatch.setattr(shared, "_MIN_CAPACITY", 1) - shots_each = 5_000 - stats = [] - for name in ["t_ghz", "midcircuit_feedforward", "injection_check"]: - for planned, component_layout, compiled in _compiled(name): - postselected = np.zeros(compiled.num_outputs, dtype=bool) - postselected[0] = True - sharing = Sharing(sharing_layout(compiled, postselected)) - patterns = f_patterns(planned.num_f, limit=3) - f_params = _f_params(component_layout, patterns, shots_each) - samples, _, _, stat = sample_shared_component( - compiled, sharing, jnp.asarray(f_params), jax.random.key(11) - ) - stats.append(stat) - samples = np.asarray(samples) - for i, f in enumerate(patterns): - rows = samples[i * shots_each : (i + 1) * shots_each] - discarded = rows[:, 0] == 1 - expected = _postselected_distribution( - plan_distribution(planned, f), postselected - ) - kept = rows[~discarded] - p_kept = 1 - expected.get(("discarded",), 0.0) - sigma = np.sqrt(shots_each * p_kept * (1 - p_kept)) - assert abs(len(kept) - shots_each * p_kept) <= 5 * sigma + 3 - assert any(stat["resumes"] > 0 for stat in stats) - - @pytest.mark.parametrize("name", ["t_ghz", "midcircuit_feedforward", "injection_check"]) -def test_capacities_and_segments_do_not_change_the_samples(name, monkeypatch): +def test_equal_keys_give_equal_samples(name): for planned, component_layout, compiled in _compiled(name): postselected = np.zeros(compiled.num_outputs, dtype=bool) postselected[0] = True layout = sharing_layout(compiled, postselected) - f_params = jnp.asarray( + f_params = np.asarray( _f_params(component_layout, f_patterns(planned.num_f, 3), 3_000) ) - key = jax.random.key(13) - reference, _, _, wide = sample_shared_component( - compiled, Sharing(layout), f_params, key - ) - segmented, _, _, _ = sample_shared_component( - compiled, Sharing(layout), f_params, key, segment=1 - ) - monkeypatch.setattr(shared, "_MIN_CAPACITY", 1) - tiny, _, _, narrow = sample_shared_component( - compiled, Sharing(layout), f_params, key, segment=1 - ) - monkeypatch.undo() - np.testing.assert_array_equal(np.asarray(segmented), np.asarray(reference)) - np.testing.assert_array_equal(np.asarray(tiny), np.asarray(reference)) - assert narrow["resumes"] > 0 or narrow["lane_capacity"] != wide["lane_capacity"] - - -def test_rare_outcomes_keep_their_probability(): - draw = jax.jit(shared._rare_drawn) - for count, q in [(1, 1e-9), (1, 3e-8), (7, 1e-7), (1000, 1e-10)]: - p_any = -np.expm1(count * np.log1p(-q)) - exact = -np.expm1(count * np.log1p(np.float32(-q))) - assert abs(p_any - exact) <= 1e-6 * exact - slots = 1 << 21 + first, _, _ = sample_shared_component(compiled, layout, f_params, 13) + second, _, _ = sample_shared_component(compiled, layout, f_params, 13) + third, _, _ = sample_shared_component(compiled, layout, f_params, 14) + np.testing.assert_array_equal(first, second) + assert not np.array_equal(first, third) + + +def test_rare_outcomes_keep_their_binomial_law(): + lanes = 1 << 21 for count, q in [(1, 1e-4), (3, 2e-3), (50, 0.3), (7, 0.5)]: - rare = np.asarray( - draw( - jax.random.key(count), - jnp.full(slots, count, dtype=jnp.int32), - jnp.full(slots, q, dtype=jnp.float32), - ) - ) + w1 = np.full(lanes, q) + ones = shared._draw_ones(np.full(lanes, count), 1 - w1, w1, np.uint64(count)) p_any = 1 - (1 - q) ** count - hits = np.count_nonzero(rare) - assert abs(hits - slots * p_any) <= 5 * np.sqrt(slots * p_any) + 3 - mean, sigma = count * q, np.sqrt(count * q * (1 - q) / slots) - assert abs(rare.mean() - mean) <= 5 * sigma + 1e-6 - assert rare.min() >= 0 and rare.max() <= count + hits = np.count_nonzero(ones) + assert abs(hits - lanes * p_any) <= 5 * np.sqrt(lanes * p_any) + 3 + mean, sigma = count * q, np.sqrt(count * q * (1 - q) / lanes) + assert abs(ones.mean() - mean) <= 5 * sigma + 1e-6 + assert ones.min() >= 0 and ones.max() <= count + + +def test_the_common_outcome_counts_the_shots_that_miss_the_rare_one(): + lanes = 1 << 18 + count, q = 40, 0.01 + w1 = np.full(lanes, 1 - q) + ones = shared._draw_ones(np.full(lanes, count), 1 - w1, w1, np.uint64(3)) + zeros = count - ones + assert abs(zeros.mean() - count * q) <= 5 * np.sqrt(count * q / lanes) def test_a_lone_shot_sees_an_outcome_of_probability_below_float32_resolution(): q = 1e-7 - slots = 1 << 22 + lanes = 1 << 22 hits = 0 - draw = jax.jit(shared._rare_drawn) for seed in range(24): - rare = draw( - jax.random.key(seed), - jnp.ones(slots, dtype=jnp.int32), - jnp.full(slots, q, dtype=jnp.float32), + w1 = np.full(lanes, q) + ones = shared._draw_ones( + np.ones(lanes, dtype=np.int64), 1 - w1, w1, np.uint64(seed) ) - hits += int(jnp.count_nonzero(rare)) - expected = 24 * slots * q + hits += int(np.count_nonzero(ones)) + expected = 24 * lanes * q assert hits > 0 assert abs(hits - expected) <= 5 * np.sqrt(expected) + 3 + + +def _coherent_surface_code(distance, rounds, angle=0.02): + """Return a surface-code memory with an R_Z over-rotation after every CX layer.""" + base = stim.Circuit.generated( + "surface_code:rotated_memory_z", + distance=distance, + rounds=rounds, + after_clifford_depolarization=1e-3, + ) + lines = [] + for line in str(base.flattened()).splitlines(): + lines.append(line) + if line.startswith("CX "): + qubits = sorted({int(t) for t in line.split()[1:]}) + lines.append(f"R_Z({angle}) " + " ".join(map(str, qubits))) + return "\n".join(lines) + + +@pytest.fixture(scope="module") +def coherent(): + """A component whose extent drops and later rises again.""" + ((planned, component_layout, compiled),) = _compiled_text( + _coherent_surface_code(5, 3) + ) + extent = np.asarray(compiled.tape.extent) + first_drop = np.flatnonzero(np.diff(extent) < 0)[0] + assert np.any(np.diff(extent[first_drop:]) > 0) + return planned, component_layout, compiled + + +def test_extents_and_reused_slots_do_not_change_the_samples(coherent, monkeypatch): + planned, component_layout, compiled = coherent + layout = sharing_layout(compiled, np.ones(compiled.num_outputs, dtype=bool)) + extent = compiled.tape.extent + + def with_extent(values): + return replace(compiled, tape=compiled.tape._replace(extent=values)) + + full = with_extent(np.full_like(extent, compiled.width)) + # NOTE: an extent is an upper bound, so raising every fifth one keeps the + # samples and makes the extent drop and rise around splits and prunes. + every_fifth = np.arange(len(extent)) % 5 == 0 + raised = with_extent(np.where(every_fifth, compiled.width, extent)) + f_params = np.asarray( + _f_params(component_layout, f_patterns(planned.num_f, 7), 300) + ) + key = 17 + reference, deviation, _ = sample_shared_component(full, layout, f_params, key) + # NOTE: NaN fills every amplitude the executor must not read. + monkeypatch.setattr(shared, "_POISON", True) + for component in (compiled, raised): + samples, other, stats = sample_shared_component( + component, layout, f_params, key + ) + np.testing.assert_array_equal(samples, reference) + assert other == deviation + assert stats["dead"] > 0 + + +@pytest.mark.parametrize("name", [*NAMES, "coherent"]) +def test_parallel_kernels_give_the_serial_samples(name, coherent, monkeypatch): + compiled_components = [coherent] if name == "coherent" else _compiled(name) + for planned, component_layout, compiled in compiled_components: + postselected = np.zeros(compiled.num_outputs, dtype=bool) + postselected[0] = True + layout = sharing_layout(compiled, postselected) + f_params = np.asarray( + _f_params(component_layout, f_patterns(planned.num_f, 3), 1_000) + ) + samples = [] + for threshold in (np.iinfo(np.int64).max, 0): + with monkeypatch.context() as patch: + patch.setattr(kernels, "PARALLEL_WORK", threshold) + rows, _, _ = sample_shared_component(compiled, layout, f_params, 19) + samples.append(rows) + np.testing.assert_array_equal(samples[1], samples[0]) + + +def test_shots_whose_lanes_all_die_are_still_sampled(): + all_died = False + for name in NAMES: + for planned, component_layout, compiled in _compiled(name): + layout = sharing_layout(compiled, np.ones(compiled.num_outputs, dtype=bool)) + f_params = np.asarray( + _f_params(component_layout, f_patterns(planned.num_f, 1)[:1], 1) + ) + for seed in range(40): + samples, _, stats = sample_shared_component( + compiled, layout, f_params, seed + ) + assert samples.shape == (1, compiled.num_outputs) + all_died |= stats["lanes"] == 0 + samples, _, _ = sample_shared_component(compiled, layout, f_params[:0], 0) + assert samples.shape == (0, compiled.num_outputs) + assert all_died + + +@pytest.mark.parametrize("name", [*NAMES, "coherent"]) +def test_every_action_stays_within_its_extent(name, coherent): + compiled_components = [coherent] if name == "coherent" else _compiled(name) + for _, _, compiled in compiled_components: + tape = compiled.tape + reach = 1 << np.asarray(tape.extent, dtype=np.int64) + support = np.asarray(tape.x_mask) | np.asarray(tape.z_mask) + assert np.all(support < reach) + measured = np.isin(np.asarray(tape.kind), (MEASURE, MEASURE_DIAG)) + assert np.all((1 << np.asarray(tape.pivot))[measured] < reach[measured]) diff --git a/test/unit/test_postselection.py b/test/unit/test_postselection.py index 70fbbb96..148f8e7e 100644 --- a/test/unit/test_postselection.py +++ b/test/unit/test_postselection.py @@ -2,6 +2,8 @@ from __future__ import annotations +from collections.abc import Callable, Iterator +from contextlib import contextmanager from unittest.mock import patch import numpy as np @@ -65,11 +67,52 @@ def _make(circuit_str: str, seed: int = 0): return Circuit(circuit_str).compile_detector_sampler(seed=seed) +def _direct_outputs(sampler, f_params: np.ndarray) -> np.ndarray: + """Return the direct outputs of ``f_params`` in their columns, False elsewhere.""" + program = sampler._program + columns = program.output_order[: len(program.direct_f_indices)] + out = np.zeros((len(f_params), program.num_outputs), dtype=np.bool_) + out[:, columns] = ( + f_params[:, program.direct_f_indices] != 0 + ) ^ program.direct_flips + return out + + def _keep(samples: np.ndarray, mask: np.ndarray) -> np.ndarray: """Boolean row mask of shots not discarded by postselection.""" return ~np.any(samples & mask, axis=1) +@contextmanager +def _record_noise(sampler) -> Iterator[Callable[[], np.ndarray]]: + """Record the noise ``sampler`` draws; yield a function returning it in shot order. + + Chunks of shots are sampled on several threads, so their draws are put back + into shot order by the order in which their seeds were drawn. + """ + channel_sampler = sampler._channel_sampler + order: list[bytes] = [] + drawn: dict[bytes, np.ndarray] = {} + block_seeds, sample = channel_sampler.block_seeds, channel_sampler.sample + + def record_seeds(n: int) -> np.ndarray: + seeds = block_seeds(n) + order.append(seeds.tobytes()) + return seeds + + def record_sample(n: int, seeds: np.ndarray | None = None) -> np.ndarray: + assert seeds is not None + batch = sample(n, seeds) + drawn[seeds.tobytes()] = batch.copy() + return batch + + with ( + patch.object(channel_sampler, "block_seeds", side_effect=record_seeds), + patch.object(channel_sampler, "sample", side_effect=record_sample), + ): + yield lambda: np.concatenate([drawn[seeds] for seeds in order]) + + # ────────────────────────── validation ─────────────────────────────────────── @@ -162,23 +205,55 @@ def test_postselection_discarded_and_surviving_rows(): assert np.all(~samples[discarded, 1]) +# Detector 0 is postselected; direct detector 2 and observable 0 are not, so a +# survivor sampled from another shot's f-params shows up in their columns. The X +# gates make the reference bits 1, and detectors 0 and 1 share an f-bit. +DIRECT_CIRCUIT = """ +X 0 2 +X_ERROR(0.5) 0 2 +R 1 +H 1 +T 1 +H 1 +M 0 1 2 +DETECTOR rec[-3] +DETECTOR rec[-2] rec[-3] +DETECTOR rec[-1] +OBSERVABLE_INCLUDE(0) rec[-1] +""" + + +@pytest.mark.parametrize("shots,batch_size", [(37, 4), (37, 1), (5, 8), (64, 7)]) +@pytest.mark.parametrize("reference", [False, True]) +def test_survivors_keep_their_own_direct_outputs(shots, batch_size, reference): + """Direct outputs come from each shot's own f-params, flips and reference.""" + sampler = _make(DIRECT_CIRCUIT, seed=11) + mask = np.array([True, False, False]) + with _record_noise(sampler) as drawn: + samples = sampler.sample( + shots, + batch_size=batch_size, + postselection_mask=mask, + append_observables=True, + use_detector_reference_sample=reference, + ) + expected = _direct_outputs(sampler, drawn()) + if reference: + expected[:, :3] ^= sampler._compute_reference_sample()[:3] + discarded = expected[:, 0] + assert np.array_equal(samples[:, [0, 2]], expected[:, [0, 2]]) + assert np.array_equal(samples[~discarded, 3], expected[~discarded, 3]) + assert not samples[discarded][:, [1, 3]].any() + + def test_postselection_direct_cols_always_equal_numpy(): """Direct detector columns match NumPy computation for every row.""" sampler = _make(MIXED_DIRECT_CIRCUIT, seed=3) mask = np.array([True, False]) - drawn: list[np.ndarray] = [] - original = sampler._channel_sampler.sample - - def capture(n: int) -> np.ndarray: - batch = original(n) - drawn.append(batch.copy()) - return batch - - with patch.object(sampler._channel_sampler, "sample", side_effect=capture): + with _record_noise(sampler) as drawn: samples = sampler.sample(8, batch_size=4, postselection_mask=mask) - f_all = np.concatenate(drawn) - expected_direct = sampler._compute_direct_outputs(f_all) + expected_direct = _direct_outputs(sampler, drawn()) direct_det = sampler._direct_detector_mask assert np.array_equal( samples[:, : sampler._num_detectors] & direct_det, @@ -189,19 +264,25 @@ def capture(n: int) -> np.ndarray: # ────────────────────────── JAX-skip behaviour ─────────────────────────────── +def _record_rows(monkeypatch, store: list[int]) -> None: + """Append the rows of every call to the program samplers to ``store``.""" + for name in ("sample_program", "sample_statevector_program"): + original = getattr(sampler_module, name) + + def spy(program, f_params, *args, _original=original, **kwargs): + store.append(f_params.shape[0]) + return _original(program, f_params, *args, **kwargs) + + monkeypatch.setattr(sampler_module, name, spy) + + def test_postselection_jax_never_called_for_all_direct_discards(monkeypatch): """When every shot is discarded by a direct detector, sample_program is never called.""" sampler = _make(ALWAYS_DISCARD_CIRCUIT, seed=0) mask = np.array([True, False]) calls: list[int] = [] - original = sampler_module.sample_program - - def spy(program, f_params, key): - calls.append(f_params.shape[0]) - return original(program, f_params, key) - - monkeypatch.setattr(sampler_module, "sample_program", spy) + _record_rows(monkeypatch, calls) sampler.sample(10, batch_size=4, postselection_mask=mask) assert calls == [] @@ -212,13 +293,8 @@ def test_postselection_jax_rows_less_than_shots(): mask = np.array([True, False]) jax_rows: list[int] = [] - original = sampler_module.sample_program - - def spy(program, f_params, key): - jax_rows.append(f_params.shape[0]) - return original(program, f_params, key) - - with patch.object(sampler_module, "sample_program", side_effect=spy): + with pytest.MonkeyPatch.context() as monkeypatch: + _record_rows(monkeypatch, jax_rows) samples = sampler.sample(32, batch_size=8, postselection_mask=mask) discarded = np.any(samples[:, : sampler._num_detectors] & mask, axis=1) @@ -228,17 +304,13 @@ def spy(program, f_params, key): def test_postselection_jax_batch_size_fixed(monkeypatch): """Every JAX call uses the same batch_size (no recompilation).""" - sampler = _make(MIXED_DIRECT_CIRCUIT, seed=4) + sampler = Circuit(MIXED_DIRECT_CIRCUIT).compile_detector_sampler( + strategy="cat5", seed=4 + ) mask = np.array([True, False]) seen: list[int] = [] - original = sampler_module.sample_program - - def spy(program, f_params, key): - seen.append(f_params.shape[0]) - return original(program, f_params, key) - - monkeypatch.setattr(sampler_module, "sample_program", spy) + _record_rows(monkeypatch, seen) sampler.sample(10, batch_size=4, postselection_mask=mask) assert seen, "expected at least one JAX call for surviving shots" @@ -251,13 +323,7 @@ def test_postselection_non_direct_mask_runs_jax_for_all(monkeypatch): mask = np.array([False, True]) # det1 is a component jax_rows: list[int] = [] - original = sampler_module.sample_program - - def spy(program, f_params, key): - jax_rows.append(f_params.shape[0]) - return original(program, f_params, key) - - monkeypatch.setattr(sampler_module, "sample_program", spy) + _record_rows(monkeypatch, jax_rows) sampler.sample(16, batch_size=8, postselection_mask=mask) assert sum(jax_rows) == 16 @@ -268,13 +334,7 @@ def test_postselection_mixed_mask_skips_jax_only_on_direct_discard(monkeypatch): mask = np.array([True, True]) jax_rows: list[int] = [] - original = sampler_module.sample_program - - def spy(program, f_params, key): - jax_rows.append(f_params.shape[0]) - return original(program, f_params, key) - - monkeypatch.setattr(sampler_module, "sample_program", spy) + _record_rows(monkeypatch, jax_rows) samples = sampler.sample(32, batch_size=8, postselection_mask=mask) direct_discarded = samples[:, 0] & mask[0] assert direct_discarded.any() @@ -296,21 +356,13 @@ def test_postselection_detector_reference_xor_before_discard_check(monkeypatch): mask = np.array([True, False]) jax_without: list[int] = [] jax_with: list[int] = [] - original = sampler_module.sample_program - - def make_spy(store: list[int]): - def spy(program, f_params, key): - store.append(f_params.shape[0]) - return original(program, f_params, key) - - return spy - s1 = _make(circuit, seed=0) - monkeypatch.setattr(sampler_module, "sample_program", make_spy(jax_without)) - s1.sample(32, batch_size=8, postselection_mask=mask) + with monkeypatch.context() as patch_rows: + _record_rows(patch_rows, jax_without) + s1.sample(32, batch_size=8, postselection_mask=mask) s2 = _make(circuit, seed=0) - monkeypatch.setattr(sampler_module, "sample_program", make_spy(jax_with)) + _record_rows(monkeypatch, jax_with) s2.sample( 32, batch_size=8, @@ -318,11 +370,15 @@ def spy(program, f_params, key): use_detector_reference_sample=True, ) - assert sum(jax_with) > sum(jax_without) + # NOTE: the reference sample itself is one more row through sample_program. + assert sum(jax_with) - 1 > sum(jax_without) -def test_postselection_observable_reference_on_jax_computed_discarded_rows(): - """Observable ref XOR applies to every row that ran JAX, including caller discards.""" +@pytest.mark.parametrize("strategy", ["cat5", "active_statevector"]) +def test_postselection_observable_reference_flips_every_sampled_row( + strategy, monkeypatch +): + """The observable reference flips every row that was sampled, but no direct discard.""" circuit = """ X 0 X_ERROR(0.5) 0 @@ -333,39 +389,25 @@ def test_postselection_observable_reference_on_jax_computed_discarded_rows(): DETECTOR rec[-1] rec[-2] OBSERVABLE_INCLUDE(0) rec[-1] """ - mask = np.array([True, True]) - kwargs = { - "batch_size": 8, - "use_observable_reference_sample": True, - "separate_observables": True, - } - # NOTE: under cat5, seed 2 samples a reference observable of 1. - sampler = Circuit(circuit).compile_detector_sampler(strategy="cat5", seed=2) - captured: dict[str, np.ndarray] = {} - original = sampler._sample_batches_with_postselection - - def capture(*args, **kwargs): - samples, reference, direct_discarded = original(*args, **kwargs) - assert reference is not None - captured["raw_obs"] = samples[:, sampler._num_detectors :].copy() - captured["direct_discarded"] = direct_discarded.copy() - captured["obs_ref"] = reference[sampler._num_detectors :].copy() - return samples, reference, direct_discarded - - sampler._sample_batches_with_postselection = capture - _dets, obs = sampler.sample(128, postselection_mask=mask, **kwargs) - - raw_obs = captured["raw_obs"] - direct_discarded = captured["direct_discarded"] - obs_ref = captured["obs_ref"] - expected = raw_obs.copy() - expected[~direct_discarded] ^= obs_ref - - assert np.all(obs_ref) - assert np.array_equal(obs, expected) - assert not np.any(obs[direct_discarded]) - ran_jax = ~direct_discarded - assert np.any(obs[ran_jax] != raw_obs[ran_jax]) + mask = np.array([True, False]) + + def observables(reference: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + sampler = Circuit(circuit).compile_detector_sampler(strategy=strategy, seed=2) + monkeypatch.setattr(sampler, "_compute_reference_sample", lambda: reference) + return sampler.sample( + 128, + batch_size=8, + postselection_mask=mask, + separate_observables=True, + use_observable_reference_sample=True, + ) + + dets, raw = observables(np.zeros(3, dtype=np.bool_)) + _, flipped = observables(np.array([False, False, True])) + direct_discarded = dets[:, 0] + assert direct_discarded.any() and (~direct_discarded).any() + np.testing.assert_array_equal(flipped[~direct_discarded], ~raw[~direct_discarded]) + assert not flipped[direct_discarded].any() and not raw[direct_discarded].any() # ────────────────────────── fully-direct fast path ─────────────────────────── @@ -375,13 +417,8 @@ def test_postselection_fully_direct_no_jax(monkeypatch): """Fully-direct circuits never call sample_program.""" sampler = _make(FULLY_DIRECT_CIRCUIT, seed=0) mask = np.array([True]) - spy = [] - - def counting_sp(program, f_params, key): - spy.append(1) - return sampler_module.sample_program(program, f_params, key) - - monkeypatch.setattr(sampler_module, "sample_program", counting_sp) + spy: list[int] = [] + _record_rows(monkeypatch, spy) result = sampler.sample(10, postselection_mask=mask) assert result.shape == (10, 1) assert spy == [] @@ -452,9 +489,9 @@ def test_postselection_reference_does_not_advance_channel_rng(): original = sampler._channel_sampler.sample calls: list[int] = [] - def spy(n: int) -> np.ndarray: + def spy(n: int, seeds: np.ndarray | None = None) -> np.ndarray: calls.append(n) - return original(n) + return original(n, seeds) with patch.object(sampler._channel_sampler, "sample", side_effect=spy): sampler._compute_reference_sample() @@ -569,30 +606,3 @@ def test_postselection_surface_code_caller_filter(): survivors = _keep(samples[:, :num_det], mask) assert survivors.any() assert not np.any(samples[survivors] & mask) - - -def test_postselection_discarded_rows_zero_direct_observable(): - """Discarded rows keep direct detector cols but zero direct observable cols.""" - circ = stim.Circuit.generated( - "surface_code:rotated_memory_x", - distance=3, - rounds=2, - after_clifford_depolarization=0.01, - ) - c = Circuit.from_stim_program(circ) - sampler = c.compile_detector_sampler(seed=0) - assert sampler._direct_output_mask[ - sampler._num_detectors : - ].any(), "expected a direct observable for this circuit" - - mask = np.zeros(c.num_detectors, dtype=np.bool_) - mask[0] = True - dets, obs = sampler.sample( - 64, - batch_size=16, - postselection_mask=mask, - separate_observables=True, - ) - discarded = np.any(dets & mask, axis=1) - assert discarded.any() and (~discarded).any() - assert not np.any(obs[discarded]) diff --git a/test/unit/test_sampler.py b/test/unit/test_sampler.py index daa136b1..cce4f74e 100644 --- a/test/unit/test_sampler.py +++ b/test/unit/test_sampler.py @@ -3,6 +3,7 @@ import numpy as np import pytest +import tsim.sampler as sampler_module from tsim.circuit import Circuit @@ -509,3 +510,70 @@ def test_detector_sampler_out_of_order_observable_indices(): assert not np.any(obs[:, 0]) assert not np.any(obs[:, 1]) assert np.all(obs[:, 2]) + + +def test_a_nan_normalization_deviation_warns(monkeypatch): + sampler = Circuit("R 0\nH 0\nT 0\nH 0\nM 0").compile_sampler( + strategy="active_statevector", seed=0 + ) + original = sampler_module.sample_shared_component + + def nan_deviation(*args): + samples, _, stats = original(*args) + return samples, float("nan"), stats + + monkeypatch.setattr(sampler_module, "sample_shared_component", nan_deviation) + with pytest.warns(UserWarning, match="not normalized"): + sampler.sample(10) + + +def test_direct_detectors_overwrite_every_column_of_their_rows(): + f = np.array([[1, 0], [0, 1], [1, 1]], dtype=np.uint8) + out = np.ones((3, 4), dtype=np.bool_) + discarded = sampler_module._write_direct_detectors( + 0, + f, + np.array([1, 0]), + np.array([False, True]), + np.array([3, 0]), + np.array([True, False]), + out, + ) + expected = np.zeros((3, 4), dtype=np.bool_) + expected[:, 3] = f[:, 1] != 0 + expected[:, 0] = (f[:, 0] != 0) ^ True + assert np.array_equal(out, expected) + assert np.array_equal(discarded, f[:, 1] != 0) + + +T_CIRCUIT = """ +H 0 +T 0 +H 0 +M 0 +DETECTOR rec[-1] +""" + + +@pytest.mark.parametrize( + ("strategy", "sharing", "backend", "batch_size", "expected"), + [ + ("cat5", False, "cpu", 10, 1), + ("active_statevector", True, "cpu", 10, 6), + ("active_statevector", True, "gpu", 10, 6), + ("active_statevector", False, "cpu", 10, 6), + # NOTE: statevector programs sample on the host whatever JAX's backend. + ("active_statevector", False, "gpu", 10, 6), + ("active_statevector", False, "cpu", 10**12, 1), + ("active_statevector", True, "cpu", 10**12, 1), + ], +) +def test_batches_running_at_once( + strategy, sharing, backend, batch_size, expected, monkeypatch +): + sampler = Circuit(T_CIRCUIT).compile_detector_sampler( + strategy=strategy, trajectory_sharing=sharing + ) + monkeypatch.setattr(sampler_module.numba, "get_num_threads", lambda: 6) + monkeypatch.setattr(sampler_module.jax, "default_backend", lambda: backend) + assert sampler._threads(batch_size) == expected diff --git a/test/unit/utils/test_kernels.py b/test/unit/utils/test_kernels.py new file mode 100644 index 00000000..f612b439 --- /dev/null +++ b/test/unit/utils/test_kernels.py @@ -0,0 +1,135 @@ +import subprocess +import sys +import textwrap +import threading +import time + +import numba +import numpy as np +import pytest + +import tsim.utils.kernels as kernels +from tsim.utils.kernels import PARALLEL_WORK, kernel, run_tasks + + +@kernel +def _square_roots(out): + for i in numba.prange(out.shape[0]): + out[i] = np.sqrt(i) + + +@kernel +def _thread_ids(out): + for i in numba.prange(out.shape[0]): + out[i] = numba.get_thread_id() + + +def test_serial_and_parallel_runs_agree(): + serial, parallel = np.zeros(1000), np.zeros(1000) + _square_roots(0, serial) + _square_roots(PARALLEL_WORK, parallel) + np.testing.assert_array_equal(serial, np.sqrt(np.arange(1000))) + np.testing.assert_array_equal(parallel, serial) + + +def test_parallel_kernels_survive_concurrent_python_threads(tmp_path): + # NOTE: without the lock numba's workqueue layer aborts the whole process, so + # the threads run in a subprocess. + script = tmp_path / "threads.py" + script.write_text(textwrap.dedent(""" + import threading + + import numba + import numpy as np + + from tsim.utils.kernels import kernel + + @kernel + def fill(out): + for i in numba.prange(out.shape[0]): + out[i] = np.sqrt(i) + + def work(): + out = np.zeros(1 << 18) + for _ in range(50): + fill(1 << 30, out) + + threads = [threading.Thread(target=work) for _ in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + """)) + command = [sys.executable, str(script)] + result = subprocess.run(command, capture_output=True) # noqa: S603 + assert result.returncode == 0, result.stderr.decode() + + +def test_run_tasks_keeps_the_order_and_runs_kernels_serially(): + def task(n): + ids = np.full(1 << 12, -1) + _thread_ids(PARALLEL_WORK, ids) + return n * n, ids + + results = run_tasks(task, range(10), threads=4) + assert [square for square, _ in results] == [n * n for n in range(10)] + assert all((ids == 0).all() for _, ids in results) + + +def test_run_tasks_with_one_thread_runs_in_the_calling_thread(): + caller = threading.get_ident() + assert run_tasks(lambda _: threading.get_ident(), range(3), 1) == [caller] * 3 + assert run_tasks(lambda n: n, [], 4) == [] + + +def test_run_tasks_raises_the_error_of_a_task_once_every_task_is_done(): + finished = [] + + def task(n): + if n == 2: + raise ValueError("task 2") + time.sleep(0.05) + finished.append(n) + + with pytest.raises(ValueError, match="task 2"): + run_tasks(task, range(4), threads=3) + assert sorted(finished) == [0, 1, 3] + + +def test_the_calling_thread_keeps_parallel_kernels_after_run_tasks(): + if numba.get_num_threads() < 2: + pytest.skip("needs two numba threads") + run_tasks(lambda _: _thread_ids(PARALLEL_WORK, np.zeros(64)), range(4), 4) + ids = np.full(1 << 12, -1) + _thread_ids(PARALLEL_WORK, ids) + assert len(np.unique(ids)) > 1 + + +def test_run_tasks_inside_a_task_runs_in_the_task(): + def task(n): + return run_tasks(lambda m: (m, threading.get_ident()), range(3), 4) + + for inner in run_tasks(task, range(8), 4): + assert [m for m, _ in inner] == [0, 1, 2] + assert len({ident for _, ident in inner}) == 1 + + +def test_run_tasks_runs_at_most_threads_tasks_at_once(): + lock, running, peak = threading.Lock(), [0], [0] + + def task(_): + with lock: + running[0] += 1 + peak[0] = max(peak[0], running[0]) + time.sleep(0.01) + with lock: + running[0] -= 1 + + run_tasks(task, range(12), 3) + assert 1 < peak[0] <= 3 + + +def test_run_tasks_reuses_one_pool_for_every_thread_count(): + for threads in range(2, 9): + run_tasks(lambda n: n, range(16), threads) + assert kernels._pool.cache_info().currsize == 1 diff --git a/uv.lock b/uv.lock index 3af73cde..f5414d79 100644 --- a/uv.lock +++ b/uv.lock @@ -178,6 +178,7 @@ dependencies = [ { name = "jax", version = "0.6.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "jax", version = "0.9.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "lxml" }, + { name = "numba" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "psutil" }, @@ -247,6 +248,7 @@ requires-dist = [ { name = "jax", extras = ["cuda12"], marker = "extra == 'cuda12'", specifier = ">=0.4.38" }, { name = "jax", extras = ["cuda13"], marker = "extra == 'cuda13'", specifier = ">=0.6.0" }, { name = "lxml", specifier = ">=5.0.0" }, + { name = "numba", specifier = ">=0.61" }, { name = "numpy", specifier = ">=1.25.0" }, { name = "psutil", specifier = ">=5.9.0" }, { name = "stim", specifier = ">=1.0.0" },