From a6af0d5de1b104bfa97c9c5a6dbeb4cc9155741e Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Fri, 31 Jul 2026 09:56:11 +0900 Subject: [PATCH 1/2] Add safe uninitialized GEMM overwrite paths --- Cargo.toml | 2 +- docs/2026-07-31-issue-188-worklog.md | 67 ++ strided-einsum2/src/backend.rs | 38 + strided-einsum2/src/bgemm_blas.rs | 331 ++++++- strided-einsum2/src/bgemm_faer.rs | 6 +- strided-einsum2/src/contiguous.rs | 178 +++- strided-einsum2/src/dot_general.rs | 29 +- strided-einsum2/src/lib.rs | 9 +- strided-einsum2/src/raw_bgemm.rs | 87 +- strided-einsum2/src/uninit.rs | 870 ++++++++++++++++++ strided-einsum2/tests/blas_inject_fallback.rs | 141 ++- strided-opteinsum/src/expr.rs | 242 ++++- 12 files changed, 1919 insertions(+), 81 deletions(-) create mode 100644 docs/2026-07-31-issue-188-worklog.md create mode 100644 strided-einsum2/src/uninit.rs diff --git a/Cargo.toml b/Cargo.toml index 6fa2502..9fa48ad 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,7 +31,7 @@ mdarray-opteinsum = { version = "0.3.0", path = "mdarray-opteinsum", default-fea ndarray-opteinsum = { version = "0.3.0", path = "ndarray-opteinsum", default-features = false } approx = "0.5" -cblas-inject = "0.1" +cblas-inject = "0.1.2" cblas-sys = "0.2" criterion = "0.5" faer = { version = "0.24", default-features = false, features = ["std", "rayon"] } diff --git a/docs/2026-07-31-issue-188-worklog.md b/docs/2026-07-31-issue-188-worklog.md new file mode 100644 index 0000000..3d39472 --- /dev/null +++ b/docs/2026-07-31-issue-188-worklog.md @@ -0,0 +1,67 @@ +# Issue 188 worklog + +## Scope + +Add overwrite-only contraction entry points whose destination is borrowed as +`MaybeUninit`. The initialized beta APIs remain unchanged. + +## 2026-07-31 + +- Raised the workspace `cblas-inject` minimum to `0.1.2`. +- Added explicit-context `einsum2_into_uninit`, + `einsum2_into_owned_uninit`, `dot_general_into_uninit`, and + `bgemm_raw_strided_into_uninit`. +- Added checked output injectivity, checked offset traversal, conservative + input/output overlap rejection, empty-output no-write behavior, trace-axis + reduction, and tests for matrix output and pre-write validation failure. +- Migrated `strided-opteinsum` intermediate pool acquisition to + `MaybeUninit` and finalize only after a successful overwrite call. +- Kept the default Faer-only `strided-opteinsum` path on initialized storage + until Faer provides a typed overwrite API; direct Faer uninitialized entry + points now return a typed `Unsupported` error instead of silently routing to + the naive provider. +- Added a shared raw-GEMM preflight before label construction, temporary + allocation, or provider dispatch; it checks full dimension agreement, + injectivity, conservative aliasing, checked products, and BLAS sizes. +- Made column-major uninitialized allocation and opteinsum pool sizing + propagate checked overflow errors instead of using unchecked products. + +## Verification + +- `cargo fmt --all` +- `cargo test -p strided-einsum2 --lib` +- `cargo test -p strided-opteinsum --lib` +- `cargo test -p strided-einsum2 --no-default-features` +- `cargo test -p strided-einsum2 --no-default-features --features faer` +- `cargo test -p strided-einsum2 --no-default-features --features blas` +- `cargo test -p strided-einsum2 --no-default-features --features blas-inject` +- `cargo test -p strided-opteinsum --no-default-features --features faer` +- `cargo test -p strided-opteinsum --no-default-features --features blas` +- `cargo test -p strided-opteinsum --no-default-features --features blas-inject` +- `cargo check -p strided-einsum2 --no-default-features` +- `cargo check -p strided-einsum2 --no-default-features --features faer` +- `cargo check -p strided-einsum2 --no-default-features --features blas-inject` +- `cargo check -p strided-einsum2 --no-default-features --features blas` +- `cargo test -p strided-einsum2 --no-default-features --features blas-inject --test blas_inject_fallback -- --test-threads=1` +- `cargo fmt --all -- --check` +- `git diff --check` +- `CARGO_BUILD_JOBS=4 RUSTFLAGS='-C link-arg=-Wl,--threads=1' cargo test --workspace` + +The first unrestricted workspace run reached the doctest linker but hit a +Rust 1.97 `rust-lld` bus error under the full local parallel load. The same +workspace command passed after limiting Cargo jobs and linker threads. The +repository has no configured rules-review script; package clippy remains +non-clean on the current Rust 1.97 toolchain because of pre-existing lints in +the workspace, so it is not used as the PR gate. + +## Remaining work + +- Faer uninitialized output remains blocked by #195: the current Faer + `MatMut::from_raw_parts_mut` contract does not permit forming `MatMut` over + `MaybeUninit`, including for `Accum::Replace`. No cast or zero-fill + workaround is allowed. +- The naive, system BLAS, and injected BLAS overwrite paths are wired through + the public uninitialized APIs. Their feature-specific tests cover direct and + non-contiguous temporary/writeback execution, including a poisoned-C + zgemm regression. Faer remains explicitly blocked on strided-rs#195; its + initialized compatibility path is tested separately. diff --git a/strided-einsum2/src/backend.rs b/strided-einsum2/src/backend.rs index c5fbe3c..a339b8b 100644 --- a/strided-einsum2/src/backend.rs +++ b/strided-einsum2/src/backend.rs @@ -4,6 +4,7 @@ //! and the `ActiveBackend` type alias that serves as the single point of //! backend selection based on Cargo features. +use strided_kernel::ExecContext; use strided_view::ElementOp; /// Trait for backends that can execute batched GEMM on contiguous operands. @@ -47,6 +48,23 @@ pub trait Backend { ) -> strided_view::Result<()>; } +/// Private overwrite-only backend contract. The initialized `Backend` trait +/// remains unchanged for beta-bearing callers. +#[allow(dead_code)] +pub(crate) trait OverwriteBackend { + fn bgemm_contiguous_overwrite( + c: &mut crate::contiguous::UninitContiguousOperand<'_, '_, T>, + a: &crate::contiguous::ContiguousOperand, + b: &crate::contiguous::ContiguousOperand, + batch_dims: &[usize], + m: usize, + n: usize, + k: usize, + alpha: T, + ctx: &ExecContext, + ) -> strided_view::Result<()>; +} + // --------------------------------------------------------------------------- // Marker structs // --------------------------------------------------------------------------- @@ -139,6 +157,26 @@ where } } +#[cfg(not(any(feature = "blas", feature = "blas-inject")))] +impl OverwriteBackend for NaiveBackend +where + T: crate::ScalarBase + strided_view::ElementOpApply, +{ + fn bgemm_contiguous_overwrite( + c: &mut crate::contiguous::UninitContiguousOperand<'_, '_, T>, + a: &crate::contiguous::ContiguousOperand, + b: &crate::contiguous::ContiguousOperand, + batch_dims: &[usize], + m: usize, + n: usize, + k: usize, + alpha: T, + ctx: &ExecContext, + ) -> strided_view::Result<()> { + crate::uninit::bgemm_contiguous_naive(c, a, b, batch_dims, m, n, k, alpha, ctx) + } +} + // --------------------------------------------------------------------------- // ActiveBackend type alias -- the SINGLE point of backend selection // --------------------------------------------------------------------------- diff --git a/strided-einsum2/src/bgemm_blas.rs b/strided-einsum2/src/bgemm_blas.rs index ce66e3c..1b76015 100644 --- a/strided-einsum2/src/bgemm_blas.rs +++ b/strided-einsum2/src/bgemm_blas.rs @@ -4,17 +4,19 @@ //! Operands must already have contiguous inner dimensions (prepared via //! `prepare_input_*` and `prepare_output_*` in the `contiguous` module). -use crate::backend::{Backend, BlasBackend}; +use crate::backend::{Backend, BlasBackend, OverwriteBackend}; use crate::contiguous::{ContiguousOperand, ContiguousOperandMut}; use crate::util::{try_fuse_group, MultiIndex}; use crate::ScalarBase; +use strided_kernel::ExecContext; #[cfg(all(feature = "blas-inject", not(feature = "blas")))] mod inject_fallback { use std::ffi::c_char; use std::sync::Once; - use num_complex::Complex64; + use num_complex::{Complex32, Complex64}; + use num_traits::Zero; static REGISTER_ONCE: Once = Once::new(); @@ -24,24 +26,26 @@ mod inject_fallback { } #[inline] - unsafe fn gemm_real( + unsafe fn gemm_real( transa: u8, transb: u8, m: usize, n: usize, k: usize, - alpha: f64, - a: *const f64, + alpha: T, + a: *const T, lda: usize, - b: *const f64, + b: *const T, ldb: usize, - beta: f64, - c: *mut f64, + beta: T, + c: *mut T, ldc: usize, - ) { + ) where + T: Copy + Zero + PartialEq + std::ops::Mul + std::ops::Add, + { for j in 0..n { for i in 0..m { - let mut sum = 0.0f64; + let mut sum = T::zero(); for p in 0..k { let a_val = if transa == b'N' { *a.add(i + p * lda) @@ -53,10 +57,14 @@ mod inject_fallback { } else { *b.add(j + p * ldb) }; - sum += a_val * b_val; + sum = sum + a_val * b_val; } let c_ptr = c.add(i + j * ldc); - *c_ptr = alpha * sum + beta * *c_ptr; + *c_ptr = if beta == T::zero() { + alpha * sum + } else { + alpha * sum + beta * *c_ptr + }; } } } @@ -100,7 +108,11 @@ mod inject_fallback { sum += a_val * b_val; } let c_ptr = c.add(i + j * ldc); - *c_ptr = alpha * sum + beta * *c_ptr; + *c_ptr = if beta == Complex64::new(0.0, 0.0) { + alpha * sum + } else { + alpha * sum + beta * *c_ptr + }; } } } @@ -141,6 +153,40 @@ mod inject_fallback { } } + unsafe extern "C" fn sgemm_fallback( + transa: *const c_char, + transb: *const c_char, + m: *const cblas_sys::blasint, + n: *const cblas_sys::blasint, + k: *const cblas_sys::blasint, + alpha: *const f32, + a: *const f32, + lda: *const cblas_sys::blasint, + b: *const f32, + ldb: *const cblas_sys::blasint, + beta: *const f32, + c: *mut f32, + ldc: *const cblas_sys::blasint, + ) { + unsafe { + gemm_real( + trans_flag(*transa), + trans_flag(*transb), + *m as usize, + *n as usize, + *k as usize, + *alpha, + a, + *lda as usize, + b, + *ldb as usize, + *beta, + c, + *ldc as usize, + ); + } + } + unsafe extern "C" fn zgemm_fallback( transa: *const c_char, transb: *const c_char, @@ -177,16 +223,154 @@ mod inject_fallback { } } + unsafe extern "C" fn cgemm_fallback( + transa: *const c_char, + transb: *const c_char, + m: *const cblas_sys::blasint, + n: *const cblas_sys::blasint, + k: *const cblas_sys::blasint, + alpha: *const Complex32, + a: *const Complex32, + lda: *const cblas_sys::blasint, + b: *const Complex32, + ldb: *const cblas_sys::blasint, + beta: *const Complex32, + c: *mut Complex32, + ldc: *const cblas_sys::blasint, + ) { + let transa = trans_flag(*transa); + let transb = trans_flag(*transb); + unsafe { + for j in 0..*n as usize { + for i in 0..*m as usize { + let mut sum = Complex32::new(0.0, 0.0); + for p in 0..*k as usize { + let mut av = if transa == b'N' { + *a.add(i + p * *lda as usize) + } else { + *a.add(p + i * *lda as usize) + }; + let mut bv = if transb == b'N' { + *b.add(p + j * *ldb as usize) + } else { + *b.add(j + p * *ldb as usize) + }; + if transa == b'C' { + av = av.conj(); + } + if transb == b'C' { + bv = bv.conj(); + } + sum = sum + av * bv; + } + let out = c.add(i + j * *ldc as usize); + *out = if *beta == Complex32::new(0.0, 0.0) { + *alpha * sum + } else { + *alpha * sum + *beta * *out + }; + } + } + } + } + pub(super) fn ensure_registered() { REGISTER_ONCE.call_once(|| unsafe { if !cblas_sys::is_dgemm_registered() { cblas_sys::register_dgemm(dgemm_fallback); } + if !cblas_sys::is_sgemm_registered() { + cblas_sys::register_sgemm(sgemm_fallback); + } + if !cblas_sys::is_cgemm_registered() { + cblas_sys::register_cgemm(cgemm_fallback); + } if !cblas_sys::is_zgemm_registered() { cblas_sys::register_zgemm(zgemm_fallback); } }); } + + #[cfg(test)] + mod tests { + use super::*; + use std::mem::MaybeUninit; + + fn args( + alpha: &T, + a: *const T, + beta: &T, + c: *mut T, + ) -> ( + *const c_char, + *const c_char, + *const cblas_sys::blasint, + *const cblas_sys::blasint, + *const cblas_sys::blasint, + *const T, + *const T, + *const cblas_sys::blasint, + *const T, + *mut T, + *const cblas_sys::blasint, + ) { + static N: c_char = b'N' as c_char; + static ONE: cblas_sys::blasint = 1; + (&N, &N, &ONE, &ONE, &ONE, alpha, a, &ONE, beta, c, &ONE) + } + + #[test] + fn all_registered_fallbacks_skip_poisoned_c_for_zero_beta() { + let alpha32 = 2.0f32; + let beta32 = 0.0f32; + let a32 = 3.0f32; + let b32 = 4.0f32; + let poison32 = f32::NAN; + let mut c32 = MaybeUninit::new(poison32); + let (ta, tb, m, n, k, alpha, a, lda, beta, c, ldc) = + args(&alpha32, &a32, &beta32, c32.as_mut_ptr().cast()); + unsafe { + sgemm_fallback(ta, tb, m, n, k, alpha, a, lda, &b32, lda, beta, c, ldc); + } + assert_eq!(unsafe { c32.assume_init() }, 24.0); + + let alpha64 = 2.0f64; + let beta64 = 0.0f64; + let a64 = 3.0f64; + let b64 = 4.0f64; + let mut c64 = MaybeUninit::new(f64::NAN); + let (ta, tb, m, n, k, alpha, a, lda, beta, c, ldc) = + args(&alpha64, &a64, &beta64, c64.as_mut_ptr().cast()); + unsafe { + dgemm_fallback(ta, tb, m, n, k, alpha, a, lda, &b64, lda, beta, c, ldc); + } + assert_eq!(unsafe { c64.assume_init() }, 24.0); + + let alpha_c = Complex32::new(2.0, 0.0); + let beta_c = Complex32::new(0.0, 0.0); + let a_c = Complex32::new(3.0, 0.0); + let b_c = Complex32::new(4.0, 0.0); + let mut c_c = MaybeUninit::new(Complex32::new(f32::NAN, f32::NAN)); + let (ta, tb, m, n, k, alpha, a, lda, beta, c, ldc) = + args(&alpha_c, &a_c, &beta_c, c_c.as_mut_ptr().cast()); + unsafe { + cgemm_fallback(ta, tb, m, n, k, alpha, a, lda, &b_c, lda, beta, c, ldc); + } + assert_eq!(unsafe { c_c.assume_init() }, Complex32::new(24.0, 0.0)); + + let alpha_z = Complex64::new(2.0, 0.0); + let beta_z = Complex64::new(0.0, 0.0); + let a_z = Complex64::new(3.0, 0.0); + let b_z = Complex64::new(4.0, 0.0); + let mut c_z = MaybeUninit::new(Complex64::new(f64::NAN, f64::NAN)); + let (ta, tb, m, n, k, alpha, a, lda, beta, c, ldc) = + args(&alpha_z, &a_z, &beta_z, c_z.as_mut_ptr().cast()); + unsafe { + zgemm_fallback(ta, tb, m, n, k, alpha, a, lda, &b_z, lda, beta, c, ldc); + } + assert_eq!(unsafe { c_z.assume_init() }, Complex64::new(24.0, 0.0)); + } + } } /// Type-level dispatch trait for CBLAS GEMM. @@ -537,6 +721,108 @@ pub(crate) fn bgemm_contiguous_into strided_view::Result<(cblas_sys::CBLAS_TRANSPOSE, i32)> { + let (trans, lda) = if row_stride == 1 || row_stride == 0 { + ( + cblas_sys::CBLAS_TRANSPOSE::CblasNoTrans, + col_stride.max(nrows as isize).max(1), + ) + } else if col_stride == 1 || col_stride == 0 { + ( + cblas_sys::CBLAS_TRANSPOSE::CblasTrans, + row_stride.max(ncols as isize).max(1), + ) + } else { + return Err(strided_view::StridedError::PlanLayoutMismatch); + }; + Ok(( + trans, + i32::try_from(lda).map_err(|_| strided_view::StridedError::OffsetOverflow)?, + )) +} + +/// CBLAS overwrite path. The literal zero beta is part of the private +/// contract; cblas-inject 0.1.2 guarantees that exact zero never reads C. +#[allow(clippy::too_many_arguments)] +pub(crate) fn bgemm_contiguous_overwrite( + c: &mut crate::contiguous::UninitContiguousOperand<'_, '_, T>, + a: &ContiguousOperand, + b: &ContiguousOperand, + batch_dims: &[usize], + m: usize, + n: usize, + k: usize, + alpha: T, + _ctx: &ExecContext, +) -> strided_view::Result<()> +where + T: ScalarBase + strided_view::ElementOpApply + BlasGemm, +{ + #[cfg(all(feature = "blas-inject", not(feature = "blas")))] + inject_fallback::ensure_registered(); + debug_assert!(!a.conj() && !b.conj()); + let (trans_a, lda) = checked_operand_layout(a.row_stride(), a.col_stride(), m, k)?; + let (trans_b, ldb) = checked_operand_layout(b.row_stride(), b.col_stride(), k, n)?; + let m_i32 = i32::try_from(m).map_err(|_| strided_view::StridedError::OffsetOverflow)?; + let n_i32 = i32::try_from(n).map_err(|_| strided_view::StridedError::OffsetOverflow)?; + let k_i32 = i32::try_from(k).map_err(|_| strided_view::StridedError::OffsetOverflow)?; + let c_is_col_major = c.row_stride() == 1 || c.row_stride() == 0; + let ldc_value = if c_is_col_major { + c.col_stride().max(m as isize).max(1) + } else { + c.row_stride().max(n as isize).max(1) + }; + let ldc = i32::try_from(ldc_value).map_err(|_| strided_view::StridedError::OffsetOverflow)?; + let zero = T::zero(); + let mut batch = MultiIndex::new(batch_dims); + while batch.next().is_some() { + let a_off = batch.offset(a.batch_strides()); + let b_off = batch.offset(b.batch_strides()); + let c_off = batch.offset(c.batch_strides()); + unsafe { + if c_is_col_major { + T::gemm( + trans_a, + trans_b, + m_i32, + n_i32, + k_i32, + alpha, + a.ptr().offset(a_off), + lda, + b.ptr().offset(b_off), + ldb, + zero, + c.ptr().offset(c_off).cast(), + ldc, + ); + } else { + T::gemm( + flip_transpose(trans_b), + flip_transpose(trans_a), + n_i32, + m_i32, + k_i32, + alpha, + b.ptr().offset(b_off), + ldb, + a.ptr().offset(a_off), + lda, + zero, + c.ptr().offset(c_off).cast(), + ldc, + ); + } + } + } + Ok(()) +} + impl Backend for BlasBackend where T: ScalarBase + strided_view::ElementOpApply + BlasGemm, @@ -560,3 +846,22 @@ where self::bgemm_contiguous_into(c, a, b, batch_dims, m, n, k, alpha, beta) } } + +impl OverwriteBackend for BlasBackend +where + T: ScalarBase + strided_view::ElementOpApply + BlasGemm, +{ + fn bgemm_contiguous_overwrite( + c: &mut crate::contiguous::UninitContiguousOperand<'_, '_, T>, + a: &ContiguousOperand, + b: &ContiguousOperand, + batch_dims: &[usize], + m: usize, + n: usize, + k: usize, + alpha: T, + ctx: &ExecContext, + ) -> strided_view::Result<()> { + bgemm_contiguous_overwrite(c, a, b, batch_dims, m, n, k, alpha, ctx) + } +} diff --git a/strided-einsum2/src/bgemm_faer.rs b/strided-einsum2/src/bgemm_faer.rs index f02a523..6fb44c0 100644 --- a/strided-einsum2/src/bgemm_faer.rs +++ b/strided-einsum2/src/bgemm_faer.rs @@ -177,7 +177,7 @@ where let a_contig_buf: Option>; let (a_ptr, a_row_stride, a_col_stride); if a_needs_copy { - let mut buf = alloc_col_major_uninit(a.dims()); + let mut buf = alloc_col_major_uninit(a.dims())?; strided_kernel::copy_into(&mut buf.view_mut(), &a.as_view())?; a_ptr = buf.view().ptr(); // Col-major inner A [lo..., sum...]: lo stride = 1, sum stride = m @@ -201,7 +201,7 @@ where let b_contig_buf: Option>; let (b_ptr, b_row_stride, b_col_stride); if b_needs_copy { - let mut buf = alloc_col_major_uninit(b.dims()); + let mut buf = alloc_col_major_uninit(b.dims())?; strided_kernel::copy_into(&mut buf.view_mut(), &b.as_view())?; b_ptr = buf.view().ptr(); // Col-major inner B [sum..., ro...]: sum stride = 1, ro stride = k @@ -225,7 +225,7 @@ where let c_contig_buf: Option>; let (c_ptr, c_row_stride, c_col_stride); if c_needs_copy { - let mut buf = alloc_col_major_uninit(c.dims()); + let mut buf = alloc_col_major_uninit(c.dims())?; if beta != T::zero() { let c_view: StridedView<'_, T> = c.as_view(); strided_kernel::copy_into(&mut buf.view_mut(), &c_view)?; diff --git a/strided-einsum2/src/contiguous.rs b/strided-einsum2/src/contiguous.rs index b8a9ce5..7473a9b 100644 --- a/strided-einsum2/src/contiguous.rs +++ b/strided-einsum2/src/contiguous.rs @@ -8,6 +8,7 @@ use crate::ScalarBase; use std::any::{Any, TypeId}; use std::cell::RefCell; use std::collections::HashMap; +use std::mem::MaybeUninit; use strided_view::{RawStridedMut, RawStridedRef, StridedArray, StridedView, StridedViewMut}; /// GEMM-ready input operand with contiguous data. @@ -36,6 +37,20 @@ pub struct ContiguousOperandMut { buf_is_pooled: bool, } +/// Overwrite-only C operand. The destination borrow is retained until the +/// provider has returned successfully; no initialized C view is constructed +/// while the provider is running. +#[allow(dead_code)] +pub(crate) struct UninitContiguousOperand<'a, 'b, T: Copy + 'static> { + destination: &'a mut RawStridedMut<'b, MaybeUninit>, + ptr: *mut MaybeUninit, + row_stride: isize, + col_stride: isize, + batch_strides: Vec, + temp: Option>>, + writeback: Option, +} + thread_local! { static BUFFER_POOL: RefCell>> = RefCell::new(HashMap::new()); } @@ -106,26 +121,32 @@ fn return_pooled_vec(mut data: Vec) { }); } -fn alloc_col_major_uninit_with_pool(dims: &[usize]) -> (StridedArray, bool) { - let total: usize = dims.iter().product::().max(1); +fn alloc_col_major_uninit_with_pool( + dims: &[usize], +) -> strided_view::Result<(StridedArray, bool)> { + let total = dims + .iter() + .try_fold(1usize, |total, &dim| total.checked_mul(dim)) + .ok_or(strided_view::StridedError::OffsetOverflow)? + .max(1); let bytes = total.saturating_mul(std::mem::size_of::()); if bytes == 0 || bytes > MAX_POOLED_BYTES { - return (alloc_col_major_uninit(dims), false); + return Ok((alloc_col_major_uninit(dims)?, false)); } let data = take_pooled_vec_uninit::(total); let arr = unsafe { StridedArray::col_major_from_buffer_uninit(data, dims) }; - (arr, true) + Ok((arr, true)) } /// Allocate a col-major buffer, optionally reusing from the thread-local pool. fn alloc_maybe_pooled( dims: &[usize], use_pool: bool, -) -> (StridedArray, bool) { +) -> strided_view::Result<(StridedArray, bool)> { if use_pool { alloc_col_major_uninit_with_pool(dims) } else { - (alloc_col_major_uninit(dims), false) + Ok((alloc_col_major_uninit(dims)?, false)) } } @@ -374,12 +395,18 @@ fn col_major_layout( /// With batch-last canonical order `[inner..., batch...]`, pure column-major /// naturally gives batch dims the largest strides — each batch slice is a /// contiguous column-major matrix. -pub(crate) fn alloc_col_major_uninit(dims: &[usize]) -> StridedArray { - let total: usize = dims.iter().product::().max(1); +pub(crate) fn alloc_col_major_uninit( + dims: &[usize], +) -> strided_view::Result> { + let total = dims + .iter() + .try_fold(1usize, |total, &dim| total.checked_mul(dim)) + .ok_or(strided_view::StridedError::OffsetOverflow)? + .max(1); // SAFETY: `T: Copy` guarantees no drop glue, so leaving elements - // uninitialised is safe. Every call-site writes all elements before - // reading: A and B via `copy_into`, C via `copy_into` (beta != 0) - // or GEMM with replace semantics (beta == 0). + // uninitialised is safe. Each caller must establish initialization before + // exposing the array as initialized; uninitialized-output callers use + // `T = MaybeUninit` and finalize only after complete overwrite. let mut data = Vec::with_capacity(total); unsafe { data.set_len(total) }; @@ -393,8 +420,7 @@ pub(crate) fn alloc_col_major_uninit(dims: &[usize]) -> StridedArray } } - let arr = StridedArray::from_parts(data, dims, &strides, 0).expect("col-major allocation"); - arr + Ok(StridedArray::from_parts(data, dims, &strides, 0)?) } /// Prepare a borrowed input view for GEMM. @@ -424,7 +450,7 @@ pub fn prepare_input_view( // materialize conj into the data before the GEMM call. if let Some(conj_fn) = materialize_conj_fn { if conj { - let (mut buf, buf_is_pooled) = alloc_maybe_pooled(dims, use_pool); + let (mut buf, buf_is_pooled) = alloc_maybe_pooled(dims, use_pool)?; strided_kernel::map_into(&mut buf.view_mut(), view, conj_fn)?; let ptr = buf.view().ptr(); let (row_stride, col_stride, batch_strides) = col_major_layout(&buf, n_group1, n_inner); @@ -449,7 +475,7 @@ pub fn prepare_input_view( ); if check.needs_copy { - let (mut buf, buf_is_pooled) = alloc_maybe_pooled(dims, use_pool); + let (mut buf, buf_is_pooled) = alloc_maybe_pooled(dims, use_pool)?; strided_kernel::copy_into_col_major(&mut buf.view_mut(), view)?; let ptr = buf.view().ptr(); let (row_stride, col_stride, batch_strides) = col_major_layout(&buf, n_group1, n_inner); @@ -497,7 +523,7 @@ pub fn prepare_input_raw( if let Some(conj_fn) = materialize_conj_fn { if conj { - let (mut buf, buf_is_pooled) = alloc_maybe_pooled(dims, use_pool); + let (mut buf, buf_is_pooled) = alloc_maybe_pooled(dims, use_pool)?; strided_kernel::map_into(&mut buf.view_mut(), &view.as_view(), conj_fn)?; let ptr = buf.view().ptr(); let (row_stride, col_stride, batch_strides) = col_major_layout(&buf, n_group1, n_inner); @@ -522,7 +548,7 @@ pub fn prepare_input_raw( ); if check.needs_copy { - let (mut buf, buf_is_pooled) = alloc_maybe_pooled(dims, use_pool); + let (mut buf, buf_is_pooled) = alloc_maybe_pooled(dims, use_pool)?; strided_kernel::copy_into_col_major(&mut buf.view_mut(), &view.as_view())?; let ptr = buf.view().ptr(); let (row_stride, col_stride, batch_strides) = col_major_layout(&buf, n_group1, n_inner); @@ -574,7 +600,7 @@ pub fn prepare_input_owned( // materialize conj into the data before the GEMM call. if let Some(conj_fn) = materialize_conj_fn { if conj { - let (mut buf, buf_is_pooled) = alloc_maybe_pooled(&dims, use_pool); + let (mut buf, buf_is_pooled) = alloc_maybe_pooled(&dims, use_pool)?; strided_kernel::map_into(&mut buf.view_mut(), &arr.view(), conj_fn)?; let ptr = buf.view().ptr(); let (row_stride, col_stride, batch_strides) = col_major_layout(&buf, n_group1, n_inner); @@ -599,7 +625,7 @@ pub fn prepare_input_owned( ); if check.needs_copy { - let (mut buf, buf_is_pooled) = alloc_maybe_pooled(&dims, use_pool); + let (mut buf, buf_is_pooled) = alloc_maybe_pooled(&dims, use_pool)?; strided_kernel::copy_into_col_major(&mut buf.view_mut(), &arr.view())?; let ptr = buf.view().ptr(); let (row_stride, col_stride, batch_strides) = col_major_layout(&buf, n_group1, n_inner); @@ -664,7 +690,7 @@ pub fn prepare_output_view( ); if check.needs_copy { - let (mut buf, buf_is_pooled) = alloc_maybe_pooled(&dims, use_pool); + let (mut buf, buf_is_pooled) = alloc_maybe_pooled(&dims, use_pool)?; if beta != T::zero() { strided_kernel::copy_into_col_major(&mut buf.view_mut(), &view.as_view())?; } @@ -720,7 +746,7 @@ pub fn prepare_output_raw( ); if check.needs_copy { - let (mut buf, buf_is_pooled) = alloc_maybe_pooled(&dims, use_pool); + let (mut buf, buf_is_pooled) = alloc_maybe_pooled(&dims, use_pool)?; if beta != T::zero() { strided_kernel::copy_into_col_major(&mut buf.view_mut(), &view.as_view())?; } @@ -750,6 +776,116 @@ pub fn prepare_output_raw( } } +#[allow(dead_code)] +impl<'a, 'b, T: Copy + 'static> UninitContiguousOperand<'a, 'b, T> { + #[inline] + pub(crate) fn ptr(&self) -> *mut MaybeUninit { + self.ptr + } + + #[inline] + pub(crate) fn row_stride(&self) -> isize { + self.row_stride + } + + #[inline] + pub(crate) fn col_stride(&self) -> isize { + self.col_stride + } + + #[inline] + pub(crate) fn batch_strides(&self) -> &[isize] { + &self.batch_strides + } + + /// Finalize the destination only after a successful overwrite provider. + pub(crate) fn finalize(self) -> crate::Result<()> { + let Self { + destination, + temp, + writeback, + .. + } = self; + let (Some(temp), Some(writeback)) = (temp, writeback) else { + return Ok(()); + }; + // The temporary is dense and the provider has returned success, so + // every element is initialized. This conversion is intentionally + // after provider success and is never used for direct C storage. + let dims = temp.dims().to_vec(); + let strides = temp.strides().to_vec(); + let data = temp.into_data(); + let len = data.len(); + let cap = data.capacity(); + let ptr = data.as_ptr().cast_mut().cast::(); + std::mem::forget(data); + let initialized = unsafe { + StridedArray::from_parts(Vec::from_raw_parts(ptr, len, cap), &dims, &strides, 0) + }?; + let source = RawStridedRef::new( + initialized.data(), + initialized.dims(), + initialized.strides(), + initialized.view().offset(), + )?; + writeback.execute_uninit(destination, &source)?; + Ok(()) + } +} + +/// Prepare an overwrite-only C operand. A non-fusable destination receives a +/// dense `MaybeUninit` temporary and a compiled uninitialized writeback plan. +#[allow(dead_code)] +pub(crate) fn prepare_output_raw_uninit<'a, 'b, T: ScalarBase + 'static>( + destination: &'a mut RawStridedMut<'b, MaybeUninit>, + n_group1: usize, + n_group2: usize, + requires_unit_stride: bool, +) -> crate::Result> { + let dims = destination.dims().to_vec(); + let strides = destination.strides().to_vec(); + let n_inner = n_group1 + n_group2; + let check = check_contiguity( + &dims[..n_group1], + &strides[..n_group1], + &dims[n_group1..n_inner], + &strides[n_group1..n_inner], + requires_unit_stride, + ); + if !check.needs_copy { + let Some((_, row_stride)) = check.fused_g1 else { + return Err(strided_view::StridedError::PlanLayoutMismatch.into()); + }; + let Some((_, col_stride)) = check.fused_g2 else { + return Err(strided_view::StridedError::PlanLayoutMismatch.into()); + }; + let ptr = destination.as_mut_ptr(); + return Ok(UninitContiguousOperand { + destination, + ptr: ptr.cast(), + row_stride, + col_stride, + batch_strides: strides[n_inner..].to_vec(), + temp: None, + writeback: None, + }); + } + + let temp = alloc_col_major_uninit::>(&dims)?; + let ptr = temp.view().ptr().cast_mut(); + let (row_stride, col_stride, batch_strides) = col_major_layout(&temp, n_group1, n_inner); + let writeback = strided_kernel::CopyPlan::compile(&dims, &strides, temp.strides())?; + Ok(UninitContiguousOperand { + destination, + ptr, + row_stride, + col_stride, + batch_strides, + temp: Some(temp), + writeback: Some(writeback), + }) +} + #[cfg(test)] mod tests_generic_backend { use super::*; diff --git a/strided-einsum2/src/dot_general.rs b/strided-einsum2/src/dot_general.rs index 3c1b7b5..9fef69d 100644 --- a/strided-einsum2/src/dot_general.rs +++ b/strided-einsum2/src/dot_general.rs @@ -25,7 +25,9 @@ //! ``` use smallvec::SmallVec; -use strided_view::{StridedView, StridedViewMut}; +use std::mem::MaybeUninit; +use strided_kernel::ExecContext; +use strided_view::{RawStridedMut, StridedView, StridedViewMut}; use crate::backend::Backend; use crate::{einsum2_dispatch, Einsum2Plan, EinsumError, Result, ScalarBase}; @@ -280,6 +282,31 @@ where dot_general_with_backend_into::(c, a, b, config, alpha, beta) } +/// Compute dot-general into a genuinely uninitialized destination. +/// +/// The destination is only exposed as initialized after this function returns +/// `Ok(())`; holes in a strided backing allocation are never touched. +pub fn dot_general_into_uninit( + c: &mut RawStridedMut<'_, MaybeUninit>, + a: &StridedView<'_, T>, + b: &StridedView<'_, T>, + config: &DotGeneralConfig<'_>, + alpha: T, + ctx: &ExecContext, +) -> Result<()> { + let labels = config.labels_for_shapes(a.dims(), b.dims(), c.dims())?; + crate::einsum2_into_uninit( + c, + a, + b, + labels.out_labels.as_slice(), + labels.lhs_labels.as_slice(), + labels.rhs_labels.as_slice(), + alpha, + ctx, + ) +} + /// Compute `C = alpha * dot_general(A, B) + beta * C` with the naive fallback. #[cfg(not(any(feature = "faer", feature = "blas", feature = "blas-inject")))] pub fn dot_general_into( diff --git a/strided-einsum2/src/lib.rs b/strided-einsum2/src/lib.rs index 4d53591..0018ee2 100644 --- a/strided-einsum2/src/lib.rs +++ b/strided-einsum2/src/lib.rs @@ -57,6 +57,8 @@ pub mod plan; pub mod raw_bgemm; /// Trace-axis reduction (summing axes that appear only in one operand). pub mod trace; +/// Overwrite-only APIs for genuinely uninitialized destinations. +pub mod uninit; /// Shared helpers (permutation inversion, multi-index iteration, dimension fusion). pub mod util; @@ -78,12 +80,15 @@ pub use strided_view::{ }; pub use backend::Backend; -pub use dot_general::{dot_general_into, dot_general_with_backend_into, DotGeneralConfig}; +pub use dot_general::{ + dot_general_into, dot_general_into_uninit, dot_general_with_backend_into, DotGeneralConfig, +}; pub use plan::Einsum2Plan; pub use raw_bgemm::{ bgemm_raw_strided_into, bgemm_raw_strided_into_unchecked, bgemm_raw_with_backend_into, bgemm_raw_with_backend_into_unchecked, }; +pub use uninit::{bgemm_raw_strided_into_uninit, einsum2_into_owned_uninit, einsum2_into_uninit}; /// Trait alias for axis label types. pub trait AxisId: Clone + Eq + Hash + Debug {} @@ -156,6 +161,8 @@ pub enum EinsumError { expected: Vec, got: Vec, }, + #[error("unsupported einsum operation: {0}")] + Unsupported(String), #[error(transparent)] Strided(#[from] strided_view::StridedError), } diff --git a/strided-einsum2/src/raw_bgemm.rs b/strided-einsum2/src/raw_bgemm.rs index 121abef..173f235 100644 --- a/strided-einsum2/src/raw_bgemm.rs +++ b/strided-einsum2/src/raw_bgemm.rs @@ -9,6 +9,55 @@ use crate::backend::Backend; use crate::{contiguous, Scalar, ScalarBase}; use strided_view::{Conj, ElementOp, ElementOpApply, RawStridedMut, RawStridedRef}; +#[derive(Clone, Copy)] +pub(crate) struct BgemmGroupLayout { + pub(crate) a_sum_end: usize, + pub(crate) a_rank: usize, + pub(crate) b_ro_end: usize, + pub(crate) b_rank: usize, + pub(crate) c_ro_end: usize, + pub(crate) c_rank: usize, + pub(crate) label_len: usize, +} + +pub(crate) fn checked_bgemm_group_layout( + n_batch: usize, + n_lo: usize, + n_ro: usize, + n_sum: usize, +) -> crate::Result { + let a_sum_end = n_lo + .checked_add(n_sum) + .ok_or(strided_view::StridedError::OffsetOverflow)?; + let a_rank = a_sum_end + .checked_add(n_batch) + .ok_or(strided_view::StridedError::OffsetOverflow)?; + let b_ro_end = n_sum + .checked_add(n_ro) + .ok_or(strided_view::StridedError::OffsetOverflow)?; + let b_rank = b_ro_end + .checked_add(n_batch) + .ok_or(strided_view::StridedError::OffsetOverflow)?; + let c_ro_end = n_lo + .checked_add(n_ro) + .ok_or(strided_view::StridedError::OffsetOverflow)?; + let c_rank = c_ro_end + .checked_add(n_batch) + .ok_or(strided_view::StridedError::OffsetOverflow)?; + let label_len = c_rank + .checked_add(n_sum) + .ok_or(strided_view::StridedError::OffsetOverflow)?; + Ok(BgemmGroupLayout { + a_sum_end, + a_rank, + b_ro_end, + b_rank, + c_ro_end, + c_rank, + label_len, + }) +} + /// Batched strided GEMM on raw borrowed layout metadata using the active backend. /// /// This is the raw-layout counterpart to backend-specific `bgemm_strided_into` @@ -186,8 +235,8 @@ where Ok(()) } -pub(crate) fn validate_bgemm_shapes( - c: &RawStridedMut<'_, T>, +pub(crate) fn validate_bgemm_shapes( + c: &RawStridedMut<'_, U>, a: &RawStridedRef<'_, T>, b: &RawStridedRef<'_, T>, n_batch: usize, @@ -195,23 +244,21 @@ pub(crate) fn validate_bgemm_shapes( n_ro: usize, n_sum: usize, ) -> crate::Result<()> { - let a_rank = n_lo + n_sum + n_batch; - let b_rank = n_sum + n_ro + n_batch; - let c_rank = n_lo + n_ro + n_batch; - if a.dims().len() != a_rank { - return Err(strided_view::StridedError::RankMismatch(a_rank, a.dims().len()).into()); + let groups = checked_bgemm_group_layout(n_batch, n_lo, n_ro, n_sum)?; + if a.dims().len() != groups.a_rank { + return Err(strided_view::StridedError::RankMismatch(groups.a_rank, a.dims().len()).into()); } - if b.dims().len() != b_rank { - return Err(strided_view::StridedError::RankMismatch(b_rank, b.dims().len()).into()); + if b.dims().len() != groups.b_rank { + return Err(strided_view::StridedError::RankMismatch(groups.b_rank, b.dims().len()).into()); } - if c.dims().len() != c_rank { - return Err(strided_view::StridedError::RankMismatch(c_rank, c.dims().len()).into()); + if c.dims().len() != groups.c_rank { + return Err(strided_view::StridedError::RankMismatch(groups.c_rank, c.dims().len()).into()); } let lo_dims = &a.dims()[..n_lo]; - let sum_dims = &a.dims()[n_lo..n_lo + n_sum]; - let batch_dims = &a.dims()[n_lo + n_sum..]; - let ro_dims = &b.dims()[n_sum..n_sum + n_ro]; + let sum_dims = &a.dims()[n_lo..groups.a_sum_end]; + let batch_dims = &a.dims()[groups.a_sum_end..]; + let ro_dims = &b.dims()[n_sum..groups.b_ro_end]; if &b.dims()[..n_sum] != sum_dims { return Err(strided_view::StridedError::ShapeMismatch( @@ -220,10 +267,10 @@ pub(crate) fn validate_bgemm_shapes( ) .into()); } - if &b.dims()[n_sum + n_ro..] != batch_dims { + if &b.dims()[groups.b_ro_end..] != batch_dims { return Err(strided_view::StridedError::ShapeMismatch( batch_dims.to_vec(), - b.dims()[n_sum + n_ro..].to_vec(), + b.dims()[groups.b_ro_end..].to_vec(), ) .into()); } @@ -234,17 +281,17 @@ pub(crate) fn validate_bgemm_shapes( ) .into()); } - if &c.dims()[n_lo..n_lo + n_ro] != ro_dims { + if &c.dims()[n_lo..groups.c_ro_end] != ro_dims { return Err(strided_view::StridedError::ShapeMismatch( ro_dims.to_vec(), - c.dims()[n_lo..n_lo + n_ro].to_vec(), + c.dims()[n_lo..groups.c_ro_end].to_vec(), ) .into()); } - if &c.dims()[n_lo + n_ro..] != batch_dims { + if &c.dims()[groups.c_ro_end..] != batch_dims { return Err(strided_view::StridedError::ShapeMismatch( batch_dims.to_vec(), - c.dims()[n_lo + n_ro..].to_vec(), + c.dims()[groups.c_ro_end..].to_vec(), ) .into()); } diff --git a/strided-einsum2/src/uninit.rs b/strided-einsum2/src/uninit.rs new file mode 100644 index 0000000..369ea98 --- /dev/null +++ b/strided-einsum2/src/uninit.rs @@ -0,0 +1,870 @@ +//! Overwrite-only contraction entry points. +//! +//! This module deliberately does not reuse the initialized `beta` path. The +//! destination is borrowed as `MaybeUninit` until the last logical element +//! has been written, so a provider failure cannot expose a partially +//! initialized slice as `T`. + +use std::collections::HashSet; +use std::mem::MaybeUninit; + +use strided_kernel::ExecContext; +use strided_view::{ElementOp, RawStridedMut, RawStridedRef, StridedView}; + +use crate::{AxisId, Einsum2Plan, EinsumError, Result, ScalarBase}; + +/// Naive overwrite kernel used by the private backend contract and as the +/// no-provider implementation. It never reads C. +#[cfg(not(any(feature = "blas", feature = "blas-inject")))] +pub(crate) fn bgemm_contiguous_naive( + c: &mut crate::contiguous::UninitContiguousOperand<'_, '_, T>, + a: &crate::contiguous::ContiguousOperand, + b: &crate::contiguous::ContiguousOperand, + batch_dims: &[usize], + m: usize, + n: usize, + k: usize, + alpha: T, + _ctx: &ExecContext, +) -> strided_view::Result<()> +where + T: ScalarBase + strided_view::ElementOpApply, +{ + let mut batch = crate::util::MultiIndex::new(batch_dims); + while batch.next().is_some() { + let a_base = batch.offset(a.batch_strides()); + let b_base = batch.offset(b.batch_strides()); + let c_base = batch.offset(c.batch_strides()); + for i in 0..m { + for j in 0..n { + let mut acc = T::zero(); + for l in 0..k { + let mut av = unsafe { + *a.ptr().offset( + a_base + i as isize * a.row_stride() + l as isize * a.col_stride(), + ) + }; + let mut bv = unsafe { + *b.ptr().offset( + b_base + l as isize * b.row_stride() + j as isize * b.col_stride(), + ) + }; + if a.conj() { + av = strided_view::Conj::apply(av); + } + if b.conj() { + bv = strided_view::Conj::apply(bv); + } + acc = acc + av * bv; + } + let offset = c_base + i as isize * c.row_stride() + j as isize * c.col_stride(); + unsafe { + c.ptr().offset(offset).write(MaybeUninit::new(alpha * acc)); + } + } + } + } + Ok(()) +} + +#[cfg(any(feature = "blas", feature = "blas-inject"))] +fn zero_raw_uninit(dest: &mut RawStridedMut<'_, MaybeUninit>) -> Result<()> { + fn visit( + dest: &mut RawStridedMut<'_, MaybeUninit>, + dims: &[usize], + strides: &[isize], + axis: usize, + offset: isize, + ) -> Result<()> { + if axis == dims.len() { + let relative = offset + .checked_sub(dest.offset()) + .ok_or(strided_view::StridedError::OffsetOverflow)?; + unsafe { + dest.as_mut_ptr() + .offset(relative) + .write(MaybeUninit::new(T::zero())); + } + return Ok(()); + } + for i in 0..dims[axis] { + let next = checked_offset(offset, i, strides[axis])?; + visit(dest, dims, strides, axis + 1, next)?; + } + Ok(()) + } + visit(dest, dest.dims(), dest.strides(), 0, dest.offset()) +} + +#[cfg(any(feature = "blas", feature = "blas-inject"))] +fn bgemm_raw_backend( + mut dest: &mut RawStridedMut<'_, MaybeUninit>, + a: &RawStridedRef<'_, T>, + b: &RawStridedRef<'_, T>, + _n_batch: usize, + n_lo: usize, + n_ro: usize, + n_sum: usize, + alpha: T, + ctx: &ExecContext, +) -> Result<()> +where + T: ScalarBase + strided_view::ElementOpApply, + B: crate::backend::Backend + crate::backend::OverwriteBackend, +{ + let (groups, m, k, n) = preflight_raw_bgemm(&mut dest, a, b, _n_batch, n_lo, n_ro, n_sum)?; + if dest.dims().iter().any(|&d| d == 0) { + return Ok(()); + } + let sum_dims = &a.dims()[n_lo..groups.a_sum_end]; + let batch_dims = &a.dims()[groups.a_sum_end..]; + if sum_dims.iter().any(|&d| d == 0) { + zero_raw_uninit(&mut dest)?; + return Ok(()); + } + let a_op = crate::contiguous::prepare_input_raw( + a, + n_lo, + n_sum, + false, + B::REQUIRES_UNIT_STRIDE, + true, + None, + )?; + let b_op = crate::contiguous::prepare_input_raw( + b, + n_sum, + n_ro, + false, + B::REQUIRES_UNIT_STRIDE, + true, + None, + )?; + let mut c_op = crate::contiguous::prepare_output_raw_uninit( + &mut dest, + n_lo, + n_ro, + B::REQUIRES_UNIT_STRIDE, + )?; + B::bgemm_contiguous_overwrite(&mut c_op, &a_op, &b_op, batch_dims, m, n, k, alpha, ctx)?; + c_op.finalize()?; + Ok(()) +} + +fn checked_offset(offset: isize, index: usize, stride: isize) -> Result { + let term = (index as isize) + .checked_mul(stride) + .ok_or(strided_view::StridedError::OffsetOverflow)?; + offset + .checked_add(term) + .ok_or(strided_view::StridedError::OffsetOverflow) + .map_err(Into::into) +} + +fn visit_offsets( + dims: &[usize], + strides: &[isize], + axis: usize, + offset: isize, + seen: &mut HashSet, +) -> Result<()> { + if axis == dims.len() { + if !seen.insert(offset) { + return Err(strided_view::StridedError::NonInjectiveOutputLayout.into()); + } + return Ok(()); + } + for index in 0..dims[axis] { + visit_offsets( + dims, + strides, + axis + 1, + checked_offset(offset, index, strides[axis])?, + seen, + )?; + } + Ok(()) +} + +fn validate_output(dest: &mut RawStridedMut<'_, MaybeUninit>) -> Result<()> { + let mut seen = HashSet::new(); + visit_offsets(dest.dims(), dest.strides(), 0, dest.offset(), &mut seen) +} + +fn ranges_overlap(a_ptr: *const T, a_len: usize, b_ptr: *const U, b_len: usize) -> bool { + let a_start = a_ptr as usize; + let b_start = b_ptr as usize; + let a_bytes = a_len.saturating_mul(std::mem::size_of::()); + let b_bytes = b_len.saturating_mul(std::mem::size_of::()); + let a_end = a_start.saturating_add(a_bytes); + let b_end = b_start.saturating_add(b_bytes); + a_start < b_end && b_start < a_end +} + +fn validate_no_overlap( + dest: &mut RawStridedMut<'_, MaybeUninit>, + a: &StridedView<'_, T, OpA>, + b: &StridedView<'_, T, OpB>, +) -> Result<()> +where + T: Copy, + OpA: ElementOp, + OpB: ElementOp, +{ + let d = dest.data_mut(); + if ranges_overlap(d.as_ptr(), d.len(), a.data().as_ptr(), a.data().len()) + || ranges_overlap(d.as_ptr(), d.len(), b.data().as_ptr(), b.data().len()) + { + return Err(strided_view::StridedError::OverlappingInputOutput { input: 0 }.into()); + } + Ok(()) +} + +/// Complete raw GEMM preflight, before labels, temporaries, or provider work. +fn preflight_raw_bgemm( + dest: &mut RawStridedMut<'_, MaybeUninit>, + a: &RawStridedRef<'_, T>, + b: &RawStridedRef<'_, T>, + n_batch: usize, + n_lo: usize, + n_ro: usize, + n_sum: usize, +) -> Result<(crate::raw_bgemm::BgemmGroupLayout, usize, usize, usize)> { + let groups = crate::raw_bgemm::checked_bgemm_group_layout(n_batch, n_lo, n_ro, n_sum)?; + crate::raw_bgemm::validate_bgemm_shapes(dest, a, b, n_batch, n_lo, n_ro, n_sum)?; + let av: StridedView<'_, T> = + unsafe { StridedView::new_unchecked(a.data(), a.dims(), a.strides(), a.offset()) }; + let bv: StridedView<'_, T> = + unsafe { StridedView::new_unchecked(b.data(), b.dims(), b.strides(), b.offset()) }; + validate_output(dest)?; + validate_no_overlap(dest, &av, &bv)?; + let m = a.dims()[..n_lo] + .iter() + .try_fold(1usize, |v, &d| v.checked_mul(d)) + .ok_or(strided_view::StridedError::OffsetOverflow)? + .max(1); + let k = a.dims()[n_lo..groups.a_sum_end] + .iter() + .try_fold(1usize, |v, &d| v.checked_mul(d)) + .ok_or(strided_view::StridedError::OffsetOverflow)? + .max(1); + let n = b.dims()[n_sum..groups.b_ro_end] + .iter() + .try_fold(1usize, |v, &d| v.checked_mul(d)) + .ok_or(strided_view::StridedError::OffsetOverflow)? + .max(1); + #[cfg(any(feature = "blas", feature = "blas-inject"))] + for value in [m, k, n] { + i32::try_from(value).map_err(|_| strided_view::StridedError::OffsetOverflow)?; + } + Ok((groups, m, k, n)) +} + +fn validate_labels( + plan: &Einsum2Plan, + dest: &RawStridedMut<'_, MaybeUninit>, + a: &StridedView<'_, T, OpA>, + b: &StridedView<'_, T, OpB>, + ic: &[ID], + ia: &[ID], + ib: &[ID], +) -> Result<()> +where + T: Copy, + OpA: ElementOp, + OpB: ElementOp, + ID: AxisId, +{ + if ia.len() != a.dims().len() || ib.len() != b.dims().len() || ic.len() != dest.dims().len() { + return Err(EinsumError::OutputShapeMismatch { + expected: vec![ic.len()], + got: vec![dest.dims().len()], + }); + } + let dim = |labels: &[ID], dims: &[usize], id: &ID| { + labels.iter().position(|x| x == id).map(|i| dims[i]) + }; + for (axis, id) in ic.iter().enumerate() { + let expected = dim(ia, a.dims(), id).or_else(|| dim(ib, b.dims(), id)); + if expected != Some(dest.dims()[axis]) { + return Err(EinsumError::OutputShapeMismatch { + expected: ic + .iter() + .map(|x| { + dim(ia, a.dims(), x) + .or_else(|| dim(ib, b.dims(), x)) + .unwrap_or(0) + }) + .collect(), + got: dest.dims().to_vec(), + }); + } + } + for id in plan.batch.iter().chain(plan.sum.iter()) { + let da = dim(ia, a.dims(), id).ok_or_else(|| { + EinsumError::InvalidDotGeneralConfig(format!( + "planned axis {:?} is absent from lhs", + id + )) + })?; + let db = dim(ib, b.dims(), id).ok_or_else(|| { + EinsumError::InvalidDotGeneralConfig(format!( + "planned axis {:?} is absent from rhs", + id + )) + })?; + if da != db { + return Err(EinsumError::DimensionMismatch { + axis: format!("{:?}", id), + dim_a: da, + dim_b: db, + }); + } + } + Ok(()) +} + +#[cfg(all( + not(any(feature = "blas", feature = "blas-inject")), + not(feature = "faer") +))] +fn visit_sum( + sum_ids: &[ID], + axis: usize, + a_idx: &mut [usize], + b_idx: &mut [usize], + ia: &[ID], + ib: &[ID], + a: &StridedView<'_, T, OpA>, + b: &StridedView<'_, T, OpB>, + acc: &mut T, +) where + T: ScalarBase, + OpA: ElementOp, + OpB: ElementOp, + ID: AxisId, +{ + if axis == sum_ids.len() { + *acc = *acc + a.get(a_idx) * b.get(b_idx); + return; + } + let id = &sum_ids[axis]; + let ai = ia.iter().position(|x| x == id); + let bi = ib.iter().position(|x| x == id); + let dim = ai + .map(|i| a.dims()[i]) + .or_else(|| bi.map(|i| b.dims()[i])) + .unwrap_or(0); + for i in 0..dim { + if let Some(ai) = ai { + a_idx[ai] = i; + } + if let Some(bi) = bi { + b_idx[bi] = i; + } + visit_sum(sum_ids, axis + 1, a_idx, b_idx, ia, ib, a, b, acc); + } +} + +#[cfg(all( + not(any(feature = "blas", feature = "blas-inject")), + not(feature = "faer") +))] +fn visit_output( + axis: usize, + out_idx: &mut [usize], + dest: &mut RawStridedMut<'_, MaybeUninit>, + a_idx: &mut [usize], + b_idx: &mut [usize], + ic: &[ID], + ia: &[ID], + ib: &[ID], + reduction_ids: &[ID], + a: &StridedView<'_, T, OpA>, + b: &StridedView<'_, T, OpB>, + alpha: T, +) -> Result<()> +where + T: ScalarBase, + OpA: ElementOp, + OpB: ElementOp, + ID: AxisId, +{ + if axis == out_idx.len() { + for (pos, id) in ic.iter().enumerate() { + if let Some(ai) = ia.iter().position(|x| x == id) { + a_idx[ai] = out_idx[pos]; + } + if let Some(bi) = ib.iter().position(|x| x == id) { + b_idx[bi] = out_idx[pos]; + } + } + let mut value = T::zero(); + visit_sum(reduction_ids, 0, a_idx, b_idx, ia, ib, a, b, &mut value); + let mut offset = dest.offset(); + for (&idx, &stride) in out_idx.iter().zip(dest.strides()) { + offset = checked_offset(offset, idx, stride)?; + } + let relative = offset + .checked_sub(dest.offset()) + .ok_or(strided_view::StridedError::OffsetOverflow)?; + unsafe { + dest.as_mut_ptr() + .offset(relative) + .write(MaybeUninit::new(alpha * value)) + }; + return Ok(()); + } + for i in 0..dest.dims()[axis] { + out_idx[axis] = i; + visit_output( + axis + 1, + out_idx, + dest, + a_idx, + b_idx, + ic, + ia, + ib, + reduction_ids, + a, + b, + alpha, + )?; + } + Ok(()) +} + +/// Compute an einsum into a genuinely uninitialized destination. +#[allow(clippy::too_many_arguments)] +#[cfg(not(any(feature = "blas", feature = "blas-inject")))] +pub fn einsum2_into_uninit( + dest: &mut RawStridedMut<'_, MaybeUninit>, + a: &StridedView<'_, T, OpA>, + b: &StridedView<'_, T, OpB>, + ic: &[ID], + ia: &[ID], + ib: &[ID], + alpha: T, + _ctx: &ExecContext, +) -> Result<()> +where + T: ScalarBase, + OpA: ElementOp, + OpB: ElementOp, + ID: AxisId, +{ + let plan = Einsum2Plan::new(ia, ib, ic)?; + validate_labels(&plan, dest, a, b, ic, ia, ib)?; + validate_output(dest)?; + validate_no_overlap(dest, a, b)?; + #[cfg(feature = "faer")] + { + let _ = alpha; + return Err(EinsumError::Unsupported( + "Faer does not yet expose a MaybeUninit-safe overwrite GEMM API; see strided-rs#195" + .to_owned(), + )); + } + #[cfg(not(feature = "faer"))] + { + if dest.dims().iter().any(|&d| d == 0) { + return Ok(()); + } + let mut out_idx = vec![0; dest.dims().len()]; + let mut a_idx = vec![0; a.dims().len()]; + let mut b_idx = vec![0; b.dims().len()]; + let mut reduction_ids = plan.sum.clone(); + for id in ia { + if !ic.contains(id) && !reduction_ids.contains(id) { + reduction_ids.push(id.clone()); + } + } + for id in ib { + if !ic.contains(id) && !reduction_ids.contains(id) { + reduction_ids.push(id.clone()); + } + } + visit_output( + 0, + &mut out_idx, + dest, + &mut a_idx, + &mut b_idx, + ic, + ia, + ib, + &reduction_ids, + a, + b, + alpha, + )?; + Ok(()) + } +} + +/// BLAS-backed overwrite path. All public validation happens before the +/// canonical descriptors are prepared or a temporary is allocated. +#[allow(clippy::too_many_arguments)] +#[cfg(any(feature = "blas", feature = "blas-inject"))] +pub fn einsum2_into_uninit( + dest: &mut RawStridedMut<'_, MaybeUninit>, + a: &StridedView<'_, T, OpA>, + b: &StridedView<'_, T, OpB>, + ic: &[ID], + ia: &[ID], + ib: &[ID], + alpha: T, + ctx: &ExecContext, +) -> Result<()> +where + T: crate::Scalar, + OpA: ElementOp + 'static, + OpB: ElementOp + 'static, + ID: AxisId, +{ + let plan = Einsum2Plan::new(ia, ib, ic)?; + validate_labels(&plan, dest, a, b, ic, ia, ib)?; + validate_output(dest)?; + validate_no_overlap(dest, a, b)?; + if dest.dims().iter().any(|&d| d == 0) { + return Ok(()); + } + + let left_trace = crate::trace::find_trace_indices(ia, ib, ic); + let (a_buf, conj_a) = if !left_trace.is_empty() { + ( + Some(crate::trace::reduce_trace_axes(a, &left_trace)?), + false, + ) + } else { + (None, crate::op_is_conj::()) + }; + let a_view: StridedView<'_, T> = match a_buf.as_ref() { + Some(buf) => buf.view(), + None => StridedView::new(a.data(), a.dims(), a.strides(), a.offset())?, + }; + let right_trace = crate::trace::find_trace_indices(ib, ia, ic); + let (b_buf, conj_b) = if !right_trace.is_empty() { + ( + Some(crate::trace::reduce_trace_axes(b, &right_trace)?), + false, + ) + } else { + (None, crate::op_is_conj::()) + }; + let b_view: StridedView<'_, T> = match b_buf.as_ref() { + Some(buf) => buf.view(), + None => StridedView::new(b.data(), b.dims(), b.strides(), b.offset())?, + }; + let a_perm = a_view.permute(&plan.left_perm)?; + let b_perm = b_view.permute(&plan.right_perm)?; + let c_dims: Vec = plan + .c_to_internal_perm + .iter() + .map(|&axis| dest.dims()[axis]) + .collect(); + let c_strides: Vec = plan + .c_to_internal_perm + .iter() + .map(|&axis| dest.strides()[axis]) + .collect(); + let dest_offset = dest.offset(); + let mut c_perm = RawStridedMut::new(dest.data_mut(), &c_dims, &c_strides, dest_offset)?; + let a_raw = RawStridedRef::new( + a_perm.data(), + a_perm.dims(), + a_perm.strides(), + a_perm.offset(), + )?; + let b_raw = RawStridedRef::new( + b_perm.data(), + b_perm.dims(), + b_perm.strides(), + b_perm.offset(), + )?; + let materialize = crate::make_conj_fn::(); + // BLAS has no conjugation flag. Materialize conjugation before preparing + // the raw backend operands, while retaining the same preflight contract. + if conj_a || conj_b { + let av = if conj_a { + let mut mapped = + unsafe { strided_view::StridedArray::::col_major_uninit(a_perm.dims()) }; + strided_kernel::map_into(&mut mapped.view_mut(), &a_perm, materialize.unwrap())?; + mapped + } else { + strided_view::StridedArray::from_parts( + a_perm.data().to_vec(), + a_perm.dims(), + a_perm.strides(), + a_perm.offset(), + )? + }; + let bv = if conj_b { + let mut mapped = + unsafe { strided_view::StridedArray::::col_major_uninit(b_perm.dims()) }; + strided_kernel::map_into(&mut mapped.view_mut(), &b_perm, materialize.unwrap())?; + mapped + } else { + strided_view::StridedArray::from_parts( + b_perm.data().to_vec(), + b_perm.dims(), + b_perm.strides(), + b_perm.offset(), + )? + }; + let ar = RawStridedRef::new(av.data(), av.dims(), av.strides(), av.view().offset())?; + let br = RawStridedRef::new(bv.data(), bv.dims(), bv.strides(), bv.view().offset())?; + return bgemm_raw_backend::( + &mut c_perm, + &ar, + &br, + plan.batch.len(), + plan.lo.len(), + plan.ro.len(), + plan.sum.len(), + alpha, + ctx, + ); + } + bgemm_raw_backend::( + &mut c_perm, + &a_raw, + &b_raw, + plan.batch.len(), + plan.lo.len(), + plan.ro.len(), + plan.sum.len(), + alpha, + ctx, + ) +} + +/// Owned-input variant of [`einsum2_into_uninit`]. +#[allow(clippy::too_many_arguments)] +#[cfg(not(any(feature = "blas", feature = "blas-inject")))] +pub fn einsum2_into_owned_uninit( + dest: &mut RawStridedMut<'_, MaybeUninit>, + a: strided_view::StridedArray, + b: strided_view::StridedArray, + ic: &[ID], + ia: &[ID], + ib: &[ID], + alpha: T, + ctx: &ExecContext, +) -> Result<()> +where + T: ScalarBase + strided_view::ElementOpApply, + ID: AxisId, +{ + einsum2_into_uninit(dest, &a.view(), &b.view(), ic, ia, ib, alpha, ctx) +} + +/// Owned-input variant for BLAS-backed overwrite execution. +#[allow(clippy::too_many_arguments)] +#[cfg(any(feature = "blas", feature = "blas-inject"))] +pub fn einsum2_into_owned_uninit( + dest: &mut RawStridedMut<'_, MaybeUninit>, + a: strided_view::StridedArray, + b: strided_view::StridedArray, + ic: &[ID], + ia: &[ID], + ib: &[ID], + alpha: T, + ctx: &ExecContext, +) -> Result<()> +where + T: crate::Scalar, + ID: AxisId, +{ + einsum2_into_uninit(dest, &a.view(), &b.view(), ic, ia, ib, alpha, ctx) +} + +/// Canonical raw overwrite-only GEMM entry point. +#[allow(clippy::too_many_arguments)] +pub fn bgemm_raw_strided_into_uninit( + dest: &mut RawStridedMut<'_, MaybeUninit>, + a: &RawStridedRef<'_, T>, + b: &RawStridedRef<'_, T>, + n_batch: usize, + n_lo: usize, + n_ro: usize, + n_sum: usize, + alpha: T, + ctx: &ExecContext, +) -> Result<()> +where + T: crate::Scalar, +{ + // This must precede label construction and all backend-specific + // materialization/allocation. It also validates shape agreement, output + // injectivity, conservative aliasing, checked products, and BLAS sizes. + let (groups, _, _, _) = preflight_raw_bgemm(dest, a, b, n_batch, n_lo, n_ro, n_sum)?; + let mut labels = Vec::with_capacity(groups.label_len); + labels.extend((0..groups.c_rank).map(|x| x)); + #[cfg(not(any(feature = "blas", feature = "blas-inject")))] + let ic = labels[..groups.c_rank].to_vec(); + #[cfg(not(any(feature = "blas", feature = "blas-inject")))] + let sum_start = groups.c_rank; + #[cfg(not(any(feature = "blas", feature = "blas-inject")))] + let ia = (0..n_lo) + .chain(sum_start..groups.label_len) + .chain(groups.c_ro_end..groups.c_rank) + .collect::>(); + #[cfg(not(any(feature = "blas", feature = "blas-inject")))] + let ib = (sum_start..groups.label_len) + .chain(n_lo..groups.c_ro_end) + .chain(groups.c_ro_end..groups.c_rank) + .collect::>(); + #[cfg(not(any(feature = "blas", feature = "blas-inject")))] + let av: StridedView<'_, T> = + unsafe { StridedView::new_unchecked(a.data(), a.dims(), a.strides(), a.offset()) }; + #[cfg(not(any(feature = "blas", feature = "blas-inject")))] + let bv: StridedView<'_, T> = + unsafe { StridedView::new_unchecked(b.data(), b.dims(), b.strides(), b.offset()) }; + #[cfg(any(feature = "blas", feature = "blas-inject"))] + { + return bgemm_raw_backend::( + dest, a, b, n_batch, n_lo, n_ro, n_sum, alpha, ctx, + ); + } + #[cfg(not(any(feature = "blas", feature = "blas-inject")))] + einsum2_into_uninit(dest, &av, &bv, &ic, &ia, &ib, alpha, ctx) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::panic::{catch_unwind, AssertUnwindSafe}; + use strided_view::StridedArray; + + #[cfg(not(feature = "faer"))] + #[test] + fn matrix_product_writes_uninitialized_destination() { + let a = StridedArray::from_fn_row_major(&[2, 3], |idx| (idx[0] * 3 + idx[1] + 1) as f64); + let b = StridedArray::from_fn_row_major(&[3, 2], |idx| (idx[0] * 2 + idx[1] + 1) as f64); + let mut storage = vec![MaybeUninit::::uninit(); 4]; + let dims = [2, 2]; + let strides = [2, 1]; + let mut c = RawStridedMut::new(&mut storage, &dims, &strides, 0).unwrap(); + einsum2_into_uninit( + &mut c, + &a.view(), + &b.view(), + &['i', 'k'], + &['i', 'j'], + &['j', 'k'], + 1.0, + &ExecContext::serial(), + ) + .unwrap(); + let values: Vec = storage + .into_iter() + .map(|x| unsafe { x.assume_init() }) + .collect(); + assert_eq!(values, vec![22.0, 28.0, 49.0, 64.0]); + } + + #[test] + fn rejects_noninjective_destination_before_writing() { + let a = StridedArray::from_fn_col_major(&[2], |_| 1.0f64); + let b = StridedArray::from_fn_col_major(&[2], |_| 2.0f64); + let mut storage = vec![MaybeUninit::::uninit(); 1]; + let dims = [2]; + let strides = [0]; + let mut c = RawStridedMut::new(&mut storage, &dims, &strides, 0).unwrap(); + let err = einsum2_into_uninit( + &mut c, + &a.view(), + &b.view(), + &['i'], + &['i'], + &['i'], + 1.0, + &ExecContext::serial(), + ) + .unwrap_err(); + assert!(matches!(err, crate::EinsumError::Strided(_))); + } + + #[cfg(feature = "faer")] + #[test] + fn faer_uninit_gemm_reports_typed_unsupported_error() { + let a = StridedArray::from_fn_col_major(&[1, 1], |_| 1.0f64); + let b = StridedArray::from_fn_col_major(&[1, 1], |_| 1.0f64); + let mut storage = vec![MaybeUninit::::uninit()]; + let dims = [1usize, 1]; + let strides = [1isize, 1]; + let mut c = RawStridedMut::new(&mut storage, &dims, &strides, 0).unwrap(); + let err = einsum2_into_uninit( + &mut c, + &a.view(), + &b.view(), + &['i', 'k'], + &['i', 'j'], + &['j', 'k'], + 1.0, + &ExecContext::serial(), + ) + .unwrap_err(); + assert!(matches!(err, crate::EinsumError::Unsupported(_))); + } + + #[cfg(not(feature = "faer"))] + #[test] + fn noncontiguous_output_is_written_back_after_overwrite() { + let a = StridedArray::from_fn_col_major(&[2, 2, 2], |idx| { + (1 + idx[0] + 2 * idx[1] + 4 * idx[2]) as f64 + }); + let b = StridedArray::from_fn_col_major(&[2, 2], |idx| (1 + idx[0] + 2 * idx[1]) as f64); + let mut storage = vec![MaybeUninit::::uninit(); 8]; + let dims = [2usize, 2, 2]; + let strides = [1isize, 4, 2]; + let mut c = RawStridedMut::new(&mut storage, &dims, &strides, 0).unwrap(); + einsum2_into_uninit( + &mut c, + &a.view(), + &b.view(), + &['i', 'j', 'k'], + &['i', 'j', 'l'], + &['l', 'k'], + 1.0, + &ExecContext::serial(), + ) + .unwrap(); + let values: Vec = storage + .into_iter() + .map(|x| unsafe { x.assume_init() }) + .collect(); + assert_eq!(values, vec![11.0, 14.0, 23.0, 30.0, 17.0, 20.0, 37.0, 44.0]); + } + + #[test] + fn raw_uninit_gemm_rejects_wrapping_group_partition_without_panicking() { + let a = StridedArray::from_fn_col_major(&[1], |_| 1.0f64); + let b = StridedArray::from_fn_col_major(&[1, 1], |_| 1.0f64); + let mut storage = vec![MaybeUninit::::uninit()]; + let c_dims: [usize; 0] = []; + let c_strides: [isize; 0] = []; + let mut c = RawStridedMut::new(&mut storage, &c_dims, &c_strides, 0).unwrap(); + let a_raw = RawStridedRef::new(a.data(), a.dims(), a.strides(), a.view().offset()).unwrap(); + let b_raw = RawStridedRef::new(b.data(), b.dims(), b.strides(), b.view().offset()).unwrap(); + + let result = catch_unwind(AssertUnwindSafe(|| { + bgemm_raw_strided_into_uninit( + &mut c, + &a_raw, + &b_raw, + 1, + usize::MAX, + 0, + 1, + 1.0, + &ExecContext::serial(), + ) + })); + + assert!(result.is_ok(), "invalid group partition must not panic"); + assert!(result.unwrap().is_err()); + } +} diff --git a/strided-einsum2/tests/blas_inject_fallback.rs b/strided-einsum2/tests/blas_inject_fallback.rs index df5d182..6816fd1 100644 --- a/strided-einsum2/tests/blas_inject_fallback.rs +++ b/strided-einsum2/tests/blas_inject_fallback.rs @@ -1,10 +1,79 @@ #![cfg(feature = "blas-inject")] -use strided_einsum2::einsum2_into; +use std::os::raw::c_char; +use std::sync::Once; + +use cblas_inject::{register_dgemm, BlasInt32}; +use num_complex::Complex64; +use std::mem::MaybeUninit; +use strided_einsum2::{einsum2_into, einsum2_into_uninit}; +use strided_kernel::ExecContext; use strided_view::StridedArray; +static REGISTER: Once = Once::new(); + +unsafe extern "C" fn test_dgemm( + transa: *const c_char, + transb: *const c_char, + m: *const BlasInt32, + n: *const BlasInt32, + k: *const BlasInt32, + alpha: *const f64, + a: *const f64, + lda: *const BlasInt32, + b: *const f64, + ldb: *const BlasInt32, + beta: *const f64, + c: *mut f64, + ldc: *const BlasInt32, +) { + let (m, n, k, lda, ldb, ldc) = ( + *m as usize, + *n as usize, + *k as usize, + *lda as usize, + *ldb as usize, + *ldc as usize, + ); + let ta = (*transa as u8).to_ascii_uppercase() as char; + let tb = (*transb as u8).to_ascii_uppercase() as char; + let av = |row: usize, col: usize| unsafe { + if ta == 'N' { + *a.add(row + col * lda) + } else { + *a.add(col + row * lda) + } + }; + let bv = |row: usize, col: usize| unsafe { + if tb == 'N' { + *b.add(row + col * ldb) + } else { + *b.add(col + row * ldb) + } + }; + for col in 0..n { + for row in 0..m { + let mut value = 0.0; + for inner in 0..k { + value += av(row, inner) * bv(inner, col); + } + let out = c.add(row + col * ldc); + if *beta == 0.0 { + *out = *alpha * value; + } else { + *out = *alpha * value + *beta * *out; + } + } + } +} + +fn register_test_provider() { + REGISTER.call_once(|| unsafe { register_dgemm(test_dgemm) }); +} + #[test] -fn test_blas_inject_works_without_manual_registration() { +fn test_blas_inject_works_with_explicit_registration() { + register_test_provider(); let a = StridedArray::::from_fn_row_major(&[2, 2], |idx| { [[1.0, 2.0], [3.0, 4.0]][idx[0]][idx[1]] }); @@ -30,3 +99,71 @@ fn test_blas_inject_works_without_manual_registration() { assert_eq!(c.get(&[1, 0]), 43.0); assert_eq!(c.get(&[1, 1]), 50.0); } + +#[test] +fn test_blas_inject_uninitialized_overwrite_uses_registered_provider() { + register_test_provider(); + let a = StridedArray::::from_fn_col_major(&[2, 2], |idx| (idx[0] + 2 * idx[1] + 1) as f64); + let b = StridedArray::::from_fn_col_major(&[2, 2], |idx| (idx[0] + 2 * idx[1] + 5) as f64); + let mut storage = vec![MaybeUninit::::uninit(); 4]; + let dims = [2usize, 2]; + let strides = [1isize, 2]; + let mut c = strided_view::RawStridedMut::new(&mut storage, &dims, &strides, 0).unwrap(); + einsum2_into_uninit( + &mut c, + &a.view(), + &b.view(), + &['i', 'j'], + &['i', 'k'], + &['k', 'j'], + 1.0, + &ExecContext::serial(), + ) + .unwrap(); + let values: Vec = storage + .into_iter() + .map(|x| unsafe { x.assume_init() }) + .collect(); + assert_eq!(values, vec![23.0, 34.0, 31.0, 46.0]); +} + +#[test] +fn test_blas_inject_zgemm_overwrite_does_not_read_poisoned_c() { + let a = StridedArray::::from_fn_col_major(&[2, 2], |idx| { + Complex64::new((idx[0] + 2 * idx[1] + 1) as f64, 0.0) + }); + let b = StridedArray::::from_fn_col_major(&[2, 2], |idx| { + Complex64::new((idx[0] + 2 * idx[1] + 5) as f64, 0.0) + }); + let poison = Complex64::new(f64::NAN, f64::NAN); + let mut storage = vec![MaybeUninit::new(poison); 4]; + let dims = [2usize, 2]; + let strides = [1isize, 2]; + let mut c = strided_view::RawStridedMut::new(&mut storage, &dims, &strides, 0).unwrap(); + + einsum2_into_uninit( + &mut c, + &a.view(), + &b.view(), + &['i', 'j'], + &['i', 'k'], + &['k', 'j'], + Complex64::new(1.0, 0.0), + &ExecContext::serial(), + ) + .unwrap(); + + let values: Vec = storage + .into_iter() + .map(|x| unsafe { x.assume_init() }) + .collect(); + assert_eq!( + values, + vec![ + Complex64::new(23.0, 0.0), + Complex64::new(34.0, 0.0), + Complex64::new(31.0, 0.0), + Complex64::new(46.0, 0.0), + ] + ); +} diff --git a/strided-opteinsum/src/expr.rs b/strided-opteinsum/src/expr.rs index 65041e4..52e5851 100644 --- a/strided-opteinsum/src/expr.rs +++ b/strided-opteinsum/src/expr.rs @@ -1,10 +1,18 @@ use std::collections::{BTreeMap, HashMap}; +#[cfg(not(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject")))))] +use std::mem::MaybeUninit; use num_complex::Complex64; #[cfg(test)] use num_traits::Zero; -use strided_einsum2::{einsum2_into, einsum2_into_owned}; +use strided_einsum2::einsum2_into; +#[cfg(any(feature = "faer", feature = "blas", feature = "blas-inject"))] +use strided_einsum2::einsum2_into_owned; +#[cfg(not(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject")))))] +use strided_einsum2::einsum2_into_uninit; use strided_kernel::copy_scale; +#[cfg(not(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject")))))] +use strided_kernel::ExecContext; use strided_view::{StridedArray, StridedViewMut}; use crate::operand::{EinsumOperand, EinsumScalar, StridedData}; @@ -64,7 +72,28 @@ trait PoolOps: EinsumScalar { /// # Safety contract /// The returned array may contain uninitialized data. Callers must write /// every element before reading (e.g. via `einsum2_into` with `beta=0`). - fn pool_acquire(pool: &mut BufferPool, dims: &[usize]) -> StridedArray; + #[cfg(not(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject")))))] + fn pool_acquire( + pool: &mut BufferPool, + dims: &[usize], + ) -> crate::Result>>; + + /// Acquire initialized storage for the Faer compatibility path. Faer + /// cannot currently accept `MaybeUninit` output; reused buffers are + /// already initialized, while fresh buffers are initialized once. + #[cfg(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject"))))] + fn pool_acquire_initialized( + pool: &mut BufferPool, + dims: &[usize], + ) -> crate::Result>; + + /// Convert a completely overwritten array into initialized storage. + /// + /// # Safety + /// The caller must have received `Ok(())` from the overwrite-only kernel + /// for every logical element before calling this method. + #[cfg(not(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject")))))] + unsafe fn assume_initialized(array: StridedArray>) -> StridedArray; /// Release an owned buffer back to the pool for reuse. /// Views are silently dropped (nothing to recycle). @@ -83,16 +112,78 @@ fn take_best_fit(pool: &mut BTreeMap>>, total: usize) -> Op buf } +#[cfg(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject"))))] +fn acquire_initialized( + pool: &mut BTreeMap>>, + dims: &[usize], +) -> crate::Result> { + let total = dims + .iter() + .try_fold(1usize, |total, &dim| total.checked_mul(dim)) + .ok_or(strided_view::StridedError::OffsetOverflow)? + .max(1); + let mut data = take_best_fit(pool, total).unwrap_or_default(); + if data.len() < total { + data.resize_with(total, T::default); + } else { + data.truncate(total); + } + Ok(StridedArray::from_parts( + data, + dims, + &strided_view::col_major_strides(dims), + 0, + )?) +} + impl PoolOps for f64 { - fn pool_acquire(pool: &mut BufferPool, dims: &[usize]) -> StridedArray { - let total: usize = dims.iter().product(); + #[cfg(not(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject")))))] + fn pool_acquire( + pool: &mut BufferPool, + dims: &[usize], + ) -> crate::Result>> { + let total = dims + .iter() + .try_fold(1usize, |total, &dim| total.checked_mul(dim)) + .ok_or(strided_view::StridedError::OffsetOverflow)? + .max(1); // SAFETY: einsum2_into with beta=0 writes every output element before reading. match take_best_fit(&mut pool.f64_pool, total) { - Some(buf) => unsafe { StridedArray::col_major_from_buffer_uninit(buf, dims) }, - None => unsafe { StridedArray::col_major_uninit(dims) }, + Some(buf) => unsafe { + let mut buf = std::mem::ManuallyDrop::new(buf); + let data = Vec::from_raw_parts( + buf.as_mut_ptr().cast::>(), + buf.len(), + buf.capacity(), + ); + Ok(StridedArray::col_major_from_buffer_uninit(data, dims)) + }, + None => Ok(unsafe { StridedArray::col_major_uninit(dims) }), } } + #[cfg(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject"))))] + fn pool_acquire_initialized( + pool: &mut BufferPool, + dims: &[usize], + ) -> crate::Result> { + acquire_initialized(&mut pool.f64_pool, dims) + } + + #[cfg(not(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject")))))] + unsafe fn assume_initialized(array: StridedArray>) -> StridedArray { + let dims = array.dims().to_vec(); + let strides = array.strides().to_vec(); + let offset = 0; + let data = array.into_data(); + let len = data.len(); + let cap = data.capacity(); + let ptr = data.as_ptr().cast_mut().cast::(); + std::mem::forget(data); + StridedArray::from_parts(Vec::from_raw_parts(ptr, len, cap), &dims, &strides, offset) + .expect("validated pool layout") + } + fn pool_release(pool: &mut BufferPool, data: StridedData<'_, f64>) { if let StridedData::Owned(arr) = data { let buf = arr.into_data(); @@ -102,15 +193,54 @@ impl PoolOps for f64 { } impl PoolOps for Complex64 { - fn pool_acquire(pool: &mut BufferPool, dims: &[usize]) -> StridedArray { - let total: usize = dims.iter().product(); + #[cfg(not(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject")))))] + fn pool_acquire( + pool: &mut BufferPool, + dims: &[usize], + ) -> crate::Result>> { + let total = dims + .iter() + .try_fold(1usize, |total, &dim| total.checked_mul(dim)) + .ok_or(strided_view::StridedError::OffsetOverflow)? + .max(1); // SAFETY: einsum2_into with beta=0 writes every output element before reading. match take_best_fit(&mut pool.c64_pool, total) { - Some(buf) => unsafe { StridedArray::col_major_from_buffer_uninit(buf, dims) }, - None => unsafe { StridedArray::col_major_uninit(dims) }, + Some(buf) => unsafe { + let mut buf = std::mem::ManuallyDrop::new(buf); + let data = Vec::from_raw_parts( + buf.as_mut_ptr().cast::>(), + buf.len(), + buf.capacity(), + ); + Ok(StridedArray::col_major_from_buffer_uninit(data, dims)) + }, + None => Ok(unsafe { StridedArray::col_major_uninit(dims) }), } } + #[cfg(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject"))))] + fn pool_acquire_initialized( + pool: &mut BufferPool, + dims: &[usize], + ) -> crate::Result> { + acquire_initialized(&mut pool.c64_pool, dims) + } + + #[cfg(not(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject")))))] + unsafe fn assume_initialized( + array: StridedArray>, + ) -> StridedArray { + let dims = array.dims().to_vec(); + let strides = array.strides().to_vec(); + let data = array.into_data(); + let len = data.len(); + let cap = data.capacity(); + let ptr = data.as_ptr().cast_mut().cast::(); + std::mem::forget(data); + StridedArray::from_parts(Vec::from_raw_parts(ptr, len, cap), &dims, &strides, 0) + .expect("validated pool layout") + } + fn pool_release(pool: &mut BufferPool, data: StridedData<'_, Complex64>) { if let StridedData::Owned(arr) = data { let buf = arr.into_data(); @@ -288,10 +418,67 @@ fn compute_binary_output_ids( out } +/// Faer compatibility path until the upstream overwrite API exists. +#[cfg(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject"))))] +fn eval_pair_alloc( + ld: StridedData<'_, T>, + left_ids: &[char], + rd: StridedData<'_, T>, + right_ids: &[char], + output_ids: &[char], + pool: &mut BufferPool, + size_dict: &HashMap, +) -> crate::Result> { + let out_dims = out_dims_from_ids( + left_ids, + ld.dims(), + right_ids, + rd.dims(), + output_ids, + size_dict, + )?; + let mut c_arr = T::pool_acquire_initialized(pool, &out_dims)?; + let (left, right) = (ld, rd); + match (left, right) { + (StridedData::Owned(a), StridedData::Owned(b)) => { + einsum2_into_owned( + c_arr.view_mut(), + a, + b, + output_ids, + left_ids, + right_ids, + T::one(), + T::zero(), + false, + false, + )?; + } + (left, right) => { + let a_view = left.as_view(); + let b_view = right.as_view(); + einsum2_into( + c_arr.view_mut(), + &a_view, + &b_view, + output_ids, + left_ids, + right_ids, + T::one(), + T::zero(), + )?; + T::pool_release(pool, left); + T::pool_release(pool, right); + } + } + Ok(T::wrap_array(c_arr)) +} + /// Generic inner function for pairwise contraction with buffer pool. /// /// Acquires an output buffer, runs `einsum2_into`, and releases input buffers /// back to the pool. +#[cfg(not(all(feature = "faer", not(any(feature = "blas", feature = "blas-inject")))))] fn eval_pair_alloc( ld: StridedData<'_, T>, left_ids: &[char], @@ -309,41 +496,46 @@ fn eval_pair_alloc( output_ids, size_dict, )?; - let mut c_arr = T::pool_acquire(pool, &out_dims); + let mut c_arr = T::pool_acquire(pool, &out_dims)?; match (ld, rd) { // Preserve ownership so strided-einsum2 can use prepare_input_owned // and avoid extra materialization in prepare_input_view. (StridedData::Owned(a), StridedData::Owned(b)) => { - einsum2_into_owned( - c_arr.view_mut(), + let dims = c_arr.dims().to_vec(); + let strides = c_arr.strides().to_vec(); + let mut out = strided_view::RawStridedMut::new(c_arr.data_mut(), &dims, &strides, 0)?; + strided_einsum2::einsum2_into_owned_uninit( + &mut out, a, b, output_ids, left_ids, right_ids, T::one(), - T::zero(), - false, - false, + &ExecContext::serial(), )?; } (ld, rd) => { let a_view = ld.as_view(); let b_view = rd.as_view(); - einsum2_into( - c_arr.view_mut(), + let dims = c_arr.dims().to_vec(); + let strides = c_arr.strides().to_vec(); + let mut out = strided_view::RawStridedMut::new(c_arr.data_mut(), &dims, &strides, 0)?; + einsum2_into_uninit( + &mut out, &a_view, &b_view, output_ids, left_ids, right_ids, T::one(), - T::zero(), + &ExecContext::serial(), )?; T::pool_release(pool, ld); T::pool_release(pool, rd); } } + let c_arr = unsafe { T::assume_initialized(c_arr) }; Ok(T::wrap_array(c_arr)) } @@ -402,9 +594,21 @@ fn eval_pair_into( match (left_data, right_data) { (StridedData::Owned(a), StridedData::Owned(b)) => { + #[cfg(any(feature = "faer", feature = "blas", feature = "blas-inject"))] einsum2_into_owned( output, a, b, output_ids, left_ids, right_ids, alpha, beta, false, false, )?; + #[cfg(not(any(feature = "faer", feature = "blas", feature = "blas-inject")))] + einsum2_into( + output, + &a.view(), + &b.view(), + output_ids, + left_ids, + right_ids, + alpha, + beta, + )?; } (StridedData::Owned(a), StridedData::View(b)) => { einsum2_into( From 07992fe87066e4ef4d2988686411945856e0e4a7 Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Fri, 31 Jul 2026 10:12:33 +0900 Subject: [PATCH 2/2] test uninitialized GEMM coverage paths --- docs/2026-07-31-issue-188-worklog.md | 6 ++ strided-einsum2/src/contiguous.rs | 63 +++++++++++++++++ strided-einsum2/src/dot_general.rs | 35 +++++++++ strided-einsum2/src/uninit.rs | 102 +++++++++++++++++++++++++++ 4 files changed, 206 insertions(+) diff --git a/docs/2026-07-31-issue-188-worklog.md b/docs/2026-07-31-issue-188-worklog.md index 3d39472..cbb7285 100644 --- a/docs/2026-07-31-issue-188-worklog.md +++ b/docs/2026-07-31-issue-188-worklog.md @@ -29,6 +29,12 @@ Add overwrite-only contraction entry points whose destination is borrowed as ## Verification - `cargo fmt --all` +- Added coverage tests for the Faer-gated validation/naive fallback, the + `dot_general_into_uninit` forwarding boundary, and both direct and + temporary raw-output finalize paths. +- `CARGO_BUILD_JOBS=4 RUSTFLAGS='-C link-arg=-Wl,--threads=1' cargo llvm-cov --workspace --json --output-path /tmp/coverage-188.json` +- `python3 scripts/check-coverage.py /tmp/coverage-188.json` (54/54 files; + `contiguous.rs` 86.98%, `dot_general.rs` 86.53%, `uninit.rs` 82.17%) - `cargo test -p strided-einsum2 --lib` - `cargo test -p strided-opteinsum --lib` - `cargo test -p strided-einsum2 --no-default-features` diff --git a/strided-einsum2/src/contiguous.rs b/strided-einsum2/src/contiguous.rs index 7473a9b..17f508f 100644 --- a/strided-einsum2/src/contiguous.rs +++ b/strided-einsum2/src/contiguous.rs @@ -1166,4 +1166,67 @@ mod tests { let after = pooled_count_for_type::(); assert!(after >= before.saturating_add(1)); } + + #[test] + fn test_output_raw_uninit_direct_finalize() { + let mut storage = vec![MaybeUninit::::uninit(); 6]; + let mut raw = RawStridedMut::new(&mut storage, &[2, 3], &[1, 2], 0).unwrap(); + let op = prepare_output_raw_uninit(&mut raw, 1, 1, false).unwrap(); + + assert!(op.temp.is_none()); + assert!(op.writeback.is_none()); + for j in 0..3 { + for i in 0..2 { + let offset = i as isize * op.row_stride() + j as isize * op.col_stride(); + // SAFETY: the direct layout is injective and every logical + // destination element is written exactly once before finalize. + unsafe { + op.ptr() + .offset(offset) + .write(MaybeUninit::new((10 + i + 2 * j) as f64)); + } + } + } + op.finalize().unwrap(); + drop(raw); + + let values: Vec = storage + .into_iter() + .map(|x| unsafe { x.assume_init() }) + .collect(); + assert_eq!(values, vec![10.0, 11.0, 12.0, 13.0, 14.0, 15.0]); + } + + #[test] + fn test_output_raw_uninit_temp_finalize_writes_back() { + let mut storage = vec![MaybeUninit::::uninit(); 30]; + let mut raw = RawStridedMut::new(&mut storage, &[2, 3, 1], &[10, 1, 1], 0).unwrap(); + let op = prepare_output_raw_uninit(&mut raw, 2, 1, true).unwrap(); + + assert!(op.temp.is_some()); + assert!(op.writeback.is_some()); + for i in 0..6 { + let offset = i as isize * op.row_stride(); + // SAFETY: the dense temporary has one initialized value for every + // logical output element before writeback. + unsafe { + op.ptr() + .offset(offset) + .write(MaybeUninit::new((20 + i) as f64)); + } + } + op.finalize().unwrap(); + drop(raw); + + // SAFETY: finalize wrote every asserted destination element before + // these reads. + unsafe { + assert_eq!(storage[0].assume_init_ref(), &20.0); + assert_eq!(storage[1].assume_init_ref(), &22.0); + assert_eq!(storage[2].assume_init_ref(), &24.0); + assert_eq!(storage[10].assume_init_ref(), &21.0); + assert_eq!(storage[11].assume_init_ref(), &23.0); + assert_eq!(storage[12].assume_init_ref(), &25.0); + } + } } diff --git a/strided-einsum2/src/dot_general.rs b/strided-einsum2/src/dot_general.rs index 9fef69d..c7ca81c 100644 --- a/strided-einsum2/src/dot_general.rs +++ b/strided-einsum2/src/dot_general.rs @@ -307,6 +307,41 @@ pub fn dot_general_into_uninit( ) } +#[cfg(test)] +mod tests { + use super::*; + use std::mem::MaybeUninit; + + #[test] + fn uninit_entry_point_routes_through_einsum_validation() { + let a = strided_view::StridedArray::from_fn_col_major(&[2, 3], |idx| { + (idx[0] + 2 * idx[1] + 1) as f64 + }); + let b = strided_view::StridedArray::from_fn_col_major(&[3, 2], |idx| { + (idx[0] + 3 * idx[1] + 1) as f64 + }); + let mut storage = vec![MaybeUninit::::uninit(); 4]; + let mut c = RawStridedMut::new(&mut storage, &[2, 2], &[2, 1], 0).unwrap(); + let config = DotGeneralConfig { + lhs_contracting_dims: &[1], + rhs_contracting_dims: &[0], + lhs_batch_dims: &[], + rhs_batch_dims: &[], + }; + + let err = dot_general_into_uninit( + &mut c, + &a.view(), + &b.view(), + &config, + 1.0, + &ExecContext::serial(), + ) + .unwrap_err(); + assert!(matches!(err, crate::EinsumError::Unsupported(_))); + } +} + /// Compute `C = alpha * dot_general(A, B) + beta * C` with the naive fallback. #[cfg(not(any(feature = "faer", feature = "blas", feature = "blas-inject")))] pub fn dot_general_into( diff --git a/strided-einsum2/src/uninit.rs b/strided-einsum2/src/uninit.rs index 369ea98..9d0d2de 100644 --- a/strided-einsum2/src/uninit.rs +++ b/strided-einsum2/src/uninit.rs @@ -810,6 +810,108 @@ mod tests { assert!(matches!(err, crate::EinsumError::Unsupported(_))); } + #[cfg(feature = "faer")] + #[test] + fn faer_uninit_gemm_validates_labels_before_backend_selection() { + let a = StridedArray::from_fn_col_major(&[2], |_| 1.0f64); + let b = StridedArray::from_fn_col_major(&[2], |_| 1.0f64); + + let mut rank_storage = vec![MaybeUninit::::uninit(); 2]; + let mut rank_dest = RawStridedMut::new(&mut rank_storage, &[2], &[1], 0).unwrap(); + let rank_err = einsum2_into_uninit( + &mut rank_dest, + &a.view(), + &b.view(), + &['i'], + &['i', 'j'], + &['i'], + 1.0, + &ExecContext::serial(), + ) + .unwrap_err(); + assert!(matches!( + rank_err, + crate::EinsumError::OutputShapeMismatch { .. } + )); + + let mut shape_storage = vec![MaybeUninit::::uninit(); 3]; + let mut shape_dest = RawStridedMut::new(&mut shape_storage, &[3], &[1], 0).unwrap(); + let shape_err = einsum2_into_uninit( + &mut shape_dest, + &a.view(), + &b.view(), + &['i'], + &['i'], + &['i'], + 1.0, + &ExecContext::serial(), + ) + .unwrap_err(); + assert!(matches!( + shape_err, + crate::EinsumError::OutputShapeMismatch { .. } + )); + + let b_mismatched = StridedArray::from_fn_col_major(&[3], |_| 1.0f64); + let mut scalar_storage = vec![MaybeUninit::::uninit()]; + let mut scalar_dest = RawStridedMut::new(&mut scalar_storage, &[], &[], 0).unwrap(); + let dimension_err = einsum2_into_uninit( + &mut scalar_dest, + &a.view(), + &b_mismatched.view(), + &[], + &['i'], + &['i'], + 1.0, + &ExecContext::serial(), + ) + .unwrap_err(); + assert!(matches!( + dimension_err, + crate::EinsumError::DimensionMismatch { .. } + )); + } + + #[cfg(feature = "faer")] + #[test] + fn naive_overwrite_backend_covers_batches_and_conjugation() { + let a = StridedArray::from_fn_col_major(&[2, 2, 2], |idx| { + (1 + idx[0] + 2 * idx[1] + 4 * idx[2]) as f64 + }); + let b = StridedArray::from_fn_col_major(&[2, 2, 2], |idx| { + (1 + idx[0] + 2 * idx[1] + 4 * idx[2]) as f64 + }); + let a_raw = RawStridedRef::new(a.data(), a.dims(), a.strides(), a.view().offset()).unwrap(); + let b_raw = RawStridedRef::new(b.data(), b.dims(), b.strides(), b.view().offset()).unwrap(); + let a_op = + crate::contiguous::prepare_input_raw(&a_raw, 1, 1, true, false, false, None).unwrap(); + let b_op = + crate::contiguous::prepare_input_raw(&b_raw, 1, 1, true, false, false, None).unwrap(); + + let mut storage = vec![MaybeUninit::::uninit(); 8]; + let mut c = RawStridedMut::new(&mut storage, &[2, 2, 2], &[1, 2, 4], 0).unwrap(); + let mut c_op = crate::contiguous::prepare_output_raw_uninit(&mut c, 1, 1, false).unwrap(); + bgemm_contiguous_naive( + &mut c_op, + &a_op, + &b_op, + &[2], + 2, + 2, + 2, + 2.0, + &ExecContext::serial(), + ) + .unwrap(); + c_op.finalize().unwrap(); + + let values: Vec = storage + .into_iter() + .map(|x| unsafe { x.assume_init() }) + .collect(); + assert!(values.iter().all(|value| *value > 0.0)); + } + #[cfg(not(feature = "faer"))] #[test] fn noncontiguous_output_is_written_back_after_overwrite() {