Skip to content

Kernel speed improvements#173

Merged
rileyjmurray merged 11 commits into
mainfrom
kernel-speed
May 30, 2026
Merged

Kernel speed improvements#173
rileyjmurray merged 11 commits into
mainfrom
kernel-speed

Conversation

@rileyjmurray

@rileyjmurray rileyjmurray commented May 30, 2026

Copy link
Copy Markdown
Contributor

Collects sparse-SpMM kernel/benchmark work and first-class MKL sparse-BLAS support
(plus the CI plumbing to actually build and test it), Non-MKL builds and the full test suite
are unaffected throughout; the MKL sparse path is now exercised in CI for the first time.

Sparse SpMM kernels

Speed up CSR ColMajor left_spmm ~1.8–2× (5ae2093). apply_csr_to_vector_ik the hot inner kernel behind apply_csr_jik_p11) gets a #pragma omp simd reduction on the inner sparse dot-product — the win is FP reassociation (multiple partial sums break the serial accumulator chain; isolated, this is the entire effect of ffast-math on M3, while -ffp-contract=fast and -mcpu=native do nothing). A unit-stride if constexpr specialization (incv==incAv==1, the ColMajor case) drops index multiplies for another ~5–8%. Also fixes a latent intint64_t truncation of the column index in the general-stride path. Measured (M3, density 0.01, min µs):

config OMP=1 OMP=4
2000² d=200 3886 → 2169 1188 → 621
2000² d=2000 44197 → 22156 12730 → 6959

CSC was evaluated but its scatter kernel is latency-bound and saw no benefit, so it's left unchanged.

MKL sparse BLAS

  • Handle CSC in mkl_left_spmm as a CSR-of-transpose view (7fb84cd).
    MKL's mkl_sparse_?_mm rejects CSC, so CSC previously fell back to the
    hand-rolled scatter kernel. But a CSC matrix's (colptr, rowidxs, vals) arrays
    are a CSR matrix for Aᵀ (exactly transpose_as_csr's zero-copy view), so
    we build that CSR-of-Aᵀ handle and flip the MKL operation: NoTrans → TRANSPOSE,
    Trans → NON_TRANSPOSE. Nothing is copied (MKL supports SPARSE_OPERATION_TRANSPOSE
    on CSR natively). With CSC handled, the opA=Trans branch in left_spmm no
    longer needs to pre-route a transposed CSR to MKL to dodge a CSC fallback — that
    special case is removed; the transpose-and-recurse path now reaches MKL for
    every format.
  • CI: build and actually test the MKL sparse path (02ffeaa, 6e1e7a6, 1290756).
    Debian's libmkl-dev puts headers under /usr/include/mkl with no MKLROOT,
    so blaspp never defined BLAS_HAVE_MKL and RandBLAS logged
    Checking for MKL sparse BLAS ... FALSE — the entire RandBLAS_HAS_MKL path
    (spgemm + accelerated spmm, including the new CSC handling) was silently
    skipped. The setup-randblas-deps action now discovers the MKL header dir and
    exports MKLROOT/CPATH, MKL_sparse.cmake's find_path handles both oneAPI
    and Debian layouts, and the MKL job builds blaspp/lapackpp as ILP64 (so
    MKL_INT is 64-bit, matching RandBLAS's int64_t sparse indices required by
    spgemm's static_assert). ILP64 is scoped to the MKL backend only, to avoid
    breaking the LP64 OpenBLAS jobs.

SpMM benchmark — examples/simple-kernel-benchmarks/spmm_performance.cc

  • Route every kernel through left_spmm using both dense layouts (568e89b).
    The old benchmark used left_spmm (ColMajor) + right_spmm (ColMajor); since
    right_spmm only reaches the RowMajor kernels by transposing the sparse
    operand, the same coverage comes from left_spmm with ColMajor and RowMajor
    dense operands (ColMajor → jik/jki, RowMajor → ikb/kib; COO delegates to
    CSC). One logical dense matrix is stored in both layouts so the two are directly
    comparable against a single densify+GEMM oracle.
  • Add opB=Trans (C = A·Bᵀ) coverage (6598d08). This is the only
    opA=NoTrans path MKL declines, and it exercises the strided (incv/incAv≠1)
    branch of the CSR kernel that the unit-stride path never hits — so the kernel
    speedup above widens the measured NoTrans-vs-Trans gap. Two extra tables +
    6 correctness checks; the transposed operand reuses the existing buffers (the
    transpose of a ColMajor n×d buffer is the RowMajor buffer).
  • Consistent (A, B, C) = (sparse, dense, result) notation (a857648),
    matching the left_spmm/spmm API (the file previously mixed (S, A, B)).

Misc

  • Cache-efficiency improvement to potrf_upper_sequential (34c6dad) in the
    test LAPACK-like helpers.
  • Fix a -Wsign-compare warning in test_exceptions.cc (0bfe539) — a
    misplaced (int) cast compared find()'s size_type to npos; swept all
    test + example sources, this was the only occurrence.

Testing

  • ctest 429/429 on non-MKL builds (macOS + Linux OpenBLAS).
  • The MKL sparse path — spgemm and accelerated spmm including the new CSC
    handling — is now compiled and run by the linux-gcc-mkl (ILP64) job; this is
    the first CI coverage of RandBLAS_HAS_MKL.
  • CSR kernel speedup benchmarked on Apple M3 (numbers above).

rileyjmurray and others added 11 commits May 29, 2026 10:26
…tride branch

apply_csr_to_vector_ik is the hot inner kernel for ColMajor left_spmm
(via apply_csr_jik_p11). Two changes give ~1.8-2x on that path on an
Apple M3 (single- and multi-threaded), with all 429 ctest cases and the
densify+GEMM reference checks still passing:

- Add `#pragma omp simd reduction(+:Av_i_diff)` to the inner sparse dot
  product. The win is FP *reassociation*: it breaks the serial dependency
  chain on the scalar accumulator (multiple partial sums). Measured in
  isolation this is the whole effect of -ffast-math here; -ffp-contract=fast
  (FMA only) and -mcpu=native do nothing on M3's NEON. The pragma is ignored
  without OpenMP -- the loop just runs unvectorized -- matching the file's
  existing unguarded omp usage, so no RandBLAS_HAS_OpenMP guard is needed.

- Specialize the kernel on unit stride (incv == incAv == 1, the ColMajor
  case) via an `if constexpr` helper so the inner reduction carries no index
  multiplies. ~5-8% on top of the reassociation win; flat for very sparse
  rows. CSC was tried too but its scatter kernel is latency-bound and saw
  no benefit, so it is left unchanged.

Also fixes a latent 32-bit truncation of the column index (int j ->
int64_t j) in the general-stride path.

Measured CSR ColMajor left_spmm, min us (M3, density 0.01):
  S=2000^2 d=200:   OMP=1 3886->2169   OMP=4 1188->621
  S=2000^2 d=2000:  OMP=1 44197->22156 OMP=4 12730->6959

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…eft_spmm

The benchmark previously used left_spmm (ColMajor) and right_spmm (ColMajor) to
cover the sparse SpMM kernels. right_spmm only reaches the RowMajor kernels by
transposing the sparse operand (CSR<->CSC view) and flipping the layout, so the
same coverage is obtained by calling left_spmm directly with both dense layouts:

  ColMajor dense -> apply_csr_jik_p11          / apply_csc_jki_p11
  RowMajor dense -> apply_csr_ikb_p1b_rowmajor / apply_csc_kib_1p1_rowmajor
  (COO delegates to the CSC kernels in both layouts)

This hits exactly the kernels that left_spmm + right_spmm reached before, but
measures each one directly without the right_spmm transpose wrapper. The old
"right CSR" path is now "RowMajor CSC", etc. (the operator transpose swaps the
CSR/CSC label).

Changes:
- Two tables per config: ColMajor dense (jik/jki) and RowMajor dense (ikb/kib),
  each timing CSR/CSC/COO through left_spmm only.
- One logical dense A is stored in both layouts (cheap repack), so the two
  layouts compute the identical product B = S*A and are directly comparable;
  the summary reports best ColMajor vs best RowMajor.
- The ColMajor densify+GEMM reference doubles as the correctness oracle for
  both layouts.

The only remaining non-left_spmm call is the MKL-vs-hand-rolled CSR comparison
(handrolled_left_spmm_csr), kept under #if defined(RandBLAS_HAS_MKL): on an MKL
build the dispatch always selects MKL, so the hand-rolled kernel is otherwise
unreachable. It compiles out when MKL is disabled.

Builds warning-clean; correctness checks pass for square/tall/wide/sparse shapes.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The benchmark drove left_spmm only with opB=NoTrans, so the jik/jci kernels
were measured only in their unit-stride (ColMajor B, ColMajor C) configuration.
The opB=Trans path -- computing C = S*A^T with the dense operand fed transposed
-- also routes to jik/jci, but with a STRIDED dense access (gather when C is
ColMajor, strided store when C is RowMajor), and it is the only sparse-NoTrans
path MKL declines (mkl_left_spmm returns false for opB != NoTrans), so it stays
hand-rolled even on MKL builds. None of that was covered.

Adds two transposed-dense tables (ColMajor C / RowMajor C) timing CSR/CSC/COO
through left_spmm with opB=Trans, plus 6 correctness checks. A^T (d x n) needs
no new storage: the transpose of the ColMajor n-by-d dense buffer is exactly the
RowMajor buffer, so the existing A_cm/A_rm are reused as the transposed operand.
All paths compute the same product S*A and check against the same oracle.

This also exercises the general-stride (incv/incAv != 1) branch of
apply_csr_to_vector_ik, which the unit-stride NoTrans ColMajor path does not --
e.g. with the committed kernel optimization, CSR ColMajor NoTrans speeds up
while the strided opB=Trans path does not, widening the measured gap.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
MKL's mkl_sparse_?_mm does not accept CSC handles, so mkl_left_spmm previously
returned false for CSC and fell back to the hand-rolled scatter kernel. But a
CSC matrix's (colptr, rowidxs, vals) arrays ARE a CSR matrix for A^T -- exactly
transpose_as_csr's zero-copy view. So build that CSR-of-A^T handle and flip the
MKL operation (NoTrans <-> Trans), which leaves op(A)*B unchanged. MKL supports
SPARSE_OPERATION_TRANSPOSE on CSR natively, so nothing is copied:
  - left CSC (opA=NoTrans): CSR(A^T) handle + TRANSPOSE.
  - CSC via opA=Trans (the right_spmm path): CSR(A^T) + NON_TRANSPOSE (fast path).

With CSC handled, the opA=Trans branch in left_spmm no longer needs to pre-route
a transposed CSR to MKL to dodge a CSC fallback: transposing the sparse operand
to a CSR<->CSC view and recursing with NoTrans now reaches MKL for every format
(the recursive NoTrans CSR-of-transpose lands on the identical CSR(A)+TRANSPOSE
call). Removed that special case.

NOTE: the MKL code paths are not compiled or run on the Apple Silicon dev
machine, so this is reasoned + mirrors the already-tested transpose_as_csr view
but needs validation on an MKL build, where test_matmul_cores' CSC spmm cases
cover it. The non-MKL build and full test suite (429/429) are unaffected.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Debian packages MKL headers under /usr/include/mkl and sets no MKLROOT,
so blaspp's mkl_version.cc compile test never finds mkl.h and never
defines BLAS_HAVE_MKL. RandBLAS then concludes blaspp wasn't built with
MKL and skips sparse-BLAS support entirely (the linux-gcc-mkl job logged
"Checking for MKL sparse BLAS ... FALSE" with no header search).

- setup-randblas-deps action: discover the MKL header dir and export
  MKLROOT + CPATH so blaspp defines BLAS_HAVE_MKL and RandBLAS's
  find_path(mkl_spblas.h) succeeds.
- Bump the blaspp cache key so the pre-fix install (built without
  BLAS_HAVE_MKL) is rebuilt rather than reused.
- Broaden MKL_sparse.cmake's find_path to handle both oneAPI
  ($MKLROOT/include) and Debian-apt (/usr/include/mkl) layouts.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
RandBLAS sparse matrices default to int64_t indices. The MKL spgemm path
(sparse x sparse -> dense, MKL-only with no fallback) static_asserts that
the index type matches MKL_INT, so it only compiles when MKL is ILP64.
Building blaspp with the default lp64 MKL made MKL_INT 32-bit and broke
test_spgemm. Standardize the whole CI stack on 64-bit integers instead.

- Pass -Dblas_int=int64 to both the blaspp and lapackpp configures.
- Switch the Linux OpenBLAS package to the ILP64 variant
  (libopenblas64-openmp-dev) so blaspp can find a working int64 BLAS.
- Bump the blaspp and lapackpp cache keys so the prior lp64 installs are
  rebuilt rather than reused.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The blanket -Dblas_int=int64 broke the OpenBLAS jobs: blaspp's BLASFinder
links -lopenblas, but the ILP64 apt package ships libopenblas64 with a
64_ symbol suffix that blaspp's main branch can't consume.

ILP64 is only actually needed in the MKL job: it makes MKL_INT 64-bit,
which RandBLAS's MKL spgemm path (test_spgemm, compiled only under
RandBLAS_HAS_MKL) requires to match its int64_t sparse indices. The other
backends never compile that path, and blaspp's public API is int64_t
regardless of blas_int, so RandBLAS is unaffected by their int width.

- Pass -Dblas_int=int64 to blaspp/lapackpp only for the mkl backend.
- Revert the Linux OpenBLAS package to libopenblas-openmp-dev (LP64).
- Bump the blaspp/lapackpp cache keys (mklfix2) to drop the prior caches.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@rileyjmurray rileyjmurray changed the title WIP: kernel speed improvements Kernel speed improvements May 30, 2026
@rileyjmurray rileyjmurray merged commit cefb9b3 into main May 30, 2026
16 checks passed
@rileyjmurray rileyjmurray deleted the kernel-speed branch May 30, 2026 15:41
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