diff --git a/.claude/skills/benchmarking-tsim/SKILL.md b/.claude/skills/benchmarking-tsim/SKILL.md index e2dfb6de..fdfc0ffb 100644 --- a/.claude/skills/benchmarking-tsim/SKILL.md +++ b/.claude/skills/benchmarking-tsim/SKILL.md @@ -82,6 +82,32 @@ uv run tsim-dev bench --side .:threads=4 --side origin/main:threads=4 - Compile time is compared too; when the sides differ only in threads, read the rate column. +## Apple-silicon Macs (Metal) + +On a Mac with an M-series chip, the default strategy samples on the GPU through Metal +whenever every component has at most 11 active qubits (k_max <= 11). A plain `uv sync` +installs it. That is what users of a MacBook get, and it is much faster than the CPU: on an +M4 Pro, 26-96x one CPU thread and 4-19x all ten performance cores, bit-packed. End-to-end +cultivation d=5 ran at 125M shots/s, 19x the ten cores. Circuits without noise whose CPU +trajectories share well, such as qv_q10, run at the CPU's rate. + +The bench sets `TSIM_DEVICE=cpu` by default, so on a Mac a plain run measures the CPU only. +Say which you measured, and on a Mac measure both: + +```bash +uv run tsim-dev bench --device gpu --ref REF # a GPU change; REF must have the Metal backend +uv run tsim-dev bench --side .:device=gpu --side .:device=cpu,threads=10 # GPU vs all cores +``` + +- The `device` column must say `metal` on the GPU side. It says `cpu` where Metal falls + back: a component wider than 11, noise beyond the GPU tables, no GPU memory. +- A reference without the Metal backend samples on the CPU even with `--device gpu`; its + rows end in `DEVICE: cpu -> metal`, and the ratio compares devices, not code. +- Compare the GPU with all CPU cores (`--threads N`, from the `NOTE` line), not one. +- The GPU's own time per kernel: set `tsim.statevector.metal.runtime.PROFILE = {}` before + a `sample` call, which then runs each dispatch alone and adds its GPU seconds under the + kernel's name (`noise`, `tape`, `gather_rows`, ...). + ## What the command controls for you - One thread, pinned to a core (`--threads N` to change; say so when you report). diff --git a/.claude/skills/profiling-tsim/SKILL.md b/.claude/skills/profiling-tsim/SKILL.md index a95fbf59..7ee9c0a4 100644 --- a/.claude/skills/profiling-tsim/SKILL.md +++ b/.claude/skills/profiling-tsim/SKILL.md @@ -38,8 +38,9 @@ Tsim is slow: performance cores by default (`numba.get_num_threads()`); the `NOTE` line says how many. Quote all-core rates with `--threads N`. - **Device.** The harness sets `TSIM_DEVICE=cpu`, so every side samples on the - CPU unless you pass `--device gpu` (or `--side REF:device=gpu`). With a GPU - backend installed, `--device gpu` measures what users get by default. The + CPU unless you pass `--device gpu` (or `--side REF:device=gpu`). On an + Apple-silicon Mac, Tsim samples on the GPU by default for k_max <= 11, 4-19x + faster than all CPU cores; there `--device gpu` measures what users get. The `device` column says where each row sampled: a GPU falls back to the CPU for components wider than it runs, and a row that ends in `DEVICE: cpu -> metal` compares a GPU with the CPU. @@ -61,7 +62,8 @@ uv run tsim-dev profile cultivation_d5 --seconds 5 The run is a fresh process, one pinned thread, warm numba cache, under a memory cap, with bit-packed output on the CPU (`--threads`, `--device` and `--sampler` as in step 2). py-spy sees the host only: on a GPU the sampling phase shows the wait for -the device, and the time inside its kernels needs the backend's own timers. +the device. For Metal's time per kernel, see "Apple-silicon Macs" in the +`benchmarking-tsim` skill (`runtime.PROFILE`). The output splits samples by phase (`compile`, `sample`) and prints three tables: - **innermost Tsim frame (kernel callers):** read this first. Jitted numba code has diff --git a/AGENTS.md b/AGENTS.md index 4d00d861..7b7df498 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -66,9 +66,10 @@ more than a handful of T gates: their compile memory grows exponentially. may change in any release; every other module is internal. - **Raw outputs:** detectors and observables are raw parities, unlike Stim's flips relative to a reference. The reference-sample flags change only that. -- **Seeds:** with `active_statevector`, a seed's samples do not depend on the thread count or - `batch_size`. Do not break this; noise blocks and per-shot draws are seeded by call and shot - index (`shots.py`). +- **Seeds:** with `active_statevector` on the CPU, a seed's samples do not depend on the + thread count or `batch_size`. Do not break this; noise blocks and per-shot draws are seeded + by call and shot index (`shots.py`). On a GPU (`statevector/metal/`) they depend on + `batch_size`, as documented, and differ from the CPU's. - **Direct outputs stay direct:** outputs that depend on the noise alone never go through a state vector. Most detectors of QEC circuits are direct, which is why the detector sampler is fast. diff --git a/CHANGELOG.md b/CHANGELOG.md index 97d6a672..1beda668 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- `strategy="active_statevector"` samples on the GPU of an Apple-silicon Mac, in + float32, whenever every component has at most 11 active qubits (Tsim now depends + on `pyobjc-framework-Metal` there). A seed's GPU samples differ from the CPU's and + depend on `batch_size`; `TSIM_DEVICE=cpu` keeps the samplers on the CPU, and + `SamplerStats.device` and the repr say where they sample. - The docs pages have a menu to copy a page or open it in ChatGPT or Claude, and the README, the install page and the agent skill point to the docs MCP server at `https://bloqade.quera.com/tsim/mcp`. The README and the install page also diff --git a/README.md b/README.md index 008c4b7b..b9c694c0 100644 --- a/README.md +++ b/README.md @@ -23,8 +23,8 @@ Or with [uv](https://docs.astral.sh/uv/): uv add bloqade-tsim ``` -The default sampling strategy runs on the CPU. To run the stabilizer-rank strategies on an NVIDIA GPU, -install the JAX CUDA extra: +The default sampling strategy runs on the CPU, and on the GPU of an Apple-silicon Mac (no extra +needed). To run the stabilizer-rank strategies on an NVIDIA GPU, install the JAX CUDA extra: ```bash pip install "bloqade-tsim[cuda13]" # or bloqade-tsim[cuda12] ``` @@ -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 Apple GPU (Metal) | Number of active qubits (at most 30) | Most circuits, especially with postselection | | `"cat5"`, `"bss"`, `"cutting"` | Stabilizer rank decomposition | CPU or GPU (JAX) | Number of non-Clifford gates | Few non-Clifford gates acting on many qubits; GPU throughput | ```python diff --git a/codecov.yml b/codecov.yml new file mode 100644 index 00000000..94b8ad40 --- /dev/null +++ b/codecov.yml @@ -0,0 +1,4 @@ +# NOTE: the Metal executor needs an Apple GPU, which CI does not have; its tests +# run on Apple silicon only. +ignore: + - "src/tsim/statevector/metal" diff --git a/docs/index.mdx b/docs/index.mdx index 21c5681a..e51a566e 100644 --- a/docs/index.mdx +++ b/docs/index.mdx @@ -39,8 +39,8 @@ samples = detector_sampler.sample(shots=100) uv add bloqade-tsim ``` -The default sampling strategy runs on the CPU. To run the stabilizer-rank strategies on an NVIDIA GPU, -install the JAX CUDA extra: +The default sampling strategy runs on the CPU, and on the GPU of an Apple-silicon Mac (no extra +needed). To run the stabilizer-rank strategies on an NVIDIA GPU, install the JAX CUDA extra: ```bash uv add "bloqade-tsim[cuda13]" @@ -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) | Number of active qubits (at most 30) | Most circuits, especially with postselection | +| `"active_statevector"` (default) | Dense state vector over the active qubits of each component, as in [Clifft](https://arxiv.org/abs/2604.27058) and [SymFT](https://arxiv.org/abs/2607.28600) | CPU (numba, multi-threaded), or an Apple GPU (Metal) | Number of active qubits (at most 30) | Most circuits, especially with postselection | | `"cat5"`, `"bss"`, `"cutting"` | Stabilizer rank decomposition | CPU or GPU (JAX) | Number of non-Clifford gates | Few non-Clifford gates acting on many qubits; GPU throughput | ```python diff --git a/docs/install.mdx b/docs/install.mdx index eeb3e10f..89138f0b 100644 --- a/docs/install.mdx +++ b/docs/install.mdx @@ -1,6 +1,6 @@ --- title: "Installation" -description: "Install Tsim with uv or pip, with optional CUDA support." +description: "Install Tsim with uv or pip, with optional GPU support." --- Tsim requires Python 3.11 or later. @@ -35,6 +35,34 @@ For GPU acceleration with CUDA: pip install "bloqade-tsim[cuda13]" ``` +## Apple GPUs + +On an Apple-silicon Mac, the default strategy (`"active_statevector"`) samples on the GPU +through Metal; the plain install includes PyObjC's Metal bindings for it. The samplers use the +GPU whenever every component of the circuit has at most 11 active qubits, and otherwise sample +on the CPU, with a warning. Postselection runs on the GPU, so heavily postselected circuits +gain the most. A sampler's repr says where it samples (`samples on metal` or `samples on +cpu`), and so does `sampler.stats.device`. + +The GPU computes outcome probabilities in float32, so a sample follows the CPU's distribution +to within float32 rounding. Float32 cannot tell a measurement outcome of probability below +about 1e-11 from rounding error, so the GPU never draws one; a circuit with a rotation small +enough to make such an outcome (an `R_Z` or `R_X` by less than about 1e-4 π) samples on the +CPU, with a warning. A seed's samples on the GPU differ from the CPU's and depend on +`batch_size`. + +### Sampling on the CPU instead + +Set `TSIM_DEVICE=cpu` to keep the samplers on the CPU, in double precision, for example to +reproduce samples from another machine: + +```bash +TSIM_DEVICE=cpu python my_script.py +``` + +Tsim reads it whenever it compiles a sampler, so `os.environ["TSIM_DEVICE"] = "cpu"` before +`compile_sampler` or `compile_detector_sampler` works too. + ## Agent skill 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..5ef984d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,8 @@ dependencies = [ "stim>=1.15.0", "stim-timeline-text>=0.1.0", "zxparam>=0.1.0", + # NOTE: the statevector executor on Apple GPUs (about 7 MB). + "pyobjc-framework-Metal>=10.0; sys_platform == 'darwin' and platform_machine == 'arm64'", ] [project.urls] @@ -167,6 +169,9 @@ extend_exclude = [ "^test/", ] +[tool.deptry.package_module_name_map] +pyobjc-framework-Metal = "Metal" + [tool.deptry.per_rule_ignores] DEP002 = ["ipywidgets"] diff --git a/skills/simulating-circuits-with-tsim/SKILL.md b/skills/simulating-circuits-with-tsim/SKILL.md index e5cfb39b..b3bfbc4c 100644 --- a/skills/simulating-circuits-with-tsim/SKILL.md +++ b/skills/simulating-circuits-with-tsim/SKILL.md @@ -189,8 +189,14 @@ yourself. default. Use `numba.set_num_threads(k)` before `sample` (or `NUMBA_NUM_THREADS=k` before importing Tsim) and leave cores free on a shared machine. Do not run several multi-threaded Tsim processes side by side; they oversubscribe the CPU. +- **Apple GPUs:** on an Apple-silicon Mac the default strategy samples on the GPU when every + component has at most 11 active qubits, often 5 to 10 times faster than all CPU cores. It + computes outcome probabilities in float32 and never draws an outcome below about 1e-11 + (circuits with rotations by less than about 1e-4 π stay on the CPU, with a warning); set + `TSIM_DEVICE=cpu` before compiling to stay on the CPU, in double precision. The sampler's repr says where it samples (`samples on metal`). - **Seeds:** a `seed` reproduces samples for the same Tsim and dependency versions and - backend. 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..816109d2 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -954,7 +954,9 @@ 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 + "bss", "cutting"). "active_statevector" samples in float32 on an + Apple GPU when every component has at most 11 active qubits + (``TSIM_DEVICE=cpu`` turns this off), else on the CPU, on ``numba.get_num_threads()`` threads. The stabilizer rank strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed @@ -964,7 +966,8 @@ def compile_sampler( non-direct components, the samples of a seed depend on neither the batch size nor the thread count. With a stabilizer-rank strategy they also need the same explicit ``batch_size``, since the default batch size - depends on the free memory. + depends on the free memory. On a GPU, they do depend on the batch + size. Returns: A CompiledMeasurementSampler that can be used to sample measurements. @@ -991,7 +994,9 @@ 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 + "bss", "cutting"). "active_statevector" samples in float32 on an + Apple GPU when every component has at most 11 active qubits + (``TSIM_DEVICE=cpu`` turns this off), else on the CPU, on ``numba.get_num_threads()`` threads. The stabilizer rank strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed @@ -1002,6 +1007,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..10d08ea4 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 import metal from tsim.statevector.program import StatevectorComponent from tsim.statevector.shots import ShotSampler from tsim.utils.cuda_helpers import copy_d2h @@ -28,6 +30,7 @@ from jax import Array as PRNGKey from tsim.circuit import Circuit + from tsim.statevector.metal.sampler import MetalShotSampler # NOTE: postselection with a stabilizer-rank strategy screens shots for direct # discards in rounds of SCREEN_CHUNKS chunks of at most SCREEN_SHOTS shots. @@ -308,6 +311,7 @@ class SamplerStats: error_channel_bits: int peak_bytes_per_shot: int compiled_bytes: int + device: str = "cpu" class _CompiledSamplerBase: @@ -331,7 +335,9 @@ 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 + "bss", "cutting"). "active_statevector" samples in float32 on an + Apple GPU when every component has at most 11 active qubits + (``TSIM_DEVICE=cpu`` turns this off), else on the CPU, on ``numba.get_num_threads()`` threads. The stabilizer rank strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed @@ -341,14 +347,18 @@ def __init__( non-direct components, the samples of a seed depend on neither the batch size nor the thread count. With a stabilizer-rank strategy they also need the same explicit ``batch_size``, since the default batch size - depends on the free memory. + depends on the free memory. On a GPU, they do depend on the batch + size. """ if seed is None: seed = secrets.randbits(128) - # NOTE: the noise, the shot words and the JAX key draw from children of the - # whole seed, so different seeds give independent samples. - noise_seed, shots_seed, key_seed = np.random.SeedSequence(seed).spawn(3) + # NOTE: the noise, the shot words, the JAX key and the GPU draw from children + # of the whole seed, so different seeds give independent samples. The first + # three are those of a spawn(3), so the CPU's samples are unchanged. + noise_seed, shots_seed, key_seed, gpu_seed = np.random.SeedSequence(seed).spawn( + 4 + ) prepared = prepare_graph(circuit, sample_detectors=sample_detectors) self._program = compile_program(prepared, mode=mode, strategy=strategy) @@ -364,15 +374,25 @@ def __init__( program = self._program # NOTE: programs of statevector components, and programs without components # under any strategy, are sampled as shot words; the others with JAX. - self._shots: ShotSampler | None = None + self._shots: ShotSampler | MetalShotSampler | None = None self._backend: _StabilizerRankBackend | None = None if strategy == "active_statevector" or not program.components: - self._shots = ShotSampler( + host = ShotSampler( program, self._channel_sampler, self._num_detectors, shots_seed, ) + self._shots = host + # NOTE: as JAX does for the stabilizer-rank strategies, a GPU is used + # whenever there is one. + if ( + mode == "sequential" + and os.environ.get("TSIM_DEVICE", "").lower() != "cpu" + ): + self._shots = metal.shot_sampler( + host, self._channel_sampler, self._num_detectors, gpu_seed + ) else: self._backend = _StabilizerRankBackend(program, key_seed) @@ -428,7 +448,7 @@ def _sample_outputs( # NOTE: only the shot words are filtered by a mask; the noise sampler # writes direct outputs straight into boolean rows otherwise. if not (self._program.components or bit_packed or mask is not None): - samples = self._shots.direct_bools(shots, flip) + samples = self._shots.direct_bools(shots, flip, batch_size) return [_take_columns(samples, c) for c in columns] arrays, deviation = self._shots.sample( shots, @@ -721,6 +741,9 @@ def stats(self) -> SamplerStats: - ``peak_bytes_per_shot`` (``int``): estimated peak memory of one shot, by which the default batch sizes are chosen. - ``compiled_bytes`` (``int``): memory of the compiled components. + - ``device`` (``str``): where ``"active_statevector"`` samples: + ``"metal"`` on an Apple GPU, else ``"cpu"`` (a stabilizer-rank + strategy runs wherever JAX does). """ components = self._program.components widths: tuple[int, ...] | None = None @@ -751,6 +774,7 @@ def stats(self) -> SamplerStats: ), peak_bytes_per_shot=peak, compiled_bytes=compiled_bytes, + device=self._shots.device if self._shots is not None else "cpu", ) def __repr__(self) -> str: @@ -803,6 +827,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 += [ @@ -850,7 +876,8 @@ def __init__( non-direct components, the samples of a seed depend on neither the batch size nor the thread count. With a stabilizer-rank strategy they also need the same explicit ``batch_size``, since the default batch size - depends on the free memory. + depends on the free memory. On a GPU, they do depend on the batch + size. """ super().__init__( @@ -951,6 +978,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/gpu/tape.py b/src/tsim/statevector/gpu/tape.py index 20e86449..22e6a97a 100644 --- a/src/tsim/statevector/gpu/tape.py +++ b/src/tsim/statevector/gpu/tape.py @@ -179,10 +179,10 @@ def rarest_branch(layout: ShotLayout) -> float: 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.""" +def threads_per_shot(layout: ShotLayout, amplitude_axes: int = AMPLITUDE_AXES) -> int: + """Return log2 of a shot's threads: one per ``2^amplitude_axes`` amplitudes, at most 32.""" width = max((c.width for c in layout.components), default=0) - return min(5, max(0, width - AMPLITUDE_AXES)) + return min(5, max(0, width - amplitude_axes)) def deviation(bits) -> float: diff --git a/src/tsim/statevector/metal/__init__.py b/src/tsim/statevector/metal/__init__.py new file mode 100644 index 00000000..541064a5 --- /dev/null +++ b/src/tsim/statevector/metal/__init__.py @@ -0,0 +1,82 @@ +"""The statevector executor on an Apple GPU, through Metal. + +The noise sampler (``noise``) and every component's tape (``codegen``, written +by ``tsim.statevector.gpu.tape``) run as Metal kernels that PyObjC compiles at +runtime (``runtime``); ``MetalShotSampler`` (``sampler``) replaces +``ShotSampler`` behind the compiled samplers whenever PyObjC sees an Apple GPU +and every component fits. Amplitudes are complex64. It needs Apple silicon, +where Tsim depends on ``pyobjc-framework-Metal``. +""" + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +from tsim.statevector.gpu import FORKED +from tsim.statevector.gpu.tape import MAX_WIDTH, RESOLVABLE, rarest_branch, supports +from tsim.statevector.metal import runtime + +if TYPE_CHECKING: + import numpy as np + + from tsim.noise.channels import ChannelSampler + from tsim.statevector.metal.sampler import MetalShotSampler + from tsim.statevector.shots import ShotSampler + +# NOTE: the warnings point at the caller of compile_sampler. +_STACKLEVEL = 5 + + +def shot_sampler( + host: ShotSampler, + channels: ChannelSampler, + num_detectors: int, + seed: np.random.SeedSequence, +) -> ShotSampler | MetalShotSampler: + """Return a Metal shot sampler of ``host``'s program when one can run it, else ``host``. + + It warns when an Apple GPU is there but cannot be used: in a process forked + after its parent used the GPU, for a component wider than ``MAX_WIDTH`` or a + rotation by an angle too small for float32 amplitudes (``RESOLVABLE``), for + noise beyond the GPU tables' limits, and when the sampler cannot be set up + for lack of GPU memory. + """ + 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) + if width > MAX_WIDTH: + reason = ( + f"a component needs {width} active qubits, more than the " + f"{MAX_WIDTH} that the GPU executor supports" + ) + else: + reason = ( + f"a rotation has a branch of weight {rarest_branch(host.layout):.1e}, " + f"below the {RESOLVABLE:.0e} that float32 amplitudes resolve" + ) + warnings.warn( + f"An Apple GPU is available, but {reason}; sampling on the CPU.", + stacklevel=_STACKLEVEL, + ) + return host + from tsim.statevector.metal.sampler import MetalShotSampler + + try: + return MetalShotSampler(host, channels, num_detectors, seed) + except MemoryError: + warnings.warn( + "The GPU is out of memory; sampling on the CPU.", stacklevel=_STACKLEVEL + ) + except ValueError as error: + # NOTE: noise beyond the GPU tables' limits (channels of more than 127 + # outcomes, outcomes that flip more than 255 words). + warnings.warn( + f"The GPU cannot sample this circuit ({error}); sampling on the CPU.", + stacklevel=_STACKLEVEL, + ) + return host diff --git a/src/tsim/statevector/metal/codegen.py b/src/tsim/statevector/metal/codegen.py new file mode 100644 index 00000000..4fb39658 --- /dev/null +++ b/src/tsim/statevector/metal/codegen.py @@ -0,0 +1,91 @@ +"""Metal source of a component's tape (see ``tsim.statevector.gpu.tape``). + +The Metal Shading Language has SIMD-group shuffles for warp shuffles, no +``goto`` (a pruned shot breaks out of a ``do``-block) and no double precision: +near the boundary ``draw_one`` refines the rarer outcome's share as a +float-float ``q + e / total`` instead. Unrolled tapes keep their amplitudes in +arrays of 32, which Metal's compiler keeps in registers, and project partner +amplitudes in pairs, which keeps fewer values live. +""" + +from __future__ import annotations + +from tsim.statevector.gpu import tape + +METAL = tape.Dialect( + inline="inline", + noinline="__attribute__((noinline))", + private="thread ", + deviation="device atomic_uint*", + constant="constant", + float2="float2", + rsqrt="rsqrt", + fmax="fmax", + u64="UL", + shuffle2="simd_shuffle_xor({value}, (ushort){lanes})", + shuffle1="simd_shuffle_xor({value}, (ushort){lanes})", + lane_mask=False, + goto=False, + # NOTE: Metal's compiler moves a thread array of 64 complex amplitudes (a + # width-11 component on 32 threads) to the stack, which costs the tape 1.8x. + amplitude_chunk=32, + pair_projections=True, +) + +PRELUDE = r""" +#include +using namespace metal; +typedef ulong u64; +#define GOLDEN 0x9E3779B97F4A7C15UL + +inline u64 mix64(u64 z) { + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9UL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBUL; + return z ^ (z >> 31); +} +// The CPU executor's draw: splitmix64 at index shot * slots + slot of stream seed. +inline u64 draw(u64 seed, u64 shot, u64 slots, u64 slot) { + return mix64(seed + (shot * slots + slot) * GOLDEN + GOLDEN); +} +inline int par(u64 v) { return popcount(v) & 1; } +// draw_one against the rarer outcome's share, which float32 keeps to full +// relative precision: a rare 0 takes the top of the range, as on the CPU +// (``shared._cut``). The share is refined to about 46 bits as q + e / total +// (Metal has no double): q correctly rounded, e its exact residual. Not +// inlined, as it is rarely taken. +__attribute__((noinline)) int draw_one_exact(u64 r, float w0, float w1, float total) { + int rare0 = w0 < w1; + float w = rare0 ? w0 : w1; + float q = precise::divide(w, total); + float e = fma(-q, total, w); + u64 below = u64(q * 0x1p64f) + u64(long(precise::divide(e, total) * 0x1p64f)); + return ((rare0 ? ~r : r) < below) ^ rare0; +} +// Whether the 64-bit draw r gives outcome 1, of weight w1 against w0. float32 +// decides it unless r / 2^64 lies within 2^-21 of the share w1 / total, which +// covers both roundings and fast division's error; there draw_one_exact does. +inline int draw_one(u64 r, float w0, float w1, float total) { + float p = w1 / total, x = float(r) * 0x1p-64f; + if (x < p * (1.f - 0x1p-21f)) return 1; + if (x > p * (1.f + 0x1p-21f)) return 0; + return draw_one_exact(r, w0, w1, total); +} +// A uniform in (0, 1] from the high 24 bits, so that log never sees 0. +inline float unit(u64 r) { return float((r >> 40) + 1) * (1.0f / 16777216.0f); } +inline float2 cmul(float2 p, float2 q) { + return float2(p.x * q.x - p.y * q.y, p.x * q.y + p.y * q.x); +} +// c a + s q and c a + i s q +inline float2 cmac(float2 a, float c, float s, float2 q) { + return float2(fma(c, a.x, s * q.x), fma(c, a.y, s * q.y)); +} +inline float2 cmac_i(float2 a, float c, float s, float2 q) { + return float2(fma(c, a.x, -s * q.y), fma(c, a.y, s * q.x)); +} +// The largest normalization deviation, as the bits of a non-negative float. +inline void deviate(device atomic_uint* deviation, float total) { + float d = fabs(total - 1.f); + if (!(d <= 1e-3f)) + atomic_fetch_max_explicit(deviation, d == d ? as_type(d) : 0x7f800000u, memory_order_relaxed); +} +""" diff --git a/src/tsim/statevector/metal/kernels.py b/src/tsim/statevector/metal/kernels.py new file mode 100644 index 00000000..147ca5bb --- /dev/null +++ b/src/tsim/statevector/metal/kernels.py @@ -0,0 +1,231 @@ +"""The Metal sampler's kernels other than the noise ones, and their source. + +``tape`` runs the components' tapes on the queued shots; ``count_bits`` and +``list_bits`` list the accepted shots from the bitmap; ``gather_rows``, +``unpack_rows`` and ``gather_float`` copy out rows, booleans and expectations; +``reset`` and ``advance`` keep a call's counters on the GPU. +""" + +from __future__ import annotations + +from tsim.statevector.gpu.noise import NoiseTables +from tsim.statevector.gpu.tape import TapeKernel +from tsim.statevector.metal.codegen import PRELUDE +from tsim.statevector.metal.noise import source as noise_source +from tsim.statevector.shots import ShotLayout + +SCAN_THREADS = 1024 +"""Threads per threadgroup of ``count_bits`` and ``list_bits``: one per bitmap word.""" + +TEMPLATE = r""" +struct TapeArgs {{ long first_shot; u64 seed; int flags; }}; + +// flags: 1 clears rejected shots from ``survivors`` and counts the accepted ones, +// 2 writes the expectations. +kernel void tape( + device u64* out [[buffer(0)]], constant TapeArgs& args [[buffer(1)]], + device const uint* queue [[buffer(2)]], device const uint* count [[buffer(3)]], + device atomic_uint* survivors [[buffer(4)]], device atomic_uint* acc_count [[buffer(5)]], + device float* expectations [[buffer(6)]], device atomic_uint* deviation [[buffer(7)]], + uint gid [[thread_position_in_grid]], uint lane [[thread_index_in_simdgroup]]) +{{ + int t = lane & {T1}; + uint q = gid >> {tb}; + if (q >= count[0]) return; + const u64 seed = args.seed; +{seeds} + uint row_index = queue[q]; + u64 shot = (u64)args.first_shot + row_index; + device u64* row = out + (long)row_index * {NW}; + u64 w[{NW}]; + for (int i = 0; i < {NW}; ++i) w[i] = row[i]; + float ev[{NO}]; + int grave = 0x7fffffff; +{calls} + if (t == 0) {{ +{stores} + if (args.flags & 2) + for (int i = 0; i < {NO}; ++i) expectations[(long)row_index * {NO} + i] = ev[i]; + }} + if (args.flags & 1) {{ + u64 hit = {hits}; + if (t == 0 && hit) + atomic_fetch_and_explicit(&survivors[row_index >> 5], ~(1u << (row_index & 31)), memory_order_relaxed); + uint ballot = uint(static_cast(simd_ballot(t == 0 && hit == 0))); + if (ballot && simd_is_first()) atomic_fetch_add_explicit(acc_count, popcount(ballot), memory_order_relaxed); + }} +}} + +// The counts of a batch: its screened survivors and its accepted shots. +kernel void reset(device uint* counts [[buffer(0)]]) {{ counts[0] = 0; counts[1] = 0; }} + +// Moves the cursor past a batch's ``*n`` rows. +kernel void advance(device uint* cursor [[buffer(0)]], device const uint* n [[buffer(1)]]) {{ + cursor[0] += *n; +}} + +inline uint block_scan(uint v, thread uint& total, threadgroup uint* warps, uint lane, uint warp, uint nwarps) {{ + uint x = simd_prefix_inclusive_sum(v); + if (lane == 31) warps[warp] = x; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (warp == 0) {{ + uint y = lane < nwarps ? warps[lane] : 0; + warps[lane] = simd_prefix_inclusive_sum(y); + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + uint before = (warp ? warps[warp - 1] : 0) + x - v; + total = warps[nwarps - 1]; + threadgroup_barrier(mem_flags::mem_threadgroup); + return before; +}} + +// The shots of the set bits of ``bits`` (``words`` words), in order: ``count_bits`` +// counts each block's {SCAN} words, ``list_bits`` lists a block's shots after +// those of the blocks before it. +kernel void count_bits( + device const uint* bits [[buffer(0)]], constant long& words [[buffer(1)]], + device uint* totals [[buffer(2)]], + uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], uint warp [[simdgroup_index_in_threadgroup]], + uint nwarps [[simdgroups_per_threadgroup]]) +{{ + threadgroup uint warps[32]; + long i = (long)bid * {SCAN} + tid; + uint total; + block_scan(i < words ? popcount(bits[i]) : 0, total, warps, lane, warp, nwarps); + if (tid == 0) totals[bid] = total; +}} + +kernel void list_bits( + device const uint* bits [[buffer(0)]], constant long& words [[buffer(1)]], + device const uint* totals [[buffer(2)]], device uint* index [[buffer(3)]], + uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]], uint warp [[simdgroup_index_in_threadgroup]], + uint nwarps [[simdgroups_per_threadgroup]]) +{{ + threadgroup uint warps[32]; + uint earlier = 0, total; + for (uint b = tid; b < bid; b += {SCAN}) earlier += totals[b]; + block_scan(earlier, earlier, warps, lane, warp, nwarps); + long i = (long)bid * {SCAN} + tid; + uint w = i < words ? bits[i] : 0; + uint at = earlier + block_scan(popcount(w), total, warps, lane, warp, nwarps); + for (; w; w &= w - 1) index[at++] = (uint)(32 * i + ctz(w)); +}} + +// The rows ``index[j]`` (or every row, without ``has_index``) for ``j < *n``, +// densely after the ``*cursor`` rows of the call's earlier batches: each row's +// first {NB} bytes, which hold its outputs bit-packed (the bits past them +// cleared). A thread writes 4 aligned bytes at a time. +kernel void gather_rows( + device const uchar* rows [[buffer(0)]], device const uint* n [[buffer(1)]], + device const uint* index [[buffer(2)]], constant int& has_index [[buffer(3)]], + device uchar* dst [[buffer(4)]], device const uint* cursor [[buffer(5)]], + uint gid [[thread_position_in_grid]], uint grid [[threads_per_grid]]) +{{ + long count = min((long)*n, (long)cursor[1] - (long)cursor[0]); + if (count <= 0) return; + long at = (long)cursor[0] * {NB}, end = at + count * {NB}; + for (long p = (at & ~3L) + 4L * gid; p < end; p += 4L * grid) {{ + // NOTE: a batch's rows span less than 4 GB (``BATCH_BYTES``). + uint k = (uint)(max(p, at) - at); + uint r = k / {NB}, b = k - r * {NB}; + uint sh = 8 * (b & 3); + if (p >= at && p + 4 <= end && b + 4 <= {NB} && (sh == 0 || (b & ~3u) + 8 <= {RB})) {{ + // NOTE: the 4 bytes lie in one row: two aligned loads at most. + long source = has_index ? (long)index[r] : r; + device const uint* w = (device const uint*)(rows + source * {RB}) + (b >> 2); + uint v = sh ? (w[0] >> sh) | (w[1] << (32 - sh)) : w[0]; + if (b + 4 == {NB}) v &= ({LAST}u << 24) | 0xFFFFFFu; + *(device uint*)(dst + p) = v; + continue; + }} + uint v = 0; + for (int j = (int)(max(p, at) - p); j < 4 && p + j < end; ++j) {{ + long source = has_index ? (long)index[r] : r; + uint x = rows[source * {RB} + b]; + if (b == {NB} - 1) x &= {LAST}; + v |= x << (8 * j); + if (++b == {NB}) {{ + b = 0; + ++r; + }} + }} + for (int j = 0; j < 4; ++j) + if (p + j >= at && p + j < end) dst[p + j] = uchar(v >> (8 * j)); + }} +}} + +// The first ``n`` packed rows of ``{NB}`` bytes, one byte per output (``{NOUT}`` +// per row), as booleans: a thread unpacks 4 outputs, which share a packed byte. +kernel void unpack_rows( + device const uchar* packed [[buffer(0)]], device uchar* dst [[buffer(1)]], + constant uint& n [[buffer(2)]], + uint gid [[thread_position_in_grid]], uint grid [[threads_per_grid]]) +{{ + const uint groups = ({NOUT} + 3) / 4; + for (uint i = gid; i < n * groups; i += grid) {{ + uint r = i / groups, g = i - r * groups; + uint b = packed[(ulong)r * {NB} + (g >> 1)] >> (4 * (g & 1)); + device uchar* out = dst + (ulong)r * {NOUT} + 4 * g; + for (uint j = 0; j < 4 && 4 * g + j < {NOUT}; ++j) out[j] = uchar((b >> j) & 1); + }} +}} + +// Observable expectations of rows ``index[j]`` (or of row ``j``, without +// ``has_index``) for ``j < *n``, densely, after the ``*cursor`` rows of the +// call's earlier batches. +kernel void gather_float( + device const float* values [[buffer(0)]], device const uint* n [[buffer(1)]], + constant int& width [[buffer(2)]], device const uint* index [[buffer(3)]], + device float* dst [[buffer(4)]], device const uint* cursor [[buffer(5)]], + constant int& has_index [[buffer(6)]], + uint gid [[thread_position_in_grid]], uint grid [[threads_per_grid]]) +{{ + long total = min((long)*n, (long)cursor[1] - (long)cursor[0]) * width; + dst += (long)*cursor * width; + for (long k = gid; k < total; k += grid) {{ + long j = k / width; int c = (int)(k - j * width); + long r = has_index ? (long)index[j] : j; + dst[k] = values[r * width + c]; + }} +}} +""" + + +def source( + layout: ShotLayout, + tables: NoiseTables, + parts: TapeKernel, + tb: int, + num_observables: int, +) -> str: + """Return the MSL source of every kernel, for tapes of ``2^tb`` threads a shot.""" + nb = -(-layout.num_outputs // 8) + return ( + PRELUDE + + noise_source(tables, layout.output_words, layout.num_outputs) + + parts.functions + + TEMPLATE.format( + T1=(1 << tb) - 1, + tb=tb, + NW=layout.num_words, + SCAN=SCAN_THREADS, + NO=num_observables, + NB=max(nb, 1), + RB=8 * layout.num_words, + NOUT=max(layout.num_outputs, 1), + LAST=hex((1 << (layout.num_outputs % 8 or 8)) - 1), + calls=parts.calls, + seeds=parts.seeds, + stores=parts.stores, + hits=parts.hits, + ) + ) + + +def names(shared_words: bool) -> list[str]: + """Return the kernels of a source; ``noise_kept`` needs ``shared_words``.""" + found = ["noise", "tape", "count_bits", "list_bits", "gather_rows"] + found += ["gather_float", "reset", "advance", "unpack_rows"] + return [*found, "noise_kept"] if shared_words else found diff --git a/src/tsim/statevector/metal/noise.py b/src/tsim/statevector/metal/noise.py new file mode 100644 index 00000000..d3d55105 --- /dev/null +++ b/src/tsim/statevector/metal/noise.py @@ -0,0 +1,394 @@ +"""Noise sampling into shot words on a Metal GPU. + +The kernels read the tables of ``tsim.statevector.gpu.noise`` (see there): a +threadgroup owns ``S`` consecutive shots, walks segments of the channel groups +with geometric skips, records the faults as events in threadgroup memory and +then XORs their outcomes' flips into the shots' words. Metal has no 64-bit +atomics, so a flip XORs the two halves of a word with 32-bit atomics, and a +threadgroup has 32 KB of threadgroup memory. Rows too wide for it keep their +words in device memory, where the faults are XORed in with device atomics. + +A block then screens its shots (constants, postselected direct detectors), +writes the survivors' rows, appends them to a queue for the tape kernel and +marks them in a bitmap of the batch's shots. Under postselection, +``noise_kept`` draws the noise that cannot change whether a shot is kept +(``gpu.noise.deferred_channels``) on the kept shots only, after the tapes. +""" + +from __future__ import annotations + +from tsim.noise.channels import ChannelSampler, WordTables +from tsim.statevector.gpu.noise import ( + NoiseTables, + noise_tables as _noise_tables, + words_fit, +) + +SHARED_BYTES = 14 * 1024 +"""Threadgroup memory of a block's words and events. + +A threadgroup may have 32 KB, but the kernel is latency bound: with 14 KB, a +GPU core keeps more blocks resident, and the noise of end-to-end cultivation +d=5 takes 7.3 instead of 11.5 ns per shot on an M4 Pro (most other circuits +gain 10-25%, distillation loses up to 15%). +""" + +_SEGMENT_PARAMS = """\ + device const int* seg_first [[buffer(2)]], device const int* seg_outcomes [[buffer(3)]], + device const int* seg_begin [[buffer(4)]], device const int* seg_end [[buffer(5)]], + device const float* seg_log_keep [[buffer(6)]], device const Record* records [[buffer(7)]], + device const Flip* overflow [[buffer(8)]], +""" + +# NOTE: the signature of both ``noise`` kernels: with the block's words in +# threadgroup memory or in device memory. +_NOISE_SIGNATURE = ( + """\ +kernel void noise( + device u64* out [[buffer(0)]], constant NoiseArgs& args [[buffer(1)]], +""" + + _SEGMENT_PARAMS + + """\ + device const u64* constants [[buffer(9)]], device const u64* screen [[buffer(10)]], + device uint* queue [[buffer(11)]], device atomic_uint* count [[buffer(12)]], + device uint* survivors [[buffer(13)]], + threadgroup uint2* events [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], + uint nt [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]]) +""" +) + +NOISE_SHARED = _NOISE_SIGNATURE + r"""{{ + threadgroup u64 sw[{NW} * {S}]; + threadgroup atomic_int num_events; + threadgroup bool passed[{S}]; + long base = (long)bid * {S}; + int ns = (int)min((long){S}, args.shots - base); + for (uint i = tid; i < {NW} * {S}; i += nt) sw[i] = 0; + if (tid == 0) atomic_store_explicit(&num_events, 0, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + Segments seg = {{seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep}}; + add_faults(sw, events, &num_events, args.max_events, ns, block_stream(args.seed, bid), + args.num_segments, seg, records, overflow, tid, nt); + // NOTE: the screen reads the words in parallel, a thread per word and shot (a + // shot's threads only ever clear its flag); the constants go into the rows. + for (uint s = tid; s < {S}; s += nt) passed[s] = (int)s < ns; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint i = tid; i < {OW} * {S}; i += nt) {{ + uint word = i / {S}, s = i - word * {S}; + if ((sw[i] ^ constants[word]) & screen[word]) passed[s] = false; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int s0 = 0; s0 < {S}; s0 += nt) {{ + int s = s0 + tid; + bool alive = s < {S} && passed[s]; + enqueue(alive, s, s0 + (tid & ~31u), base, ns, lane, args.has_survivors, queue, count, survivors); + }} + device u64* rows = out + base * {NW}; + for (int i = tid; i < ns * {NW}; i += nt) {{ + int s = i / {NW}, word = i - s * {NW}; + if (passed[s]) rows[i] = sw[word * {S} + s] ^ (word < {OW} ? constants[word] : 0); + }} +}} + +""" + +NOISE_DEVICE = _NOISE_SIGNATURE + r"""{{ + threadgroup atomic_int num_events; + threadgroup bool passed[{S}]; + long base = (long)bid * {S}; + int ns = (int)min((long){S}, args.shots - base); + device u64* rows = out + base * {NW}; + for (int i = tid; i < ns * {NW}; i += nt) rows[i] = 0; + if (tid == 0) atomic_store_explicit(&num_events, 0, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup); + Segments seg = {{seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep}}; + add_faults(rows, events, &num_events, args.max_events, ns, + block_stream(args.seed, bid), args.num_segments, seg, records, + overflow, tid, nt); + // NOTE: as in the kernel above, the screen reads the words in parallel. + for (uint s = tid; s < {S}; s += nt) passed[s] = (int)s < ns; + threadgroup_barrier(mem_flags::mem_threadgroup); + for (int i = tid; i < ns * {OW}; i += nt) {{ + int s = i / {OW}, word = i - s * {OW}; + device u64* at = rows + (long)s * {NW} + word; + u64 v = *at ^ constants[word]; + *at = v; + if (v & screen[word]) passed[s] = false; + }} + threadgroup_barrier(mem_flags::mem_device | mem_flags::mem_threadgroup); + for (int s0 = 0; s0 < {S}; s0 += nt) {{ + int s = s0 + tid; + bool alive = s < {S} && passed[s]; + enqueue(alive, s, s0 + (tid & ~31u), base, ns, lane, args.has_survivors, queue, count, survivors); + }} +}} + +""" + +NOISE_KEPT = ( + r""" +struct KeptArgs {{ u64 seed; int num_segments; int max_events; }}; + +// The deferred channels' faults on the kept shots ``index`` (``*n`` of them), +// then their returned outputs bit-packed (``pack_rows``). +kernel void noise_kept( + device const u64* out [[buffer(0)]], constant KeptArgs& args [[buffer(1)]], +""" + + _SEGMENT_PARAMS + + r""" device const uint* n [[buffer(9)]], device const uint* index [[buffer(10)]], + device uchar* dst [[buffer(11)]], device const uint* cursor [[buffer(12)]], + threadgroup uint2* events [[threadgroup(0)]], + uint tid [[thread_position_in_threadgroup]], uint bid [[threadgroup_position_in_grid]], + uint nt [[threads_per_threadgroup]]) +{{ + threadgroup u64 sw[{OW} * {S}]; + threadgroup atomic_int num_events; + long base = (long)bid * {S}; + if (base >= (long)*n) return; + // NOTE: rows past the call's capacity (cursor[1]) are not written. + int ns = (int)min(min((long){S}, (long)*n - base), (long)cursor[1] - (long)cursor[0] - base); + if (ns <= 0) return; + for (int i = tid; i < ns * {OW}; i += nt) {{ + int s = i / {OW}, word = i - s * {OW}; + sw[word * {S} + s] = out[(long)index[base + s] * {NW} + word]; + }} + if (tid == 0) atomic_store_explicit(&num_events, 0, memory_order_relaxed); + threadgroup_barrier(mem_flags::mem_threadgroup); + Segments seg = {{seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep}}; + add_faults(sw, events, &num_events, args.max_events, ns, + block_stream(args.seed, bid), args.num_segments, seg, records, + overflow, tid, nt); + pack_rows(sw, ns, (threadgroup uchar*)events, dst, ((long)*cursor + base) * {NB}, tid, nt); +}} +""" +) + +SOURCE = r""" +struct Record {{ uint alias; uint words; uint extra; uint keep; u64 m0; u64 m1; }}; +struct Flip {{ uint word; uint pad; u64 mask; }}; + +inline void xor64(threadgroup u64* sw, int index, u64 mask) {{ + threadgroup atomic_uint* halves = (threadgroup atomic_uint*)(sw + index); + uint lo = uint(mask), hi = uint(mask >> 32); + if (lo) atomic_fetch_xor_explicit(halves, lo, memory_order_relaxed); + if (hi) atomic_fetch_xor_explicit(halves + 1, hi, memory_order_relaxed); +}} + +inline void xor64(device u64* sw, int index, u64 mask) {{ + device atomic_uint* halves = (device atomic_uint*)(sw + index); + uint lo = uint(mask), hi = uint(mask >> 32); + if (lo) atomic_fetch_xor_explicit(halves, lo, memory_order_relaxed); + if (hi) atomic_fetch_xor_explicit(halves + 1, hi, memory_order_relaxed); +}} + +// The fault drawn at record ``rec`` (bits 0-23 of ``draw`` the high 24 bits of +// its 40-bit fraction, bits 24-31 its slot in the channel): the slot's outcome +// if the fraction is below the keep threshold (gpu/noise.py), else its alias's. +// The fraction's low 16 bits matter only when its high bits tie with the +// threshold's, and then come from a hash of the fault. Word ``w`` of row ``s`` +// is ``sw[w * WS + s * SS]``. +template +inline void apply_fault(W sw, int row, uint rec, uint draw, u64 stream, + device const Record* records, device const Flip* overflow) +{{ + Record r = records[rec]; + uint high = draw & 0xFFFFFF, keep = r.alias & 0x1FFFFFF; + if (high > keep || (high == keep && + (uint)(mix64(stream ^ ((u64)(rec << 8 | row) * GOLDEN)) & 0xFFFF) >= r.keep)) + r = records[rec - (draw >> 24) + (r.alias >> 25)]; + int flips = r.extra & 255; + if (flips > 0) xor64(sw, (r.words & 0xFFFF) * WS + row * SS, r.m0); + if (flips > 1) xor64(sw, (r.words >> 16) * WS + row * SS, r.m1); + for (int j = 2; j < flips; ++j) {{ + Flip f = overflow[(r.extra >> 8) + j - 2]; + xor64(sw, f.word * WS + row * SS, f.mask); + }} +}} + +// The outputs of the block's ``ns`` rows ``sw[w * S + s]``, bit-packed into +// {NB} bytes per row (the bits past them cleared), contiguously at ``dst``. +inline void pack_rows(threadgroup const u64* sw, int ns, threadgroup uchar* stage, + device uchar* dst, long at, uint tid, uint nt) +{{ + for (uint i = tid; i < {NK} * {S}; i += nt) {{ + int k = i / {S}, s = i - k * {S}; + if (s >= ns) continue; + u64 v = sw[k * {S} + s]; + if (k == {NK} - 1) v &= {LAST_WORD}; + threadgroup uchar* at = stage + s * {NB} + 8 * k; + for (int j = 0; j < 8; ++j) if (8 * k + j < {NB}) at[j] = uchar(v >> (8 * j)); + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + // NOTE: the rows go to byte ``at`` of the (aligned) buffer ``dst``: bytes up + // to a 16-byte boundary, then 16 at a time, then the rest. + int total = ns * {NB}; + int head = min(total, (int)((16 - (at & 15)) & 15)); + device uchar* d = dst + at; + for (int i = tid; i < head; i += nt) d[i] = stage[i]; + int body = (total - head) / 16; + device uint4* aligned = (device uint4*)(d + head); + if ((head & 3) == 0) {{ + threadgroup const uint* from = (threadgroup const uint*)(stage + head); + for (int i = tid; i < body; i += nt) + aligned[i] = uint4(from[4 * i], from[4 * i + 1], from[4 * i + 2], from[4 * i + 3]); + }} else {{ + for (int i = tid; i < body; i += nt) {{ + threadgroup const uchar* b = stage + head + 16 * i; + uint4 v; + for (int k = 0; k < 4; ++k) + v[k] = uint(b[4 * k]) | uint(b[4 * k + 1]) << 8 | uint(b[4 * k + 2]) << 16 | uint(b[4 * k + 3]) << 24; + aligned[i] = v; + }} + }} + for (int i = head + body * 16 + tid; i < total; i += nt) d[i] = stage[i]; +}} + +// The stream of block ``bid``. NOTE: hashed, so that a segment's stream +// ``mix64(stream ^ sg * GOLDEN)`` cannot equal another block's: with ``seed ^ bid * C`` +// alone, two (block, segment) pairs with equal XORs would share their draws. +inline u64 block_stream(u64 seed, uint bid) {{ + return mix64(seed ^ ((u64)bid * 0xD1B54A32D192ED03UL)); +}} + +struct Segments {{ + device const int* first; device const int* outcomes; device const int* begin; + device const int* end; device const float* log_keep; +}}; + +// One segment's faults. With ``FINE``, for segments of fewer than 2^-12 expected +// faults, a draw within 2^-12 of u = 1 takes 24 more bits, which ``unit`` alone +// cannot resolve: it would fire such a segment at least 2^-24 of the time. +template +inline void add_segment_faults(W sw, threadgroup uint2* events, threadgroup atomic_int* num_events, + int max_events, int ns, u64 stream, u64 state, float lk, int first, int outcomes, int idx, + int end, device const Record* records, device const Flip* overflow) +{{ + while (true) {{ + state += GOLDEN; + u64 r = mix64(state); + float skip; + if (!FINE) skip = lk < -3.0e38f ? 0.f : floor(fast::log(unit(r)) / lk); + else if ((~r >> 40) >= 4096) skip = floor(fast::log(unit(r)) / lk); + else {{ + // NOTE: 1 - u = y = (j + w) 2^-24 <= 2^-12, with w from a fresh draw; + // Metal has no log1p, and the series to y^3 is exact to ~1e-11 here. + float w = float((mix64(state ^ 0x5851F42D4C957F2DUL) >> 40) + 1) * 0x1p-24f; + float y = (float(~r >> 40) + w) * 0x1p-24f; + skip = floor(-y * (1.f + y * (0.5f + y * (1.f / 3.f))) / lk); + }} + if (skip >= (float)(end - 1 - idx)) break; + idx += 1 + (int)skip; + int ch = idx / {S}, row = idx - ch * {S}; + if (row >= ns) continue; // the last block's missing shots + // NOTE: the skip reads the high 24 bits of r, the outcome its low 40: a + // slot, and the fraction of the slot below. + u64 x = (r & 0xFFFFFFFFFFUL) * (uint)outcomes; + uint slot = (uint)(x >> 40); + uint rec = (uint)(first + ch * outcomes) + slot; + uint draw = (uint)((x >> 16) & 0xFFFFFF) | (slot << 24); + int e = atomic_fetch_add_explicit(num_events, 1, memory_order_relaxed); + if (e < max_events) events[e] = uint2((uint)row | (rec << 8), draw); + else apply_fault(sw, row, rec, draw, stream, records, overflow); + }} +}} + +template +__attribute__((noinline)) void add_fine_segment_faults(W sw, threadgroup uint2* events, + threadgroup atomic_int* num_events, int max_events, int ns, u64 stream, u64 state, float lk, + int first, int outcomes, int idx, int end, device const Record* records, + device const Flip* overflow) +{{ + add_segment_faults(sw, events, num_events, max_events, ns, 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`` (laid out as in ``apply_fault``; in device memory if ``DEV``). +template +inline void add_faults(W sw, threadgroup uint2* events, threadgroup atomic_int* num_events, + int max_events, int ns, u64 stream, int num_segments, Segments seg, + device const Record* records, device const Flip* overflow, uint tid, uint nt) +{{ + for (int sg = tid; sg < num_segments; sg += nt) {{ + float lk = seg.log_keep[sg]; + int first = seg.first[sg], outcomes = seg.outcomes[sg], end = seg.end[sg]; + int idx = seg.begin[sg] - 1; + u64 state = mix64(stream ^ ((u64)sg * GOLDEN)); + if ((float)(end - 1 - idx) * -lk < 0x1p-12f) + add_fine_segment_faults(sw, events, num_events, max_events, ns, stream, + state, lk, first, outcomes, idx, end, records, + overflow); + else + add_segment_faults(sw, events, num_events, max_events, ns, stream, + state, lk, first, outcomes, idx, end, records, + overflow); + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + int ne = min(atomic_load_explicit(num_events, memory_order_relaxed), max_events); + for (int e = tid; e < ne; e += nt) {{ + uint2 v = events[e]; + apply_fault(sw, v.x & 255, v.x >> 8, v.y, stream, records, overflow); + }} + if (DEV) threadgroup_barrier(mem_flags::mem_device); + else threadgroup_barrier(mem_flags::mem_threadgroup); +}} + +struct NoiseArgs {{ long shots; u64 seed; int num_segments; int max_events; int has_survivors; }}; + +// Marks a SIMD group's surviving shots (of the block's ``ns``, from ``first``) +// in the bitmap ``survivors`` if ``has_survivors``, and appends them to +// ``queue`` after the ``count`` queued so far. +inline void enqueue(bool alive, int s, int first, long base, int ns, uint lane, + int has_survivors, device uint* queue, device atomic_uint* count, device uint* survivors) +{{ + uint ballot = uint(static_cast(simd_ballot(alive))); + if (has_survivors && lane == 0 && first < ns) survivors[(base + first) >> 5] = ballot; + if (ballot) {{ + uint leader = ctz(ballot), b = 0; + if (lane == leader) b = atomic_fetch_add_explicit(count, popcount(ballot), memory_order_relaxed); + b = simd_broadcast(b, (ushort)leader); + if (alive) queue[b + popcount(ballot & ((1u << lane) - 1))] = (uint)(base + s); + }} +}} + +{NOISE}{KEPT}""" + + +def noise_tables( + channels: ChannelSampler, flips: WordTables, num_words: int +) -> NoiseTables: + """Return the noise kernels' tables, for Metal's threadgroup memory. + + Rows too wide for ``SHARED_BYTES`` keep their words in device memory + (``shared_words`` False), where the noise kernel XORs the faults in. + """ + 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 MSL source of the ``noise`` and ``noise_kept`` kernels. + + ``noise_kept`` packs rows of ``num_outputs`` outputs. Without + ``tables.shared_words`` there is no ``noise_kept``: the rows of a block do + not fit in threadgroup memory. + """ + packed_bytes = -(-num_outputs // 8) + if tables.shared_words: + noise, kept = NOISE_SHARED, NOISE_KEPT + else: + noise, kept = NOISE_DEVICE, "" + return SOURCE.replace("{NOISE}{KEPT}", noise + kept).format( + NW=tables.num_words, + OW=output_words, + S=tables.block_shots, + NB=max(packed_bytes, 1), + NK=max(-(-packed_bytes // 8), 1), + LAST_WORD=f"0x{(1 << (num_outputs % 64 or 64)) - 1:016X}UL", + ) diff --git a/src/tsim/statevector/metal/runtime.py b/src/tsim/statevector/metal/runtime.py new file mode 100644 index 00000000..24ea971a --- /dev/null +++ b/src/tsim/statevector/metal/runtime.py @@ -0,0 +1,367 @@ +"""Metal through PyObjC for the Metal executor, loaded on first use. + +Buffers are ``MTLStorageModeShared``: on Apple silicon the CPU and the GPU share +memory, so a buffer is also a NumPy array (``Buffer.array``) and nothing is +copied between them. +""" + +from __future__ import annotations + +import functools +import os +import threading +from dataclasses import dataclass +from types import ModuleType +from typing import Any + +import numpy as np + + +@functools.cache +def metal() -> ModuleType: + """Return PyObjC's ``Metal`` module. + + Raises: + ImportError: If ``pyobjc-framework-Metal`` is not installed. + + """ + import Metal # pyright: ignore[reportMissingImports] # Apple silicon only + + return Metal + + +@functools.cache +def device() -> Any: + """Return the system's default Metal device. + + Raises: + RuntimeError: If there is none. + + """ + global _initialized + found = metal().MTLCreateSystemDefaultDevice() + if found is None: + raise RuntimeError("no Metal device") + _initialized = True + return found + + +_initialized = False +_forked_after_init = False + + +def _after_fork_in_child() -> None: + # NOTE: Metal's device and its connections to the GPU driver do not survive a + # fork; a child of a process that used them samples on the CPU. + global _forked_after_init + if _initialized: + _forked_after_init = True + + +if hasattr(os, "register_at_fork"): # not on Windows + os.register_at_fork(after_in_child=_after_fork_in_child) + + +def forked_after_init() -> bool: + """Return whether this process was forked after its parent initialized Metal.""" + return _forked_after_init + + +@functools.cache +def queue() -> Any: + """Return the executor's command queue.""" + return device().newCommandQueue() + + +LANE_PROBE_THREADS = (128, 256, 1024) +"""Threadgroup sizes of the kernels, whose SIMD lanes the device check probes.""" + +_PROBE = r""" +#include +using namespace metal; +kernel void lanes(device uint* out [[buffer(0)]], uint gid [[thread_position_in_grid]], + uint tid [[thread_position_in_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint group [[simdgroup_index_in_threadgroup]]) +{ + out[gid] = tid | (lane << 12) | (group << 20); +} +""" + + +@functools.cache +def check_device() -> None: + """Check that the device runs SIMD groups of 32 threads, in order, as the kernels assume. + + The kernels take a thread's lane and SIMD group from its index in its + threadgroup (``tid % 32`` and ``tid / 32``): a shot's threads exchange + amplitudes by lane, and the scans list shots in lane order. + + Raises: + RuntimeError: If it does not (e.g. an Intel or AMD GPU). + + """ + if not device().supportsFamily_(metal().MTLGPUFamilyApple7): + raise RuntimeError("the Metal executor needs an Apple GPU (Apple silicon)") + probe = Library(_PROBE).pipeline("lanes") + width = int(probe.threadExecutionWidth()) + if width != 32: + raise RuntimeError( + f"the Metal executor needs SIMD groups of 32 threads, not {width}" + ) + for threads in LANE_PROBE_THREADS: + out = buffer(4 * threads, np.uint32) + encoder = Encoder() + encoder.dispatch(probe, [out], 4, threads) + wait(encoder.commit()) + if not lanes_in_order(out.array, threads): + raise RuntimeError( + "the Metal executor needs SIMD lanes in thread order, which this " + f"GPU does not keep in threadgroups of {threads} threads" + ) + + +def lanes_in_order(probed: np.ndarray, threads: int) -> bool: + """Return whether the ``lanes`` probe of threadgroups of ``threads`` found lanes in order. + + ``probed[gid]`` holds the thread's index in its threadgroup (bits 0-11), its + lane (bits 12-19) and its SIMD group (bits 20-31). + """ + probed = np.asarray(probed, np.uint32) + tid, lane, group = probed & 0xFFF, (probed >> 12) & 0xFF, probed >> 20 + return bool( + np.all(tid == np.arange(len(probed)) % threads) + and np.all(lane == tid % 32) + and np.all(group == tid // 32) + ) + + +def available() -> bool: + """Return whether PyObjC's Metal is installed and there is a usable Apple GPU. + + It is not in a process forked after its parent used the GPU. + """ + return not _forked_after_init and _device_available() + + +@functools.cache +def _device_available() -> bool: + try: + check_device() + return True + except Exception: + return False + + +@dataclass +class Buffer: + """A shared Metal buffer and a NumPy view of its memory.""" + + handle: Any + array: np.ndarray + + +POOL_BUFFERS = 2 +"""Freed output buffers kept for reuse. + +The GPU's first write to each page of a fresh buffer costs the OS a fault: for +rows of 440 bytes a call runs 1.7x longer than into a buffer that a previous +call returned (and its caller dropped). +""" + +_pool: list[Any] = [] +# NOTE: reentrant, since a buffer can be freed (``_Owner.__del__``) while this +# thread searches the pool. +_pool_lock = threading.RLock() + + +class _Owner: + """Holds a Metal buffer for the NumPy arrays over its memory (their ``base``). + + A ``recycle`` buffer goes back to the pool when they are gone. + """ + + def __init__(self, handle: Any, nbytes: int, recycle: bool = False) -> None: + self.handle = handle + self.recycle = recycle + view = handle.contents().as_buffer(nbytes) + address = np.frombuffer(view, dtype=np.uint8).ctypes.data + self.__array_interface__ = { + "shape": (nbytes,), + "typestr": "|u1", + "data": (address, False), + "version": 3, + } + + def __del__(self) -> None: + if self.recycle and _pool is not None: + with _pool_lock: + _pool.append(self.handle) + del _pool[:-POOL_BUFFERS] + + +def buffer(count: int, dtype, recycle: bool = False) -> Buffer: + """Return a shared buffer of ``count`` elements of ``dtype`` (zeroed). + + Its array keeps the Metal buffer alive, so it may outlive the ``Buffer``. + With ``recycle`` the buffer may be a freed one of up to twice the size, not + zeroed, and returns to the pool once its arrays are gone. + """ + dtype = np.dtype(dtype) + nbytes = max(int(count) * dtype.itemsize, 16) + handle = None + if recycle: + with _pool_lock: + fits = [b for b in _pool if nbytes <= b.length() <= 2 * nbytes] + # NOTE: by identity, in case a buffer freed meanwhile trimmed the pool. + if fits and any(b is fits[0] for b in _pool): + handle = fits[0] + _pool[:] = [b for b in _pool if b is not handle] + if handle is None: + handle = device().newBufferWithLength_options_( + nbytes, metal().MTLResourceStorageModeShared + ) + if handle is None: + raise MemoryError( + f"no Metal buffer of {nbytes} bytes (at most " + f"{device().maxBufferLength()}); sample fewer shots per call" + ) + owner = _Owner(handle, nbytes, recycle) + array = np.asarray(owner)[: count * dtype.itemsize].view(dtype) + return Buffer(handle, array) + + +def upload(values: np.ndarray) -> Buffer: + """Return a shared buffer holding a copy of ``values`` (flattened).""" + values = np.ascontiguousarray(values) + held = buffer(values.size, values.dtype) + held.array[:] = values.reshape(-1) + return held + + +class CompileError(RuntimeError): + """A Metal kernel that does not compile, or whose pipeline cannot be built.""" + + +class Library: + """A compiled Metal library and its kernels' pipeline states.""" + + def __init__(self, source: str) -> None: + """Compile ``source`` (Metal Shading Language) with fast math. + + Raises: + CompileError: If it does not compile. + + """ + mtl = metal() + options = mtl.MTLCompileOptions.new() + options.setFastMathEnabled_(True) + library, error = device().newLibraryWithSource_options_error_( + source, options, None + ) + if library is None: + raise CompileError(f"Metal compilation failed: {error}") + self._library = library + self._pipelines: dict[str, Any] = {} + + def pipeline(self, name: str) -> Any: + """Return the compute pipeline state of kernel ``name``. + + Raises: + CompileError: If its pipeline cannot be built. + + """ + if name not in self._pipelines: + function = self._library.newFunctionWithName_(name) + if function is None: + raise KeyError(name) + descriptor = metal().MTLComputePipelineDescriptor.new() + descriptor.setComputeFunction_(function) + descriptor.setLabel_(name) + ( + state, + error, + ) = device().newComputePipelineStateWithDescriptor_options_reflection_error_( + descriptor, 0, None, None + ) + if state is None: + raise CompileError(f"pipeline {name}: {error}") + self._pipelines[name] = state + return self._pipelines[name] + + +Argument = Buffer | tuple[Buffer, int] | np.ndarray | np.generic + +PROFILE: dict[str, float] | None = None +"""If a dict, every dispatch runs as its own command buffer and adds its GPU time +(seconds) under its kernel's name: for profiling only.""" + + +class Encoder: + """Records kernel dispatches into one command buffer.""" + + def __init__(self) -> None: + """Start a command buffer with a compute encoder.""" + self.commands = queue().commandBuffer() + self._encoder = self.commands.computeCommandEncoder() + + def dispatch( + self, + pipeline: Any, + arguments: list[Argument], + threadgroups: int, + threads: int, + threadgroup_bytes: int = 0, + ) -> None: + """Dispatch ``threadgroups`` groups of ``threads`` threads of ``pipeline``. + + Each argument binds buffer index ``i``: a buffer (or a buffer and a byte + offset), or a NumPy scalar or array passed by value. A nonzero + ``threadgroup_bytes`` binds threadgroup memory 0. + """ + mtl, encoder = metal(), self._encoder + encoder.setComputePipelineState_(pipeline) + for i, argument in enumerate(arguments): + if isinstance(argument, Buffer): + encoder.setBuffer_offset_atIndex_(argument.handle, 0, i) + elif isinstance(argument, tuple): + encoder.setBuffer_offset_atIndex_(argument[0].handle, argument[1], i) + else: + data = np.ascontiguousarray(argument).tobytes() + encoder.setBytes_length_atIndex_(data, len(data), i) + if threadgroup_bytes: + encoder.setThreadgroupMemoryLength_atIndex_( + -(-threadgroup_bytes // 16) * 16, 0 + ) + if threadgroups: + encoder.dispatchThreadgroups_threadsPerThreadgroup_( + mtl.MTLSizeMake(threadgroups, 1, 1), mtl.MTLSizeMake(threads, 1, 1) + ) + if PROFILE is not None: + name = str(pipeline.label() or "") or str(pipeline) + commands = self.commit() + PROFILE[name] = PROFILE.get(name, 0.0) + wait(commands) + self.commands = queue().commandBuffer() + self._encoder = self.commands.computeCommandEncoder() + + def commit(self) -> Any: + """End encoding and commit; return the command buffer.""" + self._encoder.endEncoding() + self.commands.commit() + return self.commands + + +def wait(commands: Any) -> float: + """Wait for a committed command buffer; return its GPU time in seconds. + + Raises: + MemoryError: If the GPU ran out of memory. + RuntimeError: If it failed otherwise. + + """ + commands.waitUntilCompleted() + error = commands.error() + if error is not None: + if error.code() == metal().MTLCommandBufferErrorOutOfMemory: + raise MemoryError(f"the GPU ran out of memory: {error}") + raise RuntimeError(f"Metal command buffer failed: {error}") + return float(commands.GPUEndTime() - commands.GPUStartTime()) diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py new file mode 100644 index 00000000..4b58dd89 --- /dev/null +++ b/src/tsim/statevector/metal/sampler.py @@ -0,0 +1,834 @@ +"""``MetalShotSampler``: ``ShotSampler`` with the noise and the tapes on a Metal GPU. + +A batch of shots is one command buffer of these kernels, which run one after +another: + + noise draws every shot's noise into its words and screens it; the screened + survivors go to a queue + tape runs every component on the queued shots, a SIMD group of ``2^tb`` + threads per shot with its amplitudes in registers; with a + postselection mask it clears the rejected shots from the screen's + bitmap of survivors + list lists the accepted shots in order from the bitmap + pack gathers the accepted rows, bit-packed (or draws the deferred noise + on them first, ``noise_kept``) + +Every buffer is shared with the CPU (unified memory), so nothing is copied: +the packing kernels write each batch's rows into the array that the call +returns, after the rows of the batches before it (a cursor in the GPU's memory +counts them). The host never waits between batches: it commits each batch's +command buffer as it encodes it, the batches reuse one set of buffers (so +Metal runs them one after another), and the GPU resets their counters itself. + +The kernels are compiled once per postselection mask (Metal caches compiled +shaders across processes). Amplitudes are complex64, so a sample follows the +CPU executor's distribution to within float32 rounding of the outcome +probabilities; the noise and the outcome draws come from other streams than +the CPU's, so a seed's samples differ from the CPU's, and they depend on the +batch size. +""" + +from __future__ import annotations + +import threading +import warnings +from dataclasses import dataclass, field +from typing import Any + +import numpy as np + +from tsim.noise.channels import ChannelSampler +from tsim.statevector.gpu import FORKED, tape +from tsim.statevector.gpu.noise import NoiseGroups, screen_pass_rate, split +from tsim.statevector.gpu.rows import select as _select +from tsim.statevector.metal import kernels, runtime +from tsim.statevector.metal.codegen import METAL +from tsim.statevector.metal.kernels import SCAN_THREADS +from tsim.statevector.metal.noise import noise_tables +from tsim.statevector.shots import ShotPlan, ShotSampler + +BATCH_SHOTS = 1 << 20 +"""Shots of one batch (one command buffer).""" + +BATCH_BYTES = 256 << 20 +"""Bytes of shot words of one batch, which bound its shots below ``BATCH_SHOTS``.""" + +NOISE_THREADS = 256 +"""Threads per threadgroup of the noise kernels.""" + +COPY_THREADS = 256 +"""Threads per threadgroup of the kernels that copy rows.""" + +MAX_COPY_THREADS = 1 << 20 +"""Most threads of a copy kernel's dispatch; each thread loops over its share.""" + +# NOTE: the words of ``_Slot.counts``. +_SURVIVORS, _ACCEPTED, _DEVIATION, _CURSOR, _CAPACITY = range(5) + +_STACKLEVEL = 6 +"""Stack level of the samplers' warnings: the caller of ``sample``.""" + +TAPE_THREADS = 128 +"""Threads per threadgroup of the tape kernel.""" + +WIDE_TAPE = 8 +"""Width from which a shot's threads hold at most 16 amplitudes each, not 32.""" + +MAX_CALL_SHOTS = 1 << 31 +"""Most shots of one GPU call; larger calls run in parts (the counters are 32-bit).""" + + +@dataclass +class _Launch: + """The device tables of one noise launch: its segments and records.""" + + segments: list[runtime.Buffer] + records: runtime.Buffer + num_segments: int + max_events: int + + +@dataclass +class _Module: + """A plan's compiled kernels and its device constants.""" + + library: runtime.Library + tb: int + has_mask: bool + screen: runtime.Buffer + plan: ShotPlan + psm_words: np.ndarray + passing: float + """Estimated share of shots that pass the screen.""" + splits: dict[bool, tuple[_Launch, _Launch] | None] = field(default_factory=dict) + + +@dataclass +class _Slot: + """The buffers that a call's batches reuse, one batch after another. + + ``counts`` holds the batch's screened survivors, its accepted shots and the + call's largest normalization deviation (as float bits), then the call's + cursor: the rows that its batches have packed, and the rows that its + output has room for. + """ + + words: runtime.Buffer + queue: runtime.Buffer + accepted: runtime.Buffer + counts: runtime.Buffer + totals: runtime.Buffer + index: runtime.Buffer + values: runtime.Buffer | None = None + + +@dataclass +class _Call: + """What a call's batches share: its kernels, buffers, noise passes and output.""" + + module: _Module + slot: _Slot + constant: runtime.Buffer + split: tuple[_Launch, _Launch] | None + expectations: bool + rows: runtime.Buffer + """The accepted rows, bit-packed.""" + values: runtime.Buffer | None + """The accepted rows' expectations, with ``expectations``.""" + + +class MetalShotSampler: + """Samples shot words of a statevector program on a Metal GPU. + + It keeps ``ShotSampler``'s interface (``sample``, ``direct_bools``, + ``reference``, ``plan``, ``layout``); ``sample`` returns only the shots that + pass postselection. ``sample_accepted`` can also return each observable's + expectation. + """ + + supports = staticmethod(tape.supports) + + def __init__( + self, + host: ShotSampler, + channels: ChannelSampler, + num_detectors: int, + seed: np.random.SeedSequence, + ) -> None: + """Upload the noise tables of ``host``'s shot words. + + ``host`` draws the reference samples, on the CPU. + + Raises: + ImportError: If PyObjC's Metal is not installed. + RuntimeError: If there is no Apple GPU. + + """ + runtime.check_device() + assert self.supports(host.layout) + self._args = (host, channels, num_detectors, seed) + # NOTE: a call's batches share the sampler's buffers, so calls from several + # threads run one at a time. + self._lock = threading.RLock() + self._host = host + layout = host.layout + self._num_detectors = num_detectors + self._observables = np.arange(num_detectors, layout.num_outputs) + self._rng = np.random.default_rng(seed) + self._channels = channels + self._noise = noise_tables(channels, layout.noise, layout.num_words) + self._overflow = runtime.upload(self._noise.overflow.view(np.uint8)) + self._all = self._launch(self._noise.groups) + self._modules: dict[bytes, _Module] = {} + self._constants: dict[bytes, runtime.Buffer] = {} + self._slots: dict[tuple[int, bool], _Slot] = {} + self._shares: dict[tuple, float] = {} + self._cpu_only = False + + @property + def device(self) -> str: + """Return where this sampler samples (``SamplerStats.device``).""" + return "cpu" if self._cpu_only else "metal" + + def __reduce__(self): + """Pickle the CPU sampler and the inputs; the GPU state is rebuilt on load.""" + return _restore, (*self._args, self._rng.bit_generator.state) + + @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 _module(self, mask: np.ndarray | None) -> _Module: + key = b"" if mask is None else np.asarray(mask, bool).tobytes() + if key in self._modules: + return self._modules[key] + plan = self.plan(mask) + layout = self.layout + tb = _threads_per_shot(layout) + parts = tape.tape_kernel(layout, plan, mask, self._num_detectors, tb, METAL) + source = kernels.source( + layout, self._noise, parts, tb, max(1, len(self._observables)) + ) + library = runtime.Library(source) + for name in kernels.names(self._noise.shared_words): + library.pipeline(name) + compiled = _Module( + library=library, + tb=tb, + has_mask=mask is not None, + screen=runtime.upload(plan.screen), + plan=plan, + psm_words=parts.psm_words, + passing=screen_pass_rate(self._channels, layout, plan), + ) + self._modules[key] = compiled + return compiled + + def _launch(self, groups: NoiseGroups) -> _Launch: + """Upload a noise launch's segments and records.""" + return _Launch( + segments=[runtime.upload(a) for a in groups.segments], + records=runtime.upload(groups.records.view(np.uint8)), + num_segments=len(groups.segments[0]), + max_events=groups.max_events, + ) + + def _split( + self, module: _Module, expectations: bool + ) -> tuple[_Launch, _Launch] | None: + """Return a postselected batch's noise passes (screened, deferred), or None. + + As ``gpu.noise.split``; with ``expectations`` the observables' bits are + read too. + """ + # NOTE: rows too wide for threadgroup memory draw all their noise at once. + if not module.has_mask or not self._noise.shared_words: + return None + if expectations not in module.splits: + found = split( + self._noise, + self._channels, + self.layout, + module.plan, + module.psm_words, + self._num_detectors if expectations else None, + ) + module.splits[expectations] = ( + None + if found is None + else (self._launch(found[0]), self._launch(found[1])) + ) + return module.splits[expectations] + + # -- execution -------------------------------------------------------------- + + def _constant(self, flip: np.ndarray | None) -> runtime.Buffer: + """Return the outputs' constants XOR the reference bits ``flip``.""" + key = b"" if flip is None else np.asarray(flip, bool).tobytes() + if key not in self._constants: + self._constants[key] = runtime.upload(self._host._constant(flip)) + return self._constants[key] + + def _batch_shots(self, shots: int, batch_size: int | None) -> int: + """Return the shots of a call's batches.""" + # NOTE: at most 2^16 words a row, so a batch stays below 512 MB, and the + # kernels' 32-bit offsets within it do not wrap. + cap = max(1 << 10, BATCH_BYTES // (8 * self.layout.num_words)) + return max(1, min(batch_size or BATCH_SHOTS, BATCH_SHOTS, cap, shots)) + + def _slot(self, shots: int, expectations: bool) -> _Slot: + """Return the buffers for batches of up to ``shots`` shots.""" + key = (shots, expectations) + slot = self._slots.get(key) + if slot is None: + nw, no = self.layout.num_words, max(1, len(self._observables)) + slot = _Slot( + words=runtime.buffer(shots * nw, np.uint64), + queue=runtime.buffer(shots, np.uint32), + accepted=runtime.buffer(-(-shots // 32), np.uint32), + counts=runtime.buffer(8, np.uint32), + totals=runtime.buffer(-(-shots // 32 // SCAN_THREADS) + 1, np.uint32), + index=runtime.buffer(shots, np.uint32), + ) + if expectations: + slot.values = runtime.buffer(shots * no, np.float32) + self._slots = {k: v for k, v in self._slots.items() if k[1] != expectations} + self._slots[key] = slot + return slot + + def _encode_batch( + self, encoder: runtime.Encoder, call: _Call, shots: int, first_shot: int + ) -> None: + """Encode a batch's noise and tape kernels. + + Only the shots that pass the screen get their words written. + """ + module, slot, constant = call.module, call.slot, call.constant + expectations = call.expectations + launch = call.split[0] if call.split else self._all + library = module.library + words = slot.words + encoder.dispatch(library.pipeline("reset"), [slot.counts], 1, 1) + seed = np.uint64(self._rng.integers(0, 2**64, dtype=np.uint64)) + noise_seed = np.uint64(self._rng.integers(0, 2**64, dtype=np.uint64)) + S = self._noise.block_shots + args = np.zeros(1, dtype=_NOISE_ARGS) + args[0] = ( + shots, + noise_seed, + launch.num_segments, + launch.max_events, + int(module.has_mask), + ) + encoder.dispatch( + library.pipeline("noise"), + [ + words, + args, + *launch.segments, + launch.records, + self._overflow, + constant, + module.screen, + slot.queue, + _count(slot, _SURVIVORS), + slot.accepted, + ], + -(-shots // S), + NOISE_THREADS, + threadgroup_bytes=8 * max(launch.max_events, 1), + ) + if self.layout.components: + tape_args = np.zeros(1, dtype=_TAPE_ARGS) + flags = int(module.has_mask) | (2 if expectations else 0) + tape_args[0] = (first_shot, seed, flags, 0) + values = slot.values if slot.values is not None else slot.counts + threads = shots << module.tb + encoder.dispatch( + library.pipeline("tape"), + [ + words, + tape_args, + slot.queue, + _count(slot, _SURVIVORS), + slot.accepted, + _count(slot, _ACCEPTED), + values, + _count(slot, _DEVIATION), + ], + -(-threads // TAPE_THREADS), + TAPE_THREADS, + ) + + def _encode_pack( + self, + encoder: runtime.Encoder, + call: _Call, + shots: int, + fixed: tuple[runtime.Buffer, int], + ) -> None: + """Encode the listing and packing of a batch's accepted rows into ``call.rows``. + + They follow the rows of the call's earlier batches (at the cursor), and + their expectations go to ``call.values``; ``fixed`` holds the batch's + shots, which a call without a mask returns. + """ + module, slot, split = call.module, call.slot, call.split + expectations, rows, values = call.expectations, call.rows, call.values + library = module.library + S = self._noise.block_shots + nb = -(-self.layout.num_outputs // 8) + no = len(self._observables) + if module.has_mask: + words = -(-shots // 32) + blocks = -(-words // SCAN_THREADS) + encoder.dispatch( + library.pipeline("count_bits"), + [slot.accepted, np.int64(words), slot.totals], + blocks, + SCAN_THREADS, + ) + encoder.dispatch( + library.pipeline("list_bits"), + [slot.accepted, np.int64(words), slot.totals, slot.index], + blocks, + SCAN_THREADS, + ) + kept = _ACCEPTED if self.layout.components else _SURVIVORS + count = _count(slot, kept) + else: + count = fixed + cursor = _count(slot, _CURSOR) + blocks = -(-shots // S) + if nb and split is not None: + launch = split[1] + kept_args = np.zeros(1, dtype=_KEPT_ARGS) + kept_args[0] = ( + np.uint64(self._rng.integers(0, 2**64, dtype=np.uint64)), + launch.num_segments, + launch.max_events, + ) + encoder.dispatch( + library.pipeline("noise_kept"), + [ + slot.words, + kept_args, + *launch.segments, + launch.records, + self._overflow, + count, + slot.index, + rows, + cursor, + ], + blocks, + NOISE_THREADS, + threadgroup_bytes=max(8 * launch.max_events, S * nb), + ) + elif nb: + threads = min(-(-shots * nb // 4) + 1, MAX_COPY_THREADS) + encoder.dispatch( + library.pipeline("gather_rows"), + [ + slot.words, + count, + slot.index, + np.int32(module.has_mask), + rows, + cursor, + ], + -(-threads // COPY_THREADS), + COPY_THREADS, + ) + if expectations and no and self.layout.components: + assert slot.values is not None and values is not None + encoder.dispatch( + library.pipeline("gather_float"), + [ + slot.values, + count, + np.int32(no), + slot.index, + values, + cursor, + np.int32(module.has_mask), + ], + COPY_THREADS, + COPY_THREADS, + ) + encoder.dispatch(library.pipeline("advance"), [cursor, count], 1, 1) + + def sample( + self, + shots: int, + columns: list[np.ndarray], + *, + bit_packed: bool, + batch_size: int | None = None, + mask: np.ndarray | None = None, + flip: np.ndarray | None = None, + ) -> tuple[list[np.ndarray], float]: + """Return the output bits ``columns`` of ``shots`` shots and the largest norm deviation. + + As ``ShotSampler.sample``: with a ``mask`` only the shots that pass + postselection are returned. ``batch_size`` bounds the shots of a batch. + Out of GPU memory, when a kernel does not compile, and in a process + forked after its parent used the GPU, the call warns and samples on the + CPU. + """ + + def on_gpu(): + rows, _, deviation = self.sample_accepted( + shots, + mask=mask, + flip=flip, + batch_size=batch_size, + bools=not bit_packed, + ) + if not bit_packed: + return [_columns(rows, c) for c in columns], deviation + num_outputs = self.layout.num_outputs + selected = [_select(rows, c, num_outputs, bit_packed) for c in columns] + return selected, deviation + + return self._or_cpu( + on_gpu, + lambda: self._host.sample( + shots, + columns, + bit_packed=bit_packed, + batch_size=batch_size, + mask=mask, + flip=flip, + ), + ) + + def _or_cpu(self, on_gpu, on_cpu): + """Return ``on_gpu()``, or warn and return ``on_cpu()`` where the GPU fails. + + It fails in a process forked after its parent used the GPU, out of GPU + memory, and when a kernel does not compile; after that the sampler + samples on the CPU. + """ + if self._cpu_only or _forked(): + return on_cpu() + try: + return on_gpu() + except MemoryError as error: + _warn_out_of_memory(error) + return on_cpu() + except runtime.CompileError as error: + self._cpu_only = True + warnings.warn( + f"{error}; this sampler samples on the CPU from now on.", + stacklevel=_STACKLEVEL, + ) + return on_cpu() + + def _capacity(self, key: tuple, module: _Module, shots: int) -> int: + """Return the rows to make room for: the expected accepted shots and a margin. + + The share comes from the previous call of ``key``, or else the screen's + estimate (which the tapes' postselection only lowers). + """ + share = self._shares.get(key, module.passing if module.has_mask else 1.0) + expected = shots * min(1.0, share) + return min(shots, int(1.05 * expected + 8 * np.sqrt(expected) + 4096)) + + def _run_accepted(self, call: _Call, shots: int, batch: int, capacity: int) -> None: + """Run a call's batches into ``call.rows`` (and ``values``) of ``capacity`` rows.""" + slot = call.slot + starts = range(0, shots, batch) + fixed = runtime.buffer(len(starts), np.uint32) + slot.counts.array[:] = 0 + slot.counts.array[_CAPACITY] = capacity + commands = [] + try: + for k, start in enumerate(starts): + n = min(batch, shots - start) + fixed.array[k] = n + encoder = runtime.Encoder() + self._encode_batch(encoder, call, n, start) + self._encode_pack(encoder, call, n, (fixed, 4 * k)) + commands.append(encoder.commit()) + except BaseException: + self._settle(commands) + raise + self._wait_all(commands) + + def _wait_all(self, commands: list[Any]) -> None: + """Wait for a call's command buffers; on any error, let them all finish first. + + Raises: + MemoryError: If one ran out of GPU memory. + RuntimeError: If one failed otherwise. + + """ + try: + for command in commands: + runtime.wait(command) + except BaseException: + self._settle(commands) + raise + + def _settle(self, commands: list[Any]) -> None: + """Drop the batch buffers and wait for ``commands``, through further interrupts. + + A call that fails or is interrupted leaves its batches on the GPU; they + must finish before the next call reuses (or the pool hands out) a buffer + that they write. + """ + self._slots.clear() + for command in commands: + while True: + try: + command.waitUntilCompleted() + break + except KeyboardInterrupt: + continue + + def sample_accepted( + self, + shots: int, + *, + mask: np.ndarray | None = None, + flip: np.ndarray | None = None, + expectations: bool = False, + batch_size: int | None = None, + bools: bool = False, + ) -> tuple[np.ndarray, np.ndarray | None, float]: + """Sample ``shots`` shots; return only those that pass postselection. + + A shot passes when none of the ``mask``'s detectors is set (every shot + passes without a mask). Rows keep their shots' order and hold every output; + with ``bools``, one boolean per output, unpacked on the GPU. + + Returns: + ``(outputs, expectations, deviation)``: the passing shots' outputs, + bit-packed as ``np.packbits(..., bitorder="little")`` would (``uint8``, + in shared memory), with ``expectations`` each observable's + expectation of ``(-1)^observable`` given the shot's history before + the measurement that settles it, or its sampled ``±1`` when a + postselected output of its component is tested at or after that + measurement (else None), and the largest normalization deviation + beyond float32 rounding. + + Raises: + MemoryError: If the GPU runs out of memory. + + """ + if shots > MAX_CALL_SHOTS: + parts = [ + self.sample_accepted( + min(MAX_CALL_SHOTS, shots - start), + mask=mask, + flip=flip, + expectations=expectations, + batch_size=batch_size, + bools=bools, + ) + for start in range(0, shots, MAX_CALL_SHOTS) + ] + found = [p[1] for p in parts if p[1] is not None] + values = np.concatenate(found) if expectations else None + return ( + np.concatenate([p[0] for p in parts]), + values, + max(p[2] for p in parts), + ) + with self._lock: + return self._sample_call( + shots, + mask=mask, + flip=flip, + expectations=expectations, + batch_size=batch_size, + bools=bools, + ) + + def _sample_call( + self, + shots: int, + *, + mask: np.ndarray | None = None, + flip: np.ndarray | None = None, + expectations: bool = False, + batch_size: int | None = None, + bools: bool = False, + ) -> tuple[np.ndarray, np.ndarray | None, float]: + """One call of ``sample_accepted``, of at most ``MAX_CALL_SHOTS`` shots.""" + module = self._module(mask) + nb, no = -(-self.layout.num_outputs // 8), len(self._observables) + if not shots: + values = np.zeros((0, no), np.float32) if expectations else None + if bools: + return np.zeros((0, self.layout.num_outputs), bool), values, 0.0 + return np.zeros((0, nb), np.uint8), values, 0.0 + constant = self._constant(flip) + split = self._split(module, expectations) + batch = self._batch_shots(shots, batch_size) + slot = self._slot(batch, expectations) + key = (module.plan, expectations) + state = self._rng.bit_generator.state + capacity = self._capacity(key, module, shots) + while True: + # NOTE: a command buffer makes all of its buffers resident, so the + # output has room for the expected rows only; a call that overflows + # it runs again, on the same streams, with room for every shot. + rows = runtime.buffer(capacity * nb, np.uint8, recycle=True) + values = ( + runtime.buffer(capacity * no, np.float32, recycle=True) + if expectations + else None + ) + call = _Call(module, slot, constant, split, expectations, rows, values) + self._run_accepted(call, shots, batch, capacity) + total = int(slot.counts.array[_CURSOR]) + if total <= capacity: + break + self._rng.bit_generator.state = state + capacity = shots + self._shares[key] = total / max(shots, 1) + # NOTE: fast math records a NaN total as NaN (not inf), which must warn. + deviation = tape.deviation(slot.counts.array[_DEVIATION]) + packed = rows.array[: total * nb].reshape(total, nb) + expected = None + if values is not None: + expected = values.array[: total * no].reshape(total, no) + if not self.layout.components: + expected[:] = 0 + self._direct_expectations(packed, expected) + if bools: + return self._unpack(module, rows, total), expected, deviation + return packed, expected, deviation + + def _unpack(self, module: _Module, rows: runtime.Buffer, total: int) -> np.ndarray: + """Return the first ``total`` packed ``rows`` as booleans, unpacked on the GPU.""" + num_outputs = self.layout.num_outputs + nb = -(-num_outputs // 8) + groups = -(-num_outputs // 4) + if total * groups >= 1 << 32: + packed = rows.array[: total * nb].reshape(total, nb) + return _select(packed, np.arange(num_outputs), num_outputs, False) + out = runtime.buffer(total * num_outputs, np.uint8, recycle=True) + if total and num_outputs: + threads = min(total * groups, MAX_COPY_THREADS) + encoder = runtime.Encoder() + encoder.dispatch( + module.library.pipeline("unpack_rows"), + [rows, out, np.uint32(total)], + -(-threads // COPY_THREADS), + COPY_THREADS, + ) + self._wait_all([encoder.commit()]) + return out.array[: total * num_outputs].reshape(total, num_outputs).view(bool) + + def _direct_expectations(self, rows: np.ndarray, values: np.ndarray) -> 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(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, + batch_size: int | None = None, + ) -> np.ndarray: + """Return every output bit of a program without components, as booleans.""" + return self._or_cpu( + lambda: self.sample_accepted( + shots, flip=flip, batch_size=batch_size, bools=True + )[0], + lambda: self._host.direct_bools(shots, flip, batch_size), + ) + + +def _forked() -> bool: + """Return whether this process was forked after Metal's setup, with a warning.""" + if not runtime.forked_after_init(): + return False + warnings.warn(FORKED, stacklevel=_STACKLEVEL) + return True + + +def _threads_per_shot(layout) -> int: + """Return log2 of a shot's threads, as ``tape.threads_per_shot`` from ``WIDE_TAPE`` on.""" + # NOTE: fewer registers per thread pay for the shuffles once tapes are wide: on + # an M4 Pro, 32 threads of 16 amplitudes run 9 qubits 1.4x and 8 qubits 1.1-1.2x + # faster than 16 threads of 32; at 7 qubits they tie, at 6 four threads of 16 + # are 10% slower than two of 32. + width = max((c.width for c in layout.components), default=0) + if width >= WIDE_TAPE: + return tape.threads_per_shot(layout, tape.AMPLITUDE_AXES - 1) + return tape.threads_per_shot(layout) + + +def _warn_out_of_memory(error: MemoryError) -> None: + warnings.warn( + f"{error}; sampling on the CPU. Fewer shots per call fit.", + stacklevel=_STACKLEVEL, + ) + + +def _count(slot: _Slot, word: int) -> tuple[runtime.Buffer, int]: + """Return word ``word`` of ``slot.counts``, as a kernel argument.""" + return slot.counts, 4 * word + + +def _restore( + host, channels, num_detectors, seed, state +) -> MetalShotSampler | ShotSampler: + """Rebuild a pickled ``MetalShotSampler``, or its CPU sampler where Metal is missing.""" + if not runtime.available(): + return host + shots = MetalShotSampler(host, channels, num_detectors, seed) + shots._rng.bit_generator.state = state + return shots + + +def _columns(rows: np.ndarray, columns: np.ndarray) -> np.ndarray: + """Return the outputs ``columns`` of boolean rows of every output, C-contiguous.""" + count = len(columns) + if count == rows.shape[1] and np.array_equal(columns, np.arange(count)): + return rows + return np.ascontiguousarray(rows[:, columns]) + + +_NOISE_ARGS = np.dtype( + [ + ("shots", " np.ndarray: + def direct_bools( + self, + shots: int, + flip: np.ndarray | None = None, + batch_size: int | None = None, + ) -> np.ndarray: """Return every output bit of a program without components, as booleans. The noise sampler writes them straight into the rows. The samples equal - those of ``sample``. ``flip`` is as in ``sample``. + those of ``sample``. ``flip`` is as in ``sample``; ``batch_size`` does not + change the samples on the CPU. """ assert self._direct_tables is not None out = np.empty((shots, self._layout.num_outputs), dtype=np.uint8) diff --git a/test/conftest.py b/test/conftest.py new file mode 100644 index 00000000..d12e57ba --- /dev/null +++ b/test/conftest.py @@ -0,0 +1,10 @@ +import pytest + +from tsim.statevector.metal import runtime as metal_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, "available", lambda: False) diff --git a/test/integration/test_metal_sampler.py b/test/integration/test_metal_sampler.py new file mode 100644 index 00000000..2c8730ce --- /dev/null +++ b/test/integration/test_metal_sampler.py @@ -0,0 +1,975 @@ +"""The Metal executor end to end, against the stabilizer-rank strategy and the CPU executor.""" + +import itertools +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 + +import numpy as np +import pytest + +from tsim.circuit import Circuit +from tsim.statevector.gpu import noise as gpu_noise, tape +from tsim.statevector.metal import ( + codegen, + kernels, + noise as metal_noise, + runtime, + sampler as metal_sampler, +) +from tsim.statevector.metal.sampler import MetalShotSampler +from tsim.statevector.program import StatevectorComponent +from tsim.statevector.shots import ShotSampler + +pytestmark = pytest.mark.skipif( + not runtime.available(), reason="needs PyObjC's Metal and an Apple 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) -> MetalShotSampler: + """Return the GPU shot sampler of ``sampler``.""" + assert isinstance(sampler._shots, MetalShotSampler) + return sampler._shots + + +def _cpu(compile, **kwargs): + """Return ``compile(**kwargs)``, a sampler compiled as on a host without a GPU.""" + with pytest.MonkeyPatch.context() as patch: + patch.setattr(runtime, "available", lambda: False) + sampler = compile(**kwargs) + assert sampler._shots is None or type(sampler._shots) is ShotSampler + return sampler + + +def _accepted(sampler, shots, mask, **kwargs): + detectors, observables = sampler.sample( + shots, separate_observables=True, postselection_mask=mask, **kwargs + ) + # NOTE: sample returns only the shots that the mask keeps. + assert mask is None or not (detectors & mask).any() + return np.concatenate([detectors, observables], axis=1) + + +def _packed(sampler, shots, mask=None): + """Return ``sample``'s bit-packed detectors then observables.""" + return sampler.sample( + shots, append_observables=True, bit_packed=True, postselection_mask=mask + ) + + +@pytest.mark.parametrize("name", sorted(CIRCUITS)) +@pytest.mark.parametrize( + ("kind", "looping"), + [("detector", False), ("measurement", False), ("detector", True)], +) +def test_metal_sampling_matches_stabilizer_rank(monkeypatch, name, kind, looping): + if looping: + # NOTE: measurements, projections and pruning in looping tapes. + monkeypatch.setattr(tape, "UNROLL_BUDGET", 0) + circuit = Circuit(CIRCUITS[name]) + shots = 40_000 + if kind == "detector": + expected = circuit.compile_detector_sampler(strategy="cat5", seed=1).sample( + shots, batch_size=10_000, append_observables=True + ) + got = _gpu(circuit.compile_detector_sampler(seed=2)).sample( + shots, append_observables=True + ) + else: + expected = circuit.compile_sampler(strategy="cat5", seed=1).sample( + shots, batch_size=10_000 + ) + got = _gpu(circuit.compile_sampler(seed=2)).sample(shots) + assert_counts_match(got, expected) + + +@pytest.mark.parametrize("name", sorted(CIRCUITS)) +@pytest.mark.parametrize("device_memory", [False, True]) +def test_metal_postselection_keeps_the_accepted_distribution( + monkeypatch, name, device_memory +): + circuit = Circuit(CIRCUITS[name]) + if device_memory: + # NOTE: too little threadgroup memory for the rows: the noise kernel + # keeps them in device memory. + monkeypatch.setattr(metal_noise, "SHARED_BYTES", 2048 + 8) + elif circuit.num_detectors == 0: + pytest.skip("no detectors to postselect") + shots = 40_000 + mask = np.ones(circuit.num_detectors, dtype=bool) if circuit.num_detectors else None + expected = _accepted(_cpu(circuit.compile_detector_sampler, seed=1), shots, mask) + sampler = _gpu(circuit.compile_detector_sampler(seed=2)) + assert _gpu_shots(sampler)._noise.shared_words != device_memory + got = _accepted(sampler, shots, mask) + rate, gpu_rate = len(expected) / shots, len(got) / shots + sigma = np.sqrt(rate * (1 - rate) * 2 / shots) + assert abs(gpu_rate - rate) <= 5 * sigma + 3 / shots + if len(got) and len(expected): + assert_counts_match(got, expected) + + +@pytest.mark.parametrize("seed", [1, 2, 3]) +def test_metal_sampling_of_random_noisy_circuits(seed): + stim_circuit = gen_stim_circuit(3, 60, p_t=3, p_x_error=0.5, p_r_z=0.5, seed=seed) + circuit = Circuit.from_stim_program(stim_circuit) + shots = 8_192 + expected = circuit.compile_sampler(strategy="cat5", seed=seed).sample( + shots, batch_size=shots + ) + got = _gpu(circuit.compile_sampler(seed=seed + 10)).sample(shots) + assert_counts_match(got, expected) + + +@pytest.mark.parametrize("reference", [False, True]) +def test_metal_reference_samples_and_pruning(reference): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + mask = np.ones(circuit.num_detectors, dtype=bool) + shots = 40_000 + kwargs = {"use_detector_reference_sample": reference} + expected = _accepted( + _cpu(circuit.compile_detector_sampler, seed=3), shots, mask, **kwargs + ) + got = _accepted( + _gpu(circuit.compile_detector_sampler(seed=3)), shots, mask, **kwargs + ) + rate, gpu_rate = len(expected) / shots, len(got) / shots + assert abs(gpu_rate - rate) <= 5 * np.sqrt(rate * (1 - rate) * 2 / shots) + assert_counts_match(got, expected) + + +def test_metal_postselection_keeps_every_column(): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + packed = _packed(_gpu(circuit.compile_detector_sampler(seed=4)), 20_000, mask) + outputs = circuit.num_detectors + circuit.num_observables + assert packed.shape[1] == -(-outputs // 8) + bits = np.unpackbits(packed, axis=1, bitorder="little") + assert not bits[:, outputs:].any() + assert not bits[:, : circuit.num_detectors][:, mask].any() + full = _accepted(_cpu(circuit.compile_detector_sampler, seed=5), 20_000, mask) + assert abs(len(bits) - len(full)) <= 5 * np.sqrt(len(full)) + assert_counts_match(bits[:, :outputs].astype(bool), full) + + +@pytest.mark.parametrize("bit_packed", [False, True]) +def test_metal_arranges_the_columns_as_the_cpu_does(bit_packed): + # NOTE: samplers of one seed draw the same shots, whichever columns they return. + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[0] = True + num_detectors = circuit.num_detectors + + def sample(**kwargs): + sampler = _gpu(circuit.compile_detector_sampler(seed=17)) + return sampler.sample( + 5_000, postselection_mask=mask, bit_packed=bit_packed, **kwargs + ) + + def bools(array, count): + if not bit_packed: + return array + return np.unpackbits(array, axis=1, count=count, bitorder="little").astype(bool) + + appended = bools(sample(append_observables=True), num_detectors + 1) + detectors = bools(sample(), num_detectors) + prepended = bools(sample(prepend_observables=True), num_detectors + 1) + separate = sample(separate_observables=True) + np.testing.assert_array_equal(detectors, appended[:, :num_detectors]) + np.testing.assert_array_equal(prepended[:, 0], appended[:, -1]) + np.testing.assert_array_equal(prepended[:, 1:], detectors) + np.testing.assert_array_equal(bools(separate[0], num_detectors), detectors) + np.testing.assert_array_equal(bools(separate[1], 1), appended[:, -1:]) + cpu = _cpu(circuit.compile_detector_sampler, seed=17).sample( + 10, postselection_mask=mask, bit_packed=bit_packed, append_observables=True + ) + gpu = sample(append_observables=True) + assert gpu.dtype == cpu.dtype + assert gpu.shape[1] == cpu.shape[1] + + +def test_direct_booleans_take_the_batch_size_as_packed_rows_do(): + # NOTE: each batch draws its own stream, so a component-free sampler's + # booleans and packed rows match only if both honor batch_size. + circuit = Circuit("X_ERROR(0.3) 0 1 2\nM 0 1 2\nDETECTOR rec[-1]\nDETECTOR rec[-2]") + + def sample(bit_packed): + sampler = _gpu(circuit.compile_detector_sampler(seed=5)) + return sampler.sample(5_000, batch_size=1_000, bit_packed=bit_packed) + + packed = np.packbits(sample(False), axis=1, bitorder="little") + np.testing.assert_array_equal(packed, sample(True)) + + +def test_metal_samples_no_shots(): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + sampler = _gpu(circuit.compile_detector_sampler(seed=1)) + assert sampler.sample(0, append_observables=True).shape == (0, 4) + assert sampler.sample(0, bit_packed=True).shape == (0, 1) + + +# NOTE: 150 noisy detectors, so the rows span three words, and two observables. +NOISY_DETECTORS = "\n".join( + [ + "R " + " ".join(map(str, range(150))), + "X_ERROR(0.3) " + " ".join(map(str, range(150))), + "M " + " ".join(map(str, range(150))), + *(f"DETECTOR rec[{i - 150}]" for i in range(150)), + "OBSERVABLE_INCLUDE(0) rec[-1] rec[-7]", + "OBSERVABLE_INCLUDE(1) rec[-150]", + ] +) + + +def test_metal_postselection_drops_the_rows_that_fail(monkeypatch): + # NOTE: the masked run must draw the unmasked run's noise, so none is deferred. + monkeypatch.setattr(gpu_noise, "DEFER_SHARE", 2.0) + circuit = Circuit(NOISY_DETECTORS) + rng = np.random.default_rng(0) + mask = np.zeros(150, dtype=bool) + # NOTE: a few masked detectors, then unmasked runs across words, then a mix. + mask[[0, 3, 4]] = True + mask[100:] = rng.random(50) < 0.1 + full = _packed(_gpu(circuit.compile_detector_sampler(seed=11)), 5_000) + rows = _packed(_gpu(circuit.compile_detector_sampler(seed=11)), 5_000, mask) + full_bits = np.unpackbits(full, axis=1, count=152, bitorder="little") + passing = ~full_bits[:, :150][:, mask].any(axis=1) + assert 0 < passing.sum() < 5_000 + np.testing.assert_array_equal(rows, full[passing]) + + +# NOTE: GHZ_WITH_DIRECT_DETECTOR's component and direct detector, then four +# noisy direct detectors that no mask below reads; one fault flips detectors 2 +# and 3 together, and the observable also reads the last noisy qubit. +UNREAD_NOISE = """ +R 0 1 2 3 4 5 6 7 +H 0 +CNOT 0 1 +CNOT 1 2 +T 0 1 2 +H 0 1 2 +DEPOLARIZE1(0.05) 0 1 2 +T 0 1 +CZ 0 1 +H 0 1 +M 0 1 2 +DETECTOR rec[-3] rec[-2] +DETECTOR rec[-2] rec[-1] +OBSERVABLE_INCLUDE(0) rec[-1] +X_ERROR(0.2) 3 +E(0.1) X3 X4 +X_ERROR(0.3) 4 5 6 7 +M 3 4 5 6 7 +DETECTOR rec[-5] +DETECTOR rec[-4] +DETECTOR rec[-3] +DETECTOR rec[-2] +DETECTOR rec[-1] +OBSERVABLE_INCLUDE(0) rec[-1] +""" + + +def test_metal_defers_the_noise_that_postselection_cannot_see(): + circuit = Circuit(UNREAD_NOISE) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + sampler = _gpu(circuit.compile_detector_sampler(seed=12)) + got = _accepted(sampler, 200_000, mask) + assert _gpu_shots(sampler)._module(mask).splits[False] is not None + full = _accepted(_cpu(circuit.compile_detector_sampler, seed=13), 200_000, mask) + assert abs(len(got) - len(full)) <= 5 * np.sqrt(len(full)) + assert_counts_match(got, full) + + +def _expectations(sampler, shots, mask=None): + """Return the passing rows and expectations of ``sample_accepted``.""" + rows, values, _ = _gpu_shots(sampler).sample_accepted( + shots, mask=mask, expectations=True + ) + assert values is not None + return rows, values + + +def test_deferred_noise_leaves_the_expectations_unbiased(): + # NOTE: with expectations, the noise on the observable's direct qubit is not + # deferred: the tape reads the observable's bit for its sign. + circuit = Circuit(UNREAD_NOISE) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + sampler = _gpu(circuit.compile_detector_sampler(seed=14)) + _, values = _expectations(sampler, 200_000, mask) + assert _gpu_shots(sampler)._module(mask).splits[True] is not None + full = _accepted(_cpu(circuit.compile_detector_sampler, seed=15), 200_000, mask) + expected = np.mean(1 - 2 * full[:, -1].astype(float)) + assert abs(values.mean() - expected) < 5 / np.sqrt(len(full)) + + +@pytest.mark.parametrize("text", [NOISY_DETECTORS, UNREAD_NOISE]) +def test_metal_batches_return_the_rows_of_one_pass(monkeypatch, text): + # NOTE: 1000-shot batches, queued on one set of buffers; without deferral, a + # masked and an unmasked run on the same batches draw the same shots. + monkeypatch.setattr(gpu_noise, "DEFER_SHARE", 2.0) + circuit = Circuit(text) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + outputs = circuit.num_detectors + circuit.num_observables + + def packed(mask=None): + sampler = _gpu(circuit.compile_detector_sampler(seed=16)) + return sampler.sample( + 10_500, + append_observables=True, + bit_packed=True, + postselection_mask=mask, + batch_size=1000, + ) + + full, rows = packed(), packed(mask) + full_bits = np.unpackbits(full, axis=1, count=outputs, bitorder="little") + passing = ~full_bits[:, : circuit.num_detectors][:, mask].any(axis=1) + assert 0 < passing.sum() < 10_500 + np.testing.assert_array_equal(rows, full[passing]) + + +def test_observable_expectations_average_to_the_sampled_observable(): + circuit = Circuit(T_STATE_X) + sampler = _gpu(circuit.compile_detector_sampler(seed=6)) + packed, values = _expectations(sampler, 200_000) + assert values.shape == (len(packed), 1) + # NOTE: given the shot's Z error, the expectation is ±cos(π/4). + np.testing.assert_allclose(np.abs(values), np.sqrt(0.5), atol=1e-5) + observed = 1 - 2 * np.unpackbits(packed, axis=1, bitorder="little")[:, 0].astype( + float + ) + sampled = _cpu(circuit.compile_detector_sampler, seed=7).sample( + 200_000, append_observables=True + )[:, 0] + expected = np.mean(1 - 2 * sampled.astype(float)) + assert abs(values.mean() - expected) < 5 * 1 / np.sqrt(200_000) + # NOTE: the returned observable bit is sampled from the same state. + assert abs(observed.mean() - values.mean()) < 5 / np.sqrt(200_000) + + +# NOTE: the T state's X outcome is copied onto qubit 1, whose postselected (and +# noisy) detector is settled after the observable's measurement. +T_STATE_X_THEN_DETECTOR = """ +R 0 1 +H 0 +T 0 +H 0 +CNOT 0 1 +X_ERROR(0.1) 1 +M 0 +OBSERVABLE_INCLUDE(0) rec[-1] +M 1 +DETECTOR rec[-1] +""" + +# NOTE: a postselected T-state detector settled before the observable's measurement. +T_STATE_X_AFTER_DETECTOR = """ +R 0 1 +H 0 1 +T 0 1 +H 1 +M 1 +DETECTOR rec[-1] +Z_ERROR(0.1) 0 +H 0 +M 0 +OBSERVABLE_INCLUDE(0) rec[-1] +""" + + +def _postselected_expectations(text, shots): + circuit = Circuit(text) + mask = np.ones(circuit.num_detectors, dtype=bool) + sampler = _gpu(circuit.compile_detector_sampler(seed=9)) + _, values = _expectations(sampler, shots, mask) + detectors, observables = _cpu(circuit.compile_detector_sampler, seed=10).sample( + shots, separate_observables=True + ) + kept = observables[~(detectors & mask).any(axis=1), 0] + return values[:, 0], np.mean(1 - 2 * kept.astype(float)), len(kept) + + +def test_expectations_stay_unbiased_under_later_postselection(): + values, expected, kept = _postselected_expectations( + T_STATE_X_THEN_DETECTOR, 200_000 + ) + # NOTE: the expectation before the measurement would average to cos(π/4). + assert abs(values.mean() - expected) < 5 * np.sqrt(1 / len(values) + 1 / kept) + np.testing.assert_array_equal(np.abs(values), 1) + + +def test_expectations_are_exact_after_earlier_postselection(): + values, expected, kept = _postselected_expectations( + T_STATE_X_AFTER_DETECTOR, 200_000 + ) + np.testing.assert_allclose(np.abs(values), np.sqrt(0.5), atol=1e-5) + assert abs(values.mean() - expected) < 5 * np.sqrt(1 / len(values) + 1 / kept) + + +def test_clifford_observables_have_expectations_of_plus_or_minus_one(): + circuit = Circuit( + "R 0 1\nX_ERROR(0.3) 0\nCNOT 0 1\nM 0 1\nDETECTOR rec[-1]\nOBSERVABLE_INCLUDE(0) rec[-2]" + ) + sampler = _gpu(circuit.compile_detector_sampler(seed=8)) + packed, values = _expectations(sampler, 10_000) + bits = np.unpackbits(packed, axis=1, bitorder="little") + np.testing.assert_array_equal(values[:, 0], 1 - 2 * bits[:, 1].astype(np.float32)) + + +@pytest.mark.parametrize("reference", [False, True]) +def test_metal_samples_programs_without_components(reference): + circuit = Circuit( + "R 0 1\nX_ERROR(0.3) 0\nCNOT 0 1\nM 0 1\nDETECTOR rec[-1]\nOBSERVABLE_INCLUDE(0) rec[-2]" + ) + kwargs = {"use_detector_reference_sample": reference} + expected = _cpu(circuit.compile_detector_sampler, seed=1).sample( + 40_000, append_observables=True, **kwargs + ) + got = _gpu(circuit.compile_detector_sampler(seed=2)).sample( + 40_000, append_observables=True, **kwargs + ) + assert_counts_match(got, expected) + measured = _gpu(circuit.compile_sampler(seed=3)).sample(40_000) + assert measured.shape == (40_000, 2) + np.testing.assert_array_equal(measured[:, 0], measured[:, 1]) + + +def test_stabilizer_rank_strategies_stay_on_jax(): + circuit = Circuit(T_STATE_X) + assert circuit.compile_detector_sampler(strategy="cat5")._shots is None + + +def _widest(circuit): + """Return the active width of the circuit's widest statevector component.""" + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) + components = circuit.compile_sampler(seed=1)._program.components + return max(c.width for c in components if isinstance(c, StatevectorComponent)) + + +def _random_unitary_circuit(width, layers, seed=0, noise=True): + """Layers of random U3 pairs and CX gates, then every qubit measured. + + With ``noise``, X errors come first, so that every rotation reads the shot's + f-bits and none of them runs before the shots split (``common_prefix``). + """ + rng = np.random.default_rng(seed) + lines = ["X_ERROR(0.2) " + " ".join(map(str, range(width)))] if noise else [] + for _ in range(layers): + perm = rng.permutation(width) + for a, b in zip(perm[::2], perm[1::2], strict=False): + for q in (a, b): + t = rng.uniform(-1, 1, 3) + lines.append(f"U3({t[0]},{t[1]},{t[2]}) {q}") + lines.append(f"CX {a} {b}") + lines.append("M " + " ".join(map(str, range(width)))) + return Circuit("\n".join(lines)) + + +def _assert_matches_cpu(circuit, width, shots=200_000): + expected = _cpu(circuit.compile_sampler, seed=1).sample(shots) + got = _gpu(circuit.compile_sampler(seed=2)).sample(shots) + np.testing.assert_allclose( + got.mean(axis=0), expected.mean(axis=0), atol=5 * 0.5 / np.sqrt(shots / 2) + ) + # NOTE: pairs of bits, which a wrong shuffle partner would scramble. + for a, b in [(0, 1), (0, width - 1), (width // 2, width - 1)]: + assert_counts_match(got[:, [a, b]], expected[:, [a, b]]) + + +@pytest.mark.parametrize(("width", "layers"), [(6, 8), (8, 8), (9, 10), (11, 12)]) +def test_metal_wide_components_match_the_cpu_executor(width, layers): + # NOTE: widths 6, 8, 9 and 11 run 2, 16, 32 and 32 threads per shot. + circuit = _random_unitary_circuit(width, layers) + assert _widest(circuit) == width + _assert_matches_cpu(circuit, width) + + +@pytest.mark.parametrize(("width", "layers"), [(6, 6), (8, 8), (11, 6)]) +def test_metal_looping_tapes_match_the_cpu_executor(monkeypatch, width, layers): + # NOTE: a tape past the unroll budget loops, applying runs of rotations together. + monkeypatch.setattr(tape, "UNROLL_BUDGET", 0) + _assert_matches_cpu(_random_unitary_circuit(width, layers, seed=width), width) + + +@pytest.mark.parametrize("width", [6, 9, 11]) +def test_metal_shots_start_after_the_common_prefix(width): + # NOTE: without noise the rotations before the first measurement run once, + # on the host; the shots load their amplitudes and run the rest. + circuit = _random_unitary_circuit(width, width + 2, seed=width, noise=False) + _assert_matches_cpu(circuit, width) + + +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"): + sampler = circuit.compile_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + + +@pytest.mark.parametrize(("probability", "shots"), [(2e-8, 1 << 30), (1e-6, 1 << 26)]) +def test_metal_draws_small_outcome_probabilities_at_their_rate(probability, shots): + # NOTE: a 24-bit draw takes any positive probability at least 2^-24 = 6e-8 of + # the time; float32 rounding gives certain outcomes such probabilities. 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)) + ones = sum(int(sampler.sample(1 << 24).sum()) for _ in range(shots >> 24)) + expected = probability * shots + assert abs(ones - expected) <= 5 * np.sqrt(expected) + 1 + + +def test_metal_reruns_a_call_whose_rows_overflow_the_output(monkeypatch): + # NOTE: room for 10 rows: the call runs again with room for every shot, on + # the same streams, so it returns the rows of a call with enough room. + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + + def packed(): + sampler = _gpu(circuit.compile_detector_sampler(seed=21)) + return sampler.sample( + 5_000, + append_observables=True, + bit_packed=True, + postselection_mask=mask, + batch_size=1000, + ) + + expected = packed() + monkeypatch.setattr(MetalShotSampler, "_capacity", lambda self, *args: 10) + rows = packed() + assert len(rows) > 10 + np.testing.assert_array_equal(rows, expected) + + +def test_metal_unmasked_expectations_after_a_masked_call(): + # NOTE: a masked call lists its accepted shots in the batch buffers; an + # unmasked call on the same buffers must still read every row's own value. + circuit = Circuit(T_STATE_X_AFTER_DETECTOR) + mask = np.ones(circuit.num_detectors, dtype=bool) + sampler = _gpu(circuit.compile_detector_sampler(seed=5)) + _expectations(sampler, 20_000, mask) + shots = _gpu_shots(sampler) + state = shots._rng.bit_generator.state + rows, values = _expectations(sampler, 20_000) + shots._rng.bit_generator.state = state + shots._slots.clear() + again, values_again = _expectations(sampler, 20_000) + np.testing.assert_array_equal(again, rows) + np.testing.assert_array_equal(values_again, values) + + +def test_metal_samples_rows_too_wide_for_threadgroup_memory(): + # NOTE: 3200 detectors are 51 words: the noise goes into device memory. + circuit = Circuit( + "\n".join( + [ + "R " + " ".join(map(str, range(3200))), + "X_ERROR(0.05) " + " ".join(map(str, range(3200))), + "M " + " ".join(map(str, range(3200))), + *(f"DETECTOR rec[{i - 3200}]" for i in range(3200)), + "OBSERVABLE_INCLUDE(0) rec[-1] rec[-2]", + ] + ) + ) + sampler = _gpu(circuit.compile_detector_sampler(seed=31)) + assert not _gpu_shots(sampler)._noise.shared_words + got = sampler.sample(20_000, append_observables=True) + np.testing.assert_allclose(got[:, :3200].mean(), 0.05, atol=5e-4) + np.testing.assert_allclose(got[:, -1].mean(), 2 * 0.05 * 0.95, atol=0.01) + np.testing.assert_array_equal(got[:, -1], got[:, -2] ^ got[:, -3]) + + +def test_metal_unrolled_width_11_tapes_match_the_cpu_executor(): + # NOTE: 64 amplitudes per thread, unrolled, in two register arrays. + circuit = _random_unitary_circuit(11, 7, seed=0) + assert _widest(circuit) == 11 + shots = _gpu_shots(_gpu(circuit.compile_sampler(seed=1))) + assert not tape.loops(shots.plan(None).programs[0], 5, 0) + _assert_matches_cpu(circuit, 11) + + +def test_metal_returned_rows_are_not_reused_while_alive(): + circuit = Circuit(NOISY_DETECTORS) + sampler = _gpu(circuit.compile_detector_sampler(seed=3)) + first = _packed(sampler, 200_000) + kept = first.copy() + second = _packed(sampler, 200_000) + np.testing.assert_array_equal(first, kept) + assert not np.shares_memory(first, second) + del first + third = _packed(sampler, 200_000) + assert not np.shares_memory(second, third) + assert not (third == second).all() + + +def test_metal_leaves_programs_without_shot_words_on_the_cpu(): + circuit = Circuit("R 0\nH 0\nT 0\nH 0\nM 0") + sampler = circuit.compile_detector_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + assert sampler.sample(10).shape == (10, 0) + + +def _sampler_in_child(_): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + sampler = Circuit(T_STATE_X).compile_sampler(seed=1) + sampler.sample(10) + return type(sampler._shots).__name__, [str(w.message) for w in caught] + + +def test_processes_forked_after_the_gpu_was_used_sample_on_the_cpu(): + # NOTE: Metal's device does not survive a fork of a process that used it. + _gpu(Circuit(T_STATE_X).compile_sampler(seed=1)).sample(10) + with multiprocessing.get_context("fork").Pool(1) as pool: + name, messages = pool.map(_sampler_in_child, [0])[0] + assert name == "ShotSampler" + assert any("forked" in m for m in messages) + + +def test_a_kernel_that_does_not_compile_samples_on_the_cpu(monkeypatch): + circuit = Circuit(T_STATE_X) + sampler = _gpu(circuit.compile_sampler(seed=1)) + compiled = [] + + def broken(*args, **kwargs): + compiled.append(1) + return "kernel void tape(" + + monkeypatch.setattr(kernels, "source", broken) + with pytest.warns(UserWarning, match="Metal compilation failed"): + got = sampler.sample(20_000) + assert sampler.stats.device == "cpu" + with warnings.catch_warnings(): + warnings.simplefilter("error") + got = np.concatenate([got, sampler.sample(20_000)]) + assert len(compiled) == 1 + expected = _cpu(circuit.compile_sampler, seed=2).sample(40_000) + assert_counts_match(got, expected) + + +def test_noise_of_tiny_probabilities_fires_at_its_rate(): + # NOTE: segments of fewer than 2^-12 expected faults take the kernels' 48-bit + # path. 24-bit skips alone would fire every segment at least 2^-24 of the + # time, ~110 more faults than the ~83 expected here; a path that never fired + # would give none. + n, p, calls = 1000, 3e-10, 44 + probabilities = [p * (1 + i / n) for i in range(n)] + text = "\n".join(f"X_ERROR({q!r}) {i}" for i, q in enumerate(probabilities)) + sampler = _gpu( + Circuit(f"{text}\nM " + " ".join(map(str, range(n)))).compile_sampler(seed=1) + ) + # NOTE: faults are rare enough that each lands in a byte of its own. + faults = sum( + np.count_nonzero(sampler.sample(1 << 22, bit_packed=True)) for _ in range(calls) + ) + expected = sum(probabilities) * calls * (1 << 22) + assert abs(faults - expected) <= 5 * np.sqrt(expected) + + +def test_rare_outcome_of_a_channel_is_drawn_at_its_rate(): + # NOTE: 1000 channels (distinct, so Tsim keeps them apart) fire half the time + # and then flip detector 0 with chance 6e-9: with 24-bit keep thresholds it + # rounded to zero; ~100 flips are expected here. + k, q, calls = 10, 3e-9, 34 + lines = ["R " + " ".join(map(str, range(k + 1)))] + subsets = [ + s for r in range(1, k + 1) for s in itertools.combinations(range(1, k + 1), r) + ][:1000] + for subset in subsets: + lines.append("E(0.5) " + " ".join(f"X{i}" for i in subset)) + lines.append(f"ELSE_CORRELATED_ERROR({q / 0.5}) X0") + lines.append("M " + " ".join(map(str, range(k + 1)))) + lines += [f"DETECTOR rec[{i - k - 1}]" for i in range(k + 1)] + sampler = _gpu(Circuit("\n".join(lines)).compile_detector_sampler(seed=1)) + flips = sum( + int(np.count_nonzero(sampler.sample(1_000_000, bit_packed=True)[:, 0] & 1)) + for _ in range(calls) + ) + expected = len(subsets) * q * calls * 1_000_000 + assert abs(flips - expected) <= 5 * np.sqrt(expected) + + +def test_rotation_with_a_branch_float32_cannot_resolve_samples_on_the_cpu(): + circuit = Circuit("R 0\nH 0\nR_Z(1e-5) 0\nH 0\nM 0\nDETECTOR rec[-1]") + with pytest.warns(UserWarning, match="rotation has a branch of weight 2.5e-10"): + sampler = circuit.compile_detector_sampler(seed=1) + assert sampler.stats.device == "cpu" + + +@pytest.mark.parametrize("p0", [3e-8, 1e-6]) +def test_rare_outcome_zero_is_drawn_at_its_rate(p0): + # NOTE: w1 / total rounds to a multiple of 2^-24 below 1 in float32; the draw + # must use w0, as the CPU's ``_cut`` does. + source = codegen.PRELUDE + r""" +kernel void zeros(constant float& w0 [[buffer(0)]], constant ulong& n [[buffer(1)]], + device atomic_uint* count [[buffer(2)]], uint gid [[thread_position_in_grid]], + uint grid [[threads_per_grid]]) +{ + uint local = 0; + for (u64 i = gid; i < n; i += grid) + local += !draw_one(mix64(i * GOLDEN + 12345), w0, 1.f, w0 + 1.f); + if (local) atomic_fetch_add_explicit(count, local, memory_order_relaxed); +} +""" + pipeline = runtime.Library(source).pipeline("zeros") + count = runtime.buffer(1, np.uint32) + n = 1 << 30 + encoder = runtime.Encoder() + encoder.dispatch(pipeline, [np.float32(p0), np.uint64(n), count], 1024, 256) + runtime.wait(encoder.commit()) + expected = p0 * n + assert abs(int(count.array[0]) - expected) <= 5 * np.sqrt(expected) + 1 + + +@pytest.mark.parametrize( + "text", + # NOTE: a quantum component (``sample``) and direct outputs only (``direct_bools``). + [GHZ_WITH_DIRECT_DETECTOR, "R 0\nX_ERROR(0.1) 0\nM 0\nDETECTOR rec[-1]"], +) +def test_calls_out_of_gpu_memory_fall_back_to_the_cpu(monkeypatch, text): + circuit = Circuit(text) + sampler = _gpu(circuit.compile_detector_sampler(seed=1)) + + def out_of_memory(*args, **kwargs): + raise MemoryError("the GPU ran out of memory") + + monkeypatch.setattr(_gpu_shots(sampler), "sample_accepted", out_of_memory) + with pytest.warns(UserWarning, match="out of memory; sampling on the CPU"): + got = sampler.sample(20_000, append_observables=True) + expected = _cpu(circuit.compile_detector_sampler, seed=2).sample( + 20_000, append_observables=True + ) + assert_counts_match(got, expected) + + +def test_samplers_out_of_gpu_memory_at_setup_run_on_the_cpu(monkeypatch): + def out_of_memory(*args, **kwargs): + raise MemoryError("no Metal buffer") + + monkeypatch.setattr(MetalShotSampler, "__init__", out_of_memory) + with pytest.warns(UserWarning, match="out of memory; sampling on the CPU"): + sampler = Circuit(T_STATE_X).compile_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + + +def test_tsim_device_cpu_keeps_the_samplers_on_the_cpu(monkeypatch): + monkeypatch.setenv("TSIM_DEVICE", "cpu") + with warnings.catch_warnings(): + warnings.simplefilter("error") + sampler = Circuit(T_STATE_X).compile_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + monkeypatch.setenv("TSIM_DEVICE", "") + _gpu(Circuit(T_STATE_X).compile_sampler(seed=1)) + + +@pytest.mark.parametrize("text", [GHZ_WITH_DIRECT_DETECTOR, NOISY_DETECTORS]) +def test_metal_unpacks_booleans_as_it_packs_rows(text): + # NOTE: samplers of one seed draw the same shots; 4 and 152 outputs, neither a + # multiple of 8. + circuit = Circuit(text) + outputs = circuit.num_detectors + circuit.num_observables + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[0] = True + for postselection_mask in (None, mask): + + def sample(bit_packed, postselection_mask=postselection_mask): + sampler = _gpu(circuit.compile_detector_sampler(seed=23)) + return sampler.sample( + 30_000, + append_observables=True, + bit_packed=bit_packed, + postselection_mask=postselection_mask, + ) + + packed, bools = sample(True), sample(False) + assert bools.dtype == np.bool_ and bools.flags.c_contiguous + unpacked = np.unpackbits(packed, axis=1, count=outputs, bitorder="little") + np.testing.assert_array_equal(bools, unpacked.astype(bool)) + detectors = _gpu(circuit.compile_detector_sampler(seed=23)).sample( + 30_000, postselection_mask=postselection_mask + ) + assert detectors.flags.c_contiguous + np.testing.assert_array_equal(detectors, bools[:, : circuit.num_detectors]) + + +def test_threads_sharing_a_gpu_sampler_get_their_own_rows(): + # NOTE: a call's batches share the sampler's buffers; four threads at once + # must not see each other's rows. + circuit = Circuit(NOISY_DETECTORS) + sampler = _gpu(circuit.compile_detector_sampler(seed=24)) + mask = np.zeros(150, dtype=bool) + mask[:3] = True + failures = [] + + def work(masked): + for _ in range(6): + rows = sampler.sample( + 20_000, + batch_size=1000, + append_observables=True, + postselection_mask=mask if masked else None, + ) + if not masked and len(rows) != 20_000: + failures.append(f"{len(rows)} rows") + if masked and rows[:, :3].any(): + failures.append("a masked detector is set") + if (rows[:, 150] != rows[:, 149] ^ rows[:, 143]).any(): + failures.append("an observable disagrees with its detectors") + + threads = [threading.Thread(target=work, args=(i % 2,)) for i in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert not failures, failures[:5] + + +def test_an_interrupted_call_leaves_the_next_call_alone(monkeypatch): + # NOTE: a slow width-11 tape and kernels compiled beforehand, so the + # interrupted call's batches are still on the GPU when the next call starts. + width = 11 + lines = [str(_random_unitary_circuit(width, 30, seed=0))] + lines += [f"DETECTOR rec[{i - width}]" for i in range(width)] + circuit = Circuit("\n".join(lines)) + sampler = _gpu(circuit.compile_detector_sampler(seed=25)) + shots = _gpu_shots(sampler) + all_detectors = np.ones(width, dtype=bool) + mask = np.zeros(width, dtype=bool) + mask[1] = True + for warm in (mask, all_detectors): + sampler.sample(1 << 17, batch_size=1 << 16, postselection_mask=warm) + encode_pack = MetalShotSampler._encode_pack + calls = [] + + def interrupt_late(self, *args, **kwargs): + calls.append(None) + if len(calls) == 20: + raise KeyboardInterrupt + return encode_pack(self, *args, **kwargs) + + monkeypatch.setattr(MetalShotSampler, "_encode_pack", interrupt_late) + with pytest.raises(KeyboardInterrupt): + sampler.sample(1 << 23, batch_size=1 << 16, postselection_mask=all_detectors) + monkeypatch.setattr(MetalShotSampler, "_encode_pack", encode_pack) + state = shots._rng.bit_generator.state + rows = sampler.sample(1 << 18, batch_size=1 << 16, postselection_mask=mask) + shots._rng.bit_generator.state = state + expected = sampler.sample(1 << 18, batch_size=1 << 16, postselection_mask=mask) + np.testing.assert_array_equal(rows, expected) + + +SAMPLER_BEFORE_FORK = None + + +def _sample_in_child(_): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + assert SAMPLER_BEFORE_FORK is not None + rows = SAMPLER_BEFORE_FORK.sample(100) + return rows.shape, [str(w.message) for w in caught] + + +def test_samplers_compiled_before_a_fork_sample_on_the_cpu_in_the_child(): + global SAMPLER_BEFORE_FORK + SAMPLER_BEFORE_FORK = _gpu(Circuit(T_STATE_X).compile_detector_sampler(seed=1)) + with multiprocessing.get_context("fork").Pool(1) as pool: + shape, messages = pool.map_async(_sample_in_child, [0]).get(timeout=60)[0] + assert shape == (100, 0) + assert any("forked" in m for m in messages) + + +def test_gpu_samplers_pickle_and_continue_their_streams(): + circuit = Circuit(NOISY_DETECTORS) + sampler = _gpu(circuit.compile_detector_sampler(seed=26)) + sampler.sample(1000) + loaded = pickle.loads(pickle.dumps(sampler)) # noqa: S301 + assert isinstance(loaded._shots, MetalShotSampler) + np.testing.assert_array_equal(loaded.sample(5000), sampler.sample(5000)) + + +def test_noise_beyond_the_gpu_tables_samples_on_the_cpu(): + # NOTE: a chain of 8 correlated errors has 255 outcomes, more than the 127 of + # a GPU noise record. + chain = ["E(0.01) X0", *(f"ELSE_CORRELATED_ERROR(0.01) X{k}" for k in range(1, 8))] + text = "\n".join( + ["R 0 1 2 3 4 5 6 7 8", "H 8", "T 8", "H 8", *chain, "M 0 1 2 3 4 5 6 7 8"] + + [f"DETECTOR rec[{k - 9}]" for k in range(9)] + ) + with pytest.warns(UserWarning, match="cannot sample this circuit"): + sampler = Circuit(text).compile_detector_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + assert sampler.sample(10).shape == (10, 9) + + +def test_calls_larger_than_the_counters_run_in_parts(monkeypatch): + # NOTE: the 32-bit counters hold MAX_CALL_SHOTS; a larger call is the calls of + # its parts. + monkeypatch.setattr(metal_sampler, "MAX_CALL_SHOTS", 1000) + mask = np.zeros(150, dtype=bool) + mask[:3] = True + circuit = Circuit(NOISY_DETECTORS) + whole = _gpu(circuit.compile_detector_sampler(seed=27)).sample( + 2500, postselection_mask=mask, bit_packed=True + ) + parts = _gpu(circuit.compile_detector_sampler(seed=27)) + expected = [ + parts.sample(n, postselection_mask=mask, bit_packed=True) + for n in (1000, 1000, 500) + ] + np.testing.assert_array_equal(whole, np.concatenate(expected)) + + +def test_the_repr_says_where_a_sampler_samples(): + circuit = Circuit(T_STATE_X) + assert "samples on metal" in repr(_gpu(circuit.compile_detector_sampler(seed=1))) + cpu = _cpu(circuit.compile_detector_sampler, seed=1) + assert "samples on cpu" in repr(cpu) and cpu.stats.device == "cpu" diff --git a/test/unit/statevector/test_metal_runtime.py b/test/unit/statevector/test_metal_runtime.py new file mode 100644 index 00000000..d9ea11e0 --- /dev/null +++ b/test/unit/statevector/test_metal_runtime.py @@ -0,0 +1,31 @@ +import numpy as np +import pytest + +from tsim.statevector.metal.runtime import lanes_in_order + + +def _probed(tid, lane, group): + return (tid | (lane << 12) | (group << 20)).astype(np.uint32) + + +@pytest.mark.parametrize("threads", [128, 1024]) +def test_lanes_in_thread_order_pass_the_probe(threads): + tid = np.tile(np.arange(threads), 3) + assert lanes_in_order(_probed(tid, tid % 32, tid // 32), threads) + + +@pytest.mark.parametrize( + "lane_of", + [ + lambda tid: (tid * 5) % 32, # lanes permuted within each SIMD group + lambda tid: (tid % 32) ^ ((tid // 32) & 1), # swapped in every other group + ], +) +def test_lanes_out_of_order_fail_the_probe(lane_of): + tid = np.tile(np.arange(256), 2) + assert not lanes_in_order(_probed(tid, lane_of(tid), tid // 32), 256) + + +def test_simd_groups_out_of_order_fail_the_probe(): + tid = np.arange(256) + assert not lanes_in_order(_probed(tid, tid % 32, 7 - tid // 32), 256) diff --git a/uv.lock b/uv.lock index a741f3e7..ff76fd32 100644 --- a/uv.lock +++ b/uv.lock @@ -130,6 +130,7 @@ dependencies = [ { name = "numba" }, { name = "numpy" }, { name = "psutil" }, + { name = "pyobjc-framework-metal", marker = "platform_machine == 'arm64' and sys_platform == 'darwin'" }, { name = "stim" }, { name = "stim-timeline-text" }, { name = "zxparam" }, @@ -191,6 +192,7 @@ requires-dist = [ { name = "numba", specifier = ">=0.61" }, { name = "numpy", specifier = ">=2.0" }, { name = "psutil", specifier = ">=5.9.0" }, + { name = "pyobjc-framework-metal", marker = "platform_machine == 'arm64' and sys_platform == 'darwin'", specifier = ">=10.0" }, { name = "stim", specifier = ">=1.15.0" }, { name = "stim-timeline-text", specifier = ">=0.1.0" }, { name = "zxparam", specifier = ">=0.1.0" }, @@ -2788,6 +2790,61 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/17/0d/eb2df17fe0bf089c43445ba759e89150a0d36fb7fcbc44e750d8d31f0dbe/pymatching-2.4.0-cp314-cp314-win_amd64.whl", hash = "sha256:5c5ce4ed245d3670ea7d62b09f8a29b105533b032a3d21fdf6df04c9b6441ff1", size = 356304, upload-time = "2026-05-22T18:09:54.3Z" }, ] +[[package]] +name = "pyobjc-core" +version = "13.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/83/58/7bab095330b8893e1ac58f3cf9710462c3242b600f8fa5db0d5dda795775/pyobjc_core-13.0.tar.gz", hash = "sha256:ffe154462d14aa4852f424e9116f17b5f8d44f78fc0f6c4727e721a5fa7177da", size = 1083172, upload-time = "2026-10-30T10:45:39.835Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/06/ad/7151cd0f01d561a6d9e761b12cdcbb5bcf4f4e65efa619bf4afcb202df28/pyobjc_core-13.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:4627348cb31c8deb2bb94702d9bcaa74869f4fef9559fd134416a6a4ca42006e", size = 5615698, upload-time = "2026-10-30T10:32:42.051Z" }, + { url = "https://files.pythonhosted.org/packages/68/ac/783ecf1c5045792fb51ad18e0131ff64e33ed9dee9921766c14e03eb27db/pyobjc_core-13.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:5818e3d8cc2677a1e767968ebfe3bc0cb7f80c1f7c19c653ee25ef1cc1b3e188", size = 6716501, upload-time = "2026-10-30T10:32:46.2Z" }, + { url = "https://files.pythonhosted.org/packages/d8/42/7b31511491d966b64183dfab325f022098e06e90c19a9cc8540de4f32588/pyobjc_core-13.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:f29f7cdb3969268716017ada9516b94f6dd1c66c2e415274b5ac1169440da1a4", size = 6734122, upload-time = "2026-10-30T10:32:50.192Z" }, + { url = "https://files.pythonhosted.org/packages/04/54/fda17a291ecdb5186f545c9a7575d02f55815e239756567f030b555b14f3/pyobjc_core-13.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:1e8752b3950ceec603bf36f57083151abed15e0c2e7c5b8c3efded9c81dc732c", size = 7031547, upload-time = "2026-10-30T10:32:55.183Z" }, + { url = "https://files.pythonhosted.org/packages/41/a5/463dc32a6c32cff57b4678298a13cd830d7ab09cfec3d7e80f20a2f414fd/pyobjc_core-13.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:fe470bd7b8163b4b8d216b1ebe5b227f01c0fa1c4e0085d7ccade269dec02bb0", size = 6729572, upload-time = "2026-10-30T10:32:59.498Z" }, + { url = "https://files.pythonhosted.org/packages/b4/11/e57bb09924e8e0dfbaa87a9cf8da13e224189f760aa0c39ab0ce9defb0d4/pyobjc_core-13.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:c497a23077ba0a42b59b879d81049eebfb0b7b67a8ca80c68b6d3683bbbd6d6f", size = 7021797, upload-time = "2026-10-30T10:33:03.702Z" }, + { url = "https://files.pythonhosted.org/packages/43/3c/0f39ddc5ff4d1b8e38134349c4f018c47d4fbba297510561429aa6c1a2c3/pyobjc_core-13.0-cp315-cp315-macosx_10_15_universal2.whl", hash = "sha256:d65d383efed5a803e43f6c67ccc9f4c94754842b2268da24319bb57a852b016f", size = 6782204, upload-time = "2026-10-30T10:33:08.125Z" }, + { url = "https://files.pythonhosted.org/packages/dd/bc/35eb48cff5d20770a2adb2bb0f01424f5dfbcc7dae8e5ecb990f6d90adf9/pyobjc_core-13.0-cp315-cp315t-macosx_10_15_universal2.whl", hash = "sha256:292462e0947a09aca3737a06c44457cb2aa58f267cfc581c16fca036fc073b45", size = 7036187, upload-time = "2026-10-30T10:33:12.36Z" }, +] + +[[package]] +name = "pyobjc-framework-cocoa" +version = "13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyobjc-core" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/30/50/55a8779d35f1f86d795264cbd6fe62e35b2cf4fb8343fc888e99268c988a/pyobjc_framework_cocoa-13.0.tar.gz", hash = "sha256:bb7a23941fc8d7866578dbc84e56140efd343a1d28510f530dd1dbf80807daf4", size = 3135308, upload-time = "2026-10-30T10:46:03.982Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f2/88/d8cece46725d7879aee3b6a047cbbb86e54aa51b31c707c7d8da6bffa331/pyobjc_framework_cocoa-13.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:16dd1e8890ed695ce12bcd367ed0ac9f3e70dd1f7a8098b6398427c162591b83", size = 391003, upload-time = "2026-10-30T10:35:14.969Z" }, + { url = "https://files.pythonhosted.org/packages/ea/44/98d626726809da517c757d77e2888d34eaf758ebad908effb8d0b60c8118/pyobjc_framework_cocoa-13.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:e1618d315939c9fae57cb858d4db496ec61f82495ce6a877398131ee494fc19d", size = 391597, upload-time = "2026-10-30T10:35:16.217Z" }, + { url = "https://files.pythonhosted.org/packages/01/b0/cb28ff4474a2f07bb9b64751112b3b40c421755297749c70810f4cadfc24/pyobjc_framework_cocoa-13.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:bb0bcba2e4e9d05af875306e0e055cfc2f152fd602d5f2f1a54490adebebcd34", size = 391670, upload-time = "2026-10-30T10:35:17.336Z" }, + { url = "https://files.pythonhosted.org/packages/54/e6/53aeaa92fd8261ba32d7faca85bef4aa210740f4f4462d0e727be6769731/pyobjc_framework_cocoa-13.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:d6fa11ea29d85d46e65ed1c9fdba3a5131ad26971bf24bb7d16aa55f06bcdc1e", size = 395676, upload-time = "2026-10-30T10:35:18.813Z" }, + { url = "https://files.pythonhosted.org/packages/b4/bb/f683ec329180e1931c8d592c882defebca8810e019cf4bc8727fa36b7d89/pyobjc_framework_cocoa-13.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:88be6d452e7531f695ce6d41428aa74bcaade9ddbc50af592419b9f2fbb7c879", size = 391936, upload-time = "2026-10-30T10:35:20.154Z" }, + { url = "https://files.pythonhosted.org/packages/c5/3e/1b24b56d3710a594dc2350ed6b321aaf0ee12fca3dafa6034bccffc8ac08/pyobjc_framework_cocoa-13.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:1a6b9f9e0ab123e51fa7df43bbc6c0414a444facd2225ad4f221a0d842c49ad6", size = 395844, upload-time = "2026-10-30T10:35:21.319Z" }, + { url = "https://files.pythonhosted.org/packages/f2/e0/b2f89d9c92e2b6513d8495d81c3a7d6042f936f51ed7dd2180615a01c980/pyobjc_framework_cocoa-13.0-cp315-cp315-macosx_10_15_universal2.whl", hash = "sha256:9282ebad57784ec18f4b7ea2e6b65bac0a6f56cc21a5c63e061fa48a9b49dd26", size = 392140, upload-time = "2026-10-30T10:35:22.562Z" }, + { url = "https://files.pythonhosted.org/packages/2c/ac/d2c8fa5d96ddf4e70ac5ddb23f24ac7861577728df342697125c1b42e159/pyobjc_framework_cocoa-13.0-cp315-cp315t-macosx_10_15_universal2.whl", hash = "sha256:f1e26ad2225d1c84973cce725f2f4aeeb4d5f653a36c8a9f0b691c5529dc7780", size = 396079, upload-time = "2026-10-30T10:35:23.821Z" }, +] + +[[package]] +name = "pyobjc-framework-metal" +version = "13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyobjc-core" }, + { name = "pyobjc-framework-cocoa" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/27/e6/f6efb0f951df22e8571facc26e7bfde7f1b06cb59bf682737c0bbbbc4593/pyobjc_framework_metal-13.0.tar.gz", hash = "sha256:be798a872021520701183f015f1b0fdadbc5ca74e4aee80d0c4952b890c90a7a", size = 244348, upload-time = "2026-10-30T10:47:02.567Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ac/36/9f2319402d376768a9d1d9ef0b902e1d3d785d52670eb5e918262c2775d1/pyobjc_framework_metal-13.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:f58b3dde81fa568a11ea4098debb827e5cdb130565efe4ce0b1d9ec48a11cd68", size = 78042, upload-time = "2026-10-30T10:40:38.412Z" }, + { url = "https://files.pythonhosted.org/packages/03/a6/a0330996fe6e1082bc0ef9251da651ac9cb00976f54d4be35dc8a977c5fe/pyobjc_framework_metal-13.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a69d760476121cd0faad8ea2f4c70708398aad97eb33fe7c186712e437ecacd7", size = 78028, upload-time = "2026-10-30T10:40:39.555Z" }, + { url = "https://files.pythonhosted.org/packages/9e/6d/12e13ce0db817a809d2fa7e948d688b3cfbe26c39d1f721c1a489daeaf7e/pyobjc_framework_metal-13.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:79d682af412582a092222f4b85b5120f05050ef6c5d5e69a199a4aa3a25a86da", size = 78046, upload-time = "2026-10-30T10:40:40.482Z" }, + { url = "https://files.pythonhosted.org/packages/76/72/f553bc777a59c674dc5b847937b2de844524e09e93dd4e760b4ab1fb9d76/pyobjc_framework_metal-13.0-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:e197115790ee469143708c45b203723ee5589ffc06f1a0006ef53139141cd472", size = 78572, upload-time = "2026-10-30T10:40:41.382Z" }, + { url = "https://files.pythonhosted.org/packages/7c/5d/2d1e4f9eb76bf338f3528aa1511a817889042061c3d847b06b3a86bb419c/pyobjc_framework_metal-13.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:26b6a4e998dc60261d01be45cb9c53c45386a6750af27d612e3fbfc1f2565fea", size = 78138, upload-time = "2026-10-30T10:40:42.316Z" }, + { url = "https://files.pythonhosted.org/packages/9a/94/9a75ef9f0df859e5ceffe0e477f96d72e1a09bdec007811836cd8b7088e7/pyobjc_framework_metal-13.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:72b9d46a05b72a54adf3183a9ac6353346f01345a8e32d1b763b6176142d43d7", size = 78719, upload-time = "2026-10-30T10:40:43.194Z" }, + { url = "https://files.pythonhosted.org/packages/dd/26/57832f2b7704c61e9581663dc36d689d083bd2d9af74602b3fd3889f90d6/pyobjc_framework_metal-13.0-cp315-cp315-macosx_10_15_universal2.whl", hash = "sha256:902f71d42aa282a65b5dbb49c0884543b4a31a17b31d3dd61aaa80e688715fdd", size = 78249, upload-time = "2026-10-30T10:40:44.165Z" }, + { url = "https://files.pythonhosted.org/packages/bf/4a/6e19d461171f1fca713b8d7569df3334d1f3d23b727c660c327ba2dab7ce/pyobjc_framework_metal-13.0-cp315-cp315t-macosx_10_15_universal2.whl", hash = "sha256:1706dcdd2002ff2bf6dc05fdcd0ea0673692359fb20130f3e8a3250a99f88503", size = 78817, upload-time = "2026-10-30T10:40:45.037Z" }, +] + [[package]] name = "pyparsing" version = "3.3.2"