diff --git a/AGENTS.md b/AGENTS.md index af91f65e..e42a67b3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -10,6 +10,8 @@ RandBLAS is a header-only C++ library for sketching in randomized linear algebra - **Web Documentation**: https://randblas.readthedocs.io/en/1.1.0/ - **Main Repository**: https://github.com/BallisticLA/RandBLAS +- **Style Guide**: `STYLE_GUIDE.md`. Improvements to this guide are always in + scope for any pull request, regardless of that pull request's primary purpose. - **DevNotes**: Critical implementation details are in `RandBLAS/DevNotes.md`, `RandBLAS/sparse_data/DevNotes.md`, and `test/DevNotes.md` ## Architecture and Code Organization diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 7227679d..e03bf723 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -5,7 +5,8 @@ reproducible random-number generation. This guide explains the project constraints and development workflow that matter when changing it. Follow the [`STYLE_GUIDE.md`](STYLE_GUIDE.md) for source and documentation -conventions. +conventions. Improvements to the style guide are always in scope for any pull +request, even when they are not required by its primary purpose. ## Before you start diff --git a/RandBLAS/DevNotes.md b/RandBLAS/DevNotes.md index 30b02218..01531843 100644 --- a/RandBLAS/DevNotes.md +++ b/RandBLAS/DevNotes.md @@ -5,26 +5,27 @@ for our user guide. * Our basic random number generation is handled by [Random123](https://github.com/DEShawResearch/random123). - We have small wrappers around Random123 code in ``RandBLAS/base.hh`` and ``RandBLAS/random_gen.hh``. + We have small wrappers around Random123 code in `RandBLAS/base.hh` and `RandBLAS/random_gen.hh`. - * ``RandBLAS/dense_skops.hh`` has code for representing and sampling dense sketching operators. + * `RandBLAS/dense_skops.hh` has code for representing and sampling dense sketching operators. The sampling code is complicated because it supports multi-threaded random (sub)matrix generation, and yet the generated (sub)matrices are the same no matter how many threads you're using. - * ``RandBLAS/sparse_skops.hh`` has code for representing and sampling sparse sketching operators. + * `RandBLAS/sparse_skops.hh` has code for representing and sampling sparse sketching operators. The sampling code has a customized method for repeatedly sampling from an index set without replacement, which is needed to quickly generate the structures used in statistically reliable - sparse sketching operators. + sparse sketching operators. See [Sparse sampling and OpenMP](#sparse-sampling-and-openmp) + for details. * [BLAS++ (aka blaspp)](https://github.com/icl-utk-edu/blaspp) is our portability layer for BLAS. We actually use very few functions in BLAS at time of writing (GEMM, SCAL, COPY, and AXPY) but we use its enumerations _everywhere_. Fast GEMM is important for sketching dense data with dense operators. - * The ``sketch_general`` functions in ``RandBLAS/skge.hh`` are the main entry point for sketching dense data. + * The `sketch_general` functions in `RandBLAS/skge.hh` are the main entry point for sketching dense data. These functions are small wrappers around functions with more BLAS-like names: - * ``lskge3`` and ``rskge3`` are basically wrappers around GEMM. - * ``lskges`` and ``rskges`` trigger an opaque call sequence that uses sparse matrix operations. + * `lskge3` and `rskge3` are basically wrappers around GEMM. + * `lskges` and `rskges` trigger an opaque call sequence that uses sparse matrix operations. * There is no widely accepted standard for sparse BLAS operations. This is a bummer because sparse matrices are super important in data science and scientific computing. In view of this, @@ -32,4 +33,37 @@ for our user guide. The abstractions can either own their associated data or just wrap existing data (say, data attached to a sparse matrix in Eigen). RandBLAS has reasonably flexible and high-performance code for multiplying a sparse matrix and a dense matrix. All code related to sparse matrices is in - ``RandBLAS/sparse_data``. See that folder's [``DevNotes.md``](sparse_data/DevNotes.md) file for details. + `RandBLAS/sparse_data`. See that folder's [`DevNotes.md`](sparse_data/DevNotes.md) file for details. + +## Sparse sampling and OpenMP + +Sparse sampling assigns randomness to logical major-axis vectors rather than to physical +threads. If the first requested vector starts at `initial_counter`, then vector `i` starts at +`initial_counter + i * vec_nnz`. The state returned by the sampling routine is computed outside +the OpenMP region by adding `num_major_axis_vectors * vec_nnz` to the initial counter. Thread +scheduling therefore cannot change either the sampled operator or the returned state. + +Products that determine addressable output sizes are checked at sampling boundaries before +inner-loop offsets are formed. Once a boundary check establishes the size, the loops reuse that +invariant instead of checking every offset. Overflow of Random123's extended-width counter is a +separate matter: its unsigned wrap-around behavior is intentional and benign, so RandBLAS does +not treat counter wrap-around as an error. + +Short-axis-sparse operators (SASOs) sample without replacement. Each active thread owns a +restored permutation of length `dim_major` and a pivot array of length `vec_nnz`; no thread +shares mutable sampling workspace. This gives an `O(T * dim_major)` permutation-workspace cost +for `T` active threads. An internal policy limits `T` by the available OpenMP threads, the number +of major-axis vectors, the amount of sampling work, and the work available to amortize each +permutation. Workspace storage is allocated before entering an OpenMP region, while each thread +initializes its own permutation. The specialized `vec_nnz == 1` path does not allocate +permutation workspace. A full-major-coordinate SASO is already packed after its per-vector sort; +only a partial window needs the serial filtering pass. + +Long-axis-sparse operators (LASOs) sample with replacement. Vector `i` first occupies a lane of +length `vec_nnz` at offset `i * vec_nnz`. A thread-private pair of hash maps merges duplicate +locations, the surviving entries are sorted by major coordinate, and a per-vector count records +the live prefix of each lane. The maps are constructed and reserved before the OpenMP region. +Allocation failures during insertion are captured inside that region and rethrown after all +threads leave it. A serial pass then packs lanes in increasing vector order. Every packed +destination precedes or equals its source, so this pass is safe in place and preserves the +canonical COO ordering without a second `O(nnz)` buffer. diff --git a/RandBLAS/base.hh b/RandBLAS/base.hh index 0a6d4f46..7dac8f1c 100644 --- a/RandBLAS/base.hh +++ b/RandBLAS/base.hh @@ -32,13 +32,17 @@ /// @file #include "RandBLAS/config.h" +#include "RandBLAS/exceptions.hh" #include "RandBLAS/random_gen.hh" #include -#include #include #include #include +#include +#include +#include +#include #if defined(RandBLAS_HAS_OpenMP) #include @@ -191,6 +195,24 @@ inline blas::Layout flipped_layout(const blas::Layout &layout_before) { return (layout_before == Layout::RowMajor) ? Layout::ColMajor : Layout::RowMajor; } +/// Return the calling thread's number, or zero when OpenMP is unavailable. +inline int randblas_get_thread_num() { +#if defined(RandBLAS_HAS_OpenMP) + return omp_get_thread_num(); +#else + return 0; +#endif +} + +/// Return the current OpenMP team size, or one when OpenMP is unavailable. +inline int randblas_get_num_threads() { +#if defined(RandBLAS_HAS_OpenMP) + return omp_get_num_threads(); +#else + return 1; +#endif +} + /** * Stores stride information for a matrix represented as a buffer. * The intended semantics for a buffer "A" and the conceptualized @@ -205,6 +227,22 @@ struct stride_64t { int64_t inter_col_stride; // step along a row }; +/// Require a submatrix window to lie within its parent matrix. +inline void validate_submat_dims( + int64_t parent_rows, int64_t parent_cols, + int64_t n_rows_sub, int64_t n_cols_sub, + int64_t ro, int64_t co +) { + randblas_require(n_rows_sub >= 0); + randblas_require(n_cols_sub >= 0); + randblas_require(ro >= 0); + randblas_require(co >= 0); + randblas_require(n_rows_sub <= parent_rows); + randblas_require(n_cols_sub <= parent_cols); + randblas_require(ro <= parent_rows - n_rows_sub); + randblas_require(co <= parent_cols - n_cols_sub); +} + inline stride_64t layout_to_strides(blas::Layout layout, int64_t ldim) { if (layout == blas::Layout::ColMajor) { return stride_64t{(int64_t) 1, ldim}; @@ -244,10 +282,10 @@ inline submat_spec_64t offset_and_ldim( ) { if (layout == blas::Layout::ColMajor) { int64_t offset = ro_s + n_rows * co_s; - return submat_spec_64t{offset, n_rows}; + return submat_spec_64t{offset, std::max(n_rows, (int64_t)1)}; } else { int64_t offset = ro_s * n_cols + co_s; - return submat_spec_64t{offset, n_cols}; + return submat_spec_64t{offset, std::max(n_cols, (int64_t)1)}; } } @@ -262,18 +300,32 @@ concept SignedInteger = (std::numeric_limits::is_signed && std::numeric_limit template inline TO safe_int_product(TI a, TI b) { - if (a == 0 || b == 0) { - return 0; + static_assert( + std::numeric_limits::digits >= std::numeric_limits::digits, + "safe_int_product requires an output type at least as wide as its input type." + ); + const TO a_out = static_cast(a); + const TO b_out = static_cast(b); + const TO min_out = std::numeric_limits::min(); + const TO max_out = std::numeric_limits::max(); + bool overflow = false; + if (b_out > 0) { + const TO min_safe_a = min_out / b_out; + const TO max_safe_a = max_out / b_out; + overflow = a_out < min_safe_a || a_out > max_safe_a; + } else if (b_out == -1) { + overflow = a_out == min_out; + } else if (b_out < -1) { + const TO min_safe_a = max_out / b_out; + const TO max_safe_a = min_out / b_out; + overflow = a_out < min_safe_a || a_out > max_safe_a; } - TO c = a * b; - TO b_check = c / a; - TO a_check = c / b; - if ((a_check != a) || (b_check != b)) { + if (overflow) { std::stringstream s; - s << "Overflow when multiplying a (=" << a << ") and b(=" << b << "), which resulted in " << c << ".\n"; + s << "Overflow when multiplying a (=" << a << ") and b (=" << b << ").\n"; throw std::overflow_error(s.str()); } - return c; + return a_out * b_out; } diff --git a/RandBLAS/dense_skops.hh b/RandBLAS/dense_skops.hh index b89fa91e..d9b73716 100644 --- a/RandBLAS/dense_skops.hh +++ b/RandBLAS/dense_skops.hh @@ -35,7 +35,9 @@ #include +#include #include +#include #include #include #include @@ -95,35 +97,48 @@ inline void copy_promote(int n, const T_IN &a, T_OUT* b) { */ template static RNGState fill_dense_submat_impl(int64_t n_cols, T* smat, int64_t n_srows, int64_t n_scols, int64_t ptr, const RNGState &seed, int64_t lda = 0) { + randblas_require(n_cols > 0); + randblas_require(n_srows > 0); + randblas_require(n_scols > 0); + randblas_require(ptr >= 0); if (lda <= 0) { lda = n_scols; } else { randblas_require(lda >= n_scols); } randblas_require(n_cols >= n_scols); + // Validate the documented output span before deriving row offsets in parallel. + (void)safe_int_product(n_srows, lda); RNG rng; using CTR_t = typename RNG::ctr_type; using KEY_t = typename RNG::key_type; const int64_t ctr_size = CTR_t::static_size; - - int64_t pad = 0; - // ^ computed such that n_cols+pad is divisible by ctr_size - if (n_cols % ctr_size != 0) { - pad = ctr_size - n_cols % ctr_size; - } - - const int64_t ptr_padded = ptr + ptr / n_cols * pad; - // ^ ptr corresponding to the padded matrix - const int64_t ctr_mat_start = ptr_padded / ctr_size; - const int64_t first_block_start = ptr_padded % ctr_size; - // ^ counter and [position within the counter's array] for index "ptr_padded". - const int64_t ctr_mat_row_end = (ptr_padded + n_scols - 1) / ctr_size; - const int64_t last_block_stop = ((ptr_padded + n_scols - 1) % ctr_size) + 1; - // ^ counter and [1 + position within the counter's array] for index "(ptr_padded + n_scols - 1)". - const int64_t ctr_inter_row_stride = (n_cols + pad) / ctr_size; + const int64_t ctr_inter_row_stride = n_cols / ctr_size + (n_cols % ctr_size != 0); // ^ number of counters between the first counter of a given row to the first counter of the next row; + const int64_t parent_row = ptr / n_cols; + const int64_t parent_col = ptr % n_cols; + const int64_t parent_row_ctr_offset = safe_int_product(parent_row, ctr_inter_row_stride); + const uint64_t ctr_mat_start_wide = static_cast(parent_row_ctr_offset) + + static_cast(parent_col / ctr_size); + const uint64_t last_row_offset = static_cast(parent_col) + + static_cast(n_scols) - 1; + const uint64_t ctr_mat_row_end_wide = static_cast(parent_row_ctr_offset) + + last_row_offset / static_cast(ctr_size); + const uint64_t max_int64 = std::numeric_limits::max(); + if (ctr_mat_start_wide > max_int64 || ctr_mat_row_end_wide > max_int64) { + throw std::overflow_error("Overflow when computing a dense submatrix's counter offsets.\n"); + } + const int64_t ctr_mat_start = static_cast(ctr_mat_start_wide); + const int64_t first_block_start = parent_col % ctr_size; + // ^ counter and [position within the counter's array] for index "ptr". + const int64_t ctr_mat_row_end = static_cast(ctr_mat_row_end_wide); + const int64_t last_block_stop = static_cast( + last_row_offset % static_cast(ctr_size) + ) + 1; + // ^ counter and [1 + position within the counter's array] for the last requested column. const bool one_block_per_row = ctr_mat_start == ctr_mat_row_end; const int64_t first_block_len = ((one_block_per_row) ? last_block_stop : ctr_size) - first_block_start; + const int64_t full_incr = safe_int_product(n_srows, ctr_inter_row_stride); CTR_t temp_c = seed.counter; temp_c.incr(ctr_mat_start); @@ -132,9 +147,7 @@ static RNGState fill_dense_submat_impl(int64_t n_cols, T* smat, int64_t n_s #pragma omp parallel for schedule(static) for (int64_t row = 0; row < n_srows; row++) { - - int64_t incr_from_c = safe_int_product(ctr_inter_row_stride, row); - + int64_t incr_from_c = ctr_inter_row_stride * row; auto c_row = c; c_row.incr(incr_from_c); auto rv = OP::generate(rng, c_row, k); @@ -148,7 +161,8 @@ static RNGState fill_dense_submat_impl(int64_t n_cols, T* smat, int64_t n_s } // middle blocks int64_t ind = first_block_len; - for (int i = 0; i < (ctr_mat_row_end - ctr_mat_start - 1); ++i) { + const int64_t n_middle_blocks = ctr_mat_row_end - ctr_mat_start - 1; + for (int64_t block = 0; block < n_middle_blocks; ++block) { c_row.incr(); rv = OP::generate(rng, c_row, k); copy_promote(ctr_size, rv, smat_row + ind); @@ -162,7 +176,7 @@ static RNGState fill_dense_submat_impl(int64_t n_cols, T* smat, int64_t n_s // find the largest counter in the counter array CTR_t max_c = c; - max_c.incr(n_srows * ctr_inter_row_stride); + max_c.incr(full_incr); return RNGState {max_c, k}; } @@ -171,11 +185,7 @@ RNGState compute_next_state(DD dist, RNGState state) { int64_t major_len = dist.dim_major; int64_t minor_len = dist.dim_minor; int64_t ctr_size = RNG::ctr_type::static_size; - int64_t pad = 0; - if (major_len % ctr_size != 0) { - pad = ctr_size - major_len % ctr_size; - } - int64_t ctr_major_axis_stride = (major_len + pad) / ctr_size; + int64_t ctr_major_axis_stride = major_len / ctr_size + (major_len % ctr_size != 0); int64_t full_incr = safe_int_product(ctr_major_axis_stride, minor_len); state.counter.incr(full_incr); return state; @@ -560,21 +570,30 @@ static_assert(SketchingOperator>); template RNGState fill_dense_unpacked(blas::Layout layout, const DenseDist &D, int64_t n_rows, int64_t n_cols, int64_t ro_s, int64_t co_s, T* buff, const RNGState &seed) { using RandBLAS::dense::fill_dense_submat_impl; - randblas_require(D.n_rows >= n_rows + ro_s); - randblas_require(D.n_cols >= n_cols + co_s); + validate_submat_dims(D.n_rows, D.n_cols, n_rows, n_cols, ro_s, co_s); + if (n_rows == 0 || n_cols == 0) { + return seed; + } + const int64_t size_mat = safe_int_product(n_rows, n_cols); blas::Layout natural_layout = D.natural_layout; int64_t ma_len = D.dim_major; - int64_t n_rows_, n_cols_, ptr; + int64_t n_rows_, n_cols_, major_offset, minor_offset; if (natural_layout == blas::Layout::ColMajor) { // operate on the transpose in row-major n_rows_ = n_cols; n_cols_ = n_rows; - ptr = ro_s + safe_int_product(co_s, ma_len); + major_offset = safe_int_product(co_s, ma_len); + minor_offset = ro_s; } else { n_rows_ = n_rows; n_cols_ = n_cols; - ptr = safe_int_product(ro_s, ma_len) + co_s; + major_offset = safe_int_product(ro_s, ma_len); + minor_offset = co_s; + } + if (minor_offset > std::numeric_limits::max() - major_offset) { + throw std::overflow_error("Overflow when computing the dense submatrix's starting offset.\n"); } + const int64_t ptr = major_offset + minor_offset; RNGState next_state{}; switch (D.family) { case ScalarDist::Gaussian: { @@ -583,14 +602,13 @@ RNGState fill_dense_unpacked(blas::Layout layout, const DenseDist &D, int64 } case ScalarDist::Uniform: { next_state = fill_dense_submat_impl(ma_len, buff, n_rows_, n_cols_, ptr, seed); - blas::scal(n_rows_ * n_cols_, (T)std::sqrt(3), buff, 1); + blas::scal(size_mat, (T)std::sqrt(3), buff, 1); break; } default: { throw std::runtime_error(std::string("Unrecognized distribution.")); } } - int64_t size_mat = n_rows * n_cols; if (layout != natural_layout) { T* flip_work = new T[size_mat]; blas::copy(size_mat, buff, 1, flip_work, 1); @@ -647,7 +665,8 @@ template void fill_dense(DenseSkOp &S) { if (S.own_memory && S.buff == nullptr) { using T = typename DenseSkOp::scalar_t; - S.buff = new T[S.n_rows * S.n_cols]; + const int64_t size_mat = safe_int_product(S.n_rows, S.n_cols); + S.buff = new T[size_mat]; } randblas_require(S.buff != nullptr); fill_dense_unpacked(S.layout, S.dist, S.n_rows, S.n_cols, 0, 0, S.buff, S.seed_state); @@ -673,10 +692,10 @@ struct BLASFriendlyOperator { template BFO submatrix_as_blackbox(const DenseSkOp &S, int64_t n_rows, int64_t n_cols, int64_t ro_s, int64_t co_s) { - randblas_require(ro_s + n_rows <= S.n_rows); - randblas_require(co_s + n_cols <= S.n_cols); + validate_submat_dims(S.n_rows, S.n_cols, n_rows, n_cols, ro_s, co_s); + const int64_t size_mat = safe_int_product(n_rows, n_cols); using T = typename DenseSkOp::scalar_t; - T *buff = new T[n_rows * n_cols]; + T *buff = new T[size_mat]; auto layout = S.layout; fill_dense_unpacked(layout, S.dist, n_rows, n_cols, ro_s, co_s, buff, S.seed_state); int64_t dim_major = S.dist.dim_major; diff --git a/RandBLAS/skge.hh b/RandBLAS/skge.hh index aab9f58c..dddc5aa4 100644 --- a/RandBLAS/skge.hh +++ b/RandBLAS/skge.hh @@ -33,6 +33,7 @@ #include "RandBLAS/random_gen.hh" #include "RandBLAS/dense_skops.hh" #include "RandBLAS/sparse_skops.hh" +#include "RandBLAS/util.hh" #include #include @@ -182,8 +183,7 @@ void lskge3( } // else, continue with the function as usual. } randblas_require( S.buff != nullptr ); - randblas_require( S.n_rows >= rows_submat_S + ro_s ); - randblas_require( S.n_cols >= cols_submat_S + co_s ); + validate_submat_dims(S.n_rows, S.n_cols, rows_submat_S, cols_submat_S, ro_s, co_s); auto [rows_A, cols_A] = dims_before_op(m, n, opA); if (layout == blas::Layout::ColMajor) { randblas_require(lda >= rows_A); @@ -335,8 +335,7 @@ void rskge3( } } randblas_require( S.buff != nullptr ); - randblas_require( S.n_rows >= rows_submat_S + ro_s ); - randblas_require( S.n_cols >= cols_submat_S + co_s ); + validate_submat_dims(S.n_rows, S.n_cols, rows_submat_S, cols_submat_S, ro_s, co_s); auto [rows_A, cols_A] = dims_before_op(m, n, opA); if (layout == blas::Layout::ColMajor) { randblas_require(lda >= rows_A); @@ -554,6 +553,7 @@ void lskges( int64_t ldb ) { auto [n_rows, n_cols] = dims_before_op(d, m, opS); + validate_submat_dims(S.n_rows, S.n_cols, n_rows, n_cols, ro_s, co_s); if (S.nnz < 0) { 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); @@ -693,6 +693,7 @@ inline void rskges( int64_t ldb ) { auto [n_rows, n_cols] = dims_before_op(n, d, opS); + validate_submat_dims(S.n_rows, S.n_cols, n_rows, n_cols, ro_s, co_s); if (S.nnz < 0) { 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); diff --git a/RandBLAS/sparse_data/csc_spmm_impl.hh b/RandBLAS/sparse_data/csc_spmm_impl.hh index 1f85e4de..481005ec 100644 --- a/RandBLAS/sparse_data/csc_spmm_impl.hh +++ b/RandBLAS/sparse_data/csc_spmm_impl.hh @@ -126,13 +126,8 @@ static void apply_csc_kib_1p1_rowmajor( #pragma omp parallel default(shared) { - #if defined(RandBLAS_HAS_OpenMP) - int t = omp_get_thread_num(); - int num_threads = omp_get_num_threads(); - #else - int t = 0; - int num_threads = 1; - #endif + const int t = randblas_get_thread_num(); + const int num_threads = randblas_get_num_threads(); int i_lower = (d * t) / num_threads; int i_upper = (d * (t + 1)) / num_threads; diff --git a/RandBLAS/sparse_data/mkl_spmm_impl.hh b/RandBLAS/sparse_data/mkl_spmm_impl.hh index ba65b2f3..12c0f5a3 100644 --- a/RandBLAS/sparse_data/mkl_spmm_impl.hh +++ b/RandBLAS/sparse_data/mkl_spmm_impl.hh @@ -46,6 +46,7 @@ #include "RandBLAS/sparse_data/coo_matrix.hh" #include "RandBLAS/sparse_data/csr_matrix.hh" #include "RandBLAS/sparse_data/csc_matrix.hh" +#include "RandBLAS/util.hh" namespace RandBLAS::sparse_data::mkl { @@ -260,9 +261,6 @@ MKLSparseHandle make_mkl_handle(const SpMat& A) { // MKL-accelerated left_spmm: C = alpha * op(A) * op(B) + beta * C // where A is sparse, B and C are dense. // -// NOTE: The caller (spmm_dispatch.hh) has already applied beta scaling to C, -// so this function is called with beta=0. -// // For COO matrices with submatrix offsets (ro_a, co_a != 0), we fall back // to the existing RandBLAS implementation since MKL doesn't support // submatrix views on COO. @@ -303,6 +301,16 @@ bool mkl_left_spmm( if (opB != blas::Op::NoTrans) return false; + // MKL rejects some valid empty sparse matrices when creating its handle. + // Handle those products directly. An empty output needs no work, while an + // empty contraction or an all-zero sparse operand leaves beta * C. + if (d == 0 || n == 0) + return true; + if (m == 0 || A.nnz == 0) { + RandBLAS::util::lascl(layout, d, n, beta, C, ldc); + return true; + } + // Build the MKL sparse handle and the operation to apply to it. // CSR / COO: the handle wraps A directly; the operation is opA. // CSC: mkl_sparse_?_mm does not accept CSC, but a CSC matrix's arrays ARE @@ -341,6 +349,9 @@ bool mkl_left_spmm( B, (MKL_INT)n, (MKL_INT)ldb, beta, C, (MKL_INT)ldc ); + } else { + // unsupported floating point type. + return false; } check_mkl_status(status, "mkl_sparse_mm"); return true; // signal: MKL handled it diff --git a/RandBLAS/sparse_data/sksp.hh b/RandBLAS/sparse_data/sksp.hh index f3eefdd8..37984ce8 100644 --- a/RandBLAS/sparse_data/sksp.hh +++ b/RandBLAS/sparse_data/sksp.hh @@ -32,6 +32,7 @@ #include "RandBLAS/base.hh" #include "RandBLAS/dense_skops.hh" #include "RandBLAS/exceptions.hh" +#include "RandBLAS/util.hh" namespace RandBLAS::sparse_data { @@ -162,10 +163,8 @@ void lsksp3( } randblas_require( S.buff != nullptr ); auto [rows_submat_A, cols_submat_A] = dims_before_op(m, n, opA); - randblas_require( A.n_rows >= rows_submat_A + ro_a ); - randblas_require( A.n_cols >= cols_submat_A + co_a ); - randblas_require( S.n_rows >= rows_submat_S + ro_s ); - randblas_require( S.n_cols >= cols_submat_S + co_s ); + validate_submat_dims(A.n_rows, A.n_cols, rows_submat_A, cols_submat_A, ro_a, co_a); + validate_submat_dims(S.n_rows, S.n_cols, rows_submat_S, cols_submat_S, ro_s, co_s); if (layout == blas::Layout::ColMajor) { randblas_require(ldb >= d); } else { @@ -306,10 +305,8 @@ void rsksp3( } randblas_require( S.buff != nullptr ); auto [rows_submat_A, cols_submat_A] = dims_before_op(m, n, opA); - randblas_require( A.n_rows >= rows_submat_A + ro_a ); - randblas_require( A.n_cols >= cols_submat_A + co_a ); - randblas_require( S.n_rows >= rows_submat_S + ro_s ); - randblas_require( S.n_cols >= cols_submat_S + co_s ); + validate_submat_dims(A.n_rows, A.n_cols, rows_submat_A, cols_submat_A, ro_a, co_a); + validate_submat_dims(S.n_rows, S.n_cols, rows_submat_S, cols_submat_S, ro_s, co_s); if (layout == blas::Layout::ColMajor) { randblas_require(ldb >= m); } else { diff --git a/RandBLAS/sparse_data/spmm_dispatch.hh b/RandBLAS/sparse_data/spmm_dispatch.hh index 30bf9a82..8df989a5 100644 --- a/RandBLAS/sparse_data/spmm_dispatch.hh +++ b/RandBLAS/sparse_data/spmm_dispatch.hh @@ -39,6 +39,7 @@ #include "RandBLAS/sparse_data/csc_spmm_impl.hh" #include "RandBLAS/sparse_data/csr_spmm_impl.hh" #include "RandBLAS/sparse_data/coo_spmm_impl.hh" +#include "RandBLAS/util.hh" #include "RandBLAS/config.h" #if defined(RandBLAS_HAS_MKL) #include "RandBLAS/sparse_data/mkl_spmm_impl.hh" @@ -69,13 +70,9 @@ void left_spmm( ) { using blas::Layout; using blas::Op; - // Applying a transposed sparse matrix reduces to the NoTrans case on a - // zero-copy transpose view (CSR<->CSC, COO<->COO). MKL, when present, - // engages on the recursive NoTrans call: mkl_left_spmm handles all three - // formats -- including CSC, which it consumes as a CSR-of-transpose view -- - // so the transposed CSR no longer needs to be pre-routed to MKL here to - // avoid a CSC fallback. if (opA == Op::Trans) { + // Applying a transposed sparse matrix reduces to the NoTrans case on a + // zero-copy transpose view (CSR<->CSC, COO<->COO). auto At = A.transpose(); left_spmm(layout, Op::NoTrans, opB, d, n, m, alpha, At, co_a, ro_a, B, ldb, beta, C, ldc); return; @@ -88,83 +85,65 @@ void left_spmm( constexpr bool is_csc = std::is_same_v>; randblas_require(is_coo || is_csr || is_csc); - if constexpr (is_coo) { - randblas_require(A.n_rows >= d); - randblas_require(A.n_cols >= m); - } else { + validate_submat_dims(A.n_rows, A.n_cols, d, m, ro_a, co_a); + if constexpr (!is_coo) { randblas_require(A.n_rows == d); randblas_require(A.n_cols == m); randblas_require(ro_a == 0); randblas_require(co_a == 0); } - // Dimensions of B, rather than \op(B) - Layout layout_C = layout; - Layout layout_opB; - int64_t rows_B, cols_B; - if (opB == Op::NoTrans) { - rows_B = m; - cols_B = n; - layout_opB = layout; - } else { - rows_B = n; - cols_B = m; - layout_opB = (layout == Layout::ColMajor) ? Layout::RowMajor : Layout::ColMajor; - } - - // Check dimensions and compute C = beta * C. - // Note: both B and C are checked based on "layout"; B is *not* checked on layout_opB. + // Check dimensions. Both B and C are checked based on "layout", even + // if we end up lying about B's layout later on to resolve a transpose. + auto [rows_B, cols_B] = dims_before_op(m, n, opB); if (layout == Layout::ColMajor) { randblas_require(ldb >= rows_B); randblas_require(ldc >= d); - for (int64_t i = 0; i < n; ++i) - RandBLAS::util::safe_scal(d, beta, &C[i*ldc]); } else { randblas_require(ldc >= n); randblas_require(ldb >= cols_B); - for (int64_t i = 0; i < d; ++i) - RandBLAS::util::safe_scal(n, beta, &C[i*ldc]); } - if (alpha == (T) 0) + if (alpha == (T) 0) { + RandBLAS::util::lascl(layout, d, n, beta, C, ldc); return; + } // Try MKL-accelerated path if available. #if defined(RandBLAS_HAS_MKL) if constexpr (sizeof(typename SpMat::index_t) == sizeof(MKL_INT)) { // mkl_left_spmm returns false if it can't handle this case - // (e.g., COO with submatrix offsets, or opB == Trans). CSC is handled - // via a CSR-of-transpose view inside mkl_left_spmm. - // Beta is already applied to C above, so pass beta=1 to MKL - // so it adds alpha*A*B to the existing (pre-scaled) C. + // (e.g., COO with submatrix offsets, or opB == Trans). bool handled = RandBLAS::sparse_data::mkl::mkl_left_spmm( layout, Op::NoTrans, opB, d, n, m, alpha, - A, ro_a, co_a, B, ldb, (T)1, C, ldc + A, ro_a, co_a, B, ldb, beta, C, ldc ); if (handled) return; } #endif - // Fallback: hand-rolled sparse kernels. + // RandBLAS-defined implementations + Layout layout_opB = (opB == Op::NoTrans) ? layout : flipped_layout(layout); + RandBLAS::util::lascl(layout, d, n, beta, C, ldc); // <-- TODO: update to perform beta scaling in SPMM kernels. if constexpr (is_coo) { 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); + apply_coo_via_csx(alpha, layout_opB, layout, 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) { + if (layout_opB == Layout::RowMajor && layout == Layout::RowMajor) { using RandBLAS::sparse_data::csc::apply_csc_kib_1p1_rowmajor; apply_csc_kib_1p1_rowmajor(alpha, n, A, B, ldb, C, ldc); } else { using RandBLAS::sparse_data::csc::apply_csc_jki_p11; - apply_csc_jki_p11(alpha, layout_opB, layout_C, n, A, B, ldb, C, ldc); + apply_csc_jki_p11(alpha, layout_opB, layout, n, A, B, ldb, C, ldc); } } else { - if (layout_opB == Layout::RowMajor && layout_C == Layout::RowMajor) { + if (layout_opB == Layout::RowMajor && layout == Layout::RowMajor) { using RandBLAS::sparse_data::csr::apply_csr_ikb_p1b_rowmajor; apply_csr_ikb_p1b_rowmajor(alpha, d, n, m, A, B, ldb, C, ldc); } else { using RandBLAS::sparse_data::csr::apply_csr_jik_p11; - apply_csr_jik_p11(alpha, layout_opB, layout_C, d, n, m, A, B, ldb, C, ldc); + apply_csr_jik_p11(alpha, layout_opB, layout, d, n, m, A, B, ldb, C, ldc); } } @@ -204,7 +183,7 @@ inline void right_spmm( using blas::Layout; using blas::Op; auto trans_opB = (opB == Op::NoTrans) ? Op::Trans : Op::NoTrans; - auto trans_layout = (layout == Layout::ColMajor) ? Layout::RowMajor : Layout::ColMajor; + auto trans_layout = flipped_layout(layout); left_spmm( trans_layout, trans_opB, opA, d, m, n, alpha, B, i_off, j_off, A, lda, beta, C, ldc diff --git a/RandBLAS/sparse_skops.hh b/RandBLAS/sparse_skops.hh index 70e4c55f..c737d0fc 100644 --- a/RandBLAS/sparse_skops.hh +++ b/RandBLAS/sparse_skops.hh @@ -36,11 +36,19 @@ #include "RandBLAS/sparse_data/spmm_dispatch.hh" #include + +#if defined(RandBLAS_HAS_OpenMP) +#include +#endif + #include #include #include +#include #include +#include #include +#include #include #define MAX(a, b) (((a) < (b)) ? (b) : (a)) @@ -48,6 +56,39 @@ namespace RandBLAS::sparse { +static inline int sparse_sampling_thread_count( + int64_t dim_major, int64_t num_major_axis_vectors, int64_t vec_nnz, bool uses_perm_work +) { + // Two tests are to this policy's constants. + // + // - fisher_yates_is_exact_at_parallel_policy_boundary + // - sparse_sampling_thread_policy_uses_available_threads + // + // Update those tests if you update this policy! +#if defined(RandBLAS_HAS_OpenMP) + int64_t num_threads = std::min( + omp_get_max_threads(), num_major_axis_vectors + ); + const int64_t useful_work = num_major_axis_vectors * vec_nnz; + if (useful_work < 1024) { + return 1; + } + if (uses_perm_work) { + const int64_t amortized_threads = std::max( + 1, useful_work / dim_major + ); + num_threads = std::min(num_threads, amortized_threads); + } + return static_cast(std::max(1, num_threads)); +#else + (void) dim_major; + (void) num_major_axis_vectors; + (void) vec_nnz; + (void) uses_perm_work; + return 1; +#endif +} + template > void _considerate_fisher_yates( const state_t &state, @@ -94,51 +135,143 @@ void _considerate_fisher_yates( return; } +template +static void sort_major_axis_vector(sint_t *idxs_major, T *vals, int64_t len) { + // Keep values paired with their major coordinates. The minor coordinate is + // constant within a vector, so it does not participate in the permutation. + // These vectors are normally short, making insertion sort allocation-free + // and inexpensive. + if (vals == nullptr) { + std::sort(idxs_major, idxs_major + len); + return; + } + for (int64_t i = 1; i < len; ++i) { + const sint_t major = idxs_major[i]; + const T val = vals[i]; + int64_t j = i - 1; + for (; j >= 0 && idxs_major[j] > major; --j) { + idxs_major[j + 1] = idxs_major[j]; + vals[j + 1] = vals[j]; + } + idxs_major[j + 1] = major; + vals[j + 1] = val; + } +} + template > -static state_t repeated_fisher_yates( +static void sample_singleton_vectors( const state_t &state, - int64_t vec_nnz, int64_t dim_major, int64_t dim_minor, sint_t *idxs_major, sint_t *idxs_minor, T *vals ) { - randblas_error_if(vec_nnz > dim_major); - if (vec_nnz == 1) { - // Sampling 1 element without replacement is the same as with replacement, - // so delegate to the cheaper i.i.d. sampler in a single batched call. - if (idxs_minor != nullptr) - std::iota(idxs_minor, idxs_minor + dim_minor, sint_t{0}); + // Sampling one location without replacement is ordinary uniform sampling. + // Give each physical thread one contiguous range of logical vectors so a + // thread can process its range with a single batched sampler call. + [[maybe_unused]] const int num_threads = sparse_sampling_thread_count( + dim_major, dim_minor, 1, false + ); + const auto base_ctr = state.counter; + #pragma omp parallel num_threads(num_threads) if(num_threads > 1) + { + const int tid = randblas_get_thread_num(); + const int team_size = randblas_get_num_threads(); + const int64_t chunk = dim_minor / team_size; + const int64_t rem = dim_minor % team_size; + const int64_t begin = tid * chunk + std::min(tid, rem); + const int64_t end = begin + chunk + (tid < rem); + auto chunk_ctr = base_ctr; + chunk_ctr.incr(begin); + state_t chunk_state{chunk_ctr, state.key}; if (vals != nullptr) { - return sample_indices_iid_uniform( - dim_major, dim_minor, idxs_major, vals, state); + sample_indices_iid_uniform( + dim_major, end - begin, idxs_major + begin, + vals + begin, chunk_state + ); } else { - return sample_indices_iid_uniform( - dim_major, dim_minor, idxs_major, state); + sample_indices_iid_uniform( + dim_major, end - begin, idxs_major + begin, chunk_state + ); + } + if (idxs_minor != nullptr) { + std::iota( + idxs_minor + begin, idxs_minor + end, static_cast(begin) + ); } } - std::vector vec_work(dim_major); - std::iota(vec_work.begin(), vec_work.end(), 0); - std::vector pivots(vec_nnz); - auto [ctr, key] = state; - for (sint_t i = 0; i < dim_minor; ++i) { - state_t state_work{ctr, state.key}; +} + + +template > +static state_t repeated_fisher_yates( + const state_t &state, + int64_t vec_nnz, + int64_t dim_major, + int64_t dim_minor, + sint_t *idxs_major, + sint_t *idxs_minor, + T *vals, + bool apply_sort +) { + randblas_error_if(vec_nnz > dim_major); + const int64_t full_incr = safe_int_product(dim_minor, vec_nnz); + if (vals != nullptr) { + randblas_require(state.len_c >= 4); + } else { + randblas_require(state.len_c >= 2); + } + auto end_ctr = state.counter; + end_ctr.incr(full_incr); + auto out = state_t{end_ctr, state.key}; + + if (vec_nnz == 1) { + sample_singleton_vectors(state, dim_major, dim_minor, idxs_major, idxs_minor, vals); + return out; + } + + // Each thread owns one reusable permutation and pivot workspace. Logical + // vector indices, rather than physical thread IDs, determine counter ranges. + const int num_threads = sparse_sampling_thread_count(dim_major, dim_minor, vec_nnz, true); + const int64_t perm_size = safe_int_product(dim_major, static_cast(num_threads)); + const int64_t pivot_size = safe_int_product(vec_nnz, static_cast(num_threads)); + std::vector perm_works(perm_size); + std::vector pivot_works(pivot_size); + + const auto base_ctr = state.counter; + auto sample_lane = [&](int64_t i, int tid) { + const int64_t offset = i * vec_nnz; + auto vec_ctr = base_ctr; + vec_ctr.incr(offset); + state_t vec_state{vec_ctr, state.key}; + sint_t *vec_major = idxs_major + offset; + sint_t *vec_minor = (idxs_minor == nullptr) ? nullptr : idxs_minor + offset; + T *vec_vals = (vals == nullptr) ? nullptr : vals + offset; + sint_t *vec_perm = perm_works.data() + tid * dim_major; + sint_t *vec_pivs = pivot_works.data() + tid * vec_nnz; _considerate_fisher_yates( - state_work, vec_nnz, dim_major, - idxs_major, vec_work.data(), pivots.data(), vals + vec_state, vec_nnz, dim_major, vec_major, vec_perm, vec_pivs, vec_vals ); - ctr.incr(vec_nnz); - idxs_major += vec_nnz; - if (idxs_minor != nullptr) { - std::fill(idxs_minor, idxs_minor + vec_nnz, i); - idxs_minor += vec_nnz; + if (vec_minor != nullptr) { + std::fill(vec_minor, vec_minor + vec_nnz, static_cast(i)); } - if (vals != nullptr) { - vals += vec_nnz; + if (apply_sort) { + sort_major_axis_vector(vec_major, vec_vals, vec_nnz); + } + }; + + #pragma omp parallel num_threads(num_threads) if(num_threads > 1) + { + const int tid = randblas_get_thread_num(); + sint_t *perm = perm_works.data() + tid * dim_major; + std::iota(perm, perm + dim_major, sint_t{0}); + #pragma omp for schedule(static) + for (int64_t i = 0; i < dim_minor; ++i) { + sample_lane(i, tid); } } - return state_t {ctr, key}; + return out; } inline double isometry_scale(Axis major_axis, int64_t vec_nnz, int64_t dim_major, int64_t dim_minor) { @@ -241,6 +374,8 @@ struct SparseDist { /// /// This constructor will raise an error if \math{\min\\{\ttt{n_rows}, \ttt{n_cols}\\} \leq 0} or if /// \math{\vecnnz} does not respect the bounds documented for the \math{\vecnnz} member. + /// It raises an overflow error if \math{\ttt{full_nnz}} cannot be represented by + /// \math{\ttt{int64_t}.} SparseDist( int64_t n_rows, int64_t n_cols, @@ -249,9 +384,9 @@ struct SparseDist { ) : n_rows(n_rows), n_cols(n_cols), major_axis(major_axis), dim_major((major_axis == Axis::Short) ? std::min(n_rows, n_cols) : std::max(n_rows, n_cols)), - dim_minor(n_rows + n_cols - dim_major), + dim_minor((major_axis == Axis::Short) ? std::max(n_rows, n_cols) : std::min(n_rows, n_cols)), isometry_scale(sparse::isometry_scale(major_axis, vec_nnz, dim_major, dim_minor)), - vec_nnz(vec_nnz), full_nnz(vec_nnz * dim_minor) + vec_nnz(vec_nnz), full_nnz(safe_int_product(vec_nnz, dim_minor)) { // argument validation randblas_require(n_rows > 0); randblas_require(n_cols > 0); @@ -287,7 +422,12 @@ struct SparseDist { /// 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.} -/// +/// +/// When RandBLAS is built with OpenMP, sampling is automatically parallelized over the +/// \math{r} blocks. The counter range for block \math{i} depends only on \math{i,} so the +/// samples and returned RNGState do not depend on the number of OpenMP threads or their +/// scheduling. +/// /// The returned RNGState should /// be used for the next call to a random sampling function whose output should be statistically /// independent from \math{\ttt{samples}.} @@ -296,7 +436,7 @@ template > inline state_t repeated_fisher_yates( int64_t k, int64_t n, int64_t r, sint_t *samples, const state_t &state ) { - return sparse::repeated_fisher_yates(state, k, n, r, samples, (sint_t*) nullptr, (double*) nullptr); + return sparse::repeated_fisher_yates(state, k, n, r, samples, (sint_t*) nullptr, (double*) nullptr, false); } template @@ -304,10 +444,7 @@ RNGState compute_next_state(SparseDist dist, RNGState state) { // Both _considerate_fisher_yates (SASO with vec_nnz > 1) and // sample_indices_iid_uniform (SASO with vec_nnz == 1, and LASO) consume // exactly one CBRNG counter increment per nonzero. - int64_t num_major_axis_vec = (dist.major_axis == Axis::Short) - ? std::max(dist.n_rows, dist.n_cols) - : std::min(dist.n_rows, dist.n_cols); - state.counter.incr(num_major_axis_vec * dist.vec_nnz); + state.counter.incr(dist.full_nnz); return state; } @@ -532,6 +669,12 @@ void laso_merge_long_axis_vector_coo_data( /// defined by \math{(\D,\ttt{seed_state}).} The submatrix is sampled directly, without /// materializing the full operator, and is returned in COO format. /// +/// The COO entries are ordered by increasing major-axis-vector index and then by +/// increasing major coordinate within each vector. When RandBLAS is built with OpenMP, +/// sampling is automatically parallelized over these vectors. The counter range for a +/// vector depends only on its logical index, so the sparse representation and returned +/// RNGState do not depend on the number of OpenMP threads or their scheduling. +/// /// 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. @@ -573,8 +716,7 @@ state_t fill_sparse_unpacked( 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); + validate_submat_dims(D.n_rows, D.n_cols, n_rows_sub, n_cols_sub, ro_s, 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 @@ -615,17 +757,20 @@ state_t fill_sparse_unpacked( // 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. + const int64_t lane_cap = safe_int_product(vec_nnz, num_major_sub); if (vals == nullptr || rows == nullptr || cols == nullptr) { - nnz = vec_nnz * num_major_sub; + nnz = lane_cap; return seed_state; } + randblas_require(seed_state.len_c >= 4); // 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); + const int64_t counter_skip = safe_int_product(num_major_off, vec_nnz); + work_state.counter.incr(counter_skip); // 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 @@ -634,70 +779,92 @@ 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 - // minor coordinates (0..num_major_sub-1); "total" is the pre-filter nnz. - int64_t total; + // Phase 1: sample each requested major-axis vector into a fixed-width output lane. + // Fixed lanes let physical threads work independently while logical vector indices + // determine counter ranges and output positions. + std::vector lane_counts; state_t end_state; if (D.major_axis == Axis::Short) { end_state = sparse::repeated_fisher_yates( - work_state, vec_nnz, dim_major, num_major_sub, idxs_major, idxs_minor, vals + work_state, vec_nnz, dim_major, num_major_sub, idxs_major, idxs_minor, vals, true ); - 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); + if (dim_major_off == 0 && dim_major_sub == dim_major) { + nnz = lane_cap; + return end_state; } } 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. - std::unordered_map loc2count{}; - 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); - // 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; + lane_counts.assign(num_major_sub, 0); + const int num_threads = sparse::sparse_sampling_thread_count( + dim_major, num_major_sub, vec_nnz, false + ); + std::vector> count_works(num_threads); + std::vector> scale_works(num_threads); + for (int tid = 0; tid < num_threads; ++tid) { + count_works[tid].reserve(vec_nnz); + scale_works[tid].reserve(vec_nnz); } + std::exception_ptr sample_error; + const auto base_ctr = work_state.counter; + #pragma omp parallel num_threads(num_threads) if(num_threads > 1) + { + const int tid = randblas_get_thread_num(); + auto &loc2count = count_works[tid]; + auto &loc2scale = scale_works[tid]; + + #pragma omp for schedule(static) + for (int64_t i = 0; i < num_major_sub; ++i) { + try { + const int64_t lane_offset = i * vec_nnz; + auto vec_ctr = base_ctr; + vec_ctr.incr(lane_offset); + state_t vec_state{vec_ctr, work_state.key}; + sint_t *vec_major = idxs_major + lane_offset; + sint_t *vec_minor = idxs_minor + lane_offset; + T *vec_vals = vals + lane_offset; + + sample_indices_iid_uniform( + dim_major, vec_nnz, vec_major, vec_vals, vec_state + ); + laso_merge_long_axis_vector_coo_data( + vec_nnz, vec_vals, vec_major, vec_minor, i, + loc2count, loc2scale + ); + const int64_t survivors = static_cast(loc2count.size()); + sparse::sort_major_axis_vector(vec_major, vec_vals, survivors); + lane_counts[i] = survivors; + } catch (...) { + #pragma omp critical(RandBLAS_laso_sampling_exception) + { + if (sample_error == nullptr) { + sample_error = std::current_exception(); + } + } + } + } + } + if (sample_error != nullptr) { + std::rethrow_exception(sample_error); + } + end_state = work_state; + end_state.counter.incr(lane_cap); } - // 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. + // Phase 2: pack lanes in increasing logical-vector order and keep only nonzeros in + // the requested major-coordinate window. Every destination precedes or equals its + // source, so this serial pass cannot overwrite an unread lane. 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++; + for (int64_t i = 0; i < num_major_sub; ++i) { + const int64_t lane_offset = i * vec_nnz; + const int64_t lane_count = (D.major_axis == Axis::Short) ? vec_nnz : lane_counts[i]; + for (int64_t j = 0; j < lane_count; ++j) { + const int64_t read = lane_offset + j; + const sint_t local_major = idxs_major[read] - static_cast(dim_major_off); + if (0 <= local_major && local_major < static_cast(dim_major_sub)) { + idxs_major[nnz] = local_major; + idxs_minor[nnz] = static_cast(i); + vals[nnz] = vals[read]; + ++nnz; + } } } return end_state; @@ -736,6 +903,10 @@ state_t fill_sparse_unpacked_nosub( /// If all reference members are are non-null, then we'll assume each of them has length /// at least \math{\ttt{S.dist.full_nnz}.} We'll proceed to populate those members /// (and \math{\ttt{S.nnz}}) with the data for the explicit representation of \math{\ttt{S}.} +/// When RandBLAS is built with OpenMP, sampling is automatically parallelized over the +/// operator's major-axis vectors. Each vector receives a counter range determined only by +/// its logical index, so the explicit representation is independent of the number of +/// OpenMP threads and their scheduling. /// On exit, \math{\ttt{S}} can be equivalently represented by /// @verbatim embed:rst:leading-slashes /// .. code:: c++ @@ -844,8 +1015,7 @@ template 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); + validate_submat_dims(S.n_rows, S.n_cols, n_rows_sub, n_cols_sub, ro_s, co_s); const SparseDist &D = S.dist; // Ask fill_sparse_unpacked (via its workspace-query mode) how large the buffers must @@ -857,20 +1027,19 @@ COOMatrix submatrix_as_coo( (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 - // (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]; + // Attach each worst-case buffer to an owning COOMatrix as soon as it is allocated, + // so a later allocation or sampling exception cannot leak an earlier buffer. 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. + COOMatrix A(n_rows_sub, n_cols_sub); // own_memory == true, null arrays. + A.vals = new T[cap]; + A.rows = new sint_t[cap]; + A.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); + fill_sparse_unpacked( + D, n_rows_sub, n_cols_sub, ro_s, co_s, nnz, A.vals, A.rows, A.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; // fill_sparse_unpacked emits each major-axis vector in sorted order, so the sampled // submatrix is CSR- or CSC-sorted; label it as such. diff --git a/RandBLAS/testing/DevNotes.md b/RandBLAS/testing/DevNotes.md index 6861caa0..d7d06c56 100644 --- a/RandBLAS/testing/DevNotes.md +++ b/RandBLAS/testing/DevNotes.md @@ -2,6 +2,12 @@ **None of the files in this directory are part of RandBLAS' public API.** +benchmarking.hh. + + Small helpers shared by performance benchmarks, including OpenMP settings + and command-line thread-list parsing. + The OpenMP helpers become serial no-ops when RandBLAS is built without OpenMP. + comparison.hh. This currently holds a single utility function for testing approximate-equality of floating point numbers. @@ -18,6 +24,13 @@ linops.hh. This file also defines functions reference_left_apply and reference_right_apply, which compute an expected answer and a componentwise error tolerance of a given matrix-matrix product. +samplers.hh. + + Baseline samplers shared by tests and benchmarks. + Everything here is discrete uniform in some sense: uniform random bits, uniform indices, + uniform subsets, or Rademacher signs. + Future baseline samplers of this kind should go in this file. + sparse_data.hh. Functions for generating (random) sparse matrices with various structures and formats. diff --git a/RandBLAS/testing/benchmarking.hh b/RandBLAS/testing/benchmarking.hh new file mode 100644 index 00000000..90ce3bbc --- /dev/null +++ b/RandBLAS/testing/benchmarking.hh @@ -0,0 +1,137 @@ +// Copyright, 2026. See LICENSE for copyright holder information. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// (1) Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// +// (2) Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// (3) Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +// POSSIBILITY OF SUCH DAMAGE. + +#pragma once + +#include "RandBLAS/config.h" + +#if defined(RandBLAS_HAS_OpenMP) +#include +#endif + +#include +#include +#include +#include +#include +#include + +namespace RandBLAS::testing { + +// A parsed list and the result of applying the benchmark's positivity rule. +struct ThreadCountsParseResult { + std::vector thread_counts; + bool valid; +}; + +// Parse and validate a comma-separated thread-count list. Empty fields are +// skipped, and an empty result is replaced with default_thread_counts. +inline ThreadCountsParseResult parse_thread_counts(const std::string &csv, const std::vector &default_thread_counts) { + std::vector thread_counts; + std::stringstream stream(csv); + std::string token; + while (std::getline(stream, token, ',')) { + if (!token.empty()) { + thread_counts.push_back(std::atoi(token.c_str())); + } + } + if (thread_counts.empty()) { + thread_counts = default_thread_counts; + } + const bool valid = std::all_of( + thread_counts.begin(), thread_counts.end(), + [](int thread_count) { return thread_count > 0; } + ); + return {std::move(thread_counts), valid}; +} + +// Return the maximum OpenMP thread count, or one in a serial build. +inline int current_threads() { +#if defined(RandBLAS_HAS_OpenMP) + return omp_get_max_threads(); +#else + return 1; +#endif +} + +// Disable dynamic teams and set the maximum OpenMP thread count. This is a +// no-op in a serial build. +inline void set_threads(int thread_count) { +#if defined(RandBLAS_HAS_OpenMP) + omp_set_dynamic(0); + omp_set_num_threads(thread_count); +#else + (void) thread_count; +#endif +} + +// Return the team size OpenMP actually provides for the requested count, or +// one in a serial build. +inline int effective_threads(int requested_threads) { +#if defined(RandBLAS_HAS_OpenMP) + int actual_threads = 1; + #pragma omp parallel num_threads(requested_threads) + { + #pragma omp single + { + actual_threads = omp_get_num_threads(); + } + } + return actual_threads; +#else + (void) requested_threads; + return 1; +#endif +} + +// Restore the OpenMP dynamic-team setting and maximum thread count when a +// benchmark scope exits. +class OpenMPSettingsGuard { +public: + OpenMPSettingsGuard() { +#if defined(RandBLAS_HAS_OpenMP) + dynamic_ = omp_get_dynamic(); + threads_ = omp_get_max_threads(); +#endif + } + + ~OpenMPSettingsGuard() { +#if defined(RandBLAS_HAS_OpenMP) + omp_set_num_threads(threads_); + omp_set_dynamic(dynamic_); +#endif + } + +private: +#if defined(RandBLAS_HAS_OpenMP) + int dynamic_; + int threads_; +#endif +}; + +} // namespace RandBLAS::testing diff --git a/RandBLAS/testing/samplers.hh b/RandBLAS/testing/samplers.hh new file mode 100644 index 00000000..779f1731 --- /dev/null +++ b/RandBLAS/testing/samplers.hh @@ -0,0 +1,218 @@ +// Copyright, 2026. See LICENSE for copyright holder information. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// (1) Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// +// (2) Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// (3) Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +// POSSIBILITY OF SUCH DAMAGE. + +#pragma once + +#include "RandBLAS/base.hh" +#include "RandBLAS/util.hh" + +#include +#include +#include +#include +#include +#include +#include + +namespace RandBLAS::testing { + +// Adapt RandBLAS' default Philox generator to the UniformRandomBitGenerator +// interface expected by the C++ standard library. +class PhiloxURBG { +public: + using result_type = uint64_t; + + // Initialize the adapter with a RandBLAS seed and counter zero. By default, + // each call consumes a new counter, matching RandBLAS' sampling kernels. + explicit PhiloxURBG(uint64_t seed, bool one_result_per_counter = true) + : state_(seed), + one_result_per_counter_(one_result_per_counter), + random_values_{}, + use_second_result_(false) {} + + // Return the smallest value produced by the adapter. + static constexpr result_type min() { + return 0; + } + + // Return the largest value produced by the adapter. + static constexpr result_type max() { + return std::numeric_limits::max(); + } + + // Draw one 64-bit value, optionally using both 64-bit pairs from each + // Philox4x32 result before advancing to the next result. + result_type operator()() { + if (use_second_result_) { + use_second_result_ = false; + return RandBLAS::promote_uint_pair(random_values_[2], random_values_[3]); + } + typename RNGState<>::generator generator; + random_values_ = generator(state_.counter, state_.key); + state_.counter.incr(); + use_second_result_ = !one_result_per_counter_; + return RandBLAS::promote_uint_pair(random_values_[0], random_values_[1]); + } + +private: + RNGState<> state_; + bool one_result_per_counter_; + typename RNGState<>::ctr_type random_values_; + bool use_second_result_; +}; + +// Use std::sample over an iota-filled vector to draw vec_nnz distinct indices +// from [0, n) for each requested vector. This scans all n candidates for every +// vector and uses O(n) workspace. +template +void sample_std_sample( + int64_t n, int64_t num_vectors, int64_t vec_nnz, int64_t *samples, RNG &rng +) { + std::vector population(n); + std::iota(population.begin(), population.end(), int64_t{0}); + for (int64_t vector = 0; vector < num_vectors; ++vector) { + std::sample( + population.begin(), population.end(), + samples + vector * vec_nnz, vec_nnz, rng + ); + } +} + +// Use the first vec_nnz steps of Fisher-Yates to draw distinct indices from +// [0, n). The identity permutation is rebuilt for every requested vector. +template +void sample_partial_fisher_yates( + int64_t n, int64_t num_vectors, int64_t vec_nnz, int64_t *samples, RNG &rng +) { + std::vector population(n); + for (int64_t vector = 0; vector < num_vectors; ++vector) { + std::iota(population.begin(), population.end(), int64_t{0}); + for (int64_t entry = 0; entry < vec_nnz; ++entry) { + std::uniform_int_distribution pick(entry, n - 1); + int64_t pivot = pick(rng); + std::swap(population[entry], population[pivot]); + samples[vector * vec_nnz + entry] = population[entry]; + } + } +} + +// Shuffle a permutation of [0, n) and copy its first vec_nnz entries for each +// requested vector. The method uses O(n) work per vector and O(n) workspace. +template +void sample_full_shuffle( + int64_t n, int64_t num_vectors, int64_t vec_nnz, int64_t *samples, RNG &rng +) { + std::vector population(n); + std::iota(population.begin(), population.end(), int64_t{0}); + for (int64_t vector = 0; vector < num_vectors; ++vector) { + std::shuffle(population.begin(), population.end(), rng); + std::copy_n(population.begin(), vec_nnz, samples + vector * vec_nnz); + } +} + +// Draw indices uniformly from [0, n), rejecting duplicates found by a linear +// scan. The expected work is O(vec_nnz^2) per vector when vec_nnz is small. +template +void sample_rejection( + int64_t n, int64_t num_vectors, int64_t vec_nnz, int64_t *samples, RNG &rng +) { + std::uniform_int_distribution distribution(0, n - 1); + for (int64_t vector = 0; vector < num_vectors; ++vector) { + int64_t *vector_samples = samples + vector * vec_nnz; + for (int64_t entry = 0; entry < vec_nnz;) { + int64_t candidate = distribution(rng); + auto duplicate = std::find(vector_samples, vector_samples + entry, candidate); + if (duplicate == vector_samples + entry) { + vector_samples[entry] = candidate; + ++entry; + } + } + } +} + +// Apply Floyd's algorithm with std::unordered_set to draw vec_nnz distinct +// indices from [0, n). Expected work and workspace are O(vec_nnz). +template +void sample_floyd( + int64_t n, int64_t num_vectors, int64_t vec_nnz, int64_t *samples, RNG &rng +) { + std::unordered_set selected_values; + selected_values.reserve(vec_nnz); + for (int64_t vector = 0; vector < num_vectors; ++vector) { + selected_values.clear(); + for (int64_t entry = 0; entry < vec_nnz; ++entry) { + int64_t upper_bound = n - vec_nnz + entry; + std::uniform_int_distribution pick(0, upper_bound); + int64_t candidate = pick(rng); + bool inserted = selected_values.insert(candidate).second; + int64_t selected = inserted ? candidate : upper_bound; + if (!inserted) { + selected_values.insert(selected); + } + samples[vector * vec_nnz + entry] = selected; + } + } +} + +// Fill SASO COO data from a sampler of distinct major coordinates. The minor +// coordinates identify vectors, values are random signs, and each vector's +// major coordinates are sorted into canonical order. +template +void fill_saso_data( + int64_t n, int64_t num_vectors, int64_t vec_nnz, bool major_is_rows, + int64_t *rows, int64_t *cols, double *values, RNG &rng, Sampler sampler +) { + int64_t *major_indices = major_is_rows ? rows : cols; + int64_t *minor_indices = major_is_rows ? cols : rows; + sampler(n, num_vectors, vec_nnz, major_indices, rng); + + for (int64_t vector = 0; vector < num_vectors; ++vector) { + int64_t block_start = vector * vec_nnz; + std::fill_n(minor_indices + block_start, vec_nnz, vector); + for (int64_t entry = 0; entry < vec_nnz; ++entry) { + int64_t offset = block_start + entry; + values[offset] = (rng() & 1) == 0 ? 1.0 : -1.0; + } + + for (int64_t entry = 1; entry < vec_nnz; ++entry) { + int64_t key = major_indices[block_start + entry]; + double value = values[block_start + entry]; + int64_t cursor = entry - 1; + while (cursor >= 0 && major_indices[block_start + cursor] > key) { + major_indices[block_start + cursor + 1] = + major_indices[block_start + cursor]; + values[block_start + cursor + 1] = values[block_start + cursor]; + --cursor; + } + major_indices[block_start + cursor + 1] = key; + values[block_start + cursor + 1] = value; + } + } +} + +} // namespace RandBLAS::testing diff --git a/RandBLAS/util.hh b/RandBLAS/util.hh index 5f392225..62733f9e 100644 --- a/RandBLAS/util.hh +++ b/RandBLAS/util.hh @@ -36,6 +36,7 @@ #include #include +#include #include #include #include @@ -61,6 +62,48 @@ void safe_scal(int64_t n, T a, T* x) { } } + +// ============================================================================= +/// \fn lascl(blas::Layout layout, int64_t m, int64_t n, T alpha, T* A, int64_t lda) +/// @verbatim embed:rst:leading-slashes +/// In-place scale a dense :math:`m \times n` matrix: +/// +/// .. math:: +/// A \leftarrow \alpha \cdot A. +/// +/// Named after LAPACK's ``?lascl`` for the convention; the signature is +/// simpler than LAPACK's (no ``CFROM`` / ``CTO`` overflow protection, no +/// matrix-type flag --- general dense only). +/// +/// Fast paths: +/// +/// - :math:`\alpha = 1`: returns immediately. +/// - :math:`\alpha = 0`: ``std::fill``-based zero out. +/// - otherwise: per-column (ColMajor) or per-row (RowMajor) ``blas::scal``. +/// @endverbatim +template +void lascl(blas::Layout layout, int64_t m, int64_t n, T alpha, T* A, int64_t lda) { + if (alpha == T(1)) return; + if (alpha == T(0)) { + if (layout == blas::Layout::ColMajor) { + for (int64_t j = 0; j < n; ++j) + std::fill(A + j * lda, A + j * lda + m, T(0)); + } else { + for (int64_t i = 0; i < m; ++i) + std::fill(A + i * lda, A + i * lda + n, T(0)); + } + return; + } + if (layout == blas::Layout::ColMajor) { + for (int64_t j = 0; j < n; ++j) + blas::scal(m, alpha, A + j * lda, 1); + } else { + for (int64_t i = 0; i < m; ++i) + blas::scal(n, alpha, A + i * lda, 1); + } +} + + template void omatcopy(int64_t m, int64_t n, const T* A, int64_t irs_a, int64_t ics_a, T* B, int64_t irs_b, int64_t ics_b) { // TODO: diff --git a/STYLE_GUIDE.md b/STYLE_GUIDE.md index ef66f32d..10ab2211 100644 --- a/STYLE_GUIDE.md +++ b/STYLE_GUIDE.md @@ -35,8 +35,11 @@ Follow the declaration you are modifying. There is no fixed line-length limit. Wrap prose, signatures, and expressions when doing so makes them easier to read; do not damage a formula or compact tabular code to hit a column count. -For a long function signature, put one logical parameter on each line and the -closing parenthesis on its own line. +Do not wrap a function signature merely because it has several parameters. +Keep the parameters together when the declaration fits comfortably and remains +easy to read. Put one logical parameter on each line only when that layout is +necessary to make a long or semantically dense signature readable. In that +case, put the closing parenthesis on its own line. ### Headers and includes @@ -110,10 +113,12 @@ Use `///` comments for declarations included in the web API reference. For a function with a few arguments, explain the contract in plain prose. Do not add a field for every parameter merely because Doxygen supports one. -For a long BLAS-like interface, put structured reStructuredText inside -`@verbatim embed:rst:leading-slashes`. -The established parameter form is `name - [direction]`, followed by indented -bullets: +Use a parameter dropdown only when an interface has many parameters whose +precise meanings have complicated relationships with one another. Put the +structured reStructuredText inside +`@verbatim embed:rst:leading-slashes`. Within such a dropdown, the established +parameter form is `name - [direction]`, followed by indented bullets. State +entry and exit behavior when the direction matters: ```cpp // ============================================================================= @@ -135,6 +140,38 @@ void mathfunc(int a, int &b) { } ``` +List a natural group of related parameter names together when one description +captures their shared role: + +```cpp +// ============================================================================= +/// Sample a matrix window into caller-owned storage. +/// +/// @verbatim embed:rst:leading-slashes +/// .. dropdown:: Full parameter descriptions +/// :animate: fade-in-slide-down +/// +/// n_rows, n_cols +/// * The dimensions of the window. +/// +/// row_offset, col_offset +/// * The position of the window in the full matrix. +/// +/// nnz +/// * On exit: the number of entries written to ``values``. +/// +/// values +/// * A caller-owned buffer with enough capacity for the requested window. +/// @endverbatim +void sample_window( + int64_t n_rows, int64_t n_cols, + int64_t row_offset, int64_t col_offset, + int64_t &nnz, double *values +) { + // ... +} +``` + Do not put a blank documentation line between `@endverbatim` and the declaration. Do not use `@tparam`, `@param`, `@return`, or `@returns` for new web-facing diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 5328b13a..f4d32fd2 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -130,6 +130,16 @@ target_link_libraries( sketch_general_performance PUBLIC RandBLAS blaspp lapackpp ) +add_executable( + saso_sampling_performance simple-kernel-benchmarks/saso_sampling_performance.cc +) +target_include_directories( + saso_sampling_performance PUBLIC ${Random123_DIR} +) +target_link_libraries( + saso_sampling_performance PUBLIC RandBLAS blaspp lapackpp +) + foreach(example_target IN ITEMS tls_dense_skop tls_sparse_skop @@ -138,6 +148,7 @@ foreach(example_target IN ITEMS slra_qrcp spmm_performance sketch_general_performance + saso_sampling_performance ) randblas_stage_runtime_dlls(${example_target}) endforeach() diff --git a/examples/simple-kernel-benchmarks/saso_sampling_performance.cc b/examples/simple-kernel-benchmarks/saso_sampling_performance.cc new file mode 100644 index 00000000..71ea6cd9 --- /dev/null +++ b/examples/simple-kernel-benchmarks/saso_sampling_performance.cc @@ -0,0 +1,767 @@ +// Copyright, 2026. See LICENSE for copyright holder information. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// (1) Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// +// (2) Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// (3) Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +// POSSIBILITY OF SUCH DAMAGE. + +// ============================================================================ +// SASO SAMPLING PERFORMANCE BENCHMARK +// ============================================================================ +// +// This benchmark compares RandBLAS::repeated_fisher_yates against plausible +// C++ implementations for sampling the support of a SASO. If dim_major = n, +// num_major_axis_vectors = r, and each vector has vec_nnz = k nonzeros, every +// method produces r independent k-subsets of {0, ..., n - 1}. +// +// SUPPORT-ONLY METHODS: +// +// * std::sample over an iota-filled vector +// * partial Fisher-Yates with a fresh length-n iota per vector +// * full std::shuffle followed by the first k entries +// * uniform draws with linear duplicate rejection +// * Floyd's algorithm with std::unordered_set +// * RandBLAS::repeated_fisher_yates +// +// The first table is a natural-library comparison: the five alternatives use +// std::mt19937_64 and RandBLAS uses its native Philox path. The second table is +// a controlled-engine comparison: a URBG adapter exposes the same Philox stream +// to the C++ alternatives. The controlled table equalizes the generator, but +// not integer range mapping. RandBLAS currently uses a 64-bit value modulo the +// active range; standard algorithms and std::uniform_int_distribution use the +// C++ library's range mapping. +// +// The end-to-end tables also construct COO data: they write minor coordinates, +// generate Rademacher values, and sort every major-axis vector by its major +// coordinate. Both wide (major coordinates are rows) and tall (major +// coordinates are columns) SASOs are covered when the shape is nonsquare. +// Output allocation and correctness checks are outside the timed region. +// Method-owned workspace allocation remains timed, matching the natural cost +// of calling each implementation as written. +// +// METRICS: +// +// * min and median wall time over repeated trials +// * minimum-time nanoseconds per generated nonzero +// * speedup relative to std::sample in the same table +// +// The comparison tables force RandBLAS to one OpenMP thread. Scaling mode times +// only RandBLAS because the competing implementations own one serial RNG +// engine. The k=1 RandBLAS row uses the library's specialized i.i.d.-uniform +// path rather than repeated Fisher-Yates. +// +// USAGE: +// +// ./saso_sampling_performance [flags] +// ./saso_sampling_performance [flags] n r k [trials] +// +// flags: +// --natural-only skip the controlled-Philox tables +// --support-only skip end-to-end COO construction +// --scaling report RandBLAS thread scaling only +// --threads=LIST requested thread counts (default 1,2,4,8) +// --help print usage +// +// EXAMPLES: +// +// ./saso_sampling_performance +// ./saso_sampling_performance 256 4096 8 20 +// ./saso_sampling_performance --natural-only --support-only 1024 8192 8 +// ./saso_sampling_performance --scaling --threads=1,2,4,8 2000 100000 8 10 +// +// ============================================================================ + +#include +#include "RandBLAS/config.h" +#include "RandBLAS/testing/benchmarking.hh" +#include "RandBLAS/testing/samplers.hh" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using RandBLAS::testing::current_threads; +using RandBLAS::testing::effective_threads; +using RandBLAS::testing::OpenMPSettingsGuard; +using RandBLAS::testing::parse_thread_counts; +using RandBLAS::testing::set_threads; + +// MARK: benchmark setup + +struct Config { + int64_t dim_major; + int64_t num_major_axis_vectors; + int64_t vec_nnz; +}; + +struct Record { + std::string label; + int64_t min_ns = 0; + int64_t median_ns = 0; + double ns_per_nonzero = -1.0; + double speedup_vs_std_sample = -1.0; + std::string notes; +}; + +struct ScalingRecord { + int requested_threads; + int threads; + int64_t min_ns; + int64_t median_ns; + double ns_per_nonzero; + double speedup; + double efficiency; +}; + +template +static std::pair run_trials(Func &&func, int64_t num_trials) { + std::vector times; + times.reserve(num_trials); + for (int64_t trial = 0; trial < num_trials; ++trial) { + auto start = std::chrono::steady_clock::now(); + func(); + auto end = std::chrono::steady_clock::now(); + times.push_back( + std::chrono::duration_cast(end - start).count() + ); + } + std::sort(times.begin(), times.end()); + return {times.front(), times[num_trials / 2]}; +} + +static std::string format_cell(double value, int precision) { + if (value < 0.0) { + return "-"; + } + std::ostringstream stream; + stream << std::fixed << std::setprecision(precision) << value; + return stream.str(); +} + +static void print_table_header(const std::string &title) { + std::cout << " " << title << "\n"; + std::cout << " " << std::left << std::setw(29) << "Implementation" + << std::right << std::setw(13) << "Min(ns)" + << std::setw(13) << "Median(ns)" + << std::setw(13) << "ns/nonzero" + << std::setw(12) << "vs sample" + << " notes\n"; + std::cout << " " << std::string(105, '-') << "\n"; +} + +static void print_table_record(const Record &record) { + std::cout << " " << std::left << std::setw(29) << record.label + << std::right << std::setw(13) << record.min_ns + << std::setw(13) << record.median_ns + << std::setw(13) << format_cell(record.ns_per_nonzero, 2) + << std::setw(12) << format_cell(record.speedup_vs_std_sample, 2) + << " " << record.notes << "\n"; +} + +static void fill_speedups(std::vector &records) { + if (records.empty() || records.front().min_ns <= 0) { + return; + } + double baseline = static_cast(records.front().min_ns); + for (Record &record : records) { + if (record.min_ns > 0) { + record.speedup_vs_std_sample = baseline + / static_cast(record.min_ns); + } + } +} + +// MARK: correctness checks + +static bool support_is_valid(const Config &config, const std::vector &samples) { + if (samples.size() != static_cast( + config.num_major_axis_vectors * config.vec_nnz + )) { + return false; + } + std::vector seen(config.dim_major, false); + for (int64_t vector = 0; vector < config.num_major_axis_vectors; ++vector) { + std::fill(seen.begin(), seen.end(), false); + for (int64_t entry = 0; entry < config.vec_nnz; ++entry) { + int64_t index = samples[vector * config.vec_nnz + entry]; + if (index < 0 || index >= config.dim_major || seen[index]) { + return false; + } + seen[index] = true; + } + } + return true; +} + +static bool saso_data_is_valid( + const Config &config, bool major_is_rows, + const std::vector &rows, const std::vector &cols, + const std::vector &values +) { + const std::vector &major = major_is_rows ? rows : cols; + const std::vector &minor = major_is_rows ? cols : rows; + if (!support_is_valid(config, major)) { + return false; + } + for (int64_t vector = 0; vector < config.num_major_axis_vectors; ++vector) { + for (int64_t entry = 0; entry < config.vec_nnz; ++entry) { + int64_t offset = vector * config.vec_nnz + entry; + if (minor[offset] != vector) { + return false; + } + if (values[offset] != -1.0 && values[offset] != 1.0) { + return false; + } + if (entry > 0 && major[offset - 1] >= major[offset]) { + return false; + } + } + } + return true; +} + +// MARK: support-only benchmarks + +template +static Record benchmark_support_method( + const std::string &label, const std::string ¬es, const Config &config, + int64_t num_trials, uint64_t seed, Sampler sampler +) { + int64_t nnz = config.num_major_axis_vectors * config.vec_nnz; + std::vector samples(nnz, -1); + RNG rng(seed); + auto sample = [&]() { + sampler( + config.dim_major, config.num_major_axis_vectors, + config.vec_nnz, samples.data(), rng + ); + }; + + sample(); + bool valid = support_is_valid(config, samples); + auto [min_ns, median_ns] = run_trials(sample, num_trials); + + Record record; + record.label = label; + record.min_ns = min_ns; + record.median_ns = median_ns; + record.ns_per_nonzero = static_cast(min_ns) / static_cast(nnz); + record.notes = valid ? notes : "FAIL: invalid support; " + notes; + return record; +} + +static Record benchmark_randblas_support( + const Config &config, int64_t num_trials, uint64_t seed +) { + int64_t nnz = config.num_major_axis_vectors * config.vec_nnz; + std::vector samples(nnz, -1); + RandBLAS::RNGState<> state(seed); + auto sample = [&]() { + state = RandBLAS::repeated_fisher_yates( + config.vec_nnz, config.dim_major, config.num_major_axis_vectors, + samples.data(), state + ); + }; + + sample(); + bool valid = support_is_valid(config, samples); + auto [min_ns, median_ns] = run_trials(sample, num_trials); + + Record record; + record.label = "RandBLAS repeated FY"; + record.min_ns = min_ns; + record.median_ns = median_ns; + record.ns_per_nonzero = static_cast(min_ns) / static_cast(nnz); + if (config.vec_nnz == 1) { + record.notes = "O(r), specialized i.i.d. path"; + } else { + record.notes = "O(n + r*k), restore k swaps"; + } + if (!valid) { + record.notes = "FAIL: invalid support; " + record.notes; + } + return record; +} + +template +static std::vector support_records( + const Config &config, int64_t num_trials, uint64_t seed +) { + using namespace RandBLAS::testing; + std::vector records; + records.push_back(benchmark_support_method( + "std::sample(iota vector)", "O(r*n), O(n) workspace", + config, num_trials, seed, + [](int64_t n, int64_t r, int64_t k, int64_t *output, auto &rng) { + sample_std_sample(n, r, k, output, rng); + } + )); + records.push_back(benchmark_support_method( + "partial FY + iota reset", "O(r*n), reset dominates for k << n", + config, num_trials, seed, + [](int64_t n, int64_t r, int64_t k, int64_t *output, auto &rng) { + sample_partial_fisher_yates(n, r, k, output, rng); + } + )); + records.push_back(benchmark_support_method( + "full std::shuffle", "O(r*n), take first k", config, num_trials, seed, + [](int64_t n, int64_t r, int64_t k, int64_t *output, auto &rng) { + sample_full_shuffle(n, r, k, output, rng); + } + )); + records.push_back(benchmark_support_method( + "draw/reject + linear find", "O(r*k^2) expected when k << n", + config, num_trials, seed, + [](int64_t n, int64_t r, int64_t k, int64_t *output, auto &rng) { + sample_rejection(n, r, k, output, rng); + } + )); + records.push_back(benchmark_support_method( + "Floyd + unordered_set", "O(r*k) expected", + config, num_trials, seed, + [](int64_t n, int64_t r, int64_t k, int64_t *output, auto &rng) { + sample_floyd(n, r, k, output, rng); + } + )); + records.push_back(benchmark_randblas_support(config, num_trials, seed)); + fill_speedups(records); + return records; +} + +// MARK: end-to-end COO benchmarks + +template +static Record benchmark_saso_data_method( + const std::string &label, const std::string ¬es, const Config &config, + bool major_is_rows, int64_t num_trials, uint64_t seed, Sampler sampler +) { + int64_t nnz = config.num_major_axis_vectors * config.vec_nnz; + std::vector rows(nnz, -1); + std::vector cols(nnz, -1); + std::vector values(nnz, 0.0); + RNG rng(seed); + auto sample = [&]() { + RandBLAS::testing::fill_saso_data( + config.dim_major, config.num_major_axis_vectors, config.vec_nnz, + major_is_rows, rows.data(), cols.data(), values.data(), rng, sampler + ); + }; + + sample(); + bool valid = saso_data_is_valid(config, major_is_rows, rows, cols, values); + auto [min_ns, median_ns] = run_trials(sample, num_trials); + + Record record; + record.label = label; + record.min_ns = min_ns; + record.median_ns = median_ns; + record.ns_per_nonzero = static_cast(min_ns) / static_cast(nnz); + record.notes = valid ? notes : "FAIL: invalid COO data; " + notes; + return record; +} + +static Record benchmark_randblas_saso_data( + const Config &config, bool major_is_rows, int64_t num_trials, uint64_t seed +) { + int64_t nnz = config.num_major_axis_vectors * config.vec_nnz; + int64_t n_rows = major_is_rows + ? config.dim_major + : config.num_major_axis_vectors; + int64_t n_cols = major_is_rows + ? config.num_major_axis_vectors + : config.dim_major; + std::vector rows(nnz, -1); + std::vector cols(nnz, -1); + std::vector values(nnz, 0.0); + RandBLAS::SparseDist distribution( + n_rows, n_cols, config.vec_nnz, RandBLAS::Axis::Short + ); + RandBLAS::RNGState<> state(seed); + int64_t sampled_nnz = nnz; + auto sample = [&]() { + sampled_nnz = nnz; + state = RandBLAS::fill_sparse_unpacked( + distribution, n_rows, n_cols, 0, 0, sampled_nnz, + values.data(), rows.data(), cols.data(), state + ); + }; + + sample(); + bool valid = sampled_nnz == nnz + && saso_data_is_valid(config, major_is_rows, rows, cols, values); + auto [min_ns, median_ns] = run_trials(sample, num_trials); + + Record record; + record.label = "RandBLAS fill unpacked"; + record.min_ns = min_ns; + record.median_ns = median_ns; + record.ns_per_nonzero = static_cast(min_ns) / static_cast(nnz); + record.notes = config.vec_nnz == 1 + ? "fused support/sign; i.i.d. path; sorted" + : "fused support/sign; restore swaps; sorted"; + if (!valid) { + record.notes = "FAIL: invalid COO data; " + record.notes; + } + return record; +} + +template +static std::vector saso_data_records( + const Config &config, bool major_is_rows, int64_t num_trials, uint64_t seed +) { + using namespace RandBLAS::testing; + std::vector records; + records.push_back(benchmark_saso_data_method( + "std::sample(iota vector)", "support + minor + sign + sort", + config, major_is_rows, num_trials, seed, + [](int64_t n, int64_t r, int64_t k, int64_t *output, auto &rng) { + sample_std_sample(n, r, k, output, rng); + } + )); + records.push_back(benchmark_saso_data_method( + "partial FY + iota reset", "support + minor + sign + sort", + config, major_is_rows, num_trials, seed, + [](int64_t n, int64_t r, int64_t k, int64_t *output, auto &rng) { + sample_partial_fisher_yates(n, r, k, output, rng); + } + )); + records.push_back(benchmark_saso_data_method( + "full std::shuffle", "support + minor + sign + sort", + config, major_is_rows, num_trials, seed, + [](int64_t n, int64_t r, int64_t k, int64_t *output, auto &rng) { + sample_full_shuffle(n, r, k, output, rng); + } + )); + records.push_back(benchmark_saso_data_method( + "draw/reject + linear find", "support + minor + sign + sort", + config, major_is_rows, num_trials, seed, + [](int64_t n, int64_t r, int64_t k, int64_t *output, auto &rng) { + sample_rejection(n, r, k, output, rng); + } + )); + records.push_back(benchmark_saso_data_method( + "Floyd + unordered_set", "support + minor + sign + sort", + config, major_is_rows, num_trials, seed, + [](int64_t n, int64_t r, int64_t k, int64_t *output, auto &rng) { + sample_floyd(n, r, k, output, rng); + } + )); + records.push_back(benchmark_randblas_saso_data( + config, major_is_rows, num_trials, seed + )); + fill_speedups(records); + return records; +} + +static void print_records( + const std::string &title, const std::vector &records +) { + print_table_header(title); + for (const Record &record : records) { + print_table_record(record); + } + std::cout << "\n"; +} + +// MARK: thread scaling + +static bool run_scaling( + const Config &config, int64_t num_trials, const std::vector &thread_counts +) { + constexpr uint64_t seed = 12345; + const int64_t nnz = config.num_major_axis_vectors * config.vec_nnz; + std::vector samples(nnz, -1); + std::vector expected_samples; + RandBLAS::RNGState<> expected_state(seed); + std::vector records; + records.reserve(thread_counts.size()); + bool exact = true; + int64_t baseline_ns = 0; + int baseline_threads = 1; + + for (int thread_count : thread_counts) { + set_threads(thread_count); + const int policy_threads = RandBLAS::sparse::sparse_sampling_thread_count( + config.dim_major, config.num_major_axis_vectors, + config.vec_nnz, config.vec_nnz > 1 + ); + const int actual_threads = effective_threads(policy_threads); + RandBLAS::RNGState<> state(seed); + auto end_state = RandBLAS::repeated_fisher_yates( + config.vec_nnz, config.dim_major, config.num_major_axis_vectors, + samples.data(), state + ); + exact = exact && support_is_valid(config, samples); + if (records.empty()) { + expected_samples = samples; + expected_state = end_state; + } else { + exact = exact + && samples == expected_samples + && end_state == expected_state; + } + + auto [min_ns, median_ns] = run_trials([&]() { + RandBLAS::RNGState<> trial_state(seed); + end_state = RandBLAS::repeated_fisher_yates( + config.vec_nnz, config.dim_major, config.num_major_axis_vectors, + samples.data(), trial_state + ); + }, num_trials); + if (records.empty()) { + baseline_ns = min_ns; + baseline_threads = actual_threads; + } + const double speedup = min_ns > 0 + ? static_cast(baseline_ns) / static_cast(min_ns) + : -1.0; + const double relative_threads = static_cast(actual_threads) + / static_cast(baseline_threads); + records.push_back({ + thread_count, actual_threads, min_ns, median_ns, + static_cast(min_ns) / static_cast(nnz), + speedup, speedup / relative_threads + }); + } + + std::cout << "=== RANDBLAS THREAD SCALING: n=" << config.dim_major + << " r=" << config.num_major_axis_vectors + << " k=" << config.vec_nnz + << ", 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::right << std::setw(7) << "Request" + << std::setw(8) << "Threads" + << std::setw(13) << "Min(ns)" + << std::setw(13) << "Median(ns)" + << std::setw(13) << "ns/nonzero" + << std::setw(11) << "Spd(min)" + << std::setw(11) << "Eff(min)" << "\n" + << " " << std::string(74, '-') << "\n"; + for (const ScalingRecord &record : records) { + std::cout << " " << std::right << std::setw(7) << record.requested_threads + << std::setw(8) << record.threads + << std::setw(13) << record.min_ns + << std::setw(13) << record.median_ns + << std::setw(13) << format_cell(record.ns_per_nonzero, 2) + << std::setw(11) << format_cell(record.speedup, 2) + << std::setw(11) << format_cell(record.efficiency, 2) << "\n"; + } + std::cout << "\n Exact output/state check: " + << (exact ? "PASS" : "FAIL") << "\n\n"; + return exact; +} + +static void run_support_tables(const Config &config, int64_t num_trials, bool include_controlled) { + constexpr uint64_t seed = 12345; + std::cout << "=== SUPPORT ONLY: n=" << config.dim_major + << " r=" << config.num_major_axis_vectors + << " k=" << config.vec_nnz + << ", trials=" << num_trials << " ===\n\n"; + print_records( + "natural C++ RNGs: std::mt19937_64 versus native RandBLAS Philox", + support_records(config, num_trials, seed) + ); + if (include_controlled) { + print_records( + "controlled engine: Philox for every implementation", + support_records(config, num_trials, seed) + ); + } +} + +static void run_saso_data_tables( + const Config &config, bool major_is_rows, int64_t num_trials, bool include_controlled +) { + constexpr uint64_t seed = 67890; + const char *shape = major_is_rows + ? "wide SASO (major coordinates are rows)" + : "tall SASO (major coordinates are columns)"; + std::cout << "=== END-TO-END COO: " << shape + << ", n=" << config.dim_major + << " r=" << config.num_major_axis_vectors + << " k=" << config.vec_nnz + << ", trials=" << num_trials << " ===\n\n"; + print_records( + "natural C++ RNGs: support, minor coordinates, signs, and sorting", + saso_data_records( + config, major_is_rows, num_trials, seed + ) + ); + if (include_controlled) { + print_records( + "controlled engine: support, minor coordinates, signs, and sorting", + saso_data_records( + config, major_is_rows, num_trials, seed + ) + ); + } +} + +// MARK: command-line interface + +static bool config_is_valid(const Config &config, int64_t num_trials) { + return config.dim_major > 0 + && config.num_major_axis_vectors >= config.dim_major + && config.vec_nnz > 0 + && config.vec_nnz <= config.dim_major + && num_trials > 0; +} + +static void print_usage(const char *program) { + std::cout << "Usage:\n" + << " " << program << " [flags]\n" + << " " << program << " [flags] n r k [trials]\n\n" + << "Constraints: 0 < k <= n <= r and trials > 0.\n" + << "Flags: --natural-only, --support-only, --scaling, " + << "--threads=1,2,4,8, --help\n"; +} + +int main(int argc, char **argv) { + OpenMPSettingsGuard openmp_settings; + bool include_controlled = true; + bool support_only = false; + bool scaling = false; + const std::vector default_thread_counts{1, 2, 4, 8}; + auto thread_config = parse_thread_counts("", default_thread_counts); + std::vector positional; + for (int arg = 1; arg < argc; ++arg) { + std::string value = argv[arg]; + if (value == "--natural-only") { + include_controlled = false; + } else if (value == "--support-only") { + support_only = true; + } else if (value == "--scaling") { + scaling = true; + } else if (value.rfind("--threads=", 0) == 0) { + thread_config = parse_thread_counts(value.substr(10), default_thread_counts); + } else if (value == "--help") { + print_usage(argv[0]); + return 0; + } else if (value.rfind("--", 0) == 0) { + std::cerr << "Unknown flag: " << value << "\n"; + print_usage(argv[0]); + return 1; + } else { + positional.push_back(value); + } + } + + if (!positional.empty() && positional.size() != 3 && positional.size() != 4) { + print_usage(argv[0]); + return 1; + } + if (!thread_config.valid) { + std::cerr << "Invalid thread list. Expected positive integers.\n"; + return 1; + } + if (!scaling) { + set_threads(1); + } + + std::cout << "\n============================================================\n" + << "SASO SAMPLING PERFORMANCE BENCHMARK\n" + << "============================================================\n"; + if (scaling) { + std::cout << "RandBLAS-only thread scaling; output allocation is not timed.\n" + << "Configured OpenMP maximum: " << current_threads() << "\n\n"; + } else { + std::cout << "Comparison mode; allocations for output arrays are not timed.\n" + << "RandBLAS is forced to one OpenMP thread.\n" + << "Speedup is relative to std::sample in the same table.\n\n"; + } + + if (!positional.empty()) { + Config config{ + std::atoll(positional[0].c_str()), std::atoll(positional[1].c_str()), + std::atoll(positional[2].c_str()) + }; + int64_t num_trials = positional.size() == 4 + ? std::atoll(positional[3].c_str()) + : 10; + if (!config_is_valid(config, num_trials)) { + std::cerr << "Invalid configuration. Expected 0 < k <= n <= r " + << "and trials > 0.\n"; + return 1; + } + if (scaling) { + return run_scaling(config, num_trials, thread_config.thread_counts) ? 0 : 2; + } + run_support_tables(config, num_trials, include_controlled); + if (!support_only) { + run_saso_data_tables( + config, true, num_trials, include_controlled + ); + if (config.num_major_axis_vectors > config.dim_major) { + run_saso_data_tables( + config, false, num_trials, include_controlled + ); + } + } + return 0; + } + + constexpr int64_t num_trials = 5; + std::vector support_configs{ + {64, 4096, 8}, + {256, 4096, 1}, + {256, 4096, 4}, + {256, 4096, 16}, + {1024, 4096, 8}, + {1024, 4096, 64}, + {4096, 4096, 8}, + }; + if (scaling) { + bool exact = true; + for (const Config &config : support_configs) { + exact = run_scaling(config, num_trials, thread_config.thread_counts) && exact; + } + return exact ? 0 : 2; + } + for (const Config &config : support_configs) { + run_support_tables(config, num_trials, include_controlled); + } + + if (!support_only) { + Config data_config{256, 4096, 8}; + run_saso_data_tables( + data_config, true, num_trials, include_controlled + ); + run_saso_data_tables( + data_config, false, num_trials, include_controlled + ); + } + return 0; +} diff --git a/examples/simple-kernel-benchmarks/sketch_general_performance.cc b/examples/simple-kernel-benchmarks/sketch_general_performance.cc index 4aba663c..94f58cf8 100644 --- a/examples/simple-kernel-benchmarks/sketch_general_performance.cc +++ b/examples/simple-kernel-benchmarks/sketch_general_performance.cc @@ -80,9 +80,7 @@ #include "RandBLAS/config.h" #include "RandBLAS/sparse_data/spmm_dispatch.hh" -#if defined(RandBLAS_HAS_OpenMP) -#include -#endif +#include "RandBLAS/testing/benchmarking.hh" using namespace std::chrono; using blas::Layout; @@ -90,26 +88,15 @@ using blas::Op; using RandBLAS::Axis; using RandBLAS::SparseDist; using RandBLAS::SparseSkOp; +using RandBLAS::testing::current_threads; +using RandBLAS::testing::effective_threads; +using RandBLAS::testing::OpenMPSettingsGuard; +using RandBLAS::testing::parse_thread_counts; +using RandBLAS::testing::set_threads; // 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) { @@ -209,6 +196,16 @@ struct OpSpec { Axis axis; }; +struct SamplingScalingRecord { + int requested_threads; + int threads; + long min_us; + long median_us; + double ns_per_nonzero; + double speedup; + double efficiency; +}; + // --------------------------------------------------------------------------- // 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. @@ -400,20 +397,122 @@ void run_right(int64_t d, int64_t m, int64_t n, std::cout << "\n"; } +static bool run_sampling_scaling( + const OpSpec &spec, const SparseDist &dist, + const RandBLAS::RNGState<> &seed_state, int num_trials, + const std::vector &threads +) { + const int64_t capacity = dist.full_nnz; + std::vector values(capacity); + std::vector rows(capacity); + std::vector cols(capacity); + std::vector expected_values; + std::vector expected_rows; + std::vector expected_cols; + auto expected_state = seed_state; + int64_t expected_nnz = -1; + std::vector scaling_records; + scaling_records.reserve(threads.size()); + long baseline_us = 0; + int baseline_threads = 1; + bool exact = true; + + for (int thread_count : threads) { + set_threads(thread_count); + const bool uses_permutation_workspace = + dist.major_axis == Axis::Short && dist.vec_nnz > 1; + const int policy_threads = RandBLAS::sparse::sparse_sampling_thread_count( + dist.dim_major, dist.dim_minor, dist.vec_nnz, + uses_permutation_workspace + ); + const int actual_threads = effective_threads(policy_threads); + int64_t sampled_nnz = -1; + auto end_state = RandBLAS::fill_sparse_unpacked( + dist, dist.n_rows, dist.n_cols, 0, 0, sampled_nnz, + values.data(), rows.data(), cols.data(), seed_state + ); + if (scaling_records.empty()) { + expected_nnz = sampled_nnz; + expected_values = values; + expected_rows = rows; + expected_cols = cols; + expected_state = end_state; + } else { + exact = exact + && sampled_nnz == expected_nnz + && std::equal( + values.begin(), values.begin() + sampled_nnz, expected_values.begin() + ) + && std::equal( + rows.begin(), rows.begin() + sampled_nnz, expected_rows.begin() + ) + && std::equal( + cols.begin(), cols.begin() + sampled_nnz, expected_cols.begin() + ) + && end_state == expected_state; + } + + auto [min_us, median_us] = run_trials([&]() { + sampled_nnz = -1; + end_state = RandBLAS::fill_sparse_unpacked( + dist, dist.n_rows, dist.n_cols, 0, 0, sampled_nnz, + values.data(), rows.data(), cols.data(), seed_state + ); + }, num_trials); + if (scaling_records.empty()) { + baseline_us = min_us; + baseline_threads = actual_threads; + } + const double speedup = min_us > 0 + ? static_cast(baseline_us) / static_cast(min_us) + : -1.0; + const double relative_threads = static_cast(actual_threads) + / static_cast(baseline_threads); + scaling_records.push_back({ + thread_count, actual_threads, min_us, median_us, + static_cast(min_us) * 1000.0 + / static_cast(sampled_nnz), + speedup, speedup / relative_threads + }); + } + + std::cout << " " << spec.label << " sampling (nnz=" << expected_nnz << ")\n" + << " " << std::right << std::setw(6) << "Req" + << std::setw(6) << "Thr" + << std::setw(10) << "Min(us)" + << std::setw(10) << "Med(us)" + << std::setw(13) << "ns/nonzero" + << std::setw(10) << "Spd(min)" + << std::setw(10) << "Eff(min)" << "\n" + << " " << std::string(65, '-') << "\n"; + for (const SamplingScalingRecord &record : scaling_records) { + std::cout << " " << std::right << std::setw(6) << record.requested_threads + << std::setw(6) << record.threads + << std::setw(10) << record.min_us + << std::setw(10) << record.median_us + << std::setw(13) << fcell(record.ns_per_nonzero, 2) + << std::setw(10) << fcell(record.speedup, 2) + << std::setw(10) << fcell(record.efficiency, 2) << "\n"; + } + std::cout << " exact output/state: " << (exact ? "PASS" : "FAIL") << "\n\n"; + return exact; +} + // --------------------------------------------------------------------------- -// 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. +// SCALING: re-run sparse sampling and the left-sketch ColMajor warm apply at +// each thread count. STREAM is recalibrated per thread count for the application +// table, 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, +bool 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 + std::cout << "=== SCALING (sampling + 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"; @@ -424,24 +523,31 @@ void run_scaling(int64_t d, int64_t m, int64_t n, const std::vector& spe rb::DenseDist DA(m, n); auto st = rb::fill_dense(DA, A_cm.data(), rb::RNGState<>(seed)); std::vector B_cm(d * n); + bool sampling_exact = true; for (const auto& spec : specs) { SparseDist dist(d, m, spec.vec_nnz, spec.axis); + sampling_exact = run_sampling_scaling( + spec, dist, st, num_trials, threads + ) && sampling_exact; 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::cout << " " << std::right << std::setw(6) << "Req" + << std::setw(6) << "Thr" + << std::setw(10) << "Min(us)" << std::setw(10) << "Spd(min)" + << std::setw(10) << "Eff(min)" << std::setw(9) << "GB/s" << std::setw(7) << "%STR" << "\n"; - std::cout << " " << std::string(48, '-') << "\n"; + std::cout << " " << std::string(56, '-') << "\n"; long base_us = 0; - int base_t = threads.empty() ? 1 : threads.front(); + int base_threads = 1; + bool first_thread_count = true; for (int t : threads) { set_threads(t); + const int actual_threads = effective_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); @@ -449,22 +555,29 @@ void run_scaling(int64_t d, int64_t m, int64_t n, const std::vector& spe 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; + if (first_thread_count) { + base_us = mn; + base_threads = actual_threads; + first_thread_count = false; + } double speedup = (mn > 0) ? (double)base_us / (double)mn : 0.0; - double eff = speedup / ((double)t / (double)base_t); + double eff = speedup + / ((double)actual_threads / (double)base_threads); 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(6) << actual_threads << std::setw(10) << mn << std::setw(10) << fcell(speedup, 2) - << std::setw(8) << fcell(eff, 2) + << std::setw(10) << fcell(eff, 2) << std::setw(9) << fcell(gbps, 1) << std::setw(7) << fcell(pct, 0) << "\n"; } std::cout << "\n"; } + return sampling_exact; } // --------------------------------------------------------------------------- @@ -569,29 +682,26 @@ std::vector sweep_specs() { }; } -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) { + OpenMPSettingsGuard openmp_settings; bool no_stream = false, scaling = false, csr_probe = false; - std::vector threads = {1, 2, 4, 8}; + const std::vector default_thread_counts{1, 2, 4, 8}; + auto thread_config = parse_thread_counts("", default_thread_counts); 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 if (s.rfind("--threads=", 0) == 0) { + thread_config = parse_thread_counts(s.substr(10), default_thread_counts); + } else pos.push_back(s); } + if (!thread_config.valid) { + std::cerr << "Invalid thread list. Expected positive integers.\n"; + return 1; + } std::cout << "\n============================================================\n"; std::cout << "SKETCH_GENERAL PERFORMANCE BENCHMARK (sparse operators)\n"; @@ -612,8 +722,10 @@ int main(int argc, char** argv) { ? 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; + return run_scaling( + sd, sm, sn, specs, trials, + thread_config.thread_counts, !no_stream + ) ? 0 : 2; } if (!no_stream) { diff --git a/examples/sparse-low-rank-approx/svd_rank1_plus_noise.cc b/examples/sparse-low-rank-approx/svd_rank1_plus_noise.cc index 07f63a22..cf322bbe 100644 --- a/examples/sparse-low-rank-approx/svd_rank1_plus_noise.cc +++ b/examples/sparse-low-rank-approx/svd_rank1_plus_noise.cc @@ -157,8 +157,13 @@ void make_signal_matrix(double signal_scale, double* u, int64_t m, double* v, in double uv_scale = 1.0 / std::sqrt((double) vec_nnz); - auto v_state = RandBLAS::sparse::repeated_fisher_yates(u_state, vec_nnz, m, 1, work_idxs, trash, work_vals); - auto next_state = RandBLAS::sparse::repeated_fisher_yates(v_state, vec_nnz, n, 1, work_idxs+vec_nnz, trash, work_vals+vec_nnz); + auto v_state = RandBLAS::sparse::repeated_fisher_yates( + u_state, vec_nnz, m, 1, work_idxs, trash, work_vals, /*apply_sort=*/false + ); + auto next_state = RandBLAS::sparse::repeated_fisher_yates( + v_state, vec_nnz, n, 1, work_idxs+vec_nnz, trash, work_vals+vec_nnz, + /*apply_sort=*/false + ); for (int j = 0; j < vec_nnz; ++j) { for (int i = 0; i < vec_nnz; ++i) { int temp = i + j*vec_nnz; diff --git a/rtd/source/tutorial/sampling_skops.rst b/rtd/source/tutorial/sampling_skops.rst index 1b21a64d..34cf63da 100644 --- a/rtd/source/tutorial/sampling_skops.rst +++ b/rtd/source/tutorial/sampling_skops.rst @@ -9,9 +9,12 @@ Sampling a sketching operator RandBLAS relies on counter-based random number generators (CBRNGs) from Random123. A CBRNG returns a random number upon being called with two integer parameters: the *counter* and the *key*. The time required for the CBRNG to return does not depend on either of these parameters. -A serial application can set the key at the outset of the program and never change it, while -parallel applications should use different keys across different threads. -Sequential calls to the CBRNG with a fixed key should use different values for the counter. +RandBLAS assigns counter values to logical matrix locations, not to physical threads. +You provide one RNGState, and RandBLAS partitions its counter space when dense or sparse +operator sampling uses the available OpenMP threads. +The resulting operator and returned RNGState therefore don't depend on the thread count or schedule. +Different keys are useful when you want statistically separate operator streams or program runs; +they aren't needed merely because an application uses multiple threads. RandBLAS doesn't expose CBRNGs directly. Instead, it exposes an abstraction of @@ -114,4 +117,3 @@ different keys, as in the following code: // ^ S1 and S2 are defined only from a mathematical perspective. // No real work is performed here. - diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 6e154929..1e9a4baf 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -64,6 +64,7 @@ if (GTest_FOUND) basic_rng/test_discrete.cc basic_rng/test_continuous.cc basic_rng/test_distortion.cc + basic_rng/test_samplers.cc ) add_executable(stat_tests ${STAT_SOURCES}) target_link_libraries(stat_tests RandBLAS GTest::GTest GTest::Main) @@ -79,7 +80,12 @@ if (GTest_FOUND) # ##################################################################### - set(META_SOURCES meta/test_lapack_like.cc meta/test_sparse_data_generators.cc meta/test_comparison.cc) + set(META_SOURCES + meta/test_benchmarking.cc + meta/test_comparison.cc + meta/test_lapack_like.cc + meta/test_sparse_data_generators.cc + ) add_executable(meta_tests ${META_SOURCES}) target_link_libraries(meta_tests RandBLAS GTest::GTest GTest::Main) randblas_stage_runtime_dlls(meta_tests) diff --git a/test/basic_rng/test_discrete.cc b/test/basic_rng/test_discrete.cc index 54c504ba..264ba248 100644 --- a/test/basic_rng/test_discrete.cc +++ b/test/basic_rng/test_discrete.cc @@ -41,7 +41,12 @@ using RandBLAS::repeated_fisher_yates; #include "RandBLAS/testing/stats.hh" #include "RandBLAS/testing/comparison.hh" +#if defined(RandBLAS_HAS_OpenMP) +#include +#endif + #include +#include #include #include #include @@ -343,6 +348,150 @@ TEST_F(TestSampleIndices, rngstate_updates_fisher_yates) { test_updated_rngstates_fisher_yates(); } +TEST_F(TestSampleIndices, fisher_yates_can_sort_major_axis_vectors) { + // Sparse-operator construction needs each major-axis vector in ascending + // coordinate order, while the public sampling function preserves the raw + // Fisher--Yates order. Generate both forms from the same state, manually + // sort each raw (coordinate, value) pair, and compare that reference with + // the opt-in sorted result. This also checks that sorting leaves the minor + // coordinate and returned RNG state unchanged. + constexpr int64_t dim_major = 29; + constexpr int64_t num_vectors = 8; + constexpr int64_t vec_nnz = 7; + constexpr int64_t nnz = num_vectors * vec_nnz; + RandBLAS::RNGState<> seed(1729); + seed.counter.incr(306); + std::vector raw_major(nnz); + std::vector raw_minor(nnz); + std::vector raw_vals(nnz); + std::vector sorted_major(nnz); + std::vector sorted_minor(nnz); + std::vector sorted_vals(nnz); + + auto raw_state = RandBLAS::sparse::repeated_fisher_yates( + seed, vec_nnz, dim_major, num_vectors, + raw_major.data(), raw_minor.data(), raw_vals.data(), false + ); + auto sorted_state = RandBLAS::sparse::repeated_fisher_yates( + seed, vec_nnz, dim_major, num_vectors, + sorted_major.data(), sorted_minor.data(), sorted_vals.data(), true + ); + + EXPECT_EQ(sorted_state, raw_state); + for (int64_t i = 0; i < num_vectors; ++i) { + const int64_t offset = i * vec_nnz; + std::vector> expected(vec_nnz); + for (int64_t j = 0; j < vec_nnz; ++j) { + expected[j] = {raw_major[offset + j], raw_vals[offset + j]}; + } + std::sort(expected.begin(), expected.end()); + for (int64_t j = 0; j < vec_nnz; ++j) { + EXPECT_EQ(sorted_major[offset + j], expected[j].first); + EXPECT_EQ(sorted_vals[offset + j], expected[j].second); + EXPECT_EQ(sorted_minor[offset + j], i); + } + } +} + + +#if defined(RandBLAS_HAS_OpenMP) +TEST_F(TestSampleIndices, fisher_yates_is_thread_count_independent) { + // Parallel sampling must reproduce the serial samples and the serial end + // state exactly. Use one thread to define the reference result, then repeat + // the same call with one, two, and four threads. The two vec_nnz values + // exercise both the single-index fast path and the general Fisher-Yates + // path, and the nonzero initial counter catches implementations that + // accidentally treat the counter as though it always starts at zero. + constexpr int64_t n = 29; + constexpr int64_t num_vectors = 2048; + constexpr std::array vec_nnz_values{1, 7}; + constexpr std::array thread_counts{1, 2, 4}; + const int saved_dynamic = omp_get_dynamic(); + const int saved_max_threads = omp_get_max_threads(); + omp_set_dynamic(0); + + for (int64_t vec_nnz : vec_nnz_values) { + RandBLAS::RNGState<> seed(1729); + seed.counter.incr(306); + std::vector expected(num_vectors * vec_nnz, -1); + + omp_set_num_threads(1); + auto expected_state = RandBLAS::repeated_fisher_yates( + vec_nnz, n, num_vectors, expected.data(), seed + ); + + for (int thread_count : thread_counts) { + std::vector actual(num_vectors * vec_nnz, -1); + omp_set_num_threads(thread_count); + auto actual_state = RandBLAS::repeated_fisher_yates( + vec_nnz, n, num_vectors, actual.data(), seed + ); + EXPECT_EQ(actual, expected); + EXPECT_EQ(actual_state, expected_state); + } + } + + omp_set_num_threads(saved_max_threads); + omp_set_dynamic(saved_dynamic); +} + +TEST_F(TestSampleIndices, fisher_yates_is_exact_at_parallel_policy_boundary) { + // General Fisher-Yates enters the parallel path when num_vectors * vec_nnz + // reaches 1024. With vec_nnz = 4, 255 vectors give 1020 units of useful + // work and stay serial, while 256 vectors give 1024 and use the four + // available threads. Compare both workloads against one-thread references + // to check exactness on either side of that policy boundary. + constexpr int64_t n = 29; + constexpr int64_t vec_nnz = 4; + constexpr std::array num_vectors_values{255, 256}; + const int saved_dynamic = omp_get_dynamic(); + const int saved_max_threads = omp_get_max_threads(); + omp_set_dynamic(0); + omp_set_num_threads(4); + + EXPECT_EQ(RandBLAS::sparse::sparse_sampling_thread_count(n, 255, vec_nnz, true), 1); + EXPECT_EQ(RandBLAS::sparse::sparse_sampling_thread_count(n, 256, vec_nnz, true), 4); + + RandBLAS::RNGState<> seed(20260822); + seed.counter.incr(173); + for (int64_t num_vectors : num_vectors_values) { + std::vector expected(num_vectors * vec_nnz, -1); + std::vector actual(num_vectors * vec_nnz, -1); + omp_set_num_threads(1); + auto expected_state = RandBLAS::repeated_fisher_yates( + vec_nnz, n, num_vectors, expected.data(), seed + ); + omp_set_num_threads(4); + auto actual_state = RandBLAS::repeated_fisher_yates( + vec_nnz, n, num_vectors, actual.data(), seed + ); + + EXPECT_EQ(actual, expected); + EXPECT_EQ(actual_state, expected_state); + } + + omp_set_num_threads(saved_max_threads); + omp_set_dynamic(saved_dynamic); +} + +TEST_F(TestSampleIndices, sparse_sampling_thread_policy_uses_available_threads) { + // The sampling policy should use the OpenMP threads available to a large + // job without paying parallel overhead for a small job. Advertise four + // threads, then check one problem above the policy's work threshold and + // one below it. + const int saved_dynamic = omp_get_dynamic(); + const int saved_max_threads = omp_get_max_threads(); + omp_set_dynamic(0); + omp_set_num_threads(4); + + EXPECT_EQ(RandBLAS::sparse::sparse_sampling_thread_count(2000, 100000, 4, true), 4); + EXPECT_EQ(RandBLAS::sparse::sparse_sampling_thread_count(2000, 100, 4, true), 1); + + omp_set_num_threads(saved_max_threads); + omp_set_dynamic(saved_dynamic); +} +#endif + TEST_F(TestSampleIndices, smoke_3_x_10) { for (uint32_t i = 0; i < 10; ++i) diff --git a/test/basic_rng/test_samplers.cc b/test/basic_rng/test_samplers.cc new file mode 100644 index 00000000..19faeb83 --- /dev/null +++ b/test/basic_rng/test_samplers.cc @@ -0,0 +1,238 @@ +// Copyright, 2026. See LICENSE for copyright holder information. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// (1) Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// +// (2) Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// (3) Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +// POSSIBILITY OF SUCH DAMAGE. + +#include "RandBLAS/testing/samplers.hh" + +#include + +#include +#include +#include +#include +#include + +// MARK: test helpers + +class TestSamplers : public ::testing::Test { +protected: + // A support sampler writes a block of vec_nnz indices for each major-axis + // vector. This helper checks that every block represents a subset of + // {0, ..., n - 1}: all indices are in range, and no index is repeated + // within a block. It does not test whether those subsets are uniform. + static void expect_valid_samples( + int64_t n, int64_t num_vectors, int64_t vec_nnz, + const std::vector &samples + ) { + ASSERT_EQ(samples.size(), static_cast(num_vectors * vec_nnz)); + for (int64_t vector = 0; vector < num_vectors; ++vector) { + std::vector seen(n, false); + for (int64_t entry = 0; entry < vec_nnz; ++entry) { + int64_t index = samples[vector * vec_nnz + entry]; + ASSERT_GE(index, 0); + ASSERT_LT(index, n); + ASSERT_FALSE(seen[index]); + seen[index] = true; + } + } + } +}; + +// MARK: support samplers + +TEST_F(TestSamplers, std_sample_produces_valid_major_axis_vectors) { + // std::sample is the most direct standard-library baseline for sampling + // without replacement. Draw several support sets from one RNG stream and + // pass the resulting blocks through the structural checks above. + constexpr int64_t n = 7; + constexpr int64_t num_vectors = 11; + constexpr int64_t vec_nnz = 4; + std::vector samples(num_vectors * vec_nnz, -1); + std::mt19937_64 rng(42); + + RandBLAS::testing::sample_std_sample( + n, num_vectors, vec_nnz, samples.data(), rng + ); + + expect_valid_samples(n, num_vectors, vec_nnz, samples); +} + +TEST_F(TestSamplers, partial_fisher_yates_produces_valid_major_axis_vectors) { + // Partial Fisher-Yates stops after selecting the requested number of + // indices rather than shuffling the full population. Check that its output + // still has the subset structure required of every major-axis vector. + constexpr int64_t n = 7; + constexpr int64_t num_vectors = 11; + constexpr int64_t vec_nnz = 4; + std::vector samples(num_vectors * vec_nnz, -1); + std::mt19937_64 rng(42); + + RandBLAS::testing::sample_partial_fisher_yates( + n, num_vectors, vec_nnz, samples.data(), rng + ); + + expect_valid_samples(n, num_vectors, vec_nnz, samples); +} + +TEST_F(TestSamplers, full_shuffle_produces_valid_major_axis_vectors) { + // This baseline shuffles all n indices and keeps the first vec_nnz entries. + // The retained prefix should therefore pass the same range and uniqueness + // checks as the samplers that avoid a full shuffle. + constexpr int64_t n = 7; + constexpr int64_t num_vectors = 11; + constexpr int64_t vec_nnz = 4; + std::vector samples(num_vectors * vec_nnz, -1); + std::mt19937_64 rng(42); + + RandBLAS::testing::sample_full_shuffle( + n, num_vectors, vec_nnz, samples.data(), rng + ); + + expect_valid_samples(n, num_vectors, vec_nnz, samples); +} + +TEST_F(TestSamplers, rejection_produces_valid_major_axis_vectors) { + // Rejection sampling draws indices until it has vec_nnz distinct values. + // Generate several blocks and check the property that makes a completed + // block valid: every entry is in range and appears only once. + constexpr int64_t n = 7; + constexpr int64_t num_vectors = 11; + constexpr int64_t vec_nnz = 4; + std::vector samples(num_vectors * vec_nnz, -1); + std::mt19937_64 rng(42); + + RandBLAS::testing::sample_rejection( + n, num_vectors, vec_nnz, samples.data(), rng + ); + + expect_valid_samples(n, num_vectors, vec_nnz, samples); +} + +TEST_F(TestSamplers, floyd_produces_valid_major_axis_vectors) { + // Floyd's algorithm constructs a subset without storing a permutation of + // all n indices. Its internal representation is different, but its output + // must satisfy the same per-vector subset contract. + constexpr int64_t n = 7; + constexpr int64_t num_vectors = 11; + constexpr int64_t vec_nnz = 4; + std::vector samples(num_vectors * vec_nnz, -1); + std::mt19937_64 rng(42); + + RandBLAS::testing::sample_floyd( + n, num_vectors, vec_nnz, samples.data(), rng + ); + + expect_valid_samples(n, num_vectors, vec_nnz, samples); +} + +// MARK: end-to-end COO sampling + +TEST_F(TestSamplers, saso_data_respects_major_axis_orientation) { + // fill_saso_data turns sampled support sets into COO matrix data. Within + // each block, the minor coordinate identifies the major-axis vector while + // the major coordinate stores that vector's sampled support. Exercise both + // orientations and check the full COO contract: valid sorted supports, + // correct vector labels, and nonzero values in {-1, +1}. + constexpr int64_t n = 7; + constexpr int64_t num_vectors = 11; + constexpr int64_t vec_nnz = 4; + constexpr int64_t nnz = num_vectors * vec_nnz; + + for (bool major_is_rows : {false, true}) { + std::vector rows(nnz, -1); + std::vector cols(nnz, -1); + std::vector values(nnz, 0.0); + std::mt19937_64 rng(42); + auto sampler = []( + int64_t sampler_n, int64_t sampler_num_vectors, + int64_t sampler_vec_nnz, int64_t *samples, auto &sampler_rng + ) { + RandBLAS::testing::sample_partial_fisher_yates( + sampler_n, sampler_num_vectors, sampler_vec_nnz, + samples, sampler_rng + ); + }; + + RandBLAS::testing::fill_saso_data( + n, num_vectors, vec_nnz, major_is_rows, + rows.data(), cols.data(), values.data(), rng, sampler + ); + + const std::vector &major = major_is_rows ? rows : cols; + const std::vector &minor = major_is_rows ? cols : rows; + expect_valid_samples(n, num_vectors, vec_nnz, major); + for (int64_t vector = 0; vector < num_vectors; ++vector) { + for (int64_t entry = 0; entry < vec_nnz; ++entry) { + int64_t offset = vector * vec_nnz + entry; + EXPECT_EQ(minor[offset], vector); + EXPECT_TRUE(values[offset] == -1.0 || values[offset] == 1.0); + if (entry > 0) { + EXPECT_LT(major[offset - 1], major[offset]); + } + } + } + } +} + +// MARK: Philox URBG adapter + +TEST_F(TestSamplers, philox_urbg_matches_randblas_stream) { + // In its default mode, PhiloxURBG returns one 64-bit value per Philox + // counter and then advances to the next counter. Generate the same counters + // directly with DefaultRNG, combine the first two 32-bit words by hand, + // and use those values as the reference stream. + constexpr uint64_t seed = 42; + RandBLAS::RNGState<> state(seed); + RandBLAS::DefaultRNG generator; + RandBLAS::testing::PhiloxURBG rng(seed); + + for (int64_t draw = 0; draw < 3; ++draw) { + auto random_values = generator(state.counter, state.key); + uint64_t expected = RandBLAS::promote_uint_pair(random_values[0], random_values[1]); + EXPECT_EQ(rng(), expected); + state.counter.incr(); + } +} + +TEST_F(TestSamplers, philox_urbg_can_use_both_results_per_counter) { + // When one-result-per-counter mode is disabled, the adapter exposes two + // 64-bit values from each four-word Philox result before advancing the + // counter. Compute both values directly and check their order. + constexpr uint64_t seed = 42; + RandBLAS::RNGState<> state(seed); + RandBLAS::DefaultRNG generator; + RandBLAS::testing::PhiloxURBG rng(seed, false); + + for (int64_t counter = 0; counter < 3; ++counter) { + auto random_values = generator(state.counter, state.key); + uint64_t first = RandBLAS::promote_uint_pair(random_values[0], random_values[1]); + uint64_t second = RandBLAS::promote_uint_pair(random_values[2], random_values[3]); + EXPECT_EQ(rng(), first); + EXPECT_EQ(rng(), second); + state.counter.incr(); + } +} diff --git a/test/datastructures/test_denseskop.cc b/test/datastructures/test_denseskop.cc index 3305d765..4168d745 100644 --- a/test/datastructures/test_denseskop.cc +++ b/test/datastructures/test_denseskop.cc @@ -37,6 +37,8 @@ #include #include +#include +#include #include #include @@ -490,3 +492,26 @@ TEST_F(TestDenseSkOpStates, compare_skopless_fill_dense_to_compute_next_state) { test_compute_next_state(key, 91, 43, sd); } } + +TEST_F(TestDenseSkOpStates, compute_next_state_avoids_padded_dimension_overflow) { + // Dense sampling consumes ceil(dim_major / counter_size) counters for each + // major-axis vector. Computing that ceiling as (dim_major + padding) / + // counter_size overflows when dim_major is INT64_MAX. Build a distribution + // with that extreme major dimension, advance a copy of the seed by the + // mathematically equivalent quotient-plus-one formula, and compare every + // word of the resulting counters. No dense matrix needs to be allocated. + using RNG = r123::Philox4x32; + constexpr int64_t max_int64 = std::numeric_limits::max(); + constexpr uint64_t ctr_size = RNG::ctr_type::static_size; + constexpr uint64_t expected_increment = static_cast(max_int64) / ctr_size + 1; + RandBLAS::DenseDist dist(max_int64, 1, RandBLAS::ScalarDist::Gaussian, RandBLAS::Axis::Long); + RandBLAS::RNGState state(0); + + auto actual = RandBLAS::dense::compute_next_state(dist, state); + auto expected = state; + expected.counter.incr(expected_increment); + + for (int i = 0; i < RNG::ctr_type::static_size; ++i) { + EXPECT_EQ(actual.counter[i], expected.counter[i]); + } +} diff --git a/test/datastructures/test_sparseskop.cc b/test/datastructures/test_sparseskop.cc index faee4004..0bc9f77c 100644 --- a/test/datastructures/test_sparseskop.cc +++ b/test/datastructures/test_sparseskop.cc @@ -31,8 +31,16 @@ #include #include #include "RandBLAS/testing/comparison.hh" + +#if defined(RandBLAS_HAS_OpenMP) +#include +#endif + #include + +#include #include +#include #include #include @@ -51,9 +59,52 @@ using RandBLAS::Axis; using RandBLAS::fill_sparse; using RandBLAS::fill_sparse_unpacked_nosub; +// A sparse sampling call has five observable outputs: the number of stored +// entries, three COO arrays, and the RNG state that follows the sample. Keep +// them together so the parallel tests can compare a complete result rather +// than checking only the sampled coordinates. +template +struct SparseSnapshot { + int64_t nnz; + std::vector vals; + std::vector rows; + std::vector cols; + RandBLAS::RNGState<> end_state; +}; -class TestSparseSkOpConstruction : public ::testing::Test -{ +// Allocate the largest COO buffers permitted by the requested window, sample +// the window, and trim those buffers to the number of entries actually written. +// Trimming matters for LASOs, where restricting a sampled vector to a window +// can remove entries and leave the output shorter than its upper bound. +template +static SparseSnapshot sample_sparse_snapshot( + const RandBLAS::SparseDist &dist, int64_t n_rows_sub, int64_t n_cols_sub, + int64_t row_offset, int64_t col_offset, + const RandBLAS::RNGState<> &seed +) { + const bool short_is_rows = dist.n_rows <= dist.n_cols; + const int64_t short_sub = short_is_rows ? n_rows_sub : n_cols_sub; + const int64_t long_sub = short_is_rows ? n_cols_sub : n_rows_sub; + const int64_t num_vectors = dist.major_axis == RandBLAS::Axis::Short + ? long_sub + : short_sub; + const int64_t capacity = dist.vec_nnz * num_vectors; + SparseSnapshot result{ + -1, std::vector(capacity), std::vector(capacity), + std::vector(capacity), seed + }; + result.end_state = RandBLAS::fill_sparse_unpacked( + dist, n_rows_sub, n_cols_sub, row_offset, col_offset, result.nnz, + result.vals.data(), result.rows.data(), result.cols.data(), seed + ); + result.vals.resize(result.nnz); + result.rows.resize(result.nnz); + result.cols.resize(result.nnz); + return result; +} + + +class TestSparseSkOpConstruction : public ::testing::Test { protected: std::vector keys{42, 0, 1}; std::vector vec_nnzs{(int64_t) 1, (int64_t) 2, (int64_t) 3, (int64_t) 7}; @@ -303,6 +354,232 @@ class TestSparseSkOpConstruction : public ::testing::Test }; +TEST_F(TestSparseSkOpConstruction, submatrix_sampling_rejects_invalid_windows) { + // Submatrix validation must reject each negative argument and every window + // that extends past the parent operator. Exercise workspace-query mode so + // the check has to run before size arithmetic, allocation, RNG counter + // updates, or output writes. The near-INT64_MAX offsets specifically catch + // addition-form bounds checks that can overflow before making a decision. + struct Window { + int64_t n_rows; + int64_t n_cols; + int64_t row_offset; + int64_t col_offset; + }; + constexpr int64_t max_int64 = std::numeric_limits::max(); + constexpr std::array invalid_windows{{ + {-1, 5, 0, 0}, {5, -1, 0, 0}, {5, 5, -1, 0}, {5, 5, 0, -1}, + {8, 5, 0, 0}, {5, 14, 0, 0}, {7, 5, 1, 0}, {5, 13, 0, 1}, + {1, 1, max_int64, 0}, {1, 1, 0, max_int64} + }}; + RandBLAS::SparseDist dist(7, 13, 2, RandBLAS::Axis::Short); + RandBLAS::RNGState<> seed(314); + RandBLAS::SparseSkOp S(dist, seed); + + for (const Window &window : invalid_windows) { + int64_t capacity = -1; + EXPECT_THROW( + RandBLAS::fill_sparse_unpacked( + dist, window.n_rows, window.n_cols, + window.row_offset, window.col_offset, capacity, + static_cast(nullptr), + static_cast(nullptr), + static_cast(nullptr), seed + ), + RandBLAS::Error + ); + EXPECT_THROW( + RandBLAS::sparse::submatrix_as_coo( + S, window.n_rows, window.n_cols, + window.row_offset, window.col_offset + ), + RandBLAS::Error + ); + } +} + +TEST_F(TestSparseSkOpConstruction, submatrix_sampling_accepts_empty_boundary_window) { + // A zero-by-zero window at the lower-right corner is inside the parent + // operator. Verify that both the workspace query and the owning-COO helper + // accept it, report no entries, and leave caller-owned sentinel storage + // untouched when the sampling overload receives nonnull buffers. + RandBLAS::SparseDist dist(7, 13, 2, RandBLAS::Axis::Short); + RandBLAS::RNGState<> seed(2718); + int64_t capacity = -1; + RandBLAS::fill_sparse_unpacked( + dist, 0, 0, 7, 13, capacity, + static_cast(nullptr), + static_cast(nullptr), + static_cast(nullptr), seed + ); + EXPECT_EQ(capacity, 0); + + double val = 1.25; + int64_t row = 23; + int64_t col = 29; + int64_t nnz = -1; + RandBLAS::fill_sparse_unpacked( + dist, 0, 0, 7, 13, nnz, &val, &row, &col, seed + ); + EXPECT_EQ(nnz, 0); + EXPECT_EQ(val, 1.25); + EXPECT_EQ(row, 23); + EXPECT_EQ(col, 29); + + RandBLAS::SparseSkOp S(dist, seed); + auto empty = RandBLAS::sparse::submatrix_as_coo(S, 0, 0, 7, 13); + EXPECT_EQ(empty.n_rows, 0); + EXPECT_EQ(empty.n_cols, 0); + EXPECT_EQ(empty.nnz, 0); +} + +#if defined(RandBLAS_HAS_OpenMP) +TEST_F(TestSparseSkOpConstruction, sampling_is_thread_count_independent) { + // Changing the OpenMP thread count must not change a sampled sparse + // operator. For SASOs and LASOs of several shapes and sparsity levels, use + // a serial full-operator sample and a serial submatrix sample as reference + // results. Repeat each call with one, two, and four threads, then compare + // every COO array, the number of entries written, and the returned RNG + // state. A nonzero initial counter also checks that work is partitioned + // relative to the supplied state rather than an implicit zero state. + struct Case { + int64_t n_rows; + int64_t n_cols; + int64_t vec_nnz; + RandBLAS::Axis major_axis; + }; + constexpr std::array cases{{ + {7, 31, 1, RandBLAS::Axis::Short}, + {7, 31, 5, RandBLAS::Axis::Short}, + {31, 7, 5, RandBLAS::Axis::Short}, + {257, 2048, 1, RandBLAS::Axis::Short}, + {257, 2048, 12, RandBLAS::Axis::Short}, + {7, 31, 1, RandBLAS::Axis::Long}, + {7, 31, 12, RandBLAS::Axis::Long}, + {31, 7, 12, RandBLAS::Axis::Long}, + {257, 257, 128, RandBLAS::Axis::Long}, + {2048, 4096, 1, RandBLAS::Axis::Long} + }}; + constexpr std::array thread_counts{1, 2, 4}; + const int saved_dynamic = omp_get_dynamic(); + const int saved_max_threads = omp_get_max_threads(); + omp_set_dynamic(0); + + RandBLAS::RNGState<> seed(20260817); + seed.counter.incr(41); + for (const Case &test_case : cases) { + RandBLAS::SparseDist dist( + test_case.n_rows, test_case.n_cols, + test_case.vec_nnz, test_case.major_axis + ); + omp_set_num_threads(1); + auto expected_full = sample_sparse_snapshot( + dist, dist.n_rows, dist.n_cols, 0, 0, seed + ); + auto expected_sub = sample_sparse_snapshot( + dist, dist.n_rows - 2, dist.n_cols - 3, 1, 2, seed + ); + + for (int thread_count : thread_counts) { + omp_set_num_threads(thread_count); + auto actual_full = sample_sparse_snapshot( + dist, dist.n_rows, dist.n_cols, 0, 0, seed + ); + auto actual_sub = sample_sparse_snapshot( + dist, dist.n_rows - 2, dist.n_cols - 3, 1, 2, seed + ); + EXPECT_EQ(actual_full.nnz, expected_full.nnz); + EXPECT_EQ(actual_full.vals, expected_full.vals); + EXPECT_EQ(actual_full.rows, expected_full.rows); + EXPECT_EQ(actual_full.cols, expected_full.cols); + EXPECT_EQ(actual_full.end_state, expected_full.end_state); + EXPECT_EQ(actual_sub.nnz, expected_sub.nnz); + EXPECT_EQ(actual_sub.vals, expected_sub.vals); + EXPECT_EQ(actual_sub.rows, expected_sub.rows); + EXPECT_EQ(actual_sub.cols, expected_sub.cols); + EXPECT_EQ(actual_sub.end_state, expected_sub.end_state); + } + } + + omp_set_num_threads(saved_max_threads); + omp_set_dynamic(saved_dynamic); +} + +TEST_F(TestSparseSkOpConstruction, parallel_saso_k1_writes_full_coo_data) { + // SASOs with one nonzero per major-axis vector use a specialized sampling + // path. Compare serial and four-thread samples entry-for-entry using float + // values and 32-bit indices, so this test checks that the fast path writes + // all three COO arrays and advances the RNG state for nondefault types. + const int saved_dynamic = omp_get_dynamic(); + const int saved_max_threads = omp_get_max_threads(); + omp_set_dynamic(0); + RandBLAS::SparseDist dist(257, 2048, 1, RandBLAS::Axis::Short); + RandBLAS::RNGState<> seed(991); + seed.counter.incr(73); + + omp_set_num_threads(1); + auto expected = sample_sparse_snapshot( + dist, dist.n_rows, dist.n_cols, 0, 0, seed + ); + omp_set_num_threads(4); + auto actual = sample_sparse_snapshot( + dist, dist.n_rows, dist.n_cols, 0, 0, seed + ); + + EXPECT_EQ(actual.nnz, expected.nnz); + EXPECT_EQ(actual.vals, expected.vals); + EXPECT_EQ(actual.rows, expected.rows); + EXPECT_EQ(actual.cols, expected.cols); + EXPECT_EQ(actual.end_state, expected.end_state); + + omp_set_num_threads(saved_max_threads); + omp_set_dynamic(saved_dynamic); +} + +TEST_F(TestSparseSkOpConstruction, parallel_sampling_rejects_two_word_rng) { + // Signed sparse sampling currently requires a counter-based generator + // result with at least four words so one counter can supply the sampled + // index and its sign. Force the parallel path with Philox2x32, which has + // only two words, and verify that both SASO and LASO sampling reject it + // with a RandBLAS error instead of reading nonexistent random words. + using TwoWordRNG = r123::Philox2x32; + const int saved_dynamic = omp_get_dynamic(); + const int saved_max_threads = omp_get_max_threads(); + omp_set_dynamic(0); + omp_set_num_threads(4); + + for (RandBLAS::Axis axis : {RandBLAS::Axis::Short, RandBLAS::Axis::Long}) { + RandBLAS::SparseDist dist(257, 2048, 4, axis); + RandBLAS::RNGState seed(17); + std::vector vals(dist.full_nnz); + std::vector rows(dist.full_nnz); + std::vector cols(dist.full_nnz); + int64_t nnz = -1; + EXPECT_THROW( + RandBLAS::fill_sparse_unpacked( + dist, dist.n_rows, dist.n_cols, 0, 0, nnz, + vals.data(), rows.data(), cols.data(), seed + ), + RandBLAS::Error + ); + + // The owning-COO convenience path allocates its output buffers before + // sampling. It must propagate the same validation error while allowing + // those partially constructed buffers to be reclaimed during unwinding. + RandBLAS::SparseSkOp S(dist, seed); + EXPECT_THROW( + RandBLAS::sparse::submatrix_as_coo( + S, dist.n_rows, dist.n_cols, 0, 0 + ), + RandBLAS::Error + ); + } + + omp_set_num_threads(saved_max_threads); + omp_set_dynamic(saved_dynamic); +} +#endif + TEST_F(TestSparseSkOpConstruction, respect_ownership) { respect_ownership(7, 20); respect_ownership(7, 20); diff --git a/test/linops/test_lskge3.cc b/test/linops/test_lskge3.cc index 28874b67..2abdb29e 100644 --- a/test/linops/test_lskge3.cc +++ b/test/linops/test_lskge3.cc @@ -30,6 +30,8 @@ #include "test/linops/linop_common.hh" #include +#include + using RandBLAS::DenseDist; using RandBLAS::DenseSkOp; using namespace test::linop_common; @@ -154,6 +156,41 @@ TEST_F(TestLSKGE3, sketch_eye_single_null) sketch_eye(seed, 200, 30, false, blas::Layout::ColMajor); } +TEST_F(TestLSKGE3, lazy_zero_contraction_scales_nonempty_output) { + DenseSkOp S(DenseDist(2, 2), 7); + double A = 0.0; + // Padding values make the expected result sensitive to both beta scaling + // and the column-major leading dimension. + std::array B{1.0, 2.0, 91.0, 3.0, 4.0, 92.0, 5.0, 6.0, 93.0}; + const std::array expected{2.0, 4.0, 91.0, 6.0, 8.0, 92.0, 10.0, 12.0, 93.0}; + + // A zero contraction dimension leaves beta * B. Keeping S lazy exercises + // the submatrix-packing path that used to reject this valid operation. + ASSERT_EQ(S.buff, nullptr); + ASSERT_NO_THROW(RandBLAS::sketch_general( + blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + 2, 3, 0, 1.0, S, 0, 0, &A, 1, 2.0, B.data(), 3 + )); + EXPECT_EQ(B, expected); + EXPECT_EQ(S.buff, nullptr); +} + +TEST_F(TestLSKGE3, lazy_empty_output_is_noop) { + DenseSkOp S(DenseDist(2, 2), 7); + const std::array A{1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; + double B = 17.0; + + // The zero row extent makes B empty, so neither packing S nor touching + // the sentinel output is necessary. + ASSERT_EQ(S.buff, nullptr); + ASSERT_NO_THROW(RandBLAS::sketch_general( + blas::Layout::ColMajor, blas::Op::NoTrans, blas::Op::NoTrans, + 0, 3, 2, 1.0, S, 0, 0, A.data(), 2, 2.0, &B, 1 + )); + EXPECT_EQ(B, 17.0); + EXPECT_EQ(S.buff, nullptr); +} + //////////////////////////////////////////////////////////////////////// // @@ -304,4 +341,3 @@ TEST_F(TestLSKGE3, submatrix_a_single) blas::Layout::ColMajor ); } - diff --git a/test/linops/test_lskges.cc b/test/linops/test_lskges.cc index ff92b0cd..989fa9e0 100644 --- a/test/linops/test_lskges.cc +++ b/test/linops/test_lskges.cc @@ -218,6 +218,25 @@ TEST_F(TestLSKGES, sketch_laso_colMajor_oneThread) } } +TEST_F(TestLSKGES, full_shape_submatrix_rejects_nonzero_offset) { + // Asking for a submatrix as large as S is valid only at offset (0, 0). + // Materialize S before the call so the test exercises the explicit-operator + // fast path, which must validate the offset before treating S as the result. + SparseDist dist(2, 3, 1, Axis::Short); + SparseSkOp S(dist, 0); + RandBLAS::fill_sparse(S); + const double A[] = {1.0, 1.0, 1.0}; + double B[] = {0.0, 0.0}; + + EXPECT_THROW( + RandBLAS::sparse::lskges( + blas::Layout::RowMajor, blas::Op::NoTrans, blas::Op::NoTrans, + 2, 1, 3, 1.0, S, 1, 0, A, 1, 0.0, B, 1 + ), + RandBLAS::Error + ); +} + #if defined (RandBLAS_HAS_OpenMP) TEST_F(TestLSKGES, sketch_saso_colMajor_fourThreads) { diff --git a/test/linops/test_rskge3.cc b/test/linops/test_rskge3.cc index d47d0ee3..cb8fcd4c 100644 --- a/test/linops/test_rskge3.cc +++ b/test/linops/test_rskge3.cc @@ -30,6 +30,8 @@ #include "test/linops/linop_common.hh" #include +#include + using namespace test::linop_common; using RandBLAS::DenseDist; using RandBLAS::DenseSkOp; @@ -150,6 +152,41 @@ TEST_F(TestRSKGE3, right_sketch_eye_single_null) sketch_eye(seed, 200, 30, false, Layout::ColMajor); } +TEST_F(TestRSKGE3, lazy_zero_contraction_scales_nonempty_output) { + DenseSkOp S(DenseDist(2, 2), 7); + double A = 0.0; + // Padding values make the expected result sensitive to both beta scaling + // and the row-major leading dimension. + std::array B{1.0, 2.0, 91.0, 3.0, 4.0, 92.0, 5.0, 6.0, 93.0}; + const std::array expected{2.0, 4.0, 91.0, 6.0, 8.0, 92.0, 10.0, 12.0, 93.0}; + + // A zero contraction dimension leaves beta * B. Keeping S lazy exercises + // the submatrix-packing path that used to reject this valid operation. + ASSERT_EQ(S.buff, nullptr); + ASSERT_NO_THROW(RandBLAS::sketch_general( + Layout::RowMajor, blas::Op::NoTrans, blas::Op::NoTrans, + 3, 2, 0, 1.0, &A, 1, S, 0, 0, 2.0, B.data(), 3 + )); + EXPECT_EQ(B, expected); + EXPECT_EQ(S.buff, nullptr); +} + +TEST_F(TestRSKGE3, lazy_empty_output_is_noop) { + DenseSkOp S(DenseDist(2, 2), 7); + const std::array A{1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; + double B = 17.0; + + // The zero column extent makes B empty, so neither packing S nor touching + // the sentinel output is necessary. + ASSERT_EQ(S.buff, nullptr); + ASSERT_NO_THROW(RandBLAS::sketch_general( + Layout::RowMajor, blas::Op::NoTrans, blas::Op::NoTrans, + 3, 0, 2, 1.0, A.data(), 2, S, 0, 0, 2.0, &B, 1 + )); + EXPECT_EQ(B, 17.0); + EXPECT_EQ(S.buff, nullptr); +} + //////////////////////////////////////////////////////////////////////// // diff --git a/test/linops/test_rskges.cc b/test/linops/test_rskges.cc index da1ab06a..3b5dfa5c 100644 --- a/test/linops/test_rskges.cc +++ b/test/linops/test_rskges.cc @@ -216,6 +216,25 @@ TEST_F(TestRSKGES, sketch_laso_colMajor_oneThread) } } +TEST_F(TestRSKGES, full_shape_submatrix_rejects_nonzero_offset) { + // Asking for a submatrix as large as S is valid only at offset (0, 0). + // Materialize S before the call so the test exercises the explicit-operator + // fast path, which must validate the offset before treating S as the result. + SparseDist dist(3, 2, 1, RandBLAS::Axis::Short); + SparseSkOp S(dist, 0); + RandBLAS::fill_sparse(S); + const double A[] = {1.0, 1.0, 1.0}; + double B[] = {0.0, 0.0}; + + EXPECT_THROW( + RandBLAS::sparse::rskges( + Layout::RowMajor, blas::Op::NoTrans, blas::Op::NoTrans, + 1, 2, 3, 1.0, A, 3, S, 1, 0, 0.0, B, 2 + ), + RandBLAS::Error + ); +} + //////////////////////////////////////////////////////////////////////// // diff --git a/test/linops/test_sketch_sparse.cc b/test/linops/test_sketch_sparse.cc index 2afa4d7f..f511c738 100644 --- a/test/linops/test_sketch_sparse.cc +++ b/test/linops/test_sketch_sparse.cc @@ -1,6 +1,8 @@ #include "test/linops/linop_common.hh" // ^ That includes a ton of stuff. +#include + using blas::Layout; using blas::Op; @@ -284,6 +286,41 @@ TEST_F(TestLSKSP3, sketch_eye_single_null) { sketch_eye(seed, 200, 30, false, blas::Layout::ColMajor); } +TEST_F(TestLSKSP3, lazy_zero_contraction_scales_nonempty_output) { + DenseSkOp S(DenseDist(2, 2), 7); + COOMatrix A(0, 3); + // Padding values make the expected result sensitive to both beta scaling + // and the row-major leading dimension. + std::array B{1.0, 2.0, 3.0, 91.0, 4.0, 5.0, 6.0, 92.0}; + const std::array expected{2.0, 4.0, 6.0, 91.0, 8.0, 10.0, 12.0, 92.0}; + + // A zero contraction dimension leaves beta * B. Keeping S lazy exercises + // the submatrix-packing path that used to reject this valid operation. + ASSERT_EQ(S.buff, nullptr); + ASSERT_NO_THROW(sketch_sparse( + Layout::RowMajor, Op::NoTrans, Op::NoTrans, + 2, 3, 0, 1.0, S, 0, 0, A, 2.0, B.data(), 4 + )); + EXPECT_EQ(B, expected); + EXPECT_EQ(S.buff, nullptr); +} + +TEST_F(TestLSKSP3, lazy_empty_output_is_noop) { + DenseSkOp S(DenseDist(2, 2), 7); + COOMatrix A(2, 3); + double B = 17.0; + + // The zero row extent makes B empty, so neither packing S nor touching + // the sentinel output is necessary. + ASSERT_EQ(S.buff, nullptr); + ASSERT_NO_THROW(sketch_sparse( + Layout::RowMajor, Op::NoTrans, Op::NoTrans, + 0, 3, 2, 1.0, S, 0, 0, A, 2.0, &B, 3 + )); + EXPECT_EQ(B, 17.0); + EXPECT_EQ(S.buff, nullptr); +} + //////////////////////////////////////////////////////////////////////// // // @@ -464,6 +501,41 @@ TEST_F(TestRSKSP3, right_sketch_eye_single_null) sketch_eye(seed, 200, 30, false, Layout::ColMajor); } +TEST_F(TestRSKSP3, lazy_zero_contraction_scales_nonempty_output) { + DenseSkOp S(DenseDist(2, 2), 7); + COOMatrix A(3, 0); + // Padding values make the expected result sensitive to both beta scaling + // and the column-major leading dimension. + std::array B{1.0, 2.0, 3.0, 91.0, 4.0, 5.0, 6.0, 92.0}; + const std::array expected{2.0, 4.0, 6.0, 91.0, 8.0, 10.0, 12.0, 92.0}; + + // A zero contraction dimension leaves beta * B. Keeping S lazy exercises + // the submatrix-packing path that used to reject this valid operation. + ASSERT_EQ(S.buff, nullptr); + ASSERT_NO_THROW(sketch_sparse( + Layout::ColMajor, Op::NoTrans, Op::NoTrans, + 3, 2, 0, 1.0, A, S, 0, 0, 2.0, B.data(), 4 + )); + EXPECT_EQ(B, expected); + EXPECT_EQ(S.buff, nullptr); +} + +TEST_F(TestRSKSP3, lazy_empty_output_is_noop) { + DenseSkOp S(DenseDist(2, 2), 7); + COOMatrix A(3, 2); + double B = 17.0; + + // The zero column extent makes B empty, so neither packing S nor touching + // the sentinel output is necessary. + ASSERT_EQ(S.buff, nullptr); + ASSERT_NO_THROW(sketch_sparse( + Layout::ColMajor, Op::NoTrans, Op::NoTrans, + 3, 0, 2, 1.0, A, S, 0, 0, 2.0, &B, 3 + )); + EXPECT_EQ(B, 17.0); + EXPECT_EQ(S.buff, nullptr); +} + //////////////////////////////////////////////////////////////////////// // diff --git a/test/linops/test_spmm/test_spmm_coo.cc b/test/linops/test_spmm/test_spmm_coo.cc index 72e2b0dd..5113923f 100644 --- a/test/linops/test_spmm/test_spmm_coo.cc +++ b/test/linops/test_spmm/test_spmm_coo.cc @@ -185,6 +185,24 @@ TEST_F(TestLeftMultiply_COO_single, submatrix_self) { } } +TEST_F(TestLeftMultiply_COO_double, submatrix_rejects_out_of_bounds_offset) { + // The requested 1-by-2 window has individually valid dimensions, but row + // offset 2 places it just beyond this 2-by-3 COO matrix. With alpha zero, + // the kernel would otherwise return without accessing data; the expected + // exception therefore proves that it validated the complete window first. + COOMatrix A(2, 3); + const double B[] = {1.0, 1.0}; + double C = 0.0; + + EXPECT_THROW( + left_spmm( + Layout::RowMajor, blas::Op::NoTrans, blas::Op::NoTrans, + 1, 1, 2, 0.0, A, 2, 0, B, 1, 0.0, &C, 1 + ), + RandBLAS::Error + ); +} + //////////////////////////////////////////////////////////////////////// // // submatrix of other operand in left-multiply @@ -430,4 +448,3 @@ TEST_F(TestRightMultiply_COO_double, trans_other_times_sparse_rowmajor) { transpose_other(key, 7, 22, 5, Layout::RowMajor, 0.10); transpose_other(key, 7, 22, 5, Layout::RowMajor, 0.80); } - diff --git a/test/meta/test_benchmarking.cc b/test/meta/test_benchmarking.cc new file mode 100644 index 00000000..cca00110 --- /dev/null +++ b/test/meta/test_benchmarking.cc @@ -0,0 +1,84 @@ +// Copyright, 2026. See LICENSE for copyright holder information. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// (1) Redistributions of source code must retain the above copyright notice, +// this list of conditions and the following disclaimer. +// +// (2) Redistributions in binary form must reproduce the above copyright +// notice, this list of conditions and the following disclaimer in the +// documentation and/or other materials provided with the distribution. +// +// (3) Neither the name of the copyright holder nor the names of its +// contributors may be used to endorse or promote products derived from +// this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +// POSSIBILITY OF SUCH DAMAGE. + +#include "RandBLAS/testing/benchmarking.hh" + +#include + +#include + +using RandBLAS::testing::parse_thread_counts; + +TEST(ParseThreadCounts, parses_positive_counts) { + // An explicit command-line list takes precedence over the benchmark's + // defaults. Check the parsed order as well as the positivity decision. + auto result = parse_thread_counts("1,3,8", {2, 4}); + + EXPECT_TRUE(result.valid); + EXPECT_EQ(result.thread_counts, (std::vector{1, 3, 8})); +} + +TEST(ParseThreadCounts, uses_explicit_defaults_for_an_empty_list) { + // Both benchmark programs initialize their thread configuration by + // parsing an empty string. The caller-supplied defaults should survive + // that path unchanged. + auto result = parse_thread_counts("", {3, 6}); + + EXPECT_TRUE(result.valid); + EXPECT_EQ(result.thread_counts, (std::vector{3, 6})); +} + +TEST(ParseThreadCounts, skips_empty_fields) { + // Empty comma-separated fields did not represent thread counts in either + // original parser. Verify that consolidating them preserves that behavior. + auto result = parse_thread_counts("1,,4,", {2, 8}); + + EXPECT_TRUE(result.valid); + EXPECT_EQ(result.thread_counts, (std::vector{1, 4})); +} + +TEST(ParseThreadCounts, rejects_nonpositive_counts) { + // OpenMP thread requests must be positive. Exercise zero and a negative + // value inside otherwise valid lists so either one invalidates the result. + auto zero_result = parse_thread_counts("1,0,4", {2, 8}); + auto negative_result = parse_thread_counts("1,-2,4", {2, 8}); + + EXPECT_FALSE(zero_result.valid); + EXPECT_FALSE(negative_result.valid); +} + +TEST(ParseThreadCounts, preserves_atoi_prefix_parsing) { + // This extraction is deliberately behavior-preserving: the old benchmark + // parsers used std::atoi, which accepts a numeric prefix and leading space. + // Use both forms so a stricter parser cannot arrive as an accidental part + // of this refactor. + auto result = parse_thread_counts("2threads, 4", {1, 8}); + + EXPECT_TRUE(result.valid); + EXPECT_EQ(result.thread_counts, (std::vector{2, 4})); +} diff --git a/test/test_exceptions.cc b/test/test_exceptions.cc index 3fa04b2a..da16530e 100644 --- a/test/test_exceptions.cc +++ b/test/test_exceptions.cc @@ -1,12 +1,25 @@ +#include "RandBLAS/base.hh" #include "RandBLAS/config.h" +#include "RandBLAS/dense_skops.hh" #include "RandBLAS/exceptions.hh" +#include "RandBLAS/sparse_skops.hh" -#include #include +#include +#include +#include +#include +#include +#include + +#if defined(RandBLAS_HAS_OpenMP) +#include +#endif + class TestExceptions : public ::testing::Test { protected: }; @@ -44,3 +57,166 @@ TEST_F(TestExceptions, randblas_error_if_msg_output) { } ASSERT_TRUE(expect_true); } + +TEST_F(TestExceptions, safe_int_product_multiplies_in_output_type) { + // Twice INT32_MAX cannot be represented by an int32_t, but it fits in the + // requested int64_t output. This catches implementations that multiply in + // the input type and cast only after the intermediate value has overflowed. + constexpr int32_t max_int32 = std::numeric_limits::max(); + constexpr int64_t expected = 2 * static_cast(max_int32); + + EXPECT_EQ((RandBLAS::safe_int_product(max_int32, 2)), expected); +} + +TEST_F(TestExceptions, safe_int_product_accepts_signed_boundary_products) { + // Overflow checks must leave every representable product alone. Exercise + // both ends of the int64_t range, negative products, and a product of two + // negative operands. These lie near cases where signed division requires + // special handling. + constexpr int64_t max_int64 = std::numeric_limits::max(); + constexpr int64_t min_int64 = std::numeric_limits::min(); + + EXPECT_EQ(RandBLAS::safe_int_product(min_int64, 1), min_int64); + EXPECT_EQ(RandBLAS::safe_int_product(max_int64, -1), -max_int64); + EXPECT_EQ(RandBLAS::safe_int_product(-2, -3), 6); +} + +TEST_F(TestExceptions, safe_int_product_rejects_output_type_overflow) { + // These products all lie outside the int64_t range. The two INT64_MIN * -1 + // cases also check that overflow is detected regardless of which operand + // contains the exceptional minimum value. + constexpr int64_t max_int64 = std::numeric_limits::max(); + constexpr int64_t min_int64 = std::numeric_limits::min(); + + EXPECT_THROW(RandBLAS::safe_int_product(max_int64, 2), std::overflow_error); + EXPECT_THROW(RandBLAS::safe_int_product(min_int64, -1), std::overflow_error); + EXPECT_THROW(RandBLAS::safe_int_product(-1, min_int64), std::overflow_error); +} + +TEST_F(TestExceptions, sparse_dist_rejects_full_nnz_overflow) { + // This is otherwise a valid short-axis distribution: each of its + // INT64_MAX major-axis vectors would have two nonzeros. Constructing it + // would therefore require full_nnz = 2 * INT64_MAX, so SparseDist should + // report the overflow before storing a wrapped buffer length. + constexpr int64_t max_int64 = std::numeric_limits::max(); + + EXPECT_THROW( + RandBLAS::SparseDist(max_int64, 2, 2, RandBLAS::Axis::Short), std::overflow_error + ); +} + +TEST_F(TestExceptions, fill_dense_rejects_allocation_size_overflow) { + // Choosing the short axis as the major axis keeps the RNG-state increment + // representable, so constructing the operator S succeeds. Materializing S + // would still require 2 * INT64_MAX entries. fill_dense should report that + // allocation count as an overflow rather than passing a wrapped size to + // new[]. + constexpr int64_t max_int64 = std::numeric_limits::max(); + RandBLAS::DenseDist dist(max_int64, 2, RandBLAS::ScalarDist::Gaussian, RandBLAS::Axis::Short); + RandBLAS::RNGState state(0); + RandBLAS::DenseSkOp S(dist, state); + + EXPECT_THROW(RandBLAS::fill_dense(S), std::overflow_error); +} + +TEST_F(TestExceptions, dense_submatrix_rejects_allocation_size_overflow) { + // submatrix_as_blackbox owns the buffer it creates. Request the full + // INT64_MAX-by-2 window from a distribution with valid individual + // dimensions and verify that the helper rejects the overflowing element + // count before it attempts the allocation. + constexpr int64_t max_int64 = std::numeric_limits::max(); + RandBLAS::DenseDist dist(max_int64, 2, RandBLAS::ScalarDist::Gaussian, RandBLAS::Axis::Short); + RandBLAS::RNGState state(0); + RandBLAS::DenseSkOp S(dist, state); + using BFO = RandBLAS::BLASFriendlyOperator; + + EXPECT_THROW( + RandBLAS::submatrix_as_blackbox(S, max_int64, 2, 0, 0), std::overflow_error + ); +} + +TEST_F(TestExceptions, fill_dense_rejects_starting_offset_overflow) { + // In this row-major 2-by-INT64_MAX matrix, the 1-by-1 window beginning at + // (1, 1) has linear offset INT64_MAX + 1. The output buffer itself has only + // one entry, so this isolates overflow in the parent-matrix offset and + // checks that fill_dense_unpacked rejects it before writing to the buffer. + constexpr int64_t max_int64 = std::numeric_limits::max(); + RandBLAS::DenseDist dist(2, max_int64, RandBLAS::ScalarDist::Gaussian, RandBLAS::Axis::Long); + RandBLAS::RNGState state(0); + double buff; + + EXPECT_THROW( + RandBLAS::fill_dense_unpacked( + blas::Layout::RowMajor, dist, 1, 1, 1, 1, &buff, state + ), + std::overflow_error + ); +} + +TEST_F(TestExceptions, validate_submat_dims_accepts_boundary_windows) { + // A submatrix may occupy its entire parent or may be empty. In particular, + // an empty window may begin at the parent's lower-right boundary because it + // does not address any entries beyond that boundary. + EXPECT_NO_THROW(RandBLAS::validate_submat_dims(5, 7, 5, 7, 0, 0)); + EXPECT_NO_THROW(RandBLAS::validate_submat_dims(5, 7, 2, 3, 3, 4)); + EXPECT_NO_THROW(RandBLAS::validate_submat_dims(5, 7, 0, 0, 5, 7)); +} + +TEST_F(TestExceptions, validate_submat_dims_rejects_invalid_windows) { + // The subtraction-form bounds check should reject malformed windows without + // first adding an extent and offset, which could overflow and make an invalid + // request appear valid. Cover each kind of bad dimension as well as INT64_MAX. + constexpr int64_t max_int64 = std::numeric_limits::max(); + + EXPECT_THROW(RandBLAS::validate_submat_dims(5, 7, -1, 1, 0, 0), RandBLAS::Error); + EXPECT_THROW(RandBLAS::validate_submat_dims(5, 7, 1, -1, 0, 0), RandBLAS::Error); + EXPECT_THROW(RandBLAS::validate_submat_dims(5, 7, 1, 1, -1, 0), RandBLAS::Error); + EXPECT_THROW(RandBLAS::validate_submat_dims(5, 7, 1, 1, 0, -1), RandBLAS::Error); + EXPECT_THROW(RandBLAS::validate_submat_dims(5, 7, 6, 1, 0, 0), RandBLAS::Error); + EXPECT_THROW(RandBLAS::validate_submat_dims(5, 7, 1, 8, 0, 0), RandBLAS::Error); + EXPECT_THROW(RandBLAS::validate_submat_dims(5, 7, 2, 3, 4, 0), RandBLAS::Error); + EXPECT_THROW(RandBLAS::validate_submat_dims(5, 7, 2, 3, 0, 5), RandBLAS::Error); + EXPECT_THROW( + RandBLAS::validate_submat_dims(5, 7, 1, 1, max_int64, max_int64), + RandBLAS::Error + ); +} + +TEST_F(TestExceptions, thread_number_helpers_report_serial_context) { + // Outside an OpenMP parallel region, every caller is the sole member of a + // one-thread team. These values are also the complete fallback contract for + // builds where RandBLAS was compiled without OpenMP. + EXPECT_EQ(RandBLAS::randblas_get_thread_num(), 0); + EXPECT_EQ(RandBLAS::randblas_get_num_threads(), 1); +} + +#if defined(RandBLAS_HAS_OpenMP) +TEST_F(TestExceptions, thread_number_helpers_report_openmp_team) { + // Give each physical OpenMP thread its own array slot, then compare the + // RandBLAS wrappers with OpenMP's direct answers. This checks both the + // thread's zero-based identity and the size of the team it belongs to. + const int orig_dynamic = omp_get_dynamic(); + const int orig_max_threads = omp_get_max_threads(); + const int requested_threads = std::min(4, orig_max_threads); + std::vector thread_nums(requested_threads, -1); + std::vector team_sizes(requested_threads, -1); + int actual_threads = 0; + + omp_set_dynamic(0); + #pragma omp parallel num_threads(requested_threads) + { + const int thread_num = omp_get_thread_num(); + thread_nums[thread_num] = RandBLAS::randblas_get_thread_num(); + team_sizes[thread_num] = RandBLAS::randblas_get_num_threads(); + #pragma omp single + actual_threads = omp_get_num_threads(); + } + omp_set_num_threads(orig_max_threads); + omp_set_dynamic(orig_dynamic); + + for (int thread_num = 0; thread_num < actual_threads; ++thread_num) { + EXPECT_EQ(thread_nums[thread_num], thread_num); + EXPECT_EQ(team_sizes[thread_num], actual_threads); + } +} +#endif