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..c4faba5 --- /dev/null +++ b/docs/2026-07-31-issue-187-worklog.md @@ -0,0 +1,110 @@ +# 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 + +Initial fixture pass did not run benchmark timing; later affinity benchmark +evidence is recorded below with exact candidate-relative upper bounds. + +## 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/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_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 new file mode 100644 index 0000000..8403536 --- /dev/null +++ b/strided-kernel/benches/issue_187_uninit_indexed.rs @@ -0,0 +1,336 @@ +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 sample( + ctx: &ExecContext, + plan: &ErasedReducePlan, + source: &ErasedRawStridedRef<'_>, + init: &mut [f64], + raw: &mut [MaybeUninit], + initialized_first: bool, +) -> (Duration, Duration) { + let initialized = || { + 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::from_uninit_slice(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::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]; + 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::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::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) + .unwrap(); + black_box(&gather_init); + start.elapsed() + }, + || { + let mut dest = ErasedRawStridedUninitMut::from_uninit_slice( + &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::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::from_slice_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::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(); + slice + .execute_uninit(&ctx, &mut dest, &operand, &starts) + .unwrap(); + black_box(&slice_raw); + start.elapsed() + }, + ); + + let update_ref = + ErasedRawStridedRef::from_slice(&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::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) + .unwrap(); + black_box(&update_init); + start.elapsed() + }, + || { + 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); + 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::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::from_slice_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::from_uninit_slice( + &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..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}; @@ -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)] @@ -315,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, @@ -393,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, @@ -485,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, @@ -574,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, @@ -811,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, @@ -1029,18 +1093,90 @@ 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 => { + let mut writer = reduce_writer::(dest)?; + dispatch_reduce::(self.op, &self.layout, ctx, &mut writer, src) + } + KernelDType::F64 => { + let mut writer = reduce_writer::(dest)?; + dispatch_reduce::(self.op, &self.layout, ctx, &mut writer, src) + } + KernelDType::I32 => { + let mut writer = reduce_writer::(dest)?; + dispatch_reduce::(self.op, &self.layout, ctx, &mut writer, src) + } + KernelDType::I64 => { + 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) + } _ => Err(StridedError::UnsupportedDType { dtype: self.dtype.label(), }), }; 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, + 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)?; + 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 +1242,8 @@ impl ErasedGatherPlan { &self.plan, self.index_dtype, dest, - operand, - start_indices, + &operand, + &start_indices, ), KernelDType::F64 => dispatch_gather_index::( &self.plan, @@ -1157,6 +1293,89 @@ impl ErasedGatherPlan { }); result } + + /// 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, + 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 +1441,8 @@ impl ErasedDynamicSlicePlan { &self.plan, self.index_dtype, dest, - operand, - starts, + &operand, + &starts, ), KernelDType::F64 => dispatch_dynamic_slice_index::( &self.plan, @@ -1273,6 +1492,88 @@ impl ErasedDynamicSlicePlan { }); result } + + /// 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, + 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 +1643,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 +1701,98 @@ 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, + 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 +1864,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,9 +1914,103 @@ impl ErasedScatterPlan { }); result } -} -fn execute_one_shot_map( + /// 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, + 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, dest: &mut ErasedRawStridedMut<'_>, input: &ErasedRawStridedRef<'_>, @@ -1919,6 +2406,40 @@ fn check_dtype(expected: KernelDType, actual: KernelDType) -> Result<()> { Ok(()) } +fn reduce_writer<'a, T>(dest: &'a mut ErasedRawStridedMut<'_>) -> Result> +where + T: KernelStorageElement, +{ + let offset = dest.offset(); + 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<'_>, +) -> Result> +where + T: KernelStorageElement, +{ + let offset = dest.offset(); + let data = dest.data_as_uninit_mut::()?; + let ptr = data.as_mut_ptr().cast::(); + let extent = data.len(); + Ok(RawReduceWriter { + ptr, + extent, + offset, + _marker: core::marker::PhantomData, + }) +} + fn check_fused_dtype(dtype: KernelDType) -> Result<()> { match dtype { KernelDType::F32 @@ -2112,6 +2633,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 + KernelStorageElement, +{ + 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 + KernelStorageElement, + I: GatherIndex + KernelStorageElement, +{ + 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 = dest.data_as_uninit_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 + KernelStorageElement, +{ + 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 + KernelStorageElement, + I: GatherIndex + KernelStorageElement, +{ + 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 = dest.data_as_uninit_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 + KernelStorageElement, +{ + 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 + KernelStorageElement, + I: GatherIndex + KernelStorageElement, +{ + 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 = dest.data_as_uninit_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 + KernelStorageElement, +{ + 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 + KernelStorageElement, + I: GatherIndex + KernelStorageElement, +{ + let indices = scatter_indices; + 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 = dest.data_as_uninit_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 +3108,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 +3146,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 +3212,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 +3235,7 @@ where reduce_dims, dest_total, reduce_total, - } => execute_reduce_axes::( + } => execute_reduce_axes::( op, ctx, dest, @@ -2467,21 +3255,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 +3278,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 +3303,7 @@ where return execute_reduce_axes_parallel( op, dest_offset_base, - dest_data, + dest, src.offset(), source_data, layout, @@ -2526,45 +3315,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 +3383,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 +3402,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 +3446,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..9daf903 --- /dev/null +++ b/strided-kernel/tests/issue_187_source_contract.rs @@ -0,0 +1,176 @@ +#[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")); + 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] +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("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 new file mode 100644 index 0000000..60759fb --- /dev/null +++ b/strided-kernel/tests/issue_187_uninit_indexed.rs @@ -0,0 +1,1671 @@ +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 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); + } +} + +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::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::from_slice_mut(&mut expected, &dd, &[1], 0).unwrap(); + plan.execute(&ExecContext::serial(), &mut init, &source, &index) + .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( + &ExecContext::max_threads(4).unwrap(), + &mut out, + &source_ptr, + &index_ptr, + ) + .unwrap(); + assert_initialized_eq(&raw, &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) + ] +); +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() { + 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::from_raw_parts( + KernelDType::Bool, + NonNull::new(bad.as_ptr() as *mut u8).unwrap(), + bad.len(), + &od, + &[1], + 0, + ) + .unwrap() + }; + let indices = [0i32]; + 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, + &operand, + &ErasedRawStridedPtr::from_ref(&index), + ); + assert!(result.is_err()); + assert_eq!(unsafe { raw[0].assume_init_ref() }, unsafe { + before[0].assume_init_ref() + }); +} + +#[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::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, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&index), + ) + .unwrap(); + assert_initialized_eq(&raw, &[false, true]); +} + +#[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::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(), + 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); + } + let _ = raw; +} + +#[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::from_slice(&operand, &od, &[1], 0).unwrap(); + let bad_indices = [0i64, 0]; + 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(), + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&index), + ) + .is_err()); + assert_eq!( + unsafe { out.data_as_uninit_mut::().unwrap()[0].assume_init_ref() }, + unsafe { before[0].assume_init_ref() } + ); +} + +#[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::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, + &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::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::::uninit(); 2]; + let mut out = + ErasedRawStridedUninitMut::from_uninit_slice(&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_initialized_eq(&raw, &expected); + let mut expected_update = operand; + let mut init_update = + ErasedRawStridedMut::from_slice_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::::uninit(); 5]; + let mut out_update = + ErasedRawStridedUninitMut::from_uninit_slice(&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_initialized_eq(&raw_update, &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::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::from_slice_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::<$ty>::uninit(); operand.len()]; + 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(&update), + &ErasedRawStridedPtr::from_ref(&start), + ) + .unwrap(); + assert_initialized_eq(&raw, &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::from_slice(&operand, &dims, &[1], 0).unwrap(); + let start = ErasedRawStridedRef::from_slice(&starts, &sd, &[1], 0).unwrap(); + let update = unsafe { + ErasedRawStridedPtr::from_raw_parts( + 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(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(), + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &update, + &ErasedRawStridedPtr::from_ref(&start), + ) + .is_err()); + assert_eq!(unsafe { raw[0].assume_init_ref() }, unsafe { + before[0].assume_init_ref() + }); +} + +#[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::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, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&update), + &ErasedRawStridedPtr::from_ref(&start), + ) + .unwrap(); + 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 { + ($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::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::from_slice_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::<$ty>::uninit(); 2]; + let mut out = + ErasedRawStridedUninitMut::from_uninit_slice(&mut raw, &dd, &[1], 0).unwrap(); + plan.execute_uninit( + &ctx, + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&start), + ) + .unwrap(); + assert_initialized_eq(&raw, &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::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::::uninit(); 3]; + let mut out = ErasedRawStridedUninitMut::from_uninit_slice(&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_initialized_eq(&raw, &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::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(), + ExecContext::max_threads(2).unwrap(), + ExecContext::max_threads(4).unwrap(), + ] { + 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, &initialized); + } + } + }; +} + +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) + ] +); +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() { + 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]), + 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] +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::from_slice(&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::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(), + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&index), + ) + .unwrap(); + 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::from_uninit_slice(&mut scalar, &[], &[], 0).unwrap(); + reduce + .execute_uninit( + &ExecContext::serial(), + &mut reduce_out, + &ErasedRawStridedPtr::from_ref(&source), + ) + .unwrap(); + 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::from_slice(&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::from_slice(&indices, &index_dims, &[1], 0).unwrap(); + let mut expected = vec![0.0f64; n]; + 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::from_uninit_slice(&mut raw, &dims, &[1], 0).unwrap(); + gather + .execute_uninit( + &ctx, + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&index), + ) + .unwrap(); + assert_initialized_eq(out.data_as_uninit_mut::().unwrap(), &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::from_slice(&starts, &start_dims, &[1], 0).unwrap(); + let mut slice_expected = vec![0.0f64; n / 2]; + 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(); + 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::from_uninit_slice(&mut raw, &slice_dims, &[1], 0).unwrap(); + slice + .execute_uninit( + &ctx, + &mut out, + &ErasedRawStridedPtr::from_ref(&source), + &ErasedRawStridedPtr::from_ref(&start), + ) + .unwrap(); + 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::from_slice(&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::from_slice_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::from_uninit_slice(&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_initialized_eq(out.data_as_uninit_mut::().unwrap(), &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::from_slice(&axis_input, &axis_dims, &axis_strides, 0).unwrap(); + let mut axis_expected = vec![0.0f64; n]; + let mut axis_init = ErasedRawStridedMut::from_slice_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::from_uninit_slice( + &mut raw, + &axis_dest_dims, + &axis_dest_strides, + 0, + ) + .unwrap(); + axis_plan + .execute_uninit(&ctx, &mut out, &ErasedRawStridedPtr::from_ref(&axis_source)) + .unwrap(); + 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 new file mode 100644 index 0000000..ab839a6 --- /dev/null +++ b/strided-kernel/tests/issue_187_uninit_reduce.rs @@ -0,0 +1,225 @@ +use core::fmt::Debug; +use core::mem::MaybeUninit; +use num_complex::{Complex32, Complex64}; +use strided_kernel::{ + ErasedRawStridedMut, ErasedRawStridedPtr, ErasedRawStridedRef, ErasedRawStridedUninitMut, + ErasedReducePlan, ExecContext, KernelDType, KernelStorageElement, ReduceOp, +}; + +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); + + 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 { + ($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::from_slice(&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::from_slice_mut(&mut expected, &[], &[], 0).unwrap(); + plan.execute(&ctx, &mut initialized, &source).unwrap(); + + let mut raw = vec![MaybeUninit::<$ty>::uninit(); 1]; + let mut uninit = + 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!(unsafe { raw[0].assume_init_ref() }, &expected[0]); + } + } + }; +} + +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::from_slice(&input, &src_dims, &src_strides, 4).unwrap(); + let mut expected = [0.0f64; 4]; + 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(0x5a as f64); 4]; + let before = raw.clone(); + let mut uninit = + ErasedRawStridedUninitMut::from_uninit_slice(&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!(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::from_slice(&input, &dims, &[1], 0).unwrap(); + let plan = ErasedReducePlan::compile(KernelDType::F64, ReduceOp::Sum, &dims, &[1]).unwrap(); + let mut raw = vec![MaybeUninit::::new(0xa5 as f64); 1]; + let before = raw.clone(); + 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(), + &mut dest, + &ErasedRawStridedPtr::from_ref(&wrong) + ) + .is_err()); + assert_eq!( + unsafe { dest.data_as_uninit_mut::().unwrap()[0].assume_init_ref() }, + unsafe { before[0].assume_init_ref() } + ); + assert!(plan + .execute_uninit( + &ExecContext::serial(), + &mut dest, + &ErasedRawStridedPtr::from_ref(&source) + ) + .is_ok()); +}