diff --git a/CLAUDE.md b/AGENTS.md similarity index 98% rename from CLAUDE.md rename to AGENTS.md index 17a47240..af91f65e 100644 --- a/CLAUDE.md +++ b/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.* diff --git a/RandBLAS/base.hh b/RandBLAS/base.hh index 4cc2ff94..f7f60c21 100644 --- a/RandBLAS/base.hh +++ b/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 diff --git a/RandBLAS/skge.hh b/RandBLAS/skge.hh index 9bbd4e37..aab9f58c 100644 --- a/RandBLAS/skge.hh +++ b/RandBLAS/skge.hh @@ -38,6 +38,7 @@ #include #include #include +#include #include #include @@ -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 +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 &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 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 +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 &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 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 &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; } diff --git a/RandBLAS/sparse_data/DevNotes.md b/RandBLAS/sparse_data/DevNotes.md index 24d52336..c575abf2 100644 --- a/RandBLAS/sparse_data/DevNotes.md +++ b/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``. - diff --git a/RandBLAS/sparse_data/base.hh b/RandBLAS/sparse_data/base.hh index 7eb4df2a..ecf98471 100644 --- a/RandBLAS/sparse_data/base.hh +++ b/RandBLAS/sparse_data/base.hh @@ -31,6 +31,7 @@ #include "RandBLAS/config.h" #include "RandBLAS/base.hh" #include +#include #ifdef __cpp_concepts #include @@ -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 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 +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 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 diff --git a/RandBLAS/sparse_data/conversions.hh b/RandBLAS/sparse_data/conversions.hh index 77a5a7c6..819e2fac 100644 --- a/RandBLAS/sparse_data/conversions.hh +++ b/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 namespace RandBLAS::sparse_data { @@ -90,6 +91,12 @@ auto coo_to_csc( const COOMatrix &coo, CSCMatrix &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 &coo, CSRMatrix &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 &coo, CSRMatrix &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 +CSRMatrix coo_to_csr_view_or_copy(const COOMatrix &coo, std::vector &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(coo.n_rows, coo.n_cols, coo.nnz, coo.vals, rowptr.data(), coo.cols, coo.index_base); + } + CSRMatrix csr(coo.n_rows, coo.n_cols); + coo_to_csr(coo, csr); + return csr; +} + +// CSC analog of coo_to_csr_view_or_copy. +template +CSCMatrix coo_to_csc_view_or_copy(const COOMatrix &coo, std::vector &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(coo.n_rows, coo.n_cols, coo.nnz, coo.vals, coo.rows, colptr.data(), coo.index_base); + } + CSCMatrix 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. diff --git a/RandBLAS/sparse_data/coo_matrix.hh b/RandBLAS/sparse_data/coo_matrix.hh index 47eded5f..33b15f88 100644 --- a/RandBLAS/sparse_data/coo_matrix.hh +++ b/RandBLAS/sparse_data/coo_matrix.hh @@ -395,7 +395,11 @@ void coo_to_dense(const COOMatrix &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; } diff --git a/RandBLAS/sparse_data/coo_spmm_impl.hh b/RandBLAS/sparse_data/coo_spmm_impl.hh index e2781453..aad21b18 100644 --- a/RandBLAS/sparse_data/coo_spmm_impl.hh +++ b/RandBLAS/sparse_data/coo_spmm_impl.hh @@ -34,6 +34,7 @@ #include "RandBLAS/sparse_data/base.hh" #include "RandBLAS/sparse_data/coo_matrix.hh" #include "RandBLAS/sparse_data/csc_spmm_impl.hh" +#include "RandBLAS/sparse_data/csr_spmm_impl.hh" #include #include #if defined(RandBLAS_HAS_OpenMP) @@ -49,8 +50,51 @@ using RandBLAS::SignedInteger; #endif +// 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; +} + +// Restrict A0 to the requested d-by-m window at (ro_a, co_a): copy the in-window nonzeros into +// a fresh, memory-owning COO, shifted to local coordinates. We over-allocate to A0.nnz (the +// worst case) so the compaction is a single pass; the result's logical nnz is the in-window +// count, and its sort label is recomputed so the conversion below can hit a fast path. template -static void apply_coo_via_csc( +static COOMatrix restrict_to_window( + const COOMatrix &A0, int64_t d, int64_t m, int64_t ro_a, int64_t co_a +) { + COOMatrix sub(d, m); + if (A0.nnz == 0) + return sub; + sub.reserve(A0.nnz); + int64_t write = 0; + for (int64_t i = 0; i < A0.nnz; ++i) { + auto r = A0.rows[i] - ro_a; + auto c = A0.cols[i] - co_a; + if (0 <= r && r < d && 0 <= c && c < m) { + sub.rows[write] = r; + sub.cols[write] = c; + sub.vals[write] = A0.vals[i]; + write += 1; + } + } + sub.nnz = write; + sub.sort = coo_arrays_determine_sort(sub.nnz, sub.rows, sub.cols); + return sub; +} + +template +static void apply_coo_via_csx( T alpha, blas::Layout layout_B, blas::Layout layout_C, @@ -67,40 +111,36 @@ static void apply_coo_via_csc( ) { randblas_require(A0.index_base == IndexBase::Zero); + // The exact d-by-m operator: A0 itself when it already has those dimensions, else an owning + // copy restricted to the window. Binding "op" by reference lets the common full-operator + // path use A0 (with its existing sort) at zero cost; "windowed" is then an empty placeholder. bool submatrix = (A0.n_rows != d) || (A0.n_cols != m); - if (submatrix || A0.sort != NonzeroSort::CSC) { - auto A1 = A0.deepcopy(); - auto new_nnz = A1.nnz; - if (submatrix) { - int64_t write = 0; - for (int64_t i = 0; i < A1.nnz; ++i) { - auto r = A1.rows[i] - ro_a; - auto c = A1.cols[i] - co_a; - if (0 <= r && r < d && 0 <= c && c < m) { - A1.rows [write] = r; - A1.cols [write] = c; - A1.vals [write] = A1.vals[i]; - write += 1; - } - } - new_nnz = write; - } - COOMatrix A2(d, m, new_nnz, A1.vals, A1.rows, A1.cols, false); - A2.sort_arrays(NonzeroSort::CSC); - apply_coo_via_csc(alpha, layout_B, layout_C, d, n, m, A2, 0, 0, B, ldb, C, ldc); - return; - } - auto colptr = new sint_t[m+1]; - sorted_idxs_to_compressed_ptr( A0.nnz, A0.cols, m, colptr ); - CSCMatrix A_csc( d, m, A0.nnz, A0.vals, A0.rows, colptr ); - if (layout_B == layout_C && layout_B == blas::Layout::RowMajor) { - using RandBLAS::sparse_data::csc::apply_csc_kib_1p1_rowmajor; - apply_csc_kib_1p1_rowmajor(alpha, n, A_csc, B, ldb, C, ldc); + COOMatrix windowed = submatrix + ? restrict_to_window(A0, d, m, ro_a, co_a) + : COOMatrix(0, 0); + const COOMatrix &op = submatrix ? windowed : A0; + if (op.nnz == 0) + return; // structurally-zero operator; C was already scaled by beta upstream. + + // Choose the compressed format by the dispatch heuristic, materialize "op" in that format + // (a zero-copy view when its sort already matches, else an O(nnz) conversion), and call the + // matching per-RHS-column kernel -- exactly the kernels left_spmm uses for a native CSR/CSC + // operand. coo_spmm_target_format only selects CSR for ColMajor/ColMajor, which is why only + // apply_csr_jik_p11 appears here; CSC is the catch-all and keeps its RowMajor axpy kernel. + // Keep this in sync with coo_spmm_target_format and with left_spmm's native dispatch. + NonzeroSort target = coo_spmm_target_format(layout_B, layout_C, d, m); + std::vector ptr; // backs a zero-copy view's ptr array; must outlive the kernel call + if (target == NonzeroSort::CSR) { + auto M = RandBLAS::sparse_data::coo_to_csr_view_or_copy(op, ptr); + csr::apply_csr_jik_p11(alpha, layout_B, layout_C, d, n, m, M, B, ldb, C, ldc); } else { - using RandBLAS::sparse_data::csc::apply_csc_jki_p11; - apply_csc_jki_p11(alpha, layout_B, layout_C, n, A_csc, B, ldb, C, ldc); + auto M = RandBLAS::sparse_data::coo_to_csc_view_or_copy(op, ptr); + if (layout_B == layout_C && layout_B == blas::Layout::RowMajor) { + csc::apply_csc_kib_1p1_rowmajor(alpha, n, M, B, ldb, C, ldc); + } else { + csc::apply_csc_jki_p11(alpha, layout_B, layout_C, n, M, B, ldb, C, ldc); + } } - delete [] colptr; return; } diff --git a/RandBLAS/sparse_data/csc_spmm_impl.hh b/RandBLAS/sparse_data/csc_spmm_impl.hh index 8d2509f2..1f85e4de 100644 --- a/RandBLAS/sparse_data/csc_spmm_impl.hh +++ b/RandBLAS/sparse_data/csc_spmm_impl.hh @@ -71,30 +71,6 @@ static void apply_csc_to_vector_ki( } } -template -static void apply_regular_csc_to_vector_ki( - T alpha, - // data for "regular CSC": CSC with fixed nnz per col, - // which obviates the requirement for colptr. - const T *vals, - const sint_t *rowidxs, - int64_t col_nnz, - // input-output vector data - int64_t len_v, - const T *v, - int64_t incv, // stride between elements of v - T *Av, // Av += A * v. - int64_t incAv // stride between elements of Av -) { - for (int64_t c = 0; c < len_v; ++c) { - T scale = alpha * v[c * incv]; - for (int64_t j = c * col_nnz; j < (c + 1) * col_nnz; ++j) { - int64_t row = rowidxs[j]; - Av[row * incAv] += (vals[j] * scale); - } - } -} - template static void apply_csc_jki_p11( T alpha, @@ -111,10 +87,6 @@ static void apply_csc_jki_p11( auto m = A.n_cols; - bool fixed_nnz_per_col = true; - for (int64_t ell = 2; (ell < m + 1) && fixed_nnz_per_col; ++ell) - fixed_nnz_per_col = (A.colptr[1] + A.colptr[ell-1]) == A.colptr[ell]; - auto s = layout_to_strides(layout_B, ldb); auto B_inter_col_stride = s.inter_col_stride; auto B_inter_row_stride = s.inter_row_stride; @@ -127,21 +99,12 @@ static void apply_csc_jki_p11( for (int64_t j = 0; j < n; j++) { const T* B_col = &B[B_inter_col_stride * j]; T* C_col = &C[C_inter_col_stride * j]; - if (fixed_nnz_per_col) { - apply_regular_csc_to_vector_ki( - alpha, - A.vals, A.rowidxs, A.colptr[1], - m, B_col, B_inter_row_stride, - C_col, C_inter_row_stride - ); - } else { - apply_csc_to_vector_ki( - alpha, - A.vals, A.rowidxs, A.colptr, - m, B_col, B_inter_row_stride, - C_col, C_inter_row_stride - ); - } + apply_csc_to_vector_ki( + alpha, + A.vals, A.rowidxs, A.colptr, + m, B_col, B_inter_row_stride, + C_col, C_inter_row_stride + ); } return; } diff --git a/RandBLAS/sparse_data/spmm_dispatch.hh b/RandBLAS/sparse_data/spmm_dispatch.hh index 2c6344fe..30bf9a82 100644 --- a/RandBLAS/sparse_data/spmm_dispatch.hh +++ b/RandBLAS/sparse_data/spmm_dispatch.hh @@ -148,8 +148,8 @@ void left_spmm( // Fallback: hand-rolled sparse kernels. if constexpr (is_coo) { - using RandBLAS::sparse_data::coo::apply_coo_via_csc; - apply_coo_via_csc(alpha, layout_opB, layout_C, d, n, m, A, ro_a, co_a, B, ldb, C, ldc); + using RandBLAS::sparse_data::coo::apply_coo_via_csx; + apply_coo_via_csx(alpha, layout_opB, layout_C, d, n, m, A, ro_a, co_a, B, ldb, C, ldc); } else if constexpr (is_csc) { if (layout_opB == Layout::RowMajor && layout_C == Layout::RowMajor) { using RandBLAS::sparse_data::csc::apply_csc_kib_1p1_rowmajor; diff --git a/RandBLAS/sparse_skops.hh b/RandBLAS/sparse_skops.hh index 9284d25d..70e4c55f 100644 --- a/RandBLAS/sparse_skops.hh +++ b/RandBLAS/sparse_skops.hh @@ -194,7 +194,7 @@ struct SparseDist { /// If \math{\ttt{major_axis} = \ttt{Long}}, then \math{\mtxx} has *at most* \math{\vecnnz} nonzero /// entries. The locations of the nonzeros are determined by sampling uniformly /// with replacement from \math{\\{0,\ldots,k-1\\}.} - /// If index \math{j} occurs in the sample \math{\ell} times, then + /// If index \math{j} occurs in the sample \math{\ell} times, then /// \math{\mtxx_j} will equal \math{\sqrt{\ell}} with probability 1/2 and /// \math{-\sqrt{\ell}} with probability 1/2. /// @@ -634,6 +634,24 @@ state_t fill_sparse_unpacked( sint_t* idxs_major = major_is_rows ? rows : cols; sint_t* idxs_minor = major_is_rows ? cols : rows; + // Sort a contiguous block of "len" nonzeros into ascending major-coordinate order, + // moving the parallel (major, vals) pair together. len is at most vec_nnz, which is + // small, so use a no-alloc insertion sort. The idxs_minor array is constant across + // these blocks, so the helper doesn't need to look at it. + auto sort_block_by_major = [](sint_t* blk_major, T* blk_vals, int64_t len) { + for (int64_t a = 1; a < len; ++a) { + sint_t key = blk_major[a]; + T v = blk_vals[a]; + int64_t c = a - 1; + for (; c >= 0 && blk_major[c] > key; --c) { + blk_major[c+1] = blk_major[c]; + blk_vals[c+1] = blk_vals[c]; + } + blk_major[c+1] = key; + blk_vals[c+1] = v; + } + }; + // Phase 1: sample the num_major_sub requested major-axis vectors directly into the // output buffers, using the same helpers (and hence the same RNG stream) as the full // operator. On exit, the first "total" entries carry full major coordinates and local @@ -645,6 +663,9 @@ state_t fill_sparse_unpacked( work_state, vec_nnz, dim_major, num_major_sub, idxs_major, idxs_minor, vals ); total = vec_nnz * num_major_sub; + for (int64_t b = 0; b < num_major_sub; ++b) { + sort_block_by_major(idxs_major + b * vec_nnz, vals + b * vec_nnz, vec_nnz); + } } else { // LASO: each major-axis vector is sampled with replacement and merged in place, // advancing through the output buffers exactly as the full operator does. @@ -658,7 +679,9 @@ state_t fill_sparse_unpacked( for (int64_t i = 0; i < num_major_sub; ++i) { end_state = sample_indices_iid_uniform(dim_major, vec_nnz, im, v, end_state); laso_merge_long_axis_vector_coo_data(vec_nnz, v, im, in, i, loc2count, loc2scale); + // The merge compacts to the (<= vec_nnz) distinct survivors. int64_t count = (int64_t) loc2count.size(); + sort_block_by_major(im, v, count); im += count; in += count; v += count; total += count; } } @@ -849,7 +872,15 @@ COOMatrix submatrix_as_coo( A.rows = rows; A.cols = cols; A.nnz = nnz; - A.sort = RandBLAS::sparse_data::NonzeroSort::None; + // fill_sparse_unpacked emits each major-axis vector in sorted order, so the sampled + // submatrix is CSR- or CSC-sorted; label it as such. + if (D.dim_major == D.dim_minor) { + // degenerate case; recompute sort order with another scan over the data. + A.sort = RandBLAS::sparse_data::coo_arrays_determine_sort(A.nnz, A.rows, A.cols); + } else { + using RandBLAS::sparse_data::NonzeroSort; + A.sort = (D.dim_major == D.n_rows) ? NonzeroSort::CSC : NonzeroSort::CSR; + } return A; } diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 0d373186..a4d0a14c 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -102,5 +102,15 @@ target_link_libraries( spmm_performance PUBLIC RandBLAS blaspp lapackpp ) +add_executable( + sketch_general_performance simple-kernel-benchmarks/sketch_general_performance.cc +) +target_include_directories( + sketch_general_performance PUBLIC ${Random123_DIR} +) +target_link_libraries( + sketch_general_performance PUBLIC RandBLAS blaspp lapackpp +) + diff --git a/examples/simple-kernel-benchmarks/sketch_general_performance.cc b/examples/simple-kernel-benchmarks/sketch_general_performance.cc new file mode 100644 index 00000000..4aba663c --- /dev/null +++ b/examples/simple-kernel-benchmarks/sketch_general_performance.cc @@ -0,0 +1,654 @@ +// Copyright, 2026. See LICENSE for copyright holder information. +// +// ============================================================================ +// SKETCH_GENERAL PERFORMANCE BENCHMARK (sparse sketching operators) +// ============================================================================ +// +// Analogous to spmm_performance.cc, but it benchmarks the full sketching call +// sketch_general() with a real SparseSkOp -- not the low-level left_spmm on a +// generic uniform-random matrix. This exercises the *structured* operators that +// spmm_performance never produces: +// +// * Wide SASO SparseDist(d, m, k, Short) -- fixed k nnz per column +// * CountSketch SparseDist(d, m, 1, Short) -- exactly one nnz per column +// * Wide LASO SparseDist(d, m, k, Long) -- <= k nnz per row, empty cols +// +// and (via the right-sketch / tall operators) the transpose reduction that maps +// a tall SASO right-sketch onto the same CSC kernel path as a wide SASO. +// +// WHY WORK-NORMALIZED METRICS: raw wall-time conflates "how much work" with "how +// efficiently it's done". A vec_nnz=8 operator has 4x the nonzeros of vec_nnz=2, +// and a wide SASO has ~m/d x the nonzeros of a wide LASO at the same vec_nnz, so +// raw times are not comparable across those axes. We therefore report: +// +// * ns/elt = time / (nnz * work_mult): nanoseconds per nonzero-output-element. +// work_mult = n for a left-sketch (each S-nonzero touches a length-n output +// row), m for a right-sketch. This IS comparable across vec_nnz and SASO/LASO; +// for an appropriately optimized kernel it stays roughly flat as vec_nnz grows. +// * GB/s (%STR) = a MODEL byte count / time, as a fraction of measured STREAM +// bandwidth. These kernels are memory-bound (arithmetic intensity ~ vec_nnz/4 +// flop/byte), so % of STREAM is the roofline ceiling: ~80% means "done", ~20% +// means headroom. The byte count is an idealized-reuse model, not a hardware +// counter (perf counters are SIP-gated on macOS) -- read it as a bound. Two +// caveats: (a) A-read traffic is bounded by the DISTINCT rows/cols of A actually +// touched (<= min(nnz, contract_dim)), so a sparse LASO reads far less than all +// of A; (b) STREAM is a DRAM ceiling, so a small cache-resident problem (e.g. a +// tiny LASO) can legitimately exceed 100% %STR -- that flags "fits in cache", +// not a measurement error. +// +// Parallel efficiency needs the SAME kernel run at multiple thread counts, so it +// is NOT free per run; it lives behind --scaling (see USAGE). +// +// The cost of a sparse sketch has three phases; we time them separately: +// 1. SAMPLE fill_sparse(S): populate (rows, cols, vals). +// 2. CONVERT the COO->CSC sort inside apply_coo_via_csx (deepcopy + re-sort, +// incurred every apply when the operator's COO is not CSC-sorted). +// 3. KERNEL apply_csc_jki_p11 (ColMajor) / apply_csc_kib_1p1_rowmajor (RowMajor). +// WARM apply (S pre-sampled; phases 2+3) is the primary number; COLD (1+2+3) and +// CONVERT-only appear in the notes of the left/ColMajor table. +// +// NOTATION (left-sketch): B(d x n) = alpha * S(d x m) * A(m x n) +// (right-sketch): B(m x d) = alpha * A(m x n) * S(n x d) +// +// USAGE: +// ./sketch_general_performance [flags] # default sweep +// ./sketch_general_performance [flags] d m n vec_nnz major [trials] # single cfg +// major: 0 = Short (SASO), 1 = Long (LASO) +// flags: +// --no-stream skip the startup STREAM calibration (%STR prints '-') +// --scaling thread-scaling mode: speedup/efficiency/%STR per thread +// --threads=1,2,4,8 thread counts for --scaling (default 1,2,4,8) +// +// EXAMPLES: +// env OMP_NUM_THREADS=4 ./sketch_general_performance +// env OMP_NUM_THREADS=4 ./sketch_general_performance 200 2000 2000 4 0 +// ./sketch_general_performance --scaling --threads=1,2,4,8 200 2000 2000 4 0 +// +// ============================================================================ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "RandBLAS/config.h" +#include "RandBLAS/sparse_data/spmm_dispatch.hh" +#if defined(RandBLAS_HAS_OpenMP) +#include +#endif + +using namespace std::chrono; +using blas::Layout; +using blas::Op; +using RandBLAS::Axis; +using RandBLAS::SparseDist; +using RandBLAS::SparseSkOp; + +// Machine bandwidth ceiling (GB/s) used as the %STR denominator; <0 disables %STR. +static double g_stream_gbps = -1.0; + +// ---- OpenMP shims (no-ops without OpenMP) ---------------------------------- +static int current_threads() { +#if defined(RandBLAS_HAS_OpenMP) + return omp_get_max_threads(); +#else + return 1; +#endif +} +static void set_threads(int t) { +#if defined(RandBLAS_HAS_OpenMP) + omp_set_num_threads(t); +#else + (void)t; +#endif +} + +// Run num_trials repetitions, return {min, median} times in microseconds. +template +std::pair run_trials(Func&& func, int num_trials) { + std::vector times; + times.reserve(num_trials); + for (int t = 0; t < num_trials; ++t) { + auto start = steady_clock::now(); + func(); + auto end = steady_clock::now(); + times.push_back(duration_cast(end - start).count()); + } + std::sort(times.begin(), times.end()); + return {times[0], times[num_trials / 2]}; +} + +// STREAM triad (a = b + scalar*c) at the current OpenMP thread count. Returns +// achievable bandwidth in GB/s (min time over reps; 3*N*8 bytes per pass). +static double measure_stream_gbps(int64_t N = (int64_t)1 << 24, int reps = 10) { + std::vector a(N), b(N), c(N); + #pragma omp parallel for schedule(static) + for (int64_t i = 0; i < N; ++i) { a[i] = 0.0; b[i] = 1.0; c[i] = 2.0; } + const double scalar = 3.0; + // warmup + #pragma omp parallel for schedule(static) + for (int64_t i = 0; i < N; ++i) a[i] = b[i] + scalar * c[i]; + auto best = std::numeric_limits::max(); + for (int r = 0; r < reps; ++r) { + auto s = steady_clock::now(); + #pragma omp parallel for schedule(static) + for (int64_t i = 0; i < N; ++i) a[i] = b[i] + scalar * c[i]; + auto e = steady_clock::now(); + best = std::min(best, duration_cast(e - s).count()); + } + volatile double sink = a[N - 1]; (void)sink; // defeat dead-store elision + double bytes = 3.0 * (double)N * sizeof(double); + return best > 0 ? bytes / (double)best / 1000.0 : -1.0; // bytes/us/1000 == GB/s +} + +// ---- per-row metrics ------------------------------------------------------- +struct Row { + std::string label; + long min_us = 0, med_us = 0; + double ns_per_elt = -1; // time / (nnz * work_mult) + double gbps = -1; // model bytes / time + double pct_stream = -1; // gbps / g_stream_gbps + std::string notes; +}; + +// work_mult = output cols touched per nonzero (n for left-sketch, m for right). +// read_dense = elements of A read; out_elems = elements of the dense output. +static void fill_metrics(Row& r, long min_us, int64_t nnz, int64_t work_mult, + int64_t read_dense, int64_t out_elems) { + int64_t work = nnz * work_mult; + r.ns_per_elt = (work > 0) ? (double)min_us * 1000.0 / (double)work : -1; + double bytes = (double)(read_dense + 2 * out_elems) * sizeof(double) + + (double)nnz * (sizeof(double) + sizeof(int64_t)); + r.gbps = (min_us > 0) ? bytes / (double)min_us / 1000.0 : -1; + r.pct_stream = (g_stream_gbps > 0 && r.gbps > 0) ? r.gbps / g_stream_gbps * 100.0 : -1; +} + +static std::string fcell(double v, int prec) { + if (v < 0) return "-"; + std::ostringstream o; o << std::fixed << std::setprecision(prec) << v; return o.str(); +} + +static void print_metric_header(const std::string& title) { + std::cout << " " << title << "\n"; + std::cout << " " << std::left << std::setw(26) << "Operator" << std::right + << std::setw(9) << "Min(us)" << std::setw(9) << "Med(us)" + << std::setw(9) << "ns/elt" << std::setw(8) << "GB/s" + << std::setw(7) << "%STR" << " notes\n"; + std::cout << " " << std::string(76, '-') << "\n"; +} + +static void print_metric_row(const Row& r) { + std::cout << " " << std::left << std::setw(26) << r.label << std::right + << std::setw(9) << r.min_us << std::setw(9) << r.med_us + << std::setw(9) << fcell(r.ns_per_elt, 3) + << std::setw(8) << fcell(r.gbps, 1) + << std::setw(7) << fcell(r.pct_stream, 0) + << " " << r.notes << "\n"; +} + +static const char* sort_name(RandBLAS::sparse_data::NonzeroSort s) { + using NS = RandBLAS::sparse_data::NonzeroSort; + switch (s) { + case NS::CSC: return "CSC"; + case NS::CSR: return "CSR"; + default: return "None"; + } +} + +// One distribution config, benchmarked across both layouts on one side. +struct OpSpec { + std::string label; + int64_t vec_nnz; + Axis axis; +}; + +// --------------------------------------------------------------------------- +// LEFT-SKETCH: B(d x n) = S(d x m) * A(m x n), S ~ SparseDist(d, m, k, axis). +// work_mult = n; read_dense(A) = m*n; out_elems(B) = d*n. +// --------------------------------------------------------------------------- +void run_left(int64_t d, int64_t m, int64_t n, + const std::vector& specs, int num_trials) { + using T = double; + namespace rb = RandBLAS; + uint64_t seed = 12345; + + std::cout << "=== LEFT-SKETCH B(" << d << "x" << n << ") = S(" << d << "x" << m + << ") * A(" << m << "x" << n << "), trials=" << num_trials << " ===\n\n"; + + std::vector A_cm(m * n), A_rm(m * n); + rb::DenseDist DA(m, n); + auto st = rb::fill_dense(DA, A_cm.data(), rb::RNGState<>(seed)); + for (int64_t i = 0; i < m; ++i) + for (int64_t j = 0; j < n; ++j) + A_rm[i * n + j] = A_cm[i + j * m]; + + std::vector B_cm(d * n), B_rm(d * n), B_ref(d * n), result(d * n); + std::vector S_dense(d * m); + + std::vector rows_cm, rows_rm; + Row oracle; oracle.label = "densify(S)+GEMM"; oracle.notes = "dense oracle"; + + for (const auto& spec : specs) { + SparseDist dist(d, m, spec.vec_nnz, spec.axis); + SparseSkOp S(dist, st); + rb::fill_sparse(S); + auto view = rb::sparse::coo_view_of_skop(S); + int64_t nnz = S.nnz; + std::string base_note = std::string(sort_name(view.sort)) + " nnz=" + std::to_string(nnz); + + // Oracle: densify S then GEMM (ColMajor); also the correctness oracle. + rb::sparse_data::coo::coo_to_dense(view, Layout::ColMajor, S_dense.data()); + auto [t_gemm, m_gemm] = run_trials([&]() { + std::fill(B_ref.begin(), B_ref.end(), 0.0); + blas::gemm(Layout::ColMajor, Op::NoTrans, Op::NoTrans, d, n, m, + 1.0, S_dense.data(), d, A_cm.data(), m, 0.0, B_ref.data(), d); + }, num_trials); + oracle.min_us = t_gemm; oracle.med_us = m_gemm; + + auto [wcm, wcm_med] = run_trials([&]() { + std::fill(B_cm.begin(), B_cm.end(), 0.0); + rb::sketch_general(Layout::ColMajor, Op::NoTrans, Op::NoTrans, d, n, m, + 1.0, S, 0, 0, A_cm.data(), m, 0.0, B_cm.data(), d); + }, num_trials); + auto [wrm, wrm_med] = run_trials([&]() { + std::fill(B_rm.begin(), B_rm.end(), 0.0); + rb::sketch_general(Layout::RowMajor, Op::NoTrans, Op::NoTrans, d, n, m, + 1.0, S, 0, 0, A_rm.data(), n, 0.0, B_rm.data(), n); + }, num_trials); + + // COLD ColMajor: fresh unsampled operator each trial (sample+convert+kernel). + auto [ccm, ccm_med] = run_trials([&]() { + SparseSkOp Sc(dist, st); + std::fill(B_cm.begin(), B_cm.end(), 0.0); + rb::sketch_general(Layout::ColMajor, Op::NoTrans, Op::NoTrans, d, n, m, + 1.0, Sc, 0, 0, A_cm.data(), m, 0.0, B_cm.data(), d); + }, num_trials); + (void)ccm_med; + // CONVERT-only: worst-case COO->CSC re-sort (deepcopy + sort). + auto [cvt, cvt_med] = run_trials([&]() { + auto cp = view.deepcopy(); + cp.sort_arrays(rb::sparse_data::NonzeroSort::CSC); + }, num_trials); + (void)cvt_med; + + // Correctness vs the ColMajor oracle (both layouts). + std::fill(result.begin(), result.end(), 0.0); + rb::sketch_general(Layout::ColMajor, Op::NoTrans, Op::NoTrans, d, n, m, + 1.0, S, 0, 0, A_cm.data(), m, 0.0, result.data(), d); + double maxdiff = 0; + for (int64_t i = 0; i < d * n; ++i) + maxdiff = std::max(maxdiff, std::abs(result[i] - B_ref[i])); + if (maxdiff > 1e-9) + std::cout << " FAIL " << spec.label << " ColMajor max|diff|=" + << std::scientific << maxdiff << std::fixed << "\n"; + std::fill(result.begin(), result.end(), 0.0); + rb::sketch_general(Layout::RowMajor, Op::NoTrans, Op::NoTrans, d, n, m, + 1.0, S, 0, 0, A_rm.data(), n, 0.0, result.data(), n); + maxdiff = 0; + for (int64_t i = 0; i < d; ++i) + for (int64_t j = 0; j < n; ++j) + maxdiff = std::max(maxdiff, std::abs(result[i * n + j] - B_ref[i + j * d])); + if (maxdiff > 1e-9) + std::cout << " FAIL " << spec.label << " RowMajor max|diff|=" + << std::scientific << maxdiff << std::fixed << "\n"; + + // A-read traffic is bounded by the DISTINCT rows of A touched (<= min(nnz,m)), + // not all m rows: a wide LASO has mostly-empty columns. Using m here would + // overcount sparse operators (and push %STR past 100%). + int64_t a_read = std::min(nnz, m) * n; + Row rc; rc.label = spec.label; rc.min_us = wcm; rc.med_us = wcm_med; + fill_metrics(rc, wcm, nnz, n, a_read, d * n); + rc.notes = base_note + " cold=" + std::to_string(ccm) + " cvt=" + std::to_string(cvt); + rows_cm.push_back(rc); + + Row rr; rr.label = spec.label; rr.min_us = wrm; rr.med_us = wrm_med; + fill_metrics(rr, wrm, nnz, n, a_read, d * n); + rr.notes = base_note; + rows_rm.push_back(rr); + } + + print_metric_header("warm apply, ColMajor dense (jik/jki kernels)"); + for (auto& r : rows_cm) print_metric_row(r); + print_metric_row(oracle); + std::cout << "\n"; + print_metric_header("warm apply, RowMajor dense (ikb/kib kernels)"); + for (auto& r : rows_rm) print_metric_row(r); + std::cout << "\n"; +} + +// --------------------------------------------------------------------------- +// RIGHT-SKETCH: B(m x d) = A(m x n) * S(n x d), S ~ SparseDist(n, d, k, axis). +// work_mult = m; read_dense(A) = m*n; out_elems(B) = m*d. +// --------------------------------------------------------------------------- +void run_right(int64_t d, int64_t m, int64_t n, + const std::vector& specs, int num_trials) { + using T = double; + namespace rb = RandBLAS; + uint64_t seed = 67890; + + std::cout << "=== RIGHT-SKETCH B(" << m << "x" << d << ") = A(" << m << "x" << n + << ") * S(" << n << "x" << d << "), trials=" << num_trials << " ===\n\n"; + + std::vector A_cm(m * n), A_rm(m * n); + rb::DenseDist DA(m, n); + auto st = rb::fill_dense(DA, A_cm.data(), rb::RNGState<>(seed)); + for (int64_t i = 0; i < m; ++i) + for (int64_t j = 0; j < n; ++j) + A_rm[i * n + j] = A_cm[i + j * m]; + + std::vector B_cm(m * d), B_rm(m * d), B_ref(m * d), result(m * d); + std::vector S_dense(n * d); + + std::vector rows_cm, rows_rm; + + for (const auto& spec : specs) { + SparseDist dist(n, d, spec.vec_nnz, spec.axis); + SparseSkOp S(dist, st); + rb::fill_sparse(S); + auto view = rb::sparse::coo_view_of_skop(S); + int64_t nnz = S.nnz; + std::string base_note = std::string(sort_name(view.sort)) + " nnz=" + std::to_string(nnz); + + rb::sparse_data::coo::coo_to_dense(view, Layout::ColMajor, S_dense.data()); + std::fill(B_ref.begin(), B_ref.end(), 0.0); + blas::gemm(Layout::ColMajor, Op::NoTrans, Op::NoTrans, m, d, n, + 1.0, A_cm.data(), m, S_dense.data(), n, 0.0, B_ref.data(), m); + + auto [wcm, wcm_med] = run_trials([&]() { + std::fill(B_cm.begin(), B_cm.end(), 0.0); + rb::sketch_general(Layout::ColMajor, Op::NoTrans, Op::NoTrans, m, d, n, + 1.0, A_cm.data(), m, S, 0.0, B_cm.data(), m); + }, num_trials); + auto [wrm, wrm_med] = run_trials([&]() { + std::fill(B_rm.begin(), B_rm.end(), 0.0); + rb::sketch_general(Layout::RowMajor, Op::NoTrans, Op::NoTrans, m, d, n, + 1.0, A_rm.data(), n, S, 0.0, B_rm.data(), d); + }, num_trials); + + std::fill(result.begin(), result.end(), 0.0); + rb::sketch_general(Layout::ColMajor, Op::NoTrans, Op::NoTrans, m, d, n, + 1.0, A_cm.data(), m, S, 0.0, result.data(), m); + double maxdiff = 0; + for (int64_t i = 0; i < m * d; ++i) + maxdiff = std::max(maxdiff, std::abs(result[i] - B_ref[i])); + if (maxdiff > 1e-9) + std::cout << " FAIL " << spec.label << " ColMajor max|diff|=" + << std::scientific << maxdiff << std::fixed << "\n"; + + // A-read traffic bounded by distinct columns of A touched (<= min(nnz,n)). + int64_t a_read = std::min(nnz, n) * m; + Row rc; rc.label = spec.label; rc.min_us = wcm; rc.med_us = wcm_med; + fill_metrics(rc, wcm, nnz, m, a_read, m * d); rc.notes = base_note; + rows_cm.push_back(rc); + Row rr; rr.label = spec.label; rr.min_us = wrm; rr.med_us = wrm_med; + fill_metrics(rr, wrm, nnz, m, a_read, m * d); rr.notes = base_note; + rows_rm.push_back(rr); + } + + print_metric_header("warm apply, ColMajor dense"); + for (auto& r : rows_cm) print_metric_row(r); + std::cout << "\n"; + print_metric_header("warm apply, RowMajor dense"); + for (auto& r : rows_rm) print_metric_row(r); + std::cout << "\n"; +} + +// --------------------------------------------------------------------------- +// SCALING: re-run the left-sketch ColMajor warm apply at each thread count and +// report speedup, parallel efficiency, and %STREAM. STREAM is recalibrated per +// thread count so %STR is apples-to-apples. Note: a memory-bound kernel that has +// already saturated bandwidth SHOULD show efficiency < 1 -- read efficiency next +// to %STR, not in isolation. +// --------------------------------------------------------------------------- +void run_scaling(int64_t d, int64_t m, int64_t n, const std::vector& specs, + int num_trials, const std::vector& threads, bool do_stream) { + using T = double; + namespace rb = RandBLAS; + uint64_t seed = 12345; + + std::cout << "=== SCALING (left-sketch, ColMajor warm) d=" << d << " m=" << m + << " n=" << n << ", trials=" << num_trials << " ===\n"; +#if !defined(RandBLAS_HAS_OpenMP) + std::cout << " (built without OpenMP -- thread sweep is a no-op)\n"; +#endif + std::cout << "\n"; + + std::vector A_cm(m * n); + rb::DenseDist DA(m, n); + auto st = rb::fill_dense(DA, A_cm.data(), rb::RNGState<>(seed)); + std::vector B_cm(d * n); + + for (const auto& spec : specs) { + SparseDist dist(d, m, spec.vec_nnz, spec.axis); + SparseSkOp S(dist, st); + rb::fill_sparse(S); + int64_t nnz = S.nnz; + + std::cout << " " << spec.label << " (nnz=" << nnz << ")\n"; + std::cout << " " << std::right << std::setw(6) << "Thr" + << std::setw(10) << "Min(us)" << std::setw(10) << "Speedup" + << std::setw(8) << "Eff" << std::setw(9) << "GB/s" + << std::setw(7) << "%STR" << "\n"; + std::cout << " " << std::string(48, '-') << "\n"; + + long base_us = 0; + int base_t = threads.empty() ? 1 : threads.front(); + for (int t : threads) { + set_threads(t); + double stream = do_stream ? measure_stream_gbps() : -1.0; + auto [mn, md] = run_trials([&]() { + std::fill(B_cm.begin(), B_cm.end(), 0.0); + rb::sketch_general(Layout::ColMajor, Op::NoTrans, Op::NoTrans, d, n, m, + 1.0, S, 0, 0, A_cm.data(), m, 0.0, B_cm.data(), d); + }, num_trials); + (void)md; + if (t == base_t) base_us = mn; + double speedup = (mn > 0) ? (double)base_us / (double)mn : 0.0; + double eff = speedup / ((double)t / (double)base_t); + double bytes = (double)(std::min(nnz, m) * n + 2 * d * n) * sizeof(double) + + (double)nnz * (sizeof(double) + sizeof(int64_t)); + double gbps = (mn > 0) ? bytes / (double)mn / 1000.0 : -1; + double pct = (stream > 0 && gbps > 0) ? gbps / stream * 100.0 : -1; + std::cout << " " << std::right << std::setw(6) << t + << std::setw(10) << mn + << std::setw(10) << fcell(speedup, 2) + << std::setw(8) << fcell(eff, 2) + << std::setw(9) << fcell(gbps, 1) + << std::setw(7) << fcell(pct, 0) << "\n"; + } + std::cout << "\n"; + } +} + +// --------------------------------------------------------------------------- +// CSR-vs-CSC KERNEL PROBE (left-sketch). Builds the 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. The question: does the CSR +// jik path -- whose ColMajor outer loop runs over d rows (all non-empty for a +// wide LASO) -- recover the gap vs the CSC jki path, whose outer loop runs over +// m mostly-empty columns? CSC kib (RowMajor) is the already-fast baseline. +// +// Finding: CSR jik (ColMajor) beats CSC jki (ColMajor) 1.1-2x for BOTH SASO and +// LASO with no regression, motivating the ColMajor "prefer CSR for wide +// operators" routing now in apply_coo_via_csx (coo_spmm_impl.hh). (A strided-axpy +// "ColMajor kib" kernel was also tried and rejected -- it regressed dense-column +// SASO ~3x.) The probe times the kernels directly, so it is unaffected by that +// routing change and remains the A/B harness for it. +// --------------------------------------------------------------------------- +void run_csr_probe(int64_t d, int64_t m, int64_t n, + const std::vector& specs, int num_trials) { + using T = double; + namespace rb = RandBLAS; + uint64_t seed = 12345; + + std::cout << "=== CSR-vs-CSC KERNEL PROBE (left-sketch) d=" << d << " m=" << m + << " n=" << n << ", trials=" << num_trials << " ===\n"; + std::cout << " kernels timed directly via left_spmm on prebuilt CSC/CSR (no COO re-sort)\n\n"; + + std::vector A_cm(m * n), A_rm(m * n); + rb::DenseDist DA(m, n); + auto st = rb::fill_dense(DA, A_cm.data(), rb::RNGState<>(seed)); + for (int64_t i = 0; i < m; ++i) + for (int64_t j = 0; j < n; ++j) + A_rm[i * n + j] = A_cm[i + j * m]; + std::vector C_cm(d * n), C_rm(d * n), B_ref(d * n), S_dense(d * m), result(d * n); + + for (const auto& spec : specs) { + SparseDist dist(d, m, spec.vec_nnz, spec.axis); + SparseSkOp S(dist, st); + rb::fill_sparse(S); + auto view = rb::sparse::coo_view_of_skop(S); + int64_t nnz = S.nnz; + auto S_csc = view.as_owning_csc(); + auto S_csr = view.as_owning_csr(); + int64_t a_read = std::min(nnz, m) * n; + + // Oracle: densify(S) + GEMM (also the correctness reference). + rb::sparse_data::coo::coo_to_dense(view, Layout::ColMajor, S_dense.data()); + std::fill(B_ref.begin(), B_ref.end(), 0.0); + blas::gemm(Layout::ColMajor, Op::NoTrans, Op::NoTrans, d, n, m, + 1.0, S_dense.data(), d, A_cm.data(), m, 0.0, B_ref.data(), d); + + std::vector rows; + auto probe = [&](const std::string& label, const auto& Sp, Layout lay) { + const T* Bp = (lay == Layout::ColMajor) ? A_cm.data() : A_rm.data(); + int64_t ldb = (lay == Layout::ColMajor) ? m : n; + T* Cp = (lay == Layout::ColMajor) ? C_cm.data() : C_rm.data(); + int64_t ldc = (lay == Layout::ColMajor) ? d : n; + auto [mn, md] = run_trials([&]() { + std::fill(Cp, Cp + d * n, 0.0); + rb::sparse_data::left_spmm(lay, Op::NoTrans, Op::NoTrans, d, n, m, + 1.0, Sp, 0, 0, Bp, ldb, 0.0, Cp, ldc); + }, num_trials); + std::fill(result.begin(), result.end(), 0.0); + rb::sparse_data::left_spmm(lay, Op::NoTrans, Op::NoTrans, d, n, m, + 1.0, Sp, 0, 0, Bp, ldb, 0.0, result.data(), ldc); + double maxdiff = 0; + if (lay == Layout::ColMajor) + for (int64_t i = 0; i < d * n; ++i) + maxdiff = std::max(maxdiff, std::abs(result[i] - B_ref[i])); + else + for (int64_t i = 0; i < d; ++i) + for (int64_t j = 0; j < n; ++j) + maxdiff = std::max(maxdiff, std::abs(result[i * n + j] - B_ref[i + j * d])); + if (maxdiff > 1e-9) + std::cout << " FAIL " << label << " max|diff|=" << std::scientific + << maxdiff << std::fixed << "\n"; + Row r; r.label = label; r.min_us = mn; r.med_us = md; + fill_metrics(r, mn, nnz, n, a_read, d * n); + r.notes = std::string(sort_name(view.sort)) + " nnz=" + std::to_string(nnz); + rows.push_back(r); + }; + + print_metric_header(spec.label); + probe("CSC jki ColMajor", S_csc, Layout::ColMajor); + probe("CSR jik ColMajor", S_csr, Layout::ColMajor); + probe("CSC kib RowMajor", S_csc, Layout::RowMajor); + probe("CSR ikb RowMajor", S_csr, Layout::RowMajor); + for (auto& r : rows) print_metric_row(r); + std::cout << "\n"; + } +} + +std::vector sweep_specs() { + return { + {"Wide SASO k=2", 2, Axis::Short}, + {"Wide SASO k=4", 4, Axis::Short}, + {"Wide SASO k=8", 8, Axis::Short}, + {"CountSketch k=1", 1, Axis::Short}, + {"Wide LASO k=2", 2, Axis::Long}, + {"Wide LASO k=4", 4, Axis::Long}, + {"Wide LASO k=8", 8, Axis::Long}, + }; +} + +static std::vector parse_threads(const std::string& csv) { + std::vector out; + std::stringstream ss(csv); + std::string tok; + while (std::getline(ss, tok, ',')) { + if (!tok.empty()) out.push_back(std::atoi(tok.c_str())); + } + if (out.empty()) out = {1, 2, 4, 8}; + return out; +} + +int main(int argc, char** argv) { + bool no_stream = false, scaling = false, csr_probe = false; + std::vector threads = {1, 2, 4, 8}; + std::vector pos; + for (int i = 1; i < argc; ++i) { + std::string s = argv[i]; + if (s == "--no-stream") no_stream = true; + else if (s == "--scaling") scaling = true; + else if (s == "--csr-probe") csr_probe = true; + else if (s.rfind("--threads=", 0) == 0) threads = parse_threads(s.substr(10)); + else pos.push_back(s); + } + + std::cout << "\n============================================================\n"; + std::cout << "SKETCH_GENERAL PERFORMANCE BENCHMARK (sparse operators)\n"; + std::cout << "============================================================\n"; + std::cout << "Threads (OMP_NUM_THREADS/max): " << current_threads() << "\n"; + + bool have_cfg = pos.size() >= 5; + int64_t d = have_cfg ? std::atoll(pos[0].c_str()) : 0; + int64_t m = have_cfg ? std::atoll(pos[1].c_str()) : 0; + int64_t n = have_cfg ? std::atoll(pos[2].c_str()) : 0; + int64_t k = have_cfg ? std::atoll(pos[3].c_str()) : 0; + Axis axis = (have_cfg && std::atoi(pos[4].c_str()) != 0) ? Axis::Long : Axis::Short; + int trials = (pos.size() > 5) ? std::atoi(pos[5].c_str()) : (have_cfg ? 20 : 10); + + if (scaling) { + // STREAM is calibrated per thread count inside run_scaling. + std::vector specs = have_cfg + ? std::vector{{"single", k, axis}} : sweep_specs(); + int64_t sd = have_cfg ? d : 200, sm = have_cfg ? m : 2000, sn = have_cfg ? n : 2000; + std::cout << "\n"; + run_scaling(sd, sm, sn, specs, trials, threads, !no_stream); + return 0; + } + + if (!no_stream) { + std::cout << "Calibrating STREAM triad..." << std::flush; + g_stream_gbps = measure_stream_gbps(); + std::cout << " " << std::fixed << std::setprecision(1) << g_stream_gbps + << " GB/s (= 100% on the %STR column)\n"; + } else { + std::cout << "STREAM calibration skipped (--no-stream); %STR shows '-'\n"; + } + std::cout << "\n"; + + if (csr_probe) { + std::vector specs = have_cfg + ? std::vector{{"single", k, axis}} : sweep_specs(); + int64_t sd = have_cfg ? d : 200, sm = have_cfg ? m : 2000, sn = have_cfg ? n : 2000; + run_csr_probe(sd, sm, sn, specs, trials); + return 0; + } + + if (have_cfg) { + std::vector specs = {{"single", k, axis}}; + run_left(d, m, n, specs, trials); + run_right(d, m, n, specs, trials); + return 0; + } + + auto specs = sweep_specs(); + std::cout << "########## SECTION 1: vary sketch size d (m=n=2000) ##########\n\n"; + for (int64_t dd : {40, 200, 500}) { + run_left(dd, 2000, 2000, specs, trials); + run_right(dd, 2000, 2000, specs, trials); + } + std::cout << "########## SECTION 2: narrow n (SpMV-ish) ##########\n\n"; + run_left(200, 2000, 1, specs, trials); + run_left(200, 2000, 10, specs, trials); + return 0; +} diff --git a/examples/simple-kernel-benchmarks/spmm_performance.cc b/examples/simple-kernel-benchmarks/spmm_performance.cc index 84c74e92..ff3ab37e 100644 --- a/examples/simple-kernel-benchmarks/spmm_performance.cc +++ b/examples/simple-kernel-benchmarks/spmm_performance.cc @@ -12,7 +12,7 @@ // // 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, via apply_coo_via_csc) +// (COO delegates to the CSC kernels in both layouts, via apply_coo_via_csx) // // A fourth combination -- op(B)=Trans (computing C = A * B^T, the dense operand // fed transposed) -- also routes to apply_csr_jik_p11 / apply_csc_jki_p11, but diff --git a/test/datastructures/test_coo_matrix.cc b/test/datastructures/test_coo_matrix.cc index 4b83b6df..4f5f5e23 100644 --- a/test/datastructures/test_coo_matrix.cc +++ b/test/datastructures/test_coo_matrix.cc @@ -98,6 +98,30 @@ class TestCOO : public ::testing::Test { return; } + template + void test_dense_sums_duplicates() { + // coo_to_dense must ACCUMULATE duplicate (row, col) records (not overwrite). + COOMatrix A(3, 3); + A.reserve(4); + // Two records at (0, 1) that must sum: 1.5 + (-0.5) = 1.0. + A.rows[0] = 0; A.cols[0] = 1; A.vals[0] = (T) 1.5; + A.rows[1] = 0; A.cols[1] = 1; A.vals[1] = (T) -0.5; + // Three-way duplicate at (2, 2): 0.25 + 0.25 + 0.5 = 1.0 (only 2 here + 1 below). + A.rows[2] = 2; A.cols[2] = 2; A.vals[2] = (T) 0.25; + A.rows[3] = 2; A.cols[3] = 2; A.vals[3] = (T) 0.75; + + std::vector dense(9, (T) -7); // sentinel; must be overwritten to 0 then summed. + coo_to_dense(A, Layout::RowMajor, dense.data()); + EXPECT_EQ(dense[0*3 + 1], (T) 1.0); // (0,1): 1.5 - 0.5 + EXPECT_EQ(dense[2*3 + 2], (T) 1.0); // (2,2): 0.25 + 0.75 + // Every other cell is exactly zero (sentinel cleared, no records). + for (int64_t i = 0; i < 3; ++i) + for (int64_t j = 0; j < 3; ++j) + if (!((i==0 && j==1) || (i==2 && j==2))) + EXPECT_EQ((dense[i*3 + j]), (T) 0.0) << "cell (" << i << "," << j << ")"; + return; + } + template void test_sort_order(int64_t n) { COOMatrix A(n, n); @@ -251,6 +275,11 @@ TEST_F(TestCOO, to_from_dense) { test_to_from_dense(10); } +TEST_F(TestCOO, dense_sums_duplicates) { + test_dense_sums_duplicates(); + test_dense_sums_duplicates(); +} + TEST_F(TestCOO, sort_order) { test_sort_order(3); test_sort_order(7); diff --git a/test/datastructures/test_sparseskop.cc b/test/datastructures/test_sparseskop.cc index 5d4a87eb..c0c5d5a6 100644 --- a/test/datastructures/test_sparseskop.cc +++ b/test/datastructures/test_sparseskop.cc @@ -33,6 +33,8 @@ #include "RandBLAS/testing/comparison.hh" #include #include +#include +#include using std::vector; using RandBLAS::RNGState; @@ -104,9 +106,10 @@ class TestSparseSkOpConstruction : public ::testing::Test } template - void proper_laso_construction(int64_t d, int64_t m, int64_t key_index, int64_t nnz_index) { + void proper_laso_construction(int64_t d, int64_t m, int64_t key_index) { using RNG = SparseSkOp::state_t::generator; - SparseDist D0(d, m, vec_nnzs[nnz_index], Axis::Long); + int64_t vec_nnz = 1; + SparseDist D0(d, m, vec_nnz, Axis::Long); SparseSkOp S0(D0, keys[key_index]); fill_sparse(S0); if (d < m) { @@ -297,6 +300,7 @@ class TestSparseSkOpConstruction : public ::testing::Test } return; } + }; TEST_F(TestSparseSkOpConstruction, respect_ownership) { @@ -421,71 +425,29 @@ TEST_F(TestSparseSkOpConstruction, SASO_Dim_15by7_int32) { TEST_F(TestSparseSkOpConstruction, LASO_Dim_7by20) { // vec_nnz=1 - proper_laso_construction(7, 20, 0, 0); - proper_laso_construction(7, 20, 1, 0); - proper_laso_construction(7, 20, 2, 0); - // // vec_nnz=2 - // proper_laso_construction(7, 20, 0, 1); - // proper_laso_construction(7, 20, 1, 1); - // proper_laso_construction(7, 20, 2, 1); - // // vec_nnz=3 - // proper_laso_construction(7, 20, 0, 2); - // proper_laso_construction(7, 20, 1, 2); - // proper_laso_construction(7, 20, 2, 2); - // // vec_nnz=7 - // proper_laso_construction(7, 20, 0, 3); - // proper_laso_construction(7, 20, 1, 3); - // proper_laso_construction(7, 20, 2, 3); + proper_laso_construction(7, 20, 0); + proper_laso_construction(7, 20, 1); + proper_laso_construction(7, 20, 2); } TEST_F(TestSparseSkOpConstruction, LASO_Dim_15by7) { // vec_nnz=1 - proper_laso_construction(15, 7, 0, 0); - proper_laso_construction(15, 7, 1, 0); - // // vec_nnz=2 - // proper_laso_construction(15, 7, 0, 1); - // proper_laso_construction(15, 7, 1, 1); - // // vec_nnz=3 - // proper_laso_construction(15, 7, 0, 2); - // proper_laso_construction(15, 7, 1, 2); - // // vec_nnz=7 - // proper_laso_construction(15, 7, 0, 3); - // proper_laso_construction(15, 7, 1, 3); + proper_laso_construction(15, 7, 0); + proper_laso_construction(15, 7, 1); } TEST_F(TestSparseSkOpConstruction, LASO_Dim_7by20_int32) { // vec_nnz=1 - proper_laso_construction(7, 20, 0, 0); - proper_laso_construction(7, 20, 1, 0); - proper_laso_construction(7, 20, 2, 0); - // // vec_nnz=2 - // proper_laso_construction(7, 20, 0, 1); - // proper_laso_construction(7, 20, 1, 1); - // proper_laso_construction(7, 20, 2, 1); - // // vec_nnz=3 - // proper_laso_construction(7, 20, 0, 2); - // proper_laso_construction(7, 20, 1, 2); - // proper_laso_construction(7, 20, 2, 2); - // // vec_nnz=7 - // proper_laso_construction(7, 20, 0, 3); - // proper_laso_construction(7, 20, 1, 3); - // proper_laso_construction(7, 20, 2, 3); + proper_laso_construction(7, 20, 0); + proper_laso_construction(7, 20, 1); + proper_laso_construction(7, 20, 2); } TEST_F(TestSparseSkOpConstruction, LASO_Dim_15by7_int32) { // vec_nnz=1 - proper_laso_construction(15, 7, 0, 0); - proper_laso_construction(15, 7, 1, 0); - // // vec_nnz=2 - // proper_laso_construction(15, 7, 0, 1); - // proper_laso_construction(15, 7, 1, 1); - // // vec_nnz=3 - // proper_laso_construction(15, 7, 0, 2); - // proper_laso_construction(15, 7, 1, 2); - // // vec_nnz=7 - // proper_laso_construction(15, 7, 0, 3); - // proper_laso_construction(15, 7, 1, 3); + proper_laso_construction(15, 7, 0); + proper_laso_construction(15, 7, 1); } diff --git a/test/linops/test_spmm/DevNotes.md b/test/linops/test_spmm/DevNotes.md index 5c675463..e46e888a 100644 --- a/test/linops/test_spmm/DevNotes.md +++ b/test/linops/test_spmm/DevNotes.md @@ -45,8 +45,7 @@ The effective layout for reading the dense operand B (`layout_opB`) is determine From there, our next steps are format-dependent. COO format ([spmm_dispatch.hh:124-126](../../../RandBLAS/sparse_data/spmm_dispatch.hh#L124-L126)): -- Always uses `apply_coo_via_csc` (converts to CSC internally) -- 1 kernel handles all 4 (layout × opB) combinations +- Always uses `apply_coo_via_csx` (converts to CSR or CSC internally, at its own discretion). CSC format ([spmm_dispatch.hh:128-134](../../../RandBLAS/sparse_data/spmm_dispatch.hh#L128-L134)): - If `layout_opB == RowMajor && layout_C == RowMajor`: `apply_csc_kib_1p1_rowmajor` @@ -114,7 +113,7 @@ For completeness, here's a grouping of tests based on the kernel they hit. | Kernel | Direct Tests | Via Transformation | |--------|--------------|-------------------| -| `apply_coo_via_csc` | All COO tests | CSR/CSC transpose_self | +| `apply_coo_via_csx` | All COO tests | CSR/CSC transpose_self | | `apply_csc_jki_p11` | CSC ColMajor (both opB), CSC RowMajor + opB=Trans | CSR transpose_self | | `apply_csc_kib_1p1_rowmajor` | CSC RowMajor + opB=NoTrans | CSR transpose_self + RowMajor | | `apply_csr_jik_p11` | CSR ColMajor (both opB), CSR RowMajor + opB=Trans | CSC transpose_self | diff --git a/test/linops/test_spmm/test_spmm_csr.cc b/test/linops/test_spmm/test_spmm_csr.cc index 2071818c..18f11fdf 100644 --- a/test/linops/test_spmm/test_spmm_csr.cc +++ b/test/linops/test_spmm/test_spmm_csr.cc @@ -469,3 +469,70 @@ TEST_F(TestPublicAPI_DenseTimesSparse_double, nontrivial_alpha_beta_rowmajor) { test_dense_times_sparse(50, 30, 40, 2.5, -1.0, Layout::RowMajor, 0.80); } + +//////////////////////////////////////////////////////////////////////// +// +// +// CSR SpMM reference checks (duplicate-colidx accumulation, variable nnz/row) +// +// +//////////////////////////////////////////////////////////////////////// + +// Compare C = alpha * A @ B against a hand-computed dense reference via the public dispatch +// apply_csr_jik_p11, exercising in particular the accumulation of a column index that is +// duplicated within a row. +static void check_csr_jik_against_reference( + int64_t d, int64_t m, int64_t n, + const std::vector& rowptr, + const std::vector& colidxs, + const std::vector& vals +) { + CSRMatrix A(d, m); + A.reserve((int64_t) vals.size()); + for (int64_t i = 0; i <= d; ++i) A.rowptr[i] = rowptr[i]; + for (size_t e = 0; e < vals.size(); ++e) { A.colidxs[e] = colidxs[e]; A.vals[e] = vals[e]; } + + std::vector B(m * n), C(d * n, 0.0), Cref(d * n, 0.0); + for (int64_t k = 0; k < m; ++k) + for (int64_t j = 0; j < n; ++j) + B[k + j*m] = (double)(k + 1) - 0.3 * (double) j; + + double alpha = 2.0; + for (int64_t i = 0; i < d; ++i) + for (int64_t e = rowptr[i]; e < rowptr[i+1]; ++e) + for (int64_t j = 0; j < n; ++j) + Cref[i + j*d] += alpha * vals[e] * B[colidxs[e] + j*m]; + + auto compare = [&](const char* which) { + auto msg = RandBLAS::testing::buffs_approx_equal( + C.data(), Cref.data(), d * n, which, __FILE__, __LINE__ + ); + if (msg.size() > 0) FAIL() << msg; + }; + + std::fill(C.begin(), C.end(), 0.0); + apply_csr_jik_p11(alpha, Layout::ColMajor, Layout::ColMajor, + d, n, m, A, B.data(), m, C.data(), d); + compare("apply_csr_jik_p11"); +} + +TEST(TestCSRSpMMReference, duplicate_colidx_accumulates) { + // d=4 rows, each with exactly 2 nonzeros. Row 1 has a duplicated column index (2, 2): + // the two records (0.5, 0.5) must accumulate to 1.0 at column 2. + check_csr_jik_against_reference( + 4, 5, 3, + {0, 2, 4, 6, 8}, + {0, 3, 2, 2, 1, 4, 0, 2}, + {1.0, -2.0, 0.5, 0.5, 3.0, -1.0, 2.0, 1.0} + ); +} + +TEST(TestCSRSpMMReference, variable_nnz_per_row) { + // Rows have 1, 2, 1, 2 nonzeros. The same hand-computed dense reference must hold. + check_csr_jik_against_reference( + 4, 5, 3, + {0, 1, 3, 4, 6}, + {2, 0, 4, 3, 1, 2}, + {1.5, -0.5, 2.0, 4.0, -1.0, 0.25} + ); +}