From a913858530b7e93843302fae8e3918e5b03d21e5 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Thu, 12 Nov 2026 12:52:30 -0500 Subject: [PATCH 1/2] perf(metal): select the requested output columns on the GPU Detector samples that return a subset or reordering of the outputs (separate_observables, prepend_observables, or detectors only) split the rows on the host: bit-packed columns not leading the rows were unpacked to booleans, indexed and repacked (rows.select), and boolean columns were copied out of every output unpacked. On distance-3 cultivation, postselected on every detector, that took 0.48 s of a 0.53 s call of 2^26 shots on an M4 Pro. A new select_rows kernel writes each requested column list into its own array, bit-packed or as booleans, one thread per output byte, in one command buffer. All outputs in order stay the rows (or unpack_rows), and leading whole bytes a view of them; calls past 32-bit indices keep the host path. Boolean arrays stay C-contiguous. M shots/s on an M4 Pro, main -> this commit (postselected as recorded; the GPU does the work, so the numba thread count does not matter): cultivation d3 (21 outputs) bit-packed booleans separate observables 147 -> 1316 97 -> 1104 prepended observables 55 -> 1508 89 -> 1249 detectors only 1312 -> 1468 101 -> 1241 appended observables 1835 -> 1863 1424 -> 1391 end2end d5 (2745 outputs) separate observables 59 -> 291 3 -> 156 prepended observables 1 -> 221 3 -> 171 detectors only 305 -> 306 3 -> 170 appended observables 308 -> 307 214 -> 213 Without postselection the same layouts gain alike (cultivation 21-102 -> 888-1187, end2end 0-7 -> 14-65), where end2end's rows of 344 bytes bound every layout near 68 M/s. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 7 ++ src/tsim/statevector/metal/kernels.py | 27 ++++++- src/tsim/statevector/metal/sampler.py | 100 +++++++++++++++++++------ test/integration/test_metal_sampler.py | 56 ++++++++++++++ 4 files changed, 166 insertions(+), 24 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 95081d37..40707380 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -51,6 +51,13 @@ 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 every layout, instead of 0.06 to + 0.14; end-to-end cultivation d=5 samples 156 to 291 million instead of 1 to 3 + million with booleans or prepended 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 diff --git a/src/tsim/statevector/metal/kernels.py b/src/tsim/statevector/metal/kernels.py index 147ca5bb..f98cc58d 100644 --- a/src/tsim/statevector/metal/kernels.py +++ b/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`` below 2^32. +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 diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index c92ce9ba..3bfd983a 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -499,18 +499,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,12 +624,15 @@ 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, @@ -657,16 +657,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 +679,7 @@ def sample_accepted( expectations=expectations, batch_size=batch_size, bools=bools, + columns=columns, ) def _sample_call( @@ -686,12 +691,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 +741,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 @@ -754,6 +769,55 @@ 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 >= 1 << 32: + selected.append(_select(packed, chosen, num_outputs, not bools)) + 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 +876,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", " Date: Thu, 12 Nov 2026 14:09:03 -0500 Subject: [PATCH 2/2] fix(metal): keep copy kernels' 32-bit loops from wrapping Review of #294: a copy kernel's thread steps by the dispatch's threads (up to 2^20), so with an output just under the 2^32-byte guard its index could pass 2^32, wrap to where it started and never end; the GPU would stall until Metal's watchdog failed the command buffer. select_rows (bit-packed column subsets, now on the GPU) and unpack_rows (already on main) fall back to the host from MAX_COPY_ITEMS = 2^32 - 2^20 bytes on. That fallback's booleans are now C-contiguous, as the GPU's are, and a test forces it. Also: the CHANGELOG's ranges match the measurements, a test comment names the shapes whose observables end inside a byte, the host helper is imported as _host_select so that it does not read as the method _select, and sample_accepted's Returns names the list it gives with columns. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 8 +++++--- src/tsim/statevector/metal/kernels.py | 2 +- src/tsim/statevector/metal/sampler.py | 21 +++++++++++++++------ test/integration/test_metal_sampler.py | 23 +++++++++++++++++++++-- 4 files changed, 42 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 40707380..388110b9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -55,9 +55,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 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 every layout, instead of 0.06 to - 0.14; end-to-end cultivation d=5 samples 156 to 291 million instead of 1 to 3 - million with booleans or prepended observables. + 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 diff --git a/src/tsim/statevector/metal/kernels.py b/src/tsim/statevector/metal/kernels.py index f98cc58d..885dff02 100644 --- a/src/tsim/statevector/metal/kernels.py +++ b/src/tsim/statevector/metal/kernels.py @@ -176,7 +176,7 @@ // 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`` below 2^32. +// 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)]], diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index 3bfd983a..eead872d 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/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) @@ -637,7 +644,8 @@ def sample_accepted( 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 @@ -753,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) @@ -797,8 +805,9 @@ def _select( selected.append(packed[:, : count // 8]) continue width = count if bools else -(-count // 8) - if total * width >= 1 << 32: - selected.append(_select(packed, chosen, num_outputs, not bools)) + 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: diff --git a/test/integration/test_metal_sampler.py b/test/integration/test_metal_sampler.py index 5f70b811..c019c6b9 100644 --- a/test/integration/test_metal_sampler.py +++ b/test/integration/test_metal_sampler.py @@ -848,8 +848,9 @@ def test_tsim_device_cpu_keeps_the_samplers_on_the_cpu(monkeypatch): ) def test_metal_selects_every_layout_of_columns(detectors, observables): # NOTE: the observables start inside a byte (20, 21, 7, 130), end inside one - # (21, 7, 130) or span two bytes (7). Samplers of one seed draw the same - # shots; the reference is every output as booleans, in order. + # (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)))}") @@ -893,6 +894,24 @@ def sample(layout, bit_packed=False, postselection_mask=postselection_mask): 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