Gather postselected shots in parallel and stop finishing discarded ones

CHANGELOG.md

@@ -121,6 +121,13 @@ 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, 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`.

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]
+        )

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]
@@ -722,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``.
 
@@ -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
 
@@ -932,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)):