perf(metal): select the requested output columns on the GPU

CHANGELOG.md

@@ -51,6 +51,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   in `CS 0 0`, raise an error naming the gate instead of its decomposition.
 
 ### Changed
+- On an Apple GPU, detector samples that return a subset or reordering of the
+  outputs (`separate_observables=True`, `prepend_observables=True`, or detectors
+  only) select their columns on the GPU instead of on the host, bit-packed or as
+  booleans. On an M4 Pro, distance-3 cultivation postselected on every detector
+  samples 1.1 to 1.5 billion shots per second in these layouts, instead of 0.06
+  to 0.15 (detectors only, bit-packed, was already fast); end-to-end cultivation
+  d=5 samples 156 to 221 million instead of 1 to 3 million with booleans or
+  prepended observables, and 291 instead of 59 million with separate bit-packed
+  observables.
 - `strategy="active_statevector"` may run a component's measurements later
   where that frees no axis that a later rotation reuses and leaves more of the
   tape before its first shot-dependent action, which a GPU backend runs once for

src/tsim/statevector/metal/kernels.py

@@ -2,7 +2,8 @@
 
 ``tape`` runs the components' tapes on the queued shots; ``count_bits`` and
 ``list_bits`` list the accepted shots from the bitmap; ``gather_rows``,
-``unpack_rows`` and ``gather_float`` copy out rows, booleans and expectations;
+``unpack_rows``, ``select_rows`` and ``gather_float`` copy out rows, booleans,
+chosen outputs and expectations;
 ``reset`` and ``advance`` keep a call's counters on the GPU.
 """
 
@@ -172,6 +173,28 @@
     }}
 }}
 
+// Outputs ``columns[k]`` (``k < count``) of the first ``n`` packed rows of ``{NB}``
+// bytes, into rows of ``width`` bytes (``args = (n, count, width, bools)``): one
+// byte per output with ``bools``, else bit-packed with zero bits past them. A
+// thread writes one byte; the host keeps ``n * width`` at most ``MAX_COPY_ITEMS``.
+kernel void select_rows(
+    device const uchar* packed [[buffer(0)]], device uchar* dst [[buffer(1)]],
+    device const uint* columns [[buffer(2)]], constant uint4& args [[buffer(3)]],
+    uint gid [[thread_position_in_grid]], uint grid [[threads_per_grid]])
+{{
+    const uint width = args.z, per = args.w ? 1 : 8;
+    for (uint i = gid; i < args.x * width; i += grid) {{
+        uint r = i / width, k = (i - r * width) * per;
+        device const uchar* row = packed + (ulong)r * {NB};
+        uint v = 0;
+        for (uint b = 0; b < per && k + b < args.y; ++b) {{
+            uint c = columns[k + b];
+            v |= ((row[c >> 3] >> (c & 7)) & 1u) << b;
+        }}
+        dst[i] = uchar(v);
+    }}
+}}
+
 // Observable expectations of rows ``index[j]`` (or of row ``j``, without
 // ``has_index``) for ``j < *n``, densely, after the ``*cursor`` rows of the
 // call's earlier batches.
@@ -227,5 +250,5 @@ def source(
 def names(shared_words: bool) -> list[str]:
     """Return the kernels of a source; ``noise_kept`` needs ``shared_words``."""
     found = ["noise", "tape", "count_bits", "list_bits", "gather_rows"]
-    found += ["gather_float", "reset", "advance", "unpack_rows"]
+    found += ["gather_float", "reset", "advance", "unpack_rows", "select_rows"]
     return [*found, "noise_kept"] if shared_words else found

src/tsim/statevector/metal/sampler.py

@@ -39,7 +39,7 @@
 from tsim.noise.channels import ChannelSampler
 from tsim.statevector.gpu import FORKED, tape
 from tsim.statevector.gpu.noise import NoiseGroups, screen_pass_rate, split
-from tsim.statevector.gpu.rows import select as _select
+from tsim.statevector.gpu.rows import select as _host_select
 from tsim.statevector.metal import kernels, runtime
 from tsim.statevector.metal.codegen import METAL
 from tsim.statevector.metal.kernels import SCAN_THREADS
@@ -61,6 +61,13 @@
 MAX_COPY_THREADS = 1 << 20
 """Most threads of a copy kernel's dispatch; each thread loops over its share."""
 
+MAX_COPY_ITEMS = (1 << 32) - MAX_COPY_THREADS
+"""Most items (bytes) of a copy kernel's dispatch, which runs a 32-bit loop.
+
+NOTE: a thread steps by the dispatch's threads, so an index must stay below 2^32
+by that much; past it, it would wrap and the loop would not end.
+"""
+
 # NOTE: the words of ``_Slot.counts``.
 _SURVIVORS, _ACCEPTED, _DEVIATION, _CURSOR, _CAPACITY = range(5)
 
@@ -499,18 +506,15 @@ def sample(
         """
 
         def on_gpu():
-            rows, _, deviation = self.sample_accepted(
+            outputs, _, deviation = self.sample_accepted(
                 shots,
                 mask=mask,
                 flip=flip,
                 batch_size=batch_size,
                 bools=not bit_packed,
+                columns=columns,
             )
-            if not bit_packed:
-                return [_columns(rows, c) for c in columns], deviation
-            num_outputs = self.layout.num_outputs
-            selected = [_select(rows, c, num_outputs, bit_packed) for c in columns]
-            return selected, deviation
+            return outputs, deviation
 
         return self._or_cpu(
             on_gpu,
@@ -627,17 +631,21 @@ def sample_accepted(
         expectations: bool = False,
         batch_size: int | None = None,
         bools: bool = False,
-    ) -> tuple[np.ndarray, np.ndarray | None, float]:
+        columns: list[np.ndarray] | None = None,
+    ) -> tuple[Any, np.ndarray | None, float]:
         """Sample ``shots`` shots; return only those that pass postselection.
 
         A shot passes when none of the ``mask``'s detectors is set (every shot
         passes without a mask). Rows keep their shots' order and hold every output;
-        with ``bools``, one boolean per output, unpacked on the GPU.
+        with ``bools``, one boolean per output, unpacked on the GPU. With
+        ``columns``, the outputs are a list of arrays instead, one per entry: its
+        outputs, in its order, selected on the GPU.
 
         Returns:
             ``(outputs, expectations, deviation)``: the passing shots' outputs,
             bit-packed as ``np.packbits(..., bitorder="little")`` would (``uint8``,
-            in shared memory), with ``expectations`` each observable's
+            in shared memory), as booleans, or with ``columns`` a list of such
+            arrays; with ``expectations`` each observable's
             expectation of ``(-1)^observable`` given the shot's history before
             the measurement that settles it, or its sampled ``±1`` when a
             postselected output of its component is tested at or after that
@@ -657,16 +665,20 @@ def sample_accepted(
                     expectations=expectations,
                     batch_size=batch_size,
                     bools=bools,
+                    columns=columns,
                 )
                 for start in range(0, shots, MAX_CALL_SHOTS)
             ]
             found = [p[1] for p in parts if p[1] is not None]
             values = np.concatenate(found) if expectations else None
-            return (
-                np.concatenate([p[0] for p in parts]),
-                values,
-                max(p[2] for p in parts),
-            )
+            if columns is None:
+                outputs = np.concatenate([p[0] for p in parts])
+            else:
+                outputs = [
+                    np.concatenate(cut)
+                    for cut in zip(*(p[0] for p in parts), strict=True)
+                ]
+            return outputs, values, max(p[2] for p in parts)
         with self._lock:
             return self._sample_call(
                 shots,
@@ -675,6 +687,7 @@ def sample_accepted(
                 expectations=expectations,
                 batch_size=batch_size,
                 bools=bools,
+                columns=columns,
             )
 
     def _sample_call(
@@ -686,12 +699,19 @@ def _sample_call(
         expectations: bool = False,
         batch_size: int | None = None,
         bools: bool = False,
-    ) -> tuple[np.ndarray, np.ndarray | None, float]:
+        columns: list[np.ndarray] | None = None,
+    ) -> tuple[Any, np.ndarray | None, float]:
         """One call of ``sample_accepted``, of at most ``MAX_CALL_SHOTS`` shots."""
         module = self._module(mask)
         nb, no = -(-self.layout.num_outputs // 8), len(self._observables)
         if not shots:
             values = np.zeros((0, no), np.float32) if expectations else None
+            if columns is not None:
+                if bools:
+                    empty = [np.zeros((0, len(c)), bool) for c in columns]
+                else:
+                    empty = [np.zeros((0, -(-len(c) // 8)), np.uint8) for c in columns]
+                return empty, values, 0.0
             if bools:
                 return np.zeros((0, self.layout.num_outputs), bool), values, 0.0
             return np.zeros((0, nb), np.uint8), values, 0.0
@@ -729,6 +749,9 @@ def _sample_call(
             if not self.layout.components:
                 expected[:] = 0
             self._direct_expectations(packed, expected)
+        if columns is not None:
+            selected = self._select(module, rows, total, columns, bools)
+            return selected, expected, deviation
         if bools:
             return self._unpack(module, rows, total), expected, deviation
         return packed, expected, deviation
@@ -738,9 +761,9 @@ def _unpack(self, module: _Module, rows: runtime.Buffer, total: int) -> np.ndarr
         num_outputs = self.layout.num_outputs
         nb = -(-num_outputs // 8)
         groups = -(-num_outputs // 4)
-        if total * groups >= 1 << 32:
+        if total * groups > MAX_COPY_ITEMS:
             packed = rows.array[: total * nb].reshape(total, nb)
-            return _select(packed, np.arange(num_outputs), num_outputs, False)
+            return _host_select(packed, np.arange(num_outputs), num_outputs, False)
         out = runtime.buffer(total * num_outputs, np.uint8, recycle=True)
         if total and num_outputs:
             threads = min(total * groups, MAX_COPY_THREADS)
@@ -754,6 +777,56 @@ def _unpack(self, module: _Module, rows: runtime.Buffer, total: int) -> np.ndarr
             self._wait_all([encoder.commit()])
         return out.array[: total * num_outputs].reshape(total, num_outputs).view(bool)
 
+    def _select(
+        self,
+        module: _Module,
+        rows: runtime.Buffer,
+        total: int,
+        columns: list[np.ndarray],
+        bools: bool,
+    ) -> list[np.ndarray]:
+        """Return the outputs ``columns`` of the first ``total`` ``rows``, an array each.
+
+        Bit-packed, or one boolean per output with ``bools``. All outputs in order
+        are the rows (or ``_unpack``), leading whole bytes a view of them, and
+        other columns are selected on the GPU.
+        """
+        num_outputs = self.layout.num_outputs
+        nb = -(-num_outputs // 8)
+        packed = rows.array[: total * nb].reshape(total, nb)
+        selected, encoder, dispatched = [], runtime.Encoder(), False
+        for chosen in columns:
+            count = len(chosen)
+            leading = np.array_equal(chosen, np.arange(count))
+            if leading and count == num_outputs:
+                selected.append(self._unpack(module, rows, total) if bools else packed)
+                continue
+            if leading and not bools and count % 8 == 0:
+                selected.append(packed[:, : count // 8])
+                continue
+            width = count if bools else -(-count // 8)
+            if total * width > MAX_COPY_ITEMS:
+                found = _host_select(packed, chosen, num_outputs, not bools)
+                selected.append(np.ascontiguousarray(found))
+                continue
+            out = runtime.buffer(total * width, np.uint8, recycle=True)
+            if total * width:
+                threads = min(total * width, MAX_COPY_THREADS)
+                arguments = np.array([total, count, width, int(bools)], np.uint32)
+                encoder.dispatch(
+                    module.library.pipeline("select_rows"),
+                    [rows, out, runtime.upload(chosen.astype(np.uint32)), arguments],
+                    -(-threads // COPY_THREADS),
+                    COPY_THREADS,
+                )
+                dispatched = True
+            array = out.array[: total * width].reshape(total, width)
+            selected.append(array.view(bool) if bools else array)
+        commands = encoder.commit()
+        if dispatched:
+            self._wait_all([commands])
+        return selected
+
     def _direct_expectations(self, rows: np.ndarray, values: np.ndarray) -> None:
         """Set the expectations of observables outside the components: ``±1``."""
         layout = self.layout
@@ -812,14 +885,6 @@ def _restore(
     return shots
 
 
-def _columns(rows: np.ndarray, columns: np.ndarray) -> np.ndarray:
-    """Return the outputs ``columns`` of boolean rows of every output, C-contiguous."""
-    count = len(columns)
-    if count == rows.shape[1] and np.array_equal(columns, np.arange(count)):
-        return rows
-    return np.ascontiguousarray(rows[:, columns])
-
-
 _NOISE_ARGS = np.dtype(
     [
         ("shots", "<i8"),

test/integration/test_metal_sampler.py

@@ -246,6 +246,12 @@ def test_metal_samples_no_shots():
     sampler = _gpu(circuit.compile_detector_sampler(seed=1))
     assert sampler.sample(0, append_observables=True).shape == (0, 4)
     assert sampler.sample(0, bit_packed=True).shape == (0, 1)
+    for bit_packed in (False, True):
+        detectors, observables = sampler.sample(
+            0, separate_observables=True, bit_packed=bit_packed
+        )
+        assert detectors.shape == (0, 3 if not bit_packed else 1)
+        assert observables.shape == (0, 1)
 
 
 # NOTE: 150 noisy detectors, so the rows span three words, and two observables.
@@ -837,6 +843,75 @@ def test_tsim_device_cpu_keeps_the_samplers_on_the_cpu(monkeypatch):
     assert type(sampler._shots) is ShotSampler
 
 
+@pytest.mark.parametrize(
+    ("detectors", "observables"), [(20, 1), (21, 3), (7, 9), (16, 8), (130, 5)]
+)
+def test_metal_selects_every_layout_of_columns(detectors, observables):
+    # NOTE: the observables start inside a byte (20, 21, 7, 130), end inside one
+    # (20, 130) or span two bytes (7); 16 + 8 lies on byte boundaries. Samplers of
+    # one seed draw the same shots; the reference is every output as booleans,
+    # in order.
+    n = detectors + observables
+    lines = [f"X_ERROR(0.3) {' '.join(map(str, range(n)))}", "H 0", "T 0", "H 0"]
+    lines.append(f"M {' '.join(map(str, range(n)))}")
+    lines += [f"DETECTOR rec[-{n - i}]" for i in range(detectors)]
+    lines += [
+        f"OBSERVABLE_INCLUDE({j}) rec[-{observables - j}]" for j in range(observables)
+    ]
+    circuit = Circuit("\n".join(lines))
+    mask = np.zeros(detectors, dtype=bool)
+    mask[0] = True
+    for postselection_mask in (None, mask):
+
+        def sample(layout, bit_packed=False, postselection_mask=postselection_mask):
+            sampler = _gpu(circuit.compile_detector_sampler(seed=25))
+            return sampler.sample(
+                20_000,
+                separate_observables=layout == "separate",
+                prepend_observables=layout == "prepend",
+                append_observables=layout == "append",
+                bit_packed=bit_packed,
+                postselection_mask=postselection_mask,
+            )
+
+        every = sample("append")
+        assert isinstance(every, np.ndarray)
+        found, events = every[:, :detectors], every[:, detectors:]
+        layouts = {
+            "separate": (found, events),
+            "prepend": (np.hstack([events, found]),),
+            "detectors": (found,),
+        }
+        for layout, expected in layouts.items():
+            for bit_packed in (True, False):
+                got = sample(layout, bit_packed)
+                got = got if isinstance(got, tuple) else (got,)
+                for array, bools in zip(got, expected, strict=True):
+                    if bit_packed:
+                        bools = np.packbits(bools, axis=1, bitorder="little")
+                    else:
+                        assert array.dtype == np.bool_ and array.flags.c_contiguous
+                    np.testing.assert_array_equal(array, bools)
+
+
+def test_outputs_too_large_for_a_copy_kernel_are_selected_on_the_host(monkeypatch):
+    # NOTE: past MAX_COPY_ITEMS bytes a copy kernel's 32-bit loop would wrap.
+    circuit = Circuit(NOISY_DETECTORS)
+
+    def sample(**kwargs):
+        sampler = _gpu(circuit.compile_detector_sampler(seed=26))
+        return sampler.sample(5_000, separate_observables=True, **kwargs)
+
+    expected = {
+        bit_packed: sample(bit_packed=bit_packed) for bit_packed in (True, False)
+    }
+    monkeypatch.setattr(metal_sampler, "MAX_COPY_ITEMS", 0)
+    for bit_packed, arrays in expected.items():
+        for got, want in zip(sample(bit_packed=bit_packed), arrays, strict=True):
+            assert got.flags.c_contiguous
+            np.testing.assert_array_equal(got, want)
+
+
 @pytest.mark.parametrize("text", [GHZ_WITH_DIRECT_DETECTOR, NOISY_DETECTORS])
 def test_metal_unpacks_booleans_as_it_packs_rows(text):
     # NOTE: samplers of one seed draw the same shots; 4 and 152 outputs, neither a