import gc import psutil import numpy as np from typing import List, Optional, Tuple # JAX is now a mandatory dependency (see registry.py) -- kept as an # import-time flag for the dual-path code below, which stays as-is. import jax import jax.numpy as jnp HAS_JAX = True # ── Flexible import with stub fallback ────────────────────────────────────── try: from simulator import DenseSVSimulator from compiler import QuantumTranspiler from gates import GATE_IDS except ModuleNotFoundError: try: from dense_evolution.simulator import DenseSVSimulator from dense_evolution.compiler import QuantumTranspiler from dense_evolution.gates import GATE_IDS except ModuleNotFoundError: class DenseSVSimulator: # type: ignore[no-redef] def __init__(self, n_qubits, **kwargs): self.n = n_qubits self.dim = 2 ** n_qubits self.dtype = np.complex128 self.sv = np.zeros(self.dim, dtype=self.dtype) self.sv[0] = 1.0 def run_circuit_jit_beast_mode(self, circuit_slice): pass def memory_mb(self) -> float: return (self.dim * np.dtype(self.dtype).itemsize) / 1_000_000 class QuantumTranspiler: # type: ignore[no-redef] @staticmethod def transpile(circuit): return circuit GATE_IDS: dict = {} # ───────────────────────────────────────────────────────────────────────────── # Helpers # ───────────────────────────────────────────────────────────────────────────── def get_dynamic_chunk(dtype_target) -> int: vm = psutil.virtual_memory() safe_ram = vm.available * 0.85 if HAS_JAX and dtype_target is jnp.complex128: bpe = 16 elif dtype_target is np.complex128: bpe = 16 else: bpe = 8 max_elements = safe_ram / bpe max_bits = int(np.floor(np.log2(max(max_elements, 2.0)))) return max(16, min(max_bits, 27)) def _dtype_for_qubits(n_qubits: int): xp = jnp if HAS_JAX else np return xp.complex64 if n_qubits > 26 else xp.complex128 # ───────────────────────────────────────────────────────────────────────────── # SafeMemoryGuard — Anti-OOM block # ───────────────────────────────────────────────────────────────────────────── class MemoryPressureError(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. """ pass class SafeMemoryGuard: """ Monitors system RAM before every high-memory operation and blocks execution if free RAM falls below ``threshold_pct`` of total physical memory. """ _WARN_MULTIPLIER = 2.0 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) def status(self) -> dict: vm = psutil.virtual_memory() available_mb = vm.available / (1024 * 1024) free_pct = vm.available / vm.total return { "total_mb" : self._total_mb, "available_mb": available_mb, "used_pct" : vm.percent, "free_pct" : free_pct * 100.0, "safe" : free_pct >= self.threshold_pct, } def check(self, context: str = "") -> None: if self.gc_before_check: gc.collect() s = self.status() tag = f"[{context}] " if context else "" free_frac = s["free_pct"] / 100.0 if not s["safe"]: raise MemoryPressureError( f"\n{'─'*60}\n" f" {tag}MEMORIA CRITICA — simulazione bloccata\n" f" Disponibile : {s['available_mb']:.0f} MB " f"({s['free_pct']:.1f}% libera)\n" f" Soglia : {self.threshold_pct * 100:.0f}% " f"({self._total_mb * self.threshold_pct:.0f} MB)\n" f" Azione : liberare RAM o ridurre n_qubits / chunk_size.\n" f"{'─'*60}" ) warn_threshold = self.threshold_pct * self._WARN_MULTIPLIER if free_frac < warn_threshold: print( f" [WARN] {tag}RAM bassa: {s['available_mb']:.0f} MB liberi " f"({s['free_pct']:.1f}%) — soglia critica al " f"{self.threshold_pct * 100:.0f}%." ) 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}" ) def __repr__(self) -> str: s = self.status() return ( f"SafeMemoryGuard(" f"threshold={self.threshold_pct*100:.0f}%, " f"available={s['available_mb']:.0f} MB / {s['free_pct']:.1f}% free, " f"safe={s['safe']})" ) # ───────────────────────────────────────────────────────────────────────────── # MemoryChunker (chunk1) # ───────────────────────────────────────────────────────────────────────────── class MemoryChunker: """ 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 """ 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 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 def memory_mb(self) -> float: """Estimated RAM per chunk in MB.""" bpe = np.dtype(self.dtype).itemsize return (self.chunk_dim * bpe) / (1024 * 1024) def __repr__(self) -> str: return ( f"MemoryChunker(n_qubits={self.n_qubits}, " f"num_chunks={self.num_chunks}, " f"chunk_dim={self.chunk_dim}, " f"chunk_size_bits={self.chunk_size_bits}, " f"dtype={self.dtype}, " f"mem_per_chunk={self.memory_mb():.2f} MB)" ) # ───────────────────────────────────────────────────────────────────────────── # Multi-chunk JIT kernel (num_chunks > 1) # ───────────────────────────────────────────────────────────────────────────── if HAS_JAX: def _build_multi_chunk_step(num_chunks: int, m: int, k: int): """ Build a jax.lax.scan-compatible step function operating on the STACKED multi-chunk representation — a (num_chunks, chunk_dim) array — instead of a flat (2**n_qubits,) statevector. Modeled directly on compiler.py's _apply_gate_fast_step (same [g_id, q1, q2, param] encoding via GATE_IDS, same lax.switch for the 1-qubit matrix, same per-gate controlled-U dispatch for the 5 two-qubit gate types) — but a gate touching a "chunk-select" qubit (index < m, the top m logical qubits that select WHICH chunk) mixes whole (chunk_dim,)-shaped ROWS instead of individual amplitudes, while a gate touching a "local" qubit (index >= m) mixes individual elements WITHIN each row in parallel. This never materializes a (2**n_qubits,) array — Chunk's whole reason to exist — because the stacked shape (num_chunks, chunk_dim) holds exactly the same total elements as the num_chunks separate per-chunk arrays it replaces. The 6 gate/qubit-location combinations below are a direct translation of the pre-JIT _apply_gate_multi's Python-loop formulas (removed once this replaced it) — verified case-by-case against DenseSVSimulator on non-chunked reference circuits before being wired in, not derived fresh. All 6 are traced unconditionally every step and selected via jnp.where on the runtime q1/q2 vs static m comparison, same "trace every branch" pattern _apply_gate_fast_step already uses for is_1q/is_2q and the 5 two-qubit sub-gates. """ chunk_dim = 1 << k def step(chunks, operation): g_id = operation[0].astype(jnp.int32) q1 = operation[1].astype(jnp.int32) q2 = operation[2].astype(jnp.int32) param = operation[3] dtype = chunks.dtype # never hardcode complex128 — see the # use_float32 bug this exact mistake # caused once already in beast-mode. inv2 = jnp.asarray(1.0 / jnp.sqrt(2.0), dtype=dtype) half_p = param * jnp.float64(0.5) cos_p = jnp.cos(half_p).astype(dtype) sin_p = jnp.sin(half_p).astype(dtype) exp_pos = jnp.exp(1j * param).astype(dtype) exp_ph4 = jnp.exp(1j * jnp.pi / 4.0).astype(dtype) exp_mh4 = jnp.exp(-1j * jnp.pi / 4.0).astype(dtype) exp_pos_half = jnp.exp(1j * half_p).astype(dtype) exp_neg_half = jnp.exp(-1j * half_p).astype(dtype) # 1-qubit gate matrix — identical table to _apply_gate_fast_step safe_gid = jnp.clip(g_id, 0, 13) g_1q = jax.lax.switch( safe_gid, [ lambda _: jnp.eye(2, dtype=dtype), lambda _: jnp.array([[inv2, inv2], [inv2, -inv2]], dtype=dtype), lambda _: jnp.array([[0.0 + 0j, 1.0 + 0j], [1.0 + 0j, 0.0 + 0j]], dtype=dtype), lambda _: jnp.array([[0.0 + 0j, -1j], [1j, 0.0 + 0j]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, -1.0 + 0j]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, 1j]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, -1j]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, exp_ph4]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, exp_mh4]], dtype=dtype), lambda _: jnp.array([[cos_p, -1j * sin_p], [-1j * sin_p, cos_p]], dtype=dtype), lambda _: jnp.array([[cos_p, -sin_p], [sin_p, cos_p]], dtype=dtype), lambda _: jnp.array([[jnp.exp(-1j * half_p), 0.0 + 0j], [0.0 + 0j, jnp.exp(1j * half_p)]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, exp_pos]], dtype=dtype), lambda _: jnp.array([[0.5 + 0.5j, 0.5 - 0.5j], [0.5 - 0.5j, 0.5 + 0.5j]], dtype=dtype), ], operand=None, ) # Controlled-U submatrix for the 5 two-qubit gate types (mat[2:,2:] # of each gate's full 4x4 form — same values _apply_gate_multi's # `U = mat[2:, 2:]` extracted from GATES/PARAMETRIC_GATES). # 20=CX->X, 21=CZ->Z, 22=CP->P(theta), 24=CY->Y, 25=CRZ->RZ(theta). two_q_idx = jnp.where(g_id == 20, 0, jnp.where(g_id == 21, 1, jnp.where(g_id == 22, 2, jnp.where(g_id == 24, 3, 4)))) U = jax.lax.switch( two_q_idx, [ lambda _: jnp.array([[0.0 + 0j, 1.0 + 0j], [1.0 + 0j, 0.0 + 0j]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, -1.0 + 0j]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, exp_pos]], dtype=dtype), lambda _: jnp.array([[0.0 + 0j, -1j], [1j, 0.0 + 0j]], dtype=dtype), lambda _: jnp.array([[exp_neg_half, 0.0 + 0j], [0.0 + 0j, exp_pos_half]], dtype=dtype), ], operand=None, ) g00, g01, g10, g11 = g_1q[0, 0], g_1q[0, 1], g_1q[1, 0], g_1q[1, 1] u00, u01, u10, u11 = U[0, 0], U[0, 1], U[1, 0], U[1, 1] # ── case 1: 1-qubit gate, LOCAL qubit (q1 >= m) ────────────── # Same do_1q amplitude-pair math as beast-mode, applied to every # chunk row in parallel (gather along axis 1, broadcasts over # axis 0 for free). def case_1q_local(_c): local_phys = (k - 1) - (q1 - m) stride = jnp.int32(1) << local_phys idx = jnp.arange(chunk_dim, dtype=jnp.int32) idx_pair = idx ^ stride mask0 = (idx & stride) == 0 amp_pair = _c[:, idx_pair] new0 = g00 * _c + g01 * amp_pair new1 = g10 * amp_pair + g11 * _c return jnp.where(mask0[None, :], new0, new1) # ── case 2: 1-qubit gate, CHUNK-SELECT qubit (q1 < m) ──────── # Same math, one level up: each "amplitude" is a whole # (chunk_dim,)-shaped row, mixed pairwise across axis 0. def case_1q_chunk(_c): stride = jnp.int32(1) << (m - 1 - q1) idxc = jnp.arange(num_chunks, dtype=jnp.int32) idxc_pair = idxc ^ stride mask0 = (idxc & stride) == 0 amp_pair = _c[idxc_pair] new0 = g00 * _c + g01 * amp_pair new1 = g10 * amp_pair + g11 * _c return jnp.where(mask0[:, None], new0, new1) # ── case 3: 2-qubit, ctrl AND tgt both LOCAL ───────────────── def case_2q_local_local(_c): ctrl_phys = (k - 1) - (q1 - m) tgt_phys = (k - 1) - (q2 - m) idx = jnp.arange(chunk_dim, dtype=jnp.int32) ctrl_bit = (idx & (jnp.int32(1) << ctrl_phys)) != 0 tgt_bit = (idx & (jnp.int32(1) << tgt_phys)) != 0 partner = idx ^ (jnp.int32(1) << tgt_phys) amp_partner = _c[:, partner] new0 = u00 * _c + u01 * amp_partner new1 = u10 * amp_partner + u11 * _c after = jnp.where(tgt_bit[None, :], new1, new0) return jnp.where(ctrl_bit[None, :], after, _c) # ── case 4: ctrl CHUNK-SELECT, tgt LOCAL ───────────────────── # Whole chunks where the chunk-index's ctrl bit is set get U # applied as a local 1-qubit gate; the rest are untouched. def case_2q_ctrl_chunk_tgt_local(_c): ctrl_stride = jnp.int32(1) << (m - 1 - q1) idxc = jnp.arange(num_chunks, dtype=jnp.int32) ctrl_set = (idxc & ctrl_stride) != 0 tgt_phys = (k - 1) - (q2 - m) idxl = jnp.arange(chunk_dim, dtype=jnp.int32) tgt_bit = (idxl & (jnp.int32(1) << tgt_phys)) != 0 partner = idxl ^ (jnp.int32(1) << tgt_phys) amp_partner = _c[:, partner] new0 = u00 * _c + u01 * amp_partner new1 = u10 * amp_partner + u11 * _c after = jnp.where(tgt_bit[None, :], new1, new0) return jnp.where(ctrl_set[:, None], after, _c) # ── case 5: ctrl LOCAL, tgt CHUNK-SELECT ───────────────────── # Pairs of chunks get mixed, but ONLY where the local ctrl bit # (same position in every chunk) is set — an elementwise mask. def case_2q_ctrl_local_tgt_chunk(_c): ctrl_phys = (k - 1) - (q1 - m) idxl = jnp.arange(chunk_dim, dtype=jnp.int32) ctrl_bit = (idxl & (jnp.int32(1) << ctrl_phys)) != 0 tgt_stride = jnp.int32(1) << (m - 1 - q2) idxc = jnp.arange(num_chunks, dtype=jnp.int32) idxc_pair = idxc ^ tgt_stride is_c0 = (idxc & tgt_stride) == 0 amp_pair = _c[idxc_pair] new_c0 = u00 * _c + u01 * amp_pair new_c1 = u10 * amp_pair + u11 * _c after = jnp.where(is_c0[:, None], new_c0, new_c1) return jnp.where(ctrl_bit[None, :], after, _c) # ── case 6: ctrl AND tgt both CHUNK-SELECT ─────────────────── def case_2q_both_chunk(_c): ctrl_stride = jnp.int32(1) << (m - 1 - q1) tgt_stride = jnp.int32(1) << (m - 1 - q2) idxc = jnp.arange(num_chunks, dtype=jnp.int32) ctrl_set = (idxc & ctrl_stride) != 0 idxc_pair = idxc ^ tgt_stride is_c0 = (idxc & tgt_stride) == 0 amp_pair = _c[idxc_pair] new_c0 = u00 * _c + u01 * amp_pair new_c1 = u10 * amp_pair + u11 * _c after = jnp.where(is_c0[:, None], new_c0, new_c1) return jnp.where(ctrl_set[:, None], after, _c) is_2q = g_id >= 20 q1_chunk = q1 < m q2_chunk = q2 < m result_2q = jnp.where( q1_chunk & q2_chunk, case_2q_both_chunk(chunks), jnp.where(q1_chunk & (~q2_chunk), case_2q_ctrl_chunk_tgt_local(chunks), jnp.where((~q1_chunk) & q2_chunk, case_2q_ctrl_local_tgt_chunk(chunks), case_2q_local_local(chunks)))) result_1q = jnp.where(q1_chunk, case_1q_chunk(chunks), case_1q_local(chunks)) new_chunks = jnp.where(is_2q, result_2q, result_1q) return new_chunks.astype(dtype), None return step # ───────────────────────────────────────────────────────────────────── # Distributed (multi-device) variant — one physical chunk per device # ───────────────────────────────────────────────────────────────────── def _build_distributed_chunk_step(num_chunks: int, m: int, k: int, axis_name: str): """Same 6-case formula set as _build_multi_chunk_step, but each device holds exactly ONE chunk row (chunk_dim,) instead of the whole (num_chunks, chunk_dim) stack living on one device/process — issue #1: distribute chunks across a device mesh (multi-GPU/ multi-host), not just multi-chunk within one process's RAM. The stacked-array formulation's "mix pairs of rows across axis 0" (cases 2/5/6, touching a chunk-select qubit) becomes real point-to-point network communication here: jax.lax.ppermute, keyed on the fixed XOR-stride pairing between chunk indices — the textbook pairwise-exchange communication pattern used by distributed statevector simulators (each device sends its local row to its stride-partner device and receives the partner's row back). Cases 1/3/4 need NO communication at all: case 1/3 are purely local (both qubits live inside this device's own chunk_dim index space), and case 4 (ctrl chunk-select, tgt local) is a decision every device can make on its OWN chunk index alone (whether ITS id has the ctrl bit set) — no data from any other device is needed to decide or to apply the local tgt gate. ppermute's `perm` argument is a communication topology and must be STATIC (known at trace time) — it cannot be built from q1/q2, which are traced values read from the scanned circuit array. Since q1/q2 only ever range over the m possible chunk-select qubit indices [0, m), every possible stride is enumerated as its own statically-built ppermute call, and jax.lax.switch (traced unconditionally, same "trace every branch" pattern used throughout this codebase) picks the right one at runtime — plus one extra identity branch for "no chunk-select qubit involved, no communication needed" (verified below to be exactly cases 1/3/4, never 2/5/6).""" chunk_dim = 1 << k def step(local_row, operation): g_id = operation[0].astype(jnp.int32) q1 = operation[1].astype(jnp.int32) q2 = operation[2].astype(jnp.int32) param = operation[3] dtype = local_row.dtype my_id = jax.lax.axis_index(axis_name).astype(jnp.int32) inv2 = jnp.asarray(1.0 / jnp.sqrt(2.0), dtype=dtype) half_p = param * jnp.float64(0.5) cos_p = jnp.cos(half_p).astype(dtype) sin_p = jnp.sin(half_p).astype(dtype) exp_pos = jnp.exp(1j * param).astype(dtype) exp_ph4 = jnp.exp(1j * jnp.pi / 4.0).astype(dtype) exp_mh4 = jnp.exp(-1j * jnp.pi / 4.0).astype(dtype) exp_pos_half = jnp.exp(1j * half_p).astype(dtype) exp_neg_half = jnp.exp(-1j * half_p).astype(dtype) safe_gid = jnp.clip(g_id, 0, 13) g_1q = jax.lax.switch( safe_gid, [ lambda _: jnp.eye(2, dtype=dtype), lambda _: jnp.array([[inv2, inv2], [inv2, -inv2]], dtype=dtype), lambda _: jnp.array([[0.0 + 0j, 1.0 + 0j], [1.0 + 0j, 0.0 + 0j]], dtype=dtype), lambda _: jnp.array([[0.0 + 0j, -1j], [1j, 0.0 + 0j]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, -1.0 + 0j]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, 1j]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, -1j]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, exp_ph4]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, exp_mh4]], dtype=dtype), lambda _: jnp.array([[cos_p, -1j * sin_p], [-1j * sin_p, cos_p]], dtype=dtype), lambda _: jnp.array([[cos_p, -sin_p], [sin_p, cos_p]], dtype=dtype), lambda _: jnp.array([[jnp.exp(-1j * half_p), 0.0 + 0j], [0.0 + 0j, jnp.exp(1j * half_p)]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, exp_pos]], dtype=dtype), lambda _: jnp.array([[0.5 + 0.5j, 0.5 - 0.5j], [0.5 - 0.5j, 0.5 + 0.5j]], dtype=dtype), ], operand=None, ) two_q_idx = jnp.where(g_id == 20, 0, jnp.where(g_id == 21, 1, jnp.where(g_id == 22, 2, jnp.where(g_id == 24, 3, 4)))) U = jax.lax.switch( two_q_idx, [ lambda _: jnp.array([[0.0 + 0j, 1.0 + 0j], [1.0 + 0j, 0.0 + 0j]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, -1.0 + 0j]], dtype=dtype), lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, exp_pos]], dtype=dtype), lambda _: jnp.array([[0.0 + 0j, -1j], [1j, 0.0 + 0j]], dtype=dtype), lambda _: jnp.array([[exp_neg_half, 0.0 + 0j], [0.0 + 0j, exp_pos_half]], dtype=dtype), ], operand=None, ) g00, g01, g10, g11 = g_1q[0, 0], g_1q[0, 1], g_1q[1, 0], g_1q[1, 1] u00, u01, u10, u11 = U[0, 0], U[0, 1], U[1, 0], U[1, 1] is_2q = g_id >= 20 q1_chunk = q1 < m q2_chunk = q2 < m # ── single point-to-point exchange, done once per step ────── # comm_qubit selects which stride to ppermute on: q2 (tgt) for # any 2-qubit gate with a chunk-select target (cases 5 and 6), # q1 for a 1-qubit gate on a chunk-select qubit (case 2), # sentinel `m` (-> identity, no network traffic) for cases # 1/3/4, which never need another device's data. needs_comm_2q = is_2q & q2_chunk needs_comm_1q = (~is_2q) & q1_chunk comm_qubit = jnp.where(needs_comm_2q, q2, jnp.where(needs_comm_1q, q1, m)) safe_comm_idx = jnp.clip(comm_qubit, 0, m) # NOTE: `_perm` MUST be bound as a default-argument value here # (evaluated eagerly, once, at lambda-creation time inside this # loop) rather than referencing `q`/`m` freely inside the # lambda body -- a free reference would be looked up at CALL # time via Python's normal late-binding closure semantics, by # which point the loop variable `q` has already reached its # final value (m-1) for every branch, silently making every # ppermute use the LAST qubit's stride regardless of which # branch was actually selected. Caught by exactly that # symptom: only the branch for q == m-1 gave correct results. ppermute_branches = [ (lambda _row, _perm=[(i, i ^ (1 << (m - 1 - q))) for i in range(num_chunks)]: jax.lax.ppermute(_row, axis_name, perm=_perm)) for q in range(m) ] + [lambda _row: _row] # identity: no chunk-select qubit involved paired_row = jax.lax.switch(safe_comm_idx, ppermute_branches, local_row) # ── case 1: 1-qubit, LOCAL qubit (q1 >= m) — no comm ───────── def case_1q_local(_row): local_phys = (k - 1) - (q1 - m) stride = jnp.int32(1) << local_phys idx = jnp.arange(chunk_dim, dtype=jnp.int32) idx_pair = idx ^ stride mask0 = (idx & stride) == 0 amp_pair = _row[idx_pair] new0 = g00 * _row + g01 * amp_pair new1 = g10 * amp_pair + g11 * _row return jnp.where(mask0, new0, new1) # ── case 2: 1-qubit, CHUNK-SELECT qubit (q1 < m) ───────────── # paired_row already fetched via ppermute above (comm_qubit=q1). def case_1q_chunk(_row): stride = jnp.int32(1) << (m - 1 - q1) mask0 = (my_id & stride) == 0 new0 = g00 * _row + g01 * paired_row new1 = g10 * paired_row + g11 * _row return jnp.where(mask0, new0, new1) # ── case 3: 2-qubit, ctrl AND tgt both LOCAL — no comm ─────── def case_2q_local_local(_row): ctrl_phys = (k - 1) - (q1 - m) tgt_phys = (k - 1) - (q2 - m) idx = jnp.arange(chunk_dim, dtype=jnp.int32) ctrl_bit = (idx & (jnp.int32(1) << ctrl_phys)) != 0 partner = idx ^ (jnp.int32(1) << tgt_phys) amp_partner = _row[partner] new0 = u00 * _row + u01 * amp_partner new1 = u10 * amp_partner + u11 * _row tgt_bit = (idx & (jnp.int32(1) << tgt_phys)) != 0 after = jnp.where(tgt_bit, new1, new0) return jnp.where(ctrl_bit, after, _row) # ── case 4: ctrl CHUNK-SELECT, tgt LOCAL — no comm needed: ─ # every device decides purely from its OWN chunk index (my_id) # whether the control bit is set, and if so applies U locally. def case_2q_ctrl_chunk_tgt_local(_row): ctrl_stride = jnp.int32(1) << (m - 1 - q1) ctrl_set = (my_id & ctrl_stride) != 0 tgt_phys = (k - 1) - (q2 - m) idx = jnp.arange(chunk_dim, dtype=jnp.int32) tgt_bit = (idx & (jnp.int32(1) << tgt_phys)) != 0 partner = idx ^ (jnp.int32(1) << tgt_phys) amp_partner = _row[partner] new0 = u00 * _row + u01 * amp_partner new1 = u10 * amp_partner + u11 * _row after = jnp.where(tgt_bit, new1, new0) return jnp.where(ctrl_set, after, _row) # ── case 5: ctrl LOCAL, tgt CHUNK-SELECT ────────────────────── # paired_row already fetched via ppermute above (comm_qubit=q2, # keyed on tgt's stride). is_c0 is a per-device scalar decision # (which side of the tgt pairing this device is on); ctrl_bit # is a per-element mask within the row (elementwise, local). def case_2q_ctrl_local_tgt_chunk(_row): ctrl_phys = (k - 1) - (q1 - m) idx = jnp.arange(chunk_dim, dtype=jnp.int32) ctrl_bit = (idx & (jnp.int32(1) << ctrl_phys)) != 0 tgt_stride = jnp.int32(1) << (m - 1 - q2) is_c0 = (my_id & tgt_stride) == 0 new_c0 = u00 * _row + u01 * paired_row new_c1 = u10 * paired_row + u11 * _row after = jnp.where(is_c0, new_c0, new_c1) return jnp.where(ctrl_bit, after, _row) # ── case 6: ctrl AND tgt both CHUNK-SELECT ──────────────────── # paired_row via ppermute keyed on tgt's stride (comm_qubit=q2); # ctrl_set and is_c0 are both per-device scalars (this device's # own chunk index bits) — no per-element masking needed at all. def case_2q_both_chunk(_row): ctrl_stride = jnp.int32(1) << (m - 1 - q1) ctrl_set = (my_id & ctrl_stride) != 0 tgt_stride = jnp.int32(1) << (m - 1 - q2) is_c0 = (my_id & tgt_stride) == 0 new_c0 = u00 * _row + u01 * paired_row new_c1 = u10 * paired_row + u11 * _row after = jnp.where(is_c0, new_c0, new_c1) return jnp.where(ctrl_set, after, _row) result_2q = jnp.where( q1_chunk & q2_chunk, case_2q_both_chunk(local_row), jnp.where(q1_chunk & (~q2_chunk), case_2q_ctrl_chunk_tgt_local(local_row), jnp.where((~q1_chunk) & q2_chunk, case_2q_ctrl_local_tgt_chunk(local_row), case_2q_local_local(local_row)))) result_1q = jnp.where(q1_chunk, case_1q_chunk(local_row), case_1q_local(local_row)) new_row = jnp.where(is_2q, result_2q, result_1q) return new_row.astype(dtype), None return step def _build_distributed_chunk_runner(num_chunks: int, m: int, k: int): """shard_map-wrapped runner: one chunk row per physical JAX device. Requires jax.device_count() >= num_chunks (v1 scope: exactly one chunk per device, the literal reading of issue #1 -- "distribuire i chunk su più device"; a hybrid scheme with several chunks per device is a possible future refinement, not attempted here). compiled_ops (the small [g_id, q1, q2, param] sequence, identical on every device) is replicated, not sharded -- P(None, None). local_row is sharded along axis 0 of the (num_chunks, chunk_dim) logical array, one (chunk_dim,) row per device -- P(axis_name, None) on input/output, so each device's shard is that one row.""" from jax.sharding import Mesh, PartitionSpec as P axis_name = 'chunks' step = _build_distributed_chunk_step(num_chunks, m, k, axis_name) devices = np.array(jax.devices()[:num_chunks]) mesh = Mesh(devices, axis_names=(axis_name,)) def run_local(local_shard, compiled_ops): # shard_map keeps the sharded axis in the local shape (size # num_chunks/mesh_size along axis 0 -- 1 in this v1 one- # chunk-per-device scope), it doesn't squeeze it away: a # (num_chunks, chunk_dim) input shards to (1, chunk_dim) per # device here, not (chunk_dim,). `step` itself works on a # clean (chunk_dim,) row -- squeeze going in, restore going out. local_row = local_shard[0] final_row, _ = jax.lax.scan(step, local_row, compiled_ops) return final_row[None, :] sharded_run = jax.shard_map( run_local, mesh=mesh, in_specs=(P(axis_name, None), P(None, None)), out_specs=P(axis_name, None), check_vma=False, ) return jax.jit(sharded_run), mesh def _build_multi_chunk_runner(num_chunks: int, m: int, k: int): """jax.jit-compiled (chunks, compiled_ops) -> final_chunks, closed over the static per-Chunk-instance geometry (num_chunks, m, k don't change across calls on the same instance).""" step = _build_multi_chunk_step(num_chunks, m, k) @jax.jit def run(chunks, compiled_ops): final, _ = jax.lax.scan(step, chunks, compiled_ops) return final return run def _compile_multi_chunk_ops(circuit: List) -> "jnp.ndarray": """Structural + GATE_IDS compilation shared by the multi-chunk JIT path — same [g_id, q1, q2, param] row format as beast-mode's own compiled ops, built via GATE_IDS instead of the old _resolve_gate's GATES/PARAMETRIC_GATES lookup. This finally aligns multi-chunk's gate coverage with beast-mode's (both silently skip a gate name not in GATE_IDS — same known, tracked behavior, see issue #4 — instead of the old _resolve_gate's NotImplementedError for e.g. ecr/iswap).""" target = QuantumTranspiler.transpile(circuit) rows = [] for cmd in target: name = cmd[0].lower() if isinstance(cmd[0], str) else str(cmd[0]).lower() if name not in GATE_IDS: continue g_id = float(GATE_IDS[name]) args = cmd[1:] if name in ('cx', 'cz', 'cp', 'cphase', 'cy', 'crz'): q1, q2 = float(args[0]), float(args[1]) param = float(args[2]) if len(args) > 2 else 0.0 rows.append([g_id, q1, q2, param]) elif args: q1 = float(args[0]) param = float(args[1]) if len(args) > 1 else 0.0 rows.append([g_id, q1, 0.0, param]) if not rows: return jnp.empty((0, 4), dtype=jnp.float64) return jnp.array(rows, dtype=jnp.float64) # ───────────────────────────────────────────────────────────────────────────── # CircuitChunker # ───────────────────────────────────────────────────────────────────────────── class CircuitChunker: """ 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%). """ def __init__( self, simulator_instance: Optional[DenseSVSimulator] = None, memory_threshold: float = 0.15, ): self.sim = simulator_instance self._guard = SafeMemoryGuard(threshold_pct=memory_threshold) 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 (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) — 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 """ 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 # ── 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).""" 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, or split back into per-chunk slices (same ascending concatenation order as the getter) for num_chunks>1.""" if self._chunk_sims is None: self._inner_sim.sv = value return chunk_dim = self._mem_chunker.chunk_dim 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.""" 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).""" 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 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]) # ── 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.""" 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_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 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"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