dense-evolution 8.3.0__py3-none-win_amd64.whl

dense_evolution/backends/chunk/core.py

@@ -0,0 +1,432 @@
+import shutil
+import tempfile
+from typing import List, Optional
+
+import numpy as np
+
+from ._engine_imports import DenseSVSimulator
+from .guard import HAS_JAX, MemoryPressureError, SafeMemoryGuard
+from .geometry import MemoryChunker
+from .circuit_chunker import CircuitChunker
+from .kernels import (
+    _build_multi_chunk_runner, _build_distributed_chunk_runner,
+    _compile_multi_chunk_ops,
+)
+from .disk_overflow import run_disk_overflow_circuit
+
+__all__ = ["Chunk", "chunk1", "chunk2", "Chunk2Incrociato"]
+
+
+# ─────────────────────────────────────────────────────────────────────────────────
+# Chunk  (chunk2 / Chunk2Incrociato)
+# ─────────────────────────────────────────────────────────────────────────────────
+
+class Chunk:
+    """
+    Anti-OOM wrapper for large-qubit simulation.
+
+    Does NOT subclass DenseSVSimulator directly — the parent __init__ allocates
+    2**n_qubits elements immediately (17 GB for 30 qubits).
+
+    For n_qubits <= chunk_size_bits (the RAM-safe budget): a single inner
+    simulator is allocated and the logical qubit count is stored separately.
+
+    For n_qubits > chunk_size_bits: num_chunks separate chunk_size_bits-qubit
+    simulators are held in RAM simultaneously (see _dispatch_multi) — as
+    many chunks as actually fit in RAM at once, checked up front via
+    SafeMemoryGuard.check_allocation before anything is allocated.
+    Benchmark attributes (num_chunks, chunk_size_bits, dtype) are
+    forwarded transparently from the embedded MemoryChunker.
+
+    Past that RAM ceiling, allow_disk_overflow=True (default False) falls
+    back to disk-backed storage instead of raising MemoryPressureError --
+    see dense_evolution/backends/chunk/disk_overflow.py (Pednault et al.
+    2019, arXiv:1910.09534) for the phased execution this uses, and
+    docs/api/chunk.md for a real verified demo and the speed trade-off
+    (never more than 1-2 chunks materialized in RAM at once, but every
+    gate now pays disk I/O -- v1 is correctness-first, not fast).
+
+    A SafeMemoryGuard fires before any simulator is instantiated
+    (pre-allocation check) and is also embedded in CircuitChunker for
+    per-slice protection during execution (n_qubits <= chunk_size_bits path).
+
+    Parameters
+    ----------
+    n_qubits          : logical qubit count of the target circuit
+    chunk_size_gates  : gate-slice size for JIT compilation (default 500)
+    memory_threshold  : free-RAM fraction below which execution is blocked
+                        (default 0.15 = 15%)
+    use_float32       : forwarded to DenseSVSimulator
+    allow_disk_overflow : fall back to disk-backed chunks instead of
+                        raising MemoryPressureError when num_chunks
+                        chunks don't fit in RAM at once (default False)
+    disk_dir          : directory for the overflow .npy files (default:
+                        a fresh tempfile.mkdtemp(), removed by close())
+    """
+
+    def __init__(
+        self,
+        n_qubits: int,
+        chunk_size_gates:    int   = 500,
+        memory_threshold:    float = 0.15,
+        use_float32:         bool  = False,
+        allow_disk_overflow: bool  = False,
+        disk_dir:            Optional[str] = None,
+    ):
+        # 1. Geometry — purely RAM-based, no JAX allocation yet
+        self._mem_chunker     = MemoryChunker(n_qubits)
+        self._guard           = SafeMemoryGuard(threshold_pct=memory_threshold)
+
+        # 2. Logical qubit count (for circuit parsing)
+        self.n                = n_qubits
+        self.chunk_size_gates = chunk_size_gates
+        self._m                = n_qubits - self._mem_chunker.chunk_size_bits  # chunk-select qubit count (0 if num_chunks==1)
+
+        self._chunk_paths  = None
+        self._disk_dir     = None
+        self._owns_disk_dir = False
+
+        if self._mem_chunker.num_chunks == 1:
+            # 3a. Pre-allocation RAM check — block here rather than inside JAX
+            safe_q = min(n_qubits, self._mem_chunker.chunk_size_bits)
+            self._guard.check(f"Chunk.__init__ — allocating {safe_q}-qubit simulator")
+
+            # 4a. Physical simulator sized to what RAM can actually hold
+            self._inner_sim = DenseSVSimulator(
+                safe_q,
+                use_float32=use_float32,
+            )
+            self._chunk_sims = None
+            self._multi_chunk_runner = None
+
+            # 5a. Circuit chunker wired to the physical simulator, with same threshold
+            self._circuit_chunker = CircuitChunker(
+                simulator_instance=self._inner_sim,
+                memory_threshold=memory_threshold,
+            )
+        else:
+            # 3b. Sized pre-allocation check: num_chunks chunk-sized simulators
+            # held in RAM at once, plus ~2 chunks of headroom for the temporary
+            # arrays the cross-chunk gate-mixing math allocates at its peak.
+            num_chunks   = self._mem_chunker.num_chunks
+            per_chunk_mb = self._mem_chunker.memory_mb()
+            required_mb  = (num_chunks + 2) * per_chunk_mb
+            try:
+                self._guard.check_allocation(
+                    required_mb,
+                    f"Chunk.__init__ — allocating {num_chunks} chunks of "
+                    f"{self._mem_chunker.chunk_size_bits} qubits each",
+                )
+            except MemoryPressureError:
+                if not allow_disk_overflow:
+                    raise
+                self._init_disk_overflow(num_chunks, disk_dir)
+                return
+
+            # 4b. num_chunks independent chunk-sized simulators. Each one's own
+            # __init__ resets it to |0...0>: only chunk 0 should carry the
+            # amplitude-1 seed for the LOGICAL |0...0>, the rest must start
+            # at all-zero (direct .sv assignment — set_state/set_initial_state
+            # reject zero-norm vectors by design).
+            self._chunk_sims = [
+                DenseSVSimulator(self._mem_chunker.chunk_size_bits, use_float32=use_float32)
+                for _ in range(num_chunks)
+            ]
+            for sim in self._chunk_sims[1:]:
+                sim.sv = sim.xp.zeros(self._mem_chunker.chunk_dim, dtype=sim.dtype)
+
+            self._inner_sim        = None
+            self._circuit_chunker  = None
+
+            # 6b. JIT runner for the multi-chunk dispatch — built once here
+            # since num_chunks/m/chunk_size_bits are fixed for this
+            # instance's whole lifetime, reused by every run_chunk() call.
+            self._multi_chunk_runner = _build_multi_chunk_runner(
+                num_chunks, self._m, self._mem_chunker.chunk_size_bits)
+
+            # 7b. Distributed (multi-device) runner — built lazily on first
+            # run_chunk_distributed() call, not here: it requires
+            # jax.device_count() >= num_chunks, which most single-process
+            # uses of Chunk will never need or satisfy.
+            self._distributed_runner = None
+            self._distributed_mesh   = None
+
+    def _init_disk_overflow(self, num_chunks: int, disk_dir: Optional[str]) -> None:
+        """Fallback storage for allow_disk_overflow=True when num_chunks
+        chunks don't fit in RAM at once -- see disk_overflow.py. Each
+        chunk becomes its own .npy file instead of a live DenseSVSimulator;
+        chunk 0 seeded to the logical |0...0>, the rest zero, same
+        convention _chunk_sims uses in the in-RAM path."""
+        self._inner_sim           = None
+        self._circuit_chunker     = None
+        self._chunk_sims          = None
+        self._multi_chunk_runner  = None
+        self._distributed_runner  = None
+        self._distributed_mesh    = None
+
+        self._owns_disk_dir = disk_dir is None
+        self._disk_dir = disk_dir or tempfile.mkdtemp(prefix="dense_evolution_chunk_")
+        dtype = self._mem_chunker.dtype
+        chunk_dim = self._mem_chunker.chunk_dim
+        paths = []
+        for i in range(num_chunks):
+            arr = np.zeros(chunk_dim, dtype=dtype)
+            if i == 0:
+                arr[0] = 1.0
+            path = f"{self._disk_dir}/chunk_{i}.npy"
+            np.save(path, arr)
+            paths.append(path)
+        self._chunk_paths = paths
+
+    def close(self) -> None:
+        """Removes the disk-overflow directory, if this Chunk created one
+        (allow_disk_overflow=True with no explicit disk_dir). Safe to call
+        even if disk overflow was never used."""
+        if self._owns_disk_dir and self._disk_dir is not None:
+            shutil.rmtree(self._disk_dir, ignore_errors=True)
+            self._disk_dir = None
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass
+
+    # ── Benchmark-facing attribute forwarding ─────────────────────
+
+    @property
+    def num_chunks(self) -> int:
+        return self._mem_chunker.num_chunks
+
+    @property
+    def chunk_size_bits(self) -> int:
+        return self._mem_chunker.chunk_size_bits
+
+    @property
+    def chunk_dim(self) -> int:
+        return self._mem_chunker.chunk_dim
+
+    @property
+    def dtype(self):
+        return self._mem_chunker.dtype
+
+    @property
+    def memory_geometry(self) -> MemoryChunker:
+        return self._mem_chunker
+
+    # ── Simulator-facing forwarding ─────────────────────────────
+
+    @property
+    def sv(self):
+        """Current statevector. For num_chunks==1, the physical (chunk-sized)
+        simulator's own array. For num_chunks>1, the chunks concatenated in
+        ascending order — valid because of the MSB-first correspondence
+        between chunk index and the top `_m` logical qubits (see
+        _dispatch_multi's docstring). For the disk-overflow path, streams
+        each chunk's .npy file off disk one at a time to build the same
+        concatenation -- this DOES materialize the full (2**n,) array in
+        RAM, unlike run_chunk() itself; only use it on a size you know
+        fits, e.g. for a final readout after a run."""
+        if self._chunk_paths is not None:
+            return np.concatenate([np.load(p) for p in self._chunk_paths])
+        if self._chunk_sims is None:
+            return self._inner_sim.sv
+        xp = self._chunk_sims[0].xp
+        return xp.concatenate([sim.sv for sim in self._chunk_sims])
+
+    @sv.setter
+    def sv(self, value):
+        """Accepts a full-length (2**n,) statevector (e.g. the output of
+        NoiseModel.apply_to_sv called on `.sv`) and writes it back through
+        to the physical storage -- the inner simulator directly for
+        num_chunks==1, split back into per-chunk slices (same ascending
+        concatenation order as the getter) for num_chunks>1, or rewritten
+        to each chunk's .npy file for the disk-overflow path."""
+        chunk_dim = self._mem_chunker.chunk_dim
+        if self._chunk_paths is not None:
+            value = np.asarray(value)
+            for i, path in enumerate(self._chunk_paths):
+                np.save(path, value[i * chunk_dim:(i + 1) * chunk_dim])
+            return
+        if self._chunk_sims is None:
+            self._inner_sim.sv = value
+            return
+        xp = self._chunk_sims[0].xp
+        value = xp.asarray(value)
+        for i, sim in enumerate(self._chunk_sims):
+            sim.sv = value[i * chunk_dim:(i + 1) * chunk_dim]
+
+    def memory_mb(self) -> float:
+        """RAM used by the physical statevector(s) in MB -- 0 for the
+        disk-overflow path (see memory_geometry.memory_mb() for the
+        per-chunk on-disk size instead, and disk_overflow.py's own
+        docstring for why nothing (2**n_qubits,)-sized, or even
+        num_chunks-chunks-sized, is ever resident in RAM at once)."""
+        if self._chunk_paths is not None:
+            return 0.0
+        if self._chunk_sims is None:
+            return self._inner_sim.memory_mb()
+        return sum(sim.memory_mb() for sim in self._chunk_sims)
+
+    def get_probabilities(self):
+        """|amplitude|^2 for every basis state.
+
+        num_chunks==1: forwards to the inner DenseSVSimulator for parity
+        with its own get_probabilities().
+
+        num_chunks>1: concatenates the RAW statevectors first and normalizes
+        ONCE over the full array — NOT each chunk's own get_probabilities()
+        (that would independently renormalize each chunk's partial mass to
+        1, summing to num_chunks overall and destroying the relative
+        weighting between chunks). Disk-overflow path: same normalization,
+        chunks streamed from their .npy files instead of live arrays."""
+        if self._chunk_paths is not None:
+            full_sv = np.concatenate([np.load(p) for p in self._chunk_paths])
+        elif self._chunk_sims is None:
+            return self._inner_sim.get_probabilities()
+        else:
+            full_sv = np.concatenate([np.array(sim.sv) for sim in self._chunk_sims])
+        probs = np.abs(full_sv) ** 2
+        probs = np.clip(probs, 0.0, 1.0)
+        total = probs.sum()
+        if total > 1e-12:
+            probs /= total
+        return probs
+
+    def get_statevector(self):
+        """Full complex statevector, num_qubits logical qubits long
+        (2**n elements). num_chunks==1: forwards to the inner
+        DenseSVSimulator. num_chunks>1: raw chunks concatenated in order
+        (see `sv` property). Disk-overflow path: streamed from the .npy
+        files, same order."""
+        if self._chunk_paths is not None:
+            return np.concatenate([np.load(p) for p in self._chunk_paths])
+        if self._chunk_sims is None:
+            return self._inner_sim.get_statevector()
+        return np.concatenate([np.array(sim.sv, dtype=sim.dtype) for sim in self._chunk_sims])
+
+    # ── Multi-chunk gate dispatch (num_chunks > 1) ───────────────────
+
+    def _dispatch_multi(self, circuit: List) -> None:
+        """Executes *circuit* against the num_chunks>1 chunk representation
+        via one jax.lax.scan call over the whole (transpiled, GATE_IDS-
+        compiled) circuit — see _build_multi_chunk_step/_build_multi_chunk_runner
+        above for the kernel, and _compile_multi_chunk_ops for the encoding.
+
+        Convention: chunk index `c` (m = self._m bits, MSB-first, same
+        n-1-qubit convention as DenseSVSimulator) equals the value of the
+        top m logical qubits (indices [0, m)); chunk_sims[c].sv holds the
+        chunk_dim amplitudes for the remaining (local) qubits [m, n). This
+        makes full_sv.reshape(num_chunks, chunk_dim)[c] == chunk_sims[c].sv
+        exactly, since NumPy's row-major reshape splits a (2,)*n tensor on
+        the leading axes first — i.e. the most-significant qubits, matching
+        this simulator's MSB-first indexing throughout. Stacking
+        chunk_sims[i].sv into one (num_chunks, chunk_dim) array before the
+        scan, and unstacking after, holds exactly the same total element
+        count as the num_chunks separate arrays it replaces — the anti-OOM
+        property this class exists for is preserved, nothing (2**n_qubits,)
+        shaped is ever materialized."""
+        compiled_ops = _compile_multi_chunk_ops(circuit)
+        xp = self._chunk_sims[0].xp
+        stacked = xp.stack([sim.sv for sim in self._chunk_sims])
+        final = self._multi_chunk_runner(stacked, compiled_ops)
+        for i, sim in enumerate(self._chunk_sims):
+            sim.sv = final[i]
+
+    # ── Distributed multi-device gate dispatch (issue #1) ────────────
+
+    def dispatch_distributed(self, circuit: List) -> None:
+        """Executes *circuit* the same way _dispatch_multi does, but with
+        each chunk pinned to its own physical JAX device instead of all
+        chunks sharing one process's RAM — see
+        _build_distributed_chunk_step/_build_distributed_chunk_runner for
+        the kernel. Requires jax.device_count() >= num_chunks (v1 scope:
+        exactly one chunk per device); raises RuntimeError otherwise
+        rather than silently falling back to the single-process path,
+        since that would silently give up the whole point of calling this
+        method instead of run_chunk().
+
+        The (num_chunks, chunk_dim) logical array is never materialized
+        on any single device here (unlike _dispatch_multi, where it's one
+        process's RAM) -- each device holds and ever sees only its own
+        (chunk_dim,) row, exchanging edge data with its stride-partner
+        device via jax.lax.ppermute inside the kernel, not through this
+        Python method."""
+        import jax
+        if self._chunk_sims is None:
+            raise RuntimeError(
+                "dispatch_distributed() requires num_chunks > 1 "
+                "(this Chunk instance fits in a single chunk)."
+            )
+        num_chunks = self._mem_chunker.num_chunks
+        available = jax.device_count()
+        if available < num_chunks:
+            raise RuntimeError(
+                f"dispatch_distributed() needs >= {num_chunks} JAX devices "
+                f"(one per chunk), only {available} available. Force extra "
+                f"CPU devices for testing via the XLA_FLAGS environment "
+                f"variable: --xla_force_host_platform_device_count=N "
+                f"(set before the process starts, JAX's device count is "
+                f"fixed at first initialization)."
+            )
+        if self._distributed_runner is None:
+            self._distributed_runner, self._distributed_mesh = _build_distributed_chunk_runner(
+                num_chunks, self._m, self._mem_chunker.chunk_size_bits)
+
+        compiled_ops = _compile_multi_chunk_ops(circuit)
+        xp = self._chunk_sims[0].xp
+        stacked = xp.stack([sim.sv for sim in self._chunk_sims])
+        final = self._distributed_runner(stacked, compiled_ops)
+        for i, sim in enumerate(self._chunk_sims):
+            sim.sv = np.asarray(final[i])
+
+    # ── Public API ───────────────────────────────────────────
+
+    def run_chunk(
+        self,
+        circuit: List,
+        chunk_size_gates: Optional[int] = None,
+    ) -> None:
+
+        if self._chunk_paths is not None:
+            compiled_ops = _compile_multi_chunk_ops(circuit)
+            run_disk_overflow_circuit(
+                self._chunk_paths, compiled_ops, self._m, self._mem_chunker.chunk_size_bits)
+            return
+        if self._chunk_sims is not None:
+            self._dispatch_multi(circuit)
+            return
+        size = chunk_size_gates if chunk_size_gates is not None else self.chunk_size_gates
+        self._circuit_chunker.split_circuit(circuit, chunk_size=size)
+
+    def run_chunk_distributed(self, circuit: List) -> None:
+        """Like run_chunk(), but dispatches across a real JAX device mesh
+        (dispatch_distributed) instead of one process's RAM — issue #1.
+        Requires jax.device_count() >= num_chunks; raises RuntimeError
+        otherwise (see dispatch_distributed's docstring for how to test
+        this with simulated multi-device CPU)."""
+        self.dispatch_distributed(circuit)
+
+    def __repr__(self) -> str:
+        s = self._guard.status()
+        safe_qubits = self._inner_sim.n if self._inner_sim is not None else self._mem_chunker.chunk_size_bits
+        storage = f"disk ({self._disk_dir})" if self._chunk_paths is not None else "ram"
+        return (
+            f"Chunk(n_qubits={self.n}, "
+            f"safe_qubits={safe_qubits}, "
+            f"num_chunks={self.num_chunks}, "
+            f"chunk_size_bits={self.chunk_size_bits}, "
+            f"storage={storage}, "
+            f"dtype={self.dtype}, "
+            f"mem_per_chunk={self.memory_mb():.1f} MB, "
+            f"ram_free={s['free_pct']:.1f}%, "
+            f"has_jax={HAS_JAX})"
+        )
+
+
+# ─────────────────────────────────────────────────────────────────────────────────
+# Backward-compatibility aliases
+# ─────────────────────────────────────────────────────────────────────────────────
+chunk1           = MemoryChunker
+chunk2           = Chunk
+Chunk2Incrociato = Chunk

dense_evolution/backends/chunk/disk_overflow.py

@@ -0,0 +1,232 @@
+"""Disk-backed statevector overflow for Chunk -- Pednault et al. 2019
+(arXiv:1910.09534, "Leveraging Secondary Storage to Simulate Deep
+54-qubit Sycamore Circuits"): when num_chunks chunks don't all fit in
+RAM at once, keep the idle ones on disk as plain .npy files and only
+materialize, as a jax.Array, the small working set an individual gate
+actually needs -- one chunk for a local gate, two (an XOR-stride pair)
+for a chunk-mixing gate -- instead of the whole (num_chunks, chunk_dim)
+stack _dispatch_multi holds in RAM for the entire circuit.
+
+v1 scope (correctness-first, matching run_chunk_distributed's own staged
+rollout): processes one chunk / one pair at a time, no batching multiple
+pairs into one call for speed. This is strictly slower per gate than the
+in-RAM path -- it exists to make otherwise-impossible sizes possible at
+all, not to compete with it on speed. See docs/api/chunk.md for the real
+distinction and a verified demo.
+
+Every gate is one of the same 6 cases kernels.py's in-RAM
+_build_multi_chunk_step already classifies. This module reuses that
+classification (not a fresh derivation) via the same needs_comm split
+_build_distributed_chunk_step's own docstring documents:
+  - 1-qubit, q1 >= m -> LOCAL: needs only this one chunk.
+  - 2-qubit, q1 >= m and q2 >= m -> LOCAL: needs only this one chunk.
+  - 2-qubit, q1 < m, q2 >= m ("ctrl chunk-select, tgt local") -> LOCAL,
+    but conditional on this chunk's OWN absolute index (no partner
+    needed -- kernels.py's _case_2q_ctrl_chunk_tgt_local already
+    computes exactly this from an explicit `idxc`).
+  - everything else (1-qubit q1 < m, or 2-qubit q2 < m) -> MIX: needs
+    exactly the XOR-stride partner chunk.
+"""
+import numpy as np
+import jax.numpy as jnp
+
+from .kernels import (
+    _gate_matrix_elements, _case_1q_local, _case_2q_local_local,
+    _case_2q_ctrl_chunk_tgt_local,
+)
+
+__all__ = ["partition_ops_into_phases", "run_disk_overflow_circuit", "LocalPhase", "ConditionalPhase", "MixPhase"]
+
+
+class LocalPhase:
+    """A maximal run of consecutive gates with q1,q2 >= m (or unused) --
+    applied to each chunk independently, no partner chunk ever needed."""
+    __slots__ = ("ops",)
+
+    def __init__(self, ops):
+        self.ops = ops  # list of (g_id, q1, q2, param) python tuples
+
+
+class ConditionalPhase:
+    """A maximal run of consecutive 2-qubit gates with ctrl chunk-select
+    (q1 < m), tgt local (q2 >= m) -- applied to each chunk independently,
+    each op conditioned on that chunk's own absolute index (see
+    _case_2q_ctrl_chunk_tgt_local). Batched the same way LocalPhase is
+    (was: strictly one gate per phase object, so consecutive
+    ConditionalPhase-eligible gates did one full disk load/save cycle per
+    chunk PER GATE instead of once for the whole run -- prog.txt point 5b)."""
+    __slots__ = ("ops",)
+
+    def __init__(self, ops):
+        self.ops = ops  # list of (g_id, q1, q2, param) python tuples
+
+
+class MixPhase:
+    """A single gate that needs a partner chunk: 1-qubit with q1 < m, or
+    2-qubit with q2 < m (covers both "ctrl local/tgt chunk" and "ctrl
+    AND tgt both chunk-select"). `stride` is the mixing qubit (q1 for the
+    1-qubit case, q2 for the 2-qubit cases) -- chunk i always pairs with
+    chunk i ^ (1 << (m - 1 - stride))."""
+    __slots__ = ("op", "stride")
+
+    def __init__(self, op, stride):
+        self.op = op
+        self.stride = stride
+
+
+def partition_ops_into_phases(compiled_ops, m: int):
+    """compiled_ops: the (n_gates, 4) [g_id, q1, q2, param] array
+    _compile_multi_chunk_ops already builds (reused as-is). Returns a
+    list of LocalPhase/ConditionalPhase/MixPhase in circuit order. Pure
+    Python -- runs once per run_chunk() call, not traced/jitted."""
+    rows = np.asarray(compiled_ops)
+    phases = []
+    pending_local = []
+    pending_conditional = []
+
+    def flush_local():
+        if pending_local:
+            phases.append(LocalPhase(list(pending_local)))
+            pending_local.clear()
+
+    def flush_conditional():
+        if pending_conditional:
+            phases.append(ConditionalPhase(list(pending_conditional)))
+            pending_conditional.clear()
+
+    for row in rows:
+        g_id, q1, q2, param = int(row[0]), int(row[1]), int(row[2]), float(row[3])
+        is_2q = g_id >= 20
+        if not is_2q:
+            if q1 < m:
+                flush_local()
+                flush_conditional()
+                phases.append(MixPhase((g_id, q1, q2, param), q1))
+            else:
+                flush_conditional()
+                pending_local.append((g_id, q1, q2, param))
+        elif q1 < m and q2 >= m:
+            flush_local()
+            pending_conditional.append((g_id, q1, q2, param))
+        elif q2 < m:
+            flush_local()
+            flush_conditional()
+            phases.append(MixPhase((g_id, q1, q2, param), q2))
+        else:
+            flush_conditional()
+            pending_local.append((g_id, q1, q2, param))
+    flush_local()
+    flush_conditional()
+    return phases
+
+
+def _run_local_phase_on_chunk(chunk_arr, ops, m: int, k: int):
+    """chunk_arr: (chunk_dim,) array for one chunk. Applies every op via
+    the same _case_1q_local/_case_2q_local_local formulas the in-RAM
+    kernel uses, as a (1, chunk_dim) batch of size 1 -- both cases never
+    reference the batch dimension's meaning, so this is exact, not an
+    approximation of the vectorized path."""
+    c = chunk_arr[None, :]
+    dtype = c.dtype
+    for g_id, q1, q2, param in ops:
+        g00, g01, g10, g11, u00, u01, u10, u11 = _gate_matrix_elements(g_id, param, dtype)
+        if g_id >= 20:
+            c = _case_2q_local_local(c, u00, u01, u10, u11, q1, q2, m, k)
+        else:
+            c = _case_1q_local(c, g00, g01, g10, g11, q1, m, k)
+    return c[0]
+
+
+def _run_conditional_phase_on_chunk(chunk_arr, ops, chunk_index: int, m: int, k: int):
+    """A run of ConditionalPhase gates, applied to one chunk whose real
+    absolute index is `chunk_index` (needed for the ctrl-bit decision;
+    see _case_2q_ctrl_chunk_tgt_local's own docstring). Same load-once/
+    apply-all/save-once shape as _run_local_phase_on_chunk."""
+    dtype = chunk_arr.dtype
+    c = chunk_arr[None, :]
+    idxc = jnp.array([chunk_index], dtype=jnp.int32)
+    for g_id, q1, q2, param in ops:
+        _, _, _, _, u00, u01, u10, u11 = _gate_matrix_elements(g_id, param, dtype)
+        c = _case_2q_ctrl_chunk_tgt_local(c, u00, u01, u10, u11, q1, q2, m, k, idxc)
+    return c[0]
+
+
+def _mix_pair(row_a, row_b, e00, e01, e10, e11):
+    """The 2x2 amplitude-mixing algebra shared by every chunk-select-
+    qubit case in kernels.py (_case_1q_chunk / _case_2q_ctrl_local_tgt_chunk
+    / _case_2q_both_chunk): row_a is the row on the "mask/is_c0 = True"
+    side, row_b its XOR-stride partner. Mirrors those cases' own
+    new0/new1 formula exactly -- new0 = e00*_c + e01*amp_pair evaluated
+    at the True-side row (_c=row_a, amp_pair=row_b); new1 = e10*amp_pair
+    + e11*_c evaluated at the False-side row, where _c=row_b and
+    amp_pair=row_a, i.e. new_b = e10*row_a + e11*row_b."""
+    new_a = e00 * row_a + e01 * row_b
+    new_b = e10 * row_a + e11 * row_b
+    return new_a, new_b
+
+
+def _run_mix_phase_on_pair(row_a, row_b, op, index_a: int, m: int, k: int):
+    """row_a is the chunk whose bit at the gate's own mixing qubit is 0,
+    row_b its XOR-stride partner (bit 1) -- the caller picks which
+    loaded chunk is "a" vs "b" by checking that bit on the real absolute
+    indices before calling this. `index_a` is row_a's real absolute
+    chunk index (needed only for case 6's ctrl-chunk-select decision)."""
+    g_id, q1, q2, param = op
+    dtype = row_a.dtype
+    g00, g01, g10, g11, u00, u01, u10, u11 = _gate_matrix_elements(g_id, param, dtype)
+    if g_id < 20:
+        # case 2: 1-qubit, chunk-select q1 -- unconditional mix.
+        return _mix_pair(row_a, row_b, g00, g01, g10, g11)
+    if q1 >= m:
+        # case 5: ctrl LOCAL (elementwise mask within the row), tgt
+        # chunk-select q2 (the mixing qubit).
+        ctrl_phys = (k - 1) - (q1 - m)
+        idxl = jnp.arange(1 << k, dtype=jnp.int32)
+        ctrl_bit = (idxl & (jnp.int32(1) << ctrl_phys)) != 0
+        new_a, new_b = _mix_pair(row_a, row_b, u00, u01, u10, u11)
+        return jnp.where(ctrl_bit, new_a, row_a), jnp.where(ctrl_bit, new_b, row_b)
+    # case 6: ctrl AND tgt both chunk-select; mixing qubit is q2, ctrl
+    # decided once from index_a's own bit at q1 (identical for both rows
+    # of the pair -- XOR-ing the tgt/q2 stride never touches the q1 bit).
+    ctrl_stride = 1 << (m - 1 - q1)
+    ctrl_set = (index_a & ctrl_stride) != 0
+    new_a, new_b = _mix_pair(row_a, row_b, u00, u01, u10, u11)
+    return (new_a, new_b) if ctrl_set else (row_a, row_b)
+
+
+def run_disk_overflow_circuit(chunk_paths, compiled_ops, m: int, k: int):
+    """Runs the compiled circuit against num_chunks chunks stored as
+    plain .npy files at `chunk_paths` (index = real chunk index),
+    mutating those files in place. Never materializes more than one
+    (LocalPhase/ConditionalPhase) or two (MixPhase) chunks as jax.Array
+    at once, regardless of num_chunks -- this is the whole point."""
+    num_chunks = len(chunk_paths)
+    phases = partition_ops_into_phases(compiled_ops, m)
+
+    for phase in phases:
+        if isinstance(phase, LocalPhase):
+            for i, path in enumerate(chunk_paths):
+                arr = jnp.asarray(np.load(path))
+                arr = _run_local_phase_on_chunk(arr, phase.ops, m, k)
+                np.save(path, np.asarray(arr))
+
+        elif isinstance(phase, ConditionalPhase):
+            for i, path in enumerate(chunk_paths):
+                arr = jnp.asarray(np.load(path))
+                arr = _run_conditional_phase_on_chunk(arr, phase.ops, i, m, k)
+                np.save(path, np.asarray(arr))
+
+        elif isinstance(phase, MixPhase):
+            stride_bits = 1 << (m - 1 - phase.stride)
+            done = set()
+            for i in range(num_chunks):
+                if i in done:
+                    continue
+                j = i ^ stride_bits
+                done.add(i)
+                done.add(j)
+                row_a = jnp.asarray(np.load(chunk_paths[i]))
+                row_b = jnp.asarray(np.load(chunk_paths[j]))
+                new_a, new_b = _run_mix_phase_on_pair(row_a, row_b, phase.op, i, m, k)
+                np.save(chunk_paths[i], np.asarray(new_a))
+                np.save(chunk_paths[j], np.asarray(new_b))