diff --git a/.claude/skills/benchmarking-tsim/SKILL.md b/.claude/skills/benchmarking-tsim/SKILL.md index fdfc0ffb..cde4e168 100644 --- a/.claude/skills/benchmarking-tsim/SKILL.md +++ b/.claude/skills/benchmarking-tsim/SKILL.md @@ -100,10 +100,25 @@ uv run tsim-dev bench --side .:device=gpu --side .:device=cpu,threads=10 # GPU ``` - 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. + back: a component wider than 11, a rotation branch below float32's resolution (an angle + below about 6e-5 π), noise beyond the GPU tables, no GPU memory. Under the default + `device="auto"` these fallbacks are silent; in a script of your own, compile with + `device="metal"` to make them raise instead. +- `TSIM_DEVICE=cpu` (the bench's default) keeps `device="auto"` on the CPU; an explicit + `device=` argument overrides it. +- The GPU's rate depends on call size much more than the CPU's: each call pays a fixed cost + (setting up the launches, waiting for the GPU to finish), and each batch of up to 2^20 shots is one + launch. The bench calibrates calls to about 0.3 s, which is enough. A GPU rate measured + with calls of 10^5 shots or fewer, or with a small `batch_size`, measures overhead, not + sampling; never report one. Leave `batch_size` unset on the GPU; it can only lower the + batch, and it changes a seed's GPU samples. - 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. +- Say whether rates are bit-packed (the default) or `--bools`, above all on a GPU: booleans + cost it 2.3-3.2x on circuits with hundreds of detectors (surface_d7_r7: 757M packed, + 237M booleans on an M4 Pro), against 1.0-1.8x on the CPU, so a boolean comparison + understates the GPU's lead. Compare Tsim with other simulators on the same output format. - 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`, ...). diff --git a/.claude/skills/profiling-tsim/SKILL.md b/.claude/skills/profiling-tsim/SKILL.md index 7ee9c0a4..8782bb30 100644 --- a/.claude/skills/profiling-tsim/SKILL.md +++ b/.claude/skills/profiling-tsim/SKILL.md @@ -31,9 +31,12 @@ Tsim is slow: sampler returns its one observable at 60.9M shots/s per core, the measurement sampler its 10 measurements at 11.6M (Apple M4 Pro). - **Output.** The harness times bit-packed rows, Tsim's own output. `--bools` - times `sample()`'s default boolean output, which adds a conversion: 1.6x on - surface_d7_r7 on the CPU. A slow boolean rate next to a fast bit-packed one - points at the conversion, not the sampler. + times `sample()`'s default boolean output, which adds a conversion: 1.6-1.8x on + surface_d7_r7 on the CPU, and 2.3-3.2x on a GPU for circuits with hundreds of + detectors (Metal on an M4 Pro: surface_d7_r7 3.2x, end2end_d3_dsurface15 2.8x, + hook_inject_T_d15_pd5 2.3x; under 1.3x for circuits with few detectors). A slow + boolean rate next to a fast bit-packed one points at the conversion, not the + sampler. - **Threads.** One pinned thread gives the rate per core. Tsim samples on all performance cores by default (`numba.get_num_threads()`); the `NOTE` line says how many. Quote all-core rates with `--threads N`. @@ -41,9 +44,11 @@ Tsim is slow: 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. + `device` column says where each row sampled: under `device="auto"` a GPU falls + back to the CPU silently (components wider than 11, tiny rotations, noise beyond + its tables), and a row that ends in `DEVICE: cpu -> metal` compares a GPU with the + CPU. A slow GPU row with small calls or a small `batch_size` is launch overhead; + see "Apple-silicon Macs" in the `benchmarking-tsim` skill. - **Compile.** `compile` includes the first use of a backend in the process (about 0.1 s for Metal's PyObjC) but not the first call, whose kernel compilation goes into `warmup_s` in the JSON. diff --git a/.claude/skills/testing-tsim-correctness/SKILL.md b/.claude/skills/testing-tsim-correctness/SKILL.md index 56985caa..0c6cf111 100644 --- a/.claude/skills/testing-tsim-correctness/SKILL.md +++ b/.claude/skills/testing-tsim-correctness/SKILL.md @@ -27,7 +27,9 @@ handful of T gates; its memory grows exponentially. ## Invariants to test when you touch them -- Same seed gives the same samples for any thread count and any `batch_size`. +- Same seed gives the same samples for any thread count and any `batch_size` on the + CPU. On a GPU the samples depend on `batch_size`; same seed and same `batch_size` + give the same samples. - Bit-packed and unpacked outputs agree. Every `sample()` call draws new shots, so compare two fresh samplers with the same seed, not two calls on one sampler. - Postselection returns only the accepted rows; compare accepted-row distributions. @@ -44,6 +46,11 @@ handful of T gates; its memory grows exponentially. test is a broken test. - Stim merges consecutive identical instructions into one with several targets, so write circuits that exercise both single- and multi-target paths. +- GPU tests: `test/conftest.py` keeps every test outside + `test/integration/test_metal_sampler.py` on the CPU. In the Metal tests compile the + GPU side with `device="metal"`, so a silent fallback cannot make the test pass on the + CPU, and the reference with `device="cpu"`; compare them with `assert_counts_match` + (GPU samples never equal CPU samples bit for bit). - Keep circuits small (a few qubits, tens of thousands of shots) so the suite stays near 90 s. Long randomized campaigns belong in a job, not in the suite: `uv run tsim-dev job start fuzz --mem 4G -- `. diff --git a/CHANGELOG.md b/CHANGELOG.md index 1beda668..95081d37 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `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 + depend on `batch_size`. The `device` argument of `compile_sampler`, + `compile_detector_sampler` and the sampler classes chooses where they sample: + `"auto"` (the default) uses the GPU when it can run the circuit and otherwise + the CPU, without a warning; `"cpu"` keeps them on the CPU; `"metal"` raises when + the GPU cannot run the circuit. `TSIM_DEVICE=cpu` keeps `"auto"` on the CPU, and `SamplerStats.device` and the repr say where they sample. - 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 diff --git a/docs/install.mdx b/docs/install.mdx index 89138f0b..f452828a 100644 --- a/docs/install.mdx +++ b/docs/install.mdx @@ -38,30 +38,47 @@ 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`. +through Metal; the plain install includes PyObjC's Metal bindings for it. Where it samples is +the `device` argument of `compile_sampler` and `compile_detector_sampler`: + +| `device` | Samples | +| --- | --- | +| `"auto"` (default) | On the GPU when it can run the circuit, else on the CPU, without a warning | +| `"cpu"` | On the CPU, in double precision | +| `"metal"` | On the GPU, or raises an error that says why it cannot | + +```python +sampler = circuit.compile_detector_sampler(device="metal") +``` + +The GPU runs circuits whose components have at most 11 active qubits. That includes Clifford +circuits, whose detectors depend on the noise alone: the GPU samples the noise as well as the +state vectors. Postselection runs on +the GPU, so heavily postselected circuits gain the most. A sampler's repr says where it +samples (`samples on metal` or `samples on cpu`), and so does `sampler.stats.device`. 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`. +enough to make such an outcome (an `R_Z` or `R_X` by less than about 6e-5 π) samples on the +CPU (or raises, with `device="metal"`). + +A seed's samples on the GPU differ from the CPU's and depend on `batch_size`, so with +`device="auto"` a seed reproduces its samples only on the same kind of machine. Use +`device="cpu"` to reproduce samples across machines. ### 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: +Set `TSIM_DEVICE=cpu` to keep the samplers of `device="auto"` on the CPU without changing +code, 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. +`compile_sampler` or `compile_detector_sampler` works too. An explicit `device` argument +overrides it. ## Agent skill diff --git a/skills/simulating-circuits-with-tsim/SKILL.md b/skills/simulating-circuits-with-tsim/SKILL.md index b3bfbc4c..9839136f 100644 --- a/skills/simulating-circuits-with-tsim/SKILL.md +++ b/skills/simulating-circuits-with-tsim/SKILL.md @@ -104,8 +104,9 @@ If `sampler.active_width` raises `AttributeError`, the installed Tsim is older t upgrade `bloqade-tsim` before going further. Before a long run, time it: one small warm-up call (the first call compiles kernels), then -one call of about 10^6 shots, and extrapolate. Small calls understate the rate. Report the -width and the measured shots per second, with the thread count, along with any plan. +one call of about 10^6 shots, and extrapolate. Small calls understate the rate, on a GPU even +more. Report the width, the measured shots per second and where it ran (`sampler.stats.device`, +and the thread count on the CPU), along with any plan. If the width is too large: postselect inside Tsim (discarded shots stop early, see below), simulate a smaller instance (lower distance or fewer rounds) and report how the width grows, @@ -176,13 +177,19 @@ yourself. - **Shots per call:** each `sample()` call has a fixed overhead. For fast circuits, ask for about 10^6 or more shots per call, limited by memory, instead of looping over small calls. Calls of 10^5 shots have been measured 2 to 4 times slower than calls of 10^7, and 100-shot - calls about 80 times slower. + calls about 80 times slower. On a GPU the fixed cost per call is larger (setting up the + launches and waiting for the GPU to finish), so ask for 10^6 to 10^7 shots per call there. - **`batch_size`:** leave it at the default, which is sized from free memory and is close to optimal. Small values (a few thousand or less) are several times slower. Lower it only to - cap memory; it does not change the samples. + cap memory; on the CPU it does not change the samples. On a GPU the default batch is up to 2^20 + shots (fewer for very wide rows), one GPU launch, all launched back to back; `batch_size` can only lower it, a small + one makes launch overhead dominate, and it changes which samples a seed gives. - **`bit_packed=True`** returns little-endian packed `uint8` rows, as in Stim, using 8 times less memory. On a distance-11 surface code memory (1,320 detectors) it was twice as fast as - booleans, at distance 5 about 1.3 times. For small circuits it gains little and is slower + booleans, at distance 5 about 1.3 times. On an Apple GPU the gap is larger, since booleans + move 8 times more bytes off the GPU: 2.3 to 3.2 times on circuits with hundreds of + detectors (surface_d7_r7: 757M against 237M shots/s on an M4 Pro). Use packed rows for + large studies and decode them packed. For small circuits it gains little and is slower with `separate_observables` (the observables are split off and re-packed); use `append_observables=True` with packed output. - **Threads:** the default strategy runs on `numba.get_num_threads()` threads, all cores by @@ -190,13 +197,29 @@ yourself. 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 + component has at most 11 active qubits, often 5 to 10 times faster than all CPU cores. That + includes Clifford circuits, whose detectors come from the noise alone: the GPU samples the + noise too. It computes outcome probabilities in float32 and never draws an outcome below about 1e-11 - (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`). + (circuits with rotations by less than about 6e-5 π stay on the CPU). Thread settings do not + affect GPU sampling. +- **`device`** (in `compile_sampler` and `compile_detector_sampler`) chooses where it samples: + `"auto"` (default) uses the GPU when it can and otherwise the CPU, silently; `"cpu"` stays + on the CPU in double precision; `"metal"` raises with the reason if the GPU cannot run the + circuit (no Apple GPU, width above 11, a tiny rotation, out of GPU memory). Check + `sampler.stats.device` (or the repr, `samples on metal`) after compiling, since `"auto"` + does not warn. Pass `device="metal"` for long runs that must not drop to the CPU unnoticed, + and `device="cpu"` to reproduce samples across machines or to compare against the GPU. + `TSIM_DEVICE=cpu` in the environment keeps `"auto"` on the CPU; it is read when a sampler + is compiled, so an unpickled sampler keeps the device choice it was compiled with. +- **Processes:** sample from one process per GPU. Threads sharing a GPU sampler run its calls + one at a time, and several processes compete for the same GPU. A process forked after its parent used the GPU cannot use + it (it samples on the CPU, or raises with `device="metal"`); start workers with the + `"spawn"` or `"forkserver"` method (macOS's default is `"spawn"`). - **Seeds:** a `seed` reproduces samples for the same Tsim and dependency versions and - backend. On the CPU, thread count and `batch_size` do not affect the samples; on a GPU, - `batch_size` does. + device. On the CPU, thread count and `batch_size` do not affect the samples; on a GPU, + `batch_size` does, and GPU samples differ from CPU samples, so with `device="auto"` a seed + reproduces only on the same kind of machine. ## Postselection @@ -286,7 +309,7 @@ tick it off: ``` Tsim results: -- [ ] Active width read; rate timed with a warm-up and a large call; thread count noted +- [ ] Active width read; rate timed with a warm-up and a large call; device and threads noted - [ ] Noiseless run: detectors used for postselection or decoding never fire - [ ] Reference flags used only if detectors and observables are deterministic - [ ] At least one check below passed (stand-in vs Stim, to_tensor, independent simulator) diff --git a/skills/simulating-circuits-with-tsim/references/qec-experiments.md b/skills/simulating-circuits-with-tsim/references/qec-experiments.md index c4ecf6d6..581bcb8d 100644 --- a/skills/simulating-circuits-with-tsim/references/qec-experiments.md +++ b/skills/simulating-circuits-with-tsim/references/qec-experiments.md @@ -56,7 +56,7 @@ sinter can call Tsim, but it costs speed and can cost correctness: - sinter's decoders expect Stim's flips relative to a noiseless reference; Tsim returns raw values, so a detector that is 1 without noise looks like an error. - Each sinter worker process compiles its own sampler and starts its own numba thread pool, so - several workers oversubscribe the CPU. + several workers oversubscribe the CPU, or queue on the one GPU of an Apple-silicon Mac. Write the short loop below instead. sinter's helpers are still useful for the statistics (`sinter.fit_binomial`) and for plotting. @@ -124,7 +124,8 @@ Notes: `max_seconds` as caps, rather than fixing the shot count in advance: logical error rates are often orders of magnitude away from the first guess. - Calls are sized from a memory budget (`max_bytes_per_call`, 1 GB by default), not a fixed - shot count: 10^6 shots of a circuit with 50,000 detectors is a 50 GB array. + shot count: 10^6 shots of a circuit with 50,000 detectors is a 50 GB array. Keep calls + large on a GPU as well: below about 10^6 shots per call its fixed cost per call shows. - `run_point` works on boolean rows, which the decoder recipes expect. For circuits with hundreds of detectors or more, packed rows use 8 times less memory and can be twice as fast; to use them, change the loop to sample with `bit_packed=True` and `append_observables=True`, diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index 816109d2..d90191b2 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -24,7 +24,7 @@ from tsim.core.graph import transform_error_basis from tsim.core.parse import parse_parametric_tag, parse_stim_circuit from tsim.core.tags import encode_t_tag -from tsim.core.types import SamplingStrategy +from tsim.core.types import Device, SamplingStrategy from tsim.noise.dem import get_detector_error_model from tsim.utils.clifford import expand_clifford_rotations, is_clifford from tsim.utils.diagram import draw_zx, render_svg @@ -947,6 +947,7 @@ def compile_sampler( *, strategy: SamplingStrategy = "active_statevector", seed: int | None = None, + device: Device = "auto", ) -> CompiledMeasurementSampler: """Compile circuit into a measurement sampler. @@ -954,11 +955,7 @@ 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 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. + "bss", "cutting"). The stabilizer rank strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed reproduces its samples only with the same versions of Tsim and its dependencies (such as NumPy and JAX) and the same backend (CPU or GPU). @@ -968,20 +965,39 @@ def compile_sampler( need the same explicit ``batch_size``, since the default batch size depends on the free memory. On a GPU, they do depend on the batch size. + device: Where "active_statevector" samples. "auto" (the default) + samples in float32 on an Apple GPU when there is one that can + run the circuit, and otherwise on the CPU, without a warning + (``TSIM_DEVICE=cpu`` keeps it on the CPU). "cpu" samples on the + CPU in float64, on ``numba.get_num_threads()`` threads. "metal" + samples on an Apple GPU or raises: without one, for a component + of more than 11 active qubits, a rotation by less than about + 6e-5 π, noise beyond the GPU's tables, or a call that runs out + of GPU memory. The stabilizer rank strategies take only "auto". Returns: A CompiledMeasurementSampler that can be used to sample measurements. + Raises: + ValueError: If ``device`` is unknown, is not "auto" under a + stabilizer-rank strategy, or is "metal" for a circuit that the + GPU executor cannot run. + RuntimeError: If ``device`` is "metal" and this process cannot use + an Apple GPU. + """ from tsim.sampler import CompiledMeasurementSampler - return CompiledMeasurementSampler(self, seed=seed, strategy=strategy) + return CompiledMeasurementSampler( + self, seed=seed, strategy=strategy, device=device + ) def compile_detector_sampler( self, *, strategy: SamplingStrategy = "active_statevector", seed: int | None = None, + device: Device = "auto", ) -> CompiledDetectorSampler: """Compile circuit into a detector sampler. @@ -994,11 +1010,7 @@ 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 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. + "bss", "cutting"). The stabilizer rank strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed reproduces its samples only with the same versions of Tsim and its dependencies (such as NumPy and JAX) and the same backend (CPU or GPU). @@ -1008,14 +1020,32 @@ def compile_detector_sampler( 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. + device: Where "active_statevector" samples. "auto" (the default) + samples in float32 on an Apple GPU when there is one that can + run the circuit, and otherwise on the CPU, without a warning + (``TSIM_DEVICE=cpu`` keeps it on the CPU). "cpu" samples on the + CPU in float64, on ``numba.get_num_threads()`` threads. "metal" + samples on an Apple GPU or raises: without one, for a component + of more than 11 active qubits, a rotation by less than about + 6e-5 π, noise beyond the GPU's tables, or a call that runs out + of GPU memory. The stabilizer rank strategies take only "auto". Returns: A CompiledDetectorSampler that can be used to sample detectors and observables. + Raises: + ValueError: If ``device`` is unknown, is not "auto" under a + stabilizer-rank strategy, or is "metal" for a circuit that the + GPU executor cannot run. + RuntimeError: If ``device`` is "metal" and this process cannot use + an Apple GPU. + """ from tsim.sampler import CompiledDetectorSampler - return CompiledDetectorSampler(self, seed=seed, strategy=strategy) + return CompiledDetectorSampler( + self, seed=seed, strategy=strategy, device=device + ) def inverse(self) -> Circuit: """Return the inverse of the circuit.""" diff --git a/src/tsim/core/types.py b/src/tsim/core/types.py index 65de7c19..d142170b 100644 --- a/src/tsim/core/types.py +++ b/src/tsim/core/types.py @@ -28,6 +28,9 @@ SamplingStrategy = Literal["active_statevector", "cat5", "bss", "cutting"] """``"active_statevector"``, or a stabilizer-rank decomposition strategy.""" +Device = Literal["auto", "cpu", "metal"] +"""Where ``"active_statevector"`` samples: ``"auto"`` picks a GPU when one can run it.""" + @dataclass(frozen=True) class SamplingGraph: diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 10d08ea4..325438b0 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -18,7 +18,12 @@ from tsim.compile.evaluate import evaluate from tsim.compile.pipeline import compile_program from tsim.core.graph import prepare_graph -from tsim.core.types import CompiledComponent, CompiledProgram, SamplingStrategy +from tsim.core.types import ( + CompiledComponent, + CompiledProgram, + Device, + SamplingStrategy, +) from tsim.noise.channels import ChannelSampler from tsim.statevector import metal from tsim.statevector.program import StatevectorComponent @@ -325,6 +330,7 @@ def __init__( mode: Literal["sequential", "joint"], strategy: SamplingStrategy = "active_statevector", seed: int | None = None, + device: Device = "auto", ): """Initialize the sampler by compiling the circuit. @@ -335,11 +341,7 @@ 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 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. + "bss", "cutting"). The stabilizer rank strategies are provisional. seed: Random seed for the sampler. If None, a random seed is drawn. A seed reproduces its samples only with the same versions of Tsim and its dependencies (such as NumPy and JAX) and the same backend (CPU or GPU). @@ -349,8 +351,31 @@ def __init__( need the same explicit ``batch_size``, since the default batch size depends on the free memory. On a GPU, they do depend on the batch size. + device: Where "active_statevector" samples. "auto" (the default) + samples in float32 on an Apple GPU when there is one that can + run the circuit, and otherwise on the CPU, without a warning + (``TSIM_DEVICE=cpu`` keeps it on the CPU). "cpu" samples on the + CPU in float64, on ``numba.get_num_threads()`` threads. "metal" + samples on an Apple GPU or raises: without one, for a component + of more than 11 active qubits, a rotation by less than about + 6e-5 π, noise beyond the GPU's tables, or a call that runs out + of GPU memory. The stabilizer rank strategies take only "auto". + + Raises: + ValueError: If ``device`` is unknown, is not "auto" under a + stabilizer-rank strategy, or is "metal" for a circuit that the + GPU executor cannot run. + RuntimeError: If ``device`` is "metal" and this process cannot use + an Apple GPU. """ + if device not in ("auto", "cpu", "metal"): + raise ValueError(f"device must be 'auto', 'cpu' or 'metal', not {device!r}") + if device != "auto" and strategy != "active_statevector": + raise ValueError( + f"device={device!r} needs strategy='active_statevector'; the " + "stabilizer-rank strategies run wherever JAX does" + ) if seed is None: seed = secrets.randbits(128) # NOTE: the noise, the shot words, the JAX key and the GPU draw from children @@ -384,14 +409,17 @@ def __init__( 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" - ): + # NOTE: as JAX does for the stabilizer-rank strategies, "auto" uses a + # GPU whenever there is one that can run the program. + if device == "auto" and os.environ.get("TSIM_DEVICE", "").lower() == "cpu": + device = "cpu" + if mode == "sequential" and device != "cpu": self._shots = metal.shot_sampler( - host, self._channel_sampler, self._num_detectors, gpu_seed + host, + self._channel_sampler, + self._num_detectors, + gpu_seed, + required=device == "metal", ) else: self._backend = _StabilizerRankBackend(program, key_seed) @@ -743,7 +771,10 @@ def stats(self) -> SamplerStats: - ``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). + strategy runs wherever JAX does). With ``device="auto"`` it is where + the sampler ended up: ``"cpu"`` after its GPU kernel failed to + compile or in a process forked after its parent used the GPU. A call + that runs out of GPU memory samples on the CPU without changing it. """ components = self._program.components widths: tuple[int, ...] | None = None @@ -861,6 +892,7 @@ def __init__( *, strategy: SamplingStrategy = "active_statevector", seed: int | None = None, + device: Device = "auto", ): """Create a measurement sampler. @@ -878,6 +910,22 @@ def __init__( need the same explicit ``batch_size``, since the default batch size depends on the free memory. On a GPU, they do depend on the batch size. + device: Where "active_statevector" samples. "auto" (the default) + samples in float32 on an Apple GPU when there is one that can + run the circuit, and otherwise on the CPU, without a warning + (``TSIM_DEVICE=cpu`` keeps it on the CPU). "cpu" samples on the + CPU in float64, on ``numba.get_num_threads()`` threads. "metal" + samples on an Apple GPU or raises: without one, for a component + of more than 11 active qubits, a rotation by less than about + 6e-5 π, noise beyond the GPU's tables, or a call that runs out + of GPU memory. The stabilizer rank strategies take only "auto". + + Raises: + ValueError: If ``device`` is unknown, is not "auto" under a + stabilizer-rank strategy, or is "metal" for a circuit that the + GPU executor cannot run. + RuntimeError: If ``device`` is "metal" and this process cannot use + an Apple GPU. """ super().__init__( @@ -886,6 +934,7 @@ def __init__( mode="sequential", seed=seed, strategy=strategy, + device=device, ) def sample( @@ -962,6 +1011,7 @@ def __init__( *, strategy: SamplingStrategy = "active_statevector", seed: int | None = None, + device: Device = "auto", ): """Create a detector sampler. @@ -979,6 +1029,22 @@ def __init__( 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. + device: Where "active_statevector" samples. "auto" (the default) + samples in float32 on an Apple GPU when there is one that can + run the circuit, and otherwise on the CPU, without a warning + (``TSIM_DEVICE=cpu`` keeps it on the CPU). "cpu" samples on the + CPU in float64, on ``numba.get_num_threads()`` threads. "metal" + samples on an Apple GPU or raises: without one, for a component + of more than 11 active qubits, a rotation by less than about + 6e-5 π, noise beyond the GPU's tables, or a call that runs out + of GPU memory. The stabilizer rank strategies take only "auto". + + Raises: + ValueError: If ``device`` is unknown, is not "auto" under a + stabilizer-rank strategy, or is "metal" for a circuit that the + GPU executor cannot run. + RuntimeError: If ``device`` is "metal" and this process cannot use + an Apple GPU. """ super().__init__( @@ -987,6 +1053,7 @@ def __init__( mode="sequential", seed=seed, strategy=strategy, + device=device, ) @overload diff --git a/src/tsim/statevector/gpu/__init__.py b/src/tsim/statevector/gpu/__init__.py index a1e4391b..e5582ea1 100644 --- a/src/tsim/statevector/gpu/__init__.py +++ b/src/tsim/statevector/gpu/__init__.py @@ -8,7 +8,7 @@ """ FORKED = ( - "This process was forked after its parent used the GPU, so it samples on the " - "CPU; start it with the 'spawn' or 'forkserver' method to sample on the GPU." + "this process was forked after its parent used the GPU; start it with the " + "'spawn' or 'forkserver' method to sample on the GPU" ) -"""The warning of a GPU sampler in a process forked after its parent used the GPU.""" +"""Why a GPU cannot sample in a process forked after its parent used the GPU.""" diff --git a/src/tsim/statevector/metal/__init__.py b/src/tsim/statevector/metal/__init__.py index 541064a5..c9e5fe50 100644 --- a/src/tsim/statevector/metal/__init__.py +++ b/src/tsim/statevector/metal/__init__.py @@ -3,17 +3,15 @@ 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, +``ShotSampler`` behind the compiled samplers of ``device="auto"`` or +``"metal"`` 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 @@ -22,10 +20,7 @@ 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 + from tsim.statevector.shots import ShotLayout, ShotSampler def shot_sampler( @@ -33,50 +28,64 @@ def shot_sampler( channels: ChannelSampler, num_detectors: int, seed: np.random.SeedSequence, + *, + required: bool, ) -> ShotSampler | MetalShotSampler: - """Return a Metal shot sampler of ``host``'s program when one can run it, else ``host``. + """Return a Metal shot sampler of ``host``'s program, or ``host`` where Metal cannot run it. + + Without ``required`` (``device="auto"``) it returns ``host`` silently; with it + (``device="metal"``) it raises instead. A program without shot words (one + with no outputs; direct outputs are shot-word bits too) has nothing to + sample on a GPU and keeps ``host`` either way. + + Raises: + RuntimeError: With ``required``, if this process cannot use an Apple GPU: + no Apple GPU, no PyObjC's Metal, or forked after its parent used the GPU. + ValueError: With ``required``, if the GPU executor cannot run the program: + a component wider than ``MAX_WIDTH``, a rotation by an angle too small + for float32 amplitudes (``RESOLVABLE``), or noise beyond the GPU + tables' limits. + MemoryError: With ``required``, if the GPU runs out of memory. - It warns when 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) + reason = runtime.unavailable() + if reason is not None: + if required: + raise RuntimeError(f"Cannot sample on Metal: {reason}.") return host - if not runtime.available() or host.layout.num_words == 0: + if host.layout.num_words == 0: return host - if not supports(host.layout): - width = max(c.width for c in host.layout.components) - 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, - ) + reason = _unsupported(host.layout) + if reason is not None: + if required: + raise ValueError(f"Cannot sample on Metal: {reason}.") return host from tsim.statevector.metal.sampler import MetalShotSampler try: - return MetalShotSampler(host, channels, num_detectors, seed) + return MetalShotSampler(host, channels, num_detectors, seed, required=required) except MemoryError: - warnings.warn( - "The GPU is out of memory; sampling on the CPU.", stacklevel=_STACKLEVEL - ) + if required: + raise except ValueError as error: # NOTE: noise beyond the GPU tables' limits (channels of more than 127 # outcomes, outcomes that flip more than 255 words). - warnings.warn( - f"The GPU cannot sample this circuit ({error}); sampling on the CPU.", - stacklevel=_STACKLEVEL, - ) + if required: + raise ValueError(f"Cannot sample on Metal: {error}.") from error return host + + +def _unsupported(layout: ShotLayout) -> str | None: + """Return why the GPU executor cannot run ``layout``, or ``None`` if it can.""" + if supports(layout): + return None + width = max(c.width for c in layout.components) + if width > MAX_WIDTH: + return ( + f"a component needs {width} active qubits, more than the " + f"{MAX_WIDTH} that the GPU executor supports" + ) + return ( + f"a rotation has a branch of weight {rarest_branch(layout):.1e}, " + f"below the {RESOLVABLE:.0e} that float32 amplitudes resolve" + ) diff --git a/src/tsim/statevector/metal/runtime.py b/src/tsim/statevector/metal/runtime.py index 24ea971a..2c573a02 100644 --- a/src/tsim/statevector/metal/runtime.py +++ b/src/tsim/statevector/metal/runtime.py @@ -16,6 +16,8 @@ import numpy as np +from tsim.statevector.gpu import FORKED + @functools.cache def metal() -> ModuleType: @@ -140,16 +142,25 @@ def available() -> bool: It is not in a process forked after its parent used the GPU. """ - return not _forked_after_init and _device_available() + return unavailable() is None + + +def unavailable() -> str | None: + """Return why this process cannot sample on Metal, or ``None`` if it can.""" + if _forked_after_init: + return FORKED + return _device_error() @functools.cache -def _device_available() -> bool: +def _device_error() -> str | None: try: check_device() - return True - except Exception: - return False + except ImportError: + return "PyObjC's Metal (pyobjc-framework-Metal) is not installed" + except Exception as error: + return str(error) or type(error).__name__ + return None @dataclass diff --git a/src/tsim/statevector/metal/sampler.py b/src/tsim/statevector/metal/sampler.py index 4b58dd89..c92ce9ba 100644 --- a/src/tsim/statevector/metal/sampler.py +++ b/src/tsim/statevector/metal/sampler.py @@ -31,7 +31,6 @@ from __future__ import annotations import threading -import warnings from dataclasses import dataclass, field from typing import Any @@ -65,9 +64,6 @@ # 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.""" @@ -154,10 +150,13 @@ def __init__( channels: ChannelSampler, num_detectors: int, seed: np.random.SeedSequence, + *, + required: bool = False, ) -> None: """Upload the noise tables of ``host``'s shot words. - ``host`` draws the reference samples, on the CPU. + ``host`` draws the reference samples, on the CPU, and the calls that the + GPU cannot run, unless ``required`` (``device="metal"``): then those raise. Raises: ImportError: If PyObjC's Metal is not installed. @@ -166,7 +165,8 @@ def __init__( """ runtime.check_device() assert self.supports(host.layout) - self._args = (host, channels, num_detectors, seed) + self._args = (host, channels, num_detectors, seed, required) + self._required = required # NOTE: a call's batches share the sampler's buffers, so calls from several # threads run one at a time. self._lock = threading.RLock() @@ -187,8 +187,14 @@ def __init__( @property def device(self) -> str: - """Return where this sampler samples (``SamplerStats.device``).""" - return "cpu" if self._cpu_only else "metal" + """Return where this sampler samples (``SamplerStats.device``). + + It is ``"cpu"`` after a kernel failed to compile, and in a process forked + after its parent used the GPU, where every call samples on the CPU. + """ + if self._cpu_only or runtime.forked_after_init(): + return "cpu" + return "metal" def __reduce__(self): """Pickle the CPU sampler and the inputs; the GPU state is rebuilt on load.""" @@ -488,8 +494,8 @@ def sample( As ``ShotSampler.sample``: with a ``mask`` only the shots that pass postselection are returned. ``batch_size`` bounds the shots of a batch. Out of GPU memory, 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. + forked after its parent used the GPU, the call samples on the CPU, or + raises if the sampler is ``required`` on the GPU (``_or_cpu``). """ def on_gpu(): @@ -519,25 +525,35 @@ def on_gpu(): ) def _or_cpu(self, on_gpu, on_cpu): - """Return ``on_gpu()``, or warn and return ``on_cpu()`` where the GPU fails. + """Return ``on_gpu()``, or ``on_cpu()`` where the GPU fails. It fails in a process forked after its parent used the GPU, out of GPU - memory, and when a kernel does not compile; after that the sampler - samples on the CPU. + memory, and when a kernel does not compile; after the last the sampler + samples on the CPU. A ``required`` sampler raises instead. + + Raises: + RuntimeError: If ``required``, in a process forked after its parent + used the GPU. + MemoryError: If ``required``, out of GPU memory. + runtime.CompileError: If ``required``, when a kernel does not compile. + """ - if self._cpu_only or _forked(): + if self._cpu_only: + return on_cpu() + if runtime.forked_after_init(): + if self._required: + raise RuntimeError(f"Cannot sample on Metal: {FORKED}.") return on_cpu() try: return on_gpu() - except MemoryError as error: - _warn_out_of_memory(error) + except MemoryError: + if self._required: + raise return on_cpu() - except runtime.CompileError as error: + except runtime.CompileError: + if self._required: + raise 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: @@ -764,14 +780,6 @@ def direct_bools( ) -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 @@ -784,26 +792,23 @@ def _threads_per_shot(layout) -> int: 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 + host, channels, num_detectors, seed, required, 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 + """Rebuild a pickled ``MetalShotSampler``, or its CPU sampler where Metal cannot run. + + As at compile time, a sampler ``required`` on the GPU raises where Metal cannot run. + """ + from tsim.statevector import metal + + shots = metal.shot_sampler(host, channels, num_detectors, seed, required=required) + if isinstance(shots, MetalShotSampler): + shots._rng.bit_generator.state = state return shots diff --git a/test/conftest.py b/test/conftest.py index d12e57ba..1435e449 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -7,4 +7,4 @@ 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) + monkeypatch.setattr(metal_runtime, "unavailable", lambda: "a test on the CPU") diff --git a/test/integration/test_metal_sampler.py b/test/integration/test_metal_sampler.py index 2c8730ce..00d0ab8d 100644 --- a/test/integration/test_metal_sampler.py +++ b/test/integration/test_metal_sampler.py @@ -75,10 +75,8 @@ def _gpu_shots(sampler) -> MetalShotSampler: def _cpu(compile, **kwargs): - """Return ``compile(**kwargs)``, a sampler compiled as on a host without a GPU.""" - with pytest.MonkeyPatch.context() as patch: - patch.setattr(runtime, "available", lambda: False) - sampler = compile(**kwargs) + """Return ``compile(**kwargs)``, a sampler on the CPU.""" + sampler = compile(device="cpu", **kwargs) assert sampler._shots is None or type(sampler._shots) is ShotSampler return sampler @@ -548,9 +546,12 @@ def test_metal_shots_start_after_the_common_prefix(width): def test_components_wider_than_the_limit_run_on_the_cpu(): circuit = _random_unitary_circuit(13, 12) assert _widest(circuit) > tape.MAX_WIDTH - with pytest.warns(UserWarning, match="sampling on the CPU"): + with warnings.catch_warnings(): + warnings.simplefilter("error") sampler = circuit.compile_sampler(seed=1) assert type(sampler._shots) is ShotSampler + with pytest.raises(ValueError, match="more than the 11 that the GPU executor"): + circuit.compile_sampler(seed=1, device="metal") @pytest.mark.parametrize(("probability", "shots"), [(2e-8, 1 << 30), (1e-6, 1 << 26)]) @@ -653,9 +654,10 @@ def test_metal_returned_rows_are_not_reused_while_alive(): 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) + for device in ("auto", "metal"): + sampler = circuit.compile_detector_sampler(seed=1, device=device) + assert type(sampler._shots) is ShotSampler + assert sampler.sample(10).shape == (10, 0) def _sampler_in_child(_): @@ -663,16 +665,23 @@ def _sampler_in_child(_): warnings.simplefilter("always") sampler = Circuit(T_STATE_X).compile_sampler(seed=1) sampler.sample(10) - return type(sampler._shots).__name__, [str(w.message) for w in caught] + assert sampler.stats.device == "cpu" + try: + Circuit(T_STATE_X).compile_sampler(seed=1, device="metal") + error = "" + except RuntimeError as raised: + error = str(raised) + return type(sampler._shots).__name__, [str(w.message) for w in caught], error def test_processes_forked_after_the_gpu_was_used_sample_on_the_cpu(): # NOTE: 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] + name, messages, error = pool.map(_sampler_in_child, [0])[0] assert name == "ShotSampler" - assert any("forked" in m for m in messages) + assert messages == [] + assert "forked" in error def test_a_kernel_that_does_not_compile_samples_on_the_cpu(monkeypatch): @@ -685,15 +694,17 @@ def broken(*args, **kwargs): 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 = sampler.sample(20_000) + assert sampler.stats.device == "cpu" got = np.concatenate([got, sampler.sample(20_000)]) assert len(compiled) == 1 expected = _cpu(circuit.compile_sampler, seed=2).sample(40_000) assert_counts_match(got, expected) + required = _gpu(circuit.compile_sampler(seed=3, device="metal")) + with pytest.raises(runtime.CompileError): + required.sample(10) def test_noise_of_tiny_probabilities_fires_at_its_rate(): @@ -740,9 +751,12 @@ def test_rare_outcome_of_a_channel_is_drawn_at_its_rate(): 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"): + with warnings.catch_warnings(): + warnings.simplefilter("error") sampler = circuit.compile_detector_sampler(seed=1) assert sampler.stats.device == "cpu" + with pytest.raises(ValueError, match=r"rotation has a branch of weight 2\.5e-10"): + circuit.compile_detector_sampler(seed=1, device="metal") @pytest.mark.parametrize("p0", [3e-8, 1e-6]) @@ -783,12 +797,17 @@ 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"): + with warnings.catch_warnings(): + warnings.simplefilter("error") got = sampler.sample(20_000, append_observables=True) expected = _cpu(circuit.compile_detector_sampler, seed=2).sample( 20_000, append_observables=True ) assert_counts_match(got, expected) + required = _gpu(circuit.compile_detector_sampler(seed=3, device="metal")) + monkeypatch.setattr(_gpu_shots(required), "sample_accepted", out_of_memory) + with pytest.raises(MemoryError, match="the GPU ran out of memory"): + required.sample(10, append_observables=True) def test_samplers_out_of_gpu_memory_at_setup_run_on_the_cpu(monkeypatch): @@ -796,9 +815,12 @@ 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"): + with warnings.catch_warnings(): + warnings.simplefilter("error") sampler = Circuit(T_STATE_X).compile_sampler(seed=1) assert type(sampler._shots) is ShotSampler + with pytest.raises(MemoryError, match="no Metal buffer"): + Circuit(T_STATE_X).compile_sampler(seed=1, device="metal") def test_tsim_device_cpu_keeps_the_samplers_on_the_cpu(monkeypatch): @@ -807,8 +829,12 @@ def test_tsim_device_cpu_keeps_the_samplers_on_the_cpu(monkeypatch): warnings.simplefilter("error") sampler = Circuit(T_STATE_X).compile_sampler(seed=1) assert type(sampler._shots) is ShotSampler + # NOTE: an explicit device overrides TSIM_DEVICE. + _gpu(Circuit(T_STATE_X).compile_sampler(seed=1, device="metal")) monkeypatch.setenv("TSIM_DEVICE", "") _gpu(Circuit(T_STATE_X).compile_sampler(seed=1)) + sampler = Circuit(T_STATE_X).compile_sampler(seed=1, device="cpu") + assert type(sampler._shots) is ShotSampler @pytest.mark.parametrize("text", [GHZ_WITH_DIRECT_DETECTOR, NOISY_DETECTORS]) @@ -915,16 +941,24 @@ def _sample_in_child(_): 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] + return ( + rows.shape, + [str(w.message) for w in caught], + SAMPLER_BEFORE_FORK.stats.device, + ) 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] + shape, messages, device = pool.map_async(_sample_in_child, [0]).get(timeout=60)[ + 0 + ] assert shape == (100, 0) - assert any("forked" in m for m in messages) + assert messages == [] + assert device == "cpu" + assert SAMPLER_BEFORE_FORK.stats.device == "metal" def test_gpu_samplers_pickle_and_continue_their_streams(): @@ -934,6 +968,8 @@ def test_gpu_samplers_pickle_and_continue_their_streams(): loaded = pickle.loads(pickle.dumps(sampler)) # noqa: S301 assert isinstance(loaded._shots, MetalShotSampler) np.testing.assert_array_equal(loaded.sample(5000), sampler.sample(5000)) + required = circuit.compile_detector_sampler(seed=26, device="metal") + assert _gpu_shots(pickle.loads(pickle.dumps(required)))._required # noqa: S301 def test_noise_beyond_the_gpu_tables_samples_on_the_cpu(): @@ -944,10 +980,13 @@ def test_noise_beyond_the_gpu_tables_samples_on_the_cpu(): ["R 0 1 2 3 4 5 6 7 8", "H 8", "T 8", "H 8", *chain, "M 0 1 2 3 4 5 6 7 8"] + [f"DETECTOR rec[{k - 9}]" for k in range(9)] ) - with pytest.warns(UserWarning, match="cannot sample this circuit"): + with warnings.catch_warnings(): + warnings.simplefilter("error") sampler = Circuit(text).compile_detector_sampler(seed=1) assert type(sampler._shots) is ShotSampler assert sampler.sample(10).shape == (10, 9) + with pytest.raises(ValueError, match="channels of up to 127 outcomes"): + Circuit(text).compile_detector_sampler(seed=1, device="metal") def test_calls_larger_than_the_counters_run_in_parts(monkeypatch): diff --git a/test/unit/test_sampler.py b/test/unit/test_sampler.py index c1220dd0..f1dc53e1 100644 --- a/test/unit/test_sampler.py +++ b/test/unit/test_sampler.py @@ -1,5 +1,6 @@ import subprocess import sys +import warnings from test.helpers.statevector import shot_sampler from unittest.mock import patch @@ -9,6 +10,7 @@ import tsim.sampler as sampler_module import tsim.statevector.shared as shared from tsim.circuit import Circuit +from tsim.statevector.shots import ShotSampler from tsim.utils.kernels import PARALLEL_WORK @@ -813,3 +815,37 @@ def test_only_the_stabilizer_rank_strategies_start_jax(): # importing Tsim and sampling with "active_statevector" must not start it. A # fresh interpreter, since other tests start it. subprocess.run([sys.executable, "-c", _JAX_BACKEND_CHECK], check=True) # noqa: S603 + + +# NOTE: a T state measured after noise: one component with shot words. +NOISY_T_STATE = "R 0\nH 0\nT 0\nZ_ERROR(0.1) 0\nH 0\nM 0\nDETECTOR rec[-1]" + + +@pytest.mark.parametrize("compile", ["compile_sampler", "compile_detector_sampler"]) +def test_device_cpu_samples_on_the_cpu(compile): + sampler = getattr(Circuit(NOISY_T_STATE), compile)(seed=1, device="cpu") + assert type(shot_sampler(sampler)) is ShotSampler + assert sampler.stats.device == "cpu" + + +def test_unknown_devices_raise(): + with pytest.raises(ValueError, match="device must be 'auto', 'cpu' or 'metal'"): + Circuit(NOISY_T_STATE).compile_sampler(device="gpu") # type: ignore[arg-type] + + +@pytest.mark.parametrize("device", ["cpu", "metal"]) +def test_stabilizer_rank_strategies_take_only_the_auto_device(device): + with pytest.raises(ValueError, match="needs strategy='active_statevector'"): + Circuit(NOISY_T_STATE).compile_sampler(strategy="cat5", device=device) + + +def test_metal_device_raises_where_metal_cannot_run(monkeypatch): + from tsim.statevector.metal import runtime + + monkeypatch.setattr(runtime, "unavailable", lambda: "no Apple GPU") + with pytest.raises(RuntimeError, match="Cannot sample on Metal: no Apple GPU"): + Circuit(NOISY_T_STATE).compile_detector_sampler(device="metal") + with warnings.catch_warnings(): + warnings.simplefilter("error") + sampler = Circuit(NOISY_T_STATE).compile_detector_sampler(seed=1) + assert sampler.stats.device == "cpu"