prism-q 0.35.0

//! CUDA C kernel sources and launchers, one submodule per backend, plus the shared
//! scratch and staging helpers. `GpuDevice::new` compiles the concatenated source and
//! pre-resolves every name in `KERNEL_NAMES`.

pub(crate) mod bts;
pub(crate) mod dense;
pub(crate) mod density;
pub(crate) mod stabilizer;

use std::sync::Arc;

use cudarc::driver::{
    CudaFunction, CudaStream, DevicePtr, DeviceRepr, LaunchConfig, PushKernelArg, ValidAsZeroBits,
};
use num_complex::Complex64;

use crate::error::{PrismError, Result};
use crate::gpu::device::GpuDevice;
use crate::gpu::memory::GpuBuffer;
use crate::gpu::{GpuContext, driver_err};

pub(super) fn stream_and_fn<'a>(
    ctx: &'a GpuContext,
    name: &'static str,
) -> Result<(&'a Arc<CudaStream>, CudaFunction)> {
    let device = ctx.device();
    Ok((device.stream()?, device.function(name)?))
}

pub(super) fn linear_cfg(block_size: u32, grid_blocks: u32) -> LaunchConfig {
    LaunchConfig {
        grid_dim: (grid_blocks, 1, 1),
        block_dim: (block_size, 1, 1),
        shared_mem_bytes: 0,
    }
}

pub(super) fn check_qubit(n: usize, qubit: usize) -> Result<()> {
    if qubit >= n {
        return Err(PrismError::InvalidQubit {
            index: qubit,
            register_size: n,
        });
    }
    Ok(())
}

/// Two distinct qubits below `n`; the error names the larger index.
pub(super) fn check_pair(n: usize, q0: usize, q1: usize) -> Result<()> {
    if q0 >= n || q1 >= n || q0 == q1 {
        return Err(PrismError::InvalidQubit {
            index: q0.max(q1),
            register_size: n,
        });
    }
    Ok(())
}

pub(super) fn launch_limit_err(op: &str, name: &str, value: usize, limit: &str) -> PrismError {
    PrismError::BackendUnsupported {
        backend: "gpu".to_string(),
        operation: format!("{op}: {name}={value} exceeds {limit} kernel limit"),
    }
}

pub(super) fn require_i32(op: &str, name: &str, value: usize) -> Result<i32> {
    i32::try_from(value).map_err(|_| launch_limit_err(op, name, value, "i32"))
}

pub(super) fn require_u32(op: &str, name: &str, value: usize) -> Result<u32> {
    u32::try_from(value).map_err(|_| launch_limit_err(op, name, value, "u32"))
}

pub(super) fn div_ceil_grid(op: &str, name: &str, value: usize, block: u32) -> Result<u32> {
    Ok(require_u32(op, name, value)?.div_ceil(block).max(1))
}

/// Grow-only device buffers reused across launches, so a small per-call upload fills a
/// resident allocation instead of allocating through `clone_htod`.
///
/// Tied to the `GpuContext` and serialized through a `Mutex`, since every launch shares
/// one CUDA stream anyway.
#[derive(Default)]
pub(crate) struct LauncherScratch {
    pub(crate) f64_a: Option<GpuBuffer<f64>>,
    pub(crate) u64_a: Option<GpuBuffer<u64>>,
    pub(crate) u64_b: Option<GpuBuffer<u64>>,
    /// Per-block partials for [`super::dense::measure_prob_one`]. Sized by the
    /// number of grid blocks at the largest qubit count seen so far.
    pub(crate) measure_partials: Option<GpuBuffer<f64>>,
    /// One-element output for the measure-prob finalize reduction.
    pub(crate) measure_result: Option<GpuBuffer<f64>>,
    /// Four-element output for the reduced-density-matrix finalize reduction.
    /// Separate from `measure_result`: readback length must match the buffer.
    pub(crate) rdm_result: Option<GpuBuffer<f64>>,
    /// The `2^n` diagonal of a density matrix, sized to the mixture width.
    pub(crate) dm_diag: Option<GpuBuffer<f64>>,
    /// Two f64s per Pauli mask for the Pauli expectation finalize, pure or
    /// mixed state. Exact length: the readback covers the whole buffer.
    pub(crate) pauli_result: Option<GpuBuffer<f64>>,
    /// Three f64s per block for [`super::dense::group_moments`]. Exact length, like
    /// `pauli_result`.
    pub(crate) moment_partials: Option<GpuBuffer<f64>>,
    /// One upload per gate for what does not fit parameter space; see [`stage_blob`].
    pub(crate) blob: BlobScratch,
}

const BLOB_ALIGN: usize = 16;

/// Device addresses of one staged blob: the `int` list at the start and the `double`
/// block at a 16-byte-aligned offset behind it.
#[derive(Clone, Copy)]
pub(crate) struct Staged {
    pub(crate) ints: u64,
    pub(crate) f64s: u64,
}

/// Host and device halves of the per-gate upload behind [`stage_blob`].
#[derive(Default)]
pub(crate) struct BlobScratch {
    host: Vec<u64>,
    device: Option<GpuBuffer<u64>>,
}

/// Pack `ints` and `f64_len` doubles written through `fill` into one upload and return
/// the device addresses of both halves.
///
/// Everything a kernel can take by value goes in parameter space instead; this is for
/// the lookup tables and metadata of the batched diagonal gates and the density-matrix
/// tables, which exceed the 4 KB parameter limit. The copy is synchronous from pageable
/// memory, so it costs one stream drain per gate, and the buffer is reused as soon as
/// the copy returns because that drain has already retired the kernel that read it.
pub(crate) fn stage_blob(
    scratch: &mut BlobScratch,
    device: &GpuDevice,
    ints: impl ExactSizeIterator<Item = i32>,
    f64_len: usize,
    fill: impl FnOnce(&mut [f64]),
) -> Result<Staged> {
    let int_len = ints.len();
    let f64_off = (int_len * size_of::<i32>()).next_multiple_of(BLOB_ALIGN);
    let words = (f64_off + f64_len * size_of::<f64>()) / size_of::<u64>();
    let host = &mut scratch.host;
    if host.len() < words {
        host.resize(words, 0);
    }
    // SAFETY: `host` holds at least `words` u64s, so both views lie inside it; the int
    // view ends at `f64_off`, where the f64 view starts, so they are disjoint; and a u64
    // buffer is aligned for both element types.
    unsafe {
        let base = host.as_mut_ptr().cast::<u8>();
        let int_dst = std::slice::from_raw_parts_mut(base.cast::<i32>(), int_len);
        for (d, v) in int_dst.iter_mut().zip(ints) {
            *d = v;
        }
        let f64_dst = std::slice::from_raw_parts_mut(base.add(f64_off).cast::<f64>(), f64_len);
        fill(f64_dst);
    }
    let buf = ensure_scratch(&mut scratch.device, device, &host[..words])?;
    let stream = device.stream()?;
    let (base, _sync) = buf.raw().device_ptr(stream);
    Ok(Staged {
        ints: base,
        f64s: base + f64_off as u64,
    })
}

/// Stage complex values as interleaved re/im f64 pairs.
pub(crate) fn stage_complex(
    scratch: &mut BlobScratch,
    device: &GpuDevice,
    values: &[num_complex::Complex64],
) -> Result<Staged> {
    stage_blob(
        scratch,
        device,
        std::iter::empty(),
        2 * values.len(),
        |dst| {
            for (pair, v) in dst.as_chunks_mut::<2>().0.iter_mut().zip(values) {
                pair[0] = v.re;
                pair[1] = v.im;
            }
        },
    )
}

/// Grow `slot` to at least `host.len()` elements and copy `host` into its start.
pub(crate) fn ensure_scratch<'a, T: DeviceRepr + ValidAsZeroBits>(
    slot: &'a mut Option<GpuBuffer<T>>,
    device: &GpuDevice,
    host: &[T],
) -> Result<&'a GpuBuffer<T>> {
    let needed = host.len().max(1);
    let realloc = match slot.as_ref() {
        Some(buf) => buf.len() < needed,
        None => true,
    };
    if realloc {
        *slot = Some(GpuBuffer::<T>::alloc_zeros(device, needed)?);
    }
    if !host.is_empty() {
        slot.as_mut().unwrap().copy_from_host(device, host)?;
    }
    Ok(slot.as_ref().unwrap())
}

/// Grow `slot` to at least `needed` elements without a host copy, for kernel outputs.
pub(crate) fn ensure_capacity<'a, T: DeviceRepr + ValidAsZeroBits>(
    slot: &'a mut Option<GpuBuffer<T>>,
    device: &GpuDevice,
    needed: usize,
) -> Result<&'a mut GpuBuffer<T>> {
    let needed = needed.max(1);
    let realloc = match slot.as_ref() {
        Some(buf) => buf.len() < needed,
        None => true,
    };
    if realloc {
        *slot = Some(GpuBuffer::<T>::alloc_zeros(device, needed)?);
    }
    Ok(slot.as_mut().unwrap())
}

/// A device buffer of exactly `len` elements in `slot`, reallocated on any size
/// change so a readback of its full length matches the request.
pub(crate) fn ensure_exact<'a>(
    slot: &'a mut Option<GpuBuffer<f64>>,
    device: &GpuDevice,
    len: usize,
) -> Result<&'a mut GpuBuffer<f64>> {
    let len = len.max(1);
    if slot.as_ref().is_none_or(|buf| buf.len() != len) {
        *slot = Some(GpuBuffer::<f64>::alloc_zeros(device, len)?);
    }
    Ok(slot.as_mut().unwrap())
}

/// Reduce the `[mask][block][re, im]` partials in `measure_partials` with
/// `dm_pauli_expect_finalize` and read back one complex sum per mask, times `scale`.
pub(super) fn finalize_pauli_sums(
    device: &GpuDevice,
    scratch: &mut LauncherScratch,
    num_masks: u32,
    blocks_per_mask: u32,
    block_size: u32,
    scale: f64,
) -> Result<Vec<Complex64>> {
    let stream = device.stream()?;
    let finalize = device.function("dm_pauli_expect_finalize")?;
    let result = ensure_exact(&mut scratch.pauli_result, device, 2 * num_masks as usize)?;
    let cfg = LaunchConfig {
        grid_dim: (num_masks, 1, 1),
        block_dim: (block_size, 1, 1),
        shared_mem_bytes: 2 * block_size * size_of::<f64>() as u32,
    };
    let mut builder = stream.launch_builder(&finalize);
    builder
        .arg(scratch.measure_partials.as_ref().unwrap().raw())
        .arg(&blocks_per_mask)
        .arg(result.raw_mut());
    // SAFETY: signature matches the kernel; one block per mask strides over that
    // mask's `blocks_per_mask` partial pairs, and `result` holds two f64s per mask.
    // Both buffers are held by the caller's scratch guard.
    unsafe {
        builder
            .launch(cfg)
            .map_err(|e| driver_err("dm_pauli_expect_finalize", e))?;
    }
    let mut host = vec![0.0_f64; 2 * num_masks as usize];
    result.copy_to_host(device, &mut host)?;
    Ok(host
        .as_chunks::<2>()
        .0
        .iter()
        .map(|pair| Complex64::new(pair[0], pair[1]) * scale)
        .collect())
}

/// Combined CUDA C source for the PTX module. A kernel added to any part also needs
/// its name in [`KERNEL_NAMES`].
pub(crate) fn kernel_source() -> String {
    let mut src = dense::kernel_source();
    src.push('\n');
    src.push_str(&stabilizer::kernel_source());
    src.push('\n');
    src.push_str(&bts::kernel_source());
    src.push('\n');
    src.push_str(density::KERNEL_SOURCE);
    src
}

/// Every kernel entry point that appears in the materialised PTX source.
///
/// `GpuDevice::new` pre-resolves each name once so gate dispatch does not pay the
/// driver lookup per launch.
pub(crate) const KERNEL_NAMES: &[&str] = &[
    // Dense statevector kernels.
    "set_initial_state",
    "apply_gate_1q",
    "apply_diagonal_1q",
    "apply_cx",
    "apply_cz",
    "apply_swap",
    "apply_parity_phase",
    "apply_cu",
    "apply_cu_phase",
    "apply_mcu",
    "apply_mcu_phase",
    "apply_fused_2q",
    "measure_prob_one",
    "measure_prob_one_finalize",
    "rdm_qubit",
    "rdm_qubit_finalize",
    "sv_pauli_expect",
    "sv_group_moments",
    "measure_collapse",
    "compute_probabilities",
    "apply_multi_fused_diagonal",
    "apply_batch_phase",
    "apply_batch_rzz",
    "apply_diagonal_batch",
    "apply_multi_fused_tiled",
    // Stabilizer tableau kernels.
    "stab_set_initial_tableau",
    "stab_apply_word_grouped",
    "stab_rowmul_words",
    "stab_measure_find_pivot",
    "stab_measure_cascade",
    "stab_measure_fixup",
    "stab_measure_deterministic",
    // Block-triangular sampling.
    "bts_sample_meas_major",
    "bts_popcount_rows",
    "bts_count_meas_major_upto8",
    "bts_count_shot_major_upto8",
    "bts_count_used_slots",
    "bts_compact_counts_upto8",
    "bts_transpose_meas_to_shot",
    "bts_apply_noise_masks_meas_major",
    "bts_generate_and_apply_noise_meas_major_by_row",
    // Density-matrix sweeps over the embedded 2n-qubit buffer.
    "dm_diagonal",
    "dm_norm_sqr",
    "dm_project",
    "dm_reset",
    "dm_conjugate",
    "dm_diagonal_sandwich",
    "dm_kraus_2q_diagonal",
    "dm_kraus_2q_dense",
    "dm_depolarizing_2q",
    "dm_outer_product",
    "dm_pauli_expect",
    "dm_pauli_expect_finalize",
];

#[cfg(test)]
mod sync_tests {
    use super::{KERNEL_NAMES, kernel_source};
    use std::collections::BTreeSet;

    /// Names following each `__global__ void` marker in the PTX source.
    fn source_entry_points(src: &str) -> BTreeSet<String> {
        const MARKER: &str = "__global__ void ";
        let mut names = BTreeSet::new();
        let mut rest = src;
        while let Some(pos) = rest.find(MARKER) {
            rest = &rest[pos + MARKER.len()..];
            let name: String = rest
                .chars()
                .take_while(|c| c.is_alphanumeric() || *c == '_')
                .collect();
            if !name.is_empty() {
                names.insert(name);
            }
        }
        names
    }

    // `KERNEL_NAMES` must match the kernels in the CUDA source exactly: a missing
    // name is never pre-resolved by `GpuDevice::new`, an orphan name is dead drift.
    // Inspects the source string only, so it needs no GPU device.
    #[test]
    fn kernel_names_match_source_entry_points() {
        let src = kernel_source();
        let in_source = source_entry_points(&src);
        let declared: BTreeSet<String> = KERNEL_NAMES.iter().map(|s| s.to_string()).collect();

        let missing: Vec<&String> = in_source.difference(&declared).collect();
        let orphan: Vec<&String> = declared.difference(&in_source).collect();

        assert!(
            missing.is_empty(),
            "kernels defined in source but absent from KERNEL_NAMES \
             (GpuDevice::new would not pre-resolve them): {missing:?}"
        );
        assert!(
            orphan.is_empty(),
            "KERNEL_NAMES entries with no matching kernel in source: {orphan:?}"
        );
    }
}