//! Select, take/Gather, GatherAxis, Arange and Scan from the pinned runtime.
//! take exports BF16/F16/F32/I32/U32 with I64 and I32 indices.
use super::super::gpu::Buffer;
use super::array::{Array, Dtype, Layout};
use super::stream::Stream;
use super::{bytes, dispatch_geometry, ops, tensor, views};

pub(super) fn select(
    condition: &Array,
    yes: &Array,
    no: &Array,
    stream: Stream,
) -> Result<Array, String> {
    let condition = ops::astype(condition, Dtype::Bool, false, stream)?;
    let dtype = yes.layout().dtype().promote(no.layout().dtype());
    let yes = ops::astype(yes, dtype, false, stream)?;
    let no = ops::astype(no, dtype, false, stream)?;
    let inputs = ops::broadcast_arrays(&[condition, yes, no], &[], stream)?;
    let layout = Layout::new(inputs[0].layout().shape(), dtype)?;
    Array::make_operation_with_inputs(stream, &inputs, layout, ops::Operation::Select)
}

pub(super) fn arange_i32(stop: i32, stream: Stream) -> Result<Array, String> {
    arange_range_i32(0, stop, stream)
}

pub(super) fn arange_u32(stop: i32, stream: Stream) -> Result<Array, String> {
    Array::make_operation_with_inputs(
        stream,
        &[],
        Layout::new(&[stop.max(0)], Dtype::U32)?,
        ops::Operation::ArangeU32,
    )
}

pub(super) fn arange_range_i32(start: i32, stop: i32, stream: Stream) -> Result<Array, String> {
    let size = i32::try_from((i64::from(stop) - i64::from(start)).max(0))
        .map_err(|_| "I32 arange size overflow")?;
    Array::make_operation_with_inputs(
        stream,
        &[],
        Layout::new(&[size], Dtype::I32)?,
        ops::Operation::ArangeI32(start),
    )
}

pub(super) fn arange_i64(stop: i32, stream: Stream) -> Result<Array, String> {
    Array::make_operation_with_inputs(
        stream,
        &[],
        Layout::new(&[stop.max(0)], Dtype::I64)?,
        ops::Operation::ArangeI64,
    )
}

pub(super) fn arange_f32(
    start: f64,
    stop: f64,
    step: f64,
    stream: Stream,
) -> Result<Array, String> {
    if !start.is_finite() || !stop.is_finite() || step.is_nan() || step == 0.0 {
        return Err("[arange] Cannot compute length or zero step".into());
    }
    if step.is_infinite() && ((step > 0.0 && start < stop) || (step < 0.0 && start > stop)) {
        return Array::new(
            &[1],
            Dtype::F32,
            super::scalar_buffer(&(start as f32).to_le_bytes())?,
        );
    }
    let size = ((stop - start) / step).ceil();
    if size > f64::from(i32::MAX) {
        return Err("[arange] Maximum size exceeded".into());
    }
    Array::make_operation_with_inputs(
        stream,
        &[],
        Layout::new(&[(size as i32).max(0)], Dtype::F32)?,
        ops::Operation::ArangeF32 { start, step },
    )
}

pub(super) fn cummax_last_i64(input: &Array, stream: Stream) -> Result<Array, String> {
    let layout = input.layout();
    if layout.dtype() != Dtype::I64 || layout.shape().is_empty() || layout.size() == 0 {
        return Err("NGram cumulative maximum requires nonempty I64 input".into());
    }
    Array::make_operation(
        stream,
        input,
        Layout::new(layout.shape(), Dtype::I64)?,
        ops::Operation::CummaxLastI64,
    )
}

pub(super) fn cumsum_last_i32(input: &Array, stream: Stream) -> Result<Array, String> {
    let layout = input.layout();
    if layout.dtype() != Dtype::I32 || layout.shape().is_empty() || layout.size() == 0 {
        return Err("sampler cumulative sum requires nonempty, non-scalar I32 input".into());
    }
    Array::make_operation(
        stream,
        input,
        Layout::new(layout.shape(), Dtype::I32)?,
        ops::Operation::CumsumLastI32,
    )
}

pub(super) fn cumsum_last_f32(
    input: &Array,
    inclusive: bool,
    stream: Stream,
) -> Result<Array, String> {
    let layout = input.layout();
    if layout.dtype() != Dtype::F32 || layout.shape().is_empty() || layout.size() == 0 {
        return Err("sampler cumulative sum requires nonempty, non-scalar F32 input".into());
    }
    Array::make_operation(
        stream,
        input,
        Layout::new(layout.shape(), Dtype::F32)?,
        ops::Operation::CumsumLastF32(inclusive),
    )
}

pub(super) fn scan_evaluate(
    input: &Array,
    output: &Array,
    operation: &str,
    inclusive: bool,
) -> Result<(), String> {
    let layout = input.layout();
    let axis = layout.shape().len() - 1;
    let input = if layout.flags().contiguous && layout.strides()[axis] != 0 {
        // Scan::eval_gpu does not use the unary allocator's +16KiB ceiling.
        if input.is_donatable() {
            output.copy_shared_buffer(
                input,
                layout.strides(),
                layout.flags(),
                input.data_size(),
                0,
            )?;
        } else {
            output.set_strided_data(
                Buffer::mtplx_bytes((input.data_size() * input.itemsize()) as u64)?,
                input.data_size(),
                layout.strides(),
                layout.flags(),
                0,
            )?;
        }
        input.clone()
    } else {
        let copy = Array::new(
            layout.shape(),
            layout.dtype(),
            Buffer::mtplx_bytes(layout.nbytes() as u64)?,
        )?;
        views::general_copy_inplace(input, &copy)?;
        let cl = copy.layout();
        output.copy_shared_buffer(&copy, cl.strides(), cl.flags(), copy.data_size(), 0)?;
        drop(cl);
        copy
    };
    drop(layout);
    let layout = input.layout();
    let size = layout.shape()[axis] as u64;
    let stride = layout.strides()[axis] as u64;
    let contiguous = stride == 1;
    let name = format!(
        "{}_scan_{}_{operation}_{}_{}",
        if contiguous { "contig" } else { "strided" },
        if inclusive { "inclusive" } else { "exclusive" },
        layout.dtype().kernel_name(),
        output.layout().dtype().kernel_name()
    );
    let max = views::kernel_max_threads(&name)?;
    let ib = input.buffer();
    let ob = output.buffer();
    let mut bindings = vec![
        tensor(0, &ib)
            .at_byte_offset(input.offset())
            .input(input.data_size()),
        tensor(1, &ob)
            .at_byte_offset(output.offset())
            .output(output.data_size()),
        bytes(2, &size),
    ];
    let blocks = stride.div_ceil(32);
    let reads = if input.itemsize() <= 4 { 4 } else { 2 };
    let elements_per_simd = reads * 32;
    let (threads, mut grid) = if contiguous {
        let threads = if size <= reads * 1024 {
            size.div_ceil(elements_per_simd) * 32
        } else if size <= reads * 2048 {
            (size / 2).div_ceil(elements_per_simd) * 32
        } else {
            u64::from(max)
        };
        (
            threads.min(u64::from(max)) as u32,
            views::large_copy_grid(&layout, size as usize)?,
        )
    } else {
        bindings.extend([bytes(3, &stride), bytes(4, &blocks)]);
        (
            (32 / reads * 32) as u32,
            views::large_copy_grid(&layout, (size * stride) as usize)?,
        )
    };
    if !contiguous {
        if u64::from(grid[0]) * blocks <= u64::from(u32::MAX) {
            grid[0] *= blocks as u32;
        } else {
            grid[1] =
                u32::try_from(u64::from(grid[1]) * blocks).map_err(|_| "scan grid overflow")?;
        }
    }
    dispatch_geometry(
        &name,
        &bindings,
        &[],
        [threads, grid[0], grid[1]],
        [threads, 1, 1],
        true,
    )
}

pub(super) fn take_along_axis(
    input: &Array,
    indices: &Array,
    axis: i32,
    stream: Stream,
) -> Result<Array, String> {
    let layout = input.layout();
    let rank = layout.shape().len();
    let axis = if axis < 0 {
        i64::from(axis) + rank as i64
    } else {
        i64::from(axis)
    };
    let axis = usize::try_from(axis).map_err(|_| "GatherAxis axis out of range")?;
    if axis >= rank || indices.layout().shape().len() != rank {
        return Err("GatherAxis axis/rank mismatch".into());
    }
    if !((matches!(
        layout.dtype(),
        Dtype::BF16 | Dtype::F16 | Dtype::F32 | Dtype::U32
    ) && matches!(indices.layout().dtype(), Dtype::I32 | Dtype::U32))
        || (matches!(layout.dtype(), Dtype::Bool | Dtype::I64)
            && indices.layout().dtype() == Dtype::I64))
    {
        return Err("GatherAxis specialization is not exported for these types".into());
    }
    drop(layout);
    let inputs = ops::broadcast_arrays(&[input.clone(), indices.clone()], &[axis], stream)?;
    let layout = Layout::new(inputs[1].layout().shape(), inputs[0].layout().dtype())?;
    Array::make_operation_with_inputs(stream, &inputs, layout, ops::Operation::GatherAxis(axis))
}

/// The one-dimensional Scatter(Sum, axes=[0]) used by logits.at[ids].add.
/// Python indexing expands updates to [index_count, 1], then scatter casts
/// them to the source dtype. It does not promote the source to Float32.
pub(super) fn scatter_add_row(
    input: &Array,
    indices: &Array,
    values: &Array,
    stream: Stream,
) -> Result<Array, String> {
    let shape = indices.layout().shape().to_vec();
    let dtype = input.layout().dtype();
    if input.layout().shape().len() != 1
        || input.layout().size() == 0
        || shape.len() != 1
        || indices.layout().dtype() != Dtype::I64
        || !matches!(dtype, Dtype::BF16 | Dtype::F16 | Dtype::F32)
    {
        return Err("penalty scatter requires one logit row and Int64 index vector".into());
    }
    let leading = values
        .layout()
        .shape()
        .iter()
        .take_while(|&&n| n == 1)
        .count();
    let values = ops::squeeze(
        values,
        Some(&(0..leading as i32).collect::<Vec<_>>()),
        stream,
    )?;
    let values = ops::broadcast_to(&values, &shape, stream)?;
    let values = ops::reshape(&values, &[shape[0], 1], stream)?;
    let values = ops::astype(&values, dtype, false, stream)?;
    let layout = Layout::new(input.layout().shape(), dtype)?;
    Array::make_operation_with_inputs(
        stream,
        &[input.clone(), indices.clone(), values],
        layout,
        ops::Operation::ScatterAddRow,
    )
}

fn scatter_add_row_evaluate(
    src: &Array,
    idx: &Array,
    upd: &Array,
    out: &Array,
) -> Result<(), String> {
    views::copy_gpu(src, out, src.layout().flags().row_contiguous)?;
    if upd.layout().size() == 0 {
        return Ok(());
    }
    let ul = upd.layout();
    let il = idx.layout();
    let ol = out.layout();
    let tag = match ol.dtype() {
        Dtype::BF16 => "bf16",
        Dtype::F16 => "f16",
        Dtype::F32 => "f32",
        _ => unreachable!(),
    };
    let large = [ol.size(), il.size(), ul.size()]
        .into_iter()
        .any(|n| n > i32::MAX as usize);
    let name = format!(
        "kernel_qwen_mtplx_scatter_{tag}_idxi64_sum_1_updc_{}_nwork1_{}",
        ul.flags().row_contiguous,
        if large { "int64_t" } else { "int" }
    );
    if views::kernel_max_threads(&name)? != 1024 {
        return Err("Scatter reference pipeline requires 1024 maximum threads".into());
    }
    let grid = [
        1,
        u32::try_from(il.size()).map_err(|_| "Scatter update grid overflow")?,
        1,
    ];
    dispatch_geometry(
        &name,
        &[
            tensor(1, &upd.buffer())
                .at_byte_offset(upd.offset())
                .input(upd.data_size()),
            tensor(2, &out.buffer())
                .at_byte_offset(out.offset())
                .output(out.data_size()),
            bytes(3, ul.shape()),
            bytes(4, ul.strides()),
            bytes(5, &2_usize),
            bytes(6, &1_usize),
            bytes(7, ol.shape()),
            bytes(8, ol.strides()),
            bytes(9, &1_usize),
            bytes(10, &[0_i32]),
            bytes(11, il.shape()),
            bytes(12, il.strides()),
            bytes(13, &[u8::from(il.flags().row_contiguous)]),
            bytes(14, &1_i32),
            bytes(15, &il.size()),
            tensor(20, &idx.buffer())
                .at_byte_offset(idx.offset())
                .input(idx.data_size()),
        ],
        &[],
        grid,
        super::block_dims(grid),
        true,
    )
}

/// Original put_along_axis factory. QSA uses Bool updates and I64 indices.
pub(super) fn put_along_axis(
    input: &Array,
    indices: &Array,
    values: &Array,
    axis: i32,
    stream: Stream,
) -> Result<Array, String> {
    let rank = input.layout().shape().len();
    let axis = if axis < 0 {
        i64::from(axis) + rank as i64
    } else {
        i64::from(axis)
    };
    let axis = usize::try_from(axis).map_err(|_| "ScatterAxis axis out of range")?;
    if axis >= rank || indices.layout().shape().len() != rank {
        return Err("ScatterAxis axis/rank mismatch".into());
    }
    if input.layout().size() == 0 {
        return Ok(input.clone());
    }
    if input.layout().dtype() != Dtype::Bool || indices.layout().dtype() != Dtype::I64 {
        return Err("ScatterAxis QSA specialization requires Bool and I64".into());
    }
    let mut upd = ops::astype(values, Dtype::Bool, false, stream)?;
    let excess = upd.layout().shape().len().saturating_sub(rank);
    if excess > 0 {
        upd = ops::squeeze(&upd, Some(&(0..excess as i32).collect::<Vec<_>>()), stream)?;
    }
    let mut inputs = ops::broadcast_arrays(&[indices.clone(), upd], &[], stream)?;
    inputs.insert(0, input.clone());
    let inputs = ops::broadcast_arrays(&inputs, &[axis], stream)?;
    let layout = Layout::new(inputs[0].layout().shape(), Dtype::Bool)?;
    Array::make_operation_with_inputs(stream, &inputs, layout, ops::Operation::ScatterAxis(axis))
}

fn scatter_evaluate(
    src: &Array,
    idx: &Array,
    upd: &Array,
    out: &Array,
    axis: usize,
) -> Result<(), String> {
    views::copy_gpu(src, out, src.layout().flags().row_contiguous)?;
    if upd.layout().size() == 0 {
        return Ok(());
    }
    let il = idx.layout();
    let ul = upd.layout();
    let sl = src.layout();
    let large = il.size() > i32::MAX as usize || sl.size() > i32::MAX as usize;
    let name = format!(
        "kernel_qwen_mtplx_scatter_axis_bool_idxi64_{}_{}{}",
        if large { "int64_t" } else { "int" },
        u8::from(ul.flags().row_contiguous),
        u8::from(il.flags().row_contiguous)
    );
    let pre = il.shape()[..axis]
        .iter()
        .map(|&n| n as u64)
        .product::<u64>();
    let post = il.shape()[axis + 1..]
        .iter()
        .map(|&n| n as u64)
        .product::<u64>();
    let grid = [
        u32::try_from(post).map_err(|_| "ScatterAxis grid overflow")?,
        il.shape()[axis] as u32,
        u32::try_from(pre).map_err(|_| "ScatterAxis grid overflow")?,
    ];
    let mut shape = il.shape().to_vec();
    let mut us = ul.strides().to_vec();
    let mut is = il.strides().to_vec();
    // Rank-one bindings are ignored by the kernel but must still be present.
    if shape.len() > 1 {
        shape.remove(axis);
        us.remove(axis);
        is.remove(axis);
    }
    dispatch_geometry(
        &name,
        &[
            tensor(0, &upd.buffer())
                .at_byte_offset(upd.offset())
                .input(upd.data_size()),
            tensor(1, &idx.buffer())
                .at_byte_offset(idx.offset())
                .input(idx.data_size()),
            tensor(2, &out.buffer())
                .at_byte_offset(out.offset())
                .output(out.data_size()),
            bytes(3, shape.as_slice()),
            bytes(4, us.as_slice()),
            bytes(5, is.as_slice()),
            bytes(6, &(sl.shape().len() - 1)),
            bytes(7, &(axis as i32)),
            bytes(8, &out.layout().shape()[axis]),
            bytes(9, &ul.strides()[axis]),
            bytes(10, &il.strides()[axis]),
        ],
        &[],
        grid,
        super::block_dims(grid),
        true,
    )
}

/// take(table, ids, axis=0): Gather preserves its unit slice axis, then squeeze.
pub(super) fn take_rows(table: &Array, ids: &Array, stream: Stream) -> Result<Array, String> {
    if table.layout().shape().len() != 2 {
        return Err("table take requires a matrix".into());
    }
    take(table, ids, 0, stream)
}

/// Original take: Gather, move the index dimensions to axis, then squeeze.
pub(super) fn take(table: &Array, ids: &Array, axis: i32, stream: Stream) -> Result<Array, String> {
    let sl = table.layout();
    let il = ids.layout();
    let rank = sl.shape().len();
    let axis = if axis < 0 {
        i64::from(axis) + rank as i64
    } else {
        i64::from(axis)
    };
    let axis = usize::try_from(axis).map_err(|_| "take axis out of range")?;
    if axis >= rank {
        return Err("take axis out of range".into());
    }
    if il.shape().len() > 4
        || !(il.dtype() == Dtype::I64 || il.dtype() == Dtype::U32 || il.dtype() == Dtype::I32)
        || !matches!(
            sl.dtype(),
            Dtype::BF16 | Dtype::F16 | Dtype::F32 | Dtype::I32 | Dtype::U32
        )
        || (sl.shape()[axis] == 0 && il.size() != 0)
    {
        return Err(
            "take requires an exported dtype/index specialization and a nonempty axis".into(),
        );
    }
    let mut shape = il.shape().to_vec();
    let idx_rank = shape.len();
    shape.extend(sl.shape());
    shape[idx_rank + axis] = 1;
    let mut gathered = Array::make_operation_with_inputs(
        stream,
        &[table.clone(), ids.clone()],
        Layout::new(&shape, sl.dtype())?,
        ops::Operation::GatherTake(axis),
    )?;
    drop(sl);
    drop(il);
    if axis != 0 {
        let axes = (idx_rank..idx_rank + axis)
            .chain(0..idx_rank)
            .chain(idx_rank + axis..shape.len())
            .map(|v| v as i32)
            .collect::<Vec<_>>();
        gathered = ops::transpose(&gathered, &axes, stream)?;
    }
    ops::squeeze(&gathered, Some(&[(idx_rank + axis) as i32]), stream)
}

fn gather_take_evaluate(
    src: &Array,
    idx: &Array,
    output: &Array,
    axis: usize,
) -> Result<(), String> {
    let nbytes = output.layout().nbytes();
    output.set_data(Buffer::mtplx_bytes(nbytes as u64)?)?;
    if output.layout().size() == 0 {
        return Ok(());
    }
    let sl = src.layout();
    let il = idx.layout();
    let tag = match sl.dtype() {
        Dtype::BF16 => "bf16",
        Dtype::F16 => "f16",
        Dtype::F32 => "f32",
        // Gather copies 32-bit payloads without value arithmetic.
        Dtype::I32 | Dtype::U32 => "u32",
        _ => unreachable!(),
    };
    let large = sl.size() > i32::MAX as usize
        || il.size() > i32::MAX as usize
        || output.layout().size() > i32::MAX as usize;
    let loc = if large { "int64_t" } else { "int" };
    let idx_tag = match il.dtype() {
        Dtype::I32 => "i32",
        Dtype::U32 => "u32",
        _ => "i64",
    };
    let mut slice = sl.shape().to_vec();
    slice[axis] = 1;
    let width = slice.iter().map(|&v| v as u64).product::<u64>();
    if sl.flags().row_contiguous
        && axis == 0
        && il.flags().row_contiguous
        && sl.strides()[0] == width as i64
    {
        let work = if width > 8 && sl.dtype().itemsize() < 4 {
            2
        } else {
            1
        };
        let grid = [
            u32::try_from(width.div_ceil(work)).map_err(|_| "table gather width overflow")?,
            u32::try_from(il.size()).map_err(|_| "table gather index overflow")?,
            1,
        ];
        let index_suffix = if il.dtype() == Dtype::U32 {
            String::new()
        } else {
            format!("idx{idx_tag}_")
        };
        return dispatch_geometry(
            &format!("kernel_qwen_mtplx_gather_front_{tag}_{index_suffix}{loc}_{work}"),
            &[
                tensor(0, &src.buffer())
                    .at_byte_offset(src.offset())
                    .input(src.data_size()),
                tensor(1, &idx.buffer())
                    .at_byte_offset(idx.offset())
                    .input(idx.data_size()),
                tensor(2, &output.buffer()).output(output.data_size()),
                bytes(3, &width),
                bytes(4, &sl.shape()[0]),
            ],
            &[],
            grid,
            super::block_dims(grid),
            true,
        );
    }
    let rank = il.shape().len() as i32;
    let dim0 = il.shape().first().copied().unwrap_or(1) as u32;
    let dim1 = if rank >= 2 {
        u32::try_from(il.size() / dim0 as usize).map_err(|_| "table gather grid overflow")?
    } else {
        1
    };
    let grid = [
        dim0,
        dim1,
        u32::try_from(width).map_err(|_| "gather slice grid overflow")?,
    ];
    let axes = [axis as i32];
    let contiguous = [u8::from(il.flags().row_contiguous)];
    let ndim = sl.shape().len();
    dispatch_geometry(
        &format!("kernel_qwen_mtplx_gather_{tag}_idx{idx_tag}_1_{rank}_{loc}"),
        &[
            tensor(0, &src.buffer())
                .at_byte_offset(src.offset())
                .input(src.data_size()),
            tensor(1, &output.buffer()).output(output.data_size()),
            bytes(2, sl.shape()),
            bytes(3, sl.strides()),
            bytes(4, &ndim),
            bytes(5, slice.as_slice()),
            bytes(6, axes.as_slice()),
            bytes(7, il.shape()),
            bytes(8, il.strides()),
            bytes(9, contiguous.as_slice()),
            bytes(10, &rank),
            tensor(20, &idx.buffer())
                .at_byte_offset(idx.offset())
                .input(idx.data_size()),
        ],
        &[],
        grid,
        super::block_dims(grid),
        true,
    )
}

pub(super) fn evaluate(
    operation: &ops::Operation,
    inputs: &[Array],
    output: &Array,
) -> Result<(), String> {
    match (operation, inputs) {
        (ops::Operation::Select, [a, b, c]) => select_evaluate([a, b, c], output),
        (ops::Operation::SearchSorted(right), [a, v]) => {
            super::sort::searchsorted_evaluate(a, v, output, *right)
        }
        (ops::Operation::GatherAxis(axis), [src, idx]) => gather_evaluate(src, idx, output, *axis),
        (ops::Operation::ScatterAxis(axis), [src, idx, upd]) => {
            scatter_evaluate(src, idx, upd, output, *axis)
        }
        (ops::Operation::ScatterAddRow, [src, idx, upd]) => {
            scatter_add_row_evaluate(src, idx, upd, output)
        }
        (ops::Operation::GatherTake(axis), [src, idx]) => {
            gather_take_evaluate(src, idx, output, *axis)
        }
        (
            ops::Operation::ArangeI32(_)
            | ops::Operation::ArangeU32
            | ops::Operation::ArangeI64
            | ops::Operation::ArangeF32 { .. },
            [],
        ) => {
            let nbytes = output.layout().nbytes();
            output.set_data(Buffer::mtplx_bytes(nbytes as u64)?)?;
            let size = output.layout().size() as u32;
            if size == 0 {
                return Ok(());
            }
            if let ops::Operation::ArangeF32 { start, step } = operation {
                // Original arange_set_scalars casts start and next before
                // subtracting: casting step alone changes large-start ranges.
                let next = (start + step) as f32;
                let start = *start as f32;
                return dispatch_geometry(
                    "arangefloat32",
                    &[
                        bytes(0, &start),
                        bytes(1, &(next - start)),
                        tensor(2, &output.buffer()).output(output.data_size()),
                    ],
                    &[],
                    [size, 1, 1],
                    [size.min(views::kernel_max_threads("arangefloat32")?), 1, 1],
                    true,
                );
            }
            let wide = output.layout().dtype() == Dtype::I64;
            let start = match operation {
                ops::Operation::ArangeI32(start) => *start,
                _ => 0,
            };
            let name = match output.layout().dtype() {
                Dtype::I64 => "arangeint64",
                Dtype::U32 => "arangeuint32",
                _ => "arangeint32",
            };
            dispatch_geometry(
                name,
                &[
                    if wide {
                        bytes(0, &0_i64)
                    } else {
                        bytes(0, &start)
                    },
                    if wide {
                        bytes(1, &1_i64)
                    } else {
                        bytes(1, &1_i32)
                    },
                    tensor(2, &output.buffer()).output(output.data_size()),
                ],
                &[],
                [size, 1, 1],
                [size.min(views::kernel_max_threads(name)?), 1, 1],
                true,
            )
        }
        _ => Err("invalid indexing primitive arity".into()),
    }
}

fn select_evaluate(inputs: [&Array; 3], output: &Array) -> Result<(), String> {
    let [a, b, c] = inputs;
    let all_scalar = inputs.iter().all(|a| a.data_size() == 1);
    let all_vector = inputs.iter().all(|a| a.layout().flags().row_contiguous)
        || inputs.iter().all(|a| a.layout().flags().col_contiguous);
    let sv = b.data_size() == 1
        && a.layout().flags().row_contiguous
        && c.layout().flags().row_contiguous;
    let vs = c.data_size() == 1
        && a.layout().flags().row_contiguous
        && b.layout().flags().row_contiguous;
    let general = !(all_scalar || all_vector || sv || vs);
    // common/ternary.h: SSS never donates; VVV tries a,b,c regardless of
    // row/column layout; mixed/general tries only row-contiguous donors.
    if all_scalar {
        let l = b.layout();
        output.set_strided_data(
            Buffer::mtplx_bytes(output.itemsize() as u64)?,
            1,
            l.strides(),
            l.flags(),
            0,
        )?;
    } else if let Some(source) = inputs
        .iter()
        .find(|a| (all_vector || a.layout().flags().row_contiguous) && views::can_donate(a, output))
    {
        let l = source.layout();
        output.copy_shared_buffer(source, l.strides(), l.flags(), source.data_size(), 0)?;
    } else if all_vector {
        let l = b.layout();
        output.set_strided_data(
            Buffer::mtplx_bytes((output.itemsize() * b.data_size()) as u64)?,
            b.data_size(),
            l.strides(),
            l.flags(),
            0,
        )?;
    } else {
        let nbytes = output.layout().nbytes();
        output.set_data(Buffer::mtplx_bytes(nbytes as u64)?)?;
    }
    if output.layout().size() == 0 {
        return Ok(());
    }
    let layouts = inputs.map(Array::layout);
    let out = output.layout();
    let (shape, strides) = views::collapse_joint(
        out.shape(),
        &[
            layouts[0].strides(),
            layouts[1].strides(),
            layouts[2].strides(),
            out.strides(),
        ],
    );
    let ndim = shape.len();
    let large = if general {
        inputs.iter().any(|a| a.data_size() > i32::MAX as usize) || out.size() > i32::MAX as usize
    } else {
        output.data_size() > i32::MAX as usize
    };
    let work = if general {
        if large { 4 } else { 2 }
    } else if output.data_size() < 65536 {
        1
    } else {
        (8 / output.itemsize()).max(1)
    };
    let prefix = if general {
        format!(
            "g{}{}",
            if ndim <= 3 {
                ndim.to_string()
            } else {
                format!("n{work}")
            },
            if large { "large" } else { "" }
        )
    } else {
        format!(
            "{}{}",
            if all_scalar || all_vector {
                "v"
            } else if sv {
                "sv"
            } else {
                "vs"
            },
            if large {
                "2"
            } else if work > 1 {
                "n"
            } else {
                ""
            }
        )
    };
    let name = format!("{prefix}_Select{}", out.dtype().kernel_name());
    let max = views::kernel_max_threads(&name)?;
    let buffers = inputs.map(Array::buffer);
    let out_buffer = output.buffer();
    let mut bindings = inputs
        .iter()
        .zip(&buffers)
        .enumerate()
        .map(|(i, (a, b))| {
            tensor(i as u32, b)
                .at_byte_offset(a.offset())
                .input(a.data_size())
        })
        .collect::<Vec<_>>();
    bindings.push(
        tensor(3, &out_buffer)
            .at_byte_offset(output.offset())
            .output(output.data_size()),
    );
    let count32 = output.data_size() as i32;
    let count64 = output.data_size() as i64;
    let (grid, group) = if general {
        if max != 1024 {
            return Err("[Metal::ternary] Must use 1024 sized block".into());
        }
        let dim0 = shape.last().copied().unwrap_or(1) as usize;
        let dim1 = shape.iter().rev().nth(1).copied().unwrap_or(1) as usize;
        if ndim > 3 {
            bindings.extend([
                bytes(4, shape.as_slice()),
                bytes(5, strides[0].as_slice()),
                bytes(6, strides[1].as_slice()),
                bytes(7, strides[2].as_slice()),
                bytes(8, &ndim),
            ]);
        } else {
            bindings.extend([
                bytes(4, strides[0].as_slice()),
                bytes(5, strides[1].as_slice()),
                bytes(6, strides[2].as_slice()),
            ]);
        }
        let grid = [
            if ndim > 3 { dim0.div_ceil(work) } else { dim0 },
            dim1,
            out.size() / (dim0 * dim1),
        ];
        let grid = grid.map(|n| u32::try_from(n).map_err(|_| "Select grid overflow"));
        let grid = [grid[0]?, grid[1]?, grid[2]?];
        (grid, super::block_dims(grid))
    } else {
        bindings.push(if large {
            bytes(4, &count64)
        } else {
            bytes(4, &count32)
        });
        let threads = output.data_size().div_ceil(work);
        let grid = if large {
            views::large_copy_grid(&out, work)?
        } else {
            [threads as u32, 1, 1]
        };
        (grid, [threads.min(max as usize) as u32, 1, 1])
    };
    dispatch_geometry(&name, &bindings, &[], grid, group, true)
}

fn gather_evaluate(src: &Array, idx: &Array, output: &Array, axis: usize) -> Result<(), String> {
    let nbytes = output.layout().nbytes();
    output.set_data(Buffer::mtplx_bytes(nbytes as u64)?)?;
    if output.layout().size() == 0 {
        return Ok(());
    }
    let sl = src.layout();
    let il = idx.layout();
    let large = sl.size() > i32::MAX as usize || il.size() > i32::MAX as usize;
    let name = format!(
        "kernel_qwen_mtplx_gather_axis_{}_{}{}_{}{}",
        match sl.dtype() {
            Dtype::I64 => "i64",
            Dtype::Bool => "bool",
            Dtype::F32 => "f32",
            Dtype::F16 => "f16",
            Dtype::U32 => "u32",
            _ => "bf16",
        },
        if il.dtype() == Dtype::I64 {
            "idxi64_"
        } else if il.dtype() == Dtype::I32 {
            "idxi32_"
        } else {
            ""
        },
        if large { "int64_t" } else { "int" },
        u8::from(sl.flags().row_contiguous),
        u8::from(il.flags().row_contiguous)
    );
    let pre = il.shape()[..axis]
        .iter()
        .map(|&n| n as usize)
        .product::<usize>();
    let post = il.shape()[axis + 1..]
        .iter()
        .map(|&n| n as usize)
        .product::<usize>();
    let grid = [
        u32::try_from(post).map_err(|_| "GatherAxis grid overflow")?,
        il.shape()[axis] as u32,
        u32::try_from(pre).map_err(|_| "GatherAxis grid overflow")?,
    ];
    let mut shape = il.shape().to_vec();
    shape.remove(axis);
    let mut src_strides = sl.strides().to_vec();
    src_strides.remove(axis);
    let mut idx_strides = il.strides().to_vec();
    idx_strides.remove(axis);
    let ndim = shape.len();
    let axis_i32 = axis as i32;
    let src_buffer = src.buffer();
    let idx_buffer = idx.buffer();
    let out = output.buffer();
    let bindings = [
        tensor(0, &src_buffer)
            .at_byte_offset(src.offset())
            .input(src.data_size()),
        tensor(1, &idx_buffer)
            .at_byte_offset(idx.offset())
            .input(idx.data_size()),
        tensor(2, &out).output(output.data_size()),
        bytes(3, shape.as_slice()),
        bytes(4, src_strides.as_slice()),
        bytes(5, idx_strides.as_slice()),
        bytes(6, &ndim),
        bytes(7, &axis_i32),
        bytes(8, &sl.shape()[axis]),
        bytes(9, &sl.strides()[axis]),
        bytes(10, &il.strides()[axis]),
    ];
    dispatch_geometry(&name, &bindings, &[], grid, super::block_dims(grid), true)
}

#[test]
#[ignore = "requires Apple Metal; pinned Bool ScatterAxis layout and broadcast paths"]
fn mtplx_scatter_axis_graph_matches_reference() {
    use super::super::{
        gpu::Context,
        qwen_mtplx_tests::{capture_dispatches, verify_typed},
    };
    use super::{
        allocator, configure_sources,
        stream::{Device, Streams},
    };
    configure_sources().unwrap();
    let _context = Context::open_qwen(0).unwrap();
    let streams = Streams::new(Some(allocator::Allocator::new().unwrap()));
    let gpu = streams.default_stream(Device::Gpu).unwrap();
    let cases = include_str!("../../../../tests/fixtures/mtplx-scatter-axis.jsonl")
        .lines()
        .map(|l| serde_json::from_str::<serde_json::Value>(l).unwrap())
        .collect::<Vec<_>>();
    assert_eq!(cases.len(), 64);
    for prepared in [false, true] {
        for case in &cases {
            super::eval::execution_tests::observed(|| {
                let shape = case["shape"]
                    .as_array()
                    .unwrap()
                    .iter()
                    .map(|v| v.as_i64().unwrap() as i32)
                    .collect::<Vec<_>>();
                let axis = case["axis"].as_u64().unwrap() as usize;
                let strided = case["strided"].as_bool().unwrap();
                let scalar = case["scalar"].as_bool().unwrap();
                let source_layout = case["src_layout"].as_str().unwrap();
                let n = shape.iter().product::<i32>();
                let boolean = |shape: &[i32], values: &[u8]| {
                    Array::new(shape, Dtype::Bool, super::scalar_buffer(values).unwrap()).unwrap()
                };
                let src = if source_layout == "scalar" {
                    ops::broadcast_to(&boolean(&[], &[0]), &shape, gpu).unwrap()
                } else {
                    let mult = if source_layout == "strided" { 2 } else { 1 };
                    let values = (0..n * mult + 1)
                        .map(|i| u8::from(i % 3 == 0))
                        .collect::<Vec<_>>();
                    let a = boolean(&[n * mult + 1], &values);
                    let offset = i32::from(source_layout == "offset");
                    let a = ops::slice(&a, &[offset], &[n * mult + offset], &[mult], gpu).unwrap();
                    ops::reshape(&a, &shape, gpu).unwrap()
                };
                let mut ishape = shape.clone();
                ishape[axis] = 2;
                let size = ishape.iter().product::<i32>();
                let mult = if strided { 2 } else { 1 };
                let values = (0..size * mult)
                    .flat_map(|i| (if i / mult % 2 == 0 { -1_i64 } else { 0 }).to_le_bytes())
                    .collect::<Vec<_>>();
                let idx = Array::new(
                    &[size * mult],
                    Dtype::I64,
                    super::scalar_buffer(&values).unwrap(),
                )
                .unwrap();
                let idx = if strided {
                    ops::slice(&idx, &[0], &[size * mult], &[2], gpu).unwrap()
                } else {
                    idx
                };
                let idx = ops::reshape(&idx, &ishape, gpu).unwrap();
                let upd = if scalar {
                    boolean(&[], &[1])
                } else {
                    boolean(
                        &ishape,
                        &(0..size).map(|i| u8::from(i % 3 == 0)).collect::<Vec<_>>(),
                    )
                };
                if prepared {
                    ops::evaluate(
                        &streams,
                        &[src.clone(), idx.clone(), upd.clone()],
                        gpu,
                        false,
                    )
                    .unwrap();
                }
                let out = put_along_axis(&src, &idx, &upd, axis as i32, gpu).unwrap();
                let (_, calls) = capture_dispatches(|| {
                    ops::evaluate(&streams, std::slice::from_ref(&out), gpu, false).unwrap()
                });
                let name = format!(
                    "kernel_qwen_mtplx_scatter_axis_bool_idxi64_int_{}{}",
                    u8::from(!scalar),
                    u8::from(!strided)
                );
                assert_eq!(calls.iter().filter(|c| c.0 == name).count(), 1, "{calls:?}");
                let read = ops::contiguous(
                    &ops::astype(&out, Dtype::F32, false, gpu).unwrap(),
                    false,
                    gpu,
                )
                .unwrap();
                ops::evaluate(&streams, std::slice::from_ref(&read), gpu, false).unwrap();
                verify_typed(
                    case["kernel"].as_str().unwrap(),
                    &[(&read.buffer(), n as u32, 4)],
                );
            });
        }
    }
    let empty = Array::new(&[0], Dtype::Bool, super::scalar_buffer(&[]).unwrap()).unwrap();
    let idx = Array::new(&[0], Dtype::I64, super::scalar_buffer(&[]).unwrap()).unwrap();
    assert!(put_along_axis(&empty, &idx, &empty, 1, gpu).is_err());
    let out = put_along_axis(&empty, &idx, &empty, 0, gpu).unwrap();
    assert_eq!(out.id(), empty.id());
    let src = Array::new(
        &[3, 5],
        Dtype::Bool,
        super::scalar_buffer(&[0; 15]).unwrap(),
    )
    .unwrap();
    let idx = Array::new(
        &[1, 2],
        Dtype::I64,
        super::scalar_buffer(&[0_i64.to_le_bytes(), (-1_i64).to_le_bytes()].concat()).unwrap(),
    )
    .unwrap();
    let upd = Array::new(
        &[1, 1, 2],
        Dtype::Bool,
        super::scalar_buffer(&[1, 0]).unwrap(),
    )
    .unwrap();
    let identity = src.buffer().data_identity();
    let out = put_along_axis(&src, &idx, &upd, -1, gpu).unwrap();
    drop(src);
    let (_, calls) = capture_dispatches(|| {
        ops::evaluate(&streams, std::slice::from_ref(&out), gpu, false).unwrap()
    });
    assert_eq!(out.buffer().data_identity(), identity);
    assert_eq!(
        calls.len(),
        1,
        "Vector donation must not insert a copy: {calls:?}"
    );
    let mut bytes = [0; 15];
    out.buffer().read(out.offset(), &mut bytes).unwrap();
    assert_eq!(bytes, [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0]);
    let wrong_upd = Array::new(
        &[2, 1, 2],
        Dtype::Bool,
        super::scalar_buffer(&[1; 4]).unwrap(),
    )
    .unwrap();
    assert!(put_along_axis(&out, &idx, &wrong_upd, -1, gpu).is_err());
    let empty_idx = Array::new(&[3, 0], Dtype::I64, super::scalar_buffer(&[]).unwrap()).unwrap();
    let empty_upd = Array::new(&[3, 0], Dtype::Bool, super::scalar_buffer(&[]).unwrap()).unwrap();
    let copied = put_along_axis(&out, &empty_idx, &empty_upd, -1, gpu).unwrap();
    let (_, calls) = capture_dispatches(|| {
        ops::evaluate(&streams, std::slice::from_ref(&copied), gpu, false).unwrap()
    });
    assert!(calls.iter().all(|c| !c.0.contains("scatter_axis")));
    copied.buffer().read(copied.offset(), &mut bytes).unwrap();
    assert_eq!(bytes, [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0]);
}

#[test]
#[ignore = "requires Apple Metal; Select layout branches, GatherAxis views and empty arguments"]
fn mtplx_graph_indexing_preserves_layouts_and_offsets() {
    super::eval::execution_tests::observed(|| {
        use super::super::gpu::Context;
        use super::{
            allocator, configure_sources,
            stream::{Device, Streams},
        };
        configure_sources().unwrap();
        let _context = Context::open_qwen(0).unwrap();
        let streams = Streams::new(Some(allocator::Allocator::new().unwrap()));
        let gpu = streams.default_stream(Device::Gpu).unwrap();
        let bf16 = |shape: &[i32], values: &[f32]| {
            Array::new(
                shape,
                Dtype::BF16,
                super::scalar_buffer(
                    &values
                        .iter()
                        .flat_map(|v| ((v.to_bits() >> 16) as u16).to_le_bytes())
                        .collect::<Vec<_>>(),
                )
                .unwrap(),
            )
            .unwrap()
        };
        let yes = bf16(&[2, 3], &[1., 2., 3., 4., 5., 6.]);
        let no = bf16(&[2, 3], &[7., 8., 9., 10., 11., 12.]);
        let nine = bf16(&[], &[9.]);
        let condition = Array::new(
            &[2, 3],
            Dtype::Bool,
            super::scalar_buffer(&[1, 0, 1, 0, 1, 0]).unwrap(),
        )
        .unwrap();
        let scalar_true =
            Array::new(&[], Dtype::Bool, super::scalar_buffer(&[1]).unwrap()).unwrap();
        let row_mask =
            Array::new(&[2, 1], Dtype::Bool, super::scalar_buffer(&[1, 0]).unwrap()).unwrap();
        let cases = [
            (select(&scalar_true, &nine, &nine, gpu).unwrap(), vec![9.]),
            (
                select(&condition, &yes, &no, gpu).unwrap(),
                vec![1., 8., 3., 10., 5., 12.],
            ),
            (
                select(&condition, &nine, &no, gpu).unwrap(),
                vec![9., 8., 9., 10., 9., 12.],
            ),
            (
                select(&condition, &yes, &nine, gpu).unwrap(),
                vec![1., 9., 3., 9., 5., 9.],
            ),
            (
                select(&row_mask, &yes, &nine, gpu).unwrap(),
                vec![1., 2., 3., 9., 9., 9.],
            ),
            (
                select(
                    &ops::transpose(&condition, &[1, 0], gpu).unwrap(),
                    &ops::transpose(&yes, &[1, 0], gpu).unwrap(),
                    &ops::transpose(&no, &[1, 0], gpu).unwrap(),
                    gpu,
                )
                .unwrap(),
                vec![1., 10., 8., 5., 3., 12.],
            ),
        ];
        let check = |array: &Array, expected: &[f32]| {
            let dense = ops::contiguous(array, false, gpu).unwrap();
            ops::evaluate(&streams, std::slice::from_ref(&dense), gpu, false).unwrap();
            let mut bytes = vec![0; dense.layout().nbytes()];
            dense.buffer().read(dense.offset(), &mut bytes).unwrap();
            let actual = bytes
                .chunks_exact(2)
                .map(|b| f32::from_bits(u32::from(u16::from_le_bytes(b.try_into().unwrap())) << 16))
                .collect::<Vec<_>>();
            assert_eq!(actual, expected);
        };
        for (array, expected) in cases {
            check(&array, &expected);
        }
        let indices = Array::new(
            &[1, 2],
            Dtype::I32,
            super::scalar_buffer(
                &[-1_i32, 0]
                    .into_iter()
                    .flat_map(i32::to_le_bytes)
                    .collect::<Vec<_>>(),
            )
            .unwrap(),
        )
        .unwrap();
        check(
            &take_along_axis(&yes, &indices, -1, gpu).unwrap(),
            &[3., 1., 6., 4.],
        );
        let reversed = ops::slice(&yes, &[0, 2], &[2, -4], &[1, -1], gpu).unwrap();
        check(
            &take_along_axis(&reversed, &indices, 1, gpu).unwrap(),
            &[1., 3., 4., 6.],
        );
        let column = ops::transpose(&yes, &[1, 0], gpu).unwrap();
        // Token IDs are I32, including when selecting packed U32 embedding
        // weights. Cover both front-gather and the strided Gather wrapper.
        for dtype in [Dtype::U32, Dtype::F16] {
            let table = ops::astype(&yes, dtype, false, gpu).unwrap();
            let gathered = take_rows(&table, &indices, gpu).unwrap();
            check(
                &ops::astype(&gathered, Dtype::BF16, false, gpu).unwrap(),
                &[4., 5., 6., 1., 2., 3.],
            );
            let table = ops::transpose(&table, &[1, 0], gpu).unwrap();
            let gathered = take(&table, &indices, 1, gpu).unwrap();
            check(
                &ops::astype(&gathered, Dtype::BF16, false, gpu).unwrap(),
                &[4., 1., 5., 2., 6., 3.],
            );
        }
        let unsigned = Array::new(
            &[2, 1],
            Dtype::U32,
            super::scalar_buffer(
                &[2_u32, 0]
                    .into_iter()
                    .flat_map(u32::to_le_bytes)
                    .collect::<Vec<_>>(),
            )
            .unwrap(),
        )
        .unwrap();
        check(
            &take_along_axis(&column, &unsigned, 0, gpu).unwrap(),
            &[3., 6., 1., 4.],
        );
        for dtype in [Dtype::F16, Dtype::F32, Dtype::U32] {
            let table = ops::astype(&yes, dtype, false, gpu).unwrap();
            let gathered = take_along_axis(&table, &indices, -1, gpu).unwrap();
            check(
                &ops::astype(&gathered, Dtype::BF16, false, gpu).unwrap(),
                &[3., 1., 6., 4.],
            );
            let table = ops::transpose(&table, &[1, 0], gpu).unwrap();
            let gathered = take_along_axis(&table, &unsigned, 0, gpu).unwrap();
            check(
                &ops::astype(&gathered, Dtype::BF16, false, gpu).unwrap(),
                &[3., 6., 1., 4.],
            );
        }
        // Rank one must pass genuinely empty shape/stride bindings through
        // the native bridge, including noncontiguous source/index views.
        let vector = ops::reshape(&yes, &[-1], gpu).unwrap();
        let vector = ops::slice(&vector, &[5], &[-7], &[-1], gpu).unwrap();
        let idx_vector = ops::reshape(&indices, &[-1], gpu).unwrap();
        let idx_vector = ops::slice(&idx_vector, &[1], &[-3], &[-1], gpu).unwrap();
        check(
            &take_along_axis(&vector, &idx_vector, 0, gpu).unwrap(),
            &[6., 1.],
        );
        let empty = arange_i32(-3, gpu).unwrap();
        ops::evaluate(&streams, std::slice::from_ref(&empty), gpu, false).unwrap();
        assert_eq!(empty.layout().size(), 0);
        assert!(take_along_axis(&yes, &indices, 2, gpu).is_err());
        assert!(take_along_axis(&yes, &idx_vector, 1, gpu).is_err());
        assert!(select(&indices, &yes, &nine, gpu).is_err());
        // QSA's validity GatherAxis also supports noncontiguous Bool/I64
        // views. Exercise all four source/index row-contiguity combinations.
        let boolean = Array::new(
            &[2, 3],
            Dtype::Bool,
            super::scalar_buffer(&[1, 1, 0, 0, 1, 0]).unwrap(),
        )
        .unwrap();
        let ids = Array::new(
            &[2, 2],
            Dtype::I64,
            super::scalar_buffer(
                &[-1_i64, 0, 1, 0]
                    .into_iter()
                    .flat_map(i64::to_le_bytes)
                    .collect::<Vec<_>>(),
            )
            .unwrap(),
        )
        .unwrap();
        for reverse_source in [false, true] {
            for reverse_indices in [false, true] {
                let src = if reverse_source {
                    ops::slice(&boolean, &[0, 2], &[2, -4], &[1, -1], gpu).unwrap()
                } else {
                    boolean.clone()
                };
                let idx = if reverse_indices {
                    ops::slice(&ids, &[0, 1], &[2, -3], &[1, -1], gpu).unwrap()
                } else {
                    ids.clone()
                };
                let out = take_along_axis(&src, &idx, -1, gpu).unwrap();
                let expected = match (reverse_source, reverse_indices) {
                    (false, false) => [0., 1., 1., 0.],
                    (false, true) => [1., 0., 0., 1.],
                    (true, false) => [1., 0., 1., 0.],
                    (true, true) => [0., 1., 0., 1.],
                };
                check(
                    &ops::astype(&out, Dtype::BF16, false, gpu).unwrap(),
                    &expected,
                );
            }
        }
        streams.clear_streams().unwrap();
    });
}