From a528e29cc1255172da923ca525ffb794c3c68a76 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 3 Nov 2026 19:44:15 -0700 Subject: [PATCH 1/3] Gather postselected shots in parallel and stop finishing discarded ones #252 made CompiledDetectorSampler.sample return only the shots that pass postselection. On the statevector path it found them with one NumPy pass over every shot's words after sampling, on one thread, and still cleared and finished discarded shots as #222 and #223 did. On B850 against main 5af2731 (before #252), postselected cells fell to 0.25x (surface d=7, 10 threads), 0.53x (cultivation d=5) and 0.17-0.31x (coherent r=1). - Each chunk now returns whether its shots passed (the screen and every component's pruning, i.e. `alive`) and compacts its survivors in place on its own thread; a numba pass joins the chunks. - A program without components screens each noise block inside the noise kernel, while it is in cache, and copies the survivors in parallel out of scratch words the sampler reuses (at most 32 MiB). - Discarded shots are no longer finished: _prune drops their lanes and clears `alive` without drawing coins or writing outputs, and the `kept`/`pruned_kept` clearing is gone. Nothing reads them any more. Samples are bit-for-bit those of main. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 6 + src/tsim/noise/channels.py | 105 +++++++++++++ src/tsim/statevector/shared.py | 90 ++--------- src/tsim/statevector/shots.py | 219 +++++++++++++++++++-------- test/helpers/statevector.py | 12 +- test/unit/statevector/test_shared.py | 22 +-- 6 files changed, 297 insertions(+), 157 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 41824cd6..8e7ba57d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -121,6 +121,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 shot's outputs when its lane dies instead of keeping graves, and both samplers pick their backend in one `_sample_outputs` method. Samples are unchanged. (#223) +- With `strategy="active_statevector"` and a postselection mask, the shots + that pass are gathered by numba kernels in the threads that sampled them, + instead of by NumPy on one thread. Discarded shots are no longer finished: + a pruned shot draws no coins, and its outputs are neither completed nor + cleared. A program without components screens each noise block while it is + in cache. Samples are unchanged. ### Deprecated - The `"pyzx"`, `"pyzx-dets"` and `"pyzx-meas"` diagram types. They still work as aliases for `"zx"`, `"zx-dets"` and `"zx-meas"` but emit a `DeprecationWarning`. diff --git a/src/tsim/noise/channels.py b/src/tsim/noise/channels.py index 4cd23516..59b3645d 100644 --- a/src/tsim/noise/channels.py +++ b/src/tsim/noise/channels.py @@ -774,6 +774,48 @@ def sample_words( *flips, ) + def sample_screened_words( + self, + out: np.ndarray, + flips: WordTables, + seeds: np.ndarray, + constant: np.ndarray, + screen: np.ndarray, + ) -> tuple[np.ndarray, np.ndarray]: + """``sample_words``, then XOR ``constant`` into each row and screen it. + + Returns whether each row has no ``screen`` bit, and the number of such + rows in each block of ``_BLOCK_SHOTS``. ``constant`` and ``screen`` cover + every word of a row. + """ + num_samples = len(out) + num_blocks = -(-num_samples // _BLOCK_SHOTS) + if len(seeds) != num_blocks: + raise ValueError( + f"{num_samples} samples need {num_blocks} block seeds, " + f"got {len(seeds)}" + ) + if not (len(constant) == len(screen) == out.shape[1]): + raise ValueError("constant and screen must cover every word of a row") + t = self._tables + keep = np.empty(num_samples, dtype=np.bool_) + kept = np.empty(num_blocks, dtype=np.int64) + _sample_screened_blocks( + num_samples * len(t.p_fire), + out, + seeds, + t.p_fire, + t.outcome_start, + t.alias_prob, + t.alias_index, + *flips, + constant, + screen, + keep, + kept, + ) + return keep, kept + class WordTables(NamedTuple): """The words each fault outcome XORs into a row. @@ -949,3 +991,66 @@ def _sample_blocks( word, mask, ) + + +@jit +def _screen_rows(rows, constant, screen, keep): + """XOR ``constant`` into each row and flag in ``keep`` the rows without a ``screen`` bit. + + ``constant`` and ``screen`` cover every word of a row. Returns the number of + rows kept. + """ + num_rows, width = rows.shape + kept = 0 + if width == 1: + # NOTE: a loop over one word would keep LLVM from vectorizing. + c, m = constant[0], screen[0] + for s in range(num_rows): + value = rows[s, 0] ^ c + rows[s, 0] = value + keep[s] = (value & m) == 0 + kept += keep[s] + return kept + for s in range(num_rows): + hit = np.uint64(0) + for w in range(width): + rows[s, w] ^= constant[w] + hit |= rows[s, w] & screen[w] + keep[s] = hit == 0 + kept += keep[s] + return kept + + +@kernel +def _sample_screened_blocks( + out, + seeds, + p_fire, + outcome_start, + alias_prob, + alias_index, + flip_start, + word, + mask, + constant, + screen, + keep, + kept, +): + """``_sample_blocks``, then each block's rows are screened by ``_screen_rows``.""" + 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, + alias_prob, + alias_index, + flip_start, + word, + mask, + ) + kept[b] = _screen_rows( + rows, constant, screen, keep[b * _BLOCK_SHOTS : (b + 1) * _BLOCK_SHOTS] + ) diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index 2d81b6fb..66097f6d 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -20,8 +20,8 @@ Pruning. After the action that sets a postselected output's last branch bit (or before the first action, for an output of f-bits only) the lanes in which it reads -1 die. Their shots get their outputs there, from the outcomes drawn so far, and -those that need a later branch bit read False. +1 die. Their shots leave ``alive`` and are not sampled further: they are +discarded, so their outputs are left unfinished. Outputs. Each shot XORs the branch part of its outputs, from its last lane, and its own fair coins: the dormant measurements that no sign and no postselected @@ -599,29 +599,6 @@ def _lane_pattern(symbols, r, branch_row, output_word, output_bit, num_words): return pattern -# NOTE: _finish_shot runs once per pruned shot, so, like _lane_key, it takes only -# the arrays it reads and not the LaneProgram. -@jit -def _finish_shot( - words, shot, pattern, hidden, out_words, coin_pattern, seed, index, slots, first -): - """Add a shot's branch parts and coins to its outputs, then clear ``hidden``. - - The coins come from the shot's draw slots ``first`` on, 64 per slot. - """ - for w in range(len(out_words)): - words[shot, out_words[w]] ^= pattern[w] - coins = np.uint64(0) - for j in range(coin_pattern.shape[0]): - if j % 64 == 0: - coins = _draw(seed, index, slots, first + j // 64) - if (coins >> np.uint64(j % 64)) & np.uint64(1): - for w in range(len(out_words)): - words[shot, out_words[w]] ^= coin_pattern[j, w] - for w in range(len(out_words)): - words[shot, out_words[w]] &= ~hidden[w] - - @jit def _prune( amps, @@ -632,67 +609,30 @@ def _prune( row, n, code_row, - words, alive, - first_shot, - seed, - slots, - program, + prune_row, tests, - action, ): - """Finish the shots of the lanes in which a tested output reads 1. + """Drop the lanes in which a tested output reads 1, and their shots. - Those shots get their outputs as they stand after ``action``, with the - outputs that need a later branch bit cleared, and leave ``live`` and - ``alive``. The other lanes move down in order, and ``code_row`` maps codes - to their new rows. Returns ``(size, n, pruned lanes)``. + Those shots leave ``live`` and ``alive``. The other lanes move down in + order, and ``code_row`` maps codes to their new rows. Returns + ``(size, n, pruned lanes)``. """ dies = np.zeros(size, dtype=np.bool_) pruned = 0 for r in range(size): for o in range(len(tests)): - if tests[o] and _row_parity(symbols, r, program.prune_row[o]): + if tests[o] and _row_parity(symbols, r, prune_row[o]): dies[r] = True pruned += 1 break if pruned == 0: return size, n, 0 - out_words = program.out_words - hidden = np.zeros(len(out_words), dtype=np.uint64) - for o in range(len(program.output_bit)): - if program.ready[o] > action: - w = program.output_word[o] - hidden[w] |= np.uint64(1) << np.uint64(program.output_bit[o]) - pattern = np.zeros((size, len(out_words)), dtype=np.uint64) - for r in range(size): - if dies[r]: - pattern[r] = _lane_pattern( - symbols, - r, - program.branch_row, - program.output_word, - program.output_bit, - len(out_words), - ) - coins = program.coin_pattern - first_coin = len(program.kind) m = 0 for i in range(n): shot, r = live[i], code_row[row[i]] if dies[r]: - _finish_shot( - words, - shot, - pattern[r], - hidden, - out_words, - coins, - seed, - first_shot + shot, - slots, - first_coin, - ) alive[shot] = False continue live[m], row[m] = shot, row[i] @@ -769,14 +709,9 @@ def run_lanes( row, n, code_row, - words, alive, - first_shot, - seed, - slots, - program, + program.prune_row, program.prune_first, - -1, ) for a in range(num_actions): @@ -911,14 +846,9 @@ def run_lanes( row, n, code_row, - words, alive, - first_shot, - seed, - slots, - program, + program.prune_row, program.prune_outputs[a], - a, ) dead += pruned diff --git a/src/tsim/statevector/shots.py b/src/tsim/statevector/shots.py index 080821dd..6885ec3d 100644 --- a/src/tsim/statevector/shots.py +++ b/src/tsim/statevector/shots.py @@ -35,6 +35,8 @@ # 2^8 its fixed costs dominate. MAX_CHUNK_SHOTS = 2**16 MIN_CHUNK_SHOTS = 2**8 +# NOTE: the most scratch memory a sampler keeps between calls. +MAX_SCRATCH_BYTES = 2**25 def memory_budget() -> int: @@ -126,22 +128,15 @@ class ShotPlan: Attributes: screen: The postselected direct outputs, as words: a shot with one of them set is discarded before the components. - kept: The outputs a discarded shot keeps (its direct detectors), as words. - pruned_kept: The outputs a shot that component ``k`` prunes keeps, as - words: ``kept`` and component ``k``'s outputs, which it clears - itself past the pruning. None for a component without postselected - outputs, which prunes no shot. programs: Each component's lane program. - rejected: The postselected detectors, as words: a shot with one of them - set was discarded, and the others are returned. + rejected: Whether any detector is postselected: if so, only the shots + that pass the screen and every component are returned. """ screen: np.ndarray - kept: np.ndarray - pruned_kept: tuple[np.ndarray | None, ...] programs: tuple[LaneProgram, ...] - rejected: np.ndarray + rejected: bool def shot_plan( @@ -152,7 +147,6 @@ def shot_plan( postselected = np.zeros(num_outputs, dtype=bool) if mask is not None: postselected[:num_detectors] = mask - is_detector = np.arange(num_outputs) < num_detectors programs = tuple( lane_program( component, key, key_start, postselected[list(component.output_indices)] @@ -161,38 +155,22 @@ def shot_plan( layout.components, layout.keys, layout.key_starts, strict=True ) ) - kept = shared.pack_bits(layout.direct & is_detector, words) - pruned_kept = [] - for component in layout.components: - outputs = list(component.output_indices) - if not postselected[outputs].any(): - pruned_kept.append(None) - continue - own = kept.copy() - for o in outputs: - own[o >> 6] |= np.uint64(1) << np.uint64(o & 63) - pruned_kept.append(own) return ShotPlan( screen=shared.pack_bits(postselected & layout.direct, words), - kept=kept, - pruned_kept=tuple(pruned_kept), programs=programs, - rejected=shared.pack_bits(postselected, words), + rejected=bool(postselected.any()), ) @jit -def _screen(words, constant, screen, kept, alive): - """XOR the constants into the outputs; a screened shot keeps only ``kept``.""" +def _screen(words, constant, screen, alive): + """XOR the constants into the outputs; a shot with a ``screen`` bit is not alive.""" for s in range(words.shape[0]): hit = np.uint64(0) for w in range(len(constant)): words[s, w] ^= constant[w] hit |= words[s, w] & screen[w] alive[s] = hit == 0 - if hit: - for w in range(len(kept)): - words[s, w] &= kept[w] @jit @@ -241,6 +219,84 @@ def _unpack(words, runs, width, out): _spread(bits, out, target) +@jit +def _compact(rows, alive, width): + """Move the first ``width`` words of each alive shot to the front, in order. + + Returns the number of alive shots. + """ + target = 0 + for s in range(rows.shape[0]): + if alive[s]: + if target != s: + for w in range(width): + rows[target, w] = rows[s, w] + target += 1 + return target + + +@kernel +def _copy_kept(words, keep, size, kept, starts, out): + """Copy the rows ``keep`` flags in each run of ``size`` rows to ``out``, in order. + + Run ``i`` keeps ``kept[i]`` rows, from row ``starts[i]`` of ``out`` on; runs + copy in parallel. + """ + width = out.shape[1] + for i in numba.prange(starts.shape[0]): + first = i * size + last = min(first + size, words.shape[0]) - 1 + while last >= first and not keep[last]: + last -= 1 + target = starts[i] + if 2 * kept[i] < last + 1 - first: + # NOTE: few kept rows: copy only them. + for s in range(first, last + 1): + if keep[s]: + for w in range(width): + out[target, w] = words[s, w] + target += 1 + elif width == 1: + # NOTE: every row up to the run's last kept one is written to + # ``target``, which only a kept row moves on, so the loop has no + # branch and never writes past the run's rows of ``out``. + for s in range(first, last + 1): + out[target, 0] = words[s, 0] + target += keep[s] + else: + for s in range(first, last + 1): + for w in range(width): + out[target, w] = words[s, w] + target += keep[s] + + +@jit +def _join(words, kept, size, width): + """Move the ``kept[i]`` rows at the front of each run of ``size`` rows together. + + Only the first ``width`` words move. Returns the number of rows kept. + """ + filled = 0 + for i in range(kept.shape[0]): + source = i * size + if source != filled: + for r in range(kept[i]): + for w in range(width): + words[filled + r, w] = words[source + r, w] + filled += kept[i] + return filled + + +def _shrink(words: np.ndarray, rows: int) -> np.ndarray: + """Return the first ``rows`` rows of ``words``, an array no view shares. + + The memory past them is given back without a copy, so that a view of the + survivors does not keep a whole call's rows alive. + """ + words.resize((rows, words.shape[1]), refcheck=False) + return words + + def bool_outputs(words: np.ndarray, columns: np.ndarray) -> np.ndarray: """Return the output bits ``columns`` of every shot as booleans.""" columns = np.asarray(columns, dtype=np.int64) @@ -318,6 +374,9 @@ def __init__( self._plans: dict[bytes, ShotPlan] = {} self._widest = max((c.width for c in self._layout.components), default=0) self._direct_tables: WordTables | None = None + # NOTE: noise words of a program without components, reused by the calls + # that screen them, since only their survivors are returned. + self._scratch: np.ndarray | None = None if not self._layout.components: f = self._layout.direct_f sets = np.flatnonzero(f >= 0) @@ -400,15 +459,7 @@ def sample( noise_seeds = self._channels.block_seeds(shots) seeds = self._next_seeds(self._seeds) - rejected = plan.rejected if plan.rejected.any() else None - def read(words: np.ndarray) -> list[np.ndarray]: - if rejected is not None: - hits = (words[:, : len(rejected)] & rejected).any(axis=1) - if hits.any(): - # NOTE: the output words alone, so that bit-packed outputs - # can be views of them. - words = words[~hits, : self._layout.output_words] return [ read_outputs(words, c, num_outputs, bit_packed=bit_packed) for c in columns @@ -464,14 +515,19 @@ def _sample_part( ) -> tuple[np.ndarray, float]: """Return the shot words of a call's shots ``start … stop-1`` and their deviation. - ``noise_seeds`` seed the call's noise blocks, so ``start`` is a whole - number of blocks, and ``seeds`` its components. + With a postselected detector, only the shots that pass postselection + are returned, and only their output words. ``noise_seeds`` seed the + call's noise blocks, so ``start`` is a whole number of blocks, and + ``seeds`` its components. """ shots = stop - start + width = self._layout.output_words + blocks = noise_seeds[start // _BLOCK_SHOTS : -(-stop // _BLOCK_SHOTS)] + if shots and not self._layout.components and plan.rejected: + return self._sample_screened(shots, blocks, plan, constant), 0.0 words = np.empty((shots, self._layout.num_words), dtype=np.uint64) if shots == 0: return words, 0.0 - blocks = noise_seeds[start // _BLOCK_SHOTS : -(-stop // _BLOCK_SHOTS)] self._channels.sample_words(words, self._layout.noise, blocks) if not self._layout.components and not (constant.any() or plan.screen.any()): # NOTE: direct outputs without constants are their noise bits. @@ -479,12 +535,55 @@ def _sample_part( threads = numba.get_num_threads() chunk, workers = self._schedule(shots, batch_size, threads) - def run(offset: int) -> list[tuple]: + def run(offset: int) -> tuple[list[tuple], int]: rows = words[offset : offset + chunk] - return self._run_chunk(rows, plan, constant, start + offset, seeds, threads) - - results = run_tasks(run, range(0, shots, chunk), workers) - return words, self._record([stats for result in results for stats in result]) + alive = np.empty(len(rows), dtype=np.bool_) + stats = self._run_chunk( + rows, alive, plan, constant, start + offset, seeds, threads + ) + # NOTE: each chunk moves its survivors to its front on its own thread. + kept = _compact(rows, alive, width) if plan.rejected else len(rows) + return stats, kept + + offsets = range(0, shots, chunk) + results = run_tasks(run, offsets, workers) + deviation = self._record([stats for result, _ in results for stats in result]) + if not plan.rejected: + return words, deviation + kept = np.array([k for _, k in results], dtype=np.int64) + filled = _join(words, kept, chunk, width) + # NOTE: the output words alone, so that bit-packed outputs can be views. + if width == words.shape[1]: + return _shrink(words, filled), deviation + return np.ascontiguousarray(words[:filled, :width]), deviation + + def _sample_screened( + self, shots: int, blocks: np.ndarray, plan: ShotPlan, constant: np.ndarray + ) -> np.ndarray: + """Return the output words of the shots that pass postselection, of ``shots``. + + For a program without components, which postselects direct detectors + only, all of them in the screen. Each noise block is screened while it is + in cache, and the survivors are copied out of words that later calls + reuse: fresh memory would cost a page fault per 4 kB. + """ + width = self._layout.num_words + assert width == self._layout.output_words + # NOTE: a call that finds the scratch taken (another thread's) allocates. + scratch, self._scratch = self._scratch, None + if scratch is None or len(scratch) < shots: + scratch = np.empty((shots, width), dtype=np.uint64) + words = scratch[:shots] + keep, kept = self._channels.sample_screened_words( + words, self._layout.noise, blocks, constant, plan.screen + ) + starts = np.zeros(len(kept), dtype=np.int64) + np.cumsum(kept[:-1], out=starts[1:]) + out = np.empty((int(kept.sum()), width), dtype=np.uint64) + _copy_kept(len(out) * width, words, keep, _BLOCK_SHOTS, kept, starts, out) + if scratch.nbytes <= MAX_SCRATCH_BYTES: + self._scratch = scratch + return out def direct_bools(self, shots: int, flip: np.ndarray | None = None) -> np.ndarray: """Return every output bit of a program without components, as booleans. @@ -507,6 +606,7 @@ def reference(self) -> tuple[np.ndarray, float]: words = np.zeros((1, self._layout.num_words), dtype=np.uint64) stats = self._run_chunk( words, + np.empty(1, dtype=np.bool_), self.plan(None), self._constant(None), 0, @@ -556,6 +656,7 @@ def _schedule( def _run_chunk( self, rows: np.ndarray, + alive: np.ndarray, plan: ShotPlan, constant: np.ndarray, first_shot: int, @@ -564,24 +665,19 @@ def _run_chunk( ) -> list[tuple]: """Screen the rows' shots and run the components on them. - A shot that a component prunes is skipped by the later components and - keeps only the outputs of ``plan.pruned_kept``, whatever the order of - the components. Alone on its thread, a chunk runs each component's - amplitude loops on ``threads`` threads. + ``alive`` receives whether each shot passed the screen and every + component's postselected outputs: a shot that one prunes is skipped by + the later ones, and its outputs are left as they are, since it is not + returned. Alone on its thread, a chunk runs each component's amplitude + loops on ``threads`` threads. """ - alive = np.empty(len(rows), dtype=np.bool_) - _screen(rows, constant, plan.screen, plan.kept, alive) + _screen(rows, constant, plan.screen, alive) stats = [] - for component, program, seed, kept in zip( - self._layout.components, - plan.programs, - seeds, - plan.pruned_kept, - strict=True, + for component, program, seed in zip( + self._layout.components, plan.programs, seeds, strict=True ): if not alive.any(): break - before = None if kept is None else alive.copy() stats.append( shared.run_lanes( len(rows) << component.width, @@ -596,11 +692,6 @@ def _run_chunk( threads, ) ) - if kept is not None: - assert before is not None - pruned = before & ~alive - if pruned.any(): - rows[pruned, : len(kept)] &= kept return stats def _record(self, stats: list[tuple]) -> float: diff --git a/test/helpers/statevector.py b/test/helpers/statevector.py index 9878bc71..7dbc3217 100644 --- a/test/helpers/statevector.py +++ b/test/helpers/statevector.py @@ -459,7 +459,9 @@ def sample_lanes( The shot words hold the component's outputs (their f-parts, constants and ``flipped`` bits) and its key, as ``tsim.statevector.shots`` lays them out, and - ``run_lanes`` adds the branch parts. Returns ``(samples, deviation, stats)``. + ``run_lanes`` adds the branch parts. Returns ``(samples, deviation, stats)``; + ``stats["alive"]`` tells which shots no pruning discarded, whose samples are + the only finished ones. """ key, program = component_program(compiled, postselected) tape = compiled.tape @@ -497,5 +499,11 @@ def sample_lanes( bits = np.unpackbits( words[:, :output_words].view(np.uint8), axis=1, bitorder="little" ) - stats = {"groups": groups, "peak lanes": peak, "lanes": size, "dead": dead} + stats = { + "groups": groups, + "peak lanes": peak, + "lanes": size, + "dead": dead, + "alive": alive, + } return bits[:, :num_outputs].astype(bool), float(deviation), stats diff --git a/test/unit/statevector/test_shared.py b/test/unit/statevector/test_shared.py index 792bda99..6f455dbb 100644 --- a/test/unit/statevector/test_shared.py +++ b/test/unit/statevector/test_shared.py @@ -123,34 +123,34 @@ def test_pruning_discards_exactly_the_postselected_failures(name, flip, share): flipped[0] = flip patterns = f_patterns(planned.num_f, limit=3) f_params = _f_params(component_layout, patterns, shots_each) - samples, _, _ = sample_lanes( + samples, _, stats = sample_lanes( compiled, f_params, 5, share, postselected, flipped ) for i, f in enumerate(patterns): rows = samples[i * shots_each : (i + 1) * shots_each] - discarded = np.asarray((rows & postselected).any(axis=1)) + discarded = ~stats["alive"][i * shots_each : (i + 1) * shots_each] expected = _postselected_distribution( plan_distribution(planned, f), postselected, flipped ) _assert_counts(rows ^ flipped, discarded, expected) +@pytest.mark.parametrize("share", [True, False]) @pytest.mark.parametrize("flip", [False, True]) @pytest.mark.parametrize("name", NAMES) -def test_pruned_shots_report_false_after_their_pruning_action(name, flip): +def test_pruning_keeps_exactly_the_shots_an_unpruned_run_accepts(name, flip, share): for planned, component_layout, compiled in _compiled(name): postselected = np.ones(compiled.num_outputs, dtype=bool) + postselected[-1] = False flipped = np.full(compiled.num_outputs, flip) f_params = _f_params(component_layout, f_patterns(planned.num_f, 3), 4_000) - rows, _, _ = sample_lanes( - compiled, f_params, 9, postselected=postselected, flipped=flipped + pruned, _, stats = sample_lanes( + compiled, f_params, 9, share, postselected, flipped ) - lanes = component_program(compiled, postselected)[1] - ready, tested = lanes.ready, lanes.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" + unpruned, _, _ = sample_lanes(compiled, f_params, 9, share, None, flipped) + accepted = ~(unpruned & postselected).any(axis=1) + np.testing.assert_array_equal(stats["alive"], accepted) + np.testing.assert_array_equal(pruned[accepted], unpruned[accepted]) @pytest.mark.parametrize("name", ["t_ghz", "midcircuit_feedforward", "injection_check"]) From a9af140a1362c9214e0dfcf29334243063c2d260 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 3 Nov 2026 19:54:19 -0700 Subject: [PATCH 2/3] Address review: scratch lock and pickling, kernel tests, stale docs - Take and return the scratch words under a lock, so that two threads on a free-threaded build cannot share them, and pickle a ShotSampler without the lock and the scratch, as main pickles. - Keep at most 16 MiB of scratch per sampler, noted in the changelog. - Warn in _shrink's docstring that no other view of the words may remain. - Unit tests for _copy_kept's branches, _join, _compact, scratch reuse across shrinking and growing calls, and pickling after a screened call. - run_lanes' docstring and the end-of-run NOTE no longer mention clearing pruned shots' outputs or _finish_shot. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 3 +- src/tsim/statevector/shared.py | 8 +-- src/tsim/statevector/shots.py | 25 +++++++- test/unit/statevector/test_shots.py | 95 +++++++++++++++++++++++++++++ 4 files changed, 123 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8e7ba57d..cd92d6b5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -126,7 +126,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 instead of by NumPy on one thread. Discarded shots are no longer finished: a pruned shot draws no coins, and its outputs are neither completed nor cleared. A program without components screens each noise block while it is - in cache. Samples are unchanged. + in cache, and reuses its noise words across calls: a sampler keeps up to + 16 MiB of them. Samples are unchanged. ### Deprecated - The `"pyzx"`, `"pyzx-dets"` and `"pyzx-meas"` diagram types. They still work as aliases for `"zx"`, `"zx-dets"` and `"zx-meas"` but emit a `DeprecationWarning`. diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index 66097f6d..f7f67444 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -662,8 +662,8 @@ def run_lanes( ``words`` holds the chunk's shot words (row ``s`` is shot ``first_shot + s``) and ``program`` is the component's ``LaneProgram``. The shots' output bits get - the component's branch parts and coins, and those that a pruning hides are - cleared. A pruned shot leaves ``alive``, so later components skip it. The + the component's branch parts and coins. A pruned shot leaves ``alive``, so + later components skip it, and its outputs are left unfinished. The amplitude loops run in up to ``4 * threads`` parallel chunks once their work reaches ``parallel_work``. @@ -862,8 +862,8 @@ def run_lanes( program.output_bit, len(out_words), ) - # NOTE: inline rather than _finish_shot: this loop runs once per live shot, - # and the call costs about 1.7x on one thread when lanes are shared. + # NOTE: inline rather than a call: this loop runs once per live shot, and a + # call costs about 1.7x on one thread when lanes are shared. for i in range(n): shot, e = live[i], code_row[row[i]] for w in range(len(out_words)): diff --git a/src/tsim/statevector/shots.py b/src/tsim/statevector/shots.py index 6885ec3d..1a0f93a9 100644 --- a/src/tsim/statevector/shots.py +++ b/src/tsim/statevector/shots.py @@ -17,6 +17,7 @@ from __future__ import annotations import sys +import threading from collections import deque from dataclasses import dataclass @@ -36,7 +37,7 @@ MAX_CHUNK_SHOTS = 2**16 MIN_CHUNK_SHOTS = 2**8 # NOTE: the most scratch memory a sampler keeps between calls. -MAX_SCRATCH_BYTES = 2**25 +MAX_SCRATCH_BYTES = 2**24 def memory_budget() -> int: @@ -292,6 +293,9 @@ def _shrink(words: np.ndarray, rows: int) -> np.ndarray: The memory past them is given back without a copy, so that a view of the survivors does not keep a whole call's rows alive. + + Warning: the resize skips numpy's reference check, so ``words`` must have no + other view left, or that view would read freed memory. """ words.resize((rows, words.shape[1]), refcheck=False) return words @@ -377,6 +381,7 @@ def __init__( # NOTE: noise words of a program without components, reused by the calls # that screen them, since only their survivors are returned. self._scratch: np.ndarray | None = None + self._scratch_lock = threading.Lock() if not self._layout.components: f = self._layout.direct_f sets = np.flatnonzero(f >= 0) @@ -557,6 +562,18 @@ def run(offset: int) -> tuple[list[tuple], int]: return _shrink(words, filled), deviation return np.ascontiguousarray(words[:filled, :width]), deviation + def __getstate__(self) -> dict: + """Pickle without the scratch words and their lock.""" + state = self.__dict__.copy() + state["_scratch"] = None + del state["_scratch_lock"] + return state + + def __setstate__(self, state: dict) -> None: + """Unpickle, with a new lock for the scratch words.""" + self.__dict__.update(state) + self._scratch_lock = threading.Lock() + def _sample_screened( self, shots: int, blocks: np.ndarray, plan: ShotPlan, constant: np.ndarray ) -> np.ndarray: @@ -570,7 +587,8 @@ def _sample_screened( width = self._layout.num_words assert width == self._layout.output_words # NOTE: a call that finds the scratch taken (another thread's) allocates. - scratch, self._scratch = self._scratch, None + with self._scratch_lock: + scratch, self._scratch = self._scratch, None if scratch is None or len(scratch) < shots: scratch = np.empty((shots, width), dtype=np.uint64) words = scratch[:shots] @@ -582,7 +600,8 @@ def _sample_screened( out = np.empty((int(kept.sum()), width), dtype=np.uint64) _copy_kept(len(out) * width, words, keep, _BLOCK_SHOTS, kept, starts, out) if scratch.nbytes <= MAX_SCRATCH_BYTES: - self._scratch = scratch + with self._scratch_lock: + self._scratch = scratch return out def direct_bools(self, shots: int, flip: np.ndarray | None = None) -> np.ndarray: diff --git a/test/unit/statevector/test_shots.py b/test/unit/statevector/test_shots.py index 5f0b0166..b347c3e2 100644 --- a/test/unit/statevector/test_shots.py +++ b/test/unit/statevector/test_shots.py @@ -248,3 +248,98 @@ def test_a_big_endian_host_is_rejected(monkeypatch): monkeypatch.setattr(sys, "byteorder", "big") with pytest.raises(NotImplementedError, match="little-endian"): Circuit(DIRECT_FLIPS).compile_detector_sampler(seed=0) + + +def _kept_rows(words, keep, size): + """The rows ``keep`` flags, and each run of ``size`` rows' count of them.""" + counts = np.add.reduceat(keep.astype(np.int64), np.arange(0, len(keep), size)) + return words[keep], counts + + +@pytest.mark.parametrize("width", [1, 3]) +@pytest.mark.parametrize("acceptance", [0.0, 0.1, 0.9, 1.0]) +def test_copy_kept_copies_the_kept_rows_in_order(width, acceptance): + rng = np.random.default_rng(int(10 * acceptance) + width) + size = 64 + words = rng.integers(0, 2**63, size=(1000, width), dtype=np.uint64) + keep = rng.random(len(words)) < acceptance + expected, kept = _kept_rows(words, keep, size) + starts = np.concatenate([[0], np.cumsum(kept)[:-1]]).astype(np.int64) + out = np.full((len(expected), width), 7, dtype=np.uint64) + for work in (0, 1 << 30): + shots_module._copy_kept(work, words, keep, size, kept, starts, out) + np.testing.assert_array_equal(out, expected) + + +def test_join_moves_each_runs_front_rows_together(): + rng = np.random.default_rng(3) + size, width = 50, 4 + words = rng.integers(0, 2**63, size=(170, width), dtype=np.uint64) + kept = np.array([0, 50, 7, 13], dtype=np.int64) + runs = [words[i * size : i * size + k].copy() for i, k in enumerate(kept)] + filled = shots_module._join(words, kept, size, width) + assert filled == kept.sum() + np.testing.assert_array_equal(words[:filled], np.vstack(runs)) + + +def test_compact_keeps_the_alive_rows_in_order(): + rng = np.random.default_rng(4) + rows = rng.integers(0, 2**63, size=(300, 3), dtype=np.uint64) + alive = rng.random(len(rows)) < 0.4 + expected = rows[alive][:, :2].copy() + kept = shots_module._compact(rows, alive, 2) + assert kept == alive.sum() + np.testing.assert_array_equal(rows[:kept, :2], expected) + + +REPETITION = """ +R 0 1 2 3 4 +X_ERROR(0.05) 0 1 2 3 4 +CNOT 0 1 1 2 2 3 3 4 +M 0 1 2 3 4 +DETECTOR rec[-1] rec[-2] +DETECTOR rec[-2] rec[-3] +DETECTOR rec[-3] rec[-4] +X 4 +M 4 +DETECTOR rec[-1] +OBSERVABLE_INCLUDE(0) rec[-2] +""" + + +@pytest.mark.parametrize("bit_packed", [False, True]) +def test_screened_calls_match_whatever_the_scratch_held(monkeypatch, bit_packed): + """Scratch reuse across shrinking and growing calls leaves the samples alone.""" + circuit = Circuit(REPETITION) + mask = np.array([True, True, False, False]) + calls = [5000, 1, 9000, 4096, 0, 12345, 3] + + def run(): + sampler = circuit.compile_detector_sampler(seed=5) + return [ + sampler.sample(n, postselection_mask=mask, bit_packed=bit_packed) + for n in calls + ] + + reused = run() + assert 0 < len(reused[0]) < calls[0] + monkeypatch.setattr(shots_module, "MAX_SCRATCH_BYTES", 0) + fresh = run() + for a, b in zip(reused, fresh, strict=True): + np.testing.assert_array_equal(a, b) + rows = reused[0] if not bit_packed else np.unpackbits(reused[0], axis=1) + assert not rows[:, :2].any() + + +def test_a_sampler_pickles_after_a_screened_call(): + import pickle + + circuit = Circuit(REPETITION) + mask = np.array([True, True, True, False]) + sampler = circuit.compile_detector_sampler(seed=6) + sampler.sample(100, postselection_mask=mask) + assert shot_sampler(sampler)._scratch is not None + copy = pickle.loads(pickle.dumps(sampler)) # noqa: S301 (a trusted round trip) + expected = sampler.sample(500, postselection_mask=mask) + assert 0 < len(expected) < 500 + np.testing.assert_array_equal(copy.sample(500, postselection_mask=mask), expected) From c99612a7befa812f1fc37ccb6a56cc2299c75ba9 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 4 Nov 2026 00:56:40 -0700 Subject: [PATCH 3/3] Address review: copy single-survivor output words, unpack test bits little-endian Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/shots.py | 3 ++- test/unit/statevector/test_shots.py | 6 +++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/tsim/statevector/shots.py b/src/tsim/statevector/shots.py index 1a0f93a9..e0b6db62 100644 --- a/src/tsim/statevector/shots.py +++ b/src/tsim/statevector/shots.py @@ -560,7 +560,8 @@ def run(offset: int) -> tuple[list[tuple], int]: # NOTE: the output words alone, so that bit-packed outputs can be views. if width == words.shape[1]: return _shrink(words, filled), deviation - return np.ascontiguousarray(words[:filled, :width]), deviation + # NOTE: a copy even of one row, which a view would not release. + return words[:filled, :width].copy(), deviation def __getstate__(self) -> dict: """Pickle without the scratch words and their lock.""" diff --git a/test/unit/statevector/test_shots.py b/test/unit/statevector/test_shots.py index b347c3e2..15cd46d3 100644 --- a/test/unit/statevector/test_shots.py +++ b/test/unit/statevector/test_shots.py @@ -327,7 +327,11 @@ def run(): fresh = run() for a, b in zip(reused, fresh, strict=True): np.testing.assert_array_equal(a, b) - rows = reused[0] if not bit_packed else np.unpackbits(reused[0], axis=1) + rows = ( + reused[0] + if not bit_packed + else np.unpackbits(reused[0], axis=1, bitorder="little") + ) assert not rows[:, :2].any()