From 4318f2c40841a27de7e7cd382a547484afedc4aa Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 13 Oct 2026 00:13:07 -0400 Subject: [PATCH 01/21] Buffer postselection survivors in a preallocated array The survivors of each chunk were appended to a Python list one row at a time and stacked back into a batch with np.stack, which cost more than the statevector simulation on cultivation d3 (about 60 ms of a 110 ms call of 262,144 shots). They now go into a two-batch array: a chunk adds at most one batch of survivors to fewer than one pending batch. Samples are bit-identical for a fixed seed. Quick benchmark, all threads: cultivation d3 2.61M -> 5.81M shots/s with trajectory sharing and 1.66M -> 2.56M with one lane per shot. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 59 ++++++++++++++++++++------------------------- 1 file changed, 26 insertions(+), 33 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 4e1374b4..ea9a3da8 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -600,41 +600,31 @@ def _sample_batches_with_postselection( 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: + # NOTE: a chunk adds at most batch_size survivors to fewer than batch_size + # pending ones, so two batches always hold the pending survivors. + num_f = self._channel_sampler.signature_matrix.shape[1] + pending_f = np.empty((2 * batch_size, num_f), dtype=np.uint8) + pending_shot = np.empty(2 * batch_size, dtype=np.int64) + pending = 0 + + def dispatch(count: int) -> None: + """Sample the first ``count`` pending survivors in one full batch.""" + nonlocal pending + pending_f[count:batch_size] = pending_f[0] self._key, subkey = jax.random.split(self._key) - jax_out = np.asarray( - sample_program(self._program, jnp.asarray(f_batch), subkey, **sharing) + samples = sample_program( + self._program, jnp.asarray(pending_f[:batch_size]), subkey, **sharing ) - 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, - ) - 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 = [] + result[pending_shot[:count]] = np.asarray(samples)[:count] + pending -= count + pending_f[:pending] = pending_f[count : count + pending] + pending_shot[:pending] = pending_shot[count : count + pending] while shot_idx < shots: chunk = min(batch_size, shots - shot_idx) @@ -651,15 +641,18 @@ def _flush(*, final: bool = False) -> None: ] 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()) + survivors = np.flatnonzero(~discarded) + end = pending + len(survivors) + pending_f[pending:end] = f_params_np[survivors] + pending_shot[pending:end] = shot_idx + survivors + pending = end shot_idx += chunk - _flush() + if pending >= batch_size: + dispatch(batch_size) - _flush(final=True) + if pending: + dispatch(pending) if xor_detector_ref and reference is not None: det_ref = reference[: self._num_detectors] From 728a0833019a1e6b1a234da95af5a8ba139ecb1e Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 13 Oct 2026 00:22:30 -0400 Subject: [PATCH 02/21] Sample noise with a parallel numba kernel ChannelSampler.sample looped over every channel in Python and XOR-ed a dense uint8 row of all f-variables into each shot an error fired on; on surface d7r7 an error flips 2.7 of 337 f-variables, and the sampler took 91% of the time. A numba kernel now walks each channel's geometric skips and flips only the f-variables of the fired outcome. Blocks of 4096 shots run in parallel, each from its own SplitMix64 stream seeded with 64 bits from the sampler's Generator, so samples depend on the seed and the number of shots but not on the number of threads. The output layout is unchanged. The stream differs from NumPy's, so one test that pinned 4 of 10 events at p=0.3 now checks the rate over 100,000 shots. New tests cover a 2-bit channel's joint distribution across blocks, a certain error reaching every row of a partial last block, and reproducibility. Quick benchmark, all threads: surface d7r7 1.55M -> 22.4M shots/s, distillation (shared) 3.6M -> 9.0M, cultivation d5 (shared) 1.5M -> 2.7M. Co-Authored-By: Claude Opus 5.5 --- pyproject.toml | 1 + src/tsim/noise/channels.py | 167 +++++++++++++--------- test/integration/test_sampler_circuits.py | 5 +- test/unit/noise/test_channels.py | 42 +++++- uv.lock | 12 +- 5 files changed, 151 insertions(+), 76 deletions(-) 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/noise/channels.py b/src/tsim/noise/channels.py index 58f32d0f..339e371a 100644 --- a/src/tsim/noise/channels.py +++ b/src/tsim/noise/channels.py @@ -4,9 +4,16 @@ from collections import defaultdict from dataclasses import dataclass +from typing import NamedTuple +import numba import numpy as np +# 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: @@ -584,61 +591,15 @@ 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 - ) - - @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. - - 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 - - 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. - - 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). - - """ - 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 - ) - xor_patterns = (bits_mask @ signature_matrix[col_ids] % 2).astype(np.uint8) - - data.append((p_fire, cond_cdf, xor_patterns)) - return data + self._tables = fire_tables(self.channels, self.signature_matrix) def sample(self, num_samples: int = 1) -> 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. + 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: num_samples: Number of samples to draw. @@ -648,24 +609,98 @@ def sample(self, num_samples: int = 1) -> np.ndarray: 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) + num_f = self.signature_matrix.shape[1] + result = np.zeros((num_samples, num_f), dtype=np.uint8) + num_blocks = -(-num_samples // _BLOCK_SHOTS) + seeds = self._rng.integers(0, 2**64, size=num_blocks, dtype=np.uint64) + _sample_blocks(result, seeds, *self._tables) + return result - 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 - positions = np.cumsum(self._rng.geometric(p_fire, size=n_draws)) - 1 - positions = positions[positions < num_samples] +class FireTables(NamedTuple): + """Geometric-skip sampling tables of the channels that can fire. - if len(positions) == 0: - continue + 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]]``. + """ - outcome_idx = np.searchsorted( - cond_cdf, self._rng.uniform(size=len(positions)) - ) - result[positions] ^= xor_pats[outcome_idx] + p_fire: np.ndarray + outcome_start: np.ndarray + cdf: np.ndarray + flip_start: np.ndarray + flip_cols: np.ndarray - return result + +def fire_tables(channels: list[Channel], signature_matrix: np.ndarray) -> FireTables: + """Return the sampling tables of ``channels``. + + 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. + + """ + 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) + fire = 1.0 - float(probs[0]) + if fire <= 1e-15 or len(probs) <= 1: + continue + p_fire.append(fire) + conditional = np.cumsum(probs[1:] / 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), + ) + + +@numba.njit(cache=True) +def _next_uniform(state: np.uint64) -> tuple[np.uint64, float]: + """Advance a SplitMix64 state; return it and a uniform float in [0, 1).""" + state += np.uint64(0x9E3779B97F4A7C15) + z = (state ^ (state >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9) + z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB) + z ^= z >> np.uint64(31) + return state, (z >> np.uint64(11)) * 2.0**-53 + + +@numba.njit(cache=True) +def _sample_block(out, state, p_fire, outcome_start, cdf, flip_start, flip_cols): + """XOR every channel's outcomes into the rows of ``out`` it fires on.""" + 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 + while True: + state, u = _next_uniform(state) + row += 1 + int(np.log1p(-u) / log_keep) + 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[row, flip_cols[j]] ^= 1 + + +@numba.njit(parallel=True, cache=True) +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/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/unit/noise/test_channels.py b/test/unit/noise/test_channels.py index 162f75a3..d93184b3 100644 --- a/test/unit/noise/test_channels.py +++ b/test/unit/noise/test_channels.py @@ -3,6 +3,7 @@ from numpy.testing import assert_allclose from tsim.noise.channels import ( + _BLOCK_SHOTS, Channel, ChannelSampler, absorb_subset_channels, @@ -10,6 +11,7 @@ correlated_error_probs, error_probs, expand_channel, + fire_tables, fold_duplicate_channel_bits, heralded_pauli_channel_1_probs, merge_identical_channels, @@ -213,9 +215,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) @@ -1047,6 +1047,42 @@ 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_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)] diff --git a/uv.lock b/uv.lock index 3af73cde..8082c79f 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.15'", @@ -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" }, @@ -1281,15 +1283,15 @@ wheels = [ [[package]] name = "griffe-kirin" -version = "0.2.0" -source = { registry = "https://pypi.org/simple" } +version = "0.1.0" +source = { registry = "https://quera.jfrog.io/artifactory/api/pypi/kirin/simple" } dependencies = [ { name = "griffe" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a5/e0/9db75de974546903ab46a41179a95243abc84ceb2ba0b4101b5016f5301e/griffe_kirin-0.2.0.tar.gz", hash = "sha256:87d4280fbb1c77e1eb3fe7b627810b219295ec74651cf216748330b379a4ded5", size = 57664, upload-time = "2025-02-21T20:41:16.862Z" } +sdist = { url = "https://quera.jfrog.io/artifactory/api/pypi/kirin/griffe-kirin/0.1.0/griffe_kirin-0.1.0.tar.gz", hash = "sha256:b5ed85a8c07da0219787ac8d57de0a4f2a6df3109fc3aab47d78a12f96dda04d" } wheels = [ - { url = "https://files.pythonhosted.org/packages/38/b2/04bf90dc9e9e80f969b562731cf1619bd2eb0d0b726ba7befd38fda38131/griffe_kirin-0.2.0-py3-none-any.whl", hash = "sha256:81d999d6e1480cbe46e5c7b82933d2eee96863840dfbac283df620d85c13c2fc", size = 4063, upload-time = "2025-02-21T20:41:14.734Z" }, + { url = "https://quera.jfrog.io/artifactory/api/pypi/kirin/griffe-kirin/0.1.0/griffe_kirin-0.1.0-py3-none-any.whl", hash = "sha256:cdbc7a08c4fa5229f8f5f9137fd4de7dde35214166d3d448cbac8778079bd10a" }, ] [[package]] From bfdcd4734f5d59491f1b43df469dd2ce1f6ec925 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 13 Oct 2026 00:31:41 -0400 Subject: [PATCH 03/21] Write direct detectors and the discard test in one numba pass With postselection, every chunk gathered the direct outputs into a zeroed array, copied its detector columns into the result and tested the postselected ones in three NumPy passes: 80 ms per 262,144 shots on cultivation d5 (47% of a call) and 466 ms on SymFT's d=7 circuit. One parallel numba pass now writes each direct detector into its result row and returns the discard mask (3.7 ms and 13 ms). Survivors are now buffered with one np.concatenate per chunk, as the review of the preallocated buffer suggested: as fast, without the two-batch invariant, and JAX never reads a buffer that is overwritten later. A new test checks that every survivor keeps its own direct detectors; the old tests could not catch misattributed rows, because their only direct detector was postselected. Samples are bit-identical for a fixed seed. Quick benchmark, all threads: distillation (shared) 9.0M -> 16.3M shots/s, cultivation d5 (shared) 2.7M -> 3.2M, SymFT d7 postselected (shared) 202k -> 241k. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 89 ++++++++++++++++++++------------- test/unit/test_postselection.py | 42 ++++++++++++++++ 2 files changed, 96 insertions(+), 35 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index ea9a3da8..c3ddbf0d 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -9,6 +9,7 @@ import jax import jax.numpy as jnp +import numba import numpy as np import psutil @@ -191,6 +192,23 @@ def sample_program( return combined +@numba.njit(parallel=True, cache=True) +def _write_direct_detectors(f, f_indices, flips, columns, postselect, reference, out): + """Write direct detector ``j`` of every row into ``out[row, columns[j]]``. + + Detector ``j`` is ``f[row, f_indices[j]] ^ flips[j]``. Returns which rows + have a postselected detector that differs from its ``reference`` bit. + """ + discarded = np.zeros(f.shape[0], dtype=np.bool_) + for row in numba.prange(f.shape[0]): + 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 != reference[j]: + discarded[row] = True + return discarded + + def _nbytes(tree: object) -> int: return sum( v.nbytes for v in jax.tree_util.tree_leaves(tree) if isinstance(v, jax.Array) @@ -576,8 +594,6 @@ def _sample_batches_with_postselection( return empty, np.zeros(num_outputs, dtype=np.bool_), empty_discarded return empty, None, empty_discarded - postselect_direct = postselection_mask & self._direct_detector_mask - if not self._program.components: samples = self._sample_direct(shots) if compute_reference: @@ -602,57 +618,60 @@ def _sample_batches_with_postselection( was_discarded = np.zeros(shots, dtype=np.bool_) shot_idx = 0 + detector = self._direct_global_indices < self._num_detectors + columns = self._direct_global_indices[detector] + direct = ( + self._direct_f_indices[detector], + self._direct_flips[detector], + columns, + postselection_mask[columns], + ( + reference[columns] + if xor_detector_ref and reference is not None + else np.zeros(len(columns), dtype=np.bool_) + ), + ) + sharing = self._sharing_arguments( postselection_mask, reference if xor_detector_ref else None ) - # NOTE: a chunk adds at most batch_size survivors to fewer than batch_size - # pending ones, so two batches always hold the pending survivors. num_f = self._channel_sampler.signature_matrix.shape[1] - pending_f = np.empty((2 * batch_size, num_f), dtype=np.uint8) - pending_shot = np.empty(2 * batch_size, dtype=np.int64) - pending = 0 - - def dispatch(count: int) -> None: - """Sample the first ``count`` pending survivors in one full batch.""" - nonlocal pending - pending_f[count:batch_size] = pending_f[0] + pending_f = np.empty((0, num_f), dtype=np.uint8) + pending_shot = np.empty(0, dtype=np.int64) + + def dispatch(f_batch: np.ndarray, shot_batch: np.ndarray) -> None: + """Sample the shots ``shot_batch`` in one full batch, padded with row 0.""" + count = len(shot_batch) + if count < batch_size: + padding = np.repeat(f_batch[:1], batch_size - count, axis=0) + f_batch = np.concatenate([f_batch, padding]) self._key, subkey = jax.random.split(self._key) samples = sample_program( - self._program, jnp.asarray(pending_f[:batch_size]), subkey, **sharing + self._program, jnp.asarray(f_batch), subkey, **sharing ) - result[pending_shot[:count]] = np.asarray(samples)[:count] - pending -= count - pending_f[:pending] = pending_f[count : count + pending] - pending_shot[:pending] = pending_shot[count : count + pending] + result[shot_batch] = np.asarray(samples)[:count] 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 - ] + discarded = _write_direct_detectors( + f_params_np, *direct, result[shot_idx : shot_idx + chunk] + ) was_discarded[shot_idx : shot_idx + chunk] = discarded survivors = np.flatnonzero(~discarded) - end = pending + len(survivors) - pending_f[pending:end] = f_params_np[survivors] - pending_shot[pending:end] = shot_idx + survivors - pending = end + pending_f = np.concatenate([pending_f, f_params_np[survivors]]) + pending_shot = np.concatenate([pending_shot, shot_idx + survivors]) shot_idx += chunk - if pending >= batch_size: - dispatch(batch_size) + if len(pending_shot) >= batch_size: + dispatch(pending_f[:batch_size], pending_shot[:batch_size]) + pending_f = pending_f[batch_size:] + pending_shot = pending_shot[batch_size:] - if pending: - dispatch(pending) + if len(pending_shot): + dispatch(pending_f, pending_shot) if xor_detector_ref and reference is not None: det_ref = reference[: self._num_detectors] diff --git a/test/unit/test_postselection.py b/test/unit/test_postselection.py index 70fbbb96..f3aa99f6 100644 --- a/test/unit/test_postselection.py +++ b/test/unit/test_postselection.py @@ -162,6 +162,48 @@ def test_postselection_discarded_and_surviving_rows(): assert np.all(~samples[discarded, 1]) +# Detector 0 is postselected; direct detector 2 is not, so a survivor that is +# sampled from another shot's f-params shows up in column 2. +TWO_DIRECT_CIRCUIT = """ +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] +""" + + +@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_detectors(shots, batch_size, reference): + """Every row's direct detectors come from that shot's own f-params.""" + sampler = _make(TWO_DIRECT_CIRCUIT, seed=11) + drawn: list[np.ndarray] = [] + original = sampler._channel_sampler.sample + + def capture(n: int) -> np.ndarray: + batch = original(n) + drawn.append(batch.copy()) + return batch + + mask = np.array([True, False, False]) + with patch.object(sampler._channel_sampler, "sample", side_effect=capture): + samples = sampler.sample( + shots, + batch_size=batch_size, + postselection_mask=mask, + use_detector_reference_sample=reference, + ) + expected = sampler._compute_direct_outputs(np.concatenate(drawn)) + if reference: + expected[:, :3] ^= sampler._compute_reference_sample()[:3] + assert np.array_equal(samples[:, [0, 2]], expected[:, [0, 2]]) + + def test_postselection_direct_cols_always_equal_numpy(): """Direct detector columns match NumPy computation for every row.""" sampler = _make(MIXED_DIRECT_CIRCUIT, seed=3) From feb18e4169a37754950f0a0d6960d97fac986f46 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 13 Oct 2026 01:04:36 -0400 Subject: [PATCH 04/21] Fix the issues the reviews of the numba kernels found - A channel whose identity probability is rounded below zero (e.g. PAULI_CHANNEL_1(0.33, 0.56, 0.11)) had a firing probability above 1, so log1p(-p) was NaN: undefined samples, a segfault on x86-64. The firing probability is now the clipped total of the non-identity outcomes, which is also exact for small probabilities. - Numba's default threading layer aborts the process when two Python threads run parallel code at once. tsim.utils.kernels.kernel compiles each kernel serially and in parallel, runs it in parallel only above a work threshold (waking the threads costs about 0.1 ms) and holds a lock around parallel calls. The parallel compilation gets its own name, since numba names cache files after the function and the two would share one entry. - Numba refuses to define a cached function where no cache directory is writable, which made `import tsim` fail on read-only installs; jit falls back to compiling without a cache. - SplitMix64 helpers move to tsim.utils.rng; geometric skips use a 64-bit uniform and return a float, inf when the uniform rounds to 1. - The direct-detector kernel takes named arguments built once per sampler. - Tests: blocks draw independent streams, a rounded identity probability, concurrent Python threads (in a subprocess), and a stronger survivor test with reference bits, flips, a shared f-bit, reordered outputs and a direct observable. A test whose circuit never reached the postselection path is removed; the reference test no longer counts the reference sample's row. - uv.lock adds only the numba lines (the earlier lock had been rewritten with a user-level index configuration). Co-Authored-By: Claude Opus 5.5 --- src/tsim/noise/channels.py | 48 +++++++++++++-------------- src/tsim/sampler.py | 34 +++++++++++-------- src/tsim/utils/kernels.py | 57 ++++++++++++++++++++++++++++++++ src/tsim/utils/rng.py | 47 ++++++++++++++++++++++++++ test/unit/noise/test_channels.py | 23 +++++++++++-- test/unit/test_postselection.py | 49 +++++++++------------------ test/unit/utils/test_kernels.py | 55 ++++++++++++++++++++++++++++++ uv.lock | 10 +++--- 8 files changed, 243 insertions(+), 80 deletions(-) create mode 100644 src/tsim/utils/kernels.py create mode 100644 src/tsim/utils/rng.py create mode 100644 test/unit/utils/test_kernels.py diff --git a/src/tsim/noise/channels.py b/src/tsim/noise/channels.py index 339e371a..664c4209 100644 --- a/src/tsim/noise/channels.py +++ b/src/tsim/noise/channels.py @@ -9,6 +9,9 @@ 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. @@ -591,7 +594,7 @@ def __init__( self._rng = np.random.default_rng( seed if seed is not None else np.random.default_rng().integers(0, 2**30) ) - self._tables = fire_tables(self.channels, self.signature_matrix) + self._tables = _fire_tables(self.channels, self.signature_matrix) def sample(self, num_samples: int = 1) -> np.ndarray: """Sample from all error channels and transform to new error basis. @@ -613,11 +616,12 @@ def sample(self, num_samples: int = 1) -> np.ndarray: result = np.zeros((num_samples, num_f), dtype=np.uint8) num_blocks = -(-num_samples // _BLOCK_SHOTS) seeds = self._rng.integers(0, 2**64, size=num_blocks, dtype=np.uint64) - _sample_blocks(result, seeds, *self._tables) + work = num_samples * len(self._tables.p_fire) + _sample_blocks(work, result, seeds, *self._tables) return result -class FireTables(NamedTuple): +class _FireTables(NamedTuple): """Geometric-skip sampling tables of the channels that can fire. Channel ``c`` fires with probability ``p_fire[c]``. Its non-identity @@ -633,9 +637,13 @@ class FireTables(NamedTuple): flip_cols: np.ndarray -def fire_tables(channels: list[Channel], signature_matrix: np.ndarray) -> FireTables: +def _fire_tables(channels: list[Channel], signature_matrix: np.ndarray) -> _FireTables: """Return the sampling tables of ``channels``. + 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. + Args: channels: The noise channels. signature_matrix: Binary matrix of shape (num_signatures, num_f) mapping @@ -648,11 +656,11 @@ def fire_tables(channels: list[Channel], signature_matrix: np.ndarray) -> FireTa outcome_start, flip_start = [0], [0] for channel in channels: probs = channel.probs.astype(np.float64) - fire = 1.0 - float(probs[0]) - if fire <= 1e-15 or len(probs) <= 1: + 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) - conditional = np.cumsum(probs[1:] / fire) cdf.extend(conditional / conditional[-1]) outcome_start.append(len(cdf)) col_ids = np.asarray(channel.unique_col_ids) @@ -661,7 +669,7 @@ def fire_tables(channels: list[Channel], signature_matrix: np.ndarray) -> FireTa for flips in bits @ signature_matrix[col_ids] % 2: flip_cols.extend(np.flatnonzero(flips).tolist()) flip_start.append(len(flip_cols)) - return FireTables( + 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), @@ -670,35 +678,25 @@ def fire_tables(channels: list[Channel], signature_matrix: np.ndarray) -> FireTa ) -@numba.njit(cache=True) -def _next_uniform(state: np.uint64) -> tuple[np.uint64, float]: - """Advance a SplitMix64 state; return it and a uniform float in [0, 1).""" - state += np.uint64(0x9E3779B97F4A7C15) - z = (state ^ (state >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9) - z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB) - z ^= z >> np.uint64(31) - return state, (z >> np.uint64(11)) * 2.0**-53 - - -@numba.njit(cache=True) +@jit def _sample_block(out, state, p_fire, outcome_start, cdf, flip_start, flip_cols): """XOR every channel's outcomes into the rows of ``out`` it fires on.""" 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 + row = -1.0 while True: - state, u = _next_uniform(state) - row += 1 + int(np.log1p(-u) / log_keep) + state, skip = next_skip(state, log_keep) + row += 1.0 + skip if row >= out.shape[0]: break - state, u = _next_uniform(state) + 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[row, flip_cols[j]] ^= 1 + out[int(row), flip_cols[j]] ^= 1 -@numba.njit(parallel=True, cache=True) +@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)): diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index c3ddbf0d..f55825dd 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -22,6 +22,7 @@ from tsim.statevector.program import STABILIZER_RANK_HINT, StatevectorComponent from tsim.statevector.shared import Sharing, sample_shared_component, sharing_layout from tsim.utils.cuda_helpers import copy_d2h +from tsim.utils.kernels import kernel if TYPE_CHECKING: from jax import Array as PRNGKey @@ -192,7 +193,7 @@ def sample_program( return combined -@numba.njit(parallel=True, cache=True) +@kernel def _write_direct_detectors(f, f_indices, flips, columns, postselect, reference, out): """Write direct detector ``j`` of every row into ``out[row, columns[j]]``. @@ -296,6 +297,10 @@ def __init__( self._direct_detector_mask = self._direct_output_mask[ : self._num_detectors ].copy() + is_detector = self._direct_global_indices < self._num_detectors + self._direct_detector_f = self._direct_f_indices[is_detector] + self._direct_detector_flips = self._direct_flips[is_detector] + self._direct_detector_columns = self._direct_global_indices[is_detector] statevector = [ c for c in self._program.components if isinstance(c, StatevectorComponent) @@ -618,18 +623,12 @@ def _sample_batches_with_postselection( was_discarded = np.zeros(shots, dtype=np.bool_) shot_idx = 0 - detector = self._direct_global_indices < self._num_detectors - columns = self._direct_global_indices[detector] - direct = ( - self._direct_f_indices[detector], - self._direct_flips[detector], - columns, - postselection_mask[columns], - ( - reference[columns] - if xor_detector_ref and reference is not None - else np.zeros(len(columns), dtype=np.bool_) - ), + columns = self._direct_detector_columns + postselect = postselection_mask[columns] + reference_bits = ( + reference[columns] + if xor_detector_ref and reference is not None + else np.zeros(len(columns), dtype=np.bool_) ) sharing = self._sharing_arguments( @@ -656,7 +655,14 @@ def dispatch(f_batch: np.ndarray, shot_batch: np.ndarray) -> None: chunk = min(batch_size, shots - shot_idx) f_params_np = self._channel_sampler.sample(chunk) discarded = _write_direct_detectors( - f_params_np, *direct, result[shot_idx : shot_idx + chunk] + chunk * len(columns), + f_params_np, + self._direct_detector_f, + self._direct_detector_flips, + columns, + postselect, + reference_bits, + result[shot_idx : shot_idx + chunk], ) was_discarded[shot_idx : shot_idx + chunk] = discarded diff --git a/src/tsim/utils/kernels.py b/src/tsim/utils/kernels.py new file mode 100644 index 00000000..d84d1e8a --- /dev/null +++ b/src/tsim/utils/kernels.py @@ -0,0 +1,57 @@ +"""Numba kernels that run in parallel only when the work pays for it.""" + +import threading +import types +from collections.abc import Callable + +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() + + +def jit(function: Callable, *, parallel: bool = False) -> Callable: + """Compile ``function`` with numba, cached on disk where numba can write. + + 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. + """ + try: + return numba.njit(cache=True, parallel=parallel)(function) + except RuntimeError: + return numba.njit(parallel=parallel)(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. 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. + renamed = types.FunctionType( + function.__code__, + function.__globals__, + function.__name__ + "_parallel", + function.__defaults__, + function.__closure__, + ) + renamed.__qualname__ = function.__qualname__ + "_parallel" + parallel = jit(renamed, parallel=True) + + def run(work: int, *args): + if work < PARALLEL_WORK: + return serial(*args) + with _LOCK: + return parallel(*args) + + run.__doc__ = function.__doc__ + return run 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/unit/noise/test_channels.py b/test/unit/noise/test_channels.py index d93184b3..e017f4fc 100644 --- a/test/unit/noise/test_channels.py +++ b/test/unit/noise/test_channels.py @@ -6,12 +6,12 @@ _BLOCK_SHOTS, Channel, ChannelSampler, + _fire_tables, absorb_subset_channels, channel_num_bits, correlated_error_probs, error_probs, expand_channel, - fire_tables, fold_duplicate_channel_bits, heralded_pauli_channel_1_probs, merge_identical_channels, @@ -215,7 +215,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._tables = fire_tables(channels, sampler.signature_matrix) + sampler._tables = _fire_tables(channels, sampler.signature_matrix) return sampler.sample(n_samples) @@ -1059,6 +1059,25 @@ def test_two_bit_channel_joint_distribution_across_blocks(self): 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( diff --git a/test/unit/test_postselection.py b/test/unit/test_postselection.py index f3aa99f6..4360122b 100644 --- a/test/unit/test_postselection.py +++ b/test/unit/test_postselection.py @@ -162,9 +162,11 @@ def test_postselection_discarded_and_surviving_rows(): assert np.all(~samples[discarded, 1]) -# Detector 0 is postselected; direct detector 2 is not, so a survivor that is -# sampled from another shot's f-params shows up in column 2. -TWO_DIRECT_CIRCUIT = """ +# 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 @@ -174,14 +176,15 @@ def test_postselection_discarded_and_surviving_rows(): 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_detectors(shots, batch_size, reference): - """Every row's direct detectors come from that shot's own f-params.""" - sampler = _make(TWO_DIRECT_CIRCUIT, seed=11) +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) drawn: list[np.ndarray] = [] original = sampler._channel_sampler.sample @@ -196,12 +199,16 @@ def capture(n: int) -> np.ndarray: shots, batch_size=batch_size, postselection_mask=mask, + append_observables=True, use_detector_reference_sample=reference, ) expected = sampler._compute_direct_outputs(np.concatenate(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(): @@ -360,7 +367,8 @@ 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(): @@ -611,30 +619,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/utils/test_kernels.py b/test/unit/utils/test_kernels.py new file mode 100644 index 00000000..d24c9368 --- /dev/null +++ b/test/unit/utils/test_kernels.py @@ -0,0 +1,55 @@ +import subprocess +import sys +import textwrap + +import numba +import numpy as np + +from tsim.utils.kernels import PARALLEL_WORK, kernel + + +@kernel +def _square_roots(out): + for i in numba.prange(out.shape[0]): + out[i] = np.sqrt(i) + + +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() diff --git a/uv.lock b/uv.lock index 8082c79f..f5414d79 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.15'", @@ -1283,15 +1283,15 @@ wheels = [ [[package]] name = "griffe-kirin" -version = "0.1.0" -source = { registry = "https://quera.jfrog.io/artifactory/api/pypi/kirin/simple" } +version = "0.2.0" +source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "griffe" }, { name = "typing-extensions" }, ] -sdist = { url = "https://quera.jfrog.io/artifactory/api/pypi/kirin/griffe-kirin/0.1.0/griffe_kirin-0.1.0.tar.gz", hash = "sha256:b5ed85a8c07da0219787ac8d57de0a4f2a6df3109fc3aab47d78a12f96dda04d" } +sdist = { url = "https://files.pythonhosted.org/packages/a5/e0/9db75de974546903ab46a41179a95243abc84ceb2ba0b4101b5016f5301e/griffe_kirin-0.2.0.tar.gz", hash = "sha256:87d4280fbb1c77e1eb3fe7b627810b219295ec74651cf216748330b379a4ded5", size = 57664, upload-time = "2025-02-21T20:41:16.862Z" } wheels = [ - { url = "https://quera.jfrog.io/artifactory/api/pypi/kirin/griffe-kirin/0.1.0/griffe_kirin-0.1.0-py3-none-any.whl", hash = "sha256:cdbc7a08c4fa5229f8f5f9137fd4de7dde35214166d3d448cbac8778079bd10a" }, + { url = "https://files.pythonhosted.org/packages/38/b2/04bf90dc9e9e80f969b562731cf1619bd2eb0d0b726ba7befd38fda38131/griffe_kirin-0.2.0-py3-none-any.whl", hash = "sha256:81d999d6e1480cbe46e5c7b82933d2eee96863840dfbac283df620d85c13c2fc", size = 4063, upload-time = "2025-02-21T20:41:14.734Z" }, ] [[package]] From d06de41c88b4e7c39f545a4891491444020d7508 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 13 Oct 2026 01:10:30 -0400 Subject: [PATCH 05/21] Run trajectory sharing in a numba executor on the CPU The JAX shared executor kept static lane tables: kernels ran over every slot of a power-of-two table (26-40% live), the loop ran in segments with host checks, resized and resumed when a table filled, and compiled again for every new table size (about a second, randomly inside sample calls). On wide components that work dominated: 88% of a coherent d5r5 call. The executor is now a Python loop over the tape that keeps the lanes in growable NumPy arrays and calls small numba kernels: each action touches only the live lanes and the 2^extent amplitudes it can reach. extent comes from the plan (lower records one plus the highest active axis), and the planner now promotes onto the lowest free axis (a free SWAP of two dormant axes), so the extent tracks the active width: 1.8x less array work on coherent d5r5, 2.0x on SymFT's d=7 circuit. Rare outcomes are drawn as binomials by geometric skips with a 64-bit uniform, so the chance of a first hit stays exact below float32 resolution, as before. Grouping, splitting, pruning and reconstruction keep their semantics, and the exact distribution tests run unchanged against the new executor. The JAX executor stays for trajectory_sharing=False. shared_breakdown.py timed stages of the JAX executor and is removed. Quick benchmark, all threads, trajectory sharing: coherent d5r5 109k -> 436k shots/s, SymFT d7 postselected 241k -> 1.35M, cultivation d5 3.2M -> 8.6M, cultivation d3 5.1M -> 8.9M, distillation 16.3M -> 17.5M. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 44 +- src/tsim/statevector/plan.py | 9 +- src/tsim/statevector/program.py | 10 + src/tsim/statevector/shared.py | 994 ++++++++++++--------------- test/unit/statevector/test_shared.py | 130 ++-- 5 files changed, 510 insertions(+), 677 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index f55825dd..3f619553 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -20,7 +20,11 @@ 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, + sample_shared_component, + sharing_layout, +) from tsim.utils.cuda_helpers import copy_d2h from tsim.utils.kernels import kernel @@ -128,7 +132,7 @@ def sample_program( program: CompiledProgram, f_params: jax.Array, key: PRNGKey, - sharing: tuple[Sharing | None, ...] | None = None, + sharing: tuple[SharingLayout | None, ...] | None = None, sharing_stats: deque[dict] | None = None, ) -> jax.Array: """Sample all outputs from a compiled program. @@ -137,7 +141,7 @@ def sample_program( 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 + sharing: One sharing layout 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. @@ -147,7 +151,7 @@ def sample_program( match the original output indices. """ - results: list[jax.Array] = [] + results: list[jax.Array | np.ndarray] = [] if program.num_outputs == 0: batch_size = f_params.shape[0] @@ -257,7 +261,7 @@ def __init__( self._key = jax.random.key(seed) self._trajectory_sharing = trajectory_sharing - self._sharing: dict[bytes, tuple[Sharing | None, ...]] = {} + self._sharing: dict[bytes, tuple[SharingLayout | None, ...]] = {} self.sharing_stats: deque[dict] = deque(maxlen=10_000) prepared = prepare_graph(circuit, sample_detectors=sample_detectors) @@ -342,22 +346,20 @@ def _sharing_arguments( 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 - ] - ), - ) + 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 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..d42f098d 100644 --- a/src/tsim/statevector/program.py +++ b/src/tsim/statevector/program.py @@ -54,6 +54,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)``. @@ -72,6 +74,7 @@ class StatevectorComponent(eqx.Module): cos: Array sin: Array pivot: Array + extent: Array branch: Array sign_const: Array sign_row: Array @@ -145,15 +148,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,6 +167,9 @@ 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) @@ -176,6 +185,7 @@ def lower(plan: Plan, layout: ComponentLayout, dtype: str) -> StatevectorCompone cos=jnp.asarray(np.cos(angle).astype(real)), sin=jnp.asarray(np.sin(angle).astype(real)), pivot=jnp.asarray(pivot), + extent=jnp.asarray(extent), branch=jnp.asarray(branch), sign_const=jnp.asarray(sign_const), sign_row=jnp.asarray(sign_row), diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index f40b1751..4b960de2 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 @@ -11,38 +11,32 @@ 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. + +Lanes live in growable arrays, and a Python loop runs the tape: each action calls +a parallel numba kernel on the live lanes and on the ``2^extent`` amplitudes the +action can touch. 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; 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 from tsim.statevector.program import ( MEASURE, MEASURE_DIAG, @@ -51,15 +45,37 @@ 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 Tape(NamedTuple): + """A component's action tape as NumPy arrays (see ``StatevectorComponent``).""" + + 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 + f_selection: np.ndarray + + +class SharingLayout(NamedTuple): + """Which symbols the shots of a lane must share, and the tape they run. Attributes: + tape: The component's actions and outputs. 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)``. @@ -75,32 +91,18 @@ class SharingLayout(eqx.Module): 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. """ - 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 + tape: Tape + key_rows: np.ndarray + lane_dormant: np.ndarray + pruned: np.ndarray + prune_outputs: np.ndarray + prune_flip: np.ndarray + lane_branches: np.ndarray + shot_dormant: np.ndarray + ready: np.ndarray def _unpack(words: np.ndarray, num_words: int) -> np.ndarray: @@ -133,13 +135,13 @@ def sharing_layout( ``(num_outputs,)``. ``None`` flips nothing. """ + tape = Tape(*(np.array(getattr(component, field)) for field in Tape._fields)) num_words = component.num_words - num_f = len(component.f_selection) - kind = np.asarray(component.kind) - branch = np.asarray(component.branch) + num_f = len(tape.f_selection) + kind, branch = tape.kind, tape.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) + sign_bits = _unpack(tape.sign_row, num_words) + output_bits = _unpack(tape.output_row, num_words) num_outputs = component.num_outputs postselected = np.asarray( np.zeros(num_outputs) if postselected is None else postselected, dtype=bool @@ -159,7 +161,7 @@ def sharing_layout( 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) + pruned = np.asarray(prune_outputs.any(axis=1)) rotation = np.isin(kind, (ROTATE_DIAG, ROTATE)) keyed_bits = np.vstack([sign_bits[rotation], output_bits[postselected]]) @@ -168,7 +170,6 @@ 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)) reaches_key = np.asarray(np.any(keyed_bits, axis=0), dtype=bool) dormant = kind == MEASURE_DORMANT @@ -180,63 +181,263 @@ def sharing_layout( shot_bits[branch[dormant & ~lane_dormant & (branch >= 0)]] = True 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, + tape=tape, + key_rows=_pack(key_bits, num_words).reshape(-1, num_words), + lane_dormant=lane_dormant, + pruned=pruned, + prune_outputs=prune_outputs, + prune_flip=flipped, + lane_branches=_pack(lane_bits, num_words), + shot_dormant=_pack(shot_bits, num_words), + ready=ready.astype(np.int64), ) -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 + + +@kernel +def _row_parities(symbols, rows): + """Return ``parity(symbols[i] & rows[r])`` for every row i and r.""" + out = np.zeros((symbols.shape[0], rows.shape[0]), dtype=np.bool_) + for i in numba.prange(symbols.shape[0]): + for r in range(rows.shape[0]): + acc = 0 + for w in range(rows.shape[1]): + acc ^= symbols[i, w] & rows[r, w] + out[i, r] = _parity(acc) + return out + + +@kernel +def _key_words(symbols, key_rows): + """Pack the parities of each shot's symbols with the key rows into words.""" + num_keys = key_rows.shape[0] + out = np.zeros((symbols.shape[0], (num_keys + 31) >> 5), dtype=np.uint32) + for s in numba.prange(symbols.shape[0]): + for k in range(num_keys): + 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) + return out + + +@kernel +def _pack_symbols(f, selection, num_words): + """Pack each shot's selected f-bits: bit j is bit j % 32 of word j // 32.""" + out = np.zeros((f.shape[0], num_words), dtype=np.uint32) + for s in numba.prange(f.shape[0]): + for j in range(len(selection)): + if f[s, selection[j]]: + out[s, j >> 5] |= np.uint32(1) << np.uint32(j & 31) + return out + + +@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 + + +@kernel +def _rotate_diag(amps, lanes, extent, z, symbols, sign_row, sign_const, factor): + """Multiply amplitude i of lane l by ``factor[sign(l) ^ parity(i & z)]``.""" + flip = _signs(symbols, lanes, sign_row, sign_const) + span = 1 << extent + for t in numba.prange(lanes * span): + lane, i = t >> extent, t & (span - 1) + amps[lane, i] *= factor[flip[lane] ^ _parity(i & z)] + + +@kernel +def _rotate_pair( + amps, lanes, extent, x, z, symbols, sign_row, sign_const, cos, coupling +): + """Rotate about ``X^x Z^z``: mix amplitudes i and ``i ^ x`` of every lane.""" + flip = _signs(symbols, lanes, sign_row, sign_const) + span = 1 << extent + low = x & -x + for t in numba.prange(lanes * span): + lane, i = t >> extent, t & (span - 1) + if i & low: + continue + j = i ^ x + c = -coupling if flip[lane] else coupling + a, b = amps[lane, i], amps[lane, j] + amps[lane, i] = cos * a + c * (-b if _parity(j & z) else b) + amps[lane, j] = cos * b + c * (-a if _parity(i & z) else a) + + +@kernel +def _weights(amps, lanes, extent, x, z, i_power): + """Return each lane's outcome weights ``(w0, w1)`` of measuring ``X^x Z^z``.""" + span = 1 << extent + w0 = np.zeros(lanes) + w1 = np.zeros(lanes) + for lane in numba.prange(lanes): + for i in range(span): + a = np.complex128(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 * np.complex128(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 + return w0, w1 + + +@kernel +def _project(amps, lanes, extent, x, z, pivot, i_power, outcome, weight): + """Project lane l onto ``outcome[l]`` of weight ``weight[l]``; 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⟩``. """ + span = 1 << extent + bit = 1 << pivot + for lane in numba.prange(lanes): + 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 + + +@kernel +def _shot_symbols(symbols, order, owner, entry_symbols, lane_branches, dormant, seed): + """Return every shot's symbols: its f-bits, its lane's branches, its coins.""" + out = np.empty_like(symbols) + num_words = symbols.shape[1] + for position in numba.prange(len(order)): + 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 + return out + + +class _Lanes: + """Live lanes in growable arrays: rows ``0 … size-1`` are live.""" + + def __init__(self, symbols, start, count, width: int, dtype: np.dtype): + self.size = len(start) + capacity = max(16, 2 * self.size) + self.amps = np.zeros((capacity, 1 << width), dtype=dtype) + self.amps[: self.size, 0] = 1 + self.symbols = np.zeros((capacity, symbols.shape[1]), dtype=np.uint32) + self.symbols[: self.size] = symbols + self.start = np.zeros(capacity, dtype=np.int64) + self.start[: self.size] = start + self.count = np.zeros(capacity, dtype=np.int64) + self.count[: self.size] = count + + def split(self, ones: np.ndarray, span: int) -> tuple[np.ndarray, np.ndarray]: + """Split every lane whose shots see both outcomes. + + Returns each lane's outcome and the lane its amplitudes came from (its + parent for a new child). + """ + size, count = self.size, self.count[: self.size] + outcome = ones == count + parents = np.flatnonzero((ones > 0) & ~outcome) + new = size + len(parents) + if new > len(self.count): + self._grow(new) + self.amps[size:new, :span] = self.amps[parents, :span] + self.symbols[size:new] = self.symbols[parents] + self.count[size:new] = ones[parents] + self.start[size:new] = self.start[parents] + count[parents] - ones[parents] + self.count[parents] -= ones[parents] + self.size = new + children = np.ones(len(parents), dtype=bool) + return np.concatenate([outcome, children]), np.concatenate( + [np.arange(size), parents] + ) - 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 prune(self, dies: np.ndarray, span: int) -> tuple[np.ndarray, ...]: + """Remove the lanes that ``dies`` marks; return their starts, counts, symbols.""" + gone, kept = np.flatnonzero(dies), np.flatnonzero(~dies) + removed = (self.start[gone], self.count[gone], self.symbols[gone]) + self.amps[: len(kept), :span] = self.amps[kept, :span] + for array in (self.symbols, self.start, self.count): + array[: len(kept)] = array[kept] + self.size = len(kept) + return removed + + def _grow(self, needed: int) -> None: + capacity = max(needed, 2 * len(self.count)) + for name in ("amps", "symbols", "start", "count"): + old = getattr(self, name) + new = np.zeros((capacity, *old.shape[1:]), dtype=old.dtype) + new[: self.size] = old[: self.size] + setattr(self, name, new) def group_on_host( @@ -245,9 +446,7 @@ def group_on_host( """Sort shots by key and return ``(order, start, count)`` of each group. 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) @@ -258,479 +457,144 @@ def group_on_host( 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. - """ - 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) - - -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) - - -def _split(lanes: Lanes, ones: Array) -> tuple[Lanes, Array, Array]: - """Split every live lane whose shots see both outcomes. - - 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) - ) - start = jnp.where( - is_child, lanes.start[source] + parent_count - parent_ones, lanes.start - ) - 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] - ) - 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 - - -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), - ) - 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 + return order, start, count + + +def _run( + layout: SharingLayout, lanes: _Lanes, rng: np.random.Generator +) -> tuple[float, list[tuple], int]: + """Run the tape on ``lanes``; return the norm deviation, dead lanes and peak.""" + tape = layout.tape + dtype = lanes.amps.dtype + dead: list[tuple] = [] + deviation, peak, region = 0.0, lanes.size, 0 + for a, kind in enumerate(tape.kind): + extent = int(tape.extent[a]) + span = 1 << extent + if extent > region: + lanes.amps[: lanes.size, 1 << region : span] = 0 + region = extent + x, z = int(tape.x_mask[a]), int(tape.z_mask[a]) + i_power = dtype.type(_I_POWERS[(x & z).bit_count() & 3]) + size = lanes.size + work = size * span + if kind in (ROTATE_DIAG, ROTATE): + sign = (lanes.symbols, tape.sign_row[a], int(tape.sign_const[a])) + cos, sin = tape.cos[a], tape.sin[a] + if kind == ROTATE_DIAG: + factor = np.array([cos - 1j * sin, cos + 1j * sin], dtype=dtype) + _rotate_diag(work, lanes.amps, size, extent, z, *sign, factor) + else: + coupling = dtype.type(-1j * sin * i_power) + _rotate_pair(work, lanes.amps, size, extent, x, z, *sign, cos, coupling) + continue + if kind == MEASURE_DORMANT: + if not layout.lane_dormant[a]: + continue + w0 = w1 = np.full(size, 0.5) + else: + w0, w1 = _weights(work, lanes.amps, size, extent, x, z, i_power) + deviation = max(deviation, float(np.max(np.abs(w0 + w1 - 1)))) + seed = rng.integers(0, 2**64, dtype=np.uint64) + ones = _draw_ones(lanes.count[:size], w0, w1, seed) + outcome, source = lanes.split(ones, span) + if kind != MEASURE_DORMANT: + weight = np.where(outcome, w1[source], w0[source]) + _project( + lanes.size * span, + lanes.amps, + lanes.size, + extent, + x, + z, + int(tape.pivot[a]), + i_power, + outcome, + weight, + ) + branch = int(tape.branch[a]) + lanes.symbols[: lanes.size, branch >> 5] |= outcome.astype(np.uint32) << ( + branch & 31 ) - 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_), - ) - - -@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) - - -def _capacity(needed: int) -> int: - return max(_MIN_CAPACITY, 1 << int(np.ceil(np.log2(max(needed, 1))))) - - -@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) + peak = max(peak, lanes.size) + if layout.pruned[a]: + values = _row_parities(0, lanes.symbols[: lanes.size], tape.output_row) + values ^= tape.output_const ^ layout.prune_flip + dies = np.any(values & layout.prune_outputs[a], axis=1) + if dies.any(): + dead.append((*lanes.prune(dies, span), a)) + return deviation, dead, peak def sample_shared_component( component: StatevectorComponent, - sharing: Sharing, - f_params: Array, - key: Array, - *, - segment: int | None = None, -) -> tuple[Array, Array, Array, dict]: + layout: SharingLayout, + f_params: jax.Array | np.ndarray, + key: jax.Array, +) -> tuple[np.ndarray, jax.Array, float, dict]: """Sample one component's outputs with shared trajectories. 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. Returns: ``(samples, next_key, 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 + tape = layout.tape + key, subkey = jax.random.split(key) + rng = np.random.default_rng(np.asarray(jax.random.bits(subkey, (4,)))) + f = np.asarray(f_params) + symbols = _pack_symbols(f.size, f, tape.f_selection, component.num_words) + work = len(symbols) * max(1, layout.key_rows.size) + keys = _key_words(work, symbols, layout.key_rows) + order, start, count = group_on_host(keys, rng) + representative = symbols[order[start]] + lanes = _Lanes( + representative, start, count, component.width, np.dtype(component.dtype) ) - 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, + deviation, dead, peak = _run(layout, lanes, rng) + + num_actions = len(tape.kind) + live = slice(0, lanes.size) + entries = [ + (lanes.start[live], lanes.count[live], lanes.symbols[live], num_actions), + *dead, + ] + entry_start = np.concatenate([e[0] for e in entries]) + entry_count = np.concatenate([e[1] for e in entries]) + entry_symbols = np.concatenate([e[2] for e in entries]) + entry_pruned = np.concatenate([np.full(len(e[0]), e[3]) for e in entries]) + # NOTE: the entries' ranges tile the sorted positions, so sorting them by + # start gives every position's entry. + by_start = np.argsort(entry_start) + owner = np.repeat(by_start, entry_count[by_start]) + coin_seed = rng.integers(0, 2**64, dtype=np.uint64) + shot_symbols = _shot_symbols( + symbols.size, + symbols, order, - jnp.asarray(start_padded), - jnp.asarray(count_padded), - jnp.asarray(groups, dtype=jnp.int32), - capacity=capacity, + owner, + entry_symbols, + layout.lane_branches, + layout.shot_dormant, + coin_seed, ) - 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) + values = _row_parities( + shot_symbols.size * len(tape.output_row), shot_symbols, tape.output_row + ) + values ^= tape.output_const + pruned_at = np.empty(len(order), dtype=np.int64) + pruned_at[order] = entry_pruned[owner] + samples = values & (layout.ready[None, :] <= pruned_at[:, None]) stats = { - "groups": groups, + "groups": len(start), "peak lanes": peak, - "lanes": used, - "dead": dead_used, - "lane_capacity": lanes.count.shape[0], - "resumes": resumes, + "lanes": lanes.size, + "dead": sum(len(e[0]) for e in dead), } - return samples, key, lanes.max_norm_deviation, stats + return samples, key, deviation, stats diff --git a/test/unit/statevector/test_shared.py b/test/unit/statevector/test_shared.py index b77b4bfb..11c4b289 100644 --- a/test/unit/statevector/test_shared.py +++ b/test/unit/statevector/test_shared.py @@ -14,11 +14,7 @@ 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, -) +from tsim.statevector.shared import sample_shared_component, sharing_layout NAMES = [ "t_ghz", @@ -81,11 +77,11 @@ def _assert_counts(rows, discarded, expected): def test_shared_sampling_matches_the_plan_distribution(name): 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) + compiled, layout, jnp.asarray(f_params), jax.random.key(7) ) assert float(deviation) < 1e-4 assert_matches_distribution( @@ -102,7 +98,7 @@ def test_mixed_noise_patterns_share_one_batch(name): permutation = np.random.default_rng(1).permutation(len(f_params)) samples, _, _, _ = sample_shared_component( compiled, - Sharing(sharing_layout(compiled)), + sharing_layout(compiled), jnp.asarray(f_params[permutation]), jax.random.key(3), ) @@ -121,11 +117,11 @@ def test_pruning_discards_exactly_the_postselected_failures(name, flip): 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) + compiled, layout, jnp.asarray(f_params), jax.random.key(5) ) samples = np.asarray(samples) for i, f in enumerate(patterns): @@ -141,51 +137,22 @@ def test_pruning_discards_exactly_the_postselected_failures(name, flip): def test_pruned_shots_report_false_after_their_pruning_action(name): for planned, component_layout, compiled in _compiled(name): postselected = np.ones(compiled.num_outputs, dtype=bool) - sharing = Sharing(sharing_layout(compiled, postselected)) + layout = sharing_layout(compiled, postselected) 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) + compiled, layout, jnp.asarray(f_params), jax.random.key(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 @@ -193,58 +160,49 @@ def test_capacities_and_segments_do_not_change_the_samples(name, monkeypatch): f_params = jnp.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 + first, next_key, _, _ = sample_shared_component( + compiled, layout, f_params, jax.random.key(13) ) - monkeypatch.setattr(shared, "_MIN_CAPACITY", 1) - tiny, _, _, narrow = sample_shared_component( - compiled, Sharing(layout), f_params, key, segment=1 + second, _, _, _ = sample_shared_component( + compiled, layout, f_params, jax.random.key(13) ) - 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 + third, _, _, _ = sample_shared_component(compiled, layout, f_params, next_key) + 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 From 8de99dd933b1a9ec3b0f8212a9f5d6ca364e3ecf Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 13 Oct 2026 01:54:28 -0400 Subject: [PATCH 06/21] Fix the issues the review of the numba shared executor found - Once postselection had pruned every lane, the next active measurement took np.max of an empty array and raised; tiny batches, high rejection and zero shots hit it. The tape now stops when no lanes are left, and the norm deviation is a NaN-propagating max with an initial value, which the sampler now also warns about. - The lane table allocated at least 16 rows of 2^width amplitudes even for a few shots and copied full rows when growing. A lane holds at least one shot, so the table is capped at the shot count, and growing copies only the amplitudes the current extent reaches. - The JAX executor's per-lane shot count and choose_outcome's count argument existed only so that a future sharing executor could split lanes; both are gone. - Tests: a component whose extent drops and rises again (a generated d5r3 surface code with coherent over-rotations), sampled with NaN in every amplitude outside the tracked region, matches a run at full width bitwise; every kernel in parallel matches the serial run bitwise; shots whose lanes all die, and zero shots; every action stays within its extent. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 3 +- src/tsim/statevector/execute.py | 19 ++--- src/tsim/statevector/shared.py | 28 ++++-- test/unit/statevector/test_execute.py | 15 ++-- test/unit/statevector/test_shared.py | 117 +++++++++++++++++++++++++- 5 files changed, 150 insertions(+), 32 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 3f619553..701524a4 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -182,7 +182,8 @@ def sample_program( "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: + # NOTE: written so that a NaN deviation also warns. + if not 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}). " diff --git a/src/tsim/statevector/execute.py b/src/tsim/statevector/execute.py index e24e02b6..3826b3f0 100644 --- a/src/tsim/statevector/execute.py +++ b/src/tsim/statevector/execute.py @@ -5,8 +5,7 @@ ``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. +``choose_outcome``. """ from __future__ import annotations @@ -39,14 +38,12 @@ class LaneState(NamedTuple): 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 @@ -57,7 +54,7 @@ def lane_bytes(component: StatevectorComponent) -> int: return amplitudes + 4 * component.num_words + component.num_outputs -def choose_outcome(p0: Array, p1: Array, count: Array, key: Array) -> Array: +def choose_outcome(p0: Array, p1: 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 @@ -65,10 +62,8 @@ def choose_outcome(p0: Array, p1: Array, count: Array, key: Array) -> Array: ``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. + probability 0 is never drawn. """ - 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 @@ -97,7 +92,6 @@ def initial_state(component: StatevectorComponent, f_bits: Array) -> LaneState: 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), ) @@ -185,7 +179,7 @@ def _sample_outcome( """ 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) + outcome = choose_outcome(weight_0 / total, weight_1 / total, key) return outcome, jnp.where(outcome, weight_1, weight_0) @@ -244,8 +238,9 @@ def _measure_diag( 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)) + lanes = state.symbols.shape[0] + half = jnp.full(lanes, 0.5, dtype=state.max_norm_deviation.dtype) + outcome = choose_outcome(half, half, jax.random.fold_in(key, a)) return _set_branch(component, state, a, outcome) diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index 4b960de2..df1d8786 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -21,8 +21,8 @@ it is a fair coin drawn per shot at the end. Lanes live in growable arrays, and a Python loop runs the tape: each action calls -a parallel numba kernel on the live lanes and on the ``2^extent`` amplitudes the -action can touch. When the extent grows, the newly touched amplitudes are +a numba kernel on the live lanes and on the ``2^extent`` amplitudes the action +can touch. 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; a pruned shot's outputs that need a branch @@ -384,11 +384,16 @@ def _shot_symbols(symbols, order, owner, entry_symbols, lane_branches, dormant, class _Lanes: - """Live lanes in growable arrays: rows ``0 … size-1`` are live.""" + """Live lanes in growable arrays: rows ``0 … size-1`` are live. + + Every lane holds at least one shot, so the arrays never need more rows than + there are shots. + """ def __init__(self, symbols, start, count, width: int, dtype: np.dtype): self.size = len(start) - capacity = max(16, 2 * self.size) + self.shots = max(1, int(np.sum(count))) + capacity = min(max(16, 2 * self.size), self.shots) self.amps = np.zeros((capacity, 1 << width), dtype=dtype) self.amps[: self.size, 0] = 1 self.symbols = np.zeros((capacity, symbols.shape[1]), dtype=np.uint32) @@ -409,7 +414,7 @@ def split(self, ones: np.ndarray, span: int) -> tuple[np.ndarray, np.ndarray]: parents = np.flatnonzero((ones > 0) & ~outcome) new = size + len(parents) if new > len(self.count): - self._grow(new) + self._grow(new, span) self.amps[size:new, :span] = self.amps[parents, :span] self.symbols[size:new] = self.symbols[parents] self.count[size:new] = ones[parents] @@ -431,9 +436,12 @@ def prune(self, dies: np.ndarray, span: int) -> tuple[np.ndarray, ...]: self.size = len(kept) return removed - def _grow(self, needed: int) -> None: - capacity = max(needed, 2 * len(self.count)) - for name in ("amps", "symbols", "start", "count"): + def _grow(self, needed: int, span: int) -> None: + capacity = min(max(needed, 2 * len(self.count)), self.shots) + amps = np.zeros((capacity, self.amps.shape[1]), dtype=self.amps.dtype) + amps[: self.size, :span] = self.amps[: self.size, :span] + self.amps = amps + for name in ("symbols", "start", "count"): old = getattr(self, name) new = np.zeros((capacity, *old.shape[1:]), dtype=old.dtype) new[: self.size] = old[: self.size] @@ -469,6 +477,8 @@ def _run( dead: list[tuple] = [] deviation, peak, region = 0.0, lanes.size, 0 for a, kind in enumerate(tape.kind): + if not lanes.size: + break extent = int(tape.extent[a]) span = 1 << extent if extent > region: @@ -494,7 +504,7 @@ def _run( w0 = w1 = np.full(size, 0.5) else: w0, w1 = _weights(work, lanes.amps, size, extent, x, z, i_power) - deviation = max(deviation, float(np.max(np.abs(w0 + w1 - 1)))) + deviation = float(np.max(np.abs(w0 + w1 - 1), initial=deviation)) seed = rng.integers(0, 2**64, dtype=np.uint64) ones = _draw_ones(lanes.count[:size], w0, w1, seed) outcome, source = lanes.split(ones, span) diff --git a/test/unit/statevector/test_execute.py b/test/unit/statevector/test_execute.py index cb30ea1f..bec80346 100644 --- a/test/unit/statevector/test_execute.py +++ b/test/unit/statevector/test_execute.py @@ -115,7 +115,7 @@ def test_measurements_match_the_reference_kernel(monkeypatch, x, z, pivot): 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) + monkeypatch.setattr(execute, "choose_outcome", lambda p0, p1, 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) @@ -131,7 +131,7 @@ def test_outputs_are_parities_of_the_symbols(monkeypatch): outputs=[AffineRow(True, frozenset({1})), AffineRow(False, frozenset({0, 1}))], ) monkeypatch.setattr( - execute, "choose_outcome", lambda p0, p1, count, key: jnp.array([False, True]) + execute, "choose_outcome", lambda p0, p1, key: jnp.array([False, True]) ) state = initial_state(component, jnp.array([[True], [True]])) out = run(component, state, jax.random.key(0), 0, 1) @@ -153,7 +153,7 @@ def test_outputs_read_symbols_beyond_the_first_word(monkeypatch): ) assert component.num_words == 2 monkeypatch.setattr( - execute, "choose_outcome", lambda p0, p1, count, key: jnp.array([False, True]) + execute, "choose_outcome", lambda p0, p1, key: jnp.array([False, True]) ) f_bits = np.zeros((2, 40), dtype=bool) f_bits[1, 3] = True @@ -225,8 +225,7 @@ def _count_outcomes(p1, dtype, draws, chunk=2**22): 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 jnp.sum(execute.choose_outcome(p0s, p1s, key)) return sum( int(ones_in_chunk(jax.random.key(seed))) for seed in range(draws // chunk) @@ -264,9 +263,9 @@ def _run_recording_probabilities(monkeypatch, component, state): recorded = [] choose = execute.choose_outcome - def recording(p0, p1, count, key): + def recording(p0, p1, key): recorded.append((np.asarray(p0, np.float64), np.asarray(p1, np.float64))) - return choose(p0, p1, count, key) + return choose(p0, p1, key) monkeypatch.setattr(execute, "choose_outcome", recording) with jax.disable_jit(): @@ -323,7 +322,7 @@ def test_an_unlikely_outcome_leaves_a_normalized_state( 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) + monkeypatch.setattr(execute, "choose_outcome", lambda p0, p1, 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 diff --git a/test/unit/statevector/test_shared.py b/test/unit/statevector/test_shared.py index 11c4b289..2bca8963 100644 --- a/test/unit/statevector/test_shared.py +++ b/test/unit/statevector/test_shared.py @@ -11,9 +11,11 @@ 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 +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 = [ @@ -28,7 +30,11 @@ 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) @@ -206,3 +212,110 @@ def test_a_lone_shot_sees_an_outcome_of_probability_below_float32_resolution(): 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.extent) + first_drop = np.flatnonzero(np.diff(extent) < 0)[0] + assert np.any(np.diff(extent[first_drop:]) > 0) + return planned, component_layout, compiled + + +class _PoisonedLanes(shared._Lanes): + """Lanes whose amplitudes outside the region the executor tracks are NaN.""" + + def __init__(self, *args): + super().__init__(*args) + self.amps[:, 1:] = np.nan + self.amps[self.size :] = np.nan + + def _grow(self, needed, span): + super()._grow(needed, span) + self.amps[: self.size, span:] = np.nan + self.amps[self.size :] = np.nan + + +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)) + width = np.full_like(layout.tape.extent, compiled.width) + full = layout._replace(tape=layout.tape._replace(extent=width)) + f_params = jnp.asarray( + _f_params(component_layout, f_patterns(planned.num_f, 7), 300) + ) + key = jax.random.key(17) + reference, _, _, _ = sample_shared_component(compiled, full, f_params, key) + monkeypatch.setattr(shared, "_Lanes", _PoisonedLanes) + poisoned, _, _, stats = sample_shared_component(compiled, layout, f_params, key) + np.testing.assert_array_equal(poisoned, reference) + 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 = jnp.asarray( + _f_params(component_layout, f_patterns(planned.num_f, 3), 1_000) + ) + key = jax.random.key(19) + serial, _, _, _ = sample_shared_component(compiled, layout, f_params, key) + monkeypatch.setattr(kernels, "PARALLEL_WORK", 0) + parallel, _, _, _ = sample_shared_component(compiled, layout, f_params, key) + monkeypatch.undo() + np.testing.assert_array_equal(parallel, serial) + + +@pytest.mark.parametrize("name", NAMES) +def test_shots_whose_lanes_all_die_are_still_sampled(name): + for planned, component_layout, compiled in _compiled(name): + layout = sharing_layout(compiled, np.ones(compiled.num_outputs, dtype=bool)) + f_params = jnp.asarray( + _f_params(component_layout, f_patterns(planned.num_f, 1)[:1], 1) + ) + for seed in range(40): + samples, _, _, _ = sample_shared_component( + compiled, layout, f_params, jax.random.key(seed) + ) + assert samples.shape == (1, compiled.num_outputs) + empty = f_params[:0] + samples, _, _, _ = sample_shared_component( + compiled, layout, empty, jax.random.key(0) + ) + assert samples.shape == (0, compiled.num_outputs) + + +@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: + reach = 1 << np.asarray(compiled.extent, dtype=np.int64) + support = np.asarray(compiled.x_mask) | np.asarray(compiled.z_mask) + assert np.all(support < reach) + measured = np.isin(np.asarray(compiled.kind), (MEASURE, MEASURE_DIAG)) + assert np.all((1 << np.asarray(compiled.pivot))[measured] < reach[measured]) From 57e670d0d6b3c0431e3e33aca049aa74f7dd0111 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 13 Oct 2026 02:09:16 -0400 Subject: [PATCH 07/21] Strengthen the shared executor's tests as their review suggested - The extent test missed a clearing bug after the first rise: the generated circuit's low-extent stretches hold only rotations, so no lane moved while the extent was low. It now also runs with every fifth extent raised to the full width (an extent is only an upper bound), poisons the lane table with NaN after every split, prune and growth, and compares the norm deviations bitwise. - The parallel-vs-serial test's baseline ran the wide calls in parallel too; it now forces every kernel serial. - The all-dead test checks that some shot's lanes did all die. - A test checks that a NaN normalization deviation warns. - The lane table's shot cap no longer needs max(1, ...), and the module docstring is reflowed. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/shared.py | 12 ++-- test/unit/statevector/test_shared.py | 97 +++++++++++++++++++--------- test/unit/test_sampler.py | 16 +++++ 3 files changed, 87 insertions(+), 38 deletions(-) diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index df1d8786..fba02304 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -22,11 +22,11 @@ Lanes live in growable arrays, and a Python loop runs the tape: each action calls a numba kernel on the live lanes and on the ``2^extent`` amplitudes the action -can touch. 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; a pruned shot's outputs that need a branch -bit set after its pruning action are False. +can touch. 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; a pruned shot's outputs that need a branch bit set +after its pruning action are False. """ from __future__ import annotations @@ -392,7 +392,7 @@ class _Lanes: def __init__(self, symbols, start, count, width: int, dtype: np.dtype): self.size = len(start) - self.shots = max(1, int(np.sum(count))) + self.shots = int(np.sum(count)) capacity = min(max(16, 2 * self.size), self.shots) self.amps = np.zeros((capacity, 1 << width), dtype=dtype) self.amps[: self.size, 0] = 1 diff --git a/test/unit/statevector/test_shared.py b/test/unit/statevector/test_shared.py index 2bca8963..f50c6e67 100644 --- a/test/unit/statevector/test_shared.py +++ b/test/unit/statevector/test_shared.py @@ -244,15 +244,31 @@ def coherent(): class _PoisonedLanes(shared._Lanes): - """Lanes whose amplitudes outside the region the executor tracks are NaN.""" + """Lanes that hold NaN wherever the executor must not read. + + After every operation, that is every row from ``size`` up and every + amplitude of a live row from ``span`` up. + """ def __init__(self, *args): super().__init__(*args) - self.amps[:, 1:] = np.nan - self.amps[self.size :] = np.nan + self._poison(1) + + def split(self, ones, span): + result = super().split(ones, span) + self._poison(span) + return result + + def prune(self, dies, span): + removed = super().prune(dies, span) + self._poison(span) + return removed def _grow(self, needed, span): super()._grow(needed, span) + self._poison(span) + + def _poison(self, span): self.amps[: self.size, span:] = np.nan self.amps[self.size :] = np.nan @@ -260,17 +276,29 @@ def _grow(self, needed, span): 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)) - width = np.full_like(layout.tape.extent, compiled.width) - full = layout._replace(tape=layout.tape._replace(extent=width)) + extent = layout.tape.extent + + def with_extent(values): + return layout._replace(tape=layout.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 = jnp.asarray( _f_params(component_layout, f_patterns(planned.num_f, 7), 300) ) key = jax.random.key(17) - reference, _, _, _ = sample_shared_component(compiled, full, f_params, key) + reference, _, deviation, _ = sample_shared_component(compiled, full, f_params, key) monkeypatch.setattr(shared, "_Lanes", _PoisonedLanes) - poisoned, _, _, stats = sample_shared_component(compiled, layout, f_params, key) - np.testing.assert_array_equal(poisoned, reference) - assert stats["dead"] > 0 + for tape in (layout, raised): + samples, _, other, stats = sample_shared_component( + compiled, tape, f_params, key + ) + np.testing.assert_array_equal(samples, reference) + assert other == deviation + assert stats["dead"] > 0 @pytest.mark.parametrize("name", [*NAMES, "coherent"]) @@ -283,31 +311,36 @@ def test_parallel_kernels_give_the_serial_samples(name, coherent, monkeypatch): f_params = jnp.asarray( _f_params(component_layout, f_patterns(planned.num_f, 3), 1_000) ) - key = jax.random.key(19) - serial, _, _, _ = sample_shared_component(compiled, layout, f_params, key) - monkeypatch.setattr(kernels, "PARALLEL_WORK", 0) - parallel, _, _, _ = sample_shared_component(compiled, layout, f_params, key) - monkeypatch.undo() - np.testing.assert_array_equal(parallel, serial) - - -@pytest.mark.parametrize("name", NAMES) -def test_shots_whose_lanes_all_die_are_still_sampled(name): - for planned, component_layout, compiled in _compiled(name): - layout = sharing_layout(compiled, np.ones(compiled.num_outputs, dtype=bool)) - f_params = jnp.asarray( - _f_params(component_layout, f_patterns(planned.num_f, 1)[:1], 1) - ) - for seed in range(40): + 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, jax.random.key(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 = jnp.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, jax.random.key(seed) + ) + assert samples.shape == (1, compiled.num_outputs) + all_died |= stats["lanes"] == 0 samples, _, _, _ = sample_shared_component( - compiled, layout, f_params, jax.random.key(seed) + compiled, layout, f_params[:0], jax.random.key(0) ) - assert samples.shape == (1, compiled.num_outputs) - empty = f_params[:0] - samples, _, _, _ = sample_shared_component( - compiled, layout, empty, jax.random.key(0) - ) - assert samples.shape == (0, compiled.num_outputs) + assert samples.shape == (0, compiled.num_outputs) + assert all_died @pytest.mark.parametrize("name", [*NAMES, "coherent"]) diff --git a/test/unit/test_sampler.py b/test/unit/test_sampler.py index daa136b1..129794e8 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,18 @@ 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_component + + def nan_deviation(component, f_params, key): + samples, key, _ = original(component, f_params, key) + return samples, key, float("nan") + + monkeypatch.setattr(sampler_module, "sample_component", nan_deviation) + with pytest.warns(UserWarning, match="not normalized"): + sampler.sample(10) From a0622089fb1ad79a3f9dbd20df27cf3dc4929256 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 13 Oct 2026 08:56:05 -0400 Subject: [PATCH 08/21] Speed up the noise model setup of large circuits On the d=5 end-to-end cultivation circuit (463 qubits, 184,039 error bits) building a sampler took 101 s, 80 s of it in ChannelSampler: - np.unique(error_transform, axis=1) sorted all 184,039 columns of the 2,821-row error transform, comparing whole columns: 66 s. unique_columns finds the 15,148 distinct columns with a dict over the bit-packed columns and sorts only those; big-endian packing keeps the lexicographic order, so the result equals np.unique's. - absorb_subset_channels compared every pair of the 21,341 channels: 14 s. A channel of at most max_bits signatures has at most 15 strict subsets, so they are looked up in a dict, and absorbed in the original order. Both give the same noise model as before, bitwise, and the same samples for a fixed seed (checked on distillation, cultivation d3/d5, coherent d3r3 and the end-to-end circuit). Building the end-to-end sampler now takes 24 s, 17 s of it in zxparam's full_reduce. Tests compare unique_columns with np.unique and the indexed absorption with the pairwise one. Co-Authored-By: Claude Opus 5.5 --- src/tsim/noise/channels.py | 66 +++++++++++++++++++++++++------- test/unit/noise/test_channels.py | 55 ++++++++++++++++++++++++++ 2 files changed, 108 insertions(+), 13 deletions(-) diff --git a/src/tsim/noise/channels.py b/src/tsim/noise/channels.py index 664c4209..0678c0ab 100644 --- a/src/tsim/noise/channels.py +++ b/src/tsim/noise/channels.py @@ -4,6 +4,7 @@ from collections import defaultdict from dataclasses import dataclass +from itertools import combinations from typing import NamedTuple import numba @@ -469,6 +470,12 @@ def absorb_subset_channels(channels: list[Channel], max_bits: int = 4) -> list[C # Sort by number of bits (largest first) for efficient processing 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() @@ -476,22 +483,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) @@ -523,6 +533,36 @@ 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)`` for a binary matrix. + + ``np.unique`` sorts every column, comparing whole columns, which takes about a + minute for the 2,821 x 184,039 error transform of a d=5 end-to-end + cultivation circuit. 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. + + Returns: + The distinct columns in lexicographic order, and the index of each + column of ``matrix`` among them. + + """ + packed = np.ascontiguousarray(np.packbits(matrix, axis=0).T) + ids: dict[bytes, int] = {} + seen = np.array( + [ids.setdefault(column.tobytes(), len(ids)) for column in packed], + dtype=np.int64, + ) + _, first = np.unique(seen, 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[seen] + + class ChannelSampler: """Samples from multiple error channels and transforms to a reduced basis. @@ -570,7 +610,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) diff --git a/test/unit/noise/test_channels.py b/test/unit/noise/test_channels.py index e017f4fc..4245a89f 100644 --- a/test/unit/noise/test_channels.py +++ b/test/unit/noise/test_channels.py @@ -20,6 +20,7 @@ pauli_channel_2_probs, reduce_null_bits, simplify_channels, + unique_columns, xor_convolve, ) @@ -751,6 +752,60 @@ 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=4): + """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("seed", range(5)) + def test_random_channels(self, seed): + rng = np.random.default_rng(seed) + channels = [] + for _ in range(60): + size = int(rng.integers(1, 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) + expected = _absorb_by_comparing_every_pair(channels) + 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)) + + class TestSimplifyChannels: """Tests for the full simplify_channels function.""" From acedb3c96a66f57a596d1a38d8a547de49f899c5 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 13 Oct 2026 09:07:38 -0400 Subject: [PATCH 09/21] Polish the noise setup code as its review suggested - unique_columns' docstring states the 1-D inverse (np.unique returned a 2-D one on numpy 2.0) and that nonzero entries count as 1, and drops benchmark figures that would go stale; `seen` is now `column_ids`. - The absorption test also draws channels without signatures and runs with max_bits 2 and 4; dropping the empty subset now fails it. - absorb_subset_channels' docstring and sort comment say what holds and why the order matters, and ChannelSampler's simplification steps are numbered 1 to 4. Co-Authored-By: Claude Opus 5.5 --- src/tsim/noise/channels.py | 25 +++++++++++++------------ test/unit/noise/test_channels.py | 11 ++++++----- 2 files changed, 19 insertions(+), 17 deletions(-) diff --git a/src/tsim/noise/channels.py b/src/tsim/noise/channels.py index 0678c0ab..39007004 100644 --- a/src/tsim/noise/channels.py +++ b/src/tsim/noise/channels.py @@ -464,10 +464,12 @@ 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 @@ -534,16 +536,15 @@ def simplify_channels( def unique_columns(matrix: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """Return ``np.unique(matrix, axis=1, return_inverse=True)`` for a binary matrix. + """Return ``np.unique(matrix, axis=1, return_inverse=True)``, with a 1-D inverse. - ``np.unique`` sorts every column, comparing whole columns, which takes about a - minute for the 2,821 x 184,039 error transform of a d=5 end-to-end - cultivation circuit. Here a dict over the bit-packed columns finds the + ``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. + 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 @@ -552,15 +553,15 @@ def unique_columns(matrix: np.ndarray) -> tuple[np.ndarray, np.ndarray]: """ packed = np.ascontiguousarray(np.packbits(matrix, axis=0).T) ids: dict[bytes, int] = {} - seen = np.array( + column_ids = np.array( [ids.setdefault(column.tobytes(), len(ids)) for column in packed], dtype=np.int64, ) - _, first = np.unique(seen, return_index=True) + _, 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[seen] + return matrix[:, first[order]], rank[column_ids] class ChannelSampler: @@ -578,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: diff --git a/test/unit/noise/test_channels.py b/test/unit/noise/test_channels.py index 4245a89f..e3b161de 100644 --- a/test/unit/noise/test_channels.py +++ b/test/unit/noise/test_channels.py @@ -752,7 +752,7 @@ 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=4): +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() @@ -775,17 +775,18 @@ def _absorb_by_comparing_every_pair(channels, max_bits=4): 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): + def test_random_channels(self, seed, max_bits): rng = np.random.default_rng(seed) channels = [] for _ in range(60): - size = int(rng.integers(1, 6)) + 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) - expected = _absorb_by_comparing_every_pair(channels) + 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) From bd743c7ad39c91bd82084defe27eba191ca20d41 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Thu, 15 Oct 2026 18:30:29 -0400 Subject: [PATCH 10/21] Sample whole batches on worker threads with trajectory sharing Parallelism used to live inside each numba kernel, with serial Python, numpy and JAX glue between kernels, so sampling kept only 3-4 of the 14 cores of an M4 Pro busy. With trajectory_sharing=True the sampler now runs whole tasks on numba.get_num_threads() Python threads, and the kernels in them run serially: - utils/kernels.py: every numba function releases the GIL; run_tasks runs tasks on a thread pool and makes kernel() pick the serial compilation on its worker threads. Numba's cache ignores nogil, so jit tags the cached name with _nogil, as kernel() already did with _parallel. - ChannelSampler.block_seeds draws the seeds sample() drew itself, so the main thread draws them in order and hands them to the workers. - _sample_batches runs its batches on the threads. With postselection the shots are screened in chunks of at most 2^14 (noise, direct detectors, survivors), two chunks per thread at a time, until every thread has a full batch of survivors; the batches of survivors then run on the threads too. Seeds and keys are drawn in order and batches still start every batch_size survivors, so the samples do not depend on the thread count, and they equal the previous ones bit for bit when batch_size is at most 2^14 or a multiple of it (checked on cultivation d3, distillation, coherent d3r3, qv_q10 and the measurement sampler at 1, 3 and 14 threads). The JAX executors still sample one batch at a time. Against the previous commit, all threads, shots/s: workload (batch, shots per call) before after cultivation_d3 (65,536, 1M) 9.8M 19.6M coherent_d3_r3 (65,536, 1M) 6.4M 12.5M distillation (65,536, 1M) 19.9M 25.8M cultivation_d5 (65,536, 1M) 8.3M 10.7M coherent_d5_r5 (65,536, 1M) 393k 734k symft_msc_d7_postselected (1M) 2.71M 3.33M qv_q10 (16,384, 262k) 766k 1.60M Calls of a few batches gain less (1.0-1.6x at 262k shots), since one call has few survivor batches and their JAX and numpy glue holds the GIL. One thread is unchanged (within 7%). Co-Authored-By: Claude Opus 5.5 --- src/tsim/circuit.py | 8 +- src/tsim/noise/channels.py | 19 ++- src/tsim/sampler.py | 157 +++++++++++++------- src/tsim/utils/kernels.py | 71 +++++++-- test/integration/test_trajectory_sharing.py | 52 +++++++ test/unit/noise/test_channels.py | 25 ++++ test/unit/test_postselection.py | 12 +- test/unit/utils/test_kernels.py | 46 +++++- 8 files changed, 313 insertions(+), 77 deletions(-) diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index d1022f40..836d11a0 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -875,7 +875,9 @@ def compile_sampler( qubits, 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, and sample up to + ``numba.get_num_threads()`` batches at once, each with its own + memory (experimental). 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 @@ -911,7 +913,9 @@ def compile_detector_sampler( qubits, 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, and sample up to + ``numba.get_num_threads()`` batches at once, each with its own + memory (experimental). 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/noise/channels.py b/src/tsim/noise/channels.py index 39007004..944c0dc3 100644 --- a/src/tsim/noise/channels.py +++ b/src/tsim/noise/channels.py @@ -637,7 +637,14 @@ def __init__( ) self._tables = _fire_tables(self.channels, self.signature_matrix) - def sample(self, num_samples: int = 1) -> np.ndarray: + 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) + + 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 @@ -647,6 +654,8 @@ def sample(self, num_samples: int = 1) -> np.ndarray: Args: num_samples: Number of samples to draw. + seeds: The blocks' seeds from ``block_seeds(num_samples)``; drawn + here if ``None``. Returns: NumPy array of shape (num_samples, num_f) with uint8 values indicating @@ -655,8 +664,14 @@ def sample(self, num_samples: int = 1) -> np.ndarray: """ num_f = self.signature_matrix.shape[1] result = np.zeros((num_samples, num_f), dtype=np.uint8) + if seeds is None: + seeds = self.block_seeds(num_samples) num_blocks = -(-num_samples // _BLOCK_SHOTS) - seeds = self._rng.integers(0, 2**64, size=num_blocks, dtype=np.uint64) + if len(seeds) != num_blocks: + raise ValueError( + f"{num_samples} samples need {num_blocks} block seeds, " + f"got {len(seeds)}" + ) work = num_samples * len(self._tables.p_fire) _sample_blocks(work, result, seeds, *self._tables) return result diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 701524a4..686e4b29 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -26,7 +26,7 @@ sharing_layout, ) from tsim.utils.cuda_helpers import copy_d2h -from tsim.utils.kernels import kernel +from tsim.utils.kernels import kernel, run_tasks if TYPE_CHECKING: from jax import Array as PRNGKey @@ -36,6 +36,12 @@ # NOTE: beyond 2^16 lanes, statevector batches cost memory without being faster. STATEVECTOR_MAX_BATCH_SIZE = 2**16 +# 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, @@ -248,7 +254,9 @@ def __init__( 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, and sample up to + ``numba.get_num_threads()`` batches at once, each with its own + memory (experimental). Raises: ValueError: If ``trajectory_sharing`` is combined with a @@ -458,6 +466,14 @@ def _memory_budget(self) -> int: available = psutil.virtual_memory().available return int(available * 0.5) # conservative + def _threads(self) -> int: + """Return how many batches run at once. + + With trajectory sharing every numba thread samples its own batches; the + JAX executors sample one batch at a time. + """ + return numba.get_num_threads() if self._trajectory_sharing else 1 + def _estimate_batch_size(self) -> int: """Estimate the largest batch size that fits in available device memory. @@ -536,26 +552,27 @@ def _sample_batches( # within existing batches (keeps shapes uniform for JIT). batch_size += 1 - batches: list[jax.Array] = [] - reference: np.ndarray | None = None - - for _ in range(num_batches): - f_params_np = self._channel_sampler.sample(batch_size) + sharing = self._sharing_arguments(None) + tasks = [] + for index in range(num_batches): + seeds = self._channel_sampler.block_seeds(batch_size) + self._key, subkey = jax.random.split(self._key) + tasks.append((index, seeds, subkey)) - if compute_reference and reference is None: + def sample_batch(task: tuple[int, np.ndarray, PRNGKey]) -> jax.Array: + index, seeds, key = task + f_params_np = self._channel_sampler.sample(batch_size, seeds) + if compute_reference and index == 0: 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) + return sample_program( + self._program, jnp.asarray(f_params_np), key, **sharing ) - if compute_reference and reference is None: - reference = np.asarray(samples[0]) - samples = samples[1:] - - batches.append(samples) + batches = run_tasks(sample_batch, tasks, self._threads()) + reference: np.ndarray | None = None + if compute_reference: + reference = np.asarray(batches[0][0]) + batches[0] = batches[0][1:] # Concatenate on device, then a single d2h. The prior # np.concatenate(batches) triggered per-batch __array__ (one d2h each) @@ -588,6 +605,8 @@ def _sample_batches_with_postselection( ``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. + Chunks of shots, and then the batches of their survivors, run on + ``_threads()`` threads at once. """ if shots < 0: raise ValueError(f"shots must be non-negative, got {shots}") @@ -624,7 +643,6 @@ def _sample_batches_with_postselection( result = np.zeros((shots, num_outputs), dtype=np.bool_) was_discarded = np.zeros(shots, dtype=np.bool_) - shot_idx = 0 columns = self._direct_detector_columns postselect = postselection_mask[columns] @@ -638,49 +656,82 @@ def _sample_batches_with_postselection( postselection_mask, reference if xor_detector_ref else None ) - 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) + chunk = min(batch_size, SCREEN_SHOTS) - def dispatch(f_batch: np.ndarray, shot_batch: np.ndarray) -> None: - """Sample the shots ``shot_batch`` in one full batch, padded with row 0.""" - count = len(shot_batch) - if count < batch_size: - padding = np.repeat(f_batch[:1], batch_size - count, axis=0) - f_batch = np.concatenate([f_batch, padding]) - self._key, subkey = jax.random.split(self._key) - samples = sample_program( - self._program, jnp.asarray(f_batch), subkey, **sharing - ) - result[shot_batch] = np.asarray(samples)[:count] + def screen(task: tuple[int, np.ndarray]) -> tuple[np.ndarray, np.ndarray]: + """Sample the chunk of shots at ``start``; return its survivors. - while shot_idx < shots: - chunk = min(batch_size, shots - shot_idx) - f_params_np = self._channel_sampler.sample(chunk) + Writes the chunk's direct detectors and which of its shots a direct + detector discards; returns the survivors' f-params and shot indices. + """ + start, seeds = task + stop = min(start + chunk, shots) + f_params_np = self._channel_sampler.sample(stop - start, seeds) discarded = _write_direct_detectors( - chunk * len(columns), + (stop - start) * len(columns), f_params_np, self._direct_detector_f, self._direct_detector_flips, columns, postselect, reference_bits, - result[shot_idx : shot_idx + chunk], + result[start:stop], ) - was_discarded[shot_idx : shot_idx + chunk] = discarded - + was_discarded[start:stop] = discarded survivors = np.flatnonzero(~discarded) - pending_f = np.concatenate([pending_f, f_params_np[survivors]]) - pending_shot = np.concatenate([pending_shot, shot_idx + survivors]) + return f_params_np[survivors], start + survivors - shot_idx += chunk - if len(pending_shot) >= batch_size: - dispatch(pending_f[:batch_size], pending_shot[:batch_size]) - pending_f = pending_f[batch_size:] - pending_shot = pending_shot[batch_size:] + def dispatch(task: tuple[np.ndarray, np.ndarray, PRNGKey]) -> None: + """Sample the shots ``shot_batch`` in one full batch, padded with row 0.""" + f_batch, shot_batch, key = task + count = len(shot_batch) + if count < batch_size: + padding = np.repeat(f_batch[:1], batch_size - count, axis=0) + f_batch = np.concatenate([f_batch, padding]) + samples = sample_program( + self._program, jnp.asarray(f_batch), key, **sharing + ) + result[shot_batch] = np.asarray(samples)[:count] - if len(pending_shot): - dispatch(pending_f, pending_shot) + threads = self._threads() + 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. Batches still start every batch_size survivors and seeds and + # keys 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 + if screened_chunks < len(starts): + ready -= ready % (threads * batch_size) + all_f, all_shot = np.concatenate(pending_f), np.concatenate(pending_shot) + batches = [] + for start in range(0, ready, batch_size): + self._key, subkey = jax.random.split(self._key) + stop = start + batch_size + batches.append((all_f[start:stop], all_shot[start:stop], subkey)) + run_tasks(dispatch, batches, threads) + pending_f, pending_shot = [all_f[ready:]], [all_shot[ready:]] + pending -= ready if xor_detector_ref and reference is not None: det_ref = reference[: self._num_detectors] @@ -805,7 +856,9 @@ 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, and sample up to + ``numba.get_num_threads()`` batches at once, each with its own + memory (experimental). 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 @@ -864,7 +917,9 @@ 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, and sample up to + ``numba.get_num_threads()`` batches at once, each with its own + memory (experimental). 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 diff --git a/src/tsim/utils/kernels.py b/src/tsim/utils/kernels.py index d84d1e8a..25c31e86 100644 --- a/src/tsim/utils/kernels.py +++ b/src/tsim/utils/kernels.py @@ -2,7 +2,9 @@ import threading import types -from collections.abc import Callable +from collections.abc import Callable, Iterable +from concurrent.futures import ThreadPoolExecutor +from typing import TypeVar import numba @@ -12,46 +14,85 @@ # 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 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") try: - return numba.njit(cache=True, parallel=parallel)(function) + return numba.njit(cache=True, parallel=parallel, nogil=True)(function) except RuntimeError: - return numba.njit(parallel=parallel)(function) + return numba.njit(parallel=parallel, 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. In the serial compilation every - ``numba.prange`` is a plain range, so both compute the same result. + ``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. - renamed = types.FunctionType( - function.__code__, - function.__globals__, - function.__name__ + "_parallel", - function.__defaults__, - function.__closure__, - ) - renamed.__qualname__ = function.__qualname__ + "_parallel" - parallel = jit(renamed, parallel=True) + parallel = jit(_renamed(function, "_parallel"), parallel=True) def run(work: int, *args): - if work < PARALLEL_WORK: + 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 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 the + tasks run in the calling thread. + """ + tasks = list(tasks) + threads = min(threads, len(tasks)) + if threads <= 1: + return [function(task) for task in tasks] + + def run(task: _Task) -> _Result: + _TASK_THREAD.active = True + return function(task) + + with ThreadPoolExecutor(threads) as pool: + return list(pool.map(run, tasks)) diff --git a/test/integration/test_trajectory_sharing.py b/test/integration/test_trajectory_sharing.py index aac80e76..b1357947 100644 --- a/test/integration/test_trajectory_sharing.py +++ b/test/integration/test_trajectory_sharing.py @@ -6,6 +6,8 @@ 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 @@ -144,3 +146,53 @@ def test_trajectory_sharing_needs_the_statevector_strategy(): Circuit(CIRCUITS["t_ghz"]).compile_detector_sampler( strategy="cat5", trajectory_sharing=True ) + + +@pytest.mark.parametrize("screen_shots", [None, 64]) +@pytest.mark.parametrize("postselect", [False, True]) +@pytest.mark.parametrize("reference", [False, True]) +def test_shared_samples_do_not_depend_on_the_thread_count( + postselect, reference, screen_shots, 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=True + ) + sampler._threads = lambda: 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: threads + return sampler.sample(2_001, batch_size=250) + + np.testing.assert_array_equal(sample(4), sample(1)) diff --git a/test/unit/noise/test_channels.py b/test/unit/noise/test_channels.py index e3b161de..79596dba 100644 --- a/test/unit/noise/test_channels.py +++ b/test/unit/noise/test_channels.py @@ -1167,3 +1167,28 @@ 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)) diff --git a/test/unit/test_postselection.py b/test/unit/test_postselection.py index 4360122b..4cbaafe7 100644 --- a/test/unit/test_postselection.py +++ b/test/unit/test_postselection.py @@ -188,8 +188,8 @@ def test_survivors_keep_their_own_direct_outputs(shots, batch_size, reference): drawn: list[np.ndarray] = [] original = sampler._channel_sampler.sample - def capture(n: int) -> np.ndarray: - batch = original(n) + def capture(n: int, seeds: np.ndarray | None = None) -> np.ndarray: + batch = original(n, seeds) drawn.append(batch.copy()) return batch @@ -218,8 +218,8 @@ def test_postselection_direct_cols_always_equal_numpy(): drawn: list[np.ndarray] = [] original = sampler._channel_sampler.sample - def capture(n: int) -> np.ndarray: - batch = original(n) + def capture(n: int, seeds: np.ndarray | None = None) -> np.ndarray: + batch = original(n, seeds) drawn.append(batch.copy()) return batch @@ -502,9 +502,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() diff --git a/test/unit/utils/test_kernels.py b/test/unit/utils/test_kernels.py index d24c9368..3f573ac0 100644 --- a/test/unit/utils/test_kernels.py +++ b/test/unit/utils/test_kernels.py @@ -1,11 +1,13 @@ import subprocess import sys import textwrap +import threading import numba import numpy as np +import pytest -from tsim.utils.kernels import PARALLEL_WORK, kernel +from tsim.utils.kernels import PARALLEL_WORK, kernel, run_tasks @kernel @@ -14,6 +16,12 @@ def _square_roots(out): 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) @@ -53,3 +61,39 @@ def work(): 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(): + def task(n): + if n == 2: + raise ValueError("task 2") + return n + + with pytest.raises(ValueError, match="task 2"): + run_tasks(task, range(4), threads=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 From 87eb174cdafc63d8e78e44763dc2a2394eec1763 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Fri, 16 Oct 2026 09:51:57 -0400 Subject: [PATCH 11/21] Put the sharing path's outputs together on the host With trajectory sharing every component runs in the numba executor on the host, but the sampler still moved each batch into JAX, gathered the direct outputs, concatenated and reordered the outputs with JAX ops, padded the last postselection survivor batch to batch_size rows (only JAX needs fixed shapes) and joined the batches on a JAX device. Now: - sample_program takes NumPy f-params and returns NumPy samples, put together by one numba kernel, _gather_outputs, that writes the direct outputs and the component outputs into their original columns. NumPy's column gathers were 4x slower than JAX's and ran under the GIL. - The sharing path passes NumPy f-params, samples the last survivor batch as it is, and joins its batches with np.concatenate. The JAX paths are unchanged. Samples without postselection are bit-identical to before. With postselection only the last survivor batch of a call differs: its padding rows were copies of row 0 that took part in the executor's grouping and random draws. A review compared the statistics with one lane per shot on 18 configurations, including calls whose every batch is partial (acceptance within |z| <= 1.6, matching outcome counts). Samples still do not depend on the thread count. Against the previous commit, all 14 threads, shots/s: workload (shots per call) before after symft_msc_d7_postselected (262k) 1.41M 2.07M symft_msc_d7_postselected (1M) 2.67M 3.64M distillation (262k) 20.1M 27.0M cultivation_d5 (262k) 9.1M 11.4M coherent_d3_r3 (1M) 12.9M 14.7M cultivation_d3 (1M) 21.6M 25.1M qv_q10 (262k) 1.51M 1.64M One thread: distillation +11%, cultivation_d5 +18%. Tests check that the host and JAX glue give equal outputs on every test circuit and on one with flipped direct detectors and observables, at 1, 300 and 70,000 rows, and that survivor batches hold exactly the surviving shots, without padding. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 108 +++++++++++++++----- test/integration/test_trajectory_sharing.py | 79 ++++++++++++++ 2 files changed, 161 insertions(+), 26 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 686e4b29..84623041 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -136,13 +136,16 @@ def sample_component( def sample_program( program: CompiledProgram, - f_params: jax.Array, + f_params: jax.Array | np.ndarray, key: PRNGKey, sharing: tuple[SharingLayout | None, ...] | None = None, sharing_stats: deque[dict] | None = None, -) -> jax.Array: +) -> jax.Array | np.ndarray: """Sample all outputs from a compiled program. + NumPy ``f_params`` give NumPy samples, put together by a numba kernel (the + shared executor works on the host); JAX ``f_params`` give JAX samples. + Args: program: The compiled program to sample from. f_params: Error parameters, shape (batch_size, num_f_params). @@ -159,11 +162,14 @@ def sample_program( """ results: list[jax.Array | np.ndarray] = [] + on_host = isinstance(f_params, np.ndarray) if program.num_outputs == 0: batch_size = f_params.shape[0] + if on_host: + return np.zeros((batch_size, 0), dtype=np.bool_) return jnp.zeros((batch_size, 0), dtype=jnp.bool_) - if len(program.direct_f_indices) > 0: + if len(program.direct_f_indices) > 0 and not on_host: direct_bits = ( f_params[:, program.direct_f_indices].astype(jnp.bool_) ^ program.direct_flips @@ -180,7 +186,7 @@ def sample_program( sharing_stats.append(stats) else: samples, key, max_norm_deviation = sample_component( - component, f_params, key + component, jnp.asarray(f_params), key ) if np.isclose(max_norm_deviation, 1): raise ValueError( @@ -198,6 +204,25 @@ def sample_program( ) results.append(samples) + if on_host: + assert isinstance(f_params, np.ndarray) + rows = f_params.shape[0] + component_bits = ( + np.concatenate(results, axis=1) + if results + else np.zeros((rows, 0), dtype=np.bool_) + ) + out = np.empty((rows, program.num_outputs), dtype=np.bool_) + _gather_outputs( + out.size, + f_params, + np.asarray(program.direct_f_indices), + np.asarray(program.direct_flips), + component_bits, + np.asarray(program.output_order), + out, + ) + return out combined = jnp.concatenate(results, axis=1) if program.output_reindex is not None: combined = combined[:, program.output_reindex] @@ -221,6 +246,21 @@ def _write_direct_detectors(f, f_indices, flips, columns, postselect, reference, 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: return sum( v.nbytes for v in jax.tree_util.tree_leaves(tree) if isinstance(v, jax.Array) @@ -287,6 +327,10 @@ def __init__( self._num_detectors = prepared.num_detectors prog = self._program + # NOTE: _gather_outputs writes column output_order[k] without bounds checks. + assert np.array_equal( + np.sort(np.asarray(prog.output_order)), np.arange(prog.num_outputs) + ) self._direct_f_indices = np.asarray(prog.direct_f_indices) self._direct_flips = np.asarray(prog.direct_flips, dtype=np.bool_) self._direct_reindex = ( @@ -559,29 +603,35 @@ def _sample_batches( self._key, subkey = jax.random.split(self._key) tasks.append((index, seeds, subkey)) - def sample_batch(task: tuple[int, np.ndarray, PRNGKey]) -> jax.Array: + def sample_batch( + task: tuple[int, np.ndarray, PRNGKey], + ) -> jax.Array | np.ndarray: index, seeds, key = task f_params_np = self._channel_sampler.sample(batch_size, seeds) if compute_reference and index == 0: f_params_np[0] = 0 - return sample_program( - self._program, jnp.asarray(f_params_np), key, **sharing - ) + f_params = f_params_np if sharing else jnp.asarray(f_params_np) + return sample_program(self._program, f_params, key, **sharing) batches = run_tasks(sample_batch, tasks, self._threads()) reference: np.ndarray | None = None if compute_reference: reference = np.asarray(batches[0][0]) batches[0] = batches[0][1:] - - # 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 sharing: + result = np.concatenate(batches)[:shots] + else: + # 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: assert reference is not None @@ -603,8 +653,9 @@ def _sample_batches_with_postselection( 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. + call. For JAX the final partial batch is padded to keep the batch size + fixed (avoiding recompilation) and the padding rows are discarded; the + shared executor samples it as it is. Chunks of shots, and then the batches of their survivors, run on ``_threads()`` threads at once. """ @@ -682,15 +733,20 @@ def screen(task: tuple[int, np.ndarray]) -> tuple[np.ndarray, np.ndarray]: return f_params_np[survivors], start + survivors def dispatch(task: tuple[np.ndarray, np.ndarray, PRNGKey]) -> None: - """Sample the shots ``shot_batch`` in one full batch, padded with row 0.""" + """Sample the shots ``shot_batch`` in one batch. + + JAX batches are padded with row 0 to ``batch_size`` rows; the shared + executor takes any number of rows. + """ f_batch, shot_batch, key = task count = len(shot_batch) - if count < batch_size: - padding = np.repeat(f_batch[:1], batch_size - count, axis=0) - f_batch = np.concatenate([f_batch, padding]) - samples = sample_program( - self._program, jnp.asarray(f_batch), key, **sharing - ) + if sharing: + samples = sample_program(self._program, f_batch, key, **sharing) + else: + if count < batch_size: + padding = np.repeat(f_batch[:1], batch_size - count, axis=0) + f_batch = np.concatenate([f_batch, padding]) + samples = sample_program(self._program, jnp.asarray(f_batch), key) result[shot_batch] = np.asarray(samples)[:count] threads = self._threads() diff --git a/test/integration/test_trajectory_sharing.py b/test/integration/test_trajectory_sharing.py index b1357947..897bfed0 100644 --- a/test/integration/test_trajectory_sharing.py +++ b/test/integration/test_trajectory_sharing.py @@ -3,6 +3,8 @@ 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 @@ -34,6 +36,36 @@ """ +# 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 _accepted(sampler, shots, mask, *, reference=False): detectors, observables = sampler.sample( shots, @@ -196,3 +228,50 @@ def sample(threads): return sampler.sample(2_001, batch_size=250) np.testing.assert_array_equal(sample(4), sample(1)) + + +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: 1 + rows = [] + original = sampler_module.sample_program + + def spy(program, f_params, key, **kwargs): + rows.append((type(f_params), f_params.shape[0])) + return original(program, f_params, key, **kwargs) + + monkeypatch.setattr(sampler_module, "sample_program", spy) + mask = np.ones(circuit.num_detectors, dtype=bool) + detectors = sampler.sample(5_003, batch_size=300, postselection_mask=mask) + + counts = [count for _, count in rows] + assert all(kind is np.ndarray for kind, _ in rows) + 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("rows", [1, 300, 70_000]) +@pytest.mark.parametrize("name", [*sorted(CIRCUITS), "direct_flips"]) +def test_host_outputs_equal_the_jax_outputs(name, rows): + text = DIRECT_FLIPS if name == "direct_flips" else CIRCUITS[name] + sampler = Circuit(text).compile_detector_sampler( + strategy="active_statevector", seed=1, trajectory_sharing=True + ) + if name == "direct_flips": + assert sampler._direct_flips.any() + assert sampler._program.output_reindex is not None + f_params = sampler._channel_sampler.sample(rows) + key = jax.random.key(2) + sharing = sampler._sharing_arguments(None) + program = sampler._program + on_host = sampler_module.sample_program(program, f_params, key, **sharing) + on_device = sampler_module.sample_program( + program, jnp.asarray(f_params), key, **sharing + ) + assert isinstance(on_host, np.ndarray) + np.testing.assert_array_equal(on_host, np.asarray(on_device)) From 412d3d0aa52f3c291eb12331e382314366f07073 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Fri, 16 Oct 2026 09:52:26 -0400 Subject: [PATCH 12/21] Thread the one-lane JAX path and trim the sharing executor's overhead Calls of 2.6M shots on 10 threads still fell short of an in-process prototype that sampled one batch per thread: without trajectory sharing the JAX batches ran one at a time, and with it the main thread zeroed large outputs, every round started fresh threads, and Python made one GIL round trip per tape action. Now: - On a CPU backend the JAX statevector executor (one lane per shot) also samples a batch per numba thread, as many as fit in the memory budget, and the default batch size leaves room for them. The caller's jax.enable_x64, which is thread-local, reaches the worker threads, so precision and samples do not depend on the thread count. Stabilizer-rank programs still sample one batch at a time. - The noise and direct-detector kernels clear their own rows, so outputs come from np.empty: np.zeros of a large array is a serial memset when malloc reuses freed memory. - run_tasks keeps one thread pool (rebuilt in a forked child; no fork hook where os has none), waits for every task before raising, and runs nested calls inline. - The shared executor applies each stretch of consecutive rotations in one kernel call (qv_q10: 692 rotations in 11 calls), with the rotation coefficients computed once per layout. Samples are bit-identical to the previous commit with and without sharing, at 1 and 14 threads (sample hashes of cultivation d3, distillation, coherent d3r3, qv_q10 and the measurement sampler). shots/s, 10 threads with calls of 2.6M shots, 1 thread with calls of 262k; the prototype and the 10 single-thread processes ran the code before these changes: sharing on 1 thread 10 threads prototype 10 processes distillation 5.89M 41.2M 41.9M 41.8M coherent_d5_r5 147k 978k 979k 973k cultivation_d5 3.73M 20.3M 17.3M 22.8M qv_q10 737k 3.72M 1.82M 5.0M surface_d7_r7 12.6M 60.7M 73.3M 71.8M symft_msc_d7, postsel. 779k 3.18M - 2.85M cultivation_d3 7.63M 31.5M 26.2M 44.9M coherent_d3_r3 5.80M 19.6M - 37.4M sharing off distillation 4.25M 36.5M 26.2M 31.4M cultivation_d5 512k 4.04M 3.14M 3.0M cultivation_d3 1.12M 5.70M 5.34M 5.9M coherent_d3_r3 457k 2.30M 2.24M 2.28M surface_d7_r7 is bound by the OS freeing each call's 883 MB result (16 ms of a 45 ms call). On cultivation_d3 and coherent_d3_r3 a quarter to two fifths of each batch is executor bookkeeping in Python that holds the GIL. Co-Authored-By: Claude Opus 5.5 --- src/tsim/circuit.py | 16 +-- src/tsim/noise/channels.py | 5 +- src/tsim/sampler.py | 92 ++++++++++------ src/tsim/statevector/shared.py | 112 +++++++++++++------- src/tsim/utils/kernels.py | 43 +++++--- test/integration/test_trajectory_sharing.py | 30 ++++-- test/unit/noise/test_channels.py | 13 +++ test/unit/test_sampler.py | 51 +++++++++ test/unit/utils/test_kernels.py | 18 +++- 9 files changed, 278 insertions(+), 102 deletions(-) diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index 836d11a0..2583bede 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -873,11 +873,11 @@ 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"). On the CPU, "active_statevector" samples 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, and sample up to - ``numba.get_num_threads()`` batches at once, each with its own - memory (experimental). + with equal histories on one lane (experimental). 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 @@ -911,11 +911,11 @@ 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"). On the CPU, "active_statevector" samples 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, and sample up to - ``numba.get_num_threads()`` batches at once, each with its own - memory (experimental). + with equal histories on one lane (experimental). 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/noise/channels.py b/src/tsim/noise/channels.py index 944c0dc3..5099e54e 100644 --- a/src/tsim/noise/channels.py +++ b/src/tsim/noise/channels.py @@ -663,7 +663,7 @@ def sample( """ num_f = self.signature_matrix.shape[1] - result = np.zeros((num_samples, num_f), dtype=np.uint8) + 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) @@ -736,7 +736,8 @@ def _fire_tables(channels: list[Channel], signature_matrix: np.ndarray) -> _Fire @jit def _sample_block(out, state, p_fire, outcome_start, cdf, flip_start, flip_cols): - """XOR every channel's outcomes into the rows of ``out`` it fires on.""" + """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] diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 84623041..e177d4cf 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -4,6 +4,7 @@ import warnings from collections import deque +from collections.abc import Callable from math import ceil from typing import TYPE_CHECKING, Literal, overload @@ -231,13 +232,14 @@ def sample_program( @kernel def _write_direct_detectors(f, f_indices, flips, columns, postselect, reference, out): - """Write direct detector ``j`` of every row into ``out[row, columns[j]]``. + """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 that differs from its ``reference`` bit. """ 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 @@ -261,6 +263,20 @@ def _gather_outputs(f, f_indices, flips, component_bits, order, out): out[row, order[num_direct + j]] = component_bits[row, j] +def _with_x64_of_caller(function: Callable) -> Callable: + """Return ``function`` run with the calling thread's JAX 64-bit mode. + + ``jax.enable_x64`` only reaches the thread that enters it, not the workers. + """ + x64 = jax.config.read("jax_enable_x64") + + def run(task): + with jax.enable_x64(x64): + return function(task) + + return run + + def _nbytes(tree: object) -> int: return sum( v.nbytes for v in jax.tree_util.tree_leaves(tree) if isinstance(v, jax.Array) @@ -289,14 +305,11 @@ 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"). - 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. + "bss", "cutting"). On the CPU, "active_statevector" samples 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, and sample up to - ``numba.get_num_threads()`` batches at once, each with its own - memory (experimental). + with equal histories on one lane (experimental). Raises: ValueError: If ``trajectory_sharing`` is combined with a @@ -362,6 +375,7 @@ def __init__( statevector = [ c for c in self._program.components if isinstance(c, StatevectorComponent) ] + self._has_statevector = bool(statevector) if statevector: peak = self._peak_bytes_per_sample() if peak > self._memory_budget(): @@ -510,21 +524,35 @@ def _memory_budget(self) -> int: available = psutil.virtual_memory().available return int(available * 0.5) # conservative - def _threads(self) -> int: - """Return how many batches run at once. + def _threads(self, batch_size: int) -> int: + """Return how many batches of ``batch_size`` shots run at once. - With trajectory sharing every numba thread samples its own batches; the - JAX executors sample one batch at a time. + Statevector components sample one batch per numba thread on the CPU: in + the shared executor, or in JAX on a CPU backend as long as the batches + fit in the memory budget. The other executors sample one batch at a time. """ - return numba.get_num_threads() if self._trajectory_sharing else 1 + if not self._has_statevector: + return 1 + threads = numba.get_num_threads() + if self._trajectory_sharing: + return threads + if jax.default_backend() != "cpu": + return 1 + fits = self._memory_budget() // (self._peak_bytes_per_sample() * batch_size) + return max(1, min(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. + JAX allocates a lane per shot in every batch that runs at once. The shared + executor allocates lanes as its trajectories split, so its estimate leaves + room for a lane per shot in one batch only. Programs with statevector + components use at most ``STATEVECTOR_MAX_BATCH_SIZE`` lanes per batch. """ - batch_size = max(1, self._memory_budget() // self._peak_bytes_per_sample()) + peak = self._peak_bytes_per_sample() + if not self._trajectory_sharing: + peak *= self._threads(1) + batch_size = max(1, self._memory_budget() // peak) if any(isinstance(c, StatevectorComponent) for c in self._program.components): batch_size = min(batch_size, STATEVECTOR_MAX_BATCH_SIZE) return batch_size @@ -613,7 +641,9 @@ def sample_batch( f_params = f_params_np if sharing else jnp.asarray(f_params_np) return sample_program(self._program, f_params, key, **sharing) - batches = run_tasks(sample_batch, tasks, self._threads()) + batches = run_tasks( + _with_x64_of_caller(sample_batch), tasks, self._threads(batch_size) + ) reference: np.ndarray | None = None if compute_reference: reference = np.asarray(batches[0][0]) @@ -692,8 +722,8 @@ def _sample_batches_with_postselection( 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_) + result = np.empty((shots, num_outputs), dtype=np.bool_) + was_discarded = np.empty(shots, dtype=np.bool_) columns = self._direct_detector_columns postselect = postselection_mask[columns] @@ -749,7 +779,7 @@ def dispatch(task: tuple[np.ndarray, np.ndarray, PRNGKey]) -> None: samples = sample_program(self._program, jnp.asarray(f_batch), key) result[shot_batch] = np.asarray(samples)[:count] - threads = self._threads() + 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)] @@ -785,7 +815,7 @@ def dispatch(task: tuple[np.ndarray, np.ndarray, PRNGKey]) -> None: self._key, subkey = jax.random.split(self._key) stop = start + batch_size batches.append((all_f[start:stop], all_shot[start:stop], subkey)) - run_tasks(dispatch, batches, threads) + run_tasks(_with_x64_of_caller(dispatch), batches, threads) pending_f, pending_shot = [all_f[ready:]], [all_shot[ready:]] pending -= ready @@ -912,9 +942,7 @@ 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, and sample up to - ``numba.get_num_threads()`` batches at once, each with its own - memory (experimental). + with equal histories on one lane (experimental). 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 @@ -938,8 +966,10 @@ def sample(self, shots: int, *, batch_size: int | None = None) -> np.ndarray: 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. + utilize VRAM for maximum performance. With ``strategy="active_statevector"`` + on the CPU, up to ``numba.get_num_threads()`` batches run at once. + NOTE: Changing the batch size will affect reproducibility even with + a fixed seed. Returns: A numpy array containing the measurement samples. @@ -973,9 +1003,7 @@ 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, and sample up to - ``numba.get_num_threads()`` batches at once, each with its own - memory (experimental). + with equal histories on one lane (experimental). 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 @@ -1045,8 +1073,10 @@ def sample( 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. + utilize VRAM for maximum performance. With ``strategy="active_statevector"`` + on the CPU, up to ``numba.get_num_threads()`` batches run at once. + 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. diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index fba02304..0b2b9138 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -31,6 +31,7 @@ from __future__ import annotations +import itertools from typing import NamedTuple import jax @@ -91,6 +92,11 @@ class SharingLayout(NamedTuple): 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,)``. + 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,)``. """ @@ -103,6 +109,9 @@ class SharingLayout(NamedTuple): 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: @@ -180,6 +189,14 @@ def sharing_layout( shot_bits = np.zeros(32 * num_words, dtype=bool) shot_bits[branch[dormant & ~lane_dormant & (branch >= 0)]] = True + dtype = np.dtype(component.dtype) + 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=dtype) + cos, sin = tape.cos, tape.sin + return SharingLayout( tape=tape, key_rows=_pack(key_bits, num_words).reshape(-1, num_words), @@ -190,6 +207,9 @@ def sharing_layout( 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).astype(dtype), + coupling=(-1j * sin * i_power).astype(dtype), ) @@ -256,32 +276,37 @@ def _signs(symbols, lanes, row, const): @kernel -def _rotate_diag(amps, lanes, extent, z, symbols, sign_row, sign_const, factor): - """Multiply amplitude i of lane l by ``factor[sign(l) ^ parity(i & z)]``.""" - flip = _signs(symbols, lanes, sign_row, sign_const) - span = 1 << extent - for t in numba.prange(lanes * span): - lane, i = t >> extent, t & (span - 1) - amps[lane, i] *= factor[flip[lane] ^ _parity(i & z)] - - -@kernel -def _rotate_pair( - amps, lanes, extent, x, z, symbols, sign_row, sign_const, cos, coupling -): - """Rotate about ``X^x Z^z``: mix amplitudes i and ``i ^ x`` of every lane.""" - flip = _signs(symbols, lanes, sign_row, sign_const) - span = 1 << extent - low = x & -x - for t in numba.prange(lanes * span): - lane, i = t >> extent, t & (span - 1) - if i & low: - continue - j = i ^ x - c = -coupling if flip[lane] else coupling - a, b = amps[lane, i], amps[lane, j] - amps[lane, i] = cos * a + c * (-b if _parity(j & z) else b) - amps[lane, j] = cos * b + c * (-a if _parity(i & z) else a) +def _rotate(amps, lanes, symbols, start, stop, rotations): + """Apply the rotations ``start … stop-1`` of the tape to every lane. + + ``rotations`` holds the tape's kinds, Pauli masks, extents, signs and cosines + and the layout's factors and couplings. Rotation ``a`` is about ``X^x Z^z`` + with a sign per lane: a diagonal one multiplies amplitude i by + ``factor[a, sign ^ parity(i & z)]``, the others mix amplitudes i and + ``i ^ x`` with ``cos[a]`` and ``±coupling[a]``. + """ + kind, x_mask, z_mask, extent, sign_row, sign_const, cos, factor, coupling = ( + rotations + ) + for a in range(start, stop): + flip = _signs(symbols, lanes, sign_row[a], np.int64(sign_const[a])) + x, z, shift = np.int64(x_mask[a]), np.int64(z_mask[a]), extent[a] + span = 1 << shift + if kind[a] == ROTATE_DIAG: + for t in numba.prange(lanes * span): + lane, i = t >> shift, t & (span - 1) + amps[lane, i] *= factor[a, flip[lane] ^ _parity(i & z)] + else: + low = x & -x + for t in numba.prange(lanes * span): + lane, i = t >> shift, t & (span - 1) + if i & low: + continue + j = i ^ x + c = -coupling[a] if flip[lane] else coupling[a] + p, q = amps[lane, i], amps[lane, j] + amps[lane, i] = cos[a] * p + c * (-q if _parity(j & z) else q) + amps[lane, j] = cos[a] * q + c * (-p if _parity(i & z) else p) @kernel @@ -471,33 +496,38 @@ def group_on_host( def _run( layout: SharingLayout, lanes: _Lanes, rng: np.random.Generator ) -> tuple[float, list[tuple], int]: - """Run the tape on ``lanes``; return the norm deviation, dead lanes and peak.""" + """Run the tape on ``lanes``; return the norm deviation, dead lanes and peak. + + Every stretch of consecutive rotations runs in one kernel call, the other + actions one by one. A stretch clears the amplitudes up to its largest extent at + once; the amplitudes below ``2^region`` are the lanes' own. + """ tape = layout.tape - dtype = lanes.amps.dtype + rotations = ( + *(tape.kind, tape.x_mask, tape.z_mask, tape.extent), + *(tape.sign_row, tape.sign_const, tape.cos, layout.factor, layout.coupling), + ) + rotation = np.isin(tape.kind, (ROTATE_DIAG, ROTATE)) + others = np.flatnonzero(~rotation) + edges = sorted({0, len(rotation), *others, *(others + 1)}) dead: list[tuple] = [] deviation, peak, region = 0.0, lanes.size, 0 - for a, kind in enumerate(tape.kind): + for a, stop in itertools.pairwise(edges): if not lanes.size: break - extent = int(tape.extent[a]) + extent = int(tape.extent[a:stop].max()) span = 1 << extent if extent > region: lanes.amps[: lanes.size, 1 << region : span] = 0 region = extent - x, z = int(tape.x_mask[a]), int(tape.z_mask[a]) - i_power = dtype.type(_I_POWERS[(x & z).bit_count() & 3]) size = lanes.size work = size * span - if kind in (ROTATE_DIAG, ROTATE): - sign = (lanes.symbols, tape.sign_row[a], int(tape.sign_const[a])) - cos, sin = tape.cos[a], tape.sin[a] - if kind == ROTATE_DIAG: - factor = np.array([cos - 1j * sin, cos + 1j * sin], dtype=dtype) - _rotate_diag(work, lanes.amps, size, extent, z, *sign, factor) - else: - coupling = dtype.type(-1j * sin * i_power) - _rotate_pair(work, lanes.amps, size, extent, x, z, *sign, cos, coupling) + if rotation[a]: + _rotate(work, lanes.amps, size, lanes.symbols, a, stop, rotations) continue + kind = tape.kind[a] + x, z = int(tape.x_mask[a]), int(tape.z_mask[a]) + i_power = layout.i_power[a] if kind == MEASURE_DORMANT: if not layout.lane_dormant[a]: continue diff --git a/src/tsim/utils/kernels.py b/src/tsim/utils/kernels.py index 25c31e86..272525a3 100644 --- a/src/tsim/utils/kernels.py +++ b/src/tsim/utils/kernels.py @@ -1,9 +1,11 @@ """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 +from concurrent.futures import ThreadPoolExecutor, wait from typing import TypeVar import numba @@ -75,6 +77,26 @@ def run(work: int, *args): return run +def _mark_task_thread() -> None: + _TASK_THREAD.active = True + + +@functools.cache +def _pool(threads: int) -> ThreadPoolExecutor: + return ThreadPoolExecutor(threads, 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]: @@ -82,17 +104,14 @@ def run_tasks( 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 the - tasks run in the calling thread. + 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) - threads = min(threads, len(tasks)) - if threads <= 1: + if min(threads, len(tasks)) <= 1 or getattr(_TASK_THREAD, "active", False): return [function(task) for task in tasks] - - def run(task: _Task) -> _Result: - _TASK_THREAD.active = True - return function(task) - - with ThreadPoolExecutor(threads) as pool: - return list(pool.map(run, tasks)) + pool = _pool(threads) + futures = [pool.submit(function, task) for task in tasks] + wait(futures) + return [future.result() for future in futures] diff --git a/test/integration/test_trajectory_sharing.py b/test/integration/test_trajectory_sharing.py index 897bfed0..5aeacd0d 100644 --- a/test/integration/test_trajectory_sharing.py +++ b/test/integration/test_trajectory_sharing.py @@ -1,5 +1,6 @@ """trajectory_sharing=True end to end, against one lane per shot.""" +import warnings from test.helpers.gen import gen_stim_circuit from test.helpers.statevector import CIRCUITS, assert_counts_match @@ -180,11 +181,12 @@ def test_trajectory_sharing_needs_the_statevector_strategy(): ) +@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_shared_samples_do_not_depend_on_the_thread_count( - postselect, reference, screen_shots, monkeypatch +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 @@ -193,9 +195,9 @@ def test_shared_samples_do_not_depend_on_the_thread_count( def sample(threads): sampler = circuit.compile_detector_sampler( - strategy="active_statevector", seed=3, trajectory_sharing=True + strategy="active_statevector", seed=3, trajectory_sharing=sharing ) - sampler._threads = lambda: threads + sampler._threads = lambda batch_size: threads return [ sampler.sample( 5_003, @@ -224,7 +226,7 @@ def sample(threads): sampler = circuit.compile_sampler( strategy="active_statevector", seed=4, trajectory_sharing=True ) - sampler._threads = lambda: threads + sampler._threads = lambda batch_size: threads return sampler.sample(2_001, batch_size=250) np.testing.assert_array_equal(sample(4), sample(1)) @@ -235,7 +237,7 @@ def test_shared_postselection_samples_the_survivors_without_padding(monkeypatch) sampler = circuit.compile_detector_sampler( strategy="active_statevector", seed=5, trajectory_sharing=True ) - sampler._threads = lambda: 1 + sampler._threads = lambda batch_size: 1 rows = [] original = sampler_module.sample_program @@ -275,3 +277,19 @@ def test_host_outputs_equal_the_jax_outputs(name, rows): ) assert isinstance(on_host, np.ndarray) np.testing.assert_array_equal(on_host, np.asarray(on_device)) + + +@pytest.mark.parametrize("sharing", [False, True]) +def test_workers_sample_in_the_callers_64_bit_mode(sharing): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + + def sample(threads): + sampler = circuit.compile_detector_sampler( + strategy="active_statevector", seed=6, trajectory_sharing=sharing + ) + sampler._threads = lambda batch_size: threads + return sampler.sample(2_001, batch_size=250, append_observables=True) + + with jax.enable_x64(True), warnings.catch_warnings(): + warnings.simplefilter("error") + np.testing.assert_array_equal(sample(3), sample(1)) diff --git a/test/unit/noise/test_channels.py b/test/unit/noise/test_channels.py index 79596dba..6acf78e1 100644 --- a/test/unit/noise/test_channels.py +++ b/test/unit/noise/test_channels.py @@ -7,6 +7,7 @@ Channel, ChannelSampler, _fire_tables, + _sample_blocks, absorb_subset_channels, channel_num_bits, correlated_error_probs, @@ -1192,3 +1193,15 @@ def test_wrong_number_of_block_seeds_raises(self): ) 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/test_sampler.py b/test/unit/test_sampler.py index 129794e8..525e3081 100644 --- a/test/unit/test_sampler.py +++ b/test/unit/test_sampler.py @@ -525,3 +525,54 @@ def nan_deviation(component, f_params, key): monkeypatch.setattr(sampler_module, "sample_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]), + np.array([False, 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), + ("active_statevector", False, "gpu", 10, 1), + ("active_statevector", False, "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 index 3f573ac0..abc8feb4 100644 --- a/test/unit/utils/test_kernels.py +++ b/test/unit/utils/test_kernels.py @@ -2,6 +2,7 @@ import sys import textwrap import threading +import time import numba import numpy as np @@ -80,14 +81,18 @@ def test_run_tasks_with_one_thread_runs_in_the_calling_thread(): assert run_tasks(lambda n: n, [], 4) == [] -def test_run_tasks_raises_the_error_of_a_task(): +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") - return n + 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(): @@ -97,3 +102,12 @@ def test_the_calling_thread_keeps_parallel_kernels_after_run_tasks(): 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 From a6769b0a168bc4c517777a839595683ca85b101b Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 19 Oct 2026 17:26:04 -0400 Subject: [PATCH 13/21] Fix the lint error and the thread-order race in two postselection tests The sampler's __init__ docstring had lost its `seed` entry (ruff D417). Two postselection tests rebuilt the drawn noise by appending each chunk's draw as it finished, but chunks are screened on several threads, so on CI's few-core runners the draws came back out of shot order. The tests now put the draws back in shot order by the order of their block seeds. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 3 ++ test/unit/test_postselection.py | 57 +++++++++++++++++++++------------ 2 files changed, 39 insertions(+), 21 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index e177d4cf..da481246 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -308,6 +308,9 @@ def __init__( "bss", "cutting"). On the CPU, "active_statevector" samples 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). diff --git a/test/unit/test_postselection.py b/test/unit/test_postselection.py index 4cbaafe7..63851611 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 @@ -70,6 +72,36 @@ def _keep(samples: np.ndarray, mask: np.ndarray) -> np.ndarray: 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 ─────────────────────────────────────── @@ -185,16 +217,8 @@ def test_postselection_discarded_and_surviving_rows(): 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) - drawn: list[np.ndarray] = [] - original = sampler._channel_sampler.sample - - def capture(n: int, seeds: np.ndarray | None = None) -> np.ndarray: - batch = original(n, seeds) - drawn.append(batch.copy()) - return batch - mask = np.array([True, False, False]) - with patch.object(sampler._channel_sampler, "sample", side_effect=capture): + with _record_noise(sampler) as drawn: samples = sampler.sample( shots, batch_size=batch_size, @@ -202,7 +226,7 @@ def capture(n: int, seeds: np.ndarray | None = None) -> np.ndarray: append_observables=True, use_detector_reference_sample=reference, ) - expected = sampler._compute_direct_outputs(np.concatenate(drawn)) + expected = sampler._compute_direct_outputs(drawn()) if reference: expected[:, :3] ^= sampler._compute_reference_sample()[:3] discarded = expected[:, 0] @@ -215,19 +239,10 @@ 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, seeds: np.ndarray | None = None) -> np.ndarray: - batch = original(n, seeds) - 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 = sampler._compute_direct_outputs(drawn()) direct_det = sampler._direct_detector_mask assert np.array_equal( samples[:, : sampler._num_detectors] & direct_det, From 87926d947ae907146755233d4c2a3bbd2821d69e Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 19 Oct 2026 17:28:07 -0400 Subject: [PATCH 14/21] Split default batches so that every thread has one Without a batch size, a call is split into at least one batch per thread of at least MIN_THREAD_BATCH shots, and postselection splits the survivors left after the last chunk the same way (JAX pads them to a power of two). An explicit batch size keeps its samples independent of the thread count. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 71 +++++++++++++-------- test/integration/test_trajectory_sharing.py | 49 ++++++++++++++ 2 files changed, 93 insertions(+), 27 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index da481246..fe821b16 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -37,6 +37,10 @@ # 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, but +# into batches of no fewer shots than this, whose fixed costs would dominate. +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. @@ -485,15 +489,24 @@ 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 at least one batch per thread (``_threads``) of no fewer than + ``MIN_THREAD_BATCH`` shots each, 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())) + per_thread = min(self._threads(1), shots // MIN_THREAD_BATCH) + return ceil(shots / max(num_batches, per_thread)) + 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] @@ -615,17 +628,13 @@ 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) - - 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 + # 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) sharing = self._sharing_arguments(None) tasks = [] @@ -690,7 +699,10 @@ def _sample_batches_with_postselection( fixed (avoiding recompilation) and the padding rows are discarded; the shared executor samples it as it is. Chunks of shots, and then the batches of their survivors, run on - ``_threads()`` threads at once. + ``_threads()`` threads at once. Without a ``batch_size``, batches are as + large as memory allows, and the survivors left after the last chunk are + split into a batch per thread (as in ``_resolve_batch_size``); JAX pads + each to a power of two, so that few batch shapes are compiled. """ if shots < 0: raise ValueError(f"shots must be non-negative, got {shots}") @@ -716,10 +728,9 @@ def _sample_batches_with_postselection( return samples, reference, np.zeros(shots, dtype=np.bool_) return samples, None, np.zeros(shots, 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() reference: np.ndarray | None = None if compute_reference: @@ -765,19 +776,19 @@ def screen(task: tuple[int, np.ndarray]) -> tuple[np.ndarray, np.ndarray]: survivors = np.flatnonzero(~discarded) return f_params_np[survivors], start + survivors - def dispatch(task: tuple[np.ndarray, np.ndarray, PRNGKey]) -> None: + def dispatch(task: tuple[np.ndarray, np.ndarray, PRNGKey, int]) -> None: """Sample the shots ``shot_batch`` in one batch. - JAX batches are padded with row 0 to ``batch_size`` rows; the shared + JAX batches are padded with row 0 to ``rows`` rows; the shared executor takes any number of rows. """ - f_batch, shot_batch, key = task + f_batch, shot_batch, key, rows = task count = len(shot_batch) if sharing: samples = sample_program(self._program, f_batch, key, **sharing) else: - if count < batch_size: - padding = np.repeat(f_batch[:1], batch_size - count, axis=0) + if count < rows: + padding = np.repeat(f_batch[:1], rows - count, axis=0) f_batch = np.concatenate([f_batch, padding]) samples = sample_program(self._program, jnp.asarray(f_batch), key) result[shot_batch] = np.asarray(samples)[:count] @@ -810,14 +821,20 @@ def dispatch(task: tuple[np.ndarray, np.ndarray, PRNGKey]) -> None: 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: + per_thread = min(threads, ready // MIN_THREAD_BATCH) + 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 = [] - for start in range(0, ready, batch_size): + for start in range(0, ready, size): self._key, subkey = jax.random.split(self._key) - stop = start + batch_size - batches.append((all_f[start:stop], all_shot[start:stop], subkey)) + stop = start + size + batch = (all_f[start:stop], all_shot[start:stop], subkey, rows) + batches.append(batch) run_tasks(_with_x64_of_caller(dispatch), batches, threads) pending_f, pending_shot = [all_f[ready:]], [all_shot[ready:]] pending -= ready diff --git a/test/integration/test_trajectory_sharing.py b/test/integration/test_trajectory_sharing.py index 5aeacd0d..8b4a13d5 100644 --- a/test/integration/test_trajectory_sharing.py +++ b/test/integration/test_trajectory_sharing.py @@ -257,6 +257,55 @@ def spy(program, f_params, key, **kwargs): assert sum(counts) == int((~detectors[:, 2]).sum()) +def _recording_rows(sampler, monkeypatch) -> list[int]: + rows = [] + original = sampler_module.sample_program + + def spy(program, f_params, key, **kwargs): + rows.append(f_params.shape[0]) + return original(program, f_params, key, **kwargs) + + monkeypatch.setattr(sampler_module, "sample_program", spy) + return rows + + +@pytest.mark.parametrize("sharing", [False, True]) +@pytest.mark.parametrize( + ("shots", "expected"), [(40_000, [10_000] * 4), (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 + if sharing: + assert sum(rows) == survivors and max(rows) == size + else: + # NOTE: JAX pads every batch to a power of two. + assert rows == [1 << (size - 1).bit_length()] * 4 + + @pytest.mark.parametrize("rows", [1, 300, 70_000]) @pytest.mark.parametrize("name", [*sorted(CIRCUITS), "direct_flips"]) def test_host_outputs_equal_the_jax_outputs(name, rows): From 832a621a28cb31358f29661850d4ac4cb439dade Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 19 Oct 2026 17:28:07 -0400 Subject: [PATCH 15/21] Run each shared batch in one numba function The shared executor's per-batch Python held the GIL for 20-30% of a batch on small circuits (the tape loop, lane split and prune in NumPy, group_on_host and the output assembly), which capped batch-per-thread sampling at 3-5 busy cores. `_execute` now samples a whole batch without the GIL: it groups the shots (a SplitMix64 Fisher-Yates shuffle, then stable sorts), runs the tape on lanes in arrays with a row per shot (only the amplitudes grow), and writes the outputs. Its parallel build splits a loop into chunks only once its work reaches PARALLEL_WORK, so a lone wide batch keeps the parallel kernels and small actions launch no threads. Shuffles and per-action seeds come from SplitMix64 instead of NumPy's Generator, so samples for a given seed change. One thread runs at the old speed (qv_q10 13.7M, SymFT d=7 postselected 1.39M against 1.29M shots/s). Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/shared.py | 675 ++++++++++++++++----------- test/unit/statevector/test_shared.py | 33 +- 2 files changed, 397 insertions(+), 311 deletions(-) diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index 0b2b9138..c868c994 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -20,18 +20,17 @@ 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. -Lanes live in growable arrays, and a Python loop runs the tape: each action calls -a numba kernel on the live lanes and on the ``2^extent`` amplitudes the action -can touch. 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; a pruned shot's outputs that need a branch bit set -after its pruning action are False. +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; a pruned shot's outputs +that need a branch bit set after its pruning action are False. """ from __future__ import annotations -import itertools from typing import NamedTuple import jax @@ -46,6 +45,7 @@ ROTATE_DIAG, StatevectorComponent, ) +import tsim.utils.kernels as kernels from tsim.utils.kernels import jit, kernel from tsim.utils.linalg import find_basis from tsim.utils.rng import GOLDEN, next_skip, splitmix64 @@ -224,43 +224,58 @@ def _parity(value): return value & 1 -@kernel -def _row_parities(symbols, rows): - """Return ``parity(symbols[i] & rows[r])`` for every row i and r.""" - out = np.zeros((symbols.shape[0], rows.shape[0]), dtype=np.bool_) - for i in numba.prange(symbols.shape[0]): - for r in range(rows.shape[0]): - acc = 0 - for w in range(rows.shape[1]): - acc ^= symbols[i, w] & rows[r, w] - out[i, r] = _parity(acc) - return out +@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) -@kernel -def _key_words(symbols, key_rows): - """Pack the parities of each shot's symbols with the key rows into words.""" - num_keys = key_rows.shape[0] - out = np.zeros((symbols.shape[0], (num_keys + 31) >> 5), dtype=np.uint32) - for s in numba.prange(symbols.shape[0]): - for k in range(num_keys): - 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 _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 -@kernel -def _pack_symbols(f, selection, num_words): - """Pack each shot's selected f-bits: bit j is bit j % 32 of word j // 32.""" - out = np.zeros((f.shape[0], num_words), dtype=np.uint32) - for s in numba.prange(f.shape[0]): +@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) - return out + + +@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 @@ -275,47 +290,33 @@ def _signs(symbols, lanes, row, const): return flip -@kernel -def _rotate(amps, lanes, symbols, start, stop, rotations): - """Apply the rotations ``start … stop-1`` of the tape to every lane. - - ``rotations`` holds the tape's kinds, Pauli masks, extents, signs and cosines - and the layout's factors and couplings. Rotation ``a`` is about ``X^x Z^z`` - with a sign per lane: a diagonal one multiplies amplitude i by - ``factor[a, sign ^ parity(i & z)]``, the others mix amplitudes i and - ``i ^ x`` with ``cos[a]`` and ``±coupling[a]``. +@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``. """ - kind, x_mask, z_mask, extent, sign_row, sign_const, cos, factor, coupling = ( - rotations - ) - for a in range(start, stop): - flip = _signs(symbols, lanes, sign_row[a], np.int64(sign_const[a])) - x, z, shift = np.int64(x_mask[a]), np.int64(z_mask[a]), extent[a] - span = 1 << shift - if kind[a] == ROTATE_DIAG: - for t in numba.prange(lanes * span): - lane, i = t >> shift, t & (span - 1) - amps[lane, i] *= factor[a, flip[lane] ^ _parity(i & z)] - else: - low = x & -x - for t in numba.prange(lanes * span): - lane, i = t >> shift, t & (span - 1) - if i & low: - continue - j = i ^ x - c = -coupling[a] if flip[lane] else coupling[a] - p, q = amps[lane, i], amps[lane, j] - amps[lane, i] = cos[a] * p + c * (-q if _parity(j & z) else q) - amps[lane, j] = cos[a] * q + c * (-p if _parity(i & z) else p) + 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) -@kernel -def _weights(amps, lanes, extent, x, z, i_power): - """Return each lane's outcome weights ``(w0, w1)`` of measuring ``X^x Z^z``.""" - span = 1 << extent - w0 = np.zeros(lanes) - w1 = np.zeros(lanes) - for lane in numba.prange(lanes): +@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 = np.complex128(amps[lane, i]) if x == 0: @@ -331,19 +332,17 @@ def _weights(amps, lanes, extent, x, z, i_power): plus, minus = a + b, a - b w0[lane] += (plus.real**2 + plus.imag**2) / 4 w1[lane] += (minus.real**2 + minus.imag**2) / 4 - return w0, w1 -@kernel -def _project(amps, lanes, extent, x, z, pivot, i_power, outcome, weight): - """Project lane l onto ``outcome[l]`` of weight ``weight[l]``; free the pivot. +@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⟩``. """ - span = 1 << extent bit = 1 << pivot - for lane in numba.prange(lanes): + for lane in range(lo, hi): if x == 0: scale = 1 / np.sqrt(max(weight[lane], _TINY)) else: @@ -391,12 +390,53 @@ def _draw_ones(count, w0, w1, state): return ones -@kernel -def _shot_symbols(symbols, order, owner, entry_symbols, lane_branches, dormant, seed): - """Return every shot's symbols: its f-bits, its lane's branches, its coins.""" - out = np.empty_like(symbols) +@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. + """ + num_shots = keys.shape[0] + noise = np.empty(num_shots, dtype=np.uint64) + for s in range(num_shots): + state, noise[s] = splitmix64(state) + order = np.argsort(noise) + 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 _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 numba.prange(len(order)): + 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]) @@ -405,165 +445,264 @@ def _shot_symbols(symbols, order, owner, entry_symbols, lane_branches, dormant, _, coin = splitmix64(seed + n * GOLDEN) word |= np.uint32(coin & np.uint64(0xFFFFFFFF)) & dormant[w] out[shot, w] = word - return out -class _Lanes: - """Live lanes in growable arrays: rows ``0 … size-1`` are live. +@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] + ) - Every lane holds at least one shot, so the arrays never need more rows than - there are shots. - """ - def __init__(self, symbols, start, count, width: int, dtype: np.dtype): - self.size = len(start) - self.shots = int(np.sum(count)) - capacity = min(max(16, 2 * self.size), self.shots) - self.amps = np.zeros((capacity, 1 << width), dtype=dtype) - self.amps[: self.size, 0] = 1 - self.symbols = np.zeros((capacity, symbols.shape[1]), dtype=np.uint32) - self.symbols[: self.size] = symbols - self.start = np.zeros(capacity, dtype=np.int64) - self.start[: self.size] = start - self.count = np.zeros(capacity, dtype=np.int64) - self.count[: self.size] = count - - def split(self, ones: np.ndarray, span: int) -> tuple[np.ndarray, np.ndarray]: - """Split every lane whose shots see both outcomes. - - Returns each lane's outcome and the lane its amplitudes came from (its - parent for a new child). - """ - size, count = self.size, self.count[: self.size] - outcome = ones == count - parents = np.flatnonzero((ones > 0) & ~outcome) - new = size + len(parents) - if new > len(self.count): - self._grow(new, span) - self.amps[size:new, :span] = self.amps[parents, :span] - self.symbols[size:new] = self.symbols[parents] - self.count[size:new] = ones[parents] - self.start[size:new] = self.start[parents] + count[parents] - ones[parents] - self.count[parents] -= ones[parents] - self.size = new - children = np.ones(len(parents), dtype=bool) - return np.concatenate([outcome, children]), np.concatenate( - [np.arange(size), parents] - ) +@kernel +def _execute(f, tape, layout, width, prototype, seed, parallel_work, threads, poison): + """Sample one batch of shots with shared trajectories. - def prune(self, dies: np.ndarray, span: int) -> tuple[np.ndarray, ...]: - """Remove the lanes that ``dies`` marks; return their starts, counts, symbols.""" - gone, kept = np.flatnonzero(dies), np.flatnonzero(~dies) - removed = (self.start[gone], self.count[gone], self.symbols[gone]) - self.amps[: len(kept), :span] = self.amps[kept, :span] - for array in (self.symbols, self.start, self.count): - array[: len(kept)] = array[kept] - self.size = len(kept) - return removed - - def _grow(self, needed: int, span: int) -> None: - capacity = min(max(needed, 2 * len(self.count)), self.shots) - amps = np.zeros((capacity, self.amps.shape[1]), dtype=self.amps.dtype) - amps[: self.size, :span] = self.amps[: self.size, :span] - self.amps = amps - for name in ("symbols", "start", "count"): - old = getattr(self, name) - new = np.zeros((capacity, *old.shape[1:]), dtype=old.dtype) - new[: self.size] = old[: self.size] - setattr(self, name, new) - - -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. + ``tape`` and ``layout`` hold the arrays of ``Tape`` and ``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``. + + Returns: + ``(samples, deviation, groups, peak, lanes, dead)``. - Ties are broken at random (a random permutation, then stable sorts), so the - 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, start, count - - -def _run( - layout: SharingLayout, lanes: _Lanes, rng: np.random.Generator -) -> tuple[float, list[tuple], int]: - """Run the tape on ``lanes``; return the norm deviation, dead lanes and peak. - - Every stretch of consecutive rotations runs in one kernel call, the other - actions one by one. A stretch clears the amplitudes up to its largest extent at - once; the amplitudes below ``2^region`` are the lanes' own. """ - tape = layout.tape - rotations = ( - *(tape.kind, tape.x_mask, tape.z_mask, tape.extent), - *(tape.sign_row, tape.sign_const, tape.cos, layout.factor, layout.coupling), + ( + kind, + x_mask, + z_mask, + cos, + pivot, + extent, + branch, + sign_const, + sign_row, + output_const, + output_row, + selection, + ) = tape + ( + key_rows, + lane_dormant, + pruned, + prune_outputs, + prune_flip, + lane_branches, + shot_dormant, + ready, + i_power, + factor, + coupling, + ) = layout + shots, num_words = f.shape[0], lane_branches.shape[0] + num_actions = kind.shape[0] + state = np.uint64(seed) + + symbols = np.zeros((shots, num_words), dtype=np.uint32) + keys = np.zeros((shots, (key_rows.shape[0] + 31) >> 5), dtype=np.uint32) + chunks = _chunks(shots * (len(selection) + key_rows.size), parallel_work, threads) + for c in numba.prange(chunks): + lo, hi = shots * c // chunks, shots * (c + 1) // chunks + _pack_range(f, selection, symbols, lo, hi) + _key_range(symbols, key_rows, keys, lo, hi) + order, start, count, state = _group(keys, state) + groups = len(start) + + lane_symbols = np.zeros((shots, num_words), dtype=np.uint32) + lane_start = np.zeros(shots, dtype=np.int64) + lane_count = np.zeros(shots, dtype=np.int64) + amps = np.zeros( + (min(max(16, 2 * groups), shots), 1 << width), dtype=prototype.dtype ) - rotation = np.isin(tape.kind, (ROTATE_DIAG, ROTATE)) - others = np.flatnonzero(~rotation) - edges = sorted({0, len(rotation), *others, *(others + 1)}) - dead: list[tuple] = [] - deviation, peak, region = 0.0, lanes.size, 0 - for a, stop in itertools.pairwise(edges): - if not lanes.size: + 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 + dead_symbols = np.zeros((shots, num_words), dtype=np.uint32) + dead_start = np.zeros(shots, dtype=np.int64) + dead_count = np.zeros(shots, dtype=np.int64) + dead_action = np.zeros(shots, dtype=np.int64) + size, dead, peak, region, deviation = groups, 0, groups, 0, 0.0 + if poison: + _poison(amps, size, 1) + + for a in range(num_actions): + if size == 0: break - extent = int(tape.extent[a:stop].max()) - span = 1 << extent - if extent > region: - lanes.amps[: lanes.size, 1 << region : span] = 0 - region = extent - size = lanes.size - work = size * span - if rotation[a]: - _rotate(work, lanes.amps, size, lanes.symbols, a, stop, rotations) + shift = extent[a] + span = 1 << shift + if shift > region: + _clear(amps, size, 1 << region, span) + region = shift + x, z = np.int64(x_mask[a]), np.int64(z_mask[a]) + if kind[a] == ROTATE_DIAG or kind[a] == ROTATE: + flip = _signs(lane_symbols, size, sign_row[a], np.int64(sign_const[a])) + work = size * span + chunks = _chunks(work, parallel_work, threads) + for c in numba.prange(chunks): + lo, hi = work * c // chunks, work * (c + 1) // chunks + _rotate_range( + amps, + lo, + hi, + shift, + flip, + kind[a] == ROTATE_DIAG, + x, + z, + cos[a], + factor[a], + coupling[a], + ) continue - kind = tape.kind[a] - x, z = int(tape.x_mask[a]), int(tape.z_mask[a]) - i_power = layout.i_power[a] - if kind == MEASURE_DORMANT: - if not layout.lane_dormant[a]: - continue - w0 = w1 = np.full(size, 0.5) + dormant = kind[a] == MEASURE_DORMANT + if dormant and not lane_dormant[a]: + continue + half = 0.5 if dormant else 0.0 + w0 = np.full(size, half) + w1 = np.full(size, half) + if not dormant: + chunks = _chunks(size * span, parallel_work, threads) + 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[a], 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 poison: + _poison(amps, size, span) + if not dormant: + chunks = _chunks(size * span, parallel_work, threads) + for c in numba.prange(chunks): + lo, hi = size * c // chunks, size * (c + 1) // chunks + _project_range( + amps, lo, hi, span, x, z, pivot[a], i_power[a], outcome, weight + ) + bit = 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 pruned[a]: + kept = 0 + for lane in range(size): + dies = False + for o in range(output_row.shape[0]): + if not prune_outputs[a, o]: + continue + acc = 0 + for w in range(num_words): + acc ^= lane_symbols[lane, w] & output_row[o, w] + if _parity(acc) ^ output_const[o] ^ prune_flip[o]: + dies = True + break + if dies: + dead_symbols[dead] = lane_symbols[lane] + dead_start[dead], dead_count[dead] = ( + lane_start[lane], + lane_count[lane], + ) + dead_action[dead] = a + dead += 1 + continue + if kept != lane: + _copy_lane(amps, kept, lane, span) + lane_symbols[kept] = lane_symbols[lane] + lane_start[kept], lane_count[kept] = ( + lane_start[lane], + lane_count[lane], + ) + kept += 1 + size = kept + if poison: + _poison(amps, size, span) + + # NOTE: the live and dead lanes' ranges tile the sorted positions. + 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: - w0, w1 = _weights(work, lanes.amps, size, extent, x, z, i_power) - deviation = float(np.max(np.abs(w0 + w1 - 1), initial=deviation)) - seed = rng.integers(0, 2**64, dtype=np.uint64) - ones = _draw_ones(lanes.count[:size], w0, w1, seed) - outcome, source = lanes.split(ones, span) - if kind != MEASURE_DORMANT: - weight = np.where(outcome, w1[source], w0[source]) - _project( - lanes.size * span, - lanes.amps, - lanes.size, - extent, - x, - z, - int(tape.pivot[a]), - i_power, - outcome, - weight, - ) - branch = int(tape.branch[a]) - lanes.symbols[: lanes.size, branch >> 5] |= outcome.astype(np.uint32) << ( - branch & 31 + 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) + 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, ) - peak = max(peak, lanes.size) - if layout.pruned[a]: - values = _row_parities(0, lanes.symbols[: lanes.size], tape.output_row) - values ^= tape.output_const ^ layout.prune_flip - dies = np.any(values & layout.prune_outputs[a], axis=1) - if dies.any(): - dead.append((*lanes.prune(dies, span), a)) - return deviation, dead, peak + 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 + + +# NOTE: tests set this to fill every amplitude the executor must not read with NaN. +_POISON = False def sample_shared_component( @@ -587,54 +726,30 @@ def sample_shared_component( """ tape = layout.tape key, subkey = jax.random.split(key) - rng = np.random.default_rng(np.asarray(jax.random.bits(subkey, (4,)))) + high, low = np.asarray(jax.random.bits(subkey, (2,), dtype=np.uint32)) + seed = (np.uint64(high) << np.uint64(32)) | np.uint64(low) f = np.asarray(f_params) - symbols = _pack_symbols(f.size, f, tape.f_selection, component.num_words) - work = len(symbols) * max(1, layout.key_rows.size) - keys = _key_words(work, symbols, layout.key_rows) - order, start, count = group_on_host(keys, rng) - representative = symbols[order[start]] - lanes = _Lanes( - representative, start, count, component.width, np.dtype(component.dtype) + tape_arrays = ( + *(tape.kind, tape.x_mask, tape.z_mask, tape.cos, tape.pivot, tape.extent), + *(tape.branch, tape.sign_const, tape.sign_row, tape.output_const), + *(tape.output_row, tape.f_selection), ) - deviation, dead, peak = _run(layout, lanes, rng) - - num_actions = len(tape.kind) - live = slice(0, lanes.size) - entries = [ - (lanes.start[live], lanes.count[live], lanes.symbols[live], num_actions), - *dead, - ] - entry_start = np.concatenate([e[0] for e in entries]) - entry_count = np.concatenate([e[1] for e in entries]) - entry_symbols = np.concatenate([e[2] for e in entries]) - entry_pruned = np.concatenate([np.full(len(e[0]), e[3]) for e in entries]) - # NOTE: the entries' ranges tile the sorted positions, so sorting them by - # start gives every position's entry. - by_start = np.argsort(entry_start) - owner = np.repeat(by_start, entry_count[by_start]) - coin_seed = rng.integers(0, 2**64, dtype=np.uint64) - shot_symbols = _shot_symbols( - symbols.size, - symbols, - order, - owner, - entry_symbols, - layout.lane_branches, - layout.shot_dormant, - coin_seed, + layout_arrays = ( + *(layout.key_rows, layout.lane_dormant, layout.pruned, layout.prune_outputs), + *(layout.prune_flip, layout.lane_branches, layout.shot_dormant, layout.ready), + *(layout.i_power, layout.factor, layout.coupling), ) - values = _row_parities( - shot_symbols.size * len(tape.output_row), shot_symbols, tape.output_row + samples, deviation, groups, peak, lanes, dead = _execute( + f.shape[0] << component.width, + f, + tape_arrays, + layout_arrays, + component.width, + np.zeros(0, dtype=component.dtype), + seed, + kernels.PARALLEL_WORK, + numba.get_num_threads(), + _POISON, ) - values ^= tape.output_const - pruned_at = np.empty(len(order), dtype=np.int64) - pruned_at[order] = entry_pruned[owner] - samples = values & (layout.ready[None, :] <= pruned_at[:, None]) - stats = { - "groups": len(start), - "peak lanes": peak, - "lanes": lanes.size, - "dead": sum(len(e[0]) for e in dead), - } - return samples, key, deviation, stats + stats = {"groups": groups, "peak lanes": peak, "lanes": lanes, "dead": dead} + return samples, key, float(deviation), stats diff --git a/test/unit/statevector/test_shared.py b/test/unit/statevector/test_shared.py index f50c6e67..16a5c6bf 100644 --- a/test/unit/statevector/test_shared.py +++ b/test/unit/statevector/test_shared.py @@ -243,36 +243,6 @@ def coherent(): return planned, component_layout, compiled -class _PoisonedLanes(shared._Lanes): - """Lanes that hold NaN wherever the executor must not read. - - After every operation, that is every row from ``size`` up and every - amplitude of a live row from ``span`` up. - """ - - def __init__(self, *args): - super().__init__(*args) - self._poison(1) - - def split(self, ones, span): - result = super().split(ones, span) - self._poison(span) - return result - - def prune(self, dies, span): - removed = super().prune(dies, span) - self._poison(span) - return removed - - def _grow(self, needed, span): - super()._grow(needed, span) - self._poison(span) - - def _poison(self, span): - self.amps[: self.size, span:] = np.nan - self.amps[self.size :] = np.nan - - 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)) @@ -291,7 +261,8 @@ def with_extent(values): ) key = jax.random.key(17) reference, _, deviation, _ = sample_shared_component(compiled, full, f_params, key) - monkeypatch.setattr(shared, "_Lanes", _PoisonedLanes) + # NOTE: NaN fills every amplitude the executor must not read. + monkeypatch.setattr(shared, "_POISON", True) for tape in (layout, raised): samples, _, other, stats = sample_shared_component( compiled, tape, f_params, key From 960cff0ae11b5e4a2161813d3761974531c7534f Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 19 Oct 2026 17:47:59 -0400 Subject: [PATCH 16/21] Sample strategy="active_statevector" without JAX The statevector strategy now runs entirely on the host: the numba executor samples both with trajectory sharing and with one lane per shot (share=False skips the grouping), so the JAX executor in execute.py and its tests are gone. StatevectorComponent is a frozen dataclass of NumPy arrays; batch seeds come from a NumPy Generator instead of JAX keys, and the memory budget is the host's. Every statevector batch gets a sharing layout, so postselected component detectors prune lanes in both modes. trajectory_sharing now defaults to True for active_statevector. The precision still follows jax_enable_x64 at compile time. Co-Authored-By: Claude Opus 5.5 --- src/tsim/circuit.py | 10 +- src/tsim/sampler.py | 329 +++++++++--------- src/tsim/statevector/execute.py | 297 ---------------- src/tsim/statevector/program.py | 73 ++-- src/tsim/statevector/shared.py | 154 +++++++-- test/integration/test_statevector_sampler.py | 16 +- test/integration/test_trajectory_sharing.py | 69 ++-- test/unit/statevector/test_execute.py | 345 ------------------- test/unit/statevector/test_shared.py | 64 ++-- test/unit/test_postselection.py | 82 ++--- test/unit/test_sampler.py | 13 +- 11 files changed, 434 insertions(+), 1018 deletions(-) delete mode 100644 src/tsim/statevector/execute.py delete mode 100644 test/unit/statevector/test_execute.py diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index 2583bede..c3066d2e 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. @@ -877,7 +877,8 @@ def compile_sampler( 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 @@ -898,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. @@ -915,7 +916,8 @@ def compile_detector_sampler( 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/sampler.py b/src/tsim/sampler.py index fe821b16..a946c158 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -4,7 +4,6 @@ import warnings from collections import deque -from collections.abc import Callable from math import ceil from typing import TYPE_CHECKING, Literal, overload @@ -19,10 +18,10 @@ 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 ( SharingLayout, + lane_bytes, sample_shared_component, sharing_layout, ) @@ -47,6 +46,8 @@ SCREEN_SHOTS = 2**14 SCREEN_CHUNKS_PER_THREAD = 2 +_GOLDEN = 0x9E3779B97F4A7C15 + def _sample_component( component: CompiledComponent, @@ -115,7 +116,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]: @@ -131,109 +132,112 @@ 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 sample_program( - program: CompiledProgram, - f_params: jax.Array | np.ndarray, - key: PRNGKey, - sharing: tuple[SharingLayout | None, ...] | None = None, - sharing_stats: deque[dict] | None = None, -) -> jax.Array | np.ndarray: - """Sample all outputs from a compiled program. +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, + ) - NumPy ``f_params`` give NumPy samples, put together by a numba kernel (the - shared executor works on the host); JAX ``f_params`` give JAX samples. + +def sample_program( + program: CompiledProgram, f_params: jax.Array, key: PRNGKey +) -> jax.Array: + """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 layout 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 | np.ndarray] = [] - - on_host = isinstance(f_params, np.ndarray) if program.num_outputs == 0: - batch_size = f_params.shape[0] - if on_host: - return np.zeros((batch_size, 0), dtype=np.bool_) - return jnp.zeros((batch_size, 0), dtype=jnp.bool_) - - if len(program.direct_f_indices) > 0 and not on_host: + 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, jnp.asarray(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 - # NOTE: written so that a NaN deviation also warns. - if not 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) - - if on_host: - assert isinstance(f_params, np.ndarray) - rows = f_params.shape[0] - component_bits = ( - np.concatenate(results, axis=1) - if results - else np.zeros((rows, 0), dtype=np.bool_) - ) - out = np.empty((rows, program.num_outputs), dtype=np.bool_) - _gather_outputs( - out.size, - f_params, - np.asarray(program.direct_f_indices), - np.asarray(program.direct_flips), - component_bits, - np.asarray(program.output_order), - out, - ) - return out 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, ...], + *, + share: bool, + direct: tuple[np.ndarray, np.ndarray, np.ndarray], + 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. + share: Whether shots with equal histories share a lane. + direct: The program's direct f-indices, direct flips and output order. + 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 * _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, *direct[:2], component_bits, direct[2], out) + return out + + @kernel def _write_direct_detectors(f, f_indices, flips, columns, postselect, reference, out): """Clear every row of ``out`` and write its direct detector ``j`` at ``columns[j]``. @@ -267,21 +271,9 @@ def _gather_outputs(f, f_indices, flips, component_bits, order, out): out[row, order[num_direct + j]] = component_bits[row, j] -def _with_x64_of_caller(function: Callable) -> Callable: - """Return ``function`` run with the calling thread's JAX 64-bit mode. - - ``jax.enable_x64`` only reaches the thread that enters it, not the workers. - """ - x64 = jax.config.read("jax_enable_x64") - - def run(task): - with jax.enable_x64(x64): - return function(task) - - return run - - def _nbytes(tree: object) -> int: + if isinstance(tree, StatevectorComponent): + return sum(v.nbytes for v in vars(tree).values() if isinstance(v, np.ndarray)) return sum( v.nbytes for v in jax.tree_util.tree_leaves(tree) if isinstance(v, jax.Array) ) @@ -298,7 +290,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. @@ -316,21 +308,26 @@ def __init__( 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) + # NOTE: statevector batches draw their seeds here instead of from the key. + self._batch_seeds = np.random.default_rng([seed, 1]) self._trajectory_sharing = trajectory_sharing - self._sharing: dict[bytes, tuple[SharingLayout | None, ...]] = {} + self._layouts_by_key: dict[bytes, tuple[SharingLayout, ...]] = {} self.sharing_stats: deque[dict] = deque(maxlen=10_000) prepared = prepare_graph(circuit, sample_detectors=sample_detectors) @@ -368,6 +365,11 @@ def __init__( self._direct_global_indices = np.asarray( prog.output_order[:n_direct], dtype=np.int32 ) + self._direct = ( + self._direct_f_indices, + self._direct_flips, + np.asarray(prog.output_order), + ) 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 @@ -394,12 +396,12 @@ def __init__( f"{STABILIZER_RANK_HINT}" ) - def _sharing_arguments( + def _layouts( self, postselection_mask: np.ndarray | None, detector_reference: np.ndarray | None = None, - ) -> dict: - """Return the extra ``sample_program`` arguments: none without sharing. + ) -> tuple[SharingLayout, ...]: + """Return the sharing layout of every statevector component. Args: postselection_mask: Postselected detectors, whose lanes are pruned. @@ -407,8 +409,6 @@ def _sharing_arguments( 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: @@ -417,30 +417,39 @@ def _sharing_arguments( 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( - ( + if key not in self._layouts_by_key: + layouts = [] + for component in self._program.components: + assert isinstance(component, StatevectorComponent) + detector = np.asarray(component.output_indices) < num_detectors + indices = np.minimum(component.output_indices, num_detectors - 1) + layouts.append( 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 - ] - ), + detector & mask[indices] if num_detectors else detector, + detector & flip[indices] if num_detectors else detector, ) - if isinstance(component, StatevectorComponent) - else None ) - for component in self._program.components - ) - return {"sharing": self._sharing[key], "sharing_stats": self.sharing_stats} + self._layouts_by_key[key] = tuple(layouts) + return self._layouts_by_key[key] + + def _sample_statevector( + self, f_params: np.ndarray, seed: int, layouts: tuple[SharingLayout, ...] + ) -> np.ndarray: + """Sample one batch of a statevector program.""" + return sample_statevector_program( + self._program, + f_params, + seed, + layouts, + share=self._trajectory_sharing, + direct=self._direct, + sharing_stats=self.sharing_stats, + ) + + def _next_seed(self) -> int: + """Return the next statevector batch seed.""" + return int(self._batch_seeds.integers(0, 2**64, dtype=np.uint64)) def _compute_direct_outputs(self, f_params_np: np.ndarray) -> np.ndarray: """Scatter direct output bits into a full (batch, num_outputs) bool array. @@ -476,6 +485,9 @@ def _compute_reference_sample(self) -> np.ndarray: f_ref = np.zeros((1, num_f), dtype=np.uint8) if not self._program.components: return self._compute_direct_outputs(f_ref)[0] + if self._has_statevector: + seed = self._next_seed() + return self._sample_statevector(f_ref, seed, self._layouts(None))[0] self._key, subkey = jax.random.split(self._key) return np.asarray( sample_program(self._program, jnp.asarray(f_ref), subkey)[0], @@ -529,9 +541,12 @@ def _peak_bytes_per_sample(self) -> int: return max(peak, 1) def _memory_budget(self) -> int: - """Return the memory one batch may use: half of the free device memory.""" + """Return the memory one batch may use: half of the free device memory. + + Statevector programs sample on the host, so their device is the host. + """ device = jax.devices()[0] - if device.platform == "gpu": + if device.platform == "gpu" and not self._has_statevector: stats = device.memory_stats() available = stats.get("bytes_limit", 8 * 1024**3) - stats.get( "bytes_in_use", 0 @@ -543,26 +558,24 @@ def _memory_budget(self) -> int: def _threads(self, batch_size: int) -> int: """Return how many batches of ``batch_size`` shots run at once. - Statevector components sample one batch per numba thread on the CPU: in - the shared executor, or in JAX on a CPU backend as long as the batches - fit in the memory budget. The other executors sample one batch at a time. + Statevector programs sample one batch per numba thread, as long as the + batches fit in the memory budget; the stabilizer-rank programs sample + one batch at a time. """ if not self._has_statevector: return 1 threads = numba.get_num_threads() if self._trajectory_sharing: return threads - if jax.default_backend() != "cpu": - return 1 fits = self._memory_budget() // (self._peak_bytes_per_sample() * batch_size) return max(1, min(threads, fits)) def _estimate_batch_size(self) -> int: """Estimate the largest batch size that fits in available device memory. - JAX allocates a lane per shot in every batch that runs at once. The shared - executor allocates lanes as its trajectories split, so its estimate leaves - room for a lane per shot in one batch only. Programs with statevector + 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. Programs with statevector components use at most ``STATEVECTOR_MAX_BATCH_SIZE`` lanes per batch. """ peak = self._peak_bytes_per_sample() @@ -636,31 +649,32 @@ def _sample_batches( ) num_batches = ceil((shots + compute_reference) / batch_size) - sharing = self._sharing_arguments(None) + statevector = self._has_statevector + layouts = self._layouts(None) if statevector else () tasks = [] for index in range(num_batches): seeds = self._channel_sampler.block_seeds(batch_size) - self._key, subkey = jax.random.split(self._key) - tasks.append((index, seeds, subkey)) + if statevector: + tasks.append((index, seeds, self._next_seed())) + else: + self._key, subkey = jax.random.split(self._key) + tasks.append((index, seeds, subkey)) - def sample_batch( - task: tuple[int, np.ndarray, PRNGKey], - ) -> jax.Array | np.ndarray: + def sample_batch(task: tuple) -> jax.Array | np.ndarray: index, seeds, key = task f_params_np = self._channel_sampler.sample(batch_size, seeds) if compute_reference and index == 0: f_params_np[0] = 0 - f_params = f_params_np if sharing else jnp.asarray(f_params_np) - return sample_program(self._program, f_params, key, **sharing) + if statevector: + return self._sample_statevector(f_params_np, key, layouts) + return sample_program(self._program, jnp.asarray(f_params_np), key) - batches = run_tasks( - _with_x64_of_caller(sample_batch), tasks, self._threads(batch_size) - ) + batches = run_tasks(sample_batch, tasks, self._threads(batch_size)) reference: np.ndarray | None = None if compute_reference: reference = np.asarray(batches[0][0]) batches[0] = batches[0][1:] - if sharing: + if statevector: result = np.concatenate(batches)[:shots] else: # Concatenate on device, then a single d2h. The prior @@ -747,8 +761,11 @@ def _sample_batches_with_postselection( else np.zeros(len(columns), dtype=np.bool_) ) - sharing = self._sharing_arguments( - postselection_mask, reference if xor_detector_ref else None + statevector = self._has_statevector + layouts = ( + self._layouts(postselection_mask, reference if xor_detector_ref else None) + if statevector + else () ) chunk = min(batch_size, SCREEN_SHOTS) @@ -776,16 +793,16 @@ def screen(task: tuple[int, np.ndarray]) -> tuple[np.ndarray, np.ndarray]: survivors = np.flatnonzero(~discarded) return f_params_np[survivors], start + survivors - def dispatch(task: tuple[np.ndarray, np.ndarray, PRNGKey, int]) -> None: + def dispatch(task: tuple) -> None: """Sample the shots ``shot_batch`` in one batch. - JAX batches are padded with row 0 to ``rows`` rows; the shared - executor takes any number of rows. + JAX batches are padded with row 0 to ``rows`` rows; statevector + batches take any number of rows. """ f_batch, shot_batch, key, rows = task count = len(shot_batch) - if sharing: - samples = sample_program(self._program, f_batch, key, **sharing) + if statevector: + samples = self._sample_statevector(f_batch, key, layouts) else: if count < rows: padding = np.repeat(f_batch[:1], rows - count, axis=0) @@ -831,11 +848,13 @@ def dispatch(task: tuple[np.ndarray, np.ndarray, PRNGKey, int]) -> None: all_f, all_shot = np.concatenate(pending_f), np.concatenate(pending_shot) batches = [] for start in range(0, ready, size): - self._key, subkey = jax.random.split(self._key) + if statevector: + key = self._next_seed() + else: + self._key, key = jax.random.split(self._key) stop = start + size - batch = (all_f[start:stop], all_shot[start:stop], subkey, rows) - batches.append(batch) - run_tasks(_with_x64_of_caller(dispatch), batches, threads) + batches.append((all_f[start:stop], all_shot[start:stop], key, rows)) + run_tasks(dispatch, batches, threads) pending_f, pending_shot = [all_f[ready:]], [all_shot[ready:]] pending -= ready @@ -953,7 +972,7 @@ def __init__( *, strategy: SamplingStrategy = "active_statevector", seed: int | None = None, - trajectory_sharing: bool = False, + trajectory_sharing: bool | None = None, ): """Create a measurement sampler. @@ -962,7 +981,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 @@ -1014,7 +1034,7 @@ def __init__( *, strategy: SamplingStrategy = "active_statevector", seed: int | None = None, - trajectory_sharing: bool = False, + trajectory_sharing: bool | None = None, ): """Create a detector sampler. @@ -1023,7 +1043,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 diff --git a/src/tsim/statevector/execute.py b/src/tsim/statevector/execute.py deleted file mode 100644 index 3826b3f0..00000000 --- a/src/tsim/statevector/execute.py +++ /dev/null @@ -1,297 +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``. -""" - -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)``. - max_norm_deviation: Largest ``|‖ψ‖² - 1|`` seen at a measurement. - - """ - - amps: Array - symbols: Array - 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, 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. - """ - 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), - 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, 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: - lanes = state.symbols.shape[0] - half = jnp.full(lanes, 0.5, dtype=state.max_norm_deviation.dtype) - outcome = choose_outcome(half, half, 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/program.py b/src/tsim/statevector/program.py index d42f098d..dacffdaa 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 -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 @@ -40,7 +38,8 @@ ) -class StatevectorComponent(eqx.Module): +@dataclass(frozen=True, eq=False) +class StatevectorComponent: """A component compiled for ``strategy="active_statevector"``: an action tape. Attributes: @@ -64,22 +63,22 @@ 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 - extent: Array - branch: Array - sign_const: Array - sign_row: Array - output_const: Array - output_row: Array + output_indices: tuple[int, ...] + width: int + dtype: str + 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 @property def num_actions(self) -> int: @@ -178,19 +177,19 @@ def lower(plan: Plan, layout: ComponentLayout, dtype: str) -> StatevectorCompone 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), - extent=jnp.asarray(extent), - 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), + f_selection=np.array(layout.f_selection, dtype=np.int32), + kind=kind, + x_mask=x_mask, + z_mask=z_mask, + cos=np.cos(angle).astype(real), + sin=np.sin(angle).astype(real), + 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, ) @@ -208,5 +207,5 @@ def compile_statevector_component( """ planned, component_layout = plan_component(component, output_indices, layout) - dtype = jax.dtypes.canonicalize_dtype(np.complex128).name + dtype = "complex128" if jax.config.read("jax_enable_x64") else "complex64" return lower(planned, component_layout, dtype) diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index c868c994..221a391e 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -33,7 +33,6 @@ from typing import NamedTuple -import jax import numba import numpy as np @@ -398,10 +397,11 @@ def _group(keys, state): shots of a group appear in uniformly random order. """ num_shots = keys.shape[0] - noise = np.empty(num_shots, dtype=np.uint64) - for s in range(num_shots): - state, noise[s] = splitmix64(state) - order = np.argsort(noise) + 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) @@ -463,7 +463,9 @@ def _output_range( @kernel -def _execute(f, tape, layout, width, prototype, seed, parallel_work, threads, poison): +def _execute( + f, tape, layout, width, prototype, seed, share, parallel_work, threads, poison +): """Sample one batch of shots with shared trajectories. ``tape`` and ``layout`` hold the arrays of ``Tape`` and ``SharingLayout``. @@ -509,16 +511,27 @@ def _execute(f, tape, layout, width, prototype, seed, parallel_work, threads, po symbols = np.zeros((shots, num_words), dtype=np.uint32) keys = np.zeros((shots, (key_rows.shape[0] + 31) >> 5), dtype=np.uint32) chunks = _chunks(shots * (len(selection) + key_rows.size), parallel_work, threads) - for c in numba.prange(chunks): - lo, hi = shots * c // chunks, shots * (c + 1) // chunks + if chunks == 1: + lo, hi = 0, shots _pack_range(f, selection, symbols, lo, hi) _key_range(symbols, key_rows, keys, lo, hi) - order, start, count, state = _group(keys, state) + else: + for c in numba.prange(chunks): + lo, hi = shots * c // chunks, shots * (c + 1) // chunks + _pack_range(f, selection, symbols, lo, hi) + _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) - lane_symbols = np.zeros((shots, num_words), dtype=np.uint32) - lane_start = np.zeros(shots, dtype=np.int64) - lane_count = np.zeros(shots, dtype=np.int64) + # 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=prototype.dtype ) @@ -526,10 +539,10 @@ def _execute(f, tape, layout, width, prototype, seed, parallel_work, threads, po lane_symbols[g] = symbols[order[start[g]]] lane_start[g], lane_count[g] = start[g], count[g] amps[g, 0] = 1 - dead_symbols = np.zeros((shots, num_words), dtype=np.uint32) - dead_start = np.zeros(shots, dtype=np.int64) - dead_count = np.zeros(shots, dtype=np.int64) - dead_action = np.zeros(shots, dtype=np.int64) + dead_symbols = np.empty((shots, num_words), dtype=np.uint32) + dead_start = np.empty(shots, dtype=np.int64) + dead_count = np.empty(shots, dtype=np.int64) + dead_action = np.empty(shots, dtype=np.int64) size, dead, peak, region, deviation = groups, 0, groups, 0, 0.0 if poison: _poison(amps, size, 1) @@ -547,8 +560,8 @@ def _execute(f, tape, layout, width, prototype, seed, parallel_work, threads, po flip = _signs(lane_symbols, size, sign_row[a], np.int64(sign_const[a])) work = size * span chunks = _chunks(work, parallel_work, threads) - for c in numba.prange(chunks): - lo, hi = work * c // chunks, work * (c + 1) // chunks + if chunks == 1: + lo, hi = 0, work _rotate_range( amps, lo, @@ -562,6 +575,22 @@ def _execute(f, tape, layout, width, prototype, seed, parallel_work, threads, po factor[a], coupling[a], ) + else: + for c in numba.prange(chunks): + lo, hi = work * c // chunks, work * (c + 1) // chunks + _rotate_range( + amps, + lo, + hi, + shift, + flip, + kind[a] == ROTATE_DIAG, + x, + z, + cos[a], + factor[a], + coupling[a], + ) continue dormant = kind[a] == MEASURE_DORMANT if dormant and not lane_dormant[a]: @@ -571,9 +600,13 @@ def _execute(f, tape, layout, width, prototype, seed, parallel_work, threads, po w1 = np.full(size, half) if not dormant: chunks = _chunks(size * span, parallel_work, threads) - for c in numba.prange(chunks): - lo, hi = size * c // chunks, size * (c + 1) // chunks + if chunks == 1: + lo, hi = 0, size _weights_range(amps, lo, hi, span, x, z, i_power[a], 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[a], w0, w1) for lane in range(size): d = abs(w0[lane] + w1[lane] - 1) # NOTE: written so that a NaN deviation sticks. @@ -610,11 +643,17 @@ def _execute(f, tape, layout, width, prototype, seed, parallel_work, threads, po _poison(amps, size, span) if not dormant: chunks = _chunks(size * span, parallel_work, threads) - for c in numba.prange(chunks): - lo, hi = size * c // chunks, size * (c + 1) // chunks + if chunks == 1: + lo, hi = 0, size _project_range( amps, lo, hi, span, x, z, pivot[a], i_power[a], 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[a], i_power[a], outcome, weight + ) bit = branch[a] if bit >= 0: for lane in range(size): @@ -679,8 +718,8 @@ def _execute(f, tape, layout, width, prototype, seed, parallel_work, threads, po 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) - for c in numba.prange(chunks): - lo, hi = shots * c // chunks, shots * (c + 1) // chunks + if chunks == 1: + lo, hi = 0, shots _shot_range( symbols, order, @@ -693,11 +732,39 @@ def _execute(f, tape, layout, width, prototype, seed, parallel_work, threads, po lo, hi, ) - for c in numba.prange(chunks): - lo, hi = shots * c // chunks, shots * (c + 1) // chunks + 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, + ) + if chunks == 1: + lo, hi = 0, shots _output_range( shot_symbols, output_row, output_const, ready, pruned_at, samples, lo, hi ) + else: + 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 @@ -705,29 +772,39 @@ def _execute(f, tape, layout, width, prototype, seed, parallel_work, threads, po _POISON = False +def lane_bytes(component: StatevectorComponent) -> int: + """Return the peak memory of one lane of ``sample_shared_component``. + + 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(component.dtype).itemsize + return amplitudes + 20 * component.num_words + 88 + component.num_outputs + + def sample_shared_component( component: StatevectorComponent, layout: SharingLayout, - f_params: jax.Array | np.ndarray, - key: jax.Array, -) -> tuple[np.ndarray, jax.Array, float, dict]: - """Sample one component's outputs with shared trajectories. + 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. layout: Its sharing layout. f_params: Error-basis bits of every shot, shape ``(B, num_f_global)``. - key: JAX random key. + 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. """ tape = layout.tape - key, subkey = jax.random.split(key) - high, low = np.asarray(jax.random.bits(subkey, (2,), dtype=np.uint32)) - seed = (np.uint64(high) << np.uint64(32)) | np.uint64(low) f = np.asarray(f_params) tape_arrays = ( *(tape.kind, tape.x_mask, tape.z_mask, tape.cos, tape.pivot, tape.extent), @@ -746,10 +823,11 @@ def sample_shared_component( layout_arrays, component.width, np.zeros(0, dtype=component.dtype), - seed, + np.uint64(seed), + share, kernels.PARALLEL_WORK, numba.get_num_threads(), _POISON, ) stats = {"groups": groups, "peak lanes": peak, "lanes": lanes, "dead": dead} - return samples, key, float(deviation), stats + return samples, float(deviation), stats diff --git a/test/integration/test_statevector_sampler.py b/test/integration/test_statevector_sampler.py index 3e6efd5b..d5c62672 100644 --- a/test/integration/test_statevector_sampler.py +++ b/test/integration/test_statevector_sampler.py @@ -15,7 +15,7 @@ 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.shared import lane_bytes from tsim.statevector.program import StatevectorComponent @@ -123,12 +123,18 @@ def test_complex128_sampling_matches_stabilizer_rank(): def _norm_warnings(monkeypatch, strategy, deviation, *, x64=False): 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) + monkeypatch.setattr(sampler_module, "sample_shared_component", drifting_shared) with jax.enable_x64(x64): sampler = Circuit(CIRCUITS["t_ghz"]).compile_detector_sampler( strategy=strategy, seed=0 @@ -176,13 +182,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 8b4a13d5..5e8313f6 100644 --- a/test/integration/test_trajectory_sharing.py +++ b/test/integration/test_trajectory_sharing.py @@ -5,7 +5,6 @@ from test.helpers.statevector import CIRCUITS, assert_counts_match import jax -import jax.numpy as jnp import numpy as np import pytest @@ -174,6 +173,13 @@ def test_shared_sampling_of_random_noisy_circuits(seed): assert_counts_match(sampler.sample(shots, batch_size=shots), expected) +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( @@ -232,43 +238,35 @@ def sample(threads): 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 - rows = [] - original = sampler_module.sample_program - - def spy(program, f_params, key, **kwargs): - rows.append((type(f_params), f_params.shape[0])) - return original(program, f_params, key, **kwargs) - - monkeypatch.setattr(sampler_module, "sample_program", spy) + counts = _recording_rows(sampler, monkeypatch) mask = np.ones(circuit.num_detectors, dtype=bool) detectors = sampler.sample(5_003, batch_size=300, postselection_mask=mask) - counts = [count for _, count in rows] - assert all(kind is np.ndarray for kind, _ in rows) 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()) -def _recording_rows(sampler, monkeypatch) -> list[int]: - rows = [] - original = sampler_module.sample_program - - def spy(program, f_params, key, **kwargs): - rows.append(f_params.shape[0]) - return original(program, f_params, key, **kwargs) - - monkeypatch.setattr(sampler_module, "sample_program", spy) - return rows - - @pytest.mark.parametrize("sharing", [False, True]) @pytest.mark.parametrize( ("shots", "expected"), [(40_000, [10_000] * 4), (6_000, [6_000])] @@ -299,37 +297,28 @@ def test_default_postselection_splits_the_survivors_over_the_threads( survivors = int((~detectors[:, 2]).sum()) size = -(-survivors // 4) assert len(rows) == 4 - if sharing: - assert sum(rows) == survivors and max(rows) == size - else: - # NOTE: JAX pads every batch to a power of two. - assert rows == [1 << (size - 1).bit_length()] * 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_host_outputs_equal_the_jax_outputs(name, rows): +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, trajectory_sharing=True + strategy="active_statevector", seed=1 ) if name == "direct_flips": assert sampler._direct_flips.any() assert sampler._program.output_reindex is not None f_params = sampler._channel_sampler.sample(rows) - key = jax.random.key(2) - sharing = sampler._sharing_arguments(None) - program = sampler._program - on_host = sampler_module.sample_program(program, f_params, key, **sharing) - on_device = sampler_module.sample_program( - program, jnp.asarray(f_params), key, **sharing - ) - assert isinstance(on_host, np.ndarray) - np.testing.assert_array_equal(on_host, np.asarray(on_device)) + samples = sampler._sample_statevector(f_params, 2, sampler._layouts(None)) + columns = sampler._direct_global_indices + expected = (f_params[:, sampler._direct_f_indices] != 0) ^ sampler._direct_flips + np.testing.assert_array_equal(samples[:, columns], expected) @pytest.mark.parametrize("sharing", [False, True]) -def test_workers_sample_in_the_callers_64_bit_mode(sharing): +def test_complex128_samples_do_not_depend_on_the_thread_count(sharing): circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) def sample(threads): diff --git a/test/unit/statevector/test_execute.py b/test/unit/statevector/test_execute.py deleted file mode 100644 index bec80346..00000000 --- a/test/unit/statevector/test_execute.py +++ /dev/null @@ -1,345 +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, 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, 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, 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) - return jnp.sum(execute.choose_outcome(p0s, p1s, 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, key): - recorded.append((np.asarray(p0, np.float64), np.asarray(p1, np.float64))) - return choose(p0, p1, 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, 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_shared.py b/test/unit/statevector/test_shared.py index 16a5c6bf..2b87d862 100644 --- a/test/unit/statevector/test_shared.py +++ b/test/unit/statevector/test_shared.py @@ -7,8 +7,6 @@ quantum_components, ) -import jax -import jax.numpy as jnp import numpy as np import pytest import stim @@ -79,15 +77,16 @@ 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): 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, layout, 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( @@ -102,11 +101,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_layout(compiled), - jnp.asarray(f_params[permutation]), - jax.random.key(3), + np.asarray(f_params[permutation]), + 3, ) samples = np.asarray(samples)[np.argsort(permutation)] for i, f in enumerate(patterns): @@ -114,9 +113,10 @@ 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) @@ -126,8 +126,8 @@ def test_pruning_discards_exactly_the_postselected_failures(name, flip): 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, layout, 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): @@ -145,8 +145,8 @@ def test_pruned_shots_report_false_after_their_pruning_action(name): postselected = np.ones(compiled.num_outputs, dtype=bool) layout = sharing_layout(compiled, postselected) f_params = _f_params(component_layout, f_patterns(planned.num_f, 3), 4_000) - samples, _, _, _ = sample_shared_component( - compiled, layout, jnp.asarray(f_params), jax.random.key(9) + samples, _, _ = sample_shared_component( + compiled, layout, np.asarray(f_params), 9 ) rows = np.asarray(samples) ready = layout.ready @@ -163,16 +163,12 @@ def test_equal_keys_give_equal_samples(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) ) - first, next_key, _, _ = sample_shared_component( - compiled, layout, f_params, jax.random.key(13) - ) - second, _, _, _ = sample_shared_component( - compiled, layout, f_params, jax.random.key(13) - ) - third, _, _, _ = sample_shared_component(compiled, layout, f_params, next_key) + 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) @@ -256,17 +252,15 @@ def with_extent(values): # 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 = jnp.asarray( + f_params = np.asarray( _f_params(component_layout, f_patterns(planned.num_f, 7), 300) ) - key = jax.random.key(17) - reference, _, deviation, _ = sample_shared_component(compiled, full, f_params, key) + key = 17 + reference, deviation, _ = sample_shared_component(compiled, full, f_params, key) # NOTE: NaN fills every amplitude the executor must not read. monkeypatch.setattr(shared, "_POISON", True) for tape in (layout, raised): - samples, _, other, stats = sample_shared_component( - compiled, tape, f_params, key - ) + samples, other, stats = sample_shared_component(compiled, tape, f_params, key) np.testing.assert_array_equal(samples, reference) assert other == deviation assert stats["dead"] > 0 @@ -279,16 +273,14 @@ def test_parallel_kernels_give_the_serial_samples(name, coherent, monkeypatch): 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), 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, jax.random.key(19) - ) + rows, _, _ = sample_shared_component(compiled, layout, f_params, 19) samples.append(rows) np.testing.assert_array_equal(samples[1], samples[0]) @@ -298,18 +290,16 @@ def test_shots_whose_lanes_all_die_are_still_sampled(): 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 = jnp.asarray( + 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, jax.random.key(seed) + 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], jax.random.key(0) - ) + samples, _, _ = sample_shared_component(compiled, layout, f_params[:0], 0) assert samples.shape == (0, compiled.num_outputs) assert all_died diff --git a/test/unit/test_postselection.py b/test/unit/test_postselection.py index 63851611..639b0feb 100644 --- a/test/unit/test_postselection.py +++ b/test/unit/test_postselection.py @@ -253,19 +253,25 @@ def test_postselection_direct_cols_always_equal_numpy(): # ────────────────────────── 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 == [] @@ -276,13 +282,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) @@ -292,17 +293,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" @@ -315,13 +312,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 @@ -332,13 +323,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() @@ -360,21 +345,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, @@ -440,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 == [] diff --git a/test/unit/test_sampler.py b/test/unit/test_sampler.py index 525e3081..90adf9ec 100644 --- a/test/unit/test_sampler.py +++ b/test/unit/test_sampler.py @@ -516,13 +516,13 @@ 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_component + original = sampler_module.sample_shared_component - def nan_deviation(component, f_params, key): - samples, key, _ = original(component, f_params, key) - return samples, key, float("nan") + def nan_deviation(*args): + samples, _, stats = original(*args) + return samples, float("nan"), stats - monkeypatch.setattr(sampler_module, "sample_component", nan_deviation) + monkeypatch.setattr(sampler_module, "sample_shared_component", nan_deviation) with pytest.warns(UserWarning, match="not normalized"): sampler.sample(10) @@ -563,7 +563,8 @@ def test_direct_detectors_overwrite_every_column_of_their_rows(): ("active_statevector", True, "cpu", 10, 6), ("active_statevector", True, "gpu", 10, 6), ("active_statevector", False, "cpu", 10, 6), - ("active_statevector", False, "gpu", 10, 1), + # NOTE: statevector programs sample on the host whatever JAX's backend. + ("active_statevector", False, "gpu", 10, 6), ("active_statevector", False, "cpu", 10**12, 1), ], ) From eba929eed2e1992914088ae27ea07e2a19fbac78 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 19 Oct 2026 17:54:54 -0400 Subject: [PATCH 17/21] Keep calls too small for a batch per thread in one batch A call splits into a batch per thread only when every batch gets MIN_THREAD_BATCH shots; otherwise it stays one batch, which the parallel build of the executor spreads over the threads. On coherent d5r5, one batch of 16,384 shots samples 84.7k shots/s against 59k as four batches. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 20 +++++++++++++------- test/integration/test_trajectory_sharing.py | 3 ++- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index a946c158..606fb495 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -36,8 +36,10 @@ # 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, but -# into batches of no fewer shots than this, whose fixed costs would dominate. +# 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 @@ -504,8 +506,8 @@ def _resolve_batch_size( """Choose a uniform batch size for ``shots`` samples. Without a ``batch_size``, batches are as large as memory allows, but split - into at least one batch per thread (``_threads``) of no fewer than - ``MIN_THREAD_BATCH`` shots each, so that every thread has work. + 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: batch_size = self._auto_batch_size(shots) @@ -516,8 +518,12 @@ def _resolve_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())) - per_thread = min(self._threads(1), shots // MIN_THREAD_BATCH) - return ceil(shots / max(num_batches, per_thread)) + 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.""" @@ -842,7 +848,7 @@ def dispatch(task: tuple) -> None: if screened_chunks < len(starts): ready -= ready % (threads * batch_size) elif split_last and ready: - per_thread = min(threads, ready // MIN_THREAD_BATCH) + 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) diff --git a/test/integration/test_trajectory_sharing.py b/test/integration/test_trajectory_sharing.py index 5e8313f6..8e6f44b2 100644 --- a/test/integration/test_trajectory_sharing.py +++ b/test/integration/test_trajectory_sharing.py @@ -269,7 +269,8 @@ def test_shared_postselection_samples_the_survivors_without_padding(monkeypatch) @pytest.mark.parametrize("sharing", [False, True]) @pytest.mark.parametrize( - ("shots", "expected"), [(40_000, [10_000] * 4), (6_000, [6_000])] + ("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( From c27c129a9ae5cde19ea46283cdf01b603b58b52b Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 19 Oct 2026 18:20:00 -0400 Subject: [PATCH 18/21] Fix the issues the review of the JAX-free statevector path found - Components sampled overlapping random numbers: SplitMix64 counts up from its state, so seeds GOLDEN steps apart gave shifted copies of one stream. The executor now starts from the seed's mixed output; a test checks that two disconnected components are independent per shot and per batch. - With sharing the default, concurrent batches ignored the memory budget; _threads now applies the same fit with and without sharing. - The integration tests compare against strategy="cat5" instead of the default, which now shares too, and the plan-distribution tests run in complex128 as well and reject any outcome of probability 0. - Grouping keys are no longer computed without sharing, lane_bytes counts every per-shot array, and the docstrings no longer mention JAX, GPUs or the deleted execute module for the statevector path. - CHANGELOG entries for this PR. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 11 +++ src/tsim/sampler.py | 79 ++++++++++++--------- src/tsim/statevector/__init__.py | 6 +- src/tsim/statevector/shared.py | 18 +++-- test/integration/test_trajectory_sharing.py | 66 +++++++++++++---- test/unit/statevector/test_shared.py | 21 +++--- test/unit/test_sampler.py | 1 + 7 files changed, 137 insertions(+), 65 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index db3e077f..9948da8a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,17 @@ 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). Precision still follows `jax_enable_x64`. Postselected + detectors of a component prune its shots with sharing off as well. (#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/src/tsim/sampler.py b/src/tsim/sampler.py index 606fb495..77b9dff7 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -27,6 +27,7 @@ ) 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 @@ -48,8 +49,6 @@ SCREEN_SHOTS = 2**14 SCREEN_CHUNKS_PER_THREAD = 2 -_GOLDEN = 0x9E3779B97F4A7C15 - def _sample_component( component: CompiledComponent, @@ -224,7 +223,7 @@ def sample_statevector_program( zip(program.components, layouts, strict=True) ): assert isinstance(component, StatevectorComponent) - component_seed = (seed + index * _GOLDEN) % 2**64 + component_seed = (seed + index * int(GOLDEN)) % 2**64 samples, max_norm_deviation, stats = sample_shared_component( component, layout, f_params, component_seed, share ) @@ -303,7 +302,7 @@ 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"). On the CPU, "active_statevector" samples up + "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 @@ -531,7 +530,7 @@ def _peak_bytes_per_sample(self) -> int: 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. + # NOTE: plus the f-params and the output records, twice over. lane = lane_bytes(component) + 2 * num_f + 2 * self._program.num_outputs peak = max(peak, lane) continue @@ -571,8 +570,7 @@ def _threads(self, batch_size: int) -> int: if not self._has_statevector: return 1 threads = numba.get_num_threads() - if self._trajectory_sharing: - return threads + # 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(threads, fits)) @@ -667,13 +665,13 @@ def _sample_batches( tasks.append((index, seeds, subkey)) def sample_batch(task: tuple) -> jax.Array | np.ndarray: - index, seeds, key = task + index, seeds, seed_or_key = task f_params_np = self._channel_sampler.sample(batch_size, seeds) if compute_reference and index == 0: f_params_np[0] = 0 if statevector: - return self._sample_statevector(f_params_np, key, layouts) - return sample_program(self._program, jnp.asarray(f_params_np), key) + return self._sample_statevector(f_params_np, seed_or_key, layouts) + return sample_program(self._program, jnp.asarray(f_params_np), seed_or_key) batches = run_tasks(sample_batch, tasks, self._threads(batch_size)) reference: np.ndarray | None = None @@ -709,15 +707,17 @@ def _sample_batches_with_postselection( 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. + """Sample with postselection, skipping the components for direct discards. 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. For JAX the final partial batch is padded to keep the batch size - fixed (avoiding recompilation) and the padding rows are discarded; the - shared executor samples it as it is. + direct detector columns and ``False`` elsewhere; no component samples + them. Survivors are buffered until a full batch of ``batch_size`` is + ready, then sampled in one call. For JAX (the stabilizer-rank + strategies) the final partial batch is padded to keep the batch size + fixed (avoiding recompilation) and the padding rows are discarded; + statevector batches are sampled as they are, and their shots discarded + by a masked component detector read ``False`` from the action that + discards them on. Chunks of shots, and then the batches of their survivors, run on ``_threads()`` threads at once. Without a ``batch_size``, batches are as large as memory allows, and the survivors left after the last chunk are @@ -805,15 +805,17 @@ def dispatch(task: tuple) -> None: JAX batches are padded with row 0 to ``rows`` rows; statevector batches take any number of rows. """ - f_batch, shot_batch, key, rows = task + f_batch, shot_batch, seed_or_key, rows = task count = len(shot_batch) if statevector: - samples = self._sample_statevector(f_batch, key, layouts) + samples = self._sample_statevector(f_batch, seed_or_key, layouts) else: if count < rows: padding = np.repeat(f_batch[:1], rows - count, axis=0) f_batch = np.concatenate([f_batch, padding]) - samples = sample_program(self._program, jnp.asarray(f_batch), key) + samples = sample_program( + self._program, jnp.asarray(f_batch), seed_or_key + ) result[shot_batch] = np.asarray(samples)[:count] threads = self._threads(batch_size) @@ -825,7 +827,8 @@ def dispatch(task: tuple) -> None: screened_chunks = 0 # NOTE: chunks are screened until every thread has a full batch of # survivors. Batches still start every batch_size survivors and seeds and - # keys are drawn in order, so the thread count doesn't change the samples. + # seeds and keys 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 = [ @@ -855,11 +858,12 @@ def dispatch(task: tuple) -> None: batches = [] for start in range(0, ready, size): if statevector: - key = self._next_seed() + seed_or_key = self._next_seed() else: - self._key, key = jax.random.split(self._key) + self._key, seed_or_key = jax.random.split(self._key) stop = start + size - batches.append((all_f[start:stop], all_shot[start:stop], key, rows)) + batch = (all_f[start:stop], all_shot[start:stop], seed_or_key, rows) + batches.append(batch) run_tasks(dispatch, batches, threads) pending_f, pending_shot = [all_f[ready:]], [all_shot[ready:]] pending -= ready @@ -1011,9 +1015,10 @@ 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. With ``strategy="active_statevector"`` - on the CPU, up to ``numba.get_num_threads()`` batches run at once. + 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. @@ -1119,9 +1124,10 @@ 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. With ``strategy="active_statevector"`` - on the CPU, up to ``numba.get_num_threads()`` batches run at once. + 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 @@ -1143,11 +1149,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. 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/shared.py b/src/tsim/statevector/shared.py index 221a391e..ce7b5109 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -506,20 +506,26 @@ def _execute( ) = layout shots, num_words = f.shape[0], lane_branches.shape[0] num_actions = kind.shape[0] - state = np.uint64(seed) + # 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) - keys = np.zeros((shots, (key_rows.shape[0] + 31) >> 5), dtype=np.uint32) - chunks = _chunks(shots * (len(selection) + key_rows.size), parallel_work, threads) + 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: lo, hi = 0, shots _pack_range(f, selection, symbols, lo, hi) - _key_range(symbols, key_rows, keys, lo, hi) + if share: + _key_range(symbols, key_rows, keys, lo, hi) else: for c in numba.prange(chunks): lo, hi = shots * c // chunks, shots * (c + 1) // chunks _pack_range(f, selection, symbols, lo, hi) - _key_range(symbols, key_rows, keys, lo, hi) + if share: + _key_range(symbols, key_rows, keys, lo, hi) if share: order, start, count, state = _group(keys, state) else: @@ -779,7 +785,7 @@ def lane_bytes(component: StatevectorComponent) -> int: grow) and the per-shot symbols, ranges and outputs. """ amplitudes = 3 * 2**component.width * np.dtype(component.dtype).itemsize - return amplitudes + 20 * component.num_words + 88 + component.num_outputs + return amplitudes + 24 * component.num_words + 128 + component.num_outputs def sample_shared_component( diff --git a/test/integration/test_trajectory_sharing.py b/test/integration/test_trajectory_sharing.py index 8e6f44b2..cc2c33b7 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.""" import warnings from test.helpers.gen import gen_stim_circuit @@ -91,8 +91,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, @@ -108,12 +111,12 @@ def test_pruning_through_the_sampler_keeps_the_accepted_distribution(reference): @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 ) @@ -122,10 +125,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( @@ -142,9 +145,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( @@ -164,7 +165,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( @@ -173,6 +174,47 @@ 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 diff --git a/test/unit/statevector/test_shared.py b/test/unit/statevector/test_shared.py index 2b87d862..69cd124a 100644 --- a/test/unit/statevector/test_shared.py +++ b/test/unit/statevector/test_shared.py @@ -27,16 +27,16 @@ ] -def _compiled(name): - yield from _compiled_text(CIRCUITS[name]) +def _compiled(name, dtype="complex64"): + yield from _compiled_text(CIRCUITS[name], dtype) -def _compiled_text(text): +def _compiled_text(text, dtype="complex64"): 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, dtype) def _f_params(component_layout, patterns, shots_each): @@ -77,11 +77,12 @@ def _assert_counts(rows, discarded, expected): ) +@pytest.mark.parametrize("dtype", ["complex64", "complex128"]) @pytest.mark.parametrize("share", [True, False]) @pytest.mark.parametrize("name", NAMES) -def test_sampling_matches_the_plan_distribution(name, share): +def test_sampling_matches_the_plan_distribution(name, share, dtype): shots = 20_000 - for planned, component_layout, compiled in _compiled(name): + for planned, component_layout, compiled in _compiled(name, dtype): layout = sharing_layout(compiled) for f in f_patterns(planned.num_f, limit=2): f_params = _f_params(component_layout, [f], shots) @@ -89,9 +90,11 @@ def test_sampling_matches_the_plan_distribution(name, share): 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) diff --git a/test/unit/test_sampler.py b/test/unit/test_sampler.py index 90adf9ec..7382d730 100644 --- a/test/unit/test_sampler.py +++ b/test/unit/test_sampler.py @@ -566,6 +566,7 @@ def test_direct_detectors_overwrite_every_column_of_their_rows(): # 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( From e6caade99c592b5a31411275a89b2c6ec96b9551 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 19 Oct 2026 18:44:45 -0400 Subject: [PATCH 19/21] Parallelize only the explicit prange loops of numba kernels numba's parallel=True also turned array expressions such as np.zeros, np.full and slice fills into parallel loops, each a thread launch, so the parallel build of the shared executor launched threads at every measurement. Kernels now enable prange only (and get a new cache name). One SymFT d=7 postselected call of 16,384 shots samples 4.40M shots/s against 1.80M, cultivation d5 10.0M against 7.08M (65,536 shots). Co-Authored-By: Claude Opus 5.5 --- src/tsim/utils/kernels.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/src/tsim/utils/kernels.py b/src/tsim/utils/kernels.py index 272525a3..ee8ba244 100644 --- a/src/tsim/utils/kernels.py +++ b/src/tsim/utils/kernels.py @@ -36,6 +36,16 @@ def _renamed(function: Callable, suffix: str) -> Callable: 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. @@ -49,9 +59,11 @@ def jit(function: Callable, *, parallel: bool = False) -> Callable: # would keep holding the GIL. function = _renamed(function, "_nogil") try: - return numba.njit(cache=True, parallel=parallel, nogil=True)(function) + options = _prange_only() if parallel else False + return numba.njit(cache=True, parallel=options, nogil=True)(function) except RuntimeError: - return numba.njit(parallel=parallel, nogil=True)(function) + options = _prange_only() if parallel else False + return numba.njit(parallel=options, nogil=True)(function) def kernel(function: Callable) -> Callable: @@ -65,7 +77,7 @@ def kernel(function: Callable) -> Callable: 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, "_parallel"), parallel=True) + parallel = jit(_renamed(function, "_prange"), parallel=True) def run(work: int, *args): if work < PARALLEL_WORK or getattr(_TASK_THREAD, "active", False): From a4af91abfb0f041355186621ed2a262a7f8b15ed Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Mon, 19 Oct 2026 19:32:15 -0400 Subject: [PATCH 20/21] Sort the imports of the shared executor and its sampler tests Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/shared.py | 2 +- test/integration/test_statevector_sampler.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index ce7b5109..cc75bb37 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -36,6 +36,7 @@ import numba import numpy as np +import tsim.utils.kernels as kernels from tsim.statevector.program import ( MEASURE, MEASURE_DIAG, @@ -44,7 +45,6 @@ ROTATE_DIAG, StatevectorComponent, ) -import tsim.utils.kernels as kernels from tsim.utils.kernels import jit, kernel from tsim.utils.linalg import find_basis from tsim.utils.rng import GOLDEN, next_skip, splitmix64 diff --git a/test/integration/test_statevector_sampler.py b/test/integration/test_statevector_sampler.py index d5c62672..550866e8 100644 --- a/test/integration/test_statevector_sampler.py +++ b/test/integration/test_statevector_sampler.py @@ -15,8 +15,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.shared import lane_bytes from tsim.statevector.program import StatevectorComponent +from tsim.statevector.shared import lane_bytes @pytest.mark.parametrize("name", sorted(CIRCUITS)) From 7e13528089e7a485af967ccc27f9f338a21d40ee Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 20 Oct 2026 10:22:05 -0400 Subject: [PATCH 21/21] Address the stack review of the statevector sampler - Postselection prunes in the statevector executor whenever a component detector is masked, not only when a direct detector is masked too. - Reference bits are XOR-ed into the outputs before a pruned shot's later outputs are cleared, so they stay False with reference samples as well. Outputs of f-bits only prune before the first action. - The sampler picks a JAX or a statevector backend once, instead of branching in every batching method, and the statevector path no longer creates JAX keys or queries JAX devices. CompiledProgram holds NumPy arrays. - Statevector amplitudes are complex128, independent of jax_enable_x64. - StatevectorComponent holds its arrays in one Tape; the executor takes the Tape and the SharingLayout as named tuples instead of positional tuples. - run_tasks uses one thread pool and a semaphore instead of a pool per thread count, and unique_columns accepts non-integer 0/1 matrices. - Smaller cleanups: fewer direct-output attributes, private sharing stats, one NaN-poisoning point for tests, docstrings and the changelog. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 5 +- src/tsim/circuit.py | 4 +- src/tsim/compile/pipeline.py | 8 +- src/tsim/core/types.py | 8 +- src/tsim/noise/channels.py | 2 +- src/tsim/sampler.py | 742 ++++++++++--------- src/tsim/statevector/program.py | 67 +- src/tsim/statevector/shared.py | 352 ++++----- src/tsim/utils/kernels.py | 20 +- test/integration/test_sampler.py | 6 +- test/integration/test_statevector_sampler.py | 49 +- test/integration/test_trajectory_sharing.py | 62 +- test/unit/noise/test_channels.py | 8 + test/unit/statevector/test_program.py | 33 +- test/unit/statevector/test_shared.py | 47 +- test/unit/test_postselection.py | 74 +- test/unit/test_sampler.py | 1 - test/unit/utils/test_kernels.py | 22 + 18 files changed, 754 insertions(+), 756 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9948da8a..8f2ac2b9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -34,8 +34,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `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). Precision still follows `jax_enable_x64`. Postselected - detectors of a component prune its shots with sharing off as well. (#199) + 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` diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index c3066d2e..98436cae 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -873,7 +873,7 @@ 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"). On the CPU, "active_statevector" samples up + "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 @@ -912,7 +912,7 @@ 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"). On the CPU, "active_statevector" samples up + "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 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 5099e54e..f8c40e33 100644 --- a/src/tsim/noise/channels.py +++ b/src/tsim/noise/channels.py @@ -551,7 +551,7 @@ def unique_columns(matrix: np.ndarray) -> tuple[np.ndarray, np.ndarray]: column of ``matrix`` among them. """ - packed = np.ascontiguousarray(np.packbits(matrix, axis=0).T) + 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], diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 77b9dff7..9cbfbae0 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -5,7 +5,7 @@ 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 @@ -197,9 +197,9 @@ def sample_statevector_program( f_params: np.ndarray, seed: int, layouts: tuple[SharingLayout, ...], + direct_flips: np.ndarray, *, share: bool, - direct: tuple[np.ndarray, np.ndarray, np.ndarray], sharing_stats: deque[dict] | None = None, ) -> np.ndarray: """Sample all outputs of a program of statevector components, without JAX. @@ -209,8 +209,8 @@ def sample_statevector_program( 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. - direct: The program's direct f-indices, direct flips and output order. sharing_stats: If given, the executor appends its lane counts. Returns: @@ -235,16 +235,24 @@ def sample_statevector_program( 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, *direct[:2], component_bits, direct[2], out) + _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, reference, out): +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 that differs from its ``reference`` bit. + have a postselected detector of 1. """ discarded = np.zeros(f.shape[0], dtype=np.bool_) for row in numba.prange(f.shape[0]): @@ -252,7 +260,7 @@ def _write_direct_detectors(f, f_indices, flips, columns, postselect, reference, 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 != reference[j]: + if postselect[j] and bit: discarded[row] = True return discarded @@ -274,12 +282,226 @@ def _gather_outputs(f, f_indices, flips, component_bits, order, out): def _nbytes(tree: object) -> int: if isinstance(tree, StatevectorComponent): - return sum(v.nbytes for v in vars(tree).values() if isinstance(v, np.ndarray)) + 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.""" @@ -323,13 +545,7 @@ def __init__( 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) - # NOTE: statevector batches draw their seeds here instead of from the key. - self._batch_seeds = np.random.default_rng([seed, 1]) self._trajectory_sharing = trajectory_sharing - self._layouts_by_key: dict[bytes, tuple[SharingLayout, ...]] = {} - 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) @@ -343,49 +559,32 @@ def __init__( self.circuit = circuit self._num_detectors = prepared.num_detectors + self._backend: _Backend = ( + _StatevectorBackend( + self._program, seed, self._num_detectors, trajectory_sharing + ) + if strategy == "active_statevector" + else _StabilizerRankBackend(self._program, seed) + ) - prog = self._program + program = self._program # NOTE: _gather_outputs writes column output_order[k] without bounds checks. assert np.array_equal( - np.sort(np.asarray(prog.output_order)), np.arange(prog.num_outputs) + np.sort(program.output_order), np.arange(program.num_outputs) ) - 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 + 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], ) - # 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)) - ) - self._direct_global_indices = np.asarray( - prog.output_order[:n_direct], dtype=np.int32 - ) - self._direct = ( - self._direct_f_indices, - self._direct_flips, - np.asarray(prog.output_order), - ) - 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() - is_detector = self._direct_global_indices < self._num_detectors - self._direct_detector_f = self._direct_f_indices[is_detector] - self._direct_detector_flips = self._direct_flips[is_detector] - self._direct_detector_columns = self._direct_global_indices[is_detector] + 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) ] - self._has_statevector = bool(statevector) if statevector: peak = self._peak_bytes_per_sample() if peak > self._memory_budget(): @@ -397,103 +596,22 @@ def __init__( f"{STABILIZER_RANK_HINT}" ) - def _layouts( - self, - postselection_mask: np.ndarray | None, - detector_reference: np.ndarray | None = None, - ) -> tuple[SharingLayout, ...]: - """Return the sharing layout of every statevector component. - - Args: - postselection_mask: Postselected detectors, whose lanes are pruned. - detector_reference: Reference bits XOR-ed into the detectors before - the postselection test, if any. - - """ - 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._layouts_by_key: - layouts = [] - for component in self._program.components: - assert isinstance(component, StatevectorComponent) - detector = np.asarray(component.output_indices) < num_detectors - indices = np.minimum(component.output_indices, num_detectors - 1) - layouts.append( - sharing_layout( - component, - detector & mask[indices] if num_detectors else detector, - detector & flip[indices] if num_detectors else detector, - ) - ) - self._layouts_by_key[key] = tuple(layouts) - return self._layouts_by_key[key] - - def _sample_statevector( - self, f_params: np.ndarray, seed: int, layouts: tuple[SharingLayout, ...] - ) -> np.ndarray: - """Sample one batch of a statevector program.""" - return sample_statevector_program( - self._program, - f_params, - seed, - layouts, - share=self._trajectory_sharing, - direct=self._direct, - sharing_stats=self.sharing_stats, - ) - - def _next_seed(self) -> int: - """Return the next statevector batch seed.""" - return int(self._batch_seeds.integers(0, 2**64, dtype=np.uint64)) - - 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] - if self._has_statevector: - seed = self._next_seed() - return self._sample_statevector(f_ref, seed, self._layouts(None))[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, @@ -527,67 +645,36 @@ def _per_thread(self, shots: int) -> int: 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, twice over. - 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. - - Statevector programs sample on the host, so their device is the host. - """ - device = jax.devices()[0] - if device.platform == "gpu" and not self._has_statevector: - 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 the memory one batch may use: half of the free device memory.""" + 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. - Statevector programs sample one batch per numba thread, as long as the - batches fit in the memory budget; the stabilizer-rank programs sample - one batch at a time. + That is the backend's limit, as long as the batches fit in the memory + budget. """ - if not self._has_statevector: - return 1 - threads = numba.get_num_threads() # 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(threads, fits)) + 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. 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. Programs with statevector - components use at most ``STATEVECTOR_MAX_BATCH_SIZE`` lanes per batch. + leaves room for a lane per shot in one batch only. The backend may cap + the batch size further. """ peak = self._peak_bytes_per_sample() if not self._trajectory_sharing: peak *= self._threads(1) batch_size = max(1, self._memory_budget() // peak) - if any(isinstance(c, StatevectorComponent) for c in self._program.components): - batch_size = min(batch_size, STATEVECTOR_MAX_BATCH_SIZE) + if self._backend.max_batch_size is not None: + batch_size = min(batch_size, self._backend.max_batch_size) return batch_size @overload @@ -627,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: @@ -652,46 +735,25 @@ def _sample_batches( shots, batch_size, compute_reference=compute_reference ) num_batches = ceil((shots + compute_reference) / batch_size) - - statevector = self._has_statevector - layouts = self._layouts(None) if statevector else () - tasks = [] - for index in range(num_batches): - seeds = self._channel_sampler.block_seeds(batch_size) - if statevector: - tasks.append((index, seeds, self._next_seed())) - else: - self._key, subkey = jax.random.split(self._key) - tasks.append((index, seeds, subkey)) + backend = self._backend + tasks = [ + (index, self._channel_sampler.block_seeds(batch_size), backend.next_seed()) + for index in range(num_batches) + ] def sample_batch(task: tuple) -> jax.Array | np.ndarray: - index, seeds, seed_or_key = task - f_params_np = self._channel_sampler.sample(batch_size, seeds) + index, block_seeds, seed = task + f_params = self._channel_sampler.sample(batch_size, block_seeds) if compute_reference and index == 0: - f_params_np[0] = 0 - if statevector: - return self._sample_statevector(f_params_np, seed_or_key, layouts) - return sample_program(self._program, jnp.asarray(f_params_np), seed_or_key) + f_params[0] = 0 + return backend.sample(f_params, seed) batches = run_tasks(sample_batch, tasks, self._threads(batch_size)) reference: np.ndarray | None = None if compute_reference: reference = np.asarray(batches[0][0]) batches[0] = batches[0][1:] - if statevector: - result = np.concatenate(batches)[:shots] - else: - # 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] + result = backend.concatenate(batches)[:shots] if compute_reference: assert reference is not None @@ -704,119 +766,67 @@ 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]: + referenced: np.ndarray | None = None, + ) -> np.ndarray: """Sample with postselection, skipping the components for direct discards. - Shots discarded by a direct masked detector are filled with their - direct detector columns and ``False`` elsewhere; no component samples - them. Survivors are buffered until a full batch of ``batch_size`` is - ready, then sampled in one call. For JAX (the stabilizer-rank - strategies) the final partial batch is padded to keep the batch size - fixed (avoiding recompilation) and the padding rows are discarded; - statevector batches are sampled as they are, and their shots discarded - by a masked component detector read ``False`` from the action that - discards them on. - Chunks of shots, and then the batches of their survivors, run on - ``_threads()`` threads at once. Without a ``batch_size``, batches are as - large as memory allows, and the survivors left after the last chunk are - split into a batch per thread (as in ``_resolve_batch_size``); JAX pads - each to a power of two, so that few batch shapes are compiled. - """ - 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}") + 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. + + 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. + """ + _check_sizes(shots, batch_size) 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 - - 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_) + return np.empty((0, num_outputs), dtype=np.bool_) split_last = batch_size is None if batch_size is None: batch_size = self._estimate_batch_size() - - reference: np.ndarray | None = None - if compute_reference: - reference = self._compute_reference_sample() - - result = np.empty((shots, num_outputs), dtype=np.bool_) - was_discarded = np.empty(shots, dtype=np.bool_) - - columns = self._direct_detector_columns + 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] - reference_bits = ( - reference[columns] - if xor_detector_ref and reference is not None - else np.zeros(len(columns), dtype=np.bool_) - ) - - statevector = self._has_statevector - layouts = ( - self._layouts(postselection_mask, reference if xor_detector_ref else None) - if statevector - else () - ) - + 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]: - """Sample the chunk of shots at ``start``; return its survivors. - - Writes the chunk's direct detectors and which of its shots a direct - detector discards; returns the survivors' f-params and shot indices. - """ - start, seeds = task + """Write the direct detectors of the chunk at ``start``; return its survivors.""" + start, block_seeds = task stop = min(start + chunk, shots) - f_params_np = self._channel_sampler.sample(stop - start, seeds) + f_params = self._channel_sampler.sample(stop - start, block_seeds) discarded = _write_direct_detectors( (stop - start) * len(columns), - f_params_np, - self._direct_detector_f, - self._direct_detector_flips, + f_params, + f_indices, + flips, columns, postselect, - reference_bits, result[start:stop], ) - was_discarded[start:stop] = discarded survivors = np.flatnonzero(~discarded) - return f_params_np[survivors], start + survivors + return f_params[survivors], start + survivors def dispatch(task: tuple) -> None: - """Sample the shots ``shot_batch`` in one batch. - - JAX batches are padded with row 0 to ``rows`` rows; statevector - batches take any number of rows. - """ - f_batch, shot_batch, seed_or_key, rows = task - count = len(shot_batch) - if statevector: - samples = self._sample_statevector(f_batch, seed_or_key, layouts) - else: - if count < rows: - padding = np.repeat(f_batch[:1], rows - count, axis=0) - f_batch = np.concatenate([f_batch, padding]) - samples = sample_program( - self._program, jnp.asarray(f_batch), seed_or_key - ) - result[shot_batch] = np.asarray(samples)[:count] + f_batch, shot_batch, seed, rows = task + samples = backend.sample( + f_batch, seed, mask=postselection_mask, flip=flip, pad_to=rows + ) + result[shot_batch] = np.asarray(samples) threads = self._threads(batch_size) num_f = self._channel_sampler.signature_matrix.shape[1] @@ -826,9 +836,8 @@ def dispatch(task: tuple) -> None: starts = range(0, shots, chunk) screened_chunks = 0 # NOTE: chunks are screened until every thread has a full batch of - # survivors. Batches still start every batch_size survivors and seeds and - # seeds and keys are drawn in order, so the thread count doesn't change - # the samples. + # 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 = [ @@ -851,45 +860,49 @@ def dispatch(task: tuple) -> None: 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 = [] - for start in range(0, ready, size): - if statevector: - seed_or_key = self._next_seed() - else: - self._key, seed_or_key = jax.random.split(self._key) - stop = start + size - batch = (all_f[start:stop], all_shot[start:stop], seed_or_key, rows) - batches.append(batch) + batches = [ + ( + all_f[start : start + size], + all_shot[start : start + size], + backend.next_seed(), + rows, + ) + for start in range(0, ready, size) + ] run_tasks(dispatch, batches, threads) pending_f, pending_shot = [all_f[ready:]], [all_shot[ready:]] pending -= ready - - 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 + 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.""" @@ -1175,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: @@ -1225,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/program.py b/src/tsim/statevector/program.py index dacffdaa..6aea2ce9 100644 --- a/src/tsim/statevector/program.py +++ b/src/tsim/statevector/program.py @@ -11,8 +11,8 @@ from collections.abc import Iterable, Sequence from dataclasses import dataclass from math import ceil +from typing import ClassVar, NamedTuple -import jax import numpy as np from zxparam import Graph @@ -29,23 +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." ) -@dataclass(frozen=True, eq=False) -class StatevectorComponent: - """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. @@ -63,9 +55,6 @@ class StatevectorComponent: """ - output_indices: tuple[int, ...] - width: int - dtype: str f_selection: np.ndarray kind: np.ndarray x_mask: np.ndarray @@ -80,26 +69,44 @@ class StatevectorComponent: 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 @@ -138,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)) @@ -172,17 +179,13 @@ def lower(plan: Plan, layout: ComponentLayout, dtype: str) -> StatevectorCompone 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, + 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).astype(real), - sin=np.sin(angle).astype(real), + cos=np.cos(angle), + sin=np.sin(angle), pivot=pivot, extent=extent, branch=branch, @@ -191,6 +194,7 @@ def lower(plan: Plan, layout: ComponentLayout, dtype: str) -> StatevectorCompone 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( @@ -198,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. @@ -207,5 +209,4 @@ def compile_statevector_component( """ planned, component_layout = plan_component(component, output_indices, layout) - dtype = "complex128" if jax.config.read("jax_enable_x64") else "complex64" - 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 cc75bb37..50ad6e91 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -6,7 +6,7 @@ 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 @@ -25,8 +25,9 @@ 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; a pruned shot's outputs -that need a branch bit set after its pruning action are False. +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 @@ -53,39 +54,22 @@ _I_POWERS = (1, 1j, -1, -1j) -class Tape(NamedTuple): - """A component's action tape as NumPy arrays (see ``StatevectorComponent``).""" - - 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 - f_selection: np.ndarray - - class SharingLayout(NamedTuple): - """Which symbols the shots of a lane must share, and the tape they run. + """Which symbols the shots of a lane must share, and which outputs prune them. Attributes: - tape: The component's actions and outputs. 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,)``. @@ -99,12 +83,12 @@ class SharingLayout(NamedTuple): """ - tape: Tape key_rows: np.ndarray lane_dormant: np.ndarray + prune_first: np.ndarray pruned: np.ndarray prune_outputs: np.ndarray - prune_flip: np.ndarray + output_const: np.ndarray lane_branches: np.ndarray shot_dormant: np.ndarray ready: np.ndarray @@ -139,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 = Tape(*(np.array(getattr(component, field)) for field in Tape._fields)) + tape = component.tape num_words = component.num_words num_f = len(tape.f_selection) kind, branch = tape.kind, tape.branch - num_actions = len(kind) + num_outputs = component.num_outputs sign_bits = _unpack(tape.sign_row, num_words) output_bits = _unpack(tape.output_row, num_words) - num_outputs = component.num_outputs + 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 = np.asarray(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]]) @@ -179,36 +155,37 @@ def sharing_layout( key_bits, _ = find_basis(key_bits.astype(np.uint8)) key_bits = np.asarray(key_bits, dtype=bool) + # 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 - dtype = np.dtype(component.dtype) 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=dtype) + i_power = np.array([_I_POWERS[q] for q in quarter_turns], dtype=np.complex128) cos, sin = tape.cos, tape.sin return SharingLayout( - tape=tape, key_rows=_pack(key_bits, num_words).reshape(-1, num_words), lane_dormant=lane_dormant, - pruned=pruned, + prune_first=postselected & (ready < 0), + pruned=np.asarray(prune_outputs.any(axis=1)), prune_outputs=prune_outputs, - prune_flip=flipped, + 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).astype(dtype), - coupling=(-1j * sin * i_power).astype(dtype), + factor=np.stack([cos - 1j * sin, cos + 1j * sin], axis=1), + coupling=-1j * sin * i_power, ) @@ -317,7 +294,7 @@ 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 = np.complex128(amps[lane, i]) + a = amps[lane, i] if x == 0: if _parity(i & z): w1[lane] += a.real**2 + a.imag**2 @@ -325,7 +302,7 @@ def _weights_range(amps, lo, hi, span, x, z, i_power, w0, w1): w0[lane] += a.real**2 + a.imag**2 else: j = i ^ x - b = i_power * np.complex128(amps[lane, j]) + b = i_power * amps[lane, j] if _parity(j & z): b = -b plus, minus = a + b, a - b @@ -420,6 +397,42 @@ def _group(keys, state): 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``. + """ + 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).""" @@ -463,49 +476,28 @@ def _output_range( @kernel -def _execute( - f, tape, layout, width, prototype, seed, share, parallel_work, threads, poison -): +def _execute(f, tape, layout, width, seed, share, parallel_work, threads, poison): """Sample one batch of shots with shared trajectories. - ``tape`` and ``layout`` hold the arrays of ``Tape`` and ``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``. + ``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``. Returns: ``(samples, deviation, groups, peak, lanes, dead)``. """ - ( - kind, - x_mask, - z_mask, - cos, - pivot, - extent, - branch, - sign_const, - sign_row, - output_const, - output_row, - selection, - ) = tape - ( - key_rows, - lane_dormant, - pruned, - prune_outputs, - prune_flip, - lane_branches, - shot_dormant, - ready, - i_power, - factor, - coupling, - ) = layout + # 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, + ) shots, num_words = f.shape[0], lane_branches.shape[0] - num_actions = kind.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)) @@ -516,10 +508,9 @@ def _execute( key_work = key_rows.size if share else 0 chunks = _chunks(shots * (len(selection) + key_work), parallel_work, threads) if chunks == 1: - lo, hi = 0, shots - _pack_range(f, selection, symbols, lo, hi) + _pack_range(f, selection, symbols, 0, shots) if share: - _key_range(symbols, key_rows, keys, lo, hi) + _key_range(symbols, key_rows, keys, 0, shots) else: for c in numba.prange(chunks): lo, hi = shots * c // chunks, shots * (c + 1) // chunks @@ -538,81 +529,78 @@ def _execute( 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=prototype.dtype - ) + 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 - dead_symbols = np.empty((shots, num_words), dtype=np.uint32) - dead_start = np.empty(shots, dtype=np.int64) - dead_count = np.empty(shots, dtype=np.int64) - dead_action = np.empty(shots, dtype=np.int64) - size, dead, peak, region, deviation = groups, 0, groups, 0, 0.0 - if poison: - _poison(amps, size, 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), + ) + size, dead = _prune( + amps, + 1, + groups, + 0, + layout.prune_first, + output_row, + output_const, + lanes, + graves, + -1, + ) + peak, region, deviation = groups, 0, 0.0 for a in range(num_actions): if size == 0: break - shift = extent[a] + 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 - x, z = np.int64(x_mask[a]), np.int64(z_mask[a]) - if kind[a] == ROTATE_DIAG or kind[a] == ROTATE: - flip = _signs(lane_symbols, size, sign_row[a], np.int64(sign_const[a])) + 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: - lo, hi = 0, work _rotate_range( - amps, - lo, - hi, - shift, - flip, - kind[a] == ROTATE_DIAG, - x, - z, - cos[a], - factor[a], - coupling[a], + 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, - kind[a] == ROTATE_DIAG, - x, - z, - cos[a], - factor[a], - coupling[a], + amps, lo, hi, shift, flip, diagonal, x, z, cos, factor, coupling ) continue - dormant = kind[a] == MEASURE_DORMANT - if dormant and not lane_dormant[a]: + 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: - lo, hi = 0, size - _weights_range(amps, lo, hi, span, x, z, i_power[a], w0, w1) + _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[a], w0, w1) + _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. @@ -645,63 +633,41 @@ def _execute( weight[child] = w1[lane] child += 1 size = new - if poison: - _poison(amps, size, span) if not dormant: chunks = _chunks(size * span, parallel_work, threads) if chunks == 1: - lo, hi = 0, size _project_range( - amps, lo, hi, span, x, z, pivot[a], i_power[a], outcome, weight + 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[a], i_power[a], outcome, weight + amps, lo, hi, span, x, z, pivot, i_power, outcome, weight ) - bit = branch[a] + 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 pruned[a]: - kept = 0 - for lane in range(size): - dies = False - for o in range(output_row.shape[0]): - if not prune_outputs[a, o]: - continue - acc = 0 - for w in range(num_words): - acc ^= lane_symbols[lane, w] & output_row[o, w] - if _parity(acc) ^ output_const[o] ^ prune_flip[o]: - dies = True - break - if dies: - dead_symbols[dead] = lane_symbols[lane] - dead_start[dead], dead_count[dead] = ( - lane_start[lane], - lane_count[lane], - ) - dead_action[dead] = a - dead += 1 - continue - if kept != lane: - _copy_lane(amps, kept, lane, span) - lane_symbols[kept] = lane_symbols[lane] - lane_start[kept], lane_count[kept] = ( - lane_start[lane], - lane_count[lane], - ) - kept += 1 - size = kept - if poison: - _poison(amps, size, span) + if layout.pruned[a]: + tests = layout.prune_outputs[a] + size, dead = _prune( + amps, + span, + size, + dead, + tests, + output_row, + output_const, + lanes, + graves, + a, + ) # 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) @@ -725,7 +691,6 @@ def _execute( 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: - lo, hi = 0, shots _shot_range( symbols, order, @@ -735,8 +700,11 @@ def _execute( shot_dormant, coins, shot_symbols, - lo, - hi, + 0, + shots, + ) + _output_range( + shot_symbols, output_row, output_const, ready, pruned_at, samples, 0, shots ) else: for c in numba.prange(chunks): @@ -753,12 +721,6 @@ def _execute( lo, hi, ) - if chunks == 1: - lo, hi = 0, shots - _output_range( - shot_symbols, output_row, output_const, ready, pruned_at, samples, lo, hi - ) - else: for c in numba.prange(chunks): lo, hi = shots * c // chunks, shots * (c + 1) // chunks _output_range( @@ -784,7 +746,7 @@ def lane_bytes(component: StatevectorComponent) -> int: 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(component.dtype).itemsize + amplitudes = 3 * 2**component.width * np.dtype(np.complex128).itemsize return amplitudes + 24 * component.num_words + 128 + component.num_outputs @@ -810,25 +772,13 @@ def sample_shared_component( ``(B, num_outputs)`` and ``stats`` counts groups and lanes. """ - tape = layout.tape f = np.asarray(f_params) - tape_arrays = ( - *(tape.kind, tape.x_mask, tape.z_mask, tape.cos, tape.pivot, tape.extent), - *(tape.branch, tape.sign_const, tape.sign_row, tape.output_const), - *(tape.output_row, tape.f_selection), - ) - layout_arrays = ( - *(layout.key_rows, layout.lane_dormant, layout.pruned, layout.prune_outputs), - *(layout.prune_flip, layout.lane_branches, layout.shot_dormant, layout.ready), - *(layout.i_power, layout.factor, layout.coupling), - ) samples, deviation, groups, peak, lanes, dead = _execute( f.shape[0] << component.width, f, - tape_arrays, - layout_arrays, + component.tape, + layout, component.width, - np.zeros(0, dtype=component.dtype), np.uint64(seed), share, kernels.PARALLEL_WORK, diff --git a/src/tsim/utils/kernels.py b/src/tsim/utils/kernels.py index ee8ba244..574a16be 100644 --- a/src/tsim/utils/kernels.py +++ b/src/tsim/utils/kernels.py @@ -58,11 +58,10 @@ def jit(function: Callable, *, parallel: bool = False) -> Callable: # 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: - options = _prange_only() if parallel else False return numba.njit(cache=True, parallel=options, nogil=True)(function) except RuntimeError: - options = _prange_only() if parallel else False return numba.njit(parallel=options, nogil=True)(function) @@ -94,8 +93,12 @@ def _mark_task_thread() -> None: @functools.cache -def _pool(threads: int) -> ThreadPoolExecutor: - return ThreadPoolExecutor(threads, initializer=_mark_task_thread) +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: @@ -123,7 +126,12 @@ def run_tasks( tasks = list(tasks) if min(threads, len(tasks)) <= 1 or getattr(_TASK_THREAD, "active", False): return [function(task) for task in tasks] - pool = _pool(threads) - futures = [pool.submit(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/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_statevector_sampler.py b/test/integration/test_statevector_sampler.py index 550866e8..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 @@ -104,24 +103,7 @@ 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 @@ -135,29 +117,20 @@ def drifting_shared(*args): monkeypatch.setattr(sampler_module, "sample_component", drifting) monkeypatch.setattr(sampler_module, "sample_shared_component", drifting_shared) - 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) + 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 diff --git a/test/integration/test_trajectory_sharing.py b/test/integration/test_trajectory_sharing.py index cc2c33b7..99b45257 100644 --- a/test/integration/test_trajectory_sharing.py +++ b/test/integration/test_trajectory_sharing.py @@ -1,10 +1,8 @@ """trajectory_sharing=True end to end, against the stabilizer-rank strategy.""" -import warnings from test.helpers.gen import gen_stim_circuit from test.helpers.statevector import CIRCUITS, assert_counts_match -import jax import numpy as np import pytest @@ -12,9 +10,8 @@ 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 @@ -66,6 +63,12 @@ """ +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, @@ -104,7 +107,7 @@ 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) @@ -350,27 +353,48 @@ def test_direct_outputs_land_in_their_columns(name, rows): sampler = Circuit(text).compile_detector_sampler( strategy="active_statevector", seed=1 ) + program = sampler._program if name == "direct_flips": - assert sampler._direct_flips.any() - assert sampler._program.output_reindex is not None + assert program.direct_flips.any() + assert program.output_reindex is not None f_params = sampler._channel_sampler.sample(rows) - samples = sampler._sample_statevector(f_params, 2, sampler._layouts(None)) - columns = sampler._direct_global_indices - expected = (f_params[:, sampler._direct_f_indices] != 0) ^ sampler._direct_flips + 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_complex128_samples_do_not_depend_on_the_thread_count(sharing): +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(threads): + def sample(reference): sampler = circuit.compile_detector_sampler( - strategy="active_statevector", seed=6, trajectory_sharing=sharing + 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, ) - sampler._threads = lambda batch_size: threads - return sampler.sample(2_001, batch_size=250, append_observables=True) - with jax.enable_x64(True), warnings.catch_warnings(): - warnings.simplefilter("error") - np.testing.assert_array_equal(sample(3), sample(1)) + 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 6acf78e1..35708b5f 100644 --- a/test/unit/noise/test_channels.py +++ b/test/unit/noise/test_channels.py @@ -807,6 +807,14 @@ def test_matches_numpy(self, rows, cols): 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.""" 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 69cd124a..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, @@ -27,16 +28,16 @@ ] -def _compiled(name, dtype="complex64"): - yield from _compiled_text(CIRCUITS[name], dtype) +def _compiled(name): + yield from _compiled_text(CIRCUITS[name]) -def _compiled_text(text, dtype="complex64"): +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, dtype) + yield planned, component_layout, lower(planned, component_layout) def _f_params(component_layout, patterns, shots_each): @@ -77,12 +78,11 @@ def _assert_counts(rows, discarded, expected): ) -@pytest.mark.parametrize("dtype", ["complex64", "complex128"]) @pytest.mark.parametrize("share", [True, False]) @pytest.mark.parametrize("name", NAMES) -def test_sampling_matches_the_plan_distribution(name, share, dtype): +def test_sampling_matches_the_plan_distribution(name, share): shots = 20_000 - for planned, component_layout, compiled in _compiled(name, dtype): + for planned, component_layout, compiled in _compiled(name): layout = sharing_layout(compiled) for f in f_patterns(planned.num_f, limit=2): f_params = _f_params(component_layout, [f], shots) @@ -135,18 +135,20 @@ def test_pruning_discards_exactly_the_postselected_failures(name, flip, 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) - layout = 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, layout, np.asarray(f_params), 9 @@ -236,7 +238,7 @@ def coherent(): ((planned, component_layout, compiled),) = _compiled_text( _coherent_surface_code(5, 3) ) - extent = np.asarray(compiled.extent) + 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 @@ -245,10 +247,10 @@ def coherent(): 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 = layout.tape.extent + extent = compiled.tape.extent def with_extent(values): - return layout._replace(tape=layout.tape._replace(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 @@ -259,11 +261,13 @@ def with_extent(values): _f_params(component_layout, f_patterns(planned.num_f, 7), 300) ) key = 17 - reference, deviation, _ = sample_shared_component(compiled, full, f_params, key) + 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 tape in (layout, raised): - samples, other, stats = sample_shared_component(compiled, tape, f_params, key) + 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 @@ -311,8 +315,9 @@ def test_shots_whose_lanes_all_die_are_still_sampled(): def test_every_action_stays_within_its_extent(name, coherent): compiled_components = [coherent] if name == "coherent" else _compiled(name) for _, _, compiled in compiled_components: - reach = 1 << np.asarray(compiled.extent, dtype=np.int64) - support = np.asarray(compiled.x_mask) | np.asarray(compiled.z_mask) + 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(compiled.kind), (MEASURE, MEASURE_DIAG)) - assert np.all((1 << np.asarray(compiled.pivot))[measured] < reach[measured]) + 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 639b0feb..148f8e7e 100644 --- a/test/unit/test_postselection.py +++ b/test/unit/test_postselection.py @@ -67,6 +67,17 @@ 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) @@ -226,7 +237,7 @@ def test_survivors_keep_their_own_direct_outputs(shots, batch_size, reference): append_observables=True, use_detector_reference_sample=reference, ) - expected = sampler._compute_direct_outputs(drawn()) + expected = _direct_outputs(sampler, drawn()) if reference: expected[:, :3] ^= sampler._compute_reference_sample()[:3] discarded = expected[:, 0] @@ -242,7 +253,7 @@ def test_postselection_direct_cols_always_equal_numpy(): with _record_noise(sampler) as drawn: samples = sampler.sample(8, batch_size=4, postselection_mask=mask) - expected_direct = sampler._compute_direct_outputs(drawn()) + expected_direct = _direct_outputs(sampler, drawn()) direct_det = sampler._direct_detector_mask assert np.array_equal( samples[:, : sampler._num_detectors] & direct_det, @@ -363,8 +374,11 @@ def test_postselection_detector_reference_xor_before_discard_check(monkeypatch): 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 @@ -375,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 ─────────────────────────── diff --git a/test/unit/test_sampler.py b/test/unit/test_sampler.py index 7382d730..cce4f74e 100644 --- a/test/unit/test_sampler.py +++ b/test/unit/test_sampler.py @@ -537,7 +537,6 @@ def test_direct_detectors_overwrite_every_column_of_their_rows(): np.array([False, True]), np.array([3, 0]), np.array([True, False]), - np.array([False, False]), out, ) expected = np.zeros((3, 4), dtype=np.bool_) diff --git a/test/unit/utils/test_kernels.py b/test/unit/utils/test_kernels.py index abc8feb4..f612b439 100644 --- a/test/unit/utils/test_kernels.py +++ b/test/unit/utils/test_kernels.py @@ -8,6 +8,7 @@ import numpy as np import pytest +import tsim.utils.kernels as kernels from tsim.utils.kernels import PARALLEL_WORK, kernel, run_tasks @@ -111,3 +112,24 @@ def task(n): 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