From 0ff4e8be1beb6070f2cb78ca6460640a94eba484 Mon Sep 17 00:00:00 2001 From: BoltzmannEntropy <91342039+BoltzmannEntropy@users.noreply.github.com> Date: Sun, 2 Aug 2026 06:29:52 +0300 Subject: [PATCH] benches: add mlxQ statevector baseline to the third-party Pauli benchmark Adds an exact (non-propagation) baseline to benches/third_party/pauli_prop: mlxQ, an MLX statevector simulator for Apple-silicon GPUs. Wherever the 2^n statevector fits, its per-step cost is independent of operator growth, so it locates the lattice size below which Pauli propagation is not worth running (~25-28 qubits against monoprop on an M1 Max, 32 GB). - backends.py: run_mlxq() following the existing per-backend pattern (lazy imports, one Trotter step + expectation per timed point); registered in MEMORY_METRICS/LABELS/THREAD_VARS and a new APPLE_BACKENDS group. - run_model.py / run_one.py: dispatch + choices for the new group. Also make _num_threads portable to macOS (os.sched_getaffinity does not exist there). - README: install and usage note; the backend is deliberately not in the uv project's dependencies since MLX wheels only exist for Apple silicon. Cross-validated against monoprop at 3x3/3x4: expectation values agree to ~2.6e-5 (monoprop's lower_atol=1e-6 truncation plus mlxQ's complex64). --- benches/third_party/README.md | 20 ++++++ benches/third_party/pauli_prop/backends.py | 77 ++++++++++++++++++++- benches/third_party/pauli_prop/run_model.py | 8 ++- benches/third_party/pauli_prop/run_one.py | 13 +++- 4 files changed, 114 insertions(+), 4 deletions(-) diff --git a/benches/third_party/README.md b/benches/third_party/README.md index bf0a1e6..4336556 100644 --- a/benches/third_party/README.md +++ b/benches/third_party/README.md @@ -43,6 +43,26 @@ Any lattice size works: `--nx/--ny` override `settings.json`, and the observable lattice (the central horizontally-adjacent bond, which is the committed `[20, 21]` at 6x6) rather than staying pinned to indices that a resize would invalidate. +### Optional: exact statevector baseline (Apple silicon) + +`--backends ... mlxq` adds [mlxQ](https://github.com/BoltzmannEntropy/Qupertino), an +MLX-based statevector simulator, as an exact non-propagation baseline. Wherever the +2^n statevector fits in memory, its per-step cost does not depend on the operator, so +it locates the lattice size below which propagation is not worth running — on an +M1 Max (32 GB) the crossover against monoprop sits at roughly 25–28 qubits, and a +36-qubit statevector (~0.5 TB) is out of reach entirely. It is not part of this uv +project's dependencies because its MLX dependency only exists on Apple silicon; +install it there with + +```bash +uv pip install "mlxq @ git+https://github.com/BoltzmannEntropy/Qupertino" +``` + +`MLXQ_METAL_KERNELS=1` selects its hand-written Metal kernels (its fastest tier; the +first step of a process then carries one-off kernel compilation, which the fixed-size +benchmark's step-0 drop already absorbs). Its memory column reports the exact +statevector footprint, and its expectation values are exact up to complex64 (~1e-6). + ### Scaling with lattice size ```bash diff --git a/benches/third_party/pauli_prop/backends.py b/benches/third_party/pauli_prop/backends.py index 447fe5c..c6883a1 100644 --- a/benches/third_party/pauli_prop/backends.py +++ b/benches/third_party/pauli_prop/backends.py @@ -41,6 +41,7 @@ "Qiskit pauli-prop": "process RSS growth (no memory accounting exposed)", "cuPauliProp (GPU)": "GPU memory pool in use", "PauliPropagation.jl": "Base.summarysize of the Pauli sum", + "mlxQ statevector": "statevector size (2^n complex64 amplitudes)", } @@ -247,12 +248,84 @@ def _pack_pauli_string( return out +def run_mlxq(settings: Settings) -> BackendResult: + """mlxQ: an exact statevector baseline on the Apple-silicon GPU (via MLX). + + Not a Pauli-propagation engine — it bounds them from below: wherever the + 2^n statevector fits, its per-step cost is size-independent of the operator, + so it locates the size under which propagation is not worth running (about + 25-28 qubits against monoprop on an M1 Max). Install on Apple silicon with + `pip install "mlxq @ git+https://github.com/BoltzmannEntropy/Qupertino"`; + MLXQ_METAL_KERNELS=1 selects its hand-written Metal kernels. Expectation + values are exact up to complex64 (~1e-6); the memory column is the exact + statevector footprint. + """ + import mlx.core as mx + from mlxq import shaders + from mlxq.gates import RX + from mlxq.sim import StateVectorSimulator + + nq = settings.num_qubits + sim = StateVectorSimulator(nq) + edges = grid_edges(settings.nx, settings.ny) + a, b = settings.observable_qubits + metal = shaders.metal_enabled() + rx_gate = RX(settings.theta_x) + + # The all-qubit RZ layer as one cached diagonal (the same construction the + # simulator uses internally for its fused ZZ layer): the phase per basis + # state is exp(-i*(theta_z/2)*sum_q s_q) with s_q = +-1 per bit. + idx = mx.arange(1 << nq, dtype=mx.uint32) + acc = mx.zeros((1 << nq,), dtype=mx.int32) + for q in range(nq): + acc = acc + (1 - 2 * ((idx >> (nq - 1 - q)) & 1).astype(mx.int32)) + ang = (-settings.theta_z / 2.0) * acc.astype(mx.float32) + z_phase = mx.cos(ang).astype(mx.complex64) + 1j * mx.sin(ang).astype(mx.complex64) + mx.eval(z_phase) + + q0, q1 = sorted((a, b)) + state_mb = (1 << nq) * 8 / 1024**2 + + def step(_step_idx: int) -> tuple[float, int, float]: + # apply_zz_layer(theta) applies exp(-i*theta*sum Z_a Z_b): Qiskit's + # RZZ(theta_zz) is theta_zz/2 in that convention. + sim.apply_zz_layer(settings.theta_zz / 2.0, edges) + sim.state = sim.state * z_phase + if metal: + sim.state = shaders.rx_layer_all(sim.state, nq, settings.theta_x) + else: + for q in range(nq): + sim.apply_single(rx_gate, q) + # = P(bits agree) - P(bits disagree). MLX is lazy; float() + # forces the whole step's evaluation, so the timer sees the real work. + t = mx.reshape(sim.state, (1 << q0, 2, 1 << (q1 - q0 - 1), 2, -1)) + p = mx.abs(t) ** 2 + expval = float( + mx.sum(p[:, 0, :, 0, :]) + + mx.sum(p[:, 1, :, 1, :]) + - mx.sum(p[:, 0, :, 1, :]) + - mx.sum(p[:, 1, :, 0, :]) + ) + if nq >= 26: + # State-sized temporaries otherwise accumulate in MLX's buffer pool + # and push a 32 GB machine into unified-memory thrashing. + mx.clear_cache() + return expval, 1 << nq, state_mb + + result = _run_steps(settings, LABELS["mlxq"], step) + result.operator_memory_mb = state_mb + return result + + # Backend name (as passed on a command line) -> runner. The Julia backend is not # here: it is a separate process driven by run_scaling.jl. CPU_BACKENDS = ("monoprop", "ppvm", "qiskit") GPU_BACKENDS = ("cupauliprop",) +# Apple-silicon GPU via MLX. Not a Pauli-propagation engine: an exact statevector +# baseline that locates the size below which propagation is not worth running. +APPLE_BACKENDS = ("mlxq",) JULIA_BACKEND = "juliapp" -ALL_BACKENDS = (*CPU_BACKENDS, *GPU_BACKENDS, JULIA_BACKEND) +ALL_BACKENDS = (*CPU_BACKENDS, *GPU_BACKENDS, *APPLE_BACKENDS, JULIA_BACKEND) # The environment variable each backend takes its thread count from. Read both ways: the # sweep driver sets these to apply a per-backend cap, and the worker reads its own to record @@ -268,6 +341,7 @@ def _pack_pauli_string( "ppvm": ("RAYON_NUM_THREADS",), "qiskit": ("OMP_NUM_THREADS", "MKL_NUM_THREADS"), "cupauliprop": ("OMP_NUM_THREADS",), + "mlxq": (), # Apple GPU: no CPU thread knob JULIA_BACKEND: ("JULIA_NUM_THREADS",), } @@ -276,5 +350,6 @@ def _pack_pauli_string( "ppvm": "QuEra ppvm", "qiskit": "Qiskit pauli-prop", "cupauliprop": "cuPauliProp (GPU)", + "mlxq": "mlxQ statevector", JULIA_BACKEND: "PauliPropagation.jl", } diff --git a/benches/third_party/pauli_prop/run_model.py b/benches/third_party/pauli_prop/run_model.py index 4b8745c..793a690 100644 --- a/benches/third_party/pauli_prop/run_model.py +++ b/benches/third_party/pauli_prop/run_model.py @@ -45,7 +45,11 @@ def main() -> None: "--backends", nargs="+", default=["monoprop", "ppvm", "qiskit", "cupauliprop"], - choices=[*backend_mod.CPU_BACKENDS, *backend_mod.GPU_BACKENDS], + choices=[ + *backend_mod.CPU_BACKENDS, + *backend_mod.GPU_BACKENDS, + *backend_mod.APPLE_BACKENDS, + ], ) args = parser.parse_args() @@ -75,6 +79,8 @@ def main() -> None: ) elif backend == "cupauliprop": results[backend] = backend_mod.run_cupauliprop(settings) + elif backend == "mlxq": + results[backend] = backend_mod.run_mlxq(settings) # Step 0's runtime is dropped: it carries one-off warm-up, and plot_results.py # expects the runtime series to be one shorter than the step range. diff --git a/benches/third_party/pauli_prop/run_one.py b/benches/third_party/pauli_prop/run_one.py index 25c948e..c0bdb52 100644 --- a/benches/third_party/pauli_prop/run_one.py +++ b/benches/third_party/pauli_prop/run_one.py @@ -48,7 +48,10 @@ def _num_threads(backend: str) -> int: slurm = os.environ.get("SLURM_CPUS_PER_TASK") if slurm: return int(slurm) - return len(os.sched_getaffinity(0)) + # sched_getaffinity does not exist on macOS, which the mlxq backend needs. + if hasattr(os, "sched_getaffinity"): + return len(os.sched_getaffinity(0)) + return os.cpu_count() or 1 def main() -> None: @@ -56,7 +59,11 @@ def main() -> None: parser.add_argument( "--backend", required=True, - choices=[*backend_mod.CPU_BACKENDS, *backend_mod.GPU_BACKENDS], + choices=[ + *backend_mod.CPU_BACKENDS, + *backend_mod.GPU_BACKENDS, + *backend_mod.APPLE_BACKENDS, + ], ) parser.add_argument("--nx", type=int, default=None) parser.add_argument("--ny", type=int, default=None) @@ -100,6 +107,8 @@ def main() -> None: result = backend_mod.run_qiskit(settings, args.max_terms) elif args.backend == "cupauliprop": result = backend_mod.run_cupauliprop(settings) + elif args.backend == "mlxq": + result = backend_mod.run_mlxq(settings) else: # unreachable: argparse constrains the choices raise SystemExit(f"unknown backend {args.backend}")