diff --git a/.claude/skills/testing-tsim-correctness/SKILL.md b/.claude/skills/testing-tsim-correctness/SKILL.md index 0c6cf111..6dfce882 100644 --- a/.claude/skills/testing-tsim-correctness/SKILL.md +++ b/.claude/skills/testing-tsim-correctness/SKILL.md @@ -46,9 +46,9 @@ handful of T gates; its memory grows exponentially. test is a broken test. - Stim merges consecutive identical instructions into one with several targets, so write circuits that exercise both single- and multi-target paths. -- GPU tests: `test/conftest.py` keeps every test outside - `test/integration/test_metal_sampler.py` on the CPU. In the Metal tests compile the - GPU side with `device="metal"`, so a silent fallback cannot make the test pass on the +- GPU tests: `test/conftest.py` keeps every test outside `test/integration/test_metal_sampler.py` + and `test_cuda_sampler.py` on the CPU. In the GPU tests compile the GPU side with + `device="metal"` or `device="cuda"`, so a silent fallback cannot make the test pass on the CPU, and the reference with `device="cpu"`; compare them with `assert_counts_match` (GPU samples never equal CPU samples bit for bit). - Keep circuits small (a few qubits, tens of thousands of shots) so the suite stays diff --git a/AGENTS.md b/AGENTS.md index 7b7df498..c154abd4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -68,7 +68,7 @@ more than a handful of T gates: their compile memory grows exponentially. a reference. The reference-sample flags change only that. - **Seeds:** with `active_statevector` on the CPU, a seed's samples do not depend on the thread count or `batch_size`. Do not break this; noise blocks and per-shot draws are seeded - by call and shot index (`shots.py`). On a GPU (`statevector/metal/`) they depend on + by call and shot index (`shots.py`). On a GPU (`statevector/metal/`, `statevector/cuda/`) they depend on `batch_size`, as documented, and differ from the CPU's. - **Direct outputs stay direct:** outputs that depend on the noise alone never go through a state vector. Most detectors of QEC circuits are direct, which is why the detector sampler is diff --git a/CHANGELOG.md b/CHANGELOG.md index 2b26a13b..e2cdfe21 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- **NVIDIA GPU sampling.** With the new `cuda` extra + (`pip install "bloqade-tsim[cuda]"`, CUDA 13), Tsim samples on NVIDIA GPUs + through CuPy, as it does on Apple GPUs: every component must have at most 11 + active qubits, `device="auto"` picks the GPU when it can run the circuit, and + `device="cuda"` raises an error explaining why it can't. On an RTX 5090, + end-to-end cultivation at d=5 samples about 1.6 billion shots per second, + postselected, in any output layout. +- The `jax-cuda12` and `jax-cuda13` extras, for the stabilizer-rank strategies + on NVIDIA GPUs with JAX. The `cuda12` and `cuda13` extras of Tsim 1.0 install + the same packages. + ## [1.1.0] - 2026-11-12 Tsim 1.1 adds a GPU backend for Apple GPUs. For example, end-to-end diff --git a/README.md b/README.md index b9c694c0..168bdc86 100644 --- a/README.md +++ b/README.md @@ -24,9 +24,10 @@ uv add bloqade-tsim ``` The default sampling strategy runs on the CPU, and on the GPU of an Apple-silicon Mac (no extra -needed). To run the stabilizer-rank strategies on an NVIDIA GPU, install the JAX CUDA extra: +needed). To sample on an NVIDIA GPU, install the extra of your strategy: ```bash -pip install "bloqade-tsim[cuda13]" # or bloqade-tsim[cuda12] +pip install "bloqade-tsim[cuda]" # the default strategy, "active_statevector" (CUDA 13) +pip install "bloqade-tsim[jax-cuda13]" # the stabilizer-rank strategies, with JAX (or jax-cuda12) ``` ### Using Tsim with coding agents @@ -82,7 +83,7 @@ Choose a strategy with the `strategy` argument of `compile_sampler` and `compile | Strategy | Method | Runs on | Cost grows with | Use it for | |---|---|---|---|---| -| `"active_statevector"` (default) | Dense state vector over the active qubits of each component, as in [Clifft](https://arxiv.org/abs/2604.27058) and [SymFT](https://arxiv.org/abs/2607.28600) | CPU (numba, multi-threaded), or an Apple GPU (Metal) | Number of active qubits (at most 30) | Most circuits, especially with postselection | +| `"active_statevector"` (default) | Dense state vector over the active qubits of each component, as in [Clifft](https://arxiv.org/abs/2604.27058) and [SymFT](https://arxiv.org/abs/2607.28600) | CPU (numba, multi-threaded), or an Apple GPU (Metal) or NVIDIA GPU (CUDA) | Number of active qubits (at most 30) | Most circuits, especially with postselection | | `"cat5"`, `"bss"`, `"cutting"` | Stabilizer rank decomposition | CPU or GPU (JAX) | Number of non-Clifford gates | Few non-Clifford gates acting on many qubits; GPU throughput | ```python diff --git a/docs/index.mdx b/docs/index.mdx index e51a566e..cacac74b 100644 --- a/docs/index.mdx +++ b/docs/index.mdx @@ -40,10 +40,10 @@ uv add bloqade-tsim ``` The default sampling strategy runs on the CPU, and on the GPU of an Apple-silicon Mac (no extra -needed). To run the stabilizer-rank strategies on an NVIDIA GPU, install the JAX CUDA extra: +needed). To sample on an NVIDIA GPU with the default strategy, install the `cuda` extra: ```bash -uv add "bloqade-tsim[cuda13]" +uv add "bloqade-tsim[cuda]" ``` See [Installation](/install) for more options. @@ -55,7 +55,7 @@ Choose a strategy with the `strategy` argument of `compile_sampler` and `compile | Strategy | Method | Runs on | Cost grows with | Use it for | |---|---|---|---|---| -| `"active_statevector"` (default) | Dense state vector over the active qubits of each component, as in [Clifft](https://arxiv.org/abs/2604.27058) and [SymFT](https://arxiv.org/abs/2607.28600) | CPU (numba, multi-threaded), or an Apple GPU (Metal) | Number of active qubits (at most 30) | Most circuits, especially with postselection | +| `"active_statevector"` (default) | Dense state vector over the active qubits of each component, as in [Clifft](https://arxiv.org/abs/2604.27058) and [SymFT](https://arxiv.org/abs/2607.28600) | CPU (numba, multi-threaded), or an Apple GPU (Metal) or NVIDIA GPU (CUDA) | Number of active qubits (at most 30) | Most circuits, especially with postselection | | `"cat5"`, `"bss"`, `"cutting"` | Stabilizer rank decomposition | CPU or GPU (JAX) | Number of non-Clifford gates | Few non-Clifford gates acting on many qubits; GPU throughput | ```python diff --git a/docs/install.mdx b/docs/install.mdx index f8938550..c1350bca 100644 --- a/docs/install.mdx +++ b/docs/install.mdx @@ -13,55 +13,60 @@ We recommend using [`uv`](https://docs.astral.sh/uv/) for dependency management. uv add bloqade-tsim ``` -For GPU acceleration with CUDA: - -```bash -# For CUDA 13 -uv add "bloqade-tsim[cuda13]" - -# For CUDA 12 -uv add "bloqade-tsim[cuda12]" -``` - ## Using pip ```bash pip install bloqade-tsim ``` -For GPU acceleration with CUDA: +## GPUs + +The default strategy (`"active_statevector"`) samples on an Apple GPU through Metal, or on an +NVIDIA GPU through CUDA. On an Apple-silicon Mac, the plain install includes PyObjC's Metal +bindings for it. For NVIDIA GPUs, install an extra; each sampling strategy has its own: + +| Extra | Accelerates | Needs | +|---|---|---| +| `cuda` | `"active_statevector"` (the default), with CuPy | CUDA 13 (driver 580 or later) | +| `jax-cuda13` | `"cat5"`, `"bss"` and `"cutting"`, with JAX | CUDA 13 | +| `jax-cuda12` | `"cat5"`, `"bss"` and `"cutting"`, with JAX | CUDA 12 | ```bash -pip install "bloqade-tsim[cuda13]" +uv add "bloqade-tsim[cuda]" +# or, with both strategies on the GPU: +uv add "bloqade-tsim[cuda,jax-cuda13]" ``` -## Apple GPUs +The `cuda12` and `cuda13` extras of Tsim 1.0 still install the JAX extras. JAX starts only +when a stabilizer-rank strategy samples, and then reserves 75% of the GPU's memory; to sample +with both strategies in one process, set `XLA_PYTHON_CLIENT_PREALLOCATE=false` before starting +Python. -On an Apple-silicon Mac, the default strategy (`"active_statevector"`) samples on the GPU -through Metal; the plain install includes PyObjC's Metal bindings for it. Where it samples is -the `device` argument of `compile_sampler` and `compile_detector_sampler`: +Where `"active_statevector"` samples is the `device` argument of `compile_sampler` and +`compile_detector_sampler`: | `device` | Samples | | --- | --- | -| `"auto"` (default) | On the GPU when it can run the circuit, else on the CPU, without a warning | +| `"auto"` (default) | On a GPU when it can run the circuit, else on the CPU, without a warning | | `"cpu"` | On the CPU, in double precision | -| `"metal"` | On the GPU, or raises an error that says why it cannot | +| `"metal"` | On an Apple GPU, or raises an error that says why it cannot | +| `"cuda"` | On an NVIDIA GPU, or raises an error that says why it cannot | ```python -sampler = circuit.compile_detector_sampler(device="metal") +sampler = circuit.compile_detector_sampler(device="cuda") ``` The GPU runs circuits whose components have at most 11 active qubits. That includes Clifford circuits, whose detectors depend on the noise alone: the GPU samples the noise as well as the state vectors. Postselection runs on the GPU, so heavily postselected circuits gain the most. -A sampler's repr says where it samples (`samples on metal` or `samples on cpu`), and so does -`sampler.stats.device`. +A sampler's repr says where it samples (`samples on metal`, `samples on cuda` or +`samples on cpu`), and so does `sampler.stats.device`. The GPU computes outcome probabilities in float32, so a sample follows the CPU's distribution to within float32 rounding. Float32 cannot tell a measurement outcome of probability below about 1e-11 from rounding error, so the GPU never draws one; a circuit with a rotation small enough to make such an outcome (an angle within about 6e-5 π of a multiple of π/2) samples on -the CPU (or raises, with `device="metal"`). +the CPU (or raises, with `device="metal"` or `"cuda"`). A seed's samples on the GPU differ from the CPU's and depend on `batch_size`, so with `device="auto"` a seed reproduces its samples only on the same kind of machine. Use diff --git a/pyproject.toml b/pyproject.toml index e7b38007..0b823c05 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,12 @@ Repository = "https://github.com/QuEraComputing/tsim" Issues = "https://github.com/QuEraComputing/tsim/issues" [project.optional-dependencies] +# The default strategy, "active_statevector", on NVIDIA GPUs (CUDA 13). +cuda = ["cupy-cuda13x[ctk]>=14.0"] +# The stabilizer-rank strategies ("cat5", "bss", "cutting") on NVIDIA GPUs, with JAX. +jax-cuda12 = ["jax[cuda12]>=0.4.38"] +jax-cuda13 = ["jax[cuda13]>=0.7.0"] +# NOTE: the names of the JAX extras up to 1.0, kept so that they still install. cuda12 = ["jax[cuda12]>=0.4.38"] cuda13 = ["jax[cuda13]>=0.7.0"] diff --git a/skills/simulating-circuits-with-tsim/SKILL.md b/skills/simulating-circuits-with-tsim/SKILL.md index 466728bb..65baf352 100644 --- a/skills/simulating-circuits-with-tsim/SKILL.md +++ b/skills/simulating-circuits-with-tsim/SKILL.md @@ -197,23 +197,26 @@ yourself. default. Use `numba.set_num_threads(k)` before `sample` (or `NUMBA_NUM_THREADS=k` before importing Tsim) and leave cores free on a shared machine. Do not run several multi-threaded Tsim processes side by side; they oversubscribe the CPU. -- **Apple GPUs:** on an Apple-silicon Mac the default strategy samples on the GPU when every - component has at most 11 active qubits, often 5 to 10 times faster than all CPU cores. That +- **GPUs:** on an Apple-silicon Mac (Metal), or on an NVIDIA GPU with the `cuda` extra, the + default strategy samples on the GPU when every component has at most 11 active qubits, + often 5 to 10 times faster than all CPU cores. That includes Clifford circuits, whose detectors come from the noise alone: the GPU samples the noise too. It computes outcome probabilities in float32 and never draws an outcome below about 1e-11 (circuits with a rotation angle within about 6e-5 π of a multiple of π/2 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, + on the CPU in double precision; `"metal"` (Apple) and `"cuda"` (NVIDIA) raise with the + reason if the GPU cannot run the circuit (no such GPU, width above 11, a tiny rotation, out + of GPU memory). Check `sampler.stats.device` (or the repr, `samples on metal` or `cuda`) + after compiling, since `"auto"` does not warn. Pass `device="metal"` or `"cuda"` for long + runs that must not drop to the CPU unnoticed, and `device="cpu"` to reproduce samples across machines or to compare against the GPU. - **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"`). + parent used the GPU cannot use it (it samples on the CPU, or raises with `device="metal"` or + `"cuda"`); start workers with the `"spawn"` or `"forkserver"` method (macOS's default is + `"spawn"`). - **Seeds:** a `seed` reproduces samples for the same Tsim and dependency versions and device. On the CPU, thread count and `batch_size` do not affect the samples; on a GPU, `batch_size` does, and GPU samples differ from CPU samples, so with `device="auto"` a seed diff --git a/src/tsim/circuit.py b/src/tsim/circuit.py index 6caf204d..f2918747 100644 --- a/src/tsim/circuit.py +++ b/src/tsim/circuit.py @@ -966,24 +966,25 @@ def compile_sampler( depends on the free memory. On a GPU, they do depend on the batch size. device: Where "active_statevector" samples. "auto" (the default) - samples in float32 on an Apple GPU when there is one that can - run the circuit, and otherwise on the CPU, without a warning. - "cpu" samples on the CPU in float64, on + samples in float32 on an Apple GPU or an NVIDIA GPU when there + is one that can run the circuit, and otherwise on the CPU, + without a warning. "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 within about 6e-5 π of a multiple of - π/2, noise beyond the GPU's tables, or a call that runs out - of GPU memory. The stabilizer rank strategies take only "auto". + GPU and "cuda" on an NVIDIA GPU (with the ``cuda`` extra), or + raises: without one, for a component of more than 11 active + qubits, a rotation within about 6e-5 π of a multiple of π/2, + noise beyond the GPU's tables, or a call that runs out of GPU + memory. The stabilizer rank strategies take only "auto". Returns: A CompiledMeasurementSampler that can be used to sample measurements. Raises: ValueError: If ``device`` is unknown, is not "auto" under a - stabilizer-rank strategy, or is "metal" for a circuit that the - GPU executor cannot run. - RuntimeError: If ``device`` is "metal" and this process cannot use - an Apple GPU. + stabilizer-rank strategy, or is "metal" or "cuda" for a circuit + that the GPU executor cannot run. + RuntimeError: If ``device`` is "metal" or "cuda" and this process + cannot use such a GPU. """ from tsim.sampler import CompiledMeasurementSampler @@ -1021,24 +1022,25 @@ def compile_detector_sampler( settings, since the default batch size depends on the free memory. On a GPU, they do depend on the batch size. device: Where "active_statevector" samples. "auto" (the default) - samples in float32 on an Apple GPU when there is one that can - run the circuit, and otherwise on the CPU, without a warning. - "cpu" samples on the CPU in float64, on + samples in float32 on an Apple GPU or an NVIDIA GPU when there + is one that can run the circuit, and otherwise on the CPU, + without a warning. "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 within about 6e-5 π of a multiple of - π/2, noise beyond the GPU's tables, or a call that runs out - of GPU memory. The stabilizer rank strategies take only "auto". + GPU and "cuda" on an NVIDIA GPU (with the ``cuda`` extra), or + raises: without one, for a component of more than 11 active + qubits, a rotation within about 6e-5 π of a multiple of π/2, + noise beyond the GPU's tables, or a call that runs out of GPU + memory. The stabilizer rank strategies take only "auto". Returns: A CompiledDetectorSampler that can be used to sample detectors and observables. Raises: ValueError: If ``device`` is unknown, is not "auto" under a - stabilizer-rank strategy, or is "metal" for a circuit that the - GPU executor cannot run. - RuntimeError: If ``device`` is "metal" and this process cannot use - an Apple GPU. + stabilizer-rank strategy, or is "metal" or "cuda" for a circuit + that the GPU executor cannot run. + RuntimeError: If ``device`` is "metal" or "cuda" and this process + cannot use such a GPU. """ from tsim.sampler import CompiledDetectorSampler diff --git a/src/tsim/core/types.py b/src/tsim/core/types.py index d142170b..efa5533b 100644 --- a/src/tsim/core/types.py +++ b/src/tsim/core/types.py @@ -28,7 +28,7 @@ SamplingStrategy = Literal["active_statevector", "cat5", "bss", "cutting"] """``"active_statevector"``, or a stabilizer-rank decomposition strategy.""" -Device = Literal["auto", "cpu", "metal"] +Device = Literal["auto", "cpu", "metal", "cuda"] """Where ``"active_statevector"`` samples: ``"auto"`` picks a GPU when one can run it.""" diff --git a/src/tsim/sampler.py b/src/tsim/sampler.py index 09d08d38..c8b7690f 100644 --- a/src/tsim/sampler.py +++ b/src/tsim/sampler.py @@ -24,7 +24,7 @@ SamplingStrategy, ) from tsim.noise.channels import ChannelSampler -from tsim.statevector import metal +from tsim.statevector import cuda, metal from tsim.statevector.program import StatevectorComponent from tsim.statevector.shots import ShotSampler from tsim.utils.cuda_helpers import copy_d2h @@ -34,8 +34,12 @@ from jax import Array as PRNGKey from tsim.circuit import Circuit + from tsim.statevector.cuda.sampler import CudaShotSampler from tsim.statevector.metal.sampler import MetalShotSampler +_GPUS = {"metal": metal, "cuda": cuda} +"""The GPU backends of ``device``, in the order that ``device="auto"`` tries them.""" + # NOTE: postselection with a stabilizer-rank strategy screens shots for direct # discards in rounds of SCREEN_CHUNKS chunks of at most SCREEN_SHOTS shots. SCREEN_SHOTS = 2**14 @@ -351,25 +355,28 @@ def __init__( depends on the free memory. On a GPU, they do depend on the batch size. device: Where "active_statevector" samples. "auto" (the default) - samples in float32 on an Apple GPU when there is one that can - run the circuit, and otherwise on the CPU, without a warning. - "cpu" samples on the CPU in float64, on + samples in float32 on an Apple GPU or an NVIDIA GPU when there + is one that can run the circuit, and otherwise on the CPU, + without a warning. "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 within about 6e-5 π of a multiple of - π/2, noise beyond the GPU's tables, or a call that runs out - of GPU memory. The stabilizer rank strategies take only "auto". + GPU and "cuda" on an NVIDIA GPU (with the ``cuda`` extra), or + raises: without one, for a component of more than 11 active + qubits, a rotation within about 6e-5 π of a multiple of π/2, + noise beyond the GPU's tables, or a call that runs out of GPU + memory. The stabilizer rank strategies take only "auto". Raises: ValueError: If ``device`` is unknown, is not "auto" under a - stabilizer-rank strategy, or is "metal" for a circuit that the - GPU executor cannot run. - RuntimeError: If ``device`` is "metal" and this process cannot use - an Apple GPU. + stabilizer-rank strategy, or is "metal" or "cuda" for a circuit + that the GPU executor cannot run. + RuntimeError: If ``device`` is "metal" or "cuda" and this process + cannot use such a GPU. """ - if device not in ("auto", "cpu", "metal"): - raise ValueError(f"device must be 'auto', 'cpu' or 'metal', not {device!r}") + if device not in ("auto", "cpu", *_GPUS): + raise ValueError( + f"device must be 'auto', 'cpu', 'metal' or 'cuda', not {device!r}" + ) if device != "auto" and strategy != "active_statevector": raise ValueError( f"device={device!r} needs strategy='active_statevector'; the " @@ -398,7 +405,7 @@ def __init__( program = self._program # NOTE: programs of statevector components, and programs without components # under any strategy, are sampled as shot words; the others with JAX. - self._shots: ShotSampler | MetalShotSampler | None = None + self._shots: ShotSampler | MetalShotSampler | CudaShotSampler | None = None self._backend: _StabilizerRankBackend | None = None if strategy == "active_statevector" or not program.components: host = ShotSampler( @@ -411,13 +418,17 @@ def __init__( # NOTE: as JAX does for the stabilizer-rank strategies, "auto" uses a # GPU whenever there is one that can run the program. if mode == "sequential" and device != "cpu": - self._shots = metal.shot_sampler( - host, - self._channel_sampler, - self._num_detectors, - gpu_seed, - required=device == "metal", - ) + # NOTE: "auto" takes the first GPU that runs the program. + for name in _GPUS if device == "auto" else (device,): + self._shots = _GPUS[name].shot_sampler( + host, + self._channel_sampler, + self._num_detectors, + gpu_seed, + required=device == name, + ) + if self._shots is not host: + break else: self._backend = _StabilizerRankBackend(program, key_seed) @@ -767,8 +778,8 @@ def stats(self) -> SamplerStats: by which the default batch sizes are chosen. - ``compiled_bytes`` (``int``): memory of the compiled components. - ``device`` (``str``): where ``"active_statevector"`` samples: - ``"metal"`` on an Apple GPU, else ``"cpu"`` (a stabilizer-rank - strategy runs wherever JAX does). With ``device="auto"`` it is where + ``"metal"`` on an Apple GPU, ``"cuda"`` on an NVIDIA GPU, else + ``"cpu"`` (a stabilizer-rank strategy runs wherever JAX does). With ``device="auto"`` it is where the sampler ended up: ``"cpu"`` after its GPU kernel failed to compile or in a process forked after its parent used the GPU. A call that runs out of GPU memory samples on the CPU without changing it. @@ -908,21 +919,22 @@ def __init__( depends on the free memory. On a GPU, they do depend on the batch size. device: Where "active_statevector" samples. "auto" (the default) - samples in float32 on an Apple GPU when there is one that can - run the circuit, and otherwise on the CPU, without a warning. - "cpu" samples on the CPU in float64, on + samples in float32 on an Apple GPU or an NVIDIA GPU when there + is one that can run the circuit, and otherwise on the CPU, + without a warning. "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 within about 6e-5 π of a multiple of - π/2, noise beyond the GPU's tables, or a call that runs out - of GPU memory. The stabilizer rank strategies take only "auto". + GPU and "cuda" on an NVIDIA GPU (with the ``cuda`` extra), or + raises: without one, for a component of more than 11 active + qubits, a rotation within about 6e-5 π of a multiple of π/2, + noise beyond the GPU's tables, or a call that runs out of GPU + memory. The stabilizer rank strategies take only "auto". Raises: ValueError: If ``device`` is unknown, is not "auto" under a - stabilizer-rank strategy, or is "metal" for a circuit that the - GPU executor cannot run. - RuntimeError: If ``device`` is "metal" and this process cannot use - an Apple GPU. + stabilizer-rank strategy, or is "metal" or "cuda" for a circuit + that the GPU executor cannot run. + RuntimeError: If ``device`` is "metal" or "cuda" and this process + cannot use such a GPU. """ super().__init__( @@ -1027,21 +1039,22 @@ def __init__( settings, since the default batch size depends on the free memory. On a GPU, they do depend on the batch size. device: Where "active_statevector" samples. "auto" (the default) - samples in float32 on an Apple GPU when there is one that can - run the circuit, and otherwise on the CPU, without a warning. - "cpu" samples on the CPU in float64, on + samples in float32 on an Apple GPU or an NVIDIA GPU when there + is one that can run the circuit, and otherwise on the CPU, + without a warning. "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 within about 6e-5 π of a multiple of - π/2, noise beyond the GPU's tables, or a call that runs out - of GPU memory. The stabilizer rank strategies take only "auto". + GPU and "cuda" on an NVIDIA GPU (with the ``cuda`` extra), or + raises: without one, for a component of more than 11 active + qubits, a rotation within about 6e-5 π of a multiple of π/2, + noise beyond the GPU's tables, or a call that runs out of GPU + memory. The stabilizer rank strategies take only "auto". Raises: ValueError: If ``device`` is unknown, is not "auto" under a - stabilizer-rank strategy, or is "metal" for a circuit that the - GPU executor cannot run. - RuntimeError: If ``device`` is "metal" and this process cannot use - an Apple GPU. + stabilizer-rank strategy, or is "metal" or "cuda" for a circuit + that the GPU executor cannot run. + RuntimeError: If ``device`` is "metal" or "cuda" and this process + cannot use such a GPU. """ super().__init__( diff --git a/src/tsim/statevector/cuda/__init__.py b/src/tsim/statevector/cuda/__init__.py new file mode 100644 index 00000000..a9d1ed7d --- /dev/null +++ b/src/tsim/statevector/cuda/__init__.py @@ -0,0 +1,80 @@ +"""The statevector executor on a CUDA GPU. + +The noise sampler (``noise``) and every component's tape (``codegen``, written +by ``tsim.statevector.gpu.tape``) run as CUDA kernels that CuPy compiles with +NVRTC (``runtime``); ``CudaShotSampler`` (``sampler``) replaces ``ShotSampler`` +behind the compiled samplers of ``device="auto"`` or ``"cuda"`` whenever CuPy +sees a GPU and every component fits. Amplitudes are complex64. It needs CuPy +and an NVIDIA GPU (``pip install bloqade-tsim[cuda]``). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from tsim.statevector.cuda import runtime +from tsim.statevector.gpu.tape import unsupported + +if TYPE_CHECKING: + import numpy as np + + from tsim.noise.channels import ChannelSampler + from tsim.statevector.cuda.sampler import CudaShotSampler + from tsim.statevector.shots import ShotSampler + + +def shot_sampler( + host: ShotSampler, + channels: ChannelSampler, + num_detectors: int, + seed: np.random.SeedSequence, + *, + required: bool, +) -> ShotSampler | CudaShotSampler: + """Return a CUDA shot sampler of ``host``'s program, or ``host`` where CUDA cannot run it. + + Without ``required`` (``device="auto"``) it returns ``host`` silently; with it + (``device="cuda"``) it raises instead. A program without shot words (one + with no outputs; direct outputs are shot-word bits too) has nothing to + sample on a GPU and keeps ``host`` either way. + + Raises: + RuntimeError: With ``required``, if this process cannot use a CUDA GPU: + no CuPy, no device, or forked after its parent used the GPU. + ValueError: With ``required``, if the GPU executor cannot run the program: + a component wider than ``MAX_WIDTH``, a rotation by an angle too small + for float32 amplitudes (``RESOLVABLE``), or noise beyond + the GPU tables' limits. + MemoryError: With ``required``, if the GPU runs out of memory. + + """ + reason = runtime.unavailable() + if reason is not None: + if required: + raise RuntimeError(f"Cannot sample on CUDA: {reason}.") + return host + if host.layout.num_words == 0: + return host + reason = unsupported(host.layout) + if reason is not None: + if required: + raise ValueError( + f"Cannot sample on CUDA: {reason}; sample with device='cpu', or " + "with device='auto', which falls back to the CPU." + ) + return host + from tsim.statevector.cuda.sampler import CudaShotSampler + + try: + return CudaShotSampler(host, channels, num_detectors, seed, required=required) + except ValueError as error: + # NOTE: noise beyond the GPU tables' limits (channels of more than 127 + # outcomes, outcomes that flip more than 255 words). + if required: + raise ValueError(f"Cannot sample on CUDA: {error}.") from error + except Exception as error: + if not runtime.out_of_memory(error): + raise + if required: + raise MemoryError(f"Cannot sample on CUDA: {error}") from error + return host diff --git a/src/tsim/statevector/cuda/codegen.py b/src/tsim/statevector/cuda/codegen.py new file mode 100644 index 00000000..4a65161e --- /dev/null +++ b/src/tsim/statevector/cuda/codegen.py @@ -0,0 +1,79 @@ +"""CUDA source of a component's tape (see ``tsim.statevector.gpu.tape``). + +Partner amplitudes on other threads come from ``__shfl_xor_sync`` among the +shot's group of lanes (``gmask``), and a pruned shot jumps to the end. Near +the boundary ``draw_one`` decides an outcome in double precision. +""" + +from __future__ import annotations + +from tsim.statevector.gpu import tape + +CUDA = tape.Dialect( + inline="__device__ __forceinline__", + noinline="__device__ __noinline__", + private="", + deviation="unsigned*", + constant="__device__ const", + float2="make_float2", + rsqrt="rsqrtf", + fmax="fmaxf", + u64="ULL", + shuffle2="shfl2({value}, {lanes}, gmask)", + shuffle1="__shfl_xor_sync(gmask, {value}, {lanes})", + lane_mask=True, + goto=True, +) + +PRELUDE = r""" +typedef unsigned long long u64; +#define GOLDEN 0x9E3779B97F4A7C15ULL + +__device__ __forceinline__ u64 mix64(u64 z) { + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL; + return z ^ (z >> 31); +} +// The CPU executor's draw: splitmix64 at index shot * slots + slot of stream seed. +__device__ __forceinline__ u64 draw(u64 seed, u64 shot, u64 slots, u64 slot) { + return mix64(seed + (shot * slots + slot) * GOLDEN + GOLDEN); +} +__device__ __forceinline__ int par(u64 v) { return __popcll(v) & 1; } +// draw_one in double precision, from the rarer outcome's share, which float32 +// keeps to full relative precision: a rare 0 takes the top of the range, as on +// the CPU (``shared._cut``). Not inlined, as it is rarely taken. +__device__ __noinline__ int draw_one_exact(u64 r, float w0, float w1, float total) { + int rare0 = w0 < w1; + u64 below = (u64)((double)(rare0 ? w0 : w1) / (double)total * 18446744073709551616.0); + return ((rare0 ? ~r : r) < below) ^ rare0; +} +// Whether the 64-bit draw r gives outcome 1, of weight w1 against w0. float32 +// decides it unless r / 2^64 lies within 2^-21 of the share w1 / total, which +// covers both roundings; there draw_one_exact does. +__device__ __forceinline__ int draw_one(u64 r, float w0, float w1, float total) { + float p = w1 / total, x = __ull2float_rz(r) * 0x1p-64f; + if (x < p * (1.f - 0x1p-21f)) return 1; + if (x > p * (1.f + 0x1p-21f)) return 0; + return draw_one_exact(r, w0, w1, total); +} +// A uniform in (0, 1] from the high 24 bits, so that logf never sees 0. +__device__ __forceinline__ float unit(u64 r) { return ((r >> 40) + 1) * (1.0f / 16777216.0f); } +__device__ __forceinline__ float2 cmul(float2 p, float2 q) { + return make_float2(p.x * q.x - p.y * q.y, p.x * q.y + p.y * q.x); +} +// c a + s q and c a + i s q +__device__ __forceinline__ float2 cmac(float2 a, float c, float s, float2 q) { + return make_float2(fmaf(c, a.x, s * q.x), fmaf(c, a.y, s * q.y)); +} +__device__ __forceinline__ float2 cmac_i(float2 a, float c, float s, float2 q) { + return make_float2(fmaf(c, a.x, -s * q.y), fmaf(c, a.y, s * q.x)); +} +__device__ __forceinline__ float2 shfl2(float2 v, int lane_xor, unsigned mask) { + return make_float2(__shfl_xor_sync(mask, v.x, lane_xor), __shfl_xor_sync(mask, v.y, lane_xor)); +} +// The largest normalization deviation, as the bits of a non-negative float. +__device__ __forceinline__ void deviate(unsigned* deviation, float total) { + float d = fabsf(total - 1.f); + if (!(d <= 1e-3f)) atomicMax(deviation, d == d ? __float_as_uint(d) : 0x7f800000u); +} +""" diff --git a/src/tsim/statevector/cuda/kernels.py b/src/tsim/statevector/cuda/kernels.py new file mode 100644 index 00000000..ee71e579 --- /dev/null +++ b/src/tsim/statevector/cuda/kernels.py @@ -0,0 +1,246 @@ +"""The CUDA sampler's kernels other than the noise ones, and their source. + +``tape`` runs the components' tapes on the queued shots; ``count_bits`` and +``list_bits`` list the accepted shots from the bitmap; ``gather_columns`` and +``gather_float`` gather rows and expectations for the copy to the host +(``gather_bytes`` the rows too wide for ``gather_columns``); +``select_rows`` picks the requested outputs out of them, bit-packed or as +booleans; ``publish`` stores a batch's counts into pinned host memory. +""" + +from __future__ import annotations + +from tsim.statevector.cuda.codegen import PRELUDE +from tsim.statevector.cuda.noise import source as noise_source +from tsim.statevector.gpu.noise import NoiseTables +from tsim.statevector.gpu.tape import TapeKernel +from tsim.statevector.shots import ShotLayout + +TAPE_THREADS = 128 +"""Threads per block of the tape kernel.""" + +COPY_THREADS = 256 +"""Threads per block of the kernels that gather rows and expectations.""" + +SCAN_THREADS = 1024 +"""Threads per block of ``count_bits`` and ``list_bits``: one per bitmap word.""" + +KERNEL_SHARED_BYTES = 48 * 1024 +"""Shared memory that a kernel may use without opting in, static and dynamic.""" + + +def gather_shots(layout: ShotLayout, block_shots: int) -> int: + """Return the rows that a ``gather_columns`` block packs, or ``0`` if not even 32 fit. + + A block stages its rows' output words and packed bytes in shared memory: at + most ``block_shots`` rows, a multiple of 32. Rows of which 32 do not fit + take ``gather_bytes`` instead. + """ + row = 8 * layout.output_words + -(-layout.num_outputs // 8) + return min(block_shots, KERNEL_SHARED_BYTES // max(row, 1) // 32 * 32) + + +TEMPLATE = r""" +extern "C" __global__ void __launch_bounds__({block}) tape( + u64* __restrict__ out, long long first_shot, u64 seed, + const unsigned* __restrict__ queue, const unsigned* __restrict__ count, + unsigned* __restrict__ survivors, unsigned* __restrict__ acc_count, + float* __restrict__ expectations, unsigned* __restrict__ deviation) +{{ + int t = threadIdx.x & {T1}; + unsigned gmask = {gmask}; +{seeds} + long long groups = ((long long)gridDim.x * blockDim.x) >> {tb}; + long long n = *count; + // NOTE: whole warps leave the loop together, so the shuffles keep their groups. + long long rounds = (n + groups - 1) / groups; + long long g0 = ((long long)blockIdx.x * blockDim.x + threadIdx.x) >> {tb}; + for (long long r = 0; r < rounds; ++r) {{ + long long q = g0 + r * groups; + bool valid = q < n; + u64 row_index = queue[valid ? q : n - 1]; + u64 shot = first_shot + row_index; + u64* row = out + row_index * {NW}; + // NOTE: the words that no component reads are dead loads, dropped. + u64 w[{NW}]; + #pragma unroll + for (int i = 0; i < {NW}; ++i) w[i] = row[i]; + float ev[{NO}]; + int grave = 0x7fffffff; +{calls} + if (valid && t == 0) {{ +{stores} + if (expectations != nullptr) {{ + #pragma unroll + for (int i = 0; i < {NO}; ++i) expectations[row_index * {NO} + i] = ev[i]; + }} + }} + if (survivors != nullptr) {{ + // NOTE: the postselected direct outputs passed the screen already; + // the shots that fail now leave the screen's bitmap. + u64 hit = {hits}; + if (valid && t == 0 && hit) atomicAnd(&survivors[row_index >> 5], ~(1u << (row_index & 31))); + unsigned ballot = __ballot_sync(0xffffffffu, valid && t == 0 && hit == 0); + if (ballot && (threadIdx.x & 31) == __ffs(ballot) - 1) atomicAdd(acc_count, __popc(ballot)); + }} + }} +}} + +// The shots of the ``words`` words of ``bits`` set, in order, as indices into +// ``index``: ``count_bits`` counts each block's {SCAN} words, ``list_bits`` lists +// a block's shots after those of the blocks before it. +__device__ __forceinline__ unsigned block_scan(unsigned v, unsigned* total) {{ + __shared__ unsigned warps[32]; + int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + unsigned x = v; + #pragma unroll + for (int o = 1; o < 32; o <<= 1) {{ unsigned y = __shfl_up_sync(0xffffffffu, x, o); if (lane >= o) x += y; }} + if (lane == 31) warps[warp] = x; + __syncthreads(); + if (warp == 0) {{ + unsigned w = lane < (blockDim.x >> 5) ? warps[lane] : 0; + #pragma unroll + for (int o = 1; o < 32; o <<= 1) {{ unsigned y = __shfl_up_sync(0xffffffffu, w, o); if (lane >= o) w += y; }} + warps[lane] = w; + }} + __syncthreads(); + unsigned before = (warp ? warps[warp - 1] : 0) + x - v; + *total = warps[31]; + __syncthreads(); + return before; +}} + +extern "C" __global__ void __launch_bounds__({SCAN}) count_bits( + const unsigned* __restrict__ bits, long long words, unsigned* __restrict__ totals) {{ + long long i = (long long)blockIdx.x * {SCAN} + threadIdx.x; + unsigned total; + block_scan(i < words ? __popc(bits[i]) : 0, &total); + if (threadIdx.x == 0) totals[blockIdx.x] = total; +}} + +extern "C" __global__ void __launch_bounds__({SCAN}) list_bits( + const unsigned* __restrict__ bits, long long words, const unsigned* __restrict__ totals, + unsigned* __restrict__ index) {{ + unsigned earlier = 0, total; + for (int b = threadIdx.x; b < blockIdx.x; b += blockDim.x) earlier += totals[b]; + block_scan(earlier, &earlier); + long long i = (long long)blockIdx.x * {SCAN} + threadIdx.x; + unsigned w = i < words ? bits[i] : 0; + unsigned at = earlier + block_scan(__popc(w), &total); + for (; w; w &= w - 1) index[at++] = (unsigned)(32 * i + __ffs(w) - 1); +}} + +// The outputs of rows ``index[j]`` (or every row), bit-packed densely +// (``pack_rows``): a block packs {GS} rows' output words. +extern "C" __global__ void __launch_bounds__({COPY}) gather_columns( + const u64* __restrict__ rows, long long n, const unsigned* __restrict__ index, + unsigned char* __restrict__ dst) {{ + __shared__ u64 sw[{OW} * {GS}]; + extern __shared__ __align__(16) unsigned char stage[]; + long long base = (long long)blockIdx.x * {GS}; + int ns = (int)min((long long){GS}, n - base); + for (int i = threadIdx.x; i < ns * {OW}; i += blockDim.x) {{ + int s = i / {OW}, w = i - s * {OW}; + long long r = index != nullptr ? (long long)index[base + s] : base + s; + sw[w * {GS} + s] = rows[r * {NW} + w]; + }} + __syncthreads(); + pack_rows<{GS}>(sw, ns, stage, dst + base * {NB}); +}} + +// As ``gather_columns``, for rows too wide to stage in shared memory: a thread +// per output byte, read from device memory. +extern "C" __global__ void gather_bytes( + const u64* __restrict__ rows, long long n, const unsigned* __restrict__ index, + unsigned char* __restrict__ dst) {{ + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n * {NB}; + i += (long long)gridDim.x * blockDim.x) {{ + long long j = i / {NB}; + int b = (int)(i - j * {NB}); + long long r = index != nullptr ? (long long)index[j] : j; + unsigned v = (unsigned)(rows[r * {NW} + (b >> 3)] >> (8 * (b & 7))) & 0xFF; + dst[i] = (unsigned char)(b == {NB} - 1 ? v & {LAST_BYTE} : v); + }} +}} + +// Outputs ``columns[k]`` (``k < count``; output ``k`` without ``columns``) of +// ``n`` packed rows of {NB} bytes, into rows of ``width`` bytes: one byte per +// output with ``bools``, else bit-packed with zero bits past them. A thread +// writes one byte. +extern "C" __global__ void select_rows( + const unsigned char* __restrict__ packed, long long n, const unsigned* __restrict__ columns, + int count, int width, int bools, unsigned char* __restrict__ dst) {{ + int per = bools ? 1 : 8; + for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; i < n * width; + i += (long long)gridDim.x * blockDim.x) {{ + long long r = i / width; + int k = (int)(i - r * width) * per; + const unsigned char* row = packed + r * {NB}; + unsigned v = 0; + for (int b = 0; b < per && k + b < count; ++b) {{ + unsigned c = columns != nullptr ? columns[k + b] : (unsigned)(k + b); + v |= ((row[c >> 3] >> (c & 7)) & 1u) << b; + }} + dst[i] = (unsigned char)v; + }} +}} + +// The batch's counts, stored into pinned host memory: a read that does not +// queue behind the copies to the host. +extern "C" __global__ void publish(const unsigned* __restrict__ counts, volatile unsigned* host) {{ + if (threadIdx.x < 3) host[threadIdx.x] = counts[threadIdx.x]; + __threadfence_system(); +}} + +// Observable expectations of rows ``index[j]``, densely. +extern "C" __global__ void gather_float(const float* __restrict__ values, long long n, int width, + const unsigned* __restrict__ index, float* __restrict__ dst) {{ + for (long long k = (long long)blockIdx.x * blockDim.x + threadIdx.x; k < n * width; + k += (long long)gridDim.x * blockDim.x) {{ + long long j = k / width; int c = (int)(k - j * width); + dst[k] = values[(long long)index[j] * width + c]; + }} +}} +""" + + +def source( + layout: ShotLayout, + tables: NoiseTables, + parts: TapeKernel, + tb: int, + num_observables: int, +) -> str: + """Return the CUDA source of every kernel, for tapes of ``2^tb`` threads a shot.""" + T = 1 << tb + # NOTE: a shot's lanes within its warp. + gmask = ( + "0xffffffffu" + if T == 32 + else f"((1u << {T}) - 1) << (threadIdx.x & 31 & ~{T - 1})" + ) + return ( + PRELUDE + + noise_source(tables, layout.output_words, layout.num_outputs) + + parts.functions + + TEMPLATE.format( + block=TAPE_THREADS, + SCAN=SCAN_THREADS, + COPY=COPY_THREADS, + T1=T - 1, + tb=tb, + gmask=gmask, + NW=layout.num_words, + OW=layout.output_words, + NO=num_observables, + # NOTE: gather_columns goes unused (and small) where 32 rows do not fit. + GS=max(gather_shots(layout, tables.block_shots), 1), + LAST_BYTE=f"0x{(1 << (layout.num_outputs % 8 or 8)) - 1:02X}u", + # NOTE: the rows keep every output, the postselected ones too. + NB=max(-(-layout.num_outputs // 8), 1), + calls=parts.calls, + seeds=parts.seeds, + stores=parts.stores, + hits=parts.hits, + ) + ) diff --git a/src/tsim/statevector/cuda/noise.py b/src/tsim/statevector/cuda/noise.py new file mode 100644 index 00000000..3c6d9cfe --- /dev/null +++ b/src/tsim/statevector/cuda/noise.py @@ -0,0 +1,341 @@ +"""Noise sampling into shot words on a CUDA GPU. + +The kernels read the tables of ``tsim.statevector.gpu.noise`` (see there): a +block owns ``S`` consecutive shots and keeps their words in shared memory, +word major, walks segments of the channel groups with geometric skips, +records the faults as events in shared memory and then XORs their outcomes' +flips into the shots' words. Recording first keeps the warps converged while +the faults are applied, and many of their table reads in flight. Rows too wide +for shared memory keep their words in device memory, where the faults are XORed +in with device atomics. + +A block then screens its shots (constants, postselected direct detectors), +writes their rows (coalesced; only the survivors' unless every row is asked +for), appends the surviving shots to a queue for the tape kernel and marks +them in a bitmap of the batch's shots. Under postselection, ``noise_kept`` +draws the noise that cannot change whether a shot is kept +(``gpu.noise.deferred_channels``) on the kept shots only, after the tapes. +""" + +from __future__ import annotations + +from tsim.noise.channels import ChannelSampler, WordTables +from tsim.statevector.gpu.noise import ( + NoiseTables, + noise_tables as _noise_tables, + words_fit, +) + +THREADS = 256 +"""Threads per block of the noise kernels.""" + +SHARED_BYTES = 46 * 1024 +"""Shared memory of a block's words and events (the events' share is set per launch).""" + +_NOISE_SIGNATURE = r""" +extern "C" __global__ void __launch_bounds__({NT}) noise( + u64* __restrict__ out, long long shots, u64 seed, int num_segments, int max_events, + const int* __restrict__ seg_first, const int* __restrict__ seg_outcomes, + const int* __restrict__ seg_begin, const int* __restrict__ seg_end, + const float* __restrict__ seg_log_keep, const Record* __restrict__ records, + const Flip* __restrict__ overflow, + const u64* __restrict__ constant, const u64* __restrict__ screen, + unsigned* __restrict__ queue, unsigned* __restrict__ count, unsigned* __restrict__ survivors) +""" + +# NOTE: the noise kernel with the block's words in shared memory, word major. +NOISE_SHARED = _NOISE_SIGNATURE + r"""{{ + __shared__ u64 sw[{NW} * {S}]; + extern __shared__ uint2 events[]; + __shared__ int num_events; + __shared__ bool passed[{S}]; + long long base = (long long)blockIdx.x * {S}; + int ns = (int)min((long long){S}, shots - base); + for (int i = threadIdx.x; i < {NW} * {S}; i += blockDim.x) sw[i] = 0; + if (threadIdx.x == 0) num_events = 0; + __syncthreads(); + add_faults<{S}, 1>(sw, events, &num_events, max_events, ns, block_stream(seed), num_segments, + seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep, records, overflow); + screen_shots<{S}, 1>(sw, ns, base, constant, screen, passed, queue, count, survivors); + __syncthreads(); + // NOTE: only the shots that pass the screen are read again. + u64* rows = out + base * {NW}; + for (int i = threadIdx.x; i < ns * {NW}; i += blockDim.x) {{ + int s = i / {NW}, word = i - s * {NW}; + if (passed[s]) rows[i] = sw[word * {S} + s]; + }} +}} +""" + +# NOTE: the noise kernel for rows too wide for shared memory: the faults are XORed +# into the rows in device memory, and the events have all of the shared memory. +NOISE_DEVICE = _NOISE_SIGNATURE + r"""{{ + extern __shared__ uint2 events[]; + __shared__ int num_events; + __shared__ bool passed[{S}]; + long long base = (long long)blockIdx.x * {S}; + int ns = (int)min((long long){S}, shots - base); + u64* rows = out + base * {NW}; + for (int i = threadIdx.x; i < ns * {NW}; i += blockDim.x) rows[i] = 0; + if (threadIdx.x == 0) num_events = 0; + __syncthreads(); + add_faults<1, {NW}>(rows, events, &num_events, max_events, ns, block_stream(seed), num_segments, + seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep, records, overflow); + screen_shots<1, {NW}>(rows, ns, base, constant, screen, passed, queue, count, survivors); +}} +""" + +NOISE_KEPT = r""" +// The deferred channels' faults on the ``n`` kept shots ``index``, then their +// returned outputs bit-packed (``pack_rows``). The deferred channels flip output +// words only, so a block loads just those. +extern "C" __global__ void __launch_bounds__({NT}) noise_kept( + const u64* __restrict__ out, long long n, const unsigned* __restrict__ index, u64 seed, + int num_segments, int max_events, + const int* __restrict__ seg_first, const int* __restrict__ seg_outcomes, + const int* __restrict__ seg_begin, const int* __restrict__ seg_end, + const float* __restrict__ seg_log_keep, const Record* __restrict__ records, + const Flip* __restrict__ overflow, + unsigned char* __restrict__ dst) +{{ + __shared__ u64 sw[{OW} * {S}]; + extern __shared__ __align__(16) uint2 events[]; + __shared__ int num_events; + long long base = (long long)blockIdx.x * {S}; + int ns = (int)min((long long){S}, n - base); + for (int i = threadIdx.x; i < ns * {OW}; i += blockDim.x) {{ + int s = i / {OW}, word = i - s * {OW}; + sw[word * {S} + s] = out[(long long)index[base + s] * {NW} + word]; + }} + if (threadIdx.x == 0) num_events = 0; + __syncthreads(); + add_faults<{S}, 1>(sw, events, &num_events, max_events, ns, block_stream(seed), num_segments, + seg_first, seg_outcomes, seg_begin, seg_end, seg_log_keep, records, overflow); + // NOTE: the events are applied, so their buffer stages the packed rows. + pack_rows<{S}>(sw, ns, (unsigned char*)events, dst + base * {NB}); +}} +""" + +SOURCE = r""" +struct __align__(32) Record {{ unsigned alias; unsigned words; unsigned extra; unsigned keep; u64 m0; u64 m1; }}; +struct __align__(16) Flip {{ unsigned word; unsigned pad; u64 mask; }}; + +// The fault drawn at record ``rec`` (bits 0-23 of ``draw`` 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]``. +// Whether a fault whose fraction's high bits tie with its record's keep +// threshold takes the alias: its low 16 bits, from a hash of the fault, are at +// least the threshold's. NOTE: rare, so out of line, which keeps the hot loop small. +__device__ __noinline__ bool alias_on_tie(u64 stream, unsigned rec, int row, unsigned keep) +{{ + return (unsigned)(mix64(stream ^ ((u64)(rec << 8 | row) * GOLDEN)) & 0xFFFF) >= keep; +}} + +template +__device__ __forceinline__ void apply_fault(u64* sw, int row, unsigned rec, unsigned draw, + u64 stream, const Record* __restrict__ records, const Flip* __restrict__ overflow) +{{ + Record r = records[rec]; + unsigned high = draw & 0xFFFFFF, keep = r.alias & 0x1FFFFFF; + if (high > keep || (high == keep && alias_on_tie(stream, rec, row, r.keep))) + r = records[rec - (draw >> 24) + (r.alias >> 25)]; + int flips = r.extra & 255; + if (flips > 0) atomicXor(&sw[(r.words & 0xFFFF) * WS + row * SS], r.m0); + if (flips > 1) atomicXor(&sw[(r.words >> 16) * WS + row * SS], r.m1); + for (int j = 2; j < flips; ++j) {{ + Flip f = overflow[(r.extra >> 8) + j - 2]; + atomicXor(&sw[f.word * WS + row * SS], f.mask); + }} +}} + +// The outputs of a block's ``ns`` rows ``sw[w * S + s]``, bit-packed into +// {NB} bytes per row (the bits past them cleared), contiguously at ``dst``. +// Consecutive threads pack one word of consecutive rows (distinct banks) into +// ``stage``, from which the block's rows go out 16 bytes at a time. +template +__device__ __forceinline__ void pack_rows(const u64* sw, int ns, unsigned char* stage, + unsigned char* __restrict__ dst) +{{ + for (int i = threadIdx.x; i < {NK} * S; i += blockDim.x) {{ + int k = i / S, s = i - k * S; + if (s >= ns) continue; + u64 v = sw[k * S + s]; + if (k == {NK} - 1) v &= {LAST_WORD}; + unsigned char* at = stage + s * {NB} + 8 * k; + #pragma unroll + for (int j = 0; j < 8; ++j) if (8 * k + j < {NB}) at[j] = (unsigned char)(v >> (8 * j)); + }} + __syncthreads(); + int total = ns * {NB}; + for (int i = threadIdx.x; i < total / 16; i += blockDim.x) ((uint4*)dst)[i] = ((const uint4*)stage)[i]; + for (int i = total / 16 * 16 + threadIdx.x; i < total; i += blockDim.x) dst[i] = stage[i]; +}} + +// The stream of this block. NOTE: hashed, so that a segment's stream +// ``mix64(stream ^ sg * GOLDEN)`` cannot equal another block's: with ``seed ^ bid * C`` +// alone, two (block, segment) pairs with equal XORs would share their draws. +__device__ __forceinline__ u64 block_stream(u64 seed) {{ + return mix64(seed ^ ((u64)blockIdx.x * 0xD1B54A32D192ED03ULL)); +}} + +// One segment's faults. With ``FINE``, for segments of fewer than 2^-12 expected +// faults, a draw within 2^-12 of u = 1 takes 24 more bits, which ``unit`` alone +// cannot resolve: it would fire such a segment at least 2^-24 of the time. +template +__device__ __forceinline__ void add_segment_faults(u64* sw, uint2* events, int* num_events, + int max_events, int ns, u64 stream, u64 state, float lk, int first, int outcomes, int idx, + int end, const Record* __restrict__ records, const Flip* __restrict__ overflow) +{{ + while (true) {{ + state += GOLDEN; + u64 r = mix64(state); + float skip; + if (!FINE) skip = lk < -3.0e38f ? 0.f : floorf(__logf(unit(r)) / lk); + else if ((~r >> 40) >= 4096) skip = floorf(__logf(unit(r)) / lk); + else {{ + // NOTE: 1 - u = (j + w) 2^-24, with w from a fresh draw. + float w = ((mix64(state ^ 0x5851F42D4C957F2DULL) >> 40) + 1) * 0x1p-24f; + skip = floorf(log1pf(-((float)(~r >> 40) + w) * 0x1p-24f) / lk); + }} + if (skip >= (float)(end - 1 - idx)) break; + idx += 1 + (int)skip; + int ch = idx / {S}, row = idx - ch * {S}; + if (row >= ns) continue; // the last block's missing shots + // NOTE: the skip reads the high 24 bits of r, the outcome its low 40: a + // slot, and the fraction of the slot below. + u64 x = (r & 0xFFFFFFFFFFULL) * (unsigned)outcomes; + unsigned slot = (unsigned)(x >> 40); + unsigned rec = (unsigned)(first + ch * outcomes) + slot; + unsigned draw = (unsigned)((x >> 16) & 0xFFFFFF) | (slot << 24); + int e = atomicAdd(num_events, 1); + if (e < max_events) events[e] = make_uint2((unsigned)row | (rec << 8), draw); + else apply_fault(sw, row, rec, draw, stream, records, overflow); + }} +}} + +template +__device__ __noinline__ void add_fine_segment_faults(u64* sw, uint2* events, int* num_events, + int max_events, int ns, u64 stream, u64 state, float lk, int first, int outcomes, int idx, + int end, const Record* __restrict__ records, const Flip* __restrict__ overflow) +{{ + add_segment_faults(sw, events, num_events, max_events, ns, 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``). +template +__device__ __forceinline__ void add_faults(u64* sw, uint2* events, int* num_events, int max_events, + int ns, u64 stream, int num_segments, + const int* __restrict__ seg_first, const int* __restrict__ seg_outcomes, + const int* __restrict__ seg_begin, const int* __restrict__ seg_end, + const float* __restrict__ seg_log_keep, const Record* __restrict__ records, + const Flip* __restrict__ overflow) +{{ + for (int sg = threadIdx.x; sg < num_segments; sg += blockDim.x) {{ + float lk = seg_log_keep[sg]; + int first = seg_first[sg], outcomes = seg_outcomes[sg], end = seg_end[sg]; + int idx = seg_begin[sg] - 1; + u64 state = mix64(stream ^ ((u64)sg * GOLDEN)); + if ((float)(end - 1 - idx) * -lk < 0x1p-12f) + add_fine_segment_faults(sw, events, num_events, max_events, ns, 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); + }} + __syncthreads(); + int ne = min(*num_events, max_events); + for (int e = threadIdx.x; e < ne; e += blockDim.x) {{ + uint2 v = events[e]; + apply_fault(sw, v.x & 255, v.x >> 8, v.y, stream, records, overflow); + }} + __syncthreads(); +}} + +// Screens the block's ``ns`` shots (words laid out as in ``apply_fault``): XORs +// the constants into their output words, marks the survivors in ``passed`` and +// in the bitmap ``survivors`` (if given), and appends them to ``queue`` after +// the ``count`` queued so far. +template +__device__ __forceinline__ void screen_shots(u64* sw, int ns, long long base, + const u64* __restrict__ constant, const u64* __restrict__ screen, bool* passed, + unsigned* __restrict__ queue, unsigned* __restrict__ count, unsigned* __restrict__ survivors) +{{ + for (int s0 = 0; s0 < {S}; s0 += blockDim.x) {{ + int s = s0 + threadIdx.x; + bool alive = false; + if (s < ns) {{ + u64 hit = 0; + #pragma unroll + for (int i = 0; i < {OW}; ++i) {{ + u64 v = sw[i * WS + s * SS] ^ constant[i]; + hit |= v & screen[i]; + sw[i * WS + s * SS] = v; + }} + alive = !hit; + }} + if (s < {S}) passed[s] = alive; + // NOTE: one atomic per warp: the leader reserves the survivors' slots. + unsigned ballot = __ballot_sync(0xffffffffu, alive); + long long word = (base + s0 + (threadIdx.x & ~31)) >> 5; + if (survivors != nullptr && (threadIdx.x & 31) == 0 && s0 + (threadIdx.x & ~31) < ns) + survivors[word] = ballot; + if (ballot) {{ + int lane = threadIdx.x & 31, leader = __ffs(ballot) - 1; + unsigned b = 0; + if (lane == leader) b = atomicAdd(count, __popc(ballot)); + b = __shfl_sync(0xffffffffu, b, leader); + if (alive) queue[b + __popc(ballot & ((1u << lane) - 1))] = (unsigned)(base + s); + }} + }} +}} +{NOISE}{KEPT}""" + + +def noise_tables( + channels: ChannelSampler, flips: WordTables, num_words: int +) -> NoiseTables: + """Return the noise kernels' tables, for CUDA's shared memory. + + Rows too wide for ``SHARED_BYTES`` keep their words in device memory + (``shared_words`` False), where the noise kernel XORs the faults in. + + Raises: + ValueError: If the noise exceeds the kernels' limits. + + """ + return _noise_tables( + channels, + flips, + num_words, + shared_bytes=SHARED_BYTES, + shared_words=words_fit(num_words, SHARED_BYTES), + ) + + +def source(tables: NoiseTables, output_words: int, num_outputs: int) -> str: + """Return the CUDA source of the ``noise`` and ``noise_kept`` kernels, and ``pack_rows``. + + ``noise_kept`` packs rows of ``num_outputs`` outputs. Without + ``tables.shared_words`` there is no ``noise_kept``: the rows of a block do + not fit in shared memory. + """ + packed_bytes = -(-num_outputs // 8) + if tables.shared_words: + noise, kept = NOISE_SHARED, NOISE_KEPT + else: + noise, kept = NOISE_DEVICE, "" + return SOURCE.replace("{NOISE}{KEPT}", noise + kept).format( + NT=THREADS, + NW=tables.num_words, + OW=output_words, + S=tables.block_shots, + NB=max(packed_bytes, 1), + NK=max(-(-packed_bytes // 8), 1), + LAST_WORD=f"0x{(1 << (num_outputs % 64 or 64)) - 1:016X}ULL", + ) diff --git a/src/tsim/statevector/cuda/runtime.py b/src/tsim/statevector/cuda/runtime.py new file mode 100644 index 00000000..8f7ce0ae --- /dev/null +++ b/src/tsim/statevector/cuda/runtime.py @@ -0,0 +1,121 @@ +"""CuPy and NVRTC for the CUDA executor, loaded on first use. + +CuPy compiles the kernels at runtime with NVRTC. A pip-installed CUDA 13 NVRTC +(``nvidia-cuda-nvrtc``) is preloaded when present: CuPy may otherwise pick up +an older system NVRTC that cannot target the GPU (CUDA 12.0 cannot target +``sm_120``, e.g. the RTX 5090). +""" + +from __future__ import annotations + +import ctypes +import functools +import importlib.util +import os +import pathlib +from types import ModuleType + +from tsim.statevector.gpu import FORKED + + +@functools.cache +def cupy() -> ModuleType: + """Return the ``cupy`` module, with a pip CUDA 13 NVRTC preloaded if installed. + + Raises: + ImportError: If CuPy is not installed. + + """ + if importlib.util.find_spec("cupy") is None: + raise ImportError("CuPy is not installed") + spec = importlib.util.find_spec("nvidia") + if spec is not None and spec.submodule_search_locations: + for root in spec.submodule_search_locations: + library = pathlib.Path(root) / "cu13" / "lib" / "libnvrtc.so.13" + if library.exists(): + ctypes.CDLL(str(library), mode=ctypes.RTLD_GLOBAL) + break + import cupy as cp + + return cp + + +_forked_after_init = False + + +def _after_fork_in_child() -> None: + # NOTE: a child forked after its parent initialized CUDA cannot use CUDA; the + # cached answer is read without touching CUDA. + global _forked_after_init + if _device_error.cache_info().currsize and _device_error() is None: + _forked_after_init = True + + +def available() -> bool: + """Return whether CuPy is installed and sees a CUDA device this process can use. + + ``CUDA_VISIBLE_DEVICES=""`` hides the devices, as in a process forked after + its parent used the GPU. + """ + return unavailable() is None + + +def unavailable() -> str | None: + """Return why this process cannot sample on CUDA, or ``None`` if it can.""" + if _forked_after_init: + return FORKED + return _device_error() + + +def forked_after_init() -> bool: + """Return whether this process was forked after its parent initialized CUDA.""" + return _forked_after_init + + +@functools.cache +def _device_error() -> str | None: + try: + cp = cupy() + if cp.cuda.runtime.getDeviceCount() > 0: + return None + return "CuPy sees no CUDA device" + except ImportError: + return ( + "CuPy is not installed (install Tsim's cuda extra: " + "pip install 'bloqade-tsim[cuda]')" + ) + except Exception as error: + return f"CuPy sees no usable CUDA device ({error})" + + +if hasattr(os, "register_at_fork"): # not on Windows + os.register_at_fork(after_in_child=_after_fork_in_child) + + +@functools.cache +def arch() -> str: + """Return the NVRTC architecture option of the current device, e.g. ``-arch=sm_120``.""" + return "-arch=sm_" + cupy().cuda.Device().compute_capability + + +@functools.cache +def multiprocessors() -> int: + """Return the number of streaming multiprocessors of the current device.""" + cuda = cupy().cuda + return int( + cuda.runtime.getDeviceProperties(cuda.Device().id)["multiProcessorCount"] + ) + + +def out_of_memory(error: BaseException) -> bool: + """Return whether ``error`` is CuPy's or CUDA's out-of-memory error.""" + cuda = cupy().cuda + if isinstance(error, cuda.memory.OutOfMemoryError): + return True + # NOTE: CUDA_ERROR_OUT_OF_MEMORY and cudaErrorMemoryAllocation are both 2. + errors = (cuda.driver.CUDADriverError, cuda.runtime.CUDARuntimeError) + return isinstance(error, errors) and getattr(error, "status", None) == 2 + + +class CompileError(RuntimeError): + """A CUDA kernel that NVRTC does not compile.""" diff --git a/src/tsim/statevector/cuda/sampler.py b/src/tsim/statevector/cuda/sampler.py new file mode 100644 index 00000000..5ab8b0d4 --- /dev/null +++ b/src/tsim/statevector/cuda/sampler.py @@ -0,0 +1,1028 @@ +"""``CudaShotSampler``: ``ShotSampler`` with the noise and the tapes on a CUDA GPU. + +A batch of shots runs as these kernels on one stream: + + noise draws every shot's noise into its words and screens it (``noise``); + the screened survivors go to a queue + tape runs every component on the queued shots, a group of ``2^tb`` threads + per shot with its amplitudes in registers (``codegen``, ``gpu.tape``); with a + postselection mask it clears the rejected shots from the screen's + bitmap of survivors, and with + expectations it writes each observable's expectation + kept under postselection, draws the noise that cannot change whether a + shot is kept on the accepted shots only (``noise_kept``), when that + noise is at least ``DEFER_SHARE`` of the faults + list lists the accepted shots in order from the bitmap (``count_bits``, + ``list_bits``) + select gathers the rows (or the accepted rows), bit-packed, and picks the + requested outputs out of them, bit-packed or as booleans, for the + copy to the host + +``sample_accepted`` runs a call as chunks of shots on two streams, so that a +chunk's copy to the host overlaps the next chunk's kernels. + +The kernels are compiled with NVRTC once per postselection mask (CuPy caches +them on disk). Amplitudes are complex64, so a sample follows the CPU executor's +distribution to within float32 rounding of the outcome probabilities; the noise +and the outcome draws come from other streams than the CPU's, so a seed's +samples differ from the CPU's, and they depend on the batch size. +""" + +from __future__ import annotations + +import math +import threading +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from tsim.noise.channels import ChannelSampler +from tsim.statevector.cuda import kernels, runtime +from tsim.statevector.cuda.codegen import CUDA +from tsim.statevector.cuda.kernels import COPY_THREADS, SCAN_THREADS, TAPE_THREADS +from tsim.statevector.cuda.noise import THREADS as NOISE_THREADS, noise_tables +from tsim.statevector.gpu import FORKED, tape +from tsim.statevector.gpu.noise import NoiseGroups, screen_pass_rate, split +from tsim.statevector.shots import ShotPlan, ShotSampler + +# NOTE: the words of a batch's ``counts``. +_SURVIVORS, _ACCEPTED, _DEVIATION = range(3) + +BATCH_SHOTS = 1 << 22 +"""Shots of one batch on the device.""" + +MIN_BATCH_SHOTS = 1 << 14 +"""Smallest batch that a call out of device memory retries with, before the CPU.""" + +CHUNK_BYTES = 6 << 20 +"""Bytes of kept rows per chunk that a call's chunks aim at: each chunk's copy to +the host overlaps the next chunk's kernels on the other stream.""" + +MAX_CHUNKS = 8 +"""Most chunks of a batch, so that a chunk's fixed costs stay small.""" + + +@dataclass +class _Launch: + """The device tables of one noise launch: its segments and records.""" + + args: list[Any] + num_segments: int + max_events: int + + +@dataclass +class _Module: + """A plan's compiled kernels and its device constants.""" + + noise: Any + noise_kept: Any + tape: Any + gather_columns: Any + gather_bytes: Any + gather_float: Any + select_rows: Any + publish: Any + count_bits: Any + list_bits: Any + screen: Any + tape_blocks: int + has_mask: bool + plan: ShotPlan + psm_words: np.ndarray + passing: float + """Estimated share of shots that pass the screen.""" + splits: dict[bool, tuple[_Launch, _Launch] | None] + """Each ``expectations`` value's noise passes (screened, deferred), or None.""" + + +@dataclass(frozen=True) +class _Output: + """One array of a call: chosen outputs of the accepted rows, ``width`` bytes a row. + + Bit-packed (zero bits past the outputs) or one boolean per output; ``whole`` + is every output in order, bit-packed: the gathered rows themselves. + """ + + columns: Any + """The outputs on the device (``uint32``), or None for the first ``count``.""" + count: int + width: int + bools: bool + whole: bool + + +def _pinned(shape: tuple[int, ...], dtype) -> np.ndarray: + """Return an uninitialized array in pinned host memory.""" + import cupyx + + # NOTE: CuPy's stub types the dtype as a float type. + return cupyx.empty_pinned(shape, dtype=dtype) # pyright: ignore[reportArgumentType] + + +def _pinned_counts(cp) -> tuple[np.ndarray, Any]: + """Return three pinned host counters and a device array over the same memory.""" + host = _pinned((3,), np.uint32) + # NOTE: with unified addressing a pinned host pointer is also a device pointer. + memory = cp.cuda.UnownedMemory(host.ctypes.data, host.nbytes, host) + device = cp.ndarray((3,), np.uint32, cp.cuda.MemoryPointer(memory, 0)) + return host, device + + +class _HostRows: + """A call's arrays of passing rows (and expectations) in pinned host memory. + + One array per output of ``widths`` bytes a row, filled chunk by chunk. + """ + + def __init__(self, widths: list[int], no: int | None) -> None: + self.widths, self.no = widths, no + self.capacity, self.total = 0, 0 + self.arrays: list[np.ndarray] = [] + self.values: np.ndarray | None = None + + def allocate(self, passed: int, of: int, shots: int) -> None: + """Reserve the rows of a call of ``shots``, ``passed`` of its first ``of`` passing.""" + rest = shots - of + share = passed / of + expected = rest * share + self._resize( + min(shots, passed + int(1.05 * expected + 6 * np.sqrt(expected) + 256)) + ) + + def grow(self, capacity: int) -> None: + """Make room for ``capacity`` rows, keeping the copied ones (no copy may be pending).""" + arrays, values, total = self.arrays, self.values, self.total + self._resize(capacity) + for new, old in zip(self.arrays, arrays, strict=True): + new[:total] = old[:total] + if values is not None and self.values is not None: + self.values[:total] = values[:total] + + def _resize(self, capacity: int) -> None: + self.capacity = capacity + self.arrays = [_pinned((capacity, w), np.uint8) for w in self.widths] + if self.no is not None: + self.values = _pinned((capacity, self.no), np.float32) + + def take(self, count: int) -> tuple[list[np.ndarray], np.ndarray | None]: + """Return the next ``count`` rows of each array (and expectations) to fill.""" + assert self.total + count <= self.capacity + lo, self.total = self.total, self.total + count + values = None if self.values is None else self.values[lo : self.total] + return [a[lo : self.total] for a in self.arrays], values + + def result(self) -> tuple[list[np.ndarray], np.ndarray | None]: + """Return the filled arrays and expectations.""" + if self.capacity == 0: + self._resize(0) + values = None if self.values is None else self.values[: self.total] + return [a[: self.total] for a in self.arrays], values + + +class CudaShotSampler: + """Samples shot words of a statevector program on a CUDA GPU. + + It keeps ``ShotSampler``'s interface (``sample``, ``direct_bools``, + ``reference``, ``plan``, ``layout``); ``sample`` copies only the shots that + pass postselection off the GPU. ``sample_accepted`` can also return each + observable's expectation. + """ + + supports = staticmethod(tape.supports) + + def __init__( + self, + host: ShotSampler, + 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, and the calls that the + GPU cannot run, unless ``required`` (``device="cuda"``): then those raise. + + Raises: + ImportError: If CuPy is not installed. + ValueError: If the noise exceeds the kernels' limits. + + """ + self._cp = runtime.cupy() + cp = self._cp + assert self.supports(host.layout) + self._args = (host, channels, num_detectors, seed, required) + self._required = required + self._cpu_only = False + 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) + n = self._noise + # NOTE: gather_columns stages a block's output words and its packed rows in + # shared memory; without room for 32 rows, gather_bytes reads them directly. + self._gather_shots = kernels.gather_shots(layout, n.block_shots) + self._record_args = [self._upload(n.overflow.view(np.uint8))] + self._all = self._launch(n.groups) + self._modules: dict[bytes | None, _Module] = {} + self._buffers: dict[str, Any] = {} + self._views: dict[tuple[str, tuple[int, ...]], Any] = {} + self._constants: dict[bytes, Any] = {} + self._columns: dict[bytes, Any] = {} + # NOTE: each buffer slot has a kernel stream and a copy stream, so that a + # slot's next kernels need not wait for its copy to the host. + # NOTE: a slot's noise kernel runs on a stream of the least priority, so the + # other slot's tape and packing, which the next copy waits for, go first. + self._streams = [ + cp.cuda.Stream(non_blocking=True, priority=-1) for _ in range(2) + ] + self._noise_streams = [cp.cuda.Stream(non_blocking=True) for _ in range(2)] + self._copies = [cp.cuda.Stream(non_blocking=True) for _ in range(2)] + self._published = [_pinned_counts(cp) for _ in range(2)] + # NOTE: lowered when a call runs out of device memory. + self._max_batch = BATCH_SHOTS + + @property + def device(self) -> str: + """Return where this sampler samples (``SamplerStats.device``). + + It is ``"cpu"`` after a kernel failed to compile, and in a process forked + after its parent used the GPU, where every call samples on the CPU. + """ + if self._cpu_only or runtime.forked_after_init(): + return "cpu" + return "cuda" + + def __reduce__(self): + """Pickle the CPU sampler and the inputs; the GPU state is rebuilt on load.""" + return _restore, (*self._args, self._rng.bit_generator.state) + + @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: + # NOTE: None, not b"", so that an empty mask is not taken for no mask. + key = None if mask is None else np.asarray(mask, bool).tobytes() + if key in self._modules: + return self._modules[key] + cp = self._cp + plan = self.plan(mask) + layout = self.layout + tb = tape.threads_per_shot(layout) + parts = tape.tape_kernel(layout, plan, mask, self._num_detectors, tb, CUDA) + source = kernels.source( + layout, self._noise, parts, tb, max(1, len(self._observables)) + ) + module = self._compile(source) + tape_fn = module.get_function("tape") + blocks_per_sm = 4 + if hasattr(cp.cuda.runtime, "occupancyMaxActiveBlocksPerMultiprocessor"): + blocks_per_sm = max( + 1, + cp.cuda.runtime.occupancyMaxActiveBlocksPerMultiprocessor( + tape_fn.ptr, TAPE_THREADS, 0 + ), + ) + compiled = _Module( + noise=module.get_function("noise"), + noise_kept=( + module.get_function("noise_kept") if self._noise.shared_words else None + ), + tape=tape_fn, + gather_columns=module.get_function("gather_columns"), + gather_bytes=module.get_function("gather_bytes"), + gather_float=module.get_function("gather_float"), + select_rows=module.get_function("select_rows"), + publish=module.get_function("publish"), + count_bits=module.get_function("count_bits"), + list_bits=module.get_function("list_bits"), + screen=self._upload(plan.screen), + tape_blocks=runtime.multiprocessors() * blocks_per_sm, + has_mask=mask is not None, + plan=plan, + psm_words=parts.psm_words, + passing=screen_pass_rate(self._channels, layout, plan), + splits={}, + ) + self._modules[key] = compiled + return compiled + + def _upload(self, array: np.ndarray) -> Any: + """Return ``array`` on the device, copied before any later kernel runs. + + ``cp.asarray`` copies on the current stream, usually the legacy default + stream, which the sampler's non-blocking streams do not wait for: a kernel + could read the array before the copy lands. Uploads are one-time, cached. + """ + device = self._cp.asarray(array) + self._cp.cuda.get_current_stream().synchronize() + return device + + def _launch(self, groups: NoiseGroups) -> _Launch: + """Upload a noise launch's segments and records.""" + return _Launch( + args=[ + *(self._upload(a) for a in groups.segments), + self._upload(groups.records.view(np.uint8)), + ], + num_segments=len(groups.segments[0]), + max_events=groups.max_events, + ) + + def _split( + self, module: _Module, expectations: bool + ) -> tuple[_Launch, _Launch] | None: + """Return a postselected batch's noise passes (screened, deferred), or None. + + As ``gpu.noise.split``; with ``expectations`` the observables' bits are + read too. Each key word is read whole, so the deferred channels flip + output words only. + """ + # NOTE: rows too wide for shared memory draw all their noise at once. + if not module.has_mask or not self._noise.shared_words: + return None + if expectations not in module.splits: + found = split( + self._noise, + self._channels, + self.layout, + module.plan, + module.psm_words, + self._num_detectors if expectations else None, + ) + module.splits[expectations] = ( + None + if found is None + else (self._launch(found[0]), self._launch(found[1])) + ) + return module.splits[expectations] + + def _compile(self, source: str): + """Compile with NVRTC; ptxas -O3 fails on some shuffled tapes, so lower levels follow. + + Raises: + runtime.CompileError: If no level compiles. + + """ + cp = self._cp + for level in ("3", "2", "1"): + options = ( + runtime.arch(), + "-std=c++17", + "-diag-suppress=177", + f"-Xptxas=-O{level}", + ) + try: + module = cp.RawModule(code=source, options=options) + module.get_function("tape") + return module + except cp.cuda.compiler.CompileException as error: + if level == "1": + raise runtime.CompileError( + f"NVRTC compilation failed: {error}" + ) from error + raise AssertionError("unreachable") + + # -- execution -------------------------------------------------------------- + + def _buffer(self, name: str, shape: tuple[int, ...], dtype) -> Any: + """Return a device buffer of at least ``shape``, reused across calls.""" + view = self._views.get((name, shape)) + if view is not None and view.dtype == np.dtype(dtype): + return view + size = int(np.prod(shape)) + held = self._buffers.get(name) + if held is None or held.size < size or held.dtype != np.dtype(dtype): + held = self._cp.empty(size, dtype=dtype) + self._buffers[name] = held + # NOTE: views of the replaced buffer must not be handed out again. + self._views = {k: v for k, v in self._views.items() if k[0] != name} + view = held[:size].reshape(shape) + self._views[(name, shape)] = view + return view + + def _constant(self, flip: np.ndarray | None) -> Any: + """Return the outputs' constants XOR the reference bits ``flip``, on the device.""" + key = b"" if flip is None else np.asarray(flip, bool).tobytes() + if key not in self._constants: + self._constants[key] = self._upload(self._host._constant(flip)) + return self._constants[key] + + def _run_batch( + self, + module: _Module, + shots: int, + first_shot: int, + constant: Any, + expectations: bool, + launch: _Launch, + slot: int = 0, + noise_stream: Any = None, + ) -> tuple[Any, Any, Any, Any]: + """Launch one batch on the current stream: ``(words, accepted, counts, values)``. + + ``accepted`` is the bitmap of the accepted shots (with a mask). + + ``launch`` holds the channels that the noise kernel draws; only the rows + of the shots that pass the screen are written. + ``counts`` holds the screened survivors, the accepted shots and the norm + deviation's bits. ``slot`` names the device buffers; ``noise_stream``, if + given, runs the noise kernel, after the current stream's work. + """ + cp = self._cp + layout = self.layout + nw, no = layout.num_words, max(1, len(self._observables)) + words = self._buffer(f"words{slot}", (shots, nw), np.uint64) + queue = self._buffer(f"queue{slot}", (shots,), np.uint32) + accepted = self._buffer(f"accepted{slot}", (-(-shots // 32),), np.uint32) + counts = self._buffer(f"counts{slot}", (3,), np.uint32) + values = ( + self._buffer(f"values{slot}", (shots, no), np.float32) + if expectations + else None + ) + current = cp.cuda.get_current_stream() + if noise_stream is not None: + noise_stream.wait_event(current.record()) + noise_stream.use() + counts.fill(0) + seed = np.uint64(self._rng.integers(0, 2**64, dtype=np.uint64)) + noise_seed = np.uint64(self._rng.integers(0, 2**64, dtype=np.uint64)) + S = self._noise.block_shots + module.noise( + ((shots + S - 1) // S,), + (NOISE_THREADS,), + ( + words, + np.int64(shots), + noise_seed, + np.int32(launch.num_segments), + np.int32(launch.max_events), + *launch.args, + *self._record_args, + constant, + module.screen, + queue, + counts[_SURVIVORS : _SURVIVORS + 1], + accepted if module.has_mask else np.uint64(0), + ), + shared_mem=8 * launch.max_events, + ) + if noise_stream is not None: + current.wait_event(noise_stream.record()) + current.use() + if layout.components: + module.tape( + (module.tape_blocks,), + (TAPE_THREADS,), + ( + words, + np.int64(first_shot), + seed, + queue, + counts[_SURVIVORS : _SURVIVORS + 1], + accepted if module.has_mask else np.uint64(0), + counts[_ACCEPTED : _ACCEPTED + 1], + values if values is not None else np.uint64(0), + counts[_DEVIATION : _DEVIATION + 1], + ), + ) + elif module.has_mask: + # NOTE: without components the screened survivors are the accepted shots. + counts[_ACCEPTED : _ACCEPTED + 1] = counts[_SURVIVORS : _SURVIVORS + 1] + return words, accepted, counts, values + + def _pack( + self, + module: _Module, + split: tuple[_Launch, _Launch] | None, + words, + count: int, + index, + ) -> Any: + """Return the ``count`` rows ``index`` (or all) on the device, as copied to the host. + + Their outputs, bit-packed; with ``split`` the deferred noise kernel adds + the deferred channels' faults first. + """ + S = self._noise.block_shots + nb = -(-self.layout.num_outputs // 8) + # NOTE: a new array per chunk, so that refilling it never waits for a copy. + dst = self._cp.empty((max(count, 1), nb), dtype=np.uint8) + if not count or not nb: + return dst + if split is not None: + blocks = (count + S - 1) // S + launch = split[1] + seed = np.uint64(self._rng.integers(0, 2**64, dtype=np.uint64)) + module.noise_kept( + (blocks,), + (NOISE_THREADS,), + ( + words, + np.int64(count), + index, + seed, + np.int32(launch.num_segments), + np.int32(launch.max_events), + *launch.args, + *self._record_args, + dst, + ), + shared_mem=max(8 * launch.max_events, S * nb), + ) + elif self._gather_shots == 0: + module.gather_bytes( + (min(-(-count * nb // COPY_THREADS), 16 * runtime.multiprocessors()),), + (COPY_THREADS,), + ( + words, + np.int64(count), + index if index is not None else np.uint64(0), + dst, + ), + ) + else: + GS = self._gather_shots + module.gather_columns( + ((count + GS - 1) // GS,), + (COPY_THREADS,), + ( + words, + np.int64(count), + index if index is not None else np.uint64(0), + dst, + ), + shared_mem=GS * nb, + ) + return dst + + def _select(self, module: _Module, rows, count: int, output: _Output) -> Any: + """Return ``output``'s outputs of the ``count`` packed ``rows``, on the device.""" + dst = self._cp.empty((count, output.width), dtype=np.uint8) + blocks = min( + -(-count * output.width // COPY_THREADS), 16 * runtime.multiprocessors() + ) + module.select_rows( + (max(blocks, 1),), + (COPY_THREADS,), + ( + rows, + np.int64(count), + output.columns if output.columns is not None else np.uint64(0), + np.int32(output.count), + np.int32(output.width), + np.int32(output.bools), + dst, + ), + ) + return dst + + def _outputs( + self, columns: list[np.ndarray], bools: bool + ) -> tuple[list[_Output], list[tuple[int, int | None]]]: + """Return the arrays that a call returning ``columns`` copies to the host. + + Also returns, for each entry of ``columns``, its array and the bytes of + each row it keeps (or None for all): bit-packed leading whole bytes are a + view of the rows, which are copied once. + """ + num_outputs = self.layout.num_outputs + nb = -(-num_outputs // 8) + found: list[_Output] = [] + picks: list[tuple[int, int | None]] = [] + rows = None + for chosen in columns: + chosen = np.asarray(chosen, dtype=np.uint32) + count = len(chosen) + in_order = bool(np.array_equal(chosen, np.arange(count))) + if in_order and not bools and (count == num_outputs or count % 8 == 0): + if rows is None: + rows = len(found) + found.append(_Output(None, num_outputs, nb, False, True)) + picks.append((rows, None if count == num_outputs else count // 8)) + continue + device = None + if not in_order: + key = chosen.tobytes() + if key not in self._columns: + self._columns[key] = self._upload(chosen) + device = self._columns[key] + found.append( + _Output( + columns=device, + count=count, + width=count if bools else -(-count // 8), + bools=bools, + whole=False, + ) + ) + picks.append((len(found) - 1, None)) + return found, picks + + def _list(self, module: _Module, bits, shots: int, count: int, slot: int) -> Any: + """Return the ``count`` shots of the bitmap ``bits`` of ``shots`` shots, in order.""" + words = -(-shots // 32) + blocks = -(-words // SCAN_THREADS) + totals = self._buffer(f"totals{slot}", (blocks,), np.uint32) + index = self._buffer(f"index{slot}", (max(count, 1),), np.uint32) + grid, threads = (blocks,), (SCAN_THREADS,) + module.count_bits(grid, threads, (bits, np.int64(words), totals)) + module.list_bits(grid, threads, (bits, np.int64(words), totals, index)) + return index[:count] + + def _or_cpu(self, on_gpu, on_cpu): + """Return ``on_gpu()``, or ``on_cpu()`` where the GPU fails. + + It fails in a process forked after its parent used the GPU, out of GPU + memory, and when a kernel does not compile; after the last the sampler + samples on the CPU. A ``required`` sampler raises instead. + + Raises: + RuntimeError: If ``required``, in a process forked after its parent + used the GPU. + MemoryError: If ``required``, out of GPU memory. + runtime.CompileError: If ``required``, when a kernel does not compile. + + """ + if self._cpu_only: + return on_cpu() + if runtime.forked_after_init(): + if self._required: + raise RuntimeError(f"Cannot sample on CUDA: {FORKED}.") + return on_cpu() + try: + return on_gpu() + except MemoryError: + if self._required: + raise + return on_cpu() + except runtime.CompileError: + if self._required: + raise + self._cpu_only = True + return on_cpu() + + def sample( + self, + shots: int, + columns: list[np.ndarray], + *, + bit_packed: bool, + batch_size: int | None = None, + mask: np.ndarray | None = None, + flip: np.ndarray | None = None, + ) -> tuple[list[np.ndarray], float]: + """Return the output bits ``columns`` of ``shots`` shots and the largest norm deviation. + + As ``ShotSampler.sample``: with a ``mask`` only the shots that pass + postselection are returned. ``batch_size`` bounds the shots of a device + batch. Out of GPU memory, when a kernel does not compile, and in a + process forked after its parent used the GPU, the call samples on the + CPU, or raises if the sampler is ``required`` on the GPU (``_or_cpu``). + """ + + def on_gpu(): + outputs, _, deviation = self._accepted( + shots, + mask=mask, + flip=flip, + batch_size=batch_size, + bools=not bit_packed, + columns=columns, + ) + return outputs, deviation + + return self._or_cpu( + on_gpu, + lambda: self._host.sample( + shots, + columns, + bit_packed=bit_packed, + batch_size=batch_size, + mask=mask, + flip=flip, + ), + ) + + def _accepted(self, shots: int, *, batch_size: int | None, **options) -> Any: + """Return ``sample_accepted(shots, **options)``, in smaller batches if the GPU is full. + + Out of device memory, the call is retried in batches half as large, down + to ``MIN_BATCH_SHOTS``. Later calls start from the batch that fit. + + Raises: + MemoryError: If even batches of ``MIN_BATCH_SHOTS`` do not fit. + + """ + batch = min(batch_size or BATCH_SHOTS, self._max_batch) + while True: + try: + return self.sample_accepted(shots, batch_size=batch, **options) + except Exception as error: + if not runtime.out_of_memory(error): + raise + self._release() + self._max_batch = max(batch // 2, MIN_BATCH_SHOTS) + if batch <= MIN_BATCH_SHOTS: + raise MemoryError( + "the GPU ran out of memory, even for batches of " + f"{MIN_BATCH_SHOTS} shots: {error}" + ) from error + batch = self._max_batch + + def _release(self) -> None: + """Wait for the device and free the buffers that calls reuse.""" + cp = self._cp + cp.cuda.Device().synchronize() + self._buffers.clear() + self._views.clear() + cp.get_default_memory_pool().free_all_blocks() + + def _synchronize(self) -> None: + """Wait for every kernel and copy of ``sample_accepted``.""" + for stream in (*self._streams, *self._copies): + stream.synchronize() + + def _chunk_shots(self, shots: int, nb: int, passing: float) -> int: + """Return a chunk's shots: about ``CHUNK_BYTES`` of rows of ``nb`` bytes. + + ``passing`` is the share of shots kept, and a batch of ``shots`` shots + splits into at most ``MAX_CHUNKS`` chunks. + """ + chunks = min(MAX_CHUNKS, max(1, math.ceil(nb * passing * shots / CHUNK_BYTES))) + return -(-shots // chunks) + + def sample_accepted( + self, + shots: int, + *, + mask: np.ndarray | None = None, + flip: np.ndarray | None = None, + expectations: bool = False, + batch_size: int | None = None, + bools: bool = False, + columns: list[np.ndarray] | None = None, + ) -> tuple[Any, np.ndarray | None, float]: + """Sample ``shots`` shots; return only those that pass postselection. + + A shot passes when none of the ``mask``'s detectors is set (every shot + passes without a mask). Rows keep their shots' order and hold every output; + with ``bools``, one boolean per output, unpacked on the GPU. With + ``columns``, the outputs are a list of arrays instead, one per entry: its + outputs, in its order, selected on the GPU. + + The shots run in chunks on two streams: each chunk's rows go to the host + while the next chunk's kernels run. The chunks hold about ``CHUNK_BYTES`` + of rows, by the screen's estimated share of passing shots, so a seed's + samples do not depend on earlier calls. + + Returns: + ``(outputs, expectations, deviation)``: the passing shots' outputs, + bit-packed as ``np.packbits(..., bitorder="little")`` would (``uint8``, + in pinned host memory), as booleans, or with ``columns`` a list of such + arrays; with ``expectations`` each observable's + expectation of ``(-1)^observable`` given the shot's history before + the measurement that settles it, or its sampled ``±1`` when a + postselected output of its component is tested at or after that + measurement (else None), and the largest normalization deviation + beyond float32 rounding. + + """ + with self._lock: + return self._sample_accepted( + shots, + mask=mask, + flip=flip, + expectations=expectations, + batch_size=batch_size, + bools=bools, + columns=columns, + ) + + def _sample_accepted( + self, + shots: int, + *, + mask: np.ndarray | None, + flip: np.ndarray | None, + expectations: bool, + batch_size: int | None, + bools: bool, + columns: list[np.ndarray] | None, + ) -> tuple[Any, np.ndarray | None, float]: + cp = self._cp + module = self._module(mask) + layout = self.layout + nb, no = -(-layout.num_outputs // 8), len(self._observables) + chosen = [np.arange(layout.num_outputs)] if columns is None else columns + outputs, picks = self._outputs(chosen, bools) + out = _HostRows([o.width for o in outputs], no if expectations else None) + if not shots: + arrays, values = out.result() + return self._arrays(arrays, picks, columns, bools), values, 0.0 + constant = self._constant(flip) + split = self._split(module, expectations) + launch = split[0] if split else self._all + batch = min(batch_size or BATCH_SHOTS, BATCH_SHOTS) + state = {"start": 0, "slot": 0} + # NOTE: by the packed rows, so that a seed's shots do not depend on the + # outputs that a call returns. + size = self._chunk_shots(min(batch, shots), nb, module.passing) + deviation = 0.0 + pending: list[tuple[int, int, int, tuple[Any, Any, Any, Any], Any]] = [] + # NOTE: the device rows in flight to the host, kept until the copies land, + # and the streams that copy them. + held: list[Any] = [] + used: set[Any] = set() + + def launch_next() -> None: + first, slot = state["start"], state["slot"] + n = min(size, shots - first) + device = self._published[slot][1] + ready = None + with self._streams[slot]: + launched = self._run_batch( + module, + n, + first, + constant, + expectations, + launch, + slot, + self._noise_streams[slot], + ) + # NOTE: no copy is in flight when the first chunk finishes, so it + # reads its counts with a copy of its own. + if first: + module.publish((1,), (32,), (launched[2], device)) + ready = self._streams[slot].record() + pending.append((slot, first, n, launched, ready)) + state["start"], state["slot"] = first + n, 1 - slot + + def finish() -> float: + slot, first, n, (words, accepted, counts, values), ready = pending.pop(0) + if ready is None: + host_counts = counts.get(stream=self._streams[slot]) + else: + ready.synchronize() + host_counts = self._published[slot][0].copy() + with self._streams[slot]: + count = n if mask is None else int(host_counts[_ACCEPTED]) + if first == 0: + out.allocate(count, n, shots) + elif out.total + count > out.capacity: + # NOTE: a rare excess of passing shots: the copies in flight land first. + self._synchronize() + out.grow(out.total + count + (shots - first - n)) + parts, values_host = out.take(count) + index = ( + None + if mask is None + else self._list(module, accepted, n, count, slot) + ) + stream = self._streams[slot] + # NOTE: a call of one chunk has nothing to overlap its copy with. + copy = stream if first + n == shots and not held else self._copies[slot] + used.add(copy) + if count and nb: + rows = self._pack(module, split, words, count, index) + held.append(rows) + found = [ + ( + rows[:count] + if o.whole + else ( + self._select(module, rows, count, o) + if o.width + else None + ) + ) + for o in outputs + ] + if copy is not stream: + copy.wait_event(stream.record()) + for part, dst in zip(parts, found, strict=True): + if dst is not None: + dst.get(stream=copy, out=part, blocking=False) + held.append(dst) + if values_host is not None and count: + if no and layout.components: + if index is None: + index = cp.arange(n, dtype=cp.uint32) + held.append(index) + dst = cp.empty((count, no), dtype=np.float32) + module.gather_float( + (runtime.multiprocessors() * 8,), + (COPY_THREADS,), + (values, np.int64(count), np.int32(no), index, dst), + ) + if copy is not stream: + copy.wait_event(stream.record()) + dst.get(stream=copy, out=values_host, blocking=False) + held.append(dst) + else: + values_host[:] = 0 + return tape.deviation(host_counts[_DEVIATION]) + + # NOTE: each chunk is launched before the previous one is finished, so its + # kernels overlap that copy. + launch_next() + while state["start"] < shots: + launch_next() + if len(pending) == 2: + deviation = max(deviation, finish()) + while pending: + deviation = max(deviation, finish()) + for stream in used: + stream.synchronize() + arrays, values = out.result() + if values is not None: + # NOTE: the rows of every output, bit-packed, which a call with + # expectations returns. + self._direct_expectations(arrays[0], values) + return self._arrays(arrays, picks, columns, bools), values, deviation + + @staticmethod + def _arrays( + arrays: list[np.ndarray], + picks: list[tuple[int, int | None]], + columns, + bools: bool, + ) -> Any: + """Return a call's outputs from its ``arrays`` and ``picks`` (``_outputs``). + + Booleans are viewed as such; without ``columns`` there is one array. + """ + found = [arrays[k] if cut is None else arrays[k][:, :cut] for k, cut in picks] + if bools: + found = [a.view(bool) for a in found] + return found[0] if columns is None else found + + def _direct_expectations(self, rows: np.ndarray, values: np.ndarray) -> None: + """Set the expectations of observables outside the components: ``±1``.""" + 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._accepted(shots, flip=flip, batch_size=batch_size, bools=True)[ + 0 + ], + lambda: self._host.direct_bools(shots, flip, batch_size), + ) + + +def _restore( + host: ShotSampler, + channels: ChannelSampler, + num_detectors: int, + seed: np.random.SeedSequence, + required: bool, + state: dict, +) -> CudaShotSampler | ShotSampler: + """Rebuild a pickled ``CudaShotSampler``, or its CPU sampler where CUDA cannot run. + + As at compile time, a sampler ``required`` on the GPU raises where CUDA cannot run. + """ + from tsim.statevector import cuda + + shots = cuda.shot_sampler(host, channels, num_detectors, seed, required=required) + if isinstance(shots, CudaShotSampler): + shots._rng.bit_generator.state = state + return shots diff --git a/src/tsim/statevector/gpu/tape.py b/src/tsim/statevector/gpu/tape.py index 22e6a97a..2ff95ea5 100644 --- a/src/tsim/statevector/gpu/tape.py +++ b/src/tsim/statevector/gpu/tape.py @@ -34,6 +34,7 @@ from __future__ import annotations +import math import re from collections.abc import Callable from dataclasses import dataclass @@ -162,6 +163,33 @@ def supports(layout: ShotLayout) -> bool: ) +def unsupported(layout: ShotLayout) -> str | None: + """Return why a GPU executor cannot run ``layout``'s components, or ``None`` if it can.""" + if supports(layout) or layout.num_words == 0: + return None + width = max(c.width for c in layout.components) + if width > MAX_WIDTH: + return ( + f"a component needs {width} active qubits, more than the " + f"{MAX_WIDTH} that the GPU executor supports" + ) + # NOTE: a rotation's branch weights are cos^2 and sin^2 of half its angle, so + # a weight w is that of an angle 2 asin(sqrt(w)) from a multiple of π. The + # tape's rotations are the circuit's without their Clifford parts (multiples + # of π/2), so the circuit's angle is as far from a multiple of π/2. + distance = _angle(rarest_branch(layout)) + return ( + f"a rotation's angle is {distance:.1e} π from a Clifford angle (a " + f"multiple of π/2), closer than the {_angle(RESOLVABLE):.1e} π that " + "float32 amplitudes resolve" + ) + + +def _angle(weight: float) -> float: + """Return the angle, in units of π, whose branch weight ``sin^2`` is ``weight``.""" + return 2 * math.asin(math.sqrt(weight)) / math.pi + + def rarest_branch(layout: ShotLayout) -> float: """Return the smallest nonzero branch weight (``cos^2``, ``sin^2``) of a rotation. diff --git a/src/tsim/statevector/metal/__init__.py b/src/tsim/statevector/metal/__init__.py index e5468af0..01314d14 100644 --- a/src/tsim/statevector/metal/__init__.py +++ b/src/tsim/statevector/metal/__init__.py @@ -11,10 +11,9 @@ from __future__ import annotations -import math from typing import TYPE_CHECKING -from tsim.statevector.gpu.tape import MAX_WIDTH, RESOLVABLE, rarest_branch, supports +from tsim.statevector.gpu.tape import unsupported from tsim.statevector.metal import runtime if TYPE_CHECKING: @@ -22,7 +21,7 @@ from tsim.noise.channels import ChannelSampler from tsim.statevector.metal.sampler import MetalShotSampler - from tsim.statevector.shots import ShotLayout, ShotSampler + from tsim.statevector.shots import ShotSampler def shot_sampler( @@ -57,7 +56,7 @@ def shot_sampler( return host if host.layout.num_words == 0: return host - reason = _unsupported(host.layout) + reason = unsupported(host.layout) if reason is not None: if required: raise ValueError( @@ -78,30 +77,3 @@ def shot_sampler( if required: raise ValueError(f"Cannot sample on Metal: {error}.") from error return host - - -def _unsupported(layout: ShotLayout) -> str | None: - """Return why the GPU executor cannot run ``layout``, or ``None`` if it can.""" - if supports(layout): - return None - width = max(c.width for c in layout.components) - if width > MAX_WIDTH: - return ( - f"a component needs {width} active qubits, more than the " - f"{MAX_WIDTH} that the GPU executor supports" - ) - # NOTE: a rotation's branch weights are cos^2 and sin^2 of half its angle, so - # a weight w is that of an angle 2 asin(sqrt(w)) from a multiple of π. The - # tape's rotations are the circuit's without their Clifford parts (multiples - # of π/2), so the circuit's angle is as far from a multiple of π/2. - distance = _angle(rarest_branch(layout)) - return ( - f"a rotation's angle is {distance:.1e} π from a Clifford angle (a " - f"multiple of π/2), closer than the {_angle(RESOLVABLE):.1e} π that " - "float32 amplitudes resolve" - ) - - -def _angle(weight: float) -> float: - """Return the angle, in units of π, whose branch weight ``sin^2`` is ``weight``.""" - return 2 * math.asin(math.sqrt(weight)) / math.pi diff --git a/test/conftest.py b/test/conftest.py index 1435e449..16e49047 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,10 +1,17 @@ import pytest +from tsim.statevector.cuda import runtime as cuda_runtime from tsim.statevector.metal import runtime as metal_runtime +GPU_TESTS = { + "test_metal_sampler.py": metal_runtime, + "test_cuda_sampler.py": cuda_runtime, +} + @pytest.fixture(autouse=True) def _cpu_samplers(request, monkeypatch): - """Sample on the CPU, except in the Metal tests, which choose per sampler.""" - if request.path.name != "test_metal_sampler.py": - monkeypatch.setattr(metal_runtime, "unavailable", lambda: "a test on the CPU") + """Sample on the CPU, except in each GPU's tests, which choose per sampler.""" + for name, runtime in GPU_TESTS.items(): + if request.path.name != name: + monkeypatch.setattr(runtime, "unavailable", lambda: "a test on the CPU") diff --git a/test/integration/test_cuda_sampler.py b/test/integration/test_cuda_sampler.py new file mode 100644 index 00000000..13e03578 --- /dev/null +++ b/test/integration/test_cuda_sampler.py @@ -0,0 +1,998 @@ +"""The CUDA 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.cuda import codegen, kernels, runtime, sampler as cuda_sampler +from tsim.statevector.cuda.sampler import CudaShotSampler +from tsim.statevector.gpu import noise as gpu_noise, tape +from tsim.statevector.program import StatevectorComponent +from tsim.statevector.shots import ShotSampler + +pytestmark = pytest.mark.skipif( + not runtime.available(), reason="needs CuPy and a CUDA GPU" +) + +# NOTE: t_ghz plus a direct detector (detector 2), as in test_trajectory_sharing. +GHZ_WITH_DIRECT_DETECTOR = """ +R 0 1 2 3 +H 0 +CNOT 0 1 +CNOT 1 2 +T 0 1 2 +H 0 1 2 +DEPOLARIZE1(0.05) 0 1 2 +T 0 1 +CZ 0 1 +H 0 1 +M 0 1 2 +DETECTOR rec[-3] rec[-2] +DETECTOR rec[-2] rec[-1] +OBSERVABLE_INCLUDE(0) rec[-1] +X_ERROR(0.2) 3 +M 3 +DETECTOR rec[-1] +""" + +# NOTE: a T state measured in X after a noisy idle: its observable's expectation +# is cos(π/4) = 0.707 before noise. +T_STATE_X = """ +R 0 +H 0 +T 0 +Z_ERROR(0.1) 0 +H 0 +M 0 +OBSERVABLE_INCLUDE(0) rec[-1] +""" + + +def _gpu(sampler): + """Return ``sampler``, which must run on the GPU.""" + _gpu_shots(sampler) + return sampler + + +def _gpu_shots(sampler) -> CudaShotSampler: + """Return the GPU shot sampler of ``sampler``.""" + assert isinstance(sampler._shots, CudaShotSampler) + return sampler._shots + + +def _cpu(compile, **kwargs): + """Return ``compile(**kwargs)``, a sampler on the CPU.""" + sampler = compile(device="cpu", **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_cuda_sampling_matches_stabilizer_rank(monkeypatch, name, kind, looping): + if looping: + # NOTE: measurements, projections and pruning in looping tapes. + monkeypatch.setattr(tape, "UNROLL_BUDGET", 0) + circuit = Circuit(CIRCUITS[name]) + shots = 40_000 + if kind == "detector": + expected = circuit.compile_detector_sampler(strategy="cat5", seed=1).sample( + shots, batch_size=10_000, append_observables=True + ) + got = _gpu(circuit.compile_detector_sampler(seed=2)).sample( + shots, append_observables=True + ) + else: + expected = circuit.compile_sampler(strategy="cat5", seed=1).sample( + shots, batch_size=10_000 + ) + got = _gpu(circuit.compile_sampler(seed=2)).sample(shots) + assert_counts_match(got, expected) + + +@pytest.mark.parametrize("name", sorted(CIRCUITS)) +def test_cuda_postselection_keeps_the_accepted_distribution(name): + circuit = Circuit(CIRCUITS[name]) + if circuit.num_detectors == 0: + pytest.skip("no detectors to postselect") + shots = 40_000 + mask = np.ones(circuit.num_detectors, dtype=bool) + expected = _accepted(_cpu(circuit.compile_detector_sampler, seed=1), shots, mask) + got = _accepted(_gpu(circuit.compile_detector_sampler(seed=2)), shots, mask) + rate, cuda_rate = len(expected) / shots, len(got) / shots + sigma = np.sqrt(rate * (1 - rate) * 2 / shots) + assert abs(cuda_rate - rate) <= 5 * sigma + 3 / shots + if len(got) and len(expected): + assert_counts_match(got, expected) + + +@pytest.mark.parametrize("seed", [1, 2, 3]) +def test_cuda_sampling_of_random_noisy_circuits(seed): + stim_circuit = gen_stim_circuit(3, 60, p_t=3, p_x_error=0.5, p_r_z=0.5, seed=seed) + circuit = Circuit.from_stim_program(stim_circuit) + shots = 8_192 + expected = circuit.compile_sampler(strategy="cat5", seed=seed).sample( + shots, batch_size=shots + ) + got = _gpu(circuit.compile_sampler(seed=seed + 10)).sample(shots) + assert_counts_match(got, expected) + + +@pytest.mark.parametrize("reference", [False, True]) +def test_cuda_reference_samples_and_pruning(reference): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + mask = np.ones(circuit.num_detectors, dtype=bool) + shots = 40_000 + kwargs = {"use_detector_reference_sample": reference} + expected = _accepted( + _cpu(circuit.compile_detector_sampler, seed=3), shots, mask, **kwargs + ) + got = _accepted( + _gpu(circuit.compile_detector_sampler(seed=3)), shots, mask, **kwargs + ) + rate, cuda_rate = len(expected) / shots, len(got) / shots + assert abs(cuda_rate - rate) <= 5 * np.sqrt(rate * (1 - rate) * 2 / shots) + assert_counts_match(got, expected) + + +def test_cuda_postselection_keeps_every_column(): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + packed = _packed(_gpu(circuit.compile_detector_sampler(seed=4)), 20_000, mask) + outputs = circuit.num_detectors + circuit.num_observables + assert packed.shape[1] == -(-outputs // 8) + bits = np.unpackbits(packed, axis=1, bitorder="little") + assert not bits[:, outputs:].any() + assert not bits[:, : circuit.num_detectors][:, mask].any() + full = _accepted(_cpu(circuit.compile_detector_sampler, seed=5), 20_000, mask) + assert abs(len(bits) - len(full)) <= 5 * np.sqrt(len(full)) + assert_counts_match(bits[:, :outputs].astype(bool), full) + + +@pytest.mark.parametrize("bit_packed", [False, True]) +def test_cuda_arranges_the_columns_as_the_cpu_does(bit_packed): + # NOTE: samplers of one seed draw the same shots, whichever columns they return. + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[0] = True + num_detectors = circuit.num_detectors + + def sample(**kwargs): + sampler = _gpu(circuit.compile_detector_sampler(seed=17)) + return sampler.sample( + 5_000, postselection_mask=mask, bit_packed=bit_packed, **kwargs + ) + + def bools(array, count): + if not bit_packed: + return array + return np.unpackbits(array, axis=1, count=count, bitorder="little").astype(bool) + + appended = bools(sample(append_observables=True), num_detectors + 1) + detectors = bools(sample(), num_detectors) + prepended = bools(sample(prepend_observables=True), num_detectors + 1) + separate = sample(separate_observables=True) + np.testing.assert_array_equal(detectors, appended[:, :num_detectors]) + np.testing.assert_array_equal(prepended[:, 0], appended[:, -1]) + np.testing.assert_array_equal(prepended[:, 1:], detectors) + np.testing.assert_array_equal(bools(separate[0], num_detectors), detectors) + np.testing.assert_array_equal(bools(separate[1], 1), appended[:, -1:]) + cpu = _cpu(circuit.compile_detector_sampler, seed=17).sample( + 10, postselection_mask=mask, bit_packed=bit_packed, append_observables=True + ) + gpu = sample(append_observables=True) + assert gpu.dtype == cpu.dtype + assert gpu.shape[1] == cpu.shape[1] + + +def test_cuda_samples_no_shots(): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + sampler = _gpu(circuit.compile_detector_sampler(seed=1)) + assert sampler.sample(0, append_observables=True).shape == (0, 4) + assert sampler.sample(0, bit_packed=True).shape == (0, 1) + for bit_packed in (False, True): + detectors, observables = sampler.sample( + 0, separate_observables=True, bit_packed=bit_packed + ) + assert detectors.shape == (0, 3 if not bit_packed else 1) + assert observables.shape == (0, 1) + + +def test_direct_booleans_take_the_batch_size_as_packed_rows_do(): + # NOTE: each batch draws its own stream, so a component-free sampler's + # booleans and packed rows match only if both honor batch_size. + circuit = Circuit("X_ERROR(0.3) 0 1 2\nM 0 1 2\nDETECTOR rec[-1]\nDETECTOR rec[-2]") + + def sample(bit_packed): + sampler = _gpu(circuit.compile_detector_sampler(seed=5)) + return sampler.sample(5_000, batch_size=1_000, bit_packed=bit_packed) + + packed = np.packbits(sample(False), axis=1, bitorder="little") + np.testing.assert_array_equal(packed, sample(True)) + + +# NOTE: 150 noisy detectors, so the rows span three words, and two observables. +NOISY_DETECTORS = "\n".join( + [ + "R " + " ".join(map(str, range(150))), + "X_ERROR(0.3) " + " ".join(map(str, range(150))), + "M " + " ".join(map(str, range(150))), + *(f"DETECTOR rec[{i - 150}]" for i in range(150)), + "OBSERVABLE_INCLUDE(0) rec[-1] rec[-7]", + "OBSERVABLE_INCLUDE(1) rec[-150]", + ] +) + + +def test_cuda_postselection_drops_the_rows_that_fail(monkeypatch): + # NOTE: the masked run must draw the unmasked run's noise, so none is deferred. + monkeypatch.setattr(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_cuda_defers_the_noise_that_postselection_cannot_see(): + circuit = Circuit(UNREAD_NOISE) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + sampler = _gpu(circuit.compile_detector_sampler(seed=12)) + got = _accepted(sampler, 200_000, mask) + assert _gpu_shots(sampler)._module(mask).splits[False] is not None + full = _accepted(_cpu(circuit.compile_detector_sampler, seed=13), 200_000, mask) + assert abs(len(got) - len(full)) <= 5 * np.sqrt(len(full)) + assert_counts_match(got, full) + + +def _expectations(sampler, shots, mask=None): + """Return the passing rows and expectations of ``sample_accepted``.""" + rows, values, _ = _gpu_shots(sampler).sample_accepted( + shots, mask=mask, expectations=True + ) + assert values is not None + return rows, values + + +def test_deferred_noise_leaves_the_expectations_unbiased(): + # NOTE: with expectations, the noise on the observable's direct qubit is not + # deferred: the tape reads the observable's bit for its sign. + circuit = Circuit(UNREAD_NOISE) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + sampler = _gpu(circuit.compile_detector_sampler(seed=14)) + _, values = _expectations(sampler, 200_000, mask) + assert _gpu_shots(sampler)._module(mask).splits[True] is not None + full = _accepted(_cpu(circuit.compile_detector_sampler, seed=15), 200_000, mask) + expected = np.mean(1 - 2 * full[:, -1].astype(float)) + assert abs(values.mean() - expected) < 5 / np.sqrt(len(full)) + + +@pytest.mark.parametrize("text", [NOISY_DETECTORS, UNREAD_NOISE]) +def test_cuda_chunks_return_the_rows_of_one_pass(monkeypatch, text): + # NOTE: 1000-shot chunks on both streams, and host rows that start with room + # for the first chunk only; without deferral, a masked and an unmasked run on + # the same chunks draw the same shots. + monkeypatch.setattr(gpu_noise, "DEFER_SHARE", 2.0) + monkeypatch.setattr( + cuda_sampler.CudaShotSampler, "_chunk_shots", lambda self, *args: 1000 + ) + monkeypatch.setattr( + cuda_sampler._HostRows, + "allocate", + lambda self, passed, of, shots: self._resize(passed), + ) + circuit = Circuit(text) + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[[0, 2]] = True + outputs = circuit.num_detectors + circuit.num_observables + full = _packed(_gpu(circuit.compile_detector_sampler(seed=16)), 10_500) + rows = _packed(_gpu(circuit.compile_detector_sampler(seed=16)), 10_500, mask) + full_bits = np.unpackbits(full, axis=1, count=outputs, bitorder="little") + passing = ~full_bits[:, : circuit.num_detectors][:, mask].any(axis=1) + assert 0 < passing.sum() < 10_500 + np.testing.assert_array_equal(rows, full[passing]) + + +def test_observable_expectations_average_to_the_sampled_observable(): + circuit = Circuit(T_STATE_X) + sampler = _gpu(circuit.compile_detector_sampler(seed=6)) + packed, values = _expectations(sampler, 200_000) + assert values.shape == (len(packed), 1) + # NOTE: given the shot's Z error, the expectation is ±cos(π/4). + np.testing.assert_allclose(np.abs(values), np.sqrt(0.5), atol=1e-5) + observed = 1 - 2 * np.unpackbits(packed, axis=1, bitorder="little")[:, 0].astype( + float + ) + sampled = _cpu(circuit.compile_detector_sampler, seed=7).sample( + 200_000, append_observables=True + )[:, 0] + expected = np.mean(1 - 2 * sampled.astype(float)) + assert abs(values.mean() - expected) < 5 * 1 / np.sqrt(200_000) + # NOTE: the returned observable bit is sampled from the same state. + assert abs(observed.mean() - values.mean()) < 5 / np.sqrt(200_000) + + +# NOTE: the T state's X outcome is copied onto qubit 1, whose postselected (and +# noisy) detector is settled after the observable's measurement. +T_STATE_X_THEN_DETECTOR = """ +R 0 1 +H 0 +T 0 +H 0 +CNOT 0 1 +X_ERROR(0.1) 1 +M 0 +OBSERVABLE_INCLUDE(0) rec[-1] +M 1 +DETECTOR rec[-1] +""" + +# NOTE: a postselected T-state detector settled before the observable's measurement. +T_STATE_X_AFTER_DETECTOR = """ +R 0 1 +H 0 1 +T 0 1 +H 1 +M 1 +DETECTOR rec[-1] +Z_ERROR(0.1) 0 +H 0 +M 0 +OBSERVABLE_INCLUDE(0) rec[-1] +""" + + +def _postselected_expectations(text, shots): + circuit = Circuit(text) + mask = np.ones(circuit.num_detectors, dtype=bool) + sampler = _gpu(circuit.compile_detector_sampler(seed=9)) + _, values = _expectations(sampler, shots, mask) + detectors, observables = _cpu(circuit.compile_detector_sampler, seed=10).sample( + shots, separate_observables=True + ) + kept = observables[~(detectors & mask).any(axis=1), 0] + return values[:, 0], np.mean(1 - 2 * kept.astype(float)), len(kept) + + +def test_expectations_stay_unbiased_under_later_postselection(): + values, expected, kept = _postselected_expectations( + T_STATE_X_THEN_DETECTOR, 200_000 + ) + # NOTE: the expectation before the measurement would average to cos(π/4). + assert abs(values.mean() - expected) < 5 * np.sqrt(1 / len(values) + 1 / kept) + np.testing.assert_array_equal(np.abs(values), 1) + + +def test_expectations_are_exact_after_earlier_postselection(): + values, expected, kept = _postselected_expectations( + T_STATE_X_AFTER_DETECTOR, 200_000 + ) + np.testing.assert_allclose(np.abs(values), np.sqrt(0.5), atol=1e-5) + assert abs(values.mean() - expected) < 5 * np.sqrt(1 / len(values) + 1 / kept) + + +def test_clifford_observables_have_expectations_of_plus_or_minus_one(): + circuit = Circuit( + "R 0 1\nX_ERROR(0.3) 0\nCNOT 0 1\nM 0 1\nDETECTOR rec[-1]\nOBSERVABLE_INCLUDE(0) rec[-2]" + ) + sampler = _gpu(circuit.compile_detector_sampler(seed=8)) + packed, values = _expectations(sampler, 10_000) + bits = np.unpackbits(packed, axis=1, bitorder="little") + np.testing.assert_array_equal(values[:, 0], 1 - 2 * bits[:, 1].astype(np.float32)) + + +@pytest.mark.parametrize("reference", [False, True]) +def test_cuda_samples_programs_without_components(reference): + circuit = Circuit( + "R 0 1\nX_ERROR(0.3) 0\nCNOT 0 1\nM 0 1\nDETECTOR rec[-1]\nOBSERVABLE_INCLUDE(0) rec[-2]" + ) + kwargs = {"use_detector_reference_sample": reference} + expected = _cpu(circuit.compile_detector_sampler, seed=1).sample( + 40_000, append_observables=True, **kwargs + ) + got = _gpu(circuit.compile_detector_sampler(seed=2)).sample( + 40_000, append_observables=True, **kwargs + ) + assert_counts_match(got, expected) + measured = _gpu(circuit.compile_sampler(seed=3)).sample(40_000) + assert measured.shape == (40_000, 2) + np.testing.assert_array_equal(measured[:, 0], measured[:, 1]) + + +def test_stabilizer_rank_strategies_stay_on_jax(): + circuit = Circuit(T_STATE_X) + assert circuit.compile_detector_sampler(strategy="cat5")._shots is None + + +def _widest(circuit): + """Return the active width of the circuit's widest statevector component.""" + components = circuit.compile_sampler(seed=1)._program.components + return max(c.width for c in components if isinstance(c, StatevectorComponent)) + + +def _random_unitary_circuit(width, layers, seed=0, noise=True): + """Layers of random U3 pairs and CX gates, then every qubit measured. + + With ``noise``, X errors come first, so that every rotation reads the shot's + f-bits and none of them runs before the shots split (``common_prefix``). + """ + rng = np.random.default_rng(seed) + lines = ["X_ERROR(0.2) " + " ".join(map(str, range(width)))] if noise else [] + for _ in range(layers): + perm = rng.permutation(width) + for a, b in zip(perm[::2], perm[1::2], strict=False): + for q in (a, b): + t = rng.uniform(-1, 1, 3) + lines.append(f"U3({t[0]},{t[1]},{t[2]}) {q}") + lines.append(f"CX {a} {b}") + lines.append("M " + " ".join(map(str, range(width)))) + return Circuit("\n".join(lines)) + + +def _assert_matches_cpu(circuit, width, shots=200_000): + expected = _cpu(circuit.compile_sampler, seed=1).sample(shots) + got = _gpu(circuit.compile_sampler(seed=2)).sample(shots) + np.testing.assert_allclose( + got.mean(axis=0), expected.mean(axis=0), atol=5 * 0.5 / np.sqrt(shots / 2) + ) + # NOTE: pairs of bits, which a wrong shuffle partner would scramble. + for a, b in [(0, 1), (0, width - 1), (width // 2, width - 1)]: + assert_counts_match(got[:, [a, b]], expected[:, [a, b]]) + + +@pytest.mark.parametrize(("width", "layers"), [(6, 8), (8, 8), (9, 10), (11, 12)]) +def test_cuda_wide_components_match_the_cpu_executor(width, layers): + # NOTE: widths 6, 8, 9 and 11 run 2, 8, 16 and 32 threads per shot. + circuit = _random_unitary_circuit(width, layers) + assert _widest(circuit) == width + _assert_matches_cpu(circuit, width) + + +@pytest.mark.parametrize(("width", "layers"), [(6, 6), (8, 8), (11, 6)]) +def test_cuda_looping_tapes_match_the_cpu_executor(monkeypatch, width, layers): + # NOTE: a tape past the unroll budget loops, applying runs of rotations together. + monkeypatch.setattr(tape, "UNROLL_BUDGET", 0) + _assert_matches_cpu(_random_unitary_circuit(width, layers, seed=width), width) + + +@pytest.mark.parametrize("width", [6, 9, 11]) +def test_cuda_shots_start_after_the_common_prefix(width): + # NOTE: without noise the rotations before the first measurement run once, + # on the host; the shots load their amplitudes and run the rest. + circuit = _random_unitary_circuit(width, width + 2, seed=width, noise=False) + _assert_matches_cpu(circuit, width) + + +def test_components_wider_than_the_limit_run_on_the_cpu(): + circuit = _random_unitary_circuit(13, 12) + assert _widest(circuit) > tape.MAX_WIDTH + with warnings.catch_warnings(): + warnings.simplefilter("error") + sampler = circuit.compile_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + with pytest.raises(ValueError, match="more than the 11 that the GPU executor"): + circuit.compile_sampler(seed=1, device="cuda") + + +def _sampler_in_child(_): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + sampler = Circuit(T_STATE_X).compile_sampler(seed=1) + sampler.sample(10) + assert sampler.stats.device == "cpu" + try: + Circuit(T_STATE_X).compile_sampler(seed=1, device="cuda") + error = "" + except RuntimeError as raised: + error = str(raised) + return type(sampler._shots).__name__, [str(w.message) for w in caught], error + + +def test_processes_forked_after_the_gpu_was_used_sample_on_the_cpu(): + # NOTE: CUDA cannot be used in a child forked after its parent initialized it. + _gpu(Circuit(T_STATE_X).compile_sampler(seed=1)).sample(10) + with multiprocessing.get_context("fork").Pool(1) as pool: + name, messages, error = pool.map(_sampler_in_child, [0])[0] + assert name == "ShotSampler" + assert messages == [] + assert "forked" in error + + +def test_noise_of_tiny_probabilities_fires_at_its_rate(): + # NOTE: 24-bit skips alone fire every segment at least 2^-24 of the time, + # here ~90 faults instead of ~0.4. + n, p = 1000, 1e-12 + text = "\n".join(f"X_ERROR({p * (1 + i / n)!r}) {i}" for i in range(n)) + sampler = _gpu( + Circuit(f"{text}\nM " + " ".join(map(str, range(n)))).compile_sampler(seed=1) + ) + # NOTE: faults are rare enough that each lands in a byte of its own. + faults = sum( + np.count_nonzero(sampler.sample(1 << 22, bit_packed=True)) for _ in range(64) + ) + assert faults <= 5 + + +@pytest.mark.parametrize("p0", [3e-8, 1e-6]) +def test_rare_outcome_zero_is_drawn_at_its_rate(p0): + # NOTE: w1 / total rounds to a multiple of 2^-24 below 1 in float32; the draw + # must use w0, as the CPU's ``_cut`` does. + cp = runtime.cupy() + source = codegen.PRELUDE + r""" +extern "C" __global__ void zeros(float w0, unsigned long long n, unsigned long long* count) { + unsigned long long local = 0; + for (u64 i = blockIdx.x * (u64)blockDim.x + threadIdx.x; i < n; i += (u64)gridDim.x * blockDim.x) + local += !draw_one(mix64(i * GOLDEN + 12345), w0, 1.f, w0 + 1.f); + atomicAdd(count, local); +} +""" + kernel = cp.RawModule( + code=source, options=(runtime.arch(), "-std=c++17") + ).get_function("zeros") + count = cp.zeros(1, dtype=cp.uint64) + n = 1 << 30 + kernel((1024,), (256,), (np.float32(p0), np.uint64(n), count)) + expected = p0 * n + assert abs(int(count.get()[0]) - expected) <= 5 * np.sqrt(expected) + 1 + + +@pytest.mark.parametrize(("probability", "shots"), [(2e-8, 1 << 30), (1e-6, 1 << 26)]) +def test_cuda_draws_small_outcome_probabilities_at_their_rate(probability, shots): + # NOTE: a 24-bit draw takes any positive probability at least 2^-24 = 6e-8 of + # the time; float32 rounding gives certain outcomes such probabilities. 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 _out_of_memory_above(monkeypatch, sampler, limit): + """Make ``sampler``'s GPU calls of batches above ``limit`` run out of memory.""" + shots = _gpu_shots(sampler) + sample_accepted = shots.sample_accepted + batches = [] + + def fake(n, *, batch_size=None, **kwargs): + batches.append(batch_size) + if batch_size is None or batch_size > limit: + cp = runtime.cupy() + raise cp.cuda.memory.OutOfMemoryError(1 << 40, 0) + return sample_accepted(n, batch_size=batch_size, **kwargs) + + monkeypatch.setattr(shots, "sample_accepted", fake) + return batches + + +def test_calls_out_of_gpu_memory_retry_in_smaller_batches(monkeypatch): + sampler = _gpu(Circuit(GHZ_WITH_DIRECT_DETECTOR).compile_detector_sampler(seed=1)) + batches = _out_of_memory_above(monkeypatch, sampler, 1 << 20) + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert sampler.sample(1000).shape == (1000, 3) + assert batches == [1 << 22, 1 << 21, 1 << 20] + # NOTE: later calls start from the batch that fit. + sampler.sample(1000) + assert batches[-1] == 1 << 20 + + +def test_calls_out_of_gpu_memory_fall_back_to_the_cpu(monkeypatch): + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + sampler = _gpu(circuit.compile_detector_sampler(seed=1)) + _out_of_memory_above(monkeypatch, sampler, 0) + with warnings.catch_warnings(): + warnings.simplefilter("error") + got = sampler.sample(20_000, append_observables=True) + expected = _cpu(circuit.compile_detector_sampler, seed=2).sample( + 20_000, append_observables=True + ) + assert_counts_match(got, expected) + required = _gpu(circuit.compile_detector_sampler(seed=3, device="cuda")) + _out_of_memory_above(monkeypatch, required, 0) + with pytest.raises(MemoryError, match="the GPU ran out of memory"): + required.sample(10, append_observables=True) + + +def test_samplers_out_of_gpu_memory_at_setup_run_on_the_cpu(monkeypatch): + cp = runtime.cupy() + + def out_of_memory(*args, **kwargs): + raise cp.cuda.memory.OutOfMemoryError(1 << 40, 0) + + monkeypatch.setattr(CudaShotSampler, "__init__", out_of_memory) + with warnings.catch_warnings(): + warnings.simplefilter("error") + sampler = Circuit(T_STATE_X).compile_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + with pytest.raises(MemoryError, match="Cannot sample on CUDA"): + Circuit(T_STATE_X).compile_sampler(seed=1, device="cuda") + + +def test_kernels_wait_for_the_tables_they_read(): + # NOTE: the tables are copied on the current (legacy default) stream, which + # the sampler's non-blocking streams do not wait for; a kernel that keeps + # that stream busy delays the copies, so the first call would read them + # before they land. + cp = runtime.cupy() + spin = cp.RawKernel( + r""" + extern "C" __global__ void spin(long long cycles) { + long long start = clock64(); + while (clock64() - start < cycles) {} + } + """, + "spin", + ) + circuit = Circuit(GHZ_WITH_DIRECT_DETECTOR) + expected = _cpu(circuit.compile_detector_sampler, seed=1).sample( + 40_000, append_observables=True + ) + for seed in range(8): + spin((1,), (1,), (np.int64(200_000_000),)) + sampler = _gpu(circuit.compile_detector_sampler(seed=seed)) + got = sampler.sample(40_000, append_observables=True) + assert_counts_match(got, expected) + + +def test_device_cpu_keeps_the_samplers_on_the_cpu(monkeypatch): + # NOTE: the environment no longer chooses the device. + monkeypatch.setenv("TSIM_DEVICE", "cpu") + _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 + assert sampler.stats.device == "cpu" + + +def test_the_repr_says_where_a_sampler_samples(): + circuit = Circuit(T_STATE_X) + gpu = _gpu(circuit.compile_detector_sampler(seed=1)) + assert "samples on cuda" in repr(gpu) and gpu.stats.device == "cuda" + cpu = _cpu(circuit.compile_detector_sampler, seed=1) + assert "samples on cpu" in repr(cpu) and cpu.stats.device == "cpu" + + +def test_programs_without_shot_words_sample_on_the_cpu(): + # NOTE: a detector sampler of a circuit without detectors, observables or noise. + for device in ("auto", "cuda"): + with warnings.catch_warnings(): + warnings.simplefilter("error") + sampler = Circuit("R 0\nH 0\nT 0\nM 0").compile_detector_sampler( + seed=1, device=device + ) + assert type(sampler._shots) is ShotSampler + assert sampler.sample(10).shape == (10, 0) + + +def test_noise_beyond_the_gpu_tables_samples_on_the_cpu(): + # 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 warnings.catch_warnings(): + warnings.simplefilter("error") + sampler = Circuit(text).compile_detector_sampler(seed=1) + assert type(sampler._shots) is ShotSampler + assert sampler.sample(10).shape == (10, 9) + with pytest.raises(ValueError, match="channels of up to 127 outcomes"): + Circuit(text).compile_detector_sampler(seed=1, device="cuda") + + +_FORKED = {} + + +def _sample_in_child(_): + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + shape = _FORKED["sampler"].sample(100).shape + return shape, [str(w.message) for w in caught], _FORKED["sampler"].stats.device + + +def test_samplers_compiled_before_a_fork_sample_on_the_cpu_in_the_child(): + sampler = _gpu(Circuit(GHZ_WITH_DIRECT_DETECTOR).compile_detector_sampler(seed=1)) + sampler.sample(10) + # NOTE: the child gets the sampler by fork, not by pickle. + _FORKED["sampler"] = sampler + try: + with multiprocessing.get_context("fork").Pool(1) as pool: + shape, messages, device = pool.map(_sample_in_child, [0])[0] + finally: + _FORKED.clear() + assert shape == (100, 3) + assert messages == [] + assert device == "cpu" + assert sampler.stats.device == "cuda" + + +def test_gpu_samplers_pickle_and_continue_their_streams(): + sampler = _gpu(Circuit(GHZ_WITH_DIRECT_DETECTOR).compile_detector_sampler(seed=26)) + sampler.sample(1000) + loaded = pickle.loads(pickle.dumps(sampler)) # noqa: S301 + assert isinstance(loaded._shots, CudaShotSampler) + np.testing.assert_array_equal(loaded.sample(5000), sampler.sample(5000)) + required = Circuit(T_STATE_X).compile_detector_sampler(seed=26, device="cuda") + assert _gpu_shots(pickle.loads(pickle.dumps(required)))._required # noqa: S301 + + +def test_an_empty_mask_is_not_taken_for_no_mask(): + sampler = _gpu(Circuit(T_STATE_X).compile_detector_sampler(seed=1)) + assert sampler.sample(1000, append_observables=True).shape == (1000, 1) + empty = np.zeros(0, dtype=bool) + kept = sampler.sample(1000, append_observables=True, postselection_mask=empty) + assert kept.shape == (1000, 1) + + +def test_cuda_samples_rows_too_wide_for_shared_memory(): + # NOTE: 12000 detectors are 188 words: the noise goes into device memory, and + # the rows are gathered without staging them in shared memory. + n = 12000 + qubits = " ".join(map(str, range(n))) + circuit = Circuit( + "\n".join( + [ + f"R {qubits}", + f"X_ERROR(0.05) {qubits}", + f"M {qubits}", + *(f"DETECTOR rec[{i - n}]" for i in range(n)), + "OBSERVABLE_INCLUDE(0) rec[-1] rec[-2]", + ] + ) + ) + sampler = _gpu(circuit.compile_detector_sampler(seed=31)) + shots = _gpu_shots(sampler) + assert not shots._noise.shared_words and shots._gather_shots == 0 + got = sampler.sample(20_000, append_observables=True) + np.testing.assert_allclose(got[:, :n].mean(), 0.05, atol=5e-4) + np.testing.assert_allclose(got[:, -1].mean(), 2 * 0.05 * 0.95, atol=0.01) + np.testing.assert_array_equal(got[:, -1], got[:, -2] ^ got[:, -3]) + packed = sampler.sample(1000, append_observables=True, bit_packed=True) + assert packed.shape == (1000, -(-(n + 1) // 8)) + assert not np.unpackbits(packed, axis=1, bitorder="little")[:, n + 1 :].any() + + +def test_concurrent_calls_on_one_sampler_are_serialized(): + sampler = _gpu(Circuit(GHZ_WITH_DIRECT_DETECTOR).compile_detector_sampler(seed=1)) + mask = np.array([False, False, True]) + results, errors = [], [] + + def work(k): + try: + for _ in range(20): + got = sampler.sample(20_000, postselection_mask=mask if k % 2 else None) + results.append((k % 2, got)) + except Exception as error: + errors.append(error) + + threads = [threading.Thread(target=work, args=(k,)) for k in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert not errors + for masked, got in results: + assert got.shape[1] == 3 + if masked: + assert not got[:, 2].any() + else: + # NOTE: detector 2 fires at rate 0.2. + assert abs(got[:, 2].mean() - 0.2) < 0.02 + + +def test_a_kernel_that_does_not_compile_samples_on_the_cpu(monkeypatch): + circuit = Circuit(T_STATE_X) + sampler = _gpu(circuit.compile_sampler(seed=1)) + compiled = [] + + def broken(*args, **kwargs): + compiled.append(1) + return 'extern "C" __global__ void tape(' + + monkeypatch.setattr(kernels, "source", broken) + with warnings.catch_warnings(): + warnings.simplefilter("error") + got = sampler.sample(20_000) + assert sampler.stats.device == "cpu" + got = np.concatenate([got, sampler.sample(20_000)]) + assert len(compiled) == 1 + expected = _cpu(circuit.compile_sampler, seed=2).sample(40_000) + assert_counts_match(got, expected) + required = _gpu(circuit.compile_sampler(seed=3, device="cuda")) + with pytest.raises(runtime.CompileError): + required.sample(10) + + +def test_rare_outcome_of_a_channel_is_drawn_at_its_rate(): + # NOTE: 1000 channels (distinct, so Tsim keeps them apart) fire half the time + # and then flip detector 0 with chance 6e-9: with 24-bit keep thresholds it + # rounded to zero; ~100 flips are expected here. + k, q, calls = 10, 3e-9, 34 + lines = ["R " + " ".join(map(str, range(k + 1)))] + subsets = [ + s for r in range(1, k + 1) for s in itertools.combinations(range(1, k + 1), r) + ][:1000] + for subset in subsets: + lines.append("E(0.5) " + " ".join(f"X{i}" for i in subset)) + lines.append(f"ELSE_CORRELATED_ERROR({q / 0.5}) X0") + lines.append("M " + " ".join(map(str, range(k + 1)))) + lines += [f"DETECTOR rec[{i - k - 1}]" for i in range(k + 1)] + sampler = _gpu(Circuit("\n".join(lines)).compile_detector_sampler(seed=1)) + flips = sum( + int(np.count_nonzero(sampler.sample(1_000_000, bit_packed=True)[:, 0] & 1)) + for _ in range(calls) + ) + expected = len(subsets) * q * calls * 1_000_000 + assert abs(flips - expected) <= 5 * np.sqrt(expected) + + +def test_rotation_with_a_branch_float32_cannot_resolve_samples_on_the_cpu(): + circuit = Circuit("R 0\nH 0\nR_Z(1e-5) 0\nH 0\nM 0\nDETECTOR rec[-1]") + with warnings.catch_warnings(): + warnings.simplefilter("error") + sampler = circuit.compile_detector_sampler(seed=1) + assert sampler.stats.device == "cpu" + message = r"angle is 1\.0e-05 π from a Clifford angle.*6\.4e-05 π.*device='cpu'" + with pytest.raises(ValueError, match=message): + circuit.compile_detector_sampler(seed=1, device="cuda") + + +@pytest.mark.parametrize( + ("detectors", "observables"), [(20, 1), (21, 3), (7, 9), (16, 8), (130, 5)] +) +def test_cuda_selects_every_layout_of_columns(detectors, observables): + # NOTE: the observables start inside a byte (20, 21, 7, 130), end inside one + # (20, 130) or span two bytes (7); 16 + 8 lies on byte boundaries. Samplers of + # one seed draw the same shots; the reference is every output as booleans, + # in order. + n = detectors + observables + lines = [f"X_ERROR(0.3) {' '.join(map(str, range(n)))}", "H 0", "T 0", "H 0"] + lines.append(f"M {' '.join(map(str, range(n)))}") + lines += [f"DETECTOR rec[-{n - i}]" for i in range(detectors)] + lines += [ + f"OBSERVABLE_INCLUDE({j}) rec[-{observables - j}]" for j in range(observables) + ] + circuit = Circuit("\n".join(lines)) + mask = np.zeros(detectors, dtype=bool) + mask[0] = True + for postselection_mask in (None, mask): + + def sample(layout, bit_packed=False, postselection_mask=postselection_mask): + sampler = _gpu(circuit.compile_detector_sampler(seed=25)) + return sampler.sample( + 20_000, + separate_observables=layout == "separate", + prepend_observables=layout == "prepend", + append_observables=layout == "append", + bit_packed=bit_packed, + postselection_mask=postselection_mask, + ) + + every = sample("append") + assert isinstance(every, np.ndarray) + found, events = every[:, :detectors], every[:, detectors:] + layouts = { + "separate": (found, events), + "prepend": (np.hstack([events, found]),), + "detectors": (found,), + } + for layout, expected in layouts.items(): + for bit_packed in (True, False): + got = sample(layout, bit_packed) + got = got if isinstance(got, tuple) else (got,) + for array, bools in zip(got, expected, strict=True): + if bit_packed: + bools = np.packbits(bools, axis=1, bitorder="little") + else: + assert array.dtype == np.bool_ and array.flags.c_contiguous + np.testing.assert_array_equal(array, bools) + + +@pytest.mark.parametrize("text", [GHZ_WITH_DIRECT_DETECTOR, NOISY_DETECTORS]) +def test_cuda_unpacks_booleans_as_it_packs_rows(text): + # NOTE: samplers of one seed draw the same shots; 4 and 152 outputs, neither a + # multiple of 8. + circuit = Circuit(text) + outputs = circuit.num_detectors + circuit.num_observables + mask = np.zeros(circuit.num_detectors, dtype=bool) + mask[0] = True + for postselection_mask in (None, mask): + + def sample(bit_packed, postselection_mask=postselection_mask): + sampler = _gpu(circuit.compile_detector_sampler(seed=23)) + return sampler.sample( + 30_000, + append_observables=True, + bit_packed=bit_packed, + postselection_mask=postselection_mask, + ) + + packed, bools = sample(True), sample(False) + assert bools.dtype == np.bool_ and bools.flags.c_contiguous + unpacked = np.unpackbits(packed, axis=1, count=outputs, bitorder="little") + np.testing.assert_array_equal(bools, unpacked.astype(bool)) + detectors = _gpu(circuit.compile_detector_sampler(seed=23)).sample( + 30_000, postselection_mask=postselection_mask + ) + assert detectors.flags.c_contiguous + np.testing.assert_array_equal(detectors, bools[:, : circuit.num_detectors]) + + +def test_returned_rows_are_not_reused_while_alive(): + circuit = Circuit(NOISY_DETECTORS) + sampler = _gpu(circuit.compile_detector_sampler(seed=3)) + first = _packed(sampler, 200_000) + kept = first.copy() + second = _packed(sampler, 200_000) + np.testing.assert_array_equal(first, kept) + assert not np.shares_memory(first, second) diff --git a/test/unit/test_sampler.py b/test/unit/test_sampler.py index f1dc53e1..bf882587 100644 --- a/test/unit/test_sampler.py +++ b/test/unit/test_sampler.py @@ -829,22 +829,29 @@ def test_device_cpu_samples_on_the_cpu(compile): def test_unknown_devices_raise(): - with pytest.raises(ValueError, match="device must be 'auto', 'cpu' or 'metal'"): + with pytest.raises( + ValueError, match="device must be 'auto', 'cpu', 'metal' or 'cuda'" + ): Circuit(NOISY_T_STATE).compile_sampler(device="gpu") # type: ignore[arg-type] -@pytest.mark.parametrize("device", ["cpu", "metal"]) +@pytest.mark.parametrize("device", ["cpu", "metal", "cuda"]) def test_stabilizer_rank_strategies_take_only_the_auto_device(device): with pytest.raises(ValueError, match="needs strategy='active_statevector'"): Circuit(NOISY_T_STATE).compile_sampler(strategy="cat5", device=device) -def test_metal_device_raises_where_metal_cannot_run(monkeypatch): - from tsim.statevector.metal import runtime +@pytest.mark.parametrize( + ("device", "name", "reason"), + [("metal", "Metal", "no Apple GPU"), ("cuda", "CUDA", "no CUDA device")], +) +def test_gpu_devices_raise_where_the_gpu_cannot_run(monkeypatch, device, name, reason): + from tsim.statevector import cuda, metal - monkeypatch.setattr(runtime, "unavailable", lambda: "no Apple GPU") - with pytest.raises(RuntimeError, match="Cannot sample on Metal: no Apple GPU"): - Circuit(NOISY_T_STATE).compile_detector_sampler(device="metal") + runtime = {"metal": metal.runtime, "cuda": cuda.runtime}[device] + monkeypatch.setattr(runtime, "unavailable", lambda: reason) + with pytest.raises(RuntimeError, match=f"Cannot sample on {name}: {reason}"): + Circuit(NOISY_T_STATE).compile_detector_sampler(device=device) with warnings.catch_warnings(): warnings.simplefilter("error") sampler = Circuit(NOISY_T_STATE).compile_detector_sampler(seed=1) diff --git a/uv.lock b/uv.lock index 24aeb1c8..0581301a 100644 --- a/uv.lock +++ b/uv.lock @@ -137,12 +137,21 @@ dependencies = [ ] [package.optional-dependencies] +cuda = [ + { name = "cupy-cuda13x", extra = ["ctk"] }, +] cuda12 = [ { name = "jax", extra = ["cuda12"] }, ] cuda13 = [ { name = "jax", extra = ["cuda13"] }, ] +jax-cuda12 = [ + { name = "jax", extra = ["cuda12"] }, +] +jax-cuda13 = [ + { name = "jax", extra = ["cuda13"] }, +] [package.dev-dependencies] dev = [ @@ -183,11 +192,14 @@ doc = [ [package.metadata] requires-dist = [ + { name = "cupy-cuda13x", extras = ["ctk"], marker = "extra == 'cuda'", specifier = ">=14.0" }, { name = "equinox", specifier = ">=0.11.5" }, { name = "ipywidgets", specifier = ">=8.0.0" }, { name = "jax", specifier = ">=0.4.38" }, { name = "jax", extras = ["cuda12"], marker = "extra == 'cuda12'", specifier = ">=0.4.38" }, + { name = "jax", extras = ["cuda12"], marker = "extra == 'jax-cuda12'", specifier = ">=0.4.38" }, { name = "jax", extras = ["cuda13"], marker = "extra == 'cuda13'", specifier = ">=0.7.0" }, + { name = "jax", extras = ["cuda13"], marker = "extra == 'jax-cuda13'", specifier = ">=0.7.0" }, { name = "lxml", specifier = ">=5.0.0" }, { name = "numba", specifier = ">=0.61" }, { name = "numpy", specifier = ">=2.0" }, @@ -197,7 +209,7 @@ requires-dist = [ { name = "stim-timeline-text", specifier = ">=0.1.0" }, { name = "zxparam", specifier = ">=0.1.0" }, ] -provides-extras = ["cuda12", "cuda13"] +provides-extras = ["cuda", "jax-cuda12", "jax-cuda13", "cuda12", "cuda13"] [package.metadata.requires-dev] dev = [ @@ -537,6 +549,87 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9e/ee/a4cf96b8ce1e566ed238f0659ac2d3f007ed1d14b181bcb684e19561a69a/coverage-7.13.5-py3-none-any.whl", hash = "sha256:34b02417cf070e173989b3db962f7ed56d2f644307b2cf9d5a0f258e13084a61", size = 211346, upload-time = "2026-03-17T10:33:15.691Z" }, ] +[[package]] +name = "cuda-pathfinder" +version = "1.8.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6a/9e/08555797c979322def080814686417ea93f66d38f5c48ebab6867a85a21e/cuda_pathfinder-1.8.6-py3-none-any.whl", hash = "sha256:bf6b3b06323cf95a11d671818392da5da23ddbbb4dfc4b77ca10622783c25980", size = 62603, upload-time = "2026-11-05T21:22:51.229Z" }, +] + +[[package]] +name = "cuda-toolkit" +version = "13.5.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e6/4a/7d75cd156b6651831473a60f5a4dfa96b5b6f1d04ccb47fcac1f8a08456e/cuda_toolkit-13.5.0-py2.py3-none-any.whl", hash = "sha256:d5b168b8284af2bad07f313b4d9fd5af0d1c7d6a26b7b019f735784ef718a737", size = 2706, upload-time = "2026-11-04T23:44:09.608Z" }, +] + +[package.optional-dependencies] +cublas = [ + { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cudart = [ + { name = "nvidia-cuda-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cufft = [ + { name = "nvidia-cufft", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +curand = [ + { name = "nvidia-curand", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cusolver = [ + { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-cusolver", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +cusparse = [ + { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] +nvrtc = [ + { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'ARM64' and sys_platform == 'win32')" }, +] + +[[package]] +name = "cupy-cuda13x" +version = "14.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder" }, + { name = "numpy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8d/3c/4716e285f8459a5707ef39d84e9a7547bf16f87c583ac4789c0763d6d986/cupy_cuda13x-14.2.0-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:a22e2f17c97940cd566f351f186d8fe04285280d052ca1512ff223febbe9979c", size = 74074375, upload-time = "2026-08-20T02:41:04.327Z" }, + { url = "https://files.pythonhosted.org/packages/ac/cb/41b6d8f1880e6f3128f8084b7622b1a3ca2d615ab77af9ef313b4aafed3d/cupy_cuda13x-14.2.0-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:f94843e54500a9c8c169135f79572bfaba946efe4f7178e834c227280c1b0c12", size = 69969199, upload-time = "2026-08-20T02:41:07.842Z" }, + { url = "https://files.pythonhosted.org/packages/f8/eb/25b8c1a415ade6479ed684c5cbff1e6f5002a92cbf26ef826df56a40cf69/cupy_cuda13x-14.2.0-cp311-cp311-win_amd64.whl", hash = "sha256:3893acd8d4dc78d113736045eaa306e496ee7bd8253bd959513ba28c91ed396d", size = 36082966, upload-time = "2026-08-20T02:41:10.723Z" }, + { url = "https://files.pythonhosted.org/packages/10/72/ad885cc21a21208b94044b34f50a840394b2e9dbb6111ed43959366aa6eb/cupy_cuda13x-14.2.0-cp311-cp311-win_arm64.whl", hash = "sha256:e0ad42a1ef35549282e2087a7c1e6ec8b7aa624a95e8fc0a60da777ca87e71de", size = 35136070, upload-time = "2026-09-01T06:41:38.284Z" }, + { url = "https://files.pythonhosted.org/packages/dd/c8/68e07ff959f6a9f93169e37a64a4c24223487102b8625191e4a566cab465/cupy_cuda13x-14.2.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:0a9b6538079c5151dcd99cf44b94cc3f1242315b7de0498b4619cef9b2a7c2c8", size = 73382974, upload-time = "2026-08-20T02:41:13.648Z" }, + { url = "https://files.pythonhosted.org/packages/7e/bf/c4c25e97552320451afdcd459d4b6e0b715459594e998b23051c06db64c3/cupy_cuda13x-14.2.0-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:9aafadbe29db3d4285026044aafeb28f2a3de4bd0df70718ec6892b7d27a4ded", size = 69439247, upload-time = "2026-08-20T02:41:17.375Z" }, + { url = "https://files.pythonhosted.org/packages/5e/88/0a4d2d075485c4a82bfcd4725300ff87d17de36dce27b36f720d286f0f70/cupy_cuda13x-14.2.0-cp312-cp312-win_amd64.whl", hash = "sha256:9e8bc8a7693c5c212cd1450d728f47e1f05aa07ae003de30415893d0daf7d243", size = 35986921, upload-time = "2026-08-20T02:41:20.253Z" }, + { url = "https://files.pythonhosted.org/packages/a4/ab/8a756f553397b8509906d4a1d8083d92d3d21cbde115504e7a5db97de61f/cupy_cuda13x-14.2.0-cp312-cp312-win_arm64.whl", hash = "sha256:0dae6d788752c987c1cd87adb903640620820feb895339e8bc51a01fba0fb804", size = 35049803, upload-time = "2026-09-01T06:41:43.442Z" }, + { url = "https://files.pythonhosted.org/packages/3c/02/3cb39774c71a16275c50b5d1e92ddcefa948ff1e69f464249b70a79b0abe/cupy_cuda13x-14.2.0-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:449af7f608dc70ce64e8c45cf5a6ea0bbef51623e7c806b1a43445ae3c76e1e0", size = 72969322, upload-time = "2026-08-20T02:41:23.311Z" }, + { url = "https://files.pythonhosted.org/packages/4e/27/216370ae8c39fa9835645be3c27513ab72f52cf4f51faf58d799fba981b6/cupy_cuda13x-14.2.0-cp313-cp313-manylinux2014_x86_64.whl", hash = "sha256:5b79a606f6639d9bed74ef7072f36756a95a804c86d7adb367b39307ac771600", size = 69048238, upload-time = "2026-08-20T02:41:26.831Z" }, + { url = "https://files.pythonhosted.org/packages/5d/75/9df3c545baad4ff2e384db1ff6c8005fa4040492c4ee5672ca8b8711d29d/cupy_cuda13x-14.2.0-cp313-cp313-win_amd64.whl", hash = "sha256:f9a2a55abf889f6df68b5f62cd5f3ec4a2f305ac57ed09becd52e16201edcd86", size = 35965490, upload-time = "2026-08-20T02:41:29.943Z" }, + { url = "https://files.pythonhosted.org/packages/5c/3d/ee85ce39c2f32e8561c81c507df458a914b45d9a0eb1a54e5396c2cfd632/cupy_cuda13x-14.2.0-cp313-cp313-win_arm64.whl", hash = "sha256:5ddaf38291f60ec20e8c7889b28e339bf17ad1b4df768b8bb5e6c9f9e445efbb", size = 35029087, upload-time = "2026-09-01T06:41:48.593Z" }, + { url = "https://files.pythonhosted.org/packages/e5/70/8c80c9e7010193afed8ccb7110d6173aa23bf3a6ce2fc2450244b0fa50a1/cupy_cuda13x-14.2.0-cp314-cp314-manylinux2014_aarch64.whl", hash = "sha256:7ac57c4e16c62f265b5ecc2b45d6bbafe9dc70a99b849ed6cfb8d16f32204589", size = 72840193, upload-time = "2026-08-20T02:41:33.146Z" }, + { url = "https://files.pythonhosted.org/packages/c2/b8/4f4c4f34fc31ab8d136ed965a919505297974629d9852c50e15bfd616281/cupy_cuda13x-14.2.0-cp314-cp314-manylinux2014_x86_64.whl", hash = "sha256:ff0bdebd1b43c0c6db53095784c787c4e4eae671356cf521c4b7482ed78a1a7e", size = 68422828, upload-time = "2026-08-20T02:41:36.527Z" }, + { url = "https://files.pythonhosted.org/packages/7d/5c/58941ece5ed0b2cc8304f8ba81b5e8334759d409306475d53c52ac99809b/cupy_cuda13x-14.2.0-cp314-cp314-win_amd64.whl", hash = "sha256:b99a8bf9d5391954c0c08675a8cb39219923fe2ee249faa3546ebbf75783a360", size = 36110882, upload-time = "2026-08-20T02:41:39.837Z" }, + { url = "https://files.pythonhosted.org/packages/e2/c5/cc7d1a9b764b3f828e09b48a73d0fdeaad15ecd2c7e50507f6d5fcf5fd1a/cupy_cuda13x-14.2.0-cp314-cp314-win_arm64.whl", hash = "sha256:c61c59ba665948c0a855926997ee96026420b628816b2e406998ec797caddb22", size = 35184806, upload-time = "2026-09-01T06:41:58.317Z" }, + { url = "https://files.pythonhosted.org/packages/6f/54/a582ac70e9e75d3fad7ae073051bcd50dd128a1f66dd275291dc5f48d997/cupy_cuda13x-14.2.0-cp314-cp314t-manylinux2014_aarch64.whl", hash = "sha256:e8bb453551deced08ec5c4576da61ebf1e2b4a45aeff48e09f47ec332f7cad3d", size = 73336756, upload-time = "2026-08-20T02:41:42.718Z" }, + { url = "https://files.pythonhosted.org/packages/f9/aa/222151ba5061d7a01efe4009efad9f9a7164e15bb2a494d9aa86efc97d4f/cupy_cuda13x-14.2.0-cp314-cp314t-manylinux2014_x86_64.whl", hash = "sha256:313500c5d415f65d388e84a9d2eb7f07f97c23c85dc5785f986e75d46d9c9693", size = 68792663, upload-time = "2026-08-20T02:41:46.079Z" }, + { url = "https://files.pythonhosted.org/packages/d4/20/13dd37dbeaec879c8348fde2bfca67d96daa56a6e63e4a5083ff61660e53/cupy_cuda13x-14.2.0-cp314-cp314t-win_amd64.whl", hash = "sha256:c517fc55502a7e3d1a8b9310abad95cec61dded5a38d328fe0220bcf7b036268", size = 37004845, upload-time = "2026-08-20T02:41:48.963Z" }, + { url = "https://files.pythonhosted.org/packages/d1/5b/b1367f79abd90a259dbb7f0603dd162c03c3a76ec1d02c2527b0c79d46ba/cupy_cuda13x-14.2.0-cp314-cp314t-win_arm64.whl", hash = "sha256:83885db963eb0b95a0757b4907a37115a1d2aebfa9cf43b99ebac8db35713d37", size = 35485280, upload-time = "2026-09-01T06:41:53.618Z" }, +] + +[package.optional-dependencies] +ctk = [ + { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "curand", "cusolver", "cusparse", "nvrtc"] }, +] + [[package]] name = "cycler" version = "0.12.1" @@ -2195,11 +2288,16 @@ wheels = [ [[package]] name = "nvidia-cublas" -version = "13.3.0.5" +version = "13.10.0.20" source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cuda-nvrtc" }, +] wheels = [ - { url = "https://files.pythonhosted.org/packages/e8/5c/08177998e1234459e46b2cdad73738b5516f84b8fa28a8379c678b95c6c0/nvidia_cublas-13.3.0.5-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:48308e7f44feb337ca24d95efdceeac5703fb5dcf9bfb0c23f1eb48015fdd8a1", size = 505057164, upload-time = "2026-03-09T09:43:27.897Z" }, - { url = "https://files.pythonhosted.org/packages/3c/7c/ae5d1751819acff18b0fac29c0a4e93d06d36cfabebe36365ddacc7c32a9/nvidia_cublas-13.3.0.5-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:366568e2dc59e6fe71ffd179f9f2a38b8b2772aed626320a64008651b1e72974", size = 403287501, upload-time = "2026-03-09T09:47:05.046Z" }, + { url = "https://files.pythonhosted.org/packages/b4/12/16628161ed8c109c8c73b38167df50cbc630a8a6e6548e2534d461526c3c/nvidia_cublas-13.10.0.20-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:dcc6d63884c53d07e267ce7e6feac17b2538cebaafec5db865736c3e28e4dc84", size = 444289903, upload-time = "2026-11-04T23:27:26.095Z" }, + { url = "https://files.pythonhosted.org/packages/5c/e3/a5a992c3cb1c49201cd4855345dd50ec9e9c075195a0e1e534ff7e914ca5/nvidia_cublas-13.10.0.20-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c5f3418c0a845232a8cabc58f6fa52cfae7d8bd3680d94c3907aca710fd48d70", size = 560089875, upload-time = "2026-11-04T23:26:43.825Z" }, + { url = "https://files.pythonhosted.org/packages/12/39/e1b1b4adb8b171894d9a89035df0596d448cf73dc403802d24dbe30cc0be/nvidia_cublas-13.10.0.20-py3-none-win_amd64.whl", hash = "sha256:c310b749eea855b88fafb906de08d6819b87216c6711640771764fea5dcd8546", size = 427284863, upload-time = "2026-11-04T23:36:18.775Z" }, + { url = "https://files.pythonhosted.org/packages/4d/89/c875dcd88abc508a98a725214a30b2c001406095855154f32ee3f0946409/nvidia_cublas-13.10.0.20-py3-none-win_arm64.whl", hash = "sha256:05929d6a64af9edce1951ff82a514d630fc1a2369bcfe1ab195f5e85266dbfed", size = 279793527, upload-time = "2026-11-04T23:41:15.348Z" }, ] [[package]] @@ -2284,11 +2382,13 @@ wheels = [ [[package]] name = "nvidia-cuda-nvrtc" -version = "13.2.51" +version = "13.5.37" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5e/21/2fd0aa5a03a8c71962d281084ac44ae7b3b6690d6163ffd7d6486fdb7aa8/nvidia_cuda_nvrtc-13.2.51-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:c88076f32cbbd26e7ebd2107d4b093dd8667e2a90b23b3273d028f3daf574d2e", size = 47019178, upload-time = "2026-03-09T09:38:11.629Z" }, - { url = "https://files.pythonhosted.org/packages/27/ce/ef85d9b59e3fcb0e44041d72de857bf4e87f772d747c603018acbb4addb1/nvidia_cuda_nvrtc-13.2.51-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8d4ddd23dbd4878eea5d970a70631d2668381fc2bdc5ec16cd5faa8c2db24d03", size = 44754448, upload-time = "2026-03-09T09:37:38.021Z" }, + { url = "https://files.pythonhosted.org/packages/a9/54/db4e7f650746265f61c05cf1f8f556dd8c6739b58c823fbc9a7d6fda3247/nvidia_cuda_nvrtc-13.5.37-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:86b4d68058807988b7d9af5b44917a0601e9c337f35ccd657bfb3c759ff8d8bc", size = 58185930, upload-time = "2026-11-04T23:23:57.537Z" }, + { url = "https://files.pythonhosted.org/packages/4e/9f/e8321dfd9dcd923c864213fed3da815da474800ead891a99b6322bdd2039/nvidia_cuda_nvrtc-13.5.37-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:36e4fa2a8ada293f5693d308e531c196ab6f91cf7a0ff479bf96a537d43a6096", size = 56419462, upload-time = "2026-11-04T23:23:44.768Z" }, + { url = "https://files.pythonhosted.org/packages/9f/6c/241c5155199a6de55bad8862213ed64b487dc9526aae817e4af31afea1a4/nvidia_cuda_nvrtc-13.5.37-py3-none-win_amd64.whl", hash = "sha256:e8e6c10b251385707ae3abd7b4ee118de6736fd7678d0f0f4eab4d759f679c21", size = 49248991, upload-time = "2026-11-04T23:35:13.589Z" }, + { url = "https://files.pythonhosted.org/packages/d9/62/b858b2e6da4d9f4ab2744dcd337dc31c75bc4db4ad35543a04ee07bb3eab/nvidia_cuda_nvrtc-13.5.37-py3-none-win_arm64.whl", hash = "sha256:d47388d9054a9f16c705d86e59d7e6c57e85beece36d23d1657833531fa913aa", size = 45119472, upload-time = "2026-11-04T23:40:19.254Z" }, ] [[package]] @@ -2302,11 +2402,13 @@ wheels = [ [[package]] name = "nvidia-cuda-runtime" -version = "13.2.51" +version = "13.5.37" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/12/ec/0fa54349c6cf17054746529a3b99c6bc0049a229ac7fe667a24a244f79fa/nvidia_cuda_runtime-13.2.51-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:dfcccf62936b211a86e3cbc5fcc30647af8e601744e987f6a5925ed045b58672", size = 2340606, upload-time = "2026-03-09T09:29:42.659Z" }, - { url = "https://files.pythonhosted.org/packages/a6/5a/b116ad2b7e574d691458ca0139ab4e9f26beed62184c85570636ce127b7f/nvidia_cuda_runtime-13.2.51-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9c43b06a52c5b9316e19abc047236932c4d5c729969918a83223c4d2a4132f9a", size = 2321923, upload-time = "2026-03-09T09:30:13.061Z" }, + { url = "https://files.pythonhosted.org/packages/d4/8d/d162f19bd57e25048b654337c6f811802abe46c3da0f7d7052a2dfd709e8/nvidia_cuda_runtime-13.5.37-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:de2dc5a7367974228b0a69588a9acc9340d3768bef1bf38bf0801eb08497cd2d", size = 2520010, upload-time = "2026-11-04T23:22:09.169Z" }, + { url = "https://files.pythonhosted.org/packages/0b/7c/14b8f5518f8951e9dca57ea80f44f8ea8bcc8671812e6cc38009f3205e57/nvidia_cuda_runtime-13.5.37-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:df2b79a8852f10d1fef8c848489c49eebb5754051bf870202a862831412b179d", size = 2504024, upload-time = "2026-11-04T23:22:14.265Z" }, + { url = "https://files.pythonhosted.org/packages/35/99/e07dd9565d2e0317c1a62ff39ae9fc0542c048f74bcfc777c8be97710347/nvidia_cuda_runtime-13.5.37-py3-none-win_amd64.whl", hash = "sha256:f225d3fa5ea6bfbb58fc4979ef72375cbe661d9107f0d8a4ff22cc151525e4a9", size = 2780254, upload-time = "2026-11-04T23:34:42.881Z" }, + { url = "https://files.pythonhosted.org/packages/8f/35/dbd62813d7cdc3a958bec65b846e8bb3ff41e7af461802f7c0e2a522d7b2/nvidia_cuda_runtime-13.5.37-py3-none-win_arm64.whl", hash = "sha256:dcbcb895bc2098b4a6784b12b0a1ceafb09359dc95f4f92ea688ac935a106877", size = 2779932, upload-time = "2026-11-04T23:39:39.147Z" }, ] [[package]] @@ -2344,14 +2446,16 @@ wheels = [ [[package]] name = "nvidia-cufft" -version = "12.2.0.37" +version = "12.5.0.14" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-nvjitlink" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/6d/4d/31158ab042b044b269019574da4430ecce8d05fec7af1d270e1ba28e3512/nvidia_cufft-12.2.0.37-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d0d762b7ece2f2a4971a5f29a6cef13f26fd87c7cd0003b48ed82d9a1d7380d7", size = 218244744, upload-time = "2026-03-09T09:48:16.661Z" }, - { url = "https://files.pythonhosted.org/packages/00/a8/d8c0a8c4c45a3904a52c9860b07fdf775ca0517df884e3d240205a42b7ff/nvidia_cufft-12.2.0.37-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1530739e18736b07f57f835664659aa99179dab7b567c581a1ec7cb6c9737662", size = 218258094, upload-time = "2026-03-09T09:48:55.172Z" }, + { url = "https://files.pythonhosted.org/packages/1d/64/a71dc77eaca1623d1ff61f84045f2d21fb50a77796551e062bb8a989326d/nvidia_cufft-12.5.0.14-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:091cde103312d6e09366122fe3f67092dcfd94cc9785029d695d6bcffaa76fdf", size = 165074034, upload-time = "2026-11-04T23:27:54.684Z" }, + { url = "https://files.pythonhosted.org/packages/c6/8f/f909edaadf9a7acf1b92f805a384b6984193bd595b45005f8e320d4170cf/nvidia_cufft-12.5.0.14-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b8be26837a5abad2b5b74f3e92dd29f622a42934b355be0961a48712cbba23ad", size = 165072555, upload-time = "2026-11-04T23:28:25.686Z" }, + { url = "https://files.pythonhosted.org/packages/94/79/8f2a7125f560b86571c0be3dbf10e33041187bc4a8cf66107541fe8e5cce/nvidia_cufft-12.5.0.14-py3-none-win_amd64.whl", hash = "sha256:78e7585159ecf37d73af1ce43a5dea7fd3360e1af71cde08f070c07966a336bc", size = 164476208, upload-time = "2026-11-04T23:36:37.446Z" }, + { url = "https://files.pythonhosted.org/packages/9c/78/ec6cff1d3ba1d743acd7d6889db148fa00b97038bce8b169e49edb7c5a1b/nvidia_cufft-12.5.0.14-py3-none-win_arm64.whl", hash = "sha256:77eaacecf83758e8e20458a199d56359d46596e3a2814af741f27d1f483381e9", size = 165385033, upload-time = "2026-11-04T23:41:34.766Z" }, ] [[package]] @@ -2366,9 +2470,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/95/f4/61e6996dd20481ee834f57a8e9dca28b1869366a135e0d42e2aa8493bdd4/nvidia_cufft_cu12-11.4.1.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c67884f2a7d276b4b80eb56a79322a95df592ae5e765cf1243693365ccab4e28", size = 200877592, upload-time = "2025-06-05T20:05:45.862Z" }, ] +[[package]] +name = "nvidia-curand" +version = "10.4.5.25" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/44/33/c8d7fa70f90bde9b88a665734ab5e0440d4749646ac4e8590482760af3c6/nvidia_curand-10.4.5.25-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:07c75d8ba14b627d1a01d2c1bfa1eece88f25d564d3aae7e6b6fb44f7bbec9c1", size = 77529633, upload-time = "2026-11-04T23:28:49.816Z" }, + { url = "https://files.pythonhosted.org/packages/12/63/636e5dc9afb2959bfc0042a357d5567358df829b66c96d9fc36d77dc0596/nvidia_curand-10.4.5.25-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:22bb226a857e5888fe6b54ae465fc4371296bcff13a35f79e19691bcf454f529", size = 77113015, upload-time = "2026-11-04T23:29:04.942Z" }, + { url = "https://files.pythonhosted.org/packages/ee/b6/5b38490991b6875ae60c4f41e190ac91cadd1e6a1324b152999cc4a3f580/nvidia_curand-10.4.5.25-py3-none-win_amd64.whl", hash = "sha256:f6a02a70f5b9715235df31a164a4bf3237a45be1600a5f435f5e4a0d8682eb22", size = 77330709, upload-time = "2026-11-04T23:36:49.798Z" }, + { url = "https://files.pythonhosted.org/packages/c7/1c/4806430849d745b8133481fe9c6ac0d47d2ca5c54f2ae34c4f64b7cb7032/nvidia_curand-10.4.5.25-py3-none-win_arm64.whl", hash = "sha256:c998546b8bd1c746f2872cc11ec4c79b006dfed2a5b0acd25e7b710522b2d7f4", size = 77108234, upload-time = "2026-11-04T23:41:50.279Z" }, +] + [[package]] name = "nvidia-cusolver" -version = "12.1.0.51" +version = "12.4.1.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-cublas" }, @@ -2376,8 +2491,10 @@ dependencies = [ { name = "nvidia-nvjitlink" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/25/09/6214d11749dfc2d1be9a9e0bcfb04069077b98f7df0ad41115da87c84700/nvidia_cusolver-12.1.0.51-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:2a0bde343f956934103bc19344584a2d7b95d03b8cca9c0d22c90ff8917bbc3c", size = 224103511, upload-time = "2026-03-09T09:51:14.704Z" }, - { url = "https://files.pythonhosted.org/packages/ec/3a/6ee9b1c6632ec9cc0339996ffb331e5a8cbedcd361f7d4d0b63d48519a28/nvidia_cusolver-12.1.0.51-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0148ba705c196075607cd9d7a856a834695b406907b1ba8ad99b8a325a463611", size = 201732826, upload-time = "2026-03-09T09:51:45.431Z" }, + { url = "https://files.pythonhosted.org/packages/80/38/566a08a71dac1eebb21fa439970eab971d372c94633c3b16b851b415a13d/nvidia_cusolver-12.4.1.2-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:1f54c8e0d5b1ebd183b4c5a79f0709fa90241e8254da381ca5b969ee97c8de32", size = 304822178, upload-time = "2026-11-04T23:29:37.266Z" }, + { url = "https://files.pythonhosted.org/packages/78/4f/c5c926dbe0100e485c3ef86022ba3f6799fc57df6f9a73aaadeb7ab368fe/nvidia_cusolver-12.4.1.2-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b3b8e02fd2f2df869369af075dc3a179d43f6f7b85e94632a7dce556697be044", size = 260003940, upload-time = "2026-11-04T23:30:56.62Z" }, + { url = "https://files.pythonhosted.org/packages/cb/ca/bdf3af194210aad2b34db2b31033698be8b443dab38f566893149142d3d0/nvidia_cusolver-12.4.1.2-py3-none-win_amd64.whl", hash = "sha256:6b7be7d7dd60ba0e91cae19357a1d8c57dae321c0df99d92792ec025d791cc88", size = 274964624, upload-time = "2026-11-04T23:37:14.787Z" }, + { url = "https://files.pythonhosted.org/packages/c4/f9/733a3c7876db2f8fb154315cd4832fcbd2e6acd74f9274521444416388a8/nvidia_cusolver-12.4.1.2-py3-none-win_arm64.whl", hash = "sha256:8ec8a68b661ecdb89f4724a993ea9fab5c8096fc14aee5452344628a11529715", size = 103368508, upload-time = "2026-11-04T23:42:04.52Z" }, ] [[package]] @@ -2396,14 +2513,16 @@ wheels = [ [[package]] name = "nvidia-cusparse" -version = "12.7.9.17" +version = "12.9.4.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "nvidia-nvjitlink" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/06/5a/6bb7fc5f9658902efebc8551b1a9265b7a5908cbf9efdabdf97bc30b168a/nvidia_cusparse-12.7.9.17-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7c5d21980dc5f064c7ba13af2250e7fb4036283f1d1943c2915573c27928c89e", size = 168894384, upload-time = "2026-03-09T09:52:23.003Z" }, - { url = "https://files.pythonhosted.org/packages/87/40/23990a83164aaec2bfeffcee87794299f3cfdbdd7ed024b2af078afb600a/nvidia_cusparse-12.7.9.17-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7fb409bc7bb85e7a95706bd1e0b502b418a026dc35823179b4dafa92f1f2f7fd", size = 150883058, upload-time = "2026-03-09T09:53:00.781Z" }, + { url = "https://files.pythonhosted.org/packages/1b/71/9cf160aff268332d07b1d46e9009c5f28d626d36bd17d9cdc462a0668384/nvidia_cusparse-12.9.4.2-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:e2ce97f185c9c4b285e4a8fbf02cb96dd1b8dd66736ebe2f0a8da5f6012a4fb7", size = 230296529, upload-time = "2026-11-04T23:31:18.146Z" }, + { url = "https://files.pythonhosted.org/packages/30/7a/9928f812629a0e871871f2692cf5618857c51bf7093c7ecf2f78d17f9255/nvidia_cusparse-12.9.4.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:072667e58c851a318738f5527a12eb225cbf058803f2ea451c722d8cf8edddec", size = 230422857, upload-time = "2026-11-04T23:31:36.722Z" }, + { url = "https://files.pythonhosted.org/packages/71/9a/54114b62952d16a67e43a696916014e1e19ed6f495b0c1755d5baa967c7f/nvidia_cusparse-12.9.4.2-py3-none-win_amd64.whl", hash = "sha256:2defba1fbe76140c9eedc9fe82d58f11a0e2b7fd2c12dc5fb1267cc076689669", size = 228382183, upload-time = "2026-11-04T23:37:44.311Z" }, + { url = "https://files.pythonhosted.org/packages/a9/4d/18ba5c89b1968c8af804b33f647b2d94307c6fa98591a84d985f7b657b12/nvidia_cusparse-12.9.4.2-py3-none-win_arm64.whl", hash = "sha256:50d7a8c47c377c44ec1db681710156a28246ddfe3bde7be801daa85aba5483cc", size = 228043702, upload-time = "2026-11-04T23:42:34.936Z" }, ] [[package]] @@ -2438,11 +2557,13 @@ wheels = [ [[package]] name = "nvidia-nvjitlink" -version = "13.2.51" +version = "13.5.37" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/12/ae/ef3c49f1918aef93b39045499bfdb0ac9fb13e1785bc83f7a1b5d58a292d/nvidia_nvjitlink-13.2.51-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:6703c9ed79301382787a23fda9a7388af0779ecbc37545e4d50c055c897694a0", size = 41370377, upload-time = "2026-03-09T09:55:51.938Z" }, - { url = "https://files.pythonhosted.org/packages/6d/44/9e42f8ad320a23dfcb9542fc6ccd06a3a3ef495a779ade5937540fe7aa59/nvidia_nvjitlink-13.2.51-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3ba8bc06d756d10b643929e1f65dd162120c3e25ace58bd6414ee7eb98e10dde", size = 39285486, upload-time = "2026-03-09T09:55:19.032Z" }, + { url = "https://files.pythonhosted.org/packages/cc/1f/2321702b8d60b4b4dd18b957f61612dbac9cb1c07e940cf6582106b7c8b8/nvidia_nvjitlink-13.5.37-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:280d7cb9649704145f074b8dbde500de04ce02d0872e066e20a83c17bcaa93eb", size = 46642491, upload-time = "2026-11-04T23:33:31.696Z" }, + { url = "https://files.pythonhosted.org/packages/24/15/dd228f798cd01d7aeaf19d872b2b20f9aeafc1a7d8b6bfaa878b30b8d6f7/nvidia_nvjitlink-13.5.37-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:893f6f9a55b2720ed94913b5ba3332e972da8c78dfa2863d0244403dbd528ffc", size = 45253761, upload-time = "2026-11-04T23:33:22.854Z" }, + { url = "https://files.pythonhosted.org/packages/a7/53/c43398492e8374880d5e01c23c4dac33ba0c3bd03515d6ecc574254c0967/nvidia_nvjitlink-13.5.37-py3-none-win_amd64.whl", hash = "sha256:4465d4e8f2212118220b580c82523c5654f93a62906bd51d29d79bdfc12376c1", size = 41467665, upload-time = "2026-11-04T23:38:44.689Z" }, + { url = "https://files.pythonhosted.org/packages/d7/4a/1ff9ab7d0c961a1b5539f8291c881bea82adbbf8ba40a9e81f9c9b3a8d5d/nvidia_nvjitlink-13.5.37-py3-none-win_arm64.whl", hash = "sha256:eb30348ae64c525180a280025dd12a089c3e2e1dc4e19cdd58f68c6edc622404", size = 38552904, upload-time = "2026-11-04T23:43:27.658Z" }, ] [[package]]