295 lines
12 KiB
Rust
295 lines
12 KiB
Rust
//! Graph unary primitives: ops.cpp and metal/unary.cpp from the pinned runtime.
|
|
use super::array::{Array, Dtype, Layout};
|
|
use super::stream::{Device, Stream};
|
|
use super::{bytes, dispatch_geometry, ops, tensor, views};
|
|
|
|
#[derive(Clone, Copy, Debug)]
|
|
pub(super) enum Unary {
|
|
LogicalNot,
|
|
Abs,
|
|
Sign,
|
|
Negative,
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
Square,
|
|
Exp,
|
|
#[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
|
|
Log,
|
|
Sigmoid,
|
|
Tanh,
|
|
Sqrt,
|
|
Rsqrt,
|
|
Sin,
|
|
Cos,
|
|
Erf,
|
|
}
|
|
|
|
pub(super) fn unary(input: &Array, op: Unary, stream: Stream) -> Result<Array, String> {
|
|
let dtype = input.layout().dtype();
|
|
if stream.device() != Device::Gpu {
|
|
return Err("CPU unary arithmetic is not ported".into());
|
|
}
|
|
if matches!(
|
|
(op, dtype),
|
|
(Unary::Negative, Dtype::Bool) | (Unary::Sigmoid, Dtype::C64)
|
|
) {
|
|
return Err(format!("{op:?} does not support {dtype:?}"));
|
|
}
|
|
let dtype = if matches!(op, Unary::LogicalNot) {
|
|
Dtype::Bool
|
|
} else if matches!(
|
|
op,
|
|
Unary::Negative | Unary::Square | Unary::Abs | Unary::Sign
|
|
) {
|
|
dtype
|
|
} else {
|
|
dtype.at_least_float()
|
|
};
|
|
let input = ops::astype(input, dtype, false, stream)?;
|
|
let output = Array::make_operation(
|
|
stream,
|
|
&input,
|
|
Layout::new(input.layout().shape(), dtype)?,
|
|
ops::Operation::Unary(op),
|
|
)?;
|
|
if matches!(op, Unary::Abs) && dtype == Dtype::C64 {
|
|
ops::astype(&output, Dtype::F32, false, stream)
|
|
} else {
|
|
Ok(output)
|
|
}
|
|
}
|
|
|
|
pub(super) fn evaluate(input: &Array, output: &Array, op: Unary) -> Result<(), String> {
|
|
views::set_unary_output_data(input, output)?;
|
|
let src = input.layout();
|
|
let dst = output.layout();
|
|
if src.size() == 0 {
|
|
return Ok(());
|
|
}
|
|
let contiguous = src.flags().contiguous;
|
|
let size = input.data_size();
|
|
let large = if contiguous {
|
|
size > u32::MAX as usize
|
|
} else {
|
|
size > i32::MAX as usize || dst.size() > i32::MAX as usize
|
|
};
|
|
let work = if contiguous {
|
|
if size < 1 << 16 {
|
|
1
|
|
} else {
|
|
(8 / src.dtype().itemsize()).max(1)
|
|
}
|
|
} else if large {
|
|
4
|
|
} else {
|
|
1
|
|
};
|
|
let prefix = if contiguous {
|
|
if large {
|
|
"v2"
|
|
} else if work > 1 {
|
|
"vn"
|
|
} else {
|
|
"v"
|
|
}
|
|
} else if large {
|
|
"gn4large"
|
|
} else {
|
|
"gn1"
|
|
};
|
|
let name = format!(
|
|
"{prefix}_{op:?}{}{}",
|
|
src.dtype().kernel_name(),
|
|
dst.dtype().kernel_name()
|
|
);
|
|
let max_threads = views::kernel_max_threads(&name)?;
|
|
let source = input.buffer();
|
|
let target = output.buffer();
|
|
let mut bindings = vec![
|
|
tensor(0, &source)
|
|
.at_byte_offset(input.offset())
|
|
.input(size),
|
|
tensor(1, &target)
|
|
.at_byte_offset(output.offset())
|
|
.output(output.data_size()),
|
|
];
|
|
if contiguous {
|
|
let n32 = size as i32;
|
|
let n64 = size as i64;
|
|
bindings.push(if large {
|
|
bytes(2, &n64)
|
|
} else {
|
|
bytes(2, &n32)
|
|
});
|
|
let grid = if large {
|
|
views::large_copy_grid(&dst, work)?
|
|
} else {
|
|
[size.div_ceil(work) as u32, 1, 1]
|
|
};
|
|
let group = [(max_threads as usize).min(size.div_ceil(work)) as u32, 1, 1];
|
|
dispatch_geometry(&name, &bindings, &[], grid, group, true)
|
|
} else {
|
|
let (shape, strides) = views::collapse_dims(&src);
|
|
let ndim = i32::try_from(shape.len()).map_err(|_| "unary rank overflow")?;
|
|
let dim0 = shape.last().copied().unwrap_or(1) as usize;
|
|
let dim1 = shape.iter().rev().nth(1).copied().unwrap_or(1) as usize;
|
|
let rest = dst.size() / (dim0 * dim1);
|
|
let grid = [dim0.div_ceil(work), dim1, rest]
|
|
.map(|d| u32::try_from(d).map_err(|_| "unary grid overflow"));
|
|
let grid = [grid[0]?, grid[1]?, grid[2]?];
|
|
if max_threads != 1024 {
|
|
return Err("[Metal::unary] Must use 1024 sized block".into());
|
|
}
|
|
bindings.extend([
|
|
bytes(2, shape.as_slice()),
|
|
bytes(3, strides.as_slice()),
|
|
bytes(4, &ndim),
|
|
]);
|
|
dispatch_geometry(&name, &bindings, &[], grid, super::block_dims(grid), true)
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "requires Apple Metal; graph unary layout, donation and kernel boundaries"]
|
|
fn mtplx_graph_unary_preserves_reference_layout_and_values() {
|
|
super::eval::execution_tests::observed(|| {
|
|
use super::super::gpu::{Buffer, Context};
|
|
use super::{allocator, configure_sources, stream::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();
|
|
// Exact dyadic values only; no host transcendental approximation.
|
|
let encode = |dtype: Dtype, values: &[f32]| -> Vec<u8> {
|
|
values
|
|
.iter()
|
|
.flat_map(|&v| {
|
|
let f = v.to_bits();
|
|
let bits = match dtype {
|
|
Dtype::BF16 => f >> 16,
|
|
Dtype::F16 => {
|
|
(f >> 16 & 0x8000)
|
|
| if f & 0x7fffffff == 0 {
|
|
0
|
|
} else {
|
|
(((f >> 23 & 255) - 112) << 10) | (f >> 13 & 1023)
|
|
}
|
|
}
|
|
Dtype::I32 => v as i32 as u32,
|
|
_ => f,
|
|
};
|
|
bits.to_le_bytes()[..dtype.itemsize()].to_vec()
|
|
})
|
|
.collect()
|
|
};
|
|
for dtype in [Dtype::F16, Dtype::BF16, Dtype::F32] {
|
|
for (op, value, expected) in [
|
|
(Unary::Negative, 4., -4.),
|
|
(Unary::Abs, -4., 4.),
|
|
(Unary::Sign, -4., -1.),
|
|
(Unary::Square, 4., 16.),
|
|
(Unary::Exp, 0., 1.),
|
|
(Unary::Sigmoid, 0., 0.5),
|
|
(Unary::Tanh, 0., 0.),
|
|
(Unary::Sqrt, 4., 2.),
|
|
(Unary::Rsqrt, 4., 0.5),
|
|
] {
|
|
for count in [0, 1, 65535, 65536, 65537] {
|
|
for keep_alias in [false, true] {
|
|
let data = encode(dtype, &vec![value; count]);
|
|
let buffer = Buffer::mtplx_bytes(data.len() as u64).unwrap();
|
|
buffer.write(0, &data).unwrap();
|
|
let input = Array::new(&[count as i32], dtype, buffer).unwrap();
|
|
let identity = input.buffer().data_identity();
|
|
let out = unary(&input, op, gpu).unwrap();
|
|
let alias = keep_alias.then(|| input.clone());
|
|
drop(input);
|
|
ops::evaluate(&streams, std::slice::from_ref(&out), gpu, false).unwrap();
|
|
assert_eq!(out.buffer().data_identity() == identity, !keep_alias);
|
|
let mut actual = vec![0; count * dtype.itemsize()];
|
|
out.buffer().read(out.offset(), &mut actual).unwrap();
|
|
assert_eq!(
|
|
actual,
|
|
encode(dtype, &vec![expected; count]),
|
|
"{op:?} {dtype:?}"
|
|
);
|
|
drop(alias);
|
|
}
|
|
}
|
|
}
|
|
for mode in 0..5 {
|
|
let values = [1., 4., 1., 4., 1., 4., 1., 4.];
|
|
let bytes = encode(dtype, &values);
|
|
let buffer = Buffer::mtplx_bytes(bytes.len() as u64).unwrap();
|
|
buffer.write(0, &bytes).unwrap();
|
|
let base = Array::new(&[2, 4], dtype, buffer).unwrap();
|
|
let view = match mode {
|
|
0 => ops::transpose(&base, &[1, 0], gpu).unwrap(),
|
|
1 => ops::slice(&base, &[0, 1], &[2, 4], &[1, 2], gpu).unwrap(),
|
|
2 => ops::slice(&base, &[1, 3], &[-3, -5], &[-1, -1], gpu).unwrap(),
|
|
3 => {
|
|
let row = ops::slice(&base, &[1, 0], &[2, 4], &[1, 1], gpu).unwrap();
|
|
ops::broadcast_to(&row, &[3, 4], gpu).unwrap()
|
|
}
|
|
_ => ops::slice(&base, &[1, 0], &[2, 4], &[1, 1], gpu).unwrap(),
|
|
};
|
|
let identity = base.buffer().data_identity();
|
|
let out = unary(&view, Unary::Square, gpu).unwrap();
|
|
drop((view, base));
|
|
ops::evaluate(&streams, std::slice::from_ref(&out), gpu, false).unwrap();
|
|
let expected: &[f32] = match mode {
|
|
0 => &[1., 16., 1., 16., 1., 16., 1., 16.],
|
|
1 => &[16., 16., 16., 16.],
|
|
2 => &[16., 1., 16., 1., 16., 1., 16., 1.],
|
|
_ => &[1., 16., 1., 16.],
|
|
};
|
|
let mut actual = vec![0; out.data_size() * dtype.itemsize()];
|
|
out.buffer().read(out.offset(), &mut actual).unwrap();
|
|
assert_eq!(actual, encode(dtype, expected), "layout {mode}");
|
|
// check_contiguity counts positive strides only: reversed
|
|
// and stepped inputs both require a fresh dense output.
|
|
assert_eq!(
|
|
out.buffer().data_identity() == identity,
|
|
!matches!(mode, 1 | 2)
|
|
);
|
|
if mode == 3 {
|
|
assert_eq!(out.layout().strides(), [0, 1]);
|
|
}
|
|
}
|
|
let oversized = Buffer::mtplx_bytes(20000).unwrap();
|
|
oversized.write(0, &encode(dtype, &[1., 4.])).unwrap();
|
|
let identity = oversized.data_identity();
|
|
let input = Array::new(&[2], dtype, oversized).unwrap();
|
|
let out = unary(&input, Unary::Square, gpu).unwrap();
|
|
drop(input);
|
|
ops::evaluate(&streams, std::slice::from_ref(&out), gpu, false).unwrap();
|
|
assert_ne!(
|
|
out.buffer().data_identity(),
|
|
identity,
|
|
"16KiB donation slack"
|
|
);
|
|
let mut actual = vec![0; 2 * dtype.itemsize()];
|
|
out.buffer().read(0, &mut actual).unwrap();
|
|
assert_eq!(actual, encode(dtype, &[1., 16.]));
|
|
}
|
|
let buffer = Buffer::mtplx_bytes(4).unwrap();
|
|
buffer.write(0, &4_i32.to_le_bytes()).unwrap();
|
|
let input = Array::new(&[], Dtype::I32, buffer).unwrap();
|
|
let out = unary(&input, Unary::Sqrt, gpu).unwrap();
|
|
ops::evaluate(&streams, std::slice::from_ref(&out), gpu, false).unwrap();
|
|
assert_eq!(out.layout().dtype(), Dtype::F32);
|
|
let mut actual = [0; 4];
|
|
out.buffer().read(0, &mut actual).unwrap();
|
|
assert_eq!(actual, 2_f32.to_le_bytes());
|
|
let squared = unary(&input, Unary::Square, gpu).unwrap();
|
|
ops::evaluate(&streams, std::slice::from_ref(&squared), gpu, false).unwrap();
|
|
assert_eq!(squared.layout().dtype(), Dtype::I32);
|
|
squared.buffer().read(0, &mut actual).unwrap();
|
|
assert_eq!(actual, 16_i32.to_le_bytes());
|
|
for (dtype, op) in [(Dtype::Bool, Unary::Negative), (Dtype::C64, Unary::Sigmoid)] {
|
|
let input = Array::new(&[0], dtype, Buffer::mtplx_bytes(0).unwrap()).unwrap();
|
|
assert!(unary(&input, op, gpu).is_err());
|
|
}
|
|
streams.clear_streams().unwrap();
|
|
});
|
|
}
|