Repository navigation
perf: batched spin-1/2 ensemble kernel for the NumPy fallback and incremental methyl time evolution - #22
Merged
Conversation
An 81-offset, 200-step, one-channel Liouville ensemble mirroring the methyl examples' refine_pulse problems, timed for value and value+gradient on the active backend (label from _accelerator._enabled), so the NumPy fallback's ensemble cost is re-measurable. Baseline on this machine: Rust 1.7 ms, NumPy 58 ms per gradient.
time_evolution_profiles in the methyl hub re-propagated every waveform prefix (O(N^2) Bloch steps per offset, 0.26 s on the NumPy fallback); it now carries the running Bloch state forward one step at a time, the same loop examples/reburp_pulse.py already uses. Trajectories are bitwise identical (tests/test_methyl_time_evolution.py pins every step against the prefix propagation on both backends) and the four methyl time-evolution PNGs are unchanged.
… fallback A private closed-form kernel that evaluates every member of a spin-1/2 Liouville-commutator ensemble (4x4 generators -i[H, .] with a Hermitian 2x2 H, vector states) in one vectorised pass: Pauli-decomposed eigenpairs, the same Daleckii-Krein sinc derivative the per-member path uses, 2x2 sandwich sweeps, and an ordered member accumulation with chunk carry under a fixed workspace budget. spin_half_inputs validates the whole ensemble first (validation preflight, 4x4 shapes, vector states, metadata, anti-Hermitian plus commutator-reconstruction gate on every expanded drift and waveform-scaled control) and returns None otherwise, so callers keep the existing path. Not dispatched yet. Parity: 1e-13 vs the per-member NumPy path, 1e-12 vs Rust, finite differences 1e-5; 81-member methyl-like gradient 8 ms vs 59 ms per-member.
…n the NumPy path ensemble_fidelity and ensemble_xy_and_gradient gain a third rung after both Rust rungs: only when the accelerator is disabled, spin_half_inputs marshals the whole ensemble and the batched kernel evaluates it in one pass (frozen entries zeroed exactly as the Rust rung does); any ineligible ensemble keeps the per-member loop. Native behaviour is unchanged, including phase-cycle ensembles that the Rust ensemble rungs decline (still one Rust call per member). tests/test_ensemble_dispatch.py pins one kernel call per evaluation on the fallback, zero on the native path, the dissipative fall-through, identical validation messages on both paths, and public-API parity with penalties and a freeze mask. The methyl refine_pulse objective (4 problems, 252 members) drops from 177 ms to 37 ms on the fallback with value/gradient parity 1.6e-14 / 2.3e-15.
…atched kernel Chunk sizes 1, 7 and 1000 give bitwise-identical values and gradients on both backends; the chunk policy follows the workspace budget (81 of 81, 1044, 17, and a single member beyond the budget still runs); the commutator-reconstruction gate is scale-relative (0.5e-12*s accepted, 2e-12*s declined for an anti-Hermitian non-commutator defect) and a Hermitian defect of 7.5e-13*s is declined by the anti-Hermitian gate first. Removing either gate in the kernel makes its own test fail (mutation receipts in the evidence log).
…al methyl time evolution The user manual's fallback cost model no longer claims a 15-20x GRAPE ensemble gap: spin-1/2 ensembles run in one batched closed-form pass (methyl_ensemble benchmark row quoted for both backends) while two-spin, density-matrix and relaxation problems still propagate member by member; the stale section 4-3 timings are dropped. CHANGELOG Unreleased gains the kernel, the incremental time evolution, the benchmark rows and the new tests.
Architecture items 3 and 4 describe the third fallback rung (_ensemble_kernel.py after the two Rust rungs, only when the accelerator is disabled) and that the kernel reuses _metadata_supported read-only. The local AGENTS.md drafts (git-excluded) carry the matching CODE MAP rows and dispatch-order convention with a passing reference check (REFS 25 STALE 0). Plan: .kbsi/plans/grape-fallback-batched.md
…at range edges The closed-form eigenvector normalisation squared the Hamiltonian components: around 1e-200 the squares underflowed to zero, which the r == 0 guard read as an exact degeneracy and replaced the evolution by the identity (wrong fidelity, no error), and around 1e155 they overflowed to NaN, although dt*H was ordinary and the per-member path was exact. The norm now comes from nested hypot and the eigenvectors are built from the unit direction, so every intermediate is O(1) at any scale. Found by the final adversarial review; test_kernel_is_scale_invariant_at_float_range_edges pins scales 1e-200..1e200 under errstate(all=raise) and fails on the previous kernel.
…-rescaled components Dividing the raw Hamiltonian components by a subnormal radius loses the unit direction (a subnormal carries few significant bits), so the batched kernel returned fidelities off by up to 1.25 for accepted finite inputs below ~1e-308 while the per-member path was exact. The direction now comes from the components rescaled by their largest magnitude and normalised with hypot, with the zero-radius member kept on the identity basis. Found by the second adversarial review round; test_kernel_direction_survives_subnormal_components pins 1e-310, 1e-320 and the smallest float in a mixed batch and fails on the previous kernel.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
NumPy-fallback follow-ups tracked after #21, executed from the reviewed plan
.kbsi/plans/grape-fallback-batched.md(deep-high review round 3 APPROVE).optimalcontrol/_ensemble_kernel.py) evaluates every member of a spin-1/2 Liouville-commutator GRAPE ensemble in one vectorised pass (Pauli eigenpairs, the same Daleckii-Krein sinc derivative, 2x2 sandwich sweeps, ordered accumulation with chunk carry under a 64 MiB workspace budget).ensemble.pydispatches to it only after both Rust rungs decline and only when the accelerator is disabled; native behaviour is byte-for-byte unchanged.time_evolution_profilescarries the running Bloch state instead of re-propagating every prefix (bitwise-identical trajectories, PNGs untouched).benchmarks/bench_grape_hotpath.py.Numbers (Apple M4 Pro)
test_below_minimum_lengthon the fallbackParity: kernel vs forced per-member path 1e-13 (value 3.5e-16 / gradient 1.8e-17 on an 18-member phase-cycled, 2-channel, 2-pair, 2-power, 3-drift ensemble), vs Rust 1e-12 (1.3e-15 / 4.5e-17), finite differences rtol 1e-5 (2.2e-8); chunk sizes 1/7/1000 bitwise identical; both eligibility gates proven by mutation (removing either makes its own test fail).
Verification
examples/outputclean (PNG hashes unchanged)Docs (user manual, CHANGELOG, CLAUDE.md) follow in the next commits of this PR.