prism-q 0.35.0

//! Exact density-matrix backend.
//!
//! # Memory layout
//!
//! `rho` for `n` qubits is a `4^n` buffer, row-major: index `(r << n) | c` holds
//! `<r|rho|c>`. That is a `2n`-qubit statevector with the ket in the high `n`
//! qubits and the bra in the low `n`, so gates reuse the statevector kernels:
//! `U` on the ket register gives `U rho`, and `conj(U)` on the bra register gives
//! `rho U^dagger`. At 16 bytes per entry the ceiling is about 14 qubits on a
//! 16 GiB host and 15 on 32 GiB. With [`DensityMatrixBackend::with_gpu`] the
//! buffer lives in VRAM, where an 11 GiB card holds 13 qubits (1 GiB at 13,
//! 4 GiB at 14 plus scratch) and every sweep runs as a kernel over the embedded
//! buffer. Explicit dispatch only; `Auto` never picks it.
//!
//! # Gate support
//!
//! Unitary evolution, probabilities, measurement with collapse, reset,
//! classically conditioned gates, exact one- and two-qubit Kraus channels
//! (`apply_1q_kraus`, `apply_2q_kraus`, and the twirled closed form
//! `apply_2q_depolarizing`), and exact `Tr(rho P)` through
//! [`Backend::pauli_expectations`]. Fusion runs, gated on the `2n` buffer width.
//! `QftBlock` is remapped onto the ket register before the left product.
//! On the host, `MultiFused` past the parallel threshold and `Multi2q` run their
//! constituents and the bra conjugates as batches in the statevector kernels,
//! `Multi2q` in the longest runs that keep list order, whenever the batch takes
//! fewer buffer passes than one constituent at a time.
//! Diagonal gates skip the two-product route: their two factors combine into one
//! table, so `Rzz` and the diagonal batches each sweep the buffer once.
//!
//! # When to prefer this backend
//!
//! - Exact noise-channel evolution: one run yields the exact mixed state,
//!   where trajectory averaging converges as `1/sqrt(shots)`. Selecting it
//!   with a noise model attached routes every `Simulate` terminal except the two
//!   gradient terminals to that one evolution.
//! - Mixed-state diagnostics: purity and exact `Tr(rho P)` observables.
//! - Continuing from a stored mixture: [`Backend::init_from_density_matrix`]
//!   takes the buffer [`DensityMatrixBackend::density_matrix`] exports.
//!
//! # When NOT to use this backend
//!
//! - Pure-state circuits; a statevector covers twice the qubits in the same
//!   memory.
//! - Qubit counts past the `4^n` ceiling; trajectory sampling on a pure-state
//!   backend scales further.
//! - Noisy circuits with mid-circuit measurement or classical conditioning.
//!   The mixture holds every branch at once, so per-shot feedback cannot be
//!   replayed from it and the noisy terminals reject those shapes.
//! - Adjoint gradients under noise. The adjoint backpropagates against a pure
//!   state and a channel has no reverse evolution to walk; the parameter-shift
//!   terminal serves the noisy gradient from the mixture instead.

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

use num_complex::Complex64;
use rand::{RngExt, SeedableRng};
use rand_chacha::ChaCha8Rng;
use smallvec::SmallVec;

use crate::backend::simd;
use crate::backend::statevector::{StatevectorBackend, insert_zero_bit, kernels};
use crate::backend::{Backend, NORM_CLAMP_MIN, reduced_density, schmidt};
use crate::circuit::{ClassicalCondition, Instruction};
use crate::error::Result;
use crate::gates::{DiagEntry, Gate, McuData, diag_entries_phase, mat_mul_4x4};
#[cfg(feature = "gpu")]
use crate::gpu::{GpuContext, GpuState, kernels::density as dk};
use crate::sim::i_pow;
use crate::sim::unified_pauli::PauliTerm;

/// Compile a one-qubit Kraus set into the 4x4 superoperator acting on the
/// `(row-bit, col-bit)` block of `rho`, where block index `i = 2*a + b` orders
/// `(row-bit a, col-bit b)`:
/// `S[2a+b][2a'+b'] = sum_k K_k[a][a'] * conj(K_k[b][b'])`.
///
/// A single-element set `[U]` gives the unitary sandwich `U rho U^dagger`.
fn block_superoperator(kraus: &[[[Complex64; 2]; 2]]) -> [[Complex64; 4]; 4] {
    let mut s = [[Complex64::new(0.0, 0.0); 4]; 4];
    for k in kraus {
        for a in 0..2 {
            for b in 0..2 {
                for ap in 0..2 {
                    for bp in 0..2 {
                        s[2 * a + b][2 * ap + bp] += k[a][ap] * k[b][bp].conj();
                    }
                }
            }
        }
    }
    s
}

/// Compile a two-qubit Kraus set into the 16x16 block superoperator
/// [`DensityMatrixBackend::apply_2q_kraus`] sweeps:
/// `S[4*tr+tc][4*trp+tcp] = sum_k K_k[tr][trp] * conj(K_k[tc][tcp])`.
fn kraus_2q_superoperator(kraus: &[[[Complex64; 4]; 4]]) -> [[Complex64; 16]; 16] {
    let mut s = [[Complex64::new(0.0, 0.0); 16]; 16];
    for k in kraus {
        for tr in 0..4 {
            for trp in 0..4 {
                let kr = k[tr][trp];
                if kr == Complex64::new(0.0, 0.0) {
                    continue;
                }
                for tc in 0..4 {
                    for tcp in 0..4 {
                        s[4 * tr + tc][4 * trp + tcp] += kr * k[tc][tcp].conj();
                    }
                }
            }
        }
    }
    s
}

/// Collapse one tile of `rho` onto the `outcome` subspace of the qubit whose
/// row and column bits in the buffer index are `rmask` and `cmask`, scaling the
/// survivors by `scale`. `base` is the tile's buffer index, and the tile must
/// not straddle `rmask`, so one test fixes its row class: a tile on the
/// eliminated row is a contiguous zero fill, and only the surviving row pays a
/// per-entry column test.
fn project_tile(
    tile: &mut [Complex64],
    base: usize,
    rmask: usize,
    cmask: usize,
    outcome: bool,
    scale: Complex64,
) {
    if ((base & rmask) != 0) != outcome {
        simd::zero_slice(tile);
        return;
    }
    let zero = Complex64::new(0.0, 0.0);
    for (j, amp) in tile.iter_mut().enumerate() {
        *amp = if (((base + j) & cmask) != 0) == outcome {
            *amp * scale
        } else {
            zero
        };
    }
}

/// Fold the row-set half of one `2 * rmask` block of `rho` onto the row-clear
/// half, then clear it. `r0` and `r1` are the two halves, or matching sub-tiles
/// of them, and `cmask` is the qubit's column bit.
fn reset_fold_pair(r0: &mut [Complex64], r1: &mut [Complex64], cmask: usize) {
    let zero = Complex64::new(0.0, 0.0);
    for (j, amp) in r0.iter_mut().enumerate() {
        *amp = if j & cmask == 0 {
            *amp + r1[j | cmask]
        } else {
            zero
        };
    }
    simd::zero_slice(r1);
}

/// Tile size for a parallel sweep over `rho` whose body reads the qubit's row
/// class from the tile's own base index. Both `rmask` and `2 * cmask` are powers
/// of two and `rmask >= 2 * cmask`, so the result is a multiple of the column run
/// that divides `rmask`: a tile can never straddle the row bit.
#[cfg(feature = "parallel")]
fn row_aligned_tile(cmask: usize, rmask: usize) -> usize {
    (cmask << 1).max(crate::backend::MIN_PAR_ELEMS).min(rmask)
}

/// Base buffer index of block `m`, expanding the four block bit positions out
/// of the compacted index. `positions` must be ascending.
#[inline(always)]
fn block_base(m: usize, positions: &[usize; 4]) -> usize {
    let mut base = m;
    for &pos in positions {
        base = insert_zero_bit(base, pos);
    }
    base
}

/// Index into the 16-entry superoperator diagonal for the amplitude at buffer
/// index `idx`: `4 * tr + tc` with `tr` read from the ket bits and `tc` from
/// the bra bits of `(q0, q1)`, matching the flat order in
/// [`DensityMatrixBackend::block_layout`].
#[inline(always)]
fn diag_slot(idx: usize, q0: usize, q1: usize, n: usize) -> usize {
    ((idx >> (q0 + n)) & 1) << 3
        | ((idx >> (q1 + n)) & 1) << 2
        | ((idx >> q0) & 1) << 1
        | ((idx >> q1) & 1)
}

fn conjugate_2x2(m: &[[Complex64; 2]; 2]) -> [[Complex64; 2]; 2] {
    [
        [m[0][0].conj(), m[0][1].conj()],
        [m[1][0].conj(), m[1][1].conj()],
    ]
}

fn conjugate_4x4(m: &[[Complex64; 4]; 4]) -> [[Complex64; 4]; 4] {
    let mut out = *m;
    for row in &mut out {
        for entry in row {
            *entry = entry.conj();
        }
    }
    out
}

/// The gate's 2x2 matrix, or `None` for anything else. `Gate::num_qubits` is
/// not a usable test here: it reports 1 for a `BatchPhase` with no phases, a
/// single-qubit `DiagonalBatch`, and `QftBlock { num: 1 }`, none of which
/// `matrix_2x2` accepts.
fn matrix_1q(gate: &Gate) -> Option<[[Complex64; 2]; 2]> {
    match gate {
        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(_) => Some(gate.matrix_2x2()),
        _ => None,
    }
}

/// The gate's 4x4 matrix over `(targets[0], targets[1])`, or `None` for a
/// variant [`Gate::matrix_4x4`] does not cover.
fn matrix_2q(gate: &Gate) -> Option<[[Complex64; 4]; 4]> {
    match gate {
        Gate::Rzz(_) | Gate::Cx | Gate::Cz | Gate::Swap | Gate::Cu(_) | Gate::Fused2q(_) => {
            Some(gate.matrix_4x4())
        }
        Gate::PauliRot(data) if data.axes.len() == 2 => Some(gate.matrix_4x4()),
        _ => None,
    }
}

/// `A (x) B` over the 16x16 block index `4 * tr + tc`, `A` on the first
/// qubit's `(row-bit, col-bit)` pair and `B` on the second's.
fn kron_block_superoperators(
    a: &[[Complex64; 4]; 4],
    b: &[[Complex64; 4]; 4],
) -> [[Complex64; 16]; 16] {
    let pair = |t: usize, u: usize, shift: usize| 2 * ((t >> shift) & 1) + ((u >> shift) & 1);
    std::array::from_fn(|row| {
        let (tr, tc) = (row >> 2, row & 3);
        std::array::from_fn(|col| {
            let (trp, tcp) = (col >> 2, col & 3);
            a[pair(tr, tc, 1)][pair(trp, tcp, 1)] * b[pair(tr, tc, 0)][pair(trp, tcp, 0)]
        })
    })
}

fn mat_mul_16x16(a: &[[Complex64; 16]; 16], b: &[[Complex64; 16]; 16]) -> [[Complex64; 16]; 16] {
    std::array::from_fn(|r| std::array::from_fn(|c| (0..16).map(|k| a[r][k] * b[k][c]).sum()))
}

const IDENTITY_4X4: [[Complex64; 4]; 4] = {
    let z = Complex64::new(0.0, 0.0);
    let o = Complex64::new(1.0, 0.0);
    [[o, z, z, z], [z, o, z, z], [z, z, o, z], [z, z, z, o]]
};

/// Narrowest register the exact noisy walk defers one-qubit maps on. Below it
/// the native one-qubit and two-qubit kernels beat the composed sweeps.
pub(crate) const MIN_DEFERRED_QUBITS: usize = 7;

/// Fewest pending maps and trailing channels a two-qubit gate must absorb for its
/// dense 16x16 sweep to beat the native kernels, by where the buffer sits. Past the
/// last-level cache every sweep streams from memory and costs about the same, so one
/// map pays. Resident in a quarter of it or more, the sweep is compute-bound and costs
/// about two one-qubit block sweeps: on an i7-6700K (8 MB L3) a one-map fold at 9
/// qubits read 10% slower than a CX or controlled phase plus its block sweep, and a
/// one-map `Rzz` fold 11% faster. Smaller still, composing the 16x16 map is itself a
/// sweep's worth of work and three maps pay for it; four lost the three-map `Rzz`
/// folds at 7 qubits.
const MIN_FOLDED_MAPS_PAST_CACHE: usize = 1;
const MIN_FOLDED_MAPS_IN_CACHE: usize = 2;
const MIN_FOLDED_MAPS_SMALL: usize = 3;

/// Last-level cache assumed when the OS reports none: the 8 MB L3 the floors above were
/// measured against.
const ASSUMED_LLC_BYTES: usize = 8 << 20;

/// The fold floor for an `num_qubits` buffer on this host, see
/// [`MIN_FOLDED_MAPS_PAST_CACHE`].
fn min_folded_maps(num_qubits: usize) -> usize {
    let bytes = size_of::<Complex64>() << (2 * num_qubits);
    let llc = crate::backend::cache::last_level_cache().map_or(ASSUMED_LLC_BYTES, |c| c.bytes);
    if bytes >= llc {
        MIN_FOLDED_MAPS_PAST_CACHE
    } else if bytes >= llc / 4 {
        MIN_FOLDED_MAPS_IN_CACHE
    } else {
        MIN_FOLDED_MAPS_SMALL
    }
}

/// One-qubit block superoperators the exact noisy walk holds back per qubit.
///
/// Gates and one-qubit channels on a qubit compose here rather than sweeping
/// the buffer one at a time, whichever instruction the channel trails. A
/// two-qubit gate folds both its qubits' pending maps and its own trailing
/// channels on the pair, one- or two-qubit, into one 16x16 sweep, and any
/// other instruction on a held qubit flushes it first. Every sweep costs about
/// the same on a buffer past the caches, so a CX with relaxation on both
/// targets and a pair depolarizing channel drops from five sweeps to one.
pub(crate) struct DeferredSuperoperators {
    pending: Vec<Option<[[Complex64; 4]; 4]>>,
    min_maps: usize,
}

/// A channel trailing a two-qubit gate, lowered for
/// [`DeferredSuperoperators::apply_2q`].
pub(crate) enum PairChannel {
    /// Block superoperator on one of the gate's two qubits.
    Block(usize, Box<[[Complex64; 4]; 4]>),
    /// 16x16 superoperator over the gate's pair, indexed for `(q0, q1)` as
    /// [`DensityMatrixBackend::apply_2q_kraus`] compiles it.
    Pair(Box<[[Complex64; 16]; 16]>),
}

impl PairChannel {
    pub(crate) fn block(qubit: usize, kraus: &[[[Complex64; 2]; 2]]) -> Self {
        Self::Block(qubit, Box::new(block_superoperator(kraus)))
    }

    /// A two-qubit Kraus set given on `(q1, q0)` when `swapped`.
    pub(crate) fn kraus_2q(kraus: &[[[Complex64; 4]; 4]], swapped: bool) -> Self {
        let s = kraus_2q_superoperator(kraus);
        if !swapped {
            return Self::Pair(Box::new(s));
        }
        let swap = |i: usize| {
            let t = |t: usize| ((t & 1) << 1) | (t >> 1);
            4 * t(i >> 2) + t(i & 3)
        };
        Self::Pair(Box::new(std::array::from_fn(|r| {
            std::array::from_fn(|c| s[swap(r)][swap(c)])
        })))
    }

    /// Symmetric two-qubit depolarizing, the 16x16 form of
    /// [`DensityMatrixBackend::apply_2q_depolarizing`]'s closed form.
    pub(crate) fn depolarizing_2q(p: f64) -> Self {
        let alpha = Complex64::new(1.0 - 16.0 * p / 15.0, 0.0);
        let beta = Complex64::new(4.0 * p / 15.0, 0.0);
        let zero = Complex64::new(0.0, 0.0);
        Self::Pair(Box::new(std::array::from_fn(|r| {
            std::array::from_fn(|c| {
                let mut e = if r == c { alpha } else { zero };
                if r % 5 == 0 && c % 5 == 0 {
                    e += beta;
                }
                e
            })
        })))
    }
}

impl DeferredSuperoperators {
    pub(crate) fn new(num_qubits: usize) -> Self {
        Self {
            pending: vec![None; num_qubits],
            min_maps: min_folded_maps(num_qubits),
        }
    }

    /// Compose `gate` onto `qubit`'s pending map. False, with nothing held,
    /// when `gate` has no 2x2 matrix.
    pub(crate) fn defer_gate(&mut self, gate: &Gate, qubit: usize) -> bool {
        let Some(mat) = matrix_1q(gate) else {
            return false;
        };
        self.defer_channel(qubit, &[mat]);
        true
    }

    /// Compose the Kraus set onto `qubit`'s pending map.
    pub(crate) fn defer_channel(&mut self, qubit: usize, kraus: &[[[Complex64; 2]; 2]]) {
        let mut s = block_superoperator(kraus);
        if let Some(prior) = &self.pending[qubit] {
            s = mat_mul_4x4(&s, prior);
        }
        self.pending[qubit] = Some(s);
    }

    /// Apply `gate` on `(q0, q1)` in one sweep, preceded by both qubits'
    /// pending maps and followed by `channels` in order, each on `q0`, `q1`,
    /// or the pair. False, with nothing applied, when `gate` has no 4x4 matrix
    /// or the fold would absorb fewer pending maps and channels than the
    /// register width calls for; one-qubit channels count once per qubit.
    pub(crate) fn apply_2q(
        &mut self,
        dm: &mut DensityMatrixBackend,
        gate: &Gate,
        q0: usize,
        q1: usize,
        channels: &[PairChannel],
    ) -> bool {
        let Some(g) = matrix_2q(gate) else {
            return false;
        };
        let on = |qubit: usize| {
            channels
                .iter()
                .any(|c| matches!(c, PairChannel::Block(q, _) if *q == qubit))
        };
        let maps = usize::from(self.pending[q0].is_some())
            + usize::from(self.pending[q1].is_some())
            + usize::from(on(q0))
            + usize::from(on(q1))
            + channels
                .iter()
                .filter(|c| matches!(c, PairChannel::Pair(_)))
                .count();
        if maps < self.min_maps {
            return false;
        }

        let before = kron_block_superoperators(
            self.pending[q0].as_ref().unwrap_or(&IDENTITY_4X4),
            self.pending[q1].as_ref().unwrap_or(&IDENTITY_4X4),
        );
        let gate_map: [[Complex64; 16]; 16] = std::array::from_fn(|row| {
            std::array::from_fn(|col| g[row >> 2][col >> 2] * g[row & 3][col & 3].conj())
        });
        let mut s = mat_mul_16x16(&gate_map, &before);
        let (mut after0, mut after1) = (None, None);
        let close = |s: &mut [[Complex64; 16]; 16], a0: &mut Option<_>, a1: &mut Option<_>| {
            if a0.is_some() || a1.is_some() {
                let after = kron_block_superoperators(
                    a0.as_ref().unwrap_or(&IDENTITY_4X4),
                    a1.as_ref().unwrap_or(&IDENTITY_4X4),
                );
                *s = mat_mul_16x16(&after, s);
                (*a0, *a1) = (None, None);
            }
        };
        for channel in channels {
            match channel {
                PairChannel::Block(qubit, b) => {
                    let target = if *qubit == q0 {
                        &mut after0
                    } else {
                        &mut after1
                    };
                    *target = Some(match target {
                        Some(prior) => mat_mul_4x4(b, prior),
                        None => **b,
                    });
                }
                PairChannel::Pair(p) => {
                    close(&mut s, &mut after0, &mut after1);
                    s = mat_mul_16x16(p, &s);
                }
            }
        }
        close(&mut s, &mut after0, &mut after1);

        self.pending[q0] = None;
        self.pending[q1] = None;
        dm.apply_2q_superoperator(q0, q1, &s);
        true
    }

    /// Sweep every pending map into the buffer.
    pub(crate) fn flush(&mut self, dm: &mut DensityMatrixBackend) {
        for (qubit, pending) in self.pending.iter_mut().enumerate() {
            if let Some(s) = pending.take() {
                dm.apply_block_superoperator_infallible(qubit, &s);
            }
        }
    }

    /// Sweep the pending maps of `qubits` into the buffer.
    pub(crate) fn flush_qubits(&mut self, dm: &mut DensityMatrixBackend, qubits: &[usize]) {
        for &qubit in qubits {
            if let Some(s) = self.pending[qubit].take() {
                dm.apply_block_superoperator_infallible(qubit, &s);
            }
        }
    }
}

/// `conj(gate)` as a native gate variant, which is what the bra register of
/// `rho -> U rho U^dagger` needs. `Cx`, `Cz`, and `Swap` are real, so they are
/// their own conjugate. `None` means the variant has no native conjugate form
/// and the caller falls back to conjugating the buffer around the gate.
fn conjugate_gate(gate: &Gate) -> Option<Gate> {
    match gate {
        Gate::Cx | Gate::Cz | Gate::Swap => Some(gate.clone()),
        Gate::Rzz(theta) => Some(Gate::Rzz(-*theta)),
        Gate::Fused2q(mat) => Some(Gate::Fused2q(Box::new(conjugate_4x4(mat)))),
        Gate::Cu(mat) => Some(Gate::Cu(Box::new(conjugate_2x2(mat)))),
        Gate::Mcu(data) => Some(Gate::Mcu(Box::new(McuData {
            mat: conjugate_2x2(&data.mat),
            num_controls: data.num_controls,
        }))),
        Gate::Unitary(data) => Some(Gate::Unitary(Box::new(data.conjugated()))),
        _ => None,
    }
}

/// The batch payload as a flat [`DiagEntry`] list, or `None` for a gate that is
/// not a diagonal batch. `BatchPhase` keeps its shared control in the
/// instruction targets rather than in the payload, so it arrives separately.
fn diagonal_batch_entries(gate: &Gate, targets: &[usize]) -> Option<Vec<DiagEntry>> {
    match gate {
        Gate::BatchPhase(data) => Some(
            data.phases
                .iter()
                .map(|&(target, phase)| DiagEntry::Phase2q {
                    q0: targets[0],
                    q1: target,
                    phase,
                })
                .collect(),
        ),
        Gate::BatchRzz(data) => Some(
            data.edges
                .iter()
                .map(|&(q0, q1, theta)| DiagEntry::Parity2q {
                    q0,
                    q1,
                    same: Complex64::from_polar(1.0, -theta / 2.0),
                    diff: Complex64::from_polar(1.0, theta / 2.0),
                })
                .collect(),
        ),
        Gate::DiagonalBatch(data) => Some(data.entries.clone()),
        _ => None,
    }
}

/// `gates` cut into the runs [`kernels::multi_2q_ordered_prefix`] allows, in order.
fn ordered_runs(
    gates: &[(usize, usize, [[Complex64; 4]; 4])],
    num_qubits: usize,
) -> impl Iterator<Item = &[(usize, usize, [[Complex64; 4]; 4])]> {
    let mut rest = gates;
    std::iter::from_fn(move || {
        if rest.is_empty() {
            return None;
        }
        let (run, tail) = rest.split_at(kernels::multi_2q_ordered_prefix(rest, num_qubits));
        rest = tail;
        Some(run)
    })
}

/// The gate with every qubit index stored inside its payload shifted onto the
/// ket register. Offsetting the instruction targets is not enough for
/// `QftBlock`, whose whole range lives in the variant. Borrows the gate when it
/// carries no such index.
///
/// `MultiFused` and `Multi2q` are absent: [`DensityMatrixBackend::apply_unitary`]
/// shifts their constituents itself. The diagonal batches are absent because
/// [`DensityMatrixBackend::apply_diagonal_sandwich`] takes them before the
/// two-product route is reached.
fn ket_register_gate(gate: &Gate, n: usize) -> Cow<'_, Gate> {
    match gate {
        Gate::QftBlock {
            start,
            num,
            inverse,
            swaps,
            big_endian,
        } => Cow::Owned(Gate::QftBlock {
            start: start + n as u8,
            num: *num,
            inverse: *inverse,
            swaps: *swaps,
            big_endian: *big_endian,
        }),
        _ => Cow::Borrowed(gate),
    }
}

/// Reject a mixture the device cannot hold before anything is allocated: the
/// `4^n` buffer at 16 bytes per entry plus the `2^n` diagonal scratch. A
/// context that cannot report free memory falls through to the allocation,
/// whose own error stays authoritative.
#[cfg(feature = "gpu")]
fn check_device_budget(context: &GpuContext, num_qubits: usize) -> Result<()> {
    if 2 * num_qubits >= usize::BITS as usize - 4 {
        return Ok(());
    }
    let Ok(free) = context.vram_available() else {
        return Ok(());
    };
    let needed = (1usize << (2 * num_qubits)) * 16 + (1usize << num_qubits) * 8;
    if needed > free {
        return Err(crate::error::PrismError::ResourceLimit(Box::new(
            crate::error::ResourceLimit {
                backend: "density_matrix-gpu".to_string(),
                operation: format!("the 4^n mixture of a {num_qubits}-qubit circuit on the device"),
                resource: crate::error::ResourceKind::DeviceBytes,
                required: needed as u128,
                limit: free as u128,
                env_var: None,
            },
        )));
    }
    Ok(())
}

/// The channel entry points are infallible by signature, so a device launch
/// that fails under one has nowhere to report and stops the run instead.
#[cfg(feature = "gpu")]
fn launched<T>(result: Result<T>) -> T {
    result.unwrap_or_else(|e| panic!("density-matrix device kernel failed: {e}"))
}

/// Exact density-matrix simulator. See the module docs for the state layout.
pub struct DensityMatrixBackend {
    num_qubits: usize,
    classical_bits: Vec<bool>,
    rng: ChaCha8Rng,
    sv: StatevectorBackend,
    /// The doubled constituent lists of a batched `MultiFused` and `Multi2q`,
    /// kept between calls.
    multi_1q_scratch: Vec<(usize, [[Complex64; 2]; 2])>,
    multi_2q_scratch: Vec<(usize, usize, [[Complex64; 4]; 4])>,
    #[cfg(feature = "gpu")]
    gpu_context: Option<Arc<GpuContext>>,
}

impl DensityMatrixBackend {
    pub fn new(seed: u64) -> Self {
        Self {
            num_qubits: 0,
            classical_bits: Vec::new(),
            rng: ChaCha8Rng::seed_from_u64(seed),
            sv: StatevectorBackend::new(seed),
            multi_1q_scratch: Vec::new(),
            multi_2q_scratch: Vec::new(),
            #[cfg(feature = "gpu")]
            gpu_context: None,
        }
    }

    /// Hold the mixture on the device bound to `context`.
    ///
    /// A device that cannot hold the `4^n` buffer is an error at
    /// [`Backend::init`], never a host fallback. A kernel launch that fails
    /// under one of the infallible channel entry points panics.
    #[cfg(feature = "gpu")]
    pub fn with_gpu(mut self, context: Arc<GpuContext>) -> Self {
        self.gpu_context = Some(context.clone());
        self.sv = self.sv.with_gpu(context);
        self
    }

    #[cfg(feature = "gpu")]
    fn device(&mut self) -> Option<(Arc<GpuContext>, &mut GpuState)> {
        let gpu = self.sv.gpu_state_mut()?;
        let ctx = gpu.context().clone();
        Some((ctx, gpu))
    }

    /// The full `4^n` buffer, row-major as the module docs lay it out. Reads
    /// back from the device when the mixture is resident there.
    pub fn density_matrix(&self) -> Result<Vec<Complex64>> {
        self.sv.export_statevector()
    }

    /// Purity `Tr(rho^2)`, equal to `1` for a pure state and less otherwise.
    pub fn purity(&self) -> f64 {
        #[cfg(feature = "gpu")]
        if let Some(gpu) = self.sv.gpu_state() {
            return launched(dk::norm_sqr(gpu.context(), gpu, self.num_qubits));
        }
        crate::backend::state_norm_sqr(&self.sv.state)
    }

    /// The host cap, or with a device attached the VRAM budget in its place.
    fn check_capacity(&self, num_qubits: usize) -> Result<()> {
        #[cfg(feature = "gpu")]
        if let Some(ctx) = &self.gpu_context {
            return check_device_budget(ctx, num_qubits);
        }
        crate::backend::check_state_allocation(
            "density_matrix",
            num_qubits,
            crate::backend::max_density_matrix_qubits(),
            crate::backend::DM_QUBIT_CAP_ENV,
        )
    }

    #[inline]
    fn dim(&self) -> usize {
        1usize << self.num_qubits
    }

    /// `rho_A[t * dim + t'] = sum_e rho[idx(t, e)][idx(t', e)]` over the
    /// `2^(n - k)` traced indices `e`, `buffer` being the `4^n` mixture in the
    /// module's layout. Once the buffer is the size of a statevector past the
    /// parallel threshold the work is split either as stripes of result rows,
    /// which repeat only the `k` bit insertions of each traced index, or as a
    /// fold over the traced index whose jobs each zero and merge a
    /// `dim * dim` accumulator; the stripe wins once the insertions cost less
    /// than that.
    fn partial_trace(&self, buffer: &[Complex64], subsystem: &[usize]) -> Vec<Complex64> {
        let d = self.dim();
        let k = subsystem.len();
        let dim = 1usize << k;
        let offsets = reduced_density::row_offsets(subsystem);
        let mut ascending = subsystem.to_vec();
        ascending.sort_unstable();
        let groups = d >> k;

        let accumulate_rows = |rows: &mut [Complex64], first: usize, base: usize| {
            for (r, out) in rows.chunks_exact_mut(dim).enumerate() {
                let row = (base | offsets[first + r]) * d + base;
                for (entry, &off) in out.iter_mut().zip(&offsets) {
                    *entry += buffer[row | off];
                }
            }
        };
        let zero = Complex64::new(0.0, 0.0);
        #[cfg(feature = "parallel")]
        if self.sv.num_qubits >= crate::backend::PARALLEL_THRESHOLD_QUBITS {
            use rayon::prelude::*;
            if groups * k < 2 * dim * dim {
                let rows = reduced_density::stripe_rows(dim);
                let mut rho = vec![zero; dim * dim];
                rho.par_chunks_mut(rows * dim)
                    .enumerate()
                    .for_each(|(stripe, out)| {
                        for e in 0..groups {
                            let base = reduced_density::traced_base(e, &ascending);
                            accumulate_rows(out, stripe * rows, base);
                        }
                    });
                return rho;
            }
            let fresh = || vec![zero; dim * dim];
            let add = |mut a: Vec<Complex64>, b: Vec<Complex64>| {
                for (x, y) in a.iter_mut().zip(&b) {
                    *x += y;
                }
                a
            };
            return (0..groups)
                .into_par_iter()
                .with_min_len(reduced_density::fold_min_len(groups, dim * dim))
                .fold(fresh, |mut acc, e| {
                    accumulate_rows(&mut acc, 0, reduced_density::traced_base(e, &ascending));
                    acc
                })
                .reduce(fresh, add);
        }
        let mut rho = vec![zero; dim * dim];
        for e in 0..groups {
            accumulate_rows(&mut rho, 0, reduced_density::traced_base(e, &ascending));
        }
        rho
    }

    fn conjugate_buffer(&mut self) -> Result<()> {
        #[cfg(feature = "gpu")]
        {
            let n = self.num_qubits;
            if let Some((ctx, gpu)) = self.device() {
                return dk::conjugate(&ctx, gpu, n);
            }
        }
        #[cfg(feature = "parallel")]
        {
            use rayon::prelude::*;
            if self.sv.state.len() >= (1 << crate::backend::PARALLEL_THRESHOLD_QUBITS) {
                self.sv
                    .state
                    .par_iter_mut()
                    .for_each(|amp| *amp = amp.conj());
                return Ok(());
            }
        }
        for amp in self.sv.state.iter_mut() {
            *amp = amp.conj();
        }
        Ok(())
    }

    /// A two-qubit matrix on the embedded buffer, on whichever side holds it.
    fn fused_2q(&mut self, q0: usize, q1: usize, mat: &[[Complex64; 4]; 4]) -> Result<()> {
        #[cfg(feature = "gpu")]
        if let Some((ctx, gpu)) = self.device() {
            return crate::gpu::kernels::dense::launch_apply_fused_2q(&ctx, gpu, q0, q1, mat);
        }
        self.sv.apply_fused_2q(q0, q1, mat);
        Ok(())
    }

    /// Apply a compiled one-qubit block superoperator in a single buffer pass.
    ///
    /// The `(row-bit, col-bit)` block of `qubit` is the two-qubit subspace
    /// `(qubit + n, qubit)` of the embedded `2n`-qubit statevector, so the
    /// statevector two-qubit kernel applies `S` directly. That kernel indexes
    /// its matrix as `2 * bit(q0) + bit(q1)`, matching the block index `2a + b`
    /// once the row bit is passed as `q0`.
    fn apply_block_superoperator(&mut self, qubit: usize, s: &[[Complex64; 4]; 4]) -> Result<()> {
        let n = self.num_qubits;
        self.fused_2q(qubit + n, qubit, s)
    }

    /// Evolve `rho -> U rho U^dagger` for the unitary `gate` on `targets`.
    ///
    /// `U` applies to the ket register (targets and payload indices offset by
    /// `n`, see [`ket_register_gate`]) for the left product `U rho`, and
    /// `conj(U)` to the bra register (indices unchanged) for the right product
    /// `rho U^dagger`. Variants with no native conjugate form fall back to
    /// conjugating the whole buffer around the gate, which costs two extra
    /// passes.
    ///
    /// One-qubit gates take [`DensityMatrixBackend::apply_1q_sandwich`], `Rzz`
    /// takes [`DensityMatrixBackend::apply_rzz_sandwich`], and the diagonal
    /// batches take [`DensityMatrixBackend::apply_diagonal_sandwich`]. All
    /// three fold both products into one pass rather than taking a conjugate
    /// form here.
    ///
    /// `Multi2q` takes [`DensityMatrixBackend::apply_multi_2q_sandwich`] on the
    /// host and one constituent at a time on the device. `MultiFused` entries
    /// are one per qubit and commute, so its list is order independent: past
    /// the parallel threshold on the host it takes
    /// [`DensityMatrixBackend::apply_multi_1q_sandwich`], and elsewhere the
    /// per-constituent one-qubit sandwich.
    fn apply_unitary(&mut self, gate: &Gate, targets: &[usize]) -> Result<()> {
        let n = self.num_qubits;

        if let Some(mat) = matrix_1q(gate) {
            return self.apply_1q_sandwich(targets[0], &mat);
        }

        match gate {
            Gate::MultiFused(data) => {
                if !self.sv.is_gpu_resident() && 2 * n >= crate::backend::PARALLEL_THRESHOLD_QUBITS
                {
                    return self.apply_multi_1q_sandwich(&data.gates, data.all_diagonal);
                }
                for (qubit, mat) in data.gates.iter() {
                    self.apply_1q_sandwich(*qubit, mat)?;
                }
                return Ok(());
            }
            Gate::Rzz(theta) => {
                return self.apply_rzz_sandwich(targets[0], targets[1], *theta);
            }
            Gate::Multi2q(data) => {
                if !self.sv.is_gpu_resident() {
                    self.apply_multi_2q_sandwich(&data.gates);
                    return Ok(());
                }
                for &(q0, q1, ref mat) in data.gates.iter() {
                    self.fused_2q(q0 + n, q1 + n, mat)?;
                }
                for &(q0, q1, ref mat) in data.gates.iter() {
                    self.fused_2q(q0, q1, &conjugate_4x4(mat))?;
                }
                return Ok(());
            }
            _ => {}
        }

        if let Some(entries) = diagonal_batch_entries(gate, targets) {
            return self.apply_diagonal_sandwich(&entries);
        }

        let ket_targets: SmallVec<[usize; 4]> = targets.iter().map(|&t| t + n).collect();
        self.sv.apply(&Instruction::Gate {
            gate: ket_register_gate(gate, n).into_owned(),
            targets: ket_targets,
        })?;

        let bra_targets: SmallVec<[usize; 4]> = targets.iter().copied().collect();
        if let Some(conjugate) = conjugate_gate(gate) {
            return self.sv.apply(&Instruction::Gate {
                gate: conjugate,
                targets: bra_targets,
            });
        }

        self.conjugate_buffer()?;
        self.sv.apply(&Instruction::Gate {
            gate: gate.clone(),
            targets: bra_targets,
        })?;
        self.conjugate_buffer()
    }

    /// Evolve `rho -> U rho U^dagger` for a one-qubit `U`.
    ///
    /// Above the parallel threshold, and always on the device, the two products
    /// compile to one block superoperator and sweep the buffer once; below it
    /// the superoperator's dense 4x4 block loses to two cheap register passes.
    fn apply_1q_sandwich(&mut self, qubit: usize, matrix: &[[Complex64; 2]; 2]) -> Result<()> {
        let n = self.num_qubits;
        if self.sv.is_gpu_resident() || 2 * n >= crate::backend::PARALLEL_THRESHOLD_QUBITS {
            let s = block_superoperator(&[*matrix]);
            return self.apply_block_superoperator(qubit, &s);
        }
        self.sv.apply_1q_matrix(qubit + n, matrix)?;
        self.sv.apply_1q_matrix(qubit, &conjugate_2x2(matrix))
    }

    /// Evolve `rho -> U rho U^dagger` for a `Multi2q` batch `U`.
    ///
    /// The statevector kernel keeps list order only for a batch that fits one
    /// subcube or the in-place tile; past that its tiered pass reorders gates
    /// across tiers, and the `+n` shift onto the ket register pushes a list
    /// fusion built to fit past it. The ket constituents and their bra
    /// conjugates act on disjoint registers and commute, so the doubled list
    /// is cut into the longest runs [`kernels::multi_2q_ordered_prefix`]
    /// allows, each one tiled pass. A split that needs no fewer passes than one
    /// per ket constituent plus the bra half keeps that route instead.
    fn apply_multi_2q_sandwich(&mut self, gates: &[(usize, usize, [[Complex64; 4]; 4])]) {
        let n = self.num_qubits;
        self.multi_2q_scratch.clear();
        self.multi_2q_scratch
            .extend(gates.iter().map(|&(q0, q1, mat)| (q0 + n, q1 + n, mat)));
        self.multi_2q_scratch.extend(
            gates
                .iter()
                .map(|&(q0, q1, ref mat)| (q0, q1, conjugate_4x4(mat))),
        );
        let num_qubits = self.sv.num_qubits;
        let (ket, bra) = self.multi_2q_scratch.split_at(gates.len());
        let bra_ordered = kernels::multi_2q_ordered_prefix(bra, num_qubits) == bra.len();
        let per_constituent = ket.len() + if bra_ordered { 1 } else { bra.len() };
        if ordered_runs(&self.multi_2q_scratch, num_qubits).count() < per_constituent {
            for run in ordered_runs(&self.multi_2q_scratch, num_qubits) {
                self.sv.apply_multi_2q(run);
            }
            return;
        }
        for &(q0, q1, ref mat) in ket {
            self.sv.apply_fused_2q(q0, q1, mat);
        }
        if bra_ordered {
            self.sv.apply_multi_2q(bra);
        } else {
            for &(q0, q1, ref mat) in bra {
                self.sv.apply_fused_2q(q0, q1, mat);
            }
        }
    }

    /// Evolve `rho -> U rho U^dagger` for a `MultiFused` batch `U`.
    ///
    /// The constituents on the ket register and their conjugates on the bra
    /// register are `2k` one-qubit gates on distinct qubits, so the doubled list
    /// can run through the statevector `MultiFused` kernel as one batch.
    /// It does when [`kernels::multi_1q_sweeps`] counts fewer passes than the
    /// `k` one-qubit sandwiches; otherwise the sandwiches run. An all-diagonal
    /// batch of two or more takes [`DensityMatrixBackend::apply_diagonal_sandwich`],
    /// one sweep.
    fn apply_multi_1q_sandwich(
        &mut self,
        gates: &[(usize, [[Complex64; 2]; 2])],
        all_diagonal: bool,
    ) -> Result<()> {
        if all_diagonal && gates.len() > 1 {
            let entries: Vec<DiagEntry> = gates
                .iter()
                .map(|&(qubit, m)| DiagEntry::Phase1q {
                    qubit,
                    d0: m[0][0],
                    d1: m[1][1],
                })
                .collect();
            return self.apply_diagonal_sandwich(&entries);
        }
        if !all_diagonal {
            let n = self.num_qubits;
            self.multi_1q_scratch.clear();
            for &(qubit, ref mat) in gates {
                self.multi_1q_scratch.push((qubit + n, *mat));
                self.multi_1q_scratch.push((qubit, conjugate_2x2(mat)));
            }
            if kernels::multi_1q_sweeps(&self.multi_1q_scratch, self.sv.num_qubits) < gates.len() {
                self.sv.apply_multi_1q(&self.multi_1q_scratch);
                return Ok(());
            }
        }
        for (qubit, mat) in gates {
            self.apply_1q_sandwich(*qubit, mat)?;
        }
        Ok(())
    }

    /// Evolve `rho -> R rho R^dagger` for `R = Rzz(theta)` in one buffer pass.
    ///
    /// Both factors are diagonal, so the entry at `(r, c)` picks up
    /// `phase(p_r) * conj(phase(p_c))` where `p` is the parity of the target
    /// pair in that register and `phase` is the statevector kernel's
    /// `[exp(-i theta/2), exp(+i theta/2)]`. Conjugate phases cancel wherever
    /// the two registers agree on parity, which is half the sixteen entry
    /// classes, and the rest carry `exp(-i theta)` or its conjugate. The
    /// generic route's two register passes therefore collapse onto the single
    /// contiguous pass [`DensityMatrixBackend::kraus_2q_diagonal_sweep`]
    /// already runs for a diagonal channel.
    fn apply_rzz_sandwich(&mut self, q0: usize, q1: usize, theta: f64) -> Result<()> {
        let phase = [
            Complex64::from_polar(1.0, -theta / 2.0),
            Complex64::from_polar(1.0, theta / 2.0),
        ];
        let mut diag = [Complex64::new(1.0, 0.0); 16];
        for tr in 0..4usize {
            for tc in 0..4usize {
                let pr = ((tr >> 1) ^ tr) & 1;
                let pc = ((tc >> 1) ^ tc) & 1;
                diag[4 * tr + tc] = phase[pr] * phase[pc].conj();
            }
        }
        self.kraus_2q_diagonal_sweep(&diag, q0, q1)
    }

    /// Evolve `rho -> D rho D^dagger` for a diagonal batch `D` in one buffer
    /// pass.
    ///
    /// A diagonal `D` scales `<r|rho|c>` by `f(r) * conj(f(c))`, where `f` is
    /// the combined phase the payload puts on a basis index. The embedded
    /// layout splits the buffer index into exactly those two operands, `r` in
    /// the high `n` bits and `c` in the low `n`, so one table of `2^n` phases
    /// serves both registers and the sweep reads it twice per amplitude. The
    /// table costs `2^n` against the `4^n` buffer it saves three passes on: the
    /// generic route has no conjugate form for these variants and pays two
    /// register passes plus two buffer conjugations. Entries are diagonal, so
    /// their order does not matter and the tier reordering cannot arise.
    fn apply_diagonal_sandwich(&mut self, entries: &[DiagEntry]) -> Result<()> {
        let n = self.num_qubits;
        let d = self.dim();
        let ket: Vec<Complex64> = (0..d).map(|r| diag_entries_phase(r, entries)).collect();

        #[cfg(feature = "gpu")]
        if let Some((ctx, gpu)) = self.device() {
            return dk::diagonal_sandwich(&ctx, gpu, n, &ket);
        }

        let bra: Vec<Complex64> = ket.iter().map(|f| f.conj()).collect();

        #[cfg(feature = "parallel")]
        if 2 * n >= crate::backend::PARALLEL_THRESHOLD_QUBITS {
            use crate::backend::MIN_PAR_ELEMS;
            use rayon::prelude::*;

            self.sv
                .state
                .par_iter_mut()
                .enumerate()
                .with_min_len(MIN_PAR_ELEMS)
                .for_each(|(idx, amp)| {
                    *amp *= ket[idx >> n] * bra[idx & (d - 1)];
                });
            return Ok(());
        }

        for (idx, amp) in self.sv.state.iter_mut().enumerate() {
            *amp *= ket[idx >> n] * bra[idx & (d - 1)];
        }
        Ok(())
    }

    /// `P(qubit = |1>) = Tr(P_1 rho)`, the sum of the diagonal `rho` entries
    /// whose row index has `qubit` set.
    fn prob_one(&self, qubit: usize) -> Result<f64> {
        let d = self.dim();
        let bit = 1usize << qubit;
        #[cfg(feature = "gpu")]
        if let Some(gpu) = self.sv.gpu_state() {
            let diag = dk::diagonal(gpu.context(), gpu, self.num_qubits)?;
            let p1: f64 = diag
                .iter()
                .enumerate()
                .filter(|(r, _)| r & bit != 0)
                .map(|(_, p)| p)
                .sum();
            return Ok(p1.clamp(0.0, 1.0));
        }
        let mut p1 = 0.0;
        for r in 0..d {
            if r & bit != 0 {
                p1 += self.sv.state[r * d + r].re;
            }
        }
        Ok(p1.clamp(0.0, 1.0))
    }

    /// Sample and project qubit `qubit`, recording the outcome in
    /// `classical_bit`. Collapses `rho -> P_m rho P_m / p_m`, which in the
    /// embedded layout zeroes every entry whose row or column disagrees with
    /// the outcome on `qubit` and rescales the survivors to unit trace.
    fn apply_measure(&mut self, qubit: usize, classical_bit: usize) -> Result<()> {
        let p1 = self.prob_one(qubit)?;
        let u: f64 = self.rng.random();
        let outcome = u < p1;
        self.classical_bits[classical_bit] = outcome;
        let p = if outcome { p1 } else { 1.0 - p1 };
        self.project(qubit, outcome, p)
    }

    /// Deterministic reset `rho -> |0><0| (x) tr_q rho`: fold the block with
    /// `qubit` set on both row and column into the block with it clear on both,
    /// then zero the three sibling entries that still touch `qubit`. The four
    /// entry classes are contiguous runs of the buffer, so the pass walks
    /// `2 * rmask` blocks rather than scattering per block base.
    fn apply_reset(&mut self, qubit: usize) -> Result<()> {
        let n = self.num_qubits;
        #[cfg(feature = "gpu")]
        if let Some((ctx, gpu)) = self.device() {
            return dk::reset(&ctx, gpu, n, qubit);
        }
        let rmask = 1usize << (qubit + n);
        let cmask = 1usize << qubit;
        let block_size = rmask << 1;

        #[cfg(feature = "parallel")]
        if 2 * n >= crate::backend::PARALLEL_THRESHOLD_QUBITS {
            use crate::backend::chunk_min_len;
            use rayon::prelude::*;

            if self.sv.state.len() / block_size >= 4 {
                self.sv
                    .state
                    .par_chunks_mut(block_size)
                    .with_min_len(chunk_min_len(block_size))
                    .for_each(|block| {
                        let (r0, r1) = block.split_at_mut(rmask);
                        reset_fold_pair(r0, r1, cmask);
                    });
                return Ok(());
            }

            let tile = row_aligned_tile(cmask, rmask);
            for block in self.sv.state.chunks_mut(block_size) {
                let (r0, r1) = block.split_at_mut(rmask);
                r0.par_chunks_mut(tile)
                    .zip(r1.par_chunks_mut(tile))
                    .for_each(|(t0, t1)| reset_fold_pair(t0, t1, cmask));
            }
            return Ok(());
        }

        for block in self.sv.state.chunks_mut(block_size) {
            let (r0, r1) = block.split_at_mut(rmask);
            reset_fold_pair(r0, r1, cmask);
        }
        Ok(())
    }

    /// Project `rho` onto the `outcome` subspace of `qubit` with outcome
    /// probability `p`, renormalizing the survivors to unit trace.
    fn project(&mut self, qubit: usize, outcome: bool, p: f64) -> Result<()> {
        let n = self.num_qubits;
        let rmask = 1usize << (qubit + n);
        let cmask = 1usize << qubit;
        let scale = Complex64::new(1.0 / p.clamp(NORM_CLAMP_MIN, 1.0), 0.0);

        #[cfg(feature = "gpu")]
        if let Some((ctx, gpu)) = self.device() {
            return dk::project(&ctx, gpu, n, qubit, outcome, scale.re);
        }

        #[cfg(feature = "parallel")]
        if 2 * n >= crate::backend::PARALLEL_THRESHOLD_QUBITS {
            use crate::backend::chunk_min_len;
            use rayon::prelude::*;

            let tile = row_aligned_tile(cmask, rmask);
            self.sv
                .state
                .par_chunks_mut(tile)
                .with_min_len(chunk_min_len(tile))
                .enumerate()
                .for_each(|(t, chunk)| {
                    project_tile(chunk, t * tile, rmask, cmask, outcome, scale);
                });
            return Ok(());
        }

        for (t, chunk) in self.sv.state.chunks_mut(rmask).enumerate() {
            project_tile(chunk, t * rmask, rmask, cmask, outcome, scale);
        }
        Ok(())
    }

    fn apply_conditional(
        &mut self,
        condition: &ClassicalCondition,
        gate: &Gate,
        targets: &[usize],
    ) -> Result<()> {
        if condition.evaluate(&self.classical_bits) {
            self.apply_unitary(gate, targets)?;
        }
        Ok(())
    }

    /// Apply `rho -> sum_k K_k rho K_k^dagger` on `qubit` in one buffer sweep,
    /// through a 4x4 superoperator compiled from the Kraus set and acting on
    /// each `(row-bit, col-bit)` block of `rho`.
    pub fn apply_1q_kraus(&mut self, qubit: usize, kraus: &[[[Complex64; 2]; 2]]) {
        let s = block_superoperator(kraus);
        self.apply_block_superoperator_infallible(qubit, &s);
    }

    /// [`DensityMatrixBackend::apply_block_superoperator`] for the entry points
    /// that carry no error channel: the host path cannot fail and a device
    /// launch failure stops the run.
    fn apply_block_superoperator_infallible(&mut self, qubit: usize, s: &[[Complex64; 4]; 4]) {
        let result = self.apply_block_superoperator(qubit, s);
        #[cfg(feature = "gpu")]
        launched(result);
        #[cfg(not(feature = "gpu"))]
        result.expect("host block superoperator is infallible");
    }

    /// Apply symmetric two-qubit depolarizing on `(q0, q1)`:
    /// `rho -> (1-p) rho + (p/15) sum_{P != I(x)I} P rho P`, summed over the 15
    /// non-identity two-qubit Paulis.
    ///
    /// The Pauli twirl `sum_P P B P = 4 Tr(B) I4` over all 16 two-qubit Paulis
    /// holds for any 4x4 `B`, so on each `(q0, q1)` block the map collapses to
    /// `B -> alpha B + beta Tr(B) I4` with `alpha = 1 - 16p/15` and
    /// `beta = 4p/15`. That is one real scale per amplitude against the 16
    /// complex multiply-accumulates the equivalent 16x16 superoperator pays in
    /// [`DensityMatrixBackend::apply_2q_kraus`].
    ///
    /// # Panics
    ///
    /// If `q0` and `q1` are equal or outside the register, as
    /// [`DensityMatrixBackend::apply_2q_kraus`] does.
    pub fn apply_2q_depolarizing(&mut self, q0: usize, q1: usize, p: f64) {
        let alpha = 1.0 - 16.0 * p / 15.0;
        let beta = 4.0 * p / 15.0;
        let (positions, flats, num_groups) = self.block_layout(q0, q1);

        #[cfg(feature = "gpu")]
        {
            let n = self.num_qubits;
            if let Some((ctx, gpu)) = self.device() {
                launched(dk::depolarizing_2q(&ctx, gpu, n, q0, q1, alpha, beta));
                return;
            }
        }

        #[cfg(feature = "parallel")]
        if 2 * self.num_qubits >= crate::backend::PARALLEL_THRESHOLD_QUBITS {
            use crate::backend::MIN_PAR_ITERS;
            use crate::backend::statevector::SendPtr;
            use rayon::prelude::*;

            let ptr = SendPtr(self.sv.state.as_mut_ptr());
            (0..num_groups)
                .into_par_iter()
                .with_min_len(MIN_PAR_ITERS)
                .for_each(move |m| {
                    let base = block_base(m, &positions);
                    // SAFETY: inserting a zero bit at each of the four block
                    // positions is a bijection from `m` onto the block bases, so
                    // the 16 offsets one iteration touches are disjoint from
                    // every other iteration's. The safe alternative needs the
                    // 16 entries as disjoint sub-slices, which they are not:
                    // they are strided across four independent bit positions.
                    unsafe {
                        let mut trace = Complex64::new(0.0, 0.0);
                        for t in 0..4 {
                            trace += ptr.load(base | flats[5 * t]);
                        }
                        let shift = trace * beta;
                        for (i, &off) in flats.iter().enumerate() {
                            let mut out = ptr.load(base | off) * alpha;
                            if i % 5 == 0 {
                                out += shift;
                            }
                            ptr.store(base | off, out);
                        }
                    }
                });
            return;
        }

        for m in 0..num_groups {
            let base = block_base(m, &positions);
            let mut trace = Complex64::new(0.0, 0.0);
            for t in 0..4 {
                trace += self.sv.state[base | flats[5 * t]];
            }
            let shift = trace * beta;
            for (i, &off) in flats.iter().enumerate() {
                let mut out = self.sv.state[base | off] * alpha;
                if i % 5 == 0 {
                    out += shift;
                }
                self.sv.state[base | off] = out;
            }
        }
    }

    /// Block traversal for a two-qubit map on `(q0, q1)`: the four buffer bit
    /// positions the block spans in ascending order, the 16 offsets from a
    /// block base indexed `4 * tr + tc`, and the number of blocks. Block index
    /// `5 * t` is the diagonal entry `tr == tc == t`.
    ///
    /// The four positions must be distinct and inside the register, or the
    /// zero-bit insertion in [`block_base`] stops being a bijection onto the
    /// block bases and the sweep reads and writes the wrong entries. Both channel
    /// entry points are public and route through here, so this is their choke
    /// point. It is not the sweep's:
    /// [`DensityMatrixBackend::apply_rzz_sandwich`] builds its table and calls
    /// the sweep directly, where equal targets yield the identity table that the
    /// two register passes it replaced also produced.
    fn block_layout(&self, q0: usize, q1: usize) -> ([usize; 4], [usize; 16], usize) {
        let n = self.num_qubits;
        assert!(
            q0 != q1 && q0 < n && q1 < n,
            "two-qubit channel on ({q0}, {q1}) needs distinct targets inside the {n}-qubit register"
        );
        let mut positions = [q0, q1, q0 + n, q1 + n];
        positions.sort_unstable();

        let mut flats = [0usize; 16];
        for tr in 0..4 {
            for tc in 0..4 {
                flats[4 * tr + tc] = (if tr & 2 != 0 { 1usize << (q0 + n) } else { 0 })
                    | (if tr & 1 != 0 { 1usize << (q1 + n) } else { 0 })
                    | (if tc & 2 != 0 { 1usize << q0 } else { 0 })
                    | (if tc & 1 != 0 { 1usize << q1 } else { 0 });
            }
        }

        let d = self.dim();
        (positions, flats, (d * d) >> 4)
    }

    /// Apply a general two-qubit channel `rho -> sum_k K_k rho K_k^dagger` on
    /// `(q0, q1)`. Each operator is indexed `K[t][t']` with
    /// `t = 2 * bit(q0) + bit(q1)`, matching [`Gate::matrix_4x4`].
    ///
    /// The Kraus set is compiled once into a 16x16 block superoperator over the
    /// four-bit `(q0-row, q1-row, q0-col, q1-col)` block, so the buffer is swept
    /// once. Block index `4*tr + tc` orders the two-qubit row value `tr` against
    /// the column value `tc`. An all-diagonal set compiles to a diagonal
    /// superoperator and is applied as one complex multiply per amplitude in a
    /// contiguous pass instead of the dense sweep. On hosts with detected AVX2
    /// and FMA the dense sweep issues each row inner product two complex terms
    /// per FMA through `simd::PreparedKraus2q`.
    ///
    /// # Panics
    ///
    /// If `q0` and `q1` are equal or outside the register. A model built through
    /// [`NoiseModel`](crate::NoiseModel) is rejected before it reaches here.
    pub fn apply_2q_kraus(&mut self, q0: usize, q1: usize, kraus: &[[[Complex64; 4]; 4]]) {
        self.apply_2q_superoperator(q0, q1, &kraus_2q_superoperator(kraus));
    }

    /// Sweep a compiled 16x16 block superoperator over `(q0, q1)`, indexed as
    /// [`DensityMatrixBackend::apply_2q_kraus`] compiles it.
    pub(crate) fn apply_2q_superoperator(
        &mut self,
        q0: usize,
        q1: usize,
        s: &[[Complex64; 16]; 16],
    ) {
        let (positions, flats, num_groups) = self.block_layout(q0, q1);

        // Off-diagonal entries of an all-diagonal set only ever accumulate
        // `kr * conj(0.0)`, so exact zero is the test.
        let zero = Complex64::new(0.0, 0.0);
        let diagonal = s
            .iter()
            .enumerate()
            .all(|(r, row)| row.iter().enumerate().all(|(c, &e)| r == c || e == zero));
        if diagonal {
            let diag = std::array::from_fn(|i| s[i][i]);
            let result = self.kraus_2q_diagonal_sweep(&diag, q0, q1);
            #[cfg(feature = "gpu")]
            launched(result);
            #[cfg(not(feature = "gpu"))]
            result.expect("host diagonal sweep is infallible");
            return;
        }

        #[cfg(feature = "gpu")]
        {
            let n = self.num_qubits;
            if let Some((ctx, gpu)) = self.device() {
                launched(dk::kraus_2q_dense(&ctx, gpu, n, q0, q1, s));
                return;
            }
        }

        #[cfg(target_arch = "x86_64")]
        if simd::has_avx2_fma() && simd::kraus_2q_wide_enabled() {
            // SAFETY: AVX2 and FMA checked above; AVX2 implies the AVX the
            // constructor requires.
            let prepared = unsafe { simd::PreparedKraus2q::new(s) };
            self.kraus_2q_sweep_wide(&prepared, &positions, &flats, num_groups);
            return;
        }

        match positions[0] {
            0 | 1 => self.kraus_2q_sweep::<1>(s, &positions, &flats, num_groups),
            _ => self.kraus_2q_sweep::<4>(s, &positions, &flats, num_groups),
        }
    }

    /// [`DensityMatrixBackend::kraus_2q_sweep`] with the row inner products
    /// issued through [`simd::PreparedKraus2q`], one block per iteration
    /// whatever the qubit pair: the vector axis is the superoperator's 16-wide
    /// `j` axis, which does not depend on `min(q0, q1)` the way the block-run
    /// width does.
    #[cfg(target_arch = "x86_64")]
    fn kraus_2q_sweep_wide(
        &mut self,
        prepared: &simd::PreparedKraus2q,
        positions: &[usize; 4],
        flats: &[usize; 16],
        num_groups: usize,
    ) {
        #[cfg(feature = "parallel")]
        if 2 * self.num_qubits >= crate::backend::PARALLEL_THRESHOLD_QUBITS {
            use crate::backend::MIN_PAR_ITERS;
            use crate::backend::statevector::SendPtr;
            use rayon::prelude::*;

            let ptr = SendPtr(self.sv.state.as_mut_ptr());
            let positions = *positions;
            let flats = *flats;
            (0..num_groups)
                .into_par_iter()
                .with_min_len(MIN_PAR_ITERS)
                .for_each(move |m| {
                    let base = block_base(m, &positions);
                    // SAFETY: AVX2 and FMA detected. Inserting a zero bit at
                    // each of the four block positions is a bijection from
                    // `m` onto the block bases, so every `base | flats[j]`
                    // stays under `4^n` and the 16 offsets one iteration
                    // touches are disjoint from every other iteration's.
                    unsafe { prepared.apply_block_ptr(ptr.as_f64_ptr(), base, &flats) };
                });
            return;
        }

        let ptr = self.sv.state.as_mut_ptr() as *mut f64;
        for m in 0..num_groups {
            let base = block_base(m, positions);
            // SAFETY: AVX2 and FMA detected. The block bases are disjoint and
            // in bounds as in the parallel arm, on one thread.
            unsafe { prepared.apply_block_ptr(ptr, base, flats) };
        }
    }

    /// Apply a diagonal block superoperator in one contiguous pass: the
    /// amplitude at `idx` is scaled by `diag[diag_slot(idx, q0, q1, n)]`, one
    /// complex multiply against the dense sweep's 16 multiply-accumulates.
    fn kraus_2q_diagonal_sweep(
        &mut self,
        diag: &[Complex64; 16],
        q0: usize,
        q1: usize,
    ) -> Result<()> {
        let n = self.num_qubits;

        #[cfg(feature = "gpu")]
        if let Some((ctx, gpu)) = self.device() {
            return dk::kraus_2q_diagonal(&ctx, gpu, n, q0, q1, diag);
        }

        #[cfg(feature = "parallel")]
        if 2 * n >= crate::backend::PARALLEL_THRESHOLD_QUBITS {
            use crate::backend::MIN_PAR_ELEMS;
            use rayon::prelude::*;

            let diag = *diag;
            self.sv
                .state
                .par_iter_mut()
                .enumerate()
                .with_min_len(MIN_PAR_ELEMS)
                .for_each(move |(idx, amp)| {
                    *amp *= diag[diag_slot(idx, q0, q1, n)];
                });
            return Ok(());
        }

        for (idx, amp) in self.sv.state.iter_mut().enumerate() {
            *amp *= diag[diag_slot(idx, q0, q1, n)];
        }
        Ok(())
    }

    /// Sweep the buffer applying the compiled 16x16 block superoperator `s`,
    /// `W` blocks per iteration.
    ///
    /// `W` must divide `2^positions[0]`. Blocks whose compacted indices form an
    /// aligned run of that length have consecutive bases, because
    /// [`insert_zero_bit`] preserves every bit below its position and no `flats`
    /// entry has a bit below `positions[0]`. Each of the 16 slots is then a
    /// contiguous run of `W` amplitudes, which amortizes the base and offset
    /// arithmetic over `W` blocks and gives the accumulator chains room to
    /// overlap. The arithmetic itself stays scalar: the crate builds at the
    /// x86-64 baseline and every wide kernel here takes its width from
    /// `#[target_feature]` dispatch in [`crate::backend::simd`], which this does
    /// not use.
    ///
    /// Only `W = 1` and `W = 4` are instantiated. `W = 1` is forced when
    /// `min(q0, q1) == 0`, where one block bit is bit 0 and no run exists, and
    /// it also serves `min(q0, q1) == 1`: a two-block step measured +0.2% and
    /// +1.9% against one, so it earned no arm of its own. Hosts with detected
    /// AVX2 and FMA route past this sweep to
    /// [`DensityMatrixBackend::kraus_2q_sweep_wide`] unless
    /// `PRISM_NO_AVX2_KRAUS` is set.
    fn kraus_2q_sweep<const W: usize>(
        &mut self,
        s: &[[Complex64; 16]; 16],
        positions: &[usize; 4],
        flats: &[usize; 16],
        num_groups: usize,
    ) {
        let zero = Complex64::new(0.0, 0.0);
        let steps = num_groups / W;

        #[cfg(feature = "parallel")]
        if 2 * self.num_qubits >= crate::backend::PARALLEL_THRESHOLD_QUBITS {
            use crate::backend::MIN_PAR_ITERS;
            use crate::backend::statevector::SendPtr;
            use rayon::prelude::*;

            let ptr = SendPtr(self.sv.state.as_mut_ptr());
            let s = *s;
            let positions = *positions;
            let flats = *flats;
            (0..steps)
                .into_par_iter()
                .with_min_len(MIN_PAR_ITERS.div_ceil(W))
                .for_each(move |step| {
                    let base = block_base(step * W, &positions);
                    let mut v = [[zero; W]; 16];
                    // SAFETY: inserting a zero bit at each of the four block
                    // positions is a bijection from the compacted index onto the
                    // block bases, so the `16 * W` offsets one iteration touches
                    // are disjoint from every other iteration's. The safe
                    // alternative needs them as disjoint sub-slices, which they
                    // are not: the 16 slots are strided across four independent
                    // bit positions.
                    unsafe {
                        for (slot, &off) in v.iter_mut().zip(flats.iter()) {
                            for (j, lane) in slot.iter_mut().enumerate() {
                                *lane = ptr.load((base | off) + j);
                            }
                        }
                        for (row, &off) in s.iter().zip(flats.iter()) {
                            let mut acc = [zero; W];
                            for (&coeff, slot) in row.iter().zip(v.iter()) {
                                for (a, &lane) in acc.iter_mut().zip(slot.iter()) {
                                    *a += coeff * lane;
                                }
                            }
                            for (j, &a) in acc.iter().enumerate() {
                                ptr.store((base | off) + j, a);
                            }
                        }
                    }
                });
            return;
        }

        for step in 0..steps {
            let base = block_base(step * W, positions);
            let mut v = [[zero; W]; 16];
            for (k, &off) in flats.iter().enumerate() {
                v[k].copy_from_slice(&self.sv.state[(base | off)..(base | off) + W]);
            }
            for (row, &off) in s.iter().zip(flats.iter()) {
                let mut acc = [zero; W];
                for (&coeff, slot) in row.iter().zip(v.iter()) {
                    for (a, &lane) in acc.iter_mut().zip(slot.iter()) {
                        *a += coeff * lane;
                    }
                }
                self.sv.state[(base | off)..(base | off) + W].copy_from_slice(&acc);
            }
        }
    }

    /// Exact `Tr(rho P)` for every joint Pauli reduced to `(xmask, zmask, num_y)`,
    /// where `P|j> = i^{num_y} * (-1)^{popcount(j & zmask)} * |j ^ xmask>`. The
    /// trace collapses to a diagonal-offset sum:
    /// `Tr(rho P) = i^{num_y} sum_j (-1)^{popcount(j & zmask)} rho[j][j ^ xmask]`.
    ///
    /// Each row of the `4^n` buffer is visited once and contributes one entry per
    /// observable, so the strided sweep is paid once instead of once per observable.
    pub fn expectations_pauli(&self, masks: &[(usize, usize, u32)]) -> Vec<f64> {
        let d = self.dim();
        #[cfg(feature = "gpu")]
        if let Some(gpu) = self.sv.gpu_state() {
            let pairs: Vec<(u64, u64)> = masks
                .iter()
                .map(|&(x, z, _)| (x as u64, z as u64))
                .collect();
            let sums = launched(dk::pauli_sums(gpu.context(), gpu, self.num_qubits, &pairs));
            return sums
                .iter()
                .zip(masks)
                .map(|(value, &(_, _, num_y))| (value * i_pow(num_y)).re)
                .collect();
        }
        let mut acc = vec![Complex64::new(0.0, 0.0); masks.len()];
        for j in 0..d {
            let row = &self.sv.state[j * d..(j + 1) * d];
            for (slot, &(xmask, zmask, _)) in acc.iter_mut().zip(masks) {
                let sign = if (j & zmask).count_ones() & 1 == 1 {
                    -1.0
                } else {
                    1.0
                };
                *slot += row[j ^ xmask] * sign;
            }
        }
        acc.iter()
            .zip(masks)
            .map(|(value, &(_, _, num_y))| (value * i_pow(num_y)).re)
            .collect()
    }
}

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

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

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

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

    fn schmidt_values(&mut self, _subsystem: &[usize]) -> Result<Vec<f64>> {
        Err(crate::error::PrismError::BackendUnsupported {
            backend: self.name().to_string(),
            operation: "Schmidt values of a mixed state".to_string(),
        })
    }

    /// Declined: the fidelity of two mixtures is not an inner product, and a
    /// mixture holds no statevector for the dense route to export.
    fn overlap_sq(&self, _other: &dyn Backend) -> Result<f64> {
        Err(crate::error::PrismError::BackendUnsupported {
            backend: self.name().to_string(),
            operation: "state overlap".to_string(),
        })
    }

    /// Partial trace of the `4^n` buffer, `2^(n + k)` reads. At `k = n` the
    /// buffer itself comes back, its rows and columns in `subsystem` order.
    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
            .sv
            .gpu_state()
            .map(|_| self.sv.export_statevector())
            .transpose()?;
        #[cfg(feature = "gpu")]
        let buffer = readback.as_deref().unwrap_or(self.sv.state.as_slice());
        #[cfg(not(feature = "gpu"))]
        let buffer = self.sv.state.as_slice();
        let mut rho = self.partial_trace(buffer, subsystem);
        reduced_density::normalize_trace(&mut rho, dim);
        Ok(rho)
    }

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

    /// With a device attached the host cap does not apply: the mixture is
    /// budgeted against free VRAM instead, before anything is allocated.
    fn init(&mut self, num_qubits: usize, num_classical_bits: usize) -> Result<()> {
        self.check_capacity(num_qubits)?;
        self.num_qubits = num_qubits;
        self.classical_bits = vec![false; num_classical_bits];
        self.sv.init(2 * num_qubits, 0)
    }

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

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

    /// The buffer becomes `rho` itself: moved in on the host, uploaded in one
    /// transfer when the mixture is device resident. The cap check of
    /// [`init`](Self::init) runs first, before the buffer is sized.
    fn init_from_density_matrix(
        &mut self,
        rho: Vec<Complex64>,
        num_classical_bits: usize,
    ) -> Result<()> {
        let num_qubits = crate::backend::validate_initial_density_matrix(&rho)?;
        self.check_capacity(num_qubits)?;
        self.num_qubits = num_qubits;
        self.classical_bits = vec![false; num_classical_bits];
        self.sv.init_from_state(rho, 0)
    }

    /// Starts from the pure mixture `|psi><psi|`, so the buffer is the `4^n`
    /// outer product of `amplitudes` with its own conjugate. The cap check in
    /// [`init`](Self::init) runs first, before that buffer is sized.
    fn init_from_amplitudes(
        &mut self,
        amplitudes: Vec<Complex64>,
        num_classical_bits: usize,
    ) -> Result<()> {
        crate::backend::validate_initial_amplitudes(&amplitudes)?;
        let num_qubits = amplitudes.len().trailing_zeros() as usize;
        self.init(num_qubits, num_classical_bits)?;

        #[cfg(feature = "gpu")]
        if let Some((ctx, gpu)) = self.device() {
            return dk::outer_product(&ctx, gpu, num_qubits, &amplitudes);
        }

        let d = amplitudes.len();
        for (row, &amp) in amplitudes.iter().enumerate() {
            let dst = &mut self.sv.state[row * d..(row + 1) * d];
            for (entry, &col) in dst.iter_mut().zip(amplitudes.iter()) {
                *entry = amp * col.conj();
            }
        }
        Ok(())
    }

    fn apply(&mut self, instruction: &Instruction) -> Result<()> {
        match instruction {
            Instruction::Gate { gate, targets } => self.apply_unitary(gate, targets),
            Instruction::Barrier { .. } => Ok(()),
            Instruction::Save { label, .. } => {
                Err(crate::backend::save_not_applied("DensityMatrix", label))
            }
            Instruction::Measure {
                qubit,
                classical_bit,
            } => self.apply_measure(*qubit, *classical_bit),
            Instruction::Reset { qubit } => self.apply_reset(*qubit),
            Instruction::Conditional {
                condition,
                gate,
                targets,
            } => self.apply_conditional(condition, gate, targets),
            Instruction::Region(region) => self.apply_region(region),
        }
    }

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

    fn probabilities(&self) -> Result<Vec<f64>> {
        let d = self.dim();
        #[cfg(feature = "gpu")]
        if let Some(gpu) = self.sv.gpu_state() {
            let mut diag = dk::diagonal(gpu.context(), gpu, self.num_qubits)?;
            for p in &mut diag {
                *p = p.max(0.0);
            }
            return Ok(diag);
        }
        let mut probs = vec![0.0f64; d];
        for (k, p) in probs.iter_mut().enumerate() {
            *p = self.sv.state[k * d + k].re.max(0.0);
        }
        Ok(probs)
    }

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

    /// `Multi2q` cuts its doubled list into runs the tiled kernel applies in
    /// order, since the ket register's shifted indices can push a batch past
    /// one tier.
    fn supports_fused_gates(&self) -> bool {
        true
    }

    /// The mixture is a `2n`-qubit statevector, so every fusion floor is
    /// reached at half the circuit width the statevector needs.
    fn fusion_state_qubits(&self, num_qubits: usize) -> usize {
        2 * num_qubits
    }

    fn qubit_probability(&self, qubit: usize) -> Result<f64> {
        self.prob_one(qubit)
    }

    fn reset(&mut self, qubit: usize) -> Result<()> {
        self.apply_reset(qubit)
    }

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

    /// `Tr(rho P_k)` per observable, the mixed-state reading of the trait's
    /// `<psi|P_k|psi>`. `rho` is trace-one by construction, so no
    /// normalization divide is needed.
    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<_>>>()?;
        Ok(self.expectations_pauli(&masks))
    }

    /// On the device the four entries come from the Pauli sums `T`, `Z`, `X`,
    /// and `Y` on `qubit`: the diagonal is `(T +- Z) / 2` and the off-diagonal
    /// pair is `(X +- Y) / 2`, where `Y` carries the row sign and no `i`.
    fn reduced_density_matrix_1q(&self, qubit: usize) -> Result<[[Complex64; 2]; 2]> {
        let n = self.num_qubits;
        let d = self.dim();
        let bit = 1usize << qubit;
        #[cfg(feature = "gpu")]
        if let Some(gpu) = self.sv.gpu_state() {
            let b = bit as u64;
            let sums = dk::pauli_sums(gpu.context(), gpu, n, &[(0, 0), (0, b), (b, 0), (b, b)])?;
            let (t, z, x, y) = (sums[0], sums[1], sums[2], sums[3]);
            return Ok([
                [(t + z) * 0.5, (x + y) * 0.5],
                [(x - y) * 0.5, (t - z) * 0.5],
            ]);
        }
        let others = 1usize << (n - 1);
        let mut r00 = Complex64::new(0.0, 0.0);
        let mut r01 = Complex64::new(0.0, 0.0);
        let mut r10 = Complex64::new(0.0, 0.0);
        let mut r11 = Complex64::new(0.0, 0.0);
        for m in 0..others {
            let base = (m & (bit - 1)) | ((m >> qubit) << (qubit + 1));
            let i1 = base | bit;
            r00 += self.sv.state[base * d + base];
            r01 += self.sv.state[base * d + i1];
            r10 += self.sv.state[i1 * d + base];
            r11 += self.sv.state[i1 * d + i1];
        }
        Ok([[r00, r01], [r10, r11]])
    }

    /// Evolve `rho -> K rho K^dagger` for an arbitrary `K`, on the same kernel
    /// selection as the one-qubit branch of `apply_unitary`. Trajectories never
    /// route here, since `supports_noisy_per_shot` excludes the density matrix in
    /// favour of `apply_1q_kraus` on the mixture; this keeps the trait method
    /// allocation-free on every backend that holds a state.
    fn apply_1q_matrix(&mut self, qubit: usize, matrix: &[[Complex64; 2]; 2]) -> Result<()> {
        self.apply_1q_sandwich(qubit, matrix)
    }
}