Sparse sampling performance benchmarks and parallelism

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

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
 

RandBLAS/DevNotes.md

@@ -5,31 +5,65 @@ 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,
    RandBLAS provides its own abstractions for sparse matrices (CSC, CSR, and COO formats).
    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.

RandBLAS/base.hh

@@ -32,13 +32,17 @@
 /// @file
 
 #include "RandBLAS/config.h"
+#include "RandBLAS/exceptions.hh"
 #include "RandBLAS/random_gen.hh"
 
 #include <blas.hh>
-#include <utility>
 #include <cstring>
 #include <cstdint>
 #include <iostream>
+#include <limits>
+#include <sstream>
+#include <stdexcept>
+#include <utility>
 
 #if defined(RandBLAS_HAS_OpenMP)
 #include <omp.h>
@@ -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<T>::is_signed && std::numeric_limit
 
 template <SignedInteger TI, SignedInteger TO = int64_t>
 inline TO safe_int_product(TI a, TI b) {
-    if (a == 0 || b == 0) {
-        return 0;
+    static_assert(
+        std::numeric_limits<TO>::digits >= std::numeric_limits<TI>::digits,
+        "safe_int_product requires an output type at least as wide as its input type."
+    );
+    const TO a_out = static_cast<TO>(a);
+    const TO b_out = static_cast<TO>(b);
+    const TO min_out = std::numeric_limits<TO>::min();
+    const TO max_out = std::numeric_limits<TO>::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;
 }