Sparse sketching speedups by rileyjmurray · Pull Request #176 · BallisticLA/RandBLAS

rileyjmurray · 2026-06-19T15:05:55Z

This branch reworks how COO-represented sparse sketching operators are applied to
dense matrices in lskges/rskges. Instead of always converting to CSC and running a
single kernel, we now choose a compressed format (CSR or CSC) per call and hand the
operator to the same per-RHS-column kernels that left_spmm/right_spmm already use
for native CSR/CSC operands. The sampling path was also adjusted so that operators
arrive at the apply step already sorted, which lets the conversion be a zero-copy view
in the common case.

New kernels and dispatching logic

Format-selection heuristic (coo_spmm_target_format in coo_spmm_impl.hh).
The per-column kernels loop the operator's outer dimension once per output column, so
a shorter outer loop is cheaper. For the ColMajor/ColMajor path with a wide
operator (d < m), we now route through the CSR jik kernel; every other case stays
on the CSC kernels. The RowMajor path keeps the bandwidth-saturating CSC kib axpy
kernel.
COO fallback rewritten (apply_coo_via_csc → apply_coo_via_csx). The COO
branch of left_spmm now materializes the operator in the heuristically-chosen format
and dispatches to the matching CSR/CSC kernel, rather than hard-coding CSC.
lskges/rskges custom dispatch. Full operators (sampled or not) go through new
_lskges_compress_and_apply_coo / _rskges_compress_and_apply_coo helpers that pick
the compressed format and apply it. The right-apply helper accounts for right_spmm
reducing to a transposed left_spmm (which flips CSR<->CSC) by requesting the opposite
format. Proper submatrices of an already-sampled operator still fall back to the generic
left_spmm/right_spmm offset handling.

Supporting infrastructure

Zero-copy compressed views (coo_to_csr_view_or_copy /
coo_to_csc_view_or_copy). When a COO is already sorted in the target order, we build
a CSR/CSC that shares the COO's value and index arrays and only allocates the pointer
array — avoiding a deep copy on the hot apply path. Otherwise they fall back to a full
conversion.
O(nnz) counting-sort transpose (counting_sort_transpose). Converting between the
two opposing compressed orders (e.g. a CSC-sorted COO into CSR) now uses a counting-sort
transpose instead of a comparison re-sort. Wired into coo_to_csr / coo_to_csc.
Sampling emits sorted nonzeros. fill_sparse_unpacked now sorts each major-axis
vector as it is sampled (small no-alloc insertion sort for SASO; post-merge sort for
LASO). submatrix_as_coo labels the result as CSR- or CSC-sorted instead of None, so
the apply path can take the zero-copy view.
Small shared helpers: flipped_layout (base.hh), flipped_nonzerosort
(sparse_data/base.hh), and a restrict_to_window helper for windowed COO operators.

Other material changes

Removed the "regular CSC" fixed-nnz-per-column kernel and its runtime probe. It was
previously found to be no faster (and sometimes slower) than the general kernel, and the
probe added a scan over the data on every call.
New benchmark: examples/.../sketch_general_performance.cc with work-normalized
metrics and a --csr-probe mode used to calibrate the dispatch heuristic.
Tests: new CSR SpMM coverage (test_spmm_csr.cc), COO-matrix tests for the
accumulate semantics, and updated sparse-skop tests for the new sort labeling.
Renamed the workspace CLAUDE.md to AGENTS.md and refreshed dev notes.

The SASO sampling path emits each major-axis vector's vec_nnz nonzeros in Fisher-Yates draw order, so a wide SASO's COO is column-ordered but NOT row-ordered within a column -- coo_arrays_determine_sort labels it None, and apply_coo_via_csc then pays a deepcopy + re-sort to CSC on every apply (only the vec_nnz==1 CountSketch path, which short-circuits to the i.i.d. sampler, escaped this). Sort each contiguous vec_nnz block by its major coordinate in the SASO branch of fill_sparse_unpacked, before Phase-2 compaction. A wide SASO now emits CSC-sorted COO (tall SASO -> CSR, which the right_spmm transpose turns into CSC under the row/col swap), so the per-apply re-sort is skipped. The operator is mathematically identical -- only the within-vector storage order changes -- so thread-independence and the RNG stream are preserved. vec_nnz is small, so a no-alloc insertion sort is used. Scoped to operator construction; the public repeated_fisher_yates index utility is untouched. All 435 ctest pass (the structural and KS tests are order-independent). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Benchmarks the full sketch_general() call with real SparseSkOps (wide SASO, CountSketch, wide LASO; left and right sketch; both layouts) -- the structured operators spmm_performance.cc never produces, so it exercises the regular-CSC path and the COO sort/re-sort behavior that sketching actually hits. Reports work-normalized metrics so results are comparable across vec_nnz and SASO/LASO: ns/elt = time/(nnz*work_mult), and model GB/s as a fraction of a measured STREAM-triad ceiling (%STR). The byte model bounds A-read traffic by min(nnz, contract_dim) so sparse LASO operators don't report impossible bandwidth. --no-stream skips calibration; --scaling sweeps thread counts and reports speedup / parallel efficiency / %STR per count. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Design docs and measured baselines for the sparse-sketching SPMM work: - sparsesketching_opt.md: proposal framed on row- vs column-structured operators (transposition is free), with candidate kernels (CountSketch, bounded-CSR for wide LASO, implicit +/-1 values) and a benchmark-first recommendation. - sparsesketching_bench_plan.md: plan for the sketch_general benchmark, plus the measured finding that motivated the sampling fix. - sketch_general_results.txt: pre-fix baseline (raw timings). - sketch_general_results_postfix.txt: post-fix run with the work-normalized metric columns. These are working/analysis artifacts; drop if they shouldn't live in-tree. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Completes the construction-time sort started for SASO. The LASO branch of fill_sparse_unpacked emits each major-axis vector's (<= vec_nnz) merged survivors in hash/sample order, so the COO is grouped by the minor axis but unordered within each vector (coo-sort = None). Sorting each merged block by its major coordinate makes a wide LASO CSR-sorted (cols ascending within each row) and a tall LASO CSC-sorted. For a wide LASO this natural order also matches the ColMajor dispatch preference (CSR for wide operators), so apply_coo_via_csc skips its per-apply re-sort -- which also makes the ColMajor CSR routing an unconditional win even at tiny nnz, where the re-sort had previously cost more than the kernel gain. Operator is mathematically unchanged; only within-vector storage order changes. All 435 ctest pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

apply_coo_via_csc previously always converted COO to CSC. The per-RHS-column kernels loop the operator's outer dimension once per output column: CSC jki's outer loop is n_cols (= m), CSR jik's is n_rows (= d). For a WIDE operator (d < m) in the ColMajor path, the shorter CSR outer loop measured 1.1-2x faster than CSC jki for both SASO and LASO, with no regression (see the --csr-probe benchmark mode). So we now sort the COO to CSR and dispatch apply_csr_jik_p11 when ColMajor and d < m; RowMajor keeps the bandwidth-saturating CSC kib axpy kernel, and tall operators keep CSC. A ColMajor strided-axpy 'CSC kib' kernel was also prototyped and rejected: it regressed dense-column SASO ~3x because ColMajor rows are strided. Conversion cost is unchanged (the COO is sorted either way). A wide SASO that sampling sorted to CSC is still re-sorted to CSR here; a follow-up will build the CSR rowptr via an O(nnz) counting-sort transpose to remove that. All 435 ctest pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Builds each operator's CSC and CSR representations once and times left_spmm directly on each (no per-call COO re-sort), in both layouts, to isolate the kernel: does CSR jik (ColMajor outer loop over d rows) beat CSC jki (outer loop over m mostly-empty columns)? CSC kib (RowMajor) is the fast baseline. This is the harness that justified the ColMajor CSR routing. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

When apply_coo_via_csc needs a format the COO isn't sorted for (e.g. a wide SASO sampled CSC-sorted, but the ColMajor path wants CSR), it previously deep-copied and ran an O(nnz log nnz) comparison re-sort. For a full operator already sorted in the OPPOSITE compressed order, build the wanted format directly with an O(nnz + dim) counting-sort transpose instead. Iterating the source in its sorted order leaves the companion index ascending within each group, so the result is properly sorted in the target order. This removes the per-apply re-sort tax that had capped wide-SASO ColMajor at ~1.1x: it now reaches ~1.45x (matching the CSR jik kernel's intrinsic advantage over CSC jki). The submatrix and unsorted (None) cases still take the general deepcopy + sort path. All 435 ctest pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…nspose) Regenerate sketch_general_results_postfix.txt to capture the combined branch state: construction-time SASO/LASO sorting, ColMajor wide -> CSR jik routing, and the O(nnz) counting-sort transpose. ColMajor warm apply vs the pre-change baseline: wide SASO ~1.45x, wide LASO 1.2-2.1x, no RowMajor regression, 0 FAIL. sketch_general_results.txt is left frozen as the pre-change baseline. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

--------------------------------------------------------------- Store LASO operators "regular"; add a gated regular-CSR kernel Regularize long-axis-major (LASO) sketching operators so every major-axis vector is stored with exactly vec_nnz records, matching the fixed-nnz structure short-axis-major (SASO) operators already have. This is a storage-only change: a coordinate sampled c times is kept as c duplicate records each scaled by 1/sqrt(c), which sum to the merged value sqrt(c)*s, so the operator is bit-for-bit the same matrix. RNG draws are unchanged, so thread-count reproducibility holds. nnz becomes exactly full_nnz. Because duplicate (row,col) records can now appear, coo_to_dense (and the test densifier) accumulate with += instead of overwriting; SpMM kernels and COO->CSR/CSC conversions already accumulate / preserve duplicates. trsm does not accept duplicates -- noted in DevNotes. Adds a regular-CSR SpMV kernel (apply_regular_csr_to_vector_ik), the CSR analogue of apply_regular_csc_to_vector_ki, plus detection in apply_csr_jik_p11. A back-to-back regular-vs-general A/B (new rows in the benchmark's --csr-probe) shows it is NEUTRAL-to-SLOWER on Apple M3 / NEON (the eliminated rowptr loads were never a bottleneck; 2-wide doubles can't exploit the fixed inner trip count). The dispatch is therefore GATED OFF by default behind RandBLAS_ENABLE_REGULAR_CSR_KERNEL; the default build keeps the general kernel (no regression). The kernel is kept in-tree for evaluation on wide-vector Arm (SVE), where a known-length predicated inner loop may pay off. Also corrects an assumption: wide SASO is CSC-regular (fixed nnz per column), not CSR-regular, so its CSR form has variable row nnz and does not reach the regular CSR kernel; only wide LASO (and tall SASO) are CSR-regular. Tests: operator-preservation (regular vs merged densify identically), regular structure / nnz==full_nnz, coo_to_dense summing, and the regular CSR kernel directly (incl. a duplicated colidx). 439/439 pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…eed to optimize this). Have "COO kernel" dispatch to either CSC or CSR (not very important to optimize).

rileyjmurray · 2026-06-29T04:07:01Z

+// Format choice for the COO->compressed SpMM dispatch. The per-RHS-column kernels loop
+// the operator's OUTER dimension once per output column: CSC jki's outer loop is n_cols
+// (= m), CSR jik's is n_rows (= d). In the ColMajor path, routing a WIDE operator (d < m)
+// through the CSR jik kernel (shorter outer loop) measured 1.1-2x faster than CSC jki for
+// both SASO and LASO, with no regression -- so we prefer CSR there. The RowMajor path keeps
+// the bandwidth-saturating CSC kib axpy kernel. See the --csr-probe mode of
+// examples/.../sketch_general_performance.cc for measurements.
+inline NonzeroSort coo_spmm_target_format(
+    blas::Layout layout_opB, blas::Layout layout_C, int64_t d, int64_t m
+) {
+    bool col_major  = (layout_opB == blas::Layout::ColMajor && layout_C == blas::Layout::ColMajor);
+    bool prefer_csr = col_major && (d < m);
+    return prefer_csr ? NonzeroSort::CSR : NonzeroSort::CSC;
+}


Reinvestigate this; probably define a separate version when sketching.

rileyjmurray and others added 10 commits June 17, 2026 22:19

remove old performance log

a95a5b4