diff --git a/.github/workflows/jouleclaw-fabrics.yml b/.github/workflows/jouleclaw-fabrics.yml index c2f139dc8..88403456c 100644 --- a/.github/workflows/jouleclaw-fabrics.yml +++ b/.github/workflows/jouleclaw-fabrics.yml @@ -76,8 +76,10 @@ jobs: echo "::error::Either the metal cfg stopped holding or tests were removed." exit 1 fi - cargo test --manifest-path "$M" -p jouleclaw-omni --features metal \ - --test metal_kernels + # Run every target, not just that one. `--test metal_kernels` left + # the crate's 46 metal-gated *lib* tests unrun: a feature's unit + # tests are as invisible to CI as its code was. + cargo test --manifest-path "$M" -p jouleclaw-omni --features metal linux-fabrics: name: ${{ matrix.feature }} (linux, RTX 4050 + Vulkan loader) diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/hal/metal/buffer.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/hal/metal/buffer.rs index ee5998da0..f743cd820 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/hal/metal/buffer.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/hal/metal/buffer.rs @@ -118,9 +118,15 @@ impl MetalBuffer { /// Get mutable contents (zero-copy for Shared storage). /// + /// Takes `&mut self`: handing out a `&mut [T]` from a `&self` lets two + /// callers hold aliasing mutable slices of the same buffer, which is + /// undefined behaviour on the Rust side regardless of what the GPU is + /// doing. The exclusive borrow rules that out; the GPU half of the + /// contract is still the caller's, below. + /// /// # Safety /// Must ensure GPU is not accessing this buffer. - pub unsafe fn as_mut_slice(&self) -> &mut [T] { + pub unsafe fn as_mut_slice(&mut self) -> &mut [T] { let ptr = self.buffer.contents().cast::(); let count = self.size / std::mem::size_of::(); unsafe { std::slice::from_raw_parts_mut(ptr, count) } @@ -128,9 +134,13 @@ impl MetalBuffer { /// Copy data to this buffer. /// + /// `&mut self` for the same reason as [`Self::as_mut_slice`]: this writes + /// through the buffer's contents pointer, and a shared reference does not + /// express that. + /// /// # Safety /// Must ensure GPU is not reading from this buffer. - pub unsafe fn copy_from_slice(&self, data: &[T]) -> Result<()> { + pub unsafe fn copy_from_slice(&mut self, data: &[T]) -> Result<()> { let bytes = bytemuck::cast_slice(data); if bytes.len() > self.size { return Err(Error::InvalidArgument { diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/hal/metal/lazy_load.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/hal/metal/lazy_load.rs index 9216629e7..8322bdc76 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/hal/metal/lazy_load.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/hal/metal/lazy_load.rs @@ -90,30 +90,35 @@ fn bf16_to_f16(bf16_bits: u16) -> u16 { } /// Convert a buffer of BF16 values to F16 values in-place or to new buffer. +/// +/// Reads and writes through `from_le_bytes`/`to_le_bytes` rather than casting +/// `&[u8]` to `&[u16]`. Both slices come from mmap'd GGUF regions at offsets +/// the file chooses, so nothing here guarantees 2-byte alignment, and +/// `slice::from_raw_parts` on a misaligned pointer is undefined behaviour -- +/// not merely a slower load. GGUF does align tensor data in practice, but +/// that is a property of the file, not of these arguments. fn convert_bf16_to_f16_buffer(src: &[u8], dst: &mut [u8]) { assert_eq!(src.len(), dst.len()); assert!(src.len().is_multiple_of(2)); - let src_u16 = unsafe { - std::slice::from_raw_parts(src.as_ptr().cast::(), src.len() / 2) - }; - let dst_u16 = unsafe { - std::slice::from_raw_parts_mut(dst.as_mut_ptr().cast::(), dst.len() / 2) - }; - - for (i, &bf16) in src_u16.iter().enumerate() { - dst_u16[i] = bf16_to_f16(bf16); + let (src_pairs, _) = src.as_chunks::<2>(); + let (dst_pairs, _) = dst.as_chunks_mut::<2>(); + for (s, d) in src_pairs.iter().zip(dst_pairs.iter_mut()) { + *d = bf16_to_f16(u16::from_le_bytes(*s)).to_le_bytes(); } } /// Convert a buffer of F32 values to F16 values in-place or to new buffer. +/// +/// Byte-wise for the same reason as `convert_bf16_to_f16_buffer`: neither +/// slice is known to be 4- or 2-byte aligned. fn convert_f32_to_f16_buffer(src: &[u8], dst: &mut [u8]) { assert_eq!(src.len() / 4, dst.len() / 2); - let src_f32 = unsafe { std::slice::from_raw_parts(src.as_ptr().cast::(), src.len() / 4) }; - let dst_f16 = unsafe { std::slice::from_raw_parts_mut(dst.as_mut_ptr().cast::(), dst.len() / 2) }; - - for (i, &f) in src_f32.iter().enumerate() { - dst_f16[i] = f16::from_f32(f); + + let (src_quads, _) = src.as_chunks::<4>(); + let (dst_pairs, _) = dst.as_chunks_mut::<2>(); + for (s, d) in src_quads.iter().zip(dst_pairs.iter_mut()) { + *d = f16::from_f32(f32::from_le_bytes(*s)).to_le_bytes(); } } @@ -325,31 +330,49 @@ impl LazyTensor { /// On Apple Silicon UMA, reads directly from shared memory. pub fn to_f32_vec(&self) -> Result> { let numel = self.shape.numel(); - unsafe { - let ptr = self.raw_buffer.contents() as *const u8; - match self.dtype { - DType::F32 => { - let f32_ptr = ptr.cast::(); - let slice = std::slice::from_raw_parts(f32_ptr, numel); - Ok(slice.to_vec()) - } - DType::F16 => { - let f16_ptr = ptr.cast::(); - let slice = std::slice::from_raw_parts(f16_ptr, numel); - Ok(slice.iter().map(|v| v.to_f32()).collect()) - } - DType::BF16 => { - let u16_ptr = ptr.cast::(); - let slice = std::slice::from_raw_parts(u16_ptr, numel); - Ok(slice.iter().map(|&bits| { - f32::from_bits(u32::from(bits) << 16) - }).collect()) - } - _ => Err(Error::internal(format!( - "Unsupported dtype {:?} for to_f32_vec", self.dtype - ))), + let width = match self.dtype { + DType::F32 => 4, + DType::F16 | DType::BF16 => 2, + other => { + return Err(Error::internal(format!( + "Unsupported dtype {other:?} for to_f32_vec" + ))) } - } + }; + + // Read as bytes and decode element-wise rather than casting the + // pointer. `contents()` is not known to carry f32 or u16 alignment: + // when the buffer was built with `new_buffer_with_bytes_no_copy` it is + // the mmap'd pointer itself, at whatever offset the GGUF file put the + // tensor. `u8` has alignment 1, so this slice is always well formed. + let bytes = unsafe { + std::slice::from_raw_parts( + self.raw_buffer.contents() as *const u8, + numel * width, + ) + }; + + Ok(match self.dtype { + DType::F32 => bytes + .as_chunks::<4>() + .0 + .iter() + .map(|&c| f32::from_le_bytes(c)) + .collect(), + DType::F16 => bytes + .as_chunks::<2>() + .0 + .iter() + .map(|&c| f16::from_le_bytes(c).to_f32()) + .collect(), + // BF16 is the top 16 bits of the F32 pattern. + _ => bytes + .as_chunks::<2>() + .0 + .iter() + .map(|&c| f32::from_bits(u32::from(u16::from_le_bytes(c)) << 16)) + .collect(), + }) } /// Advise OS this data won't be needed soon. @@ -935,13 +958,15 @@ impl LazyLoader { | metal::MTLResourceOptions::CPUCacheModeWriteCombined, ); unsafe { - let src = std::slice::from_raw_parts(data.as_ptr().cast::(), numel); + // `contents()` is a Metal allocation, which is page + // aligned, so this cast is sound. `data` is the mmap'd + // region and is not, so it is read byte-wise. let dst = std::slice::from_raw_parts_mut( new_buffer.contents().cast::(), numel, ); - for i in 0..numel { - dst[i] = f32_to_f16(src[i]); + for (d, c) in dst.iter_mut().zip(data.as_chunks::<4>().0) { + *d = f32_to_f16(f32::from_le_bytes(*c)); } } (new_buffer, DType::F16) @@ -1112,13 +1137,13 @@ impl LazyLoader { | metal::MTLResourceOptions::CPUCacheModeWriteCombined, ); unsafe { - let src = std::slice::from_raw_parts(data.as_ptr().cast::(), numel); + // As above: the Metal allocation is aligned, the mmap is not. let dst = std::slice::from_raw_parts_mut( new_buffer.contents().cast::(), numel, ); - for i in 0..numel { - dst[i] = f32_to_f16(src[i]); + for (d, c) in dst.iter_mut().zip(data.as_chunks::<4>().0) { + *d = f32_to_f16(f32::from_le_bytes(*c)); } } (new_buffer, false, GgmlType::F16) @@ -1638,3 +1663,80 @@ impl LazyLoader { #[cfg(not(feature = "metal"))] pub struct LazyModel; + +#[cfg(all(test, feature = "metal"))] +mod conversion_tests { + use super::{bf16_to_f16, convert_bf16_to_f16_buffer, convert_f32_to_f16_buffer}; + use half::f16; + + /// The buffer walk must agree with the scalar converter element for + /// element. This is what a wrong index or byte order would break. + #[test] + fn bf16_buffer_walk_matches_the_scalar_converter() { + // zero, 1.0, -1.0, pi-ish, +inf, a quiet NaN. + let bits: [u16; 6] = [0x0000, 0x3F80, 0xBF80, 0x4049, 0x7F80, 0xFFC0]; + let mut src = Vec::new(); + for b in bits { + src.extend_from_slice(&b.to_le_bytes()); + } + let mut dst = vec![0u8; src.len()]; + convert_bf16_to_f16_buffer(&src, &mut dst); + + for (i, b) in bits.iter().enumerate() { + let got = u16::from_le_bytes([dst[2 * i], dst[2 * i + 1]]); + assert_eq!(got, bf16_to_f16(*b), "element {i} (bf16 bits {b:#06x})"); + } + } + + #[test] + fn f32_buffer_walk_matches_the_scalar_converter() { + let values: [f32; 5] = [0.0, 1.0, -1.0, core::f32::consts::PI, 65504.0]; + let mut src = Vec::new(); + for v in values { + src.extend_from_slice(&v.to_le_bytes()); + } + let mut dst = vec![0u8; values.len() * 2]; + convert_f32_to_f16_buffer(&src, &mut dst); + + for (i, v) in values.iter().enumerate() { + let got = f16::from_le_bytes([dst[2 * i], dst[2 * i + 1]]); + assert_eq!(got, f16::from_f32(*v), "element {i} ({v})"); + } + } + + /// The reason these read byte-wise. The previous implementation cast + /// `&[u8]` to `&[f32]`/`&[u16]`; at an odd offset that pointer cannot + /// satisfy the alignment those types require, and `from_raw_parts` on it + /// is undefined behaviour rather than a slow read. GGUF tensor data is + /// aligned in practice, but these functions take plain slices and must + /// not depend on where their caller got them. + #[test] + fn conversion_is_correct_from_odd_offsets_on_both_sides() { + let values: [f32; 4] = [0.0, 1.0, -2.5, core::f32::consts::PI]; + + // One leading pad byte makes 4-byte (and 2-byte) alignment impossible. + let mut src_backing = vec![0u8; 1 + values.len() * 4]; + for (i, v) in values.iter().enumerate() { + src_backing[1 + i * 4..1 + i * 4 + 4].copy_from_slice(&v.to_le_bytes()); + } + let mut dst_backing = vec![0u8; 1 + values.len() * 2]; + + convert_f32_to_f16_buffer(&src_backing[1..], &mut dst_backing[1..]); + + for (i, v) in values.iter().enumerate() { + let got = f16::from_le_bytes([dst_backing[1 + i * 2], dst_backing[1 + i * 2 + 1]]); + assert_eq!(got, f16::from_f32(*v), "element {i} ({v}) from an odd offset"); + } + } + + /// A trailing partial element is dropped rather than read out of bounds. + #[test] + fn a_trailing_partial_element_is_ignored() { + // Five bytes of "f32" input: one whole element plus one stray byte. + let mut src = 1.0f32.to_le_bytes().to_vec(); + src.push(0x7F); + let mut dst = vec![0u8; 2]; + convert_f32_to_f16_buffer(&src, &mut dst); + assert_eq!(f16::from_le_bytes([dst[0], dst[1]]), f16::from_f32(1.0)); + } +} diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/architecture/canary.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/architecture/canary.rs index 428165dcf..261c9ce1b 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/architecture/canary.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/inference/architecture/canary.rs @@ -236,18 +236,33 @@ impl CanaryPipeline { let numel = shape.numel(); // LazyTensor data is already on the Metal buffer; read as f32 or f16 depending on dtype let buf = lt.buffer(); - let ptr = buf.contents() as *const u8; + // Byte-wise for the same reason as `LazyTensor::to_f32_vec`: a no-copy + // buffer's `contents()` is the mmap'd pointer, which carries no f32 or + // f16 alignment guarantee. + let width = match lt.dtype() { + DType::F32 => 4, + DType::F16 => 2, + other => { + return Err(Error::internal(format!( + "unsupported dtype {other:?} for batch norm" + ))) + } + }; + let bytes = + unsafe { std::slice::from_raw_parts(buf.contents() as *const u8, numel * width) }; match lt.dtype() { - DType::F32 => { - let f32_ptr = ptr.cast::(); - let slice = unsafe { std::slice::from_raw_parts(f32_ptr, numel) }; - Ok(slice.to_vec()) - } - DType::F16 => { - let f16_ptr = ptr.cast::(); - let slice = unsafe { std::slice::from_raw_parts(f16_ptr, numel) }; - Ok(slice.iter().map(|v| v.to_f32()).collect()) - } + DType::F32 => Ok(bytes + .as_chunks::<4>() + .0 + .iter() + .map(|&c| f32::from_le_bytes(c)) + .collect()), + DType::F16 => Ok(bytes + .as_chunks::<2>() + .0 + .iter() + .map(|&c| half::f16::from_le_bytes(c).to_f32()) + .collect()), other => Err(Error::internal(format!("unsupported dtype {other:?} for batch norm"))), } } diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/lib.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/lib.rs index 59724fe3d..bfbb17a9c 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/lib.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/lib.rs @@ -134,16 +134,6 @@ // the non-NaN operand, so the chain scrubs NaN, while `clamp` propagates it. // Converting would let a NaN timestep through a diffusion schedule. #![allow(clippy::manual_clamp)] -// These cast mmap'd GGUF tensor bytes to f32/u16 slices. GGUF aligns tensor -// data (the alignment is a header field, 32 by default), so the pointers are -// aligned in practice -- but that is an invariant of the file format, not -// something the type system is carrying. Worth replacing with a checked cast -// (bytemuck::try_cast_slice) rather than leaving as an assumption. -#![allow(clippy::cast_ptr_alignment)] -// `as_mut_slice(&self)` on a Metal buffer hands out a `&mut` from a `&`. It is -// `unsafe` and the caller carries the aliasing requirement, but the signature -// invites misuse and should take `&mut self` or return a raw pointer. -#![allow(clippy::mut_from_ref)] // Style preferences on existing code: let-else, memcpy and counter loops // written out, Debug impls that omit noisy fields, doc continuation // indentation, and signatures shaped for the caller rather than the lint.