Save inference parity implementation and evaluation harness

src/engine/metal/qwen_mtplx/unary.rs

@@ -0,0 +1,294 @@
+//! 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,
+    #[cfg_attr(not(test), expect(dead_code, reason = "Reference diagnostic API"))]
+    Tanh,
+    Sqrt,
+    Rsqrt,
+    Sin,
+    Cos,
+}
+
+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();
+    });
+}