From ea4e9e307c2e3a95e94d6fbe6d156732bca26bf5 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 10 Nov 2026 18:37:45 +0100 Subject: [PATCH 01/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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 fd8beac6e642ed6d58311c6dca805655f35d49d0 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 10 Nov 2026 13:06:24 -0700 Subject: [PATCH 09/20] feat(statevector): sample on CUDA GPUs on the shared GPU host code Rebuild #219's CUDA backend on #284's tsim.statevector.gpu: the tape comes from gpu.tape with a CUDA dialect (byte-identical to #219's codegen), the noise tables from gpu.noise and the row packing from gpu.rows. cuda/ keeps the CUDA prelude, the noise kernels' source, the runtime and the sampler. The dispatch follows the Metal backend's shape (_gpu_shots): TSIM_DEVICE=cpu keeps the samplers on the CPU, sampler.stats.device and the repr say where a sampler samples, programs without shot words stay on the CPU, and noise beyond the GPU tables' limits falls back to the CPU with a warning. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 7 +- CHANGELOG.md | 16 + README.md | 8 +- docs/index.mdx | 7 +- docs/install.mdx | 54 +- pyproject.toml | 6 + skills/simulating-circuits-with-tsim/SKILL.md | 7 +- src/tsim/circuit.py | 20 +- src/tsim/sampler.py | 94 +- src/tsim/statevector/cuda/__init__.py | 13 + src/tsim/statevector/cuda/codegen.py | 89 ++ src/tsim/statevector/cuda/noise.py | 262 +++++ src/tsim/statevector/cuda/runtime.py | 100 ++ src/tsim/statevector/cuda/sampler.py | 1005 +++++++++++++++++ test/conftest.py | 10 + test/integration/test_cuda_sampler.py | 705 ++++++++++++ uv.lock | 165 ++- 17 files changed, 2509 insertions(+), 59 deletions(-) create mode 100644 src/tsim/statevector/cuda/__init__.py create mode 100644 src/tsim/statevector/cuda/codegen.py create mode 100644 src/tsim/statevector/cuda/noise.py create mode 100644 src/tsim/statevector/cuda/runtime.py create mode 100644 src/tsim/statevector/cuda/sampler.py create mode 100644 test/conftest.py create mode 100644 test/integration/test_cuda_sampler.py diff --git a/AGENTS.md b/AGENTS.md index 4d00d861..fab765d8 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/cuda/`) 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 97d6a672..281f2850 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- With `strategy="active_statevector"`, the samplers run on a CUDA GPU when + CuPy sees one, as the stabilizer-rank strategies do with JAX: the noise + sampler and every statevector component run as CUDA kernels that CuPy + compiles with NVRTC, in float32, with no trajectory sharing. Circuits with a + component wider than 11 active qubits stay on the CPU, as does everything + under `TSIM_DEVICE=cpu` or `CUDA_VISIBLE_DEVICES=""` and in processes forked + after their parent used the GPU; a warning says so when a GPU is present, as + it does for noise beyond the GPU tables' limits. `sampler.stats.device` and + the sampler's repr say where it samples. With a + `postselection_mask`, only the shots that pass are copied off the GPU. A + seed's samples on the GPU differ from the CPU's and depend on the batch size. + A call that runs out of GPU memory retries in smaller batches, then samples + on the CPU with a warning. Install with `pip install bloqade-tsim[cuda]`. +- The `jax-cuda12` and `jax-cuda13` extras, for the stabilizer-rank strategies + on NVIDIA GPUs with JAX. The `cuda12` and `cuda13` extras of Tsim 1.0 install + the same packages. - 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..6aeda1b1 100644 --- a/README.md +++ b/README.md @@ -23,10 +23,10 @@ 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: +To sample on an NVIDIA GPU, install the extra of your strategy: ```bash -pip install "bloqade-tsim[cuda13]" # or bloqade-tsim[cuda12] +pip install "bloqade-tsim[cuda]" # the default strategy, "active_statevector" (CUDA 13) +pip install "bloqade-tsim[jax-cuda13]" # the stabilizer-rank strategies, with JAX (or jax-cuda12) ``` ### Using Tsim with coding agents @@ -82,7 +82,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 NVIDIA GPU (CuPy, at most 11 active qubits) | 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..d509c932 100644 --- a/docs/index.mdx +++ b/docs/index.mdx @@ -39,11 +39,10 @@ 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: +To sample on an NVIDIA GPU with the default strategy, install the `cuda` extra: ```bash -uv add "bloqade-tsim[cuda13]" +uv add "bloqade-tsim[cuda]" ``` See [Installation](/install) for more options. @@ -55,7 +54,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 NVIDIA GPU (CuPy, at most 11 active qubits) | 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..a09811b2 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. @@ -13,28 +13,58 @@ We recommend using [`uv`](https://docs.astral.sh/uv/) for dependency management. uv add bloqade-tsim ``` -For GPU acceleration with CUDA: +## Using pip ```bash -# For CUDA 13 -uv add "bloqade-tsim[cuda13]" - -# For CUDA 12 -uv add "bloqade-tsim[cuda12]" +pip install bloqade-tsim ``` -## Using pip +## NVIDIA GPUs + +Each sampling strategy has its own GPU extra: + +| Extra | Accelerates | Needs | +|---|---|---| +| `cuda` | `"active_statevector"` (the default), with CuPy | CUDA 13 (driver 580 or later) | +| `jax-cuda13` | `"cat5"`, `"bss"` and `"cutting"`, with JAX | CUDA 13 | +| `jax-cuda12` | `"cat5"`, `"bss"` and `"cutting"`, with JAX | CUDA 12 | ```bash -pip install bloqade-tsim +uv add "bloqade-tsim[cuda]" +# or, with both strategies on the GPU: +uv add "bloqade-tsim[cuda,jax-cuda13]" ``` -For GPU acceleration with CUDA: +The `cuda12` and `cuda13` extras of Tsim 1.0 still install the JAX extras. -```bash -pip install "bloqade-tsim[cuda13]" +### The default strategy + +With the `cuda` extra, `strategy="active_statevector"` samples on the GPU +whenever CuPy sees one and every component has at most 11 active qubits; +otherwise it samples on the CPU, with a warning. The GPU computes in float32, +and a seed's samples differ from the CPU's. With a `postselection_mask`, only +the shots that pass are copied off the GPU, so heavily postselected circuits +gain the most: + +```python +sampler = circuit.compile_detector_sampler() +kept = sampler.sample( + 1_000_000, append_observables=True, bit_packed=True, postselection_mask=mask +) ``` +A sampler's repr says where it samples (`samples on cuda` or `samples on cpu`), and so +does `sampler.stats.device`. Set `TSIM_DEVICE=cpu` (or `CUDA_VISIBLE_DEVICES=""`) to keep +the samplers on the CPU, in double precision, for example to reproduce samples from another +machine. Tsim reads `TSIM_DEVICE` whenever it compiles a sampler. + +### Both strategies in one process + +JAX starts only when a stabilizer-rank strategy samples, and then reserves 75% +of the GPU's memory. To sample with both strategies in one process, set +`XLA_PYTHON_CLIENT_PREALLOCATE=false` before starting Python, so that JAX +allocates as it goes. + ## 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..f6b2fa3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,12 @@ Repository = "https://github.com/QuEraComputing/tsim" Issues = "https://github.com/QuEraComputing/tsim/issues" [project.optional-dependencies] +# The default strategy, "active_statevector", on NVIDIA GPUs (CUDA 13). +cuda = ["cupy-cuda13x[ctk]>=14.0"] +# The stabilizer-rank strategies ("cat5", "bss", "cutting") on NVIDIA GPUs, with JAX. +jax-cuda12 = ["jax[cuda12]>=0.4.38"] +jax-cuda13 = ["jax[cuda13]>=0.7.0"] +# NOTE: the names of the JAX extras up to 1.0, kept so that they still install. cuda12 = ["jax[cuda12]>=0.4.38"] cuda13 = ["jax[cuda13]>=0.7.0"] diff --git a/skills/simulating-circuits-with-tsim/SKILL.md b/skills/simulating-circuits-with-tsim/SKILL.md index e5cfb39b..ff4352c2 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. +- **NVIDIA GPUs:** with the `cuda` extra, the default strategy samples on the GPU when every + component has at most 11 active qubits. 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 cuda`). - **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 05151114..6e54891f 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -954,9 +954,12 @@ 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 - strategies are provisional. + "bss", "cutting"). "active_statevector" samples on a CUDA GPU + when CuPy sees one 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 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). @@ -965,6 +968,7 @@ def compile_sampler( 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. + On a GPU, they do depend on the batch size. Returns: A CompiledMeasurementSampler that can be used to sample measurements. @@ -991,9 +995,12 @@ 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 - strategies are provisional. + "bss", "cutting"). "active_statevector" samples on a CUDA GPU + when CuPy sees one 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 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). @@ -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..67a36e44 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.cuda import runtime as cuda_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.cuda.sampler import CudaShotSampler # NOTE: postselection with a stabilizer-rank strategy screens shots for direct # discards in rounds of SCREEN_CHUNKS chunks of at most SCREEN_SHOTS shots. @@ -308,6 +311,64 @@ class SamplerStats: error_channel_bits: int peak_bytes_per_shot: int compiled_bytes: int + device: str = "cpu" + + +def _gpu_shots( + host: ShotSampler, + channels: ChannelSampler, + num_detectors: int, + seed: np.random.SeedSequence, +) -> ShotSampler | CudaShotSampler: + """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 cuda_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 cuda_runtime.available(): + return host + from tsim.statevector.cuda.sampler import MAX_WIDTH, CudaShotSampler + + layout = host.layout + if layout.num_words == 0: + return host + if not CudaShotSampler.supports(layout): + width = max(c.width for c in layout.components) + warnings.warn( + f"A CUDA 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 CudaShotSampler(host, channels, num_detectors, seed) + 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, + ) + except Exception as error: + if not cuda_runtime.out_of_memory(error): + raise + warnings.warn("The GPU is out of memory; sampling on the CPU.", stacklevel=5) + return host class _CompiledSamplerBase: @@ -331,9 +392,12 @@ 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 - strategies are provisional. + "bss", "cutting"). "active_statevector" samples on a CUDA GPU + when CuPy sees one 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 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). @@ -342,13 +406,17 @@ def __init__( 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. + 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, cuda_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,7 +432,7 @@ 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 | CudaShotSampler | None = None self._backend: _StabilizerRankBackend | None = None if strategy == "active_statevector" or not program.components: self._shots = ShotSampler( @@ -373,6 +441,10 @@ def __init__( self._num_detectors, shots_seed, ) + if mode == "sequential": + self._shots = _gpu_shots( + self._shots, self._channel_sampler, self._num_detectors, cuda_seed + ) else: self._backend = _StabilizerRankBackend(program, key_seed) @@ -721,6 +793,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: + ``"cuda"`` on an NVIDIA GPU, else ``"cpu"`` (a stabilizer-rank + strategy runs wherever JAX does). """ components = self._program.components widths: tuple[int, ...] | None = None @@ -751,6 +826,7 @@ def stats(self) -> SamplerStats: ), peak_bytes_per_shot=peak, compiled_bytes=compiled_bytes, + device=getattr(self._shots, "device", "cpu"), ) def __repr__(self) -> str: @@ -803,6 +879,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 += [ @@ -851,6 +929,7 @@ def __init__( 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. + On a GPU, they do depend on the batch size. """ super().__init__( @@ -951,6 +1030,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/cuda/__init__.py b/src/tsim/statevector/cuda/__init__.py new file mode 100644 index 00000000..e84ec1e8 --- /dev/null +++ b/src/tsim/statevector/cuda/__init__.py @@ -0,0 +1,13 @@ +"""The statevector executor on a CUDA GPU. + +The noise sampler (``noise``) and every component's tape (``codegen``) run as +CUDA kernels that CuPy compiles with NVRTC (``runtime``), from the host code +that the GPU backends share (``tsim.statevector.gpu``); ``CudaShotSampler`` +(``sampler``) replaces ``ShotSampler`` behind the compiled samplers whenever +CuPy sees a GPU and every component fits. Amplitudes are complex64. It needs +CuPy and an NVIDIA GPU (``pip install bloqade-tsim[cuda]``). +""" + +from tsim.statevector.cuda.runtime import available + +__all__ = ["available"] diff --git a/src/tsim/statevector/cuda/codegen.py b/src/tsim/statevector/cuda/codegen.py new file mode 100644 index 00000000..2193b62c --- /dev/null +++ b/src/tsim/statevector/cuda/codegen.py @@ -0,0 +1,89 @@ +"""CUDA source of a component's tape (see ``tsim.statevector.gpu.tape``). + +Partner amplitudes on other threads come from ``__shfl_xor_sync`` among the +shot's group of lanes (``gmask``), and a pruned shot jumps to the end. Near +the boundary ``draw_one`` decides an outcome in double precision. +""" + +from __future__ import annotations + +from tsim.statevector.gpu import tape +from tsim.statevector.shared import LaneProgram + +CUDA = tape.Dialect( + inline="__device__ __forceinline__", + noinline="__device__ __noinline__", + private="", + deviation="unsigned*", + constant="__device__ const", + float2="make_float2", + rsqrt="rsqrtf", + fmax="fmaxf", + u64="ULL", + shuffle2="shfl2({value}, {lanes}, gmask)", + shuffle1="__shfl_xor_sync(gmask, {value}, {lanes})", + lane_mask=True, + goto=True, +) + +PRELUDE = r""" +typedef unsigned long long u64; +#define GOLDEN 0x9E3779B97F4A7C15ULL + +__device__ __forceinline__ u64 mix64(u64 z) { + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + return z ^ (z >> 31); +} +// The CPU executor's draw: splitmix64 at index shot * slots + slot of stream seed. +__device__ __forceinline__ u64 draw(u64 seed, u64 shot, u64 slots, u64 slot) { + return mix64(seed + (shot * slots + slot) * GOLDEN + GOLDEN); +} +__device__ __forceinline__ int par(u64 v) { return __popcll(v) & 1; } +// draw_one in double precision, from 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``). Not inlined, as it is rarely taken. +__device__ __noinline__ int draw_one_exact(u64 r, float w0, float w1, float total) { + int rare0 = w0 < w1; + u64 below = (u64)((double)(rare0 ? w0 : w1) / (double)total * 18446744073709551616.0); + 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; there draw_one_exact does. +__device__ __forceinline__ int draw_one(u64 r, float w0, float w1, float total) { + float p = w1 / total, x = __ull2float_rz(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 logf never sees 0. +__device__ __forceinline__ float unit(u64 r) { return ((r >> 40) + 1) * (1.0f / 16777216.0f); } +__device__ __forceinline__ float2 cmul(float2 p, float2 q) { + return make_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 +__device__ __forceinline__ float2 cmac(float2 a, float c, float s, float2 q) { + return make_float2(fmaf(c, a.x, s * q.x), fmaf(c, a.y, s * q.y)); +} +__device__ __forceinline__ float2 cmac_i(float2 a, float c, float s, float2 q) { + return make_float2(fmaf(c, a.x, -s * q.y), fmaf(c, a.y, s * q.x)); +} +__device__ __forceinline__ float2 shfl2(float2 v, int lane_xor, unsigned mask) { + return make_float2(__shfl_xor_sync(mask, v.x, lane_xor), __shfl_xor_sync(mask, v.y, lane_xor)); +} +// The largest normalization deviation, as the bits of a non-negative float. +__device__ __forceinline__ void deviate(unsigned* deviation, float total) { + float d = fabsf(total - 1.f); + if (!(d <= 1e-3f)) atomicMax(deviation, d == d ? __float_as_uint(d) : 0x7f800000u); +} +""" + +hex64 = CUDA.hex64 + + +def tape_function( + program: LaneProgram, index: int, tb: int, observables: dict[int, int] +) -> str: + """Return ``component_`` in CUDA C++ (``tape.tape_function``).""" + return tape.tape_function(program, index, tb, observables, CUDA) diff --git a/src/tsim/statevector/cuda/noise.py b/src/tsim/statevector/cuda/noise.py new file mode 100644 index 00000000..27013183 --- /dev/null +++ b/src/tsim/statevector/cuda/noise.py @@ -0,0 +1,262 @@ +"""Noise sampling into shot words on a CUDA GPU. + +The kernels read the tables of ``tsim.statevector.gpu.noise`` (see there): a +block owns ``S`` consecutive shots and keeps their words in shared memory, +word major, walks segments of the channel groups with geometric skips, +records the faults as events in shared memory and then XORs their outcomes' +flips into the shots' words. Recording first keeps the warps converged while +the faults are applied, and many of their table reads in flight. + +A block then screens its shots (constants, postselected direct detectors), +writes their rows (coalesced; only the survivors' unless every row is asked +for), appends the surviving shots 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 = 46 * 1024 +"""Shared memory of a block's words and events (the events' share is set per launch).""" + +SOURCE = r""" +struct __align__(32) Record {{ unsigned alias; unsigned words; unsigned extra; unsigned pad; u64 m0; u64 m1; }}; +struct __align__(16) Flip {{ unsigned word; unsigned pad; u64 mask; }}; + +// 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. +__device__ __forceinline__ void apply_fault(u64* sw, int row, unsigned rec, unsigned draw, + const Record* __restrict__ records, const Flip* __restrict__ 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) atomicXor(&sw[(r.words & 0xFFFF) * {S} + row], r.m0); + if (flips > 1) atomicXor(&sw[(r.words >> 16) * {S} + row], r.m1); + for (int j = 2; j < flips; ++j) {{ + Flip f = overflow[(r.extra >> 8) + j - 2]; + atomicXor(&sw[f.word * {S} + row], 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``. Output word ``k`` of a row +// is the pieces ``pieces[start[k] : start[k + 1]]``, each ``(word, bit | (length +// - 1) << 6 | position << 12)``: ``length`` bits from ``bit`` of a shot word. +// Consecutive threads pack one word of consecutive rows (one piece, distinct +// banks) into ``stage``, from which the block's rows go out 16 bytes at a time. +__device__ __forceinline__ void pack_rows(const u64* sw, int ns, const uint2* __restrict__ pieces, + const unsigned* __restrict__ start, unsigned char* stage, unsigned char* __restrict__ dst) +{{ + for (int i = threadIdx.x; i < {NK} * {S}; i += blockDim.x) {{ + int k = i / {S}, s = i - k * {S}; + if (s >= ns) continue; + u64 v = 0; + for (unsigned 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 &= (1ULL << length) - 1; + v |= x << (q.y >> 12); + }} + unsigned char* at = stage + s * {NB} + 8 * k; + #pragma unroll + for (int j = 0; j < 8; ++j) if (8 * k + j < {NB}) at[j] = (unsigned char)(v >> (8 * j)); + }} + __syncthreads(); + int total = ns * {NB}; + for (int i = threadIdx.x; i < total / 16; i += blockDim.x) ((uint4*)dst)[i] = ((const uint4*)stage)[i]; + for (int i = total / 16 * 16 + threadIdx.x; i < total; i += blockDim.x) dst[i] = stage[i]; +}} + +// 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 +__device__ __forceinline__ void add_segment_faults(u64* sw, uint2* events, int* num_events, + int max_events, int ns, u64 state, float lk, int first, int outcomes, int idx, int end, + const Record* __restrict__ records, const Flip* __restrict__ overflow) +{{ + while (true) {{ + state += GOLDEN; + u64 r = mix64(state); + float skip; + if (!FINE) skip = lk < -3.0e38f ? 0.f : floorf(__logf(unit(r)) / lk); + else if ((~r >> 40) >= 4096) skip = floorf(__logf(unit(r)) / lk); + else {{ + // NOTE: 1 - u = (j + w) 2^-24, with w from a fresh draw. + float w = ((mix64(state ^ 0x5851F42D4C957F2DULL) >> 40) + 1) * 0x1p-24f; + skip = floorf(log1pf(-((float)(~r >> 40) + w) * 0x1p-24f) / 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. + unsigned long long x = (unsigned long long)((unsigned)r >> 8) * (unsigned)outcomes; + unsigned slot = (unsigned)(x >> 24); + unsigned rec = (unsigned)(first + ch * outcomes) + slot; + unsigned draw = (unsigned)(x & 0xFFFFFF) | (slot << 24); + int e = atomicAdd(num_events, 1); + if (e < max_events) events[e] = make_uint2((unsigned)row | (rec << 8), draw); + else apply_fault(sw, row, rec, draw, records, overflow); + }} +}} + +__device__ __noinline__ void add_fine_segment_faults(u64* sw, uint2* events, int* num_events, + int max_events, int ns, u64 state, float lk, int first, int outcomes, int idx, int end, + const Record* __restrict__ records, const Flip* __restrict__ 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``. +__device__ __forceinline__ void add_faults(u64* sw, uint2* events, int* num_events, int max_events, + int ns, u64 stream, int num_segments, + const int* __restrict__ seg_first, const int* __restrict__ seg_outcomes, + const int* __restrict__ seg_begin, const int* __restrict__ seg_end, + const float* __restrict__ seg_log_keep, const Record* __restrict__ records, + const Flip* __restrict__ overflow) +{{ + for (int sg = threadIdx.x; sg < num_segments; sg += blockDim.x) {{ + 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); + }} + __syncthreads(); + int ne = min(*num_events, max_events); + for (int e = threadIdx.x; e < ne; e += blockDim.x) {{ + uint2 v = events[e]; + apply_fault(sw, v.x & 255, v.x >> 8, v.y, records, overflow); + }} + __syncthreads(); +}} + +extern "C" __global__ void __launch_bounds__(256) noise( + u64* __restrict__ out, long long shots, u64 seed, int num_segments, int max_events, + const int* __restrict__ seg_first, const int* __restrict__ seg_outcomes, + const int* __restrict__ seg_begin, const int* __restrict__ seg_end, + const float* __restrict__ seg_log_keep, const Record* __restrict__ records, + const Flip* __restrict__ overflow, + const u64* __restrict__ constant, const u64* __restrict__ screen, + unsigned* __restrict__ queue, unsigned* __restrict__ count, unsigned* __restrict__ survivors) +{{ + __shared__ u64 sw[{NW} * {S}]; + extern __shared__ uint2 events[]; + __shared__ int num_events; + __shared__ bool passed[{S}]; + long long base = (long long)blockIdx.x * {S}; + int ns = (int)min((long long){S}, shots - base); + for (int i = threadIdx.x; i < {NW} * {S}; i += blockDim.x) sw[i] = 0; + if (threadIdx.x == 0) num_events = 0; + __syncthreads(); + add_faults(sw, events, &num_events, max_events, ns, seed ^ ((u64)blockIdx.x * 0xD1B54A32D192ED03ULL), + num_segments, seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep, records, + overflow); + for (int s0 = 0; s0 < {S}; s0 += blockDim.x) {{ + int s = s0 + threadIdx.x; + bool alive = false; + if (s < ns) {{ + u64 hit = 0; + #pragma unroll + for (int i = 0; i < {OW}; ++i) {{ + u64 v = sw[i * {S} + s] ^ constant[i]; + hit |= v & screen[i]; + sw[i * {S} + s] = v; + }} + alive = !hit; + }} + if (s < {S}) passed[s] = alive; + // NOTE: one atomic per warp: the leader reserves the survivors' slots. + unsigned ballot = __ballot_sync(0xffffffffu, alive); + long long word = (base + s0 + (threadIdx.x & ~31)) >> 5; + if (survivors != nullptr && (threadIdx.x & 31) == 0 && s0 + (threadIdx.x & ~31) < ns) + survivors[word] = ballot; + if (ballot) {{ + int lane = threadIdx.x & 31, leader = __ffs(ballot) - 1; + unsigned b = 0; + if (lane == leader) b = atomicAdd(count, __popc(ballot)); + b = __shfl_sync(0xffffffffu, b, leader); + if (alive) queue[b + __popc(ballot & ((1u << lane) - 1))] = (unsigned)(base + s); + }} + }} + __syncthreads(); + // NOTE: only the shots that pass the screen are read again. + u64* rows = out + base * {NW}; + for (int i = threadIdx.x; i < ns * {NW}; i += blockDim.x) {{ + int s = i / {NW}, word = i - s * {NW}; + if (passed[s]) rows[i] = sw[word * {S} + s]; + }} +}} + +// The deferred channels' faults on the ``n`` kept shots ``index``, then their +// returned outputs bit-packed (``pack_rows``). The deferred channels flip output +// words only, so a block loads just those. +extern "C" __global__ void __launch_bounds__(256) noise_kept( + const u64* __restrict__ out, long long n, const unsigned* __restrict__ index, u64 seed, + int num_segments, int max_events, + const int* __restrict__ seg_first, const int* __restrict__ seg_outcomes, + const int* __restrict__ seg_begin, const int* __restrict__ seg_end, + const float* __restrict__ seg_log_keep, const Record* __restrict__ records, + const Flip* __restrict__ overflow, + unsigned char* __restrict__ dst, const uint2* __restrict__ pieces, + const unsigned* __restrict__ piece_start) +{{ + __shared__ u64 sw[{OW} * {S}]; + extern __shared__ __align__(16) uint2 events[]; + __shared__ int num_events; + long long base = (long long)blockIdx.x * {S}; + int ns = (int)min((long long){S}, n - base); + for (int i = threadIdx.x; i < ns * {OW}; i += blockDim.x) {{ + int s = i / {OW}, word = i - s * {OW}; + sw[word * {S} + s] = out[(long long)index[base + s] * {NW} + word]; + }} + if (threadIdx.x == 0) num_events = 0; + __syncthreads(); + add_faults(sw, events, &num_events, max_events, ns, seed ^ ((u64)blockIdx.x * 0xD1B54A32D192ED03ULL), + num_segments, seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep, records, + overflow); + // NOTE: the events are applied, so their buffer stages the packed rows. + pack_rows(sw, ns, pieces, piece_start, (unsigned char*)events, dst + base * {NB}); +}} +""" + + +def noise_tables( + channels: ChannelSampler, flips: WordTables, num_words: int +) -> NoiseTables: + """Return the noise kernels' tables, for CUDA's shared memory. + + Raises: + ValueError: If the rows do not fit in ``SHARED_BYTES``, or the noise + exceeds the kernels' limits. + + """ + return _noise_tables(channels, flips, num_words, shared_bytes=SHARED_BYTES) + + +def source(tables: NoiseTables, output_words: int, packed_bytes: int) -> str: + """Return the CUDA source of the ``noise`` and ``noise_kept`` kernels. + + ``noise_kept`` packs rows of ``packed_bytes`` bytes. + """ + return SOURCE.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/cuda/runtime.py b/src/tsim/statevector/cuda/runtime.py new file mode 100644 index 00000000..185fd03c --- /dev/null +++ b/src/tsim/statevector/cuda/runtime.py @@ -0,0 +1,100 @@ +"""CuPy and NVRTC for the CUDA executor, loaded on first use. + +CuPy compiles the kernels at runtime with NVRTC. A pip-installed CUDA 13 NVRTC +(``nvidia-cuda-nvrtc``) is preloaded when present: CuPy may otherwise pick up +an older system NVRTC that cannot target the GPU (CUDA 12.0 cannot target +``sm_120``, e.g. the RTX 5090). +""" + +from __future__ import annotations + +import ctypes +import functools +import importlib.util +import os +import pathlib +from types import ModuleType + + +@functools.cache +def cupy() -> ModuleType: + """Return the ``cupy`` module, with a pip CUDA 13 NVRTC preloaded if installed. + + Raises: + ImportError: If CuPy is not installed. + + """ + if importlib.util.find_spec("cupy") is None: + raise ImportError("CuPy is not installed") + spec = importlib.util.find_spec("nvidia") + if spec is not None and spec.submodule_search_locations: + for root in spec.submodule_search_locations: + library = pathlib.Path(root) / "cu13" / "lib" / "libnvrtc.so.13" + if library.exists(): + ctypes.CDLL(str(library), mode=ctypes.RTLD_GLOBAL) + break + import cupy as cp + + return cp + + +_forked_after_init = False + + +def _after_fork_in_child() -> None: + # NOTE: a child forked after its parent initialized CUDA cannot use CUDA; the + # cached answer is read without touching CUDA. + global _forked_after_init + if _device_available.cache_info().currsize and _device_available(): + _forked_after_init = True + + +def available() -> bool: + """Return whether CuPy is installed and sees a CUDA device this process can use. + + ``CUDA_VISIBLE_DEVICES=""`` hides the devices, so that the samplers run on the + CPU, as they do in a process forked after its parent used the GPU. + """ + return not _forked_after_init and _device_available() + + +def forked_after_init() -> bool: + """Return whether this process was forked after its parent initialized CUDA.""" + return _forked_after_init + + +@functools.cache +def _device_available() -> bool: + try: + cp = cupy() + return cp.cuda.runtime.getDeviceCount() > 0 + except Exception: + return False + + +os.register_at_fork(after_in_child=_after_fork_in_child) + + +@functools.cache +def arch() -> str: + """Return the NVRTC architecture option of the current device, e.g. ``-arch=sm_120``.""" + return "-arch=sm_" + cupy().cuda.Device().compute_capability + + +@functools.cache +def multiprocessors() -> int: + """Return the number of streaming multiprocessors of the current device.""" + cuda = cupy().cuda + return int( + cuda.runtime.getDeviceProperties(cuda.Device().id)["multiProcessorCount"] + ) + + +def out_of_memory(error: BaseException) -> bool: + """Return whether ``error`` is CuPy's or CUDA's out-of-memory error.""" + cuda = cupy().cuda + if isinstance(error, cuda.memory.OutOfMemoryError): + return True + # NOTE: CUDA_ERROR_OUT_OF_MEMORY and cudaErrorMemoryAllocation are both 2. + errors = (cuda.driver.CUDADriverError, cuda.runtime.CUDARuntimeError) + return isinstance(error, errors) and getattr(error, "status", None) == 2 diff --git a/src/tsim/statevector/cuda/sampler.py b/src/tsim/statevector/cuda/sampler.py new file mode 100644 index 00000000..dc058d13 --- /dev/null +++ b/src/tsim/statevector/cuda/sampler.py @@ -0,0 +1,1005 @@ +"""``CudaShotSampler``: ``ShotSampler`` with the noise and the tapes on a CUDA GPU. + +A batch of shots runs as these kernels on one stream: + + noise draws every shot's noise into its words and screens it (``noise``); + the screened survivors go to a queue + tape runs every component on the queued shots, a group of ``2^tb`` threads + per shot with its amplitudes in registers (``codegen``, ``gpu.tape``); with a + postselection mask it clears the rejected shots from the screen's + bitmap of survivors, and with + expectations it writes each observable's expectation + kept under postselection, draws the noise that cannot change whether a + shot is kept on the accepted shots only (``noise_kept``), when that + noise is at least ``DEFER_SHARE`` of the faults + list lists the accepted shots in order from the bitmap (``count_bits``, + ``list_bits``) + select gathers the rows (or the accepted rows) for the copy to the host + +``sample_accepted`` runs a call as chunks of shots on two streams, so that a +chunk's copy to the host overlaps the next chunk's kernels. + +The kernels are compiled with NVRTC once per postselection mask (CuPy caches +them on disk). 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. +""" + +from __future__ import annotations + +import math +import warnings +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from tsim.noise.channels import ChannelSampler +from tsim.statevector.cuda import runtime +from tsim.statevector.cuda.codegen import PRELUDE, hex64, tape_function +from tsim.statevector.cuda.noise import noise_tables, source as noise_source +from tsim.statevector.gpu.noise import NoiseGroups, deferred_channels, pass_rate +from tsim.statevector.gpu.rows import pieces, select +from tsim.statevector.shared import pack_bits +from tsim.statevector.shots import ShotLayout, ShotPlan, ShotSampler + +MAX_WIDTH = 11 +"""Widest component the CUDA executor runs: 32 threads of 64 amplitudes.""" + +BATCH_SHOTS = 1 << 22 +"""Shots of one batch on the device.""" + +MIN_BATCH_SHOTS = 1 << 14 +"""Smallest batch that a call out of device memory retries with, before the CPU.""" + +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.""" + +CHUNK_BYTES = 6 << 20 +"""Bytes of kept rows per chunk that a call's chunks aim at: each chunk's copy to +the host overlaps the next chunk's kernels on the other stream.""" + +MAX_CHUNKS = 8 +"""Most chunks of a batch, so that a chunk's fixed costs stay small.""" + +_KERNELS = r""" +extern "C" __global__ void __launch_bounds__({block}) tape( + u64* __restrict__ out, long long first_shot, u64 seed, + const unsigned* __restrict__ queue, const unsigned* __restrict__ count, + unsigned* __restrict__ survivors, unsigned* __restrict__ acc_count, + float* __restrict__ expectations, unsigned* __restrict__ deviation) +{{ + int t = threadIdx.x & {T1}; + unsigned gmask = {gmask}; +{seeds} + long long groups = ((long long)gridDim.x * blockDim.x) >> {tb}; + long long n = *count; + // NOTE: whole warps leave the loop together, so the shuffles keep their groups. + long long rounds = (n + groups - 1) / groups; + long long g0 = ((long long)blockIdx.x * blockDim.x + threadIdx.x) >> {tb}; + for (long long r = 0; r < rounds; ++r) {{ + long long q = g0 + r * groups; + bool valid = q < n; + u64 row_index = queue[valid ? q : n - 1]; + u64 shot = first_shot + row_index; + u64* row = out + row_index * {NW}; + // NOTE: the words that no component reads are dead loads, dropped. + u64 w[{NW}]; + #pragma unroll + for (int i = 0; i < {NW}; ++i) w[i] = row[i]; + float ev[{NE}]; + int grave = 0x7fffffff; +{calls} + if (valid && t == 0) {{ +{stores} + if (expectations != nullptr) {{ + #pragma unroll + for (int i = 0; i < {NO}; ++i) expectations[row_index * {NO} + i] = ev[i]; + }} + }} + if (survivors != nullptr) {{ + // NOTE: the postselected direct outputs passed the screen already; + // the shots that fail now leave the screen's bitmap. + u64 hit = {hits}; + if (valid && t == 0 && hit) atomicAnd(&survivors[row_index >> 5], ~(1u << (row_index & 31))); + unsigned ballot = __ballot_sync(0xffffffffu, valid && t == 0 && hit == 0); + if (ballot && (threadIdx.x & 31) == __ffs(ballot) - 1) atomicAdd(acc_count, __popc(ballot)); + }} + }} +}} + +// The shots of the ``words`` words of ``bits`` set, in order, as indices into +// ``index``: ``count_bits`` counts each block's 1024 words, ``list_bits`` lists +// a block's shots after those of the blocks before it. +__device__ __forceinline__ unsigned block_scan(unsigned v, unsigned* total) {{ + __shared__ unsigned warps[32]; + int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + unsigned x = v; + #pragma unroll + for (int o = 1; o < 32; o <<= 1) {{ unsigned y = __shfl_up_sync(0xffffffffu, x, o); if (lane >= o) x += y; }} + if (lane == 31) warps[warp] = x; + __syncthreads(); + if (warp == 0) {{ + unsigned w = lane < (blockDim.x >> 5) ? warps[lane] : 0; + #pragma unroll + for (int o = 1; o < 32; o <<= 1) {{ unsigned y = __shfl_up_sync(0xffffffffu, w, o); if (lane >= o) w += y; }} + warps[lane] = w; + }} + __syncthreads(); + unsigned before = (warp ? warps[warp - 1] : 0) + x - v; + *total = warps[31]; + __syncthreads(); + return before; +}} + +extern "C" __global__ void __launch_bounds__(1024) count_bits( + const unsigned* __restrict__ bits, long long words, unsigned* __restrict__ totals) {{ + long long i = (long long)blockIdx.x * 1024 + threadIdx.x; + unsigned total; + block_scan(i < words ? __popc(bits[i]) : 0, &total); + if (threadIdx.x == 0) totals[blockIdx.x] = total; +}} + +extern "C" __global__ void __launch_bounds__(1024) list_bits( + const unsigned* __restrict__ bits, long long words, const unsigned* __restrict__ totals, + unsigned* __restrict__ index) {{ + unsigned earlier = 0, total; + for (int b = threadIdx.x; b < blockIdx.x; b += blockDim.x) earlier += totals[b]; + block_scan(earlier, &earlier); + long long i = (long long)blockIdx.x * 1024 + threadIdx.x; + unsigned w = i < words ? bits[i] : 0; + unsigned at = earlier + block_scan(__popc(w), &total); + for (; w; w &= w - 1) index[at++] = (unsigned)(32 * i + __ffs(w) - 1); +}} + +// The returned outputs of rows ``index[j]`` (or every row), bit-packed densely +// (``pack_rows``): a block packs {S} rows' output words. +extern "C" __global__ void __launch_bounds__(256) gather_columns( + const u64* __restrict__ rows, long long n, const unsigned* __restrict__ index, + const uint2* __restrict__ pieces, const unsigned* __restrict__ start, + unsigned char* __restrict__ dst) {{ + __shared__ u64 sw[{OW} * {S}]; + extern __shared__ __align__(16) unsigned char stage[]; + long long base = (long long)blockIdx.x * {S}; + int ns = (int)min((long long){S}, n - base); + for (int i = threadIdx.x; i < ns * {OW}; i += blockDim.x) {{ + int s = i / {OW}, w = i - s * {OW}; + long long r = index != nullptr ? (long long)index[base + s] : base + s; + sw[w * {S} + s] = rows[r * {NW} + w]; + }} + __syncthreads(); + pack_rows(sw, ns, pieces, start, stage, dst + base * {NB}); +}} + +// The batch's counts, stored into pinned host memory: a read that does not +// queue behind the copies to the host. +extern "C" __global__ void publish(const unsigned* __restrict__ counts, volatile unsigned* host) {{ + if (threadIdx.x < 3) host[threadIdx.x] = counts[threadIdx.x]; + __threadfence_system(); +}} + +// Observable expectations of rows ``index[j]``, densely. +extern "C" __global__ void gather_float(const float* __restrict__ values, long long n, int width, + const unsigned* __restrict__ index, float* __restrict__ dst) {{ + for (long long k = (long long)blockIdx.x * blockDim.x + threadIdx.x; k < n * width; + k += (long long)gridDim.x * blockDim.x) {{ + long long j = k / width; int c = (int)(k - j * width); + dst[k] = values[(long long)index[j] * width + c]; + }} +}} +""" + + +@dataclass +class _Launch: + """The device tables of one noise launch: its segments and records.""" + + args: list[Any] + num_segments: int + max_events: int + + +@dataclass +class _Module: + """A plan's compiled kernels and its device constants.""" + + noise: Any + noise_kept: Any + tape: Any + gather_columns: Any + gather_float: Any + publish: Any + count_bits: Any + list_bits: Any + screen: Any + tape_blocks: int + has_mask: bool + columns: np.ndarray + """The outputs that ``sample_accepted`` returns: all of them, in order.""" + pieces: Any + piece_start: Any + plan: ShotPlan + psm_words: np.ndarray + passing: float + """Estimated share of shots that pass the screen.""" + splits: dict[bool, tuple[_Launch, _Launch] | None] + """Each ``expectations`` value's noise passes (screened, deferred), or None.""" + + +def _pinned(shape: tuple[int, ...], dtype) -> np.ndarray: + """Return an uninitialized array in pinned host memory.""" + import cupyx + + # NOTE: CuPy's stub types the dtype as a float type. + return cupyx.empty_pinned(shape, dtype=dtype) # pyright: ignore[reportArgumentType] + + +def _pinned_counts(cp) -> tuple[np.ndarray, Any]: + """Return three pinned host counters and a device array over the same memory.""" + host = _pinned((3,), np.uint32) + # NOTE: with unified addressing a pinned host pointer is also a device pointer. + memory = cp.cuda.UnownedMemory(host.ctypes.data, host.nbytes, host) + device = cp.ndarray((3,), np.uint32, cp.cuda.MemoryPointer(memory, 0)) + return host, device + + +class _HostRows: + """A call's passing rows (and expectations) in pinned host memory, chunk by chunk.""" + + def __init__(self, nb: int, no: int | None) -> None: + self.nb, self.no = nb, no + self.capacity, self.total = 0, 0 + self.rows: np.ndarray | None = None + self.values: np.ndarray | None = None + + def allocate(self, passed: int, of: int, shots: int) -> None: + """Reserve the rows of a call of ``shots``, ``passed`` of its first ``of`` passing.""" + rest = shots - of + share = passed / of + expected = rest * share + self._resize( + min(shots, passed + int(1.05 * expected + 6 * np.sqrt(expected) + 256)) + ) + + def grow(self, capacity: int) -> None: + """Make room for ``capacity`` rows, keeping the copied ones (no copy may be pending).""" + rows, values, total = self.rows, self.values, self.total + self._resize(capacity) + assert self.rows is not None and rows is not None + self.rows[:total] = rows[:total] + if values is not None and self.values is not None: + self.values[:total] = values[:total] + + def _resize(self, capacity: int) -> None: + self.capacity = capacity + self.rows = _pinned((capacity, self.nb), np.uint8) + if self.no is not None: + self.values = _pinned((capacity, self.no), np.float32) + + def take(self, count: int) -> tuple[np.ndarray, np.ndarray | None]: + """Return the next ``count`` rows (and expectations) to fill.""" + assert self.rows is not None and self.total + count <= self.capacity + lo, self.total = self.total, self.total + count + values = None if self.values is None else self.values[lo : self.total] + return self.rows[lo : self.total], values + + def result(self) -> tuple[np.ndarray, np.ndarray | None]: + """Return the filled rows and expectations.""" + if self.rows is None: + self._resize(0) + assert self.rows is not None + values = None if self.values is None else self.values[: self.total] + return self.rows[: self.total], values + + +class CudaShotSampler: + """Samples shot words of a statevector program on a CUDA GPU. + + It keeps ``ShotSampler``'s interface (``sample``, ``direct_bools``, + ``reference``, ``plan``, ``layout``); ``sample`` copies only the shots that + pass postselection off the GPU. ``sample_accepted`` can also return each + observable's expectation. + """ + + device = "cuda" + + @staticmethod + def supports(layout: ShotLayout) -> bool: + """Return whether every component of ``layout`` fits in ``MAX_WIDTH`` qubits.""" + return 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 CuPy is not installed. + + """ + self._cp = runtime.cupy() + cp = self._cp + 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) + n = self._noise + self._record_args = [self._upload(n.overflow.view(np.uint8))] + self._all = self._launch(n.groups) + self._modules: dict[bytes, _Module] = {} + self._buffers: dict[str, Any] = {} + self._views: dict[tuple[str, tuple[int, ...]], Any] = {} + self._constants: dict[bytes, Any] = {} + # NOTE: each buffer slot has a kernel stream and a copy stream, so that a + # slot's next kernels need not wait for its copy to the host. + # NOTE: a slot's noise kernel runs on a stream of the least priority, so the + # other slot's tape and packing, which the next copy waits for, go first. + self._streams = [ + cp.cuda.Stream(non_blocking=True, priority=-1) for _ in range(2) + ] + self._noise_streams = [cp.cuda.Stream(non_blocking=True) for _ in range(2)] + self._copies = [cp.cuda.Stream(non_blocking=True) for _ in range(2)] + self._published = [_pinned_counts(cp) for _ in range(2)] + # NOTE: lowered when a call runs out of device memory. + self._max_batch = BATCH_SHOTS + + @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 32 amplitudes, at most 32.""" + width = max((c.width for c in self.layout.components), default=0) + return min(5, max(0, width - 5)) + + 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] + cp = self._cp + plan = self.plan(mask) + layout = self.layout + tb = self._threads_per_shot() + T = 1 << tb + block = 128 + 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, gmask, ev, deviation));" + for i in range(len(plan.programs)) + ) + gmask = ( + "0xffffffffu" + if T == 32 + else f"((1u << {T}) - 1) << (threadIdx.x & 31 & ~{T - 1})" + ) + 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) + # 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}) + 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] + ) + source = ( + PRELUDE + + noise_source(self._noise, layout.output_words, -(-len(columns) // 8)) + + functions + + _KERNELS.format( + block=block, + T1=T - 1, + tb=tb, + gmask=gmask, + 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(-(-len(columns) // 8), 1), + calls=calls, + # NOTE: each component's stream, from the call's seed. + seeds="\n".join( + f" const u64 seed_{i} = mix64(seed + {i + 1} * GOLDEN);" + for i in range(len(plan.programs)) + ), + stores=stores, + hits=hits or "0", + ) + ) + module = self._compile(source) + tape = module.get_function("tape") + blocks_per_sm = 4 + if hasattr(cp.cuda.runtime, "occupancyMaxActiveBlocksPerMultiprocessor"): + blocks_per_sm = max( + 1, + cp.cuda.runtime.occupancyMaxActiveBlocksPerMultiprocessor( + tape.ptr, block, 0 + ), + ) + found, piece_start = pieces(columns) + compiled = _Module( + noise=module.get_function("noise"), + noise_kept=module.get_function("noise_kept"), + tape=tape, + gather_columns=module.get_function("gather_columns"), + gather_float=module.get_function("gather_float"), + publish=module.get_function("publish"), + count_bits=module.get_function("count_bits"), + list_bits=module.get_function("list_bits"), + screen=self._upload(plan.screen), + tape_blocks=runtime.multiprocessors() * blocks_per_sm, + has_mask=mask is not None, + columns=columns, + pieces=self._upload(found), + piece_start=self._upload(piece_start), + plan=plan, + psm_words=psm_words, + passing=self._passing(plan), + splits={}, + ) + self._modules[key] = compiled + return compiled + + def _upload(self, array: np.ndarray) -> Any: + """Return ``array`` on the device, copied before any later kernel runs. + + ``cp.asarray`` copies on the current stream, usually the legacy default + stream, which the sampler's non-blocking streams do not wait for: a kernel + could read the array before the copy lands. Uploads are one-time, cached. + """ + device = self._cp.asarray(array) + self._cp.cuda.get_current_stream().synchronize() + return device + + def _passing(self, plan: ShotPlan) -> float: + """Return the estimated share of shots that pass the screen of ``plan``.""" + layout = self.layout + screen = np.zeros(layout.num_words, dtype=np.uint64) + screen[: layout.output_words] = plan.screen + return pass_rate(self._channels, layout.noise, screen) if screen.any() else 1.0 + + def _launch(self, groups: NoiseGroups) -> _Launch: + """Upload a noise launch's segments and records.""" + return _Launch( + args=[ + *(self._upload(a) for a in groups.segments), + self._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. Each key word is read + whole, so the deferred channels flip output words only. + """ + if not module.has_mask: + 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 _compile(self, source: str): + """Compile with NVRTC; ptxas -O3 fails on some shuffled tapes, so lower levels follow.""" + cp = self._cp + for level in ("3", "2", "1"): + options = ( + runtime.arch(), + "-std=c++17", + "-diag-suppress=177", + f"-Xptxas=-O{level}", + ) + try: + module = cp.RawModule(code=source, options=options) + module.get_function("tape") + return module + except cp.cuda.compiler.CompileException: + if level == "1": + raise + raise AssertionError("unreachable") + + # -- execution -------------------------------------------------------------- + + def _buffer(self, name: str, shape: tuple[int, ...], dtype) -> Any: + """Return a device buffer of at least ``shape``, reused across calls.""" + view = self._views.get((name, shape)) + if view is not None and view.dtype == np.dtype(dtype): + return view + size = int(np.prod(shape)) + held = self._buffers.get(name) + if held is None or held.size < size or held.dtype != np.dtype(dtype): + held = self._cp.empty(size, dtype=dtype) + self._buffers[name] = held + # NOTE: views of the replaced buffer must not be handed out again. + self._views = {k: v for k, v in self._views.items() if k[0] != name} + view = held[:size].reshape(shape) + self._views[(name, shape)] = view + return view + + def _constant(self, flip: np.ndarray | None) -> Any: + """Return the outputs' constants XOR the reference bits ``flip``, on the device.""" + key = b"" if flip is None else np.asarray(flip, bool).tobytes() + if key not in self._constants: + self._constants[key] = self._upload(self._host._constant(flip)) + return self._constants[key] + + def _run_batch( + self, + module: _Module, + shots: int, + first_shot: int, + constant: Any, + expectations: bool, + launch: _Launch, + slot: int = 0, + noise_stream: Any = None, + ) -> tuple[Any, Any, Any, Any]: + """Launch one batch on the current stream: ``(words, accepted, counts, values)``. + + ``accepted`` is the bitmap of the accepted shots (with a mask). + + ``launch`` holds the channels that the noise kernel draws; only the rows + of the shots that pass the screen are written. + ``counts`` holds the screened survivors, the accepted shots and the norm + deviation's bits. ``slot`` names the device buffers; ``noise_stream``, if + given, runs the noise kernel, after the current stream's work. + """ + cp = self._cp + layout = self.layout + nw, no = layout.num_words, max(1, len(self._observables)) + words = self._buffer(f"words{slot}", (shots, nw), np.uint64) + queue = self._buffer(f"queue{slot}", (shots,), np.uint32) + accepted = self._buffer(f"accepted{slot}", (-(-shots // 32),), np.uint32) + counts = self._buffer(f"counts{slot}", (3,), np.uint32) + values = ( + self._buffer(f"values{slot}", (shots, no), np.float32) + if expectations + else None + ) + current = cp.cuda.get_current_stream() + if noise_stream is not None: + noise_stream.wait_event(current.record()) + noise_stream.use() + counts.fill(0) + 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 + module.noise( + ((shots + S - 1) // S,), + (256,), + ( + words, + np.int64(shots), + noise_seed, + np.int32(launch.num_segments), + np.int32(launch.max_events), + *launch.args, + *self._record_args, + constant, + module.screen, + queue, + counts[0:1], + accepted if module.has_mask else np.uint64(0), + ), + shared_mem=8 * launch.max_events, + ) + if noise_stream is not None: + current.wait_event(noise_stream.record()) + current.use() + if layout.components: + module.tape( + (module.tape_blocks,), + (128,), + ( + words, + np.int64(first_shot), + seed, + queue, + counts[0:1], + accepted if module.has_mask else np.uint64(0), + counts[1:2], + values if values is not None else np.uint64(0), + counts[2:3], + ), + ) + elif module.has_mask: + # NOTE: without components the screened survivors are the accepted shots. + counts[1:2] = counts[0:1] + return words, accepted, counts, values + + def _pack( + self, + module: _Module, + split: tuple[_Launch, _Launch] | None, + words, + count: int, + index, + ) -> Any: + """Return the ``count`` rows ``index`` (or all) on the device, as copied to the host. + + Their outputs ``module.columns``, bit-packed; with ``split`` the deferred + noise kernel adds the deferred channels' faults first. + """ + S = self._noise.block_shots + m = len(module.columns) + nb = -(-m // 8) + # NOTE: a new array per chunk, so that refilling it never waits for a copy. + dst = self._cp.empty((max(count, 1), nb), dtype=np.uint8) + if not count or not nb: + return dst + blocks = (count + S - 1) // S + if split is not None: + launch = split[1] + seed = np.uint64(self._rng.integers(0, 2**64, dtype=np.uint64)) + module.noise_kept( + (blocks,), + (256,), + ( + words, + np.int64(count), + index, + seed, + np.int32(launch.num_segments), + np.int32(launch.max_events), + *launch.args, + *self._record_args, + dst, + module.pieces, + module.piece_start, + ), + shared_mem=max(8 * launch.max_events, S * nb), + ) + else: + module.gather_columns( + (blocks,), + (256,), + ( + words, + np.int64(count), + index if index is not None else np.uint64(0), + module.pieces, + module.piece_start, + dst, + ), + shared_mem=S * nb, + ) + return dst + + def _list(self, module: _Module, bits, shots: int, count: int, slot: int) -> Any: + """Return the ``count`` shots of the bitmap ``bits`` of ``shots`` shots, in order.""" + words = -(-shots // 32) + blocks = -(-words // 1024) + totals = self._buffer(f"totals{slot}", (blocks,), np.uint32) + index = self._buffer(f"index{slot}", (max(count, 1),), np.uint32) + module.count_bits((blocks,), (1024,), (bits, np.int64(words), totals)) + module.list_bits((blocks,), (1024,), (bits, np.int64(words), totals, index)) + return index[:count] + + def _deviation(self, counts: np.ndarray) -> float: + """Return a batch's normalization deviation beyond float32 rounding, else 0.""" + bits = np.asarray(counts[2:3], dtype=np.uint32) + deviation = float(bits.view(np.float32)[0]) + return deviation if deviation > DEVIATION_TOLERANCE else 0.0 + + 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 device + batch. + """ + accepted = self._sample_or_none(shots, mask, flip, batch_size) + if accepted is None: + return self._host.sample( + shots, + columns, + bit_packed=bit_packed, + batch_size=batch_size, + mask=mask, + flip=flip, + ) + rows, deviation = accepted + num_outputs = self.layout.num_outputs + return [select(rows, c, num_outputs, bit_packed) for c in columns], deviation + + def _sample_or_none( + self, + shots: int, + mask: np.ndarray | None, + flip: np.ndarray | None, + batch_size: int | None, + ) -> tuple[np.ndarray, float] | None: + """Return ``sample_accepted``'s rows and deviation, or None if the GPU is full. + + Out of device memory, the call is retried in batches half as large, down + to ``MIN_BATCH_SHOTS``; past that it warns and returns None, for the CPU. + Later calls start from the batch that fit. + """ + batch = min(batch_size or BATCH_SHOTS, self._max_batch) + while True: + try: + rows, _, deviation = self.sample_accepted( + shots, mask=mask, flip=flip, batch_size=batch + ) + return rows, deviation + except Exception as error: + if not runtime.out_of_memory(error): + raise + self._release() + self._max_batch = max(batch // 2, MIN_BATCH_SHOTS) + if batch <= MIN_BATCH_SHOTS: + warnings.warn( + "The GPU ran out of memory, even for batches of " + f"{MIN_BATCH_SHOTS} shots; sampling on the CPU.", + stacklevel=5, + ) + return None + batch = self._max_batch + + def _release(self) -> None: + """Wait for the device and free the buffers that calls reuse.""" + cp = self._cp + cp.cuda.Device().synchronize() + self._buffers.clear() + self._views.clear() + cp.get_default_memory_pool().free_all_blocks() + + def _synchronize(self) -> None: + """Wait for every kernel and copy of ``sample_accepted``.""" + for stream in (*self._streams, *self._copies): + stream.synchronize() + + def _chunk_shots(self, shots: int, nb: int, passing: float) -> int: + """Return a chunk's shots: about ``CHUNK_BYTES`` of rows of ``nb`` bytes. + + ``passing`` is the share of shots kept, and a batch of ``shots`` shots + splits into at most ``MAX_CHUNKS`` chunks. + """ + chunks = min(MAX_CHUNKS, max(1, math.ceil(nb * passing * shots / CHUNK_BYTES))) + return -(-shots // chunks) + + def sample_accepted( + self, + shots: int, + *, + mask: np.ndarray | None = None, + flip: np.ndarray | None = None, + expectations: bool = False, + batch_size: int | None = None, + ) -> 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. + + The shots run in chunks on two streams: each chunk's rows go to the host + while the next chunk's kernels run. The chunks hold about ``CHUNK_BYTES`` + of rows, by the screen's estimated share of passing shots, so a seed's + samples do not depend on earlier calls. + + Returns: + ``(outputs, expectations, deviation)``: the passing shots' outputs, + bit-packed as ``np.packbits(..., bitorder="little")`` would (``uint8``, + in pinned host 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. + + """ + cp = self._cp + module = self._module(mask) + layout = self.layout + nb, no = -(-len(module.columns) // 8), len(self._observables) + if not shots: + values = np.zeros((0, no), np.float32) if expectations else None + return np.zeros((0, nb), np.uint8), values, 0.0 + constant = self._constant(flip) + split = self._split(module, expectations) + launch = split[0] if split else self._all + batch = min(batch_size or BATCH_SHOTS, BATCH_SHOTS) + out = _HostRows(nb, no if expectations else None) + state = {"start": 0, "slot": 0} + size = self._chunk_shots(min(batch, shots), nb, module.passing) + deviation = 0.0 + pending: list[tuple[int, int, int, tuple[Any, Any, Any, Any], Any]] = [] + # NOTE: the device rows in flight to the host, kept until the copies land, + # and the streams that copy them. + held: list[Any] = [] + used: set[Any] = set() + + def launch_next() -> None: + first, slot = state["start"], state["slot"] + n = min(size, shots - first) + device = self._published[slot][1] + ready = None + with self._streams[slot]: + launched = self._run_batch( + module, + n, + first, + constant, + expectations, + launch, + slot, + self._noise_streams[slot], + ) + # NOTE: no copy is in flight when the first chunk finishes, so it + # reads its counts with a copy of its own. + if first: + module.publish((1,), (32,), (launched[2], device)) + ready = self._streams[slot].record() + pending.append((slot, first, n, launched, ready)) + state["start"], state["slot"] = first + n, 1 - slot + + def finish() -> float: + slot, first, n, (words, accepted, counts, values), ready = pending.pop(0) + if ready is None: + host_counts = counts.get(stream=self._streams[slot]) + else: + ready.synchronize() + host_counts = self._published[slot][0].copy() + with self._streams[slot]: + count = n if mask is None else int(host_counts[1]) + if first == 0: + out.allocate(count, n, shots) + elif out.total + count > out.capacity: + # NOTE: a rare excess of passing shots: the copies in flight land first. + self._synchronize() + out.grow(out.total + count + (shots - first - n)) + rows_host, values_host = out.take(count) + index = ( + None + if mask is None + else self._list(module, accepted, n, count, slot) + ) + stream = self._streams[slot] + # NOTE: a call of one chunk has nothing to overlap its copy with. + copy = stream if first + n == shots and not held else self._copies[slot] + used.add(copy) + if count and nb: + dst = self._pack(module, split, words, count, index) + if copy is not stream: + copy.wait_event(stream.record()) + dst[:count].get(stream=copy, out=rows_host, blocking=False) + held.append(dst) + if values_host is not None and count: + if no and layout.components: + if index is None: + index = cp.arange(n, dtype=cp.uint32) + dst = cp.empty((count, no), dtype=np.float32) + module.gather_float( + (runtime.multiprocessors() * 8,), + (256,), + (values, np.int64(count), np.int32(no), index, dst), + ) + if copy is not stream: + copy.wait_event(stream.record()) + dst.get(stream=copy, out=values_host, blocking=False) + held.append(dst) + else: + values_host[:] = 0 + return self._deviation(host_counts) + + # NOTE: each chunk is launched before the previous one is finished, so its + # kernels overlap that copy. + launch_next() + while state["start"] < shots: + launch_next() + if len(pending) == 2: + deviation = max(deviation, finish()) + while pending: + deviation = max(deviation, finish()) + for stream in used: + stream.synchronize() + rows, values = out.result() + if values is None: + return rows, None, deviation + self._direct_expectations(rows, values, module.columns) + return rows, values, deviation + + def _direct_expectations( + self, rows: np.ndarray, values: np.ndarray, columns: np.ndarray + ) -> None: + """Set the expectations of observables outside the components: ``±1``. + + ``columns`` are the outputs that ``rows`` hold. + """ + 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.""" + accepted = self._sample_or_none(shots, None, flip, None) + if accepted is None: + return self._host.direct_bools(shots, flip) + rows, _ = accepted + num_outputs = self.layout.num_outputs + return select(rows, np.arange(num_outputs), num_outputs, False) diff --git a/test/conftest.py b/test/conftest.py new file mode 100644 index 00000000..83da5ce7 --- /dev/null +++ b/test/conftest.py @@ -0,0 +1,10 @@ +import pytest + +from tsim.statevector.cuda import runtime + + +@pytest.fixture(autouse=True) +def _cpu_samplers(request, monkeypatch): + """Sample on the CPU, except in the CUDA tests, which choose per sampler.""" + if request.path.name != "test_cuda_sampler.py": + monkeypatch.setattr(runtime, "available", lambda: False) diff --git a/test/integration/test_cuda_sampler.py b/test/integration/test_cuda_sampler.py new file mode 100644 index 00000000..c3d17e17 --- /dev/null +++ b/test/integration/test_cuda_sampler.py @@ -0,0 +1,705 @@ +"""The CUDA executor end to end, against the stabilizer-rank strategy and the CPU executor.""" + +import multiprocessing +import warnings +from test.helpers.gen import gen_stim_circuit +from test.helpers.statevector import CIRCUITS, assert_counts_match + +import numpy as np +import pytest + +from tsim.circuit import Circuit +from tsim.statevector.cuda import available, codegen, runtime, sampler as cuda_sampler +from tsim.statevector.cuda.sampler import CudaShotSampler +from tsim.statevector.gpu import tape +from tsim.statevector.program import StatevectorComponent +from tsim.statevector.shots import ShotSampler + +pytestmark = pytest.mark.skipif(not available(), reason="needs CuPy and a CUDA GPU") + +# NOTE: t_ghz plus a direct detector (detector 2), as in test_trajectory_sharing. +GHZ_WITH_DIRECT_DETECTOR = """ +R 0 1 2 3 +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 +M 3 +DETECTOR rec[-1] +""" + +# NOTE: a T state measured in X after a noisy idle: its observable's expectation +# is cos(π/4) = 0.707 before noise. +T_STATE_X = """ +R 0 +H 0 +T 0 +Z_ERROR(0.1) 0 +H 0 +M 0 +OBSERVABLE_INCLUDE(0) rec[-1] +""" + + +def _gpu(sampler): + """Return ``sampler``, which must run on the GPU.""" + _gpu_shots(sampler) + return sampler + + +def _gpu_shots(sampler) -> CudaShotSampler: + """Return the GPU shot sampler of ``sampler``.""" + assert isinstance(sampler._shots, CudaShotSampler) + 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_cuda_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_cuda_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_cuda_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, cuda_rate = len(expected) / shots, len(got) / shots + sigma = np.sqrt(rate * (1 - rate) * 2 / shots) + assert abs(cuda_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_cuda_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_cuda_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, cuda_rate = len(expected) / shots, len(got) / shots + assert abs(cuda_rate - rate) <= 5 * np.sqrt(rate * (1 - rate) * 2 / shots) + assert_counts_match(got, expected) + + +def test_cuda_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_cuda_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_cuda_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_cuda_postselection_drops_the_rows_that_fail(monkeypatch): + # NOTE: the masked run must draw the unmasked run's noise, so none is deferred. + monkeypatch.setattr(cuda_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_cuda_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_cuda_chunks_return_the_rows_of_one_pass(monkeypatch, text): + # NOTE: 1000-shot chunks on both streams, and host rows that start with room + # for the first chunk only; without deferral, a masked and an unmasked run on + # the same chunks draw the same shots. + monkeypatch.setattr(cuda_sampler, "DEFER_SHARE", 2.0) + monkeypatch.setattr( + cuda_sampler.CudaShotSampler, "_chunk_shots", lambda self, *args: 1000 + ) + monkeypatch.setattr( + cuda_sampler._HostRows, + "allocate", + lambda self, passed, of, shots: self._resize(passed), + ) + circuit = Circuit(text) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + outputs = circuit.num_detectors + circuit.num_observables + full = _packed(_gpu(circuit.compile_detector_sampler(seed=16)), 10_500) + rows = _packed(_gpu(circuit.compile_detector_sampler(seed=16)), 10_500, 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_cuda_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.""" + 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_cuda_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_cuda_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_cuda_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_cuda_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) > cuda_sampler.MAX_WIDTH + with pytest.warns(UserWarning, match="sampling on the CPU"): + sampler = circuit.compile_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + + +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: CUDA cannot be used in a child forked after its parent initialized 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 ~90 faults instead of ~0.4. + 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(64) + ) + 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. + cp = runtime.cupy() + source = codegen.PRELUDE + r""" +extern "C" __global__ void zeros(float w0, unsigned long long n, unsigned long long* count) { + unsigned long long local = 0; + for (u64 i = blockIdx.x * (u64)blockDim.x + threadIdx.x; i < n; i += (u64)gridDim.x * blockDim.x) + local += !draw_one(mix64(i * GOLDEN + 12345), w0, 1.f, w0 + 1.f); + atomicAdd(count, local); +} +""" + kernel = cp.RawModule( + code=source, options=(runtime.arch(), "-std=c++17") + ).get_function("zeros") + count = cp.zeros(1, dtype=cp.uint64) + n = 1 << 30 + kernel((1024,), (256,), (np.float32(p0), np.uint64(n), count)) + expected = p0 * n + assert abs(int(count.get()[0]) - expected) <= 5 * np.sqrt(expected) + 1 + + +@pytest.mark.parametrize(("probability", "shots"), [(1e-10, 1 << 28), (1e-6, 1 << 26)]) +def test_cuda_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 _out_of_memory_above(monkeypatch, sampler, limit): + """Make ``sampler``'s GPU calls of batches above ``limit`` run out of memory.""" + shots = _gpu_shots(sampler) + sample_accepted = shots.sample_accepted + batches = [] + + def fake(n, *, batch_size=None, **kwargs): + batches.append(batch_size) + if batch_size is None or batch_size > limit: + cp = runtime.cupy() + raise cp.cuda.memory.OutOfMemoryError(1 << 40, 0) + return sample_accepted(n, batch_size=batch_size, **kwargs) + + monkeypatch.setattr(shots, "sample_accepted", fake) + return batches + + +def test_calls_out_of_gpu_memory_retry_in_smaller_batches(monkeypatch): + sampler = _gpu(Circuit(GHZ_WITH_DIRECT_DETECTOR).compile_detector_sampler(seed=1)) + batches = _out_of_memory_above(monkeypatch, sampler, 1 << 20) + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert sampler.sample(1000).shape == (1000, 3) + assert batches == [1 << 22, 1 << 21, 1 << 20] + # NOTE: later calls start from the batch that fit. + sampler.sample(1000) + assert batches[-1] == 1 << 20 + + +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)) + _out_of_memory_above(monkeypatch, sampler, 0) + with pytest.warns(UserWarning, match="out of memory"): + 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_kernels_wait_for_the_tables_they_read(): + # NOTE: the tables are copied on the current (legacy default) stream, which + # the sampler's non-blocking streams do not wait for; a kernel that keeps + # that stream busy delays the copies, so the first call would read them + # before they land. + cp = runtime.cupy() + spin = cp.RawKernel( + r""" + extern "C" __global__ void spin(long long cycles) { + long long start = clock64(); + while (clock64() - start < cycles) {} + } + """, + "spin", + ) + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + expected = _cpu(circuit.compile_detector_sampler, seed=1).sample( + 40_000, append_observables=True + ) + for seed in range(8): + spin((1,), (1,), (np.int64(200_000_000),)) + sampler = _gpu(circuit.compile_detector_sampler(seed=seed)) + got = sampler.sample(40_000, append_observables=True) + assert_counts_match(got, expected) + + +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)) + + +def test_the_repr_says_where_a_sampler_samples(): + circuit = Circuit(T_STATE_X) + gpu = _gpu(circuit.compile_detector_sampler(seed=1)) + assert "samples on cuda" in repr(gpu) and gpu.stats.device == "cuda" + cpu = _cpu(circuit.compile_detector_sampler, seed=1) + assert "samples on cpu" in repr(cpu) and cpu.stats.device == "cpu" + + +def test_programs_without_shot_words_sample_on_the_cpu(): + # NOTE: a detector sampler of a circuit without detectors, observables or noise. + with warnings.catch_warnings(): + warnings.simplefilter("error") + sampler = Circuit("R 0\nH 0\nT 0\nM 0").compile_detector_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + assert sampler.sample(10).shape == (10, 0) + + +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) diff --git a/uv.lock b/uv.lock index a741f3e7..6db4bb59 100644 --- a/uv.lock +++ b/uv.lock @@ -136,12 +136,21 @@ dependencies = [ ] [package.optional-dependencies] +cuda = [ + { name = "cupy-cuda13x", extra = ["ctk"] }, +] cuda12 = [ { name = "jax", extra = ["cuda12"] }, ] cuda13 = [ { name = "jax", extra = ["cuda13"] }, ] +jax-cuda12 = [ + { name = "jax", extra = ["cuda12"] }, +] +jax-cuda13 = [ + { name = "jax", extra = ["cuda13"] }, +] [package.dev-dependencies] dev = [ @@ -182,11 +191,14 @@ doc = [ [package.metadata] requires-dist = [ + { name = "cupy-cuda13x", extras = ["ctk"], marker = "extra == 'cuda'", specifier = ">=14.0" }, { name = "equinox", specifier = ">=0.11.5" }, { name = "ipywidgets", specifier = ">=8.0.0" }, { name = "jax", specifier = ">=0.4.38" }, { name = "jax", extras = ["cuda12"], marker = "extra == 'cuda12'", specifier = ">=0.4.38" }, + { name = "jax", extras = ["cuda12"], marker = "extra == 'jax-cuda12'", specifier = ">=0.4.38" }, { name = "jax", extras = ["cuda13"], marker = "extra == 'cuda13'", specifier = ">=0.7.0" }, + { name = "jax", extras = ["cuda13"], marker = "extra == 'jax-cuda13'", specifier = ">=0.7.0" }, { name = "lxml", specifier = ">=5.0.0" }, { name = "numba", specifier = ">=0.61" }, { name = "numpy", specifier = ">=2.0" }, @@ -195,7 +207,7 @@ requires-dist = [ { name = "stim-timeline-text", specifier = ">=0.1.0" }, { name = "zxparam", specifier = ">=0.1.0" }, ] -provides-extras = ["cuda12", "cuda13"] +provides-extras = ["cuda", "jax-cuda12", "jax-cuda13", "cuda12", "cuda13"] [package.metadata.requires-dev] dev = [ @@ -535,6 +547,87 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9e/ee/a4cf96b8ce1e566ed238f0659ac2d3f007ed1d14b181bcb684e19561a69a/coverage-7.13.5-py3-none-any.whl", hash = "sha256:34b02417cf070e173989b3db962f7ed56d2f644307b2cf9d5a0f258e13084a61", size = 211346, upload-time = "2026-03-17T10:33:15.691Z" }, ] +[[package]] +name = "cuda-pathfinder" +version = "1.8.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6a/9e/08555797c979322def080814686417ea93f66d38f5c48ebab6867a85a21e/cuda_pathfinder-1.8.6-py3-none-any.whl", hash = "sha256:bf6b3b06323cf95a11d671818392da5da23ddbbb4dfc4b77ca10622783c25980", size = 62603, upload-time = "2026-11-05T21:22:51.229Z" }, +] + +[[package]] +name = "cuda-toolkit" +version = "13.5.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e6/4a/7d75cd156b6651831473a60f5a4dfa96b5b6f1d04ccb47fcac1f8a08456e/cuda_toolkit-13.5.0-py2.py3-none-any.whl", hash = "sha256:d5b168b8284af2bad07f313b4d9fd5af0d1c7d6a26b7b019f735784ef718a737", size = 2706, upload-time = "2026-11-04T23:44:09.608Z" }, +] + +[package.optional-dependencies] +cublas = [ + { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cudart = [ + { name = "nvidia-cuda-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cufft = [ + { name = "nvidia-cufft", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +curand = [ + { name = "nvidia-curand", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cusolver = [ + { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-cusolver", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cusparse = [ + { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +nvrtc = [ + { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] + +[[package]] +name = "cupy-cuda13x" +version = "14.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder" }, + { name = "numpy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/3c/4716e285f8459a5707ef39d84e9a7547bf16f87c583ac4789c0763d6d986/cupy_cuda13x-14.2.0-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:a22e2f17c97940cd566f351f186d8fe04285280d052ca1512ff223febbe9979c", size = 74074375, upload-time = "2026-08-20T02:41:04.327Z" }, + { url = "https://files.pythonhosted.org/packages/ac/cb/41b6d8f1880e6f3128f8084b7622b1a3ca2d615ab77af9ef313b4aafed3d/cupy_cuda13x-14.2.0-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:f94843e54500a9c8c169135f79572bfaba946efe4f7178e834c227280c1b0c12", size = 69969199, upload-time = "2026-08-20T02:41:07.842Z" }, + { url = "https://files.pythonhosted.org/packages/f8/eb/25b8c1a415ade6479ed684c5cbff1e6f5002a92cbf26ef826df56a40cf69/cupy_cuda13x-14.2.0-cp311-cp311-win_amd64.whl", hash = "sha256:3893acd8d4dc78d113736045eaa306e496ee7bd8253bd959513ba28c91ed396d", size = 36082966, upload-time = "2026-08-20T02:41:10.723Z" }, + { url = "https://files.pythonhosted.org/packages/10/72/ad885cc21a21208b94044b34f50a840394b2e9dbb6111ed43959366aa6eb/cupy_cuda13x-14.2.0-cp311-cp311-win_arm64.whl", hash = "sha256:e0ad42a1ef35549282e2087a7c1e6ec8b7aa624a95e8fc0a60da777ca87e71de", size = 35136070, upload-time = "2026-09-01T06:41:38.284Z" }, + { url = "https://files.pythonhosted.org/packages/dd/c8/68e07ff959f6a9f93169e37a64a4c24223487102b8625191e4a566cab465/cupy_cuda13x-14.2.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:0a9b6538079c5151dcd99cf44b94cc3f1242315b7de0498b4619cef9b2a7c2c8", size = 73382974, upload-time = "2026-08-20T02:41:13.648Z" }, + { url = "https://files.pythonhosted.org/packages/7e/bf/c4c25e97552320451afdcd459d4b6e0b715459594e998b23051c06db64c3/cupy_cuda13x-14.2.0-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:9aafadbe29db3d4285026044aafeb28f2a3de4bd0df70718ec6892b7d27a4ded", size = 69439247, upload-time = "2026-08-20T02:41:17.375Z" }, + { url = "https://files.pythonhosted.org/packages/5e/88/0a4d2d075485c4a82bfcd4725300ff87d17de36dce27b36f720d286f0f70/cupy_cuda13x-14.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:9e8bc8a7693c5c212cd1450d728f47e1f05aa07ae003de30415893d0daf7d243", size = 35986921, upload-time = "2026-08-20T02:41:20.253Z" }, + { url = "https://files.pythonhosted.org/packages/a4/ab/8a756f553397b8509906d4a1d8083d92d3d21cbde115504e7a5db97de61f/cupy_cuda13x-14.2.0-cp312-cp312-win_arm64.whl", hash = "sha256:0dae6d788752c987c1cd87adb903640620820feb895339e8bc51a01fba0fb804", size = 35049803, upload-time = "2026-09-01T06:41:43.442Z" }, + { url = "https://files.pythonhosted.org/packages/3c/02/3cb39774c71a16275c50b5d1e92ddcefa948ff1e69f464249b70a79b0abe/cupy_cuda13x-14.2.0-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:449af7f608dc70ce64e8c45cf5a6ea0bbef51623e7c806b1a43445ae3c76e1e0", size = 72969322, upload-time = "2026-08-20T02:41:23.311Z" }, + { url = "https://files.pythonhosted.org/packages/4e/27/216370ae8c39fa9835645be3c27513ab72f52cf4f51faf58d799fba981b6/cupy_cuda13x-14.2.0-cp313-cp313-manylinux2014_x86_64.whl", hash = "sha256:5b79a606f6639d9bed74ef7072f36756a95a804c86d7adb367b39307ac771600", size = 69048238, upload-time = "2026-08-20T02:41:26.831Z" }, + { url = "https://files.pythonhosted.org/packages/5d/75/9df3c545baad4ff2e384db1ff6c8005fa4040492c4ee5672ca8b8711d29d/cupy_cuda13x-14.2.0-cp313-cp313-win_amd64.whl", hash = "sha256:f9a2a55abf889f6df68b5f62cd5f3ec4a2f305ac57ed09becd52e16201edcd86", size = 35965490, upload-time = "2026-08-20T02:41:29.943Z" }, + { url = "https://files.pythonhosted.org/packages/5c/3d/ee85ce39c2f32e8561c81c507df458a914b45d9a0eb1a54e5396c2cfd632/cupy_cuda13x-14.2.0-cp313-cp313-win_arm64.whl", hash = "sha256:5ddaf38291f60ec20e8c7889b28e339bf17ad1b4df768b8bb5e6c9f9e445efbb", size = 35029087, upload-time = "2026-09-01T06:41:48.593Z" }, + { url = "https://files.pythonhosted.org/packages/e5/70/8c80c9e7010193afed8ccb7110d6173aa23bf3a6ce2fc2450244b0fa50a1/cupy_cuda13x-14.2.0-cp314-cp314-manylinux2014_aarch64.whl", hash = "sha256:7ac57c4e16c62f265b5ecc2b45d6bbafe9dc70a99b849ed6cfb8d16f32204589", size = 72840193, upload-time = "2026-08-20T02:41:33.146Z" }, + { url = "https://files.pythonhosted.org/packages/c2/b8/4f4c4f34fc31ab8d136ed965a919505297974629d9852c50e15bfd616281/cupy_cuda13x-14.2.0-cp314-cp314-manylinux2014_x86_64.whl", hash = "sha256:ff0bdebd1b43c0c6db53095784c787c4e4eae671356cf521c4b7482ed78a1a7e", size = 68422828, upload-time = "2026-08-20T02:41:36.527Z" }, + { url = "https://files.pythonhosted.org/packages/7d/5c/58941ece5ed0b2cc8304f8ba81b5e8334759d409306475d53c52ac99809b/cupy_cuda13x-14.2.0-cp314-cp314-win_amd64.whl", hash = "sha256:b99a8bf9d5391954c0c08675a8cb39219923fe2ee249faa3546ebbf75783a360", size = 36110882, upload-time = "2026-08-20T02:41:39.837Z" }, + { url = "https://files.pythonhosted.org/packages/e2/c5/cc7d1a9b764b3f828e09b48a73d0fdeaad15ecd2c7e50507f6d5fcf5fd1a/cupy_cuda13x-14.2.0-cp314-cp314-win_arm64.whl", hash = "sha256:c61c59ba665948c0a855926997ee96026420b628816b2e406998ec797caddb22", size = 35184806, upload-time = "2026-09-01T06:41:58.317Z" }, + { url = "https://files.pythonhosted.org/packages/6f/54/a582ac70e9e75d3fad7ae073051bcd50dd128a1f66dd275291dc5f48d997/cupy_cuda13x-14.2.0-cp314-cp314t-manylinux2014_aarch64.whl", hash = "sha256:e8bb453551deced08ec5c4576da61ebf1e2b4a45aeff48e09f47ec332f7cad3d", size = 73336756, upload-time = "2026-08-20T02:41:42.718Z" }, + { url = "https://files.pythonhosted.org/packages/f9/aa/222151ba5061d7a01efe4009efad9f9a7164e15bb2a494d9aa86efc97d4f/cupy_cuda13x-14.2.0-cp314-cp314t-manylinux2014_x86_64.whl", hash = "sha256:313500c5d415f65d388e84a9d2eb7f07f97c23c85dc5785f986e75d46d9c9693", size = 68792663, upload-time = "2026-08-20T02:41:46.079Z" }, + { url = "https://files.pythonhosted.org/packages/d4/20/13dd37dbeaec879c8348fde2bfca67d96daa56a6e63e4a5083ff61660e53/cupy_cuda13x-14.2.0-cp314-cp314t-win_amd64.whl", hash = "sha256:c517fc55502a7e3d1a8b9310abad95cec61dded5a38d328fe0220bcf7b036268", size = 37004845, upload-time = "2026-08-20T02:41:48.963Z" }, + { url = "https://files.pythonhosted.org/packages/d1/5b/b1367f79abd90a259dbb7f0603dd162c03c3a76ec1d02c2527b0c79d46ba/cupy_cuda13x-14.2.0-cp314-cp314t-win_arm64.whl", hash = "sha256:83885db963eb0b95a0757b4907a37115a1d2aebfa9cf43b99ebac8db35713d37", size = 35485280, upload-time = "2026-09-01T06:41:53.618Z" }, +] + +[package.optional-dependencies] +ctk = [ + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "curand", "cusolver", "cusparse", "nvrtc"] }, +] + [[package]] name = "cycler" version = "0.12.1" @@ -2193,11 +2286,16 @@ wheels = [ [[package]] name = "nvidia-cublas" -version = "13.3.0.5" +version = "13.10.0.20" source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-nvrtc" }, +] wheels = [ - { url = "https://files.pythonhosted.org/packages/e8/5c/08177998e1234459e46b2cdad73738b5516f84b8fa28a8379c678b95c6c0/nvidia_cublas-13.3.0.5-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:48308e7f44feb337ca24d95efdceeac5703fb5dcf9bfb0c23f1eb48015fdd8a1", size = 505057164, upload-time = "2026-03-09T09:43:27.897Z" }, - { url = "https://files.pythonhosted.org/packages/3c/7c/ae5d1751819acff18b0fac29c0a4e93d06d36cfabebe36365ddacc7c32a9/nvidia_cublas-13.3.0.5-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:366568e2dc59e6fe71ffd179f9f2a38b8b2772aed626320a64008651b1e72974", size = 403287501, upload-time = "2026-03-09T09:47:05.046Z" }, + { url = "https://files.pythonhosted.org/packages/b4/12/16628161ed8c109c8c73b38167df50cbc630a8a6e6548e2534d461526c3c/nvidia_cublas-13.10.0.20-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:dcc6d63884c53d07e267ce7e6feac17b2538cebaafec5db865736c3e28e4dc84", size = 444289903, upload-time = "2026-11-04T23:27:26.095Z" }, + { url = "https://files.pythonhosted.org/packages/5c/e3/a5a992c3cb1c49201cd4855345dd50ec9e9c075195a0e1e534ff7e914ca5/nvidia_cublas-13.10.0.20-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c5f3418c0a845232a8cabc58f6fa52cfae7d8bd3680d94c3907aca710fd48d70", size = 560089875, upload-time = "2026-11-04T23:26:43.825Z" }, + { url = "https://files.pythonhosted.org/packages/12/39/e1b1b4adb8b171894d9a89035df0596d448cf73dc403802d24dbe30cc0be/nvidia_cublas-13.10.0.20-py3-none-win_amd64.whl", hash = "sha256:c310b749eea855b88fafb906de08d6819b87216c6711640771764fea5dcd8546", size = 427284863, upload-time = "2026-11-04T23:36:18.775Z" }, + { url = "https://files.pythonhosted.org/packages/4d/89/c875dcd88abc508a98a725214a30b2c001406095855154f32ee3f0946409/nvidia_cublas-13.10.0.20-py3-none-win_arm64.whl", hash = "sha256:05929d6a64af9edce1951ff82a514d630fc1a2369bcfe1ab195f5e85266dbfed", size = 279793527, upload-time = "2026-11-04T23:41:15.348Z" }, ] [[package]] @@ -2282,11 +2380,13 @@ wheels = [ [[package]] name = "nvidia-cuda-nvrtc" -version = "13.2.51" +version = "13.5.37" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5e/21/2fd0aa5a03a8c71962d281084ac44ae7b3b6690d6163ffd7d6486fdb7aa8/nvidia_cuda_nvrtc-13.2.51-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:c88076f32cbbd26e7ebd2107d4b093dd8667e2a90b23b3273d028f3daf574d2e", size = 47019178, upload-time = "2026-03-09T09:38:11.629Z" }, - { url = "https://files.pythonhosted.org/packages/27/ce/ef85d9b59e3fcb0e44041d72de857bf4e87f772d747c603018acbb4addb1/nvidia_cuda_nvrtc-13.2.51-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8d4ddd23dbd4878eea5d970a70631d2668381fc2bdc5ec16cd5faa8c2db24d03", size = 44754448, upload-time = "2026-03-09T09:37:38.021Z" }, + { url = "https://files.pythonhosted.org/packages/a9/54/db4e7f650746265f61c05cf1f8f556dd8c6739b58c823fbc9a7d6fda3247/nvidia_cuda_nvrtc-13.5.37-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:86b4d68058807988b7d9af5b44917a0601e9c337f35ccd657bfb3c759ff8d8bc", size = 58185930, upload-time = "2026-11-04T23:23:57.537Z" }, + { url = "https://files.pythonhosted.org/packages/4e/9f/e8321dfd9dcd923c864213fed3da815da474800ead891a99b6322bdd2039/nvidia_cuda_nvrtc-13.5.37-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:36e4fa2a8ada293f5693d308e531c196ab6f91cf7a0ff479bf96a537d43a6096", size = 56419462, upload-time = "2026-11-04T23:23:44.768Z" }, + { url = "https://files.pythonhosted.org/packages/9f/6c/241c5155199a6de55bad8862213ed64b487dc9526aae817e4af31afea1a4/nvidia_cuda_nvrtc-13.5.37-py3-none-win_amd64.whl", hash = "sha256:e8e6c10b251385707ae3abd7b4ee118de6736fd7678d0f0f4eab4d759f679c21", size = 49248991, upload-time = "2026-11-04T23:35:13.589Z" }, + { url = "https://files.pythonhosted.org/packages/d9/62/b858b2e6da4d9f4ab2744dcd337dc31c75bc4db4ad35543a04ee07bb3eab/nvidia_cuda_nvrtc-13.5.37-py3-none-win_arm64.whl", hash = "sha256:d47388d9054a9f16c705d86e59d7e6c57e85beece36d23d1657833531fa913aa", size = 45119472, upload-time = "2026-11-04T23:40:19.254Z" }, ] [[package]] @@ -2300,11 +2400,13 @@ wheels = [ [[package]] name = "nvidia-cuda-runtime" -version = "13.2.51" +version = "13.5.37" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/12/ec/0fa54349c6cf17054746529a3b99c6bc0049a229ac7fe667a24a244f79fa/nvidia_cuda_runtime-13.2.51-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:dfcccf62936b211a86e3cbc5fcc30647af8e601744e987f6a5925ed045b58672", size = 2340606, upload-time = "2026-03-09T09:29:42.659Z" }, - { url = "https://files.pythonhosted.org/packages/a6/5a/b116ad2b7e574d691458ca0139ab4e9f26beed62184c85570636ce127b7f/nvidia_cuda_runtime-13.2.51-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9c43b06a52c5b9316e19abc047236932c4d5c729969918a83223c4d2a4132f9a", size = 2321923, upload-time = "2026-03-09T09:30:13.061Z" }, + { url = "https://files.pythonhosted.org/packages/d4/8d/d162f19bd57e25048b654337c6f811802abe46c3da0f7d7052a2dfd709e8/nvidia_cuda_runtime-13.5.37-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:de2dc5a7367974228b0a69588a9acc9340d3768bef1bf38bf0801eb08497cd2d", size = 2520010, upload-time = "2026-11-04T23:22:09.169Z" }, + { url = "https://files.pythonhosted.org/packages/0b/7c/14b8f5518f8951e9dca57ea80f44f8ea8bcc8671812e6cc38009f3205e57/nvidia_cuda_runtime-13.5.37-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:df2b79a8852f10d1fef8c848489c49eebb5754051bf870202a862831412b179d", size = 2504024, upload-time = "2026-11-04T23:22:14.265Z" }, + { url = "https://files.pythonhosted.org/packages/35/99/e07dd9565d2e0317c1a62ff39ae9fc0542c048f74bcfc777c8be97710347/nvidia_cuda_runtime-13.5.37-py3-none-win_amd64.whl", hash = "sha256:f225d3fa5ea6bfbb58fc4979ef72375cbe661d9107f0d8a4ff22cc151525e4a9", size = 2780254, upload-time = "2026-11-04T23:34:42.881Z" }, + { url = "https://files.pythonhosted.org/packages/8f/35/dbd62813d7cdc3a958bec65b846e8bb3ff41e7af461802f7c0e2a522d7b2/nvidia_cuda_runtime-13.5.37-py3-none-win_arm64.whl", hash = "sha256:dcbcb895bc2098b4a6784b12b0a1ceafb09359dc95f4f92ea688ac935a106877", size = 2779932, upload-time = "2026-11-04T23:39:39.147Z" }, ] [[package]] @@ -2342,14 +2444,16 @@ wheels = [ [[package]] name = "nvidia-cufft" -version = "12.2.0.37" +version = "12.5.0.14" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-nvjitlink" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/6d/4d/31158ab042b044b269019574da4430ecce8d05fec7af1d270e1ba28e3512/nvidia_cufft-12.2.0.37-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d0d762b7ece2f2a4971a5f29a6cef13f26fd87c7cd0003b48ed82d9a1d7380d7", size = 218244744, upload-time = "2026-03-09T09:48:16.661Z" }, - { url = "https://files.pythonhosted.org/packages/00/a8/d8c0a8c4c45a3904a52c9860b07fdf775ca0517df884e3d240205a42b7ff/nvidia_cufft-12.2.0.37-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1530739e18736b07f57f835664659aa99179dab7b567c581a1ec7cb6c9737662", size = 218258094, upload-time = "2026-03-09T09:48:55.172Z" }, + { url = "https://files.pythonhosted.org/packages/1d/64/a71dc77eaca1623d1ff61f84045f2d21fb50a77796551e062bb8a989326d/nvidia_cufft-12.5.0.14-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:091cde103312d6e09366122fe3f67092dcfd94cc9785029d695d6bcffaa76fdf", size = 165074034, upload-time = "2026-11-04T23:27:54.684Z" }, + { url = "https://files.pythonhosted.org/packages/c6/8f/f909edaadf9a7acf1b92f805a384b6984193bd595b45005f8e320d4170cf/nvidia_cufft-12.5.0.14-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b8be26837a5abad2b5b74f3e92dd29f622a42934b355be0961a48712cbba23ad", size = 165072555, upload-time = "2026-11-04T23:28:25.686Z" }, + { url = "https://files.pythonhosted.org/packages/94/79/8f2a7125f560b86571c0be3dbf10e33041187bc4a8cf66107541fe8e5cce/nvidia_cufft-12.5.0.14-py3-none-win_amd64.whl", hash = "sha256:78e7585159ecf37d73af1ce43a5dea7fd3360e1af71cde08f070c07966a336bc", size = 164476208, upload-time = "2026-11-04T23:36:37.446Z" }, + { url = "https://files.pythonhosted.org/packages/9c/78/ec6cff1d3ba1d743acd7d6889db148fa00b97038bce8b169e49edb7c5a1b/nvidia_cufft-12.5.0.14-py3-none-win_arm64.whl", hash = "sha256:77eaacecf83758e8e20458a199d56359d46596e3a2814af741f27d1f483381e9", size = 165385033, upload-time = "2026-11-04T23:41:34.766Z" }, ] [[package]] @@ -2364,9 +2468,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/95/f4/61e6996dd20481ee834f57a8e9dca28b1869366a135e0d42e2aa8493bdd4/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c67884f2a7d276b4b80eb56a79322a95df592ae5e765cf1243693365ccab4e28", size = 200877592, upload-time = "2025-06-05T20:05:45.862Z" }, ] +[[package]] +name = "nvidia-curand" +version = "10.4.5.25" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/44/33/c8d7fa70f90bde9b88a665734ab5e0440d4749646ac4e8590482760af3c6/nvidia_curand-10.4.5.25-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:07c75d8ba14b627d1a01d2c1bfa1eece88f25d564d3aae7e6b6fb44f7bbec9c1", size = 77529633, upload-time = "2026-11-04T23:28:49.816Z" }, + { url = "https://files.pythonhosted.org/packages/12/63/636e5dc9afb2959bfc0042a357d5567358df829b66c96d9fc36d77dc0596/nvidia_curand-10.4.5.25-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:22bb226a857e5888fe6b54ae465fc4371296bcff13a35f79e19691bcf454f529", size = 77113015, upload-time = "2026-11-04T23:29:04.942Z" }, + { url = "https://files.pythonhosted.org/packages/ee/b6/5b38490991b6875ae60c4f41e190ac91cadd1e6a1324b152999cc4a3f580/nvidia_curand-10.4.5.25-py3-none-win_amd64.whl", hash = "sha256:f6a02a70f5b9715235df31a164a4bf3237a45be1600a5f435f5e4a0d8682eb22", size = 77330709, upload-time = "2026-11-04T23:36:49.798Z" }, + { url = "https://files.pythonhosted.org/packages/c7/1c/4806430849d745b8133481fe9c6ac0d47d2ca5c54f2ae34c4f64b7cb7032/nvidia_curand-10.4.5.25-py3-none-win_arm64.whl", hash = "sha256:c998546b8bd1c746f2872cc11ec4c79b006dfed2a5b0acd25e7b710522b2d7f4", size = 77108234, upload-time = "2026-11-04T23:41:50.279Z" }, +] + [[package]] name = "nvidia-cusolver" -version = "12.1.0.51" +version = "12.4.1.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-cublas" }, @@ -2374,8 +2489,10 @@ dependencies = [ { name = "nvidia-nvjitlink" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/25/09/6214d11749dfc2d1be9a9e0bcfb04069077b98f7df0ad41115da87c84700/nvidia_cusolver-12.1.0.51-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:2a0bde343f956934103bc19344584a2d7b95d03b8cca9c0d22c90ff8917bbc3c", size = 224103511, upload-time = "2026-03-09T09:51:14.704Z" }, - { url = "https://files.pythonhosted.org/packages/ec/3a/6ee9b1c6632ec9cc0339996ffb331e5a8cbedcd361f7d4d0b63d48519a28/nvidia_cusolver-12.1.0.51-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0148ba705c196075607cd9d7a856a834695b406907b1ba8ad99b8a325a463611", size = 201732826, upload-time = "2026-03-09T09:51:45.431Z" }, + { url = "https://files.pythonhosted.org/packages/80/38/566a08a71dac1eebb21fa439970eab971d372c94633c3b16b851b415a13d/nvidia_cusolver-12.4.1.2-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:1f54c8e0d5b1ebd183b4c5a79f0709fa90241e8254da381ca5b969ee97c8de32", size = 304822178, upload-time = "2026-11-04T23:29:37.266Z" }, + { url = "https://files.pythonhosted.org/packages/78/4f/c5c926dbe0100e485c3ef86022ba3f6799fc57df6f9a73aaadeb7ab368fe/nvidia_cusolver-12.4.1.2-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b3b8e02fd2f2df869369af075dc3a179d43f6f7b85e94632a7dce556697be044", size = 260003940, upload-time = "2026-11-04T23:30:56.62Z" }, + { url = "https://files.pythonhosted.org/packages/cb/ca/bdf3af194210aad2b34db2b31033698be8b443dab38f566893149142d3d0/nvidia_cusolver-12.4.1.2-py3-none-win_amd64.whl", hash = "sha256:6b7be7d7dd60ba0e91cae19357a1d8c57dae321c0df99d92792ec025d791cc88", size = 274964624, upload-time = "2026-11-04T23:37:14.787Z" }, + { url = "https://files.pythonhosted.org/packages/c4/f9/733a3c7876db2f8fb154315cd4832fcbd2e6acd74f9274521444416388a8/nvidia_cusolver-12.4.1.2-py3-none-win_arm64.whl", hash = "sha256:8ec8a68b661ecdb89f4724a993ea9fab5c8096fc14aee5452344628a11529715", size = 103368508, upload-time = "2026-11-04T23:42:04.52Z" }, ] [[package]] @@ -2394,14 +2511,16 @@ wheels = [ [[package]] name = "nvidia-cusparse" -version = "12.7.9.17" +version = "12.9.4.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-nvjitlink" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/06/5a/6bb7fc5f9658902efebc8551b1a9265b7a5908cbf9efdabdf97bc30b168a/nvidia_cusparse-12.7.9.17-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7c5d21980dc5f064c7ba13af2250e7fb4036283f1d1943c2915573c27928c89e", size = 168894384, upload-time = "2026-03-09T09:52:23.003Z" }, - { url = "https://files.pythonhosted.org/packages/87/40/23990a83164aaec2bfeffcee87794299f3cfdbdd7ed024b2af078afb600a/nvidia_cusparse-12.7.9.17-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7fb409bc7bb85e7a95706bd1e0b502b418a026dc35823179b4dafa92f1f2f7fd", size = 150883058, upload-time = "2026-03-09T09:53:00.781Z" }, + { url = "https://files.pythonhosted.org/packages/1b/71/9cf160aff268332d07b1d46e9009c5f28d626d36bd17d9cdc462a0668384/nvidia_cusparse-12.9.4.2-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e2ce97f185c9c4b285e4a8fbf02cb96dd1b8dd66736ebe2f0a8da5f6012a4fb7", size = 230296529, upload-time = "2026-11-04T23:31:18.146Z" }, + { url = "https://files.pythonhosted.org/packages/30/7a/9928f812629a0e871871f2692cf5618857c51bf7093c7ecf2f78d17f9255/nvidia_cusparse-12.9.4.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:072667e58c851a318738f5527a12eb225cbf058803f2ea451c722d8cf8edddec", size = 230422857, upload-time = "2026-11-04T23:31:36.722Z" }, + { url = "https://files.pythonhosted.org/packages/71/9a/54114b62952d16a67e43a696916014e1e19ed6f495b0c1755d5baa967c7f/nvidia_cusparse-12.9.4.2-py3-none-win_amd64.whl", hash = "sha256:2defba1fbe76140c9eedc9fe82d58f11a0e2b7fd2c12dc5fb1267cc076689669", size = 228382183, upload-time = "2026-11-04T23:37:44.311Z" }, + { url = "https://files.pythonhosted.org/packages/a9/4d/18ba5c89b1968c8af804b33f647b2d94307c6fa98591a84d985f7b657b12/nvidia_cusparse-12.9.4.2-py3-none-win_arm64.whl", hash = "sha256:50d7a8c47c377c44ec1db681710156a28246ddfe3bde7be801daa85aba5483cc", size = 228043702, upload-time = "2026-11-04T23:42:34.936Z" }, ] [[package]] @@ -2436,11 +2555,13 @@ wheels = [ [[package]] name = "nvidia-nvjitlink" -version = "13.2.51" +version = "13.5.37" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/12/ae/ef3c49f1918aef93b39045499bfdb0ac9fb13e1785bc83f7a1b5d58a292d/nvidia_nvjitlink-13.2.51-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:6703c9ed79301382787a23fda9a7388af0779ecbc37545e4d50c055c897694a0", size = 41370377, upload-time = "2026-03-09T09:55:51.938Z" }, - { url = "https://files.pythonhosted.org/packages/6d/44/9e42f8ad320a23dfcb9542fc6ccd06a3a3ef495a779ade5937540fe7aa59/nvidia_nvjitlink-13.2.51-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3ba8bc06d756d10b643929e1f65dd162120c3e25ace58bd6414ee7eb98e10dde", size = 39285486, upload-time = "2026-03-09T09:55:19.032Z" }, + { url = "https://files.pythonhosted.org/packages/cc/1f/2321702b8d60b4b4dd18b957f61612dbac9cb1c07e940cf6582106b7c8b8/nvidia_nvjitlink-13.5.37-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:280d7cb9649704145f074b8dbde500de04ce02d0872e066e20a83c17bcaa93eb", size = 46642491, upload-time = "2026-11-04T23:33:31.696Z" }, + { url = "https://files.pythonhosted.org/packages/24/15/dd228f798cd01d7aeaf19d872b2b20f9aeafc1a7d8b6bfaa878b30b8d6f7/nvidia_nvjitlink-13.5.37-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:893f6f9a55b2720ed94913b5ba3332e972da8c78dfa2863d0244403dbd528ffc", size = 45253761, upload-time = "2026-11-04T23:33:22.854Z" }, + { url = "https://files.pythonhosted.org/packages/a7/53/c43398492e8374880d5e01c23c4dac33ba0c3bd03515d6ecc574254c0967/nvidia_nvjitlink-13.5.37-py3-none-win_amd64.whl", hash = "sha256:4465d4e8f2212118220b580c82523c5654f93a62906bd51d29d79bdfc12376c1", size = 41467665, upload-time = "2026-11-04T23:38:44.689Z" }, + { url = "https://files.pythonhosted.org/packages/d7/4a/1ff9ab7d0c961a1b5539f8291c881bea82adbbf8ba40a9e81f9c9b3a8d5d/nvidia_nvjitlink-13.5.37-py3-none-win_arm64.whl", hash = "sha256:eb30348ae64c525180a280025dd12a089c3e2e1dc4e19cdd58f68c6edc622404", size = 38552904, upload-time = "2026-11-04T23:43:27.658Z" }, ] [[package]] From bfa5063695eaa8d791fb3d1019b63d58215e7f51 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Tue, 10 Nov 2026 13:36:39 -0700 Subject: [PATCH 10/20] fix(statevector): fix CUDA sampler issues found in review - Rows of 97 to 176 output words overflowed the shared memory of the gather and noise_kept kernels (CUDA_ERROR_INVALID_VALUE); such circuits now sample on the CPU with a warning. - Calls on one sampler from several threads are serialized, as on Metal. - An empty postselection mask no longer shares the module of no mask, which returned no rows. - A sampler compiled before a fork samples on the CPU in the child, with a warning, instead of failing to initialize CUDA. - GPU samplers pickle: the inputs and the stream state are kept, and the GPU state is rebuilt on load (the CPU sampler where CUDA is missing). Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/cuda/sampler.py | 79 ++++++++++++++++++++++++-- test/integration/test_cuda_sampler.py | 82 +++++++++++++++++++++++++++ 2 files changed, 157 insertions(+), 4 deletions(-) diff --git a/src/tsim/statevector/cuda/sampler.py b/src/tsim/statevector/cuda/sampler.py index dc058d13..1fab3d1b 100644 --- a/src/tsim/statevector/cuda/sampler.py +++ b/src/tsim/statevector/cuda/sampler.py @@ -29,6 +29,7 @@ from __future__ import annotations import math +import threading import warnings from dataclasses import dataclass from typing import Any @@ -45,6 +46,8 @@ from tsim.statevector.shots import ShotLayout, ShotPlan, ShotSampler MAX_WIDTH = 11 +KERNEL_SHARED_BYTES = 48 * 1024 +"""Shared memory that a kernel may use without opting in, static and dynamic.""" """Widest component the CUDA executor runs: 32 threads of 64 amplitudes.""" BATCH_SHOTS = 1 << 22 @@ -331,6 +334,8 @@ def __init__( self._cp = runtime.cupy() cp = self._cp assert self.supports(host.layout) + self._args = (host, channels, num_detectors, seed) + self._lock = threading.RLock() self._host = host layout = host.layout self._num_detectors = num_detectors @@ -339,9 +344,17 @@ def __init__( self._channels = channels self._noise = noise_tables(channels, layout.noise, layout.num_words) n = self._noise + # NOTE: gather_columns and noise_kept stage a block's output words and its + # packed rows in shared memory. + staged = 8 * layout.output_words * n.block_shots + staged += n.block_shots * -(-layout.num_outputs // 8) + if staged > KERNEL_SHARED_BYTES: + raise ValueError( + f"{layout.num_outputs} outputs do not fit the GPU kernels' shared memory" + ) self._record_args = [self._upload(n.overflow.view(np.uint8))] self._all = self._launch(n.groups) - self._modules: dict[bytes, _Module] = {} + self._modules: dict[bytes | None, _Module] = {} self._buffers: dict[str, Any] = {} self._views: dict[tuple[str, tuple[int, ...]], Any] = {} self._constants: dict[bytes, Any] = {} @@ -358,6 +371,10 @@ def __init__( # NOTE: lowered when a call runs out of device memory. self._max_batch = BATCH_SHOTS + 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.""" @@ -383,7 +400,8 @@ def _threads_per_shot(self) -> int: return min(5, max(0, width - 5)) def _module(self, mask: np.ndarray | None) -> _Module: - key = b"" if mask is None else np.asarray(mask, bool).tobytes() + # NOTE: None, not b"", so that an empty mask is not taken for no mask. + key = None if mask is None else np.asarray(mask, bool).tobytes() if key in self._modules: return self._modules[key] cp = self._cp @@ -768,7 +786,10 @@ def sample( postselection are returned. ``batch_size`` bounds the shots of a device batch. """ - accepted = self._sample_or_none(shots, mask, flip, batch_size) + accepted = None + if not _forked(): + with self._lock: + accepted = self._sample_or_none(shots, mask, flip, batch_size) if accepted is None: return self._host.sample( shots, @@ -868,6 +889,24 @@ def sample_accepted( beyond float32 rounding. """ + with self._lock: + return self._sample_accepted( + shots, + mask=mask, + flip=flip, + expectations=expectations, + batch_size=batch_size, + ) + + def _sample_accepted( + self, + shots: int, + *, + mask: np.ndarray | None, + flip: np.ndarray | None, + expectations: bool, + batch_size: int | None, + ) -> tuple[np.ndarray, np.ndarray | None, float]: cp = self._cp module = self._module(mask) layout = self.layout @@ -948,6 +987,7 @@ def finish() -> float: if no and layout.components: if index is None: index = cp.arange(n, dtype=cp.uint32) + held.append(index) dst = cp.empty((count, no), dtype=np.float32) module.gather_float( (runtime.multiprocessors() * 8,), @@ -997,9 +1037,40 @@ 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.""" - accepted = self._sample_or_none(shots, None, flip, None) + accepted = None + if not _forked(): + with self._lock: + accepted = self._sample_or_none(shots, None, flip, None) if accepted is None: return self._host.direct_bools(shots, flip) rows, _ = accepted num_outputs = self.layout.num_outputs return select(rows, np.arange(num_outputs), num_outputs, False) + + +def _forked() -> bool: + """Return whether this process was forked after its parent used CUDA, with a warning.""" + if not runtime.forked_after_init(): + return False + 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 True + + +def _restore( + host: ShotSampler, + channels: ChannelSampler, + num_detectors: int, + seed: np.random.SeedSequence, + state: dict, +) -> CudaShotSampler | ShotSampler: + """Rebuild a pickled ``CudaShotSampler``, or its CPU sampler where CUDA is missing.""" + if not runtime.available(): + return host + shots = CudaShotSampler(host, channels, num_detectors, seed) + shots._rng.bit_generator.state = state + return shots diff --git a/test/integration/test_cuda_sampler.py b/test/integration/test_cuda_sampler.py index c3d17e17..74e3ea16 100644 --- a/test/integration/test_cuda_sampler.py +++ b/test/integration/test_cuda_sampler.py @@ -1,6 +1,8 @@ """The CUDA 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 @@ -703,3 +705,83 @@ def test_noise_beyond_the_gpu_tables_samples_on_the_cpu(): sampler = Circuit(text).compile_detector_sampler(seed=1) assert type(sampler._shots) is ShotSampler assert sampler.sample(10).shape == (10, 9) + + +_FORKED = {} + + +def _sample_in_child(_): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + shape = _FORKED["sampler"].sample(100).shape + return shape, [str(w.message) for w in caught] + + +def test_samplers_compiled_before_a_fork_sample_on_the_cpu_in_the_child(): + sampler = _gpu(Circuit(GHZ_WITH_DIRECT_DETECTOR).compile_detector_sampler(seed=1)) + sampler.sample(10) + # NOTE: the child gets the sampler by fork, not by pickle. + _FORKED["sampler"] = sampler + try: + with multiprocessing.get_context("fork").Pool(1) as pool: + shape, messages = pool.map(_sample_in_child, [0])[0] + finally: + _FORKED.clear() + assert shape == (100, 3) + assert any("forked" in m for m in messages) + + +def test_gpu_samplers_pickle_and_continue_their_streams(): + sampler = _gpu(Circuit(GHZ_WITH_DIRECT_DETECTOR).compile_detector_sampler(seed=26)) + sampler.sample(1000) + loaded = pickle.loads(pickle.dumps(sampler)) # noqa: S301 + assert isinstance(loaded._shots, CudaShotSampler) + np.testing.assert_array_equal(loaded.sample(5000), sampler.sample(5000)) + + +def test_an_empty_mask_is_not_taken_for_no_mask(): + sampler = _gpu(Circuit(T_STATE_X).compile_detector_sampler(seed=1)) + assert sampler.sample(1000, append_observables=True).shape == (1000, 1) + empty = np.zeros(0, dtype=bool) + kept = sampler.sample(1000, append_observables=True, postselection_mask=empty) + assert kept.shape == (1000, 1) + + +def test_outputs_beyond_the_kernels_shared_memory_sample_on_the_cpu(): + # NOTE: 6200 outputs need 97 words, whose packed rows do not fit in shared + # memory beside them. + n = 6200 + qubits = " ".join(map(str, range(n))) + circuit = Circuit(f"R {qubits}\nX_ERROR(0.01) {qubits}\nM {qubits}") + with pytest.warns(UserWarning, match="cannot sample this circuit"): + sampler = circuit.compile_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + assert sampler.sample(10, bit_packed=True).shape == (10, n // 8) + + +def test_concurrent_calls_on_one_sampler_are_serialized(): + sampler = _gpu(Circuit(GHZ_WITH_DIRECT_DETECTOR).compile_detector_sampler(seed=1)) + mask = np.array([False, False, True]) + results, errors = [], [] + + def work(k): + try: + for _ in range(20): + got = sampler.sample(20_000, postselection_mask=mask if k % 2 else None) + results.append((k % 2, got)) + except Exception as error: + errors.append(error) + + threads = [threading.Thread(target=work, args=(k,)) for k in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert not errors + for masked, got in results: + assert got.shape[1] == 3 + if masked: + assert not got[:, 2].any() + else: + # NOTE: detector 2 fires at rate 0.2. + assert abs(got[:, 2].mean() - 0.2) < 0.02 From d832e6656c10df30a3ede7ec15ca5171301e2d9a Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 06:22:18 -0700 Subject: [PATCH 11/20] refactor(cuda): pack whole shot words, without piece tables The sampler always returns every output in order, so pack_rows copies output word k and clears the bits past the last output, instead of reading a piece table per word (gpu.rows.pieces is gone). Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/cuda/noise.py | 35 ++++++++------------ src/tsim/statevector/cuda/sampler.py | 48 ++++++++-------------------- 2 files changed, 28 insertions(+), 55 deletions(-) diff --git a/src/tsim/statevector/cuda/noise.py b/src/tsim/statevector/cuda/noise.py index 27013183..f90b9e65 100644 --- a/src/tsim/statevector/cuda/noise.py +++ b/src/tsim/statevector/cuda/noise.py @@ -43,26 +43,18 @@ }} }} -// The returned outputs of the block's ``ns`` rows ``sw[w * S + s]``, bit-packed -// into {NB} bytes per row, contiguously at ``dst``. Output word ``k`` of a row -// is the pieces ``pieces[start[k] : start[k + 1]]``, each ``(word, bit | (length -// - 1) << 6 | position << 12)``: ``length`` bits from ``bit`` of a shot word. -// Consecutive threads pack one word of consecutive rows (one piece, distinct -// banks) into ``stage``, from which the block's rows go out 16 bytes at a time. -__device__ __forceinline__ void pack_rows(const u64* sw, int ns, const uint2* __restrict__ pieces, - const unsigned* __restrict__ start, unsigned char* stage, unsigned char* __restrict__ dst) +// 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``. +// Consecutive threads pack one word of consecutive rows (distinct banks) into +// ``stage``, from which the block's rows go out 16 bytes at a time. +__device__ __forceinline__ void pack_rows(const u64* sw, int ns, unsigned char* stage, + unsigned char* __restrict__ dst) {{ for (int i = threadIdx.x; i < {NK} * {S}; i += blockDim.x) {{ int k = i / {S}, s = i - k * {S}; if (s >= ns) continue; - u64 v = 0; - for (unsigned 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 &= (1ULL << length) - 1; - v |= x << (q.y >> 12); - }} + u64 v = sw[k * {S} + s]; + if (k == {NK} - 1) v &= {LAST_WORD}; unsigned char* at = stage + s * {NB} + 8 * k; #pragma unroll for (int j = 0; j < 8; ++j) if (8 * k + j < {NB}) at[j] = (unsigned char)(v >> (8 * j)); @@ -212,8 +204,7 @@ const int* __restrict__ seg_begin, const int* __restrict__ seg_end, const float* __restrict__ seg_log_keep, const Record* __restrict__ records, const Flip* __restrict__ overflow, - unsigned char* __restrict__ dst, const uint2* __restrict__ pieces, - const unsigned* __restrict__ piece_start) + unsigned char* __restrict__ dst) {{ __shared__ u64 sw[{OW} * {S}]; extern __shared__ __align__(16) uint2 events[]; @@ -230,7 +221,7 @@ num_segments, seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep, records, overflow); // NOTE: the events are applied, so their buffer stages the packed rows. - pack_rows(sw, ns, pieces, piece_start, (unsigned char*)events, dst + base * {NB}); + pack_rows(sw, ns, (unsigned char*)events, dst + base * {NB}); }} """ @@ -248,15 +239,17 @@ def noise_tables( return _noise_tables(channels, flips, num_words, shared_bytes=SHARED_BYTES) -def source(tables: NoiseTables, output_words: int, packed_bytes: int) -> str: +def source(tables: NoiseTables, output_words: int, num_outputs: int) -> str: """Return the CUDA source of the ``noise`` and ``noise_kept`` kernels. - ``noise_kept`` packs rows of ``packed_bytes`` bytes. + ``noise_kept`` packs rows of ``num_outputs`` outputs. """ + packed_bytes = -(-num_outputs // 8) return SOURCE.format( NW=tables.num_words, OW=output_words, 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}ULL", ) diff --git a/src/tsim/statevector/cuda/sampler.py b/src/tsim/statevector/cuda/sampler.py index 1fab3d1b..caa9d2f5 100644 --- a/src/tsim/statevector/cuda/sampler.py +++ b/src/tsim/statevector/cuda/sampler.py @@ -41,7 +41,7 @@ from tsim.statevector.cuda.codegen import PRELUDE, hex64, tape_function from tsim.statevector.cuda.noise import noise_tables, source as noise_source from tsim.statevector.gpu.noise import NoiseGroups, deferred_channels, pass_rate -from tsim.statevector.gpu.rows import pieces, select +from tsim.statevector.gpu.rows import select from tsim.statevector.shared import pack_bits from tsim.statevector.shots import ShotLayout, ShotPlan, ShotSampler @@ -160,11 +160,10 @@ for (; w; w &= w - 1) index[at++] = (unsigned)(32 * i + __ffs(w) - 1); }} -// The returned outputs of rows ``index[j]`` (or every row), bit-packed densely +// The outputs of rows ``index[j]`` (or every row), bit-packed densely // (``pack_rows``): a block packs {S} rows' output words. extern "C" __global__ void __launch_bounds__(256) gather_columns( const u64* __restrict__ rows, long long n, const unsigned* __restrict__ index, - const uint2* __restrict__ pieces, const unsigned* __restrict__ start, unsigned char* __restrict__ dst) {{ __shared__ u64 sw[{OW} * {S}]; extern __shared__ __align__(16) unsigned char stage[]; @@ -176,7 +175,7 @@ sw[w * {S} + s] = rows[r * {NW} + w]; }} __syncthreads(); - pack_rows(sw, ns, pieces, start, stage, dst + base * {NB}); + pack_rows(sw, ns, stage, dst + base * {NB}); }} // The batch's counts, stored into pinned host memory: a read that does not @@ -222,10 +221,6 @@ class _Module: screen: Any tape_blocks: int has_mask: bool - columns: np.ndarray - """The outputs that ``sample_accepted`` returns: all of them, in order.""" - pieces: Any - piece_start: Any plan: ShotPlan psm_words: np.ndarray passing: float @@ -435,8 +430,6 @@ 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) # 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}) @@ -446,7 +439,7 @@ def _module(self, mask: np.ndarray | None) -> _Module: ) source = ( PRELUDE - + noise_source(self._noise, layout.output_words, -(-len(columns) // 8)) + + noise_source(self._noise, layout.output_words, layout.num_outputs) + functions + _KERNELS.format( block=block, @@ -458,7 +451,8 @@ def _module(self, mask: np.ndarray | None) -> _Module: NO=max(1, len(self._observables)), NE=max(1, len(self._observables)), S=self._noise.block_shots, - NB=max(-(-len(columns) // 8), 1), + # NOTE: the rows keep every output, the postselected ones too. + NB=max(-(-layout.num_outputs // 8), 1), calls=calls, # NOTE: each component's stream, from the call's seed. seeds="\n".join( @@ -479,7 +473,6 @@ def _module(self, mask: np.ndarray | None) -> _Module: tape.ptr, block, 0 ), ) - found, piece_start = pieces(columns) compiled = _Module( noise=module.get_function("noise"), noise_kept=module.get_function("noise_kept"), @@ -492,9 +485,6 @@ def _module(self, mask: np.ndarray | None) -> _Module: screen=self._upload(plan.screen), tape_blocks=runtime.multiprocessors() * blocks_per_sm, has_mask=mask is not None, - columns=columns, - pieces=self._upload(found), - piece_start=self._upload(piece_start), plan=plan, psm_words=psm_words, passing=self._passing(plan), @@ -706,12 +696,11 @@ def _pack( ) -> Any: """Return the ``count`` rows ``index`` (or all) on the device, as copied to the host. - Their outputs ``module.columns``, bit-packed; with ``split`` the deferred - noise kernel adds the deferred channels' faults first. + Their outputs, bit-packed; with ``split`` the deferred noise kernel adds + the deferred channels' faults first. """ S = self._noise.block_shots - m = len(module.columns) - nb = -(-m // 8) + nb = -(-self.layout.num_outputs // 8) # NOTE: a new array per chunk, so that refilling it never waits for a copy. dst = self._cp.empty((max(count, 1), nb), dtype=np.uint8) if not count or not nb: @@ -733,8 +722,6 @@ def _pack( *launch.args, *self._record_args, dst, - module.pieces, - module.piece_start, ), shared_mem=max(8 * launch.max_events, S * nb), ) @@ -746,8 +733,6 @@ def _pack( words, np.int64(count), index if index is not None else np.uint64(0), - module.pieces, - module.piece_start, dst, ), shared_mem=S * nb, @@ -910,7 +895,7 @@ def _sample_accepted( cp = self._cp module = self._module(mask) layout = self.layout - nb, no = -(-len(module.columns) // 8), len(self._observables) + nb, no = -(-layout.num_outputs // 8), len(self._observables) if not shots: values = np.zeros((0, no), np.float32) if expectations else None return np.zeros((0, nb), np.uint8), values, 0.0 @@ -1016,22 +1001,17 @@ def finish() -> float: rows, values = out.result() if values is None: return rows, None, deviation - self._direct_expectations(rows, values, module.columns) + self._direct_expectations(rows, values) return rows, values, deviation - def _direct_expectations( - self, rows: np.ndarray, values: np.ndarray, columns: np.ndarray - ) -> None: - """Set the expectations of observables outside the components: ``±1``. - - ``columns`` are the outputs that ``rows`` hold. - """ + 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 5b354f6a49fd0cb8d2721ba83096460bebd6bd1f Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 06:26:15 -0700 Subject: [PATCH 12/20] refactor(cuda): use the shared GPU helpers The tape kernel's glue, the noise split, the screen's pass rate, the deviation, the width limit and the fork warning now come from tsim.statevector.gpu; the samples of a seed are unchanged. codegen's tape_function and hex64 pass-throughs, the duplicate NE argument and a misplaced docstring go. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 11 +- src/tsim/statevector/cuda/codegen.py | 10 -- src/tsim/statevector/cuda/sampler.py | 151 +++++++------------------- test/integration/test_cuda_sampler.py | 8 +- 4 files changed, 46 insertions(+), 134 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 67a36e44..e6bde5a0 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -21,6 +21,7 @@ from tsim.core.types import CompiledComponent, CompiledProgram, SamplingStrategy from tsim.noise.channels import ChannelSampler from tsim.statevector.cuda import runtime as cuda_runtime +from tsim.statevector.gpu import FORKED from tsim.statevector.program import StatevectorComponent from tsim.statevector.shots import ShotSampler from tsim.utils.cuda_helpers import copy_d2h @@ -332,16 +333,12 @@ def _gpu_shots( if os.environ.get("TSIM_DEVICE", "").lower() == "cpu": return host if cuda_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(FORKED, stacklevel=5) return host if not cuda_runtime.available(): return host - from tsim.statevector.cuda.sampler import MAX_WIDTH, CudaShotSampler + from tsim.statevector.cuda.sampler import CudaShotSampler + from tsim.statevector.gpu.tape import MAX_WIDTH layout = host.layout if layout.num_words == 0: diff --git a/src/tsim/statevector/cuda/codegen.py b/src/tsim/statevector/cuda/codegen.py index 2193b62c..4a65161e 100644 --- a/src/tsim/statevector/cuda/codegen.py +++ b/src/tsim/statevector/cuda/codegen.py @@ -8,7 +8,6 @@ from __future__ import annotations from tsim.statevector.gpu import tape -from tsim.statevector.shared import LaneProgram CUDA = tape.Dialect( inline="__device__ __forceinline__", @@ -78,12 +77,3 @@ if (!(d <= 1e-3f)) atomicMax(deviation, d == d ? __float_as_uint(d) : 0x7f800000u); } """ - -hex64 = CUDA.hex64 - - -def tape_function( - program: LaneProgram, index: int, tb: int, observables: dict[int, int] -) -> str: - """Return ``component_`` in CUDA C++ (``tape.tape_function``).""" - return tape.tape_function(program, index, tb, observables, CUDA) diff --git a/src/tsim/statevector/cuda/sampler.py b/src/tsim/statevector/cuda/sampler.py index caa9d2f5..5041a1bc 100644 --- a/src/tsim/statevector/cuda/sampler.py +++ b/src/tsim/statevector/cuda/sampler.py @@ -38,17 +38,15 @@ from tsim.noise.channels import ChannelSampler from tsim.statevector.cuda import runtime -from tsim.statevector.cuda.codegen import PRELUDE, hex64, tape_function +from tsim.statevector.cuda.codegen import CUDA, PRELUDE from tsim.statevector.cuda.noise import noise_tables, source as noise_source -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 -from tsim.statevector.shared import pack_bits -from tsim.statevector.shots import ShotLayout, ShotPlan, ShotSampler +from tsim.statevector.shots import ShotPlan, ShotSampler -MAX_WIDTH = 11 KERNEL_SHARED_BYTES = 48 * 1024 """Shared memory that a kernel may use without opting in, static and dynamic.""" -"""Widest component the CUDA executor runs: 32 threads of 64 amplitudes.""" BATCH_SHOTS = 1 << 22 """Shots of one batch on the device.""" @@ -56,13 +54,6 @@ MIN_BATCH_SHOTS = 1 << 14 """Smallest batch that a call out of device memory retries with, before the CPU.""" -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.""" - CHUNK_BYTES = 6 << 20 """Bytes of kept rows per chunk that a call's chunks aim at: each chunk's copy to the host overlaps the next chunk's kernels on the other stream.""" @@ -95,7 +86,7 @@ u64 w[{NW}]; #pragma unroll for (int i = 0; i < {NW}; ++i) w[i] = row[i]; - float ev[{NE}]; + float ev[{NO}]; int grave = 0x7fffffff; {calls} if (valid && t == 0) {{ @@ -306,10 +297,7 @@ class CudaShotSampler: device = "cuda" - @staticmethod - def supports(layout: ShotLayout) -> bool: - """Return whether every component of ``layout`` fits in ``MAX_WIDTH`` qubits.""" - return all(c.width <= MAX_WIDTH for c in layout.components) + supports = staticmethod(tape.supports) def __init__( self, @@ -389,11 +377,6 @@ def reference(self) -> tuple[np.ndarray, float]: # -- compilation ---------------------------------------------------------- - def _threads_per_shot(self) -> int: - """Return log2 of the threads per shot: one per 32 amplitudes, at most 32.""" - width = max((c.width for c in self.layout.components), default=0) - return min(5, max(0, width - 5)) - def _module(self, mask: np.ndarray | None) -> _Module: # NOTE: None, not b"", so that an empty mask is not taken for no mask. key = None if mask is None else np.asarray(mask, bool).tobytes() @@ -402,45 +385,19 @@ def _module(self, mask: np.ndarray | None) -> _Module: cp = self._cp plan = self.plan(mask) layout = self.layout - tb = self._threads_per_shot() + tb = tape.threads_per_shot(layout) T = 1 << tb block = 128 - 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, gmask, ev, deviation));" - for i in range(len(plan.programs)) - ) + parts = tape.tape_kernel(layout, plan, mask, self._num_detectors, tb, CUDA) gmask = ( "0xffffffffu" if T == 32 else f"((1u << {T}) - 1) << (threadIdx.x & 31 & ~{T - 1})" ) - 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}) - 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] - ) source = ( PRELUDE + noise_source(self._noise, layout.output_words, layout.num_outputs) - + functions + + parts.functions + _KERNELS.format( block=block, T1=T - 1, @@ -449,34 +406,29 @@ def _module(self, mask: np.ndarray | None) -> _Module: 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, # NOTE: the rows keep every output, the postselected ones too. NB=max(-(-layout.num_outputs // 8), 1), - calls=calls, - # NOTE: each component's stream, from the call's seed. - seeds="\n".join( - f" const u64 seed_{i} = mix64(seed + {i + 1} * 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, ) ) module = self._compile(source) - tape = module.get_function("tape") + tape_fn = module.get_function("tape") blocks_per_sm = 4 if hasattr(cp.cuda.runtime, "occupancyMaxActiveBlocksPerMultiprocessor"): blocks_per_sm = max( 1, cp.cuda.runtime.occupancyMaxActiveBlocksPerMultiprocessor( - tape.ptr, block, 0 + tape_fn.ptr, block, 0 ), ) compiled = _Module( noise=module.get_function("noise"), noise_kept=module.get_function("noise_kept"), - tape=tape, + tape=tape_fn, gather_columns=module.get_function("gather_columns"), gather_float=module.get_function("gather_float"), publish=module.get_function("publish"), @@ -486,8 +438,8 @@ def _module(self, mask: np.ndarray | None) -> _Module: tape_blocks=runtime.multiprocessors() * blocks_per_sm, has_mask=mask is not None, plan=plan, - psm_words=psm_words, - passing=self._passing(plan), + psm_words=parts.psm_words, + passing=screen_pass_rate(self._channels, layout, plan), splits={}, ) self._modules[key] = compiled @@ -504,13 +456,6 @@ def _upload(self, array: np.ndarray) -> Any: self._cp.cuda.get_current_stream().synchronize() return device - def _passing(self, plan: ShotPlan) -> float: - """Return the estimated share of shots that pass the screen of ``plan``.""" - layout = self.layout - screen = np.zeros(layout.num_words, dtype=np.uint64) - screen[: layout.output_words] = plan.screen - return pass_rate(self._channels, layout.noise, screen) if screen.any() else 1.0 - def _launch(self, groups: NoiseGroups) -> _Launch: """Upload a noise launch's segments and records.""" return _Launch( @@ -527,36 +472,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. Each key word is read - whole, so the deferred channels flip output words only. + As ``gpu.noise.split``; with ``expectations`` the observables' bits are + read too. Each key word is read whole, so the deferred channels flip + output words only. """ if not module.has_mask: 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 + 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] def _compile(self, source: str): """Compile with NVRTC; ptxas -O3 fails on some shuffled tapes, so lower levels follow.""" @@ -749,12 +685,6 @@ def _list(self, module: _Module, bits, shots: int, count: int, slot: int) -> Any module.list_bits((blocks,), (1024,), (bits, np.int64(words), totals, index)) return index[:count] - def _deviation(self, counts: np.ndarray) -> float: - """Return a batch's normalization deviation beyond float32 rounding, else 0.""" - bits = np.asarray(counts[2:3], dtype=np.uint32) - deviation = float(bits.view(np.float32)[0]) - return deviation if deviation > DEVIATION_TOLERANCE else 0.0 - def sample( self, shots: int, @@ -985,7 +915,7 @@ def finish() -> float: held.append(dst) else: values_host[:] = 0 - return self._deviation(host_counts) + return tape.deviation(host_counts[2]) # NOTE: each chunk is launched before the previous one is finished, so its # kernels overlap that copy. @@ -1032,12 +962,7 @@ def _forked() -> bool: """Return whether this process was forked after its parent used CUDA, with a warning.""" if not runtime.forked_after_init(): return False - 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(FORKED, stacklevel=5) return True diff --git a/test/integration/test_cuda_sampler.py b/test/integration/test_cuda_sampler.py index 74e3ea16..7809402d 100644 --- a/test/integration/test_cuda_sampler.py +++ b/test/integration/test_cuda_sampler.py @@ -13,7 +13,7 @@ from tsim.circuit import Circuit from tsim.statevector.cuda import available, codegen, runtime, sampler as cuda_sampler from tsim.statevector.cuda.sampler import CudaShotSampler -from tsim.statevector.gpu import tape +from tsim.statevector.gpu import noise as gpu_noise, tape from tsim.statevector.program import StatevectorComponent from tsim.statevector.shots import ShotSampler @@ -232,7 +232,7 @@ def test_cuda_samples_no_shots(): def test_cuda_postselection_drops_the_rows_that_fail(monkeypatch): # NOTE: the masked run must draw the unmasked run's noise, so none is deferred. - monkeypatch.setattr(cuda_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) @@ -318,7 +318,7 @@ def test_cuda_chunks_return_the_rows_of_one_pass(monkeypatch, text): # NOTE: 1000-shot chunks on both streams, and host rows that start with room # for the first chunk only; without deferral, a masked and an unmasked run on # the same chunks draw the same shots. - monkeypatch.setattr(cuda_sampler, "DEFER_SHARE", 2.0) + monkeypatch.setattr(gpu_noise, "DEFER_SHARE", 2.0) monkeypatch.setattr( cuda_sampler.CudaShotSampler, "_chunk_shots", lambda self, *args: 1000 ) @@ -526,7 +526,7 @@ def test_cuda_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) > cuda_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 From bce660639344539c6164a14ed750e48bde0b0042 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 06:27:36 -0700 Subject: [PATCH 13/20] refactor(cuda): let the CUDA package choose its sampler tsim/sampler.py knew CUDA's width limit, its fork state, its out-of-memory errors and the no-words rule; cuda.shot_sampler now decides, and sampler.py keeps only the TSIM_DEVICE check and one call. ShotSampler says device = "cpu" itself. Co-Authored-By: Claude Opus 5.5 --- src/tsim/sampler.py | 67 ++++----------------------- src/tsim/statevector/cuda/__init__.py | 67 ++++++++++++++++++++++++++- src/tsim/statevector/shots.py | 3 ++ 3 files changed, 78 insertions(+), 59 deletions(-) diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index e6bde5a0..55ceae12 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.cuda import runtime as cuda_runtime -from tsim.statevector.gpu import FORKED +from tsim.statevector import cuda from tsim.statevector.program import StatevectorComponent from tsim.statevector.shots import ShotSampler from tsim.utils.cuda_helpers import copy_d2h @@ -315,59 +314,6 @@ class SamplerStats: device: str = "cpu" -def _gpu_shots( - host: ShotSampler, - channels: ChannelSampler, - num_detectors: int, - seed: np.random.SeedSequence, -) -> ShotSampler | CudaShotSampler: - """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 cuda_runtime.forked_after_init(): - warnings.warn(FORKED, stacklevel=5) - return host - if not cuda_runtime.available(): - return host - from tsim.statevector.cuda.sampler import CudaShotSampler - from tsim.statevector.gpu.tape import MAX_WIDTH - - layout = host.layout - if layout.num_words == 0: - return host - if not CudaShotSampler.supports(layout): - width = max(c.width for c in layout.components) - warnings.warn( - f"A CUDA 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 CudaShotSampler(host, channels, num_detectors, seed) - 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, - ) - except Exception as error: - if not cuda_runtime.out_of_memory(error): - raise - 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.""" @@ -438,8 +384,13 @@ def __init__( self._num_detectors, shots_seed, ) - 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 = cuda.shot_sampler( self._shots, self._channel_sampler, self._num_detectors, cuda_seed ) else: @@ -823,7 +774,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/cuda/__init__.py b/src/tsim/statevector/cuda/__init__.py index e84ec1e8..fa2fcf89 100644 --- a/src/tsim/statevector/cuda/__init__.py +++ b/src/tsim/statevector/cuda/__init__.py @@ -8,6 +8,71 @@ CuPy and an NVIDIA GPU (``pip install bloqade-tsim[cuda]``). """ +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +from tsim.statevector.cuda import runtime from tsim.statevector.cuda.runtime import available +from tsim.statevector.gpu import FORKED +from tsim.statevector.gpu.tape import MAX_WIDTH, supports + +if TYPE_CHECKING: + import numpy as np + + from tsim.noise.channels import ChannelSampler + from tsim.statevector.cuda.sampler import CudaShotSampler + from tsim.statevector.shots import ShotSampler + +__all__ = ["available", "shot_sampler"] + +# 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 | CudaShotSampler: + """Return a CUDA shot sampler of ``host``'s program when one can run it, else ``host``. + + It warns when a CUDA 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"A CUDA 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.cuda.sampler import CudaShotSampler -__all__ = ["available"] + try: + return CudaShotSampler(host, channels, num_detectors, seed) + 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, + ) + except Exception as error: + if not runtime.out_of_memory(error): + raise + warnings.warn( + "The GPU is out of memory; sampling on the CPU.", stacklevel=_STACKLEVEL + ) + return host 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 23841cdedc9a058bf876c7e820aa919407e04881 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 06:28:41 -0700 Subject: [PATCH 14/20] refactor(cuda): name the thread counts and the counts words The kernels' __launch_bounds__ and their launches spelled out 128, 256 and 1024 separately, and the counts buffer was sliced by bare indices. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/cuda/noise.py | 8 +++- src/tsim/statevector/cuda/sampler.py | 68 ++++++++++++++++++---------- 2 files changed, 49 insertions(+), 27 deletions(-) diff --git a/src/tsim/statevector/cuda/noise.py b/src/tsim/statevector/cuda/noise.py index f90b9e65..89ca10c7 100644 --- a/src/tsim/statevector/cuda/noise.py +++ b/src/tsim/statevector/cuda/noise.py @@ -20,6 +20,9 @@ from tsim.noise.channels import ChannelSampler, WordTables from tsim.statevector.gpu.noise import NoiseTables, noise_tables as _noise_tables +THREADS = 256 +"""Threads per block of the noise kernels.""" + SHARED_BYTES = 46 * 1024 """Shared memory of a block's words and events (the events' share is set per launch).""" @@ -137,7 +140,7 @@ __syncthreads(); }} -extern "C" __global__ void __launch_bounds__(256) noise( +extern "C" __global__ void __launch_bounds__({NT}) noise( u64* __restrict__ out, long long shots, u64 seed, int num_segments, int max_events, const int* __restrict__ seg_first, const int* __restrict__ seg_outcomes, const int* __restrict__ seg_begin, const int* __restrict__ seg_end, @@ -197,7 +200,7 @@ // The deferred channels' faults on the ``n`` kept shots ``index``, then their // returned outputs bit-packed (``pack_rows``). The deferred channels flip output // words only, so a block loads just those. -extern "C" __global__ void __launch_bounds__(256) noise_kept( +extern "C" __global__ void __launch_bounds__({NT}) noise_kept( const u64* __restrict__ out, long long n, const unsigned* __restrict__ index, u64 seed, int num_segments, int max_events, const int* __restrict__ seg_first, const int* __restrict__ seg_outcomes, @@ -246,6 +249,7 @@ def source(tables: NoiseTables, output_words: int, num_outputs: int) -> str: """ packed_bytes = -(-num_outputs // 8) return SOURCE.format( + NT=THREADS, NW=tables.num_words, OW=output_words, S=tables.block_shots, diff --git a/src/tsim/statevector/cuda/sampler.py b/src/tsim/statevector/cuda/sampler.py index 5041a1bc..a99a9537 100644 --- a/src/tsim/statevector/cuda/sampler.py +++ b/src/tsim/statevector/cuda/sampler.py @@ -39,12 +39,28 @@ from tsim.noise.channels import ChannelSampler from tsim.statevector.cuda import runtime from tsim.statevector.cuda.codegen import CUDA, PRELUDE -from tsim.statevector.cuda.noise import noise_tables, source as noise_source +from tsim.statevector.cuda.noise import ( + THREADS as NOISE_THREADS, + noise_tables, + source as noise_source, +) 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 from tsim.statevector.shots import ShotPlan, ShotSampler +TAPE_THREADS = 128 +"""Threads per block of the tape kernel.""" + +COPY_THREADS = 256 +"""Threads per block of the kernels that gather rows and expectations.""" + +SCAN_THREADS = 1024 +"""Threads per block of ``count_bits`` and ``list_bits``: one per bitmap word.""" + +# NOTE: the words of a batch's ``counts``. +_SURVIVORS, _ACCEPTED, _DEVIATION = range(3) + KERNEL_SHARED_BYTES = 48 * 1024 """Shared memory that a kernel may use without opting in, static and dynamic.""" @@ -108,7 +124,7 @@ }} // The shots of the ``words`` words of ``bits`` set, in order, as indices into -// ``index``: ``count_bits`` counts each block's 1024 words, ``list_bits`` lists +// ``index``: ``count_bits`` counts each block's {SCAN} words, ``list_bits`` lists // a block's shots after those of the blocks before it. __device__ __forceinline__ unsigned block_scan(unsigned v, unsigned* total) {{ __shared__ unsigned warps[32]; @@ -131,21 +147,21 @@ return before; }} -extern "C" __global__ void __launch_bounds__(1024) count_bits( +extern "C" __global__ void __launch_bounds__({SCAN}) count_bits( const unsigned* __restrict__ bits, long long words, unsigned* __restrict__ totals) {{ - long long i = (long long)blockIdx.x * 1024 + threadIdx.x; + long long i = (long long)blockIdx.x * {SCAN} + threadIdx.x; unsigned total; block_scan(i < words ? __popc(bits[i]) : 0, &total); if (threadIdx.x == 0) totals[blockIdx.x] = total; }} -extern "C" __global__ void __launch_bounds__(1024) list_bits( +extern "C" __global__ void __launch_bounds__({SCAN}) list_bits( const unsigned* __restrict__ bits, long long words, const unsigned* __restrict__ totals, unsigned* __restrict__ index) {{ unsigned earlier = 0, total; for (int b = threadIdx.x; b < blockIdx.x; b += blockDim.x) earlier += totals[b]; block_scan(earlier, &earlier); - long long i = (long long)blockIdx.x * 1024 + threadIdx.x; + long long i = (long long)blockIdx.x * {SCAN} + threadIdx.x; unsigned w = i < words ? bits[i] : 0; unsigned at = earlier + block_scan(__popc(w), &total); for (; w; w &= w - 1) index[at++] = (unsigned)(32 * i + __ffs(w) - 1); @@ -153,7 +169,7 @@ // The outputs of rows ``index[j]`` (or every row), bit-packed densely // (``pack_rows``): a block packs {S} rows' output words. -extern "C" __global__ void __launch_bounds__(256) gather_columns( +extern "C" __global__ void __launch_bounds__({COPY}) gather_columns( const u64* __restrict__ rows, long long n, const unsigned* __restrict__ index, unsigned char* __restrict__ dst) {{ __shared__ u64 sw[{OW} * {S}]; @@ -387,7 +403,6 @@ def _module(self, mask: np.ndarray | None) -> _Module: layout = self.layout tb = tape.threads_per_shot(layout) T = 1 << tb - block = 128 parts = tape.tape_kernel(layout, plan, mask, self._num_detectors, tb, CUDA) gmask = ( "0xffffffffu" @@ -399,7 +414,9 @@ def _module(self, mask: np.ndarray | None) -> _Module: + noise_source(self._noise, layout.output_words, layout.num_outputs) + parts.functions + _KERNELS.format( - block=block, + block=TAPE_THREADS, + SCAN=SCAN_THREADS, + COPY=COPY_THREADS, T1=T - 1, tb=tb, gmask=gmask, @@ -422,7 +439,7 @@ def _module(self, mask: np.ndarray | None) -> _Module: blocks_per_sm = max( 1, cp.cuda.runtime.occupancyMaxActiveBlocksPerMultiprocessor( - tape_fn.ptr, block, 0 + tape_fn.ptr, TAPE_THREADS, 0 ), ) compiled = _Module( @@ -581,7 +598,7 @@ def _run_batch( S = self._noise.block_shots module.noise( ((shots + S - 1) // S,), - (256,), + (NOISE_THREADS,), ( words, np.int64(shots), @@ -593,7 +610,7 @@ def _run_batch( constant, module.screen, queue, - counts[0:1], + counts[_SURVIVORS : _SURVIVORS + 1], accepted if module.has_mask else np.uint64(0), ), shared_mem=8 * launch.max_events, @@ -604,22 +621,22 @@ def _run_batch( if layout.components: module.tape( (module.tape_blocks,), - (128,), + (TAPE_THREADS,), ( words, np.int64(first_shot), seed, queue, - counts[0:1], + counts[_SURVIVORS : _SURVIVORS + 1], accepted if module.has_mask else np.uint64(0), - counts[1:2], + counts[_ACCEPTED : _ACCEPTED + 1], values if values is not None else np.uint64(0), - counts[2:3], + counts[_DEVIATION : _DEVIATION + 1], ), ) elif module.has_mask: # NOTE: without components the screened survivors are the accepted shots. - counts[1:2] = counts[0:1] + counts[_ACCEPTED : _ACCEPTED + 1] = counts[_SURVIVORS : _SURVIVORS + 1] return words, accepted, counts, values def _pack( @@ -647,7 +664,7 @@ def _pack( seed = np.uint64(self._rng.integers(0, 2**64, dtype=np.uint64)) module.noise_kept( (blocks,), - (256,), + (NOISE_THREADS,), ( words, np.int64(count), @@ -664,7 +681,7 @@ def _pack( else: module.gather_columns( (blocks,), - (256,), + (COPY_THREADS,), ( words, np.int64(count), @@ -678,11 +695,12 @@ def _pack( def _list(self, module: _Module, bits, shots: int, count: int, slot: int) -> Any: """Return the ``count`` shots of the bitmap ``bits`` of ``shots`` shots, in order.""" words = -(-shots // 32) - blocks = -(-words // 1024) + blocks = -(-words // SCAN_THREADS) totals = self._buffer(f"totals{slot}", (blocks,), np.uint32) index = self._buffer(f"index{slot}", (max(count, 1),), np.uint32) - module.count_bits((blocks,), (1024,), (bits, np.int64(words), totals)) - module.list_bits((blocks,), (1024,), (bits, np.int64(words), totals, index)) + grid, threads = (blocks,), (SCAN_THREADS,) + module.count_bits(grid, threads, (bits, np.int64(words), totals)) + module.list_bits(grid, threads, (bits, np.int64(words), totals, index)) return index[:count] def sample( @@ -875,7 +893,7 @@ def finish() -> float: ready.synchronize() host_counts = self._published[slot][0].copy() with self._streams[slot]: - count = n if mask is None else int(host_counts[1]) + count = n if mask is None else int(host_counts[_ACCEPTED]) if first == 0: out.allocate(count, n, shots) elif out.total + count > out.capacity: @@ -906,7 +924,7 @@ def finish() -> float: dst = cp.empty((count, no), dtype=np.float32) module.gather_float( (runtime.multiprocessors() * 8,), - (256,), + (COPY_THREADS,), (values, np.int64(count), np.int32(no), index, dst), ) if copy is not stream: @@ -915,7 +933,7 @@ def finish() -> float: held.append(dst) else: values_host[:] = 0 - return tape.deviation(host_counts[2]) + return tape.deviation(host_counts[_DEVIATION]) # NOTE: each chunk is launched before the previous one is finished, so its # kernels overlap that copy. From 1d3b8eea591ef7e40ffd24daf58a2cab4b063673 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 06:29:44 -0700 Subject: [PATCH 15/20] refactor(cuda): move the kernels' CUDA C++ to cuda/kernels.py The kernel template and the source assembly leave sampler.py, so _module only plans, compiles and looks up the kernels. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/cuda/kernels.py | 190 +++++++++++++++++++++++++++ src/tsim/statevector/cuda/sampler.py | 176 +------------------------ 2 files changed, 196 insertions(+), 170 deletions(-) create mode 100644 src/tsim/statevector/cuda/kernels.py diff --git a/src/tsim/statevector/cuda/kernels.py b/src/tsim/statevector/cuda/kernels.py new file mode 100644 index 00000000..c47b8bee --- /dev/null +++ b/src/tsim/statevector/cuda/kernels.py @@ -0,0 +1,190 @@ +"""The CUDA 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_columns`` and +``gather_float`` gather rows and expectations for the copy to the host; +``publish`` stores a batch's counts into pinned host memory. +""" + +from __future__ import annotations + +from tsim.statevector.cuda.codegen import PRELUDE +from tsim.statevector.cuda.noise import source as noise_source +from tsim.statevector.gpu.noise import NoiseTables +from tsim.statevector.gpu.tape import TapeKernel +from tsim.statevector.shots import ShotLayout + +TAPE_THREADS = 128 +"""Threads per block of the tape kernel.""" + +COPY_THREADS = 256 +"""Threads per block of the kernels that gather rows and expectations.""" + +SCAN_THREADS = 1024 +"""Threads per block of ``count_bits`` and ``list_bits``: one per bitmap word.""" + +TEMPLATE = r""" +extern "C" __global__ void __launch_bounds__({block}) tape( + u64* __restrict__ out, long long first_shot, u64 seed, + const unsigned* __restrict__ queue, const unsigned* __restrict__ count, + unsigned* __restrict__ survivors, unsigned* __restrict__ acc_count, + float* __restrict__ expectations, unsigned* __restrict__ deviation) +{{ + int t = threadIdx.x & {T1}; + unsigned gmask = {gmask}; +{seeds} + long long groups = ((long long)gridDim.x * blockDim.x) >> {tb}; + long long n = *count; + // NOTE: whole warps leave the loop together, so the shuffles keep their groups. + long long rounds = (n + groups - 1) / groups; + long long g0 = ((long long)blockIdx.x * blockDim.x + threadIdx.x) >> {tb}; + for (long long r = 0; r < rounds; ++r) {{ + long long q = g0 + r * groups; + bool valid = q < n; + u64 row_index = queue[valid ? q : n - 1]; + u64 shot = first_shot + row_index; + u64* row = out + row_index * {NW}; + // NOTE: the words that no component reads are dead loads, dropped. + u64 w[{NW}]; + #pragma unroll + for (int i = 0; i < {NW}; ++i) w[i] = row[i]; + float ev[{NO}]; + int grave = 0x7fffffff; +{calls} + if (valid && t == 0) {{ +{stores} + if (expectations != nullptr) {{ + #pragma unroll + for (int i = 0; i < {NO}; ++i) expectations[row_index * {NO} + i] = ev[i]; + }} + }} + if (survivors != nullptr) {{ + // NOTE: the postselected direct outputs passed the screen already; + // the shots that fail now leave the screen's bitmap. + u64 hit = {hits}; + if (valid && t == 0 && hit) atomicAnd(&survivors[row_index >> 5], ~(1u << (row_index & 31))); + unsigned ballot = __ballot_sync(0xffffffffu, valid && t == 0 && hit == 0); + if (ballot && (threadIdx.x & 31) == __ffs(ballot) - 1) atomicAdd(acc_count, __popc(ballot)); + }} + }} +}} + +// The shots of the ``words`` words of ``bits`` set, in order, as indices into +// ``index``: ``count_bits`` counts each block's {SCAN} words, ``list_bits`` lists +// a block's shots after those of the blocks before it. +__device__ __forceinline__ unsigned block_scan(unsigned v, unsigned* total) {{ + __shared__ unsigned warps[32]; + int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + unsigned x = v; + #pragma unroll + for (int o = 1; o < 32; o <<= 1) {{ unsigned y = __shfl_up_sync(0xffffffffu, x, o); if (lane >= o) x += y; }} + if (lane == 31) warps[warp] = x; + __syncthreads(); + if (warp == 0) {{ + unsigned w = lane < (blockDim.x >> 5) ? warps[lane] : 0; + #pragma unroll + for (int o = 1; o < 32; o <<= 1) {{ unsigned y = __shfl_up_sync(0xffffffffu, w, o); if (lane >= o) w += y; }} + warps[lane] = w; + }} + __syncthreads(); + unsigned before = (warp ? warps[warp - 1] : 0) + x - v; + *total = warps[31]; + __syncthreads(); + return before; +}} + +extern "C" __global__ void __launch_bounds__({SCAN}) count_bits( + const unsigned* __restrict__ bits, long long words, unsigned* __restrict__ totals) {{ + long long i = (long long)blockIdx.x * {SCAN} + threadIdx.x; + unsigned total; + block_scan(i < words ? __popc(bits[i]) : 0, &total); + if (threadIdx.x == 0) totals[blockIdx.x] = total; +}} + +extern "C" __global__ void __launch_bounds__({SCAN}) list_bits( + const unsigned* __restrict__ bits, long long words, const unsigned* __restrict__ totals, + unsigned* __restrict__ index) {{ + unsigned earlier = 0, total; + for (int b = threadIdx.x; b < blockIdx.x; b += blockDim.x) earlier += totals[b]; + block_scan(earlier, &earlier); + long long i = (long long)blockIdx.x * {SCAN} + threadIdx.x; + unsigned w = i < words ? bits[i] : 0; + unsigned at = earlier + block_scan(__popc(w), &total); + for (; w; w &= w - 1) index[at++] = (unsigned)(32 * i + __ffs(w) - 1); +}} + +// The outputs of rows ``index[j]`` (or every row), bit-packed densely +// (``pack_rows``): a block packs {S} rows' output words. +extern "C" __global__ void __launch_bounds__({COPY}) gather_columns( + const u64* __restrict__ rows, long long n, const unsigned* __restrict__ index, + unsigned char* __restrict__ dst) {{ + __shared__ u64 sw[{OW} * {S}]; + extern __shared__ __align__(16) unsigned char stage[]; + long long base = (long long)blockIdx.x * {S}; + int ns = (int)min((long long){S}, n - base); + for (int i = threadIdx.x; i < ns * {OW}; i += blockDim.x) {{ + int s = i / {OW}, w = i - s * {OW}; + long long r = index != nullptr ? (long long)index[base + s] : base + s; + sw[w * {S} + s] = rows[r * {NW} + w]; + }} + __syncthreads(); + pack_rows(sw, ns, stage, dst + base * {NB}); +}} + +// The batch's counts, stored into pinned host memory: a read that does not +// queue behind the copies to the host. +extern "C" __global__ void publish(const unsigned* __restrict__ counts, volatile unsigned* host) {{ + if (threadIdx.x < 3) host[threadIdx.x] = counts[threadIdx.x]; + __threadfence_system(); +}} + +// Observable expectations of rows ``index[j]``, densely. +extern "C" __global__ void gather_float(const float* __restrict__ values, long long n, int width, + const unsigned* __restrict__ index, float* __restrict__ dst) {{ + for (long long k = (long long)blockIdx.x * blockDim.x + threadIdx.x; k < n * width; + k += (long long)gridDim.x * blockDim.x) {{ + long long j = k / width; int c = (int)(k - j * width); + dst[k] = values[(long long)index[j] * width + c]; + }} +}} +""" + + +def source( + layout: ShotLayout, + tables: NoiseTables, + parts: TapeKernel, + tb: int, + num_observables: int, +) -> str: + """Return the CUDA source of every kernel, for tapes of ``2^tb`` threads a shot.""" + T = 1 << tb + # NOTE: a shot's lanes within its warp. + gmask = ( + "0xffffffffu" + if T == 32 + else f"((1u << {T}) - 1) << (threadIdx.x & 31 & ~{T - 1})" + ) + return ( + PRELUDE + + noise_source(tables, layout.output_words, layout.num_outputs) + + parts.functions + + TEMPLATE.format( + block=TAPE_THREADS, + SCAN=SCAN_THREADS, + COPY=COPY_THREADS, + T1=T - 1, + tb=tb, + gmask=gmask, + NW=layout.num_words, + OW=layout.output_words, + NO=num_observables, + S=tables.block_shots, + # NOTE: the rows keep every output, the postselected ones too. + NB=max(-(-layout.num_outputs // 8), 1), + calls=parts.calls, + seeds=parts.seeds, + stores=parts.stores, + hits=parts.hits, + ) + ) diff --git a/src/tsim/statevector/cuda/sampler.py b/src/tsim/statevector/cuda/sampler.py index a99a9537..5af1c6f7 100644 --- a/src/tsim/statevector/cuda/sampler.py +++ b/src/tsim/statevector/cuda/sampler.py @@ -37,27 +37,15 @@ import numpy as np from tsim.noise.channels import ChannelSampler -from tsim.statevector.cuda import runtime -from tsim.statevector.cuda.codegen import CUDA, PRELUDE -from tsim.statevector.cuda.noise import ( - THREADS as NOISE_THREADS, - noise_tables, - source as noise_source, -) +from tsim.statevector.cuda import kernels, runtime +from tsim.statevector.cuda.codegen import CUDA +from tsim.statevector.cuda.kernels import COPY_THREADS, SCAN_THREADS, TAPE_THREADS +from tsim.statevector.cuda.noise import THREADS as NOISE_THREADS, noise_tables 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 from tsim.statevector.shots import ShotPlan, ShotSampler -TAPE_THREADS = 128 -"""Threads per block of the tape kernel.""" - -COPY_THREADS = 256 -"""Threads per block of the kernels that gather rows and expectations.""" - -SCAN_THREADS = 1024 -"""Threads per block of ``count_bits`` and ``list_bits``: one per bitmap word.""" - # NOTE: the words of a batch's ``counts``. _SURVIVORS, _ACCEPTED, _DEVIATION = range(3) @@ -77,132 +65,6 @@ MAX_CHUNKS = 8 """Most chunks of a batch, so that a chunk's fixed costs stay small.""" -_KERNELS = r""" -extern "C" __global__ void __launch_bounds__({block}) tape( - u64* __restrict__ out, long long first_shot, u64 seed, - const unsigned* __restrict__ queue, const unsigned* __restrict__ count, - unsigned* __restrict__ survivors, unsigned* __restrict__ acc_count, - float* __restrict__ expectations, unsigned* __restrict__ deviation) -{{ - int t = threadIdx.x & {T1}; - unsigned gmask = {gmask}; -{seeds} - long long groups = ((long long)gridDim.x * blockDim.x) >> {tb}; - long long n = *count; - // NOTE: whole warps leave the loop together, so the shuffles keep their groups. - long long rounds = (n + groups - 1) / groups; - long long g0 = ((long long)blockIdx.x * blockDim.x + threadIdx.x) >> {tb}; - for (long long r = 0; r < rounds; ++r) {{ - long long q = g0 + r * groups; - bool valid = q < n; - u64 row_index = queue[valid ? q : n - 1]; - u64 shot = first_shot + row_index; - u64* row = out + row_index * {NW}; - // NOTE: the words that no component reads are dead loads, dropped. - u64 w[{NW}]; - #pragma unroll - for (int i = 0; i < {NW}; ++i) w[i] = row[i]; - float ev[{NO}]; - int grave = 0x7fffffff; -{calls} - if (valid && t == 0) {{ -{stores} - if (expectations != nullptr) {{ - #pragma unroll - for (int i = 0; i < {NO}; ++i) expectations[row_index * {NO} + i] = ev[i]; - }} - }} - if (survivors != nullptr) {{ - // NOTE: the postselected direct outputs passed the screen already; - // the shots that fail now leave the screen's bitmap. - u64 hit = {hits}; - if (valid && t == 0 && hit) atomicAnd(&survivors[row_index >> 5], ~(1u << (row_index & 31))); - unsigned ballot = __ballot_sync(0xffffffffu, valid && t == 0 && hit == 0); - if (ballot && (threadIdx.x & 31) == __ffs(ballot) - 1) atomicAdd(acc_count, __popc(ballot)); - }} - }} -}} - -// The shots of the ``words`` words of ``bits`` set, in order, as indices into -// ``index``: ``count_bits`` counts each block's {SCAN} words, ``list_bits`` lists -// a block's shots after those of the blocks before it. -__device__ __forceinline__ unsigned block_scan(unsigned v, unsigned* total) {{ - __shared__ unsigned warps[32]; - int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; - unsigned x = v; - #pragma unroll - for (int o = 1; o < 32; o <<= 1) {{ unsigned y = __shfl_up_sync(0xffffffffu, x, o); if (lane >= o) x += y; }} - if (lane == 31) warps[warp] = x; - __syncthreads(); - if (warp == 0) {{ - unsigned w = lane < (blockDim.x >> 5) ? warps[lane] : 0; - #pragma unroll - for (int o = 1; o < 32; o <<= 1) {{ unsigned y = __shfl_up_sync(0xffffffffu, w, o); if (lane >= o) w += y; }} - warps[lane] = w; - }} - __syncthreads(); - unsigned before = (warp ? warps[warp - 1] : 0) + x - v; - *total = warps[31]; - __syncthreads(); - return before; -}} - -extern "C" __global__ void __launch_bounds__({SCAN}) count_bits( - const unsigned* __restrict__ bits, long long words, unsigned* __restrict__ totals) {{ - long long i = (long long)blockIdx.x * {SCAN} + threadIdx.x; - unsigned total; - block_scan(i < words ? __popc(bits[i]) : 0, &total); - if (threadIdx.x == 0) totals[blockIdx.x] = total; -}} - -extern "C" __global__ void __launch_bounds__({SCAN}) list_bits( - const unsigned* __restrict__ bits, long long words, const unsigned* __restrict__ totals, - unsigned* __restrict__ index) {{ - unsigned earlier = 0, total; - for (int b = threadIdx.x; b < blockIdx.x; b += blockDim.x) earlier += totals[b]; - block_scan(earlier, &earlier); - long long i = (long long)blockIdx.x * {SCAN} + threadIdx.x; - unsigned w = i < words ? bits[i] : 0; - unsigned at = earlier + block_scan(__popc(w), &total); - for (; w; w &= w - 1) index[at++] = (unsigned)(32 * i + __ffs(w) - 1); -}} - -// The outputs of rows ``index[j]`` (or every row), bit-packed densely -// (``pack_rows``): a block packs {S} rows' output words. -extern "C" __global__ void __launch_bounds__({COPY}) gather_columns( - const u64* __restrict__ rows, long long n, const unsigned* __restrict__ index, - unsigned char* __restrict__ dst) {{ - __shared__ u64 sw[{OW} * {S}]; - extern __shared__ __align__(16) unsigned char stage[]; - long long base = (long long)blockIdx.x * {S}; - int ns = (int)min((long long){S}, n - base); - for (int i = threadIdx.x; i < ns * {OW}; i += blockDim.x) {{ - int s = i / {OW}, w = i - s * {OW}; - long long r = index != nullptr ? (long long)index[base + s] : base + s; - sw[w * {S} + s] = rows[r * {NW} + w]; - }} - __syncthreads(); - pack_rows(sw, ns, stage, dst + base * {NB}); -}} - -// The batch's counts, stored into pinned host memory: a read that does not -// queue behind the copies to the host. -extern "C" __global__ void publish(const unsigned* __restrict__ counts, volatile unsigned* host) {{ - if (threadIdx.x < 3) host[threadIdx.x] = counts[threadIdx.x]; - __threadfence_system(); -}} - -// Observable expectations of rows ``index[j]``, densely. -extern "C" __global__ void gather_float(const float* __restrict__ values, long long n, int width, - const unsigned* __restrict__ index, float* __restrict__ dst) {{ - for (long long k = (long long)blockIdx.x * blockDim.x + threadIdx.x; k < n * width; - k += (long long)gridDim.x * blockDim.x) {{ - long long j = k / width; int c = (int)(k - j * width); - dst[k] = values[(long long)index[j] * width + c]; - }} -}} -""" - @dataclass class _Launch: @@ -402,35 +264,9 @@ 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, CUDA) - gmask = ( - "0xffffffffu" - if T == 32 - else f"((1u << {T}) - 1) << (threadIdx.x & 31 & ~{T - 1})" - ) - source = ( - PRELUDE - + noise_source(self._noise, layout.output_words, layout.num_outputs) - + parts.functions - + _KERNELS.format( - block=TAPE_THREADS, - SCAN=SCAN_THREADS, - COPY=COPY_THREADS, - T1=T - 1, - tb=tb, - gmask=gmask, - NW=layout.num_words, - OW=layout.output_words, - NO=max(1, len(self._observables)), - S=self._noise.block_shots, - # NOTE: the rows keep every output, the postselected ones too. - NB=max(-(-layout.num_outputs // 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)) ) module = self._compile(source) tape_fn = module.get_function("tape") From e211b596f0a885c56fe66e442dccf812a092af65 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 06:30:42 -0700 Subject: [PATCH 16/20] test(cuda): one parametrized comparison with cat5 The detector, measurement and looping-tape comparisons repeated one body; the same 24 cases run. Co-Authored-By: Claude Opus 5.5 --- test/integration/test_cuda_sampler.py | 53 ++++++++++----------------- 1 file changed, 20 insertions(+), 33 deletions(-) diff --git a/test/integration/test_cuda_sampler.py b/test/integration/test_cuda_sampler.py index 7809402d..3791582c 100644 --- a/test/integration/test_cuda_sampler.py +++ b/test/integration/test_cuda_sampler.py @@ -91,26 +91,28 @@ def _packed(sampler, shots, mask=None): @pytest.mark.parametrize("name", sorted(CIRCUITS)) -def test_cuda_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_cuda_measurement_sampling_matches_stabilizer_rank(name): +@pytest.mark.parametrize( + ("kind", "looping"), + [("detector", False), ("measurement", False), ("detector", True)], +) +def test_cuda_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) @@ -509,21 +511,6 @@ def test_cuda_shots_start_after_the_common_prefix(width): _assert_matches_cpu(circuit, width) -@pytest.mark.parametrize("name", sorted(CIRCUITS)) -def test_cuda_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 From 70e55b896ea9d8d5cac53dbf2622d41779ae0db1 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Wed, 11 Nov 2026 06:30:53 -0700 Subject: [PATCH 17/20] docs(cuda): shorten the changelog entry and the strategy docstrings The changelog entry described the kernels, and three docstrings repeated a six-line paragraph; the install page keeps the details. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 20 +++++++------------- src/tsim/circuit.py | 22 ++++++++++------------ src/tsim/sampler.py | 11 +++++------ 3 files changed, 22 insertions(+), 31 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 281f2850..0191ba35 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,19 +8,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added -- With `strategy="active_statevector"`, the samplers run on a CUDA GPU when - CuPy sees one, as the stabilizer-rank strategies do with JAX: the noise - sampler and every statevector component run as CUDA kernels that CuPy - compiles with NVRTC, in float32, with no trajectory sharing. Circuits with a - component wider than 11 active qubits stay on the CPU, as does everything - under `TSIM_DEVICE=cpu` or `CUDA_VISIBLE_DEVICES=""` and in processes forked - after their parent used the GPU; a warning says so when a GPU is present, as - it does for noise beyond the GPU tables' limits. `sampler.stats.device` and - the sampler's repr say where it samples. With a - `postselection_mask`, only the shots that pass are copied off the GPU. A - seed's samples on the GPU differ from the CPU's and depend on the batch size. - A call that runs out of GPU memory retries in smaller batches, then samples - on the CPU with a warning. Install with `pip install bloqade-tsim[cuda]`. +- `strategy="active_statevector"` samples on a CUDA GPU, in float32, whenever + CuPy sees one and every component has at most 11 active qubits (install with + `pip install bloqade-tsim[cuda]`). A seed's GPU samples differ from the CPU's and + depend on the batch size; `TSIM_DEVICE=cpu` keeps the samplers on the CPU, and + `SamplerStats.device` and the repr say where they sample. The samplers warn and + fall back to the CPU for wider components, noise beyond the GPU tables' limits, + forked processes and calls that run out of GPU memory even in small batches. - The `jax-cuda12` and `jax-cuda13` extras, for the stabilizer-rank strategies on NVIDIA GPUs with JAX. The `cuda12` and `cuda13` extras of Tsim 1.0 install the same packages. diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index 6e54891f..4103df60 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -954,12 +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 a CUDA GPU - when CuPy sees one 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 the sampler on the CPU, in double precision. The - stabilizer rank strategies are provisional. + "bss", "cutting"). "active_statevector" samples in float32 on a + CUDA 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). @@ -995,12 +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 a CUDA GPU - when CuPy sees one 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 the sampler on the CPU, in double precision. The - stabilizer rank strategies are provisional. + "bss", "cutting"). "active_statevector" samples in float32 on a + CUDA 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 55ceae12..bf071e68 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -335,12 +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 a CUDA GPU - when CuPy sees one 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 the sampler on the CPU, in double precision. The - stabilizer rank strategies are provisional. + "bss", "cutting"). "active_statevector" samples in float32 on a + CUDA 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 75621386a163566767ea9cfe2dd939040820e8b3 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Thu, 12 Nov 2026 15:08:16 -0700 Subject: [PATCH 18/20] feat(cuda): sample with the device argument and Metal's fixes CUDA now behaves as the Metal backend does on main: - device="cuda" samples on an NVIDIA GPU or raises with the reason, and device="auto" tries Metal, then CUDA, then the CPU, without warnings. Kernel compile failures (runtime.CompileError) move a sampler to the CPU; out of memory, forked processes and pickles follow 'required'. - The noise kernels draw outcomes from 40-bit keep thresholds, hash each block's stream, and keep rows too wide for shared memory in device memory; gather_bytes packs rows too wide to stage, so no output count falls back. - Requested output columns (separate, prepended, detectors only) and booleans are selected on the GPU (select_rows) before the copy to the host. - The reason a layout cannot run on a GPU is shared (gpu.tape.unsupported). Co-Authored-By: Claude Opus 5.5 --- src/tsim/circuit.py | 52 ++-- src/tsim/core/types.py | 2 +- src/tsim/sampler.py | 114 ++++--- src/tsim/statevector/cuda/__init__.py | 77 +++-- src/tsim/statevector/cuda/kernels.py | 74 ++++- src/tsim/statevector/cuda/noise.py | 271 +++++++++++------ src/tsim/statevector/cuda/runtime.py | 39 ++- src/tsim/statevector/cuda/sampler.py | 395 ++++++++++++++++++------- src/tsim/statevector/gpu/tape.py | 16 + src/tsim/statevector/metal/__init__.py | 22 +- test/conftest.py | 13 +- test/integration/test_cuda_sampler.py | 290 ++++++++++++++++-- test/unit/test_sampler.py | 21 +- 13 files changed, 983 insertions(+), 403 deletions(-) diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index d90191b2..2eeb0239 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -966,24 +966,26 @@ def compile_sampler( depends on the free memory. On a GPU, they do depend on the batch size. device: Where "active_statevector" samples. "auto" (the default) - samples in float32 on an Apple GPU when there is one that can - run the circuit, and otherwise on the CPU, without a warning - (``TSIM_DEVICE=cpu`` keeps it on the CPU). "cpu" samples on the - CPU in float64, on ``numba.get_num_threads()`` threads. "metal" - samples on an Apple GPU or raises: without one, for a component - of more than 11 active qubits, a rotation by less than about - 6e-5 π, noise beyond the GPU's tables, or a call that runs out - of GPU memory. The stabilizer rank strategies take only "auto". + samples in float32 on an Apple GPU or an NVIDIA GPU when there + is one that can run the circuit, and otherwise on the CPU, + without a warning (``TSIM_DEVICE=cpu`` keeps it on the CPU). + "cpu" samples on the CPU in float64, on + ``numba.get_num_threads()`` threads. "metal" samples on an Apple + GPU and "cuda" on an NVIDIA GPU (with the ``cuda`` extra), or + raises: without one, for a component of more than 11 active + qubits, a rotation by less than about 6e-5 π, noise beyond the + GPU's tables, or a call that runs out of GPU memory. The + stabilizer rank strategies take only "auto". Returns: A CompiledMeasurementSampler that can be used to sample measurements. Raises: ValueError: If ``device`` is unknown, is not "auto" under a - stabilizer-rank strategy, or is "metal" for a circuit that the - GPU executor cannot run. - RuntimeError: If ``device`` is "metal" and this process cannot use - an Apple GPU. + stabilizer-rank strategy, or is "metal" or "cuda" for a circuit + that the GPU executor cannot run. + RuntimeError: If ``device`` is "metal" or "cuda" and this process + cannot use such a GPU. """ from tsim.sampler import CompiledMeasurementSampler @@ -1021,24 +1023,26 @@ def compile_detector_sampler( settings, since the default batch size depends on the free memory. On a GPU, they do depend on the batch size. device: Where "active_statevector" samples. "auto" (the default) - samples in float32 on an Apple GPU when there is one that can - run the circuit, and otherwise on the CPU, without a warning - (``TSIM_DEVICE=cpu`` keeps it on the CPU). "cpu" samples on the - CPU in float64, on ``numba.get_num_threads()`` threads. "metal" - samples on an Apple GPU or raises: without one, for a component - of more than 11 active qubits, a rotation by less than about - 6e-5 π, noise beyond the GPU's tables, or a call that runs out - of GPU memory. The stabilizer rank strategies take only "auto". + samples in float32 on an Apple GPU or an NVIDIA GPU when there + is one that can run the circuit, and otherwise on the CPU, + without a warning (``TSIM_DEVICE=cpu`` keeps it on the CPU). + "cpu" samples on the CPU in float64, on + ``numba.get_num_threads()`` threads. "metal" samples on an Apple + GPU and "cuda" on an NVIDIA GPU (with the ``cuda`` extra), or + raises: without one, for a component of more than 11 active + qubits, a rotation by less than about 6e-5 π, noise beyond the + GPU's tables, or a call that runs out of GPU memory. The + stabilizer rank strategies take only "auto". Returns: A CompiledDetectorSampler that can be used to sample detectors and observables. Raises: ValueError: If ``device`` is unknown, is not "auto" under a - stabilizer-rank strategy, or is "metal" for a circuit that the - GPU executor cannot run. - RuntimeError: If ``device`` is "metal" and this process cannot use - an Apple GPU. + stabilizer-rank strategy, or is "metal" or "cuda" for a circuit + that the GPU executor cannot run. + RuntimeError: If ``device`` is "metal" or "cuda" and this process + cannot use such a GPU. """ from tsim.sampler import CompiledDetectorSampler diff --git a/src/tsim/core/types.py b/src/tsim/core/types.py index d142170b..efa5533b 100644 --- a/src/tsim/core/types.py +++ b/src/tsim/core/types.py @@ -28,7 +28,7 @@ SamplingStrategy = Literal["active_statevector", "cat5", "bss", "cutting"] """``"active_statevector"``, or a stabilizer-rank decomposition strategy.""" -Device = Literal["auto", "cpu", "metal"] +Device = Literal["auto", "cpu", "metal", "cuda"] """Where ``"active_statevector"`` samples: ``"auto"`` picks a GPU when one can run it.""" diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 325438b0..b8c10a94 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -25,7 +25,7 @@ SamplingStrategy, ) from tsim.noise.channels import ChannelSampler -from tsim.statevector import metal +from tsim.statevector import cuda, metal from tsim.statevector.program import StatevectorComponent from tsim.statevector.shots import ShotSampler from tsim.utils.cuda_helpers import copy_d2h @@ -35,8 +35,12 @@ from jax import Array as PRNGKey from tsim.circuit import Circuit + from tsim.statevector.cuda.sampler import CudaShotSampler from tsim.statevector.metal.sampler import MetalShotSampler +_GPUS = {"metal": metal, "cuda": cuda} +"""The GPU backends of ``device``, in the order that ``device="auto"`` tries them.""" + # NOTE: postselection with a stabilizer-rank strategy screens shots for direct # discards in rounds of SCREEN_CHUNKS chunks of at most SCREEN_SHOTS shots. SCREEN_SHOTS = 2**14 @@ -352,25 +356,29 @@ def __init__( depends on the free memory. On a GPU, they do depend on the batch size. device: Where "active_statevector" samples. "auto" (the default) - samples in float32 on an Apple GPU when there is one that can - run the circuit, and otherwise on the CPU, without a warning - (``TSIM_DEVICE=cpu`` keeps it on the CPU). "cpu" samples on the - CPU in float64, on ``numba.get_num_threads()`` threads. "metal" - samples on an Apple GPU or raises: without one, for a component - of more than 11 active qubits, a rotation by less than about - 6e-5 π, noise beyond the GPU's tables, or a call that runs out - of GPU memory. The stabilizer rank strategies take only "auto". + samples in float32 on an Apple GPU or an NVIDIA GPU when there + is one that can run the circuit, and otherwise on the CPU, + without a warning (``TSIM_DEVICE=cpu`` keeps it on the CPU). + "cpu" samples on the CPU in float64, on + ``numba.get_num_threads()`` threads. "metal" samples on an Apple + GPU and "cuda" on an NVIDIA GPU (with the ``cuda`` extra), or + raises: without one, for a component of more than 11 active + qubits, a rotation by less than about 6e-5 π, noise beyond the + GPU's tables, or a call that runs out of GPU memory. The + stabilizer rank strategies take only "auto". Raises: ValueError: If ``device`` is unknown, is not "auto" under a - stabilizer-rank strategy, or is "metal" for a circuit that the - GPU executor cannot run. - RuntimeError: If ``device`` is "metal" and this process cannot use - an Apple GPU. + stabilizer-rank strategy, or is "metal" or "cuda" for a circuit + that the GPU executor cannot run. + RuntimeError: If ``device`` is "metal" or "cuda" and this process + cannot use such a GPU. """ - if device not in ("auto", "cpu", "metal"): - raise ValueError(f"device must be 'auto', 'cpu' or 'metal', not {device!r}") + if device not in ("auto", "cpu", *_GPUS): + raise ValueError( + f"device must be 'auto', 'cpu', 'metal' or 'cuda', not {device!r}" + ) if device != "auto" and strategy != "active_statevector": raise ValueError( f"device={device!r} needs strategy='active_statevector'; the " @@ -399,7 +407,7 @@ 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 | MetalShotSampler | None = None + self._shots: ShotSampler | MetalShotSampler | CudaShotSampler | None = None self._backend: _StabilizerRankBackend | None = None if strategy == "active_statevector" or not program.components: host = ShotSampler( @@ -414,13 +422,17 @@ def __init__( if device == "auto" and os.environ.get("TSIM_DEVICE", "").lower() == "cpu": device = "cpu" if mode == "sequential" and device != "cpu": - self._shots = metal.shot_sampler( - host, - self._channel_sampler, - self._num_detectors, - gpu_seed, - required=device == "metal", - ) + # NOTE: "auto" takes the first GPU that runs the program. + for name in _GPUS if device == "auto" else (device,): + self._shots = _GPUS[name].shot_sampler( + host, + self._channel_sampler, + self._num_detectors, + gpu_seed, + required=device == name, + ) + if self._shots is not host: + break else: self._backend = _StabilizerRankBackend(program, key_seed) @@ -770,8 +782,8 @@ def stats(self) -> SamplerStats: 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). With ``device="auto"`` it is where + ``"metal"`` on an Apple GPU, ``"cuda"`` on an NVIDIA GPU, else + ``"cpu"`` (a stabilizer-rank strategy runs wherever JAX does). With ``device="auto"`` it is where the sampler ended up: ``"cpu"`` after its GPU kernel failed to compile or in a process forked after its parent used the GPU. A call that runs out of GPU memory samples on the CPU without changing it. @@ -911,21 +923,23 @@ def __init__( depends on the free memory. On a GPU, they do depend on the batch size. device: Where "active_statevector" samples. "auto" (the default) - samples in float32 on an Apple GPU when there is one that can - run the circuit, and otherwise on the CPU, without a warning - (``TSIM_DEVICE=cpu`` keeps it on the CPU). "cpu" samples on the - CPU in float64, on ``numba.get_num_threads()`` threads. "metal" - samples on an Apple GPU or raises: without one, for a component - of more than 11 active qubits, a rotation by less than about - 6e-5 π, noise beyond the GPU's tables, or a call that runs out - of GPU memory. The stabilizer rank strategies take only "auto". + samples in float32 on an Apple GPU or an NVIDIA GPU when there + is one that can run the circuit, and otherwise on the CPU, + without a warning (``TSIM_DEVICE=cpu`` keeps it on the CPU). + "cpu" samples on the CPU in float64, on + ``numba.get_num_threads()`` threads. "metal" samples on an Apple + GPU and "cuda" on an NVIDIA GPU (with the ``cuda`` extra), or + raises: without one, for a component of more than 11 active + qubits, a rotation by less than about 6e-5 π, noise beyond the + GPU's tables, or a call that runs out of GPU memory. The + stabilizer rank strategies take only "auto". Raises: ValueError: If ``device`` is unknown, is not "auto" under a - stabilizer-rank strategy, or is "metal" for a circuit that the - GPU executor cannot run. - RuntimeError: If ``device`` is "metal" and this process cannot use - an Apple GPU. + stabilizer-rank strategy, or is "metal" or "cuda" for a circuit + that the GPU executor cannot run. + RuntimeError: If ``device`` is "metal" or "cuda" and this process + cannot use such a GPU. """ super().__init__( @@ -1030,21 +1044,23 @@ def __init__( settings, since the default batch size depends on the free memory. On a GPU, they do depend on the batch size. device: Where "active_statevector" samples. "auto" (the default) - samples in float32 on an Apple GPU when there is one that can - run the circuit, and otherwise on the CPU, without a warning - (``TSIM_DEVICE=cpu`` keeps it on the CPU). "cpu" samples on the - CPU in float64, on ``numba.get_num_threads()`` threads. "metal" - samples on an Apple GPU or raises: without one, for a component - of more than 11 active qubits, a rotation by less than about - 6e-5 π, noise beyond the GPU's tables, or a call that runs out - of GPU memory. The stabilizer rank strategies take only "auto". + samples in float32 on an Apple GPU or an NVIDIA GPU when there + is one that can run the circuit, and otherwise on the CPU, + without a warning (``TSIM_DEVICE=cpu`` keeps it on the CPU). + "cpu" samples on the CPU in float64, on + ``numba.get_num_threads()`` threads. "metal" samples on an Apple + GPU and "cuda" on an NVIDIA GPU (with the ``cuda`` extra), or + raises: without one, for a component of more than 11 active + qubits, a rotation by less than about 6e-5 π, noise beyond the + GPU's tables, or a call that runs out of GPU memory. The + stabilizer rank strategies take only "auto". Raises: ValueError: If ``device`` is unknown, is not "auto" under a - stabilizer-rank strategy, or is "metal" for a circuit that the - GPU executor cannot run. - RuntimeError: If ``device`` is "metal" and this process cannot use - an Apple GPU. + stabilizer-rank strategy, or is "metal" or "cuda" for a circuit + that the GPU executor cannot run. + RuntimeError: If ``device`` is "metal" or "cuda" and this process + cannot use such a GPU. """ super().__init__( diff --git a/src/tsim/statevector/cuda/__init__.py b/src/tsim/statevector/cuda/__init__.py index fa2fcf89..4c0aa739 100644 --- a/src/tsim/statevector/cuda/__init__.py +++ b/src/tsim/statevector/cuda/__init__.py @@ -1,22 +1,19 @@ """The statevector executor on a CUDA GPU. -The noise sampler (``noise``) and every component's tape (``codegen``) run as -CUDA kernels that CuPy compiles with NVRTC (``runtime``), from the host code -that the GPU backends share (``tsim.statevector.gpu``); ``CudaShotSampler`` -(``sampler``) replaces ``ShotSampler`` behind the compiled samplers whenever -CuPy sees a GPU and every component fits. Amplitudes are complex64. It needs -CuPy and an NVIDIA GPU (``pip install bloqade-tsim[cuda]``). +The noise sampler (``noise``) and every component's tape (``codegen``, written +by ``tsim.statevector.gpu.tape``) run as CUDA kernels that CuPy compiles with +NVRTC (``runtime``); ``CudaShotSampler`` (``sampler``) replaces ``ShotSampler`` +behind the compiled samplers of ``device="auto"`` or ``"cuda"`` whenever CuPy +sees a GPU and every component fits. Amplitudes are complex64. It needs CuPy +and an NVIDIA GPU (``pip install bloqade-tsim[cuda]``). """ from __future__ import annotations -import warnings from typing import TYPE_CHECKING from tsim.statevector.cuda import runtime -from tsim.statevector.cuda.runtime import available -from tsim.statevector.gpu import FORKED -from tsim.statevector.gpu.tape import MAX_WIDTH, supports +from tsim.statevector.gpu.tape import unsupported if TYPE_CHECKING: import numpy as np @@ -25,54 +22,56 @@ from tsim.statevector.cuda.sampler import CudaShotSampler from tsim.statevector.shots import ShotSampler -__all__ = ["available", "shot_sampler"] - -# 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, + *, + required: bool, ) -> ShotSampler | CudaShotSampler: - """Return a CUDA shot sampler of ``host``'s program when one can run it, else ``host``. + """Return a CUDA shot sampler of ``host``'s program, or ``host`` where CUDA cannot run it. + + Without ``required`` (``device="auto"``) it returns ``host`` silently; with it + (``device="cuda"``) it raises instead. A program without shot words (one + with no outputs; direct outputs are shot-word bits too) has nothing to + sample on a GPU and keeps ``host`` either way. + + Raises: + RuntimeError: With ``required``, if this process cannot use a CUDA GPU: + no CuPy, no device, or forked after its parent used the GPU. + ValueError: With ``required``, if the GPU executor cannot run the program: + a component wider than ``MAX_WIDTH``, a rotation by an angle too small + for float32 amplitudes (``RESOLVABLE``), or noise beyond + the GPU tables' limits. + MemoryError: With ``required``, if the GPU runs out of memory. - It warns when a CUDA 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) + reason = runtime.unavailable() + if reason is not None: + if required: + raise RuntimeError(f"Cannot sample on CUDA: {reason}.") return host - if not runtime.available() or host.layout.num_words == 0: + if 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"A CUDA 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, - ) + reason = unsupported(host.layout) + if reason is not None: + if required: + raise ValueError(f"Cannot sample on CUDA: {reason}.") return host from tsim.statevector.cuda.sampler import CudaShotSampler try: - return CudaShotSampler(host, channels, num_detectors, seed) + return CudaShotSampler(host, channels, num_detectors, seed, required=required) 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, - ) + if required: + raise ValueError(f"Cannot sample on CUDA: {error}.") from error except Exception as error: if not runtime.out_of_memory(error): raise - warnings.warn( - "The GPU is out of memory; sampling on the CPU.", stacklevel=_STACKLEVEL - ) + if required: + raise MemoryError(f"Cannot sample on CUDA: {error}") from error return host diff --git a/src/tsim/statevector/cuda/kernels.py b/src/tsim/statevector/cuda/kernels.py index c47b8bee..ee71e579 100644 --- a/src/tsim/statevector/cuda/kernels.py +++ b/src/tsim/statevector/cuda/kernels.py @@ -2,8 +2,10 @@ ``tape`` runs the components' tapes on the queued shots; ``count_bits`` and ``list_bits`` list the accepted shots from the bitmap; ``gather_columns`` and -``gather_float`` gather rows and expectations for the copy to the host; -``publish`` stores a batch's counts into pinned host memory. +``gather_float`` gather rows and expectations for the copy to the host +(``gather_bytes`` the rows too wide for ``gather_columns``); +``select_rows`` picks the requested outputs out of them, bit-packed or as +booleans; ``publish`` stores a batch's counts into pinned host memory. """ from __future__ import annotations @@ -23,6 +25,21 @@ SCAN_THREADS = 1024 """Threads per block of ``count_bits`` and ``list_bits``: one per bitmap word.""" +KERNEL_SHARED_BYTES = 48 * 1024 +"""Shared memory that a kernel may use without opting in, static and dynamic.""" + + +def gather_shots(layout: ShotLayout, block_shots: int) -> int: + """Return the rows that a ``gather_columns`` block packs, or ``0`` if not even 32 fit. + + A block stages its rows' output words and packed bytes in shared memory: at + most ``block_shots`` rows, a multiple of 32. Rows of which 32 do not fit + take ``gather_bytes`` instead. + """ + row = 8 * layout.output_words + -(-layout.num_outputs // 8) + return min(block_shots, KERNEL_SHARED_BYTES // max(row, 1) // 32 * 32) + + TEMPLATE = r""" extern "C" __global__ void __launch_bounds__({block}) tape( u64* __restrict__ out, long long first_shot, u64 seed, @@ -114,21 +131,58 @@ }} // The outputs of rows ``index[j]`` (or every row), bit-packed densely -// (``pack_rows``): a block packs {S} rows' output words. +// (``pack_rows``): a block packs {GS} rows' output words. extern "C" __global__ void __launch_bounds__({COPY}) gather_columns( const u64* __restrict__ rows, long long n, const unsigned* __restrict__ index, unsigned char* __restrict__ dst) {{ - __shared__ u64 sw[{OW} * {S}]; + __shared__ u64 sw[{OW} * {GS}]; extern __shared__ __align__(16) unsigned char stage[]; - long long base = (long long)blockIdx.x * {S}; - int ns = (int)min((long long){S}, n - base); + long long base = (long long)blockIdx.x * {GS}; + int ns = (int)min((long long){GS}, n - base); for (int i = threadIdx.x; i < ns * {OW}; i += blockDim.x) {{ int s = i / {OW}, w = i - s * {OW}; long long r = index != nullptr ? (long long)index[base + s] : base + s; - sw[w * {S} + s] = rows[r * {NW} + w]; + sw[w * {GS} + s] = rows[r * {NW} + w]; }} __syncthreads(); - pack_rows(sw, ns, stage, dst + base * {NB}); + pack_rows<{GS}>(sw, ns, stage, dst + base * {NB}); +}} + +// As ``gather_columns``, for rows too wide to stage in shared memory: a thread +// per output byte, read from device memory. +extern "C" __global__ void gather_bytes( + const u64* __restrict__ rows, long long n, const unsigned* __restrict__ index, + unsigned char* __restrict__ dst) {{ + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n * {NB}; + i += (long long)gridDim.x * blockDim.x) {{ + long long j = i / {NB}; + int b = (int)(i - j * {NB}); + long long r = index != nullptr ? (long long)index[j] : j; + unsigned v = (unsigned)(rows[r * {NW} + (b >> 3)] >> (8 * (b & 7))) & 0xFF; + dst[i] = (unsigned char)(b == {NB} - 1 ? v & {LAST_BYTE} : v); + }} +}} + +// Outputs ``columns[k]`` (``k < count``; output ``k`` without ``columns``) of +// ``n`` packed rows of {NB} bytes, into rows of ``width`` bytes: one byte per +// output with ``bools``, else bit-packed with zero bits past them. A thread +// writes one byte. +extern "C" __global__ void select_rows( + const unsigned char* __restrict__ packed, long long n, const unsigned* __restrict__ columns, + int count, int width, int bools, unsigned char* __restrict__ dst) {{ + int per = bools ? 1 : 8; + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n * width; + i += (long long)gridDim.x * blockDim.x) {{ + long long r = i / width; + int k = (int)(i - r * width) * per; + const unsigned char* row = packed + r * {NB}; + unsigned v = 0; + for (int b = 0; b < per && k + b < count; ++b) {{ + unsigned c = columns != nullptr ? columns[k + b] : (unsigned)(k + b); + v |= ((row[c >> 3] >> (c & 7)) & 1u) << b; + }} + dst[i] = (unsigned char)v; + }} }} // The batch's counts, stored into pinned host memory: a read that does not @@ -179,7 +233,9 @@ def source( NW=layout.num_words, OW=layout.output_words, NO=num_observables, - S=tables.block_shots, + # NOTE: gather_columns goes unused (and small) where 32 rows do not fit. + GS=max(gather_shots(layout, tables.block_shots), 1), + LAST_BYTE=f"0x{(1 << (layout.num_outputs % 8 or 8)) - 1:02X}u", # NOTE: the rows keep every output, the postselected ones too. NB=max(-(-layout.num_outputs // 8), 1), calls=parts.calls, diff --git a/src/tsim/statevector/cuda/noise.py b/src/tsim/statevector/cuda/noise.py index 89ca10c7..381239fe 100644 --- a/src/tsim/statevector/cuda/noise.py +++ b/src/tsim/statevector/cuda/noise.py @@ -5,7 +5,9 @@ word major, walks segments of the channel groups with geometric skips, records the faults as events in shared memory and then XORs their outcomes' flips into the shots' words. Recording first keeps the warps converged while -the faults are applied, and many of their table reads in flight. +the faults are applied, and many of their table reads in flight. Rows too wide +for shared memory 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 their rows (coalesced; only the survivors' unless every row is asked @@ -18,7 +20,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, +) THREADS = 256 """Threads per block of the noise kernels.""" @@ -26,37 +32,130 @@ SHARED_BYTES = 46 * 1024 """Shared memory of a block's words and events (the events' share is set per launch).""" +_NOISE_SIGNATURE = r""" +extern "C" __global__ void __launch_bounds__({NT}) noise( + u64* __restrict__ out, long long shots, u64 seed, int num_segments, int max_events, + const int* __restrict__ seg_first, const int* __restrict__ seg_outcomes, + const int* __restrict__ seg_begin, const int* __restrict__ seg_end, + const float* __restrict__ seg_log_keep, const Record* __restrict__ records, + const Flip* __restrict__ overflow, + const u64* __restrict__ constant, const u64* __restrict__ screen, + unsigned* __restrict__ queue, unsigned* __restrict__ count, unsigned* __restrict__ survivors) +""" + +# NOTE: the noise kernel with the block's words in shared memory, word major. +NOISE_SHARED = _NOISE_SIGNATURE + r"""{{ + __shared__ u64 sw[{NW} * {S}]; + extern __shared__ uint2 events[]; + __shared__ int num_events; + __shared__ bool passed[{S}]; + long long base = (long long)blockIdx.x * {S}; + int ns = (int)min((long long){S}, shots - base); + for (int i = threadIdx.x; i < {NW} * {S}; i += blockDim.x) sw[i] = 0; + if (threadIdx.x == 0) num_events = 0; + __syncthreads(); + add_faults<{S}, 1>(sw, events, &num_events, max_events, ns, block_stream(seed), num_segments, + seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep, records, overflow); + screen_shots<{S}, 1>(sw, ns, base, constant, screen, passed, queue, count, survivors); + __syncthreads(); + // NOTE: only the shots that pass the screen are read again. + u64* rows = out + base * {NW}; + for (int i = threadIdx.x; i < ns * {NW}; i += blockDim.x) {{ + int s = i / {NW}, word = i - s * {NW}; + if (passed[s]) rows[i] = sw[word * {S} + s]; + }} +}} +""" + +# NOTE: the noise kernel for rows too wide for shared memory: the faults are XORed +# into the rows in device memory, and the events have all of the shared memory. +NOISE_DEVICE = _NOISE_SIGNATURE + r"""{{ + extern __shared__ uint2 events[]; + __shared__ int num_events; + __shared__ bool passed[{S}]; + long long base = (long long)blockIdx.x * {S}; + int ns = (int)min((long long){S}, shots - base); + u64* rows = out + base * {NW}; + for (int i = threadIdx.x; i < ns * {NW}; i += blockDim.x) rows[i] = 0; + if (threadIdx.x == 0) num_events = 0; + __syncthreads(); + add_faults<1, {NW}>(rows, events, &num_events, max_events, ns, block_stream(seed), num_segments, + seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep, records, overflow); + screen_shots<1, {NW}>(rows, ns, base, constant, screen, passed, queue, count, survivors); +}} +""" + +NOISE_KEPT = r""" +// The deferred channels' faults on the ``n`` kept shots ``index``, then their +// returned outputs bit-packed (``pack_rows``). The deferred channels flip output +// words only, so a block loads just those. +extern "C" __global__ void __launch_bounds__({NT}) noise_kept( + const u64* __restrict__ out, long long n, const unsigned* __restrict__ index, u64 seed, + int num_segments, int max_events, + const int* __restrict__ seg_first, const int* __restrict__ seg_outcomes, + const int* __restrict__ seg_begin, const int* __restrict__ seg_end, + const float* __restrict__ seg_log_keep, const Record* __restrict__ records, + const Flip* __restrict__ overflow, + unsigned char* __restrict__ dst) +{{ + __shared__ u64 sw[{OW} * {S}]; + extern __shared__ __align__(16) uint2 events[]; + __shared__ int num_events; + long long base = (long long)blockIdx.x * {S}; + int ns = (int)min((long long){S}, n - base); + for (int i = threadIdx.x; i < ns * {OW}; i += blockDim.x) {{ + int s = i / {OW}, word = i - s * {OW}; + sw[word * {S} + s] = out[(long long)index[base + s] * {NW} + word]; + }} + if (threadIdx.x == 0) num_events = 0; + __syncthreads(); + add_faults<{S}, 1>(sw, events, &num_events, max_events, ns, block_stream(seed), num_segments, + seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep, records, overflow); + // NOTE: the events are applied, so their buffer stages the packed rows. + pack_rows<{S}>(sw, ns, (unsigned char*)events, dst + base * {NB}); +}} +""" + SOURCE = r""" -struct __align__(32) Record {{ unsigned alias; unsigned words; unsigned extra; unsigned pad; u64 m0; u64 m1; }}; +struct __align__(32) Record {{ unsigned alias; unsigned words; unsigned extra; unsigned keep; u64 m0; u64 m1; }}; struct __align__(16) Flip {{ unsigned word; unsigned pad; u64 mask; }}; -// 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. +// 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 __device__ __forceinline__ void apply_fault(u64* sw, int row, unsigned rec, unsigned draw, - const Record* __restrict__ records, const Flip* __restrict__ overflow) + u64 stream, const Record* __restrict__ records, const Flip* __restrict__ overflow) {{ Record r = records[rec]; - if ((draw & 0xFFFFFF) >= (r.alias & 0x1FFFFFF)) r = records[rec - (draw >> 24) + (r.alias >> 25)]; + unsigned high = draw & 0xFFFFFF, keep = r.alias & 0x1FFFFFF; + if (high > keep || (high == keep && + (unsigned)(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) atomicXor(&sw[(r.words & 0xFFFF) * {S} + row], r.m0); - if (flips > 1) atomicXor(&sw[(r.words >> 16) * {S} + row], r.m1); + if (flips > 0) atomicXor(&sw[(r.words & 0xFFFF) * WS + row * SS], r.m0); + if (flips > 1) atomicXor(&sw[(r.words >> 16) * WS + row * SS], r.m1); for (int j = 2; j < flips; ++j) {{ Flip f = overflow[(r.extra >> 8) + j - 2]; - atomicXor(&sw[f.word * {S} + row], f.mask); + atomicXor(&sw[f.word * WS + row * SS], f.mask); }} }} -// The outputs of the block's ``ns`` rows ``sw[w * S + s]``, bit-packed into +// The outputs of a block's ``ns`` rows ``sw[w * S + s]``, bit-packed into // {NB} bytes per row (the bits past them cleared), contiguously at ``dst``. // Consecutive threads pack one word of consecutive rows (distinct banks) into // ``stage``, from which the block's rows go out 16 bytes at a time. +template __device__ __forceinline__ void pack_rows(const u64* sw, int ns, unsigned char* stage, unsigned char* __restrict__ dst) {{ - for (int i = threadIdx.x; i < {NK} * {S}; i += blockDim.x) {{ - int k = i / {S}, s = i - k * {S}; + for (int i = threadIdx.x; i < {NK} * S; i += blockDim.x) {{ + int k = i / S, s = i - k * S; if (s >= ns) continue; - u64 v = sw[k * {S} + s]; + u64 v = sw[k * S + s]; if (k == {NK} - 1) v &= {LAST_WORD}; unsigned char* at = stage + s * {NB} + 8 * k; #pragma unroll @@ -68,13 +167,20 @@ for (int i = total / 16 * 16 + threadIdx.x; i < total; i += blockDim.x) dst[i] = stage[i]; }} +// The stream of this block. 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. +__device__ __forceinline__ u64 block_stream(u64 seed) {{ + return mix64(seed ^ ((u64)blockIdx.x * 0xD1B54A32D192ED03ULL)); +}} + // 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 +template __device__ __forceinline__ void add_segment_faults(u64* sw, uint2* events, int* num_events, - int max_events, int ns, u64 state, float lk, int first, int outcomes, int idx, int end, - const Record* __restrict__ records, const Flip* __restrict__ overflow) + int max_events, int ns, u64 stream, u64 state, float lk, int first, int outcomes, int idx, + int end, const Record* __restrict__ records, const Flip* __restrict__ overflow) {{ while (true) {{ state += GOLDEN; @@ -91,27 +197,30 @@ 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. - unsigned long long x = (unsigned long long)((unsigned)r >> 8) * (unsigned)outcomes; - unsigned slot = (unsigned)(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 & 0xFFFFFFFFFFULL) * (unsigned)outcomes; + unsigned slot = (unsigned)(x >> 40); unsigned rec = (unsigned)(first + ch * outcomes) + slot; - unsigned draw = (unsigned)(x & 0xFFFFFF) | (slot << 24); + unsigned draw = (unsigned)((x >> 16) & 0xFFFFFF) | (slot << 24); int e = atomicAdd(num_events, 1); if (e < max_events) events[e] = make_uint2((unsigned)row | (rec << 8), draw); - else apply_fault(sw, row, rec, draw, records, overflow); + else apply_fault(sw, row, rec, draw, stream, records, overflow); }} }} +template __device__ __noinline__ void add_fine_segment_faults(u64* sw, uint2* events, int* num_events, - int max_events, int ns, u64 state, float lk, int first, int outcomes, int idx, int end, - const Record* __restrict__ records, const Flip* __restrict__ overflow) + int max_events, int ns, u64 stream, u64 state, float lk, int first, int outcomes, int idx, + int end, const Record* __restrict__ records, const Flip* __restrict__ 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 them into ``sw``. +// 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``). +template __device__ __forceinline__ void add_faults(u64* sw, uint2* events, int* num_events, int max_events, int ns, u64 stream, int num_segments, const int* __restrict__ seg_first, const int* __restrict__ seg_outcomes, @@ -125,42 +234,31 @@ 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); }} __syncthreads(); int ne = min(*num_events, max_events); for (int e = threadIdx.x; e < ne; e += blockDim.x) {{ 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); }} __syncthreads(); }} -extern "C" __global__ void __launch_bounds__({NT}) noise( - u64* __restrict__ out, long long shots, u64 seed, int num_segments, int max_events, - const int* __restrict__ seg_first, const int* __restrict__ seg_outcomes, - const int* __restrict__ seg_begin, const int* __restrict__ seg_end, - const float* __restrict__ seg_log_keep, const Record* __restrict__ records, - const Flip* __restrict__ overflow, - const u64* __restrict__ constant, const u64* __restrict__ screen, +// Screens the block's ``ns`` shots (words laid out as in ``apply_fault``): XORs +// the constants into their output words, marks the survivors in ``passed`` and +// in the bitmap ``survivors`` (if given), and appends them to ``queue`` after +// the ``count`` queued so far. +template +__device__ __forceinline__ void screen_shots(u64* sw, int ns, long long base, + const u64* __restrict__ constant, const u64* __restrict__ screen, bool* passed, unsigned* __restrict__ queue, unsigned* __restrict__ count, unsigned* __restrict__ survivors) {{ - __shared__ u64 sw[{NW} * {S}]; - extern __shared__ uint2 events[]; - __shared__ int num_events; - __shared__ bool passed[{S}]; - long long base = (long long)blockIdx.x * {S}; - int ns = (int)min((long long){S}, shots - base); - for (int i = threadIdx.x; i < {NW} * {S}; i += blockDim.x) sw[i] = 0; - if (threadIdx.x == 0) num_events = 0; - __syncthreads(); - add_faults(sw, events, &num_events, max_events, ns, seed ^ ((u64)blockIdx.x * 0xD1B54A32D192ED03ULL), - num_segments, seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep, records, - overflow); for (int s0 = 0; s0 < {S}; s0 += blockDim.x) {{ int s = s0 + threadIdx.x; bool alive = false; @@ -168,9 +266,9 @@ u64 hit = 0; #pragma unroll for (int i = 0; i < {OW}; ++i) {{ - u64 v = sw[i * {S} + s] ^ constant[i]; + u64 v = sw[i * WS + s * SS] ^ constant[i]; hit |= v & screen[i]; - sw[i * {S} + s] = v; + sw[i * WS + s * SS] = v; }} alive = !hit; }} @@ -188,45 +286,8 @@ if (alive) queue[b + __popc(ballot & ((1u << lane) - 1))] = (unsigned)(base + s); }} }} - __syncthreads(); - // NOTE: only the shots that pass the screen are read again. - u64* rows = out + base * {NW}; - for (int i = threadIdx.x; i < ns * {NW}; i += blockDim.x) {{ - int s = i / {NW}, word = i - s * {NW}; - if (passed[s]) rows[i] = sw[word * {S} + s]; - }} -}} - -// The deferred channels' faults on the ``n`` kept shots ``index``, then their -// returned outputs bit-packed (``pack_rows``). The deferred channels flip output -// words only, so a block loads just those. -extern "C" __global__ void __launch_bounds__({NT}) noise_kept( - const u64* __restrict__ out, long long n, const unsigned* __restrict__ index, u64 seed, - int num_segments, int max_events, - const int* __restrict__ seg_first, const int* __restrict__ seg_outcomes, - const int* __restrict__ seg_begin, const int* __restrict__ seg_end, - const float* __restrict__ seg_log_keep, const Record* __restrict__ records, - const Flip* __restrict__ overflow, - unsigned char* __restrict__ dst) -{{ - __shared__ u64 sw[{OW} * {S}]; - extern __shared__ __align__(16) uint2 events[]; - __shared__ int num_events; - long long base = (long long)blockIdx.x * {S}; - int ns = (int)min((long long){S}, n - base); - for (int i = threadIdx.x; i < ns * {OW}; i += blockDim.x) {{ - int s = i / {OW}, word = i - s * {OW}; - sw[word * {S} + s] = out[(long long)index[base + s] * {NW} + word]; - }} - if (threadIdx.x == 0) num_events = 0; - __syncthreads(); - add_faults(sw, events, &num_events, max_events, ns, seed ^ ((u64)blockIdx.x * 0xD1B54A32D192ED03ULL), - num_segments, seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep, records, - overflow); - // NOTE: the events are applied, so their buffer stages the packed rows. - pack_rows(sw, ns, (unsigned char*)events, dst + base * {NB}); }} -""" +{NOISE}{KEPT}""" def noise_tables( @@ -234,21 +295,35 @@ def noise_tables( ) -> NoiseTables: """Return the noise kernels' tables, for CUDA's shared memory. + Rows too wide for ``SHARED_BYTES`` keep their words in device memory + (``shared_words`` False), where the noise kernel XORs the faults in. + Raises: - ValueError: If the rows do not fit in ``SHARED_BYTES``, or the noise - exceeds the kernels' limits. + ValueError: If the noise exceeds the kernels' limits. """ - return _noise_tables(channels, flips, num_words, shared_bytes=SHARED_BYTES) + return _noise_tables( + channels, + flips, + num_words, + shared_bytes=SHARED_BYTES, + shared_words=words_fit(num_words, SHARED_BYTES), + ) def source(tables: NoiseTables, output_words: int, num_outputs: int) -> str: - """Return the CUDA source of the ``noise`` and ``noise_kept`` kernels. + """Return the CUDA source of the ``noise`` and ``noise_kept`` kernels, and ``pack_rows``. - ``noise_kept`` packs rows of ``num_outputs`` outputs. + ``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 shared memory. """ packed_bytes = -(-num_outputs // 8) - return SOURCE.format( + if tables.shared_words: + noise, kept = NOISE_SHARED, NOISE_KEPT + else: + noise, kept = NOISE_DEVICE, "" + return SOURCE.replace("{NOISE}{KEPT}", noise + kept).format( NT=THREADS, NW=tables.num_words, OW=output_words, diff --git a/src/tsim/statevector/cuda/runtime.py b/src/tsim/statevector/cuda/runtime.py index 185fd03c..8f7ce0ae 100644 --- a/src/tsim/statevector/cuda/runtime.py +++ b/src/tsim/statevector/cuda/runtime.py @@ -15,6 +15,8 @@ import pathlib from types import ModuleType +from tsim.statevector.gpu import FORKED + @functools.cache def cupy() -> ModuleType: @@ -45,17 +47,24 @@ def _after_fork_in_child() -> None: # NOTE: a child forked after its parent initialized CUDA cannot use CUDA; the # cached answer is read without touching CUDA. global _forked_after_init - if _device_available.cache_info().currsize and _device_available(): + if _device_error.cache_info().currsize and _device_error() is None: _forked_after_init = True def available() -> bool: """Return whether CuPy is installed and sees a CUDA device this process can use. - ``CUDA_VISIBLE_DEVICES=""`` hides the devices, so that the samplers run on the - CPU, as they do in a process forked after its parent used the GPU. + ``CUDA_VISIBLE_DEVICES=""`` hides the devices, as in a process forked after + its parent used the GPU. """ - return not _forked_after_init and _device_available() + return unavailable() is None + + +def unavailable() -> str | None: + """Return why this process cannot sample on CUDA, or ``None`` if it can.""" + if _forked_after_init: + return FORKED + return _device_error() def forked_after_init() -> bool: @@ -64,15 +73,23 @@ def forked_after_init() -> bool: @functools.cache -def _device_available() -> bool: +def _device_error() -> str | None: try: cp = cupy() - return cp.cuda.runtime.getDeviceCount() > 0 - except Exception: - return False + if cp.cuda.runtime.getDeviceCount() > 0: + return None + return "CuPy sees no CUDA device" + except ImportError: + return ( + "CuPy is not installed (install Tsim's cuda extra: " + "pip install 'bloqade-tsim[cuda]')" + ) + except Exception as error: + return f"CuPy sees no usable CUDA device ({error})" -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) @functools.cache @@ -98,3 +115,7 @@ def out_of_memory(error: BaseException) -> bool: # NOTE: CUDA_ERROR_OUT_OF_MEMORY and cudaErrorMemoryAllocation are both 2. errors = (cuda.driver.CUDADriverError, cuda.runtime.CUDARuntimeError) return isinstance(error, errors) and getattr(error, "status", None) == 2 + + +class CompileError(RuntimeError): + """A CUDA kernel that NVRTC does not compile.""" diff --git a/src/tsim/statevector/cuda/sampler.py b/src/tsim/statevector/cuda/sampler.py index 5af1c6f7..13318ad8 100644 --- a/src/tsim/statevector/cuda/sampler.py +++ b/src/tsim/statevector/cuda/sampler.py @@ -14,7 +14,9 @@ noise is at least ``DEFER_SHARE`` of the faults list lists the accepted shots in order from the bitmap (``count_bits``, ``list_bits``) - select gathers the rows (or the accepted rows) for the copy to the host + select gathers the rows (or the accepted rows), bit-packed, and picks the + requested outputs out of them, bit-packed or as booleans, for the + copy to the host ``sample_accepted`` runs a call as chunks of shots on two streams, so that a chunk's copy to the host overlaps the next chunk's kernels. @@ -23,14 +25,13 @@ them on disk). 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. +samples differ from the CPU's, and they depend on the batch size. """ from __future__ import annotations import math import threading -import warnings from dataclasses import dataclass from typing import Any @@ -43,15 +44,11 @@ from tsim.statevector.cuda.noise import THREADS as NOISE_THREADS, noise_tables 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 from tsim.statevector.shots import ShotPlan, ShotSampler # NOTE: the words of a batch's ``counts``. _SURVIVORS, _ACCEPTED, _DEVIATION = range(3) -KERNEL_SHARED_BYTES = 48 * 1024 -"""Shared memory that a kernel may use without opting in, static and dynamic.""" - BATCH_SHOTS = 1 << 22 """Shots of one batch on the device.""" @@ -83,7 +80,9 @@ class _Module: noise_kept: Any tape: Any gather_columns: Any + gather_bytes: Any gather_float: Any + select_rows: Any publish: Any count_bits: Any list_bits: Any @@ -98,6 +97,22 @@ class _Module: """Each ``expectations`` value's noise passes (screened, deferred), or None.""" +@dataclass(frozen=True) +class _Output: + """One array of a call: chosen outputs of the accepted rows, ``width`` bytes a row. + + Bit-packed (zero bits past the outputs) or one boolean per output; ``whole`` + is every output in order, bit-packed: the gathered rows themselves. + """ + + columns: Any + """The outputs on the device (``uint32``), or None for the first ``count``.""" + count: int + width: int + bools: bool + whole: bool + + def _pinned(shape: tuple[int, ...], dtype) -> np.ndarray: """Return an uninitialized array in pinned host memory.""" import cupyx @@ -116,12 +131,15 @@ def _pinned_counts(cp) -> tuple[np.ndarray, Any]: class _HostRows: - """A call's passing rows (and expectations) in pinned host memory, chunk by chunk.""" + """A call's arrays of passing rows (and expectations) in pinned host memory. + + One array per output of ``widths`` bytes a row, filled chunk by chunk. + """ - def __init__(self, nb: int, no: int | None) -> None: - self.nb, self.no = nb, no + def __init__(self, widths: list[int], no: int | None) -> None: + self.widths, self.no = widths, no self.capacity, self.total = 0, 0 - self.rows: np.ndarray | None = None + self.arrays: list[np.ndarray] = [] self.values: np.ndarray | None = None def allocate(self, passed: int, of: int, shots: int) -> None: @@ -135,33 +153,32 @@ def allocate(self, passed: int, of: int, shots: int) -> None: def grow(self, capacity: int) -> None: """Make room for ``capacity`` rows, keeping the copied ones (no copy may be pending).""" - rows, values, total = self.rows, self.values, self.total + arrays, values, total = self.arrays, self.values, self.total self._resize(capacity) - assert self.rows is not None and rows is not None - self.rows[:total] = rows[:total] + for new, old in zip(self.arrays, arrays, strict=True): + new[:total] = old[:total] if values is not None and self.values is not None: self.values[:total] = values[:total] def _resize(self, capacity: int) -> None: self.capacity = capacity - self.rows = _pinned((capacity, self.nb), np.uint8) + self.arrays = [_pinned((capacity, w), np.uint8) for w in self.widths] if self.no is not None: self.values = _pinned((capacity, self.no), np.float32) - def take(self, count: int) -> tuple[np.ndarray, np.ndarray | None]: - """Return the next ``count`` rows (and expectations) to fill.""" - assert self.rows is not None and self.total + count <= self.capacity + def take(self, count: int) -> tuple[list[np.ndarray], np.ndarray | None]: + """Return the next ``count`` rows of each array (and expectations) to fill.""" + assert self.total + count <= self.capacity lo, self.total = self.total, self.total + count values = None if self.values is None else self.values[lo : self.total] - return self.rows[lo : self.total], values + return [a[lo : self.total] for a in self.arrays], values - def result(self) -> tuple[np.ndarray, np.ndarray | None]: - """Return the filled rows and expectations.""" - if self.rows is None: + def result(self) -> tuple[list[np.ndarray], np.ndarray | None]: + """Return the filled arrays and expectations.""" + if self.capacity == 0: self._resize(0) - assert self.rows is not None values = None if self.values is None else self.values[: self.total] - return self.rows[: self.total], values + return [a[: self.total] for a in self.arrays], values class CudaShotSampler: @@ -173,8 +190,6 @@ class CudaShotSampler: observable's expectation. """ - device = "cuda" - supports = staticmethod(tape.supports) def __init__( @@ -183,19 +198,25 @@ def __init__( channels: ChannelSampler, num_detectors: int, seed: np.random.SeedSequence, + *, + required: bool = False, ) -> None: """Upload the noise tables of ``host``'s shot words. - ``host`` draws the reference samples, on the CPU. + ``host`` draws the reference samples, on the CPU, and the calls that the + GPU cannot run, unless ``required`` (``device="cuda"``): then those raise. Raises: ImportError: If CuPy is not installed. + ValueError: If the noise exceeds the kernels' limits. """ self._cp = runtime.cupy() cp = self._cp assert self.supports(host.layout) - self._args = (host, channels, num_detectors, seed) + self._args = (host, channels, num_detectors, seed, required) + self._required = required + self._cpu_only = False self._lock = threading.RLock() self._host = host layout = host.layout @@ -205,20 +226,16 @@ def __init__( self._channels = channels self._noise = noise_tables(channels, layout.noise, layout.num_words) n = self._noise - # NOTE: gather_columns and noise_kept stage a block's output words and its - # packed rows in shared memory. - staged = 8 * layout.output_words * n.block_shots - staged += n.block_shots * -(-layout.num_outputs // 8) - if staged > KERNEL_SHARED_BYTES: - raise ValueError( - f"{layout.num_outputs} outputs do not fit the GPU kernels' shared memory" - ) + # NOTE: gather_columns stages a block's output words and its packed rows in + # shared memory; without room for 32 rows, gather_bytes reads them directly. + self._gather_shots = kernels.gather_shots(layout, n.block_shots) self._record_args = [self._upload(n.overflow.view(np.uint8))] self._all = self._launch(n.groups) self._modules: dict[bytes | None, _Module] = {} self._buffers: dict[str, Any] = {} self._views: dict[tuple[str, tuple[int, ...]], Any] = {} self._constants: dict[bytes, Any] = {} + self._columns: dict[bytes, Any] = {} # NOTE: each buffer slot has a kernel stream and a copy stream, so that a # slot's next kernels need not wait for its copy to the host. # NOTE: a slot's noise kernel runs on a stream of the least priority, so the @@ -232,6 +249,17 @@ def __init__( # NOTE: lowered when a call runs out of device memory. self._max_batch = BATCH_SHOTS + @property + def device(self) -> str: + """Return where this sampler samples (``SamplerStats.device``). + + It is ``"cpu"`` after a kernel failed to compile, and in a process forked + after its parent used the GPU, where every call samples on the CPU. + """ + if self._cpu_only or runtime.forked_after_init(): + return "cpu" + return "cuda" + 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) @@ -280,10 +308,14 @@ def _module(self, mask: np.ndarray | None) -> _Module: ) compiled = _Module( noise=module.get_function("noise"), - noise_kept=module.get_function("noise_kept"), + noise_kept=( + module.get_function("noise_kept") if self._noise.shared_words else None + ), tape=tape_fn, gather_columns=module.get_function("gather_columns"), + gather_bytes=module.get_function("gather_bytes"), gather_float=module.get_function("gather_float"), + select_rows=module.get_function("select_rows"), publish=module.get_function("publish"), count_bits=module.get_function("count_bits"), list_bits=module.get_function("list_bits"), @@ -329,7 +361,8 @@ def _split( read too. Each key word is read whole, so the deferred channels flip output words only. """ - if not module.has_mask: + # NOTE: rows too wide for shared memory draw all their noise at once. + if not module.has_mask or not self._noise.shared_words: return None if expectations not in module.splits: found = split( @@ -348,7 +381,12 @@ def _split( return module.splits[expectations] def _compile(self, source: str): - """Compile with NVRTC; ptxas -O3 fails on some shuffled tapes, so lower levels follow.""" + """Compile with NVRTC; ptxas -O3 fails on some shuffled tapes, so lower levels follow. + + Raises: + runtime.CompileError: If no level compiles. + + """ cp = self._cp for level in ("3", "2", "1"): options = ( @@ -361,9 +399,11 @@ def _compile(self, source: str): module = cp.RawModule(code=source, options=options) module.get_function("tape") return module - except cp.cuda.compiler.CompileException: + except cp.cuda.compiler.CompileException as error: if level == "1": - raise + raise runtime.CompileError( + f"NVRTC compilation failed: {error}" + ) from error raise AssertionError("unreachable") # -- execution -------------------------------------------------------------- @@ -494,8 +534,8 @@ def _pack( dst = self._cp.empty((max(count, 1), nb), dtype=np.uint8) if not count or not nb: return dst - blocks = (count + S - 1) // S if split is not None: + blocks = (count + S - 1) // S launch = split[1] seed = np.uint64(self._rng.integers(0, 2**64, dtype=np.uint64)) module.noise_kept( @@ -514,9 +554,21 @@ def _pack( ), shared_mem=max(8 * launch.max_events, S * nb), ) + elif self._gather_shots == 0: + module.gather_bytes( + (min(-(-count * nb // COPY_THREADS), 16 * runtime.multiprocessors()),), + (COPY_THREADS,), + ( + words, + np.int64(count), + index if index is not None else np.uint64(0), + dst, + ), + ) else: + GS = self._gather_shots module.gather_columns( - (blocks,), + ((count + GS - 1) // GS,), (COPY_THREADS,), ( words, @@ -524,10 +576,56 @@ def _pack( index if index is not None else np.uint64(0), dst, ), - shared_mem=S * nb, + shared_mem=GS * nb, ) return dst + def _select(self, module: _Module, rows, count: int, output: _Output) -> Any: + """Return ``output``'s outputs of the ``count`` packed ``rows``, on the device.""" + dst = self._cp.empty((count, output.width), dtype=np.uint8) + blocks = min( + -(-count * output.width // COPY_THREADS), 16 * runtime.multiprocessors() + ) + module.select_rows( + (max(blocks, 1),), + (COPY_THREADS,), + ( + rows, + np.int64(count), + output.columns if output.columns is not None else np.uint64(0), + np.int32(output.count), + np.int32(output.width), + np.int32(output.bools), + dst, + ), + ) + return dst + + def _outputs(self, columns: list[np.ndarray], bools: bool) -> list[_Output]: + """Return the arrays of a call that returns ``columns``, one array each.""" + num_outputs = self.layout.num_outputs + found = [] + for chosen in columns: + chosen = np.asarray(chosen, dtype=np.uint32) + count = len(chosen) + in_order = bool(np.array_equal(chosen, np.arange(count))) + device = None + if not in_order: + key = chosen.tobytes() + if key not in self._columns: + self._columns[key] = self._upload(chosen) + device = self._columns[key] + found.append( + _Output( + columns=device, + count=count, + width=count if bools else -(-count // 8), + bools=bools, + whole=in_order and count == num_outputs and not bools, + ) + ) + return found + def _list(self, module: _Module, bits, shots: int, count: int, slot: int) -> Any: """Return the ``count`` shots of the bitmap ``bits`` of ``shots`` shots, in order.""" words = -(-shots // 32) @@ -539,6 +637,38 @@ def _list(self, module: _Module, bits, shots: int, count: int, slot: int) -> Any module.list_bits(grid, threads, (bits, np.int64(words), totals, index)) return index[:count] + def _or_cpu(self, on_gpu, on_cpu): + """Return ``on_gpu()``, or ``on_cpu()`` where the GPU fails. + + It fails in a process forked after its parent used the GPU, out of GPU + memory, and when a kernel does not compile; after the last the sampler + samples on the CPU. A ``required`` sampler raises instead. + + Raises: + RuntimeError: If ``required``, in a process forked after its parent + used the GPU. + MemoryError: If ``required``, out of GPU memory. + runtime.CompileError: If ``required``, when a kernel does not compile. + + """ + if self._cpu_only: + return on_cpu() + if runtime.forked_after_init(): + if self._required: + raise RuntimeError(f"Cannot sample on CUDA: {FORKED}.") + return on_cpu() + try: + return on_gpu() + except MemoryError: + if self._required: + raise + return on_cpu() + except runtime.CompileError: + if self._required: + raise + self._cpu_only = True + return on_cpu() + def sample( self, shots: int, @@ -553,57 +683,58 @@ def sample( As ``ShotSampler.sample``: with a ``mask`` only the shots that pass postselection are returned. ``batch_size`` bounds the shots of a device - batch. + batch. Out of GPU memory, when a kernel does not compile, and in a + process forked after its parent used the GPU, the call samples on the + CPU, or raises if the sampler is ``required`` on the GPU (``_or_cpu``). """ - accepted = None - if not _forked(): - with self._lock: - accepted = self._sample_or_none(shots, mask, flip, batch_size) - if accepted is None: - return self._host.sample( + + def on_gpu(): + outputs, _, deviation = self._accepted( + shots, + mask=mask, + flip=flip, + batch_size=batch_size, + bools=not bit_packed, + columns=columns, + ) + return outputs, 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, - ) - rows, deviation = accepted - num_outputs = self.layout.num_outputs - return [select(rows, c, num_outputs, bit_packed) for c in columns], deviation + ), + ) - def _sample_or_none( - self, - shots: int, - mask: np.ndarray | None, - flip: np.ndarray | None, - batch_size: int | None, - ) -> tuple[np.ndarray, float] | None: - """Return ``sample_accepted``'s rows and deviation, or None if the GPU is full. + def _accepted(self, shots: int, *, batch_size: int | None, **options) -> Any: + """Return ``sample_accepted(shots, **options)``, in smaller batches if the GPU is full. Out of device memory, the call is retried in batches half as large, down - to ``MIN_BATCH_SHOTS``; past that it warns and returns None, for the CPU. - Later calls start from the batch that fit. + to ``MIN_BATCH_SHOTS``. Later calls start from the batch that fit. + + Raises: + MemoryError: If even batches of ``MIN_BATCH_SHOTS`` do not fit. + """ batch = min(batch_size or BATCH_SHOTS, self._max_batch) while True: try: - rows, _, deviation = self.sample_accepted( - shots, mask=mask, flip=flip, batch_size=batch - ) - return rows, deviation + return self.sample_accepted(shots, batch_size=batch, **options) except Exception as error: if not runtime.out_of_memory(error): raise self._release() self._max_batch = max(batch // 2, MIN_BATCH_SHOTS) if batch <= MIN_BATCH_SHOTS: - warnings.warn( - "The GPU ran out of memory, even for batches of " - f"{MIN_BATCH_SHOTS} shots; sampling on the CPU.", - stacklevel=5, - ) - return None + raise MemoryError( + "the GPU ran out of memory, even for batches of " + f"{MIN_BATCH_SHOTS} shots: {error}" + ) from error batch = self._max_batch def _release(self) -> None: @@ -636,11 +767,16 @@ def sample_accepted( flip: np.ndarray | None = None, expectations: bool = False, batch_size: int | None = None, - ) -> tuple[np.ndarray, np.ndarray | None, float]: + bools: bool = False, + columns: list[np.ndarray] | None = None, + ) -> tuple[Any, np.ndarray | None, float]: """Sample ``shots`` shots; return only those that pass postselection. A shot passes when none of the ``mask``'s detectors is set (every shot - passes without a mask). Rows keep their shots' order and hold every output. + passes without a mask). Rows keep their shots' order and hold every output; + with ``bools``, one boolean per output, unpacked on the GPU. With + ``columns``, the outputs are a list of arrays instead, one per entry: its + outputs, in its order, selected on the GPU. The shots run in chunks on two streams: each chunk's rows go to the host while the next chunk's kernels run. The chunks hold about ``CHUNK_BYTES`` @@ -650,7 +786,8 @@ def sample_accepted( Returns: ``(outputs, expectations, deviation)``: the passing shots' outputs, bit-packed as ``np.packbits(..., bitorder="little")`` would (``uint8``, - in pinned host memory), with ``expectations`` each observable's + in pinned host memory), as booleans, or with ``columns`` a list of such + arrays; with ``expectations`` each observable's expectation of ``(-1)^observable`` given the shot's history before the measurement that settles it, or its sampled ``±1`` when a postselected output of its component is tested at or after that @@ -665,6 +802,8 @@ def sample_accepted( flip=flip, expectations=expectations, batch_size=batch_size, + bools=bools, + columns=columns, ) def _sample_accepted( @@ -675,20 +814,26 @@ def _sample_accepted( flip: np.ndarray | None, expectations: bool, batch_size: int | None, - ) -> tuple[np.ndarray, np.ndarray | None, float]: + bools: bool, + columns: list[np.ndarray] | None, + ) -> tuple[Any, np.ndarray | None, float]: cp = self._cp module = self._module(mask) layout = self.layout nb, no = -(-layout.num_outputs // 8), len(self._observables) + chosen = [np.arange(layout.num_outputs)] if columns is None else columns + outputs = self._outputs(chosen, bools) + out = _HostRows([o.width for o in outputs], no if expectations else None) if not shots: - values = np.zeros((0, no), np.float32) if expectations else None - return np.zeros((0, nb), np.uint8), values, 0.0 + arrays, values = out.result() + return self._arrays(arrays, columns, bools), values, 0.0 constant = self._constant(flip) split = self._split(module, expectations) launch = split[0] if split else self._all batch = min(batch_size or BATCH_SHOTS, BATCH_SHOTS) - out = _HostRows(nb, no if expectations else None) state = {"start": 0, "slot": 0} + # NOTE: by the packed rows, so that a seed's shots do not depend on the + # outputs that a call returns. size = self._chunk_shots(min(batch, shots), nb, module.passing) deviation = 0.0 pending: list[tuple[int, int, int, tuple[Any, Any, Any, Any], Any]] = [] @@ -736,7 +881,7 @@ def finish() -> float: # NOTE: a rare excess of passing shots: the copies in flight land first. self._synchronize() out.grow(out.total + count + (shots - first - n)) - rows_host, values_host = out.take(count) + parts, values_host = out.take(count) index = ( None if mask is None @@ -747,11 +892,26 @@ def finish() -> float: copy = stream if first + n == shots and not held else self._copies[slot] used.add(copy) if count and nb: - dst = self._pack(module, split, words, count, index) + rows = self._pack(module, split, words, count, index) + held.append(rows) + found = [ + ( + rows[:count] + if o.whole + else ( + self._select(module, rows, count, o) + if o.width + else None + ) + ) + for o in outputs + ] if copy is not stream: copy.wait_event(stream.record()) - dst[:count].get(stream=copy, out=rows_host, blocking=False) - held.append(dst) + for part, dst in zip(parts, found, strict=True): + if dst is not None: + dst.get(stream=copy, out=part, blocking=False) + held.append(dst) if values_host is not None and count: if no and layout.components: if index is None: @@ -782,11 +942,19 @@ def finish() -> float: deviation = max(deviation, finish()) for stream in used: stream.synchronize() - rows, values = out.result() - if values is None: - return rows, None, deviation - self._direct_expectations(rows, values) - return rows, values, deviation + arrays, values = out.result() + if values is not None: + # NOTE: the rows of every output, bit-packed, which a call with + # expectations returns. + self._direct_expectations(arrays[0], values) + return self._arrays(arrays, columns, bools), values, deviation + + @staticmethod + def _arrays(arrays: list[np.ndarray], columns, bools: bool) -> Any: + """Return a call's arrays: booleans viewed as such, one array without ``columns``.""" + if bools: + arrays = [a.view(bool) for a in arrays] + return arrays[0] if columns is None else arrays def _direct_expectations(self, rows: np.ndarray, values: np.ndarray) -> None: """Set the expectations of observables outside the components: ``±1``.""" @@ -799,25 +967,19 @@ 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.""" - accepted = None - if not _forked(): - with self._lock: - accepted = self._sample_or_none(shots, None, flip, None) - if accepted is None: - return self._host.direct_bools(shots, flip) - rows, _ = accepted - num_outputs = self.layout.num_outputs - return select(rows, np.arange(num_outputs), num_outputs, False) - - -def _forked() -> bool: - """Return whether this process was forked after its parent used CUDA, with a warning.""" - if not runtime.forked_after_init(): - return False - warnings.warn(FORKED, stacklevel=5) - return True + return self._or_cpu( + lambda: self._accepted(shots, flip=flip, batch_size=batch_size, bools=True)[ + 0 + ], + lambda: self._host.direct_bools(shots, flip, batch_size), + ) def _restore( @@ -825,11 +987,16 @@ def _restore( channels: ChannelSampler, num_detectors: int, seed: np.random.SeedSequence, + required: bool, state: dict, ) -> CudaShotSampler | ShotSampler: - """Rebuild a pickled ``CudaShotSampler``, or its CPU sampler where CUDA is missing.""" - if not runtime.available(): - return host - shots = CudaShotSampler(host, channels, num_detectors, seed) - shots._rng.bit_generator.state = state + """Rebuild a pickled ``CudaShotSampler``, or its CPU sampler where CUDA cannot run. + + As at compile time, a sampler ``required`` on the GPU raises where CUDA cannot run. + """ + from tsim.statevector import cuda + + shots = cuda.shot_sampler(host, channels, num_detectors, seed, required=required) + if isinstance(shots, CudaShotSampler): + shots._rng.bit_generator.state = state return shots diff --git a/src/tsim/statevector/gpu/tape.py b/src/tsim/statevector/gpu/tape.py index 22e6a97a..4ec08f28 100644 --- a/src/tsim/statevector/gpu/tape.py +++ b/src/tsim/statevector/gpu/tape.py @@ -162,6 +162,22 @@ def supports(layout: ShotLayout) -> bool: ) +def unsupported(layout: ShotLayout) -> str | None: + """Return why a GPU executor cannot run ``layout``'s components, or ``None`` if it can.""" + if supports(layout) or layout.num_words == 0: + return None + width = max(c.width for c in layout.components) + if width > MAX_WIDTH: + return ( + f"a component needs {width} active qubits, more than the " + f"{MAX_WIDTH} that the GPU executor supports" + ) + return ( + f"a rotation has a branch of weight {rarest_branch(layout):.1e}, " + f"below the {RESOLVABLE:.0e} that float32 amplitudes resolve" + ) + + def rarest_branch(layout: ShotLayout) -> float: """Return the smallest nonzero branch weight (``cos^2``, ``sin^2``) of a rotation. diff --git a/src/tsim/statevector/metal/__init__.py b/src/tsim/statevector/metal/__init__.py index c9e5fe50..a8f2f18f 100644 --- a/src/tsim/statevector/metal/__init__.py +++ b/src/tsim/statevector/metal/__init__.py @@ -12,7 +12,7 @@ from typing import TYPE_CHECKING -from tsim.statevector.gpu.tape import MAX_WIDTH, RESOLVABLE, rarest_branch, supports +from tsim.statevector.gpu.tape import unsupported from tsim.statevector.metal import runtime if TYPE_CHECKING: @@ -20,7 +20,7 @@ from tsim.noise.channels import ChannelSampler from tsim.statevector.metal.sampler import MetalShotSampler - from tsim.statevector.shots import ShotLayout, ShotSampler + from tsim.statevector.shots import ShotSampler def shot_sampler( @@ -55,7 +55,7 @@ def shot_sampler( return host if host.layout.num_words == 0: return host - reason = _unsupported(host.layout) + reason = unsupported(host.layout) if reason is not None: if required: raise ValueError(f"Cannot sample on Metal: {reason}.") @@ -73,19 +73,3 @@ def shot_sampler( if required: raise ValueError(f"Cannot sample on Metal: {error}.") from error return host - - -def _unsupported(layout: ShotLayout) -> str | None: - """Return why the GPU executor cannot run ``layout``, or ``None`` if it can.""" - if supports(layout): - return None - width = max(c.width for c in layout.components) - if width > MAX_WIDTH: - return ( - f"a component needs {width} active qubits, more than the " - f"{MAX_WIDTH} that the GPU executor supports" - ) - return ( - f"a rotation has a branch of weight {rarest_branch(layout):.1e}, " - f"below the {RESOLVABLE:.0e} that float32 amplitudes resolve" - ) diff --git a/test/conftest.py b/test/conftest.py index 1435e449..16e49047 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,10 +1,17 @@ import pytest +from tsim.statevector.cuda import runtime as cuda_runtime from tsim.statevector.metal import runtime as metal_runtime +GPU_TESTS = { + "test_metal_sampler.py": metal_runtime, + "test_cuda_sampler.py": cuda_runtime, +} + @pytest.fixture(autouse=True) def _cpu_samplers(request, monkeypatch): - """Sample on the CPU, except in the Metal tests, which choose per sampler.""" - if request.path.name != "test_metal_sampler.py": - monkeypatch.setattr(metal_runtime, "unavailable", lambda: "a test on the CPU") + """Sample on the CPU, except in each GPU's tests, which choose per sampler.""" + for name, runtime in GPU_TESTS.items(): + if request.path.name != name: + monkeypatch.setattr(runtime, "unavailable", lambda: "a test on the CPU") diff --git a/test/integration/test_cuda_sampler.py b/test/integration/test_cuda_sampler.py index 3791582c..1c3a4b93 100644 --- a/test/integration/test_cuda_sampler.py +++ b/test/integration/test_cuda_sampler.py @@ -1,5 +1,6 @@ """The CUDA executor end to end, against the stabilizer-rank strategy and the CPU executor.""" +import itertools import multiprocessing import pickle import threading @@ -11,13 +12,15 @@ import pytest from tsim.circuit import Circuit -from tsim.statevector.cuda import available, codegen, runtime, sampler as cuda_sampler +from tsim.statevector.cuda import codegen, kernels, runtime, sampler as cuda_sampler from tsim.statevector.cuda.sampler import CudaShotSampler from tsim.statevector.gpu import noise as gpu_noise, tape from tsim.statevector.program import StatevectorComponent from tsim.statevector.shots import ShotSampler -pytestmark = pytest.mark.skipif(not available(), reason="needs CuPy and a CUDA GPU") +pytestmark = pytest.mark.skipif( + not runtime.available(), reason="needs CuPy and a CUDA GPU" +) # NOTE: t_ghz plus a direct detector (detector 2), as in test_trajectory_sharing. GHZ_WITH_DIRECT_DETECTOR = """ @@ -66,10 +69,8 @@ def _gpu_shots(sampler) -> CudaShotSampler: 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) + """Return ``compile(**kwargs)``, a sampler on the CPU.""" + sampler = compile(device="cpu", **kwargs) assert sampler._shots is None or type(sampler._shots) is ShotSampler return sampler @@ -217,6 +218,25 @@ def test_cuda_samples_no_shots(): sampler = _gpu(circuit.compile_detector_sampler(seed=1)) assert sampler.sample(0, append_observables=True).shape == (0, 4) assert sampler.sample(0, bit_packed=True).shape == (0, 1) + for bit_packed in (False, True): + detectors, observables = sampler.sample( + 0, separate_observables=True, bit_packed=bit_packed + ) + assert detectors.shape == (0, 3 if not bit_packed else 1) + assert observables.shape == (0, 1) + + +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)) # NOTE: 150 noisy detectors, so the rows span three words, and two observables. @@ -514,9 +534,12 @@ def test_cuda_shots_start_after_the_common_prefix(width): def test_components_wider_than_the_limit_run_on_the_cpu(): circuit = _random_unitary_circuit(13, 12) assert _widest(circuit) > tape.MAX_WIDTH - with pytest.warns(UserWarning, match="sampling on the CPU"): + with warnings.catch_warnings(): + warnings.simplefilter("error") sampler = circuit.compile_sampler(seed=1) assert type(sampler._shots) is ShotSampler + with pytest.raises(ValueError, match="more than the 11 that the GPU executor"): + circuit.compile_sampler(seed=1, device="cuda") def _sampler_in_child(_): @@ -524,16 +547,23 @@ def _sampler_in_child(_): 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] + assert sampler.stats.device == "cpu" + try: + Circuit(T_STATE_X).compile_sampler(seed=1, device="cuda") + error = "" + except RuntimeError as raised: + error = str(raised) + return type(sampler._shots).__name__, [str(w.message) for w in caught], error def test_processes_forked_after_the_gpu_was_used_sample_on_the_cpu(): # NOTE: CUDA cannot be used in a child forked after its parent initialized 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] + name, messages, error = pool.map(_sampler_in_child, [0])[0] assert name == "ShotSampler" - assert any("forked" in m for m in messages) + assert messages == [] + assert "forked" in error def test_noise_of_tiny_probabilities_fires_at_its_rate(): @@ -574,10 +604,11 @@ def test_rare_outcome_zero_is_drawn_at_its_rate(p0): assert abs(int(count.get()[0]) - expected) <= 5 * np.sqrt(expected) + 1 -@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_cuda_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)) @@ -619,12 +650,32 @@ 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)) _out_of_memory_above(monkeypatch, sampler, 0) - with pytest.warns(UserWarning, match="out of memory"): + with warnings.catch_warnings(): + warnings.simplefilter("error") 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) + required = _gpu(circuit.compile_detector_sampler(seed=3, device="cuda")) + _out_of_memory_above(monkeypatch, required, 0) + with pytest.raises(MemoryError, match="the GPU ran out of memory"): + required.sample(10, append_observables=True) + + +def test_samplers_out_of_gpu_memory_at_setup_run_on_the_cpu(monkeypatch): + cp = runtime.cupy() + + def out_of_memory(*args, **kwargs): + raise cp.cuda.memory.OutOfMemoryError(1 << 40, 0) + + monkeypatch.setattr(CudaShotSampler, "__init__", out_of_memory) + with warnings.catch_warnings(): + warnings.simplefilter("error") + sampler = Circuit(T_STATE_X).compile_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + with pytest.raises(MemoryError, match="Cannot sample on CUDA"): + Circuit(T_STATE_X).compile_sampler(seed=1, device="cuda") def test_kernels_wait_for_the_tables_they_read(): @@ -659,8 +710,12 @@ def test_tsim_device_cpu_keeps_the_samplers_on_the_cpu(monkeypatch): warnings.simplefilter("error") sampler = Circuit(T_STATE_X).compile_sampler(seed=1) assert type(sampler._shots) is ShotSampler + # NOTE: an explicit device overrides TSIM_DEVICE. + _gpu(Circuit(T_STATE_X).compile_sampler(seed=1, device="cuda")) monkeypatch.setenv("TSIM_DEVICE", "") _gpu(Circuit(T_STATE_X).compile_sampler(seed=1)) + sampler = Circuit(T_STATE_X).compile_sampler(seed=1, device="cpu") + assert type(sampler._shots) is ShotSampler def test_the_repr_says_where_a_sampler_samples(): @@ -673,11 +728,14 @@ def test_the_repr_says_where_a_sampler_samples(): def test_programs_without_shot_words_sample_on_the_cpu(): # NOTE: a detector sampler of a circuit without detectors, observables or noise. - with warnings.catch_warnings(): - warnings.simplefilter("error") - sampler = Circuit("R 0\nH 0\nT 0\nM 0").compile_detector_sampler(seed=1) - assert type(sampler._shots) is ShotSampler - assert sampler.sample(10).shape == (10, 0) + for device in ("auto", "cuda"): + with warnings.catch_warnings(): + warnings.simplefilter("error") + sampler = Circuit("R 0\nH 0\nT 0\nM 0").compile_detector_sampler( + seed=1, device=device + ) + assert type(sampler._shots) is ShotSampler + assert sampler.sample(10).shape == (10, 0) def test_noise_beyond_the_gpu_tables_samples_on_the_cpu(): @@ -688,10 +746,13 @@ def test_noise_beyond_the_gpu_tables_samples_on_the_cpu(): ["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"): + with warnings.catch_warnings(): + warnings.simplefilter("error") sampler = Circuit(text).compile_detector_sampler(seed=1) assert type(sampler._shots) is ShotSampler assert sampler.sample(10).shape == (10, 9) + with pytest.raises(ValueError, match="channels of up to 127 outcomes"): + Circuit(text).compile_detector_sampler(seed=1, device="cuda") _FORKED = {} @@ -701,7 +762,7 @@ def _sample_in_child(_): with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") shape = _FORKED["sampler"].sample(100).shape - return shape, [str(w.message) for w in caught] + return shape, [str(w.message) for w in caught], _FORKED["sampler"].stats.device def test_samplers_compiled_before_a_fork_sample_on_the_cpu_in_the_child(): @@ -711,11 +772,13 @@ def test_samplers_compiled_before_a_fork_sample_on_the_cpu_in_the_child(): _FORKED["sampler"] = sampler try: with multiprocessing.get_context("fork").Pool(1) as pool: - shape, messages = pool.map(_sample_in_child, [0])[0] + shape, messages, device = pool.map(_sample_in_child, [0])[0] finally: _FORKED.clear() assert shape == (100, 3) - assert any("forked" in m for m in messages) + assert messages == [] + assert device == "cpu" + assert sampler.stats.device == "cuda" def test_gpu_samplers_pickle_and_continue_their_streams(): @@ -724,6 +787,8 @@ def test_gpu_samplers_pickle_and_continue_their_streams(): loaded = pickle.loads(pickle.dumps(sampler)) # noqa: S301 assert isinstance(loaded._shots, CudaShotSampler) np.testing.assert_array_equal(loaded.sample(5000), sampler.sample(5000)) + required = Circuit(T_STATE_X).compile_detector_sampler(seed=26, device="cuda") + assert _gpu_shots(pickle.loads(pickle.dumps(required)))._required # noqa: S301 def test_an_empty_mask_is_not_taken_for_no_mask(): @@ -734,16 +799,32 @@ def test_an_empty_mask_is_not_taken_for_no_mask(): assert kept.shape == (1000, 1) -def test_outputs_beyond_the_kernels_shared_memory_sample_on_the_cpu(): - # NOTE: 6200 outputs need 97 words, whose packed rows do not fit in shared - # memory beside them. - n = 6200 +def test_cuda_samples_rows_too_wide_for_shared_memory(): + # NOTE: 12000 detectors are 188 words: the noise goes into device memory, and + # the rows are gathered without staging them in shared memory. + n = 12000 qubits = " ".join(map(str, range(n))) - circuit = Circuit(f"R {qubits}\nX_ERROR(0.01) {qubits}\nM {qubits}") - with pytest.warns(UserWarning, match="cannot sample this circuit"): - sampler = circuit.compile_sampler(seed=1) - assert type(sampler._shots) is ShotSampler - assert sampler.sample(10, bit_packed=True).shape == (10, n // 8) + circuit = Circuit( + "\n".join( + [ + f"R {qubits}", + f"X_ERROR(0.05) {qubits}", + f"M {qubits}", + *(f"DETECTOR rec[{i - n}]" for i in range(n)), + "OBSERVABLE_INCLUDE(0) rec[-1] rec[-2]", + ] + ) + ) + sampler = _gpu(circuit.compile_detector_sampler(seed=31)) + shots = _gpu_shots(sampler) + assert not shots._noise.shared_words and shots._gather_shots == 0 + got = sampler.sample(20_000, append_observables=True) + np.testing.assert_allclose(got[:, :n].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]) + packed = sampler.sample(1000, append_observables=True, bit_packed=True) + assert packed.shape == (1000, -(-(n + 1) // 8)) + assert not np.unpackbits(packed, axis=1, bitorder="little")[:, n + 1 :].any() def test_concurrent_calls_on_one_sampler_are_serialized(): @@ -772,3 +853,150 @@ def work(k): else: # NOTE: detector 2 fires at rate 0.2. assert abs(got[:, 2].mean() - 0.2) < 0.02 + + +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 'extern "C" __global__ void tape(' + + monkeypatch.setattr(kernels, "source", broken) + with warnings.catch_warnings(): + warnings.simplefilter("error") + got = sampler.sample(20_000) + assert sampler.stats.device == "cpu" + 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) + required = _gpu(circuit.compile_sampler(seed=3, device="cuda")) + with pytest.raises(runtime.CompileError): + required.sample(10) + + +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 warnings.catch_warnings(): + warnings.simplefilter("error") + sampler = circuit.compile_detector_sampler(seed=1) + assert sampler.stats.device == "cpu" + with pytest.raises(ValueError, match=r"rotation has a branch of weight 2\.5e-10"): + circuit.compile_detector_sampler(seed=1, device="cuda") + + +@pytest.mark.parametrize( + ("detectors", "observables"), [(20, 1), (21, 3), (7, 9), (16, 8), (130, 5)] +) +def test_cuda_selects_every_layout_of_columns(detectors, observables): + # NOTE: the observables start inside a byte (20, 21, 7, 130), end inside one + # (20, 130) or span two bytes (7); 16 + 8 lies on byte boundaries. Samplers of + # one seed draw the same shots; the reference is every output as booleans, + # in order. + n = detectors + observables + lines = [f"X_ERROR(0.3) {' '.join(map(str, range(n)))}", "H 0", "T 0", "H 0"] + lines.append(f"M {' '.join(map(str, range(n)))}") + lines += [f"DETECTOR rec[-{n - i}]" for i in range(detectors)] + lines += [ + f"OBSERVABLE_INCLUDE({j}) rec[-{observables - j}]" for j in range(observables) + ] + circuit = Circuit("\n".join(lines)) + mask = np.zeros(detectors, dtype=bool) + mask[0] = True + for postselection_mask in (None, mask): + + def sample(layout, bit_packed=False, postselection_mask=postselection_mask): + sampler = _gpu(circuit.compile_detector_sampler(seed=25)) + return sampler.sample( + 20_000, + separate_observables=layout == "separate", + prepend_observables=layout == "prepend", + append_observables=layout == "append", + bit_packed=bit_packed, + postselection_mask=postselection_mask, + ) + + every = sample("append") + assert isinstance(every, np.ndarray) + found, events = every[:, :detectors], every[:, detectors:] + layouts = { + "separate": (found, events), + "prepend": (np.hstack([events, found]),), + "detectors": (found,), + } + for layout, expected in layouts.items(): + for bit_packed in (True, False): + got = sample(layout, bit_packed) + got = got if isinstance(got, tuple) else (got,) + for array, bools in zip(got, expected, strict=True): + if bit_packed: + bools = np.packbits(bools, axis=1, bitorder="little") + else: + assert array.dtype == np.bool_ and array.flags.c_contiguous + np.testing.assert_array_equal(array, bools) + + +@pytest.mark.parametrize("text", [GHZ_WITH_DIRECT_DETECTOR, NOISY_DETECTORS]) +def test_cuda_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]) + + +def test_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) diff --git a/test/unit/test_sampler.py b/test/unit/test_sampler.py index f1dc53e1..bf882587 100644 --- a/test/unit/test_sampler.py +++ b/test/unit/test_sampler.py @@ -829,22 +829,29 @@ def test_device_cpu_samples_on_the_cpu(compile): def test_unknown_devices_raise(): - with pytest.raises(ValueError, match="device must be 'auto', 'cpu' or 'metal'"): + with pytest.raises( + ValueError, match="device must be 'auto', 'cpu', 'metal' or 'cuda'" + ): Circuit(NOISY_T_STATE).compile_sampler(device="gpu") # type: ignore[arg-type] -@pytest.mark.parametrize("device", ["cpu", "metal"]) +@pytest.mark.parametrize("device", ["cpu", "metal", "cuda"]) def test_stabilizer_rank_strategies_take_only_the_auto_device(device): with pytest.raises(ValueError, match="needs strategy='active_statevector'"): Circuit(NOISY_T_STATE).compile_sampler(strategy="cat5", device=device) -def test_metal_device_raises_where_metal_cannot_run(monkeypatch): - from tsim.statevector.metal import runtime +@pytest.mark.parametrize( + ("device", "name", "reason"), + [("metal", "Metal", "no Apple GPU"), ("cuda", "CUDA", "no CUDA device")], +) +def test_gpu_devices_raise_where_the_gpu_cannot_run(monkeypatch, device, name, reason): + from tsim.statevector import cuda, metal - monkeypatch.setattr(runtime, "unavailable", lambda: "no Apple GPU") - with pytest.raises(RuntimeError, match="Cannot sample on Metal: no Apple GPU"): - Circuit(NOISY_T_STATE).compile_detector_sampler(device="metal") + runtime = {"metal": metal.runtime, "cuda": cuda.runtime}[device] + monkeypatch.setattr(runtime, "unavailable", lambda: reason) + with pytest.raises(RuntimeError, match=f"Cannot sample on {name}: {reason}"): + Circuit(NOISY_T_STATE).compile_detector_sampler(device=device) with warnings.catch_warnings(): warnings.simplefilter("error") sampler = Circuit(NOISY_T_STATE).compile_detector_sampler(seed=1) From 9371e02c0abee1325f8a8e20b103524c13ed19b9 Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Thu, 12 Nov 2026 15:08:16 -0700 Subject: [PATCH 19/20] docs: document the CUDA backend beside Metal Co-Authored-By: Claude Opus 5.5 --- .../skills/testing-tsim-correctness/SKILL.md | 6 +-- AGENTS.md | 2 +- CHANGELOG.md | 7 +++ README.md | 7 +-- docs/index.mdx | 6 +-- docs/install.mdx | 51 ++++++++++--------- skills/simulating-circuits-with-tsim/SKILL.md | 16 +++--- 7 files changed, 55 insertions(+), 40 deletions(-) diff --git a/.claude/skills/testing-tsim-correctness/SKILL.md b/.claude/skills/testing-tsim-correctness/SKILL.md index 0c6cf111..6dfce882 100644 --- a/.claude/skills/testing-tsim-correctness/SKILL.md +++ b/.claude/skills/testing-tsim-correctness/SKILL.md @@ -46,9 +46,9 @@ handful of T gates; its memory grows exponentially. test is a broken test. - Stim merges consecutive identical instructions into one with several targets, so write circuits that exercise both single- and multi-target paths. -- GPU tests: `test/conftest.py` keeps every test outside - `test/integration/test_metal_sampler.py` on the CPU. In the Metal tests compile the - GPU side with `device="metal"`, so a silent fallback cannot make the test pass on the +- GPU tests: `test/conftest.py` keeps every test outside `test/integration/test_metal_sampler.py` + and `test_cuda_sampler.py` on the CPU. In the GPU tests compile the GPU side with + `device="metal"` or `device="cuda"`, so a silent fallback cannot make the test pass on the CPU, and the reference with `device="cpu"`; compare them with `assert_counts_match` (GPU samples never equal CPU samples bit for bit). - Keep circuits small (a few qubits, tens of thousands of shots) so the suite stays diff --git a/AGENTS.md b/AGENTS.md index 7b7df498..c154abd4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -68,7 +68,7 @@ more than a handful of T gates: their compile memory grows exponentially. a reference. The reference-sample flags change only that. - **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 + by call and shot index (`shots.py`). On a GPU (`statevector/metal/`, `statevector/cuda/`) 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 diff --git a/CHANGELOG.md b/CHANGELOG.md index 388110b9..814b2b53 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 the CPU, without a warning; `"cpu"` keeps them on the CPU; `"metal"` raises when the GPU cannot run the circuit. `TSIM_DEVICE=cpu` keeps `"auto"` on the CPU, and `SamplerStats.device` and the repr say where they sample. +- `strategy="active_statevector"` also samples on NVIDIA GPUs, through CUDA, with + the new `cuda` extra (`pip install bloqade-tsim[cuda]`, CUDA 13). It runs the + same circuits as Metal and behaves the same way: `device="auto"` picks it when + CuPy sees a GPU, and `device="cuda"` raises when the GPU cannot run the circuit. +- The `jax-cuda12` and `jax-cuda13` extras, for the stabilizer-rank strategies + on NVIDIA GPUs with JAX. The `cuda12` and `cuda13` extras of Tsim 1.0 install + the same packages. - 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 b9c694c0..168bdc86 100644 --- a/README.md +++ b/README.md @@ -24,9 +24,10 @@ uv add bloqade-tsim ``` 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: +needed). To sample on an NVIDIA GPU, install the extra of your strategy: ```bash -pip install "bloqade-tsim[cuda13]" # or bloqade-tsim[cuda12] +pip install "bloqade-tsim[cuda]" # the default strategy, "active_statevector" (CUDA 13) +pip install "bloqade-tsim[jax-cuda13]" # the stabilizer-rank strategies, with JAX (or jax-cuda12) ``` ### Using Tsim with coding agents @@ -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), or an Apple GPU (Metal) | 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) or NVIDIA GPU (CUDA) | 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 e51a566e..cacac74b 100644 --- a/docs/index.mdx +++ b/docs/index.mdx @@ -40,10 +40,10 @@ uv add bloqade-tsim ``` 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: +needed). To sample on an NVIDIA GPU with the default strategy, install the `cuda` extra: ```bash -uv add "bloqade-tsim[cuda13]" +uv add "bloqade-tsim[cuda]" ``` See [Installation](/install) for more options. @@ -55,7 +55,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), or an Apple GPU (Metal) | 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) or NVIDIA GPU (CUDA) | 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 f452828a..dc6ef4f0 100644 --- a/docs/install.mdx +++ b/docs/install.mdx @@ -13,55 +13,60 @@ We recommend using [`uv`](https://docs.astral.sh/uv/) for dependency management. uv add bloqade-tsim ``` -For GPU acceleration with CUDA: - -```bash -# For CUDA 13 -uv add "bloqade-tsim[cuda13]" - -# For CUDA 12 -uv add "bloqade-tsim[cuda12]" -``` - ## Using pip ```bash pip install bloqade-tsim ``` -For GPU acceleration with CUDA: +## GPUs + +The default strategy (`"active_statevector"`) samples on an Apple GPU through Metal, or on an +NVIDIA GPU through CUDA. On an Apple-silicon Mac, the plain install includes PyObjC's Metal +bindings for it. For NVIDIA GPUs, install an extra; each sampling strategy has its own: + +| Extra | Accelerates | Needs | +|---|---|---| +| `cuda` | `"active_statevector"` (the default), with CuPy | CUDA 13 (driver 580 or later) | +| `jax-cuda13` | `"cat5"`, `"bss"` and `"cutting"`, with JAX | CUDA 13 | +| `jax-cuda12` | `"cat5"`, `"bss"` and `"cutting"`, with JAX | CUDA 12 | ```bash -pip install "bloqade-tsim[cuda13]" +uv add "bloqade-tsim[cuda]" +# or, with both strategies on the GPU: +uv add "bloqade-tsim[cuda,jax-cuda13]" ``` -## Apple GPUs +The `cuda12` and `cuda13` extras of Tsim 1.0 still install the JAX extras. JAX starts only +when a stabilizer-rank strategy samples, and then reserves 75% of the GPU's memory; to sample +with both strategies in one process, set `XLA_PYTHON_CLIENT_PREALLOCATE=false` before starting +Python. -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. Where it samples is -the `device` argument of `compile_sampler` and `compile_detector_sampler`: +Where `"active_statevector"` samples is the `device` argument of `compile_sampler` and +`compile_detector_sampler`: | `device` | Samples | | --- | --- | -| `"auto"` (default) | On the GPU when it can run the circuit, else on the CPU, without a warning | +| `"auto"` (default) | On a GPU when it can run the circuit, else on the CPU, without a warning | | `"cpu"` | On the CPU, in double precision | -| `"metal"` | On the GPU, or raises an error that says why it cannot | +| `"metal"` | On an Apple GPU, or raises an error that says why it cannot | +| `"cuda"` | On an NVIDIA GPU, or raises an error that says why it cannot | ```python -sampler = circuit.compile_detector_sampler(device="metal") +sampler = circuit.compile_detector_sampler(device="cuda") ``` The GPU runs circuits whose components have at most 11 active qubits. That includes Clifford circuits, whose detectors depend on the noise alone: the GPU samples the noise as well as the -state vectors. 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`. +state vectors. Postselection runs on the GPU, so heavily postselected circuits gain the most. +A sampler's repr says where it samples (`samples on metal`, `samples on cuda` 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. 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 6e-5 π) samples on the -CPU (or raises, with `device="metal"`). +CPU (or raises, with `device="metal"` or `"cuda"`). A seed's samples on the GPU differ from the CPU's and depend on `batch_size`, so with `device="auto"` a seed reproduces its samples only on the same kind of machine. Use diff --git a/skills/simulating-circuits-with-tsim/SKILL.md b/skills/simulating-circuits-with-tsim/SKILL.md index 9839136f..627a0cc0 100644 --- a/skills/simulating-circuits-with-tsim/SKILL.md +++ b/skills/simulating-circuits-with-tsim/SKILL.md @@ -196,8 +196,9 @@ 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. That +- **GPUs:** on an Apple-silicon Mac (Metal), or on an NVIDIA GPU with the `cuda` extra, 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. That includes Clifford circuits, whose detectors come from the noise alone: the GPU samples the noise too. It computes outcome probabilities in float32 and never draws an outcome below about 1e-11 @@ -205,16 +206,17 @@ yourself. affect GPU sampling. - **`device`** (in `compile_sampler` and `compile_detector_sampler`) chooses where it samples: `"auto"` (default) uses the GPU when it can and otherwise the CPU, silently; `"cpu"` stays - on the CPU in double precision; `"metal"` raises with the reason if the GPU cannot run the - circuit (no Apple GPU, width above 11, a tiny rotation, out of GPU memory). Check - `sampler.stats.device` (or the repr, `samples on metal`) after compiling, since `"auto"` - does not warn. Pass `device="metal"` for long runs that must not drop to the CPU unnoticed, + on the CPU in double precision; `"metal"` (Apple) and `"cuda"` (NVIDIA) raise with the + reason if the GPU cannot run the circuit (no such GPU, width above 11, a tiny rotation, out + of GPU memory). Check `sampler.stats.device` (or the repr, `samples on metal` or `cuda`) + after compiling, since `"auto"` does not warn. Pass `device="metal"` or `"cuda"` for long + runs that must not drop to the CPU unnoticed, and `device="cpu"` to reproduce samples across machines or to compare against the GPU. `TSIM_DEVICE=cpu` in the environment keeps `"auto"` on the CPU; it is read when a sampler is compiled, so an unpickled sampler keeps the device choice it was compiled with. - **Processes:** sample from one process per GPU. Threads sharing a GPU sampler run its calls one at a time, and several processes compete for the same GPU. A process forked after its parent used the GPU cannot use - it (it samples on the CPU, or raises with `device="metal"`); start workers with the + it (it samples on the CPU, or raises with `device="metal"` or `"cuda"`); start workers with the `"spawn"` or `"forkserver"` method (macOS's default is `"spawn"`). - **Seeds:** a `seed` reproduces samples for the same Tsim and dependency versions and device. On the CPU, thread count and `batch_size` do not affect the samples; on a GPU, From 86bd80b8a0dc4678535a6b88d7a5a0c8bd6a858f Mon Sep 17 00:00:00 2001 From: Rafael Haenel Date: Thu, 12 Nov 2026 16:19:27 -0700 Subject: [PATCH 20/20] perf(cuda): keep leading whole bytes a view of the rows; tie check out of line Bit-packed outputs in order that end on a byte boundary (detectors only, the detectors of separate_observables) are a view of the copied rows again, as on Metal, and the 40-bit keep threshold's rare tie check is __noinline__. Co-Authored-By: Claude Opus 5.5 --- src/tsim/statevector/cuda/noise.py | 11 ++++-- src/tsim/statevector/cuda/sampler.py | 50 +++++++++++++++++++++------- 2 files changed, 47 insertions(+), 14 deletions(-) diff --git a/src/tsim/statevector/cuda/noise.py b/src/tsim/statevector/cuda/noise.py index 381239fe..3c6d9cfe 100644 --- a/src/tsim/statevector/cuda/noise.py +++ b/src/tsim/statevector/cuda/noise.py @@ -126,14 +126,21 @@ // 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]``. +// Whether a fault whose fraction's high bits tie with its record's keep +// threshold takes the alias: its low 16 bits, from a hash of the fault, are at +// least the threshold's. NOTE: rare, so out of line, which keeps the hot loop small. +__device__ __noinline__ bool alias_on_tie(u64 stream, unsigned rec, int row, unsigned keep) +{{ + return (unsigned)(mix64(stream ^ ((u64)(rec << 8 | row) * GOLDEN)) & 0xFFFF) >= keep; +}} + template __device__ __forceinline__ void apply_fault(u64* sw, int row, unsigned rec, unsigned draw, u64 stream, const Record* __restrict__ records, const Flip* __restrict__ overflow) {{ Record r = records[rec]; unsigned high = draw & 0xFFFFFF, keep = r.alias & 0x1FFFFFF; - if (high > keep || (high == keep && - (unsigned)(mix64(stream ^ ((u64)(rec << 8 | row) * GOLDEN)) & 0xFFFF) >= r.keep)) + if (high > keep || (high == keep && alias_on_tie(stream, rec, row, r.keep))) r = records[rec - (draw >> 24) + (r.alias >> 25)]; int flips = r.extra & 255; if (flips > 0) atomicXor(&sw[(r.words & 0xFFFF) * WS + row * SS], r.m0); diff --git a/src/tsim/statevector/cuda/sampler.py b/src/tsim/statevector/cuda/sampler.py index 13318ad8..5ab8b0d4 100644 --- a/src/tsim/statevector/cuda/sampler.py +++ b/src/tsim/statevector/cuda/sampler.py @@ -601,14 +601,30 @@ def _select(self, module: _Module, rows, count: int, output: _Output) -> Any: ) return dst - def _outputs(self, columns: list[np.ndarray], bools: bool) -> list[_Output]: - """Return the arrays of a call that returns ``columns``, one array each.""" + def _outputs( + self, columns: list[np.ndarray], bools: bool + ) -> tuple[list[_Output], list[tuple[int, int | None]]]: + """Return the arrays that a call returning ``columns`` copies to the host. + + Also returns, for each entry of ``columns``, its array and the bytes of + each row it keeps (or None for all): bit-packed leading whole bytes are a + view of the rows, which are copied once. + """ num_outputs = self.layout.num_outputs - found = [] + nb = -(-num_outputs // 8) + found: list[_Output] = [] + picks: list[tuple[int, int | None]] = [] + rows = None for chosen in columns: chosen = np.asarray(chosen, dtype=np.uint32) count = len(chosen) in_order = bool(np.array_equal(chosen, np.arange(count))) + if in_order and not bools and (count == num_outputs or count % 8 == 0): + if rows is None: + rows = len(found) + found.append(_Output(None, num_outputs, nb, False, True)) + picks.append((rows, None if count == num_outputs else count // 8)) + continue device = None if not in_order: key = chosen.tobytes() @@ -621,10 +637,11 @@ def _outputs(self, columns: list[np.ndarray], bools: bool) -> list[_Output]: count=count, width=count if bools else -(-count // 8), bools=bools, - whole=in_order and count == num_outputs and not bools, + whole=False, ) ) - return found + picks.append((len(found) - 1, None)) + return found, picks def _list(self, module: _Module, bits, shots: int, count: int, slot: int) -> Any: """Return the ``count`` shots of the bitmap ``bits`` of ``shots`` shots, in order.""" @@ -822,11 +839,11 @@ def _sample_accepted( layout = self.layout nb, no = -(-layout.num_outputs // 8), len(self._observables) chosen = [np.arange(layout.num_outputs)] if columns is None else columns - outputs = self._outputs(chosen, bools) + outputs, picks = self._outputs(chosen, bools) out = _HostRows([o.width for o in outputs], no if expectations else None) if not shots: arrays, values = out.result() - return self._arrays(arrays, columns, bools), values, 0.0 + return self._arrays(arrays, picks, columns, bools), values, 0.0 constant = self._constant(flip) split = self._split(module, expectations) launch = split[0] if split else self._all @@ -947,14 +964,23 @@ def finish() -> float: # NOTE: the rows of every output, bit-packed, which a call with # expectations returns. self._direct_expectations(arrays[0], values) - return self._arrays(arrays, columns, bools), values, deviation + return self._arrays(arrays, picks, columns, bools), values, deviation @staticmethod - def _arrays(arrays: list[np.ndarray], columns, bools: bool) -> Any: - """Return a call's arrays: booleans viewed as such, one array without ``columns``.""" + def _arrays( + arrays: list[np.ndarray], + picks: list[tuple[int, int | None]], + columns, + bools: bool, + ) -> Any: + """Return a call's outputs from its ``arrays`` and ``picks`` (``_outputs``). + + Booleans are viewed as such; without ``columns`` there is one array. + """ + found = [arrays[k] if cut is None else arrays[k][:, :cut] for k, cut in picks] if bools: - arrays = [a.view(bool) for a in arrays] - return arrays[0] if columns is None else arrays + found = [a.view(bool) for a in found] + return found[0] if columns is None else found def _direct_expectations(self, rows: np.ndarray, values: np.ndarray) -> None: """Set the expectations of observables outside the components: ``±1``."""