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..a1e4391b --- /dev/null +++ b/src/tsim/statevector/gpu/__init__.py @@ -0,0 +1,14 @@ +"""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 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.""" diff --git a/src/tsim/statevector/gpu/noise.py b/src/tsim/statevector/gpu/noise.py new file mode 100644 index 00000000..15e8b8d6 --- /dev/null +++ b/src/tsim/statevector/gpu/noise.py @@ -0,0 +1,386 @@ +"""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's high bits, bits 25-31 alias slot + words: the first two flipped words (16 bits each) + extra: bits 0-7 number of flips, bits 8-31 start of the rest in the overflow + keep: bits 0-15 keep threshold's low bits + masks: the first two flip masks + +The keep threshold is the slot's keep probability times ``2^KEEP_BITS``: a fault +draws its slot and a ``KEEP_BITS``-bit fraction, and keeps the slot's outcome if +the fraction is below it (``keep_threshold``). + +The flips past an outcome's first two are 16-byte ``(word, mask)`` pairs, in +order. The samples follow ``ChannelSampler``'s distribution, not its streams: +every block draws from its own counter-based stream. + +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 +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. +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 +# 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 +# NOTE: an outcome's share of its channel is exact to 2^-40 (24 bits left +# outcomes below ~1e-8 impossible), and a nonzero share never rounds to zero. +KEEP_BITS = 40 +DEFER_SHARE = 0.25 +"""Share of the expected faults that postselection cannot see, from which they +are drawn on the accepted shots only, after the tapes (``split``).""" + +_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(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 + # 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)) + 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: + 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 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. + + 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 _keep(probability: float) -> int: + """Return a slot's keep threshold for its keep ``probability``: at least 1 if positive.""" + keep = min(round(float(probability) * (1 << KEEP_BITS)), 1 << KEEP_BITS) + return max(keep, 1) if probability > 0 else 0 + + +def keep_threshold(record: np.void) -> int: + """Return a record's keep threshold, as the kernels read it from its fields.""" + return ((int(record["alias"]) & 0x1FFFFFF) << 16) | int(record["keep"]) + + +def noise_tables( + channels: ChannelSampler, + flips: WordTables, + 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 or 256 events 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 + 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) + 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 << KEEP_BITS if hi - lo == 1 else _keep(fire.alias_prob[k]) + records[k]["alias"] = (keep >> 16) | ((int(fire.alias_index[k]) - lo) << 25) + records[k]["keep"] = keep & 0xFFFF + j0, j1 = int(flips.start[k]), int(flips.start[k + 1]) + words = flips.word[j0:j1].tolist() + masks = flips.mask[j0:j1].tolist() + 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..5f1ee337 --- /dev/null +++ b/src/tsim/statevector/gpu/rows.py @@ -0,0 +1,26 @@ +"""The rows of outputs that the GPU backends return.""" + +from __future__ import annotations + +import numpy as np + + +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..20e86449 --- /dev/null +++ b/src/tsim/statevector/gpu/tape.py @@ -0,0 +1,881 @@ +"""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 +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, 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 + +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, which the GPU cannot tell from a genuine outcome probability that small. +""" + +RESOLVABLE = 1e-8 +"""Smallest branch weight of a rotation (``cos^2`` or ``sin^2`` of its angle) that +a GPU executor runs: a smaller one could give an outcome a genuine probability +below ``CERTAIN``, which the GPU would never draw.""" + +_STRUCTURAL = 1e-20 + +UNROLL_BUDGET = 12_000 +"""Amplitude updates (local spans summed over the actions) up to which a tape unrolls. + +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 (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. + + """ + + 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 __post_init__(self) -> None: + """Check that the chunk size is a positive power of two.""" + chunk = self.amplitude_chunk + assert chunk is None or (chunk > 0 and not chunk & (chunk - 1)) + + 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 supports(layout: ShotLayout) -> bool: + """Return whether a GPU executor runs ``layout``. + + It needs shot words, components of at most ``MAX_WIDTH`` qubits and rotations + whose branches the GPU resolves (``rarest_branch`` at least ``RESOLVABLE``). A + layout without shot words (no outputs and no noise) stays on the CPU. + """ + return ( + layout.num_words > 0 + and all(c.width <= MAX_WIDTH for c in layout.components) + and rarest_branch(layout) >= RESOLVABLE + ) + + +def rarest_branch(layout: ShotLayout) -> float: + """Return the smallest nonzero branch weight (``cos^2``, ``sin^2``) of a rotation. + + It is 1 without rotations; a weight below ``_STRUCTURAL`` (of a Clifford + angle, in float64) counts as zero. + """ + rarest = 1.0 + for component in layout.components: + t = component.tape + rotations = (t.kind == ROTATE) | (t.kind == ROTATE_DIAG) + weights = np.concatenate([t.cos[rotations] ** 2, t.sin[rotations] ** 2]) + weights = weights[weights >= _STRUCTURAL] + if len(weights): + rarest = min(rarest, float(weights.min())) + return rarest + + +def threads_per_shot(layout: ShotLayout) -> int: + """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" + + +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..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 @@ -133,7 +133,43 @@ 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 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 + 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..61c561dd --- /dev/null +++ b/test/unit/statevector/test_gpu.py @@ -0,0 +1,626 @@ +"""The host side of the GPU backends: noise tables, row packing and the tape generator.""" + +import ctypes +import re +import shutil +import subprocess +from test.helpers.statevector import CIRCUITS + +import numpy as np +import pytest + +from tsim.circuit import Circuit +from tsim.noise.channels import ChannelSampler +from tsim.statevector import shared +from tsim.statevector.gpu import noise, rows, tape +from tsim.statevector.shots import ShotSampler + +# 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 ``KEEP_BITS``-bit fraction, takes the slot's + alias past its threshold and XORs the outcome's flips into its row. + """ + rng = np.random.default_rng(seed) + 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 << noise.KEEP_BITS)) + at = int(first[k]) + channel * n + slot + record = groups.records[at] + if fraction >= noise.keep_threshold(record): + record = groups.records[at - slot + (int(record["alias"]) >> 25)] + extra = int(record["extra"]) + words = [int(record["words"]) & 0xFFFF, int(record["words"]) >> 16] + 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 + + +@pytest.mark.parametrize( + "channel", + [ + "HERALDED_PAULI_CHANNEL_1(0.1, 1e-10, 0, 0) 0", + "PAULI_CHANNEL_2(0.1, 0, 0, 1e-10, 0.05, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0) 0 1", + "PAULI_CHANNEL_2(0.3, 0, 0, 0.2, 1e-12, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0) 0 1", + "DEPOLARIZE2(0.03) 0 1", + ], +) +def test_noise_tables_keep_every_outcome_probability(channel): + # NOTE: with 24-bit keep thresholds, an outcome below ~1e-8 of its channel + # rounded to zero, and the GPU never drew it. + detectors = 3 if channel.startswith("HERALDED") else 2 + text = f"R 0 1\n{channel}\nM 0 1\n" + "\n".join( + f"DETECTOR rec[{-i}]" for i in range(1, detectors + 1) + ) + host, channels = _shots(text) + tables = noise.noise_tables( + channels, host.layout.noise, host.layout.num_words, 14 * 1024 + ) + fire, full = channels._tables, 1 << noise.KEEP_BITS + lo, hi = int(fire.outcome_start[0]), int(fire.outcome_start[1]) + n = hi - lo + assert n > 1 + drawn, exact = np.zeros(n), np.zeros(n) + for k in range(n): + keep = noise.keep_threshold(tables.records[lo + k]) + alias = int(tables.records[lo + k]["alias"]) >> 25 + drawn[k] += keep / (n * full) + drawn[alias] += (full - keep) / (n * full) + exact[k] += fire.alias_prob[lo + k] / n + exact[fire.alias_index[lo + k] - lo] += (1 - fire.alias_prob[lo + k]) / n + assert ((drawn > 0) == (exact > 0)).all() + np.testing.assert_allclose(drawn, exact, rtol=1e-12, atol=2.0**-40) + + +def test_noise_tables_size_blocks_for_their_memory(): + host, channels = _shots(NOISY) + layout = host.layout + 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 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) + 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(): + 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_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 + 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(count - 3), + 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") + ) + + +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) + + +GXX = shutil.which("g++") + +# NOTE: a tape in C++ for one thread per shot (tb = 0): no shuffles. ``draw_one`` +# follows the CPU executor's ``shared._cut`` in double. +CPP_DIALECT = tape.Dialect( + inline="static inline", + noinline="__attribute__((noinline)) static", + private="", + deviation="unsigned*", + constant="static const", + float2="float2", + rsqrt="rsqrt", + fmax="std::fmax", + u64="ULL", + shuffle2="{value}", + shuffle1="{value}", + lane_mask=False, + goto=False, + amplitude_chunk=32, + pair_projections=True, +) + +CPP_PRELUDE = r""" +#include +#include +typedef unsigned long long u64; +struct float2 { + float x, y; + float2() = default; + float2(float x_, float y_) : x(x_), y(y_) {} +}; +#define GOLDEN 0x9E3779B97F4A7C15ULL +static inline u64 mix64(u64 z) { + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + return z ^ (z >> 31); +} +static inline u64 draw(u64 seed, u64 shot, u64 slots, u64 slot) { + return mix64(seed + (shot * slots + slot) * GOLDEN + GOLDEN); +} +static inline int par(u64 v) { return __builtin_popcountll(v) & 1; } +static inline u64 threshold(double p) { + double scaled = p * 18446744073709551616.0; + return scaled < 1.8446744073709550e19 ? (u64)scaled : 0xFFFFFFFFFFFFF800ULL; +} +static inline int draw_one(u64 r, float w0, float w1, float) { + double total = (double)w0 + (double)w1; + if (w1 <= w0) return r < threshold(w1 / total); + return r < 0ULL - threshold(w0 / total); +} +static inline float rsqrt(float x) { return 1.f / std::sqrt(x); } +static inline float2 cmul(float2 p, float2 q) { + return float2(p.x * q.x - p.y * q.y, p.x * q.y + p.y * q.x); +} +static inline float2 cmac(float2 a, float c, float s, float2 q) { + return float2(std::fmaf(c, a.x, s * q.x), std::fmaf(c, a.y, s * q.y)); +} +static inline float2 cmac_i(float2 a, float c, float s, float2 q) { + return float2(std::fmaf(c, a.x, -s * q.y), std::fmaf(c, a.y, s * q.x)); +} +static inline void deviate(unsigned* deviation, float total) { + float d = std::fabs(total - 1.f); + unsigned bits; + std::memcpy(&bits, &d, 4); + if (!(d <= 1e-3f) && bits > *deviation) *deviation = bits; +} +""" + + +def _compile_tapes(sources: list[str], directory) -> list: + """Compile tape functions ``component_0`` into ``run_(words, shots, nwords, seed, grave, ev, nev, deviation)``.""" + parts = [CPP_PRELUDE] + for k, source in enumerate(sources): + parts.append( + f"namespace v{k} {{\n{source}\n" + f'extern "C" void run_{k}(u64* words, long shots, long nwords, u64 seed,' + " int* grave, float* ev, long nev, unsigned* deviation) {\n" + " for (long s = 0; s < shots; ++s)\n" + " grave[s] = component_0(words + s * nwords, (u64)s, seed, 0," + " nev ? ev + s * nev : nullptr, deviation);\n}\n}" + ) + cpp, lib = directory / "tapes.cpp", directory / "tapes.so" + cpp.write_text("\n".join(parts)) + # NOTE: -O0 compiles the unrolled tapes several times faster; ``cmac`` and + # ``cmac_i`` fuse explicitly, so the arithmetic does not depend on it. + assert GXX is not None + flags = ["-std=c++17", "-O0", "-ffp-contract=off", "-w", "-shared", "-fPIC"] + subprocess.run([GXX, *flags, "-o", str(lib), str(cpp)], check=True) # noqa: S603 + dll = ctypes.CDLL(str(lib)) + functions = [] + for k in range(len(sources)): + f = getattr(dll, f"run_{k}") + f.restype = None + functions.append(f) + return functions + + +def _run_tape(f, words, seed, nev): + words = words.copy() + shots, nwords = words.shape + grave = np.empty(shots, np.int32) + ev = np.zeros((shots, max(nev, 1)), np.float32) + deviation = np.zeros(1, np.uint32) + f( + words.ctypes.data_as(ctypes.c_void_p), + ctypes.c_long(shots), + ctypes.c_long(nwords), + ctypes.c_uint64(seed), + grave.ctypes.data_as(ctypes.c_void_p), + ev.ctypes.data_as(ctypes.c_void_p), + ctypes.c_long(nev), + deviation.ctypes.data_as(ctypes.c_void_p), + ) + assert tape.deviation(deviation[0]) == 0.0 + return words, grave, ev[:, :nev] + + +def _looping_circuit(n=5, layers=5, seed=0): + """A width-``n`` component with every kind of action that a looping tape runs. + + T gates (diagonal rotations), U3 (others), X, Z and Pauli-product mid-circuit + measurements with feedforward, noise, postselected detectors, and two + observables. + """ + rng = np.random.default_rng(seed) + lines = ["R " + " ".join(map(str, range(n))), "H " + " ".join(map(str, range(n)))] + for layer in range(layers): + for q in range(n): + t = rng.uniform(-1, 1, 3) + lines.append(f"U3({t[0]:.4f},{t[1]:.4f},{t[2]:.4f}) {q}") + lines.append("T " + " ".join(map(str, range(0, n, 2)))) + perm = rng.permutation(n) + lines.append("CX " + " ".join(map(str, perm[: n & ~1]))) + lines.append(f"DEPOLARIZE1(0.02) {rng.integers(n)}") + q = int(rng.integers(n)) + if layer % 3 == 0: + lines += [f"MX {q}", f"CZ rec[-1] {(q + 1) % n}"] + elif layer % 3 == 1: + lines += [f"MPP X{q}*Z{(q + 2) % n}*Y{(q + 4) % n}"] + else: + lines += [f"M {q}", f"CX rec[-1] {(q + 3) % n}"] + lines.append("M " + " ".join(map(str, range(n)))) + lines += [ + f"DETECTOR rec[-{n + 1}]", + "DETECTOR rec[-1] rec[-2]", + "OBSERVABLE_INCLUDE(0) rec[-3]", + "OBSERVABLE_INCLUDE(1) rec[-4] rec[-5]", + ] + return "\n".join(lines) + + +@pytest.mark.skipif(GXX is None, reason="needs g++") +@pytest.mark.parametrize( + "text", + [_looping_circuit(), CIRCUITS["midcircuit_feedforward"], CIRCUITS["t_ghz"]], + ids=["looping", "feedforward", "t_ghz"], +) +def test_looping_and_unrolled_tapes_run_as_the_cpu_executor( + text, tmp_path, monkeypatch +): + host, _ = _shots(text) + layout = host.layout + num_detectors = Circuit(text).num_detectors + plan = host.plan(np.ones(num_detectors, bool)) + (component,) = layout.components + (program,) = plan.programs + observables = { + o: index - num_detectors + for o, index in enumerate(component.output_indices) + if index >= num_detectors + } + assert observables + # NOTE: unrolled in each dialect variant, and looping. + variants = [ + CPP_DIALECT, + tape.Dialect(**{**CPP_DIALECT.__dict__, "goto": True}), + tape.Dialect(**{**CPP_DIALECT.__dict__, "amplitude_chunk": None}), + tape.Dialect(**{**CPP_DIALECT.__dict__, "pair_projections": False}), + ] + assert not tape.loops(program, 0, 0) + sources = [tape.tape_function(program, 0, 0, observables, d) for d in variants] + monkeypatch.setattr(tape, "UNROLL_BUDGET", 0) + assert tape.loops(program, 0, 0) + looping = tape.tape_function(program, 0, 0, observables, CPP_DIALECT) + assert "#pragma unroll 1" in looping + *unrolled, looped = _compile_tapes([*sources, looping], tmp_path) + + shots, seed = 2000, 0x1234_5678_9ABC_DEF0 + rng = np.random.default_rng(7) + words = rng.integers(0, 2**63, (shots, layout.num_words), dtype=np.uint64) + words <<= np.uint64(1) + words |= rng.integers(0, 2, words.shape, dtype=np.uint64) + expected = words.copy() + alive = np.ones(shots, np.bool_) + shared.run_lanes( + shots << program.width, + expected, + alive, + 0, + np.uint64(seed), + True, + program, + False, + 1 << 60, + 1, + ) + assert 0 < alive.sum() < shots + got, grave, ev = _run_tape(looped, words, seed, len(observables)) + np.testing.assert_array_equal(grave == np.iinfo(np.int32).max, alive) + np.testing.assert_array_equal(got[alive], expected[alive]) + assert np.abs(ev).max() > 0 + for f in unrolled: + other, other_grave, other_ev = _run_tape(f, words, seed, len(observables)) + np.testing.assert_array_equal(other_grave, grave) + np.testing.assert_array_equal(other, got) + np.testing.assert_allclose(other_ev, ev, atol=1e-5) + + +@pytest.mark.parametrize(("p", "qubits"), [(2e-8, 192), (1e-9, 3200)]) +def test_tiny_probability_segments_take_the_kernels_fine_path(p, qubits): + # NOTE: a 24-bit skip reaches only some trials of a segment of noise below + # 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 + + +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))) + + +@pytest.mark.parametrize("angle, runs", [(0.25, True), (1e-4, True), (1e-5, False)]) +def test_gpu_runs_only_rotations_whose_rare_branches_it_draws(angle, runs): + # NOTE: R_Z(1e-5) between Hadamards gives outcome 1 a probability of 2.5e-10, + # too close to what float32 rounding leaves a certain outcome (CERTAIN). + host, _ = _shots(f"R 0\nH 0\nR_Z({angle}) 0\nH 0\nM 0\nDETECTOR rec[-1]") + assert tape.supports(host.layout) == runs diff --git a/test/unit/statevector/test_program.py b/test/unit/statevector/test_program.py index 673134d8..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 @@ -52,6 +54,76 @@ 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 _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, 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