Sparse sketching speedups

CLAUDE.md → AGENTS.md

@@ -1,4 +1,4 @@
-# RandBLAS Project Guide for Claude
+# RandBLAS Project Guide for coding agents
 
 ## Project Overview
 
@@ -244,7 +244,7 @@ GitHub Actions workflows test:
 
 All CI tests must pass before merging.
 
-## Working with Claude on RandBLAS
+## Working on RandBLAS with agents
 
 ### Preferred Workflow
 
@@ -316,7 +316,3 @@ All CI tests must pass before merging.
 - GitHub Issues: https://github.com/BallisticLA/RandBLAS/issues
 - Documentation: https://randblas.readthedocs.io/
 - Contact: Project maintainers listed in repository
-
----
-
-*This CLAUDE.md file was created to help Claude Code understand the RandBLAS project structure, conventions, and workflows. Update it as the project evolves.*

RandBLAS/base.hh

@@ -188,6 +188,11 @@ std::ostream &operator<<(
     return out;
 }
 
+inline blas::Layout flipped_layout(const blas::Layout &layout_before) {
+    using blas::Layout;
+    return (layout_before == Layout::RowMajor) ? Layout::ColMajor : Layout::RowMajor;
+}
+
 /**
  * Stores stride information for a matrix represented as a buffer.
  * The intended semantics for a buffer "A" and the conceptualized

RandBLAS/skge.hh

@@ -38,6 +38,7 @@
 #include <stdio.h>
 #include <stdexcept>
 #include <string>
+#include <vector>
 
 #include <cmath>
 #include <typeinfo>
@@ -359,6 +360,78 @@ void rskge3(
 
 namespace RandBLAS::sparse {
 
+// Apply a COOMatrix-represented (possibly transposed) operator `op(S)` on the LEFT
+// by instantiating it in the compressed format chosen by a heuristic, then handing
+// that to left_spmm.
+//
+// Assumes S is in CSC or CSR sort order (never None).
+//
+template <typename T, SignedInteger sint_t>
+void _lskges_compress_and_apply_coo(
+    blas::Layout layout, blas::Op opS, blas::Op opA, int64_t d, int64_t n, int64_t m,
+    T alpha, const sparse_data::COOMatrix<T, sint_t> &S,
+    const T *A, int64_t lda, T beta, T *B, int64_t ldb
+) {
+    using namespace RandBLAS::sparse_data;
+    using blas::Layout; using blas::Op;
+    if (opS == Op::Trans) {
+        auto St = transpose_as_coo(S, /*share_memory=*/true);
+        _lskges_compress_and_apply_coo(layout, Op::NoTrans, opA, d, n, m, alpha, St, A, lda, beta, B, ldb);
+        return;
+    }
+    // left_spmm sees opB = opA (the dense operand's op); recover its post-op layout.
+    Layout layout_opB = (opA == Op::NoTrans) ? layout : flipped_layout(layout);
+    NonzeroSort fmt = coo::coo_spmm_target_format(layout_opB, layout, d, m);
+    // The view borrows S's buffers; S and "ptr" outlive the left_spmm call.
+    std::vector<sint_t> ptr;
+    if (fmt == NonzeroSort::CSR) {
+        auto M = coo_to_csr_view_or_copy(S, ptr);
+        left_spmm(layout, Op::NoTrans, opA, d, n, m, alpha, M, 0, 0, A, lda, beta, B, ldb);
+    } else {
+        auto M = coo_to_csc_view_or_copy(S, ptr);
+        left_spmm(layout, Op::NoTrans, opA, d, n, m, alpha, M, 0, 0, A, lda, beta, B, ldb);
+    }
+    return;
+}
+
+// Apply a COOMatrix-represented (possibly transposed) operator `op(S)` on the RIGHT,
+// analogously to _lskges_compress_and_apply_coo.
+//
+//      right_spmm reduces to a transposed left_spmm, which flips the operand CSR<->CSC.
+//      To get the desired format after that transposition we instantiate S's compressed
+//      representation in the OPPOSITE of the hueristic's requested format.
+//
+// Assumes S is in CSC or CSR sort order (never None).
+//
+template <typename T, SignedInteger sint_t>
+void _rskges_compress_and_apply_coo(
+    blas::Layout layout, blas::Op opS, blas::Op opA, int64_t m, int64_t d, int64_t n,
+    T alpha, const T *A, int64_t lda, const sparse_data::COOMatrix<T, sint_t> &S,
+    T beta, T *B, int64_t ldb
+) {
+    using namespace RandBLAS::sparse_data;
+    using blas::Layout; using blas::Op;
+    if (opS == Op::Trans) {
+        auto St = transpose_as_coo(S, /*share_memory=*/true);
+        _rskges_compress_and_apply_coo(layout, Op::NoTrans, opA, m, d, n, alpha, A, lda, St, beta, B, ldb);
+        return;
+    }
+    Layout trans_layout = flipped_layout(layout);
+    Layout layout_opB   = (opA == Op::NoTrans) ? trans_layout : layout;
+    NonzeroSort reduced = coo::coo_spmm_target_format(layout_opB, trans_layout, d, n);
+    NonzeroSort fmt = (reduced == NonzeroSort::CSR) ? NonzeroSort::CSC : NonzeroSort::CSR;
+    // The view borrows S's buffers; S and "ptr" outlive the right_spmm call.
+    std::vector<sint_t> ptr;
+    if (fmt == NonzeroSort::CSR) {
+        auto M = coo_to_csr_view_or_copy(S, ptr);
+        right_spmm(layout, opA, Op::NoTrans, m, d, n, alpha, A, lda, M, 0, 0, beta, B, ldb);
+    } else {
+        auto M = coo_to_csc_view_or_copy(S, ptr);
+        right_spmm(layout, opA, Op::NoTrans, m, d, n, alpha, A, lda, M, 0, 0, beta, B, ldb);
+    }
+    return;
+}
+
 // MARK: LSKGES
 
 // =============================================================================
@@ -469,7 +542,7 @@ void lskges(
     blas::Op opA,
     int64_t d, // B is d-by-n
     int64_t n, // \op(A) is m-by-n
-    int64_t m, // \op(\mtxS) is d-by-m
+    int64_t m, // \op(submat(S)) is d-by-m
     T alpha,
     const SparseSkOp<T,RNG,sint_t> &S,
     int64_t ro_s,
@@ -480,19 +553,20 @@ void lskges(
     T *B,
     int64_t ldb
 ) {
+    auto [n_rows, n_cols] = dims_before_op(d, m, opS);
     if (S.nnz < 0) {
-        // The operator hasn't been sampled yet. Generate ONLY the needed submatrix of S
-        // (op(submat(S)) is d-by-m). submat(S) is d-by-m if opS == NoTrans, else m-by-d.
-        // We then call left_spmm with offsets (0,0); it applies opS itself.
-        auto Ssub = submatrix_as_coo(S,
-            (opS == blas::Op::NoTrans) ? d : m,
-            (opS == blas::Op::NoTrans) ? m : d,
-            ro_s, co_s);
-        left_spmm(layout, opS, opA, d, n, m, alpha, Ssub, 0, 0, A, lda, beta, B, ldb);
+        auto Ssub = submatrix_as_coo(S, n_rows, n_cols, ro_s, co_s);
+        _lskges_compress_and_apply_coo(layout, opS, opA, d, n, m, alpha, Ssub, A, lda, beta, B, ldb);
         return;
     }
+    bool full_operator = (S.n_rows == n_rows && S.n_cols == n_cols);
     auto Scoo = coo_view_of_skop(S);
-    left_spmm(layout, opS, opA, d, n, m, alpha, Scoo, ro_s, co_s, A, lda, beta, B, ldb);
+    if (full_operator) {
+        _lskges_compress_and_apply_coo(layout, opS, opA, d, n, m, alpha, Scoo, A, lda, beta, B, ldb);
+    } else {
+        // A proper submatrix of an already-sampled operator: let left_spmm handle the offsets.
+        left_spmm(layout, opS, opA, d, n, m, alpha, Scoo, ro_s, co_s, A, lda, beta, B, ldb);
+    }
     return;
 }
 
@@ -606,7 +680,7 @@ inline void rskges(
     blas::Op opA,
     blas::Op opS,
     int64_t m, // B is m-by-d
-    int64_t d, // op(S) is n-by-d
+    int64_t d, // op(submat(S)) is n-by-d
     int64_t n, // op(A) is m-by-n
     T alpha,
     const T *A,
@@ -618,21 +692,20 @@ inline void rskges(
     T *B,
     int64_t ldb
 ) { 
+    auto [n_rows, n_cols] = dims_before_op(n, d, opS);
     if (S.nnz < 0) {
-        // The operator hasn't been sampled yet. Generate ONLY the needed submatrix of S
-        // (op(submat(S)) is n-by-d). submat(S) is n-by-d if opS == NoTrans, else d-by-n.
-        // We then call right_spmm with offsets (0,0); it applies opS itself.
-        auto Ssub = submatrix_as_coo(S,
-            (opS == blas::Op::NoTrans) ? n : d,
-            (opS == blas::Op::NoTrans) ? d : n,
-            ro_s, co_s);
-        right_spmm(layout, opA, opS, m, d, n, alpha, A, lda, Ssub, 0, 0, beta, B, ldb);
+        auto Ssub = submatrix_as_coo(S, n_rows, n_cols, ro_s, co_s);
+        _rskges_compress_and_apply_coo(layout, opS, opA, m, d, n, alpha, A, lda, Ssub, beta, B, ldb);
         return;
     }
+    bool full_operator = (S.n_rows == n_rows && S.n_cols == n_cols);
     auto Scoo = coo_view_of_skop(S);
-    right_spmm(
-        layout, opA, opS, m, d, n, alpha, A, lda, Scoo, ro_s, co_s, beta, B, ldb
-    );
+    if (full_operator) {
+        _rskges_compress_and_apply_coo(layout, opS, opA, m, d, n, alpha, A, lda, Scoo, beta, B, ldb);
+    } else {
+        // A proper submatrix of an already-sampled operator: let right_spmm handle the offsets.
+        right_spmm(layout, opA, opS, m, d, n, alpha, A, lda, Scoo, ro_s, co_s, beta, B, ldb);
+    }
     return;
 }
 

RandBLAS/sparse_data/DevNotes.md

@@ -46,17 +46,23 @@ Sketching dense data with a sparse operator is typically handled with ``sketch_g
 which is defined in ``skge.hh``.
 
 If we call this function with a SparseSkOp object, ``S``, we'd immediately get routed to
-either ``lskges`` or ``rskges``. Here's what would happen after we entered one of those functions:
+either ``lskges`` or ``rskges``. Both are defined in ``skge.hh``. Here's what happens once
+we're inside one of those functions.
 
- 1. If necessary, we'd sample the defining data of ``S`` with ``RandBLAS::fill_sparse(S)``.
+ 1. We get a COO view of ``S`` (sampling its defining data first, if that hasn't happened yet).
 
- 2. We'd obtain a lightweight view of ``S`` as a COOMatrix, and we'd pass that matrix to ``left_spmm``
-    (if inside ``lskges``) or ``right_spmm`` (if inside ``rskges``).
+ 2. If we've been asked for a *proper submatrix* of an already-sampled operator, we hand the
+    COO view and the ``(ro_s, co_s)`` offsets straight to ``[left/right]_spmm`` and let that
+    function deal with the offsets. Otherwise, we have the full operator, and we proceed to compress it.
 
+ 3. Compression and application happens in ``_[l/r]skges_compress_and_apply_coo``. These functions
+    normalize-away the transposition and use a heuristic to choose the compressed format (CSR or CSC).
+    This heuristic can differ from that used if we had called `[left/right]_spmm` directly on
+    a COO matrix.
 
 ## Sketching sparse data with dense operators
 
-If we call ``sketch_sparse`` with a DenseSkOp, ``S``, and a sparse matrix, ``A``, then  we'll get routed to either
+If we call ``sketch_sparse`` with a DenseSkOp, ``S``, and a sparse matrix, ``A``, then we'll get routed to either
 ``lsksp3`` or ``rsksp3``.
 
 From there, we'll do the following.
@@ -72,4 +78,3 @@ From there, we'll do the following.
     Note that the ``l`` and ``r`` in the ``[l/r]sksp3`` function names
     get matched to opposite sides for ``[left/right]_spmm``! This is because all the fancy abstractions in ``S`` have been stripped away by this point in the call sequence, so the "side" that we emphasize in function names changes
     from emphasizing ``S`` to emphasizing ``A``.
-

RandBLAS/sparse_data/base.hh

@@ -31,6 +31,7 @@
 #include "RandBLAS/config.h"
 #include "RandBLAS/base.hh"
 #include <blas.hh>
+#include <vector>
 
 #ifdef __cpp_concepts
 #include <concepts>
@@ -117,6 +118,16 @@ enum class NonzeroSort : char {
     None = 'N'
 };
 
+inline NonzeroSort flipped_nonzerosort(const NonzeroSort &s) {
+    if (s == NonzeroSort::CSC) {
+        return NonzeroSort::CSR;
+    } else if (s == NonzeroSort::CSR) {
+        return NonzeroSort::CSC;
+    } else {
+        return NonzeroSort::None;
+    }
+}
+
 template <SignedInteger sint_t>
 static inline bool increasing_by_csr(sint_t i0, sint_t j0, sint_t i1, sint_t j1) {
     if (i0 > i1) {
@@ -300,6 +311,31 @@ void compressed_ptr_to_sorted_idxs(int64_t num_comp, sint_t1* ptr, int64_t len_i
     }
 }
 
+// Build a compressed layout (CSR or CSC) from a COO that is already sorted in the
+// OPPOSITE compressed order, via an O(nnz + n_groups) counting-sort transpose.
+//
+//   group_idx[nnz] : axis to compress   (rows for a CSR target, cols for CSC)
+//   other_idx[nnz] : the companion axis (cols for CSR, rows for CSC)
+//   n_groups       : number of groups (n_rows for CSR, n_cols for CSC)
+//
+// Outputs (caller-allocated): ptr[n_groups+1], out_other[nnz], out_vals[nnz].
+template <typename T, SignedInteger sint_t>
+void counting_sort_transpose(
+    int64_t nnz, const sint_t* group_idx, const sint_t* other_idx, const T* vals,
+    int64_t n_groups, sint_t* ptr, sint_t* out_other, T* out_vals
+) {
+    for (int64_t g = 0; g <= n_groups; ++g) ptr[g] = 0;
+    for (int64_t e = 0; e < nnz;       ++e) ptr[group_idx[e] + 1] += 1;
+    for (int64_t g = 1; g <= n_groups; ++g) ptr[g] += ptr[g-1];
+    std::vector<sint_t> cursor(ptr, ptr + n_groups);   // running write position per group
+    for (int64_t e = 0; e < nnz; ++e) {
+        sint_t g = group_idx[e];
+        sint_t pos = cursor[g]++;
+        out_other[pos] = other_idx[e];
+        out_vals[pos]  = vals[e];
+    }
+}
+
 // MARK: SparseMatrix
 
 #ifdef __cpp_concepts

RandBLAS/sparse_data/conversions.hh

@@ -34,6 +34,7 @@
 #include "RandBLAS/sparse_data/coo_matrix.hh"
 #include "RandBLAS/sparse_data/csr_matrix.hh"
 #include "RandBLAS/sparse_data/csc_matrix.hh"
+#include <vector>
 
 
 namespace RandBLAS::sparse_data {
@@ -90,6 +91,12 @@ auto coo_to_csc( const COOMatrix<T, sint_t1> &coo, CSCMatrix<T,sint_t2> &csc ) {
         std::copy( coo.vals, coo.vals + coo.nnz, csc.vals    );
         std::copy( coo.rows, coo.rows + coo.nnz, csc.rowidxs );
         return;
+    } else if (coo.sort == NonzeroSort::CSR) {
+        csc.reserve(coo.nnz);
+        counting_sort_transpose(
+            coo.nnz, coo.cols, coo.rows, coo.vals, coo.n_cols, csc.colptr, csc.rowidxs, csc.vals
+        );
+        return;
     } else {
         auto coo_copy = coo.deepcopy();
         coo_copy.sort_arrays(NonzeroSort::CSC);
@@ -112,6 +119,14 @@ void coo_to_csr( const COOMatrix<T, sint_t1> &coo, CSRMatrix<T,sint_t2> &csr ) {
         std::copy( coo.vals, coo.vals + coo.nnz, csr.vals    );
         std::copy( coo.cols, coo.cols + coo.nnz, csr.colidxs );
         return;
+    } else if (coo.sort == NonzeroSort::CSC) {
+        // Opposing order: group the CSC-sorted records by row with an O(nnz)
+        // counting-sort transpose instead of a comparison re-sort.
+        csr.reserve(coo.nnz);
+        counting_sort_transpose(
+            coo.nnz, coo.rows, coo.cols, coo.vals, coo.n_rows, csr.rowptr, csr.colidxs, csr.vals
+        );
+        return;
     } else {
         auto coo_copy = coo.deepcopy();
         coo_copy.sort_arrays(NonzeroSort::CSR);
@@ -120,6 +135,37 @@ void coo_to_csr( const COOMatrix<T, sint_t1> &coo, CSRMatrix<T,sint_t2> &csr ) {
     }
 }
 
+// Represent `coo` as a CSR matrix without copying its structural nonzeros when possible.
+//
+//  If coo is already CSR-sorted (or empty), the result `csr` is a ZERO-COPY view that shares
+//  coo's vals and cols arrays, and has `csr.rowptr` backed by the input `rowptr` std::vector.
+//  If coo is not CSR-sorted, the result `csr` is a new matrix from coo_to_csr.
+//
+template <typename T, SignedInteger sint_t>
+CSRMatrix<T, sint_t> coo_to_csr_view_or_copy(const COOMatrix<T, sint_t> &coo, std::vector<sint_t> &rowptr) {
+    if (coo.nnz == 0 || coo.sort == NonzeroSort::CSR) {
+        rowptr.resize(coo.n_rows + 1);
+        sorted_idxs_to_compressed_ptr(coo.nnz, coo.rows, coo.n_rows, rowptr.data());
+        return CSRMatrix<T, sint_t>(coo.n_rows, coo.n_cols, coo.nnz, coo.vals, rowptr.data(), coo.cols, coo.index_base);
+    }
+    CSRMatrix<T, sint_t> csr(coo.n_rows, coo.n_cols);
+    coo_to_csr(coo, csr);
+    return csr;
+}
+
+// CSC analog of coo_to_csr_view_or_copy.
+template <typename T, SignedInteger sint_t>
+CSCMatrix<T, sint_t> coo_to_csc_view_or_copy(const COOMatrix<T, sint_t> &coo, std::vector<sint_t> &colptr) {
+    if (coo.nnz == 0 || coo.sort == NonzeroSort::CSC) {
+        colptr.resize(coo.n_cols + 1);
+        sorted_idxs_to_compressed_ptr(coo.nnz, coo.cols, coo.n_cols, colptr.data());
+        return CSCMatrix<T, sint_t>(coo.n_rows, coo.n_cols, coo.nnz, coo.vals, coo.rows, colptr.data(), coo.index_base);
+    }
+    CSCMatrix<T, sint_t> csc(coo.n_rows, coo.n_cols);
+    coo_to_csc(coo, csc);
+    return csc;
+}
+
 // MARK: transposes
 //
 // These functions' return values should be treated as const if share_memory = true.

RandBLAS/sparse_data/coo_matrix.hh

@@ -395,7 +395,11 @@ void coo_to_dense(const COOMatrix<T> &spmat, int64_t stride_row, int64_t stride_
             i -= 1;
             j -= 1;
         }
-        MAT(i, j) = spmat.vals[ell];
+        // Accumulate (not overwrite): a COO may carry several records at the same
+        // (i, j), which must sum to that entry's true value. The matrix was zeroed
+        // above, so += is correct, and for a deduplicated COO (each (i,j) once) it
+        // is identical to an assignment.
+        MAT(i, j) += spmat.vals[ell];
     }
     return;
 }