From ea4e9e307c2e3a95e94d6fbe6d156732bca26bf5 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 10 Nov 2026 18:37:45 +0100 Subject: [PATCH 01/38] feat(statevector): add the host code that GPU backends share The CUDA and Metal backends of the statevector executor share their host side: `statevector/gpu/` holds the noise kernels' tables (`noise`), a tape generator that writes a component's tape in either language through a small `Dialect` (`tape`), and the packing of returned rows (`rows`). Two changes to the statevector plans come with it. `shared.common_prefix` returns the leading actions of a tape that every shot runs alike, which a GPU backend runs once. `plan_component` also plans each component with its measurements deferred (`schedule(defer=True)`) and keeps that plan when it is no wider and leaves less shot-dependent work. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 5 + src/tsim/statevector/gpu/__init__.py | 8 + src/tsim/statevector/gpu/noise.py | 289 ++++++++++ src/tsim/statevector/gpu/rows.py | 52 ++ src/tsim/statevector/gpu/tape.py | 737 +++++++++++++++++++++++++ src/tsim/statevector/program.py | 30 +- src/tsim/statevector/schedule.py | 18 +- src/tsim/statevector/shared.py | 34 ++ test/unit/statevector/test_gpu.py | 262 +++++++++ test/unit/statevector/test_program.py | 22 + test/unit/statevector/test_schedule.py | 37 +- test/unit/statevector/test_shared.py | 45 ++ 12 files changed, 1528 insertions(+), 11 deletions(-) create mode 100644 src/tsim/statevector/gpu/__init__.py create mode 100644 src/tsim/statevector/gpu/noise.py create mode 100644 src/tsim/statevector/gpu/rows.py create mode 100644 src/tsim/statevector/gpu/tape.py create mode 100644 test/unit/statevector/test_gpu.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4499852b..97d6a672 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -42,6 +42,11 @@ 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 +- `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 + all shots. The samples follow the same distribution, but some circuits' seeds + give other samples than in 1.0. - `ChannelSampler.signature_matrix` is now a `tsim.utils.linalg.SparseGF2Matrix` instead of a dense uint8 array; `np.asarray` gives the dense matrix. diff --git a/src/tsim/statevector/gpu/__init__.py b/src/tsim/statevector/gpu/__init__.py new file mode 100644 index 00000000..edf5a6c6 --- /dev/null +++ b/src/tsim/statevector/gpu/__init__.py @@ -0,0 +1,8 @@ +"""What the statevector executor's GPU backends share. + +A GPU backend runs the noise sampler and every component's tape as kernels of +its own language. The host side is common to them: the noise tables that the +noise kernels read (``noise``), the tape generator, which writes a component's +tape in any of their languages (``tape``), and the packing of the rows that +the kernels return (``rows``). +""" diff --git a/src/tsim/statevector/gpu/noise.py b/src/tsim/statevector/gpu/noise.py new file mode 100644 index 00000000..3c78a3c5 --- /dev/null +++ b/src/tsim/statevector/gpu/noise.py @@ -0,0 +1,289 @@ +"""The host tables of the GPU noise kernels. + +A noise kernel's block owns ``S`` consecutive shots and keeps their words in +fast memory (or, rows too wide for it, in device memory). The channels with +equal fire probability and equal numbers of outcomes form a group, and within +a block the group's channels make ``S * len(group)`` Bernoulli trials, channel +major (trial ``i`` is channel ``group[i / S]`` on shot ``i % S``). The host cuts +the groups into *segments* of about ``FAULTS_PER_SEGMENT`` expected faults. In +phase 1 the threads walk the segments with geometric skips and record each +fault as an *event*: its row, the record of the alias slot it draws and the +draw's fraction; phase 2 applies the events: it keeps the slot's outcome or +takes its alias and XORs the outcome's flips into the words. + +Each launch lays out its channels' outcomes in group order, so a group's +channel ``c`` has its records at ``first + c * outcomes`` and a fault reads one +aligned 32-byte record (and its alias's, in the same channel): + + alias: bits 0-24 keep threshold (probability * 2^24), bits 25-31 alias slot + words: the first two flipped words (16 bits each) + extra: bits 0-7 number of flips, bits 8-31 start of the rest in the overflow + masks: the first two flip masks + +The flips past an outcome's first two are 16-byte ``(word, mask)`` pairs, in +order. The samples follow ``ChannelSampler``'s distribution, not its streams: +every block draws from its own counter-based stream. + +Under postselection the channels can be drawn in two passes. Those with an +outcome that flips a bit the screen, the tapes or the postselection read go +first; the others (``deferred_channels``) cannot change whether a shot is kept, +so they can be drawn after the tapes, on the kept shots only. Channels are +independent and their flips commute, so the kept shots' distribution is +unchanged. +""" + +from __future__ import annotations + +import itertools +from dataclasses import dataclass, replace + +import numpy as np + +from tsim.noise.channels import ChannelSampler, WordTables + +FAULTS_PER_SEGMENT = 4.0 +# NOTE: rows take 8 bits of an event. +MAX_BLOCK_SHOTS = 256 + +_RECORD = np.dtype( + [ + ("alias", " NoiseGroups: + """Return the groups of the channels ``select`` (a boolean mask).""" + return _groups(self, select) + + +def _groups(tables: NoiseTables, select: np.ndarray) -> NoiseGroups: + p_fire, outcome_start, S = tables.p_fire, tables.outcome_start, tables.block_shots + outcomes = np.diff(outcome_start) + faults = float(p_fire[select].sum()) + max_events = min( + max(256, int(2 * faults * S + 256)), + (tables.shared_bytes - 8 * tables.num_words * S * tables.shared_words) // 8, + ) + # NOTE: channels of equal fire probability and outcomes form a group, so a + # group's channel c has its records at first + c * outcomes. + chosen = np.flatnonzero(select) + order = chosen[np.lexsort((outcomes[chosen], p_fire[chosen]))] + p_sorted, n_sorted = p_fire[order], outcomes[order] + change = (np.diff(p_sorted, prepend=-1.0) != 0) | ( + np.diff(n_sorted, prepend=-1) != 0 + ) + starts = np.flatnonzero(change) + group_start = np.append(starts, len(order)) + record_start = np.concatenate([[0], np.cumsum(n_sorted)]) + first, per, begin, end, log_keep = [], [], [], [], [] + for g in range(len(starts)): + lo, hi = int(group_start[g]), int(group_start[g + 1]) + p = float(p_sorted[lo]) + lk = -np.inf if p >= 1.0 else float(np.log1p(-p)) + trials = S * (hi - lo) + # NOTE: float32 skips are exact below 2^24 trials only. + pieces = max(-(-trials >> 24), round(trials * min(p, 1.0) / FAULTS_PER_SEGMENT)) + cuts = np.linspace(0, trials, pieces + 1).astype(np.int64) + for a, b in itertools.pairwise(cuts): + if b > a: + first.append(int(record_start[lo])) + per.append(int(n_sorted[lo])) + begin.append(a) + end.append(b) + log_keep.append(lk) + if max(end, default=0) >= 1 << 31: + raise ValueError("too many noise channels for the GPU noise sampler") + # NOTE: the segments with the most expected faults first, so the threads that + # take a second segment take short ones. + expected = (np.array(end) - np.array(begin)) * -np.expm1(np.array(log_keep)) + by_work = np.argsort(-expected, kind="stable") + # NOTE: each chosen channel's outcomes, in group order. + within = np.arange(int(record_start[-1])) - np.repeat(record_start[:-1], n_sorted) + records = tables.records[np.repeat(outcome_start[order], n_sorted) + within] + return NoiseGroups( + max_events=max_events, + segments=( + np.asarray(first, np.int32)[by_work], + np.asarray(per, np.int32)[by_work], + np.asarray(begin, np.int32)[by_work], + np.asarray(end, np.int32)[by_work], + np.asarray(log_keep, np.float32)[by_work], + ), + records=records, + faults=faults, + ) + + +def _outcome_hits(flips: WordTables, read: np.ndarray) -> np.ndarray: + """Return which outcomes flip a bit of ``read`` (as shot words).""" + hit = (read[flips.word] & flips.mask) != 0 + # NOTE: hits per outcome as differences of running counts. + per_flip = np.concatenate([[0], np.cumsum(hit)]) + return per_flip[flips.start[1:]] - per_flip[flips.start[:-1]] > 0 + + +def deferred_channels( + channels: ChannelSampler, flips: WordTables, read: np.ndarray +) -> np.ndarray: + """Return which channels flip none of the bits ``read`` (as shot words).""" + hits = np.concatenate([[0], np.cumsum(_outcome_hits(flips, read))]) + start = channels._tables.outcome_start + return hits[start[1:]] == hits[start[:-1]] + + +def pass_rate(channels: ChannelSampler, flips: WordTables, read: np.ndarray) -> float: + """Return the share of shots whose faults flip no bit of ``read``, to first order. + + Each channel's chance to flip one is its fire probability times its hitting + outcomes' probabilities (from its alias table); flips that cancel are not + counted. + """ + fire = channels._tables + n = np.diff(fire.outcome_start) + per = np.repeat(n, n).astype(np.float64) + q = fire.alias_prob / per + np.add.at(q, fire.alias_index, (1 - fire.alias_prob) / per) + hits = np.concatenate([[0.0], np.cumsum(q * _outcome_hits(flips, read))]) + start = fire.outcome_start + p = np.clip(fire.p_fire * (hits[start[1:]] - hits[start[:-1]]), 0.0, 1.0) + return float(np.exp(np.log1p(-np.minimum(p, 1 - 1e-12)).sum())) + + +def noise_tables( + channels: ChannelSampler, + flips: WordTables, + num_words: int, + shared_bytes: int, + shared_words: bool = True, +) -> NoiseTables: + """Return the noise kernels' tables for rows of ``num_words`` shot words. + + A block's words and events share ``shared_bytes`` of fast memory, or + without ``shared_words`` the words are in device memory and the events + have all of it. + + Raises: + ValueError: If the rows have no words, their words do not fit, or the + noise exceeds the kernels' limits. + + """ + fire = channels._tables + if len(fire.p_fire) >= 1 << 24: + raise ValueError("the GPU noise sampler supports fewer than 2^24 channels") + if num_words >= 1 << 16: + raise ValueError("the GPU noise sampler supports fewer than 2^16 shot words") + if num_words == 0: + raise ValueError("the GPU noise sampler needs at least one shot word") + block_shots = max( + 32, min(MAX_BLOCK_SHOTS, (shared_bytes // 2 // (8 * num_words)) // 32 * 32) + ) + if not shared_words: + block_shots = MAX_BLOCK_SHOTS + elif 8 * num_words * block_shots > shared_bytes - 8 * 256: + raise ValueError( + f"{num_words} shot words do not fit the noise kernel's shared memory" + ) + outcome_start = fire.outcome_start + if len(fire.alias_prob) >= 1 << 24: + raise ValueError("the GPU noise sampler supports fewer than 2^24 outcomes") + records = np.zeros(len(fire.alias_prob), dtype=_RECORD) + overflow: list[tuple[int, int]] = [] + for c in range(len(fire.p_fire)): + lo, hi = int(outcome_start[c]), int(outcome_start[c + 1]) + if hi - lo > 127: + raise ValueError( + "the GPU noise sampler supports channels of up to 127 outcomes" + ) + for k in range(lo, hi): + keep = ( + 1 << 24 + if hi - lo == 1 + else min(round(float(fire.alias_prob[k]) * (1 << 24)), 1 << 24) + ) + records[k]["alias"] = keep | ((int(fire.alias_index[k]) - lo) << 25) + j0, j1 = int(flips.start[k]), int(flips.start[k + 1]) + words = flips.word[j0:j1].tolist() + masks = flips.mask[j0:j1].tolist() + records[k]["words"] = (words[0] if words else 0) | ( + (words[1] if len(words) > 1 else 0) << 16 + ) + records[k]["m0"] = masks[0] if masks else 0 + records[k]["m1"] = masks[1] if len(masks) > 1 else 0 + records[k]["extra"] = len(words) | (len(overflow) << 8) + overflow.extend(zip(words[2:], masks[2:], strict=True)) + if len(overflow) >= 1 << 24: + raise ValueError("too many noise flips for the GPU noise sampler") + pairs = np.zeros(len(overflow) + 1, dtype=_FLIP) + if overflow: + pairs["word"][:-1] = [w for w, _ in overflow] + pairs["mask"][:-1] = [m for _, m in overflow] + empty = np.empty(0, np.int32) + tables = NoiseTables( + block_shots=block_shots, + groups=NoiseGroups(0, (empty, empty, empty, empty, empty), records[:0], 0.0), + records=records, + overflow=pairs, + p_fire=fire.p_fire, + outcome_start=outcome_start, + num_words=num_words, + shared_bytes=shared_bytes, + shared_words=shared_words, + ) + every = np.ones(len(fire.p_fire), dtype=bool) + return replace(tables, groups=_groups(tables, every)) diff --git a/src/tsim/statevector/gpu/rows.py b/src/tsim/statevector/gpu/rows.py new file mode 100644 index 00000000..75e1eb0e --- /dev/null +++ b/src/tsim/statevector/gpu/rows.py @@ -0,0 +1,52 @@ +"""The rows of outputs that the GPU backends return, and how their kernels pack them.""" + +from __future__ import annotations + +import numpy as np + + +def pieces(columns: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Return how a kernel packs the outputs ``columns`` of shot words: ``(pieces, start)``. + + Output word ``k`` holds ``columns[64 k : 64 k + 64]``; its pieces are + ``pieces[start[k] : start[k + 1]]``, each a run of them that is consecutive + in one shot word, as ``(word, bit | (length - 1) << 6 | position << 12)``. + """ + found, start = [], [0] + for k in range(0, len(columns), 64): + word = columns[k : k + 64].tolist() + i = 0 + while i < len(word): + j = i + while ( + j + 1 < len(word) + and word[j + 1] == word[j] + 1 + and word[j + 1] >> 6 == word[i] >> 6 + ): + j += 1 + c = word[i] + found.append((c >> 6, (c & 63) | (j - i) << 6 | i << 12)) + i = j + 1 + start.append(len(found)) + return np.asarray(found or [(0, 0)], np.uint32), np.asarray(start, np.uint32) + + +def select( + rows: np.ndarray, columns: np.ndarray, num_outputs: int, bit_packed: bool +) -> np.ndarray: + """Return the outputs ``columns`` of rows of every output, as ``read_outputs`` does. + + The rows are bit-packed as ``np.packbits`` does, with zero bits past the + outputs; leading outputs in order are a view of them, or of whole bytes. + """ + count = len(columns) + if bit_packed and np.array_equal(columns, np.arange(count)): + packed = rows[:, : -(-count // 8)] + if count == num_outputs or count % 8 == 0: + return packed + packed = packed.copy() + packed[:, -1] &= np.uint8((1 << (count % 8)) - 1) + return packed + bits = np.unpackbits(rows, axis=1, count=num_outputs, bitorder="little") + bits = bits.view(np.bool_)[:, columns] + return np.packbits(bits, axis=1, bitorder="little") if bit_packed else bits diff --git a/src/tsim/statevector/gpu/tape.py b/src/tsim/statevector/gpu/tape.py new file mode 100644 index 00000000..139b9329 --- /dev/null +++ b/src/tsim/statevector/gpu/tape.py @@ -0,0 +1,737 @@ +"""GPU source of a component's tape: straight-line FP32 code over its amplitudes. + +A shot's ``2^W`` amplitudes are spread over ``T = 2^tb`` threads: the lowest +``tb`` axes are thread axes, so amplitude ``g`` is local ``g >> tb`` of thread +``g & (T - 1)``, held in registers. A Pauli with X on a thread axis reads its +partner amplitudes with a shuffle among the group. Every update is +written per amplitude ``g`` and its partner ``g ^ x``: a rotation is +``a_g ← cos a_g + κ s (-1)^{par((g^x)&z)} a_{g^x}``, with ``κ`` 1 or i, ``s`` its +signed strength. Masks, angles and amplitude indices are constants of the +generated code, so a tape compiles to a few thousand fused multiply-adds. +The same tape is written in each backend's language (``Dialect``); a backend's +prelude defines the helpers that it calls (``draw``, ``draw_one``, ``cmul``, +``cmac``, ``cmac_i``, ``par``, ``deviate``). + +The tape function follows ``shared.run_lanes`` for one shot on its own lane: its +lane symbols are the key, the postselected f-parts and its branch bits; it draws +a measurement's outcome as the CPU executor does, by its seed, shot and action; +a postselected output that reads 1 ends it (its *grave* is that action); at the +end it adds its branch parts, coins and hidden outputs to the shot's words. +The actions before the first one that reads a lane symbol run alike for every +shot (``shared.common_prefix``): they run once, on the host, and every shot +starts from the amplitudes they leave. + +With ``expectations``, the tape also returns, for each observable output, the +expectation of ``(-1)^output`` given the shot's history before the measurement +that sets the observable's last branch bit ``b``: +``(-1)^(output ⊕ b) (w0 - w1) / (w0 + w1)``. A fair coin's is 0 and an output +without branch bits has ``±1``. The expectation ignores postselection that +follows it, so when a postselected output of the component is tested at or +after that measurement, the tape returns the sampled ``(-1)^output`` instead: +the expectation's average over the accepted shots would be biased, the sample's +is not. +""" + +from __future__ import annotations + +import re +from collections.abc import Callable +from dataclasses import dataclass + +import numpy as np + +from tsim.statevector.program import MEASURE_DORMANT, ROTATE, ROTATE_DIAG +from tsim.statevector.shared import LaneProgram, common_prefix + +TINY = 1e-30 + +CERTAIN = 1e-11 +"""Outcome weight, relative to the total, below which an outcome is impossible. + +float32 amplitudes leave a certain outcome's other branch a weight of ~1e-14 to +1e-12; genuine outcome probabilities of tsim's circuits (T gates, Pauli noise, +small-angle rotations) stay above ~1e-10. +""" + +UNROLL_BUDGET = 12_000 +"""Amplitude updates (local spans summed over the actions) up to which a tape unrolls. + +A longer tape loops over its amplitudes instead: they then live in local memory +(cached in L1), but it compiles in seconds instead of minutes. +""" + +BLOCK_AXES = 4 +"""Local axes that a looping tape's run of rotations may mix in one pass. + +A looping tape applies consecutive rotations whose X masks share these axes +together, subcube by subcube, with the subcube's amplitudes in registers, so +it reads and writes the amplitudes once per run instead of once per rotation. +""" + + +@dataclass(frozen=True) +class Dialect: + """How a GPU language spells the parts of a tape that differ between them. + + Attributes: + inline: Qualifiers of the component function, inlined into its kernel. + noinline: Qualifiers of a looping tape's rotation passes. + private: Address space of pointers to a thread's own arrays. + deviation: Type of the pointer to the largest normalization deviation. + constant: Qualifiers of a program-scope constant array. + float2: Constructor of a ``float2``. + rsqrt: Reciprocal square root of a float. + fmax: Larger of two floats. + u64: Suffix of a 64-bit unsigned literal. + shuffle2: A ``float2`` ``{value}`` from the lane ``{lanes}`` (an XOR) away. + shuffle1: A ``float`` ``{value}`` from the lane ``{lanes}`` away. + lane_mask: Whether the shuffles take the group's lane mask ``gmask``, + which the tape functions then pass on. + goto: Whether a pruned shot jumps to the end; otherwise it breaks out + of a ``do``-block around the actions. + amplitude_chunk: Amplitudes per array of an unrolled tape, or None for + one array: a compiler may move a long array to the stack. + pair_projections: Whether a projection reads and updates each pair of + partner amplitudes at once, instead of every partner first: fewer + values are live. + + """ + + inline: str + noinline: str + private: str + deviation: str + constant: str + float2: str + rsqrt: str + fmax: str + u64: str + shuffle2: str + shuffle1: str + lane_mask: bool + goto: bool + amplitude_chunk: int | None = None + pair_projections: bool = False + + def hex64(self, value) -> str: + """Return an unsigned 64-bit literal.""" + return f"0x{int(value):016X}{self.u64}" + + def row_parity(self, row: np.ndarray) -> str: + """Return the parity of ``sym & row``, or ``0``.""" + terms = [f"(sym[{w}] & {self.hex64(m)})" for w, m in enumerate(row) if int(m)] + return f"par({' ^ '.join(terms)})" if terms else "0" + + +def f32(value: float) -> str: + """Return a float literal of ``value`` rounded to float32.""" + return repr(float(np.float32(value))) + "f" + + +def _par(value: int) -> int: + return bin(value).count("1") & 1 + + +def loops(program: LaneProgram, tb: int, start: int) -> bool: + """Return whether the tape from action ``start`` loops instead of unrolling. + + One runtime index puts the whole amplitude array in local memory, so a tape + either unrolls every action or loops every wide one. + """ + extent = program.extent[start:].astype(np.int64) + return int((1 << np.maximum(extent - tb, 0)).sum()) > UNROLL_BUDGET + + +class _Tape: + """Emits one component's actions for a thread group of ``2^tb`` threads.""" + + def __init__( + self, + program: LaneProgram, + tb: int, + emit: Callable[[str], None], + dialect: Dialect, + name: str = "run", + start: int = 0, + ): + self.d = dialect + self.p = program + self.tb = tb + self.T = 1 << tb + self.emit = emit + self.name = name + self.helpers: list[str] = [] + self.loop = loops(program, tb, start) + + def shfl(self, expr: str, xt: int) -> str: + return expr if xt == 0 else self.d.shuffle2.format(value=expr, lanes=xt) + + def f2(self, x: str, y: str) -> str: + return f"{self.d.float2}({x}, {y})" + + def span(self, a: int) -> int: + return 1 << max(int(self.p.extent[a]) - self.tb, 0) + + def masks(self, a: int) -> tuple[int, int, int, int]: + x, z, T = int(self.p.x_mask[a]), int(self.p.z_mask[a]), self.T + return x & (T - 1), x >> self.tb, z & (T - 1), z >> self.tb + + def rotation(self, a: int) -> None: + p, emit = self.p, self.emit + xt, xl, zt, zl = self.masks(a) + emit(f" f = {self.d.row_parity(p.sign_row[a])} ^ {int(p.sign_const[a])};") + if int(p.kind[a]) == ROTATE_DIAG: + g0, g1 = p.factor[a] + tpar = f" ^ par((u64)(t & {zt}))" if zt else "" + emit(f" {{ int ff = f{tpar};") + emit( + f" float2 h0 = ff ? {self.d.float2}({f32(g1.real)}, {f32(g1.imag)}) : {self.d.float2}({f32(g0.real)}, {f32(g0.imag)});" + ) + emit( + f" float2 h1 = ff ? {self.d.float2}({f32(g0.real)}, {f32(g0.imag)}) : {self.d.float2}({f32(g1.real)}, {f32(g1.imag)});" + ) + for loc in range(self.span(a)): + emit( + f" a[{loc}] = cmul(a[{loc}], {'h1' if _par(loc & zl) else 'h0'});" + ) + emit(" }") + return + coupling = p.coupling[a] + imaginary = coupling.real == 0 + strength = coupling.imag if imaginary else coupling.real + c = f32(p.cos[a]) + tpar = f" ^ par((u64)((t ^ {xt}) & {zt}))" if zt else "" + emit(f" s = (f{tpar}) ? {f32(-strength)} : {f32(strength)};") + update = "cmac_i" if imaginary else "cmac" + + def sign(loc: int) -> str: + return "-s" if _par((loc ^ xl) & zl) else "s" + + if xl: + low = xl & -xl + for loc in range(self.span(a)): + if loc & low: + continue + m = loc ^ xl + emit( + f" {{ float2 ql = {self.shfl(f'a[{m}]', xt)}, qm = {self.shfl(f'a[{loc}]', xt)};" + ) + emit( + f" a[{loc}] = {update}(a[{loc}], {c}, {sign(loc)}, ql); a[{m}] = {update}(a[{m}], {c}, {sign(m)}, qm); }}" + ) + else: + for loc in range(self.span(a)): + emit( + f" {{ float2 q = {self.shfl(f'a[{loc}]', xt)}; a[{loc}] = {update}(a[{loc}], {c}, {sign(loc)}, q); }}" + ) + + def weights(self, a: int) -> None: + """Set ``w0``, ``w1``: the outcome weights of a splitting measurement.""" + p, emit = self.p, self.emit + if int(p.kind[a]) == MEASURE_DORMANT: + emit(" w0 = 0.5f; w1 = 0.5f;") + return + x = int(p.x_mask[a]) + xt, xl, zt, zl = self.masks(a) + emit(" w0 = 0.f; w1 = 0.f;") + if self.loop: + self._weights_loop(a) + elif x == 0: + emit(f" {{ int tp = {f'par((u64)(t & {zt}))' if zt else '0'};") + for loc in range(self.span(a)): + emit( + f" {{ float m2 = a[{loc}].x * a[{loc}].x + a[{loc}].y * a[{loc}].y; if (tp ^ {_par(loc & zl)}) w1 += m2; else w0 += m2; }}" + ) + emit(" }") + else: + ip = p.i_power[a] + emit(f" {{ int tp = {f'par((u64)((t ^ {xt}) & {zt}))' if zt else '0'};") + emit( + f" float2 ipw = tp ? {self.d.float2}({f32(-ip.real)}, {f32(-ip.imag)}) : {self.d.float2}({f32(ip.real)}, {f32(ip.imag)});" + ) + for loc in range(self.span(a)): + m = loc ^ xl + sg = "-1.f" if _par(m & zl) else "1.f" + emit( + f" {{ float2 b = cmul({self.shfl(f'a[{m}]', xt)}, ipw); b.x *= {sg}; b.y *= {sg};" + ) + emit( + f" float px = a[{loc}].x + b.x, py = a[{loc}].y + b.y, mx = a[{loc}].x - b.x, my = a[{loc}].y - b.y;" + ) + emit( + " w0 += 0.25f * (px * px + py * py); w1 += 0.25f * (mx * mx + my * my); }" + ) + emit(" }") + if self.tb: + emit( + f" for (int o = 1; o < {self.T}; o <<= 1) {{ w0 += {self.d.shuffle1.format(value='w0', lanes='o')}; w1 += {self.d.shuffle1.format(value='w1', lanes='o')}; }}" + ) + + def project(self, a: int) -> None: + """Project onto ``outcome`` (``weight`` set) and free the pivot.""" + p, emit, tb = self.p, self.emit, self.tb + x = int(p.x_mask[a]) + xt, xl, zt, zl = self.masks(a) + pivot = int(p.pivot[a]) + ip = p.i_power[a] + emit( + f" scale = {self.d.rsqrt}({self.d.fmax}({'weight' if x == 0 else '2.f * weight'}, {TINY}f));" + ) + emit(f" {{ int lowhalf = {f'!((t >> {pivot}) & 1)' if pivot < tb else '1'};") + lbit = (1 << (pivot - tb)) if pivot >= tb else 0 + span = self.span(a) + if self.loop: + self._project_loop(a, lbit) + elif x == 0: + bt = (1 << pivot) if pivot < tb else 0 + emit(f" int tp = {f'par((u64)(t & {zt}))' if zt else '0'};") + for loc in range(span): + if lbit and (loc & lbit): + emit(f" a[{loc}] = {self.d.float2}(0.f, 0.f);") + continue + hi = loc | lbit + emit( + f" {{ float2 h = {self.shfl(f'a[{hi}]', bt)}; int keep = (tp ^ {_par(loc & zl)}) == outcome;" + ) + emit( + f" a[{loc}] = lowhalf ? (keep ? {self.d.float2}(a[{loc}].x * scale, a[{loc}].y * scale) : {self.d.float2}(h.x * scale, h.y * scale)) : {self.d.float2}(0.f, 0.f);" + ) + if lbit: + emit(f" a[{hi}] = {self.d.float2}(0.f, 0.f);") + emit(" }") + else: + emit(f" int tp = {f'par((u64)((t ^ {xt}) & {zt}))' if zt else '0'};") + emit( + f" float2 ipw = (tp ^ outcome) ? {self.d.float2}({f32(-ip.real)}, {f32(-ip.imag)}) : {self.d.float2}({f32(ip.real)}, {f32(ip.imag)});" + ) + if self.d.pair_projections: + self._project_pairs(a, lbit) + emit(" }") + return + # NOTE: every partner is read before any amplitude changes. + for loc in range(span): + emit(f" float2 q{loc} = {self.shfl(f'a[{loc ^ xl}]', xt)};") + for loc in range(span): + if lbit and (loc & lbit): + emit(f" a[{loc}] = {self.d.float2}(0.f, 0.f);") + continue + sg = "-1.f" if _par((loc ^ xl) & zl) else "1.f" + emit( + f" {{ float2 b = cmul(q{loc}, ipw); a[{loc}] = lowhalf ? {self.d.float2}((a[{loc}].x + {sg} * b.x) * scale, (a[{loc}].y + {sg} * b.y) * scale) : {self.d.float2}(0.f, 0.f); }}" + ) + emit(" }") + + def _project_pairs(self, a: int, lbit: int) -> None: + """Project, unrolled, one pair of partner amplitudes ``l``, ``l ^ xl`` at a time. + + An amplitude's update reads only its partner, so each pair is read + before either changes, and few values are live. + """ + emit = self.emit + xt, xl, _, zl = self.masks(a) + for loc in range(self.span(a)): + m = loc ^ xl + if m < loc: + continue + pair = [loc] if m == loc else [loc, m] + reads = ", ".join(f"q{j} = {self.shfl(f'a[{j ^ xl}]', xt)}" for j in pair) + emit(f" {{ float2 {reads};") + for j in pair: + if lbit and (j & lbit): + emit(f" a[{j}] = {self.f2('0.f', '0.f')};") + continue + sg = "-1.f" if _par((j ^ xl) & zl) else "1.f" + x_new = f"(a[{j}].x + {sg} * b.x) * scale" + y_new = f"(a[{j}].y + {sg} * b.y) * scale" + emit( + f" {{ float2 b = cmul(q{j}, ipw); a[{j}] = lowhalf ? {self.f2(x_new, y_new)} : {self.f2('0.f', '0.f')}; }}" + ) + emit(" }") + + # NOTE: the loops index ``a`` at runtime, so wide tapes keep their + # amplitudes in local memory (cached in L1) and compile quickly. + + def run(self, a: int) -> tuple[list[int], int]: + """Return the rotations from ``a`` that one pass applies, and the next action. + + The run ends at a splitting measurement or where its X masks would mix + more than ``BLOCK_AXES`` local axes; measurements that do not split + change no amplitude. + """ + p, rotations, axes, b = self.p, [], 0, a + while b < len(p.kind): + kind = int(p.kind[b]) + if kind in (ROTATE, ROTATE_DIAG): + joined = axes | (int(p.x_mask[b]) >> self.tb) + if rotations and bin(joined).count("1") > BLOCK_AXES: + break + rotations.append(b) + axes = joined + elif p.splits[b]: + break + b += 1 + return rotations, b + + def rotation_run(self, rotations: list[int]) -> None: + """Apply ``rotations`` in one pass over the subcubes of their X axes. + + The pass is a function of its own, so that the compiler takes a long tape's + runs one by one; the caller passes each rotation's signed strength (or + diagonal factors), which depend on the shot. + """ + p, emit = self.p, self.emit + axes = 0 + for a in rotations: + axes |= int(p.x_mask[a]) >> self.tb + bits = [i for i in range(axes.bit_length()) if axes >> i & 1] + # NOTE: the subcube's amplitude j is local base + offset[j]. + offset = [ + sum(1 << b for i, b in enumerate(bits) if j >> i & 1) + for j in range(1 << len(bits)) + ] + index = {o: j for j, o in enumerate(offset)} + span = max(self.span(a) for a in rotations) + name = f"{self.name}_{rotations[0]}" + mask_param = "unsigned gmask, " if self.d.lane_mask else "" + mask_arg = "gmask, " if self.d.lane_mask else "" + params, args = [], [] + for a in rotations: + xt, xl, zt, zl = self.masks(a) + f = f"({self.d.row_parity(p.sign_row[a])} ^ {int(p.sign_const[a])})" + if int(p.kind[a]) == ROTATE_DIAG: + g0, g1 = p.factor[a] + tpar = f" ^ par((u64)(t & {zt}))" if zt else "" + params += [f"float2 h0_{a}", f"float2 h1_{a}"] + v0 = f"{self.d.float2}({f32(g0.real)}, {f32(g0.imag)})" + v1 = f"{self.d.float2}({f32(g1.real)}, {f32(g1.imag)})" + args += [f"({f}{tpar}) ? {v1} : {v0}", f"({f}{tpar}) ? {v0} : {v1}"] + else: + coupling = p.coupling[a] + strength = coupling.imag if coupling.real == 0 else coupling.real + tpar = f" ^ par((u64)((t ^ {xt}) & {zt}))" if zt else "" + params.append(f"float s{a}") + args.append(f"({f}{tpar}) ? {f32(-strength)} : {f32(strength)}") + body = [ + f"{self.d.noinline} void {name}({self.d.private}float2* a, int t, {mask_param}{', '.join(params)}) {{", + f" #pragma unroll 1\n for (int base = 0; base < {span}; base = ((base | {axes}) + 1) & ~{axes}) {{", + f" float2 r[{len(offset)}];", + ] + body += [f" r[{j}] = a[base + {o}];" for j, o in enumerate(offset)] + for a in rotations: + xt, xl, zt, zl = self.masks(a) + # NOTE: a local index's sign parity is its subcube's (base) and its offset's. + outside = zl & ~axes + pb = f"par((u64)(base & {outside}))" if outside else "0" + if int(p.kind[a]) == ROTATE_DIAG: + body.append(f" {{ int pb = {pb};") + body += [ + f" r[{j}] = cmul(r[{j}], (pb ^ {_par(o & zl)}) ? h1_{a} : h0_{a});" + for j, o in enumerate(offset) + ] + body.append(" }") + continue + coupling = p.coupling[a] + update = "cmac_i" if coupling.real == 0 else "cmac" + c = f32(p.cos[a]) + body.append(f" {{ float sb = {pb} ? -s{a} : s{a};") + sign = ["-sb" if _par((o ^ xl) & zl) else "sb" for o in offset] + for j, o in enumerate(offset): + m = o ^ xl + if xl and m < o: + continue + if xl: + k = index[m] + body.append( + f" {{ float2 ql = {self.shfl(f'r[{k}]', xt)}, qm = {self.shfl(f'r[{j}]', xt)};" + ) + body.append( + f" r[{j}] = {update}(r[{j}], {c}, {sign[j]}, ql); r[{k}] = {update}(r[{k}], {c}, {sign[k]}, qm); }}" + ) + else: + body.append( + f" {{ float2 q = {self.shfl(f'r[{j}]', xt)}; r[{j}] = {update}(r[{j}], {c}, {sign[j]}, q); }}" + ) + body.append(" }") + body += [f" a[base + {o}] = r[{j}];" for j, o in enumerate(offset)] + body += [" }", "}"] + self.helpers.append("\n".join(body)) + emit( + f" // actions {rotations[0]}-{rotations[-1]}: {len(rotations)} rotations on local axes {axes:#x}" + ) + emit(f" {name}(a, t, {mask_arg}{', '.join(args)});") + + def _weights_loop(self, a: int) -> None: + emit = self.emit + x = int(self.p.x_mask[a]) + xt, xl, zt, zl = self.masks(a) + span = self.span(a) + if x == 0: + emit(f" {{ int tp = {f'par((u64)(t & {zt}))' if zt else '0'};") + emit(f" #pragma unroll 1\n for (int l = 0; l < {span}; ++l) {{") + emit(" float m2 = a[l].x * a[l].x + a[l].y * a[l].y;") + emit(f" if (tp ^ par((u64)(l & {zl}))) w1 += m2; else w0 += m2;") + emit(" }") + emit(" }") + return + ip = self.p.i_power[a] + emit(f" {{ int tp = {f'par((u64)((t ^ {xt}) & {zt}))' if zt else '0'};") + emit( + f" float2 ipw = tp ? {self.d.float2}({f32(-ip.real)}, {f32(-ip.imag)}) : {self.d.float2}({f32(ip.real)}, {f32(ip.imag)});" + ) + emit(f" #pragma unroll 1\n for (int l = 0; l < {span}; ++l) {{") + emit(f" int m = l ^ {xl};") + emit(f" float2 b = cmul({self.shfl('a[m]', xt)}, ipw);") + emit(f" if (par((u64)(m & {zl}))) {{ b.x = -b.x; b.y = -b.y; }}") + emit( + " float px = a[l].x + b.x, py = a[l].y + b.y, mx = a[l].x - b.x, my = a[l].y - b.y;" + ) + emit( + " w0 += 0.25f * (px * px + py * py); w1 += 0.25f * (mx * mx + my * my);" + ) + emit(" }") + emit(" }") + + def _project_loop(self, a: int, lbit: int) -> None: + emit, tb = self.emit, self.tb + x = int(self.p.x_mask[a]) + xt, xl, zt, zl = self.masks(a) + pivot = int(self.p.pivot[a]) + span = self.span(a) + if x == 0: + bt = (1 << pivot) if pivot < tb else 0 + emit(f" int tp = {f'par((u64)(t & {zt}))' if zt else '0'};") + emit(f" #pragma unroll 1\n for (int l = 0; l < {span}; ++l) {{") + emit(f" if ({lbit} & l) continue;") + emit( + f" float2 h = {self.shfl(f'a[l | {lbit}]', bt)}; int keep = (tp ^ par((u64)(l & {zl}))) == outcome;" + ) + emit( + f" a[l] = lowhalf ? (keep ? {self.d.float2}(a[l].x * scale, a[l].y * scale) : {self.d.float2}(h.x * scale, h.y * scale)) : {self.d.float2}(0.f, 0.f);" + ) + emit(" }") + if lbit: + emit( + f" #pragma unroll 1\n for (int l = 0; l < {span}; ++l) if (l & {lbit}) a[l] = {self.d.float2}(0.f, 0.f);" + ) + return + ip = self.p.i_power[a] + emit(f" int tp = {f'par((u64)((t ^ {xt}) & {zt}))' if zt else '0'};") + emit( + f" float2 ipw = (tp ^ outcome) ? {self.d.float2}({f32(-ip.real)}, {f32(-ip.imag)}) : {self.d.float2}({f32(ip.real)}, {f32(ip.imag)});" + ) + emit(f" float2 q[{span}];") + emit( + f" #pragma unroll 1\n for (int l = 0; l < {span}; ++l) q[l] = {self.shfl(f'a[l ^ {xl}]', xt)};" + ) + emit(f" #pragma unroll 1\n for (int l = 0; l < {span}; ++l) {{") + emit( + f" if ({lbit} & l) {{ a[l] = {self.d.float2}(0.f, 0.f); continue; }}" + ) + emit( + f" float2 b = cmul(q[l], ipw); float sg = par((u64)((l ^ {xl}) & {zl})) ? -1.f : 1.f;" + ) + emit( + f" a[l] = lowhalf ? {self.d.float2}((a[l].x + sg * b.x) * scale, (a[l].y + sg * b.y) * scale) : {self.d.float2}(0.f, 0.f);" + ) + emit(" }") + + def prune(self, tests: np.ndarray, at: int) -> None: + conds = [self.d.row_parity(self.p.prune_row[o]) for o in np.flatnonzero(tests)] + conds = [c for c in conds if c != "0"] + if conds: + leave = "goto done" if self.d.goto else "break" + self.emit(f" if ({' | '.join(conds)}) {{ grave = {at}; {leave}; }}") + + +def tape_function( + program: LaneProgram, + index: int, + tb: int, + observables: dict[int, int], + dialect: Dialect, +) -> str: + """Return ``component_``: one shot's run of a component's tape. + + Args: + program: The component's lane program. + index: The component's index (names the function). + tb: Log2 of the threads per shot. + observables: Slot of each of the component's observable outputs + (output slot in the program -> column of the expectations). + dialect: The language to write it in. + + """ + lines: list[str] = [] + emit = lines.append + p = program + A = len(p.kind) + n = 1 << max(p.width - tb, 0) + start, amps = common_prefix(p) + d = dialect + tape = _Tape(p, tb, emit, d, f"run{index}", start) + loaded = 0 + if start: + # NOTE: the prefix leaves amplitudes outside its last action's span zero. + loaded = tape.span(start - 1) + values = ", ".join( + f"{{{f32(v.real)}, {f32(v.imag)}}}" for v in amps[: loaded << tb] + ) + tape.helpers.append( + f"{d.constant} float2 start{index}[{loaded << tb}] = {{{values}}};" + ) + num_coins = p.coin_pattern.shape[0] + slots = A + (num_coins + 63) // 64 + ready = p.ready + # NOTE: an expectation is exact only if no postselected output of the component + # is tested at or after its measurement (other components and direct outputs + # depend on it only through the f-bits, which it is conditioned on). + exact = {o: ready[o] < 0 or not p.pruned[ready[o] :].any() for o in observables} + # NOTE: the measurements whose weights an observable's expectation reads. + watched = {int(ready[o]) for o in observables if ready[o] >= 0 and exact[o]} + + mask = "unsigned gmask, " if d.lane_mask else "" + emit( + f"{d.inline} int component_{index}({d.private}u64* words, u64 shot, u64 seed, int t, {mask}{d.private}float* ev, {d.deviation} deviation) {{" + ) + emit(f" u64 sym[{p.lane_words}];") + for w in range(p.key_words): + emit(f" sym[{w}] = words[{p.key_start + w}];") + for j, (pw, pm) in enumerate(zip(p.ps_word, p.ps_mask, strict=True)): + emit(f" sym[{p.key_words + j}] = words[{int(pw)}] & {d.hex64(pm)};") + for w in range(p.key_words + len(p.ps_word), p.lane_words): + emit(f" sym[{w}] = 0;") + chunk = d.amplitude_chunk + split = chunk is not None and not tape.loop and n > chunk + first = len(lines) + if split: + assert chunk is not None + # NOTE: unrolled, every index is a constant: ``_split_amplitudes`` sends + # ``a[i]`` to chunk ``i / chunk``. + emit(f" float2 {', '.join(f'a{c}[{chunk}]' for c in range(n // chunk))};") + for i in range(loaded, n): + emit(f" a[{i}] = {d.float2}(0.f, 0.f);") + for i in range(loaded if start else 0): + emit(f" a[{i}] = start{index}[{i << tb} | t];") + else: + emit(f" float2 a[{n}];") + emit( + f" #pragma unroll\n for (int i = {loaded}; i < {n}; ++i) a[i] = {d.float2}(0.f, 0.f);" + ) + if start: + emit( + f" #pragma unroll\n for (int i = 0; i < {loaded}; ++i) a[i] = start{index}[(i << {tb}) | t];" + ) + if not start: + emit(" if (t == 0) a[0].x = 1.f;") + emit(" int grave = 0x7fffffff;") + emit(" float w0, w1, total, weight, scale, s; int outcome, f;") + for a in sorted(watched): + emit(f" float ev_{a} = 0.f; int b_{a} = 0;") + if not d.goto: + emit(" do {") + if p.prune_first.any(): + tape.prune(p.prune_first, -1) + a = start - 1 + while a + 1 < A: + a += 1 + kind = int(p.kind[a]) + if kind in (ROTATE, ROTATE_DIAG) and tape.loop: + rotations, after = tape.run(a) + tape.rotation_run(rotations) + a = after - 1 + continue + emit(f" // action {a}: kind {kind}") + if kind in (ROTATE, ROTATE_DIAG): + tape.rotation(a) + continue + if not p.splits[a]: + continue + tape.weights(a) + emit(" total = w0 + w1;") + if kind != MEASURE_DORMANT: + emit(" if (t == 0) deviate(deviation, total);") + # NOTE: float32 rounding leaves an outcome that is certain a weight of up to + # ~1e-12 of the total; it counts as impossible (CERTAIN), and the others + # are drawn against the full 64-bit draw, as the CPU executor draws them. + emit( + f" if (w1 <= {CERTAIN}f * total) outcome = 0; else if (w0 <= {CERTAIN}f * total) outcome = 1;" + ) + emit( + f" else outcome = draw_one(draw(seed, shot, {slots}, {a}), w0, w1, total);" + ) + emit(" weight = outcome ? w1 : w0;") + if a in watched: + emit( + f" ev_{a} = total > {TINY}f ? (w0 - w1) / total : 0.f; b_{a} = outcome;" + ) + if kind != MEASURE_DORMANT: + tape.project(a) + bit = int(p.branch[a]) + if bit >= 0: + emit(f" if (outcome) sym[{bit >> 6}] |= 1{d.u64} << {bit & 63};") + if p.pruned[a]: + tape.prune(p.prune_outputs[a], a) + emit("done:" if d.goto else " } while (false);") + for o in range(len(p.output_bit)): + w = int(p.out_words[p.output_word[o]]) + parity = d.row_parity(p.branch_row[o]) + if parity != "0": + emit( + f" if ({parity}) words[{w}] ^= {d.hex64(1 << int(p.output_bit[o]))};" + ) + if num_coins: + emit(" { u64 coins = 0;") + for j in range(num_coins): + if j % 64 == 0: + emit(f" coins = draw(seed, shot, {slots}, {A + j // 64});") + flips = [ + (int(p.out_words[w]), int(m)) + for w, m in enumerate(p.coin_pattern[j]) + if int(m) + ] + if flips: + body = " ".join(f"words[{w}] ^= {d.hex64(m)};" for w, m in flips) + emit(f" if ((coins >> {j % 64}) & 1) {{ {body} }}") + emit(" }") + for o in range(len(p.output_bit)): + if ready[o] >= 0: + w = int(p.out_words[p.output_word[o]]) + emit( + f" if (grave < {int(ready[o])}) words[{w}] &= ~{d.hex64(1 << int(p.output_bit[o]))};" + ) + if observables: + emit(" if (ev != nullptr) {") + for o, column in observables.items(): + w = int(p.out_words[p.output_word[o]]) + bit = f"(int)((words[{w}] >> {int(p.output_bit[o])}) & 1)" + r = int(ready[o]) + if r >= 0 and not p.splits[r]: + # NOTE: a fair coin that no postselected output reads. + value = "0.f" + elif r < 0 or not exact[o]: + value = f"({bit} ? -1.f : 1.f)" + elif int(p.kind[r]) == MEASURE_DORMANT: + value = "0.f" + else: + value = f"((({bit}) ^ b_{r}) ? -ev_{r} : ev_{r})" + emit(f" ev[{column}] = grave < {max(r, 0)} ? 0.f : {value};") + emit(" }") + emit(" return grave;") + emit("}") + if split: + assert chunk is not None + lines[first:] = [_split_amplitudes(line, chunk) for line in lines[first:]] + return "\n\n".join([*tape.helpers, "\n".join(lines)]) + + +_AMPLITUDE = re.compile(r"\ba\[(\d+)\]") + + +def _split_amplitudes(line: str, chunk: int) -> str: + """Return ``line`` with each ``a[i]`` in its chunk's array.""" + + def place(match: re.Match) -> str: + c, i = divmod(int(match.group(1)), chunk) + return f"a{c}[{i}]" + + return _AMPLITUDE.sub(place, line) diff --git a/src/tsim/statevector/program.py b/src/tsim/statevector/program.py index 9eac40c4..7743ad61 100644 --- a/src/tsim/statevector/program.py +++ b/src/tsim/statevector/program.py @@ -133,7 +133,35 @@ def plan_component( f"component with outputs {outputs}, more than {MAX_WIDTH}. " f"{STABILIZER_RANK_HINT}" ) - return plan(ops, scheduled), component_layout + planned = plan(ops, scheduled) + # NOTE: measurements run as early as they can, which frees axes for later + # rotations. Where that frees none that a rotation reuses, running them late + # leaves more of the tape before the first action that depends on the shot, + # which runs once for all shots (``shared.common_prefix``). + deferred = schedule(ops, defer=True) + if deferred.width <= scheduled.width: + later = plan(ops, deferred) + if _shot_work(later) < _shot_work(planned): + planned = later + return planned, component_layout + + +def _shot_work(planned: Plan) -> int: + """Return the amplitudes that the actions from the first shot-dependent one touch. + + That is the first measurement or rotation whose sign reads a symbol. An action + touches the amplitudes below its extent (``lower``). + """ + work, active, shared = 0, 0, True + for action in planned.actions: + if isinstance(action, Rotate): + active |= action.x | action.z + shared = shared and isinstance(action, Rotate) and not action.sign.terms + if not shared: + work += 1 << active.bit_length() + if isinstance(action, MeasureActive): + active &= ~(1 << action.pivot) + return work def lower(plan: Plan, layout: ComponentLayout) -> StatevectorComponent: diff --git a/src/tsim/statevector/schedule.py b/src/tsim/statevector/schedule.py index 8d95e827..4e92812d 100644 --- a/src/tsim/statevector/schedule.py +++ b/src/tsim/statevector/schedule.py @@ -136,8 +136,13 @@ class Schedule: width: int -def schedule(ops: PauliOps) -> Schedule: - """Return an order of ``ops`` that keeps the active width small.""" +def schedule(ops: PauliOps, defer: bool = False) -> Schedule: + """Return an order of ``ops`` that keeps the active width small. + + With ``defer``, a measurement waits until the rotations left to run before it + would all widen the array, or until the end: it runs after every rotation + that it does not need to precede. + """ num_ops = ops.num_ops before = np.triu(anticommutation(ops), 1) ancestors = _ancestors(before) @@ -176,12 +181,16 @@ def run(j: int) -> None: while progressed: progressed = False for j in np.flatnonzero(~done): - if waiting[j]: + if waiting[j] or (defer and is_measurement[j]): continue keeps_width = not group.anticommuting(ops.x[j], ops.z[j]).any() if is_measurement[j] or keeps_width: run(int(j)) progressed = True + if defer and not progressed: + for j in np.flatnonzero(~done & is_measurement & (waiting == 0)): + run(int(j)) + progressed = True pending = np.flatnonzero(~done & is_measurement) if len(pending) == 0: for j in np.flatnonzero(~done): @@ -196,5 +205,6 @@ def run(j: int) -> None: best, lowest = k, highest for j in chains[best]: run(int(j)) - run(int(pending[best])) + if not defer: + run(int(pending[best])) return Schedule(order=np.asarray(order, dtype=np.int32), width=width) diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index f7f67444..66a61f97 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -321,6 +321,40 @@ def lane_bytes(component: StatevectorComponent, key: KeyBasis) -> int: return amplitudes + 2 * 8 * max_lane_words(component, key) + 64 +def common_prefix(program: LaneProgram) -> tuple[int, np.ndarray]: + """Return the leading actions that run alike on every lane, and the amplitudes after. + + They end at the first rotation whose sign reads a lane symbol or the first + measurement that splits lanes; measurements that do not split change no + amplitude. + """ + amps = np.zeros((1, 1 << program.width), dtype=np.complex128) + amps[0, 0] = 1 + for a, kind in enumerate(program.kind): + if kind not in (ROTATE_DIAG, ROTATE): + if program.splits[a]: + return a, amps[0] + continue + if program.sign_row[a].any(): + return a, amps[0] + span = 1 << int(program.extent[a]) + diagonal = kind == ROTATE_DIAG + _rotate_range( + amps, + 0, + span if diagonal else span >> 1, + span, + np.array([program.sign_const[a]], dtype=np.int64), + diagonal, + np.int64(program.x_mask[a]), + np.int64(program.z_mask[a]), + program.cos[a], + program.factor[a], + program.coupling[a], + ) + return len(program.kind), amps[0] + + @intrinsic def _ctpop(typingctx, value): def codegen(context, builder, signature, args): diff --git a/test/unit/statevector/test_gpu.py b/test/unit/statevector/test_gpu.py new file mode 100644 index 00000000..868a4e69 --- /dev/null +++ b/test/unit/statevector/test_gpu.py @@ -0,0 +1,262 @@ +"""The host side of the GPU backends: noise tables, row packing and the tape generator.""" + +import re +from test.helpers.statevector import CIRCUITS + +import numpy as np +import pytest + +from tsim.circuit import Circuit +from tsim.noise.channels import ChannelSampler +from tsim.statevector.gpu import noise, rows, tape +from tsim.statevector.shots import ShotSampler + +# NOTE: 200 detectors (four words), a fault that flips all four words (past the +# two that a record holds), two-qubit depolarizing noise (15 outcomes, alias +# tables) and an uneven single-qubit Pauli channel. +NOISY = "\n".join( + [ + "R " + " ".join(map(str, range(200))), + "E(0.05) X0 X70 X140 X199", + "DEPOLARIZE2(0.03) 3 4 130 131", + "PAULI_CHANNEL_1(0.02, 0.07, 0.01) 5 66", + "H 5 66", + "X_ERROR(0.01) " + " ".join(map(str, range(10, 200, 3))), + "M " + " ".join(map(str, range(200))), + *(f"DETECTOR rec[{i - 200}]" for i in range(200)), + "OBSERVABLE_INCLUDE(0) rec[-1] rec[-2]", + ] +) + + +def _shots(text: str) -> tuple[ShotSampler, ChannelSampler]: + sampler = Circuit(text).compile_detector_sampler(seed=1) + assert isinstance(sampler._shots, ShotSampler) + return sampler._shots, sampler._channel_sampler + + +def _model_words(tables: noise.NoiseTables, shots: int, seed: int) -> np.ndarray: + """Draw shot words from the tables as a noise kernel reads them. + + Each segment's trials fire independently with probability ``1 - e^log_keep``; + a fault draws a slot of its channel and a 24-bit fraction, takes the slot's + alias past its threshold and XORs the outcome's flips into its row. + """ + rng = np.random.default_rng(seed) + S, groups = tables.block_shots, tables.groups + first, per, begin, end, log_keep = groups.segments + out = np.zeros((shots, tables.num_words), dtype=np.uint64) + for base in range(0, shots, S): + ns = min(S, shots - base) + for k in range(len(first)): + p = -np.expm1(np.float64(log_keep[k])) + trials = np.arange(begin[k], end[k]) + fired = trials[rng.random(len(trials)) < p] + for trial in fired: + channel, row = divmod(int(trial), S) + if row >= ns: + continue + n = int(per[k]) + slot = int(rng.integers(n)) + fraction = int(rng.integers(1 << 24)) + at = int(first[k]) + channel * n + slot + record = groups.records[at] + if fraction >= int(record["alias"]) & 0x1FFFFFF: + record = groups.records[at - slot + (int(record["alias"]) >> 25)] + extra = int(record["extra"]) + words = [int(record["words"]) & 0xFFFF, int(record["words"]) >> 16] + masks = [record["m0"], record["m1"]] + for j in range(extra & 255): + if j < 2: + word, mask = words[j], masks[j] + else: + flip = tables.overflow[(extra >> 8) + j - 2] + word, mask = int(flip["word"]), flip["mask"] + out[base + row, word] ^= np.uint64(mask) + return out + + +def _bits(words: np.ndarray) -> np.ndarray: + return np.unpackbits(words.view(np.uint8), axis=1, bitorder="little") + + +@pytest.mark.parametrize("shared_words", [True, False]) +def test_noise_tables_draw_the_channel_samplers_distribution(shared_words): + host, channels = _shots(NOISY) + layout = host.layout + tables = noise.noise_tables( + channels, layout.noise, layout.num_words, 14 * 1024, shared_words + ) + assert (tables.overflow["word"][:-1] < layout.num_words).all() + assert len(tables.overflow) > 1 + shots = 20_000 + got = _bits(_model_words(tables, shots, seed=2)) + expected = np.zeros((shots, layout.num_words), dtype=np.uint64) + channels.sample_words(expected, layout.noise) + expected = _bits(expected) + p = (got.mean(0) + expected.mean(0)) / 2 + sigma = np.sqrt(np.maximum(p * (1 - p), 1 / shots) * 2 / shots) + assert (np.abs(got.mean(0) - expected.mean(0)) <= 5 * sigma).all() + # NOTE: the correlated fault flips detectors 0, 70, 140 and 199 together. + joint = [0, 70, 140, 199] + for bits in (got, expected): + together = bits[:, joint].all(axis=1).mean() + assert abs(together - 0.05 * (1 - 0.01) ** 2) < 0.01 + + +def test_noise_tables_size_blocks_for_their_memory(): + host, channels = _shots(NOISY) + layout = host.layout + shared = noise.noise_tables(channels, layout.noise, layout.num_words, 14 * 1024) + assert shared.shared_words + assert 8 * layout.num_words * shared.block_shots <= 14 * 1024 - 8 * 256 + assert 8 * shared.groups.max_events <= ( + 14 * 1024 - 8 * layout.num_words * shared.block_shots + ) + wide = noise.noise_tables( + channels, layout.noise, layout.num_words, 2048, shared_words=False + ) + assert wide.block_shots == noise.MAX_BLOCK_SHOTS + assert wide.groups.max_events <= 2048 // 8 + with pytest.raises(ValueError, match="do not fit"): + noise.noise_tables(channels, layout.noise, layout.num_words, 2048) + with pytest.raises(ValueError, match="at least one shot word"): + noise.noise_tables(channels, layout.noise, 0, 14 * 1024) + + +def test_pass_rate_and_deferred_channels(): + host, channels = _shots(NOISY) + layout = host.layout + read = np.zeros(layout.num_words, dtype=np.uint64) + # NOTE: detectors 0 and 5: the correlated fault and the uneven Pauli channel. + read[0] = np.uint64(1 | 1 << 5) + deferred = noise.deferred_channels(channels, layout.noise, read) + hits = (read[layout.noise.word] & layout.noise.mask) != 0 + start = channels._tables.outcome_start + for c in range(len(start) - 1): + flips = layout.noise.start[start[c]], layout.noise.start[start[c + 1]] + assert deferred[c] == (not hits[flips[0] : flips[1]].any()) + assert 0 < deferred.sum() < len(deferred) + shots = 200_000 + words = np.zeros((shots, layout.num_words), dtype=np.uint64) + channels.sample_words(words, layout.noise) + passed = ((words & read) == 0).all(axis=1).mean() + rate = noise.pass_rate(channels, layout.noise, read) + assert abs(passed - rate) < 5 * np.sqrt(rate * (1 - rate) / shots) + tables = noise.noise_tables(channels, layout.noise, layout.num_words, 14 * 1024) + first, later = tables.subset(~deferred), tables.subset(deferred) + assert np.isclose(first.faults + later.faults, tables.groups.faults) + + +def test_select_and_pieces_read_the_columns_of_rows(): + rng = np.random.default_rng(3) + num_outputs, count = 150, 40 + bits = rng.integers(0, 2, (64, num_outputs)).astype(np.uint8) + packed = np.packbits(bits, axis=1, bitorder="little") + for columns in ( + np.arange(count), + np.arange(num_outputs), + np.sort(rng.choice(num_outputs, count, replace=False)), + ): + got = rows.select(packed, columns, num_outputs, bit_packed=False) + np.testing.assert_array_equal(got, bits[:, columns].astype(bool)) + got = rows.select(packed, columns, num_outputs, bit_packed=True) + np.testing.assert_array_equal( + got, np.packbits(bits[:, columns], axis=1, bitorder="little") + ) + # NOTE: pack the columns from 64-bit shot words, as a kernel does. + words = np.zeros((64, 3), dtype=np.uint64) + words.view(np.uint8)[:, : packed.shape[1]] = packed + found, start = rows.pieces(columns) + out = np.zeros((64, -(-len(columns) // 64)), dtype=np.uint64) + for k in range(out.shape[1]): + for word, code in found[start[k] : start[k + 1]].tolist(): + bit, length, at = code & 63, ((code >> 6) & 63) + 1, code >> 12 + run = (words[:, word] >> np.uint64(bit)) & np.uint64((1 << length) - 1) + out[:, k] |= run << np.uint64(at) + unpacked = np.unpackbits(out.view(np.uint8), axis=1, bitorder="little") + np.testing.assert_array_equal(unpacked[:, : len(columns)], bits[:, columns]) + + +TEST_DIALECT = tape.Dialect( + inline="inline", + noinline="noinline", + private="thread ", + deviation="device atomic_uint*", + constant="constant", + float2="float2", + rsqrt="rsqrt", + fmax="fmax", + u64="UL", + shuffle2="shuffle({value}, {lanes})", + shuffle1="shuffle({value}, {lanes})", + lane_mask=False, + goto=False, + amplitude_chunk=32, + pair_projections=True, +) + + +def _programs(text, mask=True): + host, _ = _shots(text) + num_detectors = Circuit(text).num_detectors + plan = host.plan(np.ones(num_detectors, bool) if mask and num_detectors else None) + return list(plan.programs) + + +def _wide_unrolled(): + rng = np.random.default_rng(0) + lines = ["X_ERROR(0.2) " + " ".join(map(str, range(11)))] + for _ in range(7): + perm = rng.permutation(11) + for a, b in zip(perm[::2], perm[1::2], strict=False): + for q in (a, b): + t = rng.uniform(-1, 1, 3) + lines.append(f"U3({t[0]},{t[1]},{t[2]}) {q}") + lines.append(f"CX {a} {b}") + lines.append("M " + " ".join(map(str, range(11)))) + lines += [f"DETECTOR rec[{i - 11}]" for i in range(11)] + return "\n".join(lines) + + +def test_tapes_break_out_of_a_block_without_goto(): + programs = _programs(CIRCUITS["t_ghz"]) + assert any(p.pruned.any() or p.prune_first.any() for p in programs) + for i, program in enumerate(programs): + source = tape.tape_function(program, i, 0, {}, TEST_DIALECT) + assert "goto" not in source and "done:" not in source + assert source.count(" do {") == source.count("} while (false);") == 1 + if program.pruned.any(): + assert "break; }" in source + goto = tape.tape_function( + program, i, 0, {}, tape.Dialect(**{**TEST_DIALECT.__dict__, "goto": True}) + ) + assert "do {" not in goto and "done:" in goto + + +def test_unrolled_tapes_split_their_amplitudes_into_chunks(): + (program,) = _programs(_wide_unrolled(), mask=False) + assert program.width == 11 and not tape.loops(program, 5, 0) + source = tape.tape_function(program, 0, 5, {}, TEST_DIALECT) + assert re.search(r"\ba\[", source) is None + assert "float2 a0[32], a1[32];" in source and "a1[31]" in source + whole = tape.tape_function( + program, + 0, + 5, + {}, + tape.Dialect(**{**TEST_DIALECT.__dict__, "amplitude_chunk": None}), + ) + assert "float2 a[64];" in whole + + +def test_projections_read_pairs_or_every_partner_first(): + (program,) = _programs(_wide_unrolled(), mask=False) + unpaired = tape.Dialect(**{**TEST_DIALECT.__dict__, "pair_projections": False}) + paired = tape.tape_function(program, 0, 5, {}, TEST_DIALECT) + whole = tape.tape_function(program, 0, 5, {}, unpaired) + # NOTE: a non-diagonal projection loads its partners as q: all of them + # before it updates any, or a pair in a block of its own. + assert re.search(r"\n float2 q63 = ", whole) + assert not re.search(r"\n float2 q\d+ = ", paired) + assert re.search(r"\{ float2 q\d+ = .*, q\d+ = ", paired) diff --git a/test/unit/statevector/test_program.py b/test/unit/statevector/test_program.py index 673134d8..40f4560e 100644 --- a/test/unit/statevector/test_program.py +++ b/test/unit/statevector/test_program.py @@ -52,6 +52,28 @@ def test_plans_match_stabilizer_rank(name, sample_detectors): ) +def test_measurements_that_free_no_reused_axis_run_after_every_rotation(): + # NOTE: random two-qubit layers, then every qubit measured: measuring early + # would free no axis that a later rotation reuses, so the measurements stay + # at the end and every rotation runs before the shots split. + rng = np.random.default_rng(0) + lines = [] + for _ in range(8): + perm = rng.permutation(6) + for a, b in zip(perm[::2], perm[1::2], strict=True): + lines += [ + f"U3({t[0]},{t[1]},{t[2]}) {q}" + for q in (a, b) + for t in [rng.uniform(-1, 1, 3)] + ] + lines.append(f"CX {a} {b}") + lines.append("M " + " ".join(map(str, range(6)))) + (component,), layout = quantum_components("\n".join(lines), sample_detectors=False) + planned, _ = plan_component(component, output_indices(component), layout) + rotations = [isinstance(action, Rotate) for action in planned.actions] + assert rotations == sorted(rotations, reverse=True) + + def test_lower_encodes_every_action(): planned = Plan( width=2, diff --git a/test/unit/statevector/test_schedule.py b/test/unit/statevector/test_schedule.py index a10304f0..7fa1e86f 100644 --- a/test/unit/statevector/test_schedule.py +++ b/test/unit/statevector/test_schedule.py @@ -99,9 +99,32 @@ def test_widening_runs_the_measurement_with_the_lowest_peak_first(): assert lower_second.width == 2 -def test_schedules_keep_every_anticommuting_pair_in_order(): +def _measurement_that_a_later_rotation_need_not_follow(): + # Rotate about X0 and X1, measure Z0 and Z1, rotate about Z1: the last + # rotation commutes with the measurement of Z1 and keeps the width. + return PauliOps( + kind=np.array([ROTATION] * 3 + [MEASUREMENT] * 2, dtype=np.uint8), + x=_bits([1, 0], [0, 1], [0, 0], [0, 0], [0, 0]), + z=_bits([0, 0], [0, 0], [0, 1], [1, 0], [0, 1]), + angle=np.array([0.3, 0.3, 0.3, 0.0, 0.0]), + sign=np.zeros((5, 1), dtype=bool), + slot=np.array([-1, -1, -1, 0, 1], dtype=np.int32), + ) + + +def test_deferred_measurements_wait_for_the_rotations_that_keep_the_width(): + ops = _measurement_that_a_later_rotation_need_not_follow() + eager = schedule(ops) + assert eager.order.tolist() == [0, 3, 1, 4, 2] + deferred = schedule(ops, defer=True) + assert deferred.order.tolist() == [0, 3, 1, 2, 4] + assert eager.width == deferred.width == 1 + + +@pytest.mark.parametrize("defer", [False, True]) +def test_schedules_keep_every_anticommuting_pair_in_order(defer): for ops in component_ops(): - order = schedule(ops).order + order = schedule(ops, defer=defer).order assert sorted(order.tolist()) == list(range(ops.num_ops)) position = np.empty(ops.num_ops, dtype=int) position[order] = np.arange(ops.num_ops) @@ -109,18 +132,20 @@ def test_schedules_keep_every_anticommuting_pair_in_order(): assert np.all(position[first] < position[second]) -def test_schedules_preserve_the_distribution(): +@pytest.mark.parametrize("defer", [False, True]) +def test_schedules_preserve_the_distribution(defer): for ops in component_ops(): - reordered = ops.take(schedule(ops).order) + reordered = ops.take(schedule(ops, defer=defer).order) for f in f_patterns(ops.num_f): assert_distributions_close( ops_distribution(reordered, f), ops_distribution(ops, f) ) -def test_width_is_the_largest_structural_width_of_the_order(): +@pytest.mark.parametrize("defer", [False, True]) +def test_width_is_the_largest_structural_width_of_the_order(defer): for ops in component_ops(): - scheduled = schedule(ops) + scheduled = schedule(ops, defer=defer) group = StabilizerGroup(ops.num_lines) widths = [0] for j in scheduled.order: diff --git a/test/unit/statevector/test_shared.py b/test/unit/statevector/test_shared.py index 6f455dbb..8e9ee445 100644 --- a/test/unit/statevector/test_shared.py +++ b/test/unit/statevector/test_shared.py @@ -18,6 +18,7 @@ import stim import tsim.statevector.shared as shared +from tsim.statevector.plan import Rotate from tsim.statevector.program import MEASURE, MEASURE_DIAG, lower, plan_component NAMES = [ @@ -390,3 +391,47 @@ def test_every_action_stays_within_its_extent(name, coherent): assert np.all(support < reach) measured = np.isin(np.asarray(tape.kind), (MEASURE, MEASURE_DIAG)) assert np.all((1 << np.asarray(tape.pivot))[measured] < reach[measured]) + + +def _noiseless_rotations(width=6, layers=8): + """Random U3 pairs and CX gates, then every qubit measured: a long prefix.""" + rng = np.random.default_rng(width) + lines = [] + for _ in range(layers): + perm = rng.permutation(width) + for a, b in zip(perm[::2], perm[1::2], strict=False): + for q in (a, b): + t = rng.uniform(-1, 1, 3) + lines.append(f"U3({t[0]},{t[1]},{t[2]}) {q}") + lines.append(f"CX {a} {b}") + lines.append("M " + " ".join(map(str, range(width)))) + lines += [f"DETECTOR rec[{i - width}]" for i in range(width)] + return "\n".join(lines) + + +@pytest.mark.parametrize("name", [*NAMES, "long_prefix"]) +def test_the_common_prefix_runs_the_leading_shared_rotations(name): + text = _noiseless_rotations() if name == "long_prefix" else CIRCUITS[name] + starts = [] + for planned, _, compiled in _compiled_text(text): + _, program = component_program(compiled) + assert len(program.kind) == len(planned.actions) + start, amps = shared.common_prefix(program) + starts.append(start) + psi = np.zeros(1 << program.width, dtype=complex) + psi[0] = 1 + for a in range(start): + action = planned.actions[a] + if not isinstance(action, Rotate): + assert not program.splits[a] + continue + assert not program.sign_row[a].any() + psi = reference_rotate( + psi, action.x, action.z, action.angle, bool(program.sign_const[a]) + ) + np.testing.assert_allclose(amps, psi, atol=1e-12) + # NOTE: the prefix ends at the first action that depends on the shot. + if start < len(program.kind): + assert program.splits[start] or program.sign_row[start].any() + if name == "long_prefix": + assert max(starts) >= 20 From 1c3a3d8d0d45869a814b26286dd81568f42d505c Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 10 Nov 2026 11:34:39 -0700 Subject: [PATCH 02/38] fix(statevector): bound the GPU noise tables, chunk sizes and shot-work model - noise_tables raises when an outcome flips more than 255 words (the record's 8-bit count would wrap) and when the events alone do not fit shared_bytes without shared words (MIN_EVENTS is now a named constant). - Dialect rejects an amplitude_chunk that is not a positive power of two, which would declare fewer chunk arrays than the tape indexes. - _shot_work ends the shared prefix where shared.common_prefix does: an unread dormant measurement is a fair coin and no longer counts as shot work. Co-Authored-By: Claude Sonnet 5.5 --- src/tsim/statevector/gpu/noise.py | 20 ++++++++--- src/tsim/statevector/gpu/tape.py | 11 ++++-- src/tsim/statevector/program.py | 12 +++++-- test/unit/statevector/test_gpu.py | 43 +++++++++++++++++++++++ test/unit/statevector/test_program.py | 50 +++++++++++++++++++++++++++ 5 files changed, 128 insertions(+), 8 deletions(-) diff --git a/src/tsim/statevector/gpu/noise.py b/src/tsim/statevector/gpu/noise.py index 3c78a3c5..c366c8ae 100644 --- a/src/tsim/statevector/gpu/noise.py +++ b/src/tsim/statevector/gpu/noise.py @@ -44,6 +44,10 @@ FAULTS_PER_SEGMENT = 4.0 # NOTE: rows take 8 bits of an event. MAX_BLOCK_SHOTS = 256 +# NOTE: a block's event buffer holds at least this many events. +MIN_EVENTS = 256 +# NOTE: a record's flip count has 8 bits. +MAX_FLIPS = 255 _RECORD = np.dtype( [ @@ -117,7 +121,7 @@ def _groups(tables: NoiseTables, select: np.ndarray) -> NoiseGroups: outcomes = np.diff(outcome_start) faults = float(p_fire[select].sum()) max_events = min( - max(256, int(2 * faults * S + 256)), + max(MIN_EVENTS, int(2 * faults * S + MIN_EVENTS)), (tables.shared_bytes - 8 * tables.num_words * S * tables.shared_words) // 8, ) # NOTE: channels of equal fire probability and outcomes form a group, so a @@ -219,8 +223,8 @@ def noise_tables( have all of it. Raises: - ValueError: If the rows have no words, their words do not fit, or the - noise exceeds the kernels' limits. + ValueError: If the rows have no words, their words or 256 events do not + fit, or the noise exceeds the kernels' limits. """ fire = channels._tables @@ -235,11 +239,19 @@ def noise_tables( ) if not shared_words: block_shots = MAX_BLOCK_SHOTS - elif 8 * num_words * block_shots > shared_bytes - 8 * 256: + if shared_bytes < 8 * MIN_EVENTS: + raise ValueError( + f"the noise kernel's shared memory holds fewer than {MIN_EVENTS} events" + ) + elif 8 * num_words * block_shots > shared_bytes - 8 * MIN_EVENTS: raise ValueError( f"{num_words} shot words do not fit the noise kernel's shared memory" ) outcome_start = fire.outcome_start + if len(flips.start) > 1 and int(np.diff(flips.start).max()) > MAX_FLIPS: + raise ValueError( + f"the GPU noise sampler supports outcomes of up to {MAX_FLIPS} flipped words" + ) if len(fire.alias_prob) >= 1 << 24: raise ValueError("the GPU noise sampler supports fewer than 2^24 outcomes") records = np.zeros(len(fire.alias_prob), dtype=_RECORD) diff --git a/src/tsim/statevector/gpu/tape.py b/src/tsim/statevector/gpu/tape.py index 139b9329..6401943c 100644 --- a/src/tsim/statevector/gpu/tape.py +++ b/src/tsim/statevector/gpu/tape.py @@ -89,8 +89,9 @@ class Dialect: which the tape functions then pass on. goto: Whether a pruned shot jumps to the end; otherwise it breaks out of a ``do``-block around the actions. - amplitude_chunk: Amplitudes per array of an unrolled tape, or None for - one array: a compiler may move a long array to the stack. + amplitude_chunk: Amplitudes per array of an unrolled tape (a power of + two, so the arrays are full), or None for one array: a compiler + may move a long array to the stack. pair_projections: Whether a projection reads and updates each pair of partner amplitudes at once, instead of every partner first: fewer values are live. @@ -113,6 +114,12 @@ class Dialect: amplitude_chunk: int | None = None pair_projections: bool = False + def __post_init__(self) -> None: + """Check the chunk size.""" + chunk = self.amplitude_chunk + if chunk is not None and (chunk <= 0 or chunk & (chunk - 1)): + raise ValueError("amplitude_chunk must be a positive power of two") + def hex64(self, value) -> str: """Return an unsigned 64-bit literal.""" return f"0x{int(value):016X}{self.u64}" diff --git a/src/tsim/statevector/program.py b/src/tsim/statevector/program.py index 7743ad61..3fda4d29 100644 --- a/src/tsim/statevector/program.py +++ b/src/tsim/statevector/program.py @@ -17,7 +17,7 @@ from tsim.core.params import ParamLayout from tsim.statevector.pbc import ComponentLayout, read_pauli_ops -from tsim.statevector.plan import MeasureActive, Plan, Rotate, plan +from tsim.statevector.plan import MeasureActive, MeasureDormant, Plan, Rotate, plan from tsim.statevector.schedule import schedule ROTATE_DIAG = 0 @@ -149,11 +149,19 @@ def plan_component( def _shot_work(planned: Plan) -> int: """Return the amplitudes that the actions from the first shot-dependent one touch. - That is the first measurement or rotation whose sign reads a symbol. An action + That is the first rotation whose sign reads a symbol, measurement of the array + or dormant measurement whose outcome a rotation's sign reads, where + ``shared.common_prefix`` ends without postselection. The other dormant + measurements are fair coins, which touch no amplitude; any other action touches the amplitudes below its extent (``lower``). """ + read = frozenset().union( + *(a.sign.terms for a in planned.actions if isinstance(a, Rotate)) + ) work, active, shared = 0, 0, True for action in planned.actions: + if isinstance(action, MeasureDormant) and action.branch not in read: + continue if isinstance(action, Rotate): active |= action.x | action.z shared = shared and isinstance(action, Rotate) and not action.sign.terms diff --git a/test/unit/statevector/test_gpu.py b/test/unit/statevector/test_gpu.py index 868a4e69..43507798 100644 --- a/test/unit/statevector/test_gpu.py +++ b/test/unit/statevector/test_gpu.py @@ -122,6 +122,41 @@ def test_noise_tables_size_blocks_for_their_memory(): noise.noise_tables(channels, layout.noise, layout.num_words, 2048) with pytest.raises(ValueError, match="at least one shot word"): noise.noise_tables(channels, layout.noise, 0, 14 * 1024) + # NOTE: without the words, the memory still has to hold 256 events. + with pytest.raises(ValueError, match="fewer than 256 events"): + noise.noise_tables( + channels, layout.noise, layout.num_words, 1024, shared_words=False + ) + + +def test_noise_tables_reject_an_outcome_whose_flip_count_overflows(): + # NOTE: a record counts its flips in 8 bits: one fault that flips a detector + # in each of 256 words, then one that flips 255. + def tables(num_words): + targets = " ".join(f"X{64 * w}" for w in range(num_words)) + text = "\n".join( + [ + f"E(0.05) {targets}", + "M " + " ".join(str(64 * w) for w in range(num_words)), + *( + f"DETECTOR rec[{w - num_words}]\nSHIFT_COORDS(1)" + + "\nDETECTOR" * 63 + for w in range(num_words) + ), + ] + ) + host, channels = _shots(text) + layout = host.layout + assert np.diff(layout.noise.start).max() == num_words + return noise.noise_tables( + channels, layout.noise, layout.num_words, 14 * 1024, shared_words=False + ) + + with pytest.raises(ValueError, match="up to 255 flipped words"): + tables(256) + fits = tables(255) + assert int(fits.records["extra"][0]) & 255 == 255 + assert len(fits.overflow) == 253 + 1 def test_pass_rate_and_deferred_channels(): @@ -250,6 +285,14 @@ def test_unrolled_tapes_split_their_amplitudes_into_chunks(): assert "float2 a[64];" in whole +@pytest.mark.parametrize("chunk", [0, -32, 48]) +def test_dialects_reject_chunks_that_would_leave_a_partial_array(chunk): + # NOTE: an unrolled tape of 64 amplitudes in chunks of 48 would declare one + # array and index a second. + with pytest.raises(ValueError, match="power of two"): + tape.Dialect(**{**TEST_DIALECT.__dict__, "amplitude_chunk": chunk}) + + def test_projections_read_pairs_or_every_partner_first(): (program,) = _programs(_wide_unrolled(), mask=False) unpaired = tape.Dialect(**{**TEST_DIALECT.__dict__, "pair_projections": False}) diff --git a/test/unit/statevector/test_program.py b/test/unit/statevector/test_program.py index 40f4560e..d861779c 100644 --- a/test/unit/statevector/test_program.py +++ b/test/unit/statevector/test_program.py @@ -1,6 +1,7 @@ from test.helpers.statevector import ( CIRCUITS, assert_distributions_close, + component_program, f_patterns, output_indices, plan_distribution, @@ -12,6 +13,7 @@ import pytest import tsim.statevector.program as program +import tsim.statevector.shared as shared from tsim.circuit import Circuit from tsim.compile.pipeline import classify_direct, compile_program from tsim.core.graph import prepare_graph @@ -74,6 +76,54 @@ def test_measurements_that_free_no_reused_axis_run_after_every_rotation(): assert rotations == sorted(rotations, reverse=True) +def _work_past_the_common_prefix(planned: Plan, layout: ComponentLayout) -> int: + """Amplitudes that the actions past ``shared.common_prefix`` touch, coins aside.""" + _, lanes = component_program(lower(planned, layout)) + start, _ = shared.common_prefix(lanes) + coin = (lanes.kind == MEASURE_DORMANT) & ~lanes.splits + return sum(1 << int(e) for e in lanes.extent[start:][~coin[start:]]) + + +def test_shot_work_starts_where_the_common_prefix_ends(): + free = AffineRow(False, frozenset()) + rotations = (Rotate(1, 0, 0.3, free), Rotate(2, 1, 0.2, free)) + measure = MeasureActive(0, 1, 0, branch=1) + outputs = (AffineRow(False, frozenset({0})), AffineRow(False, frozenset({1}))) + # NOTE: an unread dormant measurement is a fair coin, which ends no prefix. + coin_first = Plan(2, 0, (MeasureDormant(0), *rotations, measure), outputs) + coin_last = Plan( + 2, 0, (*rotations, MeasureActive(0, 1, 0, branch=0), MeasureDormant(1)), outputs + ) + # NOTE: a rotation's sign reads this one, so the shots split at it. + read = Plan( + 2, + 0, + ( + MeasureDormant(0), + rotations[0], + Rotate(2, 1, 0.2, AffineRow(False, frozenset({0}))), + measure, + ), + outputs, + ) + assert program._shot_work(coin_first) == program._shot_work(coin_last) == 4 + assert program._shot_work(read) == 1 + 2 + 4 + 4 + layout = ComponentLayout((), {}, (0, 1)) + for planned in (coin_first, coin_last, read): + work = _work_past_the_common_prefix(planned, layout) + assert program._shot_work(planned) == work + + +@pytest.mark.parametrize("name", sorted(CIRCUITS)) +def test_shot_work_matches_the_common_prefix_of_every_plan(name): + components, layout = quantum_components(CIRCUITS[name], sample_detectors=True) + for component in components: + indices = output_indices(component) + planned, component_layout = plan_component(component, indices, layout) + work = _work_past_the_common_prefix(planned, component_layout) + assert program._shot_work(planned) == work + + def test_lower_encodes_every_action(): planned = Plan( width=2, From 44264cb7082e9c284d8d8cb9735af27d862d9cae Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 10 Nov 2026 20:46:42 +0100 Subject: [PATCH 03/38] fix(statevector): draw noise of tiny probabilities from 48-bit skips Below p ~ 2^-24, a kernel's 24-bit geometric skip reaches only some trials of a segment, so a group of many such channels fired them at biased per-channel rates (0.93 to 1.17 of their rates at p = 1e-9 on 3200 channels; the total rate was right). Such groups are now cut into segments of fewer than 2^-13 expected faults, which the kernels draw from 48 bits (their fine path). Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/gpu/noise.py | 8 ++++++++ test/unit/statevector/test_gpu.py | 20 ++++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/src/tsim/statevector/gpu/noise.py b/src/tsim/statevector/gpu/noise.py index c366c8ae..70016c08 100644 --- a/src/tsim/statevector/gpu/noise.py +++ b/src/tsim/statevector/gpu/noise.py @@ -35,6 +35,7 @@ from __future__ import annotations import itertools +import math from dataclasses import dataclass, replace import numpy as np @@ -48,6 +49,11 @@ MIN_EVENTS = 256 # NOTE: a record's flip count has 8 bits. MAX_FLIPS = 255 +# NOTE: a kernel draws a skip from 24 bits of a uniform, which leave trials that no +# skip reaches once p is below about 2^-24; a segment of fewer than 2^-12 expected +# faults is drawn from 48 bits instead (the kernels' fine path). +TINY_P = 2.0**-23 +FINE_FAULTS = 2.0**-13 _RECORD = np.dtype( [ @@ -143,6 +149,8 @@ def _groups(tables: NoiseTables, select: np.ndarray) -> NoiseGroups: trials = S * (hi - lo) # NOTE: float32 skips are exact below 2^24 trials only. pieces = max(-(-trials >> 24), round(trials * min(p, 1.0) / FAULTS_PER_SEGMENT)) + if p < TINY_P: + pieces = max(pieces, math.ceil(trials * -lk / FINE_FAULTS)) cuts = np.linspace(0, trials, pieces + 1).astype(np.int64) for a, b in itertools.pairwise(cuts): if b > a: diff --git a/test/unit/statevector/test_gpu.py b/test/unit/statevector/test_gpu.py index 43507798..f898436e 100644 --- a/test/unit/statevector/test_gpu.py +++ b/test/unit/statevector/test_gpu.py @@ -303,3 +303,23 @@ def test_projections_read_pairs_or_every_partner_first(): assert re.search(r"\n float2 q63 = ", whole) assert not re.search(r"\n float2 q\d+ = ", paired) assert re.search(r"\{ float2 q\d+ = .*, q\d+ = ", paired) + + +@pytest.mark.parametrize(("p", "qubits"), [(2e-8, 192), (1e-9, 3200)]) +def test_tiny_probability_segments_take_the_kernels_fine_path(p, qubits): + # NOTE: a 24-bit skip reaches only some trials of a segment of noise below + # p ~ 2^-24; segments of fewer than 2^-12 expected faults draw from 48 bits. + text = "\n".join( + [f"X_ERROR({p!r}) " + " ".join(map(str, range(qubits)))] + + ["M " + " ".join(map(str, range(qubits)))] + + [f"DETECTOR rec[{i - qubits}]" for i in range(qubits)] + ) + host, channels = _shots(text) + layout = host.layout + tables = noise.noise_tables(channels, layout.noise, layout.num_words, 46 * 1024) + _, _, begin, end, log_keep = tables.groups.segments + tiny = -log_keep < noise.TINY_P + assert tiny.any() + expected = (end - begin).astype(np.float32) * -log_keep + assert (expected[tiny] < 2.0**-12).all() + assert (end - begin).sum() == tables.block_shots * qubits From 0bff456d245e1fd3572e733b5a16d29775e32cfa Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 03:56:54 -0700 Subject: [PATCH 04/38] refactor(statevector): assert the tape dialects' chunk size Both dialects are internal constants, so a ValueError and its tests guard nothing a user can reach. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/gpu/tape.py | 5 ++--- test/unit/statevector/test_gpu.py | 8 -------- 2 files changed, 2 insertions(+), 11 deletions(-) diff --git a/src/tsim/statevector/gpu/tape.py b/src/tsim/statevector/gpu/tape.py index 6401943c..ce41dc3a 100644 --- a/src/tsim/statevector/gpu/tape.py +++ b/src/tsim/statevector/gpu/tape.py @@ -115,10 +115,9 @@ class Dialect: pair_projections: bool = False def __post_init__(self) -> None: - """Check the chunk size.""" + """Check that the chunk size is a positive power of two.""" chunk = self.amplitude_chunk - if chunk is not None and (chunk <= 0 or chunk & (chunk - 1)): - raise ValueError("amplitude_chunk must be a positive power of two") + assert chunk is None or (chunk > 0 and not chunk & (chunk - 1)) def hex64(self, value) -> str: """Return an unsigned 64-bit literal.""" diff --git a/test/unit/statevector/test_gpu.py b/test/unit/statevector/test_gpu.py index f898436e..3c86643f 100644 --- a/test/unit/statevector/test_gpu.py +++ b/test/unit/statevector/test_gpu.py @@ -285,14 +285,6 @@ def test_unrolled_tapes_split_their_amplitudes_into_chunks(): assert "float2 a[64];" in whole -@pytest.mark.parametrize("chunk", [0, -32, 48]) -def test_dialects_reject_chunks_that_would_leave_a_partial_array(chunk): - # NOTE: an unrolled tape of 64 amplitudes in chunks of 48 would declare one - # array and index a second. - with pytest.raises(ValueError, match="power of two"): - tape.Dialect(**{**TEST_DIALECT.__dict__, "amplitude_chunk": chunk}) - - def test_projections_read_pairs_or_every_partner_first(): (program,) = _programs(_wide_unrolled(), mask=False) unpaired = tape.Dialect(**{**TEST_DIALECT.__dict__, "pair_projections": False}) From 710ddf2f55371d9e0d2a8580a4619e59ff4f61da Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 03:56:58 -0700 Subject: [PATCH 05/38] refactor(statevector): drop rows.pieces Both backends return every output, in order, so each packed word is a whole shot word and the piece tables only encoded the identity. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/gpu/__init__.py | 4 ++-- src/tsim/statevector/gpu/rows.py | 28 +--------------------------- test/unit/statevector/test_gpu.py | 14 +------------- 3 files changed, 4 insertions(+), 42 deletions(-) diff --git a/src/tsim/statevector/gpu/__init__.py b/src/tsim/statevector/gpu/__init__.py index edf5a6c6..5bc5402d 100644 --- a/src/tsim/statevector/gpu/__init__.py +++ b/src/tsim/statevector/gpu/__init__.py @@ -3,6 +3,6 @@ A GPU backend runs the noise sampler and every component's tape as kernels of its own language. The host side is common to them: the noise tables that the noise kernels read (``noise``), the tape generator, which writes a component's -tape in any of their languages (``tape``), and the packing of the rows that -the kernels return (``rows``). +tape in any of their languages (``tape``), and the selection of outputs from +the rows that the kernels return (``rows``). """ diff --git a/src/tsim/statevector/gpu/rows.py b/src/tsim/statevector/gpu/rows.py index 75e1eb0e..5f1ee337 100644 --- a/src/tsim/statevector/gpu/rows.py +++ b/src/tsim/statevector/gpu/rows.py @@ -1,36 +1,10 @@ -"""The rows of outputs that the GPU backends return, and how their kernels pack them.""" +"""The rows of outputs that the GPU backends return.""" from __future__ import annotations import numpy as np -def pieces(columns: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """Return how a kernel packs the outputs ``columns`` of shot words: ``(pieces, start)``. - - Output word ``k`` holds ``columns[64 k : 64 k + 64]``; its pieces are - ``pieces[start[k] : start[k + 1]]``, each a run of them that is consecutive - in one shot word, as ``(word, bit | (length - 1) << 6 | position << 12)``. - """ - found, start = [], [0] - for k in range(0, len(columns), 64): - word = columns[k : k + 64].tolist() - i = 0 - while i < len(word): - j = i - while ( - j + 1 < len(word) - and word[j + 1] == word[j] + 1 - and word[j + 1] >> 6 == word[i] >> 6 - ): - j += 1 - c = word[i] - found.append((c >> 6, (c & 63) | (j - i) << 6 | i << 12)) - i = j + 1 - start.append(len(found)) - return np.asarray(found or [(0, 0)], np.uint32), np.asarray(start, np.uint32) - - def select( rows: np.ndarray, columns: np.ndarray, num_outputs: int, bit_packed: bool ) -> np.ndarray: diff --git a/test/unit/statevector/test_gpu.py b/test/unit/statevector/test_gpu.py index 3c86643f..016f30b0 100644 --- a/test/unit/statevector/test_gpu.py +++ b/test/unit/statevector/test_gpu.py @@ -183,7 +183,7 @@ def test_pass_rate_and_deferred_channels(): assert np.isclose(first.faults + later.faults, tables.groups.faults) -def test_select_and_pieces_read_the_columns_of_rows(): +def test_select_reads_the_columns_of_rows(): rng = np.random.default_rng(3) num_outputs, count = 150, 40 bits = rng.integers(0, 2, (64, num_outputs)).astype(np.uint8) @@ -199,18 +199,6 @@ def test_select_and_pieces_read_the_columns_of_rows(): np.testing.assert_array_equal( got, np.packbits(bits[:, columns], axis=1, bitorder="little") ) - # NOTE: pack the columns from 64-bit shot words, as a kernel does. - words = np.zeros((64, 3), dtype=np.uint64) - words.view(np.uint8)[:, : packed.shape[1]] = packed - found, start = rows.pieces(columns) - out = np.zeros((64, -(-len(columns) // 64)), dtype=np.uint64) - for k in range(out.shape[1]): - for word, code in found[start[k] : start[k + 1]].tolist(): - bit, length, at = code & 63, ((code >> 6) & 63) + 1, code >> 12 - run = (words[:, word] >> np.uint64(bit)) & np.uint64((1 << length) - 1) - out[:, k] |= run << np.uint64(at) - unpacked = np.unpackbits(out.view(np.uint8), axis=1, bitorder="little") - np.testing.assert_array_equal(unpacked[:, : len(columns)], bits[:, columns]) TEST_DIALECT = tape.Dialect( From df60b740621322517b72d6a9993dfbe52a872f74 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 03:58:09 -0700 Subject: [PATCH 06/38] refactor(statevector): share the noise split and screen pass rate Both GPU backends built the read mask of a postselected batch, chose the deferred channels and estimated the screen's pass rate with the same code; gpu.noise now has split, screen_pass_rate and words_fit (the shared-memory fit test that Metal repeated with literals). Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/gpu/noise.py | 61 +++++++++++++++++++++++++++++++ test/unit/statevector/test_gpu.py | 23 ++++++++++++ 2 files changed, 84 insertions(+) diff --git a/src/tsim/statevector/gpu/noise.py b/src/tsim/statevector/gpu/noise.py index 70016c08..c308fdf0 100644 --- a/src/tsim/statevector/gpu/noise.py +++ b/src/tsim/statevector/gpu/noise.py @@ -37,10 +37,15 @@ import itertools import math from dataclasses import dataclass, replace +from typing import TYPE_CHECKING import numpy as np from tsim.noise.channels import ChannelSampler, WordTables +from tsim.statevector.shared import pack_bits + +if TYPE_CHECKING: + from tsim.statevector.shots import ShotLayout, ShotPlan FAULTS_PER_SEGMENT = 4.0 # NOTE: rows take 8 bits of an event. @@ -54,6 +59,9 @@ # faults is drawn from 48 bits instead (the kernels' fine path). TINY_P = 2.0**-23 FINE_FAULTS = 2.0**-13 +DEFER_SHARE = 0.25 +"""Share of the expected faults that postselection cannot see, from which they +are drawn on the accepted shots only, after the tapes (``split``).""" _RECORD = np.dtype( [ @@ -199,6 +207,59 @@ def deferred_channels( return hits[start[1:]] == hits[start[:-1]] +def split( + tables: NoiseTables, + channels: ChannelSampler, + layout: ShotLayout, + plan: ShotPlan, + psm_words: np.ndarray, + observables_from: int | None = None, +) -> tuple[NoiseGroups, NoiseGroups] | None: + """Return a postselected batch's noise passes ``(screened, deferred)``, or None. + + The deferred channels flip none of the bits that the screen, the tapes or + the postselection ``psm_words`` read: the postselected outputs, the keys, + and, with ``observables_from``, the components' outputs from that index on + (the observables, whose expectations the tapes return). They are drawn + apart when they are at least ``DEFER_SHARE`` of the expected faults. + """ + read = np.zeros(layout.num_words, dtype=np.uint64) + read[: layout.output_words] = psm_words | plan.screen + for program in plan.programs: + start = program.key_start + read[start : start + program.key_words] = ~np.uint64(0) + np.bitwise_or.at(read, program.ps_word, program.ps_mask.astype(np.uint64)) + if observables_from is not None: + bits = np.zeros(64 * layout.output_words, dtype=bool) + for component in layout.components: + for index in component.output_indices: + if index >= observables_from: + bits[index] = True + read[: layout.output_words] |= pack_bits(bits, layout.output_words) + deferred = deferred_channels(channels, layout.noise, read) + later = tables.subset(deferred) + if later.faults > 0 and later.faults >= DEFER_SHARE * tables.groups.faults: + return tables.subset(~deferred), later + return None + + +def screen_pass_rate( + channels: ChannelSampler, layout: ShotLayout, plan: ShotPlan +) -> float: + """Return the estimated share of shots that pass ``plan``'s screen.""" + screen = np.zeros(layout.num_words, dtype=np.uint64) + screen[: layout.output_words] = plan.screen + return pass_rate(channels, layout.noise, screen) if screen.any() else 1.0 + + +def words_fit(num_words: int, shared_bytes: int) -> bool: + """Return whether rows of ``num_words`` words fit a block's ``shared_bytes``. + + The smallest block (32 shots) must leave room for ``MIN_EVENTS`` events. + """ + return 8 * 32 * num_words <= shared_bytes - 8 * MIN_EVENTS + + def pass_rate(channels: ChannelSampler, flips: WordTables, read: np.ndarray) -> float: """Return the share of shots whose faults flip no bit of ``read``, to first order. diff --git a/test/unit/statevector/test_gpu.py b/test/unit/statevector/test_gpu.py index 016f30b0..c7ab5b5a 100644 --- a/test/unit/statevector/test_gpu.py +++ b/test/unit/statevector/test_gpu.py @@ -117,6 +117,8 @@ def test_noise_tables_size_blocks_for_their_memory(): channels, layout.noise, layout.num_words, 2048, shared_words=False ) assert wide.block_shots == noise.MAX_BLOCK_SHOTS + assert noise.words_fit(layout.num_words, 14 * 1024) + assert not noise.words_fit(layout.num_words, 2048) assert wide.groups.max_events <= 2048 // 8 with pytest.raises(ValueError, match="do not fit"): noise.noise_tables(channels, layout.noise, layout.num_words, 2048) @@ -183,6 +185,27 @@ def test_pass_rate_and_deferred_channels(): assert np.isclose(first.faults + later.faults, tables.groups.faults) +def test_split_defers_the_channels_that_postselection_cannot_see(): + host, channels = _shots(NOISY) + layout = host.layout + tables = noise.noise_tables(channels, layout.noise, layout.num_words, 14 * 1024) + # NOTE: postselect detector 0 only: every channel but those that flip it waits. + mask = np.zeros(200, bool) + mask[0] = True + plan = host.plan(mask) + psm = np.zeros(layout.output_words, np.uint64) + psm[0] = 1 + found = noise.split(tables, channels, layout, plan, psm) + assert found is not None + first, later = found + assert 0 < first.faults < later.faults + assert np.isclose(first.faults + later.faults, tables.groups.faults) + # NOTE: reading every output leaves nothing to defer. + every = np.full(layout.output_words, ~np.uint64(0)) + assert noise.split(tables, channels, layout, plan, every) is None + assert noise.screen_pass_rate(channels, layout, host.plan(None)) == 1.0 + + def test_select_reads_the_columns_of_rows(): rng = np.random.default_rng(3) num_outputs, count = 150, 40 From 41503fd43ed381af4748db7c438445a37b4803cb Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 03:58:59 -0700 Subject: [PATCH 07/38] refactor(statevector): share the tape kernel's glue and limits Both GPU backends wrote the same component calls, seeds, word stores and postselection hits around the tape functions, and repeated MAX_WIDTH, the threads per shot, the deviation tolerance and the layout check. gpu.tape now has tape_kernel, threads_per_shot, deviation and supports. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/gpu/tape.py | 111 +++++++++++++++++++++++++++++- test/unit/statevector/test_gpu.py | 31 +++++++++ 2 files changed, 141 insertions(+), 1 deletion(-) diff --git a/src/tsim/statevector/gpu/tape.py b/src/tsim/statevector/gpu/tape.py index ce41dc3a..341f822b 100644 --- a/src/tsim/statevector/gpu/tape.py +++ b/src/tsim/statevector/gpu/tape.py @@ -37,11 +37,24 @@ import re from collections.abc import Callable from dataclasses import dataclass +from typing import TYPE_CHECKING import numpy as np from tsim.statevector.program import MEASURE_DORMANT, ROTATE, ROTATE_DIAG -from tsim.statevector.shared import LaneProgram, common_prefix +from tsim.statevector.shared import LaneProgram, common_prefix, pack_bits + +if TYPE_CHECKING: + from tsim.statevector.shots import ShotLayout, ShotPlan + +MAX_WIDTH = 11 +"""Widest component the GPU executors run: 32 threads of 64 amplitudes.""" + +AMPLITUDE_AXES = 5 +"""Log2 of the amplitudes per thread from which a shot takes more threads.""" + +DEVIATION_TOLERANCE = 1e-3 +"""Normalization deviation that float32 amplitudes tolerate without a warning.""" TINY = 1e-30 @@ -129,6 +142,102 @@ def row_parity(self, row: np.ndarray) -> str: return f"par({' ^ '.join(terms)})" if terms else "0" +def supports(layout: ShotLayout) -> bool: + """Return whether a GPU executor runs ``layout``. + + It needs shot words and components of at most ``MAX_WIDTH`` qubits. A layout without shot words (no outputs and no noise) stays on the CPU. + """ + return layout.num_words > 0 and all(c.width <= MAX_WIDTH for c in layout.components) + + +def threads_per_shot(layout: ShotLayout) -> int: + """Return log2 of a shot's threads: one per ``2^AMPLITUDE_AXES`` amplitudes, at most 32.""" + width = max((c.width for c in layout.components), default=0) + return min(5, max(0, width - AMPLITUDE_AXES)) + + +def deviation(bits) -> float: + """Return the normalization deviation of a call (the bits of a float32), or 0. + + It is 0 within ``DEVIATION_TOLERANCE``; a NaN stays NaN, which must warn. + """ + value = float(np.asarray(bits, np.uint32).reshape(1).view(np.float32)[0]) + return 0.0 if value <= DEVIATION_TOLERANCE else value + + +@dataclass(frozen=True) +class TapeKernel: + """The parts of a backend's tape kernel that a plan fills in. + + Attributes: + functions: Every component's tape function, ``component_``. + seeds: The components' streams ``seed_``, from the call's ``seed``. + calls: The calls of the tape functions, each lowering ``grave``. + stores: The stores of the words that the tapes write, ``row[k] = w[k]``. + hits: The postselected bits among those words, OR-ed (``0`` if none). + psm_words: The postselection mask, as output words. + + """ + + functions: str + seeds: str + calls: str + stores: str + hits: str + psm_words: np.ndarray + + +def tape_kernel( + layout: ShotLayout, + plan: ShotPlan, + mask: np.ndarray | None, + num_detectors: int, + tb: int, + dialect: Dialect, +) -> TapeKernel: + """Return the tape kernel's parts for ``plan`` of postselection ``mask``. + + The outputs from ``num_detectors`` on are observables, whose expectations + the tapes return in that order. + """ + functions = [] + for i, (program, component) in enumerate( + zip(plan.programs, layout.components, strict=True) + ): + observables = { + o: index - num_detectors + for o, index in enumerate(component.output_indices) + if index >= num_detectors + } + functions.append(tape_function(program, i, tb, observables, dialect)) + lanes = "t, gmask" if dialect.lane_mask else "t" + count = len(plan.programs) + bits = np.zeros(64 * layout.output_words, dtype=bool) + if mask is not None: + bits[: len(mask)] = mask + psm_words = pack_bits(bits, layout.output_words) + # NOTE: the tapes write only their outputs' words, and test the postselected + # outputs among them. + written = sorted({int(w) for p in plan.programs for w in p.out_words}) + hits = [ + f"(w[{w}] & {dialect.hex64(psm_words[w])})" for w in written if psm_words[w] + ] + return TapeKernel( + functions="\n\n".join(functions), + seeds="\n".join( + f" const u64 seed_{i} = mix64(seed + {i + 1}{dialect.u64} * GOLDEN);" + for i in range(count) + ), + calls="\n".join( + f" grave = min(grave, component_{i}(w, shot, seed_{i}, {lanes}, ev, deviation));" + for i in range(count) + ), + stores="\n".join(f" row[{w}] = w[{w}];" for w in written), + hits=" | ".join(hits) or "0", + psm_words=psm_words, + ) + + def f32(value: float) -> str: """Return a float literal of ``value`` rounded to float32.""" return repr(float(np.float32(value))) + "f" diff --git a/test/unit/statevector/test_gpu.py b/test/unit/statevector/test_gpu.py index c7ab5b5a..ff6fb8f4 100644 --- a/test/unit/statevector/test_gpu.py +++ b/test/unit/statevector/test_gpu.py @@ -326,3 +326,34 @@ def test_tiny_probability_segments_take_the_kernels_fine_path(p, qubits): expected = (end - begin).astype(np.float32) * -log_keep assert (expected[tiny] < 2.0**-12).all() assert (end - begin).sum() == tables.block_shots * qubits + + +def test_tape_kernel_fills_in_every_component(): + text = CIRCUITS["two_components"] + host, _ = _shots(text) + layout = host.layout + num_detectors = Circuit(text).num_detectors + mask = np.ones(num_detectors, bool) + plan = host.plan(mask) + tb = tape.threads_per_shot(layout) + kernel = tape.tape_kernel(layout, plan, mask, num_detectors, tb, TEST_DIALECT) + count = len(plan.programs) + assert count == len(layout.components) > 1 + assert kernel.calls.count("grave = min(grave, component_") == count + assert kernel.seeds.count("const u64 seed_") == count + assert all(f"component_{i}(" in kernel.functions for i in range(count)) + assert kernel.hits != "0" and kernel.psm_words.any() + unmasked = tape.tape_kernel( + layout, host.plan(None), None, num_detectors, tb, TEST_DIALECT + ) + assert unmasked.hits == "0" and not unmasked.psm_words.any() + assert tape.supports(layout) + + +def test_deviation_reports_only_what_float32_rounding_does_not_explain(): + def bits(value): + return np.float32(value).view(np.uint32) + + assert tape.deviation(bits(1e-4)) == 0.0 + assert tape.deviation(bits(0.01)) == pytest.approx(0.01) + assert np.isnan(tape.deviation(np.uint32(0x7FC00000))) From f25fcfcf962386aeff0635cfabc6ca072c7fc580 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 03:59:05 -0700 Subject: [PATCH 08/38] refactor(statevector): share the GPU samplers' fork warning Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/gpu/__init__.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/tsim/statevector/gpu/__init__.py b/src/tsim/statevector/gpu/__init__.py index 5bc5402d..a1e4391b 100644 --- a/src/tsim/statevector/gpu/__init__.py +++ b/src/tsim/statevector/gpu/__init__.py @@ -6,3 +6,9 @@ tape in any of their languages (``tape``), and the selection of outputs from the rows that the kernels return (``rows``). """ + +FORKED = ( + "This process was forked after its parent used the GPU, so it samples on the " + "CPU; start it with the 'spawn' or 'forkserver' method to sample on the GPU." +) +"""The warning of a GPU sampler in a process forked after its parent used the GPU.""" From c9b6ed659daeba7bf3d8128ae54bc63a7c451a4b Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 10 Nov 2026 19:51:56 +0100 Subject: [PATCH 09/38] feat(statevector): sample on Apple GPUs through Metal Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 9 + README.md | 7 +- docs/index.mdx | 7 +- docs/install.mdx | 18 +- pyproject.toml | 7 + src/tsim/circuit.py | 18 +- src/tsim/sampler.py | 82 +- src/tsim/statevector/metal/__init__.py | 10 + src/tsim/statevector/metal/codegen.py | 101 +++ src/tsim/statevector/metal/noise.py | 379 ++++++++++ src/tsim/statevector/metal/runtime.py | 309 ++++++++ src/tsim/statevector/metal/sampler.py | 992 +++++++++++++++++++++++++ test/conftest.py | 10 + test/integration/test_metal_sampler.py | 792 ++++++++++++++++++++ uv.lock | 61 +- 15 files changed, 2780 insertions(+), 22 deletions(-) create mode 100644 src/tsim/statevector/metal/__init__.py create mode 100644 src/tsim/statevector/metal/codegen.py create mode 100644 src/tsim/statevector/metal/noise.py create mode 100644 src/tsim/statevector/metal/runtime.py create mode 100644 src/tsim/statevector/metal/sampler.py create mode 100644 test/conftest.py create mode 100644 test/integration/test_metal_sampler.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 97d6a672..1061e45b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- `strategy="active_statevector"` samples on the GPU of an Apple-silicon Mac + when the `metal` extra is installed (`pip install "bloqade-tsim[metal]"`, which + installs PyObjC's Metal bindings) and every component has at most 11 active + qubits. The noise sampler and every statevector component run as Metal + kernels, in float32; postselection runs on the GPU, and only the kept shots' + rows are written. A seed's samples on the GPU differ from the CPU's and depend + on `batch_size`. The samplers fall back to the CPU, with a warning, for wider + components, when the GPU is out of memory and in a process forked after its + parent used the GPU. `TSIM_DEVICE=cpu` keeps them on the CPU. - The docs pages have a menu to copy a page or open it in ChatGPT or Claude, and the README, the install page and the agent skill point to the docs MCP server at `https://bloqade.quera.com/tsim/mcp`. The README and the install page also diff --git a/README.md b/README.md index 008c4b7b..2a04c640 100644 --- a/README.md +++ b/README.md @@ -23,8 +23,9 @@ Or with [uv](https://docs.astral.sh/uv/): uv add bloqade-tsim ``` -The default sampling strategy runs on the CPU. To run the stabilizer-rank strategies on an NVIDIA GPU, -install the JAX CUDA extra: +The default sampling strategy runs on the CPU, or on the GPU of an Apple-silicon Mac with the +`metal` extra (`pip install "bloqade-tsim[metal]"`). To run the stabilizer-rank strategies on an +NVIDIA GPU, install the JAX CUDA extra: ```bash pip install "bloqade-tsim[cuda13]" # or bloqade-tsim[cuda12] ``` @@ -82,7 +83,7 @@ Choose a strategy with the `strategy` argument of `compile_sampler` and `compile | Strategy | Method | Runs on | Cost grows with | Use it for | |---|---|---|---|---| -| `"active_statevector"` (default) | Dense state vector over the active qubits of each component, as in [Clifft](https://arxiv.org/abs/2604.27058) and [SymFT](https://arxiv.org/abs/2607.28600) | CPU (numba, multi-threaded) | Number of active qubits (at most 30) | Most circuits, especially with postselection | +| `"active_statevector"` (default) | Dense state vector over the active qubits of each component, as in [Clifft](https://arxiv.org/abs/2604.27058) and [SymFT](https://arxiv.org/abs/2607.28600) | CPU (numba, multi-threaded), or an Apple GPU (Metal) | Number of active qubits (at most 30) | Most circuits, especially with postselection | | `"cat5"`, `"bss"`, `"cutting"` | Stabilizer rank decomposition | CPU or GPU (JAX) | Number of non-Clifford gates | Few non-Clifford gates acting on many qubits; GPU throughput | ```python diff --git a/docs/index.mdx b/docs/index.mdx index 21c5681a..571ef1dc 100644 --- a/docs/index.mdx +++ b/docs/index.mdx @@ -39,8 +39,9 @@ samples = detector_sampler.sample(shots=100) uv add bloqade-tsim ``` -The default sampling strategy runs on the CPU. To run the stabilizer-rank strategies on an NVIDIA GPU, -install the JAX CUDA extra: +The default sampling strategy runs on the CPU, or on the GPU of an Apple-silicon Mac with the +`metal` extra (`pip install "bloqade-tsim[metal]"`). To run the stabilizer-rank strategies on an +NVIDIA GPU, install the JAX CUDA extra: ```bash uv add "bloqade-tsim[cuda13]" @@ -55,7 +56,7 @@ Choose a strategy with the `strategy` argument of `compile_sampler` and `compile | Strategy | Method | Runs on | Cost grows with | Use it for | |---|---|---|---|---| -| `"active_statevector"` (default) | Dense state vector over the active qubits of each component, as in [Clifft](https://arxiv.org/abs/2604.27058) and [SymFT](https://arxiv.org/abs/2607.28600) | CPU (numba, multi-threaded) | Number of active qubits (at most 30) | Most circuits, especially with postselection | +| `"active_statevector"` (default) | Dense state vector over the active qubits of each component, as in [Clifft](https://arxiv.org/abs/2604.27058) and [SymFT](https://arxiv.org/abs/2607.28600) | CPU (numba, multi-threaded), or an Apple GPU (Metal) | Number of active qubits (at most 30) | Most circuits, especially with postselection | | `"cat5"`, `"bss"`, `"cutting"` | Stabilizer rank decomposition | CPU or GPU (JAX) | Number of non-Clifford gates | Few non-Clifford gates acting on many qubits; GPU throughput | ```python diff --git a/docs/install.mdx b/docs/install.mdx index eeb3e10f..60ce201b 100644 --- a/docs/install.mdx +++ b/docs/install.mdx @@ -1,6 +1,6 @@ --- title: "Installation" -description: "Install Tsim with uv or pip, with optional CUDA support." +description: "Install Tsim with uv or pip, with optional GPU support." --- Tsim requires Python 3.11 or later. @@ -35,6 +35,22 @@ For GPU acceleration with CUDA: pip install "bloqade-tsim[cuda13]" ``` +## Apple GPUs + +On an Apple-silicon Mac, the default strategy (`"active_statevector"`) can sample on the GPU +through Metal. It needs PyObjC's Metal bindings, which the `metal` extra installs: + +```bash +uv add "bloqade-tsim[metal]" # or: pip install "bloqade-tsim[metal]" +``` + +With it installed, the samplers use the GPU whenever every component of the circuit has at +most 11 active qubits; otherwise they sample on the CPU, with a warning. The GPU computes in +float32, so a sample follows the CPU's distribution to within float32 rounding, and a seed's +samples differ from the CPU's (and depend on `batch_size`). Postselection runs on the GPU, so +heavily postselected circuits gain the most. Set `TSIM_DEVICE=cpu` to keep the samplers on the +CPU, for example to reproduce samples from another machine. + ## Agent skill Tsim ships an agent skill that teaches coding agents (Claude Code, Codex, Cursor and others) diff --git a/pyproject.toml b/pyproject.toml index ff9159ce..39b4f7d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,10 @@ Issues = "https://github.com/QuEraComputing/tsim/issues" [project.optional-dependencies] cuda12 = ["jax[cuda12]>=0.4.38"] cuda13 = ["jax[cuda13]>=0.7.0"] +# The statevector executor on Apple GPUs (Apple silicon only). +metal = [ + "pyobjc-framework-Metal>=10.0; sys_platform == 'darwin' and platform_machine == 'arm64'", +] [build-system] requires = ["hatchling"] @@ -167,6 +171,9 @@ extend_exclude = [ "^test/", ] +[tool.deptry.package_module_name_map] +pyobjc-framework-Metal = "Metal" + [tool.deptry.per_rule_ignores] DEP002 = ["ipywidgets"] diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index 05151114..22546a99 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -954,8 +954,11 @@ def compile_sampler( strategy: "active_statevector" (the default) to simulate each non-direct component on a dense state vector over its active qubits, or a stabilizer rank decomposition strategy ("cat5", - "bss", "cutting"). "active_statevector" samples on the CPU, on - ``numba.get_num_threads()`` threads. The stabilizer rank + "bss", "cutting"). "active_statevector" samples on an Apple GPU + when PyObjC's Metal sees one (the ``metal`` extra) and every + component has at most 11 active qubits (in float32), and + otherwise on the CPU, on ``numba.get_num_threads()`` threads; + ``TSIM_DEVICE=cpu`` keeps it on the CPU. The stabilizer rank strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed reproduces its samples only with the same versions of Tsim and its @@ -964,7 +967,8 @@ def compile_sampler( non-direct components, the samples of a seed depend on neither the batch size nor the thread count. With a stabilizer-rank strategy they also need the same explicit ``batch_size``, since the default batch size - depends on the free memory. + depends on the free memory. On a GPU, they do depend on the batch + size. Returns: A CompiledMeasurementSampler that can be used to sample measurements. @@ -991,8 +995,11 @@ def compile_detector_sampler( strategy: "active_statevector" (the default) to simulate each non-direct component on a dense state vector over its active qubits, or a stabilizer rank decomposition strategy ("cat5", - "bss", "cutting"). "active_statevector" samples on the CPU, on - ``numba.get_num_threads()`` threads. The stabilizer rank + "bss", "cutting"). "active_statevector" samples on an Apple GPU + when PyObjC's Metal sees one (the ``metal`` extra) and every + component has at most 11 active qubits (in float32), and + otherwise on the CPU, on ``numba.get_num_threads()`` threads; + ``TSIM_DEVICE=cpu`` keeps it on the CPU. The stabilizer rank strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed reproduces its samples only with the same versions of Tsim and its @@ -1002,6 +1009,7 @@ def compile_detector_sampler( size nor the thread count. With a stabilizer-rank strategy they also need the same explicit ``batch_size`` and the same reference sample settings, since the default batch size depends on the free memory. + On a GPU, they do depend on the batch size. Returns: A CompiledDetectorSampler that can be used to sample detectors and observables. diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 205affc5..3b6cb413 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -2,6 +2,7 @@ from __future__ import annotations +import os import secrets import warnings from dataclasses import dataclass @@ -19,6 +20,7 @@ from tsim.core.graph import prepare_graph from tsim.core.types import CompiledComponent, CompiledProgram, SamplingStrategy from tsim.noise.channels import ChannelSampler +from tsim.statevector.metal import runtime as metal_runtime from tsim.statevector.program import StatevectorComponent from tsim.statevector.shots import ShotSampler from tsim.utils.cuda_helpers import copy_d2h @@ -28,6 +30,7 @@ from jax import Array as PRNGKey from tsim.circuit import Circuit + from tsim.statevector.metal.sampler import MetalShotSampler # NOTE: postselection with a stabilizer-rank strategy screens shots for direct # discards in rounds of SCREEN_CHUNKS chunks of at most SCREEN_SHOTS shots. @@ -310,6 +313,53 @@ class SamplerStats: compiled_bytes: int +def _gpu_shots( + host: ShotSampler, + channels: ChannelSampler, + num_detectors: int, + seed: np.random.SeedSequence, +) -> ShotSampler | MetalShotSampler: + """Return a GPU shot sampler of ``host``'s program when one can run it, else ``host``. + + It warns when a GPU is there but cannot be used: in a process forked after its + parent used the GPU, for a component wider than the GPU executor runs, and + when the GPU sampler cannot be set up for lack of GPU memory. + """ + # NOTE: as JAX does for the stabilizer-rank strategies, a GPU is used whenever + # there is one. + if os.environ.get("TSIM_DEVICE", "").lower() == "cpu": + return host + if metal_runtime.forked_after_init(): + warnings.warn( + "This process was forked after its parent used the GPU, so it samples " + "on the CPU; start it with the 'spawn' or 'forkserver' method to " + "sample on the GPU.", + stacklevel=5, + ) + return host + if not metal_runtime.available(): + return host + from tsim.statevector.metal.sampler import MAX_WIDTH, MetalShotSampler + + layout = host.layout + if layout.num_words == 0: + return host + if not MetalShotSampler.supports(layout): + width = max(c.width for c in layout.components) + warnings.warn( + f"An Apple GPU is available, but a component needs {width} active " + f"qubits, more than the {MAX_WIDTH} that the GPU executor supports; " + "sampling on the CPU.", + stacklevel=5, + ) + return host + try: + return MetalShotSampler(host, channels, num_detectors, seed) + except MemoryError: + warnings.warn("The GPU is out of memory; sampling on the CPU.", stacklevel=5) + return host + + class _CompiledSamplerBase: """Base class for compiled samplers with common initialization logic.""" @@ -331,8 +381,11 @@ def __init__( strategy: "active_statevector" (the default) to simulate each non-direct component on a dense state vector over its active qubits, or a stabilizer rank decomposition strategy ("cat5", - "bss", "cutting"). "active_statevector" samples on the CPU, on - ``numba.get_num_threads()`` threads. The stabilizer rank + "bss", "cutting"). "active_statevector" samples on an Apple GPU + when PyObjC's Metal sees one (the ``metal`` extra) and every + component has at most 11 active qubits (in float32), and + otherwise on the CPU, on ``numba.get_num_threads()`` threads; + ``TSIM_DEVICE=cpu`` keeps it on the CPU. The stabilizer rank strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed reproduces its samples only with the same versions of Tsim and its @@ -341,14 +394,18 @@ def __init__( non-direct components, the samples of a seed depend on neither the batch size nor the thread count. With a stabilizer-rank strategy they also need the same explicit ``batch_size``, since the default batch size - depends on the free memory. + depends on the free memory. On a GPU, they do depend on the batch + size. """ if seed is None: seed = secrets.randbits(128) - # NOTE: the noise, the shot words and the JAX key draw from children of the - # whole seed, so different seeds give independent samples. - noise_seed, shots_seed, key_seed = np.random.SeedSequence(seed).spawn(3) + # NOTE: the noise, the shot words, the JAX key and the GPU draw from children + # of the whole seed, so different seeds give independent samples. The first + # three are those of a spawn(3), so the CPU's samples are unchanged. + noise_seed, shots_seed, key_seed, gpu_seed = np.random.SeedSequence(seed).spawn( + 4 + ) prepared = prepare_graph(circuit, sample_detectors=sample_detectors) self._program = compile_program(prepared, mode=mode, strategy=strategy) @@ -364,15 +421,20 @@ def __init__( program = self._program # NOTE: programs of statevector components, and programs without components # under any strategy, are sampled as shot words; the others with JAX. - self._shots: ShotSampler | None = None + self._shots: ShotSampler | MetalShotSampler | None = None self._backend: _StabilizerRankBackend | None = None if strategy == "active_statevector" or not program.components: - self._shots = ShotSampler( + host = ShotSampler( program, self._channel_sampler, self._num_detectors, shots_seed, ) + self._shots = host + if mode == "sequential": + self._shots = _gpu_shots( + host, self._channel_sampler, self._num_detectors, gpu_seed + ) else: self._backend = _StabilizerRankBackend(program, key_seed) @@ -850,7 +912,8 @@ def __init__( non-direct components, the samples of a seed depend on neither the batch size nor the thread count. With a stabilizer-rank strategy they also need the same explicit ``batch_size``, since the default batch size - depends on the free memory. + depends on the free memory. On a GPU, they do depend on the batch + size. """ super().__init__( @@ -951,6 +1014,7 @@ def __init__( size nor the thread count. With a stabilizer-rank strategy they also need the same explicit ``batch_size`` and the same reference sample settings, since the default batch size depends on the free memory. + On a GPU, they do depend on the batch size. """ super().__init__( diff --git a/src/tsim/statevector/metal/__init__.py b/src/tsim/statevector/metal/__init__.py new file mode 100644 index 00000000..bf4c629f --- /dev/null +++ b/src/tsim/statevector/metal/__init__.py @@ -0,0 +1,10 @@ +"""The statevector executor on an Apple GPU, through Metal. + +The noise sampler (``noise``) and every component's tape (``codegen``, written +by ``tsim.statevector.gpu.tape``) run as Metal kernels that PyObjC compiles at +runtime (``runtime``); ``MetalShotSampler`` (``sampler``) replaces +``ShotSampler`` behind the compiled samplers whenever PyObjC sees an Apple GPU +and every component fits. Amplitudes are complex64. It needs +``pyobjc-framework-Metal`` (``pip install bloqade-tsim[metal]``) and Apple +silicon. +""" diff --git a/src/tsim/statevector/metal/codegen.py b/src/tsim/statevector/metal/codegen.py new file mode 100644 index 00000000..833eda5f --- /dev/null +++ b/src/tsim/statevector/metal/codegen.py @@ -0,0 +1,101 @@ +"""Metal source of a component's tape (see ``tsim.statevector.gpu.tape``). + +The Metal Shading Language has SIMD-group shuffles for warp shuffles, no +``goto`` (a pruned shot breaks out of a ``do``-block) and no double precision: +near the boundary ``draw_one`` refines the rarer outcome's share as a +float-float ``q + e / total`` instead. Unrolled tapes keep their amplitudes in +arrays of 32, which Metal's compiler keeps in registers, and project partner +amplitudes in pairs, which keeps fewer values live. +""" + +from __future__ import annotations + +from tsim.statevector.gpu import tape +from tsim.statevector.shared import LaneProgram + +METAL = tape.Dialect( + inline="inline", + noinline="__attribute__((noinline))", + private="thread ", + deviation="device atomic_uint*", + constant="constant", + float2="float2", + rsqrt="rsqrt", + fmax="fmax", + u64="UL", + shuffle2="simd_shuffle_xor({value}, (ushort){lanes})", + shuffle1="simd_shuffle_xor({value}, (ushort){lanes})", + lane_mask=False, + goto=False, + # NOTE: Metal's compiler moves a thread array of 64 complex amplitudes (a + # width-11 component on 32 threads) to the stack, which costs the tape 1.8x. + amplitude_chunk=32, + pair_projections=True, +) + +PRELUDE = r""" +#include +using namespace metal; +typedef ulong u64; +#define GOLDEN 0x9E3779B97F4A7C15UL + +inline u64 mix64(u64 z) { + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9UL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBUL; + return z ^ (z >> 31); +} +// The CPU executor's draw: splitmix64 at index shot * slots + slot of stream seed. +inline u64 draw(u64 seed, u64 shot, u64 slots, u64 slot) { + return mix64(seed + (shot * slots + slot) * GOLDEN + GOLDEN); +} +inline int par(u64 v) { return popcount(v) & 1; } +// draw_one against the rarer outcome's share, which float32 keeps to full +// relative precision: a rare 0 takes the top of the range, as on the CPU +// (``shared._cut``). The share is refined to about 46 bits as q + e / total +// (Metal has no double): q correctly rounded, e its exact residual. Not +// inlined, as it is rarely taken. +__attribute__((noinline)) int draw_one_exact(u64 r, float w0, float w1, float total) { + int rare0 = w0 < w1; + float w = rare0 ? w0 : w1; + float q = precise::divide(w, total); + float e = fma(-q, total, w); + u64 below = u64(q * 0x1p64f) + u64(long(precise::divide(e, total) * 0x1p64f)); + return ((rare0 ? ~r : r) < below) ^ rare0; +} +// Whether the 64-bit draw r gives outcome 1, of weight w1 against w0. float32 +// decides it unless r / 2^64 lies within 2^-21 of the share w1 / total, which +// covers both roundings and fast division's error; there draw_one_exact does. +inline int draw_one(u64 r, float w0, float w1, float total) { + float p = w1 / total, x = float(r) * 0x1p-64f; + if (x < p * (1.f - 0x1p-21f)) return 1; + if (x > p * (1.f + 0x1p-21f)) return 0; + return draw_one_exact(r, w0, w1, total); +} +// A uniform in (0, 1] from the high 24 bits, so that log never sees 0. +inline float unit(u64 r) { return float((r >> 40) + 1) * (1.0f / 16777216.0f); } +inline float2 cmul(float2 p, float2 q) { + return float2(p.x * q.x - p.y * q.y, p.x * q.y + p.y * q.x); +} +// c a + s q and c a + i s q +inline float2 cmac(float2 a, float c, float s, float2 q) { + return float2(fma(c, a.x, s * q.x), fma(c, a.y, s * q.y)); +} +inline float2 cmac_i(float2 a, float c, float s, float2 q) { + return float2(fma(c, a.x, -s * q.y), fma(c, a.y, s * q.x)); +} +// The largest normalization deviation, as the bits of a non-negative float. +inline void deviate(device atomic_uint* deviation, float total) { + float d = fabs(total - 1.f); + if (!(d <= 1e-3f)) + atomic_fetch_max_explicit(deviation, d == d ? as_type(d) : 0x7f800000u, memory_order_relaxed); +} +""" + +hex64 = METAL.hex64 + + +def tape_function( + program: LaneProgram, index: int, tb: int, observables: dict[int, int] +) -> str: + """Return ``component_`` in Metal (``tape.tape_function``).""" + return tape.tape_function(program, index, tb, observables, METAL) diff --git a/src/tsim/statevector/metal/noise.py b/src/tsim/statevector/metal/noise.py new file mode 100644 index 00000000..9d6d7c95 --- /dev/null +++ b/src/tsim/statevector/metal/noise.py @@ -0,0 +1,379 @@ +"""Noise sampling into shot words on a Metal GPU. + +The kernels read the tables of ``tsim.statevector.gpu.noise`` (see there): a +threadgroup owns ``S`` consecutive shots, walks segments of the channel groups +with geometric skips, records the faults as events in threadgroup memory and +then XORs their outcomes' flips into the shots' words. Metal has no 64-bit +atomics, so a flip XORs the two halves of a word with 32-bit atomics, and a +threadgroup has 32 KB of threadgroup memory. Rows too wide for it keep their +words in device memory, where the faults are XORed in with device atomics. + +A block then screens its shots (constants, postselected direct detectors), +writes the survivors' rows, appends them to a queue for the tape kernel and +marks them in a bitmap of the batch's shots. Under postselection, +``noise_kept`` draws the noise that cannot change whether a shot is kept +(``gpu.noise.deferred_channels``) on the kept shots only, after the tapes. +""" + +from __future__ import annotations + +from tsim.noise.channels import ChannelSampler, WordTables +from tsim.statevector.gpu.noise import NoiseTables, noise_tables as _noise_tables + +SHARED_BYTES = 14 * 1024 +"""Threadgroup memory of a block's words and events. + +A threadgroup may have 32 KB, but the kernel is latency bound: with 14 KB, a +GPU core keeps more blocks resident, and the noise of end-to-end cultivation +d=5 takes 7.3 instead of 11.5 ns per shot on an M4 Pro (most other circuits +gain 10-25%, distillation loses up to 15%). +""" + +NOISE_SHARED = r""" +kernel void noise( + device u64* out [[buffer(0)]], constant NoiseArgs& args [[buffer(1)]], + device const int* seg_first [[buffer(2)]], device const int* seg_outcomes [[buffer(3)]], + device const int* seg_begin [[buffer(4)]], device const int* seg_end [[buffer(5)]], + device const float* seg_log_keep [[buffer(6)]], device const Record* records [[buffer(7)]], + device const Flip* overflow [[buffer(8)]], + device const u64* constants [[buffer(9)]], device const u64* screen [[buffer(10)]], + device uint* queue [[buffer(11)]], device atomic_uint* count [[buffer(12)]], + device uint* survivors [[buffer(13)]], + threadgroup uint2* events [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], + uint nt [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]]) +{{ + threadgroup u64 sw[{NW} * {S}]; + threadgroup atomic_int num_events; + threadgroup bool passed[{S}]; + long base = (long)bid * {S}; + int ns = (int)min((long){S}, args.shots - base); + for (uint i = tid; i < {NW} * {S}; i += nt) sw[i] = 0; + if (tid == 0) atomic_store_explicit(&num_events, 0, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + Segments seg = {{seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep}}; + add_faults(sw, events, &num_events, args.max_events, ns, args.seed ^ ((u64)bid * 0xD1B54A32D192ED03UL), + args.num_segments, seg, records, overflow, tid, nt); + for (int s0 = 0; s0 < {S}; s0 += nt) {{ + int s = s0 + tid; + bool alive = false; + if (s < ns) {{ + u64 hit = 0; + for (int i = 0; i < {OW}; ++i) {{ + u64 v = sw[i * {S} + s] ^ constants[i]; + hit |= v & screen[i]; + sw[i * {S} + s] = v; + }} + alive = !hit; + }} + if (s < {S}) passed[s] = alive; + uint ballot = uint(static_cast(simd_ballot(alive))); + int warp_first = s0 + (tid & ~31u); + if (args.has_survivors && lane == 0 && warp_first < ns) + survivors[(base + warp_first) >> 5] = ballot; + if (ballot) {{ + uint leader = ctz(ballot), b = 0; + if (lane == leader) b = atomic_fetch_add_explicit(count, popcount(ballot), memory_order_relaxed); + b = simd_broadcast(b, (ushort)leader); + if (alive) queue[b + popcount(ballot & ((1u << lane) - 1))] = (uint)(base + s); + }} + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + device u64* rows = out + base * {NW}; + for (int i = tid; i < ns * {NW}; i += nt) {{ + int s = i / {NW}, word = i - s * {NW}; + if (passed[s]) rows[i] = sw[word * {S} + s]; + }} +}} + +""" + +NOISE_DEVICE = r""" +kernel void noise( + device u64* out [[buffer(0)]], constant NoiseArgs& args [[buffer(1)]], + device const int* seg_first [[buffer(2)]], device const int* seg_outcomes [[buffer(3)]], + device const int* seg_begin [[buffer(4)]], device const int* seg_end [[buffer(5)]], + device const float* seg_log_keep [[buffer(6)]], device const Record* records [[buffer(7)]], + device const Flip* overflow [[buffer(8)]], + device const u64* constants [[buffer(9)]], device const u64* screen [[buffer(10)]], + device uint* queue [[buffer(11)]], device atomic_uint* count [[buffer(12)]], + device uint* survivors [[buffer(13)]], + threadgroup uint2* events [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], + uint nt [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]]) +{{ + threadgroup atomic_int num_events; + long base = (long)bid * {S}; + int ns = (int)min((long){S}, args.shots - base); + device u64* rows = out + base * {NW}; + for (int i = tid; i < ns * {NW}; i += nt) rows[i] = 0; + if (tid == 0) atomic_store_explicit(&num_events, 0, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup); + Segments seg = {{seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep}}; + add_faults(rows, events, &num_events, args.max_events, ns, + args.seed ^ ((u64)bid * 0xD1B54A32D192ED03UL), args.num_segments, seg, records, + overflow, tid, nt); + for (int s0 = 0; s0 < {S}; s0 += nt) {{ + int s = s0 + tid; + bool alive = false; + if (s < ns) {{ + u64 hit = 0; + device u64* row = rows + (long)s * {NW}; + for (int i = 0; i < {OW}; ++i) {{ + u64 v = row[i] ^ constants[i]; + hit |= v & screen[i]; + row[i] = v; + }} + alive = !hit; + }} + uint ballot = uint(static_cast(simd_ballot(alive))); + int warp_first = s0 + (tid & ~31u); + if (args.has_survivors && lane == 0 && warp_first < ns) + survivors[(base + warp_first) >> 5] = ballot; + if (ballot) {{ + uint leader = ctz(ballot), b = 0; + if (lane == leader) b = atomic_fetch_add_explicit(count, popcount(ballot), memory_order_relaxed); + b = simd_broadcast(b, (ushort)leader); + if (alive) queue[b + popcount(ballot & ((1u << lane) - 1))] = (uint)(base + s); + }} + }} +}} + +""" + +NOISE_KEPT = r""" +struct KeptArgs {{ u64 seed; int num_segments; int max_events; }}; + +// The deferred channels' faults on the kept shots ``index`` (``*n`` of them), +// then their returned outputs bit-packed (``pack_rows``). +kernel void noise_kept( + device const u64* out [[buffer(0)]], constant KeptArgs& args [[buffer(1)]], + device const int* seg_first [[buffer(2)]], device const int* seg_outcomes [[buffer(3)]], + device const int* seg_begin [[buffer(4)]], device const int* seg_end [[buffer(5)]], + device const float* seg_log_keep [[buffer(6)]], device const Record* records [[buffer(7)]], + device const Flip* overflow [[buffer(8)]], + device const uint* n [[buffer(9)]], device const uint* index [[buffer(10)]], + device uchar* dst [[buffer(11)]], device const uint2* pieces [[buffer(12)]], + device const uint* piece_start [[buffer(13)]], device const uint* cursor [[buffer(14)]], + threadgroup uint2* events [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], + uint nt [[threads_per_threadgroup]]) +{{ + threadgroup u64 sw[{OW} * {S}]; + threadgroup atomic_int num_events; + long base = (long)bid * {S}; + if (base >= (long)*n) return; + // NOTE: rows past the call's capacity (cursor[1]) are not written. + int ns = (int)min(min((long){S}, (long)*n - base), (long)cursor[1] - (long)cursor[0] - base); + if (ns <= 0) return; + for (int i = tid; i < ns * {OW}; i += nt) {{ + int s = i / {OW}, word = i - s * {OW}; + sw[word * {S} + s] = out[(long)index[base + s] * {NW} + word]; + }} + if (tid == 0) atomic_store_explicit(&num_events, 0, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + Segments seg = {{seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep}}; + add_faults(sw, events, &num_events, args.max_events, ns, + args.seed ^ ((u64)bid * 0xD1B54A32D192ED03UL), args.num_segments, seg, records, + overflow, tid, nt); + pack_rows(sw, ns, pieces, piece_start, (threadgroup uchar*)events, dst, ((long)*cursor + base) * {NB}, tid, nt); +}} +""" + +SOURCE = r""" +struct Record {{ uint alias; uint words; uint extra; uint pad; u64 m0; u64 m1; }}; +struct Flip {{ uint word; uint pad; u64 mask; }}; + +inline void xor64(threadgroup u64* sw, int index, u64 mask) {{ + threadgroup atomic_uint* halves = (threadgroup atomic_uint*)(sw + index); + uint lo = uint(mask), hi = uint(mask >> 32); + if (lo) atomic_fetch_xor_explicit(halves, lo, memory_order_relaxed); + if (hi) atomic_fetch_xor_explicit(halves + 1, hi, memory_order_relaxed); +}} + +inline void xor64(device u64* sw, int index, u64 mask) {{ + device atomic_uint* halves = (device atomic_uint*)(sw + index); + uint lo = uint(mask), hi = uint(mask >> 32); + if (lo) atomic_fetch_xor_explicit(halves, lo, memory_order_relaxed); + if (hi) atomic_fetch_xor_explicit(halves + 1, hi, memory_order_relaxed); +}} + +// The fault drawn at record ``rec`` (bits 0-23 of ``draw`` its fraction, bits +// 24-31 its slot in the channel): the slot's outcome or its alias's. Word ``w`` +// of row ``s`` is ``sw[w * WS + s * SS]``. +template +inline void apply_fault(W sw, int row, uint rec, uint draw, + device const Record* records, device const Flip* overflow) +{{ + Record r = records[rec]; + if ((draw & 0xFFFFFF) >= (r.alias & 0x1FFFFFF)) r = records[rec - (draw >> 24) + (r.alias >> 25)]; + int flips = r.extra & 255; + if (flips > 0) xor64(sw, (r.words & 0xFFFF) * WS + row * SS, r.m0); + if (flips > 1) xor64(sw, (r.words >> 16) * WS + row * SS, r.m1); + for (int j = 2; j < flips; ++j) {{ + Flip f = overflow[(r.extra >> 8) + j - 2]; + xor64(sw, f.word * WS + row * SS, f.mask); + }} +}} + +// The returned outputs of the block's ``ns`` rows ``sw[w * S + s]``, bit-packed +// into {NB} bytes per row, contiguously at ``dst`` (``gpu.rows.pieces``). +inline void pack_rows(threadgroup const u64* sw, int ns, device const uint2* pieces, + device const uint* start, threadgroup uchar* stage, device uchar* dst, long at, uint tid, uint nt) +{{ + for (uint i = tid; i < {NK} * {S}; i += nt) {{ + int k = i / {S}, s = i - k * {S}; + if (s >= ns) continue; + u64 v = 0; + for (uint p = start[k]; p < start[k + 1]; ++p) {{ + uint2 q = pieces[p]; + u64 x = sw[q.x * {S} + s] >> (q.y & 63); + int length = ((q.y >> 6) & 63) + 1; + if (length < 64) x &= (1UL << length) - 1; + v |= x << (q.y >> 12); + }} + threadgroup uchar* at = stage + s * {NB} + 8 * k; + for (int j = 0; j < 8; ++j) if (8 * k + j < {NB}) at[j] = uchar(v >> (8 * j)); + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + // NOTE: the rows go to byte ``at`` of the (aligned) buffer ``dst``: bytes up + // to a 16-byte boundary, then 16 at a time, then the rest. + int total = ns * {NB}; + int head = min(total, (int)((16 - (at & 15)) & 15)); + device uchar* d = dst + at; + for (int i = tid; i < head; i += nt) d[i] = stage[i]; + int body = (total - head) / 16; + device uint4* aligned = (device uint4*)(d + head); + if ((head & 3) == 0) {{ + threadgroup const uint* from = (threadgroup const uint*)(stage + head); + for (int i = tid; i < body; i += nt) + aligned[i] = uint4(from[4 * i], from[4 * i + 1], from[4 * i + 2], from[4 * i + 3]); + }} else {{ + for (int i = tid; i < body; i += nt) {{ + threadgroup const uchar* b = stage + head + 16 * i; + uint4 v; + for (int k = 0; k < 4; ++k) + v[k] = uint(b[4 * k]) | uint(b[4 * k + 1]) << 8 | uint(b[4 * k + 2]) << 16 | uint(b[4 * k + 3]) << 24; + aligned[i] = v; + }} + }} + for (int i = head + body * 16 + tid; i < total; i += nt) d[i] = stage[i]; +}} + +struct Segments {{ + device const int* first; device const int* outcomes; device const int* begin; + device const int* end; device const float* log_keep; +}}; + +// One segment's faults. With ``FINE``, for segments of fewer than 2^-12 expected +// faults, a draw within 2^-12 of u = 1 takes 24 more bits, which ``unit`` alone +// cannot resolve: it would fire such a segment at least 2^-24 of the time. +template +inline void add_segment_faults(W sw, threadgroup uint2* events, threadgroup atomic_int* num_events, + int max_events, int ns, u64 state, float lk, int first, int outcomes, int idx, int end, + device const Record* records, device const Flip* overflow) +{{ + while (true) {{ + state += GOLDEN; + u64 r = mix64(state); + float skip; + if (!FINE) skip = lk < -3.0e38f ? 0.f : floor(fast::log(unit(r)) / lk); + else if ((~r >> 40) >= 4096) skip = floor(fast::log(unit(r)) / lk); + else {{ + // NOTE: 1 - u = y = (j + w) 2^-24 <= 2^-12, with w from a fresh draw; + // Metal has no log1p, and the series to y^3 is exact to ~1e-11 here. + float w = float((mix64(state ^ 0x5851F42D4C957F2DUL) >> 40) + 1) * 0x1p-24f; + float y = (float(~r >> 40) + w) * 0x1p-24f; + skip = floor(-y * (1.f + y * (0.5f + y * (1.f / 3.f))) / lk); + }} + if (skip >= (float)(end - 1 - idx)) break; + idx += 1 + (int)skip; + int ch = idx / {S}, row = idx - ch * {S}; + if (row >= ns) continue; // the last block's missing shots + // NOTE: the skip reads the high 24 bits of r, the outcome its low 32: + // a slot from the high 24 of those, the fraction of the slot below. + u64 x = (u64)((uint)r >> 8) * (uint)outcomes; + uint slot = (uint)(x >> 24); + uint rec = (uint)(first + ch * outcomes) + slot; + uint draw = (uint)(x & 0xFFFFFF) | (slot << 24); + int e = atomic_fetch_add_explicit(num_events, 1, memory_order_relaxed); + if (e < max_events) events[e] = uint2((uint)row | (rec << 8), draw); + else apply_fault(sw, row, rec, draw, records, overflow); + }} +}} + +template +__attribute__((noinline)) void add_fine_segment_faults(W sw, threadgroup uint2* events, + threadgroup atomic_int* num_events, int max_events, int ns, u64 state, float lk, int first, + int outcomes, int idx, int end, device const Record* records, device const Flip* overflow) +{{ + add_segment_faults(sw, events, num_events, max_events, ns, state, lk, first, + outcomes, idx, end, records, overflow); +}} + +// Phases 1 and 2: draw the segments' faults on the block's ``ns`` shots and XOR +// them into ``sw`` (laid out as in ``apply_fault``; in device memory if ``DEV``). +template +inline void add_faults(W sw, threadgroup uint2* events, threadgroup atomic_int* num_events, + int max_events, int ns, u64 stream, int num_segments, Segments seg, + device const Record* records, device const Flip* overflow, uint tid, uint nt) +{{ + for (int sg = tid; sg < num_segments; sg += nt) {{ + float lk = seg.log_keep[sg]; + int first = seg.first[sg], outcomes = seg.outcomes[sg], end = seg.end[sg]; + int idx = seg.begin[sg] - 1; + u64 state = mix64(stream ^ ((u64)sg * GOLDEN)); + if ((float)(end - 1 - idx) * -lk < 0x1p-12f) + add_fine_segment_faults(sw, events, num_events, max_events, ns, state, lk, + first, outcomes, idx, end, records, overflow); + else + add_segment_faults(sw, events, num_events, max_events, ns, state, + lk, first, outcomes, idx, end, records, overflow); + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + int ne = min(atomic_load_explicit(num_events, memory_order_relaxed), max_events); + for (int e = tid; e < ne; e += nt) {{ + uint2 v = events[e]; + apply_fault(sw, v.x & 255, v.x >> 8, v.y, records, overflow); + }} + if (DEV) threadgroup_barrier(mem_flags::mem_device); + else threadgroup_barrier(mem_flags::mem_threadgroup); +}} + +struct NoiseArgs {{ long shots; u64 seed; int num_segments; int max_events; int has_survivors; }}; + +{NOISE}{KEPT}""" + + +def noise_tables( + channels: ChannelSampler, flips: WordTables, num_words: int +) -> NoiseTables: + """Return the noise kernels' tables, for Metal's threadgroup memory. + + Rows too wide for ``SHARED_BYTES`` keep their words in device memory + (``shared_words`` False), where the noise kernel XORs the faults in. + """ + if 8 * 32 * num_words <= SHARED_BYTES - 8 * 256: + return _noise_tables(channels, flips, num_words, shared_bytes=SHARED_BYTES) + return _noise_tables( + channels, flips, num_words, shared_bytes=SHARED_BYTES, shared_words=False + ) + + +def source(tables: NoiseTables, output_words: int, packed_bytes: int) -> str: + """Return the MSL source of the ``noise`` and ``noise_kept`` kernels. + + Without ``tables.shared_words`` there is no ``noise_kept``: the rows of a + block do not fit in threadgroup memory. + """ + if tables.shared_words: + noise, kept = NOISE_SHARED, NOISE_KEPT + else: + noise, kept = NOISE_DEVICE, "" + return SOURCE.replace("{NOISE}{KEPT}", noise + kept).format( + NW=tables.num_words, + OW=output_words, + S=tables.block_shots, + NB=max(packed_bytes, 1), + NK=max(-(-packed_bytes // 8), 1), + ) diff --git a/src/tsim/statevector/metal/runtime.py b/src/tsim/statevector/metal/runtime.py new file mode 100644 index 00000000..51ac30f3 --- /dev/null +++ b/src/tsim/statevector/metal/runtime.py @@ -0,0 +1,309 @@ +"""Metal through PyObjC for the Metal executor, loaded on first use. + +Buffers are ``MTLStorageModeShared``: on Apple silicon the CPU and the GPU share +memory, so a buffer is also a NumPy array (``Buffer.array``) and nothing is +copied between them. +""" + +from __future__ import annotations + +import functools +import os +from dataclasses import dataclass +from types import ModuleType +from typing import Any + +import numpy as np + + +@functools.cache +def metal() -> ModuleType: + """Return PyObjC's ``Metal`` module. + + Raises: + ImportError: If ``pyobjc-framework-Metal`` is not installed. + + """ + import Metal + + return Metal + + +@functools.cache +def device() -> Any: + """Return the system's default Metal device. + + Raises: + RuntimeError: If there is none. + + """ + global _initialized + found = metal().MTLCreateSystemDefaultDevice() + if found is None: + raise RuntimeError("no Metal device") + _initialized = True + return found + + +_initialized = False +_forked_after_init = False + + +def _after_fork_in_child() -> None: + # NOTE: Metal's device and its connections to the GPU driver do not survive a + # fork; a child of a process that used them samples on the CPU. + global _forked_after_init + if _initialized: + _forked_after_init = True + + +os.register_at_fork(after_in_child=_after_fork_in_child) + + +def forked_after_init() -> bool: + """Return whether this process was forked after its parent initialized Metal.""" + return _forked_after_init + + +@functools.cache +def queue() -> Any: + """Return the executor's command queue.""" + return device().newCommandQueue() + + +@functools.cache +def check_simd_width() -> None: + """Check that the device runs SIMD groups of 32 threads, as the kernels assume. + + Raises: + RuntimeError: If it does not (e.g. an Intel or AMD GPU). + + """ + if not device().supportsFamily_(metal().MTLGPUFamilyApple7): + raise RuntimeError("the Metal executor needs an Apple GPU (Apple silicon)") + probe = Library( + "#include \nusing namespace metal;\n" + "kernel void probe(device uint* out [[buffer(0)]]) { out[0] = 0; }" + ).pipeline("probe") + width = int(probe.threadExecutionWidth()) + if width != 32: + raise RuntimeError( + f"the Metal executor needs SIMD groups of 32 threads, not {width}" + ) + + +def available() -> bool: + """Return whether PyObjC's Metal is installed and there is a usable Apple GPU. + + It is not in a process forked after its parent used the GPU. + """ + return not _forked_after_init and _device_available() + + +@functools.cache +def _device_available() -> bool: + try: + check_simd_width() + return True + except Exception: + return False + + +@dataclass +class Buffer: + """A shared Metal buffer and a NumPy view of its memory.""" + + handle: Any + array: np.ndarray + + +POOL_BUFFERS = 2 +"""Freed output buffers kept for reuse. + +The GPU's first write to each page of a fresh buffer costs the OS a fault: for +rows of 440 bytes a call runs 1.7x longer than into a buffer that a previous +call returned (and its caller dropped). +""" + +_pool: list[Any] = [] + + +class _Owner: + """Holds a Metal buffer for the NumPy arrays over its memory (their ``base``). + + A ``recycle`` buffer goes back to the pool when they are gone. + """ + + def __init__(self, handle: Any, nbytes: int, recycle: bool = False) -> None: + self.handle = handle + self.recycle = recycle + view = handle.contents().as_buffer(nbytes) + address = np.frombuffer(view, dtype=np.uint8).ctypes.data + self.__array_interface__ = { + "shape": (nbytes,), + "typestr": "|u1", + "data": (address, False), + "version": 3, + } + + def __del__(self) -> None: + if self.recycle and _pool is not None: + _pool.append(self.handle) + del _pool[:-POOL_BUFFERS] + + +def buffer(count: int, dtype, recycle: bool = False) -> Buffer: + """Return a shared buffer of ``count`` elements of ``dtype`` (zeroed). + + Its array keeps the Metal buffer alive, so it may outlive the ``Buffer``. + With ``recycle`` the buffer may be a freed one of up to twice the size, not + zeroed, and returns to the pool once its arrays are gone. + """ + dtype = np.dtype(dtype) + nbytes = max(int(count) * dtype.itemsize, 16) + handle = None + if recycle: + for k, pooled in enumerate(_pool): + if nbytes <= pooled.length() <= 2 * nbytes: + handle = _pool.pop(k) + break + if handle is None: + handle = device().newBufferWithLength_options_( + nbytes, metal().MTLResourceStorageModeShared + ) + if handle is None: + raise MemoryError( + f"no Metal buffer of {nbytes} bytes (at most " + f"{device().maxBufferLength()}); sample fewer shots per call" + ) + owner = _Owner(handle, nbytes, recycle) + array = np.asarray(owner)[: count * dtype.itemsize].view(dtype) + return Buffer(handle, array) + + +def upload(values: np.ndarray) -> Buffer: + """Return a shared buffer holding a copy of ``values`` (flattened).""" + values = np.ascontiguousarray(values) + held = buffer(values.size, values.dtype) + held.array[:] = values.reshape(-1) + return held + + +class Library: + """A compiled Metal library and its kernels' pipeline states.""" + + def __init__(self, source: str, fast_math: bool = True) -> None: + """Compile ``source`` (Metal Shading Language). + + Raises: + RuntimeError: If it does not compile. + + """ + mtl = metal() + options = mtl.MTLCompileOptions.new() + options.setFastMathEnabled_(fast_math) + library, error = device().newLibraryWithSource_options_error_( + source, options, None + ) + if library is None: + raise RuntimeError(f"Metal compilation failed: {error}") + self._library = library + self._pipelines: dict[str, Any] = {} + + def pipeline(self, name: str) -> Any: + """Return the compute pipeline state of kernel ``name``.""" + if name not in self._pipelines: + function = self._library.newFunctionWithName_(name) + if function is None: + raise KeyError(name) + descriptor = metal().MTLComputePipelineDescriptor.new() + descriptor.setComputeFunction_(function) + descriptor.setLabel_(name) + ( + state, + error, + ) = device().newComputePipelineStateWithDescriptor_options_reflection_error_( + descriptor, 0, None, None + ) + if state is None: + raise RuntimeError(f"pipeline {name}: {error}") + self._pipelines[name] = state + return self._pipelines[name] + + +Argument = Buffer | tuple[Buffer, int] | np.ndarray | np.generic + +PROFILE: dict[str, float] | None = None +"""If a dict, every dispatch runs as its own command buffer and adds its GPU time +(seconds) under its kernel's name: for profiling only.""" + + +class Encoder: + """Records kernel dispatches into one command buffer.""" + + def __init__(self) -> None: + """Start a command buffer with a compute encoder.""" + self.commands = queue().commandBuffer() + self._encoder = self.commands.computeCommandEncoder() + + def dispatch( + self, + pipeline: Any, + arguments: list[Argument], + threadgroups: int, + threads: int, + threadgroup_bytes: int = 0, + ) -> None: + """Dispatch ``threadgroups`` groups of ``threads`` threads of ``pipeline``. + + Each argument binds buffer index ``i``: a buffer (or a buffer and a byte + offset), or a NumPy scalar or array passed by value. A nonzero + ``threadgroup_bytes`` binds threadgroup memory 0. + """ + mtl, encoder = metal(), self._encoder + encoder.setComputePipelineState_(pipeline) + for i, argument in enumerate(arguments): + if isinstance(argument, Buffer): + encoder.setBuffer_offset_atIndex_(argument.handle, 0, i) + elif isinstance(argument, tuple): + encoder.setBuffer_offset_atIndex_(argument[0].handle, argument[1], i) + else: + data = np.ascontiguousarray(argument).tobytes() + encoder.setBytes_length_atIndex_(data, len(data), i) + if threadgroup_bytes: + encoder.setThreadgroupMemoryLength_atIndex_( + -(-threadgroup_bytes // 16) * 16, 0 + ) + if threadgroups: + encoder.dispatchThreadgroups_threadsPerThreadgroup_( + mtl.MTLSizeMake(threadgroups, 1, 1), mtl.MTLSizeMake(threads, 1, 1) + ) + if PROFILE is not None: + name = str(pipeline.label() or "") or str(pipeline) + commands = self.commit() + PROFILE[name] = PROFILE.get(name, 0.0) + wait(commands) + self.commands = queue().commandBuffer() + self._encoder = self.commands.computeCommandEncoder() + + def commit(self) -> Any: + """End encoding and commit; return the command buffer.""" + self._encoder.endEncoding() + self.commands.commit() + return self.commands + + +def wait(commands: Any) -> float: + """Wait for a committed command buffer; return its GPU time in seconds. + + Raises: + MemoryError: If the GPU ran out of memory. + RuntimeError: If it failed otherwise. + + """ + commands.waitUntilCompleted() + error = commands.error() + if error is not None: + if error.code() == metal().MTLCommandBufferErrorOutOfMemory: + raise MemoryError(f"the GPU ran out of memory: {error}") + raise RuntimeError(f"Metal command buffer failed: {error}") + return float(commands.GPUEndTime() - commands.GPUStartTime()) diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py new file mode 100644 index 00000000..c48b1a5b --- /dev/null +++ b/src/tsim/statevector/metal/sampler.py @@ -0,0 +1,992 @@ +"""``MetalShotSampler``: ``ShotSampler`` with the noise and the tapes on a Metal GPU. + +A batch of shots is one command buffer of these kernels, which run one after +another: + + noise draws every shot's noise into its words and screens it; the screened + survivors go to a queue + tape runs every component on the queued shots, a SIMD group of ``2^tb`` + threads per shot with its amplitudes in registers; with a + postselection mask it clears the rejected shots from the screen's + bitmap of survivors + list lists the accepted shots in order from the bitmap + pack gathers the accepted rows, bit-packed (or draws the deferred noise + on them first, ``noise_kept``) + +Every buffer is shared with the CPU (unified memory), so nothing is copied: +the packing kernels write each batch's rows into the array that the call +returns, after the rows of the batches before it (a cursor in the GPU's memory +counts them). The host never waits between batches: it commits each batch's +command buffer as it encodes it, the batches reuse one set of buffers (so +Metal runs them one after another), and the GPU resets their counters itself. + +The kernels are compiled once per postselection mask (Metal caches compiled +shaders across processes). Amplitudes are complex64, so a sample follows the +CPU executor's distribution to within float32 rounding of the outcome +probabilities; the noise and the outcome draws come from other streams than +the CPU's, so a seed's samples differ from the CPU's, and they depend on the +batch size. +""" + +from __future__ import annotations + +import time +import warnings +from dataclasses import dataclass, field +from typing import Any + +import numpy as np + +from tsim.noise.channels import ChannelSampler +from tsim.statevector.gpu.noise import NoiseGroups, deferred_channels, pass_rate +from tsim.statevector.gpu.rows import pieces as _pieces, select as _select +from tsim.statevector.metal import runtime +from tsim.statevector.metal.codegen import PRELUDE, hex64, tape_function +from tsim.statevector.metal.noise import noise_tables, source as noise_source +from tsim.statevector.shared import pack_bits +from tsim.statevector.shots import ShotLayout, ShotPlan, ShotSampler + +MAX_WIDTH = 11 +"""Widest component the Metal executor runs: 32 threads of 64 amplitudes.""" + +BATCH_SHOTS = 1 << 20 +"""Shots of one batch (one command buffer).""" + +AMPLITUDE_AXES = 5 +"""Log2 of the amplitudes per thread from which a shot takes more threads.""" + +BATCH_BYTES = 256 << 20 +"""Bytes of shot words of one batch, which bound its shots below ``BATCH_SHOTS``.""" + +NOISE_THREADS = 256 +"""Threads per threadgroup of the noise kernel.""" + +TAPE_THREADS = 128 +"""Threads per threadgroup of the tape kernel.""" + +FAST_MATH = True +"""Whether the kernels compile with fast math.""" + +DEVIATION_TOLERANCE = 1e-3 +"""Normalization deviation that float32 amplitudes tolerate without a warning.""" + +DEFER_SHARE = 0.25 +"""Share of the expected faults that postselection cannot see, from which they +are drawn on the accepted shots only, after the tapes.""" + +_KERNELS = r""" +struct TapeArgs {{ long first_shot; u64 seed; int flags; }}; + +// flags: 1 clears rejected shots from ``survivors`` and counts the accepted ones, +// 2 writes the expectations. +kernel void tape( + device u64* out [[buffer(0)]], constant TapeArgs& args [[buffer(1)]], + device const uint* queue [[buffer(2)]], device const uint* count [[buffer(3)]], + device atomic_uint* survivors [[buffer(4)]], device atomic_uint* acc_count [[buffer(5)]], + device float* expectations [[buffer(6)]], device atomic_uint* deviation [[buffer(7)]], + uint gid [[thread_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) +{{ + int t = lane & {T1}; + uint q = gid >> {tb}; + if (q >= count[0]) return; + const u64 seed = args.seed; +{seeds} + uint row_index = queue[q]; + u64 shot = (u64)args.first_shot + row_index; + device u64* row = out + (long)row_index * {NW}; + u64 w[{NW}]; + for (int i = 0; i < {NW}; ++i) w[i] = row[i]; + float ev[{NE}]; + int grave = 0x7fffffff; +{calls} + if (t == 0) {{ +{stores} + if (args.flags & 2) + for (int i = 0; i < {NO}; ++i) expectations[(long)row_index * {NO} + i] = ev[i]; + }} + if (args.flags & 1) {{ + u64 hit = {hits}; + if (t == 0 && hit) + atomic_fetch_and_explicit(&survivors[row_index >> 5], ~(1u << (row_index & 31)), memory_order_relaxed); + uint ballot = uint(static_cast(simd_ballot(t == 0 && hit == 0))); + if (ballot && simd_is_first()) atomic_fetch_add_explicit(acc_count, popcount(ballot), memory_order_relaxed); + }} +}} + +// The counts of a batch: its screened survivors and its accepted shots. +kernel void reset(device uint* counts [[buffer(0)]]) {{ counts[0] = 0; counts[1] = 0; }} + +// Moves the cursor past a batch's ``*n`` rows. +kernel void advance(device uint* cursor [[buffer(0)]], device const uint* n [[buffer(1)]]) {{ + cursor[0] += *n; +}} + +inline uint block_scan(uint v, thread uint& total, threadgroup uint* warps, uint lane, uint warp, uint nwarps) {{ + uint x = simd_prefix_inclusive_sum(v); + if (lane == 31) warps[warp] = x; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (warp == 0) {{ + uint y = lane < nwarps ? warps[lane] : 0; + warps[lane] = simd_prefix_inclusive_sum(y); + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + uint before = (warp ? warps[warp - 1] : 0) + x - v; + total = warps[nwarps - 1]; + threadgroup_barrier(mem_flags::mem_threadgroup); + return before; +}} + +// The shots of the set bits of ``bits`` (``words`` words), in order: ``count_bits`` +// counts each block's 1024 words, ``list_bits`` lists a block's shots after +// those of the blocks before it. +kernel void count_bits( + device const uint* bits [[buffer(0)]], constant long& words [[buffer(1)]], + device uint* totals [[buffer(2)]], + uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], uint warp [[simdgroup_index_in_threadgroup]], + uint nwarps [[simdgroups_per_threadgroup]]) +{{ + threadgroup uint warps[32]; + long i = (long)bid * 1024 + tid; + uint total; + block_scan(i < words ? popcount(bits[i]) : 0, total, warps, lane, warp, nwarps); + if (tid == 0) totals[bid] = total; +}} + +kernel void list_bits( + device const uint* bits [[buffer(0)]], constant long& words [[buffer(1)]], + device const uint* totals [[buffer(2)]], device uint* index [[buffer(3)]], + uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], uint warp [[simdgroup_index_in_threadgroup]], + uint nwarps [[simdgroups_per_threadgroup]]) +{{ + threadgroup uint warps[32]; + uint earlier = 0, total; + for (uint b = tid; b < bid; b += 1024) earlier += totals[b]; + block_scan(earlier, earlier, warps, lane, warp, nwarps); + long i = (long)bid * 1024 + tid; + uint w = i < words ? bits[i] : 0; + uint at = earlier + block_scan(popcount(w), total, warps, lane, warp, nwarps); + for (; w; w &= w - 1) index[at++] = (uint)(32 * i + ctz(w)); +}} + +// The rows ``index[j]`` (or every row, without ``has_index``) for ``j < *n``, +// densely after the ``*cursor`` rows of the call's earlier batches: each row's +// first {NB} bytes, which hold its outputs bit-packed (the bits past them +// cleared). A thread writes 4 aligned bytes at a time. +kernel void gather_rows( + device const uchar* rows [[buffer(0)]], device const uint* n [[buffer(1)]], + device const uint* index [[buffer(2)]], constant int& has_index [[buffer(3)]], + device uchar* dst [[buffer(4)]], device const uint* cursor [[buffer(5)]], + uint gid [[thread_position_in_grid]], uint grid [[threads_per_grid]]) +{{ + long count = min((long)*n, (long)cursor[1] - (long)cursor[0]); + if (count <= 0) return; + long at = (long)cursor[0] * {NB}, end = at + count * {NB}; + for (long p = (at & ~3L) + 4L * gid; p < end; p += 4L * grid) {{ + // NOTE: a batch's rows span less than 4 GB (``BATCH_BYTES``). + uint k = (uint)(max(p, at) - at); + uint r = k / {NB}, b = k - r * {NB}; + uint sh = 8 * (b & 3); + if (p >= at && p + 4 <= end && b + 4 <= {NB} && (sh == 0 || (b & ~3u) + 8 <= {RB})) {{ + // NOTE: the 4 bytes lie in one row: two aligned loads at most. + long source = has_index ? (long)index[r] : r; + device const uint* w = (device const uint*)(rows + source * {RB}) + (b >> 2); + uint v = sh ? (w[0] >> sh) | (w[1] << (32 - sh)) : w[0]; + if (b + 4 == {NB}) v &= ({LAST}u << 24) | 0xFFFFFFu; + *(device uint*)(dst + p) = v; + continue; + }} + uint v = 0; + for (int j = (int)(max(p, at) - p); j < 4 && p + j < end; ++j) {{ + long source = has_index ? (long)index[r] : r; + uint x = rows[source * {RB} + b]; + if (b == {NB} - 1) x &= {LAST}; + v |= x << (8 * j); + if (++b == {NB}) {{ + b = 0; + ++r; + }} + }} + for (int j = 0; j < 4; ++j) + if (p + j >= at && p + j < end) dst[p + j] = uchar(v >> (8 * j)); + }} +}} + +// The first ``n`` packed rows of ``{NB}`` bytes, one byte per output (``{NOUT}`` +// per row), as booleans: a thread unpacks 4 outputs, which share a packed byte. +kernel void unpack_rows( + device const uchar* packed [[buffer(0)]], device uchar* dst [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint grid [[threads_per_grid]]) +{{ + const uint groups = ({NOUT} + 3) / 4; + for (uint i = gid; i < n * groups; i += grid) {{ + uint r = i / groups, g = i - r * groups; + uint b = packed[(ulong)r * {NB} + (g >> 1)] >> (4 * (g & 1)); + device uchar* out = dst + (ulong)r * {NOUT} + 4 * g; + for (uint j = 0; j < 4 && 4 * g + j < {NOUT}; ++j) out[j] = uchar((b >> j) & 1); + }} +}} + +// 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. +kernel void gather_float( + device const float* values [[buffer(0)]], device const uint* n [[buffer(1)]], + constant int& width [[buffer(2)]], device const uint* index [[buffer(3)]], + device float* dst [[buffer(4)]], device const uint* cursor [[buffer(5)]], + constant int& has_index [[buffer(6)]], + uint gid [[thread_position_in_grid]], uint grid [[threads_per_grid]]) +{{ + long total = min((long)*n, (long)cursor[1] - (long)cursor[0]) * width; + dst += (long)*cursor * width; + for (long k = gid; k < total; k += grid) {{ + long j = k / width; int c = (int)(k - j * width); + long r = has_index ? (long)index[j] : j; + dst[k] = values[r * width + c]; + }} +}} +""" + + +@dataclass +class _Launch: + """The device tables of one noise launch: its segments and records.""" + + segments: list[runtime.Buffer] + records: runtime.Buffer + num_segments: int + max_events: int + + +@dataclass +class _Module: + """A plan's compiled kernels and its device constants.""" + + library: runtime.Library + tb: int + has_mask: bool + columns: np.ndarray + """The outputs that ``sample_accepted`` returns: all of them, in order.""" + pieces: runtime.Buffer + piece_start: runtime.Buffer + screen: runtime.Buffer + plan: ShotPlan + psm_words: np.ndarray + compile_seconds: float + passing: float + """Estimated share of shots that pass the screen.""" + splits: dict[bool, tuple[_Launch, _Launch] | None] = field(default_factory=dict) + + +@dataclass +class _Slot: + """The buffers that a call's batches reuse, one batch after another. + + ``counts`` holds the batch's screened survivors, its accepted shots and the + call's largest normalization deviation (as float bits), then the call's + cursor: the rows that its batches have packed, and the rows that its + output has room for. + """ + + shots: int + words: runtime.Buffer + queue: runtime.Buffer + accepted: runtime.Buffer + counts: runtime.Buffer + totals: runtime.Buffer + index: runtime.Buffer + values: runtime.Buffer | None = None + + +class MetalShotSampler: + """Samples shot words of a statevector program on a Metal GPU. + + It keeps ``ShotSampler``'s interface (``sample``, ``direct_bools``, + ``reference``, ``plan``, ``layout``); ``sample`` returns only the shots that + pass postselection. ``sample_accepted`` can also return each observable's + expectation. + """ + + @staticmethod + def supports(layout: ShotLayout) -> bool: + """Return whether every component of ``layout`` fits in ``MAX_WIDTH`` qubits. + + A layout without shot words (no outputs and no noise) stays on the CPU. + """ + return layout.num_words > 0 and all( + c.width <= MAX_WIDTH for c in layout.components + ) + + def __init__( + self, + host: ShotSampler, + channels: ChannelSampler, + num_detectors: int, + seed: np.random.SeedSequence, + ) -> None: + """Upload the noise tables of ``host``'s shot words. + + ``host`` draws the reference samples, on the CPU. + + Raises: + ImportError: If PyObjC's Metal is not installed. + RuntimeError: If there is no Apple GPU. + + """ + runtime.check_simd_width() + assert self.supports(host.layout) + self._host = host + layout = host.layout + self._num_detectors = num_detectors + self._observables = np.arange(num_detectors, layout.num_outputs) + self._rng = np.random.default_rng(seed) + self._channels = channels + self._noise = noise_tables(channels, layout.noise, layout.num_words) + self._overflow = runtime.upload(self._noise.overflow.view(np.uint8)) + self._all = self._launch(self._noise.groups) + self._modules: dict[bytes, _Module] = {} + self._constants: dict[bytes, runtime.Buffer] = {} + self._slots: dict[tuple[int, bool], _Slot] = {} + self._shares: dict[tuple, float] = {} + self.gpu_seconds = 0.0 + """GPU time of the command buffers so far (for benchmarks).""" + + @property + def layout(self): + """Return the shot words' layout.""" + return self._host.layout + + def plan(self, mask: np.ndarray | None) -> ShotPlan: + """Return the plan for a postselection ``mask``.""" + return self._host.plan(mask) + + def bytes_per_shot(self) -> int: + """Return the host memory of one shot's words.""" + return 8 * self.layout.num_words + + def reference(self) -> tuple[np.ndarray, float]: + """Return a noiseless sample of every output (on the CPU executor).""" + return self._host.reference() + + # -- compilation ---------------------------------------------------------- + + def _threads_per_shot(self) -> int: + """Return log2 of the threads per shot: one per ``2^AMPLITUDE_AXES`` amplitudes.""" + width = max((c.width for c in self.layout.components), default=0) + return min(5, max(0, width - AMPLITUDE_AXES)) + + def _module(self, mask: np.ndarray | None) -> _Module: + key = b"" if mask is None else np.asarray(mask, bool).tobytes() + if key in self._modules: + return self._modules[key] + plan = self.plan(mask) + layout = self.layout + tb = self._threads_per_shot() + T = 1 << tb + observables: list[dict[int, int]] = [] + for component in layout.components: + slots = { + o: int(np.searchsorted(self._observables, index)) + for o, index in enumerate(component.output_indices) + if index >= self._num_detectors + } + observables.append(slots) + functions = "\n\n".join( + tape_function(program, i, tb, observables[i]) + for i, program in enumerate(plan.programs) + ) + calls = "\n".join( + f" grave = min(grave, component_{i}(w, shot, seed_{i}, t, ev, deviation));" + for i in range(len(plan.programs)) + ) + bits = np.zeros(64 * layout.output_words, dtype=bool) + if mask is not None: + bits[: len(mask)] = mask + psm_words = pack_bits(bits, layout.output_words) + # NOTE: the rows keep the postselected outputs, as the CPU's do. + columns = np.arange(layout.num_outputs, dtype=np.int32) + written = sorted({int(w) for p in plan.programs for w in p.out_words}) + stores = "\n".join(f" row[{w}] = w[{w}];" for w in written) + hits = " | ".join( + f"(w[{w}] & {hex64(psm_words[w])})" for w in written if psm_words[w] + ) + nb = -(-len(columns) // 8) + source = ( + PRELUDE + + noise_source(self._noise, layout.output_words, nb) + + functions + + _KERNELS.format( + T1=T - 1, + tb=tb, + NW=layout.num_words, + OW=layout.output_words, + NO=max(1, len(self._observables)), + NE=max(1, len(self._observables)), + S=self._noise.block_shots, + NB=max(nb, 1), + RB=8 * layout.num_words, + NOUT=max(layout.num_outputs, 1), + LAST=hex((1 << (layout.num_outputs % 8 or 8)) - 1), + calls=calls, + seeds="\n".join( + f" const u64 seed_{i} = mix64(seed + {i + 1}UL * GOLDEN);" + for i in range(len(plan.programs)) + ), + stores=stores, + hits=hits or "0", + ) + ) + start = time.perf_counter() + library = runtime.Library(source, fast_math=FAST_MATH) + names = ["noise", "tape", "count_bits", "list_bits", "gather_rows"] + names += ["gather_float", "reset", "advance", "unpack_rows"] + if self._noise.shared_words: + names.append("noise_kept") + for name in names: + library.pipeline(name) + compile_seconds = time.perf_counter() - start + pieces, piece_start = _pieces(columns) + compiled = _Module( + library=library, + tb=tb, + has_mask=mask is not None, + columns=columns, + pieces=runtime.upload(pieces), + piece_start=runtime.upload(piece_start), + screen=runtime.upload(plan.screen), + plan=plan, + psm_words=psm_words, + compile_seconds=compile_seconds, + passing=self.passing(mask), + ) + self._modules[key] = compiled + return compiled + + def _launch(self, groups: NoiseGroups) -> _Launch: + """Upload a noise launch's segments and records.""" + return _Launch( + segments=[runtime.upload(a) for a in groups.segments], + records=runtime.upload(groups.records.view(np.uint8)), + num_segments=len(groups.segments[0]), + max_events=groups.max_events, + ) + + def _split( + self, module: _Module, expectations: bool + ) -> tuple[_Launch, _Launch] | None: + """Return a postselected batch's noise passes (screened, deferred), or None. + + The deferred channels flip none of the bits that the screen, the tapes or + the postselection read: the postselected outputs, the keys, and with + ``expectations`` the components' observables. They are drawn apart when + they are at least ``DEFER_SHARE`` of the expected faults. + """ + # NOTE: rows too wide for threadgroup memory draw all their noise at once. + if not module.has_mask or not self._noise.shared_words: + return None + if expectations in module.splits: + return module.splits[expectations] + layout = self.layout + read = np.zeros(layout.num_words, dtype=np.uint64) + read[: layout.output_words] = module.psm_words | module.plan.screen + for program in module.plan.programs: + start = program.key_start + read[start : start + program.key_words] = ~np.uint64(0) + np.bitwise_or.at(read, program.ps_word, program.ps_mask.astype(np.uint64)) + if expectations: + bits = np.zeros(64 * layout.output_words, dtype=bool) + for component in layout.components: + for index in component.output_indices: + if index >= self._num_detectors: + bits[index] = True + read[: layout.output_words] |= pack_bits(bits, layout.output_words) + deferred = deferred_channels(self._channels, layout.noise, read) + later = self._noise.subset(deferred) + split = None + if later.faults > 0 and later.faults >= DEFER_SHARE * self._noise.groups.faults: + split = (self._launch(self._noise.subset(~deferred)), self._launch(later)) + module.splits[expectations] = split + return split + + def passing(self, mask: np.ndarray | None) -> float: + """Return the estimated share of shots that pass the screen of ``mask``.""" + layout = self.layout + screen = np.zeros(layout.num_words, dtype=np.uint64) + screen[: layout.output_words] = self.plan(mask).screen + return pass_rate(self._channels, layout.noise, screen) if screen.any() else 1.0 + + # -- execution -------------------------------------------------------------- + + def _constant(self, flip: np.ndarray | None) -> runtime.Buffer: + """Return the outputs' constants XOR the reference bits ``flip``.""" + key = b"" if flip is None else np.asarray(flip, bool).tobytes() + if key not in self._constants: + self._constants[key] = runtime.upload(self._host._constant(flip)) + return self._constants[key] + + def _batch_shots(self, shots: int, batch_size: int | None) -> int: + """Return the shots of a call's batches.""" + cap = max(1 << 16, BATCH_BYTES // (8 * self.layout.num_words)) + return max(1, min(batch_size or BATCH_SHOTS, BATCH_SHOTS, cap, shots)) + + def _slot(self, shots: int, expectations: bool) -> _Slot: + """Return the buffers for batches of up to ``shots`` shots.""" + key = (shots, expectations) + slot = self._slots.get(key) + if slot is None: + nw, no = self.layout.num_words, max(1, len(self._observables)) + slot = _Slot( + shots=shots, + words=runtime.buffer(shots * nw, np.uint64), + queue=runtime.buffer(shots, np.uint32), + accepted=runtime.buffer(-(-shots // 32), np.uint32), + counts=runtime.buffer(8, np.uint32), + totals=runtime.buffer(-(-shots // 32 // 1024) + 1, np.uint32), + index=runtime.buffer(shots, np.uint32), + ) + if expectations: + slot.values = runtime.buffer(shots * no, np.float32) + self._slots = {k: v for k, v in self._slots.items() if k[1] != expectations} + self._slots[key] = slot + return slot + + def _encode_batch( + self, + encoder: runtime.Encoder, + module: _Module, + slot: _Slot, + shots: int, + first_shot: int, + constant: runtime.Buffer, + expectations: bool, + launch: _Launch, + ) -> None: + """Encode a batch's noise and tape kernels. + + Only the shots that pass the screen get their words written. + """ + library = module.library + words = slot.words + encoder.dispatch(library.pipeline("reset"), [slot.counts], 1, 1) + seed = np.uint64(self._rng.integers(0, 2**64, dtype=np.uint64)) + noise_seed = np.uint64(self._rng.integers(0, 2**64, dtype=np.uint64)) + S = self._noise.block_shots + args = np.zeros(1, dtype=_NOISE_ARGS) + args[0] = ( + shots, + noise_seed, + launch.num_segments, + launch.max_events, + int(module.has_mask), + ) + encoder.dispatch( + library.pipeline("noise"), + [ + words, + args, + *launch.segments, + launch.records, + self._overflow, + constant, + module.screen, + slot.queue, + (slot.counts, 0), + slot.accepted, + ], + -(-shots // S), + NOISE_THREADS, + threadgroup_bytes=8 * max(launch.max_events, 1), + ) + if self.layout.components: + tape_args = np.zeros(1, dtype=_TAPE_ARGS) + flags = int(module.has_mask) | (2 if expectations else 0) + tape_args[0] = (first_shot, seed, flags, 0) + values = slot.values if slot.values is not None else slot.counts + threads = shots << module.tb + encoder.dispatch( + library.pipeline("tape"), + [ + words, + tape_args, + slot.queue, + (slot.counts, 0), + slot.accepted, + (slot.counts, 4), + values, + (slot.counts, 8), + ], + -(-threads // TAPE_THREADS), + TAPE_THREADS, + ) + + def _encode_pack( + self, + encoder: runtime.Encoder, + module: _Module, + slot: _Slot, + shots: int, + split: tuple[_Launch, _Launch] | None, + expectations: bool, + rows: runtime.Buffer, + values: runtime.Buffer | None, + fixed: tuple[runtime.Buffer, int], + ) -> None: + """Encode the listing and packing of a batch's accepted rows into ``rows``. + + They follow the rows of the call's earlier batches (at the cursor), and + their expectations go to ``values``; ``fixed`` holds the batch's shots, + which a call without a mask returns. + """ + library = module.library + S = self._noise.block_shots + nb = -(-len(module.columns) // 8) + no = len(self._observables) + if module.has_mask: + words = -(-shots // 32) + blocks = -(-words // 1024) + encoder.dispatch( + library.pipeline("count_bits"), + [slot.accepted, np.int64(words), slot.totals], + blocks, + 1024, + ) + encoder.dispatch( + library.pipeline("list_bits"), + [slot.accepted, np.int64(words), slot.totals, slot.index], + blocks, + 1024, + ) + count = (slot.counts, 4) if self.layout.components else (slot.counts, 0) + else: + count = fixed + cursor = (slot.counts, 12) + blocks = -(-shots // S) + if nb and split is not None: + launch = split[1] + kept_args = np.zeros(1, dtype=_KEPT_ARGS) + kept_args[0] = ( + np.uint64(self._rng.integers(0, 2**64, dtype=np.uint64)), + launch.num_segments, + launch.max_events, + ) + encoder.dispatch( + library.pipeline("noise_kept"), + [ + slot.words, + kept_args, + *launch.segments, + launch.records, + self._overflow, + count, + slot.index, + rows, + module.pieces, + module.piece_start, + cursor, + ], + blocks, + 256, + threadgroup_bytes=max(8 * launch.max_events, S * nb), + ) + elif nb: + threads = min(-(-shots * nb // 4) + 1, 1 << 20) + encoder.dispatch( + library.pipeline("gather_rows"), + [ + slot.words, + count, + slot.index, + np.int32(module.has_mask), + rows, + cursor, + ], + -(-threads // 256), + 256, + ) + if expectations and no and self.layout.components: + assert slot.values is not None and values is not None + encoder.dispatch( + library.pipeline("gather_float"), + [ + slot.values, + count, + np.int32(no), + slot.index, + values, + cursor, + np.int32(module.has_mask), + ], + 256, + 256, + ) + encoder.dispatch(library.pipeline("advance"), [cursor, count], 1, 1) + + def _deviation(self, counts: np.ndarray) -> float: + """Return a batch's normalization deviation beyond float32 rounding, else 0.""" + deviation = float(counts[2:3].view(np.float32)[0]) + # NOTE: fast math records a NaN total as NaN (not inf), which must warn. + return 0.0 if deviation <= DEVIATION_TOLERANCE else deviation + + def _wait(self, commands: Any) -> None: + self.gpu_seconds += runtime.wait(commands) + + def sample( + self, + shots: int, + columns: list[np.ndarray], + *, + bit_packed: bool, + batch_size: int | None = None, + mask: np.ndarray | None = None, + flip: np.ndarray | None = None, + ) -> tuple[list[np.ndarray], float]: + """Return the output bits ``columns`` of ``shots`` shots and the largest norm deviation. + + As ``ShotSampler.sample``: with a ``mask`` only the shots that pass + postselection are returned. ``batch_size`` bounds the shots of a batch. + Out of GPU memory, the call warns and samples on the CPU. + """ + try: + rows, _, deviation = self.sample_accepted( + shots, + mask=mask, + flip=flip, + batch_size=batch_size, + bools=not bit_packed, + ) + except MemoryError as error: + warnings.warn( + f"{error}; sampling on the CPU. Fewer shots per call fit.", + stacklevel=4, + ) + return self._host.sample( + shots, + columns, + bit_packed=bit_packed, + batch_size=batch_size, + mask=mask, + flip=flip, + ) + if not bit_packed: + return [_columns(rows, c) for c in columns], deviation + num_outputs = self.layout.num_outputs + return [_select(rows, c, num_outputs, bit_packed) for c in columns], deviation + + def _capacity(self, key: tuple, module: _Module, shots: int) -> int: + """Return the rows to make room for: the expected accepted shots and a margin. + + The share comes from the previous call of ``key``, or else the screen's + estimate (which the tapes' postselection only lowers). + """ + share = self._shares.get(key, module.passing if module.has_mask else 1.0) + expected = shots * min(1.0, share) + return min(shots, int(1.05 * expected + 8 * np.sqrt(expected) + 4096)) + + def _run_accepted( + self, + module: _Module, + slot: _Slot, + shots: int, + batch: int, + constant: runtime.Buffer, + expectations: bool, + split: tuple[_Launch, _Launch] | None, + rows: runtime.Buffer, + values: runtime.Buffer | None, + capacity: int, + ) -> None: + """Run a call's batches into ``rows`` (and ``values``) of ``capacity`` rows.""" + launch = split[0] if split else self._all + starts = range(0, shots, batch) + fixed = runtime.buffer(len(starts), np.uint32) + slot.counts.array[:] = 0 + slot.counts.array[4] = capacity + commands = [] + for k, start in enumerate(starts): + n = min(batch, shots - start) + fixed.array[k] = n + encoder = runtime.Encoder() + self._encode_batch( + encoder, module, slot, n, start, constant, expectations, launch + ) + self._encode_pack( + encoder, + module, + slot, + n, + split, + expectations, + rows, + values, + (fixed, 4 * k), + ) + commands.append(encoder.commit()) + self._wait_all(commands) + + def _wait_all(self, commands: list[Any]) -> None: + """Wait for a call's command buffers, all of them even if one fails. + + Raises: + RuntimeError: If one failed; the batch buffers are then dropped. + + """ + try: + for command in commands: + self._wait(command) + except BaseException: + for command in commands: + command.waitUntilCompleted() + self._slots.clear() + raise + + def sample_accepted( + self, + shots: int, + *, + mask: np.ndarray | None = None, + flip: np.ndarray | None = None, + expectations: bool = False, + batch_size: int | None = None, + bools: bool = False, + ) -> tuple[np.ndarray, 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. + + 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 + 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 + measurement (else None), and the largest normalization deviation + beyond float32 rounding. + + Raises: + MemoryError: If the GPU runs out of memory. + + """ + module = self._module(mask) + nb, no = -(-len(module.columns) // 8), len(self._observables) + if not shots: + values = np.zeros((0, no), np.float32) if expectations else None + if bools: + return np.zeros((0, self.layout.num_outputs), bool), values, 0.0 + return np.zeros((0, nb), np.uint8), values, 0.0 + constant = self._constant(flip) + split = self._split(module, expectations) + batch = self._batch_shots(shots, batch_size) + slot = self._slot(batch, expectations) + key = (module.plan, expectations) + state = self._rng.bit_generator.state + capacity = self._capacity(key, module, shots) + while True: + # NOTE: a command buffer makes all of its buffers resident, so the + # output has room for the expected rows only; a call that overflows + # it runs again, on the same streams, with room for every shot. + rows = runtime.buffer(capacity * nb, np.uint8, recycle=True) + values = ( + runtime.buffer(capacity * no, np.float32, recycle=True) + if expectations + else None + ) + self._run_accepted( + module, + slot, + shots, + batch, + constant, + expectations, + split, + rows, + values, + capacity, + ) + total = int(slot.counts.array[3]) + if total <= capacity: + break + self._rng.bit_generator.state = state + capacity = shots + self._shares[key] = total / max(shots, 1) + deviation = self._deviation(slot.counts.array) + packed = rows.array[: total * nb].reshape(total, nb) + expected = None + if values is not None: + expected = values.array[: total * no].reshape(total, no) + if not self.layout.components: + expected[:] = 0 + self._direct_expectations(packed, expected, module.columns) + if bools: + return self._unpack(module, rows, total), expected, deviation + return packed, expected, deviation + + def _unpack(self, module: _Module, rows: runtime.Buffer, total: int) -> np.ndarray: + """Return the first ``total`` packed ``rows`` as booleans, unpacked on the GPU.""" + num_outputs = self.layout.num_outputs + nb = -(-num_outputs // 8) + groups = -(-num_outputs // 4) + if total * groups >= 1 << 32: + packed = rows.array[: total * nb].reshape(total, nb) + return _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, 1 << 20) + encoder = runtime.Encoder() + encoder.dispatch( + module.library.pipeline("unpack_rows"), + [rows, out, np.uint32(total)], + -(-threads // 256), + 256, + ) + self._wait(encoder.commit()) + return out.array[: total * num_outputs].reshape(total, num_outputs).view(bool) + + def _direct_expectations( + self, rows: np.ndarray, values: np.ndarray, columns: np.ndarray + ) -> None: + """Set the expectations of observables outside the components: ``±1``.""" + layout = self.layout + in_components = {i for c in layout.components for i in c.output_indices} + for column, index in enumerate(self._observables): + if int(index) in in_components: + continue + k = int(np.searchsorted(columns, index)) + bit = (rows[:, k >> 3] >> np.uint8(k & 7)) & np.uint8(1) + values[:, column] = 1.0 - 2.0 * bit.astype(np.float32) + + def direct_bools(self, shots: int, flip: np.ndarray | None = None) -> np.ndarray: + """Return every output bit of a program without components, as booleans.""" + rows, _, _ = self.sample_accepted(shots, flip=flip, bools=True) + return rows + + +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", " MetalShotSampler: + """Return the GPU shot sampler of ``sampler``.""" + assert isinstance(sampler._shots, MetalShotSampler) + return sampler._shots + + +def _cpu(compile, **kwargs): + """Return ``compile(**kwargs)``, a sampler compiled as on a host without a GPU.""" + with pytest.MonkeyPatch.context() as patch: + patch.setattr(runtime, "available", lambda: False) + sampler = compile(**kwargs) + assert sampler._shots is None or type(sampler._shots) is ShotSampler + return sampler + + +def _accepted(sampler, shots, mask, **kwargs): + detectors, observables = sampler.sample( + shots, separate_observables=True, postselection_mask=mask, **kwargs + ) + # NOTE: sample returns only the shots that the mask keeps. + assert mask is None or not (detectors & mask).any() + return np.concatenate([detectors, observables], axis=1) + + +def _packed(sampler, shots, mask=None): + """Return ``sample``'s bit-packed detectors then observables.""" + return sampler.sample( + shots, append_observables=True, bit_packed=True, postselection_mask=mask + ) + + +@pytest.mark.parametrize("name", sorted(CIRCUITS)) +def test_metal_detector_sampling_matches_stabilizer_rank(name): + circuit = Circuit(CIRCUITS[name]) + shots = 40_000 + expected = circuit.compile_detector_sampler(strategy="cat5", seed=1).sample( + shots, batch_size=10_000, append_observables=True + ) + got = _gpu(circuit.compile_detector_sampler(seed=2)).sample( + shots, append_observables=True + ) + assert_counts_match(got, expected) + + +@pytest.mark.parametrize("name", sorted(CIRCUITS)) +def test_metal_measurement_sampling_matches_stabilizer_rank(name): + circuit = Circuit(CIRCUITS[name]) + shots = 40_000 + expected = circuit.compile_sampler(strategy="cat5", seed=1).sample( + shots, batch_size=10_000 + ) + got = _gpu(circuit.compile_sampler(seed=2)).sample(shots) + assert_counts_match(got, expected) + + +@pytest.mark.parametrize("name", sorted(CIRCUITS)) +def test_metal_postselection_keeps_the_accepted_distribution(name): + circuit = Circuit(CIRCUITS[name]) + if circuit.num_detectors == 0: + pytest.skip("no detectors to postselect") + shots = 40_000 + mask = np.ones(circuit.num_detectors, dtype=bool) + expected = _accepted(_cpu(circuit.compile_detector_sampler, seed=1), shots, mask) + got = _accepted(_gpu(circuit.compile_detector_sampler(seed=2)), shots, mask) + rate, gpu_rate = len(expected) / shots, len(got) / shots + sigma = np.sqrt(rate * (1 - rate) * 2 / shots) + assert abs(gpu_rate - rate) <= 5 * sigma + 3 / shots + if len(got) and len(expected): + assert_counts_match(got, expected) + + +@pytest.mark.parametrize("seed", [1, 2, 3]) +def test_metal_sampling_of_random_noisy_circuits(seed): + stim_circuit = gen_stim_circuit(3, 60, p_t=3, p_x_error=0.5, p_r_z=0.5, seed=seed) + circuit = Circuit.from_stim_program(stim_circuit) + shots = 8_192 + expected = circuit.compile_sampler(strategy="cat5", seed=seed).sample( + shots, batch_size=shots + ) + got = _gpu(circuit.compile_sampler(seed=seed + 10)).sample(shots) + assert_counts_match(got, expected) + + +@pytest.mark.parametrize("reference", [False, True]) +def test_metal_reference_samples_and_pruning(reference): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + mask = np.ones(circuit.num_detectors, dtype=bool) + shots = 40_000 + kwargs = {"use_detector_reference_sample": reference} + expected = _accepted( + _cpu(circuit.compile_detector_sampler, seed=3), shots, mask, **kwargs + ) + got = _accepted( + _gpu(circuit.compile_detector_sampler(seed=3)), shots, mask, **kwargs + ) + rate, gpu_rate = len(expected) / shots, len(got) / shots + assert abs(gpu_rate - rate) <= 5 * np.sqrt(rate * (1 - rate) * 2 / shots) + assert_counts_match(got, expected) + + +def test_metal_postselection_keeps_every_column(): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + packed = _packed(_gpu(circuit.compile_detector_sampler(seed=4)), 20_000, mask) + outputs = circuit.num_detectors + circuit.num_observables + assert packed.shape[1] == -(-outputs // 8) + bits = np.unpackbits(packed, axis=1, bitorder="little") + assert not bits[:, outputs:].any() + assert not bits[:, : circuit.num_detectors][:, mask].any() + full = _accepted(_cpu(circuit.compile_detector_sampler, seed=5), 20_000, mask) + assert abs(len(bits) - len(full)) <= 5 * np.sqrt(len(full)) + assert_counts_match(bits[:, :outputs].astype(bool), full) + + +@pytest.mark.parametrize("bit_packed", [False, True]) +def test_metal_arranges_the_columns_as_the_cpu_does(bit_packed): + # NOTE: samplers of one seed draw the same shots, whichever columns they return. + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[0] = True + num_detectors = circuit.num_detectors + + def sample(**kwargs): + sampler = _gpu(circuit.compile_detector_sampler(seed=17)) + return sampler.sample( + 5_000, postselection_mask=mask, bit_packed=bit_packed, **kwargs + ) + + def bools(array, count): + if not bit_packed: + return array + return np.unpackbits(array, axis=1, count=count, bitorder="little").astype(bool) + + appended = bools(sample(append_observables=True), num_detectors + 1) + detectors = bools(sample(), num_detectors) + prepended = bools(sample(prepend_observables=True), num_detectors + 1) + separate = sample(separate_observables=True) + np.testing.assert_array_equal(detectors, appended[:, :num_detectors]) + np.testing.assert_array_equal(prepended[:, 0], appended[:, -1]) + np.testing.assert_array_equal(prepended[:, 1:], detectors) + np.testing.assert_array_equal(bools(separate[0], num_detectors), detectors) + np.testing.assert_array_equal(bools(separate[1], 1), appended[:, -1:]) + cpu = _cpu(circuit.compile_detector_sampler, seed=17).sample( + 10, postselection_mask=mask, bit_packed=bit_packed, append_observables=True + ) + gpu = sample(append_observables=True) + assert gpu.dtype == cpu.dtype + assert gpu.shape[1] == cpu.shape[1] + + +def test_metal_samples_no_shots(): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + 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) + + +# NOTE: 150 noisy detectors, so the rows span three words, and two observables. +NOISY_DETECTORS = "\n".join( + [ + "R " + " ".join(map(str, range(150))), + "X_ERROR(0.3) " + " ".join(map(str, range(150))), + "M " + " ".join(map(str, range(150))), + *(f"DETECTOR rec[{i - 150}]" for i in range(150)), + "OBSERVABLE_INCLUDE(0) rec[-1] rec[-7]", + "OBSERVABLE_INCLUDE(1) rec[-150]", + ] +) + + +def test_metal_postselection_drops_the_rows_that_fail(monkeypatch): + # NOTE: the masked run must draw the unmasked run's noise, so none is deferred. + monkeypatch.setattr(metal_sampler, "DEFER_SHARE", 2.0) + circuit = Circuit(NOISY_DETECTORS) + rng = np.random.default_rng(0) + mask = np.zeros(150, dtype=bool) + # NOTE: a few masked detectors, then unmasked runs across words, then a mix. + mask[[0, 3, 4]] = True + mask[100:] = rng.random(50) < 0.1 + full = _packed(_gpu(circuit.compile_detector_sampler(seed=11)), 5_000) + rows = _packed(_gpu(circuit.compile_detector_sampler(seed=11)), 5_000, mask) + full_bits = np.unpackbits(full, axis=1, count=152, bitorder="little") + passing = ~full_bits[:, :150][:, mask].any(axis=1) + assert 0 < passing.sum() < 5_000 + np.testing.assert_array_equal(rows, full[passing]) + + +# NOTE: GHZ_WITH_DIRECT_DETECTOR's component and direct detector, then four +# noisy direct detectors that no mask below reads; one fault flips detectors 2 +# and 3 together, and the observable also reads the last noisy qubit. +UNREAD_NOISE = """ +R 0 1 2 3 4 5 6 7 +H 0 +CNOT 0 1 +CNOT 1 2 +T 0 1 2 +H 0 1 2 +DEPOLARIZE1(0.05) 0 1 2 +T 0 1 +CZ 0 1 +H 0 1 +M 0 1 2 +DETECTOR rec[-3] rec[-2] +DETECTOR rec[-2] rec[-1] +OBSERVABLE_INCLUDE(0) rec[-1] +X_ERROR(0.2) 3 +E(0.1) X3 X4 +X_ERROR(0.3) 4 5 6 7 +M 3 4 5 6 7 +DETECTOR rec[-5] +DETECTOR rec[-4] +DETECTOR rec[-3] +DETECTOR rec[-2] +DETECTOR rec[-1] +OBSERVABLE_INCLUDE(0) rec[-1] +""" + + +def test_metal_defers_the_noise_that_postselection_cannot_see(): + circuit = Circuit(UNREAD_NOISE) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + sampler = _gpu(circuit.compile_detector_sampler(seed=12)) + got = _accepted(sampler, 200_000, mask) + assert _gpu_shots(sampler)._module(mask).splits[False] is not None + full = _accepted(_cpu(circuit.compile_detector_sampler, seed=13), 200_000, mask) + assert abs(len(got) - len(full)) <= 5 * np.sqrt(len(full)) + assert_counts_match(got, full) + + +def _expectations(sampler, shots, mask=None): + """Return the passing rows and expectations of ``sample_accepted``.""" + rows, values, _ = _gpu_shots(sampler).sample_accepted( + shots, mask=mask, expectations=True + ) + assert values is not None + return rows, values + + +def test_deferred_noise_leaves_the_expectations_unbiased(): + # NOTE: with expectations, the noise on the observable's direct qubit is not + # deferred: the tape reads the observable's bit for its sign. + circuit = Circuit(UNREAD_NOISE) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + sampler = _gpu(circuit.compile_detector_sampler(seed=14)) + _, values = _expectations(sampler, 200_000, mask) + assert _gpu_shots(sampler)._module(mask).splits[True] is not None + full = _accepted(_cpu(circuit.compile_detector_sampler, seed=15), 200_000, mask) + expected = np.mean(1 - 2 * full[:, -1].astype(float)) + assert abs(values.mean() - expected) < 5 / np.sqrt(len(full)) + + +@pytest.mark.parametrize("text", [NOISY_DETECTORS, UNREAD_NOISE]) +def test_metal_batches_return_the_rows_of_one_pass(monkeypatch, text): + # NOTE: 1000-shot batches, queued on one set of buffers; without deferral, a + # masked and an unmasked run on the same batches draw the same shots. + monkeypatch.setattr(metal_sampler, "DEFER_SHARE", 2.0) + circuit = Circuit(text) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + outputs = circuit.num_detectors + circuit.num_observables + + def packed(mask=None): + sampler = _gpu(circuit.compile_detector_sampler(seed=16)) + return sampler.sample( + 10_500, + append_observables=True, + bit_packed=True, + postselection_mask=mask, + batch_size=1000, + ) + + full, rows = packed(), packed(mask) + full_bits = np.unpackbits(full, axis=1, count=outputs, bitorder="little") + passing = ~full_bits[:, : circuit.num_detectors][:, mask].any(axis=1) + assert 0 < passing.sum() < 10_500 + np.testing.assert_array_equal(rows, full[passing]) + + +def test_observable_expectations_average_to_the_sampled_observable(): + circuit = Circuit(T_STATE_X) + sampler = _gpu(circuit.compile_detector_sampler(seed=6)) + packed, values = _expectations(sampler, 200_000) + assert values.shape == (len(packed), 1) + # NOTE: given the shot's Z error, the expectation is ±cos(π/4). + np.testing.assert_allclose(np.abs(values), np.sqrt(0.5), atol=1e-5) + observed = 1 - 2 * np.unpackbits(packed, axis=1, bitorder="little")[:, 0].astype( + float + ) + sampled = _cpu(circuit.compile_detector_sampler, seed=7).sample( + 200_000, append_observables=True + )[:, 0] + expected = np.mean(1 - 2 * sampled.astype(float)) + assert abs(values.mean() - expected) < 5 * 1 / np.sqrt(200_000) + # NOTE: the returned observable bit is sampled from the same state. + assert abs(observed.mean() - values.mean()) < 5 / np.sqrt(200_000) + + +# NOTE: the T state's X outcome is copied onto qubit 1, whose postselected (and +# noisy) detector is settled after the observable's measurement. +T_STATE_X_THEN_DETECTOR = """ +R 0 1 +H 0 +T 0 +H 0 +CNOT 0 1 +X_ERROR(0.1) 1 +M 0 +OBSERVABLE_INCLUDE(0) rec[-1] +M 1 +DETECTOR rec[-1] +""" + +# NOTE: a postselected T-state detector settled before the observable's measurement. +T_STATE_X_AFTER_DETECTOR = """ +R 0 1 +H 0 1 +T 0 1 +H 1 +M 1 +DETECTOR rec[-1] +Z_ERROR(0.1) 0 +H 0 +M 0 +OBSERVABLE_INCLUDE(0) rec[-1] +""" + + +def _postselected_expectations(text, shots): + circuit = Circuit(text) + mask = np.ones(circuit.num_detectors, dtype=bool) + sampler = _gpu(circuit.compile_detector_sampler(seed=9)) + _, values = _expectations(sampler, shots, mask) + detectors, observables = _cpu(circuit.compile_detector_sampler, seed=10).sample( + shots, separate_observables=True + ) + kept = observables[~(detectors & mask).any(axis=1), 0] + return values[:, 0], np.mean(1 - 2 * kept.astype(float)), len(kept) + + +def test_expectations_stay_unbiased_under_later_postselection(): + values, expected, kept = _postselected_expectations( + T_STATE_X_THEN_DETECTOR, 200_000 + ) + # NOTE: the expectation before the measurement would average to cos(π/4). + assert abs(values.mean() - expected) < 5 * np.sqrt(1 / len(values) + 1 / kept) + np.testing.assert_array_equal(np.abs(values), 1) + + +def test_expectations_are_exact_after_earlier_postselection(): + values, expected, kept = _postselected_expectations( + T_STATE_X_AFTER_DETECTOR, 200_000 + ) + np.testing.assert_allclose(np.abs(values), np.sqrt(0.5), atol=1e-5) + assert abs(values.mean() - expected) < 5 * np.sqrt(1 / len(values) + 1 / kept) + + +def test_clifford_observables_have_expectations_of_plus_or_minus_one(): + circuit = Circuit( + "R 0 1\nX_ERROR(0.3) 0\nCNOT 0 1\nM 0 1\nDETECTOR rec[-1]\nOBSERVABLE_INCLUDE(0) rec[-2]" + ) + sampler = _gpu(circuit.compile_detector_sampler(seed=8)) + packed, values = _expectations(sampler, 10_000) + bits = np.unpackbits(packed, axis=1, bitorder="little") + np.testing.assert_array_equal(values[:, 0], 1 - 2 * bits[:, 1].astype(np.float32)) + + +@pytest.mark.parametrize("reference", [False, True]) +def test_metal_samples_programs_without_components(reference): + circuit = Circuit( + "R 0 1\nX_ERROR(0.3) 0\nCNOT 0 1\nM 0 1\nDETECTOR rec[-1]\nOBSERVABLE_INCLUDE(0) rec[-2]" + ) + kwargs = {"use_detector_reference_sample": reference} + expected = _cpu(circuit.compile_detector_sampler, seed=1).sample( + 40_000, append_observables=True, **kwargs + ) + got = _gpu(circuit.compile_detector_sampler(seed=2)).sample( + 40_000, append_observables=True, **kwargs + ) + assert_counts_match(got, expected) + measured = _gpu(circuit.compile_sampler(seed=3)).sample(40_000) + assert measured.shape == (40_000, 2) + np.testing.assert_array_equal(measured[:, 0], measured[:, 1]) + + +def test_stabilizer_rank_strategies_stay_on_jax(): + circuit = Circuit(T_STATE_X) + assert circuit.compile_detector_sampler(strategy="cat5")._shots is None + + +def _widest(circuit): + """Return the active width of the circuit's widest statevector component.""" + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + components = circuit.compile_sampler(seed=1)._program.components + return max(c.width for c in components if isinstance(c, StatevectorComponent)) + + +def _random_unitary_circuit(width, layers, seed=0, noise=True): + """Layers of random U3 pairs and CX gates, then every qubit measured. + + With ``noise``, X errors come first, so that every rotation reads the shot's + f-bits and none of them runs before the shots split (``common_prefix``). + """ + rng = np.random.default_rng(seed) + lines = ["X_ERROR(0.2) " + " ".join(map(str, range(width)))] if noise else [] + for _ in range(layers): + perm = rng.permutation(width) + for a, b in zip(perm[::2], perm[1::2], strict=False): + for q in (a, b): + t = rng.uniform(-1, 1, 3) + lines.append(f"U3({t[0]},{t[1]},{t[2]}) {q}") + lines.append(f"CX {a} {b}") + lines.append("M " + " ".join(map(str, range(width)))) + return Circuit("\n".join(lines)) + + +def _assert_matches_cpu(circuit, width, shots=200_000): + expected = _cpu(circuit.compile_sampler, seed=1).sample(shots) + got = _gpu(circuit.compile_sampler(seed=2)).sample(shots) + np.testing.assert_allclose( + got.mean(axis=0), expected.mean(axis=0), atol=5 * 0.5 / np.sqrt(shots / 2) + ) + # NOTE: pairs of bits, which a wrong shuffle partner would scramble. + for a, b in [(0, 1), (0, width - 1), (width // 2, width - 1)]: + assert_counts_match(got[:, [a, b]], expected[:, [a, b]]) + + +@pytest.mark.parametrize(("width", "layers"), [(6, 8), (8, 8), (9, 10), (11, 12)]) +def test_metal_wide_components_match_the_cpu_executor(width, layers): + # NOTE: widths 6, 8, 9 and 11 run 2, 8, 16 and 32 threads per shot. + circuit = _random_unitary_circuit(width, layers) + assert _widest(circuit) == width + _assert_matches_cpu(circuit, width) + + +@pytest.mark.parametrize(("width", "layers"), [(6, 6), (8, 8), (11, 6)]) +def test_metal_looping_tapes_match_the_cpu_executor(monkeypatch, width, layers): + # NOTE: a tape past the unroll budget loops, applying runs of rotations together. + monkeypatch.setattr(tape, "UNROLL_BUDGET", 0) + _assert_matches_cpu(_random_unitary_circuit(width, layers, seed=width), width) + + +@pytest.mark.parametrize("width", [6, 9, 11]) +def test_metal_shots_start_after_the_common_prefix(width): + # NOTE: without noise the rotations before the first measurement run once, + # on the host; the shots load their amplitudes and run the rest. + circuit = _random_unitary_circuit(width, width + 2, seed=width, noise=False) + _assert_matches_cpu(circuit, width) + + +@pytest.mark.parametrize("name", sorted(CIRCUITS)) +def test_metal_looping_tapes_match_stabilizer_rank(monkeypatch, name): + # NOTE: measurements, projections and pruning in looping tapes. + monkeypatch.setattr(tape, "UNROLL_BUDGET", 0) + circuit = Circuit(CIRCUITS[name]) + shots = 40_000 + expected = circuit.compile_detector_sampler(strategy="cat5", seed=1).sample( + shots, batch_size=10_000, append_observables=True + ) + got = _gpu(circuit.compile_detector_sampler(seed=2)).sample( + shots, append_observables=True + ) + assert_counts_match(got, expected) + + +def test_components_wider_than_the_limit_run_on_the_cpu(): + circuit = _random_unitary_circuit(13, 12) + assert _widest(circuit) > metal_sampler.MAX_WIDTH + with pytest.warns(UserWarning, match="sampling on the CPU"): + sampler = circuit.compile_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + + +@pytest.mark.parametrize(("probability", "shots"), [(1e-10, 1 << 28), (1e-6, 1 << 26)]) +def test_metal_draws_small_outcome_probabilities_at_their_rate(probability, shots): + # NOTE: a 24-bit draw takes any positive probability at least 2^-24 = 6e-8 of + # the time; float32 rounding gives certain outcomes such probabilities. + theta = float(2 * np.arcsin(np.sqrt(probability)) / np.pi) + circuit = Circuit(f"R 0\nI[R_X(theta={theta!r}*pi)] 0\nM 0") + sampler = _gpu(circuit.compile_sampler(seed=3)) + ones = sum(int(sampler.sample(1 << 24).sum()) for _ in range(shots >> 24)) + expected = probability * shots + assert abs(ones - expected) <= 5 * np.sqrt(expected) + 1 + + +def test_metal_reruns_a_call_whose_rows_overflow_the_output(monkeypatch): + # NOTE: room for 10 rows: the call runs again with room for every shot, on + # the same streams, so it returns the rows of a call with enough room. + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + + def packed(): + sampler = _gpu(circuit.compile_detector_sampler(seed=21)) + return sampler.sample( + 5_000, + append_observables=True, + bit_packed=True, + postselection_mask=mask, + batch_size=1000, + ) + + expected = packed() + monkeypatch.setattr(MetalShotSampler, "_capacity", lambda self, *args: 10) + rows = packed() + assert len(rows) > 10 + np.testing.assert_array_equal(rows, expected) + + +def test_metal_unmasked_expectations_after_a_masked_call(): + # NOTE: a masked call lists its accepted shots in the batch buffers; an + # unmasked call on the same buffers must still read every row's own value. + circuit = Circuit(T_STATE_X_AFTER_DETECTOR) + mask = np.ones(circuit.num_detectors, dtype=bool) + sampler = _gpu(circuit.compile_detector_sampler(seed=5)) + _expectations(sampler, 20_000, mask) + shots = _gpu_shots(sampler) + state = shots._rng.bit_generator.state + rows, values = _expectations(sampler, 20_000) + shots._rng.bit_generator.state = state + shots._slots.clear() + again, values_again = _expectations(sampler, 20_000) + np.testing.assert_array_equal(again, rows) + np.testing.assert_array_equal(values_again, values) + + +def test_metal_reports_a_nan_normalization(): + sampler = _gpu(Circuit(T_STATE_X).compile_detector_sampler(seed=1)) + counts = np.zeros(4, dtype=np.uint32) + counts[2] = 0x7FC00000 + assert np.isnan(_gpu_shots(sampler)._deviation(counts)) + + +def test_metal_samples_rows_too_wide_for_threadgroup_memory(): + # NOTE: 3200 detectors are 51 words: the noise goes into device memory. + circuit = Circuit( + "\n".join( + [ + "R " + " ".join(map(str, range(3200))), + "X_ERROR(0.05) " + " ".join(map(str, range(3200))), + "M " + " ".join(map(str, range(3200))), + *(f"DETECTOR rec[{i - 3200}]" for i in range(3200)), + "OBSERVABLE_INCLUDE(0) rec[-1] rec[-2]", + ] + ) + ) + sampler = _gpu(circuit.compile_detector_sampler(seed=31)) + assert not _gpu_shots(sampler)._noise.shared_words + got = sampler.sample(20_000, append_observables=True) + np.testing.assert_allclose(got[:, :3200].mean(), 0.05, atol=5e-4) + np.testing.assert_allclose(got[:, -1].mean(), 2 * 0.05 * 0.95, atol=0.01) + np.testing.assert_array_equal(got[:, -1], got[:, -2] ^ got[:, -3]) + + +@pytest.mark.parametrize("name", sorted(CIRCUITS)) +def test_metal_noise_in_device_memory_matches_the_cpu(monkeypatch, name): + monkeypatch.setattr(metal_noise, "SHARED_BYTES", 2048 + 8) + circuit = Circuit(CIRCUITS[name]) + shots = 40_000 + mask = np.ones(circuit.num_detectors, dtype=bool) if circuit.num_detectors else None + expected = _accepted(_cpu(circuit.compile_detector_sampler, seed=1), shots, mask) + sampler = _gpu(circuit.compile_detector_sampler(seed=2)) + assert not _gpu_shots(sampler)._noise.shared_words + got = _accepted(sampler, shots, mask) + rate, gpu_rate = len(expected) / shots, len(got) / shots + assert ( + abs(gpu_rate - rate) <= 5 * np.sqrt(rate * (1 - rate) * 2 / shots) + 3 / shots + ) + if len(got) and len(expected): + assert_counts_match(got, expected) + + +def test_metal_unrolled_width_11_tapes_match_the_cpu_executor(): + # NOTE: 64 amplitudes per thread, unrolled, in two register arrays. + circuit = _random_unitary_circuit(11, 7, seed=0) + assert _widest(circuit) == 11 + shots = _gpu_shots(_gpu(circuit.compile_sampler(seed=1))) + assert not tape.loops(shots.plan(None).programs[0], 5, 0) + _assert_matches_cpu(circuit, 11) + + +def test_metal_returned_rows_are_not_reused_while_alive(): + circuit = Circuit(NOISY_DETECTORS) + sampler = _gpu(circuit.compile_detector_sampler(seed=3)) + first = _packed(sampler, 200_000) + kept = first.copy() + second = _packed(sampler, 200_000) + np.testing.assert_array_equal(first, kept) + assert not np.shares_memory(first, second) + del first + third = _packed(sampler, 200_000) + assert not np.shares_memory(second, third) + assert not (third == second).all() + + +def test_metal_leaves_programs_without_shot_words_on_the_cpu(): + circuit = Circuit("R 0\nH 0\nT 0\nH 0\nM 0") + sampler = circuit.compile_detector_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + assert sampler.sample(10).shape == (10, 0) + + +def _sampler_in_child(_): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + sampler = Circuit(T_STATE_X).compile_sampler(seed=1) + sampler.sample(10) + return type(sampler._shots).__name__, [str(w.message) for w in caught] + + +def test_processes_forked_after_the_gpu_was_used_sample_on_the_cpu(): + # NOTE: Metal's device does not survive a fork of a process that used it. + _gpu(Circuit(T_STATE_X).compile_sampler(seed=1)).sample(10) + with multiprocessing.get_context("fork").Pool(1) as pool: + name, messages = pool.map(_sampler_in_child, [0])[0] + assert name == "ShotSampler" + assert any("forked" in m for m in messages) + + +def test_noise_of_tiny_probabilities_fires_at_its_rate(): + # NOTE: 24-bit skips alone fire every segment at least 2^-24 of the time, + # here ~350 faults instead of ~0.3. + n, p = 1000, 1e-12 + text = "\n".join(f"X_ERROR({p * (1 + i / n)!r}) {i}" for i in range(n)) + sampler = _gpu( + Circuit(f"{text}\nM " + " ".join(map(str, range(n)))).compile_sampler(seed=1) + ) + # NOTE: faults are rare enough that each lands in a byte of its own. + faults = sum( + np.count_nonzero(sampler.sample(1 << 22, bit_packed=True)) for _ in range(44) + ) + assert faults <= 5 + + +@pytest.mark.parametrize("p0", [3e-8, 1e-6]) +def test_rare_outcome_zero_is_drawn_at_its_rate(p0): + # NOTE: w1 / total rounds to a multiple of 2^-24 below 1 in float32; the draw + # must use w0, as the CPU's ``_cut`` does. + source = codegen.PRELUDE + r""" +kernel void zeros(constant float& w0 [[buffer(0)]], constant ulong& n [[buffer(1)]], + device atomic_uint* count [[buffer(2)]], uint gid [[thread_position_in_grid]], + uint grid [[threads_per_grid]]) +{ + uint local = 0; + for (u64 i = gid; i < n; i += grid) + local += !draw_one(mix64(i * GOLDEN + 12345), w0, 1.f, w0 + 1.f); + if (local) atomic_fetch_add_explicit(count, local, memory_order_relaxed); +} +""" + pipeline = runtime.Library(source).pipeline("zeros") + count = runtime.buffer(1, np.uint32) + n = 1 << 30 + encoder = runtime.Encoder() + encoder.dispatch(pipeline, [np.float32(p0), np.uint64(n), count], 1024, 256) + runtime.wait(encoder.commit()) + expected = p0 * n + assert abs(int(count.array[0]) - expected) <= 5 * np.sqrt(expected) + 1 + + +def test_calls_out_of_gpu_memory_fall_back_to_the_cpu(monkeypatch): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + sampler = _gpu(circuit.compile_detector_sampler(seed=1)) + + def out_of_memory(*args, **kwargs): + raise MemoryError("the GPU ran out of memory") + + monkeypatch.setattr(_gpu_shots(sampler), "sample_accepted", out_of_memory) + with pytest.warns(UserWarning, match="out of memory; sampling on the CPU"): + got = sampler.sample(20_000, append_observables=True) + expected = _cpu(circuit.compile_detector_sampler, seed=2).sample( + 20_000, append_observables=True + ) + assert_counts_match(got, expected) + + +def test_samplers_out_of_gpu_memory_at_setup_run_on_the_cpu(monkeypatch): + def out_of_memory(*args, **kwargs): + raise MemoryError("no Metal buffer") + + monkeypatch.setattr(MetalShotSampler, "__init__", out_of_memory) + with pytest.warns(UserWarning, match="out of memory; sampling on the CPU"): + sampler = Circuit(T_STATE_X).compile_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + + +def test_tsim_device_cpu_keeps_the_samplers_on_the_cpu(monkeypatch): + monkeypatch.setenv("TSIM_DEVICE", "cpu") + with warnings.catch_warnings(): + warnings.simplefilter("error") + sampler = Circuit(T_STATE_X).compile_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + monkeypatch.setenv("TSIM_DEVICE", "") + _gpu(Circuit(T_STATE_X).compile_sampler(seed=1)) + + +@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 + # multiple of 8. + circuit = Circuit(text) + outputs = circuit.num_detectors + circuit.num_observables + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[0] = True + for postselection_mask in (None, mask): + + def sample(bit_packed, postselection_mask=postselection_mask): + sampler = _gpu(circuit.compile_detector_sampler(seed=23)) + return sampler.sample( + 30_000, + append_observables=True, + bit_packed=bit_packed, + postselection_mask=postselection_mask, + ) + + packed, bools = sample(True), sample(False) + assert bools.dtype == np.bool_ and bools.flags.c_contiguous + unpacked = np.unpackbits(packed, axis=1, count=outputs, bitorder="little") + np.testing.assert_array_equal(bools, unpacked.astype(bool)) + detectors = _gpu(circuit.compile_detector_sampler(seed=23)).sample( + 30_000, postselection_mask=postselection_mask + ) + assert detectors.flags.c_contiguous + np.testing.assert_array_equal(detectors, bools[:, : circuit.num_detectors]) diff --git a/uv.lock b/uv.lock index a741f3e7..c9dd6b3f 100644 --- a/uv.lock +++ b/uv.lock @@ -142,6 +142,9 @@ cuda12 = [ cuda13 = [ { name = "jax", extra = ["cuda13"] }, ] +metal = [ + { name = "pyobjc-framework-metal", marker = "platform_machine == 'arm64' and sys_platform == 'darwin'" }, +] [package.dev-dependencies] dev = [ @@ -191,11 +194,12 @@ requires-dist = [ { name = "numba", specifier = ">=0.61" }, { name = "numpy", specifier = ">=2.0" }, { name = "psutil", specifier = ">=5.9.0" }, + { name = "pyobjc-framework-metal", marker = "platform_machine == 'arm64' and sys_platform == 'darwin' and extra == 'metal'", specifier = ">=10.0" }, { name = "stim", specifier = ">=1.15.0" }, { name = "stim-timeline-text", specifier = ">=0.1.0" }, { name = "zxparam", specifier = ">=0.1.0" }, ] -provides-extras = ["cuda12", "cuda13"] +provides-extras = ["cuda12", "cuda13", "metal"] [package.metadata.requires-dev] dev = [ @@ -2788,6 +2792,61 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/17/0d/eb2df17fe0bf089c43445ba759e89150a0d36fb7fcbc44e750d8d31f0dbe/pymatching-2.4.0-cp314-cp314-win_amd64.whl", hash = "sha256:5c5ce4ed245d3670ea7d62b09f8a29b105533b032a3d21fdf6df04c9b6441ff1", size = 356304, upload-time = "2026-05-22T18:09:54.3Z" }, ] +[[package]] +name = "pyobjc-core" +version = "13.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/83/58/7bab095330b8893e1ac58f3cf9710462c3242b600f8fa5db0d5dda795775/pyobjc_core-13.0.tar.gz", hash = "sha256:ffe154462d14aa4852f424e9116f17b5f8d44f78fc0f6c4727e721a5fa7177da", size = 1083172, upload-time = "2026-10-30T10:45:39.835Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/06/ad/7151cd0f01d561a6d9e761b12cdcbb5bcf4f4e65efa619bf4afcb202df28/pyobjc_core-13.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:4627348cb31c8deb2bb94702d9bcaa74869f4fef9559fd134416a6a4ca42006e", size = 5615698, upload-time = "2026-10-30T10:32:42.051Z" }, + { url = "https://files.pythonhosted.org/packages/68/ac/783ecf1c5045792fb51ad18e0131ff64e33ed9dee9921766c14e03eb27db/pyobjc_core-13.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:5818e3d8cc2677a1e767968ebfe3bc0cb7f80c1f7c19c653ee25ef1cc1b3e188", size = 6716501, upload-time = "2026-10-30T10:32:46.2Z" }, + { url = "https://files.pythonhosted.org/packages/d8/42/7b31511491d966b64183dfab325f022098e06e90c19a9cc8540de4f32588/pyobjc_core-13.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:f29f7cdb3969268716017ada9516b94f6dd1c66c2e415274b5ac1169440da1a4", size = 6734122, upload-time = "2026-10-30T10:32:50.192Z" }, + { url = "https://files.pythonhosted.org/packages/04/54/fda17a291ecdb5186f545c9a7575d02f55815e239756567f030b555b14f3/pyobjc_core-13.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:1e8752b3950ceec603bf36f57083151abed15e0c2e7c5b8c3efded9c81dc732c", size = 7031547, upload-time = "2026-10-30T10:32:55.183Z" }, + { url = "https://files.pythonhosted.org/packages/41/a5/463dc32a6c32cff57b4678298a13cd830d7ab09cfec3d7e80f20a2f414fd/pyobjc_core-13.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:fe470bd7b8163b4b8d216b1ebe5b227f01c0fa1c4e0085d7ccade269dec02bb0", size = 6729572, upload-time = "2026-10-30T10:32:59.498Z" }, + { url = "https://files.pythonhosted.org/packages/b4/11/e57bb09924e8e0dfbaa87a9cf8da13e224189f760aa0c39ab0ce9defb0d4/pyobjc_core-13.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:c497a23077ba0a42b59b879d81049eebfb0b7b67a8ca80c68b6d3683bbbd6d6f", size = 7021797, upload-time = "2026-10-30T10:33:03.702Z" }, + { url = "https://files.pythonhosted.org/packages/43/3c/0f39ddc5ff4d1b8e38134349c4f018c47d4fbba297510561429aa6c1a2c3/pyobjc_core-13.0-cp315-cp315-macosx_10_15_universal2.whl", hash = "sha256:d65d383efed5a803e43f6c67ccc9f4c94754842b2268da24319bb57a852b016f", size = 6782204, upload-time = "2026-10-30T10:33:08.125Z" }, + { url = "https://files.pythonhosted.org/packages/dd/bc/35eb48cff5d20770a2adb2bb0f01424f5dfbcc7dae8e5ecb990f6d90adf9/pyobjc_core-13.0-cp315-cp315t-macosx_10_15_universal2.whl", hash = "sha256:292462e0947a09aca3737a06c44457cb2aa58f267cfc581c16fca036fc073b45", size = 7036187, upload-time = "2026-10-30T10:33:12.36Z" }, +] + +[[package]] +name = "pyobjc-framework-cocoa" +version = "13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyobjc-core" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/30/50/55a8779d35f1f86d795264cbd6fe62e35b2cf4fb8343fc888e99268c988a/pyobjc_framework_cocoa-13.0.tar.gz", hash = "sha256:bb7a23941fc8d7866578dbc84e56140efd343a1d28510f530dd1dbf80807daf4", size = 3135308, upload-time = "2026-10-30T10:46:03.982Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f2/88/d8cece46725d7879aee3b6a047cbbb86e54aa51b31c707c7d8da6bffa331/pyobjc_framework_cocoa-13.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:16dd1e8890ed695ce12bcd367ed0ac9f3e70dd1f7a8098b6398427c162591b83", size = 391003, upload-time = "2026-10-30T10:35:14.969Z" }, + { url = "https://files.pythonhosted.org/packages/ea/44/98d626726809da517c757d77e2888d34eaf758ebad908effb8d0b60c8118/pyobjc_framework_cocoa-13.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:e1618d315939c9fae57cb858d4db496ec61f82495ce6a877398131ee494fc19d", size = 391597, upload-time = "2026-10-30T10:35:16.217Z" }, + { url = "https://files.pythonhosted.org/packages/01/b0/cb28ff4474a2f07bb9b64751112b3b40c421755297749c70810f4cadfc24/pyobjc_framework_cocoa-13.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:bb0bcba2e4e9d05af875306e0e055cfc2f152fd602d5f2f1a54490adebebcd34", size = 391670, upload-time = "2026-10-30T10:35:17.336Z" }, + { url = "https://files.pythonhosted.org/packages/54/e6/53aeaa92fd8261ba32d7faca85bef4aa210740f4f4462d0e727be6769731/pyobjc_framework_cocoa-13.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:d6fa11ea29d85d46e65ed1c9fdba3a5131ad26971bf24bb7d16aa55f06bcdc1e", size = 395676, upload-time = "2026-10-30T10:35:18.813Z" }, + { url = "https://files.pythonhosted.org/packages/b4/bb/f683ec329180e1931c8d592c882defebca8810e019cf4bc8727fa36b7d89/pyobjc_framework_cocoa-13.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:88be6d452e7531f695ce6d41428aa74bcaade9ddbc50af592419b9f2fbb7c879", size = 391936, upload-time = "2026-10-30T10:35:20.154Z" }, + { url = "https://files.pythonhosted.org/packages/c5/3e/1b24b56d3710a594dc2350ed6b321aaf0ee12fca3dafa6034bccffc8ac08/pyobjc_framework_cocoa-13.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:1a6b9f9e0ab123e51fa7df43bbc6c0414a444facd2225ad4f221a0d842c49ad6", size = 395844, upload-time = "2026-10-30T10:35:21.319Z" }, + { url = "https://files.pythonhosted.org/packages/f2/e0/b2f89d9c92e2b6513d8495d81c3a7d6042f936f51ed7dd2180615a01c980/pyobjc_framework_cocoa-13.0-cp315-cp315-macosx_10_15_universal2.whl", hash = "sha256:9282ebad57784ec18f4b7ea2e6b65bac0a6f56cc21a5c63e061fa48a9b49dd26", size = 392140, upload-time = "2026-10-30T10:35:22.562Z" }, + { url = "https://files.pythonhosted.org/packages/2c/ac/d2c8fa5d96ddf4e70ac5ddb23f24ac7861577728df342697125c1b42e159/pyobjc_framework_cocoa-13.0-cp315-cp315t-macosx_10_15_universal2.whl", hash = "sha256:f1e26ad2225d1c84973cce725f2f4aeeb4d5f653a36c8a9f0b691c5529dc7780", size = 396079, upload-time = "2026-10-30T10:35:23.821Z" }, +] + +[[package]] +name = "pyobjc-framework-metal" +version = "13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyobjc-core" }, + { name = "pyobjc-framework-cocoa" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/27/e6/f6efb0f951df22e8571facc26e7bfde7f1b06cb59bf682737c0bbbbc4593/pyobjc_framework_metal-13.0.tar.gz", hash = "sha256:be798a872021520701183f015f1b0fdadbc5ca74e4aee80d0c4952b890c90a7a", size = 244348, upload-time = "2026-10-30T10:47:02.567Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ac/36/9f2319402d376768a9d1d9ef0b902e1d3d785d52670eb5e918262c2775d1/pyobjc_framework_metal-13.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:f58b3dde81fa568a11ea4098debb827e5cdb130565efe4ce0b1d9ec48a11cd68", size = 78042, upload-time = "2026-10-30T10:40:38.412Z" }, + { url = "https://files.pythonhosted.org/packages/03/a6/a0330996fe6e1082bc0ef9251da651ac9cb00976f54d4be35dc8a977c5fe/pyobjc_framework_metal-13.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a69d760476121cd0faad8ea2f4c70708398aad97eb33fe7c186712e437ecacd7", size = 78028, upload-time = "2026-10-30T10:40:39.555Z" }, + { url = "https://files.pythonhosted.org/packages/9e/6d/12e13ce0db817a809d2fa7e948d688b3cfbe26c39d1f721c1a489daeaf7e/pyobjc_framework_metal-13.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:79d682af412582a092222f4b85b5120f05050ef6c5d5e69a199a4aa3a25a86da", size = 78046, upload-time = "2026-10-30T10:40:40.482Z" }, + { url = "https://files.pythonhosted.org/packages/76/72/f553bc777a59c674dc5b847937b2de844524e09e93dd4e760b4ab1fb9d76/pyobjc_framework_metal-13.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:e197115790ee469143708c45b203723ee5589ffc06f1a0006ef53139141cd472", size = 78572, upload-time = "2026-10-30T10:40:41.382Z" }, + { url = "https://files.pythonhosted.org/packages/7c/5d/2d1e4f9eb76bf338f3528aa1511a817889042061c3d847b06b3a86bb419c/pyobjc_framework_metal-13.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:26b6a4e998dc60261d01be45cb9c53c45386a6750af27d612e3fbfc1f2565fea", size = 78138, upload-time = "2026-10-30T10:40:42.316Z" }, + { url = "https://files.pythonhosted.org/packages/9a/94/9a75ef9f0df859e5ceffe0e477f96d72e1a09bdec007811836cd8b7088e7/pyobjc_framework_metal-13.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:72b9d46a05b72a54adf3183a9ac6353346f01345a8e32d1b763b6176142d43d7", size = 78719, upload-time = "2026-10-30T10:40:43.194Z" }, + { url = "https://files.pythonhosted.org/packages/dd/26/57832f2b7704c61e9581663dc36d689d083bd2d9af74602b3fd3889f90d6/pyobjc_framework_metal-13.0-cp315-cp315-macosx_10_15_universal2.whl", hash = "sha256:902f71d42aa282a65b5dbb49c0884543b4a31a17b31d3dd61aaa80e688715fdd", size = 78249, upload-time = "2026-10-30T10:40:44.165Z" }, + { url = "https://files.pythonhosted.org/packages/bf/4a/6e19d461171f1fca713b8d7569df3334d1f3d23b727c660c327ba2dab7ce/pyobjc_framework_metal-13.0-cp315-cp315t-macosx_10_15_universal2.whl", hash = "sha256:1706dcdd2002ff2bf6dc05fdcd0ea0673692359fb20130f3e8a3250a99f88503", size = 78817, upload-time = "2026-10-30T10:40:45.037Z" }, +] + [[package]] name = "pyparsing" version = "3.3.2" From f3ab6707a43d15993628b522e5f0cedaabec83f9 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 10 Nov 2026 20:16:12 +0100 Subject: [PATCH 10/38] fix(metal): make GPU samplers safe to share, interrupt, fork and pickle - Calls on one sampler run one at a time (they share its batch buffers), and the buffer pool is locked. - A failed or interrupted call waits for the batches it committed, through further interrupts, and drops the batch buffers, so the next call does not return its rows. - A sampler used in a process forked after Metal's setup samples on the CPU, with a warning, instead of hanging. - Samplers pickle: the CPU sampler and the inputs, with the GPU sampler rebuilt on load (or the CPU one where Metal is missing). - Noise beyond the GPU tables' limits, and running out of GPU memory for programs without components, fall back to the CPU with a warning. - Calls of more than 2^31 shots run in parts (the GPU counters are 32-bit). - `SamplerStats.device` says where a sampler samples. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 3 +- src/tsim/sampler.py | 26 ++-- src/tsim/statevector/metal/runtime.py | 25 +++- src/tsim/statevector/metal/sampler.py | 187 +++++++++++++++++++------ test/integration/test_metal_sampler.py | 142 +++++++++++++++++++ 5 files changed, 328 insertions(+), 55 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1061e45b..111cf347 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,7 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 rows are written. A seed's samples on the GPU differ from the CPU's and depend on `batch_size`. The samplers fall back to the CPU, with a warning, for wider components, when the GPU is out of memory and in a process forked after its - parent used the GPU. `TSIM_DEVICE=cpu` keeps them on the CPU. + parent used the GPU. `TSIM_DEVICE=cpu` keeps them on the CPU, and + `SamplerStats.device` says where a sampler samples. - The docs pages have a menu to copy a page or open it in ChatGPT or Claude, and the README, the install page and the agent skill point to the docs MCP server at `https://bloqade.quera.com/tsim/mcp`. The README and the install page also diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 3b6cb413..63a56304 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -311,6 +311,7 @@ class SamplerStats: error_channel_bits: int peak_bytes_per_shot: int compiled_bytes: int + device: str = "cpu" def _gpu_shots( @@ -322,20 +323,16 @@ def _gpu_shots( """Return a GPU shot sampler of ``host``'s program when one can run it, else ``host``. It warns when a GPU is there but cannot be used: in a process forked after its - parent used the GPU, for a component wider than the GPU executor runs, and - when the GPU sampler cannot be set up for lack of GPU memory. + parent used the GPU, for a component wider than the GPU executor runs, for + noise beyond its tables' limits, and when the GPU sampler cannot be set up + for lack of GPU memory. """ # NOTE: as JAX does for the stabilizer-rank strategies, a GPU is used whenever # there is one. if os.environ.get("TSIM_DEVICE", "").lower() == "cpu": return host if metal_runtime.forked_after_init(): - warnings.warn( - "This process was forked after its parent used the GPU, so it samples " - "on the CPU; start it with the 'spawn' or 'forkserver' method to " - "sample on the GPU.", - stacklevel=5, - ) + warnings.warn(metal_runtime.FORKED, stacklevel=5) return host if not metal_runtime.available(): return host @@ -357,7 +354,14 @@ def _gpu_shots( return MetalShotSampler(host, channels, num_detectors, seed) except MemoryError: warnings.warn("The GPU is out of memory; sampling on the CPU.", stacklevel=5) - return host + except ValueError as error: + # NOTE: noise beyond the GPU tables' limits (channels of more than 127 + # outcomes, outcomes that flip more than 255 words). + warnings.warn( + f"The GPU cannot sample this circuit ({error}); sampling on the CPU.", + stacklevel=5, + ) + return host class _CompiledSamplerBase: @@ -783,6 +787,9 @@ def stats(self) -> SamplerStats: - ``peak_bytes_per_shot`` (``int``): estimated peak memory of one shot, by which the default batch sizes are chosen. - ``compiled_bytes`` (``int``): memory of the compiled components. + - ``device`` (``str``): where ``"active_statevector"`` samples: + ``"metal"`` on an Apple GPU, else ``"cpu"`` (a stabilizer-rank + strategy runs wherever JAX does). """ components = self._program.components widths: tuple[int, ...] | None = None @@ -813,6 +820,7 @@ def stats(self) -> SamplerStats: ), peak_bytes_per_shot=peak, compiled_bytes=compiled_bytes, + device=getattr(self._shots, "device", "cpu"), ) def __repr__(self) -> str: diff --git a/src/tsim/statevector/metal/runtime.py b/src/tsim/statevector/metal/runtime.py index 51ac30f3..abf0d298 100644 --- a/src/tsim/statevector/metal/runtime.py +++ b/src/tsim/statevector/metal/runtime.py @@ -9,6 +9,7 @@ import functools import os +import threading from dataclasses import dataclass from types import ModuleType from typing import Any @@ -60,6 +61,12 @@ def _after_fork_in_child() -> None: os.register_at_fork(after_in_child=_after_fork_in_child) +FORKED = ( + "This process was forked after its parent used the GPU, so it samples on the " + "CPU; start it with the 'spawn' or 'forkserver' method to sample on the GPU." +) + + def forked_after_init() -> bool: """Return whether this process was forked after its parent initialized Metal.""" return _forked_after_init @@ -126,6 +133,9 @@ class Buffer: """ _pool: list[Any] = [] +# NOTE: reentrant, since a buffer can be freed (``_Owner.__del__``) while this +# thread searches the pool. +_pool_lock = threading.RLock() class _Owner: @@ -148,8 +158,9 @@ def __init__(self, handle: Any, nbytes: int, recycle: bool = False) -> None: def __del__(self) -> None: if self.recycle and _pool is not None: - _pool.append(self.handle) - del _pool[:-POOL_BUFFERS] + with _pool_lock: + _pool.append(self.handle) + del _pool[:-POOL_BUFFERS] def buffer(count: int, dtype, recycle: bool = False) -> Buffer: @@ -163,10 +174,12 @@ def buffer(count: int, dtype, recycle: bool = False) -> Buffer: nbytes = max(int(count) * dtype.itemsize, 16) handle = None if recycle: - for k, pooled in enumerate(_pool): - if nbytes <= pooled.length() <= 2 * nbytes: - handle = _pool.pop(k) - break + with _pool_lock: + fits = [b for b in _pool if nbytes <= b.length() <= 2 * nbytes] + # NOTE: by identity, in case a buffer freed meanwhile trimmed the pool. + if fits and any(b is fits[0] for b in _pool): + handle = fits[0] + _pool[:] = [b for b in _pool if b is not handle] if handle is None: handle = device().newBufferWithLength_options_( nbytes, metal().MTLResourceStorageModeShared diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index c48b1a5b..fef356f2 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -30,6 +30,7 @@ from __future__ import annotations +import threading import time import warnings from dataclasses import dataclass, field @@ -70,6 +71,9 @@ DEVIATION_TOLERANCE = 1e-3 """Normalization deviation that float32 amplitudes tolerate without a warning.""" +MAX_CALL_SHOTS = 1 << 31 +"""Most shots of one GPU call; larger calls run in parts (the counters are 32-bit).""" + DEFER_SHARE = 0.25 """Share of the expected faults that postselection cannot see, from which they are drawn on the accepted shots only, after the tapes.""" @@ -337,6 +341,10 @@ def __init__( """ runtime.check_simd_width() assert self.supports(host.layout) + self._args = (host, channels, num_detectors, seed) + # NOTE: a call's batches share the sampler's buffers, so calls from several + # threads run one at a time. + self._lock = threading.RLock() self._host = host layout = host.layout self._num_detectors = num_detectors @@ -353,6 +361,13 @@ def __init__( self.gpu_seconds = 0.0 """GPU time of the command buffers so far (for benchmarks).""" + device = "metal" + """Where this sampler samples (``SamplerStats.device``).""" + + def __reduce__(self): + """Pickle the CPU sampler and the inputs; the GPU state is rebuilt on load.""" + return _restore, (*self._args, self._rng.bit_generator.state) + @property def layout(self): """Return the shot words' layout.""" @@ -746,8 +761,22 @@ def sample( As ``ShotSampler.sample``: with a ``mask`` only the shots that pass postselection are returned. ``batch_size`` bounds the shots of a batch. - Out of GPU memory, the call warns and samples on the CPU. + Out of GPU memory, and in a process forked after its parent used the + GPU, the call warns and samples on the CPU. """ + + def on_cpu(): + return self._host.sample( + shots, + columns, + bit_packed=bit_packed, + batch_size=batch_size, + mask=mask, + flip=flip, + ) + + if _forked(): + return on_cpu() try: rows, _, deviation = self.sample_accepted( shots, @@ -757,18 +786,8 @@ def sample( bools=not bit_packed, ) except MemoryError as error: - warnings.warn( - f"{error}; sampling on the CPU. Fewer shots per call fit.", - stacklevel=4, - ) - return self._host.sample( - shots, - columns, - bit_packed=bit_packed, - batch_size=batch_size, - mask=mask, - flip=flip, - ) + _warn_out_of_memory(error) + return on_cpu() if not bit_packed: return [_columns(rows, c) for c in columns], deviation num_outputs = self.layout.num_outputs @@ -804,43 +823,62 @@ def _run_accepted( slot.counts.array[:] = 0 slot.counts.array[4] = capacity commands = [] - for k, start in enumerate(starts): - n = min(batch, shots - start) - fixed.array[k] = n - encoder = runtime.Encoder() - self._encode_batch( - encoder, module, slot, n, start, constant, expectations, launch - ) - self._encode_pack( - encoder, - module, - slot, - n, - split, - expectations, - rows, - values, - (fixed, 4 * k), - ) - commands.append(encoder.commit()) + try: + for k, start in enumerate(starts): + n = min(batch, shots - start) + fixed.array[k] = n + encoder = runtime.Encoder() + self._encode_batch( + encoder, module, slot, n, start, constant, expectations, launch + ) + self._encode_pack( + encoder, + module, + slot, + n, + split, + expectations, + rows, + values, + (fixed, 4 * k), + ) + commands.append(encoder.commit()) + except BaseException: + self._settle(commands) + raise self._wait_all(commands) def _wait_all(self, commands: list[Any]) -> None: - """Wait for a call's command buffers, all of them even if one fails. + """Wait for a call's command buffers; on any error, let them all finish first. Raises: - RuntimeError: If one failed; the batch buffers are then dropped. + MemoryError: If one ran out of GPU memory. + RuntimeError: If one failed otherwise. """ try: for command in commands: self._wait(command) except BaseException: - for command in commands: - command.waitUntilCompleted() - self._slots.clear() + self._settle(commands) raise + def _settle(self, commands: list[Any]) -> None: + """Drop the batch buffers and wait for ``commands``, through further interrupts. + + A call that fails or is interrupted leaves its batches on the GPU; they + must finish before the next call reuses (or the pool hands out) a buffer + that they write. + """ + self._slots.clear() + for command in commands: + while True: + try: + command.waitUntilCompleted() + break + except KeyboardInterrupt: + continue + def sample_accepted( self, shots: int, @@ -871,6 +909,46 @@ def sample_accepted( MemoryError: If the GPU runs out of memory. """ + if shots > MAX_CALL_SHOTS: + parts = [ + self.sample_accepted( + min(MAX_CALL_SHOTS, shots - start), + mask=mask, + flip=flip, + expectations=expectations, + batch_size=batch_size, + bools=bools, + ) + 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), + ) + with self._lock: + return self._sample_call( + shots, + mask=mask, + flip=flip, + expectations=expectations, + batch_size=batch_size, + bools=bools, + ) + + def _sample_call( + self, + shots: int, + *, + mask: np.ndarray | None = None, + flip: np.ndarray | None = None, + expectations: bool = False, + batch_size: int | None = None, + bools: bool = False, + ) -> tuple[np.ndarray, np.ndarray | None, float]: + """One call of ``sample_accepted``, of at most ``MAX_CALL_SHOTS`` shots.""" module = self._module(mask) nb, no = -(-len(module.columns) // 8), len(self._observables) if not shots: @@ -943,7 +1021,7 @@ def _unpack(self, module: _Module, rows: runtime.Buffer, total: int) -> np.ndarr -(-threads // 256), 256, ) - self._wait(encoder.commit()) + self._wait_all([encoder.commit()]) return out.array[: total * num_outputs].reshape(total, num_outputs).view(bool) def _direct_expectations( @@ -961,10 +1039,41 @@ def _direct_expectations( def direct_bools(self, shots: int, flip: np.ndarray | None = None) -> np.ndarray: """Return every output bit of a program without components, as booleans.""" - rows, _, _ = self.sample_accepted(shots, flip=flip, bools=True) + if _forked(): + return self._host.direct_bools(shots, flip) + try: + rows, _, _ = self.sample_accepted(shots, flip=flip, bools=True) + except MemoryError as error: + _warn_out_of_memory(error) + return self._host.direct_bools(shots, flip) return rows +def _forked() -> bool: + """Return whether this process was forked after Metal's setup, with a warning.""" + if not runtime.forked_after_init(): + return False + warnings.warn(runtime.FORKED, stacklevel=5) + return True + + +def _warn_out_of_memory(error: MemoryError) -> None: + warnings.warn( + f"{error}; sampling on the CPU. Fewer shots per call fit.", stacklevel=5 + ) + + +def _restore( + host, channels, num_detectors, seed, state +) -> MetalShotSampler | ShotSampler: + """Rebuild a pickled ``MetalShotSampler``, or its CPU sampler where Metal is missing.""" + if not runtime.available(): + return host + shots = MetalShotSampler(host, channels, num_detectors, seed) + shots._rng.bit_generator.state = state + 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) diff --git a/test/integration/test_metal_sampler.py b/test/integration/test_metal_sampler.py index a8e85ae3..ab8aa52d 100644 --- a/test/integration/test_metal_sampler.py +++ b/test/integration/test_metal_sampler.py @@ -1,6 +1,8 @@ """The Metal executor end to end, against the stabilizer-rank strategy and the CPU executor.""" import multiprocessing +import pickle +import threading import warnings from test.helpers.gen import gen_stim_circuit from test.helpers.statevector import CIRCUITS, assert_counts_match @@ -790,3 +792,143 @@ def sample(bit_packed, postselection_mask=postselection_mask): ) assert detectors.flags.c_contiguous np.testing.assert_array_equal(detectors, bools[:, : circuit.num_detectors]) + + +def test_threads_sharing_a_gpu_sampler_get_their_own_rows(): + # NOTE: a call's batches share the sampler's buffers; four threads at once + # must not see each other's rows. + circuit = Circuit(NOISY_DETECTORS) + sampler = _gpu(circuit.compile_detector_sampler(seed=24)) + mask = np.zeros(150, dtype=bool) + mask[:3] = True + failures = [] + + def work(masked): + for _ in range(6): + rows = sampler.sample( + 20_000, + batch_size=1000, + append_observables=True, + postselection_mask=mask if masked else None, + ) + if not masked and len(rows) != 20_000: + failures.append(f"{len(rows)} rows") + if masked and rows[:, :3].any(): + failures.append("a masked detector is set") + if (rows[:, 150] != rows[:, 149] ^ rows[:, 143]).any(): + failures.append("an observable disagrees with its detectors") + + threads = [threading.Thread(target=work, args=(i % 2,)) for i in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert not failures, failures[:5] + + +def test_an_interrupted_call_leaves_the_next_call_alone(monkeypatch): + # NOTE: a slow width-11 tape and kernels compiled beforehand, so the + # interrupted call's batches are still on the GPU when the next call starts. + width = 11 + lines = [str(_random_unitary_circuit(width, 30, seed=0))] + lines += [f"DETECTOR rec[{i - width}]" for i in range(width)] + circuit = Circuit("\n".join(lines)) + sampler = _gpu(circuit.compile_detector_sampler(seed=25)) + shots = _gpu_shots(sampler) + all_detectors = np.ones(width, dtype=bool) + mask = np.zeros(width, dtype=bool) + mask[1] = True + for warm in (mask, all_detectors): + sampler.sample(1 << 17, batch_size=1 << 16, postselection_mask=warm) + encode_pack = MetalShotSampler._encode_pack + calls = [] + + def interrupt_late(self, *args, **kwargs): + calls.append(None) + if len(calls) == 20: + raise KeyboardInterrupt + return encode_pack(self, *args, **kwargs) + + monkeypatch.setattr(MetalShotSampler, "_encode_pack", interrupt_late) + with pytest.raises(KeyboardInterrupt): + sampler.sample(1 << 23, batch_size=1 << 16, postselection_mask=all_detectors) + monkeypatch.setattr(MetalShotSampler, "_encode_pack", encode_pack) + state = shots._rng.bit_generator.state + rows = sampler.sample(1 << 18, batch_size=1 << 16, postselection_mask=mask) + shots._rng.bit_generator.state = state + expected = sampler.sample(1 << 18, batch_size=1 << 16, postselection_mask=mask) + np.testing.assert_array_equal(rows, expected) + + +SAMPLER_BEFORE_FORK = None + + +def _sample_in_child(_): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert SAMPLER_BEFORE_FORK is not None + rows = SAMPLER_BEFORE_FORK.sample(100) + return rows.shape, [str(w.message) for w in caught] + + +def test_samplers_compiled_before_a_fork_sample_on_the_cpu_in_the_child(): + global SAMPLER_BEFORE_FORK + SAMPLER_BEFORE_FORK = _gpu(Circuit(T_STATE_X).compile_detector_sampler(seed=1)) + with multiprocessing.get_context("fork").Pool(1) as pool: + shape, messages = pool.map_async(_sample_in_child, [0]).get(timeout=60)[0] + assert shape == (100, 0) + assert any("forked" in m for m in messages) + + +def test_gpu_samplers_pickle_and_continue_their_streams(): + circuit = Circuit(NOISY_DETECTORS) + sampler = _gpu(circuit.compile_detector_sampler(seed=26)) + sampler.sample(1000) + loaded = pickle.loads(pickle.dumps(sampler)) # noqa: S301 + assert isinstance(loaded._shots, MetalShotSampler) + np.testing.assert_array_equal(loaded.sample(5000), sampler.sample(5000)) + + +def test_noise_beyond_the_gpu_tables_samples_on_the_cpu(): + # NOTE: a chain of 8 correlated errors has 255 outcomes, more than the 127 of + # a GPU noise record. + chain = ["E(0.01) X0", *(f"ELSE_CORRELATED_ERROR(0.01) X{k}" for k in range(1, 8))] + text = "\n".join( + ["R 0 1 2 3 4 5 6 7 8", "H 8", "T 8", "H 8", *chain, "M 0 1 2 3 4 5 6 7 8"] + + [f"DETECTOR rec[{k - 9}]" for k in range(9)] + ) + with pytest.warns(UserWarning, match="cannot sample this circuit"): + sampler = Circuit(text).compile_detector_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + assert sampler.sample(10).shape == (10, 9) + + +def test_direct_outputs_out_of_gpu_memory_fall_back_to_the_cpu(monkeypatch): + circuit = Circuit("R 0\nX_ERROR(0.1) 0\nM 0\nDETECTOR rec[-1]") + sampler = _gpu(circuit.compile_detector_sampler(seed=1)) + + def out_of_memory(*args, **kwargs): + raise MemoryError("the GPU ran out of memory") + + monkeypatch.setattr(_gpu_shots(sampler), "sample_accepted", out_of_memory) + with pytest.warns(UserWarning, match="out of memory; sampling on the CPU"): + rows = sampler.sample(20_000) + assert rows.shape == (20_000, 1) and abs(rows.mean() - 0.1) < 0.01 + + +def test_calls_larger_than_the_counters_run_in_parts(monkeypatch): + # NOTE: the 32-bit counters hold MAX_CALL_SHOTS; a larger call is the calls of + # its parts. + monkeypatch.setattr(metal_sampler, "MAX_CALL_SHOTS", 1000) + mask = np.zeros(150, dtype=bool) + mask[:3] = True + circuit = Circuit(NOISY_DETECTORS) + whole = _gpu(circuit.compile_detector_sampler(seed=27)).sample( + 2500, postselection_mask=mask, bit_packed=True + ) + parts = _gpu(circuit.compile_detector_sampler(seed=27)) + expected = [ + parts.sample(n, postselection_mask=mask, bit_packed=True) + for n in (1000, 1000, 500) + ] + np.testing.assert_array_equal(whole, np.concatenate(expected)) From cfbd2fcf51b27842c1d6f9e87825a53e4f34c5ba Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 10 Nov 2026 20:35:30 +0100 Subject: [PATCH 11/38] fix(metal): let pyright pass where PyObjC's Metal is not installed Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/metal/runtime.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tsim/statevector/metal/runtime.py b/src/tsim/statevector/metal/runtime.py index abf0d298..c0b99815 100644 --- a/src/tsim/statevector/metal/runtime.py +++ b/src/tsim/statevector/metal/runtime.py @@ -25,7 +25,7 @@ def metal() -> ModuleType: ImportError: If ``pyobjc-framework-Metal`` is not installed. """ - import Metal + import Metal # pyright: ignore[reportMissingImports] # the metal extra, macOS only return Metal From a8896a1c0ff1ee0b5ade3535d4b166453c3623af Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 10 Nov 2026 20:48:15 +0100 Subject: [PATCH 12/38] feat(metal): sample on Apple GPUs by default, and say where a sampler samples - On Apple silicon, Tsim depends on pyobjc-framework-Metal (about 7 MB), so the plain install samples on the GPU; the `metal` extra is gone. - The samplers' repr says where they sample (`samples on metal`). - The docs, the user skill and AGENTS.md say that the GPU computes outcome probabilities in float32, that its seeds depend on `batch_size`, and that `TSIM_DEVICE=cpu` keeps the samplers on the CPU. - Review fixes: `os.register_at_fork` is guarded (Windows has none), and batches of wide rows stay below 512 MB, so the kernels' 32-bit offsets do not wrap. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 7 ++--- CHANGELOG.md | 10 +++---- README.md | 5 ++-- docs/index.mdx | 5 ++-- docs/install.mdx | 27 ++++++++++++------- pyproject.toml | 6 ++--- skills/simulating-circuits-with-tsim/SKILL.md | 7 ++++- src/tsim/circuit.py | 18 +++++++------ src/tsim/sampler.py | 11 +++++--- src/tsim/statevector/metal/__init__.py | 5 ++-- src/tsim/statevector/metal/runtime.py | 5 ++-- src/tsim/statevector/metal/sampler.py | 4 ++- test/integration/test_metal_sampler.py | 7 +++++ uv.lock | 8 +++--- 14 files changed, 74 insertions(+), 51 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 4d00d861..7b7df498 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -66,9 +66,10 @@ more than a handful of T gates: their compile memory grows exponentially. may change in any release; every other module is internal. - **Raw outputs:** detectors and observables are raw parities, unlike Stim's flips relative to a reference. The reference-sample flags change only that. -- **Seeds:** with `active_statevector`, a seed's samples do not depend on the thread count or - `batch_size`. Do not break this; noise blocks and per-shot draws are seeded by call and shot - index (`shots.py`). +- **Seeds:** with `active_statevector` on the CPU, a seed's samples do not depend on the + thread count or `batch_size`. Do not break this; noise blocks and per-shot draws are seeded + by call and shot index (`shots.py`). On a GPU (`statevector/metal/`) they depend on + `batch_size`, as documented, and differ from the CPU's. - **Direct outputs stay direct:** outputs that depend on the noise alone never go through a state vector. Most detectors of QEC circuits are direct, which is why the detector sampler is fast. diff --git a/CHANGELOG.md b/CHANGELOG.md index 111cf347..7427f4cb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,15 +9,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - `strategy="active_statevector"` samples on the GPU of an Apple-silicon Mac - when the `metal` extra is installed (`pip install "bloqade-tsim[metal]"`, which - installs PyObjC's Metal bindings) and every component has at most 11 active - qubits. The noise sampler and every statevector component run as Metal + whenever every component has at most 11 active qubits. On Apple silicon, Tsim + now depends on PyObjC's Metal bindings (`pyobjc-framework-Metal`, about 7 MB to + download); other platforms install nothing new. The noise sampler and every statevector component run as Metal kernels, in float32; postselection runs on the GPU, and only the kept shots' rows are written. A seed's samples on the GPU differ from the CPU's and depend on `batch_size`. The samplers fall back to the CPU, with a warning, for wider components, when the GPU is out of memory and in a process forked after its - parent used the GPU. `TSIM_DEVICE=cpu` keeps them on the CPU, and - `SamplerStats.device` says where a sampler samples. + parent used the GPU. `TSIM_DEVICE=cpu` keeps them on the CPU. The samplers' + repr (`samples on metal`) and `SamplerStats.device` say where they sample. - The docs pages have a menu to copy a page or open it in ChatGPT or Claude, and the README, the install page and the agent skill point to the docs MCP server at `https://bloqade.quera.com/tsim/mcp`. The README and the install page also diff --git a/README.md b/README.md index 2a04c640..b9c694c0 100644 --- a/README.md +++ b/README.md @@ -23,9 +23,8 @@ Or with [uv](https://docs.astral.sh/uv/): uv add bloqade-tsim ``` -The default sampling strategy runs on the CPU, or on the GPU of an Apple-silicon Mac with the -`metal` extra (`pip install "bloqade-tsim[metal]"`). To run the stabilizer-rank strategies on an -NVIDIA GPU, install the JAX CUDA extra: +The default sampling strategy runs on the CPU, and on the GPU of an Apple-silicon Mac (no extra +needed). To run the stabilizer-rank strategies on an NVIDIA GPU, install the JAX CUDA extra: ```bash pip install "bloqade-tsim[cuda13]" # or bloqade-tsim[cuda12] ``` diff --git a/docs/index.mdx b/docs/index.mdx index 571ef1dc..e51a566e 100644 --- a/docs/index.mdx +++ b/docs/index.mdx @@ -39,9 +39,8 @@ samples = detector_sampler.sample(shots=100) uv add bloqade-tsim ``` -The default sampling strategy runs on the CPU, or on the GPU of an Apple-silicon Mac with the -`metal` extra (`pip install "bloqade-tsim[metal]"`). To run the stabilizer-rank strategies on an -NVIDIA GPU, install the JAX CUDA extra: +The default sampling strategy runs on the CPU, and on the GPU of an Apple-silicon Mac (no extra +needed). To run the stabilizer-rank strategies on an NVIDIA GPU, install the JAX CUDA extra: ```bash uv add "bloqade-tsim[cuda13]" diff --git a/docs/install.mdx b/docs/install.mdx index 60ce201b..3879f918 100644 --- a/docs/install.mdx +++ b/docs/install.mdx @@ -37,19 +37,28 @@ pip install "bloqade-tsim[cuda13]" ## Apple GPUs -On an Apple-silicon Mac, the default strategy (`"active_statevector"`) can sample on the GPU -through Metal. It needs PyObjC's Metal bindings, which the `metal` extra installs: +On an Apple-silicon Mac, the default strategy (`"active_statevector"`) samples on the GPU +through Metal; the plain install includes PyObjC's Metal bindings for it. The samplers use the +GPU whenever every component of the circuit has at most 11 active qubits, and otherwise sample +on the CPU, with a warning. Postselection runs on the GPU, so heavily postselected circuits +gain the most. A sampler's repr says where it samples (`samples on metal` or `samples on +cpu`), and so does `sampler.stats.device`. + +The GPU computes outcome probabilities in float32, so a sample follows the CPU's distribution +to within float32 rounding. A seed's samples on the GPU differ from the CPU's and depend on +`batch_size`. + +### Sampling on the CPU instead + +Set `TSIM_DEVICE=cpu` to keep the samplers on the CPU, in double precision, for example to +reproduce samples from another machine: ```bash -uv add "bloqade-tsim[metal]" # or: pip install "bloqade-tsim[metal]" +TSIM_DEVICE=cpu python my_script.py ``` -With it installed, the samplers use the GPU whenever every component of the circuit has at -most 11 active qubits; otherwise they sample on the CPU, with a warning. The GPU computes in -float32, so a sample follows the CPU's distribution to within float32 rounding, and a seed's -samples differ from the CPU's (and depend on `batch_size`). Postselection runs on the GPU, so -heavily postselected circuits gain the most. Set `TSIM_DEVICE=cpu` to keep the samplers on the -CPU, for example to reproduce samples from another machine. +Tsim reads it whenever it compiles a sampler, so `os.environ["TSIM_DEVICE"] = "cpu"` before +`compile_sampler` or `compile_detector_sampler` works too. ## Agent skill diff --git a/pyproject.toml b/pyproject.toml index 39b4f7d7..5ef984d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,8 @@ dependencies = [ "stim>=1.15.0", "stim-timeline-text>=0.1.0", "zxparam>=0.1.0", + # NOTE: the statevector executor on Apple GPUs (about 7 MB). + "pyobjc-framework-Metal>=10.0; sys_platform == 'darwin' and platform_machine == 'arm64'", ] [project.urls] @@ -54,10 +56,6 @@ Issues = "https://github.com/QuEraComputing/tsim/issues" [project.optional-dependencies] cuda12 = ["jax[cuda12]>=0.4.38"] cuda13 = ["jax[cuda13]>=0.7.0"] -# The statevector executor on Apple GPUs (Apple silicon only). -metal = [ - "pyobjc-framework-Metal>=10.0; sys_platform == 'darwin' and platform_machine == 'arm64'", -] [build-system] requires = ["hatchling"] diff --git a/skills/simulating-circuits-with-tsim/SKILL.md b/skills/simulating-circuits-with-tsim/SKILL.md index e5cfb39b..e7023925 100644 --- a/skills/simulating-circuits-with-tsim/SKILL.md +++ b/skills/simulating-circuits-with-tsim/SKILL.md @@ -189,8 +189,13 @@ yourself. default. Use `numba.set_num_threads(k)` before `sample` (or `NUMBA_NUM_THREADS=k` before importing Tsim) and leave cores free on a shared machine. Do not run several multi-threaded Tsim processes side by side; they oversubscribe the CPU. +- **Apple GPUs:** on an Apple-silicon Mac the default strategy samples on the GPU when every + component has at most 11 active qubits, often 5 to 10 times faster than all CPU cores. It + computes outcome probabilities in float32; set `TSIM_DEVICE=cpu` before compiling to stay on + the CPU, in double precision. The sampler's repr says where it samples (`samples on metal`). - **Seeds:** a `seed` reproduces samples for the same Tsim and dependency versions and - backend. Thread count and `batch_size` do not affect the samples. + backend. On the CPU, thread count and `batch_size` do not affect the samples; on a GPU, + `batch_size` does. ## Postselection diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index 22546a99..0b64c357 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -955,11 +955,12 @@ def compile_sampler( non-direct component on a dense state vector over its active qubits, or a stabilizer rank decomposition strategy ("cat5", "bss", "cutting"). "active_statevector" samples on an Apple GPU - when PyObjC's Metal sees one (the ``metal`` extra) and every + when Metal sees one (Apple silicon) and every component has at most 11 active qubits (in float32), and - otherwise on the CPU, on ``numba.get_num_threads()`` threads; - ``TSIM_DEVICE=cpu`` keeps it on the CPU. The stabilizer rank - strategies are provisional. + otherwise on the CPU, on ``numba.get_num_threads()`` threads. + The GPU's outcome probabilities are float32; ``TSIM_DEVICE=cpu`` + keeps the sampler on the CPU, in double precision. The + stabilizer rank strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed reproduces its samples only with the same versions of Tsim and its dependencies (such as NumPy and JAX) and the same backend (CPU or GPU). @@ -996,11 +997,12 @@ def compile_detector_sampler( non-direct component on a dense state vector over its active qubits, or a stabilizer rank decomposition strategy ("cat5", "bss", "cutting"). "active_statevector" samples on an Apple GPU - when PyObjC's Metal sees one (the ``metal`` extra) and every + when Metal sees one (Apple silicon) and every component has at most 11 active qubits (in float32), and - otherwise on the CPU, on ``numba.get_num_threads()`` threads; - ``TSIM_DEVICE=cpu`` keeps it on the CPU. The stabilizer rank - strategies are provisional. + otherwise on the CPU, on ``numba.get_num_threads()`` threads. + The GPU's outcome probabilities are float32; ``TSIM_DEVICE=cpu`` + keeps the sampler on the CPU, in double precision. The + stabilizer rank strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed reproduces its samples only with the same versions of Tsim and its dependencies (such as NumPy and JAX) and the same backend (CPU or GPU). diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 63a56304..d8bc35b6 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -386,11 +386,12 @@ def __init__( non-direct component on a dense state vector over its active qubits, or a stabilizer rank decomposition strategy ("cat5", "bss", "cutting"). "active_statevector" samples on an Apple GPU - when PyObjC's Metal sees one (the ``metal`` extra) and every + when Metal sees one (Apple silicon) and every component has at most 11 active qubits (in float32), and - otherwise on the CPU, on ``numba.get_num_threads()`` threads; - ``TSIM_DEVICE=cpu`` keeps it on the CPU. The stabilizer rank - strategies are provisional. + otherwise on the CPU, on ``numba.get_num_threads()`` threads. + The GPU's outcome probabilities are float32; ``TSIM_DEVICE=cpu`` + keeps the sampler on the CPU, in double precision. The + stabilizer rank strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed reproduces its samples only with the same versions of Tsim and its dependencies (such as NumPy and JAX) and the same backend (CPU or GPU). @@ -873,6 +874,8 @@ def _format_bytes(n: int) -> str: f"(width ≤ {stats.active_width}, " f"{stats.num_actions} actions)" ) + if self._shots is not None: + parts.append(f"samples on {stats.device}") if c_graphs: parts.append(f"{sum(c_graphs)} graphs") parts += [ diff --git a/src/tsim/statevector/metal/__init__.py b/src/tsim/statevector/metal/__init__.py index bf4c629f..ab4273a6 100644 --- a/src/tsim/statevector/metal/__init__.py +++ b/src/tsim/statevector/metal/__init__.py @@ -4,7 +4,6 @@ by ``tsim.statevector.gpu.tape``) run as Metal kernels that PyObjC compiles at runtime (``runtime``); ``MetalShotSampler`` (``sampler``) replaces ``ShotSampler`` behind the compiled samplers whenever PyObjC sees an Apple GPU -and every component fits. Amplitudes are complex64. It needs -``pyobjc-framework-Metal`` (``pip install bloqade-tsim[metal]``) and Apple -silicon. +and every component fits. Amplitudes are complex64. It needs Apple silicon, +where Tsim depends on ``pyobjc-framework-Metal``. """ diff --git a/src/tsim/statevector/metal/runtime.py b/src/tsim/statevector/metal/runtime.py index c0b99815..e73034ff 100644 --- a/src/tsim/statevector/metal/runtime.py +++ b/src/tsim/statevector/metal/runtime.py @@ -25,7 +25,7 @@ def metal() -> ModuleType: ImportError: If ``pyobjc-framework-Metal`` is not installed. """ - import Metal # pyright: ignore[reportMissingImports] # the metal extra, macOS only + import Metal # pyright: ignore[reportMissingImports] # Apple silicon only return Metal @@ -58,7 +58,8 @@ def _after_fork_in_child() -> None: _forked_after_init = True -os.register_at_fork(after_in_child=_after_fork_in_child) +if hasattr(os, "register_at_fork"): # not on Windows + os.register_at_fork(after_in_child=_after_fork_in_child) FORKED = ( diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index fef356f2..aab061b3 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -543,7 +543,9 @@ def _constant(self, flip: np.ndarray | None) -> runtime.Buffer: def _batch_shots(self, shots: int, batch_size: int | None) -> int: """Return the shots of a call's batches.""" - cap = max(1 << 16, BATCH_BYTES // (8 * self.layout.num_words)) + # NOTE: at most 2^16 words a row, so a batch stays below 512 MB, and the + # kernels' 32-bit offsets within it do not wrap. + cap = max(1 << 10, BATCH_BYTES // (8 * self.layout.num_words)) return max(1, min(batch_size or BATCH_SHOTS, BATCH_SHOTS, cap, shots)) def _slot(self, shots: int, expectations: bool) -> _Slot: diff --git a/test/integration/test_metal_sampler.py b/test/integration/test_metal_sampler.py index ab8aa52d..8ec3d167 100644 --- a/test/integration/test_metal_sampler.py +++ b/test/integration/test_metal_sampler.py @@ -932,3 +932,10 @@ def test_calls_larger_than_the_counters_run_in_parts(monkeypatch): for n in (1000, 1000, 500) ] np.testing.assert_array_equal(whole, np.concatenate(expected)) + + +def test_the_repr_says_where_a_sampler_samples(): + circuit = Circuit(T_STATE_X) + assert "samples on metal" in repr(_gpu(circuit.compile_detector_sampler(seed=1))) + cpu = _cpu(circuit.compile_detector_sampler, seed=1) + assert "samples on cpu" in repr(cpu) and cpu.stats.device == "cpu" diff --git a/uv.lock b/uv.lock index c9dd6b3f..ff76fd32 100644 --- a/uv.lock +++ b/uv.lock @@ -130,6 +130,7 @@ dependencies = [ { name = "numba" }, { name = "numpy" }, { name = "psutil" }, + { name = "pyobjc-framework-metal", marker = "platform_machine == 'arm64' and sys_platform == 'darwin'" }, { name = "stim" }, { name = "stim-timeline-text" }, { name = "zxparam" }, @@ -142,9 +143,6 @@ cuda12 = [ cuda13 = [ { name = "jax", extra = ["cuda13"] }, ] -metal = [ - { name = "pyobjc-framework-metal", marker = "platform_machine == 'arm64' and sys_platform == 'darwin'" }, -] [package.dev-dependencies] dev = [ @@ -194,12 +192,12 @@ requires-dist = [ { name = "numba", specifier = ">=0.61" }, { name = "numpy", specifier = ">=2.0" }, { name = "psutil", specifier = ">=5.9.0" }, - { name = "pyobjc-framework-metal", marker = "platform_machine == 'arm64' and sys_platform == 'darwin' and extra == 'metal'", specifier = ">=10.0" }, + { name = "pyobjc-framework-metal", marker = "platform_machine == 'arm64' and sys_platform == 'darwin'", specifier = ">=10.0" }, { name = "stim", specifier = ">=1.15.0" }, { name = "stim-timeline-text", specifier = ">=0.1.0" }, { name = "zxparam", specifier = ">=0.1.0" }, ] -provides-extras = ["cuda12", "cuda13", "metal"] +provides-extras = ["cuda12", "cuda13"] [package.metadata.requires-dev] dev = [ From 3d6ad1a010a9d2e78a57ce9cc9945240334462e2 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 10 Nov 2026 21:01:49 +0100 Subject: [PATCH 13/38] docs(skills): say how to benchmark Tsim on Apple-silicon Macs Co-Authored-By: Claude Opus 5.5 --- .claude/skills/benchmarking-tsim/SKILL.md | 26 +++++++++++++++++++++++ .claude/skills/profiling-tsim/SKILL.md | 8 ++++--- 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/.claude/skills/benchmarking-tsim/SKILL.md b/.claude/skills/benchmarking-tsim/SKILL.md index 9aa77409..16bfa12a 100644 --- a/.claude/skills/benchmarking-tsim/SKILL.md +++ b/.claude/skills/benchmarking-tsim/SKILL.md @@ -82,6 +82,32 @@ uv run tsim-dev bench --side .:threads=4 --side origin/main:threads=4 - Compile time is compared too; when the sides differ only in threads, read the rate column. +## Apple-silicon Macs (Metal) + +On a Mac with an M-series chip, the default strategy samples on the GPU through Metal +whenever every component has at most 11 active qubits (k_max <= 11). A plain `uv sync` +installs it. That is what users of a MacBook get, and it is much faster than the CPU: on an +M4 Pro, 26-96x one CPU thread and 4-19x all ten performance cores, bit-packed. End-to-end +cultivation d=5 ran at 125M shots/s, 19x the ten cores. Circuits without noise whose CPU +trajectories share well, such as qv_q10, run at the CPU's rate. + +The bench sets `TSIM_DEVICE=cpu` by default, so on a Mac a plain run measures the CPU only. +Say which you measured, and on a Mac measure both: + +```bash +uv run tsim-dev bench --device gpu --ref REF # a GPU change; REF must have the Metal backend +uv run tsim-dev bench --side .:device=gpu --side .:device=cpu,threads=10 # GPU vs all cores +``` + +- The `device` column must say `metal` on the GPU side. It says `cpu` where Metal falls + back: a component wider than 11, noise beyond the GPU tables, no GPU memory. +- A reference without the Metal backend samples on the CPU even with `--device gpu`; its + rows end in `DEVICE: cpu -> metal`, and the ratio compares devices, not code. +- Compare the GPU with all CPU cores (`--threads N`, from the `NOTE` line), not one. +- The GPU's own time per kernel: set `tsim.statevector.metal.runtime.PROFILE = {}` before + a `sample` call, which then runs each dispatch alone and adds its GPU seconds under the + kernel's name (`noise`, `tape`, `gather_rows`, ...). + ## What the command controls for you - One thread, pinned to a core (`--threads N` to change; say so when you report). diff --git a/.claude/skills/profiling-tsim/SKILL.md b/.claude/skills/profiling-tsim/SKILL.md index a95fbf59..7ee9c0a4 100644 --- a/.claude/skills/profiling-tsim/SKILL.md +++ b/.claude/skills/profiling-tsim/SKILL.md @@ -38,8 +38,9 @@ Tsim is slow: performance cores by default (`numba.get_num_threads()`); the `NOTE` line says how many. Quote all-core rates with `--threads N`. - **Device.** The harness sets `TSIM_DEVICE=cpu`, so every side samples on the - CPU unless you pass `--device gpu` (or `--side REF:device=gpu`). With a GPU - backend installed, `--device gpu` measures what users get by default. The + CPU unless you pass `--device gpu` (or `--side REF:device=gpu`). On an + Apple-silicon Mac, Tsim samples on the GPU by default for k_max <= 11, 4-19x + faster than all CPU cores; there `--device gpu` measures what users get. The `device` column says where each row sampled: a GPU falls back to the CPU for components wider than it runs, and a row that ends in `DEVICE: cpu -> metal` compares a GPU with the CPU. @@ -61,7 +62,8 @@ uv run tsim-dev profile cultivation_d5 --seconds 5 The run is a fresh process, one pinned thread, warm numba cache, under a memory cap, with bit-packed output on the CPU (`--threads`, `--device` and `--sampler` as in step 2). py-spy sees the host only: on a GPU the sampling phase shows the wait for -the device, and the time inside its kernels needs the backend's own timers. +the device. For Metal's time per kernel, see "Apple-silicon Macs" in the +`benchmarking-tsim` skill (`runtime.PROFILE`). The output splits samples by phase (`compile`, `sample`) and prints three tables: - **innermost Tsim frame (kernel callers):** read this first. Jitted numba code has From 23e2f7f74e71e93f6e8a3c5de7c8b4207a1816fe Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 04:00:37 -0700 Subject: [PATCH 14/38] refactor(metal): pack whole shot words, without piece tables The sampler always returns every output in order, so noise_kept's pack_rows copies output word k and clears the bits past the last output, instead of reading a piece table per word. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/metal/noise.py | 32 ++++++++++++--------------- src/tsim/statevector/metal/sampler.py | 31 ++++++++------------------ 2 files changed, 23 insertions(+), 40 deletions(-) diff --git a/src/tsim/statevector/metal/noise.py b/src/tsim/statevector/metal/noise.py index 9d6d7c95..9eb2c74e 100644 --- a/src/tsim/statevector/metal/noise.py +++ b/src/tsim/statevector/metal/noise.py @@ -153,8 +153,7 @@ device const float* seg_log_keep [[buffer(6)]], device const Record* records [[buffer(7)]], device const Flip* overflow [[buffer(8)]], device const uint* n [[buffer(9)]], device const uint* index [[buffer(10)]], - device uchar* dst [[buffer(11)]], device const uint2* pieces [[buffer(12)]], - device const uint* piece_start [[buffer(13)]], device const uint* cursor [[buffer(14)]], + device uchar* dst [[buffer(11)]], device const uint* cursor [[buffer(12)]], threadgroup uint2* events [[threadgroup(0)]], uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], uint nt [[threads_per_threadgroup]]) @@ -176,7 +175,7 @@ add_faults(sw, events, &num_events, args.max_events, ns, args.seed ^ ((u64)bid * 0xD1B54A32D192ED03UL), args.num_segments, seg, records, overflow, tid, nt); - pack_rows(sw, ns, pieces, piece_start, (threadgroup uchar*)events, dst, ((long)*cursor + base) * {NB}, tid, nt); + pack_rows(sw, ns, (threadgroup uchar*)events, dst, ((long)*cursor + base) * {NB}, tid, nt); }} """ @@ -216,22 +215,16 @@ }} }} -// The returned outputs of the block's ``ns`` rows ``sw[w * S + s]``, bit-packed -// into {NB} bytes per row, contiguously at ``dst`` (``gpu.rows.pieces``). -inline void pack_rows(threadgroup const u64* sw, int ns, device const uint2* pieces, - device const uint* start, threadgroup uchar* stage, device uchar* dst, long at, uint tid, uint nt) +// The outputs of the block's ``ns`` rows ``sw[w * S + s]``, bit-packed into +// {NB} bytes per row (the bits past them cleared), contiguously at ``dst``. +inline void pack_rows(threadgroup const u64* sw, int ns, threadgroup uchar* stage, + device uchar* dst, long at, uint tid, uint nt) {{ for (uint i = tid; i < {NK} * {S}; i += nt) {{ int k = i / {S}, s = i - k * {S}; if (s >= ns) continue; - u64 v = 0; - for (uint p = start[k]; p < start[k + 1]; ++p) {{ - uint2 q = pieces[p]; - u64 x = sw[q.x * {S} + s] >> (q.y & 63); - int length = ((q.y >> 6) & 63) + 1; - if (length < 64) x &= (1UL << length) - 1; - v |= x << (q.y >> 12); - }} + u64 v = sw[k * {S} + s]; + if (k == {NK} - 1) v &= {LAST_WORD}; threadgroup uchar* at = stage + s * {NB} + 8 * k; for (int j = 0; j < 8; ++j) if (8 * k + j < {NB}) at[j] = uchar(v >> (8 * j)); }} @@ -360,12 +353,14 @@ def noise_tables( ) -def source(tables: NoiseTables, output_words: int, packed_bytes: int) -> str: +def source(tables: NoiseTables, output_words: int, num_outputs: int) -> str: """Return the MSL source of the ``noise`` and ``noise_kept`` kernels. - Without ``tables.shared_words`` there is no ``noise_kept``: the rows of a - block do not fit in threadgroup memory. + ``noise_kept`` packs rows of ``num_outputs`` outputs. Without + ``tables.shared_words`` there is no ``noise_kept``: the rows of a block do + not fit in threadgroup memory. """ + packed_bytes = -(-num_outputs // 8) if tables.shared_words: noise, kept = NOISE_SHARED, NOISE_KEPT else: @@ -376,4 +371,5 @@ def source(tables: NoiseTables, output_words: int, packed_bytes: int) -> str: S=tables.block_shots, NB=max(packed_bytes, 1), NK=max(-(-packed_bytes // 8), 1), + LAST_WORD=f"0x{(1 << (num_outputs % 64 or 64)) - 1:016X}UL", ) diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index aab061b3..4c622f70 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -40,7 +40,7 @@ from tsim.noise.channels import ChannelSampler from tsim.statevector.gpu.noise import NoiseGroups, deferred_channels, pass_rate -from tsim.statevector.gpu.rows import pieces as _pieces, select as _select +from tsim.statevector.gpu.rows import select as _select from tsim.statevector.metal import runtime from tsim.statevector.metal.codegen import PRELUDE, hex64, tape_function from tsim.statevector.metal.noise import noise_tables, source as noise_source @@ -271,10 +271,6 @@ class _Module: library: runtime.Library tb: int has_mask: bool - columns: np.ndarray - """The outputs that ``sample_accepted`` returns: all of them, in order.""" - pieces: runtime.Buffer - piece_start: runtime.Buffer screen: runtime.Buffer plan: ShotPlan psm_words: np.ndarray @@ -420,17 +416,16 @@ def _module(self, mask: np.ndarray | None) -> _Module: if mask is not None: bits[: len(mask)] = mask psm_words = pack_bits(bits, layout.output_words) - # NOTE: the rows keep the postselected outputs, as the CPU's do. - columns = np.arange(layout.num_outputs, dtype=np.int32) written = sorted({int(w) for p in plan.programs for w in p.out_words}) stores = "\n".join(f" row[{w}] = w[{w}];" for w in written) hits = " | ".join( f"(w[{w}] & {hex64(psm_words[w])})" for w in written if psm_words[w] ) - nb = -(-len(columns) // 8) + # NOTE: the rows keep every output, the postselected ones too, as the CPU's do. + nb = -(-layout.num_outputs // 8) source = ( PRELUDE - + noise_source(self._noise, layout.output_words, nb) + + noise_source(self._noise, layout.output_words, layout.num_outputs) + functions + _KERNELS.format( T1=T - 1, @@ -462,14 +457,10 @@ def _module(self, mask: np.ndarray | None) -> _Module: for name in names: library.pipeline(name) compile_seconds = time.perf_counter() - start - pieces, piece_start = _pieces(columns) compiled = _Module( library=library, tb=tb, has_mask=mask is not None, - columns=columns, - pieces=runtime.upload(pieces), - piece_start=runtime.upload(piece_start), screen=runtime.upload(plan.screen), plan=plan, psm_words=psm_words, @@ -658,7 +649,7 @@ def _encode_pack( """ library = module.library S = self._noise.block_shots - nb = -(-len(module.columns) // 8) + nb = -(-self.layout.num_outputs // 8) no = len(self._observables) if module.has_mask: words = -(-shots // 32) @@ -699,8 +690,6 @@ def _encode_pack( count, slot.index, rows, - module.pieces, - module.piece_start, cursor, ], blocks, @@ -952,7 +941,7 @@ def _sample_call( ) -> tuple[np.ndarray, np.ndarray | None, float]: """One call of ``sample_accepted``, of at most ``MAX_CALL_SHOTS`` shots.""" module = self._module(mask) - nb, no = -(-len(module.columns) // 8), len(self._observables) + 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 bools: @@ -1000,7 +989,7 @@ def _sample_call( expected = values.array[: total * no].reshape(total, no) if not self.layout.components: expected[:] = 0 - self._direct_expectations(packed, expected, module.columns) + self._direct_expectations(packed, expected) if bools: return self._unpack(module, rows, total), expected, deviation return packed, expected, deviation @@ -1026,16 +1015,14 @@ 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 _direct_expectations( - self, rows: np.ndarray, values: np.ndarray, columns: np.ndarray - ) -> None: + def _direct_expectations(self, rows: np.ndarray, values: np.ndarray) -> None: """Set the expectations of observables outside the components: ``±1``.""" layout = self.layout in_components = {i for c in layout.components for i in c.output_indices} for column, index in enumerate(self._observables): if int(index) in in_components: continue - k = int(np.searchsorted(columns, index)) + k = int(index) bit = (rows[:, k >> 3] >> np.uint8(k & 7)) & np.uint8(1) values[:, column] = 1.0 - 2.0 * bit.astype(np.float32) From 29861fb794e68637581e4677b8a321d70b28346b Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 04:01:12 -0700 Subject: [PATCH 15/38] refactor(metal): drop write-only timers and unused kernel arguments Nothing read compile_seconds, gpu_seconds or _Slot.shots; the tape kernel template never used OW or S, and NE always equalled NO; FAST_MATH was never toggled. runtime.PROFILE stays: the benchmarking skill uses it. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/metal/runtime.py | 6 +++--- src/tsim/statevector/metal/sampler.py | 24 +++--------------------- 2 files changed, 6 insertions(+), 24 deletions(-) diff --git a/src/tsim/statevector/metal/runtime.py b/src/tsim/statevector/metal/runtime.py index e73034ff..74fca304 100644 --- a/src/tsim/statevector/metal/runtime.py +++ b/src/tsim/statevector/metal/runtime.py @@ -206,8 +206,8 @@ def upload(values: np.ndarray) -> Buffer: class Library: """A compiled Metal library and its kernels' pipeline states.""" - def __init__(self, source: str, fast_math: bool = True) -> None: - """Compile ``source`` (Metal Shading Language). + def __init__(self, source: str) -> None: + """Compile ``source`` (Metal Shading Language) with fast math. Raises: RuntimeError: If it does not compile. @@ -215,7 +215,7 @@ def __init__(self, source: str, fast_math: bool = True) -> None: """ mtl = metal() options = mtl.MTLCompileOptions.new() - options.setFastMathEnabled_(fast_math) + options.setFastMathEnabled_(True) library, error = device().newLibraryWithSource_options_error_( source, options, None ) diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index 4c622f70..477d9eea 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -31,7 +31,6 @@ from __future__ import annotations import threading -import time import warnings from dataclasses import dataclass, field from typing import Any @@ -65,9 +64,6 @@ TAPE_THREADS = 128 """Threads per threadgroup of the tape kernel.""" -FAST_MATH = True -"""Whether the kernels compile with fast math.""" - DEVIATION_TOLERANCE = 1e-3 """Normalization deviation that float32 amplitudes tolerate without a warning.""" @@ -100,7 +96,7 @@ device u64* row = out + (long)row_index * {NW}; u64 w[{NW}]; for (int i = 0; i < {NW}; ++i) w[i] = row[i]; - float ev[{NE}]; + float ev[{NO}]; int grave = 0x7fffffff; {calls} if (t == 0) {{ @@ -274,7 +270,6 @@ class _Module: screen: runtime.Buffer plan: ShotPlan psm_words: np.ndarray - compile_seconds: float passing: float """Estimated share of shots that pass the screen.""" splits: dict[bool, tuple[_Launch, _Launch] | None] = field(default_factory=dict) @@ -290,7 +285,6 @@ class _Slot: output has room for. """ - shots: int words: runtime.Buffer queue: runtime.Buffer accepted: runtime.Buffer @@ -354,8 +348,6 @@ def __init__( self._constants: dict[bytes, runtime.Buffer] = {} self._slots: dict[tuple[int, bool], _Slot] = {} self._shares: dict[tuple, float] = {} - self.gpu_seconds = 0.0 - """GPU time of the command buffers so far (for benchmarks).""" device = "metal" """Where this sampler samples (``SamplerStats.device``).""" @@ -431,10 +423,7 @@ def _module(self, mask: np.ndarray | None) -> _Module: T1=T - 1, tb=tb, NW=layout.num_words, - OW=layout.output_words, NO=max(1, len(self._observables)), - NE=max(1, len(self._observables)), - S=self._noise.block_shots, NB=max(nb, 1), RB=8 * layout.num_words, NOUT=max(layout.num_outputs, 1), @@ -448,15 +437,13 @@ def _module(self, mask: np.ndarray | None) -> _Module: hits=hits or "0", ) ) - start = time.perf_counter() - library = runtime.Library(source, fast_math=FAST_MATH) + library = runtime.Library(source) names = ["noise", "tape", "count_bits", "list_bits", "gather_rows"] names += ["gather_float", "reset", "advance", "unpack_rows"] if self._noise.shared_words: names.append("noise_kept") for name in names: library.pipeline(name) - compile_seconds = time.perf_counter() - start compiled = _Module( library=library, tb=tb, @@ -464,7 +451,6 @@ def _module(self, mask: np.ndarray | None) -> _Module: screen=runtime.upload(plan.screen), plan=plan, psm_words=psm_words, - compile_seconds=compile_seconds, passing=self.passing(mask), ) self._modules[key] = compiled @@ -546,7 +532,6 @@ def _slot(self, shots: int, expectations: bool) -> _Slot: if slot is None: nw, no = self.layout.num_words, max(1, len(self._observables)) slot = _Slot( - shots=shots, words=runtime.buffer(shots * nw, np.uint64), queue=runtime.buffer(shots, np.uint32), accepted=runtime.buffer(-(-shots // 32), np.uint32), @@ -735,9 +720,6 @@ def _deviation(self, counts: np.ndarray) -> float: # NOTE: fast math records a NaN total as NaN (not inf), which must warn. return 0.0 if deviation <= DEVIATION_TOLERANCE else deviation - def _wait(self, commands: Any) -> None: - self.gpu_seconds += runtime.wait(commands) - def sample( self, shots: int, @@ -849,7 +831,7 @@ def _wait_all(self, commands: list[Any]) -> None: """ try: for command in commands: - self._wait(command) + runtime.wait(command) except BaseException: self._settle(commands) raise From f4c848dea1c0b9fe939f7ea9ea13cd2b29ed8f8d Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 04:02:03 -0700 Subject: [PATCH 16/38] refactor(metal): use the shared GPU helpers The tape kernel's glue, the noise split, the screen's pass rate, the shared-memory fit test, the deviation, the width limit and the fork warning now come from tsim.statevector.gpu; the generated kernels and dispatches are unchanged. codegen's tape_function and hex64 pass-throughs go. Co-Authored-By: Claude Opus 5.5 --- src/tsim/core/parse.py | 1 - src/tsim/noise/channels.py | 6 +- src/tsim/sampler.py | 6 +- src/tsim/statevector/metal/codegen.py | 10 -- src/tsim/statevector/metal/noise.py | 14 ++- src/tsim/statevector/metal/runtime.py | 6 -- src/tsim/statevector/metal/sampler.py | 143 ++++++------------------- src/tsim/statevector/shared.py | 5 +- test/integration/test_metal_sampler.py | 15 +-- 9 files changed, 57 insertions(+), 149 deletions(-) diff --git a/src/tsim/core/parse.py b/src/tsim/core/parse.py index 2e5c540a..d981c1df 100644 --- a/src/tsim/core/parse.py +++ b/src/tsim/core/parse.py @@ -171,7 +171,6 @@ def parse_stim_circuit(stim_circuit: stim.Circuit) -> GraphBuilder: name = instruction.name if name == "SHIFT_COORDS": - # TODO: handle visualization annotations in ZX diagrams continue diff --git a/src/tsim/noise/channels.py b/src/tsim/noise/channels.py index 0dfb1e87..050a846f 100644 --- a/src/tsim/noise/channels.py +++ b/src/tsim/noise/channels.py @@ -740,8 +740,7 @@ def sample_words( 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)}" + f"{num_samples} samples need {num_blocks} block seeds, got {len(seeds)}" ) t = self._tables _sample_blocks( @@ -773,8 +772,7 @@ def sample_screened_words( 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)}" + f"{num_samples} samples need {num_blocks} block seeds, got {len(seeds)}" ) if not (len(constant) == len(screen) == out.shape[1]): raise ValueError("constant and screen must cover every word of a row") diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index d8bc35b6..265adb0c 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -20,6 +20,7 @@ from tsim.core.graph import prepare_graph from tsim.core.types import CompiledComponent, CompiledProgram, SamplingStrategy from tsim.noise.channels import ChannelSampler +from tsim.statevector.gpu import FORKED from tsim.statevector.metal import runtime as metal_runtime from tsim.statevector.program import StatevectorComponent from tsim.statevector.shots import ShotSampler @@ -332,11 +333,12 @@ def _gpu_shots( if os.environ.get("TSIM_DEVICE", "").lower() == "cpu": return host if metal_runtime.forked_after_init(): - warnings.warn(metal_runtime.FORKED, stacklevel=5) + warnings.warn(FORKED, stacklevel=5) return host if not metal_runtime.available(): return host - from tsim.statevector.metal.sampler import MAX_WIDTH, MetalShotSampler + from tsim.statevector.gpu.tape import MAX_WIDTH + from tsim.statevector.metal.sampler import MetalShotSampler layout = host.layout if layout.num_words == 0: diff --git a/src/tsim/statevector/metal/codegen.py b/src/tsim/statevector/metal/codegen.py index 833eda5f..4fb39658 100644 --- a/src/tsim/statevector/metal/codegen.py +++ b/src/tsim/statevector/metal/codegen.py @@ -11,7 +11,6 @@ from __future__ import annotations from tsim.statevector.gpu import tape -from tsim.statevector.shared import LaneProgram METAL = tape.Dialect( inline="inline", @@ -90,12 +89,3 @@ atomic_fetch_max_explicit(deviation, d == d ? as_type(d) : 0x7f800000u, memory_order_relaxed); } """ - -hex64 = METAL.hex64 - - -def tape_function( - program: LaneProgram, index: int, tb: int, observables: dict[int, int] -) -> str: - """Return ``component_`` in Metal (``tape.tape_function``).""" - return tape.tape_function(program, index, tb, observables, METAL) diff --git a/src/tsim/statevector/metal/noise.py b/src/tsim/statevector/metal/noise.py index 9eb2c74e..fac5432c 100644 --- a/src/tsim/statevector/metal/noise.py +++ b/src/tsim/statevector/metal/noise.py @@ -18,7 +18,11 @@ from __future__ import annotations from tsim.noise.channels import ChannelSampler, WordTables -from tsim.statevector.gpu.noise import NoiseTables, noise_tables as _noise_tables +from tsim.statevector.gpu.noise import ( + NoiseTables, + noise_tables as _noise_tables, + words_fit, +) SHARED_BYTES = 14 * 1024 """Threadgroup memory of a block's words and events. @@ -346,10 +350,12 @@ def noise_tables( Rows too wide for ``SHARED_BYTES`` keep their words in device memory (``shared_words`` False), where the noise kernel XORs the faults in. """ - if 8 * 32 * num_words <= SHARED_BYTES - 8 * 256: - return _noise_tables(channels, flips, num_words, shared_bytes=SHARED_BYTES) return _noise_tables( - channels, flips, num_words, shared_bytes=SHARED_BYTES, shared_words=False + channels, + flips, + num_words, + shared_bytes=SHARED_BYTES, + shared_words=words_fit(num_words, SHARED_BYTES), ) diff --git a/src/tsim/statevector/metal/runtime.py b/src/tsim/statevector/metal/runtime.py index 74fca304..1e2d3cbc 100644 --- a/src/tsim/statevector/metal/runtime.py +++ b/src/tsim/statevector/metal/runtime.py @@ -62,12 +62,6 @@ def _after_fork_in_child() -> None: os.register_at_fork(after_in_child=_after_fork_in_child) -FORKED = ( - "This process was forked after its parent used the GPU, so it samples on the " - "CPU; start it with the 'spawn' or 'forkserver' method to sample on the GPU." -) - - def forked_after_init() -> bool: """Return whether this process was forked after its parent initialized Metal.""" return _forked_after_init diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index 477d9eea..22e49ebf 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -38,23 +38,17 @@ import numpy as np from tsim.noise.channels import ChannelSampler -from tsim.statevector.gpu.noise import NoiseGroups, deferred_channels, pass_rate +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.metal import runtime -from tsim.statevector.metal.codegen import PRELUDE, hex64, tape_function +from tsim.statevector.metal.codegen import METAL, PRELUDE from tsim.statevector.metal.noise import noise_tables, source as noise_source -from tsim.statevector.shared import pack_bits -from tsim.statevector.shots import ShotLayout, ShotPlan, ShotSampler - -MAX_WIDTH = 11 -"""Widest component the Metal executor runs: 32 threads of 64 amplitudes.""" +from tsim.statevector.shots import ShotPlan, ShotSampler BATCH_SHOTS = 1 << 20 """Shots of one batch (one command buffer).""" -AMPLITUDE_AXES = 5 -"""Log2 of the amplitudes per thread from which a shot takes more threads.""" - BATCH_BYTES = 256 << 20 """Bytes of shot words of one batch, which bound its shots below ``BATCH_SHOTS``.""" @@ -64,16 +58,9 @@ TAPE_THREADS = 128 """Threads per threadgroup of the tape kernel.""" -DEVIATION_TOLERANCE = 1e-3 -"""Normalization deviation that float32 amplitudes tolerate without a warning.""" - MAX_CALL_SHOTS = 1 << 31 """Most shots of one GPU call; larger calls run in parts (the counters are 32-bit).""" -DEFER_SHARE = 0.25 -"""Share of the expected faults that postselection cannot see, from which they -are drawn on the accepted shots only, after the tapes.""" - _KERNELS = r""" struct TapeArgs {{ long first_shot; u64 seed; int flags; }}; @@ -303,15 +290,7 @@ class MetalShotSampler: expectation. """ - @staticmethod - def supports(layout: ShotLayout) -> bool: - """Return whether every component of ``layout`` fits in ``MAX_WIDTH`` qubits. - - A layout without shot words (no outputs and no noise) stays on the CPU. - """ - return layout.num_words > 0 and all( - c.width <= MAX_WIDTH for c in layout.components - ) + supports = staticmethod(tape.supports) def __init__( self, @@ -375,50 +354,21 @@ def reference(self) -> tuple[np.ndarray, float]: # -- compilation ---------------------------------------------------------- - def _threads_per_shot(self) -> int: - """Return log2 of the threads per shot: one per ``2^AMPLITUDE_AXES`` amplitudes.""" - width = max((c.width for c in self.layout.components), default=0) - return min(5, max(0, width - AMPLITUDE_AXES)) - def _module(self, mask: np.ndarray | None) -> _Module: key = b"" if mask is None else np.asarray(mask, bool).tobytes() if key in self._modules: return self._modules[key] plan = self.plan(mask) layout = self.layout - tb = self._threads_per_shot() + tb = tape.threads_per_shot(layout) T = 1 << tb - observables: list[dict[int, int]] = [] - for component in layout.components: - slots = { - o: int(np.searchsorted(self._observables, index)) - for o, index in enumerate(component.output_indices) - if index >= self._num_detectors - } - observables.append(slots) - functions = "\n\n".join( - tape_function(program, i, tb, observables[i]) - for i, program in enumerate(plan.programs) - ) - calls = "\n".join( - f" grave = min(grave, component_{i}(w, shot, seed_{i}, t, ev, deviation));" - for i in range(len(plan.programs)) - ) - bits = np.zeros(64 * layout.output_words, dtype=bool) - if mask is not None: - bits[: len(mask)] = mask - psm_words = pack_bits(bits, layout.output_words) - written = sorted({int(w) for p in plan.programs for w in p.out_words}) - stores = "\n".join(f" row[{w}] = w[{w}];" for w in written) - hits = " | ".join( - f"(w[{w}] & {hex64(psm_words[w])})" for w in written if psm_words[w] - ) + parts = tape.tape_kernel(layout, plan, mask, self._num_detectors, tb, METAL) # NOTE: the rows keep every output, the postselected ones too, as the CPU's do. nb = -(-layout.num_outputs // 8) source = ( PRELUDE + noise_source(self._noise, layout.output_words, layout.num_outputs) - + functions + + parts.functions + _KERNELS.format( T1=T - 1, tb=tb, @@ -428,13 +378,10 @@ def _module(self, mask: np.ndarray | None) -> _Module: RB=8 * layout.num_words, NOUT=max(layout.num_outputs, 1), LAST=hex((1 << (layout.num_outputs % 8 or 8)) - 1), - calls=calls, - seeds="\n".join( - f" const u64 seed_{i} = mix64(seed + {i + 1}UL * GOLDEN);" - for i in range(len(plan.programs)) - ), - stores=stores, - hits=hits or "0", + calls=parts.calls, + seeds=parts.seeds, + stores=parts.stores, + hits=parts.hits, ) ) library = runtime.Library(source) @@ -450,8 +397,8 @@ def _module(self, mask: np.ndarray | None) -> _Module: has_mask=mask is not None, screen=runtime.upload(plan.screen), plan=plan, - psm_words=psm_words, - passing=self.passing(mask), + psm_words=parts.psm_words, + passing=screen_pass_rate(self._channels, layout, plan), ) self._modules[key] = compiled return compiled @@ -470,44 +417,27 @@ def _split( ) -> tuple[_Launch, _Launch] | None: """Return a postselected batch's noise passes (screened, deferred), or None. - The deferred channels flip none of the bits that the screen, the tapes or - the postselection read: the postselected outputs, the keys, and with - ``expectations`` the components' observables. They are drawn apart when - they are at least ``DEFER_SHARE`` of the expected faults. + As ``gpu.noise.split``; with ``expectations`` the observables' bits are + read too. """ # NOTE: rows too wide for threadgroup memory draw all their noise at once. if not module.has_mask or not self._noise.shared_words: return None - if expectations in module.splits: - return module.splits[expectations] - layout = self.layout - read = np.zeros(layout.num_words, dtype=np.uint64) - read[: layout.output_words] = module.psm_words | module.plan.screen - for program in module.plan.programs: - start = program.key_start - read[start : start + program.key_words] = ~np.uint64(0) - np.bitwise_or.at(read, program.ps_word, program.ps_mask.astype(np.uint64)) - if expectations: - bits = np.zeros(64 * layout.output_words, dtype=bool) - for component in layout.components: - for index in component.output_indices: - if index >= self._num_detectors: - bits[index] = True - read[: layout.output_words] |= pack_bits(bits, layout.output_words) - deferred = deferred_channels(self._channels, layout.noise, read) - later = self._noise.subset(deferred) - split = None - if later.faults > 0 and later.faults >= DEFER_SHARE * self._noise.groups.faults: - split = (self._launch(self._noise.subset(~deferred)), self._launch(later)) - module.splits[expectations] = split - return split - - def passing(self, mask: np.ndarray | None) -> float: - """Return the estimated share of shots that pass the screen of ``mask``.""" - layout = self.layout - screen = np.zeros(layout.num_words, dtype=np.uint64) - screen[: layout.output_words] = self.plan(mask).screen - return pass_rate(self._channels, layout.noise, screen) if screen.any() else 1.0 + if expectations not in module.splits: + found = split( + self._noise, + self._channels, + self.layout, + module.plan, + module.psm_words, + self._num_detectors if expectations else None, + ) + module.splits[expectations] = ( + None + if found is None + else (self._launch(found[0]), self._launch(found[1])) + ) + return module.splits[expectations] # -- execution -------------------------------------------------------------- @@ -714,12 +644,6 @@ def _encode_pack( ) encoder.dispatch(library.pipeline("advance"), [cursor, count], 1, 1) - def _deviation(self, counts: np.ndarray) -> float: - """Return a batch's normalization deviation beyond float32 rounding, else 0.""" - deviation = float(counts[2:3].view(np.float32)[0]) - # NOTE: fast math records a NaN total as NaN (not inf), which must warn. - return 0.0 if deviation <= DEVIATION_TOLERANCE else deviation - def sample( self, shots: int, @@ -964,7 +888,8 @@ def _sample_call( self._rng.bit_generator.state = state capacity = shots self._shares[key] = total / max(shots, 1) - deviation = self._deviation(slot.counts.array) + # NOTE: fast math records a NaN total as NaN (not inf), which must warn. + deviation = tape.deviation(slot.counts.array[2]) packed = rows.array[: total * nb].reshape(total, nb) expected = None if values is not None: @@ -1024,7 +949,7 @@ def _forked() -> bool: """Return whether this process was forked after Metal's setup, with a warning.""" if not runtime.forked_after_init(): return False - warnings.warn(runtime.FORKED, stacklevel=5) + warnings.warn(FORKED, stacklevel=5) return True diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index 66a61f97..1d0717fc 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -763,8 +763,9 @@ def run_lanes( if kind in (ROTATE_DIAG, ROTATE): flip = np.empty(size, dtype=np.int64) for r in range(size): - flip[r] = _row_parity(symbols, r, program.sign_row[a]) ^ ( - program.sign_const[a] + flip[r] = ( + _row_parity(symbols, r, program.sign_row[a]) + ^ (program.sign_const[a]) ) diagonal = kind == ROTATE_DIAG cos, factor, coupling = ( diff --git a/test/integration/test_metal_sampler.py b/test/integration/test_metal_sampler.py index 8ec3d167..db1d831f 100644 --- a/test/integration/test_metal_sampler.py +++ b/test/integration/test_metal_sampler.py @@ -11,7 +11,7 @@ import pytest from tsim.circuit import Circuit -from tsim.statevector.gpu import tape +from tsim.statevector.gpu import noise as gpu_noise, tape from tsim.statevector.metal import ( codegen, noise as metal_noise, @@ -239,7 +239,7 @@ def test_metal_samples_no_shots(): def test_metal_postselection_drops_the_rows_that_fail(monkeypatch): # NOTE: the masked run must draw the unmasked run's noise, so none is deferred. - monkeypatch.setattr(metal_sampler, "DEFER_SHARE", 2.0) + monkeypatch.setattr(gpu_noise, "DEFER_SHARE", 2.0) circuit = Circuit(NOISY_DETECTORS) rng = np.random.default_rng(0) mask = np.zeros(150, dtype=bool) @@ -324,7 +324,7 @@ def test_deferred_noise_leaves_the_expectations_unbiased(): def test_metal_batches_return_the_rows_of_one_pass(monkeypatch, text): # NOTE: 1000-shot batches, queued on one set of buffers; without deferral, a # masked and an unmasked run on the same batches draw the same shots. - monkeypatch.setattr(metal_sampler, "DEFER_SHARE", 2.0) + monkeypatch.setattr(gpu_noise, "DEFER_SHARE", 2.0) circuit = Circuit(text) mask = np.zeros(circuit.num_detectors, dtype=bool) mask[[0, 2]] = True @@ -536,7 +536,7 @@ def test_metal_looping_tapes_match_stabilizer_rank(monkeypatch, name): def test_components_wider_than_the_limit_run_on_the_cpu(): circuit = _random_unitary_circuit(13, 12) - assert _widest(circuit) > metal_sampler.MAX_WIDTH + assert _widest(circuit) > tape.MAX_WIDTH with pytest.warns(UserWarning, match="sampling on the CPU"): sampler = circuit.compile_sampler(seed=1) assert type(sampler._shots) is ShotSampler @@ -595,13 +595,6 @@ def test_metal_unmasked_expectations_after_a_masked_call(): np.testing.assert_array_equal(values_again, values) -def test_metal_reports_a_nan_normalization(): - sampler = _gpu(Circuit(T_STATE_X).compile_detector_sampler(seed=1)) - counts = np.zeros(4, dtype=np.uint32) - counts[2] = 0x7FC00000 - assert np.isnan(_gpu_shots(sampler)._deviation(counts)) - - def test_metal_samples_rows_too_wide_for_threadgroup_memory(): # NOTE: 3200 detectors are 51 words: the noise goes into device memory. circuit = Circuit( From d61bad2c6a6f63803b81b6835d6eb7c3e42a8da4 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 04:02:35 -0700 Subject: [PATCH 17/38] refactor(metal): name the counts words and the dispatch sizes The kernels' counts buffer was addressed by bare byte offsets and indices, and 256, 1024, 1 << 20 and stacklevel=5 were spelled out at each use. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/metal/sampler.py | 81 +++++++++++++++++---------- 1 file changed, 52 insertions(+), 29 deletions(-) diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index 22e49ebf..701a0375 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -53,7 +53,22 @@ """Bytes of shot words of one batch, which bound its shots below ``BATCH_SHOTS``.""" NOISE_THREADS = 256 -"""Threads per threadgroup of the noise kernel.""" +"""Threads per threadgroup of the noise kernels.""" + +COPY_THREADS = 256 +"""Threads per threadgroup of the kernels that copy rows.""" + +MAX_COPY_THREADS = 1 << 20 +"""Most threads of a copy kernel's dispatch; each thread loops over its share.""" + +SCAN_THREADS = 1024 +"""Threads per threadgroup of ``count_bits`` and ``list_bits``: one per bitmap word.""" + +# NOTE: the words of ``_Slot.counts``. +_SURVIVORS, _ACCEPTED, _DEVIATION, _CURSOR, _CAPACITY = range(5) + +_STACKLEVEL = 5 +"""Stack level of the samplers' warnings: the caller of ``sample``.""" TAPE_THREADS = 128 """Threads per threadgroup of the tape kernel.""" @@ -124,7 +139,7 @@ }} // The shots of the set bits of ``bits`` (``words`` words), in order: ``count_bits`` -// counts each block's 1024 words, ``list_bits`` lists a block's shots after +// counts each block's {SCAN} words, ``list_bits`` lists a block's shots after // those of the blocks before it. kernel void count_bits( device const uint* bits [[buffer(0)]], constant long& words [[buffer(1)]], @@ -134,7 +149,7 @@ uint nwarps [[simdgroups_per_threadgroup]]) {{ threadgroup uint warps[32]; - long i = (long)bid * 1024 + tid; + long i = (long)bid * {SCAN} + tid; uint total; block_scan(i < words ? popcount(bits[i]) : 0, total, warps, lane, warp, nwarps); if (tid == 0) totals[bid] = total; @@ -149,9 +164,9 @@ {{ threadgroup uint warps[32]; uint earlier = 0, total; - for (uint b = tid; b < bid; b += 1024) earlier += totals[b]; + for (uint b = tid; b < bid; b += {SCAN}) earlier += totals[b]; block_scan(earlier, earlier, warps, lane, warp, nwarps); - long i = (long)bid * 1024 + tid; + long i = (long)bid * {SCAN} + tid; uint w = i < words ? bits[i] : 0; uint at = earlier + block_scan(popcount(w), total, warps, lane, warp, nwarps); for (; w; w &= w - 1) index[at++] = (uint)(32 * i + ctz(w)); @@ -373,6 +388,7 @@ def _module(self, mask: np.ndarray | None) -> _Module: T1=T - 1, tb=tb, NW=layout.num_words, + SCAN=SCAN_THREADS, NO=max(1, len(self._observables)), NB=max(nb, 1), RB=8 * layout.num_words, @@ -466,7 +482,7 @@ def _slot(self, shots: int, expectations: bool) -> _Slot: queue=runtime.buffer(shots, np.uint32), accepted=runtime.buffer(-(-shots // 32), np.uint32), counts=runtime.buffer(8, np.uint32), - totals=runtime.buffer(-(-shots // 32 // 1024) + 1, np.uint32), + totals=runtime.buffer(-(-shots // 32 // SCAN_THREADS) + 1, np.uint32), index=runtime.buffer(shots, np.uint32), ) if expectations: @@ -515,7 +531,7 @@ def _encode_batch( constant, module.screen, slot.queue, - (slot.counts, 0), + _count(slot, _SURVIVORS), slot.accepted, ], -(-shots // S), @@ -534,11 +550,11 @@ def _encode_batch( words, tape_args, slot.queue, - (slot.counts, 0), + _count(slot, _SURVIVORS), slot.accepted, - (slot.counts, 4), + _count(slot, _ACCEPTED), values, - (slot.counts, 8), + _count(slot, _DEVIATION), ], -(-threads // TAPE_THREADS), TAPE_THREADS, @@ -568,23 +584,24 @@ def _encode_pack( no = len(self._observables) if module.has_mask: words = -(-shots // 32) - blocks = -(-words // 1024) + blocks = -(-words // SCAN_THREADS) encoder.dispatch( library.pipeline("count_bits"), [slot.accepted, np.int64(words), slot.totals], blocks, - 1024, + SCAN_THREADS, ) encoder.dispatch( library.pipeline("list_bits"), [slot.accepted, np.int64(words), slot.totals, slot.index], blocks, - 1024, + SCAN_THREADS, ) - count = (slot.counts, 4) if self.layout.components else (slot.counts, 0) + kept = _ACCEPTED if self.layout.components else _SURVIVORS + count = _count(slot, kept) else: count = fixed - cursor = (slot.counts, 12) + cursor = _count(slot, _CURSOR) blocks = -(-shots // S) if nb and split is not None: launch = split[1] @@ -608,11 +625,11 @@ def _encode_pack( cursor, ], blocks, - 256, + NOISE_THREADS, threadgroup_bytes=max(8 * launch.max_events, S * nb), ) elif nb: - threads = min(-(-shots * nb // 4) + 1, 1 << 20) + threads = min(-(-shots * nb // 4) + 1, MAX_COPY_THREADS) encoder.dispatch( library.pipeline("gather_rows"), [ @@ -623,8 +640,8 @@ def _encode_pack( rows, cursor, ], - -(-threads // 256), - 256, + -(-threads // COPY_THREADS), + COPY_THREADS, ) if expectations and no and self.layout.components: assert slot.values is not None and values is not None @@ -639,8 +656,8 @@ def _encode_pack( cursor, np.int32(module.has_mask), ], - 256, - 256, + COPY_THREADS, + COPY_THREADS, ) encoder.dispatch(library.pipeline("advance"), [cursor, count], 1, 1) @@ -718,7 +735,7 @@ def _run_accepted( starts = range(0, shots, batch) fixed = runtime.buffer(len(starts), np.uint32) slot.counts.array[:] = 0 - slot.counts.array[4] = capacity + slot.counts.array[_CAPACITY] = capacity commands = [] try: for k, start in enumerate(starts): @@ -882,14 +899,14 @@ def _sample_call( values, capacity, ) - total = int(slot.counts.array[3]) + total = int(slot.counts.array[_CURSOR]) if total <= capacity: break self._rng.bit_generator.state = state capacity = shots self._shares[key] = total / max(shots, 1) # NOTE: fast math records a NaN total as NaN (not inf), which must warn. - deviation = tape.deviation(slot.counts.array[2]) + deviation = tape.deviation(slot.counts.array[_DEVIATION]) packed = rows.array[: total * nb].reshape(total, nb) expected = None if values is not None: @@ -911,13 +928,13 @@ def _unpack(self, module: _Module, rows: runtime.Buffer, total: int) -> np.ndarr return _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, 1 << 20) + threads = min(total * groups, MAX_COPY_THREADS) encoder = runtime.Encoder() encoder.dispatch( module.library.pipeline("unpack_rows"), [rows, out, np.uint32(total)], - -(-threads // 256), - 256, + -(-threads // COPY_THREADS), + COPY_THREADS, ) self._wait_all([encoder.commit()]) return out.array[: total * num_outputs].reshape(total, num_outputs).view(bool) @@ -949,16 +966,22 @@ def _forked() -> bool: """Return whether this process was forked after Metal's setup, with a warning.""" if not runtime.forked_after_init(): return False - warnings.warn(FORKED, stacklevel=5) + warnings.warn(FORKED, stacklevel=_STACKLEVEL) return True def _warn_out_of_memory(error: MemoryError) -> None: warnings.warn( - f"{error}; sampling on the CPU. Fewer shots per call fit.", stacklevel=5 + f"{error}; sampling on the CPU. Fewer shots per call fit.", + stacklevel=_STACKLEVEL, ) +def _count(slot: _Slot, word: int) -> tuple[runtime.Buffer, int]: + """Return word ``word`` of ``slot.counts``, as a kernel argument.""" + return slot.counts, 4 * word + + def _restore( host, channels, num_detectors, seed, state ) -> MetalShotSampler | ShotSampler: From ec61a399cdb1227c6e07bd10019881333d6c209b Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 04:03:32 -0700 Subject: [PATCH 18/38] refactor(metal): share the noise kernels' signature and queue code Both noise kernels and noise_kept repeated the segment buffers, and both noise kernels the same ballot, bitmap and queue block; those are now one signature string and an enqueue function. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/metal/noise.py | 85 ++++++++++++++--------------- 1 file changed, 41 insertions(+), 44 deletions(-) diff --git a/src/tsim/statevector/metal/noise.py b/src/tsim/statevector/metal/noise.py index fac5432c..a7bd39dc 100644 --- a/src/tsim/statevector/metal/noise.py +++ b/src/tsim/statevector/metal/noise.py @@ -33,20 +33,32 @@ gain 10-25%, distillation loses up to 15%). """ -NOISE_SHARED = r""" -kernel void noise( - device u64* out [[buffer(0)]], constant NoiseArgs& args [[buffer(1)]], +_SEGMENT_PARAMS = """\ device const int* seg_first [[buffer(2)]], device const int* seg_outcomes [[buffer(3)]], device const int* seg_begin [[buffer(4)]], device const int* seg_end [[buffer(5)]], device const float* seg_log_keep [[buffer(6)]], device const Record* records [[buffer(7)]], device const Flip* overflow [[buffer(8)]], +""" + +# NOTE: the signature of both ``noise`` kernels: with the block's words in +# threadgroup memory or in device memory. +_NOISE_SIGNATURE = ( + """\ +kernel void noise( + device u64* out [[buffer(0)]], constant NoiseArgs& args [[buffer(1)]], +""" + + _SEGMENT_PARAMS + + """\ device const u64* constants [[buffer(9)]], device const u64* screen [[buffer(10)]], device uint* queue [[buffer(11)]], device atomic_uint* count [[buffer(12)]], device uint* survivors [[buffer(13)]], threadgroup uint2* events [[threadgroup(0)]], uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], uint nt [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]]) -{{ +""" +) + +NOISE_SHARED = _NOISE_SIGNATURE + r"""{{ threadgroup u64 sw[{NW} * {S}]; threadgroup atomic_int num_events; threadgroup bool passed[{S}]; @@ -71,16 +83,7 @@ alive = !hit; }} if (s < {S}) passed[s] = alive; - uint ballot = uint(static_cast(simd_ballot(alive))); - int warp_first = s0 + (tid & ~31u); - if (args.has_survivors && lane == 0 && warp_first < ns) - survivors[(base + warp_first) >> 5] = ballot; - if (ballot) {{ - uint leader = ctz(ballot), b = 0; - if (lane == leader) b = atomic_fetch_add_explicit(count, popcount(ballot), memory_order_relaxed); - b = simd_broadcast(b, (ushort)leader); - if (alive) queue[b + popcount(ballot & ((1u << lane) - 1))] = (uint)(base + s); - }} + enqueue(alive, s, s0 + (tid & ~31u), base, ns, lane, args.has_survivors, queue, count, survivors); }} threadgroup_barrier(mem_flags::mem_threadgroup); device u64* rows = out + base * {NW}; @@ -92,20 +95,7 @@ """ -NOISE_DEVICE = r""" -kernel void noise( - device u64* out [[buffer(0)]], constant NoiseArgs& args [[buffer(1)]], - device const int* seg_first [[buffer(2)]], device const int* seg_outcomes [[buffer(3)]], - device const int* seg_begin [[buffer(4)]], device const int* seg_end [[buffer(5)]], - device const float* seg_log_keep [[buffer(6)]], device const Record* records [[buffer(7)]], - device const Flip* overflow [[buffer(8)]], - device const u64* constants [[buffer(9)]], device const u64* screen [[buffer(10)]], - device uint* queue [[buffer(11)]], device atomic_uint* count [[buffer(12)]], - device uint* survivors [[buffer(13)]], - threadgroup uint2* events [[threadgroup(0)]], - uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], - uint nt [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]]) -{{ +NOISE_DEVICE = _NOISE_SIGNATURE + r"""{{ threadgroup atomic_int num_events; long base = (long)bid * {S}; int ns = (int)min((long){S}, args.shots - base); @@ -130,33 +120,23 @@ }} alive = !hit; }} - uint ballot = uint(static_cast(simd_ballot(alive))); - int warp_first = s0 + (tid & ~31u); - if (args.has_survivors && lane == 0 && warp_first < ns) - survivors[(base + warp_first) >> 5] = ballot; - if (ballot) {{ - uint leader = ctz(ballot), b = 0; - if (lane == leader) b = atomic_fetch_add_explicit(count, popcount(ballot), memory_order_relaxed); - b = simd_broadcast(b, (ushort)leader); - if (alive) queue[b + popcount(ballot & ((1u << lane) - 1))] = (uint)(base + s); - }} + enqueue(alive, s, s0 + (tid & ~31u), base, ns, lane, args.has_survivors, queue, count, survivors); }} }} """ -NOISE_KEPT = r""" +NOISE_KEPT = ( + r""" struct KeptArgs {{ u64 seed; int num_segments; int max_events; }}; // The deferred channels' faults on the kept shots ``index`` (``*n`` of them), // then their returned outputs bit-packed (``pack_rows``). kernel void noise_kept( device const u64* out [[buffer(0)]], constant KeptArgs& args [[buffer(1)]], - device const int* seg_first [[buffer(2)]], device const int* seg_outcomes [[buffer(3)]], - device const int* seg_begin [[buffer(4)]], device const int* seg_end [[buffer(5)]], - device const float* seg_log_keep [[buffer(6)]], device const Record* records [[buffer(7)]], - device const Flip* overflow [[buffer(8)]], - device const uint* n [[buffer(9)]], device const uint* index [[buffer(10)]], +""" + + _SEGMENT_PARAMS + + r""" device const uint* n [[buffer(9)]], device const uint* index [[buffer(10)]], device uchar* dst [[buffer(11)]], device const uint* cursor [[buffer(12)]], threadgroup uint2* events [[threadgroup(0)]], uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], @@ -182,6 +162,7 @@ pack_rows(sw, ns, (threadgroup uchar*)events, dst, ((long)*cursor + base) * {NB}, tid, nt); }} """ +) SOURCE = r""" struct Record {{ uint alias; uint words; uint extra; uint pad; u64 m0; u64 m1; }}; @@ -339,6 +320,22 @@ struct NoiseArgs {{ long shots; u64 seed; int num_segments; int max_events; int has_survivors; }}; +// Marks a SIMD group's surviving shots (of the block's ``ns``, from ``first``) +// in the bitmap ``survivors`` if ``has_survivors``, and appends them to +// ``queue`` after the ``count`` queued so far. +inline void enqueue(bool alive, int s, int first, long base, int ns, uint lane, + int has_survivors, device uint* queue, device atomic_uint* count, device uint* survivors) +{{ + uint ballot = uint(static_cast(simd_ballot(alive))); + if (has_survivors && lane == 0 && first < ns) survivors[(base + first) >> 5] = ballot; + if (ballot) {{ + uint leader = ctz(ballot), b = 0; + if (lane == leader) b = atomic_fetch_add_explicit(count, popcount(ballot), memory_order_relaxed); + b = simd_broadcast(b, (ushort)leader); + if (alive) queue[b + popcount(ballot & ((1u << lane) - 1))] = (uint)(base + s); + }} +}} + {NOISE}{KEPT}""" From b0b32e709b1bfe3222c7c3f82d473b0b2bd5b479 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 04:04:41 -0700 Subject: [PATCH 19/38] refactor(metal): move the kernels' MSL to metal/kernels.py The 170-line MSL template and the source assembly leave sampler.py, so _module only plans, compiles and uploads. sample and direct_bools share one fallback to the CPU (_or_cpu) for forked processes and GPU out-of-memory. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/metal/kernels.py | 231 ++++++++++++++++++++++ src/tsim/statevector/metal/sampler.py | 269 ++++---------------------- 2 files changed, 269 insertions(+), 231 deletions(-) create mode 100644 src/tsim/statevector/metal/kernels.py diff --git a/src/tsim/statevector/metal/kernels.py b/src/tsim/statevector/metal/kernels.py new file mode 100644 index 00000000..147ca5bb --- /dev/null +++ b/src/tsim/statevector/metal/kernels.py @@ -0,0 +1,231 @@ +"""The Metal sampler's kernels other than the noise ones, and their source. + +``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; +``reset`` and ``advance`` keep a call's counters on the GPU. +""" + +from __future__ import annotations + +from tsim.statevector.gpu.noise import NoiseTables +from tsim.statevector.gpu.tape import TapeKernel +from tsim.statevector.metal.codegen import PRELUDE +from tsim.statevector.metal.noise import source as noise_source +from tsim.statevector.shots import ShotLayout + +SCAN_THREADS = 1024 +"""Threads per threadgroup of ``count_bits`` and ``list_bits``: one per bitmap word.""" + +TEMPLATE = r""" +struct TapeArgs {{ long first_shot; u64 seed; int flags; }}; + +// flags: 1 clears rejected shots from ``survivors`` and counts the accepted ones, +// 2 writes the expectations. +kernel void tape( + device u64* out [[buffer(0)]], constant TapeArgs& args [[buffer(1)]], + device const uint* queue [[buffer(2)]], device const uint* count [[buffer(3)]], + device atomic_uint* survivors [[buffer(4)]], device atomic_uint* acc_count [[buffer(5)]], + device float* expectations [[buffer(6)]], device atomic_uint* deviation [[buffer(7)]], + uint gid [[thread_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) +{{ + int t = lane & {T1}; + uint q = gid >> {tb}; + if (q >= count[0]) return; + const u64 seed = args.seed; +{seeds} + uint row_index = queue[q]; + u64 shot = (u64)args.first_shot + row_index; + device u64* row = out + (long)row_index * {NW}; + u64 w[{NW}]; + for (int i = 0; i < {NW}; ++i) w[i] = row[i]; + float ev[{NO}]; + int grave = 0x7fffffff; +{calls} + if (t == 0) {{ +{stores} + if (args.flags & 2) + for (int i = 0; i < {NO}; ++i) expectations[(long)row_index * {NO} + i] = ev[i]; + }} + if (args.flags & 1) {{ + u64 hit = {hits}; + if (t == 0 && hit) + atomic_fetch_and_explicit(&survivors[row_index >> 5], ~(1u << (row_index & 31)), memory_order_relaxed); + uint ballot = uint(static_cast(simd_ballot(t == 0 && hit == 0))); + if (ballot && simd_is_first()) atomic_fetch_add_explicit(acc_count, popcount(ballot), memory_order_relaxed); + }} +}} + +// The counts of a batch: its screened survivors and its accepted shots. +kernel void reset(device uint* counts [[buffer(0)]]) {{ counts[0] = 0; counts[1] = 0; }} + +// Moves the cursor past a batch's ``*n`` rows. +kernel void advance(device uint* cursor [[buffer(0)]], device const uint* n [[buffer(1)]]) {{ + cursor[0] += *n; +}} + +inline uint block_scan(uint v, thread uint& total, threadgroup uint* warps, uint lane, uint warp, uint nwarps) {{ + uint x = simd_prefix_inclusive_sum(v); + if (lane == 31) warps[warp] = x; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (warp == 0) {{ + uint y = lane < nwarps ? warps[lane] : 0; + warps[lane] = simd_prefix_inclusive_sum(y); + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + uint before = (warp ? warps[warp - 1] : 0) + x - v; + total = warps[nwarps - 1]; + threadgroup_barrier(mem_flags::mem_threadgroup); + return before; +}} + +// The shots of the set bits of ``bits`` (``words`` words), in order: ``count_bits`` +// counts each block's {SCAN} words, ``list_bits`` lists a block's shots after +// those of the blocks before it. +kernel void count_bits( + device const uint* bits [[buffer(0)]], constant long& words [[buffer(1)]], + device uint* totals [[buffer(2)]], + uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], uint warp [[simdgroup_index_in_threadgroup]], + uint nwarps [[simdgroups_per_threadgroup]]) +{{ + threadgroup uint warps[32]; + long i = (long)bid * {SCAN} + tid; + uint total; + block_scan(i < words ? popcount(bits[i]) : 0, total, warps, lane, warp, nwarps); + if (tid == 0) totals[bid] = total; +}} + +kernel void list_bits( + device const uint* bits [[buffer(0)]], constant long& words [[buffer(1)]], + device const uint* totals [[buffer(2)]], device uint* index [[buffer(3)]], + uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], uint warp [[simdgroup_index_in_threadgroup]], + uint nwarps [[simdgroups_per_threadgroup]]) +{{ + threadgroup uint warps[32]; + uint earlier = 0, total; + for (uint b = tid; b < bid; b += {SCAN}) earlier += totals[b]; + block_scan(earlier, earlier, warps, lane, warp, nwarps); + long i = (long)bid * {SCAN} + tid; + uint w = i < words ? bits[i] : 0; + uint at = earlier + block_scan(popcount(w), total, warps, lane, warp, nwarps); + for (; w; w &= w - 1) index[at++] = (uint)(32 * i + ctz(w)); +}} + +// The rows ``index[j]`` (or every row, without ``has_index``) for ``j < *n``, +// densely after the ``*cursor`` rows of the call's earlier batches: each row's +// first {NB} bytes, which hold its outputs bit-packed (the bits past them +// cleared). A thread writes 4 aligned bytes at a time. +kernel void gather_rows( + device const uchar* rows [[buffer(0)]], device const uint* n [[buffer(1)]], + device const uint* index [[buffer(2)]], constant int& has_index [[buffer(3)]], + device uchar* dst [[buffer(4)]], device const uint* cursor [[buffer(5)]], + uint gid [[thread_position_in_grid]], uint grid [[threads_per_grid]]) +{{ + long count = min((long)*n, (long)cursor[1] - (long)cursor[0]); + if (count <= 0) return; + long at = (long)cursor[0] * {NB}, end = at + count * {NB}; + for (long p = (at & ~3L) + 4L * gid; p < end; p += 4L * grid) {{ + // NOTE: a batch's rows span less than 4 GB (``BATCH_BYTES``). + uint k = (uint)(max(p, at) - at); + uint r = k / {NB}, b = k - r * {NB}; + uint sh = 8 * (b & 3); + if (p >= at && p + 4 <= end && b + 4 <= {NB} && (sh == 0 || (b & ~3u) + 8 <= {RB})) {{ + // NOTE: the 4 bytes lie in one row: two aligned loads at most. + long source = has_index ? (long)index[r] : r; + device const uint* w = (device const uint*)(rows + source * {RB}) + (b >> 2); + uint v = sh ? (w[0] >> sh) | (w[1] << (32 - sh)) : w[0]; + if (b + 4 == {NB}) v &= ({LAST}u << 24) | 0xFFFFFFu; + *(device uint*)(dst + p) = v; + continue; + }} + uint v = 0; + for (int j = (int)(max(p, at) - p); j < 4 && p + j < end; ++j) {{ + long source = has_index ? (long)index[r] : r; + uint x = rows[source * {RB} + b]; + if (b == {NB} - 1) x &= {LAST}; + v |= x << (8 * j); + if (++b == {NB}) {{ + b = 0; + ++r; + }} + }} + for (int j = 0; j < 4; ++j) + if (p + j >= at && p + j < end) dst[p + j] = uchar(v >> (8 * j)); + }} +}} + +// The first ``n`` packed rows of ``{NB}`` bytes, one byte per output (``{NOUT}`` +// per row), as booleans: a thread unpacks 4 outputs, which share a packed byte. +kernel void unpack_rows( + device const uchar* packed [[buffer(0)]], device uchar* dst [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint grid [[threads_per_grid]]) +{{ + const uint groups = ({NOUT} + 3) / 4; + for (uint i = gid; i < n * groups; i += grid) {{ + uint r = i / groups, g = i - r * groups; + uint b = packed[(ulong)r * {NB} + (g >> 1)] >> (4 * (g & 1)); + device uchar* out = dst + (ulong)r * {NOUT} + 4 * g; + for (uint j = 0; j < 4 && 4 * g + j < {NOUT}; ++j) out[j] = uchar((b >> j) & 1); + }} +}} + +// 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. +kernel void gather_float( + device const float* values [[buffer(0)]], device const uint* n [[buffer(1)]], + constant int& width [[buffer(2)]], device const uint* index [[buffer(3)]], + device float* dst [[buffer(4)]], device const uint* cursor [[buffer(5)]], + constant int& has_index [[buffer(6)]], + uint gid [[thread_position_in_grid]], uint grid [[threads_per_grid]]) +{{ + long total = min((long)*n, (long)cursor[1] - (long)cursor[0]) * width; + dst += (long)*cursor * width; + for (long k = gid; k < total; k += grid) {{ + long j = k / width; int c = (int)(k - j * width); + long r = has_index ? (long)index[j] : j; + dst[k] = values[r * width + c]; + }} +}} +""" + + +def source( + layout: ShotLayout, + tables: NoiseTables, + parts: TapeKernel, + tb: int, + num_observables: int, +) -> str: + """Return the MSL source of every kernel, for tapes of ``2^tb`` threads a shot.""" + nb = -(-layout.num_outputs // 8) + return ( + PRELUDE + + noise_source(tables, layout.output_words, layout.num_outputs) + + parts.functions + + TEMPLATE.format( + T1=(1 << tb) - 1, + tb=tb, + NW=layout.num_words, + SCAN=SCAN_THREADS, + NO=num_observables, + NB=max(nb, 1), + RB=8 * layout.num_words, + NOUT=max(layout.num_outputs, 1), + LAST=hex((1 << (layout.num_outputs % 8 or 8)) - 1), + calls=parts.calls, + seeds=parts.seeds, + stores=parts.stores, + hits=parts.hits, + ) + ) + + +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"] + 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 701a0375..9c8bcabb 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -41,9 +41,10 @@ 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.metal import runtime -from tsim.statevector.metal.codegen import METAL, PRELUDE -from tsim.statevector.metal.noise import noise_tables, source as noise_source +from tsim.statevector.metal import kernels, runtime +from tsim.statevector.metal.codegen import METAL +from tsim.statevector.metal.kernels import SCAN_THREADS +from tsim.statevector.metal.noise import noise_tables from tsim.statevector.shots import ShotPlan, ShotSampler BATCH_SHOTS = 1 << 20 @@ -61,13 +62,10 @@ MAX_COPY_THREADS = 1 << 20 """Most threads of a copy kernel's dispatch; each thread loops over its share.""" -SCAN_THREADS = 1024 -"""Threads per threadgroup of ``count_bits`` and ``list_bits``: one per bitmap word.""" - # NOTE: the words of ``_Slot.counts``. _SURVIVORS, _ACCEPTED, _DEVIATION, _CURSOR, _CAPACITY = range(5) -_STACKLEVEL = 5 +_STACKLEVEL = 6 """Stack level of the samplers' warnings: the caller of ``sample``.""" TAPE_THREADS = 128 @@ -76,181 +74,6 @@ MAX_CALL_SHOTS = 1 << 31 """Most shots of one GPU call; larger calls run in parts (the counters are 32-bit).""" -_KERNELS = r""" -struct TapeArgs {{ long first_shot; u64 seed; int flags; }}; - -// flags: 1 clears rejected shots from ``survivors`` and counts the accepted ones, -// 2 writes the expectations. -kernel void tape( - device u64* out [[buffer(0)]], constant TapeArgs& args [[buffer(1)]], - device const uint* queue [[buffer(2)]], device const uint* count [[buffer(3)]], - device atomic_uint* survivors [[buffer(4)]], device atomic_uint* acc_count [[buffer(5)]], - device float* expectations [[buffer(6)]], device atomic_uint* deviation [[buffer(7)]], - uint gid [[thread_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) -{{ - int t = lane & {T1}; - uint q = gid >> {tb}; - if (q >= count[0]) return; - const u64 seed = args.seed; -{seeds} - uint row_index = queue[q]; - u64 shot = (u64)args.first_shot + row_index; - device u64* row = out + (long)row_index * {NW}; - u64 w[{NW}]; - for (int i = 0; i < {NW}; ++i) w[i] = row[i]; - float ev[{NO}]; - int grave = 0x7fffffff; -{calls} - if (t == 0) {{ -{stores} - if (args.flags & 2) - for (int i = 0; i < {NO}; ++i) expectations[(long)row_index * {NO} + i] = ev[i]; - }} - if (args.flags & 1) {{ - u64 hit = {hits}; - if (t == 0 && hit) - atomic_fetch_and_explicit(&survivors[row_index >> 5], ~(1u << (row_index & 31)), memory_order_relaxed); - uint ballot = uint(static_cast(simd_ballot(t == 0 && hit == 0))); - if (ballot && simd_is_first()) atomic_fetch_add_explicit(acc_count, popcount(ballot), memory_order_relaxed); - }} -}} - -// The counts of a batch: its screened survivors and its accepted shots. -kernel void reset(device uint* counts [[buffer(0)]]) {{ counts[0] = 0; counts[1] = 0; }} - -// Moves the cursor past a batch's ``*n`` rows. -kernel void advance(device uint* cursor [[buffer(0)]], device const uint* n [[buffer(1)]]) {{ - cursor[0] += *n; -}} - -inline uint block_scan(uint v, thread uint& total, threadgroup uint* warps, uint lane, uint warp, uint nwarps) {{ - uint x = simd_prefix_inclusive_sum(v); - if (lane == 31) warps[warp] = x; - threadgroup_barrier(mem_flags::mem_threadgroup); - if (warp == 0) {{ - uint y = lane < nwarps ? warps[lane] : 0; - warps[lane] = simd_prefix_inclusive_sum(y); - }} - threadgroup_barrier(mem_flags::mem_threadgroup); - uint before = (warp ? warps[warp - 1] : 0) + x - v; - total = warps[nwarps - 1]; - threadgroup_barrier(mem_flags::mem_threadgroup); - return before; -}} - -// The shots of the set bits of ``bits`` (``words`` words), in order: ``count_bits`` -// counts each block's {SCAN} words, ``list_bits`` lists a block's shots after -// those of the blocks before it. -kernel void count_bits( - device const uint* bits [[buffer(0)]], constant long& words [[buffer(1)]], - device uint* totals [[buffer(2)]], - uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], - uint lane [[thread_index_in_simdgroup]], uint warp [[simdgroup_index_in_threadgroup]], - uint nwarps [[simdgroups_per_threadgroup]]) -{{ - threadgroup uint warps[32]; - long i = (long)bid * {SCAN} + tid; - uint total; - block_scan(i < words ? popcount(bits[i]) : 0, total, warps, lane, warp, nwarps); - if (tid == 0) totals[bid] = total; -}} - -kernel void list_bits( - device const uint* bits [[buffer(0)]], constant long& words [[buffer(1)]], - device const uint* totals [[buffer(2)]], device uint* index [[buffer(3)]], - uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], - uint lane [[thread_index_in_simdgroup]], uint warp [[simdgroup_index_in_threadgroup]], - uint nwarps [[simdgroups_per_threadgroup]]) -{{ - threadgroup uint warps[32]; - uint earlier = 0, total; - for (uint b = tid; b < bid; b += {SCAN}) earlier += totals[b]; - block_scan(earlier, earlier, warps, lane, warp, nwarps); - long i = (long)bid * {SCAN} + tid; - uint w = i < words ? bits[i] : 0; - uint at = earlier + block_scan(popcount(w), total, warps, lane, warp, nwarps); - for (; w; w &= w - 1) index[at++] = (uint)(32 * i + ctz(w)); -}} - -// The rows ``index[j]`` (or every row, without ``has_index``) for ``j < *n``, -// densely after the ``*cursor`` rows of the call's earlier batches: each row's -// first {NB} bytes, which hold its outputs bit-packed (the bits past them -// cleared). A thread writes 4 aligned bytes at a time. -kernel void gather_rows( - device const uchar* rows [[buffer(0)]], device const uint* n [[buffer(1)]], - device const uint* index [[buffer(2)]], constant int& has_index [[buffer(3)]], - device uchar* dst [[buffer(4)]], device const uint* cursor [[buffer(5)]], - uint gid [[thread_position_in_grid]], uint grid [[threads_per_grid]]) -{{ - long count = min((long)*n, (long)cursor[1] - (long)cursor[0]); - if (count <= 0) return; - long at = (long)cursor[0] * {NB}, end = at + count * {NB}; - for (long p = (at & ~3L) + 4L * gid; p < end; p += 4L * grid) {{ - // NOTE: a batch's rows span less than 4 GB (``BATCH_BYTES``). - uint k = (uint)(max(p, at) - at); - uint r = k / {NB}, b = k - r * {NB}; - uint sh = 8 * (b & 3); - if (p >= at && p + 4 <= end && b + 4 <= {NB} && (sh == 0 || (b & ~3u) + 8 <= {RB})) {{ - // NOTE: the 4 bytes lie in one row: two aligned loads at most. - long source = has_index ? (long)index[r] : r; - device const uint* w = (device const uint*)(rows + source * {RB}) + (b >> 2); - uint v = sh ? (w[0] >> sh) | (w[1] << (32 - sh)) : w[0]; - if (b + 4 == {NB}) v &= ({LAST}u << 24) | 0xFFFFFFu; - *(device uint*)(dst + p) = v; - continue; - }} - uint v = 0; - for (int j = (int)(max(p, at) - p); j < 4 && p + j < end; ++j) {{ - long source = has_index ? (long)index[r] : r; - uint x = rows[source * {RB} + b]; - if (b == {NB} - 1) x &= {LAST}; - v |= x << (8 * j); - if (++b == {NB}) {{ - b = 0; - ++r; - }} - }} - for (int j = 0; j < 4; ++j) - if (p + j >= at && p + j < end) dst[p + j] = uchar(v >> (8 * j)); - }} -}} - -// The first ``n`` packed rows of ``{NB}`` bytes, one byte per output (``{NOUT}`` -// per row), as booleans: a thread unpacks 4 outputs, which share a packed byte. -kernel void unpack_rows( - device const uchar* packed [[buffer(0)]], device uchar* dst [[buffer(1)]], - constant uint& n [[buffer(2)]], - uint gid [[thread_position_in_grid]], uint grid [[threads_per_grid]]) -{{ - const uint groups = ({NOUT} + 3) / 4; - for (uint i = gid; i < n * groups; i += grid) {{ - uint r = i / groups, g = i - r * groups; - uint b = packed[(ulong)r * {NB} + (g >> 1)] >> (4 * (g & 1)); - device uchar* out = dst + (ulong)r * {NOUT} + 4 * g; - for (uint j = 0; j < 4 && 4 * g + j < {NOUT}; ++j) out[j] = uchar((b >> j) & 1); - }} -}} - -// 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. -kernel void gather_float( - device const float* values [[buffer(0)]], device const uint* n [[buffer(1)]], - constant int& width [[buffer(2)]], device const uint* index [[buffer(3)]], - device float* dst [[buffer(4)]], device const uint* cursor [[buffer(5)]], - constant int& has_index [[buffer(6)]], - uint gid [[thread_position_in_grid]], uint grid [[threads_per_grid]]) -{{ - long total = min((long)*n, (long)cursor[1] - (long)cursor[0]) * width; - dst += (long)*cursor * width; - for (long k = gid; k < total; k += grid) {{ - long j = k / width; int c = (int)(k - j * width); - long r = has_index ? (long)index[j] : j; - dst[k] = values[r * width + c]; - }} -}} -""" - @dataclass class _Launch: @@ -376,36 +199,12 @@ def _module(self, mask: np.ndarray | None) -> _Module: plan = self.plan(mask) layout = self.layout tb = tape.threads_per_shot(layout) - T = 1 << tb parts = tape.tape_kernel(layout, plan, mask, self._num_detectors, tb, METAL) - # NOTE: the rows keep every output, the postselected ones too, as the CPU's do. - nb = -(-layout.num_outputs // 8) - source = ( - PRELUDE - + noise_source(self._noise, layout.output_words, layout.num_outputs) - + parts.functions - + _KERNELS.format( - T1=T - 1, - tb=tb, - NW=layout.num_words, - SCAN=SCAN_THREADS, - NO=max(1, len(self._observables)), - NB=max(nb, 1), - RB=8 * layout.num_words, - NOUT=max(layout.num_outputs, 1), - LAST=hex((1 << (layout.num_outputs % 8 or 8)) - 1), - calls=parts.calls, - seeds=parts.seeds, - stores=parts.stores, - hits=parts.hits, - ) + source = kernels.source( + layout, self._noise, parts, tb, max(1, len(self._observables)) ) library = runtime.Library(source) - names = ["noise", "tape", "count_bits", "list_bits", "gather_rows"] - names += ["gather_float", "reset", "advance", "unpack_rows"] - if self._noise.shared_words: - names.append("noise_kept") - for name in names: + for name in kernels.names(self._noise.shared_words): library.pipeline(name) compiled = _Module( library=library, @@ -679,33 +478,45 @@ def sample( GPU, the call warns and samples on the CPU. """ - def on_cpu(): - return self._host.sample( + def on_gpu(): + rows, _, deviation = self.sample_accepted( + shots, + mask=mask, + flip=flip, + batch_size=batch_size, + bools=not bit_packed, + ) + 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 self._or_cpu( + on_gpu, + lambda: self._host.sample( shots, columns, bit_packed=bit_packed, batch_size=batch_size, mask=mask, flip=flip, - ) + ), + ) + + def _or_cpu(self, on_gpu, on_cpu): + """Return ``on_gpu()``, or warn and return ``on_cpu()`` where the GPU fails. + It fails in a process forked after its parent used the GPU, and out of + GPU memory. + """ if _forked(): return on_cpu() try: - rows, _, deviation = self.sample_accepted( - shots, - mask=mask, - flip=flip, - batch_size=batch_size, - bools=not bit_packed, - ) + return on_gpu() except MemoryError as error: _warn_out_of_memory(error) return on_cpu() - if not bit_packed: - return [_columns(rows, c) for c in columns], deviation - num_outputs = self.layout.num_outputs - return [_select(rows, c, num_outputs, bit_packed) for c in columns], deviation def _capacity(self, key: tuple, module: _Module, shots: int) -> int: """Return the rows to make room for: the expected accepted shots and a margin. @@ -952,14 +763,10 @@ def _direct_expectations(self, rows: np.ndarray, values: np.ndarray) -> None: def direct_bools(self, shots: int, flip: np.ndarray | None = None) -> np.ndarray: """Return every output bit of a program without components, as booleans.""" - if _forked(): - return self._host.direct_bools(shots, flip) - try: - rows, _, _ = self.sample_accepted(shots, flip=flip, bools=True) - except MemoryError as error: - _warn_out_of_memory(error) - return self._host.direct_bools(shots, flip) - return rows + return self._or_cpu( + lambda: self.sample_accepted(shots, flip=flip, bools=True)[0], + lambda: self._host.direct_bools(shots, flip), + ) def _forked() -> bool: From 1d88e5c0d65e4ec6ae512b7767a91ef6633e0c95 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 04:05:16 -0700 Subject: [PATCH 20/38] refactor(metal): pass a call's shared state as one _Call _run_accepted took 10 parameters, _encode_pack 9 and _encode_batch 8, most of them the same per call; they now take a _Call and the batch's own values. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/metal/sampler.py | 91 ++++++++++----------------- 1 file changed, 32 insertions(+), 59 deletions(-) diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index 9c8bcabb..cbf0a8c0 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -119,6 +119,21 @@ class _Slot: values: runtime.Buffer | None = None +@dataclass +class _Call: + """What a call's batches share: its kernels, buffers, noise passes and output.""" + + module: _Module + slot: _Slot + constant: runtime.Buffer + split: tuple[_Launch, _Launch] | None + expectations: bool + rows: runtime.Buffer + """The accepted rows, bit-packed.""" + values: runtime.Buffer | None + """The accepted rows' expectations, with ``expectations``.""" + + class MetalShotSampler: """Samples shot words of a statevector program on a Metal GPU. @@ -291,20 +306,15 @@ def _slot(self, shots: int, expectations: bool) -> _Slot: return slot def _encode_batch( - self, - encoder: runtime.Encoder, - module: _Module, - slot: _Slot, - shots: int, - first_shot: int, - constant: runtime.Buffer, - expectations: bool, - launch: _Launch, + self, encoder: runtime.Encoder, call: _Call, shots: int, first_shot: int ) -> None: """Encode a batch's noise and tape kernels. Only the shots that pass the screen get their words written. """ + module, slot, constant = call.module, call.slot, call.constant + expectations = call.expectations + launch = call.split[0] if call.split else self._all library = module.library words = slot.words encoder.dispatch(library.pipeline("reset"), [slot.counts], 1, 1) @@ -362,21 +372,18 @@ def _encode_batch( def _encode_pack( self, encoder: runtime.Encoder, - module: _Module, - slot: _Slot, + call: _Call, shots: int, - split: tuple[_Launch, _Launch] | None, - expectations: bool, - rows: runtime.Buffer, - values: runtime.Buffer | None, fixed: tuple[runtime.Buffer, int], ) -> None: - """Encode the listing and packing of a batch's accepted rows into ``rows``. + """Encode the listing and packing of a batch's accepted rows into ``call.rows``. They follow the rows of the call's earlier batches (at the cursor), and - their expectations go to ``values``; ``fixed`` holds the batch's shots, - which a call without a mask returns. + their expectations go to ``call.values``; ``fixed`` holds the batch's + shots, which a call without a mask returns. """ + module, slot, split = call.module, call.slot, call.split + expectations, rows, values = call.expectations, call.rows, call.values library = module.library S = self._noise.block_shots nb = -(-self.layout.num_outputs // 8) @@ -528,21 +535,9 @@ def _capacity(self, key: tuple, module: _Module, shots: int) -> int: expected = shots * min(1.0, share) return min(shots, int(1.05 * expected + 8 * np.sqrt(expected) + 4096)) - def _run_accepted( - self, - module: _Module, - slot: _Slot, - shots: int, - batch: int, - constant: runtime.Buffer, - expectations: bool, - split: tuple[_Launch, _Launch] | None, - rows: runtime.Buffer, - values: runtime.Buffer | None, - capacity: int, - ) -> None: - """Run a call's batches into ``rows`` (and ``values``) of ``capacity`` rows.""" - launch = split[0] if split else self._all + def _run_accepted(self, call: _Call, shots: int, batch: int, capacity: int) -> None: + """Run a call's batches into ``call.rows`` (and ``values``) of ``capacity`` rows.""" + slot = call.slot starts = range(0, shots, batch) fixed = runtime.buffer(len(starts), np.uint32) slot.counts.array[:] = 0 @@ -553,20 +548,8 @@ def _run_accepted( n = min(batch, shots - start) fixed.array[k] = n encoder = runtime.Encoder() - self._encode_batch( - encoder, module, slot, n, start, constant, expectations, launch - ) - self._encode_pack( - encoder, - module, - slot, - n, - split, - expectations, - rows, - values, - (fixed, 4 * k), - ) + self._encode_batch(encoder, call, n, start) + self._encode_pack(encoder, call, n, (fixed, 4 * k)) commands.append(encoder.commit()) except BaseException: self._settle(commands) @@ -698,18 +681,8 @@ def _sample_call( if expectations else None ) - self._run_accepted( - module, - slot, - shots, - batch, - constant, - expectations, - split, - rows, - values, - capacity, - ) + call = _Call(module, slot, constant, split, expectations, rows, values) + self._run_accepted(call, shots, batch, capacity) total = int(slot.counts.array[_CURSOR]) if total <= capacity: break From d5c48c8a39d63b18d43b5654b1a04b2af31ac7e4 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 04:06:42 -0700 Subject: [PATCH 21/38] refactor(metal): let the Metal package choose its sampler tsim/sampler.py knew Metal's width limit, its fork state, its error types and the no-words rule; metal.shot_sampler now decides, and sampler.py keeps only the TSIM_DEVICE check and one call, so a second GPU backend adds one line there. ShotSampler says device = "cpu" itself. Co-Authored-By: Claude Opus 5.5 --- src/tsim/core/parse.py | 1 + src/tsim/noise/channels.py | 6 ++- src/tsim/sampler.py | 65 ++++---------------------- src/tsim/statevector/metal/__init__.py | 64 +++++++++++++++++++++++++ src/tsim/statevector/shared.py | 5 +- src/tsim/statevector/shots.py | 3 ++ 6 files changed, 83 insertions(+), 61 deletions(-) diff --git a/src/tsim/core/parse.py b/src/tsim/core/parse.py index d981c1df..2e5c540a 100644 --- a/src/tsim/core/parse.py +++ b/src/tsim/core/parse.py @@ -171,6 +171,7 @@ def parse_stim_circuit(stim_circuit: stim.Circuit) -> GraphBuilder: name = instruction.name if name == "SHIFT_COORDS": + # TODO: handle visualization annotations in ZX diagrams continue diff --git a/src/tsim/noise/channels.py b/src/tsim/noise/channels.py index 050a846f..0dfb1e87 100644 --- a/src/tsim/noise/channels.py +++ b/src/tsim/noise/channels.py @@ -740,7 +740,8 @@ def sample_words( num_blocks = -(-num_samples // _BLOCK_SHOTS) if len(seeds) != num_blocks: raise ValueError( - f"{num_samples} samples need {num_blocks} block seeds, got {len(seeds)}" + f"{num_samples} samples need {num_blocks} block seeds, " + f"got {len(seeds)}" ) t = self._tables _sample_blocks( @@ -772,7 +773,8 @@ def sample_screened_words( num_blocks = -(-num_samples // _BLOCK_SHOTS) if len(seeds) != num_blocks: raise ValueError( - f"{num_samples} samples need {num_blocks} block seeds, got {len(seeds)}" + 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") diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 265adb0c..35b31d38 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -20,8 +20,7 @@ from tsim.core.graph import prepare_graph from tsim.core.types import CompiledComponent, CompiledProgram, SamplingStrategy from tsim.noise.channels import ChannelSampler -from tsim.statevector.gpu import FORKED -from tsim.statevector.metal import runtime as metal_runtime +from tsim.statevector import metal from tsim.statevector.program import StatevectorComponent from tsim.statevector.shots import ShotSampler from tsim.utils.cuda_helpers import copy_d2h @@ -315,57 +314,6 @@ class SamplerStats: device: str = "cpu" -def _gpu_shots( - host: ShotSampler, - channels: ChannelSampler, - num_detectors: int, - seed: np.random.SeedSequence, -) -> ShotSampler | MetalShotSampler: - """Return a GPU shot sampler of ``host``'s program when one can run it, else ``host``. - - It warns when a GPU is there but cannot be used: in a process forked after its - parent used the GPU, for a component wider than the GPU executor runs, for - noise beyond its tables' limits, and when the GPU sampler cannot be set up - for lack of GPU memory. - """ - # NOTE: as JAX does for the stabilizer-rank strategies, a GPU is used whenever - # there is one. - if os.environ.get("TSIM_DEVICE", "").lower() == "cpu": - return host - if metal_runtime.forked_after_init(): - warnings.warn(FORKED, stacklevel=5) - return host - if not metal_runtime.available(): - return host - from tsim.statevector.gpu.tape import MAX_WIDTH - from tsim.statevector.metal.sampler import MetalShotSampler - - layout = host.layout - if layout.num_words == 0: - return host - if not MetalShotSampler.supports(layout): - width = max(c.width for c in layout.components) - warnings.warn( - f"An Apple GPU is available, but a component needs {width} active " - f"qubits, more than the {MAX_WIDTH} that the GPU executor supports; " - "sampling on the CPU.", - stacklevel=5, - ) - return host - try: - return MetalShotSampler(host, channels, num_detectors, seed) - except MemoryError: - warnings.warn("The GPU is out of memory; sampling on the CPU.", stacklevel=5) - except ValueError as error: - # NOTE: noise beyond the GPU tables' limits (channels of more than 127 - # outcomes, outcomes that flip more than 255 words). - warnings.warn( - f"The GPU cannot sample this circuit ({error}); sampling on the CPU.", - stacklevel=5, - ) - return host - - class _CompiledSamplerBase: """Base class for compiled samplers with common initialization logic.""" @@ -438,8 +386,13 @@ def __init__( shots_seed, ) self._shots = host - if mode == "sequential": - self._shots = _gpu_shots( + # NOTE: as JAX does for the stabilizer-rank strategies, a GPU is used + # whenever there is one. + if ( + mode == "sequential" + and os.environ.get("TSIM_DEVICE", "").lower() != "cpu" + ): + self._shots = metal.shot_sampler( host, self._channel_sampler, self._num_detectors, gpu_seed ) else: @@ -823,7 +776,7 @@ def stats(self) -> SamplerStats: ), peak_bytes_per_shot=peak, compiled_bytes=compiled_bytes, - device=getattr(self._shots, "device", "cpu"), + device=self._shots.device if self._shots is not None else "cpu", ) def __repr__(self) -> str: diff --git a/src/tsim/statevector/metal/__init__.py b/src/tsim/statevector/metal/__init__.py index ab4273a6..39592d15 100644 --- a/src/tsim/statevector/metal/__init__.py +++ b/src/tsim/statevector/metal/__init__.py @@ -7,3 +7,67 @@ and every component fits. Amplitudes are complex64. It needs Apple silicon, where Tsim depends on ``pyobjc-framework-Metal``. """ + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +from tsim.statevector.gpu import FORKED +from tsim.statevector.gpu.tape import MAX_WIDTH, supports +from tsim.statevector.metal import runtime + +if TYPE_CHECKING: + import numpy as np + + from tsim.noise.channels import ChannelSampler + from tsim.statevector.metal.sampler import MetalShotSampler + from tsim.statevector.shots import ShotSampler + +# NOTE: the warnings point at the caller of compile_sampler. +_STACKLEVEL = 5 + + +def shot_sampler( + host: ShotSampler, + channels: ChannelSampler, + num_detectors: int, + seed: np.random.SeedSequence, +) -> ShotSampler | MetalShotSampler: + """Return a Metal shot sampler of ``host``'s program when one can run it, else ``host``. + + It warns when an Apple GPU is there but cannot be used: in a process forked + after its parent used the GPU, for a component wider than ``MAX_WIDTH``, for + noise beyond the GPU tables' limits, and when the sampler cannot be set up + for lack of GPU memory. + """ + if runtime.forked_after_init(): + warnings.warn(FORKED, stacklevel=_STACKLEVEL) + return host + if not runtime.available() or host.layout.num_words == 0: + return host + if not supports(host.layout): + width = max(c.width for c in host.layout.components) + warnings.warn( + f"An Apple GPU is available, but a component needs {width} active " + f"qubits, more than the {MAX_WIDTH} that the GPU executor supports; " + "sampling on the CPU.", + stacklevel=_STACKLEVEL, + ) + return host + from tsim.statevector.metal.sampler import MetalShotSampler + + try: + return MetalShotSampler(host, channels, num_detectors, seed) + except MemoryError: + warnings.warn( + "The GPU is out of memory; sampling on the CPU.", stacklevel=_STACKLEVEL + ) + except ValueError as error: + # NOTE: noise beyond the GPU tables' limits (channels of more than 127 + # outcomes, outcomes that flip more than 255 words). + warnings.warn( + f"The GPU cannot sample this circuit ({error}); sampling on the CPU.", + stacklevel=_STACKLEVEL, + ) + return host diff --git a/src/tsim/statevector/shared.py b/src/tsim/statevector/shared.py index 1d0717fc..66a61f97 100644 --- a/src/tsim/statevector/shared.py +++ b/src/tsim/statevector/shared.py @@ -763,9 +763,8 @@ def run_lanes( if kind in (ROTATE_DIAG, ROTATE): flip = np.empty(size, dtype=np.int64) for r in range(size): - flip[r] = ( - _row_parity(symbols, r, program.sign_row[a]) - ^ (program.sign_const[a]) + flip[r] = _row_parity(symbols, r, program.sign_row[a]) ^ ( + program.sign_const[a] ) diagonal = kind == ROTATE_DIAG cos, factor, coupling = ( diff --git a/src/tsim/statevector/shots.py b/src/tsim/statevector/shots.py index b88bd04f..cbea80f8 100644 --- a/src/tsim/statevector/shots.py +++ b/src/tsim/statevector/shots.py @@ -348,6 +348,9 @@ def read_outputs( class ShotSampler: """Samples shot words of a program whose components are statevector ones.""" + device = "cpu" + """Where this sampler samples (``SamplerStats.device``).""" + def __init__( self, program: CompiledProgram, From 9aa1cf3f1afb6cd36d54ed459d583cce2898c366 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 04:08:37 -0700 Subject: [PATCH 22/38] test(metal): parametrize the tests that repeated one body The detector, measurement and looping comparisons with cat5 are one test, as are the postselection checks in threadgroup and device memory, and the two out-of-memory fallbacks; the same cases run. Co-Authored-By: Claude Opus 5.5 --- test/integration/test_metal_sampler.py | 110 +++++++++---------------- 1 file changed, 40 insertions(+), 70 deletions(-) diff --git a/test/integration/test_metal_sampler.py b/test/integration/test_metal_sampler.py index db1d831f..ba73cf1b 100644 --- a/test/integration/test_metal_sampler.py +++ b/test/integration/test_metal_sampler.py @@ -98,38 +98,49 @@ def _packed(sampler, shots, mask=None): @pytest.mark.parametrize("name", sorted(CIRCUITS)) -def test_metal_detector_sampling_matches_stabilizer_rank(name): - circuit = Circuit(CIRCUITS[name]) - shots = 40_000 - expected = circuit.compile_detector_sampler(strategy="cat5", seed=1).sample( - shots, batch_size=10_000, append_observables=True - ) - got = _gpu(circuit.compile_detector_sampler(seed=2)).sample( - shots, append_observables=True - ) - assert_counts_match(got, expected) - - -@pytest.mark.parametrize("name", sorted(CIRCUITS)) -def test_metal_measurement_sampling_matches_stabilizer_rank(name): +@pytest.mark.parametrize( + ("kind", "looping"), + [("detector", False), ("measurement", False), ("detector", True)], +) +def test_metal_sampling_matches_stabilizer_rank(monkeypatch, name, kind, looping): + if looping: + # NOTE: measurements, projections and pruning in looping tapes. + monkeypatch.setattr(tape, "UNROLL_BUDGET", 0) circuit = Circuit(CIRCUITS[name]) shots = 40_000 - expected = circuit.compile_sampler(strategy="cat5", seed=1).sample( - shots, batch_size=10_000 - ) - got = _gpu(circuit.compile_sampler(seed=2)).sample(shots) + if kind == "detector": + expected = circuit.compile_detector_sampler(strategy="cat5", seed=1).sample( + shots, batch_size=10_000, append_observables=True + ) + got = _gpu(circuit.compile_detector_sampler(seed=2)).sample( + shots, append_observables=True + ) + else: + expected = circuit.compile_sampler(strategy="cat5", seed=1).sample( + shots, batch_size=10_000 + ) + got = _gpu(circuit.compile_sampler(seed=2)).sample(shots) assert_counts_match(got, expected) @pytest.mark.parametrize("name", sorted(CIRCUITS)) -def test_metal_postselection_keeps_the_accepted_distribution(name): +@pytest.mark.parametrize("device_memory", [False, True]) +def test_metal_postselection_keeps_the_accepted_distribution( + monkeypatch, name, device_memory +): circuit = Circuit(CIRCUITS[name]) - if circuit.num_detectors == 0: + if device_memory: + # NOTE: too little threadgroup memory for the rows: the noise kernel + # keeps them in device memory. + monkeypatch.setattr(metal_noise, "SHARED_BYTES", 2048 + 8) + elif circuit.num_detectors == 0: pytest.skip("no detectors to postselect") shots = 40_000 - mask = np.ones(circuit.num_detectors, dtype=bool) + mask = np.ones(circuit.num_detectors, dtype=bool) if circuit.num_detectors else None expected = _accepted(_cpu(circuit.compile_detector_sampler, seed=1), shots, mask) - got = _accepted(_gpu(circuit.compile_detector_sampler(seed=2)), shots, mask) + sampler = _gpu(circuit.compile_detector_sampler(seed=2)) + assert _gpu_shots(sampler)._noise.shared_words != device_memory + got = _accepted(sampler, shots, mask) rate, gpu_rate = len(expected) / shots, len(got) / shots sigma = np.sqrt(rate * (1 - rate) * 2 / shots) assert abs(gpu_rate - rate) <= 5 * sigma + 3 / shots @@ -519,21 +530,6 @@ def test_metal_shots_start_after_the_common_prefix(width): _assert_matches_cpu(circuit, width) -@pytest.mark.parametrize("name", sorted(CIRCUITS)) -def test_metal_looping_tapes_match_stabilizer_rank(monkeypatch, name): - # NOTE: measurements, projections and pruning in looping tapes. - monkeypatch.setattr(tape, "UNROLL_BUDGET", 0) - circuit = Circuit(CIRCUITS[name]) - shots = 40_000 - expected = circuit.compile_detector_sampler(strategy="cat5", seed=1).sample( - shots, batch_size=10_000, append_observables=True - ) - got = _gpu(circuit.compile_detector_sampler(seed=2)).sample( - shots, append_observables=True - ) - assert_counts_match(got, expected) - - def test_components_wider_than_the_limit_run_on_the_cpu(): circuit = _random_unitary_circuit(13, 12) assert _widest(circuit) > tape.MAX_WIDTH @@ -616,24 +612,6 @@ def test_metal_samples_rows_too_wide_for_threadgroup_memory(): np.testing.assert_array_equal(got[:, -1], got[:, -2] ^ got[:, -3]) -@pytest.mark.parametrize("name", sorted(CIRCUITS)) -def test_metal_noise_in_device_memory_matches_the_cpu(monkeypatch, name): - monkeypatch.setattr(metal_noise, "SHARED_BYTES", 2048 + 8) - circuit = Circuit(CIRCUITS[name]) - shots = 40_000 - mask = np.ones(circuit.num_detectors, dtype=bool) if circuit.num_detectors else None - expected = _accepted(_cpu(circuit.compile_detector_sampler, seed=1), shots, mask) - sampler = _gpu(circuit.compile_detector_sampler(seed=2)) - assert not _gpu_shots(sampler)._noise.shared_words - got = _accepted(sampler, shots, mask) - rate, gpu_rate = len(expected) / shots, len(got) / shots - assert ( - abs(gpu_rate - rate) <= 5 * np.sqrt(rate * (1 - rate) * 2 / shots) + 3 / shots - ) - if len(got) and len(expected): - assert_counts_match(got, expected) - - def test_metal_unrolled_width_11_tapes_match_the_cpu_executor(): # NOTE: 64 amplitudes per thread, unrolled, in two register arrays. circuit = _random_unitary_circuit(11, 7, seed=0) @@ -721,8 +699,13 @@ def test_rare_outcome_zero_is_drawn_at_its_rate(p0): assert abs(int(count.array[0]) - expected) <= 5 * np.sqrt(expected) + 1 -def test_calls_out_of_gpu_memory_fall_back_to_the_cpu(monkeypatch): - circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) +@pytest.mark.parametrize( + "text", + # NOTE: a quantum component (``sample``) and direct outputs only (``direct_bools``). + [GHZ_WITH_DIRECT_DETECTOR, "R 0\nX_ERROR(0.1) 0\nM 0\nDETECTOR rec[-1]"], +) +def test_calls_out_of_gpu_memory_fall_back_to_the_cpu(monkeypatch, text): + circuit = Circuit(text) sampler = _gpu(circuit.compile_detector_sampler(seed=1)) def out_of_memory(*args, **kwargs): @@ -896,19 +879,6 @@ def test_noise_beyond_the_gpu_tables_samples_on_the_cpu(): assert sampler.sample(10).shape == (10, 9) -def test_direct_outputs_out_of_gpu_memory_fall_back_to_the_cpu(monkeypatch): - circuit = Circuit("R 0\nX_ERROR(0.1) 0\nM 0\nDETECTOR rec[-1]") - sampler = _gpu(circuit.compile_detector_sampler(seed=1)) - - def out_of_memory(*args, **kwargs): - raise MemoryError("the GPU ran out of memory") - - monkeypatch.setattr(_gpu_shots(sampler), "sample_accepted", out_of_memory) - with pytest.warns(UserWarning, match="out of memory; sampling on the CPU"): - rows = sampler.sample(20_000) - assert rows.shape == (20_000, 1) and abs(rows.mean() - 0.1) < 0.01 - - def test_calls_larger_than_the_counters_run_in_parts(monkeypatch): # NOTE: the 32-bit counters hold MAX_CALL_SHOTS; a larger call is the calls of # its parts. From 82d3ed1f048ad87d3200b7962fa01ff19ea3a1fd Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 04:08:57 -0700 Subject: [PATCH 23/38] docs(metal): shorten the changelog entry and the strategy docstrings The changelog entry described the kernels, and three docstrings repeated a seven-line paragraph; the install page keeps the details. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 15 +++++---------- src/tsim/circuit.py | 24 ++++++++++-------------- src/tsim/sampler.py | 12 +++++------- 3 files changed, 20 insertions(+), 31 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7427f4cb..1beda668 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,16 +8,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added -- `strategy="active_statevector"` samples on the GPU of an Apple-silicon Mac - whenever every component has at most 11 active qubits. On Apple silicon, Tsim - now depends on PyObjC's Metal bindings (`pyobjc-framework-Metal`, about 7 MB to - download); other platforms install nothing new. The noise sampler and every statevector component run as Metal - kernels, in float32; postselection runs on the GPU, and only the kept shots' - rows are written. A seed's samples on the GPU differ from the CPU's and depend - on `batch_size`. The samplers fall back to the CPU, with a warning, for wider - components, when the GPU is out of memory and in a process forked after its - parent used the GPU. `TSIM_DEVICE=cpu` keeps them on the CPU. The samplers' - repr (`samples on metal`) and `SamplerStats.device` say where they sample. +- `strategy="active_statevector"` samples on the GPU of an Apple-silicon Mac, in + float32, whenever every component has at most 11 active qubits (Tsim now depends + on `pyobjc-framework-Metal` there). A seed's GPU samples differ from the CPU's and + depend on `batch_size`; `TSIM_DEVICE=cpu` keeps the samplers on the CPU, and + `SamplerStats.device` and the repr say where they sample. - The docs pages have a menu to copy a page or open it in ChatGPT or Claude, and the README, the install page and the agent skill point to the docs MCP server at `https://bloqade.quera.com/tsim/mcp`. The README and the install page also diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index 0b64c357..816109d2 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -954,13 +954,11 @@ def compile_sampler( strategy: "active_statevector" (the default) to simulate each non-direct component on a dense state vector over its active qubits, or a stabilizer rank decomposition strategy ("cat5", - "bss", "cutting"). "active_statevector" samples on an Apple GPU - when Metal sees one (Apple silicon) and every - component has at most 11 active qubits (in float32), and - otherwise on the CPU, on ``numba.get_num_threads()`` threads. - The GPU's outcome probabilities are float32; ``TSIM_DEVICE=cpu`` - keeps the sampler on the CPU, in double precision. The - stabilizer rank strategies are provisional. + "bss", "cutting"). "active_statevector" samples in float32 on an + Apple GPU when every component has at most 11 active qubits + (``TSIM_DEVICE=cpu`` turns this off), else on the CPU, on + ``numba.get_num_threads()`` threads. The stabilizer rank + strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed reproduces its samples only with the same versions of Tsim and its dependencies (such as NumPy and JAX) and the same backend (CPU or GPU). @@ -996,13 +994,11 @@ def compile_detector_sampler( strategy: "active_statevector" (the default) to simulate each non-direct component on a dense state vector over its active qubits, or a stabilizer rank decomposition strategy ("cat5", - "bss", "cutting"). "active_statevector" samples on an Apple GPU - when Metal sees one (Apple silicon) and every - component has at most 11 active qubits (in float32), and - otherwise on the CPU, on ``numba.get_num_threads()`` threads. - The GPU's outcome probabilities are float32; ``TSIM_DEVICE=cpu`` - keeps the sampler on the CPU, in double precision. The - stabilizer rank strategies are provisional. + "bss", "cutting"). "active_statevector" samples in float32 on an + Apple GPU when every component has at most 11 active qubits + (``TSIM_DEVICE=cpu`` turns this off), else on the CPU, on + ``numba.get_num_threads()`` threads. The stabilizer rank + strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed reproduces its samples only with the same versions of Tsim and its dependencies (such as NumPy and JAX) and the same backend (CPU or GPU). diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 35b31d38..f8a37b7a 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -335,13 +335,11 @@ def __init__( strategy: "active_statevector" (the default) to simulate each non-direct component on a dense state vector over its active qubits, or a stabilizer rank decomposition strategy ("cat5", - "bss", "cutting"). "active_statevector" samples on an Apple GPU - when Metal sees one (Apple silicon) and every - component has at most 11 active qubits (in float32), and - otherwise on the CPU, on ``numba.get_num_threads()`` threads. - The GPU's outcome probabilities are float32; ``TSIM_DEVICE=cpu`` - keeps the sampler on the CPU, in double precision. The - stabilizer rank strategies are provisional. + "bss", "cutting"). "active_statevector" samples in float32 on an + Apple GPU when every component has at most 11 active qubits + (``TSIM_DEVICE=cpu`` turns this off), else on the CPU, on + ``numba.get_num_threads()`` threads. The stabilizer rank + strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed reproduces its samples only with the same versions of Tsim and its dependencies (such as NumPy and JAX) and the same backend (CPU or GPU). From ec5d77c0d5af812159959cba6a27ced610873c21 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 09:56:23 -0500 Subject: [PATCH 24/38] perf(metal): screen a noise block's shots in parallel The screen ran one thread per shot over every output word, so 32 of a block's 256 threads walked 44 words in sequence on end-to-end cultivation d=5. A thread per word and shot now clears the shot's flag, and the constants go into the rows as they are written. The noise kernel of end2end_d5_dsurface15 takes 2.27 instead of 2.95 ms per million shots on an M4 Pro (1.2x the throughput). Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/metal/noise.py | 48 ++++++++++++++--------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/src/tsim/statevector/metal/noise.py b/src/tsim/statevector/metal/noise.py index a7bd39dc..5b1db22d 100644 --- a/src/tsim/statevector/metal/noise.py +++ b/src/tsim/statevector/metal/noise.py @@ -70,26 +70,24 @@ Segments seg = {{seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep}}; add_faults(sw, events, &num_events, args.max_events, ns, args.seed ^ ((u64)bid * 0xD1B54A32D192ED03UL), args.num_segments, seg, records, overflow, tid, nt); + // NOTE: the screen reads the words in parallel, a thread per word and shot (a + // shot's threads only ever clear its flag); the constants go into the rows. + for (uint s = tid; s < {S}; s += nt) passed[s] = (int)s < ns; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint i = tid; i < {OW} * {S}; i += nt) {{ + uint word = i / {S}, s = i - word * {S}; + if ((sw[i] ^ constants[word]) & screen[word]) passed[s] = false; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); for (int s0 = 0; s0 < {S}; s0 += nt) {{ int s = s0 + tid; - bool alive = false; - if (s < ns) {{ - u64 hit = 0; - for (int i = 0; i < {OW}; ++i) {{ - u64 v = sw[i * {S} + s] ^ constants[i]; - hit |= v & screen[i]; - sw[i * {S} + s] = v; - }} - alive = !hit; - }} - if (s < {S}) passed[s] = alive; + bool alive = s < {S} && passed[s]; enqueue(alive, s, s0 + (tid & ~31u), base, ns, lane, args.has_survivors, queue, count, survivors); }} - threadgroup_barrier(mem_flags::mem_threadgroup); device u64* rows = out + base * {NW}; for (int i = tid; i < ns * {NW}; i += nt) {{ int s = i / {NW}, word = i - s * {NW}; - if (passed[s]) rows[i] = sw[word * {S} + s]; + if (passed[s]) rows[i] = sw[word * {S} + s] ^ (word < {OW} ? constants[word] : 0); }} }} @@ -97,6 +95,7 @@ NOISE_DEVICE = _NOISE_SIGNATURE + r"""{{ threadgroup atomic_int num_events; + threadgroup bool passed[{S}]; long base = (long)bid * {S}; int ns = (int)min((long){S}, args.shots - base); device u64* rows = out + base * {NW}; @@ -107,19 +106,20 @@ add_faults(rows, events, &num_events, args.max_events, ns, args.seed ^ ((u64)bid * 0xD1B54A32D192ED03UL), args.num_segments, seg, records, overflow, tid, nt); + // NOTE: as in the kernel above, the screen reads the words in parallel. + for (uint s = tid; s < {S}; s += nt) passed[s] = (int)s < ns; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int i = tid; i < ns * {OW}; i += nt) {{ + int s = i / {OW}, word = i - s * {OW}; + device u64* at = rows + (long)s * {NW} + word; + u64 v = *at ^ constants[word]; + *at = v; + if (v & screen[word]) passed[s] = false; + }} + threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup); for (int s0 = 0; s0 < {S}; s0 += nt) {{ int s = s0 + tid; - bool alive = false; - if (s < ns) {{ - u64 hit = 0; - device u64* row = rows + (long)s * {NW}; - for (int i = 0; i < {OW}; ++i) {{ - u64 v = row[i] ^ constants[i]; - hit |= v & screen[i]; - row[i] = v; - }} - alive = !hit; - }} + bool alive = s < {S} && passed[s]; enqueue(alive, s, s0 + (tid & ~31u), base, ns, lane, args.has_survivors, queue, count, survivors); }} }} From 5f77544032030fc9fbe221c941877e7d302b3002 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 09:56:29 -0500 Subject: [PATCH 25/38] perf(metal): give wide tapes 16 amplitudes per thread From 8 qubits on, a shot's threads hold at most 16 amplitudes instead of 32: fewer registers per thread pay for the extra shuffles. On an M4 Pro tricycle_6_3_2 and unrotated_H_injection (9 qubits) run 1.4x, the 8-qubit circuits 1.1-1.2x faster; 7 qubits tie and 6 lose 10%, so narrower tapes keep 32. CUDA keeps tape.threads_per_shot as it was. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/gpu/tape.py | 6 +++--- src/tsim/statevector/metal/sampler.py | 17 ++++++++++++++++- test/integration/test_metal_sampler.py | 2 +- 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/src/tsim/statevector/gpu/tape.py b/src/tsim/statevector/gpu/tape.py index 341f822b..2e24c518 100644 --- a/src/tsim/statevector/gpu/tape.py +++ b/src/tsim/statevector/gpu/tape.py @@ -150,10 +150,10 @@ def supports(layout: ShotLayout) -> bool: return layout.num_words > 0 and all(c.width <= MAX_WIDTH for c in layout.components) -def threads_per_shot(layout: ShotLayout) -> int: - """Return log2 of a shot's threads: one per ``2^AMPLITUDE_AXES`` amplitudes, at most 32.""" +def threads_per_shot(layout: ShotLayout, amplitude_axes: int = AMPLITUDE_AXES) -> int: + """Return log2 of a shot's threads: one per ``2^amplitude_axes`` amplitudes, at most 32.""" width = max((c.width for c in layout.components), default=0) - return min(5, max(0, width - AMPLITUDE_AXES)) + return min(5, max(0, width - amplitude_axes)) def deviation(bits) -> float: diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index cbf0a8c0..c651c0b3 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -71,6 +71,9 @@ TAPE_THREADS = 128 """Threads per threadgroup of the tape kernel.""" +WIDE_TAPE = 8 +"""Width from which a shot's threads hold at most 16 amplitudes each, not 32.""" + MAX_CALL_SHOTS = 1 << 31 """Most shots of one GPU call; larger calls run in parts (the counters are 32-bit).""" @@ -213,7 +216,7 @@ def _module(self, mask: np.ndarray | None) -> _Module: return self._modules[key] plan = self.plan(mask) layout = self.layout - tb = tape.threads_per_shot(layout) + tb = _threads_per_shot(layout) parts = tape.tape_kernel(layout, plan, mask, self._num_detectors, tb, METAL) source = kernels.source( layout, self._noise, parts, tb, max(1, len(self._observables)) @@ -750,6 +753,18 @@ def _forked() -> bool: return True +def _threads_per_shot(layout) -> int: + """Return log2 of a shot's threads, as ``tape.threads_per_shot`` from ``WIDE_TAPE`` on.""" + # NOTE: fewer registers per thread pay for the shuffles once tapes are wide: on + # an M4 Pro, 32 threads of 16 amplitudes run 9 qubits 1.4x and 8 qubits 1.1-1.2x + # faster than 16 threads of 32; at 7 qubits they tie, at 6 four threads of 16 + # are 10% slower than two of 32. + width = max((c.width for c in layout.components), default=0) + if width >= WIDE_TAPE: + return tape.threads_per_shot(layout, tape.AMPLITUDE_AXES - 1) + return tape.threads_per_shot(layout) + + def _warn_out_of_memory(error: MemoryError) -> None: warnings.warn( f"{error}; sampling on the CPU. Fewer shots per call fit.", diff --git a/test/integration/test_metal_sampler.py b/test/integration/test_metal_sampler.py index ba73cf1b..9400365a 100644 --- a/test/integration/test_metal_sampler.py +++ b/test/integration/test_metal_sampler.py @@ -509,7 +509,7 @@ def _assert_matches_cpu(circuit, width, shots=200_000): @pytest.mark.parametrize(("width", "layers"), [(6, 8), (8, 8), (9, 10), (11, 12)]) def test_metal_wide_components_match_the_cpu_executor(width, layers): - # NOTE: widths 6, 8, 9 and 11 run 2, 8, 16 and 32 threads per shot. + # NOTE: widths 6, 8, 9 and 11 run 2, 16, 32 and 32 threads per shot. circuit = _random_unitary_circuit(width, layers) assert _widest(circuit) == width _assert_matches_cpu(circuit, width) From c7a825d763e06470c30b9112aaea526181e3d95e Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 10:30:48 -0500 Subject: [PATCH 26/38] ci: do not fail the build when the Codecov upload fails Codecov deactivated this repository and now rejects every upload with HTTP 400, which failed the 3.12 job and cancelled the others. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8f816df5..626e53c9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -41,7 +41,7 @@ jobs: uses: codecov/codecov-action@v7 with: files: coverage.xml # optional - fail_ci_if_error: true # optional (default = false) + fail_ci_if_error: false # NOTE: Codecov deactivated this repository use_pypi: true verbose: true # optional (default = false) token: ${{ secrets.CODECOV_TOKEN }} # required From 44623192e77f24b179173a30c69d41163fe43c07 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 10:58:16 -0500 Subject: [PATCH 27/38] fix(metal): check that SIMD lanes follow thread order before using the GPU The kernels take a thread's lane and SIMD group from its index in its threadgroup: a shot's threads exchange amplitudes by lane, and the scans list shots in lane order. The device check only checked the SIMD width, so a GPU that laid lanes out otherwise would have sampled wrongly without an error. A probe kernel now records every thread's lane and group in threadgroups of 128, 256 and 1024 threads (the kernels' sizes); where they are out of order, Metal counts as unavailable and Tsim samples on the CPU. check_simd_width becomes check_device. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/metal/runtime.py | 55 ++++++++++++++++++--- src/tsim/statevector/metal/sampler.py | 2 +- test/unit/statevector/test_metal_runtime.py | 31 ++++++++++++ 3 files changed, 80 insertions(+), 8 deletions(-) create mode 100644 test/unit/statevector/test_metal_runtime.py diff --git a/src/tsim/statevector/metal/runtime.py b/src/tsim/statevector/metal/runtime.py index 1e2d3cbc..b3f144d4 100644 --- a/src/tsim/statevector/metal/runtime.py +++ b/src/tsim/statevector/metal/runtime.py @@ -73,9 +73,28 @@ def queue() -> Any: return device().newCommandQueue() +LANE_PROBE_THREADS = (128, 256, 1024) +"""Threadgroup sizes of the kernels, whose SIMD lanes the device check probes.""" + +_PROBE = r""" +#include +using namespace metal; +kernel void lanes(device uint* out [[buffer(0)]], uint gid [[thread_position_in_grid]], + uint tid [[thread_position_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint group [[simdgroup_index_in_threadgroup]]) +{ + out[gid] = tid | (lane << 12) | (group << 20); +} +""" + + @functools.cache -def check_simd_width() -> None: - """Check that the device runs SIMD groups of 32 threads, as the kernels assume. +def check_device() -> None: + """Check that the device runs SIMD groups of 32 threads, in order, as the kernels assume. + + The kernels take a thread's lane and SIMD group from its index in its + threadgroup (``tid % 32`` and ``tid / 32``): a shot's threads exchange + amplitudes by lane, and the scans list shots in lane order. Raises: RuntimeError: If it does not (e.g. an Intel or AMD GPU). @@ -83,15 +102,37 @@ def check_simd_width() -> None: """ if not device().supportsFamily_(metal().MTLGPUFamilyApple7): raise RuntimeError("the Metal executor needs an Apple GPU (Apple silicon)") - probe = Library( - "#include \nusing namespace metal;\n" - "kernel void probe(device uint* out [[buffer(0)]]) { out[0] = 0; }" - ).pipeline("probe") + probe = Library(_PROBE).pipeline("lanes") width = int(probe.threadExecutionWidth()) if width != 32: raise RuntimeError( f"the Metal executor needs SIMD groups of 32 threads, not {width}" ) + for threads in LANE_PROBE_THREADS: + out = buffer(4 * threads, np.uint32) + encoder = Encoder() + encoder.dispatch(probe, [out], 4, threads) + wait(encoder.commit()) + if not lanes_in_order(out.array, threads): + raise RuntimeError( + "the Metal executor needs SIMD lanes in thread order, which this " + f"GPU does not keep in threadgroups of {threads} threads" + ) + + +def lanes_in_order(probed: np.ndarray, threads: int) -> bool: + """Return whether the ``lanes`` probe of threadgroups of ``threads`` found lanes in order. + + ``probed[gid]`` holds the thread's index in its threadgroup (bits 0-11), its + lane (bits 12-19) and its SIMD group (bits 20-31). + """ + probed = np.asarray(probed, np.uint32) + tid, lane, group = probed & 0xFFF, (probed >> 12) & 0xFF, probed >> 20 + return bool( + np.all(tid == np.arange(len(probed)) % threads) + and np.all(lane == tid % 32) + and np.all(group == tid // 32) + ) def available() -> bool: @@ -105,7 +146,7 @@ def available() -> bool: @functools.cache def _device_available() -> bool: try: - check_simd_width() + check_device() return True except Exception: return False diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index c651c0b3..ad4acb21 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -164,7 +164,7 @@ def __init__( RuntimeError: If there is no Apple GPU. """ - runtime.check_simd_width() + runtime.check_device() assert self.supports(host.layout) self._args = (host, channels, num_detectors, seed) # NOTE: a call's batches share the sampler's buffers, so calls from several diff --git a/test/unit/statevector/test_metal_runtime.py b/test/unit/statevector/test_metal_runtime.py new file mode 100644 index 00000000..d9ea11e0 --- /dev/null +++ b/test/unit/statevector/test_metal_runtime.py @@ -0,0 +1,31 @@ +import numpy as np +import pytest + +from tsim.statevector.metal.runtime import lanes_in_order + + +def _probed(tid, lane, group): + return (tid | (lane << 12) | (group << 20)).astype(np.uint32) + + +@pytest.mark.parametrize("threads", [128, 1024]) +def test_lanes_in_thread_order_pass_the_probe(threads): + tid = np.tile(np.arange(threads), 3) + assert lanes_in_order(_probed(tid, tid % 32, tid // 32), threads) + + +@pytest.mark.parametrize( + "lane_of", + [ + lambda tid: (tid * 5) % 32, # lanes permuted within each SIMD group + lambda tid: (tid % 32) ^ ((tid // 32) & 1), # swapped in every other group + ], +) +def test_lanes_out_of_order_fail_the_probe(lane_of): + tid = np.tile(np.arange(256), 2) + assert not lanes_in_order(_probed(tid, lane_of(tid), tid // 32), 256) + + +def test_simd_groups_out_of_order_fail_the_probe(): + tid = np.arange(256) + assert not lanes_in_order(_probed(tid, tid % 32, 7 - tid // 32), 256) From f6b033654186fb8574ae44103916052027fc0cd3 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 10:58:37 -0500 Subject: [PATCH 28/38] fix(metal): sample on the CPU when a kernel does not compile A kernel compiles on a sampler's first call with each postselection mask; a compile error raised out of sample(). Library now raises CompileError, on which the sampler warns and samples on the CPU, as out of GPU memory, and stays there: it neither recompiles nor warns on each call, and its stats say cpu. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/metal/runtime.py | 17 +++++++++++++---- src/tsim/statevector/metal/sampler.py | 26 +++++++++++++++++++------- test/integration/test_metal_sampler.py | 22 ++++++++++++++++++++++ 3 files changed, 54 insertions(+), 11 deletions(-) diff --git a/src/tsim/statevector/metal/runtime.py b/src/tsim/statevector/metal/runtime.py index b3f144d4..24ea971a 100644 --- a/src/tsim/statevector/metal/runtime.py +++ b/src/tsim/statevector/metal/runtime.py @@ -238,6 +238,10 @@ def upload(values: np.ndarray) -> Buffer: return held +class CompileError(RuntimeError): + """A Metal kernel that does not compile, or whose pipeline cannot be built.""" + + class Library: """A compiled Metal library and its kernels' pipeline states.""" @@ -245,7 +249,7 @@ def __init__(self, source: str) -> None: """Compile ``source`` (Metal Shading Language) with fast math. Raises: - RuntimeError: If it does not compile. + CompileError: If it does not compile. """ mtl = metal() @@ -255,12 +259,17 @@ def __init__(self, source: str) -> None: source, options, None ) if library is None: - raise RuntimeError(f"Metal compilation failed: {error}") + raise CompileError(f"Metal compilation failed: {error}") self._library = library self._pipelines: dict[str, Any] = {} def pipeline(self, name: str) -> Any: - """Return the compute pipeline state of kernel ``name``.""" + """Return the compute pipeline state of kernel ``name``. + + Raises: + CompileError: If its pipeline cannot be built. + + """ if name not in self._pipelines: function = self._library.newFunctionWithName_(name) if function is None: @@ -275,7 +284,7 @@ def pipeline(self, name: str) -> Any: descriptor, 0, None, None ) if state is None: - raise RuntimeError(f"pipeline {name}: {error}") + raise CompileError(f"pipeline {name}: {error}") self._pipelines[name] = state return self._pipelines[name] diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index ad4acb21..4d19bc9c 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -183,9 +183,12 @@ def __init__( self._constants: dict[bytes, runtime.Buffer] = {} self._slots: dict[tuple[int, bool], _Slot] = {} self._shares: dict[tuple, float] = {} + self._cpu_only = False - device = "metal" - """Where this sampler samples (``SamplerStats.device``).""" + @property + def device(self) -> str: + """Return where this sampler samples (``SamplerStats.device``).""" + return "cpu" if self._cpu_only else "metal" def __reduce__(self): """Pickle the CPU sampler and the inputs; the GPU state is rebuilt on load.""" @@ -484,8 +487,9 @@ def sample( As ``ShotSampler.sample``: with a ``mask`` only the shots that pass postselection are returned. ``batch_size`` bounds the shots of a batch. - Out of GPU memory, and in a process forked after its parent used the - GPU, the call warns and samples on the CPU. + Out of GPU memory, when a kernel does not compile, and in a process + forked after its parent used the GPU, the call warns and samples on the + CPU. """ def on_gpu(): @@ -517,16 +521,24 @@ def on_gpu(): def _or_cpu(self, on_gpu, on_cpu): """Return ``on_gpu()``, or warn and return ``on_cpu()`` where the GPU fails. - It fails in a process forked after its parent used the GPU, and out of - GPU memory. + It fails in a process forked after its parent used the GPU, out of GPU + memory, and when a kernel does not compile; after that the sampler + samples on the CPU. """ - if _forked(): + if self._cpu_only or _forked(): return on_cpu() try: return on_gpu() except MemoryError as error: _warn_out_of_memory(error) return on_cpu() + except runtime.CompileError as error: + self._cpu_only = True + warnings.warn( + f"{error}; this sampler samples on the CPU from now on.", + stacklevel=_STACKLEVEL, + ) + return on_cpu() def _capacity(self, key: tuple, module: _Module, shots: int) -> int: """Return the rows to make room for: the expected accepted shots and a margin. diff --git a/test/integration/test_metal_sampler.py b/test/integration/test_metal_sampler.py index 9400365a..ed813903 100644 --- a/test/integration/test_metal_sampler.py +++ b/test/integration/test_metal_sampler.py @@ -14,6 +14,7 @@ from tsim.statevector.gpu import noise as gpu_noise, tape from tsim.statevector.metal import ( codegen, + kernels, noise as metal_noise, runtime, sampler as metal_sampler, @@ -659,6 +660,27 @@ def test_processes_forked_after_the_gpu_was_used_sample_on_the_cpu(): assert any("forked" in m for m in messages) +def test_a_kernel_that_does_not_compile_samples_on_the_cpu(monkeypatch): + circuit = Circuit(T_STATE_X) + sampler = _gpu(circuit.compile_sampler(seed=1)) + compiled = [] + + def broken(*args, **kwargs): + compiled.append(1) + return "kernel void tape(" + + monkeypatch.setattr(kernels, "source", broken) + with pytest.warns(UserWarning, match="Metal compilation failed"): + got = sampler.sample(20_000) + assert sampler.stats.device == "cpu" + with warnings.catch_warnings(): + warnings.simplefilter("error") + got = np.concatenate([got, sampler.sample(20_000)]) + assert len(compiled) == 1 + expected = _cpu(circuit.compile_sampler, seed=2).sample(40_000) + assert_counts_match(got, expected) + + def test_noise_of_tiny_probabilities_fires_at_its_rate(): # NOTE: 24-bit skips alone fire every segment at least 2^-24 of the time, # here ~350 faults instead of ~0.3. From a796624ea4c9f2373cf71c040492b803598f06fa Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 10:58:43 -0500 Subject: [PATCH 29/38] fix(metal): hash each noise block's stream A segment's stream was mix64(seed ^ bid * C ^ sg * GOLDEN): two (block, segment) pairs with equal XORs shared their draws. The block's stream is now hashed first, so streams of different blocks are unrelated. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/metal/noise.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/tsim/statevector/metal/noise.py b/src/tsim/statevector/metal/noise.py index 5b1db22d..d0455a7f 100644 --- a/src/tsim/statevector/metal/noise.py +++ b/src/tsim/statevector/metal/noise.py @@ -68,7 +68,7 @@ if (tid == 0) atomic_store_explicit(&num_events, 0, memory_order_relaxed); threadgroup_barrier(mem_flags::mem_threadgroup); Segments seg = {{seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep}}; - add_faults(sw, events, &num_events, args.max_events, ns, args.seed ^ ((u64)bid * 0xD1B54A32D192ED03UL), + add_faults(sw, events, &num_events, args.max_events, ns, block_stream(args.seed, bid), args.num_segments, seg, records, overflow, tid, nt); // NOTE: the screen reads the words in parallel, a thread per word and shot (a // shot's threads only ever clear its flag); the constants go into the rows. @@ -104,7 +104,7 @@ threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup); Segments seg = {{seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep}}; add_faults(rows, events, &num_events, args.max_events, ns, - args.seed ^ ((u64)bid * 0xD1B54A32D192ED03UL), args.num_segments, seg, records, + block_stream(args.seed, bid), args.num_segments, seg, records, overflow, tid, nt); // NOTE: as in the kernel above, the screen reads the words in parallel. for (uint s = tid; s < {S}; s += nt) passed[s] = (int)s < ns; @@ -157,7 +157,7 @@ threadgroup_barrier(mem_flags::mem_threadgroup); Segments seg = {{seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep}}; add_faults(sw, events, &num_events, args.max_events, ns, - args.seed ^ ((u64)bid * 0xD1B54A32D192ED03UL), args.num_segments, seg, records, + block_stream(args.seed, bid), args.num_segments, seg, records, overflow, tid, nt); pack_rows(sw, ns, (threadgroup uchar*)events, dst, ((long)*cursor + base) * {NB}, tid, nt); }} @@ -238,6 +238,13 @@ for (int i = head + body * 16 + tid; i < total; i += nt) d[i] = stage[i]; }} +// The stream of block ``bid``. NOTE: hashed, so that a segment's stream +// ``mix64(stream ^ sg * GOLDEN)`` cannot equal another block's: with ``seed ^ bid * C`` +// alone, two (block, segment) pairs with equal XORs would share their draws. +inline u64 block_stream(u64 seed, uint bid) {{ + return mix64(seed ^ ((u64)bid * 0xD1B54A32D192ED03UL)); +}} + struct Segments {{ device const int* first; device const int* outcomes; device const int* begin; device const int* end; device const float* log_keep; From 981a69417a6b777c726c5e21c9e60928f4aa2f9c Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 10:58:48 -0500 Subject: [PATCH 30/38] test(metal): check that noise of tiny probabilities fires at its rate The test only bounded the faults from above, so it passed if the 48-bit path never fired. Probabilities near 3e-10 keep the channels on that path and give ~83 expected faults, which the test now checks both ways: with the 48-bit path off, 24-bit skips add ~110 faults and it fails. Co-Authored-By: Claude Opus 5.5 --- test/integration/test_metal_sampler.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/test/integration/test_metal_sampler.py b/test/integration/test_metal_sampler.py index ed813903..a8e528cb 100644 --- a/test/integration/test_metal_sampler.py +++ b/test/integration/test_metal_sampler.py @@ -682,18 +682,22 @@ def broken(*args, **kwargs): def test_noise_of_tiny_probabilities_fires_at_its_rate(): - # NOTE: 24-bit skips alone fire every segment at least 2^-24 of the time, - # here ~350 faults instead of ~0.3. - n, p = 1000, 1e-12 - text = "\n".join(f"X_ERROR({p * (1 + i / n)!r}) {i}" for i in range(n)) + # NOTE: segments of fewer than 2^-12 expected faults take the kernels' 48-bit + # path. 24-bit skips alone would fire every segment at least 2^-24 of the + # time, ~110 more faults than the ~83 expected here; a path that never fired + # would give none. + n, p, calls = 1000, 3e-10, 44 + probabilities = [p * (1 + i / n) for i in range(n)] + text = "\n".join(f"X_ERROR({q!r}) {i}" for i, q in enumerate(probabilities)) sampler = _gpu( Circuit(f"{text}\nM " + " ".join(map(str, range(n)))).compile_sampler(seed=1) ) # NOTE: faults are rare enough that each lands in a byte of its own. faults = sum( - np.count_nonzero(sampler.sample(1 << 22, bit_packed=True)) for _ in range(44) + np.count_nonzero(sampler.sample(1 << 22, bit_packed=True)) for _ in range(calls) ) - assert faults <= 5 + expected = sum(probabilities) * calls * (1 << 22) + assert abs(faults - expected) <= 5 * np.sqrt(expected) @pytest.mark.parametrize("p0", [3e-8, 1e-6]) From b17457cff3581900a011f78595a729ba5751edf7 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 10:58:54 -0500 Subject: [PATCH 31/38] docs(metal): say that the GPU never draws outcomes below about 1e-11 Co-Authored-By: Claude Opus 5.5 --- docs/install.mdx | 5 +++-- skills/simulating-circuits-with-tsim/SKILL.md | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/docs/install.mdx b/docs/install.mdx index 3879f918..966ea637 100644 --- a/docs/install.mdx +++ b/docs/install.mdx @@ -45,8 +45,9 @@ gain the most. A sampler's repr says where it samples (`samples on metal` or `sa cpu`), and so does `sampler.stats.device`. The GPU computes outcome probabilities in float32, so a sample follows the CPU's distribution -to within float32 rounding. A seed's samples on the GPU differ from the CPU's and depend on -`batch_size`. +to within float32 rounding. A measurement outcome with probability below about 1e-11 is never +drawn on the GPU: float32 cannot tell it from rounding error. A seed's samples on the GPU +differ from the CPU's and depend on `batch_size`. ### Sampling on the CPU instead diff --git a/skills/simulating-circuits-with-tsim/SKILL.md b/skills/simulating-circuits-with-tsim/SKILL.md index e7023925..c8773a69 100644 --- a/skills/simulating-circuits-with-tsim/SKILL.md +++ b/skills/simulating-circuits-with-tsim/SKILL.md @@ -191,8 +191,8 @@ yourself. Tsim processes side by side; they oversubscribe the CPU. - **Apple GPUs:** on an Apple-silicon Mac the default strategy samples on the GPU when every component has at most 11 active qubits, often 5 to 10 times faster than all CPU cores. It - computes outcome probabilities in float32; set `TSIM_DEVICE=cpu` before compiling to stay on - the CPU, in double precision. The sampler's repr says where it samples (`samples on metal`). + computes outcome probabilities in float32 and never draws an outcome below about 1e-11; set + `TSIM_DEVICE=cpu` before compiling to stay on the CPU, in double precision. The sampler's repr says where it samples (`samples on metal`). - **Seeds:** a `seed` reproduces samples for the same Tsim and dependency versions and backend. On the CPU, thread count and `batch_size` do not affect the samples; on a GPU, `batch_size` does. From 3923b4409cf228e139fdf1e2a564b7af6ef76567 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 14:57:14 -0500 Subject: [PATCH 32/38] Revert "ci: do not fail the build when the Codecov upload fails" Codecov accepts uploads again, so a failed upload fails the build again. This reverts commit c7a825d7. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 626e53c9..8f816df5 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -41,7 +41,7 @@ jobs: uses: codecov/codecov-action@v7 with: files: coverage.xml # optional - fail_ci_if_error: false # NOTE: Codecov deactivated this repository + fail_ci_if_error: true # optional (default = false) use_pypi: true verbose: true # optional (default = false) token: ${{ secrets.CODECOV_TOKEN }} # required From 405f28f4da64b6de8841fb91984af777951d3803 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 15:13:13 -0700 Subject: [PATCH 33/38] test(statevector): run looping and unrolled GPU tapes against the CPU executor The looping tape path (UNROLL_BUDGET exceeded) was never generated in CI. Compile each tape in a C++ dialect for one thread per shot with g++ and check that looping and unrolled tapes (in each dialect variant) agree with each other and with shared.run_lanes on random shot words. Also cover rows.select's partial last byte. Co-Authored-By: Claude Opus 5.5 --- test/unit/statevector/test_gpu.py | 223 ++++++++++++++++++++++++++++++ 1 file changed, 223 insertions(+) diff --git a/test/unit/statevector/test_gpu.py b/test/unit/statevector/test_gpu.py index ff6fb8f4..d2cb7062 100644 --- a/test/unit/statevector/test_gpu.py +++ b/test/unit/statevector/test_gpu.py @@ -1,6 +1,9 @@ """The host side of the GPU backends: noise tables, row packing and the tape generator.""" +import ctypes import re +import shutil +import subprocess from test.helpers.statevector import CIRCUITS import numpy as np @@ -8,6 +11,7 @@ from tsim.circuit import Circuit from tsim.noise.channels import ChannelSampler +from tsim.statevector import shared from tsim.statevector.gpu import noise, rows, tape from tsim.statevector.shots import ShotSampler @@ -213,6 +217,7 @@ def test_select_reads_the_columns_of_rows(): packed = np.packbits(bits, axis=1, bitorder="little") for columns in ( np.arange(count), + np.arange(count - 3), np.arange(num_outputs), np.sort(rng.choice(num_outputs, count, replace=False)), ): @@ -308,6 +313,224 @@ def test_projections_read_pairs_or_every_partner_first(): assert re.search(r"\{ float2 q\d+ = .*, q\d+ = ", paired) +GXX = shutil.which("g++") + +# NOTE: a tape in C++ for one thread per shot (tb = 0): no shuffles. ``draw_one`` +# follows the CPU executor's ``shared._cut`` in double. +CPP_DIALECT = tape.Dialect( + inline="static inline", + noinline="__attribute__((noinline)) static", + private="", + deviation="unsigned*", + constant="static const", + float2="float2", + rsqrt="rsqrt", + fmax="std::fmax", + u64="ULL", + shuffle2="{value}", + shuffle1="{value}", + lane_mask=False, + goto=False, + amplitude_chunk=32, + pair_projections=True, +) + +CPP_PRELUDE = r""" +#include +#include +typedef unsigned long long u64; +struct float2 { + float x, y; + float2() = default; + float2(float x_, float y_) : x(x_), y(y_) {} +}; +#define GOLDEN 0x9E3779B97F4A7C15ULL +static inline u64 mix64(u64 z) { + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + return z ^ (z >> 31); +} +static inline u64 draw(u64 seed, u64 shot, u64 slots, u64 slot) { + return mix64(seed + (shot * slots + slot) * GOLDEN + GOLDEN); +} +static inline int par(u64 v) { return __builtin_popcountll(v) & 1; } +static inline u64 threshold(double p) { + double scaled = p * 18446744073709551616.0; + return scaled < 1.8446744073709550e19 ? (u64)scaled : 0xFFFFFFFFFFFFF800ULL; +} +static inline int draw_one(u64 r, float w0, float w1, float) { + double total = (double)w0 + (double)w1; + if (w1 <= w0) return r < threshold(w1 / total); + return r < 0ULL - threshold(w0 / total); +} +static inline float rsqrt(float x) { return 1.f / std::sqrt(x); } +static inline float2 cmul(float2 p, float2 q) { + return float2(p.x * q.x - p.y * q.y, p.x * q.y + p.y * q.x); +} +static inline float2 cmac(float2 a, float c, float s, float2 q) { + return float2(std::fmaf(c, a.x, s * q.x), std::fmaf(c, a.y, s * q.y)); +} +static inline float2 cmac_i(float2 a, float c, float s, float2 q) { + return float2(std::fmaf(c, a.x, -s * q.y), std::fmaf(c, a.y, s * q.x)); +} +static inline void deviate(unsigned* deviation, float total) { + float d = std::fabs(total - 1.f); + unsigned bits; + std::memcpy(&bits, &d, 4); + if (!(d <= 1e-3f) && bits > *deviation) *deviation = bits; +} +""" + + +def _compile_tapes(sources: list[str], directory) -> list: + """Compile tape functions ``component_0`` into ``run_(words, shots, nwords, seed, grave, ev, nev, deviation)``.""" + parts = [CPP_PRELUDE] + for k, source in enumerate(sources): + parts.append( + f"namespace v{k} {{\n{source}\n" + f'extern "C" void run_{k}(u64* words, long shots, long nwords, u64 seed,' + " int* grave, float* ev, long nev, unsigned* deviation) {\n" + " for (long s = 0; s < shots; ++s)\n" + " grave[s] = component_0(words + s * nwords, (u64)s, seed, 0," + " nev ? ev + s * nev : nullptr, deviation);\n}\n}" + ) + cpp, lib = directory / "tapes.cpp", directory / "tapes.so" + cpp.write_text("\n".join(parts)) + # NOTE: -O0 compiles the unrolled tapes several times faster; ``cmac`` and + # ``cmac_i`` fuse explicitly, so the arithmetic does not depend on it. + assert GXX is not None + flags = ["-std=c++17", "-O0", "-ffp-contract=off", "-w", "-shared", "-fPIC"] + subprocess.run([GXX, *flags, "-o", str(lib), str(cpp)], check=True) # noqa: S603 + dll = ctypes.CDLL(str(lib)) + functions = [] + for k in range(len(sources)): + f = getattr(dll, f"run_{k}") + f.restype = None + functions.append(f) + return functions + + +def _run_tape(f, words, seed, nev): + words = words.copy() + shots, nwords = words.shape + grave = np.empty(shots, np.int32) + ev = np.zeros((shots, max(nev, 1)), np.float32) + deviation = np.zeros(1, np.uint32) + f( + words.ctypes.data_as(ctypes.c_void_p), + ctypes.c_long(shots), + ctypes.c_long(nwords), + ctypes.c_uint64(seed), + grave.ctypes.data_as(ctypes.c_void_p), + ev.ctypes.data_as(ctypes.c_void_p), + ctypes.c_long(nev), + deviation.ctypes.data_as(ctypes.c_void_p), + ) + assert tape.deviation(deviation[0]) == 0.0 + return words, grave, ev[:, :nev] + + +def _looping_circuit(n=5, layers=5, seed=0): + """A width-``n`` component with every kind of action that a looping tape runs. + + T gates (diagonal rotations), U3 (others), X, Z and Pauli-product mid-circuit + measurements with feedforward, noise, postselected detectors, and two + observables. + """ + rng = np.random.default_rng(seed) + lines = ["R " + " ".join(map(str, range(n))), "H " + " ".join(map(str, range(n)))] + for layer in range(layers): + for q in range(n): + t = rng.uniform(-1, 1, 3) + lines.append(f"U3({t[0]:.4f},{t[1]:.4f},{t[2]:.4f}) {q}") + lines.append("T " + " ".join(map(str, range(0, n, 2)))) + perm = rng.permutation(n) + lines.append("CX " + " ".join(map(str, perm[: n & ~1]))) + lines.append(f"DEPOLARIZE1(0.02) {rng.integers(n)}") + q = int(rng.integers(n)) + if layer % 3 == 0: + lines += [f"MX {q}", f"CZ rec[-1] {(q + 1) % n}"] + elif layer % 3 == 1: + lines += [f"MPP X{q}*Z{(q + 2) % n}*Y{(q + 4) % n}"] + else: + lines += [f"M {q}", f"CX rec[-1] {(q + 3) % n}"] + lines.append("M " + " ".join(map(str, range(n)))) + lines += [ + f"DETECTOR rec[-{n + 1}]", + "DETECTOR rec[-1] rec[-2]", + "OBSERVABLE_INCLUDE(0) rec[-3]", + "OBSERVABLE_INCLUDE(1) rec[-4] rec[-5]", + ] + return "\n".join(lines) + + +@pytest.mark.skipif(GXX is None, reason="needs g++") +@pytest.mark.parametrize( + "text", + [_looping_circuit(), CIRCUITS["midcircuit_feedforward"], CIRCUITS["t_ghz"]], + ids=["looping", "feedforward", "t_ghz"], +) +def test_looping_and_unrolled_tapes_run_as_the_cpu_executor( + text, tmp_path, monkeypatch +): + host, _ = _shots(text) + layout = host.layout + num_detectors = Circuit(text).num_detectors + plan = host.plan(np.ones(num_detectors, bool)) + (component,) = layout.components + (program,) = plan.programs + observables = { + o: index - num_detectors + for o, index in enumerate(component.output_indices) + if index >= num_detectors + } + assert observables + # NOTE: unrolled in each dialect variant, and looping. + variants = [ + CPP_DIALECT, + tape.Dialect(**{**CPP_DIALECT.__dict__, "goto": True}), + tape.Dialect(**{**CPP_DIALECT.__dict__, "amplitude_chunk": None}), + tape.Dialect(**{**CPP_DIALECT.__dict__, "pair_projections": False}), + ] + assert not tape.loops(program, 0, 0) + sources = [tape.tape_function(program, 0, 0, observables, d) for d in variants] + monkeypatch.setattr(tape, "UNROLL_BUDGET", 0) + assert tape.loops(program, 0, 0) + looping = tape.tape_function(program, 0, 0, observables, CPP_DIALECT) + assert "#pragma unroll 1" in looping + *unrolled, looped = _compile_tapes([*sources, looping], tmp_path) + + shots, seed = 2000, 0x1234_5678_9ABC_DEF0 + rng = np.random.default_rng(7) + words = rng.integers(0, 2**63, (shots, layout.num_words), dtype=np.uint64) + words <<= np.uint64(1) + words |= rng.integers(0, 2, words.shape, dtype=np.uint64) + expected = words.copy() + alive = np.ones(shots, np.bool_) + shared.run_lanes( + shots << program.width, + expected, + alive, + 0, + np.uint64(seed), + True, + program, + False, + 1 << 60, + 1, + ) + assert 0 < alive.sum() < shots + got, grave, ev = _run_tape(looped, words, seed, len(observables)) + np.testing.assert_array_equal(grave == np.iinfo(np.int32).max, alive) + np.testing.assert_array_equal(got[alive], expected[alive]) + assert np.abs(ev).max() > 0 + for f in unrolled: + other, other_grave, other_ev = _run_tape(f, words, seed, len(observables)) + np.testing.assert_array_equal(other_grave, grave) + np.testing.assert_array_equal(other, got) + np.testing.assert_allclose(other_ev, ev, atol=1e-5) + + @pytest.mark.parametrize(("p", "qubits"), [(2e-8, 192), (1e-9, 3200)]) def test_tiny_probability_segments_take_the_kernels_fine_path(p, qubits): # NOTE: a 24-bit skip reaches only some trials of a segment of noise below From 4d15b336cd93774726e82c68d578b6b2ef9edccc Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 22:49:37 -0500 Subject: [PATCH 34/38] ci: leave the Metal executor out of Codecov CI has no Apple GPU, so the Metal tests do not run there and the executor's lines always count as uncovered. Co-Authored-By: Claude Opus 5.5 --- codecov.yml | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 codecov.yml diff --git a/codecov.yml b/codecov.yml new file mode 100644 index 00000000..94b8ad40 --- /dev/null +++ b/codecov.yml @@ -0,0 +1,4 @@ +# NOTE: the Metal executor needs an Apple GPU, which CI does not have; its tests +# run on Apple silicon only. +ignore: + - "src/tsim/statevector/metal" From 94cb838211d13dded594a97fd8934cad1eaf35a9 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 23:16:46 -0500 Subject: [PATCH 35/38] fix(statevector): keep rare outcomes of a GPU noise channel A slot's keep threshold had 24 bits, so an outcome below about 1e-8 of its channel rounded to zero and the GPU never drew it, e.g. the X branch of HERALDED_PAULI_CHANNEL_1(0.1, 1e-10, 0, 0). Thresholds now have 40 bits (the low 16 in the record's spare field), and a nonzero one never rounds to zero. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/gpu/noise.py | 32 +++++++++++++++++------ test/unit/statevector/test_gpu.py | 42 ++++++++++++++++++++++++++++--- 2 files changed, 63 insertions(+), 11 deletions(-) diff --git a/src/tsim/statevector/gpu/noise.py b/src/tsim/statevector/gpu/noise.py index c308fdf0..15e8b8d6 100644 --- a/src/tsim/statevector/gpu/noise.py +++ b/src/tsim/statevector/gpu/noise.py @@ -15,11 +15,16 @@ channel ``c`` has its records at ``first + c * outcomes`` and a fault reads one aligned 32-byte record (and its alias's, in the same channel): - alias: bits 0-24 keep threshold (probability * 2^24), bits 25-31 alias slot + alias: bits 0-24 keep threshold's high bits, bits 25-31 alias slot words: the first two flipped words (16 bits each) extra: bits 0-7 number of flips, bits 8-31 start of the rest in the overflow + keep: bits 0-15 keep threshold's low bits masks: the first two flip masks +The keep threshold is the slot's keep probability times ``2^KEEP_BITS``: a fault +draws its slot and a ``KEEP_BITS``-bit fraction, and keeps the slot's outcome if +the fraction is below it (``keep_threshold``). + The flips past an outcome's first two are 16-byte ``(word, mask)`` pairs, in order. The samples follow ``ChannelSampler``'s distribution, not its streams: every block draws from its own counter-based stream. @@ -59,6 +64,9 @@ # faults is drawn from 48 bits instead (the kernels' fine path). TINY_P = 2.0**-23 FINE_FAULTS = 2.0**-13 +# NOTE: an outcome's share of its channel is exact to 2^-40 (24 bits left +# outcomes below ~1e-8 impossible), and a nonzero share never rounds to zero. +KEEP_BITS = 40 DEFER_SHARE = 0.25 """Share of the expected faults that postselection cannot see, from which they are drawn on the accepted shots only, after the tapes (``split``).""" @@ -68,7 +76,7 @@ ("alias", " return float(np.exp(np.log1p(-np.minimum(p, 1 - 1e-12)).sum())) +def _keep(probability: float) -> int: + """Return a slot's keep threshold for its keep ``probability``: at least 1 if positive.""" + keep = min(round(float(probability) * (1 << KEEP_BITS)), 1 << KEEP_BITS) + return max(keep, 1) if probability > 0 else 0 + + +def keep_threshold(record: np.void) -> int: + """Return a record's keep threshold, as the kernels read it from its fields.""" + return ((int(record["alias"]) & 0x1FFFFFF) << 16) | int(record["keep"]) + + def noise_tables( channels: ChannelSampler, flips: WordTables, @@ -332,12 +351,9 @@ def noise_tables( "the GPU noise sampler supports channels of up to 127 outcomes" ) for k in range(lo, hi): - keep = ( - 1 << 24 - if hi - lo == 1 - else min(round(float(fire.alias_prob[k]) * (1 << 24)), 1 << 24) - ) - records[k]["alias"] = keep | ((int(fire.alias_index[k]) - lo) << 25) + keep = 1 << KEEP_BITS if hi - lo == 1 else _keep(fire.alias_prob[k]) + records[k]["alias"] = (keep >> 16) | ((int(fire.alias_index[k]) - lo) << 25) + records[k]["keep"] = keep & 0xFFFF j0, j1 = int(flips.start[k]), int(flips.start[k + 1]) words = flips.word[j0:j1].tolist() masks = flips.mask[j0:j1].tolist() diff --git a/test/unit/statevector/test_gpu.py b/test/unit/statevector/test_gpu.py index d2cb7062..28442858 100644 --- a/test/unit/statevector/test_gpu.py +++ b/test/unit/statevector/test_gpu.py @@ -43,7 +43,7 @@ def _model_words(tables: noise.NoiseTables, shots: int, seed: int) -> np.ndarray """Draw shot words from the tables as a noise kernel reads them. Each segment's trials fire independently with probability ``1 - e^log_keep``; - a fault draws a slot of its channel and a 24-bit fraction, takes the slot's + a fault draws a slot of its channel and a ``KEEP_BITS``-bit fraction, takes the slot's alias past its threshold and XORs the outcome's flips into its row. """ rng = np.random.default_rng(seed) @@ -62,10 +62,10 @@ def _model_words(tables: noise.NoiseTables, shots: int, seed: int) -> np.ndarray continue n = int(per[k]) slot = int(rng.integers(n)) - fraction = int(rng.integers(1 << 24)) + fraction = int(rng.integers(1 << noise.KEEP_BITS)) at = int(first[k]) + channel * n + slot record = groups.records[at] - if fraction >= int(record["alias"]) & 0x1FFFFFF: + if fraction >= noise.keep_threshold(record): record = groups.records[at - slot + (int(record["alias"]) >> 25)] extra = int(record["extra"]) words = [int(record["words"]) & 0xFFFF, int(record["words"]) >> 16] @@ -108,6 +108,42 @@ def test_noise_tables_draw_the_channel_samplers_distribution(shared_words): assert abs(together - 0.05 * (1 - 0.01) ** 2) < 0.01 +@pytest.mark.parametrize( + "channel", + [ + "HERALDED_PAULI_CHANNEL_1(0.1, 1e-10, 0, 0) 0", + "PAULI_CHANNEL_2(0.1, 0, 0, 1e-10, 0.05, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0) 0 1", + "PAULI_CHANNEL_2(0.3, 0, 0, 0.2, 1e-12, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0) 0 1", + "DEPOLARIZE2(0.03) 0 1", + ], +) +def test_noise_tables_keep_every_outcome_probability(channel): + # NOTE: with 24-bit keep thresholds, an outcome below ~1e-8 of its channel + # rounded to zero, and the GPU never drew it. + detectors = 3 if channel.startswith("HERALDED") else 2 + text = f"R 0 1\n{channel}\nM 0 1\n" + "\n".join( + f"DETECTOR rec[{-i}]" for i in range(1, detectors + 1) + ) + host, channels = _shots(text) + tables = noise.noise_tables( + channels, host.layout.noise, host.layout.num_words, 14 * 1024 + ) + fire, full = channels._tables, 1 << noise.KEEP_BITS + lo, hi = int(fire.outcome_start[0]), int(fire.outcome_start[1]) + n = hi - lo + assert n > 1 + drawn, exact = np.zeros(n), np.zeros(n) + for k in range(n): + keep = noise.keep_threshold(tables.records[lo + k]) + alias = int(tables.records[lo + k]["alias"]) >> 25 + drawn[k] += keep / (n * full) + drawn[alias] += (full - keep) / (n * full) + exact[k] += fire.alias_prob[lo + k] / n + exact[fire.alias_index[lo + k] - lo] += (1 - fire.alias_prob[lo + k]) / n + assert ((drawn > 0) == (exact > 0)).all() + np.testing.assert_allclose(drawn, exact, rtol=1e-12, atol=2.0**-40) + + def test_noise_tables_size_blocks_for_their_memory(): host, channels = _shots(NOISY) layout = host.layout From 1da9703dea78d991fe02a71b36b5dea6ebf053a2 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 23:16:52 -0500 Subject: [PATCH 36/38] fix(statevector): keep rotations with rare branches off the GPU The GPU takes an outcome below 1e-11 of the total weight as float32 residue and never draws it, which also dropped genuine branches of small-angle rotations. A layout with a rotation branch below 1e-8 (R_Z by less than about 1e-4 pi) now stays on the CPU. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/gpu/tape.py | 37 +++++++++++++++++++++++++++---- test/unit/statevector/test_gpu.py | 8 +++++++ 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/src/tsim/statevector/gpu/tape.py b/src/tsim/statevector/gpu/tape.py index 341f822b..20e86449 100644 --- a/src/tsim/statevector/gpu/tape.py +++ b/src/tsim/statevector/gpu/tape.py @@ -62,10 +62,16 @@ """Outcome weight, relative to the total, below which an outcome is impossible. float32 amplitudes leave a certain outcome's other branch a weight of ~1e-14 to -1e-12; genuine outcome probabilities of tsim's circuits (T gates, Pauli noise, -small-angle rotations) stay above ~1e-10. +1e-12, which the GPU cannot tell from a genuine outcome probability that small. """ +RESOLVABLE = 1e-8 +"""Smallest branch weight of a rotation (``cos^2`` or ``sin^2`` of its angle) that +a GPU executor runs: a smaller one could give an outcome a genuine probability +below ``CERTAIN``, which the GPU would never draw.""" + +_STRUCTURAL = 1e-20 + UNROLL_BUDGET = 12_000 """Amplitude updates (local spans summed over the actions) up to which a tape unrolls. @@ -145,9 +151,32 @@ def row_parity(self, row: np.ndarray) -> str: def supports(layout: ShotLayout) -> bool: """Return whether a GPU executor runs ``layout``. - It needs shot words and components of at most ``MAX_WIDTH`` qubits. A layout without shot words (no outputs and no noise) stays on the CPU. + It needs shot words, components of at most ``MAX_WIDTH`` qubits and rotations + whose branches the GPU resolves (``rarest_branch`` at least ``RESOLVABLE``). A + layout without shot words (no outputs and no noise) stays on the CPU. + """ + return ( + layout.num_words > 0 + and all(c.width <= MAX_WIDTH for c in layout.components) + and rarest_branch(layout) >= RESOLVABLE + ) + + +def rarest_branch(layout: ShotLayout) -> float: + """Return the smallest nonzero branch weight (``cos^2``, ``sin^2``) of a rotation. + + It is 1 without rotations; a weight below ``_STRUCTURAL`` (of a Clifford + angle, in float64) counts as zero. """ - return layout.num_words > 0 and all(c.width <= MAX_WIDTH for c in layout.components) + rarest = 1.0 + for component in layout.components: + t = component.tape + rotations = (t.kind == ROTATE) | (t.kind == ROTATE_DIAG) + weights = np.concatenate([t.cos[rotations] ** 2, t.sin[rotations] ** 2]) + weights = weights[weights >= _STRUCTURAL] + if len(weights): + rarest = min(rarest, float(weights.min())) + return rarest def threads_per_shot(layout: ShotLayout) -> int: diff --git a/test/unit/statevector/test_gpu.py b/test/unit/statevector/test_gpu.py index 28442858..61c561dd 100644 --- a/test/unit/statevector/test_gpu.py +++ b/test/unit/statevector/test_gpu.py @@ -616,3 +616,11 @@ def bits(value): assert tape.deviation(bits(1e-4)) == 0.0 assert tape.deviation(bits(0.01)) == pytest.approx(0.01) assert np.isnan(tape.deviation(np.uint32(0x7FC00000))) + + +@pytest.mark.parametrize("angle, runs", [(0.25, True), (1e-4, True), (1e-5, False)]) +def test_gpu_runs_only_rotations_whose_rare_branches_it_draws(angle, runs): + # NOTE: R_Z(1e-5) between Hadamards gives outcome 1 a probability of 2.5e-10, + # too close to what float32 rounding leaves a certain outcome (CERTAIN). + host, _ = _shots(f"R 0\nH 0\nR_Z({angle}) 0\nH 0\nM 0\nDETECTOR rec[-1]") + assert tape.supports(host.layout) == runs From 97fd0b46d129fda3782734af6307fef1148ae161 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 23:22:52 -0500 Subject: [PATCH 37/38] fix(metal): draw a noise outcome from 40-bit keep thresholds A fault draws its slot and a 40-bit fraction from the low 40 bits of its draw and compares the fraction with the 40-bit keep thresholds of the shared host tables, so an outcome below about 1e-8 of its channel is no longer impossible. The event keeps the fraction's high 24 bits; the low 16 only matter when those tie with the threshold's (a chance of 2^-24) and come from a hash of the fault then, so the noise kernel costs the same (tsim-dev bench --device gpu, quick suite: no circuit slower than 5%). A circuit with a rotation branch that float32 cannot resolve samples on the CPU, with a warning that says so. Co-Authored-By: Claude Opus 5.5 --- docs/install.mdx | 8 ++- skills/simulating-circuits-with-tsim/SKILL.md | 3 +- src/tsim/statevector/metal/__init__.py | 19 +++++-- src/tsim/statevector/metal/noise.py | 55 +++++++++++-------- test/integration/test_metal_sampler.py | 36 +++++++++++- 5 files changed, 87 insertions(+), 34 deletions(-) diff --git a/docs/install.mdx b/docs/install.mdx index 966ea637..89138f0b 100644 --- a/docs/install.mdx +++ b/docs/install.mdx @@ -45,9 +45,11 @@ gain the most. A sampler's repr says where it samples (`samples on metal` or `sa cpu`), and so does `sampler.stats.device`. The GPU computes outcome probabilities in float32, so a sample follows the CPU's distribution -to within float32 rounding. A measurement outcome with probability below about 1e-11 is never -drawn on the GPU: float32 cannot tell it from rounding error. A seed's samples on the GPU -differ from the CPU's and depend on `batch_size`. +to within float32 rounding. Float32 cannot tell a measurement outcome of probability below +about 1e-11 from rounding error, so the GPU never draws one; a circuit with a rotation small +enough to make such an outcome (an `R_Z` or `R_X` by less than about 1e-4 π) samples on the +CPU, with a warning. A seed's samples on the GPU differ from the CPU's and depend on +`batch_size`. ### Sampling on the CPU instead diff --git a/skills/simulating-circuits-with-tsim/SKILL.md b/skills/simulating-circuits-with-tsim/SKILL.md index c8773a69..b3bfbc4c 100644 --- a/skills/simulating-circuits-with-tsim/SKILL.md +++ b/skills/simulating-circuits-with-tsim/SKILL.md @@ -191,7 +191,8 @@ yourself. Tsim processes side by side; they oversubscribe the CPU. - **Apple GPUs:** on an Apple-silicon Mac the default strategy samples on the GPU when every component has at most 11 active qubits, often 5 to 10 times faster than all CPU cores. It - computes outcome probabilities in float32 and never draws an outcome below about 1e-11; set + computes outcome probabilities in float32 and never draws an outcome below about 1e-11 + (circuits with rotations by less than about 1e-4 π stay on the CPU, with a warning); set `TSIM_DEVICE=cpu` before compiling to stay on the CPU, in double precision. The sampler's repr says where it samples (`samples on metal`). - **Seeds:** a `seed` reproduces samples for the same Tsim and dependency versions and backend. On the CPU, thread count and `batch_size` do not affect the samples; on a GPU, diff --git a/src/tsim/statevector/metal/__init__.py b/src/tsim/statevector/metal/__init__.py index 39592d15..541064a5 100644 --- a/src/tsim/statevector/metal/__init__.py +++ b/src/tsim/statevector/metal/__init__.py @@ -14,7 +14,7 @@ from typing import TYPE_CHECKING from tsim.statevector.gpu import FORKED -from tsim.statevector.gpu.tape import MAX_WIDTH, supports +from tsim.statevector.gpu.tape import MAX_WIDTH, RESOLVABLE, rarest_branch, supports from tsim.statevector.metal import runtime if TYPE_CHECKING: @@ -37,7 +37,8 @@ def shot_sampler( """Return a Metal shot sampler of ``host``'s program when one can run it, else ``host``. It warns when an Apple GPU is there but cannot be used: in a process forked - after its parent used the GPU, for a component wider than ``MAX_WIDTH``, for + after its parent used the GPU, for a component wider than ``MAX_WIDTH`` or a + rotation by an angle too small for float32 amplitudes (``RESOLVABLE``), for noise beyond the GPU tables' limits, and when the sampler cannot be set up for lack of GPU memory. """ @@ -48,10 +49,18 @@ def shot_sampler( return host if not supports(host.layout): width = max(c.width for c in host.layout.components) + if width > MAX_WIDTH: + reason = ( + f"a component needs {width} active qubits, more than the " + f"{MAX_WIDTH} that the GPU executor supports" + ) + else: + reason = ( + f"a rotation has a branch of weight {rarest_branch(host.layout):.1e}, " + f"below the {RESOLVABLE:.0e} that float32 amplitudes resolve" + ) warnings.warn( - f"An Apple GPU is available, but a component needs {width} active " - f"qubits, more than the {MAX_WIDTH} that the GPU executor supports; " - "sampling on the CPU.", + f"An Apple GPU is available, but {reason}; sampling on the CPU.", stacklevel=_STACKLEVEL, ) return host diff --git a/src/tsim/statevector/metal/noise.py b/src/tsim/statevector/metal/noise.py index d0455a7f..d3d55105 100644 --- a/src/tsim/statevector/metal/noise.py +++ b/src/tsim/statevector/metal/noise.py @@ -165,7 +165,7 @@ ) SOURCE = r""" -struct Record {{ uint alias; uint words; uint extra; uint pad; u64 m0; u64 m1; }}; +struct Record {{ uint alias; uint words; uint extra; uint keep; u64 m0; u64 m1; }}; struct Flip {{ uint word; uint pad; u64 mask; }}; inline void xor64(threadgroup u64* sw, int index, u64 mask) {{ @@ -182,15 +182,21 @@ if (hi) atomic_fetch_xor_explicit(halves + 1, hi, memory_order_relaxed); }} -// The fault drawn at record ``rec`` (bits 0-23 of ``draw`` its fraction, bits -// 24-31 its slot in the channel): the slot's outcome or its alias's. Word ``w`` -// of row ``s`` is ``sw[w * WS + s * SS]``. +// The fault drawn at record ``rec`` (bits 0-23 of ``draw`` the high 24 bits of +// its 40-bit fraction, bits 24-31 its slot in the channel): the slot's outcome +// if the fraction is below the keep threshold (gpu/noise.py), else its alias's. +// The fraction's low 16 bits matter only when its high bits tie with the +// threshold's, and then come from a hash of the fault. Word ``w`` of row ``s`` +// is ``sw[w * WS + s * SS]``. template -inline void apply_fault(W sw, int row, uint rec, uint draw, +inline void apply_fault(W sw, int row, uint rec, uint draw, u64 stream, device const Record* records, device const Flip* overflow) {{ Record r = records[rec]; - if ((draw & 0xFFFFFF) >= (r.alias & 0x1FFFFFF)) r = records[rec - (draw >> 24) + (r.alias >> 25)]; + uint high = draw & 0xFFFFFF, keep = r.alias & 0x1FFFFFF; + if (high > keep || (high == keep && + (uint)(mix64(stream ^ ((u64)(rec << 8 | row) * GOLDEN)) & 0xFFFF) >= r.keep)) + r = records[rec - (draw >> 24) + (r.alias >> 25)]; int flips = r.extra & 255; if (flips > 0) xor64(sw, (r.words & 0xFFFF) * WS + row * SS, r.m0); if (flips > 1) xor64(sw, (r.words >> 16) * WS + row * SS, r.m1); @@ -255,8 +261,8 @@ // cannot resolve: it would fire such a segment at least 2^-24 of the time. template inline void add_segment_faults(W sw, threadgroup uint2* events, threadgroup atomic_int* num_events, - int max_events, int ns, u64 state, float lk, int first, int outcomes, int idx, int end, - device const Record* records, device const Flip* overflow) + int max_events, int ns, u64 stream, u64 state, float lk, int first, int outcomes, int idx, + int end, device const Record* records, device const Flip* overflow) {{ while (true) {{ state += GOLDEN; @@ -275,25 +281,26 @@ idx += 1 + (int)skip; int ch = idx / {S}, row = idx - ch * {S}; if (row >= ns) continue; // the last block's missing shots - // NOTE: the skip reads the high 24 bits of r, the outcome its low 32: - // a slot from the high 24 of those, the fraction of the slot below. - u64 x = (u64)((uint)r >> 8) * (uint)outcomes; - uint slot = (uint)(x >> 24); + // NOTE: the skip reads the high 24 bits of r, the outcome its low 40: a + // slot, and the fraction of the slot below. + u64 x = (r & 0xFFFFFFFFFFUL) * (uint)outcomes; + uint slot = (uint)(x >> 40); uint rec = (uint)(first + ch * outcomes) + slot; - uint draw = (uint)(x & 0xFFFFFF) | (slot << 24); + uint draw = (uint)((x >> 16) & 0xFFFFFF) | (slot << 24); int e = atomic_fetch_add_explicit(num_events, 1, memory_order_relaxed); if (e < max_events) events[e] = uint2((uint)row | (rec << 8), draw); - else apply_fault(sw, row, rec, draw, records, overflow); + else apply_fault(sw, row, rec, draw, stream, records, overflow); }} }} template __attribute__((noinline)) void add_fine_segment_faults(W sw, threadgroup uint2* events, - threadgroup atomic_int* num_events, int max_events, int ns, u64 state, float lk, int first, - int outcomes, int idx, int end, device const Record* records, device const Flip* overflow) + threadgroup atomic_int* num_events, int max_events, int ns, u64 stream, u64 state, float lk, + int first, int outcomes, int idx, int end, device const Record* records, + device const Flip* overflow) {{ - add_segment_faults(sw, events, num_events, max_events, ns, state, lk, first, - outcomes, idx, end, records, overflow); + add_segment_faults(sw, events, num_events, max_events, ns, stream, state, lk, + first, outcomes, idx, end, records, overflow); }} // Phases 1 and 2: draw the segments' faults on the block's ``ns`` shots and XOR @@ -309,17 +316,19 @@ int idx = seg.begin[sg] - 1; u64 state = mix64(stream ^ ((u64)sg * GOLDEN)); if ((float)(end - 1 - idx) * -lk < 0x1p-12f) - add_fine_segment_faults(sw, events, num_events, max_events, ns, state, lk, - first, outcomes, idx, end, records, overflow); + add_fine_segment_faults(sw, events, num_events, max_events, ns, stream, + state, lk, first, outcomes, idx, end, records, + overflow); else - add_segment_faults(sw, events, num_events, max_events, ns, state, - lk, first, outcomes, idx, end, records, overflow); + add_segment_faults(sw, events, num_events, max_events, ns, stream, + state, lk, first, outcomes, idx, end, records, + overflow); }} threadgroup_barrier(mem_flags::mem_threadgroup); int ne = min(atomic_load_explicit(num_events, memory_order_relaxed), max_events); for (int e = tid; e < ne; e += nt) {{ uint2 v = events[e]; - apply_fault(sw, v.x & 255, v.x >> 8, v.y, records, overflow); + apply_fault(sw, v.x & 255, v.x >> 8, v.y, stream, records, overflow); }} if (DEV) threadgroup_barrier(mem_flags::mem_device); else threadgroup_barrier(mem_flags::mem_threadgroup); diff --git a/test/integration/test_metal_sampler.py b/test/integration/test_metal_sampler.py index a8e528cb..76edfe02 100644 --- a/test/integration/test_metal_sampler.py +++ b/test/integration/test_metal_sampler.py @@ -1,5 +1,6 @@ """The Metal executor end to end, against the stabilizer-rank strategy and the CPU executor.""" +import itertools import multiprocessing import pickle import threading @@ -539,10 +540,11 @@ def test_components_wider_than_the_limit_run_on_the_cpu(): assert type(sampler._shots) is ShotSampler -@pytest.mark.parametrize(("probability", "shots"), [(1e-10, 1 << 28), (1e-6, 1 << 26)]) +@pytest.mark.parametrize(("probability", "shots"), [(2e-8, 1 << 30), (1e-6, 1 << 26)]) def test_metal_draws_small_outcome_probabilities_at_their_rate(probability, shots): # NOTE: a 24-bit draw takes any positive probability at least 2^-24 = 6e-8 of - # the time; float32 rounding gives certain outcomes such probabilities. + # the time; float32 rounding gives certain outcomes such probabilities. 2e-8 + # is about the smallest a rotation may give on the GPU (tape.RESOLVABLE). theta = float(2 * np.arcsin(np.sqrt(probability)) / np.pi) circuit = Circuit(f"R 0\nI[R_X(theta={theta!r}*pi)] 0\nM 0") sampler = _gpu(circuit.compile_sampler(seed=3)) @@ -700,6 +702,36 @@ def test_noise_of_tiny_probabilities_fires_at_its_rate(): assert abs(faults - expected) <= 5 * np.sqrt(expected) +def test_rare_outcome_of_a_channel_is_drawn_at_its_rate(): + # NOTE: 1000 channels (distinct, so Tsim keeps them apart) fire half the time + # and then flip detector 0 with chance 6e-9: with 24-bit keep thresholds it + # rounded to zero; ~100 flips are expected here. + k, q, calls = 10, 3e-9, 34 + lines = ["R " + " ".join(map(str, range(k + 1)))] + subsets = [ + s for r in range(1, k + 1) for s in itertools.combinations(range(1, k + 1), r) + ][:1000] + for subset in subsets: + lines.append("E(0.5) " + " ".join(f"X{i}" for i in subset)) + lines.append(f"ELSE_CORRELATED_ERROR({q / 0.5}) X0") + lines.append("M " + " ".join(map(str, range(k + 1)))) + lines += [f"DETECTOR rec[{i - k - 1}]" for i in range(k + 1)] + sampler = _gpu(Circuit("\n".join(lines)).compile_detector_sampler(seed=1)) + flips = sum( + int(np.count_nonzero(sampler.sample(1_000_000, bit_packed=True)[:, 0] & 1)) + for _ in range(calls) + ) + expected = len(subsets) * q * calls * 1_000_000 + assert abs(flips - expected) <= 5 * np.sqrt(expected) + + +def test_rotation_with_a_branch_float32_cannot_resolve_samples_on_the_cpu(): + circuit = Circuit("R 0\nH 0\nR_Z(1e-5) 0\nH 0\nM 0\nDETECTOR rec[-1]") + with pytest.warns(UserWarning, match="rotation has a branch of weight 2.5e-10"): + sampler = circuit.compile_detector_sampler(seed=1) + assert sampler.stats.device == "cpu" + + @pytest.mark.parametrize("p0", [3e-8, 1e-6]) def test_rare_outcome_zero_is_drawn_at_its_rate(p0): # NOTE: w1 / total rounds to a multiple of 2^-24 below 1 in float32; the draw From 02e24acaf58967140f3be831b96494191d393c66 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 23:54:06 -0500 Subject: [PATCH 38/38] fix(metal): honor batch_size for component-free boolean samples A sampler without components returns booleans through direct_bools, which dropped batch_size; on the GPU each batch draws its own stream, so booleans and packed rows of one seed and batch size differed. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 2 +- src/tsim/statevector/metal/sampler.py | 13 ++++++++++--- src/tsim/statevector/shots.py | 10 ++++++++-- test/integration/test_metal_sampler.py | 13 +++++++++++++ 4 files changed, 32 insertions(+), 6 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index f8a37b7a..10d08ea4 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -448,7 +448,7 @@ def _sample_outputs( # NOTE: only the shot words are filtered by a mask; the noise sampler # writes direct outputs straight into boolean rows otherwise. if not (self._program.components or bit_packed or mask is not None): - samples = self._shots.direct_bools(shots, flip) + samples = self._shots.direct_bools(shots, flip, batch_size) return [_take_columns(samples, c) for c in columns] arrays, deviation = self._shots.sample( shots, diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index 4d19bc9c..4b58dd89 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -749,11 +749,18 @@ def _direct_expectations(self, rows: np.ndarray, values: np.ndarray) -> None: bit = (rows[:, k >> 3] >> np.uint8(k & 7)) & np.uint8(1) values[:, column] = 1.0 - 2.0 * bit.astype(np.float32) - def direct_bools(self, shots: int, flip: np.ndarray | None = None) -> np.ndarray: + def direct_bools( + self, + shots: int, + flip: np.ndarray | None = None, + batch_size: int | None = None, + ) -> np.ndarray: """Return every output bit of a program without components, as booleans.""" return self._or_cpu( - lambda: self.sample_accepted(shots, flip=flip, bools=True)[0], - lambda: self._host.direct_bools(shots, flip), + lambda: self.sample_accepted( + shots, flip=flip, batch_size=batch_size, bools=True + )[0], + lambda: self._host.direct_bools(shots, flip, batch_size), ) diff --git a/src/tsim/statevector/shots.py b/src/tsim/statevector/shots.py index cbea80f8..9f6ebedb 100644 --- a/src/tsim/statevector/shots.py +++ b/src/tsim/statevector/shots.py @@ -609,11 +609,17 @@ def _sample_screened( self._scratch = scratch return out - def direct_bools(self, shots: int, flip: np.ndarray | None = None) -> np.ndarray: + def direct_bools( + self, + shots: int, + flip: np.ndarray | None = None, + batch_size: int | None = None, + ) -> np.ndarray: """Return every output bit of a program without components, as booleans. The noise sampler writes them straight into the rows. The samples equal - those of ``sample``. ``flip`` is as in ``sample``. + those of ``sample``. ``flip`` is as in ``sample``; ``batch_size`` does not + change the samples on the CPU. """ assert self._direct_tables is not None out = np.empty((shots, self._layout.num_outputs), dtype=np.uint8) diff --git a/test/integration/test_metal_sampler.py b/test/integration/test_metal_sampler.py index 76edfe02..2c8730ce 100644 --- a/test/integration/test_metal_sampler.py +++ b/test/integration/test_metal_sampler.py @@ -230,6 +230,19 @@ def bools(array, count): assert gpu.shape[1] == cpu.shape[1] +def test_direct_booleans_take_the_batch_size_as_packed_rows_do(): + # NOTE: each batch draws its own stream, so a component-free sampler's + # booleans and packed rows match only if both honor batch_size. + circuit = Circuit("X_ERROR(0.3) 0 1 2\nM 0 1 2\nDETECTOR rec[-1]\nDETECTOR rec[-2]") + + def sample(bit_packed): + sampler = _gpu(circuit.compile_detector_sampler(seed=5)) + return sampler.sample(5_000, batch_size=1_000, bit_packed=bit_packed) + + packed = np.packbits(sample(False), axis=1, bitorder="little") + np.testing.assert_array_equal(packed, sample(True)) + + def test_metal_samples_no_shots(): circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) sampler = _gpu(circuit.compile_detector_sampler(seed=1))