From fadb8e997a0fbd9ddb0b2fda27d3f6d84f196b3b Mon Sep 17 00:00:00 2001 From: Riley Murray Date: Sat, 30 May 2026 18:32:34 -0700 Subject: [PATCH 1/3] initial implementation of sparse sketching operator lazy submatrix generation --- RandBLAS/skge.hh | 22 ++- RandBLAS/sparse_skops.hh | 229 +++++++++++++++++++------ test/datastructures/test_sparseskop.cc | 121 +++++++++++++ test/linops/test_lskges.cc | 72 ++++++++ test/linops/test_rskges.cc | 74 +++++++- 5 files changed, 460 insertions(+), 58 deletions(-) diff --git a/RandBLAS/skge.hh b/RandBLAS/skge.hh index bbf910ab..9bbd4e37 100644 --- a/RandBLAS/skge.hh +++ b/RandBLAS/skge.hh @@ -481,9 +481,14 @@ void lskges( int64_t ldb ) { if (S.nnz < 0) { - SparseSkOp shallowcopy(S.dist, S.seed_state); // shallowcopy.own_memory = true. - fill_sparse(shallowcopy); - lskges(layout, opS, opA, d, n, m, alpha, shallowcopy, ro_s, co_s, A, lda, beta, B, ldb); + // 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); return; } auto Scoo = coo_view_of_skop(S); @@ -614,9 +619,14 @@ inline void rskges( int64_t ldb ) { if (S.nnz < 0) { - SparseSkOp shallowcopy(S.dist, S.seed_state); // shallowcopy.own_memory = true. - fill_sparse(shallowcopy); - rskges(layout, opA, opS, m, d, n, alpha, A, lda, shallowcopy, ro_s, co_s, beta, B, ldb); + // 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); return; } auto Scoo = coo_view_of_skop(S); diff --git a/RandBLAS/sparse_skops.hh b/RandBLAS/sparse_skops.hh index 7a1afdc8..c0b72e3c 100644 --- a/RandBLAS/sparse_skops.hh +++ b/RandBLAS/sparse_skops.hh @@ -94,7 +94,6 @@ void _considerate_fisher_yates( return; } - template > static state_t repeated_fisher_yates( const state_t &state, @@ -155,7 +154,7 @@ inline double isometry_scale(Axis major_axis, int64_t vec_nnz, int64_t dim_major namespace RandBLAS { // Forward declaration of SparseSkOp. It's returnable by -// SparseDist.sample(), but its definition involves DenseDist. +// SparseDist.sample(), but its definition involves SparseDist. template struct SparseSkOp; @@ -287,7 +286,7 @@ struct SparseDist { /// This function is used for sampling a sequence of \math{k} elements uniformly /// without replacement from the index set \math{\\{0,\ldots,n-1\\}.} It uses a special /// implementation of Fisher-Yates shuffling to produce \math{r} such samples in \math{O(n + rk)} time. -/// These samples are stored by writing to \math{\ttt{samples}} in \math{r} blocks of length \math{k.} +/// These samples are stored by writing to \math{\ttt{samples}} in \math{r} blocks of length \math{k.} /// /// The returned RNGState should /// be used for the next call to a random sampling function whose output should be statistically @@ -529,69 +528,147 @@ void laso_merge_long_axis_vector_coo_data( /// .. |Dfullnnz| mathmacro:: {\mathcal{D}\mathtt{.full\_nnz}} /// /// @endverbatim -/// This function is the underlying implementation of fill_sparse(SparseSkOp &S). -/// It has no allocation stage and it skips checks for null pointers. +/// Samples ONLY the \math{\ttt{n_rows_sub} \times \ttt{n_cols_sub}} submatrix of the +/// operator defined by \math{(\D,\ttt{seed_state})} whose upper-left corner sits at +/// \math{(\ttt{ro_s},\ttt{co_s})}, without ever materializing the full operator. This +/// function has no allocation stage and it skips checks for null pointers. /// -/// On entry, \math{(\vals,\rows,\cols)} are arrays of length at least \math{\Dfullnnz.} -/// On exit, the first \math{\ttt{nnz}} entries of these arrays contain the data for -/// a COO sparse matrix representation of the SparseSkOp -/// defined by \math{(\D,\ttt{seed_state)}.} +/// On entry, \math{(\vals,\rows,\cols)} are arrays of length at least +/// \math{\ttt{D.vec_nnz} \cdot \ttt{num_major_sub}}, where \math{\ttt{num_major_sub}} is +/// the number of major-axis vectors that intersect the requested submatrix (this is the +/// worst-case nonzero count for the submatrix). On exit, the first \math{\ttt{nnz}} +/// entries of these arrays contain the data for a COO sparse matrix representation of +/// \math{\submat(\mtxS)}, with row/column indices already shifted into +/// \math{[0,\ttt{n_rows_sub}) \times [0,\ttt{n_cols_sub})}. +/// +/// The data produced by this function equals, entry-for-entry as a sparse matrix, the +/// \math{[\ttt{ro_s}:\ttt{ro_s}+\ttt{n_rows_sub},\,\ttt{co_s}:\ttt{co_s}+\ttt{n_cols_sub}]} +/// submatrix of \math{\ttt{fill_sparse_unpacked_nosub}(\D,\ldots)}. /// -/// Note: the "nosub" suffix in this function name reflects how it has no ability -/// to sample submatrices of sparse sketching operators. A future release of -/// RandBLAS will add a function called "fill_sparse_unpacked()" with capabilities -/// that are completely analogous to fill_dense_unpacked(). -/// template -state_t fill_sparse_unpacked_nosub( +state_t fill_sparse_unpacked( const SparseDist &D, - int64_t &nnz, T* vals, sint_t* rows, sint_t *cols, + int64_t n_rows_sub, int64_t n_cols_sub, + int64_t ro_s, int64_t co_s, + int64_t &nnz, T* vals, sint_t* rows, sint_t* cols, const state_t &seed_state ) { + randblas_require(D.n_rows >= n_rows_sub + ro_s); + randblas_require(D.n_cols >= n_cols_sub + co_s); + + // An operator sampled from D is built by drawing D.dim_minor major-axis vectors, + // each a length-(D.dim_major) sparse vector with vec_nnz nonzeros. Below we call the + // length of a major-axis vector "dim_major" and (in the submatrix variables) the + // count of such vectors "num_major" -- so the full operator's num_major == D.dim_minor. int64_t dim_major = D.dim_major; - int64_t dim_minor = D.dim_minor; + int64_t vec_nnz = D.vec_nnz; + + // Map the submatrix's (row, col) offsets/extents onto the (short, long) axes, + // exactly as fill_sparse_unpacked_nosub chooses idxs_short/idxs_long. + bool short_is_rows = (D.n_rows <= D.n_cols); + int64_t short_off, short_sub, long_off, long_sub; + if (short_is_rows) { + short_off = ro_s; short_sub = n_rows_sub; long_off = co_s; long_sub = n_cols_sub; + } else { + short_off = co_s; short_sub = n_cols_sub; long_off = ro_s; long_sub = n_rows_sub; + } - sint_t *idxs_short = (D.n_rows <= D.n_cols) ? rows : cols; - sint_t *idxs_long = (D.n_rows <= D.n_cols) ? cols : rows; - int64_t vec_nnz = D.vec_nnz; + // The major axis is the short axis for SASO and the long axis for LASO; the count + // of major-axis vectors runs along the other axis. The submatrix keeps a window of + // dim_major_sub coordinates (starting at dim_major_off) WITHIN each kept vector, and + // keeps num_major_sub of the vectors (starting at vector index num_major_off). + int64_t dim_major_off, dim_major_sub, num_major_off, num_major_sub; + bool major_is_rows; + if (D.major_axis == Axis::Short) { + dim_major_off = short_off; dim_major_sub = short_sub; + num_major_off = long_off; num_major_sub = long_sub; + major_is_rows = short_is_rows; + } else { + dim_major_off = long_off; dim_major_sub = long_sub; + num_major_off = short_off; num_major_sub = short_sub; + major_is_rows = !short_is_rows; + } + // Skip the RNG counter past the num_major_off major-axis vectors we don't need. + // Both the Fisher-Yates path (vec_nnz > 1) and the i.i.d.-uniform path (vec_nnz == 1 + // and LASO) consume exactly vec_nnz counter increments per major-axis vector, so the + // skip amount is uniform. + state_t work_state = seed_state; + work_state.counter.incr(num_major_off * vec_nnz); + + // Identify which output array holds the major-axis coordinate and which holds the + // minor-axis coordinate (the index of the major-axis vector). We sample directly + // into these buffers (no scratch space) and then compact in place, so they must each + // have capacity >= vec_nnz * num_major_sub. + sint_t* idxs_major = major_is_rows ? rows : cols; + sint_t* idxs_minor = major_is_rows ? cols : rows; + + // 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 + // minor coordinates (0..num_major_sub-1); "total" is the pre-filter nnz. + int64_t total; + state_t end_state; if (D.major_axis == Axis::Short) { - auto state = sparse::repeated_fisher_yates( - seed_state, vec_nnz, dim_major, dim_minor, idxs_short, idxs_long, vals + end_state = sparse::repeated_fisher_yates( + work_state, vec_nnz, dim_major, num_major_sub, idxs_major, idxs_minor, vals ); - nnz = vec_nnz * dim_minor; - return state; + total = vec_nnz * num_major_sub; } else { - // We're long-axis major. - - // We don't sample all at once since we might need to merge duplicate entries - // in each long-axis vector. The way we do this is different than the - // standard COOMatrix convention of just adding entries together. - - // We begin by defining some datastructures that we repeatedly pass to a helper function. - // See the comments in the helper function for info on what these guys mean. + // LASO: each major-axis vector is sampled with replacement and merged in place, + // advancing through the output buffers exactly as the full operator does. std::unordered_map loc2count{}; - std::unordered_map loc2scale{}; - int64_t total_nnz = 0; - auto state = seed_state; - for (int64_t i = 0; i < dim_minor; ++i) { - state = sample_indices_iid_uniform(dim_major, vec_nnz, idxs_long, vals, state); - // ^ That writes directly so S.vals and either S.rows or S.cols. - // The new values might need to be changed if there are duplicates in idxs_long. - // We have a helper function for this since it's a tedious process. - // The helper function also sets whichever of S.rows or S.cols wasn't populated. - laso_merge_long_axis_vector_coo_data( - vec_nnz, vals, idxs_long, idxs_short, i, loc2count, loc2scale - ); - int64_t count = loc2count.size(); - vals += count; - idxs_long += count; - idxs_short += count; - total_nnz += count; + std::unordered_map loc2scale{}; + sint_t* im = idxs_major; + sint_t* in = idxs_minor; + T* v = vals; + total = 0; + end_state = work_state; + 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); + int64_t count = (int64_t) loc2count.size(); + im += count; in += count; v += count; total += count; } - nnz = total_nnz; - return state; } + + // Phase 2: compact in place, keeping only nonzeros whose major coordinate lands in + // the window [dim_major_off, dim_major_off + dim_major_sub) and shifting those + // coordinates down to local indices. The write index nnz never exceeds the read + // index k, so reading and writing the same buffers is safe. + nnz = 0; + for (int64_t k = 0; k < total; ++k) { + sint_t mc = idxs_major[k] - (sint_t) dim_major_off; + if (0 <= mc && mc < (sint_t) dim_major_sub) { + idxs_major[nnz] = mc; + idxs_minor[nnz] = idxs_minor[k]; + vals[nnz] = vals[k]; + nnz++; + } + } + return end_state; +} + +// ============================================================================= +/// This function is the underlying implementation of fill_sparse(SparseSkOp &S). +/// It has no allocation stage and it skips checks for null pointers. +/// +/// On entry, \math{(\vals,\rows,\cols)} are arrays of length at least \math{\Dfullnnz.} +/// On exit, the first \math{\ttt{nnz}} entries of these arrays contain the data for +/// a COO sparse matrix representation of the SparseSkOp +/// defined by \math{(\D,\ttt{seed_state)}.} +/// +/// The function fill_sparse_unpacked lets you perform a similar operation that +/// only returns data for a submatrix of a SparseSkOp. +template +state_t fill_sparse_unpacked_nosub( + const SparseDist &D, + int64_t &nnz, T* vals, sint_t* rows, sint_t *cols, + const state_t &seed_state +) { + return fill_sparse_unpacked( + D, D.n_rows, D.n_cols, 0, 0, nnz, vals, rows, cols, seed_state + ); } @@ -687,7 +764,9 @@ void print_sparse(SparseSkOp const &S0) { namespace RandBLAS::sparse { using RandBLAS::SparseSkOp; +using RandBLAS::SparseDist; using RandBLAS::Axis; +using RandBLAS::fill_sparse_unpacked; using RandBLAS::sparse_data::COOMatrix; template @@ -697,5 +776,53 @@ COOMatrix coo_view_of_skop(const SparseSkOp &S) { return A; } +// ============================================================================= +/// Allocate and return a memory-owning COOMatrix holding ONLY the +/// \math{\ttt{n_rows_sub} \times \ttt{n_cols_sub}} submatrix of \math{\mtxS} whose +/// upper-left corner sits at \math{(\ttt{ro_s},\ttt{co_s})}, without ever materializing +/// the full operator. This is the sparse analog of submatrix_as_blackbox(). +/// +/// The returned COOMatrix has dimensions \math{(\ttt{n_rows_sub},\ttt{n_cols_sub})}, +/// zero-based indexing, and indices that have already been shifted to local coordinates. +/// Because its dimensions exactly match the submatrix, passing it to left_spmm/right_spmm +/// with offsets (0,0) takes the no-extract fast path. The returned object owns its memory +/// (own_memory == true); its destructor frees the buffers. +/// +template +COOMatrix submatrix_as_coo( + const SparseSkOp &S, int64_t n_rows_sub, int64_t n_cols_sub, int64_t ro_s, int64_t co_s +) { + randblas_require(ro_s + n_rows_sub <= S.n_rows); + randblas_require(co_s + n_cols_sub <= S.n_cols); + const SparseDist &D = S.dist; + + // Worst-case nnz of the submatrix = vec_nnz * (number of major-axis vectors that + // intersect the submatrix). num_major_sub is that count; derive it via the same axis + // mapping used in fill_sparse_unpacked. + bool short_is_rows = (D.n_rows <= D.n_cols); + int64_t short_sub = short_is_rows ? n_rows_sub : n_cols_sub; + int64_t long_sub = short_is_rows ? n_cols_sub : n_rows_sub; + int64_t num_major_sub = (D.major_axis == Axis::Short) ? long_sub : short_sub; + int64_t cap = D.vec_nnz * num_major_sub; + + // Allocate the worst-case buffers, sample only the requested submatrix, and attach + // the buffers to an owning COOMatrix. We use the standard ctor + manual attach + // (rather than reserve()) because the submatrix may be empty (cap or actual nnz == 0) + // and reserve() rejects arg_nnz <= 0. + T* vals = new T[cap]; + sint_t* rows = new sint_t[cap]; + sint_t* cols = new sint_t[cap]; + int64_t nnz = 0; + fill_sparse_unpacked(D, n_rows_sub, n_cols_sub, ro_s, co_s, nnz, vals, rows, cols, S.seed_state); + + COOMatrix A(n_rows_sub, n_cols_sub); // own_memory == true, null arrays. + A.vals = vals; + A.rows = rows; + A.cols = cols; + A.nnz = nnz; + A.sort = RandBLAS::sparse_data::NonzeroSort::None; + return A; +} + } // end namespace RandBLAS::sparse diff --git a/test/datastructures/test_sparseskop.cc b/test/datastructures/test_sparseskop.cc index 6d5f474f..5d4a87eb 100644 --- a/test/datastructures/test_sparseskop.cc +++ b/test/datastructures/test_sparseskop.cc @@ -206,6 +206,97 @@ class TestSparseSkOpConstruction : public ::testing::Test } return; } + + // Equivalence oracle for submatrix sampling: for a given (ro_s, co_s, sub dims), + // fill_sparse_unpacked(...) must equal -- entry-for-entry as a sparse matrix -- the + // corresponding submatrix of the FULL operator from fill_sparse_unpacked_nosub(...). + // We compare by densifying both into dense n_rows_sub-by-n_cols_sub buffers, which is + // order-independent (important since SASO/LASO entry order may differ). + template + void fill_unpacked_sub_one( + Axis major_axis, int64_t vec_nnz, int64_t d, int64_t m, + int64_t ro_s, int64_t co_s, int64_t n_rows_sub, int64_t n_cols_sub, uint32_t key + ) { + SparseDist D {d, m, vec_nnz, major_axis}; + RNGState seed(key); + + // Full operator via the no-submatrix path. + int64_t full_nnz_out = -1; + vector full_vals(D.full_nnz); + vector full_rows(D.full_nnz); + vector full_cols(D.full_nnz); + fill_sparse_unpacked_nosub( + D, full_nnz_out, full_vals.data(), full_rows.data(), full_cols.data(), seed + ); + + // Submatrix via the new path. + // Worst-case nnz = vec_nnz * (#major-axis vectors intersecting the submatrix). + bool short_is_rows = (d <= m); + int64_t long_sub = short_is_rows ? n_cols_sub : n_rows_sub; + int64_t short_sub = short_is_rows ? n_rows_sub : n_cols_sub; + int64_t minor_sub = (major_axis == Axis::Short) ? long_sub : short_sub; + int64_t cap = vec_nnz * minor_sub; + int64_t sub_nnz_out = -1; + vector sub_vals(cap > 0 ? cap : 1); + vector sub_rows(cap > 0 ? cap : 1); + vector sub_cols(cap > 0 ? cap : 1); + RandBLAS::fill_sparse_unpacked( + D, n_rows_sub, n_cols_sub, ro_s, co_s, + sub_nnz_out, sub_vals.data(), sub_rows.data(), sub_cols.data(), seed + ); + ASSERT_GE(sub_nnz_out, 0); + ASSERT_LE(sub_nnz_out, cap); + + // Densify the oracle (restrict-after-sampling). + vector dense_oracle(n_rows_sub * n_cols_sub, (T) 0); + for (int64_t i = 0; i < full_nnz_out; ++i) { + int64_t r = full_rows[i]; + int64_t c = full_cols[i]; + if (ro_s <= r && r < ro_s + n_rows_sub && co_s <= c && c < co_s + n_cols_sub) { + dense_oracle[(r - ro_s) * n_cols_sub + (c - co_s)] += full_vals[i]; + } + } + // Densify the submatrix (restrict-while-sampling). + vector dense_sub(n_rows_sub * n_cols_sub, (T) 0); + for (int64_t i = 0; i < sub_nnz_out; ++i) { + int64_t r = sub_rows[i]; + int64_t c = sub_cols[i]; + ASSERT_GE(r, 0); ASSERT_LT(r, n_rows_sub); + ASSERT_GE(c, 0); ASSERT_LT(c, n_cols_sub); + dense_sub[r * n_cols_sub + c] += sub_vals[i]; + } + std::string msg = RandBLAS::testing::buffs_approx_equal( + dense_sub.data(), dense_oracle.data(), n_rows_sub * n_cols_sub, + __PRETTY_FUNCTION__, __FILE__, __LINE__ + ); + if (msg.size() > 0) { + FAIL() << msg; + } + return; + } + + template + void fill_unpacked_sub(Axis major_axis, int64_t vec_nnz, int64_t d, int64_t m) { + for (auto key : keys) { + // full-size (regression guard: must reproduce nosub exactly) + fill_unpacked_sub_one(major_axis, vec_nnz, d, m, 0, 0, d, m, key); + // corner offsets / partial windows + fill_unpacked_sub_one(major_axis, vec_nnz, d, m, 1, 0, d - 1, m, key); + fill_unpacked_sub_one(major_axis, vec_nnz, d, m, 0, 1, d, m - 1, key); + fill_unpacked_sub_one(major_axis, vec_nnz, d, m, 1, 1, d - 1, m - 1, key); + fill_unpacked_sub_one(major_axis, vec_nnz, d, m, 2, 3, d - 3, m - 4, key); + // single row / single column + fill_unpacked_sub_one(major_axis, vec_nnz, d, m, 0, 0, 1, m, key); + fill_unpacked_sub_one(major_axis, vec_nnz, d, m, 0, 0, d, 1, key); + fill_unpacked_sub_one(major_axis, vec_nnz, d, m, d - 1, 0, 1, m, key); + fill_unpacked_sub_one(major_axis, vec_nnz, d, m, 0, m - 1, d, 1, key); + // 1x1 windows at a few positions (some likely empty for SASO) + fill_unpacked_sub_one(major_axis, vec_nnz, d, m, 0, 0, 1, 1, key); + fill_unpacked_sub_one(major_axis, vec_nnz, d, m, d - 1, m - 1, 1, 1, key); + fill_unpacked_sub_one(major_axis, vec_nnz, d, m, d / 2, m / 2, 1, 1, key); + } + return; + } }; TEST_F(TestSparseSkOpConstruction, respect_ownership) { @@ -228,6 +319,36 @@ TEST_F(TestSparseSkOpConstruction, fill_unpacked_nosub_laso) { unpacked_nosub({20,10,7,Axis::Long}); } +TEST_F(TestSparseSkOpConstruction, fill_unpacked_sub_saso) { + for (int64_t vnz : {(int64_t) 1, (int64_t) 2, (int64_t) 4, (int64_t) 7}) { + fill_unpacked_sub(Axis::Short, vnz, 7, 13); // wide + fill_unpacked_sub(Axis::Short, vnz, 13, 7); // tall + fill_unpacked_sub(Axis::Short, vnz, 10, 10); // square + } +} + +TEST_F(TestSparseSkOpConstruction, fill_unpacked_sub_laso) { + for (int64_t vnz : {(int64_t) 1, (int64_t) 2, (int64_t) 4, (int64_t) 7}) { + fill_unpacked_sub(Axis::Long, vnz, 7, 13); // wide + fill_unpacked_sub(Axis::Long, vnz, 13, 7); // tall + fill_unpacked_sub(Axis::Long, vnz, 10, 10); // square + } +} + +TEST_F(TestSparseSkOpConstruction, fill_unpacked_sub_saso_int32) { + for (int64_t vnz : {(int64_t) 1, (int64_t) 2, (int64_t) 4, (int64_t) 7}) { + fill_unpacked_sub(Axis::Short, vnz, 7, 13); + fill_unpacked_sub(Axis::Short, vnz, 13, 7); + } +} + +TEST_F(TestSparseSkOpConstruction, fill_unpacked_sub_laso_int32) { + for (int64_t vnz : {(int64_t) 1, (int64_t) 2, (int64_t) 4, (int64_t) 7}) { + fill_unpacked_sub(Axis::Long, vnz, 7, 13); + fill_unpacked_sub(Axis::Long, vnz, 13, 7); + } +} + //////////////////////////////////////////////////////////////////////// // // diff --git a/test/linops/test_lskges.cc b/test/linops/test_lskges.cc index bcb4cee3..d4d85d0d 100644 --- a/test/linops/test_lskges.cc +++ b/test/linops/test_lskges.cc @@ -574,3 +574,75 @@ TEST_F(TestLSKGES, nontrivial_scales_rowmajor2) double beta = -1.0; alpha_beta(0, alpha, beta, 21, 4, blas::Layout::RowMajor); } + + +//////////////////////////////////////////////////////////////////////// +// +// +// Submatrix codepath (S.nnz < 0): generate only the needed submatrix. +// Check it matches the result of first filling the full operator. +// +// +//////////////////////////////////////////////////////////////////////// + +class TestLSKGES_SubmatrixPath : public ::testing::Test { + protected: + virtual void SetUp(){}; + virtual void TearDown(){}; + + template + void run(Axis major_axis, blas::Op opS, blas::Layout layout, int64_t vec_nnz) { + // op(submat(S)) is d1-by-m1; submat(S) lives at (S_ro, S_co) of the larger D0. + int64_t d1 = 5, m1 = 8, n = 6; + int64_t S_ro = 2, S_co = 3; + int64_t big_rows, big_cols; + if (opS == blas::Op::NoTrans) { big_rows = d1 + S_ro + 1; big_cols = m1 + S_co + 2; } + else { big_rows = m1 + S_ro + 1; big_cols = d1 + S_co + 2; } + SparseDist D0 {big_rows, big_cols, vec_nnz, major_axis}; + RandBLAS::RNGState seed((uint32_t) 7); + + // Dense input A: op(A) is m1-by-n with opA = NoTrans, so A is m1-by-n. + int64_t lda = (layout == blas::Layout::ColMajor) ? m1 : n; + std::vector A(m1 * n); + RandBLAS::RNGState a_state((uint32_t) 99); + RandBLAS::DenseDist DA {m1, n}; + RandBLAS::fill_dense(DA, A.data(), a_state); + + int64_t ldb = (layout == blas::Layout::ColMajor) ? d1 : n; + std::vector B_old(d1 * n, (T) 0); + std::vector B_new(d1 * n, (T) 0); + + // Old path: fill the FULL operator, then sketch its submatrix. + SparseSkOp S_old(D0, seed); + RandBLAS::fill_sparse(S_old); + RandBLAS::sketch_general(layout, opS, blas::Op::NoTrans, d1, n, m1, + (T) 1.0, S_old, S_ro, S_co, A.data(), lda, (T) 0.0, B_old.data(), ldb); + + // New path: do NOT fill (S.nnz < 0), so lskges builds only the submatrix. + SparseSkOp S_new(D0, seed); + RandBLAS::sketch_general(layout, opS, blas::Op::NoTrans, d1, n, m1, + (T) 1.0, S_new, S_ro, S_co, A.data(), lda, (T) 0.0, B_new.data(), ldb); + + auto msg = RandBLAS::testing::buffs_approx_equal( + B_new.data(), B_old.data(), d1 * n, __PRETTY_FUNCTION__, __FILE__, __LINE__ + ); + if (msg.size() > 0) { + FAIL() << msg; + } + return; + } +}; + +TEST_F(TestLSKGES_SubmatrixPath, equivalence_to_full_fill) { + using L = blas::Layout; + using O = blas::Op; + for (auto ax : {Axis::Short, Axis::Long}) { + for (auto opS : {O::NoTrans, O::Trans}) { + for (auto layout : {L::ColMajor, L::RowMajor}) { + for (int64_t vnz : {(int64_t) 1, (int64_t) 2, (int64_t) 3}) { + run(ax, opS, layout, vnz); + } + } + } + } +} diff --git a/test/linops/test_rskges.cc b/test/linops/test_rskges.cc index 46a14a7c..56d0510b 100644 --- a/test/linops/test_rskges.cc +++ b/test/linops/test_rskges.cc @@ -374,7 +374,7 @@ TEST_F(TestRSKGES, subset_cols_s_rowmajor1) ); } -TEST_F(TestRSKGES, subset_cols_s_rowmajor2) +TEST_F(TestRSKGES, subset_cols_s_rowmajor2) { for (uint32_t seed : {0}) submatrix_S(seed, @@ -385,3 +385,75 @@ TEST_F(TestRSKGES, subset_cols_s_rowmajor2) Layout::RowMajor ); } + + +//////////////////////////////////////////////////////////////////////// +// +// +// Submatrix codepath (S.nnz < 0): generate only the needed submatrix. +// Check it matches the result of first filling the full operator. +// +// +//////////////////////////////////////////////////////////////////////// + +class TestRSKGES_SubmatrixPath : public ::testing::Test { + protected: + virtual void SetUp(){}; + virtual void TearDown(){}; + + template + void run(RandBLAS::Axis major_axis, blas::Op opS, blas::Layout layout, int64_t vec_nnz) { + // op(submat(S)) is n-by-d1; submat(S) lives at (S_ro, S_co) of the larger D0. + int64_t m = 6, d1 = 5, n = 8; + int64_t S_ro = 2, S_co = 3; + int64_t big_rows, big_cols; + if (opS == blas::Op::NoTrans) { big_rows = n + S_ro + 1; big_cols = d1 + S_co + 2; } + else { big_rows = d1 + S_ro + 1; big_cols = n + S_co + 2; } + SparseDist D0 {big_rows, big_cols, vec_nnz, major_axis}; + RandBLAS::RNGState seed((uint32_t) 7); + + // Dense input A: op(A) is m-by-n with opA = NoTrans, so A is m-by-n. + int64_t lda = (layout == blas::Layout::ColMajor) ? m : n; + std::vector A(m * n); + RandBLAS::RNGState a_state((uint32_t) 99); + RandBLAS::DenseDist DA {m, n}; + RandBLAS::fill_dense(DA, A.data(), a_state); + + int64_t ldb = (layout == blas::Layout::ColMajor) ? m : d1; + std::vector B_old(m * d1, (T) 0); + std::vector B_new(m * d1, (T) 0); + + // Old path: fill the FULL operator, then sketch its submatrix. + SparseSkOp S_old(D0, seed); + RandBLAS::fill_sparse(S_old); + RandBLAS::sketch_general(layout, blas::Op::NoTrans, opS, m, d1, n, + (T) 1.0, A.data(), lda, S_old, S_ro, S_co, (T) 0.0, B_old.data(), ldb); + + // New path: do NOT fill (S.nnz < 0), so rskges builds only the submatrix. + SparseSkOp S_new(D0, seed); + RandBLAS::sketch_general(layout, blas::Op::NoTrans, opS, m, d1, n, + (T) 1.0, A.data(), lda, S_new, S_ro, S_co, (T) 0.0, B_new.data(), ldb); + + auto msg = RandBLAS::testing::buffs_approx_equal( + B_new.data(), B_old.data(), m * d1, __PRETTY_FUNCTION__, __FILE__, __LINE__ + ); + if (msg.size() > 0) { + FAIL() << msg; + } + return; + } +}; + +TEST_F(TestRSKGES_SubmatrixPath, equivalence_to_full_fill) { + using L = blas::Layout; + using O = blas::Op; + for (auto ax : {RandBLAS::Axis::Short, RandBLAS::Axis::Long}) { + for (auto opS : {O::NoTrans, O::Trans}) { + for (auto layout : {L::ColMajor, L::RowMajor}) { + for (int64_t vnz : {(int64_t) 1, (int64_t) 2, (int64_t) 3}) { + run(ax, opS, layout, vnz); + } + } + } + } +} From 6383c63c12d371fbf5ca781e6b37fb2f4c724589 Mon Sep 17 00:00:00 2001 From: Riley Murray Date: Sat, 30 May 2026 19:05:02 -0700 Subject: [PATCH 2/3] add workspace query to fill_sparse_unpacked --- RandBLAS/sparse_skops.hh | 53 ++++++++++++++++++++++++++++------------ 1 file changed, 37 insertions(+), 16 deletions(-) diff --git a/RandBLAS/sparse_skops.hh b/RandBLAS/sparse_skops.hh index c0b72e3c..8e091f56 100644 --- a/RandBLAS/sparse_skops.hh +++ b/RandBLAS/sparse_skops.hh @@ -531,14 +531,21 @@ void laso_merge_long_axis_vector_coo_data( /// Samples ONLY the \math{\ttt{n_rows_sub} \times \ttt{n_cols_sub}} submatrix of the /// operator defined by \math{(\D,\ttt{seed_state})} whose upper-left corner sits at /// \math{(\ttt{ro_s},\ttt{co_s})}, without ever materializing the full operator. This -/// function has no allocation stage and it skips checks for null pointers. +/// function has no allocation stage. /// -/// On entry, \math{(\vals,\rows,\cols)} are arrays of length at least -/// \math{\ttt{D.vec_nnz} \cdot \ttt{num_major_sub}}, where \math{\ttt{num_major_sub}} is -/// the number of major-axis vectors that intersect the requested submatrix (this is the -/// worst-case nonzero count for the submatrix). On exit, the first \math{\ttt{nnz}} -/// entries of these arrays contain the data for a COO sparse matrix representation of -/// \math{\submat(\mtxS)}, with row/column indices already shifted into +/// Workspace query. If any of \math{(\vals,\rows,\cols)} is null, then this +/// function performs no sampling: it writes the required length of each array to +/// \math{\ttt{nnz}} and returns \math{\ttt{seed_state}} unchanged. Use this to size your +/// buffers from \math{(\D,\ttt{n_rows_sub},\ttt{n_cols_sub},\ttt{ro_s},\ttt{co_s})} +/// without having to derive the worst-case nonzero count yourself. +/// +/// On entry (when \math{(\vals,\rows,\cols)} are all non-null), they must be arrays of +/// length at least \math{\ttt{D.vec_nnz} \cdot \ttt{num_major_sub}}, where +/// \math{\ttt{num_major_sub}} is the number of major-axis vectors that intersect the +/// requested submatrix (this is the worst-case nonzero count for the submatrix, and it is +/// exactly the value returned by the workspace query above). On exit, the first +/// \math{\ttt{nnz}} entries of these arrays contain the data for a COO sparse matrix +/// representation of \math{\submat(\mtxS)}, with row/column indices already shifted into /// \math{[0,\ttt{n_rows_sub}) \times [0,\ttt{n_cols_sub})}. /// /// The data produced by this function equals, entry-for-entry as a sparse matrix, the @@ -589,6 +596,17 @@ state_t fill_sparse_unpacked( major_is_rows = !short_is_rows; } + // Workspace query. If any of (vals, rows, cols) is null, we do not sample: we just + // report the required array length in nnz and return. The worst-case nonzero count + // for the requested submatrix is vec_nnz * num_major_sub (every nonzero of each + // sampled major-axis vector could land inside the window). Callers can use this to + // size (vals, rows, cols) from (D, n_rows_sub, n_cols_sub, ro_s, co_s) alone, rather + // than reconstructing the axis mapping themselves. + if (vals == nullptr || rows == nullptr || cols == nullptr) { + nnz = vec_nnz * num_major_sub; + return seed_state; + } + // Skip the RNG counter past the num_major_off major-axis vectors we don't need. // Both the Fisher-Yates path (vec_nnz > 1) and the i.i.d.-uniform path (vec_nnz == 1 // and LASO) consume exactly vec_nnz counter increments per major-axis vector, so the @@ -651,7 +669,7 @@ state_t fill_sparse_unpacked( // ============================================================================= /// This function is the underlying implementation of fill_sparse(SparseSkOp &S). -/// It has no allocation stage and it skips checks for null pointers. +/// It has no allocation stage. /// /// On entry, \math{(\vals,\rows,\cols)} are arrays of length at least \math{\Dfullnnz.} /// On exit, the first \math{\ttt{nnz}} entries of these arrays contain the data for @@ -666,6 +684,9 @@ state_t fill_sparse_unpacked_nosub( int64_t &nnz, T* vals, sint_t* rows, sint_t *cols, const state_t &seed_state ) { + randblas_require( vals != nullptr ); + randblas_require( rows != nullptr ); + randblas_require( cols != nullptr ); return fill_sparse_unpacked( D, D.n_rows, D.n_cols, 0, 0, nnz, vals, rows, cols, seed_state ); @@ -796,14 +817,14 @@ COOMatrix submatrix_as_coo( randblas_require(co_s + n_cols_sub <= S.n_cols); const SparseDist &D = S.dist; - // Worst-case nnz of the submatrix = vec_nnz * (number of major-axis vectors that - // intersect the submatrix). num_major_sub is that count; derive it via the same axis - // mapping used in fill_sparse_unpacked. - bool short_is_rows = (D.n_rows <= D.n_cols); - int64_t short_sub = short_is_rows ? n_rows_sub : n_cols_sub; - int64_t long_sub = short_is_rows ? n_cols_sub : n_rows_sub; - int64_t num_major_sub = (D.major_axis == Axis::Short) ? long_sub : short_sub; - int64_t cap = D.vec_nnz * num_major_sub; + // Ask fill_sparse_unpacked (via its workspace-query mode) how large the buffers must + // be, rather than reconstructing the axis mapping here. cap is the worst-case nonzero + // count for the requested submatrix. + int64_t cap; + fill_sparse_unpacked( + D, n_rows_sub, n_cols_sub, ro_s, co_s, cap, + (T*) nullptr, (sint_t*) nullptr, (sint_t*) nullptr, S.seed_state + ); // Allocate the worst-case buffers, sample only the requested submatrix, and attach // the buffers to an owning COOMatrix. We use the standard ctor + manual attach From 474b4fc5bad0b81171ddf2b97b5e8cbf4896d6f5 Mon Sep 17 00:00:00 2001 From: Riley Murray Date: Sat, 30 May 2026 19:46:59 -0700 Subject: [PATCH 3/3] Add submatrix sketching for SparseSkOp Sample only the requested submatrix of a sparse sketching operator without materializing the full operator, mirroring the DenseSkOp machinery. - fill_sparse_unpacked(D, n_rows_sub, n_cols_sub, ro_s, co_s, ...): samples the requested submatrix directly into COO buffers. Supports a null-pointer workspace query that reports the required buffer length in nnz. - submatrix_as_coo(...): sparse analog of submatrix_as_blackbox(); uses the workspace query to size its buffers. - lskges/rskges generate only the needed submatrix on the S.nnz < 0 path and call left_spmm/right_spmm with offsets (0,0). - fill_sparse now routes through fill_sparse_unpacked. fill_sparse_unpacked_nosub is retained only for backward compatibility in the 1.x series (removal in 2.0). - RTD: document fill_sparse_unpacked, drop fill_sparse_unpacked_nosub. Thread-independence and exact RNG reproducibility are preserved. 435/435 tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- RandBLAS/sparse_skops.hh | 74 +++++++++++--------- rtd/source/api_reference/skops_and_dists.rst | 2 +- 2 files changed, 42 insertions(+), 34 deletions(-) diff --git a/RandBLAS/sparse_skops.hh b/RandBLAS/sparse_skops.hh index 8e091f56..9284d25d 100644 --- a/RandBLAS/sparse_skops.hh +++ b/RandBLAS/sparse_skops.hh @@ -525,33 +525,46 @@ void laso_merge_long_axis_vector_coo_data( /// .. |vals| mathmacro:: \mathtt{vals} /// .. |rows| mathmacro:: \mathtt{rows} /// .. |cols| mathmacro:: \mathtt{cols} -/// .. |Dfullnnz| mathmacro:: {\mathcal{D}\mathtt{.full\_nnz}} /// /// @endverbatim -/// Samples ONLY the \math{\ttt{n_rows_sub} \times \ttt{n_cols_sub}} submatrix of the -/// operator defined by \math{(\D,\ttt{seed_state})} whose upper-left corner sits at -/// \math{(\ttt{ro_s},\ttt{co_s})}, without ever materializing the full operator. This -/// function has no allocation stage. +/// Sample the \math{\ttt{n_rows_sub} \times \ttt{n_cols_sub}} submatrix of \math{\mtxS} +/// whose upper-left corner is at \math{(\ttt{ro_s},\ttt{co_s}),} where \math{\mtxS} is +/// defined by \math{(\D,\ttt{seed_state}).} The submatrix is sampled directly, without +/// materializing the full operator, and is returned in COO format. /// -/// Workspace query. If any of \math{(\vals,\rows,\cols)} is null, then this -/// function performs no sampling: it writes the required length of each array to -/// \math{\ttt{nnz}} and returns \math{\ttt{seed_state}} unchanged. Use this to size your -/// buffers from \math{(\D,\ttt{n_rows_sub},\ttt{n_cols_sub},\ttt{ro_s},\ttt{co_s})} -/// without having to derive the worst-case nonzero count yourself. +/// If any of \math{(\vals,\rows,\cols)} is null, then no sampling occurs: the required +/// length of each output array is written to \math{\ttt{nnz},} and \math{\ttt{seed_state}} +/// is returned unchanged. Use this "workspace query" to size the output arrays. /// -/// On entry (when \math{(\vals,\rows,\cols)} are all non-null), they must be arrays of -/// length at least \math{\ttt{D.vec_nnz} \cdot \ttt{num_major_sub}}, where -/// \math{\ttt{num_major_sub}} is the number of major-axis vectors that intersect the -/// requested submatrix (this is the worst-case nonzero count for the submatrix, and it is -/// exactly the value returned by the workspace query above). On exit, the first -/// \math{\ttt{nnz}} entries of these arrays contain the data for a COO sparse matrix -/// representation of \math{\submat(\mtxS)}, with row/column indices already shifted into -/// \math{[0,\ttt{n_rows_sub}) \times [0,\ttt{n_cols_sub})}. +/// This function is the sparse analog of fill_dense_unpacked(). /// -/// The data produced by this function equals, entry-for-entry as a sparse matrix, the -/// \math{[\ttt{ro_s}:\ttt{ro_s}+\ttt{n_rows_sub},\,\ttt{co_s}:\ttt{co_s}+\ttt{n_cols_sub}]} -/// submatrix of \math{\ttt{fill_sparse_unpacked_nosub}(\D,\ldots)}. +/// @verbatim embed:rst:leading-slashes +/// .. dropdown:: Full parameter descriptions +/// :animate: fade-in-slide-down +/// +/// D +/// - A SparseDist that defines the full operator :math:`\mtxS.` +/// +/// n_rows_sub, n_cols_sub +/// - The number of rows and columns in the submatrix to sample. +/// +/// ro_s, co_s +/// - The row and column offsets of the submatrix as a part of :math:`\mtxS.` +/// +/// nnz +/// - On exit: the number of nonzeros written to the output arrays. +/// - For a workspace query: the required length of each output array. /// +/// vals, rows, cols +/// - Output buffers for the COO data of the submatrix, with indices shifted into +/// :math:`[0,\ttt{n_rows_sub}) \times [0,\ttt{n_cols_sub}).` +/// - Each must have length at least the value reported by a workspace query. +/// - Pass any of them as null to perform a workspace query instead of sampling. +/// +/// seed_state +/// - A CBRNG state used to define :math:`\mtxS.` +/// +/// @endverbatim template state_t fill_sparse_unpacked( const SparseDist &D, @@ -571,7 +584,7 @@ state_t fill_sparse_unpacked( int64_t vec_nnz = D.vec_nnz; // Map the submatrix's (row, col) offsets/extents onto the (short, long) axes, - // exactly as fill_sparse_unpacked_nosub chooses idxs_short/idxs_long. + // matching how the full operator assigns its short/long index arrays. bool short_is_rows = (D.n_rows <= D.n_cols); int64_t short_off, short_sub, long_off, long_sub; if (short_is_rows) { @@ -668,16 +681,11 @@ state_t fill_sparse_unpacked( } // ============================================================================= -/// This function is the underlying implementation of fill_sparse(SparseSkOp &S). -/// It has no allocation stage. -/// -/// On entry, \math{(\vals,\rows,\cols)} are arrays of length at least \math{\Dfullnnz.} -/// On exit, the first \math{\ttt{nnz}} entries of these arrays contain the data for -/// a COO sparse matrix representation of the SparseSkOp -/// defined by \math{(\D,\ttt{seed_state)}.} -/// -/// The function fill_sparse_unpacked lets you perform a similar operation that -/// only returns data for a submatrix of a SparseSkOp. +// DEPRECATED: retained only for backward compatibility in the RandBLAS 1.x release +// series, and scheduled for removal in RandBLAS 2. Use fill_sparse_unpacked with +// ro_s = co_s = 0 and the full operator dimensions instead. It writes the COO data for +// the operator (D, seed_state) into the first nnz entries of (vals, rows, cols), which +// must have length at least D.full_nnz. template state_t fill_sparse_unpacked_nosub( const SparseDist &D, @@ -725,7 +733,7 @@ void fill_sparse(SparseSkOp &S) { randblas_require(S.rows != nullptr); randblas_require(S.cols != nullptr); randblas_require(S.vals != nullptr); - fill_sparse_unpacked_nosub(S.dist, S.nnz, S.vals, S.rows, S.cols, S.seed_state); + fill_sparse_unpacked(S.dist, S.dist.n_rows, S.dist.n_cols, 0, 0, S.nnz, S.vals, S.rows, S.cols, S.seed_state); // ^ We ignore the return value from that function call. return; } diff --git a/rtd/source/api_reference/skops_and_dists.rst b/rtd/source/api_reference/skops_and_dists.rst index 40bc2b13..8a7c35e5 100644 --- a/rtd/source/api_reference/skops_and_dists.rst +++ b/rtd/source/api_reference/skops_and_dists.rst @@ -105,7 +105,7 @@ Sparse sketching, with CountSketch *et al.* .. doxygenfunction:: RandBLAS::fill_sparse(SparseSkOp &S) :project: RandBLAS - .. doxygenfunction:: RandBLAS::fill_sparse_unpacked_nosub(const SparseDist &D, int64_t &nnz, T* vals, sint_t* rows, sint_t* cols, const state_t &seed_state) + .. doxygenfunction:: RandBLAS::fill_sparse_unpacked :project: RandBLAS