From 2052b0f48f09ecd997d9f17722d940cd57a2e946 Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Thu, 29 Oct 2026 20:09:50 +0900 Subject: [PATCH] matmul: overwrite a MaybeUninit destination (#136) matmul_with_conj_uninit(dst: MatUninitMut, lhs, conj_lhs, rhs, conj_rhs, alpha, par) computes alpha * op(lhs) * op(rhs) into storage whose elements may be uninitialized and returns it as an initialized MatMut, for callers (strided-rs#198) that otherwise fill the output before an overwriting product. Requested upstream as sarah-quinones/faer-rs#269; the view type and the first tests follow the same author's upstream PR #270. The destination's prior contents are never read and no reference to an unwritten element is formed. The native types on x86_64 with AVX2 or AVX-512 go to t4a-private-gemm-x86 with Accum::Replace (on a team, through gemm_team); other scalars write each dot product through ptr::write; an empty depth writes zeros. The gemm and nano-gemm crates never see uninitialized storage (gemm's GEMV forms a slice of the destination, its AMX kernel copies partial tiles, nano-gemm scales by beta), so elsewhere the native types are zero-filled first. Existing MatMut APIs are unchanged. Co-Authored-By: Claude Opus 5.5 --- PATCHSET.md | 1 + docs/worklogs/2026-10-29-matmul-uninit.md | 37 ++ faer/src/lib.rs | 4 +- faer/src/linalg/matmul/mod.rs | 450 +++++++++++++++++++++- faer/src/mat/matmut.rs | 149 +++++++ faer/src/mat/mod.rs | 21 +- 6 files changed, 657 insertions(+), 5 deletions(-) create mode 100644 docs/worklogs/2026-10-29-matmul-uninit.md diff --git a/PATCHSET.md b/PATCHSET.md index b2cb793c..49c85612 100644 --- a/PATCHSET.md +++ b/PATCHSET.md @@ -133,6 +133,7 @@ gate has changed and two kinds of adoption are now in the ledger. | Team dispatch cost read from threadteam | `utils::cost::team_section_cost` and `MEASURED_MAX_WIDTH` (faer's own linear extrapolation above width 8, #97) are deleted: the extrapolation lives in threadteam (`WidthCosts` 0.1.1, linear above a row's largest anchor; 0.1.2, a measured Zen2 team row for `x86_64-avx2` at widths 2 to 64). `ParShape::dispatch_cost` reads `DispatchCosts::cost(Team, w)`; `par_pays_time`, `team::kernel::kernel_width` and the tiled team product's floor ask `TeamExec::dispatch_cost(w)`. `kernel_width` searches from the full width down instead of stopping at the first narrow width that does not pay (measured rows fall with the width). Requirement `threadteam = "0.1.2"`. Decisions change at widths 7 and 8 for rows whose largest anchor is 6, at width 8 on `x86_64-avx512` (its 0.98 us anchor) and at every width on `x86_64-avx2`. | EXPERIMENT (pending the maintainer's acceptance decision) | tensor4all/threadteam-rs#6, tensor4all/faer#91 | No timing (gate inputs only; acceptance run left to the maintainer). Tests: `utils::cost::tests::the_team_cost_is_threadteams_table_at_every_width`, `utils::team::tests::kernel_width_finds_the_widest_paying_width_on_a_falling_row`. Record: [worklog](docs/worklogs/2026-10-29-team-cost-from-threadteam.md). | | Every region at the widest width that pays | The full-width gates (`par_pays`, `ParThreshold::region_pays`, `ParShape::pays_full` at the call sites, the block apply's `n_tasks`) chose between the token's whole degree and one thread, so on a wide team a region that pays at 8 or 16 ran serially. Every region now asks `utils::cost::par_width` / `ParThreshold::region_width` for the widest width that pays (every width searched, falling rows included; the GEMM kernel's `kernel_width` is the same rule) and runs on a sub-team of that width (`utils::thread::with_width`, a split into one child: no worker dispatch). Splits whose task count is capped (triangular solve jobs of 64 columns, block apply columns, `upgrade_householder_factor`'s rows per participant) search inside their cap. The factorizations that latched `par = Par::Seq` for the rest of the call (Hessenberg, tridiagonal, unblocked bidiagonal, full pivoting LU, Bunch-Kaufman, column pivoting QR) decide per step; `Never` still makes the whole call sequential; the post-loop compact-factor upgrades of the Hessenberg and tridiagonal reductions now ask their own gate (they ran under the last step's latch, i.e. sequentially). The blocked bidiagonal panel gate is a floor (sequential or the whole `par`), so the panel's products choose their own widths. The Bunch-Kaufman blocked swap loop, split over the whole degree without a gate, now has one (row swaps at the LU swap rate). `Fixed` keeps its rules (the whole degree above the floor). `par_pays`, `region_pays` and the Seq-only GEMM predicate are removed; `par-trace` records the chosen width. | EXPERIMENT (pending the maintainer's acceptance run on broccoli / evo) | tensor4all/faer#91, tensor4all/faer#95 | No timing here (performance evidence is the monitor's: broccoli 1/8/16/32/64T, evo). Tests: `utils::team::tests::{par_width_takes_the_widest_width_that_pays, with_width_runs_on_a_sub_team_of_that_width, the_families_narrow_their_regions_and_match_sequential}`, the rewritten `region_rules_follow_the_serial_time`, `fixed_and_never_keep_their_meaning`, `par_width_reads_the_degree_of_the_token`. Record: [worklog](docs/worklogs/2026-10-29-width-per-region.md). | | No heap allocation inside a parallel region (#111) | Every region on the team route allocated on the heap. `for_each_child`, behind every `for_each` and `for_each_with_par`, collected all the items into a `Vec`, one entry per column for `par_col_iter_mut`, and then a slot `Vec`. The full-pivoting LU step and both Bunch-Kaufman full-pivoting updates built an `nthreads`-long vector of pivot candidates per step, and `join_sized` boxed its two closures. The allocations were suspected of causing #109's floor-level loss (`full_piv_lu f64 128` at 8T on evo, +6.6 to +10.4%). The evo re-measurement of this change shows they were not: the cell is at 1.0025 of `main` with them gone. The change stands on its own as a cleanup of the mechanism. Now faer's dispatch and these workspaces allocate nothing. An iterator adapter that allocates in its own `split_at`, such as rayon's `map_with`, still does. `for_each_child` splits the iterator's producer on the leader into the blocks (`split_at`, remainder-relative) and each block's leader iterates its piece. Up to 64 blocks, with a producer of at most 256 bytes, are dealt from an inline array in one fork; wider or heavier splits halve recursively, one section per level. The pivot candidates come from the call's scratch: `lu_in_place_scratch` and `cholesky_in_place_scratch` grow by `par.degree()` entries. `join_sized` takes two typed slots. | KEEP (evo: no regression; `full_piv_lu f64 128` 8T at 1.0025 of `main`) | tensor4all/faer#111, tensor4all/faer#95 | Test `faer/tests/no_step_alloc.rs`: its own binary with a process-wide counting allocator. On teams of 4, 8 and 70 (the last takes the recursive path), with the regions forced parallel and a warm-up first, it counts zero allocations for full-pivoting LU 96, Bunch-Kaufman with full pivoting 96 and `join_sized`, and checks that each one dispatched. On `origin/main` the same test counts 190 allocations in one full-pivoting LU 96 call at width 4. Record: [worklog](docs/worklogs/2026-10-29-no-region-allocation.md). | +| Matrix product into uninitialized storage (#136) | `matmul_with_conj_uninit(dst: MatUninitMut, lhs, conj_lhs, rhs, conj_rhs, alpha, par) -> MatMut` computes `alpha * op(lhs) * op(rhs)` into a destination of `MaybeUninit` elements and returns it as an initialized `MatMut`, for callers (strided-rs#198) that would otherwise fill the output before an overwriting product. Requested upstream as sarah-quinones/faer-rs#269; the view type `MatUninitMut` (`UninitMut`) and the first tests come from shinaoka's upstream PR #270 (same author). The destination's prior contents are never read and no `&T`/`&mut T` to an unwritten element is formed. The native types on x86_64 with AVX2 or AVX-512 go to `t4a-private-gemm-x86` with `Accum::Replace` (on a team, through `gemm_team`); other scalars write each dot product through `ptr::write`; an empty depth writes zeros. The `gemm` and `nano-gemm` crates never see uninitialized storage (`gemm`'s GEMV forms a `&mut [T]` of the destination, its Apple AMX kernel copies partial tiles of it before checking `alpha`, and `nano-gemm` scales it by `beta`), so on other targets the native types are zero-filled first and multiplied as by `matmul_with_conj`. Any strides, negative included; a batch is one view per slot. A panic leaves the storage `MaybeUninit` (possibly partly written) with nothing exposed as initialized. Existing `MatMut` APIs are unchanged. | EXPERIMENT | tensor4all/faer#136 | Tests in `linalg::matmul::tests`: NaN-prefilled destinations with padded leading dimensions for f32/f64/c32/c64, row- and column-major, conjugations, at `1x1x1` up to `130x70x90` and across the kernel's 512 depth block (`9x600x7`); the padding must stay NaN and the product must match `matmul_with_conj`. Also a team, a batched buffer, negative strides, an empty depth, a matrix-vector shape and `fx128`. Miri runs the seven small tests; six of them use `MaybeUninit::uninit()` storage (the batched one is NaN-filled), and under Miri the native types take the zero-fill route. No uninitialized read was reported. The NaN and team tests are too large for Miri. Indicative only (primerose, shared, one shape, best of 30, three runs): `2048x64x2048` f64 `Par::Seq` into existing storage is 0.89 to 0.93 of a zero fill followed by `Accum::Replace`. The saving is the output-sized fill, so it shrinks as `k` grows. Not timed on a quiet host. Record: [worklog](docs/worklogs/2026-10-29-matmul-uninit.md). | ## Ground rules diff --git a/docs/worklogs/2026-10-29-matmul-uninit.md b/docs/worklogs/2026-10-29-matmul-uninit.md new file mode 100644 index 00000000..1f4b7f37 --- /dev/null +++ b/docs/worklogs/2026-10-29-matmul-uninit.md @@ -0,0 +1,37 @@ +# Matrix product into uninitialized storage (#136) + +## Goal + +strided-rs#198 holds output buffers as `MaybeUninit` and wants a product that overwrites +them without first filling them, which the upstream request sarah-quinones/faer-rs#269 asks +for. This fork adds `matmul_with_conj_uninit` taking a `MatUninitMut` and returning the +initialized `MatMut`. The view type and the first tests come from shinaoka's upstream PR #270. +Nothing was posted upstream. + +## Decisions + +- **Kernel routes.** Overwriting is not enough; the destination must never be read and never + be borrowed as `&T`/`&mut T` before it is written. + - `t4a-private-gemm-x86` with `Accum::Replace` meets both. It writes every element in its + first depth block (later blocks add to what that block wrote) and loads the destination + only when accumulating. It is the route for the native types on x86_64 with AVX2 or + AVX-512, directly or through `gemm_team` on a team. + - The `gemm` crate with `read_dst = false` was the first plan for the other targets. It was + dropped after the independent review traced two exceptions in `gemm-common` 0.19: its GEMV + builds a `&mut [T]` over the destination before writing, and its experimental Apple AMX + kernel copies the destination into a staging tile before it checks `alpha`. + - `nano-gemm` scales the destination by `beta`, so it reads it. + - On those targets the native types are therefore zero-filled through raw writes and then + multiplied by the ordinary path. This is correct everywhere, and gains nothing there. +- **Other scalars.** One dot product per element, written with `ptr::write`. A team splits + the columns. +- **Panics.** The `MatMut` is formed only after the whole write. A panic leaves the caller's + storage `MaybeUninit`, possibly partly written, and never exposed as initialized. + +## Limitations + +- Miri cannot run the SIMD kernel. Under Miri the native types take the zero-fill route, so + Miri checks the generic and fill paths, and the NaN-prefill tests check that the x86 kernel + does not carry old contents into the product. Neither proves the absence of a discarded read + inside the kernel; that rests on reading its epilogues. +- The timing is one shape on a shared host (in `PATCHSET.md`), not a quiet-host measurement. diff --git a/faer/src/lib.rs b/faer/src/lib.rs index 57378722..9fd387cf 100644 --- a/faer/src/lib.rs +++ b/faer/src/lib.rs @@ -1153,7 +1153,7 @@ pub type cx128 = traits::cx128; /// `Complex` pub type fx128 = traits::fx128; pub use col::{Col, ColMut, ColRef}; -pub use mat::{Mat, MatMut, MatRef}; +pub use mat::{Mat, MatMut, MatRef, MatUninitMut}; pub use row::{Row, RowMut, RowRef}; #[allow(unused_imports, dead_code)] mod internal_prelude { @@ -1253,7 +1253,7 @@ pub mod prelude { pub use super::prelude_sp::*; pub use super::{Par, ParShape, Scale, c32, c64, col, mat, row, unzip, zip}; pub use col::{Col, ColMut, ColRef}; - pub use mat::{Mat, MatMut, MatRef}; + pub use mat::{Mat, MatMut, MatRef, MatUninitMut}; pub use reborrow::{IntoConst, Reborrow, ReborrowMut}; pub use row::{Row, RowMut, RowRef}; /// see [`Default`] diff --git a/faer/src/linalg/matmul/mod.rs b/faer/src/linalg/matmul/mod.rs index 5213db2b..4c1b14a3 100644 --- a/faer/src/linalg/matmul/mod.rs +++ b/faer/src/linalg/matmul/mod.rs @@ -2,7 +2,7 @@ use super::temp_mat_scratch; use crate::col::ColRef; use crate::internal_prelude::*; -use crate::mat::{MatMut, MatRef}; +use crate::mat::{MatMut, MatRef, MatUninitMut}; use crate::row::RowRef; use crate::utils::bound::Dim; use crate::utils::simd::SimdCtx; @@ -1610,6 +1610,269 @@ fn matmul_imp<'M, 'N, 'K, T: ComplexField>( }, } } +/// `alpha * sum_p op(lhs[i, p]) * op(rhs[p, j])` written into the uninitialized slot `(i, j)` of +/// `dst`: read through raw pointers, written with `ptr::write`, no reference to the slot. +/// +/// # Safety +/// The pointers and strides describe valid `m x k`, `k x n` and `m x n` matrices with `i < m`, +/// `j < n`; `dst` does not overlap the operands. +#[inline(always)] +unsafe fn matmul_uninit_one( + dst: *mut MaybeUninit, + (dst_rs, dst_cs): (isize, isize), + lhs: *const T, + (lhs_rs, lhs_cs): (isize, isize), + rhs: *const T, + (rhs_rs, rhs_cs): (isize, isize), + k: usize, + (i, j): (usize, usize), + (conj_lhs, conj_rhs): (Conj, Conj), + alpha: &T, +) { + let mut acc: T = zero(); + for p in 0..k { + let l = unsafe { &*lhs.offset(i as isize * lhs_rs + p as isize * lhs_cs) }; + let r = unsafe { &*rhs.offset(p as isize * rhs_rs + j as isize * rhs_cs) }; + acc = &conj_lhs.apply_rt(l) * &conj_rhs.apply_rt(r) + acc; + } + unsafe { + dst.offset(i as isize * dst_rs + j as isize * dst_cs) + .write(MaybeUninit::new(alpha * &acc)); + } +} + +/// The body of [`matmul_with_conj_uninit`], on dynamic shapes: every logical element of `dst` is +/// written through raw pointers before it returns, and its prior contents are never read. +fn matmul_uninit_imp( + dst: MatUninitMut<'_, T>, + lhs: MatRef<'_, T>, + conj_lhs: Conj, + rhs: MatRef<'_, T>, + conj_rhs: Conj, + alpha: T, + par: Par, +) { + let m = dst.nrows(); + let n = dst.ncols(); + let k = lhs.ncols(); + precondition(m, n, lhs.nrows(), k, rhs.nrows(), rhs.ncols()); + if m == 0 || n == 0 { + return; + } + let ptr = dst.as_ptr_mut(); + let (dst_rs, dst_cs) = (dst.row_stride(), dst.col_stride()); + // zeros through raw writes, after which the destination is initialized + let fill_zero = || { + for j in 0..n { + for i in 0..m { + unsafe { + ptr.offset(i as isize * dst_rs + j as isize * dst_cs) + .write(MaybeUninit::new(zero())); + } + } + } + }; + if k == 0 { + // an empty product is zero; the kernels below write nothing when the depth is empty + fill_zero(); + return; + } + // The native types on x86_64 with AVX2 or AVX-512: `t4a-private-gemm-x86` with + // `Accum::Replace`, which writes every element in its first depth block and never loads the + // destination before that (on a `Par::Team`, `gemm_team` takes the same kernel). The `gemm` + // and `nano-gemm` crates are not given uninitialized storage: `gemm`'s GEMV forms a + // `&mut [T]` of the destination and its Apple AMX kernel copies partial tiles of it before + // checking `alpha`, and `nano-gemm` scales it by `beta`. + #[cfg(all(target_arch = "x86_64", feature = "std"))] + { + macro_rules! gemm_call { + ($kind:ident, $ty:ty) => { + unsafe { + use private_gemm_x86::*; + let feat = if std::arch::is_x86_feature_detected!("avx512f") { + Some(InstrSet::Avx512) + } else if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma") { + Some(InstrSet::Avx256) + } else { + None + }; + if let Some(feat) = feat { + let alpha: $ty = core::mem::transmute_copy(&alpha); + let dst = ptr as *mut $ty; + let (lhs_p, rhs_p) = (lhs.as_ptr() as *const $ty, rhs.as_ptr() as *const $ty); + #[cfg(feature = "threadteam")] + if let Par::Team(team) = par { + // the same instruction set test as `gemm_team`'s, so it takes the kernel + crate::utils::team::gemm_team::<$ty>( + team, + m, + n, + k, + crate::utils::cost::gemm_ns_per_mac::(), + dst, + dst_cs, + dst_rs, + false, + lhs_p, + lhs.col_stride(), + lhs.row_stride(), + rhs_p, + rhs.col_stride(), + rhs.row_stride(), + core::mem::zeroed(), + alpha, + conj_lhs == Conj::Yes, + conj_rhs == Conj::Yes, + ); + return; + } + gemm( + DType::$kind, + IType::U64, + feat, + m, + n, + k, + dst as *mut (), + dst_rs, + dst_cs, + core::ptr::null(), + core::ptr::null(), + DstKind::Full, + Accum::Replace, + lhs_p as *const (), + lhs.row_stride(), + lhs.col_stride(), + conj_lhs == Conj::Yes, + core::ptr::null(), + 0, + rhs_p as *const (), + rhs.row_stride(), + rhs.col_stride(), + conj_rhs == Conj::Yes, + &raw const alpha as *const (), + 1, + ); + return; + } + } + }; + } + if const { T::IS_NATIVE_F64 } { + gemm_call!(F64, f64); + } + if const { T::IS_NATIVE_C64 } { + gemm_call!(C64, num_complex::Complex); + } + if const { T::IS_NATIVE_F32 } { + gemm_call!(F32, f32); + } + if const { T::IS_NATIVE_C32 } { + gemm_call!(C32, num_complex::Complex); + } + } + if const { T::IS_NATIVE_F64 || T::IS_NATIVE_C64 || T::IS_NATIVE_F32 || T::IS_NATIVE_C32 } { + // a native type without that kernel: initialize, then the ordinary product, which + // overwrites the zeros without reading them (the zero pass is the cost the kernel route + // avoids) + fill_zero(); + let dst = unsafe { MatMut::from_raw_parts_mut(ptr as *mut T, m, n, dst_rs, dst_cs) }; + matmul_with_conj(dst, Accum::Replace, lhs, conj_lhs, rhs, conj_rhs, alpha, par); + return; + } + // any other scalar: one dot product per element, by columns (a column per task on a team) + let (dst_w, lhs_w, rhs_w) = ( + crate::utils::thread::Ptr(ptr), + crate::utils::thread::Ptr(lhs.as_ptr() as *mut T), + crate::utils::thread::Ptr(rhs.as_ptr() as *mut T), + ); + let lhs_s = (lhs.row_stride(), lhs.col_stride()); + let rhs_s = (rhs.row_stride(), rhs.col_stride()); + let alpha = α + let column = move |j: usize| { + let (dst_w, lhs_w, rhs_w) = (dst_w, lhs_w, rhs_w); + for i in 0..m { + // SAFETY: the shapes were checked by `precondition`; `dst` is exclusive and does not + // overlap the operands (the caller's contract); each column is written by one task + unsafe { + matmul_uninit_one(dst_w.0, (dst_rs, dst_cs), lhs_w.0, lhs_s, rhs_w.0, rhs_s, k, (i, j), (conj_lhs, conj_rhs), alpha) + } + } + }; + match par { + Par::Seq => (0..n).for_each(column), + #[cfg(feature = "threadteam")] + par => { + use rayon::prelude::*; + let degree = par.degree(); + crate::utils::thread::for_each(par, degree, (0..n).into_par_iter(), column); + }, + #[cfg(not(feature = "threadteam"))] + _ => (0..n).for_each(column), + } +} + +/// computes `alpha * op(lhs) * op(rhs)` into a destination whose elements may be uninitialized, +/// and returns it as an initialized [`MatMut`]. +/// +/// This is the overwrite-only entry for callers that hold `MaybeUninit` storage and do not +/// want an output-sized initialization pass before the product (upstream request +/// sarah-quinones/faer-rs#269). The destination's prior contents are never read, no `&T` or +/// `&mut T` to one of its elements is formed before the element is written, and every logical +/// element is written before the function returns. Only then is the returned [`MatMut`] formed. +/// If the call panics, the storage stays `MaybeUninit` (some elements may have been written) +/// and nothing is exposed as initialized. +/// +/// The output-sized initialization pass is skipped for the native types on x86_64 with AVX2 or +/// AVX-512, and for other scalar types. On other targets the native types are zero-filled first +/// and then multiplied as by [`matmul_with_conj`]. +/// +/// The parallelism is the explicit `par`. The operands must not overlap the destination, and the +/// destination's elements must not overlap each other: the safe slice constructors of +/// [`MatUninitMut`] check this, and the raw constructor requires it. +/// +/// # example +/// +/// ``` +/// use core::mem::MaybeUninit; +/// use faer::linalg::matmul::matmul_with_conj_uninit; +/// use faer::{mat, Conj, MatUninitMut, Par}; +/// +/// let lhs = mat![[1.0_f64, 2.0], [3.0, 4.0]]; +/// let rhs = mat![[5.0_f64, 6.0], [7.0, 8.0]]; +/// let mut storage = [MaybeUninit::::uninit(); 4]; +/// let dst = MatUninitMut::from_column_major_slice_mut(&mut storage, 2, 2); +/// let result = matmul_with_conj_uninit(dst, lhs.as_ref(), Conj::No, rhs.as_ref(), Conj::No, 1.0, Par::Seq); +/// assert_eq!(result[(0, 0)], 19.0); +/// ``` +#[track_caller] +#[inline] +pub fn matmul_with_conj_uninit<'a, T: ComplexField, M: Shape, N: Shape, K: Shape, RStride: Stride, CStride: Stride>( + dst: MatUninitMut<'a, T, M, N, RStride, CStride>, + lhs: impl AsMatRef, + conj_lhs: Conj, + rhs: impl AsMatRef, + conj_rhs: Conj, + alpha: T, + par: Par, +) -> MatMut<'a, T, M, N, RStride, CStride> { + let ptr = dst.as_ptr_mut(); + let (nrows, ncols) = (dst.nrows(), dst.ncols()); + let (row_stride, col_stride) = (dst.row_stride(), dst.col_stride()); + matmul_uninit_imp( + dst.as_dyn(), + lhs.as_mat_ref().as_dyn(), + conj_lhs, + rhs.as_mat_ref().as_dyn(), + conj_rhs, + alpha, + par, + ); + // SAFETY: `matmul_uninit_imp` returned, so every logical element has been written; the view + // keeps the destination's exclusive borrow `'a` + unsafe { MatMut::from_raw_parts_mut(ptr.cast(), nrows, ncols, row_stride, col_stride) } +} + #[track_caller] fn precondition( dst_nrows: M, @@ -1805,7 +2068,8 @@ mod tests { use super::*; use crate::mat::{Mat, MatMut, MatRef}; use crate::stats::prelude::*; - use crate::{assert, c32}; + use crate::{assert, c32, c64, fx128}; + use core::mem::MaybeUninit; use std::num::NonZeroUsize; /// Below the flop floor `matmul` runs the sequential kernel whatever budget /// the caller passes, and above it the kernel threads: both must compute the @@ -2002,6 +2266,188 @@ mod tests { } } } + #[test] + fn test_matmul_with_conj_uninit() { + let lhs = mat![[1.0_f64, 2.0], [3.0, 4.0]]; + let rhs = mat![[5.0_f64, 6.0], [7.0, 8.0]]; + let mut data = [MaybeUninit::::uninit(); 4]; + let dst = MatUninitMut::from_column_major_slice_mut(&mut data, 2, 2); + let result = matmul_with_conj_uninit(dst, lhs.as_ref(), Conj::No, rhs.as_ref(), Conj::No, 1.0, Par::Seq); + assert!(result[(0, 0)] == 19.0); + assert!(result[(1, 0)] == 43.0); + assert!(result[(0, 1)] == 22.0); + assert!(result[(1, 1)] == 50.0); + } + + /// A matrix-vector product into uninitialized storage: the shape for which the `gemm` crate + /// would take its GEMV path, which forms a slice of the destination (run it under Miri). + #[test] + fn test_matmul_with_conj_uninit_matrix_vector() { + let lhs = mat![[1.0_f64, 2.0, 3.0], [4.0, 5.0, 6.0]]; + let rhs = mat![[1.0_f64], [-1.0], [2.0]]; + let mut data = [MaybeUninit::::uninit(); 2]; + let dst = MatUninitMut::from_column_major_slice_mut(&mut data, 2, 1); + let result = matmul_with_conj_uninit(dst, lhs.as_ref(), Conj::No, rhs.as_ref(), Conj::No, 1.0, Par::Seq); + assert!(result[(0, 0)] == 5.0); + assert!(result[(1, 0)] == 11.0); + } + + #[test] + fn test_matmul_with_conj_uninit_strided_and_conjugated() { + let lhs = mat![[c64::new(1.0, 2.0), c64::new(2.0, -1.0)], [c64::new(3.0, 1.0), c64::new(-2.0, 4.0)],]; + let rhs = mat![[c64::new(5.0, -1.0), c64::new(1.0, 3.0)], [c64::new(2.0, 2.0), c64::new(-4.0, 1.0)],]; + let mut expected = Mat::::zeros(2, 2); + matmul_with_conj(expected.as_mut(), Accum::Replace, lhs.as_ref(), Conj::Yes, rhs.as_ref(), Conj::No, c64::new(2.0, 0.0), Par::Seq); + let mut data = [MaybeUninit::::uninit(); 7]; + let dst = unsafe { MatUninitMut::from_raw_parts_mut(data.as_mut_ptr(), 2, 2, 3, 1) }; + let result = matmul_with_conj_uninit(dst, lhs.as_ref(), Conj::Yes, rhs.as_ref(), Conj::No, c64::new(2.0, 0.0), Par::Seq); + for j in 0..2 { + for i in 0..2 { + assert!((result[(i, j)] - expected[(i, j)]).norm() < 1e-12); + } + } + } + + #[test] + fn test_matmul_with_conj_uninit_negative_destination_stride() { + let lhs = mat![[1.0_f64, 2.0], [3.0, 4.0]]; + let rhs = mat![[5.0_f64, 6.0], [7.0, 8.0]]; + let mut data = [MaybeUninit::::uninit(); 4]; + let dst = unsafe { MatUninitMut::from_raw_parts_mut(data.as_mut_ptr().add(1), 2, 2, -1, 2) }; + let result = matmul_with_conj_uninit(dst, lhs.as_ref(), Conj::No, rhs.as_ref(), Conj::No, 1.0, Par::Seq); + assert!(result[(0, 0)] == 19.0); + assert!(result[(1, 0)] == 43.0); + assert!(result[(0, 1)] == 22.0); + assert!(result[(1, 1)] == 50.0); + } + + #[test] + fn test_matmul_with_conj_uninit_empty_inner_dimension() { + let lhs = Mat::::zeros(2, 0); + let rhs = Mat::::zeros(0, 3); + let mut data = [MaybeUninit::::uninit(); 6]; + let dst = MatUninitMut::from_row_major_slice_mut(&mut data, 2, 3); + let result = matmul_with_conj_uninit(dst, lhs.as_ref(), Conj::No, rhs.as_ref(), Conj::No, 2.0, Par::Seq); + for i in 0..2 { + for j in 0..3 { + assert!(result[(i, j)] == 0.0); + } + } + } + + #[test] + fn test_matmul_with_conj_uninit_generic_scalar() { + let from = |x: f64| ::from_f64(x); + let lhs = Mat::::from_fn(2, 2, |i, j| from((i + 2 * j + 1) as f64)); + let rhs = Mat::::from_fn(2, 2, |i, j| from((3 * i + j + 1) as f64)); + let mut data = [MaybeUninit::::uninit(); 4]; + let dst = MatUninitMut::from_column_major_slice_mut(&mut data, 2, 2); + let result = matmul_with_conj_uninit(dst, lhs.as_ref(), Conj::No, rhs.as_ref(), Conj::No, from(2.0), Par::Seq); + assert!(result[(0, 0)] == from(26.0)); + assert!(result[(1, 0)] == from(36.0)); + assert!(result[(0, 1)] == from(34.0)); + assert!(result[(1, 1)] == from(48.0)); + } + + /// The native kernels at sizes below and above the small-matrix and tiling thresholds: the destination starts as NaN, which a kernel that + /// read it would carry into the product, and its padding rows must be left untouched. + fn check_uninit_native( + from: impl Fn(f64, f64) -> T, + dist: impl Fn(T, T) -> f64, + tol: f64, + ) { + let nan = ::from_f64(f64::NAN); + for (m, k, n) in [(1, 1, 1), (3, 4, 5), (17, 23, 31), (64, 64, 64), (130, 70, 90), (5, 300, 7), (9, 600, 7)] { + for (conj_lhs, conj_rhs) in [(Conj::No, Conj::No), (Conj::Yes, Conj::No), (Conj::No, Conj::Yes)] { + for row_major in [false, true] { + let lhs = Mat::::from_fn(m, k, |i, j| from(((i * 7 + j * 3) % 11) as f64 - 5.0, ((i + 2 * j) % 5) as f64 - 2.0)); + let rhs = Mat::::from_fn(k, n, |i, j| from(((i * 5 + j) % 13) as f64 - 6.0, ((3 * i + j) % 7) as f64 - 3.0)); + let alpha = from(0.5, -1.5); + let mut expected = Mat::::zeros(m, n); + matmul_with_conj(expected.as_mut(), Accum::Replace, lhs.as_ref(), conj_lhs, rhs.as_ref(), conj_rhs, alpha, Par::Seq); + // a padded leading dimension: three NaN elements between columns (or rows) + let (rs, cs, len) = if row_major { (n as isize + 3, 1, m * (n + 3)) } else { (1, m as isize + 3, n * (m + 3)) }; + let mut data = alloc::vec![MaybeUninit::new(nan); len]; + let dst = unsafe { MatUninitMut::::from_raw_parts_mut(data.as_mut_ptr(), m, n, rs, cs) }; + let result = matmul_with_conj_uninit(dst, lhs.as_ref(), conj_lhs, rhs.as_ref(), conj_rhs, alpha, Par::Seq); + for j in 0..n { + for i in 0..m { + let err = dist(result[(i, j)], expected[(i, j)]); + ::core::assert!(err <= tol * k as f64 * 200.0, "{m}x{k}x{n} ({i}, {j}): {:?} vs {:?}", result[(i, j)], expected[(i, j)]); + } + } + let inner = if row_major { n } else { m }; + for (p, x) in data.iter().enumerate() { + if p % (inner + 3) >= inner { + let x = unsafe { x.assume_init() }; + ::core::assert!(x != x, "{m}x{k}x{n}: padding element {p} was written"); + } + } + } + } + } + } + + #[test] + fn test_matmul_with_conj_uninit_native_never_reads_destination() { + check_uninit_native(|x, _| x as f32, |a: f32, b| (a - b).abs() as f64, 1e-6); + check_uninit_native(|x, _| x, |a: f64, b| (a - b).abs(), 1e-14); + check_uninit_native(|x, y| c32::new(x as f32, y as f32), |a: c32, b| (a - b).norm() as f64, 1e-6); + check_uninit_native(|x, y| c64::new(x, y), |a: c64, b| (a - b).norm(), 1e-14); + } + + #[cfg(feature = "threadteam")] + #[test] + fn test_matmul_with_conj_uninit_team() { + crate::utils::team::with_team(3, |mut par| { + { + let lhs = Mat::::from_fn(130, 70, |i, j| ((i * 7 + j * 3) % 11) as f64 - 5.0); + let rhs = Mat::::from_fn(70, 90, |i, j| ((i * 5 + j) % 13) as f64 - 6.0); + let mut expected = Mat::::zeros(130, 90); + matmul_with_conj(expected.as_mut(), Accum::Replace, lhs.as_ref(), Conj::No, rhs.as_ref(), Conj::No, 1.0, Par::Seq); + let mut data = alloc::vec![MaybeUninit::new(f64::NAN); 130 * 90]; + let dst = MatUninitMut::from_column_major_slice_mut(&mut data, 130, 90); + let result = matmul_with_conj_uninit(dst, lhs.as_ref(), Conj::No, rhs.as_ref(), Conj::No, 1.0, par.reborrow()); + assert!(result == expected); + } + { + let from = |x: f64| ::from_f64(x); + let lhs = Mat::::from_fn(9, 6, |i, j| from((i + 2 * j) as f64)); + let rhs = Mat::::from_fn(6, 11, |i, j| from((3 * i + j) as f64)); + let mut expected = Mat::::zeros(9, 11); + matmul_with_conj(expected.as_mut(), Accum::Replace, lhs.as_ref(), Conj::No, rhs.as_ref(), Conj::No, from(1.0), Par::Seq); + let mut data = alloc::vec![MaybeUninit::::uninit(); 9 * 11]; + let dst = MatUninitMut::from_column_major_slice_mut(&mut data, 9, 11); + let result = matmul_with_conj_uninit(dst, lhs.as_ref(), Conj::No, rhs.as_ref(), Conj::No, from(1.0), par.reborrow()); + assert!(result == expected); + } + }); + } + + /// A batched destination: one buffer holding `batch` matrices at a batch stride, each filled + /// through its own strided view. + #[test] + fn test_matmul_with_conj_uninit_batched() { + let (batch, m, k, n) = (3, 5, 4, 6); + let batch_stride = m * n + 2; + let mut data = alloc::vec![MaybeUninit::new(f64::NAN); batch * batch_stride]; + for b in 0..batch { + let lhs = Mat::::from_fn(m, k, |i, j| (i + j + b) as f64); + let rhs = Mat::::from_fn(k, n, |i, j| (i * j + b) as f64 - 2.0); + let mut expected = Mat::::zeros(m, n); + matmul_with_conj(expected.as_mut(), Accum::Replace, lhs.as_ref(), Conj::No, rhs.as_ref(), Conj::No, 1.0, Par::Seq); + // row major inside each batch slot + let dst = unsafe { MatUninitMut::::from_raw_parts_mut(data.as_mut_ptr().add(b * batch_stride), m, n, n as isize, 1) }; + let result = matmul_with_conj_uninit(dst, lhs.as_ref(), Conj::No, rhs.as_ref(), Conj::No, 1.0, Par::Seq); + assert!(result == expected); + } + for b in 0..batch { + for p in m * n..batch_stride { + assert!(unsafe { data[b * batch_stride + p].assume_init() }.is_nan()); + } + } + } + fn matmul_with_conj_fallback( acc: MatMut<'_, T>, a: MatRef<'_, T>, diff --git a/faer/src/mat/matmut.rs b/faer/src/mat/matmut.rs index 12c9d9f9..4fdd827e 100644 --- a/faer/src/mat/matmut.rs +++ b/faer/src/mat/matmut.rs @@ -2,6 +2,7 @@ use super::{MatRef, *}; use crate::internal_prelude::*; use crate::utils::bound::{Dim, Partition}; use crate::{Conj, ContiguousFwd, Idx, IdxInc}; +use core::mem::MaybeUninit; use equator::assert; use faer_traits::ComplexField; @@ -19,6 +20,154 @@ pub struct Mut< pub(super) imp: MatView, pub(super) __marker: PhantomData<&'a mut T>, } + +/// write-only mutable matrix view over possibly uninitialized storage. +/// +/// This type deliberately does not implement any operation that forms a +/// reference to `T`. It is used by overwrite-only kernels, which write through +/// the raw pointer and return an initialized view only after a full write. +pub struct UninitMut< + 'a, + T, + Rows = usize, + Cols = usize, + RStride = isize, + CStride = isize, +> { + pub(super) imp: MatView, Rows, Cols, RStride, CStride>, + pub(super) __marker: PhantomData<&'a mut MaybeUninit>, +} + +unsafe impl Sync + for UninitMut<'_, T, Rows, Cols, RStride, CStride> +{ +} +unsafe impl Send + for UninitMut<'_, T, Rows, Cols, RStride, CStride> +{ +} + +impl<'a, T, Rows: Shape, Cols: Shape, RStride: Stride, CStride: Stride> + MatUninitMut<'a, T, Rows, Cols, RStride, CStride> +{ + /// creates a write-only matrix view from a raw pointer. + /// + /// # safety + /// the pointed-to memory must contain the complete matrix described by the + /// dimensions and strides, be properly aligned for `MaybeUninit`, and + /// remain exclusively borrowed for `'a`. Matrix elements must not overlap + /// each other. Unlike [`MatMut::from_raw_parts_mut`], the elements do not + /// need to be initialized. + #[inline] + #[track_caller] + pub const unsafe fn from_raw_parts_mut( + ptr: *mut MaybeUninit, + nrows: Rows, + ncols: Cols, + row_stride: RStride, + col_stride: CStride, + ) -> Self { + generic::Mat(UninitMut { + imp: MatView { + ptr: NonNull::new_unchecked(ptr), + nrows, + ncols, + row_stride, + col_stride, + }, + __marker: PhantomData, + }) + } + + /// returns the raw pointer to the first matrix element. + #[inline] + pub fn as_ptr_mut(&self) -> *mut MaybeUninit { + self.imp.ptr.as_ptr() + } + + /// returns the number of rows. + #[inline] + pub fn nrows(&self) -> Rows { + self.imp.nrows + } + + /// returns the number of columns. + #[inline] + pub fn ncols(&self) -> Cols { + self.imp.ncols + } + + /// returns the row stride in elements. + #[inline] + pub fn row_stride(&self) -> RStride { + self.imp.row_stride + } + + /// returns the column stride in elements. + #[inline] + pub fn col_stride(&self) -> CStride { + self.imp.col_stride + } + + /// converts this view to dynamic dimensions and strides. + #[inline] + pub fn as_dyn(self) -> MatUninitMut<'a, T, usize, usize, isize, isize> { + let ptr = self.as_ptr_mut(); + let nrows = self.nrows().unbound(); + let ncols = self.ncols().unbound(); + let row_stride = self.row_stride().element_stride(); + let col_stride = self.col_stride().element_stride(); + unsafe { + MatUninitMut::from_raw_parts_mut( + ptr, nrows, ncols, row_stride, col_stride, + ) + } + } +} + +impl<'a, T, Rows: Shape, Cols: Shape> + MatUninitMut<'a, T, Rows, Cols, isize, isize> +{ + /// creates a column-major write-only view over a `MaybeUninit` slice. + #[inline] + #[track_caller] + pub fn from_column_major_slice_mut( + slice: &'a mut [MaybeUninit], + nrows: Rows, + ncols: Cols, + ) -> Self { + from_slice_assert(nrows.unbound(), ncols.unbound(), slice.len()); + unsafe { + Self::from_raw_parts_mut( + slice.as_mut_ptr(), + nrows, + ncols, + 1, + nrows.unbound() as isize, + ) + } + } + + /// creates a row-major write-only view over a `MaybeUninit` slice. + #[inline] + #[track_caller] + pub fn from_row_major_slice_mut( + slice: &'a mut [MaybeUninit], + nrows: Rows, + ncols: Cols, + ) -> Self { + from_slice_assert(nrows.unbound(), ncols.unbound(), slice.len()); + unsafe { + Self::from_raw_parts_mut( + slice.as_mut_ptr(), + nrows, + ncols, + ncols.unbound() as isize, + 1, + ) + } + } +} #[repr(transparent)] pub(crate) struct SyncCell(T); unsafe impl Sync for SyncCell {} diff --git a/faer/src/mat/mod.rs b/faer/src/mat/mod.rs index a8e819a9..6dbe85dc 100644 --- a/faer/src/mat/mod.rs +++ b/faer/src/mat/mod.rs @@ -48,6 +48,7 @@ pub(crate) mod matmut; pub(crate) mod matown; pub(crate) mod matref; pub use matmut::Mut; +pub use matmut::UninitMut; pub use matown::Own; pub use matref::Ref; /// heap allocated resizable matrix, similar to a 2d [`alloc::vec::Vec`] @@ -132,7 +133,9 @@ pub type MatRef< /// which allows us to mutably borrow a `MatMut` to obtain another `MatMut` for /// the lifetime of the borrow. it's also similarly possible to immutably borrow /// a `MatMut` to obtain a `MatRef` for the lifetime of the borrow, using -/// [`reborrow::Reborrow`] ``` +/// [`reborrow::Reborrow`] +/// +/// ``` /// use faer::{Mat, MatMut, MatRef}; /// use reborrow::*; /// fn takes_matmut(view: MatMut<'_, f64>) {} @@ -153,6 +156,22 @@ pub type MatMut< CStride = isize, > = generic::Mat>; +/// mutable view over a matrix whose elements may be uninitialized. +/// +/// Unlike [`MatMut`], this type does not expose element references. It is +/// intended for APIs that overwrite every logical destination element before +/// the storage is converted into an initialized matrix. The only operations +/// available on this type are metadata queries and raw-pointer access for +/// write-only kernels. +pub type MatUninitMut< + 'a, + T, + Rows = usize, + Cols = usize, + RStride = isize, + CStride = isize, +> = generic::Mat>; + #[doc(hidden)] /// generic `Mat` wrapper pub mod generic {