Skip to content

Fix wrong results of the matmul fast path for multi-chunk layouts - #782

Open
FabianIsensee wants to merge 3 commits into
Blosc:mainfrom
FabianIsensee:fix-matmul-first-chunk
Open

FabianIsensee wants to merge 3 commits into
Blosc:mainfrom
FabianIsensee:fix-matmul-first-chunk

Conversation

@FabianIsensee

Copy link
Copy Markdown
Contributor

Depends on #774 (stacked; its commit shows here until #774 is merged).

The prefilter-based fast path of blosc2.matmul() (2-D float32/float64 operands whose chunks and blocks pass _matmul_can_use_fast_path()) gives wrong results in most layouts with more than one chunk, in 4.14.1 and on main. Three separate problems:

  1. aux_matmul() read the operand chunks (src[i] = ndarr.sc.data[chunk_idx]) before its first_run branch moved nchunkA/nchunkB to the output chunk's row and column, so the first K step of every output chunk multiplied A's and B's first chunks. Only layouts with a single output chunk came out right. Example: 200x200 float64 random integers, chunks (100, 100), blocks (50, 50) for A, B and the output: max abs error 226 vs NumPy.
  2. The kernel assumes that A, B and the output share one chunk grid (output chunk rows = A's, output chunk columns = B's, same chunks along K), but the gate only checks blocks. Other chunk shapes gave wrong results, raised in the prefilter, or crashed: A chunks (100, 50) with B chunks (100, 100) walks B's chunk index past sc.data and segfaults.
  3. The kernel multiplies whole blocks, so when K is not a multiple of the chunk shape it also multiplies A's padding columns with B's padding rows. That padding is not always zero: blosc2.full()/ones() fill it with the value, and resize() keeps the cut values. blosc2.matmul(ones((170, 170)), full((170, 170), 3)) with chunks (200, 200) and blocks (50, 50) is off by 90 in every element; this one also hits single-chunk layouts.

The existing tests use ones()/full() with shapes that are multiples of the chunks, which hides 1 and 3. Over 400 random gate-passing layouts with integer data, main gave 6 exact results, 159 wrong ones, 223 exceptions and 12 segfaults. Layouts chosen by default (no chunks/blocks given) only take the fast path when every array is a single chunk, so this needs explicitly chosen chunks/blocks, as in bench/ndarray/matmul_path_compare.py.

This PR:

  • computes the first chunk and block indices before the chunk loop (a move of the existing lines; they only depend on the shapes);
  • makes _matmul_can_use_fast_path() require the shared chunk grid; other layouts take the chunked path (the same 16 default single-chunk layouts as before still take the fast path);
  • (second commit, separable but stacked on the first) stops the block loop at the first block that lies entirely in the padding along K, and zeroes A's padding columns and B's padding rows of a block that ends in it (both, so NaN in one operand's padding cannot leak through 0 * NaN).

Tests in tests/ndarray/test_linalg.py (after #774's matmul tests, using its PREFILTER_NTHREADS), with random small-integer data compared exactly against NumPy, 1 and 4 threads (threaded cases skip on WASM):

  • test_matmul_fast_path_multichunk: more than one chunk along rows, columns, K, all three with partial last chunks, non-square chunks; float32 and float64. 16 of 20 fail without this PR (the K-only layout has a single output chunk).
  • test_matmul_fast_path_needs_shared_chunk_grid: three mismatched layouts take the chunked path. Without this PR two give wrong results and in the third the prefilter fails (fallback warning).
  • test_matmul_fast_path_ignores_padding: operands with NaN in their padding (shrunk with resize()), with a partial last K block and with padding-only K blocks. All 16 fail without this PR.

Full test suite: same results as #774 plus the 39 new tests (only the 2 b2view tests that need aiohttp fail locally, on both).

Performance of the fast path against #774 (EPYC 9554, 16 pinned cores, 1 and 8 threads, alternating rounds, #774 built twice for the noise floor; the same picture as an earlier 8-round A/B against main):

  • K a multiple of the chunk (512², 1024², 2048², float32/float64): no measurable change (0.993-1.011, apart from two 8-thread float32 cells where Make the matmul, miniexpr and Menudet prefilters unable to need the GIL #774 against itself moves just as much, 0.85-0.88).
  • 1100² with chunks 512, blocks 128 (the last K chunk has 3 blocks of padding only): 0.48-0.78x the time, since those blocks are now skipped.
  • 1000² with chunks 256, blocks 64: 1.03x (1 thread) and 1.07-1.10x (8 threads). This comes from compressing the output, not from the kernel: with the output stored uncompressed (clevel=0) it is within 1.5%, and a 1024² control with the same chunk/block grid is unchanged (0.99-1.02). Before, the wrong first step left non-zero values in the output's padding rows; now they are zeros, so those output blocks compress (cratio 1.53 -> 1.61 for float32), which costs a bit more than giving up on incompressible blocks.

Not in this PR: when the fast path fails after the prefilter is set, matmul() runs the chunked fallback while the prefilter is still attached to the output, so the fallback's writes run the prefilter again and fail the same way (this was the cause of most exceptions above). With the gate fixed this only remains reachable through real prefilter failures (corrupt chunk, out of memory); removing the prefilter before the fallback would fix it (a follow-up after #774).

I used Claude Code (Anthropic) to find the causes, write the fix and the tests, and run the layout probes and benchmarks; I reviewed the changes.

🤖 Generated with Claude Code

FabianIsensee and others added 3 commits November 6, 2026 12:47
C-Blosc2 runs prefilters on its worker threads when nthreads > 1 (and on the
calling thread without the GIL when compressions release it), but the error
paths of the miniexpr and matmul prefilters raised Python exceptions: Cython
took the GIL for the raise and the wrapper then called PyErr_Occurred()
without it.  With the caller holding the GIL, any error (e.g. an operand
chunk that fails to decompress) deadlocked matmul and legacy-evaluator
expressions and reductions with nthreads > 1; without it, it segfaulted.

Declare miniexpr_prefilter, matmul_prefilter and their helpers
`noexcept nogil` and remove every `raise` (Cython would still wrap one in
an implicit `with gil`); use cdivision there and in the Menudet prefilter
(the zero-division checks took the GIL as well).  aux_miniexpr() records
its error in its worker slot and returns -1; _store_data() and
_prefilter_data() raise it (C-Blosc2 ignores the return code of a
prefilter with a disposable output, i.e. of reductions).  aux_matmul()
frees what it allocated and returns -1.

These prefilters, and the workers of parallel native graph maps, decode
their inputs themselves, so operands with a codec or filter registered
from Python ran that callback without the GIL (segfault; the Menudet
reduction even with one thread).  The routing gates (matmul fast path,
_set_pref_expr(), _set_pref_matmul(), Menudet input eligibility, native
graph layout) now decline such operands, so they take another path.

matmul_prefilter raised "Unsupported dtype" for dtypes other than float and
int; blosc2.matmul() never sent those, but it sent float16 and longdouble
(kind "f"), which the kernels computed as float64: wrong results, and for
float16 reads and writes past the block buffers.  The fast path now takes
float32/float64 only, and _set_pref_matmul() checks the dtypes before the
prefilter is set.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DiPFztdpnKMJkzrjT36BuP
aux_matmul() read src[i] = ndarr.sc.data[chunk_idx] before its first_run
branch moved nchunkA/nchunkB to the output chunk's row and column, so
the first K step of every output chunk multiplied A's and B's first
chunks.  Only layouts with one output chunk came out right; with
non-constant data (the tests used ones()/full()) a 200x200 float64
product with 100x100 chunks was off by up to 226.

Compute the first chunk and block indices before the chunk loop; they
depend only on the shapes.

The kernel also assumes that A, B and the output share one chunk grid
(output chunk rows = A's, output chunk columns = B's, same chunks along
K), but _matmul_can_use_fast_path() did not check that.  Such layouts
gave wrong results, failed in the prefilter (and in the chunked
fallback, which still runs with the prefilter set) or crashed reading
past the operand's chunk list.  Require the shared grid in the gate;
other layouts take the chunked path.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DiPFzsfE5i74cBSxZN13M8
The matmul prefilter multiplies whole blocks.  When K is not a multiple
of the chunk shape, the last blocks along K reach into the chunks'
padding, and A's padding columns times B's padding rows were added to
the result.  The padding is not always zero: blosc2.full()/ones() fill
it with the value, and resize() or partial writes into such arrays keep
old values there.  ones((170, 170)) @ full((170, 170), 3) with 200x200
chunks and 50x50 blocks was off by 90 in every element (main and
4.14.1), and NaN in the padding turned every element into NaN.

Stop at the first block that starts past K (only padding follows), and
zero A's padding columns and B's padding rows in a block that ends past
K.  Both are zeroed so that NaN or inf in one operand's padding cannot
leak through 0 * NaN.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DiPFzsfE5i74cBSxZN13M8

This branch has not been deployed

No deployments
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.

1 participant