Changes to repeated_fisher_yates and sample_indices_iid_uniform

CLAUDE.md

@@ -104,7 +104,7 @@ The sparse triangular solve function (`trsm`) handles transposition similarly:
 - `lskgeX`: **Left** sketching, **ge**neral (dense) data, variant X
 - `rskgeX`: **Right** sketching, **ge**neral (dense) data, variant X
 - `lskges`: Left sketching with **s**parse operator
-- `lskge3`: Left sketching with dense operator (calls GEMM, "3" for 3-argument GEMM-like)
+- `lskge3`: Left sketching with dense operator (calls GEMM, a "BLAS 3" operation)
 - `lsksp3`: Left sketching **sp**arse data (where "left" refers to operator position)
 
 **Counterintuitive detail**: In `lsksp3` and `rsksp3`, the "left/right" refers to the operator's position. But these functions call `right_spmm`/`left_spmm` respectively, where "left/right" refers to the sparse data matrix position. See `sparse_data/DevNotes.md` lines 59-74.

RandBLAS/sparse_skops.hh

@@ -48,10 +48,6 @@
 
 namespace RandBLAS::sparse {
 
-inline std::uint64_t promote_uint_pair(std::uint32_t &a, std::uint32_t &b) {
-    return static_cast<std::uint64_t>(a) + (static_cast<std::uint64_t>(b) << 32);
-}
-
 template <typename T, SignedInteger sint_t, typename state_t = RNGState<DefaultRNG>>
 void _considerate_fisher_yates(
     const state_t &state,
@@ -110,6 +106,19 @@ static state_t repeated_fisher_yates(
     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});
+        if (vals != nullptr) {
+            return sample_indices_iid_uniform<T, sint_t, true>(
+                dim_major, dim_minor, idxs_major, vals, state);
+        } else {
+            return sample_indices_iid_uniform<sint_t>(
+                dim_major, dim_minor, idxs_major, state);
+        }
+    }
     std::vector<sint_t> vec_work(dim_major);
     std::iota(vec_work.begin(), vec_work.end(), 0);
     std::vector<sint_t> pivots(vec_nnz);
@@ -293,20 +302,13 @@ inline state_t repeated_fisher_yates(
 
 template <typename RNG = DefaultRNG>
 RNGState<RNG> compute_next_state(SparseDist dist, RNGState<RNG> state) {
-    int64_t num_mavec, incrs_per_mavec;
-    if (dist.major_axis == Axis::Short) {
-        num_mavec = std::max(dist.n_rows, dist.n_cols);
-        incrs_per_mavec = dist.vec_nnz;
-        // ^ SASOs don't try to be frugal with CBRNG increments.
-        //   See repeated_fisher_yates.
-    } else {
-        num_mavec = std::min(dist.n_rows, dist.n_cols);
-        incrs_per_mavec = (int64_t) std::ceil((double) dist.vec_nnz / ((double) state.len_c/2));
-        // ^ LASOs do try to be frugal with CBRNG increments.
-        //   See sample_indices_iid_uniform.
-    }
-    int64_t full_incr = num_mavec * incrs_per_mavec;
-    state.counter.incr(full_incr);
+    // 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);
     return state;
 }
 

RandBLAS/util.hh

@@ -513,35 +513,31 @@ state_t sample_indices_iid(int64_t n, const T* cdf, int64_t k, sint_t* samples,
     return state_t(ctr, key);
 }
 
+inline std::uint64_t promote_uint_pair(std::uint32_t a, std::uint32_t b) {
+    return static_cast<std::uint64_t>(a) + (static_cast<std::uint64_t>(b) << 32);
+}
+
 template <typename T, SignedInteger sint_t, bool WriteRademachers = true, typename state_t = RNGState<DefaultRNG>>
 state_t sample_indices_iid_uniform(int64_t n, int64_t k, sint_t* samples, T* rademachers, const state_t &state) {
-    using RNG = typename state_t::generator;
-    auto [ctr, key] = state;
-    RNG gen;
-    auto rv_array = r123ext::uneg11::generate(gen, ctr, key);
-    int64_t len_c = static_cast<int64_t>(state.len_c);
     if constexpr (WriteRademachers) {
-        len_c = 2*(len_c/2);
-        // ^ round down to the nearest multiple of two.
+        randblas_require(state.len_c >= 4);
+    } else {
+        randblas_require(state.len_c >= 2);
     }
-    int64_t rv_index = 0;
-    double dN = (double) n;
+    using RNG = typename state_t::generator;
+    RNG gen;
+    auto ctr = state.counter;
+    auto key = state.key;
+    std::uint64_t n_64 = static_cast<std::uint64_t>(n);
     for (int64_t i = 0; i < k; ++i) {
-        auto random_unif01 = uneg11_to_u01<double>(rv_array[rv_index]);
-        sint_t sample_index = (sint_t) dN * random_unif01;
-        samples[i] = sample_index;
-        rv_index += 1;
+        auto rv = gen(ctr, key);
+        ctr.incr();
+        std::uint64_t s = promote_uint_pair(rv[0], rv[1]);
+        samples[i] = static_cast<sint_t>(s % n_64);
         if constexpr (WriteRademachers) {
-            rademachers[i] = (rv_array[rv_index] >= 0) ? (T) 1 : (T) -1;
-            rv_index += 1;
-        }
-        if (rv_index == len_c) {
-            ctr.incr(1);
-            rv_array = r123ext::uneg11::generate(gen, ctr, key);
-            rv_index = 0;
+            rademachers[i] = (rv[2] % 2 == 0) ? (T) 1 : (T) -1;
         }
     }
-    if (0 < rv_index) ctr.incr(1);
     return state_t(ctr, key);
 }