prism-q 0.35.0

//! Full state-vector simulation backend.
//!
//! Stores the complete 2^n amplitude vector and applies gates via direct
//! index manipulation. This is the reference backend. The qubit cap is derived
//! from detected system memory at dispatch time; `PRISM_MAX_SV_QUBITS`
//! overrides it.
//!
//! # Memory layout
//!
//! State is a contiguous `Vec<Complex64>` of length 2^n, indexed by computational
//! basis state in standard binary order (qubit 0 = least significant bit).
//!
//! # Gate support
//!
//! The full gate set, including MCU, every fused or batched variant, a native
//! `Gate::QftBlock` kernel on the CPU path for blocks starting at qubit 0, and native
//! `Gate::PauliRot` and `Gate::Unitary` kernels on the CPU path. The device
//! path declines `Gate::Unitary` by name.
//!
//! # When to prefer this backend
//!
//! - Exact amplitudes for arbitrary circuits within the memory cap.
//! - Dense entanglement, where structured backends lose their advantage.
//! - Reference results when validating other backends.
//!
//! # When NOT to use this backend
//!
//! - Qubit counts past the memory cap (dispatch picks sparse or MPS instead).
//! - Clifford-only circuits, where the stabilizer tableau is exact at any size.
//! - Non-entangling or sparse-entanglement circuits, where product or factored
//!   states are exponentially cheaper.
//!
//! # Hot-path design
//!
//! - Single-qubit gates: iterate 2^(n-1) pairs with stride 2^target.
//! - CX/CZ/SWAP: specialized routines avoid materializing a 4×4 matrix.
//! - QFT twiddle tables allocate once outside the per-amplitude loop and are
//!   bounded by cache policy.
//!
//! # Threading and SIMD
//!
//! The pair-iteration loops are embarrassingly parallel. With the `parallel`
//! feature at `PARALLEL_THRESHOLD_QUBITS` (15) and above, each kernel splits
//! across Rayon with `par_chunks_mut` / `split_at_mut`; a few hot kernels use
//! raw pointers after proving disjoint access.
//!
//! The single-qubit kernel runs through the shared `simd` module: matrix
//! entries are precomputed as broadcast pairs `[re, re]` / `[im, im]`, and
//! runtime dispatch picks AVX2+FMA (256-bit), FMA (128-bit), SSE2, or scalar.
//! Two-qubit gate and measurement loops use the SIMD bulk helpers
//! (`negate_slice`, `swap_slices`, `norm_sqr_sum`, `zero_slice`).

pub(crate) mod kernels;
#[cfg(test)]
mod tests;

use num_complex::Complex64;
use rand::SeedableRng;
use rand_chacha::ChaCha8Rng;

#[cfg(feature = "gpu")]
use std::sync::Arc;

use crate::backend::{Backend, dense_statevector_len, statevector_probability_len};
use crate::backend::{reduced_density, schmidt, simd};
use crate::circuit::Instruction;
#[cfg(feature = "gpu")]
use crate::circuit::{QftTextbookStep, qft_textbook_steps};
use crate::error::Result;
use crate::gates::Gate;
use crate::sim::unified_pauli::PauliTerm;

#[cfg(feature = "gpu")]
use crate::gpu::{GpuContext, GpuState};

#[cfg(feature = "parallel")]
use rayon::prelude::*;

#[cfg(feature = "parallel")]
pub(crate) use super::{MIN_PAR_ELEMS, PARALLEL_THRESHOLD_QUBITS};

/// Width from which `init` zeroes the state with the thread pool, which spreads the first
/// touch of every page of a fresh buffer: 157 ms to 63 ms at 26 qubits. At 16 qubits
/// waking the pool cost more than the fill and read +6% on a loop of fresh runs.
#[cfg(feature = "parallel")]
const PARALLEL_INIT_QUBITS: usize = 20;

/// Amplitudes one pool task zeroes on `init`, 1 MiB.
#[cfg(feature = "parallel")]
const INIT_ZERO_CHUNK: usize = 1 << 16;

/// Per-block partial sums for a 1q reduced density matrix: `(p0, p1, r)` over
/// one `2 * 2^qubit` block whose low half holds the qubit-0 amplitudes.
#[inline(always)]
pub(super) fn rdm_block_sums(block: &[Complex64], half: usize) -> (f64, f64, Complex64) {
    let (lo, hi) = block.split_at(half);
    rdm_pair_sums(lo, hi)
}

/// [`rdm_block_sums`] over an already split pair, so a sub-tile of one block
/// contributes without re-deriving the halves.
#[inline(always)]
pub(super) fn rdm_pair_sums(lo: &[Complex64], hi: &[Complex64]) -> (f64, f64, Complex64) {
    let mut p0 = 0.0f64;
    let mut p1 = 0.0f64;
    let mut r = Complex64::new(0.0, 0.0);
    for (&a0, &a1) in lo.iter().zip(hi.iter()) {
        p0 += a0.norm_sqr();
        p1 += a1.norm_sqr();
        r += a1 * a0.conj();
    }
    (p0, p1, r)
}

/// Sum of `(p0, p1, r)` triples, the reduction [`rdm_pair_sums`] folds under.
#[cfg(feature = "parallel")]
#[inline(always)]
pub(super) fn rdm_sum_add(
    a: (f64, f64, Complex64),
    b: (f64, f64, Complex64),
) -> (f64, f64, Complex64) {
    (a.0 + b.0, a.1 + b.1, a.2 + b.2)
}

/// Reject a device statevector that cannot fit the currently free VRAM.
///
/// Guards the hard GPU path before the allocation is attempted, so an
/// over-VRAM circuit surfaces a budget error naming the request and the free
/// memory instead of the raw driver failure. A context that cannot report
/// free memory (stub device) falls through to the allocation attempt, whose
/// own error stays authoritative; so does a lost race against another process
/// allocating between this check and the allocation.
/// Apply `exp(-i θ P / 2)` on the device through its CNOT-ladder lowering.
#[cfg(feature = "gpu")]
fn launch_pauli_rotation_gpu(
    ctx: &GpuContext,
    gpu: &mut GpuState,
    theta: f64,
    targets: &[usize],
    axes: &[crate::sim::unified_pauli::PauliAxis],
) -> Result<()> {
    use crate::gpu::kernels::dense as k;

    let mut result = Ok(());
    crate::circuit::pauli_rotation_lowering(theta, targets, axes, |step, tgts| {
        if result.is_err() {
            return;
        }
        result = match &step {
            Gate::Cx => k::launch_apply_cx(ctx, gpu, tgts[0], tgts[1]),
            _ => {
                let mat = step.matrix_2x2();
                if step.is_diagonal_1q() {
                    k::launch_apply_diagonal_1q(ctx, gpu, tgts[0], mat[0][0], mat[1][1])
                } else {
                    k::launch_apply_gate_1q(ctx, gpu, tgts[0], mat)
                }
            }
        };
    });
    result
}

#[cfg(feature = "gpu")]
fn check_device_budget(context: &GpuContext, num_qubits: usize) -> crate::error::Result<()> {
    if num_qubits >= usize::BITS as usize - 4 {
        return Ok(());
    }
    let Ok(free) = context.vram_available() else {
        return Ok(());
    };
    let needed = (1usize << num_qubits) * crate::gpu::STATEVECTOR_BYTES_PER_AMPLITUDE;
    if needed > free {
        return Err(crate::error::PrismError::ResourceLimit(Box::new(
            crate::error::ResourceLimit {
                backend: "statevector-gpu".to_string(),
                operation: format!("the {num_qubits}-qubit state on the device"),
                resource: crate::error::ResourceKind::DeviceBytes,
                required: needed as u128,
                limit: free as u128,
                env_var: None,
            },
        )));
    }
    Ok(())
}

/// Insert a zero bit at `bit_pos`, shifting all higher bits left by one.
///
/// Maps a compact iteration index to a state-vector index with a gap at
/// `bit_pos`. Chaining multiple calls (in ascending `bit_pos` order) creates
/// gaps at all controlled-gate qubit positions.
#[inline(always)]
pub(crate) fn insert_zero_bit(val: usize, bit_pos: usize) -> usize {
    let lo = val & ((1 << bit_pos) - 1);
    let hi = val >> bit_pos;
    (hi << (bit_pos + 1)) | lo
}

/// Wrapper to send a raw pointer across Rayon threads.
///
/// Callers must partition the state so each thread touches disjoint elements.
/// The mask-based index bijection establishes that for controlled-gate kernels.
#[cfg(feature = "parallel")]
#[derive(Copy, Clone)]
pub(crate) struct SendPtr(pub(crate) *mut Complex64);

#[cfg(feature = "parallel")]
// SAFETY: SendPtr is only used by kernels that partition the state into
// disjoint mutable indices before entering Rayon work.
unsafe impl Send for SendPtr {}
#[cfg(feature = "parallel")]
// SAFETY: Sharing the pointer wrapper is sound because mutation happens only
// through disjoint indices established by each caller's index bijection.
unsafe impl Sync for SendPtr {}

#[cfg(feature = "parallel")]
impl SendPtr {
    #[inline(always)]
    pub(crate) unsafe fn load(self, idx: usize) -> Complex64 {
        // SAFETY: same contract as the enclosing unsafe fn.
        unsafe { *self.0.add(idx) }
    }

    #[inline(always)]
    pub(crate) unsafe fn store(self, idx: usize, val: Complex64) {
        // SAFETY: same contract as the enclosing unsafe fn.
        unsafe {
            *self.0.add(idx) = val;
        }
    }

    #[inline(always)]
    pub(crate) fn as_f64_ptr(self) -> *mut f64 {
        self.0 as *mut f64
    }

    #[inline(always)]
    pub(crate) fn as_complex_ptr(self) -> *mut Complex64 {
        self.0
    }
}

pub struct StatevectorBackend {
    pub(crate) num_qubits: usize,
    pub(crate) state: Vec<Complex64>,
    pub(crate) classical_bits: Vec<bool>,
    pub(crate) rng: ChaCha8Rng,
    pub(crate) pending_norm: f64,
    pub(crate) batch_tables: kernels::BatchTableScratch,
    #[cfg(feature = "gpu")]
    gpu_context: Option<Arc<GpuContext>>,
    #[cfg(feature = "gpu")]
    gpu_state: Option<GpuState>,
    #[cfg(feature = "gpu")]
    gpu_soft: bool,
}

/// Kill switch for the native `QftBlock` FFT path and the QFT recognition that feeds it.
///
/// Set `PRISM_NO_QFT_BLOCK=1` to force textbook expansion for A/B runs.
#[inline]
fn qft_block_enabled() -> bool {
    use std::sync::OnceLock;
    static ENABLED: OnceLock<bool> = OnceLock::new();
    *ENABLED.get_or_init(|| std::env::var_os("PRISM_NO_QFT_BLOCK").is_none())
}

impl StatevectorBackend {
    pub fn new(seed: u64) -> Self {
        Self {
            num_qubits: 0,
            state: Vec::new(),
            classical_bits: Vec::new(),
            rng: ChaCha8Rng::seed_from_u64(seed),
            pending_norm: 1.0,
            batch_tables: kernels::BatchTableScratch::default(),
            #[cfg(feature = "gpu")]
            gpu_context: None,
            #[cfg(feature = "gpu")]
            gpu_state: None,
            #[cfg(feature = "gpu")]
            gpu_soft: false,
        }
    }

    /// Opt into GPU acceleration using the given shared execution context.
    ///
    /// When set, [`Backend::init`] allocates a device-resident state instead of a host
    /// `Vec<Complex64>` and gate application routes through GPU kernels. A device
    /// allocation failure is propagated as an error.
    #[cfg(feature = "gpu")]
    pub fn with_gpu(mut self, context: Arc<GpuContext>) -> Self {
        self.gpu_context = Some(context);
        self
    }

    /// Opt into GPU acceleration but fall back to the host state if the device
    /// allocation fails at [`Backend::init`], instead of erroring. Used by the
    /// `Auto` GPU dispatch path, where an unfit or racing device must degrade to
    /// CPU rather than surface a failure.
    #[cfg(feature = "gpu")]
    pub fn with_gpu_auto(mut self, context: Arc<GpuContext>) -> Self {
        self.gpu_context = Some(context);
        self.gpu_soft = true;
        self
    }

    #[cfg(feature = "gpu")]
    fn apply_gpu(&mut self, instruction: &Instruction) -> Result<()> {
        debug_assert!(
            self.gpu_state.is_some(),
            "apply_gpu called without gpu_state (callers must check self.gpu_state.is_some() first)"
        );
        // Every instruction routes to a kernel once `gpu_state` is Some: an
        // explicit `with_gpu(ctx)` skips the dispatch-level crossover, which
        // `BackendKind::StatevectorGpu` applies per sub-block through
        // `gpu_min_qubits()`. Multi2q launches one kernel per subgate.
        match instruction {
            Instruction::Gate { gate, targets } => self.dispatch_gate_gpu(gate, targets),
            Instruction::Measure {
                qubit,
                classical_bit,
            } => self.apply_measure_gpu(*qubit, *classical_bit),
            Instruction::Reset { qubit } => self.apply_reset_gpu(*qubit),
            Instruction::Barrier { .. } => Ok(()),
            Instruction::Save { label, .. } => {
                Err(crate::backend::save_not_applied("Statevector", label))
            }
            Instruction::Conditional {
                condition,
                gate,
                targets,
            } => {
                if condition.evaluate(&self.classical_bits) {
                    self.dispatch_gate_gpu(gate, targets)
                } else {
                    Ok(())
                }
            }
            Instruction::Region(region) => self.apply_region(region),
        }
    }

    #[cfg(feature = "gpu")]
    fn dispatch_gate_gpu(&mut self, gate: &Gate, targets: &[usize]) -> Result<()> {
        use crate::gpu::kernels::dense as k;

        let gpu = self
            .gpu_state
            .as_mut()
            .expect("dispatch_gate_gpu called without gpu_state");
        let ctx = gpu.context().clone();

        match gate {
            Gate::Rzz(theta) => k::launch_apply_rzz(&ctx, gpu, targets[0], targets[1], *theta),
            Gate::Cx => k::launch_apply_cx(&ctx, gpu, targets[0], targets[1]),
            Gate::Cz => k::launch_apply_cz(&ctx, gpu, targets[0], targets[1]),
            Gate::Swap => k::launch_apply_swap(&ctx, gpu, targets[0], targets[1]),
            Gate::Cu(mat) => {
                if let Some(phase) = gate.controlled_phase() {
                    k::launch_apply_cu_phase(&ctx, gpu, targets[0], targets[1], phase)
                } else {
                    k::launch_apply_cu(&ctx, gpu, targets[0], targets[1], **mat)
                }
            }
            Gate::Mcu(data) => {
                let num_ctrl = data.num_controls as usize;
                let controls = &targets[..num_ctrl];
                let target = targets[num_ctrl];
                if let Some(phase) = gate.controlled_phase() {
                    k::launch_apply_mcu_phase(&ctx, gpu, controls, target, phase)
                } else {
                    k::launch_apply_mcu(&ctx, gpu, controls, target, data.mat)
                }
            }
            Gate::BatchPhase(data) => {
                let control = targets[0];
                k::launch_apply_batch_phase(&ctx, gpu, control, &data.phases)
            }
            Gate::BatchRzz(data) => k::launch_apply_batch_rzz(&ctx, gpu, &data.edges),
            Gate::DiagonalBatch(data) => k::launch_apply_diagonal_batch(&ctx, gpu, &data.entries),
            Gate::QftBlock {
                start,
                num,
                inverse,
                swaps,
                big_endian,
            } => {
                let h = Gate::H.matrix_2x2();
                for step in qft_textbook_steps(
                    *start as usize,
                    *num as usize,
                    *inverse,
                    *swaps,
                    *big_endian,
                ) {
                    match step {
                        QftTextbookStep::Hadamard(q) => k::launch_apply_gate_1q(&ctx, gpu, q, h)?,
                        QftTextbookStep::CPhase {
                            control,
                            target,
                            theta,
                        } => k::launch_apply_cu_phase(
                            &ctx,
                            gpu,
                            control,
                            target,
                            Complex64::from_polar(1.0, theta),
                        )?,
                        QftTextbookStep::Swap(a, b) => k::launch_apply_swap(&ctx, gpu, a, b)?,
                    }
                }
                Ok(())
            }
            Gate::PauliRot(data) => {
                launch_pauli_rotation_gpu(&ctx, gpu, data.theta, targets, &data.axes)
            }
            Gate::MultiPauliRot(data) => {
                let mut result = Ok(());
                data.for_each_rotation(|theta, targets, axes| {
                    if result.is_ok() {
                        result = launch_pauli_rotation_gpu(&ctx, gpu, theta, targets, axes);
                    }
                });
                result
            }
            Gate::Unitary(_) => Err(crate::error::PrismError::BackendUnsupported {
                backend: "statevector (device-resident)".to_string(),
                operation: format!("dense multi-qubit gate `{}`", gate.name()),
            }),
            Gate::MultiFused(data) => {
                if data.all_diagonal {
                    k::launch_apply_multi_fused_diagonal(&ctx, gpu, &data.gates)
                } else {
                    k::launch_apply_multi_fused_nondiag(&ctx, gpu, &data.gates)
                }
            }
            Gate::Fused2q(mat) => k::launch_apply_fused_2q(&ctx, gpu, targets[0], targets[1], mat),
            Gate::Multi2q(data) => {
                for &(q0, q1, mat) in &data.gates {
                    k::launch_apply_fused_2q(&ctx, gpu, q0, q1, &mat)?;
                }
                Ok(())
            }
            Gate::Id
            | Gate::X
            | Gate::Y
            | Gate::Z
            | Gate::H
            | Gate::S
            | Gate::Sdg
            | Gate::T
            | Gate::Tdg
            | Gate::SX
            | Gate::SXdg
            | Gate::Rx(_)
            | Gate::Ry(_)
            | Gate::Rz(_)
            | Gate::P(_)
            | Gate::Fused(_) => {
                let mat = gate.matrix_2x2();
                if gate.is_diagonal_1q() {
                    k::launch_apply_diagonal_1q(&ctx, gpu, targets[0], mat[0][0], mat[1][1])
                } else {
                    k::launch_apply_gate_1q(&ctx, gpu, targets[0], mat)
                }
            }
        }
    }

    #[cfg(feature = "gpu")]
    fn apply_measure_gpu(&mut self, qubit: usize, classical_bit: usize) -> Result<()> {
        use crate::gpu::kernels::dense as k;
        use rand::RngExt;

        let gpu = self
            .gpu_state
            .as_mut()
            .expect("apply_measure_gpu called without gpu_state");
        let ctx = gpu.context().clone();

        let prob_one = k::measure_prob_one(&ctx, gpu, qubit)?;
        let u: f64 = self.rng.random();
        let outcome = u < prob_one;
        self.classical_bits[classical_bit] = outcome;

        k::measure_collapse(&ctx, gpu, qubit, outcome)?;

        let inv_sqrt = crate::backend::measurement_inv_norm(outcome, prob_one);
        gpu.set_pending_norm(gpu.pending_norm() * inv_sqrt);
        Ok(())
    }

    #[cfg(feature = "gpu")]
    fn apply_reset_gpu(&mut self, qubit: usize) -> Result<()> {
        use crate::gpu::kernels::dense as k;
        use rand::RngExt;

        let gpu = self
            .gpu_state
            .as_mut()
            .expect("apply_reset_gpu called without gpu_state");
        let ctx = gpu.context().clone();

        let prob_one = k::measure_prob_one(&ctx, gpu, qubit)?;
        let u: f64 = self.rng.random();
        let outcome = u < prob_one;

        let gpu = self
            .gpu_state
            .as_mut()
            .expect("apply_reset_gpu called without gpu_state");
        k::measure_collapse(&ctx, gpu, qubit, outcome)?;
        let inv_norm = crate::backend::measurement_inv_norm(outcome, prob_one);
        gpu.set_pending_norm(gpu.pending_norm() * inv_norm);

        if outcome {
            self.dispatch_gate_gpu(&Gate::X, &[qubit])?;
        }
        Ok(())
    }

    #[cfg(feature = "gpu")]
    fn apply_1q_matrix_gpu(&mut self, qubit: usize, matrix: &[[Complex64; 2]; 2]) -> Result<()> {
        use crate::gpu::kernels::dense as k;

        let gpu = self
            .gpu_state
            .as_mut()
            .expect("apply_1q_matrix_gpu called without gpu_state");
        let ctx = gpu.context().clone();

        let is_diagonal = matrix[0][1].norm() < 1e-14 && matrix[1][0].norm() < 1e-14;
        if is_diagonal {
            k::launch_apply_diagonal_1q(&ctx, gpu, qubit, matrix[0][0], matrix[1][1])
        } else {
            k::launch_apply_gate_1q(&ctx, gpu, qubit, *matrix)
        }
    }

    /// Read-only access to the raw amplitude vector.
    ///
    /// After measurements, amplitudes may be un-normalized due to deferred
    /// normalization. Call [`export_statevector`](Self::export_statevector)
    /// for a properly normalized copy.
    pub fn state_vector(&self) -> &[Complex64] {
        &self.state
    }

    pub(crate) fn probability_scale(&self) -> f64 {
        self.pending_norm * self.pending_norm
    }

    /// Overwrite this backend's host state with `source`'s, reusing the amplitude
    /// buffer when it already holds enough. The RNG stays this backend's own.
    pub(crate) fn copy_state_from(&mut self, source: &Self) {
        self.num_qubits = source.num_qubits;
        self.pending_norm = source.pending_norm;
        self.classical_bits.clone_from(&source.classical_bits);
        self.state.clone_from(&source.state);
    }

    /// Probabilities of the host state, skipping the dense-output cap.
    ///
    /// For a caller whose vector is a working buffer sized by a state it
    /// already holds, where the output cap does not apply: the distributed
    /// backend builds its rank-local sampling CDF from its own slice. Anything
    /// handing the vector back to a user goes through
    /// [`probabilities`](crate::backend::Backend::probabilities).
    pub(crate) fn host_probability_vector(&self) -> Vec<f64> {
        let norm_sq = self.probability_scale();
        let mut probs = vec![0.0_f64; self.state.len()];

        #[cfg(feature = "parallel")]
        if self.num_qubits >= PARALLEL_THRESHOLD_QUBITS {
            let src_chunks = self.state.par_chunks(MIN_PAR_ELEMS);
            let dst_chunks = probs.par_chunks_mut(MIN_PAR_ELEMS);
            if norm_sq == 1.0 {
                src_chunks.zip(dst_chunks).for_each(|(s, d)| {
                    simd::norm_sqr_to_slice(s, d);
                });
            } else {
                src_chunks.zip(dst_chunks).for_each(|(s, d)| {
                    simd::norm_sqr_to_slice_scaled(s, d, norm_sq);
                });
            }
            return probs;
        }

        if norm_sq == 1.0 {
            simd::norm_sqr_to_slice(&self.state, &mut probs);
        } else {
            simd::norm_sqr_to_slice_scaled(&self.state, &mut probs, norm_sq);
        }
        probs
    }

    /// Whether the amplitudes live in device memory. When true, the host
    /// `state_vector()` slice is empty; read the state through
    /// [`probabilities`](crate::backend::Backend::probabilities) or
    /// [`export_statevector`](crate::backend::Backend::export_statevector).
    #[cfg(feature = "gpu")]
    pub(crate) fn is_gpu_resident(&self) -> bool {
        self.gpu_state.is_some()
    }

    #[cfg(not(feature = "gpu"))]
    pub(crate) fn is_gpu_resident(&self) -> bool {
        false
    }

    /// `sum_j conj(psi[j ^ xmask]) psi[j] (-1)^{popcount(j & zmask)}` per mask
    /// pair, reduced on the device; `None` when the state lives on the host.
    /// Normalization is applied, so an appended `(0, 0)` mask returns `<psi|psi>`.
    #[cfg(feature = "gpu")]
    pub(crate) fn gpu_pauli_sums(&self, masks: &[(u64, u64)]) -> Option<Result<Vec<Complex64>>> {
        let gpu = self.gpu_state.as_ref()?;
        Some(crate::gpu::kernels::dense::pauli_sums(
            gpu.context(),
            gpu,
            masks,
        ))
    }

    #[cfg(not(feature = "gpu"))]
    pub(crate) fn gpu_pauli_sums(&self, _masks: &[(u64, u64)]) -> Option<Result<Vec<Complex64>>> {
        None
    }

    /// Group moments on a device-resident state through
    /// [`group_moments`](crate::gpu::kernels::dense::group_moments). `None` when the
    /// state lives on the host or a rotated copy does not fit on the device.
    #[cfg(feature = "gpu")]
    pub(crate) fn group_moments_on_device(
        &self,
        groups: &[(Vec<u64>, Vec<f64>, u64, u64)],
    ) -> Option<Result<Vec<(f64, f64)>>> {
        let gpu = self.gpu_state.as_ref()?;
        crate::gpu::kernels::dense::group_moments(gpu.context(), gpu, groups).transpose()
    }

    #[cfg(not(feature = "gpu"))]
    pub(crate) fn group_moments_on_device(
        &self,
        _groups: &[(Vec<u64>, Vec<f64>, u64, u64)],
    ) -> Option<Result<Vec<(f64, f64)>>> {
        None
    }

    /// [`pauli_expectations_from_masks`](crate::sim::pauli_expectations_from_masks)
    /// evaluated on a device-resident state: one reduction launch over every
    /// mask plus an appended identity mask that supplies the norm, so nothing
    /// but `16 * (masks.len() + 1)` bytes leaves the card. `None` when the
    /// state lives on the host.
    pub(crate) fn pauli_expectations_on_device(
        &self,
        masks: &[(usize, usize, u32)],
    ) -> Option<Result<Vec<f64>>> {
        if !self.is_gpu_resident() {
            return None;
        }
        let request: Vec<(u64, u64)> = masks
            .iter()
            .map(|&(xmask, zmask, _)| (xmask as u64, zmask as u64))
            .chain(std::iter::once((0, 0)))
            .collect();
        let sums = match self.gpu_pauli_sums(&request)? {
            Ok(sums) => sums,
            Err(e) => return Some(Err(e)),
        };
        let norm = sums[masks.len()].re;
        if norm == 0.0 {
            return Some(Ok(vec![0.0; masks.len()]));
        }
        Some(Ok(masks
            .iter()
            .zip(&sums)
            .map(|(&(_, _, num_y), sum)| (sum * crate::sim::i_pow(num_y)).re / norm)
            .collect()))
    }

    #[cfg(feature = "gpu")]
    pub(crate) fn gpu_state(&self) -> Option<&GpuState> {
        self.gpu_state.as_ref()
    }

    #[cfg(feature = "gpu")]
    pub(crate) fn gpu_state_mut(&mut self) -> Option<&mut GpuState> {
        self.gpu_state.as_mut()
    }

    /// Initialize the backend from a pre-computed state vector.
    ///
    /// Accepts ownership of the amplitude vector, bypassing the default |0...0⟩
    /// initialization. The vector length must be a power of 2. Normalization is
    /// not checked: the adjoint gradient seeds its lambda register with `H|ψ⟩`,
    /// whose norm is not 1, and reads raw amplitudes. A start state coming from
    /// a user goes through
    /// [`init_from_amplitudes`](crate::backend::Backend::init_from_amplitudes),
    /// which validates first.
    pub fn init_from_state(
        &mut self,
        state: Vec<Complex64>,
        num_classical_bits: usize,
    ) -> crate::error::Result<()> {
        #[cfg(feature = "parallel")]
        crate::backend::init_thread_pool();

        let dim = state.len();
        if !dim.is_power_of_two() || dim < 2 {
            return Err(crate::error::PrismError::InvalidParameter {
                message: format!(
                    "state vector length must be a power of 2 and >= 2, got {}",
                    dim
                ),
            });
        }
        self.num_qubits = dim.trailing_zeros() as usize;
        self.pending_norm = 1.0;
        crate::backend::init_classical_bits(&mut self.classical_bits, num_classical_bits);

        #[cfg(feature = "gpu")]
        if let Some(ctx) = self.gpu_context.clone() {
            if !self.gpu_soft {
                check_device_budget(&ctx, self.num_qubits)?;
            }
            match GpuState::from_host_amplitudes(ctx, &state) {
                Ok(gpu) => {
                    self.state.clear();
                    self.gpu_state = Some(gpu);
                    return Ok(());
                }
                Err(e) if !self.gpu_soft => return Err(e),
                Err(_) => {
                    self.gpu_context = None;
                    self.gpu_state = None;
                }
            }
        }

        self.state = state;
        Ok(())
    }

    /// Apply a gate to the host state without constructing an owned
    /// [`Instruction`]. The distributed backend dispatches its local gates
    /// through this to avoid deep-cloning boxed gate payloads per application.
    #[inline(always)]
    pub(crate) fn dispatch_gate(&mut self, gate: &Gate, targets: &[usize]) {
        match gate {
            Gate::Rzz(theta) => self.apply_rzz(targets[0], targets[1], *theta),
            Gate::Cx => self.apply_cx(targets[0], targets[1]),
            Gate::Cz => self.apply_cz(targets[0], targets[1]),
            Gate::Swap => self.apply_swap(targets[0], targets[1]),
            Gate::Cu(mat) => {
                if let Some(phase) = gate.controlled_phase() {
                    self.apply_cu_phase(targets[0], targets[1], phase);
                } else {
                    self.apply_cu(targets[0], targets[1], **mat);
                }
            }
            Gate::Mcu(data) => {
                let num_ctrl = data.num_controls as usize;
                if let Some(phase) = gate.controlled_phase() {
                    self.apply_mcu_phase(&targets[..num_ctrl], targets[num_ctrl], phase);
                } else {
                    self.apply_mcu(&targets[..num_ctrl], targets[num_ctrl], data.mat);
                }
            }
            Gate::BatchPhase(data) => {
                self.apply_batch_phase(targets[0], &data.phases);
            }
            &Gate::QftBlock {
                start,
                num,
                inverse,
                swaps,
                big_endian,
            } => {
                if start == 0 {
                    self.apply_qft_block(num as usize, inverse, swaps, big_endian);
                } else {
                    self.apply_qft_block_textbook(
                        start as usize,
                        num as usize,
                        inverse,
                        swaps,
                        big_endian,
                    );
                }
            }
            Gate::BatchRzz(data) => {
                self.apply_batch_rzz(&data.edges);
            }
            Gate::DiagonalBatch(data) => {
                self.apply_diagonal_batch(&data.entries);
            }
            Gate::PauliRot(data) => {
                self.apply_pauli_rot(targets, data.theta, &data.axes);
            }
            Gate::MultiPauliRot(data) => {
                self.apply_multi_pauli_rot(&data.rotations);
            }
            Gate::Unitary(data) => {
                self.apply_unitary(targets, data.matrix());
            }
            Gate::MultiFused(data) => {
                if data.all_diagonal {
                    self.apply_multi_1q_diagonal(&data.gates);
                } else {
                    self.apply_multi_1q(&data.gates);
                }
            }
            Gate::Fused2q(mat) => {
                self.apply_fused_2q(targets[0], targets[1], mat);
            }
            Gate::Multi2q(data) => {
                self.apply_multi_2q(&data.gates);
            }
            Gate::Id
            | Gate::X
            | Gate::Y
            | Gate::Z
            | Gate::H
            | Gate::S
            | Gate::Sdg
            | Gate::T
            | Gate::Tdg
            | Gate::SX
            | Gate::SXdg
            | Gate::Rx(_)
            | Gate::Ry(_)
            | Gate::Rz(_)
            | Gate::P(_)
            | Gate::Fused(_) => {
                let mat = gate.matrix_2x2();
                if gate.is_diagonal_1q() {
                    self.apply_diagonal_gate(targets[0], mat[0][0], mat[1][1]);
                } else {
                    self.apply_single_gate(targets[0], mat);
                }
            }
        }
    }
}

impl crate::backend::sealed::Sealed for StatevectorBackend {}

impl Backend for StatevectorBackend {
    fn name(&self) -> &'static str {
        "statevector"
    }

    fn as_any(&self) -> Option<&dyn std::any::Any> {
        Some(self)
    }

    fn resolved(&self) -> crate::sim::ResolvedBackend {
        crate::sim::ResolvedBackend::Statevector
    }

    fn placement(&self) -> crate::sim::Placement {
        if self.is_gpu_resident() {
            crate::sim::Placement::Device
        } else {
            crate::sim::Placement::Host
        }
    }

    fn supports_qft_block(&self) -> bool {
        if !qft_block_enabled() {
            return false;
        }
        #[cfg(feature = "gpu")]
        {
            self.gpu_context.is_none()
        }
        #[cfg(not(feature = "gpu"))]
        {
            true
        }
    }

    fn supports_pauli_rotation(&self) -> bool {
        #[cfg(feature = "gpu")]
        {
            self.gpu_context.is_none()
        }
        #[cfg(not(feature = "gpu"))]
        {
            true
        }
    }

    fn init(&mut self, num_qubits: usize, num_classical_bits: usize) -> Result<()> {
        #[cfg(feature = "parallel")]
        crate::backend::init_thread_pool();

        self.num_qubits = num_qubits;
        self.pending_norm = 1.0;
        crate::backend::init_classical_bits(&mut self.classical_bits, num_classical_bits);

        // The device path is budgeted by VRAM, not by the host cap: the host
        // vector is never allocated while the state lives on the device.
        #[cfg(feature = "gpu")]
        if let Some(ctx) = self.gpu_context.clone() {
            if !self.gpu_soft {
                check_device_budget(&ctx, num_qubits)?;
            }
            match GpuState::new(ctx, num_qubits) {
                Ok(state) => {
                    self.state.clear();
                    self.gpu_state = Some(state);
                    return Ok(());
                }
                Err(e) if !self.gpu_soft => return Err(e),
                Err(_) => {
                    self.gpu_context = None;
                    self.gpu_state = None;
                }
            }
        }

        crate::backend::check_state_allocation(
            "statevector",
            num_qubits,
            crate::backend::max_statevector_qubits(),
            "PRISM_MAX_SV_QUBITS",
        )?;

        let dim = 1usize << num_qubits;
        let zero = Complex64::new(0.0, 0.0);
        let reuse = self.state.len() == dim;
        if !reuse {
            self.state = Vec::new();
        }
        #[cfg(feature = "parallel")]
        if num_qubits >= PARALLEL_INIT_QUBITS {
            if !reuse {
                self.state = Vec::with_capacity(dim);
            }
            let base = SendPtr(self.state.as_mut_ptr());
            (0..dim / INIT_ZERO_CHUNK).into_par_iter().for_each(|c| {
                // SAFETY: the buffer has capacity for `dim` elements and each chunk index
                // addresses its own `INIT_ZERO_CHUNK` of them. A per-chunk `fill` wrote
                // through the cache and re-zeroed a 2^28 state in 342 ms against 226 ms
                // for one thread's memset; `write_bytes` keeps memset's streaming stores.
                unsafe {
                    std::ptr::write_bytes(
                        base.as_complex_ptr().add(c * INIT_ZERO_CHUNK),
                        0,
                        INIT_ZERO_CHUNK,
                    );
                }
            });
            // SAFETY: every element below `dim` was written above, and all-zero bytes are
            // the value 0 + 0i.
            unsafe { self.state.set_len(dim) };
            self.state[0] = Complex64::new(1.0, 0.0);
            return Ok(());
        }
        if reuse {
            self.state.fill(zero);
        } else {
            self.state = vec![zero; dim];
        }
        self.state[0] = Complex64::new(1.0, 0.0);
        Ok(())
    }

    fn supports_initial_state(&self) -> bool {
        true
    }

    fn init_from_amplitudes(
        &mut self,
        amplitudes: Vec<Complex64>,
        num_classical_bits: usize,
    ) -> Result<()> {
        crate::backend::validate_initial_amplitudes(&amplitudes)?;
        self.init_from_state(amplitudes, num_classical_bits)
    }

    fn apply(&mut self, instruction: &Instruction) -> Result<()> {
        #[cfg(feature = "gpu")]
        if self.gpu_state.is_some() {
            return self.apply_gpu(instruction);
        }
        match instruction {
            Instruction::Gate { gate, targets } => self.dispatch_gate(gate, targets),
            Instruction::Measure {
                qubit,
                classical_bit,
            } => {
                self.apply_measure(*qubit, *classical_bit);
            }
            Instruction::Reset { qubit } => {
                self.apply_reset(*qubit);
            }
            Instruction::Barrier { .. } => {}
            Instruction::Save { label, .. } => {
                return Err(crate::backend::save_not_applied("Statevector", label));
            }
            Instruction::Conditional {
                condition,
                gate,
                targets,
            } => {
                if condition.evaluate(&self.classical_bits) {
                    self.dispatch_gate(gate, targets);
                }
            }
            Instruction::Region(region) => self.apply_region(region)?,
        }
        Ok(())
    }

    fn classical_results(&self) -> &[bool] {
        &self.classical_bits
    }

    fn probabilities(&self) -> Result<Vec<f64>> {
        #[cfg(feature = "gpu")]
        if let Some(gpu) = self.gpu_state.as_ref() {
            return gpu.probabilities();
        }
        statevector_probability_len(self.name(), self.num_qubits)?;
        Ok(self.host_probability_vector())
    }

    fn num_qubits(&self) -> usize {
        self.num_qubits
    }

    fn export_statevector(&self) -> crate::error::Result<Vec<Complex64>> {
        #[cfg(feature = "gpu")]
        if let Some(gpu) = self.gpu_state.as_ref() {
            return gpu.export_statevector();
        }
        dense_statevector_len(self.name(), "statevector export", self.num_qubits)?;
        if self.pending_norm == 1.0 {
            return Ok(self.state.clone());
        }
        let s = Complex64::new(self.pending_norm, 0.0);
        Ok(self.state.iter().map(|&c| c * s).collect())
    }

    fn qubit_probability(&self, qubit: usize) -> crate::error::Result<f64> {
        #[cfg(feature = "gpu")]
        if let Some(gpu) = self.gpu_state.as_ref() {
            return crate::gpu::kernels::dense::measure_prob_one(gpu.context(), gpu, qubit);
        }
        Ok(self.qubit_probability_one(qubit))
    }

    fn reduced_density_matrix_1q(&self, qubit: usize) -> Result<[[Complex64; 2]; 2]> {
        #[cfg(feature = "gpu")]
        if let Some(gpu) = self.gpu_state.as_ref() {
            return crate::gpu::kernels::dense::reduced_density_matrix_1q(
                gpu.context(),
                gpu,
                qubit,
            );
        }
        Ok(self.reduced_density_matrix_one(qubit))
    }

    fn supports_two_qubit_kraus(&self) -> bool {
        #[cfg(feature = "gpu")]
        if self.gpu_state.is_some() {
            return false;
        }
        true
    }

    fn reduced_density_matrix_2q(&mut self, q0: usize, q1: usize) -> Result<[[Complex64; 4]; 4]> {
        assert_ne!(q0, q1, "reduced_density_matrix_2q needs distinct qubits");
        #[cfg(feature = "gpu")]
        if self.gpu_state.is_some() {
            return Err(crate::error::PrismError::BackendUnsupported {
                backend: "statevector (device-resident)".to_string(),
                operation: "reduced_density_matrix_2q".to_string(),
            });
        }
        Ok(self.reduced_density_matrix_two(q0, q1))
    }

    fn supports_pauli_expectation(&self) -> bool {
        true
    }

    /// One pass over the amplitudes per observable group, in place.
    ///
    /// The result divides by the norm, so `pending_norm` needs no application:
    /// it scales bra and ket alike. A device-resident state reduces on the card
    /// and reads back one sum per observable.
    fn pauli_expectations(&self, observables: &[Vec<PauliTerm>]) -> Result<Vec<f64>> {
        let masks = observables
            .iter()
            .map(|observable| crate::sim::pauli_masks(observable, self.num_qubits))
            .collect::<Result<Vec<_>>>()?;
        if let Some(values) = self.pauli_expectations_on_device(&masks) {
            return values;
        }
        Ok(crate::sim::pauli_expectations_from_masks(&self.state, &masks).0)
    }

    fn schmidt_values(&mut self, subsystem: &[usize]) -> Result<Vec<f64>> {
        schmidt::validate_subsystem(subsystem, self.num_qubits)?;
        #[cfg(feature = "gpu")]
        if let Some(gpu) = self.gpu_state.as_ref() {
            let state = gpu.export_statevector()?;
            return schmidt::dense_schmidt_values(self.name(), &state, self.num_qubits, subsystem);
        }
        schmidt::dense_schmidt_values(self.name(), &self.state, self.num_qubits, subsystem)
    }

    /// Partial trace over the complement at `2^(n + k)` multiply-adds, scaled
    /// to trace one afterwards, which covers the `pending_norm` too. A
    /// device-resident state is read back first.
    fn reduced_density_matrix(&mut self, subsystem: &[usize]) -> Result<Vec<Complex64>> {
        schmidt::validate_qubit_set(subsystem, self.num_qubits)?;
        let dim = reduced_density::reduced_density_side(self.name(), subsystem.len())?;
        #[cfg(feature = "gpu")]
        let readback = self
            .gpu_state
            .as_ref()
            .map(|gpu| gpu.export_statevector())
            .transpose()?;
        #[cfg(feature = "gpu")]
        let state = readback.as_deref().unwrap_or(self.state.as_slice());
        #[cfg(not(feature = "gpu"))]
        let state = self.state.as_slice();
        let mut rho = reduced_density::dense_reduced_density(state, self.num_qubits, subsystem);
        reduced_density::normalize_trace(&mut rho, dim);
        Ok(rho)
    }

    fn reset(&mut self, qubit: usize) -> Result<()> {
        #[cfg(feature = "gpu")]
        if self.gpu_state.is_some() {
            return self.apply_reset_gpu(qubit);
        }
        self.apply_reset(qubit);
        Ok(())
    }

    fn apply_1q_matrix(&mut self, qubit: usize, matrix: &[[Complex64; 2]; 2]) -> Result<()> {
        #[cfg(feature = "gpu")]
        if self.gpu_state.is_some() {
            return self.apply_1q_matrix_gpu(qubit, matrix);
        }
        let is_diagonal = matrix[0][1].norm() < 1e-14 && matrix[1][0].norm() < 1e-14;
        if is_diagonal {
            self.apply_diagonal_gate(qubit, matrix[0][0], matrix[1][1]);
        } else {
            self.apply_single_gate(qubit, *matrix);
        }
        Ok(())
    }
}