1356 lines
47 KiB
Rust
1356 lines
47 KiB
Rust
//! 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, ©)?;
|
|
let cl = copy.layout();
|
|
output.copy_shared_buffer(©, 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();
|
|
});
|
|
}
|