prism-q 0.35.0

//! Block-triangular sampling on the GPU: a device port of the CPU `sample_bts_meas_major`
//! fed by a host-generated random-bits pool, so both paths share one RNG contract, plus
//! device noise application, shot-major transpose, and outcome counting.
//!
//! The sampling kernel reads the precomputed CSR parity matrix and runs one thread per
//! (measurement, 64-shot batch) pair, emitting packed outcomes in measurement-major layout.

use std::sync::Arc;

use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};

use crate::error::Result;
use crate::gpu::GpuBuffer;
use crate::sim::compiled::parity::SparseParity;
use crate::sim::compiled::rng::Xoshiro256PlusPlus;
use crate::sim::compiled::shot_tail_mask;

use super::super::GpuContext;
use super::{div_ceil_grid, driver_err, linear_cfg, require_i32, require_u32, stream_and_fn};

const SAMPLE_BLOCK_SIZE: u32 = 128;
const NOISE_BLOCK_SIZE: u32 = 128;
const POPCOUNT_BLOCK_SIZE: u32 = 256;
const COUNT_MEAS_BATCH_SIZE: u32 = 64;
const COUNT_HASH_BLOCK_SIZE: u32 = 256;
const COUNT_COMPACT_BLOCK_SIZE: u32 = 256;

/// Maximum packed measurement words per shot key for the device-side count
/// kernels; wider keys fall back to host counting.
pub(crate) const GPU_COUNTS_MAX_WORDS: usize = 8;

/// Shots per device chunk, which bounds peak memory at any shot count. Matches the CPU
/// `BTS_BATCH_SHOTS` boundary.
const CHUNK_SHOTS: usize = 65_536;

const KERNEL_SOURCE: &str = include_str!("bts.cu");

pub(crate) fn kernel_source() -> String {
    KERNEL_SOURCE.to_string()
}

/// Batch count rounded up to a power of two in `[32, NOISE_BLOCK_SIZE]`, so small tiles
/// do not leave half the block idle.
fn noise_block_threads(chunk_s_words: usize) -> u32 {
    let desired = chunk_s_words.max(1).next_power_of_two();
    desired.clamp(32, NOISE_BLOCK_SIZE as usize) as u32
}

/// Reusable device and host storage for GPU BTS sampling.
///
/// The sparse parity CSR arrays and packed reference bits are uploaded once per
/// compiled sampler. `random_bits` and the chunk output buffer grow to the
/// largest chunk requested, then stay resident across repeated sampling calls.
pub(crate) struct GpuBtsCache {
    num_meas: usize,
    col_indices_dev: GpuBuffer<u32>,
    row_offsets_dev: GpuBuffer<u32>,
    ref_bits_dev: GpuBuffer<u64>,
    random_bits_dev: GpuBuffer<u64>,
    chunk_output_dev: GpuBuffer<u64>,
    random_bits_host: Vec<u64>,
    chunk_output_host: Vec<u64>,
}

impl GpuBtsCache {
    pub(crate) fn new(ctx: &GpuContext, sparse: &SparseParity, ref_bits: &[u64]) -> Result<Self> {
        let device = ctx.device();
        let mut col_indices_dev =
            GpuBuffer::<u32>::alloc_zeros(device, sparse.col_indices.len().max(1))?;
        let mut row_offsets_dev =
            GpuBuffer::<u32>::alloc_zeros(device, sparse.row_offsets.len().max(1))?;
        let mut ref_bits_dev = GpuBuffer::<u64>::alloc_zeros(device, ref_bits.len().max(1))?;

        if !sparse.col_indices.is_empty() {
            col_indices_dev.copy_from_host(device, &sparse.col_indices)?;
        }
        row_offsets_dev.copy_from_host(device, &sparse.row_offsets)?;
        if !ref_bits.is_empty() {
            ref_bits_dev.copy_from_host(device, ref_bits)?;
        }

        Ok(Self {
            num_meas: sparse.num_rows,
            col_indices_dev,
            row_offsets_dev,
            ref_bits_dev,
            random_bits_dev: GpuBuffer::<u64>::alloc_zeros(device, 1)?,
            chunk_output_dev: GpuBuffer::<u64>::alloc_zeros(device, 1)?,
            random_bits_host: Vec::new(),
            chunk_output_host: Vec::new(),
        })
    }

    fn ensure_chunk_capacity(
        &mut self,
        ctx: &GpuContext,
        random_len: usize,
        output_len: usize,
    ) -> Result<()> {
        let device = ctx.device();
        let random_needed = random_len.max(1);
        if self.random_bits_dev.len() < random_needed {
            self.random_bits_dev = GpuBuffer::<u64>::alloc_zeros(device, random_needed)?;
        }
        if self.chunk_output_dev.len() < output_len.max(1) {
            self.chunk_output_dev = GpuBuffer::<u64>::alloc_zeros(device, output_len.max(1))?;
        }
        if self.random_bits_host.len() < random_needed {
            self.random_bits_host.resize(random_needed, 0);
        }
        if self.chunk_output_host.len() < output_len.max(1) {
            self.chunk_output_host.resize(output_len.max(1), 0);
        }
        Ok(())
    }

    fn fill_random_bits(
        &mut self,
        rng: &mut Xoshiro256PlusPlus,
        rank: usize,
        chunk_shots: usize,
        chunk_s_words: usize,
    ) {
        // Parallel path: seed one fresh xoshiro stream per worker from the
        // master rng, matching the CPU BTS batched path in
        // `src/sim/compiled/bts.rs`. Partitioning by `current_num_threads()`
        // means a seed reproduces only at a fixed thread count; the serial
        // fallback below `MIN_PAR_DRAWS` reproduces outright.
        let required = chunk_s_words * rank;
        #[cfg(not(feature = "parallel"))]
        let _ = required;
        #[cfg(feature = "parallel")]
        {
            const MIN_PAR_DRAWS: usize = 16_384;
            let num_threads = rayon::current_num_threads();
            if num_threads > 1 && required >= MIN_PAR_DRAWS {
                let batches_per_thread = chunk_s_words.div_ceil(num_threads);
                if batches_per_thread > 0 {
                    let thread_seeds: Vec<[u64; 4]> = (0..num_threads)
                        .map(|_| {
                            [
                                rng.next_u64(),
                                rng.next_u64(),
                                rng.next_u64(),
                                rng.next_u64(),
                            ]
                        })
                        .collect();
                    let tail_mask = shot_tail_mask(chunk_shots);
                    let last_batch = chunk_s_words - 1;
                    use rayon::prelude::*;
                    self.random_bits_host[..required]
                        .par_chunks_mut(batches_per_thread * rank)
                        .enumerate()
                        .for_each(|(tid, slab)| {
                            let mut trng = Xoshiro256PlusPlus::from_seeds(thread_seeds[tid]);
                            let first_batch = tid * batches_per_thread;
                            let batches_here = slab.len() / rank;
                            for b in 0..batches_here {
                                let start = b * rank;
                                let end = start + rank;
                                for word in &mut slab[start..end] {
                                    *word = trng.next_u64();
                                }
                                let absolute_batch = first_batch + b;
                                if absolute_batch == last_batch && tail_mask != u64::MAX {
                                    for word in &mut slab[start..end] {
                                        *word &= tail_mask;
                                    }
                                }
                            }
                        });
                    return;
                }
            }
        }

        for batch in 0..chunk_s_words {
            let start = batch * rank;
            let end = start + rank;
            for word in &mut self.random_bits_host[start..end] {
                *word = rng.next_u64();
            }
            if batch == chunk_s_words - 1 {
                let mask = shot_tail_mask(chunk_shots);
                if mask != u64::MAX {
                    for word in &mut self.random_bits_host[start..end] {
                        *word &= mask;
                    }
                }
            }
        }
    }

    /// Upload the filled random bits and launch `bts_sample_meas_major` into `out` (the
    /// chunk output buffer when `None`), laid out as `num_meas` rows of `out_stride_words`
    /// words with the chunk starting at word `out_word_offset` of each row.
    #[allow(clippy::too_many_arguments)]
    fn launch_sample_chunk(
        &mut self,
        stream: &Arc<CudaStream>,
        func: &CudaFunction,
        rank: usize,
        chunk_shots: usize,
        out_stride_words: usize,
        out_word_offset: usize,
        out: Option<&mut GpuBuffer<u64>>,
    ) -> Result<()> {
        let num_meas = self.num_meas;
        let chunk_s_words = chunk_shots.div_ceil(64);
        let chunk_random_len = chunk_s_words * rank;
        {
            let mut random_bits_view = self
                .random_bits_dev
                .raw_mut()
                .slice_mut(0..chunk_random_len.max(1));
            stream
                .memcpy_htod(
                    &self.random_bits_host[..chunk_random_len],
                    &mut random_bits_view,
                )
                .map_err(|e| driver_err("upload bts random_bits", e))?;
        }

        let num_meas_i = require_i32("bts_sample_meas_major", "num_meas", num_meas)?;
        let chunk_s_words_i = require_i32("bts_sample_meas_major", "chunk_s_words", chunk_s_words)?;
        let rank_i = require_i32("bts_sample_meas_major", "rank", rank)?;
        let out_stride_words_i = require_i32(
            "bts_sample_meas_major",
            "out_stride_words",
            out_stride_words,
        )?;
        let out_word_offset_i =
            require_i32("bts_sample_meas_major", "out_word_offset", out_word_offset)?;
        let tail_mask_u64 = shot_tail_mask(chunk_shots);
        let batch_blocks = div_ceil_grid(
            "bts_sample_meas_major",
            "chunk_s_words",
            chunk_s_words,
            SAMPLE_BLOCK_SIZE,
        )?;
        let num_meas_grid = require_u32("bts_sample_meas_major", "num_meas", num_meas)?;
        let cfg = LaunchConfig {
            grid_dim: (num_meas_grid, batch_blocks, 1),
            block_dim: (SAMPLE_BLOCK_SIZE, 1, 1),
            shared_mem_bytes: 0,
        };

        let out = out.unwrap_or(&mut self.chunk_output_dev);
        let mut out_view = out
            .raw_mut()
            .slice_mut(0..(num_meas * out_stride_words).max(1));
        let random_bits_dev = self.random_bits_dev.raw().slice(0..chunk_random_len.max(1));
        let mut builder = stream.launch_builder(func);
        builder
            .arg(self.col_indices_dev.raw())
            .arg(self.row_offsets_dev.raw())
            .arg(self.ref_bits_dev.raw())
            .arg(&random_bits_dev)
            .arg(&num_meas_i)
            .arg(&chunk_s_words_i)
            .arg(&rank_i)
            .arg(&out_stride_words_i)
            .arg(&out_word_offset_i)
            .arg(&tail_mask_u64)
            .arg(&mut out_view);
        // SAFETY: kernel signature matches the call shape. Each thread writes one unique
        // output word, addressed by measurement row and chunk-local batch index, inside
        // the `num_meas * out_stride_words` view. All input buffers are sized to cover
        // the access pattern above.
        unsafe {
            builder
                .launch(cfg)
                .map_err(|e| driver_err("bts_sample_meas_major", e))?;
        }
        Ok(())
    }
}

fn next_power_of_two_or_max(n: usize) -> usize {
    n.checked_next_power_of_two().unwrap_or(usize::MAX)
}

fn highest_power_of_two_at_most(n: usize) -> usize {
    if n == 0 {
        0
    } else {
        1usize << (usize::BITS as usize - 1 - n.leading_zeros() as usize)
    }
}

fn counts_slot_bytes(m_words: usize) -> usize {
    std::mem::size_of::<u32>() + std::mem::size_of::<u64>() + m_words * std::mem::size_of::<u64>()
}

fn counts_outcome_bound(num_shots: usize, rank: usize) -> usize {
    if rank >= usize::BITS as usize {
        num_shots
    } else {
        num_shots.min(1usize << rank)
    }
}

fn count_raw_transfer_bytes(len_words: usize) -> usize {
    len_words.saturating_mul(std::mem::size_of::<u64>())
}

fn count_compact_transfer_bound_bytes(num_shots: usize, m_words: usize, rank: usize) -> usize {
    counts_outcome_bound(num_shots, rank)
        .saturating_mul(m_words + 1)
        .saturating_mul(std::mem::size_of::<u64>())
}

fn should_try_device_count(
    num_shots: usize,
    m_words: usize,
    rank: usize,
    raw_transfer_bytes: usize,
) -> bool {
    if raw_transfer_bytes == 0 {
        return false;
    }
    count_compact_transfer_bound_bytes(num_shots, m_words, rank) < raw_transfer_bytes
}

fn plan_count_table_slots(
    ctx: &GpuContext,
    num_shots: usize,
    m_words: usize,
    rank: usize,
) -> Result<Option<usize>> {
    if m_words == 0 || m_words > GPU_COUNTS_MAX_WORDS {
        return Ok(None);
    }

    let available = ctx.vram_available()?;
    let slot_bytes = counts_slot_bytes(m_words);
    let budget = available / 2;
    let max_slots_fit = highest_power_of_two_at_most(budget / slot_bytes);
    if max_slots_fit < 64 || max_slots_fit > u32::MAX as usize + 1 {
        return Ok(None);
    }

    let distinct_bound = counts_outcome_bound(num_shots, rank);
    let target_slots = next_power_of_two_or_max(distinct_bound.saturating_mul(2).max(64));
    let table_slots = target_slots.min(max_slots_fit);
    if table_slots < 64 || table_slots > i32::MAX as usize {
        return Ok(None);
    }

    Ok(Some(table_slots))
}

/// Sample `num_shots` BTS shots on the GPU with the `sample_bts_meas_major` contract:
/// measurement-major, `num_meas * num_shots.div_ceil(64)` words, and deterministic rows
/// (empty `row_cols`) carrying only the `ref_bits` flip.
///
/// Shots stream through the device in `CHUNK_SHOTS` chunks. The parity CSR and
/// `ref_bits` stay resident in `cache`; only random bits and chunk output move.
pub(crate) fn launch_bts_sample(
    ctx: &GpuContext,
    rng: &mut Xoshiro256PlusPlus,
    rank: usize,
    num_shots: usize,
    cache: &mut GpuBtsCache,
) -> Result<Vec<u64>> {
    let num_meas = cache.num_meas;
    let s_words = num_shots.div_ceil(64);
    if num_meas == 0 || num_shots == 0 || rank == 0 {
        return Ok(vec![0u64; num_meas * s_words]);
    }

    let (stream, func) = stream_and_fn(ctx, "bts_sample_meas_major")?;

    let mut output = vec![0u64; num_meas * s_words];
    let mut shots_done = 0usize;
    while shots_done < num_shots {
        let chunk_shots = (num_shots - shots_done).min(CHUNK_SHOTS);
        let chunk_s_words = chunk_shots.div_ceil(64);
        let chunk_random_len = chunk_s_words * rank;
        let chunk_output_len = num_meas * chunk_s_words;
        cache.ensure_chunk_capacity(ctx, chunk_random_len, chunk_output_len)?;
        cache.fill_random_bits(rng, rank, chunk_shots, chunk_s_words);

        cache.launch_sample_chunk(stream, &func, rank, chunk_shots, chunk_s_words, 0, None)?;

        stream
            .memcpy_dtoh(
                &cache
                    .chunk_output_dev
                    .raw()
                    .slice(0..chunk_output_len.max(1)),
                &mut cache.chunk_output_host[..chunk_output_len.max(1)],
            )
            .map_err(|e| driver_err("bts output dtoh", e))?;

        let chunk_word_offset = shots_done / 64;
        for m in 0..num_meas {
            let dst_start = m * s_words + chunk_word_offset;
            let src_start = m * chunk_s_words;
            output[dst_start..dst_start + chunk_s_words]
                .copy_from_slice(&cache.chunk_output_host[src_start..src_start + chunk_s_words]);
        }

        shots_done += chunk_shots;
    }

    Ok(output)
}

/// Same chunked contract as [`launch_bts_sample`], but the meas-major result
/// (length `num_meas * s_words` u64s) stays resident on device to feed the
/// on-device transpose and count paths.
pub(crate) fn launch_bts_sample_device(
    ctx: &GpuContext,
    rng: &mut Xoshiro256PlusPlus,
    rank: usize,
    num_shots: usize,
    cache: &mut GpuBtsCache,
) -> Result<GpuBuffer<u64>> {
    let num_meas = cache.num_meas;
    let s_words = num_shots.div_ceil(64);
    let output_len = num_meas * s_words;
    if num_meas == 0 || num_shots == 0 || rank == 0 {
        return GpuBuffer::<u64>::alloc_zeros(ctx.device(), output_len.max(1));
    }

    let device = ctx.device();
    let (stream, func) = stream_and_fn(ctx, "bts_sample_meas_major")?;
    let mut output_dev = GpuBuffer::<u64>::alloc_zeros(device, output_len.max(1))?;

    let mut shots_done = 0usize;
    while shots_done < num_shots {
        let chunk_shots = (num_shots - shots_done).min(CHUNK_SHOTS);
        let chunk_s_words = chunk_shots.div_ceil(64);
        let chunk_random_len = chunk_s_words * rank;
        cache.ensure_chunk_capacity(ctx, chunk_random_len, 1)?;
        cache.fill_random_bits(rng, rank, chunk_shots, chunk_s_words);

        cache.launch_sample_chunk(
            stream,
            &func,
            rank,
            chunk_shots,
            s_words,
            shots_done / 64,
            Some(&mut output_dev),
        )?;

        shots_done += chunk_shots;
    }

    Ok(output_dev)
}

/// Bit-transpose a device-resident meas-major shot matrix into shot-major
/// layout on device. Result length is `num_shots * m_words` u64s.
pub(crate) fn launch_bts_transpose_meas_to_shot(
    ctx: &GpuContext,
    meas_major: &GpuBuffer<u64>,
    num_meas: usize,
    num_shots: usize,
    s_words: usize,
    m_words: usize,
) -> Result<GpuBuffer<u64>> {
    let device = ctx.device();
    let shot_major_len = (num_shots * m_words).max(1);
    let mut shot_major = GpuBuffer::<u64>::alloc_zeros(device, shot_major_len)?;
    if num_shots == 0 || num_meas == 0 || m_words == 0 || s_words == 0 {
        return Ok(shot_major);
    }

    let stream = device.stream()?;
    let func = device.function("bts_transpose_meas_to_shot")?;
    let num_meas_i = require_i32("bts_transpose_meas_to_shot", "num_meas", num_meas)?;
    let num_shots_i = require_i32("bts_transpose_meas_to_shot", "num_shots", num_shots)?;
    let s_words_i = require_i32("bts_transpose_meas_to_shot", "s_words", s_words)?;
    let m_words_i = require_i32("bts_transpose_meas_to_shot", "m_words", m_words)?;
    let m_words_grid = require_u32("bts_transpose_meas_to_shot", "m_words", m_words)?;
    let s_words_grid = require_u32("bts_transpose_meas_to_shot", "s_words", s_words)?;
    let cfg = LaunchConfig {
        grid_dim: (m_words_grid, s_words_grid, 1),
        block_dim: (64, 1, 1),
        shared_mem_bytes: 0,
    };

    let mut builder = stream.launch_builder(&func);
    builder
        .arg(meas_major.raw())
        .arg(&num_meas_i)
        .arg(&num_shots_i)
        .arg(&s_words_i)
        .arg(&m_words_i)
        .arg(shot_major.raw_mut());
    // SAFETY: grid covers every (m_word, batch) tile exactly once; each tile's
    // 64 threads produce 64 disjoint shot-major u64s. Shared tile is block-local.
    unsafe {
        builder
            .launch(cfg)
            .map_err(|e| driver_err("bts_transpose_meas_to_shot", e))?;
    }
    Ok(shot_major)
}

/// Target buffer and event CSR arguments of the host-mask noise kernel.
pub(crate) struct NoiseApplyBase<'a> {
    pub meas_major: &'a mut GpuBuffer<u64>,
    pub num_meas: usize,
    pub s_words: usize,
    pub word_offset: usize,
    pub chunk_s_words: usize,
    pub num_events: usize,
    pub x_row_offsets: &'a GpuBuffer<u32>,
    pub x_row_indices: &'a GpuBuffer<u32>,
    pub z_row_offsets: &'a GpuBuffer<u32>,
    pub z_row_indices: &'a GpuBuffer<u32>,
}

/// [`NoiseApplyBase`] plus per-event X/Z flip masks pre-generated on the host,
/// as opposed to the on-device generation of [`NoiseDeviceGenApplyByRow`].
pub(crate) struct NoiseMaskApply<'a> {
    pub base: NoiseApplyBase<'a>,
    pub x_masks: &'a GpuBuffer<u64>,
    pub z_masks: &'a GpuBuffer<u64>,
}

/// Arguments for the per-(row, batch) fused noise kernel. Each thread walks
/// one row's entry list in the row-major event CSR and XORs the generated
/// masks into a register, avoiding the cross-block `atomicXor` contention an
/// event-major launch pays on rows touched by many events. `row_event_offsets`
/// is length `num_meas + 1`; each `row_event_entries` entry is
/// `event << 2 | flag` (bit 0 = X, bit 1 = Z).
pub(crate) struct NoiseDeviceGenApplyByRow<'a> {
    pub meas_major: &'a mut GpuBuffer<u64>,
    pub num_meas: usize,
    pub s_words: usize,
    pub word_offset: usize,
    pub chunk_s_words: usize,
    pub row_event_offsets: &'a GpuBuffer<u32>,
    pub row_event_entries: &'a GpuBuffer<u32>,
    pub event_thresholds: &'a GpuBuffer<u64>,
    pub master_seed: u64,
    pub batch_offset: u64,
}

pub(crate) fn generate_and_apply_noise_masks_meas_major_by_row(
    ctx: &GpuContext,
    args: NoiseDeviceGenApplyByRow<'_>,
) -> Result<()> {
    let NoiseDeviceGenApplyByRow {
        meas_major,
        num_meas,
        s_words,
        word_offset,
        chunk_s_words,
        row_event_offsets,
        row_event_entries,
        event_thresholds,
        master_seed,
        batch_offset,
    } = args;
    if num_meas == 0 || s_words == 0 || chunk_s_words == 0 {
        return Ok(());
    }

    let (stream, func) = stream_and_fn(ctx, "bts_generate_and_apply_noise_meas_major_by_row")?;
    let num_meas_i = require_i32(
        "bts_generate_and_apply_noise_meas_major_by_row",
        "num_meas",
        num_meas,
    )?;
    let s_words_i = require_i32(
        "bts_generate_and_apply_noise_meas_major_by_row",
        "s_words",
        s_words,
    )?;
    let word_offset_i = require_i32(
        "bts_generate_and_apply_noise_meas_major_by_row",
        "word_offset",
        word_offset,
    )?;
    let chunk_s_words_i = require_i32(
        "bts_generate_and_apply_noise_meas_major_by_row",
        "chunk_s_words",
        chunk_s_words,
    )?;
    let block_threads = noise_block_threads(chunk_s_words);
    let batch_blocks = div_ceil_grid(
        "bts_generate_and_apply_noise_meas_major_by_row",
        "chunk_s_words",
        chunk_s_words,
        block_threads,
    )?;
    let num_meas_grid = require_u32(
        "bts_generate_and_apply_noise_meas_major_by_row",
        "num_meas",
        num_meas,
    )?;
    let cfg = LaunchConfig {
        grid_dim: (num_meas_grid, batch_blocks, 1),
        block_dim: (block_threads, 1, 1),
        shared_mem_bytes: 0,
    };

    let mut builder = stream.launch_builder(&func);
    builder
        .arg(meas_major.raw_mut())
        .arg(&num_meas_i)
        .arg(&s_words_i)
        .arg(&word_offset_i)
        .arg(row_event_offsets.raw())
        .arg(row_event_entries.raw())
        .arg(event_thresholds.raw())
        .arg(&chunk_s_words_i)
        .arg(&master_seed)
        .arg(&batch_offset);
    // SAFETY: one thread per (row, batch); register accumulator replaces the
    // cross-block atomic path. No two threads write the same output index.
    unsafe {
        builder
            .launch(cfg)
            .map_err(|e| driver_err("bts_generate_and_apply_noise_meas_major_by_row", e))?;
    }
    Ok(())
}

/// One thread per (event, batch) pair, with atomic XOR on rows shared
/// between events.
pub(crate) fn apply_noise_masks_meas_major(
    ctx: &GpuContext,
    args: NoiseMaskApply<'_>,
) -> Result<()> {
    let NoiseMaskApply {
        base:
            NoiseApplyBase {
                meas_major,
                num_meas,
                s_words,
                word_offset,
                chunk_s_words,
                num_events,
                x_row_offsets,
                x_row_indices,
                z_row_offsets,
                z_row_indices,
            },
        x_masks,
        z_masks,
    } = args;
    if num_meas == 0 || s_words == 0 || chunk_s_words == 0 || num_events == 0 {
        return Ok(());
    }

    let (stream, func) = stream_and_fn(ctx, "bts_apply_noise_masks_meas_major")?;
    let num_meas_i = require_i32("bts_apply_noise_masks_meas_major", "num_meas", num_meas)?;
    let s_words_i = require_i32("bts_apply_noise_masks_meas_major", "s_words", s_words)?;
    let word_offset_i = require_i32(
        "bts_apply_noise_masks_meas_major",
        "word_offset",
        word_offset,
    )?;
    let chunk_s_words_i = require_i32(
        "bts_apply_noise_masks_meas_major",
        "chunk_s_words",
        chunk_s_words,
    )?;
    let num_events_i = require_i32("bts_apply_noise_masks_meas_major", "num_events", num_events)?;
    let block_threads = noise_block_threads(chunk_s_words);
    let batch_blocks = div_ceil_grid(
        "bts_apply_noise_masks_meas_major",
        "chunk_s_words",
        chunk_s_words,
        block_threads,
    )?;
    let num_events_grid =
        require_u32("bts_apply_noise_masks_meas_major", "num_events", num_events)?;
    let cfg = LaunchConfig {
        grid_dim: (num_events_grid, batch_blocks, 1),
        block_dim: (block_threads, 1, 1),
        shared_mem_bytes: 0,
    };

    let mut builder = stream.launch_builder(&func);
    builder
        .arg(meas_major.raw_mut())
        .arg(&num_meas_i)
        .arg(&s_words_i)
        .arg(&word_offset_i)
        .arg(x_row_offsets.raw())
        .arg(x_row_indices.raw())
        .arg(z_row_offsets.raw())
        .arg(z_row_indices.raw())
        .arg(x_masks.raw())
        .arg(z_masks.raw())
        .arg(&chunk_s_words_i)
        .arg(&num_events_i);
    // SAFETY: each thread owns one (event, batch) pair and uses atomic XOR
    // updates when multiple events touch the same measurement row.
    unsafe {
        builder
            .launch(cfg)
            .map_err(|e| driver_err("bts_apply_noise_masks_meas_major", e))?;
    }
    Ok(())
}

/// Sample `num_shots` BTS shots in shot-major layout, `num_shots * m_words` words,
/// transposing on device instead of through the host `into_shot_major_data()`.
pub(crate) fn launch_bts_sample_shot_major_host(
    ctx: &GpuContext,
    rng: &mut Xoshiro256PlusPlus,
    rank: usize,
    num_shots: usize,
    cache: &mut GpuBtsCache,
) -> Result<Vec<u64>> {
    let num_meas = cache.num_meas;
    let m_words = num_meas.div_ceil(64);
    let s_words = num_shots.div_ceil(64);
    let shot_major_len = num_shots * m_words;
    if num_meas == 0 || num_shots == 0 || rank == 0 {
        return Ok(vec![0u64; shot_major_len]);
    }

    let meas_major = launch_bts_sample_device(ctx, rng, rank, num_shots, cache)?;
    let shot_major =
        launch_bts_transpose_meas_to_shot(ctx, &meas_major, num_meas, num_shots, s_words, m_words)?;

    let mut host = vec![0u64; shot_major_len];
    shot_major.copy_to_host(ctx.device(), &mut host)?;
    Ok(host)
}

pub(crate) fn count_meas_major_marginals(
    ctx: &GpuContext,
    meas_major: &GpuBuffer<u64>,
    num_meas: usize,
    num_shots: usize,
    s_words: usize,
) -> Result<Vec<u64>> {
    if num_meas == 0 || num_shots == 0 {
        return Ok(vec![0u64; num_meas]);
    }

    let device = ctx.device();
    let (stream, func) = stream_and_fn(ctx, "bts_popcount_rows")?;
    let mut row_counts = GpuBuffer::<u64>::alloc_zeros(device, num_meas.max(1))?;
    let mut host_counts = vec![0u64; num_meas];

    let num_meas_i = require_i32("bts_popcount_rows", "num_meas", num_meas)?;
    let s_words_i = require_i32("bts_popcount_rows", "s_words", s_words)?;
    let num_meas_grid = require_u32("bts_popcount_rows", "num_meas", num_meas)?;
    let tail_mask_u64 = shot_tail_mask(num_shots);
    let cfg = linear_cfg(POPCOUNT_BLOCK_SIZE, num_meas_grid);

    let mut builder = stream.launch_builder(&func);
    builder
        .arg(meas_major.raw())
        .arg(&num_meas_i)
        .arg(&s_words_i)
        .arg(&tail_mask_u64)
        .arg(row_counts.raw_mut());
    // SAFETY: one block owns one measurement row and writes one unique count.
    unsafe {
        builder
            .launch(cfg)
            .map_err(|e| driver_err("bts_popcount_rows", e))?;
    }

    row_counts.copy_to_host(device, &mut host_counts)?;
    Ok(host_counts)
}

fn count_used_slots(
    ctx: &GpuContext,
    slot_states: &GpuBuffer<u32>,
    table_slots: usize,
) -> Result<usize> {
    if table_slots == 0 {
        return Ok(0);
    }

    let device = ctx.device();
    let (stream, func) = stream_and_fn(ctx, "bts_count_used_slots")?;
    let mut used_out = GpuBuffer::<u32>::alloc_zeros(device, 1)?;
    let table_capacity_i = require_i32("bts_count_used_slots", "table_slots", table_slots)?;
    let blocks = div_ceil_grid(
        "bts_count_used_slots",
        "table_slots",
        table_slots,
        COUNT_HASH_BLOCK_SIZE,
    )?;
    let cfg = linear_cfg(COUNT_HASH_BLOCK_SIZE, blocks);

    let mut builder = stream.launch_builder(&func);
    builder
        .arg(slot_states.raw())
        .arg(&table_capacity_i)
        .arg(used_out.raw_mut());
    // SAFETY: each block scans a disjoint subset of indices and contributes
    // its local sum through one atomic add into `used_out`.
    unsafe {
        builder
            .launch(cfg)
            .map_err(|e| driver_err("bts_count_used_slots", e))?;
    }

    let mut used_host = [0u32; 1];
    used_out.copy_to_host(device, &mut used_host)?;
    Ok(used_host[0] as usize)
}

/// Open-addressed device hash table the count kernels fill: `slots` entries of `m_words`
/// key words, a count, and a state word each, plus a one-word overflow flag.
struct CountTable {
    slots: usize,
    keys: GpuBuffer<u64>,
    counts: GpuBuffer<u64>,
    states: GpuBuffer<u32>,
    overflow: GpuBuffer<u32>,
}

impl CountTable {
    /// `None` when [`plan_count_table_slots`] finds no table size that fits.
    fn alloc(
        ctx: &GpuContext,
        num_shots: usize,
        m_words: usize,
        rank: usize,
    ) -> Result<Option<Self>> {
        let Some(slots) = plan_count_table_slots(ctx, num_shots, m_words, rank)? else {
            return Ok(None);
        };
        let device = ctx.device();
        Ok(Some(Self {
            slots,
            keys: GpuBuffer::<u64>::alloc_zeros(device, (slots * m_words).max(1))?,
            counts: GpuBuffer::<u64>::alloc_zeros(device, slots.max(1))?,
            states: GpuBuffer::<u32>::alloc_zeros(device, slots.max(1))?,
            overflow: GpuBuffer::<u32>::alloc_zeros(device, 1)?,
        }))
    }

    /// Compact the filled table on device; `None` when the fill overflowed it.
    fn finish(
        self,
        ctx: &GpuContext,
        m_words: usize,
        raw_transfer_bytes: usize,
    ) -> Result<Option<std::collections::HashMap<Vec<u64>, u64>>> {
        let mut overflow_host = [0u32; 1];
        self.overflow
            .copy_to_host(ctx.device(), &mut overflow_host)?;
        if overflow_host[0] != 0 {
            return Ok(None);
        }
        compact_count_table(
            ctx,
            &self.keys,
            &self.counts,
            &self.states,
            self.slots,
            m_words,
            raw_transfer_bytes,
        )
    }
}

fn compact_count_table(
    ctx: &GpuContext,
    slot_keys: &GpuBuffer<u64>,
    slot_counts: &GpuBuffer<u64>,
    slot_states: &GpuBuffer<u32>,
    table_slots: usize,
    m_words: usize,
    raw_transfer_bytes: usize,
) -> Result<Option<std::collections::HashMap<Vec<u64>, u64>>> {
    let used_slots = count_used_slots(ctx, slot_states, table_slots)?;
    if used_slots == 0 {
        return Ok(Some(std::collections::HashMap::new()));
    }

    let compact_transfer_bytes = used_slots
        .saturating_mul(m_words + 1)
        .saturating_mul(std::mem::size_of::<u64>());
    if compact_transfer_bytes >= raw_transfer_bytes {
        return Ok(None);
    }

    let device = ctx.device();
    let stream = device.stream()?;
    let mut out_keys = GpuBuffer::<u64>::alloc_zeros(device, (used_slots * m_words).max(1))?;
    let mut out_counts = GpuBuffer::<u64>::alloc_zeros(device, used_slots.max(1))?;
    let mut out_len = GpuBuffer::<u32>::alloc_zeros(device, 1)?;
    let compact_func = device.function("bts_compact_counts_upto8")?;
    let m_words_i = require_i32("bts_compact_counts_upto8", "m_words", m_words)?;
    let table_capacity_i = require_i32("bts_compact_counts_upto8", "table_slots", table_slots)?;
    let compact_blocks = div_ceil_grid(
        "bts_compact_counts_upto8",
        "table_slots",
        table_slots,
        COUNT_COMPACT_BLOCK_SIZE,
    )?;
    let compact_cfg = linear_cfg(COUNT_COMPACT_BLOCK_SIZE, compact_blocks);

    let mut compact_builder = stream.launch_builder(&compact_func);
    compact_builder
        .arg(slot_states.raw())
        .arg(slot_keys.raw())
        .arg(slot_counts.raw())
        .arg(&m_words_i)
        .arg(&table_capacity_i)
        .arg(out_keys.raw_mut())
        .arg(out_counts.raw_mut())
        .arg(out_len.raw_mut());
    // SAFETY: each occupied slot compacts into a unique output index obtained
    // from an atomic counter.
    unsafe {
        compact_builder
            .launch(compact_cfg)
            .map_err(|e| driver_err("bts_compact_counts_upto8", e))?;
    }

    let mut out_len_host = [0u32; 1];
    out_len.copy_to_host(device, &mut out_len_host)?;
    let used_len = out_len_host[0] as usize;
    if used_len == 0 {
        return Ok(Some(std::collections::HashMap::new()));
    }

    let mut keys_host = vec![0u64; used_len * m_words];
    let mut counts_host = vec![0u64; used_len];
    out_keys.copy_to_host(device, &mut keys_host)?;
    out_counts.copy_to_host(device, &mut counts_host)?;

    let mut counts = std::collections::HashMap::with_capacity(used_len);
    for idx in 0..used_len {
        let key = keys_host[idx * m_words..(idx + 1) * m_words].to_vec();
        counts.insert(key, counts_host[idx]);
    }
    Ok(Some(counts))
}

/// Histogram shot-major records via an open-addressed device hash table.
/// `Ok(None)` defers to the host path (transfer heuristic, table sizing, or
/// table overflow).
pub(crate) fn try_count_shot_major(
    ctx: &GpuContext,
    shot_major: &GpuBuffer<u64>,
    num_shots: usize,
    m_words: usize,
    rank: usize,
    raw_transfer_bytes: usize,
) -> Result<Option<std::collections::HashMap<Vec<u64>, u64>>> {
    if num_shots == 0 || m_words == 0 {
        return Ok(Some(std::collections::HashMap::new()));
    }
    if !should_try_device_count(num_shots, m_words, rank, raw_transfer_bytes) {
        return Ok(None);
    }

    let Some(mut table) = CountTable::alloc(ctx, num_shots, m_words, rank)? else {
        return Ok(None);
    };

    let device = ctx.device();
    let stream = device.stream()?;
    let func = device.function("bts_count_shot_major_upto8")?;
    let num_shots_i = require_i32("bts_count_shot_major_upto8", "num_shots", num_shots)?;
    let m_words_i = require_i32("bts_count_shot_major_upto8", "m_words", m_words)?;
    let table_mask_u32 =
        require_u32("bts_count_shot_major_upto8", "table_slots", table.slots)?.saturating_sub(1);
    let blocks = div_ceil_grid(
        "bts_count_shot_major_upto8",
        "num_shots",
        num_shots,
        COUNT_HASH_BLOCK_SIZE,
    )?;
    let cfg = linear_cfg(COUNT_HASH_BLOCK_SIZE, blocks);

    let mut builder = stream.launch_builder(&func);
    builder
        .arg(shot_major.raw())
        .arg(&num_shots_i)
        .arg(&m_words_i)
        .arg(table.keys.raw_mut())
        .arg(table.counts.raw_mut())
        .arg(table.states.raw_mut())
        .arg(&table_mask_u32)
        .arg(table.overflow.raw());
    // SAFETY: each thread handles at most one shot-major outcome and uses
    // atomics to claim or update hash slots.
    unsafe {
        builder
            .launch(cfg)
            .map_err(|e| driver_err("bts_count_shot_major_upto8", e))?;
    }

    table.finish(ctx, m_words, raw_transfer_bytes)
}

/// Transpose to shot-major on device and count there, falling back to the
/// direct meas-major count kernel. `Ok(None)` defers to the host path.
pub(crate) fn try_count_meas_major(
    ctx: &GpuContext,
    meas_major: &GpuBuffer<u64>,
    num_meas: usize,
    num_shots: usize,
    m_words: usize,
    s_words: usize,
    rank: usize,
) -> Result<Option<std::collections::HashMap<Vec<u64>, u64>>> {
    if num_shots == 0 || num_meas == 0 {
        return Ok(Some(std::collections::HashMap::new()));
    }
    if m_words == 0 || m_words > GPU_COUNTS_MAX_WORDS {
        return Ok(None);
    }

    let raw_transfer_bytes = count_raw_transfer_bytes(num_meas.saturating_mul(s_words));
    if !should_try_device_count(num_shots, m_words, rank, raw_transfer_bytes) {
        return Ok(None);
    }

    if let Ok(shot_major) =
        launch_bts_transpose_meas_to_shot(ctx, meas_major, num_meas, num_shots, s_words, m_words)
        && let Some(counts) = try_count_shot_major(
            ctx,
            &shot_major,
            num_shots,
            m_words,
            rank,
            raw_transfer_bytes,
        )?
    {
        return Ok(Some(counts));
    }

    let Some(mut table) = CountTable::alloc(ctx, num_shots, m_words, rank)? else {
        return Ok(None);
    };

    let device = ctx.device();
    let stream = device.stream()?;
    let func = device.function("bts_count_meas_major_upto8")?;
    let num_meas_i = require_i32("bts_count_meas_major_upto8", "num_meas", num_meas)?;
    let num_shots_i = require_i32("bts_count_meas_major_upto8", "num_shots", num_shots)?;
    let s_words_i = require_i32("bts_count_meas_major_upto8", "s_words", s_words)?;
    let m_words_i = require_i32("bts_count_meas_major_upto8", "m_words", m_words)?;
    let table_mask_u32 =
        require_u32("bts_count_meas_major_upto8", "table_slots", table.slots)?.saturating_sub(1);
    let batches = div_ceil_grid(
        "bts_count_meas_major_upto8",
        "num_shots",
        num_shots,
        COUNT_MEAS_BATCH_SIZE,
    )?;
    let cfg = linear_cfg(COUNT_MEAS_BATCH_SIZE, batches.max(1));

    let mut builder = stream.launch_builder(&func);
    builder
        .arg(meas_major.raw())
        .arg(&num_meas_i)
        .arg(&num_shots_i)
        .arg(&s_words_i)
        .arg(&m_words_i)
        .arg(table.keys.raw_mut())
        .arg(table.counts.raw_mut())
        .arg(table.states.raw_mut())
        .arg(&table_mask_u32)
        .arg(table.overflow.raw());
    // SAFETY: each thread handles at most one logical shot and uses atomics to
    // claim or update hash slots.
    unsafe {
        builder
            .launch(cfg)
            .map_err(|e| driver_err("bts_count_meas_major_upto8", e))?;
    }

    table.finish(ctx, m_words, raw_transfer_bytes)
}