Chunk (large-scale, anti-OOM)

chunk

MemoryPressureError

Bases: RuntimeError

Raised when available system RAM drops below the configured safety threshold. Catches the condition before the allocator attempts and crashes with jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED.

SafeMemoryGuard

SafeMemoryGuard(
    threshold_pct: float = 0.15,
    gc_before_check: bool = True,
)

Monitors system RAM before every high-memory operation and blocks execution if free RAM falls below threshold_pct of total physical memory.

Source code in dense_evolution/chunk.py

def __init__(self, threshold_pct: float = 0.15, gc_before_check: bool = True):
    if not 0.0 < threshold_pct < 1.0:
        raise ValueError(f"threshold_pct must be in (0, 1), got {threshold_pct}")
    self.threshold_pct   = threshold_pct
    self.gc_before_check = gc_before_check
    self._total_mb       = psutil.virtual_memory().total / (1024 * 1024)

check_allocation

check_allocation(
    required_mb: float, context: str = ""
) -> None

Like check(), but sized: verifies that required_mb can be allocated while still leaving threshold_pct free afterwards — check() alone only looks at RAM free right now, independent of what's about to be allocated (needed for Chunk's multi-chunk path, where several chunk-sized simulators are held in RAM at once).

Source code in dense_evolution/chunk.py

def check_allocation(self, required_mb: float, context: str = "") -> None:
    """Like check(), but sized: verifies that *required_mb* can be
    allocated while still leaving threshold_pct free afterwards —
    check() alone only looks at RAM free *right now*, independent of
    what's about to be allocated (needed for Chunk's multi-chunk path,
    where several chunk-sized simulators are held in RAM at once)."""
    if self.gc_before_check:
        gc.collect()

    s   = self.status()
    tag = f"[{context}] " if context else ""
    available_after_mb = s["available_mb"] - required_mb
    free_frac_after = available_after_mb / self._total_mb if self._total_mb > 0 else 0.0

    if available_after_mb < 0 or free_frac_after < self.threshold_pct:
        raise MemoryPressureError(
            f"\n{'─'*60}\n"
            f"  {tag}MEMORIA INSUFFICIENTE per l'allocazione richiesta\n"
            f"  Richiesti    : {required_mb:.0f} MB\n"
            f"  Disponibile  : {s['available_mb']:.0f} MB ({s['free_pct']:.1f}% libera)\n"
            f"  Dopo alloc.  : {available_after_mb:.0f} MB ({free_frac_after * 100:.1f}% libera)\n"
            f"  Soglia       : {self.threshold_pct * 100:.0f}% libera dopo l'allocazione\n"
            f"  Azione       : ridurre n_qubits o liberare RAM. Il chunking su\n"
            f"                 disco per overflow oltre la RAM disponibile non\n"
            f"                 e' implementato — vedi CHANGELOG.\n"
            f"{'─'*60}"
        )

MemoryChunker

MemoryChunker(n_qubits: int)

Geometry calculator for chunked simulation.

Attributes

n_qubits int — requested logical qubit count dtype — numpy/jax dtype for the statevector chunk_size_bits int — safe qubit-width that fits in RAM num_chunks int — number of statevector chunks required chunk_dim int — 2 ** chunk_size_bits

Source code in dense_evolution/chunk.py

def __init__(self, n_qubits: int):
    self.n_qubits        = n_qubits
    self.dtype           = _dtype_for_qubits(n_qubits)
    self.chunk_size_bits = get_dynamic_chunk(self.dtype)

    if self.n_qubits <= self.chunk_size_bits:
        self.num_chunks = 1
        self.chunk_dim  = 2 ** self.n_qubits
    else:
        self.num_chunks = 2 ** (self.n_qubits - self.chunk_size_bits)
        self.chunk_dim  = 2 ** self.chunk_size_bits

geometry

geometry() -> Tuple[int, int, int]

(num_chunks, chunk_dim, chunk_size_bits)

Source code in dense_evolution/chunk.py

def geometry(self) -> Tuple[int, int, int]:
    """(num_chunks, chunk_dim, chunk_size_bits)"""
    return self.num_chunks, self.chunk_dim, self.chunk_size_bits

memory_mb

memory_mb() -> float

Estimated RAM per chunk in MB.

Source code in dense_evolution/chunk.py

def memory_mb(self) -> float:
    """Estimated RAM per chunk in MB."""
    bpe = np.dtype(self.dtype).itemsize
    return (self.chunk_dim * bpe) / (1024 * 1024)

CircuitChunker

CircuitChunker(
    simulator_instance: Optional[DenseSVSimulator] = None,
    memory_threshold: float = 0.15,
)

Transpile a circuit once, then execute it in fixed-size gate-slices so XLA sees the same trace shape on every compilation.

A SafeMemoryGuard is checked before every slice — if RAM drops below 15% the current slice is aborted with MemoryPressureError before JAX attempts the allocation.

Parameters

simulator_instance : DenseSVSimulator Physical simulator (sized to safe_qubits, not logical n_qubits). memory_threshold : float Passed to SafeMemoryGuard. Default 0.15 (15%).

Source code in dense_evolution/chunk.py

def __init__(
    self,
    simulator_instance: Optional[DenseSVSimulator] = None,
    memory_threshold: float = 0.15,
):
    self.sim   = simulator_instance
    self._guard = SafeMemoryGuard(threshold_pct=memory_threshold)

split_circuit

split_circuit(circuit: List, chunk_size: int = 500) -> None

Execute circuit in slices of chunk_size gates.

Raises

RuntimeError if no simulator instance is attached. MemoryPressureError if RAM drops below threshold before a slice.

Source code in dense_evolution/chunk.py

def split_circuit(self, circuit: List, chunk_size: int = 500) -> None:
    """
    Execute *circuit* in slices of *chunk_size* gates.

    Raises
    ------
    RuntimeError        if no simulator instance is attached.
    MemoryPressureError if RAM drops below threshold before a slice.
    """
    if self.sim is None:
        raise RuntimeError(
            "CircuitChunker: no simulator instance attached. "
            "Pass simulator_instance= at construction or assign .sim."
        )

    target: List = QuantumTranspiler.transpile(circuit)
    n_slices     = (len(target) + chunk_size - 1) // chunk_size

    for idx, i in enumerate(range(0, len(target), chunk_size)):
        # ── Anti-OOM check before every slice ───────────────────────────
        self._guard.check(f"slice {idx + 1}/{n_slices}")
        self.sim.run_circuit_jit_beast_mode(target[i : i + chunk_size])

Chunk

Chunk(
    n_qubits: int,
    chunk_size_gates: int = 500,
    memory_threshold: float = 0.15,
    use_gpu: bool = False,
    use_float32: bool = False,
)

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) — no disk/memmap paging, so this only covers a moderate overflow beyond the safe budget (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.

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_gpu : forwarded to DenseSVSimulator use_float32 : forwarded to DenseSVSimulator

Source code in dense_evolution/chunk.py

def __init__(
    self,
    n_qubits: int,
    chunk_size_gates:  int   = 500,
    memory_threshold:  float = 0.15,
    use_gpu:           bool  = False,
    use_float32:       bool  = False,
):
    # 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)

    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_gpu=use_gpu,
            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
        self._guard.check_allocation(
            required_mb,
            f"Chunk.__init__ — allocating {num_chunks} chunks of "
            f"{self._mem_chunker.chunk_size_bits} qubits each",
        )

        # 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_gpu=use_gpu, 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

sv property writable

sv

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).

memory_mb

memory_mb() -> float

RAM used by the physical statevector(s) in MB.

Source code in dense_evolution/chunk.py

def memory_mb(self) -> float:
    """RAM used by the physical statevector(s) in MB."""
    if self._chunk_sims is None:
        return self._inner_sim.memory_mb()
    return sum(sim.memory_mb() for sim in self._chunk_sims)

get_probabilities

get_probabilities()

|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).

Source code in dense_evolution/chunk.py

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)."""
    if self._chunk_sims is None:
        return self._inner_sim.get_probabilities()
    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

get_statevector

get_statevector()

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).

Source code in dense_evolution/chunk.py

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)."""
    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])

dispatch_distributed

dispatch_distributed(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.

Source code in dense_evolution/chunk.py

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."""
    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])

run_chunk_distributed

run_chunk_distributed(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).

Source code in dense_evolution/chunk.py

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)