From e4edd51f49fe3aef239978073085149a7668619b Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 17:09:37 -0400 Subject: [PATCH 1/2] jouleclaw-omni: read mmap'd GGUF bytes without assuming their alignment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `clippy::cast_ptr_alignment` flagged eight casts from `*const u8` to more-strictly-aligned pointers. The previous commit allowed the lint crate-wide with a note that GGUF aligns tensor data in practice. That note was true and the reasoning was still wrong: these functions take plain `&[u8]` slices and a `metal::Buffer`, and neither carries the alignment the cast asserts. let src_f32 = std::slice::from_raw_parts(src.as_ptr().cast::(), src.len() / 4); `src` is an mmap'd region at whatever offset the file chose. Misaligned `from_raw_parts` is undefined behaviour, not a slow load, and the compiler is entitled to assume it cannot happen. The three that looked safest were the least safe to reason about: `contents()` on a buffer built with `new_buffer_with_bytes_no_copy` is the mmap'd pointer itself, so "it's a Metal allocation, those are page aligned" does not hold for them either. All eight now decode byte-wise through `from_le_bytes` over `as_chunks::()`, which needs no unsafe alignment claim at all -- `u8` has alignment 1. The two staging loops keep one cast each, for the destination, which really is a fresh Metal allocation; the comment says which side is which and why. The crate-level allow is gone, and `clippy -- -D warnings` passes without it. Four unit tests, because the rewrite moved every byte index and a wrong one would be silent: - the bf16 and f32 buffer walks must agree with the scalar converters element for element, over zero, ±1, pi, +inf, a NaN and f16's max; - conversion from an odd offset on BOTH sides, which is the case the old code could not express without UB; - a trailing partial element is dropped rather than read past the end. The metal CI job also now runs every target under `--features metal`. It ran only `--test metal_kernels`, which left the crate's 46 metal-gated lib tests unrun -- a feature's unit tests being as invisible to CI as its code was is the same mistake one level down. That job now runs 48 lib tests and the 26 kernel tests. Still outstanding from the same sweep: `mut_from_ref` on `MetalBuffer::as_mut_slice(&self)`, which hands out a `&mut` from a `&`. Co-Authored-By: Claude Opus 5 --- .github/workflows/jouleclaw-fabrics.yml | 6 +- .../jouleclaw-omni/src/hal/metal/lazy_load.rs | 190 ++++++++++++++---- .../src/inference/architecture/canary.rs | 37 +++- .../crates/jouleclaw-omni/src/lib.rs | 6 - 4 files changed, 176 insertions(+), 63 deletions(-) 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/lazy_load.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/hal/metal/lazy_load.rs index 9216629e7..895dabd97 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, 3.141_592_7, 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, 3.141_592_7]; + + // 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..e93cb82bd 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/lib.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/lib.rs @@ -134,12 +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. -- 2.54.0 (Apple Git-157) From fdce8f2faa89f772c2050907fadbe813f89dc80c Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 20 Oct 2026 17:43:22 -0400 Subject: [PATCH 2/2] MetalBuffer: a method that writes to the buffer takes &mut self `clippy::mut_from_ref` flagged pub unsafe fn as_mut_slice(&self) -> &mut [T] which hands out a mutable slice from a shared reference. #329 allowed the lint, reasoning that the method is `unsafe` so the caller carries the requirement. That conflated two different requirements. The documented one -- "must ensure GPU is not accessing this buffer" -- is genuinely the caller's. The one the signature creates is not: two callers holding a `&MetalBuffer` can each take a `&mut [T]` over the same bytes, and aliasing `&mut` is undefined behaviour on the Rust side whatever the GPU is doing. Both now take `&mut self`: `as_mut_slice`, and `copy_from_slice`, which has the same shape and was not flagged only because it does not hand the slice back. The doc comments say which half of the contract the borrow settles and which is still the caller's. This costs nothing. `MetalBuffer::as_slice`, `as_mut_slice` and `copy_from_slice` have no callers anywhere in the workspace, and `MetalBuffer` itself is used nowhere outside jouleclaw-omni, so there is no caller churn to review -- which is why this landed here rather than waiting for an API decision. I had assumed otherwise before checking. With this and the alignment commit, both soundness allows added in #329 are gone and `clippy -- -D warnings` passes without either. One thing the sweep caught in my own work: the new conversion tests used 3.141_592_7 as a literal, and `approx_constant` -- which I deliberately did not allow, being correctness-adjacent -- failed the build until they used `core::f32::consts::PI`. Co-Authored-By: Claude Opus 5 --- .../crates/jouleclaw-omni/src/hal/metal/buffer.rs | 14 ++++++++++++-- .../jouleclaw-omni/src/hal/metal/lazy_load.rs | 4 ++-- .../jouleclaw-rs/crates/jouleclaw-omni/src/lib.rs | 4 ---- 3 files changed, 14 insertions(+), 8 deletions(-) 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 895dabd97..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 @@ -1690,7 +1690,7 @@ mod conversion_tests { #[test] fn f32_buffer_walk_matches_the_scalar_converter() { - let values: [f32; 5] = [0.0, 1.0, -1.0, 3.141_592_7, 65504.0]; + 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()); @@ -1712,7 +1712,7 @@ mod conversion_tests { /// 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, 3.141_592_7]; + 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]; diff --git a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/lib.rs b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/lib.rs index e93cb82bd..bfbb17a9c 100644 --- a/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/lib.rs +++ b/jouleclaw/jouleclaw-rs/crates/jouleclaw-omni/src/lib.rs @@ -134,10 +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)] -// `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. -- 2.54.0 (Apple Git-157)