Skip to content

matmul: overwrite a MaybeUninit destination (#136) - #153

Merged
shinaoka merged 1 commit into
mainfrom
feat/matmul-uninit
Oct 29, 2026
Merged

shinaoka merged 1 commit into
mainfrom
feat/matmul-uninit

Conversation

@shinaoka

Copy link
Copy Markdown
Member

Closes #136.

Summary

This adds faer::linalg::matmul::matmul_with_conj_uninit, an overwrite-only matrix product whose destination may be uninitialized storage:

pub fn matmul_with_conj_uninit<'a, T, M, N, K, RS, CS>(
    dst: MatUninitMut<'a, T, M, N, RS, CS>,
    lhs: impl AsMatRef<T = T, Rows = M, Cols = K>, conj_lhs: Conj,
    rhs: impl AsMatRef<T = T, Rows = K, Cols = N>, conj_rhs: Conj,
    alpha: T, par: Par,
) -> MatMut<'a, T, M, N, RS, CS>

It was requested upstream as sarah-quinones#269 and is needed by strided-rs#198. The view type MatUninitMut (mat::UninitMut) and the first tests come from shinaoka's upstream PR sarah-quinones#270 (branch feat/maybe-uninit-matmul, same author, MIT). The product routing is new and uses this fork's kernels. Nothing was posted upstream.

Requirements → implementation

requirement how
never read prior C Native types on x86_64 with AVX2/AVX-512 go to t4a-private-gemm-x86 with Accum::Replace, which writes every element in the first depth block and loads the destination only when accumulating; on a Par::Team they go through gemm_team, which takes the same kernel. Other scalars write each dot product with ptr::write. An empty depth writes zeros. The gemm and nano-gemm crates never see uninitialized storage (see the review below), so on other targets the native types are zero-filled through raw writes and then multiplied by the ordinary path.
no &T/&mut T to unwritten elements UninitMut exposes only metadata and the raw pointer; every write goes through a raw pointer (ptr::write on the generic path). The MatMut is formed from the raw parts only after the call has written every element.
strided and batched destinations any row and column strides, negative ones included. A batch is one MatUninitMut per slot (from_raw_parts_mut at the batch offset), as in the test.
every element initialized, MatMut returned yes; shapes are checked before any write.
drop safety after an error or panic the caller keeps its MaybeUninit<T> storage; no element is exposed as initialized unless the call returns.
explicit Par par argument; Par::Team reaches gemm_team, and the generic path is split by columns on the team.
existing MatMut APIs unchanged only additions: UninitMut, MatUninitMut, matmul_with_conj_uninit.

Tests

All tests are in linalg::matmul::tests.

  • From the upstream PR: basic, strided with conjugation, negative destination stride, empty inner dimension, fx128. The upstream PR's Rayon test is replaced by a team test, since Par::Rayon no longer exists here.
  • New, test_matmul_with_conj_uninit_native_never_reads_destination:
    • covers f32/f64/c32/c64, from 1x1x1 to 130x70x90, plus 5x300x7 and 9x600x7, which crosses the kernel's 512 depth block, row- and column-major, and the conjugation pairs;
    • the destination is prefilled with NaN and its leading dimension is padded by three elements;
    • the product must match matmul_with_conj, and the padding must still be NaN afterwards.
  • New, _team: f64 130x70x90 and fx128 on a team of 3.
  • New, _matrix_vector: 2x3 · 3x1 into uninit() storage, the shape where the gemm crate takes its GEMV.
  • New, _batched: three row-major slots in one buffer; the padding between slots stays untouched.
  • Miri (cargo +nightly miri test) runs the seven small tests, six of them on MaybeUninit::uninit() storage, and reports no uninitialized read. Under Miri the native types take the zero-fill route, because Miri cannot run the SIMD kernel. The NaN and team tests are too large for Miri.

Performance (indicative)

On primerose, a shared host, I timed one shape: 2048x64x2048, f64, Par::Seq, into existing storage, best of 30, three runs. The new call took 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. This is not a quiet-host measurement.

Verification

  • cargo test -p t4a-faer --lib uninit: 9 passed after the review fixes. Before them, --lib linalg::matmul gave 16 passed (1 ignored, as on main), and the new doctest passes.
  • Record: docs/worklogs/2026-10-29-matmul-uninit.md.
  • cargo check -p t4a-faer --no-default-features and --no-default-features --features std compile with no warning from this change.
  • Independent post-review (gpt-6.1-sol). I confirmed each finding against the source:
    • BLOCKER: gemm-common 0.19's mixed_gemv_colmajor builds a &mut [T] over the destination before writing it. The first version sent the non-x86 fallback through the gemm crate.
    • IMPORTANT: its experimental Apple AMX kernel copies partial destination tiles before it checks alpha.
    • Fix for both: the gemm crate no longer receives uninitialized storage; those targets zero-fill first.
    • MINOR: the docs promised "written exactly once" and "storage as it was" after a panic. They now promise no read of prior contents and say a panic may leave the storage partly written.
    • MINOR: the experiment record was missing; it is added.
    • MINOR: the duplicated column loop and the test adapter were removed.

🤖 Generated with Claude Code

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#269; the view type
and the first tests follow the same author's upstream PR sarah-quinones#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 <noreply@anthropic.com>
@shinaoka
shinaoka merged commit fd46421 into main Oct 29, 2026
7 checks passed
@shinaoka
shinaoka deleted the feat/matmul-uninit branch October 29, 2026 11:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

matmul: a MaybeUninit-safe overwrite destination (upstream faer-rs#269)

1 participant