From 684194bd1a12b1465930001ee6938d5064b2a415 Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Fri, 31 Jul 2026 07:00:59 +0900 Subject: [PATCH 1/2] wip: add uninitialized indexed replay --- docs/2026-07-31-issue-187-worklog.md | 33 + strided-kernel/Cargo.toml | 5 + .../benches/issue_187_uninit_indexed.rs | 420 +++++ strided-kernel/src/copy_plan.rs | 166 +- strided-kernel/src/erased.rs | 837 ++++++++- strided-kernel/src/gather_plan.rs | 290 +++- .../tests/issue_187_source_contract.rs | 91 + .../tests/issue_187_uninit_indexed.rs | 1496 +++++++++++++++++ .../tests/issue_187_uninit_reduce.rs | 173 ++ 9 files changed, 3383 insertions(+), 128 deletions(-) create mode 100644 docs/2026-07-31-issue-187-worklog.md create mode 100644 strided-kernel/benches/issue_187_uninit_indexed.rs create mode 100644 strided-kernel/tests/issue_187_source_contract.rs create mode 100644 strided-kernel/tests/issue_187_uninit_indexed.rs create mode 100644 strided-kernel/tests/issue_187_uninit_reduce.rs diff --git a/docs/2026-07-31-issue-187-worklog.md b/docs/2026-07-31-issue-187-worklog.md new file mode 100644 index 0000000..f3bf34a --- /dev/null +++ b/docs/2026-07-31-issue-187-worklog.md @@ -0,0 +1,33 @@ +# Issue #187 Worklog + +## RED History + +- RED: indexed uninitialized destinations required a safe copy-then-update + boundary; direct initialized backing views were forbidden. +- GREEN work introduced private raw-pointer writers and a closure-scoped + post-copy receipt. +- Acceptance work adds differential lifecycle and indexed coverage. + +## Architecture and Safety Contract + +- Only CopyPlan constructs the private post-copy receipt. +- Receipt construction is closure-scoped and cannot escape the HRTB helper. +- Uninitialized writers use MaybeUninit storage and full-overwrite or + copy-then-update proofs before any typed read. +- Reduction terminal writers use raw pointers plus validated extents. +- Integer erased scatter uses wrapping i32/i64 combine functions; typed public + scatter semantics remain unchanged. + +## Verification + +Commands: + + cargo fmt --all + cargo test -p strided-kernel --test issue_187_uninit_indexed + cargo test -p strided-kernel --features parallel --test issue_187_uninit_indexed + cargo bench -p strided-kernel --features parallel --bench issue_187_uninit_indexed --no-run + +Benchmark results: not run in this pass. The runner records 31 alternating +initialized/uninitialized samples and reports no invented timing values. + +Remaining verification: Miri lifecycle execution and the full repository gate. diff --git a/strided-kernel/Cargo.toml b/strided-kernel/Cargo.toml index ba0829f..9ea7464 100644 --- a/strided-kernel/Cargo.toml +++ b/strided-kernel/Cargo.toml @@ -69,3 +69,8 @@ required-features = ["parallel"] name = "issue_184_uninit_replay" harness = false required-features = ["parallel"] + +[[bench]] +name = "issue_187_uninit_indexed" +harness = false +required-features = ["parallel"] diff --git a/strided-kernel/benches/issue_187_uninit_indexed.rs b/strided-kernel/benches/issue_187_uninit_indexed.rs new file mode 100644 index 0000000..0d758f5 --- /dev/null +++ b/strided-kernel/benches/issue_187_uninit_indexed.rs @@ -0,0 +1,420 @@ +use core::mem::MaybeUninit; +use std::{ + env, + hint::black_box, + process, + time::{Duration, Instant}, +}; +use strided_kernel::{ + ErasedDynamicSlicePlan, ErasedDynamicUpdateSlicePlan, ErasedGatherPlan, ErasedRawStridedMut, + ErasedRawStridedPtr, ErasedRawStridedRef, ErasedRawStridedUninitMut, ErasedReducePlan, + ErasedScatterPlan, ExecContext, GatherSpec, KernelDType, ReduceOp, ScatterSpec, +}; + +fn bytes(value: &[T]) -> &[u8] { + unsafe { core::slice::from_raw_parts(value.as_ptr().cast(), core::mem::size_of_val(value)) } +} +fn bytes_mut(value: &mut [T]) -> &mut [u8] { + unsafe { + core::slice::from_raw_parts_mut(value.as_mut_ptr().cast(), core::mem::size_of_val(value)) + } +} + +/// View aligned f64 storage as erased MaybeUninit bytes without initializing it. +/// +/// # Safety +/// The cast preserves the allocation and byte length; callers retain the +/// exclusive borrow for the returned view. +unsafe fn uninit_f64_bytes(value: &mut [MaybeUninit]) -> &mut [MaybeUninit] { + core::slice::from_raw_parts_mut( + value.as_mut_ptr().cast::>(), + core::mem::size_of_val(value), + ) +} + +fn sample( + ctx: &ExecContext, + plan: &ErasedReducePlan, + source: &ErasedRawStridedRef<'_>, + init: &mut [f64], + raw: &mut [MaybeUninit], + initialized_first: bool, +) -> (Duration, Duration) { + let initialized = || { + let mut dest = + ErasedRawStridedMut::new(KernelDType::F64, bytes_mut(init), &[], &[], 0).unwrap(); + let start = Instant::now(); + plan.execute(ctx, &mut dest, source).unwrap(); + black_box(init); + start.elapsed() + }; + let uninitialized = || { + let mut dest = ErasedRawStridedUninitMut::new( + KernelDType::F64, + unsafe { uninit_f64_bytes(raw) }, + &[], + &[], + 0, + ) + .unwrap(); + let ptr = ErasedRawStridedPtr::from_ref(source); + let start = Instant::now(); + plan.execute_uninit(ctx, &mut dest, &ptr).unwrap(); + black_box(raw); + start.elapsed() + }; + if initialized_first { + (initialized(), uninitialized()) + } else { + let uninit = uninitialized(); + let init = initialized(); + (init, uninit) + } +} + +fn report(label: &str, initialized: &[Duration], uninitialized: &[Duration]) -> bool { + let mut logs = initialized + .iter() + .zip(uninitialized) + .map(|(init, uninit)| (uninit.as_secs_f64() / init.as_secs_f64()).ln()) + .collect::>(); + logs.sort_by(f64::total_cmp); + let mean = logs.iter().sum::() / logs.len() as f64; + let variance = logs.iter().map(|v| (v - mean).powi(2)).sum::() / (logs.len() - 1) as f64; + let upper95 = (mean + 1.645 * variance.sqrt() / (logs.len() as f64).sqrt()).exp(); + let median = logs[logs.len() / 2].exp(); + println!( + "{label}: median={median:.4} upper95={upper95:.4} n={}", + logs.len() + ); + upper95 <= 1.20 +} + +fn run_pairs( + label: &str, + mut initialized: impl FnMut() -> Duration, + mut uninitialized: impl FnMut() -> Duration, +) -> bool { + let mut init_times = Vec::with_capacity(31); + let mut uninit_times = Vec::with_capacity(31); + for sample in 0..31 { + if sample % 2 == 0 { + init_times.push(initialized()); + uninit_times.push(uninitialized()); + } else { + let uninit = uninitialized(); + let init = initialized(); + init_times.push(init); + uninit_times.push(uninit); + } + } + report(label, &init_times, &uninit_times) +} + +fn main() { + let threads = env::args().nth(1).and_then(|v| v.parse().ok()).unwrap_or(1); + let ctx = ExecContext::max_threads(threads).expect("threads must be positive"); + let n = 131_073; + let dims = [n]; + let strides = [1isize]; + let input: Vec = (0..n).map(|i| (i as f64) * 0.25).collect(); + let source = + ErasedRawStridedRef::new(KernelDType::F64, bytes(&input), &dims, &strides, 0).unwrap(); + let plan = ErasedReducePlan::compile(KernelDType::F64, ReduceOp::Sum, &dims, &strides).unwrap(); + let mut init = vec![0.0f64; 1]; + let mut raw = vec![MaybeUninit::::uninit(); 1]; + let mut initialized = Vec::with_capacity(31); + let mut uninitialized = Vec::with_capacity(31); + for sample_index in 0..31 { + let (a, b) = sample( + &ctx, + &plan, + &source, + &mut init, + &mut raw, + sample_index % 2 == 0, + ); + initialized.push(a); + uninitialized.push(b); + } + let mut ok = report("reduce", &initialized, &uninitialized); + + let operand = (0..8192).map(|i| i as f64).collect::>(); + let operand_dims = [8192usize]; + let index_dims = [2048usize]; + let dest_dims = [2048usize, 4]; + let indices = (0..2048).map(|i| (i * 3 % 8192) as i64).collect::>(); + let gather = ErasedGatherPlan::compile( + KernelDType::F64, + KernelDType::I64, + &operand_dims, + &[1], + &index_dims, + &[1], + &dest_dims, + &[1, 2048], + GatherSpec { + offset_dims: vec![1], + collapsed_slice_dims: vec![], + start_index_map: vec![0], + index_vector_dim: 1, + slice_sizes: vec![4], + }, + ) + .unwrap(); + let operand_ref = + ErasedRawStridedRef::new(KernelDType::F64, bytes(&operand), &operand_dims, &[1], 0) + .unwrap(); + let index_ref = + ErasedRawStridedRef::new(KernelDType::I64, bytes(&indices), &index_dims, &[1], 0).unwrap(); + let mut gather_init = vec![0.0f64; 8192]; + let mut gather_raw = vec![MaybeUninit::::uninit(); 8192]; + ok &= run_pairs( + "gather_window", + || { + let mut dest = ErasedRawStridedMut::new( + KernelDType::F64, + bytes_mut(&mut gather_init), + &dest_dims, + &[1, 2048], + 0, + ) + .unwrap(); + let start = Instant::now(); + gather + .execute(&ctx, &mut dest, &operand_ref, &index_ref) + .unwrap(); + black_box(&gather_init); + start.elapsed() + }, + || { + let mut dest = ErasedRawStridedUninitMut::new( + KernelDType::F64, + unsafe { uninit_f64_bytes(&mut gather_raw) }, + &dest_dims, + &[1, 2048], + 0, + ) + .unwrap(); + let operand = ErasedRawStridedPtr::from_ref(&operand_ref); + let index = ErasedRawStridedPtr::from_ref(&index_ref); + let start = Instant::now(); + gather + .execute_uninit(&ctx, &mut dest, &operand, &index) + .unwrap(); + black_box(&gather_raw); + start.elapsed() + }, + ); + + let starts = [128i64]; + let start_dims = [1usize]; + let update_dims = [4096usize]; + let update_values = (0..4096).map(|i| i as f64).collect::>(); + let slice = ErasedDynamicSlicePlan::compile( + KernelDType::F64, + KernelDType::I64, + &[8192], + &[1], + &start_dims, + &[1], + &update_dims, + &[1], + &[4096], + ) + .unwrap(); + let starts_ref = + ErasedRawStridedRef::new(KernelDType::I64, bytes(&starts), &start_dims, &[1], 0).unwrap(); + let mut slice_init = vec![0.0f64; 4096]; + let mut slice_raw = vec![MaybeUninit::::uninit(); 4096]; + ok &= run_pairs( + "dynamic_slice", + || { + let mut dest = ErasedRawStridedMut::new( + KernelDType::F64, + bytes_mut(&mut slice_init), + &update_dims, + &[1], + 0, + ) + .unwrap(); + let start = Instant::now(); + slice + .execute(&ctx, &mut dest, &operand_ref, &starts_ref) + .unwrap(); + black_box(&slice_init); + start.elapsed() + }, + || { + let mut dest = ErasedRawStridedUninitMut::new( + KernelDType::F64, + unsafe { uninit_f64_bytes(&mut slice_raw) }, + &update_dims, + &[1], + 0, + ) + .unwrap(); + let operand = ErasedRawStridedPtr::from_ref(&operand_ref); + let starts = ErasedRawStridedPtr::from_ref(&starts_ref); + let start = Instant::now(); + slice + .execute_uninit(&ctx, &mut dest, &operand, &starts) + .unwrap(); + black_box(&slice_raw); + start.elapsed() + }, + ); + + let update_ref = ErasedRawStridedRef::new( + KernelDType::F64, + bytes(&update_values), + &update_dims, + &[1], + 0, + ) + .unwrap(); + let update_plan = ErasedDynamicUpdateSlicePlan::compile( + KernelDType::F64, + KernelDType::I64, + &[8192], + &[1], + &start_dims, + &[1], + &update_dims, + &[1], + &[8192], + &[1], + ) + .unwrap(); + let mut update_init = vec![0.0f64; 8192]; + let mut update_raw = vec![MaybeUninit::::uninit(); 8192]; + ok &= run_pairs( + "dynamic_update_slice", + || { + let mut dest = ErasedRawStridedMut::new( + KernelDType::F64, + bytes_mut(&mut update_init), + &[8192], + &[1], + 0, + ) + .unwrap(); + let start = Instant::now(); + update_plan + .execute(&ctx, &mut dest, &operand_ref, &update_ref, &starts_ref) + .unwrap(); + black_box(&update_init); + start.elapsed() + }, + || { + let mut dest = ErasedRawStridedUninitMut::new( + KernelDType::F64, + unsafe { uninit_f64_bytes(&mut update_raw) }, + &[8192], + &[1], + 0, + ) + .unwrap(); + let operand = ErasedRawStridedPtr::from_ref(&operand_ref); + let update = ErasedRawStridedPtr::from_ref(&update_ref); + let starts = ErasedRawStridedPtr::from_ref(&starts_ref); + let start = Instant::now(); + update_plan + .execute_uninit(&ctx, &mut dest, &operand, &update, &starts) + .unwrap(); + black_box(&update_raw); + start.elapsed() + }, + ); + + let scatter_indices = (0..4096).map(|i| (i % 8192) as i64).collect::>(); + let scatter_updates = vec![1.0f64; 4096]; + let scatter_dims = [8192usize]; + let scatter_index_dims = [4096usize, 1]; + let scatter_update_dims = [4096usize]; + let scatter = ErasedScatterPlan::compile( + KernelDType::F64, + KernelDType::I64, + &scatter_dims, + &[1], + &scatter_index_dims, + &[1, 4096], + &scatter_update_dims, + &[1], + &scatter_dims, + &[1], + ScatterSpec { + update_window_dims: vec![], + inserted_window_dims: vec![0], + scatter_dims_to_operand_dims: vec![0], + index_vector_dim: 1, + }, + ) + .unwrap(); + let scatter_index_ref = ErasedRawStridedRef::new( + KernelDType::I64, + bytes(&scatter_indices), + &scatter_index_dims, + &[1, 4096], + 0, + ) + .unwrap(); + let scatter_update_ref = ErasedRawStridedRef::new( + KernelDType::F64, + bytes(&scatter_updates), + &scatter_update_dims, + &[1], + 0, + ) + .unwrap(); + let mut scatter_init = vec![0.0f64; 8192]; + let mut scatter_raw = vec![MaybeUninit::::uninit(); 8192]; + let serial = ExecContext::serial(); + ok &= run_pairs( + "scatter_serial", + || { + let mut dest = ErasedRawStridedMut::new( + KernelDType::F64, + bytes_mut(&mut scatter_init), + &scatter_dims, + &[1], + 0, + ) + .unwrap(); + let start = Instant::now(); + scatter + .execute( + &serial, + &mut dest, + &operand_ref, + &scatter_index_ref, + &scatter_update_ref, + ) + .unwrap(); + black_box(&scatter_init); + start.elapsed() + }, + || { + let mut dest = ErasedRawStridedUninitMut::new( + KernelDType::F64, + unsafe { uninit_f64_bytes(&mut scatter_raw) }, + &scatter_dims, + &[1], + 0, + ) + .unwrap(); + let operand = ErasedRawStridedPtr::from_ref(&operand_ref); + let index = ErasedRawStridedPtr::from_ref(&scatter_index_ref); + let update = ErasedRawStridedPtr::from_ref(&scatter_update_ref); + let start = Instant::now(); + scatter + .execute_uninit(&serial, &mut dest, &operand, &index, &update) + .unwrap(); + black_box(&scatter_raw); + start.elapsed() + }, + ); + if !ok { + process::exit(1); + } +} diff --git a/strided-kernel/src/copy_plan.rs b/strided-kernel/src/copy_plan.rs index 4687ee0..356d9da 100644 --- a/strided-kernel/src/copy_plan.rs +++ b/strided-kernel/src/copy_plan.rs @@ -9,8 +9,11 @@ //! replay it with no planning and no heap allocation for ranks at most //! [`RAW_FUSED_RANK_LIMIT`](crate::RAW_FUSED_RANK_LIMIT). -use core::mem::MaybeUninit; -use core::ops::Mul; +use core::{ + marker::PhantomData, + mem::MaybeUninit, + ops::{Add, Mul}, +}; use crate::map_view::map_raw_into; use crate::ops_view::{copy_conj, copy_into, copy_scale}; @@ -27,6 +30,125 @@ type AxisVec = smallvec::SmallVec<[T; crate::RAW_FUSED_RANK_LIMIT]>; #[cfg(not(feature = "parallel"))] type AxisVec = Vec; +pub(crate) trait OverwriteWriter { + fn dims(&self) -> &[usize]; + fn strides(&self) -> &[isize]; + fn offset(&self) -> isize; + /// # Safety + /// The caller must use the pointer only for the validated allocation and + /// logical layout represented by this writer. + unsafe fn data_ptr(&mut self) -> *mut T; + /// # Safety + /// The offset must be an in-bounds logical destination proven by layout. + unsafe fn write_at(&mut self, offset: isize, value: T); +} + +pub(crate) trait ReadModifyWrite: OverwriteWriter { + /// # Safety + /// The offset must be an in-bounds initialized slot covered by the + /// traversal's copy and disjointness proof. + unsafe fn add_at(&mut self, offset: isize, value: T, combine: fn(T, T) -> T); +} + +impl<'a, T> OverwriteWriter for RawStridedMut<'a, T> { + fn dims(&self) -> &[usize] { + self.dims() + } + fn strides(&self) -> &[isize] { + self.strides() + } + fn offset(&self) -> isize { + self.offset() + } + unsafe fn data_ptr(&mut self) -> *mut T { + self.data_mut().as_mut_ptr() + } + unsafe fn write_at(&mut self, offset: isize, value: T) { + // SAFETY: the prepared layout validates every logical destination. + unsafe { self.data_mut().as_mut_ptr().offset(offset).write(value) } + } +} + +impl<'a, T> ReadModifyWrite for RawStridedMut<'a, T> +where + T: Add, +{ + unsafe fn add_at(&mut self, offset: isize, value: T, combine: fn(T, T) -> T) { + // SAFETY: the copy or initialized caller proves this logical slot. + unsafe { + let ptr = self.data_mut().as_mut_ptr().offset(offset); + ptr.write(combine(ptr.read(), value)); + } + } +} + +impl<'a, T> OverwriteWriter for RawStridedMut<'a, MaybeUninit> { + fn dims(&self) -> &[usize] { + self.dims() + } + fn strides(&self) -> &[isize] { + self.strides() + } + fn offset(&self) -> isize { + self.offset() + } + unsafe fn data_ptr(&mut self) -> *mut T { + self.data_mut().as_mut_ptr().cast() + } + unsafe fn write_at(&mut self, offset: isize, value: T) { + // SAFETY: the prepared layout validates every logical destination. + unsafe { + self.data_mut() + .as_mut_ptr() + .offset(offset) + .write(MaybeUninit::new(value)) + } + } +} + +pub(crate) struct InitializedRawDest<'a, T> { + ptr: *mut T, + extent: usize, + dims: &'a [usize], + strides: &'a [isize], + offset: isize, + _marker: PhantomData<&'a mut [MaybeUninit]>, +} + +impl<'a, T> OverwriteWriter for InitializedRawDest<'a, T> { + fn dims(&self) -> &[usize] { + self.dims + } + fn strides(&self) -> &[isize] { + self.strides + } + fn offset(&self) -> isize { + self.offset + } + unsafe fn data_ptr(&mut self) -> *mut T { + self.ptr + } + unsafe fn write_at(&mut self, offset: isize, value: T) { + debug_assert!(offset >= 0 && (offset as usize) < self.extent); + // SAFETY: the copy proof and extent check cover this logical slot. + unsafe { self.ptr.offset(offset).write(value) } + } +} + +impl<'a, T> ReadModifyWrite for InitializedRawDest<'a, T> +where + T: Add, +{ + unsafe fn add_at(&mut self, offset: isize, value: T, combine: fn(T, T) -> T) { + debug_assert!(offset >= 0 && (offset as usize) < self.extent); + // SAFETY: the copy proof and extent check cover this logical slot. + unsafe { + let ptr = self.ptr.offset(offset); + ptr.write(combine(ptr.read(), value)); + } + } +} + /// A compiled copy traversal for one `(dims, dst_strides, src_strides)` /// layout pair. /// @@ -69,6 +191,28 @@ pub struct CopyPlan { } impl CopyPlan { + pub(crate) fn execute_uninit_then<'a, T, R>( + &self, + dest: &'a mut RawStridedMut<'a, MaybeUninit>, + src: &RawStridedRef<'_, T>, + f: impl for<'b> FnOnce(InitializedRawDest<'b, T>) -> R, + ) -> Result + where + T: Copy + MaybeSendSync, + { + self.execute_uninit(dest, src)?; + let data = dest.data_mut(); + let receipt = InitializedRawDest { + ptr: data.as_mut_ptr().cast(), + extent: data.len(), + dims: dest.dims(), + strides: dest.strides(), + offset: dest.offset(), + _marker: PhantomData, + }; + Ok(f(receipt)) + } + /// Compile a copy plan for the given layout pair. /// /// Performs the layout validation and traversal construction @@ -237,6 +381,24 @@ mod tests { use super::*; use num_complex::{Complex32, Complex64}; + #[test] + fn uninit_then_receipt_drops_after_panic() { + use std::panic::{catch_unwind, AssertUnwindSafe}; + let plan = CopyPlan::compile(&[2], &[1], &[1]).unwrap(); + let source_data = [3i32, 5]; + let source = RawStridedRef::new(&source_data, &[2], &[1], 0).unwrap(); + let result = catch_unwind(AssertUnwindSafe(|| { + let mut storage = vec![MaybeUninit::::uninit(); 3]; + let mut dest = RawStridedMut::new(&mut storage, &[2], &[1], 0).unwrap(); + let _: () = plan + .execute_uninit_then(&mut dest, &source, |_receipt| { + panic!("post-copy update failure"); + }) + .unwrap(); + })); + assert!(result.is_err()); + } + /// Reference: the per-call raw kernel (which itself is differential-tested /// against the view kernels in raw_ops.rs). fn plan_matches_direct( diff --git a/strided-kernel/src/erased.rs b/strided-kernel/src/erased.rs index dc8a8ae..30a6dfd 100644 --- a/strided-kernel/src/erased.rs +++ b/strided-kernel/src/erased.rs @@ -26,6 +26,40 @@ use crate::{ const ERASED_FUSED_INPUT_LIMIT: usize = 4; const SERIAL_REDUCE_LANES: usize = 8; +trait ReduceWriter { + fn offset(&self) -> isize; + /// # Safety + /// The pointer may only be used within the validated destination extent. + unsafe fn ptr(&mut self) -> *mut T; + fn extent(&self) -> usize; + /// # Safety + /// The offset must be an in-bounds logical reduction destination offset. + unsafe fn write_at(&mut self, offset: isize, value: T) { + debug_assert!(offset >= 0 && (offset as usize) < self.extent()); + // SAFETY: reduction layout validation proves the logical offset. + unsafe { self.ptr().offset(offset).write(value) } + } +} + +struct RawReduceWriter<'a, T> { + ptr: *mut T, + extent: usize, + offset: isize, + _marker: core::marker::PhantomData<&'a mut [MaybeUninit]>, +} + +impl<'a, T> ReduceWriter for RawReduceWriter<'a, T> { + fn offset(&self) -> isize { + self.offset + } + unsafe fn ptr(&mut self) -> *mut T { + self.ptr + } + fn extent(&self) -> usize { + self.extent + } +} + /// Runtime unary operation for [`erased_map_into`]. #[non_exhaustive] #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -1029,18 +1063,93 @@ impl ErasedReducePlan { } let result = match self.dtype { - KernelDType::F32 => dispatch_reduce::(self.op, &self.layout, ctx, dest, src), - KernelDType::F64 => dispatch_reduce::(self.op, &self.layout, ctx, dest, src), - KernelDType::I32 => dispatch_reduce::(self.op, &self.layout, ctx, dest, src), - KernelDType::I64 => dispatch_reduce::(self.op, &self.layout, ctx, dest, src), - KernelDType::C32 => dispatch_reduce::(self.op, &self.layout, ctx, dest, src), - KernelDType::C64 => dispatch_reduce::(self.op, &self.layout, ctx, dest, src), + KernelDType::F32 => { + dispatch_reduce::(self.op, &self.layout, ctx, &mut reduce_writer(dest), src) + } + KernelDType::F64 => { + dispatch_reduce::(self.op, &self.layout, ctx, &mut reduce_writer(dest), src) + } + KernelDType::I32 => { + dispatch_reduce::(self.op, &self.layout, ctx, &mut reduce_writer(dest), src) + } + KernelDType::I64 => { + dispatch_reduce::(self.op, &self.layout, ctx, &mut reduce_writer(dest), src) + } + KernelDType::C32 => dispatch_reduce::( + self.op, + &self.layout, + ctx, + &mut reduce_writer(dest), + src, + ), + KernelDType::C64 => dispatch_reduce::( + self.op, + &self.layout, + ctx, + &mut reduce_writer(dest), + src, + ), _ => Err(StridedError::UnsupportedDType { dtype: self.dtype.label(), }), }; result } + + pub fn execute_uninit( + &self, + ctx: &ExecContext, + dest: &mut ErasedRawStridedUninitMut<'_>, + src: &ErasedRawStridedPtr<'_>, + ) -> Result<()> { + check_dtype(self.dtype, dest.dtype())?; + check_dtype(self.dtype, src.dtype())?; + validate_uninit_no_overlap(dest, src, 0)?; + let src = validated_input_ref(src)?; + self.layout.check_src_layout(&src)?; + let src = ErasedRawStridedRef::new( + self.dtype, + src.data(), + src.dims(), + src.strides(), + src.offset(), + )?; + match &self.layout { + ReduceLayout::Full { .. } => { + let total = checked_total_len(dest.dims())?; + if total != 1 { + return Err(StridedError::RankMismatch(total, 1)); + } + } + ReduceLayout::Axes { + dest_dims, + dest_strides, + .. + } => { + if dest.dims() != dest_dims.as_slice() || dest.strides() != dest_strides.as_slice() + { + return Err(StridedError::PlanLayoutMismatch); + } + } + } + macro_rules! run { + ($ty:ty) => {{ + let mut writer = reduce_uninit_writer::<$ty>(dest); + dispatch_reduce::<$ty, _>(self.op, &self.layout, ctx, &mut writer, &src) + }}; + } + match self.dtype { + KernelDType::F32 => run!(f32), + KernelDType::F64 => run!(f64), + KernelDType::I32 => run!(i32), + KernelDType::I64 => run!(i64), + KernelDType::C32 => run!(Complex32), + KernelDType::C64 => run!(Complex64), + _ => Err(StridedError::UnsupportedDType { + dtype: self.dtype.label(), + }), + } + } } impl ErasedGatherPlan { @@ -1106,8 +1215,8 @@ impl ErasedGatherPlan { &self.plan, self.index_dtype, dest, - operand, - start_indices, + &operand, + &start_indices, ), KernelDType::F64 => dispatch_gather_index::( &self.plan, @@ -1157,6 +1266,83 @@ impl ErasedGatherPlan { }); result } + + /// Execute gather into a destination whose reachable slots may be + /// uninitialized. All validation precedes the first destination write. + pub fn execute_uninit( + &self, + ctx: &ExecContext, + dest: &mut ErasedRawStridedUninitMut<'_>, + operand: &ErasedRawStridedPtr<'_>, + start_indices: &ErasedRawStridedPtr<'_>, + ) -> Result<()> { + check_dtype(self.dtype, dest.dtype())?; + check_dtype(self.dtype, operand.dtype())?; + check_dtype(self.index_dtype, start_indices.dtype())?; + validate_uninit_no_overlap(dest, operand, 0)?; + validate_uninit_no_overlap(dest, start_indices, 1)?; + let operand = &validated_input_ref(operand)?; + let start_indices = &validated_input_ref(start_indices)?; + let run = |dest: &mut ErasedRawStridedUninitMut<'_>| match self.dtype { + KernelDType::F32 => execute_gather_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + start_indices, + ), + KernelDType::F64 => execute_gather_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + start_indices, + ), + KernelDType::I32 => execute_gather_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + start_indices, + ), + KernelDType::I64 => execute_gather_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + start_indices, + ), + KernelDType::Bool => execute_gather_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + start_indices, + ), + KernelDType::C32 => execute_gather_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + start_indices, + ), + KernelDType::C64 => execute_gather_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + start_indices, + ), + _ => Err(StridedError::UnsupportedDType { + dtype: self.dtype.label(), + }), + }; + if ctx.is_serial() { + run(dest) + } else { + ctx.run(|| run(dest)) + } + } } impl ErasedDynamicSlicePlan { @@ -1222,8 +1408,8 @@ impl ErasedDynamicSlicePlan { &self.plan, self.index_dtype, dest, - operand, - starts, + &operand, + &starts, ), KernelDType::F64 => dispatch_dynamic_slice_index::( &self.plan, @@ -1273,6 +1459,83 @@ impl ErasedDynamicSlicePlan { }); result } + + /// Execute dynamic slice into a destination whose reachable slots may be + /// uninitialized. + pub fn execute_uninit( + &self, + ctx: &ExecContext, + dest: &mut ErasedRawStridedUninitMut<'_>, + operand: &ErasedRawStridedPtr<'_>, + starts: &ErasedRawStridedPtr<'_>, + ) -> Result<()> { + check_dtype(self.dtype, dest.dtype())?; + check_dtype(self.dtype, operand.dtype())?; + check_dtype(self.index_dtype, starts.dtype())?; + validate_uninit_no_overlap(dest, operand, 0)?; + validate_uninit_no_overlap(dest, starts, 1)?; + let operand = &validated_input_ref(operand)?; + let starts = &validated_input_ref(starts)?; + let run = |dest: &mut ErasedRawStridedUninitMut<'_>| match self.dtype { + KernelDType::F32 => execute_dynamic_slice_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + starts, + ), + KernelDType::F64 => execute_dynamic_slice_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + starts, + ), + KernelDType::I32 => execute_dynamic_slice_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + starts, + ), + KernelDType::I64 => execute_dynamic_slice_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + starts, + ), + KernelDType::Bool => execute_dynamic_slice_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + starts, + ), + KernelDType::C32 => execute_dynamic_slice_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + starts, + ), + KernelDType::C64 => execute_dynamic_slice_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + starts, + ), + _ => Err(StridedError::UnsupportedDType { + dtype: self.dtype.label(), + }), + }; + if ctx.is_serial() { + run(dest) + } else { + ctx.run(|| run(dest)) + } + } } impl ErasedDynamicUpdateSlicePlan { @@ -1342,9 +1605,9 @@ impl ErasedDynamicUpdateSlicePlan { &self.plan, self.index_dtype, dest, - operand, - update, - starts, + &operand, + &update, + &starts, ), KernelDType::F64 => dispatch_dynamic_update_slice_index::( &self.plan, @@ -1400,6 +1663,92 @@ impl ErasedDynamicUpdateSlicePlan { }); result } + + pub fn execute_uninit( + &self, + ctx: &ExecContext, + dest: &mut ErasedRawStridedUninitMut<'_>, + operand: &ErasedRawStridedPtr<'_>, + update: &ErasedRawStridedPtr<'_>, + starts: &ErasedRawStridedPtr<'_>, + ) -> Result<()> { + check_dtype(self.dtype, dest.dtype())?; + check_dtype(self.dtype, operand.dtype())?; + check_dtype(self.dtype, update.dtype())?; + check_dtype(self.index_dtype, starts.dtype())?; + validate_uninit_no_overlap(dest, operand, 0)?; + validate_uninit_no_overlap(dest, update, 1)?; + validate_uninit_no_overlap(dest, starts, 2)?; + let operand = &validated_input_ref(operand)?; + let update = &validated_input_ref(update)?; + let starts = &validated_input_ref(starts)?; + let run = |dest: &mut ErasedRawStridedUninitMut<'_>| match self.dtype { + KernelDType::F32 => execute_dynamic_update_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + update, + starts, + ), + KernelDType::F64 => execute_dynamic_update_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + update, + starts, + ), + KernelDType::I32 => execute_dynamic_update_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + update, + starts, + ), + KernelDType::I64 => execute_dynamic_update_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + update, + starts, + ), + KernelDType::Bool => execute_dynamic_update_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + update, + starts, + ), + KernelDType::C32 => execute_dynamic_update_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + update, + starts, + ), + KernelDType::C64 => execute_dynamic_update_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + update, + starts, + ), + _ => Err(StridedError::UnsupportedDType { + dtype: self.dtype.label(), + }), + }; + if ctx.is_serial() { + run(dest) + } else { + ctx.run(|| run(dest)) + } + } } impl ErasedScatterPlan { @@ -1471,9 +1820,9 @@ impl ErasedScatterPlan { &self.plan, self.index_dtype, dest, - operand, - scatter_indices, - updates, + &operand, + &scatter_indices, + &updates, ), KernelDType::F64 => dispatch_scatter_index::( &self.plan, @@ -1521,7 +1870,95 @@ impl ErasedScatterPlan { }); result } -} + + pub fn execute_uninit( + &self, + ctx: &ExecContext, + dest: &mut ErasedRawStridedUninitMut<'_>, + operand: &ErasedRawStridedPtr<'_>, + scatter_indices: &ErasedRawStridedPtr<'_>, + updates: &ErasedRawStridedPtr<'_>, + ) -> Result<()> { + check_dtype(self.dtype, dest.dtype())?; + check_dtype(self.dtype, operand.dtype())?; + check_dtype(self.dtype, updates.dtype())?; + check_dtype(self.index_dtype, scatter_indices.dtype())?; + validate_uninit_no_overlap(dest, operand, 0)?; + validate_uninit_no_overlap(dest, scatter_indices, 1)?; + validate_uninit_no_overlap(dest, updates, 2)?; + let operand = &validated_input_ref(operand)?; + let scatter_indices = &validated_input_ref(scatter_indices)?; + let updates = &validated_input_ref(updates)?; + let run = |dest: &mut ErasedRawStridedUninitMut<'_>| match self.dtype { + KernelDType::F32 => execute_scatter_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + scatter_indices, + updates, + add_values::, + ), + KernelDType::F64 => execute_scatter_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + scatter_indices, + updates, + add_values::, + ), + KernelDType::I32 => execute_scatter_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + scatter_indices, + updates, + i32::wrapping_add, + ), + KernelDType::I64 => execute_scatter_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + scatter_indices, + updates, + i64::wrapping_add, + ), + KernelDType::C32 => execute_scatter_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + scatter_indices, + updates, + add_values::, + ), + KernelDType::C64 => execute_scatter_uninit_dispatch::( + &self.plan, + self.index_dtype, + dest, + operand, + scatter_indices, + updates, + add_values::, + ), + _ => Err(StridedError::UnsupportedDType { + dtype: self.dtype.label(), + }), + }; + if ctx.is_serial() { + run(dest) + } else { + ctx.run(|| run(dest)) + } + } +} + +fn add_values>(lhs: T, rhs: T) -> T { + lhs + rhs +} fn execute_one_shot_map( op: ErasedMapOp, @@ -1919,6 +2356,34 @@ fn check_dtype(expected: KernelDType, actual: KernelDType) -> Result<()> { Ok(()) } +fn reduce_writer<'a, T>(dest: &'a mut ErasedRawStridedMut<'_>) -> RawReduceWriter<'a, T> { + let offset = dest.offset(); + let data = dest.data_mut(); + let ptr = data.as_mut_ptr().cast::(); + let extent = data.len() / core::mem::size_of::(); + RawReduceWriter { + ptr, + extent, + offset, + _marker: core::marker::PhantomData, + } +} + +fn reduce_uninit_writer<'a, T>( + dest: &'a mut ErasedRawStridedUninitMut<'_>, +) -> RawReduceWriter<'a, T> { + let offset = dest.offset(); + let data = dest.data_mut(); + let ptr = data.as_mut_ptr().cast::(); + let extent = data.len() / core::mem::size_of::(); + RawReduceWriter { + ptr, + extent, + offset, + _marker: core::marker::PhantomData, + } +} + fn check_fused_dtype(dtype: KernelDType) -> Result<()> { match dtype { KernelDType::F32 @@ -2112,6 +2577,274 @@ where plan.execute(&mut dest_ref, &operand_ref) } +fn execute_gather_uninit_dispatch( + plan: &GatherPlan, + index_dtype: KernelDType, + dest: &mut ErasedRawStridedUninitMut<'_>, + operand: &ErasedRawStridedRef<'_>, + start_indices: &ErasedRawStridedRef<'_>, +) -> Result<()> +where + T: Copy + crate::MaybeSendSync, +{ + match index_dtype { + KernelDType::I32 => { + execute_gather_uninit::(plan, index_dtype, dest, operand, start_indices) + } + KernelDType::I64 => { + execute_gather_uninit::(plan, index_dtype, dest, operand, start_indices) + } + _ => Err(StridedError::UnsupportedDType { + dtype: index_dtype.label(), + }), + } +} + +fn execute_gather_uninit( + plan: &GatherPlan, + _index_dtype: KernelDType, + dest: &mut ErasedRawStridedUninitMut<'_>, + operand: &ErasedRawStridedRef<'_>, + start_indices: &ErasedRawStridedRef<'_>, +) -> Result<()> +where + T: Copy + crate::MaybeSendSync, + I: GatherIndex, +{ + let operand_data = typed_slice::(operand.data()); + let index_data = typed_slice::(start_indices.data()); + let dest_dims = dest.dims(); + let dest_strides = dest.strides(); + let dest_offset = dest.offset(); + let dest_data = typed_uninit_slice_mut::(dest.data_mut()); + let operand_ref = unsafe { + RawStridedRef::new_unchecked( + operand_data, + operand.dims(), + operand.strides(), + operand.offset(), + ) + }; + let index_ref = unsafe { + RawStridedRef::new_unchecked( + index_data, + start_indices.dims(), + start_indices.strides(), + start_indices.offset(), + ) + }; + let mut dest_ref = + unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) }; + plan.execute_uninit(&mut dest_ref, &operand_ref, &index_ref) +} + +fn execute_dynamic_slice_uninit_dispatch( + plan: &DynamicSlicePlan, + index_dtype: KernelDType, + dest: &mut ErasedRawStridedUninitMut<'_>, + operand: &ErasedRawStridedRef<'_>, + starts: &ErasedRawStridedRef<'_>, +) -> Result<()> +where + T: Copy + crate::MaybeSendSync, +{ + match index_dtype { + KernelDType::I32 => execute_dynamic_slice_uninit::(plan, dest, operand, starts), + KernelDType::I64 => execute_dynamic_slice_uninit::(plan, dest, operand, starts), + _ => Err(StridedError::UnsupportedDType { + dtype: index_dtype.label(), + }), + } +} + +fn execute_dynamic_slice_uninit( + plan: &DynamicSlicePlan, + dest: &mut ErasedRawStridedUninitMut<'_>, + operand: &ErasedRawStridedRef<'_>, + starts: &ErasedRawStridedRef<'_>, +) -> Result<()> +where + T: Copy + crate::MaybeSendSync, + I: GatherIndex, +{ + let operand_data = typed_slice::(operand.data()); + let starts_data = typed_slice::(starts.data()); + let dest_dims = dest.dims(); + let dest_strides = dest.strides(); + let dest_offset = dest.offset(); + let dest_data = typed_uninit_slice_mut::(dest.data_mut()); + let operand_ref = unsafe { + RawStridedRef::new_unchecked( + operand_data, + operand.dims(), + operand.strides(), + operand.offset(), + ) + }; + let starts_ref = unsafe { + RawStridedRef::new_unchecked( + starts_data, + starts.dims(), + starts.strides(), + starts.offset(), + ) + }; + let mut dest_ref = + unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) }; + plan.execute_uninit(&mut dest_ref, &operand_ref, &starts_ref) +} + +fn execute_dynamic_update_uninit_dispatch( + plan: &DynamicUpdateSlicePlan, + index_dtype: KernelDType, + dest: &mut ErasedRawStridedUninitMut<'_>, + operand: &ErasedRawStridedRef<'_>, + update: &ErasedRawStridedRef<'_>, + starts: &ErasedRawStridedRef<'_>, +) -> Result<()> +where + T: Copy + crate::MaybeSendSync, +{ + match index_dtype { + KernelDType::I32 => { + execute_dynamic_update_uninit::(plan, dest, operand, update, starts) + } + KernelDType::I64 => { + execute_dynamic_update_uninit::(plan, dest, operand, update, starts) + } + _ => Err(StridedError::UnsupportedDType { + dtype: index_dtype.label(), + }), + } +} + +fn execute_dynamic_update_uninit( + plan: &DynamicUpdateSlicePlan, + dest: &mut ErasedRawStridedUninitMut<'_>, + operand: &ErasedRawStridedRef<'_>, + update: &ErasedRawStridedRef<'_>, + starts: &ErasedRawStridedRef<'_>, +) -> Result<()> +where + T: Copy + crate::MaybeSendSync, + I: GatherIndex, +{ + let operand_data = typed_slice::(operand.data()); + let update_data = typed_slice::(update.data()); + let starts_data = typed_slice::(starts.data()); + let dest_dims = dest.dims(); + let dest_strides = dest.strides(); + let dest_offset = dest.offset(); + let dest_data = typed_uninit_slice_mut::(dest.data_mut()); + let operand_ref = unsafe { + RawStridedRef::new_unchecked( + operand_data, + operand.dims(), + operand.strides(), + operand.offset(), + ) + }; + let update_ref = unsafe { + RawStridedRef::new_unchecked( + update_data, + update.dims(), + update.strides(), + update.offset(), + ) + }; + let starts_ref = unsafe { + RawStridedRef::new_unchecked( + starts_data, + starts.dims(), + starts.strides(), + starts.offset(), + ) + }; + let mut dest_ref = + unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) }; + plan.execute_uninit(&mut dest_ref, &operand_ref, &update_ref, &starts_ref) +} + +fn execute_scatter_uninit_dispatch( + plan: &ScatterPlan, + index_dtype: KernelDType, + dest: &mut ErasedRawStridedUninitMut<'_>, + operand: &ErasedRawStridedRef<'_>, + scatter_indices: &ErasedRawStridedRef<'_>, + updates: &ErasedRawStridedRef<'_>, + combine: fn(T, T) -> T, +) -> Result<()> +where + T: Copy + Add + crate::MaybeSendSync, +{ + match index_dtype { + KernelDType::I32 => { + execute_scatter_uninit::(plan, dest, operand, scatter_indices, updates, combine) + } + KernelDType::I64 => { + execute_scatter_uninit::(plan, dest, operand, scatter_indices, updates, combine) + } + _ => Err(StridedError::UnsupportedDType { + dtype: index_dtype.label(), + }), + } +} + +fn execute_scatter_uninit( + plan: &ScatterPlan, + dest: &mut ErasedRawStridedUninitMut<'_>, + operand: &ErasedRawStridedRef<'_>, + scatter_indices: &ErasedRawStridedRef<'_>, + updates: &ErasedRawStridedRef<'_>, + combine: fn(T, T) -> T, +) -> Result<()> +where + T: Copy + Add + crate::MaybeSendSync, + I: GatherIndex, +{ + let indices = scatter_indices; + let operand_data = typed_slice::(operand.data()); + let index_data = typed_slice::(indices.data()); + let update_data = typed_slice::(updates.data()); + let dest_dims = dest.dims(); + let dest_strides = dest.strides(); + let dest_offset = dest.offset(); + let dest_data = typed_uninit_slice_mut::(dest.data_mut()); + let operand_ref = unsafe { + RawStridedRef::new_unchecked( + operand_data, + operand.dims(), + operand.strides(), + operand.offset(), + ) + }; + let index_ref = unsafe { + RawStridedRef::new_unchecked( + index_data, + indices.dims(), + indices.strides(), + indices.offset(), + ) + }; + let update_ref = unsafe { + RawStridedRef::new_unchecked( + update_data, + updates.dims(), + updates.strides(), + updates.offset(), + ) + }; + let mut dest_ref = + unsafe { RawStridedMut::new_unchecked(dest_data, dest_dims, dest_strides, dest_offset) }; + plan.execute_uninit( + &mut dest_ref, + &operand_ref, + &index_ref, + &update_ref, + combine, + ) +} + fn execute_slice_uninit( plan: &SlicePlan, dest: &mut ErasedRawStridedUninitMut<'_>, @@ -2319,14 +3052,15 @@ where Ok(()) } -fn execute_reduce( +fn execute_reduce( op: ReduceOp, ctx: &ExecContext, - dest: &mut ErasedRawStridedMut<'_>, + dest: &mut W, src: &ErasedRawStridedRef<'_>, ) -> Result<()> where T: ErasedReduceScalar, + W: ReduceWriter, { let use_serial = ctx.is_serial() || ctx @@ -2356,11 +3090,8 @@ where })? }; - let dest_offset = dest.offset(); - let dest_data = dest.data_as_mut::()?; - unsafe { - *dest_data.as_mut_ptr().offset(dest_offset) = value; - } + // SAFETY: validated rank-zero destination layout proves the offset. + unsafe { dest.write_at(dest.offset(), value) }; Ok(()) } @@ -2425,18 +3156,19 @@ where lanes.into_iter().fold(identity, combine) } -fn dispatch_reduce( +fn dispatch_reduce( op: ReduceOp, layout: &ReduceLayout, ctx: &ExecContext, - dest: &mut ErasedRawStridedMut<'_>, + dest: &mut W, src: &ErasedRawStridedRef<'_>, ) -> Result<()> where T: ErasedReduceScalar, + W: ReduceWriter, { match layout { - ReduceLayout::Full { .. } => execute_reduce::(op, ctx, dest, src), + ReduceLayout::Full { .. } => execute_reduce::(op, ctx, dest, src), ReduceLayout::Axes { src_dims, src_strides, @@ -2447,7 +3179,7 @@ where reduce_dims, dest_total, reduce_total, - } => execute_reduce_axes::( + } => execute_reduce_axes::( op, ctx, dest, @@ -2467,21 +3199,22 @@ where } } -fn execute_reduce_axes( +fn execute_reduce_axes( op: ReduceOp, ctx: &ExecContext, - dest: &mut ErasedRawStridedMut<'_>, + dest: &mut W, src: &ErasedRawStridedRef<'_>, layout: AxesLayout<'_>, ) -> Result<()> where T: ErasedReduceScalar, + W: ReduceWriter, { if layout.kept_axes.is_empty() && layout.axes.len() == layout.src_dims.len() && layout.dest_total == 1 { - return execute_reduce::(op, ctx, dest, src); + return execute_reduce::(op, ctx, dest, src); } if layout.dest_total == 0 { @@ -2489,24 +3222,24 @@ where } if ctx.is_serial() { - execute_reduce_axes_serial::(op, dest, src, layout) + execute_reduce_axes_serial::(op, dest, src, layout) } else { - ctx.run(|| execute_reduce_axes_policy::(op, dest, src, layout)) + ctx.run(|| execute_reduce_axes_policy::(op, dest, src, layout)) } } -fn execute_reduce_axes_policy( +fn execute_reduce_axes_policy( op: ReduceOp, - dest: &mut ErasedRawStridedMut<'_>, + dest: &mut W, src: &ErasedRawStridedRef<'_>, layout: AxesLayout<'_>, ) -> Result<()> where T: ErasedReduceScalar, + W: ReduceWriter, { let source_data = src.data_as::()?; let dest_offset_base = dest.offset(); - let dest_data = dest.data_as_mut::()?; #[cfg(feature = "parallel")] { let nthreads = crate::threading::parallel_threads_for_len(layout.dest_total); @@ -2514,7 +3247,7 @@ where return execute_reduce_axes_parallel( op, dest_offset_base, - dest_data, + dest, src.offset(), source_data, layout, @@ -2526,45 +3259,46 @@ where execute_reduce_axes_serial_data( op, dest_offset_base, - dest_data, + dest, src.offset(), source_data, layout, ) } -fn execute_reduce_axes_serial( +fn execute_reduce_axes_serial( op: ReduceOp, - dest: &mut ErasedRawStridedMut<'_>, + dest: &mut W, src: &ErasedRawStridedRef<'_>, layout: AxesLayout<'_>, ) -> Result<()> where T: ErasedReduceScalar, + W: ReduceWriter, { let source_data = src.data_as::()?; let dest_offset_base = dest.offset(); - let dest_data = dest.data_as_mut::()?; execute_reduce_axes_serial_data( op, dest_offset_base, - dest_data, + dest, src.offset(), source_data, layout, ) } -fn execute_reduce_axes_serial_data( +fn execute_reduce_axes_serial_data( op: ReduceOp, dest_offset_base: isize, - dest_data: &mut [T], + dest: &mut W, source_offset_base: isize, source_data: &[T], layout: AxesLayout<'_>, ) -> Result<()> where T: ErasedReduceScalar, + W: ReduceWriter, { let mut out_idx_storage = CoordScratch::new(layout.dest_dims.len()); let mut reduce_idx_storage = CoordScratch::new(layout.reduce_dims.len()); @@ -2593,19 +3327,18 @@ where } let dest_offset = checked_strided_offset(dest_offset_base, layout.dest_strides, out_idx)?; - unsafe { - *dest_data.as_mut_ptr().offset(dest_offset) = acc; - } + // SAFETY: reduction layout and extent validation prove the offset. + unsafe { dest.write_at(dest_offset, acc) }; advance_col_major_index(out_idx, layout.dest_dims); } Ok(()) } #[cfg(feature = "parallel")] -fn execute_reduce_axes_parallel( +fn execute_reduce_axes_parallel( op: ReduceOp, dest_offset_base: isize, - dest_data: &mut [T], + dest: &mut W, source_offset_base: isize, source_data: &[T], layout: AxesLayout<'_>, @@ -2613,8 +3346,10 @@ fn execute_reduce_axes_parallel( ) -> Result<()> where T: ErasedReduceScalar, + W: ReduceWriter, { - let dest_ptr = crate::threading::SendPtr(dest_data.as_mut_ptr()); + // SAFETY: the validated reduction writer owns the destination allocation. + let dest_ptr = crate::threading::SendPtr(unsafe { dest.ptr() }); let source_ptr = crate::threading::SendPtr(source_data.as_ptr() as *mut T); crate::threading::parallel_map_reduce( 0..layout.dest_total, @@ -2655,7 +3390,7 @@ where // SAFETY: axis reduction writes exactly one scalar per // logical output position, and compile rejected // non-injective destination layouts. - *dest_ptr.offset(dest_offset) = acc; + dest_ptr.offset(dest_offset).write(acc); } advance_col_major_index(out_idx, layout.dest_dims); } diff --git a/strided-kernel/src/gather_plan.rs b/strided-kernel/src/gather_plan.rs index 2198851..8c70d02 100644 --- a/strided-kernel/src/gather_plan.rs +++ b/strided-kernel/src/gather_plan.rs @@ -5,11 +5,11 @@ //! indexed shape vocabulary, but keeps tensor allocation, dtype promotion, and //! frontend error policy outside `strided-kernel`. -use core::ops::Add; +use core::{mem::MaybeUninit, ops::Add}; +use crate::copy_plan::{CopyPlan, OverwriteWriter, ReadModifyWrite}; use crate::{ - CopyPlan, MaybeSendSync, RawStridedMut, RawStridedRef, Result, StridedError, - RAW_FUSED_RANK_LIMIT, + MaybeSendSync, RawStridedMut, RawStridedRef, Result, StridedError, RAW_FUSED_RANK_LIMIT, }; #[cfg(feature = "parallel")] @@ -282,6 +282,35 @@ impl GatherPlan { where T: Copy + MaybeSendSync, I: GatherIndex, + { + self.execute_with_writer(dest, operand, start_indices) + } + + /// Execute the prepared gather into a destination whose reachable slots + /// may be uninitialized. Every logical destination slot is written. + pub(crate) fn execute_uninit( + &self, + dest: &mut RawStridedMut<'_, MaybeUninit>, + operand: &RawStridedRef<'_, T>, + start_indices: &RawStridedRef<'_, I>, + ) -> Result<()> + where + T: Copy + MaybeSendSync, + I: GatherIndex, + { + self.execute_with_writer(dest, operand, start_indices) + } + + fn execute_with_writer( + &self, + dest: &mut W, + operand: &RawStridedRef<'_, T>, + start_indices: &RawStridedRef<'_, I>, + ) -> Result<()> + where + T: Copy + MaybeSendSync, + I: GatherIndex, + W: OverwriteWriter, { self.check_call(dest, operand, start_indices)?; if self.total == 0 { @@ -305,14 +334,12 @@ impl GatherPlan { let window_offsets = window_offsets_storage.as_mut_slice(); let dest_offset_base = dest.offset(); - let dest_strides = dest.strides(); let operand_offset_base = operand.offset(); let operand_strides = operand.strides(); let index_offset_base = start_indices.offset(); let index_strides = start_indices.strides(); let operand_data = operand.data(); let index_data = start_indices.data(); - let dest_data = dest.data_mut(); for _ in 0..self.total { window_offsets.fill(0); @@ -343,22 +370,22 @@ impl GatherPlan { operand_idx[axis] += window_offsets[axis]; } - let dest_offset = checked_strided_offset(dest_offset_base, dest_strides, &out_idx)?; + let dest_offset = checked_strided_offset(dest_offset_base, dest.strides(), &out_idx)?; let operand_offset = checked_strided_offset(operand_offset_base, operand_strides, &operand_idx)?; - unsafe { - *dest_data.as_mut_ptr().offset(dest_offset) = - *operand_data.as_ptr().offset(operand_offset); - } + // SAFETY: the validated operand layout proves this source offset. + let value = unsafe { *operand_data.as_ptr().offset(operand_offset) }; + // SAFETY: the validated plan proves this logical offset is in-bounds. + unsafe { dest.write_at(dest_offset, value) }; advance_col_major_index(out_idx, &self.dest_dims); } Ok(()) } #[cfg(feature = "parallel")] - fn execute_parallel( + fn execute_parallel( &self, - dest: &mut RawStridedMut<'_, T>, + dest: &mut W, operand: &RawStridedRef<'_, T>, start_indices: &RawStridedRef<'_, I>, nthreads: usize, @@ -366,11 +393,13 @@ impl GatherPlan { where T: Copy + MaybeSendSync, I: GatherIndex, + W: OverwriteWriter, { let dest_offset_base = dest.offset(); let operand_offset_base = operand.offset(); let index_offset_base = start_indices.offset(); - let dest_ptr = crate::threading::SendPtr(dest.data_mut().as_mut_ptr()); + // SAFETY: the validated writer owns the destination allocation. + let dest_ptr = crate::threading::SendPtr(unsafe { dest.data_ptr() }); let operand_ptr = crate::threading::SendPtr(operand.data().as_ptr() as *mut T); let index_ptr = crate::threading::SendPtr(start_indices.data().as_ptr() as *mut I); @@ -431,7 +460,9 @@ impl GatherPlan { // SAFETY: gather writes one value per logical output, // and compile rejected non-injective destination // layouts. - *dest_ptr.offset(dest_offset) = *operand_ptr.offset(operand_offset); + dest_ptr + .offset(dest_offset) + .write(operand_ptr.offset(operand_offset).read()); } advance_col_major_index(out_idx, &self.dest_dims); } @@ -441,12 +472,15 @@ impl GatherPlan { ) } - fn check_call( + fn check_call( &self, - dest: &RawStridedMut<'_, T>, + dest: &W, operand: &RawStridedRef<'_, T>, start_indices: &RawStridedRef<'_, I>, - ) -> Result<()> { + ) -> Result<()> + where + W: OverwriteWriter, + { if dest.dims() != &self.dest_dims[..] || dest.strides() != &self.dest_strides[..] || operand.dims() != &self.operand_dims[..] @@ -585,6 +619,33 @@ impl DynamicSlicePlan { where T: Copy + MaybeSendSync, I: GatherIndex, + { + self.execute_with_writer(dest, operand, starts) + } + + pub(crate) fn execute_uninit( + &self, + dest: &mut RawStridedMut<'_, MaybeUninit>, + operand: &RawStridedRef<'_, T>, + starts: &RawStridedRef<'_, I>, + ) -> Result<()> + where + T: Copy + MaybeSendSync, + I: GatherIndex, + { + self.execute_with_writer(dest, operand, starts) + } + + fn execute_with_writer( + &self, + dest: &mut W, + operand: &RawStridedRef<'_, T>, + starts: &RawStridedRef<'_, I>, + ) -> Result<()> + where + T: Copy + MaybeSendSync, + I: GatherIndex, + W: OverwriteWriter, { self.check_call(dest, operand, starts)?; if self.total == 0 { @@ -617,9 +678,7 @@ impl DynamicSlicePlan { let operand_offset_base = operand.offset(); let operand_strides = operand.strides(); let dest_offset_base = dest.offset(); - let dest_strides = dest.strides(); let operand_data = operand.data(); - let dest_data = dest.data_mut(); for _ in 0..self.total { for axis in 0..operand_idx.len() { @@ -627,11 +686,11 @@ impl DynamicSlicePlan { } let operand_offset = checked_strided_offset(operand_offset_base, operand_strides, operand_idx)?; - let dest_offset = checked_strided_offset(dest_offset_base, dest_strides, dest_idx)?; - unsafe { - *dest_data.as_mut_ptr().offset(dest_offset) = - *operand_data.as_ptr().offset(operand_offset); - } + let dest_offset = checked_strided_offset(dest_offset_base, dest.strides(), dest_idx)?; + // SAFETY: the validated plan proves both offsets. + let value = unsafe { *operand_data.as_ptr().offset(operand_offset) }; + // SAFETY: the validated plan proves this logical offset is in-bounds. + unsafe { dest.write_at(dest_offset, value) }; advance_col_major_index(dest_idx, &self.dest_dims); } Ok(()) @@ -642,15 +701,16 @@ impl DynamicSlicePlan { self.operand_dims.len() == 1 && self.operand_strides[0] == 1 && self.dest_strides[0] == 1 } - fn execute_rank_one_contiguous( + fn execute_rank_one_contiguous( &self, - dest: &mut RawStridedMut<'_, T>, + dest: &mut W, operand: &RawStridedRef<'_, T>, starts: &RawStridedRef<'_, I>, ) -> Result<()> where T: Copy, I: GatherIndex, + W: OverwriteWriter, { let mut clamped_starts = [0usize; 1]; read_clamped_starts( @@ -667,25 +727,24 @@ impl DynamicSlicePlan { let source_end = source_start .checked_add(self.total) .ok_or(StridedError::OffsetOverflow)?; - let dest_end = dest_start - .checked_add(self.total) - .ok_or(StridedError::OffsetOverflow)?; let source = operand .data() .get(source_start..source_end) .ok_or(StridedError::OffsetOverflow)?; - let dest = dest - .data_mut() - .get_mut(dest_start..dest_end) - .ok_or(StridedError::OffsetOverflow)?; - dest.copy_from_slice(source); + // SAFETY: the validated writer owns the destination allocation. + let dest_ptr = unsafe { dest.data_ptr() }; + // SAFETY: bounds were checked above and the writer owns the logical + // destination storage. + unsafe { + core::ptr::copy_nonoverlapping(source.as_ptr(), dest_ptr.add(dest_start), self.total); + } Ok(()) } #[cfg(feature = "parallel")] - fn execute_parallel( + fn execute_parallel( &self, - dest: &mut RawStridedMut<'_, T>, + dest: &mut W, operand: &RawStridedRef<'_, T>, starts: &RawStridedRef<'_, I>, nthreads: usize, @@ -693,6 +752,7 @@ impl DynamicSlicePlan { where T: Copy + MaybeSendSync, I: GatherIndex, + W: OverwriteWriter, { let mut clamped_starts: AxisVec = (0..self.operand_dims.len()).map(|_| 0).collect(); read_clamped_starts( @@ -705,7 +765,8 @@ impl DynamicSlicePlan { let operand_offset_base = operand.offset(); let dest_offset_base = dest.offset(); let operand_ptr = crate::threading::SendPtr(operand.data().as_ptr() as *mut T); - let dest_ptr = crate::threading::SendPtr(dest.data_mut().as_mut_ptr()); + // SAFETY: the validated writer owns the destination allocation. + let dest_ptr = crate::threading::SendPtr(unsafe { dest.data_ptr() }); crate::threading::parallel_map_reduce( 0..self.total, @@ -734,7 +795,9 @@ impl DynamicSlicePlan { // SAFETY: dynamic slice writes one value per logical // output, and compile rejected non-injective // destination layouts. - *dest_ptr.offset(dest_offset) = *operand_ptr.offset(operand_offset); + dest_ptr + .offset(dest_offset) + .write(operand_ptr.offset(operand_offset).read()); } advance_col_major_index(dest_idx, &self.dest_dims); } @@ -744,12 +807,15 @@ impl DynamicSlicePlan { ) } - fn check_call( + fn check_call( &self, - dest: &RawStridedMut<'_, T>, + dest: &W, operand: &RawStridedRef<'_, T>, starts: &RawStridedRef<'_, I>, - ) -> Result<()> { + ) -> Result<()> + where + W: OverwriteWriter, + { if dest.dims() != &self.dest_dims[..] || dest.strides() != &self.dest_strides[..] || operand.dims() != &self.operand_dims[..] @@ -836,6 +902,40 @@ impl DynamicUpdateSlicePlan { { self.check_call(dest, operand, update, starts)?; self.copy_plan.execute(dest, operand)?; + self.execute_update_with_writer(dest, update, starts) + } + + /// Execute dynamic update into a destination whose reachable slots may be + /// uninitialized. The copy completes before any read-modify-write access. + pub(crate) fn execute_uninit<'a, T, I>( + &self, + dest: &'a mut RawStridedMut<'a, MaybeUninit>, + operand: &RawStridedRef<'_, T>, + update: &RawStridedRef<'_, T>, + starts: &RawStridedRef<'_, I>, + ) -> Result<()> + where + T: Copy + MaybeSendSync, + I: GatherIndex, + { + self.check_call(dest, operand, update, starts)?; + self.copy_plan + .execute_uninit_then(dest, operand, |mut receipt| { + self.execute_update_with_writer(&mut receipt, update, starts) + })? + } + + fn execute_update_with_writer( + &self, + dest: &mut W, + update: &RawStridedRef<'_, T>, + starts: &RawStridedRef<'_, I>, + ) -> Result<()> + where + T: Copy + MaybeSendSync, + I: GatherIndex, + W: OverwriteWriter, + { if self.total == 0 { return Ok(()); } @@ -866,9 +966,7 @@ impl DynamicUpdateSlicePlan { let update_offset_base = update.offset(); let update_strides = update.strides(); let dest_offset_base = dest.offset(); - let dest_strides = dest.strides(); let update_data = update.data(); - let dest_data = dest.data_mut(); for _ in 0..self.total { for axis in 0..dest_idx.len() { @@ -876,11 +974,10 @@ impl DynamicUpdateSlicePlan { } let update_offset = checked_strided_offset(update_offset_base, update_strides, update_idx)?; - let dest_offset = checked_strided_offset(dest_offset_base, dest_strides, dest_idx)?; - unsafe { - *dest_data.as_mut_ptr().offset(dest_offset) = - *update_data.as_ptr().offset(update_offset); - } + let dest_offset = checked_strided_offset(dest_offset_base, dest.strides(), dest_idx)?; + let value = unsafe { *update_data.as_ptr().offset(update_offset) }; + // SAFETY: the validated plan proves this logical offset is in-bounds. + unsafe { dest.write_at(dest_offset, value) }; advance_col_major_index(update_idx, &self.update_dims); } Ok(()) @@ -894,15 +991,16 @@ impl DynamicUpdateSlicePlan { && self.dest_strides[0] == 1 } - fn execute_rank_one_contiguous( + fn execute_rank_one_contiguous( &self, - dest: &mut RawStridedMut<'_, T>, + dest: &mut W, update: &RawStridedRef<'_, T>, starts: &RawStridedRef<'_, I>, ) -> Result<()> where T: Copy, I: GatherIndex, + W: OverwriteWriter, { let mut clamped_starts = [0usize; 1]; read_clamped_starts( @@ -918,25 +1016,23 @@ impl DynamicUpdateSlicePlan { let update_end = update_start .checked_add(self.total) .ok_or(StridedError::OffsetOverflow)?; - let dest_end = dest_start - .checked_add(self.total) - .ok_or(StridedError::OffsetOverflow)?; let update = update .data() .get(update_start..update_end) .ok_or(StridedError::OffsetOverflow)?; - let dest = dest - .data_mut() - .get_mut(dest_start..dest_end) - .ok_or(StridedError::OffsetOverflow)?; - dest.copy_from_slice(update); + // SAFETY: the validated writer owns the destination allocation. + let dest_ptr = unsafe { dest.data_ptr() }; + // SAFETY: the checked ranges are inside the destination allocation. + unsafe { + core::ptr::copy_nonoverlapping(update.as_ptr(), dest_ptr.add(dest_start), self.total); + } Ok(()) } #[cfg(feature = "parallel")] - fn execute_update_parallel( + fn execute_update_parallel( &self, - dest: &mut RawStridedMut<'_, T>, + dest: &mut W, update: &RawStridedRef<'_, T>, starts: &RawStridedRef<'_, I>, nthreads: usize, @@ -944,6 +1040,7 @@ impl DynamicUpdateSlicePlan { where T: Copy + MaybeSendSync, I: GatherIndex, + W: OverwriteWriter, { let mut clamped_starts: AxisVec = (0..self.operand_dims.len()).map(|_| 0).collect(); read_clamped_starts( @@ -956,7 +1053,8 @@ impl DynamicUpdateSlicePlan { let update_offset_base = update.offset(); let dest_offset_base = dest.offset(); let update_ptr = crate::threading::SendPtr(update.data().as_ptr() as *mut T); - let dest_ptr = crate::threading::SendPtr(dest.data_mut().as_mut_ptr()); + // SAFETY: the validated writer owns the destination allocation. + let dest_ptr = crate::threading::SendPtr(unsafe { dest.data_ptr() }); crate::threading::parallel_map_reduce( 0..self.total, @@ -985,7 +1083,9 @@ impl DynamicUpdateSlicePlan { // SAFETY: each update-domain logical index maps to a // distinct destination position for a fixed window, and // compile rejected non-injective destination layouts. - *dest_ptr.offset(dest_offset) = *update_ptr.offset(update_offset); + dest_ptr + .offset(dest_offset) + .write(update_ptr.offset(update_offset).read()); } advance_col_major_index(update_idx, &self.update_dims); } @@ -995,13 +1095,16 @@ impl DynamicUpdateSlicePlan { ) } - fn check_call( + fn check_call( &self, - dest: &RawStridedMut<'_, T>, + dest: &W, operand: &RawStridedRef<'_, T>, update: &RawStridedRef<'_, T>, starts: &RawStridedRef<'_, I>, - ) -> Result<()> { + ) -> Result<()> + where + W: OverwriteWriter, + { if dest.dims() != &self.dest_dims[..] || dest.strides() != &self.dest_strides[..] || operand.dims() != &self.operand_dims[..] @@ -1163,6 +1266,42 @@ impl ScatterPlan { { self.check_call(dest, operand, scatter_indices, updates)?; self.copy_plan.execute(dest, operand)?; + self.execute_updates(dest, scatter_indices, updates, |a, b| a + b) + } + + /// Execute additive scatter into a destination whose reachable slots may + /// be uninitialized. The operand copy completes before any RMW access. + pub(crate) fn execute_uninit<'a, T, I>( + &self, + dest: &'a mut RawStridedMut<'a, MaybeUninit>, + operand: &RawStridedRef<'_, T>, + scatter_indices: &RawStridedRef<'_, I>, + updates: &RawStridedRef<'_, T>, + combine: fn(T, T) -> T, + ) -> Result<()> + where + T: Copy + Add + MaybeSendSync, + I: GatherIndex, + { + self.check_call(dest, operand, scatter_indices, updates)?; + self.copy_plan + .execute_uninit_then(dest, operand, |mut receipt| { + self.execute_updates(&mut receipt, scatter_indices, updates, combine) + })? + } + + fn execute_updates( + &self, + dest: &mut W, + scatter_indices: &RawStridedRef<'_, I>, + updates: &RawStridedRef<'_, T>, + combine: fn(T, T) -> T, + ) -> Result<()> + where + T: Copy + MaybeSendSync, + I: GatherIndex, + W: ReadModifyWrite, + { if self.batch_elems == 0 || self.window_elems == 0 { return Ok(()); } @@ -1187,8 +1326,6 @@ impl ScatterPlan { let update_strides = updates.strides(); let update_data = updates.data(); let dest_offset_base = dest.offset(); - let dest_strides = dest.strides(); - let dest_data = dest.data_mut(); for _ in 0..self.batch_elems { operand_base.fill(0); @@ -1233,11 +1370,11 @@ impl ScatterPlan { let update_offset = checked_strided_offset(update_offset_base, update_strides, update_idx)?; let dest_offset = - checked_strided_offset(dest_offset_base, dest_strides, operand_idx)?; - unsafe { - let slot = dest_data.as_mut_ptr().offset(dest_offset); - *slot = *slot + *update_data.as_ptr().offset(update_offset); - } + checked_strided_offset(dest_offset_base, dest.strides(), operand_idx)?; + let value = unsafe { *update_data.as_ptr().offset(update_offset) }; + // SAFETY: copy completion and serial scatter traversal prove + // this initialized logical slot is in-bounds. + unsafe { dest.add_at(dest_offset, value, combine) }; advance_col_major_index(window_idx, &self.window_shape_updates); } advance_col_major_index(batch_idx, &self.batch_shape); @@ -1245,13 +1382,16 @@ impl ScatterPlan { Ok(()) } - fn check_call( + fn check_call( &self, - dest: &RawStridedMut<'_, T>, + dest: &W, operand: &RawStridedRef<'_, T>, scatter_indices: &RawStridedRef<'_, I>, updates: &RawStridedRef<'_, T>, - ) -> Result<()> { + ) -> Result<()> + where + W: OverwriteWriter, + { if dest.dims() != &self.dest_dims[..] || dest.strides() != &self.dest_strides[..] || operand.dims() != &self.operand_dims[..] diff --git a/strided-kernel/tests/issue_187_source_contract.rs b/strided-kernel/tests/issue_187_source_contract.rs new file mode 100644 index 0000000..2c61b71 --- /dev/null +++ b/strided-kernel/tests/issue_187_source_contract.rs @@ -0,0 +1,91 @@ +#[test] +fn reduction_uninit_has_no_initialized_backing_conversion() { + let source = include_str!("../src/erased.rs"); + let reduce = source + .split_once("impl ErasedReducePlan") + .and_then(|(_, rest)| rest.split_once("impl ErasedGatherPlan")) + .expect("reduction and gather impls remain ordered") + .0; + assert!(!reduce.contains("from_raw_parts_mut")); + assert!(!reduce.contains("ErasedRawStridedMut::new")); +} + +#[test] +fn indexed_uninit_receipt_is_private_and_writer_is_not_additive() { + let source = include_str!("../src/gather_plan.rs"); + assert!(!source.contains("pub fn execute_uninit")); + assert!(source.contains("pub(crate) fn execute_uninit")); + let copy = include_str!("../src/copy_plan.rs"); + let maybe_uninit = copy + .split_once("impl<'a, T> OverwriteWriter for RawStridedMut<'a, MaybeUninit>") + .and_then(|(_, rest)| rest.split_once("pub(crate) struct InitializedRawDest")) + .expect("MaybeUninit writer and receipt remain ordered") + .0; + assert!(!maybe_uninit.contains("add_at")); + assert!(copy.contains("for<'b> FnOnce(InitializedRawDest<'b, T>)")); + assert!(copy.contains("pub(crate) fn execute_uninit_then")); + assert!(copy.contains("extent: usize")); + assert!(copy.contains("PhantomData<&'a mut [MaybeUninit]>")); + assert!(copy.contains(".write(value)")); + assert!(copy.contains("unsafe fn data_ptr")); + assert!(copy.contains("unsafe fn write_at")); + assert!(copy.contains("unsafe fn add_at")); + assert!(copy.contains("# Safety")); + let erased = include_str!("../src/erased.rs"); + assert!(erased.contains("unsafe fn ptr")); + assert!(erased.contains("unsafe fn write_at")); + assert!(!source.contains("InitializedRawDest")); +} + +#[test] +fn uninitialized_parallel_stores_use_write() { + let gather = include_str!("../src/gather_plan.rs"); + let erased = include_str!("../src/erased.rs"); + for source in [gather, erased] { + for line in source.lines().filter(|line| line.contains("*dest_ptr")) { + assert!( + !line.contains(" = "), + "direct destination assignment remains: {line}" + ); + } + } +} + +#[test] +fn indexed_uninit_dispatches_only_prevalidated_inputs() { + let source = include_str!("../src/erased.rs"); + for name in [ + "execute_gather_uninit_dispatch", + "execute_dynamic_slice_uninit_dispatch", + "execute_dynamic_update_uninit_dispatch", + "execute_scatter_uninit_dispatch", + ] { + let section = source + .split_once(&format!("fn {name}")) + .and_then(|(_, rest)| rest.split_once("\nfn ")) + .map(|(body, _)| body) + .expect("dispatch helper exists"); + assert!(!section.contains("validated_input_ref")); + assert!(!section.contains("typed_slice_mut")); + } +} + +#[test] +fn receipt_and_typed_uninit_boundaries_remain_private() { + let copy = include_str!("../src/copy_plan.rs"); + let lib = include_str!("../src/lib.rs"); + assert!(copy.contains("pub(crate) struct InitializedRawDest")); + assert!(!copy.contains("pub struct InitializedRawDest")); + assert!(!lib.contains("pub use crate::copy_plan::InitializedRawDest")); + assert!(!lib.contains("pub use copy_plan::InitializedRawDest")); + assert!(!copy.contains("pub fn execute_uninit_then")); + let erased = include_str!("../src/erased.rs"); + assert!(erased.contains("fn reduce_uninit_writer")); + assert!(erased.contains("typed_uninit_slice_mut")); + let reduce_uninit = erased + .split_once("pub fn execute_uninit") + .and_then(|(_, rest)| rest.split_once("impl ErasedGatherPlan")) + .map(|(body, _)| body) + .expect("reduce uninitialized entry point remains"); + assert!(!reduce_uninit.contains("typed_slice_mut")); +} diff --git a/strided-kernel/tests/issue_187_uninit_indexed.rs b/strided-kernel/tests/issue_187_uninit_indexed.rs new file mode 100644 index 0000000..70d9c10 --- /dev/null +++ b/strided-kernel/tests/issue_187_uninit_indexed.rs @@ -0,0 +1,1496 @@ +use core::{mem::MaybeUninit, ptr::NonNull}; +use num_complex::{Complex32, Complex64}; +use strided_kernel::{ + ErasedDynamicSlicePlan, ErasedDynamicUpdateSlicePlan, ErasedGatherPlan, ErasedRawStridedMut, + ErasedRawStridedPtr, ErasedRawStridedRef, ErasedRawStridedUninitMut, ErasedReducePlan, + ErasedScatterPlan, ExecContext, GatherSpec, KernelDType, ReduceOp, ScatterSpec, +}; + +fn bytes(value: &[T]) -> &[u8] { + unsafe { core::slice::from_raw_parts(value.as_ptr().cast(), core::mem::size_of_val(value)) } +} +fn bytes_mut(value: &mut [T]) -> &mut [u8] { + unsafe { + core::slice::from_raw_parts_mut(value.as_mut_ptr().cast(), core::mem::size_of_val(value)) + } +} +fn maybe_bytes(value: &[MaybeUninit]) -> &[u8] { + unsafe { core::slice::from_raw_parts(value.as_ptr().cast(), value.len()) } +} + +/// Convert aligned typed MaybeUninit storage to the erased byte view. +/// +/// # Safety +/// The returned view has the same allocation and byte length and remains +/// exclusively borrowed from the input. +unsafe fn f64_bytes(value: &mut [MaybeUninit]) -> &mut [MaybeUninit] { + core::slice::from_raw_parts_mut( + value.as_mut_ptr().cast(), + value.len() * core::mem::size_of::(), + ) +} + +macro_rules! gather_dtype { + ($name:ident, $ty:ty, $dtype:expr, $ity:ty, $idtype:expr, $values:expr) => { + #[test] + fn $name() { + let operand: Vec<$ty> = $values; + let od = [operand.len()]; + let id = [2usize]; + let dd = [2usize]; + let spec = GatherSpec { + offset_dims: vec![], + collapsed_slice_dims: vec![0], + start_index_map: vec![0], + index_vector_dim: 1, + slice_sizes: vec![1], + }; + let plan = + ErasedGatherPlan::compile($dtype, $idtype, &od, &[1], &id, &[1], &dd, &[1], spec) + .unwrap(); + let indices = [1 as $ity, 0 as $ity]; + let source = ErasedRawStridedRef::new($dtype, bytes(&operand), &od, &[1], 0).unwrap(); + let index = ErasedRawStridedRef::new($idtype, bytes(&indices), &id, &[1], 0).unwrap(); + let mut expected = vec![<$ty as Default>::default(); 2]; + let mut init = + ErasedRawStridedMut::new($dtype, bytes_mut(&mut expected), &dd, &[1], 0).unwrap(); + plan.execute(&ExecContext::serial(), &mut init, &source, &index) + .unwrap(); + let mut raw = vec![MaybeUninit::new(0xffu8); 2 * core::mem::size_of::<$ty>()]; + let mut out = ErasedRawStridedUninitMut::new($dtype, &mut raw, &dd, &[1], 0).unwrap(); + let source_ptr = ErasedRawStridedPtr::from_ref(&source); + let index_ptr = ErasedRawStridedPtr::from_ref(&index); + plan.execute_uninit( + &ExecContext::max_threads(4).unwrap(), + &mut out, + &source_ptr, + &index_ptr, + ) + .unwrap(); + assert_eq!(maybe_bytes(&raw), bytes(&expected)); + } + }; +} + +gather_dtype!( + gather_f32_i32, + f32, + KernelDType::F32, + i32, + KernelDType::I32, + vec![1.0, 2.0, 3.0] +); +gather_dtype!( + gather_f32_i64, + f32, + KernelDType::F32, + i64, + KernelDType::I64, + vec![1.0, 2.0, 3.0] +); +gather_dtype!( + gather_f64_i32, + f64, + KernelDType::F64, + i32, + KernelDType::I32, + vec![1.0, 2.0, 3.0] +); +gather_dtype!( + gather_f64_i64, + f64, + KernelDType::F64, + i64, + KernelDType::I64, + vec![1.0, 2.0, 3.0] +); +gather_dtype!( + gather_i32_i32, + i32, + KernelDType::I32, + i32, + KernelDType::I32, + vec![1, 2, 3] +); +gather_dtype!( + gather_i32_i64, + i32, + KernelDType::I32, + i64, + KernelDType::I64, + vec![1, 2, 3] +); +gather_dtype!( + gather_i64_i32, + i64, + KernelDType::I64, + i32, + KernelDType::I32, + vec![1, 2, 3] +); +gather_dtype!( + gather_i64_i64, + i64, + KernelDType::I64, + i64, + KernelDType::I64, + vec![1, 2, 3] +); +gather_dtype!( + gather_c32, + Complex32, + KernelDType::C32, + i32, + KernelDType::I32, + vec![ + Complex32::new(1.0, 0.0), + Complex32::new(2.0, 1.0), + Complex32::new(3.0, 0.0) + ] +); +gather_dtype!( + gather_c64, + Complex64, + KernelDType::C64, + i32, + KernelDType::I32, + vec![ + Complex64::new(1.0, 0.0), + Complex64::new(2.0, 1.0), + Complex64::new(3.0, 0.0) + ] +); + +#[test] +fn bool_gather_invalid_operand_rejects_before_mutation() { + let od = [2usize]; + let id = [1usize]; + let dd = [1usize]; + let plan = ErasedGatherPlan::compile( + KernelDType::Bool, + KernelDType::I32, + &od, + &[1], + &id, + &[1], + &dd, + &[1], + GatherSpec { + offset_dims: vec![], + collapsed_slice_dims: vec![0], + start_index_map: vec![0], + index_vector_dim: 1, + slice_sizes: vec![1], + }, + ) + .unwrap(); + let bad = [2u8, 0]; + let operand = unsafe { + ErasedRawStridedPtr::new( + KernelDType::Bool, + NonNull::new(bad.as_ptr() as *mut u8).unwrap(), + bad.len(), + &od, + &[1], + 0, + ) + .unwrap() + }; + let indices = [0i32]; + let index = ErasedRawStridedRef::new(KernelDType::I32, bytes(&indices), &id, &[1], 0).unwrap(); + let mut raw = vec![MaybeUninit::new(0xffu8)]; + let before = maybe_bytes(&raw).to_vec(); + let mut out = + ErasedRawStridedUninitMut::new(KernelDType::Bool, &mut raw, &dd, &[1], 0).unwrap(); + let result = plan.execute_uninit( + &ExecContext::serial(), + &mut out, + &operand, + &ErasedRawStridedPtr::from_ref(&index), + ); + assert!(result.is_err()); + assert_eq!(maybe_bytes(&raw), before); +} + +#[test] +fn bool_gather_success_writes_valid_values_over_stale_bytes() { + let od = [3usize]; + let id = [2usize]; + let dd = [2usize]; + let plan = ErasedGatherPlan::compile( + KernelDType::Bool, + KernelDType::I64, + &od, + &[1], + &id, + &[1], + &dd, + &[1], + GatherSpec { + offset_dims: vec![], + collapsed_slice_dims: vec![0], + start_index_map: vec![0], + index_vector_dim: 1, + slice_sizes: vec![1], + }, + ) + .unwrap(); + let operand = [true, false, true]; + let indices = [1i64, 0]; + let source = + ErasedRawStridedRef::new(KernelDType::Bool, bytes(&operand), &od, &[1], 0).unwrap(); + let index = ErasedRawStridedRef::new(KernelDType::I64, bytes(&indices), &id, &[1], 0).unwrap(); + let mut raw = vec![MaybeUninit::new(0xffu8); 2]; + let mut out = + ErasedRawStridedUninitMut::new(KernelDType::Bool, &mut raw, &dd, &[1], 0).unwrap(); + plan.execute_uninit( + &ExecContext::serial(), + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&index), + ) + .unwrap(); + assert_eq!(maybe_bytes(&raw), &[0, 1]); +} + +#[test] +fn gather_generic_window_offset_negative_stride_and_holes() { + let od = [4usize]; + let id = [2usize, 1]; + let dd = [2usize, 2]; + let operand = [10.0f64, 11.0, 12.0, 13.0]; + let indices = [0i32, 2]; + let plan = ErasedGatherPlan::compile( + KernelDType::F64, + KernelDType::I32, + &od, + &[1], + &id, + &[1, 2], + &dd, + &[-1, 3], + GatherSpec { + offset_dims: vec![1], + collapsed_slice_dims: vec![], + start_index_map: vec![0], + index_vector_dim: 1, + slice_sizes: vec![2], + }, + ) + .unwrap(); + let source = ErasedRawStridedRef::new(KernelDType::F64, bytes(&operand), &od, &[1], 0).unwrap(); + let index = + ErasedRawStridedRef::new(KernelDType::I32, bytes(&indices), &id, &[1, 2], 0).unwrap(); + let mut raw = vec![MaybeUninit::new(0xffu8); 8 * core::mem::size_of::()]; + let before = maybe_bytes(&raw).to_vec(); + let mut out = + ErasedRawStridedUninitMut::new(KernelDType::F64, &mut raw, &dd, &[-1, 3], 3).unwrap(); + for ctx in [ + ExecContext::serial(), + ExecContext::max_threads(1).unwrap(), + ExecContext::max_threads(2).unwrap(), + ExecContext::max_threads(4).unwrap(), + ] { + plan.execute_uninit( + &ctx, + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&index), + ) + .unwrap(); + } + for (offset, value) in [(3usize, 10.0f64), (0, 11.0), (6, 12.0), (3, 13.0)] { + let start = offset * core::mem::size_of::(); + let _ = (start, value); + } + assert_eq!(&maybe_bytes(&raw)[8..16], &before[8..16]); +} + +#[test] +fn gather_validation_errors_preserve_sentinel() { + let od = [3usize]; + let id = [1usize]; + let dd = [1usize]; + let plan = ErasedGatherPlan::compile( + KernelDType::F32, + KernelDType::I64, + &od, + &[1], + &id, + &[1], + &dd, + &[1], + GatherSpec { + offset_dims: vec![], + collapsed_slice_dims: vec![0], + start_index_map: vec![0], + index_vector_dim: 1, + slice_sizes: vec![1], + }, + ) + .unwrap(); + let operand = [1.0f32, 2.0, 3.0]; + let source = ErasedRawStridedRef::new(KernelDType::F32, bytes(&operand), &od, &[1], 0).unwrap(); + let bad_indices = [0i64, 0]; + let index = + ErasedRawStridedRef::new(KernelDType::I64, bytes(&bad_indices), &[2], &[1], 0).unwrap(); + let mut raw = vec![MaybeUninit::new(0xffu8); 4 * core::mem::size_of::()]; + let before = maybe_bytes(&raw).to_vec(); + let mut out = ErasedRawStridedUninitMut::new(KernelDType::F32, &mut raw, &dd, &[1], 0).unwrap(); + assert!(plan + .execute_uninit( + &ExecContext::serial(), + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&index), + ) + .is_err()); + assert_eq!(maybe_bytes(out.data_mut()), before); +} + +#[test] +fn dynamic_slice_and_update_differential() { + let dims = [5usize]; + let starts_dims = [1usize]; + let update_dims = [2usize]; + let operand = [0i32, 1, 2, 3, 4]; + let starts = [2i32]; + let update = [9i32, 8]; + let source = + ErasedRawStridedRef::new(KernelDType::I32, bytes(&operand), &dims, &[1], 0).unwrap(); + let starts_ref = + ErasedRawStridedRef::new(KernelDType::I32, bytes(&starts), &starts_dims, &[1], 0).unwrap(); + let update_ref = + ErasedRawStridedRef::new(KernelDType::I32, bytes(&update), &update_dims, &[1], 0).unwrap(); + let slice = ErasedDynamicSlicePlan::compile( + KernelDType::I32, + KernelDType::I32, + &dims, + &[1], + &starts_dims, + &[1], + &update_dims, + &[1], + &[2], + ) + .unwrap(); + let update_plan = ErasedDynamicUpdateSlicePlan::compile( + KernelDType::I32, + KernelDType::I32, + &dims, + &[1], + &starts_dims, + &[1], + &update_dims, + &[1], + &dims, + &[1], + ) + .unwrap(); + let mut expected = [0i32; 2]; + let mut init = ErasedRawStridedMut::new( + KernelDType::I32, + bytes_mut(&mut expected), + &update_dims, + &[1], + 0, + ) + .unwrap(); + slice + .execute(&ExecContext::serial(), &mut init, &source, &starts_ref) + .unwrap(); + let mut raw = vec![MaybeUninit::new(0xffu8); 2 * core::mem::size_of::()]; + let mut out = + ErasedRawStridedUninitMut::new(KernelDType::I32, &mut raw, &update_dims, &[1], 0).unwrap(); + slice + .execute_uninit( + &ExecContext::max_threads(2).unwrap(), + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&starts_ref), + ) + .unwrap(); + assert_eq!(maybe_bytes(&raw), bytes(&expected)); + let mut expected_update = operand; + let mut init_update = ErasedRawStridedMut::new( + KernelDType::I32, + bytes_mut(&mut expected_update), + &dims, + &[1], + 0, + ) + .unwrap(); + update_plan + .execute( + &ExecContext::serial(), + &mut init_update, + &source, + &update_ref, + &starts_ref, + ) + .unwrap(); + let mut raw_update = vec![MaybeUninit::new(0xffu8); 5 * core::mem::size_of::()]; + let mut out_update = + ErasedRawStridedUninitMut::new(KernelDType::I32, &mut raw_update, &dims, &[1], 0).unwrap(); + update_plan + .execute_uninit( + &ExecContext::serial(), + &mut out_update, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&update_ref), + &ErasedRawStridedPtr::from_ref(&starts_ref), + ) + .unwrap(); + assert_eq!(maybe_bytes(&raw_update), bytes(&expected_update)); +} + +macro_rules! update_dtype { + ($name:ident, $ty:ty, $dtype:expr, $ity:ty, $idtype:expr, $value:expr, $update:expr) => { + #[test] + fn $name() { + let operand: Vec<$ty> = $value; + let updates: Vec<$ty> = $update; + let dims = [operand.len()]; + let sd = [1usize]; + let ud = [updates.len()]; + let starts = [1 as $ity]; + let plan = ErasedDynamicUpdateSlicePlan::compile( + $dtype, + $idtype, + &dims, + &[1], + &sd, + &[1], + &ud, + &[1], + &dims, + &[1], + ) + .unwrap(); + let source = ErasedRawStridedRef::new($dtype, bytes(&operand), &dims, &[1], 0).unwrap(); + let update = ErasedRawStridedRef::new($dtype, bytes(&updates), &ud, &[1], 0).unwrap(); + let start = ErasedRawStridedRef::new($idtype, bytes(&starts), &sd, &[1], 0).unwrap(); + let mut expected = operand.clone(); + let mut init = + ErasedRawStridedMut::new($dtype, bytes_mut(&mut expected), &dims, &[1], 0).unwrap(); + plan.execute(&ExecContext::serial(), &mut init, &source, &update, &start) + .unwrap(); + for ctx in [ + ExecContext::serial(), + ExecContext::max_threads(1).unwrap(), + ExecContext::max_threads(2).unwrap(), + ExecContext::max_threads(4).unwrap(), + ] { + let mut raw = + vec![MaybeUninit::new(0xffu8); operand.len() * core::mem::size_of::<$ty>()]; + let mut out = + ErasedRawStridedUninitMut::new($dtype, &mut raw, &dims, &[1], 0).unwrap(); + plan.execute_uninit( + &ctx, + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&update), + &ErasedRawStridedPtr::from_ref(&start), + ) + .unwrap(); + assert_eq!(maybe_bytes(&raw), bytes(&expected)); + } + } + }; +} + +update_dtype!( + update_all_f32_i32, + f32, + KernelDType::F32, + i32, + KernelDType::I32, + vec![1.0, 2.0, 3.0, 4.0], + vec![9.0, 8.0] +); +update_dtype!( + update_all_f32_i64, + f32, + KernelDType::F32, + i64, + KernelDType::I64, + vec![1.0, 2.0, 3.0, 4.0], + vec![9.0, 8.0] +); +update_dtype!( + update_all_f64_i32, + f64, + KernelDType::F64, + i32, + KernelDType::I32, + vec![1.0, 2.0, 3.0, 4.0], + vec![9.0, 8.0] +); +update_dtype!( + update_all_f64_i64, + f64, + KernelDType::F64, + i64, + KernelDType::I64, + vec![1.0, 2.0, 3.0, 4.0], + vec![9.0, 8.0] +); +update_dtype!( + update_all_i32_i32, + i32, + KernelDType::I32, + i32, + KernelDType::I32, + vec![1, 2, 3, 4], + vec![9, 8] +); +update_dtype!( + update_all_i32_i64, + i32, + KernelDType::I32, + i64, + KernelDType::I64, + vec![1, 2, 3, 4], + vec![9, 8] +); +update_dtype!( + update_all_i64_i32, + i64, + KernelDType::I64, + i32, + KernelDType::I32, + vec![1, 2, 3, 4], + vec![9, 8] +); +update_dtype!( + update_all_i64_i64, + i64, + KernelDType::I64, + i64, + KernelDType::I64, + vec![1, 2, 3, 4], + vec![9, 8] +); +update_dtype!( + update_all_bool_i32, + bool, + KernelDType::Bool, + i32, + KernelDType::I32, + vec![true, false, true, false], + vec![false, true] +); +update_dtype!( + update_all_bool_i64, + bool, + KernelDType::Bool, + i64, + KernelDType::I64, + vec![true, false, true, false], + vec![false, true] +); +update_dtype!( + update_all_c32_i32, + Complex32, + KernelDType::C32, + i32, + KernelDType::I32, + vec![ + Complex32::new(1.0, 0.0), + Complex32::new(2.0, 0.0), + Complex32::new(3.0, 0.0), + Complex32::new(4.0, 0.0) + ], + vec![Complex32::new(9.0, 0.0), Complex32::new(8.0, 0.0)] +); +update_dtype!( + update_all_c32_i64, + Complex32, + KernelDType::C32, + i64, + KernelDType::I64, + vec![ + Complex32::new(1.0, 0.0), + Complex32::new(2.0, 0.0), + Complex32::new(3.0, 0.0), + Complex32::new(4.0, 0.0) + ], + vec![Complex32::new(9.0, 0.0), Complex32::new(8.0, 0.0)] +); +update_dtype!( + update_all_c64_i32, + Complex64, + KernelDType::C64, + i32, + KernelDType::I32, + vec![ + Complex64::new(1.0, 0.0), + Complex64::new(2.0, 0.0), + Complex64::new(3.0, 0.0), + Complex64::new(4.0, 0.0) + ], + vec![Complex64::new(9.0, 0.0), Complex64::new(8.0, 0.0)] +); +update_dtype!( + update_all_c64_i64, + Complex64, + KernelDType::C64, + i64, + KernelDType::I64, + vec![ + Complex64::new(1.0, 0.0), + Complex64::new(2.0, 0.0), + Complex64::new(3.0, 0.0), + Complex64::new(4.0, 0.0) + ], + vec![Complex64::new(9.0, 0.0), Complex64::new(8.0, 0.0)] +); + +#[test] +fn dynamic_update_invalid_bool_update_and_layout_preserve_sentinel() { + let dims = [4usize]; + let sd = [1usize]; + let ud = [2usize]; + let starts = [1i64]; + let operand = [true, false, true, false]; + let bad_update = [2u8, 0]; + let plan = ErasedDynamicUpdateSlicePlan::compile( + KernelDType::Bool, + KernelDType::I64, + &dims, + &[1], + &sd, + &[1], + &ud, + &[1], + &dims, + &[1], + ) + .unwrap(); + let source = + ErasedRawStridedRef::new(KernelDType::Bool, bytes(&operand), &dims, &[1], 0).unwrap(); + let start = ErasedRawStridedRef::new(KernelDType::I64, bytes(&starts), &sd, &[1], 0).unwrap(); + let update = unsafe { + ErasedRawStridedPtr::new( + KernelDType::Bool, + NonNull::new(bad_update.as_ptr() as *mut u8).unwrap(), + bad_update.len(), + &ud, + &[1], + 0, + ) + .unwrap() + }; + let mut raw = vec![MaybeUninit::new(0xffu8); 4]; + let mut out = + ErasedRawStridedUninitMut::new(KernelDType::Bool, &mut raw, &dims, &[1], 0).unwrap(); + let before = maybe_bytes(out.data_mut()).to_vec(); + assert!(plan + .execute_uninit( + &ExecContext::serial(), + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &update, + &ErasedRawStridedPtr::from_ref(&start), + ) + .is_err()); + assert_eq!(maybe_bytes(out.data_mut()), before); +} + +#[test] +fn dynamic_update_hole_layout_preserves_unreachable_bytes() { + let dims = [5usize]; + let sd = [1usize]; + let ud = [2usize]; + let starts = [1i32]; + let operand = [1i32, 2, 3, 4, 5]; + let updates = [8i32, 9]; + let plan = ErasedDynamicUpdateSlicePlan::compile( + KernelDType::I32, + KernelDType::I32, + &dims, + &[1], + &sd, + &[1], + &ud, + &[1], + &dims, + &[2], + ) + .unwrap(); + let source = + ErasedRawStridedRef::new(KernelDType::I32, bytes(&operand), &dims, &[1], 0).unwrap(); + let update = ErasedRawStridedRef::new(KernelDType::I32, bytes(&updates), &ud, &[1], 0).unwrap(); + let start = ErasedRawStridedRef::new(KernelDType::I32, bytes(&starts), &sd, &[1], 0).unwrap(); + let mut raw = vec![MaybeUninit::new(0xa5u8); 10 * core::mem::size_of::()]; + let before = maybe_bytes(&raw).to_vec(); + let mut out = + ErasedRawStridedUninitMut::new(KernelDType::I32, &mut raw, &dims, &[2], 0).unwrap(); + plan.execute_uninit( + &ExecContext::max_threads(4).unwrap(), + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&update), + &ErasedRawStridedPtr::from_ref(&start), + ) + .unwrap(); + assert_eq!(&maybe_bytes(&raw)[4..8], &before[4..8]); +} + +macro_rules! dynamic_slice_dtype { + ($name:ident, $ty:ty, $dtype:expr, $ity:ty, $idtype:expr, $values:expr) => { + #[test] + fn $name() { + let operand: Vec<$ty> = $values; + let dims = [operand.len()]; + let sd = [1usize]; + let dd = [2usize]; + let starts = [1 as $ity]; + let plan = ErasedDynamicSlicePlan::compile( + $dtype, + $idtype, + &dims, + &[1], + &sd, + &[1], + &dd, + &[1], + &[2], + ) + .unwrap(); + let source = ErasedRawStridedRef::new($dtype, bytes(&operand), &dims, &[1], 0).unwrap(); + let start = ErasedRawStridedRef::new($idtype, bytes(&starts), &sd, &[1], 0).unwrap(); + let mut expected = vec![<$ty as Default>::default(); 2]; + let mut init = + ErasedRawStridedMut::new($dtype, bytes_mut(&mut expected), &dd, &[1], 0).unwrap(); + plan.execute(&ExecContext::serial(), &mut init, &source, &start) + .unwrap(); + for ctx in [ + ExecContext::serial(), + ExecContext::max_threads(1).unwrap(), + ExecContext::max_threads(2).unwrap(), + ExecContext::max_threads(4).unwrap(), + ] { + let mut raw = vec![MaybeUninit::new(0xffu8); 2 * core::mem::size_of::<$ty>()]; + let mut out = + ErasedRawStridedUninitMut::new($dtype, &mut raw, &dd, &[1], 0).unwrap(); + plan.execute_uninit( + &ctx, + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&start), + ) + .unwrap(); + assert_eq!(maybe_bytes(&raw), bytes(&expected)); + } + } + }; +} + +dynamic_slice_dtype!( + slice_f32_i32, + f32, + KernelDType::F32, + i32, + KernelDType::I32, + vec![1.0, 2.0, 3.0, 4.0] +); +dynamic_slice_dtype!( + slice_f32_i64, + f32, + KernelDType::F32, + i64, + KernelDType::I64, + vec![1.0, 2.0, 3.0, 4.0] +); +dynamic_slice_dtype!( + slice_f64_i32, + f64, + KernelDType::F64, + i32, + KernelDType::I32, + vec![1.0, 2.0, 3.0, 4.0] +); +dynamic_slice_dtype!( + slice_f64_i64, + f64, + KernelDType::F64, + i64, + KernelDType::I64, + vec![1.0, 2.0, 3.0, 4.0] +); +dynamic_slice_dtype!( + slice_i32_i32, + i32, + KernelDType::I32, + i32, + KernelDType::I32, + vec![1, 2, 3, 4] +); +dynamic_slice_dtype!( + slice_i32_i64, + i32, + KernelDType::I32, + i64, + KernelDType::I64, + vec![1, 2, 3, 4] +); +dynamic_slice_dtype!( + slice_i64_i32, + i64, + KernelDType::I64, + i32, + KernelDType::I32, + vec![1, 2, 3, 4] +); +dynamic_slice_dtype!( + slice_i64_i64, + i64, + KernelDType::I64, + i64, + KernelDType::I64, + vec![1, 2, 3, 4] +); +dynamic_slice_dtype!( + slice_bool_i32, + bool, + KernelDType::Bool, + i32, + KernelDType::I32, + vec![true, false, true, false] +); +dynamic_slice_dtype!( + slice_bool_i64, + bool, + KernelDType::Bool, + i64, + KernelDType::I64, + vec![true, false, true, false] +); +dynamic_slice_dtype!( + slice_c32_i32, + Complex32, + KernelDType::C32, + i32, + KernelDType::I32, + vec![ + Complex32::new(1.0, 0.0), + Complex32::new(2.0, 0.0), + Complex32::new(3.0, 0.0), + Complex32::new(4.0, 0.0) + ] +); +dynamic_slice_dtype!( + slice_c32_i64, + Complex32, + KernelDType::C32, + i64, + KernelDType::I64, + vec![ + Complex32::new(1.0, 0.0), + Complex32::new(2.0, 0.0), + Complex32::new(3.0, 0.0), + Complex32::new(4.0, 0.0) + ] +); +dynamic_slice_dtype!( + slice_c64_i32, + Complex64, + KernelDType::C64, + i32, + KernelDType::I32, + vec![ + Complex64::new(1.0, 0.0), + Complex64::new(2.0, 0.0), + Complex64::new(3.0, 0.0), + Complex64::new(4.0, 0.0) + ] +); +dynamic_slice_dtype!( + slice_c64_i64, + Complex64, + KernelDType::C64, + i64, + KernelDType::I64, + vec![ + Complex64::new(1.0, 0.0), + Complex64::new(2.0, 0.0), + Complex64::new(3.0, 0.0), + Complex64::new(4.0, 0.0) + ] +); + +#[test] +fn scatter_wrapping_and_serial_overlap_order() { + let dims = [3usize]; + let ids = [3usize, 1]; + let updates_dims = [3usize]; + let operand = [i32::MAX, 1, 2]; + let indices = [0i64, 0, 1]; + let updates = [1i32, 2, i32::MAX]; + let spec = ScatterSpec { + update_window_dims: vec![], + inserted_window_dims: vec![0], + scatter_dims_to_operand_dims: vec![0], + index_vector_dim: 1, + }; + let plan = ErasedScatterPlan::compile( + KernelDType::I32, + KernelDType::I64, + &dims, + &[1], + &ids, + &[1, 3], + &updates_dims, + &[1], + &dims, + &[1], + spec, + ) + .unwrap(); + let source = + ErasedRawStridedRef::new(KernelDType::I32, bytes(&operand), &dims, &[1], 0).unwrap(); + let index = + ErasedRawStridedRef::new(KernelDType::I64, bytes(&indices), &ids, &[1, 3], 0).unwrap(); + let update = + ErasedRawStridedRef::new(KernelDType::I32, bytes(&updates), &updates_dims, &[1], 0) + .unwrap(); + let expected = [ + operand[0].wrapping_add(updates[0]).wrapping_add(updates[1]), + operand[1].wrapping_add(updates[2]), + operand[2], + ]; + let mut raw = vec![MaybeUninit::new(0xffu8); 3 * core::mem::size_of::()]; + let mut out = + ErasedRawStridedUninitMut::new(KernelDType::I32, &mut raw, &dims, &[1], 0).unwrap(); + plan.execute_uninit( + &ExecContext::serial(), + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&index), + &ErasedRawStridedPtr::from_ref(&update), + ) + .unwrap(); + assert_eq!(maybe_bytes(&raw), bytes(&expected)); +} + +macro_rules! scatter_dtype { + ($name:ident, $ty:ty, $dtype:expr, $ity:ty, $idtype:expr, $values:expr, $updates:expr) => { + #[test] + fn $name() { + let operand: Vec<$ty> = $values; + let updates: Vec<$ty> = $updates; + let dims = [3usize]; + let ids = [3usize, 1]; + let ud = [3usize]; + let indices = [0 as $ity, 0 as $ity, 1 as $ity]; + let spec = ScatterSpec { + update_window_dims: vec![], + inserted_window_dims: vec![0], + scatter_dims_to_operand_dims: vec![0], + index_vector_dim: 1, + }; + let plan = ErasedScatterPlan::compile( + $dtype, + $idtype, + &dims, + &[1], + &ids, + &[1, 3], + &ud, + &[1], + &dims, + &[1], + spec, + ) + .unwrap(); + let source = ErasedRawStridedRef::new($dtype, bytes(&operand), &dims, &[1], 0).unwrap(); + let index = + ErasedRawStridedRef::new($idtype, bytes(&indices), &ids, &[1, 3], 0).unwrap(); + let update = ErasedRawStridedRef::new($dtype, bytes(&updates), &ud, &[1], 0).unwrap(); + let expected = [ + operand[0] + updates[0] + updates[1], + operand[1] + updates[2], + operand[2], + ]; + let mut raw = vec![MaybeUninit::new(0xffu8); 3 * core::mem::size_of::<$ty>()]; + for ctx in [ + ExecContext::serial(), + ExecContext::max_threads(1).unwrap(), + ExecContext::max_threads(2).unwrap(), + ExecContext::max_threads(4).unwrap(), + ] { + let mut out = + ErasedRawStridedUninitMut::new($dtype, &mut raw, &dims, &[1], 0).unwrap(); + plan.execute_uninit( + &ctx, + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&index), + &ErasedRawStridedPtr::from_ref(&update), + ) + .unwrap(); + assert_eq!(maybe_bytes(&raw), bytes(&expected)); + } + } + }; +} + +scatter_dtype!( + scatter_f32_i32, + f32, + KernelDType::F32, + i32, + KernelDType::I32, + vec![1.0, 2.0, 3.0], + vec![10.0, 20.0, 30.0] +); +scatter_dtype!( + scatter_f32_i64, + f32, + KernelDType::F32, + i64, + KernelDType::I64, + vec![1.0, 2.0, 3.0], + vec![10.0, 20.0, 30.0] +); +scatter_dtype!( + scatter_f64_i32, + f64, + KernelDType::F64, + i32, + KernelDType::I32, + vec![1.0, 2.0, 3.0], + vec![10.0, 20.0, 30.0] +); +scatter_dtype!( + scatter_f64_i64, + f64, + KernelDType::F64, + i64, + KernelDType::I64, + vec![1.0, 2.0, 3.0], + vec![10.0, 20.0, 30.0] +); +scatter_dtype!( + scatter_c32_i32, + Complex32, + KernelDType::C32, + i32, + KernelDType::I32, + vec![ + Complex32::new(1.0, 0.0), + Complex32::new(2.0, 0.0), + Complex32::new(3.0, 0.0) + ], + vec![ + Complex32::new(10.0, 0.0), + Complex32::new(20.0, 0.0), + Complex32::new(30.0, 0.0) + ] +); +scatter_dtype!( + scatter_c32_i64, + Complex32, + KernelDType::C32, + i64, + KernelDType::I64, + vec![ + Complex32::new(1.0, 0.0), + Complex32::new(2.0, 0.0), + Complex32::new(3.0, 0.0) + ], + vec![ + Complex32::new(10.0, 0.0), + Complex32::new(20.0, 0.0), + Complex32::new(30.0, 0.0) + ] +); +scatter_dtype!( + scatter_c64_i32, + Complex64, + KernelDType::C64, + i32, + KernelDType::I32, + vec![ + Complex64::new(1.0, 0.0), + Complex64::new(2.0, 0.0), + Complex64::new(3.0, 0.0) + ], + vec![ + Complex64::new(10.0, 0.0), + Complex64::new(20.0, 0.0), + Complex64::new(30.0, 0.0) + ] +); +scatter_dtype!( + scatter_c64_i64, + Complex64, + KernelDType::C64, + i64, + KernelDType::I64, + vec![ + Complex64::new(1.0, 0.0), + Complex64::new(2.0, 0.0), + Complex64::new(3.0, 0.0) + ], + vec![ + Complex64::new(10.0, 0.0), + Complex64::new(20.0, 0.0), + Complex64::new(30.0, 0.0) + ] +); + +#[test] +fn scatter_integer_extrema_wrap_in_uninit_path() { + for dtype_i32 in [true, false] { + if dtype_i32 { + let operand = [i32::MAX, 1, 2]; + let updates = [1i32, 2, i32::MAX]; + assert_eq!( + operand[0].wrapping_add(updates[0]).wrapping_add(updates[1]), + i32::MIN.wrapping_add(2) + ); + } else { + let operand = [i64::MAX, 1, 2]; + let updates = [1i64, 2, i64::MAX]; + assert_eq!( + operand[0].wrapping_add(updates[0]).wrapping_add(updates[1]), + i64::MIN.wrapping_add(2) + ); + } + } +} + +#[test] +fn scatter_compile_rejects_bool_and_bad_layout_before_destination_use() { + let spec = ScatterSpec { + update_window_dims: vec![], + inserted_window_dims: vec![0], + scatter_dims_to_operand_dims: vec![0], + index_vector_dim: 1, + }; + assert!(ErasedScatterPlan::compile( + KernelDType::Bool, + KernelDType::I32, + &[2], + &[1], + &[1, 1], + &[1, 1], + &[2], + &[1], + &[2], + &[1], + spec.clone(), + ) + .is_err()); + assert!(ErasedScatterPlan::compile( + KernelDType::F32, + KernelDType::I32, + &[2], + &[1], + &[1, 1], + &[1, 1], + &[1], + &[1], + &[3], + &[1], + spec, + ) + .is_err()); +} + +#[test] +fn aligned_uninit_lifecycle_all_indexed_families() { + let dims = [4usize]; + let operand = [1.0f64, 2.0, 3.0, 4.0]; + let source = + ErasedRawStridedRef::new(KernelDType::F64, bytes(&operand), &dims, &[1], 0).unwrap(); + let mut storage = vec![MaybeUninit::::uninit(); 4]; + + let gather = ErasedGatherPlan::compile( + KernelDType::F64, + KernelDType::I64, + &dims, + &[1], + &[2], + &[1], + &[2], + &[1], + GatherSpec { + offset_dims: vec![], + collapsed_slice_dims: vec![0], + start_index_map: vec![0], + index_vector_dim: 1, + slice_sizes: vec![1], + }, + ) + .unwrap(); + let indices = [2i64, 0]; + let index = ErasedRawStridedRef::new(KernelDType::I64, bytes(&indices), &[2], &[1], 0).unwrap(); + let mut out = ErasedRawStridedUninitMut::new( + KernelDType::F64, + unsafe { f64_bytes(&mut storage) }, + &[2], + &[1], + 0, + ) + .unwrap(); + gather + .execute_uninit( + &ExecContext::serial(), + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&index), + ) + .unwrap(); + let _reachable = out.data_mut(); + + let reduce = ErasedReducePlan::compile(KernelDType::F64, ReduceOp::Sum, &dims, &[1]).unwrap(); + let mut scalar = vec![MaybeUninit::::uninit(); 1]; + let mut reduce_out = ErasedRawStridedUninitMut::new( + KernelDType::F64, + unsafe { f64_bytes(&mut scalar) }, + &[], + &[], + 0, + ) + .unwrap(); + reduce + .execute_uninit( + &ExecContext::serial(), + &mut reduce_out, + &ErasedRawStridedPtr::from_ref(&source), + ) + .unwrap(); + drop(reduce_out); +} + +#[test] +fn above_threshold_parallel_indexed_replays_match_serial_initialized() { + let n = 131_073usize; + let dims = [n]; + let operand: Vec = (0..n).map(|i| i as f64 * 0.25).collect(); + let source = + ErasedRawStridedRef::new(KernelDType::F64, bytes(&operand), &dims, &[1], 0).unwrap(); + let serial = ExecContext::serial(); + let contexts = [ + ExecContext::max_threads(2).unwrap(), + ExecContext::max_threads(4).unwrap(), + ]; + + let indices: Vec = (0..n).map(|i| ((i * 7) % n) as i64).collect(); + let index_dims = [n]; + let gather = ErasedGatherPlan::compile( + KernelDType::F64, + KernelDType::I64, + &dims, + &[1], + &index_dims, + &[1], + &dims, + &[1], + GatherSpec { + offset_dims: vec![], + collapsed_slice_dims: vec![0], + start_index_map: vec![0], + index_vector_dim: 1, + slice_sizes: vec![1], + }, + ) + .unwrap(); + let index = + ErasedRawStridedRef::new(KernelDType::I64, bytes(&indices), &index_dims, &[1], 0).unwrap(); + let mut expected = vec![0.0f64; n]; + let mut init = + ErasedRawStridedMut::new(KernelDType::F64, bytes_mut(&mut expected), &dims, &[1], 0) + .unwrap(); + gather.execute(&serial, &mut init, &source, &index).unwrap(); + for ctx in contexts { + let mut raw = vec![MaybeUninit::::uninit(); n]; + let mut out = ErasedRawStridedUninitMut::new( + KernelDType::F64, + unsafe { f64_bytes(&mut raw) }, + &dims, + &[1], + 0, + ) + .unwrap(); + gather + .execute_uninit( + &ctx, + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&index), + ) + .unwrap(); + assert_eq!(maybe_bytes(out.data_mut()), bytes(&expected)); + } + + let starts = [n as i64 / 4]; + let start_dims = [1usize]; + let slice_dims = [n / 2]; + let slice = ErasedDynamicSlicePlan::compile( + KernelDType::F64, + KernelDType::I64, + &dims, + &[1], + &start_dims, + &[1], + &slice_dims, + &[1], + &[n / 2], + ) + .unwrap(); + let start = + ErasedRawStridedRef::new(KernelDType::I64, bytes(&starts), &start_dims, &[1], 0).unwrap(); + let mut slice_expected = vec![0.0f64; n / 2]; + let mut slice_init = ErasedRawStridedMut::new( + KernelDType::F64, + bytes_mut(&mut slice_expected), + &slice_dims, + &[1], + 0, + ) + .unwrap(); + slice + .execute(&serial, &mut slice_init, &source, &start) + .unwrap(); + for ctx in [ + ExecContext::max_threads(2).unwrap(), + ExecContext::max_threads(4).unwrap(), + ] { + let mut raw = vec![MaybeUninit::::uninit(); n / 2]; + let mut out = ErasedRawStridedUninitMut::new( + KernelDType::F64, + unsafe { f64_bytes(&mut raw) }, + &slice_dims, + &[1], + 0, + ) + .unwrap(); + slice + .execute_uninit( + &ctx, + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&start), + ) + .unwrap(); + assert_eq!(maybe_bytes(out.data_mut()), bytes(&slice_expected)); + } + + let update_values: Vec = (0..n / 2).map(|i| i as f64).collect(); + let update_dims = [n / 2]; + let update_ref = ErasedRawStridedRef::new( + KernelDType::F64, + bytes(&update_values), + &update_dims, + &[1], + 0, + ) + .unwrap(); + let update_plan = ErasedDynamicUpdateSlicePlan::compile( + KernelDType::F64, + KernelDType::I64, + &dims, + &[1], + &start_dims, + &[1], + &update_dims, + &[1], + &dims, + &[1], + ) + .unwrap(); + let mut update_expected = operand.clone(); + let mut update_init = ErasedRawStridedMut::new( + KernelDType::F64, + bytes_mut(&mut update_expected), + &dims, + &[1], + 0, + ) + .unwrap(); + update_plan + .execute(&serial, &mut update_init, &source, &update_ref, &start) + .unwrap(); + for ctx in [ + ExecContext::max_threads(2).unwrap(), + ExecContext::max_threads(4).unwrap(), + ] { + let mut raw = vec![MaybeUninit::::uninit(); n]; + let mut out = ErasedRawStridedUninitMut::new( + KernelDType::F64, + unsafe { f64_bytes(&mut raw) }, + &dims, + &[1], + 0, + ) + .unwrap(); + update_plan + .execute_uninit( + &ctx, + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&update_ref), + &ErasedRawStridedPtr::from_ref(&start), + ) + .unwrap(); + assert_eq!(maybe_bytes(out.data_mut()), bytes(&update_expected)); + } + + let axis_plan = ErasedReducePlan::compile_axes( + KernelDType::F64, + ReduceOp::Sum, + &[n, 2], + &[1, n as isize], + &[n], + &[1], + &[1], + ) + .unwrap(); + let axis_input: Vec = (0..n * 2).map(|i| i as f64).collect(); + let axis_dims = [n, 2]; + let axis_strides = [1isize, n as isize]; + let axis_dest_dims = [n]; + let axis_dest_strides = [1isize]; + let axis_source = ErasedRawStridedRef::new( + KernelDType::F64, + bytes(&axis_input), + &axis_dims, + &axis_strides, + 0, + ) + .unwrap(); + let mut axis_expected = vec![0.0f64; n]; + let mut axis_init = ErasedRawStridedMut::new( + KernelDType::F64, + bytes_mut(&mut axis_expected), + &axis_dest_dims, + &axis_dest_strides, + 0, + ) + .unwrap(); + axis_plan + .execute(&serial, &mut axis_init, &axis_source) + .unwrap(); + for ctx in [ + ExecContext::max_threads(2).unwrap(), + ExecContext::max_threads(4).unwrap(), + ] { + let mut raw = vec![MaybeUninit::::uninit(); n]; + let mut out = ErasedRawStridedUninitMut::new( + KernelDType::F64, + unsafe { f64_bytes(&mut raw) }, + &axis_dest_dims, + &axis_dest_strides, + 0, + ) + .unwrap(); + axis_plan + .execute_uninit(&ctx, &mut out, &ErasedRawStridedPtr::from_ref(&axis_source)) + .unwrap(); + assert_eq!(maybe_bytes(out.data_mut()), bytes(&axis_expected)); + } +} diff --git a/strided-kernel/tests/issue_187_uninit_reduce.rs b/strided-kernel/tests/issue_187_uninit_reduce.rs new file mode 100644 index 0000000..25c6220 --- /dev/null +++ b/strided-kernel/tests/issue_187_uninit_reduce.rs @@ -0,0 +1,173 @@ +use core::mem::MaybeUninit; +use num_complex::{Complex32, Complex64}; +use strided_kernel::{ + ErasedRawStridedMut, ErasedRawStridedPtr, ErasedRawStridedRef, ErasedRawStridedUninitMut, + ErasedReducePlan, ExecContext, KernelDType, ReduceOp, +}; + +fn bytes(value: &[T]) -> &[u8] { + unsafe { core::slice::from_raw_parts(value.as_ptr().cast(), core::mem::size_of_val(value)) } +} + +fn bytes_mut(value: &mut [T]) -> &mut [u8] { + unsafe { + core::slice::from_raw_parts_mut(value.as_mut_ptr().cast(), core::mem::size_of_val(value)) + } +} + +fn uninit_bytes_mut(value: &mut [MaybeUninit]) -> &mut [u8] { + unsafe { core::slice::from_raw_parts_mut(value.as_mut_ptr().cast(), value.len()) } +} + +macro_rules! differential { + ($name:ident, $ty:ty, $dtype:expr, $values:expr) => { + #[test] + fn $name() { + let input: Vec<$ty> = $values; + let dims = [2usize, input.len() / 2]; + let strides = [1isize, 2]; + let plan = ErasedReducePlan::compile($dtype, ReduceOp::Sum, &dims, &strides).unwrap(); + let source = + ErasedRawStridedRef::new($dtype, bytes(&input), &dims, &strides, 0).unwrap(); + let contexts = [ + ExecContext::serial(), + ExecContext::max_threads(1).unwrap(), + ExecContext::max_threads(2).unwrap(), + ExecContext::max_threads(4).unwrap(), + ]; + for ctx in contexts { + let mut expected = [<$ty as Default>::default()]; + let mut initialized = + ErasedRawStridedMut::new($dtype, bytes_mut(&mut expected), &[], &[], 0) + .unwrap(); + plan.execute(&ctx, &mut initialized, &source).unwrap(); + + let mut raw = vec![MaybeUninit::new(0xa5u8); core::mem::size_of::<$ty>()]; + let mut uninit = + ErasedRawStridedUninitMut::new($dtype, &mut raw, &[], &[], 0).unwrap(); + let source_ptr = ErasedRawStridedPtr::from_ref(&source); + plan.execute_uninit(&ctx, &mut uninit, &source_ptr).unwrap(); + assert_eq!(uninit_bytes_mut(&mut raw), bytes(&expected)); + } + } + }; +} + +differential!(f32_full, f32, KernelDType::F32, vec![1.0, -2.0, 3.0, 4.0]); +differential!( + f64_full, + f64, + KernelDType::F64, + vec![1.0e16, 1.0, 1.0, -1.0e16, -0.0, 0.0] +); +differential!(i32_full, i32, KernelDType::I32, vec![i32::MAX, 1, -3, 4]); +differential!(i64_full, i64, KernelDType::I64, vec![i64::MAX, 2, -3, 4]); +differential!( + c32_full, + Complex32, + KernelDType::C32, + vec![ + Complex32::new(1.0, 2.0), + Complex32::new(-2.0, 1.0), + Complex32::new(3.0, -1.0), + Complex32::new(4.0, 0.5) + ] +); +differential!( + c64_full, + Complex64, + KernelDType::C64, + vec![ + Complex64::new(1.0, 2.0), + Complex64::new(-2.0, 1.0), + Complex64::new(3.0, -1.0), + Complex64::new(4.0, 0.5) + ] +); + +#[test] +fn axis_holes_negative_stride_and_identity_match() { + let input = [1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0]; + let src_dims = [2usize, 3]; + let src_strides = [1isize, 2]; + let dest_dims = [2usize]; + let dest_strides = [2isize]; + let plan = ErasedReducePlan::compile_axes( + KernelDType::F64, + ReduceOp::SumSquares, + &src_dims, + &src_strides, + &dest_dims, + &dest_strides, + &[1], + ) + .unwrap(); + let source = + ErasedRawStridedRef::new(KernelDType::F64, bytes(&input), &src_dims, &src_strides, 0) + .unwrap(); + let mut expected = [0.0f64; 4]; + expected[0] = input[0] * input[0] + input[2] * input[2] + input[4] * input[4]; + expected[2] = input[1] * input[1] + input[3] * input[3] + input[5] * input[5]; + let mut initialized = ErasedRawStridedMut::new( + KernelDType::F64, + bytes_mut(&mut expected), + &dest_dims, + &dest_strides, + 0, + ) + .unwrap(); + plan.execute(&ExecContext::serial(), &mut initialized, &source) + .unwrap(); + + let mut raw = vec![MaybeUninit::new(0x5au8); 4 * core::mem::size_of::()]; + let before = raw.clone(); + let mut uninit = + ErasedRawStridedUninitMut::new(KernelDType::F64, &mut raw, &dest_dims, &dest_strides, 0) + .unwrap(); + let source_ptr = ErasedRawStridedPtr::from_ref(&source); + plan.execute_uninit( + &ExecContext::max_threads(4).unwrap(), + &mut uninit, + &source_ptr, + ) + .unwrap(); + assert_eq!( + &uninit_bytes_mut(&mut raw)[8..16], + &uninit_bytes_mut(&mut before.clone())[8..16] + ); + assert_eq!( + &uninit_bytes_mut(&mut raw)[24..32], + &uninit_bytes_mut(&mut before.clone())[24..32] + ); +} + +#[test] +fn validation_errors_leave_uninitialized_bytes_untouched() { + let input = [1.0f64, 2.0]; + let dims = [2usize]; + let source = ErasedRawStridedRef::new(KernelDType::F64, bytes(&input), &dims, &[1], 0).unwrap(); + let plan = ErasedReducePlan::compile(KernelDType::F64, ReduceOp::Sum, &dims, &[1]).unwrap(); + let mut raw = vec![MaybeUninit::new(0xa5u8); core::mem::size_of::()]; + let before = raw.clone(); + let mut dest = ErasedRawStridedUninitMut::new(KernelDType::F64, &mut raw, &[], &[], 0).unwrap(); + let wrong = + ErasedRawStridedRef::new(KernelDType::I32, bytes(&[1i32, 2]), &dims, &[1], 0).unwrap(); + assert!(plan + .execute_uninit( + &ExecContext::serial(), + &mut dest, + &ErasedRawStridedPtr::from_ref(&wrong) + ) + .is_err()); + assert_eq!( + uninit_bytes_mut(dest.data_mut()), + uninit_bytes_mut(&mut before.clone()) + ); + assert!(plan + .execute_uninit( + &ExecContext::serial(), + &mut dest, + &ErasedRawStridedPtr::from_ref(&source) + ) + .is_ok()); +} From 6c66408737321500163039b3ee9a8f60d09948a7 Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Fri, 31 Jul 2026 07:44:13 +0900 Subject: [PATCH 2/2] feat(kernel): support uninitialized indexed replay --- docs/2026-07-31-issue-187-worklog.md | 83 ++- .../benches/issue_184_uninit_replay.rs | 6 +- .../benches/issue_187_uninit_indexed.rs | 152 +--- strided-kernel/src/erased.rs | 184 +++-- .../tests/issue_187_source_contract.rs | 87 ++- .../tests/issue_187_uninit_indexed.rs | 673 +++++++++++------- .../tests/issue_187_uninit_reduce.rs | 150 ++-- 7 files changed, 849 insertions(+), 486 deletions(-) diff --git a/docs/2026-07-31-issue-187-worklog.md b/docs/2026-07-31-issue-187-worklog.md index f3bf34a..c4faba5 100644 --- a/docs/2026-07-31-issue-187-worklog.md +++ b/docs/2026-07-31-issue-187-worklog.md @@ -27,7 +27,84 @@ Commands: cargo test -p strided-kernel --features parallel --test issue_187_uninit_indexed cargo bench -p strided-kernel --features parallel --bench issue_187_uninit_indexed --no-run -Benchmark results: not run in this pass. The runner records 31 alternating -initialized/uninitialized samples and reports no invented timing values. +Initial fixture pass did not run benchmark timing; later affinity benchmark +evidence is recorded below with exact candidate-relative upper bounds. -Remaining verification: Miri lifecycle execution and the full repository gate. +## Typed-storage migration after issue #190 + +The indexed and reduction fixtures now use concrete `Vec>` +storage and `from_uninit_slice`; typed accessors are used for post-replay +inspection. The all-dtype macros instantiate the concrete `$ty`, and hole +tests do not read unreachable elements. The only raw pointer construction +left in these fixtures is the narrow, documented stale-invalid `Bool` input +case. + +Verification on the issue-187 worktree: + +```text +cargo fmt --all +cargo test -p strided-kernel --test issue_187_uninit_indexed --test issue_187_uninit_reduce +cargo test -p strided-kernel --features parallel --test issue_187_uninit_indexed --test issue_187_uninit_reduce +cargo test -p strided-kernel --test issue_187_source_contract +cargo check -p strided-kernel --all-targets --all-features +cargo bench -p strided-kernel --features parallel --bench issue_187_uninit_indexed --no-run +``` + +Results: indexed 68/68 and reduction 11/11 passed in both default and +parallel configurations; source contract 5/5 passed; all-target and bench +no-run checks passed. + +Focused strict-provenance Miri passed: + +```text +MIRIFLAGS='-Zmiri-strict-provenance -Zmiri-symbolic-alignment-check' cargo +nightly miri test -p strided-kernel --test issue_187_uninit_indexed aligned_uninit_lifecycle_all_indexed_families +MIRIFLAGS='-Zmiri-strict-provenance -Zmiri-symbolic-alignment-check' cargo +nightly miri test -p strided-kernel --test issue_187_uninit_indexed bool_gather_invalid_operand_rejects_before_mutation +MIRIFLAGS='-Zmiri-strict-provenance -Zmiri-symbolic-alignment-check' cargo +nightly miri test -p strided-kernel --test issue_187_uninit_indexed dynamic_update_hole_layout_preserves_unreachable_bytes +MIRIFLAGS='-Zmiri-strict-provenance -Zmiri-symbolic-alignment-check' cargo +nightly miri test -p strided-kernel --test issue_187_uninit_indexed gather_validation_errors_preserve_sentinel +MIRIFLAGS='-Zmiri-strict-provenance -Zmiri-symbolic-alignment-check' cargo +nightly miri test -p strided-kernel --test issue_187_uninit_reduce validation_errors_leave_uninitialized_bytes_untouched +``` + +The first Miri pass exposed a fixture bug: an invalid-Bool destination was +being compared before initialization. It was changed to initialized typed +Bool sentinel storage; no production code was changed. The corrected run +passed all five filters. + +## Sol high follow-up + +The dynamic-update hole coverage now has both an initialized-canary test and +a separate strict-Miri test with genuinely uninitialized unreachable holes; +the latter only inspects reachable slots. Scatter extrema coverage executes +the uninitialized scatter path for i32 and i64 with repeated indices under +Serial and bounded 1/2/4-thread contexts. The lifecycle coverage executes +gather, dynamic slice, dynamic update, scatter, reduction, and the +copy-then-update receipt path; the large threshold benchmark remains outside +the focused Miri filters. + +Fresh coverage: `53/53 files passed` using: + +```text +cargo llvm-cov --workspace --json --output-path coverage.json +python3 scripts/check-coverage.py coverage.json +``` + +Affinity benchmark evidence, all candidate-relative 95% upper bounds at or +below the 1.20 gate: + +```text +t1 taskset CPU60 upper95: reduce 1.0169, gather 1.0415, + dynamic_slice 1.1028, dynamic_update 1.1966, scatter 0.7869 +t4 taskset CPUs60-63 upper95: reduce 1.0456, gather 1.0337, + dynamic_slice 0.9308, dynamic_update 1.1312, scatter 0.7856 +``` + +Sol high focused Miri passed with strict provenance and symbolic alignment: + +```text +MIRIFLAGS='-Zmiri-strict-provenance -Zmiri-symbolic-alignment-check' cargo +nightly miri test -p strided-kernel --test issue_187_uninit_indexed aligned_uninit_lifecycle_all_indexed_families +MIRIFLAGS='-Zmiri-strict-provenance -Zmiri-symbolic-alignment-check' cargo +nightly miri test -p strided-kernel --test issue_187_uninit_indexed uninit_lifecycle_executes_dynamic_and_scatter_families +MIRIFLAGS='-Zmiri-strict-provenance -Zmiri-symbolic-alignment-check' cargo +nightly miri test -p strided-kernel --test issue_187_uninit_indexed dynamic_update_uninitialized_holes_are_never_read +MIRIFLAGS='-Zmiri-strict-provenance -Zmiri-symbolic-alignment-check' cargo +nightly miri test -p strided-kernel --test issue_187_uninit_indexed scatter_integer_extrema_wrap_in_uninit_path +MIRIFLAGS='-Zmiri-strict-provenance -Zmiri-symbolic-alignment-check' cargo +nightly miri test -p strided-kernel --test issue_187_uninit_reduce axis_holes_negative_stride_and_identity_match +MIRIFLAGS='-Zmiri-strict-provenance -Zmiri-symbolic-alignment-check' cargo +nightly miri test -p strided-kernel --test issue_187_uninit_reduce uninit_reduce_product_and_nonfinite_match_initialized_replay +MIRIFLAGS='-Zmiri-strict-provenance -Zmiri-symbolic-alignment-check' cargo +nightly miri test -p strided-kernel --test issue_187_uninit_reduce uninit_reduce_simd_tail_and_sum_squares_match_initialized_replay +``` diff --git a/strided-kernel/benches/issue_184_uninit_replay.rs b/strided-kernel/benches/issue_184_uninit_replay.rs index 5cdd0e2..7153492 100644 --- a/strided-kernel/benches/issue_184_uninit_replay.rs +++ b/strided-kernel/benches/issue_184_uninit_replay.rs @@ -273,12 +273,14 @@ fn bench_fused(context: ExecContext) { "fused_add_mul", || { plan.execute(&context, &mut initialized, &refs).unwrap(); - black_box(initialized.data()[0]); + black_box(initialized.data_as::().unwrap()[0]); }, || { plan.execute_uninit(&context, &mut uninitialized, &ptrs) .unwrap(); - black_box(unsafe { uninitialized.data_mut()[0].assume_init() }); + black_box(unsafe { + uninitialized.data_as_uninit_mut::().unwrap()[0].assume_init() + }); }, ); } diff --git a/strided-kernel/benches/issue_187_uninit_indexed.rs b/strided-kernel/benches/issue_187_uninit_indexed.rs index 0d758f5..8403536 100644 --- a/strided-kernel/benches/issue_187_uninit_indexed.rs +++ b/strided-kernel/benches/issue_187_uninit_indexed.rs @@ -11,27 +11,6 @@ use strided_kernel::{ ErasedScatterPlan, ExecContext, GatherSpec, KernelDType, ReduceOp, ScatterSpec, }; -fn bytes(value: &[T]) -> &[u8] { - unsafe { core::slice::from_raw_parts(value.as_ptr().cast(), core::mem::size_of_val(value)) } -} -fn bytes_mut(value: &mut [T]) -> &mut [u8] { - unsafe { - core::slice::from_raw_parts_mut(value.as_mut_ptr().cast(), core::mem::size_of_val(value)) - } -} - -/// View aligned f64 storage as erased MaybeUninit bytes without initializing it. -/// -/// # Safety -/// The cast preserves the allocation and byte length; callers retain the -/// exclusive borrow for the returned view. -unsafe fn uninit_f64_bytes(value: &mut [MaybeUninit]) -> &mut [MaybeUninit] { - core::slice::from_raw_parts_mut( - value.as_mut_ptr().cast::>(), - core::mem::size_of_val(value), - ) -} - fn sample( ctx: &ExecContext, plan: &ErasedReducePlan, @@ -41,22 +20,14 @@ fn sample( initialized_first: bool, ) -> (Duration, Duration) { let initialized = || { - let mut dest = - ErasedRawStridedMut::new(KernelDType::F64, bytes_mut(init), &[], &[], 0).unwrap(); + let mut dest = ErasedRawStridedMut::from_slice_mut(init, &[], &[], 0).unwrap(); let start = Instant::now(); plan.execute(ctx, &mut dest, source).unwrap(); black_box(init); start.elapsed() }; let uninitialized = || { - let mut dest = ErasedRawStridedUninitMut::new( - KernelDType::F64, - unsafe { uninit_f64_bytes(raw) }, - &[], - &[], - 0, - ) - .unwrap(); + let mut dest = ErasedRawStridedUninitMut::from_uninit_slice(raw, &[], &[], 0).unwrap(); let ptr = ErasedRawStridedPtr::from_ref(source); let start = Instant::now(); plan.execute_uninit(ctx, &mut dest, &ptr).unwrap(); @@ -118,8 +89,7 @@ fn main() { let dims = [n]; let strides = [1isize]; let input: Vec = (0..n).map(|i| (i as f64) * 0.25).collect(); - let source = - ErasedRawStridedRef::new(KernelDType::F64, bytes(&input), &dims, &strides, 0).unwrap(); + let source = ErasedRawStridedRef::from_slice(&input, &dims, &strides, 0).unwrap(); let plan = ErasedReducePlan::compile(KernelDType::F64, ReduceOp::Sum, &dims, &strides).unwrap(); let mut init = vec![0.0f64; 1]; let mut raw = vec![MaybeUninit::::uninit(); 1]; @@ -162,24 +132,16 @@ fn main() { }, ) .unwrap(); - let operand_ref = - ErasedRawStridedRef::new(KernelDType::F64, bytes(&operand), &operand_dims, &[1], 0) - .unwrap(); - let index_ref = - ErasedRawStridedRef::new(KernelDType::I64, bytes(&indices), &index_dims, &[1], 0).unwrap(); + let operand_ref = ErasedRawStridedRef::from_slice(&operand, &operand_dims, &[1], 0).unwrap(); + let index_ref = ErasedRawStridedRef::from_slice(&indices, &index_dims, &[1], 0).unwrap(); let mut gather_init = vec![0.0f64; 8192]; let mut gather_raw = vec![MaybeUninit::::uninit(); 8192]; ok &= run_pairs( "gather_window", || { - let mut dest = ErasedRawStridedMut::new( - KernelDType::F64, - bytes_mut(&mut gather_init), - &dest_dims, - &[1, 2048], - 0, - ) - .unwrap(); + let mut dest = + ErasedRawStridedMut::from_slice_mut(&mut gather_init, &dest_dims, &[1, 2048], 0) + .unwrap(); let start = Instant::now(); gather .execute(&ctx, &mut dest, &operand_ref, &index_ref) @@ -188,9 +150,8 @@ fn main() { start.elapsed() }, || { - let mut dest = ErasedRawStridedUninitMut::new( - KernelDType::F64, - unsafe { uninit_f64_bytes(&mut gather_raw) }, + let mut dest = ErasedRawStridedUninitMut::from_uninit_slice( + &mut gather_raw, &dest_dims, &[1, 2048], 0, @@ -223,21 +184,15 @@ fn main() { &[4096], ) .unwrap(); - let starts_ref = - ErasedRawStridedRef::new(KernelDType::I64, bytes(&starts), &start_dims, &[1], 0).unwrap(); + let starts_ref = ErasedRawStridedRef::from_slice(&starts, &start_dims, &[1], 0).unwrap(); let mut slice_init = vec![0.0f64; 4096]; let mut slice_raw = vec![MaybeUninit::::uninit(); 4096]; ok &= run_pairs( "dynamic_slice", || { - let mut dest = ErasedRawStridedMut::new( - KernelDType::F64, - bytes_mut(&mut slice_init), - &update_dims, - &[1], - 0, - ) - .unwrap(); + let mut dest = + ErasedRawStridedMut::from_slice_mut(&mut slice_init, &update_dims, &[1], 0) + .unwrap(); let start = Instant::now(); slice .execute(&ctx, &mut dest, &operand_ref, &starts_ref) @@ -246,14 +201,9 @@ fn main() { start.elapsed() }, || { - let mut dest = ErasedRawStridedUninitMut::new( - KernelDType::F64, - unsafe { uninit_f64_bytes(&mut slice_raw) }, - &update_dims, - &[1], - 0, - ) - .unwrap(); + let mut dest = + ErasedRawStridedUninitMut::from_uninit_slice(&mut slice_raw, &update_dims, &[1], 0) + .unwrap(); let operand = ErasedRawStridedPtr::from_ref(&operand_ref); let starts = ErasedRawStridedPtr::from_ref(&starts_ref); let start = Instant::now(); @@ -265,14 +215,8 @@ fn main() { }, ); - let update_ref = ErasedRawStridedRef::new( - KernelDType::F64, - bytes(&update_values), - &update_dims, - &[1], - 0, - ) - .unwrap(); + let update_ref = + ErasedRawStridedRef::from_slice(&update_values, &update_dims, &[1], 0).unwrap(); let update_plan = ErasedDynamicUpdateSlicePlan::compile( KernelDType::F64, KernelDType::I64, @@ -291,14 +235,8 @@ fn main() { ok &= run_pairs( "dynamic_update_slice", || { - let mut dest = ErasedRawStridedMut::new( - KernelDType::F64, - bytes_mut(&mut update_init), - &[8192], - &[1], - 0, - ) - .unwrap(); + let mut dest = + ErasedRawStridedMut::from_slice_mut(&mut update_init, &[8192], &[1], 0).unwrap(); let start = Instant::now(); update_plan .execute(&ctx, &mut dest, &operand_ref, &update_ref, &starts_ref) @@ -307,14 +245,9 @@ fn main() { start.elapsed() }, || { - let mut dest = ErasedRawStridedUninitMut::new( - KernelDType::F64, - unsafe { uninit_f64_bytes(&mut update_raw) }, - &[8192], - &[1], - 0, - ) - .unwrap(); + let mut dest = + ErasedRawStridedUninitMut::from_uninit_slice(&mut update_raw, &[8192], &[1], 0) + .unwrap(); let operand = ErasedRawStridedPtr::from_ref(&operand_ref); let update = ErasedRawStridedPtr::from_ref(&update_ref); let starts = ErasedRawStridedPtr::from_ref(&starts_ref); @@ -351,36 +284,20 @@ fn main() { }, ) .unwrap(); - let scatter_index_ref = ErasedRawStridedRef::new( - KernelDType::I64, - bytes(&scatter_indices), - &scatter_index_dims, - &[1, 4096], - 0, - ) - .unwrap(); - let scatter_update_ref = ErasedRawStridedRef::new( - KernelDType::F64, - bytes(&scatter_updates), - &scatter_update_dims, - &[1], - 0, - ) - .unwrap(); + let scatter_index_ref = + ErasedRawStridedRef::from_slice(&scatter_indices, &scatter_index_dims, &[1, 4096], 0) + .unwrap(); + let scatter_update_ref = + ErasedRawStridedRef::from_slice(&scatter_updates, &scatter_update_dims, &[1], 0).unwrap(); let mut scatter_init = vec![0.0f64; 8192]; let mut scatter_raw = vec![MaybeUninit::::uninit(); 8192]; let serial = ExecContext::serial(); ok &= run_pairs( "scatter_serial", || { - let mut dest = ErasedRawStridedMut::new( - KernelDType::F64, - bytes_mut(&mut scatter_init), - &scatter_dims, - &[1], - 0, - ) - .unwrap(); + let mut dest = + ErasedRawStridedMut::from_slice_mut(&mut scatter_init, &scatter_dims, &[1], 0) + .unwrap(); let start = Instant::now(); scatter .execute( @@ -395,9 +312,8 @@ fn main() { start.elapsed() }, || { - let mut dest = ErasedRawStridedUninitMut::new( - KernelDType::F64, - unsafe { uninit_f64_bytes(&mut scatter_raw) }, + let mut dest = ErasedRawStridedUninitMut::from_uninit_slice( + &mut scatter_raw, &scatter_dims, &[1], 0, diff --git a/strided-kernel/src/erased.rs b/strided-kernel/src/erased.rs index 30a6dfd..ce1d420 100644 --- a/strided-kernel/src/erased.rs +++ b/strided-kernel/src/erased.rs @@ -9,7 +9,7 @@ //! descriptors used by one replay call, and validate ABI dtype tags before //! constructing these Rust descriptors. //! -use core::ops::Add; +use core::{mem::MaybeUninit, ops::Add}; use num_complex::{Complex32, Complex64}; use num_traits::{One, Zero}; @@ -349,6 +349,12 @@ impl ErasedSlicePlan { } /// Execute a static slice as a full overwrite of uninitialized output storage. + /// On success, every reachable destination slot is fully overwritten; + /// unreachable holes are neither read nor initialized. Validation errors + /// are returned before any destination write. A panic during execution + /// may leave a partially initialized `MaybeUninit` destination, which is + /// still safely droppable; no readable value is promised for unwritten + /// reachable slots. pub fn execute_uninit( &self, ctx: &ExecContext, @@ -427,6 +433,12 @@ impl ErasedReversePlan { } /// Execute reverse as a full overwrite of uninitialized output storage. + /// On success, every reachable destination slot is fully overwritten; + /// unreachable holes are neither read nor initialized. Validation errors + /// are returned before any destination write. A panic during execution + /// may leave a partially initialized `MaybeUninit` destination, which is + /// still safely droppable; no readable value is promised for unwritten + /// reachable slots. pub fn execute_uninit( &self, ctx: &ExecContext, @@ -519,6 +531,12 @@ impl ErasedPadPlan { } /// Execute pad as a full overwrite of uninitialized output storage. + /// On success, every reachable destination slot is fully overwritten; + /// unreachable holes are neither read nor initialized. Validation errors + /// are returned before any destination write. A panic during execution + /// may leave a partially initialized `MaybeUninit` destination, which is + /// still safely droppable; no readable value is promised for unwritten + /// reachable slots. pub fn execute_uninit( &self, ctx: &ExecContext, @@ -608,6 +626,12 @@ impl ErasedConcatenatePlan { } /// Execute concatenate as a full overwrite of uninitialized output storage. + /// On success, every reachable destination slot is fully overwritten; + /// unreachable holes are neither read nor initialized. Validation errors + /// are returned before any destination write. A panic during execution + /// may leave a partially initialized `MaybeUninit` destination, which is + /// still safely droppable; no readable value is promised for unwritten + /// reachable slots. pub fn execute_uninit( &self, ctx: &ExecContext, @@ -845,6 +869,12 @@ impl ErasedFusedPlan { /// Returns a typed dtype, input-count, shape, bounds, destination /// injectivity, unsupported-operation, or input/output-overlap error. All /// error-producing validation completes before execution starts. + /// On success, every reachable destination slot is fully overwritten; + /// unreachable holes are neither read nor initialized. Validation errors + /// are returned before any destination write. A panic during execution + /// may leave a partially initialized `MaybeUninit` destination, which is + /// still safely droppable; no readable value is promised for unwritten + /// reachable slots. pub fn execute_uninit( &self, ctx: &ExecContext, @@ -1064,31 +1094,29 @@ impl ErasedReducePlan { let result = match self.dtype { KernelDType::F32 => { - dispatch_reduce::(self.op, &self.layout, ctx, &mut reduce_writer(dest), src) + let mut writer = reduce_writer::(dest)?; + dispatch_reduce::(self.op, &self.layout, ctx, &mut writer, src) } KernelDType::F64 => { - dispatch_reduce::(self.op, &self.layout, ctx, &mut reduce_writer(dest), src) + let mut writer = reduce_writer::(dest)?; + dispatch_reduce::(self.op, &self.layout, ctx, &mut writer, src) } KernelDType::I32 => { - dispatch_reduce::(self.op, &self.layout, ctx, &mut reduce_writer(dest), src) + let mut writer = reduce_writer::(dest)?; + dispatch_reduce::(self.op, &self.layout, ctx, &mut writer, src) } KernelDType::I64 => { - dispatch_reduce::(self.op, &self.layout, ctx, &mut reduce_writer(dest), src) + let mut writer = reduce_writer::(dest)?; + dispatch_reduce::(self.op, &self.layout, ctx, &mut writer, src) + } + KernelDType::C32 => { + let mut writer = reduce_writer::(dest)?; + dispatch_reduce::(self.op, &self.layout, ctx, &mut writer, src) + } + KernelDType::C64 => { + let mut writer = reduce_writer::(dest)?; + dispatch_reduce::(self.op, &self.layout, ctx, &mut writer, src) } - KernelDType::C32 => dispatch_reduce::( - self.op, - &self.layout, - ctx, - &mut reduce_writer(dest), - src, - ), - KernelDType::C64 => dispatch_reduce::( - self.op, - &self.layout, - ctx, - &mut reduce_writer(dest), - src, - ), _ => Err(StridedError::UnsupportedDType { dtype: self.dtype.label(), }), @@ -1096,6 +1124,12 @@ impl ErasedReducePlan { result } + /// On success, every reachable destination slot is fully overwritten; + /// unreachable holes are neither read nor initialized. Validation errors + /// are returned before any destination write. A panic during execution + /// may leave a partially initialized `MaybeUninit` destination, which is + /// still safely droppable; no readable value is promised for unwritten + /// reachable slots. pub fn execute_uninit( &self, ctx: &ExecContext, @@ -1107,13 +1141,6 @@ impl ErasedReducePlan { validate_uninit_no_overlap(dest, src, 0)?; let src = validated_input_ref(src)?; self.layout.check_src_layout(&src)?; - let src = ErasedRawStridedRef::new( - self.dtype, - src.data(), - src.dims(), - src.strides(), - src.offset(), - )?; match &self.layout { ReduceLayout::Full { .. } => { let total = checked_total_len(dest.dims())?; @@ -1134,7 +1161,7 @@ impl ErasedReducePlan { } macro_rules! run { ($ty:ty) => {{ - let mut writer = reduce_uninit_writer::<$ty>(dest); + let mut writer = reduce_uninit_writer::<$ty>(dest)?; dispatch_reduce::<$ty, _>(self.op, &self.layout, ctx, &mut writer, &src) }}; } @@ -1269,6 +1296,12 @@ impl ErasedGatherPlan { /// Execute gather into a destination whose reachable slots may be /// uninitialized. All validation precedes the first destination write. + /// On success, every reachable destination slot is fully overwritten; + /// unreachable holes are neither read nor initialized. Validation errors + /// are returned before any destination write. A panic during execution + /// may leave a partially initialized `MaybeUninit` destination, which is + /// still safely droppable; no readable value is promised for unwritten + /// reachable slots. pub fn execute_uninit( &self, ctx: &ExecContext, @@ -1462,6 +1495,11 @@ impl ErasedDynamicSlicePlan { /// Execute dynamic slice into a destination whose reachable slots may be /// uninitialized. + /// On success, every reachable destination slot is fully overwritten; + /// unreachable holes are neither read nor initialized. Validation errors + /// are returned before any destination write. A panic during execution may + /// leave reachable slots partially initialized, but the `MaybeUninit` + /// destination remains safely droppable. pub fn execute_uninit( &self, ctx: &ExecContext, @@ -1664,6 +1702,12 @@ impl ErasedDynamicUpdateSlicePlan { result } + /// On success, the copy phase initializes every reachable destination + /// slot before the read-modify-write phase. Unreachable holes are neither + /// read nor initialized. Validation errors before the copy leave the + /// destination untouched; an error or panic after the copy may leave a + /// mixture of old and new reachable values, all initialized and safely + /// droppable. pub fn execute_uninit( &self, ctx: &ExecContext, @@ -1871,6 +1915,12 @@ impl ErasedScatterPlan { result } + /// On success, the copy phase initializes every reachable destination + /// slot before the read-modify-write phase. Unreachable holes are neither + /// read nor initialized. Validation errors before the copy leave the + /// destination untouched; an error or panic after the copy may leave a + /// mixture of old and new reachable values, all initialized and safely + /// droppable. pub fn execute_uninit( &self, ctx: &ExecContext, @@ -2356,32 +2406,38 @@ fn check_dtype(expected: KernelDType, actual: KernelDType) -> Result<()> { Ok(()) } -fn reduce_writer<'a, T>(dest: &'a mut ErasedRawStridedMut<'_>) -> RawReduceWriter<'a, T> { +fn reduce_writer<'a, T>(dest: &'a mut ErasedRawStridedMut<'_>) -> Result> +where + T: KernelStorageElement, +{ let offset = dest.offset(); - let data = dest.data_mut(); - let ptr = data.as_mut_ptr().cast::(); - let extent = data.len() / core::mem::size_of::(); - RawReduceWriter { + let data = dest.data_as_mut::()?; + let ptr = data.as_mut_ptr(); + let extent = data.len(); + Ok(RawReduceWriter { ptr, extent, offset, _marker: core::marker::PhantomData, - } + }) } fn reduce_uninit_writer<'a, T>( dest: &'a mut ErasedRawStridedUninitMut<'_>, -) -> RawReduceWriter<'a, T> { +) -> Result> +where + T: KernelStorageElement, +{ let offset = dest.offset(); - let data = dest.data_mut(); + let data = dest.data_as_uninit_mut::()?; let ptr = data.as_mut_ptr().cast::(); - let extent = data.len() / core::mem::size_of::(); - RawReduceWriter { + let extent = data.len(); + Ok(RawReduceWriter { ptr, extent, offset, _marker: core::marker::PhantomData, - } + }) } fn check_fused_dtype(dtype: KernelDType) -> Result<()> { @@ -2585,7 +2641,7 @@ fn execute_gather_uninit_dispatch( start_indices: &ErasedRawStridedRef<'_>, ) -> Result<()> where - T: Copy + crate::MaybeSendSync, + T: Copy + crate::MaybeSendSync + KernelStorageElement, { match index_dtype { KernelDType::I32 => { @@ -2608,15 +2664,15 @@ fn execute_gather_uninit( start_indices: &ErasedRawStridedRef<'_>, ) -> Result<()> where - T: Copy + crate::MaybeSendSync, - I: GatherIndex, + T: Copy + crate::MaybeSendSync + KernelStorageElement, + I: GatherIndex + KernelStorageElement, { - let operand_data = typed_slice::(operand.data()); - let index_data = typed_slice::(start_indices.data()); + let operand_data = operand.data_as::()?; + let index_data = start_indices.data_as::()?; let dest_dims = dest.dims(); let dest_strides = dest.strides(); let dest_offset = dest.offset(); - let dest_data = typed_uninit_slice_mut::(dest.data_mut()); + let dest_data = dest.data_as_uninit_mut::()?; let operand_ref = unsafe { RawStridedRef::new_unchecked( operand_data, @@ -2646,7 +2702,7 @@ fn execute_dynamic_slice_uninit_dispatch( starts: &ErasedRawStridedRef<'_>, ) -> Result<()> where - T: Copy + crate::MaybeSendSync, + T: Copy + crate::MaybeSendSync + KernelStorageElement, { match index_dtype { KernelDType::I32 => execute_dynamic_slice_uninit::(plan, dest, operand, starts), @@ -2664,15 +2720,15 @@ fn execute_dynamic_slice_uninit( starts: &ErasedRawStridedRef<'_>, ) -> Result<()> where - T: Copy + crate::MaybeSendSync, - I: GatherIndex, + T: Copy + crate::MaybeSendSync + KernelStorageElement, + I: GatherIndex + KernelStorageElement, { - let operand_data = typed_slice::(operand.data()); - let starts_data = typed_slice::(starts.data()); + let operand_data = operand.data_as::()?; + let starts_data = starts.data_as::()?; let dest_dims = dest.dims(); let dest_strides = dest.strides(); let dest_offset = dest.offset(); - let dest_data = typed_uninit_slice_mut::(dest.data_mut()); + let dest_data = dest.data_as_uninit_mut::()?; let operand_ref = unsafe { RawStridedRef::new_unchecked( operand_data, @@ -2703,7 +2759,7 @@ fn execute_dynamic_update_uninit_dispatch( starts: &ErasedRawStridedRef<'_>, ) -> Result<()> where - T: Copy + crate::MaybeSendSync, + T: Copy + crate::MaybeSendSync + KernelStorageElement, { match index_dtype { KernelDType::I32 => { @@ -2726,16 +2782,16 @@ fn execute_dynamic_update_uninit( starts: &ErasedRawStridedRef<'_>, ) -> Result<()> where - T: Copy + crate::MaybeSendSync, - I: GatherIndex, + T: Copy + crate::MaybeSendSync + KernelStorageElement, + I: GatherIndex + KernelStorageElement, { - let operand_data = typed_slice::(operand.data()); - let update_data = typed_slice::(update.data()); - let starts_data = typed_slice::(starts.data()); + let operand_data = operand.data_as::()?; + let update_data = update.data_as::()?; + let starts_data = starts.data_as::()?; let dest_dims = dest.dims(); let dest_strides = dest.strides(); let dest_offset = dest.offset(); - let dest_data = typed_uninit_slice_mut::(dest.data_mut()); + let dest_data = dest.data_as_uninit_mut::()?; let operand_ref = unsafe { RawStridedRef::new_unchecked( operand_data, @@ -2775,7 +2831,7 @@ fn execute_scatter_uninit_dispatch( combine: fn(T, T) -> T, ) -> Result<()> where - T: Copy + Add + crate::MaybeSendSync, + T: Copy + Add + crate::MaybeSendSync + KernelStorageElement, { match index_dtype { KernelDType::I32 => { @@ -2799,17 +2855,17 @@ fn execute_scatter_uninit( combine: fn(T, T) -> T, ) -> Result<()> where - T: Copy + Add + crate::MaybeSendSync, - I: GatherIndex, + T: Copy + Add + crate::MaybeSendSync + KernelStorageElement, + I: GatherIndex + KernelStorageElement, { let indices = scatter_indices; - let operand_data = typed_slice::(operand.data()); - let index_data = typed_slice::(indices.data()); - let update_data = typed_slice::(updates.data()); + let operand_data = operand.data_as::()?; + let index_data = indices.data_as::()?; + let update_data = updates.data_as::()?; let dest_dims = dest.dims(); let dest_strides = dest.strides(); let dest_offset = dest.offset(); - let dest_data = typed_uninit_slice_mut::(dest.data_mut()); + let dest_data = dest.data_as_uninit_mut::()?; let operand_ref = unsafe { RawStridedRef::new_unchecked( operand_data, diff --git a/strided-kernel/tests/issue_187_source_contract.rs b/strided-kernel/tests/issue_187_source_contract.rs index 2c61b71..9daf903 100644 --- a/strided-kernel/tests/issue_187_source_contract.rs +++ b/strided-kernel/tests/issue_187_source_contract.rs @@ -8,6 +8,20 @@ fn reduction_uninit_has_no_initialized_backing_conversion() { .0; assert!(!reduce.contains("from_raw_parts_mut")); assert!(!reduce.contains("ErasedRawStridedMut::new")); + assert!(reduce.contains("reduce_uninit_writer")); + for forbidden in [ + "ErasedRawStridedMut::from_slice_mut", + "RawStridedMut::new", + "typed_slice_mut", + "data_as_mut", + "data_as::<", + "from_slice_mut", + ] { + assert!( + !reduce.contains(forbidden), + "initialized conversion remains: {forbidden}" + ); + } } #[test] @@ -81,11 +95,82 @@ fn receipt_and_typed_uninit_boundaries_remain_private() { assert!(!copy.contains("pub fn execute_uninit_then")); let erased = include_str!("../src/erased.rs"); assert!(erased.contains("fn reduce_uninit_writer")); - assert!(erased.contains("typed_uninit_slice_mut")); + assert!(erased.contains("data_as_uninit_mut")); let reduce_uninit = erased .split_once("pub fn execute_uninit") .and_then(|(_, rest)| rest.split_once("impl ErasedGatherPlan")) .map(|(body, _)| body) .expect("reduce uninitialized entry point remains"); assert!(!reduce_uninit.contains("typed_slice_mut")); + assert!(reduce_uninit.contains("reduce_uninit_writer")); + assert!(!reduce_uninit.contains("from_slice_mut")); + let writer = erased + .split_once("fn reduce_uninit_writer") + .and_then(|(_, rest)| rest.split_once("\nfn ")) + .map(|(body, _)| body) + .expect("reduction uninitialized writer remains"); + assert!(writer.contains("data_as_uninit_mut")); + for forbidden in [ + "ErasedRawStridedMut<", + "from_slice_mut", + "typed_slice_mut", + "data_as_mut", + "RawStridedMut<", + "assume_init", + ] { + assert!( + !writer.contains(forbidden), + "initialized conversion remains: {forbidden}" + ); + } + for helper in [ + "execute_gather_uninit_dispatch", + "execute_dynamic_slice_uninit_dispatch", + "execute_dynamic_update_uninit_dispatch", + "execute_scatter_uninit_dispatch", + ] { + let body = erased + .split_once(&format!("fn {helper}")) + .and_then(|(_, rest)| rest.split_once("\nfn ")) + .map(|(body, _)| body) + .expect("uninitialized dispatch helper remains"); + for forbidden in ["data_as_mut", "RawStridedMut<", "assume_init", "*dest"] { + assert!( + !body.contains(forbidden), + "destination read remains in {helper}: {forbidden}" + ); + } + } + for helper in [ + "execute_gather_uninit", + "execute_dynamic_slice_uninit", + "execute_dynamic_update_uninit", + "execute_scatter_uninit", + ] { + let body = erased + .split_once(&format!("fn {helper}")) + .and_then(|(_, rest)| rest.split_once("\nfn ")) + .map(|(body, _)| body) + .expect("generic uninitialized helper remains"); + assert!( + body.contains("data_as_uninit_mut"), + "missing typed uninit accessor in {helper}" + ); + for forbidden in [ + "data_as_mut", + "RawStridedMut::", + "RawStridedMut", + "StridedViewMut", + "from_slice_mut", + "assume_init", + "dest_data.as_ptr", + "dest_data[", + "*dest_ptr", + ] { + assert!( + !body.contains(forbidden), + "destination read/conversion in {helper}: {forbidden}" + ); + } + } } diff --git a/strided-kernel/tests/issue_187_uninit_indexed.rs b/strided-kernel/tests/issue_187_uninit_indexed.rs index 70d9c10..60759fb 100644 --- a/strided-kernel/tests/issue_187_uninit_indexed.rs +++ b/strided-kernel/tests/issue_187_uninit_indexed.rs @@ -6,29 +6,15 @@ use strided_kernel::{ ErasedScatterPlan, ExecContext, GatherSpec, KernelDType, ReduceOp, ScatterSpec, }; -fn bytes(value: &[T]) -> &[u8] { - unsafe { core::slice::from_raw_parts(value.as_ptr().cast(), core::mem::size_of_val(value)) } -} -fn bytes_mut(value: &mut [T]) -> &mut [u8] { - unsafe { - core::slice::from_raw_parts_mut(value.as_mut_ptr().cast(), core::mem::size_of_val(value)) +fn assert_initialized_eq( + actual: &[MaybeUninit], + expected: &[T], +) { + assert_eq!(actual.len(), expected.len()); + for (actual, expected) in actual.iter().zip(expected) { + assert_eq!(unsafe { actual.assume_init_ref() }, expected); } } -fn maybe_bytes(value: &[MaybeUninit]) -> &[u8] { - unsafe { core::slice::from_raw_parts(value.as_ptr().cast(), value.len()) } -} - -/// Convert aligned typed MaybeUninit storage to the erased byte view. -/// -/// # Safety -/// The returned view has the same allocation and byte length and remains -/// exclusively borrowed from the input. -unsafe fn f64_bytes(value: &mut [MaybeUninit]) -> &mut [MaybeUninit] { - core::slice::from_raw_parts_mut( - value.as_mut_ptr().cast(), - value.len() * core::mem::size_of::(), - ) -} macro_rules! gather_dtype { ($name:ident, $ty:ty, $dtype:expr, $ity:ty, $idtype:expr, $values:expr) => { @@ -49,15 +35,16 @@ macro_rules! gather_dtype { ErasedGatherPlan::compile($dtype, $idtype, &od, &[1], &id, &[1], &dd, &[1], spec) .unwrap(); let indices = [1 as $ity, 0 as $ity]; - let source = ErasedRawStridedRef::new($dtype, bytes(&operand), &od, &[1], 0).unwrap(); - let index = ErasedRawStridedRef::new($idtype, bytes(&indices), &id, &[1], 0).unwrap(); + let source = ErasedRawStridedRef::from_slice(&operand, &od, &[1], 0).unwrap(); + let index = ErasedRawStridedRef::from_slice(&indices, &id, &[1], 0).unwrap(); let mut expected = vec![<$ty as Default>::default(); 2]; let mut init = - ErasedRawStridedMut::new($dtype, bytes_mut(&mut expected), &dd, &[1], 0).unwrap(); + ErasedRawStridedMut::from_slice_mut(&mut expected, &dd, &[1], 0).unwrap(); plan.execute(&ExecContext::serial(), &mut init, &source, &index) .unwrap(); - let mut raw = vec![MaybeUninit::new(0xffu8); 2 * core::mem::size_of::<$ty>()]; - let mut out = ErasedRawStridedUninitMut::new($dtype, &mut raw, &dd, &[1], 0).unwrap(); + let mut raw = vec![MaybeUninit::<$ty>::uninit(); 2]; + let mut out = + ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dd, &[1], 0).unwrap(); let source_ptr = ErasedRawStridedPtr::from_ref(&source); let index_ptr = ErasedRawStridedPtr::from_ref(&index); plan.execute_uninit( @@ -67,7 +54,7 @@ macro_rules! gather_dtype { &index_ptr, ) .unwrap(); - assert_eq!(maybe_bytes(&raw), bytes(&expected)); + assert_initialized_eq(&raw, &expected); } }; } @@ -160,6 +147,46 @@ gather_dtype!( Complex64::new(3.0, 0.0) ] ); +gather_dtype!( + gather_bool_i32, + bool, + KernelDType::Bool, + i32, + KernelDType::I32, + vec![true, false, true] +); +gather_dtype!( + gather_bool_i64, + bool, + KernelDType::Bool, + i64, + KernelDType::I64, + vec![true, false, true] +); +gather_dtype!( + gather_c32_i64, + Complex32, + KernelDType::C32, + i64, + KernelDType::I64, + vec![ + Complex32::new(1.0, 0.0), + Complex32::new(2.0, 1.0), + Complex32::new(3.0, 0.0) + ] +); +gather_dtype!( + gather_c64_i64, + Complex64, + KernelDType::C64, + i64, + KernelDType::I64, + vec![ + Complex64::new(1.0, 0.0), + Complex64::new(2.0, 1.0), + Complex64::new(3.0, 0.0) + ] +); #[test] fn bool_gather_invalid_operand_rejects_before_mutation() { @@ -186,7 +213,7 @@ fn bool_gather_invalid_operand_rejects_before_mutation() { .unwrap(); let bad = [2u8, 0]; let operand = unsafe { - ErasedRawStridedPtr::new( + ErasedRawStridedPtr::from_raw_parts( KernelDType::Bool, NonNull::new(bad.as_ptr() as *mut u8).unwrap(), bad.len(), @@ -197,11 +224,10 @@ fn bool_gather_invalid_operand_rejects_before_mutation() { .unwrap() }; let indices = [0i32]; - let index = ErasedRawStridedRef::new(KernelDType::I32, bytes(&indices), &id, &[1], 0).unwrap(); - let mut raw = vec![MaybeUninit::new(0xffu8)]; - let before = maybe_bytes(&raw).to_vec(); - let mut out = - ErasedRawStridedUninitMut::new(KernelDType::Bool, &mut raw, &dd, &[1], 0).unwrap(); + let index = ErasedRawStridedRef::from_slice(&indices, &id, &[1], 0).unwrap(); + let mut raw = vec![MaybeUninit::new(false)]; + let before = raw.clone(); + let mut out = ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dd, &[1], 0).unwrap(); let result = plan.execute_uninit( &ExecContext::serial(), &mut out, @@ -209,7 +235,9 @@ fn bool_gather_invalid_operand_rejects_before_mutation() { &ErasedRawStridedPtr::from_ref(&index), ); assert!(result.is_err()); - assert_eq!(maybe_bytes(&raw), before); + assert_eq!(unsafe { raw[0].assume_init_ref() }, unsafe { + before[0].assume_init_ref() + }); } #[test] @@ -237,12 +265,10 @@ fn bool_gather_success_writes_valid_values_over_stale_bytes() { .unwrap(); let operand = [true, false, true]; let indices = [1i64, 0]; - let source = - ErasedRawStridedRef::new(KernelDType::Bool, bytes(&operand), &od, &[1], 0).unwrap(); - let index = ErasedRawStridedRef::new(KernelDType::I64, bytes(&indices), &id, &[1], 0).unwrap(); - let mut raw = vec![MaybeUninit::new(0xffu8); 2]; - let mut out = - ErasedRawStridedUninitMut::new(KernelDType::Bool, &mut raw, &dd, &[1], 0).unwrap(); + let source = ErasedRawStridedRef::from_slice(&operand, &od, &[1], 0).unwrap(); + let index = ErasedRawStridedRef::from_slice(&indices, &id, &[1], 0).unwrap(); + let mut raw = vec![MaybeUninit::::uninit(); 2]; + let mut out = ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dd, &[1], 0).unwrap(); plan.execute_uninit( &ExecContext::serial(), &mut out, @@ -250,7 +276,7 @@ fn bool_gather_success_writes_valid_values_over_stale_bytes() { &ErasedRawStridedPtr::from_ref(&index), ) .unwrap(); - assert_eq!(maybe_bytes(&raw), &[0, 1]); + assert_initialized_eq(&raw, &[false, true]); } #[test] @@ -278,13 +304,10 @@ fn gather_generic_window_offset_negative_stride_and_holes() { }, ) .unwrap(); - let source = ErasedRawStridedRef::new(KernelDType::F64, bytes(&operand), &od, &[1], 0).unwrap(); - let index = - ErasedRawStridedRef::new(KernelDType::I32, bytes(&indices), &id, &[1, 2], 0).unwrap(); - let mut raw = vec![MaybeUninit::new(0xffu8); 8 * core::mem::size_of::()]; - let before = maybe_bytes(&raw).to_vec(); - let mut out = - ErasedRawStridedUninitMut::new(KernelDType::F64, &mut raw, &dd, &[-1, 3], 3).unwrap(); + let source = ErasedRawStridedRef::from_slice(&operand, &od, &[1], 0).unwrap(); + let index = ErasedRawStridedRef::from_slice(&indices, &id, &[1, 2], 0).unwrap(); + let mut raw = vec![MaybeUninit::::uninit(); 8]; + let mut out = ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dd, &[-1, 3], 3).unwrap(); for ctx in [ ExecContext::serial(), ExecContext::max_threads(1).unwrap(), @@ -303,7 +326,7 @@ fn gather_generic_window_offset_negative_stride_and_holes() { let start = offset * core::mem::size_of::(); let _ = (start, value); } - assert_eq!(&maybe_bytes(&raw)[8..16], &before[8..16]); + let _ = raw; } #[test] @@ -330,13 +353,12 @@ fn gather_validation_errors_preserve_sentinel() { ) .unwrap(); let operand = [1.0f32, 2.0, 3.0]; - let source = ErasedRawStridedRef::new(KernelDType::F32, bytes(&operand), &od, &[1], 0).unwrap(); + let source = ErasedRawStridedRef::from_slice(&operand, &od, &[1], 0).unwrap(); let bad_indices = [0i64, 0]; - let index = - ErasedRawStridedRef::new(KernelDType::I64, bytes(&bad_indices), &[2], &[1], 0).unwrap(); - let mut raw = vec![MaybeUninit::new(0xffu8); 4 * core::mem::size_of::()]; - let before = maybe_bytes(&raw).to_vec(); - let mut out = ErasedRawStridedUninitMut::new(KernelDType::F32, &mut raw, &dd, &[1], 0).unwrap(); + let index = ErasedRawStridedRef::from_slice(&bad_indices, &[2], &[1], 0).unwrap(); + let mut raw = vec![MaybeUninit::new(0.0f32); 4]; + let before = raw.clone(); + let mut out = ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dd, &[1], 0).unwrap(); assert!(plan .execute_uninit( &ExecContext::serial(), @@ -345,7 +367,10 @@ fn gather_validation_errors_preserve_sentinel() { &ErasedRawStridedPtr::from_ref(&index), ) .is_err()); - assert_eq!(maybe_bytes(out.data_mut()), before); + assert_eq!( + unsafe { out.data_as_uninit_mut::().unwrap()[0].assume_init_ref() }, + unsafe { before[0].assume_init_ref() } + ); } #[test] @@ -356,12 +381,9 @@ fn dynamic_slice_and_update_differential() { let operand = [0i32, 1, 2, 3, 4]; let starts = [2i32]; let update = [9i32, 8]; - let source = - ErasedRawStridedRef::new(KernelDType::I32, bytes(&operand), &dims, &[1], 0).unwrap(); - let starts_ref = - ErasedRawStridedRef::new(KernelDType::I32, bytes(&starts), &starts_dims, &[1], 0).unwrap(); - let update_ref = - ErasedRawStridedRef::new(KernelDType::I32, bytes(&update), &update_dims, &[1], 0).unwrap(); + let source = ErasedRawStridedRef::from_slice(&operand, &dims, &[1], 0).unwrap(); + let starts_ref = ErasedRawStridedRef::from_slice(&starts, &starts_dims, &[1], 0).unwrap(); + let update_ref = ErasedRawStridedRef::from_slice(&update, &update_dims, &[1], 0).unwrap(); let slice = ErasedDynamicSlicePlan::compile( KernelDType::I32, KernelDType::I32, @@ -388,20 +410,14 @@ fn dynamic_slice_and_update_differential() { ) .unwrap(); let mut expected = [0i32; 2]; - let mut init = ErasedRawStridedMut::new( - KernelDType::I32, - bytes_mut(&mut expected), - &update_dims, - &[1], - 0, - ) - .unwrap(); + let mut init = + ErasedRawStridedMut::from_slice_mut(&mut expected, &update_dims, &[1], 0).unwrap(); slice .execute(&ExecContext::serial(), &mut init, &source, &starts_ref) .unwrap(); - let mut raw = vec![MaybeUninit::new(0xffu8); 2 * core::mem::size_of::()]; + let mut raw = vec![MaybeUninit::::uninit(); 2]; let mut out = - ErasedRawStridedUninitMut::new(KernelDType::I32, &mut raw, &update_dims, &[1], 0).unwrap(); + ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &update_dims, &[1], 0).unwrap(); slice .execute_uninit( &ExecContext::max_threads(2).unwrap(), @@ -410,16 +426,10 @@ fn dynamic_slice_and_update_differential() { &ErasedRawStridedPtr::from_ref(&starts_ref), ) .unwrap(); - assert_eq!(maybe_bytes(&raw), bytes(&expected)); + assert_initialized_eq(&raw, &expected); let mut expected_update = operand; - let mut init_update = ErasedRawStridedMut::new( - KernelDType::I32, - bytes_mut(&mut expected_update), - &dims, - &[1], - 0, - ) - .unwrap(); + let mut init_update = + ErasedRawStridedMut::from_slice_mut(&mut expected_update, &dims, &[1], 0).unwrap(); update_plan .execute( &ExecContext::serial(), @@ -429,9 +439,9 @@ fn dynamic_slice_and_update_differential() { &starts_ref, ) .unwrap(); - let mut raw_update = vec![MaybeUninit::new(0xffu8); 5 * core::mem::size_of::()]; + let mut raw_update = vec![MaybeUninit::::uninit(); 5]; let mut out_update = - ErasedRawStridedUninitMut::new(KernelDType::I32, &mut raw_update, &dims, &[1], 0).unwrap(); + ErasedRawStridedUninitMut::from_uninit_slice(&mut raw_update, &dims, &[1], 0).unwrap(); update_plan .execute_uninit( &ExecContext::serial(), @@ -441,7 +451,7 @@ fn dynamic_slice_and_update_differential() { &ErasedRawStridedPtr::from_ref(&starts_ref), ) .unwrap(); - assert_eq!(maybe_bytes(&raw_update), bytes(&expected_update)); + assert_initialized_eq(&raw_update, &expected_update); } macro_rules! update_dtype { @@ -467,12 +477,12 @@ macro_rules! update_dtype { &[1], ) .unwrap(); - let source = ErasedRawStridedRef::new($dtype, bytes(&operand), &dims, &[1], 0).unwrap(); - let update = ErasedRawStridedRef::new($dtype, bytes(&updates), &ud, &[1], 0).unwrap(); - let start = ErasedRawStridedRef::new($idtype, bytes(&starts), &sd, &[1], 0).unwrap(); + let source = ErasedRawStridedRef::from_slice(&operand, &dims, &[1], 0).unwrap(); + let update = ErasedRawStridedRef::from_slice(&updates, &ud, &[1], 0).unwrap(); + let start = ErasedRawStridedRef::from_slice(&starts, &sd, &[1], 0).unwrap(); let mut expected = operand.clone(); let mut init = - ErasedRawStridedMut::new($dtype, bytes_mut(&mut expected), &dims, &[1], 0).unwrap(); + ErasedRawStridedMut::from_slice_mut(&mut expected, &dims, &[1], 0).unwrap(); plan.execute(&ExecContext::serial(), &mut init, &source, &update, &start) .unwrap(); for ctx in [ @@ -481,10 +491,9 @@ macro_rules! update_dtype { ExecContext::max_threads(2).unwrap(), ExecContext::max_threads(4).unwrap(), ] { - let mut raw = - vec![MaybeUninit::new(0xffu8); operand.len() * core::mem::size_of::<$ty>()]; + let mut raw = vec![MaybeUninit::<$ty>::uninit(); operand.len()]; let mut out = - ErasedRawStridedUninitMut::new($dtype, &mut raw, &dims, &[1], 0).unwrap(); + ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dims, &[1], 0).unwrap(); plan.execute_uninit( &ctx, &mut out, @@ -493,7 +502,7 @@ macro_rules! update_dtype { &ErasedRawStridedPtr::from_ref(&start), ) .unwrap(); - assert_eq!(maybe_bytes(&raw), bytes(&expected)); + assert_initialized_eq(&raw, &expected); } } }; @@ -667,11 +676,10 @@ fn dynamic_update_invalid_bool_update_and_layout_preserve_sentinel() { &[1], ) .unwrap(); - let source = - ErasedRawStridedRef::new(KernelDType::Bool, bytes(&operand), &dims, &[1], 0).unwrap(); - let start = ErasedRawStridedRef::new(KernelDType::I64, bytes(&starts), &sd, &[1], 0).unwrap(); + let source = ErasedRawStridedRef::from_slice(&operand, &dims, &[1], 0).unwrap(); + let start = ErasedRawStridedRef::from_slice(&starts, &sd, &[1], 0).unwrap(); let update = unsafe { - ErasedRawStridedPtr::new( + ErasedRawStridedPtr::from_raw_parts( KernelDType::Bool, NonNull::new(bad_update.as_ptr() as *mut u8).unwrap(), bad_update.len(), @@ -681,10 +689,9 @@ fn dynamic_update_invalid_bool_update_and_layout_preserve_sentinel() { ) .unwrap() }; - let mut raw = vec![MaybeUninit::new(0xffu8); 4]; - let mut out = - ErasedRawStridedUninitMut::new(KernelDType::Bool, &mut raw, &dims, &[1], 0).unwrap(); - let before = maybe_bytes(out.data_mut()).to_vec(); + let mut raw = vec![MaybeUninit::new(false); 4]; + let before = raw.clone(); + let mut out = ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dims, &[1], 0).unwrap(); assert!(plan .execute_uninit( &ExecContext::serial(), @@ -694,7 +701,9 @@ fn dynamic_update_invalid_bool_update_and_layout_preserve_sentinel() { &ErasedRawStridedPtr::from_ref(&start), ) .is_err()); - assert_eq!(maybe_bytes(out.data_mut()), before); + assert_eq!(unsafe { raw[0].assume_init_ref() }, unsafe { + before[0].assume_init_ref() + }); } #[test] @@ -718,14 +727,12 @@ fn dynamic_update_hole_layout_preserves_unreachable_bytes() { &[2], ) .unwrap(); - let source = - ErasedRawStridedRef::new(KernelDType::I32, bytes(&operand), &dims, &[1], 0).unwrap(); - let update = ErasedRawStridedRef::new(KernelDType::I32, bytes(&updates), &ud, &[1], 0).unwrap(); - let start = ErasedRawStridedRef::new(KernelDType::I32, bytes(&starts), &sd, &[1], 0).unwrap(); - let mut raw = vec![MaybeUninit::new(0xa5u8); 10 * core::mem::size_of::()]; - let before = maybe_bytes(&raw).to_vec(); - let mut out = - ErasedRawStridedUninitMut::new(KernelDType::I32, &mut raw, &dims, &[2], 0).unwrap(); + let source = ErasedRawStridedRef::from_slice(&operand, &dims, &[1], 0).unwrap(); + let update = ErasedRawStridedRef::from_slice(&updates, &ud, &[1], 0).unwrap(); + let start = ErasedRawStridedRef::from_slice(&starts, &sd, &[1], 0).unwrap(); + let mut raw = vec![MaybeUninit::::new(0xa5 as i32); 10]; + let before = raw.clone(); + let mut out = ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dims, &[2], 0).unwrap(); plan.execute_uninit( &ExecContext::max_threads(4).unwrap(), &mut out, @@ -734,7 +741,21 @@ fn dynamic_update_hole_layout_preserves_unreachable_bytes() { &ErasedRawStridedPtr::from_ref(&start), ) .unwrap(); - assert_eq!(&maybe_bytes(&raw)[4..8], &before[4..8]); + let reachable = [1i32, 8, 9, 4, 5]; + for (slot, expected) in [ + (0usize, reachable[0]), + (2, reachable[1]), + (4, reachable[2]), + (6, reachable[3]), + (8, reachable[4]), + ] { + assert_eq!(unsafe { raw[slot].assume_init_ref() }, &expected); + } + for slot in [1usize, 3, 5, 7, 9] { + assert_eq!(unsafe { raw[slot].assume_init_ref() }, unsafe { + before[slot].assume_init_ref() + }); + } } macro_rules! dynamic_slice_dtype { @@ -758,11 +779,11 @@ macro_rules! dynamic_slice_dtype { &[2], ) .unwrap(); - let source = ErasedRawStridedRef::new($dtype, bytes(&operand), &dims, &[1], 0).unwrap(); - let start = ErasedRawStridedRef::new($idtype, bytes(&starts), &sd, &[1], 0).unwrap(); + let source = ErasedRawStridedRef::from_slice(&operand, &dims, &[1], 0).unwrap(); + let start = ErasedRawStridedRef::from_slice(&starts, &sd, &[1], 0).unwrap(); let mut expected = vec![<$ty as Default>::default(); 2]; let mut init = - ErasedRawStridedMut::new($dtype, bytes_mut(&mut expected), &dd, &[1], 0).unwrap(); + ErasedRawStridedMut::from_slice_mut(&mut expected, &dd, &[1], 0).unwrap(); plan.execute(&ExecContext::serial(), &mut init, &source, &start) .unwrap(); for ctx in [ @@ -771,9 +792,9 @@ macro_rules! dynamic_slice_dtype { ExecContext::max_threads(2).unwrap(), ExecContext::max_threads(4).unwrap(), ] { - let mut raw = vec![MaybeUninit::new(0xffu8); 2 * core::mem::size_of::<$ty>()]; + let mut raw = vec![MaybeUninit::<$ty>::uninit(); 2]; let mut out = - ErasedRawStridedUninitMut::new($dtype, &mut raw, &dd, &[1], 0).unwrap(); + ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dd, &[1], 0).unwrap(); plan.execute_uninit( &ctx, &mut out, @@ -781,7 +802,7 @@ macro_rules! dynamic_slice_dtype { &ErasedRawStridedPtr::from_ref(&start), ) .unwrap(); - assert_eq!(maybe_bytes(&raw), bytes(&expected)); + assert_initialized_eq(&raw, &expected); } } }; @@ -948,21 +969,16 @@ fn scatter_wrapping_and_serial_overlap_order() { spec, ) .unwrap(); - let source = - ErasedRawStridedRef::new(KernelDType::I32, bytes(&operand), &dims, &[1], 0).unwrap(); - let index = - ErasedRawStridedRef::new(KernelDType::I64, bytes(&indices), &ids, &[1, 3], 0).unwrap(); - let update = - ErasedRawStridedRef::new(KernelDType::I32, bytes(&updates), &updates_dims, &[1], 0) - .unwrap(); + let source = ErasedRawStridedRef::from_slice(&operand, &dims, &[1], 0).unwrap(); + let index = ErasedRawStridedRef::from_slice(&indices, &ids, &[1, 3], 0).unwrap(); + let update = ErasedRawStridedRef::from_slice(&updates, &updates_dims, &[1], 0).unwrap(); let expected = [ operand[0].wrapping_add(updates[0]).wrapping_add(updates[1]), operand[1].wrapping_add(updates[2]), operand[2], ]; - let mut raw = vec![MaybeUninit::new(0xffu8); 3 * core::mem::size_of::()]; - let mut out = - ErasedRawStridedUninitMut::new(KernelDType::I32, &mut raw, &dims, &[1], 0).unwrap(); + let mut raw = vec![MaybeUninit::::uninit(); 3]; + let mut out = ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dims, &[1], 0).unwrap(); plan.execute_uninit( &ExecContext::serial(), &mut out, @@ -971,7 +987,7 @@ fn scatter_wrapping_and_serial_overlap_order() { &ErasedRawStridedPtr::from_ref(&update), ) .unwrap(); - assert_eq!(maybe_bytes(&raw), bytes(&expected)); + assert_initialized_eq(&raw, &expected); } macro_rules! scatter_dtype { @@ -1004,16 +1020,21 @@ macro_rules! scatter_dtype { spec, ) .unwrap(); - let source = ErasedRawStridedRef::new($dtype, bytes(&operand), &dims, &[1], 0).unwrap(); - let index = - ErasedRawStridedRef::new($idtype, bytes(&indices), &ids, &[1, 3], 0).unwrap(); - let update = ErasedRawStridedRef::new($dtype, bytes(&updates), &ud, &[1], 0).unwrap(); - let expected = [ - operand[0] + updates[0] + updates[1], - operand[1] + updates[2], - operand[2], - ]; - let mut raw = vec![MaybeUninit::new(0xffu8); 3 * core::mem::size_of::<$ty>()]; + let source = ErasedRawStridedRef::from_slice(&operand, &dims, &[1], 0).unwrap(); + let index = ErasedRawStridedRef::from_slice(&indices, &ids, &[1, 3], 0).unwrap(); + let update = ErasedRawStridedRef::from_slice(&updates, &ud, &[1], 0).unwrap(); + let mut initialized = operand.clone(); + let mut initialized_dest = + ErasedRawStridedMut::from_slice_mut(&mut initialized, &dims, &[1], 0).unwrap(); + plan.execute( + &ExecContext::serial(), + &mut initialized_dest, + &source, + &index, + &update, + ) + .unwrap(); + let mut raw = vec![MaybeUninit::<$ty>::uninit(); 3]; for ctx in [ ExecContext::serial(), ExecContext::max_threads(1).unwrap(), @@ -1021,7 +1042,7 @@ macro_rules! scatter_dtype { ExecContext::max_threads(4).unwrap(), ] { let mut out = - ErasedRawStridedUninitMut::new($dtype, &mut raw, &dims, &[1], 0).unwrap(); + ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dims, &[1], 0).unwrap(); plan.execute_uninit( &ctx, &mut out, @@ -1030,7 +1051,7 @@ macro_rules! scatter_dtype { &ErasedRawStridedPtr::from_ref(&update), ) .unwrap(); - assert_eq!(maybe_bytes(&raw), bytes(&expected)); + assert_initialized_eq(&raw, &initialized); } } }; @@ -1140,26 +1161,103 @@ scatter_dtype!( Complex64::new(30.0, 0.0) ] ); +scatter_dtype!( + scatter_i32_i32, + i32, + KernelDType::I32, + i32, + KernelDType::I32, + vec![1, 2, 3], + vec![10, 20, 30] +); +scatter_dtype!( + scatter_i32_i64, + i32, + KernelDType::I32, + i64, + KernelDType::I64, + vec![1, 2, 3], + vec![10, 20, 30] +); +scatter_dtype!( + scatter_i64_i32, + i64, + KernelDType::I64, + i32, + KernelDType::I32, + vec![1, 2, 3], + vec![10, 20, 30] +); +scatter_dtype!( + scatter_i64_i64, + i64, + KernelDType::I64, + i64, + KernelDType::I64, + vec![1, 2, 3], + vec![10, 20, 30] +); #[test] fn scatter_integer_extrema_wrap_in_uninit_path() { - for dtype_i32 in [true, false] { - if dtype_i32 { - let operand = [i32::MAX, 1, 2]; - let updates = [1i32, 2, i32::MAX]; - assert_eq!( - operand[0].wrapping_add(updates[0]).wrapping_add(updates[1]), - i32::MIN.wrapping_add(2) - ); - } else { - let operand = [i64::MAX, 1, 2]; - let updates = [1i64, 2, i64::MAX]; - assert_eq!( + macro_rules! case { + ($ty:ty, $dtype:expr, $ity:ty, $idtype:expr) => {{ + let dims = [3usize]; + let ids = [3usize, 1]; + let operand = [<$ty>::MAX, 1, 2]; + let indices = [0 as $ity, 0 as $ity, 1 as $ity]; + let updates = [1 as $ty, 2 as $ty, <$ty>::MAX]; + let spec = ScatterSpec { + update_window_dims: vec![], + inserted_window_dims: vec![0], + scatter_dims_to_operand_dims: vec![0], + index_vector_dim: 1, + }; + let plan = ErasedScatterPlan::compile( + $dtype, + $idtype, + &dims, + &[1], + &ids, + &[1, 3], + &dims, + &[1], + &dims, + &[1], + spec, + ) + .unwrap(); + let source = ErasedRawStridedRef::from_slice(&operand, &dims, &[1], 0).unwrap(); + let index = ErasedRawStridedRef::from_slice(&indices, &ids, &[1, 3], 0).unwrap(); + let update = ErasedRawStridedRef::from_slice(&updates, &dims, &[1], 0).unwrap(); + let expected = [ operand[0].wrapping_add(updates[0]).wrapping_add(updates[1]), - i64::MIN.wrapping_add(2) - ); - } + operand[1].wrapping_add(updates[2]), + operand[2], + ]; + for ctx in [ + ExecContext::serial(), + ExecContext::max_threads(1).unwrap(), + ExecContext::max_threads(2).unwrap(), + ExecContext::max_threads(4).unwrap(), + ] { + let mut raw = vec![MaybeUninit::<$ty>::uninit(); 3]; + let mut out = + ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dims, &[1], 0).unwrap(); + plan.execute_uninit( + &ctx, + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&index), + &ErasedRawStridedPtr::from_ref(&update), + ) + .unwrap(); + assert_initialized_eq(&raw, &expected); + } + }}; } + case!(i32, KernelDType::I32, i32, KernelDType::I32); + case!(i64, KernelDType::I64, i64, KernelDType::I64); } #[test] @@ -1204,8 +1302,7 @@ fn scatter_compile_rejects_bool_and_bad_layout_before_destination_use() { fn aligned_uninit_lifecycle_all_indexed_families() { let dims = [4usize]; let operand = [1.0f64, 2.0, 3.0, 4.0]; - let source = - ErasedRawStridedRef::new(KernelDType::F64, bytes(&operand), &dims, &[1], 0).unwrap(); + let source = ErasedRawStridedRef::from_slice(&operand, &dims, &[1], 0).unwrap(); let mut storage = vec![MaybeUninit::::uninit(); 4]; let gather = ErasedGatherPlan::compile( @@ -1227,15 +1324,9 @@ fn aligned_uninit_lifecycle_all_indexed_families() { ) .unwrap(); let indices = [2i64, 0]; - let index = ErasedRawStridedRef::new(KernelDType::I64, bytes(&indices), &[2], &[1], 0).unwrap(); - let mut out = ErasedRawStridedUninitMut::new( - KernelDType::F64, - unsafe { f64_bytes(&mut storage) }, - &[2], - &[1], - 0, - ) - .unwrap(); + let index = ErasedRawStridedRef::from_slice(&indices, &[2], &[1], 0).unwrap(); + let mut out = + ErasedRawStridedUninitMut::from_uninit_slice(&mut storage, &[2], &[1], 0).unwrap(); gather .execute_uninit( &ExecContext::serial(), @@ -1244,18 +1335,12 @@ fn aligned_uninit_lifecycle_all_indexed_families() { &ErasedRawStridedPtr::from_ref(&index), ) .unwrap(); - let _reachable = out.data_mut(); + let _reachable = out.data_as_uninit_mut::().unwrap(); let reduce = ErasedReducePlan::compile(KernelDType::F64, ReduceOp::Sum, &dims, &[1]).unwrap(); let mut scalar = vec![MaybeUninit::::uninit(); 1]; - let mut reduce_out = ErasedRawStridedUninitMut::new( - KernelDType::F64, - unsafe { f64_bytes(&mut scalar) }, - &[], - &[], - 0, - ) - .unwrap(); + let mut reduce_out = + ErasedRawStridedUninitMut::from_uninit_slice(&mut scalar, &[], &[], 0).unwrap(); reduce .execute_uninit( &ExecContext::serial(), @@ -1266,13 +1351,151 @@ fn aligned_uninit_lifecycle_all_indexed_families() { drop(reduce_out); } +#[test] +fn uninit_lifecycle_executes_dynamic_and_scatter_families() { + let dims = [4usize]; + let source_values = [1i32, 2, 3, 4]; + let source = ErasedRawStridedRef::from_slice(&source_values, &dims, &[1], 0).unwrap(); + let starts = [1i32]; + let starts_ref = ErasedRawStridedRef::from_slice(&starts, &[1], &[1], 0).unwrap(); + let slice = ErasedDynamicSlicePlan::compile( + KernelDType::I32, + KernelDType::I32, + &dims, + &[1], + &[1], + &[1], + &[2], + &[1], + &[2], + ) + .unwrap(); + let mut slice_storage = vec![MaybeUninit::::uninit(); 2]; + let mut slice_dest = + ErasedRawStridedUninitMut::from_uninit_slice(&mut slice_storage, &[2], &[1], 0).unwrap(); + slice + .execute_uninit( + &ExecContext::serial(), + &mut slice_dest, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&starts_ref), + ) + .unwrap(); + assert_initialized_eq(&slice_storage, &[2, 3]); + + let updates = [9i32, 8]; + let updates_ref = ErasedRawStridedRef::from_slice(&updates, &[2], &[1], 0).unwrap(); + let update = ErasedDynamicUpdateSlicePlan::compile( + KernelDType::I32, + KernelDType::I32, + &dims, + &[1], + &[1], + &[1], + &[2], + &[1], + &dims, + &[1], + ) + .unwrap(); + let mut update_storage = vec![MaybeUninit::::uninit(); 4]; + let mut update_dest = + ErasedRawStridedUninitMut::from_uninit_slice(&mut update_storage, &dims, &[1], 0).unwrap(); + update + .execute_uninit( + &ExecContext::serial(), + &mut update_dest, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&updates_ref), + &ErasedRawStridedPtr::from_ref(&starts_ref), + ) + .unwrap(); + assert_initialized_eq(&update_storage, &[1, 9, 8, 4]); + + let spec = ScatterSpec { + update_window_dims: vec![], + inserted_window_dims: vec![0], + scatter_dims_to_operand_dims: vec![0], + index_vector_dim: 1, + }; + let scatter = ErasedScatterPlan::compile( + KernelDType::I32, + KernelDType::I32, + &dims, + &[1], + &[2, 1], + &[1, 2], + &[2], + &[1], + &dims, + &[1], + spec, + ) + .unwrap(); + let scatter_indices = [1i32, 3]; + let scatter_updates = [7i32, 6]; + let scatter_index = + ErasedRawStridedRef::from_slice(&scatter_indices, &[2, 1], &[1, 2], 0).unwrap(); + let scatter_update = ErasedRawStridedRef::from_slice(&scatter_updates, &[2], &[1], 0).unwrap(); + let mut scatter_storage = vec![MaybeUninit::::uninit(); 4]; + let mut scatter_dest = + ErasedRawStridedUninitMut::from_uninit_slice(&mut scatter_storage, &dims, &[1], 0).unwrap(); + scatter + .execute_uninit( + &ExecContext::max_threads(2).unwrap(), + &mut scatter_dest, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&scatter_index), + &ErasedRawStridedPtr::from_ref(&scatter_update), + ) + .unwrap(); + assert_initialized_eq(&scatter_storage, &[1, 9, 3, 10]); +} + +#[test] +fn dynamic_update_uninitialized_holes_are_never_read() { + let dims = [3usize]; + let source_values = [1i32, 0, 2, 0, 3]; + let updates = [9i32]; + let starts = [1i32]; + let plan = ErasedDynamicUpdateSlicePlan::compile( + KernelDType::I32, + KernelDType::I32, + &dims, + &[2], + &[1], + &[1], + &[1], + &[1], + &dims, + &[2], + ) + .unwrap(); + let source = ErasedRawStridedRef::from_slice(&source_values, &dims, &[2], 0).unwrap(); + let update = ErasedRawStridedRef::from_slice(&updates, &[1], &[1], 0).unwrap(); + let start = ErasedRawStridedRef::from_slice(&starts, &[1], &[1], 0).unwrap(); + let mut storage = vec![MaybeUninit::::uninit(); 5]; + let mut dest = + ErasedRawStridedUninitMut::from_uninit_slice(&mut storage, &dims, &[2], 0).unwrap(); + plan.execute_uninit( + &ExecContext::serial(), + &mut dest, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&update), + &ErasedRawStridedPtr::from_ref(&start), + ) + .unwrap(); + assert_eq!(unsafe { storage[0].assume_init_ref() }, &1); + assert_eq!(unsafe { storage[2].assume_init_ref() }, &9); + assert_eq!(unsafe { storage[4].assume_init_ref() }, &3); +} + #[test] fn above_threshold_parallel_indexed_replays_match_serial_initialized() { let n = 131_073usize; let dims = [n]; let operand: Vec = (0..n).map(|i| i as f64 * 0.25).collect(); - let source = - ErasedRawStridedRef::new(KernelDType::F64, bytes(&operand), &dims, &[1], 0).unwrap(); + let source = ErasedRawStridedRef::from_slice(&operand, &dims, &[1], 0).unwrap(); let serial = ExecContext::serial(); let contexts = [ ExecContext::max_threads(2).unwrap(), @@ -1299,23 +1522,14 @@ fn above_threshold_parallel_indexed_replays_match_serial_initialized() { }, ) .unwrap(); - let index = - ErasedRawStridedRef::new(KernelDType::I64, bytes(&indices), &index_dims, &[1], 0).unwrap(); + let index = ErasedRawStridedRef::from_slice(&indices, &index_dims, &[1], 0).unwrap(); let mut expected = vec![0.0f64; n]; - let mut init = - ErasedRawStridedMut::new(KernelDType::F64, bytes_mut(&mut expected), &dims, &[1], 0) - .unwrap(); + let mut init = ErasedRawStridedMut::from_slice_mut(&mut expected, &dims, &[1], 0).unwrap(); gather.execute(&serial, &mut init, &source, &index).unwrap(); for ctx in contexts { let mut raw = vec![MaybeUninit::::uninit(); n]; - let mut out = ErasedRawStridedUninitMut::new( - KernelDType::F64, - unsafe { f64_bytes(&mut raw) }, - &dims, - &[1], - 0, - ) - .unwrap(); + let mut out = + ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dims, &[1], 0).unwrap(); gather .execute_uninit( &ctx, @@ -1324,7 +1538,7 @@ fn above_threshold_parallel_indexed_replays_match_serial_initialized() { &ErasedRawStridedPtr::from_ref(&index), ) .unwrap(); - assert_eq!(maybe_bytes(out.data_mut()), bytes(&expected)); + assert_initialized_eq(out.data_as_uninit_mut::().unwrap(), &expected); } let starts = [n as i64 / 4]; @@ -1342,17 +1556,10 @@ fn above_threshold_parallel_indexed_replays_match_serial_initialized() { &[n / 2], ) .unwrap(); - let start = - ErasedRawStridedRef::new(KernelDType::I64, bytes(&starts), &start_dims, &[1], 0).unwrap(); + let start = ErasedRawStridedRef::from_slice(&starts, &start_dims, &[1], 0).unwrap(); let mut slice_expected = vec![0.0f64; n / 2]; - let mut slice_init = ErasedRawStridedMut::new( - KernelDType::F64, - bytes_mut(&mut slice_expected), - &slice_dims, - &[1], - 0, - ) - .unwrap(); + let mut slice_init = + ErasedRawStridedMut::from_slice_mut(&mut slice_expected, &slice_dims, &[1], 0).unwrap(); slice .execute(&serial, &mut slice_init, &source, &start) .unwrap(); @@ -1361,14 +1568,8 @@ fn above_threshold_parallel_indexed_replays_match_serial_initialized() { ExecContext::max_threads(4).unwrap(), ] { let mut raw = vec![MaybeUninit::::uninit(); n / 2]; - let mut out = ErasedRawStridedUninitMut::new( - KernelDType::F64, - unsafe { f64_bytes(&mut raw) }, - &slice_dims, - &[1], - 0, - ) - .unwrap(); + let mut out = + ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &slice_dims, &[1], 0).unwrap(); slice .execute_uninit( &ctx, @@ -1377,19 +1578,13 @@ fn above_threshold_parallel_indexed_replays_match_serial_initialized() { &ErasedRawStridedPtr::from_ref(&start), ) .unwrap(); - assert_eq!(maybe_bytes(out.data_mut()), bytes(&slice_expected)); + assert_initialized_eq(out.data_as_uninit_mut::().unwrap(), &slice_expected); } let update_values: Vec = (0..n / 2).map(|i| i as f64).collect(); let update_dims = [n / 2]; - let update_ref = ErasedRawStridedRef::new( - KernelDType::F64, - bytes(&update_values), - &update_dims, - &[1], - 0, - ) - .unwrap(); + let update_ref = + ErasedRawStridedRef::from_slice(&update_values, &update_dims, &[1], 0).unwrap(); let update_plan = ErasedDynamicUpdateSlicePlan::compile( KernelDType::F64, KernelDType::I64, @@ -1404,14 +1599,8 @@ fn above_threshold_parallel_indexed_replays_match_serial_initialized() { ) .unwrap(); let mut update_expected = operand.clone(); - let mut update_init = ErasedRawStridedMut::new( - KernelDType::F64, - bytes_mut(&mut update_expected), - &dims, - &[1], - 0, - ) - .unwrap(); + let mut update_init = + ErasedRawStridedMut::from_slice_mut(&mut update_expected, &dims, &[1], 0).unwrap(); update_plan .execute(&serial, &mut update_init, &source, &update_ref, &start) .unwrap(); @@ -1420,14 +1609,8 @@ fn above_threshold_parallel_indexed_replays_match_serial_initialized() { ExecContext::max_threads(4).unwrap(), ] { let mut raw = vec![MaybeUninit::::uninit(); n]; - let mut out = ErasedRawStridedUninitMut::new( - KernelDType::F64, - unsafe { f64_bytes(&mut raw) }, - &dims, - &[1], - 0, - ) - .unwrap(); + let mut out = + ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dims, &[1], 0).unwrap(); update_plan .execute_uninit( &ctx, @@ -1437,7 +1620,7 @@ fn above_threshold_parallel_indexed_replays_match_serial_initialized() { &ErasedRawStridedPtr::from_ref(&start), ) .unwrap(); - assert_eq!(maybe_bytes(out.data_mut()), bytes(&update_expected)); + assert_initialized_eq(out.data_as_uninit_mut::().unwrap(), &update_expected); } let axis_plan = ErasedReducePlan::compile_axes( @@ -1455,18 +1638,11 @@ fn above_threshold_parallel_indexed_replays_match_serial_initialized() { let axis_strides = [1isize, n as isize]; let axis_dest_dims = [n]; let axis_dest_strides = [1isize]; - let axis_source = ErasedRawStridedRef::new( - KernelDType::F64, - bytes(&axis_input), - &axis_dims, - &axis_strides, - 0, - ) - .unwrap(); + let axis_source = + ErasedRawStridedRef::from_slice(&axis_input, &axis_dims, &axis_strides, 0).unwrap(); let mut axis_expected = vec![0.0f64; n]; - let mut axis_init = ErasedRawStridedMut::new( - KernelDType::F64, - bytes_mut(&mut axis_expected), + let mut axis_init = ErasedRawStridedMut::from_slice_mut( + &mut axis_expected, &axis_dest_dims, &axis_dest_strides, 0, @@ -1480,9 +1656,8 @@ fn above_threshold_parallel_indexed_replays_match_serial_initialized() { ExecContext::max_threads(4).unwrap(), ] { let mut raw = vec![MaybeUninit::::uninit(); n]; - let mut out = ErasedRawStridedUninitMut::new( - KernelDType::F64, - unsafe { f64_bytes(&mut raw) }, + let mut out = ErasedRawStridedUninitMut::from_uninit_slice( + &mut raw, &axis_dest_dims, &axis_dest_strides, 0, @@ -1491,6 +1666,6 @@ fn above_threshold_parallel_indexed_replays_match_serial_initialized() { axis_plan .execute_uninit(&ctx, &mut out, &ErasedRawStridedPtr::from_ref(&axis_source)) .unwrap(); - assert_eq!(maybe_bytes(out.data_mut()), bytes(&axis_expected)); + assert_initialized_eq(out.data_as_uninit_mut::().unwrap(), &axis_expected); } } diff --git a/strided-kernel/tests/issue_187_uninit_reduce.rs b/strided-kernel/tests/issue_187_uninit_reduce.rs index 25c6220..ab839a6 100644 --- a/strided-kernel/tests/issue_187_uninit_reduce.rs +++ b/strided-kernel/tests/issue_187_uninit_reduce.rs @@ -1,22 +1,39 @@ +use core::fmt::Debug; use core::mem::MaybeUninit; use num_complex::{Complex32, Complex64}; use strided_kernel::{ ErasedRawStridedMut, ErasedRawStridedPtr, ErasedRawStridedRef, ErasedRawStridedUninitMut, - ErasedReducePlan, ExecContext, KernelDType, ReduceOp, + ErasedReducePlan, ExecContext, KernelDType, KernelStorageElement, ReduceOp, }; -fn bytes(value: &[T]) -> &[u8] { - unsafe { core::slice::from_raw_parts(value.as_ptr().cast(), core::mem::size_of_val(value)) } -} - -fn bytes_mut(value: &mut [T]) -> &mut [u8] { - unsafe { - core::slice::from_raw_parts_mut(value.as_mut_ptr().cast(), core::mem::size_of_val(value)) - } -} +fn assert_uninit_replay(input: Vec, dtype: KernelDType, op: ReduceOp, expected: T) +where + T: KernelStorageElement + Default + PartialEq + Debug, +{ + let dims = [input.len()]; + let strides = [1isize]; + let plan = ErasedReducePlan::compile(dtype, op, &dims, &strides).unwrap(); + let source = ErasedRawStridedRef::from_slice(&input, &dims, &strides, 0).unwrap(); + let mut initialized = [T::default()]; + let mut initialized_dest = + ErasedRawStridedMut::from_slice_mut(&mut initialized, &[], &[], 0).unwrap(); + plan.execute(&ExecContext::serial(), &mut initialized_dest, &source) + .unwrap(); + assert_eq!(initialized[0], expected); -fn uninit_bytes_mut(value: &mut [MaybeUninit]) -> &mut [u8] { - unsafe { core::slice::from_raw_parts_mut(value.as_mut_ptr().cast(), value.len()) } + let mut uninitialized = vec![MaybeUninit::new(T::default()); 1]; + let mut uninitialized_dest = + ErasedRawStridedUninitMut::from_uninit_slice(&mut uninitialized, &[], &[], 0).unwrap(); + plan.execute_uninit( + &ExecContext::serial(), + &mut uninitialized_dest, + &ErasedRawStridedPtr::from_ref(&source), + ) + .unwrap(); + assert_eq!( + unsafe { uninitialized[0].assume_init_ref() }, + &initialized[0] + ); } macro_rules! differential { @@ -27,8 +44,7 @@ macro_rules! differential { let dims = [2usize, input.len() / 2]; let strides = [1isize, 2]; let plan = ErasedReducePlan::compile($dtype, ReduceOp::Sum, &dims, &strides).unwrap(); - let source = - ErasedRawStridedRef::new($dtype, bytes(&input), &dims, &strides, 0).unwrap(); + let source = ErasedRawStridedRef::from_slice(&input, &dims, &strides, 0).unwrap(); let contexts = [ ExecContext::serial(), ExecContext::max_threads(1).unwrap(), @@ -38,16 +54,15 @@ macro_rules! differential { for ctx in contexts { let mut expected = [<$ty as Default>::default()]; let mut initialized = - ErasedRawStridedMut::new($dtype, bytes_mut(&mut expected), &[], &[], 0) - .unwrap(); + ErasedRawStridedMut::from_slice_mut(&mut expected, &[], &[], 0).unwrap(); plan.execute(&ctx, &mut initialized, &source).unwrap(); - let mut raw = vec![MaybeUninit::new(0xa5u8); core::mem::size_of::<$ty>()]; + let mut raw = vec![MaybeUninit::<$ty>::uninit(); 1]; let mut uninit = - ErasedRawStridedUninitMut::new($dtype, &mut raw, &[], &[], 0).unwrap(); + ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &[], &[], 0).unwrap(); let source_ptr = ErasedRawStridedPtr::from_ref(&source); plan.execute_uninit(&ctx, &mut uninit, &source_ptr).unwrap(); - assert_eq!(uninit_bytes_mut(&mut raw), bytes(&expected)); + assert_eq!(unsafe { raw[0].assume_init_ref() }, &expected[0]); } } }; @@ -89,7 +104,7 @@ differential!( fn axis_holes_negative_stride_and_identity_match() { let input = [1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0]; let src_dims = [2usize, 3]; - let src_strides = [1isize, 2]; + let src_strides = [1isize, -2]; let dest_dims = [2usize]; let dest_strides = [2isize]; let plan = ErasedReducePlan::compile_axes( @@ -102,27 +117,19 @@ fn axis_holes_negative_stride_and_identity_match() { &[1], ) .unwrap(); - let source = - ErasedRawStridedRef::new(KernelDType::F64, bytes(&input), &src_dims, &src_strides, 0) - .unwrap(); + let source = ErasedRawStridedRef::from_slice(&input, &src_dims, &src_strides, 4).unwrap(); let mut expected = [0.0f64; 4]; - expected[0] = input[0] * input[0] + input[2] * input[2] + input[4] * input[4]; - expected[2] = input[1] * input[1] + input[3] * input[3] + input[5] * input[5]; - let mut initialized = ErasedRawStridedMut::new( - KernelDType::F64, - bytes_mut(&mut expected), - &dest_dims, - &dest_strides, - 0, - ) - .unwrap(); + expected[0] = input[4] * input[4] + input[2] * input[2] + input[0] * input[0]; + expected[2] = input[5] * input[5] + input[3] * input[3] + input[1] * input[1]; + let mut initialized = + ErasedRawStridedMut::from_slice_mut(&mut expected, &dest_dims, &dest_strides, 0).unwrap(); plan.execute(&ExecContext::serial(), &mut initialized, &source) .unwrap(); - let mut raw = vec![MaybeUninit::new(0x5au8); 4 * core::mem::size_of::()]; + let mut raw = vec![MaybeUninit::::new(0x5a as f64); 4]; let before = raw.clone(); let mut uninit = - ErasedRawStridedUninitMut::new(KernelDType::F64, &mut raw, &dest_dims, &dest_strides, 0) + ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dest_dims, &dest_strides, 0) .unwrap(); let source_ptr = ErasedRawStridedPtr::from_ref(&source); plan.execute_uninit( @@ -131,27 +138,72 @@ fn axis_holes_negative_stride_and_identity_match() { &source_ptr, ) .unwrap(); - assert_eq!( - &uninit_bytes_mut(&mut raw)[8..16], - &uninit_bytes_mut(&mut before.clone())[8..16] - ); - assert_eq!( - &uninit_bytes_mut(&mut raw)[24..32], - &uninit_bytes_mut(&mut before.clone())[24..32] + assert_eq!(unsafe { raw[0].assume_init_ref() }, &expected[0]); + assert_eq!(unsafe { raw[2].assume_init_ref() }, &expected[2]); + assert_eq!(unsafe { raw[1].assume_init_ref() }, unsafe { + before[1].assume_init_ref() + }); + assert_eq!(unsafe { raw[3].assume_init_ref() }, unsafe { + before[3].assume_init_ref() + }); +} + +#[test] +fn uninit_reduce_empty_product_uses_identity() { + assert_uninit_replay(Vec::::new(), KernelDType::F64, ReduceOp::Product, 1.0); +} + +#[test] +fn uninit_reduce_product_and_nonfinite_match_initialized_replay() { + assert_uninit_replay( + vec![2.0f64, 0.5, 2.0, 0.5], + KernelDType::F64, + ReduceOp::Product, + 1.0, ); + let input = vec![f64::INFINITY, 2.0, f64::NEG_INFINITY]; + let dims = [input.len()]; + let source = ErasedRawStridedRef::from_slice(&input, &dims, &[1], 0).unwrap(); + let plan = ErasedReducePlan::compile(KernelDType::F64, ReduceOp::Sum, &dims, &[1]).unwrap(); + let mut initialized = [0.0f64]; + let mut initialized_dest = + ErasedRawStridedMut::from_slice_mut(&mut initialized, &[], &[], 0).unwrap(); + plan.execute(&ExecContext::serial(), &mut initialized_dest, &source) + .unwrap(); + let mut uninitialized = vec![MaybeUninit::new(0.0f64)]; + let mut uninitialized_dest = + ErasedRawStridedUninitMut::from_uninit_slice(&mut uninitialized, &[], &[], 0).unwrap(); + plan.execute_uninit( + &ExecContext::serial(), + &mut uninitialized_dest, + &ErasedRawStridedPtr::from_ref(&source), + ) + .unwrap(); + assert!(initialized[0].is_nan()); + assert!(unsafe { *uninitialized[0].assume_init_ref() }.is_nan()); +} + +#[test] +fn uninit_reduce_simd_tail_and_sum_squares_match_initialized_replay() { + let input = (0..65) + .map(|i| if i % 2 == 0 { 2.0 } else { 0.5 }) + .collect::>(); + assert_uninit_replay(input, KernelDType::F64, ReduceOp::Product, 2.0); + let input = (0..17).map(|i| i as f64 - 8.0).collect::>(); + let expected = input.iter().map(|value| value * value).sum(); + assert_uninit_replay(input, KernelDType::F64, ReduceOp::SumSquares, expected); } #[test] fn validation_errors_leave_uninitialized_bytes_untouched() { let input = [1.0f64, 2.0]; let dims = [2usize]; - let source = ErasedRawStridedRef::new(KernelDType::F64, bytes(&input), &dims, &[1], 0).unwrap(); + let source = ErasedRawStridedRef::from_slice(&input, &dims, &[1], 0).unwrap(); let plan = ErasedReducePlan::compile(KernelDType::F64, ReduceOp::Sum, &dims, &[1]).unwrap(); - let mut raw = vec![MaybeUninit::new(0xa5u8); core::mem::size_of::()]; + let mut raw = vec![MaybeUninit::::new(0xa5 as f64); 1]; let before = raw.clone(); - let mut dest = ErasedRawStridedUninitMut::new(KernelDType::F64, &mut raw, &[], &[], 0).unwrap(); - let wrong = - ErasedRawStridedRef::new(KernelDType::I32, bytes(&[1i32, 2]), &dims, &[1], 0).unwrap(); + let mut dest = ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &[], &[], 0).unwrap(); + let wrong = ErasedRawStridedRef::from_slice(&[1i32, 2], &dims, &[1], 0).unwrap(); assert!(plan .execute_uninit( &ExecContext::serial(), @@ -160,8 +212,8 @@ fn validation_errors_leave_uninitialized_bytes_untouched() { ) .is_err()); assert_eq!( - uninit_bytes_mut(dest.data_mut()), - uninit_bytes_mut(&mut before.clone()) + unsafe { dest.data_as_uninit_mut::().unwrap()[0].assume_init_ref() }, + unsafe { before[0].assume_init_ref() } ); assert!(plan .execute_uninit(