diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..40a9cc5 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,55 @@ +name: Publish to PyPI + +# Single publication path. PyPI Trusted Publishing is bound to this file name +# (publish.yml) and the "pypi" environment. Tests must pass and the release tag +# must match the pyproject version before anything is uploaded. +on: + release: + types: [published] + workflow_dispatch: + +permissions: + contents: read + +jobs: + test: + uses: ./.github/workflows/tests.yml + + build: + needs: test + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + - name: Check release tag matches package version + if: github.event_name == 'release' + run: | + version=$(python -c "import tomllib; print(tomllib.load(open('pyproject.toml', 'rb'))['project']['version'])") + if [ "${GITHUB_REF_NAME#v}" != "$version" ]; then + echo "::error::Release tag $GITHUB_REF_NAME does not match pyproject version $version" + exit 1 + fi + - name: Build and check distributions + run: | + python -m pip install --upgrade build twine + python -m build + twine check --strict dist/* + - uses: actions/upload-artifact@v4 + with: + name: python-package-distributions + path: dist/ + + publish: + needs: build + runs-on: ubuntu-latest + environment: pypi + permissions: + id-token: write + steps: + - uses: actions/download-artifact@v4 + with: + name: python-package-distributions + path: dist/ + - uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..8a65f95 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,26 @@ +name: Tests + +on: + workflow_call: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.10", "3.11", "3.12"] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install + run: | + python -m pip install --upgrade pip + pip install -e ".[test]" + - name: Test + run: pytest tests/ -v diff --git a/CHANGELOG.md b/CHANGELOG.md index 2cac1c9..b8f104f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,19 @@ # Changelog +## 0.2.0 — 2026-10-27 + +- Added `westquant_cudaq.batched`: batched state-vector simulation of many small + circuits (NumPy on CPU, CuPy on GPU) with Pauli-trajectory noise, readout + error, shots and Parquet output with provenance. Validated against Qiskit + `Statevector` and an exact density-matrix reference. +- Added the `wq-batch` command (`run`, `plan`, `bench`) and Slurm scripts for + A100 clusters in `scripts/hpc/`. See `docs/BATCHED.md`. +- New optional extras `batched` and `gpu`. `import westquant_cudaq` stays + dependency-free; `westquant_cudaq.batched` is imported lazily. +- The module became a package (`westquant_cudaq/`); the public 0.1.0 API is + unchanged and its tests still pass. +- CI: tests workflow and a test-gated PyPI Trusted Publishing workflow. + ## 0.1.0 — 2026-09-29 - Added framework-neutral workload profiles and representation candidates. diff --git a/CITATION.cff b/CITATION.cff index 311bbb8..905d2c3 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -9,5 +9,5 @@ authors: repository-code: "https://github.com/WestQuantOpen/westquant-cudaq" url: "https://github.com/WestQuantOpen/westquant-cudaq" license: Apache-2.0 -version: "0.1.0" -date-released: "2026-09-29" +version: "0.2.0" +date-released: "2026-10-27" diff --git a/README.md b/README.md index b9ee5d4..01630ef 100644 --- a/README.md +++ b/README.md @@ -106,6 +106,19 @@ WestQuant never executes an unverified representation candidate. Constraints for `execute_cudaq` imports CUDA-Q lazily and supports kernel execution, sampling, and observation. No CUDA dependency is imported when using only planning and passport APIs. +## Batched GPU simulation for data generation (new in 0.2.0) + +`westquant_cudaq.batched` simulates thousands of small circuits per GPU batch (one batched matrix product per gate +step), with optional Pauli-trajectory noise, readout error, shots and Parquet output with full provenance. It includes +a `wq-batch` command and Slurm scripts for A100 clusters (`scripts/hpc/`). + +```bash +pip install "westquant-cudaq[batched,gpu]" +wq-batch run --manifest shard.jsonl --out shard.parquet --noise noise.json --trajectories 200 --shots 1000 +``` + +See [docs/BATCHED.md](docs/BATCHED.md). + ## Relationship to WestQuant Open - `westquant-core`: WQIR, RepGraph, plugin contracts, and generic search. diff --git a/docs/BATCHED.md b/docs/BATCHED.md new file mode 100644 index 0000000..e431c68 --- /dev/null +++ b/docs/BATCHED.md @@ -0,0 +1,85 @@ +# Batched GPU simulation (`westquant_cudaq.batched`) + +Throughput-oriented state-vector simulation of **many small circuits** (about 2–24 qubits) for dataset generation. +One batch holds thousands of circuits of the same width; each gate step is one batched matrix product for every +row that shares the same target qubits, so the GPU stays busy even when each circuit is tiny. + +Status: v0.2.0. Validated against Qiskit `Statevector` and an exact density-matrix reference on CPU (NumPy). +The GPU path runs the same array code through CuPy (cuBLAS batched GEMM); **run `bench_a100.sbatch` once on the +cluster before a campaign** – it runs the test suite on the GPU and measures throughput per qubit count. + +## Install + +```bash +pip install "westquant-cudaq[batched,gpu]" # CUDA 12 + CuPy; CPU-only: [batched] +``` + +## Python + +```python +from westquant_cudaq.batched import BatchRunner, PauliNoiseModel, from_qiskit, get_backend + +irs = [from_qiskit(qc, fusion="1q", circuit_id=f"c{i}", meta={"expected": "0101"}) for i, qc in enumerate(circuits)] +noise = PauliNoiseModel.from_calibration({"gate_error_1q": 5e-4, "gate_error_2q": 8e-3, "readout_error": 0.015}) +runner = BatchRunner(backend=get_backend("cupy"), noise=noise, trajectories=200, shots=1000) +runner.run_to_parquet(irs, "out.parquet") +``` + +`fusion`: `"full"` for ideal runs (fewest kernel calls), `"1q"` when noise is used (exact error placement). +The CLI picks this automatically. + +## HPC workflow (Slurm, A100) + +Submit everything **from the repository root** (the scripts locate themselves via `$SLURM_SUBMIT_DIR`). + +1. **Edit `scripts/hpc/env.sh` once** (modules, venv path) and fill in `--account/--partition` in the `.sbatch` files. +2. **Validate and benchmark on one GPU:** `sbatch scripts/hpc/bench_a100.sbatch $PWD` +3. **Shard the circuits:** + `python scripts/hpc/make_manifest.py --circuits circuits.parquet --out $WORK/wq/shards --max-circuits 20000` + (input: JSONL/Parquet with `circuit_id`, `qasm`, optional `family`, `expected`, `meta`; or `--qasm-dir`). + Check the plan without a GPU: `wq-batch plan --manifest $WORK/wq/shards/shard_00000.jsonl --gpu-mem-gb 64`. +4. **Run**, either as an array (one GPU per task) + `sbatch --array=0-15 scripts/hpc/run_array_a100.sbatch $WORK/wq/shards $WORK/wq/out --noise noise.json --trajectories 200 --shots 1000` + or on whole 4-GPU nodes + `sbatch --nodes=4 scripts/hpc/run_node_4xa100.sbatch $WORK/wq/shards $WORK/wq/out --noise noise.json --trajectories 200`. + Shard *i* goes to rank *i mod K*. Finished shards are skipped, so after a timeout you resubmit the same command. +5. **Merge and check:** + `python scripts/hpc/merge_shards.py --shards $WORK/wq/shards --out-dir $WORK/wq/out --dataset $WORK/wq/dataset` + → `dataset/data/n_qubits=*/…parquet` and `dataset/campaign_summary.json`; exit code 1 if anything is missing. + +## Tuning on A100 + +* **dtype:** `complex64` (default) halves memory and, since the kernels are memory-bound, typically gives close to 2× throughput versus `complex128`; the tests use + 2e-5 tolerance. Use `complex128` for validation runs. +* **Batch size** is set from free GPU memory: rows ≈ 0.6 × free / (2ⁿ × 8 B × 5). `--safety` changes the 0.6. + Noisy runs use `1 + trajectories` rows per circuit, so fewer circuits fit per batch. +* **Where batching pays:** small n (≲ 14) is limited by per-step host work and kernel launches – batching gives the + largest gain there. Somewhere around n ≈ 20–24 a single circuit fills the GPU by itself; there, per-circuit simulation with CUDA-Q/cuStateVec is the better tool (measure the crossover with `bench_a100.sbatch`). +* **Homogeneous batches are fastest.** Circuits with identical layouts (parameter sweeps, transpilation variants of + the same layout, trajectories) need one kernel call per gate and no gather/scatter. `make_manifest.py` sorts by + size; the runner sorts by structure within a batch. +* **CPU cores** parse QASM in parallel (`--workers`, default `$SLURM_CPUS_PER_TASK`). If parsing dominates + (`parse_s` in the summary), raise `--cpus-per-task` or pre-convert circuits once. +* If GPU utilisation is low for small n, run 2 processes per GPU (CUDA MPS) so one parses while the other simulates. + +## Noise model and limits + +Pauli trajectories: after each gate a random Pauli error with the gate's depolarizing probability; readout flips +per clbit applied exactly to the averaged distribution. Average gate infidelity r is converted with +p = r(d+1)/d. Not modelled (yet): idle decoherence (T1/T2), crosstalk, leakage, coherent errors, mid-circuit +measurement and reset. Statistical error of the noisy distribution ~ 1/√trajectories. + +## Output columns (one row per circuit) + +`circuit_id, n_qubits, n_measured, n_ops, n_2q, structure_key, fusion, meta` · `ideal_/noisy_top_bitstrings`, +`*_top_probs`, `*_probs` (full distribution when ≤ 10 measured bits), `*_success_prob` (if `expected` given), +`hellinger_fidelity`, `counts` (if `--shots`) · provenance: `run_id, westquant_cudaq, qiskit, gpu_name, dtype, +noise_digest, trajectories, seed, host, slurm_*`. Schema: `westquant_cudaq.batched.record_schema()`. + +## Roadmap + +* v0.3: cuStateVec batched backend (`custatevecApplyMatrixBatched`) behind the same API, compared against the CuPy + path on A100; device-side random numbers for trajectories; CUDA Graphs for repeated layouts. +* v0.3: idle T1/T2 noise (Pauli-twirled thermal relaxation) from backend calibration. +* Records map 1:1 onto the WQBench `hardware_runs`/`transitions` tables; link each campaign to an + `ExecutionPassport`. diff --git a/pyproject.toml b/pyproject.toml index 83832ec..45782a4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,13 +4,13 @@ build-backend = "setuptools.build_meta" [project] name = "westquant-cudaq" -version = "0.1.0" -description = "Representation-aware CUDA-Q execution planning, GPU/QPU allocation, and reproducibility" +version = "0.2.0" +description = "Representation-aware CUDA-Q execution planning, GPU/QPU allocation, reproducibility, and batched GPU simulation for data generation" readme = "README.md" requires-python = ">=3.10" license = "Apache-2.0" authors = [{name = "WestQuant Open"}] -keywords = ["quantum", "cuda-q", "cuquantum", "optimization", "scheduling", "gpu", "qpu"] +keywords = ["quantum", "cuda-q", "cuquantum", "optimization", "scheduling", "gpu", "qpu", "simulation", "dataset"] classifiers = [ "Development Status :: 3 - Alpha", "Intended Audience :: Science/Research", @@ -25,17 +25,24 @@ dependencies = [] [project.optional-dependencies] cudaq = ["cudaq>=0.12"] westquant = ["westquant-core>=0.2.0", "westquant-bridges>=0.1.0", "westquant-orchestrator>=0.1.0", "westquant-qcsc>=0.1.0"] -test = ["pytest>=8"] +# Batched simulator: CPU path (NumPy) + circuit import + Parquet output +batched = ["numpy>=1.24", "qiskit>=1.0", "pyarrow>=14"] +# GPU path for the batched simulator (CUDA 12; use cupy-cuda11x on CUDA 11, cupy-rocm-* on AMD) +gpu = ["cupy-cuda12x>=13"] +test = ["pytest>=8", "numpy>=1.24", "qiskit>=1.0", "pyarrow>=14"] + +[project.scripts] +wq-batch = "westquant_cudaq.batched.cli:main" [project.urls] Homepage = "https://github.com/WestQuantOpen/westquant-cudaq" Documentation = "https://github.com/WestQuantOpen/westquant-cudaq#readme" Issues = "https://github.com/WestQuantOpen/westquant-cudaq/issues" -[tool.setuptools] -py-modules = ["westquant_cudaq"] +[tool.setuptools.packages.find] +include = ["westquant_cudaq*"] [tool.pytest.ini_options] addopts = "-q" -testpaths = ["."] +testpaths = ["tests"] python_files = ["test_*.py"] diff --git a/scripts/hpc/bench_a100.sbatch b/scripts/hpc/bench_a100.sbatch new file mode 100755 index 0000000..a05755b --- /dev/null +++ b/scripts/hpc/bench_a100.sbatch @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# Validate and benchmark on one A100 before a campaign: GPU test suite, then throughput per qubit count. +# sbatch scripts/hpc/bench_a100.sbatch /path/to/westquant-cudaq +#SBATCH --job-name=wq-bench +#SBATCH --account=YOUR_ACCOUNT +#SBATCH --partition=YOUR_GPU_PARTITION +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:1 +#SBATCH --cpus-per-task=8 +#SBATCH --time=01:00:00 +#SBATCH --output=wq-bench_%j.out +set -euo pipefail +REPO=${1:-.} +HERE=${WQ_HPC_DIR:-$SLURM_SUBMIT_DIR/scripts/hpc} +source "$HERE/env.sh" +nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv +cd "$REPO" +WQ_TEST_BACKEND=cupy python -m pytest -q tests/test_batched.py # same tests as on CPU, now through CuPy +for dt in complex64 complex128; do + echo "== heterogeneous random circuits, $dt" + wq-batch bench --backend cupy --dtype $dt --nmin 4 --nmax 20 --nstep 2 --circuits 4096 --depth 40 --compare-qiskit 32 + echo "== same structure (parameter-sweep regime), $dt" + wq-batch bench --backend cupy --dtype $dt --nmin 4 --nmax 20 --nstep 4 --circuits 4096 --depth 40 --same-structure +done diff --git a/scripts/hpc/env.sh b/scripts/hpc/env.sh new file mode 100755 index 0000000..82c4509 --- /dev/null +++ b/scripts/hpc/env.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +# Edit once per cluster. Sourced by the Slurm scripts. +# Example for an NVIDIA A100 system with CUDA 12 modules (names differ per centre: check `module avail cuda python`). +# module load python/3.11 cuda/12.2 +# One-time setup (login node): +# python -m venv $WORK/venvs/wq && source $WORK/venvs/wq/bin/activate +# pip install "westquant-cudaq[batched,gpu]" # or: pip install -e /path/to/westquant-cudaq[batched,gpu] +# python -c "import cupy; print(cupy.cuda.runtime.getDeviceCount())" # run inside a GPU job to verify +source "${WQ_VENV:-$WORK/venvs/wq}/bin/activate" +export OMP_NUM_THREADS=${SLURM_CPUS_PER_TASK:-8} +export CUPY_CACHE_DIR=${CUPY_CACHE_DIR:-$WORK/.cupy_cache} # kernel cache on shared FS, not $HOME quota diff --git a/scripts/hpc/make_manifest.py b/scripts/hpc/make_manifest.py new file mode 100755 index 0000000..bed9af2 --- /dev/null +++ b/scripts/hpc/make_manifest.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Split a circuit collection into GPU-sized, cost-balanced manifest shards. + +Input (one of): + --qasm-dir DIR every *.qasm file is one circuit (circuit_id = file stem) + --circuits FILE JSONL or Parquet with at least circuit_id and qasm (other columns are kept) + +Shards are sorted by qubit count, so a shard holds circuits of similar size and the simulator fills its batches well. +Cost per circuit ~ 2**n * (number of gate lines); shards are cut when they reach --target-cost or --max-circuits. + +Writes OUT/shard_00000.jsonl ... and OUT/index.json (used by the Slurm scripts and merge_shards.py). + + python make_manifest.py --circuits all_circuits.parquet --out shards --max-circuits 20000 +""" +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path + +QREG = re.compile(r"qreg\s+\w+\s*\[\s*(\d+)\s*\]") +QUBIT3 = re.compile(r"qubit\s*\[\s*(\d+)\s*\]") + + +def n_qubits(qasm: str) -> int: + regs = [int(m) for m in QREG.findall(qasm)] or [int(m) for m in QUBIT3.findall(qasm)] + if regs: + return sum(regs) + from qiskit import qasm2 + return qasm2.loads(qasm, custom_instructions=qasm2.LEGACY_CUSTOM_INSTRUCTIONS).num_qubits + + +def gate_lines(qasm: str) -> int: + return sum(1 for line in qasm.splitlines() + if line.strip() and not line.lstrip().startswith(("OPENQASM", "include", "qreg", "creg", "qubit", "bit", + "measure", "barrier", "//"))) + + +def load(args) -> list[dict]: + if args.qasm_dir: + return [{"circuit_id": p.stem, "qasm": p.read_text()} for p in sorted(Path(args.qasm_dir).glob("*.qasm"))] + path = args.circuits + if path.endswith(".parquet"): + import pyarrow.parquet as pq + return pq.read_table(path).to_pylist() + with open(path) as f: + return [json.loads(line) for line in f if line.strip()] + + +def main() -> None: + ap = argparse.ArgumentParser() + src = ap.add_mutually_exclusive_group(required=True) + src.add_argument("--qasm-dir") + src.add_argument("--circuits") + ap.add_argument("--out", required=True) + ap.add_argument("--max-circuits", type=int, default=20000) + ap.add_argument("--target-cost", type=float, default=2e9, help="sum of 2**n * gates per shard") + ap.add_argument("--max-qubits", type=int, default=26, help="skip larger circuits (state-vector limit per row)") + args = ap.parse_args() + + items = load(args) + seen, rows, skipped = set(), [], [] + for it in items: + cid = str(it["circuit_id"]) + if cid in seen: + raise SystemExit(f"duplicate circuit_id {cid}") + seen.add(cid) + n = n_qubits(it["qasm"]) + if n > args.max_qubits: + skipped.append(cid) + continue + rows.append((n, (2 ** n) * max(1, gate_lines(it["qasm"])), it)) + rows.sort(key=lambda r: (r[0], r[1])) + + out = Path(args.out); out.mkdir(parents=True, exist_ok=True) + shards, cur, cost = [], [], 0.0 + for n, c, it in rows: + if cur and (len(cur) >= args.max_circuits or cost + c > args.target_cost): + shards.append(cur); cur, cost = [], 0.0 + cur.append((n, c, it)); cost += c + if cur: + shards.append(cur) + + index = {"n_circuits": len(rows), "skipped_too_large": skipped, "shards": []} + for i, sh in enumerate(shards): + name = f"shard_{i:05d}.jsonl" + with open(out / name, "w") as f: + for _n, _c, it in sh: + f.write(json.dumps(it) + "\n") + index["shards"].append({"file": name, "circuits": len(sh), "n_min": sh[0][0], "n_max": sh[-1][0], + "cost": sum(c for _n, c, _it in sh)}) + (out / "index.json").write_text(json.dumps(index, indent=1)) + print(f"{len(rows)} circuits -> {len(shards)} shards in {out} ({len(skipped)} skipped > {args.max_qubits} qubits)") + + +if __name__ == "__main__": + main() diff --git a/scripts/hpc/merge_shards.py b/scripts/hpc/merge_shards.py new file mode 100644 index 0000000..53ddd44 --- /dev/null +++ b/scripts/hpc/merge_shards.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +"""Merge shard outputs into one dataset and check the campaign is complete. + + python merge_shards.py --shards $WORK/wqbench/shards --out-dir $WORK/wqbench/out --dataset $WORK/wqbench/dataset + +* checks every shard in index.json has a Parquet output and the record count equals the manifest minus parse errors; +* writes a Hive-partitioned dataset to DATASET/data (partition: n_qubits) with zstd compression; +* writes DATASET/campaign_summary.json (throughput per shard, parse errors, missing shards, provenance set). +Exit code 1 if shards are missing or counts disagree, so it can gate the next pipeline step. +""" +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + + +def main() -> int: + import pyarrow as pa + import pyarrow.dataset as ds + import pyarrow.parquet as pq + + ap = argparse.ArgumentParser() + ap.add_argument("--shards", required=True, help="directory with index.json from make_manifest.py") + ap.add_argument("--out-dir", required=True, help="directory with shard_*.parquet from the Slurm jobs") + ap.add_argument("--dataset", required=True, help="output dataset directory") + a = ap.parse_args() + + index = json.loads((Path(a.shards) / "index.json").read_text()) + out_dir = Path(a.out_dir) + tables, report, missing, mismatched = [], [], [], [] + for sh in index["shards"]: + stem = Path(sh["file"]).stem + pq_path, sm_path = out_dir / f"{stem}.parquet", out_dir / f"{stem}.parquet.summary.json" + if not pq_path.exists(): + missing.append(stem) + continue + t = pq.read_table(pq_path) + summ = json.loads(sm_path.read_text()) if sm_path.exists() else {} + expected = sh["circuits"] - summ.get("parse_errors", 0) + if t.num_rows != expected: + mismatched.append({"shard": stem, "rows": t.num_rows, "expected": expected}) + tables.append(t) + report.append({"shard": stem, "rows": t.num_rows, "parse_errors": summ.get("parse_errors"), + "circuits_per_s": summ.get("circuits_per_s"), "gpu": summ.get("gpu_name")}) + + if tables: + full = pa.concat_tables(tables, promote_options="default") + dup = full.num_rows - len(set(full.column("circuit_id").to_pylist())) + ds.write_dataset(full, str(Path(a.dataset) / "data"), format="parquet", partitioning=["n_qubits"], partitioning_flavor="hive", + existing_data_behavior="overwrite_or_ignore", + file_options=ds.ParquetFileFormat().make_write_options(compression="zstd")) + prov = {k: sorted(set(map(str, full.column(k).to_pylist()))) for k in + ("run_id", "westquant_cudaq", "qiskit", "gpu_name", "dtype", "noise_digest") if k in full.column_names} + else: + dup, prov = 0, {} + summary = {"shards_total": len(index["shards"]), "shards_done": len(tables), "missing": missing, + "count_mismatch": mismatched, "duplicate_circuit_ids": dup, + "records": sum(r["rows"] for r in report), "provenance": prov, "per_shard": report} + Path(a.dataset).mkdir(parents=True, exist_ok=True) + (Path(a.dataset) / "campaign_summary.json").write_text(json.dumps(summary, indent=1)) + print(json.dumps({k: v for k, v in summary.items() if k != "per_shard"}, indent=1)) + return 1 if (missing or mismatched or dup) else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/hpc/run_array_a100.sbatch b/scripts/hpc/run_array_a100.sbatch new file mode 100755 index 0000000..3aa02c4 --- /dev/null +++ b/scripts/hpc/run_array_a100.sbatch @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# Slurm job array: one A100 per array task, each task works through every N-th shard. +# +# python scripts/hpc/make_manifest.py --circuits circuits.parquet --out $WORK/wqbench/shards +# sbatch --array=0-15 scripts/hpc/run_array_a100.sbatch $WORK/wqbench/shards $WORK/wqbench/out \ +# --noise noise.json --trajectories 200 --shots 1000 +# +# --array=0-(K-1) starts K GPU tasks; shard i goes to task i % K. Choose K <= number of shards. +# Adapt account/partition/QoS and the environment block to your centre (Leonardo Booster example in comments). +#SBATCH --job-name=wq-batch +#SBATCH --account=YOUR_ACCOUNT +#SBATCH --partition=YOUR_GPU_PARTITION # Leonardo: boost_usr_prod +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:1 # one A100 per array task +#SBATCH --cpus-per-task=8 # QASM parsing runs in parallel on these cores +#SBATCH --mem=60G +#SBATCH --time=04:00:00 +#SBATCH --output=wq-batch_%A_%a.out +set -euo pipefail +SHARD_DIR=$1; OUT_DIR=$2; shift 2 + +# Slurm runs a copy of this script from its spool dir, so locate the repo via the submit dir: +HERE=${WQ_HPC_DIR:-$SLURM_SUBMIT_DIR/scripts/hpc} # submit from the repo root, or export WQ_HPC_DIR +source "$HERE/env.sh" # loads modules and activates the venv (edit env.sh once) + +K=${SLURM_ARRAY_TASK_COUNT:-1} +RANK=${SLURM_ARRAY_TASK_ID:-0} +export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0} +bash "$HERE/worker.sh" "$SHARD_DIR" "$OUT_DIR" "$RANK" "$K" --workers "${SLURM_CPUS_PER_TASK:-8}" "$@" diff --git a/scripts/hpc/run_node_4xa100.sbatch b/scripts/hpc/run_node_4xa100.sbatch new file mode 100755 index 0000000..425cc43 --- /dev/null +++ b/scripts/hpc/run_node_4xa100.sbatch @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# Whole nodes with 4 x A100 each (e.g. Leonardo Booster). One process per GPU; ranks split the shards. +# +# sbatch --nodes=2 scripts/hpc/run_node_4xa100.sbatch $WORK/wqbench/shards $WORK/wqbench/out --dtype complex64 +# +# 2 nodes -> 8 ranks; shard i goes to rank i % 8. Resubmit the same command after a timeout: done shards are skipped. +#SBATCH --job-name=wq-batch-node +#SBATCH --account=YOUR_ACCOUNT +#SBATCH --partition=YOUR_GPU_PARTITION # Leonardo: boost_usr_prod +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=4 +#SBATCH --gpus-per-task=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=0 # all node memory +#SBATCH --time=08:00:00 +#SBATCH --output=wq-batch-node_%j.out +set -euo pipefail +SHARD_DIR=$1; OUT_DIR=$2; shift 2 +HERE=${WQ_HPC_DIR:-$SLURM_SUBMIT_DIR/scripts/hpc} # submit from the repo root, or export WQ_HPC_DIR +source "$HERE/env.sh" +STRIDE=$((SLURM_NNODES * SLURM_NTASKS_PER_NODE)) +# srun gives every task its own GPU (CUDA_VISIBLE_DEVICES has one entry per task, so --device 0 is correct) +srun --kill-on-bad-exit=0 bash -c "bash '$HERE/worker.sh' '$SHARD_DIR' '$OUT_DIR' \$SLURM_PROCID $STRIDE --workers $SLURM_CPUS_PER_TASK $*" diff --git a/scripts/hpc/worker.sh b/scripts/hpc/worker.sh new file mode 100755 index 0000000..5c6295c --- /dev/null +++ b/scripts/hpc/worker.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +# Run a strided subset of manifest shards on ONE GPU. +# worker.sh SHARD_DIR OUT_DIR RANK STRIDE [extra wq-batch run args...] +# Shard i is processed by the rank with i % STRIDE == RANK. Finished shards (output exists) are skipped, +# so a job that hits its time limit can simply be resubmitted. +set -euo pipefail +SHARD_DIR=$1; OUT_DIR=$2; RANK=$3; STRIDE=$4; shift 4 +mkdir -p "$OUT_DIR" "$OUT_DIR/logs" +mapfile -t SHARDS < <(ls "$SHARD_DIR"/shard_*.jsonl | sort) +echo "[worker $RANK/$STRIDE] host=$(hostname) gpu=${CUDA_VISIBLE_DEVICES:-unset} shards=${#SHARDS[@]}" +for ((i = RANK; i < ${#SHARDS[@]}; i += STRIDE)); do + shard=${SHARDS[$i]} + name=$(basename "$shard" .jsonl) + out="$OUT_DIR/$name.parquet" + if [[ -f "$out" ]]; then echo "skip $name"; continue; fi + echo "[$(date +%T)] start $name" + wq-batch run --manifest "$shard" --out "$out" --backend cupy --device 0 "$@" \ + > "$OUT_DIR/logs/$name.log" 2>&1 || echo "FAILED $name (see logs/$name.log)" + echo "[$(date +%T)] done $name" +done diff --git a/tests/test_batched.py b/tests/test_batched.py new file mode 100644 index 0000000..87fead9 --- /dev/null +++ b/tests/test_batched.py @@ -0,0 +1,162 @@ +"""Correctness of the batched simulator against Qiskit and against an exact density-matrix reference. + +Runs on CPU (NumPy backend). The GPU path executes the same array code through CuPy; run this file on the A100 with +``WQ_TEST_BACKEND=cupy pytest tests/test_batched.py`` to validate it there. +""" +import json +import os + +import numpy as np +import pytest + +qiskit = pytest.importorskip("qiskit") +from qiskit import QuantumCircuit # noqa: E402 +from qiskit.circuit.random import random_circuit # noqa: E402 +from qiskit.quantum_info import Statevector # noqa: E402 + +from westquant_cudaq.batched import (BatchRunner, PauliNoiseModel, from_qiskit, get_backend, # noqa: E402 + simulate_batch) +from westquant_cudaq.batched.noise import PAULI_1Q, PAULI_2Q # noqa: E402 + +BACKEND = get_backend(os.environ.get("WQ_TEST_BACKEND", "numpy")) +TOL = 2e-5 # complex64 + + +def _reference(qc: QuantumCircuit, measured): + sv = Statevector(qc.remove_final_measurements(inplace=False)) + return sv.probabilities(qargs=list(measured) if measured is not None else None) + + +def _measured_qubits(qc): + pairs = sorted((qc.find_bit(i.clbits[0]).index, qc.find_bit(i.qubits[0]).index) + for i in qc.data if i.operation.name == "measure") + return tuple(q for _c, q in pairs) if pairs else None + + +@pytest.mark.parametrize("fusion", ["none", "1q", "full"]) +def test_heterogeneous_batch_matches_qiskit(fusion): + rng = np.random.default_rng(0) + for n in (1, 2, 3, 5): + qcs = [] + for i in range(24): + qc = random_circuit(n, depth=12, max_operands=min(3, n), seed=int(rng.integers(1 << 30))) + if i % 3 == 0 and n > 1: # permuted, partial measurement + perm = rng.permutation(n)[: max(1, n - 1)] + qc2 = QuantumCircuit(n, len(perm)); qc2.compose(qc, inplace=True) + for c, q in enumerate(perm): + qc2.measure(int(q), c) + qc = qc2 + qcs.append(qc) + irs = [from_qiskit(qc, fusion=fusion, circuit_id=str(i)) for i, qc in enumerate(qcs)] + groups = {} + for qc, ir in zip(qcs, irs): + groups.setdefault(ir.measured, []).append((qc, ir)) + for meas, items in groups.items(): + res = simulate_batch([ir for _q, ir in items], BACKEND) + for j, (qc, ir) in enumerate(items): + np.testing.assert_allclose(res.ideal[j], _reference(qc, meas), atol=TOL) + + +def test_same_structure_fast_path_and_kernel_count(): + base = random_circuit(6, depth=20, max_operands=2, seed=7) + ir = from_qiskit(base, fusion="full") + res = simulate_batch([ir] * 64, BACKEND) + ref = _reference(base, None) + np.testing.assert_allclose(res.ideal, np.tile(ref, (64, 1)), atol=TOL) + assert res.stats["kernel_calls"] == len(ir.ops) # one batched call per op, independent of batch size + + +def test_zero_noise_trajectories_equal_ideal(): + qc = random_circuit(4, depth=10, max_operands=2, seed=3) + ir = from_qiskit(qc, fusion="1q") + res = simulate_batch([ir, ir], BACKEND, noise=PauliNoiseModel(0.0, 0.0), trajectories=50) + np.testing.assert_allclose(res.noisy, res.ideal, atol=TOL) + + +def test_full_fusion_rejected_with_noise(): + ir = from_qiskit(random_circuit(3, 6, max_operands=2, seed=1), fusion="full") + with pytest.raises(ValueError): + simulate_batch([ir], BACKEND, noise=PauliNoiseModel(0.01, 0.02), trajectories=4) + + +def _density_reference(ir, noise): + n = ir.num_qubits + d = 2 ** n + rho = np.zeros((d, d), complex); rho[0, 0] = 1 + + def embed(op_mat, qubits): + # full operator via Statevector-free tensor construction + full = np.zeros((d, d), complex) + for col in range(d): + vec = np.zeros(d, complex); vec[col] = 1 + t = vec.reshape((2,) * n) + axes = [n - 1 - q for q in reversed(qubits)] + t = np.moveaxis(t, axes, list(range(n - len(qubits), n))) + shp = t.shape + t = (t.reshape(-1, 2 ** len(qubits)) @ op_mat.T).reshape(shp) + full[:, col] = np.moveaxis(t, list(range(n - len(qubits), n)), axes).reshape(d) + return full + + for op in ir.ops: + U = embed(op.matrix, op.qubits) + rho = U @ rho @ U.conj().T + if len(op.qubits) == 1: + p = noise.p1_of(op.qubits[0], op.count); table = PAULI_1Q + else: + p = noise.p2_of(*op.qubits); table = PAULI_2Q + if p: + mix = sum(embed(P, op.qubits) @ rho @ embed(P, op.qubits).conj().T for P in table) / len(table) + rho = (1 - p) * rho + p * mix + return np.real(np.diag(rho)) + + +def test_depolarizing_trajectories_converge_to_density_matrix(): + qc = QuantumCircuit(3) + qc.h(0); qc.rx(0.7, 1); qc.cx(0, 1); qc.ry(1.1, 2); qc.cz(1, 2); qc.t(0); qc.s(0); qc.cx(2, 0); qc.h(1) + ir = from_qiskit(qc, fusion="1q") + noise = PauliNoiseModel(p1=0.04, p2={"0,1": 0.12, "1,2": 0.08, "default": 0.1}) + ref = _density_reference(ir, noise) + res = simulate_batch([ir], BACKEND, noise=noise, trajectories=40000, rng=np.random.default_rng(5)) + np.testing.assert_allclose(res.noisy[0], ref, atol=6e-3) + + +def test_readout_error_is_applied_per_clbit(): + qc = QuantumCircuit(2, 2); qc.x(0); qc.measure(0, 0); qc.measure(1, 1) # true outcome "01" + ir = from_qiskit(qc) + noise = PauliNoiseModel(0.0, 0.0, readout={0: [0.1, 0.2], 1: [0.3, 0.0]}) + res = simulate_batch([ir], BACKEND, noise=noise, trajectories=2) + # clbit0: true 1 -> P(1)=0.8 ; clbit1: true 0 -> P(1)=0.3 ; index = c0 + 2 c1 + expect = np.array([0.2 * 0.7, 0.8 * 0.7, 0.2 * 0.3, 0.8 * 0.3]) + np.testing.assert_allclose(res.noisy[0], expect, atol=1e-6) + + +def test_runner_records_and_parquet(tmp_path): + pq = pytest.importorskip("pyarrow.parquet") + qcs = [] + for i in range(10): + qc = QuantumCircuit(3, 3); qc.h(0); qc.cx(0, 1); qc.cx(1, 2); qc.measure([0, 1, 2], [0, 1, 2]); qcs.append(qc) + irs = [from_qiskit(qc, fusion="1q", circuit_id=f"ghz{i}", meta={"expected": ["000", "111"]}) for i, qc in enumerate(qcs)] + runner = BatchRunner(backend=BACKEND, noise=PauliNoiseModel(0.01, 0.03, 0.02), trajectories=200, shots=500) + out = tmp_path / "s.parquet" + info = runner.run_to_parquet(irs, str(out)) + t = pq.read_table(out).to_pylist() + assert info["records"] == 10 and len(t) == 10 + assert abs(t[0]["ideal_success_prob"] - 1.0) < 1e-5 + assert 0.8 < t[0]["noisy_success_prob"] < 1.0 + assert sum(json.loads(t[0]["counts"]).values()) == 500 + + +def test_cli_run_on_manifest(tmp_path): + from qiskit import qasm2 + + from westquant_cudaq.batched.cli import main + man = tmp_path / "m.jsonl" + with open(man, "w") as f: + for i in range(6): + qc = random_circuit(4, 8, max_operands=2, measure=True, seed=i) + f.write(json.dumps({"circuit_id": f"r{i}", "qasm": qasm2.dumps(qc), "family": "random"}) + "\n") + f.write(json.dumps({"circuit_id": "bad", "qasm": "not qasm"}) + "\n") + out = tmp_path / "o.parquet" + assert main(["run", "--manifest", str(man), "--out", str(out), "--backend", "numpy", "--workers", "1"]) == 0 + summary = json.loads((tmp_path / "o.parquet.summary.json").read_text()) + assert summary["circuits"] == 6 and summary["parse_errors"] == 1 diff --git a/test_westquant_cudaq.py b/tests/test_westquant_cudaq.py similarity index 100% rename from test_westquant_cudaq.py rename to tests/test_westquant_cudaq.py diff --git a/westquant_cudaq.py b/westquant_cudaq/__init__.py similarity index 97% rename from westquant_cudaq.py rename to westquant_cudaq/__init__.py index b175373..c534665 100644 --- a/westquant_cudaq.py +++ b/westquant_cudaq/__init__.py @@ -12,7 +12,7 @@ from enum import Enum from typing import Any, Callable, Iterable, Mapping, Protocol, Sequence -__version__ = "0.1.0" +__version__ = "0.2.0" class Engine(str, Enum): @@ -474,3 +474,17 @@ def _package_version(name: str) -> str | None: "normalize_workload", "search_representations", ] + + +def __getattr__(name: str) -> Any: + """Lazy access to the GPU batch simulator (``westquant_cudaq.batched``). + + Importing :mod:`westquant_cudaq` stays dependency-free; the batched engine (NumPy, optional CuPy, + optional Qiskit) is only imported when it is first used. + """ + if name == "batched": + import importlib + + return importlib.import_module(".batched", __name__) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + diff --git a/westquant_cudaq/batched/__init__.py b/westquant_cudaq/batched/__init__.py new file mode 100644 index 0000000..4428c95 --- /dev/null +++ b/westquant_cudaq/batched/__init__.py @@ -0,0 +1,21 @@ +"""Batched GPU simulation of many small circuits (throughput-oriented data generation). + +Quick start:: + + from westquant_cudaq.batched import BatchRunner, from_qiskit, get_backend + + irs = [from_qiskit(qc, circuit_id=f"c{i}") for i, qc in enumerate(circuits)] + runner = BatchRunner(backend=get_backend("auto")) # CuPy on GPU, NumPy otherwise + runner.run_to_parquet(irs, "out.parquet") + +Command line: ``wq-batch run --manifest shard.jsonl --out shard.parquet`` (see ``docs/BATCHED.md``). +""" +from .backend import Backend, get_backend +from .engine import BatchedStatevector, BatchResult, simulate_batch +from .ir import CircuitIR, Op, from_qasm, from_qiskit, fuse_ops +from .noise import PauliNoiseModel +from .runner import BatchRunner, hellinger_fidelity, max_rows, plan_batches, record_schema + +__all__ = ["Backend", "BatchResult", "BatchRunner", "BatchedStatevector", "CircuitIR", "Op", "PauliNoiseModel", + "from_qasm", "from_qiskit", "fuse_ops", "get_backend", "hellinger_fidelity", "max_rows", "plan_batches", + "record_schema", "simulate_batch"] diff --git a/westquant_cudaq/batched/backend.py b/westquant_cudaq/batched/backend.py new file mode 100644 index 0000000..ec8f357 --- /dev/null +++ b/westquant_cudaq/batched/backend.py @@ -0,0 +1,67 @@ +"""Array backend selection for the batched simulator. + +The engine is written against the NumPy array API and runs unchanged on CuPy. On an NVIDIA GPU the heavy +operation (one batched matrix product per gate) is dispatched by CuPy to cuBLAS batched GEMM, so thousands of +small state vectors are updated in a single kernel launch. + +``prefer="auto"`` uses CuPy when it imports and sees a device, otherwise NumPy (useful for tests and for laptops). +""" +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from typing import Any + +import numpy as np + + +@dataclass +class Backend: + xp: Any # numpy or cupy module + name: str # "numpy" | "cupy" + device: int | None = None + info: dict = field(default_factory=dict) + + def asnumpy(self, a): + return self.xp.asnumpy(a) if self.name == "cupy" else np.asarray(a) + + def asarray(self, a, dtype=None): + return self.xp.asarray(a, dtype=dtype) + + def synchronize(self) -> None: + if self.name == "cupy": + self.xp.cuda.Stream.null.synchronize() + + def free_bytes(self) -> int: + """Free device memory (GPU) or a conservative host budget (CPU).""" + if self.name == "cupy": + free, _total = self.xp.cuda.runtime.memGetInfo() + # CuPy's memory pool may hold freed blocks; count them as available. + pool = self.xp.get_default_memory_pool() + return int(free + pool.free_bytes()) + env = os.environ.get("WQ_BATCH_HOST_BYTES") + return int(env) if env else 4 * 2**30 + + +def get_backend(prefer: str = "auto", device: int | None = None) -> Backend: + """Return a CuPy backend if possible (``prefer`` in {"auto", "cupy"}), else NumPy. + + ``prefer="cupy"`` raises if CuPy or a GPU is missing, so HPC jobs fail loudly instead of silently running on CPU. + """ + if prefer in ("auto", "cupy"): + try: + import cupy as cp # noqa: WPS433 + + n_dev = cp.cuda.runtime.getDeviceCount() + if n_dev == 0: + raise RuntimeError("no CUDA device visible") + dev = 0 if device is None else int(device) + cp.cuda.Device(dev).use() + props = cp.cuda.runtime.getDeviceProperties(dev) + name = props["name"].decode() if isinstance(props["name"], bytes) else str(props["name"]) + return Backend(cp, "cupy", dev, {"gpu_name": name, "gpu_total_bytes": int(props["totalGlobalMem"]), + "cupy": cp.__version__}) + except Exception as exc: # noqa: BLE001 + if prefer == "cupy": + raise RuntimeError(f"CuPy GPU backend requested but unavailable: {exc}") from exc + return Backend(np, "numpy", None, {"numpy": np.__version__}) diff --git a/westquant_cudaq/batched/cli.py b/westquant_cudaq/batched/cli.py new file mode 100644 index 0000000..22e01b4 --- /dev/null +++ b/westquant_cudaq/batched/cli.py @@ -0,0 +1,195 @@ +"""Command line: ``wq-batch run | plan | bench``. + +Manifest = JSON Lines or Parquet with one circuit per row:: + + {"circuit_id": "qaoa-8-000123", "qasm": "OPENQASM 2.0; ...", "family": "qaoa", "expected": "01011010", "meta": {...}} + +``expected`` (optional) is a classical bitstring (Qiskit order, clbit 0 rightmost) or a list of them; it enables +``*_success_prob``. Everything in ``meta`` plus ``family`` is copied to the record's ``meta`` column. +""" +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import sys +import time +from concurrent.futures import ProcessPoolExecutor + +import numpy as np + + +def _load_manifest(path: str) -> list[dict]: + if path.endswith(".parquet"): + import pyarrow.parquet as pq + + return pq.read_table(path).to_pylist() + with open(path) as f: + return [json.loads(line) for line in f if line.strip()] + + +def _parse_one(args): + item, fusion = args + from .ir import from_qasm + + meta = dict(item.get("meta") or {}) + for k in ("family", "expected", "source"): + if item.get(k) is not None: + meta[k] = item[k] + try: + return from_qasm(item["qasm"], fusion=fusion, circuit_id=str(item["circuit_id"]), meta=meta), None + except Exception as exc: # noqa: BLE001 -- one bad circuit must not kill a shard; it is reported + return None, f"{item.get('circuit_id')}: {type(exc).__name__}: {exc}" + + +def parse_manifest(path: str, fusion: str, workers: int): + items = _load_manifest(path) + jobs = [(it, fusion) for it in items] + if workers > 1: + with ProcessPoolExecutor(max_workers=workers) as ex: + out = list(ex.map(_parse_one, jobs, chunksize=64)) + else: + out = [_parse_one(j) for j in jobs] + circuits = [c for c, e in out if c is not None] + errors = [e for c, e in out if e is not None] + return circuits, errors + + +def _seed_for(path: str, base: int) -> int: + return base ^ int(hashlib.sha256(os.path.basename(path).encode()).hexdigest()[:8], 16) + + +def cmd_run(a) -> int: + from .backend import get_backend + from .noise import PauliNoiseModel + from .runner import BatchRunner + + if os.path.exists(a.out) and not a.overwrite: + print(f"[wq-batch] {a.out} exists, skipping (use --overwrite)") + return 0 + noise = PauliNoiseModel.from_json(a.noise) if a.noise else None + fusion = a.fusion if a.fusion != "auto" else ("1q" if noise and a.trajectories > 0 else "full") + t0 = time.perf_counter() + circuits, errors = parse_manifest(a.manifest, fusion, a.workers) + t_parse = time.perf_counter() - t0 + backend = get_backend(a.backend, a.device) + runner = BatchRunner(backend=backend, dtype=a.dtype, noise=noise, trajectories=a.trajectories, shots=a.shots, + topk=a.topk, seed=_seed_for(a.manifest, a.seed), safety=a.safety) + info = runner.run_to_parquet(circuits, a.out) + summary = {"manifest": a.manifest, "out": a.out, "circuits": len(circuits), "parse_errors": len(errors), + "parse_s": round(t_parse, 2), "sim_s": round(info["wall_s"], 2), + "circuits_per_s": round(len(circuits) / max(info["wall_s"], 1e-9), 1), "backend": backend.name, + **backend.info} + with open(a.out + ".summary.json", "w") as f: + json.dump({**summary, "errors": errors[:1000]}, f, indent=2) + print(json.dumps(summary)) + return 0 if not errors or not a.strict else 2 + + +def cmd_plan(a) -> int: + from .backend import Backend + from .runner import BYTES_FACTOR, plan_batches + + circuits, errors = parse_manifest(a.manifest, "full", a.workers) + fake = Backend(np, "plan", None, {}) + budget = int(a.gpu_mem_gb * 2**30 * a.safety) + rows_per = 1 + a.trajectories + + def cap(n): + return max(1, min(1 << 16, budget // ((2 ** n) * np.dtype(a.dtype).itemsize * BYTES_FACTOR))) + batches = plan_batches(circuits, rows_per, cap) + by_n: dict[int, list[int]] = {} + for b in batches: + by_n.setdefault(circuits[b[0]].num_qubits, []).append(len(b)) + print(f"{len(circuits)} circuits, {len(errors)} parse errors, {len(batches)} batches") + for n in sorted(by_n): + sizes = by_n[n] + print(f" n={n:2d}: {len(sizes):5d} batches, circuits/batch max {max(sizes)}, rows/batch max {max(sizes) * rows_per}" + f", state bytes/batch {max(sizes) * rows_per * 2**n * np.dtype(a.dtype).itemsize / 2**30:.2f} GiB") + del fake + return 0 + + +def cmd_bench(a) -> int: + from qiskit.circuit.random import random_circuit + from qiskit.quantum_info import Statevector + + from .backend import get_backend + from .ir import from_qiskit + from .runner import BatchRunner + + backend = get_backend(a.backend, a.device) + print(json.dumps({"backend": backend.name, **backend.info})) + for n in range(a.nmin, a.nmax + 1, a.nstep): + qcs = [random_circuit(n, a.depth, max_operands=2, seed=1000 * n + i) for i in range(a.circuits)] + if a.same_structure: # parameter-sweep regime: one layout, different angles + base = qcs[0] + qcs = [base.assign_parameters([]) for _ in range(a.circuits)] + irs = [from_qiskit(qc, fusion="full", circuit_id=f"b{n}-{i}") for i, qc in enumerate(qcs)] + runner = BatchRunner(backend=backend, dtype=a.dtype, topk=4) + t0 = time.perf_counter() + recs = list(runner.run(irs)) + dt = time.perf_counter() - t0 + line = {"n": n, "circuits": len(recs), "sim_s": round(dt, 3), "circuits_per_s": round(len(recs) / dt, 1), + "batches": len({r["batch_index"] for r in recs})} + if a.compare_qiskit: + k = min(a.compare_qiskit, len(qcs)) + t1 = time.perf_counter() + for qc in qcs[:k]: + Statevector(qc.remove_final_measurements(inplace=False)).probabilities() + line["qiskit_statevector_circuits_per_s"] = round(k / (time.perf_counter() - t1), 1) + print(json.dumps(line), flush=True) + return 0 + + +def main(argv=None) -> int: + p = argparse.ArgumentParser(prog="wq-batch", description="Batched GPU simulation of many small circuits") + sub = p.add_subparsers(dest="cmd", required=True) + + r = sub.add_parser("run", help="simulate one manifest shard and write Parquet") + r.add_argument("--manifest", required=True) + r.add_argument("--out", required=True) + r.add_argument("--noise", help="noise JSON (PauliNoiseModel dict or calibration with gate_error_1q/2q)") + r.add_argument("--trajectories", type=int, default=0) + r.add_argument("--shots", type=int, default=0) + r.add_argument("--dtype", default="complex64", choices=["complex64", "complex128"]) + r.add_argument("--backend", default="auto", choices=["auto", "cupy", "numpy"]) + r.add_argument("--device", type=int, default=None) + r.add_argument("--workers", type=int, default=int(os.environ.get("SLURM_CPUS_PER_TASK", "4"))) + r.add_argument("--fusion", default="auto", choices=["auto", "full", "1q", "none"]) + r.add_argument("--topk", type=int, default=16) + r.add_argument("--seed", type=int, default=20261027) + r.add_argument("--safety", type=float, default=0.6, help="fraction of free device memory to use") + r.add_argument("--overwrite", action="store_true") + r.add_argument("--strict", action="store_true", help="exit 2 if any circuit failed to parse") + r.set_defaults(fn=cmd_run) + + pl = sub.add_parser("plan", help="show batch plan for a manifest without a GPU") + pl.add_argument("--manifest", required=True) + pl.add_argument("--gpu-mem-gb", type=float, default=64.0) + pl.add_argument("--dtype", default="complex64") + pl.add_argument("--trajectories", type=int, default=0) + pl.add_argument("--safety", type=float, default=0.6) + pl.add_argument("--workers", type=int, default=4) + pl.set_defaults(fn=cmd_plan) + + b = sub.add_parser("bench", help="throughput on random circuits") + b.add_argument("--nmin", type=int, default=4) + b.add_argument("--nmax", type=int, default=16) + b.add_argument("--nstep", type=int, default=2) + b.add_argument("--circuits", type=int, default=2048) + b.add_argument("--depth", type=int, default=30) + b.add_argument("--dtype", default="complex64") + b.add_argument("--backend", default="auto", choices=["auto", "cupy", "numpy"]) + b.add_argument("--device", type=int, default=None) + b.add_argument("--same-structure", action="store_true") + b.add_argument("--compare-qiskit", type=int, default=0, help="also time Qiskit Statevector on this many circuits") + b.set_defaults(fn=cmd_bench) + + a = p.parse_args(argv) + return a.fn(a) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/westquant_cudaq/batched/engine.py b/westquant_cudaq/batched/engine.py new file mode 100644 index 0000000..49e00f0 --- /dev/null +++ b/westquant_cudaq/batched/engine.py @@ -0,0 +1,186 @@ +"""Batched state-vector simulation of many small circuits on one device. + +All circuits in a batch have the same number of qubits n. The state of the batch is one array ``psi`` of shape +``(rows, 2**n)``. Each row is one circuit (ideal) or one noise trajectory of a circuit. A gate step applies, for every +group of rows that share the same target qubits, a *batched* matrix product + + psi[rows] <- psi[rows] @ U[rows]^T (per-row 2x2 or 4x4 matrices, one cuBLAS batched GEMM on GPU) + +so the cost of a step is one kernel launch per distinct target tuple, not one per circuit. Rows that share their whole +structure (same targets at every step: parameter sweeps, trajectories, transpilation variants of a fixed layout) are +updated without any gather or scatter. +""" +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from typing import Sequence + +import numpy as np + +from .backend import Backend +from .ir import CircuitIR +from .noise import PAULI_1Q, PAULI_2Q, PauliNoiseModel + + +@dataclass +class BatchResult: + """Device-side results reduced to host arrays, one row per *circuit* (not per trajectory).""" + ideal: np.ndarray # (C, 2**m) ideal output distribution over measured bits + noisy: np.ndarray | None # (C, 2**m) trajectory-averaged + readout noise, or None + n_measured: int + stats: dict = field(default_factory=dict) + + +class BatchedStatevector: + def __init__(self, backend: Backend, num_qubits: int, rows: int, dtype: str = "complex64"): + xp = backend.xp + self.be, self.xp, self.n, self.rows = backend, xp, num_qubits, rows + self.cdtype = np.dtype(dtype) + self.psi = xp.zeros((rows, 2 ** num_qubits), dtype=self.cdtype) + self.psi[:, 0] = 1.0 + self.kernel_calls = 0 + + def apply(self, qubits: Sequence[int], mats, idx=None) -> None: + """Apply per-row matrices ``mats`` (rows_sel, D, D) or one matrix (D, D) to ``qubits`` on rows ``idx``.""" + xp, n, k = self.xp, self.n, len(qubits) + sub = self.psi if idx is None else self.psi[idx] + b = sub.shape[0] + if b == 0: + return + t = sub.reshape((b,) + (2,) * n) + src = [1 + (n - 1 - q) for q in reversed(qubits)] # most significant target first + dst = list(range(n + 1 - k, n + 1)) + t = xp.moveaxis(t, src, dst) + shp = t.shape + t = t.reshape(b, -1, 2 ** k) + if mats.ndim == 2: + t = t @ mats.T + else: + t = xp.matmul(t, xp.swapaxes(mats, 1, 2)) + t = xp.moveaxis(t.reshape(shp), dst, src).reshape(b, 2 ** n) + if idx is None: + self.psi = xp.ascontiguousarray(t) + else: + self.psi[idx] = t + self.kernel_calls += 1 + + +def _marginal_axes(n: int, measured: Sequence[int] | None): + """Axes of the (rows, 2,...,2) probability tensor to keep, ordered so the flattened index equals the classical + register value (clbit i = measured[i], clbit 0 least significant).""" + meas = list(range(n)) if measured is None else list(measured) + keep = [1 + (n - 1 - q) for q in reversed(meas)] + drop = tuple(a for a in range(1, n + 1) if a not in keep) + return meas, keep, drop + + +def simulate_batch(circuits: Sequence[CircuitIR], backend: Backend, *, dtype: str = "complex64", + noise: PauliNoiseModel | None = None, trajectories: int = 0, + rng: np.random.Generator | None = None) -> BatchResult: + """Simulate ``circuits`` (same qubit count and same measured-qubit count) in one batch. + + With ``noise`` and ``trajectories = T > 0`` each circuit occupies ``1 + T`` rows: row 0 ideal, rows 1..T noisy. + """ + if not circuits: + raise ValueError("empty batch") + xp = backend.xp + n = circuits[0].num_qubits + if any(c.num_qubits != n for c in circuits): + raise ValueError("all circuits in a batch must have the same number of qubits") + noisy = noise is not None and trajectories > 0 + if noisy and any(c.fusion == "full" for c in circuits) and not noise.is_trivial: + raise ValueError("noisy simulation needs circuits converted with fusion='1q' or 'none' " + "(full fusion moves 1q gates across 2q gates and would misplace errors)") + rng = rng or np.random.default_rng() + T = trajectories if noisy else 0 + R = 1 + T + C = len(circuits) + sv = BatchedStatevector(backend, n, C * R, dtype) + t0 = time.perf_counter() + + # ---- one pass over all ops on the host: target codes, matrices and error probabilities as dense arrays + S = max(len(c.ops) for c in circuits) + code = np.full((C, S), -1, dtype=np.int64) # -1 = no op at this step (circuit finished) + M1 = np.zeros((C, S, 2, 2), dtype=sv.cdtype) + M2 = np.zeros((C, S, 4, 4), dtype=sv.cdtype) + perr = np.zeros((C, S), dtype=np.float64) + for ci, c in enumerate(circuits): + for s, op in enumerate(c.ops): + if len(op.qubits) == 1: + code[ci, s] = op.qubits[0] + M1[ci, s] = op.matrix + if T: + perr[ci, s] = noise.p1_of(op.qubits[0], op.count) + else: + a, b = op.qubits + code[ci, s] = 1024 + a * 512 + b # unique code per ordered pair + M2[ci, s] = op.matrix + if T: + perr[ci, s] = noise.p2_of(a, b) + M1d, M2d = xp.asarray(M1), xp.asarray(M2) # one host-to-device copy per batch + rows_of = np.arange(C * R).reshape(C, R) + noisy_rows = rows_of[:, 1:] + P1d = xp.asarray(PAULI_1Q.astype(sv.cdtype)) if T else None + P2d = xp.asarray(PAULI_2Q.astype(sv.cdtype)) if T else None + + for s in range(S): + col = code[:, s] + uniq, inv = np.unique(col, return_inverse=True) + for g, cval in enumerate(uniq): + if cval < 0: + continue + cis = np.nonzero(inv == g)[0] + qubits = (int(cval),) if cval < 1024 else ((int(cval) - 1024) // 512, (int(cval) - 1024) % 512) + src = M1d if len(qubits) == 1 else M2d + if cis.size == C: + mats = src[:, s] + idx = None + else: + cis_d = xp.asarray(cis) + mats = src[cis_d, s] + idx = xp.asarray(rows_of[cis].reshape(-1)) + if R > 1: + mats = xp.repeat(mats, R, axis=0) + sv.apply(qubits, mats, idx) + + if T: # Pauli errors on the noisy rows of these circuits + ps = perr[cis, s] + if np.any(ps > 0): + hit = rng.random((cis.size, T)) < ps[:, None] + if hit.any(): + sel = noisy_rows[cis][hit] + table = P1d if len(qubits) == 1 else P2d + which = rng.integers(0, table.shape[0], size=sel.size) + sv.apply(qubits, table[xp.asarray(which)], xp.asarray(sel)) + + # ---- reduce on device: probabilities -> marginal over measured qubits -> average trajectories + meas_sets = {c.measured for c in circuits} + if len(meas_sets) != 1: + raise ValueError("all circuits in a batch must measure the same qubits; plan_batches groups them") + meas, keep, drop = _marginal_axes(n, circuits[0].measured) + m = len(meas) + probs = (sv.psi.real ** 2 + sv.psi.imag ** 2).reshape((C * R,) + (2,) * n) + if drop: + probs = probs.sum(axis=drop) + kept_sorted = sorted(keep) + order = [kept_sorted.index(a) + 1 for a in keep] + probs = xp.transpose(probs, [0] + order).reshape(C, R, 2 ** m) + ideal = probs[:, 0, :] + noisy_p = None + if T: + noisy_p = probs[:, 1:, :].mean(axis=1) + ro = [noise.readout_of(q) for q in meas] + if any(a or b for a, b in ro): + t = noisy_p.reshape((C,) + (2,) * m) + for i, (p01, p10) in enumerate(ro): # clbit i sits on axis m - i + if not (p01 or p10): + continue + conf = xp.asarray(np.array([[1 - p01, p10], [p01, 1 - p10]], dtype=t.dtype)) + ax = m - i + t = xp.moveaxis(xp.tensordot(t, conf, axes=([ax], [1])), -1, ax) + noisy_p = t.reshape(C, 2 ** m) + backend.synchronize() + stats = {"wall_s": time.perf_counter() - t0, "kernel_calls": sv.kernel_calls, "rows": C * R, "steps": S, + "n_qubits": n} + return BatchResult(backend.asnumpy(ideal), None if noisy_p is None else backend.asnumpy(noisy_p), m, stats) diff --git a/westquant_cudaq/batched/ir.py b/westquant_cudaq/batched/ir.py new file mode 100644 index 0000000..7a16371 --- /dev/null +++ b/westquant_cudaq/batched/ir.py @@ -0,0 +1,143 @@ +"""Minimal circuit IR for batched simulation: a list of dense 1- and 2-qubit unitaries. + +Conventions (identical to Qiskit): +* qubit 0 is the least significant bit of a basis-state index; +* the matrix of a gate on ``qubits=(a, b)`` acts on the index ``bit(a) + 2*bit(b)``, i.e. ``b`` is the more + significant qubit and the matrix equals ``kron(U_b, U_a)`` for a product gate. + +``from_qiskit`` fuses gates (see :func:`fuse_ops`). On a GPU every op is one batched kernel call for the whole batch, +so fusion directly raises throughput. +""" +from __future__ import annotations + +import hashlib +from dataclasses import dataclass, field +from typing import Any, Iterable, Sequence + +import numpy as np + +_I2 = np.eye(2, dtype=np.complex128) + + +@dataclass(frozen=True) +class Op: + qubits: tuple[int, ...] # 1 or 2 qubits; matrix index ordering as in the module docstring + matrix: np.ndarray # (2**k, 2**k) complex128 + name: str = "u" + count: int = 1 # number of original gates fused into this op (used to compose noise) + + +@dataclass +class CircuitIR: + num_qubits: int + ops: list[Op] + measured: tuple[int, ...] | None = None # qubits read out (in classical-bit order); None = all qubits + circuit_id: str = "" + meta: dict[str, Any] = field(default_factory=dict) + fusion: str = "full" # "full" | "1q" | "none" (see fuse_ops) + + @property + def n_two_qubit(self) -> int: + return sum(1 for op in self.ops if len(op.qubits) == 2) + + def structure_key(self) -> str: + """Hash of (n, sequence of target tuples). Circuits with equal keys batch without any gather/scatter.""" + h = hashlib.blake2b(digest_size=12) + h.update(str(self.num_qubits).encode()) + for op in self.ops: + h.update(b"|" + ",".join(map(str, op.qubits)).encode()) + return h.hexdigest() + + +def fuse_ops(num_qubits: int, ops: Iterable[Op], mode: str = "full") -> list[Op]: + """Exact gate fusion. + + * ``"1q"``: merge consecutive single-qubit gates on the same qubit (``count`` = number merged). Safe with + depolarizing noise, because a depolarizing channel commutes with single-qubit unitaries on its qubit. + * ``"full"``: additionally absorb pending single-qubit gates into the next two-qubit gate on that qubit. Fewest + kernel launches; use for ideal (noise-free) simulation only. + * ``"none"``: no fusion. + """ + ops = list(ops) + if mode == "none": + return ops + pending: dict[int, tuple[np.ndarray, int]] = {} + out: list[Op] = [] + + def flush(q: int) -> None: + if q in pending: + m, c = pending.pop(q) + out.append(Op((q,), m, "fused1", c)) + + for op in ops: + if len(op.qubits) == 1: + q = op.qubits[0] + m, c = pending.get(q, (_I2, 0)) + pending[q] = (op.matrix @ m, c + op.count) + elif len(op.qubits) == 2: + a, b = op.qubits + if mode == "1q": + flush(a); flush(b) + out.append(op) + continue + pa, pb = pending.pop(a, None), pending.pop(b, None) + if pa is None and pb is None: + out.append(op) + else: + pre = np.kron(_I2 if pb is None else pb[0], _I2 if pa is None else pa[0]) + out.append(Op((a, b), op.matrix @ pre, "fused2", op.count)) + else: + raise ValueError(f"only 1- and 2-qubit ops are supported, got {op.qubits}") + for q in sorted(pending): + flush(q) + return out + + +def from_qiskit(qc, *, fusion: str = "full", circuit_id: str = "", meta: dict | None = None, + basis_gates: Sequence[str] = ("u", "cx")) -> CircuitIR: + """Convert a Qiskit ``QuantumCircuit`` (final measurements allowed) to :class:`CircuitIR`. + + Gates on more than two qubits are decomposed with ``transpile(..., basis_gates, optimization_level=0)``. + Mid-circuit measurements, resets and classical control are rejected: use a trajectory simulator for those. + """ + from qiskit import transpile + from qiskit.quantum_info import Operator + + measured: list[tuple[int, int]] = [] + for inst in qc.data: + if inst.operation.name == "measure": + measured.append((qc.find_bit(inst.clbits[0]).index, qc.find_bit(inst.qubits[0]).index)) + base = qc.remove_final_measurements(inplace=False) + if any(i.operation.name in ("measure", "reset") or getattr(i.operation, "condition", None) is not None + for i in base.data): + raise ValueError("mid-circuit measurement, reset or classical control is not supported") + if any(i.operation.num_qubits > 2 for i in base.data if i.operation.name != "barrier"): + base = transpile(base, basis_gates=list(basis_gates), optimization_level=0) + + ops: list[Op] = [] + for inst in base.data: + name = inst.operation.name + if name in ("barrier", "delay") or name.startswith("save_"): + continue + qs = tuple(base.find_bit(q).index for q in inst.qubits) + if name == "id": + continue + mat = np.asarray(Operator(inst.operation).data, dtype=np.complex128) + ops.append(Op(qs, mat, name)) + ops = fuse_ops(base.num_qubits, ops, fusion) + meas = tuple(q for _c, q in sorted(measured)) if measured else None + return CircuitIR(base.num_qubits, ops, meas, circuit_id, dict(meta or {}), fusion) + + +def from_qasm(text: str, **kw) -> CircuitIR: + """Parse OpenQASM 2 (built-in) or OpenQASM 3 (needs ``qiskit-qasm3-import``) and convert.""" + head = text.lstrip()[:40].upper() + if head.startswith("OPENQASM 3"): + from qiskit import qasm3 + + qc = qasm3.loads(text) + else: + from qiskit import qasm2 + + qc = qasm2.loads(text, custom_instructions=qasm2.LEGACY_CUSTOM_INSTRUCTIONS) + return from_qiskit(qc, **kw) diff --git a/westquant_cudaq/batched/noise.py b/westquant_cudaq/batched/noise.py new file mode 100644 index 0000000..6d700c2 --- /dev/null +++ b/westquant_cudaq/batched/noise.py @@ -0,0 +1,114 @@ +"""Pauli (depolarizing) trajectory noise for the batched simulator. + +Model, per gate location: +* 1-qubit gate on q: with probability p1(q) a uniformly random Pauli in {X, Y, Z} is applied after the gate; +* 2-qubit gate on (a, b): with probability p2(a, b) a uniformly random non-identity two-qubit Pauli (15 choices); +* readout: bit flips with p(1|0) and p(0|1) per measured qubit, applied to the averaged distribution (exact, linear). + +Averaging many trajectories converges to the depolarizing-channel density matrix. A fused single-qubit op that +replaced ``m`` gates receives the exactly composed error probability (depolarizing channels commute with +single-qubit unitaries on their own qubit, so the composition is exact). + +``from_calibration`` converts average gate *infidelities* r (what most backends report) to Pauli error probabilities +via p = r (d + 1) / d with d = 2 (one qubit) or d = 4 (two qubits). +""" +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from typing import Any, Mapping + +import numpy as np + +_X = np.array([[0, 1], [1, 0]], dtype=np.complex128) +_Y = np.array([[0, -1j], [1j, 0]], dtype=np.complex128) +_Z = np.array([[1, 0], [0, -1]], dtype=np.complex128) +_I = np.eye(2, dtype=np.complex128) +PAULI_1Q = np.stack([_X, _Y, _Z]) # (3, 2, 2) +PAULI_2Q = np.stack([np.kron(pb, pa) for pb in (_I, _X, _Y, _Z) for pa in (_I, _X, _Y, _Z)][1:]) # (15, 4, 4) + + +def _lookup(table: float | Mapping, key, default: float = 0.0) -> float: + if isinstance(table, (int, float)): + return float(table) + for k in (key, str(key), ",".join(map(str, key)) if isinstance(key, tuple) else None): + if k is not None and k in table: + return float(table[k]) + if isinstance(key, tuple) and len(key) == 2: + rev = (key[1], key[0]) + for k in (rev, ",".join(map(str, rev))): + if k in table: + return float(table[k]) + return float(table.get("default", default)) if isinstance(table, Mapping) else default + + +@dataclass +class PauliNoiseModel: + p1: float | dict = 0.0 + p2: float | dict = 0.0 + readout: float | dict = 0.0 # symmetric flip prob, or {q: [p(1|0), p(0|1)]} + label: str = "pauli" + extra: dict = field(default_factory=dict) + + # ---- per-location probabilities + def p1_of(self, q: int, count: int = 1) -> float: + p = _lookup(self.p1, q) + if count <= 1 or p == 0.0: + return p + f = (1.0 - 4.0 * p / 3.0) ** count # depolarizing shrink factor composes multiplicatively + return 0.75 * (1.0 - f) + + def p2_of(self, a: int, b: int) -> float: + return _lookup(self.p2, (a, b)) + + def readout_of(self, q: int) -> tuple[float, float]: + if isinstance(self.readout, (int, float)): + return float(self.readout), float(self.readout) + v = self.readout.get(q, self.readout.get(str(q), self.readout.get("default", 0.0))) + if isinstance(v, (int, float)): + return float(v), float(v) + return float(v[0]), float(v[1]) + + @property + def is_trivial(self) -> bool: + def zero(t): + return (t == 0) if isinstance(t, (int, float)) else all( + (np.all(np.asarray(v, dtype=float) == 0)) for v in t.values()) + return zero(self.p1) and zero(self.p2) + + # ---- (de)serialisation + def to_dict(self) -> dict[str, Any]: + def enc(t): + if isinstance(t, (int, float)): + return t + return {(",".join(map(str, k)) if isinstance(k, tuple) else str(k)): v for k, v in t.items()} + return {"p1": enc(self.p1), "p2": enc(self.p2), "readout": enc(self.readout), "label": self.label, + "extra": self.extra} + + def digest(self) -> str: + return hashlib.sha256(json.dumps(self.to_dict(), sort_keys=True).encode()).hexdigest()[:16] + + @classmethod + def from_dict(cls, d: Mapping[str, Any]) -> "PauliNoiseModel": + return cls(d.get("p1", 0.0), d.get("p2", 0.0), d.get("readout", 0.0), d.get("label", "pauli"), + dict(d.get("extra", {}))) + + @classmethod + def from_json(cls, path: str) -> "PauliNoiseModel": + with open(path) as f: + d = json.load(f) + return cls.from_calibration(d) if "gate_error_1q" in d or "gate_error_2q" in d else cls.from_dict(d) + + @classmethod + def from_calibration(cls, cal: Mapping[str, Any], label: str = "calibration") -> "PauliNoiseModel": + """``cal`` keys: ``gate_error_1q`` {q: r}, ``gate_error_2q`` {"a,b": r}, ``readout_error`` {q: r or [p01, p10]}. + + Scalars are accepted for any key (uniform error).""" + def conv(t, d): + k = (d + 1) / d + if isinstance(t, (int, float)): + return min(1.0, float(t) * k) + return {kk: min(1.0, float(v) * k) for kk, v in t.items()} + return cls(conv(cal.get("gate_error_1q", 0.0), 2), conv(cal.get("gate_error_2q", 0.0), 4), + cal.get("readout_error", 0.0), label, {"source": cal.get("source", "")}) diff --git a/westquant_cudaq/batched/runner.py b/westquant_cudaq/batched/runner.py new file mode 100644 index 0000000..bda03dd --- /dev/null +++ b/westquant_cudaq/batched/runner.py @@ -0,0 +1,164 @@ +"""Batch planning and record generation for data campaigns. + +``BatchRunner`` takes many :class:`CircuitIR`, groups them into device-sized batches (same qubit count and measured +register; sorted by structure so identical layouts share a batch), simulates, and yields one flat record per circuit, +ready for Parquet. Records carry the provenance needed to rebuild the run (versions, device, seeds, noise digest, +Slurm ids). +""" +from __future__ import annotations + +import json +import os +import platform +import socket +import time +import uuid +from dataclasses import dataclass, field +from typing import Iterable, Iterator, Sequence + +import numpy as np + +from .backend import Backend, get_backend +from .engine import simulate_batch +from .ir import CircuitIR +from .noise import PauliNoiseModel + +BYTES_FACTOR = 5 # psi + gathered rows + two transposed temporaries + probabilities, per row + + +def max_rows(backend: Backend, n: int, dtype: str = "complex64", safety: float = 0.6, cap: int = 1 << 16) -> int: + per_row = (2 ** n) * np.dtype(dtype).itemsize * BYTES_FACTOR + return int(max(1, min(cap, backend.free_bytes() * safety // per_row))) + + +def plan_batches(circuits: Sequence[CircuitIR], rows_per_circuit: int, cap_rows) -> list[list[int]]: + """Indices grouped by (n, measured qubits), sorted by structure key, chunked to ``cap_rows(n) // rows_per_circuit``.""" + groups: dict[tuple, list[int]] = {} + for i, c in enumerate(circuits): + groups.setdefault((c.num_qubits, c.measured), []).append(i) + batches: list[list[int]] = [] + for (n, _meas), idx in sorted(groups.items(), key=lambda kv: kv[0][0]): + idx.sort(key=lambda i: (circuits[i].structure_key(), len(circuits[i].ops))) + per = max(1, cap_rows(n) // rows_per_circuit) + batches += [idx[k:k + per] for k in range(0, len(idx), per)] + return batches + + +def record_schema(): + """Fixed Parquet schema so every shard merges cleanly (absent optional fields are null).""" + import pyarrow as pa + + s, i32, i64, f64 = pa.string(), pa.int32(), pa.int64(), pa.float64() + fields = [("circuit_id", s), ("n_qubits", i32), ("n_measured", i32), ("n_ops", i32), ("n_2q", i32), + ("structure_key", s), ("fusion", s), ("batch_index", i32), ("batch_size", i32), ("batch_rows", i32), + ("batch_kernel_calls", i32), ("wall_s_share", f64), ("meta", s), ("run_id", s), ("westquant_cudaq", s), + ("qiskit", s), ("array_backend", s), ("gpu_name", s), ("gpu_total_bytes", i64), ("cupy", s), ("numpy", s), + ("dtype", s), ("noise_digest", s), ("noise_label", s), ("trajectories", i32), ("shots", i32), + ("seed", i64), ("host", s), ("python", s), ("slurm_job_id", s), ("slurm_array_task_id", s), + ("slurm_procid", s)] + for tag in ("ideal", "noisy"): + fields += [(f"{tag}_top_bitstrings", pa.list_(s)), (f"{tag}_top_probs", pa.list_(f64)), + (f"{tag}_probs", pa.list_(pa.float32())), (f"{tag}_success_prob", f64)] + fields += [("counts", s), ("hellinger_fidelity", f64)] + return pa.schema(fields) + + +def _bitstring(index: int, m: int) -> str: + return format(index, f"0{m}b") if m else "" + + +def hellinger_fidelity(p: np.ndarray, q: np.ndarray) -> np.ndarray: + return np.sum(np.sqrt(np.clip(p, 0, None) * np.clip(q, 0, None)), axis=-1) ** 2 + + +@dataclass +class BatchRunner: + backend: Backend = field(default_factory=get_backend) + dtype: str = "complex64" + noise: PauliNoiseModel | None = None + trajectories: int = 0 + shots: int = 0 + topk: int = 16 + store_full_up_to: int = 10 # store the full distribution when measured bits <= this + seed: int = 1234 + safety: float = 0.6 + run_id: str = field(default_factory=lambda: uuid.uuid4().hex[:12]) + + def provenance(self) -> dict: + from .. import __version__ as wqc_version + try: + import qiskit + qv = qiskit.__version__ + except Exception: # noqa: BLE001 + qv = None + return {"run_id": self.run_id, "westquant_cudaq": wqc_version, "qiskit": qv, "array_backend": self.backend.name, + **{k: v for k, v in self.backend.info.items()}, "dtype": self.dtype, + "noise_digest": self.noise.digest() if self.noise else None, + "noise_label": self.noise.label if self.noise else None, "trajectories": self.trajectories, + "shots": self.shots, "seed": self.seed, "host": socket.gethostname(), "python": platform.python_version(), + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), "slurm_array_task_id": os.environ.get("SLURM_ARRAY_TASK_ID"), + "slurm_procid": os.environ.get("SLURM_PROCID")} + + def run(self, circuits: Sequence[CircuitIR]) -> Iterator[dict]: + rng = np.random.default_rng(self.seed) + noisy = self.noise is not None and self.trajectories > 0 + R = 1 + (self.trajectories if noisy else 0) + prov = self.provenance() + batches = plan_batches(circuits, R, lambda n: max_rows(self.backend, n, self.dtype, self.safety)) + for b_i, idx in enumerate(batches): + batch = [circuits[i] for i in idx] + res = simulate_batch(batch, self.backend, dtype=self.dtype, noise=self.noise if noisy else None, + trajectories=self.trajectories if noisy else 0, rng=rng) + m = res.n_measured + share = res.stats["wall_s"] / len(batch) + fid = hellinger_fidelity(res.ideal, res.noisy) if res.noisy is not None else None + for j, c in enumerate(batch): + rec = {"circuit_id": c.circuit_id, "n_qubits": c.num_qubits, "n_measured": m, "n_ops": len(c.ops), + "n_2q": c.n_two_qubit, "structure_key": c.structure_key(), "fusion": c.fusion, + "batch_index": b_i, "batch_size": len(batch), "batch_rows": res.stats["rows"], + "batch_kernel_calls": res.stats["kernel_calls"], "wall_s_share": share, + "meta": json.dumps(c.meta, sort_keys=True, default=str), **prov} + for tag, dist in (("ideal", res.ideal[j]), ("noisy", None if res.noisy is None else res.noisy[j])): + if dist is None: + continue + k = min(self.topk, dist.size) + top = np.argsort(-dist)[:k] + rec[f"{tag}_top_bitstrings"] = [_bitstring(int(t), m) for t in top] + rec[f"{tag}_top_probs"] = [float(dist[t]) for t in top] + if m <= self.store_full_up_to: + rec[f"{tag}_probs"] = dist.astype(np.float32).tolist() + exp = c.meta.get("expected") + if exp: + exps = [exp] if isinstance(exp, str) else list(exp) + rec[f"{tag}_success_prob"] = float(sum(dist[int(e, 2)] for e in exps)) + if self.shots and (tag == "noisy" or res.noisy is None): + p = np.clip(dist.astype(np.float64), 0, None) + counts = rng.multinomial(self.shots, p / p.sum()) + nz = np.nonzero(counts)[0] + rec["counts"] = json.dumps({_bitstring(int(t), m): int(counts[t]) for t in nz}) + if fid is not None: + rec["hellinger_fidelity"] = float(fid[j]) + yield rec + + def run_to_parquet(self, circuits: Sequence[CircuitIR], out_path: str, flush_every: int = 20000) -> dict: + """Stream records to ``out_path`` (written to ``out_path + '.tmp'`` then renamed: restart-safe).""" + import pyarrow as pa + import pyarrow.parquet as pq + + tmp = out_path + ".tmp" + schema = record_schema() + unknown = set() + writer = pq.ParquetWriter(tmp, schema, compression="zstd") + buf, n_rec, t0 = [], 0, time.perf_counter() + for rec in self.run(circuits): + unknown |= set(rec) - set(schema.names) + buf.append(rec) + if len(buf) >= flush_every: + writer.write_table(pa.Table.from_pylist(buf, schema=schema)); n_rec += len(buf); buf = [] + if buf: + writer.write_table(pa.Table.from_pylist(buf, schema=schema)); n_rec += len(buf) + writer.close() + os.replace(tmp, out_path) + if unknown: + raise KeyError(f"record fields missing from record_schema(): {sorted(unknown)}") + return {"records": n_rec, "wall_s": time.perf_counter() - t0, "out": out_path}