tune.rs

//! Bounded, opt-in kernel selection. See [`crate::Session::tune_with`].
//!
//! The search space is deliberately small: scalar and native-f32 cooperative
//! tiles for unpacked dense matmuls, typed subgroup/workgroup tiles and K
//! partitions, shape-specialized scalar convolutions,
//! workgroup geometry and cross-lane reduction for the K-split GEMV family, and
//! cached-attention split/combine sequences. Precision remains unchanged.
//! Measurements use synthetic, private scratch, not a live training step.
//! Explicit split-K probes measure complete sequences without installing them.
//! Scalar tile and GEMV shape changes leave packed-weight decoding intact, so
//! reduced-storage kernels participate without changing their arithmetic.

use crate::codegen::CoopConfig;
use crate::compile::{Dispatch, ShaderEntry};
use serde::{Deserialize, Serialize};
use std::time::Duration;

/// Complete key within the supported contiguous-row-major, non-aliasing domain.
/// No winner is transferred between shapes, directions, or memory placements.
#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct TuneClass {
    pub shader: ShaderEntry,
    pub m: u32,
    pub n: u32,
    pub k: u32,
    /// Physical NCHW convolution contract; absent for dense matrices.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub conv2d: Option<TuneConv2d>,
    pub requires_full_precision: bool,
    /// Storage format of B. Part of the key because it changes both the
    /// kernel and the bytes behind the binding, so no winner transfers
    /// between an f32 weight and a packed one of the same extents.
    #[serde(default)]
    pub weight_format: crate::compile::WeightFormat,
    /// The dispatch quantizes its activation and uses integer dot products.
    ///
    /// Part of the key even though the weight format is already here,
    /// because the format no longer identifies the kernel: an ordinary Q4_0
    /// GEMV and an activation-quantized one of the same extents read the same
    /// bytes and compute different things. Coalescing them would measure the
    /// first member's variants and install the winner on both, leaving the
    /// other with a shape it has no pipeline for.
    #[serde(default)]
    pub gemv_int_dot: bool,
    /// RmsNorm is folded into this GEMV. A different kernel (second
    /// reduction + extra binding), so winners do not transfer from the
    /// unfused class of the same extents.
    #[serde(default)]
    pub gemv_rmsnorm: bool,
    /// `f32::to_bits` of the folded norm's epsilon. Ignored unless
    /// [`Self::gemv_rmsnorm`].
    #[serde(default)]
    pub gemv_rmsnorm_eps_bits: u32,
    /// The output holds K partitions that a following SumRows adds. Such a
    /// product keeps its reduction: candidates write two or more partitions
    /// into the partial buffer, whose size is the output binding's.
    #[serde(default)]
    pub split_k: bool,
    /// Operand, multiplication and output types of a checked typed product.
    /// Part of the key: they change the kernel and the bytes behind each
    /// binding, so no winner transfers between precisions.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub matrix_types: Option<crate::codegen::MatrixTypes>,
    /// Automatic F32 storage with F16 multiplication, already permitted by
    /// the session policy. Separate from explicit types for cache fallback.
    #[serde(default)]
    pub cooperative_f16: bool,
    /// Placement of A, B, addend (false if absent), and output, respectively.
    pub device_local: [bool; 4],
    /// Declared bytes for A, B, optional addend, output, in binding order.
    /// Candidates cannot borrow unused capacity from an aliased allocation.
    pub binding_bytes: Vec<usize>,
}

/// Exact convolution shape, including both spatial padding dimensions.
/// Small and large scalar tiles share this key, but directions never do.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct TuneConv2d {
    pub batch: u32,
    pub in_channels: u32,
    pub in_h: u32,
    pub in_w: u32,
    pub out_channels: u32,
    pub kernel_h: u32,
    pub kernel_w: u32,
    pub stride: u32,
    pub padding_h: u32,
    pub out_h: u32,
    pub out_w: u32,
    pub padding_w: u32,
}

impl TuneConv2d {
    fn from_params(params: &[u32]) -> Option<Self> {
        let &[
            batch,
            in_channels,
            in_h,
            in_w,
            out_channels,
            kernel_h,
            kernel_w,
            stride,
            padding_h,
            out_h,
            out_w,
            padding_w,
        ] = params
        else {
            return None;
        };
        let shape = Self {
            batch,
            in_channels,
            in_h,
            in_w,
            out_channels,
            kernel_h,
            kernel_w,
            stride,
            padding_h,
            out_h,
            out_w,
            padding_w,
        };
        if [
            batch,
            in_channels,
            in_h,
            in_w,
            out_channels,
            kernel_h,
            kernel_w,
            stride,
            out_h,
            out_w,
        ]
        .contains(&0)
            || stride > i32::MAX as u32
        {
            return None;
        }
        for (input, kernel, padding, output) in [
            (in_h, kernel_h, padding_h, out_h),
            (in_w, kernel_w, padding_w, out_w),
        ] {
            let padded = input.checked_add(padding.checked_mul(2)?)?;
            // Both gather directions convert coordinates to signed shader integers.
            if padded > i32::MAX as u32
                || kernel > i32::MAX as u32
                || padded
                    .checked_sub(kernel)?
                    .checked_div(stride)?
                    .checked_add(1)?
                    != output
            {
                return None;
            }
        }
        shape.elements()?;
        Some(shape)
    }

    /// Input, weight, upstream gradient: physical elements, not an im2col matrix.
    fn elements(self) -> Option<[u32; 3]> {
        Some([
            self.batch
                .checked_mul(self.in_channels)?
                .checked_mul(self.in_h)?
                .checked_mul(self.in_w)?,
            self.out_channels
                .checked_mul(self.in_channels)?
                .checked_mul(self.kernel_h)?
                .checked_mul(self.kernel_w)?,
            self.batch
                .checked_mul(self.out_channels)?
                .checked_mul(self.out_h)?
                .checked_mul(self.out_w)?,
        ])
    }
}

/// The GEMV group a shader entry belongs to, if any.
pub(crate) fn gemv_group(entry: &ShaderEntry) -> Option<crate::codegen::ShaderGroup> {
    use crate::codegen::ShaderGroup;
    match *entry {
        ShaderEntry::MatMulGemv => Some(ShaderGroup::MatMulGemv),
        ShaderEntry::MatMulGemvAdd => Some(ShaderGroup::MatMulGemvAdd),
        ShaderEntry::MatMulGemvBT => Some(ShaderGroup::MatMulGemvBT),
        ShaderEntry::MatMulGemvBTAdd => Some(ShaderGroup::MatMulGemvBTAdd),
        _ => None,
    }
}

/// Implementations with identical bindings and logical extents.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum MatmulTile {
    Tile16,
    Tile32,
    Tile64,
    Scalar(crate::codegen::ScalarMatmulShape),
    /// Scalar convolution with immutable parameters and native integer division.
    /// Spatial tiles, K-stage sizes and column assignment are measured.
    SpecializedConv {
        tile_size: u32,
        k_tile: u32,
        /// Give adjacent lanes adjacent columns of the register tile.
        #[serde(default)]
        interleave_columns: bool,
    },
    /// Hardware-native f32 operands/accumulators, never f16 staging.
    CooperativeF32 {
        tile_size: u32,
    },
    CooperativeF16 {
        tile_size: u32,
    },
    CooperativeWorkgroup {
        shape: crate::codegen::WorkgroupMatmulShape,
        splits: u32,
    },
    /// Aligned native 16x16 f32 products sharing a larger workgroup tile.
    /// More than one K partition writes partials that a following SumRows
    /// adds, which only a split product has.
    CooperativeTiled {
        shape: crate::codegen::CooperativeMatmulShape,
        splits: u32,
    },
    /// Checked typed cooperative tiles at a K partition count. The class key
    /// holds the types. More than one partition writes F32 partials that a
    /// following SumRows adds, which only a split product has.
    CooperativeChecked {
        splits: u32,
    },
    /// A K-split GEMV at one workgroup width and cross-lane reduction.
    ///
    /// Keeps the shader entry and buffer layout. Width and reduction apply
    /// to every storage format; contiguous row grouping also changes the
    /// workgroup count for transposed-B kernels.
    Gemv(crate::codegen::GemvShape),
}

fn conv_launched_tile(shader: &ShaderEntry) -> Option<u32> {
    match *shader {
        ShaderEntry::Conv2dGemm16
        | ShaderEntry::Conv2dGradInputGemm16
        | ShaderEntry::Conv2dGradWeightGemm16
        | ShaderEntry::Conv2dGradWeightGemmSplit16 => Some(16),
        ShaderEntry::Conv2dGemmSmall
        | ShaderEntry::Conv2dGradInputGemmSmall
        | ShaderEntry::Conv2dGradWeightGemmSmall
        | ShaderEntry::Conv2dGradWeightGemmSplitSmall => Some(32),
        ShaderEntry::Conv2dGemm
        | ShaderEntry::Conv2dGradInputGemm
        | ShaderEntry::Conv2dGradWeightGemm
        | ShaderEntry::Conv2dGradWeightGemmSplit => Some(64),
        _ => None,
    }
}

fn conv_tile_entry(entry: &ShaderEntry, tile: u32) -> Option<ShaderEntry> {
    let split = matches!(
        *entry,
        ShaderEntry::Conv2dGradWeightGemmSplit
            | ShaderEntry::Conv2dGradWeightGemmSplitSmall
            | ShaderEntry::Conv2dGradWeightGemmSplit16
    );
    enum Family {
        Basic,
        Input,
        Weight,
    }
    let family = match *entry {
        ShaderEntry::Conv2dGemm | ShaderEntry::Conv2dGemmSmall | ShaderEntry::Conv2dGemm16 => {
            Family::Basic
        }
        ShaderEntry::Conv2dGradInputGemm
        | ShaderEntry::Conv2dGradInputGemmSmall
        | ShaderEntry::Conv2dGradInputGemm16 => Family::Input,
        ShaderEntry::Conv2dGradWeightGemm
        | ShaderEntry::Conv2dGradWeightGemmSmall
        | ShaderEntry::Conv2dGradWeightGemm16
        | ShaderEntry::Conv2dGradWeightGemmSplit
        | ShaderEntry::Conv2dGradWeightGemmSplitSmall
        | ShaderEntry::Conv2dGradWeightGemmSplit16 => Family::Weight,
        _ => return None,
    };
    Some(match (family, split, tile) {
        (Family::Basic, _, 16) => ShaderEntry::Conv2dGemm16,
        (Family::Basic, _, 32) => ShaderEntry::Conv2dGemmSmall,
        (Family::Basic, _, _) => ShaderEntry::Conv2dGemm,
        (Family::Input, _, 16) => ShaderEntry::Conv2dGradInputGemm16,
        (Family::Input, _, 32) => ShaderEntry::Conv2dGradInputGemmSmall,
        (Family::Input, _, _) => ShaderEntry::Conv2dGradInputGemm,
        (Family::Weight, false, 16) => ShaderEntry::Conv2dGradWeightGemm16,
        (Family::Weight, false, 32) => ShaderEntry::Conv2dGradWeightGemmSmall,
        (Family::Weight, false, _) => ShaderEntry::Conv2dGradWeightGemm,
        (Family::Weight, true, 16) => ShaderEntry::Conv2dGradWeightGemmSplit16,
        (Family::Weight, true, 32) => ShaderEntry::Conv2dGradWeightGemmSplitSmall,
        (Family::Weight, true, _) => ShaderEntry::Conv2dGradWeightGemmSplit,
    })
}

impl MatmulTile {
    pub(crate) fn native_cooperative(config: Option<&CoopConfig>) -> Option<Self> {
        let config = config?;
        if config.use_f16_input || config.compensated || !matches!(config.tile_size, 8 | 16) {
            return None;
        }
        Some(Self::CooperativeF32 {
            tile_size: config.tile_size,
        })
    }

    pub(crate) fn selected(dispatch: &Dispatch, config: Option<&CoopConfig>) -> Option<Self> {
        use crate::compile::Kernel;
        let small = dispatch.use_small_tiles();
        if let Some(group) = gemv_group(&dispatch.shader) {
            match dispatch.kernel {
                Kernel::Default => Some(Self::Gemv(crate::codegen::GemvShape::initial(group))),
                Kernel::Gemv { shape, .. } => Some(Self::Gemv(shape)),
                _ => None,
            }
        } else if let Some(launched) = conv_launched_tile(&dispatch.shader) {
            match dispatch.kernel {
                Kernel::Default | Kernel::SmallTile => Some(match launched {
                    16 => Self::Tile16,
                    32 => Self::Tile32,
                    _ => Self::Tile64,
                }),
                Kernel::SpecializedConv {
                    k_tile,
                    interleave_columns,
                } => Some(Self::SpecializedConv {
                    tile_size: launched,
                    k_tile,
                    interleave_columns,
                }),
                Kernel::Cooperative => Self::native_cooperative(config),
                _ => None,
            }
        } else {
            match dispatch.kernel {
                Kernel::Default | Kernel::SmallTile => {
                    Some(if small { Self::Tile32 } else { Self::Tile64 })
                }
                Kernel::ScalarMatmul(shape) => Some(Self::Scalar(shape)),
                Kernel::SpecializedConv {
                    k_tile,
                    interleave_columns,
                } => Some(Self::SpecializedConv {
                    tile_size: if small { 32 } else { 64 },
                    k_tile,
                    interleave_columns,
                }),
                Kernel::Cooperative => config
                    .filter(|c| c.use_f16_input && !c.compensated)
                    .map(|c| Self::CooperativeF16 {
                        tile_size: c.tile_size,
                    })
                    .or_else(|| Self::native_cooperative(config)),
                Kernel::CooperativeWorkgroup { shape, splits, .. } => {
                    Some(Self::CooperativeWorkgroup { shape, splits })
                }
                Kernel::CooperativeTiled { shape, splits }
                    if config.is_some_and(|c| {
                        c.tile_size == 16 && !c.use_f16_input && !c.compensated
                    }) =>
                {
                    Some(Self::CooperativeTiled { shape, splits })
                }
                Kernel::CooperativeTyped { splits, .. } => {
                    Some(Self::CooperativeChecked { splits })
                }
                _ => None,
            }
        }
    }

    pub(crate) fn coop_config(self) -> Option<CoopConfig> {
        match self {
            Self::CooperativeF16 { tile_size } => Some(CoopConfig {
                tile_size,
                use_f16_input: true,
                compensated: false,
            }),
            Self::CooperativeTiled { .. } => Some(CoopConfig {
                tile_size: 16,
                use_f16_input: false,
                compensated: false,
            }),
            Self::CooperativeF32 { tile_size } => Some(CoopConfig {
                tile_size,
                use_f16_input: false,
                compensated: false,
            }),
            _ => None,
        }
    }

    pub(crate) fn apply(self, dispatch: &mut Dispatch, class: &TuneClass) {
        self.configure(dispatch);
        dispatch.workgroups = self.workgroups(class);
    }

    /// K partitions a split tile writes for its SumRows; `None` when unsplit.
    pub(crate) fn partitions(self) -> Option<u32> {
        match self {
            Self::CooperativeTiled { splits, .. }
            | Self::CooperativeWorkgroup { splits, .. }
            | Self::CooperativeChecked { splits } => Some(splits),
            _ => None,
        }
    }

    /// Kernels whose measured-to-predicted ratio is shared across classes.
    pub(crate) fn family(self) -> TileFamily {
        match self {
            Self::Tile16
            | Self::Tile32
            | Self::Tile64
            | Self::Scalar(_)
            | Self::SpecializedConv { .. } => TileFamily::Scalar,
            Self::CooperativeF32 { .. }
            | Self::CooperativeF16 { .. }
            | Self::CooperativeWorkgroup { .. }
            | Self::CooperativeTiled { .. }
            | Self::CooperativeChecked { .. } => TileFamily::Cooperative,
            Self::Gemv(_) => TileFamily::Gemv,
        }
    }

    pub(crate) fn configure(self, dispatch: &mut Dispatch) {
        if let Self::CooperativeWorkgroup { shape, splits } = self {
            let types = match dispatch.kernel {
                crate::compile::Kernel::CooperativeTyped { types, .. } => Some(types),
                crate::compile::Kernel::CooperativeWorkgroup { types, .. } => types,
                _ => None,
            };
            dispatch.kernel = crate::compile::Kernel::CooperativeWorkgroup {
                shape,
                types,
                splits,
            };
            return;
        }
        if let Self::CooperativeChecked { splits } = self {
            let types = match dispatch.kernel {
                crate::compile::Kernel::CooperativeTyped { types, .. }
                | crate::compile::Kernel::CooperativeWorkgroup {
                    types: Some(types), ..
                } => types,
                _ => unreachable!("checked tiles apply to typed products"),
            };
            dispatch.kernel = crate::compile::Kernel::CooperativeTyped { types, splits };
            return;
        }
        if let Self::Gemv(shape) = self {
            dispatch.kernel = crate::compile::Kernel::Gemv {
                shape,
                integer_dot: dispatch.gemv_int_dot(),
            };
            return;
        }
        dispatch.shader = self.shader(&dispatch.shader);
        dispatch.kernel = match self {
            Self::Tile16 => crate::compile::Kernel::Default,
            Self::Tile32
                if !matches!(
                    dispatch.shader,
                    ShaderEntry::Conv2dGemmSmall
                        | ShaderEntry::Conv2dGradInputGemmSmall
                        | ShaderEntry::Conv2dGradWeightGemmSmall
                ) =>
            {
                crate::compile::Kernel::SmallTile
            }
            Self::Tile32 | Self::Tile64 => crate::compile::Kernel::Default,
            Self::Scalar(shape) => crate::compile::Kernel::ScalarMatmul(shape),
            Self::CooperativeF32 { .. } | Self::CooperativeF16 { .. } => {
                crate::compile::Kernel::Cooperative
            }
            Self::CooperativeTiled { shape, splits } => {
                crate::compile::Kernel::CooperativeTiled { shape, splits }
            }
            Self::SpecializedConv {
                k_tile,
                interleave_columns,
                ..
            } => crate::compile::Kernel::SpecializedConv {
                k_tile,
                interleave_columns,
            },
            Self::Gemv(_) | Self::CooperativeChecked { .. } | Self::CooperativeWorkgroup { .. } => {
                unreachable!()
            }
        };
    }

    pub(crate) fn shader(self, entry: &ShaderEntry) -> ShaderEntry {
        let tile = match self {
            Self::Tile16 | Self::SpecializedConv { tile_size: 16, .. } => 16,
            Self::Tile32 | Self::SpecializedConv { tile_size: 32, .. } => 32,
            Self::Tile64 | Self::SpecializedConv { tile_size: 64, .. } => 64,
            Self::SpecializedConv { tile_size, .. } => tile_size,
            _ => return entry.clone(),
        };
        conv_tile_entry(entry, tile).unwrap_or_else(|| entry.clone())
    }

    pub(crate) fn workgroups(self, class: &TuneClass) -> [u32; 3] {
        let tile = match self {
            Self::Gemv(shape) => return class.gemv_workgroups(shape),
            Self::Tile16 => 16,
            Self::Tile32 => 32,
            Self::Tile64 => 64,
            Self::Scalar(shape) => {
                return [
                    class.n.div_ceil(shape.cols()),
                    class.m.div_ceil(shape.rows()),
                    class.batch_dispatches(),
                ];
            }
            Self::SpecializedConv { tile_size, .. } => tile_size,
            Self::CooperativeTiled { shape, splits } => {
                return [class.m / 64, class.n / shape.columns, splits];
            }
            Self::CooperativeWorkgroup { shape, splits } => {
                return [class.m / shape.rows, class.n / shape.columns, splits];
            }
            // Checked tiles use X for rows, Y for columns, over 32x32 blocks.
            Self::CooperativeChecked { splits } => {
                return [class.m.div_ceil(32), class.n.div_ceil(32), splits];
            }
            Self::CooperativeF32 { .. } | Self::CooperativeF16 { .. } => {
                let tile = self.coop_config().unwrap().matmul_output_tile();
                // Cooperative tiles use X for rows, Y for columns.
                return [class.m.div_ceil(tile), class.n.div_ceil(tile), 1];
            }
        };
        // Scalar tiled kernels use X for columns, Y for rows. Derive exact
        // geometry; doubling rounded 64-tile counts overdispatches edges.
        [
            class.n.div_ceil(tile),
            class.m.div_ceil(tile),
            class.batch_dispatches(),
        ]
    }

    pub(crate) fn buffer_sizes(self, class: &TuneClass) -> Option<Vec<usize>> {
        if let Self::CooperativeWorkgroup { shape, splits } = self {
            let types = class.effective_matrix_types()?;
            if !matches!(types.multiply, crate::DType::F16 | crate::DType::BF16)
                || class.conv2d.is_some()
                || !class.shader.is_matmul()
                || class.weight_format != crate::compile::WeightFormat::F32
                || !shape.fits(class.m, class.n, class.k)
                || splits == 0
                || splits > class.k / shape.k_stage
                || (splits > 1) != class.split_k
                || (splits > 1 && (class.has_addend() || types.output != crate::DType::F32))
            {
                return None;
            }
            let mut sizes = class.buffer_sizes()?;
            *sizes.last_mut()? = sizes.last()?.checked_mul(splits as usize)?;
            return crate::compile::workgroups_within_portable_limits(self.workgroups(class))
                .then_some(sizes);
        }
        if matches!(self, Self::CooperativeF16 { .. }) != class.cooperative_f16 {
            return None;
        }
        if let Some(types) = class.matrix_types {
            // Workgroup candidates were handled above; subgroup candidates
            // preserve storage too, with F32 partials when split.
            let Self::CooperativeChecked { splits } = self else {
                return None;
            };
            if splits == 0
                || (splits > 1) != class.split_k
                || (splits > 1
                    && (splits > class.k.div_ceil(32)
                        || class.has_addend()
                        || types.output != crate::DType::F32))
            {
                return None;
            }
            let mut sizes = class.buffer_sizes()?;
            *sizes.last_mut()? = sizes.last()?.checked_mul(splits as usize)?;
            return crate::compile::workgroups_within_portable_limits(self.workgroups(class))
                .then_some(sizes);
        }
        if matches!(self, Self::CooperativeChecked { .. })
            || (class.split_k && !matches!(self, Self::CooperativeTiled { .. }))
        {
            return None;
        }
        if let Self::Scalar(shape) = self {
            if class.conv2d.is_some()
                || gemv_group(&class.shader).is_some()
                || class.weight_format.is_quantized()
                || !shape.legal()
            {
                return None;
            }
        }
        if let Self::Gemv(shape) = self {
            if gemv_group(&class.shader).is_none()
                || !matches!(shape.threads, 32 | 64 | 128 | 256)
                || !matches!(shape.bt_rows, 1 | 2 | 4)
                || !shape.valid_columns()
                || (shape.column_groups > 1
                    && (!matches!(
                        class.shader,
                        ShaderEntry::MatMulGemv | ShaderEntry::MatMulGemvAdd
                    ) || !matches!(
                        class.weight_format,
                        crate::compile::WeightFormat::F32 | crate::compile::WeightFormat::F16
                    )))
            {
                return None;
            }
            let sizes = class.buffer_sizes()?;
            // The default forward GEMV workgroup count is N/4.
            // Vocab-width decode (N=262144) needs 65536 groups — one past the
            // Vulkan *minimum* maxComputeWorkGroupCount. The live plan already
            // dispatches that on devices that advertise more; applying the
            // portable 65535 cap here would drop the kernel that dominates
            // Gemma 4 logits from `tune_with`.
            if self.workgroups(class).contains(&0) {
                return None;
            }
            return Some(sizes);
        }
        if let Self::SpecializedConv {
            tile_size, k_tile, ..
        } = self
        {
            if class.conv2d.is_none()
                || !matches!(tile_size, 16 | 32 | 64)
                || !matches!(k_tile, 16 | 32)
            {
                return None;
            }
        }
        let mut sizes = class.buffer_sizes()?;
        if let Self::CooperativeTiled { shape, splits } = self {
            if class.conv2d.is_some()
                || !class.shader.is_matmul()
                || class.weight_format != crate::compile::WeightFormat::F32
                || !shape.fits_dimensions(class.m, class.n, class.k)
                || splits == 0
                || splits > class.k / shape.k_stage
                || (splits > 1) != class.split_k
            {
                return None;
            }
            // Each partition writes a whole output of partial sums.
            *sizes.last_mut()? = sizes.last()?.checked_mul(splits as usize)?;
        }
        if let Self::CooperativeF32 { tile_size } | Self::CooperativeF16 { tile_size } = self {
            use crate::kernels::matmul;
            if class.conv2d.is_some() || !matches!(tile_size, 8 | 16) {
                return None;
            }
            let problem = matmul::Problem::plain(
                matmul::Layout::of(class.shader.shader_group())?,
                [class.m, class.n, class.k],
                class.weight_format.uses_reduced_storage(),
            );
            // The same legality session construction applies.
            let geometry = matmul::cooperative_geometry(&problem, &self.coop_config()?).ok()?;
            *sizes.last_mut()? = geometry.output_bytes;
            if geometry.pad_addend {
                sizes[2] = geometry.output_bytes;
            }
        }
        if !crate::compile::workgroups_within_portable_limits(self.workgroups(class)) {
            return None;
        }
        Some(sizes)
    }

    pub(crate) fn fits(self, class: &TuneClass) -> bool {
        let Some(sizes) = self.buffer_sizes(class) else {
            return false;
        };
        sizes.len() == class.binding_bytes.len()
            && sizes
                .iter()
                .zip(&class.binding_bytes)
                .all(|(required, available)| required <= available)
    }
}

impl TuneClass {
    pub(crate) fn effective_matrix_types(&self) -> Option<crate::codegen::MatrixTypes> {
        self.matrix_types.or_else(|| {
            self.cooperative_f16
                .then_some(crate::codegen::MatrixTypes::F32_AS_F16)
        })
    }

    pub(crate) fn from_dispatch(dispatch: &Dispatch, config: Option<&CoopConfig>) -> Option<Self> {
        let addend = matches!(
            dispatch.shader,
            ShaderEntry::FusedMatMulAdd
                | ShaderEntry::FusedMatMulATAdd
                | ShaderEntry::FusedMatMulBTAdd
                | ShaderEntry::MatMulGemvAdd
                | ShaderEntry::MatMulGemvBTAdd
        );
        if !matches!(
            dispatch.shader,
            ShaderEntry::MatMul
                | ShaderEntry::FusedMatMulAdd
                | ShaderEntry::MatMulAT
                | ShaderEntry::MatMulBT
                | ShaderEntry::FusedMatMulATAdd
                | ShaderEntry::FusedMatMulBTAdd
                | ShaderEntry::MatMulGemv
                | ShaderEntry::MatMulGemvAdd
                | ShaderEntry::MatMulGemvBT
                | ShaderEntry::MatMulGemvBTAdd
                | ShaderEntry::Conv2dGemm
                | ShaderEntry::Conv2dGemmSmall
                | ShaderEntry::Conv2dGemm16
                | ShaderEntry::Conv2dGradInputGemm
                | ShaderEntry::Conv2dGradInputGemmSmall
                | ShaderEntry::Conv2dGradInputGemm16
                | ShaderEntry::Conv2dGradWeightGemm
                | ShaderEntry::Conv2dGradWeightGemmSmall
                | ShaderEntry::Conv2dGradWeightGemm16
        ) || dispatch.use_coop_compensated()
            || (dispatch.use_coop() && dispatch.weight_format.uses_reduced_storage())
            || dispatch.horizontal_batch >= 2
            || dispatch.matmul_prologue.is_some()
            || dispatch.matmul_epilogue.is_some()
            || !dispatch.extra_outputs.is_empty()
            || dispatch.pointwise().is_some()
            || dispatch.reduction().is_some()
            || dispatch.input_buffers.len() != if addend { 3 } else { 2 }
            || dispatch.schedule_locked
        {
            return None;
        }
        // Tiled and typed products dispatch one Z layer per K partition.
        let (partitions, matrix_types) = match dispatch.kernel {
            crate::compile::Kernel::CooperativeTiled { splits, .. } => (splits, None),
            crate::compile::Kernel::CooperativeTyped { types, splits } => (splits, Some(types)),
            crate::compile::Kernel::CooperativeWorkgroup { types, splits, .. } => (splits, types),
            _ => (1, None),
        };
        let shader = match dispatch.shader {
            ShaderEntry::Conv2dGemmSmall | ShaderEntry::Conv2dGemm16 => ShaderEntry::Conv2dGemm,
            ShaderEntry::Conv2dGradInputGemmSmall | ShaderEntry::Conv2dGradInputGemm16 => {
                ShaderEntry::Conv2dGradInputGemm
            }
            ShaderEntry::Conv2dGradWeightGemmSmall | ShaderEntry::Conv2dGradWeightGemm16 => {
                ShaderEntry::Conv2dGradWeightGemm
            }
            _ => dispatch.shader.clone(),
        };
        let conv2d = if matches!(
            shader,
            ShaderEntry::Conv2dGemm
                | ShaderEntry::Conv2dGradInputGemm
                | ShaderEntry::Conv2dGradWeightGemm
        ) {
            if dispatch.use_coop()
                || dispatch.use_small_tiles()
                || dispatch.weight_format.uses_reduced_storage()
            {
                return None;
            }
            Some(TuneConv2d::from_params(&dispatch.params)?)
        } else {
            if dispatch.conv_k_tile().is_some()
                || dispatch.workgroups[2] != partitions
                || dispatch.params.len() != 4
                || dispatch.params[3] != 0
            {
                return None;
            }
            None
        };
        let (m, n, k) = if let Some(s) = conv2d {
            let kernel = s.kernel_h.checked_mul(s.kernel_w)?;
            if shader == ShaderEntry::Conv2dGemm {
                (
                    s.out_channels,
                    s.out_h.checked_mul(s.out_w)?,
                    s.in_channels.checked_mul(kernel)?,
                )
            } else if shader == ShaderEntry::Conv2dGradInputGemm {
                let n = s.in_h.checked_mul(s.in_w)?;
                let k = s.out_channels.checked_mul(kernel)?;
                (s.in_channels, n, k)
            } else {
                let n = s.in_channels.checked_mul(kernel)?;
                let spatial = s.out_h.checked_mul(s.out_w)?;
                let k = s.batch.checked_mul(spatial)?;
                (s.out_channels, n, k)
            }
        } else if matches!(
            shader,
            ShaderEntry::MatMul
                | ShaderEntry::FusedMatMulAdd
                // The forward GEMV pair carries `[m, k, n]` like the tiled
                // forward matmul; only the transposed form is `[m, n, k]`.
                | ShaderEntry::MatMulGemv
                | ShaderEntry::MatMulGemvAdd
        ) {
            (dispatch.params[0], dispatch.params[2], dispatch.params[1])
        } else {
            (dispatch.params[0], dispatch.params[1], dispatch.params[2])
        };
        // Scalar K loops advance by at most 32, including the final padded tile.
        if m == 0
            || n == 0
            || k == 0
            || k > u32::MAX - 31
            || m.div_ceil(32) > 65_535
            || n.div_ceil(32) > 65_535
        {
            return None;
        }
        let class = Self {
            shader,
            m,
            n,
            k,
            conv2d,
            requires_full_precision: dispatch.requires_full_precision,
            weight_format: dispatch.weight_format,
            gemv_int_dot: dispatch.gemv_int_dot(),
            gemv_rmsnorm: dispatch.gemv_rmsnorm.is_some(),
            gemv_rmsnorm_eps_bits: dispatch
                .gemv_rmsnorm
                .as_ref()
                .map(|rn| rn.eps_bits)
                .unwrap_or(0),
            split_k: partitions > 1,
            matrix_types,
            cooperative_f16: matches!(
                dispatch.kernel,
                crate::compile::Kernel::CooperativeWorkgroup { types: None, .. }
            ) || (dispatch.kernel == crate::compile::Kernel::Cooperative
                && config.is_some_and(|c| c.use_f16_input && !c.compensated)),
            device_local: [false; 4],
            binding_bytes: Vec::new(),
        };
        let initial = MatmulTile::selected(dispatch, config)?;
        // A packed B's byte count comes from the plan, and the caller fills
        // `binding_bytes` in after this returns, so its extents cannot be
        // checked here. `collect_classes` settles them with `fits` once the
        // bytes are known; the geometry check below still applies to every
        // class.
        let extents_derivable = !class.weight_format.uses_reduced_storage();
        ((!extents_derivable || initial.buffer_sizes(&class).is_some())
            && dispatch.workgroups == initial.workgroups(&class))
        .then_some(class)
    }

    pub(crate) fn has_addend(&self) -> bool {
        matches!(
            self.shader,
            ShaderEntry::FusedMatMulAdd
                | ShaderEntry::FusedMatMulATAdd
                | ShaderEntry::FusedMatMulBTAdd
                | ShaderEntry::MatMulGemvAdd
                | ShaderEntry::MatMulGemvBTAdd
        )
    }

    pub(crate) fn batch_dispatches(&self) -> u32 {
        if matches!(
            self.shader,
            ShaderEntry::Conv2dGemm | ShaderEntry::Conv2dGradInputGemm
        ) {
            self.conv2d.expect("convolution class").batch
        } else {
            1
        }
    }

    pub(crate) fn output_elements(&self) -> usize {
        self.m as usize * self.n as usize * self.batch_dispatches() as usize
    }

    /// Predicted milliseconds of one member dispatch running `tile`, scaled by
    /// a measured-to-predicted `correction` for the tile's family.
    pub(crate) fn predict_ms(
        &self,
        profile: &MachineProfile,
        tile: MatmulTile,
        correction: f64,
    ) -> Option<f64> {
        let flops = 2.0
            * f64::from(self.m)
            * f64::from(self.n)
            * f64::from(self.k)
            * f64::from(self.batch_dispatches());
        let bytes = self.binding_bytes.iter().sum::<usize>() as f64;
        let workgroups = tile
            .workgroups(self)
            .iter()
            .map(|&count| u64::from(count))
            .product();
        profile
            .predict_ms(flops, bytes, workgroups, None)
            .map(|ms| ms * correction)
    }

    /// A, B, optional addend, output sizes. Reject u32 shader-index overflow,
    /// not just host allocation overflow.
    pub(crate) fn buffer_sizes(&self) -> Option<Vec<usize>> {
        let bytes = |a: u32, b: u32| usize::try_from(a.checked_mul(b)?).ok()?.checked_mul(4);
        if let Some(types) = self.matrix_types {
            // Logical extents: an odd count of halfwords is not word-padded.
            let typed = |a: u32, b: u32, ty| {
                usize::try_from(a.checked_mul(b)?)
                    .ok()?
                    .checked_mul(crate::codegen::MatrixTypes::bytes(ty))
            };
            let mut sizes = vec![
                typed(self.m, self.k, types.a)?,
                typed(self.k, self.n, types.b)?,
            ];
            if self.has_addend() {
                sizes.push(bytes(self.m, self.n)?);
            }
            sizes.push(typed(self.m, self.n, types.output)?);
            return Some(sizes);
        }
        if let Some(s) = self.conv2d {
            let [input, weight, upstream] = s.elements()?;
            if self.shader == ShaderEntry::Conv2dGemm {
                return Some(vec![
                    bytes(input, 1)?,
                    bytes(weight, 1)?,
                    bytes(upstream, 1)?,
                ]);
            }
            let (b, output) = if self.shader == ShaderEntry::Conv2dGradInputGemm {
                (weight, input)
            } else {
                (input, weight)
            };
            return Some(vec![bytes(upstream, 1)?, bytes(b, 1)?, bytes(output, 1)?]);
        }
        // A packed B holds blocks, not a `k * n` f32 grid, so its byte count
        // comes from the plan rather than the extents. Everything else in the
        // class is still f32 and still derived.
        let b_bytes = if self.weight_format.uses_reduced_storage() {
            *self.binding_bytes.get(1)?
        } else {
            bytes(self.k, self.n)?
        };
        let mut sizes = vec![bytes(self.m, self.k)?, b_bytes];
        if self.gemv_rmsnorm {
            sizes.push(bytes(self.m, self.k)?);
        }
        if self.has_addend() {
            sizes.push(bytes(self.m, self.n)?);
        }
        sizes.push(bytes(self.m, self.n)?);
        Some(sizes)
    }

    /// Workgroups per output vec4, or per group of transposed B rows.
    pub(crate) fn gemv_workgroups(&self, shape: crate::codegen::GemvShape) -> [u32; 3] {
        if matches!(
            self.shader,
            ShaderEntry::MatMulGemvBT | ShaderEntry::MatMulGemvBTAdd
        ) {
            crate::compile::row_gemv_workgroups(self.n.div_ceil(shape.bt_rows))
        } else {
            crate::compile::gemv_workgroups(self.n, shape)
        }
    }

    /// A bounded tournament over scalar layouts and native f32 where it fits.
    /// Capability/precision/padding are legality;
    /// occupancy and the static large-shape veto do not exclude candidates.
    pub(crate) fn challengers(
        &self,
        initial: MatmulTile,
        config: Option<&CoopConfig>,
    ) -> Vec<MatmulTile> {
        if let MatmulTile::Gemv(shape) = initial {
            use crate::codegen::{GemvReduction, GemvShape};
            let mut out = Vec::new();
            for column_groups in [4, 8] {
                out.push(MatmulTile::Gemv(GemvShape {
                    threads: 32,
                    reduction: GemvReduction::Tree,
                    bt_rows: 1,
                    column_groups,
                }));
            }
            let rows: &[u32] = if matches!(
                self.shader,
                ShaderEntry::MatMulGemvBT | ShaderEntry::MatMulGemvBTAdd
            ) {
                &GemvShape::BT_ROWS
            } else {
                &[1]
            };
            for &bt_rows in rows {
                for reduction in [GemvReduction::Subgroup, GemvReduction::Tree] {
                    for threads in [256, 128, 64, 32] {
                        out.push(MatmulTile::Gemv(GemvShape {
                            threads,
                            reduction,
                            bt_rows,
                            column_groups: 1,
                        }));
                    }
                }
            }
            return out
                .into_iter()
                .filter(|&tile| tile != MatmulTile::Gemv(shape) && tile.fits(self))
                .collect();
        }
        if self.conv2d.is_some() {
            let mut out = Vec::new();
            for tile in [16u32, 32, 64] {
                out.push(match tile {
                    16 => MatmulTile::Tile16,
                    32 => MatmulTile::Tile32,
                    _ => MatmulTile::Tile64,
                });
                for k_tile in [16, 32] {
                    for interleave_columns in [false, true] {
                        if tile == 16 && interleave_columns {
                            continue;
                        }
                        out.push(MatmulTile::SpecializedConv {
                            tile_size: tile,
                            k_tile,
                            interleave_columns,
                        });
                    }
                }
            }
            return out
                .into_iter()
                .filter(|&tile| tile != initial && tile.fits(self))
                .collect();
        }
        if self.effective_matrix_types().is_some() {
            let mut candidates = Vec::new();
            if self.matrix_types.is_some() {
                candidates.extend(
                    [1, 2, 4, 8, 16].map(|splits| MatmulTile::CooperativeChecked { splits }),
                );
            } else if let Some(config) = config.filter(|c| c.use_f16_input && !c.compensated) {
                candidates.push(MatmulTile::CooperativeF16 {
                    tile_size: config.tile_size,
                });
            }
            // Keep the incumbent's partition count first, then explore other
            // counts only when the complete SumRows/buffer contract permits.
            let splits = match initial {
                MatmulTile::CooperativeChecked { splits }
                | MatmulTile::CooperativeWorkgroup { splits, .. } => splits,
                _ => 1,
            };
            candidates.extend(
                crate::codegen::WorkgroupMatmulShape::all()
                    .map(|shape| MatmulTile::CooperativeWorkgroup { shape, splits }),
            );
            if let MatmulTile::CooperativeWorkgroup { shape, .. } = initial {
                candidates.extend(
                    [2, 4, 8, 16].map(|splits| MatmulTile::CooperativeWorkgroup { shape, splits }),
                );
            }
            return candidates
                .into_iter()
                .filter(|&tile| tile != initial && tile.fits(self))
                .collect();
        }
        if self.split_k {
            // A split product keeps its SumRows: other shapes at the same
            // partition count, and this shape at counts its buffer holds.
            let MatmulTile::CooperativeTiled { shape, splits } = initial else {
                return Vec::new();
            };
            return crate::codegen::CooperativeMatmulShape::all()
                .map(|other| MatmulTile::CooperativeTiled {
                    shape: other,
                    splits,
                })
                .chain([2, 4, 8, 16].map(|count| MatmulTile::CooperativeTiled {
                    shape,
                    splits: count,
                }))
                .filter(|&tile| tile != initial && tile.fits(self))
                .collect();
        }
        let mut candidates: Vec<_> = [
            Some(MatmulTile::Tile64),
            Some(MatmulTile::Tile32),
            MatmulTile::native_cooperative(config)
                .filter(|_| !self.weight_format.uses_reduced_storage()),
        ]
        .into_iter()
        .flatten()
        .collect();
        if !self.weight_format.is_quantized() {
            candidates.retain(|tile| matches!(tile, MatmulTile::CooperativeF32 { .. }));
            if config.is_some_and(|c| c.tile_size == 16 && !c.use_f16_input && !c.compensated) {
                candidates.extend(
                    crate::codegen::CooperativeMatmulShape::all()
                        .map(|shape| MatmulTile::CooperativeTiled { shape, splits: 1 }),
                );
            }
            for interleave_columns in [false, true] {
                for k_stage in [32, 16, 8] {
                    for tile_size in [64, 32] {
                        candidates.push(MatmulTile::Scalar(crate::codegen::ScalarMatmulShape {
                            tile_size,
                            tile_n: 0,
                            k_stage,
                            interleave_columns,
                            unroll_k: false,
                        }));
                    }
                }
            }
        }
        candidates
            .into_iter()
            .filter(|&tile| tile != initial && tile.fits(self))
            .collect()
    }
}

/// Restrict an experiment without changing any candidate's legality or precision.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub enum TuneScope {
    /// Historical search domain, also used for missing report settings.
    #[default]
    Dense,
    /// Scalar input and weight gradients only; no forward or cooperative convolution.
    ConvDerivatives,
    /// Scalar forward, input gradient and weight gradient convolutions.
    Convolution,
    All,
}

impl TuneScope {
    pub(crate) fn includes(self, class: &TuneClass) -> bool {
        match self {
            Self::All => true,
            Self::Dense => class.conv2d.is_none(),
            Self::Convolution => class.conv2d.is_some(),
            Self::ConvDerivatives => {
                class.conv2d.is_some() && class.shader != ShaderEntry::Conv2dGemm
            }
        }
    }
}

/// Placement of the private upload/readback buffer, never the kernel bindings.
/// Its default preserves historical reports; new [`TuneOptions`] use Download.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub enum TuneStaging {
    /// Original bidirectional, GPU-preferred shared staging.
    #[default]
    Shared,
    /// Host-read-optimized staging; still used for uploads as well as readbacks.
    Download,
}

/// Lifetime of one private staging buffer, never of bindings or qualified data.
/// Its default preserves historical reports; new [`TuneOptions`] use SameSize.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub enum TuneStagingReuse {
    /// Allocate and destroy staging for every comparison (historical policy).
    #[default]
    Fresh,
    /// Keep one exact-size buffer between comparisons within this tuning call.
    /// Release on a size change and before returning; never retain extra capacity.
    SameSize,
}

/// Kernels that share a measured-to-predicted ratio during one search.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub(crate) enum TileFamily {
    Scalar,
    Cooperative,
    Gemv,
}

/// How a class's challengers are compared with its incumbent.
/// Its default preserves historical reports; new [`TuneOptions`] race.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub enum TuneStrategy {
    /// One challenger at a time against the latest winner, each qualified
    /// before it is timed (historical policy).
    #[default]
    Sequential,
    /// Time the incumbent and every challenger together in shared rounds,
    /// drop challengers that cannot pass the decision rule, and qualify only
    /// a finalist that would be installed.
    Race,
}

/// Where a search's [`MachineProfile`] comes from.
/// Its default preserves historical reports; new [`TuneOptions`] measure.
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub enum TuneCalibration {
    /// No profile: classes keep their structural order, with no time-share cutoff.
    #[default]
    Off,
    /// Measure the machine within a bounded part of the budget, reusing the
    /// context's profile. Larger allowances can fill missing measurements.
    Measure,
    /// Use this profile, measured earlier or supplied by the caller.
    Profile(MachineProfile),
}

/// General properties of the machine that shape a search: how to time it and
/// roughly how long work takes there. Measured with Meganeura's own kernels,
/// never looked up by device. Predictions from it only order and bound work;
/// measurements still select every winner. `None` marks a measurement the
/// calibration budget did not reach.
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct MachineProfile {
    /// Measured with GPU timestamps rather than the host clock.
    pub gpu_timestamps: bool,
    /// Host wall time of an empty submission and wait.
    pub submit_round_trip: Option<Duration>,
    /// GPU time a timed compute pass adds around a trivial dispatch.
    pub pass_overhead: Option<Duration>,
    /// Time one more trivial dispatch adds inside a pass.
    pub dispatch_overhead: Option<Duration>,
    /// Achieved streaming bandwidth, in bytes per second.
    pub bandwidth: Option<f64>,
    /// Achieved FLOP/s of the ordinary f32 matmul kernel, scaled to a fully
    /// occupied machine.
    pub matmul_flops: Option<f64>,
    /// Concurrent 256-invocation workgroups, each with 16 KiB of workgroup
    /// memory, before a fixed-work kernel slows down.
    pub parallel_width: Option<u32>,
    /// Cumulative wall time spent measuring and completing this profile.
    pub elapsed: Duration,
}

impl MachineProfile {
    /// Share of the machine a dispatch of `workgroups` keeps busy on average,
    /// counting the idle slots of its last wave.
    pub(crate) fn wave_utilization(&self, workgroups: u64) -> f64 {
        let Some(width) = self.parallel_width.filter(|&width| width != 0) else {
            return 1.0;
        };
        if workgroups == 0 {
            return 1.0;
        }
        let width = u64::from(width);
        workgroups as f64 / (workgroups.div_ceil(width) as f64 * width as f64)
    }

    /// Roofline estimate of one dispatch in milliseconds: the slower of
    /// arithmetic at `flops_rate` (the calibrated matmul rate when `None`),
    /// slowed by idle wave slots, and `bytes` at the calibrated bandwidth,
    /// plus dispatch overhead. `None` without a bandwidth measurement.
    pub(crate) fn predict_ms(
        &self,
        flops: f64,
        bytes: f64,
        workgroups: u64,
        flops_rate: Option<f64>,
    ) -> Option<f64> {
        let traffic = bytes
            / self
                .bandwidth
                .filter(|&rate| rate.is_finite() && rate > 0.0)?;
        let arithmetic = match flops_rate.or(self.matmul_flops) {
            Some(rate) if rate.is_finite() && rate > 0.0 && flops > 0.0 => {
                flops / (rate * self.wave_utilization(workgroups))
            }
            _ => 0.0,
        };
        let overhead = self.dispatch_overhead.unwrap_or_default().as_secs_f64();
        let ms = (traffic.max(arithmetic) + overhead) * 1e3;
        (ms.is_finite() && ms > 0.0).then_some(ms)
    }
}

/// Resource bounds and decision policy for [`crate::Session::tune_with`].
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct TuneOptions {
    /// Defaults to All. Historical missing settings retain Dense.
    #[serde(default)]
    pub scope: TuneScope,
    pub max_classes: usize,
    /// Comparison/calibration buffers including staging and a conservative
    /// 1 MiB encoder-scratch allowance per command buffer. Excludes pipelines
    /// and opaque driver command/descriptor allocations.
    pub max_scratch_bytes: usize,
    /// Defaults to Download. Does not alter scratch binding placement,
    /// validation or kernel candidates.
    /// Missing historical settings deserialize to the original shared staging.
    #[serde(default)]
    pub staging: TuneStaging,
    /// Defaults to SameSize. Reuses only private staging, without changing
    /// validation or scratch bounds.
    /// Historical missing settings retain Fresh.
    #[serde(default)]
    pub staging_reuse: TuneStagingReuse,
    /// Soft wall-clock deadline, including compilation and qualification.
    /// An in-flight driver call, GPU submission, or validation cannot be preempted.
    pub max_time: Duration,
    pub warmup_runs: u32,
    /// Complete baseline/candidate pairs required for a decision. Races use
    /// shared rounds in paired forward/reverse order.
    pub sample_pairs: usize,
    /// Separate, barrier-delimited dispatches in each timed submission;
    /// complete sequences for split-K and cached-attention measurements.
    /// This is an upper bound when `target_sample_time` is set.
    pub dispatches_per_sample: u32,
    /// Use warmup timings to aim for this duration in the faster variant,
    /// capped by `dispatches_per_sample`. Both variants use the same count.
    /// None, or no warmup, retains the fixed dispatch count. Historical
    /// reports default to None; new searches target two milliseconds.
    #[serde(default)]
    pub target_sample_time: Option<Duration>,
    /// Required fractional improvement, in addition to a noise margin.
    pub min_improvement: f64,
    /// Defaults to Race. Both strategies apply the same decision rule and
    /// qualification; historical missing settings retain Sequential.
    #[serde(default)]
    pub strategy: TuneStrategy,
    /// Defaults to Measure. With a profile, classes are visited in order of
    /// predicted time, and visiting stops once the classes left are predicted
    /// to total less than 0.1% of the plan. Historical missing settings retain Off.
    #[serde(default)]
    pub calibration: TuneCalibration,
}

impl Default for TuneOptions {
    fn default() -> Self {
        Self {
            scope: TuneScope::All,
            max_classes: 8,
            max_scratch_bytes: 64 * 1024 * 1024,
            staging: TuneStaging::Download,
            staging_reuse: TuneStagingReuse::SameSize,
            max_time: Duration::from_secs(2),
            warmup_runs: 1,
            sample_pairs: 6,
            dispatches_per_sample: 16,
            target_sample_time: Some(Duration::from_millis(2)),
            min_improvement: 0.05,
            strategy: TuneStrategy::Race,
            calibration: TuneCalibration::Measure,
        }
    }
}

/// An invalid search configuration; rejected before touching the GPU.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TuneError(pub &'static str);

impl std::fmt::Display for TuneError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.0)
    }
}

impl std::error::Error for TuneError {}

impl TuneOptions {
    pub(crate) fn validate(&self) -> Result<(), TuneError> {
        if !(4..=64).contains(&self.sample_pairs) {
            return Err(TuneError("sample_pairs must be between 4 and 64"));
        }
        if !(1..=256).contains(&self.dispatches_per_sample) {
            return Err(TuneError("dispatches_per_sample must be between 1 and 256"));
        }
        if self.warmup_runs > 32 {
            return Err(TuneError("warmup_runs must not exceed 32"));
        }
        if self.target_sample_time.is_some_and(|time| time.is_zero()) {
            return Err(TuneError("target_sample_time must be positive when set"));
        }
        if !self.min_improvement.is_finite() || !(0.0..1.0).contains(&self.min_improvement) {
            return Err(TuneError("min_improvement must be finite and in [0, 1)"));
        }
        if let TuneCalibration::Profile(ref profile) = self.calibration {
            if [profile.bandwidth, profile.matmul_flops]
                .into_iter()
                .flatten()
                .any(|rate| !rate.is_finite() || rate <= 0.0)
                || profile.parallel_width == Some(0)
            {
                return Err(TuneError(
                    "machine profile rates and width must be positive and finite",
                ));
            }
        }
        Ok(())
    }

    pub(crate) fn sample_repetitions(&self, warmup_ms: [f64; 2]) -> u32 {
        let Some(target) = self.target_sample_time else {
            return self.dispatches_per_sample;
        };
        if warmup_ms.iter().any(|ms| !ms.is_finite() || *ms <= 0.0) {
            return self.dispatches_per_sample;
        }
        let faster_ms = warmup_ms[0].min(warmup_ms[1]);
        let count = (target.as_secs_f64() * 1000.0 / faster_ms).ceil() as u32;
        count.clamp(1, self.dispatches_per_sample)
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum TuneDecision {
    FasterCandidate,
    KeepBaseline,
    /// Dropped from a race without qualification: its samples could no
    /// longer pass the decision rule, or other challengers were faster.
    Eliminated,
    InvalidOutput,
    InvalidTiming,
    TimeBudget,
    ScratchLimit,
    DeviceMemoryBudget,
    ShaderRejected,
}

/// Non-overlapping host wall-time phases within one candidate comparison.
/// `None` means the phase was not reached; an early exit records partial time.
/// Their sum excludes final decision bookkeeping. Historical reports did not
/// measure cleanup; missing measurements are not zero cost.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct TunePhaseTimes {
    /// Legality/budget checks, pipeline setup, scratch allocation and bindings.
    /// Includes [`TuneOutcome::compile_time`]; do not add that timer again.
    pub preparation: Option<Duration>,
    /// Nested accounting, already inside `preparation`.
    #[serde(default)]
    pub preparation_breakdown: Option<TunePreparationTimes>,
    /// Input generation, uploads, trial dispatches, readbacks and CPU checks.
    pub qualification: Option<Duration>,
    /// Nested accounting, already inside `qualification`; never add it twice.
    /// `None` means an older report or qualification was not reached.
    #[serde(default)]
    pub qualification_breakdown: Option<TuneQualificationTimes>,
    /// Warming both variants on the ordinary-magnitude qualification inputs.
    pub warmup: Option<Duration>,
    /// Paired timing loop, including incomplete/discarded pairs and host checks.
    /// Not the sum of accepted per-dispatch samples or GPU timestamp duration.
    pub sampling: Option<Duration>,
    /// Destruction of private comparison resources, including early exits.
    /// Retained staging is released separately in [`TuneReport::final_cleanup`].
    #[serde(default)]
    pub cleanup: Option<Duration>,
}

/// Disjoint host wall times within preparation, not GPU timestamps.
/// `None` means historical/unreached work; repeated operations accumulate.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct TunePreparationTimes {
    pub checks: Option<Duration>,
    /// Same measurement as [`TuneOutcome::compile_time`]; do not add it twice.
    pub pipelines: Option<Duration>,
    /// Candidate binding allocations with matching memory placement.
    pub buffers: Option<Duration>,
    /// Private staging allocation/reuse and release of a previous size.
    pub staging: Option<Duration>,
    pub encoder: Option<Duration>,
    /// Dispatch cloning, geometry and pipeline lookup, not GPU execution.
    pub bindings: Option<Duration>,
}

/// Accumulated, disjoint host wall times within qualification, not GPU timestamps.
/// Repeated operations add to each field; `None` means the operation was not
/// reached. Early returns preserve partial time. Bookkeeping and destruction
/// outside these scopes remain in the enclosing qualification time.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct TuneQualificationTimes {
    /// Deterministic inputs and padding.
    pub input_preparation: Option<Duration>,
    /// CPU copy into mapped staging memory.
    pub upload_host_copy: Option<Duration>,
    /// Upload encoding, transfer, submission and wait.
    pub upload_transfer: Option<Duration>,
    /// GPU fill of output and split-partial NaN sentinels, including the wait.
    #[serde(default)]
    pub output_clear: Option<Duration>,
    /// Candidate encoding, dispatch, submission and wait.
    pub dispatch: Option<Duration>,
    /// Readback encoding, transfer, submission and wait.
    pub readback_transfer: Option<Duration>,
    /// Host vector allocation and CPU copy out of mapped staging memory.
    pub readback_host_copy: Option<Duration>,
    /// Full-output finite/parity scans and sampled f64 reference dots.
    pub validation: Option<Duration>,
}

/// Evidence for one candidate comparison within an exact class. Times are
/// scratch milliseconds per dispatch (per complete sequence for split-K),
/// never whole-step latency. Sequential comparisons use batched wall times;
/// races use GPU timestamps when `TuneRace::gpu_timestamps` is set, otherwise
/// batched wall times.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct TuneOutcome<Class = TuneClass, Choice = MatmulTile> {
    pub class: Class,
    pub dispatches: usize,
    pub initial: Choice,
    pub candidate: Choice,
    pub selected: Choice,
    /// Experimental two-pass dW challenger; baseline is unsplit with the same
    /// tile. Times are per complete sequence. `decision` reports the result;
    /// no live plan is changed and `selected` still describes only the tile.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub candidate_split_k: Option<u32>,
    pub decision: TuneDecision,
    pub qualified: bool,
    /// Shader rejection or the candidate/input pattern that failed qualification.
    pub failure: Option<String>,
    /// Total comparison cost, including pipeline setup, qualification and timing.
    pub elapsed: Duration,
    pub compile_time: Duration,
    /// Phase accounting is absent in reports written before this instrumentation.
    /// Never reinterpret missing historical measurements as zero-cost phases.
    #[serde(default)]
    pub phase_times: Option<TunePhaseTimes>,
    /// Actual private allocation requests; absent for historical/skipped work.
    #[serde(default)]
    pub scratch: Option<TuneScratchUsage>,
    pub baseline_ms: Vec<f64>,
    pub candidate_ms: Vec<f64>,
    /// Actual dispatch/sequence count used for each paired kernel sample.
    /// None for historical reports, whole-program trials, or skipped sampling.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub sample_repetitions: Option<u32>,
    pub baseline_median_ms: Option<f64>,
    pub candidate_median_ms: Option<f64>,
    /// Twice the median absolute deviation of paired time differences.
    /// A conservative selection guard, not a statistical confidence interval.
    pub noise_margin_ms: Option<f64>,
}

impl<Class, Choice: Copy> TuneOutcome<Class, Choice> {
    pub(crate) fn new(class: Class, dispatches: usize, initial: Choice, candidate: Choice) -> Self {
        Self {
            class,
            dispatches,
            initial,
            candidate,
            selected: initial,
            candidate_split_k: None,
            decision: TuneDecision::KeepBaseline,
            qualified: false,
            failure: None,
            elapsed: Duration::ZERO,
            compile_time: Duration::ZERO,
            phase_times: Some(Default::default()),
            scratch: None,
            baseline_ms: Vec::new(),
            candidate_ms: Vec::new(),
            sample_repetitions: None,
            baseline_median_ms: None,
            candidate_median_ms: None,
            noise_margin_ms: None,
        }
    }
}

/// Per-comparison buffer requests, not driver heap sizes or peak VRAM.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TuneScratchUsage {
    pub binding_bytes: Vec<usize>,
    pub staging_bytes: usize,
    pub staging_reused: bool,
    /// The identical synthetic B operand was already resident for this class.
    #[serde(default)]
    pub input_b_reused: bool,
    /// Both input patterns reuse full, previously qualified incumbent outputs
    /// from this call, with the same class, tile and physical binding sizes.
    #[serde(default)]
    pub baseline_reused: bool,
}

/// Private scratch lifetime and peak simultaneous buffer budget in one call.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct TuneScratchStats {
    pub staging_allocations: usize,
    pub staging_reuses: usize,
    pub staging_releases: usize,
    /// Binding requests, staging and encoder-scratch allowances; excludes
    /// pipelines and opaque driver allocations. Older reports omit encoders.
    pub peak_bytes: usize,
    /// Must be zero when `tune_with` returns, including bounded/failed searches.
    pub retained_staging_bytes: usize,
    #[serde(default)]
    pub input_b_allocations: usize,
    #[serde(default)]
    pub input_b_reuses: usize,
    #[serde(default)]
    pub input_b_releases: usize,
    /// Must be zero when tuning returns, including bounded/failed searches.
    #[serde(default)]
    pub retained_input_bytes: usize,
}

#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct TuneReport {
    pub options: TuneOptions,
    pub outcomes: Vec<TuneOutcome>,
    pub eligible_classes: usize,
    /// Live search visits a bounded candidate set per class; explicit split-K
    /// probes accept up to four counts against the same unsplit control.
    pub visited_classes: usize,
    /// Qualified selections reused inside one calibrated build, never globally.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub reused_classes: Vec<(TuneClass, MatmulTile)>,
    pub excluded_dispatches: usize,
    pub class_limit_reached: bool,
    pub time_budget_exhausted: bool,
    pub elapsed: Duration,
    /// Release of retained scratch and unused convolution pipelines, within
    /// total `elapsed`. None when cleanup was not run or recorded.
    #[serde(default)]
    pub final_cleanup: Option<Duration>,
    #[serde(default)]
    pub scratch: Option<TuneScratchStats>,
    /// The profile that ordered this search, if any. Measuring it counts
    /// toward `elapsed`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub machine_profile: Option<MachineProfile>,
    /// Calibration charged to this call, zero when a supplied/cached profile
    /// was reused without further measurement. `MachineProfile::elapsed`
    /// describes the cumulative measurement time.
    #[serde(default)]
    pub calibration_time: Duration,
    /// Predicted time of one plan execution, from the machine profile.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub predicted_plan_ms: Option<f64>,
    /// Visiting stopped because the classes left were predicted to total less
    /// than 0.1% of the plan.
    #[serde(default)]
    pub share_cutoff: bool,
    /// One record per raced class; challengers still report `outcomes`.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub races: Vec<TuneRace>,
}

/// One class race: its incumbent and challengers timed together. Durations are
/// host wall time; per-dispatch times come from GPU timestamps when
/// `gpu_timestamps` is set and from the host clock otherwise.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct TuneRace {
    pub class: TuneClass,
    pub dispatches: usize,
    pub initial: MatmulTile,
    pub selected: MatmulTile,
    /// Candidates that compiled and were timed, incumbent included.
    pub candidates: usize,
    /// Timed rounds after the warmup round.
    pub rounds: usize,
    pub gpu_timestamps: bool,
    /// Per-dispatch prediction that ordered this class, before correction.
    pub predicted_ms: Option<f64>,
    /// Measured median per dispatch of the incumbent.
    pub incumbent_ms: Option<f64>,
    /// Every challenger reached a reusable decision. False after a deadline,
    /// resource limit, invalid timing or failure of the incumbent control.
    pub complete: bool,
    /// Parallel compilation of every candidate's pipelines.
    pub compile_time: Duration,
    /// Checks, scratch allocation and input upload, excluding compilation.
    pub preparation: Duration,
    /// Warmup and timed rounds.
    pub racing: Duration,
    /// Qualification of finalists, if any would be installed.
    pub qualification: Duration,
    pub elapsed: Duration,
}

pub(crate) fn median(values: &[f64]) -> f64 {
    let mut sorted = values.to_vec();
    sorted.sort_by(f64::total_cmp);
    let mid = sorted.len() / 2;
    if sorted.len().is_multiple_of(2) {
        sorted[mid - 1] / 2.0 + sorted[mid] / 2.0
    } else {
        sorted[mid]
    }
}

/// Reusable two-candidate runner: the callback owns timing and deadline checks.
/// `false` is the baseline, `true` the alternative. Partial pairs never win.
pub(crate) fn measure_pairs(
    pairs: usize,
    mut measure: impl FnMut(bool) -> Option<f64>,
) -> (Vec<f64>, Vec<f64>) {
    let mut baseline = Vec::new();
    let mut candidate = Vec::new();
    for pair in 0..pairs {
        let order = if pair % 2 == 0 {
            [false, true]
        } else {
            [true, false]
        };
        let Some(first) = measure(order[0]) else {
            break;
        };
        let Some(second) = measure(order[1]) else {
            break;
        };
        let times = if order[0] {
            [second, first]
        } else {
            [first, second]
        };
        baseline.push(times[0]);
        candidate.push(times[1]);
    }
    (baseline, candidate)
}

pub(crate) fn decide<Class, Choice: Copy>(
    outcome: &mut TuneOutcome<Class, Choice>,
    options: &TuneOptions,
) {
    if outcome.baseline_ms.len() != options.sample_pairs
        || outcome.candidate_ms.len() != options.sample_pairs
    {
        outcome.decision = TuneDecision::TimeBudget;
        return;
    }
    if outcome
        .baseline_ms
        .iter()
        .chain(&outcome.candidate_ms)
        .any(|x| !x.is_finite() || *x <= 0.0)
    {
        outcome.decision = TuneDecision::InvalidTiming;
        return;
    }
    let baseline = median(&outcome.baseline_ms);
    let candidate = median(&outcome.candidate_ms);
    let differences: Vec<_> = outcome
        .baseline_ms
        .iter()
        .zip(&outcome.candidate_ms)
        .map(|(b, c)| b - c)
        .collect();
    let gain = median(&differences);
    let deviations: Vec<_> = differences.iter().map(|d| (d - gain).abs()).collect();
    let noise = 2.0 * median(&deviations);
    outcome.baseline_median_ms = Some(baseline);
    outcome.candidate_median_ms = Some(candidate);
    outcome.noise_margin_ms = Some(noise);
    if outcome.qualified
        && baseline - candidate > baseline * options.min_improvement + noise
        && gain > baseline * options.min_improvement + noise
    {
        outcome.selected = outcome.candidate;
        outcome.decision = TuneDecision::FasterCandidate;
    } else {
        outcome.decision = TuneDecision::KeepBaseline;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn dispatch() -> Dispatch {
        Dispatch {
            params: vec![33, 17, 65, 0],
            workgroups: [2, 1, 1],
            input_buffers: vec![crate::compile::BufferRef(0), crate::compile::BufferRef(1)],
            output_buffer: crate::compile::BufferRef(2),
            ..Default::default()
        }
    }

    #[test]
    fn rectangular_choices_preserve_types_geometry_and_binding_capacity() {
        use crate::{
            DType,
            codegen::{MatrixTypes, WorkgroupMatmulShape},
            compile::Kernel,
        };
        for types in [
            MatrixTypes::BF16,
            MatrixTypes::reduced(DType::F16),
            MatrixTypes {
                output: DType::BF16,
                ..MatrixTypes::BF16
            },
            MatrixTypes {
                a: DType::F32,
                ..MatrixTypes::BF16
            },
        ] {
            let mut class = class(256, 512, 256);
            class.matrix_types = Some(types);
            class.binding_bytes = class.buffer_sizes().unwrap();
            let initial = MatmulTile::CooperativeChecked { splits: 1 };
            let candidates = class.challengers(initial, None);
            for shape in WorkgroupMatmulShape::all() {
                let tile = MatmulTile::CooperativeWorkgroup { shape, splits: 1 };
                assert!(candidates.contains(&tile));
                let original = Dispatch {
                    params: vec![256, 256, 512, 0],
                    kernel: Kernel::CooperativeTyped { types, splits: 1 },
                    workgroups: [8, 16, 1],
                    ..dispatch()
                };
                let mut chosen = original.clone();
                tile.apply(&mut chosen, &class);
                assert_eq!(
                    chosen.workgroups,
                    [256 / shape.rows, 512 / shape.columns, 1]
                );
                assert_eq!(
                    chosen.kernel,
                    Kernel::CooperativeWorkgroup {
                        shape,
                        types: Some(types),
                        splits: 1
                    }
                );
                let saved = ron::to_string(&chosen).unwrap();
                assert_eq!(ron::from_str::<Dispatch>(&saved).unwrap(), chosen);
                initial.apply(&mut chosen, &class);
                assert_eq!(chosen, original);
                let mut undersized = class.clone();
                *undersized.binding_bytes.last_mut().unwrap() -= 1;
                assert!(!tile.fits(&undersized));
                let mut ragged = class.clone();
                ragged.n += 1;
                assert!(!tile.fits(&ragged));
            }
        }
        let mut automatic = class(256, 512, 256);
        let native = CoopConfig {
            tile_size: 16,
            use_f16_input: false,
            compensated: false,
        };
        assert!(
            automatic
                .challengers(MatmulTile::Tile64, Some(&native))
                .iter()
                .all(|tile| !matches!(
                    tile,
                    MatmulTile::CooperativeWorkgroup { .. } | MatmulTile::CooperativeF16 { .. }
                ))
        );
        automatic.cooperative_f16 = true;
        let reduced = CoopConfig {
            use_f16_input: true,
            ..native
        };
        let legacy = MatmulTile::CooperativeF16 { tile_size: 16 };
        assert!(legacy.fits(&automatic));
        assert!(!MatmulTile::Tile64.fits(&automatic));
        assert!(
            automatic
                .challengers(legacy, Some(&reduced))
                .iter()
                .all(|tile| matches!(tile, MatmulTile::CooperativeWorkgroup { .. }))
        );
    }

    #[test]
    fn activation_quantization_splits_the_tuning_class() {
        let ordinary = Dispatch {
            shader: ShaderEntry::MatMulGemv,
            params: vec![1, 256, 16, 0],
            workgroups: [4, 1, 1],
            weight_format: crate::compile::WeightFormat::Q40,
            input_buffers: vec![crate::compile::BufferRef(0), crate::compile::BufferRef(1)],
            output_buffer: crate::compile::BufferRef(2),
            ..Default::default()
        };
        let quantized = Dispatch {
            kernel: crate::compile::Kernel::Gemv {
                shape: crate::codegen::GemvShape::initial(crate::codegen::ShaderGroup::MatMulGemv),
                integer_dot: true,
            },
            ..ordinary.clone()
        };

        let a = TuneClass::from_dispatch(&ordinary, None).expect("ordinary Q4_0 GEMV");
        let b = TuneClass::from_dispatch(&quantized, None).expect("int-dot Q4_0 GEMV");

        assert_eq!(a.weight_format, b.weight_format);
        assert_eq!((a.m, a.n, a.k), (b.m, b.n, b.k));
        assert_eq!(a.shader, b.shader);
        assert_ne!(a, b, "the two kernels would coalesce into one class");
        assert!(!a.gemv_int_dot && b.gemv_int_dot);
    }

    #[test]
    fn every_gemv_storage_format_reaches_every_shape() {
        use crate::DType;
        use crate::codegen::{GemvReduction, GemvShape};

        const K: usize = 512;
        const N: usize = 256;
        for dtype in [
            DType::F32,
            DType::F16,
            DType::Q4_0,
            DType::Q8_0,
            DType::Q40,
            DType::Q4K,
            DType::Q6K,
            DType::Q5K,
            DType::Q3K,
        ] {
            let mut graph = crate::Graph::new();
            let a = graph.input("a", &[1, K]);
            let b = match dtype {
                DType::F32 => graph.parameter("b", &[K, N]),
                DType::F16 => graph.parameter_f16("b", &[K, N]),
                DType::Q4_0 => graph.parameter_q4("b", &[K, N]),
                DType::Q8_0 => graph.parameter_q8("b", &[K, N]),
                DType::Q40 => graph.parameter_q40("b", &[K, N]),
                DType::Q4K => graph.parameter_q4k("b", &[K, N]),
                DType::Q6K => graph.parameter_q6k("b", &[K, N]),
                DType::Q5K => graph.parameter_q5k("b", &[K, N]),
                DType::Q3K => graph.parameter_q3k("b", &[K, N]),
                other => unreachable!("unlisted GEMV format {other:?}"),
            };
            let y = graph.matmul(a, b);
            graph.set_outputs(vec![y]);
            let plan = crate::compile::compile_with(&graph, &Default::default());
            let dispatch = plan
                .dispatches
                .iter()
                .find(|d| d.shader == ShaderEntry::MatMulGemv)
                .unwrap_or_else(|| panic!("{dtype:?}: no K-split GEMV"));
            let mut class = TuneClass::from_dispatch(dispatch, None)
                .unwrap_or_else(|| panic!("{dtype:?}: GEMV is not tunable"));
            class.binding_bytes = dispatch
                .input_buffers
                .iter()
                .chain(std::iter::once(&dispatch.output_buffer))
                .map(|buffer| plan.buffers[buffer.0 as usize])
                .collect();
            assert_eq!(
                class.weight_format,
                crate::compile::WeightFormat::from_dtype(dtype)
            );

            let initial = MatmulTile::selected(dispatch, None).unwrap();
            let mut seen = class.challengers(initial, None);
            seen.push(initial);
            for threads in GemvShape::WIDTHS {
                for reduction in [GemvReduction::Tree, GemvReduction::Subgroup] {
                    let shape = GemvShape {
                        threads,
                        reduction,
                        bt_rows: 1,
                        column_groups: 1,
                    };
                    assert!(
                        seen.contains(&MatmulTile::Gemv(shape)),
                        "{dtype:?}: missing {shape:?} from {seen:?}"
                    );
                }
            }
        }
    }

    /// RMSNorm folded into a packed Q4_0 GEMV is its own search class, and
    /// the subgroup-reduction shapes are legal challengers. Without this,
    /// `tune_with` would skip the kernels that dominate Gemma 4 decode.
    #[test]
    fn fused_packed_q40_gemv_is_searchable_including_subgroup() {
        use crate::codegen::{GemvReduction, GemvShape};

        const K: usize = 256;
        const N: usize = 64;
        let mut graph = crate::Graph::new();
        let x = graph.input("x", &[1, K]);
        let nw = graph.parameter("nw", &[K]);
        let h = graph.rms_norm(x, nw, 1e-5);
        let w = graph.parameter_q40("w", &[K, N]);
        let y = graph.matmul(h, w);
        graph.set_outputs(vec![y]);
        let mut plan = crate::compile::compile_with(&graph, &Default::default());
        crate::compile::fuse_rmsnorm_into_gemv(&mut plan);
        let dispatch = plan
            .dispatches
            .iter()
            .find(|d| d.shader == ShaderEntry::MatMulGemv && d.gemv_rmsnorm.is_some())
            .expect("packed GEMV did not absorb the RmsNorm");
        assert!(plan.dispatches.iter().all(|d| d.reduction().is_none()));
        let mut class =
            TuneClass::from_dispatch(dispatch, None).expect("fused packed Q4_0 GEMV is tunable");
        assert!(class.gemv_rmsnorm);
        assert_eq!(class.weight_format, crate::compile::WeightFormat::Q40);
        assert_eq!((class.m, class.n, class.k), (1, N as u32, K as u32));
        let mut bindings: Vec<_> = dispatch.input_buffers.clone();
        bindings.push(dispatch.gemv_rmsnorm.as_ref().unwrap().weight);
        bindings.push(dispatch.output_buffer);
        class.binding_bytes = bindings
            .iter()
            .map(|b| plan.buffers[b.0 as usize])
            .collect();
        let initial = MatmulTile::selected(dispatch, None).unwrap();
        let challengers = class.challengers(initial, None);
        assert!(
            challengers.iter().any(|tile| matches!(
                tile,
                MatmulTile::Gemv(GemvShape {
                    reduction: GemvReduction::Subgroup,
                    ..
                })
            )),
            "fused packed GEMV has no subgroup challenger: {challengers:?}"
        );
        assert!(
            MatmulTile::Gemv(GemvShape {
                threads: 64,
                reduction: GemvReduction::Subgroup,
                bt_rows: 1,
                column_groups: 1,
            })
            .fits(&class),
            "64-wide subgroup fused Q4_0 GEMV must fit the class scratch"
        );
    }

    /// Gemma 4's tied vocab projection is Q8, M=1, N=262144. That is 65536
    /// workgroups, which the portable 65535 cap used to reject, so `tune_with`
    /// never measured subgroup variants of the costliest decode kernel.
    #[test]
    fn vocab_width_q8_rmsnorm_gemv_is_searchable_past_portable_workgroup_min() {
        use crate::codegen::{GemvReduction, GemvShape};

        const K: u32 = 1536;
        const N: u32 = 262_144;
        let dispatch = Dispatch {
            shader: ShaderEntry::MatMulGemv,
            params: vec![1, K, N, 0],
            workgroups: [N / 4, 1, 1],
            weight_format: crate::compile::WeightFormat::Q8,
            gemv_rmsnorm: Some(crate::compile::GemvRmsNorm {
                weight: crate::compile::BufferRef(2),
                eps_bits: 1e-6f32.to_bits(),
            }),
            input_buffers: vec![crate::compile::BufferRef(0), crate::compile::BufferRef(1)],
            output_buffer: crate::compile::BufferRef(3),
            ..Default::default()
        };
        let mut class = TuneClass::from_dispatch(&dispatch, None)
            .expect("vocab-width fused Q8 GEMV must be a tune class");
        assert!(class.gemv_rmsnorm);
        assert_eq!(class.weight_format, crate::compile::WeightFormat::Q8);
        assert_eq!((class.m, class.n, class.k), (1, N, K));
        class.binding_bytes = vec![K as usize * 4, 428_000_000, K as usize * 4, N as usize * 4];
        let initial = MatmulTile::selected(&dispatch, None).unwrap();
        assert!(
            initial.fits(&class),
            "vocab-width GEMV must survive the workgroup-count check"
        );
        assert!(
            class.challengers(initial, None).iter().any(|tile| matches!(
                tile,
                MatmulTile::Gemv(GemvShape {
                    reduction: GemvReduction::Subgroup,
                    ..
                })
            )),
            "vocab-width fused Q8 GEMV has no subgroup challenger"
        );
    }

    #[test]
    fn complete_geometry_handles_edges_in_both_directions() {
        let mut d = dispatch();
        let mut class = TuneClass::from_dispatch(&d, None).unwrap();
        assert_eq!((class.m, class.n, class.k), (33, 65, 17));
        class.binding_bytes = class.buffer_sizes().unwrap();
        for candidate in class.challengers(MatmulTile::Tile64, None) {
            candidate.apply(&mut d, &class);
            assert_eq!(MatmulTile::selected(&d, None), Some(candidate));
            assert_eq!(d.use_small_tiles(), d.workgroups == [3, 2, 1]);
            let roundtrip: Dispatch =
                serde_json::from_str(&serde_json::to_string(&d).unwrap()).unwrap();
            assert_eq!(roundtrip, d);
        }
        MatmulTile::Tile32.apply(&mut d, &class);
        assert!(d.scalar_matmul().is_none());
        assert_eq!(d.workgroups, [3, 2, 1]);
        assert!(d.use_small_tiles());
        MatmulTile::Tile64.apply(&mut d, &class);
        assert_eq!(d.workgroups, [2, 1, 1]);
        assert!(!d.use_small_tiles());
        d.shader = ShaderEntry::MatMulGemvBT;
        d.params = vec![1, 262_145, 4, 0];
        d.workgroups = crate::compile::row_gemv_workgroups(262_145);
        let class = TuneClass::from_dispatch(&d, None).unwrap();
        for bt_rows in crate::codegen::GemvShape::BT_ROWS {
            let shape = crate::codegen::GemvShape {
                threads: 64,
                reduction: crate::codegen::GemvReduction::Subgroup,
                bt_rows,
                column_groups: 1,
            };
            MatmulTile::Gemv(shape).apply(&mut d, &class);
            assert_eq!(
                d.workgroups,
                crate::compile::row_gemv_workgroups(262_145u32.div_ceil(bt_rows))
            );
            assert_eq!(TuneClass::from_dispatch(&d, None), Some(class.clone()));
        }
    }

    #[test]
    fn class_separates_shape_direction_precision_and_placement() {
        let d = dispatch();
        let base = TuneClass::from_dispatch(&d, None).unwrap();
        let mut changed = base.clone();
        changed.k += 1;
        assert_ne!(base, changed);
        changed = base.clone();
        changed.device_local[0] = true;
        assert_ne!(base, changed);
        changed = base.clone();
        changed.requires_full_precision = true;
        assert_ne!(base, changed);
        changed = base.clone();
        changed.binding_bytes = vec![4096; 3];
        assert_ne!(base, changed);
        for shader in [
            ShaderEntry::MatMulAT,
            ShaderEntry::MatMulBT,
            ShaderEntry::FusedMatMulATAdd,
            ShaderEntry::FusedMatMulBTAdd,
        ] {
            let fused = matches!(
                shader,
                ShaderEntry::FusedMatMulATAdd | ShaderEntry::FusedMatMulBTAdd
            );
            let mut transposed = d.clone();
            transposed.shader = shader;
            transposed.params = vec![33, 65, 17, 0];
            if fused {
                transposed.input_buffers.push(crate::compile::BufferRef(3));
            }
            let class = TuneClass::from_dispatch(&transposed, None).unwrap();
            assert_eq!((class.m, class.n, class.k), (33, 65, 17));
            assert_eq!(class.has_addend(), fused);
            assert_eq!(
                class.buffer_sizes().unwrap().len(),
                if fused { 4 } else { 3 }
            );
            assert_ne!(base, class);
        }
    }

    #[test]
    fn unsupported_modifiers_never_enter_the_search() {
        let base = dispatch();
        let mut variants = vec![base.clone(); 10];
        variants[0].kernel = crate::compile::Kernel::Cooperative;
        variants[1].horizontal_batch = 2;
        variants[2].kernel = crate::compile::Kernel::CooperativeCompensated;
        variants[3].extra_outputs.push(crate::compile::BufferRef(3));
        variants[4].workgroups[2] = 2;
        variants[5].shader = ShaderEntry::MatMulGemv;
        variants[6].params[0] = 0;
        variants[7].params[3] = 1;
        variants[8].params[0] = 32 * 65_535 + 1;
        variants[9].workgroups = [1, 2, 1]; // wrong row/column geometry
        for d in variants {
            assert!(TuneClass::from_dispatch(&d, None).is_none());
        }
        let mut class = TuneClass::from_dispatch(&base, None).unwrap();
        class.m = u32::MAX;
        assert!(class.buffer_sizes().is_none());
    }

    fn conv_dispatch(shader: ShaderEntry) -> Dispatch {
        let mut d = dispatch();
        d.shader = shader;
        d.kernel = crate::compile::Kernel::Default;
        d.params = vec![2, 3, 7, 9, 5, 3, 2, 2, 0, 3, 5, 1];
        d.workgroups = if matches!(
            d.shader,
            ShaderEntry::Conv2dGemm | ShaderEntry::Conv2dGradInputGemm
        ) {
            [1, 1, 2]
        } else {
            [1, 1, 1]
        };
        d
    }

    #[test]
    fn conv_tiles_preserve_complete_nchw_contracts() {
        for shader in [
            ShaderEntry::Conv2dGemm,
            ShaderEntry::Conv2dGradInputGemm,
            ShaderEntry::Conv2dGradWeightGemm,
        ] {
            let mut d = conv_dispatch(shader.clone());
            let original = d.clone();
            let mut class = TuneClass::from_dispatch(&d, None).unwrap();
            let dx = shader == ShaderEntry::Conv2dGradInputGemm;
            let forward = shader == ShaderEntry::Conv2dGemm;
            assert_eq!(
                (class.m, class.n, class.k),
                if forward {
                    (5, 15, 18)
                } else if dx {
                    (3, 63, 30)
                } else {
                    (5, 18, 30)
                }
            );
            class.binding_bytes = class.buffer_sizes().unwrap();
            assert_eq!(
                class.binding_bytes,
                if forward {
                    vec![1512, 360, 600]
                } else if dx {
                    vec![600, 360, 1512]
                } else {
                    vec![600, 1512, 360]
                }
            );
            assert_eq!(
                class.output_elements(),
                if forward {
                    150
                } else if dx {
                    378
                } else {
                    90
                }
            );
            let candidates = class.challengers(MatmulTile::Tile64, Some(&native_config(16)));
            // Plain tiles and specialized K stages; 32/64-wide tiles also
            // challenge column interleaving. Exclude the incumbent.
            assert_eq!(candidates.len(), 12);
            for candidate in candidates {
                candidate.apply(&mut d, &class);
                assert_eq!(MatmulTile::selected(&d, None), Some(candidate));
                let mut restored = TuneClass::from_dispatch(&d, None).unwrap();
                restored.binding_bytes = class.binding_bytes.clone();
                assert_eq!(restored, class);
            }
            assert!(!MatmulTile::CooperativeF32 { tile_size: 16 }.fits(&class));
            MatmulTile::Tile32.apply(&mut d, &class);
            assert_eq!(
                d.shader,
                if forward {
                    ShaderEntry::Conv2dGemmSmall
                } else if dx {
                    ShaderEntry::Conv2dGradInputGemmSmall
                } else {
                    ShaderEntry::Conv2dGradWeightGemmSmall
                }
            );
            assert_eq!(
                d.workgroups,
                if forward {
                    [1, 1, 2]
                } else if dx {
                    [2, 1, 2]
                } else {
                    [1, 1, 1]
                }
            );
            assert!(!d.use_small_tiles() && !d.use_coop());
            assert_eq!(MatmulTile::selected(&d, None), Some(MatmulTile::Tile32));
            let small = TuneClass::from_dispatch(&d, None).unwrap();
            let mut expected = class.clone();
            expected.binding_bytes.clear();
            assert_eq!(small, expected);
            MatmulTile::Tile64.apply(&mut d, &class);
            assert_eq!(d, original);
            for binding in 0..3 {
                let mut short = class.clone();
                short.binding_bytes[binding] -= 4;
                assert!(!MatmulTile::Tile32.fits(&short));
                assert!(!MatmulTile::Tile64.fits(&short));
            }
            let mut swapped_kernel = class.clone();
            let shape = swapped_kernel.conv2d.as_mut().unwrap();
            std::mem::swap(&mut shape.kernel_h, &mut shape.kernel_w);
            assert_ne!(swapped_kernel, class); // same GEMM shape is insufficient
        }
    }

    #[test]
    fn malformed_or_unsupported_convolutions_are_excluded() {
        for shader in [
            ShaderEntry::Conv2dGradInputGemm,
            ShaderEntry::Conv2dGradWeightGemm,
        ] {
            let base = conv_dispatch(shader);
            let mut variants = vec![base; 17];
            variants[0].kernel = crate::compile::Kernel::SmallTile;
            variants[1].kernel = crate::compile::Kernel::Cooperative;
            variants[2].kernel = crate::compile::Kernel::CooperativeCompensated;
            variants[3].params.pop();
            variants[4].params[7] = 0;
            variants[5].params[9] += 1;
            variants[6].params[11] = u32::MAX;
            variants[7].params[2] = u32::MAX;
            variants[8].workgroups[0] += 1;
            variants[9].workgroups[2] += 1;
            variants[10].params[0] = 0;
            variants[11].params[1] = u32::MAX;
            variants[12].weight_format = crate::compile::WeightFormat::F16;
            variants[13].shader = ShaderEntry::Conv2dGradInputGemmCoopGen(3, 2, 2);
            variants[14].kernel = crate::compile::Kernel::SpecializedConv {
                k_tile: 7,
                interleave_columns: false,
            };
            variants[15].input_buffers.pop();
            variants[16].params[5] = 100;
            for d in variants {
                assert!(TuneClass::from_dispatch(&d, None).is_none(), "{d:?}");
            }
        }
        let mut d = conv_dispatch(ShaderEntry::Conv2dGradInputGemm);
        d.params[0] = 65_536;
        d.workgroups[2] = 65_536;
        assert!(TuneClass::from_dispatch(&d, None).is_none());
    }

    #[test]
    fn convolution_indexing_does_not_require_exact_f32_reciprocals() {
        for divisor in [41, 47, 55] {
            assert_eq!((divisor as f32 * (1.0 / divisor as f32)) as u32, 0);
            for shader in [
                ShaderEntry::Conv2dGradInputGemm,
                ShaderEntry::Conv2dGradWeightGemm,
            ] {
                let mut d = conv_dispatch(shader);
                d.params = vec![2, 3, 1, divisor, 5, 1, 1, 1, 0, 1, divisor, 0];
                assert!(TuneClass::from_dispatch(&d, None).is_some(), "{d:?}");
            }
        }
        // No f32 integer-domain limit, but a padded K loop must not wrap u32.
        let mut d = conv_dispatch(ShaderEntry::Conv2dGradWeightGemm);
        for (batch, width, admitted) in [(2, (1 << 23) + 1, true), (65_535, 65_537, false)] {
            d.params = vec![batch, 1, 1, width, 1, 1, 1, 1, 0, 1, width, 0];
            assert_eq!(TuneClass::from_dispatch(&d, None).is_some(), admitted);
        }
    }

    #[test]
    fn scope_and_convolution_contracts_preserve_historical_reports() {
        assert_eq!(TuneOptions::default().scope, TuneScope::All);
        let dense = class(3, 5, 7);
        let conv = TuneClass::from_dispatch(&conv_dispatch(ShaderEntry::Conv2dGradInputGemm), None)
            .unwrap();
        for scope in [
            TuneScope::Dense,
            TuneScope::ConvDerivatives,
            TuneScope::Convolution,
            TuneScope::All,
        ] {
            assert_eq!(
                scope.includes(&dense),
                matches!(scope, TuneScope::Dense | TuneScope::All)
            );
            assert_eq!(
                scope.includes(&conv),
                matches!(
                    scope,
                    TuneScope::ConvDerivatives | TuneScope::Convolution | TuneScope::All
                )
            );
            let options = TuneOptions {
                scope,
                ..Default::default()
            };
            let mut json = serde_json::to_value(&options).unwrap();
            assert_eq!(
                serde_json::from_value::<TuneOptions>(json.clone())
                    .unwrap()
                    .scope,
                scope
            );
            json.as_object_mut().unwrap().remove("scope");
            assert_eq!(
                serde_json::from_value::<TuneOptions>(json).unwrap().scope,
                TuneScope::Dense
            );
        }
        let old = serde_json::to_value(&dense).unwrap();
        assert!(old.get("conv2d").is_none());
        assert_eq!(serde_json::from_value::<TuneClass>(old).unwrap(), dense);
        assert_eq!(
            serde_json::from_value::<TuneClass>(serde_json::to_value(&conv).unwrap()).unwrap(),
            conv
        );
    }

    #[test]
    fn trials_alternate_and_discard_incomplete_pairs() {
        let mut calls = Vec::new();
        let (baseline, candidate) = measure_pairs(4, |alternative| {
            calls.push(alternative);
            (calls.len() < 6).then_some(if alternative { 2.0 } else { 3.0 })
        });
        assert_eq!(calls, [false, true, true, false, false, true]);
        assert_eq!(baseline, [3.0, 3.0]);
        assert_eq!(candidate, [2.0, 2.0]);
    }

    fn outcome() -> TuneOutcome {
        TuneOutcome {
            class: TuneClass::from_dispatch(&dispatch(), None).unwrap(),
            dispatches: 3,
            initial: MatmulTile::Tile32,
            candidate: MatmulTile::Tile64,
            selected: MatmulTile::Tile32,
            candidate_split_k: None,
            decision: TuneDecision::KeepBaseline,
            qualified: true,
            failure: None,
            elapsed: Duration::ZERO,
            compile_time: Duration::ZERO,
            phase_times: None,
            scratch: None,
            baseline_ms: vec![10.0; 6],
            candidate_ms: vec![8.0; 6],
            sample_repetitions: Some(16),
            baseline_median_ms: None,
            candidate_median_ms: None,
            noise_margin_ms: None,
        }
    }

    #[test]
    fn decision_requires_complete_valid_qualified_low_noise_gain() {
        let options = TuneOptions::default();
        let mut win = outcome();
        decide(&mut win, &options);
        assert_eq!(win.selected, MatmulTile::Tile64);
        assert_eq!(win.decision, TuneDecision::FasterCandidate);
        let mut reverse = outcome();
        reverse.initial = MatmulTile::Tile64;
        reverse.candidate = MatmulTile::Tile32;
        reverse.selected = MatmulTile::Tile64;
        decide(&mut reverse, &options);
        assert_eq!(reverse.selected, MatmulTile::Tile32);
        for candidate in [vec![9.6; 6], vec![6.0, 10.0, 6.0, 10.0, 6.0, 10.0]] {
            let mut noisy = outcome();
            noisy.candidate_ms = candidate;
            decide(&mut noisy, &options);
            assert_eq!(noisy.selected, noisy.initial);
        }
        let mut incomplete = outcome();
        incomplete.candidate_ms.pop();
        decide(&mut incomplete, &options);
        assert_eq!(incomplete.decision, TuneDecision::TimeBudget);
        for value in [f64::NAN, f64::INFINITY, 0.0, -1.0] {
            let mut invalid = outcome();
            invalid.candidate_ms[0] = value;
            decide(&mut invalid, &options);
            assert_eq!(invalid.decision, TuneDecision::InvalidTiming);
        }
        let mut invalid = outcome();
        invalid.qualified = false;
        decide(&mut invalid, &options);
        assert_eq!(invalid.selected, invalid.initial);
    }

    #[test]
    fn configuration_validation_is_cpu_only() {
        assert!(TuneOptions::default().validate().is_ok());
        for options in [
            TuneOptions {
                sample_pairs: 0,
                ..Default::default()
            },
            TuneOptions {
                dispatches_per_sample: 0,
                ..Default::default()
            },
            TuneOptions {
                min_improvement: f64::NAN,
                ..Default::default()
            },
            TuneOptions {
                target_sample_time: Some(Duration::ZERO),
                ..Default::default()
            },
        ] {
            assert!(options.validate().is_err());
        }
    }

    #[test]
    fn adaptive_samples_keep_a_common_bounded_count() {
        let mut options = TuneOptions::default();
        assert_eq!(options.sample_repetitions([0.01, 0.02]), 16);
        assert_eq!(options.sample_repetitions([0.5, 2.0]), 4);
        assert_eq!(options.sample_repetitions([2.0, 0.5]), 4);
        assert_eq!(options.sample_repetitions([5.0, 3.0]), 1);
        for invalid in [0.0, -1.0, f64::NAN, f64::INFINITY] {
            assert_eq!(options.sample_repetitions([invalid, 3.0]), 16);
            assert_eq!(options.sample_repetitions([3.0, invalid]), 16);
        }
        options.dispatches_per_sample = 2;
        assert_eq!(options.sample_repetitions([0.5, 2.0]), 2);
        options.target_sample_time = None;
        assert_eq!(options.sample_repetitions([5.0, 3.0]), 2);
    }

    #[test]
    fn report_round_trips_settings_and_raw_evidence() {
        let report = TuneReport {
            outcomes: vec![outcome()],
            eligible_classes: 1,
            ..Default::default()
        };
        let json = serde_json::to_string_pretty(&report).unwrap();
        let restored: TuneReport = serde_json::from_str(&json).unwrap();
        assert_eq!(
            restored.outcomes[0].baseline_ms,
            report.outcomes[0].baseline_ms
        );
        assert_eq!(restored.options.max_time, report.options.max_time);
        assert_eq!(
            restored.options.target_sample_time,
            report.options.target_sample_time
        );
        assert_eq!(restored.outcomes[0].class, report.outcomes[0].class);
        assert_eq!(restored.outcomes[0].sample_repetitions, Some(16));
        let mut legacy = serde_json::to_value(&report).unwrap();
        legacy["options"]
            .as_object_mut()
            .unwrap()
            .remove("target_sample_time");
        legacy["outcomes"][0]
            .as_object_mut()
            .unwrap()
            .remove("sample_repetitions");
        let restored: TuneReport = serde_json::from_value(legacy).unwrap();
        assert_eq!(restored.options.target_sample_time, None);
        assert_eq!(restored.outcomes[0].sample_repetitions, None);
    }

    #[test]
    fn phase_times_distinguish_legacy_missing_unreached_and_measured() {
        let mut outcome = outcome();
        outcome.phase_times = Some(TunePhaseTimes {
            preparation: Some(Duration::from_millis(3)),
            preparation_breakdown: Some(TunePreparationTimes {
                buffers: Some(Duration::from_millis(1)),
                ..Default::default()
            }),
            qualification: Some(Duration::ZERO),
            cleanup: Some(Duration::from_millis(1)),
            qualification_breakdown: Some(TuneQualificationTimes {
                upload_host_copy: Some(Duration::ZERO),
                ..Default::default()
            }),
            ..Default::default()
        });
        let mut json = serde_json::to_value(&outcome).unwrap();
        let restored: TuneOutcome = serde_json::from_value(json.clone()).unwrap();
        assert_eq!(restored.phase_times, outcome.phase_times);
        for key in ["preparation_breakdown", "cleanup"] {
            json["phase_times"].as_object_mut().unwrap().remove(key);
        }
        let older: TuneOutcome = serde_json::from_value(json.clone()).unwrap();
        assert_eq!(older.phase_times.unwrap().preparation_breakdown, None);
        assert_eq!(older.phase_times.unwrap().cleanup, None);
        json["phase_times"]
            .as_object_mut()
            .unwrap()
            .remove("qualification_breakdown");
        let older: TuneOutcome = serde_json::from_value(json.clone()).unwrap();
        assert_eq!(older.phase_times.unwrap().qualification_breakdown, None);
        json.as_object_mut().unwrap().remove("phase_times");
        let legacy: TuneOutcome = serde_json::from_value(json).unwrap();
        assert_eq!(legacy.phase_times, None);
    }

    #[test]
    fn staging_round_trips_and_missing_historical_policy_is_shared() {
        assert_eq!(TuneOptions::default().staging, TuneStaging::Download);
        assert_eq!(TuneStaging::default(), TuneStaging::Shared);
        for staging in [TuneStaging::Shared, TuneStaging::Download] {
            let options = TuneOptions {
                staging,
                ..Default::default()
            };
            let mut value = serde_json::to_value(&options).unwrap();
            let restored: TuneOptions = serde_json::from_value(value.clone()).unwrap();
            assert_eq!(restored.staging, staging);
            value.as_object_mut().unwrap().remove("staging");
            let historical: TuneOptions = serde_json::from_value(value).unwrap();
            assert_eq!(historical.staging, TuneStaging::Shared);
        }
    }

    #[test]
    fn missing_historical_reuse_and_scratch_accounting_stay_distinct() {
        assert_eq!(
            TuneOptions::default().staging_reuse,
            TuneStagingReuse::SameSize
        );
        assert_eq!(TuneStagingReuse::default(), TuneStagingReuse::Fresh);
        for staging_reuse in [TuneStagingReuse::Fresh, TuneStagingReuse::SameSize] {
            let mut value = serde_json::to_value(TuneReport {
                options: TuneOptions {
                    staging_reuse,
                    ..Default::default()
                },
                scratch: Some(Default::default()),
                final_cleanup: Some(Duration::ZERO),
                ..Default::default()
            })
            .unwrap();
            let restored: TuneReport = serde_json::from_value(value.clone()).unwrap();
            assert_eq!(restored.options.staging_reuse, staging_reuse);
            assert!(restored.scratch.is_some());
            value["options"]
                .as_object_mut()
                .unwrap()
                .remove("staging_reuse");
            for key in ["scratch", "final_cleanup"] {
                value.as_object_mut().unwrap().remove(key);
            }
            let old: TuneReport = serde_json::from_value(value).unwrap();
            assert_eq!(old.options.staging_reuse, TuneStagingReuse::Fresh);
            assert_eq!(old.scratch, None);
            assert_eq!(old.final_cleanup, None);
        }
        let mut value = serde_json::to_value(outcome()).unwrap();
        value.as_object_mut().unwrap().remove("scratch");
        let old: TuneOutcome = serde_json::from_value(value).unwrap();
        assert_eq!(old.scratch, None);
    }

    fn native_config(tile_size: u32) -> CoopConfig {
        CoopConfig {
            tile_size,
            use_f16_input: false,
            compensated: false,
        }
    }

    fn class(m: u32, n: u32, k: u32) -> TuneClass {
        let mut class = TuneClass {
            weight_format: crate::compile::WeightFormat::F32,
            gemv_int_dot: false,
            shader: ShaderEntry::MatMul,
            m,
            n,
            k,
            conv2d: None,
            requires_full_precision: true,
            device_local: [false; 4],
            binding_bytes: Vec::new(),
            ..Default::default()
        };
        class.binding_bytes = class.buffer_sizes().unwrap();
        class
    }

    #[test]
    fn roofline_takes_the_slower_side_and_charges_idle_wave_slots() {
        let profile = MachineProfile {
            bandwidth: Some(1e9),
            matmul_flops: Some(1e12),
            parallel_width: Some(100),
            dispatch_overhead: Some(Duration::from_micros(10)),
            ..Default::default()
        };
        // 2 MB at 1 GB/s outweighs 1 GFLOP at 1 TFLOP/s on a full machine.
        let traffic = profile.predict_ms(1e9, 2e6, 1000, None).unwrap();
        assert!((traffic - 2.01).abs() < 1e-9, "{traffic}");
        // A tenth of the machine is busy: arithmetic takes ten times longer.
        let idle = profile.predict_ms(1e9, 0.0, 10, None).unwrap();
        assert!((idle - 10.01).abs() < 1e-9, "{idle}");
        // 150 workgroups fill one wave and half of another.
        assert!((profile.wave_utilization(150) - 0.75).abs() < 1e-12);
        assert_eq!(profile.wave_utilization(0), 1.0);
        // An explicit rate replaces the calibrated matmul rate.
        let faster = profile.predict_ms(2e9, 0.0, 1000, Some(1e13)).unwrap();
        assert!((faster - 0.21).abs() < 1e-9, "{faster}");
        // No bandwidth measurement: no prediction.
        assert_eq!(
            MachineProfile::default().predict_ms(1.0, 1.0, 1, None),
            None
        );
    }

    #[test]
    fn class_predictions_follow_tile_workgroups_and_correction() {
        let profile = MachineProfile {
            bandwidth: Some(1e12),
            matmul_flops: Some(1e12),
            parallel_width: Some(1024),
            ..Default::default()
        };
        let class = class(512, 512, 512);
        // Tile64 launches 64 workgroups, Tile32 256: both under one wave.
        let large = class.predict_ms(&profile, MatmulTile::Tile64, 1.0).unwrap();
        let small = class.predict_ms(&profile, MatmulTile::Tile32, 1.0).unwrap();
        assert!(large > small, "{large} vs {small}");
        let corrected = class.predict_ms(&profile, MatmulTile::Tile64, 2.0).unwrap();
        assert!((corrected - 2.0 * large).abs() < 1e-12);
    }

    #[test]
    fn supplied_profiles_reject_nonfinite_or_nonpositive_rates() {
        for rate in [0.0, -1.0, f64::NAN, f64::INFINITY] {
            let options = TuneOptions {
                calibration: TuneCalibration::Profile(MachineProfile {
                    bandwidth: Some(rate),
                    ..Default::default()
                }),
                ..Default::default()
            };
            assert!(options.validate().is_err(), "{rate}");
        }
        let profile = MachineProfile {
            parallel_width: Some(100),
            ..Default::default()
        };
        assert!(profile.wave_utilization(u64::MAX).is_finite());
    }

    #[test]
    fn new_options_race_but_historical_options_keep_the_original_policy() {
        let options = TuneOptions::default();
        assert_eq!(options.strategy, TuneStrategy::Race);
        assert_eq!(options.calibration, TuneCalibration::Measure);
        let mut historical = serde_json::to_value(&options).unwrap();
        let fields = historical.as_object_mut().unwrap();
        fields.remove("strategy");
        fields.remove("calibration");
        let historical: TuneOptions = serde_json::from_value(historical).unwrap();
        assert_eq!(historical.strategy, TuneStrategy::Sequential);
        assert_eq!(historical.calibration, TuneCalibration::Off);
        // A supplied profile round-trips, so it can be measured once and reused.
        let supplied = TuneCalibration::Profile(MachineProfile {
            bandwidth: Some(1e11),
            parallel_width: Some(64),
            ..Default::default()
        });
        let text = serde_json::to_string(&supplied).unwrap();
        assert_eq!(
            serde_json::from_str::<TuneCalibration>(&text).unwrap(),
            supplied
        );
    }

    #[test]
    fn split_tiles_report_partitions_and_families() {
        assert_eq!(
            MatmulTile::CooperativeChecked { splits: 4 }.partitions(),
            Some(4)
        );
        assert_eq!(MatmulTile::Tile64.partitions(), None);
        assert_eq!(MatmulTile::Tile32.family(), TileFamily::Scalar);
        assert_eq!(
            MatmulTile::CooperativeF32 { tile_size: 16 }.family(),
            TileFamily::Cooperative
        );
    }

    #[test]
    fn native_candidates_ignore_profitability_but_obey_capability_and_precision() {
        let config = native_config(8);
        let native = MatmulTile::CooperativeF32 { tile_size: 8 };
        // One workgroup, and then a dimension beyond the old 1024 veto.
        for class in [class(32, 32, 8), class(2048, 32, 17)] {
            assert!(
                class
                    .challengers(MatmulTile::Tile64, Some(&config))
                    .contains(&native)
            );
            assert!(
                !class
                    .challengers(MatmulTile::Tile64, None)
                    .contains(&native)
            );
            let f16 = CoopConfig {
                use_f16_input: true,
                ..config
            };
            assert!(
                !class
                    .challengers(MatmulTile::Tile64, Some(&f16))
                    .contains(&native)
            );
            let compensated = CoopConfig {
                compensated: true,
                ..config
            };
            assert!(
                !class
                    .challengers(MatmulTile::Tile64, Some(&compensated))
                    .contains(&native)
            );
        }
        assert!(MatmulTile::native_cooperative(Some(&native_config(4))).is_none());
    }

    #[test]
    fn native_padding_and_addend_capacity_are_legality_not_allocation_requests() {
        let native = MatmulTile::CooperativeF32 { tile_size: 8 };
        let mut class = class(33, 32, 17);
        assert!(!native.fits(&class));
        let required = native.buffer_sizes(&class).unwrap();
        assert_eq!(required, [33 * 17 * 4, 17 * 32 * 4, 64 * 32 * 4]);
        class.binding_bytes = required;
        assert!(native.fits(&class));
        class.shader = ShaderEntry::FusedMatMulAdd;
        class.binding_bytes = class.buffer_sizes().unwrap();
        class.binding_bytes[3] = 64 * 32 * 4;
        // Native 8x8 stages an edge addend with checked scalar loads;
        // unlike the direct-load path, its external input need not grow.
        assert_eq!(native.buffer_sizes(&class).unwrap()[2], 33 * 32 * 4);
        assert!(native.fits(&class));
        class.n = 17;
        assert!(native.buffer_sizes(&class).is_none());
        class.n = 32;
        class.k = 1;
        assert!(native.buffer_sizes(&class).is_none());
        class.shader = ShaderEntry::MatMulBT;
        assert!(native.buffer_sizes(&class).is_some());
    }

    #[test]
    fn native_geometry_and_kernel_move_together() {
        let config = native_config(8);
        let native = MatmulTile::CooperativeF32 { tile_size: 8 };
        let class = class(32, 64, 17);
        let mut d = dispatch();
        d.params = vec![class.m, class.k, class.n, 0];
        native.apply(&mut d, &class);
        assert_eq!(d.workgroups, [1, 2, 1]);
        assert!(d.use_coop() && !d.use_coop_compensated() && !d.use_small_tiles());
        assert!(TuneClass::from_dispatch(&d, Some(&config)).is_some());
        assert!(TuneClass::from_dispatch(&d, None).is_none());
        MatmulTile::Tile32.apply(&mut d, &class);
        assert_eq!(d.workgroups, [2, 1, 1]);
        assert!(!d.use_coop() && d.use_small_tiles());
        assert!(TuneClass::from_dispatch(&d, Some(&config)).is_some());
    }

    #[test]
    fn an_incomplete_next_challenger_keeps_the_completed_winner() {
        let options = TuneOptions::default();
        let mut first = outcome();
        decide(&mut first, &options);
        let mut next = outcome();
        next.initial = first.selected;
        next.selected = first.selected;
        next.candidate = MatmulTile::CooperativeF32 { tile_size: 8 };
        next.candidate_ms.pop();
        decide(&mut next, &options);
        assert_eq!(next.selected, first.selected);
        assert_eq!(next.decision, TuneDecision::TimeBudget);
        next.candidate_ms.push(8.0);
        decide(&mut next, &options);
        assert_eq!(next.selected, next.candidate);
    }
}