diff --git a/CHANGELOG.md b/CHANGELOG.md index 95081d37..388110b9 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/tsim/statevector/metal/kernels.py b/src/tsim/statevector/metal/kernels.py index 147ca5bb..885dff02 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`` 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 diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index c92ce9ba..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) @@ -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", "